From c47af4a7eab6173c8248dec0497d174cd3fe3250 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 18:04:08 -0700 Subject: [PATCH 1/4] Add NVFP4 four-over-six grouped GEMM for MoE training Grouped counterpart of four_over_six_mm for routed-expert layers (A (M, K) token groups x B (E, N, K) expert weights with group-end offsets), following TransformerEngine's GroupedLinear semantics: - per-tensor activations quantize with per-group global scales: the group amaxes expand to a per-row amax vector so one four_over_six_quantize call is bitwise identical to quantizing each split separately, and the forward is one F.scaled_grouped_mm with per-group second-level scales (bound-aware, not the 448-hardcoded helper). - row-scaled activations run a per-group loop of the dense four-over-six GEMM (FP32 output scaled by raw per-row amaxes) - the same shape TransformerEngine's general_grouped_gemm gives row-scaled NVFP4 at ea1a165d. - backward supports only the high_precision and dequantized overrides (TransformerEngine rejects four-over-six group quantization, so there is no grouped quantized backward): bf16 torch._grouped_mm on the saved originals or on dequantizations of the rowwise fprop operands. - ragged token groups zero-pad to 128-row alignment before quantization via pad_token_groups; padded rows quantize to zero codes and are sliced from the output. Validation (GB200, TE 2.19 devel container): 17 new tests pass, all bitwise (atol=rtol=0) - group-expanded amax == per-split quantize, grouped forward == dense per-group GEMMs (bitwise at the tested shapes for both weight blocks), row-scaled grouped == dense loop, both backward modes == manual grouped GEMMs on original/dequantized operands, ragged+padded == aligned construction, and the miles NVFP4 RL recipe point (row-scaled + MSE + bound 256 + 1x16 weights + dequantized backward). Full nvfp4_training suite: 461 passed / 79 skipped. Co-Authored-By: Claude Fable 5 --- .../test_four_over_six_grouped.py | 313 +++++++++++++ .../nvfp4_training/four_over_six_grouped.py | 440 ++++++++++++++++++ 2 files changed, 753 insertions(+) create mode 100644 test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py create mode 100644 torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py diff --git a/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py b/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py new file mode 100644 index 0000000000..8c1afdc9c4 --- /dev/null +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py @@ -0,0 +1,313 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + + +import pytest +import torch + +from torchao.float8.float8_utils import compute_error +from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + four_over_six_dequantize, + four_over_six_linear, + four_over_six_quantize, +) +from torchao.prototype.moe_training.nvfp4_training.four_over_six_grouped import ( + four_over_six_grouped_mm, +) +from torchao.utils import is_sm_at_least_100, torch_version_at_least + +_skip_no_sm100 = pytest.mark.skipif( + not ( + torch.cuda.is_available() + and is_sm_at_least_100() + and torch_version_at_least("2.10.0") + ), + reason="requires SM100+ and PyTorch 2.10+ (FP4 scaled_grouped_mm)", +) + + +def _make_grouped_inputs(group_sizes, K, N, seed=0, device="cuda"): + """Packed activations, stacked expert weights, and end offsets.""" + torch.manual_seed(seed) + M = sum(group_sizes) + E = len(group_sizes) + A = torch.randn(M, K, dtype=torch.bfloat16, device=device) + B = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.1 + offs = torch.tensor(group_sizes, dtype=torch.int32, device=device).cumsum( + 0, dtype=torch.int32 + ) + return A, B, offs + + +@_skip_no_sm100 +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +def test_group_expanded_amax_matches_per_split_quantize(err_mode, e4m3_scale_bound): + """One quantize call with group-expanded amaxes == a per-split loop.""" + group_sizes = [128, 384, 256] + A, _, offs = _make_grouped_inputs(group_sizes, K=256, N=128) + group_amax = torch.stack( + [ + A[start:end].abs().amax().to(torch.float32) + for start, end in zip([0, *offs.tolist()[:-1]], offs.tolist()) + ] + ) + expanded = group_amax.repeat_interleave( + torch.tensor(group_sizes, device=A.device) + ) + codes, scales = four_over_six_quantize( + A, expanded, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + start = 0 + for g, end in enumerate(offs.tolist()): + split_codes, split_scales = four_over_six_quantize( + A[start:end].contiguous(), + group_amax[g], + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + torch.testing.assert_close(codes[start:end], split_codes, atol=0, rtol=0) + torch.testing.assert_close( + scales[start:end].view(torch.uint8), + split_scales.view(torch.uint8), + atol=0, + rtol=0, + ) + start = end + + +@_skip_no_sm100 +@pytest.mark.parametrize("weight_block", ["16x16", "1x16"]) +def test_per_tensor_grouped_forward_matches_dense_loop(weight_block): + """Grouped forward vs dense four_over_six GEMMs per 128-aligned group. + + The quantized operands are bitwise-identical by construction (pinned by + the amax-expansion test above); the GEMM outputs are compared bitwise + and fall back to an SQNR bound if the grouped and dense kernels reduce + in different orders. + """ + group_sizes = [128, 256, 128] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N) + y = four_over_six_grouped_mm(A, B, offs, weight_block=weight_block) + + start = 0 + refs = [] + for e, end in enumerate(offs.tolist()): + refs.append( + four_over_six_linear( + A[start:end].contiguous(), + B[e], + None, + "mae", + 256, + False, + "high_precision", + weight_block, + ) + ) + start = end + y_ref = torch.cat(refs) + if not torch.equal(y, y_ref): + sqnr = compute_error(y_ref.float(), y.float()) + assert sqnr > 85.0, f"grouped vs dense-loop forward SQNR {sqnr:.1f} dB" + print(f"\ngrouped GEMM reduction differs from dense: SQNR {sqnr:.1f} dB") + + +@_skip_no_sm100 +def test_row_scaled_grouped_forward_matches_dense_loop(): + """Row-scaled grouped forward is the per-group dense loop by construction.""" + group_sizes = [128, 256, 128] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=384) + y = four_over_six_grouped_mm(A, B, offs, row_scaled_activation=True) + + start = 0 + refs = [] + for e, end in enumerate(offs.tolist()): + refs.append( + four_over_six_linear( + A[start:end].contiguous(), B[e], None, "mae", 256, True + ) + ) + start = end + torch.testing.assert_close(y, torch.cat(refs), atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +def test_grouped_backward_high_precision(row_scaled_activation): + """dx/dw are bf16 grouped GEMMs on the original operands.""" + group_sizes = [128, 256, 128] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=384) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, B, offs, row_scaled_activation=row_scaled_activation + ) + dy = torch.randn_like(y) + y.backward(dy) + + dx_ref = torch._grouped_mm( + dy, B.detach(), offs=offs, out_dtype=torch.bfloat16 + ) + dw_ref = torch._grouped_mm( + dy.transpose(-2, -1), A.detach(), offs=offs, out_dtype=torch.bfloat16 + ) + torch.testing.assert_close(A.grad, dx_ref, atol=0, rtol=0) + torch.testing.assert_close(B.grad, dw_ref, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +@pytest.mark.parametrize("weight_block", ["16x16", "1x16"]) +def test_grouped_backward_dequantized(row_scaled_activation, weight_block): + """dx/dw are bf16 grouped GEMMs on dequantized fprop operands.""" + group_sizes = [128, 256, 128] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, + B, + offs, + err_mode="mse", + row_scaled_activation=row_scaled_activation, + weight_block=weight_block, + backward_override="dequantized", + ) + dy = torch.randn_like(y) + y.backward(dy) + + A_hp, B_hp = A.detach(), B.detach() + if row_scaled_activation: + x_amax = A_hp.abs().amax(dim=1).to(torch.float32) + else: + group_amax = [] + start = 0 + for end in offs.tolist(): + group_amax.append(A_hp[start:end].abs().amax().to(torch.float32)) + start = end + x_amax = torch.stack(group_amax).repeat_interleave( + torch.tensor(group_sizes, device=A.device) + ) + x_codes, x_scales = four_over_six_quantize(A_hp, x_amax, err_mode="mse") + x_dq = four_over_six_dequantize(x_codes, x_scales, x_amax) + w_dq = [] + for e in range(B.shape[0]): + w_amax = B_hp[e].abs().amax().to(torch.float32) + w_codes, w_scales = four_over_six_quantize( + B_hp[e], w_amax, block=weight_block, err_mode="mse" + ) + w_dq.append(four_over_six_dequantize(w_codes, w_scales, w_amax)) + w_dq = torch.stack(w_dq) + + dx_ref = torch._grouped_mm(dy, w_dq, offs=offs, out_dtype=torch.bfloat16) + dw_ref = torch._grouped_mm( + dy.transpose(-2, -1), x_dq, offs=offs, out_dtype=torch.bfloat16 + ) + torch.testing.assert_close(A.grad, dx_ref, atol=0, rtol=0) + torch.testing.assert_close(B.grad, dw_ref, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +def test_grouped_padding_matches_aligned(row_scaled_activation): + """Unaligned groups with padding == an aligned construction, per group.""" + K, N = 256, 384 + aligned_sizes = [128, 256, 128] + ragged_sizes = [100, 220, 77] + A_al, B, offs_al = _make_grouped_inputs(aligned_sizes, K=K, N=N, seed=3) + # Ragged view: the first rows of each aligned group, so every ragged + # group's rows (and hence its amax and quantization) exist verbatim in + # the aligned run. + ragged_rows = [] + start = 0 + for size, ragged in zip(aligned_sizes, ragged_sizes): + ragged_rows.append(A_al[start : start + ragged]) + start += size + A_rg = torch.cat(ragged_rows).contiguous() + offs_rg = torch.tensor( + ragged_sizes, dtype=torch.int32, device=A_al.device + ).cumsum(0, dtype=torch.int32) + + y_rg = four_over_six_grouped_mm( + A_rg, + B, + offs_rg, + row_scaled_activation=row_scaled_activation, + pad_token_groups_for_grouped_mm=True, + ) + assert y_rg.shape == (sum(ragged_sizes), N) + + # Reference: dense per-group forward on the ragged rows padded to 128. + start = 0 + for e, ragged in enumerate(ragged_sizes): + rows = A_rg[start : start + ragged] + padded = torch.zeros(128 * ((ragged + 127) // 128), K, dtype=rows.dtype, device=rows.device) + padded[:ragged] = rows + if row_scaled_activation: + ref = four_over_six_linear(padded, B[e], None, "mae", 256, True) + else: + # Per-tensor group scale comes from the real rows' amax; the + # zero padding rows cannot change it. + ref = four_over_six_linear( + padded, B[e], None, "mae", 256, False, "high_precision" + ) + torch.testing.assert_close( + y_rg[start : start + ragged], ref[:ragged], atol=0, rtol=0 + ) + start += ragged + + +@_skip_no_sm100 +def test_grouped_validation(): + group_sizes = [128, 128] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=128) + with pytest.raises(ValueError, match="no quantized backward"): + four_over_six_grouped_mm(A, B, offs, backward_override="quantized") + with pytest.raises(ValueError, match="1D int32"): + four_over_six_grouped_mm(A, B, offs.to(torch.int64)) + with pytest.raises(ValueError, match="one group-end offset per expert"): + four_over_six_grouped_mm(A, B, offs[:1]) + with pytest.raises(ValueError, match="must be 2D"): + four_over_six_grouped_mm(A.unsqueeze(0), B, offs) + with pytest.raises(ValueError, match="divisible by 128"): + four_over_six_grouped_mm(A[:, :144], B[:, :, :144].contiguous(), offs) + with pytest.raises(ValueError, match="weight_block"): + four_over_six_grouped_mm(A, B, offs, weight_block="8x8") + + +@_skip_no_sm100 +def test_grouped_miles_recipe_point(): + """The miles NVFP4 RL recipe: row-scaled + MSE + bound 256 + 1x16 weights + + dequantized backward, on ragged token groups.""" + group_sizes = [100, 220, 77] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=384, seed=7) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, + B, + offs, + err_mode="mse", + e4m3_scale_bound=256, + row_scaled_activation=True, + weight_block="1x16", + backward_override="dequantized", + pad_token_groups_for_grouped_mm=True, + ) + assert y.shape == (sum(group_sizes), 384) + y.backward(torch.randn_like(y)) + assert A.grad is not None and A.grad.shape == A.shape + assert B.grad is not None and B.grad.shape == B.shape + sqnr = compute_error( + torch._grouped_mm( + A.detach(), B.detach().transpose(-2, -1), offs=offs + ).float(), + y.float(), + ) + assert sqnr > 14.0, f"quantization noise floor too high: {sqnr:.1f} dB" diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py new file mode 100644 index 0000000000..f4f5ed9197 --- /dev/null +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py @@ -0,0 +1,440 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Differentiable NVFP4 four-over-six grouped GEMM for MoE training. + +Grouped counterpart of ``four_over_six_mm`` for routed-expert layers: +``A`` holds token groups packed along dim 0, ``B`` holds one weight matrix +per expert, and ``offs`` marks each group's end row. The recipe follows +TransformerEngine's GroupedLinear semantics, where every group is quantized +as its own tensor: + +* per-tensor activations (the default): each token group gets its own + global scale from that group's amax. The group amaxes are expanded to a + per-row amax vector so the whole packed tensor quantizes in one + ``four_over_six_quantize`` call — bitwise identical to quantizing each + group separately, because the quantizer derives every row's scale chain + from that row's amax entry. The forward GEMM is one + ``F.scaled_grouped_mm`` with per-group second-level scales. +* row-scaled activations: one global scale per token row. TransformerEngine + has no fused row-scaled NVFP4 grouped GEMM — its ``general_grouped_gemm`` + runs a per-group loop of dense GEMMs — so the forward here is the same + loop over the dense four-over-six GEMM (FP32 output scaled by the raw + per-row amaxes, then the bf16 cast). + +Weights always quantize per expert with per-tensor scales +(``weight_block`` selects 16x16 tiles or 1x16 blocks, as in the dense op). + +Gradients never quantize with four-over-six, and TransformerEngine rejects +four-over-six group quantization outright, so the grouped backward supports +only the high-precision and dequantized overrides of ``four_over_six_mm``: + +* ``"high_precision"`` (the default): bf16 grouped GEMMs on the saved + original operands; +* ``"dequantized"``: bf16 grouped GEMMs on dequantizations of the rowwise + operands the forward consumed — the RL train/inference-consistency mode. + +``"quantized"`` raises. Requires K % 128 == 0 and N % 128 == 0; token +groups must be 128-row aligned unless ``pad_token_groups_for_grouped_mm`` +is set, which zero-pads each group to the next 128 multiple before +quantization (zero rows quantize to zero codes and are sliced away from the +output). +""" + +from typing import Optional + +import torch +import torch.nn.functional as F + +from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + FP4_E2M1_MAX, + _global_decode_scale, + _scaled_mm_nvfp4, + four_over_six_dequantize, + four_over_six_quantize, +) +from torchao.prototype.moe_training.nvfp4_training.group_hadamard_utils import ( + _DEVICE_ASSERTS, +) +from torchao.prototype.moe_training.utils import ( + conditional_nostrict_trace, + pad_token_groups, + unpad_token_groups, +) +from torchao.prototype.mx_formats.utils import to_blocked +from torchao.quantization.quantize_.common import KernelPreference +from torchao.utils import is_sm_at_least_100 + +_ALIGNMENT = 128 +_SCALE_RECIPE = [F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise] +_SWIZZLE = [F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE] + +__all__ = ["four_over_six_grouped_mm"] + + +@conditional_nostrict_trace +def four_over_six_grouped_mm( + A: torch.Tensor, + B: torch.Tensor, + offs: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + err_mode: str = "mae", + e4m3_scale_bound: int = 256, + row_scaled_activation: bool = False, + weight_block: str = "16x16", + backward_override: Optional[str] = None, + pad_token_groups_for_grouped_mm: bool = False, +) -> torch.Tensor: + """Quantize grouped activations and expert weights with four-over-six. + + ``A`` has shape ``(M, K)``, ``B`` has shape ``(E, N, K)``, and ``offs`` + contains the cumulative row-end offset for each expert. Knobs match + ``four_over_six_mm``; see the module docstring for the grouped-specific + backward and alignment semantics. + """ + output = _FourOverSixGroupedMM.apply( + A, + B, + offs, + err_mode, + e4m3_scale_bound, + row_scaled_activation, + weight_block, + backward_override, + pad_token_groups_for_grouped_mm, + ) + if bias is not None: + output = output + bias.to(output.dtype) + return output + + +def _expand_group_amax( + row_amax: torch.Tensor, group_end_offsets: torch.Tensor, num_experts: int +) -> torch.Tensor: + """Per-row amax vector holding each row's group amax. + + Rows past the final offset (the pad-helper's over-allocated tail) take + the last group's amax; they are all-zero and never enter the GEMM. + """ + group_idx = torch.searchsorted( + group_end_offsets, + torch.arange(row_amax.shape[0], device=row_amax.device, dtype=torch.int32), + right=True, + ).clamp_(max=num_experts - 1) + group_amax = torch.zeros( + num_experts, dtype=torch.float32, device=row_amax.device + ).scatter_reduce_( + 0, group_idx, row_amax.to(torch.float32), reduce="amax", include_self=True + ) + return group_amax[group_idx], group_amax + + +def _quantize_expert_weights( + weight: torch.Tensor, + weight_amax: torch.Tensor, + weight_block: str, + err_mode: str, + e4m3_scale_bound: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-expert four-over-six quantization of a stacked (E, N, K) weight.""" + codes = [] + scales = [] + for e in range(weight.shape[0]): + expert_codes, expert_scales = four_over_six_quantize( + weight[e], + weight_amax[e], + block=weight_block, + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + codes.append(expert_codes) + scales.append(expert_scales) + return torch.stack(codes), torch.stack(scales) + + +def _dequantize_expert_weights( + codes: torch.Tensor, + scales: torch.Tensor, + weight_amax: torch.Tensor, + e4m3_scale_bound: int, +) -> torch.Tensor: + """Dequantize stacked per-expert codes back to a bf16 (E, N, K) weight. + + Flattening experts along rows and expanding each expert's amax over its + rows reproduces the per-expert scalar dequantization exactly — the + decode chain reads one amax entry per row either way. + """ + num_experts, N = codes.shape[0], codes.shape[1] + row_amax = weight_amax.to(torch.float32).repeat_interleave(N) + flat = four_over_six_dequantize( + codes.reshape(num_experts * N, -1), + scales.reshape(num_experts * N, -1), + row_amax, + e4m3_scale_bound=e4m3_scale_bound, + ) + return flat.view(num_experts, N, -1) + + +class _FourOverSixGroupedMM(torch.autograd.Function): + """NVFP4 four-over-six grouped forward with override-only backward.""" + + @staticmethod + def forward( + ctx, + input_act: torch.Tensor, + weight: torch.Tensor, + group_end_offsets: torch.Tensor, + err_mode: str, + e4m3_scale_bound: int, + row_scaled_activation: bool, + weight_block: str, + backward_override: Optional[str], + pad_token_groups_for_grouped_mm: bool, + ) -> torch.Tensor: + if input_act.ndim != 2: + raise ValueError(f"input_act must be 2D, got {input_act.ndim}D") + if weight.ndim != 3: + raise ValueError(f"weight must be 3D, got {weight.ndim}D") + if group_end_offsets is None: + raise ValueError("offs is required for NVFP4 grouped GEMM") + if group_end_offsets.ndim != 1 or group_end_offsets.dtype != torch.int32: + raise ValueError("offs must be a 1D int32 tensor") + if not group_end_offsets.is_contiguous(): + raise ValueError("offs must be contiguous") + if group_end_offsets.numel() != weight.shape[0]: + raise ValueError("offs must contain one group-end offset per expert") + if not (input_act.is_cuda and weight.is_cuda and group_end_offsets.is_cuda): + raise ValueError("input_act, weight, and offs must be CUDA tensors") + if not (input_act.device == weight.device == group_end_offsets.device): + raise ValueError("all tensor arguments must be on the same device") + if not is_sm_at_least_100(): + raise NotImplementedError( + "NVFP4 four-over-six grouped GEMM requires SM100+" + ) + if backward_override is None: + backward_override = "high_precision" + if backward_override not in ("high_precision", "dequantized"): + if backward_override == "quantized": + raise ValueError( + "grouped four-over-six has no quantized backward; use " + "'high_precision' or 'dequantized'" + ) + raise ValueError( + f"backward_override must be 'high_precision' or 'dequantized', " + f"got {backward_override!r}" + ) + if weight_block not in ("1x16", "16x16"): + raise ValueError( + f"weight_block must be '1x16' or '16x16', got {weight_block!r}" + ) + + num_tokens, K = input_act.shape + num_experts, N, weight_K = weight.shape + if weight_K != K: + raise ValueError( + f"input and weight contraction dimensions differ: {K} and {weight_K}" + ) + if K % _ALIGNMENT != 0 or N % _ALIGNMENT != 0: + raise ValueError( + f"K and N must be divisible by {_ALIGNMENT}; got K={K}, N={N}" + ) + if _DEVICE_ASSERTS: + group_sizes = torch.diff( + group_end_offsets, prepend=group_end_offsets.new_zeros(1) + ) + torch.ops.aten._assert_async.msg( + torch.all(group_sizes > 0), "offs must describe non-empty groups" + ) + torch.ops.aten._assert_async.msg( + group_end_offsets[-1] == num_tokens, + "the final group-end offset must equal A.shape[0]", + ) + if not pad_token_groups_for_grouped_mm: + torch.ops.aten._assert_async.msg( + torch.all(group_sizes % _ALIGNMENT == 0), + "every token group must be 128-row aligned when padding is disabled", + ) + + input_act = input_act.to(torch.bfloat16).contiguous() + weight = weight.to(torch.bfloat16).contiguous() + original_input = input_act + + padded_group_start_offsets = None + if pad_token_groups_for_grouped_mm: + input_act, padded_group_start_offsets, padded_group_end_offsets = ( + pad_token_groups( + input_act, + group_end_offsets, + alignment_size=_ALIGNMENT, + kernel_preference=KernelPreference.TRITON, + ) + ) + else: + padded_group_end_offsets = group_end_offsets + + row_amax = input_act.abs().amax(dim=1) + group_amax = None + if row_scaled_activation: + x_amax = row_amax.to(torch.float32) + else: + x_amax, group_amax = _expand_group_amax( + row_amax, padded_group_end_offsets, num_experts + ) + weight_amax = weight.abs().amax(dim=(1, 2)).to(torch.float32) + + x_codes, x_scales = four_over_six_quantize( + input_act, + x_amax, + block="1x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_codes, w_scales = _quantize_expert_weights( + weight, weight_amax, weight_block, err_mode, e4m3_scale_bound + ) + w_global = _global_decode_scale(weight_amax, e4m3_scale_bound) + + if row_scaled_activation: + # TransformerEngine has no fused row-scaled NVFP4 grouped GEMM; + # its general_grouped_gemm loops dense GEMMs per group, and so + # does this: FP32 output with the constant 1/(6*bound) factor in + # the per-tensor slot, scaled by the raw per-row amaxes. + x_global = torch.full( + (), + 1.0 / (FP4_E2M1_MAX * float(e4m3_scale_bound)), + dtype=torch.float32, + device=input_act.device, + ) + group_bounds = torch.stack( + ( + padded_group_end_offsets + - torch.diff( + padded_group_end_offsets, + prepend=padded_group_end_offsets.new_zeros(1), + ), + padded_group_end_offsets, + ), + dim=1, + ).tolist() + output = input_act.new_zeros(input_act.shape[0], N) + for e, (start, end) in enumerate(group_bounds): + if start == end: + continue + group_out = _scaled_mm_nvfp4( + x_codes[start:end], + x_scales[start:end], + x_global, + w_codes[e].t(), + w_scales[e], + w_global[e], + torch.float32, + ) + output[start:end] = ( + group_out * x_amax[start:end].view(-1, 1) + ).to(torch.bfloat16) + else: + output = F.scaled_grouped_mm( + x_codes.view(torch.float4_e2m1fn_x2), + w_codes.view(torch.float4_e2m1fn_x2).transpose(-2, -1), + # scaled_grouped_mm consumes swizzled scale bytes viewed at the + # logical 2D shape (the layout the group quantize kernels + # return); the view needs the 128-row alignment enforced above. + scale_a=[ + to_blocked(x_scales).view(x_scales.shape), + _global_decode_scale(group_amax, e4m3_scale_bound), + ], + scale_recipe_a=_SCALE_RECIPE, + scale_b=[ + torch.stack([to_blocked(s) for s in w_scales]).reshape( + num_experts, -1 + ), + w_global, + ], + scale_recipe_b=_SCALE_RECIPE, + swizzle_a=_SWIZZLE, + swizzle_b=_SWIZZLE, + offs=padded_group_end_offsets, + output_dtype=torch.bfloat16, + ) + + if pad_token_groups_for_grouped_mm: + output = unpad_token_groups( + output, + group_end_offsets, + padded_group_start_offsets, + num_tokens, + alignment_size=_ALIGNMENT, + kernel_preference=KernelPreference.TRITON, + ) + + if backward_override == "high_precision": + ctx.save_for_backward(original_input, weight, group_end_offsets) + else: + if padded_group_start_offsets is None: + padded_group_start_offsets = group_end_offsets.new_zeros(0) + ctx.save_for_backward( + x_codes, + x_scales, + x_amax, + w_codes, + w_scales, + weight_amax, + group_end_offsets, + padded_group_start_offsets, + ) + ctx.backward_override = backward_override + ctx.e4m3_scale_bound = e4m3_scale_bound + ctx.pad_token_groups_for_grouped_mm = pad_token_groups_for_grouped_mm + ctx.num_tokens = num_tokens + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_output = grad_output.to(torch.bfloat16).contiguous() + + if ctx.backward_override == "high_precision": + input_act, weight, group_end_offsets = ctx.saved_tensors + else: + ( + x_codes, + x_scales, + x_amax, + w_codes, + w_scales, + weight_amax, + group_end_offsets, + padded_group_start_offsets, + ) = ctx.saved_tensors + input_act = four_over_six_dequantize( + x_codes, x_scales, x_amax, e4m3_scale_bound=ctx.e4m3_scale_bound + ) + if ctx.pad_token_groups_for_grouped_mm: + input_act = unpad_token_groups( + input_act, + group_end_offsets, + padded_group_start_offsets, + ctx.num_tokens, + alignment_size=_ALIGNMENT, + kernel_preference=KernelPreference.TRITON, + ) + weight = _dequantize_expert_weights( + w_codes, w_scales, weight_amax, ctx.e4m3_scale_bound + ) + + grad_input = torch._grouped_mm( + grad_output, + weight, + offs=group_end_offsets, + out_dtype=torch.bfloat16, + ) + grad_weight = torch._grouped_mm( + grad_output.transpose(-2, -1), + input_act, + offs=group_end_offsets, + out_dtype=torch.bfloat16, + ) + return grad_input, grad_weight, None, None, None, None, None, None, None From 2c3ab70862de6cd736072c42044e1a602c469c3d Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 20:53:29 -0700 Subject: [PATCH 2/4] Pin the pure-torch pad path for four-over-six grouped alignment The fused pad/unpad CUDA kernels reject alignment_size != 32 and more than 32 groups (STD_TORCH_CHECK in mxfp8_extension.cpp), so any torchao build that ships them fails the 128-row-aligned ragged-group path on the first padded call. KernelPreference.EMULATED selects the torch implementation, which handles 128-row alignment and any expert count. --- .../moe_training/nvfp4_training/four_over_six_grouped.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py index f4f5ed9197..54f0f8f470 100644 --- a/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py @@ -265,12 +265,15 @@ def forward( padded_group_start_offsets = None if pad_token_groups_for_grouped_mm: + # The fused pad/unpad CUDA kernels only accept alignment_size 32 + # and at most 32 groups; this op needs 128-row alignment with any + # expert count, so it pins the pure-torch path. input_act, padded_group_start_offsets, padded_group_end_offsets = ( pad_token_groups( input_act, group_end_offsets, alignment_size=_ALIGNMENT, - kernel_preference=KernelPreference.TRITON, + kernel_preference=KernelPreference.EMULATED, ) ) else: @@ -368,7 +371,7 @@ def forward( padded_group_start_offsets, num_tokens, alignment_size=_ALIGNMENT, - kernel_preference=KernelPreference.TRITON, + kernel_preference=KernelPreference.EMULATED, ) if backward_override == "high_precision": @@ -419,7 +422,7 @@ def backward(ctx, grad_output: torch.Tensor): padded_group_start_offsets, ctx.num_tokens, alignment_size=_ALIGNMENT, - kernel_preference=KernelPreference.TRITON, + kernel_preference=KernelPreference.EMULATED, ) weight = _dequantize_expert_weights( w_codes, w_scales, weight_amax, ctx.e4m3_scale_bound From b67687f2a51f9d50f2372f8b8c24a3abe40fcc02 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 21:37:14 -0700 Subject: [PATCH 3/4] Register four_over_six_dequantize as a custom op Inductor codegen of the fused unpack + broadcast-scale dequantize graph miscompiles the low-nibble lane (torch 2.14 nightly; the helpers compile correctly in isolation, the fused whole-function graph does not; fresh contiguous inputs reproduce it). The dequantized backward's contract is bitwise parity with the fprop operands, so the decode must keep eager numerics under compile: registering it as an opaque custom op with a fake impl does that, matching the CuTe DSL quantize op's pattern. --- .../nvfp4_training/four_over_six.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py index a0ea577766..a69fd99360 100644 --- a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py @@ -314,6 +314,26 @@ def four_over_six_dequantize( f"global_amax must be a scalar or a ({rows},) row vector, " f"got shape {tuple(global_amax.shape)}" ) + return _four_over_six_dequantize_op( + codes, scales, global_amax, e4m3_scale_bound, out_dtype + ) + + +# Registered as a custom op so torch.compile keeps the eager decode: inductor +# codegen of the fused unpack + broadcast-scale graph miscompiles the +# low-nibble lane (torch 2.14 nightly), and the dequantized backward's whole +# contract is bitwise parity with the fprop operands. +@torch.library.custom_op("torchao::four_over_six_dequantize", mutates_args=()) +def _four_over_six_dequantize_op( + codes: torch.Tensor, + scales: torch.Tensor, + global_amax: torch.Tensor, + e4m3_scale_bound: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + rows, packed_cols = codes.shape + cols = packed_cols * 2 + row_scaled = global_amax.dim() == 1 and global_amax.numel() == rows values = f4_unpacked_to_f32(unpack_uint4(codes)).view(rows, cols // 16, 16) amax = global_amax.to(torch.float32) if row_scaled: @@ -330,6 +350,11 @@ def four_over_six_dequantize( return (values * decode_scale.unsqueeze(-1)).to(out_dtype).view(rows, cols) +@_four_over_six_dequantize_op.register_fake +def _(codes, scales, global_amax, e4m3_scale_bound, out_dtype): + return codes.new_empty((codes.shape[0], codes.shape[1] * 2), dtype=out_dtype) + + def _standard_rtne_quantize( x: torch.Tensor, global_amax: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: From b701e80e545e0b5c9935b5e1f598fed441fec253 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 21:37:14 -0700 Subject: [PATCH 4/4] Test ragged dequantized backward and compile coverage - test_grouped_backward_dequantized_ragged: bitwise dx/dw for unaligned groups + padding + dequantized backward (the torchtitan grouped-experts hook's production composition, previously only shape-checked). - test_linear_compile_backward_overrides / test_linear_compile_weight_block_1x16: fullgraph compile of the override backwards and 1x16 weights, bitwise vs eager (caught the inductor dequantize miscompile the previous commit pins). - test_grouped_compile: fullgraph compile of the per-tensor grouped op, forward and backward, following the mxfp8 grouped test's direct-compile pattern; skips on torch builds whose nonstrict_trace rejects autograd.Function outputs. Row-scaled grouped is eager-only by design. --- .../nvfp4_training/test_four_over_six.py | 53 +++++++++ .../test_four_over_six_grouped.py | 106 ++++++++++++++++++ 2 files changed, 159 insertions(+) diff --git a/test/prototype/moe_training/nvfp4_training/test_four_over_six.py b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py index 1e37a8adb3..851571df17 100644 --- a/test/prototype/moe_training/nvfp4_training/test_four_over_six.py +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py @@ -613,3 +613,56 @@ def fn(x, w): y_eager = fn(x, w) y_compiled = torch.compile(fn, fullgraph=True)(x, w) torch.testing.assert_close(y_compiled, y_eager, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.skipif(not _cutedsl_available, reason="requires the CuTe DSL runtime") +@pytest.mark.parametrize( + "backward_override", ["high_precision", "dequantized"] +) +def test_linear_compile_backward_overrides(backward_override): + """fullgraph compile of the override backwards, bitwise vs eager. + + The quantize stays an opaque custom op under compile; the override + backwards are bf16 GEMMs (on original or dequantized operands), so + compiled gradients must match eager exactly. + """ + torch.manual_seed(0) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(384, 256, dtype=torch.bfloat16, device="cuda") * 0.1 + + def fn(x, w): + return four_over_six_linear( + x, w, None, "mae", 256, False, backward_override + ) + + x_e = x.clone().requires_grad_(True) + w_e = w.clone().requires_grad_(True) + y_eager = fn(x_e, w_e) + dy = torch.randn_like(y_eager) + y_eager.backward(dy) + + x_c = x.clone().requires_grad_(True) + w_c = w.clone().requires_grad_(True) + y_compiled = torch.compile(fn, fullgraph=True)(x_c, w_c) + y_compiled.backward(dy) + + torch.testing.assert_close(y_compiled, y_eager, atol=0, rtol=0) + torch.testing.assert_close(x_c.grad, x_e.grad, atol=0, rtol=0) + torch.testing.assert_close(w_c.grad, w_e.grad, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.skipif(not _cutedsl_available, reason="requires the CuTe DSL runtime") +def test_linear_compile_weight_block_1x16(): + """fullgraph compile of the 1x16-weight forward, bitwise vs eager.""" + torch.manual_seed(0) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(384, 256, dtype=torch.bfloat16, device="cuda") * 0.1 + + def fn(x, w): + return four_over_six_linear(x, w, None, "mae", 256, False, None, "1x16") + + y_eager = fn(x, w) + y_compiled = torch.compile(fn, fullgraph=True)(x, w) + torch.testing.assert_close(y_compiled, y_eager, atol=0, rtol=0) diff --git a/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py b/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py index 8c1afdc9c4..89047578d1 100644 --- a/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py @@ -311,3 +311,109 @@ def test_grouped_miles_recipe_point(): y.float(), ) assert sqnr > 14.0, f"quantization noise floor too high: {sqnr:.1f} dB" + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +@pytest.mark.parametrize("weight_block", ["16x16", "1x16"]) +def test_grouped_backward_dequantized_ragged(row_scaled_activation, weight_block): + """Ragged groups + padding + dequantized backward, value-checked. + + This is the composition the torchtitan grouped-experts hook ships; + the padded rows quantize to zeros and are unpadded away before the + backward GEMMs, so the reference can quantize the ragged rows directly. + """ + group_sizes = [100, 220, 77] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N, seed=7) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, + B, + offs, + err_mode="mse", + row_scaled_activation=row_scaled_activation, + weight_block=weight_block, + backward_override="dequantized", + pad_token_groups_for_grouped_mm=True, + ) + dy = torch.randn_like(y) + y.backward(dy) + + A_hp, B_hp = A.detach(), B.detach() + if row_scaled_activation: + x_amax = A_hp.abs().amax(dim=1).to(torch.float32) + else: + # Group amaxes come from the real rows; zero padding cannot raise them. + group_amax = [] + start = 0 + for end in offs.tolist(): + group_amax.append(A_hp[start:end].abs().amax().to(torch.float32)) + start = end + x_amax = torch.stack(group_amax).repeat_interleave( + torch.tensor(group_sizes, device=A.device) + ) + x_codes, x_scales = four_over_six_quantize(A_hp, x_amax, err_mode="mse") + x_dq = four_over_six_dequantize(x_codes, x_scales, x_amax) + w_dq = [] + for e in range(B.shape[0]): + w_amax = B_hp[e].abs().amax().to(torch.float32) + w_codes, w_scales = four_over_six_quantize( + B_hp[e], w_amax, block=weight_block, err_mode="mse" + ) + w_dq.append(four_over_six_dequantize(w_codes, w_scales, w_amax)) + w_dq = torch.stack(w_dq) + + dx_ref = torch._grouped_mm(dy, w_dq, offs=offs, out_dtype=torch.bfloat16) + dw_ref = torch._grouped_mm( + dy.transpose(-2, -1), x_dq, offs=offs, out_dtype=torch.bfloat16 + ) + torch.testing.assert_close(A.grad, dx_ref, atol=0, rtol=0) + torch.testing.assert_close(B.grad, dw_ref, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize( + "backward_override", ["high_precision", "dequantized"] +) +def test_grouped_compile(backward_override): + """fullgraph compile of the per-tensor grouped op, forward and backward. + + The op is nonstrict-traced under compile, so eager numerics carry over + bitwise. Row-scaled mode is exempt: its per-group dense-GEMM loop reads + the offsets on the host, which fullgraph capture cannot express (the + titan converter rejects that combination at config time). + """ + group_sizes = [128, 256, 128] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=384, seed=11) + + # Compile the decorated op directly (the mxfp8 grouped test's pattern); + # calling a nonstrict-traced function from a compiled frame is rejected. + A_e = A.clone().requires_grad_(True) + B_e = B.clone().requires_grad_(True) + y_eager = four_over_six_grouped_mm( + A_e, B_e, offs, err_mode="mse", backward_override=backward_override + ) + dy = torch.randn_like(y_eager) + y_eager.backward(dy) + + A_c = A.clone().requires_grad_(True) + B_c = B.clone().requires_grad_(True) + try: + y_compiled = torch.compile(four_over_six_grouped_mm, fullgraph=True)( + A_c, B_c, offs, err_mode="mse", backward_override=backward_override + ) + except torch._dynamo.exc.Unsupported as e: + if "nonstrict_trace" in str(e): + pytest.skip( + "this torch build rejects autograd.Function outputs from " + "nonstrict_trace-ed functions (the mxfp8 grouped compile " + "test's pattern); coverage resumes on builds that accept it" + ) + raise + y_compiled.backward(dy) + + torch.testing.assert_close(y_compiled, y_eager, atol=0, rtol=0) + torch.testing.assert_close(A_c.grad, A_e.grad, atol=0, rtol=0) + torch.testing.assert_close(B_c.grad, B_e.grad, atol=0, rtol=0)