plamo2.py 37 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
8
9
10

if TYPE_CHECKING:
    from vllm.attention.backends.abstract import AttentionBackend
Shinichi Hemmi's avatar
Shinichi Hemmi committed
11
12
13

import torch
from torch import nn
14
from transformers import PretrainedConfig
Shinichi Hemmi's avatar
Shinichi Hemmi committed
15
16
17

from vllm.attention.backends.abstract import AttentionMetadata
from vllm.attention.layer import Attention
18
from vllm.compilation.decorators import support_torch_compile
19
from vllm.config import VllmConfig, get_current_vllm_config
20
21
from vllm.distributed import divide, get_tensor_model_parallel_world_size
from vllm.distributed.parallel_state import get_pp_group
22
23
from vllm.forward_context import ForwardContext, get_forward_context
from vllm.model_executor.custom_op import CustomOp
24
from vllm.model_executor.layers.activation import SiluAndMul
Shinichi Hemmi's avatar
Shinichi Hemmi committed
25
from vllm.model_executor.layers.layernorm import RMSNorm
26
27
28
29
30
31
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
32
from vllm.model_executor.layers.logits_processor import LogitsProcessor
33
34
from vllm.model_executor.layers.mamba.abstract import MambaBase
from vllm.model_executor.layers.mamba.mamba_utils import (
35
36
37
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
38
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
39
40
41
42
    causal_conv1d_fn,
    causal_conv1d_update,
)
from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update
43
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
44
45
    mamba_chunk_scan_combined_varlen,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
46
47
48
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 (
49
50
51
52
    DEFAULT_VOCAB_PADDING_SIZE,
    ParallelLMHead,
    VocabParallelEmbedding,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
53
from vllm.model_executor.model_loader.weight_utils import (
54
55
56
57
58
    composed_weight_loader,
    default_weight_loader,
    sharded_weight_loader,
)
from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid, SupportsPP
59
from vllm.model_executor.models.utils import (
60
61
62
63
64
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
65
66
from vllm.model_executor.utils import set_weight_attrs
from vllm.sequence import IntermediateTensors
67
from vllm.utils.torch_utils import direct_register_custom_op
68
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
Shinichi Hemmi's avatar
Shinichi Hemmi committed
69
70
71
72
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


# 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)


102
103
104
# Adapted from:
# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
# transformers.models.mamba.modeling_mamba.MambaMixer
105
106
@CustomOp.register(name="plamo2_mamba_mixer")
class Plamo2MambaMixer(MambaBase, CustomOp):
107
    def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
108
        super().__init__()
109
        self.config = vllm_config.model_config.hf_config
110
111
        self.cache_config = vllm_config.cache_config
        self.model_config = vllm_config.model_config
112
113
114
115
        self.quant_config = vllm_config.quant_config
        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
116
117
118
        self.intermediate_size = (
            self.config.mamba_num_heads * self.config.hidden_size_per_head
        )
119
120
121
        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
122
123
124
125
        self.time_step_rank = max(64, self.hidden_size // 16)
        self.conv1d = ColumnParallelLinear(
            input_size=self.conv_kernel_size,
            output_size=self.intermediate_size,
126
127
128
            bias=False,
            prefix=f"{prefix}.conv1d",
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
129
130
131
132
133
134
135
136
137
138
        )
        # 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,
139
140
            bias=False,
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
141
            prefix=f"{prefix}.in_proj",
142
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
143
144
145
146
147
148
        )
        # 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,
149
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
150
            prefix=f"{prefix}.bcdt_proj",
151
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
152
153
154
155
156
157
158
159
        )
        # 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,
160
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
161
            prefix=f"{prefix}.dt_proj",
162
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
163
164
165
166
        )

        self.A = nn.Parameter(
            torch.empty(
167
                divide(self.num_heads, self.tp_size),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
168
                dtype=torch.float32,
169
170
            )
        )
171
        self.D = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size)))
172
        self.dt_bias = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size)))
Shinichi Hemmi's avatar
Shinichi Hemmi committed
173
174
175

        set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)})
        a_weight_loader = composed_weight_loader(
176
177
            sharded_weight_loader(0), lambda x: -torch.exp(x.float())
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
178
        set_weight_attrs(self.A, {"weight_loader": a_weight_loader})
179
        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
Shinichi Hemmi's avatar
Shinichi Hemmi committed
180
181
182
183

        self.out_proj = RowParallelLinear(
            self.intermediate_size,
            self.hidden_size,
184
            bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
185
            input_is_parallel=True,
186
            quant_config=self.quant_config,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
187
            prefix=f"{prefix}.out_proj",
188
            return_bias=False,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
189
190
191
192
        )
        # The activation function is fixed to SiLU.
        self.activation = "silu"

193
194
195
        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)
196

197
198
        self.chunk_size = self.config.mamba_chunk_size

199
200
201
202
203
204
205
        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"
206
207
208

        self.prefix = prefix

209
210
211
212
213
214
215
216
217
218
219
220
221
222
    def _project_ssm_parameters(self, hidden_states):
        ssm_parameters = self.bcdt_proj(hidden_states)
        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
223

224
225
226
227
228
229
230
231
    def forward_native(
        self,
        hidden_states: torch.Tensor,
        output: torch.Tensor,
        **kwargs,
    ):
        pass

Shinichi Hemmi's avatar
Shinichi Hemmi committed
232
233
234
    def forward(
        self,
        hidden_states: torch.Tensor,
235
        output: torch.Tensor,
Shinichi Hemmi's avatar
Shinichi Hemmi committed
236
        **kwargs,
237
    ):
238
239
240
241
242
        torch.ops.vllm.plamo2_mamba_mixer(
            hidden_states,
            output,
            self.prefix,
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
243

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

        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
270
            query_start_loc_p = attn_metadata.query_start_loc_p
271
272
            cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
            last_chunk_indices_p = attn_metadata.last_chunk_indices_p
273

Shinichi Hemmi's avatar
Shinichi Hemmi committed
274
        # 1. Gated MLP's linear projection
275
276
        projected_states = self.in_proj(hidden_states)
        gate, hidden_states = projected_states.chunk(2, dim=-1)
Shinichi Hemmi's avatar
Shinichi Hemmi committed
277
278

        # 2. Convolution sequence transformation
279
280
281
        conv_weights = self.conv1d.weight.view(
            self.conv1d.weight.size(0), self.conv1d.weight.size(2)
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
282

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

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

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

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

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

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

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

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

            # 3. State Space Model sequence transformation
412
413
414
            A = self.A[:, None, ...][:, :, None].expand(
                -1, self.head_dim, self.config.mamba_d_state
            )
415
416
417
418
419
420
            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(
421
422
                -1, self.num_heads // self.tp_size, self.head_dim
            )
423
424

            # - the hidden is reshaped into (bs, num_heads, head_dim)
425
            # - ssm_state's slots will be selected
426
            #   using state_indices_tensor_d
427
428

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

        # 4. Final linear projection
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
        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 get_attn_backend(self) -> type["AttentionBackend"]:
472
473
        from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend

474
475
476
477
478
479
480
481
482
483
        return Mamba2AttentionBackend


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]
484
    self.forward_cuda(hidden_states=hidden_states, output=output)
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500


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


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

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


class Plamo2AttentionMixer(nn.Module):
538
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
Shinichi Hemmi's avatar
Shinichi Hemmi committed
539
        super().__init__()
540
541
542
        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
543
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
        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,
570
            prefix=f"{prefix}.qkv_proj",
Shinichi Hemmi's avatar
Shinichi Hemmi committed
571
        )
572
573
574
575
576
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
577
            prefix=f"{prefix}.o_proj",
578
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
579

580
581
582
583
        self.rope_theta = config.rope_theta if hasattr(config, "rope_theta") else 10000
        self.rope_scaling = (
            config.rope_scaling if hasattr(config, "rope_scaling") else None
        )
584
585
        max_position = config.max_position_embeddings
        if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
586
587
588
            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
589
590
591
592

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

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

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

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

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

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

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


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

        config = vllm_config.model_config.hf_config

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

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

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

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


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

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

826
827
828
        self.model = Plamo2Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
829
830
831
832
833
834
835
836
837
838
839
840
841
        self.vocab_size = self.config.vocab_size
        self.unpadded_vocab_size = self.config.vocab_size
        num_embeddings = ((self.vocab_size + 15) // 16) * 16
        self.lm_head = ParallelLMHead(
            num_embeddings,
            self.config.hidden_size,
            org_num_embeddings=self.config.vocab_size,
            padding_size=DEFAULT_VOCAB_PADDING_SIZE,
            prefix=f"{prefix}.lm_head",
        )
        if self.config.tie_word_embeddings:
            self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

842
843
844
        self.logits_processor = LogitsProcessor(
            self.unpadded_vocab_size, self.config.vocab_size
        )
845
        self.make_empty_intermediate_tensors = (
846
847
            self.model.make_empty_intermediate_tensors
        )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
848

849
850
851
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

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

865
866
867
868
869
870
871
872
873
    @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
874
        )
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890

    @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
891
        intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head
892
893
894
895
896
897
898
899
900

        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
901
902
903
904
905
        )

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

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

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

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

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

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

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