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 14, 05:30 UTC
[V1] [Hybrid] Support using float32 for state in Hybrid Models (Mamba2, Mamba1, Minimax) (#22928)
Signed-off-by: Daniel Afrimi <danielafrimi8@gmail.com> Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Daniel Afrimi <danielafrimi8@gmail.com> Co-authored-by: Burkhard Ringlein <ngl@zurich.ibm.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>
Thomas Parnell committed
Aug 15, 2025 at 14:57 UTC
75531a6c134282f940c86461b3c40996b4136793
23 files changed
+467
-87
tests/models/language/generation/test_hybrid.py
+62
@@ -431,3 +431,65 @@ def test_full_cuda_graph(
431
name_0="hf" if hf_outputs is not None else "vllm-v0",
432
name_1="vllm-v1",
433
)
434
+
435
+
436
+@pytest.mark.parametrize("model", ["Zyphra/Zamba2-1.2B-instruct"])
437
+@pytest.mark.parametrize("max_tokens", [64])
438
+@pytest.mark.parametrize("num_logprobs", [5])
439
+def test_fp32_state(
440
+ hf_runner,
441
+ vllm_runner,
442
+ example_prompts,
443
+ monkeypatch,
444
+ model: str,
445
+ max_tokens: int,
446
+ num_logprobs: int,
447
+) -> None:
448
+
449
+ try:
450
+ model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
451
+ model_info.check_available_online(on_fail="skip")
452
+ model_info.check_transformers_version(on_fail="skip")
453
+ except ValueError:
454
+ pass
455
+
456
+ with hf_runner(model) as hf_model:
457
+ if model not in HF_UNSUPPORTED_MODELS:
458
+ hf_outputs = hf_model.generate_greedy_logprobs_limit(
459
+ example_prompts, max_tokens, num_logprobs)
460
+ else:
461
+ hf_outputs = None
462
+
463
+ with vllm_runner(model,
464
+ max_num_seqs=MAX_NUM_SEQS,
465
+ mamba_ssm_cache_dtype="float32") as vllm_model:
466
+ vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
467
+ example_prompts, max_tokens, num_logprobs)
468
+
469
+ with monkeypatch.context() as m:
470
+ m.setenv("VLLM_USE_V1", "1")
471
+ if model in HYBRID_MODELS:
472
+ # required due to reorder_batch behaviour
473
+ m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
474
+ with vllm_runner(model,
475
+ max_num_seqs=MAX_NUM_SEQS,
476
+ mamba_ssm_cache_dtype="float32",
477
+ enable_prefix_caching=False) as vllm_model:
478
+ vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
479
+ example_prompts, max_tokens, num_logprobs)
480
+
481
+ if hf_outputs is not None:
482
+ check_logprobs_close(
483
+ outputs_0_lst=hf_outputs,
484
+ outputs_1_lst=vllm_v0_outputs,
485
+ name_0="hf",
486
+ name_1="vllm-v0",
487
+ )
488
+
489
+ ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs
490
+ check_logprobs_close(
491
+ outputs_0_lst=ref_outputs,
492
+ outputs_1_lst=vllm_v1_outputs,
493
+ name_0="hf" if hf_outputs is not None else "vllm-v0",
494
+ name_1="vllm-v1",
495
+ )
tests/v1/worker/test_gpu_model_runner.py
+2
@@ -772,6 +772,8 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
772
head_dim=hf_config.mamba_d_head,
773
rms_norm_eps=hf_config.rms_norm_eps,
774
activation=hf_config.hidden_act,
775
+ cache_config=cache_config,
776
+ model_config=model_config,
777
prefix=key,
778
)
779
# suppress var not used error
vllm/config/__init__.py
+1
-1
@@ -29,7 +29,7 @@ from typing_extensions import Self, assert_never, runtime_checkable
29
30
import vllm.envs as envs
31
from vllm import version
32
-from vllm.config.cache import (BlockSize, CacheConfig, CacheDType,
32
+from vllm.config.cache import (BlockSize, CacheConfig, CacheDType, MambaDType,
33
PrefixCachingHashAlgo)
34
from vllm.config.compilation import (CompilationConfig, CompilationLevel,
35
PassConfig)
vllm/config/cache.py
+12
@@ -23,6 +23,7 @@ logger = init_logger(__name__)
23
24
BlockSize = Literal[1, 8, 16, 32, 64, 128]
25
CacheDType = Literal["auto", "fp8", "fp8_e4m3", "fp8_e5m2", "fp8_inc"]
26
+MambaDType = Literal["auto", "float32"]
27
PrefixCachingHashAlgo = Literal["builtin", "sha256", "sha256_cbor_64bit"]
28
29
@@ -93,6 +94,15 @@ class CacheConfig:
94
""" Optional override for mamba page size; used by hybrid mamba/attention
95
models to ensure exact alignment with attention page size."""
96
97
+ mamba_cache_dtype: MambaDType = "auto"
98
+ """The data type to use for the Mamba cache (both the conv as well as the
99
+ ssm state). If set to 'auto', the data type will be inferred from the model
100
+ config."""
101
+ mamba_ssm_cache_dtype: MambaDType = "auto"
102
+ """The data type to use for the Mamba cache (ssm state only, conv state will
103
+ still be controlled by mamba_cache_dtype). If set to 'auto', the data type
104
+ for the ssm state will be determined by mamba_cache_dtype."""
105
+
106
# Will be set after profiling.
107
num_gpu_blocks: Optional[int] = field(default=None, init=False)
108
"""The number of blocks to allocate for GPU memory."""
@@ -123,6 +133,8 @@ class CacheConfig:
133
"""
134
factors: list[Any] = []
135
factors.append(self.cache_dtype)
136
+ factors.append(self.mamba_cache_dtype)
137
+ factors.append(self.mamba_ssm_cache_dtype)
138
# `cpu_offload_gb` does not use `torch.compile` yet.
139
hash_str = hashlib.md5(str(factors).encode(),
140
usedforsecurity=False).hexdigest()
vllm/engine/arg_utils.py
+14
-6
@@ -27,12 +27,12 @@ from vllm.config import (BlockSize, CacheConfig, CacheDType, CompilationConfig,
27
DeviceConfig, DistributedExecutorBackend,
28
GuidedDecodingBackend, HfOverrides, KVEventsConfig,
29
KVTransferConfig, LoadConfig, LogprobsMode,
30
- LoRAConfig, ModelConfig, ModelDType, ModelImpl,
31
- MultiModalConfig, ObservabilityConfig, ParallelConfig,
32
- PoolerConfig, PrefixCachingHashAlgo, RunnerOption,
33
- SchedulerConfig, SchedulerPolicy, SpeculativeConfig,
34
- TaskOption, TokenizerMode, VllmConfig, get_attr_docs,
35
- get_field)
30
+ LoRAConfig, MambaDType, ModelConfig, ModelDType,
31
+ ModelImpl, MultiModalConfig, ObservabilityConfig,
32
+ ParallelConfig, PoolerConfig, PrefixCachingHashAlgo,
33
+ RunnerOption, SchedulerConfig, SchedulerPolicy,
34
+ SpeculativeConfig, TaskOption, TokenizerMode,
35
+ VllmConfig, get_attr_docs, get_field)
36
from vllm.logger import init_logger
37
from vllm.platforms import CpuArchEnum, current_platform
38
from vllm.plugins import load_general_plugins
@@ -422,6 +422,8 @@ class EngineArgs:
422
override_attention_dtype: str = ModelConfig.override_attention_dtype
423
424
calculate_kv_scales: bool = CacheConfig.calculate_kv_scales
425
+ mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype
426
+ mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype
427
428
additional_config: dict[str, Any] = \
429
get_field(VllmConfig, "additional_config")
@@ -694,6 +696,10 @@ class EngineArgs:
696
**cache_kwargs["calculate_kv_scales"])
697
cache_group.add_argument("--kv-sharing-fast-prefill",
698
**cache_kwargs["kv_sharing_fast_prefill"])
699
+ cache_group.add_argument("--mamba-cache-dtype",
700
+ **cache_kwargs["mamba_cache_dtype"])
701
+ cache_group.add_argument("--mamba-ssm-cache-dtype",
702
+ **cache_kwargs["mamba_ssm_cache_dtype"])
703
704
# Multimodal related configs
705
multimodal_kwargs = get_kwargs(MultiModalConfig)
@@ -1105,6 +1111,8 @@ class EngineArgs:
1111
cpu_offload_gb=self.cpu_offload_gb,
1112
calculate_kv_scales=self.calculate_kv_scales,
1113
kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
1114
+ mamba_cache_dtype=self.mamba_cache_dtype,
1115
+ mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,
1116
)
1117
1118
ray_runtime_env = None
vllm/model_executor/layers/mamba/mamba_mixer.py
+15
-2
@@ -9,7 +9,7 @@ from torch.nn.parameter import Parameter
9
10
from vllm import envs
11
from vllm.attention.backends.abstract import AttentionMetadata
12
-from vllm.config import get_current_vllm_config
12
+from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
13
from vllm.distributed.parallel_state import (
14
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
15
from vllm.forward_context import ForwardContext, get_forward_context
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
20
RowParallelLinear)
21
from vllm.model_executor.layers.mamba.abstract import MambaBase
22
from vllm.model_executor.layers.mamba.mamba_utils import (
23
- MambaStateShapeCalculator)
23
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
24
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
25
causal_conv1d_fn, causal_conv1d_update)
26
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
@@ -56,6 +56,8 @@ class MambaMixer(MambaBase, CustomOp):
56
rms_norm_eps: float = 1e-5,
57
activation="silu",
58
is_lora_enabled: bool = False,
59
+ model_config: Optional[ModelConfig] = None,
60
+ cache_config: Optional[CacheConfig] = None,
61
prefix: str = ""):
62
super().__init__()
63
self.time_step_rank = time_step_rank
@@ -153,6 +155,8 @@ class MambaMixer(MambaBase, CustomOp):
155
# The inner tuple is (conv_state, ssm_state)
156
self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
157
158
+ self.model_config = model_config
159
+ self.cache_config = cache_config
160
self.prefix = prefix
161
162
def _ssm_transform(
@@ -369,6 +373,15 @@ class MambaMixer(MambaBase, CustomOp):
373
374
return out
375
376
+ def get_state_dtype(self) -> tuple[torch.dtype]:
377
+ assert self.model_config is not None
378
+ assert self.cache_config is not None
379
+ return MambaStateDtypeCalculator.mamba1_state_dtype(
380
+ self.model_config.dtype,
381
+ self.cache_config.mamba_cache_dtype,
382
+ self.cache_config.mamba_ssm_cache_dtype,
383
+ )
384
+
385
def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
386
return MambaStateShapeCalculator.mamba1_state_shape(
387
tp_world_size=get_tensor_model_parallel_world_size(),
vllm/model_executor/layers/mamba/mamba_mixer2.py
+31
-20
@@ -8,7 +8,7 @@ from torch import nn
8
9
from vllm import envs
10
from vllm.attention.backends.abstract import AttentionMetadata
11
-from vllm.config import get_current_vllm_config
11
+from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
12
from vllm.distributed import (divide, get_tensor_model_parallel_rank,
13
get_tensor_model_parallel_world_size,
14
tensor_model_parallel_all_gather,
@@ -21,7 +21,7 @@ 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
- MambaStateShapeCalculator)
24
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
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.layernorm_gated import rms_norm_gated
@@ -218,23 +218,23 @@ class MambaMixer2(MambaBase, CustomOp):
218
**selective** state spaces)
219
"""
220
221
- def __init__(
222
- self,
223
- hidden_size: int,
224
- ssm_state_size: int,
225
- conv_kernel_size: int,
226
- intermediate_size: int,
227
- use_conv_bias: bool,
228
- use_bias: bool,
229
- n_groups: int = 1,
230
- num_heads: int = 128,
231
- head_dim: int = 64,
232
- rms_norm_eps: float = 1e-5,
233
- activation: str = "silu",
234
- use_rms_norm: bool = True,
235
- quant_config: Optional[QuantizationConfig] = None,
236
- prefix: str = "",
237
- ):
221
+ def __init__(self,
222
+ hidden_size: int,
223
+ ssm_state_size: int,
224
+ conv_kernel_size: int,
225
+ intermediate_size: int,
226
+ use_conv_bias: bool,
227
+ use_bias: bool,
228
+ n_groups: int = 1,
229
+ num_heads: int = 128,
230
+ head_dim: int = 64,
231
+ rms_norm_eps: float = 1e-5,
232
+ activation: str = "silu",
233
+ use_rms_norm: bool = True,
234
+ model_config: Optional[ModelConfig] = None,
235
+ cache_config: Optional[CacheConfig] = None,
236
+ quant_config: Optional[QuantizationConfig] = None,
237
+ prefix: str = ""):
238
super().__init__()
239
240
# For TP, the sharding plan is as follows:
@@ -417,6 +417,8 @@ class MambaMixer2(MambaBase, CustomOp):
417
# The inner tuple is (conv_state, ssm_state)
418
self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
419
420
+ self.model_config = model_config
421
+ self.cache_config = cache_config
422
self.prefix = prefix
423
424
def forward_native(
@@ -670,7 +672,7 @@ class MambaMixer2(MambaBase, CustomOp):
672
dt_limit=(0.0, float("inf")),
673
out=preallocated_ssm_out_p.view(1, num_prefill_tokens, -1,
674
self.head_dim),
673
- )
675
+ state_dtype=ssm_state.dtype)
676
677
# update ssm states
678
# - varlen state is a (num_prefills, nheads, headdim, dstate) tensor
@@ -732,6 +734,15 @@ class MambaMixer2(MambaBase, CustomOp):
734
# 5. Final linear projection
735
output[:num_actual_tokens], _ = self.out_proj(hidden_states)
736
737
+ def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
738
+ assert self.model_config is not None
739
+ assert self.cache_config is not None
740
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
741
+ self.model_config.dtype,
742
+ self.cache_config.mamba_cache_dtype,
743
+ self.cache_config.mamba_ssm_cache_dtype,
744
+ )
745
+
746
def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
747
return MambaStateShapeCalculator.mamba2_state_shape(
748
intermediate_size=self.intermediate_size,
vllm/model_executor/layers/mamba/mamba_utils.py
+52
@@ -1,6 +1,58 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Union
4
+
5
+import torch
6
+
7
+from vllm.config import MambaDType, ModelDType
8
from vllm.distributed import divide
9
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, get_kv_cache_torch_dtype
10
+
11
+
12
+class MambaStateDtypeCalculator:
13
+
14
+ @classmethod
15
+ def linear_attention_state_dtype(
16
+ cls,
17
+ model_dtype: Union[ModelDType, torch.dtype],
18
+ mamba_cache_dtype: MambaDType,
19
+ ) -> tuple[torch.dtype, ...]:
20
+ # TODO (tdoublep) requires testing
21
+ if mamba_cache_dtype == "float32":
22
+ raise ValueError("fp32 state for minimax is not yet supported")
23
+ state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
24
+ return (state_dtype, )
25
+
26
+ @classmethod
27
+ def mamba1_state_dtype(
28
+ cls,
29
+ model_dtype: Union[ModelDType, torch.dtype],
30
+ mamba_cache_dtype: MambaDType,
31
+ mamba_ssm_cache_dtype: MambaDType,
32
+ ) -> tuple[torch.dtype, ...]:
33
+ # TODO (tdoublep) requires kernel changes
34
+ if mamba_cache_dtype == "float32" or mamba_ssm_cache_dtype == "float32":
35
+ raise ValueError("fp32 state for mamba1 is not yet supported")
36
+ else:
37
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
38
+ model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype)
39
+
40
+ @classmethod
41
+ def mamba2_state_dtype(
42
+ cls,
43
+ model_dtype: Union[ModelDType, torch.dtype],
44
+ mamba_cache_dtype: MambaDType,
45
+ mamba_ssm_cache_dtype: MambaDType,
46
+ ) -> tuple[torch.dtype, ...]:
47
+ conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype,
48
+ model_dtype)
49
+ if mamba_ssm_cache_dtype == "auto":
50
+ temporal_state_dtype = conv_state_dtype
51
+ else:
52
+ temporal_state_dtype = (
53
+ STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype])
54
+
55
+ return (conv_state_dtype, temporal_state_dtype)
56
57
58
class MambaStateShapeCalculator:
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+7
-3
@@ -41,6 +41,7 @@ def _mamba_chunk_scan_combined_fwd(x,
41
cu_seqlens=None,
42
dt_softplus=False,
43
dt_limit=(0.0, float("inf")),
44
+ state_dtype=None,
45
out=None):
46
assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2"
47
batch, seqlen, nheads, headdim = x.shape
@@ -118,7 +119,7 @@ def _mamba_chunk_scan_combined_fwd(x,
119
if initial_states is not None else None,
120
seq_idx=seq_idx,
121
chunk_size=chunk_size,
121
- out_dtype=C.dtype,
122
+ out_dtype=state_dtype if state_dtype is not None else C.dtype,
123
is_cont_batched=cu_seqlens is not None)
124
states, final_states = (rearrange(t, "... (p n) -> ... p n", n=dstate)
125
for t in [states, final_states])
@@ -189,7 +190,8 @@ def mamba_chunk_scan_combined(x,
190
dt_limit=(0.0, float("inf")),
191
out=None,
192
return_final_states=False,
192
- return_varlen_states=False):
193
+ return_varlen_states=False,
194
+ state_dtype=None):
195
"""
196
Argument:
197
x: (batch, seqlen, nheads, headdim)
@@ -206,6 +208,7 @@ def mamba_chunk_scan_combined(x,
208
cu_seqlens: (num_sequences + 1) or None, only used if return_varlen_states is True
209
dt_softplus: Whether to apply softplus to dt
210
out: Preallocated output tensor
211
+ state_dtype: The data type of the ssm state
212
"""
213
214
if not return_varlen_states:
@@ -229,7 +232,8 @@ def mamba_chunk_scan_combined(x,
232
cu_seqlens=cu_seqlens,
233
dt_softplus=dt_softplus,
234
dt_limit=dt_limit,
232
- out=out)
235
+ out=out,
236
+ state_dtype=state_dtype)
237
if not return_varlen_states:
238
if not return_final_states:
239
return
vllm/model_executor/models/bamba.py
+25
-4
@@ -12,7 +12,7 @@ from transformers import BambaConfig
12
from vllm import envs
13
from vllm.attention.layer import Attention
14
from vllm.compilation.decorators import support_torch_compile
15
-from vllm.config import CacheConfig, VllmConfig
15
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
16
from vllm.distributed import get_tensor_model_parallel_world_size
17
from vllm.distributed.parallel_state import get_pp_group
18
from vllm.forward_context import get_forward_context
@@ -26,7 +26,7 @@ from vllm.model_executor.layers.mamba.mamba2_metadata import (
26
Mamba2Metadata, prepare_mamba2_metadata)
27
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
28
from vllm.model_executor.layers.mamba.mamba_utils import (
29
- MambaStateShapeCalculator)
29
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
30
from vllm.model_executor.layers.quantization import QuantizationConfig
31
from vllm.model_executor.layers.rotary_embedding import get_rope
32
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -83,6 +83,7 @@ class BambaMixerDecoderLayer(nn.Module):
83
def __init__(self,
84
config: BambaConfig,
85
layer_idx: int,
86
+ model_config: Optional[ModelConfig] = None,
87
cache_config: Optional[CacheConfig] = None,
88
quant_config: Optional[QuantizationConfig] = None,
89
prefix: str = "") -> None:
@@ -100,6 +101,8 @@ class BambaMixerDecoderLayer(nn.Module):
101
head_dim=config.mamba_d_head,
102
rms_norm_eps=config.rms_norm_eps,
103
activation=config.hidden_act,
104
+ model_config=model_config,
105
+ cache_config=cache_config,
106
quant_config=quant_config,
107
prefix=f"{prefix}.mixer")
108
@@ -138,6 +141,7 @@ class BambaAttentionDecoderLayer(nn.Module):
141
self,
142
config: BambaConfig,
143
layer_idx: int,
144
+ model_config: Optional[ModelConfig] = None,
145
cache_config: Optional[CacheConfig] = None,
146
quant_config: Optional[QuantizationConfig] = None,
147
prefix: str = "",
@@ -266,6 +270,7 @@ class BambaModel(nn.Module):
270
super().__init__()
271
272
config: BambaConfig = vllm_config.model_config.hf_config
273
+ model_config = vllm_config.model_config
274
cache_config = vllm_config.cache_config
275
quant_config = vllm_config.quant_config
276
lora_config = vllm_config.lora_config
@@ -289,6 +294,7 @@ class BambaModel(nn.Module):
294
return layer_class(
295
config,
296
layer_idx,
297
+ model_config,
298
cache_config,
299
quant_config=quant_config,
300
prefix=prefix,
@@ -437,6 +443,18 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
443
}
444
embedding_padding_modules = ["lm_head"]
445
446
+ @classmethod
447
+ def get_mamba_state_dtype_from_config(
448
+ cls,
449
+ vllm_config: "VllmConfig",
450
+ ) -> tuple[torch.dtype, torch.dtype]:
451
+
452
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
453
+ vllm_config.model_config.dtype,
454
+ vllm_config.cache_config.mamba_cache_dtype,
455
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
456
+ )
457
+
458
@classmethod
459
def get_mamba_state_shape_from_config(
460
cls,
@@ -528,10 +546,13 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
546
mamba_state_shape = \
547
self.get_mamba_state_shape_from_config(
548
self.vllm_config, use_v1=False)
549
+ mamba_state_dtype = \
550
+ self.get_mamba_state_dtype_from_config(
551
+ self.vllm_config)
552
self.mamba_cache = MambaCacheManager(self.vllm_config,
532
- self.lm_head.weight.dtype,
553
num_mamba_layers,
534
- *mamba_state_shape)
554
+ *mamba_state_shape,
555
+ *mamba_state_dtype)
556
557
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
558
vllm/model_executor/models/config.py
+1
-1
@@ -318,7 +318,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
318
# get mamba page size
319
mamba_page_size = MambaSpec(
320
shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
321
- dtype=kv_cache_dtype,
321
+ dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
322
block_size=model_config.max_model_len,
323
).page_size_bytes
324
vllm/model_executor/models/falcon_h1.py
+25
-4
@@ -11,7 +11,7 @@ from transformers import FalconH1Config
11
from vllm import envs
12
from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
-from vllm.config import CacheConfig, VllmConfig
14
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
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
@@ -25,7 +25,7 @@ from vllm.model_executor.layers.mamba.mamba2_metadata import (
25
Mamba2Metadata, prepare_mamba2_metadata)
26
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
27
from vllm.model_executor.layers.mamba.mamba_utils import (
28
- MambaStateShapeCalculator)
28
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
29
from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.model_executor.layers.rotary_embedding import get_rope
31
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -85,6 +85,7 @@ class FalconH1SSMDecoderLayer(nn.Module):
85
def __init__(
86
self,
87
config: FalconH1Config,
88
+ model_config: Optional[ModelConfig] = None,
89
cache_config: Optional[CacheConfig] = None,
90
quant_config: Optional[QuantizationConfig] = None,
91
prefix: str = "",
@@ -108,6 +109,8 @@ class FalconH1SSMDecoderLayer(nn.Module):
109
head_dim=config.mamba_d_head,
110
rms_norm_eps=config.rms_norm_eps,
111
activation=config.hidden_act,
112
+ model_config=model_config,
113
+ cache_config=cache_config,
114
quant_config=quant_config,
115
use_rms_norm=config.mamba_rms_norm,
116
prefix=f"{prefix}.mixer",
@@ -317,6 +320,7 @@ class FalconH1ParallelHybrid(nn.Module):
320
self,
321
config: FalconH1Config,
322
layer_idx: int,
323
+ model_config: Optional[ModelConfig] = None,
324
cache_config: Optional[CacheConfig] = None,
325
quant_config: Optional[QuantizationConfig] = None,
326
prefix: str = "",
@@ -339,6 +343,7 @@ class FalconH1ParallelHybrid(nn.Module):
343
# Instantiate the SSM branch
344
self.mamba = FalconH1SSMDecoderLayer(
345
config=config,
346
+ model_config=model_config,
347
cache_config=cache_config,
348
quant_config=quant_config,
349
prefix=ssm_prefix,
@@ -408,6 +413,7 @@ class FalconH1Model(nn.Module):
413
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
414
super().__init__()
415
config: FalconH1Config = vllm_config.model_config.hf_config
416
+ model_config = vllm_config.model_config
417
cache_config = vllm_config.cache_config
418
quant_config = vllm_config.quant_config
419
lora_config = vllm_config.lora_config
@@ -435,6 +441,7 @@ class FalconH1Model(nn.Module):
441
return layer_class(
442
config,
443
layer_idx,
444
+ model_config,
445
cache_config,
446
quant_config=quant_config,
447
prefix=prefix,
@@ -519,6 +526,18 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
526
}
527
embedding_padding_modules = ["lm_head"]
528
529
+ @classmethod
530
+ def get_mamba_state_dtype_from_config(
531
+ cls,
532
+ vllm_config: "VllmConfig",
533
+ ) -> tuple[torch.dtype, torch.dtype]:
534
+
535
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
536
+ vllm_config.model_config.dtype,
537
+ vllm_config.cache_config.mamba_cache_dtype,
538
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
539
+ )
540
+
541
@classmethod
542
def get_mamba_state_shape_from_config(
543
cls,
@@ -624,12 +643,14 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
643
mamba_state_shape = \
644
self.get_mamba_state_shape_from_config(
645
self.vllm_config, use_v1=False)
646
+ mamba_state_dtype = \
647
+ self.get_mamba_state_dtype_from_config(
648
+ self.vllm_config)
649
self.mamba_cache = MambaCacheManager(
650
self.vllm_config,
629
- self.lm_head.weight.dtype if hasattr(
630
- self.lm_head, 'weight') else torch.bfloat16,
651
self.config.num_hidden_layers,
652
*mamba_state_shape,
653
+ *mamba_state_dtype,
654
)
655
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
656
vllm/model_executor/models/granitemoehybrid.py
+26
-4
@@ -12,7 +12,7 @@ from transformers import GraniteMoeHybridConfig
12
from vllm import envs
13
from vllm.attention.layer import Attention
14
from vllm.compilation.decorators import support_torch_compile
15
-from vllm.config import CacheConfig, VllmConfig
15
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
16
from vllm.distributed import get_tensor_model_parallel_world_size
17
from vllm.distributed.parallel_state import get_pp_group
18
from vllm.forward_context import get_forward_context
@@ -24,7 +24,7 @@ from vllm.model_executor.layers.mamba.mamba2_metadata import (
24
Mamba2Metadata, prepare_mamba2_metadata)
25
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
26
from vllm.model_executor.layers.mamba.mamba_utils import (
27
- MambaStateShapeCalculator)
27
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
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 (
@@ -50,6 +50,7 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):
50
def __init__(self,
51
config: GraniteMoeHybridConfig,
52
layer_idx: int,
53
+ model_config: Optional[ModelConfig] = None,
54
cache_config: Optional[CacheConfig] = None,
55
quant_config: Optional[QuantizationConfig] = None,
56
prefix: str = "") -> None:
@@ -70,6 +71,8 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):
71
head_dim=config.mamba_d_head,
72
rms_norm_eps=config.rms_norm_eps,
73
activation=config.hidden_act,
74
+ model_config=model_config,
75
+ cache_config=cache_config,
76
quant_config=quant_config,
77
prefix=f"{prefix}.mixer")
78
@@ -137,6 +140,7 @@ class GraniteMoeHybridAttentionDecoderLayer(nn.Module):
140
self,
141
config: GraniteMoeHybridConfig,
142
layer_idx: int,
143
+ model_config: Optional[ModelConfig] = None,
144
cache_config: Optional[CacheConfig] = None,
145
quant_config: Optional[QuantizationConfig] = None,
146
prefix: str = "",
@@ -217,6 +221,7 @@ class GraniteMoeHybridAttention(nn.Module):
221
def __init__(
222
self,
223
config: GraniteMoeHybridConfig,
224
+ model_config: Optional[ModelConfig] = None,
225
cache_config: Optional[CacheConfig] = None,
226
quant_config: Optional[QuantizationConfig] = None,
227
prefix: str = "",
@@ -316,6 +321,7 @@ class GraniteMoeHybridModel(nn.Module):
321
super().__init__()
322
323
config = vllm_config.model_config.hf_config
324
+ model_config = vllm_config.model_config
325
cache_config = vllm_config.cache_config
326
quant_config = vllm_config.quant_config
327
lora_config = vllm_config.lora_config
@@ -340,6 +346,7 @@ class GraniteMoeHybridModel(nn.Module):
346
return layer_class(
347
config,
348
layer_idx,
349
+ model_config,
350
cache_config,
351
quant_config=quant_config,
352
prefix=prefix,
@@ -527,6 +534,18 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
534
}
535
embedding_padding_modules = ["lm_head"]
536
537
+ @classmethod
538
+ def get_mamba_state_dtype_from_config(
539
+ cls,
540
+ vllm_config: "VllmConfig",
541
+ ) -> tuple[torch.dtype, torch.dtype]:
542
+
543
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
544
+ vllm_config.model_config.dtype,
545
+ vllm_config.cache_config.mamba_cache_dtype,
546
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
547
+ )
548
+
549
@classmethod
550
def get_mamba_state_shape_from_config(
551
cls,
@@ -625,10 +644,13 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
644
mamba_state_shape = \
645
self.get_mamba_state_shape_from_config(
646
self.vllm_config, use_v1=False)
647
+ mamba_state_dtype = \
648
+ self.get_mamba_state_dtype_from_config(
649
+ self.vllm_config)
650
self.mamba_cache = MambaCacheManager(self.vllm_config,
629
- self.model_config.dtype,
651
num_mamba_layers,
631
- *mamba_state_shape)
652
+ *mamba_state_shape,
653
+ *mamba_state_dtype)
654
655
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
656
vllm/model_executor/models/jamba.py
+24
-4
@@ -10,7 +10,7 @@ from transformers import JambaConfig
10
11
from vllm import envs
12
from vllm.attention.layer import Attention
13
-from vllm.config import CacheConfig, VllmConfig
13
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -21,7 +21,7 @@ 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.mamba_mixer import MambaMixer
23
from vllm.model_executor.layers.mamba.mamba_utils import (
24
- MambaStateShapeCalculator)
24
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
25
from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
26
from vllm.model_executor.layers.quantization import QuantizationConfig
27
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -94,6 +94,7 @@ class JambaMambaDecoderLayer(nn.Module):
94
def __init__(self,
95
config: JambaConfig,
96
layer_idx: int,
97
+ model_config: Optional[ModelConfig] = None,
98
cache_config: Optional[CacheConfig] = None,
99
quant_config: Optional[QuantizationConfig] = None,
100
is_lora_enabled: Optional[bool] = False,
@@ -114,6 +115,8 @@ class JambaMambaDecoderLayer(nn.Module):
115
rms_norm_eps=config.rms_norm_eps,
116
activation=config.hidden_act,
117
is_lora_enabled = self.is_lora_enabled,
118
+ model_config=model_config,
119
+ cache_config=cache_config,
120
prefix=f"{prefix}.mixer",
121
)
122
@@ -164,6 +167,7 @@ class JambaAttentionDecoderLayer(nn.Module):
167
def __init__(self,
168
config: JambaConfig,
169
layer_idx: int,
170
+ model_config: Optional[ModelConfig] = None,
171
cache_config: Optional[CacheConfig] = None,
172
quant_config: Optional[QuantizationConfig] = None,
173
prefix: str = "",
@@ -280,6 +284,7 @@ class JambaModel(nn.Module):
284
super().__init__()
285
286
config = vllm_config.model_config.hf_config
287
+ model_config = vllm_config.model_config
288
cache_config = vllm_config.cache_config
289
quant_config = vllm_config.quant_config
290
lora_config = vllm_config.lora_config
@@ -304,6 +309,7 @@ class JambaModel(nn.Module):
309
config.layers_block_type[layer_idx]]
310
return layer_class(config,
311
layer_idx,
312
+ model_config,
313
cache_config,
314
quant_config=quant_config,
315
prefix=prefix,
@@ -520,9 +526,11 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
526
self.vllm_config.parallel_config, LayerBlockType.mamba)
527
state_shape = self.get_mamba_state_shape_from_config(
528
self.vllm_config)
529
+ state_dtype = self.get_mamba_state_dtype_from_config(
530
+ self.vllm_config)
531
self.mamba_cache = MambaCacheManager(self.vllm_config,
524
- self.lm_head.weight.dtype,
525
- num_layers, *state_shape)
532
+ num_layers, *state_shape,
533
+ *state_dtype)
534
535
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
536
@@ -537,6 +545,18 @@ class JambaForCausalLM(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
548
+ @classmethod
549
+ def get_mamba_state_dtype_from_config(
550
+ cls,
551
+ vllm_config: "VllmConfig",
552
+ ) -> tuple[torch.dtype, torch.dtype]:
553
+
554
+ return MambaStateDtypeCalculator.mamba1_state_dtype(
555
+ vllm_config.model_config.dtype,
556
+ vllm_config.cache_config.mamba_cache_dtype,
557
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
558
+ )
559
+
560
@classmethod
561
def get_mamba_state_shape_from_config(
562
cls,
vllm/model_executor/models/mamba.py
+23
-4
@@ -9,13 +9,13 @@ from torch import nn
9
from transformers import MambaConfig
10
11
from vllm import envs
12
-from vllm.config import CacheConfig, VllmConfig
12
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
13
from vllm.distributed.parallel_state import get_pp_group
14
from vllm.model_executor.layers.layernorm import RMSNorm
15
from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
17
from vllm.model_executor.layers.mamba.mamba_utils import (
18
- MambaStateShapeCalculator)
18
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
19
from vllm.model_executor.layers.quantization.base_config import (
20
QuantizationConfig)
21
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -40,6 +40,7 @@ class MambaDecoderLayer(nn.Module):
40
41
def __init__(self,
42
config: MambaConfig,
43
+ model_config: Optional[ModelConfig] = None,
44
cache_config: Optional[CacheConfig] = None,
45
quant_config: Optional[QuantizationConfig] = None,
46
is_lora_enabled: Optional[bool] = False,
@@ -61,6 +62,8 @@ class MambaDecoderLayer(nn.Module):
62
rms_norm_eps=mixer_rms_eps,
63
activation=config.hidden_act,
64
is_lora_enabled=self.is_lora_enabled,
65
+ model_config=model_config,
66
+ cache_config=cache_config,
67
prefix=f"{prefix}.mixer")
68
69
self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
@@ -88,6 +91,7 @@ class MambaModel(nn.Module):
91
super().__init__()
92
93
config = vllm_config.model_config.hf_config
94
+ model_config = vllm_config.model_config
95
cache_config = vllm_config.cache_config
96
quant_config = vllm_config.quant_config
97
lora_config = vllm_config.lora_config
@@ -108,6 +112,7 @@ class MambaModel(nn.Module):
112
self.start_layer, self.end_layer, self.layers = make_layers(
113
config.num_hidden_layers,
114
lambda prefix: MambaDecoderLayer(config,
115
+ model_config=model_config,
116
cache_config=cache_config,
117
quant_config=quant_config,
118
is_lora_enabled=is_lora_enabled,
@@ -243,9 +248,11 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
248
self.vllm_config.parallel_config, LayerBlockType.mamba)
249
state_shape = self.get_mamba_state_shape_from_config(
250
self.vllm_config)
251
+ state_dtype = self.get_mamba_state_dtype_from_config(
252
+ self.vllm_config)
253
self.mamba_cache = MambaCacheManager(self.vllm_config,
247
- self.lm_head.weight.dtype,
248
- num_layers, *state_shape)
254
+ num_layers, *state_shape,
255
+ *state_dtype)
256
257
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
258
@@ -254,6 +261,18 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
261
262
return hidden_states
263
264
+ @classmethod
265
+ def get_mamba_state_dtype_from_config(
266
+ cls,
267
+ vllm_config: "VllmConfig",
268
+ ) -> tuple[torch.dtype, torch.dtype]:
269
+
270
+ return MambaStateDtypeCalculator.mamba1_state_dtype(
271
+ vllm_config.model_config.dtype,
272
+ vllm_config.cache_config.mamba_cache_dtype,
273
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
274
+ )
275
+
276
@classmethod
277
def get_mamba_state_shape_from_config(
278
cls,
vllm/model_executor/models/mamba2.py
+30
-6
@@ -11,7 +11,7 @@ from transformers import MambaConfig
11
from vllm import envs
12
from vllm.attention.backends.abstract import AttentionMetadata
13
from vllm.compilation.decorators import support_torch_compile
14
-from vllm.config import VllmConfig
14
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
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.layernorm import RMSNorm
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.mamba.mamba2_metadata import (
20
Mamba2Metadata, prepare_mamba2_metadata)
21
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
22
from vllm.model_executor.layers.mamba.mamba_utils import (
23
- MambaStateShapeCalculator)
23
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
24
from vllm.model_executor.layers.quantization.base_config import (
25
QuantizationConfig)
26
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -45,6 +45,8 @@ class Mamba2DecoderLayer(nn.Module):
45
46
def __init__(self,
47
config: MambaConfig,
48
+ model_config: Optional[ModelConfig] = None,
49
+ cache_config: Optional[CacheConfig] = None,
50
quant_config: Optional[QuantizationConfig] = None,
51
prefix: str = "") -> None:
52
super().__init__()
@@ -62,6 +64,8 @@ class Mamba2DecoderLayer(nn.Module):
64
head_dim=config.head_dim,
65
rms_norm_eps=config.layer_norm_epsilon,
66
activation=config.hidden_act,
67
+ model_config=model_config,
68
+ cache_config=cache_config,
69
quant_config=quant_config,
70
prefix=f"{prefix}.mixer")
71
@@ -93,6 +97,8 @@ class Mamba2Model(nn.Module):
97
super().__init__()
98
99
config = vllm_config.model_config.hf_config
100
+ model_config = vllm_config.model_config
101
+ cache_config = vllm_config.cache_config
102
quant_config = vllm_config.quant_config
103
lora_config = vllm_config.lora_config
104
is_lora_enabled = bool(lora_config)
@@ -112,8 +118,11 @@ class Mamba2Model(nn.Module):
118
119
self.start_layer, self.end_layer, self.layers = make_layers(
120
config.num_hidden_layers,
115
- lambda prefix: Mamba2DecoderLayer(
116
- config, quant_config=quant_config, prefix=prefix),
121
+ lambda prefix: Mamba2DecoderLayer(config,
122
+ model_config=model_config,
123
+ cache_config=cache_config,
124
+ quant_config=quant_config,
125
+ prefix=prefix),
126
prefix=f"{prefix}.layers")
127
128
self.norm_f = RMSNorm(config.hidden_size,
@@ -200,6 +209,18 @@ class Mamba2Model(nn.Module):
209
210
class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
211
212
+ @classmethod
213
+ def get_mamba_state_dtype_from_config(
214
+ cls,
215
+ vllm_config: "VllmConfig",
216
+ ) -> tuple[torch.dtype, torch.dtype]:
217
+
218
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
219
+ vllm_config.model_config.dtype,
220
+ vllm_config.cache_config.mamba_cache_dtype,
221
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
222
+ )
223
+
224
@classmethod
225
def get_mamba_state_shape_from_config(
226
cls,
@@ -290,10 +311,13 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
311
mamba_state_shape = \
312
self.get_mamba_state_shape_from_config(
313
self.vllm_config, use_v1=False)
314
+ mamba_state_dtype = \
315
+ self.get_mamba_state_dtype_from_config(
316
+ self.vllm_config)
317
self.mamba_cache = MambaCacheManager(self.vllm_config,
294
- self.lm_head.weight.dtype,
318
num_mamba_layers,
296
- *mamba_state_shape)
319
+ *mamba_state_shape,
320
+ *mamba_state_dtype)
321
322
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
323
else:
vllm/model_executor/models/mamba_cache.py
+10
-5
@@ -24,9 +24,14 @@ class MambaCacheParams:
24
25
class MambaCacheManager(ConstantSizeCache):
26
27
- def __init__(self, vllm_config: VllmConfig, dtype: torch.dtype,
28
- num_mamba_layers: int, conv_state_shape: tuple[int, int],
29
- temporal_state_shape: tuple[int, int]):
27
+ def __init__(self, vllm_config: VllmConfig, num_mamba_layers: int,
28
+ conv_state_shape: tuple[int, int],
29
+ temporal_state_shape: tuple[int, int],
30
+ conv_state_dtype: torch.dtype,
31
+ temporal_state_dtype: torch.dtype):
32
+
33
+ self.conv_state_dtype = conv_state_dtype
34
+ self.temporal_state_dtype = temporal_state_dtype
35
36
# Determine max batch size to set size of MambaCache
37
max_batch_size = vllm_config.scheduler_config.max_num_seqs
@@ -40,11 +45,11 @@ class MambaCacheManager(ConstantSizeCache):
45
assert conv_state_shape[0] > conv_state_shape[1]
46
conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
47
(conv_state_shape[1], conv_state_shape[0]),
43
- dtype=dtype,
48
+ dtype=self.conv_state_dtype,
49
device="cuda").transpose(-1, -2)
50
temporal_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
51
temporal_state_shape,
47
- dtype=dtype,
52
+ dtype=self.temporal_state_dtype,
53
device="cuda")
54
55
self._mamba_cache = (conv_state, temporal_state)
vllm/model_executor/models/minimax_text_01.py
+31
-3
@@ -16,7 +16,8 @@ from transformers import MiniMaxConfig
16
17
from vllm import envs
18
from vllm.attention import Attention, AttentionMetadata
19
-from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
19
+from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
20
+ get_current_vllm_config)
21
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
22
from vllm.distributed.parallel_state import (
23
get_pp_group, get_tensor_model_parallel_rank,
@@ -36,7 +37,7 @@ 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.abstract import MambaBase
39
from vllm.model_executor.layers.mamba.mamba_utils import (
39
- MambaStateShapeCalculator)
40
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
41
from vllm.model_executor.layers.quantization.base_config import (
42
QuantizationConfig)
43
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -338,6 +339,12 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
339
def mamba_type(self) -> str:
340
return "linear_attention"
341
342
+ def get_state_dtype(self) -> tuple[torch.dtype]:
343
+ return MambaStateDtypeCalculator.linear_attention_state_dtype(
344
+ self.model_config.dtype,
345
+ self.cache_config.mamba_cache_dtype,
346
+ )
347
+
348
def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
349
return MambaStateShapeCalculator.linear_attention_state_shape(
350
num_heads=self.num_heads,
@@ -353,6 +360,8 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
360
max_position: int,
361
block_size: int,
362
num_hidden_layer: int,
363
+ model_config: Optional[ModelConfig] = None,
364
+ cache_config: Optional[CacheConfig] = None,
365
quant_config: Optional[QuantizationConfig] = None,
366
layer_idx: int = 0,
367
linear_layer_idx: int = 0,
@@ -374,6 +383,8 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
383
self.tp_heads = self.total_num_heads // self.tp_size
384
self.qkv_size = self.num_heads * self.head_dim
385
self.tp_hidden = self.head_dim * self.tp_heads
386
+ self.model_config = model_config
387
+ self.cache_config = cache_config
388
self.prefix = prefix
389
390
self.qkv_proj = ColumnParallelLinear(
@@ -657,6 +668,7 @@ class MiniMaxText01DecoderLayer(nn.Module):
668
def __init__(
669
self,
670
config: MiniMaxConfig,
671
+ model_config: Optional[ModelConfig] = None,
672
cache_config: Optional[CacheConfig] = None,
673
quant_config: Optional[QuantizationConfig] = None,
674
expert_num: int = 1,
@@ -693,6 +705,8 @@ class MiniMaxText01DecoderLayer(nn.Module):
705
max_position=max_position_embeddings,
706
block_size=config.block if hasattr(config, "block") else 256,
707
num_hidden_layer=config.num_hidden_layers,
708
+ model_config=model_config,
709
+ cache_config=cache_config,
710
quant_config=quant_config,
711
layer_idx=self._ilayer,
712
linear_layer_idx=linear_layer_id,
@@ -861,6 +875,7 @@ class MiniMaxText01Model(nn.Module):
875
def __init__(
876
self,
877
config: MiniMaxConfig,
878
+ model_config: Optional[ModelConfig] = None,
879
quant_config: Optional[QuantizationConfig] = None,
880
cache_config: Optional[CacheConfig] = None,
881
scheduler_config=None,
@@ -910,6 +925,7 @@ class MiniMaxText01Model(nn.Module):
925
decoder_kwargs = {
926
"quant_config": quant_config,
927
"layer_id": layer_idx,
928
+ "model_config": model_config,
929
"cache_config": cache_config
930
}
931
@@ -1111,8 +1127,9 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
1127
self.config.max_model_len = vllm_config.model_config.max_model_len
1128
self.model = MiniMaxText01Model(
1129
self.config,
1114
- quant_config,
1130
+ model_config=vllm_config.model_config,
1131
cache_config=vllm_config.cache_config,
1132
+ quant_config=quant_config,
1133
scheduler_config=vllm_config.scheduler_config,
1134
prefix=maybe_prefix(prefix, "model"))
1135
if get_pp_group().is_last_rank:
@@ -1409,6 +1426,17 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
1426
load_basic_weight(name, loaded_weight, self)
1427
return loaded_params
1428
1429
+ @classmethod
1430
+ def get_mamba_state_dtype_from_config(
1431
+ cls,
1432
+ vllm_config: "VllmConfig",
1433
+ ) -> tuple[torch.dtype, torch.dtype]:
1434
+
1435
+ return MambaStateDtypeCalculator.linear_attention_state_dtype(
1436
+ vllm_config.model_config.dtype,
1437
+ vllm_config.cache_config.mamba_cache_dtype,
1438
+ )
1439
+
1440
@classmethod
1441
def get_mamba_state_shape_from_config(
1442
cls,
vllm/model_executor/models/nemotron_h.py
+28
-4
@@ -26,7 +26,7 @@ from torch import nn
26
from vllm import envs
27
from vllm.attention.layer import Attention
28
from vllm.compilation.decorators import support_torch_compile
29
-from vllm.config import CacheConfig, VllmConfig
29
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
30
from vllm.distributed import get_tensor_model_parallel_world_size
31
from vllm.distributed.parallel_state import get_pp_group
32
from vllm.forward_context import get_forward_context
@@ -40,7 +40,7 @@ from vllm.model_executor.layers.mamba.mamba2_metadata import (
40
Mamba2Metadata, prepare_mamba2_metadata)
41
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
42
from vllm.model_executor.layers.mamba.mamba_utils import (
43
- MambaStateShapeCalculator)
43
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
44
from vllm.model_executor.layers.quantization import QuantizationConfig
45
from vllm.model_executor.layers.vocab_parallel_embedding import (
46
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
@@ -110,6 +110,7 @@ class NemotronHMLPDecoderLayer(nn.Module):
110
self,
111
config: NemotronHConfig,
112
layer_idx: int,
113
+ model_config: Optional[ModelConfig] = None,
114
cache_config: Optional[CacheConfig] = None,
115
quant_config: Optional[QuantizationConfig] = None,
116
prefix: str = "",
@@ -149,6 +150,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
150
self,
151
config: NemotronHConfig,
152
layer_idx: int,
153
+ model_config: Optional[ModelConfig] = None,
154
cache_config: Optional[CacheConfig] = None,
155
quant_config: Optional[QuantizationConfig] = None,
156
prefix: str = "",
@@ -167,6 +169,8 @@ class NemotronHMambaDecoderLayer(nn.Module):
169
head_dim=config.mamba_head_dim,
170
rms_norm_eps=config.rms_norm_eps,
171
activation=config.mamba_hidden_act,
172
+ model_config=model_config,
173
+ cache_config=cache_config,
174
quant_config=quant_config,
175
prefix=f"{prefix}.mixer",
176
)
@@ -198,6 +202,7 @@ class NemotronHAttention(nn.Module):
202
self,
203
config: NemotronHConfig,
204
layer_idx: int,
205
+ model_config: Optional[ModelConfig] = None,
206
cache_config: Optional[CacheConfig] = None,
207
quant_config: Optional[QuantizationConfig] = None,
208
prefix: str = "",
@@ -270,6 +275,7 @@ class NemotronHAttentionDecoderLayer(nn.Module):
275
self,
276
config: NemotronHConfig,
277
layer_idx: int,
278
+ model_config: Optional[ModelConfig] = None,
279
cache_config: Optional[CacheConfig] = None,
280
quant_config: Optional[QuantizationConfig] = None,
281
prefix: str = "",
@@ -279,6 +285,7 @@ class NemotronHAttentionDecoderLayer(nn.Module):
285
self.mixer = NemotronHAttention(
286
config,
287
layer_idx,
288
+ model_config,
289
cache_config,
290
quant_config,
291
prefix=f"{prefix}.mixer",
@@ -317,6 +324,7 @@ class NemotronHModel(nn.Module):
324
super().__init__()
325
326
config: NemotronHConfig = vllm_config.model_config.hf_config
327
+ model_config = vllm_config.model_config
328
cache_config = vllm_config.cache_config
329
quant_config = vllm_config.quant_config
330
lora_config = vllm_config.lora_config
@@ -340,6 +348,7 @@ class NemotronHModel(nn.Module):
348
return layer_class(
349
config,
350
layer_idx,
351
+ model_config,
352
cache_config,
353
quant_config=quant_config,
354
prefix=prefix,
@@ -478,6 +487,18 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
487
}
488
embedding_padding_modules = ["lm_head"]
489
490
+ @classmethod
491
+ def get_mamba_state_dtype_from_config(
492
+ cls,
493
+ vllm_config: "VllmConfig",
494
+ ) -> tuple[torch.dtype, torch.dtype]:
495
+
496
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
497
+ vllm_config.model_config.dtype,
498
+ vllm_config.cache_config.mamba_cache_dtype,
499
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
500
+ )
501
+
502
@classmethod
503
def get_mamba_state_shape_from_config(
504
cls,
@@ -569,10 +590,13 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
590
mamba_state_shape = \
591
self.get_mamba_state_shape_from_config(
592
self.vllm_config, use_v1=False)
593
+ mamba_state_dtype = \
594
+ self.get_mamba_state_dtype_from_config(
595
+ self.vllm_config)
596
self.mamba_cache = MambaCacheManager(self.vllm_config,
573
- self.lm_head.weight.dtype,
597
num_mamba_layers,
575
- *mamba_state_shape)
598
+ *mamba_state_shape,
599
+ *mamba_state_dtype)
600
601
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
602
vllm/model_executor/models/zamba2.py
+33
-5
@@ -18,7 +18,7 @@ from transformers import Zamba2Config
18
from vllm import envs
19
from vllm.attention.layer import Attention
20
from vllm.compilation.decorators import support_torch_compile
21
-from vllm.config import CacheConfig, VllmConfig
21
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
22
from vllm.distributed import get_tensor_model_parallel_world_size
23
from vllm.forward_context import get_forward_context
24
from vllm.model_executor.layers.activation import GeluAndMul
@@ -33,7 +33,7 @@ from vllm.model_executor.layers.mamba.mamba2_metadata import (
33
Mamba2Metadata, prepare_mamba2_metadata)
34
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
35
from vllm.model_executor.layers.mamba.mamba_utils import (
36
- MambaStateShapeCalculator)
36
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
37
from vllm.model_executor.layers.quantization import QuantizationConfig
38
from vllm.model_executor.layers.rotary_embedding import get_rope
39
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -478,6 +478,8 @@ class Zamba2MambaDecoderLayer(nn.Module):
478
479
def __init__(self,
480
config: Zamba2Config,
481
+ model_config: Optional[ModelConfig] = None,
482
+ cache_config: Optional[CacheConfig] = None,
483
quant_config: Optional[QuantizationConfig] = None,
484
prefix: str = "") -> None:
485
"""Initialize the Mamba decoder layer.
@@ -502,6 +504,8 @@ class Zamba2MambaDecoderLayer(nn.Module):
504
config.n_mamba_heads,
505
rms_norm_eps=config.rms_norm_eps,
506
activation="silu",
507
+ model_config=model_config,
508
+ cache_config=cache_config,
509
quant_config=quant_config,
510
prefix=f"{prefix}.mixer")
511
@@ -578,6 +582,8 @@ class Zamba2HybridLayer(nn.Module):
582
shared_transformer: Zamba2AttentionDecoderLayer,
583
config: Zamba2Config,
584
block_idx: int,
585
+ model_config: Optional[ModelConfig] = None,
586
+ cache_config: Optional[CacheConfig] = None,
587
quant_config: Optional[QuantizationConfig] = None,
588
prefix: str = "",
589
) -> None:
@@ -596,6 +602,8 @@ class Zamba2HybridLayer(nn.Module):
602
bias=False,
603
quant_config=quant_config)
604
self.mamba_decoder = Zamba2MambaDecoderLayer(config,
605
+ model_config=model_config,
606
+ cache_config=cache_config,
607
quant_config=quant_config,
608
prefix=prefix)
609
@@ -669,6 +677,7 @@ class Zamba2Model(nn.Module):
677
super().__init__()
678
679
config = vllm_config.model_config.hf_config
680
+ model_config = vllm_config.model_config
681
cache_config = vllm_config.cache_config
682
quant_config = vllm_config.quant_config
683
lora_config = vllm_config.lora_config
@@ -718,11 +727,15 @@ class Zamba2Model(nn.Module):
727
Zamba2HybridLayer(block,
728
config,
729
block_idx,
721
- quant_config,
730
+ model_config=model_config,
731
+ cache_config=cache_config,
732
+ quant_config=quant_config,
733
prefix=prefix))
734
else:
735
layers.append(
736
Zamba2MambaDecoderLayer(config,
737
+ model_config=model_config,
738
+ cache_config=cache_config,
739
quant_config=quant_config,
740
prefix=prefix))
741
self.layers = nn.ModuleList(layers)
@@ -848,6 +861,18 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
861
"1.weight": "B.weight",
862
})
863
864
+ @classmethod
865
+ def get_mamba_state_dtype_from_config(
866
+ cls,
867
+ vllm_config: "VllmConfig",
868
+ ) -> tuple[torch.dtype, torch.dtype]:
869
+
870
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
871
+ vllm_config.model_config.dtype,
872
+ vllm_config.cache_config.mamba_cache_dtype,
873
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
874
+ )
875
+
876
@classmethod
877
def get_mamba_state_shape_from_config(
878
cls,
@@ -966,10 +991,13 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
991
mamba_state_shape = \
992
self.get_mamba_state_shape_from_config(
993
self.vllm_config, use_v1=False)
994
+ mamba_state_dtype = \
995
+ self.get_mamba_state_dtype_from_config(
996
+ self.vllm_config)
997
self.mamba_cache = MambaCacheManager(self.vllm_config,
970
- self.lm_head.weight.dtype,
998
num_mamba_layers,
972
- *mamba_state_shape)
999
+ *mamba_state_shape,
1000
+ *mamba_state_dtype)
1001
1002
# Get cache parameters for current run
1003
mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
vllm/utils/__init__.py
+1
@@ -173,6 +173,7 @@ CYAN = '\033[1;36m'
173
RESET = '\033[0;0m'
174
175
STR_DTYPE_TO_TORCH_DTYPE = {
176
+ "float32": torch.float32,
177
"half": torch.half,
178
"bfloat16": torch.bfloat16,
179
"float": torch.float,
vllm/v1/kv_cache_interface.py
+4
-3
@@ -182,14 +182,15 @@ class SlidingWindowSpec(AttentionSpec):
182
@dataclass(frozen=True)
183
class MambaSpec(KVCacheSpec):
184
shapes: tuple[tuple[int, ...], ...]
185
- dtype: torch.dtype
185
+ dtypes: tuple[torch.dtype]
186
page_size_padded: Optional[int] = None
187
mamba_type: str = "mamba2"
188
189
@property
190
def page_size_bytes(self) -> int:
191
- num_elements = sum(prod(shape) for shape in self.shapes)
192
- page_size = num_elements * get_dtype_size(self.dtype)
191
+ page_size = sum(
192
+ prod(shape) * get_dtype_size(dtype)
193
+ for (shape, dtype) in zip(self.shapes, self.dtypes))
194
if self.page_size_padded is not None:
195
assert self.page_size_padded >= page_size
196
return self.page_size_padded
vllm/v1/worker/gpu_model_runner.py
+10
-8
@@ -2884,23 +2884,25 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2884
elif isinstance(kv_cache_spec, MambaSpec):
2885
has_mamba = True
2886
raw_tensor = kv_cache_raw_tensors[layer_name]
2887
- dtype = kv_cache_spec.dtype
2888
- num_element_per_page = (kv_cache_spec.page_size_bytes //
2889
- get_dtype_size(dtype))
2887
state_tensors = []
2891
- storage_offset = 0
2892
- for shape in kv_cache_spec.shapes:
2888
+ storage_offset_bytes = 0
2889
+ for (shape, dtype) in zip(kv_cache_spec.shapes,
2890
+ kv_cache_spec.dtypes):
2891
+ dtype_size = get_dtype_size(dtype)
2892
+ num_element_per_page = (
2893
+ kv_cache_spec.page_size_bytes // dtype_size)
2894
target_shape = (num_blocks, *shape)
2895
stride = torch.empty(target_shape).stride()
2896
target_stride = (num_element_per_page, *stride[1:])
2897
+ assert storage_offset_bytes % dtype_size == 0
2898
tensor = torch.as_strided(
2899
raw_tensor.view(dtype),
2900
size=target_shape,
2901
stride=target_stride,
2900
- storage_offset=storage_offset,
2902
+ storage_offset=storage_offset_bytes // dtype_size,
2903
)
2904
state_tensors.append(tensor)
2903
- storage_offset += stride[0]
2905
+ storage_offset_bytes += stride[0] * dtype_size
2906
2907
kv_caches[layer_name] = state_tensors
2908
else:
@@ -3087,7 +3089,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3089
for layer_name, mamba_module in mamba_layers.items():
3090
kv_cache_spec[layer_name] = MambaSpec(
3091
shapes=mamba_module.get_state_shape(),
3090
- dtype=self.kv_cache_dtype,
3092
+ dtypes=mamba_module.get_state_dtype(),
3093
block_size=max_model_len,
3094
page_size_padded=page_size_padded,
3095
mamba_type=mamba_module.mamba_type)