[moe training] Add backward_override to the MXFP8 grouped GEMM

[wolfcomos]

Independent PR (no stack). Companion TorchTitan converter PR: wolfcomos/torchtitan#15.

Summary

Adds a backward_override knob to the MXFP8 grouped GEMM
(_to_mxfp8_then_scaled_grouped_mm / _MXFP8GroupedMM) so the backward
computation can be ablated without touching the forward. This is the op-level
support needed for recipe ablations that hold the quantized forward fixed and
vary only where gradients come from (e.g. RL fine-tuning recipes that need the
training backward to be consistent with a quantized inference forward).

The knob is threaded through MXFP8TrainingOpConfig and the
_quantize_then_scaled_grouped_mm dispatcher, so model-conversion callers
(e.g. torchtitan converters) reach it the same way they reach the existing
MXFP8 knobs.

A second commit fixes a latent crash class on the default quantized backward:
the CUDA dim1 cast kernel checks input.is_contiguous(), and autograd can
hand backward an expanded (stride-0) or otherwise non-contiguous
grad_output view (out.sum().backward() produces one). The fix
materializes with .contiguous() on both the dim1-cast wgrad operands and
the dgrad CuTe DSL quantize input (whose stride-0 handling is unverified) —
a no-op for the contiguous tensors these paths normally see — keeping the
kernel choices unchanged for existing callers.

Design notes

• None/"quantized" leave the default quantized backward byte-for-byte
  unchanged; "quantized" is just an accepted alias for None.
• "high_precision" saves the unpadded high-precision operands and computes
  both gradients with plain torch._grouped_mm calls (upper-bound gradient
  quality arm).
• "dequantized" differentiates the quantized-forward function itself: it
  saves the quantized forward operands and dequantizes them for the same two
  grouped GEMMs. This is the arm that matters for RL train/inference
  consistency — the gradients correspond to the function the quantized
  forward actually computes.
• The dequantized arm quantizes activations with the Triton dim0 kernel
  because the backward needs logical (pre-blocking) scales; the CuTe DSL 1x32
  kernel returns blocked-only scales and no unswizzler exists. The forward
  output is unchanged (see Numerics evidence).
• backward_override participates in MXFP8TrainingOpConfig.__eq__/__hash__
  so overridden configs never alias the default config (the config is a
  registered pytree constant used under dynamo nonstrict trace).
• The default forward/backward paths are intentionally untouched: no kernel
  choices or availability gates change in this PR beyond the .contiguous()
  fix described above (kept as its own commit).
• The override arms are bf16-only: they quantize saved operands with the
  bf16-only dim0 Triton cast, so non-bf16 operands are rejected at the
  op-level validation point rather than deep inside the kernel.

Numerics evidence

• The Triton dim0 quantizer used by the dequantized arm's forward produces
  qdata bitwise identical to the CuTe DSL 1x32 rceil kernel used by the stock
  forward; test_forward_bitwise_identical_across_overrides pins that the
  forward output is torch.equal across all override arms.
• test_high_precision_backward_matches_grouped_mm_reference pins the
  high-precision arm bitwise against the same two-torch._grouped_mm
  formulation.

Test plan

CPU (no GPU; module-level skip hygiene):

python -m pytest test/prototype/moe_training/test_mxfp8_backward_override.py -q -rs
SKIPPED [1] test/prototype/moe_training/test_mxfp8_backward_override.py:13: MXFP8 grouped GEMM backward overrides require SM100
1 skipped in 10.58s

Import/threading check (CPU): MXFP8TrainingOpConfig eq/hash asymmetry with
and without backward_override, backward_override present in the
_to_mxfp8_then_scaled_grouped_mm signature and in the
_quantize_then_scaled_grouped_mm dispatcher — verified in a CPU-only
container.

Lint: ruff check + ruff format --check (ruff 0.11.6, the version pinned by
.github/workflows/ruff_linter.yml) are green on all touched files.

GPU results come from the assembled stack tip on GB200, since ao CI has no
SM100 runner (the GPU legs run SM86 collection and an SM90 curated list). The
override arms and the forward-bitwise test passed there on the exploratory
branch, and the override arms also have end-to-end GB200 evidence from RL
recipe-ablation runs. Two caveats on this exact branch: (1) these tests gate
on torchao::mx_block_rearrange_2d_M_groups, a compiled CUDA extension —
source-checkout environments (and CI without SM100 runners) skip the whole
module, so a per-branch GPU run needs a built torchao on SM100; (2) the
exploratory branch ran the quantized wgrad dim1 casts with the Triton kernel,
while this PR keeps the default CUDA cast kernel. The stride-0 crash class is
closed by construction on both backward paths (the same .contiguous()
materialization on wgrad and dgrad), and test_noncontiguous_grad_output_wgrad
pins it wherever the suite runs with a built torchao. The tests additionally
require the CuTe DSL runtime (nvidia-cutlass-dsl) for the stock forward, like
the sibling test_mxfp8_grouped_mm.py.

Known follow-ups

• backward_override is a free string that overlaps semantically with
  wgrad_with_hp ("high_precision" subsumes it for both gradients); there
  is no interaction validation or Literal typing yet. We kept the op-level
  assert as the single validation point for now and can tighten to Literal
  plus an interaction check in a follow-up if preferred.
• "dequantized" parity has only been exercised with RCEIL scale calculation;
  FLOOR is reachable but untested.
• pad_token_groups_for_grouped_mm=True combined with the overrides is
  untested (the override backward feeds unpadded offsets where the default
  uses padded ones).

Draft on the fork to stage the upstream submission; supersedes the exploratory stack (#7/#8/#9)

Review pass (2026-08-30)

• The .contiguous() materialization now also covers the dgrad CuTe DSL
  quantize input (the e2e sum().backward() leg drives it with the same
  stride-0 gradient the wgrad fix pins); the fix commit message says so.
• The override arms assert bf16 operands at the op-level validation point
  (they quantize saved operands with the bf16-only dim0 Triton cast).
• _backward_override_grads_tuple spells out the per-argument Nones like
  the default return (no magic arity constant).
• Test tolerance comment reworded: ~5% observed, 0.15 is a deliberately loose
  bound against seed sensitivity.

[wolfcomos]

Closing to tidy the fork; the branch is preserved for a future upstream cut of the MXFP8 backward_override.