[ROCm][Perf] Add Fused Shared Expert (FSE) support for GLM-4.5/6/7 (#44313)
Signed-off-by: Olga Miroshnichenko <olga.miroshnichenko@amd.com> Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com> Co-authored-by: Mehdi Ghanimifard <mghanimi@amd.com> Co-authored-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com>
Olga Miroshnichenko committed
Jun 28, 2026 at 10:04 UTC
c7ca0bccae667934c29c654544131cdab046adfd
2 files changed
+254
-105
vllm/model_executor/models/glm4_moe.py
+138
-63
@@ -32,6 +32,7 @@ import torch
32
from torch import nn
33
from transformers.models.glm4_moe import Glm4MoeConfig
34
35
+from vllm._aiter_ops import rocm_aiter_ops
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
38
from vllm.distributed import (
@@ -168,7 +169,16 @@ class Glm4MoE(nn.Module):
169
self.physical_expert_start + self.n_local_physical_experts
170
)
171
171
- if config.n_shared_experts is not None:
172
+ # AITER fused shared-expert (FSE) gate; mirrors the deepseek_v2.py
173
+ # pattern (see Glm4MoE / FusedMoE wiring there).
174
+ self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
175
+ self.is_fusion_moe_shared_experts_enabled = (
176
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
177
+ )
178
+
179
+ if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
180
+ self.shared_experts = None
181
+ else:
182
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
183
self.shared_experts = Glm4MoeMLP(
184
hidden_size=config.hidden_size,
@@ -178,8 +188,6 @@ class Glm4MoE(nn.Module):
188
reduce_results=False,
189
prefix=f"{prefix}.shared_experts",
190
)
181
- else:
182
- self.shared_experts = None
191
192
self.experts = FusedMoE(
193
shared_experts=self.shared_experts,
@@ -194,12 +202,18 @@ class Glm4MoE(nn.Module):
202
topk_group=config.topk_group,
203
prefix=f"{prefix}.experts",
204
scoring_func="sigmoid",
205
+ # aiter applies routed_scaling_factor internally; see deepseek_v2.py.
206
routed_scaling_factor=self.routed_scaling_factor,
198
- apply_routed_scale_to_output=True,
207
+ apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled,
208
e_score_correction_bias=self.gate.e_score_correction_bias,
209
enable_eplb=self.enable_eplb,
210
num_redundant_experts=self.n_redundant_experts,
211
router_logits_dtype=torch.float32,
212
+ n_shared_experts=(
213
+ config.n_shared_experts
214
+ if self.is_fusion_moe_shared_experts_enabled
215
+ else None
216
+ ),
217
)
218
219
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -469,15 +483,25 @@ class Glm4MoeModel(nn.Module):
483
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
484
# Params for weights, fp8 weight scales, fp8 activation scales
485
# (param_name, weight_name, expert_id, shard_id)
486
+ # FSE widens the mapping by n_shared_experts slots; see deepseek_v2.py.
487
+ num_experts = self.config.n_routed_experts
488
+ if (
489
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
490
+ and self.config.n_shared_experts
491
+ ):
492
+ num_experts += self.config.n_shared_experts
493
return fused_moe_make_expert_params_mapping(
494
self,
495
ckpt_gate_proj_name="gate_proj",
496
ckpt_down_proj_name="down_proj",
497
ckpt_up_proj_name="up_proj",
477
- num_experts=self.config.n_routed_experts,
498
+ num_experts=num_experts,
499
)
500
501
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
502
+ rocm_aiter_moe_shared_expert_enabled = (
503
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
504
+ )
505
stacked_params_mapping = [
506
# (param_name, shard_name, shard_id)
507
("qkv_proj", "q_proj", "q"),
@@ -494,6 +518,11 @@ class Glm4MoeModel(nn.Module):
518
spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
519
if spec_layer is not None:
520
continue
521
+
522
+ is_fusion_moe_shared_experts_layer = (
523
+ rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
524
+ )
525
+
526
for param_name, weight_name, shard_id in stacked_params_mapping:
527
# Skip non-stacked layers and experts (experts handled below).
528
if weight_name not in name:
@@ -506,6 +535,8 @@ class Glm4MoeModel(nn.Module):
535
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
536
if ("mlp.experts." in name) and name not in params_dict:
537
continue
538
+ if is_fusion_moe_shared_experts_layer:
539
+ continue
540
541
name = name.replace(weight_name, param_name)
542
# Skip loading extra bias for GPTQ models.
@@ -527,65 +558,109 @@ class Glm4MoeModel(nn.Module):
558
break
559
else:
560
is_expert_weight = False
530
- for mapping in expert_params_mapping:
531
- param_name, weight_name, expert_id, shard_id = mapping
532
- if weight_name not in name:
533
- continue
534
-
535
- # Anyway, this is an expert weight and should not be
536
- # attempted to load as other weights later
537
- is_expert_weight = True
538
-
539
- # Do not modify `name` since the loop may continue here
540
- # Instead, create a new variable
541
- name_mapped = name.replace(weight_name, param_name)
542
-
543
- if is_pp_missing_parameter(name_mapped, self):
544
- continue
545
-
546
- param = params_dict[name_mapped]
547
- # We should ask the weight loader to return success or not
548
- # here since otherwise we may skip experts with other
549
- # available replicas.
550
- weight_loader = typing.cast(
551
- Callable[..., bool], param.weight_loader
552
- )
553
- success = weight_loader(
554
- param,
555
- loaded_weight,
556
- name_mapped,
557
- shard_id=shard_id,
558
- expert_id=expert_id,
559
- return_success=True,
560
- )
561
- if success:
562
- name = name_mapped
563
- break
564
- else:
565
- if is_expert_weight:
566
- # We've checked that this is an expert weight
567
- # However it's not mapped locally to this rank
568
- # So we simply skip it
569
- continue
570
-
571
- # Skip loading extra bias for GPTQ models.
572
- if name.endswith(".bias") and name not in params_dict:
573
- continue
574
-
575
- # Remapping the name of FP8 kv-scale.
576
- name = maybe_remap_kv_scale_name(name, params_dict)
577
- if name is None:
578
- continue
579
-
580
- if is_pp_missing_parameter(name, self):
581
- continue
582
-
583
- param = params_dict[name]
584
- weight_loader = getattr(
585
- param, "weight_loader", default_weight_loader
561
+
562
+ # FSE: split a widened mlp.shared_experts tensor into
563
+ # n_shared_experts chunks; see deepseek_v2.py for details.
564
+ num_chunks = 1
565
+ split_dim = 0
566
+ chunk_size = 0
567
+ if is_fusion_moe_shared_experts_layer:
568
+ num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
569
+ split_dim = (
570
+ 1
571
+ if ("down_proj.weight" in name and loaded_weight.ndim > 1)
572
+ else 0
573
)
587
- weight_loader(param, loaded_weight)
588
- loaded_params.add(name)
574
+ total = loaded_weight.shape[split_dim]
575
+ if total % num_chunks != 0:
576
+ raise ValueError(
577
+ f"FSE shared-expert weight {name} has dim "
578
+ f"{total} along axis {split_dim} which is not "
579
+ f"divisible by n_shared_experts={num_chunks}."
580
+ )
581
+ chunk_size = total // num_chunks
582
+
583
+ for j in range(num_chunks):
584
+ chunk_name = name
585
+ weight_to_load = loaded_weight
586
+
587
+ if is_fusion_moe_shared_experts_layer:
588
+ chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
589
+ if loaded_weight.ndim == 1:
590
+ weight_to_load = loaded_weight[chunk_slice]
591
+ elif split_dim == 0:
592
+ weight_to_load = loaded_weight[chunk_slice, :]
593
+ else:
594
+ weight_to_load = loaded_weight[:, chunk_slice]
595
+ # Synthesize an expert-style name for expert mapping.
596
+ chunk_name = name.replace(
597
+ "mlp.shared_experts",
598
+ f"mlp.experts.{self.config.n_routed_experts + j}",
599
+ )
600
+
601
+ for mapping in expert_params_mapping:
602
+ param_name, weight_name, expert_id, shard_id = mapping
603
+ if weight_name not in chunk_name:
604
+ continue
605
+
606
+ # Anyway, this is an expert weight and should not be
607
+ # attempted to load as other weights later
608
+ is_expert_weight = True
609
+
610
+ # Do not modify `name` since the loop may continue here
611
+ # Instead, create a new variable
612
+ name_mapped = chunk_name.replace(weight_name, param_name)
613
+
614
+ if is_pp_missing_parameter(name_mapped, self):
615
+ continue
616
+
617
+ param = params_dict[name_mapped]
618
+ # We should ask the weight loader to return success
619
+ # or not here since otherwise we may skip experts
620
+ # with other available replicas.
621
+ weight_loader = typing.cast(
622
+ Callable[..., bool], param.weight_loader
623
+ )
624
+ success = weight_loader(
625
+ param,
626
+ weight_to_load,
627
+ name_mapped,
628
+ shard_id=shard_id,
629
+ expert_id=expert_id,
630
+ return_success=True,
631
+ )
632
+ if success:
633
+ if not is_fusion_moe_shared_experts_layer:
634
+ name = name_mapped
635
+ else:
636
+ loaded_params.add(name_mapped)
637
+ break
638
+ else:
639
+ if is_expert_weight:
640
+ # We've checked that this is an expert weight
641
+ # However it's not mapped locally to this rank
642
+ # So we simply skip it
643
+ continue
644
+
645
+ # Skip loading extra bias for GPTQ models.
646
+ if name.endswith(".bias") and name not in params_dict:
647
+ continue
648
+
649
+ # Remapping the name of FP8 kv-scale.
650
+ name = maybe_remap_kv_scale_name(name, params_dict)
651
+ if name is None:
652
+ continue
653
+
654
+ if is_pp_missing_parameter(name, self):
655
+ continue
656
+
657
+ param = params_dict[name]
658
+ weight_loader = getattr(
659
+ param, "weight_loader", default_weight_loader
660
+ )
661
+ weight_loader(param, loaded_weight)
662
+ if name is not None and not is_fusion_moe_shared_experts_layer:
663
+ loaded_params.add(name)
664
665
return loaded_params
666
vllm/model_executor/models/glm4_moe_mtp.py
+116
-42
@@ -24,12 +24,14 @@
24
"""Inference-only GLM-4.5, GLM-4.6, GLM-4.7 MTP
25
model compatible with HuggingFace weights."""
26
27
-from collections.abc import Iterable
27
+import typing
28
+from collections.abc import Callable, Iterable
29
30
import torch
31
import torch.nn as nn
32
from transformers import PretrainedConfig
33
34
+from vllm._aiter_ops import rocm_aiter_ops
35
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
36
from vllm.model_executor.layers.fused_moe import (
37
MoERunner,
@@ -237,6 +239,10 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
239
return self.model.compute_logits(hidden_states, spec_step_idx)
240
241
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
242
+ # FSE weight loading mirrors glm4_moe.py / deepseek_mtp.py.
243
+ rocm_aiter_moe_shared_expert_enabled = (
244
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
245
+ )
246
stacked_params_mapping = [
247
# (param_name, shard_name, shard_id)
248
("qkv_proj", "q_proj", "q"),
@@ -248,12 +254,15 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
254
255
# Params for weights, fp8 weight scales, fp8 activation scales
256
# (param_name, weight_name, expert_id, shard_id)
257
+ num_experts = self.config.n_routed_experts
258
+ if rocm_aiter_moe_shared_expert_enabled and self.config.n_shared_experts:
259
+ num_experts += self.config.n_shared_experts
260
expert_params_mapping = fused_moe_make_expert_params_mapping(
261
self,
262
ckpt_gate_proj_name="gate_proj",
263
ckpt_down_proj_name="down_proj",
264
ckpt_up_proj_name="up_proj",
256
- num_experts=self.config.n_routed_experts,
265
+ num_experts=num_experts,
266
)
267
268
params_dict = dict(self.named_parameters())
@@ -269,6 +278,11 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
278
if spec_layer is None:
279
continue
280
name = self._rewrite_spec_layer_name(spec_layer, name)
281
+
282
+ is_fusion_moe_shared_experts_layer = (
283
+ rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
284
+ )
285
+
286
for param_name, weight_name, shard_id in stacked_params_mapping:
287
# Skip non-stacked layers and experts (experts handled below).
288
if weight_name not in name:
@@ -281,6 +295,8 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
295
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
296
if ("mlp.experts." in name) and name not in params_dict:
297
continue
298
+ if is_fusion_moe_shared_experts_layer:
299
+ continue
300
name = name.replace(weight_name, param_name)
301
# Skip loading extra bias for GPTQ models.
302
if name.endswith(".bias") and name not in params_dict:
@@ -291,47 +307,105 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
307
weight_loader(param, loaded_weight, shard_id)
308
break
309
else:
294
- for mapping in expert_params_mapping:
295
- param_name, weight_name, expert_id, shard_id = mapping
296
- if weight_name not in name:
297
- continue
298
- name = name.replace(weight_name, param_name)
299
-
300
- param = params_dict[name]
301
- weight_loader = param.weight_loader
302
- weight_loader(
303
- param,
304
- loaded_weight,
305
- name,
306
- shard_id=shard_id,
307
- expert_id=expert_id,
308
- )
309
- break
310
- else:
311
- # Skip loading extra bias for GPTQ models.
312
- if name.endswith(".bias") and name not in params_dict:
313
- continue
314
- # Some checkpoints include weight scale tensors for the
315
- # LM head even when the quantized head isn't built. Skip
316
- # them if the model does not expose a matching parameter
317
- # to avoid KeyError during load.
318
- if name.endswith(".weight_scale") and name not in params_dict:
319
- continue
320
-
321
- # According to DeepSeek-V3 Technical Report, MTP modules
322
- # shares embedding layer. We only load the first weights.
323
- if (
324
- spec_layer != self.model.mtp_start_layer_idx
325
- and ".layers" not in name
326
- ):
327
- continue
328
-
329
- param = params_dict[name]
330
- weight_loader = getattr(
331
- param, "weight_loader", default_weight_loader
310
+ # FSE: split a widened mlp.shared_experts tensor into
311
+ # n_shared_experts chunks; see deepseek_v2.py for details.
312
+ num_chunks = 1
313
+ split_dim = 0
314
+ chunk_size = 0
315
+ if is_fusion_moe_shared_experts_layer:
316
+ num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
317
+ split_dim = (
318
+ 1
319
+ if ("down_proj.weight" in name and loaded_weight.ndim > 1)
320
+ else 0
321
)
333
- weight_loader(param, loaded_weight)
334
- loaded_params.add(name)
322
+ total = loaded_weight.shape[split_dim]
323
+ if total % num_chunks != 0:
324
+ raise ValueError(
325
+ f"FSE shared-expert weight {name} has dim "
326
+ f"{total} along axis {split_dim} which is "
327
+ f"not divisible by "
328
+ f"n_shared_experts={num_chunks}."
329
+ )
330
+ chunk_size = total // num_chunks
331
+
332
+ for j in range(num_chunks):
333
+ chunk_name = name
334
+ weight_to_load = loaded_weight
335
+
336
+ if is_fusion_moe_shared_experts_layer:
337
+ chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
338
+ if loaded_weight.ndim == 1:
339
+ weight_to_load = loaded_weight[chunk_slice]
340
+ elif split_dim == 0:
341
+ weight_to_load = loaded_weight[chunk_slice, :]
342
+ else:
343
+ weight_to_load = loaded_weight[:, chunk_slice]
344
+ chunk_name = name.replace(
345
+ "mlp.shared_experts",
346
+ f"mlp.experts.{self.config.n_routed_experts + j}",
347
+ )
348
+
349
+ is_expert_weight = False
350
+ for mapping in expert_params_mapping:
351
+ param_name, weight_name, expert_id, shard_id = mapping
352
+ if weight_name not in chunk_name:
353
+ continue
354
+
355
+ is_expert_weight = True
356
+ name_mapped = chunk_name.replace(weight_name, param_name)
357
+
358
+ param = params_dict[name_mapped]
359
+ # Use return_success so we don't blindly mark
360
+ # remote-expert replicas as loaded on this rank.
361
+ weight_loader = typing.cast(
362
+ Callable[..., bool], param.weight_loader
363
+ )
364
+ success = weight_loader(
365
+ param,
366
+ weight_to_load,
367
+ name_mapped,
368
+ shard_id=shard_id,
369
+ expert_id=expert_id,
370
+ return_success=True,
371
+ )
372
+ if success:
373
+ if not is_fusion_moe_shared_experts_layer:
374
+ name = name_mapped
375
+ else:
376
+ loaded_params.add(name_mapped)
377
+ break
378
+ else:
379
+ if is_expert_weight:
380
+ # Expert weight not local to this rank; skip.
381
+ continue
382
+
383
+ # Skip loading extra bias for GPTQ models.
384
+ if name.endswith(".bias") and name not in params_dict:
385
+ continue
386
+ # Some checkpoints include weight scale tensors for
387
+ # the LM head even when the quantized head isn't
388
+ # built. Skip them if the model does not expose a
389
+ # matching parameter to avoid KeyError during load.
390
+ if name.endswith(".weight_scale") and name not in params_dict:
391
+ continue
392
+
393
+ # According to DeepSeek-V3 Technical Report, MTP
394
+ # modules share the embedding layer. We only load
395
+ # the first weights.
396
+ if (
397
+ spec_layer != self.model.mtp_start_layer_idx
398
+ and ".layers" not in name
399
+ ):
400
+ continue
401
+
402
+ param = params_dict[name]
403
+ weight_loader = getattr(
404
+ param, "weight_loader", default_weight_loader
405
+ )
406
+ weight_loader(param, loaded_weight)
407
+ if not is_fusion_moe_shared_experts_layer:
408
+ loaded_params.add(name)
409
return loaded_params
410
411
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: