plamo2.py 37.2 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
Shinichi Hemmi's avatar
Shinichi Hemmi committed
7
8
9

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

12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import VllmConfig, get_current_vllm_config
14
15
from vllm.distributed import divide, get_tensor_model_parallel_world_size
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.forward_context import ForwardContext, get_forward_context
17
from vllm.model_executor.custom_op import PluggableLayer
18
from vllm.model_executor.layers.activation import SiluAndMul
19
from vllm.model_executor.layers.attention import Attention
Shinichi Hemmi's avatar
Shinichi Hemmi committed
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
22
23
24
25
26
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
27
from vllm.model_executor.layers.logits_processor import LogitsProcessor
28
29
from vllm.model_executor.layers.mamba.abstract import MambaBase
from vllm.model_executor.layers.mamba.mamba_utils import (
30
31
    MambaStateCopyFunc,
    MambaStateCopyFuncCalculator,
32
33
34
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
35
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
36
37
38
39
    causal_conv1d_fn,
    causal_conv1d_update,
)
from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update
40
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
41
42
    mamba_chunk_scan_combined_varlen,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
43
44
45
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 (
46
47
48
    ParallelLMHead,
    VocabParallelEmbedding,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
49
from vllm.model_executor.model_loader.weight_utils import (
50
51
52
53
    composed_weight_loader,
    default_weight_loader,
    sharded_weight_loader,
)
54
55
56
57
58
59
from vllm.model_executor.models.interfaces import (
    HasInnerState,
    IsHybrid,
    SupportsLoRA,
    SupportsPP,
)
60
from vllm.model_executor.models.utils import (
61
    AutoWeightsLoader,
62
63
64
65
66
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
67
from vllm.model_executor.utils import set_weight_attrs
68
from vllm.platforms import current_platform
Shinichi Hemmi's avatar
Shinichi Hemmi committed
69
from vllm.sequence import IntermediateTensors
70
from vllm.utils.torch_utils import direct_register_custom_op
71
from vllm.v1.attention.backend import AttentionMetadata
72
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
Shinichi Hemmi's avatar
Shinichi Hemmi committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105


# Only used for type hinting.
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:
    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)


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

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

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

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

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

201
202
203
        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)
204

205
206
        self.chunk_size = self.config.mamba_chunk_size

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

        self.prefix = prefix

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

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

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

        if attn_metadata is not None:
            assert isinstance(attn_metadata, dict)
            attn_metadata = attn_metadata[self.prefix]
            assert isinstance(attn_metadata, Mamba2AttentionMetadata)
            self_kv_cache = self.kv_cache[forward_context.virtual_engine]
            # conv_state = (..., dim, width-1) yet contiguous along 'dim'
            conv_state = self_kv_cache[0].transpose(-1, -2)
            ssm_state = self_kv_cache[1]
            state_indices_tensor = attn_metadata.state_indices_tensor
            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
274
            query_start_loc_p = attn_metadata.query_start_loc_p
275
276
            cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
            last_chunk_indices_p = attn_metadata.last_chunk_indices_p
277

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

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

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

302
303
        # Separate prefill and decode by splitting varlen input
        # Split along token dimension
304
305
306
307
308
        hidden_states_d, hidden_states_p = torch.split(
            hidden_states[:num_actual_tokens],
            [num_decodes, num_prefill_tokens],
            dim=0,
        )
309
310
311
        gate_d, gate_p = torch.split(
            gate[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0
        )
312
313
314
315
316
317
        # Split along batch dimension
        state_indices_tensor_d, state_indices_tensor_p = torch.split(
            state_indices_tensor,
            [num_decodes, num_prefills],
            dim=0,
        )
318

319
320
321
322
323
        # Preallocate output tensor to avoid memcpy cost for merging prefill
        # and decode outputs
        preallocated_ssm_out = torch.empty(
            [
                num_prefill_tokens + num_decodes,
324
                (self.num_heads // self.tp_size) * self.head_dim,
325
326
327
328
            ],
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
329
330
331
332
333
        preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
            preallocated_ssm_out,
            [num_decodes, num_prefill_tokens],
            dim=0,
        )
334
335
336
337
338

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

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

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

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

412
413
414
415
416
417
418
            # 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()

419
420
421
            B, C, dt = self._project_ssm_parameters(hidden_states_d)

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

            # - the hidden is reshaped into (bs, num_heads, head_dim)
435
            # - ssm_state's slots will be selected
436
            #   using state_indices_tensor_d
437
438

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

        # 4. Final linear projection
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
485
486
487
488
        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]
489
    self.forward_impl(hidden_states=hidden_states, output=output)
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505


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
506
507
508
509
510
511


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

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
537
538
539
        h = self.gate_up_proj(hidden_states)
        h = self.act(h)
        return self.down_proj(h)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
540
541
542


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

585
586
        max_position = config.max_position_embeddings
        if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
587
588
589
            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
590
591
592

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

        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)
632
633
634
635
636
637
638
639

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

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

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

724
        self.start_layer, self.end_layer, self.layers = make_layers(
725
726
            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
727
728
729
730
731

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


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

        config = vllm_config.model_config.hf_config

        self.config = config
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            prefix=f"{prefix}.embed_tokens",
        )
759
760
761
762
        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
763
764
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

765
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
766
767
        return self.embed_tokens(input_ids)

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


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

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

823
824
825
        self.model = Plamo2Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
826
827
        self.vocab_size = self.config.vocab_size
        self.lm_head = ParallelLMHead(
828
            self.vocab_size,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
829
830
831
832
833
834
            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)

835
        self.logits_processor = LogitsProcessor(
836
            config.vocab_size, self.config.vocab_size
837
        )
838
        self.make_empty_intermediate_tensors = (
839
840
            self.model.make_empty_intermediate_tensors
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
841

842
843
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
844

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

858
859
860
861
862
863
864
865
866
    @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
867
        )
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883

    @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
884
        intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head
885
886
887
888
889
890
891
892
893

        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
894
895
        )

896
897
898
899
    @classmethod
    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
        return MambaStateCopyFuncCalculator.mamba2_state_copy_func()

Shinichi Hemmi's avatar
Shinichi Hemmi committed
900
901
902
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
903
    ) -> torch.Tensor | None:
904
        logits = self.logits_processor(self.lm_head, hidden_states)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
905
906
        return logits

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

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

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

Shinichi Hemmi's avatar
Shinichi Hemmi committed
973
974
975
976
977
978
979
980
981
982
983
984
            # 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

985
986
987
988
            # Skip layers on other devices.
            if is_pp_missing_parameter(name, self):
                continue

Shinichi Hemmi's avatar
Shinichi Hemmi committed
989
            param = params_dict[name]
990
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
991
            weight_loader(param, loaded_weight)