For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo trait

KVCacheT

Trait for different KVCache types and implementations.

Represents a single (key or value) cache.

Implemented traits

AnyType, Copyable, Deinitable, DevicePassable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

cache_lengths_engine

comptime cache_lengths_engine

dtype

comptime dtype

Engine

comptime Engine

kv_params

comptime kv_params

page_size_

comptime page_size_

quantization_enabled

comptime quantization_enabled = False

quantization_granularity

comptime quantization_granularity = Int(1)

scale_dtype

comptime scale_dtype

Required methods

block_paged_storage

def block_paged_storage[tile_size: Int](self, batch_idx: Int, start_tok_idx: Int, head_idx: Int, head_dim_idx: Int = Int(0)) -> Self.Engine.StorageType[True, MutUnsafeAnyOrigin, Self.dtype, MutAnyOrigin, AddressSpace.GENERIC]

Returns:

_Self.Engine.StorageType[True, MutUnsafeAnyOrigin, _Self.dtype, MutAnyOrigin, AddressSpace.GENERIC]

cache_lengths_nd

def cache_lengths_nd(self) -> TileTensor[.uint32, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=Self.cache_lengths_engine]

Returns the cache lengths as a TileTensor.

Returns:

TileTensor[.uint32, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=_Self.cache_lengths_engine]

cache_length

def cache_length(self, batch_idx: Int) -> Int

Returns the length of the cache for a given batch index.

Returns:

Int

load

def load[width: Int, output_dtype: DType = _Self.dtype](self, bs: Int, head_idx: Int, tok_idx: Int, head_dim_idx: Int) -> SIMD[output_dtype, width]

Loads an element from the given index.

Returns:

SIMD[output_dtype, width]

store

def store(self, bs: Int, head_idx: Int, tok_idx: Int, head_dim_idx: Int, val: SIMD[Self.dtype])

Stores an element at the given index.

store_scale

def store_scale[scales_dtype: DType = _Self.scale_dtype, width: Int = Int(1)](self, bs: Int, head_idx: Int, tok_idx: Int, head_dim_idx: Int, scales: SIMD[scales_dtype, width])

Stores the quantization scales at the given index.

load_scale

def load_scale[width: Int](self, bs: Int, head_idx: Int, tok_idx: Int, head_dim_idx: Int) -> SIMD[Self.scale_dtype, width]

Loads the quantization scales from the given index.

Returns:

SIMD[_Self.scale_dtype, width]

load_quantized

def load_quantized[width: Int](self, bs: Int, head_idx: Int, tok_idx: Int, head_dim_idx: Int) -> SIMD[Self.dtype, width]

Loads a quantized element from the given index.

Returns:

SIMD[_Self.dtype, width]

empty_cache

def empty_cache(self) -> Bool

Returns true if the cache_lengths for all requests is 0, false otherwise.

Returns:

Bool

max_prompt_length

def max_prompt_length(self) -> UInt32

Returns the maximum sequence length across all batches of the current request.

Returns:

UInt32

max_context_length

def max_context_length(self) -> UInt32

Returns the maximum cache length used across all batches of the current request.

Returns:

UInt32

block_paged_ptr

def block_paged_ptr[tile_size: Int](self, batch_idx: Int, start_tok_idx: Int, head_idx: Int, head_dim_idx: Int = Int(0)) -> Pointer[Scalar[Self.dtype], MutAnyOrigin]

Returns a pointer to the KVCache block at the given index.

Paged KVCache implementations must have a block_size which is a multiple of the and greater than the layout's first dimension.

Returns:

Pointer[Scalar[_Self.dtype], MutAnyOrigin]

scales_block_paged_ptr

def scales_block_paged_ptr(self, batch_idx: Int, start_tok_idx: Int, head_idx: Int, head_dim_idx: Int = Int(0)) -> Pointer[Scalar[Self.scale_dtype], MutAnyOrigin]

Returns a pointer to the scales block at the requested indices.

Returns:

Pointer[Scalar[_Self.scale_dtype], MutAnyOrigin]

scales_raw_ptr

def scales_raw_ptr(self) -> Pointer[Scalar[Self.scale_dtype], MutAnyOrigin]

Returns the base pointer to the scales tensor.

For PagedKVCache with quantization enabled, this returns the raw base pointer of the scales TileTensor. For caches without quantization, returns a null pointer.

Returns:

Pointer[Scalar[_Self.scale_dtype], MutAnyOrigin]

max_tile_size

static def max_tile_size() -> Int

Returns the maximum tile size for the KVCache.

Returns:

Int

num_kv_rows

def num_kv_rows(self) -> Int

Returns the total number of virtual rows in this KV cache view.

For paged caches this accounts for the paging stride: (total_blocks - 1) * stride + page_size.

Returns:

Int

row_idx

def row_idx(self, batch_idx: UInt32, start_tok_idx: UInt32) -> UInt32

Returns the row idx when viewing the memory as a matrix.

Returns:

UInt32

scale_row_idx

def scale_row_idx(self, batch_idx: UInt32, start_tok_idx: UInt32) -> UInt32

Returns the row idx of a token's scale in the SCALE pool.

Deliberately not derivable from row_idx: a paged cache indexes its scales through scales_lookup_table, a separate member that only DEFAULTS to lookup_table, and the scale pool carries its own block stride. Reusing row_idx reads the wrong block whenever a caller supplies a distinct scales LUT.

Returns:

UInt32

get_tma_row

def get_tma_row(self, encoded_index: Int32) -> Int32

Convert an encoded sparse index to a physical TMA row.

For paged caches the encoded index is physical_block * page_size + offset and this method returns physical_block * stride + offset. Non-paged caches return the encoded index unchanged.

Returns:

Int32

create_paged_tma_tile

def create_paged_tma_tile[swizzle_mode: TensorMapSwizzle, *, BN: Int, BK: Int = padded_depth[_Self.dtype, swizzle_mode, _Self.kv_params.head_size]()](self, ctx: DeviceContext) -> TMATensorTile[Self.dtype, Int(4), Index[Int, Int, Int, Int](Int(1), BN, Int(1), BK)]

Creates the split (block, row_in_block) TMA tile for this cache.

The counterpart of :meth:create_tma_tile, addressed by :meth:kv_tma_coords. The flat descriptor folds the block into the row, and that product spans the whole allocation rather than this cache's share of it, so it leaves the signed 32-bit TMA coordinate on a large pool.

Returns:

TMATensorTile[_Self.dtype, Int(4), Index[Int, Int, Int, Int](Int(1), BN, Int(1), BK)]

create_tma_tile

def create_tma_tile[swizzle_mode: TensorMapSwizzle, *, BN: Int, BK: Int = padded_depth[_Self.dtype, swizzle_mode, _Self.kv_params.head_size](), fold_chunks: Int = Int(1), row_major: Bool = False](self, ctx: DeviceContext) -> TMATensorTile[Self.dtype, Int(3), _padded_shape[Int(3), Self.dtype, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode](), _ragged_shape[Int(3), Self.dtype, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode]()]

Creates a TMA tile for this KV cache. This is useful for k-major MMA operations where we don't need to mask any extra rows.

fold_chunks >= 2 builds a depth-chunk-folded descriptor (SM100); 1 (default) keeps the original 3D descriptor. row_major=True (with fold_chunks >= 2) builds the rank-5 chunk-inner box (one TMA per multi-atom-row page); False builds the rank-4 chunk-outer box.

Returns:

TMATensorTile[_Self.dtype, Int(3), _padded_shape[Int(3), _Self.dtype, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode](), _ragged_shape[Int(3), _Self.dtype, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode]()]

create_index_scale_tma_tile

def create_index_scale_tma_tile[TILE: Int](self, ctx: DeviceContext) -> TMATensorTile[Self.scale_dtype, Int(2), Index[Int, Int](Int(1), flat_scale_window[_Self.scale_dtype, TILE]())]

Creates a flat TMA descriptor over the SCALE pool.

The box is TILE scales plus one 16-byte alignment unit of slack, and is only well defined when a token owns exactly one scale -- implementations assert that. The caller's obligation is that TILE divides page_size, the same one the K descriptor already carries.

Returns:

TMATensorTile[_Self.scale_dtype, Int(2), Index[Int, Int](Int(1), flat_scale_window[_Self.scale_dtype, TILE]())]

create_rope_tma_tile

def create_rope_tma_tile[swizzle_mode: TensorMapSwizzle, *, BN: Int, BK: Int, padded_depth: Int](self, ctx: DeviceContext) -> TMATensorTile[.bfloat16, Int(3), _padded_shape[Int(3), DType.bfloat16, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode](), _ragged_shape[Int(3), DType.bfloat16, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode]()]

Creates a BF16 TMA tile for the rope portion of the KV cache.

For the per-tensor rope-aware layout, each token row in the KV cache is stored as padded_depth FP8 bytes (content) followed by BK BF16 elements (rope). This method returns a TMA descriptor that points at the rope data starting at byte offset padded_depth within each row, reinterpreted as BF16.

Returns:

TMATensorTile[.bfloat16, Int(3), _padded_shape[Int(3), DType.bfloat16, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode](), _ragged_shape[Int(3), DType.bfloat16, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode]()]

create_gather4_tma_tile

def create_gather4_tma_tile[*, tile_height: Int = Int(4), tile_width: Int, tile_stride: Int = tile_width, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, tma_dtype: DType = _Self.dtype, l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE](self, ctx: DeviceContext) -> TMATensorTile[tma_dtype, Int(2), IndexList(tile_height, _gather4_box_width[tma_dtype, tile_width, swizzle_mode](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[tma_dtype, tile_width, swizzle_mode](), __list_literal__=NoneType(None))]

Creates a 2D TMA gather4 descriptor for this KV cache.

The descriptor views the KV cache as a flat 2D matrix of [num_kv_rows, tile_width] and is configured for gather4 operations that load 4 non-contiguous rows per TMA instruction. The box width is derived from the swizzle mode; for SWIZZLE_NONE it equals tile_width.

The tile_height parameter records the full tile height (e.g. 64 rows) in the returned TMATensorTile.tile_shape. The hardware descriptor shape stays (1, box_width) as required by TMA gather4.

When tma_dtype differs from Self.dtype, the underlying data pointer is bitcast to tma_dtype at descriptor creation time. This allows, for example, creating an INT64/SWIZZLE_NONE descriptor over FP8 data for linear SMEM layout.

Parameters:

• ​tile_height (Int): Number of rows in the tile. Must be a multiple of 4. Defaults to 4 for backward compatibility.
• ​tile_width (Int): Number of elements per row to load (box width) in tma_dtype elements.
• ​tile_stride (Int): Row stride in elements in global memory. Defaults to tile_width. Use a larger value when the global row is wider than the portion to load.
• ​swizzle_mode (TensorMapSwizzle): TMA swizzle mode for shared memory access pattern. Defaults to SWIZZLE_NONE.
• ​tma_dtype (DType): The data type used for the TMA descriptor. Defaults to Self.dtype. When different, the pointer is bitcast.
• ​l2_promotion (TensorMapL2Promotion): L2 cache promotion hint for TMA loads. Defaults to NONE.

Args:

• ​ctx (DeviceContext): The CUDA device context used to create the TMA descriptor.

Returns:

TMATensorTile[tma_dtype, Int(2), IndexList(tile_height, _gather4_box_width[tma_dtype, tile_width, swizzle_mode](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[tma_dtype, tile_width, swizzle_mode](), __list_literal__=NoneType(None))]: A TMATensorTile with box width derived from the swizzle mode.

create_rope_gather4_tma_tile

def create_rope_gather4_tma_tile[*, tile_height: Int = Int(4), tile_width: Int, padded_depth: Int, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE](self, ctx: DeviceContext) -> TMATensorTile[.bfloat16, Int(2), IndexList(tile_height, _gather4_box_width[DType.bfloat16, tile_width, swizzle_mode](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.bfloat16, tile_width, swizzle_mode](), __list_literal__=NoneType(None))]

Creates a BF16 gather4 TMA descriptor for the rope portion of the KV cache.

For the per-tensor rope-aware layout each token row is stored as padded_depth FP8 bytes (content) followed by BF16 rope elements. This method offsets the base pointer by padded_depth bytes, reinterprets as BF16, and creates a gather4 TMA descriptor with tile_width BF16 elements per row.

Parameters:

• ​tile_height (Int): Number of rows in the tile. Must be a multiple of 4.
• ​tile_width (Int): Number of BF16 elements per row in global memory.
• ​padded_depth (Int): Byte offset from row start to the rope data.
• ​swizzle_mode (TensorMapSwizzle): TMA swizzle mode for shared memory access pattern.
• ​l2_promotion (TensorMapL2Promotion): L2 cache promotion hint for TMA loads. Defaults to NONE.

Args:

• ​ctx (DeviceContext): The CUDA device context used to create the TMA descriptor.

Returns:

TMATensorTile[.bfloat16, Int(2), IndexList(tile_height, _gather4_box_width[DType.bfloat16, tile_width, swizzle_mode](), __list_literal__=NoneType(None)), IndexList(Int(1), _gather4_box_width[DType.bfloat16, tile_width, swizzle_mode](), __list_literal__=NoneType(None))]: A BF16 TMATensorTile configured for gather4.

Provided methods

kv_tma_coords

def kv_tma_coords(self, batch_idx: UInt32, tok_idx: UInt32) -> Tuple[Int32, Int32]

The (row_in_block, block) coordinate of a token's KV tile.

Defaults to the flat form -- the whole pool as one block -- which is what a contiguous cache wants. A paged cache overrides it: its pool spans the entire KV allocation, so folding the block into the row gives a coordinate that grows with total cache memory and leaves the signed 32-bit TMA coordinate. Mirrors :meth:scale_tma_coords.

Args:

• ​batch_idx (UInt32): Batch entry to address.
• ​tok_idx (UInt32): Token within the entry.

Returns:

Tuple[Int32, Int32]: The row within the block, and the block.

scale_tma_coords

def scale_tma_coords(self, batch_idx: UInt32, start_tok_idx: UInt32) -> Tuple[Int32, Int32]

The (row, block) coordinate of a token's scale tile.

Defaults to the flat form -- the whole pool as one block -- which is what a contiguous scale buffer wants. A paged cache overrides it: its pool spans the entire KV allocation, so folding the block into the row gives a coordinate that grows with total cache memory and leaves the signed 32-bit TMA coordinate.

Args:

• ​batch_idx (UInt32): Batch entry to address.
• ​start_tok_idx (UInt32): First token of the tile, within the entry.

Returns:

Tuple[Int32, Int32]: The row, and the block.

populate

def populate[BN: Int, base_alignment: Int, pair_cta: Bool = False, is_leader: Bool = True](self, batch_idx: UInt32, base_kv_row: UInt32) -> PagedRowIndices[BN, _Self.page_size_, pair_cta, is_leader]

Populate a full PagedRowIndices[BN, ...] for a BN-row tile.

base_alignment is a comptime promise that base_kv_row % base_alignment == 0 at runtime, typically mask.start_column_alignment[...](). The PagedKVCache override uses it to pick the largest legal SIMD chunk for its LUT vector load and to skip the intra-page divmod when base_alignment % page_size == 0.

Default: scalar loop over num_pages calls to row_idx. The PagedKVCache override replaces this with a single aligned SIMD load against the lookup table.

Returns:

PagedRowIndices[BN, _Self.page_size_, pair_cta, is_leader]