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 batchSize,
int32_t deltaLen,
int32_t numKVHeads,
int32_t headDim,
int32_t maxSeqLen,
int32_t rotaryDim,
int32_t cosSinBatch,
int32_t cosSinSeqLen,
cudaStream_t stream,
)#

Launch the DFlash target KV cache update kernel.

Applies RoPE to k_delta and writes k_rope + v_delta into the combined KV cache 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 – [B, 2, numKVHeads, maxSeqLen, headDim] FP16 (in/out)

  • cosSinCache – [cosSinBatch, cosSinSeqLen, rotaryDim] FP32

  • deltaStartPositions – [B] INT32

  • batchSize – batch size

  • deltaLen – number of delta tokens per batch

  • numKVHeads – number of KV heads

  • headDim – head dimension

  • maxSeqLen – KV cache capacity (seq dim)

  • rotaryDim – rotary embedding dimension

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

  • cosSinSeqLen – cos/sin cache sequence length

  • stream – CUDA stream

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

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

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

  • 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 blockSize,
cudaStream_t stream,
)#

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

verifyTokenIds[b, 0] = lastAcceptedTokens[b], verifyTokenIds[b, j] = draftTokenIds[b, j] for j >= 1. 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, BS] INT32 — DFlash draft argmax token IDs

  • verifyTokenIds – [B, BS] INT32 — output base verify token IDs

  • verifyTreeMask – [B, BS, BS] INT8 — output EAGLE-style causal tree mask

  • batchSize – batch size

  • blockSize – DFlash block size

  • stream – CUDA stream