A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 00:45 UTC
Fix Plamo 2/3 & LFM2 for Transformers v5 (#38090)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 25, 2026 at 12:29 UTC
d6bb2a9d9a3453b86e40bc7b37d39ab7bd6f0969
3 files changed
+60
-55
vllm/model_executor/models/lfm2.py
+10
-8
@@ -52,7 +52,7 @@ class Lfm2MLP(nn.Module):
52
def __init__(
53
self,
54
dim: int,
55
- ff_dim: int,
55
+ intermediate_size: int,
56
multiple_of: int,
57
auto_adjust_ff_dim: bool,
58
ffn_dim_multiplier: float | None,
@@ -61,21 +61,23 @@ class Lfm2MLP(nn.Module):
61
):
62
super().__init__()
63
if auto_adjust_ff_dim:
64
- ff_dim = int(2 * ff_dim / 3)
64
+ intermediate_size = int(2 * intermediate_size / 3)
65
# custom dim factor multiplier
66
if ffn_dim_multiplier is not None:
67
- ff_dim = int(ffn_dim_multiplier * ff_dim)
68
- ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
67
+ intermediate_size = int(ffn_dim_multiplier * intermediate_size)
68
+ intermediate_size = multiple_of * (
69
+ (intermediate_size + multiple_of - 1) // multiple_of
70
+ )
71
72
self.w13 = MergedColumnParallelLinear(
73
input_size=dim,
72
- output_sizes=[ff_dim] * 2,
74
+ output_sizes=[intermediate_size] * 2,
75
bias=False,
76
quant_config=quant_config,
77
prefix=f"{prefix}.w13",
78
)
79
self.w2 = RowParallelLinear(
78
- input_size=ff_dim,
80
+ input_size=intermediate_size,
81
output_size=dim,
82
bias=False,
83
quant_config=quant_config,
@@ -212,7 +214,7 @@ class Lfm2AttentionDecoderLayer(nn.Module):
214
215
self.feed_forward = Lfm2MLP(
216
dim=config.block_dim,
215
- ff_dim=config.block_ff_dim,
217
+ intermediate_size=config.intermediate_size,
218
multiple_of=config.block_multiple_of,
219
auto_adjust_ff_dim=config.block_auto_adjust_ff_dim,
220
ffn_dim_multiplier=config.block_ffn_dim_multiplier,
@@ -262,7 +264,7 @@ class Lfm2ShortConvDecoderLayer(nn.Module):
264
265
self.feed_forward = Lfm2MLP(
266
dim=config.block_dim,
265
- ff_dim=config.block_ff_dim,
267
+ intermediate_size=config.intermediate_size,
268
multiple_of=config.block_multiple_of,
269
auto_adjust_ff_dim=config.block_auto_adjust_ff_dim,
270
ffn_dim_multiplier=config.block_ffn_dim_multiplier,
vllm/model_executor/models/plamo2.py
+26
-24
@@ -4,6 +4,7 @@
4
5
from collections.abc import Iterable
6
from itertools import islice
7
+from typing import TYPE_CHECKING
8
9
import torch
10
from torch import nn
@@ -71,30 +72,31 @@ from vllm.utils.torch_utils import direct_register_custom_op
72
from vllm.v1.attention.backend import AttentionMetadata
73
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
74
74
-
75
# Only used for type hinting.
76
-class Plamo2Config(PretrainedConfig): # type: ignore
77
- model_type: str = "plamo2"
78
-
79
- hidden_size: int
80
- num_hidden_layers: int
81
- rms_norm_eps: float
82
- # Attention
83
- num_attention_heads: int
84
- hidden_size_per_head: int
85
- num_key_value_heads: int
86
- # Mamba
87
- mamba_d_state: int
88
- mamba_d_conv: int
89
- mamba_num_heads: int
90
- mamba_step: int
91
- # MLP
92
- intermediate_size: int
93
- # Tokenizer
94
- vocab_size: int
95
-
96
-
97
-def is_mamba(config: Plamo2Config, i: int) -> bool:
76
+if TYPE_CHECKING:
77
+
78
+ class Plamo2Config(PretrainedConfig): # type: ignore
79
+ model_type: str = "plamo2"
80
+
81
+ hidden_size: int
82
+ num_hidden_layers: int
83
+ rms_norm_eps: float
84
+ # Attention
85
+ num_attention_heads: int
86
+ hidden_size_per_head: int
87
+ num_key_value_heads: int
88
+ # Mamba
89
+ mamba_d_state: int
90
+ mamba_d_conv: int
91
+ mamba_num_heads: int
92
+ mamba_step: int
93
+ # MLP
94
+ intermediate_size: int
95
+ # Tokenizer
96
+ vocab_size: int
97
+
98
+
99
+def is_mamba(config: "Plamo2Config", i: int) -> bool:
100
assert config.mamba_step > 1
101
102
if config.num_hidden_layers <= (config.mamba_step // 2):
@@ -502,7 +504,7 @@ direct_register_custom_op(
504
class DenseMLP(nn.Module):
505
def __init__(
506
self,
505
- config: Plamo2Config,
507
+ config: "Plamo2Config",
508
quant_config: QuantizationConfig | None = None,
509
prefix: str = "",
510
) -> None:
vllm/model_executor/models/plamo3.py
+24
-23
@@ -4,7 +4,7 @@
4
5
from collections.abc import Iterable
6
from itertools import islice
7
-from typing import Any
7
+from typing import TYPE_CHECKING, Any
8
9
import torch
10
from torch import nn
@@ -46,28 +46,29 @@ from vllm.model_executor.models.utils import (
46
from vllm.model_executor.utils import set_weight_attrs
47
from vllm.sequence import IntermediateTensors
48
49
-
49
# Only used for type hinting.
51
-class Plamo3Config(PretrainedConfig): # type: ignore
52
- model_type: str = "plamo3"
53
-
54
- hidden_size: int
55
- num_hidden_layers: int
56
- rms_norm_eps: float
57
- # Attention
58
- num_attention_heads: int
59
- head_dim: int
60
- num_key_value_heads: int
61
- # vllm rename `sliding_window` attr to `interleaved_sliding_window`
62
- # if `sliding_window` is list
63
- interleaved_sliding_window: list[int | None]
64
- sliding_window_pattern: int
65
- rope_parameters: dict[str, Any]
66
- rope_local_theta: int
67
- # MLP
68
- intermediate_size: int
69
- # Tokenizer
70
- vocab_size: int
50
+if TYPE_CHECKING:
51
+
52
+ class Plamo3Config(PretrainedConfig): # type: ignore
53
+ model_type: str = "plamo3"
54
+
55
+ hidden_size: int
56
+ num_hidden_layers: int
57
+ rms_norm_eps: float
58
+ # Attention
59
+ num_attention_heads: int
60
+ head_dim: int
61
+ num_key_value_heads: int
62
+ # vllm rename `sliding_window` attr to `interleaved_sliding_window`
63
+ # if `sliding_window` is list
64
+ interleaved_sliding_window: list[int | None]
65
+ sliding_window_pattern: int
66
+ rope_parameters: dict[str, Any]
67
+ rope_local_theta: int
68
+ # MLP
69
+ intermediate_size: int
70
+ # Tokenizer
71
+ vocab_size: int
72
73
74
def rms_norm_weight_loader(offset: float) -> LoaderFunction:
@@ -80,7 +81,7 @@ def rms_norm_weight_loader(offset: float) -> LoaderFunction:
81
class DenseMLP(nn.Module):
82
def __init__(
83
self,
83
- config: Plamo3Config,
84
+ config: "Plamo3Config",
85
quant_config: QuantizationConfig | None = None,
86
prefix: str = "",
87
) -> None: