[NolanHo]

What this changes

Makes CustomRecipe usable beyond plain Linear, and completes the MXFP4 (block-32 / E8M0) path in
the reference implementation. Without these, a custom quantizer can never be used for MoE experts and
the reference MXFP4 branch cannot run at all.

general_gemm already dispatches to a custom quantizer's qgemm; neither grouped GEMM entry does,
so GroupedLinear — the MoE expert path — goes straight to nvte_grouped_gemm with native tensors.

grep -c 'is_custom|custom_gemm|qgemm'  linear.py          -> 2
                                       grouped_linear.py -> 0

Changes, each tied to an observed failure

pytorch/custom_recipes/reference_nvfp4.py

1. The pow_2_scales (MXFP4) branch is unfinished. It validates the configuration and then falls
   through to the common tail, which raises:
   UnboundLocalError: cannot access local variable 'row_input' where it is not associated with a value
   The source's own # TODO(etsykunov): Fix bug where global_amax_row and global_amax_col are not defined points at the same hole. This prepares row_input / col_input and both amax values on
   that branch.
2. _split_quantization can hand this a 0-element split, where torch.max() without a reduction
   dim raises RuntimeError: max(): Expected reduction dim to be specified for input.numel() == 0.
   Guarded.
3. size(*args, **kwargs) ignores its argument and always returns the full logical shape, so
   x.size(0) yields a torch.Size. GroupedLinear._validate_or_alloc_output does
   torch.empty((rows, cols)) with rows = x.size(0) and dies with
   TypeError: empty(): argument 'size' failed to unpack the object at pos 2 with error "type must be tuple of ints, but got torch.Size".
   Changed to size(dim=None), returning that dimension when asked. The no-argument behaviour is
   unchanged.
4. update_usage refuses a placeholder. A tensor with neither data nor data_t populated has
   nothing to generate from, yet TE raises
   RuntimeError: Cannot generate FP4 data, even from FP4 data transpose. Treated as a no-op instead.
5. No quantize_(). module/base.py::quantize_weight calls workspace.quantize_(...) when the
   workspace exposes it and otherwise falls back to tex.quantize(...), a C++ entry that cannot drive
   a Python reference quantizer — the workspace stays empty and the GEMM then fails on data is None.
   Implemented in Python.

pytorch/cpp_extensions/gemm.py

6. Custom-recipe dispatch for both grouped GEMM entry points.
   • general_grouped_gemm (list form): per expert, the activation rows are B[i][0:m_splits[i]], the
     weight is A[i], and the result goes to out[0][offset : offset+m_splits[i]]. The operand roles
     are taken from GroupedLinear's own call site (grouped_linear.py), which passes
     (weights_fp8, inputmats, [out], single_output=True, m_splits=...).
   • general_grouped_gemm_for_grouped_tensor (GroupedTensorStorage form): groups are sliced from
     offsets / scale_inv_offsets / rowwise_data / rowwise_scale_inv.

Both refuse contracts they cannot honour explicitly (GELU outputs, bias gradients, D_dtype, output
quantisation, grad) rather than producing a silently wrong result.

Testing

Measured on 8× sm_103 (B300), EP8 + CP8, with an E2M1/block-32/E8M0 CUTLASS kernel
(KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103) bound through CustomRecipe.qgemm:

• the custom qgemm matches this reference implementation bitwise on real quantised operands
  (max|Δ| = 0.0, reference magnitude ≈ 12.6) for the fprop direction;
• each change above is a fix for a failure reproduced end-to-end; the error strings quoted here are
  the observed messages.

Not exercised by this PR: grad/dgrad/wgrad for the grouped paths (explicitly refused), and
row_scaled_nvfp4 in the MXFP4 configuration (rejected upstream already).

[greptile-apps]

Confidence Score: 5/5

[Medium impact] No newly introduced blocking issue was found in the changes since the previous review.

Summary

Completes the reference MXFP4 quantization path and adds custom-recipe calls for grouped forward GEMMs.

• Grouped GEMMs now run each expert through its custom quantizer.
• Reference MXFP4 quantization now prepares both data directions.
• Custom quantized weights now refresh in place and report their logical dimensions.
• Reference FP4 GEMM now adds bias before converting the result.

_{Reviews (13) · Last reviewed commit: "fix(custom_recipes): reach the grouped-t..." · Reviewed by Greptile}

[ptrendx]

@NolanHo Please sign every commit in this PR. The DCO check is failing because commits are missing a Signed-off-by trailer. See the contribution guide.

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[NolanHo]

@ptrendx Sign-off is in place on every commit and DCO is green on the current head. I also worked
through the Greptile findings and added tests for the custom grouped paths — would you have a look?