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 rotaryDim,
int32_t cosSinBatch,
int32_t cosSinSeqLen,
int32_t maxBatch,
int32_t kvCapacity,
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 – Two-pool NHD KV pool [2, maxBatch, kvCapacity, numKVHeads, headDim] FP16 (in/out). DFlash is identity-only (it opts out of KV-cache reuse), so this writes directly at the request’s own contiguous slot — no page table needed.

  • cosSinCache – [cosSinBatch, cosSinSeqLen, rotaryDim] FP32

  • deltaStartPositions – [B] INT32

  • batchSize – ACTIVE batch size (number of requests to process this call; <= maxBatch)

  • 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

  • maxBatch – ALLOCATION batch (outer dim of each K/V half in kvCache); sizes the V-pool offset (= maxBatch*kvCapacity*numKVHeads*headDim). NOT the active batchSize.

  • kvCapacity – ALLOCATION per-slot token capacity (capPadded); sizes each request’s slot stride and doubles as the OOB write guard (positions >= kvCapacity are dropped).

  • stream – CUDA stream

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

void trt_edgellm::kernel::checkDFlashPageTableIdentity(
int32_t const *hostKRow,
int32_t slot,
int32_t maxPagesPerSeq
)#

Assert that a page table’s K-half row for slot is the static identity range [slot*maxPagesPerSeq, (slot+1)*maxPagesPerSeq) that DFlash’s contiguous target-KV update assumes. DFlash is identity-only (it opts out of KV-cache reuse, so its update never resolves token offsets through the page table); this makes that assumption an explicit, checkable guard instead of relying on a runtime-level opt-out to keep DFlash slots unmapped.

Parameters:

  • hostKRow – [maxPagesPerSeq] host K page ids for slot (e.g. KVPageTable::hostRow(slot))

  • slot – Batch slot whose row is being checked

  • maxPagesPerSeq – Number of logical pages per slot (row length)

Throws:

std::runtime_error – if any entry deviates from the identity range

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 KV pool’s PADDED per-slot capacity (capPadded, a multiple of kTOKENS_PER_PAGE); the real, configured maximum sequence length is <= kvCapacity (padding only rounds up). The RoPE cache is sized to that real (unpadded) maximum directly (see RopeCache::getOrCreate), so cosSinSeqLen can legitimately be smaller than kvCapacity whenever the configured capacity isn’t already page-aligned (e.g. maxKVCacheCapacity=4000 -> kvCapacity=4096, cosSinSeqLen=4000) — comparing cosSinSeqLen < kvCapacity (the pre-fix check) therefore false-positives on any non-page-aligned config. The correct invariant is the other direction: cosSinSeqLen must never EXCEED kvCapacity, since kvCapacity is provably an upper bound on every real position (padding never shrinks capacity).

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