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] - Mamba1 Attention Metadata (#21249)
Signed-off-by: asafg <asafg@ai21.com> Co-authored-by: asafg <asafg@ai21.com>
Asaf Joseph Gardin committed
Aug 7, 2025 at 03:03 UTC
46a13949d5f64e4a40bac3cb30eb0f867074f741
19 files changed
+377
-171
csrc/mamba/mamba_ssm/selective_scan.h
+3
@@ -45,6 +45,9 @@ struct SSMParamsBase {
45
index_t out_d_stride;
46
index_t out_z_batch_stride;
47
index_t out_z_d_stride;
48
+ index_t ssm_states_batch_stride;
49
+ index_t ssm_states_dim_stride;
50
+ index_t ssm_states_dstate_stride;
51
52
// Common data pointers.
53
void *__restrict__ A_ptr;
csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+14
-4
@@ -132,8 +132,10 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
132
input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + sequence_start_index * params.B_batch_stride + group_id * params.B_group_stride;
133
weight_t *C = reinterpret_cast<weight_t *>(params.C_ptr) + dim_id * kNRows * params.C_d_stride;
134
input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + sequence_start_index * params.C_batch_stride + group_id * params.C_group_stride;
135
- input_t *ssm_states = reinterpret_cast<input_t *>(params.ssm_states_ptr) + (cache_index * params.dim + dim_id * kNRows) * params.dstate;
136
-
135
+ input_t *ssm_states = reinterpret_cast<input_t *>(params.ssm_states_ptr) +
136
+ cache_index * params.ssm_states_batch_stride +
137
+ dim_id * kNRows * params.ssm_states_dim_stride;
138
+
139
float D_val[kNRows] = {0};
140
if (params.D_ptr != nullptr) {
141
#pragma unroll
@@ -248,7 +250,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
250
}
251
// Initialize running total
252
251
- scan_t running_prefix = chunk > 0 ? smem_running_prefix[state_idx + r * MAX_DSTATE] : make_float2(1.0, has_initial_state ? float(ssm_states[state_idx]): 0.0);
253
+ scan_t running_prefix = chunk > 0 ? smem_running_prefix[state_idx + r * MAX_DSTATE] : make_float2(1.0, has_initial_state ? float(ssm_states[state_idx * params.ssm_states_dstate_stride]): 0.0);
254
255
SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
256
typename Ktraits::BlockScanT(smem_scan).InclusiveScan(
@@ -259,7 +261,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
261
if (threadIdx.x == 0) {
262
smem_running_prefix[state_idx] = prefix_op.running_prefix;
263
if (chunk == n_chunks - 1) {
262
- ssm_states[state_idx] = input_t(prefix_op.running_prefix.y);
264
+ ssm_states[state_idx * params.ssm_states_dstate_stride] = input_t(prefix_op.running_prefix.y);
265
}
266
}
267
#pragma unroll
@@ -481,6 +483,10 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
483
params.out_batch_stride = out.stride(1);
484
params.out_d_stride = out.stride(0);
485
486
+ params.ssm_states_batch_stride = ssm_states.stride(0);
487
+ params.ssm_states_dim_stride = ssm_states.stride(1);
488
+ params.ssm_states_dstate_stride = ssm_states.stride(2);
489
+
490
}
491
else{
492
if (!is_variable_B) {
@@ -509,6 +515,10 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
515
}
516
params.out_batch_stride = out.stride(0);
517
params.out_d_stride = out.stride(1);
518
+
519
+ params.ssm_states_batch_stride = ssm_states.stride(0);
520
+ params.ssm_states_dim_stride = ssm_states.stride(1);
521
+ params.ssm_states_dstate_stride = ssm_states.stride(2);
522
}
523
}
524
docs/models/supported_models.md
+2
-2
@@ -370,9 +370,9 @@ th {
370
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ | ✅︎ |
371
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
372
| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ | ✅︎ |
373
-| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ | |
373
+| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ | ✅︎ |
374
| `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA, Yi | `meta-llama/Meta-Llama-3.1-405B-Instruct`, `meta-llama/Meta-Llama-3.1-70B`, `meta-llama/Meta-Llama-3-70B-Instruct`, `meta-llama/Llama-2-70b-hf`, `01-ai/Yi-34B`, etc. | ✅︎ | ✅︎ | ✅︎ |
375
-| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ | |
375
+| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ | ✅︎ |
376
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ | ✅︎ |
377
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ | ✅︎ |
378
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ | ✅︎ |
docs/usage/v1_guide.md
+5
-7
@@ -83,7 +83,7 @@ based on assigned priority, with FCFS as a tie-breaker), configurable via the
83
| **Decoder-only Models** | <nobr>🚀 Optimized</nobr> |
84
| **Encoder-Decoder Models** | <nobr>🟠 Delayed</nobr> |
85
| **Embedding Models** | <nobr>🟢 Functional</nobr> |
86
-| **Mamba Models** | <nobr>🟢 (Mamba-2), 🟡 (Mamba-1)</nobr> |
86
+| **Mamba Models** | <nobr>🟢 (Mamba-2), 🟢 (Mamba-1)</nobr> |
87
| **Multimodal Models** | <nobr>🟢 Functional</nobr> |
88
89
vLLM V1 currently excludes model architectures with the `SupportsV0Only` protocol.
@@ -104,13 +104,11 @@ to enable simultaneous generation and embedding using the same engine instance i
104
105
#### Mamba Models
106
107
-Models using selective state-space mechanisms instead of standard transformer attention are partially supported.
108
-Models that use Mamba-2 layers (e.g., `Mamba2ForCausalLM`) are supported, but models that use older Mamba-1 layers
109
-(e.g., `MambaForCausalLM`, `JambaForCausalLM`) are not yet supported. Please note that these models currently require
110
-disabling prefix caching in V1.
107
+Models using selective state-space mechanisms instead of standard transformer attention are supported.
108
+Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`) are supported. Please note that these models currently require disabling prefix caching in V1. Additionally, Mamba-1 models require `enforce_eager=True`.
109
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
110
+Models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
111
+`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`). Please note that
112
these models currently require disabling prefix caching and using the FlashInfer attention backend in V1.
113
114
#### Encoder-Decoder Models
tests/models/language/generation/test_hybrid.py
+2
@@ -53,6 +53,8 @@ HF_UNSUPPORTED_MODELS = [
53
]
54
55
V1_SUPPORTED_MODELS = [
56
+ "state-spaces/mamba-130m-hf",
57
+ "ai21labs/Jamba-tiny-dev",
58
"mistralai/Mamba-Codestral-7B-v0.1",
59
"ibm-ai-platform/Bamba-9B-v1",
60
"Zyphra/Zamba2-1.2B-instruct",
tests/v1/test_oracle.py
-1
@@ -12,7 +12,6 @@ from vllm.engine.async_llm_engine import AsyncLLMEngine
12
UNSUPPORTED_MODELS_V1 = [
13
"openai/whisper-large-v3", # transcription
14
"facebook/bart-large-cnn", # encoder decoder
15
- "state-spaces/mamba-130m-hf", # mamba1
15
]
16
17
MODEL = "meta-llama/Llama-3.2-1B-Instruct"
vllm/model_executor/layers/mamba/mamba_mixer.py
+108
-36
@@ -1,30 +1,37 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from typing import Optional
5
+
6
import torch
7
from torch import nn
8
from torch.nn.parameter import Parameter
9
8
-from vllm.attention.backends.abstract import AttentionMetadata
10
+from vllm import envs
11
+from vllm.config import get_current_vllm_config
12
from vllm.distributed.parallel_state import (
13
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
11
-from vllm.forward_context import get_forward_context
14
+from vllm.forward_context import ForwardContext, get_forward_context
15
from vllm.model_executor.custom_op import CustomOp
16
from vllm.model_executor.layers.layernorm import RMSNorm
17
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
18
MergedColumnParallelLinear,
19
RowParallelLinear)
20
+from vllm.model_executor.layers.mamba.abstract import MambaBase
21
+from vllm.model_executor.layers.mamba.mamba_utils import (
22
+ MambaStateShapeCalculator)
23
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
24
causal_conv1d_fn, causal_conv1d_update)
25
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
26
selective_scan_fn, selective_state_update)
27
from vllm.model_executor.models.mamba_cache import MambaCacheParams
28
from vllm.model_executor.utils import set_weight_attrs
29
+from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
30
31
32
# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
33
@CustomOp.register("mamba_mixer")
27
-class MambaMixer(CustomOp):
34
+class MambaMixer(MambaBase, CustomOp):
35
"""
36
Compute ∆, A, B, C, and D the state space parameters and compute
37
the `contextualized_states`. A, D are input independent
@@ -47,13 +54,16 @@ class MambaMixer(CustomOp):
54
rms_norm_has_weight: bool = True,
55
rms_norm_eps: float = 1e-5,
56
activation="silu",
50
- is_lora_enabled: bool = False):
57
+ is_lora_enabled: bool = False,
58
+ prefix: str = ""):
59
super().__init__()
60
self.time_step_rank = time_step_rank
61
self.ssm_state_size = ssm_state_size
62
self.use_rms_norm = use_rms_norm
63
self.activation = activation
64
self.is_lora_enabled = is_lora_enabled
65
+ self.conv_kernel_size = conv_kernel_size
66
+ self.intermediate_size = intermediate_size
67
68
self.conv1d = ColumnParallelLinear(
69
input_size=conv_kernel_size,
@@ -131,14 +141,62 @@ class MambaMixer(CustomOp):
141
has_weight=rms_norm_has_weight,
142
) if use_rms_norm else None
143
134
- def forward_native(self, hidden_states: torch.Tensor,
135
- conv_state: torch.Tensor, ssm_state: torch.Tensor):
144
+ if envs.VLLM_USE_V1:
145
+ compilation_config = get_current_vllm_config().compilation_config
146
+ if prefix in compilation_config.static_forward_context:
147
+ raise ValueError(f"Duplicate layer name: {prefix}")
148
+ compilation_config.static_forward_context[prefix] = self
149
+ # The outer list is for v0 PP virtual engine. Though this code path
150
+ # only runs for v1, we have to do this to unify with the interface
151
+ # of Attention + v0 PP.
152
+ # The inner tuple is (conv_state, ssm_state)
153
+ self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
154
+
155
+ self.prefix = prefix
156
+
157
+ def forward(self,
158
+ hidden_states: torch.Tensor,
159
+ mamba_cache_params: Optional[MambaCacheParams] = None):
160
+ if not envs.VLLM_USE_V1:
161
+ return CustomOp.forward(self, hidden_states, mamba_cache_params)
162
+ else:
163
+ return self.forward_cuda(hidden_states, mamba_cache_params)
164
+
165
+ def forward_native(self,
166
+ hidden_states: torch.Tensor,
167
+ mamba_cache_params: Optional[MambaCacheParams] = None):
168
pass
169
138
- def forward_cuda(self, hidden_states: torch.Tensor,
139
- mamba_cache_params: MambaCacheParams):
170
+ def forward_cuda(self,
171
+ hidden_states: torch.Tensor,
172
+ mamba_cache_params: Optional[MambaCacheParams] = None):
173
+
174
+ forward_context: ForwardContext = get_forward_context()
175
+ attn_metadata = forward_context.attn_metadata
176
+
177
+ if envs.VLLM_USE_V1:
178
+ if attn_metadata is not None:
179
+ assert isinstance(attn_metadata, dict)
180
+ attn_metadata = attn_metadata[self.prefix]
181
+ mamba1_metadata = attn_metadata
182
+ assert isinstance(mamba1_metadata, Mamba1AttentionMetadata)
183
+ query_start_loc = mamba1_metadata.query_start_loc
184
+ state_indices_tensor = mamba1_metadata.state_indices_tensor
185
+ self_kv_cache = self.kv_cache[forward_context.virtual_engine]
186
+ conv_state = self_kv_cache[0].transpose(-1, -2)
187
+ ssm_state = self_kv_cache[1]
188
+ has_initial_state = mamba1_metadata.has_initial_states
189
+ context_lens_tensor = mamba1_metadata.context_lens_tensor
190
+ else:
191
+ assert mamba_cache_params is not None
192
+ conv_state = mamba_cache_params.conv_state
193
+ ssm_state = mamba_cache_params.ssm_state
194
+ state_indices_tensor = mamba_cache_params.state_indices_tensor
195
+ query_start_loc = attn_metadata.query_start_loc
196
+ context_lens_tensor = attn_metadata.context_lens_tensor
197
141
- attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
198
+ if context_lens_tensor is not None:
199
+ has_initial_state = context_lens_tensor > 0
200
201
# 1. Gated MLP's linear projection
202
projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)
@@ -148,8 +206,12 @@ class MambaMixer(CustomOp):
206
conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
207
self.conv1d.weight.size(2))
208
151
- if attn_metadata.query_start_loc is not None \
152
- and attn_metadata.context_lens_tensor is not None:
209
+ if envs.VLLM_USE_V1 and attn_metadata is None:
210
+ # V1 profile run
211
+ hidden_states = hidden_states.contiguous()
212
+ return self.out_proj(hidden_states.transpose(-2, -1))[0]
213
+
214
+ if query_start_loc is not None and context_lens_tensor is not None:
215
# |---------- N-1 iteration --------|
216
# |---------------- N iteration ---------------------|
217
# |- tokenA -|......................|-- newTokens ---|
@@ -161,18 +223,18 @@ class MambaMixer(CustomOp):
223
conv_weights,
224
bias=self.conv1d.bias,
225
activation=self.activation,
164
- conv_states=mamba_cache_params.conv_state,
165
- has_initial_state=attn_metadata.context_lens_tensor > 0,
166
- cache_indices=mamba_cache_params.state_indices_tensor,
167
- query_start_loc=attn_metadata.query_start_loc)
226
+ conv_states=conv_state,
227
+ has_initial_state=has_initial_state,
228
+ cache_indices=state_indices_tensor,
229
+ query_start_loc=query_start_loc)
230
else:
231
hidden_states = causal_conv1d_update(
232
hidden_states.transpose(0, 1),
171
- mamba_cache_params.conv_state,
233
+ conv_state,
234
conv_weights,
235
self.conv1d.bias,
236
self.activation,
175
- conv_state_indices=mamba_cache_params.state_indices_tensor)
237
+ conv_state_indices=state_indices_tensor)
238
hidden_states = hidden_states.transpose(0, 1)
239
240
# 3. State Space Model sequence transformation
@@ -203,11 +265,10 @@ class MambaMixer(CustomOp):
265
time_proj_bias = (self.dt_proj.bias.float() if hasattr(
266
self.dt_proj, "bias") else None)
267
206
- if attn_metadata.query_start_loc is not None \
207
- and attn_metadata.context_lens_tensor is not None:
268
+ if query_start_loc is not None and context_lens_tensor is not None:
269
scan_outputs = selective_scan_fn(
270
hidden_states,
210
- mamba_cache_params.ssm_state,
271
+ ssm_state,
272
discrete_time_step,
273
self.A,
274
B.transpose(-2, -1),
@@ -216,24 +277,23 @@ class MambaMixer(CustomOp):
277
gate,
278
time_proj_bias,
279
delta_softplus=True,
219
- cache_indices=mamba_cache_params.state_indices_tensor,
220
- has_initial_state=attn_metadata.context_lens_tensor > 0,
221
- query_start_loc=attn_metadata.query_start_loc)
280
+ cache_indices=state_indices_tensor,
281
+ has_initial_state=has_initial_state,
282
+ query_start_loc=query_start_loc)
283
else:
284
scan_outputs = torch.empty_like(hidden_states.transpose(0, 1))
224
- selective_state_update(
225
- mamba_cache_params.ssm_state,
226
- hidden_states.transpose(0, 1),
227
- discrete_time_step.transpose(0, 1),
228
- self.A,
229
- B,
230
- C,
231
- self.D,
232
- gate.transpose(0, 1),
233
- time_proj_bias,
234
- dt_softplus=True,
235
- state_batch_indices=mamba_cache_params.state_indices_tensor,
236
- out=scan_outputs)
285
+ selective_state_update(ssm_state,
286
+ hidden_states.transpose(0, 1),
287
+ discrete_time_step.transpose(0, 1),
288
+ self.A,
289
+ B,
290
+ C,
291
+ self.D,
292
+ gate.transpose(0, 1),
293
+ time_proj_bias,
294
+ dt_softplus=True,
295
+ state_batch_indices=state_indices_tensor,
296
+ out=scan_outputs)
297
scan_outputs = scan_outputs.transpose(0, 1)
298
299
# 4. Final linear projection
@@ -245,3 +305,15 @@ class MambaMixer(CustomOp):
305
contextualized_states = self.out_proj(
306
scan_outputs.transpose(-2, -1))[0]
307
return contextualized_states
308
+
309
+ def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
310
+ return MambaStateShapeCalculator.mamba1_state_shape(
311
+ tp_world_size=get_tensor_model_parallel_world_size(),
312
+ intermediate_size=self.intermediate_size,
313
+ state_size=self.ssm_state_size,
314
+ conv_kernel=self.conv_kernel_size,
315
+ )
316
+
317
+ @property
318
+ def mamba_type(self) -> str:
319
+ return "mamba1"
vllm/model_executor/layers/mamba/mamba_mixer2.py
+4
-3
@@ -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
- extra_groups_for_head_shards, get_mamba_state_shape)
24
+ 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
@@ -278,8 +278,9 @@ class MambaMixer2(MambaBase, CustomOp):
278
# - for TP we shard conv_dim by sharding on n_groups,
279
# - but if n_groups cannot divide tp_size, we need to
280
# extend some extra groups
281
- self.n_groups = n_groups + extra_groups_for_head_shards(
281
+ groups = MambaStateShapeCalculator.extra_groups_for_head_shards(
282
n_groups, self.tp_size)
283
+ self.n_groups = n_groups + groups
284
285
self.conv_dim = intermediate_size + 2 * self.n_groups * ssm_state_size
286
self.conv1d = ColumnParallelLinear(
@@ -732,7 +733,7 @@ class MambaMixer2(MambaBase, CustomOp):
733
output[:num_actual_tokens], _ = self.out_proj(hidden_states)
734
735
def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
735
- return get_mamba_state_shape(
736
+ return MambaStateShapeCalculator.mamba2_state_shape(
737
intermediate_size=self.intermediate_size,
738
tp_world_size=get_tensor_model_parallel_world_size(),
739
n_groups=self.n_groups,
vllm/model_executor/layers/mamba/mamba_utils.py
+67
-50
@@ -3,53 +3,70 @@
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
6
+class MambaStateShapeCalculator:
7
+
8
+ @classmethod
9
+ def mamba1_state_shape(
10
+ cls,
11
+ tp_world_size: int,
12
+ intermediate_size: int,
13
+ state_size: int,
14
+ conv_kernel: int,
15
+ use_v1: bool = True,
16
+ ) -> tuple[tuple[int, int], tuple[int, int]]:
17
+ conv_state_shape = (divide(intermediate_size,
18
+ tp_world_size), conv_kernel - 1)
19
+
20
+ temporal_state_shape = (divide(intermediate_size,
21
+ tp_world_size), state_size)
22
+
23
+ # In V0, the conv_state shape was swapped during allocation in
24
+ # MambaCacheManager, but in V1 it needs to be determined here at the
25
+ # calculation level
26
+ if use_v1:
27
+ conv_state_shape = conv_state_shape[1], conv_state_shape[0]
28
+
29
+ return conv_state_shape, temporal_state_shape
30
+
31
+ @classmethod
32
+ def mamba2_state_shape(
33
+ cls,
34
+ tp_world_size: int,
35
+ intermediate_size: int,
36
+ n_groups: int,
37
+ num_heads: int,
38
+ head_dim: int,
39
+ state_size: int,
40
+ conv_kernel: int,
41
+ use_v1: bool = True,
42
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
43
+ # if n_groups is not divisible by world_size, need to extend the shards
44
+ # to ensure all groups needed by a head is sharded along with it
45
+ n_groups = n_groups + cls.extra_groups_for_head_shards(
46
+ n_groups, tp_world_size)
47
+ # heads and n_groups are TP-ed
48
+ conv_dim = intermediate_size + 2 * n_groups * state_size
49
+
50
+ # contiguous along 'dim' axis
51
+ conv_state_shape = (conv_kernel - 1, divide(conv_dim, tp_world_size))
52
+ if not use_v1:
53
+ conv_state_shape = conv_state_shape[1], conv_state_shape[0]
54
+
55
+ # These are not TP-ed as they depend on A, dt_bias, D
56
+ # - they are typically small
57
+ # e.g., (h_heads, head_dim, state_size) = (128, 64, 128)
58
+ temporal_state_shape = (divide(num_heads,
59
+ tp_world_size), head_dim, state_size)
60
+ return conv_state_shape, temporal_state_shape
61
+
62
+ @classmethod
63
+ def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int):
64
+ """Compute the increase in group numbers to account for
65
+ replication in order to accompany the head shards."""
66
+
67
+ # in the case ngoups % tp_size == 0, this will be zero
68
+ if ngroups % tp_size == 0:
69
+ return 0
70
+
71
+ # for n_groups == 1, this is exactly tp_size - n_groups
72
+ return tp_size - ngroups
vllm/model_executor/models/bamba.py
+3
-2
@@ -25,7 +25,8 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
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 get_mamba_state_shape
28
+from vllm.model_executor.layers.mamba.mamba_utils import (
29
+ 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 (
@@ -457,7 +458,7 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
458
hf_config = vllm_config.model_config.hf_config
459
intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
460
460
- return get_mamba_state_shape(
461
+ return MambaStateShapeCalculator.mamba2_state_shape(
462
intermediate_size=intermediate_size,
463
tp_world_size=parallel_config.tensor_parallel_size,
464
n_groups=hf_config.mamba_n_groups,
vllm/model_executor/models/falcon_h1.py
+3
-2
@@ -24,7 +24,8 @@ 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 MambaMixer2
27
-from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
27
+from vllm.model_executor.layers.mamba.mamba_utils import (
28
+ 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 (
@@ -543,7 +544,7 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
544
if hf_config.mamba_d_ssm is None else
545
hf_config.mamba_d_ssm)
546
546
- return get_mamba_state_shape(
547
+ return MambaStateShapeCalculator.mamba2_state_shape(
548
intermediate_size=intermediate_size,
549
tp_world_size=parallel_config.tensor_parallel_size,
550
n_groups=hf_config.mamba_n_groups,
vllm/model_executor/models/granitemoehybrid.py
+3
-2
@@ -23,7 +23,8 @@ 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 MambaMixer2
26
-from vllm.model_executor.layers.mamba.mamba_utils import get_mamba_state_shape
26
+from vllm.model_executor.layers.mamba.mamba_utils import (
27
+ 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 (
@@ -547,7 +548,7 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
548
hf_config = vllm_config.model_config.hf_config
549
intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
550
550
- return get_mamba_state_shape(
551
+ return MambaStateShapeCalculator.mamba2_state_shape(
552
intermediate_size=intermediate_size,
553
tp_world_size=parallel_config.tensor_parallel_size,
554
n_groups=hf_config.mamba_n_groups,
vllm/model_executor/models/jamba.py
+37
-25
@@ -8,6 +8,7 @@ import torch
8
from torch import nn
9
from transformers import JambaConfig
10
11
+from vllm import envs
12
from vllm.attention.layer import Attention
13
from vllm.config import CacheConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -19,6 +20,8 @@ from vllm.model_executor.layers.linear import (QKVParallelLinear,
20
RowParallelLinear)
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)
25
from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
26
PoolingType)
27
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -32,8 +35,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
35
from vllm.sequence import IntermediateTensors
36
from vllm.utils import LayerBlockType
37
35
-from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,
36
- SupportsV0Only)
38
+from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
39
from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
40
make_empty_intermediate_tensors_factory, make_layers,
41
maybe_prefix)
@@ -112,7 +114,8 @@ class JambaMambaDecoderLayer(nn.Module):
114
use_rms_norm=True,
115
rms_norm_eps=config.rms_norm_eps,
116
activation=config.hidden_act,
115
- is_lora_enabled = self.is_lora_enabled
117
+ is_lora_enabled = self.is_lora_enabled,
118
+ prefix=f"{prefix}.mixer",
119
)
120
121
num_experts = config.layers_num_experts[layer_idx]
@@ -344,7 +347,8 @@ class JambaModel(nn.Module):
347
layer_mamba_cache_params = None
348
if isinstance(layer, JambaAttentionDecoderLayer):
349
kv_cache_index += 1
347
- if isinstance(layer, JambaMambaDecoderLayer):
350
+ if isinstance(layer,
351
+ JambaMambaDecoderLayer) and mamba_cache_params:
352
current_state_layer = mamba_cache_index
353
layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
354
current_state_layer)
@@ -442,7 +446,7 @@ class JambaModel(nn.Module):
446
447
448
class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
445
- IsHybrid, SupportsV0Only):
449
+ IsHybrid):
450
hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={
451
".self_attn.": ".",
452
".A_log": ".A"
@@ -509,14 +513,19 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
513
intermediate_tensors: Optional[IntermediateTensors] = None,
514
inputs_embeds: Optional[torch.Tensor] = None,
515
**kwargs):
512
- if self.mamba_cache is None:
513
- num_mamba_layers = self.model_config.get_num_layers_by_block_type(
514
- self.vllm_config.parallel_config, LayerBlockType.mamba)
515
- self.mamba_cache = MambaCacheManager(
516
- self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,
517
- *self._get_mamba_cache_shape())
518
-
519
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
516
+ # NOTE: mamba_cache_params is not needed for v1
517
+ mamba_cache_params = None
518
+ if not envs.VLLM_USE_V1:
519
+ if self.mamba_cache is None:
520
+ num_layers = self.model_config.get_num_layers_by_block_type(
521
+ self.vllm_config.parallel_config, LayerBlockType.mamba)
522
+ state_shape = self.get_mamba_state_shape_from_config(
523
+ self.vllm_config)
524
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
525
+ self.lm_head.weight.dtype,
526
+ num_layers, *state_shape)
527
+
528
+ mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
529
530
hidden_states = self.model(input_ids, positions, mamba_cache_params,
531
intermediate_tensors, inputs_embeds)
@@ -529,19 +538,22 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
538
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
539
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
540
532
- def _get_mamba_cache_shape(
533
- self) -> tuple[tuple[int, int], tuple[int, int]]:
534
- world_size = get_tensor_model_parallel_world_size()
535
- hidden_size = self.config.hidden_size
536
- conv_state_shape = (
537
- self.config.mamba_expand * hidden_size // world_size,
538
- self.config.mamba_d_conv - 1,
539
- )
540
- temporal_state_shape = (
541
- self.config.mamba_expand * hidden_size // world_size,
542
- self.config.mamba_d_state,
541
+ @classmethod
542
+ def get_mamba_state_shape_from_config(
543
+ cls,
544
+ vllm_config: "VllmConfig",
545
+ ) -> tuple[tuple[int, int], tuple[int, int]]:
546
+ parallel_config = vllm_config.parallel_config
547
+ hf_config = vllm_config.model_config.hf_config
548
+ hidden_size = hf_config.hidden_size
549
+
550
+ return MambaStateShapeCalculator.mamba1_state_shape(
551
+ tp_world_size=parallel_config.tensor_parallel_size,
552
+ intermediate_size=hf_config.mamba_expand * hidden_size,
553
+ state_size=hf_config.mamba_d_state,
554
+ conv_kernel=hf_config.mamba_d_conv,
555
+ use_v1=envs.VLLM_USE_V1,
556
)
544
- return conv_state_shape, temporal_state_shape
557
558
def compute_logits(
559
self,
vllm/model_executor/models/mamba.py
+46
-31
@@ -8,20 +8,21 @@ import torch
8
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.distributed import get_tensor_model_parallel_world_size
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)
19
from vllm.model_executor.layers.quantization.base_config import (
20
QuantizationConfig)
21
from vllm.model_executor.layers.vocab_parallel_embedding import (
22
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
23
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24
from vllm.model_executor.models.interfaces import (HasInnerState,
23
- IsAttentionFree, SupportsPP,
24
- SupportsV0Only)
25
+ IsAttentionFree, SupportsPP)
26
from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
27
MambaCacheParams)
28
from vllm.model_executor.sampling_metadata import SamplingMetadata
@@ -41,7 +42,8 @@ class MambaDecoderLayer(nn.Module):
42
config: MambaConfig,
43
cache_config: Optional[CacheConfig] = None,
44
quant_config: Optional[QuantizationConfig] = None,
44
- is_lora_enabled: Optional[bool] = False) -> None:
45
+ is_lora_enabled: Optional[bool] = False,
46
+ prefix: str = "") -> None:
47
super().__init__()
48
self.config = config
49
self.is_falcon_mamba = config.model_type == "falcon_mamba"
@@ -58,7 +60,8 @@ class MambaDecoderLayer(nn.Module):
60
rms_norm_has_weight=not self.is_falcon_mamba,
61
rms_norm_eps=mixer_rms_eps,
62
activation=config.hidden_act,
61
- is_lora_enabled=self.is_lora_enabled)
63
+ is_lora_enabled=self.is_lora_enabled,
64
+ prefix=f"{prefix}.mixer")
65
66
self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
67
@@ -107,7 +110,8 @@ class MambaModel(nn.Module):
110
lambda prefix: MambaDecoderLayer(config,
111
cache_config=cache_config,
112
quant_config=quant_config,
110
- is_lora_enabled=is_lora_enabled),
113
+ is_lora_enabled=is_lora_enabled,
114
+ prefix=prefix),
115
prefix=f"{prefix}.layers")
116
117
self.norm_f = RMSNorm(config.hidden_size,
@@ -123,7 +127,7 @@ class MambaModel(nn.Module):
127
self,
128
input_ids: torch.Tensor,
129
positions: torch.Tensor,
126
- mamba_cache_params: MambaCacheParams,
130
+ mamba_cache_params: Optional[MambaCacheParams] = None,
131
intermediate_tensors: Optional[IntermediateTensors] = None,
132
inputs_embeds: Optional[torch.Tensor] = None,
133
) -> torch.Tensor:
@@ -140,12 +144,17 @@ class MambaModel(nn.Module):
144
145
for i in range(self.start_layer, self.end_layer):
146
layer = self.layers[i]
147
+
148
+ layer_cache_params = None
149
+ if mamba_cache_params is not None:
150
+ layer_cache_params = mamba_cache_params.at_layer_idx(
151
+ i - self.start_layer)
152
+
153
hidden_states, residual = layer(
154
positions=positions,
155
hidden_states=hidden_states,
156
residual=residual,
147
- mamba_cache_params=mamba_cache_params.at_layer_idx(
148
- i - self.start_layer))
157
+ mamba_cache_params=layer_cache_params)
158
if not get_pp_group().is_last_rank:
159
return IntermediateTensors({
160
"hidden_states": hidden_states,
@@ -176,8 +185,7 @@ class MambaModel(nn.Module):
185
return loaded_params
186
187
179
-class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP,
180
- SupportsV0Only):
188
+class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
189
190
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
191
config = vllm_config.model_config.hf_config
@@ -227,20 +235,40 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP,
235
intermediate_tensors: Optional[IntermediateTensors] = None,
236
inputs_embeds: Optional[torch.Tensor] = None,
237
**kwargs):
230
- if self.mamba_cache is None:
231
- num_mamba_layers = self.model_config.get_num_layers_by_block_type(
232
- self.vllm_config.parallel_config, LayerBlockType.mamba)
233
- self.mamba_cache = MambaCacheManager(
234
- self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,
235
- *self._get_mamba_cache_shape())
238
237
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
239
+ mamba_cache_params = None
240
+ if not envs.VLLM_USE_V1:
241
+ if self.mamba_cache is None:
242
+ num_layers = self.model_config.get_num_layers_by_block_type(
243
+ self.vllm_config.parallel_config, LayerBlockType.mamba)
244
+ state_shape = self.get_mamba_state_shape_from_config(
245
+ self.vllm_config)
246
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
247
+ self.lm_head.weight.dtype,
248
+ num_layers, *state_shape)
249
+
250
+ mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
251
252
hidden_states = self.backbone(input_ids, positions, mamba_cache_params,
253
intermediate_tensors, inputs_embeds)
254
255
return hidden_states
256
257
+ @classmethod
258
+ def get_mamba_state_shape_from_config(
259
+ cls,
260
+ vllm_config: "VllmConfig",
261
+ ) -> tuple[tuple[int, int], tuple[int, int]]:
262
+ parallel_config = vllm_config.parallel_config
263
+ hf_config = vllm_config.model_config.hf_config
264
+
265
+ return MambaStateShapeCalculator.mamba1_state_shape(
266
+ tp_world_size=parallel_config.tensor_parallel_size,
267
+ intermediate_size=hf_config.intermediate_size,
268
+ state_size=hf_config.state_size,
269
+ conv_kernel=hf_config.conv_kernel,
270
+ use_v1=envs.VLLM_USE_V1)
271
+
272
def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
273
return self.mamba_cache.copy_inputs_before_cuda_graphs(
274
input_buffers, **kwargs)
@@ -248,19 +276,6 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP,
276
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
277
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
278
251
- def _get_mamba_cache_shape(
252
- self) -> tuple[tuple[int, int], tuple[int, int]]:
253
- world_size = get_tensor_model_parallel_world_size()
254
- conv_state_shape = (
255
- self.config.intermediate_size // world_size,
256
- self.config.conv_kernel - 1,
257
- )
258
- temporal_state_shape = (
259
- self.config.intermediate_size // world_size,
260
- self.config.state_size,
261
- )
262
- return conv_state_shape, temporal_state_shape
263
-
279
def compute_logits(self, hidden_states: torch.Tensor,
280
sampling_metadata: SamplingMetadata) -> torch.Tensor:
281
logits = self.logits_processor(self.lm_head, hidden_states,
vllm/model_executor/models/mamba2.py
+3
-2
@@ -19,7 +19,8 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
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 get_mamba_state_shape
22
+from vllm.model_executor.layers.mamba.mamba_utils import (
23
+ MambaStateShapeCalculator)
24
from vllm.model_executor.layers.quantization.base_config import (
25
QuantizationConfig)
26
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -220,7 +221,7 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
221
hf_config = vllm_config.model_config.hf_config
222
intermediate_size = hf_config.expand * hf_config.hidden_size
223
223
- return get_mamba_state_shape(
224
+ return MambaStateShapeCalculator.mamba2_state_shape(
225
intermediate_size=intermediate_size,
226
tp_world_size=parallel_config.tensor_parallel_size,
227
n_groups=hf_config.n_groups,
vllm/model_executor/models/nemotron_h.py
+3
-2
@@ -39,7 +39,8 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
39
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 get_mamba_state_shape
42
+from vllm.model_executor.layers.mamba.mamba_utils import (
43
+ 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)
@@ -482,7 +483,7 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
483
hf_config = vllm_config.model_config.hf_config
484
intermediate_size = hf_config.expand * hf_config.hidden_size
485
485
- return get_mamba_state_shape(
486
+ return MambaStateShapeCalculator.mamba2_state_shape(
487
intermediate_size=intermediate_size,
488
tp_world_size=parallel_config.tensor_parallel_size,
489
n_groups=hf_config.n_groups,
vllm/model_executor/models/zamba2.py
+3
-2
@@ -32,7 +32,8 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
32
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 get_mamba_state_shape
35
+from vllm.model_executor.layers.mamba.mamba_utils import (
36
+ 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 (
@@ -869,7 +870,7 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
870
hf_config = vllm_config.model_config.hf_config
871
intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
872
872
- return get_mamba_state_shape(
873
+ return MambaStateShapeCalculator.mamba2_state_shape(
874
intermediate_size=intermediate_size,
875
tp_world_size=parallel_config.tensor_parallel_size,
876
n_groups=hf_config.mamba_ngroups,
vllm/v1/attention/backends/mamba1_attn.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+from typing import ClassVar
6
+
7
+import torch
8
+
9
+from vllm.attention.backends.abstract import AttentionBackend
10
+from vllm.config import VllmConfig
11
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
12
+ CommonAttentionMetadata)
13
+from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
14
+
15
+
16
+class Mamba1AttentionBackend(AttentionBackend):
17
+
18
+ @staticmethod
19
+ def get_builder_cls() -> type["Mamba1AttentionMetadataBuilder"]:
20
+ return Mamba1AttentionMetadataBuilder
21
+
22
+
23
+@dataclass
24
+class Mamba1AttentionMetadata:
25
+ query_start_loc: torch.Tensor
26
+ context_lens_tensor: torch.Tensor
27
+ state_indices_tensor: torch.Tensor
28
+ has_initial_states: torch.Tensor
29
+
30
+
31
+class Mamba1AttentionMetadataBuilder(
32
+ AttentionMetadataBuilder[Mamba1AttentionMetadata]):
33
+
34
+ reorder_batch_threshold: ClassVar[int] = 1
35
+
36
+ def __init__(
37
+ self,
38
+ kv_cache_spec: AttentionSpec,
39
+ vllm_config: VllmConfig,
40
+ device: torch.device,
41
+ layer_names: list[str],
42
+ ):
43
+ assert isinstance(kv_cache_spec, MambaSpec)
44
+ self.kv_cache_spec = kv_cache_spec
45
+ self.device = device
46
+ self.vllm_config = vllm_config
47
+ self.layer_names = layer_names
48
+
49
+ def build(
50
+ self,
51
+ common_prefix_len: int,
52
+ common_attn_metadata: CommonAttentionMetadata,
53
+ fast_build: bool = False,
54
+ ) -> Mamba1AttentionMetadata:
55
+ query_start_loc = common_attn_metadata.query_start_loc
56
+
57
+ state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
58
+ context_lens_tensor = common_attn_metadata.num_computed_tokens_cpu.to(
59
+ query_start_loc.device)
60
+ has_initial_states = (context_lens_tensor > 0)
61
+
62
+ return Mamba1AttentionMetadata(
63
+ query_start_loc=query_start_loc,
64
+ context_lens_tensor=context_lens_tensor,
65
+ has_initial_states=has_initial_states,
66
+ state_indices_tensor=state_indices_tensor,
67
+ )
vllm/v1/attention/backends/mamba_selectors.py
+4
@@ -1,10 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from vllm.attention.backends.abstract import AttentionBackend
4
+from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
5
from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
6
7
8
def get_mamba_attn_backend(mamba_type: str) -> type[AttentionBackend]:
9
+ if mamba_type == "mamba1":
10
+ return Mamba1AttentionBackend
11
+
12
if mamba_type == "mamba2":
13
return Mamba2AttentionBackend
14