Dflash Runtime Kernels#

void trt_edgellm::kernel::launchDFlashTargetKVCacheUpdate(
half const *kDelta,
half const *vDelta,
half *kvCache,
float const *cosSinCache,
int32_t const *deltaStartPositions,
int32_t const *deltaLengths,
int32_t const *pageTable,
int32_t batchSize,
int32_t deltaLen,
int32_t numKVHeads,
int32_t headDim,
int32_t rotaryDim,
int32_t cosSinBatch,
int32_t cosSinSeqLen,
int32_t numPages,
int32_t maxPagesPerSeq,
cudaStream_t stream
)#

Launch the DFlash target KV cache update kernel.

Applies RoPE to k_delta and writes k_rope + v_delta into the paged KV pool at positions [deltaStart, deltaStart + deltaLen) for each batch element.

Parameters:

  • kDelta – [B, deltaLen, numKVHeads, headDim] FP16, k_normed, no RoPE

  • vDelta – [B, deltaLen, numKVHeads, headDim] FP16

  • kvCache – Paged KV pool [2, numPages, kTOKENS_PER_PAGE, numKVHeads, headDim] FP16 (in/out).

  • cosSinCache – [cosSinBatch, cosSinSeqLen, rotaryDim] FP32

  • deltaStartPositions – [B] INT32

  • 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.

  • batchSize – ACTIVE batch size

  • deltaLen – number of delta tokens per batch

  • numKVHeads – number of KV heads

  • headDim – head dimension

  • rotaryDim – rotary embedding dimension

  • cosSinBatch – cos/sin cache batch size (1 or B)

  • cosSinSeqLen – cos/sin cache sequence length

  • numPages – Number of physical pages in each KV plane

  • maxPagesPerSeq – Logical pages per sequence; positions outside this capacity are skipped

  • stream – CUDA stream

  • deltaLengths – [B] INT32, per-batch delta lengths (skip t >= deltaLengths[b])

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,
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::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::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