diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index fc7fc25ff29..d52da1f0a40 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -299,3 +299,4 @@ steps: - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py + - pytest -v -s kernels/moe/test_deepep_v2_moe.py diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 875b65ff084..53fea072555 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -19,6 +19,7 @@ from vllm.utils.flashinfer import ( has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, ) +from vllm.utils.import_utils import has_deep_ep_v2 from vllm.utils.network_utils import get_open_port from ..utils import init_test_distributed_environment @@ -194,6 +195,10 @@ requires_ptrace = pytest.mark.skipif( not _has_sys_ptrace(), reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)", ) +requires_deep_ep_v2 = pytest.mark.skipif( + not has_deep_ep_v2(), + reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4", +) # NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have # their own @requires_two_sided / @requires_one_sided decorators, and @@ -856,3 +861,76 @@ def _one_sided_data_worker(rank, world_size): def test_one_sided_dispatch_combine(world_size): """Test FlashInfer one-sided dispatch/combine with actual data flow.""" _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size) + + +# --------------------------------------------------------------------------- +# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle +# --------------------------------------------------------------------------- +# +# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using +# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4. +# +# Uses EP group because the DeepEP v2 manager is constructed with an +# EP-scoped communicator in production. With tp=world_size the EP group +# spans all ranks. +# --------------------------------------------------------------------------- + + +def _deepep_v2_lifecycle_worker(rank, world_size): + from vllm.distributed.device_communicators.all2all import ( + DeepEPV2All2AllManager, + ) + + cpu_group = get_ep_group().cpu_group + manager = DeepEPV2All2AllManager(cpu_group) + + assert manager.rank == rank + assert manager.world_size == world_size + assert manager._num_sms is None + + hidden_size = 7168 + num_experts = world_size * 32 + num_topk = 8 + max_tokens = 256 + + handle_kwargs = dict( + num_max_tokens_per_rank=max_tokens, + hidden=hidden_size, + num_topk=num_topk, + num_experts=num_experts, + use_fp8_dispatch=False, + ) + + handle = manager.get_handle(handle_kwargs) + assert handle is not None + assert manager._num_sms is not None + assert manager._num_sms > 0 + + torch.distributed.barrier() + + # get_handle again with same args should return cached handle + handle2 = manager.get_handle(dict(handle_kwargs)) + assert handle2 is handle + + torch.distributed.barrier() + + # Destroy clears the cache + manager.destroy() + assert len(manager.handle_cache._cache) == 0 + + torch.distributed.barrier() + + # Re-create after destroy + handle3 = manager.get_handle(dict(handle_kwargs)) + assert handle3 is not None + + torch.distributed.barrier() + manager.destroy() + + +@requires_multi_gpu +@requires_deep_ep_v2 +@pytest.mark.parametrize("world_size", [2]) +def test_deepep_v2_manager_lifecycle(world_size): + """Test DeepEP v2 ElasticBuffer manager init, caching, and destroy.""" + _spawn_workers(_deepep_v2_lifecycle_worker, world_size) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 7a54dd4ca37..646cff4c2d2 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from types import SimpleNamespace from typing import Any import torch @@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.utils.import_utils import ( has_aiter, has_deep_ep, + has_deep_ep_v2, has_deep_gemm, has_mori, ) @@ -150,6 +152,10 @@ class Config: make env data for vllm launch. """ vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace( + enforce_eager=True, + is_moe=True, + ) vllm_config.parallel_config.data_parallel_size = self.world_size vllm_config.parallel_config.enable_expert_parallel = True @@ -239,6 +245,10 @@ class Config: or info.backend == "deepep_low_latency" ) + def needs_deep_ep_v2(self): + info = prepare_finalize_info(self.prepare_finalize_type) + return info.backend == "deepep_v2" + def needs_aiter(self): info = expert_info(self.fused_experts_type) return info.needs_aiter @@ -315,6 +325,8 @@ class Config: # Check dependencies (turn into asserts?) if self.needs_deep_ep() and not has_deep_ep(): return False, "Needs DeepEP, but DeepEP not available." + if self.needs_deep_ep_v2() and not has_deep_ep_v2(): + return False, "Needs DeepEP v2, but DeepEP v2 not available." if self.needs_deep_gemm() and not has_deep_gemm(): return False, "Needs DeepGEMM, but DeepGEMM not available." if self.needs_aiter() and not has_aiter(): # noqa: SIM103 @@ -657,6 +669,58 @@ def make_modular_kernel( return modular_kernel +def _maybe_convert_weights_for_experts( + config: Config, + rank_weights: WeightTensors, +) -> WeightTensors: + """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK).""" + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + convert_to_fp8_moe_kernel_format, + ) + + fe_type = config.fused_experts_type + fe_name = getattr(fe_type, "__name__", "") + + backend: Fp8MoeBackend | None = None + if fe_name == "TrtLlmFp8ExpertsModular": + backend = Fp8MoeBackend.FLASHINFER_TRTLLM + elif fe_name == "FlashInferExperts": + backend = Fp8MoeBackend.FLASHINFER_CUTLASS + + if backend is None or not rank_weights.is_quantized(): + return rank_weights + + mock_layer = SimpleNamespace( + weight_block_size=config.quant_block_shape, + moe_config=SimpleNamespace( + is_act_and_mul=True, + intermediate_size_per_partition=config.N, + ), + activation=SimpleNamespace(is_gated=True), + ) + + w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=backend, + layer=mock_layer, + w13=rank_weights.w1, + w2=rank_weights.w2, + w13_scale=rank_weights.w1_scale, + w2_scale=rank_weights.w2_scale, + w13_input_scale=None, + w2_input_scale=None, + ) + + return WeightTensors( + w1=w1, + w2=w2, + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_gs=rank_weights.w1_gs, + w2_gs=rank_weights.w2_gs, + ) + + def run_modular_kernel( pgi: ProcessGroupInfo, vllm_config: VllmConfig, @@ -669,6 +733,7 @@ def run_modular_kernel( # weights for rank rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts) + rank_weights = _maybe_convert_weights_for_experts(config, rank_weights) if config.quant_dtype == "nvfp4": gscale = _make_gscale(config.num_local_experts) @@ -714,6 +779,8 @@ def run_modular_kernel( [num_tokens] * config.world_size, device="cuda", dtype=torch.int ) + torch.distributed.barrier() + with set_forward_context( None, vllm_config, diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 78ee8084d90..7f1924bac5a 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -36,10 +36,12 @@ from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import ( has_flashinfer_cutlass_fused_moe, has_flashinfer_nvlink_one_sided, + has_flashinfer_trtllm_fused_moe, ) from vllm.utils.import_utils import ( has_aiter, has_deep_ep, + has_deep_ep_v2, has_deep_gemm, has_mori, ) @@ -216,6 +218,19 @@ if has_deep_ep() and not current_platform.has_device_capability(100): backend="deepep_low_latency", ) +if has_deep_ep_v2() and current_platform.has_device_capability(100): + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import ( + DeepEPV2PrepareAndFinalize, + ) + + register_prepare_and_finalize( + DeepEPV2PrepareAndFinalize, + standard_format, + common_float_types, + blocked_quantization_support=True, + backend="deepep_v2", + ) + if has_mori(): from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import ( MoriPrepareAndFinalize, @@ -289,6 +304,18 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability blocked_quantization_support=False, ) +if has_flashinfer_trtllm_fused_moe() and current_platform.has_device_capability(100): + from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( + TrtLlmFp8ExpertsModular, + ) + + register_experts( + TrtLlmFp8ExpertsModular, + standard_format, + fp8_types, + blocked_quantization_support=True, + ) + if has_aiter(): from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( AiterExperts, diff --git a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py index 61826199771..d7394d72cb3 100644 --- a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py +++ b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py @@ -106,7 +106,7 @@ def _worker_parallel_launch( if vllm_config is not None: cpu_group = _set_vllm_config(vllm_config, world_size, rank, local_rank) - try: + def _run_worker(): worker( ProcessGroupInfo( world_size=world_size, @@ -121,6 +121,13 @@ def _worker_parallel_launch( *args, **worker_kwargs, ) + + try: + if vllm_config is not None: + with set_current_vllm_config(vllm_config): + _run_worker() + else: + _run_worker() except Exception as ex: print(ex) traceback.print_exc() diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py index 1663e562966..bb2f9efc7c4 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -15,7 +15,7 @@ from torch.distributed import ProcessGroup from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage] from typing_extensions import ParamSpec -from vllm.utils.import_utils import has_deep_ep +from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2 from vllm.utils.network_utils import get_open_port if has_deep_ep(): @@ -26,6 +26,11 @@ if has_deep_ep(): DeepEPLLPrepareAndFinalize, ) +if has_deep_ep_v2(): + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import ( + DeepEPV2PrepareAndFinalize, + ) + ## Parallel Processes Utils P = ParamSpec("P") @@ -55,11 +60,10 @@ def _worker_parallel_launch( torch.accelerator.set_device_index(local_rank) device = torch.device("cuda", local_rank) torch.distributed.init_process_group( - backend="cpu:gloo,cuda:nccl", + backend="nccl", init_method=init_method, rank=rank, world_size=world_size, - device_id=device, ) barrier = torch.tensor([rank], device=device) torch.distributed.all_reduce(barrier) @@ -200,3 +204,42 @@ def make_deepep_a2a( assert deepep_ll_args is not None return make_deepep_ll_a2a(pg, pgi, deepep_ll_args, q_dtype, block_shape) + + +@dataclasses.dataclass +class DeepEPV2Args: + num_local_experts: int + num_experts: int + num_topk: int + hidden_size: int + max_tokens_per_rank: int + use_fp8_dispatch: bool + + +def make_deepep_v2_a2a( + pg: ProcessGroup, + pgi: ProcessGroupInfo, + dp_size: int, + v2_args: DeepEPV2Args, + use_cudagraph: bool = False, +): + import deep_ep + + buffer = deep_ep.ElasticBuffer( + group=pg, + num_max_tokens_per_rank=v2_args.max_tokens_per_rank, + hidden=v2_args.hidden_size, + num_topk=v2_args.num_topk, + use_fp8_dispatch=v2_args.use_fp8_dispatch, + explicitly_destroy=True, + ) + return DeepEPV2PrepareAndFinalize( + buffer=buffer, + num_dispatchers=pgi.world_size, + dp_size=dp_size, + rank_expert_offset=pgi.rank * v2_args.num_local_experts, + num_experts=v2_args.num_experts, + num_topk=v2_args.num_topk, + use_fp8_dispatch=v2_args.use_fp8_dispatch, + use_cudagraph=use_cudagraph, + ) diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py new file mode 100644 index 00000000000..93b7c136605 --- /dev/null +++ b/tests/kernels/moe/test_deepep_v2_moe.py @@ -0,0 +1,542 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Test DeepEP v2 (ElasticBuffer) dispatch-combine logic. +Compares against a pure-PyTorch reference MoE implementation. +""" + +import dataclasses + +import pytest +import torch.distributed +from torch.distributed import ProcessGroup + +from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights +from tests.kernels.utils import torch_experts +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe import TritonExperts +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel +from vllm.utils.import_utils import has_deep_ep_v2 +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.worker.workspace import init_workspace_manager + +from ...utils import multi_gpu_test +from .parallel_utils import ProcessGroupInfo, parallel_launch + +if has_deep_ep_v2(): + from .parallel_utils import DeepEPV2Args, make_deepep_v2_a2a + +requires_deep_ep_v2 = pytest.mark.skipif( + not has_deep_ep_v2(), + reason="Requires DeepEP v2 (ElasticBuffer)", +) + + +@dataclasses.dataclass +class TestConfig: + dtype: torch.dtype + topk: int + m: int + k: int + n: int + num_experts: int + + +@dataclasses.dataclass +class TestTensors: + rank_tokens: torch.Tensor + rank_token_scales: torch.Tensor | None + topk: torch.Tensor + topk_weights: torch.Tensor + config: TestConfig + + @staticmethod + def make(config: TestConfig) -> "TestTensors": + assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn] + token_dtype = ( + torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype + ) + rank_tokens = ( + torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10 + ) + + topk = torch.stack( + [ + torch.randperm(config.num_experts, device="cuda")[: config.topk] + for _ in range(config.m) + ] + ).to(dtype=torch.int64) + topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda") + return TestTensors( + rank_tokens=rank_tokens, + rank_token_scales=None, + topk=topk, + topk_weights=topk_weights, + config=config, + ) + + +def make_modular_kernel( + pg: ProcessGroup, + pgi: ProcessGroupInfo, + dp_size: int, + hidden_size: int, + num_experts: int, + num_local_experts: int, + topk: int, + q_dtype: torch.dtype | None, + use_fp8_dispatch: bool, + quant_config: FusedMoEQuantConfig, + use_cudagraph: bool = False, +) -> FusedMoEKernel: + v2_args = DeepEPV2Args( + num_local_experts=num_local_experts, + num_experts=num_experts, + num_topk=topk, + hidden_size=hidden_size, + max_tokens_per_rank=8192, + use_fp8_dispatch=use_fp8_dispatch, + ) + + a2a = make_deepep_v2_a2a( + pg=pg, + pgi=pgi, + dp_size=dp_size, + v2_args=v2_args, + use_cudagraph=use_cudagraph, + ) + + moe_config = make_dummy_moe_config( + num_experts=num_local_experts, + experts_per_token=topk, + hidden_dim=hidden_size, + ) + + fused_experts = TritonExperts( + moe_config=moe_config, + quant_config=quant_config, + ) + + mk = FusedMoEKernel( + prepare_finalize=a2a, + fused_experts=fused_experts, + inplace=False, + ) + return mk + + +def deepep_v2_moe_impl( + pg: ProcessGroup, + pgi: ProcessGroupInfo, + dp_size: int, + test_tensors: TestTensors, + w1: torch.Tensor, + w2: torch.Tensor, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + num_experts: int, + topk: int, + use_fp8_dispatch: bool, + per_act_token_quant: bool, +) -> torch.Tensor: + num_local_experts = w1.size(0) + + def build_expert_map(): + expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32) + s = pgi.rank * num_local_experts + e = s + num_local_experts + expert_map[s:e] = torch.tensor(list(range(num_local_experts))) + device = torch.accelerator.current_device_index() + return expert_map.to(device=device, dtype=torch.int32) + + is_quantized = w1.dtype == torch.float8_e4m3fn + q_dtype = torch.float8_e4m3fn if is_quantized else None + + quant_config = FusedMoEQuantConfig.make( + q_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + per_act_token_quant=per_act_token_quant, + a1_scale=test_tensors.rank_token_scales, + ) + + hidden_size = test_tensors.rank_tokens.size(1) + + mk: FusedMoEKernel = make_modular_kernel( + pg, + pgi, + dp_size, + hidden_size, + num_experts, + num_local_experts, + topk, + q_dtype, + use_fp8_dispatch, + quant_config, + ) + + out = mk.apply( + hidden_states=test_tensors.rank_tokens, + w1=w1, + w2=w2, + topk_weights=test_tensors.topk_weights, + topk_ids=test_tensors.topk, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=build_expert_map(), + apply_router_weight_on_input=False, + ) + + return out + + +def _deep_ep_v2_moe( + pgi: ProcessGroupInfo, + dp_size: int, + config: TestConfig, + w1: torch.Tensor, + w2: torch.Tensor, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + use_fp8_dispatch: bool, + per_act_token_quant: bool, +): + device = torch.device(f"cuda:{pgi.local_rank}") + init_workspace_manager(device) + + is_quantized = w1.dtype == torch.float8_e4m3fn + device_idx = torch.accelerator.current_device_index() + w1 = w1.to(device=device_idx) + w2 = w2.to(device=device_idx) + if is_quantized: + assert w1_scale is not None and w2_scale is not None + w1_scale = w1_scale.to(device=device_idx) + w2_scale = w2_scale.to(device=device_idx) + + pg = torch.distributed.new_group(list(range(pgi.world_size))) + test_tensors = TestTensors.make(config) + + with set_current_vllm_config(VllmConfig()): + # Reference + q_dtype = torch.float8_e4m3fn if is_quantized else None + torch_combined = torch_experts( + test_tensors.rank_tokens, + w1, + w2, + test_tensors.topk_weights, + test_tensors.topk, + w1_scale=w1_scale, + w2_scale=w2_scale, + quant_dtype=q_dtype, + per_act_token_quant=per_act_token_quant, + ) + + # Splice experts for this rank + num_local_experts = config.num_experts // pgi.world_size + e_start = num_local_experts * pgi.rank + e_end = e_start + num_local_experts + w1_ep = w1[e_start:e_end] + w2_ep = w2[e_start:e_end] + + w1_scale_ep, w2_scale_ep = None, None + if is_quantized: + w1_scale_ep = w1_scale[e_start:e_end] # type: ignore + w2_scale_ep = w2_scale[e_start:e_end] # type: ignore + + deepep_combined = deepep_v2_moe_impl( + pg, + pgi, + dp_size, + test_tensors, + w1_ep, + w2_ep, + w1_scale_ep, + w2_scale_ep, + config.num_experts, + config.topk, + use_fp8_dispatch, + per_act_token_quant, + ) + + torch.testing.assert_close( + torch_combined, + deepep_combined, + atol=6e-2, + rtol=6e-2, + ) + + +MNKs = [ + (1, 256, 256), + (2, 256, 512), + (3, 1024, 2048), + (32, 256, 1024), + (45, 512, 2048), + (64, 1024, 1024), + (222, 1024, 2048), +] + +DTYPES = [torch.bfloat16, torch.float8_e4m3fn] + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("m,n,k", MNKs) +@pytest.mark.parametrize("num_experts", [32]) +@pytest.mark.parametrize("topk", [6]) +@pytest.mark.parametrize("world_dp_size", [(2, 1)]) +@multi_gpu_test(num_gpus=2) +@requires_deep_ep_v2 +def test_deep_ep_v2_moe( + dtype: torch.dtype, + m: int, + n: int, + k: int, + num_experts: int, + topk: int, + world_dp_size: tuple[int, int], + workspace_init, +): + per_act_token_quant = False + use_fp8_dispatch = False + + set_random_seed(7) + world_size, dp_size = world_dp_size + config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts) + + quant_dtype = dtype if dtype == torch.float8_e4m3fn else None + (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights( + num_experts, + n, + k, + quant_dtype=quant_dtype, + per_out_ch_quant=True, + ) + + parallel_launch( + world_size, + _deep_ep_v2_moe, + dp_size, + config, + w1, + w2, + w1_scale, + w2_scale, + use_fp8_dispatch, + per_act_token_quant, + ) + + +def _deep_ep_v2_moe_cudagraph( + pgi: ProcessGroupInfo, + dp_size: int, + config: TestConfig, + w1: torch.Tensor, + w2: torch.Tensor, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, +): + """Worker function: verify DeepEP v2 + TrtLLM FP8 with do_expand=False.""" + import tempfile + + from vllm.distributed import ( + init_distributed_environment, + initialize_model_parallel, + ) + + device = torch.device(f"cuda:{pgi.local_rank}") + init_workspace_manager(device) + + pg = torch.distributed.new_group(list(range(pgi.world_size))) + test_tensors = TestTensors.make(config) + num_local_experts = config.num_experts // pgi.world_size + hidden_size = config.k + + # Create FP8 weights directly, then dequantize for bf16 reference. + w1_fp8 = torch.randn( + (config.num_experts, 2 * config.n, config.k), + device="cuda", + dtype=torch.bfloat16, + ).to(torch.float8_e4m3fn) + w2_fp8 = torch.randn( + (config.num_experts, config.k, config.n), + device="cuda", + dtype=torch.bfloat16, + ).to(torch.float8_e4m3fn) + w1_ref = w1_fp8.to(torch.bfloat16) + w2_ref = w2_fp8.to(torch.bfloat16) + + from vllm.config import KernelConfig + + vllm_cfg = VllmConfig() + vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm") + + with set_current_vllm_config(vllm_cfg): + # Initialize vLLM parallel state (needed by FusedMoE layer) + temp_file = tempfile.mktemp() + init_distributed_environment( + world_size=pgi.world_size, + rank=pgi.rank, + distributed_init_method=f"file://{temp_file}", + local_rank=pgi.local_rank, + backend="nccl", + ) + initialize_model_parallel(tensor_model_parallel_size=1) + # Reference MoE using dequantized bf16 weights + torch_combined = torch_experts( + test_tensors.rank_tokens, + w1_ref, + w2_ref, + test_tensors.topk_weights, + test_tensors.topk, + ) + + # Use the production pipeline: make_fused_moe_layer creates + # a FusedMoE layer, quantizes weights, runs + # process_weights_after_loading (TrtLLM W31 swap + BlockMajorK + # shuffle), and selects the kernel. + # Quantize weights using production helper, EP-slice, then + # convert to TrtLLM format. + from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves + from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( + TrtLlmFp8ExpertsModular, + ) + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + convert_to_fp8_moe_kernel_format, + ) + + block_shape = [128, 128] + qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape) + + # EP-slice before format conversion + e_start = num_local_experts * pgi.rank + e_end = e_start + num_local_experts + w1_ep = qw.w13_weight[e_start:e_end] + w2_ep = qw.w2_weight[e_start:e_end] + assert qw.w13_weight_scale is not None + assert qw.w2_weight_scale is not None + w1_scale_ep = qw.w13_weight_scale[e_start:e_end] + w2_scale_ep = qw.w2_weight_scale[e_start:e_end] + + # Convert to TrtLLM format (W31 swap + BlockMajorK shuffle) + class _MockLayer: + weight_block_size = block_shape + + class moe_config: + is_act_and_mul = True + intermediate_size_per_partition = config.n + + class activation: + is_gated = True + + w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = convert_to_fp8_moe_kernel_format( + fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM, + layer=_MockLayer(), + w13=w1_ep, + w2=w2_ep, + w13_scale=w1_scale_ep, + w2_scale=w2_scale_ep, + w13_input_scale=None, + w2_input_scale=None, + ) + + # Build TrtLLM expert with correct EP params + quant_config = FusedMoEQuantConfig.make( + torch.float8_e4m3fn, + block_shape=block_shape, + w1_scale=w1_scale_ep, + w2_scale=w2_scale_ep, + ) + moe_config = make_dummy_moe_config( + num_experts=num_local_experts, + experts_per_token=config.topk, + hidden_dim=hidden_size, + intermediate_size=config.n, + ) + fused_experts = TrtLlmFp8ExpertsModular( + moe_config=moe_config, + quant_config=quant_config, + ) + + v2_args = DeepEPV2Args( + num_local_experts=num_local_experts, + num_experts=config.num_experts, + num_topk=config.topk, + hidden_size=hidden_size, + max_tokens_per_rank=8192, + use_fp8_dispatch=False, + ) + a2a = make_deepep_v2_a2a( + pg=pg, + pgi=pgi, + dp_size=dp_size, + v2_args=v2_args, + use_cudagraph=True, + ) + mk_kernel = FusedMoEKernel( + prepare_finalize=a2a, + fused_experts=fused_experts, + inplace=False, + ) + + for _ in range(3): + out = mk_kernel.apply( + hidden_states=test_tensors.rank_tokens, + w1=w1_ep, + w2=w2_ep, + topk_weights=test_tensors.topk_weights, + topk_ids=test_tensors.topk, + activation=MoEActivation.SILU, + global_num_experts=config.num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + + torch.testing.assert_close( + torch_combined, + out, + atol=6e-2, + rtol=6e-2, + ) + + +@pytest.mark.parametrize("m,n,k", [(32, 256, 1024)]) +@pytest.mark.parametrize("num_experts", [32]) +@pytest.mark.parametrize("topk", [6]) +@pytest.mark.parametrize("world_dp_size", [(2, 1)]) +@multi_gpu_test(num_gpus=2) +@requires_deep_ep_v2 +def test_deep_ep_v2_moe_cudagraph( + m: int, + n: int, + k: int, + num_experts: int, + topk: int, + world_dp_size: tuple[int, int], + workspace_init, +): + set_random_seed(7) + world_size, dp_size = world_dp_size + config = TestConfig( + dtype=torch.float8_e4m3fn, + topk=topk, + m=m, + k=k, + n=n, + num_experts=num_experts, + ) + + parallel_launch( + world_size, + _deep_ep_v2_moe_cudagraph, + dp_size, + config, + None, # weights created inside worker + None, + None, + None, + ) diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index f61aa868581..94beef897d0 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -8,7 +8,8 @@ set -ex # --nvshmem-ver NVSHMEM version CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} -DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"} +DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"d4f41e4e93"} + NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"} # Default supports both CUDA 12 and 13 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace} MODE=${MODE:-install} diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index f7fb8ec42a3..2904f40f8a4 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -42,6 +42,7 @@ All2AllBackend = Literal[ "pplx", "deepep_high_throughput", "deepep_low_latency", + "deepep_v2", "mori_high_throughput", "mori_low_latency", "nixl_ep", diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index c0055f4bbe1..fd1c826322c 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -15,7 +15,7 @@ from vllm.utils.flashinfer import ( has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, ) -from vllm.utils.import_utils import has_deep_ep, has_mori +from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori from .base_device_communicator import All2AllManagerBase, Cache @@ -863,3 +863,66 @@ class MoriAll2AllManager(All2AllManagerBase): mori_kwargs, self._make_handle ) return handle + + +class DeepEPV2All2AllManager(All2AllManagerBase): + """ + All2All communication based on DeepEP v2 ElasticBuffer (unified API). + Uses NCCL Gin backend with analytical SM calculation. + """ + + def __init__(self, cpu_group, tcp_store_group=None, device_group=None): + assert has_deep_ep_v2(), ( + "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 " + "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4." + ) + super().__init__(cpu_group, tcp_store_group) + self._device_group = device_group + self.handle_cache = Cache() + self._num_sms: int | None = None + + def _make_all2all_kwargs( + self, + num_max_tokens_per_rank: int, + hidden: int, + num_topk: int, + use_fp8_dispatch: bool, + ) -> dict: + return dict( + group=self._device_group + if self._device_group is not None + else self.cpu_group, + num_max_tokens_per_rank=num_max_tokens_per_rank, + hidden=hidden, + num_topk=num_topk, + use_fp8_dispatch=use_fp8_dispatch, + allow_hybrid_mode=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE, + prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP, + allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION), + explicitly_destroy=True, + ) + + def get_handle(self, kwargs): + import deep_ep # type: ignore[import-not-found] + + num_experts = kwargs.pop("num_experts", 256) + buffer_kwargs = self._make_all2all_kwargs(**kwargs) + logger.debug("DeepEP v2 all2all args %s", buffer_kwargs) + handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create( + buffer_kwargs, deep_ep.ElasticBuffer + ) + if self._num_sms is None: + self._num_sms = handle.get_theoretical_num_sms( + num_experts=num_experts, + num_topk=kwargs["num_topk"], + ) + return handle + + def max_sms_used(self) -> int | None: + return self._num_sms + + def destroy(self): + with self.handle_cache._lock: + for _, handle in self.handle_cache._cache.items(): + handle.destroy() + self.handle_cache._cache.clear() diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index ff4929d8415..12c425021b2 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -146,6 +146,14 @@ class CudaCommunicator(DeviceCommunicatorBase): self.all2all_manager = MoriAll2AllManager( self.cpu_group, self.all2all_backend ) + elif self.all2all_backend == "deepep_v2": + from .all2all import DeepEPV2All2AllManager + + self.all2all_manager = DeepEPV2All2AllManager( + self.cpu_group, + tcp_store_group, + device_group=self.device_group, + ) elif self.all2all_backend == "nixl_ep": from .all2all import NixlEPAll2AllManager diff --git a/vllm/envs.py b/vllm/envs.py index 8f4e18d2235..a264b3129df 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -246,6 +246,9 @@ if TYPE_CHECKING: VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024 VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False + VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE: bool = True + VLLM_DEEPEP_V2_PREFER_OVERLAP: bool = False + VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION: bool = False VLLM_DBO_COMM_SMS: int = 20 VLLM_PATTERN_MATCH_DEBUG: str | None = None VLLM_DEBUG_DUMP_PATH: str | None = None @@ -1840,6 +1843,18 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool( int(os.getenv("VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL", "0")) ), + # DeepEP v2: enable two-tier NVLink+RDMA hybrid mode + "VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE": lambda: bool( + int(os.getenv("VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE", "0")) + ), + # DeepEP v2: use fewer SMs at slight throughput cost + "VLLM_DEEPEP_V2_PREFER_OVERLAP": lambda: bool( + int(os.getenv("VLLM_DEEPEP_V2_PREFER_OVERLAP", "0")) + ), + # DeepEP v2: trade precision for transfer size in combine + "VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION": lambda: bool( + int(os.getenv("VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION", "0")) + ), # The number of SMs/CUs to allocate for communication kernels when # running DBO; the rest will be allocated to compute. # Default: 20 on CUDA (SMs), 64 on ROCm (CUs). diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 6d482214643..1351e87b5b5 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -29,7 +29,12 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two FlashInferNVLinkTwoSidedPrepareAndFinalize, ) from vllm.platforms import current_platform -from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep +from vllm.utils.import_utils import ( + has_deep_ep, + has_deep_ep_v2, + has_mori, + has_nixl_ep, +) logger = init_logger(__name__) @@ -40,6 +45,8 @@ if current_platform.is_cuda_alike(): DEEPEP_QUANT_BLOCK_SHAPE, DeepEPLLPrepareAndFinalize, ) + if has_deep_ep_v2(): + from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize if has_mori(): from .prepare_finalize.mori import MoriPrepareAndFinalize if has_nixl_ep(): @@ -94,6 +101,11 @@ def maybe_roundup_layer_hidden_size( hidden_size ) + if moe_parallel_config.use_deepep_v2_kernels: + hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size( + hidden_size, act_dtype + ) + if moe_parallel_config.use_nixl_ep_kernels: hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size( hidden_size @@ -193,6 +205,36 @@ def maybe_make_prepare_finalize( physical_to_global=physical_to_global, local_expert_global_ids=local_expert_global_ids, ) + elif moe.use_deepep_v2_kernels: + assert moe.dp_size == all2all_manager.dp_world_size + + use_fp8_dispatch = ( + quant_config is not None + and quant_config.quant_dtype == current_platform.fp8_dtype() + and quant_config.is_block_quantized + ) + all_to_all_args = dict( + num_max_tokens_per_rank=moe.max_num_tokens, + hidden=moe.hidden_dim, + num_topk=moe.experts_per_token, + num_experts=moe.num_experts, + use_fp8_dispatch=use_fp8_dispatch, + ) + handle = all2all_manager.get_handle(all_to_all_args) + vllm_config = get_current_vllm_config() + use_cudagraph = not vllm_config.model_config.enforce_eager + + prepare_finalize = DeepEPV2PrepareAndFinalize( + buffer=handle, + num_dispatchers=all2all_manager.world_size, + dp_size=all2all_manager.dp_world_size, + rank_expert_offset=all2all_manager.rank * moe.num_local_experts, + num_experts=moe.num_experts, + num_topk=moe.experts_per_token, + use_fp8_dispatch=use_fp8_dispatch, + use_cudagraph=use_cudagraph, + ) + elif moe.use_mori_kernels: assert quant_config is not None diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b289b85f17a..1b063559b8d 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1070,6 +1070,10 @@ class FusedMoEParallelConfig: def use_nixl_ep_kernels(self): return self.use_all2all_kernels and self.all2all_backend == "nixl_ep" + @property + def use_deepep_v2_kernels(self): + return self.use_all2all_kernels and self.all2all_backend == "deepep_v2" + @staticmethod def flatten_tp_across_dp_and_pcp( tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int @@ -1394,6 +1398,10 @@ class FusedMoEConfig: def use_nixl_ep_kernels(self): return self.moe_parallel_config.use_nixl_ep_kernels + @property + def use_deepep_v2_kernels(self): + return self.moe_parallel_config.use_deepep_v2_kernels + @property def needs_round_robin_routing_tables(self): return self.moe_parallel_config.needs_round_robin_routing_tables diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 9230fea6e5c..257bfeee5d3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -94,6 +94,14 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): Fp8 TRTLLM-Gen MoE kernels. Supports modular interface. """ + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + or moe_parallel_config.use_ag_rs_all2all_kernels + or moe_parallel_config.use_deepep_v2_kernels + ) and not moe_parallel_config.enable_eplb + @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, @@ -193,7 +201,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, - top_k=self.topk, + top_k=topk_ids.size(1), n_group=None, topk_group=None, intermediate_size=self.intermediate_size_per_partition, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index ef740667f1a..1099368e107 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -85,6 +85,18 @@ def _get_priority_backends( def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None: backends.insert(0, backends.pop(backends.index(backend))) + # With DeepEP v2 contiguous layout (do_expand=False), tensors are + # worst-case allocated with padding. TrtLLM's tile-level skipping + # avoids wasted compute on padding rows; other backends process all rows. + if ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and moe_config.moe_parallel_config.use_deepep_v2_kernels + and activation_key == kFp8Dynamic128Sym + and weight_key == kFp8Static128BlockSym + ): + _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.FLASHINFER_TRTLLM) + # On Hopper for Block Fp8, prefer Triton for TP and FI CUTLASS for EP. if ( current_platform.is_cuda() diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py new file mode 100644 index 00000000000..6495e1203e0 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import deep_ep +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.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + TopKWeightAndReduceDelegate, +) +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.utils.math_utils import round_up +from vllm.v1.worker.ubatching import ( + dbo_current_ubatch_id, +) + + +class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): + """ + Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API). + + Supports two modes controlled by the `use_cudagraph` constructor arg: + + **Decode mode (use_cudagraph=True):** + - do_expand=False, do_cpu_sync=False + - Tokens returned in original order with recv_topk_idx (global IDs) + - Worst-case tensor allocation; padding rows zeroed via + handle.psum_num_recv_tokens_per_scaleup_rank + - Fully cudagraph-capturable + - Expert kernel sorts internally (expert_tokens_meta=None) + + **Prefill mode (use_cudagraph=False):** + - do_expand=True, do_cpu_sync=True + - Per-expert-contiguous layout; exact memory allocation + - Saves GPU memory (no worst-case allocation) + - Not cudagraph-capturable (CPU polling), but prefill doesn't + use cudagraphs anyway + - Provides expert_tokens_meta for efficient batched expert kernels + + Both modes use async_with_compute_stream=False (synchronous from + caller's perspective). The ElasticBuffer handles comm internally. + """ + + @staticmethod + def maybe_roundup_layer_hidden_size(hidden_size: int, dtype: torch.dtype) -> int: + hidden_size_bytes = hidden_size * dtype.itemsize + xfer_atom_size = 512 # 32 * 16 (size(int4)) + if hidden_size_bytes % xfer_atom_size == 0: + return hidden_size + + hidden_size_bytes = round_up(hidden_size_bytes, xfer_atom_size) + return hidden_size_bytes // dtype.itemsize + + def __init__( + self, + buffer: deep_ep.ElasticBuffer, + num_dispatchers: int, + dp_size: int, + rank_expert_offset: int, + num_experts: int, + num_topk: int, + use_fp8_dispatch: bool = False, + use_cudagraph: bool = False, + ): + super().__init__() + self.buffer = buffer + self.num_dispatchers_ = num_dispatchers + self.dp_size = dp_size + self.rank_expert_offset = rank_expert_offset + self.num_experts = num_experts + self.num_topk = num_topk + self.use_fp8_dispatch = use_fp8_dispatch + self.use_cudagraph = use_cudagraph + + # DBO microbatching: one handle slot per micro-batch. + self.handles: list[deep_ep.EPHandle | None] = [None, None] + + def num_dispatchers(self) -> int: + return self.num_dispatchers_ + + def output_is_reduced(self) -> bool: + return True + + @property + def activation_format(self) -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def max_num_tokens_per_rank(self) -> int | None: + return None + + def topk_indices_dtype(self) -> torch.dtype | None: + return torch.int64 + + def _do_dispatch( + self, + tokens: torch.Tensor, + token_scales: torch.Tensor | None, + rank_topk_ids: torch.Tensor, + rank_topk_weights: torch.Tensor, + num_experts: int, + a1_scale: torch.Tensor | None, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool, + ) -> Callable: + has_scales = token_scales is not None + + token_data = tokens + if has_scales: + token_data = (tokens, token_scales) + + # Decode: do_expand=False + do_cpu_sync=False (cudagraph-safe) + # Prefill: do_expand=True + do_cpu_sync=True (memory-efficient) + do_expand = not self.use_cudagraph + do_cpu_sync = not self.use_cudagraph + + ( + recv_x, + recv_topk_idx, + recv_topk_weights, + handle, + event, + ) = self.buffer.dispatch( + x=token_data, + topk_idx=rank_topk_ids, + topk_weights=rank_topk_weights, + num_experts=num_experts, + do_expand=do_expand, + do_cpu_sync=do_cpu_sync, + async_with_compute_stream=False, + ) + + a2a_idx = dbo_current_ubatch_id() + self.handles[a2a_idx] = handle + + return lambda: self._receiver( + event, + has_scales, + recv_x, + recv_topk_idx, + num_experts, + handle.num_recv_tokens_per_expert_list, + recv_topk_weights, + handle.psum_num_recv_tokens_per_scaleup_rank, + a1_scale, + quant_config, + defer_input_quant=defer_input_quant, + ) + + def _receiver( + self, + event: deep_ep.EventOverlap, + has_scales: bool, + recv_x: tuple[torch.Tensor, torch.Tensor] | torch.Tensor, + recv_topk_idx: torch.Tensor | None, + num_experts: int, + recv_expert_num_tokens: list[int], + recv_topk_weights: torch.Tensor | None, + psum_recv_per_rank: torch.Tensor, + a1_scale: torch.Tensor | None, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool, + ) -> mk.PrepareResultType: + if event.event is not None: + event.current_stream_wait() + + if isinstance(recv_x, tuple): + expert_x, expert_x_scale = recv_x + else: + expert_x, expert_x_scale = recv_x, None + + if recv_topk_idx is None: + # do_expand=True (prefill mode): build topk_ids from + # per-expert token counts. + total_tokens = sum(recv_expert_num_tokens) + if total_tokens > 0: + recv_topk_idx = torch.empty( + total_tokens, + dtype=torch.int64, + device=expert_x.device, + ) + offset = 0 + for i, count in enumerate(recv_expert_num_tokens): + if count > 0: + recv_topk_idx[offset : offset + count].fill_( + i + self.rank_expert_offset + ) + offset += count + else: + recv_topk_idx = torch.empty( + 0, + dtype=torch.int64, + device=expert_x.device, + ) + recv_topk_idx = recv_topk_idx.unsqueeze(1) + else: + # do_expand=False (decode/cudagraph mode): recv_topk_idx has + # LOCAL expert IDs (-1 for non-local and padding rows). + # Convert valid local IDs to global. Rows with -1 are + # skipped by expert kernels (TrtLLM tile-level skipping, + # DeepGemm is_computation_valid), so no need to zero + # hidden states, scales, or weights for padding rows. + valid_mask = recv_topk_idx >= 0 + recv_topk_idx = torch.where( + valid_mask, + recv_topk_idx + self.rank_expert_offset, + recv_topk_idx, + ) + + # Reshape recv_topk_weights to match recv_topk_idx shape [N, 1] + if recv_topk_weights is not None and recv_topk_weights.ndim == 1: + recv_topk_weights = recv_topk_weights.unsqueeze(1) + + expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( + recv_expert_num_tokens, + device=expert_x.device, + ) + + if not quant_config.is_block_quantized and not defer_input_quant: + expert_x_scale = None + if expert_x.numel() != 0: + expert_x, expert_x_scale = moe_kernel_quantize_input( + expert_x, + a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=False, + block_shape=quant_config.block_shape, + is_scale_swizzled=quant_config.is_scale_swizzled, + ) + + return ( + expert_x, + expert_x_scale, + expert_tokens_meta, + recv_topk_idx, + recv_topk_weights, + ) + + def supports_async(self) -> bool: + return True + + def prepare_async( + 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, + defer_input_quant: bool = False, + ) -> mk.ReceiverType: + if apply_router_weight_on_input: + topk = topk_ids.size(1) + assert topk == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + a1 = a1 * topk_weights.to(a1.dtype) + + if quant_config.is_block_quantized and not defer_input_quant: + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + quant_config.a1_scale, + quant_dtype=quant_config.quant_dtype, + per_act_token_quant=quant_config.per_act_token_quant, + block_shape=quant_config.block_shape, + ) + if a1q_scale is not None and a1q_scale.numel() == 1: + a1q_scale = a1q_scale.view(1, 1) + a1_post_scale = None + else: + a1q = a1 + a1q_scale = None + a1_post_scale = ( + quant_config.a1_gscale + if quant_config.quant_dtype == "nvfp4" + else quant_config.a1_scale + ) + + return self._do_dispatch( + tokens=a1q, + token_scales=a1q_scale, + rank_topk_ids=topk_ids, + rank_topk_weights=topk_weights, + num_experts=num_experts, + a1_scale=a1_post_scale, + quant_config=quant_config, + defer_input_quant=defer_input_quant, + ) + + 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, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + receiver = self.prepare_async( + a1, + topk_weights, + topk_ids, + num_experts, + expert_map, + apply_router_weight_on_input, + quant_config, + defer_input_quant, + ) + return receiver() + + 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, + do_async: bool, + ) -> Callable | None: + a2a_idx = dbo_current_ubatch_id() + handle = self.handles[a2a_idx] + assert handle is not None + + if fused_expert_output.numel() != 0: + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): + weight_and_reduce_impl = TopKWeightAndReduceContiguous() + fused_expert_output = weight_and_reduce_impl.apply( + output=None, + fused_expert_output=fused_expert_output, + topk_weights=topk_weights, + topk_ids=topk_ids, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + if fused_expert_output.dtype != torch.bfloat16: + raise ValueError( + f"DeepEP v2 combine requires bfloat16 input, " + f"got {fused_expert_output.dtype}" + ) + + combined_x, _, event = self.buffer.combine( + x=fused_expert_output, + handle=handle, + topk_weights=None, + async_with_compute_stream=False, + ) + + output.copy_(combined_x, non_blocking=True) + return None + + def finalize_async( + 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, + ) -> Callable: + self._finalize( + output, + fused_expert_output, + topk_weights, + topk_ids, + apply_router_weight_on_input, + weight_and_reduce_impl, + False, + ) + return lambda: None + + 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, + ) -> None: + self._finalize( + output, + fused_expert_output, + topk_weights, + topk_ids, + apply_router_weight_on_input, + weight_and_reduce_impl, + False, + ) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index b5643f0c022..00540192d2f 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -1038,6 +1038,7 @@ class RoutedExperts(PluggableLayer): "e_score_correction_bias", "w13_input_scale", "w2_input_scale", + "hash_indices_table", } # Parameters of non-expert submodules that live inside runner (RoutedExperts). diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index 78e2a3c5cad..44b5e0237af 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -168,6 +168,14 @@ def fused_topk_bias( hash_indices_table: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, ): + # The topk kernel dispatches dtype based on topk_ids (set by + # indices_type) and assumes input_tokens/hash_indices_table match. + if indices_type is not None: + if input_tokens is not None and input_tokens.dtype != indices_type: + input_tokens = input_tokens.to(dtype=indices_type) + if hash_indices_table is not None and hash_indices_table.dtype != indices_type: + hash_indices_table = hash_indices_table.to(dtype=indices_type) + if not rocm_aiter_ops.is_fused_moe_enabled(): assert hidden_states.size(0) == gating_output.size(0), ( "Number of tokens mismatch" diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index c138901b6fd..1b2a8a74bdc 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -489,6 +489,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.w13_precision_config = None self.w2_precision_config = None + @property + def supports_eplb(self) -> bool: + return True + @property def skip_forward_padding(self) -> bool: # SM100_FI_MXFP4_MXFP8_TRTLLM supports padding with mxfp8 quant diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 48207a7be9a..364754f9d77 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -566,7 +566,7 @@ class DeepseekV4MoE(nn.Module): if self.use_mega_moe: self._init_mega_moe_experts(vllm_config, config, prefix) else: - self._init_fused_moe_experts(config, quant_config, prefix) + self._init_fused_moe_experts(vllm_config, config, quant_config, prefix) def _init_mega_moe_experts( self, @@ -612,22 +612,27 @@ class DeepseekV4MoE(nn.Module): def _init_fused_moe_experts( self, + vllm_config: VllmConfig, config, quant_config, prefix: str, ) -> None: + parallel_config = vllm_config.parallel_config self.tp_rank = get_tensor_model_parallel_rank() - assert config.n_routed_experts % self.tp_size == 0 - self.n_local_experts = config.n_routed_experts // self.tp_size - self.experts_start_idx = self.tp_rank * self.n_local_experts - self.experts_end_idx = self.experts_start_idx + self.n_local_experts - - self.n_redundant_experts = 0 + eplb_config = parallel_config.eplb_config + self.n_redundant_experts = eplb_config.num_redundant_experts self.n_shared_experts = config.n_shared_experts or 0 self.n_logical_experts = self.n_routed_experts - self.n_physical_experts = self.n_logical_experts - self.n_local_physical_experts = self.n_local_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + assert self.n_physical_experts % self.tp_size == 0, ( + f"n_physical_experts={self.n_physical_experts} must be divisible by " + f"tp_size={self.tp_size}. Adjust num_redundant_experts." + ) + self.n_local_physical_experts = self.n_physical_experts // self.tp_size + self.n_local_experts = self.n_local_physical_experts + self.experts_start_idx = self.tp_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts self.physical_expert_start = self.experts_start_idx self.physical_expert_end = self.experts_end_idx @@ -647,6 +652,8 @@ class DeepseekV4MoE(nn.Module): hash_indices_table=self.gate.tid2eid, swiglu_limit=self.swiglu_limit, router_logits_dtype=torch.float32, + enable_eplb=parallel_config.enable_eplb, + num_redundant_experts=eplb_config.num_redundant_experts, ) def forward( diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 63723a8f9ef..043798a584b 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -417,6 +417,61 @@ def has_deep_ep() -> bool: return _has_module("deep_ep") +DEEPEP_V2_MIN_NCCL_VERSION_RAW = 23004 # 2.30.4 + + +def _get_runtime_nccl_version() -> int | None: + """Get the runtime NCCL version by loading the actual library. + + Returns the raw version int (e.g. 23004 for 2.30.4), or None on failure. + torch.cuda.nccl.version() is a compile-time constant from the PyTorch + wheel and does not reflect a separately installed NCCL. + """ + import ctypes + + try: + from vllm.utils.nccl import find_nccl_library + + lib = ctypes.CDLL(find_nccl_library()) + version = ctypes.c_int() + lib.ncclGetVersion(ctypes.byref(version)) + return version.value + except Exception: + return None + + +def _format_nccl_raw_version(raw: int) -> str: + s = str(raw) + return f"{s[0]}.{s[1:3].lstrip('0') or '0'}.{s[3:].lstrip('0') or '0'}" + + +def has_deep_ep_v2() -> bool: + """Whether deep_ep with ElasticBuffer (v2 API) is available. + + Requires both the ElasticBuffer class in the deep_ep module and + NCCL >= 2.30.4 (GIN backend), checked against the runtime library. + """ + if not _has_module("deep_ep"): + return False + import deep_ep # type: ignore[import-not-found] + + if not hasattr(deep_ep, "ElasticBuffer"): + return False + try: + nccl_ver = _get_runtime_nccl_version() + if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW: + logger.info_once( + "DeepEP v2 requires NCCL >= %s but found %s. " + "deepep_v2 backend will not be available.", + _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW), + _format_nccl_raw_version(nccl_ver) if nccl_ver else "unknown", + ) + return False + except Exception: + return False + return True + + def has_deep_gemm() -> bool: """Whether the optional `deep_gemm` package is available.