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

Fp4WeightLoader

struct Fp4WeightLoader[in_type: DType, a_layout: TensorLayout, packed_layout: TensorLayout, scale_layout: TensorLayout]

Owner of the packed-FP4 weight -> dequant -> {register B | SMEM} transition.

Holds the activation a, packed weight, and block-scale TileTensor views plus the (M, N, K) geometry, and exposes bounds-aware loads that do all addressing via TileTensor indexing (t[i, j][0]) and width-loads (t.load[width=W](Coord(...))) -- no raw pointer arithmetic. in_type is the dequant target (bf16). Views are held with AnyOrigin (the kernel args outlive the K-loop; this loader is a local, not a struct field, so the field-origin restriction does not apply).

Parameters

• ​in_type (DType): Dequant + activation dtype (bf16).
• ​a_layout (TensorLayout): Layout of the activation (M, K) view.
• ​packed_layout (TensorLayout): Layout of the packed weight (N, K//2) view.
• ​scale_layout (TensorLayout): Layout of the block scales (N, ceil(K/16)) view.

Fields

• ​a (TileTensor[in_type, a_layout, ImmUntrackedOrigin]):
• ​packed (TileTensor[DType.uint8, packed_layout, ImmUntrackedOrigin]):
• ​scales (TileTensor[DType.float8_e4m3fn, scale_layout, ImmUntrackedOrigin]):
• ​M (Int):
• ​N (Int):
• ​K (Int):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable

comptime members

SF

comptime SF = NVFP4_SF_VECTOR_SIZE

Methods

from_kernel_args

static def from_kernel_args(a: TileTensor[in_type, a_layout, ImmutAnyOrigin], packed: TileTensor[DType.uint8, packed_layout, ImmutAnyOrigin], scales: TileTensor[DType.float8_e4m3fn, scale_layout, ImmutAnyOrigin], M: Int, N: Int, K: Int) -> Self

Build the loader from the kernel's AnyOrigin tensor args.

Rebases each view onto ImmUntrackedOrigin (the field-origin rule; the args outlive the K-loop), preserving layout/shape/stride.

Args:

• ​a (TileTensor[in_type, a_layout, ImmutAnyOrigin]): Bf16 activation TileTensor view with shape (M, K).
• ​packed (TileTensor[DType.uint8, packed_layout, ImmutAnyOrigin]): Packed FP4 weight TileTensor view with shape (N, K//2) (uint8, lo-nibble first).
• ​scales (TileTensor[DType.float8_e4m3fn, scale_layout, ImmutAnyOrigin]): FP8-E4M3 block-scale TileTensor view with shape (N, ceil(K/16)).
• ​M (Int): Number of rows in the activation and the output.
• ​N (Int): Number of columns in the output and rows of the weight.
• ​K (Int): Reduction dimension; inner size of a and the weight.

load_a_frag

def load_a_frag[bounded: Bool](self, arow0: Int, k0: Int, a_rc: Array[IndexList[Int(2)], Int(8)]) -> SIMD[in_type, SIMDLength(8)]

This lane's 8-element A fragment for the 16x16 A tile at (arow0, k0).

a_rc[i] is the lane's _apple_frag_layout (row, col) for element i. On the interior (bounded=False) every row is in-bounds and the load is unconditional (branch-free). On an edge tile (bounded=True) rows past M zero-fill -- the ragged-M A over-read fix (a zero A element multiplies into nothing, matching the dense reference). K is always in-bounds (callers pass K % 16 == 0 tile-aligned strips).

Parameters:

• ​bounded (Bool): Whether to bounds-check rows against M (edge tile).

Args:

• ​arow0 (Int): Absolute M-origin (row) of the 16x16 A tile.
• ​k0 (Int): Absolute K-origin of the current K strip (a multiple of 16).
• ​a_rc (Array[IndexList[Int(2)], Int(8)]): This lane's per-element (row, col) coords in the A tile (the _apple_frag_layout map, 8 entries).

Returns:

SIMD[in_type, SIMDLength(8)]

decode_b_frag_regc

def decode_b_frag_regc(self, bcol0: Int, k0: Int, rb: Int, cb: Int) -> SIMD[in_type, SIMDLength(16)]

Decode this lane's 16-element transpose_right register B fragment.

The fragment reads 4 N-rows (rb, rb+8, rb+16, rb+24), each contributing 4 CONTIGUOUS K values (cb .. cb+3, cb a multiple of 4). Those 4 K nibbles are one 2-byte packed load, and (because cb+3 < 16) share one FP8 block scale. All addressing is TileTensor width-loads / indexing. Interior fast path: callers pass tile-aligned N and K % 16 == 0, so the 4 N-rows and the K run are always in-bounds.

Args:

• ​bcol0 (Int): Absolute N-origin of the simdgroup's 16x32 B tile in the packed weight.
• ​k0 (Int): Absolute K-origin of the current K strip (a multiple of 16).
• ​rb (Int): This lane's base N-row offset within the B tile (frag_row_base(lane), range 0..7).
• ​cb (Int): This lane's base K-column offset within the 16x16 sub-fragment (frag_col_base(lane), a multiple of 4).

Returns:

SIMD[in_type, SIMDLength(16)]

decode_strip_to_smem

def decode_strip_to_smem[b_view_origin: MutOrigin, b_view_layout: TensorLayout, b_view_addr: AddressSpace, //, bytes_per_thread: Int, cols_per_thread: Int](self, b_view: TileTensor[in_type, b_view_layout, b_view_origin, address_space=b_view_addr], n_abs: Int, n_local: Int, k0: Int, col0: Int)

Cooperatively decode this thread's cols_per_thread bf16 cols of the (BN, BK) weight strip into the SMEM view b_view.

This thread owns N-row n_local (absolute n_abs), a CONTIGUOUS bytes_per_thread-byte packed run at bf16 col col0 (so adjacent threads read adjacent packed bytes -- the coalesced-K load), decoded with ONE FP8 block scale (cols_per_thread <= 16 = one 16-block). The packed run is a single TileTensor width-load; the SMEM write is a TileTensor width-store into b_view (no raw pointer arithmetic). All-in-bounds interior strip (callers pass tile-aligned N, K % BK == 0).

Parameters:

• ​b_view_origin (MutOrigin): Origin of the b_view SMEM store target.
• ​b_view_layout (TensorLayout): Layout of the b_view SMEM store target.
• ​b_view_addr (AddressSpace): Address space of the b_view SMEM store target.
• ​bytes_per_thread (Int): Packed bytes this thread loads in one width-load (each byte yields two bf16 columns).
• ​cols_per_thread (Int): Number of bf16 columns this thread decodes and stores (2 * bytes_per_thread; must be <= 16 for one scale block).

Args:

• ​b_view (TileTensor[in_type, b_view_layout, b_view_origin, address_space=b_view_addr]): SMEM TileTensor view of the (BN, BK) decoded weight strip to store into.
• ​n_abs (Int): Absolute N-row index in the packed weight for this thread's strip (threadgroup N-origin + local row).
• ​n_local (Int): Local N-row index within the b_view strip (the SMEM store coordinate).
• ​k0 (Int): Absolute K-origin of the current strip (a multiple of BK).
• ​col0 (Int): First bf16 column this thread decodes in its row's strip (a multiple of COLS_PER_THREAD).

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit