Add _grouped_mm dispatch handler to Float8Tensor

[xiaowangintel]

Summary

This PR adds an aten._grouped_mm.default dispatch handler to Float8Tensor, enabling FP8 inference for MoE (Mixture of Experts) models that use torch._grouped_mm for expert computation. This follows the same pattern established by NVFP4 in #4316.

Motivation

MoE architectures (e.g., DeepSeek, Qwen3-MoE) store expert weights as 3D tensors (E, N, K) and use torch._grouped_mm in the forward pass. With this PR, users can quantize expert weights to FP8 via quantize_() using existing configs (Float8WeightOnlyConfig, Float8DynamicActivationFloat8WeightConfig) and get automatic dispatch through the _grouped_mm handler — no model code changes needed.

Related: RFC #4355, NVFP4 reference #4316

Design

The handler supports two modes based on weight granularity:

RowWise-scaled weights (preferred path):

1. Quantize activation to FP8 — via act_quant_kwargs (dynamic act+weight) or Float8Tensor.from_hp() with PerRow (weight-only)
2. Try F.scaled_grouped_mm (public API, backed by aten::_scaled_grouped_mm_v2)
3. On failure (e.g., unsupported device), fall back to dequantize both → _grouped_mm

Non-RowWise weights (PerTensor, PerGroup, etc.):

• Dequantize weight only → _grouped_mm with bf16

Test plan

3 test cases added to TestFloat8Tensor in test_float8_tensor.py:

Test	Config	Granularity	Path exercised
test_fp8_grouped_mm_dynamic_act_weight	Float8DynamicActivationFloat8WeightConfig	PerRow	scaled_grouped_mm
test_fp8_grouped_mm_weight_only	Float8WeightOnlyConfig	PerRow	scaled_grouped_mm
test_fp8_grouped_mm_dequantize_roundtrip	Float8WeightOnlyConfig	PerTensor	Direct dequant -> grouped_mm

Parameterized over (E, K, N) shapes. SQNR thresholds: weight ≥ 25, output ≥ 20.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4390

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• /easycla is not responding

✅ No Failures

As of commit 6cb1fd2 with merge base 8bbe1a5 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vkuzo]

the code here should honor what the user specified in the config (either weight-only or dynamic quant), see _float8_addmm_impl in this file for an example

otherwise dequantize weight to bf16

we should not do this, as this silently changes what the user requested. If there is a path which is unsupported (for example, per-tensor quant), throw an exception.

[vkuzo]

move to top of file, make sure to gate by the relevant pytorch version if needed

[vkuzo]

there should be no need to detect transposed here, the only syntax which we should support is grouped_mm(act, weight.transpose(-2, -1), ...)

[vkuzo]

high level approach looks good! I had some commments inline. Also, could you please rebase to get onto a revision with green CI.

[vkuzo]

@claude review this PR

[claude]

Claude finished @vkuzo's task in 2m 15s —— View job

---

PR Review: Add _grouped_mm dispatch handler to Float8Tensor

• Read changed files
• Analyze implementation
• Review tests
• Post feedback

---

Overall the approach is sound and follows the established pattern from the NVFP4 _grouped_mm handler. However, I agree with @vkuzo's inline comments — there are a few design issues that should be addressed before merging. Here's my full review:

---

Issues

1. Import should be at module level (not inside function body)

File: torchao/quantization/quantize_/workflows/float8/float8_tensor.py:1099

from torch.nn.functional import ScalingType, scaled_grouped_mm

This import is inside the function body. Per @vkuzo's comment, move it to the top of the file and gate it behind the relevant PyTorch version check (similar to how other version-gated imports work in this codebase). The ScalingType and scaled_grouped_mm APIs were introduced in torch 2.8+, so a version guard is appropriate.

Fix this →

2. Non-RowWise path should raise an error, not silently dequantize

File: torchao/quantization/quantize_/workflows/float8/float8_tensor.py:1117-1119

# Non-RowWise: dequantize weight only
if not is_b_rowwise:
    return torch._grouped_mm(mat_a, mat_b.dequantize(), offs=offs)

As @vkuzo noted, silently dequantizing when the user configured a non-RowWise granularity (like PerTensor) changes the behavior without informing the user. The user expects quantized execution, but gets bf16 matmul instead. This should raise a NotImplementedError or similar error clearly stating that only PerRow granularity is supported for _grouped_mm.

This also means test_fp8_grouped_mm_dequantize_roundtrip (which uses PerTensor) would need to be updated — it should either become an assertRaises test or be removed.

Fix this →

3. Remove transpose detection logic — only support the canonical grouped_mm(act, weight.T) form

File: torchao/quantization/quantize_/workflows/float8/float8_tensor.py:1106-1115

# Detect transposed mat_b and check RowWise scaling
is_b_transposed = mat_b.qdata.stride(-2) < mat_b.qdata.stride(-1)
if is_b_transposed:
    orig_bs = list(mat_b.block_size)
    orig_bs[-2], orig_bs[-1] = orig_bs[-1], orig_bs[-2]
    orig_sh = list(mat_b.shape)
    orig_sh[-2], orig_sh[-1] = orig_sh[-1], orig_sh[-2]
    is_b_rowwise = tuple(orig_bs) == (1,) * (len(orig_bs) - 1) + (orig_sh[-1],)
else:
    is_b_rowwise = _is_rowwise_scaled(mat_b)

Per @vkuzo's review comment, there should be no need to detect transposition here. The only supported syntax is grouped_mm(act, weight.transpose(-2, -1), ...), which is exactly how GroupedMMModel uses it in the test. The NVFP4 reference implementation (nvfp4_grouped_mm) simply asserts is_transposed is True and proceeds.

I'd suggest simplifying to:

is_b_transposed = mat_b.qdata.stride(-2) < mat_b.qdata.stride(-1)
assert is_b_transposed, (
    "_grouped_mm only supports the form: grouped_mm(act, weight.transpose(-2, -1), ...)"
)
is_b_rowwise = _is_rowwise_scaled(mat_b)  # check on the transposed view

Note: _is_rowwise_scaled checks block_size against shape, so you need to verify it works correctly on the transposed Float8Tensor. If not, you may need a simple manual check similar to what NVFP4 does.

Fix this →

4. Fallback with warnings.warn on scaled path failure is problematic

File: torchao/quantization/quantize_/workflows/float8/float8_tensor.py:1144-1147

except (RuntimeError, NotImplementedError) as e:
    warnings.warn(
        f"scaled_grouped_mm failed: {e}, falling back to dequant + _grouped_mm"
    )

Catching broad exceptions and silently falling back to bf16 matmul can mask real bugs and gives the user degraded performance without them realizing. This is the same concern as issue #2. If scaled_grouped_mm fails, it's better to let the error propagate so the user knows their hardware/config doesn't support this path, rather than silently running in bf16.

If a fallback is truly desired for specific known cases, it should be opt-in (e.g., via a config flag), not the default behavior.

5. Weight-only test uses the scaled path instead of dequant path

File: test/quantization/quantize_/workflows/float8/test_float8_tensor.py:1560

The docstring says "dequant path" but the test uses Float8WeightOnlyConfig(granularity=PerRow()). With PerRow granularity, the handler will enter the RowWise branch and call Float8Tensor.from_hp(mat_a, ...) to quantize activation, then try scaled_grouped_mm. This exercises the scaled path, not the dequant path. The docstring is misleading.

---

Minor Suggestions

1. Missing offs assertion: Unlike the NVFP4 handler which asserts offs is not None, this handler allows offs=None to propagate. Consider asserting it since _grouped_mm semantics require offsets for the 2D+3D case.

2. scale.squeeze(-1) correctness: At line 1137-1139, mat_a_q.scale.squeeze(-1) and b_scale.squeeze(-1) are passed to scaled_grouped_mm. Verify the scale shapes are correct for the 3D case — for a weight of shape (E, K, N) transposed to (E, N, K), the per-row scale should be (E, N) after squeeze. Double-check this matches what scaled_grouped_mm expects.

3. Consider adding torch.no_grad() context: The handler performs quantization of activations (Float8Tensor.from_hp). In inference contexts this is fine, but if someone accidentally hits this path during training, it would be clearer to assert not mat_a.requires_grad or document this is inference-only.

---

Summary

The PR implements a useful feature for MoE FP8 inference. The main issues are:

1. Move import to top-level with version gate
2. Don't silently dequantize for unsupported granularities — raise an error
3. Don't detect transpose — just assert it (follow NVFP4 pattern)
4. Don't silently fall back on scaled_grouped_mm failure — let it fail loudly

After addressing these (which align with @vkuzo's inline comments), this should be in good shape.

---

[xiaowangintel]

@claude Analyze the failure cases above.

[claude]

Claude finished @xiaowangintel's task in 2m 13s —— View job

---

Analysis of CI Failure Cases

• Gather context on the CI failures
• Read the changed files to understand the code
• Analyze root cause of exit code 139 (SIGSEGV)
• Provide findings and recommendations

---

Findings

Both failures are caused by the same root issue: a segmentation fault in an existing PT2E x86 inductor fusion test — NOT in your new _grouped_mm code.

CPU Nightly Failure

• Crashing test: test/quantization/pt2e/test_x86inductor_fusion.py::TestPatternMatcher::test_fp8_qlinear_add_cpu_use_relu_False_mixed_bf16_False_dynamic_False_reshape_a_False_M_1_inplace_add_False
• Error: Fatal Python error: Segmentation fault at line 2117 in test_x86inductor_fusion.py
• Stack trace: crash occurs inside torch._ops.__call__ → torch._functorch._aot_autograd.runtime_wrappers.__torch_dispatch__ → inductor compiled code

CUDA Nightly Failure

• Crashing test: test/quantization/pt2e/test_x86inductor_fusion.py::TestPatternMatcher::test_fp8_qlinear_cpu (different test, same test file)
• Error: Same Fatal Python error: Segmentation fault at line 1615 in test_x86inductor_fusion.py
• Stack trace: identical pattern — crash in inductor-compiled FP8 quantized linear execution

Key Observations

1. Not caused by this PR's changes. The segfaults occur in test/quantization/pt2e/test_x86inductor_fusion.py, which is a pre-existing test file that this PR does not modify. This PR only changes:

   • torchao/quantization/quantize_/workflows/float8/float8_tensor.py (implementation)
   • test/quantization/quantize_/workflows/float8/test_float8_tensor.py (new tests)

2. Nightly-specific issue. The test (CPU 2.11) and test (CUDA 2.11) jobs both passed successfully. The crashes only occur in the nightly builds, indicating this is a bug in PyTorch nightly (likely an unstable inductor codegen path for FP8 operations on x86).

3. The crash is in PyTorch's inductor/AOT autograd runtime, not in torchao code. The stack traces show the crash at:

   torch._ops.__call__ → torch._functorch._aot_autograd.runtime_wrappers.__torch_dispatch__
   

   inside inductor-generated kernel code (.py files in /tmp/).

4. This is a known type of issue — PyTorch nightly FP8 inductor support for x86 CPU is still stabilizing. The test names (test_fp8_qlinear_add_cpu, test_fp8_qlinear_cpu) suggest these are FP8 quantized linear fusion tests under x86 inductor, which are known to be fragile on nightly.

Recommendation

These failures are unrelated to your PR. As @vkuzo suggested, rebasing onto a newer main commit (once the nightly stabilizes or the upstream fix lands) should resolve this. You can verify by checking if the same tests fail on main without your changes.

To confirm independently, you could also run just your new tests in isolation:

pytest test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k "grouped_mm" -v

---

[xiaowangintel]

@vkuzo Your suggested changes have been incorporated. Please review them again.

[vkuzo]

@claude review

[claude]

Claude finished @vkuzo's task in 5m 16s —— View job

---

PR Review: Add _grouped_mm dispatch handler to Float8Tensor

• Read changed files and gather context
• Analyze implementation against previous review feedback
• Check test coverage
• Post review findings

---

The PR author has addressed the main feedback from the previous review. Here's my assessment of the current state:

Previous Feedback — Status

Feedback item	Status
Move import to top of file	Done (line 13)
Non-RowWise should raise, not silently dequant	Done (line 1119-1122)
Remove transpose detection, just assert it	Done (line 1108-1109)
Remove silent fallback on scaled_grouped_mm failure	Done (removed entirely)
Assert offs is not None	Done (line 1106)

The implementation is clean and follows the pattern from both _float8_addmm_impl (for weight-only dequant path) and the NVFP4 handler (for the assert-transposed pattern).

---

Remaining Issues

1. Granularity check uses activation granularity as a proxy for weight granularity

File: torchao/quantization/quantize_/workflows/float8/float8_tensor.py:1118

if not isinstance(act_quant_kwargs.granularity, PerRow):
    raise NotImplementedError(...)

This checks the activation granularity, relying on the fact that Float8DynamicActivationFloat8WeightConfig normalizes a single granularity argument into the same value for both activation and weight. This works for the common case, but if someone were to pass granularity=[PerRow(), PerTensor()] (PerRow activation + PerTensor weight), the check would pass but scaled_grouped_mm would receive incorrectly-shaped weight scales.

Consider adding an explicit check on the weight's block_size to be defensive:

# After the act_quant_kwargs.granularity check, also verify weight is row-scaled
# mat_b is transposed, so its original block_size was [1, 1, K], now [1, K, 1]
assert mat_b.block_size[-1] == 1 and mat_b.block_size[-2] == mat_b.shape[-2], (
    f"Expected row-wise scaled weight, got block_size={mat_b.block_size}"
)

This is a minor/defensive point — the current code works correctly for all supported configs.

2. Missing version guard on torch.nn.functional.ScalingType import

File: torchao/quantization/quantize_/workflows/float8/float8_tensor.py:13

from torch.nn.functional import ScalingType, scaled_grouped_mm

ScalingType and scaled_grouped_mm in torch.nn.functional are relatively new APIs. The NVFP4 handler avoids this by importing inside the function body (lazy import). The mx_tensor.py prototype imports it at the top without a guard, and CI passes on torch 2.11, so this appears safe for the currently supported versions.

However, if torchao still supports torch < 2.8 in some configurations, this would break at import time. @vkuzo — is this acceptable given the minimum supported torch version for this file?

3. b_scale transpose + squeeze assumes specific scale layout

File: torchao/quantization/quantize_/workflows/float8/float8_tensor.py:1125-1126

b_scale = mat_b.scale.transpose(-2, -1)
...
scale_b=b_scale.squeeze(-1),

After the Float8Tensor transpose handler runs:

• mat_b.scale has shape (E, 1, N) (original (E, N, 1) transposed)
• .transpose(-2, -1) → (E, N, 1)
• .squeeze(-1) → (E, N)

This produces the correct shape for scaled_grouped_mm with ScalingType.RowWise on a (E, K, N) weight matrix. The logic is correct but the double-transpose (once in the Float8Tensor handler, once here) is a bit subtle. A brief inline comment on why the transpose is needed would help future readers, though I defer to your style preferences.

---

Looks Good

• Weight-only path (line 1114-1115): Correctly follows _float8_addmm_impl pattern — when act_quant_kwargs is None, dequantize weight and run bf16 _grouped_mm.
• Test coverage: Three tests cover the dynamic act+weight path, weight-only dequant path, and the error case for unsupported granularity. SQNR thresholds are reasonable.
• GroupedMMModel test fixture properly uses weight.transpose(-2, -1) matching real MoE usage patterns.
• No silent fallbacks: The handler either succeeds via scaled_grouped_mm or raises — no hidden performance degradation.

---

Summary

The PR is in good shape after the revisions. The main remaining question is whether issue #1 (activation granularity as proxy for weight granularity) warrants a defensive check, and whether the top-level import (#2) needs a version guard. Neither is a blocker for correctness in the currently supported configurations.

---