Cache quantized weights for full AC

[MarioSieg]

• Cache quantized weights
• Turn off a2a by default for mxfp8
• Some minor fixes

---

Note

Medium Risk
Monkey-patching torchao’s mx_mm autograd path affects all MXFP8 linear forwards/backwards; incorrect cache invalidation could cause silent wrong numerics, though tests cover checkpoint and inplace-update cases.

Overview
MXFP8 + full activation checkpointing: Recompute forward was re-quantizing weights on every checkpoint segment. This PR patches torchao’s mx_mm to cache each weight’s dim0 MXFP8 quantization on the tensor (invalidated via weight._version), reuse it on cache hits, and clear the cache after backward so training stays bit-identical to an uncheckpointed run.

Defaults and MoE: MXFP8Config.enable_a2a now defaults to false (docs explain when mxfp8 all-to-all helps vs hurts). The non-FP8 MoE grouped-GEMM path only casts expert weights to bf16 when they aren’t already bf16.

Tests: New Blackwell GPU tests cover cache hit/miss, inplace weight updates, per-tensor isolation, checkpoint grad parity, and bf16 sanity checks.

^{Reviewed by Cursor Bugbot for commit 3508fea. Bugbot is set up for automated code reviews on this repo. Configure here.}

[garrett361]

I don't really understand what this is doing. Seems like we're monkey-patching the forward of an autograd fn to change what's cached, but if so, why don't we need to patch the backward, too? How does the bwd know to consume the altered cached tensors?

[garrett361]

Would also like to see some unit testing regarding numerics vs a basic Linear layer, as @samsja said

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 505d99b. Configure here.}