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 13, 02:00 UTC
[Model] Add Afmoe architecture implementation (#28332)
Signed-off-by: Maziyar Panahi <maziyar.panahi@iscpif.fr> Signed-off-by: Pranav <veldurthipranav@gmail.com> Co-authored-by: Maziyar Panahi <maziyar.panahi@iscpif.fr>
Pranav committed
Nov 17, 2025 at 15:11 UTC
f77bce001a6261da0661f0192c8cddd1ca453750
7 files changed
+804
docs/models/supported_models.md
+1
@@ -351,6 +351,7 @@ th {
351
352
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
353
|--------------|--------|-------------------|----------------------|---------------------------|
354
+| `AfmoeForCausalLM` | Afmoe | TBA | ✅︎ | ✅︎ |
355
| `ApertusForCausalLM` | Apertus | `swiss-ai/Apertus-8B-2509`, `swiss-ai/Apertus-70B-Instruct-2509`, etc. | ✅︎ | ✅︎ |
356
| `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ |
357
| `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+4
@@ -173,6 +173,10 @@ class _HfExamplesInfo:
173
174
_TEXT_GENERATION_EXAMPLE_MODELS = {
175
# [Decoder-only]
176
+ "AfmoeForCausalLM": _HfExamplesInfo(
177
+ "arcee-ai/Trinity-Nano",
178
+ is_available_online=False,
179
+ ),
180
"ApertusForCausalLM": _HfExamplesInfo("swiss-ai/Apertus-8B-Instruct-2509"),
181
"AquilaModel": _HfExamplesInfo("BAAI/AquilaChat-7B", trust_remote_code=True),
182
"AquilaForCausalLM": _HfExamplesInfo("BAAI/AquilaChat2-7B", trust_remote_code=True),
vllm/model_executor/models/afmoe.py
new
+711
@@ -0,0 +1,711 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only AfMoE model compatible with HuggingFace weights."""
4
+
5
+import typing
6
+from collections.abc import Callable, Iterable
7
+from itertools import islice
8
+from typing import Any
9
+
10
+import torch
11
+from torch import nn
12
+
13
+from vllm.attention import Attention, AttentionType
14
+from vllm.compilation.decorators import support_torch_compile
15
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
16
+from vllm.distributed import (
17
+ get_ep_group,
18
+ get_pp_group,
19
+ get_tensor_model_parallel_world_size,
20
+)
21
+from vllm.logger import init_logger
22
+from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
23
+from vllm.model_executor.layers.layernorm import RMSNorm
24
+from vllm.model_executor.layers.linear import (
25
+ ColumnParallelLinear,
26
+ QKVParallelLinear,
27
+ RowParallelLinear,
28
+)
29
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
30
+from vllm.model_executor.layers.quantization import QuantizationConfig
31
+from vllm.model_executor.layers.rotary_embedding import get_rope
32
+from vllm.model_executor.layers.vocab_parallel_embedding import (
33
+ ParallelLMHead,
34
+ VocabParallelEmbedding,
35
+)
36
+from vllm.model_executor.model_loader.weight_utils import (
37
+ default_weight_loader,
38
+ maybe_remap_kv_scale_name,
39
+)
40
+from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
41
+from vllm.model_executor.models.llama import LlamaMLP as AfmoeMLP
42
+from vllm.model_executor.models.utils import (
43
+ AutoWeightsLoader,
44
+ PPMissingLayer,
45
+ WeightsMapper,
46
+ extract_layer_index,
47
+ is_pp_missing_parameter,
48
+ make_empty_intermediate_tensors_factory,
49
+ make_layers,
50
+ maybe_prefix,
51
+)
52
+from vllm.sequence import IntermediateTensors
53
+
54
+logger = init_logger(__name__)
55
+
56
+
57
+class AfmoeMoE(nn.Module):
58
+ def __init__(
59
+ self,
60
+ config, # AfmoeConfig
61
+ quant_config: QuantizationConfig | None = None,
62
+ prefix: str = "",
63
+ enable_eplb: bool = False,
64
+ ):
65
+ super().__init__()
66
+ self.tp_size = get_tensor_model_parallel_world_size()
67
+ self.route_scale = config.route_scale
68
+ self.score_func = config.score_func
69
+ self.route_norm = config.route_norm
70
+
71
+ self.ep_group = get_ep_group().device_group
72
+ self.ep_rank = self.ep_group.rank()
73
+ self.ep_size = self.ep_group.size()
74
+ self.n_routed_experts: int = config.num_experts
75
+ self.n_shared_experts: int = config.num_shared_experts
76
+
77
+ if config.hidden_act != "silu":
78
+ raise ValueError(
79
+ f"Unsupported activation: {config.hidden_act}. "
80
+ "Only silu is supported for now."
81
+ )
82
+
83
+ # Router gate
84
+ self.gate = nn.Linear(
85
+ config.hidden_size,
86
+ config.num_experts,
87
+ bias=False,
88
+ dtype=torch.float32,
89
+ )
90
+ self.expert_bias = nn.Parameter(
91
+ torch.empty(config.num_experts, dtype=torch.float32)
92
+ )
93
+
94
+ # Load balancing settings
95
+ vllm_config = get_current_vllm_config()
96
+ eplb_config = vllm_config.parallel_config.eplb_config
97
+ self.enable_eplb = enable_eplb
98
+
99
+ self.n_redundant_experts = eplb_config.num_redundant_experts
100
+ self.n_logical_experts = self.n_routed_experts
101
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
102
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
103
+
104
+ self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
105
+ self.physical_expert_end = (
106
+ self.physical_expert_start + self.n_local_physical_experts
107
+ )
108
+
109
+ self.shared_experts = None
110
+ # Shared experts
111
+ if config.num_shared_experts > 0:
112
+ intermediate_size = config.moe_intermediate_size * config.num_shared_experts
113
+ self.shared_experts = AfmoeMLP(
114
+ hidden_size=config.hidden_size,
115
+ intermediate_size=intermediate_size,
116
+ hidden_act=config.hidden_act,
117
+ quant_config=quant_config,
118
+ reduce_results=False,
119
+ prefix=f"{prefix}.shared_experts",
120
+ )
121
+
122
+ # Routed experts using SharedFusedMoE
123
+ self.experts = SharedFusedMoE(
124
+ shared_experts=self.shared_experts,
125
+ num_experts=config.num_experts,
126
+ top_k=config.num_experts_per_tok,
127
+ hidden_size=config.hidden_size,
128
+ intermediate_size=config.moe_intermediate_size,
129
+ reduce_results=False,
130
+ renormalize=self.route_norm if self.score_func == "sigmoid" else False,
131
+ quant_config=quant_config,
132
+ use_grouped_topk=True,
133
+ num_expert_group=config.n_group,
134
+ topk_group=config.topk_group,
135
+ prefix=f"{prefix}.experts",
136
+ scoring_func=self.score_func,
137
+ routed_scaling_factor=self.route_scale,
138
+ e_score_correction_bias=self.expert_bias,
139
+ enable_eplb=self.enable_eplb,
140
+ num_redundant_experts=self.n_redundant_experts,
141
+ )
142
+
143
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
144
+ num_tokens, hidden_dim = hidden_states.shape
145
+ hidden_states = hidden_states.view(-1, hidden_dim)
146
+
147
+ router_logits = self.gate(hidden_states.to(dtype=torch.float32))
148
+
149
+ fused_moe_out = self.experts(
150
+ hidden_states=hidden_states, router_logits=router_logits
151
+ )
152
+
153
+ if self.shared_experts is not None:
154
+ shared_output, final_hidden_states = fused_moe_out
155
+ final_hidden_states = final_hidden_states + shared_output
156
+ else:
157
+ final_hidden_states = fused_moe_out
158
+ if self.tp_size > 1:
159
+ final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
160
+ final_hidden_states
161
+ )
162
+
163
+ return final_hidden_states.view(num_tokens, hidden_dim)
164
+
165
+
166
+class AfmoeAttention(nn.Module):
167
+ def __init__(
168
+ self,
169
+ config, # AfmoeConfig
170
+ layer_idx: int,
171
+ hidden_size: int,
172
+ num_heads: int,
173
+ num_kv_heads: int,
174
+ rope_theta: float = 10000,
175
+ rope_scaling: dict[str, Any] | None = None,
176
+ max_position_embeddings: int = 131072,
177
+ head_dim: int | None = None,
178
+ rms_norm_eps: float = 1e-05,
179
+ cache_config: CacheConfig | None = None,
180
+ quant_config: QuantizationConfig | None = None,
181
+ prefix: str = "",
182
+ attn_type: str = AttentionType.DECODER,
183
+ ) -> None:
184
+ super().__init__()
185
+ self.layer_idx = layer_idx
186
+ self.hidden_size = hidden_size
187
+ tp_size = get_tensor_model_parallel_world_size()
188
+ self.total_num_heads = num_heads
189
+ assert self.total_num_heads % tp_size == 0
190
+ self.num_heads = self.total_num_heads // tp_size
191
+ self.total_num_kv_heads = num_kv_heads
192
+ if self.total_num_kv_heads >= tp_size:
193
+ # Number of KV heads is greater than TP size, so we partition
194
+ # the KV heads across multiple tensor parallel GPUs.
195
+ assert self.total_num_kv_heads % tp_size == 0
196
+ else:
197
+ # Number of KV heads is less than TP size, so we replicate
198
+ # the KV heads across multiple tensor parallel GPUs.
199
+ assert tp_size % self.total_num_kv_heads == 0
200
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
201
+ self.head_dim = head_dim or (hidden_size // self.total_num_heads)
202
+ self.q_size = self.num_heads * self.head_dim
203
+ self.kv_size = self.num_kv_heads * self.head_dim
204
+ self.scaling = self.head_dim**-0.5
205
+ self.rope_theta = rope_theta
206
+ self.max_position_embeddings = max_position_embeddings
207
+
208
+ # Check if this is a local attention layer
209
+ self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention"
210
+ self.sliding_window = config.sliding_window if self.is_local_attention else None
211
+
212
+ self.qkv_proj = QKVParallelLinear(
213
+ self.hidden_size,
214
+ self.head_dim,
215
+ self.total_num_heads,
216
+ self.total_num_kv_heads,
217
+ bias=False,
218
+ quant_config=quant_config,
219
+ prefix=f"{prefix}.qkv_proj",
220
+ )
221
+
222
+ self.o_proj = RowParallelLinear(
223
+ self.total_num_heads * self.head_dim,
224
+ self.hidden_size,
225
+ bias=False,
226
+ quant_config=quant_config,
227
+ prefix=f"{prefix}.o_proj",
228
+ )
229
+
230
+ # Gating projection
231
+ self.gate_proj = ColumnParallelLinear(
232
+ hidden_size,
233
+ self.total_num_heads * self.head_dim,
234
+ bias=False,
235
+ quant_config=quant_config,
236
+ prefix=f"{prefix}.gate_proj",
237
+ )
238
+
239
+ # Q/K normalization
240
+ self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
241
+ self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
242
+
243
+ # Only create rotary embeddings for local attention
244
+ if self.is_local_attention:
245
+ self.rotary_emb = get_rope(
246
+ self.head_dim,
247
+ rotary_dim=self.head_dim,
248
+ max_position=max_position_embeddings,
249
+ base=rope_theta,
250
+ rope_scaling=rope_scaling,
251
+ is_neox_style=True,
252
+ )
253
+ else:
254
+ self.rotary_emb = None
255
+
256
+ self.attn = Attention(
257
+ self.num_heads,
258
+ self.head_dim,
259
+ self.scaling,
260
+ num_kv_heads=self.num_kv_heads,
261
+ cache_config=cache_config,
262
+ quant_config=quant_config,
263
+ per_layer_sliding_window=self.sliding_window,
264
+ prefix=f"{prefix}.attn",
265
+ attn_type=attn_type,
266
+ )
267
+
268
+ def forward(
269
+ self,
270
+ positions: torch.Tensor,
271
+ hidden_states: torch.Tensor,
272
+ ) -> torch.Tensor:
273
+ qkv, _ = self.qkv_proj(hidden_states)
274
+ gate, _ = self.gate_proj(hidden_states)
275
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
276
+
277
+ # Apply Q/K normalization
278
+ q = self.q_norm(q.reshape(-1, self.num_heads, self.head_dim)).reshape(q.shape)
279
+ k = self.k_norm(k.reshape(-1, self.num_kv_heads, self.head_dim)).reshape(
280
+ k.shape
281
+ )
282
+
283
+ # Apply rotary embeddings only for local attention
284
+ if self.is_local_attention and self.rotary_emb is not None:
285
+ q, k = self.rotary_emb(positions, q, k)
286
+
287
+ attn_output = self.attn(q, k, v)
288
+
289
+ # Apply gating
290
+ attn_output = attn_output * torch.sigmoid(gate)
291
+ output, _ = self.o_proj(attn_output)
292
+ return output
293
+
294
+
295
+class AfmoeDecoderLayer(nn.Module):
296
+ def __init__(
297
+ self,
298
+ config, # AfmoeConfig
299
+ cache_config: CacheConfig | None = None,
300
+ quant_config: QuantizationConfig | None = None,
301
+ prefix: str = "",
302
+ enable_eplb: bool = False,
303
+ ) -> None:
304
+ super().__init__()
305
+ self.hidden_size = config.hidden_size
306
+ rope_theta = getattr(config, "rope_theta", 10000)
307
+ rope_scaling = getattr(config, "rope_scaling", None)
308
+ if rope_scaling is not None and getattr(
309
+ config, "original_max_position_embeddings", None
310
+ ):
311
+ rope_scaling["original_max_position_embeddings"] = (
312
+ config.original_max_position_embeddings
313
+ )
314
+ max_position_embeddings = getattr(config, "max_position_embeddings", 131072)
315
+
316
+ # DecoderLayers are created with `make_layers` which passes the prefix
317
+ # with the layer's index.
318
+ self.layer_idx = extract_layer_index(prefix)
319
+
320
+ self.self_attn = AfmoeAttention(
321
+ config=config,
322
+ layer_idx=self.layer_idx,
323
+ hidden_size=self.hidden_size,
324
+ num_heads=config.num_attention_heads,
325
+ num_kv_heads=config.num_key_value_heads,
326
+ rope_theta=rope_theta,
327
+ rope_scaling=rope_scaling,
328
+ max_position_embeddings=max_position_embeddings,
329
+ head_dim=config.head_dim,
330
+ rms_norm_eps=config.rms_norm_eps,
331
+ cache_config=cache_config,
332
+ quant_config=quant_config,
333
+ prefix=f"{prefix}.self_attn",
334
+ )
335
+
336
+ # MoE or dense FFN
337
+ self.moe_enabled = self.layer_idx >= config.num_dense_layers
338
+ if self.moe_enabled:
339
+ self.mlp = AfmoeMoE(
340
+ config=config,
341
+ quant_config=quant_config,
342
+ prefix=f"{prefix}.mlp",
343
+ enable_eplb=enable_eplb,
344
+ )
345
+ else:
346
+ self.mlp = AfmoeMLP(
347
+ hidden_size=config.hidden_size,
348
+ intermediate_size=config.intermediate_size,
349
+ hidden_act=config.hidden_act,
350
+ quant_config=quant_config,
351
+ prefix=f"{prefix}.mlp",
352
+ )
353
+
354
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
355
+ self.post_attention_layernorm = RMSNorm(
356
+ config.hidden_size, eps=config.rms_norm_eps
357
+ )
358
+ self.pre_mlp_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
359
+ self.post_mlp_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
360
+
361
+ def forward(
362
+ self,
363
+ positions: torch.Tensor,
364
+ hidden_states: torch.Tensor,
365
+ residual: torch.Tensor | None,
366
+ ) -> tuple[torch.Tensor, torch.Tensor]:
367
+ if residual is None:
368
+ residual = hidden_states
369
+ hidden_states = self.input_layernorm(hidden_states)
370
+ else:
371
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
372
+
373
+ hidden_states = self.self_attn(
374
+ positions=positions,
375
+ hidden_states=hidden_states,
376
+ )
377
+ hidden_states = self.post_attention_layernorm(hidden_states) # attn norm b
378
+
379
+ # Fully Connected
380
+ hidden_states, residual = self.pre_mlp_layernorm( # ffn norm a
381
+ hidden_states, residual
382
+ )
383
+ hidden_states = self.mlp(hidden_states)
384
+ hidden_states = self.post_mlp_layernorm(hidden_states) # ffn norm b
385
+
386
+ return hidden_states, residual
387
+
388
+
389
+@support_torch_compile(
390
+ dynamic_arg_dims={
391
+ "input_ids": 0,
392
+ "positions": -1,
393
+ "intermediate_tensors": 0,
394
+ "inputs_embeds": 0,
395
+ }
396
+)
397
+class AfmoeModel(nn.Module):
398
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
399
+ super().__init__()
400
+
401
+ config = vllm_config.model_config.hf_config
402
+ cache_config = vllm_config.cache_config
403
+ quant_config = vllm_config.quant_config
404
+ enable_eplb = vllm_config.parallel_config.enable_eplb
405
+ self.config = config
406
+
407
+ self.vocab_size = config.vocab_size
408
+ self.mup_enabled = config.mup_enabled
409
+
410
+ if get_pp_group().is_first_rank:
411
+ self.embed_tokens = VocabParallelEmbedding(
412
+ config.vocab_size, config.hidden_size, prefix=f"{prefix}.embed_tokens"
413
+ )
414
+ else:
415
+ self.embed_tokens = PPMissingLayer()
416
+
417
+ self.start_layer, self.end_layer, self.layers = make_layers(
418
+ config.num_hidden_layers,
419
+ lambda prefix: AfmoeDecoderLayer(
420
+ config=config,
421
+ cache_config=cache_config,
422
+ quant_config=quant_config,
423
+ prefix=prefix,
424
+ enable_eplb=enable_eplb,
425
+ ),
426
+ prefix=f"{prefix}.layers",
427
+ )
428
+
429
+ if get_pp_group().is_last_rank:
430
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
431
+ else:
432
+ self.norm = PPMissingLayer()
433
+
434
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
435
+ ["hidden_states", "residual"], config.hidden_size
436
+ )
437
+
438
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
439
+ return self.embed_tokens(input_ids)
440
+
441
+ def forward(
442
+ self,
443
+ input_ids: torch.Tensor,
444
+ positions: torch.Tensor,
445
+ intermediate_tensors: IntermediateTensors | None = None,
446
+ inputs_embeds: torch.Tensor | None = None,
447
+ ) -> torch.Tensor | IntermediateTensors:
448
+ if get_pp_group().is_first_rank:
449
+ if inputs_embeds is not None:
450
+ hidden_states = inputs_embeds
451
+ else:
452
+ hidden_states = self.embed_input_ids(input_ids)
453
+
454
+ # Apply muP input scaling if enabled
455
+ if self.mup_enabled:
456
+ hidden_states = hidden_states * (self.config.hidden_size**0.5)
457
+
458
+ residual = None
459
+ else:
460
+ assert intermediate_tensors is not None
461
+ hidden_states = intermediate_tensors["hidden_states"]
462
+ residual = intermediate_tensors["residual"]
463
+
464
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
465
+ hidden_states, residual = layer(positions, hidden_states, residual)
466
+
467
+ if not get_pp_group().is_last_rank:
468
+ return IntermediateTensors(
469
+ {"hidden_states": hidden_states, "residual": residual}
470
+ )
471
+
472
+ hidden_states, _ = self.norm(hidden_states, residual)
473
+ return hidden_states
474
+
475
+ def make_empty_intermediate_tensors(
476
+ self, batch_size: int, dtype: torch.dtype, device: torch.device
477
+ ) -> IntermediateTensors:
478
+ return IntermediateTensors(
479
+ {
480
+ "hidden_states": torch.zeros(
481
+ (batch_size, self.config.hidden_size), dtype=dtype, device=device
482
+ ),
483
+ "residual": torch.zeros(
484
+ (batch_size, self.config.hidden_size), dtype=dtype, device=device
485
+ ),
486
+ }
487
+ )
488
+
489
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
490
+ # Params for weights, fp8 weight scales, fp8 activation scales
491
+ # (param_name, weight_name, expert_id, shard_id)
492
+ return SharedFusedMoE.make_expert_params_mapping(
493
+ ckpt_gate_proj_name="gate_proj",
494
+ ckpt_down_proj_name="down_proj",
495
+ ckpt_up_proj_name="up_proj",
496
+ num_experts=self.config.num_experts,
497
+ )
498
+
499
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
500
+ stacked_params_mapping = [
501
+ # (param_name, shard_name, shard_id)
502
+ ("qkv_proj", "q_proj", "q"),
503
+ ("qkv_proj", "k_proj", "k"),
504
+ ("qkv_proj", "v_proj", "v"),
505
+ ("gate_up_proj", "gate_proj", 0),
506
+ ("gate_up_proj", "up_proj", 1),
507
+ ]
508
+
509
+ params_dict = dict(self.named_parameters())
510
+ loaded_params: set[str] = set()
511
+ expert_params_mapping = self.get_expert_mapping()
512
+
513
+ for name, loaded_weight in weights:
514
+ for param_name, weight_name, shard_id in stacked_params_mapping:
515
+ # Skip non-stacked layers and experts (experts handled below).
516
+ if (weight_name not in name) or ("self_attn.gate_proj" in name):
517
+ continue
518
+ # We have mlp.experts[0].gate_proj in the checkpoint.
519
+ # Since we handle the experts below in expert_params_mapping,
520
+ # we need to skip here BEFORE we update the name, otherwise
521
+ # name will be updated to mlp.experts[0].gate_up_proj, which
522
+ # will then be updated below in expert_params_mapping
523
+ # for mlp.experts[0].gate_gate_up_proj, which breaks load.
524
+ if ("mlp.experts." in name) and name not in params_dict:
525
+ continue
526
+
527
+ name = name.replace(weight_name, param_name)
528
+ # Skip loading extra bias for GPTQ models.
529
+ if name.endswith(".bias") and name not in params_dict:
530
+ continue
531
+ if is_pp_missing_parameter(name, self):
532
+ continue
533
+
534
+ param = params_dict[name]
535
+ weight_loader = param.weight_loader
536
+ weight_loader(param, loaded_weight, shard_id)
537
+ break
538
+ else:
539
+ is_expert_weight = False
540
+ for mapping in expert_params_mapping:
541
+ param_name, weight_name, expert_id, shard_id = mapping
542
+ if weight_name not in name:
543
+ continue
544
+
545
+ # Anyway, this is an expert weight and should not be
546
+ # attempted to load as other weights later
547
+ is_expert_weight = True
548
+
549
+ # Do not modify `name` since the loop may continue here
550
+ # Instead, create a new variable
551
+ name_mapped = name.replace(weight_name, param_name)
552
+
553
+ if is_pp_missing_parameter(name_mapped, self):
554
+ continue
555
+
556
+ param = params_dict[name_mapped]
557
+ # We should ask the weight loader to return success or not
558
+ # here since otherwise we may skip experts with other
559
+ # available replicas.
560
+ weight_loader = typing.cast(
561
+ Callable[..., bool], param.weight_loader
562
+ )
563
+ success = weight_loader(
564
+ param,
565
+ loaded_weight,
566
+ name_mapped,
567
+ shard_id=shard_id,
568
+ expert_id=expert_id,
569
+ return_success=True,
570
+ )
571
+ if success:
572
+ name = name_mapped
573
+ break
574
+ else:
575
+ if is_expert_weight:
576
+ # We've checked that this is an expert weight
577
+ # However it's not mapped locally to this rank
578
+ # So we simply skip it
579
+ continue
580
+
581
+ # Skip loading extra bias for GPTQ models.
582
+ if name.endswith(".bias") and name not in params_dict:
583
+ continue
584
+
585
+ # Remapping the name of FP8 kv-scale.
586
+ name = maybe_remap_kv_scale_name(name, params_dict)
587
+ if name is None:
588
+ continue
589
+
590
+ if is_pp_missing_parameter(name, self):
591
+ continue
592
+
593
+ param = params_dict[name]
594
+ weight_loader = getattr(
595
+ param, "weight_loader", default_weight_loader
596
+ )
597
+ weight_loader(param, loaded_weight)
598
+ loaded_params.add(name)
599
+
600
+ return loaded_params
601
+
602
+
603
+class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
604
+ packed_modules_mapping = {
605
+ "qkv_proj": [
606
+ "q_proj",
607
+ "k_proj",
608
+ "v_proj",
609
+ ],
610
+ "gate_up_proj": [
611
+ "gate_proj",
612
+ "up_proj",
613
+ ],
614
+ }
615
+
616
+ hf_to_vllm_mapper = WeightsMapper(
617
+ orig_to_new_suffix={
618
+ ".router.gate.weight": ".gate.weight",
619
+ },
620
+ )
621
+
622
+ fall_back_to_pt_during_load = False
623
+
624
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
625
+ super().__init__()
626
+ config = vllm_config.model_config.hf_config
627
+ quant_config = vllm_config.quant_config
628
+ self.config = config
629
+ self.quant_config = quant_config
630
+ self.model = AfmoeModel(
631
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
632
+ )
633
+ if get_pp_group().is_last_rank:
634
+ self.lm_head = ParallelLMHead(
635
+ config.vocab_size, config.hidden_size, quant_config=quant_config
636
+ )
637
+ else:
638
+ self.lm_head = PPMissingLayer()
639
+ self.logits_processor = LogitsProcessor(config.vocab_size)
640
+ self.make_empty_intermediate_tensors = (
641
+ self.model.make_empty_intermediate_tensors
642
+ )
643
+ self.expert_weights = []
644
+
645
+ # Set MoE hyperparameters
646
+ self.num_moe_layers = config.num_hidden_layers - config.num_dense_layers
647
+ self.num_expert_groups = config.n_group
648
+
649
+ self.moe_layers: list[SharedFusedMoE] = []
650
+ example_moe = None
651
+ for layer in self.model.layers:
652
+ if isinstance(layer, PPMissingLayer):
653
+ continue
654
+
655
+ assert isinstance(layer, AfmoeDecoderLayer)
656
+ if layer.moe_enabled:
657
+ example_moe = layer.mlp
658
+ self.moe_layers.append(layer.mlp.experts)
659
+
660
+ if example_moe is None and self.num_moe_layers > 0:
661
+ raise RuntimeError("No AfmoeMoE layer found in model.layers.")
662
+
663
+ if example_moe is not None:
664
+ self.num_logical_experts = example_moe.n_logical_experts
665
+ self.num_physical_experts = example_moe.n_physical_experts
666
+ self.num_local_physical_experts = example_moe.n_local_physical_experts
667
+ self.num_routed_experts = example_moe.n_routed_experts
668
+ self.num_shared_experts = example_moe.n_shared_experts
669
+ self.num_redundant_experts = example_moe.n_redundant_experts
670
+
671
+ def set_eplb_state(
672
+ self,
673
+ expert_load_view: torch.Tensor,
674
+ logical_to_physical_map: torch.Tensor,
675
+ logical_replica_count: torch.Tensor,
676
+ ) -> None:
677
+ for layer_idx, layer in enumerate(self.moe_layers):
678
+ # Register the expert weights.
679
+ self.expert_weights.append(layer.get_expert_weights())
680
+ layer.set_eplb_state(
681
+ moe_layer_idx=layer_idx,
682
+ expert_load_view=expert_load_view,
683
+ logical_to_physical_map=logical_to_physical_map,
684
+ logical_replica_count=logical_replica_count,
685
+ )
686
+
687
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
688
+ return self.model.embed_input_ids(input_ids)
689
+
690
+ def forward(
691
+ self,
692
+ input_ids: torch.Tensor,
693
+ positions: torch.Tensor,
694
+ intermediate_tensors: IntermediateTensors | None = None,
695
+ inputs_embeds: torch.Tensor | None = None,
696
+ ) -> torch.Tensor | IntermediateTensors:
697
+ hidden_states = self.model(
698
+ input_ids, positions, intermediate_tensors, inputs_embeds
699
+ )
700
+ return hidden_states
701
+
702
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
703
+ logits = self.logits_processor(self.lm_head, hidden_states)
704
+ return logits
705
+
706
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
707
+ loader = AutoWeightsLoader(self)
708
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
709
+
710
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
711
+ return self.model.get_expert_mapping()
vllm/model_executor/models/registry.py
+1
@@ -56,6 +56,7 @@ logger = init_logger(__name__)
56
57
_TEXT_GENERATION_MODELS = {
58
# [Decoder-only]
59
+ "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"),
60
"ApertusForCausalLM": ("apertus", "ApertusForCausalLM"),
61
"AquilaModel": ("llama", "LlamaForCausalLM"),
62
"AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2
vllm/transformers_utils/config.py
+1
@@ -77,6 +77,7 @@ class LazyConfigDict(dict):
77
78
79
_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
80
+ afmoe="AfmoeConfig",
81
chatglm="ChatGLMConfig",
82
deepseek_vl_v2="DeepseekVLV2Config",
83
deepseek_v32=DeepseekV3Config,
vllm/transformers_utils/configs/__init__.py
+2
@@ -7,6 +7,7 @@ Model configs may be defined in this directory for the following reasons:
7
- There is a need to override the existing config to support vLLM.
8
"""
9
10
+from vllm.transformers_utils.configs.afmoe import AfmoeConfig
11
from vllm.transformers_utils.configs.chatglm import ChatGLMConfig
12
from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
13
from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig
@@ -40,6 +41,7 @@ from vllm.transformers_utils.configs.step3_vl import (
41
from vllm.transformers_utils.configs.ultravox import UltravoxConfig
42
43
__all__ = [
44
+ "AfmoeConfig",
45
"ChatGLMConfig",
46
"DeepseekVLV2Config",
47
"DotsOCRConfig",
vllm/transformers_utils/configs/afmoe.py
new
+84
@@ -0,0 +1,84 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from transformers.configuration_utils import PretrainedConfig
5
+
6
+
7
+class AfmoeConfig(PretrainedConfig):
8
+ model_type = "afmoe"
9
+
10
+ def __init__(
11
+ self,
12
+ vocab_size: int = 200_192,
13
+ hidden_size: int = 2048,
14
+ intermediate_size: int = 6144,
15
+ moe_intermediate_size: int = 1408,
16
+ num_hidden_layers: int = 32,
17
+ num_dense_layers: int = 1,
18
+ num_attention_heads: int = 16,
19
+ num_key_value_heads: int | None = None,
20
+ head_dim: int = 128,
21
+ hidden_act: str = "silu",
22
+ max_position_embeddings: int = 131072,
23
+ initializer_range: float = 0.02,
24
+ rms_norm_eps: float = 1e-5,
25
+ use_cache: bool = True,
26
+ tie_word_embeddings: bool = False,
27
+ rope_theta: float = 10000.0,
28
+ rope_scaling: dict | None = None,
29
+ num_experts: int = 64,
30
+ num_experts_per_tok: int = 6,
31
+ num_shared_experts: int = 2,
32
+ num_expert_groups: int = 1,
33
+ num_limited_groups: int = 1,
34
+ score_func: str = "sigmoid",
35
+ route_norm: bool = True,
36
+ route_scale: float = 1.0,
37
+ global_attn_every_n_layers: int = 4,
38
+ sliding_window: int = 2048,
39
+ layer_types: list[str] | None = None,
40
+ attention_dropout: float = 0.0,
41
+ mup_enabled: bool = False,
42
+ n_group: int = 1,
43
+ topk_group: int = 1,
44
+ **kwargs,
45
+ ):
46
+ self.vocab_size = vocab_size
47
+ self.hidden_size = hidden_size
48
+ self.intermediate_size = intermediate_size
49
+ self.num_hidden_layers = num_hidden_layers
50
+ self.num_dense_layers = num_dense_layers
51
+ self.num_attention_heads = num_attention_heads
52
+ self.num_key_value_heads = num_key_value_heads or num_attention_heads
53
+ self.head_dim = head_dim
54
+ self.hidden_act = hidden_act
55
+ self.max_position_embeddings = max_position_embeddings
56
+ self.initializer_range = initializer_range
57
+ self.rms_norm_eps = rms_norm_eps
58
+ self.use_cache = use_cache
59
+ self.rope_theta = rope_theta
60
+ self.rope_scaling = rope_scaling
61
+
62
+ self.moe_intermediate_size = moe_intermediate_size
63
+ self.num_experts = num_experts
64
+ self.num_experts_per_tok = num_experts_per_tok
65
+ self.num_shared_experts = num_shared_experts
66
+ self.num_expert_groups = num_expert_groups
67
+ self.num_limited_groups = num_limited_groups
68
+ self.score_func = score_func
69
+ self.route_norm = route_norm
70
+ self.route_scale = route_scale
71
+
72
+ self.global_attn_every_n_layers = global_attn_every_n_layers
73
+ self.sliding_window = sliding_window
74
+ self.layer_types = layer_types
75
+ self.attention_dropout = attention_dropout
76
+
77
+ self.mup_enabled = mup_enabled
78
+ self.n_group = n_group
79
+ self.topk_group = topk_group
80
+
81
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
82
+
83
+
84
+__all__ = ["AfmoeConfig"]