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 6, 23:00 UTC
[Bugfix][Sparse MLA] report indexer CG support properly (#36519)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Mar 10, 2026 at 12:14 UTC
9095cbbfb6f68f3f7abc7f55c74768e9f7b1d0a7
1 file changed
+15
-3
vllm/v1/attention/backends/mla/indexer.py
+15
-3
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
-from typing import ClassVar
4
5
import torch
6
@@ -25,6 +24,7 @@ from vllm.v1.attention.backends.utils import (
24
split_decodes_and_prefills,
25
split_prefill_chunks,
26
)
27
+from vllm.v1.kv_cache_interface import AttentionSpec
28
from vllm.v1.worker.cp_utils import get_total_cp_world_size
29
30
logger = init_logger(__name__)
@@ -202,10 +202,22 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
202
203
204
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
205
- _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
206
-
205
reorder_batch_threshold: int = 1
206
207
+ @classmethod
208
+ def get_cudagraph_support(
209
+ cls,
210
+ vllm_config: VllmConfig,
211
+ kv_cache_spec: AttentionSpec,
212
+ ) -> AttentionCGSupport:
213
+ if not is_deep_gemm_supported():
214
+ logger.warning_once(
215
+ "DeepGEMM is not available. Disabling CUDA graph support "
216
+ "for sparse attention indexer. This may reduce performance.",
217
+ )
218
+ return AttentionCGSupport.NEVER
219
+ return AttentionCGSupport.UNIFORM_BATCH
220
+
221
def __init__(self, *args, **kwargs):
222
super().__init__(*args, **kwargs)
223
scheduler_config = self.vllm_config.scheduler_config