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]

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. DType.uint8 with nibble packing).
• ​scales_dtype (DType): MX scale dtype (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.

Fields

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

Implemented traits

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

comptime members

alignment

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

device_type

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

dispatch_smem_size

comptime dispatch_smem_size = 0

dispatch_wait_tile_shape

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

elems_per_byte

comptime elems_per_byte = Int(2) if (quant_dtype == DType.uint8) else Int(1)

group_size

comptime group_size = MXFP4_SF_VECTOR_SIZE

hid_dim

comptime hid_dim = _hid_dim

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 = AddressSpace.GENERIC](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 = 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))