From 8b3827ea33e01cf8da50e2558339e47743989925 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Tue, 9 Jun 2026 05:05:10 -0700 Subject: [PATCH] Revert "[WideEP] Integrate DeepEP v2 (#41183)" This reverts commit e2f993dc4116ba96ea10ec40b7bd2ee45b27b2b8. --- .buildkite/test_areas/kernels.yaml | 1 - tests/distributed/test_mnnvl_alltoall.py | 78 --- .../moe/modular_kernel_tools/common.py | 67 --- .../moe/modular_kernel_tools/mk_objects.py | 27 - .../modular_kernel_tools/parallel_utils.py | 9 +- tests/kernels/moe/parallel_utils.py | 49 +- tests/kernels/moe/test_deepep_v2_moe.py | 542 ------------------ tools/ep_kernels/install_python_libraries.sh | 3 +- vllm/config/parallel.py | 1 - .../device_communicators/all2all.py | 65 +-- .../device_communicators/cuda_communicator.py | 8 - vllm/envs.py | 15 - .../layers/fused_moe/all2all_utils.py | 44 +- .../model_executor/layers/fused_moe/config.py | 8 - .../fused_moe/experts/trtllm_fp8_moe.py | 10 +- .../layers/fused_moe/oracle/fp8.py | 12 - .../fused_moe/prepare_finalize/deepep_v2.py | 394 ------------- .../layers/fused_moe/routed_experts.py | 1 - .../router/fused_topk_bias_router.py | 8 - .../layers/quantization/mxfp4.py | 4 - vllm/models/deepseek_v4/nvidia/model.py | 25 +- vllm/utils/import_utils.py | 55 -- 22 files changed, 17 insertions(+), 1409 deletions(-) delete mode 100644 tests/kernels/moe/test_deepep_v2_moe.py delete mode 100644 vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index d52da1f0a40d..fc7fc25ff29c 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -299,4 +299,3 @@ 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 53fea072555a..875b65ff084c 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -19,7 +19,6 @@ 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 @@ -195,10 +194,6 @@ class _AttnMeta: 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 @@ -861,76 +856,3 @@ 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 646cff4c2d22..7a54dd4ca374 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -1,7 +1,6 @@ # 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 @@ -44,7 +43,6 @@ from vllm.utils.import_utils import ( has_aiter, has_deep_ep, - has_deep_ep_v2, has_deep_gemm, has_mori, ) @@ -152,10 +150,6 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: 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 @@ -245,10 +239,6 @@ def needs_deep_ep(self): 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 @@ -325,8 +315,6 @@ def is_valid(self) -> tuple[bool, str | None]: # 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 @@ -669,58 +657,6 @@ 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, @@ -733,7 +669,6 @@ 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) @@ -779,8 +714,6 @@ 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 7f1924bac5ab..78ee8084d907 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -36,12 +36,10 @@ 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, ) @@ -218,19 +216,6 @@ def expert_info(kind) -> ExpertInfo: 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, @@ -304,18 +289,6 @@ def expert_info(kind) -> ExpertInfo: 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 d7394d72cb3f..61826199771a 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) - def _run_worker(): + try: worker( ProcessGroupInfo( world_size=world_size, @@ -121,13 +121,6 @@ def _run_worker(): *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 bb2f9efc7c48..1663e562966e 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -15,7 +15,7 @@ from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage] from typing_extensions import ParamSpec -from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2 +from vllm.utils.import_utils import has_deep_ep from vllm.utils.network_utils import get_open_port if has_deep_ep(): @@ -26,11 +26,6 @@ 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") @@ -60,10 +55,11 @@ def _worker_parallel_launch( torch.accelerator.set_device_index(local_rank) device = torch.device("cuda", local_rank) torch.distributed.init_process_group( - backend="nccl", + backend="cpu:gloo,cuda: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) @@ -204,42 +200,3 @@ 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 deleted file mode 100644 index 93b7c136605b..000000000000 --- a/tests/kernels/moe/test_deepep_v2_moe.py +++ /dev/null @@ -1,542 +0,0 @@ -# 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 94beef897d08..f61aa868581e 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -8,8 +8,7 @@ set -ex # --nvshmem-ver NVSHMEM version CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} -DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"d4f41e4e93"} - +DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"} 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 2904f40f8a45..f7fb8ec42a38 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -42,7 +42,6 @@ "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 fd1c826322c3..c0055f4bbe1f 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -15,7 +15,7 @@ has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, ) -from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori +from vllm.utils.import_utils import has_deep_ep, has_mori from .base_device_communicator import All2AllManagerBase, Cache @@ -863,66 +863,3 @@ def get_handle(self, kwargs): 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 12c425021b21..ff4929d8415d 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -146,14 +146,6 @@ def __init__( 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 a264b3129df6..8f4e18d2235d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -246,9 +246,6 @@ 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 @@ -1843,18 +1840,6 @@ def _resolve_rust_frontend_path() -> str | None: "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 1351e87b5b51..6d4822146437 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -29,12 +29,7 @@ FlashInferNVLinkTwoSidedPrepareAndFinalize, ) from vllm.platforms import current_platform -from vllm.utils.import_utils import ( - has_deep_ep, - has_deep_ep_v2, - has_mori, - has_nixl_ep, -) +from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep logger = init_logger(__name__) @@ -45,8 +40,6 @@ 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(): @@ -101,11 +94,6 @@ 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 @@ -205,36 +193,6 @@ 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 1b063559b8d7..b289b85f17ae 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1070,10 +1070,6 @@ def use_mori_kernels(self): 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 @@ -1398,10 +1394,6 @@ def use_ag_rs_all2all_kernels(self): 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 257bfeee5d3e..9230fea6e5ca 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,14 +94,6 @@ 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, @@ -201,7 +193,7 @@ def apply( gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, - top_k=topk_ids.size(1), + top_k=self.topk, 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 1099368e1076..ef740667f1a0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -85,18 +85,6 @@ 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 deleted file mode 100644 index 6495e1203e0a..000000000000 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ /dev/null @@ -1,394 +0,0 @@ -# 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 00540192d2fa..b5643f0c0224 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -1038,7 +1038,6 @@ def _maybe_make_contiguous( "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 44b5e0237af0..78e2a3c5cadf 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,14 +168,6 @@ 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 1b2a8a74bdcb..c138901b6fd8 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -489,10 +489,6 @@ def __init__(self, moe: FusedMoEConfig): 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 364754f9d77c..48207a7be9a3 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -566,7 +566,7 @@ def __init__( if self.use_mega_moe: self._init_mega_moe_experts(vllm_config, config, prefix) else: - self._init_fused_moe_experts(vllm_config, config, quant_config, prefix) + self._init_fused_moe_experts(config, quant_config, prefix) def _init_mega_moe_experts( self, @@ -612,27 +612,22 @@ def _init_mega_moe_experts( 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 - 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_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.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 + 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.physical_expert_start = self.experts_start_idx self.physical_expert_end = self.experts_end_idx @@ -652,8 +647,6 @@ def _init_fused_moe_experts( 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 043798a584b8..63723a8f9ef4 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -417,61 +417,6 @@ 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.