Dflash Runtime Kernels#

void trt_edgellm::kernel::launchDFlashTargetKVCacheUpdate( half const *kDelta, half const *vDelta, half *kvCache, float const *tokenAlignedCosSin, int32_t const *deltaPositions, int32_t const *deltaTokenToSequence, int32_t const *pageTable, int32_t numDeltaTokens, int32_t batchSize, int32_t numKVHeads, int32_t headDim, int32_t rotaryDim, int32_t numPages, int32_t maxPagesPerSeq, cudaStream_t stream )#

Launch the DFlash target KV cache update kernel.

Applies RoPE to token-major K deltas and writes K/V deltas into their owning sequences’ paged KV cache.

Parameters:

kDelta – [numDeltaTokens, numKVHeads, headDim] FP16, k_normed, no RoPE
vDelta – [numDeltaTokens, numKVHeads, headDim] FP16
kvCache – Paged KV pool [2, numPages, kTOKENS_PER_PAGE, numKVHeads, headDim] FP16 (in/out).
tokenAlignedCosSin – [numDeltaTokens, rotaryDim] FP32
deltaPositions – [numDeltaTokens] absolute KV positions; negative values are padding
deltaTokenToSequence – [numDeltaTokens] owning sequence indices; negative values are padding
pageTable – [B, 2, maxPagesPerSeq] canonical page ids: K in [0, numPages), V in [numPages, 2 * numPages). Unmapped or out-of-plane ids skip that cache plane.
numDeltaTokens – Physical token-major row count
batchSize – ACTIVE batch size
numKVHeads – number of KV heads
headDim – head dimension
rotaryDim – rotary embedding dimension
numPages – Number of physical pages in each KV plane
maxPagesPerSeq – Logical pages per sequence; positions outside this capacity are skipped
stream – CUDA stream

void trt_edgellm::kernel::checkDFlashRopeCapacity( int32_t cosSinSeqLen, int32_t kvCapacity )#

Validate that a RoPE cos/sin cache covers every position DFlash’s target-KV update can write.

kvCapacity is the page-aligned upper bound on every writable position. The RoPE cache tracks the unpadded configured maximum and may therefore be smaller; reject only cosSinSeqLen > kvCapacity.

Parameters:

cosSinSeqLen – Sequence length of the bound rope_cos_sin cache
kvCapacity – KV pool’s padded per-slot capacity (capPadded)

Throws:

std::runtime_error – if cosSinSeqLen > kvCapacity

void trt_edgellm::kernel::launchDFlashPrepareProposalInputs( int32_t const *oldDraftCacheLengths, int32_t const *deltaLengths, int32_t blockSize, int32_t *packedAttentionMask, int32_t *attentionPosId, int32_t *contextLengths, int32_t *positions, int32_t *queryStartOffsets, int32_t *queryLengths, int32_t *pastLengths, int32_t *attentionSequenceLengths, int32_t const *stateIndices, bool causalProposalMask, int32_t batchSize, cudaStream_t stream )#

Launch kernel to prepare DFlash proposal attention inputs.

Computes target_len_after_delta = oldDraftCacheLengths[b] + deltaLen, then sets: attention_pos_id[b, i] = target_len_after_delta + i context_lengths[b] = target_len_after_delta + blockSize packed_attention_mask: full non-causal within proposal block, or causal rows when requested

Parameters:

oldDraftCacheLengths – [B] INT32 — draft cache lengths BEFORE delta (GPU)
deltaLengths – [B] INT32 — per-batch delta token count (GPU)
blockSize – DFlash block size (BS)
packedAttentionMask – [B, BS, divUp(BS,32)] INT32 — output
attentionPosId – [B, BS] INT32 — output
contextLengths – [B] INT32 — output
causalProposalMask – true => row i attends only to proposal positions [0, i]
batchSize – batch size
stream – CUDA stream

void trt_edgellm::kernel::launchDFlashPrepareDeltaMetadata( int32_t const *oldDraftCacheLengths, int32_t const *deltaLengths, int32_t deltaWidth, int32_t *deltaPositions, int32_t *deltaTokenToSequence, int32_t batchSize, cudaStream_t stream )#

void trt_edgellm::kernel::launchDFlashGatherDeltaRope( float const *source, float *output, int32_t const *deltaPositions, int32_t const *deltaTokenToSequence, int32_t const *stateIndices, int32_t numDeltaTokens, int32_t batchSize, int32_t sourceRows, int32_t cacheCapacity, int32_t rotaryDim, cudaStream_t stream )#

void trt_edgellm::kernel::launchPrepareSpecRaggedMetadata( int32_t const *attentionPositions, int32_t const *committedPastLengths, int32_t const *validCounts, int32_t queryWidth, int32_t *positions, int32_t *queryStartOffsets, int32_t *queryLengths, int32_t *pastLengths, int32_t *attentionSequenceLengths, int32_t *treeParentIds, int32_t *treeDepths, bool synthesizeLinearTree, int32_t batchSize, cudaStream_t stream )#

void trt_edgellm::kernel::launchDFlashPrepareBaseVerifyInputs( int32_t const *baseKVCacheLengths, int32_t verifySize, int32_t *packedAttentionMask, int32_t *attentionPosId, int64_t *selectTokenIndices, int32_t *contextLengths, int32_t batchSize, cudaStream_t stream )#

Launch kernel to prepare DFlash base verification attention inputs.

DFlash verifies a linear block, so the base tree mask is always causal: token i attends to proposal tokens [0, i]. This writes the packed INT32 mask consumed by AttentionPlugin directly, without materializing an intermediate unpacked [B, BS, BS] INT8 mask.

Parameters:

baseKVCacheLengths – [B] INT32 — committed base cache lengths (GPU)
verifySize – DFlash verify block size (BS)
packedAttentionMask – [B, BS, divUp(BS,32)] INT32 — output
attentionPosId – [B, BS] INT32 — output
selectTokenIndices – [B, BS] INT64 — output
contextLengths – [B] INT32 — output
batchSize – batch size
stream – CUDA stream

void trt_edgellm::kernel::launchDFlashBuildLinearTreeMetadata( int32_t *treeParentIds, int32_t *treeDepths, int32_t batchSize, int32_t verifySize, cudaStream_t stream )#: Populate tree metadata for a linear verification path. Node zero is the root; every later node has the previous node as its parent.

void trt_edgellm::kernel::launchDFlashBuildLinearVerifyInputs( int32_t const *lastAcceptedTokens, int32_t const *draftTokenIds, int32_t *verifyTokenIds, int8_t *verifyTreeMask, int32_t batchSize, int32_t proposalLen, int32_t draftTokenStride, int32_t verifySize, cudaStream_t stream )#

Launch kernel to build DFlash linear verification inputs for EAGLE accept.

verifyTokenIds[b, 0] = lastAcceptedTokens[b]; verifyTokenIds[b, j + 1] = draftTokenIds[b, j], j in [0, proposalLen). verifyTreeMask is an unpacked causal tree mask where row i attends to [0, i].

Parameters:

lastAcceptedTokens – [B] INT32 — last committed token per batch
draftTokenIds – [B, draftTokenStride] INT32 — DFlash draft argmax token IDs
verifyTokenIds – [B, verifySize] INT32 — output base verify token IDs
verifyTreeMask – [B, verifySize, verifySize] INT8 — output EAGLE-style causal tree mask
batchSize – batch size
proposalLen – number of proposal tokens copied after the anchor
draftTokenStride – row stride of draftTokenIds, usually dflashBlockSize
verifySize – base verification input size, must equal proposalLen + 1
stream – CUDA stream