KV Cache Manager

class KVCacheManager

Per-layer KV cache manager that supports heterogeneous head configurations across layers. Each attention layer gets its own independently-sized tensor in slot-shaped NHD layout with shape [2, maxBatchSize, capPadded, numKVHeads_i, headDim_i], where the K/V split (the leading 2) is OUTERMOST and capPadded = ceil(maxSequenceLength / kTOKENS_PER_PAGE) * kTOKENS_PER_PAGE. This same buffer is a bind-time reinterpretation of a page pool [2, numPages, kTOKENS_PER_PAGE, numKVHeads_i, headDim_i]; see kPoolPtr()/vPoolPtr(). This replaces the monolithic LinearKVCache allocation when layers have different numKVHeads or headDim.

Public Functions

KVCacheManager() noexcept = default

Default constructor.

KVCacheManager(Config const &config, cudaStream_t stream)

Construct and initialize per-layer KV cache.

Allocates one device tensor per attention layer. Once allocated, memory won’t be reallocated. Determines whether all layers share the same numKVHeads and headDim (uniform mode).

Parameters:

• config – Cache configuration with per-layer configs

• stream – CUDA stream for allocation

Throws:

std::runtime_error – if config is invalid or data type is unsupported

~KVCacheManager() noexcept

Destructor.

KVCacheManager(KVCacheManager const&) = delete

Deleted copy constructor to avoid large data copy.

KVCacheManager &operator=(KVCacheManager const&) = delete

Deleted copy assignment to avoid large data copy.

Returns:

Reference to this

KVCacheManager(KVCacheManager&&) noexcept

Move constructor.

KVCacheManager &operator=(KVCacheManager&&) noexcept

Move assignment operator.

Returns:

Reference to this

rt::Tensor &getCombinedKVCache(int32_t attnLayerIdx)

Get the combined KVCache for the given attention layer.

Note

This slot-shaped view is available only when numPages() equals the minimum active pages. With extra retained pages, its declared V-half stride would point into the K pool; use getCombinedKVCachePoolView(), getSeparateKVCache(), kPoolPtr(), or vPoolPtr().

Parameters:

attnLayerIdx – The index of the attention layer.

Returns:

A reference to the tensor with shape [2, maxBatchSize, capPadded, numKVHeads_i, headDim_i].

rt::Tensor &getCombinedKVCachePoolView(

int32_t attnLayerIdx

) noexcept

Get the pool-shaped view of the given attention layer’s combined KVCache — same device allocation as getCombinedKVCache(), reinterpreted as the AttentionPlugin’s paged-pool binding contract [2, numPages(), kTOKENS_PER_PAGE, numKVHeads_i, headDim_i]. This is the tensor that must be bound to the engine’s past/present_key_values_i inputs so dims.d[1] == numPages (see AttentionPlugin::enqueue).

Parameters:

attnLayerIdx – The index of the attention layer.

Returns:

A reference to the pool-view tensor.

std::pair<rt::Tensor, rt::Tensor> getSeparateKVCache(

int32_t attnLayerIdx

) const noexcept

Get the K-half and V-half of the given attention layer’s pool as separate tensor views.

Parameters:

attnLayerIdx – The index of the attention layer.

Returns:

{kView, vView}, each shaped [maxBatchSize, capPadded, numKVHeads_i, headDim_i].

int32_t maxCapPadded() const noexcept

Get the padded per-slot token capacity.

Returns:

capPadded = ceil(maxSequenceLength / kTOKENS_PER_PAGE) * kTOKENS_PER_PAGE.

int32_t numPages() const noexcept

Get the total number of pages spanned by the pool.

Returns:

Config::numPages if non-zero, else the minimum active pages (maxBatchSize * maxCapPadded() / kTOKENS_PER_PAGE).

void *kPoolPtr(int32_t attnLayerIdx) const noexcept

Get the K-half page-pool base pointer for the given attention layer, i.e. the base of the combined tensor reinterpreted as [numPages(), kTOKENS_PER_PAGE, numKVHeads_i, headDim_i].

Parameters:

attnLayerIdx – The index of the attention layer.

Returns:

Device pointer to the K pool.

void *vPoolPtr(int32_t attnLayerIdx) const noexcept

Get the V-half page-pool base pointer for the given attention layer, i.e. kPoolPtr(attnLayerIdx) offset by numPages() * kTOKENS_PER_PAGE * numKVHeads_i * headDim_i elements.

Parameters:

attnLayerIdx – The index of the attention layer.

Returns:

Device pointer to the V pool.

KVLayerConfig const &getLayerConfig(

int32_t attnLayerIdx

) const noexcept

Get the layer configuration for the given attention layer.

Parameters:

attnLayerIdx – The index of the attention layer.

Returns:

The KVLayerConfig for this layer.

int32_t numLayers() const noexcept

Get the number of attention layers.

Returns:

Number of attention layers

bool isUniform() const noexcept

Check if all layers have the same numKVHeads and headDim.

Returns:

True if all layers are uniform

Config const &getConfig() const noexcept

Get cache configuration.

Returns:

Cache configuration

struct KVLayerConfig

Per-layer KV head configuration for heterogeneous models.

Public Members

int32_t numKVHeads = {}

Number of key-value heads for this layer.

int32_t headDim = {}

Head dimension for this layer.