NVFP4-CSF: expand the next layer's scales during its attention

[voipmonitor]

Layers of an NVFP4-CSF model share one scale scratch. Every eager prefill call expanded its experts' scales into it serially, right before its MoE. In a GLM-5.3-Flash glm53-tp2 A4 prefill trace this took 4.1% of the GPU time, with nothing overlapping it.

After a layer's MoE, Nvfp4CsfMoEMethod now expands the following layer's scales on a side stream, ordered after that MoE. The following layer waits for it and binds with scales_expanded=True (b12x expand_scales, local-inference-lab/b12x#477). The expansion therefore overlaps the following layer's attention instead of delaying its MoE. Scope and safety:

• Only eager steps whose calls expand prefetch: A4 prefill rows, or A16 calls above b12x's stage-read limit. Decode-only steps (which read compressed scales per stage) and CUDA graph capture are untouched.
• A prefetch is used only by the next layer of the same forward pass. Any pending expansion is waited for before the scratch is used.
• VLLM_B12X_CSF_SCALE_PREFETCH=0 turns it off. It applies to W4A16 NVFP4-CSF layers (GLM-5.3-Flash).

Evidence

glm53-tp2 (stored QAD CSF checkpoint, TP2/DCP2 MTP3, DMA all-reduce) on two RTX PRO 6000 Max-Q, prefetch off vs on (two runs each):

	2K	8K	32K
A4 prefill, off	9,254 / 9,169	10,083 / 9,981	10,391 / 10,232
A4 prefill, on	9,234 / 9,300	10,298 / 10,209	10,690 / 10,549
A16 prefill, off	7,310 / 7,238	7,713 / 7,630	7,954 / 7,895
A16 prefill, on	7,311 / 7,245	7,811 / 7,715	8,129 / 7,987

Decode is unchanged (C1/C8 171/542 → 177/518 with A4 and 169/526 → 171/540 with A16, within acceptance noise). Needle-checksum with the prefetch, 3 salted runs each: A4 495, 490, 492 (492, 492, 490 before); A16 498, 495, 497.

Tests

tests/quantization/test_nvfp4_csf.py, tests/kernels/moe/test_b12x_a4_prefill.py: 25 passed (beta image with this branch overlaid, one RTX PRO 6000). New test_scale_prefetch_expands_the_next_layer_for_this_forward_only, with real CUDA streams, checks that:

• layer N's MoE runs without the flag and then expands layer N+1 on the side stream;
• layer N+1 skips its expansion;
• a prefetch left from another forward pass is waited for but not reused;
• decode-only steps expand nothing.

ruff check / ruff format --check: clean. No other PR changes this path. AI assistance (Claude) was used; the change was reviewed and tested as described.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• dev/*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: d586200f-b9bd-40bc-9e51-ecf7d9fa3b4a

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[voipmonitor]

TP4 check (glm53-flash profile, stored QAD CSF checkpoint, A16 prefill as in the profile default, four RTX PRO 6000 Max-Q, two runs each): 8K/32K prefill 11,335/11,969 and 11,233/11,881 tok/s with VLLM_B12X_CSF_SCALE_PREFETCH=0, 11,467/12,120 and 11,337/11,994 with it on (+1.0%/+1.1%). Decode C8/C16 832/1,179 vs 841/1,187.