[hao-aaron]

nvfp4_mega_moe took l2_activation_scale as a float kernel argument, so a CUDA graph replays the value it was captured with. vLLM computes this scale from the checkpoint's down-projection input_scale (the max over experts). When an RL weight update changes that scale (for example verl's NVFP4 W4A4 QAT, which sends the observed activation scales on every sync), a captured graph keeps quantizing the routed L2 input with the old scale while l2_alpha, a device vector, already holds the new product. vLLM tracks this as a known limitation of reloading weights under CUDA graphs (vllm-project/vllm#64692).

Change

• l2_activation_scale is now an optional one-element CUDA float32 tensor, like l1_alpha and l2_alpha. The epilogue reads it once, after cudaGridDependencySynchronize, so a replay sees values written after capture. None means a unit scale.
• API change: the op schema's float l2_activation_scale=1.0 becomes Tensor? l2_activation_scale=None in the same position, and the Python API rejects a float. Callers that pass a float must pass torch.full((1,), scale, device='cuda') instead. vLLM's call site changes with its pin bump. The tensor must be float32, one element, on y's device. The old positive and finite check on the float is gone, since checking a device value would need a sync.
• Docstring and AI/nvfp4-mega-moe.md describe the tensor form.

Tests

On GB200 (SM100), 4 GPUs, torch 2.14 / CUDA 13.0:

• tests/test_nvfp4_mega_moe.py passes the scale as a tensor and gains check_graph_replay_reads_scales, run for every NVFP4 shape: capture a graph, double the scale and l2_alpha in place, and require the replay to match an eager call with the doubled values. A replay that kept the captured scale would double the output.
• --check-api also rejects a float scale, a float64 tensor and a two-element tensor.
• All pass with:
  • --tokens 1 17 129 --hidden 512 --intermediate-hidden 256 --experts 8 --check-api
  • the same shapes with --input-scale 0.5 --weight-scale 0.25 --mid-scale 0.125
  • default (MiniMax-M3) shapes, --tokens 1 64 1024 --mid-scale 0.25 --weight-scale 0.5
  • small shapes with --x-scales --mid-scale 0.25
  • small shapes with default (unit, None) scales
• Negative control: capturing the graph with a copy of the scale made before capture fails the replay check, 497/512 elements mismatched.

Performance. --benchmark, MiniMax-M3 shapes, --mid-scale 0.25, one node, max over ranks of the median kernel time, mean of 2 passes. dev passes the scale as a float, this PR as a tensor:

tokens/rank	dev	this PR
1	62.6 µs	62.6 µs
8	155.9 µs	155.7 µs
64	188.7 µs	188.6 µs
512	200.4 µs	200.7 µs
2048	429.3 µs	426.2 µs

vLLM follow-up

After this merges, vLLM bumps its DeepGEMM pin and passes a device buffer refilled in place from DeepGemmMegaMoEExperts (and Step-5's MegaMoE). Its Triton prepare kernel, which computes l2_alpha, reads the same buffer.

🤖 Generated with Claude Code

https://claude.ai/code/session_01SeFSSrpw7yTpLz2C2sHDKA