From ada91f071047185db1d37e47ba49d8d116130567 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 8 Jun 2026 22:22:34 +0000 Subject: [PATCH] feat(deepseek): add --cast_mxfp4_to_nvfp4 to deepseek_v4 quantize step Add a closed-form, bit-exact MXFP4 -> NVFP4 routed-expert weight cast to examples/deepseek/deepseek_v4/quantize_to_nvfp4.py via a --cast_mxfp4_to_nvfp4 flag. Pins scale_2 = 2^(k_max-8) and each per-block E4M3 scale to 2^(k_j-m) from the source E8M0 scales, so the NVFP4 nibbles equal the source MXFP4 nibbles bit-for-bit for every in-range block. w1/w3 share one scale_2 for the fused GEMM1; activation input_scale still comes from --amax_path calibration. Hoist the shared closed-form numerics (mxfp4_to_nvfp4_global_amax, mxfp4_to_nvfp4_per_block_amax, and the E2M1/E4M3/E8M0 constants) out of the GPT-OSS example cast (examples/llm_ptq/cast_mxfp4_to_nvfp4.py, PR #1372) into modelopt.torch.quantization.utils.numeric_utils, so both the GPT-OSS and DeepSeek-V4 cast paths import them from the library. Their unit tests move to tests/unit/torch/quantization/test_numeric_utils.py; the example test keeps the cast-specific cases (quantizer naming, build_amax_map, apply_to_model). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + examples/deepseek/README.md | 22 +++ .../deepseek/deepseek_v4/quantize_to_nvfp4.py | 168 ++++++++++++++++- examples/llm_ptq/cast_mxfp4_to_nvfp4.py | 156 +-------------- .../torch/quantization/utils/numeric_utils.py | 177 ++++++++++++++++++ .../llm_ptq/test_cast_mxfp4_to_nvfp4.py | 141 -------------- .../torch/quantization/test_numeric_utils.py | 161 ++++++++++++++++ 7 files changed, 532 insertions(+), 294 deletions(-) create mode 100644 modelopt/torch/quantization/utils/numeric_utils.py create mode 100644 tests/unit/torch/quantization/test_numeric_utils.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index da02b315f67..125566673ac 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -32,6 +32,7 @@ Changelog - Add quantization examples for the Megatron-Bridge framework: post-training quantization (`quantize.py `_), export to a deployable HuggingFace checkpoint (`export.py `_), and Quantization Aware Distillation (extend existing `distill.py `_). - Add end-to-end tutorial for Minitron pruning + two-phase distillation (80B @ 8K + 20B @ 32K long-context = 100B tokens) + FP8 PTQ + vLLM deployment for Nemotron-3-Nano-30B-A3B-BF16 (MoE + Mamba-Transformer hybrid) → Pruned 22B/A3.0B active params, along with data blend preparation steps (with tool-calling data) and detailed pruning / data-blend / long-context ablations. See `examples/pruning/minitron/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md `_ for details. - Add ``--cast_mxfp4_to_nvfp4`` flag to ``examples/llm_ptq/hf_ptq.py`` for closed-form, bit-exact MXFP4 → NVFP4 weight conversion. Supports the GPT-OSS family (``openai/gpt-oss-20b``, ``openai/gpt-oss-120b``). See `examples/llm_ptq/README.md `__ for usage. +- Add ``--cast_mxfp4_to_nvfp4`` flag to ``examples/deepseek/deepseek_v4/quantize_to_nvfp4.py`` for closed-form, bit-exact MXFP4 → NVFP4 conversion of DeepSeek V4 routed-expert weights (mirrors the GPT-OSS cast; w1/w3 share one per-tensor ``scale_2`` for the fused GEMM1). Activation ``input_scale`` still comes from ``--amax_path`` calibration. - DeepSeek PTQ (``examples/deepseek/ptq.py``) now defaults to native top-k calibration with post-hoc per-layer peer-max sync of expert ``input_quantizer.amax``; the all-experts path is preserved behind ``--calib_all_experts``. - Add NVFP4 W4A16 weight-only quantization (``w4a16_nvfp4``): FP4 weights with group_size=16, BF16 activations, no calibration forward pass required. Use ``mtq.W4A16_NVFP4_CFG`` or ``--qformat w4a16_nvfp4`` in ``hf_ptq.py``. vLLM deployment support is in progress. - Add Megatron Core export/import mapping for Qwen3-VL (``Qwen3VLForConditionalGeneration``) vision-language models. The mapping handles the ``model.language_model.`` weight prefix used by Qwen3-VL. diff --git a/examples/deepseek/README.md b/examples/deepseek/README.md index a049d820fae..201997bb0ea 100644 --- a/examples/deepseek/README.md +++ b/examples/deepseek/README.md @@ -174,3 +174,25 @@ python deepseek_v4/quantize_to_nvfp4.py \ The output includes an updated `model.safetensors.index.json`, a `config.json` with `quantization_config.moe_quant_algo = "NVFP4"`, and `hf_quant_config.json` describing the mixed NVFP4 expert layers. + +When the source routed experts are MXFP4 (as in the V4 release), add +`--cast_mxfp4_to_nvfp4` for a lossless weight conversion — recommended over the +default lossy dequant/re-quant path. See below. + +#### Lossless MXFP4 → NVFP4 weight cast (`--cast_mxfp4_to_nvfp4`) + +The routed experts in the source checkpoint are already MXFP4 (E2M1 nibbles + +a power-of-two E8M0 scale per 32-element block). Without the flag, the export +dequantizes them to BF16 and re-quantizes to NVFP4 using the calibrated +per-tensor weight amax, which re-derives the per-block scales from the data and +is therefore lossy. With `--cast_mxfp4_to_nvfp4`, the per-tensor `scale_2` is +pinned to `2^(k_max - 8)` and each per-block E4M3 scale to `2^(k_j - m)` straight +from the source E8M0 scales, so `per_block_scale * scale_2 = 2^k_j` and the NVFP4 +nibbles equal the source MXFP4 nibbles bit-for-bit (for every block whose `k_j` +lands in E4M3's representable window; the rare out-of-range block falls back to a +data-derived scale). The flag only affects routed-expert **weights** — activation +`input_scale` still comes from `${AMAX}` calibration — and the run prints a +`[cast] lossless MXFP4->NVFP4 blocks: …` summary. This mirrors the GPTOSS cast in +[`examples/llm_ptq/cast_mxfp4_to_nvfp4.py`](../llm_ptq/cast_mxfp4_to_nvfp4.py); the +V4 twist is that w1/w3 share one `scale_2` (fused GEMM1), so `k_max` is taken over +both projections. diff --git a/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py b/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py index 6906df7ab45..0f4d0198676 100644 --- a/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py +++ b/examples/deepseek/deepseek_v4/quantize_to_nvfp4.py @@ -63,6 +63,20 @@ for the same projection. If no calibrated expert exists for that projection, export fails. +Lossless weight cast (``--cast_mxfp4_to_nvfp4``): the source routed experts are +already MXFP4 (E2M1 nibbles + a power-of-two E8M0 scale per 32-element block). +By default this script dequantizes them to BF16 and re-quantizes to NVFP4 with +the calibrated per-tensor weight amax, which re-derives per-block scales from +the data and is therefore lossy. With ``--cast_mxfp4_to_nvfp4`` we instead pin +``scale_2 = 2^(k_max - 8)`` and the per-block E4M3 scale to ``2^(k_j - m)`` +straight from the source E8M0 scales, so ``per_block_scale * scale_2 = 2^k_j`` +and the NVFP4 nibbles equal the source MXFP4 nibbles bit-for-bit (for every +block whose ``k_j`` lands in E4M3's representable window). The flag only affects +routed-expert *weights*; activation ``input_scale`` still comes from +``--amax_path`` calibration. This mirrors the GPTOSS cast in +``examples/llm_ptq/cast_mxfp4_to_nvfp4.py`` (PR #1372); the V4 twist is that +w1/w3 share one ``scale_2`` (fused GEMM1), so ``k_max`` is taken over both. + Usage (single compute node, CPU-default; dequant+requant math is cheap relative to shard I/O): @@ -91,6 +105,17 @@ from modelopt.torch.quantization.qtensor import MXFP4QTensor, NVFP4QTensor +# Closed-form MXFP4 -> NVFP4 numerics shared with the GPT-OSS cast (PR #1372). +from modelopt.torch.quantization.utils.numeric_utils import ( + E2M1_MAX, + E4M3_KMAX, + E4M3_KMIN, + E4M3_MAX, + E8M0_BIAS, + mxfp4_to_nvfp4_global_amax, + mxfp4_to_nvfp4_per_block_amax, +) + # Routed-expert weights in regular MoE layers. MTP experts remain in source format. _EXPERT_WEIGHT_RE = re.compile(r"^layers\.\d+\.ffn\.experts\.\d+\.w[123]\.weight$") _EXPERT_PROJ_RE = re.compile(r"^(?Players\.\d+\.ffn\.experts)\.\d+\.w[123]$") @@ -233,6 +258,98 @@ def _quantize_weight_nvfp4( return q_tensor._quantized_data, weight_scale, weight_scale_2, synthesized +# --------------------------------------------------------------------------- +# Lossless MXFP4 -> NVFP4 weight cast (``--cast_mxfp4_to_nvfp4``). +# +# NVFP4 uses the same E2M1 nibble grid as MXFP4 with 16-element blocks and a +# two-level scale ``per_block_scale (E4M3) * scale_2 (fp32)``. Pinning +# ``scale_2 = 2^m`` (``m = k_max - 8``) and ``per_block_scale = 2^(k_j - m)`` +# makes ``per_block_scale * scale_2 = 2^k_j`` exactly, so each NVFP4 nibble +# equals the source MXFP4 nibble verbatim — bit-exact for every block whose +# ``k_j`` lands in E4M3's window (``k_max - k_j <= 17``). The closed-form +# per-block amax and the format constants are reused from the GPT-OSS cast +# (``cast_mxfp4_to_nvfp4``, PR #1372); the V4 twist is that w1/w3 share one +# ``scale_2`` (fused GEMM1), so ``k_max`` is taken over both projections. +# --------------------------------------------------------------------------- +_NVFP4_BLOCK = 16 # NVFP4 block size (elements) +_MXFP4_BYTES_PER_BLOCK = 16 # 32 E2M1 nibbles packed 2-per-byte + + +def _kmax_from_mxfp4_scale(mxfp4_scale: torch.Tensor, device: str = "cpu") -> int: + """Largest non-zero E8M0 exponent ``k_j = e8m0 - 127`` (0 if all-zero). + + Delegates to the GPT-OSS cast's ``k_max`` logic, which excludes the + all-zero sentinel (``e8m0 == 0`` => ``k == -127``). + """ + e8m0 = mxfp4_scale.to(device).contiguous().view(torch.uint8) + return mxfp4_to_nvfp4_global_amax(e8m0)[1]["k_max"] + + +def _build_w13_kmax_overrides(f, expert_weight_keys: list[str], device: str) -> dict[str, int]: + """Shared ``k_max`` per w1/w3 pair so the fused GEMM1 gets one ``scale_2``.""" + groups: dict[str, dict[str, str]] = defaultdict(dict) + for key in expert_weight_keys: + expert_path = key[: -len(".weight")] + base, proj = expert_path.rsplit(".", 1) + if proj in {"w1", "w3"}: + groups[base][proj] = expert_path + + overrides: dict[str, int] = {} + for paths in groups.values(): + if "w1" not in paths or "w3" not in paths: + continue + k1 = _kmax_from_mxfp4_scale(f.get_tensor(paths["w1"] + ".scale"), device) + k3 = _kmax_from_mxfp4_scale(f.get_tensor(paths["w3"] + ".scale"), device) + shared = max(k1, k3) + overrides[paths["w1"]] = shared + overrides[paths["w3"]] = shared + return overrides + + +def _quantize_weight_nvfp4_lossless( + mxfp4_weight: torch.Tensor, + mxfp4_scale: torch.Tensor, + k_max: int, + device: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Closed-form bit-exact MXFP4 -> NVFP4 weight conversion. + + Pins ``scale_2 = 2^(k_max - 8)`` and the per-block E4M3 scale to + ``2^(k_j - m)`` so the NVFP4 nibbles equal the source MXFP4 nibbles for + every in-range block. ``k_max`` is shared across w1/w3 (fused GEMM1), so it + is passed in rather than derived per tensor. The closed-form per-block amax + (``6 * 2^k_j`` in range, data-derived out of range) is independent of + ``k_max``, so we reuse the GPT-OSS helper directly. Returns + ``(packed, weight_scale, weight_scale_2, n_blocks, n_lossless)``. + """ + bf16 = _dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device) + e8m0 = mxfp4_scale.to(bf16.device).contiguous().view(torch.uint8) # (out, nblk32) + packed = mxfp4_weight.to(bf16.device).contiguous().view(torch.uint8) # (out, nblk32*16) + blocks = packed.view(*packed.shape[:-1], e8m0.shape[-1], _MXFP4_BYTES_PER_BLOCK) + per_block_amax = mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) # (out, nblk16) fp32 + + m = k_max - E4M3_KMAX + weight_scale_2 = torch.tensor(2.0**m, dtype=torch.float32, device=bf16.device).reshape(()) + per_block_scale = ( + (per_block_amax / (E2M1_MAX * weight_scale_2)) + .clamp(min=2**-9, max=E4M3_MAX) + .to(torch.float8_e4m3fn) + ) + + # Lossless accounting against the (possibly shared) k_max. A block is lossy + # only if k_max - k_j > 17; all-zero blocks (e8m0 == 0) reconstruct to 0 + # regardless of scale and so are always lossless. + k = e8m0.to(torch.int32) - E8M0_BIAS + lossless = (k >= (k_max - (E4M3_KMAX - E4M3_KMIN))) | (e8m0 == 0) + n_blocks = k.numel() + n_lossless = int(lossless.sum().item()) + + q_tensor, weight_scale, _ = NVFP4QTensor.quantize( + bf16, _NVFP4_BLOCK, per_block_scale, weight_scale_2, try_tensorrt=False + ) + return q_tensor._quantized_data, weight_scale, weight_scale_2, n_blocks, n_lossless + + def _build_w13_weight_amax_overrides( f, expert_weight_keys: list[str], @@ -279,6 +396,7 @@ def convert_shard( input_fallback: dict[str, torch.Tensor], device: str, stats: dict[str, int], + cast: bool = False, ) -> tuple[list[str], list[str]]: """Rewrite one HF-style shard and return index deltas.""" out: dict[str, torch.Tensor] = {} @@ -289,9 +407,16 @@ def convert_shard( all_keys = list(f.keys()) expert_weight_keys = [k for k in all_keys if _EXPERT_WEIGHT_RE.match(k)] expert_weight_key_set = set(expert_weight_keys) - w13_weight_amax, w13_synth_paths = _build_w13_weight_amax_overrides( - f, expert_weight_keys, amax, device - ) + if cast: + # Closed-form weight cast derives scales from the source E8M0 + # exponents, not from calibrated weight amax. w1/w3 share k_max. + w13_kmax = _build_w13_kmax_overrides(f, expert_weight_keys, device) + w13_weight_amax, w13_synth_paths = {}, set() + else: + w13_kmax = {} + w13_weight_amax, w13_synth_paths = _build_w13_weight_amax_overrides( + f, expert_weight_keys, amax, device + ) scale_siblings = { k[: -len(".weight")] + ".scale" for k in expert_weight_keys @@ -335,9 +460,22 @@ def convert_shard( w = f.get_tensor(key) s = f.get_tensor(scale_key) - packed, weight_scale, weight_scale_2, weight_synth = _quantize_weight_nvfp4( - w, s, weight_amax, device=device - ) + if cast: + k_max = w13_kmax.get(expert_path) + if k_max is None: + k_max = _kmax_from_mxfp4_scale(s, device) + packed, weight_scale, weight_scale_2, n_blk, n_lossless = ( + _quantize_weight_nvfp4_lossless(w, s, k_max, device) + ) + weight_synth = False + stats["cast_blocks_total"] += n_blk + stats["cast_blocks_lossless"] += n_lossless + if n_lossless < n_blk: + stats[f"cast_oor_tensors_{block_kind}"] += 1 + else: + packed, weight_scale, weight_scale_2, weight_synth = _quantize_weight_nvfp4( + w, s, weight_amax, device=device + ) input_scale = _amax_to_nvfp4_scale_2(input_amax).to(weight_scale_2.device) out[key] = packed.cpu() @@ -607,6 +745,17 @@ def main(): action="store_true", help="replace an existing non-empty output checkpoint directory", ) + p.add_argument( + "--cast_mxfp4_to_nvfp4", + action="store_true", + help=( + "losslessly cast the source MXFP4 routed-expert weights to NVFP4 " + "(pin scale_2 = 2^(k_max-8) and per-block scale = 2^(k_j-m) from the " + "source E8M0 scales) instead of dequant + re-quant with calibrated " + "weight amax. Only affects weights; input_scale still comes from " + "--amax_path calibration." + ), + ) args = p.parse_args() _validate_paths(args.source_ckpt, args.output_ckpt) @@ -639,6 +788,7 @@ def main(): input_fallback, args.device, stats, + args.cast_mxfp4_to_nvfp4, ) shard_updates[src.name] = (added, removed) @@ -647,6 +797,12 @@ def main(): for k in sorted(stats.keys()): _log(f" {k:40s} {stats[k]}") + if args.cast_mxfp4_to_nvfp4: + tot = stats.get("cast_blocks_total", 0) + loss = stats.get("cast_blocks_lossless", 0) + pct = 100.0 * loss / tot if tot else 100.0 + _log(f"[cast] lossless MXFP4->NVFP4 blocks: {loss}/{tot} ({pct:.4f}%)") + quantized: set[str] = set() for _added, _removed in shard_updates.values(): for a in _added: diff --git a/examples/llm_ptq/cast_mxfp4_to_nvfp4.py b/examples/llm_ptq/cast_mxfp4_to_nvfp4.py index 90b8521e5d7..e8e86966f3b 100644 --- a/examples/llm_ptq/cast_mxfp4_to_nvfp4.py +++ b/examples/llm_ptq/cast_mxfp4_to_nvfp4.py @@ -41,6 +41,12 @@ from safetensors import safe_open from modelopt.torch.quantization.nn.modules.tensor_quantizer import NVFP4StaticQuantizer +from modelopt.torch.quantization.utils.numeric_utils import ( + E2M1_MAX, + E8M0_BIAS, + mxfp4_to_nvfp4_global_amax, + mxfp4_to_nvfp4_per_block_amax, +) @contextmanager @@ -62,150 +68,6 @@ def read(key: str, shard: Path) -> torch.Tensor: yield read -E8M0_BIAS = 127 # E8M0 stores k_j as uint8 with bias 127 -E2M1_MAX = 6.0 -E4M3_MAX = 448.0 -E4M3_KMAX = 8 -E4M3_KMIN = -9 # E4M3 represents 2^k exactly for k in [-9, 8] -# E2M1 magnitude grid indexed by the low 3 bits of an FP4 nibble. -_E2M1_MAGNITUDE = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] -# Cache of the E2M1 magnitude lookup table per (device, dtype) so we don't -# rebuild it for every layer in a batched cast. -_E2M1_MAG_CACHE: "dict[tuple, torch.Tensor]" = {} - - -def _e2m1_magnitude_table(device: torch.device, dtype: torch.dtype = torch.float32) -> torch.Tensor: - """Return ``_E2M1_MAGNITUDE`` as a tensor on the requested device, cached.""" - key = (device, dtype) - cached = _E2M1_MAG_CACHE.get(key) - if cached is None: - cached = torch.tensor(_E2M1_MAGNITUDE, dtype=dtype, device=device) - _E2M1_MAG_CACHE[key] = cached - return cached - - -def compute_global_amax_for_scales(e8m0_scales: torch.Tensor) -> tuple[float, dict]: - """Closed-form per-tensor ``global_amax``: ``m = k_max - 8``, ``global_amax = 6 * 448 * 2^m``. - - Args: - e8m0_scales: uint8 tensor of E8M0 scales for one MXFP4 source layer. - - Returns: - global_amax: scalar (float) — pins NVFP4 scale_2 to 2^m. - info: diagnostic dict with k_min, k_max, m, lossless-block stats. - """ - # k_j = e8m0 - 127. MXFP4 quantize emits e8m0=0 (=> k=-127) for all-zero - # blocks; treat those as "ignore me" when computing k_max. - k = e8m0_scales.to(torch.int32) - E8M0_BIAS - nonzero_mask = e8m0_scales > 0 - if nonzero_mask.any(): - k_nonzero = k[nonzero_mask] - k_min = int(k_nonzero.min().item()) - k_max = int(k_nonzero.max().item()) - else: - k_min = k_max = 0 - - m = k_max - E4M3_KMAX - global_amax = E2M1_MAX * E4M3_MAX * float(2.0**m) - - # A block is lossless under this cast iff k_max - k_j <= 17 (its k_j - m sits - # in E4M3's [-9, 8] window). All-zero blocks are trivially lossless because - # their reconstruction is 0 regardless of the snapped scale. - n_total = e8m0_scales.numel() - in_range = (k >= (k_max - 17)) | (~nonzero_mask) - n_lossless = int(in_range.sum().item()) - pct_lossless = 100.0 * n_lossless / n_total if n_total else 100.0 - - return global_amax, { - "k_min": k_min, - "k_max": k_max, - "m": m, - "n_total_blocks": n_total, - "n_lossless_blocks": n_lossless, - "pct_lossless": pct_lossless, - "n_zero_blocks": int((~nonzero_mask).sum().item()), - } - - -def compute_per_block_amax_for_mxfp4( - blocks: torch.Tensor, e8m0_scales: torch.Tensor -) -> torch.Tensor: - """Hybrid per-NVFP4-block amax for MXFP4 -> NVFP4 cast. - - Each MXFP4 block of 32 elements has one E8M0 exponent ``k_j``. Two cases - based on whether ``k_j`` fits in NVFP4's E4M3 scale grid (with - ``m = k_max - 8`` chosen by ``compute_global_amax_for_scales``): - - - **In-range** (``k_j - m`` in ``[-9, 8]``): ``6 * 2^k_j`` (closed-form - ideal). The resulting per-block scale ``2^(k_j - m)`` is exactly - representable in E4M3 — no rounding loss — and - ``round_to_E2M1(value / 2^k_j)`` yields the original MXFP4 nibble - verbatim. Bit-exact reconstruction. - - - **Out of range** (``|k_j - m| > 8/9``): ``max_nibble * 2^k_j``, i.e. - ``max(|w_block|)`` where ``w`` is the MXFP4-dequantized block. This is - the data-derived per-block amax. The per-block scale will still get - clamped at the E4M3 boundary, but data-derived amax keeps the post-clamp - scale closer to the block's actual magnitude than the closed-form ideal - would, which reduces re-bucketing error for OOR blocks where - ``max_nibble < 6``. - - Two NVFP4 blocks of 16 share each MXFP4 block's ``k_j``, so the result is - expanded by ``repeat_interleave(2, dim=-1)``. - - Args: - blocks: uint8 tensor of packed E2M1 nibbles, shape - ``(..., num_mxfp4_blocks, 16)`` (16 bytes per 32-element MXFP4 block). - e8m0_scales: uint8 tensor of E8M0 scales, shape - ``(..., num_mxfp4_blocks)``. - - Returns: - float32 tensor of shape ``(..., 2 * num_mxfp4_blocks)``. - """ - if blocks.shape[-1] != 16 or blocks.shape[:-1] != e8m0_scales.shape: - raise ValueError( - f"shape mismatch: blocks {tuple(blocks.shape)} " - "(expected (..., num_mxfp4_blocks, 16)) " - f"vs scales {tuple(e8m0_scales.shape)}" - ) - - k = e8m0_scales.to(torch.int32) - E8M0_BIAS # (..., num_mxfp4_blocks) - pow2_k = torch.exp2(k.float()) - closed_form_ideal = E2M1_MAX * pow2_k # (..., num_mxfp4_blocks) - - # ``m = k_max - 8`` over non-zero blocks. Compute via masked ``amax`` so - # ``m`` stays a 0-d tensor and we avoid a GPU->CPU sync just to get a - # Python int. All-zero scales fall through with the -E8M0_BIAS sentinel, - # which leaves every block trivially in-range (closed_form_ideal == 0 there). - nonzero = e8m0_scales > 0 - sentinel = torch.full_like(k, -E8M0_BIAS) - k_max = torch.where(nonzero, k, sentinel).amax() - delta = k - (k_max - E4M3_KMAX) - in_range = (delta >= E4M3_KMIN) & (delta <= E4M3_KMAX) - - # Fast path: if every block fits E4M3's [-9, 8] window the per-block amax - # is just the closed-form ideal, and we can skip the per-byte nibble scan - # over the block tensor (which is 16x larger than the scales). For typical - # MXFP4 checkpoints (e.g. gpt-oss-20b) this is the only path ever taken. - if bool(in_range.all()): - return closed_form_ideal.repeat_interleave(2, dim=-1) - - # OOR fallback: data-derived per-block amax = max(|w_block|) after MXFP4 - # dequant = ``max_nibble * 2^k_j``. The MXFP4 nibble is sign-magnitude with - # sign in bit 3 and magnitude index in bits 0-2; we extract per-byte - # magnitudes, take the byte-wise max, then reduce across the 16 bytes to - # get the largest magnitude index in the 32-element block. - low = blocks & 0x07 - high = (blocks >> 4) & 0x07 - max_idx = torch.maximum(low, high).amax(dim=-1).long() - max_nibble = _e2m1_magnitude_table(blocks.device)[max_idx] - data_derived = max_nibble * pow2_k - - per_block_amax_mxfp4 = torch.where(in_range, closed_form_ideal, data_derived) - # Each MXFP4 block of 32 splits into two NVFP4 blocks of 16 sharing k_j. - return per_block_amax_mxfp4.repeat_interleave(2, dim=-1) - - def quantizer_name_from_blocks_key(blocks_key: str) -> str: """Map ``_blocks`` -> ``_weight_quantizer``. @@ -282,7 +144,7 @@ def build_amax_map(checkpoint_dir: str | Path) -> dict[str, dict]: for tensor_key, shard in sorted(scales_keys.items()): scales = read(tensor_key, shard) - global_amax, info = compute_global_amax_for_scales(scales) + global_amax, info = mxfp4_to_nvfp4_global_amax(scales) blocks_key = tensor_key[: -len("_scales")] + "_blocks" qname = quantizer_name_from_blocks_key(blocks_key) @@ -357,7 +219,7 @@ def apply_to_model( for tensor_key, shard in sorted(scales_keys.items()): scales = read(tensor_key, shard) - global_amax_value, info = compute_global_amax_for_scales(scales) + global_amax_value, info = mxfp4_to_nvfp4_global_amax(scales) n_total_layers += 1 if info["pct_lossless"] >= 100.0: n_lossless_layers += 1 @@ -410,7 +272,7 @@ def apply_to_model( ) else: blocks = read(blocks_key, blocks_shard) - per_block_amax = compute_per_block_amax_for_mxfp4(blocks, scales).to( + per_block_amax = mxfp4_to_nvfp4_per_block_amax(blocks, scales).to( dtype=torch.float32, device=device ) # Numel must match — calibration may store ``_amax`` flat (e.g. (N, 1)) diff --git a/modelopt/torch/quantization/utils/numeric_utils.py b/modelopt/torch/quantization/utils/numeric_utils.py new file mode 100644 index 00000000000..b48a1ff5b28 --- /dev/null +++ b/modelopt/torch/quantization/utils/numeric_utils.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Closed-form numerics for the bit-exact MXFP4 -> NVFP4 weight cast. + +These helpers turn an MXFP4 source layer's E8M0 block scales into the per-tensor +``global_amax`` and per-NVFP4-block ``amax`` that pin NVFP4's two-level scale so +the cast reproduces the source MXFP4 weights bit-for-bit (see PR #1372 for the +derivation). They are pure tensor math with no model or checkpoint dependencies, +shared by the GPT-OSS (``examples/llm_ptq``) and DeepSeek-V4 +(``examples/deepseek``) PTQ cast paths. +""" + +import torch + +__all__ = [ + "E2M1_MAX", + "E4M3_KMAX", + "E4M3_KMIN", + "E4M3_MAX", + "E8M0_BIAS", + "mxfp4_to_nvfp4_global_amax", + "mxfp4_to_nvfp4_per_block_amax", +] + +E8M0_BIAS = 127 # E8M0 stores k_j as uint8 with bias 127 +E2M1_MAX = 6.0 +E4M3_MAX = 448.0 +E4M3_KMAX = 8 +E4M3_KMIN = -9 # E4M3 represents 2^k exactly for k in [-9, 8] +# E2M1 magnitude grid indexed by the low 3 bits of an FP4 nibble. +_E2M1_MAGNITUDE = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +# Cache of the E2M1 magnitude lookup table per (device, dtype) so we don't +# rebuild it for every layer in a batched cast. +_E2M1_MAG_CACHE: "dict[tuple, torch.Tensor]" = {} + + +def _e2m1_magnitude_table(device: torch.device, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Return ``_E2M1_MAGNITUDE`` as a tensor on the requested device, cached.""" + key = (device, dtype) + cached = _E2M1_MAG_CACHE.get(key) + if cached is None: + cached = torch.tensor(_E2M1_MAGNITUDE, dtype=dtype, device=device) + _E2M1_MAG_CACHE[key] = cached + return cached + + +def mxfp4_to_nvfp4_global_amax(e8m0_scales: torch.Tensor) -> tuple[float, dict]: + """Closed-form per-tensor ``global_amax``: ``m = k_max - 8``, ``global_amax = 6 * 448 * 2^m``. + + Args: + e8m0_scales: uint8 tensor of E8M0 scales for one MXFP4 source layer. + + Returns: + global_amax: scalar (float) — pins NVFP4 scale_2 to 2^m. + info: diagnostic dict with k_min, k_max, m, lossless-block stats. + """ + # k_j = e8m0 - 127. MXFP4 quantize emits e8m0=0 (=> k=-127) for all-zero + # blocks; treat those as "ignore me" when computing k_max. + k = e8m0_scales.to(torch.int32) - E8M0_BIAS + nonzero_mask = e8m0_scales > 0 + if nonzero_mask.any(): + k_nonzero = k[nonzero_mask] + k_min = int(k_nonzero.min().item()) + k_max = int(k_nonzero.max().item()) + else: + k_min = k_max = 0 + + m = k_max - E4M3_KMAX + global_amax = E2M1_MAX * E4M3_MAX * float(2.0**m) + + # A block is lossless under this cast iff k_max - k_j <= 17 (its k_j - m sits + # in E4M3's [-9, 8] window). All-zero blocks are trivially lossless because + # their reconstruction is 0 regardless of the snapped scale. + n_total = e8m0_scales.numel() + in_range = (k >= (k_max - 17)) | (~nonzero_mask) + n_lossless = int(in_range.sum().item()) + pct_lossless = 100.0 * n_lossless / n_total if n_total else 100.0 + + return global_amax, { + "k_min": k_min, + "k_max": k_max, + "m": m, + "n_total_blocks": n_total, + "n_lossless_blocks": n_lossless, + "pct_lossless": pct_lossless, + "n_zero_blocks": int((~nonzero_mask).sum().item()), + } + + +def mxfp4_to_nvfp4_per_block_amax(blocks: torch.Tensor, e8m0_scales: torch.Tensor) -> torch.Tensor: + """Hybrid per-NVFP4-block amax for MXFP4 -> NVFP4 cast. + + Each MXFP4 block of 32 elements has one E8M0 exponent ``k_j``. Two cases + based on whether ``k_j`` fits in NVFP4's E4M3 scale grid (with + ``m = k_max - 8`` chosen by ``mxfp4_to_nvfp4_global_amax``): + + - **In-range** (``k_j - m`` in ``[-9, 8]``): ``6 * 2^k_j`` (closed-form + ideal). The resulting per-block scale ``2^(k_j - m)`` is exactly + representable in E4M3 — no rounding loss — and + ``round_to_E2M1(value / 2^k_j)`` yields the original MXFP4 nibble + verbatim. Bit-exact reconstruction. + + - **Out of range** (``|k_j - m| > 8/9``): ``max_nibble * 2^k_j``, i.e. + ``max(|w_block|)`` where ``w`` is the MXFP4-dequantized block. This is + the data-derived per-block amax. The per-block scale will still get + clamped at the E4M3 boundary, but data-derived amax keeps the post-clamp + scale closer to the block's actual magnitude than the closed-form ideal + would, which reduces re-bucketing error for OOR blocks where + ``max_nibble < 6``. + + Two NVFP4 blocks of 16 share each MXFP4 block's ``k_j``, so the result is + expanded by ``repeat_interleave(2, dim=-1)``. + + Args: + blocks: uint8 tensor of packed E2M1 nibbles, shape + ``(..., num_mxfp4_blocks, 16)`` (16 bytes per 32-element MXFP4 block). + e8m0_scales: uint8 tensor of E8M0 scales, shape + ``(..., num_mxfp4_blocks)``. + + Returns: + float32 tensor of shape ``(..., 2 * num_mxfp4_blocks)``. + """ + if blocks.shape[-1] != 16 or blocks.shape[:-1] != e8m0_scales.shape: + raise ValueError( + f"shape mismatch: blocks {tuple(blocks.shape)} " + "(expected (..., num_mxfp4_blocks, 16)) " + f"vs scales {tuple(e8m0_scales.shape)}" + ) + + k = e8m0_scales.to(torch.int32) - E8M0_BIAS # (..., num_mxfp4_blocks) + pow2_k = torch.exp2(k.float()) + closed_form_ideal = E2M1_MAX * pow2_k # (..., num_mxfp4_blocks) + + # ``m = k_max - 8`` over non-zero blocks. Compute via masked ``amax`` so + # ``m`` stays a 0-d tensor and we avoid a GPU->CPU sync just to get a + # Python int. All-zero scales fall through with the -E8M0_BIAS sentinel, + # which leaves every block trivially in-range (closed_form_ideal == 0 there). + nonzero = e8m0_scales > 0 + sentinel = torch.full_like(k, -E8M0_BIAS) + k_max = torch.where(nonzero, k, sentinel).amax() + delta = k - (k_max - E4M3_KMAX) + in_range = (delta >= E4M3_KMIN) & (delta <= E4M3_KMAX) + + # Fast path: if every block fits E4M3's [-9, 8] window the per-block amax + # is just the closed-form ideal, and we can skip the per-byte nibble scan + # over the block tensor (which is 16x larger than the scales). For typical + # MXFP4 checkpoints (e.g. gpt-oss-20b) this is the only path ever taken. + if bool(in_range.all()): + return closed_form_ideal.repeat_interleave(2, dim=-1) + + # OOR fallback: data-derived per-block amax = max(|w_block|) after MXFP4 + # dequant = ``max_nibble * 2^k_j``. The MXFP4 nibble is sign-magnitude with + # sign in bit 3 and magnitude index in bits 0-2; we extract per-byte + # magnitudes, take the byte-wise max, then reduce across the 16 bytes to + # get the largest magnitude index in the 32-element block. + low = blocks & 0x07 + high = (blocks >> 4) & 0x07 + max_idx = torch.maximum(low, high).amax(dim=-1).long() + max_nibble = _e2m1_magnitude_table(blocks.device)[max_idx] + data_derived = max_nibble * pow2_k + + per_block_amax_mxfp4 = torch.where(in_range, closed_form_ideal, data_derived) + # Each MXFP4 block of 32 splits into two NVFP4 blocks of 16 sharing k_j. + return per_block_amax_mxfp4.repeat_interleave(2, dim=-1) diff --git a/tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py b/tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py index b6f8c3de123..c6446d27b93 100644 --- a/tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py +++ b/tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py @@ -32,137 +32,6 @@ import cast_mxfp4_to_nvfp4 as cast -# ---------- compute_global_amax_for_scales ---------------------------------- - - -def test_global_amax_basic_in_range(): - """Mixed in-range scales: m = k_max - 8, global_amax = 6*448*2^m, lossless = 100%.""" - # k values in [-3, 3] (spread = 6), all blocks lossless. - k = torch.tensor([0, -3, 3, 1, -1, 2], dtype=torch.int32) - e8m0 = (k + cast.E8M0_BIAS).to(torch.uint8) - - global_amax, info = cast.compute_global_amax_for_scales(e8m0) - assert info["k_min"] == -3 - assert info["k_max"] == 3 - assert info["m"] == 3 - 8 # k_max - 8 = -5 - expected = 6.0 * 448.0 * 2.0 ** info["m"] - assert global_amax == pytest.approx(expected) - assert info["n_total_blocks"] == 6 - assert info["n_lossless_blocks"] == 6 - assert info["pct_lossless"] == pytest.approx(100.0) - assert info["n_zero_blocks"] == 0 - - -def test_global_amax_with_zero_blocks(): - """Zero (e8m0=0, k=-127) blocks should be ignored when computing k_max.""" - e8m0 = torch.tensor([0, 0, 130, 125], dtype=torch.uint8) # ks: -127, -127, 3, -2 - global_amax, info = cast.compute_global_amax_for_scales(e8m0) - assert info["k_max"] == 3 # ignores zero blocks - assert info["n_zero_blocks"] == 2 - # Both nonzero blocks satisfy k_max - k_j <= 17, plus zero blocks count as - # lossless because their reconstruction is 0 regardless of scale. - assert info["n_lossless_blocks"] == 4 - - -def test_global_amax_with_oor_blocks(): - """A block 18 powers below k_max is OOR (k_max - k = 18 > 17).""" - # k values: 5, 5, -13 → spread = 18, last block is OOR. - k = torch.tensor([5, 5, -13], dtype=torch.int32) - e8m0 = (k + cast.E8M0_BIAS).to(torch.uint8) - _, info = cast.compute_global_amax_for_scales(e8m0) - assert info["k_max"] == 5 - assert info["n_total_blocks"] == 3 - assert info["n_lossless_blocks"] == 2 # the k=-13 block is OOR - - -def test_global_amax_all_zero(): - """All-zero scales should not crash; k_max defaults to 0.""" - e8m0 = torch.zeros(4, dtype=torch.uint8) - global_amax, info = cast.compute_global_amax_for_scales(e8m0) - assert info["k_min"] == 0 and info["k_max"] == 0 - assert info["n_zero_blocks"] == 4 - # All blocks count as "lossless" (their dequant is 0 regardless of scale). - assert info["n_lossless_blocks"] == 4 - - -# ---------- compute_per_block_amax_for_mxfp4 -------------------------------- - - -def _make_blocks_with_max_nibble(num_blocks: int, max_idx_per_block: list[int]) -> torch.Tensor: - """Build a (num_blocks, 16) byte tensor where block i has E2M1 magnitude - index ``max_idx_per_block[i]`` as its largest nibble; other nibbles are 0. - - Magnitude index goes in the low 3 bits of one nibble; we place it in the - high nibble of byte 0 (so the first byte = (max_idx << 4)). Every other - nibble is 0, so the block-wise max is exactly ``max_idx_per_block[i]``. - """ - assert len(max_idx_per_block) == num_blocks - blocks = torch.zeros((num_blocks, 16), dtype=torch.uint8) - for i, idx in enumerate(max_idx_per_block): - assert 0 <= idx < 8 - blocks[i, 0] = (idx & 0x07) << 4 - return blocks - - -def test_per_block_amax_in_range_returns_closed_form(): - """Every block in-range -> 6 * 2^k_j, regardless of actual nibble content.""" - # k = [0, -2, 4]; k_max = 4, k_min = -2, spread 6 (in-range). - k = torch.tensor([0, -2, 4], dtype=torch.int32) - e8m0 = (k + cast.E8M0_BIAS).to(torch.uint8) - # Blocks have varying max_nibbles, but in-range path ignores them. - blocks = _make_blocks_with_max_nibble(3, [3, 7, 1]) # max nibbles: 1.5, 6, 0.5 - - out = cast.compute_per_block_amax_for_mxfp4(blocks, e8m0) - expected_mxfp4 = 6.0 * torch.exp2(k.float()) # ignores max_nibble - expected_nvfp4 = expected_mxfp4.repeat_interleave(2, dim=-1) - assert torch.allclose(out, expected_nvfp4) - - -def test_per_block_amax_oor_uses_data_derived(): - """OOR blocks should use ``max_nibble * 2^k_j`` (data-derived).""" - # k_max=10 → m=2. OOR-low blocks have k_j - m < -9, i.e. k_j < -7. - k = torch.tensor([10, -10], dtype=torch.int32) # second is OOR-low - e8m0 = (k + cast.E8M0_BIAS).to(torch.uint8) - # Block 0 max nibble idx 7 (value 6); block 1 max nibble idx 4 (value 2). - blocks = _make_blocks_with_max_nibble(2, [7, 4]) - - out = cast.compute_per_block_amax_for_mxfp4(blocks, e8m0) - - # Block 0 (in-range): 6 * 2^10 = 6144. - # Block 1 (OOR): 2 * 2^-10 (max_nibble=2 since idx=4 -> 2.0). - expected_mxfp4 = torch.tensor([6.0 * 2**10, 2.0 * 2**-10], dtype=torch.float32) - expected_nvfp4 = expected_mxfp4.repeat_interleave(2, dim=-1) - assert torch.allclose(out, expected_nvfp4) - - -def test_per_block_amax_doubles_last_dim(): - """Two NVFP4 blocks per MXFP4 block share the same per-block amax.""" - e8m0 = torch.tensor([130, 124], dtype=torch.uint8) # ks: 3, -3 - blocks = _make_blocks_with_max_nibble(2, [7, 7]) # in-range - out = cast.compute_per_block_amax_for_mxfp4(blocks, e8m0) - assert out.shape == (4,) - # Each pair of consecutive entries should be equal. - assert out[0] == out[1] - assert out[2] == out[3] - - -def test_per_block_amax_preserves_leading_dims(): - """Leading dims (E, F, ...) flow through unchanged; only last dim doubles.""" - # shape (E=2, F=3, num_mxfp4_blocks=4) - e8m0 = torch.full((2, 3, 4), 128, dtype=torch.uint8) # all k=1, in-range - blocks = torch.zeros((2, 3, 4, 16), dtype=torch.uint8) - out = cast.compute_per_block_amax_for_mxfp4(blocks, e8m0) - assert out.shape == (2, 3, 8) - - -def test_per_block_amax_shape_mismatch_raises(): - """Mismatched leading dims should raise ``ValueError``.""" - blocks = torch.zeros((4, 16), dtype=torch.uint8) - e8m0 = torch.zeros(3, dtype=torch.uint8) # different num_blocks - with pytest.raises(ValueError, match="shape mismatch"): - cast.compute_per_block_amax_for_mxfp4(blocks, e8m0) - - # ---------- quantizer_name_from_blocks_key ---------------------------------- @@ -272,16 +141,6 @@ def test_build_amax_map_no_scales_raises(tmp_path): cast.build_amax_map(empty) -# ---------- magnitude table cache ------------------------------------------ - - -def test_e2m1_magnitude_table_cached_per_device(): - t1 = cast._e2m1_magnitude_table(torch.device("cpu")) - t2 = cast._e2m1_magnitude_table(torch.device("cpu")) - assert t1 is t2 # cached: same object - assert t1.tolist() == cast._E2M1_MAGNITUDE - - # ---------- apply_to_model end-to-end (mock model) --------------------------- diff --git a/tests/unit/torch/quantization/test_numeric_utils.py b/tests/unit/torch/quantization/test_numeric_utils.py new file mode 100644 index 00000000000..fb18a994a74 --- /dev/null +++ b/tests/unit/torch/quantization/test_numeric_utils.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for ``modelopt.torch.quantization.utils.numeric_utils`` — the +closed-form MXFP4 -> NVFP4 cast numerics.""" + +import pytest +import torch + +from modelopt.torch.quantization.utils import numeric_utils as nu + +# ---------- mxfp4_to_nvfp4_global_amax -------------------------------------- + + +def test_global_amax_basic_in_range(): + """Mixed in-range scales: m = k_max - 8, global_amax = 6*448*2^m, lossless = 100%.""" + # k values in [-3, 3] (spread = 6), all blocks lossless. + k = torch.tensor([0, -3, 3, 1, -1, 2], dtype=torch.int32) + e8m0 = (k + nu.E8M0_BIAS).to(torch.uint8) + + global_amax, info = nu.mxfp4_to_nvfp4_global_amax(e8m0) + assert info["k_min"] == -3 + assert info["k_max"] == 3 + assert info["m"] == 3 - 8 # k_max - 8 = -5 + expected = 6.0 * 448.0 * 2.0 ** info["m"] + assert global_amax == pytest.approx(expected) + assert info["n_total_blocks"] == 6 + assert info["n_lossless_blocks"] == 6 + assert info["pct_lossless"] == pytest.approx(100.0) + assert info["n_zero_blocks"] == 0 + + +def test_global_amax_with_zero_blocks(): + """Zero (e8m0=0, k=-127) blocks should be ignored when computing k_max.""" + e8m0 = torch.tensor([0, 0, 130, 125], dtype=torch.uint8) # ks: -127, -127, 3, -2 + global_amax, info = nu.mxfp4_to_nvfp4_global_amax(e8m0) + assert info["k_max"] == 3 # ignores zero blocks + assert info["n_zero_blocks"] == 2 + # Both nonzero blocks satisfy k_max - k_j <= 17, plus zero blocks count as + # lossless because their reconstruction is 0 regardless of scale. + assert info["n_lossless_blocks"] == 4 + + +def test_global_amax_with_oor_blocks(): + """A block 18 powers below k_max is OOR (k_max - k = 18 > 17).""" + # k values: 5, 5, -13 → spread = 18, last block is OOR. + k = torch.tensor([5, 5, -13], dtype=torch.int32) + e8m0 = (k + nu.E8M0_BIAS).to(torch.uint8) + _, info = nu.mxfp4_to_nvfp4_global_amax(e8m0) + assert info["k_max"] == 5 + assert info["n_total_blocks"] == 3 + assert info["n_lossless_blocks"] == 2 # the k=-13 block is OOR + + +def test_global_amax_all_zero(): + """All-zero scales should not crash; k_max defaults to 0.""" + e8m0 = torch.zeros(4, dtype=torch.uint8) + global_amax, info = nu.mxfp4_to_nvfp4_global_amax(e8m0) + assert info["k_min"] == 0 and info["k_max"] == 0 + assert info["n_zero_blocks"] == 4 + # All blocks count as "lossless" (their dequant is 0 regardless of scale). + assert info["n_lossless_blocks"] == 4 + + +# ---------- mxfp4_to_nvfp4_per_block_amax ----------------------------------- + + +def _make_blocks_with_max_nibble(num_blocks: int, max_idx_per_block: list[int]) -> torch.Tensor: + """Build a (num_blocks, 16) byte tensor where block i has E2M1 magnitude + index ``max_idx_per_block[i]`` as its largest nibble; other nibbles are 0. + + Magnitude index goes in the low 3 bits of one nibble; we place it in the + high nibble of byte 0 (so the first byte = (max_idx << 4)). Every other + nibble is 0, so the block-wise max is exactly ``max_idx_per_block[i]``. + """ + assert len(max_idx_per_block) == num_blocks + blocks = torch.zeros((num_blocks, 16), dtype=torch.uint8) + for i, idx in enumerate(max_idx_per_block): + assert 0 <= idx < 8 + blocks[i, 0] = (idx & 0x07) << 4 + return blocks + + +def test_per_block_amax_in_range_returns_closed_form(): + """Every block in-range -> 6 * 2^k_j, regardless of actual nibble content.""" + # k = [0, -2, 4]; k_max = 4, k_min = -2, spread 6 (in-range). + k = torch.tensor([0, -2, 4], dtype=torch.int32) + e8m0 = (k + nu.E8M0_BIAS).to(torch.uint8) + # Blocks have varying max_nibbles, but in-range path ignores them. + blocks = _make_blocks_with_max_nibble(3, [3, 7, 1]) # max nibbles: 1.5, 6, 0.5 + + out = nu.mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) + expected_mxfp4 = 6.0 * torch.exp2(k.float()) # ignores max_nibble + expected_nvfp4 = expected_mxfp4.repeat_interleave(2, dim=-1) + assert torch.allclose(out, expected_nvfp4) + + +def test_per_block_amax_oor_uses_data_derived(): + """OOR blocks should use ``max_nibble * 2^k_j`` (data-derived).""" + # k_max=10 → m=2. OOR-low blocks have k_j - m < -9, i.e. k_j < -7. + k = torch.tensor([10, -10], dtype=torch.int32) # second is OOR-low + e8m0 = (k + nu.E8M0_BIAS).to(torch.uint8) + # Block 0 max nibble idx 7 (value 6); block 1 max nibble idx 4 (value 2). + blocks = _make_blocks_with_max_nibble(2, [7, 4]) + + out = nu.mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) + + # Block 0 (in-range): 6 * 2^10 = 6144. + # Block 1 (OOR): 2 * 2^-10 (max_nibble=2 since idx=4 -> 2.0). + expected_mxfp4 = torch.tensor([6.0 * 2**10, 2.0 * 2**-10], dtype=torch.float32) + expected_nvfp4 = expected_mxfp4.repeat_interleave(2, dim=-1) + assert torch.allclose(out, expected_nvfp4) + + +def test_per_block_amax_doubles_last_dim(): + """Two NVFP4 blocks per MXFP4 block share the same per-block amax.""" + e8m0 = torch.tensor([130, 124], dtype=torch.uint8) # ks: 3, -3 + blocks = _make_blocks_with_max_nibble(2, [7, 7]) # in-range + out = nu.mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) + assert out.shape == (4,) + # Each pair of consecutive entries should be equal. + assert out[0] == out[1] + assert out[2] == out[3] + + +def test_per_block_amax_preserves_leading_dims(): + """Leading dims (E, F, ...) flow through unchanged; only last dim doubles.""" + # shape (E=2, F=3, num_mxfp4_blocks=4) + e8m0 = torch.full((2, 3, 4), 128, dtype=torch.uint8) # all k=1, in-range + blocks = torch.zeros((2, 3, 4, 16), dtype=torch.uint8) + out = nu.mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) + assert out.shape == (2, 3, 8) + + +def test_per_block_amax_shape_mismatch_raises(): + """Mismatched leading dims should raise ``ValueError``.""" + blocks = torch.zeros((4, 16), dtype=torch.uint8) + e8m0 = torch.zeros(3, dtype=torch.uint8) # different num_blocks + with pytest.raises(ValueError, match="shape mismatch"): + nu.mxfp4_to_nvfp4_per_block_amax(blocks, e8m0) + + +# ---------- magnitude table cache ------------------------------------------ + + +def test_e2m1_magnitude_table_cached_per_device(): + t1 = nu._e2m1_magnitude_table(torch.device("cpu")) + t2 = nu._e2m1_magnitude_table(torch.device("cpu")) + assert t1 is t2 # cached: same object + assert t1.tolist() == nu._E2M1_MAGNITUDE