Apply Rope Write KV#

void trt_edgellm::kernel::launchApplyRopeWriteKV(
rt::Tensor const &cosSinCache,
rt::OptionalInputTensor kvCacheEndLens,
rt::Tensor &q,
rt::Tensor &k,
rt::Tensor const &v,
rt::Tensor &kvCache,
float kScale,
float vScale,
cudaStream_t stream,
bool writeKInPlace,
int32_t const *pageTable,
int32_t maxPagesPerSeq
)#

Launch kernel to apply RoPE positional encoding to Q/K and write K/V to KVCache.

Paged-KV bad-page semantics: the write kernels resolve each token’s physical page through pageTable; a negative page id (kBAD_PAGE_INDEX) means the position is unmapped and the write is silently skipped (same contract as other paged-KV runtimes, e.g. vLLM’s slot_mapping < 0). The runtime is responsible for never presenting an unmapped page for a live position: identity tables cover every slot by construction, and reuse-mode table builders enforce coverage on the host at build time.

Parameters:

  • cosSinCache – [in] FP32 type tensor with layout of [cosSinCacheBatchSize, cosSinCacheSeqLen, rotaryDim]

  • kvCacheEndLens – [in] Optional INT32 type tensor with layout of [batchSize], the end position of KVCache after writing. When nullopt, KVCache is written from the start (prefill without prior cache).

  • q – [inout] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hq, headDim]

  • k – [inout] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hkv, headDim]

  • v – [in] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hkv, headDim]

  • kvCache – [out] FP16/FP8 type tensor with layout of [batchSize, 2, Hkv, kvCacheCapacity, headDim]

  • kScale – [in] K dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • vScale – [in] V dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • stream – [in] CUDA stream to launch the kernel

  • writeKInPlace – [in] Controls whether roped K is additionally written back to the K tensor in-place, on top of always being written to kvCache. Set to true for the initial prefill path (SEPARATE_Q_K_V) where the downstream FMHA kernel reads Q, K, V as separate contiguous tensors rather than from the KV cache. In this case K must contain the roped result. Set to false (default) for chunked prefill with KV cache reuse, where FMHA reads KV from the transposed KV cache, and for all decoding paths (vanilla / tree), where the XQA kernel reads KV from the cache.

  • pageTable – [in] INT32 device page table [batchSize, 2, maxPagesPerSeq]; row b*2+0 carries K page ids, row b*2+1 the derived V page ids. When non-null, token t of slot b writes K to kvCache[pageTable[(b*2+0)*maxPagesPerSeq + t/128]][t%128][h][d] and V via the row’s V half (both addressing the same flat kvCache buffer, reinterpreted as [nPages, 128, Hkv, D]). A negative entry skips the write for that half. Pass nullptr to use the legacy [B, 2, Hkv, S, D] addressing instead.

  • maxPagesPerSeq – [in] Page-table inner dimension (number of page slots per sequence). Only used when pageTable != nullptr. the check (e.g. in unit tests that don’t care).

Throws:

std::runtime_error – if tensor shape or data type is incorrect

void trt_edgellm::kernel::launchApplyRopeWriteKVTreeDecoding(
rt::Tensor const &cosSinCache,
rt::Tensor const &kvCacheEndLens,
rt::Tensor const &tokenPosIds,
rt::Tensor &q,
rt::Tensor &k,
rt::Tensor const &v,
rt::Tensor &kvCache,
float kScale,
float vScale,
cudaStream_t stream,
int32_t const *pageTable,
int32_t maxPagesPerSeq
)#

Launch the kernel when we are performing tree attention for speculative decoding.

Note

We won’t overwrite K/V tensor in this case but we use Tensor& signature to reduce duplicate code.

Parameters:

  • cosSinCache – [in] FP32 type tensor with layout of [cosSinCacheBatchSize, cosSinCacheSeqLen, rotaryDim]

  • kvCacheEndLens – [in] INT32 type tensor with layout of [batchSize], the end position of KVCache after writing.

  • tokenPosIds – [in] INT32 type tensor with layout of [batchSize, runtimeSeqLen], the position of token within sequence.

  • q – [inout] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hq, headDim]

  • k – [in] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hkv, headDim]

  • v – [in] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hkv, headDim]

  • kvCache – [out] FP16/FP8 type tensor with layout of [batchSize, 2, Hkv, kvCacheCapacity, headDim], write KVCache from the end position.

  • kScale – [in] K dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • vScale – [in] V dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • stream – [in] CUDA stream to launch the kernel

  • pageTable – [in] Optional INT32 device page table [batchSize, 2, maxPagesPerSeq]. See launchApplyRopeWriteKV.

  • maxPagesPerSeq – [in] Page-table inner dimension. Only used when pageTable != nullptr.

Throws:

std::runtime_error – if tensor shape or data type is incorrect

void trt_edgellm::kernel::launchApplyRopeWriteKVSplitQKV(
rt::Tensor const &cosSinCache,
rt::Tensor const &kvCacheEndLens,
rt::Tensor &q,
rt::Tensor const &k,
rt::Tensor const &v,
rt::Tensor &kvCache,
float kScale,
float vScale,
cudaStream_t stream,
int32_t const *pageTable,
int32_t maxPagesPerSeq,
void *fp8QOut = nullptr,
float qScale = 1.0f
)#

Launch kernel to apply RoPE to Q, apply RoPE to K and write K/V to KVCache.

Optimized for the CuTe DSL FMHA path: applies RoPE to Q, writes roped K and V into KV cache [B, 2, H_kv, S, D]. Does NOT write roped K back to the K input tensor.

When fp8QOut is non-null (FP8 KV cache path), the roped Q is quantized to FP8 and written to the provided output buffer. The original FP16 Q tensor is NOT modified. The downstream FP8 FMHA kernel reads Q from fp8QOut and K/V from the KV cache directly.

When fp8QOut is null (FP16 path), RoPE is applied to Q in-place in the FP16 Q tensor.

Parameters:

  • cosSinCache – [in] FP32 type tensor with layout of [cosSinCacheBatchSize, cosSinCacheSeqLen, rotaryDim]

  • kvCacheEndLens – [in] INT32 type tensor with layout of [batchSize], the end position of KVCache after writing.

  • q – [inout] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hq, headDim]. RoPE applied in-place when fp8QOut is null.

  • k – [in] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hkv, headDim]

  • v – [in] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hkv, headDim]

  • kvCache – [out] FP16/FP8 type tensor with layout of [batchSize, 2, Hkv, kvCacheCapacity, headDim]

  • kScale – [in] K dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • vScale – [in] V dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • stream – [in] CUDA stream to launch the kernel

  • pageTable – [in] Optional INT32 device page table [batchSize, 2, maxPagesPerSeq]. See launchApplyRopeWriteKV.

  • maxPagesPerSeq – [in] Page-table inner dimension. Only used when pageTable != nullptr.

  • fp8QOut – [out] Optional FP8 output buffer for roped Q [batchSize, runtimeSeqLen, Hq, headDim]. When non-null, roped Q is quantized to FP8 E4M3 and stored here. Pass nullptr for FP16 in-place RoPE.

  • qScale – [in] Q dequant scale (quant→orig). Only used when fp8QOut is non-null.

void trt_edgellm::kernel::launchApplyRopeFromPackedToSplit(
rt::Tensor const &cosSinCache,
rt::OptionalInputTensor kvCacheEndLens,
rt::OptionalInputTensor tokenPosIds,
rt::Tensor const &packedQKV,
rt::Tensor &qScratch,
rt::Tensor &kvCache,
float kScale,
float vScale,
cudaStream_t stream,
int32_t const *pageTable,
int32_t maxPagesPerSeq,
void *kScratchOut = nullptr,
void *vScratchOut = nullptr,
void *fp8QOut = nullptr,
float qScale = 1.0f,
half const *qNormGamma = nullptr,
half const *kNormGamma = nullptr,
float rmsNormEps = 1e-6f,
rt::OptionalInputTensor cuQSeqLens = std::nullopt
)#

Launch kernel to read a packed QKV tensor, apply RoPE to Q and K, write roped Q to a split scratch tensor, and always write roped K and V to KVCache. Optionally also mirrors roped K and V to separate scratch tensors for the SEPARATE_Q_K_V FMHA path.

Packed-input variant of launchApplyRopeWriteKV — one fused QKV tensor in:

  • NORMAL_PREFILL (SEPARATE_Q_K_V FMHA): pass non-null kScratchOut / vScratchOut.

  • CHUNKED_PREFILL / decode: pass nullptr — K/V are read back from the cache.

  • Tree decoding: pass tokenPosIds (-1 = padding token, no cache write).

  • CuTeDSL + FP8: pass fp8QOut for FP8 roped Q; otherwise qScratch gets FP16 Q.

Parameters:

  • cosSinCache – [in] FP32 tensor [cosSinCacheBatchSize, cosSinCacheSeqLen, rotaryDim]

  • kvCacheEndLens – [in] Optional INT32 tensor [batchSize] — KV cache end position after insertion. Pass nullopt for prefill without prior cache (starts at position 0).

  • tokenPosIds – [in] Optional INT32 tensor [batchSize, runtimeSeqLen] for tree decoding. Position -1 marks padding tokens whose Q is zeroed and K/V writes are skipped.

  • packedQKV – [in] FP16 tensor [batchSize, runtimeSeqLen, Hq+2*Hkv, headDim], read-only.

  • qScratch – [out] FP16 tensor [batchSize, runtimeSeqLen, Hq, headDim] — roped Q output (unless fp8QOut is non-null, in which case this is unused).

  • kvCache – [out] FP16/FP8 tensor [batchSize, 2, Hkv, kvCacheCapacity, headDim] — K/V written here.

  • kScale – [in] K dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • vScale – [in] V dequant scale (quant→orig). Use 1.0f for FP16 KV cache.

  • stream – [in] CUDA stream.

  • pageTable – [in] Optional INT32 device page table [batchSize, 2, maxPagesPerSeq]. See launchApplyRopeWriteKV. Pass nullptr for the legacy [B, 2, Hkv, S, D] addressing.

  • maxPagesPerSeq – [in] Page-table inner dimension. Only used when pageTable != nullptr.

  • kScratchOut – [out] Optional FP16 buffer [batchSize, runtimeSeqLen, Hkv, headDim] — mirrored roped K. Pass nullptr if downstream does not need scratch K.

  • vScratchOut – [out] Optional FP16 buffer [batchSize, runtimeSeqLen, Hkv, headDim] — mirrored V. Pass nullptr if downstream does not need scratch V.

  • fp8QOut – [out] Optional FP8 buffer [batchSize, runtimeSeqLen, Hq, headDim] — FP8-quantized roped Q. Pass nullptr for FP16 Q via qScratch.

  • qScale – [in] Q dequant scale (quant→orig). Only used when fp8QOut is non-null.

  • qNormGamma – [in] Optional FP16 device pointer [headDim] for per-head RMSNorm gamma applied to Q BEFORE RoPE. When non-null, qk_norm is computed inside this kernel via warp-shuffle reduction across the headDim/vec_size threads of blockDim.x.

  • kNormGamma – [in] Optional FP16 device pointer [headDim] for per-head RMSNorm gamma applied to K BEFORE RoPE. Same conventions as qNormGamma. V is never RMSNormed.

  • rmsNormEps – [in] Epsilon for the RMSNorm formula. Ignored when both gamma pointers are null.

  • cuQSeqLens – [in] Optional INT32 tensor [batchSize + 1] carrying actual cumulative Q lengths for ragged prefill. Rows at or beyond the actual per-batch length have Q zeroed and skip all K/V writes.

Throws:

std::runtime_error – if tensor shape or data type is incorrect.

void trt_edgellm::kernel::launchApplyRopeQOnly(
rt::Tensor const &cosSinCache,
rt::Tensor const &kvCacheEndLens,
rt::Tensor &q,
cudaStream_t stream
)#

Launch kernel to apply RoPE to Q only (no KV write).

Used for shared-KV layers where Q still needs positional encoding but the KV cache belongs to a donor layer and must not be modified.

Parameters:

  • cosSinCache – [in] FP32 type tensor with layout of [cosSinCacheBatchSize, cosSinCacheSeqLen, rotaryDim]

  • kvCacheEndLens – [in] INT32 type tensor with layout of [batchSize], used to compute RoPE position.

  • q – [inout] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hq, headDim]. RoPE applied in-place.

  • stream – [in] CUDA stream to launch the kernel

void trt_edgellm::kernel::launchApplyRopeQOnlyTreeDecoding(
rt::Tensor const &cosSinCache,
rt::Tensor const &tokenPosIds,
rt::Tensor &q,
cudaStream_t stream
)#

Launch kernel to apply RoPE to Q only, using per-token position IDs (tree decoding).

For shared-KV layers during tree/speculative decoding, each candidate token has its own position in the tree. RoPE is applied to Q using these explicit position IDs. No KV cache write is performed (the donor layer’s cache is already populated).

Parameters:

  • cosSinCache – [in] FP32 type tensor with layout of [cosSinCacheBatchSize, cosSinCacheSeqLen, rotaryDim]

  • tokenPosIds – [in] INT32 type tensor with layout of [batchSize, runtimeSeqLen], per-token position IDs.

  • q – [inout] FP16 type tensor with layout of [batchSize, runtimeSeqLen, Hq, headDim]. RoPE applied in-place.

  • stream – [in] CUDA stream to launch the kernel