@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1317034379

Fork 0

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

[ROCm][FEAT] Fuse DeepSeek shared experts into AITER fused_moe ops (#24097)

Signed-off-by: chenjun <junchen2@amd.com> Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com> Co-authored-by: valarLip <103567126+valarLip@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

kliuae committed Oct 16, 2025 at 10:41 UTC 1317034379522f75b5876ff5b5136ef312c32b37

8 files changed +352 -88

tests/distributed/test_expert_placement.py

+3 -3

@@ -85,7 +85,7 @@ def test_expert_placement_various_sizes(expert_placement_strategy, world_size):

85 else:

86 expected_test_local = base_experts

88 - test_local_experts, test_expert_map = determine_expert_map(

88 + test_local_experts, test_expert_map, _ = determine_expert_map(

89 ep_size=test_ep_size,

90 ep_rank=ep_rank,

91 global_num_experts=test_global_experts,

@@ -116,7 +116,7 @@ def test_expert_placement_edge_cases(expert_placement_strategy, world_size):

116 """Test edge cases for round_robin expert placement."""

117

118 # Test case 1: ep_size = 1 (should return None for expert_map)

119 - local_num_experts, expert_map = determine_expert_map(

119 + local_num_experts, expert_map, _ = determine_expert_map(

120 ep_size=1,

121 ep_rank=0,

122 global_num_experts=8,

@@ -217,7 +217,7 @@ def test_determine_expert_map_comprehensive():

217 expected_local,

218 expected_map_pattern,

219 ) in test_cases:

220 - local_num_experts, expert_map = determine_expert_map(

220 + local_num_experts, expert_map, _ = determine_expert_map(

221 ep_size=ep_size,

222 ep_rank=ep_rank,

223 global_num_experts=global_num_experts,

tests/kernels/moe/test_moe_permute_unpermute.py

+1 -1

@@ -217,7 +217,7 @@ def test_moe_permute_unpermute(

217 expert_map = None

218 n_local_expert = n_expert

219 if ep_size != 1:

220 - n_local_expert, expert_map = determine_expert_map(ep_size, ep_rank, n_expert)

220 + n_local_expert, expert_map, _ = determine_expert_map(ep_size, ep_rank, n_expert)

221 expert_map = expert_map.cuda()

222 start_expert = n_local_expert * ep_rank

223 current_platform.seed_everything(0)

vllm/envs.py

@@ -113,6 +113,7 @@ if TYPE_CHECKING:

113 VLLM_ROCM_USE_TRITON_ROPE: bool = False

114 VLLM_ROCM_USE_AITER_FP8BMM: bool = True

115 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False

116 + VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = True

117 VLLM_ROCM_USE_SKINNY_GEMM: bool = True

118 VLLM_ROCM_FP8_PADDING: bool = True

119 VLLM_ROCM_MOE_PADDING: bool = True

@@ -914,6 +915,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

915 os.getenv("VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION", "False").lower()

916 in ("true", "1")

917 ),

918 + # Whether to use aiter fusion shared experts ops.

919 + # By default is enabled.

920 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS": lambda: (

921 + os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "True").lower()

922 + in ("true", "1")

923 + ),

924 # use rocm skinny gemms

925 "VLLM_ROCM_USE_SKINNY_GEMM": lambda: (

926 os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1")

vllm/model_executor/layers/fused_moe/layer.py

+105 -11

@@ -5,6 +5,7 @@ from abc import abstractmethod

5 from collections.abc import Callable, Iterable

6 from contextlib import nullcontext

7 from enum import Enum

8 +from functools import partial

9 from typing import Literal, get_args, overload

11 import torch

@@ -12,7 +13,7 @@ import torch.nn.functional as F

13 from torch.nn.parameter import UninitializedParameter

15 import vllm.envs as envs

15 -from vllm.config import get_current_vllm_config

16 +from vllm.config import VllmConfig, get_current_vllm_config

17 from vllm.config.parallel import ExpertPlacementStrategy

18 from vllm.distributed import (

19 get_dp_group,

@@ -39,6 +40,8 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

40 FusedMoEPrepareAndFinalize,

41 )

42 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

43 + init_aiter_topK_meta_data,

44 + is_rocm_aiter_fusion_shared_expert_enabled,

45 is_rocm_aiter_moe_enabled,

46 )

47 from vllm.model_executor.layers.fused_moe.routing_simulator import RoutingSimulator

@@ -87,7 +90,7 @@ else:

91 if is_rocm_aiter_moe_enabled():

92 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501

90 - rocm_aiter_grouped_topk as grouped_topk,

93 + rocm_aiter_grouped_topk as grouped_topk_aiter,

94 )

95 else:

96 from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk

@@ -634,6 +637,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

637 global_num_experts=global_num_experts,

638 zero_expert_num=zero_expert_num,

639 zero_expert_type=zero_expert_type,

640 + num_fused_shared_experts=layer.num_fused_shared_experts,

641 )

642

643 if self.rocm_aiter_moe_enabled:

@@ -860,7 +864,8 @@ def determine_expert_map(

864 ep_rank: int,

865 global_num_experts: int,

866 expert_placement_strategy: ExpertPlacementStrategy = "linear",

863 -) -> tuple[int, torch.Tensor | None]:

867 + num_fused_shared_experts: int = 0,

868 +) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:

869 """

870 Calculates how many experts should be assigned to each rank for EP and

871 creates a mapping from global to local expert index. Experts are

@@ -882,10 +887,16 @@ def determine_expert_map(

887 (global_num_experts,) mapping from global to local index.

888 Contains -1 for experts not assigned to the current rank.

889 Returns None if ep_size is 1.

890 + - expert_mask (Optional[torch.Tensor]): A tensor of shape

891 + (global_num_experts + num_fused_shared_experts + 1,)

892 + containing 1 for experts assigned to the current rank

893 + and 0 for sentinel.

894 + Returns None if ep_size is 1.

895 + Used only when AITER MOE is enabled.

896 """

897 assert ep_size > 0

898 if ep_size == 1:

888 - return (global_num_experts, None)

899 + return (global_num_experts, None, None)

900

901 # Distribute experts as evenly as possible to each rank.

902 base_experts = global_num_experts // ep_size

@@ -914,7 +925,26 @@ def determine_expert_map(

925 f"'{expert_placement_strategy}', expected one of "

926 f"{get_args(ExpertPlacementStrategy)}"

927 )

917 - return (local_num_experts, expert_map)

928 +

929 + expert_mask = None

930 + if is_rocm_aiter_moe_enabled():

931 + expert_mask = torch.ones(

932 + (global_num_experts + num_fused_shared_experts + 1,), dtype=torch.int32

933 + )

934 + expert_mask[-1] = 0

935 + expert_mask[:global_num_experts] = expert_map > -1

936 + expert_map = torch.cat(

937 + (

938 + expert_map,

939 + torch.tensor(

940 + [local_num_experts + i for i in range(num_fused_shared_experts)],

941 + dtype=torch.int32,

942 + ),

943 + ),

944 + dim=0,

945 + )

946 +

947 + return (local_num_experts, expert_map, expert_mask)

948

949

950 def get_compressed_expert_map(expert_map: torch.Tensor) -> str:

@@ -1040,6 +1070,7 @@ class FusedMoE(CustomOp):

1070 zero_expert_num: int | None = 0,

1071 zero_expert_type: str | None = None,

1072 expert_mapping: list[tuple[str, str, int, str]] | None = None,

1073 + n_shared_experts: int | None = None,

1074 ):

1075 super().__init__()

1076 if params_dtype is None:

@@ -1096,6 +1127,22 @@ class FusedMoE(CustomOp):

1127 self.logical_to_physical_map: torch.Tensor | None = None

1128 self.logical_replica_count: torch.Tensor | None = None

1129

1130 + # ROCm aiter shared experts fusion

1131 + self.num_fused_shared_experts = (

1132 + n_shared_experts

1133 + if n_shared_experts is not None

1134 + and is_rocm_aiter_fusion_shared_expert_enabled()

1135 + else 0

1136 + )

1137 + if (

1138 + not is_rocm_aiter_fusion_shared_expert_enabled()

1139 + and self.num_fused_shared_experts != 0

1140 + ):

1141 + raise ValueError(

1142 + "n_shared_experts is only supported on ROCm aiter when "

1143 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled"

1144 + )

1145 +

1146 # Determine expert maps

1147 if self.use_ep:

1148 if self.enable_eplb:

@@ -1129,14 +1176,16 @@ class FusedMoE(CustomOp):

1176 expert_placement_strategy = "linear"

1177

1178 self.expert_map: torch.Tensor | None

1132 - local_num_experts, expert_map = determine_expert_map(

1179 + local_num_experts, expert_map, expert_mask = determine_expert_map(

1180 ep_size=self.ep_size,

1181 ep_rank=self.ep_rank,

1182 global_num_experts=self.global_num_experts,

1183 expert_placement_strategy=expert_placement_strategy,

1184 + num_fused_shared_experts=self.num_fused_shared_experts,

1185 )

1186 self.local_num_experts = local_num_experts

1187 self.register_buffer("expert_map", expert_map)

1188 + self.register_buffer("expert_mask", expert_mask)

1189 logger.info_once(

1190 "[EP Rank %s/%s] Expert parallelism is enabled. Expert "

1191 "placement strategy: %s. Local/global"

@@ -1150,10 +1199,18 @@ class FusedMoE(CustomOp):

1199 get_compressed_expert_map(self.expert_map),

1200 )

1201 else:

1153 - self.local_num_experts, self.expert_map = (self.global_num_experts, None)

1202 + self.local_num_experts, self.expert_map, self.expert_mask = (

1203 + self.global_num_experts,

1204 + None,

1205 + None,

1206 + )

1207

1208 self.top_k = top_k

1209

1210 + self._init_aiter_shared_experts_topK_buffer(

1211 + vllm_config=vllm_config, dp_size=dp_size_

1212 + )

1213 +

1214 assert intermediate_size % self.tp_size == 0

1215 self.hidden_size = hidden_size

1216 self.intermediate_size_per_partition = intermediate_size // self.tp_size

@@ -1327,13 +1384,18 @@ class FusedMoE(CustomOp):

1384 # ep_size and ep_rank should already be updated

1385 assert self.expert_map is not None

1386 with self.expert_map.device:

1330 - local_num_experts, expert_map = determine_expert_map(

1387 + local_num_experts, expert_map, expert_mask = determine_expert_map(

1388 ep_size=self.ep_size,

1389 ep_rank=self.ep_rank,

1390 global_num_experts=self.global_num_experts,

1391 + num_fused_shared_experts=self.num_fused_shared_experts,

1392 )

1393 self.local_num_experts = local_num_experts

1394 self.register_buffer("expert_map", expert_map)

1395 + self.register_buffer("expert_mask", expert_mask)

1396 + self._init_aiter_shared_experts_topK_buffer(

1397 + vllm_config=get_current_vllm_config(), dp_size=get_dp_group().world_size

1398 + )

1399

1400 def _load_per_tensor_weight_scale(

1401 self,

@@ -1504,6 +1566,24 @@ class FusedMoE(CustomOp):

1566 return expert_id

1567 return self.expert_map[expert_id].item()

1568

1569 + def _init_aiter_shared_experts_topK_buffer(

1570 + self, vllm_config: VllmConfig, dp_size: int

1571 + ):

1572 + if is_rocm_aiter_fusion_shared_expert_enabled():

1573 + if self.num_fused_shared_experts > 0:

1574 + init_aiter_topK_meta_data(

1575 + n_routed_experts=self.global_num_experts,

1576 + n_shared_experts=self.num_fused_shared_experts,

1577 + top_k=self.top_k,

1578 + tp_rank=self.ep_rank if self.use_ep else self.tp_rank,

1579 + tp_size=self.ep_size if self.use_ep else self.tp_size,

1580 + shared_experts_score=1.0,

1581 + max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens

1582 + * dp_size,

1583 + is_EP=self.use_ep,

1584 + )

1585 + self.local_num_experts += self.num_fused_shared_experts

1586 +

1587 @overload

1588 def weight_loader(

1589 self,

@@ -1866,6 +1946,7 @@ class FusedMoE(CustomOp):

1946 global_num_experts: int | None = None,

1947 zero_expert_num: int | None = None,

1948 zero_expert_type: str | None = None,

1949 + num_fused_shared_experts: int = 0,

1950 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

1951 """

1952 Route the input hidden states to the top-k experts based on the

@@ -1900,7 +1981,16 @@ class FusedMoE(CustomOp):

1981 if use_grouped_topk:

1982 assert topk_group is not None

1983 assert num_expert_group is not None

1903 - topk_weights, topk_ids = grouped_topk(

1984 + if is_rocm_aiter_moe_enabled():

1985 + if not is_rocm_aiter_fusion_shared_expert_enabled():

1986 + assert num_fused_shared_experts == 0

1987 + grouped_topk_impl = partial(

1988 + grouped_topk_aiter,

1989 + num_fused_shared_experts=num_fused_shared_experts,

1990 + )

1991 + else:

1992 + grouped_topk_impl = grouped_topk

1993 + topk_weights, topk_ids = grouped_topk_impl(

1994 hidden_states=hidden_states,

1995 gating_output=router_logits,

1996 topk=top_k,

@@ -2119,7 +2209,9 @@ class FusedMoE(CustomOp):

2209 renormalize=self.renormalize,

2210 use_grouped_topk=self.use_grouped_topk,

2211 global_num_experts=self.global_num_experts,

2122 - expert_map=self.expert_map,

2212 + expert_map=self.expert_map

2213 + if not is_rocm_aiter_moe_enabled()

2214 + else self.expert_mask,

2215 topk_group=self.topk_group,

2216 num_expert_group=self.num_expert_group,

2217 custom_routing_function=self.custom_routing_function,

@@ -2244,7 +2336,9 @@ class FusedMoE(CustomOp):

2336 renormalize=self.renormalize,

2337 use_grouped_topk=self.use_grouped_topk,

2338 global_num_experts=self.global_num_experts,

2247 - expert_map=self.expert_map,

2339 + expert_map=self.expert_map

2340 + if not is_rocm_aiter_moe_enabled()

2341 + else self.expert_mask,

2342 topk_group=self.topk_group,

2343 num_expert_group=self.num_expert_group,

2344 custom_routing_function=self.custom_routing_function,

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+93 -6

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from enum import IntEnum

4 -from functools import cache

4 +from functools import cache, lru_cache

6 import torch

@@ -46,6 +46,69 @@ def is_rocm_aiter_moe_enabled() -> bool:

46 )

49 +@cache

50 +def is_rocm_aiter_fusion_shared_expert_enabled() -> bool:

51 + return (

52 + envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS and is_rocm_aiter_moe_enabled()

53 + )

54 +

55 +

56 +aiter_topK_meta_data = None

57 +

58 +

59 +@lru_cache(maxsize=1)

60 +def init_aiter_topK_meta_data(

61 + n_routed_experts: int,

62 + n_shared_experts: int,

63 + top_k: int,

64 + tp_rank: int,

65 + tp_size: int,

66 + shared_experts_score: float = 1.0,

67 + max_num_tokens: int = 32768,

68 + is_EP: bool = False,

69 +):

70 + global aiter_topK_meta_data

71 + fake_expertid = n_routed_experts + n_shared_experts

72 +

73 + # all layers reuse same buffer

74 + # This extra element when EP is enabled is used as a sentinel

75 + # to mask out shared expert processing for tokens not owned by

76 + # the current EP rank. This is necessary to avoid double-processing

77 + # of shared experts.

78 + total_topk_ids = torch.empty(

79 + (max_num_tokens, top_k + n_shared_experts + is_EP),

80 + dtype=torch.int32,

81 + device="cuda",

82 + )

83 + ns_topk_ids, s_topk_ids = total_topk_ids.split(

84 + [top_k, n_shared_experts + is_EP], dim=1

85 + )

86 + shared_expert_ids = [n_routed_experts + i for i in range(n_shared_experts + is_EP)]

87 + if is_EP:

88 + s_topk_ids_list = [

89 + [fake_expertid] * (n_shared_experts + is_EP)

90 + ] * max_num_tokens

91 + for i in range(tp_rank, max_num_tokens, tp_size):

92 + s_topk_ids_list[i] = shared_expert_ids

93 + else:

94 + s_topk_ids_list = [

95 + list(range(n_routed_experts, fake_expertid))

96 + ] * max_num_tokens

97 + s_topk_ids[:] = torch.tensor(s_topk_ids_list, dtype=torch.int32, device="cuda")

98 +

99 + total_topk_weights = torch.empty(

100 + (max_num_tokens, top_k + n_shared_experts + is_EP),

101 + dtype=torch.float32,

102 + device="cuda",

103 + )

104 + ns_topk_weights, s_topk_weights = total_topk_weights.split(

105 + [top_k, n_shared_experts + is_EP], dim=1

106 + )

107 + s_topk_weights.fill_(shared_experts_score)

108 + assert aiter_topK_meta_data is None, "AITER topK meta data is already initialized"

109 + aiter_topK_meta_data = (total_topk_weights, total_topk_ids)

110 +

111 +

112 def rocm_aiter_asm_moe_tkw1_impl(

113 hidden_states: torch.Tensor,

114 w1: torch.Tensor,

@@ -300,11 +363,33 @@ def rocm_aiter_grouped_topk(

363 scoring_func: str = "softmax",

364 routed_scaling_factor: float = 1.0,

365 e_score_correction_bias: torch.Tensor | None = None,

366 + num_fused_shared_experts: int = 0,

367 ) -> tuple[torch.Tensor, torch.Tensor]:

368 token = hidden_states.shape[0]

369 device = hidden_states.device

306 - topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device)

307 - topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device)

370 + if is_rocm_aiter_fusion_shared_expert_enabled() and num_fused_shared_experts > 0:

371 + assert aiter_topK_meta_data is not None, (

372 + "AITER topK meta data is not initialized. "

373 + "Please ensure that init_aiter_topK_meta_data "

374 + "is called before this function."

375 + )

376 + total_topk_weights, total_topk_ids = aiter_topK_meta_data

377 + assert total_topk_weights.shape[0] >= token, (

378 + f"AITER topK meta data support {total_topk_weights.shape[0]} "

379 + f"tokens which is determined by max_num_batched_tokens, "

380 + f"but got {token} tokens now."

381 + )

382 + total_topk_weights = total_topk_weights[:token]

383 + total_topk_ids = total_topk_ids[:token]

384 + topk_weights, _ = total_topk_weights.split(

385 + [topk, total_topk_weights.shape[1] - topk], dim=1

386 + )

387 + topk_ids, _ = total_topk_ids.split(

388 + [topk, total_topk_ids.shape[1] - topk], dim=1

389 + )

390 + else:

391 + topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device)

392 + topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device)

393

394 if e_score_correction_bias is not None:

395 torch.ops.vllm.rocm_aiter_biased_grouped_topk(

@@ -315,6 +400,7 @@ def rocm_aiter_grouped_topk(

400 num_expert_group,

401 topk_group,

402 renormalize,

403 + routed_scaling_factor=routed_scaling_factor,

404 )

405 else:

406 assert scoring_func == "softmax" or scoring_func == "sigmoid"

@@ -326,10 +412,11 @@ def rocm_aiter_grouped_topk(

412 topk_group,

413 renormalize,

414 scoring_func,

415 + routed_scaling_factor=routed_scaling_factor,

416 )

417

331 - if routed_scaling_factor != 1.0:

332 - topk_weights = topk_weights * routed_scaling_factor

418 + if is_rocm_aiter_fusion_shared_expert_enabled() and num_fused_shared_experts > 0:

419 + return total_topk_weights, total_topk_ids

420 return topk_weights, topk_ids

421

422

@@ -354,7 +441,7 @@ def rocm_aiter_fused_experts(

441 topk_weights = topk_weights.to(torch.float32)

442 topk_ids = topk_ids.to(torch.int32)

443

357 - expert_mask = (expert_map > -1).to(torch.int32) if expert_map is not None else None

444 + expert_mask = expert_map if expert_map is not None else None

445

446 # w8a8 per-channel quantization

447 if (

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -1056,6 +1056,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1056 routed_scaling_factor=routed_scaling_factor,

1057 e_score_correction_bias=e_score_correction_bias,

1058 indices_type=self.topk_indices_dtype,

1059 + num_fused_shared_experts=layer.num_fused_shared_experts,

1060 )

1061

1062 per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN

vllm/model_executor/layers/quantization/fp8.py

@@ -1169,6 +1169,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1169 global_num_experts=global_num_experts,

1170 zero_expert_num=zero_expert_num,

1171 zero_expert_type=zero_expert_type,

1172 + num_fused_shared_experts=layer.num_fused_shared_experts,

1173 )

1174

1175 #

vllm/model_executor/models/deepseek_v2.py

+141 -67

@@ -50,6 +50,10 @@ from vllm.logger import init_logger

50 from vllm.model_executor.layers.activation import SiluAndMul

51 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

52 from vllm.model_executor.layers.fused_moe import SharedFusedMoE

53 +from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

54 + is_rocm_aiter_fusion_shared_expert_enabled,

55 + is_rocm_aiter_moe_enabled,

56 +)

57 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm

58 from vllm.model_executor.layers.linear import (

59 ColumnParallelLinear,

@@ -203,7 +207,10 @@ class DeepseekV2MoE(nn.Module):

207 self.physical_expert_start + self.n_local_physical_experts

208 )

209

206 - if config.n_shared_experts is None:

210 + if (

211 + config.n_shared_experts is None

212 + or is_rocm_aiter_fusion_shared_expert_enabled()

213 + ):

214 self.shared_experts = None

215 else:

216 intermediate_size = config.moe_intermediate_size * config.n_shared_experts

@@ -233,11 +240,17 @@ class DeepseekV2MoE(nn.Module):

240 prefix=f"{prefix}.experts",

241 scoring_func=config.scoring_func,

242 # we do scaling outside, set factor to 1.0 to avoid double mul

236 - routed_scaling_factor=1.0,

243 + # aiter applies routed_scaling_factor internally

244 + routed_scaling_factor=1.0

245 + if not is_rocm_aiter_moe_enabled()

246 + else self.routed_scaling_factor,

247 e_score_correction_bias=self.gate.e_score_correction_bias,

248 enable_eplb=self.enable_eplb,

249 num_redundant_experts=self.n_redundant_experts,

250 is_sequence_parallel=self.is_sequence_parallel,

251 + n_shared_experts=config.n_shared_experts

252 + if is_rocm_aiter_fusion_shared_expert_enabled()

253 + else None,

254 )

255

256 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -258,16 +271,15 @@ class DeepseekV2MoE(nn.Module):

271 hidden_states=hidden_states, router_logits=router_logits

272 )

273

261 - if self.shared_experts is not None:

262 - shared_output, final_hidden_states = fused_moe_out

263 - else:

264 - shared_output = None

265 - final_hidden_states = fused_moe_out

274 + shared_output, final_hidden_states = fused_moe_out

275 + if self.shared_experts is None:

276 + assert shared_output is None

277

278 # Fix FP16 overflow

279 # See DeepseekV2DecoderLayer for more details.

280 if hidden_states.dtype != torch.float16:

270 - final_hidden_states *= self.routed_scaling_factor

281 + if not is_rocm_aiter_moe_enabled():

282 + final_hidden_states *= self.routed_scaling_factor

283 elif self.shared_experts is not None:

284 assert shared_output is not None

285 shared_output *= 1.0 / self.routed_scaling_factor

@@ -1316,7 +1328,12 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR

1328 ckpt_gate_proj_name="gate_proj",

1329 ckpt_down_proj_name="down_proj",

1330 ckpt_up_proj_name="up_proj",

1319 - num_experts=self.config.n_routed_experts,

1331 + num_experts=self.config.n_routed_experts

1332 + + (

1333 + self.config.n_shared_experts

1334 + if is_rocm_aiter_fusion_shared_expert_enabled()

1335 + else 0

1336 + ),

1337 num_redundant_experts=self.num_redundant_experts,

1338 )

1339

@@ -1330,6 +1347,11 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR

1347 if spec_layer is not None:

1348 continue # skip spec decode layers for main model

1349

1350 + is_fuse_shared_experts_layer = (

1351 + is_rocm_aiter_fusion_shared_expert_enabled()

1352 + and ("mlp.shared_experts" in name)

1353 + )

1354 +

1355 for param_name, weight_name, shard_id in stacked_params_mapping:

1356 # Skip non-stacked layers and experts (experts handled below).

1357 if weight_name not in name:

@@ -1342,6 +1364,8 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR

1364 # for mlp.experts[0].gate_gate_up_proj, which breaks load.

1365 if ("mlp.experts." in name) and name not in params_dict:

1366 continue

1367 + if is_fuse_shared_experts_layer:

1368 + continue

1369 name_mapped = name.replace(weight_name, param_name)

1370

1371 # QKV fusion is optional, fall back to normal

@@ -1366,65 +1390,115 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR

1390 break

1391 else:

1392 is_expert_weight = False

1369 - for mapping in expert_params_mapping:

1370 - param_name, weight_name, expert_id, shard_id = mapping

1371 - if weight_name not in name:

1372 - continue

1373 -

1374 - # Anyway, this is an expert weight and should not be

1375 - # attempted to load as other weights later

1376 - is_expert_weight = True

1377 -

1378 - # Do not modify `name` since the loop may continue here

1379 - # Instead, create a new variable

1380 - name_mapped = name.replace(weight_name, param_name)

1381 -

1382 - if is_pp_missing_parameter(name_mapped, self):

1383 - continue

1384 -

1385 - param = params_dict[name_mapped]

1386 - # We should ask the weight loader to return success or not

1387 - # here since otherwise we may skip experts with other

1388 - # available replicas.

1389 - weight_loader = typing.cast(

1390 - Callable[..., bool], param.weight_loader

1391 - )

1392 - success = weight_loader(

1393 - param,

1394 - loaded_weight,

1395 - name_mapped,

1396 - shard_id=shard_id,

1397 - expert_id=expert_id,

1398 - return_success=True,

1399 - )

1400 - if success:

1401 - name = name_mapped

1402 - break

1403 - else:

1404 - if is_expert_weight:

1405 - # We've checked that this is an expert weight

1406 - # However it's not mapped locally to this rank

1407 - # So we simply skip it

1408 - continue

1409 -

1410 - # Skip loading extra bias for GPTQ models.

1411 - if name.endswith(".bias") and name not in params_dict:

1412 - continue

1413 -

1414 - # Remapping the name of FP8 kv-scale.

1415 - name = maybe_remap_kv_scale_name(name, params_dict)

1416 - if name is None:

1417 - continue

1418 -

1419 - if is_pp_missing_parameter(name, self):

1420 - continue

1421 -

1422 - param = params_dict[name]

1423 - weight_loader = getattr(

1424 - param, "weight_loader", default_weight_loader

1393 +

1394 + # Special handling: when AITER fusion_shared_experts is enabled,

1395 + # checkpoints may provide a single widened shared_experts tensor

1396 + # without explicit expert indices

1397 + # (e.g. ...mlp.shared_experts.gate_proj.weight).

1398 + # For models with multiple shared experts, split that tensor

1399 + # evenly into per-shared-expert slices and load them into

1400 + # appended expert slots mlp.experts.{n_routed_experts + j}.*

1401 + # accordingly.

1402 + num_chunks = 1

1403 + if is_fuse_shared_experts_layer:

1404 + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1

1405 + # Determine split axis based on op type

1406 + # gate/up: ColumnParallel → split along dim 0

1407 + # down: RowParallel → split along dim 1

1408 + split_dim = 1 if "down_proj.weight" in name else 0

1409 + total = loaded_weight.shape[split_dim]

1410 + assert total % num_chunks == 0, (

1411 + f"Shared expert weight dim {total} "

1412 + f"not divisible by num_chunks {num_chunks}"

1413 )

1426 - weight_loader(param, loaded_weight)

1427 - loaded_params.add(name)

1414 + chunk_size = total // num_chunks

1415 +

1416 + for j in range(num_chunks):

1417 + chunk_name = name

1418 + weight_to_load = loaded_weight

1419 +

1420 + if is_fuse_shared_experts_layer:

1421 + if split_dim == 0:

1422 + weight_to_load = loaded_weight[

1423 + j * chunk_size : (j + 1) * chunk_size, :

1424 + ]

1425 + else:

1426 + weight_to_load = loaded_weight[

1427 + :, j * chunk_size : (j + 1) * chunk_size

1428 + ]

1429 + # Synthesize an expert-style name so expert mapping

1430 + # can route it

1431 + chunk_name = name.replace(

1432 + "mlp.shared_experts",

1433 + f"mlp.experts.{self.config.n_routed_experts + j}",

1434 + )

1435 +

1436 + # Use expert_params_mapping to locate the destination

1437 + # param and delegate to its expert-aware weight_loader

1438 + # with expert_id.

1439 + for mapping in expert_params_mapping:

1440 + param_name, weight_name, expert_id, shard_id = mapping

1441 + if weight_name not in chunk_name:

1442 + continue

1443 +

1444 + # Anyway, this is an expert weight and should not be

1445 + # attempted to load as other weights later

1446 + is_expert_weight = True

1447 +

1448 + # Do not modify `name` since the loop may continue here

1449 + # Instead, create a new variable

1450 + name_mapped = chunk_name.replace(weight_name, param_name)

1451 +

1452 + if is_pp_missing_parameter(name_mapped, self):

1453 + continue

1454 +

1455 + param = params_dict[name_mapped]

1456 + # We should ask the weight loader to return success or

1457 + # not here since otherwise we may skip experts with

1458 + # other available replicas.

1459 + weight_loader = typing.cast(

1460 + Callable[..., bool], param.weight_loader

1461 + )

1462 + success = weight_loader(

1463 + param,

1464 + weight_to_load,

1465 + name_mapped,

1466 + shard_id=shard_id,

1467 + expert_id=expert_id,

1468 + return_success=True,

1469 + )

1470 + if success:

1471 + if not is_fuse_shared_experts_layer:

1472 + name = name_mapped

1473 + else:

1474 + loaded_params.add(name_mapped)

1475 + break

1476 + else:

1477 + if is_expert_weight:

1478 + # We've checked that this is an expert weight

1479 + # However it's not mapped locally to this rank

1480 + # So we simply skip it

1481 + continue

1482 +

1483 + # Skip loading extra bias for GPTQ models.

1484 + if name.endswith(".bias") and name not in params_dict:

1485 + continue

1486 +

1487 + # Remapping the name of FP8 kv-scale.

1488 + name = maybe_remap_kv_scale_name(name, params_dict)

1489 + if name is None:

1490 + continue

1491 +

1492 + if is_pp_missing_parameter(name, self):

1493 + continue

1494 +

1495 + param = params_dict[name]

1496 + weight_loader = getattr(

1497 + param, "weight_loader", default_weight_loader

1498 + )

1499 + weight_loader(param, loaded_weight)

1500 + if not is_fuse_shared_experts_layer:

1501 + loaded_params.add(name)

1502

1503 return loaded_params

1504