[Model] Remove Plamo2 (#49729)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 24, 2026 at 21:54 UTC
84d26b9ee3dde1991442539dcf5fc5b26aa08f7e
13 files changed
+6
-1014
.buildkite/test-amd.yaml
+2
-4
@@ -81,10 +81,8 @@
81
# the above test.) Also run if model initialization test file is modified. #
82
# * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model #
83
# source is modified, or when specified test files are modified. #
84
-# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to #
85
-# run plamo2 model in vLLM. #
86
-# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d) #
87
-# and to run plamo2 model in vLLM. #
84
+# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d). #
85
+# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d). #
86
# * [Multi-Modal Models (Standard) 1-4]: #
87
# - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. #
88
# * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. #
.buildkite/test_areas/models_language.yaml
-2
@@ -63,7 +63,6 @@ steps:
63
- tests/models/language/generation
64
commands:
65
# Install fast path packages for testing against transformers
66
- # Note: also needed to run plamo2 model in vLLM
66
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
67
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
68
# Shard the hybrid language model tests that are numerically stable on Hopper.
@@ -105,7 +104,6 @@ steps:
104
- tests/models/language/generation
105
commands:
106
# Install fast path packages for testing against transformers
108
- # Note: also needed to run plamo2 model in vLLM
107
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
108
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
109
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
docs/design/custom_op.md
-2
@@ -122,8 +122,6 @@ For example:
122
123
--8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm"
124
125
---8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer"
126
-
125
--8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv"
126
```
127
docs/models/supported_models.md
-1
@@ -435,7 +435,6 @@ th {
435
| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
436
| `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
437
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
438
-| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ |
438
| `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ |
439
| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ |
440
| `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ |
docs/usage/v1_guide.md
+1
-1
@@ -126,7 +126,7 @@ Models using selective state-space mechanisms instead of standard transformer at
126
Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`, `FalconMambaForCausalLM`) are supported.
127
128
Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g.,
129
-`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
129
+`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`).
130
131
Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).
132
tests/distributed/test_pipeline_parallel.py
-1
@@ -121,7 +121,6 @@ TEXT_GENERATION_MODELS = {
121
"ibm/PowerMoE-3b": PPTestSettings.fast(),
122
"internlm/internlm2-chat-7b": PPTestSettings.fast(),
123
"ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),
124
- "pfnet/plamo-2-1b": PPTestSettings.fast(),
124
"pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(),
125
"meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(),
126
# Tests TransformersForCausalLM
tests/models/language/generation/test_hybrid.py
-2
@@ -35,7 +35,6 @@ SSM_MODELS = [
35
36
HYBRID_MODELS = [
37
"ai21labs/Jamba-tiny-dev",
38
- "pfnet/plamo-2-1b",
38
"Zyphra/Zamba2-1.2B-instruct",
39
"ibm-granite/granite-4.0-tiny-preview",
40
"tiiuae/Falcon-H1-0.5B-Base",
@@ -50,7 +49,6 @@ HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = {
49
50
FULL_CUDA_GRAPH_MODELS = [
51
"ai21labs/Jamba-tiny-dev",
53
- "pfnet/plamo-2-1b",
52
"Zyphra/Zamba2-1.2B-instruct",
53
]
54
tests/models/registry.py
-4
@@ -484,10 +484,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
484
"PhiMoEForCausalLM": _HfExamplesInfo(
485
"microsoft/Phi-3.5-MoE-instruct", trust_remote_code=True
486
),
487
- "Plamo2ForCausalLM": _HfExamplesInfo(
488
- "pfnet/plamo-2-1b",
489
- trust_remote_code=True,
490
- ),
487
"Plamo3ForCausalLM": _HfExamplesInfo(
488
"pfnet/plamo-3-nict-2b-base",
489
trust_remote_code=True,
tests/quantization/test_experts_int8.py
+1
-1
@@ -12,7 +12,7 @@ from tests.quantization.utils import is_quant_method_supported
12
13
from ..models.registry import HF_EXAMPLE_MODELS
14
15
-MODELS = ["ai21labs/Jamba-tiny-random", "pfnet/plamo-2-1b"]
15
+MODELS = ["ai21labs/Jamba-tiny-random"]
16
17
18
@pytest.mark.skipif(
vllm/config/compilation.py
-1
@@ -768,7 +768,6 @@ class CompilationConfig:
768
"vllm::mamba_mixer",
769
"vllm::short_conv",
770
"vllm::linear_attention",
771
- "vllm::plamo2_mamba_mixer",
771
"vllm::qwen_gdn_attention_core",
772
"vllm::gdn_attention_core_xpu",
773
"vllm::olmo_hybrid_gdn_full_forward",
vllm/config/model.py
+1
-2
@@ -1448,7 +1448,7 @@ class ModelConfig:
1448
"""
1449
Returns the mamba chunk size if it exists
1450
"""
1451
- # used by e.g. Bamba, FalconH1, Granite, PLaMo2
1451
+ # used by e.g. Bamba, FalconH1, Granite
1452
chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None)
1453
if chunk_size is None:
1454
# used by e.g. Mamba2, NemotronH, Zamba
@@ -2022,7 +2022,6 @@ _FLOAT16_NOT_SUPPORTED_MODELS = {
2022
"gemma2": "Numerical instability. Please use bfloat16 or float32 instead.",
2023
"gemma3": "Numerical instability. Please use bfloat16 or float32 instead.",
2024
"gemma3_text": "Numerical instability. Please use bfloat16 or float32 instead.",
2025
- "plamo2": "Numerical instability. Please use bfloat16 or float32 instead.",
2025
"glm4": "Numerical instability. Please use bfloat16 or float32 instead.",
2026
}
2027
vllm/model_executor/models/plamo2.py
deleted
-992
@@ -1,992 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Inference-only PLaMo2 model."""
4
-
5
-from collections.abc import Iterable
6
-from itertools import islice
7
-from typing import TYPE_CHECKING
8
-
9
-import torch
10
-from torch import nn
11
-from transformers import PretrainedConfig
12
-
13
-from vllm.compilation.decorators import support_torch_compile
14
-from vllm.config import VllmConfig, get_current_vllm_config
15
-from vllm.distributed import divide, get_tensor_model_parallel_world_size
16
-from vllm.distributed.parallel_state import get_pp_group
17
-from vllm.forward_context import ForwardContext, get_forward_context
18
-from vllm.model_executor.custom_op import PluggableLayer
19
-from vllm.model_executor.layers.activation import SiluAndMul
20
-from vllm.model_executor.layers.attention import Attention
21
-from vllm.model_executor.layers.layernorm import RMSNorm
22
-from vllm.model_executor.layers.linear import (
23
- ColumnParallelLinear,
24
- MergedColumnParallelLinear,
25
- QKVParallelLinear,
26
- RowParallelLinear,
27
-)
28
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
29
-from vllm.model_executor.layers.mamba.abstract import MambaBase
30
-from vllm.model_executor.layers.mamba.mamba_utils import (
31
- MambaStateCopyFunc,
32
- MambaStateCopyFuncCalculator,
33
- MambaStateDtypeCalculator,
34
- MambaStateShapeCalculator,
35
- is_conv_state_dim_first,
36
-)
37
-from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
38
- causal_conv1d_fn,
39
- causal_conv1d_update,
40
-)
41
-from vllm.model_executor.layers.mamba.ops.ssd_combined import (
42
- mamba_chunk_scan_combined_varlen,
43
-)
44
-from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update
45
-from vllm.model_executor.layers.quantization import QuantizationConfig
46
-from vllm.model_executor.layers.rotary_embedding import get_rope
47
-from vllm.model_executor.layers.vocab_parallel_embedding import (
48
- ParallelLMHead,
49
- VocabParallelEmbedding,
50
-)
51
-from vllm.model_executor.model_loader.weight_utils import (
52
- composed_weight_loader,
53
- default_weight_loader,
54
- sharded_weight_loader,
55
-)
56
-from vllm.model_executor.models.interfaces import (
57
- HasInnerState,
58
- IsHybrid,
59
- SupportsLoRA,
60
- SupportsPP,
61
-)
62
-from vllm.model_executor.models.utils import (
63
- AutoWeightsLoader,
64
- is_pp_missing_parameter,
65
- make_empty_intermediate_tensors_factory,
66
- make_layers,
67
- maybe_prefix,
68
-)
69
-from vllm.model_executor.utils import set_weight_attrs
70
-from vllm.platforms import current_platform
71
-from vllm.sequence import IntermediateTensors
72
-from vllm.utils.torch_utils import direct_register_custom_op
73
-from vllm.v1.attention.backend import AttentionMetadata
74
-from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
75
-from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
76
-
77
-# Only used for type hinting.
78
-if TYPE_CHECKING:
79
-
80
- class Plamo2Config(PretrainedConfig): # type: ignore
81
- model_type: str = "plamo2"
82
-
83
- hidden_size: int
84
- num_hidden_layers: int
85
- rms_norm_eps: float
86
- # Attention
87
- num_attention_heads: int
88
- hidden_size_per_head: int
89
- num_key_value_heads: int
90
- # Mamba
91
- mamba_d_state: int
92
- mamba_d_conv: int
93
- mamba_num_heads: int
94
- mamba_step: int
95
- # MLP
96
- intermediate_size: int
97
- # Tokenizer
98
- vocab_size: int
99
-
100
-
101
-def is_mamba(config: "Plamo2Config", i: int) -> bool:
102
- assert config.mamba_step > 1
103
-
104
- if config.num_hidden_layers <= (config.mamba_step // 2):
105
- # use attention in last layer
106
- return i != config.num_hidden_layers - 1
107
- return (i % config.mamba_step) != (config.mamba_step // 2)
108
-
109
-
110
-# Adapted from:
111
-# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
112
-# transformers.models.mamba.modeling_mamba.MambaMixer
113
-# --8<-- [start:plamo2_mamba_mixer]
114
-@PluggableLayer.register("plamo2_mamba_mixer")
115
-class Plamo2MambaMixer(MambaBase, PluggableLayer):
116
- # --8<-- [end:plamo2_mamba_mixer]
117
-
118
- def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None:
119
- super().__init__()
120
- self.config = vllm_config.model_config.hf_config
121
- self.cache_config = vllm_config.cache_config
122
- self.model_config = vllm_config.model_config
123
- self.quant_config = vllm_config.quant_config
124
- self.is_lora_enabled = bool(vllm_config.lora_config)
125
- self.hidden_size = self.config.hidden_size
126
- self.ssm_state_size = self.config.mamba_d_state
127
- self.conv_kernel_size = self.config.mamba_d_conv
128
- self.intermediate_size = (
129
- self.config.mamba_num_heads * self.config.hidden_size_per_head
130
- )
131
- self.tp_size = get_tensor_model_parallel_world_size()
132
- self.head_dim = self.config.hidden_size_per_head
133
- self.num_heads = self.config.mamba_num_heads
134
- self.time_step_rank = max(64, self.hidden_size // 16)
135
- self.conv1d = ColumnParallelLinear(
136
- input_size=self.conv_kernel_size,
137
- output_size=self.intermediate_size,
138
- bias=False,
139
- prefix=f"{prefix}.conv1d",
140
- return_bias=False,
141
- )
142
- # unsqueeze to fit conv1d weights shape into the linear weights shape.
143
- # Can't do this in `weight_loader` since it already exists in
144
- # `ColumnParallelLinear` and `set_weight_attrs`
145
- # doesn't allow to override it
146
- self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
147
-
148
- self.in_proj = MergedColumnParallelLinear(
149
- self.hidden_size,
150
- [self.intermediate_size] * 2,
151
- bias=False,
152
- quant_config=self.quant_config,
153
- prefix=f"{prefix}.in_proj",
154
- return_bias=False,
155
- )
156
- # selective projection used to make dt, B and C input dependent
157
- self.bcdt_proj = RowParallelLinear(
158
- self.intermediate_size,
159
- self.time_step_rank + self.ssm_state_size * 2,
160
- bias=False,
161
- quant_config=self.quant_config,
162
- prefix=f"{prefix}.bcdt_proj",
163
- return_bias=False,
164
- )
165
- # time step projection (discretization) -
166
- # In the forward we need to apply dt_proj without the bias,
167
- # as the bias is added in the selective scan kernel.
168
- self.dt_proj = ColumnParallelLinear(
169
- self.time_step_rank,
170
- self.num_heads,
171
- bias=False,
172
- quant_config=self.quant_config,
173
- prefix=f"{prefix}.dt_proj",
174
- return_bias=False,
175
- )
176
-
177
- self.A = nn.Parameter(
178
- torch.empty(
179
- divide(self.num_heads, self.tp_size),
180
- dtype=torch.float32,
181
- )
182
- )
183
- self.D = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size)))
184
- self.dt_bias = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size)))
185
-
186
- set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)})
187
- a_weight_loader = composed_weight_loader(
188
- sharded_weight_loader(0), lambda x: -torch.exp(x.float())
189
- )
190
- set_weight_attrs(self.A, {"weight_loader": a_weight_loader})
191
- set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
192
-
193
- self.out_proj = RowParallelLinear(
194
- self.intermediate_size,
195
- self.hidden_size,
196
- bias=False,
197
- input_is_parallel=True,
198
- quant_config=self.quant_config,
199
- prefix=f"{prefix}.out_proj",
200
- return_bias=False,
201
- )
202
- # The activation function is fixed to SiLU.
203
- self.activation = "silu"
204
-
205
- self.dt_norm = RMSNorm(self.time_step_rank, eps=self.config.rms_norm_eps)
206
- self.B_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps)
207
- self.C_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps)
208
-
209
- self.chunk_size = self.config.mamba_chunk_size
210
-
211
- compilation_config = get_current_vllm_config().compilation_config
212
- if prefix in compilation_config.static_forward_context:
213
- raise ValueError(f"Duplicate layer name: {prefix}")
214
- compilation_config.static_forward_context[prefix] = self
215
- # The tuple is (conv_state, ssm_state)
216
- self.kv_cache = (torch.tensor([]), torch.tensor([]))
217
- assert self.chunk_size != -1, "chunk_size must be set for v1"
218
-
219
- self.prefix = prefix
220
-
221
- def _project_ssm_parameters(self, hidden_states):
222
- if self.is_lora_enabled:
223
- # Lora kernel requires contiguous tensor.
224
- ssm_parameters = self.bcdt_proj(hidden_states.contiguous())
225
- else:
226
- ssm_parameters = self.bcdt_proj(hidden_states)
227
- B, C, time_step = torch.split(
228
- ssm_parameters,
229
- [self.ssm_state_size, self.ssm_state_size, self.time_step_rank],
230
- dim=-1,
231
- )
232
-
233
- # vllm._custom_ops.rms_norm requires contiguous input tensors.
234
- time_step = self.dt_norm(time_step.contiguous())
235
- B = self.B_norm(B.contiguous())
236
- C = self.C_norm(C.contiguous())
237
- dt = self.dt_proj(time_step)
238
- return B, C, dt
239
-
240
- def forward(
241
- self,
242
- hidden_states: torch.Tensor,
243
- output: torch.Tensor,
244
- **kwargs,
245
- ):
246
- torch.ops.vllm.plamo2_mamba_mixer(
247
- hidden_states,
248
- output,
249
- self.prefix,
250
- )
251
-
252
- def forward_impl(
253
- self,
254
- hidden_states: torch.Tensor,
255
- output: torch.Tensor,
256
- **kwargs,
257
- ):
258
- forward_context = get_forward_context()
259
- # attn_metadata contains metadata necessary for the mamba2 triton
260
- # kernels to operate in continuous batching and in chunked prefill
261
- # modes; they are computed at top-level model forward since they
262
- # stay the same and reused for all mamba layers in the same iteration
263
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
264
-
265
- if attn_metadata is not None:
266
- assert isinstance(attn_metadata, dict)
267
- attn_metadata = attn_metadata[self.prefix]
268
- assert isinstance(attn_metadata, Mamba2AttentionMetadata)
269
- self_kv_cache = self.kv_cache
270
- # conv_state = (..., dim, width-1) yet contiguous along 'dim'
271
- # conv_state must be (..., dim, width-1) for the conv kernels.
272
- # DS layout stores it that way directly; SD layout needs a transpose.
273
- conv_state = (
274
- self_kv_cache[0]
275
- if is_conv_state_dim_first()
276
- else self_kv_cache[0].transpose(-1, -2)
277
- )
278
- ssm_state = self_kv_cache[1]
279
- state_indices_tensor_p = attn_metadata.state_indices_tensor_p
280
- state_indices_tensor_d = attn_metadata.state_indices_tensor_d
281
- has_initial_states_p = attn_metadata.has_initial_states_p
282
- prep_initial_states = attn_metadata.prep_initial_states
283
- chunk_size = attn_metadata.chunk_size
284
- seq_idx_p = attn_metadata.seq_idx_p
285
- query_start_loc_p = attn_metadata.query_start_loc_p
286
- cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
287
- last_chunk_indices_p = attn_metadata.last_chunk_indices_p
288
-
289
- # 1. Gated MLP's linear projection
290
- projected_states = self.in_proj(hidden_states)
291
- gate, hidden_states = projected_states.chunk(2, dim=-1)
292
-
293
- # 2. Convolution sequence transformation
294
- conv_weights = self.conv1d.weight.view(
295
- self.conv1d.weight.size(0), self.conv1d.weight.size(2)
296
- )
297
-
298
- if attn_metadata is None:
299
- # profile run
300
- hidden_states = (
301
- hidden_states.transpose(0, 1).clone().transpose(0, 1)
302
- ).contiguous()
303
- output[:] = self.out_proj(hidden_states)
304
- return
305
-
306
- num_prefills = attn_metadata.num_prefills # request count
307
- num_decodes = attn_metadata.num_decode_tokens # token count (=request)
308
- num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
309
- has_prefill = num_prefills > 0
310
- has_decode = num_decodes > 0
311
- num_actual_tokens = num_prefill_tokens + num_decodes
312
-
313
- # Separate prefill and decode by splitting varlen input
314
- # Split along token dimension
315
- hidden_states_d, hidden_states_p = torch.split(
316
- hidden_states[:num_actual_tokens],
317
- [num_decodes, num_prefill_tokens],
318
- dim=0,
319
- )
320
- gate_d, gate_p = torch.split(
321
- gate[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0
322
- )
323
- # Preallocate output tensor to avoid memcpy cost for merging prefill
324
- # and decode outputs
325
- preallocated_ssm_out = torch.empty(
326
- [
327
- num_prefill_tokens + num_decodes,
328
- (self.num_heads // self.tp_size) * self.head_dim,
329
- ],
330
- dtype=hidden_states.dtype,
331
- device=hidden_states.device,
332
- )
333
- preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
334
- preallocated_ssm_out,
335
- [num_decodes, num_prefill_tokens],
336
- dim=0,
337
- )
338
-
339
- # Process prefill requests
340
- if has_prefill:
341
- # 2. Convolution sequence transformation
342
- # - "cache_indices" updates the conv_state cache in positions
343
- # pointed to by "state_indices_tensor_p"
344
- x = hidden_states_p.transpose(0, 1) # this is the form that causal-conv see
345
- hidden_states_p = causal_conv1d_fn(
346
- x,
347
- conv_weights,
348
- self.conv1d.bias,
349
- activation=self.activation,
350
- conv_states=conv_state,
351
- has_initial_state=has_initial_states_p,
352
- cache_indices=state_indices_tensor_p,
353
- metadata=attn_metadata,
354
- query_start_loc=query_start_loc_p,
355
- )
356
- hidden_states_p = hidden_states_p.transpose(0, 1)
357
- hidden_states_p = hidden_states_p[:num_prefill_tokens]
358
- # In some instances, the following `bcdt_proj` op
359
- # requires contiguous inputs
360
- # (e.g. if the Marlin kernel is used).
361
- hidden_states_p = hidden_states_p.contiguous()
362
-
363
- B, C, dt = self._project_ssm_parameters(hidden_states_p)
364
-
365
- # 3. State Space Model sequence transformation
366
- initial_states = None
367
- if has_initial_states_p is not None and prep_initial_states:
368
- # making a copy of the states
369
- initial_states = torch.where(
370
- has_initial_states_p[:, None, None, None],
371
- ssm_state[state_indices_tensor_p],
372
- 0,
373
- )
374
-
375
- varlen_state = mamba_chunk_scan_combined_varlen(
376
- hidden_states_p.view(
377
- num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim
378
- ),
379
- dt,
380
- self.A,
381
- B.view(num_prefill_tokens, 1, -1),
382
- C.view(num_prefill_tokens, 1, -1),
383
- chunk_size=chunk_size,
384
- D=self.D,
385
- z=gate_p.view(
386
- num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim
387
- ),
388
- dt_bias=self.dt_bias,
389
- seq_idx=seq_idx_p,
390
- cu_seqlens=query_start_loc_p,
391
- cu_chunk_seqlens=cu_chunk_seqlen_p,
392
- last_chunk_indices=last_chunk_indices_p,
393
- initial_states=initial_states,
394
- dt_softplus=True,
395
- dt_limit=(0.0, float("inf")),
396
- out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim),
397
- state_dtype=ssm_state.dtype,
398
- )
399
-
400
- # update ssm states
401
- # - varlen state is a (batch, nheads, headdim, dstate) tensor
402
- ssm_state[state_indices_tensor_p] = varlen_state
403
-
404
- # Process decode requests
405
- if has_decode:
406
- # 2. Convolution sequence transformation
407
- hidden_states_d = causal_conv1d_update(
408
- hidden_states_d,
409
- conv_state,
410
- conv_weights,
411
- self.conv1d.bias,
412
- self.activation,
413
- conv_state_indices=state_indices_tensor_d,
414
- )
415
-
416
- # ROCm: Ensure contiguous tensor for bcdt_proj linear layer.
417
- # causal_conv1d_update returns a non-contiguous view (stride 8192
418
- # instead of 4096 for shape [batch, 4096]), causing incorrect GEMM
419
- # results when batch > 1 on ROCm.
420
- if current_platform.is_rocm():
421
- hidden_states_d = hidden_states_d.contiguous()
422
-
423
- B, C, dt = self._project_ssm_parameters(hidden_states_d)
424
-
425
- # 3. State Space Model sequence transformation
426
- A = self.A[:, None, ...][:, :, None].expand(
427
- -1, self.head_dim, self.config.mamba_d_state
428
- )
429
- dt = dt[:, :, None].expand(-1, -1, self.head_dim)
430
- dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
431
- D = self.D[:, None, ...].expand(-1, self.head_dim)
432
- B = B.unsqueeze(1)
433
- C = C.unsqueeze(1)
434
- hidden_states_d = hidden_states_d.view(
435
- -1, self.num_heads // self.tp_size, self.head_dim
436
- )
437
-
438
- # - the hidden is reshaped into (bs, num_heads, head_dim)
439
- # - ssm_state's slots will be selected
440
- # using state_indices_tensor_d
441
-
442
- # NOTE: final output is an in-place update of out tensor
443
- selective_state_update(
444
- ssm_state,
445
- hidden_states_d,
446
- dt,
447
- A,
448
- B,
449
- C,
450
- D,
451
- dt_bias,
452
- z=gate_d.reshape(num_decodes, -1, self.head_dim),
453
- dt_softplus=True,
454
- state_batch_indices=state_indices_tensor_d,
455
- out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),
456
- )
457
-
458
- # 4. Final linear projection
459
- output[:num_actual_tokens] = self.out_proj(preallocated_ssm_out)
460
-
461
- def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
462
- assert self.model_config is not None
463
- assert self.cache_config is not None
464
- return MambaStateDtypeCalculator.mamba2_state_dtype(
465
- self.model_config.dtype,
466
- self.cache_config.mamba_cache_dtype,
467
- self.cache_config.mamba_ssm_cache_dtype,
468
- )
469
-
470
- def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
471
- return MambaStateShapeCalculator.mamba2_state_shape(
472
- intermediate_size=self.intermediate_size,
473
- tp_world_size=get_tensor_model_parallel_world_size(),
474
- n_groups=0,
475
- num_heads=self.num_heads,
476
- head_dim=self.head_dim,
477
- state_size=self.ssm_state_size,
478
- conv_kernel=self.conv_kernel_size,
479
- )
480
-
481
- @property
482
- def mamba_type(self) -> MambaAttentionBackendEnum:
483
- return MambaAttentionBackendEnum.MAMBA2
484
-
485
-
486
-def plamo2_mamba_mixer(
487
- hidden_states: torch.Tensor,
488
- output: torch.Tensor,
489
- layer_name: str,
490
-) -> None:
491
- forward_context: ForwardContext = get_forward_context()
492
- self = forward_context.no_compile_layers[layer_name]
493
- self.forward_impl(hidden_states=hidden_states, output=output)
494
-
495
-
496
-def plamo2_mamba_mixer_fake(
497
- hidden_states: torch.Tensor,
498
- output: torch.Tensor,
499
- layer_name: str,
500
-) -> None:
501
- return
502
-
503
-
504
-direct_register_custom_op(
505
- op_name="plamo2_mamba_mixer",
506
- op_func=plamo2_mamba_mixer,
507
- mutates_args=["output"],
508
- fake_impl=plamo2_mamba_mixer_fake,
509
-)
510
-
511
-
512
-class DenseMLP(nn.Module):
513
- def __init__(
514
- self,
515
- config: "Plamo2Config",
516
- quant_config: QuantizationConfig | None = None,
517
- prefix: str = "",
518
- ) -> None:
519
- super().__init__()
520
- self.hidden_size = config.hidden_size
521
- self.intermediate_size = config.intermediate_size
522
- self.gate_up_proj = MergedColumnParallelLinear(
523
- self.hidden_size,
524
- [self.intermediate_size] * 2,
525
- bias=False,
526
- prefix=f"{prefix}.gate_up_proj",
527
- quant_config=quant_config,
528
- return_bias=False,
529
- )
530
- self.act = SiluAndMul()
531
- self.down_proj = RowParallelLinear(
532
- self.intermediate_size,
533
- self.hidden_size,
534
- bias=False,
535
- prefix=f"{prefix}.down_proj",
536
- quant_config=quant_config,
537
- return_bias=False,
538
- )
539
-
540
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
541
- h = self.gate_up_proj(hidden_states)
542
- h = self.act(h)
543
- return self.down_proj(h)
544
-
545
-
546
-class Plamo2AttentionMixer(nn.Module):
547
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
548
- super().__init__()
549
- config = vllm_config.model_config.hf_config
550
- cache_config = vllm_config.cache_config
551
- quant_config = vllm_config.quant_config
552
- self.hidden_size = config.hidden_size
553
- tp_size = get_tensor_model_parallel_world_size()
554
- self.total_num_heads = config.num_attention_heads
555
- assert self.total_num_heads % tp_size == 0
556
- self.num_heads = self.total_num_heads // tp_size
557
- self.total_num_kv_heads = config.num_key_value_heads
558
- if self.total_num_kv_heads >= tp_size:
559
- # Number of KV heads is greater than TP size, so we partition
560
- # the KV heads across multiple tensor parallel GPUs.
561
- assert self.total_num_kv_heads % tp_size == 0
562
- else:
563
- # Number of KV heads is less than TP size, so we replicate
564
- # the KV heads across multiple tensor parallel GPUs.
565
- assert tp_size % self.total_num_kv_heads == 0
566
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
567
- self.head_dim = config.hidden_size_per_head
568
- self.q_size = self.num_heads * self.head_dim
569
- self.kv_size = self.num_kv_heads * self.head_dim
570
- self.scaling = self.head_dim**-0.5
571
-
572
- self.qkv_proj = QKVParallelLinear(
573
- config.hidden_size,
574
- self.head_dim,
575
- self.total_num_heads,
576
- self.total_num_kv_heads,
577
- bias=False,
578
- quant_config=quant_config,
579
- prefix=f"{prefix}.qkv_proj",
580
- )
581
- self.o_proj = RowParallelLinear(
582
- self.total_num_heads * self.head_dim,
583
- config.hidden_size,
584
- bias=False,
585
- quant_config=quant_config,
586
- prefix=f"{prefix}.o_proj",
587
- )
588
-
589
- max_position = config.max_position_embeddings
590
- if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
591
- vllm_config.model_config.max_model_len, int
592
- ):
593
- max_position = min(max_position, vllm_config.model_config.max_model_len)
594
-
595
- self.rotary_emb = get_rope(
596
- self.head_dim,
597
- max_position=max_position,
598
- rope_parameters=config.rope_parameters,
599
- )
600
- self.q_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps)
601
- self.q_norm.weight = torch.nn.Parameter(
602
- torch.ones((self.num_heads, config.hidden_size_per_head))
603
- )
604
- set_weight_attrs(
605
- self.q_norm.weight, {"weight_loader": sharded_weight_loader(0)}
606
- )
607
- self.k_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps)
608
- self.k_norm.weight = torch.nn.Parameter(
609
- torch.ones((self.num_kv_heads, config.hidden_size_per_head))
610
- )
611
- # Tensor-parallelism shards the K norm weights to the tp ranks
612
- # in a head-wise manner. This approach does not work if there is only
613
- # a single KV head, as is the case for PLaMo 2-1B.
614
- if self.total_num_kv_heads != 1:
615
- set_weight_attrs(
616
- self.k_norm.weight, {"weight_loader": sharded_weight_loader(0)}
617
- )
618
-
619
- self.attn = Attention(
620
- self.num_heads,
621
- self.head_dim,
622
- self.scaling,
623
- num_kv_heads=self.num_kv_heads,
624
- cache_config=cache_config,
625
- prefix=f"{prefix}.attn",
626
- )
627
-
628
- def forward(
629
- self,
630
- positions: torch.Tensor,
631
- hidden_states: torch.Tensor,
632
- **kwargs,
633
- ) -> torch.Tensor:
634
- qkv, _ = self.qkv_proj(hidden_states)
635
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
636
-
637
- q_shape = q.shape
638
- q = q.reshape(q_shape[:-1] + self.q_norm.weight.shape)
639
- q = self.q_norm.forward_native(q).reshape(q_shape)
640
- k_shape = k.shape
641
- k = k.reshape(k_shape[:-1] + self.k_norm.weight.shape)
642
- k = self.k_norm.forward_native(k).reshape(k_shape)
643
-
644
- q, k = self.rotary_emb(positions, q, k)
645
- attn_output = self.attn(q, k, v)
646
- output, _ = self.o_proj(attn_output)
647
- return output
648
-
649
-
650
-class Plamo2DecoderLayer(nn.Module):
651
- def __init__(
652
- self, vllm_config: VllmConfig, layer_idx: int, prefix: str = "", **kwargs
653
- ) -> None:
654
- super().__init__()
655
- config = vllm_config.model_config.hf_config
656
- quant_config = vllm_config.quant_config
657
-
658
- self.is_mamba = is_mamba(config, layer_idx)
659
- if self.is_mamba:
660
- self.mixer = Plamo2MambaMixer(
661
- vllm_config=vllm_config, prefix=f"{prefix}.mixer"
662
- )
663
- else:
664
- self.mixer = Plamo2AttentionMixer(
665
- vllm_config=vllm_config, prefix=f"{prefix}.mixer"
666
- )
667
-
668
- self.mlp = DenseMLP(
669
- config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
670
- )
671
- self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
672
- self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
673
- self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
674
- self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
675
-
676
- def forward(
677
- self,
678
- positions: torch.Tensor,
679
- hidden_states: torch.Tensor,
680
- residual: torch.Tensor | None,
681
- **kwargs,
682
- ):
683
- if residual is None:
684
- residual = hidden_states
685
- hidden_states = self.pre_mixer_norm(hidden_states)
686
- else:
687
- hidden_states, residual = self.pre_mixer_norm(hidden_states, residual)
688
-
689
- if self.is_mamba:
690
- # Plamo2MambaMixer writes output to this tensor
691
- output = torch.empty_like(hidden_states)
692
- mixer_kwargs = {
693
- "output": output,
694
- }
695
- else:
696
- mixer_kwargs = {
697
- "positions": positions,
698
- }
699
- hidden_states = self.mixer(
700
- hidden_states=hidden_states,
701
- **mixer_kwargs,
702
- )
703
- if self.is_mamba:
704
- hidden_states = output
705
- hidden_states = self.post_mixer_norm(hidden_states)
706
- # Fully Connected
707
- hidden_states, residual = self.pre_mlp_norm(hidden_states, residual)
708
- hidden_states = self.mlp(hidden_states)
709
- hidden_states = self.post_mlp_norm(hidden_states)
710
- return hidden_states, residual
711
-
712
-
713
-class Plamo2Decoder(torch.nn.Module):
714
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
715
- super().__init__()
716
- config = vllm_config.model_config.hf_config
717
- extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}
718
-
719
- def get_layer(prefix: str):
720
- layer_idx = int(prefix.rsplit(".", 1)[1])
721
- return Plamo2DecoderLayer(
722
- vllm_config=vllm_config,
723
- layer_idx=layer_idx,
724
- prefix=prefix,
725
- **extra_kwargs,
726
- )
727
-
728
- self.start_layer, self.end_layer, self.layers = make_layers(
729
- config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
730
- )
731
-
732
- def forward(
733
- self,
734
- positions: torch.Tensor,
735
- hidden_states: torch.Tensor,
736
- residual: torch.Tensor | None,
737
- ) -> torch.Tensor:
738
- for layer in islice(self.layers, self.start_layer, self.end_layer):
739
- hidden_states, residual = layer(
740
- positions=positions,
741
- hidden_states=hidden_states,
742
- residual=residual,
743
- )
744
- return hidden_states, residual
745
-
746
-
747
-@support_torch_compile
748
-class Plamo2Model(torch.nn.Module):
749
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
750
- super().__init__()
751
-
752
- config = vllm_config.model_config.hf_config
753
-
754
- self.config = config
755
- self.vocab_size = config.vocab_size
756
-
757
- self.embed_tokens = VocabParallelEmbedding(
758
- self.vocab_size,
759
- config.hidden_size,
760
- prefix=f"{prefix}.embed_tokens",
761
- )
762
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
763
- ["hidden_states", "residual"], config.hidden_size
764
- )
765
- self.layers = Plamo2Decoder(vllm_config=vllm_config, prefix=f"{prefix}.layers")
766
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
767
-
768
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
769
- return self.embed_tokens(input_ids)
770
-
771
- def forward(
772
- self,
773
- input_ids: torch.Tensor | None,
774
- positions: torch.Tensor,
775
- intermediate_tensors: IntermediateTensors | None = None,
776
- inputs_embeds: torch.Tensor | None = None,
777
- ) -> torch.Tensor:
778
- if get_pp_group().is_first_rank:
779
- if inputs_embeds is not None:
780
- hidden_states = inputs_embeds
781
- else:
782
- hidden_states = self.embed_input_ids(input_ids)
783
- residual = None
784
- else:
785
- assert intermediate_tensors is not None
786
- hidden_states = intermediate_tensors["hidden_states"]
787
- residual = intermediate_tensors["residual"]
788
-
789
- hidden_states, residual = self.layers(
790
- positions=positions,
791
- hidden_states=hidden_states,
792
- residual=residual,
793
- )
794
- if not get_pp_group().is_last_rank:
795
- return IntermediateTensors(
796
- {"hidden_states": hidden_states, "residual": residual}
797
- )
798
- hidden_states, _ = self.norm(hidden_states, residual)
799
- return hidden_states
800
-
801
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
802
- params_dict = dict(self.named_parameters())
803
- loaded_params: set[str] = set()
804
- for name, loaded_weight in weights:
805
- # Update the weight names to be compatible with the vllm version
806
- # of the model.
807
- # Do not change the order of the replacements.
808
- replacements = {
809
- # Rename incompatible weight names.
810
- ".A_log": ".A",
811
- ".B_norm_weight": ".B_norm.weight",
812
- ".C_norm_weight": ".C_norm.weight",
813
- ".dt_norm_weight": ".dt_norm.weight",
814
- ".q_weight": ".q_norm.weight",
815
- ".k_weight": ".k_norm.weight",
816
- }
817
- # Apply replacements based on the defined mappings
818
- for old, new in replacements.items():
819
- if old in name:
820
- name = name.replace(old, new)
821
-
822
- # Reshape the in_proj weights to match the shape expected
823
- # by MergedColumnParallelLinear.
824
- # This works both for unquantized weights and
825
- # for quantized weights.
826
- # In the quantized case, the weights are already transposed.
827
- # Also, in addition to the quantized weights,
828
- # the zero points and scales have to be reshaped as well.
829
- # Packing should not be affected by this.
830
- if (
831
- ".mixer.in_proj.weight" in name
832
- or "mixer.in_proj.qweight" in name
833
- or "mixer.in_proj.scales" in name
834
- or "mixer.in_proj.qzeros" in name
835
- ):
836
- if "mixer.in_proj.weight" in name:
837
- loaded_weight = loaded_weight.transpose(0, 1)
838
- # for weight:
839
- # loaded_weight.shape[0] == self.config.hidden_size
840
- # for qweight:
841
- # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor # noqa
842
- # for scales and qzeros:
843
- # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size # noqa
844
- loaded_weight = loaded_weight.reshape(
845
- loaded_weight.shape[0], self.config.mamba_num_heads, -1
846
- )
847
- gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1)
848
- gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1)
849
- hidden_states_weight = hidden_states_weight.reshape(
850
- loaded_weight.shape[0], -1
851
- )
852
- loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1)
853
- if "mixer.in_proj.weight" in name:
854
- loaded_weight = loaded_weight.transpose(0, 1)
855
-
856
- # Offset parameter with vllm's RMSNorm haven't been supported yet.
857
- if ".pre_mixer_norm" in name:
858
- loaded_weight += 1.0
859
- elif ".post_mixer_norm" in name:
860
- loaded_weight += 1.0 / 5
861
- elif ".pre_mlp_norm" in name:
862
- loaded_weight += 1.0
863
- elif ".post_mlp_norm" in name:
864
- loaded_weight += 1.0 / (5**1.5)
865
- elif name == "norm.weight":
866
- loaded_weight += 1.0
867
-
868
- # Skip layers on other devices.
869
- if is_pp_missing_parameter(name, self):
870
- continue
871
-
872
- param = params_dict[name]
873
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
874
- weight_loader(param, loaded_weight)
875
- loaded_params.add(name)
876
- return loaded_params
877
-
878
-
879
-class Plamo2ForCausalLM(
880
- torch.nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid
881
-):
882
- packed_modules_mapping = {
883
- "qkv_proj": ["qkv_proj"],
884
- "gate_up_proj": ["gate_up_proj"],
885
- "in_proj": ["in_proj"],
886
- }
887
-
888
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
889
- super().__init__()
890
- config = vllm_config.model_config.hf_config
891
- scheduler_config = vllm_config.scheduler_config
892
-
893
- self.config = config
894
- self.vllm_config = vllm_config
895
- self.model_config = vllm_config.model_config
896
- self.scheduler_config = scheduler_config
897
-
898
- # ModelConfig.get_head_size assumes head_dim is set or calculated as
899
- # hidden_size // num_attention_heads. However, this is not always
900
- # the case for PLaMo2, as indicated by the FIXME comment.
901
- self.config.head_dim = self.config.hidden_size_per_head
902
-
903
- self.model = Plamo2Model(
904
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
905
- )
906
- self.vocab_size = self.config.vocab_size
907
- self.lm_head = ParallelLMHead(
908
- self.vocab_size,
909
- self.config.hidden_size,
910
- prefix=f"{prefix}.lm_head",
911
- )
912
- if self.config.tie_word_embeddings:
913
- self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
914
-
915
- self.logits_processor = LogitsProcessor(
916
- config.vocab_size, self.config.vocab_size
917
- )
918
- self.make_empty_intermediate_tensors = (
919
- self.model.make_empty_intermediate_tensors
920
- )
921
-
922
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
923
- return self.model.embed_input_ids(input_ids)
924
-
925
- def forward(
926
- self,
927
- input_ids: torch.Tensor | None,
928
- positions: torch.Tensor,
929
- intermediate_tensors: IntermediateTensors | None = None,
930
- inputs_embeds: torch.Tensor | None = None,
931
- **kwargs,
932
- ):
933
- hidden_states = self.model(
934
- input_ids, positions, intermediate_tensors, inputs_embeds
935
- )
936
- return hidden_states
937
-
938
- @classmethod
939
- def get_mamba_state_dtype_from_config(
940
- cls,
941
- vllm_config: "VllmConfig",
942
- ) -> tuple[torch.dtype, torch.dtype]:
943
- return MambaStateDtypeCalculator.mamba2_state_dtype(
944
- vllm_config.model_config.dtype,
945
- vllm_config.cache_config.mamba_cache_dtype,
946
- vllm_config.cache_config.mamba_ssm_cache_dtype,
947
- )
948
-
949
- @classmethod
950
- def get_mamba_state_shape_from_config(
951
- cls,
952
- vllm_config: "VllmConfig",
953
- ) -> tuple[tuple[int, int], tuple[int, int, int]]:
954
- """Calculate shapes for Mamba's convolutional and state caches.
955
- Args:
956
- vllm_config: vLLM config
957
- Returns:
958
- Tuple containing:
959
- - conv_state_shape: Shape for convolutional state cache
960
- - temporal_state_shape: Shape for state space model cache
961
- """
962
- parallel_config = vllm_config.parallel_config
963
- hf_config = vllm_config.model_config.hf_config
964
- intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head
965
-
966
- return MambaStateShapeCalculator.mamba2_state_shape(
967
- intermediate_size=intermediate_size,
968
- tp_world_size=parallel_config.tensor_parallel_size,
969
- n_groups=0,
970
- num_heads=hf_config.mamba_num_heads,
971
- head_dim=hf_config.hidden_size_per_head,
972
- state_size=hf_config.mamba_d_state,
973
- conv_kernel=hf_config.mamba_d_conv,
974
- )
975
-
976
- @classmethod
977
- def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
978
- return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
979
-
980
- def compute_logits(
981
- self,
982
- hidden_states: torch.Tensor,
983
- ) -> torch.Tensor | None:
984
- logits = self.logits_processor(self.lm_head, hidden_states)
985
- return logits
986
-
987
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
988
- loader = AutoWeightsLoader(
989
- self,
990
- skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
991
- )
992
- return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
-1
@@ -191,7 +191,6 @@ _TEXT_GENERATION_MODELS = {
191
"PhiForCausalLM": ("phi", "PhiForCausalLM"),
192
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
193
"PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
194
- "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
194
"Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
195
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
196
"Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
@@ -761,6 +760,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = {
760
"TeleChatForCausalLM": "0.25.0",
761
"PersimmonForCausalLM": "0.25.0",
762
"FuyuForCausalLM": "0.25.0",
763
+ "Plamo2ForCausalLM": "0.26.0",
764
}
765
766
_OOT_SUPPORTED_MODELS = {