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, 23:45 UTC
[Model] Support TP/PP/mamba2 kernel for PLaMo2 (#19674)
Signed-off-by: Shinichi Hemmi <shemmi@preferred.jp> Signed-off-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> Co-authored-by: Calvin Metzger <metzger@preferred.jp> Co-authored-by: Sixue Wang <cecilwang@preferred.jp>
Shinichi Hemmi committed
Jul 28, 2025 at 14:00 UTC
c7ffe93d9c136c14b05c96d647fbcc4d8268a5f9
4 files changed
+365
-213
docs/models/supported_models.md
+1
-1
@@ -389,7 +389,7 @@ th {
389
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
390
| `Phi4FlashForCausalLM` | Phi-4-mini-flash-reasoning | `microsoft/microsoft/Phi-4-mini-instruct`, etc. | | | |
391
| `PersimmonForCausalLM` | Persimmon | `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc. | | ✅︎ | ✅︎ |
392
-| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | | | |
392
+| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | | ✅︎ | |
393
| `QWenLMHeadModel` | Qwen | `Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |
394
| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ | ✅︎ |
395
| `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |
tests/distributed/test_pipeline_parallel.py
+1
@@ -175,6 +175,7 @@ TEXT_GENERATION_MODELS = {
175
"internlm/internlm2-chat-7b": PPTestSettings.fast(),
176
"inceptionai/jais-13b-chat": PPTestSettings.fast(),
177
"ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),
178
+ "pfnet/plamo-2-1b": PPTestSettings.fast(),
179
"meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(),
180
# Tests TransformersForCausalLM
181
"hmellor/Ilama-3.2-1B": PPTestSettings.fast(),
tests/quantization/test_experts_int8.py
+1
-1
@@ -9,7 +9,7 @@ import pytest
9
10
from tests.quantization.utils import is_quant_method_supported
11
12
-MODELS = ["ai21labs/Jamba-tiny-random"]
12
+MODELS = ["ai21labs/Jamba-tiny-random", "pfnet/plamo-2-1b"]
13
14
15
@pytest.mark.skipif(not is_quant_method_supported("experts_int8"),
vllm/model_executor/models/plamo2.py
+362
-211
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Inference-only PLaMo2 model."""
4
-import math
4
from collections.abc import Iterable
5
from typing import Optional
6
@@ -11,30 +10,40 @@ from transformers import PretrainedConfig, PreTrainedModel
10
11
from vllm.attention.backends.abstract import AttentionMetadata
12
from vllm.attention.layer import Attention
14
-from vllm.config import CacheConfig, VllmConfig
15
-from vllm.distributed import get_tensor_model_parallel_world_size
13
+from vllm.compilation.decorators import support_torch_compile
14
+from vllm.config import VllmConfig
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 get_forward_context
18
+from vllm.model_executor.layers.activation import SiluAndMul
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
21
MergedColumnParallelLinear,
22
QKVParallelLinear,
23
RowParallelLinear)
24
from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
+from vllm.model_executor.layers.mamba.mamba2_metadata import (
26
+ Mamba2Metadata, prepare_mamba2_metadata)
27
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
28
causal_conv1d_fn, causal_conv1d_update)
29
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
26
- selective_scan_fn, selective_state_update)
30
+ selective_state_update)
31
+from vllm.model_executor.layers.mamba.ops.ssd_combined import (
32
+ mamba_chunk_scan_combined)
33
from vllm.model_executor.layers.quantization import QuantizationConfig
34
from vllm.model_executor.layers.rotary_embedding import get_rope
35
+from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
36
from vllm.model_executor.layers.vocab_parallel_embedding import (
37
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
38
from vllm.model_executor.model_loader.weight_utils import (
39
composed_weight_loader, default_weight_loader, sharded_weight_loader)
40
from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,
34
- SupportsV0Only)
41
+ SupportsPP, SupportsV0Only)
42
from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
43
MambaCacheParams)
37
-from vllm.model_executor.models.utils import maybe_prefix
44
+from vllm.model_executor.models.utils import (
45
+ is_pp_missing_parameter, make_empty_intermediate_tensors_factory,
46
+ make_layers, maybe_prefix)
47
from vllm.model_executor.sampling_metadata import SamplingMetadata
48
from vllm.model_executor.utils import set_weight_attrs
49
from vllm.sequence import IntermediateTensors
@@ -77,17 +86,6 @@ class Plamo2PreTrainedModel(PreTrainedModel): # type: ignore
86
module.weight.data[module.padding_idx].zero_()
87
88
80
-def get_initial_dt_bias(num_heads: int) -> torch.Tensor:
81
- dt_min = 0.001
82
- dt_max = 0.1
83
- dt = torch.exp(
84
- torch.rand(num_heads) * (math.log(dt_max) - math.log(dt_min)) +
85
- math.log(dt_min))
86
- dt = torch.clamp(dt, 1e-4)
87
- inv_dt = dt + torch.log(-torch.expm1(-dt))
88
- return inv_dt
89
-
90
-
89
def is_mamba(config: Plamo2Config, i: int) -> bool:
90
assert config.mamba_step > 1
91
@@ -97,52 +95,36 @@ def is_mamba(config: Plamo2Config, i: int) -> bool:
95
return (i % config.mamba_step) != (config.mamba_step // 2)
96
97
100
-# TODO(Shinichi): Replace this with RMSNorm.
101
-def _rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor,
102
- eps: float) -> torch.Tensor:
103
- input_shape = hidden_states.shape
104
- hidden_states = hidden_states.reshape(input_shape[:-1] + weight.shape)
105
- input_dtype = hidden_states.dtype
106
- hidden_states = hidden_states.to(torch.float32)
107
- variance = hidden_states.pow(2).mean(-1, keepdim=True)
108
- hidden_states = hidden_states * torch.rsqrt(variance + eps)
109
- hidden_states = hidden_states.to(input_dtype)
110
- hidden_states = weight * hidden_states
111
- return hidden_states.reshape(input_shape)
112
-
113
-
114
-def _swiglu(h: torch.Tensor) -> torch.Tensor:
115
- h0, h1 = h.chunk(2, dim=-1)
116
- return torch.nn.functional.silu(h0) * h1
117
-
118
-
119
-# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
98
+# Adapted from:
99
+# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
100
+# transformers.models.mamba.modeling_mamba.MambaMixer
101
class Plamo2MambaMixer(nn.Module):
121
- # TODO(Shinichi): Rebase on Mamba2 implementation.
102
103
def __init__(self,
124
- config: Plamo2Config,
125
- cache_config: CacheConfig,
126
- quant_config: QuantizationConfig,
127
- max_model_len: int,
104
+ vllm_config: VllmConfig,
105
+ *,
106
prefix: str = "",
107
**kwargs) -> None:
108
super().__init__()
131
- self.config = config
132
- self.hidden_size = config.hidden_size
133
- self.ssm_state_size = config.mamba_d_state
134
- self.conv_kernel_size = config.mamba_d_conv
135
- self.intermediate_size = (config.mamba_num_heads *
136
- config.hidden_size_per_head)
137
- self.hidden_size_per_head = config.hidden_size_per_head
138
- self.num_heads = config.mamba_num_heads
109
+ self.config = vllm_config.model_config.hf_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
113
+ self.conv_kernel_size = self.config.mamba_d_conv
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()
117
+ self.intermediate_size_per_tp_worker = \
118
+ self.intermediate_size // self.tp_size
119
+ self.head_dim = self.config.hidden_size_per_head
120
+ self.num_heads = self.config.mamba_num_heads
121
self.time_step_rank = max(64, self.hidden_size // 16)
140
- self.use_conv_bias = False
141
- self.use_bias = False
122
self.conv1d = ColumnParallelLinear(
123
input_size=self.conv_kernel_size,
124
output_size=self.intermediate_size,
145
- bias=self.use_conv_bias,
125
+ bias=False,
126
+ prefix=f"{prefix}.conv1d",
127
+ return_bias=False,
128
)
129
# unsqueeze to fit conv1d weights shape into the linear weights shape.
130
# Can't do this in `weight_loader` since it already exists in
@@ -153,15 +135,19 @@ class Plamo2MambaMixer(nn.Module):
135
self.in_proj = MergedColumnParallelLinear(
136
self.hidden_size,
137
[self.intermediate_size] * 2,
156
- bias=self.use_bias,
138
+ bias=False,
139
+ quant_config=self.quant_config,
140
prefix=f"{prefix}.in_proj",
141
+ return_bias=False,
142
)
143
# selective projection used to make dt, B and C input dependent
144
self.bcdt_proj = RowParallelLinear(
145
self.intermediate_size,
146
self.time_step_rank + self.ssm_state_size * 2,
147
bias=False,
148
+ quant_config=self.quant_config,
149
prefix=f"{prefix}.bcdt_proj",
150
+ return_bias=False,
151
)
152
# time step projection (discretization) -
153
# In the forward we need to apply dt_proj without the bias,
@@ -170,154 +156,224 @@ class Plamo2MambaMixer(nn.Module):
156
self.time_step_rank,
157
self.num_heads,
158
bias=False,
159
+ quant_config=self.quant_config,
160
prefix=f"{prefix}.dt_proj",
161
+ return_bias=False,
162
)
175
- self.dt_bias = torch.nn.Parameter(get_initial_dt_bias(self.num_heads))
163
177
- tp_size = get_tensor_model_parallel_world_size()
164
self.A = nn.Parameter(
165
torch.empty(
180
- self.intermediate_size // tp_size,
181
- self.ssm_state_size,
166
+ divide(self.num_heads, self.tp_size),
167
dtype=torch.float32,
168
))
184
- self.D = nn.Parameter(torch.ones(self.intermediate_size // tp_size))
169
+ self.D = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size)))
170
+ self.dt_bias = nn.Parameter(
171
+ torch.ones(divide(self.num_heads, self.tp_size)))
172
173
set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)})
174
a_weight_loader = composed_weight_loader(
175
sharded_weight_loader(0), lambda x: -torch.exp(x.float()))
176
set_weight_attrs(self.A, {"weight_loader": a_weight_loader})
177
+ set_weight_attrs(self.dt_bias,
178
+ {"weight_loader": sharded_weight_loader(0)})
179
180
self.out_proj = RowParallelLinear(
181
self.intermediate_size,
182
self.hidden_size,
194
- bias=self.use_bias,
183
+ bias=False,
184
input_is_parallel=True,
185
+ quant_config=self.quant_config,
186
prefix=f"{prefix}.out_proj",
187
+ return_bias=False,
188
)
189
# The activation function is fixed to SiLU.
190
self.activation = "silu"
191
201
- self.dt_norm = RMSNorm(self.time_step_rank, eps=config.rms_norm_eps)
202
- self.B_norm = RMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
203
- self.C_norm = RMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)
192
+ self.dt_norm = RMSNorm(self.time_step_rank,
193
+ eps=self.config.rms_norm_eps)
194
+ self.B_norm = RMSNorm(self.ssm_state_size,
195
+ eps=self.config.rms_norm_eps)
196
+ self.C_norm = RMSNorm(self.ssm_state_size,
197
+ eps=self.config.rms_norm_eps)
198
+
199
+ def _project_ssm_parameters(self, hidden_states):
200
+ ssm_parameters = self.bcdt_proj(hidden_states)
201
+ B, C, time_step = torch.split(
202
+ ssm_parameters,
203
+ [self.ssm_state_size, self.ssm_state_size, self.time_step_rank],
204
+ dim=-1,
205
+ )
206
+
207
+ # vllm._custom_ops.rms_norm requires contiguous input tensors.
208
+ time_step = self.dt_norm(time_step.contiguous())
209
+ B = self.B_norm(B.contiguous())
210
+ C = self.C_norm(C.contiguous())
211
+ dt = self.dt_proj(time_step)
212
+ return B, C, dt
213
214
def forward(
215
self,
216
hidden_states: torch.Tensor,
217
mamba_cache_params: MambaCacheParams,
218
+ mamba2_metadata: Mamba2Metadata,
219
**kwargs,
220
) -> torch.Tensor:
221
222
+ # mamba2_metadata contains metadata necessary for the mamba2 triton
223
+ # kernels to operate in continuous batching and in chunked prefill
224
+ # modes; they are computed at top-level model forward since they
225
+ # stay the same and reused for all mamba layers in the same iteration
226
attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
227
228
+ num_prefills = attn_metadata.num_prefills # request count
229
+ num_decodes = attn_metadata.num_decode_tokens # token count (=request)
230
+ num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
231
+ has_prefill = num_prefills > 0
232
+ has_decode = num_decodes > 0
233
+
234
# 1. Gated MLP's linear projection
215
- projected_states = self.in_proj(hidden_states)[0]
216
- # Reshaping the projected states as in modeling_plamo.py.
217
- length = len(hidden_states)
218
- projected_states = projected_states.reshape(length, self.num_heads, -1)
219
- gate, hidden_states = torch.split(
220
- projected_states,
221
- [self.hidden_size_per_head, self.hidden_size_per_head],
222
- dim=-1)
223
- hidden_states = hidden_states.reshape(length, -1).transpose(0, 1)
224
- gate = gate.reshape(length, -1).transpose(0, 1)
235
+ projected_states = self.in_proj(hidden_states)
236
+ gate, hidden_states = projected_states.chunk(2, dim=-1)
237
238
# 2. Convolution sequence transformation
239
conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
240
self.conv1d.weight.size(2))
241
230
- if attn_metadata.query_start_loc is not None \
231
- and attn_metadata.context_lens_tensor is not None:
232
- # |---------- N-1 iteration --------|
233
- # |---------------- N iteration ---------------------|
234
- # |- tokenA -|......................|-- newTokens ---|
235
- # |---------- context_len ----------|
236
- # |-------------------- seq_len ---------------------|
237
- # |-- query_len ---|
238
- hidden_states = causal_conv1d_fn(
239
- hidden_states,
242
+ # Separate prefill and decode by splitting varlen input
243
+ # Split along token dimension
244
+ hidden_states_p, hidden_states_d = torch.split(
245
+ hidden_states,
246
+ [num_prefill_tokens, num_decodes],
247
+ dim=0,
248
+ )
249
+ gate_p, gate_d = torch.split(gate, [num_prefill_tokens, num_decodes],
250
+ dim=0)
251
+ # Split along batch dimension
252
+ state_indices_tensor_p, state_indices_tensor_d = torch.split(
253
+ mamba_cache_params.state_indices_tensor,
254
+ [num_prefills, num_decodes],
255
+ dim=0,
256
+ )
257
+ query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills + 1]
258
+ if has_prefill else None)
259
+
260
+ ssd_output_list = []
261
+
262
+ # Process prefill requests
263
+ if has_prefill:
264
+ # 2. Convolution sequence transformation
265
+ # - "cache_indices" updates the conv_state cache in positions
266
+ # pointed to by "mamba_cache_params.state_indices_tensor"
267
+ hidden_states_p = causal_conv1d_fn(
268
+ hidden_states_p.transpose(0, 1),
269
conv_weights,
270
self.conv1d.bias,
271
activation=self.activation,
272
conv_states=mamba_cache_params.conv_state,
244
- has_initial_state=attn_metadata.context_lens_tensor > 0,
245
- cache_indices=mamba_cache_params.state_indices_tensor,
246
- query_start_loc=attn_metadata.query_start_loc)
247
- else:
248
- hidden_states = causal_conv1d_update(
249
- hidden_states.transpose(0, 1),
273
+ has_initial_state=mamba2_metadata.has_initial_states,
274
+ cache_indices=state_indices_tensor_p,
275
+ query_start_loc=query_start_loc_p)
276
+ hidden_states_p = hidden_states_p.transpose(0, 1)
277
+ hidden_states_p = hidden_states_p[:num_prefill_tokens]
278
+ # In some instances, the following `bcdt_proj` op
279
+ # requires contiguous inputs
280
+ # (e.g. if the Marlin kernel is used).
281
+ hidden_states_p = hidden_states_p.contiguous()
282
+
283
+ B, C, dt = self._project_ssm_parameters(hidden_states_p)
284
+
285
+ # 3. State Space Model sequence transformation
286
+ initial_states = None
287
+ if (mamba2_metadata.has_initial_states is not None
288
+ and mamba2_metadata.prep_initial_states):
289
+ # making a copy of the states
290
+ initial_states = torch.where(
291
+ mamba2_metadata.has_initial_states[:, None, None, None],
292
+ mamba_cache_params.ssm_state[state_indices_tensor_p], 0)
293
+ scan_output, varlen_state = mamba_chunk_scan_combined(
294
+ hidden_states_p.view(1, num_prefill_tokens,
295
+ self.num_heads // self.tp_size,
296
+ self.head_dim),
297
+ dt.unsqueeze(0),
298
+ self.A,
299
+ B.view(1, num_prefill_tokens, 1, -1),
300
+ C.view(1, num_prefill_tokens, 1, -1),
301
+ chunk_size=mamba2_metadata.chunk_size,
302
+ D=self.D,
303
+ z=gate_p.view(1, num_prefill_tokens,
304
+ self.num_heads // self.tp_size, self.head_dim),
305
+ dt_bias=self.dt_bias,
306
+ seq_idx=mamba2_metadata.seq_idx,
307
+ chunk_indices=mamba2_metadata.chunk_indices,
308
+ chunk_offsets=mamba2_metadata.chunk_offsets,
309
+ cu_seqlens=attn_metadata.query_start_loc[:num_prefills + 1],
310
+ initial_states=initial_states,
311
+ return_varlen_states=True,
312
+ return_final_states=False,
313
+ dt_softplus=True,
314
+ dt_limit=(0.0, float("inf")),
315
+ )
316
+
317
+ # update ssm states
318
+ # - varlen state is a (batch, nheads, headdim, dstate) tensor
319
+ mamba_cache_params.ssm_state[state_indices_tensor_p] = varlen_state
320
+
321
+ # - reshape
322
+ ssd_output_list.append(scan_output.view(num_prefill_tokens, -1))
323
+
324
+ # Process decode requests
325
+ if has_decode:
326
+ # 2. Convolution sequence transformation
327
+ hidden_states_d = causal_conv1d_update(
328
+ hidden_states_d,
329
mamba_cache_params.conv_state,
330
conv_weights,
331
self.conv1d.bias,
332
self.activation,
254
- conv_state_indices=mamba_cache_params.state_indices_tensor)
255
- hidden_states = hidden_states.transpose(0, 1)
256
-
257
- # 3. State Space Model sequence transformation
258
- # 3.a. input varying initialization of time_step, B and C
259
- ssm_parameters = self.bcdt_proj(hidden_states.transpose(-2, -1))[0]
260
-
261
- # Splitting the ssm_parameters as in modeling_plamo.py.
262
- B, C, time_step = torch.split(
263
- ssm_parameters,
264
- [self.ssm_state_size, self.ssm_state_size, self.time_step_rank],
265
- dim=-1,
266
- )
267
- time_step = self.dt_norm(time_step.contiguous())
268
- B = self.B_norm(B.contiguous())
269
- C = self.C_norm(C.contiguous())
270
-
271
- discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)
272
- # 3.c perform the recurrence y ← SSM(A, B, C)(x)
273
- time_proj_bias = (self.dt_bias.float() if hasattr(
274
- self.dt_proj, "bias") else None)
275
-
276
- # Broadcasting as in modeling_plamo.py.
277
- discrete_time_step = discrete_time_step.transpose(
278
- 0, 1)[..., None].expand(-1, -1, self.hidden_size_per_head)
279
- discrete_time_step = discrete_time_step.reshape(
280
- -1, self.intermediate_size).transpose(0, 1)
281
- time_proj_bias = time_proj_bias[...,
282
- None].expand(-1,
283
- self.hidden_size_per_head)
284
- time_proj_bias = time_proj_bias.reshape(self.intermediate_size)
285
-
286
- if attn_metadata.query_start_loc is not None \
287
- and attn_metadata.context_lens_tensor is not None:
288
- scan_outputs = selective_scan_fn(
289
- hidden_states,
290
- mamba_cache_params.ssm_state,
291
- discrete_time_step,
292
- self.A,
293
- B.transpose(-2, -1),
294
- C.transpose(-2, -1),
295
- self.D.float(),
296
- gate,
297
- time_proj_bias,
298
- delta_softplus=True,
299
- cache_indices=mamba_cache_params.state_indices_tensor,
300
- has_initial_state=attn_metadata.context_lens_tensor > 0,
301
- query_start_loc=attn_metadata.query_start_loc)
302
- else:
303
- scan_outputs = selective_state_update(
333
+ conv_state_indices=state_indices_tensor_d)
334
+
335
+ B, C, dt = self._project_ssm_parameters(hidden_states_d)
336
+
337
+ # 3. State Space Model sequence transformation
338
+ A = self.A[:, None, ...][:, :,
339
+ None].expand(-1, self.head_dim,
340
+ self.config.mamba_d_state)
341
+ dt = dt[:, :, None].expand(-1, -1, self.head_dim)
342
+ dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
343
+ D = self.D[:, None, ...].expand(-1, self.head_dim)
344
+ B = B.unsqueeze(1)
345
+ C = C.unsqueeze(1)
346
+ hidden_states_d = hidden_states_d.view(
347
+ -1, self.num_heads // self.tp_size, self.head_dim)
348
+
349
+ # - the hidden is reshaped into (bs, num_heads, head_dim)
350
+ # - mamba_cache_params.ssm_state's slots will be selected
351
+ # using state_indices_tensor_d
352
+
353
+ hidden_states_d = selective_state_update(
354
mamba_cache_params.ssm_state,
305
- hidden_states.transpose(0, 1),
306
- discrete_time_step.transpose(0, 1),
307
- self.A,
355
+ hidden_states_d,
356
+ dt,
357
+ A,
358
B,
359
C,
310
- self.D,
311
- gate.transpose(0, 1),
312
- time_proj_bias,
360
+ D,
361
+ z=gate_d.reshape(num_decodes, -1, self.head_dim),
362
+ dt_bias=dt_bias,
363
dt_softplus=True,
314
- state_batch_indices=mamba_cache_params.state_indices_tensor)
315
- scan_outputs = scan_outputs.transpose(0, 1)
364
+ state_batch_indices=state_indices_tensor_d,
365
+ )
366
+ assert self.num_heads % self.tp_size == 0
367
+ ssd_output_list.append(
368
+ hidden_states_d.view(-1, (self.num_heads // self.tp_size) *
369
+ self.head_dim))
370
+
371
+ # Merge prefill and decode outputs before passing to MLP
372
+ hidden_states = torch.vstack(ssd_output_list)
373
374
# 4. Final linear projection
318
- contextualized_states = self.out_proj(scan_outputs.transpose(-2,
319
- -1))[0]
320
- return contextualized_states
375
+ out = self.out_proj(hidden_states)
376
+ return out
377
378
379
class DenseMLP(nn.Module):
@@ -332,33 +388,39 @@ class DenseMLP(nn.Module):
388
self.hidden_size = config.hidden_size
389
self.intermediate_size = config.intermediate_size
390
self.gate_up_proj = MergedColumnParallelLinear(
335
- self.hidden_size, [self.intermediate_size] * 2,
391
+ self.hidden_size,
392
+ [self.intermediate_size] * 2,
393
bias=False,
394
prefix=f"{prefix}.gate_up_proj",
338
- quant_config=quant_config)
395
+ quant_config=quant_config,
396
+ return_bias=False,
397
+ )
398
+ self.act = SiluAndMul()
399
self.down_proj = RowParallelLinear(self.intermediate_size,
400
self.hidden_size,
401
bias=False,
402
prefix=f"{prefix}.down_proj",
343
- quant_config=quant_config)
403
+ quant_config=quant_config,
404
+ return_bias=False)
405
406
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
346
- h = self.gate_up_proj(hidden_states)[0]
347
- h = _swiglu(h)
348
- output, _ = self.down_proj(h)
349
- return output # type: ignore
407
+ h = self.gate_up_proj(hidden_states)
408
+ h = self.act(h)
409
+ return self.down_proj(h)
410
411
412
+@support_torch_compile
413
class Plamo2AttentionMixer(nn.Module):
414
415
def __init__(self,
355
- config: Plamo2Config,
356
- cache_config: CacheConfig,
357
- quant_config: QuantizationConfig,
358
- max_model_len: int | None = None,
416
+ *,
417
+ vllm_config: VllmConfig,
418
prefix: str = "",
419
**kwargs) -> None:
420
super().__init__()
421
+ config = vllm_config.model_config.hf_config
422
+ cache_config = vllm_config.cache_config
423
+ quant_config = vllm_config.quant_config
424
self.hidden_size = config.hidden_size
425
tp_size = get_tensor_model_parallel_world_size()
426
self.total_num_heads = config.num_attention_heads
@@ -396,19 +458,35 @@ class Plamo2AttentionMixer(nn.Module):
458
"rope_theta") else 10000
459
self.rope_scaling = config.rope_scaling if hasattr(
460
config, "rope_scaling") else None
461
+ max_position = config.max_position_embeddings
462
+ if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
463
+ vllm_config.model_config.max_model_len, int):
464
+ max_position = min(max_position,
465
+ vllm_config.model_config.max_model_len)
466
400
- assert max_model_len is not None, "max_model_len must be provided"
467
self.rotary_emb = get_rope(
468
self.head_dim,
469
rotary_dim=self.head_dim,
404
- max_position=max_model_len,
470
+ max_position=max_position,
471
base=self.rope_theta,
472
rope_scaling=self.rope_scaling,
473
)
408
- self.q_weight = torch.nn.Parameter(
474
+ self.q_norm = RMSNorm(config.hidden_size_per_head,
475
+ eps=config.rms_norm_eps)
476
+ self.q_norm.weight = torch.nn.Parameter(
477
torch.ones((self.num_heads, config.hidden_size_per_head)))
410
- self.k_weight = torch.nn.Parameter(
478
+ set_weight_attrs(self.q_norm.weight,
479
+ {"weight_loader": sharded_weight_loader(0)})
480
+ self.k_norm = RMSNorm(config.hidden_size_per_head,
481
+ eps=config.rms_norm_eps)
482
+ self.k_norm.weight = torch.nn.Parameter(
483
torch.ones((self.num_kv_heads, config.hidden_size_per_head)))
484
+ # Tensor-parallelism shards the K norm weights to the tp ranks
485
+ # in a head-wise manner. This approach does not work if there is only
486
+ # a single KV head, as is the case for PLaMo 2-1B.
487
+ if self.total_num_kv_heads != 1:
488
+ set_weight_attrs(self.k_norm.weight,
489
+ {"weight_loader": sharded_weight_loader(0)})
490
491
self.attn = Attention(
492
self.num_heads,
@@ -423,13 +501,18 @@ class Plamo2AttentionMixer(nn.Module):
501
self,
502
positions: torch.Tensor,
503
hidden_states: torch.Tensor,
426
- residual: Optional[torch.Tensor],
504
**kwargs,
505
) -> torch.Tensor:
506
qkv, _ = self.qkv_proj(hidden_states)
507
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
431
- q = _rms_norm(q, self.q_weight, 1e-6)
432
- k = _rms_norm(k, self.k_weight, 1e-6)
508
+
509
+ q_shape = q.shape
510
+ q = q.reshape(q_shape[:-1] + self.q_norm.weight.shape)
511
+ q = self.q_norm.forward_native(q).reshape(q_shape)
512
+ k_shape = k.shape
513
+ k = k.reshape(k_shape[:-1] + self.k_norm.weight.shape)
514
+ k = self.k_norm.forward_native(k).reshape(k_shape)
515
+
516
q, k = self.rotary_emb(positions, q, k)
517
attn_output = self.attn(q, k, v)
518
output, _ = self.o_proj(attn_output)
@@ -441,27 +524,18 @@ class Plamo2DecoderLayer(nn.Module):
524
def __init__(self,
525
vllm_config: VllmConfig,
526
layer_idx: int,
444
- max_model_len: int | None = None,
527
prefix: str = "",
528
**kwargs) -> None:
529
super().__init__()
530
config = vllm_config.model_config.hf_config
449
- cache_config = vllm_config.cache_config
531
quant_config = vllm_config.quant_config
451
- max_model_len = vllm_config.scheduler_config.max_model_len
532
533
self.is_mamba = is_mamba(config, layer_idx)
534
if self.is_mamba:
455
- self.mixer = Plamo2MambaMixer(config=config,
456
- cache_config=cache_config,
457
- quant_config=quant_config,
458
- max_model_len=max_model_len,
535
+ self.mixer = Plamo2MambaMixer(vllm_config=vllm_config,
536
prefix=f"{prefix}.mixer")
537
else:
461
- self.mixer = Plamo2AttentionMixer(config=config,
462
- cache_config=cache_config,
463
- quant_config=quant_config,
464
- max_model_len=max_model_len,
538
+ self.mixer = Plamo2AttentionMixer(vllm_config=vllm_config,
539
prefix=f"{prefix}.mixer")
540
541
self.mlp = DenseMLP(config=config,
@@ -482,6 +556,7 @@ class Plamo2DecoderLayer(nn.Module):
556
hidden_states: torch.Tensor,
557
residual: Optional[torch.Tensor],
558
mamba_cache_params: MambaCacheParams,
559
+ mamba2_metadata: Mamba2Metadata,
560
**kwargs,
561
):
562
if residual is None:
@@ -491,10 +566,12 @@ class Plamo2DecoderLayer(nn.Module):
566
hidden_states, residual = self.pre_mixer_norm(
567
hidden_states, residual)
568
494
- hidden_states = self.mixer(positions=positions,
495
- hidden_states=hidden_states,
496
- residual=residual,
497
- mamba_cache_params=mamba_cache_params)
569
+ hidden_states = self.mixer(
570
+ positions=positions,
571
+ hidden_states=hidden_states,
572
+ mamba_cache_params=mamba_cache_params,
573
+ mamba2_metadata=mamba2_metadata,
574
+ )
575
hidden_states = self.post_mixer_norm(hidden_states)
576
# Fully Connected
577
hidden_states, residual = self.pre_mlp_norm(hidden_states, residual)
@@ -507,14 +584,18 @@ class Plamo2Decoder(torch.nn.Module):
584
585
def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
586
super().__init__()
510
- num_hidden_layers = vllm_config.model_config.hf_config.num_hidden_layers
587
+ config = vllm_config.model_config.hf_config
588
+ extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}
589
+
590
+ def get_layer(prefix: str):
591
+ layer_idx = int(prefix.rsplit(".", 1)[1])
592
+ return Plamo2DecoderLayer(vllm_config=vllm_config,
593
+ layer_idx=layer_idx,
594
+ prefix=prefix,
595
+ **extra_kwargs)
596
512
- self.layers = nn.ModuleList([
513
- Plamo2DecoderLayer(vllm_config=vllm_config,
514
- layer_idx=i,
515
- prefix=f"{prefix}.layers.{i}")
516
- for i in range(num_hidden_layers)
517
- ])
597
+ self.start_layer, self.end_layer, self.layers = make_layers(
598
+ config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers")
599
600
def forward(
601
self,
@@ -522,9 +603,10 @@ class Plamo2Decoder(torch.nn.Module):
603
hidden_states: torch.Tensor,
604
residual: Optional[torch.Tensor],
605
mamba_cache_params: MambaCacheParams,
606
+ mamba2_metadata: Mamba2Metadata,
607
) -> torch.Tensor:
608
mamba_cache_index = 0
527
- for layer in self.layers:
609
+ for layer in self.layers[self.start_layer:self.end_layer]:
610
layer_mamba_cache_params = None
611
if layer.is_mamba:
612
layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
@@ -535,7 +617,9 @@ class Plamo2Decoder(torch.nn.Module):
617
positions=positions,
618
hidden_states=hidden_states,
619
residual=residual,
538
- mamba_cache_params=layer_mamba_cache_params)
620
+ mamba_cache_params=layer_mamba_cache_params,
621
+ mamba2_metadata=mamba2_metadata,
622
+ )
623
return hidden_states, residual
624
625
@@ -557,10 +641,16 @@ class Plamo2Model(Plamo2PreTrainedModel):
641
org_num_embeddings=config.vocab_size,
642
prefix=f"{prefix}.embed_tokens",
643
)
644
+ self.make_empty_intermediate_tensors = (
645
+ make_empty_intermediate_tensors_factory(
646
+ ["hidden_states", "residual"], config.hidden_size))
647
self.layers = Plamo2Decoder(vllm_config, prefix=f"{prefix}.layers")
648
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
649
self.post_init()
650
651
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
652
+ return self.embed_tokens(input_ids)
653
+
654
def forward(
655
self,
656
input_ids: torch.Tensor,
@@ -569,21 +659,41 @@ class Plamo2Model(Plamo2PreTrainedModel):
659
intermediate_tensors: Optional[IntermediateTensors] = None,
660
inputs_embeds: Optional[torch.Tensor] = None,
661
) -> torch.Tensor:
572
- # TODO(Shinichi): Implement pipeline parallelism.
573
- hidden_states = self.embed_tokens(input_ids)
574
- residual = None
662
+ if get_pp_group().is_first_rank:
663
+ if inputs_embeds is not None:
664
+ hidden_states = inputs_embeds
665
+ else:
666
+ hidden_states = self.get_input_embeddings(input_ids)
667
+ residual = None
668
+ else:
669
+ assert intermediate_tensors is not None
670
+ hidden_states = intermediate_tensors["hidden_states"]
671
+ residual = intermediate_tensors["residual"]
672
+
673
+ attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
674
+ mamba2_metadata = prepare_mamba2_metadata(
675
+ chunk_size=self.config.mamba_chunk_size,
676
+ attn_metadata=attn_metadata,
677
+ )
678
679
hidden_states, residual = self.layers(
680
positions=positions,
681
hidden_states=hidden_states,
682
residual=residual,
580
- mamba_cache_params=mamba_cache_params)
683
+ mamba_cache_params=mamba_cache_params,
684
+ mamba2_metadata=mamba2_metadata,
685
+ )
686
+ if not get_pp_group().is_last_rank:
687
+ return IntermediateTensors({
688
+ "hidden_states": hidden_states,
689
+ "residual": residual
690
+ })
691
hidden_states, _ = self.norm(hidden_states, residual)
692
return hidden_states
693
694
585
-class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,
586
- SupportsV0Only):
695
+class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, SupportsPP,
696
+ IsHybrid, SupportsV0Only):
697
packed_modules_mapping = {
698
"qkv_proj": [
699
"q_proj",
@@ -629,10 +739,15 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,
739
740
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
741
self.config.vocab_size)
632
-
742
+ self.sampler = get_sampler()
743
+ self.make_empty_intermediate_tensors = (
744
+ self.model.make_empty_intermediate_tensors)
745
# Initialize weights and apply final processing
746
self.post_init()
747
748
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
749
+ return self.model.get_input_embeddings(input_ids)
750
+
751
def forward(self,
752
input_ids: torch.Tensor,
753
positions: torch.Tensor,
@@ -661,7 +776,7 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,
776
return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
777
778
def _get_mamba_cache_shape(
664
- self) -> tuple[tuple[int, int], tuple[int, int]]:
779
+ self) -> tuple[tuple[int, int], tuple[int, int, int]]:
780
world_size = get_tensor_model_parallel_world_size()
781
hidden_size = (self.config.mamba_num_heads *
782
self.config.hidden_size_per_head)
@@ -670,7 +785,8 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,
785
self.config.mamba_d_conv - 1,
786
)
787
temporal_state_shape = (
673
- hidden_size // world_size,
788
+ divide(self.config.mamba_num_heads, world_size),
789
+ self.config.hidden_size_per_head,
790
self.config.mamba_d_state,
791
)
792
return conv_state_shape, temporal_state_shape
@@ -684,6 +800,14 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,
800
sampling_metadata)
801
return logits
802
803
+ def sample(
804
+ self,
805
+ logits: Optional[torch.Tensor],
806
+ sampling_metadata: SamplingMetadata,
807
+ ) -> Optional[SamplerOutput]:
808
+ next_tokens = self.sampler(logits, sampling_metadata)
809
+ return next_tokens
810
+
811
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
812
params_dict = dict(self.named_parameters())
813
for name, loaded_weight in weights:
@@ -703,23 +827,46 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,
827
".B_norm_weight": ".B_norm.weight",
828
".C_norm_weight": ".C_norm.weight",
829
".dt_norm_weight": ".dt_norm.weight",
830
+ ".q_weight": ".q_norm.weight",
831
+ ".k_weight": ".k_norm.weight",
832
}
833
# Apply replacements based on the defined mappings
834
for old, new in replacements.items():
835
if old in name:
836
name = name.replace(old, new)
837
712
- # Broadcast the loaded weight to match the model's parameter shape.
713
- if ".A" in name:
714
- loaded_weight = loaded_weight[:, None, None].expand(
715
- -1, self.config.hidden_size_per_head,
716
- self.config.mamba_d_state)
838
+ # Reshape the in_proj weights to match the shape expected
839
+ # by MergedColumnParallelLinear.
840
+ # This works both for unquantized weights and
841
+ # for quantized weights.
842
+ # In the quantized case, the weights are already transposed.
843
+ # Also, in addition to the quantized weights,
844
+ # the zero points and scales have to be reshaped as well.
845
+ # Packing should not be affected by this.
846
+ if ".mixer.in_proj.weight" in name \
847
+ or "mixer.in_proj.qweight" in name \
848
+ or "mixer.in_proj.scales" in name \
849
+ or "mixer.in_proj.qzeros" in name:
850
+ if "mixer.in_proj.weight" in name:
851
+ loaded_weight = loaded_weight.transpose(0, 1)
852
+ # for weight:
853
+ # loaded_weight.shape[0] == self.config.hidden_size
854
+ # for qweight:
855
+ # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor # noqa
856
+ # for scales and qzeros:
857
+ # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size # noqa
858
loaded_weight = loaded_weight.reshape(
718
- -1, self.config.mamba_d_state)
719
- elif ".D" in name:
720
- loaded_weight = loaded_weight[:, None].expand(
721
- -1, self.config.hidden_size_per_head)
722
- loaded_weight = loaded_weight.reshape(-1)
859
+ loaded_weight.shape[0], self.config.mamba_num_heads, -1)
860
+ gate_weight, hidden_states_weight = loaded_weight.chunk(2,
861
+ dim=-1)
862
+ gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1)
863
+ hidden_states_weight = hidden_states_weight.reshape(
864
+ loaded_weight.shape[0], -1)
865
+ loaded_weight = torch.cat([gate_weight, hidden_states_weight],
866
+ dim=-1)
867
+ if "mixer.in_proj.weight" in name:
868
+ loaded_weight = loaded_weight.transpose(0, 1)
869
+
870
# Offset parameter with vllm's RMSNorm haven't been supported yet.
871
if ".pre_mixer_norm" in name:
872
loaded_weight += 1.0
@@ -732,6 +879,10 @@ class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,
879
elif "model.norm.weight" in name:
880
loaded_weight += 1.0
881
882
+ # Skip layers on other devices.
883
+ if is_pp_missing_parameter(name, self):
884
+ continue
885
+
886
param = params_dict[name]
887
weight_loader = getattr(param, "weight_loader",
888
default_weight_loader)