[Dev] Allow FP8/FP4 with AbsorbedMLA up-projection recompute

[niyunsheng]

Summary

AbsorbedMLASelfAttention asserted that quantization was unsupported whenever "mla_up_proj" selective recompute was enabled, blocking FP8 and FP4 training for any model that needs that recompute to fit in memory. The assertion arrived with the initial absorbed-MLA implementation (#3193 / #3198) as a not-yet-validated placeholder rather than a guard for a known failure.

The quantized replay is safe for the same reason it is in MLASelfAttention, which runs the identical CheckpointWithoutOutput pattern with no such assertion (multi_latent_attention.py):

• The checkpoint records the forward recipe and amax state (ctx.fp8 /
  ctx.fp8_recipe under activation_recompute_forward(recompute_phase=False)) and replays under the recorded fp8_autocast with recompute_phase=True.
• The only quantized operation inside qkv_up_proj_and_rope_apply is the Q up projection — the same module set MLA already replays.
• The absorption einsum reads the K up-projection weight directly rather than through the checkpoint context. That weight is a persistent parameter: under
  fp8_param_gather it is remapped to TE quantized storage that outlives the backward pass (param_and_grad_buffer.py), and the only path that frees param storage is explicit weight CPU offload. So the forward and the replay observe the same values.

DeepSeekV4HybridAttention already runs the same pattern without the assertion.

Test plan

Added test_fp8_up_proj_recompute_parity, which drives two identically-initialized modules — one recomputing the up projection, one not — over the same inputs under get_fp8_context, and requires bitwise-identical (atol=0, rtol=0) outputs and per-parameter gradients. Recompute is meant to be an exact replay, so any drift would indicate a different quantization scale or a stale weight; a tolerance-based comparison would hide exactly the failures this guards against.
A bf16 reference run guards the parity check against passing vacuously: the same weights go through an fp8=False module, and the FP8 output must track it closely and differ from it bitwise — which is only possible if quantization actually engaged.

Coverage: tensorwise / mxfp8 / blockwise recipes x sbhd / thd layouts x combined / split K-V up-projection module specs.

Run on GB200 (mxfp8 needs Blackwell, so nothing is capability-skipped there):

pytest tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py
====================== 76 passed, 558 warnings in 52.58s =======================

76 = 64 pre-existing test_functionality cases (unchanged, no BF16 regression) + 12 new FP8 cases, 0 skipped.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[buptzyb]

/ok to test 510d14c

[kunlunl]

/claude strict-review

[claude]

Strict review passed — no significant issues found. LGTM

Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1 (non-blocking)

Verified the change is safe:

• The removed assertion was a not-yet-validated placeholder, not a guard for a known failure.
• The new comment's two load-bearing claims check out against the source: the only quantized op inside qkv_up_proj_and_rope_apply is the Q up projection (linear_q_up_proj/linear_q_proj), and the K up-projection weight is read directly via _get_kv_up_weights() (a persistent parameter), not through the checkpoint context — so the forward and the recompute replay observe identical values.
• The identical CheckpointWithoutOutput(fp8=quantization) pattern already ships in MLASelfAttention (multi_latent_attention.py:981) with no assertion — and there it quantizes both the Q and KV up-projection GEMMs, making the absorbed path a strict subset of an already-validated one.
• The atol=0, rtol=0 bitwise-parity test is the appropriate contract: recompute is meant to be an exact replay, so any quantization-scale or stale-weight drift would fail it. Coverage spans tensorwise/delayed/mxfp8 × sbhd/thd × combined/split KV up-proj.

Non-blocking suggestion (test coverage): The code path also enables FP4 (self.config.fp8 or self.config.fp4) and the PR title says "FP8/FP4," but the test only exercises FP8 recipes; the test's guarantee is effectively FP8-scoped. FP4 shares the exact same code path and requires TE ≥2.7 + capable hardware, so deferring it is reasonable — worth a note in the test docstring or a follow-up.

Risk: Low. A single-line assertion removal plus an accurate explanatory comment, mirroring an existing validated path and backed by a strong bitwise-parity test.

[niyunsheng]

Non-blocking suggestion (test coverage): The code path also enables FP4 (self.config.fp8 or self.config.fp4) and the PR title says "FP8/FP4," but the test only exercises FP8 recipes; the test's guarantee is effectively FP8-scoped. FP4 shares the exact same code path and requires TE ≥2.7 + capable hardware, so deferring it is reasonable — worth a note in the test docstring or a follow-up.

Good catch — I will add nvfp4 coverage since I have Blackwell available.

[FDecaYed]

LGTM

[FDecaYed]

/ok to test 713fcab

[niyunsheng]

/ok to test c5d3cdd

[niyunsheng]

/ok to test 6648492

[niyunsheng]

It failed again. Are you sure it's not related to your changes?

@buptzyb CI is green now. Root cause was an import side effect: the test imported helpers from tests.unit_tests.determinism.utils, which transitively imports torch.testing._internal.common_utils and freezes torch backend global flags (disable_global_flags()), deterministically breaking test_te_layers_batch_invariant.py later in the same session. Fixed in 6648492 by inlining the two RNG helpers so torch.testing._internal is never loaded in this bucket.