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 5, 19:45 UTC
[HMA]Fix corner case when hybrid page_size can not be evenly divided issue (blk_size=64,tp=4) (#37467)
Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Chendi.Xue <chendi.xue@intel.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Chendi.Xue committed
Mar 30, 2026 at 11:47 UTC
3b1dbaad4e59742ab5a5e4fd8f4ccc4f0e83c996
12 files changed
+214
-180
tests/v1/worker/test_gpu_model_runner.py
+2
@@ -851,6 +851,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
851
vllm_ctx = vllm_config.compilation_config.static_forward_context
852
853
runner = GPUModelRunner(vllm_config, DEVICE)
854
+ current_platform.update_block_size_for_backend(vllm_config)
855
kv_cache_spec = runner.get_kv_cache_spec()
856
857
available_memory = 5 * GiB_bytes
@@ -1306,6 +1307,7 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
1307
assert fwd_context is not None
1308
1309
runner = GPUModelRunner(vllm_config, DEVICE)
1310
+ current_platform.update_block_size_for_backend(vllm_config)
1311
kv_cache_spec = runner.get_kv_cache_spec()
1312
1313
available_memory = 5 * GiB_bytes
vllm/config/cache.py
+5
@@ -38,6 +38,8 @@ class CacheConfig:
38
Accepts None (meaning "use default"). After construction, always int."""
39
user_specified_block_size: bool = field(default=False, init=False)
40
"""Whether block_size was explicitly provided. Derived automatically."""
41
+ user_specified_mamba_block_size: bool = field(default=False, init=False)
42
+ """Whether mamba_block_size was explicitly provided. Derived automatically."""
43
gpu_memory_utilization: float = Field(default=0.9, gt=0, le=1)
44
"""The fraction of GPU memory to be used for the model executor, which can
45
range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory
@@ -182,6 +184,7 @@ class CacheConfig:
184
"cpu_kvcache_space_bytes",
185
"mamba_page_size_padded",
186
"user_specified_block_size",
187
+ "user_specified_mamba_block_size",
188
"_block_size_resolved",
189
# Post-init/derived counters
190
"num_gpu_blocks",
@@ -214,6 +217,8 @@ class CacheConfig:
217
object.__setattr__(self, "block_size", self.DEFAULT_BLOCK_SIZE)
218
else:
219
object.__setattr__(self, "user_specified_block_size", True)
220
+ if self.mamba_block_size is not None:
221
+ object.__setattr__(self, "user_specified_mamba_block_size", True)
222
return self
223
224
@field_validator("calculate_kv_scales", mode="after")
vllm/model_executor/models/config.py
+7
-142
@@ -1,15 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from copy import deepcopy
4
-from math import lcm
4
from typing import TYPE_CHECKING
5
6
from vllm.logger import init_logger
8
-from vllm.model_executor.models import ModelRegistry
9
-from vllm.utils.math_utils import cdiv, round_up
10
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
11
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
12
-from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, MLAAttentionSpec
7
+from vllm.utils.math_utils import round_up
8
9
if TYPE_CHECKING:
10
from vllm.config import ModelConfig, VllmConfig
@@ -104,11 +99,11 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
99
@classmethod
100
def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
101
"""
107
- Ensure that page size of attention layers is greater than or
108
- equal to the mamba layers. If not, automatically set the attention
109
- block size to ensure that it is. If the attention page size is
110
- strictly greater than the mamba page size, we pad the mamba page size
111
- to make them equal.
102
+ Perform early validation and setup for hybrid attention/mamba models.
103
+
104
+ Block size alignment with mamba page sizes is handled later by
105
+ Platform.update_block_size_for_backend(), which runs after model
106
+ layers are constructed and the attention backend is known.
107
108
Args:
109
vllm_config: vLLM Config
@@ -118,6 +113,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
113
# Disable calculate_kv_scales for hybrid models: uninitialized
114
# recurrent state corrupts scales during the calibration pass.
115
# See issue: https://github.com/vllm-project/vllm/issues/37554
116
+
117
if cache_config.calculate_kv_scales:
118
logger.warning(
119
"Disabling calculate_kv_scales for hybrid model '%s'. "
@@ -129,140 +125,9 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
125
)
126
cache_config.calculate_kv_scales = False
127
132
- # Save the user input before it gets modified by MambaModelConfig
133
- mamba_block_size = cache_config.mamba_block_size
128
# Enable FULL_AND_PIECEWISE by default
129
MambaModelConfig.verify_and_update_config(vllm_config)
130
137
- attention_config = vllm_config.attention_config
138
- cache_config = vllm_config.cache_config
139
- model_config = vllm_config.model_config
140
- parallel_config = vllm_config.parallel_config
141
-
142
- if cache_config.cache_dtype == "auto":
143
- kv_cache_dtype = model_config.dtype
144
- else:
145
- kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
146
-
147
- # get attention page size (for 1 token)
148
- # Attention backend constraints:
149
- # - FlashAttention (FA) requires block size to be multiple of 16
150
- # - MLA (Multi-head Latent Attention) requires larger alignment:
151
- # * CUTLASS_MLA backend: kernel_block_size 128 alignment
152
- # * Other MLA backends: kernel_block_size 64 alignment
153
- if model_config.use_mla:
154
- use_cutlass_mla = (
155
- attention_config.backend == AttentionBackendEnum.CUTLASS_MLA
156
- )
157
- kernel_block_alignment_size = 128 if use_cutlass_mla else 64
158
- attn_page_size_1_token = MLAAttentionSpec(
159
- block_size=1,
160
- num_kv_heads=model_config.get_num_kv_heads(parallel_config),
161
- head_size=model_config.get_head_size(),
162
- dtype=kv_cache_dtype,
163
- ).page_size_bytes
164
- else:
165
- kernel_block_alignment_size = 16
166
- attn_page_size_1_token = FullAttentionSpec(
167
- block_size=1,
168
- num_kv_heads=model_config.get_num_kv_heads(parallel_config),
169
- head_size=model_config.get_head_size(),
170
- dtype=kv_cache_dtype,
171
- ).page_size_bytes
172
-
173
- model_cls, _ = ModelRegistry.resolve_model_cls(
174
- model_config.architecture,
175
- model_config=model_config,
176
- )
177
-
178
- # get mamba page size
179
- mamba_page_size = MambaSpec(
180
- shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
181
- dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
182
- block_size=-1, # block_size doesn't matter for mamba page size
183
- ).page_size_bytes
184
-
185
- # Model may be marked as is_hybrid
186
- # but mamba is skipped via config,
187
- # return directly
188
- if mamba_page_size == 0:
189
- return
190
-
191
- if cache_config.mamba_cache_mode == "all":
192
- # With prefix caching, select attention block size to
193
- # optimize for mamba kernel performance
194
-
195
- # Mamba2 SSD kernel uses a chunk_size, e.g. 256
196
- # Align the block to the kernel: use lowest multiple of chunk_size
197
- # of attention tokens that would fit mamba_page_size:
198
- # e.g. for mamba page size = 788kB
199
- # attn_1_token = 2kB -> fits ~394 tokens
200
- # then round up to a multiple of 256 -> 512 tokens
201
- # End result:
202
- # attn_block_size = 512
203
- # mamba_block_size = 512 (aligned to a multiple of chunk_size)
204
- # TODO(tdoublep): this constraint can be relaxed fairly
205
- # easily by changing the way we layout chunks in the
206
- # mamba2 kernels.
207
-
208
- base_chunk_size = mamba_block_size or model_config.get_mamba_chunk_size()
209
- attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
210
- chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
211
- attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
212
- cache_config.mamba_block_size = attn_block_size
213
- else:
214
- # Without prefix caching, select minimum valid attention block size
215
- # to minimize mamba state padding
216
-
217
- # Calculate minimum attention block size that satisfies both:
218
- # 1. Backend alignment requirements (kernel_block_alignment_size)
219
- # 2. Mamba page size compatibility (attn_page_size >= mamba_page_size)
220
- attn_block_size = kernel_block_alignment_size * cdiv(
221
- mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token
222
- )
223
-
224
- # override attention block size if it is too small,
225
- # even if the user has explicitly set it
226
- if cache_config.block_size < attn_block_size:
227
- cache_config.block_size = attn_block_size
228
- logger.info(
229
- "Setting attention block size to %d tokens "
230
- "to ensure that attention page size is >= mamba page size.",
231
- attn_block_size,
232
- )
233
-
234
- # By default, mamba block size will be set to max_model_len.
235
- # When enabling prefix caching and using align mamba cache
236
- # mode, we align mamba block size to the block size as the
237
- # basic granularity for prefix caching.
238
- if cache_config.mamba_cache_mode == "align":
239
- cache_config.mamba_block_size = cache_config.block_size
240
-
241
- # compute new attention page size
242
- attn_page_size = cache_config.block_size * attn_page_size_1_token
243
-
244
- assert attn_page_size >= mamba_page_size
245
-
246
- if attn_page_size == mamba_page_size:
247
- # don't need to pad mamba page size
248
- return
249
-
250
- # pad mamba page size to exactly match attention
251
- if (
252
- cache_config.mamba_page_size_padded is None
253
- or cache_config.mamba_page_size_padded != attn_page_size
254
- ):
255
- cache_config.mamba_page_size_padded = attn_page_size
256
- mamba_padding_pct = (
257
- 100 * (attn_page_size - mamba_page_size) / mamba_page_size
258
- )
259
- logger.info(
260
- "Padding mamba page size by %.2f%% to ensure "
261
- "that mamba page size and attention page size are "
262
- "exactly equal.",
263
- mamba_padding_pct,
264
- )
265
-
131
132
class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
133
@staticmethod
vllm/platforms/interface.py
+169
-28
@@ -21,6 +21,7 @@ if TYPE_CHECKING:
21
from vllm.pooling_params import PoolingParams
22
from vllm.sampling_params import SamplingParams
23
from vllm.utils.argparse_utils import FlexibleArgumentParser
24
+ from vllm.v1.attention.backend import AttentionBackend
25
from vllm.v1.attention.selector import AttentionSelectorConfig
26
else:
27
FlexibleArgumentParser = object
@@ -423,55 +424,195 @@ class Platform:
424
"""
425
pass
426
427
+ @classmethod
428
+ def _find_non_ssm_backend(
429
+ cls, vllm_config: "VllmConfig"
430
+ ) -> "type[AttentionBackend] | None":
431
+ """Find the first non-SSM attention backend from model layers."""
432
+ from vllm.config.vllm import get_layers_from_vllm_config
433
+ from vllm.model_executor.layers.attention_layer_base import (
434
+ AttentionLayerBase,
435
+ )
436
+
437
+ attn_layers = get_layers_from_vllm_config(
438
+ vllm_config,
439
+ AttentionLayerBase, # type: ignore[type-abstract]
440
+ )
441
+ for layer in attn_layers.values():
442
+ b = layer.get_attn_backend()
443
+ if not b.is_ssm():
444
+ return b
445
+ return None
446
+
447
@classmethod
448
def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
449
"""
450
Ensure block_size is compatible with the attention backend.
451
+ For hybrid models, also aligns block_size with mamba page sizes.
452
"""
453
from vllm.config.cache import CacheConfig
454
+ from vllm.config.vllm import set_current_vllm_config
455
456
cache_config = vllm_config.cache_config
434
- if cache_config.user_specified_block_size:
435
- # User specified --block-size; keep it.
436
- return
437
-
457
model_config = vllm_config.model_config
458
+
459
# model_config may be None during testing.
440
- # Skip hybrid models — their block_size is managed by
441
- # HybridAttentionMambaModelConfig.
442
- if model_config is None or model_config.is_hybrid:
443
- cache_config.block_size = CacheConfig.DEFAULT_BLOCK_SIZE
460
+ if not model_config:
461
return
462
446
- from vllm.config.vllm import (
447
- get_layers_from_vllm_config,
448
- set_current_vllm_config,
449
- )
450
- from vllm.model_executor.layers.attention_layer_base import (
451
- AttentionLayerBase,
463
+ backend_cls = cls._find_non_ssm_backend(vllm_config)
464
+ if backend_cls is None:
465
+ return
466
+
467
+ # Phase 1: Pick block size from backend (skip if user set --block-size)
468
+ if not cache_config.user_specified_block_size:
469
+ with set_current_vllm_config(vllm_config):
470
+ preferred = backend_cls.get_preferred_block_size(
471
+ CacheConfig.DEFAULT_BLOCK_SIZE
472
+ )
473
+ if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
474
+ logger.info(
475
+ "Setting kv cache block size to %d for %s backend.",
476
+ preferred,
477
+ backend_cls.get_name(),
478
+ )
479
+ cache_config.block_size = preferred
480
+
481
+ # Phase 2: Align block/mamba sizes for hybrid models
482
+ # (may override user settings).
483
+ if model_config.is_hybrid:
484
+ cls._align_hybrid_block_size(vllm_config, backend_cls)
485
+
486
+ @classmethod
487
+ def _align_hybrid_block_size(
488
+ cls,
489
+ vllm_config: "VllmConfig",
490
+ backend_cls: "type[AttentionBackend]",
491
+ ) -> None:
492
+ """
493
+ For hybrid attention/mamba models, ensure that the attention page
494
+ size is >= the mamba page size, and pad the mamba page size to match.
495
+ """
496
+ from math import lcm
497
+
498
+ from vllm.config.vllm import set_current_vllm_config
499
+ from vllm.model_executor.models import ModelRegistry
500
+ from vllm.utils.math_utils import cdiv
501
+ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
502
+ from vllm.v1.attention.backend import MultipleOf
503
+ from vllm.v1.kv_cache_interface import (
504
+ FullAttentionSpec,
505
+ MambaSpec,
506
+ MLAAttentionSpec,
507
)
508
454
- attn_layers = get_layers_from_vllm_config(
455
- vllm_config,
456
- AttentionLayerBase, # type: ignore[type-abstract]
509
+ cache_config = vllm_config.cache_config
510
+ model_config = vllm_config.model_config
511
+ parallel_config = vllm_config.parallel_config
512
+
513
+ if cache_config.cache_dtype == "auto":
514
+ kv_cache_dtype = model_config.dtype
515
+ else:
516
+ kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
517
+
518
+ # Compute attention page size for 1 token
519
+ if model_config.use_mla:
520
+ attn_page_size_1_token = MLAAttentionSpec(
521
+ block_size=1,
522
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
523
+ head_size=model_config.get_head_size(),
524
+ dtype=kv_cache_dtype,
525
+ ).page_size_bytes
526
+ else:
527
+ attn_page_size_1_token = FullAttentionSpec(
528
+ block_size=1,
529
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
530
+ head_size=model_config.get_head_size(),
531
+ dtype=kv_cache_dtype,
532
+ ).page_size_bytes
533
+
534
+ # Compute mamba page size
535
+ model_cls, _ = ModelRegistry.resolve_model_cls(
536
+ model_config.architecture,
537
+ model_config=model_config,
538
)
458
- if not attn_layers:
459
- cache_config.block_size = CacheConfig.DEFAULT_BLOCK_SIZE
539
+ mamba_page_size = MambaSpec(
540
+ shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
541
+ dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
542
+ block_size=-1,
543
+ ).page_size_bytes
544
+
545
+ if mamba_page_size == 0:
546
return
547
462
- first_layer = next(iter(attn_layers.values()))
463
- backend_cls = first_layer.get_attn_backend()
548
+ # mamba_block_size here should either be user specified value or None
549
+ mamba_block_size = (
550
+ cache_config.mamba_block_size
551
+ if cache_config.user_specified_mamba_block_size
552
+ else None
553
+ )
554
+
555
+ # Get kernel block alignment from the backend's supported sizes
556
with set_current_vllm_config(vllm_config):
465
- preferred = backend_cls.get_preferred_block_size(
466
- CacheConfig.DEFAULT_BLOCK_SIZE
557
+ kernel_block_alignment_size = max(
558
+ min(
559
+ s.base if isinstance(s, MultipleOf) else s
560
+ for s in backend_cls.get_supported_kernel_block_sizes()
561
+ ),
562
+ cache_config.block_size,
563
+ )
564
+
565
+ if cache_config.mamba_cache_mode == "all":
566
+ # With prefix caching, align to mamba chunk size for kernel perf
567
+ # TODO(tdoublep): this constraint can be relaxed fairly
568
+ # easily by changing the way we layout chunks in the
569
+ # mamba2 kernels.
570
+ base_chunk_size = mamba_block_size or model_config.get_mamba_chunk_size()
571
+ assert base_chunk_size is not None
572
+ attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
573
+ chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
574
+ attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
575
+ cache_config.mamba_block_size = attn_block_size
576
+ else:
577
+ # Without prefix caching, use minimum block size that satisfies
578
+ # both backend alignment and mamba page size compatibility
579
+ attn_block_size = kernel_block_alignment_size * cdiv(
580
+ mamba_page_size,
581
+ kernel_block_alignment_size * attn_page_size_1_token,
582
+ )
583
+
584
+ if cache_config.block_size < attn_block_size:
585
+ cache_config.block_size = attn_block_size
586
+ logger.info(
587
+ "Setting attention block size to %d tokens "
588
+ "to ensure that attention page size is >= mamba page size.",
589
+ attn_block_size,
590
+ )
591
+
592
+ if cache_config.mamba_cache_mode == "align":
593
+ cache_config.mamba_block_size = cache_config.block_size
594
+
595
+ # Pad mamba page size to exactly match attention page size
596
+ attn_page_size = cache_config.block_size * attn_page_size_1_token
597
+ assert attn_page_size >= mamba_page_size
598
+
599
+ if attn_page_size == mamba_page_size:
600
+ return
601
+
602
+ if (
603
+ cache_config.mamba_page_size_padded is None
604
+ or cache_config.mamba_page_size_padded != attn_page_size
605
+ ):
606
+ cache_config.mamba_page_size_padded = attn_page_size
607
+ mamba_padding_pct = (
608
+ 100 * (attn_page_size - mamba_page_size) / mamba_page_size
609
)
468
- if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
610
logger.info(
470
- "Setting kv cache block size to %d for %s backend.",
471
- preferred,
472
- backend_cls.get_name(),
611
+ "Padding mamba page size by %.2f%% to ensure "
612
+ "that mamba page size and attention page size are "
613
+ "exactly equal.",
614
+ mamba_padding_pct,
615
)
474
- cache_config.block_size = preferred
616
617
@classmethod
618
def verify_model_arch(cls, model_arch: str) -> None:
vllm/platforms/xpu.py
-10
@@ -160,11 +160,7 @@ class XPUPlatform(Platform):
160
161
@classmethod
162
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
163
- cache_config = vllm_config.cache_config
163
parallel_config = vllm_config.parallel_config
165
- # in V1(or with chunked prefill) block_size is 64
166
- if cache_config and not cache_config.user_specified_block_size:
167
- cache_config.block_size = 64
164
165
# lazy import to avoid circular import
166
from vllm.config import CUDAGraphMode
@@ -221,12 +217,6 @@ class XPUPlatform(Platform):
217
# ref. https://openucx.readthedocs.io/en/master/faq.html
218
os.environ["UCX_MEMTYPE_CACHE"] = "n"
219
224
- @classmethod
225
- def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
226
- # TODO: XPU still sets block_size in check_and_update_config.
227
- # Move that logic here so block_size is chosen by the backend.
228
- pass
229
-
220
@classmethod
221
def support_hybrid_kv_cache(cls) -> bool:
222
return True
vllm/v1/attention/backend.py
+4
@@ -311,6 +311,10 @@ class AttentionBackend(ABC):
311
def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
312
return None
313
314
+ @classmethod
315
+ def is_ssm(cls) -> bool:
316
+ return False
317
+
318
319
class AttentionMetadata:
320
pass
vllm/v1/attention/backends/flash_attn.py
+7
@@ -43,6 +43,7 @@ from vllm.config import (
43
from vllm.config.cache import CacheDType
44
from vllm.distributed.parallel_state import get_dcp_group
45
from vllm.logger import init_logger
46
+from vllm.platforms import current_platform
47
from vllm.platforms.interface import DeviceCapability
48
from vllm.utils.math_utils import cdiv, round_up
49
from vllm.v1.attention.backend import (
@@ -90,6 +91,12 @@ class FlashAttentionBackend(AttentionBackend):
91
92
forward_includes_kv_cache_update: bool = False
93
94
+ @classmethod
95
+ def get_preferred_block_size(cls, default_block_size: int) -> int:
96
+ if current_platform.is_xpu():
97
+ return max(default_block_size, 64)
98
+ return super().get_preferred_block_size(default_block_size)
99
+
100
@staticmethod
101
def get_name() -> str:
102
return "FLASH_ATTN"
vllm/v1/attention/backends/gdn_attn.py
+4
@@ -31,6 +31,10 @@ class GDNAttentionBackend(AttentionBackend):
31
def get_builder_cls() -> type["GDNAttentionMetadataBuilder"]:
32
return GDNAttentionMetadataBuilder
33
34
+ @classmethod
35
+ def is_ssm(cls) -> bool:
36
+ return True
37
+
38
39
@dataclass
40
class GDNAttentionMetadata:
vllm/v1/attention/backends/linear_attn.py
+4
@@ -27,6 +27,10 @@ class LinearAttentionBackend(AttentionBackend):
27
def get_builder_cls() -> type["LinearAttentionMetadataBuilder"]:
28
return LinearAttentionMetadataBuilder
29
30
+ @classmethod
31
+ def is_ssm(cls) -> bool:
32
+ return True
33
+
34
35
@dataclass
36
class LinearAttentionMetadata:
vllm/v1/attention/backends/mamba1_attn.py
+4
@@ -20,6 +20,10 @@ class Mamba1AttentionBackend(AttentionBackend):
20
def get_builder_cls() -> type["Mamba1AttentionMetadataBuilder"]:
21
return Mamba1AttentionMetadataBuilder
22
23
+ @classmethod
24
+ def is_ssm(cls) -> bool:
25
+ return True
26
+
27
28
@dataclass
29
class Mamba1AttentionMetadata(BaseMambaAttentionMetadata):
vllm/v1/attention/backends/mamba2_attn.py
+4
@@ -96,6 +96,10 @@ class Mamba2AttentionBackend(AttentionBackend):
96
def get_builder_cls() -> type["Mamba2AttentionMetadataBuilder"]:
97
return Mamba2AttentionMetadataBuilder
98
99
+ @classmethod
100
+ def is_ssm(cls) -> bool:
101
+ return True
102
+
103
104
@dataclass
105
class Mamba2AttentionMetadata(BaseMambaAttentionMetadata):
vllm/v1/attention/backends/short_conv_attn.py
+4
@@ -18,6 +18,10 @@ class ShortConvAttentionBackend(AttentionBackend):
18
def get_builder_cls() -> type["ShortConvAttentionMetadataBuilder"]:
19
return ShortConvAttentionMetadataBuilder
20
21
+ @classmethod
22
+ def is_ssm(cls) -> bool:
23
+ return True
24
+
25
26
@dataclass
27
class ShortConvAttentionMetadata(BaseMambaAttentionMetadata):