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

QRegisterBuffer

struct QRegisterBuffer[dtype: DType, mma_shape: IndexList[Int(3)], WM: Int, WN: Int, BN: Int, BK: Int, depth: Int, thread_rows: Int, thread_cols: Int]

Holds the Q query tile in register memory for gfx950 attention MMAs.

Loads the per-warp [WM, depth] Q sub-tile from DRAM into a row-major register TileTensor tiled into BK-wide strips, with partial-tile zero padding when depth is not a multiple of BK. Exposes MMA-sized sub-tiles via mma_tile, in-place scaling via scale, and zeroing via zero.

Parameters​

  • ​dtype (DType): Element dtype of the Q tile stored in the register TileTensor.
  • ​mma_shape (IndexList[Int(3)]): MFMA instruction shape [M, N, K] used for the QK MMA; sets the A-operand fragment size and M-direction MMA count.
  • ​WM (Int): Warp tile size along the M (query) dimension; each warp owns a [WM, depth] sub-tile of Q.
  • ​WN (Int): Warp tile size along the N dimension; used with BN to compute warp coordinates via get_warp_coords[BN, WN].
  • ​BN (Int): N block dimension of the attention tile; used with WN to compute the warp's row coordinate via get_warp_coords.
  • ​BK (Int): K block dimension; Q is tiled into BK-wide strips in register memory.
  • ​depth (Int): Head dimension of the Q tile; the per-warp Q sub-tile spans [WM, depth].
  • ​thread_rows (Int): Per-thread fragment rows for the RegTileLoader col-major thread distribution.
  • ​thread_cols (Int): Per-thread fragment columns for the RegTileLoader col-major thread distribution.

Fields​

  • ​reg_tile (QRegisterBuffer[dtype, mma_shape, WM, WN, BN, BK, depth, thread_rows, thread_cols].RegType):

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

input_frag_size​

comptime input_frag_size = num_matrix_reg[mma_shape[Int(0)], mma_shape[Int(2)]]()

mma_dtype​

comptime mma_dtype = dtype

MMA_K​

comptime MMA_K = mma_shape[Int(2)]

MMA_M​

comptime MMA_M = mma_shape[Int(0)]

num_k_tiles​

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

num_mmas​

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

num_tiles​

comptime num_tiles = ceildiv(depth, BK)

reg_dtype​

comptime reg_dtype = dtype

reg_layout​

comptime reg_layout = row_major[(Int((mul ceildiv(BK, mma_shape[Int(2)]), ceildiv(WM, mma_shape[Int(0)]))) * QRegisterBuffer[dtype, mma_shape, WM, WN, BN, BK, depth, thread_rows, thread_cols].num_tiles), num_matrix_reg[mma_shape[Int(0)], mma_shape[Int(2)]]()]()

RegType​

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

Methods​

__init__​

def __init__[q_tile_layout: TensorLayout](out self, q_tile: TileTensor[dtype, q_tile_layout, ImmutAnyOrigin])

Load Q tile from DRAM into registers via buffer_load intrinsics.

Each warp loads its [WM, depth] sub-tile using col-major thread distribution (matching get_warp_layout[mma_shape]), then tiles it into BK-wide strips stored in register memory.

Parameters:

  • ​q_tile_layout (TensorLayout): Compile-time TensorLayout of the input q_tile DRAM tile.

Args:

mma_tile​

def mma_tile[tile_idx: Int, k_idx: Int](self) -> TileTensor[dtype, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

Return MMA-sized sub-tile for the given tile and k indices.

Parameters:

  • ​tile_idx (Int): Index of the BK-wide Q strip within the depth dimension.
  • ​k_idx (Int): Index of the MMA-K strip within the BK-wide strip.

Returns:

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

scale​

def scale[accum_type: DType](self, scale_factor: Scalar[accum_type])

Scale all Q register elements in-place.

Casts bf16 -> f32, multiplies by scale_factor, casts back to bf16. Used for pre-scaling Q by (1/sqrt(d) * log2e) so that QK matmul produces already-scaled scores, eliminating scale from the hot loop.

Parameters:

  • ​accum_type (DType): Accumulator dtype used for the intermediate multiply before casting back to dtype.

Args:

  • ​scale_factor (Scalar[accum_type]): Multiplicative factor applied to every Q element in-place.

zero​

def zero(self)

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact