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

PRegisterBuffer

struct PRegisterBuffer[accum_type_: DType, dtype: DType, BM: Int, BN: Int, BK: Int, WM: Int, WN: Int, num_m_mmas: Int, num_n_mmas: Int, output_frag_size: Int, shared_memory_backed: Bool, mma_shape: IndexList[Int(3)], tr_load_enabled: Bool = False, num_stages: Int = Int(1), p_swizzle: Optional[Swizzle] = None, raw_fp8_cast: Bool = False]

Holds the P attention-score tile across register and shared memory.

Stores the softmax P scores as a multi-stage register TileTensor (accumulator dtype) and, when shared_memory_backed, mirrors them into a BM x BN shared-memory region carved into BM x BK blocks. Provides mma_tile to cast and interleave P fragments into the MMA-operand dtype and layout, copy_to_shared to spill register tiles to SMEM, and per-stage zero and stage_tile accessors.

Parameters​

  • ​accum_type_ (DType): Accumulator dtype used to store P scores in the register TileTensor, float32.
  • ​dtype (DType): MMA-operand dtype that P is cast to before the PV MMA; the SMEM tile stores elements in this dtype.
  • ​BM (Int): M block dimension of the P tile; each SMEM block is a [BM, BK] row-major region.
  • ​BN (Int): N block dimension of the P tile; the SMEM region is BM x BN carved into BN // BK blocks.
  • ​BK (Int): K block dimension of the P tile; each SMEM block is [BM, BK].
  • ​WM (Int): Warp tile size along the M dimension; each warp owns a [WM, BK] slice of a SMEM block.
  • ​WN (Int): Warp tile size along the N dimension; used to compute warp coordinates via get_warp_coords[BN, WN].
  • ​num_m_mmas (Int): Number of MMA tiles along the M dimension of the P tile.
  • ​num_n_mmas (Int): Number of MMA tiles along the N dimension of the P tile.
  • ​output_frag_size (Int): Per-lane element width of one MMA output fragment in the register tile.
  • ​shared_memory_backed (Bool): When True, mirror P scores into a BM x BN SMEM region and read PV operands from SMEM; when False, keep P in registers only.
  • ​mma_shape (IndexList[Int(3)]): MFMA instruction shape [M, N, K] used for the PV MMA.
  • ​tr_load_enabled (Bool): When True, use the tiled-register-load path in mma_tile (cast and slice or join from the stage tile) instead of the interleave path (defaults to False).
  • ​num_stages (Int): Number of pipeline stages in the register tile; the staging dimension is num_stages * num_n_mmas * num_m_mmas (defaults to 1).
  • ​p_swizzle (Optional[Swizzle]): Optional SMEM swizzle pattern used to spread P rows across LDS banks in copy_to_shared and get_mma_tile_shared (defaults to None).
  • ​raw_fp8_cast (Bool): When True, use raw v_cvt_pk_fp8_f32 without the compiler's clamp + NaN-scrub wrapper for f32 to fp8 casts; only safe when inputs are bounded in (0, 1] (defaults to False).

Fields​

  • ​reg_tile (PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, tr_load_enabled, num_stages, p_swizzle, raw_fp8_cast].RegType):
  • ​smem_tile (PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, tr_load_enabled, num_stages, p_swizzle, raw_fp8_cast].SmemTileType):

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

BlockSmemType​

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

input_frag_size​

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

mma_dtype​

comptime mma_dtype = dtype

MmaTileType​

comptime MmaTileType = TileTensor[PRegisterBuffer[accum_type_, dtype, BM, BN, BK, WM, WN, num_m_mmas, num_n_mmas, output_frag_size, shared_memory_backed, mma_shape, tr_load_enabled, num_stages, p_swizzle, raw_fp8_cast].mma_dtype, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

reg_dtype​

comptime reg_dtype = accum_type_

reg_layout​

comptime reg_layout = row_major[((num_stages * num_n_mmas) * num_m_mmas), output_frag_size]()

RegType​

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

SmemTileType​

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

stage_layout​

comptime stage_layout = row_major[(num_n_mmas * num_m_mmas), output_frag_size]()

StageTileType​

comptime StageTileType = TileTensor[accum_type_, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

Methods​

__init__​

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

get_mma_tile_shared​

def get_mma_tile_shared[tile_idx: Int, k_idx: Int](self) -> Self.MmaTileType

Returns:

Self.MmaTileType

stage_tile​

def stage_tile[stage: Int = Int(0)](self) -> Self.StageTileType

Return the TileTensor sub-tile for the given pipeline stage.

Parameters:

  • ​stage (Int): Pipeline stage index into the register TileTensor (defaults to 0).

Returns:

Self.StageTileType

mma_tile​

def mma_tile[tile_idx: Int, k_idx: Int, stage: Int = Int(0)](self) -> Self.MmaTileType

TileTensor MMA operand with cast+interleave via SIMD whole-vector ops.

Converts f32 accumulator rows to bf16 MMA fragments using SIMD cast, interleave, and slice; no per-element [j] indexing needed.

Parameters:

  • ​tile_idx (Int): Index of the P tile along the N dimension; selects the SMEM block when shared_memory_backed.
  • ​k_idx (Int): Index of the MMA-K strip within the tile.
  • ​stage (Int): Pipeline stage index of the source register tile (defaults to 0).

Returns:

Self.MmaTileType

zero​

def zero[stage: Int](self)

copy_to_shared​

def copy_to_shared(self)

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact