KV Cache Utils Kernels#

Warning

doxygenfunction: Unable to resolve function “trt_edgellm::kernel::incrementLengthTensor” with arguments None in doxygen xml output for project “TensorRT Edge-LLM” from directory: ../cpp_docs/xml. Potential matches: 

- void incrementLengthTensor(rt::Tensor &lengthTensor, int32_t increment, cudaStream_t stream)
- void incrementLengthTensor(rt::Tensor &lengthTensor, rt::Tensor const &newIncrementTensor, cudaStream_t stream)
void trt_edgellm::kernel::instantiateKVCacheLayerFromTensor(
rt::Tensor &dstKVCacheLayer,
rt::Tensor const &srcKVCacheTensor,
int32_t batchIdx,
cudaStream_t stream
)#

Single-layer variant: instantiate KV cache for one layer from a saved tensor.

Parameters:

  • dstKVCacheLayer – [inout] [maxBatchSize, 2, numKVHeads, maxSequenceLength, headDim]

  • srcKVCacheTensor – [in] [2, numKVHeads, sequenceLength, headDim]

  • batchIdx – [in] Target batch index in the destination buffer

  • stream – [in] CUDA stream

void trt_edgellm::kernel::saveKVCacheLayerIntoTensor(
rt::Tensor &dstKVCacheTensor,
rt::Tensor const &srcKVCacheLayer,
int32_t batchIdx,
cudaStream_t stream
)#

Single-layer variant: save KV cache for one layer into a tensor.

Parameters:

  • dstKVCacheTensor – [out] [2, numKVHeads, sequenceLength, headDim]

  • srcKVCacheLayer – [in] [maxBatchSize, 2, numKVHeads, maxSequenceLength, headDim]

  • batchIdx – [in] Source batch index in the buffer

  • stream – [in] CUDA stream

void trt_edgellm::kernel::saveKVCacheBatched(
KVLayerInfo const *srcLayerInfos,
KVLayerInfo const *dstLayerInfos,
int32_t numLayers,
int32_t headDim,
int32_t kvPoolPages,
int32_t batchIdx,
int32_t sequenceLength,
cudaStream_t stream
)#

Batched save: copy multiple layers’ KV cache into per-layer tensors in a single launch. All layers must share the same headDim. srcLayerInfos[i].data points to a two-pool NHD [2, maxBatch, capPadded, numKVHeads_i, headDim] pool; dstLayerInfos[i].data points to a [2, seqLen, numKVHeads_i, headDim] saved tensor (K plane then V plane).

Parameters:

  • srcLayerInfos – [numLayers] GPU array — source cache pools

  • dstLayerInfos – [numLayers] GPU array — destination saved tensors

  • numLayers – Number of layers in this batch

  • headDim – Head dimension (same for all layers)

  • kvPoolPages – Physical K-page count of the source cache (V-half offset = kvPoolPages*128*H*D)

  • batchIdx – Batch index to save from

  • sequenceLength – Number of tokens to copy

  • stream – CUDA stream

void trt_edgellm::kernel::instantiateKVCacheBatched(
KVLayerInfo const *dstLayerInfos,
KVLayerInfo const *srcLayerInfos,
int32_t numLayers,
int32_t headDim,
int32_t kvPoolPages,
int32_t batchIdx,
int32_t sequenceLength,
cudaStream_t stream
)#

Batched restore: load multiple layers’ KV cache from per-layer tensors in a single launch. All layers must share the same headDim. srcLayerInfos[i].data points to a [2, seqLen, numKVHeads_i, headDim] saved tensor (K plane then V plane); dstLayerInfos[i].data points to a two-pool NHD [2, maxBatch, capPadded, numKVHeads_i, headDim] pool.

Parameters:

  • dstLayerInfos – [numLayers] GPU array — destination cache pools

  • srcLayerInfos – [numLayers] GPU array — source saved tensors

  • numLayers – Number of layers in this batch

  • headDim – Head dimension (same for all layers)

  • kvPoolPages – Physical K-page count of the destination cache (V-half offset = kvPoolPages*128*H*D)

  • batchIdx – Batch index to restore into

  • sequenceLength – Number of tokens to copy

  • stream – CUDA stream

void trt_edgellm::kernel::gatherPagedKVToSplit(
void const *pool,
void *kDst,
void *vDst,
int32_t const *pageTable,
int32_t const *kvSeqLens,
int32_t maxPagesPerSeq,
int32_t batchSize,
int32_t seqLen,
int32_t numKVHeads,
int32_t headDim,
size_t elemSize,
bool dequantFp8,
float kScale,
float vScale,
cudaStream_t stream
)#

Gathers logical pages 0..ceil(seqLen/128) of every slot from a paged K/V page pool into dense split K/V workspaces, for FMHA-v2 FP8, padding, and vision-block consumers that require a contiguous [B, seqLen, H, D] FP16 view. The destination is ALWAYS FP16 (half): an FP8 pool is dequantized with the K/V scales, so downstream dataPointer<half>() consumers never reinterpret FP8 bytes as half.

pool is a single flat page array (the Task-1 [2, maxBatch, capPadded, H, D] allocation reinterpreted as pages); per KVPageTable’s convention, V page ids are always K page id + numPages, so both halves index directly into the same pool base &#8212; there is no separate V-half pointer.

Bad-page semantics: any page-table entry of -1 (unmapped) zero-fills its destination span, whether it lies beyond the slot’s live range (the legal padding tail) or inside it (an upstream table bug &#8212; the runtime guarantees mapped coverage for live positions, so an in-range -1 cannot occur by construction; if it ever does, the gather fails soft with zeros instead of reading a wild address).

Parameters:

  • pool – [in] Page pool, logically [2*numPages, 128, numKVHeads, headDim] (K pages first, V pages at page id +numPages per KVPageTable’s kernel view).

  • kDst – [out] Destination K tensor (FP16), [batchSize, seqLen, numKVHeads, headDim].

  • vDst – [out] Destination V tensor (FP16), [batchSize, seqLen, numKVHeads, headDim].

  • pageTable – [in] Device page table, [batchSize, 2, maxPagesPerSeq]; row 0 = K page ids, row 1 = V page ids. A negative id means “unallocated”.

  • kvSeqLens – [in] Device [batchSize] per-slot live KV length; distinguishes an in-range unmapped page (violation) from the legal padding tail (zero-fill).

  • maxPagesPerSeq – [in] Number of logical pages per slot in pageTable’s row stride.

  • batchSize – [in] Number of slots.

  • seqLen – [in] Destination token extent per slot (padded capacity spanned by the dst).

  • numKVHeads – [in] Number of KV heads.

  • headDim – [in] Head dimension.

  • elemSize – [in] Size in bytes of one POOL element (2 for FP16, 1 for FP8) &#8212; ignored when dequantFp8 is true (pool is FP8, dst is FP16).

  • dequantFp8 – [in] If true, treat the pool as FP8 e4m3 and dequantize to FP16 with kScale/vScale.

  • kScale – [in] K dequant scale (dequant value = fp8 * kScale). Unused when dequantFp8 false.

  • vScale – [in] V dequant scale. Unused when dequantFp8 false.

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

Throws:

std::runtime_error – if ceil(seqLen/128) exceeds maxPagesPerSeq.