plamo2.py 37.1 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
35
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
36
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
37
38
39
40
    causal_conv1d_fn,
    causal_conv1d_update,
)
from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update
41
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
42
43
    mamba_chunk_scan_combined_varlen,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
44
45
46
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 (
47
48
49
    ParallelLMHead,
    VocabParallelEmbedding,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
50
from vllm.model_executor.model_loader.weight_utils import (
51
52
53
54
    composed_weight_loader,
    default_weight_loader,
    sharded_weight_loader,
)
55
56
57
58
59
60
from vllm.model_executor.models.interfaces import (
    HasInnerState,
    IsHybrid,
    SupportsLoRA,
    SupportsPP,
)
61
from vllm.model_executor.models.utils import (
62
    AutoWeightsLoader,
63
64
65
66
67
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
68
from vllm.model_executor.utils import set_weight_attrs
69
from vllm.platforms import current_platform
Shinichi Hemmi's avatar
Shinichi Hemmi committed
70
from vllm.sequence import IntermediateTensors
71
from vllm.utils.torch_utils import direct_register_custom_op
72
from vllm.v1.attention.backend import AttentionMetadata
73
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
Shinichi Hemmi's avatar
Shinichi Hemmi committed
74
75

# Only used for type hinting.
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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
100
101
102
103
104
105
106
107
    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)


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

116
    def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
117
        super().__init__()
118
        self.config = vllm_config.model_config.hf_config
119
120
        self.cache_config = vllm_config.cache_config
        self.model_config = vllm_config.model_config
121
        self.quant_config = vllm_config.quant_config
122
        self.is_lora_enabled = bool(vllm_config.lora_config)
123
124
125
        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
126
127
128
        self.intermediate_size = (
            self.config.mamba_num_heads * self.config.hidden_size_per_head
        )
129
130
131
        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
132
133
134
135
        self.time_step_rank = max(64, self.hidden_size // 16)
        self.conv1d = ColumnParallelLinear(
            input_size=self.conv_kernel_size,
            output_size=self.intermediate_size,
136
137
138
            bias=False,
            prefix=f"{prefix}.conv1d",
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
139
140
141
142
143
144
145
146
147
148
        )
        # 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,
149
150
            bias=False,
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
151
            prefix=f"{prefix}.in_proj",
152
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
153
154
155
156
157
158
        )
        # 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,
159
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
160
            prefix=f"{prefix}.bcdt_proj",
161
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
162
163
164
165
166
167
168
169
        )
        # 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,
170
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
171
            prefix=f"{prefix}.dt_proj",
172
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
173
174
175
176
        )

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

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

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

203
204
205
        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)
206

207
208
        self.chunk_size = self.config.mamba_chunk_size

209
210
211
212
213
214
215
        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"
216
217
218

        self.prefix = prefix

219
    def _project_ssm_parameters(self, hidden_states):
220
221
222
223
224
        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)
225
226
227
228
229
230
231
232
233
234
235
236
        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
237
238
239
240

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

250
    def forward_impl(
251
252
253
254
255
256
        self,
        hidden_states: torch.Tensor,
        output: torch.Tensor,
        **kwargs,
    ):
        forward_context = get_forward_context()
257
        # attn_metadata contains metadata necessary for the mamba2 triton
258
259
260
        # 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
261
        attn_metadata: AttentionMetadata = forward_context.attn_metadata
262
263
264
265
266

        if attn_metadata is not None:
            assert isinstance(attn_metadata, dict)
            attn_metadata = attn_metadata[self.prefix]
            assert isinstance(attn_metadata, Mamba2AttentionMetadata)
267
            self_kv_cache = self.kv_cache
268
269
270
            # conv_state = (..., dim, width-1) yet contiguous along 'dim'
            conv_state = self_kv_cache[0].transpose(-1, -2)
            ssm_state = self_kv_cache[1]
271
272
            state_indices_tensor_p = attn_metadata.state_indices_tensor_p
            state_indices_tensor_d = attn_metadata.state_indices_tensor_d
273
274
275
276
            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
277
            query_start_loc_p = attn_metadata.query_start_loc_p
278
279
            cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
            last_chunk_indices_p = attn_metadata.last_chunk_indices_p
280

Shinichi Hemmi's avatar
Shinichi Hemmi committed
281
        # 1. Gated MLP's linear projection
282
283
        projected_states = self.in_proj(hidden_states)
        gate, hidden_states = projected_states.chunk(2, dim=-1)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
284
285

        # 2. Convolution sequence transformation
286
287
288
        conv_weights = self.conv1d.weight.view(
            self.conv1d.weight.size(0), self.conv1d.weight.size(2)
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
289

290
291
        if attn_metadata is None:
            # profile run
292
293
294
            hidden_states = (
                hidden_states.transpose(0, 1).clone().transpose(0, 1)
            ).contiguous()
295
296
297
298
299
300
301
302
303
304
            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

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

        # Process prefill requests
        if has_prefill:
            # 2. Convolution sequence transformation
            # - "cache_indices" updates the conv_state cache in positions
335
            #   pointed to by "state_indices_tensor_p"
336
            x = hidden_states_p.transpose(0, 1)  # this is the form that causal-conv see
337
            hidden_states_p = causal_conv1d_fn(
338
                x,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
339
340
341
                conv_weights,
                self.conv1d.bias,
                activation=self.activation,
342
343
                conv_states=conv_state,
                has_initial_state=has_initial_states_p,
344
                cache_indices=state_indices_tensor_p,
345
                metadata=attn_metadata,
346
347
                query_start_loc=query_start_loc_p,
            )
348
349
350
351
352
353
354
355
356
357
358
            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
359
            if has_initial_states_p is not None and prep_initial_states:
360
                # making a copy of the states
361
362
                initial_states = torch.where(
                    has_initial_states_p[:, None, None, None],
363
364
365
                    ssm_state[state_indices_tensor_p],
                    0,
                )
366

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

            # update ssm states
            # - varlen state is a (batch, nheads, headdim, dstate) tensor
394
            ssm_state[state_indices_tensor_p] = varlen_state
395
396
397
398
399
400

        # Process decode requests
        if has_decode:
            # 2. Convolution sequence transformation
            hidden_states_d = causal_conv1d_update(
                hidden_states_d,
401
                conv_state,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
402
403
404
                conv_weights,
                self.conv1d.bias,
                self.activation,
405
406
                conv_state_indices=state_indices_tensor_d,
            )
407

408
409
410
411
412
413
414
            # 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()

415
416
417
            B, C, dt = self._project_ssm_parameters(hidden_states_d)

            # 3. State Space Model sequence transformation
418
419
420
            A = self.A[:, None, ...][:, :, None].expand(
                -1, self.head_dim, self.config.mamba_d_state
            )
421
422
423
424
425
426
            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(
427
428
                -1, self.num_heads // self.tp_size, self.head_dim
            )
429
430

            # - the hidden is reshaped into (bs, num_heads, head_dim)
431
            # - ssm_state's slots will be selected
432
            #   using state_indices_tensor_d
433
434

            # NOTE: final output is an in-place update of out tensor
435
            selective_state_update(
436
                ssm_state,
437
438
439
                hidden_states_d,
                dt,
                A,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
440
441
                B,
                C,
442
443
444
                D,
                z=gate_d.reshape(num_decodes, -1, self.head_dim),
                dt_bias=dt_bias,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
445
                dt_softplus=True,
446
                state_batch_indices=state_indices_tensor_d,
447
                out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),
448
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
449
450

        # 4. Final linear projection
451
452
453
454
455
456
457
458
459
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
        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]
485
    self.forward_impl(hidden_states=hidden_states, output=output)
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501


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
502
503
504
505
506


class DenseMLP(nn.Module):
    def __init__(
        self,
507
        config: "Plamo2Config",
508
        quant_config: QuantizationConfig | None = None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
509
510
511
512
513
514
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_up_proj = MergedColumnParallelLinear(
515
516
            self.hidden_size,
            [self.intermediate_size] * 2,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
517
518
            bias=False,
            prefix=f"{prefix}.gate_up_proj",
519
520
521
522
            quant_config=quant_config,
            return_bias=False,
        )
        self.act = SiluAndMul()
523
524
525
526
527
528
529
530
        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
531
532

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
533
534
535
        h = self.gate_up_proj(hidden_states)
        h = self.act(h)
        return self.down_proj(h)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
536
537
538


class Plamo2AttentionMixer(nn.Module):
539
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
540
        super().__init__()
541
542
543
        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
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
        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,
571
            prefix=f"{prefix}.qkv_proj",
Shinichi Hemmi's avatar
Shinichi Hemmi committed
572
        )
573
574
575
576
577
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
578
            prefix=f"{prefix}.o_proj",
579
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
580

581
582
        max_position = config.max_position_embeddings
        if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
583
584
585
            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
586
587
588

        self.rotary_emb = get_rope(
            self.head_dim,
589
            max_position=max_position,
590
            rope_parameters=config.rope_parameters,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
591
        )
592
        self.q_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps)
593
        self.q_norm.weight = torch.nn.Parameter(
594
595
596
597
598
599
            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)
600
        self.k_norm.weight = torch.nn.Parameter(
601
602
            torch.ones((self.num_kv_heads, config.hidden_size_per_head))
        )
603
604
605
606
        # 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:
607
608
609
            set_weight_attrs(
                self.k_norm.weight, {"weight_loader": sharded_weight_loader(0)}
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627

        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)
628
629
630
631
632
633
634
635

        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
636
637
638
639
640
641
642
        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):
643
644
645
    def __init__(
        self, vllm_config: VllmConfig, layer_idx: int, prefix: str = "", **kwargs
    ) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
646
647
648
649
650
651
        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:
652
653
654
            self.mixer = Plamo2MambaMixer(
                vllm_config=vllm_config, prefix=f"{prefix}.mixer"
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
655
        else:
656
657
658
659
660
661
662
663
664
665
666
            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
667
668
669
670
671

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
672
        residual: torch.Tensor | None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
673
674
675
676
677
678
        **kwargs,
    ):
        if residual is None:
            residual = hidden_states
            hidden_states = self.pre_mixer_norm(hidden_states)
        else:
679
            hidden_states, residual = self.pre_mixer_norm(hidden_states, residual)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
680

681
682
683
684
685
686
687
688
689
690
        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,
            }
691
692
        hidden_states = self.mixer(
            hidden_states=hidden_states,
693
            **mixer_kwargs,
694
        )
695
696
        if self.is_mamba:
            hidden_states = output
Shinichi Hemmi's avatar
Shinichi Hemmi committed
697
698
699
700
701
702
703
704
705
        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):
706
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
707
        super().__init__()
708
709
710
711
712
        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])
713
714
715
716
717
718
            return Plamo2DecoderLayer(
                vllm_config=vllm_config,
                layer_idx=layer_idx,
                prefix=prefix,
                **extra_kwargs,
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
719

720
        self.start_layer, self.end_layer, self.layers = make_layers(
721
722
            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
723
724
725
726
727

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
728
        residual: torch.Tensor | None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
729
    ) -> torch.Tensor:
730
        for layer in islice(self.layers, self.start_layer, self.end_layer):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
731
732
733
734
            hidden_states, residual = layer(
                positions=positions,
                hidden_states=hidden_states,
                residual=residual,
735
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
736
737
738
        return hidden_states, residual


739
740
@support_torch_compile
class Plamo2Model(torch.nn.Module):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
741
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
742
        super().__init__()
Shinichi Hemmi's avatar
Shinichi Hemmi committed
743
744
745
746
747
748
749
750
751
752
753

        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",
        )
754
755
756
757
        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
758
759
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

760
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
761
762
        return self.embed_tokens(input_ids)

Shinichi Hemmi's avatar
Shinichi Hemmi committed
763
764
    def forward(
        self,
765
        input_ids: torch.Tensor | None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
766
        positions: torch.Tensor,
767
768
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
769
    ) -> torch.Tensor:
770
771
772
773
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
774
                hidden_states = self.embed_input_ids(input_ids)
775
776
777
778
779
780
            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
781
782
783
784
        hidden_states, residual = self.layers(
            positions=positions,
            hidden_states=hidden_states,
            residual=residual,
785
786
        )
        if not get_pp_group().is_last_rank:
787
788
789
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
790
791
792
793
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


794
795
796
class Plamo2ForCausalLM(
    torch.nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid
):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
797
    packed_modules_mapping = {
798
799
800
        "qkv_proj": ["qkv_proj"],
        "gate_up_proj": ["gate_up_proj"],
        "in_proj": ["in_proj"],
Shinichi Hemmi's avatar
Shinichi Hemmi committed
801
802
803
    }

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
804
        super().__init__()
Shinichi Hemmi's avatar
Shinichi Hemmi committed
805
806
807
808
809
810
811
812
813
814
815
816
817
        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

818
819
820
        self.model = Plamo2Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
821
822
        self.vocab_size = self.config.vocab_size
        self.lm_head = ParallelLMHead(
823
            self.vocab_size,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
824
825
826
827
828
829
            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)

830
        self.logits_processor = LogitsProcessor(
831
            config.vocab_size, self.config.vocab_size
832
        )
833
        self.make_empty_intermediate_tensors = (
834
835
            self.model.make_empty_intermediate_tensors
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
836

837
838
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
839

840
841
    def forward(
        self,
842
        input_ids: torch.Tensor | None,
843
        positions: torch.Tensor,
844
845
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
846
847
848
849
850
        **kwargs,
    ):
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
851
852
        return hidden_states

853
854
855
856
857
858
859
860
861
    @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
862
        )
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878

    @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
879
        intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head
880
881
882
883
884
885
886
887
888

        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
889
890
        )

891
892
893
894
    @classmethod
    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
        return MambaStateCopyFuncCalculator.mamba2_state_copy_func()

Shinichi Hemmi's avatar
Shinichi Hemmi committed
895
896
897
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
898
    ) -> torch.Tensor | None:
899
        logits = self.logits_processor(self.lm_head, hidden_states)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
900
901
        return logits

902
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
Shinichi Hemmi's avatar
Shinichi Hemmi committed
903
904
905
906
907
908
909
        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
910
911
912
913
914
915
            # 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
916
917
918
919
920
921
922
923
924
925

            # 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",
926
927
                ".q_weight": ".q_norm.weight",
                ".k_weight": ".k_norm.weight",
Shinichi Hemmi's avatar
Shinichi Hemmi committed
928
929
930
931
932
933
            }
            # Apply replacements based on the defined mappings
            for old, new in replacements.items():
                if old in name:
                    name = name.replace(old, new)

934
935
936
937
938
939
940
941
            # 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.
942
943
944
945
946
947
            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
            ):
948
949
950
951
952
953
954
955
                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
956
                loaded_weight = loaded_weight.reshape(
957
958
959
                    loaded_weight.shape[0], self.config.mamba_num_heads, -1
                )
                gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1)
960
961
                gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1)
                hidden_states_weight = hidden_states_weight.reshape(
962
963
964
                    loaded_weight.shape[0], -1
                )
                loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1)
965
966
967
                if "mixer.in_proj.weight" in name:
                    loaded_weight = loaded_weight.transpose(0, 1)

Shinichi Hemmi's avatar
Shinichi Hemmi committed
968
969
970
971
972
973
974
975
976
977
978
979
            # 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

980
981
982
983
            # Skip layers on other devices.
            if is_pp_missing_parameter(name, self):
                continue

Shinichi Hemmi's avatar
Shinichi Hemmi committed
984
            param = params_dict[name]
985
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
986
            weight_loader(param, loaded_weight)