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

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 9, 19:30 UTC

[AMD][ROCm] MoRI EP: a high-performance all2all backend (#28664)

Signed-off-by: Alex Sun <alex.s@amd.com>

Alex Sun committed Jan 22, 2026 at 16:33 UTC 49a12622678906727dd9d1cb157d5e14ab24d91b

16 files changed +397 -9

tests/kernels/moe/modular_kernel_tools/cli_args.py

+1 -1

@@ -141,7 +141,7 @@ def make_config(args: argparse.Namespace) -> Config:

141

142 quant_config = None

143 if args.quant_dtype is not None:

144 - quant_config = FusedMoEQuantConfig(

144 + quant_config = FusedMoEQuantConfig.make(

145 quant_dtype=args.quant_dtype,

146 per_act_token_quant=args.per_token_quantized_activations,

147 per_out_ch_quant=args.per_channel_quantized_weights,

tests/kernels/moe/modular_kernel_tools/common.py

+19 -1

@@ -28,7 +28,13 @@ from vllm.model_executor.layers.fused_moe.config import (

28 FusedMoEQuantConfig,

29 RoutingMethodType,

30 )

31 -from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx

31 +from vllm.utils.import_utils import (

32 + has_aiter,

33 + has_deep_ep,

34 + has_deep_gemm,

35 + has_mori,

36 + has_pplx,

37 +)

39 from .mk_objects import (

40 TestMoEQuantConfig,

@@ -211,6 +217,14 @@ class Config:

217 or info.backend == "deepep_low_latency"

218 )

219

220 + def needs_aiter(self):

221 + info = expert_info(self.fused_experts_type)

222 + return info.needs_aiter

223 +

224 + def needs_mori(self):

225 + info = prepare_finalize_info(self.prepare_finalize_type)

226 + return info.backend == "mori"

227 +

228 def all2all_backend(self):

229 info = prepare_finalize_info(self.prepare_finalize_type)

230 return info.backend

@@ -278,6 +292,10 @@ class Config:

292 return False, "Needs DeepGEMM, but DeepGEMM not available."

293 if self.needs_pplx() and not has_pplx(): # noqa: SIM103

294 return False, "Needs PPLX, but PPLX not available."

295 + if self.needs_aiter() and not has_aiter(): # noqa: SIM103

296 + return False, "Needs Aiter, but Aiter not available."

297 + if self.needs_mori() and not has_mori(): # noqa: SIM103

298 + return False, "Needs MoRI, but MoRI not available."

299

300 return True, None

301

tests/kernels/moe/modular_kernel_tools/mk_objects.py

+48 -1

@@ -37,7 +37,13 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

37 from vllm.platforms import current_platform

38 from vllm.utils.deep_gemm import is_deep_gemm_supported

39 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

40 -from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx

40 +from vllm.utils.import_utils import (

41 + has_aiter,

42 + has_deep_ep,

43 + has_deep_gemm,

44 + has_mori,

45 + has_pplx,

46 +)

49 @dataclass

@@ -66,6 +72,7 @@ class ExpertInfo:

72 supports_expert_map: bool

73 needs_matching_quant: bool = False

74 needs_deep_gemm: bool = False

75 + needs_aiter: bool = False

78 PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize, PrepareFinalizeInfo] = {}

@@ -126,6 +133,7 @@ def register_experts(

133 supports_expert_map: bool,

134 needs_matching_quant: bool = False,

135 needs_deep_gemm: bool = False,

136 + needs_aiter: bool = False,

137 ):

138 global EXPERT_INFO

139 global MK_FUSED_EXPERT_TYPES

@@ -139,6 +147,7 @@ def register_experts(

147 supports_expert_map,

148 needs_matching_quant,

149 needs_deep_gemm,

150 + needs_aiter,

151 )

152

153 MK_FUSED_EXPERT_TYPES.append(kind)

@@ -218,6 +227,20 @@ if has_deep_ep() and not current_platform.has_device_capability(100):

227 backend="deepep_low_latency",

228 )

229

230 +if has_mori():

231 + from vllm.model_executor.layers.fused_moe.mori_prepare_finalize import (

232 + MoriPrepareAndFinalize,

233 + )

234 +

235 + register_prepare_and_finalize(

236 + MoriPrepareAndFinalize,

237 + standard_format,

238 + fp8_types,

239 + blocked_quantization_support=True,

240 + backend="mori",

241 + supports_apply_weight_on_input=False,

242 + )

243 +

244 if has_pplx():

245 from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (

246 PplxPrepareAndFinalize,

@@ -261,6 +284,25 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability

284 )

285 else:

286 FlashInferCutlassMoEPrepareAndFinalize = None

287 + FlashInferExperts = None

288 +

289 +

290 +if has_aiter():

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

292 + AiterExperts,

293 + )

294 +

295 + register_experts(

296 + AiterExperts,

297 + standard_format,

298 + fp8_types,

299 + blocked_quantization_support=True,

300 + supports_chunking=True,

301 + supports_expert_map=True,

302 + needs_aiter=True,

303 + )

304 +else:

305 + AiterExperts = None

306

307 if has_deep_gemm() and is_deep_gemm_supported():

308 register_experts(

@@ -316,6 +358,9 @@ if cutlass_fp8_supported():

358 supports_chunking=False,

359 supports_expert_map=False,

360 )

361 +else:

362 + CutlassBatchedExpertsFp8 = None

363 + CutlassExpertsFp8 = None

364

365 if cutlass_fp4_supported():

366 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp4

@@ -328,6 +373,8 @@ if cutlass_fp4_supported():

373 supports_chunking=True,

374 supports_expert_map=False,

375 )

376 +else:

377 + CutlassExpertsFp4 = None

378

379 MK_QUANT_CONFIGS: list[TestMoEQuantConfig | None] = [

380 None,

vllm/_aiter_ops.py

+12

@@ -79,6 +79,8 @@ def _rocm_aiter_fused_moe_impl(

79 w2_scale: torch.Tensor | None = None,

80 a1_scale: torch.Tensor | None = None,

81 a2_scale: torch.Tensor | None = None,

82 + num_local_tokens: torch.Tensor | None = None,

83 + output_dtype: torch.dtype | None = None,

84 ) -> torch.Tensor:

85 from aiter import ActivationType, QuantType

86 from aiter.fused_moe import fused_moe

@@ -100,6 +102,8 @@ def _rocm_aiter_fused_moe_impl(

102 w2_scale,

103 a1_scale,

104 a2_scale,

105 + num_local_tokens=num_local_tokens,

106 + dtype=output_dtype,

107 )

108

109

@@ -117,7 +121,11 @@ def _rocm_aiter_fused_moe_fake(

121 w2_scale: torch.Tensor | None = None,

122 a1_scale: torch.Tensor | None = None,

123 a2_scale: torch.Tensor | None = None,

124 + num_local_tokens: torch.Tensor | None = None,

125 + output_dtype: torch.dtype | None = None,

126 ) -> torch.Tensor:

127 + if output_dtype is not None:

128 + return torch.empty_like(hidden_states, dtype=output_dtype)

129 return torch.empty_like(hidden_states)

130

131

@@ -1236,6 +1244,8 @@ class rocm_aiter_ops:

1244 w2_scale: torch.Tensor | None = None,

1245 a1_scale: torch.Tensor | None = None,

1246 a2_scale: torch.Tensor | None = None,

1247 + num_local_tokens: torch.Tensor | None = None,

1248 + output_dtype: torch.dtype | None = None,

1249 ) -> torch.Tensor:

1250 return torch.ops.vllm.rocm_aiter_fused_moe(

1251 hidden_states,

@@ -1251,6 +1261,8 @@ class rocm_aiter_ops:

1261 w2_scale,

1262 a1_scale,

1263 a2_scale,

1264 + num_local_tokens,

1265 + output_dtype,

1266 )

1267

1268 @staticmethod

vllm/config/parallel.py

@@ -43,6 +43,7 @@ All2AllBackend = Literal[

43 "pplx",

44 "deepep_high_throughput",

45 "deepep_low_latency",

46 + "mori",

47 "allgather_reducescatter",

48 "flashinfer_all2allv",

49 ]

@@ -158,6 +159,7 @@ class ParallelConfig:

159 - "pplx": Use pplx kernels\n

160 - "deepep_high_throughput": Use deepep high-throughput kernels\n

161 - "deepep_low_latency": Use deepep low-latency kernels\n

162 + - "mori": Use mori kernels\n

163 - "flashinfer_all2allv": Use flashinfer alltoallv kernels for mnnvl"""

164

165 max_parallel_loading_workers: int | None = None

@@ -443,6 +445,7 @@ class ParallelConfig:

445 "naive",

446 "deepep_high_throughput",

447 "deepep_low_latency",

448 + "mori",

449 )

450 and self.enable_expert_parallel

451 and self.tensor_parallel_size > 1

vllm/distributed/device_communicators/all2all.py

+94 -1

@@ -10,7 +10,7 @@ from vllm.distributed import get_dp_group, get_ep_group

10 from vllm.forward_context import get_forward_context

11 from vllm.logger import init_logger

12 from vllm.utils.flashinfer import has_flashinfer_all2all

13 -from vllm.utils.import_utils import has_deep_ep, has_pplx

13 +from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx

15 from .base_device_communicator import All2AllManagerBase, Cache

@@ -507,3 +507,96 @@ class FlashInferAllToAllManager(All2AllManagerBase):

507 self.prepare_workspace_tensor = None

508 self.mapping = None

509 self.initialized = False

510 +

511 +

512 +class MoriAll2AllManager(All2AllManagerBase):

513 + def __init__(self, cpu_group):

514 + assert has_mori(), (

515 + "MoRI kernels not found. Please follow https://github.com/ROCm/mori/blob/main/README.md"

516 + " to install MoRI kernels."

517 + ) # noqa

518 + import mori

519 +

520 + super().__init__(cpu_group)

521 + self.handle_cache = Cache()

522 +

523 + torch._C._distributed_c10d._register_process_group("mori", cpu_group)

524 + mori.shmem.shmem_torch_process_group_init("mori")

525 +

526 + def _make_all2all_kwargs(

527 + self,

528 + rank: int,

529 + num_ep_ranks: int,

530 + input_dtype: torch.dtype,

531 + quant_dtype: torch.dtype,

532 + token_hidden_size: int,

533 + scale_dim: int,

534 + scale_type_size: int,

535 + max_num_tokens_per_dp_rank: int,

536 + num_local_experts: int,

537 + num_experts_per_token: int,

538 + ):

539 + import mori # type: ignore[import-not-found]

540 +

541 + from vllm.platforms.rocm import on_gfx942, on_gfx950

542 +

543 + assert on_gfx942() or on_gfx950(), (

544 + "mori currently only support arch gfx942 and gfx950"

545 + )

546 +

547 + if not self.internode:

548 + # single node

549 + kernel_type = mori.ops.EpDispatchCombineKernelType.IntraNode

550 + rdma_block_num = 0

551 + warp_num_per_block = 16

552 + block_num = 80

553 + else:

554 + # multi node

555 + kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1

556 + if on_gfx942():

557 + warp_num_per_block = 16

558 + block_num = 32

559 + rdma_block_num = 16

560 + elif on_gfx950():

561 + warp_num_per_block = 8

562 + block_num = 64

563 + rdma_block_num = 32

564 + else:

565 + raise NotImplementedError(

566 + "mori currently only support arch gfx942 and gfx950"

567 + )

568 +

569 + return dict(

570 + rank=rank,

571 + world_size=num_ep_ranks,

572 + data_type=quant_dtype,

573 + hidden_dim=token_hidden_size,

574 + scale_dim=scale_dim,

575 + scale_type_size=scale_type_size,

576 + max_token_type_size=input_dtype.itemsize,

577 + max_num_inp_token_per_rank=max_num_tokens_per_dp_rank,

578 + num_experts_per_rank=num_local_experts,

579 + num_experts_per_token=num_experts_per_token,

580 + warp_num_per_block=warp_num_per_block,

581 + block_num=block_num,

582 + kernel_type=kernel_type,

583 + rdma_block_num=rdma_block_num,

584 + gpu_per_node=min(8, num_ep_ranks),

585 + )

586 +

587 + def _make_handle(self, **kwargs):

588 + import mori # type: ignore[import-not-found]

589 +

590 + mori_config = mori.ops.EpDispatchCombineConfig(**kwargs)

591 + handle = mori.ops.EpDispatchCombineOp(mori_config)

592 + return handle

593 +

594 + def get_handle(self, kwargs):

595 + import mori # type: ignore[import-not-found]

596 +

597 + mori_kwargs = self._make_all2all_kwargs(**kwargs)

598 + logger.debug("MoRI all2all args %s", mori_kwargs)

599 + handle: mori.ops.EpDispatchCombineOp = self.handle_cache.get_or_create(

600 + mori_kwargs, self._make_handle

601 + )

602 + return handle

vllm/distributed/device_communicators/cuda_communicator.py

@@ -110,6 +110,10 @@ class CudaCommunicator(DeviceCommunicatorBase):

110 from .all2all import DeepEPLLAll2AllManager

111

112 self.all2all_manager = DeepEPLLAll2AllManager(self.cpu_group)

113 + elif self.all2all_backend == "mori":

114 + from .all2all import MoriAll2AllManager

115 +

116 + self.all2all_manager = MoriAll2AllManager(self.cpu_group)

117 elif self.all2all_backend == "flashinfer_all2allv":

118 from .all2all import FlashInferAllToAllManager

119

vllm/envs.py

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

187 "pplx",

188 "deepep_high_throughput",

189 "deepep_low_latency",

190 + "mori",

191 "allgather_reducescatter",

192 "flashinfer_all2allv",

193 ] = "allgather_reducescatter"

@@ -1298,6 +1299,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1299 # - "pplx": use pplx kernels

1300 # - "deepep_high_throughput", use deepep high-throughput kernels

1301 # - "deepep_low_latency", use deepep low-latency kernels

1302 + # - "mori", use MoRI kernels

1303 # - "flashinfer_all2allv", use flashinfer alltoallv kernels for mnnvl

1304 "VLLM_ALL2ALL_BACKEND": env_with_choices(

1305 "VLLM_ALL2ALL_BACKEND",

@@ -1307,6 +1309,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1309 "pplx",

1310 "deepep_high_throughput",

1311 "deepep_low_latency",

1312 + "mori",

1313 "allgather_reducescatter",

1314 "flashinfer_all2allv",

1315 ],

vllm/model_executor/layers/fused_moe/__init__.py

@@ -88,6 +88,9 @@ if HAS_TRITON:

88 fused_experts,

89 get_config_file_name,

90 )

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

92 + AiterExperts,

93 + )

94 from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (

95 fused_topk,

96 )

@@ -99,6 +102,7 @@ if HAS_TRITON:

102 )

103

104 __all__ += [

105 + "AiterExperts",

106 "fused_topk",

107 "fused_experts",

108 "get_config_file_name",

vllm/model_executor/layers/fused_moe/all2all_utils.py

+34 -1

@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

16 FusedMoEPrepareAndFinalize,

17 )

18 from vllm.platforms import current_platform

19 -from vllm.utils.import_utils import has_deep_ep, has_pplx

19 +from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx

21 if current_platform.is_cuda_alike():

22 if has_pplx():

@@ -30,6 +30,8 @@ if current_platform.is_cuda_alike():

30 DEEPEP_QUANT_BLOCK_SHAPE,

31 DeepEPLLPrepareAndFinalize,

32 )

33 + if has_mori():

34 + from .mori_prepare_finalize import MoriPrepareAndFinalize

37 def maybe_roundup_layer_hidden_size(

@@ -169,5 +171,36 @@ def maybe_make_prepare_finalize(

171 physical_to_global=physical_to_global,

172 local_expert_global_ids=local_expert_global_ids,

173 )

174 + elif moe.use_mori_kernels:

175 + assert quant_config is not None

176 +

177 + # Note: We may want to use FP8 dispatch just to reduce

178 + # data movement.

179 + use_fp8_dispatch = (

180 + quant_config.is_per_act_token or quant_config.is_block_quantized

181 + )

182 + # For PTPC (per token per channel) quant, the scale dim for each token is 1

183 + # For 1x128 quant, the scale dim for each token is hidden_dim // 128

184 + scale_dim = 1 if quant_config.is_per_act_token else moe.hidden_dim // 128

185 + all_to_all_args = dict(

186 + rank=all2all_manager.rank,

187 + num_ep_ranks=all2all_manager.world_size,

188 + quant_dtype=quant_config.quant_dtype,

189 + token_hidden_size=moe.hidden_dim,

190 + scale_dim=scale_dim,

191 + scale_type_size=torch.float32.itemsize,

192 + max_num_tokens_per_dp_rank=moe.max_num_tokens,

193 + input_dtype=moe.in_dtype,

194 + num_local_experts=moe.num_experts // all2all_manager.world_size,

195 + num_experts_per_token=moe.experts_per_token,

196 + )

197 + handle = all2all_manager.get_handle(all_to_all_args)

198 +

199 + prepare_finalize = MoriPrepareAndFinalize(

200 + handle,

201 + max_tokens_per_rank=moe.max_num_tokens,

202 + num_dispatchers=all2all_manager.world_size,

203 + use_fp8_dispatch=use_fp8_dispatch,

204 + )

205

206 return prepare_finalize

vllm/model_executor/layers/fused_moe/config.py

@@ -893,6 +893,10 @@ class FusedMoEParallelConfig:

893 self.all2all_backend in ["naive", "allgather_reducescatter"]

894 )

895

896 + @property

897 + def use_mori_kernels(self):

898 + return self.use_all2all_kernels and self.all2all_backend == "mori"

899 +

900 @staticmethod

901 def flatten_tp_across_dp_and_pcp(

902 tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int

@@ -1136,6 +1140,10 @@ class FusedMoEConfig:

1140 def use_deepep_ll_kernels(self):

1141 return self.moe_parallel_config.use_deepep_ll_kernels

1142

1143 + @property

1144 + def use_mori_kernels(self):

1145 + return self.moe_parallel_config.use_mori_kernels

1146 +

1147 @property

1148 def use_flashinfer_cutlass_kernels(self):

1149 """

vllm/model_executor/layers/fused_moe/layer.py

+13

@@ -570,6 +570,14 @@ class FusedMoE(CustomOp):

570 self.moe_config_use_flashinfer_cutlass_kernels = (

571 self.moe_config.use_flashinfer_cutlass_kernels

572 )

573 + if self.use_mori_kernels:

574 + assert self.rocm_aiter_fmoe_enabled, (

575 + "Mori needs to be used with aiter fused_moe for now."

576 + )

577 + assert not self.aiter_fmoe_shared_expert_enabled, (

578 + "Mori does not support fusion shared expert now. "

579 + "Turn it off by setting VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0"

580 + )

581

582 self.quant_config = quant_config

583

@@ -712,6 +720,10 @@ class FusedMoE(CustomOp):

720 def use_deepep_ll_kernels(self):

721 return self.moe_parallel_config.use_deepep_ll_kernels

722

723 + @property

724 + def use_mori_kernels(self):

725 + return self.moe_parallel_config.use_mori_kernels

726 +

727 @property

728 def use_flashinfer_cutlass_kernels(self):

729 return (

@@ -729,6 +741,7 @@ class FusedMoE(CustomOp):

741 return (

742 self.moe_parallel_config.use_pplx_kernels

743 or self.moe_parallel_config.use_deepep_ll_kernels

744 + or self.moe_parallel_config.use_mori_kernels

745 or (self.dp_size > 1 and self.use_flashinfer_cutlass_kernels)

746 ) and envs.VLLM_ENABLE_MOE_DP_CHUNK

747

vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py new

+121

@@ -0,0 +1,121 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +import mori

5 +import torch

6 +

7 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 +from vllm.logger import init_logger

9 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

10 +from vllm.platforms import current_platform

11 +

12 +logger = init_logger(__name__)

13 +

14 +

15 +class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

16 + """

17 + Prepare/Finalize using MoRI kernels.

18 + """

19 +

20 + def __init__(

21 + self,

22 + mori_op: mori.ops.EpDispatchCombineOp,

23 + max_tokens_per_rank: int,

24 + num_dispatchers: int,

25 + use_fp8_dispatch: bool = False,

26 + ):

27 + super().__init__()

28 + self.mori_op = mori_op

29 + self.num_dispatchers_ = num_dispatchers

30 + self.max_tokens_per_rank = max_tokens_per_rank

31 + self.use_fp8_dispatch = use_fp8_dispatch

32 +

33 + @property

34 + def activation_format(self) -> mk.FusedMoEActivationFormat:

35 + return mk.FusedMoEActivationFormat.Standard

36 +

37 + def output_is_reduced(self) -> bool:

38 + return True

39 +

40 + def num_dispatchers(self):

41 + return self.num_dispatchers_

42 +

43 + def max_num_tokens_per_rank(self) -> int | None:

44 + return self.max_tokens_per_rank

45 +

46 + def topk_indices_dtype(self) -> torch.dtype | None:

47 + return torch.int32

48 +

49 + def supports_async(self) -> bool:

50 + return False

51 +

52 + def prepare(

53 + self,

54 + a1: torch.Tensor,

55 + topk_weights: torch.Tensor,

56 + topk_ids: torch.Tensor,

57 + num_experts: int,

58 + expert_map: torch.Tensor | None,

59 + apply_router_weight_on_input: bool,

60 + quant_config: FusedMoEQuantConfig,

61 + ) -> mk.PrepareResultType:

62 + """

63 + Returns a tuple of:

64 + - quantized + dispatched a.

65 + - Optional quantized + dispatched a1_scales.

66 + - Optional ExpertTokensMetadata containing gpu/cpu tensors

67 + as big as the number of local experts with the information about the

68 + number of tokens assigned to each local expert.

69 + - Optional dispatched expert topk IDs

70 + - Optional dispatched expert topk weight

71 + """

72 + assert not apply_router_weight_on_input, (

73 + "mori does not support apply_router_weight_on_input=True now."

74 + )

75 + scale = None

76 + if self.use_fp8_dispatch:

77 + from aiter import QuantType, get_hip_quant

78 +

79 + if quant_config.is_block_quantized:

80 + quant_func = get_hip_quant(QuantType.per_1x128)

81 + a1, scale = quant_func(a1, quant_dtype=current_platform.fp8_dtype())

82 + elif quant_config.is_per_act_token:

83 + quant_func = get_hip_quant(QuantType.per_Token)

84 + a1, scale = quant_func(a1, quant_dtype=current_platform.fp8_dtype())

85 +

86 + (

87 + dispatch_a1,

88 + dispatch_weights,

89 + dispatch_scale,

90 + dispatch_ids,

91 + dispatch_recv_token_num,

92 + ) = self.mori_op.dispatch(a1, topk_weights, scale, topk_ids)

93 +

94 + expert_tokens_meta = mk.ExpertTokensMetadata(

95 + expert_num_tokens=dispatch_recv_token_num, expert_num_tokens_cpu=None

96 + )

97 +

98 + return (

99 + dispatch_a1,

100 + dispatch_scale,

101 + expert_tokens_meta,

102 + dispatch_ids,

103 + dispatch_weights,

104 + )

105 +

106 + def finalize(

107 + self,

108 + output: torch.Tensor,

109 + fused_expert_output: torch.Tensor,

110 + topk_weights: torch.Tensor,

111 + topk_ids: torch.Tensor,

112 + apply_router_weight_on_input: bool,

113 + weight_and_reduce_impl: mk.TopKWeightAndReduce,

114 + ) -> None:

115 + num_token = output.shape[0]

116 + result = self.mori_op.combine(

117 + fused_expert_output,

118 + None,

119 + topk_ids,

120 + )[0]

121 + output.copy_(result[:num_token])

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+17 -4

@@ -188,6 +188,9 @@ def rocm_aiter_fused_experts(

188 apply_router_weight_on_input: bool = False,

189 expert_map: torch.Tensor | None = None,

190 quant_config: FusedMoEQuantConfig | None = None,

191 + a1q_scale: torch.Tensor | None = None,

192 + num_local_tokens: torch.Tensor | None = None,

193 + output_dtype: torch.dtype | None = None,

194 ) -> torch.Tensor:

195 if quant_config is None:

196 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

@@ -216,6 +219,9 @@ def rocm_aiter_fused_experts(

219 assert topk_weights.shape[-1] == 1, (

220 "Only support topk=1 when `apply_router_weight_on_input` is True"

221 )

222 + assert num_local_tokens is None, (

223 + "AITER tkw1 kernel does not support `num_local_tokens`"

224 + )

225

226 return rocm_aiter_ops.asm_moe_tkw1(

227 hidden_states,

@@ -272,9 +278,11 @@ def rocm_aiter_fused_experts(

278 activation_method=activation_method,

279 w1_scale=quant_config.w1_scale,

280 w2_scale=quant_config.w2_scale,

275 - a1_scale=quant_config.a1_scale,

281 + a1_scale=quant_config.a1_scale if a1q_scale is None else a1q_scale,

282 a2_scale=quant_config.a2_scale,

283 doweight_stage1=apply_router_weight_on_input,

284 + num_local_tokens=num_local_tokens,

285 + output_dtype=output_dtype,

286 )

287

288

@@ -370,9 +378,12 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

378 # TODO(rob): rocm_aiter_fused_experts uses self.quant_config's

379 # a_scales for static quantization. Update this to fit better

380 # with the interface once all quant integrations are complete.

373 - assert a1q_scale is None

381 assert a2_scale == self.quant_config.a2_scale

375 - assert expert_tokens_meta is None

382 +

383 + if expert_tokens_meta is not None:

384 + num_local_tokens = expert_tokens_meta.expert_num_tokens

385 + else:

386 + num_local_tokens = None

387

388 result = rocm_aiter_fused_experts(

389 hidden_states=hidden_states,

@@ -384,6 +395,8 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

395 apply_router_weight_on_input=apply_router_weight_on_input,

396 expert_map=expert_map,

397 quant_config=self.quant_config,

398 + a1q_scale=a1q_scale,

399 + num_local_tokens=num_local_tokens,

400 + output_dtype=output.dtype,

401 )

388 - assert result.shape == output.shape

402 output.copy_(result)

vllm/platforms/rocm.py

@@ -106,6 +106,12 @@ def on_gfx9() -> bool:

106 return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])

107

108

109 +@cache

110 +def on_gfx942() -> bool:

111 + GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName

112 + return any(arch in GPU_ARCH for arch in ["gfx942"])

113 +

114 +

115 @cache

116 def on_gfx950() -> bool:

117 GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName

vllm/utils/import_utils.py

+10

@@ -451,3 +451,13 @@ def has_helion() -> bool:

451 # use helion...

452 """

453 return _has_module("helion")

454 +

455 +

456 +def has_aiter() -> bool:

457 + """Whether the optional `aiter` package is available."""

458 + return _has_module("aiter")

459 +

460 +

461 +def has_mori() -> bool:

462 + """Whether the optional `mori` package is available."""

463 + return _has_module("mori")