test(mxfp6): cover BF16 sensitive-layer routing at the real 57-block shape

[gphuang]

Summary

Every existing MXFP6 spec test runs with sensitive_layers_enabled=False, so the enabled path — the one that keeps trailing transformer blocks in BF16 while the rest run MXFP6 — has no coverage at all. An off-by-one there is silent: the model builds, trains, and converges while a different set of blocks than the one disclosed is running in BF16.

This adds TestMXFP6SensitiveLayerRouting, which exercises that path at Flux 12B's real 19 + 38 block shape, because the boundary is the thing being tested.

What it covers

• Exactly the last four, and exactly the last eight, of 57 blocks leave the MXFP6 backend.
• sensitive_layers_start and sensitive_layers_end route independently of each other.
• With the feature disabled, all 57 blocks stay on MXFP6 — the assumption the other spec tests rely on, checked at the real layer count rather than at two layers.
• Retained blocks reach a BF16 backend, not merely a non-MXFP6 one. sensitive_layer_precision also accepts tw_fp8, so asserting "not MXFP6" alone would pass on a config that quietly retained those blocks in FP8.

Test plan

• Four test methods (five cases, one parametrized) in tests/unit_tests/backends/megatron/test_primus_turbo_mxfp6_local.py, guarded by requires_mxfp6 and following the SimpleNamespace spec-test pattern already used in that file
• No line exceeds the repo's 110-column black setting
• Not executed locally — the authoring node has no torch or ROCm, so CI is the first run of these tests. Holding this PR in draft until CI reports.
• black / pre-commit run (unavailable on the authoring node)

Written while planning a precision-placement experiment that depends on this routing being correct, so the coverage gap is worth closing regardless of that experiment's outcome.

[gphuang]

Note: CI cannot execute these tests yet, and the lint failure is pre-existing

Author review of the full diff. No blocking findings in the change itself; two CI facts are worth recording before anyone reads the red X.

The code-lint failure is not from this PR. The job runs pre-commit over all files, and the base branch's own last run reformats the identical four files with the identical 4 files reformatted, 1265 files left unchanged. This branch adds no new file to that set. I could not run black on the authoring node to confirm my added lines are individually clean, since it has no pip; the added lines do stay under the repo's 110-column setting.

That pre-existing failure is also what blocks execution. build-docker needs code-lint, and run-unittest-torch needs build-docker, so the new tests are skipped rather than run. Draft status alone would skip them too. These tests have therefore never executed anywhere. Marking the PR ready will not change that while the base branch is lint-dirty.

Plan: run them directly with pytest inside the MXFP6 container on the next GPU node, and report the result here before taking the PR out of draft.

[gphuang]

Result: the five routing tests pass on real hardware

Ran them inside the v26.5 MXFP6 dev container on an MI355X node — 5 passed in 6.48s. This is the first time the test has executed anywhere: Primus CI cannot build the docker image on this base branch, so the torch unit-test job skips.

That closes the gap the PR was opened for. Before this run, the claim that sensitive_layers_end retains exactly the last 4 or last 8 of Flux 12B's 57 blocks in BF16 rested on reading get_flux_layer_spec; it is now executed. The disabled case and the "retained blocks are BF16, not another low precision" case pass too, so the feature is safe to switch on in a real recipe.

No code changes since the last push. The code-lint failure is pre-existing on the base branch, which fails the same check on the same four files.