[Bugfix][Attention] Recompute FlashInfer BMM scales when a weight reload changes the KV scales

[aoshen02]

Purpose

FlashInferImpl caches its TRTLLM BMM scales the first time forward runs:

if self.bmm1_scale is None:
    self.bmm1_scale = self.scale * layer._q_scale_float * layer._k_scale_float
if self.bmm2_scale is None:
    self.bmm2_scale = layer._v_scale_float

The cache is never invalidated, so a weight reload that changes the k/v scales after
the first forward keeps the old ones. With an FP8 KV cache every attention layer then
dequantizes K/V with the wrong scale and the model emits garbage.

The common way to hit this is the RL weight-sync path:

1. The engine starts with load_format=dummy, so k/v scales are 1.0 while profiling,
   warmup and CUDA-graph capture run.
2. The first layerwise refit loads the real checkpoint, including its static FP8
   k/v scales.
3. kv_cache.py updates layer._k_scale_float and _v_scale_float, but
   impl.bmm1_scale and impl.bmm2_scale keep the dummy values.

Two more pieces of the same problem:

• trtllm-gen decode runs inside full CUDA graphs, and host-scalar scales are baked in at
  capture time, so recomputing the floats is not enough for graph replay.
• With attention + output-quant fusion (fuse_attn_quant, on at -O3 for quantized
  models), layer._o_scale_float (the next quant op's input scale) is read once and
  never refreshed either.

Changes

vllm/v1/attention/backends/flashinfer.py:

• bmm1_scale / bmm2_scale are computed on every forward from the layer's
  _q/_k/_v_scale_float (a few float multiplies, as get_xqa_bmm1_scale already
  does). No cache, so nothing to invalidate.
• The fused FP8-output factor stays bmm2_scale / layer._o_scale_float, applied every
  forward.
• trtllm_gen_decode_bmm_scales(): while a CUDA graph is being captured and the KV cache
  is quantized, trtllm-gen decode gets device scales computed from the layer's
  _q_scale, _k_scale, _v_scale buffers (and / output_scale for fused FP8
  output). A reload updates those buffers in place, so graphs captured before the
  reload replay with the new scales. Outside capture (eager, piecewise) it returns the
  host floats, so no extra kernels are launched per step. DCP decode keeps host scalars.

vllm/model_executor/layers/attention/attention.py:

• Attention.process_weights_after_loading resets _o_scale_float, so the next eager
  forward re-reads the fused output scale after a reload. Layerwise reload calls this
  for every attention layer, including layers whose own scales were not reloaded.

vllm/v1/attention/backends/mla/{flashinfer_mla,flashinfer_mla_sparse,tokenspeed_mla}.py:

• Same lazy cache; the scales are now computed on every call, so eager and piecewise
  execution pick up reloaded scales. Full-graph replay of these MLA decode kernels still
  uses capture-time host scalars (see "Not covered").

tests/v1/attention/test_attention_backends.py:

• test_flashinfer_trtllm_gen_decode_scales_follow_weight_reload drives
  FlashInferImpl.forward (FP8 KV, fused FP8 output) through a stubbed
  trtllm_batch_decode_with_kv_cache and checks the scales the kernel receives:
  host floats before and after a reload (k/v and o scales changed, o refreshed via
  Attention.process_weights_after_loading), and device tensors computed from the
  scale buffers while capturing. Uses the file's FlashInfer skip marker.

Not covered (follow-ups)

All of these keep capture-time scales in full CUDA graphs after a reload that changes
the scales; eager and piecewise execution are correct. Each needs FlashInfer to accept
device scales on that entry point (or a graph recapture), not a vLLM-side change:

• Non-trtllm FlashInfer decode (BatchDecodeWithPagedKVCacheWrapper.run). It takes
  q/k/v_scale as Python floats and folds them into sm_scale on the host, so a
  captured graph keeps them. This path is used when trtllm-gen decode is unavailable
  (head geometry rejected by can_use_trtllm_attention, or non-SM100 hardware).
  Workaround: cudagraph_mode=PIECEWISE (or eager) when reloading FP8 KV scales on such
  configs.
• DCP decode (decode_context_parallel_size > 1): FlashInfer's DCP path takes host
  scalars only.
• XQA decode (SM90).
• The MLA trtllm decode kernels (FlashInfer MLA, sparse MLA, Tokenspeed MLA).
• NVFP4 fused output: o_sf_scale is a host float in
  trtllm_batch_decode_with_kv_cache; it is refreshed for eager/piecewise after a
  reload but not inside captured decode graphs.
• Native reload_weights() rejects load_format=dummy
  (NotImplementedError: Model reloading with dummy format), so the repro uses the
  layerwise reload API, the same path RL frameworks use.

Duplicate check

Run on 2026-10-05 against vllm-project/vllm, issues and PRs, all states. Queries:
bmm1_scale, k_scale reload, kv scale weight update flashinfer,
layerwise reload kv scale, _k_scale_float reload.

No issue or PR covers stale FlashInfer BMM scales after a reload. Related but different:

• [RFC] Weight Reload Correctness for RL vllm-project/vllm#48312 (RFC: Weight Reload Correctness for RL): MoE runtime-constant staleness; no
  attention or KV-scale item.
• [RFC] Sleep/Wake Correctness for RL vllm-project/vllm#48310 (RFC: Sleep/Wake Correctness for RL), [Bugfix][sleepmode][fp8 kv cache]: Fix FP8 KV cache + sleep(level=2) gibberish output vllm-project/vllm#28783, [Bugfix][Quantization] Stop sleep(level=2) from zeroing compressed-tensors KV scales vllm-project/vllm#57163: KV scales lost across
  sleep(level=2). This PR is about values cached on the impl, not the layer buffers.
• [RL][Reload] Preserve graph-visible runtime state across weight reloads vllm-project/vllm#49789 (graph-visible runtime state across reloads) and [Model][Core] Register every persistent device tensor as a parameter or buffer vllm-project/vllm#60009 (register persistent
  device tensors): tensor storage, not host-scalar caches.
• [Bugfix] Do not re-apply q/k/v scales on FlashInfer TRTLLM BF16-Q paths vllm-project/vllm#48016 and [Bug]: disable_flashinfer_q_quantization + fp8 KV + spec decode silently corrupts sliding-window models on SM100 (trtllm-gen); BF16-Q paths also mis-apply q/k/v scales vllm-project/vllm#47847: q/k/v scales misapplied on BF16-Q TRTLLM paths, a different bug.

Test plan and results

Environment: GB200 (SM100), image nemotron-final-vllm-538344b48f-...sqsh
(vLLM 538344b, torch 2.13.0a0, FlashInfer 0.7.0.post1). 538344b has the same code
as this PR's base in every touched file except tokenspeed_mla.py (not exercised). The
five changed source files of this PR were bind-mounted over the installed package
(the sixth changed file is the test). The same job also mounted two Mamba2 files from a
separate review fix (ssd_emit.py, mamba2_attn.py: a stride guard and a TP>1 reject,
both inactive at TP=1 with the paged Mamba buffers). A byte compare of all 2588
vllm/**/*.py against that source tree reported 0 mismatches.

Unit tests (GPU 0, HF_HUB_OFFLINE=1):

pytest tests/v1/attention/test_attention_backends.py -k trtllm_gen_decode_scales_follow_weight_reload
  -> 1 passed
pytest tests/v1/attention/test_attention_backends.py -k "flashinfer and not correctness"
  -> 24 passed, 1 skipped, 8 failed

The 8 failures are test_flashinfer_cross_layer_layout[*]. All 8 fail while building
ModelConfig('meta-llama/Meta-Llama-3-8B') in tests/v1/attention/utils.py (that config
is not in the offline HF cache on the node), before any attention code runs, so they
say nothing about this change. They need a rerun with HF access.

NLL repro: Nemotron-3.5-Lightning-30B-A3B NVFP4 (ModelOpt mixed precision, static FP8
k/v scales on 6 attention layers), TP1, kv_cache_dtype=fp8_e4m3, FlashInfer attention
with trtllm-gen decode, VLLM_BATCH_INVARIANT=0. Start LLM, generate, then through
collective_rpc run initialize_layerwise_reload, model.load_weights on every
checkpoint tensor and finalize_layerwise_processing. Reset the prefix cache, then
measure prompt NLL on two fixed texts and greedy-decode "Question: what is 12*13?
Answer:".

init	mode	patch	prompt NLL after refit	greedy
dummy	CUDA graphs	none (earlier run)	11.40 / 11.72 (random ≈ 11.78)	garbage
dummy	CUDA graphs	this PR	1.0116 / 0.7437	" 156"
checkpoint	CUDA graphs	this PR	1.0116 / 0.7484	" 156"
dummy	eager	this PR	0.9874 / 0.8749	" 156"

• With the patch, dummy init plus refit decodes correctly under CUDA graphs. The first
  prompt's NLL matches checkpoint init. The second differs by 0.005, the same
  load-vs-reload gap the earlier revision showed (0.744 vs 0.750); it is not caused by
  the scale caching.
• This model's attention o_proj is not quantized, so the fused-output (o scale) path
  is covered only by the unit test.
• The MLA backend changes are not exercised by these runs (no MLA model was tested).

pre-commit run --files <changed files>: all hooks pass.

Throughput was measured on the first revision only (paired BI=0 serve legs, decode TPOT
within noise). This revision removes the per-step device-scale kernels outside capture
and has not been re-benchmarked.

AI assistance

AI assistance (Claude) was used to localize the bug, write the patch, the test and this
description. The human submitter must review every changed line and rerun the tests
before this goes anywhere.

• Human submitter has reviewed every changed line and run the tests.