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 14, 05:30 UTC
[Model] Lfm2Moe (#26344)
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Paul Pak committed
Oct 8, 2025 at 01:03 UTC
320feae6f506097c47b6b41a634a6197512cffc1
8 files changed
+967
-7
docs/models/supported_models.md
+1
@@ -390,6 +390,7 @@ th {
390
| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ | ✅︎ |
391
| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ | ✅︎ |
392
| `Lfm2ForCausalLM` | LFM2 | `LiquidAI/LFM2-1.2B`, `LiquidAI/LFM2-700M`, `LiquidAI/LFM2-350M`, etc. | ✅︎ | ✅︎ | ✅︎ |
393
+| `Lfm2MoeForCausalLM` | LFM2MoE | `LiquidAI/LFM2-8B-A1B-preview`, etc. | ✅︎ | ✅︎ | ✅︎ |
394
| `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA, Yi | `meta-llama/Meta-Llama-3.1-405B-Instruct`, `meta-llama/Meta-Llama-3.1-70B`, `meta-llama/Meta-Llama-3-70B-Instruct`, `meta-llama/Llama-2-70b-hf`, `01-ai/Yi-34B`, etc. | ✅︎ | ✅︎ | ✅︎ |
395
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ | ✅︎ |
396
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ | ✅︎ |
tests/models/registry.py
+3
@@ -321,6 +321,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
321
"Lfm2ForCausalLM": _HfExamplesInfo(
322
"LiquidAI/LFM2-1.2B", min_transformers_version="4.54"
323
),
324
+ "Lfm2MoeForCausalLM": _HfExamplesInfo(
325
+ "LiquidAI/LFM2-8B-A1B", min_transformers_version="4.58"
326
+ ),
327
"LlamaForCausalLM": _HfExamplesInfo(
328
"meta-llama/Llama-3.2-1B-Instruct",
329
extras={
vllm/model_executor/models/lfm2.py
+2
-7
@@ -71,14 +71,14 @@ class Lfm2MLP(nn.Module):
71
output_sizes=[ff_dim] * 2,
72
bias=False,
73
quant_config=quant_config,
74
- prefix=f"{prefix}.gate_up_proj",
74
+ prefix=f"{prefix}.w1",
75
)
76
self.w2 = RowParallelLinear(
77
input_size=ff_dim,
78
output_size=dim,
79
bias=False,
80
quant_config=quant_config,
81
- prefix=f"{prefix}.down_proj",
81
+ prefix=f"{prefix}.w2",
82
)
83
self.act_fn = SiluAndMul()
84
@@ -484,17 +484,12 @@ class Lfm2ForCausalLM(
484
quant_config = vllm_config.quant_config
485
cache_config = vllm_config.cache_config
486
lora_config = vllm_config.lora_config
487
- scheduler_config = vllm_config.scheduler_config
487
assert not cache_config.enable_prefix_caching, (
488
"Lfm2 currently does not support prefix caching"
489
)
490
491
super().__init__()
492
self.config = config
494
- self.vllm_config = vllm_config
495
- self.scheduler_config = scheduler_config
496
- self.model_config = vllm_config.model_config
497
-
493
self.model = Lfm2Model(
494
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
495
)
vllm/model_executor/models/lfm2_moe.py
new
+797
@@ -0,0 +1,797 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
+from typing import Any, Optional
5
+
6
+import torch
7
+import torch.nn as nn
8
+
9
+from vllm.attention import Attention
10
+from vllm.compilation.decorators import support_torch_compile
11
+from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
12
+from vllm.distributed import (
13
+ get_ep_group,
14
+ get_pp_group,
15
+ get_tensor_model_parallel_world_size,
16
+)
17
+from vllm.model_executor.layers.activation import SiluAndMul
18
+from vllm.model_executor.layers.fused_moe import FusedMoE
19
+from vllm.model_executor.layers.layernorm import RMSNorm
20
+from vllm.model_executor.layers.linear import (
21
+ MergedColumnParallelLinear,
22
+ QKVParallelLinear,
23
+ ReplicatedLinear,
24
+ RowParallelLinear,
25
+)
26
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
27
+from vllm.model_executor.layers.mamba.mamba_utils import (
28
+ MambaStateDtypeCalculator,
29
+ MambaStateShapeCalculator,
30
+)
31
+from vllm.model_executor.layers.mamba.short_conv import ShortConv
32
+from vllm.model_executor.layers.quantization import QuantizationConfig
33
+from vllm.model_executor.layers.rotary_embedding import get_rope
34
+from vllm.model_executor.layers.vocab_parallel_embedding import (
35
+ DEFAULT_VOCAB_PADDING_SIZE,
36
+ ParallelLMHead,
37
+ VocabParallelEmbedding,
38
+)
39
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
40
+from vllm.sequence import IntermediateTensors
41
+from vllm.transformers_utils.configs import Lfm2MoeConfig
42
+
43
+from .interfaces import (
44
+ HasInnerState,
45
+ IsHybrid,
46
+ MixtureOfExperts,
47
+ SupportsLoRA,
48
+ SupportsPP,
49
+ SupportsQuant,
50
+)
51
+from .utils import (
52
+ AutoWeightsLoader,
53
+ PPMissingLayer,
54
+ extract_layer_index,
55
+ is_pp_missing_parameter,
56
+ make_empty_intermediate_tensors_factory,
57
+ make_layers,
58
+ maybe_prefix,
59
+)
60
+
61
+
62
+class Lfm2MoeMlp(nn.Module):
63
+ def __init__(
64
+ self,
65
+ dim: int,
66
+ ff_dim: int,
67
+ quant_config: Optional[QuantizationConfig] = None,
68
+ prefix: str = "",
69
+ ):
70
+ super().__init__()
71
+ self.w1 = MergedColumnParallelLinear(
72
+ input_size=dim,
73
+ output_sizes=[ff_dim] * 2,
74
+ bias=False,
75
+ quant_config=quant_config,
76
+ prefix=f"{prefix}.w1",
77
+ )
78
+ self.w2 = RowParallelLinear(
79
+ input_size=ff_dim,
80
+ output_size=dim,
81
+ bias=False,
82
+ quant_config=quant_config,
83
+ prefix=f"{prefix}.w2",
84
+ )
85
+ self.act_fn = SiluAndMul()
86
+
87
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
88
+ gate_up, _ = self.w1(x)
89
+ x = self.act_fn(gate_up)
90
+ x, _ = self.w2(x)
91
+ return x
92
+
93
+
94
+class Lfm2MoeSparseMoeBlock(nn.Module):
95
+ def __init__(
96
+ self,
97
+ config: Lfm2MoeConfig,
98
+ quant_config: Optional[QuantizationConfig] = None,
99
+ prefix: str = "",
100
+ enable_eplb: bool = False,
101
+ ):
102
+ super().__init__()
103
+ self.tp_size = get_tensor_model_parallel_world_size()
104
+ self.routed_scaling_factor = config.routed_scaling_factor
105
+
106
+ self.ep_group = get_ep_group().device_group
107
+ self.ep_rank = self.ep_group.rank()
108
+ self.ep_size = self.ep_group.size()
109
+ self.n_routed_experts = config.num_experts
110
+
111
+ if self.tp_size > self.n_routed_experts:
112
+ raise ValueError(
113
+ f"Tensor parallel size {self.tp_size} is greater than "
114
+ f"the number of experts {self.n_routed_experts}."
115
+ )
116
+
117
+ # Load balancing settings.
118
+ vllm_config = get_current_vllm_config()
119
+ eplb_config = vllm_config.parallel_config.eplb_config
120
+ self.enable_eplb = enable_eplb
121
+
122
+ self.n_logical_experts = self.n_routed_experts
123
+ self.n_redundant_experts = eplb_config.num_redundant_experts
124
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
125
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
126
+
127
+ self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
128
+ self.physical_expert_end = (
129
+ self.physical_expert_start + self.n_local_physical_experts
130
+ )
131
+
132
+ self.gate = ReplicatedLinear(
133
+ config.hidden_size,
134
+ config.num_experts,
135
+ bias=False,
136
+ quant_config=quant_config,
137
+ prefix=f"{prefix}.gate",
138
+ )
139
+ if config.use_expert_bias:
140
+ self.gate.e_score_correction_bias = nn.Parameter(
141
+ torch.empty(self.n_routed_experts, dtype=torch.float32)
142
+ )
143
+ else:
144
+ self.gate.e_score_correction_bias = None
145
+
146
+ self.experts = FusedMoE(
147
+ num_experts=self.n_routed_experts,
148
+ top_k=config.num_experts_per_tok,
149
+ hidden_size=config.hidden_size,
150
+ intermediate_size=config.moe_intermediate_size,
151
+ reduce_results=False,
152
+ renormalize=config.norm_topk_prob,
153
+ quant_config=quant_config,
154
+ use_grouped_topk=True, # needed for softmax score func
155
+ num_expert_group=1,
156
+ topk_group=1,
157
+ prefix=f"{prefix}.experts",
158
+ enable_eplb=self.enable_eplb,
159
+ num_redundant_experts=self.n_redundant_experts,
160
+ scoring_func="sigmoid",
161
+ e_score_correction_bias=self.gate.e_score_correction_bias,
162
+ )
163
+
164
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
165
+ orig_shape = hidden_states.shape
166
+ hidden_dim = hidden_states.shape[-1]
167
+ hidden_states = hidden_states.view(-1, hidden_dim)
168
+
169
+ # router_logits: (num_tokens, n_experts)
170
+ router_logits, _ = self.gate(hidden_states)
171
+ final_hidden_states = (
172
+ self.experts(hidden_states=hidden_states, router_logits=router_logits)
173
+ * self.routed_scaling_factor
174
+ )
175
+
176
+ if self.tp_size > 1:
177
+ final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
178
+ final_hidden_states
179
+ )
180
+
181
+ return final_hidden_states.view(orig_shape)
182
+
183
+
184
+class Lfm2MoeAttention(nn.Module):
185
+ def __init__(
186
+ self,
187
+ config: Lfm2MoeConfig,
188
+ layer_idx: int,
189
+ hidden_size: int,
190
+ num_heads: int,
191
+ num_kv_heads: int,
192
+ rope_theta: float = 10000,
193
+ rope_scaling: Optional[dict[str, Any]] = None,
194
+ max_position_embeddings: int = 8192,
195
+ cache_config: Optional[CacheConfig] = None,
196
+ quant_config: Optional[QuantizationConfig] = None,
197
+ prefix: str = "",
198
+ ) -> None:
199
+ super().__init__()
200
+ self.layer_idx = layer_idx
201
+ self.hidden_size = hidden_size
202
+ self.num_kv_heads = num_kv_heads
203
+ tp_size = get_tensor_model_parallel_world_size()
204
+ self.total_num_heads = num_heads
205
+ assert self.total_num_heads % tp_size == 0
206
+ self.num_heads = self.total_num_heads // tp_size
207
+ self.total_num_kv_heads = num_kv_heads
208
+ if self.total_num_kv_heads >= tp_size:
209
+ # Number of KV heads is greater than TP size, so we partition
210
+ # the KV heads across multiple tensor parallel GPUs.
211
+ assert self.total_num_kv_heads % tp_size == 0
212
+ else:
213
+ # Number of KV heads is less than TP size, so we replicate
214
+ # the KV heads across multiple tensor parallel GPUs.
215
+ assert tp_size % self.total_num_kv_heads == 0
216
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
217
+ self.head_dim = self.hidden_size // self.total_num_heads
218
+
219
+ self.q_size = self.num_heads * self.head_dim
220
+ self.kv_size = self.num_kv_heads * self.head_dim
221
+ self.scaling = self.head_dim**-0.5
222
+ self.rope_theta = rope_theta
223
+ self.max_position_embeddings = max_position_embeddings
224
+
225
+ self.qkv_proj = QKVParallelLinear(
226
+ hidden_size=self.hidden_size,
227
+ head_size=self.head_dim,
228
+ total_num_heads=self.total_num_heads,
229
+ total_num_kv_heads=self.total_num_kv_heads,
230
+ bias=False,
231
+ quant_config=quant_config,
232
+ prefix=f"{prefix}.qkv_proj",
233
+ )
234
+ self.out_proj = RowParallelLinear(
235
+ input_size=self.total_num_heads * self.head_dim,
236
+ output_size=self.hidden_size,
237
+ bias=False,
238
+ quant_config=quant_config,
239
+ prefix=f"{prefix}.out_proj",
240
+ )
241
+ self.rotary_emb = get_rope(
242
+ self.head_dim,
243
+ rotary_dim=self.head_dim,
244
+ max_position=self.max_position_embeddings,
245
+ base=self.rope_theta,
246
+ rope_scaling=rope_scaling,
247
+ is_neox_style=True,
248
+ )
249
+ self.attn = Attention(
250
+ self.num_heads,
251
+ self.head_dim,
252
+ self.scaling,
253
+ num_kv_heads=self.num_kv_heads,
254
+ cache_config=cache_config,
255
+ prefix=f"{prefix}.attn",
256
+ )
257
+ self.q_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
258
+ self.k_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
259
+
260
+ def forward(
261
+ self,
262
+ positions: torch.Tensor,
263
+ hidden_states: torch.Tensor,
264
+ ) -> torch.Tensor:
265
+ n_tokens, _ = hidden_states.shape
266
+ qkv, _ = self.qkv_proj(hidden_states)
267
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
268
+ q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous()
269
+ k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous()
270
+ q = self.q_layernorm(q)
271
+ k = self.k_layernorm(k)
272
+ q, k = self.rotary_emb(positions, q, k)
273
+ q = q.view(n_tokens, self.num_heads * self.head_dim)
274
+ k = k.view(n_tokens, self.num_kv_heads * self.head_dim)
275
+ attn_output = self.attn(q, k, v)
276
+ output, _ = self.out_proj(attn_output)
277
+ return output
278
+
279
+
280
+class Lfm2MoeAttentionDecoderLayer(nn.Module):
281
+ def __init__(
282
+ self,
283
+ config: Lfm2MoeConfig,
284
+ layer_idx: int,
285
+ model_config: Optional[ModelConfig] = None,
286
+ cache_config: Optional[CacheConfig] = None,
287
+ quant_config: Optional[QuantizationConfig] = None,
288
+ prefix: str = "",
289
+ enable_eplb: bool = False,
290
+ ) -> None:
291
+ super().__init__()
292
+ self.prefix = prefix
293
+ self.config = config
294
+ self.layer_idx = layer_idx
295
+
296
+ rope_theta = getattr(config, "rope_theta", 10000)
297
+ rope_scaling = getattr(config, "rope_scaling", None)
298
+ if rope_scaling is not None and getattr(
299
+ config, "original_max_position_embeddings", None
300
+ ):
301
+ rope_scaling["original_max_position_embeddings"] = (
302
+ config.original_max_position_embeddings
303
+ )
304
+ max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
305
+
306
+ self.self_attn = Lfm2MoeAttention(
307
+ config=config,
308
+ layer_idx=layer_idx,
309
+ hidden_size=config.hidden_size,
310
+ num_heads=config.num_attention_heads,
311
+ num_kv_heads=config.num_key_value_heads,
312
+ rope_theta=rope_theta,
313
+ rope_scaling=rope_scaling,
314
+ max_position_embeddings=max_position_embeddings,
315
+ cache_config=cache_config,
316
+ quant_config=quant_config,
317
+ prefix=f"{prefix}.self_attn",
318
+ )
319
+
320
+ if layer_idx < config.num_dense_layers:
321
+ self.feed_forward = Lfm2MoeMlp(
322
+ dim=config.hidden_size,
323
+ ff_dim=config.intermediate_size,
324
+ quant_config=quant_config,
325
+ prefix=f"{prefix}.feed_forward",
326
+ )
327
+ else:
328
+ self.feed_forward = Lfm2MoeSparseMoeBlock(
329
+ config=config,
330
+ quant_config=quant_config,
331
+ prefix=f"{prefix}.feed_forward",
332
+ enable_eplb=enable_eplb,
333
+ )
334
+
335
+ self.operator_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
336
+ self.ffn_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
337
+
338
+ def forward(
339
+ self,
340
+ positions: torch.Tensor,
341
+ hidden_states: torch.Tensor,
342
+ residual: Optional[torch.Tensor],
343
+ **kwargs,
344
+ ) -> tuple[torch.Tensor, torch.Tensor]:
345
+ if residual is None:
346
+ residual = hidden_states
347
+ hidden_states = self.operator_norm(hidden_states)
348
+ else:
349
+ hidden_states, residual = self.operator_norm(hidden_states, residual)
350
+ hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
351
+ hidden_states, residual = self.ffn_norm(hidden_states, residual)
352
+ return self.feed_forward(hidden_states), residual
353
+
354
+
355
+class Lfm2MoeShortConvDecoderLayer(nn.Module):
356
+ def __init__(
357
+ self,
358
+ config: Lfm2MoeConfig,
359
+ layer_idx: int,
360
+ model_config: Optional[ModelConfig] = None,
361
+ cache_config: Optional[CacheConfig] = None,
362
+ quant_config: Optional[QuantizationConfig] = None,
363
+ prefix: str = "",
364
+ enable_eplb: bool = False,
365
+ ) -> None:
366
+ super().__init__()
367
+ self.layer_idx = layer_idx
368
+ self.conv = ShortConv(
369
+ config=config,
370
+ dim=config.hidden_size,
371
+ layer_idx=layer_idx,
372
+ model_config=model_config,
373
+ cache_config=cache_config,
374
+ prefix=f"{prefix}.conv",
375
+ )
376
+
377
+ if layer_idx < config.num_dense_layers:
378
+ self.feed_forward = Lfm2MoeMlp(
379
+ dim=config.hidden_size,
380
+ ff_dim=config.intermediate_size,
381
+ quant_config=quant_config,
382
+ prefix=f"{prefix}.feed_forward",
383
+ )
384
+ else:
385
+ self.feed_forward = Lfm2MoeSparseMoeBlock(
386
+ config=config,
387
+ quant_config=quant_config,
388
+ prefix=f"{prefix}.feed_forward",
389
+ enable_eplb=enable_eplb,
390
+ )
391
+
392
+ self.operator_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
393
+ self.ffn_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
394
+
395
+ def forward(
396
+ self,
397
+ hidden_states: torch.Tensor,
398
+ residual: Optional[torch.Tensor],
399
+ **kwargs,
400
+ ):
401
+ if residual is None:
402
+ residual = hidden_states
403
+ hidden_states = self.operator_norm(hidden_states)
404
+ else:
405
+ hidden_states, residual = self.operator_norm(hidden_states, residual)
406
+ output = torch.empty_like(hidden_states)
407
+ self.conv(
408
+ hidden_states,
409
+ output,
410
+ )
411
+ hidden_states, residual = self.ffn_norm(output, residual)
412
+ hidden_states = self.feed_forward(hidden_states)
413
+ return hidden_states, residual
414
+
415
+
416
+@support_torch_compile
417
+class Lfm2MoeModel(nn.Module):
418
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
419
+ super().__init__()
420
+
421
+ config = vllm_config.model_config.hf_config
422
+ model_config = vllm_config.model_config
423
+ cache_config = vllm_config.cache_config
424
+ quant_config = vllm_config.quant_config
425
+ lora_config = vllm_config.lora_config
426
+ parallel_config = vllm_config.parallel_config
427
+ enable_eplb = parallel_config.enable_eplb
428
+ eplb_config = parallel_config.eplb_config
429
+ self.num_redundant_experts = eplb_config.num_redundant_experts
430
+
431
+ self.config = config
432
+ lora_vocab = (
433
+ (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
434
+ if lora_config
435
+ else 0
436
+ )
437
+ self.vocab_size = config.vocab_size + lora_vocab
438
+ self.org_vocab_size = config.vocab_size
439
+
440
+ self.embed_tokens = VocabParallelEmbedding(
441
+ self.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size
442
+ )
443
+
444
+ def get_layer(prefix: str):
445
+ layer_idx = extract_layer_index(prefix)
446
+ is_attn = self.config.layer_types[layer_idx] == "full_attention"
447
+ layer_class = (
448
+ Lfm2MoeAttentionDecoderLayer
449
+ if is_attn
450
+ else Lfm2MoeShortConvDecoderLayer
451
+ )
452
+ return layer_class(
453
+ config,
454
+ layer_idx,
455
+ model_config,
456
+ cache_config,
457
+ quant_config=quant_config,
458
+ prefix=prefix,
459
+ enable_eplb=enable_eplb,
460
+ )
461
+
462
+ self.start_layer, self.end_layer, self.layers = make_layers(
463
+ config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
464
+ )
465
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
466
+ ["hidden_states", "residual"], config.hidden_size
467
+ )
468
+
469
+ if get_pp_group().is_last_rank:
470
+ self.embedding_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
471
+ else:
472
+ self.embedding_norm = PPMissingLayer()
473
+
474
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
475
+ return self.embed_tokens(input_ids)
476
+
477
+ def forward(
478
+ self,
479
+ input_ids: torch.Tensor,
480
+ positions: torch.Tensor,
481
+ intermediate_tensors: Optional[IntermediateTensors] = None,
482
+ inputs_embeds: Optional[torch.Tensor] = None,
483
+ ) -> torch.Tensor:
484
+ if get_pp_group().is_first_rank:
485
+ if inputs_embeds is not None:
486
+ hidden_states = inputs_embeds
487
+ else:
488
+ hidden_states = self.get_input_embeddings(input_ids)
489
+ residual = None
490
+ else:
491
+ assert intermediate_tensors is not None
492
+ hidden_states = intermediate_tensors["hidden_states"]
493
+ residual = intermediate_tensors["residual"]
494
+
495
+ for layer in self.layers[self.start_layer : self.end_layer]:
496
+ hidden_states, residual = layer(
497
+ positions=positions,
498
+ hidden_states=hidden_states,
499
+ residual=residual,
500
+ )
501
+ if not get_pp_group().is_last_rank:
502
+ return IntermediateTensors(
503
+ {"hidden_states": hidden_states, "residual": residual}
504
+ )
505
+ hidden_states, _ = self.embedding_norm(hidden_states, residual)
506
+ return hidden_states
507
+
508
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
509
+ return FusedMoE.make_expert_params_mapping(
510
+ ckpt_gate_proj_name="w1",
511
+ ckpt_down_proj_name="w2",
512
+ ckpt_up_proj_name="w3",
513
+ num_experts=self.config.num_experts,
514
+ num_redundant_experts=self.num_redundant_experts,
515
+ )
516
+
517
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
518
+ stacked_params_mapping = [
519
+ (".qkv_proj", ".q_proj", "q"),
520
+ (".qkv_proj", ".k_proj", "k"),
521
+ (".qkv_proj", ".v_proj", "v"),
522
+ (".w1", ".w1", 0),
523
+ (".w1", ".w3", 1),
524
+ ]
525
+ params_dict = dict(self.named_parameters())
526
+ loaded_params: set[str] = set()
527
+ expert_params_mapping = self.get_expert_mapping()
528
+ for name, loaded_weight in weights:
529
+ if "expert_bias" in name:
530
+ name = name.replace("expert_bias", "gate.e_score_correction_bias")
531
+
532
+ for param_name, weight_name, shard_id in stacked_params_mapping:
533
+ # Skip non-stacked layers and experts (experts handled below).
534
+ if weight_name not in name:
535
+ continue
536
+
537
+ if ("feed_forward.experts." in name) and name not in params_dict:
538
+ continue
539
+ name = name.replace(weight_name, param_name)
540
+ # Skip loading extra bias for GPTQ models.
541
+ if (
542
+ name.endswith(".bias") or name.endswith("_bias")
543
+ ) and name not in params_dict:
544
+ continue
545
+ # Skip layers on other devices.
546
+ if is_pp_missing_parameter(name, self):
547
+ continue
548
+
549
+ param = params_dict[name]
550
+ weight_loader = param.weight_loader
551
+ weight_loader(param, loaded_weight, shard_id)
552
+ break
553
+ else:
554
+ for mapping in expert_params_mapping:
555
+ param_name, weight_name, expert_id, shard_id = mapping
556
+
557
+ if weight_name not in name:
558
+ continue
559
+
560
+ name = name.replace(weight_name, param_name)
561
+ # Skip layers on other devices.
562
+ if is_pp_missing_parameter(name, self):
563
+ continue
564
+
565
+ # Skip loading extra bias for GPTQ models.
566
+ if (
567
+ name.endswith(".bias") or name.endswith("_bias")
568
+ ) and name not in params_dict:
569
+ continue
570
+ param = params_dict[name]
571
+
572
+ weight_loader = param.weight_loader
573
+ weight_loader(
574
+ param,
575
+ loaded_weight,
576
+ name,
577
+ shard_id=shard_id,
578
+ expert_id=expert_id,
579
+ )
580
+ break
581
+ else:
582
+ # Skip loading extra bias for GPTQ models.
583
+ if (
584
+ name.endswith(".bias") or name.endswith("_bias")
585
+ ) and name not in params_dict:
586
+ continue
587
+ # Skip layers on other devices.
588
+ if is_pp_missing_parameter(name, self):
589
+ continue
590
+ param = params_dict[name]
591
+ weight_loader = getattr(
592
+ param, "weight_loader", default_weight_loader
593
+ )
594
+ weight_loader(param, loaded_weight)
595
+ loaded_params.add(name)
596
+ return loaded_params
597
+
598
+
599
+class Lfm2MoeForCausalLM(
600
+ nn.Module,
601
+ HasInnerState,
602
+ SupportsLoRA,
603
+ SupportsPP,
604
+ IsHybrid,
605
+ SupportsQuant,
606
+ MixtureOfExperts,
607
+):
608
+ packed_modules_mapping = {
609
+ "qkv_proj": [
610
+ "q_proj",
611
+ "k_proj",
612
+ "v_proj",
613
+ ],
614
+ "w1": [
615
+ "w1",
616
+ "w3",
617
+ ],
618
+ }
619
+
620
+ # LoRA specific attributes
621
+ embedding_modules = {
622
+ "embed_tokens": "input_embeddings",
623
+ "lm_head": "output_embeddings",
624
+ }
625
+ embedding_padding_modules = ["lm_head"]
626
+
627
+ @classmethod
628
+ def get_mamba_state_dtype_from_config(
629
+ cls,
630
+ vllm_config: "VllmConfig",
631
+ ) -> tuple[torch.dtype, ...]:
632
+ return MambaStateDtypeCalculator.short_conv_state_dtype(
633
+ vllm_config.model_config.dtype,
634
+ vllm_config.cache_config.mamba_cache_dtype,
635
+ )
636
+
637
+ @classmethod
638
+ def get_mamba_state_shape_from_config(
639
+ cls,
640
+ vllm_config: "VllmConfig",
641
+ ) -> tuple[tuple[int, int]]:
642
+ """Calculate shapes for LFM2's convolutional cache.
643
+
644
+ Args:
645
+ vllm_config: vLLM config
646
+
647
+ Returns:
648
+ Tuple containing:
649
+ - conv_state_shape: Shape for convolutional state cache
650
+ """
651
+ parallel_config = vllm_config.parallel_config
652
+ hf_config = vllm_config.model_config.hf_config
653
+
654
+ return MambaStateShapeCalculator.short_conv_state_shape(
655
+ tp_world_size=parallel_config.tensor_parallel_size,
656
+ intermediate_size=hf_config.hidden_size,
657
+ conv_kernel=hf_config.conv_L_cache,
658
+ )
659
+
660
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
661
+ config = vllm_config.model_config.hf_config
662
+ quant_config = vllm_config.quant_config
663
+ cache_config = vllm_config.cache_config
664
+ lora_config = vllm_config.lora_config
665
+ assert not cache_config.enable_prefix_caching, (
666
+ "Lfm2Moe currently does not support prefix caching"
667
+ )
668
+
669
+ super().__init__()
670
+ self.config = config
671
+ self.model = Lfm2MoeModel(
672
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
673
+ )
674
+
675
+ if get_pp_group().is_last_rank:
676
+ self.unpadded_vocab_size = self.config.vocab_size
677
+ if lora_config:
678
+ self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
679
+
680
+ self.lm_head = ParallelLMHead(
681
+ self.unpadded_vocab_size,
682
+ config.hidden_size,
683
+ org_num_embeddings=config.vocab_size,
684
+ padding_size=(
685
+ DEFAULT_VOCAB_PADDING_SIZE
686
+ # We need bigger padding if using lora for kernel
687
+ # compatibility
688
+ if not lora_config
689
+ else lora_config.lora_vocab_padding_size
690
+ ),
691
+ quant_config=quant_config,
692
+ prefix=maybe_prefix(prefix, "lm_head"),
693
+ )
694
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
695
+ else:
696
+ self.lm_head = PPMissingLayer()
697
+
698
+ self.logits_processor = LogitsProcessor(
699
+ self.unpadded_vocab_size, config.vocab_size
700
+ )
701
+
702
+ self.make_empty_intermediate_tensors = (
703
+ self.model.make_empty_intermediate_tensors
704
+ )
705
+
706
+ # Set MoE hyperparameters
707
+ self.expert_weights = []
708
+
709
+ self.moe_layers: list[FusedMoE] = []
710
+ example_layer = None
711
+ for layer in self.model.layers:
712
+ if isinstance(layer, PPMissingLayer):
713
+ continue
714
+
715
+ assert isinstance(
716
+ layer, (Lfm2MoeAttentionDecoderLayer, Lfm2MoeShortConvDecoderLayer)
717
+ )
718
+ if isinstance(layer.feed_forward, Lfm2MoeSparseMoeBlock):
719
+ example_layer = layer.feed_forward
720
+ self.moe_layers.append(layer.feed_forward.experts)
721
+
722
+ if example_layer is None:
723
+ raise RuntimeError(
724
+ "No Lfm2MoeSparseMoeBlock layer found in the model.layers."
725
+ )
726
+
727
+ self.num_moe_layers = len(self.moe_layers)
728
+ self.num_expert_groups = 1
729
+ self.num_shared_experts = 0
730
+ self.num_logical_experts = example_layer.n_logical_experts
731
+ self.num_physical_experts = example_layer.n_physical_experts
732
+ self.num_local_physical_experts = example_layer.n_local_physical_experts
733
+ self.num_routed_experts = example_layer.n_routed_experts
734
+ self.num_redundant_experts = example_layer.n_redundant_experts
735
+
736
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
737
+ return self.model.get_input_embeddings(input_ids)
738
+
739
+ def set_eplb_state(
740
+ self,
741
+ expert_load_view: torch.Tensor,
742
+ logical_to_physical_map: torch.Tensor,
743
+ logical_replica_count: torch.Tensor,
744
+ ) -> None:
745
+ for layer_idx, layer in enumerate(self.moe_layers):
746
+ # Register the expert weights.
747
+ self.expert_weights.append(layer.get_expert_weights())
748
+ layer.set_eplb_state(
749
+ moe_layer_idx=layer_idx,
750
+ expert_load_view=expert_load_view,
751
+ logical_to_physical_map=logical_to_physical_map,
752
+ logical_replica_count=logical_replica_count,
753
+ )
754
+
755
+ def update_physical_experts_metadata(
756
+ self,
757
+ num_physical_experts: int,
758
+ num_local_physical_experts: int,
759
+ ) -> None:
760
+ assert self.num_local_physical_experts == num_local_physical_experts
761
+ self.num_physical_experts = num_physical_experts
762
+ self.num_local_physical_experts = num_local_physical_experts
763
+ self.num_redundant_experts = num_physical_experts - self.num_logical_experts
764
+ for layer in self.model.layers:
765
+ if isinstance(layer.feed_forward, Lfm2MoeSparseMoeBlock):
766
+ moe = layer.feed_forward
767
+ moe.n_local_physical_experts = num_local_physical_experts
768
+ moe.n_physical_experts = num_physical_experts
769
+ moe.n_redundant_experts = self.num_redundant_experts
770
+ moe.experts.update_expert_map()
771
+
772
+ def forward(
773
+ self,
774
+ input_ids: torch.Tensor,
775
+ positions: torch.Tensor,
776
+ intermediate_tensors: Optional[IntermediateTensors] = None,
777
+ inputs_embeds: Optional[torch.Tensor] = None,
778
+ **kwargs,
779
+ ) -> torch.Tensor:
780
+ hidden_states = self.model(
781
+ input_ids, positions, intermediate_tensors, inputs_embeds
782
+ )
783
+ return hidden_states
784
+
785
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
786
+ logits = self.logits_processor(self.lm_head, hidden_states)
787
+ return logits
788
+
789
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
790
+ loader = AutoWeightsLoader(
791
+ self,
792
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
793
+ )
794
+ return loader.load_weights(weights)
795
+
796
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
797
+ return self.model.get_expert_mapping()
vllm/model_executor/models/registry.py
+1
@@ -119,6 +119,7 @@ _TEXT_GENERATION_MODELS = {
119
"JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
120
"JambaForCausalLM": ("jamba", "JambaForCausalLM"),
121
"Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
122
+ "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
123
"LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
124
"Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),
125
# For decapoda-research/llama-*
vllm/transformers_utils/config.py
+1
@@ -91,6 +91,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
91
step3_vl="Step3VLConfig",
92
step3_text="Step3TextConfig",
93
qwen3_next="Qwen3NextConfig",
94
+ lfm2_moe="Lfm2MoeConfig",
95
)
96
97
_CONFIG_ATTRS_MAPPING: dict[str, str] = {
vllm/transformers_utils/configs/__init__.py
+2
@@ -19,6 +19,7 @@ from vllm.transformers_utils.configs.eagle import EAGLEConfig
19
from vllm.transformers_utils.configs.falcon import RWConfig
20
from vllm.transformers_utils.configs.jais import JAISConfig
21
from vllm.transformers_utils.configs.kimi_vl import KimiVLConfig
22
+from vllm.transformers_utils.configs.lfm2_moe import Lfm2MoeConfig
23
from vllm.transformers_utils.configs.medusa import MedusaConfig
24
from vllm.transformers_utils.configs.midashenglm import MiDashengLMConfig
25
from vllm.transformers_utils.configs.mlp_speculator import MLPSpeculatorConfig
@@ -46,6 +47,7 @@ __all__ = [
47
"EAGLEConfig",
48
"RWConfig",
49
"JAISConfig",
50
+ "Lfm2MoeConfig",
51
"MedusaConfig",
52
"MiDashengLMConfig",
53
"MLPSpeculatorConfig",
vllm/transformers_utils/configs/lfm2_moe.py
new
+160
@@ -0,0 +1,160 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
+
5
+from transformers.configuration_utils import PretrainedConfig
6
+
7
+
8
+class Lfm2MoeConfig(PretrainedConfig):
9
+ r"""
10
+ This is the configuration class to store the configuration of a [`Lfm2MoeModel`]. It is used to instantiate a LFM2 Moe
11
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
12
+ defaults will yield a similar configuration to that of the LFM2-8B-A1B model.
13
+ e.g. [LiquidAI/LFM2-8B-A1B](https://huggingface.co/LiquidAI/LFM2-8B-A1B)
14
+
15
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
16
+ documentation from [`PretrainedConfig`] for more information.
17
+
18
+
19
+ Args:
20
+ vocab_size (`int`, *optional*, defaults to 65536):
21
+ Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
22
+ `inputs_ids` passed when calling [`Lfm2Model`]
23
+ hidden_size (`int`, *optional*, defaults to 2048):
24
+ Dimension of the hidden representations.
25
+ intermediate_size (`int`, *optional*, defaults to 7168):
26
+ Dimension of the MLP representations.
27
+ moe_intermediate_size (`int`, *optional*, defaults to 1792):
28
+ Intermediate size of the routed expert.
29
+ num_hidden_layers (`int`, *optional*, defaults to 32):
30
+ Number of hidden layers in the Transformer decoder.
31
+ pad_token_id (`int`, *optional*, defaults to 0):
32
+ Padding token id.
33
+ bos_token_id (`int`, *optional*, defaults to 1):
34
+ Beginning of stream token id.
35
+ eos_token_id (`int`, *optional*, defaults to 2):
36
+ End of stream token id.
37
+ tie_word_embeddings (`bool`, *optional*, defaults to `True`):
38
+ Whether to tie weight embeddings
39
+ rope_theta (`float`, *optional*, defaults to 1000000.0):
40
+ The base period of the RoPE embeddings.
41
+ max_position_embeddings (`int`, *optional*, defaults to 128000):
42
+ The maximum sequence length that this model might ever be used with.
43
+ use_cache (`bool`, *optional*, defaults to `True`):
44
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
45
+ relevant if `config.is_decoder=True`.
46
+ norm_eps (`float`, *optional*, defaults to 1e-05):
47
+ The epsilon used by the rms normalization layers.
48
+ num_attention_heads (`int`, *optional*, defaults to 32):
49
+ Number of attention heads for each attention layer in the Transformer decoder.
50
+ num_key_value_heads (`int`, *optional*, defaults to 8):
51
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
52
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
53
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
54
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
55
+ by meanpooling all the original heads within that group. For more details, check out [this
56
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
57
+ `num_attention_heads`.
58
+ conv_bias (`bool`, *optional*, defaults to `False`):
59
+ Whether to use bias in the conv layers.
60
+ conv_L_cache (`int`, *optional*, defaults to 3):
61
+ L_cache dim in the conv layers.
62
+ num_dense_layers (`int`, *optional*, defaults to 2):
63
+ Number of dense Lfm2MoeMLP layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head).
64
+ num_experts_per_tok (`int`, *optional*, defaults to 4):
65
+ Number of selected experts.
66
+ num_experts (`int`, *optional*, defaults to 32):
67
+ Number of routed experts.
68
+ use_expert_bias (`bool`, *optional*, defaults to `True`):
69
+ Whether to use the expert bias on the routing weights.
70
+ routed_scaling_factor (`float`, *optional*, defaults to 1.0):
71
+ Scaling factor for routed experts in MoE models.
72
+ norm_topk_prob (`bool`, *optional*, defaults to `True`):
73
+ Whether to normalize the topk probabilities.
74
+ layer_types (`Optional`, *optional*):
75
+ Type of each layers.
76
+
77
+ ```python
78
+ >>> from transformers import Lfm2MoeModel, Lfm2MoeConfig
79
+
80
+ >>> # Initializing a LFM2 Moe model
81
+ >>> configuration = Lfm2MoeConfig()
82
+
83
+ >>> # Initializing a model from the LFM2-8B-A1B style configuration
84
+ >>> model = Lfm2MoeModel(configuration)
85
+
86
+ >>> # Accessing the model configuration
87
+ >>> configuration = model.config
88
+ ```""" # noqa: E501
89
+
90
+ model_type = "lfm2_moe"
91
+ keys_to_ignore_at_inference = ["past_key_values"]
92
+
93
+ def __init__(
94
+ self,
95
+ vocab_size: int = 65536,
96
+ hidden_size: int = 2048,
97
+ intermediate_size: int = 7168,
98
+ moe_intermediate_size: int = 1792,
99
+ num_hidden_layers: int = 32,
100
+ pad_token_id: int = 0,
101
+ bos_token_id: int = 1,
102
+ eos_token_id: int = 2,
103
+ tie_word_embeddings: bool = True,
104
+ rope_theta: float = 1000000.0,
105
+ max_position_embeddings: int = 128_000,
106
+ use_cache: bool = True,
107
+ norm_eps: float = 0.00001,
108
+ num_attention_heads: int = 32,
109
+ num_key_value_heads: int = 8,
110
+ conv_bias: bool = False,
111
+ conv_L_cache: int = 3,
112
+ num_dense_layers: int = 2,
113
+ num_experts_per_tok: int = 4,
114
+ num_experts: int = 32,
115
+ use_expert_bias: bool = True,
116
+ routed_scaling_factor: float = 1.0,
117
+ norm_topk_prob: bool = True,
118
+ layer_types: Optional[list[str]] = None,
119
+ **kwargs,
120
+ ):
121
+ self.vocab_size = vocab_size
122
+ self.hidden_size = hidden_size
123
+ self.intermediate_size = intermediate_size
124
+ self.num_hidden_layers = num_hidden_layers
125
+ self.rope_theta = rope_theta
126
+ self.max_position_embeddings = max_position_embeddings
127
+ self.use_cache = use_cache
128
+ self.norm_eps = norm_eps
129
+
130
+ # attn operator config
131
+ self.num_attention_heads = num_attention_heads
132
+ self.num_key_value_heads = num_key_value_heads
133
+
134
+ # custom operator config
135
+ self.conv_bias = conv_bias
136
+ self.conv_L_cache = conv_L_cache
137
+
138
+ # moe config
139
+ self.num_dense_layers = num_dense_layers
140
+ self.moe_intermediate_size = moe_intermediate_size
141
+ self.num_experts_per_tok = num_experts_per_tok
142
+ self.num_experts = num_experts
143
+ self.use_expert_bias = use_expert_bias
144
+ self.routed_scaling_factor = routed_scaling_factor
145
+ self.norm_topk_prob = norm_topk_prob
146
+ self.layer_types = layer_types
147
+
148
+ tie_word_embeddings = kwargs.get(
149
+ "tie_embedding", tie_word_embeddings
150
+ ) # to fit original config keys
151
+ super().__init__(
152
+ pad_token_id=pad_token_id,
153
+ bos_token_id=bos_token_id,
154
+ eos_token_id=eos_token_id,
155
+ tie_word_embeddings=tie_word_embeddings,
156
+ **kwargs,
157
+ )
158
+
159
+
160
+__all__ = ["Lfm2MoeConfig"]