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

MXFP4TokenFormat

struct MXFP4TokenFormat[fp4_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​

  • ​fp4_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 (MXFP4TokenFormat[_hid_dim, _top_k, _alignment, fuse_a_scale_preshuffle=fuse_a_scale_preshuffle].TensorType):
  • ​output_scales (MXFP4TokenFormat[_hid_dim, _top_k, _alignment, fuse_a_scale_preshuffle=fuse_a_scale_preshuffle].ScalesTensorType):
  • ​max_padded_M (Int):

Implemented traits​

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

comptime members​

alignment​

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

device_type​

comptime device_type = MXFP4TokenFormat[_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))

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[fp4_dtype, output_layout, MutUntrackedOrigin]

top_k​

comptime top_k = _top_k

Methods​

__init__​

def __init__(output_tokens: TileTensor[fp4_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

fp4_quant_size​

static def fp4_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, _safe=False], src_p: Pointer[Scalar[src_type], address_space=src_p.address_space, _safe=False], 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, _safe=False], token_index: Int, expert_slot: Int = Int(0), expert_start: Int = Int(0))

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact