A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 13:15 UTC
[Bugfix] Ensure calculated KV scales are applied in attention. (#27232)
Signed-off-by: adabeyta <aabeyta@redhat.com>
Adrian Abeyta committed
Nov 10, 2025 at 17:42 UTC
a5a790eea6035760c71eae1861c1e5f369bc6d08
4 files changed
+29
-36
.buildkite/test-pipeline.yaml
+5
-2
@@ -471,8 +471,8 @@ steps:
471
- vllm/
472
- tests/compile
473
commands:
474
- - pytest -v -s compile/test_full_graph.py
475
- # Limit to no custom ops to reduce running time
474
+ - pytest -v -s compile/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
475
+ # Limit to no custom ops to reduce running time
476
# Wrap with quotes to escape yaml and avoid starting -k string with a -
477
- "pytest -v -s compile/test_fusions_e2e.py -k 'TRITON and -quant_fp8'"
478
@@ -951,10 +951,13 @@ steps:
951
- vllm/model_executor/layers/activation.py
952
- vllm/model_executor/layers/quantization/input_quant_fp8.py
953
- tests/compile/test_fusions_e2e.py
954
+ - tests/compile/test_full_graph.py
955
commands:
956
- nvidia-smi
957
# Run all e2e fusion tests
958
- pytest -v -s tests/compile/test_fusions_e2e.py
959
+ # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
960
+ - pytest -v -s tests/compile/test_full_graph.py::test_fp8_kv_scale_compile
961
962
- label: Blackwell GPT-OSS Eval
963
timeout_in_minutes: 60
tests/compile/test_full_graph.py
+8
-2
@@ -183,8 +183,14 @@ def test_custom_compile_config(
183
"compilation_mode",
184
[CompilationMode.NONE, CompilationMode.VLLM_COMPILE],
185
)
186
-def test_fp8_kv_scale_compile(compilation_mode: int):
187
- model = "Qwen/Qwen2-0.5B"
186
+@pytest.mark.parametrize(
187
+ "model",
188
+ [
189
+ "Qwen/Qwen2-0.5B", # Standard attention model
190
+ "deepseek-ai/DeepSeek-V2-Lite", # MLA (Multi-head Latent Attention) model
191
+ ],
192
+)
193
+def test_fp8_kv_scale_compile(compilation_mode: int, model: str):
194
model_kwargs = {
195
"quantization": "fp8",
196
"kv_cache_dtype": "fp8_e4m3",
vllm/attention/layer.py
+7
-22
@@ -745,6 +745,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
745
k_pe: torch.Tensor,
746
output_shape: torch.Size | None = None,
747
) -> torch.Tensor:
748
+ if self.calculate_kv_scales:
749
+ torch.ops.vllm.maybe_calc_kv_scales(q, kv_c_normed, k_pe, self.layer_name)
750
+
751
if self.use_direct_call:
752
forward_context: ForwardContext = get_forward_context()
753
attn_metadata = forward_context.attn_metadata
@@ -752,12 +755,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
755
attn_metadata = attn_metadata[self.layer_name]
756
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
757
755
- # Mirror Attention.forward scale calculation path
756
- if self.calculate_kv_scales and getattr(
757
- attn_metadata, "enable_kv_scales_calculation", False
758
- ):
759
- self.calc_kv_scales(q, kv_c_normed, k_pe)
760
-
758
if self.attn_backend.accept_output_buffer:
759
output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
760
self.impl.forward(
@@ -786,14 +783,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
783
)
784
return output
785
else:
789
- # We can still access forward context to check calculation flag
790
- if self.calculate_kv_scales:
791
- forward_context = get_forward_context()
792
- attn_metadata = forward_context.attn_metadata
793
- if isinstance(attn_metadata, dict):
794
- attn_metadata = attn_metadata[self.layer_name]
795
- if getattr(attn_metadata, "enable_kv_scales_calculation", False):
796
- self.calc_kv_scales(q, kv_c_normed, k_pe)
786
return torch.ops.vllm.unified_mla_attention(
787
q,
788
kv_c_normed,
@@ -881,17 +870,13 @@ def maybe_calc_kv_scales(
870
layer_name: str,
871
) -> None:
872
forward_context: ForwardContext = get_forward_context()
884
- attn_metadata = forward_context.attn_metadata
885
-
886
- if isinstance(attn_metadata, dict):
887
- attn_metadata = attn_metadata[layer_name]
873
+ self = forward_context.no_compile_layers[layer_name]
874
889
- if attn_metadata is None or not getattr(
890
- attn_metadata, "enable_kv_scales_calculation", False
891
- ):
875
+ # Only calculate if the layer's calculate_kv_scales flag is True
876
+ # This flag gets set to False after the first forward pass
877
+ if not self.calculate_kv_scales:
878
return
879
894
- self = forward_context.no_compile_layers[layer_name]
880
self.calc_kv_scales(query, key, value)
881
882
vllm/v1/worker/gpu_model_runner.py
+9
-10
@@ -279,6 +279,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
279
# This will be overridden in load_model()
280
self.is_multimodal_pruning_enabled = False
281
self.max_model_len = model_config.max_model_len
282
+
283
+ # Always set to false after the first forward pass
284
+ self.calculate_kv_scales = self.cache_config.calculate_kv_scales
285
self.dcp_world_size = self.parallel_config.decode_context_parallel_size
286
self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group
287
self.max_num_tokens = scheduler_config.max_num_batched_tokens
@@ -2625,16 +2628,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2628
)
2629
2630
# Set cudagraph mode to none if calc_kv_scales is true.
2628
- if attn_metadata is not None:
2629
- metadata_list = (
2630
- attn_metadata.values()
2631
- if isinstance(attn_metadata, dict)
2632
- else [attn_metadata]
2633
- )
2634
- if any(
2635
- getattr(m, "enable_kv_scales_calculation", False) for m in metadata_list
2636
- ):
2637
- cudagraph_runtime_mode = CUDAGraphMode.NONE
2631
+ # KV scales calculation involves dynamic operations that are incompatible
2632
+ # with CUDA graph capture.
2633
+ if self.calculate_kv_scales:
2634
+ cudagraph_runtime_mode = CUDAGraphMode.NONE
2635
+ # Mark KV scales as calculated after the first forward pass
2636
+ self.calculate_kv_scales = False
2637
2638
# Run the model.
2639
# Use persistent buffers for CUDA graphs.