falcon_h1.py 23.6 KB
Newer Older
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
3
"""Inference-only FalconH1 model."""
4

Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
5
from collections.abc import Iterable
6
from itertools import islice
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
7
8
9
10
11
12

import torch
from torch import nn
from transformers import FalconH1Config

from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, ModelConfig, VllmConfig
15
from vllm.distributed import get_tensor_model_parallel_world_size
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
16
17
18
from vllm.distributed.parallel_state import get_pp_group
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
19
20
21
22
23
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
24
from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
26
from vllm.model_executor.layers.mamba.mamba_utils import (
27
28
29
    MambaStateDtypeCalculator,
    MambaStateShapeCalculator,
)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
30
31
32
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 (
33
34
35
    ParallelLMHead,
    VocabParallelEmbedding,
)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
36
37
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.sequence import IntermediateTensors
38
from vllm.transformers_utils.config import set_default_rope_theta
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
39

40
41
42
43
44
45
46
from .interfaces import (
    HasInnerState,
    IsHybrid,
    SupportsLoRA,
    SupportsMambaPrefixCaching,
    SupportsPP,
)
47
48
49
50
51
52
53
from .utils import (
    PPMissingLayer,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
54
55
56
57
58
59


class FalconH1MLP(nn.Module):
    def __init__(
        self,
        config: FalconH1Config,
60
        quant_config: QuantizationConfig | None = None,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
61
        bias: bool = False,
62
        prefix: str = "",
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
63
64
65
66
67
68
69
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=config.hidden_size,
            output_sizes=[config.intermediate_size] * 2,
            bias=bias,
            quant_config=quant_config,
70
            prefix=f"{prefix}.gate_up_proj",
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
71
72
73
74
75
76
        )
        self.down_proj = RowParallelLinear(
            input_size=config.intermediate_size,
            output_size=config.hidden_size,
            bias=bias,
            quant_config=quant_config,
77
            prefix=f"{prefix}.down_proj",
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
78
79
80
81
82
        )
        self.tp_size = get_tensor_model_parallel_world_size()
        self.intermediate_size = config.intermediate_size
        self.gate_multiplier, self.down_multiplier = config.mlp_multipliers
        if config.hidden_act != "silu":
83
84
85
86
            raise ValueError(
                f"Unsupported activation: {config.hidden_act}. "
                "Only silu is supported for now."
            )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
87
88
89
90
        self.act_fn = SiluAndMul()

    def forward(self, x):
        x, _ = self.gate_up_proj(x)
91
        x[:, : self.intermediate_size // self.tp_size] *= self.gate_multiplier
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
92
93
94
95
96
97
98
99
100
101
        x = self.act_fn(x)
        x, _ = self.down_proj(x)
        x = x * self.down_multiplier
        return x


class FalconH1SSMDecoderLayer(nn.Module):
    def __init__(
        self,
        config: FalconH1Config,
102
103
104
        model_config: ModelConfig | None = None,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
105
        prefix: str = "",
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
106
107
108
109
110
    ) -> None:
        super().__init__()
        self.config = config
        self.tp_size = get_tensor_model_parallel_world_size()

111
112
113
114
115
        self.d_ssm = (
            int(config.mamba_expand * config.hidden_size)
            if config.mamba_d_ssm is None
            else config.mamba_d_ssm
        )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
116
117
118
119
120
121
122
123
124
125
126
127
128

        self.mamba = MambaMixer2(
            hidden_size=config.hidden_size,
            ssm_state_size=config.mamba_d_state,
            conv_kernel_size=config.mamba_d_conv,
            intermediate_size=self.d_ssm,
            use_conv_bias=config.mamba_conv_bias,
            use_bias=config.mamba_proj_bias,
            n_groups=config.mamba_n_groups,
            num_heads=config.mamba_n_heads,
            head_dim=config.mamba_d_head,
            rms_norm_eps=config.rms_norm_eps,
            activation=config.hidden_act,
129
130
            model_config=model_config,
            cache_config=cache_config,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
131
132
            quant_config=quant_config,
            use_rms_norm=config.mamba_rms_norm,
133
            prefix=f"{prefix}.mixer",
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
134
135
136
137
138
139
140
141
        )
        # n_groups is overridden later by `MambaMixer2`
        self.groups_time_state_size = self.mamba.n_groups * config.mamba_d_state
        self.zxbcdt_multipliers = config.ssm_multipliers
        self._init_mup_vector()

    def _init_mup_vector(self):
        """
142
143
        Non learnable per-block scaling vector composed of element-wise
        multipliersapplied to each separate contiguous block of the output
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
144
145
146
147
148
149
        of the linear projection (in_proj) before further processing
        (gating, convolution, SSM):

            - Z block:  [0 : d_ssm]                      → zxbcdt_multipliers[0]
            - X block:  [d_ssm : 2 * d_ssm]              → zxbcdt_multipliers[1]
            - B block:  [2 * d_ssm : 2 * d_ssm + G * S]  → zxbcdt_multipliers[2]
150
            - C block:  [2 * d_ssm + G * S : 2 * d_ssm + 2 * G * S]
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
151
152
153
154
155
156
157
158
159
                        → zxbcdt_multipliers[3]
            - dt block: [2 * d_ssm + 2 * G * S : end]    → zxbcdt_multipliers[4]

        where:
            - d_ssm:     Dimension of state-space model latent
            - G:         Number of groups (n_groups)
            - S:         SSM state size per group
            - All indices are divided by tp_size to support tensor parallelism
        """
160
161
162
        vector_shape = (
            2 * self.d_ssm + 2 * self.groups_time_state_size + self.config.mamba_n_heads
        ) // self.tp_size
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
163
164
        mup_vector = torch.ones(1, vector_shape)
        # Z vector 0 -> d_ssm
165
        mup_vector[:, : self.d_ssm // self.tp_size] *= self.zxbcdt_multipliers[0]
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
166
        # X vector d_ssm -> 2 * d_ssm
167
168
169
        mup_vector[
            :, (self.d_ssm // self.tp_size) : (2 * self.d_ssm // self.tp_size)
        ] *= self.zxbcdt_multipliers[1]
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
170
171
172
        # B vector 2 * d_ssm -> 2 * d_ssm + (n_group * d_state)
        mup_vector[
            :,
173
174
175
176
            (2 * self.d_ssm) // self.tp_size : (
                2 * self.d_ssm + self.groups_time_state_size
            )
            // self.tp_size,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
177
178
179
180
181
        ] *= self.zxbcdt_multipliers[2]
        # C vector 2 * d_ssm + (n_group * d_state)
        # -> 2 * d_ssm + 2 * (n_group * d_state)
        mup_vector[
            :,
182
183
184
185
            (2 * self.d_ssm + self.groups_time_state_size) // self.tp_size : (
                2 * self.d_ssm + 2 * self.groups_time_state_size
            )
            // self.tp_size,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
186
187
188
189
190
        ] *= self.zxbcdt_multipliers[3]
        # dt vector 2 * d_ssm + 2 * (n_group * d_state)
        # -> 2 * d_ssm + 2 * (n_group * d_state) + n_heads
        mup_vector[
            :,
191
            (2 * self.d_ssm + 2 * self.groups_time_state_size) // self.tp_size :,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
192
193
194
195
196
197
198
        ] *= self.zxbcdt_multipliers[4]

        self.register_buffer("mup_vector", mup_vector, persistent=False)

    def forward(
        self,
        hidden_states: torch.Tensor,
199
        residual: torch.Tensor | None,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
200
201
        **kwargs,
    ):
202
        output = self.mamba(
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
203
204
205
            hidden_states,
            mup_vector=self.mup_vector,
        )
206
        return output, residual
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
207
208
209
210
211
212


class FalconH1AttentionDecoderLayer(nn.Module):
    def __init__(
        self,
        config: FalconH1Config,
213
214
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
215
216
217
        prefix: str = "",
    ) -> None:
        super().__init__()
218
        set_default_rope_theta(config, default_theta=1e11)
219
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
        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)
235
236
237
238
239
        self.head_dim = (
            config.hidden_size // self.total_num_heads
            if getattr(config, "head_dim", None) is None
            else config.head_dim
        )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
240
241
242
243
244
        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.max_position_embeddings = max_position_embeddings

245
        if hasattr(config, "attn_rotary_emb"):
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
246
247
248
249
250
251
252
253
            rotary_dim = config.attn_rotary_emb  # for backward compatibility
        else:
            rotary_dim = self.head_dim  # default

        self.rotary_emb = get_rope(
            head_size=self.head_dim,
            rotary_dim=rotary_dim,
            max_position=max_position_embeddings,
254
            rope_parameters=config.rope_parameters,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
            is_neox_style=True,
            dtype=None,  # see impl of get_rope
        )

        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,
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

        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",
        )
        self.key_multiplier = config.key_multiplier

    def self_attention(
        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)
        k = k * self.key_multiplier

        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
305
        residual: torch.Tensor | None,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
        **kwargs,
    ):
        hidden_states = self.self_attention(
            positions=positions,
            hidden_states=hidden_states,
        )
        return hidden_states, residual


class FalconH1ParallelHybrid(nn.Module):
    """
    A hybrid decoder layer for FalconH1 where the input is processed
    in parallel through both the self-attention branch and the SSM (Mamba)
    branch. Their outputs are then summed to produce the final hidden state.

    This layer uses:
      - FalconH1AttentionDecoderLayer for the multi-head self-attention branch.
      - FalconH1SSMDecoderLayer for the state-space (Mamba) branch.
    """

    def __init__(
        self,
        config: FalconH1Config,
        layer_idx: int,
330
331
332
        model_config: ModelConfig | None = None,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
333
334
335
        prefix: str = "",
    ) -> None:
        super().__init__()
336

Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
337
338
339
340
341
342
343
        # Instantiate the attention branch
        self.self_attn = FalconH1AttentionDecoderLayer(
            config=config,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=prefix,
        )
344
345
346
347
348
349

        # In V1 all attention/ssm layers must have
        # different index in prefix
        ssm_layer_idx = config.num_hidden_layers + layer_idx
        ssm_prefix = prefix.split(".")[0] + f".{ssm_layer_idx}"

Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
350
351
352
        # Instantiate the SSM branch
        self.mamba = FalconH1SSMDecoderLayer(
            config=config,
353
            model_config=model_config,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
354
355
            cache_config=cache_config,
            quant_config=quant_config,
356
            prefix=ssm_prefix,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
357
358
359
360
361
362
363
        )
        self.ssm_out_multiplier = config.ssm_out_multiplier
        self.ssm_in_multiplier = config.ssm_in_multiplier

        self.attention_in_multiplier = config.attention_in_multiplier
        self.attn_out_multiplier = config.attention_out_multiplier

364
        self.feed_forward = FalconH1MLP(config, prefix=f"{prefix}.feed_forward")
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
365

366
367
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        **kwargs,
    ):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        # Process input through the attention branch.
        # FalconH1AttentionDecoderLayer expects positions, hidden_states,
        # kv_cache, attn_metadata, and residual.
        attn_hidden, _ = self.self_attn(
            positions=positions,
            hidden_states=hidden_states * self.attention_in_multiplier,
            residual=residual,
            **kwargs,
        )

        # Process input through the SSM branch.
        # FalconH1SSMDecoderLayer expects hidden_states, attn_metadata,
389
        # residual, and sequence_idx.
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
390
391
392
393
394
395
396
397
398
        ssm_hidden, _ = self.mamba(
            hidden_states=hidden_states * self.ssm_in_multiplier,
            residual=residual,
            **kwargs,
        )
        # Sum the outputs from both branches.
        # We assume both branches produce outputs of the same
        # dimensionality (config.hidden_size).
        hidden_states = (attn_hidden * self.attn_out_multiplier) + (
399
400
            ssm_hidden * self.ssm_out_multiplier
        )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
401
402
403
404
405
406
407
408
409
410
411
        hidden_states = hidden_states + residual

        # feed-forward
        residual = hidden_states
        hidden_states = self.pre_ff_layernorm(hidden_states)
        hidden_states = self.feed_forward(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states


412
@support_torch_compile
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
413
414
415
416
class FalconH1Model(nn.Module):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config: FalconH1Config = vllm_config.model_config.hf_config
417
        model_config = vllm_config.model_config
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
418
419
420
421
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

        self.config = config
422
423
424

        self.vocab_size = config.vocab_size

Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
            )
            self.embedding_multiplier = config.embedding_multiplier
        else:
            self.embed_tokens = PPMissingLayer()
            self.embedding_multiplier = 1.0

        def get_layer(prefix: str):
            layer_idx = int(prefix.rsplit(".", 1)[1])
            layer_class = FalconH1ParallelHybrid
            return layer_class(
                config,
                layer_idx,
441
                model_config,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
442
443
444
445
446
447
                cache_config,
                quant_config=quant_config,
                prefix=prefix,
            )

        self.start_layer, self.end_layer, self.layers = make_layers(
448
449
450
451
452
            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
        )
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
453
        if get_pp_group().is_last_rank:
454
            self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
455
456
457
        else:
            self.final_layernorm = PPMissingLayer()

458
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
459
460
461
462
463
464
        return self.embed_tokens(input_ids)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
465
466
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
467
468
469
470
471
    ) -> torch.Tensor:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds * self.embedding_multiplier
            else:
472
                hidden_states = (
473
                    self.embed_input_ids(input_ids) * self.embedding_multiplier
474
                )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
475
476
477
478
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]

479
        for layer in islice(self.layers, self.start_layer, self.end_layer):
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
480
481
482
483
484
            hidden_states = layer(
                positions=positions,
                hidden_states=hidden_states,
            )
        if not get_pp_group().is_last_rank:
485
486
487
488
489
            return IntermediateTensors(
                {
                    "hidden_states": hidden_states,
                }
            )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
490
491
492
493
        hidden_states = self.final_layernorm(hidden_states)
        return hidden_states


494
495
496
497
498
499
500
501
class FalconH1ForCausalLM(
    nn.Module,
    HasInnerState,
    SupportsLoRA,
    SupportsPP,
    IsHybrid,
    SupportsMambaPrefixCaching,
):
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
502
503
504
505
506
507
508
509
510
511
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }

    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }

512
513
514
515
516
517
518
519
520
521
522
    @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,
        )

523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
    @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

541
542
543
544
545
        intermediate_size = (
            int(hf_config.mamba_expand * hf_config.hidden_size)
            if hf_config.mamba_d_ssm is None
            else hf_config.mamba_d_ssm
        )
546

547
        return MambaStateShapeCalculator.mamba2_state_shape(
548
549
550
551
552
553
554
555
556
            intermediate_size=intermediate_size,
            tp_world_size=parallel_config.tensor_parallel_size,
            n_groups=hf_config.mamba_n_groups,
            num_heads=hf_config.mamba_n_heads,
            head_dim=hf_config.mamba_d_head,
            state_size=hf_config.mamba_d_state,
            conv_kernel=hf_config.mamba_d_conv,
        )

Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
557
558
559
560
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        config = vllm_config.model_config.hf_config
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
561

Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
562
563
564
565
566
567
568
        scheduler_config = vllm_config.scheduler_config

        self.quant_config = vllm_config.quant_config

        super().__init__()
        self.config = config
        self.scheduler_config = scheduler_config
569
570
571
        self.model = FalconH1Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
572
        self.tie_word_embeddings = config.tie_word_embeddings
573

Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
574
575
        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(
576
                config.vocab_size,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
577
                config.hidden_size,
578
                prefix=maybe_prefix(prefix, "lm_head"),
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
579
580
581
            )
            self.lm_head_multiplier = config.lm_head_multiplier
            if self.tie_word_embeddings:
582
                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
583
584
585
            # Used to track and store by the Mamba cache between steps.

            self.logits_processor = LogitsProcessor(
586
                config.vocab_size,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
587
588
589
590
591
592
593
                config.vocab_size,
                scale=config.lm_head_multiplier,
            )
        else:
            self.lm_head = PPMissingLayer()

        self.make_empty_intermediate_tensors = (
594
595
            self.model.make_empty_intermediate_tensors
        )
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
596

597
598
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
599
600
601
602
603

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
604
605
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
606
607
608
609
610
611
612
613
614
615
616
617
618
619
        **kwargs,
    ):
        hidden_states = self.model(
            input_ids,
            positions,
            intermediate_tensors,
            inputs_embeds,
        )

        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
620
    ) -> torch.Tensor | None:
621
        logits = self.logits_processor(self.lm_head, hidden_states)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
622
623
624

        return logits

625
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue

            if "A_log" in name:
                name = name.replace("A_log", "A")

            if "mamba" in name:
                name = name.replace("mamba", "mamba.mamba")

            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue

                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                # Skip layers on other devices.
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if is_pp_missing_parameter(name, self):
                    continue
                if self.tie_word_embeddings and "lm_head" in name:
                    continue

                param = params_dict[name]
672
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
673
674
675
676
677
678
                weight_loader(param, loaded_weight)
            loaded_params.add(name)

        if self.tie_word_embeddings:
            loaded_params.add("lm_head.weight")
        return loaded_params