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 struct

KVCacheIterator

struct KVCacheIterator[cache_t: MHAOperand, tile_size: Int, kv_num_heads: Int, depth: Int, cache_depth: Int = depth, head_dim_offset: Int = Int(0)]

TileTensor-based DRAM tile iterator.

Returns a TileTensor with Scalar for the row dimension (valid-row count) and ComptimeInt for depth and strides. No RuntimeLayout storage.

When cache_depth != depth, the DRAM stride uses cache_depth (e.g., MLA K_rope reads 64 columns from a 576-wide cache row). head_dim_offset shifts the column start (e.g., skip to rope portion at column 512).

Parameters

• ​cache_t (MHAOperand): MHA operand wrapping the KV cache pointer and dtype.
• ​tile_size (Int): Number of rows loaded per DRAM tile.
• ​kv_num_heads (Int): Number of KV heads in the cache.
• ​depth (Int): Per-head depth read from each cache row (head dimension).
• ​cache_depth (Int): Row stride of the paged KV cache in elements (defaults to depth).
• ​head_dim_offset (Int): Column offset within each cache row, used to skip to a sub-region such as the rope portion (defaults to 0).

Fields

• ​cache (cache_t):
• ​end (Int):
• ​tile_start_row (Int):
• ​batch_idx (Int):
• ​kv_head_idx (Int):

Implemented traits

AnyType, Deinitable, Movable

comptime members

GmemTileLayout

comptime GmemTileLayout = Layout[*?, *?]

GmemTileType

comptime GmemTileType = TileTensor[cache_t.dtype, Layout[*?, *?], ImmutAnyOrigin]

Methods

__init__

def __init__(out self, cache: cache_t, batch_idx: Int, kv_head_idx: Int, end: Int)

next_tile

def next_tile(mut self) -> Self.GmemTileType

Returns a TileTensor for the next DRAM tile.

Returns:

Self.GmemTileType

increment

def increment(mut self)