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 15, 02:15 UTC
[V1] [Hybrid] Refactor mamba state shape calculation; enable V1 via cli (#20840)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Jul 15, 2025 at 13:04 UTC
3534c39a2037beb3cf000a2c2f25c58fc0b4b6f4
14 files changed
+441
-353
docs/usage/v1_guide.md
+1
-2
@@ -112,8 +112,7 @@ enforcing eager mode and disabling prefix caching in V1.
112
Models that combine Mamba-2 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
113
`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`). Please note that
114
these models currently require enforcing eager mode, disabling prefix caching, and using the FlashInfer attention
115
-backend in V1. It is also necessary to pass a non-standard block size for attention layers (this is not possible
116
-using the `vllm serve` CLI yet).
115
+backend in V1.
116
117
#### Encoder-Decoder Models
118
tests/models/language/generation/test_hybrid.py
+1
-15
@@ -61,14 +61,6 @@ V1_SUPPORTED_MODELS = [
61
"tiiuae/Falcon-H1-0.5B-Base",
62
]
63
64
-ATTN_BLOCK_SIZES = {
65
- "ibm-ai-platform/Bamba-9B-v1": 528,
66
- "Zyphra/Zamba2-1.2B-instruct": 80,
67
- "nvidia/Nemotron-H-8B-Base-8K": 528,
68
- "ibm-granite/granite-4.0-tiny-preview": 400,
69
- "tiiuae/Falcon-H1-0.5B-Base": 800,
70
-}
71
-
64
# Avoid OOM
65
MAX_NUM_SEQS = 4
66
@@ -105,11 +97,6 @@ def test_models(
97
example_prompts, max_tokens, num_logprobs)
98
99
if model in V1_SUPPORTED_MODELS:
108
- if model in HYBRID_MODELS and model in ATTN_BLOCK_SIZES:
109
- block_size = ATTN_BLOCK_SIZES[model]
110
- else:
111
- block_size = 16
112
-
100
with monkeypatch.context() as m:
101
m.setenv("VLLM_USE_V1", "1")
102
if model in HYBRID_MODELS:
@@ -118,8 +105,7 @@ def test_models(
105
with vllm_runner(model,
106
max_num_seqs=MAX_NUM_SEQS,
107
enforce_eager=True,
121
- enable_prefix_caching=False,
122
- block_size=block_size) as vllm_model:
108
+ enable_prefix_caching=False) as vllm_model:
109
vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
110
example_prompts, max_tokens, num_logprobs)
111
else:
vllm/config.py
+8
-1
@@ -1630,6 +1630,9 @@ class CacheConfig:
1630
checkpoint if available. Otherwise, the scales will default to 1.0."""
1631
cpu_kvcache_space_bytes: Optional[int] = None
1632
"""(CPU backend only) CPU key-value cache space."""
1633
+ mamba_page_size_padded: Optional[int] = None
1634
+ """ Optional override for mamba page size; used by hybrid mamba/attention
1635
+ models to ensure exact alignment with attention page size."""
1636
1637
# Will be set after profiling.
1638
num_gpu_blocks: Optional[int] = field(default=None, init=False)
@@ -4882,11 +4885,15 @@ class VllmConfig:
4885
if architecture is None:
4886
return
4887
4885
- from vllm.model_executor.models.config import MODELS_CONFIG_MAP
4888
+ from vllm.model_executor.models.config import (
4889
+ MODELS_CONFIG_MAP, HybridAttentionMambaModelConfig)
4890
cls = MODELS_CONFIG_MAP.get(architecture, None)
4891
if cls is not None:
4892
cls.verify_and_update_config(self)
4893
4894
+ if self.model_config.is_hybrid:
4895
+ HybridAttentionMambaModelConfig.verify_and_update_config(self)
4896
+
4897
if self.model_config.task == "classify":
4898
# Maybe convert ForCausalLM into ForSequenceClassification model.
4899
from vllm.model_executor.models.adapters import (
vllm/model_executor/layers/mamba/mamba_mixer2.py
+10
-38
@@ -20,6 +20,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
20
from vllm.model_executor.layers.mamba.abstract import MambaBase
21
from vllm.model_executor.layers.mamba.mamba2_metadata import (Mamba2Metadata,
22
update_metadata)
23
+from vllm.model_executor.layers.mamba.mamba_utils import (
24
+ extra_groups_for_head_shards, get_mamba_state_shape)
25
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
26
causal_conv1d_fn, causal_conv1d_update)
27
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
@@ -146,18 +148,6 @@ class Mixer2RMSNormGated(CustomOp):
148
return out
149
150
149
-def extra_groups_for_head_shards(ngroups: int, tp_size: int):
150
- """Compute the increase in group numbers to account for
151
- replication in order to accompany the head shards."""
152
-
153
- # in the case ngoups % tp_size == 0, this will be zero
154
- if ngroups % tp_size == 0:
155
- return 0
156
-
157
- # for n_groups == 1, this is exactly tp_size - n_groups
158
- return tp_size - ngroups
159
-
160
-
151
def mamba_v2_sharded_weight_loader(
152
shard_spec: list[tuple[int, int, float]],
153
tp_size: int,
@@ -707,30 +697,12 @@ class MambaMixer2(MambaBase, CustomOp):
697
return out
698
699
def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
710
- world_size = get_tensor_model_parallel_world_size()
711
-
712
- conv_state_shape, temporal_state_shape = None, None
713
-
714
- # if n_groups is not divisible by world_size, need to extend the shards
715
- # to ensure all groups needed by a head is sharded along with it
716
- n_groups = (self.n_groups +
717
- extra_groups_for_head_shards(self.n_groups, world_size))
718
-
719
- # - heads and n_groups are TP-ed
720
- conv_dim = (self.intermediate_size +
721
- 2 * n_groups * self.ssm_state_size)
722
- # contiguous along 'dim' axis
723
- conv_state_shape = (
724
- self.conv_kernel_size - 1,
725
- divide(conv_dim, world_size),
726
- )
727
-
728
- # These are not TP-ed as they depend on A, dt_bias, D
729
- # - they are typically small
730
- # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
731
- temporal_state_shape = (
732
- divide(self.num_heads, world_size),
733
- self.head_dim,
734
- self.ssm_state_size,
700
+ return get_mamba_state_shape(
701
+ intermediate_size=self.intermediate_size,
702
+ tp_world_size=get_tensor_model_parallel_world_size(),
703
+ n_groups=self.n_groups,
704
+ num_heads=self.num_heads,
705
+ head_dim=self.head_dim,
706
+ state_size=self.ssm_state_size,
707
+ conv_kernel=self.conv_kernel_size,
708
)
736
- return conv_state_shape, temporal_state_shape
vllm/model_executor/layers/mamba/mamba_utils.py
new
+55
@@ -0,0 +1,55 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from vllm.distributed import divide
4
+
5
+
6
+def extra_groups_for_head_shards(ngroups: int, tp_size: int):
7
+ """Compute the increase in group numbers to account for
8
+ replication in order to accompany the head shards."""
9
+
10
+ # in the case ngoups % tp_size == 0, this will be zero
11
+ if ngroups % tp_size == 0:
12
+ return 0
13
+
14
+ # for n_groups == 1, this is exactly tp_size - n_groups
15
+ return tp_size - ngroups
16
+
17
+
18
+def get_mamba_state_shape(
19
+ intermediate_size: int,
20
+ tp_world_size: int,
21
+ n_groups: int,
22
+ num_heads: int,
23
+ head_dim: int,
24
+ state_size: int,
25
+ conv_kernel: int,
26
+ use_v1: bool = True,
27
+) -> tuple[tuple[int, int], tuple[int, int, int]]:
28
+ """ Get the shape of mamba state."""
29
+
30
+ # if n_groups is not divisible by world_size, need to extend the shards
31
+ # to ensure all groups needed by a head is sharded along with it
32
+ n_groups = (n_groups +
33
+ extra_groups_for_head_shards(n_groups, tp_world_size))
34
+
35
+ # - heads and n_groups are TP-ed
36
+ conv_dim = (intermediate_size + 2 * n_groups * state_size)
37
+ # contiguous along 'dim' axis
38
+ conv_state_shape = (
39
+ conv_kernel - 1,
40
+ divide(conv_dim, tp_world_size),
41
+ )
42
+
43
+ if not use_v1:
44
+ conv_state_shape = (conv_state_shape[1], conv_state_shape[0])
45
+
46
+ # These are not TP-ed as they depend on A, dt_bias, D
47
+ # - they are typically small
48
+ # e.g., (h_heads, head_dim, state_size) = (128, 64, 128)
49
+ temporal_state_shape = (
50
+ divide(num_heads, tp_world_size),
51
+ head_dim,
52
+ state_size,
53
+ )
54
+
55
+ return conv_state_shape, temporal_state_shape
vllm/model_executor/models/bamba.py
+42
-39
@@ -12,7 +12,7 @@ from transformers import BambaConfig
12
from vllm import envs
13
from vllm.attention.layer import Attention
14
from vllm.config import CacheConfig, VllmConfig
15
-from vllm.distributed import divide, get_tensor_model_parallel_world_size
15
+from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.distributed.parallel_state import get_pp_group
17
from vllm.forward_context import get_forward_context
18
from vllm.model_executor.layers.activation import SiluAndMul
@@ -23,8 +23,8 @@ from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
23
from vllm.model_executor.layers.logits_processor import LogitsProcessor
24
from vllm.model_executor.layers.mamba.mamba2_metadata import (
25
Mamba2Metadata, prepare_mamba2_metadata)
26
-from vllm.model_executor.layers.mamba.mamba_mixer2 import (
27
- MambaMixer2, extra_groups_for_head_shards)
26
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
27
+from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
28
from vllm.model_executor.layers.quantization import QuantizationConfig
29
from vllm.model_executor.layers.rotary_embedding import get_rope
30
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -435,6 +435,38 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
435
}
436
embedding_padding_modules = ["lm_head"]
437
438
+ @classmethod
439
+ def get_mamba_state_shape_from_config(
440
+ cls,
441
+ vllm_config: "VllmConfig",
442
+ use_v1: bool = True,
443
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
444
+ """Calculate shapes for Mamba's convolutional and state caches.
445
+
446
+ Args:
447
+ vllm_config: vLLM config
448
+ use_v1: Get shapes for V1 (or V0)
449
+
450
+ Returns:
451
+ Tuple containing:
452
+ - conv_state_shape: Shape for convolutional state cache
453
+ - temporal_state_shape: Shape for state space model cache
454
+ """
455
+ parallel_config = vllm_config.parallel_config
456
+ hf_config = vllm_config.model_config.hf_config
457
+ intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
458
+
459
+ return get_mamba_state_shape(
460
+ intermediate_size=intermediate_size,
461
+ tp_world_size=parallel_config.tensor_parallel_size,
462
+ n_groups=hf_config.mamba_n_groups,
463
+ num_heads=hf_config.mamba_n_heads,
464
+ head_dim=hf_config.mamba_d_head,
465
+ state_size=hf_config.mamba_d_state,
466
+ conv_kernel=hf_config.mamba_d_conv,
467
+ use_v1=use_v1,
468
+ )
469
+
470
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
471
config = vllm_config.model_config.hf_config
472
self.vllm_config = vllm_config
@@ -491,10 +523,13 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
523
self.vllm_config.parallel_config,
524
LayerBlockType.mamba
525
)
494
-
495
- self.mamba_cache = MambaCacheManager(
496
- self.vllm_config, self.lm_head.weight.dtype,
497
- num_mamba_layers, *self._get_mamba_cache_shape())
526
+ mamba_state_shape = \
527
+ self.get_mamba_state_shape_from_config(
528
+ self.vllm_config, use_v1=False)
529
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
530
+ self.lm_head.weight.dtype,
531
+ num_mamba_layers,
532
+ *mamba_state_shape)
533
534
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
535
@@ -510,38 +545,6 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
545
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
546
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
547
513
- def _get_mamba_cache_shape(
514
- self) -> tuple[tuple[int, int], tuple[int, int]]:
515
- world_size = get_tensor_model_parallel_world_size()
516
- hidden_size = self.config.hidden_size
517
-
518
- conv_state_shape, temporal_state_shape = None, None
519
-
520
- intermediate_size = self.config.mamba_expand * hidden_size
521
-
522
- # if n_groups is not divisible by world_size, need to extend the shards
523
- # to ensure all groups needed by a head is sharded along with it
524
- n_groups = (self.config.mamba_n_groups + extra_groups_for_head_shards(
525
- self.config.mamba_n_groups, world_size))
526
-
527
- # - heads and n_groups are TP-ed
528
- conv_dim = (intermediate_size +
529
- 2 * n_groups * self.config.mamba_d_state)
530
- conv_state_shape = (
531
- divide(conv_dim, world_size),
532
- self.config.mamba_d_conv - 1,
533
- )
534
-
535
- # These are not TP-ed as they depend on A, dt_bias, D
536
- # - they are typically small
537
- # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
538
- temporal_state_shape = (
539
- divide(self.config.mamba_n_heads, world_size),
540
- self.config.mamba_d_head,
541
- self.config.mamba_d_state,
542
- )
543
- return conv_state_shape, temporal_state_shape
544
-
548
def compute_logits(
549
self,
550
hidden_states: torch.Tensor,
vllm/model_executor/models/config.py
+90
@@ -3,9 +3,14 @@
3
from copy import deepcopy
4
from typing import TYPE_CHECKING
5
6
+import vllm.envs as envs
7
from vllm.logger import init_logger
8
+from vllm.model_executor.models import ModelRegistry
9
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
10
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
11
12
if TYPE_CHECKING:
13
+
14
from vllm.config import VllmConfig
15
16
logger = init_logger(__name__)
@@ -200,6 +205,91 @@ class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
205
}
206
207
208
+class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
209
+
210
+ @classmethod
211
+ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
212
+ """
213
+ Ensure that page size of attention layers is greater than or
214
+ equal to the mamba layers. If not, automatically set the attention
215
+ block size to ensure that it is. If the attention page size is
216
+ strictly greater than the mamba page size, we pad the mamba page size
217
+ to make them equal.
218
+
219
+ Args:
220
+ vllm_config: vLLM Config
221
+ """
222
+
223
+ if not envs.VLLM_USE_V1:
224
+ return
225
+
226
+ cache_config = vllm_config.cache_config
227
+ model_config = vllm_config.model_config
228
+ parallel_config = vllm_config.parallel_config
229
+
230
+ if cache_config.cache_dtype == "auto":
231
+ kv_cache_dtype = model_config.dtype
232
+ else:
233
+ kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
234
+
235
+ # get attention page size (for 1 token)
236
+ attn_page_size_1_token = FullAttentionSpec(
237
+ block_size=1,
238
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
239
+ head_size=model_config.get_head_size(),
240
+ dtype=kv_cache_dtype,
241
+ use_mla=model_config.use_mla).page_size_bytes
242
+
243
+ model_cls = ModelRegistry.resolve_model_cls(
244
+ model_config._model_info.architecture)[0]
245
+
246
+ # get mamba page size
247
+ mamba_page_size = MambaSpec(
248
+ shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
249
+ dtype=kv_cache_dtype,
250
+ block_size=model_config.max_model_len,
251
+ ).page_size_bytes
252
+
253
+ # some attention backends (e.g. FA) only support setting
254
+ # block size to multiple of 16, so let's suggest a value
255
+ # that would work (note: FA is currently not compatible
256
+ # with mamba layers, use FlashInfer instead).
257
+ attn_block_size = 16 * cdiv(mamba_page_size,
258
+ 16 * attn_page_size_1_token)
259
+
260
+ # override attention block size if either (a) the
261
+ # user has not set it or (b) the user has set it
262
+ # too small.
263
+ if (cache_config.block_size is None
264
+ or cache_config.block_size < attn_block_size):
265
+ cache_config.block_size = attn_block_size
266
+ logger.info(
267
+ "Setting attention block size to %d tokens "
268
+ "to ensure that attention page size is >= mamba page size.",
269
+ attn_block_size)
270
+
271
+ # compute new attention page size
272
+ attn_page_size = \
273
+ cache_config.block_size * attn_page_size_1_token
274
+
275
+ assert attn_page_size >= mamba_page_size
276
+
277
+ if attn_page_size == mamba_page_size:
278
+ # don't need to pad mamba page size
279
+ return
280
+
281
+ # pad mamba page size to exactly match attention
282
+ if (cache_config.mamba_page_size_padded is None
283
+ or cache_config.mamba_page_size_padded != attn_page_size):
284
+ cache_config.mamba_page_size_padded = (attn_page_size)
285
+ mamba_padding_pct = 100 * (attn_page_size -
286
+ mamba_page_size) / mamba_page_size
287
+ logger.info(
288
+ "Padding mamba page size by %.2f%% to ensure "
289
+ "that mamba page size and attention page size are "
290
+ "exactly equal.", mamba_padding_pct)
291
+
292
+
293
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
294
"GteModel": SnowflakeGteNewModelConfig,
295
"GteNewModel": GteNewModelConfig,
vllm/model_executor/models/falcon_h1.py
+43
-37
@@ -11,7 +11,7 @@ from transformers import FalconH1Config
11
from vllm import envs
12
from vllm.attention.layer import Attention
13
from vllm.config import CacheConfig, VllmConfig
14
-from vllm.distributed import divide, get_tensor_model_parallel_world_size
14
+from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.forward_context import get_forward_context
17
from vllm.model_executor.layers.activation import SiluAndMul
@@ -22,8 +22,8 @@ from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
22
from vllm.model_executor.layers.logits_processor import LogitsProcessor
23
from vllm.model_executor.layers.mamba.mamba2_metadata import (
24
Mamba2Metadata, prepare_mamba2_metadata)
25
-from vllm.model_executor.layers.mamba.mamba_mixer2 import (
26
- MambaMixer2, extra_groups_for_head_shards)
25
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
26
+from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
27
from vllm.model_executor.layers.quantization import QuantizationConfig
28
from vllm.model_executor.layers.rotary_embedding import get_rope
29
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -514,6 +514,42 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
514
}
515
embedding_padding_modules = ["lm_head"]
516
517
+ @classmethod
518
+ def get_mamba_state_shape_from_config(
519
+ cls,
520
+ vllm_config: "VllmConfig",
521
+ use_v1: bool = True,
522
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
523
+ """Calculate shapes for Mamba's convolutional and state caches.
524
+
525
+ Args:
526
+ vllm_config: vLLM config
527
+ use_v1: Get shapes for V1 (or V0)
528
+
529
+ Returns:
530
+ Tuple containing:
531
+ - conv_state_shape: Shape for convolutional state cache
532
+ - temporal_state_shape: Shape for state space model cache
533
+ """
534
+ parallel_config = vllm_config.parallel_config
535
+ hf_config = vllm_config.model_config.hf_config
536
+
537
+ intermediate_size = (int(hf_config.mamba_expand *
538
+ hf_config.hidden_size)
539
+ if hf_config.mamba_d_ssm is None else
540
+ hf_config.mamba_d_ssm)
541
+
542
+ return get_mamba_state_shape(
543
+ intermediate_size=intermediate_size,
544
+ tp_world_size=parallel_config.tensor_parallel_size,
545
+ n_groups=hf_config.mamba_n_groups,
546
+ num_heads=hf_config.mamba_n_heads,
547
+ head_dim=hf_config.mamba_d_head,
548
+ state_size=hf_config.mamba_d_state,
549
+ conv_kernel=hf_config.mamba_d_conv,
550
+ use_v1=use_v1,
551
+ )
552
+
553
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
554
config = vllm_config.model_config.hf_config
555
self.vllm_config = vllm_config
@@ -580,12 +616,15 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
616
mamba_cache_params = None
617
if not envs.VLLM_USE_V1:
618
if self.mamba_cache is None:
619
+ mamba_state_shape = \
620
+ self.get_mamba_state_shape_from_config(
621
+ self.vllm_config, use_v1=False)
622
self.mamba_cache = MambaCacheManager(
623
self.vllm_config,
624
self.lm_head.weight.dtype if hasattr(
625
self.lm_head, 'weight') else torch.bfloat16,
626
self.config.num_hidden_layers,
588
- *self._get_mamba_cache_shape(),
627
+ *mamba_state_shape,
628
)
629
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
630
@@ -606,39 +645,6 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
645
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
646
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
647
609
- def _get_mamba_cache_shape(
610
- self) -> tuple[tuple[int, int], tuple[int, int]]:
611
- world_size = get_tensor_model_parallel_world_size()
612
- hidden_size = self.config.hidden_size
613
-
614
- conv_state_shape, temporal_state_shape = None, None
615
-
616
- intermediate_size = (int(self.config.mamba_expand *
617
- hidden_size) if self.config.mamba_d_ssm
618
- is None else self.config.mamba_d_ssm)
619
-
620
- # if n_groups is not divisible by world_size, need to extend the shards
621
- # to ensure all groups needed by a head is sharded along with it
622
- n_groups = self.config.mamba_n_groups + extra_groups_for_head_shards(
623
- self.config.mamba_n_groups, world_size)
624
-
625
- # - heads and n_groups are TP-ed
626
- conv_dim = intermediate_size + 2 * n_groups * self.config.mamba_d_state
627
- conv_state_shape = (
628
- divide(conv_dim, world_size),
629
- self.config.mamba_d_conv - 1,
630
- )
631
-
632
- # These are not TP-ed as they depend on A, dt_bias, D
633
- # - they are typically small
634
- # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
635
- temporal_state_shape = (
636
- divide(self.config.mamba_n_heads, world_size),
637
- self.config.mamba_d_head,
638
- self.config.mamba_d_state,
639
- )
640
- return conv_state_shape, temporal_state_shape
641
-
648
def compute_logits(
649
self,
650
hidden_states: torch.Tensor,
vllm/model_executor/models/granitemoehybrid.py
+42
-38
@@ -12,7 +12,7 @@ from transformers import GraniteMoeHybridConfig
12
from vllm import envs
13
from vllm.attention.layer import Attention
14
from vllm.config import CacheConfig, VllmConfig
15
-from vllm.distributed import divide, get_tensor_model_parallel_world_size
15
+from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.distributed.parallel_state import get_pp_group
17
from vllm.forward_context import get_forward_context
18
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -21,8 +21,8 @@ from vllm.model_executor.layers.linear import (QKVParallelLinear,
21
from vllm.model_executor.layers.logits_processor import LogitsProcessor
22
from vllm.model_executor.layers.mamba.mamba2_metadata import (
23
Mamba2Metadata, prepare_mamba2_metadata)
24
-from vllm.model_executor.layers.mamba.mamba_mixer2 import (
25
- MambaMixer2, extra_groups_for_head_shards)
24
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
25
+from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
26
from vllm.model_executor.layers.quantization import QuantizationConfig
27
from vllm.model_executor.layers.rotary_embedding import get_rope
28
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -524,6 +524,38 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
524
}
525
embedding_padding_modules = ["lm_head"]
526
527
+ @classmethod
528
+ def get_mamba_state_shape_from_config(
529
+ cls,
530
+ vllm_config: "VllmConfig",
531
+ use_v1: bool = True,
532
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
533
+ """Calculate shapes for Mamba's convolutional and state caches.
534
+
535
+ Args:
536
+ vllm_config: vLLM config
537
+ use_v1: Get shapes for V1 (or V0)
538
+
539
+ Returns:
540
+ Tuple containing:
541
+ - conv_state_shape: Shape for convolutional state cache
542
+ - temporal_state_shape: Shape for state space model cache
543
+ """
544
+ parallel_config = vllm_config.parallel_config
545
+ hf_config = vllm_config.model_config.hf_config
546
+ intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
547
+
548
+ return get_mamba_state_shape(
549
+ intermediate_size=intermediate_size,
550
+ tp_world_size=parallel_config.tensor_parallel_size,
551
+ n_groups=hf_config.mamba_n_groups,
552
+ num_heads=hf_config.mamba_n_heads,
553
+ head_dim=hf_config.mamba_d_head,
554
+ state_size=hf_config.mamba_d_state,
555
+ conv_kernel=hf_config.mamba_d_conv,
556
+ use_v1=use_v1,
557
+ )
558
+
559
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
560
super().__init__()
561
@@ -587,9 +619,13 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
619
self.model_config.get_num_layers_by_block_type(
620
self.vllm_config.parallel_config,
621
LayerBlockType.mamba))
590
- self.mamba_cache = MambaCacheManager(
591
- self.vllm_config, self.model_config.dtype,
592
- num_mamba_layers, *self._get_mamba_cache_shape())
622
+ mamba_state_shape = \
623
+ self.get_mamba_state_shape_from_config(
624
+ self.vllm_config, use_v1=False)
625
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
626
+ self.model_config.dtype,
627
+ num_mamba_layers,
628
+ *mamba_state_shape)
629
630
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
631
@@ -605,38 +641,6 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
641
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
642
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
643
608
- def _get_mamba_cache_shape(
609
- self) -> tuple[tuple[int, int], tuple[int, int]]:
610
- world_size = get_tensor_model_parallel_world_size()
611
- hidden_size = self.config.hidden_size
612
-
613
- conv_state_shape, temporal_state_shape = None, None
614
-
615
- intermediate_size = self.config.mamba_expand * hidden_size
616
-
617
- # if n_groups is not divisible by world_size, need to extend the shards
618
- # to ensure all groups needed by a head is sharded along with it
619
- n_groups = (self.config.mamba_n_groups + extra_groups_for_head_shards(
620
- self.config.mamba_n_groups, world_size))
621
-
622
- # - heads and n_groups are TP-ed
623
- conv_dim = (intermediate_size +
624
- 2 * n_groups * self.config.mamba_d_state)
625
- conv_state_shape = (
626
- divide(conv_dim, world_size),
627
- self.config.mamba_d_conv - 1,
628
- )
629
-
630
- # These are not TP-ed as they depend on A, dt_bias, D
631
- # - they are typically small
632
- # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
633
- temporal_state_shape = (
634
- divide(self.config.mamba_n_heads, world_size),
635
- self.config.mamba_d_head,
636
- self.config.mamba_d_state,
637
- )
638
- return conv_state_shape, temporal_state_shape
639
-
644
def compute_logits(
645
self,
646
hidden_states: torch.Tensor,
vllm/model_executor/models/interfaces.py
+20
@@ -22,6 +22,7 @@ from .interfaces_base import is_pooling_model
22
23
if TYPE_CHECKING:
24
from vllm.attention import AttentionMetadata
25
+ from vllm.config import VllmConfig
26
from vllm.model_executor.models.utils import WeightsMapper
27
from vllm.sequence import IntermediateTensors
28
@@ -481,6 +482,25 @@ class IsHybrid(Protocol):
482
, also indicates that the model's hf_config has
483
'layers_block_type' """
484
485
+ @classmethod
486
+ def get_mamba_state_shape_from_config(
487
+ cls,
488
+ vllm_config: "VllmConfig",
489
+ use_v1: bool = True,
490
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
491
+ """Calculate shapes for Mamba's convolutional and state caches.
492
+
493
+ Args:
494
+ vllm_config: vLLM config
495
+ use_v1: Get shapes for V1 (or V0)
496
+
497
+ Returns:
498
+ Tuple containing:
499
+ - conv_state_shape: Shape for convolutional state cache
500
+ - temporal_state_shape: Shape for state space model cache
501
+ """
502
+ ...
503
+
504
505
@runtime_checkable
506
class _IsHybridType(Protocol):
vllm/model_executor/models/mamba2.py
+41
-39
@@ -11,15 +11,14 @@ from transformers import MambaConfig
11
from vllm import envs
12
from vllm.attention.backends.abstract import AttentionMetadata
13
from vllm.config import VllmConfig
14
-from vllm.distributed import divide, get_tensor_model_parallel_world_size
14
from vllm.distributed.parallel_state import get_pp_group
15
from vllm.forward_context import get_forward_context
16
from vllm.model_executor.layers.layernorm import RMSNorm
17
from vllm.model_executor.layers.logits_processor import LogitsProcessor
18
from vllm.model_executor.layers.mamba.mamba2_metadata import (
19
Mamba2Metadata, prepare_mamba2_metadata)
21
-from vllm.model_executor.layers.mamba.mamba_mixer2 import (
22
- MambaMixer2, extra_groups_for_head_shards)
20
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
21
+from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
22
from vllm.model_executor.layers.quantization.base_config import (
23
QuantizationConfig)
24
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -198,6 +197,38 @@ class Mamba2Model(nn.Module):
197
198
class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
199
200
+ @classmethod
201
+ def get_mamba_state_shape_from_config(
202
+ cls,
203
+ vllm_config: "VllmConfig",
204
+ use_v1: bool = True,
205
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
206
+ """Calculate shapes for Mamba's convolutional and state caches.
207
+
208
+ Args:
209
+ vllm_config: vLLM config
210
+ use_v1: Get shapes for V1 (or V0)
211
+
212
+ Returns:
213
+ Tuple containing:
214
+ - conv_state_shape: Shape for convolutional state cache
215
+ - temporal_state_shape: Shape for state space model cache
216
+ """
217
+ parallel_config = vllm_config.parallel_config
218
+ hf_config = vllm_config.model_config.hf_config
219
+ intermediate_size = hf_config.expand * hf_config.hidden_size
220
+
221
+ return get_mamba_state_shape(
222
+ intermediate_size=intermediate_size,
223
+ tp_world_size=parallel_config.tensor_parallel_size,
224
+ n_groups=hf_config.n_groups,
225
+ num_heads=hf_config.num_heads,
226
+ head_dim=hf_config.head_dim,
227
+ state_size=hf_config.state_size,
228
+ conv_kernel=hf_config.conv_kernel,
229
+ use_v1=use_v1,
230
+ )
231
+
232
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
233
config = vllm_config.model_config.hf_config
234
cache_config = vllm_config.cache_config
@@ -253,9 +284,13 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
284
self.model_config.get_num_layers_by_block_type(
285
self.vllm_config.parallel_config,
286
LayerBlockType.mamba))
256
- self.mamba_cache = MambaCacheManager(
257
- self.vllm_config, self.lm_head.weight.dtype,
258
- num_mamba_layers, *self._get_mamba_cache_shape())
287
+ mamba_state_shape = \
288
+ self.get_mamba_state_shape_from_config(
289
+ self.vllm_config, use_v1=False)
290
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
291
+ self.lm_head.weight.dtype,
292
+ num_mamba_layers,
293
+ *mamba_state_shape)
294
295
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
296
else:
@@ -274,39 +309,6 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
309
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
310
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
311
277
- def _get_mamba_cache_shape(
278
- self) -> tuple[tuple[int, int], tuple[int, int]]:
279
- world_size = get_tensor_model_parallel_world_size()
280
-
281
- conv_state_shape, temporal_state_shape = None, None
282
-
283
- intermediate_size = getattr(
284
- self.config, "intermediate_size",
285
- self.config.expand * self.config.hidden_size)
286
-
287
- # if n_groups is not divisible by world_size, need to extend the shards
288
- # to ensure all groups needed by a head is sharded along with it
289
- n_groups = (
290
- self.config.n_groups +
291
- extra_groups_for_head_shards(self.config.n_groups, world_size))
292
-
293
- # - heads and n_groups are TP-ed
294
- conv_dim = (intermediate_size + 2 * n_groups * self.config.state_size)
295
- conv_state_shape = (
296
- divide(conv_dim, world_size),
297
- self.config.conv_kernel - 1,
298
- )
299
-
300
- # These are not TP-ed as they depend on A, dt_bias, D
301
- # - they are typically small
302
- # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
303
- temporal_state_shape = (
304
- divide(self.config.num_heads, world_size),
305
- self.config.head_dim,
306
- self.config.state_size,
307
- )
308
- return conv_state_shape, temporal_state_shape
309
-
312
def compute_logits(self, hidden_states: torch.Tensor,
313
sampling_metadata: SamplingMetadata) -> torch.Tensor:
314
logits = self.logits_processor(self.lm_head, hidden_states,
vllm/model_executor/models/nemotron_h.py
+42
-40
@@ -26,7 +26,7 @@ from torch import nn
26
from vllm import envs
27
from vllm.attention.layer import Attention
28
from vllm.config import CacheConfig, VllmConfig
29
-from vllm.distributed import divide, get_tensor_model_parallel_world_size
29
+from vllm.distributed import get_tensor_model_parallel_world_size
30
from vllm.distributed.parallel_state import get_pp_group
31
from vllm.forward_context import get_forward_context
32
from vllm.model_executor.layers.activation import ReLUSquaredActivation
@@ -37,8 +37,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
37
from vllm.model_executor.layers.logits_processor import LogitsProcessor
38
from vllm.model_executor.layers.mamba.mamba2_metadata import (
39
Mamba2Metadata, prepare_mamba2_metadata)
40
-from vllm.model_executor.layers.mamba.mamba_mixer2 import (
41
- MambaMixer2, extra_groups_for_head_shards)
40
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
41
+from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
42
from vllm.model_executor.layers.quantization import QuantizationConfig
43
from vllm.model_executor.layers.vocab_parallel_embedding import (
44
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
@@ -459,6 +459,38 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
459
}
460
embedding_padding_modules = ["lm_head"]
461
462
+ @classmethod
463
+ def get_mamba_state_shape_from_config(
464
+ cls,
465
+ vllm_config: "VllmConfig",
466
+ use_v1: bool = True,
467
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
468
+ """Calculate shapes for Mamba's convolutional and state caches.
469
+
470
+ Args:
471
+ vllm_config: vLLM config
472
+ use_v1: Get shapes for V1 (or V0)
473
+
474
+ Returns:
475
+ Tuple containing:
476
+ - conv_state_shape: Shape for convolutional state cache
477
+ - temporal_state_shape: Shape for state space model cache
478
+ """
479
+ parallel_config = vllm_config.parallel_config
480
+ hf_config = vllm_config.model_config.hf_config
481
+ intermediate_size = hf_config.expand * hf_config.hidden_size
482
+
483
+ return get_mamba_state_shape(
484
+ intermediate_size=intermediate_size,
485
+ tp_world_size=parallel_config.tensor_parallel_size,
486
+ n_groups=hf_config.n_groups,
487
+ num_heads=hf_config.mamba_num_heads,
488
+ head_dim=hf_config.mamba_head_dim,
489
+ state_size=hf_config.ssm_state_size,
490
+ conv_kernel=hf_config.conv_kernel,
491
+ use_v1=use_v1,
492
+ )
493
+
494
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
495
config = vllm_config.model_config.hf_config
496
self.vllm_config = vllm_config
@@ -515,10 +547,13 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
547
self.vllm_config.parallel_config,
548
LayerBlockType.mamba
549
)
518
-
519
- self.mamba_cache = MambaCacheManager(
520
- self.vllm_config, self.lm_head.weight.dtype,
521
- num_mamba_layers, *self._get_mamba_cache_shape())
550
+ mamba_state_shape = \
551
+ self.get_mamba_state_shape_from_config(
552
+ self.vllm_config, use_v1=False)
553
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
554
+ self.lm_head.weight.dtype,
555
+ num_mamba_layers,
556
+ *mamba_state_shape)
557
558
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
559
@@ -534,39 +569,6 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
569
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
570
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
571
537
- def _get_mamba_cache_shape(
538
- self) -> tuple[tuple[int, int], tuple[int, int]]:
539
- world_size = get_tensor_model_parallel_world_size()
540
- hidden_size = self.config.hidden_size
541
-
542
- conv_state_shape, temporal_state_shape = None, None
543
-
544
- intermediate_size = self.config.expand * hidden_size
545
-
546
- # if n_groups is not divisible by world_size, need to extend the shards
547
- # to ensure all groups needed by a head is sharded along with it
548
- n_groups = (
549
- self.config.n_groups +
550
- extra_groups_for_head_shards(self.config.n_groups, world_size))
551
-
552
- # - heads and n_groups are TP-ed
553
- conv_dim = (intermediate_size +
554
- 2 * n_groups * self.config.ssm_state_size)
555
- conv_state_shape = (
556
- divide(conv_dim, world_size),
557
- self.config.conv_kernel - 1,
558
- )
559
-
560
- # These are not TP-ed as they depend on A, dt_bias, D
561
- # - they are typically small
562
- # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
563
- temporal_state_shape = (
564
- divide(self.config.mamba_num_heads, world_size),
565
- self.config.mamba_head_dim,
566
- self.config.ssm_state_size,
567
- )
568
- return conv_state_shape, temporal_state_shape
569
-
572
def compute_logits(
573
self,
574
hidden_states: torch.Tensor,
vllm/model_executor/models/zamba2.py
+43
-49
@@ -18,7 +18,7 @@ from transformers import Zamba2Config
18
from vllm import envs
19
from vllm.attention.layer import Attention
20
from vllm.config import CacheConfig, VllmConfig
21
-from vllm.distributed import divide, get_tensor_model_parallel_world_size
21
+from vllm.distributed import get_tensor_model_parallel_world_size
22
from vllm.forward_context import get_forward_context
23
from vllm.model_executor.layers.activation import GeluAndMul
24
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -30,8 +30,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
30
from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
from vllm.model_executor.layers.mamba.mamba2_metadata import (
32
Mamba2Metadata, prepare_mamba2_metadata)
33
-from vllm.model_executor.layers.mamba.mamba_mixer2 import (
34
- MambaMixer2, extra_groups_for_head_shards)
33
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
34
+from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
35
from vllm.model_executor.layers.quantization import QuantizationConfig
36
from vllm.model_executor.layers.rotary_embedding import get_rope
37
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -843,6 +843,39 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
843
"1.weight": "B.weight",
844
})
845
846
+ @classmethod
847
+ def get_mamba_state_shape_from_config(
848
+ cls,
849
+ vllm_config: "VllmConfig",
850
+ use_v1: bool = True,
851
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
852
+ """Calculate shapes for Mamba's convolutional and state caches.
853
+
854
+ Args:
855
+ vllm_config: vLLM config
856
+ use_v1: Get shapes for V1 (or V0)
857
+
858
+ Returns:
859
+ Tuple containing:
860
+ - conv_state_shape: Shape for convolutional state cache
861
+ - temporal_state_shape: Shape for state space model cache
862
+ """
863
+
864
+ parallel_config = vllm_config.parallel_config
865
+ hf_config = vllm_config.model_config.hf_config
866
+ intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
867
+
868
+ return get_mamba_state_shape(
869
+ intermediate_size=intermediate_size,
870
+ tp_world_size=parallel_config.tensor_parallel_size,
871
+ n_groups=hf_config.mamba_ngroups,
872
+ num_heads=hf_config.n_mamba_heads,
873
+ head_dim=hf_config.mamba_headdim,
874
+ state_size=hf_config.mamba_d_state,
875
+ conv_kernel=hf_config.mamba_d_conv,
876
+ use_v1=use_v1,
877
+ )
878
+
879
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
880
"""Initialize the Zamba2 model for causal language modeling.
881
@@ -925,9 +958,13 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
958
if not envs.VLLM_USE_V1:
959
if self.mamba_cache is None:
960
num_mamba_layers = self.config.num_hidden_layers
928
- self.mamba_cache = MambaCacheManager(
929
- self.vllm_config, self.lm_head.weight.dtype,
930
- num_mamba_layers, *self._get_mamba_cache_shape())
961
+ mamba_state_shape = \
962
+ self.get_mamba_state_shape_from_config(
963
+ self.vllm_config, use_v1=False)
964
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
965
+ self.lm_head.weight.dtype,
966
+ num_mamba_layers,
967
+ *mamba_state_shape)
968
969
# Get cache parameters for current run
970
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
@@ -968,49 +1005,6 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
1005
"""
1006
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
1007
971
- def _get_mamba_cache_shape(
972
- self) -> tuple[tuple[int, int], tuple[int, int]]:
973
- """Calculate shapes for Mamba's convolutional and state caches.
974
-
975
- Returns:
976
- Tuple containing:
977
- - conv_state_shape: Shape for convolutional state cache
978
- - temporal_state_shape: Shape for state space model cache
979
- """
980
- world_size = get_tensor_model_parallel_world_size()
981
-
982
- intermediate_size = self.config.mamba_expand * self.config.hidden_size
983
-
984
- # Extend groups if needed to ensure all groups needed by a head
985
- # are sharded together
986
-
987
- # if n_groups is not divisible by world_size, need to extend the shards
988
- # to ensure all groups needed by a head is sharded along with it
989
- n_groups = (self.config.mamba_ngroups + extra_groups_for_head_shards(
990
- self.config.mamba_ngroups, world_size))
991
-
992
- # Calculate conv state shape (includes groups)
993
- # - heads and n_groups are TP-ed
994
- conv_dim = (intermediate_size +
995
- 2 * n_groups * self.config.mamba_d_state)
996
- conv_state_shape = (
997
- divide(conv_dim, world_size),
998
- self.config.mamba_d_conv - 1,
999
- )
1000
-
1001
- # Calculate temporal state shape (per-head states)
1002
- # These are not TP-ed as they depend on A, dt_bias, D
1003
- # - they are typically small
1004
- # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
1005
- temporal_state_shape = (
1006
- divide(divide(intermediate_size, self.config.mamba_headdim),
1007
- world_size),
1008
- self.config.mamba_headdim,
1009
- self.config.mamba_d_state,
1010
- )
1011
-
1012
- return conv_state_shape, temporal_state_shape
1013
-
1008
def compute_logits(
1009
self,
1010
hidden_states: torch.Tensor,
vllm/v1/worker/gpu_model_runner.py
+3
-55
@@ -42,7 +42,7 @@ from vllm.pooling_params import PoolingParams
42
from vllm.sampling_params import SamplingType
43
from vllm.sequence import IntermediateTensors
44
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
45
- GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
45
+ GiB_bytes, LazyLoader, async_tensor_h2d,
46
check_use_alibi, get_dtype_size,
47
is_pin_memory_available, round_up)
48
from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
@@ -2648,9 +2648,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2648
"Prefix caching is not supported for Mamba yet.")
2649
max_model_len = self.vllm_config.model_config.max_model_len
2650
2651
- page_size_padded = self._maybe_pad_mamba_page_size(
2652
- attn_layers, mamba_layers, kv_cache_spec, max_model_len,
2653
- block_size)
2651
+ page_size_padded = (
2652
+ self.vllm_config.cache_config.mamba_page_size_padded)
2653
2654
# Set block_size to max_model_len, so that mamba model will always
2655
# have only one block in the KV cache.
@@ -2662,54 +2661,3 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2661
page_size_padded=page_size_padded)
2662
2663
return kv_cache_spec
2665
-
2666
- def _maybe_pad_mamba_page_size(
2667
- self,
2668
- attn_layers: dict[str, Attention],
2669
- mamba_layers: dict[str, MambaBase],
2670
- kv_cache_spec: dict[str, KVCacheSpec],
2671
- max_model_len: int,
2672
- block_size: int,
2673
- ) -> Optional[int]:
2674
- """
2675
- Ensure that page size of attention KV cache groups is greater than or
2676
- equal to the mamba KV cache groups. If not, we suggest to the user
2677
- how to set the attention block size to ensure that it is.
2678
-
2679
- If the attention page size is strictly greater than the mamba page size,
2680
- we pad the mamba page size to make them equal.
2681
-
2682
- Args:
2683
- attn_layers: Attention layers
2684
- mamba_layers: Mamba layers
2685
- kv_cache_spec: KV cache spec (populated with attention layers)
2686
-
2687
- Returns:
2688
- Optional[int]: Mamba page size with padding (None if no padding).
2689
- """
2690
-
2691
- if len(attn_layers) == 0:
2692
- return None
2693
-
2694
- attn_layer_name = next(iter(attn_layers))
2695
- attn_page_size = kv_cache_spec[attn_layer_name].page_size_bytes
2696
- mamba_layer_name = next(iter(mamba_layers))
2697
- mamba_page_size = MambaSpec(
2698
- shapes=mamba_layers[mamba_layer_name].get_state_shape(),
2699
- dtype=self.kv_cache_dtype,
2700
- block_size=max_model_len).page_size_bytes
2701
- if attn_page_size < mamba_page_size:
2702
- # attention page size (for 16 tokens)
2703
- attn_page_size_16 = 16 * attn_page_size // block_size
2704
- # some attention backends (e.g. FA) only support setting
2705
- # block size to multiple of 16, so let's suggest a value
2706
- # that would work (note: FA is currently not compatible
2707
- # with mamba layers, use FlashInfer instead).
2708
- suggest_attn_block_size = 16 * cdiv(mamba_page_size,
2709
- attn_page_size_16)
2710
- raise ValueError(
2711
- "Attention block size should be increased to at least "
2712
- f"{suggest_attn_block_size} in order to match "
2713
- "the mamba page size")
2714
-
2715
- return attn_page_size