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

MXTokenFormat

struct MXTokenFormat[quant_dtype: DType, scales_dtype: DType, output_layout: TensorLayout, scales_layout: TensorLayout, //, _hid_dim: Int, _top_k: Int, _alignment: Int = Int(0), *, fuse_a_scale_preshuffle: Bool = False, mx_format: MXFormat = MXFormat.from_dtype[quant_dtype]()]

Token format for MX (microscaling) FP4 quantization.

Packs each token as FP4 values with block-wise MX scale factors (one float8_e8m0fnu scale per MXFP4_SF_VECTOR_SIZE elements). Achieves the highest compression ratio of all built-in token formats. Setting fuse_a_scale_preshuffle=True folds the grouped-matmul scale preshuffle into the dispatch-wait copy path (KS224 up-projection fusion).

Parameters

• ​quant_dtype (DType): FP4 element dtype (e.g. .uint8 with nibble packing).
• ​scales_dtype (DType): MX scale dtype (.float8_e8m0fnu).
• ​output_layout (TensorLayout): Layout of the FP4 output TileTensor.
• ​scales_layout (TensorLayout): Layout of the scale output TileTensor.
• ​_hid_dim (Int): Hidden dimension; must be divisible by the group size (MXFP4_SF_VECTOR_SIZE).
• ​_top_k (Int): Number of experts each token is routed to.
• ​_alignment (Int): Override for the byte alignment of the wire buffer; 0 selects get_device_alignment().
• ​fuse_a_scale_preshuffle (Bool): When True, fuses the per-expert scale preshuffle into the tile copy for reduced memory traffic.
• ​mx_format (MXFormat): Which MX element encoding the wire payload holds, defaulting to whatever quant_dtype denotes. Pass a six-bit quant_dtype (float6_e2m3fn/float6_e3m2fn) and the format follows; the legacy uint8 spelling is ambiguous -- MXFP4 and MXFP6 both travel as packed uint8 with float8_e8m0fnu scales -- so it resolves to FP4 and an MXFP6 caller stuck on uint8 must name the format here.

Fields

• ​output_tokens (MXTokenFormat[_hid_dim, _top_k, _alignment, fuse_a_scale_preshuffle=fuse_a_scale_preshuffle, mx_format=mx_format].TensorType):
• ​output_scales (MXTokenFormat[_hid_dim, _top_k, _alignment, fuse_a_scale_preshuffle=fuse_a_scale_preshuffle, mx_format=mx_format].ScalesTensorType):
• ​max_padded_M (Int64):

Implemented traits

AnyType, Copyable, Deinitable, DevicePassable, ImplicitlyCopyable, Movable, RegisterPassable, TokenFormat, TrivialRegisterPassable

comptime members

alignment

comptime alignment = _alignment if _alignment.__bool__() else get_device_alignment()

bits_per_element

comptime bits_per_element = mx_format.bits_per_element()

device_type

comptime device_type = MXTokenFormat[_hid_dim, _top_k, _alignment, fuse_a_scale_preshuffle=fuse_a_scale_preshuffle, mx_format=mx_format]

dispatch_smem_size

comptime dispatch_smem_size = 0

dispatch_wait_tile_shape

comptime dispatch_wait_tile_shape = Tuple(Int(32), Int(1))

fp6_fmt

comptime fp6_fmt = mx_format.fp6_format() if mx_format.is_fp6() else FP6Format.E2M3

group_size

comptime group_size = MXFP4_SF_VECTOR_SIZE

hid_dim

comptime hid_dim = _hid_dim

is_fp6

comptime is_fp6 = (mx_format.bits_per_element() == Int(6))

ScalesTensorType

comptime ScalesTensorType = TileTensor[scales_dtype, scales_layout, MutUntrackedOrigin]

TensorType

comptime TensorType = TileTensor[quant_dtype, output_layout, MutUntrackedOrigin]

top_k

comptime top_k = _top_k

Methods

__init__

def __init__(output_tokens: TileTensor[quant_dtype, output_layout, address_space=output_tokens.address_space, linear_idx_type=output_tokens.linear_idx_type], output_scales: TileTensor[scales_dtype, scales_layout, address_space=output_scales.address_space, linear_idx_type=output_scales.linear_idx_type], max_padded_M: Int = Int(0)) -> Self

get_type_name

static def get_type_name() -> String

Returns:

String

quant_size

static def quant_size() -> Int

Returns:

Int

scales_size

static def scales_size() -> Int

Returns:

Int

token_size

static def token_size() -> Int

Returns:

Int

scales_offset

static def scales_offset() -> Int

Returns:

Int

copy_token_to_send_buf

static def copy_token_to_send_buf[src_type: DType, block_size: Int, buf_addr_space: AddressSpace = .GENERIC, thread_base: Int = Int(0)](buf_p: Pointer[UInt8, address_space=buf_addr_space], src_p: Pointer[Scalar[src_type], address_space=src_p.address_space], input_scale: Float32)

copy_msg_to_output_tensor

def copy_msg_to_output_tensor[buf_addr_space: AddressSpace = .GENERIC](self, buf_p: Pointer[UInt8, address_space=buf_addr_space], token_index: Int, expert_slot: Int = Int(0), expert_start: Int = Int(0))