[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: 1/5

[Medium impact] The PR extends custom quantizer behavior across grouped forward calls and reference NVFP4 use.

The PR is not ready to merge because grouped calls can still raise, omit bias, or compute the wrong backward result.

Findings

1.  Existing grouped calls now raise ▶
2.  NVFP4 grouped forward still fails ▶
3.  Grouped outputs lose their bias ▶
4.  Backward runs the forward calculation ▶
5.  Skipped updates still quantize weights ▶

Summary

The PR completes the reference MXFP4 quantization branch and adds custom-recipe calls to both grouped GEMM entries.

• Per-expert grouped GEMMs now call custom quantizers.
• Grouped-tensor forward GEMMs now call custom quantizers.
• MXFP4 reference quantization now prepares both input directions.
• Reference NVFP4 tensors support in-place updates and dimension sizing.
• Custom quantizers now quantize weights without reusing cached workspaces.

_{Reviews (7) · Last reviewed commit: "docs(custom_recipes): record the per-gro..." · Reviewed by Greptile}

[greptile-apps]

Existing grouped calls now raise

is_custom rejects plain GroupedTensorStorage objects and lists with an AssertionError. The existing grouped forward passes a plain GroupedTensorStorage as its output, so even non-custom grouped GEMMs now fail before reaching the native implementation. Check grouped objects through their quantizers and inspect list members instead of passing these containers to is_custom.

[greptile-apps]

NVFP4 grouped forward still fails

This call omits qresult_x and qresult_w. NVFP4QuantizerRef.qgemm requires both for non-pow2 NVFP4 because they carry the global scales. A grouped forward using the supplied NVFP4 reference recipe therefore raises instead of computing its output. Pass act and w through these arguments, as custom_gemm does.

[greptile-apps]

Grouped outputs lose their bias

The helper accepts use_bias=True for forward, but passes None as the bias and never adds bias[i] afterward. GroupedLinear relies on this call to add its biases, so custom recipes that otherwise run return silently wrong outputs. Pass each bias to qgemm, apply it afterward, or explicitly reject this option.

[greptile-apps]

Backward runs the forward calculation

This guard rejects only grad and use_bias, but the helper always runs GEMMType.FPROP on rowwise weights. GroupedLinear requests input gradients with grad=True, layout="NN", and no bias. For square weights, this can pass every shape check while computing the gradient with the wrong weight direction. Reject grad=True and non-TN layouts until those paths are implemented.

[greptile-apps]

Skipped updates still quantize weights

quantize_ discards noop_flag and always quantizes. The weight-cache caller passes this flag to skip updates on later microbatches during CUDA graph replay. Ignoring it repeats weight quantization on every replay and removes that cache benefit. Honor the flag without adding a CPU read that would prevent graph capture.

[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.