MXFP4 Dequant
gpt-oss
Medium
Microscaling FP4: dequantize 4-bit expert weights using a shared per-block FP8 scale.
Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.
Overview
Section 2.1 of the model card: MoE weights are post-trained to MXFP4 at (4.25) bits per parameter so the 120b checkpoint fits on one 80GB GPU and the 20b checkpoint in about 16GB. Attention, embeddings, the unembedding, router, and biases stay BF16. The README and gpt_oss/torch/weights.py describe the on-disk layout the official checkpoints actually use.
MXFP4 (OCP Microscaling Formats 1.0) stores many 4-bit element codes that share one scale. gpt-oss uses E2M1 codes packed two per byte, and one E8M0-style scale per block of 32 elements. The educational loader dequantizes those blocks to BF16 before the MoE matmuls. The Triton path can multiply without materializing a full BF16 expert matrix. Either way the decoded real value is the same.
The paper does not specify packing order, LUT, or the scale bias. Those are official-code details.
How it works
A stored MXFP4 tensor is two arrays:
tensor.blocks:uint8, last axis length (B). Each byte holds two FP4 codes.tensor.scales: one scale per block along that last axis. The reference treats the scale as an integer exponent with bias 127 (E8M0 / FP8 exponent).
Block geometry: 32 FP4 values occupy 16 bytes (BYTES_PER_BLOCK = 16). The last dimension of blocks is those 16 bytes; scales.shape == blocks.shape[:-1]. After decode, the last axis expands (B\to 2B), i.e. 16 bytes become 32 floats, and the prefix dimensions are unchanged. A final view fuses the group axis so the weight is a dense [..., in_or_out] matrix the MoE einsum expects.
Each nibble is an index into the E2M1 lookup table used in weights.py:
\mathrm{LUT}[0..15]=\bigl[+0,+\tfrac12,+1,+\tfrac32,+2,+3,+4,+6,\;-0,-\tfrac12,-1,-\tfrac32,-2,-3,-4,-6\bigr].
For a byte (b) the official order is low nibble first, high nibble second:
n_{\mathrm{lo}}=b\bmod 16,\qquad n_{\mathrm{hi}}=\lfloor b/16\rfloor.
If the block’s scale byte (or stored scale value) is (s), every decoded code in that block is multiplied by the same power of two:
w = \mathrm{LUT}[n]\cdot 2^{s-127}.
The loader writes this as scales.to(int32) - 127 followed by torch.ldexp(values, exp). A scale of 127 is unity; 128 doubles; 126 halves. The scale is per block, not per row and not per tensor. Neighboring groups of 32 weights have their own (s).
Only mlp1_weight and mlp2_weight use this pair of names (*.blocks, *.scales) in PARAM_NAME_MAP. mlp1_bias, mlp2_bias, and gate are plain BF16 tensors. After dequant, Transformer.from_checkpoint may shard the BF16 matrix across ranks; the comment in model.py notes it would be cheaper to shard before upcast, but the reference does it after.
The “4.25 bits” figure in the card is the MXFP4 average: 4 bits per element plus a shared 8-bit scale every 32 elements, (4+8/32=4.25).
Official code
gpt_oss/torch/weights.py: FP4_VALUES, BYTES_PER_BLOCK, Checkpoint._get_mxfp4_tensor (chunked) and _get_mxfp4_tensor_copy (short). README section “Precision format” states the two-tensor layout and that scaling is along the last dimension. The Metal conversion script also treats scales as biased exponents (UE8_OFFSET).
Watch-outs
- Low nibble then high nibble. Decoding high-first swaps every adjacent pair and scrambles the SwiGLU even/odd channels on
mlp1. - Subtract 127 from the scale before
ldexp. Using (2^{s}) or (2^{s-126}) (a different FP8 bias) mis-scales every expert. - One scale per 32-value group. Broadcasting a single scale across a whole expert row, or across the last two dims, is not MXFP4 as stored.
- Do not MXFP4-decode attention, embeddings, or the router. Those keys are absent from the blocks/scales map and are already BF16.
Sources
- Paper: gpt-oss-120b & gpt-oss-20b Model Card, §2.1
- Format: Open Compute Project, Microscaling Formats (MX) Specification v1.0, 2023
- Code: openai/gpt-oss
gpt_oss/torch/weights.py(Checkpoint._get_mxfp4_tensor); README “Precision format”