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

tests/v1/attention/test_attention_backends.py

@@ -1014,6 +1014,116 @@ def mock_xqa(**kwargs):
     }
 
 
+@pytest.mark.skipif(
+    AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
+    reason="FlashInfer is not available.",
+)
+def test_flashinfer_trtllm_gen_decode_scales_follow_weight_reload(monkeypatch):
+    """A weight reload after the first forward (e.g. an RL refit of a
+    dummy-initialized engine) changes the k/v and fused-output scales. Eager
+    decode must see the new host scales; decode captured into a CUDA graph must
+    read device scales from the buffers the reload updates in place."""
+    from types import SimpleNamespace
+
+    from vllm.model_executor.layers.attention import Attention
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    fp8 = flashinfer_backend.FP8_DTYPE
+    impl = object.__new__(flashinfer_backend.FlashInferImpl)
+    impl.scale = 0.5
+    impl.kv_cache_dtype = "fp8"
+    impl.is_kvcache_nvfp4 = False
+    impl.head_size = 16
+    impl.dcp_world_size = 1
+    impl.o_sf_scale = None
+    impl.window_left = -1
+    impl.sinks = None
+    impl.need_to_return_lse_for_decode = False
+    impl._sinks_source = None
+
+    layer = SimpleNamespace(
+        impl=impl,
+        quant_config=None,
+        _q_scale=torch.tensor(1.0),
+        _k_scale=torch.tensor(1.0),
+        _v_scale=torch.tensor(1.0),
+        _prob_scale=torch.tensor(1.0),
+        _q_scale_float=1.0,
+        _k_scale_float=1.0,
+        _v_scale_float=1.0,
+        _o_scale_float=None,
+    )
+    output_scale = torch.tensor(0.5)
+    attn_metadata = flashinfer_backend.FlashInferMetadata(
+        kv_cache_layout=KVCacheLayout.LBHNC,
+        num_actual_tokens=2,
+        slot_mapping=torch.empty(0, dtype=torch.int64),
+        q_data_type_prefill=fp8,
+        q_data_type_decode=fp8,
+        num_decodes=2,
+        num_decode_tokens=2,
+        num_prefills=0,
+        num_prefill_tokens=0,
+        causal=True,
+        prefill=None,
+        decode=flashinfer_backend.FlashInferTrtllmAPIDecode(
+            kernel=flashinfer_backend.FlashInferDecodeKernel.TRTLLM_GEN,
+            block_tables=torch.zeros((2, 1), dtype=torch.int32),
+            seq_lens=torch.tensor([8, 8], dtype=torch.int32),
+            max_seq_len=8,
+            q_len_per_req=1,
+        ),
+        use_cascade=False,
+        cascade_wrapper=None,
+    )
+
+    seen = {}
+
+    def mock_decode(**kwargs):
+        seen["bmm1_scale"] = kwargs["bmm1_scale"]
+        seen["bmm2_scale"] = kwargs["bmm2_scale"]
+
+    monkeypatch.setattr(
+        flashinfer_backend,
+        "_get_trtllm_workspace_buffer",
+        lambda: torch.empty(1, dtype=torch.uint8),
+    )
+    monkeypatch.setattr(
+        flashinfer_backend, "trtllm_batch_decode_with_kv_cache", mock_decode
+    )
+
+    def decode():
+        query = torch.zeros((2, 1, 16), dtype=fp8)
+        impl.forward(
+            layer,
+            query,
+            None,
+            None,
+            torch.zeros((1, 1, 1, 32), dtype=torch.uint8),
+            attn_metadata,
+            torch.empty_like(query),
+            output_scale=output_scale,
+        )
+        return seen["bmm1_scale"], seen["bmm2_scale"]
+
+    assert decode() == (0.5, 2.0)
+
+    # A reload runs Attention.process_weights_after_loading, updates the scale
+    # buffers in place and replaces the host floats.
+    Attention.process_weights_after_loading(layer, torch.bfloat16)
+    layer._k_scale.fill_(0.25)
+    layer._v_scale.fill_(0.5)
+    layer._k_scale_float, layer._v_scale_float = 0.25, 0.5
+    output_scale.fill_(0.125)
+    assert decode() == (0.125, 4.0)
+
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    bmm1_scale, bmm2_scale = decode()
+    assert isinstance(bmm1_scale, torch.Tensor)
+    assert isinstance(bmm2_scale, torch.Tensor)
+    assert (bmm1_scale.item(), bmm2_scale.item()) == (0.125, 4.0)
+
+
 @pytest.mark.skipif(
     AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
     reason="FlashInfer is not available.",

vllm/model_executor/layers/attention/attention.py

@@ -588,6 +588,9 @@ def extra_repr(self) -> str:
 
     def process_weights_after_loading(self, act_dtype: torch.dtype):
         self.impl.process_weights_after_loading(act_dtype)
+        # A weight reload can change the next quant op's input scale; the next
+        # forward re-reads it.
+        self._o_scale_float = None
 
         # If we should not load quant weights, we initialize the scales to 1.0
         # as the default value. See [Note: Register q/k/v/prob scales in state dict]

vllm/v1/attention/backends/flashinfer.py

@@ -1987,6 +1987,29 @@ def process_weights_after_loading(self, act_dtype: torch.dtype):
         else:
             self.sinks.copy_(source_sinks)
 
+    def trtllm_gen_decode_bmm_scales(
+        self,
+        layer: torch.nn.Module,
+        output: torch.Tensor,
+        output_scale: torch.Tensor | None,
+    ) -> tuple[float | torch.Tensor, float | torch.Tensor]:
+        """BMM scales for trtllm-gen decode.
+
+        Host scalars are baked into a captured CUDA graph, so while capturing
+        with a quantized KV cache compute them on device from the scale buffers
+        a weight reload updates in place.
+        """
+        if not (
+            is_quantized_kv_cache(self.kv_cache_dtype)
+            and torch.cuda.is_current_stream_capturing()
+        ):
+            return self.bmm1_scale, self.bmm2_scale
+        bmm1_scale = layer._q_scale * layer._k_scale * self.scale
+        bmm2_scale = layer._v_scale
+        if output_scale is not None and output.dtype == FP8_DTYPE:
+            bmm2_scale = bmm2_scale / output_scale
+        return bmm1_scale, bmm2_scale
+
     def get_xqa_bmm1_scale(self, layer: torch.nn.Module, q_data_type: torch.dtype):
         bmm1_scale = self.scale
         if is_quantized_kv_cache(self.kv_cache_dtype):
@@ -2055,15 +2078,12 @@ def forward(
             # Profiling run.
             return output.fill_(0)
 
-        if self.bmm1_scale is None:
-            self.bmm1_scale = self.scale
-            if is_quantized_kv_cache(self.kv_cache_dtype):
-                self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float
-
-        if self.bmm2_scale is None:
-            self.bmm2_scale = 1.0
-            if is_quantized_kv_cache(self.kv_cache_dtype):
-                self.bmm2_scale *= layer._v_scale_float
+        # Recomputed every forward: a weight reload can change the KV scales.
+        self.bmm1_scale = self.scale
+        self.bmm2_scale = 1.0
+        if is_quantized_kv_cache(self.kv_cache_dtype):
+            self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float
+            self.bmm2_scale *= layer._v_scale_float
 
         prefill_use_trtllm = isinstance(attn_metadata.prefill, TRTLLMPrefill)
         decode_kernel = (
@@ -2102,15 +2122,15 @@ def forward(
             else:
                 raise ValueError(f"Unsupported output dtype: {output.dtype}")
 
-            # TRTLLM attn kernel requires to scale to pass as a host scalar,
-            # store the o scale as a host scalar in warmup run with cuda graph
-            # not enabled
+            # Host copy of the o scale for the kernels that take host scalars,
+            # read in an eager forward (warmup, or the first one after a weight
+            # reload resets it in Attention.process_weights_after_loading).
             if layer._o_scale_float is None:
                 layer._o_scale_float = output_scale.cpu().item()
-                if output.dtype == FP8_DTYPE:
-                    self.bmm2_scale = self.bmm2_scale / layer._o_scale_float
-                elif output.dtype == FP4_DTYPE:
+                if output.dtype == FP4_DTYPE:
                     self.o_sf_scale = layer._o_scale_float
+            if output.dtype == FP8_DTYPE:
+                self.bmm2_scale = self.bmm2_scale / layer._o_scale_float
 
         # IMPORTANT!
         # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
@@ -2611,6 +2631,12 @@ def forward(
                         device=decode_query.device,
                     )
 
+                # DCP decode takes host scalars only.
+                bmm1_scale, bmm2_scale = (
+                    (self.bmm1_scale, self.bmm2_scale)
+                    if use_dcp
+                    else self.trtllm_gen_decode_bmm_scales(layer, output, output_scale)
+                )
                 trtllm_batch_decode_with_kv_cache(
                     query=decode_query,
                     kv_cache=(
@@ -2620,8 +2646,8 @@ def forward(
                     block_tables=block_tables_decode,
                     seq_lens=seq_lens_decode,
                     max_seq_len=attn_metadata.decode.max_seq_len,
-                    bmm1_scale=self.bmm1_scale,
-                    bmm2_scale=self.bmm2_scale,
+                    bmm1_scale=bmm1_scale,
+                    bmm2_scale=bmm2_scale,
                     window_left=self.window_left,
                     sinks=self.sinks,
                     o_sf_scale=self.o_sf_scale,

vllm/v1/attention/backends/mla/flashinfer_mla.py

@@ -368,15 +368,12 @@ def forward_mqa(
         else:
             q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
 
-        if self.bmm1_scale is None:
-            self.bmm1_scale = self.scale
-            if is_quantized_kv_cache(self.kv_cache_dtype):
-                self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float
-
-        if self.bmm2_scale is None:
-            self.bmm2_scale = 1.0
-            if is_quantized_kv_cache(self.kv_cache_dtype):
-                self.bmm2_scale *= layer._k_scale_float
+        # Recomputed every call: a weight reload can change the KV scales.
+        self.bmm1_scale = self.scale
+        self.bmm2_scale = 1.0
+        if is_quantized_kv_cache(self.kv_cache_dtype):
+            self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float
+            self.bmm2_scale *= layer._k_scale_float
 
         return_lse = self.need_to_return_lse_for_decode
         workspace_buffer = _get_workspace_buffer(return_lse)

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -605,14 +605,12 @@ def _prepare_mqa_kernel(
         if self._workspace_buffer is None:
             self._workspace_buffer = _get_workspace_buffer(device)
 
-        if self.bmm1_scale is None:
-            self.bmm1_scale = self.scale
-            if is_quantized_kv_cache(self.kv_cache_dtype):
-                self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float
-        if self.bmm2_scale is None:
-            self.bmm2_scale = 1.0
-            if is_quantized_kv_cache(self.kv_cache_dtype):
-                self.bmm2_scale *= layer._k_scale_float
+        # Recomputed every call: a weight reload can change the KV scales.
+        self.bmm1_scale = self.scale
+        self.bmm2_scale = 1.0
+        if is_quantized_kv_cache(self.kv_cache_dtype):
+            self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float
+            self.bmm2_scale *= layer._k_scale_float
 
     def autotune_hisparse_decode(self, layer: AttentionLayer) -> None:
         """Autotune the largest legal HiSparse decode batch."""

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -280,15 +280,13 @@ def forward_mqa(
         else:
             q = q.view(num_decodes, -1, q.shape[-2], q.shape[-1])
 
-        if self.softmax_scale is None:
-            # FP8 KV cache is mandatory for this backend, so q_scale/k_scale
-            # always apply. softmax_scale is bmm1; output_scale is bmm2 — both
-            # required to recover the correct attention output from the FP8
-            # KV cache (V is stored as V_real/k_scale).
-            self.softmax_scale = (
-                self.scale * layer._q_scale_float * layer._k_scale_float
-            )
-            self.output_scale = layer._k_scale_float
+        # FP8 KV cache is mandatory for this backend, so q_scale/k_scale
+        # always apply. softmax_scale is bmm1; output_scale is bmm2 — both
+        # required to recover the correct attention output from the FP8
+        # KV cache (V is stored as V_real/k_scale). Recomputed every call: a
+        # weight reload can change the KV scales.
+        self.softmax_scale = self.scale * layer._q_scale_float * layer._k_scale_float
+        self.output_scale = layer._k_scale_float
 
         if self._workspace_buffer is None:
             # Parallelism can change the runtime query head count.