IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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, ImplicitlyDeletable

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)

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact