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

QRegisterBufferRDNA

struct QRegisterBufferRDNA[dtype: DType, mma_shape: IndexList[Int(3)], k_group_size: Int, WM: Int, WN: Int, BN: Int, BK: Int, depth: Int]

Q register buffer: loads each warp's (WM, depth) Q sub-tile into BK-strip MMA fragments at construction.

Parameters​

  • ​dtype (DType): Element type of the Q tensor and register fragments.
  • ​mma_shape (IndexList[Int(3)]): (M, N, K) shape of a single MMA instruction as an IndexList[3].
  • ​k_group_size (Int): Grouping factor for the MMA K dimension.
  • ​WM (Int): M (query) tile height owned by each warp, in elements.
  • ​WN (Int): N tile width owned by each warp, in elements.
  • ​BN (Int): Total N (key) tile width across all warps, in elements.
  • ​BK (Int): K strip width loaded per pipeline stage, in elements.
  • ​depth (Int): Total head depth of the Q matrix in DRAM, in elements.

Fields​

  • ​reg_tile (QRegisterBufferRDNA[dtype, mma_shape, k_group_size, WM, WN, BN, BK, depth].RegisterTileType):

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

mma_dtype​

comptime mma_dtype = dtype

MMA_K​

comptime MMA_K = mma_shape[Int(2)]

MMA_M​

comptime MMA_M = mma_shape[Int(0)]

mma_tile_layout​

comptime mma_tile_layout = row_major[ceildiv(WM, mma_shape[Int(0)]), Int(16)]()

MMATileType​

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

num_k_tiles​

comptime num_k_tiles = ceildiv(BK, Int((mul mma_shape[Int(2)], k_group_size)))

num_mmas​

comptime num_mmas = ceildiv(WM, mma_shape[Int(0)])

num_tiles​

comptime num_tiles = (depth // BK)

rdna_frag_size​

comptime rdna_frag_size = RDNA_AB_FRAG_SIZE

reg_dtype​

comptime reg_dtype = dtype

reg_tile_layout​

comptime reg_tile_layout = row_major[(Int((mul ceildiv(BK, Int((mul mma_shape[Int(2)], k_group_size))), ceildiv(WM, mma_shape[Int(0)]))) * QRegisterBufferRDNA[dtype, mma_shape, k_group_size, WM, WN, BN, BK, depth].num_tiles), Int(16)]()

RegisterTileType​

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

simd_width​

comptime simd_width = simd_width_of[dtype]()

Methods​

__init__​

def __init__[q_layout: TensorLayout](out self, tensor: TileTensor[dtype, q_layout, ImmutAnyOrigin], valid_rows: Int)

Load each warp's Q sub-tile from DRAM into register MMA fragments.

Parameters:

  • ​q_layout (TensorLayout): TensorLayout of the Q tensor in DRAM (inferred).

Args:

get_dtype​

static def get_dtype() -> DType

Returns:

DType

get_mma_tile​

def get_mma_tile[tile_idx: Int, k_idx: Int](self) -> Self.MMATileType

MMA fragment for the tile_idx-th depth tile, k_idx-th K strip within it.

Parameters:

  • ​tile_idx (Int): Index of the depth tile along the depth axis.
  • ​k_idx (Int): Index of the K strip within the depth tile.

Returns:

Self.MMATileType

get_reg_tile​

def get_reg_tile[stage: Int = Int(0)](self) -> Self.RegisterTileType

Returns:

Self.RegisterTileType

zero​

def zero(self)

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact