diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 633e23eea33e..48deb86caadc 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -41,6 +41,7 @@ th { | flashinfer4 | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferCutlassMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferCutlassMoEPrepareAndFinalize] | | MoEPrepareAndFinalizeNoEP5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] | | BatchedPrepareAndFinalize5 | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] | +| MoriPrepareAndFinalize7 | standard | fp88 | G(128),A,T8 | N | Y | [`MoriPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.mori_prepare_finalize.MoriPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 @@ -49,6 +50,8 @@ th { 4. Controlled by different env vars (`VLLM_FLASHINFER_MOE_BACKEND` "throughput" or "latency") 5. This is a no-op dispatcher that can be used to pair with any modular experts to produce a modular kernel that runs w/o dispatch or combine. These cannot be selected via environment variable. These are generally use for testing or adapting an expert subclass to the `fused_experts` API. 6. This depends on the experts implementation. + 7. Currently, MoRI supports low-latency mode only. + 8. This depends on the experts implementation, currently mori supports aiter. --- @@ -117,4 +120,5 @@ The following table shows "families" of modular kernels that are intended to wor |----------------------------------|------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------| | deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,
`TritonExperts`,
`TritonOrDeepGemmExperts`,
`CutlassExpertsFp8`,
`MarlinExperts` | | deepep_low_latency,
pplx | `DeepEPLLPrepareAndFinalize`,
`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`BatchedTritonOrDeepGemmExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts`| -| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` | +| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` | +| mori | `MoriPrepareAndFinalize` | `AiterMoriExperts` | diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 013ef3c1f5c3..de22d10bcefb 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import os +from pathlib import Path from typing import Any +import psutil import torch import torch.distributed as dist @@ -10,7 +14,7 @@ from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.utils.flashinfer import has_flashinfer_all2all -from vllm.utils.import_utils import has_deep_ep, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx from .base_device_communicator import All2AllManagerBase, Cache @@ -488,3 +492,342 @@ def cleanup(self): self.prepare_workspace_tensor = None self.mapping = None self.initialized = False + + +class MoriAll2AllManager(All2AllManagerBase): + """ + All2All communication based on mori kernels. + """ + + def __init__(self, cpu_group): + assert has_mori(), "Please install mori from ROCm/mori github." + + super().__init__(cpu_group) + self.handle_cache = Cache() + self.config = None + self._shmem_initialized = False + + self.json_config = None + config_path = envs.VLLM_MORI_CONFIG_PATH + if config_path: + self.json_config = self._load_mori_config_from_json(config_path) + + # Delay mori shmem initialization until first use + logger.debug("[rank %s] MoriAll2AllManager created", self.rank) + + def _ensure_shmem_initialized(self): + """Initialize mori's shared memory system lazily""" + if self._shmem_initialized: + return + + import torch.distributed as dist + from mori.shmem import shmem_torch_process_group_init + + try: + # Check if we have a valid backend + backend = dist.get_backend() + if backend is None: + raise RuntimeError("No valid distributed backend found") + + logger.debug( + "[rank %s] PyTorch distributed ready with backend: %s", + self.rank, + backend, + ) + + assert self.cpu_group is not None, "No CPU group is given to mori" + ppid = psutil.Process(os.getpid()).ppid() + group_name = f"mori_shmem_group_{ppid}" + + try: + import torch._C._distributed_c10d as c10d + + # Register the process group + c10d._register_process_group(group_name, self.cpu_group) + logger.debug( + "[rank %s] Registered proc group %s", self.rank, group_name + ) + + # Initialize mori shmem with the registered group + shmem_torch_process_group_init(group_name) + logger.debug("[rank %s] torch proc group shmem init success", self.rank) + self._shmem_initialized = True + return + + except Exception as torch_error: + raise RuntimeError( + "torch process group initialization failed" + ) from torch_error + + except Exception as e: + raise RuntimeError("mori shmem initialization failed") from e + + def _load_mori_config_from_json(self, json_path: str) -> dict | None: + """ + Load mori configuration parameters from JSON file. + + Supports both flat and hierarchical schema: + + Flat schema: + { + "warp_num_per_block": 8, + "block_num": 80, + } + + Hierarchical schema (dispatch/combine specific): + { + "global": { + "warp_num_per_block": 8, + "block_num": 80, + }, + "dispatch": { + "warp_num_per_block": 16, + "block_num": 160 + }, + "combine": { + "warp_num_per_block": 4, + "block_num": 40 + } + } + + Args: + json_path: Path to JSON configuration file + + Returns: + Dictionary of configuration parameters, or None if file doesn't exist + + Raises: + ValueError: If JSON is invalid or contains unsupported parameters + """ + if not json_path: + return None + + json_file = Path(json_path) + if not json_file.exists(): + logger.warning( + "[rank %d] Mori config file not found: %s", self.rank, json_path + ) + return None + + try: + with open(json_file) as f: + config = json.load(f) + + # Valid parameter keys + valid_param_keys = { + "warp_num_per_block", + "block_num", + } + + is_hierarchical = any( + key in config for key in ["global", "dispatch", "combine"] + ) + + if is_hierarchical: + valid_top_keys = {"global", "dispatch", "combine"} + invalid_keys = set(config.keys()) - valid_top_keys + if invalid_keys: + raise ValueError( + f"Invalid top-level keys: {invalid_keys}. " + f"Valid keys: {valid_top_keys}" + ) + + # Validate each section + for section in ["global", "dispatch", "combine"]: + if section in config: + section_config = config[section] + if not isinstance(section_config, dict): + raise ValueError(f"'{section}' must be a dictionary") + + invalid_keys = set(section_config.keys()) - valid_param_keys + if invalid_keys: + raise ValueError( + f"Invalid keys in '{section}': {invalid_keys}. " + f"Valid keys: {valid_param_keys}" + ) + else: + invalid_keys = set(config.keys()) - valid_param_keys + if invalid_keys: + raise ValueError( + f"Invalid config keys: {invalid_keys}. " + f"Valid keys: {valid_param_keys}" + ) + + return config + + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON in mori config file {json_path}") from e + except Exception as e: + raise ValueError(f"Error loading mori config from {json_path}") from e + + def _make_mori_config( + self, + max_num_tokens: int, + num_local_experts: int, + experts_per_token: int, + hidden_dim: int, + scale_dim: int, + scale_type_size: int, + data_type: torch.dtype = torch.bfloat16, + quant_dtype: torch.dtype | None = None, + ): + """ + Create mori EpDispatchCombineConfig. + + Args: + max_num_tokens: Maximum number of tokens per DP rank + num_local_experts: Number of local experts + experts_per_token: Number of experts per token (topk) + hidden_dim: Hidden dimension size + scale_dim: Scale dimension for quantization + scale_type_size: Scale type size for quantization + data_type: Tensor data type + quant_dtype: Quantization data type (optional) + """ + from mori.ops import EpDispatchCombineConfig + from mori.ops.dispatch_combine import EpDispatchCombineKernelType + + from vllm.platforms import current_platform + + assert quant_dtype is None or quant_dtype == current_platform.fp8_dtype() + + # Default values (can be overridden by JSON) + warp_num_per_block = 8 + block_num = 80 + + # Override with JSON config if provided + if self.json_config is not None: + is_hierarchical = any( + key in self.json_config for key in ["global", "dispatch", "combine"] + ) + + global_config = self.json_config + if is_hierarchical and "global" in global_config: + global_config = self.json_config["global"] + + warp_num_per_block = global_config.get( + "warp_num_per_block", warp_num_per_block + ) + block_num = global_config.get("block_num", block_num) + + config = EpDispatchCombineConfig( + data_type=data_type if quant_dtype is None else quant_dtype, + rank=self.rank, + world_size=self.world_size, + hidden_dim=hidden_dim, + max_num_inp_token_per_rank=max_num_tokens, + num_experts_per_rank=num_local_experts, + num_experts_per_token=experts_per_token, + max_token_type_size=data_type.itemsize, + # Performance tuning parameters + warp_num_per_block=warp_num_per_block, + block_num=block_num, + # Quantization support + scale_dim=scale_dim, + scale_type_size=scale_type_size, + # Determine kernel type based on topology + kernel_type=( + EpDispatchCombineKernelType.InterNode + if self.internode + else EpDispatchCombineKernelType.IntraNode + ), + ) + + return config + + def get_handle(self, kwargs): + """ + Get or create mori operation handle. + Args: + kwargs: Dictionary with keys: + - max_num_tokens: Maximum tokens per DP rank + - num_local_experts: Number of local experts + - experts_per_token: Number of experts per token (topk) + - hidden_dim: Hidden dimension size + - data_type: Tensor data type (optional, default bfloat16) + - scale_dim: Scale dimension (optional) + - scale_type_size: Scale type size (optional) + - ubatch_id: Microbatch ID (optional) + """ + # Ensure shmem is initialized before creating handles + self._ensure_shmem_initialized() + + def create_mori_handle( + max_num_tokens: int, + num_local_experts: int, + experts_per_token: int, + hidden_dim: int, + scale_dim: int, + scale_type_size: int, + data_type: torch.dtype = torch.bfloat16, + quant_dtype: torch.dtype | None = None, + ): + from mori.ops import EpDispatchCombineOp + + config = self._make_mori_config( + max_num_tokens=max_num_tokens, + num_local_experts=num_local_experts, + experts_per_token=experts_per_token, + hidden_dim=hidden_dim, + scale_dim=scale_dim, + scale_type_size=scale_type_size, + data_type=data_type, + quant_dtype=quant_dtype, + ) + op = EpDispatchCombineOp(config) + logger.debug( + "[rank %s] Created mori handle with config: tokens=%d, experts=%d," + " topk=%d, hidden_dim=%d", + self.dp_rank, + max_num_tokens, + num_local_experts, + experts_per_token, + hidden_dim, + ) + return op + + return self.handle_cache.get_or_create(kwargs, create_mori_handle) + + def dispatch( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + is_sequence_parallel: bool = False, + ): + raise NotImplementedError + + def combine( + self, + hidden_states: torch.Tensor, + is_sequence_parallel: bool = False, + ): + raise NotImplementedError + + def destroy(self): + """Clean up mori resources""" + try: + # Clear operation handle cache + with self.handle_cache._lock: + for _, handle in self.handle_cache._cache.items(): + handle.destroy() + + # finalize mori shared memory if it was initialized + if self._shmem_initialized: + try: + from mori.shmem import shmem_finalize + + # Check if shmem is actually active before finalizing + shmem_finalize() + logger.debug("[rank %s] mori shmem finalize", self.dp_rank) + except Exception as shmem_error: + logger.debug( + "[rank %s] shmem finalize failed " + "(may not have been active): %s", + self.dp_rank, + shmem_error, + ) + + logger.debug("[rank %s] mori resources cleaned up", self.dp_rank) + + except Exception as e: + logger.warning("[rank %s] mori cleanup fail: %s", self.dp_rank, e) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 2e878eef908a..1a2242add599 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -114,6 +114,12 @@ def __init__( from .all2all import FlashInferAllToAllManager self.all2all_manager = FlashInferAllToAllManager(self.cpu_group) + logger.info("Using Flashinfer all2allv manager.") + elif self.all2all_backend == "mori": + from .all2all import MoriAll2AllManager + + self.all2all_manager = MoriAll2AllManager(self.cpu_group) + logger.info("Using Mori all2all manager.") else: raise ValueError(f"Unknown all2all backend: {self.all2all_backend}") diff --git a/vllm/envs.py b/vllm/envs.py index 0c45f93ec057..cf008ea2375e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -163,6 +163,7 @@ VLLM_ALL2ALL_BACKEND: Literal[ "naive", "pplx", + "mori", "deepep_high_throughput", "deepep_low_latency", "allgather_reducescatter", @@ -217,6 +218,7 @@ VLLM_NCCL_INCLUDE_PATH: str | None = None VLLM_USE_FBGEMM: bool = False VLLM_GC_DEBUG: str = "" + VLLM_MORI_CONFIG_PATH: str | None = None VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False @@ -1153,6 +1155,7 @@ def get_vllm_port() -> int | None: # - "allgather_reducescatter": all2all implementation based on allgather and # reducescatter # - "pplx": use pplx kernels + # - "mori": use mori kernels (currently, only low-latency is supported) # - "deepep_high_throughput", use deepep high-throughput kernels # - "deepep_low_latency", use deepep low-latency kernels # - "flashinfer_all2allv", use flashinfer alltoallv kernels for mnnvl @@ -1162,6 +1165,7 @@ def get_vllm_port() -> int | None: [ "naive", "pplx", + "mori", "deepep_high_throughput", "deepep_low_latency", "allgather_reducescatter", @@ -1404,6 +1408,9 @@ def get_vllm_port() -> int | None: # - VLLM_GC_DEBUG='{"top_objects":5}': enable GC debugger with # top 5 collected objects "VLLM_GC_DEBUG": lambda: os.getenv("VLLM_GC_DEBUG", ""), + # Path to JSON configuration file for mori all2all parameters + # If set, mori will use parameters from this JSON file instead of defaults + "VLLM_MORI_CONFIG_PATH": lambda: os.getenv("VLLM_MORI_CONFIG_PATH", None), # Disables parallel execution of shared_experts via separate cuda stream "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: os.getenv( "VLLM_DISABLE_SHARED_EXPERTS_STREAM", False diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index cb31045971bd..cb16cb39ce6f 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -17,6 +17,7 @@ ) from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE from vllm.model_executor.layers.fused_moe.utils import activation_without_mul +from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON _config: dict[str, Any] | None = None @@ -102,3 +103,10 @@ def _raise_exception(method: str): fused_topk = lambda *args, **kwargs: _raise_exception("fused_topk") fused_experts = lambda *args, **kwargs: _raise_exception("fused_experts") + +if current_platform.is_rocm(): + from vllm.model_executor.layers.fused_moe.aiter_mori_experts import AiterMoriExperts + + __all__ += [ + "AiterMoriExperts", + ] diff --git a/vllm/model_executor/layers/fused_moe/aiter_mori_experts.py b/vllm/model_executor/layers/fused_moe/aiter_mori_experts.py new file mode 100644 index 000000000000..d9bbd6ac3ed5 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/aiter_mori_experts.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Aiter-based expert processing for Mori integration. +""" + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( + rocm_aiter_fused_experts, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) + + +class AiterMoriExperts(mk.FusedMoEPermuteExpertsUnpermute): + """ + Aiter-based expert processing that works with Mori dispatch/combine. + + This class bridges Mori's all2all communication with Aiter's optimized + expert computation kernels for AMD GPUs. + """ + + def __init__( + self, + max_num_tokens: int, + quant_config: FusedMoEQuantConfig, + ): + from vllm.platforms.rocm import on_mi3xx + + if not on_mi3xx(): + raise RuntimeError("AiterMoriExperts should be used on AMD mi3xx GPUs") + + super().__init__( + quant_config=quant_config, + ) + self.max_num_tokens = max_num_tokens + + @property + def activation_formats( + self, + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]: + """Aiter expects Standard format for both input and output.""" + return ( + mk.FusedMoEActivationFormat.Standard, + mk.FusedMoEActivationFormat.Standard, + ) + + def supports_chunking(self) -> bool: + """Aiter kernels support chunking.""" + return True + + def supports_expert_map(self) -> bool: + """Aiter kernels support expert mapping.""" + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + """Aiter handles weight and reduce internally.""" + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """ + Aiter kernels manage memory internally, so minimal workspace is needed. + """ + workspace1 = (M, K) + workspace2 = (0,) # No intermediate workspace needed + output_shape = (M, K) + return (workspace1, workspace2, output_shape) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + """ + Process expert computation using Aiter kernels. + Works with pre-dispatched tokens from Mori all2all. + """ + if expert_tokens_meta is not None: + expert_num_tokens = expert_tokens_meta.expert_num_tokens + else: + expert_num_tokens = None + + # Call Aiter fused MoE expert processing + result = rocm_aiter_fused_experts( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + expert_map=expert_map, + expert_num_tokens=expert_num_tokens, + output_dtype=output.dtype, + quant_config=self.quant_config, + a1q_scale=a1q_scale, + ) + + # Copy result to output tensor + output.copy_(result) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 5403d4e62f85..2baa36a05208 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -683,6 +683,10 @@ def use_deepep_ht_kernels(self): def use_deepep_ll_kernels(self): return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency" + @property + def use_mori_kernels(self): + return self.use_all2all_kernels and envs.VLLM_ALL2ALL_BACKEND == "mori" + @staticmethod def flatten_tp_across_dp( tp_size: int, dp_size: int, dp_rank: int @@ -873,6 +877,10 @@ def use_deepep_ht_kernels(self): def use_deepep_ll_kernels(self): return self.moe_parallel_config.use_deepep_ll_kernels + @property + def use_mori_kernels(self): + return self.moe_parallel_config.use_mori_kernels + @property def use_flashinfer_cutlass_kernels(self): """ diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c144aa23e46e..af062b8c2ed9 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -56,7 +56,7 @@ from vllm.platforms import current_platform from vllm.platforms.interface import CpuArchEnum from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe -from vllm.utils.import_utils import has_deep_ep, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx from vllm.utils.math_utils import cdiv, round_up from vllm.utils.torch_utils import current_stream, direct_register_custom_op from vllm.v1.worker.ubatching import dbo_current_ubatch_id @@ -76,6 +76,8 @@ DEEPEP_QUANT_BLOCK_SHAPE, DeepEPLLPrepareAndFinalize, ) + if has_mori(): + from .mori_prepare_finalize import MoriPrepareAndFinalize else: fused_experts = None # type: ignore FusedMoEPermuteExpertsUnpermute = object # type: ignore @@ -99,6 +101,7 @@ def _eplb_map_to_physical_and_record( ) else: from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk + if current_platform.is_tpu(): from .moe_pallas import fused_moe as fused_moe_pallas else: @@ -234,6 +237,47 @@ def _maybe_make_prepare_finalize( num_dispatchers=all2all_manager.world_size, use_fp8_dispatch=use_fp8_dispatch, ) + elif moe.use_mori_kernels: + use_fp8_dispatch = ( + quant_config is not None + and quant_config.quant_dtype == current_platform.fp8_dtype() + ) + scale_dim = 0 + scale_type_size = 0 + quant_dtype = None + if use_fp8_dispatch: + assert quant_config is not None + temp = quant_config.scale_shape( + moe.max_num_tokens, + moe.hidden_dim, + ) + if temp is not None: + scale_dim = temp[-1] + scale_type_size = ( + torch.float32.itemsize + ) # aiter quantization uses float32 scale + quant_dtype = quant_config.quant_dtype + + all_to_all_args = dict( + max_num_tokens=moe.max_num_tokens, + num_local_experts=moe.num_local_experts, + experts_per_token=moe.experts_per_token, + hidden_dim=moe.hidden_dim, + data_type=moe.in_dtype, + quant_dtype=quant_dtype, + scale_dim=scale_dim, + scale_type_size=scale_type_size, + ) + handle = all2all_manager.get_handle(all_to_all_args) + + prepare_finalize = MoriPrepareAndFinalize( + handle, + max_num_tokens=moe.max_num_tokens, + num_local_experts=moe.num_local_experts, + num_dispatchers=all2all_manager.world_size, + use_fp8_dispatch=use_fp8_dispatch, + json_config=all2all_manager.json_config, + ) return prepare_finalize @@ -400,6 +444,14 @@ def select_gemm_impl( num_dispatchers=prepare_finalize.num_dispatchers(), quant_config=self.moe_quant_config, ) + elif self.moe.use_mori_kernels and is_rocm_aiter_moe_enabled(): + from vllm.model_executor.layers.fused_moe import AiterMoriExperts + + logger.debug("AiterMoriExperts for Mori integration %s", self.moe) + return AiterMoriExperts( + max_num_tokens=self.moe.max_num_tokens, + quant_config=self.moe_quant_config, + ) else: logger.debug("TritonExperts %s", self.moe) return TritonExperts(self.moe_quant_config) @@ -1393,6 +1445,10 @@ def use_deepep_ht_kernels(self): def use_deepep_ll_kernels(self): return self.moe_parallel_config.use_deepep_ll_kernels + @property + def use_mori_kernels(self): + return self.moe_parallel_config.use_mori_kernels + @property def use_flashinfer_cutlass_kernels(self): return ( @@ -1406,6 +1462,7 @@ def use_dp_chunking(self) -> bool: return ( self.moe_parallel_config.use_pplx_kernels or self.moe_parallel_config.use_deepep_ll_kernels + or self.moe_parallel_config.use_mori_kernels or (self.dp_size > 1 and self.use_flashinfer_cutlass_kernels) ) diff --git a/vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py new file mode 100644 index 000000000000..f16821152b23 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +mori prepare and finalize module for expert parallelism. +Migration from DeepEP to mori for AMD GPU support. +""" + +from typing import Any + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig + +logger = init_logger(__name__) + + +class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): + """ + Prepare/Finalize using mori kernels for AMD GPU expert parallelism. + + This class handles the dispatch and combine operations for + expert parallelism using the mori library, which provides optimized + All2All communication primitives for AMD GPUs. + """ + + def __init__( + self, + handle: Any, # mori EpDispatchCombineOp from MoriAll2AllManager + max_num_tokens: int, + num_local_experts: int, + num_dispatchers: int, + use_fp8_dispatch: bool = False, + json_config: dict | None = None, + ): + """ + Initialize MoriPrepareAndFinalize. + + Args: + handle: mori EpDispatchCombineOp instance from All2AllManager + max_num_tokens: Maximum number of tokens per rank + num_local_experts: Number of experts on this rank + num_dispatchers: Number of dispatcher ranks (world size) + use_fp8_dispatch: Whether to use FP8 quantization during dispatch + json_config: Optional JSON configuration with operation-specific parameters + """ + super().__init__() + assert max_num_tokens > 0 + assert num_local_experts > 0 + + self.handle = handle # mori EpDispatchCombineOp + self.max_num_tokens = max_num_tokens + self.num_local_experts = num_local_experts + self.num_dispatchers_ = num_dispatchers + self.use_fp8_dispatch = use_fp8_dispatch + + # Extract dispatch and combine specific parameters from JSON config + self.dispatch_kwargs = {} + self.combine_kwargs = {} + + if json_config: + # Extract dispatch-specific parameters + if "dispatch" in json_config: + dispatch_config = json_config["dispatch"] + + if "block_num" in dispatch_config: + self.dispatch_kwargs["block_num"] = dispatch_config["block_num"] + if "warp_num_per_block" in dispatch_config: + self.dispatch_kwargs["warp_per_block"] = dispatch_config[ + "warp_num_per_block" + ] + + # Extract combine-specific parameters + if "combine" in json_config: + combine_config = json_config["combine"] + + if "block_num" in combine_config: + self.combine_kwargs["block_num"] = combine_config["block_num"] + if "warp_num_per_block" in combine_config: + self.combine_kwargs["warp_per_block"] = combine_config[ + "warp_num_per_block" + ] + + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return self.max_num_tokens + + def topk_indices_dtype(self) -> torch.dtype | None: + return torch.int32 + + def num_dispatchers(self) -> int: + return self.num_dispatchers_ + + def output_is_reduced(self) -> bool: + return True + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + ) -> mk.PrepareResultType: + """ + Prepare inputs for mori dispatch operation. + Supports pre-dispatch quantization to reduce communication overhead. + + Args: + a1: Input hidden states [num_tokens, hidden_dim] + topk_weights: Top-k routing weights [num_experts, experts_per_token] + topk_ids: Top-k expert indices [num_experts, experts_per_token] + apply_router_weight_on_input: Whether to apply router weight + quant_config: Quantization config + + Returns: + Tuple of (dispatched_x, batched_scales, expert_tokens_meta, + dispatch_indices, dispatch_weights) + where dispatched_x is in Standard format (2D tensor) + """ + # Pre-dispatch quantization to reduce communication overhead + dispatch_input = a1 + scales = None + + if self.use_fp8_dispatch: + from aiter import QuantType, get_hip_quant + + if quant_config.block_shape is not None: + assert not apply_router_weight_on_input, ( + "apply_router_weight_on_input is not supported for block scaled moe" + ) + quant_type = QuantType.per_1x128 + else: + quant_type = QuantType.per_Tensor + + quant_func = get_hip_quant(quant_type) + + dispatch_input, scales = quant_func( + a1, + quant_dtype=quant_config.quant_dtype, + ) + + ( + dispatch_output, + dispatch_weights, + dispatch_scales, + dispatch_indices, + dispatch_recv_num_token, + ) = self.handle.dispatch( + input=dispatch_input, + weights=topk_weights, + scales=scales, + indices=topk_ids, + **self.dispatch_kwargs, # Apply dispatch-specific parameters from JSON + ) + + expert_tokens_meta = mk.ExpertTokensMetadata( + expert_num_tokens=dispatch_recv_num_token, + expert_num_tokens_cpu=None, + ) + + return ( + dispatch_output, + dispatch_scales, + expert_tokens_meta, + dispatch_indices, + dispatch_weights, + ) + + def finalize( + self, + output: torch.Tensor, + fused_expert_output: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + extra_finalize_args: dict | None = None, + ) -> None: + """ + Finalize expert outputs using mori combine operation. + + Args: + output: Output tensor to write results [num_original_tokens, + hidden_dim] + fused_expert_output: Expert output activations in Standard format + (2D tensor) + topk_weights: Original top-k weights + topk_ids: Original top-k indices + """ + assert self.handle is not None + + num_original_tokens = output.size(0) # Original number of tokens + + combined_output, combined_weights = self.handle.combine( + input=fused_expert_output, + weights=topk_weights, + indices=topk_ids, + **self.combine_kwargs, # Apply combine-specific parameters from JSON + ) + + output.copy_( + combined_output[:num_original_tokens], + non_blocking=True, + ) + + def __repr__(self) -> str: + return ( + f"MoriPrepareAndFinalize(max_tokens={self.max_num_tokens}, " + f"num_local_experts={self.num_local_experts}, " + f"num_dispatchers={self.num_dispatchers_})" + ) diff --git a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py index e18514ad43f6..7a8e26b9ec4a 100644 --- a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py @@ -281,10 +281,18 @@ def rocm_aiter_fused_moe_impl( w2_scale: torch.Tensor | None = None, a1_scale: torch.Tensor | None = None, a2_scale: torch.Tensor | None = None, + expert_num_tokens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, ) -> torch.Tensor: from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe + # Check if input is already pre-quantized (from mori dispatch) + input_is_pre_quantized = ( + a1_scale is not None and hidden_states.dtype == current_platform.fp8_dtype() + ) + dtype = output_dtype if input_is_pre_quantized else None + activation = ActivationType(activation_method) quant_type = QuantType(quant_method) @@ -302,6 +310,8 @@ def rocm_aiter_fused_moe_impl( w2_scale, a1_scale, a2_scale, + num_local_tokens=expert_num_tokens, + dtype=dtype, ) @@ -319,6 +329,8 @@ def rocm_aiter_fused_moe_fake( w2_scale: torch.Tensor | None = None, a1_scale: torch.Tensor | None = None, a2_scale: torch.Tensor | None = None, + expert_num_tokens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, ) -> torch.Tensor: return torch.empty_like(hidden_states) @@ -434,7 +446,10 @@ def rocm_aiter_fused_experts( activation: str = "silu", apply_router_weight_on_input: bool = False, expert_map: torch.Tensor | None = None, + expert_num_tokens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, quant_config: FusedMoEQuantConfig | None = None, + a1q_scale: torch.Tensor | None = None, ) -> torch.Tensor: if quant_config is None: quant_config = FUSED_MOE_UNQUANTIZED_CONFIG @@ -518,9 +533,11 @@ def rocm_aiter_fused_experts( activation_method=activation_method, w1_scale=quant_config.w1_scale, w2_scale=quant_config.w2_scale, - a1_scale=quant_config.a1_scale, + a1_scale=quant_config.a1_scale if a1q_scale is None else a1q_scale, a2_scale=quant_config.a2_scale, doweight_stage1=apply_router_weight_on_input, + expert_num_tokens=expert_num_tokens, + output_dtype=output_dtype, ) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index e5681cb85625..6a527f899567 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -31,6 +31,9 @@ ) from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod +from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( + is_rocm_aiter_moe_enabled, +) from vllm.model_executor.layers.linear import ( LinearBase, LinearMethodBase, @@ -1089,8 +1092,7 @@ def process_weights_after_loading(self, layer: Module) -> None: def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None: if ( - self.rocm_aiter_moe_enabled - or self.use_marlin + self.use_marlin or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM ): return None @@ -1109,13 +1111,12 @@ def select_gemm_impl( layer: torch.nn.Module, ) -> FusedMoEPermuteExpertsUnpermute: from vllm.model_executor.layers.fused_moe import ( + AiterMoriExperts, BatchedTritonOrDeepGemmExperts, TritonOrDeepGemmExperts, ) - assert not self.use_marlin and not self.rocm_aiter_moe_enabled, ( - "Marlin and ROCm AITER are not supported with all2all yet." - ) + assert not self.use_marlin, "Marlin is not supported with all2all yet." assert self.moe_quant_config is not None @@ -1139,6 +1140,12 @@ def select_gemm_impl( quant_config=self.moe_quant_config, allow_deep_gemm=self.allow_deep_gemm, ) + elif self.moe.use_mori_kernels and is_rocm_aiter_moe_enabled(): + logger.debug("AiterMoriExperts for Mori integration %s", self.moe) + return AiterMoriExperts( + max_num_tokens=self.moe.max_num_tokens, + quant_config=self.moe_quant_config, + ) elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS: experts = select_cutlass_fp8_gemm_impl( self.moe, @@ -1292,13 +1299,11 @@ def apply( # can override fused_experts or cutlass but not rocm or marlin. # topk_weights, topk_ids, zero_expert_result = select_result - - if self.rocm_aiter_moe_enabled: + if self.rocm_aiter_moe_enabled and self.fused_experts is None: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501 rocm_aiter_fused_experts, ) - assert self.fused_experts is None result = rocm_aiter_fused_experts( x, layer.w13_weight, @@ -1337,7 +1342,7 @@ def apply( w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=True, + inplace=not self.moe.use_mori_kernels, activation=activation, global_num_experts=global_num_experts, apply_router_weight_on_input=apply_router_weight_on_input, diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 65f588b52e5e..5c94ca0b300f 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -347,6 +347,11 @@ def has_deep_ep() -> bool: return _has_module("deep_ep") +def has_mori() -> bool: + """Whether the optional `mori` package is available.""" + return _has_module("mori") + + def has_deep_gemm() -> bool: """Whether the optional `deep_gemm` package is available.""" return _has_module("deep_gemm")