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

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, zero_partial_tile_pad: Bool = True, clamp_to_parent: Bool = False]

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.
• ​zero_partial_tile_pad (Bool): Whether to zero the invalid tail of a partial BK strip after loading it.
• ​clamp_to_parent (Bool): Whether to bound the Q load's buffer descriptor by the parent tile rather than this warp's own slab. .tile[] yields comptime shapes, so a slab-built descriptor clamps nothing and a warp entirely past the live rows reads past the tensor. The decode fold turns it on because a padded width reaches that state on every launch; prefill can too but pays for the bound in address math, so it keeps the default.

Fields

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

Implemented traits

AnyType, Deinitable, Movable

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, zero_partial_tile_pad, clamp_to_parent].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:

• ​q_tile (TileTensor[dtype, q_tile_layout, ImmutAnyOrigin]): The full Q tile as a DRAM TileTensor.

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)