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] - Full CUDA and Piecewise CUDA Graphs Support (#23035)
Signed-off-by: asafg <asafg@ai21.com> Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com> Co-authored-by: asafg <asafg@ai21.com>
Asaf Joseph Gardin committed
Aug 21, 2025 at 06:08 UTC
3663870c72da246d81d8bd8f5c059890fb3f3f5d
9 files changed
+154
-87
docs/usage/v1_guide.md
+1
-1
@@ -107,7 +107,7 @@ to enable simultaneous generation and embedding using the same engine instance i
107
#### Mamba Models
108
109
Models using selective state-space mechanisms instead of standard transformer attention are supported.
110
-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`.
110
+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.
111
112
Models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
113
`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`). Please note that
tests/models/language/generation/test_hybrid.py
+6
-14
@@ -54,16 +54,14 @@ V1_SUPPORTED_MODELS = [
54
"tiiuae/Falcon-H1-0.5B-Base",
55
]
56
57
-# Avoid OOM
58
-MAX_NUM_SEQS = 4
59
-
60
-# Once we add support for FCG in Mamba1, this list will be removed and tests
61
-# all test cases will use enforce_eager=False
62
-ENFORCE_EAGER_MODELS_V1 = [
63
- "state-spaces/mamba-130m-hf",
57
+FULL_CUDA_GRAPH_MODELS = [
58
"ai21labs/Jamba-tiny-dev",
59
+ "Zyphra/Zamba2-1.2B-instruct",
60
]
61
62
+# Avoid OOM
63
+MAX_NUM_SEQS = 4
64
+
65
66
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
67
@pytest.mark.parametrize("max_tokens", [64])
@@ -101,19 +99,13 @@ def test_models(
99
example_prompts, max_tokens, num_logprobs)
100
101
if model in V1_SUPPORTED_MODELS:
104
- enforce_eager = False
102
with monkeypatch.context() as m:
103
m.setenv("VLLM_USE_V1", "1")
104
if model in HYBRID_MODELS:
105
# required due to reorder_batch behaviour
106
m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
110
-
111
- if model in ENFORCE_EAGER_MODELS_V1:
112
- enforce_eager = True
113
-
107
with vllm_runner(model,
108
max_num_seqs=MAX_NUM_SEQS,
116
- enforce_eager=enforce_eager,
109
enable_prefix_caching=False) as vllm_model:
110
vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
111
example_prompts, max_tokens, num_logprobs)
@@ -373,7 +365,7 @@ def test_distributed_correctness(
365
)
366
367
376
-@pytest.mark.parametrize("model", ["Zyphra/Zamba2-1.2B-instruct"])
368
+@pytest.mark.parametrize("model", FULL_CUDA_GRAPH_MODELS)
369
@pytest.mark.parametrize("max_tokens", [64])
370
@pytest.mark.parametrize("num_logprobs", [5])
371
def test_full_cuda_graph(
vllm/config/compilation.py
+1
@@ -336,6 +336,7 @@ class CompilationConfig:
336
"vllm.unified_attention",
337
"vllm.unified_attention_with_output",
338
"vllm.mamba_mixer2",
339
+ "vllm.mamba_mixer",
340
]
341
342
def compute_hash(self) -> str:
vllm/model_executor/layers/mamba/mamba_mixer.py
+57
-9
@@ -27,6 +27,8 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
27
selective_scan_fn, selective_state_update)
28
from vllm.model_executor.models.mamba_cache import MambaCacheParams
29
from vllm.model_executor.utils import set_weight_attrs
30
+from vllm.platforms import current_platform
31
+from vllm.utils import direct_register_custom_op
32
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
33
34
@@ -183,22 +185,26 @@ class MambaMixer(MambaBase, CustomOp):
185
186
def forward(self,
187
hidden_states: torch.Tensor,
188
+ output: torch.Tensor,
189
mamba_cache_params: Optional[MambaCacheParams] = None):
190
if not envs.VLLM_USE_V1:
188
- return CustomOp.forward(self, hidden_states, mamba_cache_params)
191
+ CustomOp.forward(self, hidden_states, output, mamba_cache_params)
192
else:
190
- return self.forward_cuda(
193
+ torch.ops.vllm.mamba_mixer(
194
hidden_states,
192
- mamba_cache_params,
195
+ output,
196
+ self.prefix,
197
)
198
199
def forward_native(self,
200
hidden_states: torch.Tensor,
201
+ output: torch.Tensor,
202
mamba_cache_params: Optional[MambaCacheParams] = None):
203
pass
204
205
def forward_cuda(self,
206
hidden_states: torch.Tensor,
207
+ output: torch.Tensor,
208
mamba_cache_params: Optional[MambaCacheParams] = None):
209
"""
210
Run the Mamba-1 SSM pipeline.
@@ -237,6 +243,7 @@ class MambaMixer(MambaBase, CustomOp):
243
conv_state = self_kv_cache[0].transpose(-1, -2)
244
ssm_state = self_kv_cache[1]
245
has_initial_states = mamba1_metadata.has_initial_states
246
+ num_padded_decodes = mamba1_metadata.num_padded_decodes
247
else:
248
assert isinstance(attn_metadata, AttentionMetadata)
249
assert mamba_cache_params is not None
@@ -248,6 +255,7 @@ class MambaMixer(MambaBase, CustomOp):
255
has_initial_states = None
256
if context_lens_tensor is not None:
257
has_initial_states = context_lens_tensor > 0
258
+ num_padded_decodes = attn_metadata.num_decode_tokens
259
260
# 1. Gated MLP's linear projection
261
projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)
@@ -267,6 +275,7 @@ class MambaMixer(MambaBase, CustomOp):
275
num_decodes = attn_metadata.num_decode_tokens # token count (=request)
276
has_prefill = num_prefill_tokens > 0
277
has_decode = num_decode_tokens > 0
278
+ num_actual_tokens = num_prefill_tokens + num_decode_tokens
279
280
prefill_decode_split = split_batch_to_prefill_and_decode(
281
hidden_states_BC,
@@ -278,6 +287,7 @@ class MambaMixer(MambaBase, CustomOp):
287
num_decode_tokens,
288
num_prefills,
289
num_decodes,
290
+ num_padded_decodes,
291
)
292
hidden_states_BC_p = prefill_decode_split.hidden_states_BC_p
293
hidden_states_BC_d = prefill_decode_split.hidden_states_BC_d
@@ -371,7 +381,7 @@ class MambaMixer(MambaBase, CustomOp):
381
else:
382
out = self.out_proj(scan_outputs_combined.transpose(-2, -1))[0]
383
374
- return out
384
+ output[:num_actual_tokens] = out
385
386
def get_state_dtype(self) -> tuple[torch.dtype]:
387
assert self.model_config is not None
@@ -421,18 +431,27 @@ def split_batch_to_prefill_and_decode(
431
num_decode_tokens: int,
432
num_prefills: int,
433
num_decodes: int,
434
+ num_padded_decodes: int,
435
) -> PrefillDecodeSplit:
436
+ num_actual_tokens = num_prefill_tokens + num_padded_decodes
437
+
438
if envs.VLLM_USE_V1:
439
# In v1, decode tokens come first, then prefill tokens.
440
hidden_states_BC_d, hidden_states_BC_p = torch.split(
428
- hidden_states_BC, [num_decode_tokens, num_prefill_tokens], dim=-1)
429
- gate_d, gate_p = torch.split(gate,
430
- [num_decode_tokens, num_prefill_tokens],
441
+ hidden_states_BC[..., :num_actual_tokens],
442
+ [num_padded_decodes, num_prefill_tokens],
443
+ dim=-1)
444
+ gate_d, gate_p = torch.split(gate[..., :num_actual_tokens],
445
+ [num_padded_decodes, num_prefill_tokens],
446
dim=-1)
447
+
448
+ # num_padded_decodes accounts for CUDA graph padding when applicable
449
state_indices_tensor_d, state_indices_tensor_p = torch.split(
433
- state_indices_tensor, [num_decodes, num_prefills], dim=0)
450
+ state_indices_tensor[:num_padded_decodes + num_prefills],
451
+ [num_padded_decodes, num_prefills],
452
+ dim=0)
453
query_start_loc_p = (query_start_loc[-num_prefills - 1:] -
435
- num_decodes if num_prefills > 0 else None)
454
+ num_padded_decodes if num_prefills > 0 else None)
455
has_initial_states_p = has_initial_states[-num_prefills:] if (
456
has_initial_states is not None and num_prefills > 0) else None
457
else:
@@ -459,3 +478,32 @@ def split_batch_to_prefill_and_decode(
478
query_start_loc_p=query_start_loc_p,
479
has_initial_states_p=has_initial_states_p,
480
)
481
+
482
+
483
+def mamba_mixer(
484
+ hidden_states: torch.Tensor,
485
+ output: torch.Tensor,
486
+ layer_name: str,
487
+) -> None:
488
+ forward_context: ForwardContext = get_forward_context()
489
+ self = forward_context.no_compile_layers[layer_name]
490
+ self.forward_cuda(hidden_states=hidden_states,
491
+ output=output,
492
+ mamba_cache_params=None)
493
+
494
+
495
+def mamba_mixer_fake(
496
+ hidden_states: torch.Tensor,
497
+ output: torch.Tensor,
498
+ layer_name: str,
499
+) -> None:
500
+ return
501
+
502
+
503
+direct_register_custom_op(
504
+ op_name="mamba_mixer",
505
+ op_func=mamba_mixer,
506
+ mutates_args=["output"],
507
+ fake_impl=mamba_mixer_fake,
508
+ dispatch_key=current_platform.dispatch_key,
509
+)
vllm/model_executor/models/jamba.py
+5
-3
@@ -10,6 +10,7 @@ from transformers import JambaConfig
10
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, ModelConfig, VllmConfig
15
from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.distributed.parallel_state import get_pp_group
@@ -154,10 +155,10 @@ class JambaMambaDecoderLayer(nn.Module):
155
hidden_states, residual = self.input_layernorm(
156
hidden_states, residual)
157
157
- hidden_states = self.mamba(hidden_states, mamba_cache_params)
158
+ output = torch.empty_like(hidden_states)
159
+ self.mamba(hidden_states, output, mamba_cache_params)
160
# Fully Connected
159
- hidden_states, residual = self.pre_ff_layernorm(
160
- hidden_states, residual)
161
+ hidden_states, residual = self.pre_ff_layernorm(output, residual)
162
hidden_states = self.feed_forward(hidden_states)
163
return hidden_states, residual
164
@@ -278,6 +279,7 @@ ALL_DECODER_LAYER_TYPES = {
279
}
280
281
282
+@support_torch_compile
283
class JambaModel(nn.Module):
284
285
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
vllm/model_executor/models/mamba.py
+5
-2
@@ -9,6 +9,7 @@ from torch import nn
9
from transformers import MambaConfig
10
11
from vllm import envs
12
+from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed.parallel_state import get_pp_group
15
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -81,10 +82,12 @@ class MambaDecoderLayer(nn.Module):
82
else:
83
hidden_states, residual = self.norm(hidden_states, residual)
84
84
- hidden_states = self.mixer(hidden_states, mamba_cache_params)
85
- return hidden_states, residual
85
+ output = torch.empty_like(hidden_states)
86
+ self.mixer(hidden_states, output, mamba_cache_params)
87
+ return output, residual
88
89
90
+@support_torch_compile
91
class MambaModel(nn.Module):
92
93
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
vllm/v1/attention/backends/mamba1_attn.py
+17
-20
@@ -2,16 +2,16 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import dataclass
5
-from typing import ClassVar, Optional
5
+from typing import Optional
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,
10
+from vllm.attention.backends.utils import PAD_SLOT_ID
11
+from vllm.v1.attention.backends.mamba_attn import (
12
+ BaseMambaAttentionMetadataBuilder)
13
+from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
14
split_decodes_and_prefills)
14
-from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
15
16
17
class Mamba1AttentionBackend(AttentionBackend):
@@ -31,24 +31,11 @@ class Mamba1AttentionMetadata:
31
num_prefill_tokens: int
32
num_decodes: int
33
num_decode_tokens: int
34
+ num_padded_decodes: int
35
36
37
class Mamba1AttentionMetadataBuilder(
37
- AttentionMetadataBuilder[Mamba1AttentionMetadata]):
38
- reorder_batch_threshold: ClassVar[int] = 1
39
-
40
- def __init__(
41
- self,
42
- kv_cache_spec: AttentionSpec,
43
- vllm_config: VllmConfig,
44
- device: torch.device,
45
- layer_names: list[str],
46
- ):
47
- assert isinstance(kv_cache_spec, MambaSpec)
48
- self.kv_cache_spec = kv_cache_spec
49
- self.device = device
50
- self.vllm_config = vllm_config
51
- self.layer_names = layer_names
38
+ BaseMambaAttentionMetadataBuilder[Mamba1AttentionMetadata]):
39
40
def build(
41
self,
@@ -67,9 +54,18 @@ class Mamba1AttentionMetadataBuilder(
54
decode_threshold=1))
55
56
has_initial_states = None
57
+ padded_decodes = num_decodes
58
59
if num_prefills > 0:
60
has_initial_states = context_lens_tensor > 0
61
+ elif (num_decodes > 0 and num_decodes <= self.decode_cudagraph_max_bs
62
+ and self.compilation_config.full_cuda_graph):
63
+ state_indices_for_decode = state_indices_tensor[:num_decodes]
64
+ padded_decodes = self.vllm_config.pad_for_cudagraph(num_decodes)
65
+ self.state_indices_tensor[:num_decodes].copy_(
66
+ state_indices_for_decode, non_blocking=True)
67
+ state_indices_tensor = self.state_indices_tensor[:padded_decodes]
68
+ state_indices_tensor[num_decodes:] = PAD_SLOT_ID
69
70
return Mamba1AttentionMetadata(
71
query_start_loc=query_start_loc,
@@ -80,4 +76,5 @@ class Mamba1AttentionMetadataBuilder(
76
num_prefill_tokens=num_prefill_tokens,
77
num_decodes=num_decodes,
78
num_decode_tokens=num_decode_tokens,
79
+ num_padded_decodes=padded_decodes,
80
)
vllm/v1/attention/backends/mamba2_attn.py
+7
-38
@@ -2,18 +2,18 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import math
4
from dataclasses import dataclass
5
-from typing import ClassVar, Optional
5
+from typing import Optional
6
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionBackend
10
from vllm.attention.backends.utils import PAD_SLOT_ID
11
from vllm.config import VllmConfig
12
-from vllm.v1.attention.backends.utils import (AttentionCGSupport,
13
- AttentionMetadataBuilder,
14
- CommonAttentionMetadata,
12
+from vllm.v1.attention.backends.mamba_attn import (
13
+ BaseMambaAttentionMetadataBuilder)
14
+from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
15
split_decodes_and_prefills)
16
-from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
16
+from vllm.v1.kv_cache_interface import AttentionSpec
17
18
19
def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,
@@ -88,29 +88,14 @@ class Mamba2AttentionMetadata:
88
89
90
class Mamba2AttentionMetadataBuilder(
91
- AttentionMetadataBuilder[Mamba2AttentionMetadata]):
92
- cudagraph_support: ClassVar[AttentionCGSupport] = \
93
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
94
-
95
- reorder_batch_threshold: ClassVar[int] = 1
91
+ BaseMambaAttentionMetadataBuilder[Mamba2AttentionMetadata]):
92
93
def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
94
vllm_config: VllmConfig, device: torch.device):
99
- assert isinstance(kv_cache_spec, MambaSpec)
100
- self.kv_cache_spec = kv_cache_spec
95
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
96
self.chunk_size = vllm_config.model_config.get_mamba_chunk_size()
102
- self.vllm_config = vllm_config
103
- self.compilation_config = vllm_config.compilation_config
97
assert self.chunk_size is not None, (
98
"chunk_size needs to be set in the model config for Mamba2 models")
106
- self.decode_cudagraph_max_bs = min(
107
- self.vllm_config.scheduler_config.max_num_seqs,
108
- self.compilation_config.max_capture_size)
109
- self.state_indices_tensor = torch.empty(
110
- (self.decode_cudagraph_max_bs, ),
111
- dtype=torch.int32,
112
- device=device,
113
- )
99
100
def build(self,
101
common_prefix_len: int,
@@ -187,19 +172,3 @@ class Mamba2AttentionMetadataBuilder(
172
state_indices_tensor=state_indices_tensor,
173
)
174
return attn_metadata
190
-
191
- def build_for_cudagraph_capture(
192
- self, common_attn_metadata: CommonAttentionMetadata):
193
- """
194
- This method builds the metadata for full cudagraph capture.
195
- Currently, only decode is supported for full cudagraphs with Mamba.
196
- """
197
- m = common_attn_metadata
198
-
199
- assert m.num_reqs == m.num_actual_tokens, \
200
- "Mamba only supports decode-only full CUDAGraph capture. " \
201
- "Make sure all cudagraph capture sizes <= max_num_seq."
202
-
203
- m.max_query_len = 1 # decode-only
204
-
205
- return self.build(0, m)
vllm/v1/attention/backends/mamba_attn.py
new
+55
@@ -0,0 +1,55 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import abc
5
+from typing import ClassVar, TypeVar
6
+
7
+import torch
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.v1.attention.backends.utils import (AttentionCGSupport,
11
+ AttentionMetadataBuilder,
12
+ CommonAttentionMetadata)
13
+from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
14
+
15
+M = TypeVar("M")
16
+
17
+
18
+class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
19
+ reorder_batch_threshold: ClassVar[int] = 1
20
+ cudagraph_support: ClassVar[AttentionCGSupport] = \
21
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
22
+
23
+ def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
24
+ vllm_config: VllmConfig, device: torch.device):
25
+ assert isinstance(kv_cache_spec, MambaSpec)
26
+ self.kv_cache_spec = kv_cache_spec
27
+ self.device = device
28
+ self.vllm_config = vllm_config
29
+ self.layer_names = layer_names
30
+
31
+ self.compilation_config = vllm_config.compilation_config
32
+ self.decode_cudagraph_max_bs = min(
33
+ self.vllm_config.scheduler_config.max_num_seqs,
34
+ self.compilation_config.max_capture_size)
35
+ self.state_indices_tensor = torch.empty(
36
+ (self.decode_cudagraph_max_bs, ),
37
+ dtype=torch.int32,
38
+ device=device,
39
+ )
40
+
41
+ def build_for_cudagraph_capture(
42
+ self, common_attn_metadata: CommonAttentionMetadata) -> M:
43
+ """
44
+ This method builds the metadata for full cudagraph capture.
45
+ Currently, only decode is supported for full cudagraphs with Mamba.
46
+ """
47
+ m = common_attn_metadata
48
+
49
+ assert m.num_reqs == m.num_actual_tokens, \
50
+ "Mamba only supports decode-only full CUDAGraph capture. " \
51
+ "Make sure all cudagraph capture sizes <= max_num_seq."
52
+
53
+ m.max_query_len = 1 # decode-only
54
+
55
+ return self.build(0, m)
\ No newline at end of file