plamo2.py 37.6 KB
Newer Older
Shinichi Hemmi's avatar
Shinichi Hemmi committed
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
Shinichi Hemmi's avatar
Shinichi Hemmi committed
3
"""Inference-only PLaMo2 model."""
4

5
from collections.abc import Iterable
6
from itertools import islice
7
from typing import TYPE_CHECKING
Shinichi Hemmi's avatar
Shinichi Hemmi committed
8
9
10

import torch
from torch import nn
11
from transformers import PretrainedConfig
Shinichi Hemmi's avatar
Shinichi Hemmi committed
12

13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import VllmConfig, get_current_vllm_config
15
16
from vllm.distributed import divide, get_tensor_model_parallel_world_size
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
Shinichi Hemmi's avatar
Shinichi Hemmi committed
21
from vllm.model_executor.layers.layernorm import RMSNorm
22
23
24
25
26
27
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
28
from vllm.model_executor.layers.logits_processor import LogitsProcessor
29
30
from vllm.model_executor.layers.mamba.abstract import MambaBase
from vllm.model_executor.layers.mamba.mamba_utils import (
31
32
    MambaStateCopyFunc,
    MambaStateCopyFuncCalculator,
33
34
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
35
    is_conv_state_dim_first,
36
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
37
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
38
39
40
41
    causal_conv1d_fn,
    causal_conv1d_update,
)
from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update
42
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
43
44
    mamba_chunk_scan_combined_varlen,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
45
46
47
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
48
49
50
    ParallelLMHead,
    VocabParallelEmbedding,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
51
from vllm.model_executor.model_loader.weight_utils import (
52
53
54
55
    composed_weight_loader,
    default_weight_loader,
    sharded_weight_loader,
)
56
57
58
59
60
61
from vllm.model_executor.models.interfaces import (
    HasInnerState,
    IsHybrid,
    SupportsLoRA,
    SupportsPP,
)
62
from vllm.model_executor.models.utils import (
63
    AutoWeightsLoader,
64
65
66
67
68
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
69
from vllm.model_executor.utils import set_weight_attrs
70
from vllm.platforms import current_platform
Shinichi Hemmi's avatar
Shinichi Hemmi committed
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
Shinichi Hemmi's avatar
Shinichi Hemmi committed
75
76

# Only used for type hinting.
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
if TYPE_CHECKING:

    class Plamo2Config(PretrainedConfig):  # type: ignore
        model_type: str = "plamo2"

        hidden_size: int
        num_hidden_layers: int
        rms_norm_eps: float
        # Attention
        num_attention_heads: int
        hidden_size_per_head: int
        num_key_value_heads: int
        # Mamba
        mamba_d_state: int
        mamba_d_conv: int
        mamba_num_heads: int
        mamba_step: int
        # MLP
        intermediate_size: int
        # Tokenizer
        vocab_size: int


def is_mamba(config: "Plamo2Config", i: int) -> bool:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
101
102
103
104
105
106
107
108
    assert config.mamba_step > 1

    if config.num_hidden_layers <= (config.mamba_step // 2):
        # use attention in last layer
        return i != config.num_hidden_layers - 1
    return (i % config.mamba_step) != (config.mamba_step // 2)


109
110
111
# Adapted from:
# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
# transformers.models.mamba.modeling_mamba.MambaMixer
112
# --8<-- [start:plamo2_mamba_mixer]
113
114
@PluggableLayer.register("plamo2_mamba_mixer")
class Plamo2MambaMixer(MambaBase, PluggableLayer):
115
116
    # --8<-- [end:plamo2_mamba_mixer]

117
    def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
118
        super().__init__()
119
        self.config = vllm_config.model_config.hf_config
120
121
        self.cache_config = vllm_config.cache_config
        self.model_config = vllm_config.model_config
122
        self.quant_config = vllm_config.quant_config
123
        self.is_lora_enabled = bool(vllm_config.lora_config)
124
125
126
        self.hidden_size = self.config.hidden_size
        self.ssm_state_size = self.config.mamba_d_state
        self.conv_kernel_size = self.config.mamba_d_conv
127
128
129
        self.intermediate_size = (
            self.config.mamba_num_heads * self.config.hidden_size_per_head
        )
130
131
132
        self.tp_size = get_tensor_model_parallel_world_size()
        self.head_dim = self.config.hidden_size_per_head
        self.num_heads = self.config.mamba_num_heads
Shinichi Hemmi's avatar
Shinichi Hemmi committed
133
134
135
136
        self.time_step_rank = max(64, self.hidden_size // 16)
        self.conv1d = ColumnParallelLinear(
            input_size=self.conv_kernel_size,
            output_size=self.intermediate_size,
137
138
139
            bias=False,
            prefix=f"{prefix}.conv1d",
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
140
141
142
143
144
145
146
147
148
149
        )
        # unsqueeze to fit conv1d weights shape into the linear weights shape.
        # Can't do this in `weight_loader` since it already exists in
        # `ColumnParallelLinear` and `set_weight_attrs`
        # doesn't allow to override it
        self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

        self.in_proj = MergedColumnParallelLinear(
            self.hidden_size,
            [self.intermediate_size] * 2,
150
151
            bias=False,
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
152
            prefix=f"{prefix}.in_proj",
153
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
154
155
156
157
158
159
        )
        # selective projection used to make dt, B and C input dependent
        self.bcdt_proj = RowParallelLinear(
            self.intermediate_size,
            self.time_step_rank + self.ssm_state_size * 2,
            bias=False,
160
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
161
            prefix=f"{prefix}.bcdt_proj",
162
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
163
164
165
166
167
168
169
170
        )
        # time step projection (discretization) -
        # In the forward we need to apply dt_proj without the bias,
        # as the bias is added in the selective scan kernel.
        self.dt_proj = ColumnParallelLinear(
            self.time_step_rank,
            self.num_heads,
            bias=False,
171
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
172
            prefix=f"{prefix}.dt_proj",
173
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
174
175
176
177
        )

        self.A = nn.Parameter(
            torch.empty(
178
                divide(self.num_heads, self.tp_size),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
179
                dtype=torch.float32,
180
181
            )
        )
182
        self.D = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size)))
183
        self.dt_bias = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size)))
Shinichi Hemmi's avatar
Shinichi Hemmi committed
184
185
186

        set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)})
        a_weight_loader = composed_weight_loader(
187
188
            sharded_weight_loader(0), lambda x: -torch.exp(x.float())
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
189
        set_weight_attrs(self.A, {"weight_loader": a_weight_loader})
190
        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
Shinichi Hemmi's avatar
Shinichi Hemmi committed
191
192
193
194

        self.out_proj = RowParallelLinear(
            self.intermediate_size,
            self.hidden_size,
195
            bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
196
            input_is_parallel=True,
197
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
198
            prefix=f"{prefix}.out_proj",
199
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
200
201
202
203
        )
        # The activation function is fixed to SiLU.
        self.activation = "silu"

204
205
206
        self.dt_norm = RMSNorm(self.time_step_rank, eps=self.config.rms_norm_eps)
        self.B_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps)
        self.C_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps)
207

208
209
        self.chunk_size = self.config.mamba_chunk_size

210
211
212
213
214
215
216
        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self
        # The tuple is (conv_state, ssm_state)
        self.kv_cache = (torch.tensor([]), torch.tensor([]))
        assert self.chunk_size != -1, "chunk_size must be set for v1"
217
218
219

        self.prefix = prefix

220
    def _project_ssm_parameters(self, hidden_states):
221
222
223
224
225
        if self.is_lora_enabled:
            #  Lora kernel requires contiguous tensor.
            ssm_parameters = self.bcdt_proj(hidden_states.contiguous())
        else:
            ssm_parameters = self.bcdt_proj(hidden_states)
226
227
228
229
230
231
232
233
234
235
236
237
        B, C, time_step = torch.split(
            ssm_parameters,
            [self.ssm_state_size, self.ssm_state_size, self.time_step_rank],
            dim=-1,
        )

        # vllm._custom_ops.rms_norm requires contiguous input tensors.
        time_step = self.dt_norm(time_step.contiguous())
        B = self.B_norm(B.contiguous())
        C = self.C_norm(C.contiguous())
        dt = self.dt_proj(time_step)
        return B, C, dt
Shinichi Hemmi's avatar
Shinichi Hemmi committed
238
239
240
241

    def forward(
        self,
        hidden_states: torch.Tensor,
242
        output: torch.Tensor,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
243
        **kwargs,
244
    ):
245
246
247
248
249
        torch.ops.vllm.plamo2_mamba_mixer(
            hidden_states,
            output,
            self.prefix,
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
250

251
    def forward_impl(
252
253
254
255
256
257
        self,
        hidden_states: torch.Tensor,
        output: torch.Tensor,
        **kwargs,
    ):
        forward_context = get_forward_context()
258
        # attn_metadata contains metadata necessary for the mamba2 triton
259
260
261
        # kernels to operate in continuous batching and in chunked prefill
        # modes; they are computed at top-level model forward since they
        # stay the same and reused for all mamba layers in the same iteration
262
        attn_metadata: AttentionMetadata = forward_context.attn_metadata
263
264
265
266
267

        if attn_metadata is not None:
            assert isinstance(attn_metadata, dict)
            attn_metadata = attn_metadata[self.prefix]
            assert isinstance(attn_metadata, Mamba2AttentionMetadata)
268
            self_kv_cache = self.kv_cache
269
            # conv_state = (..., dim, width-1) yet contiguous along 'dim'
270
271
272
273
274
275
276
            # conv_state must be (..., dim, width-1) for the conv kernels.
            # DS layout stores it that way directly; SD layout needs a transpose.
            conv_state = (
                self_kv_cache[0]
                if is_conv_state_dim_first()
                else self_kv_cache[0].transpose(-1, -2)
            )
277
            ssm_state = self_kv_cache[1]
278
279
            state_indices_tensor_p = attn_metadata.state_indices_tensor_p
            state_indices_tensor_d = attn_metadata.state_indices_tensor_d
280
281
282
283
            has_initial_states_p = attn_metadata.has_initial_states_p
            prep_initial_states = attn_metadata.prep_initial_states
            chunk_size = attn_metadata.chunk_size
            seq_idx_p = attn_metadata.seq_idx_p
284
            query_start_loc_p = attn_metadata.query_start_loc_p
285
286
            cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
            last_chunk_indices_p = attn_metadata.last_chunk_indices_p
287

Shinichi Hemmi's avatar
Shinichi Hemmi committed
288
        # 1. Gated MLP's linear projection
289
290
        projected_states = self.in_proj(hidden_states)
        gate, hidden_states = projected_states.chunk(2, dim=-1)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
291
292

        # 2. Convolution sequence transformation
293
294
295
        conv_weights = self.conv1d.weight.view(
            self.conv1d.weight.size(0), self.conv1d.weight.size(2)
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
296

297
298
        if attn_metadata is None:
            # profile run
299
300
301
            hidden_states = (
                hidden_states.transpose(0, 1).clone().transpose(0, 1)
            ).contiguous()
302
303
304
305
306
307
308
309
310
311
            output[:] = self.out_proj(hidden_states)
            return

        num_prefills = attn_metadata.num_prefills  # request count
        num_decodes = attn_metadata.num_decode_tokens  # token count (=request)
        num_prefill_tokens = attn_metadata.num_prefill_tokens  # token count
        has_prefill = num_prefills > 0
        has_decode = num_decodes > 0
        num_actual_tokens = num_prefill_tokens + num_decodes

312
313
        # Separate prefill and decode by splitting varlen input
        # Split along token dimension
314
315
316
317
318
        hidden_states_d, hidden_states_p = torch.split(
            hidden_states[:num_actual_tokens],
            [num_decodes, num_prefill_tokens],
            dim=0,
        )
319
320
321
        gate_d, gate_p = torch.split(
            gate[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0
        )
322
323
324
325
326
        # Preallocate output tensor to avoid memcpy cost for merging prefill
        # and decode outputs
        preallocated_ssm_out = torch.empty(
            [
                num_prefill_tokens + num_decodes,
327
                (self.num_heads // self.tp_size) * self.head_dim,
328
329
330
331
            ],
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
332
333
334
335
336
        preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
            preallocated_ssm_out,
            [num_decodes, num_prefill_tokens],
            dim=0,
        )
337
338
339
340
341

        # Process prefill requests
        if has_prefill:
            # 2. Convolution sequence transformation
            # - "cache_indices" updates the conv_state cache in positions
342
            #   pointed to by "state_indices_tensor_p"
343
            x = hidden_states_p.transpose(0, 1)  # this is the form that causal-conv see
344
            hidden_states_p = causal_conv1d_fn(
345
                x,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
346
347
348
                conv_weights,
                self.conv1d.bias,
                activation=self.activation,
349
350
                conv_states=conv_state,
                has_initial_state=has_initial_states_p,
351
                cache_indices=state_indices_tensor_p,
352
                metadata=attn_metadata,
353
354
                query_start_loc=query_start_loc_p,
            )
355
356
357
358
359
360
361
362
363
364
365
            hidden_states_p = hidden_states_p.transpose(0, 1)
            hidden_states_p = hidden_states_p[:num_prefill_tokens]
            # In some instances, the following `bcdt_proj` op
            # requires contiguous inputs
            # (e.g. if the Marlin kernel is used).
            hidden_states_p = hidden_states_p.contiguous()

            B, C, dt = self._project_ssm_parameters(hidden_states_p)

            # 3. State Space Model sequence transformation
            initial_states = None
366
            if has_initial_states_p is not None and prep_initial_states:
367
                # making a copy of the states
368
369
                initial_states = torch.where(
                    has_initial_states_p[:, None, None, None],
370
371
372
                    ssm_state[state_indices_tensor_p],
                    0,
                )
373

374
            varlen_state = mamba_chunk_scan_combined_varlen(
375
376
377
                hidden_states_p.view(
                    num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim
                ),
378
                dt,
379
                self.A,
380
381
                B.view(num_prefill_tokens, 1, -1),
                C.view(num_prefill_tokens, 1, -1),
382
                chunk_size=chunk_size,
383
                D=self.D,
384
385
386
                z=gate_p.view(
                    num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim
                ),
387
                dt_bias=self.dt_bias,
388
389
                seq_idx=seq_idx_p,
                cu_seqlens=query_start_loc_p,
390
391
                cu_chunk_seqlens=cu_chunk_seqlen_p,
                last_chunk_indices=last_chunk_indices_p,
392
393
394
                initial_states=initial_states,
                dt_softplus=True,
                dt_limit=(0.0, float("inf")),
395
                out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim),
396
                state_dtype=ssm_state.dtype,
397
398
399
400
            )

            # update ssm states
            # - varlen state is a (batch, nheads, headdim, dstate) tensor
401
            ssm_state[state_indices_tensor_p] = varlen_state
402
403
404
405
406
407

        # Process decode requests
        if has_decode:
            # 2. Convolution sequence transformation
            hidden_states_d = causal_conv1d_update(
                hidden_states_d,
408
                conv_state,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
409
410
411
                conv_weights,
                self.conv1d.bias,
                self.activation,
412
413
                conv_state_indices=state_indices_tensor_d,
            )
414

415
416
417
418
419
420
421
            # ROCm: Ensure contiguous tensor for bcdt_proj linear layer.
            # causal_conv1d_update returns a non-contiguous view (stride 8192
            # instead of 4096 for shape [batch, 4096]), causing incorrect GEMM
            # results when batch > 1 on ROCm.
            if current_platform.is_rocm():
                hidden_states_d = hidden_states_d.contiguous()

422
423
424
            B, C, dt = self._project_ssm_parameters(hidden_states_d)

            # 3. State Space Model sequence transformation
425
426
427
            A = self.A[:, None, ...][:, :, None].expand(
                -1, self.head_dim, self.config.mamba_d_state
            )
428
429
430
431
432
433
            dt = dt[:, :, None].expand(-1, -1, self.head_dim)
            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
            D = self.D[:, None, ...].expand(-1, self.head_dim)
            B = B.unsqueeze(1)
            C = C.unsqueeze(1)
            hidden_states_d = hidden_states_d.view(
434
435
                -1, self.num_heads // self.tp_size, self.head_dim
            )
436
437

            # - the hidden is reshaped into (bs, num_heads, head_dim)
438
            # - ssm_state's slots will be selected
439
            #   using state_indices_tensor_d
440
441

            # NOTE: final output is an in-place update of out tensor
442
            selective_state_update(
443
                ssm_state,
444
445
446
                hidden_states_d,
                dt,
                A,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
447
448
                B,
                C,
449
450
451
                D,
                z=gate_d.reshape(num_decodes, -1, self.head_dim),
                dt_bias=dt_bias,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
452
                dt_softplus=True,
453
                state_batch_indices=state_indices_tensor_d,
454
                out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),
455
456
                enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding,
                cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds,
457
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
458
459

        # 4. Final linear projection
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
        output[:num_actual_tokens] = self.out_proj(preallocated_ssm_out)

    def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
        assert self.model_config is not None
        assert self.cache_config is not None
        return MambaStateDtypeCalculator.mamba2_state_dtype(
            self.model_config.dtype,
            self.cache_config.mamba_cache_dtype,
            self.cache_config.mamba_ssm_cache_dtype,
        )

    def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
        return MambaStateShapeCalculator.mamba2_state_shape(
            intermediate_size=self.intermediate_size,
            tp_world_size=get_tensor_model_parallel_world_size(),
            n_groups=0,
            num_heads=self.num_heads,
            head_dim=self.head_dim,
            state_size=self.ssm_state_size,
            conv_kernel=self.conv_kernel_size,
        )

    @property
    def mamba_type(self) -> str:
        return "mamba2"


def plamo2_mamba_mixer(
    hidden_states: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
) -> None:
    forward_context: ForwardContext = get_forward_context()
    self = forward_context.no_compile_layers[layer_name]
494
    self.forward_impl(hidden_states=hidden_states, output=output)
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510


def plamo2_mamba_mixer_fake(
    hidden_states: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
) -> None:
    return


direct_register_custom_op(
    op_name="plamo2_mamba_mixer",
    op_func=plamo2_mamba_mixer,
    mutates_args=["output"],
    fake_impl=plamo2_mamba_mixer_fake,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
511
512
513
514
515


class DenseMLP(nn.Module):
    def __init__(
        self,
516
        config: "Plamo2Config",
517
        quant_config: QuantizationConfig | None = None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
518
519
520
521
522
523
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_up_proj = MergedColumnParallelLinear(
524
525
            self.hidden_size,
            [self.intermediate_size] * 2,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
526
527
            bias=False,
            prefix=f"{prefix}.gate_up_proj",
528
529
530
531
            quant_config=quant_config,
            return_bias=False,
        )
        self.act = SiluAndMul()
532
533
534
535
536
537
538
539
        self.down_proj = RowParallelLinear(
            self.intermediate_size,
            self.hidden_size,
            bias=False,
            prefix=f"{prefix}.down_proj",
            quant_config=quant_config,
            return_bias=False,
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
540
541

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
542
543
544
        h = self.gate_up_proj(hidden_states)
        h = self.act(h)
        return self.down_proj(h)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
545
546
547


class Plamo2AttentionMixer(nn.Module):
548
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
549
        super().__init__()
550
551
552
        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
Shinichi Hemmi's avatar
Shinichi Hemmi committed
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = config.num_key_value_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = config.hidden_size_per_head
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5

        self.qkv_proj = QKVParallelLinear(
            config.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
            quant_config=quant_config,
580
            prefix=f"{prefix}.qkv_proj",
Shinichi Hemmi's avatar
Shinichi Hemmi committed
581
        )
582
583
584
585
586
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
587
            prefix=f"{prefix}.o_proj",
588
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
589

590
591
        max_position = config.max_position_embeddings
        if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
592
593
594
            vllm_config.model_config.max_model_len, int
        ):
            max_position = min(max_position, vllm_config.model_config.max_model_len)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
595
596
597

        self.rotary_emb = get_rope(
            self.head_dim,
598
            max_position=max_position,
599
            rope_parameters=config.rope_parameters,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
600
        )
601
        self.q_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps)
602
        self.q_norm.weight = torch.nn.Parameter(
603
604
605
606
607
608
            torch.ones((self.num_heads, config.hidden_size_per_head))
        )
        set_weight_attrs(
            self.q_norm.weight, {"weight_loader": sharded_weight_loader(0)}
        )
        self.k_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps)
609
        self.k_norm.weight = torch.nn.Parameter(
610
611
            torch.ones((self.num_kv_heads, config.hidden_size_per_head))
        )
612
613
614
615
        # Tensor-parallelism shards the K norm weights to the tp ranks
        # in a head-wise manner. This approach does not work if there is only
        # a single KV head, as is the case for PLaMo 2-1B.
        if self.total_num_kv_heads != 1:
616
617
618
            set_weight_attrs(
                self.k_norm.weight, {"weight_loader": sharded_weight_loader(0)}
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636

        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            prefix=f"{prefix}.attn",
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
637
638
639
640
641
642
643
644

        q_shape = q.shape
        q = q.reshape(q_shape[:-1] + self.q_norm.weight.shape)
        q = self.q_norm.forward_native(q).reshape(q_shape)
        k_shape = k.shape
        k = k.reshape(k_shape[:-1] + self.k_norm.weight.shape)
        k = self.k_norm.forward_native(k).reshape(k_shape)

Shinichi Hemmi's avatar
Shinichi Hemmi committed
645
646
647
648
649
650
651
        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output


class Plamo2DecoderLayer(nn.Module):
652
653
654
    def __init__(
        self, vllm_config: VllmConfig, layer_idx: int, prefix: str = "", **kwargs
    ) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
655
656
657
658
659
660
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config

        self.is_mamba = is_mamba(config, layer_idx)
        if self.is_mamba:
661
662
663
            self.mixer = Plamo2MambaMixer(
                vllm_config=vllm_config, prefix=f"{prefix}.mixer"
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
664
        else:
665
666
667
668
669
670
671
672
673
674
675
            self.mixer = Plamo2AttentionMixer(
                vllm_config=vllm_config, prefix=f"{prefix}.mixer"
            )

        self.mlp = DenseMLP(
            config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
        )
        self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
676
677
678
679
680

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
681
        residual: torch.Tensor | None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
682
683
684
685
686
687
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.pre_mixer_norm(hidden_states)
        else:
688
            hidden_states, residual = self.pre_mixer_norm(hidden_states, residual)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
689

690
691
692
693
694
695
696
697
698
699
        if self.is_mamba:
            # Plamo2MambaMixer writes output to this tensor
            output = torch.empty_like(hidden_states)
            mixer_kwargs = {
                "output": output,
            }
        else:
            mixer_kwargs = {
                "positions": positions,
            }
700
701
        hidden_states = self.mixer(
            hidden_states=hidden_states,
702
            **mixer_kwargs,
703
        )
704
705
        if self.is_mamba:
            hidden_states = output
Shinichi Hemmi's avatar
Shinichi Hemmi committed
706
707
708
709
710
711
712
713
714
        hidden_states = self.post_mixer_norm(hidden_states)
        # Fully Connected
        hidden_states, residual = self.pre_mlp_norm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.post_mlp_norm(hidden_states)
        return hidden_states, residual


class Plamo2Decoder(torch.nn.Module):
715
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
716
        super().__init__()
717
718
719
720
721
        config = vllm_config.model_config.hf_config
        extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)}

        def get_layer(prefix: str):
            layer_idx = int(prefix.rsplit(".", 1)[1])
722
723
724
725
726
727
            return Plamo2DecoderLayer(
                vllm_config=vllm_config,
                layer_idx=layer_idx,
                prefix=prefix,
                **extra_kwargs,
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
728

729
        self.start_layer, self.end_layer, self.layers = make_layers(
730
731
            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
732
733
734
735
736

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
737
        residual: torch.Tensor | None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
738
    ) -> torch.Tensor:
739
        for layer in islice(self.layers, self.start_layer, self.end_layer):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
740
741
742
743
            hidden_states, residual = layer(
                positions=positions,
                hidden_states=hidden_states,
                residual=residual,
744
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
745
746
747
        return hidden_states, residual


748
749
@support_torch_compile
class Plamo2Model(torch.nn.Module):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
750
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
751
        super().__init__()
Shinichi Hemmi's avatar
Shinichi Hemmi committed
752
753
754
755
756
757
758
759
760
761
762

        config = vllm_config.model_config.hf_config

        self.config = config
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            prefix=f"{prefix}.embed_tokens",
        )
763
764
765
766
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
        self.layers = Plamo2Decoder(vllm_config=vllm_config, prefix=f"{prefix}.layers")
Shinichi Hemmi's avatar
Shinichi Hemmi committed
767
768
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

769
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
770
771
        return self.embed_tokens(input_ids)

Shinichi Hemmi's avatar
Shinichi Hemmi committed
772
773
    def forward(
        self,
774
        input_ids: torch.Tensor | None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
775
        positions: torch.Tensor,
776
777
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
778
    ) -> torch.Tensor:
779
780
781
782
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
783
                hidden_states = self.embed_input_ids(input_ids)
784
785
786
787
788
789
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

Shinichi Hemmi's avatar
Shinichi Hemmi committed
790
791
792
793
        hidden_states, residual = self.layers(
            positions=positions,
            hidden_states=hidden_states,
            residual=residual,
794
795
        )
        if not get_pp_group().is_last_rank:
796
797
798
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
799
800
801
802
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


803
804
805
class Plamo2ForCausalLM(
    torch.nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid
):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
806
    packed_modules_mapping = {
807
808
809
        "qkv_proj": ["qkv_proj"],
        "gate_up_proj": ["gate_up_proj"],
        "in_proj": ["in_proj"],
Shinichi Hemmi's avatar
Shinichi Hemmi committed
810
811
812
    }

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
813
        super().__init__()
Shinichi Hemmi's avatar
Shinichi Hemmi committed
814
815
816
817
818
819
820
821
822
823
824
825
826
        config = vllm_config.model_config.hf_config
        scheduler_config = vllm_config.scheduler_config

        self.config = config
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
        self.scheduler_config = scheduler_config

        # ModelConfig.get_head_size assumes head_dim is set or calculated as
        # hidden_size // num_attention_heads. However, this is not always
        # the case for PLaMo2, as indicated by the FIXME comment.
        self.config.head_dim = self.config.hidden_size_per_head

827
828
829
        self.model = Plamo2Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
830
831
        self.vocab_size = self.config.vocab_size
        self.lm_head = ParallelLMHead(
832
            self.vocab_size,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
833
834
835
836
837
838
            self.config.hidden_size,
            prefix=f"{prefix}.lm_head",
        )
        if self.config.tie_word_embeddings:
            self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

839
        self.logits_processor = LogitsProcessor(
840
            config.vocab_size, self.config.vocab_size
841
        )
842
        self.make_empty_intermediate_tensors = (
843
844
            self.model.make_empty_intermediate_tensors
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
845

846
847
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
848

849
850
    def forward(
        self,
851
        input_ids: torch.Tensor | None,
852
        positions: torch.Tensor,
853
854
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
855
856
857
858
859
        **kwargs,
    ):
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
860
861
        return hidden_states

862
863
864
865
866
867
868
869
870
    @classmethod
    def get_mamba_state_dtype_from_config(
        cls,
        vllm_config: "VllmConfig",
    ) -> tuple[torch.dtype, torch.dtype]:
        return MambaStateDtypeCalculator.mamba2_state_dtype(
            vllm_config.model_config.dtype,
            vllm_config.cache_config.mamba_cache_dtype,
            vllm_config.cache_config.mamba_ssm_cache_dtype,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
871
        )
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887

    @classmethod
    def get_mamba_state_shape_from_config(
        cls,
        vllm_config: "VllmConfig",
    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
        """Calculate shapes for Mamba's convolutional and state caches.
        Args:
            vllm_config: vLLM config
        Returns:
            Tuple containing:
            - conv_state_shape: Shape for convolutional state cache
            - temporal_state_shape: Shape for state space model cache
        """
        parallel_config = vllm_config.parallel_config
        hf_config = vllm_config.model_config.hf_config
888
        intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head
889
890
891
892
893
894
895
896
897

        return MambaStateShapeCalculator.mamba2_state_shape(
            intermediate_size=intermediate_size,
            tp_world_size=parallel_config.tensor_parallel_size,
            n_groups=0,
            num_heads=hf_config.mamba_num_heads,
            head_dim=hf_config.hidden_size_per_head,
            state_size=hf_config.mamba_d_state,
            conv_kernel=hf_config.mamba_d_conv,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
898
899
        )

900
901
902
903
    @classmethod
    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
        return MambaStateCopyFuncCalculator.mamba2_state_copy_func()

Shinichi Hemmi's avatar
Shinichi Hemmi committed
904
905
906
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
907
    ) -> torch.Tensor | None:
908
        logits = self.logits_processor(self.lm_head, hidden_states)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
909
910
        return logits

911
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
912
913
914
915
916
917
918
        params_dict = dict(self.named_parameters())
        for name, loaded_weight in weights:
            # Both tie_word_embeddings=True and lm_head.weight in the safetensor
            # at the same time causes dict key access error.
            if name == "lm_head.weight" and self.config.tie_word_embeddings:
                assert "lm_head.weight" not in params_dict
                continue
919
920
921
922
923
924
            # Same workaround as AutoWeightsLoader for GPTQModel
            if any(
                substr in name
                for substr in AutoWeightsLoader.ROTARY_EMBEDS_UNUSED_WEIGHTS
            ):
                continue
Shinichi Hemmi's avatar
Shinichi Hemmi committed
925
926
927
928
929
930
931
932
933
934

            # Update the weight names to be compatible with the vllm version
            # of the model.
            # Do not change the order of the replacements.
            replacements = {
                # Rename incompatible weight names.
                ".A_log": ".A",
                ".B_norm_weight": ".B_norm.weight",
                ".C_norm_weight": ".C_norm.weight",
                ".dt_norm_weight": ".dt_norm.weight",
935
936
                ".q_weight": ".q_norm.weight",
                ".k_weight": ".k_norm.weight",
Shinichi Hemmi's avatar
Shinichi Hemmi committed
937
938
939
940
941
942
            }
            # Apply replacements based on the defined mappings
            for old, new in replacements.items():
                if old in name:
                    name = name.replace(old, new)

943
944
945
946
947
948
949
950
            # Reshape the in_proj weights to match the shape expected
            # by MergedColumnParallelLinear.
            # This works both for unquantized weights and
            # for quantized weights.
            # In the quantized case, the weights are already transposed.
            # Also, in addition to the quantized weights,
            # the zero points and scales have to be reshaped as well.
            # Packing should not be affected by this.
951
952
953
954
955
956
            if (
                ".mixer.in_proj.weight" in name
                or "mixer.in_proj.qweight" in name
                or "mixer.in_proj.scales" in name
                or "mixer.in_proj.qzeros" in name
            ):
957
958
959
960
961
962
963
964
                if "mixer.in_proj.weight" in name:
                    loaded_weight = loaded_weight.transpose(0, 1)
                # for weight:
                # loaded_weight.shape[0] == self.config.hidden_size
                # for qweight:
                # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor  # noqa
                # for scales and qzeros:
                # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size  # noqa
Shinichi Hemmi's avatar
Shinichi Hemmi committed
965
                loaded_weight = loaded_weight.reshape(
966
967
968
                    loaded_weight.shape[0], self.config.mamba_num_heads, -1
                )
                gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1)
969
970
                gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1)
                hidden_states_weight = hidden_states_weight.reshape(
971
972
973
                    loaded_weight.shape[0], -1
                )
                loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1)
974
975
976
                if "mixer.in_proj.weight" in name:
                    loaded_weight = loaded_weight.transpose(0, 1)

Shinichi Hemmi's avatar
Shinichi Hemmi committed
977
978
979
980
981
982
983
984
985
986
987
988
            # Offset parameter with vllm's RMSNorm haven't been supported yet.
            if ".pre_mixer_norm" in name:
                loaded_weight += 1.0
            elif ".post_mixer_norm" in name:
                loaded_weight += 1.0 / 5
            elif ".pre_mlp_norm" in name:
                loaded_weight += 1.0
            elif ".post_mlp_norm" in name:
                loaded_weight += 1.0 / (5**1.5)
            elif "model.norm.weight" in name:
                loaded_weight += 1.0

989
990
991
992
            # Skip layers on other devices.
            if is_pp_missing_parameter(name, self):
                continue

Shinichi Hemmi's avatar
Shinichi Hemmi committed
993
            param = params_dict[name]
994
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
995
            weight_loader(param, loaded_weight)