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, 04:30 UTC
[Bugfix] Fix ChunkedLocalAttention CUDA Graph setting (#28739)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Nov 14, 2025 at 17:14 UTC
bf3ffb61e61525cce5fdec8a249f8114a0c0bfcc
1 file changed
+16
-3
vllm/attention/layers/chunked_local_attention.py
+16
-3
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import functools
4
-from typing import ClassVar
4
5
import torch
6
@@ -12,11 +11,16 @@ from vllm.config.vllm import VllmConfig
11
from vllm.model_executor.layers.quantization import QuantizationConfig
12
from vllm.v1.attention.backends.utils import (
13
AttentionCGSupport,
14
+ AttentionMetadataBuilder,
15
CommonAttentionMetadata,
16
make_local_attention_virtual_batches,
17
subclass_attention_backend,
18
)
19
-from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, KVCacheSpec
19
+from vllm.v1.kv_cache_interface import (
20
+ AttentionSpec,
21
+ ChunkedLocalAttentionSpec,
22
+ KVCacheSpec,
23
+)
24
25
from ..layer import Attention
26
@@ -30,9 +34,18 @@ def create_chunked_local_attention_backend(
34
prefix = f"ChunkedLocalAttention_{attention_chunk_size}_{block_size}_"
35
36
underlying_builder = underlying_attn_backend.get_builder_cls()
37
+ assert issubclass(underlying_builder, AttentionMetadataBuilder)
38
39
class ChunkedLocalAttentionBuilder(underlying_builder): # type: ignore
35
- _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
40
+ @classmethod
41
+ def get_cudagraph_support(
42
+ cls: type["AttentionMetadataBuilder"],
43
+ vllm_config: VllmConfig,
44
+ kv_cache_spec: AttentionSpec,
45
+ ) -> AttentionCGSupport:
46
+ # Explicit override in case the underlying builder specialized this getter.
47
+ # @override omitted only because of mypy limitation due to type variable.
48
+ return AttentionCGSupport.NEVER
49
50
def build(
51
self,