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

DecodeKVBuffer

struct DecodeKVBuffer[dtype: DType, kv_tile_layout: TensorLayout, //, config: KVBufferConfig, tensor_core_mma: TiledTensorCore[tensor_core_mma.out_type, tensor_core_mma.in_type, tensor_core_mma.shape, tensor_core_mma.group_size, tensor_core_mma.transpose_b], swizzle: Optional[Swizzle], BN: Int, WN: Int, BK: Int, depth: Int, num_threads: Int, num_stages: Int = Int(1), token_gen: Bool = False]

Double-buffered register-staged KV buffer for the AMD decode mirror path.

Stages DRAM tiles through a register load_tile (optionally multi-stage) before copying to SMEM and loading into the MFMA mma_tile. The KVBufferConfig parameter selects the K (transpose) or V (non-transpose) tiling, iterator axis, and warp-tile coordinate so a single implementation serves both operands.

Parameters

• ​dtype (DType): Element type of the KV cache tiles (inferred).
• ​kv_tile_layout (TensorLayout): Runtime layout of the DRAM KV tile used to build gmem_tile (inferred).
• ​config (KVBufferConfig): K or V role config selecting warp tile, block tile, iterator axis, and warp-tile coordinate.
• ​tensor_core_mma (TiledTensorCore[tensor_core_mma.out_type, tensor_core_mma.in_type, tensor_core_mma.shape, tensor_core_mma.group_size, tensor_core_mma.transpose_b]): MFMA descriptor (shape, group size, transpose_b, out type) driving SMEM-to-register loads.
• ​swizzle (Optional[Swizzle]): LDS bank-conflict swizzle for SMEM loads and stores, or None.
• ​BN (Int): Number of keys per DRAM tile (block tile extent along the key axis).
• ​WN (Int): Number of keys per warp tile (warp tile extent along the key axis).
• ​BK (Int): Strip width along the MMA K dimension of one DRAM-to-SMEM load.
• ​depth (Int): Per-head depth of the KV cache (head dimension).
• ​num_threads (Int): Number of threads in the block driving DMA and register distribution.
• ​num_stages (Int): Number of register buffer stages for double buffering (defaults to 1).
• ​token_gen (Bool): Selects the token-generation thread layout for DMA distribution (defaults to False).

Fields

• ​load_tile (DecodeKVBuffer[config, tensor_core_mma, swizzle, BN, WN, BK, depth, num_threads, num_stages, token_gen].LoadTile):
• ​mma_tile (DecodeKVBuffer[config, tensor_core_mma, swizzle, BN, WN, BK, depth, num_threads, num_stages, token_gen].MmaTile):
• ​smem_tile (DecodeKVBuffer[config, tensor_core_mma, swizzle, BN, WN, BK, depth, num_threads, num_stages, token_gen].SmemTile):
• ​gmem_tile (DecodeKVBuffer[config, tensor_core_mma, swizzle, BN, WN, BK, depth, num_threads, num_stages, token_gen].GmemTileType):
• ​reg_loader (DecodeKVBuffer[config, tensor_core_mma, swizzle, BN, WN, BK, depth, num_threads, num_stages, token_gen].RegLoaderType):
• ​tile_idx (Int):
• ​load_tile_id (Int):

Implemented traits

AnyType, Deinitable, Movable

comptime members

GmemTileType

comptime GmemTileType = TileTensor[dtype, kv_tile_layout, ImmutAnyOrigin]

LoadTile

comptime LoadTile = TileTensor[dtype, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

MMA_K

comptime MMA_K = shape[Int(2)]

MMA_N

comptime MMA_N = shape[Int(1)]

MmaTile

comptime MmaTile = TileTensor[dtype, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

num_k_tiles

comptime num_k_tiles = ceildiv(BK, Int((mul shape[Int(2)], group_size)))

num_mmas

comptime num_mmas = ceildiv(config.wsize, shape[Int(1)])

num_warps_n

comptime num_warps_n = (BN // WN)

RegLoaderType

comptime RegLoaderType = RegTileLoader[dtype, row_major[(Int((mul min(num_threads, (Int((mul config.btile_dim0, config.btile_dim1)) // simd_width_of[dtype]())), simd_width_of[dtype]())) // config.btile_dim1) if token_gen else (num_threads // Int(4)), (config.btile_dim1 // simd_width_of[dtype]()) if token_gen else Int(4)](), num_threads]

simd_width

comptime simd_width = simd_width_of[dtype]()

SmemTile

comptime SmemTile = TileTensor[dtype, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]

wtile_dim0

comptime wtile_dim0 = config.wtile_dim0

wtile_dim1

comptime wtile_dim1 = config.wtile_dim1

Methods

__init__

def __init__(out self, gmem_tile: TileTensor[dtype, kv_tile_layout, ImmutAnyOrigin], smem_tile: TileTensor[dtype, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED])

load_from_dram

def load_from_dram(mut self)

get_mma_tile

def get_mma_tile(self) -> TileTensor[dtype, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.LOCAL]

Returns:

TileTensor[dtype, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.LOCAL]

copy_to_shared

def copy_to_shared[tile_id: Int = Int(0)](self)

load_from_shared

def load_from_shared[k_mma: Int](self)