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] v1 engine + full CUDA graph support for PLaMo2 (#23998)
Signed-off-by: Hemmi Shinichi <shemmi@preferred.jp> Signed-off-by: nopperl <54780682+nopperl@users.noreply.github.com> Co-authored-by: Hemmi Shinichi <shemmi@preferred.jp> Co-authored-by: Thomas Parnell <tom.parnell@gmail.com>
nopperl committed
Sep 4, 2025 at 00:24 UTC
fa4311d85f3a5f6d445f20429ca6d38122908eeb
6 files changed
+350
-126
docs/models/supported_models.md
+1
-1
@@ -395,7 +395,7 @@ th {
395
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
396
| `Phi4FlashForCausalLM` | Phi-4-mini-flash-reasoning | `microsoft/microsoft/Phi-4-mini-instruct`, etc. | | | |
397
| `PersimmonForCausalLM` | Persimmon | `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc. | | ✅︎ | ✅︎ |
398
-| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | | ✅︎ | |
398
+| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | | ✅︎ | ✅︎ |
399
| `QWenLMHeadModel` | Qwen | `Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |
400
| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ | ✅︎ |
401
| `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |
docs/usage/v1_guide.md
+1
-1
@@ -110,7 +110,7 @@ Models using selective state-space mechanisms instead of standard transformer at
110
Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`,`FalconMambaForCausalLM`) are supported.
111
112
Hybrid 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`).
113
+`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
114
115
Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`).
116
tests/models/language/generation/test_hybrid.py
+3
-2
@@ -25,8 +25,7 @@ SSM_MODELS = [
25
26
HYBRID_MODELS = [
27
"ai21labs/Jamba-tiny-dev",
28
- # skipping until vLLM implementation issues are resolved
29
- # "pfnet/plamo-2-1b",
28
+ "pfnet/plamo-2-1b",
29
"Zyphra/Zamba2-1.2B-instruct",
30
"hmellor/tiny-random-BambaForCausalLM",
31
"ibm-granite/granite-4.0-tiny-preview",
@@ -37,6 +36,7 @@ HYBRID_MODELS = [
36
V1_SUPPORTED_MODELS = [
37
"state-spaces/mamba-130m-hf",
38
"ai21labs/Jamba-tiny-dev",
39
+ "pfnet/plamo-2-1b",
40
"yujiepan/mamba2-codestral-v0.1-tiny-random",
41
"Zyphra/Zamba2-1.2B-instruct",
42
"hmellor/tiny-random-BambaForCausalLM",
@@ -47,6 +47,7 @@ V1_SUPPORTED_MODELS = [
47
48
FULL_CUDA_GRAPH_MODELS = [
49
"ai21labs/Jamba-tiny-dev",
50
+ "pfnet/plamo-2-1b",
51
"Zyphra/Zamba2-1.2B-instruct",
52
]
53
tests/models/registry.py
-2
@@ -287,8 +287,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
287
"PhiMoEForCausalLM": _HfExamplesInfo("microsoft/Phi-3.5-MoE-instruct",
288
trust_remote_code=True),
289
"Plamo2ForCausalLM": _HfExamplesInfo("pfnet/plamo-2-1b",
290
- max_transformers_version="4.53",
291
- transformers_version_reason="vLLM impl inherits PreTrainedModel and clashes with get_input_embeddings", # noqa: E501
290
trust_remote_code=True),
291
"QWenLMHeadModel": _HfExamplesInfo("Qwen/Qwen-7B-Chat",
292
max_transformers_version="4.53",
vllm/config/compilation.py
+1
@@ -340,6 +340,7 @@ class CompilationConfig:
340
"vllm.mamba_mixer",
341
"vllm.short_conv",
342
"vllm.linear_attention",
343
+ "vllm.plamo2_mamba_mixer",
344
]
345
346
def compute_hash(self) -> str:
vllm/model_executor/models/plamo2.py
+344
-120
@@ -3,19 +3,24 @@
3
"""Inference-only PLaMo2 model."""
4
from collections.abc import Iterable
5
from itertools import islice
6
-from typing import Optional
6
+from typing import TYPE_CHECKING, Optional
7
+
8
+if TYPE_CHECKING:
9
+ from vllm.attention.backends.abstract import AttentionBackend
10
11
import torch
12
from torch import nn
10
-from transformers import PretrainedConfig, PreTrainedModel
13
+from transformers import PretrainedConfig
14
15
+from vllm import envs
16
from vllm.attention.backends.abstract import AttentionMetadata
17
from vllm.attention.layer import Attention
18
from vllm.compilation.decorators import support_torch_compile
15
-from vllm.config import VllmConfig
19
+from vllm.config import VllmConfig, get_current_vllm_config
20
from vllm.distributed import divide, get_tensor_model_parallel_world_size
21
from vllm.distributed.parallel_state import get_pp_group
18
-from vllm.forward_context import get_forward_context
22
+from vllm.forward_context import ForwardContext, get_forward_context
23
+from vllm.model_executor.custom_op import CustomOp
24
from vllm.model_executor.layers.activation import SiluAndMul
25
from vllm.model_executor.layers.layernorm import RMSNorm
26
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
@@ -23,8 +28,11 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
28
QKVParallelLinear,
29
RowParallelLinear)
30
from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
+from vllm.model_executor.layers.mamba.abstract import MambaBase
32
from vllm.model_executor.layers.mamba.mamba2_metadata import (
27
- Mamba2Metadata, prepare_mamba2_metadata)
33
+ Mamba2Metadata, prepare_mamba2_metadata, update_metadata)
34
+from vllm.model_executor.layers.mamba.mamba_utils import (
35
+ MambaStateDtypeCalculator, MambaStateShapeCalculator)
36
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
37
causal_conv1d_fn, causal_conv1d_update)
38
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
@@ -39,7 +47,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
47
from vllm.model_executor.model_loader.weight_utils import (
48
composed_weight_loader, default_weight_loader, sharded_weight_loader)
49
from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,
42
- SupportsPP, SupportsV0Only)
50
+ SupportsPP)
51
from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
52
MambaCacheParams)
53
from vllm.model_executor.models.utils import (
@@ -47,8 +55,10 @@ from vllm.model_executor.models.utils import (
55
make_layers, maybe_prefix)
56
from vllm.model_executor.sampling_metadata import SamplingMetadata
57
from vllm.model_executor.utils import set_weight_attrs
58
+from vllm.platforms import current_platform
59
from vllm.sequence import IntermediateTensors
51
-from vllm.utils import LayerBlockType
60
+from vllm.utils import LayerBlockType, direct_register_custom_op
61
+from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
62
63
64
# Only used for type hinting.
@@ -73,20 +83,6 @@ class Plamo2Config(PretrainedConfig): # type: ignore
83
vocab_size: int
84
85
76
-class Plamo2PreTrainedModel(PreTrainedModel): # type: ignore
77
-
78
- def _init_weights(self, module: torch.nn.Module) -> None:
79
- std = 0.02
80
- if isinstance(module, nn.Linear):
81
- module.weight.data.normal_(mean=0.0, std=std)
82
- if module.bias is not None:
83
- module.bias.data.zero_()
84
- elif isinstance(module, nn.Embedding):
85
- module.weight.data.normal_(mean=0.0, std=std)
86
- if module.padding_idx is not None:
87
- module.weight.data[module.padding_idx].zero_()
88
-
89
-
86
def is_mamba(config: Plamo2Config, i: int) -> bool:
87
assert config.mamba_step > 1
88
@@ -99,7 +95,8 @@ def is_mamba(config: Plamo2Config, i: int) -> bool:
95
# Adapted from:
96
# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
97
# transformers.models.mamba.modeling_mamba.MambaMixer
102
-class Plamo2MambaMixer(nn.Module):
98
+@CustomOp.register(name="plamo2_mamba_mixer")
99
+class Plamo2MambaMixer(MambaBase, CustomOp):
100
101
def __init__(self,
102
vllm_config: VllmConfig,
@@ -108,6 +105,8 @@ class Plamo2MambaMixer(nn.Module):
105
**kwargs) -> None:
106
super().__init__()
107
self.config = vllm_config.model_config.hf_config
108
+ self.cache_config = vllm_config.cache_config
109
+ self.model_config = vllm_config.model_config
110
self.quant_config = vllm_config.quant_config
111
self.hidden_size = self.config.hidden_size
112
self.ssm_state_size = self.config.mamba_d_state
@@ -115,8 +114,6 @@ class Plamo2MambaMixer(nn.Module):
114
self.intermediate_size = (self.config.mamba_num_heads *
115
self.config.hidden_size_per_head)
116
self.tp_size = get_tensor_model_parallel_world_size()
118
- self.intermediate_size_per_tp_worker = \
119
- self.intermediate_size // self.tp_size
117
self.head_dim = self.config.hidden_size_per_head
118
self.num_heads = self.config.mamba_num_heads
119
self.time_step_rank = max(64, self.hidden_size // 16)
@@ -197,6 +194,22 @@ class Plamo2MambaMixer(nn.Module):
194
self.C_norm = RMSNorm(self.ssm_state_size,
195
eps=self.config.rms_norm_eps)
196
197
+ self.chunk_size = self.config.mamba_chunk_size
198
+
199
+ if envs.VLLM_USE_V1:
200
+ compilation_config = get_current_vllm_config().compilation_config
201
+ if prefix in compilation_config.static_forward_context:
202
+ raise ValueError(f"Duplicate layer name: {prefix}")
203
+ compilation_config.static_forward_context[prefix] = self
204
+ # The outer list is for v0 PP virtual engine. Though this code path
205
+ # only runs for v1, we have to do this to unify with the interface
206
+ # of Attention + v0 PP.
207
+ # The inner tuple is (conv_state, ssm_state)
208
+ self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
209
+ assert self.chunk_size != -1, "chunk_size must be set for v1"
210
+
211
+ self.prefix = prefix
212
+
213
def _project_ssm_parameters(self, hidden_states):
214
ssm_parameters = self.bcdt_proj(hidden_states)
215
B, C, time_step = torch.split(
@@ -212,25 +225,76 @@ class Plamo2MambaMixer(nn.Module):
225
dt = self.dt_proj(time_step)
226
return B, C, dt
227
228
+ def forward_native(
229
+ self,
230
+ hidden_states: torch.Tensor,
231
+ output: torch.Tensor,
232
+ mamba_cache_params: MambaCacheParams,
233
+ mamba2_metadata: Mamba2Metadata,
234
+ **kwargs,
235
+ ):
236
+ pass
237
+
238
def forward(
239
self,
240
hidden_states: torch.Tensor,
241
+ output: torch.Tensor,
242
mamba_cache_params: MambaCacheParams,
243
mamba2_metadata: Mamba2Metadata,
244
**kwargs,
221
- ) -> torch.Tensor:
245
+ ):
246
+ if not envs.VLLM_USE_V1:
247
+ CustomOp.forward(self, hidden_states, output, mamba_cache_params,
248
+ mamba2_metadata)
249
+ else:
250
+ torch.ops.vllm.plamo2_mamba_mixer(
251
+ hidden_states,
252
+ output,
253
+ self.prefix,
254
+ )
255
+
256
+ def forward_cuda(
257
+ self,
258
+ hidden_states: torch.Tensor,
259
+ output: torch.Tensor,
260
+ mamba_cache_params: MambaCacheParams,
261
+ mamba2_metadata: Mamba2Metadata,
262
+ **kwargs,
263
+ ):
264
265
+ forward_context = get_forward_context()
266
# mamba2_metadata contains metadata necessary for the mamba2 triton
267
# kernels to operate in continuous batching and in chunked prefill
268
# modes; they are computed at top-level model forward since they
269
# stay the same and reused for all mamba layers in the same iteration
227
- attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
228
-
229
- num_prefills = attn_metadata.num_prefills # request count
230
- num_decodes = attn_metadata.num_decode_tokens # token count (=request)
231
- num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
232
- has_prefill = num_prefills > 0
233
- has_decode = num_decodes > 0
270
+ attn_metadata: AttentionMetadata = forward_context.attn_metadata
271
+ if envs.VLLM_USE_V1:
272
+ if attn_metadata is not None:
273
+ assert isinstance(attn_metadata, dict)
274
+ attn_metadata = attn_metadata[self.prefix]
275
+ mamba2_metadata = attn_metadata
276
+ assert isinstance(attn_metadata, Mamba2AttentionMetadata)
277
+ self_kv_cache = self.kv_cache[forward_context.virtual_engine]
278
+ # conv_state = (..., dim, width-1) yet contiguous along 'dim'
279
+ conv_state = self_kv_cache[0].transpose(-1, -2)
280
+ ssm_state = self_kv_cache[1]
281
+ state_indices_tensor = attn_metadata.state_indices_tensor
282
+ has_initial_states_p = attn_metadata.has_initial_states_p
283
+ prep_initial_states = attn_metadata.prep_initial_states
284
+ chunk_size = attn_metadata.chunk_size
285
+ seq_idx_p = attn_metadata.seq_idx_p
286
+ chunk_indices_p = attn_metadata.chunk_indices_p
287
+ chunk_offsets_p = attn_metadata.chunk_offsets_p
288
+ else:
289
+ conv_state = mamba_cache_params.conv_state
290
+ ssm_state = mamba_cache_params.ssm_state
291
+ state_indices_tensor = mamba_cache_params.state_indices_tensor
292
+ has_initial_states_p = mamba2_metadata.has_initial_states
293
+ prep_initial_states = mamba2_metadata.prep_initial_states
294
+ chunk_size = mamba2_metadata.chunk_size
295
+ seq_idx_p = mamba2_metadata.seq_idx
296
+ chunk_indices_p = mamba2_metadata.chunk_indices
297
+ chunk_offsets_p = mamba2_metadata.chunk_offsets
298
299
# 1. Gated MLP's linear projection
300
projected_states = self.in_proj(hidden_states)
@@ -240,23 +304,59 @@ class Plamo2MambaMixer(nn.Module):
304
conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
305
self.conv1d.weight.size(2))
306
307
+ if envs.VLLM_USE_V1 and attn_metadata is None:
308
+ # V1 profile run
309
+ hidden_states = (hidden_states.transpose(0, 1).clone().transpose(
310
+ 0, 1)).contiguous()
311
+ output[:] = self.out_proj(hidden_states)
312
+ return
313
+
314
+ num_prefills = attn_metadata.num_prefills # request count
315
+ num_decodes = attn_metadata.num_decode_tokens # token count (=request)
316
+ num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
317
+ has_prefill = num_prefills > 0
318
+ has_decode = num_decodes > 0
319
+ num_actual_tokens = num_prefill_tokens + num_decodes
320
+
321
+ # NOTE: V0 put prefill before decode, v1 puts decode before prefill
322
# Separate prefill and decode by splitting varlen input
323
# Split along token dimension
245
- hidden_states_p, hidden_states_d = torch.split(
246
- hidden_states,
247
- [num_prefill_tokens, num_decodes],
248
- dim=0,
249
- )
250
- gate_p, gate_d = torch.split(gate, [num_prefill_tokens, num_decodes],
251
- dim=0)
252
- # Split along batch dimension
253
- state_indices_tensor_p, state_indices_tensor_d = torch.split(
254
- mamba_cache_params.state_indices_tensor,
255
- [num_prefills, num_decodes],
256
- dim=0,
257
- )
258
- query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills + 1]
259
- if has_prefill else None)
324
+ if envs.VLLM_USE_V1:
325
+ hidden_states_d, hidden_states_p = torch.split(
326
+ hidden_states[:num_actual_tokens],
327
+ [num_decodes, num_prefill_tokens],
328
+ dim=0,
329
+ )
330
+ gate_d, gate_p = torch.split(gate[:num_actual_tokens],
331
+ [num_decodes, num_prefill_tokens],
332
+ dim=0)
333
+ # Split along batch dimension
334
+ state_indices_tensor_d, state_indices_tensor_p = torch.split(
335
+ state_indices_tensor,
336
+ [num_decodes, num_prefills],
337
+ dim=0,
338
+ )
339
+ query_start_loc_p = (
340
+ attn_metadata.query_start_loc[-num_prefills - 1:] -
341
+ num_decodes if has_prefill else None)
342
+ else:
343
+ hidden_states_p, hidden_states_d = torch.split(
344
+ hidden_states,
345
+ [num_prefill_tokens, num_decodes],
346
+ dim=0,
347
+ )
348
+ gate_p, gate_d = torch.split(gate,
349
+ [num_prefill_tokens, num_decodes],
350
+ dim=0)
351
+ # Split along batch dimension
352
+ state_indices_tensor_p, state_indices_tensor_d = torch.split(
353
+ state_indices_tensor,
354
+ [num_prefills, num_decodes],
355
+ dim=0,
356
+ )
357
+ query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills +
358
+ 1]
359
+ if has_prefill else None)
360
361
# Preallocate output tensor to avoid memcpy cost for merging prefill
362
# and decode outputs
@@ -268,25 +368,38 @@ class Plamo2MambaMixer(nn.Module):
368
dtype=hidden_states.dtype,
369
device=hidden_states.device,
370
)
271
- preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split(
272
- preallocated_ssm_out,
273
- [num_prefill_tokens, num_decodes],
274
- dim=0,
275
- )
371
+ if envs.VLLM_USE_V1:
372
+ preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
373
+ preallocated_ssm_out,
374
+ [num_decodes, num_prefill_tokens],
375
+ dim=0,
376
+ )
377
+ else:
378
+ preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split(
379
+ preallocated_ssm_out,
380
+ [num_prefill_tokens, num_decodes],
381
+ dim=0,
382
+ )
383
384
# Process prefill requests
385
if has_prefill:
386
# 2. Convolution sequence transformation
387
# - "cache_indices" updates the conv_state cache in positions
281
- # pointed to by "mamba_cache_params.state_indices_tensor"
388
+ # pointed to by "state_indices_tensor"
389
+ x = hidden_states_p.transpose(
390
+ 0, 1) # this is the form that causal-conv see
391
+ if mamba2_metadata.cu_seqlen is None:
392
+ mamba2_metadata = update_metadata(x, query_start_loc_p,
393
+ mamba2_metadata)
394
hidden_states_p = causal_conv1d_fn(
283
- hidden_states_p.transpose(0, 1),
395
+ x,
396
conv_weights,
397
self.conv1d.bias,
398
activation=self.activation,
287
- conv_states=mamba_cache_params.conv_state,
288
- has_initial_state=mamba2_metadata.has_initial_states,
399
+ conv_states=conv_state,
400
+ has_initial_state=has_initial_states_p,
401
cache_indices=state_indices_tensor_p,
402
+ metadata=mamba2_metadata,
403
query_start_loc=query_start_loc_p)
404
hidden_states_p = hidden_states_p.transpose(0, 1)
405
hidden_states_p = hidden_states_p[:num_prefill_tokens]
@@ -299,12 +412,16 @@ class Plamo2MambaMixer(nn.Module):
412
413
# 3. State Space Model sequence transformation
414
initial_states = None
302
- if (mamba2_metadata.has_initial_states is not None
303
- and mamba2_metadata.prep_initial_states):
415
+ if has_initial_states_p is not None and prep_initial_states:
416
# making a copy of the states
305
- initial_states = torch.where(
306
- mamba2_metadata.has_initial_states[:, None, None, None],
307
- mamba_cache_params.ssm_state[state_indices_tensor_p], 0)
417
+ if envs.VLLM_USE_V1:
418
+ initial_states = torch.where(
419
+ has_initial_states_p[:, None, None, None],
420
+ ssm_state[state_indices_tensor_p], 0)
421
+ else:
422
+ initial_states = torch.where(
423
+ has_initial_states_p[:num_prefills, None, None, None],
424
+ ssm_state[state_indices_tensor_p], 0)
425
varlen_state = mamba_chunk_scan_combined(
426
hidden_states_p.view(1, num_prefill_tokens,
427
self.num_heads // self.tp_size,
@@ -313,15 +430,15 @@ class Plamo2MambaMixer(nn.Module):
430
self.A,
431
B.view(1, num_prefill_tokens, 1, -1),
432
C.view(1, num_prefill_tokens, 1, -1),
316
- chunk_size=mamba2_metadata.chunk_size,
433
+ chunk_size=chunk_size,
434
D=self.D,
435
z=gate_p.view(1, num_prefill_tokens,
436
self.num_heads // self.tp_size, self.head_dim),
437
dt_bias=self.dt_bias,
321
- seq_idx=mamba2_metadata.seq_idx,
322
- chunk_indices=mamba2_metadata.chunk_indices,
323
- chunk_offsets=mamba2_metadata.chunk_offsets,
324
- cu_seqlens=attn_metadata.query_start_loc[:num_prefills + 1],
438
+ seq_idx=seq_idx_p,
439
+ chunk_indices=chunk_indices_p,
440
+ chunk_offsets=chunk_offsets_p,
441
+ cu_seqlens=query_start_loc_p,
442
initial_states=initial_states,
443
return_varlen_states=True,
444
return_final_states=False,
@@ -329,18 +446,19 @@ class Plamo2MambaMixer(nn.Module):
446
dt_limit=(0.0, float("inf")),
447
out=preallocated_ssm_out_p.view(1, num_prefill_tokens, -1,
448
self.head_dim),
449
+ state_dtype=ssm_state.dtype,
450
)
451
452
# update ssm states
453
# - varlen state is a (batch, nheads, headdim, dstate) tensor
336
- mamba_cache_params.ssm_state[state_indices_tensor_p] = varlen_state
454
+ ssm_state[state_indices_tensor_p] = varlen_state
455
456
# Process decode requests
457
if has_decode:
458
# 2. Convolution sequence transformation
459
hidden_states_d = causal_conv1d_update(
460
hidden_states_d,
343
- mamba_cache_params.conv_state,
461
+ conv_state,
462
conv_weights,
463
self.conv1d.bias,
464
self.activation,
@@ -363,8 +481,10 @@ class Plamo2MambaMixer(nn.Module):
481
# - the hidden is reshaped into (bs, num_heads, head_dim)
482
# - mamba_cache_params.ssm_state's slots will be selected
483
# using state_indices_tensor_d
484
+
485
+ # NOTE: final output is an in-place update of out tensor
486
selective_state_update(
367
- mamba_cache_params.ssm_state,
487
+ ssm_state,
488
hidden_states_d,
489
dt,
490
A,
@@ -378,11 +498,68 @@ class Plamo2MambaMixer(nn.Module):
498
out=preallocated_ssm_out_d.view(num_decodes, -1,
499
self.head_dim),
500
)
381
- assert self.num_heads % self.tp_size == 0
501
502
# 4. Final linear projection
384
- out = self.out_proj(preallocated_ssm_out)
385
- return out
503
+ output[:num_actual_tokens] = self.out_proj(preallocated_ssm_out)
504
+
505
+ def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
506
+ assert self.model_config is not None
507
+ assert self.cache_config is not None
508
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
509
+ self.model_config.dtype,
510
+ self.cache_config.mamba_cache_dtype,
511
+ self.cache_config.mamba_ssm_cache_dtype,
512
+ )
513
+
514
+ def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
515
+ return MambaStateShapeCalculator.mamba2_state_shape(
516
+ intermediate_size=self.intermediate_size,
517
+ tp_world_size=get_tensor_model_parallel_world_size(),
518
+ n_groups=0,
519
+ num_heads=self.num_heads,
520
+ head_dim=self.head_dim,
521
+ state_size=self.ssm_state_size,
522
+ conv_kernel=self.conv_kernel_size,
523
+ )
524
+
525
+ @property
526
+ def mamba_type(self) -> str:
527
+ return "mamba2"
528
+
529
+ def get_attn_backend(self) -> type["AttentionBackend"]:
530
+ from vllm.v1.attention.backends.mamba2_attn import (
531
+ Mamba2AttentionBackend)
532
+ return Mamba2AttentionBackend
533
+
534
+
535
+def plamo2_mamba_mixer(
536
+ hidden_states: torch.Tensor,
537
+ output: torch.Tensor,
538
+ layer_name: str,
539
+) -> None:
540
+ forward_context: ForwardContext = get_forward_context()
541
+ self = forward_context.no_compile_layers[layer_name]
542
+ self.forward_cuda(hidden_states=hidden_states,
543
+ output=output,
544
+ mamba_cache_params=None,
545
+ mamba2_metadata=None)
546
+
547
+
548
+def plamo2_mamba_mixer_fake(
549
+ hidden_states: torch.Tensor,
550
+ output: torch.Tensor,
551
+ layer_name: str,
552
+) -> None:
553
+ return
554
+
555
+
556
+direct_register_custom_op(
557
+ op_name="plamo2_mamba_mixer",
558
+ op_func=plamo2_mamba_mixer,
559
+ mutates_args=["output"],
560
+ fake_impl=plamo2_mamba_mixer_fake,
561
+ dispatch_key=current_platform.dispatch_key,
562
+)
563
564
565
class DenseMLP(nn.Module):
@@ -418,7 +595,6 @@ class DenseMLP(nn.Module):
595
return self.down_proj(h)
596
597
421
-@support_torch_compile
598
class Plamo2AttentionMixer(nn.Module):
599
600
def __init__(self,
@@ -575,12 +751,24 @@ class Plamo2DecoderLayer(nn.Module):
751
hidden_states, residual = self.pre_mixer_norm(
752
hidden_states, residual)
753
754
+ if self.is_mamba:
755
+ # Plamo2MambaMixer writes output to this tensor
756
+ output = torch.empty_like(hidden_states)
757
+ mixer_kwargs = {
758
+ "output": output,
759
+ "mamba_cache_params": mamba_cache_params,
760
+ "mamba2_metadata": mamba2_metadata,
761
+ }
762
+ else:
763
+ mixer_kwargs = {
764
+ "positions": positions,
765
+ }
766
hidden_states = self.mixer(
579
- positions=positions,
767
hidden_states=hidden_states,
581
- mamba_cache_params=mamba_cache_params,
582
- mamba2_metadata=mamba2_metadata,
768
+ **mixer_kwargs,
769
)
770
+ if self.is_mamba:
771
+ hidden_states = output
772
hidden_states = self.post_mixer_norm(hidden_states)
773
# Fully Connected
774
hidden_states, residual = self.pre_mlp_norm(hidden_states, residual)
@@ -591,7 +779,7 @@ class Plamo2DecoderLayer(nn.Module):
779
780
class Plamo2Decoder(torch.nn.Module):
781
594
- def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
782
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
783
super().__init__()
784
config = vllm_config.model_config.hf_config
785
extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}
@@ -617,7 +805,7 @@ class Plamo2Decoder(torch.nn.Module):
805
mamba_cache_index = 0
806
for layer in islice(self.layers, self.start_layer, self.end_layer):
807
layer_mamba_cache_params = None
620
- if layer.is_mamba:
808
+ if layer.is_mamba and mamba_cache_params is not None:
809
layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
810
mamba_cache_index)
811
mamba_cache_index += 1
@@ -632,10 +820,11 @@ class Plamo2Decoder(torch.nn.Module):
820
return hidden_states, residual
821
822
635
-class Plamo2Model(Plamo2PreTrainedModel):
823
+@support_torch_compile
824
+class Plamo2Model(torch.nn.Module):
825
826
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
638
- super().__init__(vllm_config.model_config.hf_config)
827
+ super().__init__()
828
829
config = vllm_config.model_config.hf_config
830
@@ -653,9 +842,9 @@ class Plamo2Model(Plamo2PreTrainedModel):
842
self.make_empty_intermediate_tensors = (
843
make_empty_intermediate_tensors_factory(
844
["hidden_states", "residual"], config.hidden_size))
656
- self.layers = Plamo2Decoder(vllm_config, prefix=f"{prefix}.layers")
845
+ self.layers = Plamo2Decoder(vllm_config=vllm_config,
846
+ prefix=f"{prefix}.layers")
847
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
658
- self.post_init()
848
849
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
850
return self.embed_tokens(input_ids)
@@ -679,11 +868,16 @@ class Plamo2Model(Plamo2PreTrainedModel):
868
hidden_states = intermediate_tensors["hidden_states"]
869
residual = intermediate_tensors["residual"]
870
682
- attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
683
- mamba2_metadata = prepare_mamba2_metadata(
684
- chunk_size=self.config.mamba_chunk_size,
685
- attn_metadata=attn_metadata,
686
- )
871
+ if not envs.VLLM_USE_V1:
872
+ attn_metadata: AttentionMetadata = get_forward_context(
873
+ ).attn_metadata
874
+ mamba2_metadata = prepare_mamba2_metadata(
875
+ chunk_size=self.config.mamba_chunk_size,
876
+ attn_metadata=attn_metadata,
877
+ )
878
+ else:
879
+ # v1 get mamba2_metadata from forward_context
880
+ mamba2_metadata = None
881
882
hidden_states, residual = self.layers(
883
positions=positions,
@@ -701,8 +895,7 @@ class Plamo2Model(Plamo2PreTrainedModel):
895
return hidden_states
896
897
704
-class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, SupportsPP,
705
- IsHybrid, SupportsV0Only):
898
+class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):
899
packed_modules_mapping = {
900
"qkv_proj": [
901
"q_proj",
@@ -712,12 +905,10 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, SupportsPP,
905
}
906
907
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
908
+ super().__init__()
909
config = vllm_config.model_config.hf_config
910
scheduler_config = vllm_config.scheduler_config
717
- assert not vllm_config.cache_config.enable_prefix_caching, \
718
- "PLaMo2 currently does not support prefix caching"
911
720
- super().__init__(config)
912
self.config = config
913
self.vllm_config = vllm_config
914
self.model_config = vllm_config.model_config
@@ -751,8 +942,6 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, SupportsPP,
942
self.sampler = get_sampler()
943
self.make_empty_intermediate_tensors = (
944
self.model.make_empty_intermediate_tensors)
754
- # Initialize weights and apply final processing
755
- self.post_init()
945
946
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
947
return self.model.get_input_embeddings(input_ids)
@@ -763,19 +952,27 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, SupportsPP,
952
intermediate_tensors: Optional[IntermediateTensors] = None,
953
inputs_embeds: Optional[torch.Tensor] = None,
954
**kwargs):
766
- if self.mamba_cache is None:
767
- num_mamba_layers = self.model_config.get_num_layers_by_block_type(
768
- self.vllm_config.parallel_config, LayerBlockType.mamba)
769
-
770
- self.mamba_cache = MambaCacheManager(
771
- self.vllm_config,
772
- num_mamba_layers,
773
- *self._get_mamba_cache_shape(),
774
- self.lm_head.weight.dtype,
775
- self.lm_head.weight.dtype,
776
- )
777
-
778
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
955
+ if not envs.VLLM_USE_V1:
956
+ if self.mamba_cache is None:
957
+ num_mamba_layers = (
958
+ self.model_config.get_num_layers_by_block_type(
959
+ self.vllm_config.parallel_config,
960
+ LayerBlockType.mamba))
961
+
962
+ mamba_state_shape = self.get_mamba_state_shape_from_config(
963
+ self.vllm_config, use_v1=False)
964
+ mamba_state_dtype = \
965
+ self.get_mamba_state_dtype_from_config(
966
+ self.vllm_config)
967
+ self.mamba_cache = MambaCacheManager(self.vllm_config,
968
+ num_mamba_layers,
969
+ *mamba_state_shape,
970
+ *mamba_state_dtype)
971
+
972
+ mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
973
+ else:
974
+ # NOTE: mamba_cache_params is not needed for v1
975
+ mamba_cache_params = None
976
977
hidden_states = self.model(input_ids, positions, mamba_cache_params,
978
intermediate_tensors, inputs_embeds)
@@ -788,21 +985,48 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, SupportsPP,
985
def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
986
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
987
791
- def _get_mamba_cache_shape(
792
- self) -> tuple[tuple[int, int], tuple[int, int, int]]:
793
- world_size = get_tensor_model_parallel_world_size()
794
- hidden_size = (self.config.mamba_num_heads *
795
- self.config.hidden_size_per_head)
796
- conv_state_shape = (
797
- hidden_size // world_size,
798
- self.config.mamba_d_conv - 1,
988
+ @classmethod
989
+ def get_mamba_state_dtype_from_config(
990
+ cls,
991
+ vllm_config: "VllmConfig",
992
+ ) -> tuple[torch.dtype, torch.dtype]:
993
+
994
+ return MambaStateDtypeCalculator.mamba2_state_dtype(
995
+ vllm_config.model_config.dtype,
996
+ vllm_config.cache_config.mamba_cache_dtype,
997
+ vllm_config.cache_config.mamba_ssm_cache_dtype,
998
)
800
- temporal_state_shape = (
801
- divide(self.config.mamba_num_heads, world_size),
802
- self.config.hidden_size_per_head,
803
- self.config.mamba_d_state,
999
+
1000
+ @classmethod
1001
+ def get_mamba_state_shape_from_config(
1002
+ cls,
1003
+ vllm_config: "VllmConfig",
1004
+ use_v1: bool = True,
1005
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
1006
+ """Calculate shapes for Mamba's convolutional and state caches.
1007
+ Args:
1008
+ vllm_config: vLLM config
1009
+ use_v1: Get shapes for V1 (or V0)
1010
+ Returns:
1011
+ Tuple containing:
1012
+ - conv_state_shape: Shape for convolutional state cache
1013
+ - temporal_state_shape: Shape for state space model cache
1014
+ """
1015
+ parallel_config = vllm_config.parallel_config
1016
+ hf_config = vllm_config.model_config.hf_config
1017
+ intermediate_size =\
1018
+ hf_config.mamba_num_heads * hf_config.hidden_size_per_head
1019
+
1020
+ return MambaStateShapeCalculator.mamba2_state_shape(
1021
+ intermediate_size=intermediate_size,
1022
+ tp_world_size=parallel_config.tensor_parallel_size,
1023
+ n_groups=0,
1024
+ num_heads=hf_config.mamba_num_heads,
1025
+ head_dim=hf_config.hidden_size_per_head,
1026
+ state_size=hf_config.mamba_d_state,
1027
+ conv_kernel=hf_config.mamba_d_conv,
1028
+ use_v1=use_v1,
1029
)
805
- return conv_state_shape, temporal_state_shape
1030
1031
def compute_logits(
1032
self,