Batch Evict Kernels#

struct KVLayerInfo#

Per-layer KV cache metadata for batched kernel operations.

Public Members

void *data#

Base of this layer’s physical K-then-V page-pool allocation.

int32_t numKVHeads#

Number of KV heads for this layer.

int32_t maxSeqLen#

Per-row token capacity of this layer’s pool (capPadded)

int32_t maxBatch#

Allocation batch (outer dim of each K/V half); needed to compute the V-half offset (= maxBatch*maxSeqLen*H*D) for kernels that address a single KVLayerInfo without a separately-passed maxBatchSize parameter.

void trt_edgellm::kernel::compactKVCacheSingleLayer(
rt::Tensor &kvCacheLayer,
rt::Tensor const &batchMapping,
rt::Tensor const &kvCacheLengths,
rt::Tensor &dstKVCacheLengths,
int32_t oldActiveBatch,
int32_t newActiveBatch,
bool updateLengths,
cudaStream_t stream
)#

Compact a single layer’s KV cache by removing evicted batches.

Single-layer variant of compactKVCache for per-layer heterogeneous KV cache.

Parameters:

  • kvCacheLayer – [maxBatch, 2, numKVHeads, maxSeq, headDim] single-layer buffer (in/out)

  • batchMapping – [oldActiveBatch] GPU tensor, mapping[i] = newBatchIdx or -1 (evict)

  • kvCacheLengths – [maxBatch] GPU tensor of sequence lengths (const input)

  • dstKVCacheLengths – [maxBatch] GPU tensor for compacted lengths (output, may alias kvCacheLengths)

  • oldActiveBatch – Number of batches before eviction

  • newActiveBatch – Number of batches after eviction

  • updateLengths – If true, update dstKVCacheLengths (only first layer should do this)

  • stream – CUDA stream

void trt_edgellm::kernel::compactTensorBatch(
rt::Tensor const &src,
rt::Tensor const &batchMapping,
rt::Tensor &dst,
int32_t oldActiveBatch,
int32_t newActiveBatch,
cudaStream_t stream
)#

Generic tensor compaction along batch dimension.

This kernel compacts a tensor by removing evicted batches.

Note

Assumes batch dimension is the first dimension (dim 0)

Note

For in-place operation, pass the same tensor as both src and dst

Parameters:

  • src – Source tensor (const input)

  • batchMapping – [oldActiveBatch] GPU tensor (const input), mapping[i] = newBatchIdx or -1

  • dst – Destination tensor (output, can be same as src for in-place operation)

  • oldActiveBatch – Number of batches before eviction

  • newActiveBatch – Number of batches after eviction

  • stream – CUDA stream

Throws:

std::runtime_error – if tensors are not located on the GPU, or tensor shapes are invalid

void trt_edgellm::kernel::compactKVCacheBatched(
KVLayerInfo const *layerInfos,
rt::Tensor const &batchMapping,
rt::Tensor const &liveLengths,
int32_t numLayers,
int32_t headDim,
int32_t kvPoolPages,
nvinfer1::DataType kvCacheType,
int32_t oldActiveBatch,
int32_t newActiveBatch,
cudaStream_t stream
)#

Batched in-place KV pool compaction across a headDim group, moving only live prefixes.

One grouped launch covers every layer in layerInfos (all sharing headDim), K and V halves of the NHD pool included. For each moved row only the contiguous live prefix (liveLengths[oldBatchIdx] * numKVHeads * headDim elements) is copied with vectorized loads/stores; padding beyond the live length is left untouched. Scheduled CTAs are proportional to the number of layers (not the allocated capacity), and identity (oldActiveBatch == newActiveBatch) or all-evicted (newActiveBatch == 0) calls return without launching.

Parameters:

  • layerInfos – [numLayers] GPU array of KVLayerInfo for one headDim group

  • batchMapping – [oldActiveBatch] GPU tensor (const input), mapping[i] = newBatchIdx or -1

  • liveLengths – [oldActiveBatch] GPU INT32 tensor (const input), live token length per old batch slot

  • numLayers – Number of layers in this group

  • headDim – Head dimension shared by all layers in this group

  • kvPoolPages – Physical K-page count; determines the V-half offset

  • kvCacheType – KV pool dtype (kHALF or kFP8); selects the copy element type

  • oldActiveBatch – Number of batches before eviction

  • newActiveBatch – Number of batches after eviction

  • stream – CUDA stream

Throws:

std::invalid_argument – for unsupported kvCacheType