Ddtree Kernels#

struct DDTreeBuildInputs#

Tensor inputs consumed by ddtreeBuild().

Public Members

rt::Tensor const &draftLogits#: [batch, dflashBlockSize, vocabSize] draft logits.

rt::Tensor const &rootTokenIds#: [batch] last accepted token ids.

rt::Tensor const &baseLengths#: [batch] committed base lengths before verification.

rt::Tensor const *draftVocabMappingTable = {nullptr}#: Optional [vocabSize] reduced-to-full draft vocab mapping. When present, DDTree selects candidates in the reduced draft vocab space but emits full-vocab token ids for base verify and accept.

struct DDTreeBuildOutputs#

Tensor outputs produced by ddtreeBuild().

Public Members

rt::Tensor &nodeTokenIds#: [batch, verifySize] flattened tree token ids.

rt::Tensor &nodeDepths#: [batch, verifySize] flattened tree depths.

rt::Tensor &parentIds#: [batch, verifySize] parent node indices.

rt::Tensor &nodeScores#: [batch, verifySize] prefix log-prob scores.

rt::Tensor &validCounts#: [batch] valid node counts.

rt::Tensor &verifyTokenIds#: [batch, verifySize] base verify token ids.

rt::Tensor &verifyPositionIds#: [batch, verifySize] base verify position ids.

rt::Tensor &packedAncestorMask#: [batch, verifySize, ceil(verifySize / 32)] packed mask.

rt::Tensor &ancestorMask#: [batch, verifySize, verifySize] unpacked accept mask.

rt::Tensor &contextLengths#: [batch] base context lengths for verify.

rt::Tensor &selectTokenIndices#: [batch, verifySize] selected token indices.

struct DDTreeBuildParams#

Parameters for ddtreeBuild().

Public Members

DDTreeBuildInputs inputs#: Required input tensors.

DDTreeBuildOutputs outputs#: Required output tensors.

int32_t candidateTopK#: Per-depth candidateTopK from DFlash draftingTopK.

void *workspace#: Temporary workspace from getDDTreeBuildWorkspaceSize().

size_t workspaceSize#: Workspace size in bytes.

cudaStream_t stream#: CUDA stream for kernel launches.

size_t trt_edgellm::kernel::getDDTreeBuildWorkspaceSize( int32_t batchSize, int32_t dflashBlockSize, int32_t verifySize, int32_t vocabSize, int32_t candidateTopK, )#

Returns temporary workspace size required by ddtreeBuild().

The workspace stores candidate token ids and log probabilities for DFlash candidateTopK selection.

void trt_edgellm::kernel::ddtreeBuild( DDTreeBuildParams const &params, )#

Build a prefix-closed DDTree from one DFlash draft logits pass.

The tree is flattened in score-prioritized order. Node 0 is the committed root token. For each proposal depth, the builder first takes the top candidateTopK tokens from the corresponding DFlash draft-logits row, then repeatedly appends the highest-scoring available child whose parent is already in the tree. This keeps the verify tree prefix-closed: every emitted node can trace a valid path back to the root.

Example with candidateTopK = 2 and verifySize = 6:

node:        0     1     2     3     4     5
depth:       0     1     1     2     2     3
parent:     -1     0     0     1     2     3

In this example, nodes 1 and 2 are the best depth-1 children of the root. Nodes 3 and 4 extend different parents, while node 5 extends node 3. The packed/unpacked ancestor masks allow each verify node to attend only to the tokens on its own root-to-node path.

Inputs: draftLogits [GPU, Float]: [batch, dflashBlockSize, vocabSize]. rootTokenIds [GPU, Int32]: last accepted token for each batch, [batch]. baseLengths [GPU, Int32]: committed base length before verify, [batch]. candidateTopK: DFlash DDTree candidateTopK, wired from draftingTopK. draftVocabMappingTable: optional reduced-to-full draft vocab mapping, [vocabSize]. workspace: temporary storage with size from getDDTreeBuildWorkspaceSize().

Outputs: nodeTokenIds [GPU, Int32]: flattened tree token ids, [batch, verifySize]. Node 0 is root. nodeDepths [GPU, Int32]: flattened tree depths, [batch, verifySize]. Root depth is 0. parentIds [GPU, Int32]: flattened parent node indices, [batch, verifySize]. Root and padding use -1. nodeScores [GPU, Float]: prefix log-prob scores, [batch, verifySize]. Padding uses -inf. validCounts [GPU, Int32]: valid node count per batch, [batch]. verifyTokenIds [GPU, Int32]: base verify token ids, [batch, verifySize]. verifyPositionIds [GPU, Int32]: base verify position ids, baseLengths + nodeDepths, [batch, verifySize]. packedAncestorMask [GPU, Int32]: packed tree attention mask, [batch, verifySize, ceil(verifySize / 32)]. ancestorMask [GPU, Int8]: unpacked EAGLE-style tree attention mask, [batch, verifySize, verifySize]. contextLengths [GPU, Int32]: baseLengths + verifySize, [batch]. selectTokenIndices [GPU, Int64]: [0, 1, …, verifySize - 1] per batch, [batch, verifySize].

The DDTree builder runs on GPU with fixed-size scratch. It does not provide a CPU runtime fallback.