diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index afc8fc49a2..c755c64365 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -155,5 +155,14 @@ steps: commands: - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py - pytest -v -s kernels/moe/test_deepep_moe.py - - pytest -v -s kernels/moe/test_pplx_cutlass_moe.py - # - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main + +- label: Kernels Fp4 MoE Test (B200) + timeout_in_minutes: 60 + device: b200 + num_devices: 1 + optional: true + commands: + - pytest -v -s kernels/moe/test_cutedsl_moe.py + - pytest -v -s kernels/moe/test_flashinfer_moe.py + - pytest -v -s kernels/moe/test_nvfp4_moe.py + - pytest -v -s kernels/moe/test_ocp_mx_moe.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 55127a514f..39714b8461 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -725,7 +725,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # CUTLASS MoE kernels # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works - # on Hopper). get_cutlass_(pplx_)moe_mm_data should only be compiled + # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled # if it's possible to compile MoE kernels that use its output. cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) diff --git a/csrc/ops.h b/csrc/ops.h index 5e2b475fa8..690342b37c 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -269,13 +269,13 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2, const int64_t n, const int64_t k, const bool swap_ab); -void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, const int64_t n, - const int64_t k); +void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets, + torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, + const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, + const int64_t padded_m, const int64_t n, + const int64_t k); void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a, torch::Tensor const& b, diff --git a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu index eae500cb63..41cf170a24 100644 --- a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu @@ -263,12 +263,10 @@ void get_cutlass_moe_mm_data_caller( } template -__global__ void compute_pplx_data(int32_t* expert_offsets, - int32_t* problem_sizes1, - int32_t* problem_sizes2, - const int32_t* __restrict__ expert_num_tokens, - const int padded_m, const int n, - const int k) { +__global__ void compute_batched_moe_data( + int32_t* expert_offsets, int32_t* problem_sizes1, int32_t* problem_sizes2, + const int32_t* __restrict__ expert_num_tokens, const int padded_m, + const int n, const int k) { int expert_idx = threadIdx.x; expert_offsets[expert_idx] = expert_idx * padded_m; @@ -289,24 +287,22 @@ __global__ void compute_pplx_data(int32_t* expert_offsets, } } -void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, - const int64_t n, const int64_t k) { +void get_cutlass_batched_moe_mm_data_caller( + torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k) { auto stream = at::cuda::getCurrentCUDAStream(expert_offsets.device().index()); if (num_local_experts * padded_m > SWAP_AB_THRESHOLD) { - compute_pplx_data<<<1, num_local_experts, 0, stream>>>( + compute_batched_moe_data<<<1, num_local_experts, 0, stream>>>( static_cast(expert_offsets.data_ptr()), static_cast(problem_sizes1.data_ptr()), static_cast(problem_sizes2.data_ptr()), static_cast(expert_num_tokens.data_ptr()), padded_m, n, k); } else { - compute_pplx_data<<<1, num_local_experts, 0, stream>>>( + compute_batched_moe_data<<<1, num_local_experts, 0, stream>>>( static_cast(expert_offsets.data_ptr()), static_cast(problem_sizes1.data_ptr()), static_cast(problem_sizes2.data_ptr()), diff --git a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 82ccc19608..d6e82f1db9 100644 --- a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -82,13 +82,11 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller( torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2, const int64_t n, const int64_t k, const bool swap_ab); -void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, - const int64_t n, const int64_t k); +void get_cutlass_batched_moe_mm_data_caller( + torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k); #endif void cutlass_scaled_mm_azp_sm75(torch::Tensor& c, torch::Tensor const& a, @@ -319,29 +317,30 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( version_num, ". Required capability: 90, 100, or 120"); } -void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, const int64_t n, - const int64_t k) { +void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets, + torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, + const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, + const int64_t padded_m, const int64_t n, + const int64_t k) { // This function currently gets compiled only if we have a valid cutlass moe // mm to run it for. int32_t version_num = get_sm_version_num(); #if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120) - get_cutlass_pplx_moe_mm_data_caller(expert_offsets, problem_sizes1, - problem_sizes2, expert_num_tokens, - num_local_experts, padded_m, n, k); + get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1, + problem_sizes2, expert_num_tokens, + num_local_experts, padded_m, n, k); return; #endif - TORCH_CHECK_NOT_IMPLEMENTED( - false, - "No compiled get_cutlass_pplx_moe_mm_data: no cutlass_scaled_mm kernel " - "for CUDA device capability: ", - version_num, ". Required capability: 90, 100, or 120"); + TORCH_CHECK_NOT_IMPLEMENTED(false, + "No compiled get_cutlass_batched_moe_mm_data: no " + "cutlass_scaled_mm kernel " + "for CUDA device capability: ", + version_num, + ". Required capability: 90, 100, or 120"); } void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 39b6bc98a8..8be30b2091 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -489,19 +489,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { &get_cutlass_moe_mm_problem_sizes_from_expert_offsets); // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM and PPLX. It takes expert_num_tokens and non_zero_expert_idxs + // GEMM in batched expert format. It takes expert_num_tokens // as an input, and computes expert_offsets (token start indices of each // expert). In addition to this, it computes problem sizes for each expert's // multiplication used by the two mms called from fused MoE operation. ops.def( - "get_cutlass_pplx_moe_mm_data(Tensor! expert_offsets, " + "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " " Tensor! problem_sizes1, " " Tensor! problem_sizes2, " " Tensor expert_num_tokens, " " int num_local_experts, int padded_m, " " int n, int k) -> ()"); - ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA, - &get_cutlass_pplx_moe_mm_data); + ops.impl("get_cutlass_batched_moe_mm_data", torch::kCUDA, + &get_cutlass_batched_moe_mm_data); // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) ops.def( diff --git a/docker/Dockerfile b/docker/Dockerfile index 717f27b6b2..495a480b75 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -308,7 +308,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ #################### CSRC BUILD IMAGE #################### #################### EXTENSIONS BUILD IMAGE #################### -# Build DeepGEMM, pplx-kernels, DeepEP - runs in PARALLEL with csrc-build +# Build DeepGEMM, DeepEP - runs in PARALLEL with csrc-build # This stage is independent and doesn't affect csrc cache FROM base AS extensions-build ARG CUDA_VERSION @@ -335,10 +335,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Ensure the wheel dir exists so COPY won't fail when DeepGEMM is skipped RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped -# Build pplx-kernels and DeepEP wheels +# Build DeepEP wheels COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management -ARG PPLX_COMMIT_HASH=12cecfd ARG DEEPEP_COMMIT_HASH=73b6ea4 ARG NVSHMEM_VER RUN --mount=type=cache,target=/root/.cache/uv \ @@ -347,7 +346,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ /tmp/install_python_libraries.sh \ --workspace /tmp/ep_kernels_workspace \ --mode wheel \ - ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \ ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} \ ${NVSHMEM_VER:+--nvshmem-ver "$NVSHMEM_VER"} && \ find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete @@ -676,7 +674,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH -# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage +# Install EP kernels wheels (DeepEP) that have been built in the `build` stage RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \ --mount=type=cache,target=/root/.cache/uv \ uv pip install --system ep_kernels/dist/*.whl --verbose \ diff --git a/docker/versions.json b/docker/versions.json index 24f4b6e7d1..fa090c10c4 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -52,9 +52,6 @@ "DEEPGEMM_GIT_REF": { "default": "477618cd51baffca09c4b0b87e97c03fe827ef03" }, - "PPLX_COMMIT_HASH": { - "default": "12cecfd" - }, "DEEPEP_COMMIT_HASH": { "default": "73b6ea4" }, diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index 975df8ba29..9db356cdf5 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -15,7 +15,7 @@ Based on the format of the input activations, FusedMoE implementations are broad The input activation format completely depends on the All2All Dispatch being used. * In the Contiguous variant, the All2All Dispatch returns the activations as a contiguous tensor of shape (M, K) along with TopK Ids and TopK weights of shape (M, num_topk). Look at `DeepEPHTPrepareAndFinalize` for an example. -* In the Batched variant, the All2All Dispatch returns the activations as a tensor of shape (num_experts, max_tokens, K). Here, the activations/tokens that subscribe to the same expert are batched together. Note that not all entries of the tensor are valid. The activations tensor is typically accompanied by an `expert_num_tokens` tensor of size `num_experts`, where `expert_num_tokens[i]` indicates the number of valid tokens that subscribe to the ith expert. Look at `PplxPrepareAndFinalize` or `DeepEPLLPrepareAndFinalize` for an example. +* In the Batched variant, the All2All Dispatch returns the activations as a tensor of shape (num_experts, max_tokens, K). Here, the activations/tokens that subscribe to the same expert are batched together. Note that not all entries of the tensor are valid. The activations tensor is typically accompanied by an `expert_num_tokens` tensor of size `num_experts`, where `expert_num_tokens[i]` indicates the number of valid tokens that subscribe to the ith expert. Look at `DeepEPLLPrepareAndFinalize` for an example. The FusedMoE operation is generally made of multiple operations, in both the Contiguous and Batched variants, as described in the diagrams below @@ -132,7 +132,6 @@ class FusedMoEModularKernel: Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & Combine implementation / kernel. For example, -* PplxPrepareAndFinalize type is backed by Pplx All2All kernels, * DeepEPHTPrepareAndFinalize type is backed by DeepEP High-Throughput All2All kernels, and * DeepEPLLPrepareAndFinalize type is backed by DeepEP Low-Latency All2All kernels. @@ -229,7 +228,7 @@ Doing this will add the new implementation to the test suite. ### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` Compatibility The unit test file [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script. -Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` +Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts` As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` compatibility. When invoked with incompatible types, the script will error. @@ -238,7 +237,7 @@ with incompatible types, the script will error. Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py) The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible `FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` types. -Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` +Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts` ## FusedMoEPrepareAndFinalize Implementations diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 04ceeede3e..ac5acb66bd 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -33,7 +33,6 @@ th { | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass | |---------|--------------------|--------------|---------------|-------|-----------------------|-----------| | naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] | -| pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] | | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] | | flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] | @@ -68,7 +67,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes. There are a number of MoE experts kernel implementations for different quantization types and architectures. Most follow the general API of the base Triton [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts] function. Many have modular kernel adapters, so they can be used with compatible all2all backends. This table lists each experts kernel and its particular properties. -Each kernel must be provided with one of the supported input activation formats. Some flavors of kernels support both standard and batched formats through different entry points, e.g. `TritonExperts` and `BatchedTritonExperts`. Batched format kernels are currently only needed for matching with certain all2all backends, e.g. `pplx` and `DeepEPLLPrepareAndFinalize`. +Each kernel must be provided with one of the supported input activation formats. Some flavors of kernels support both standard and batched formats through different entry points, e.g. `TritonExperts` and `BatchedTritonExperts`. Batched format kernels are currently only needed for matching with certain all2all backends, e.g. `DeepEPLLPrepareAndFinalize`. Similar to the backend kernels, each experts kernel only supports certain quantization formats. For non-modular experts, the activations will be in the original type and quantized internally by the kernel. Modular experts will expect the activations to already be in the quantized format. Both types of experts will yield outputs in the original activation type. @@ -110,5 +109,5 @@ The following table shows "families" of modular kernels that are intended to wor | backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses | |---------|-----------------------------------------|----------------------------------------------| | deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,
`TritonExperts`,
`TritonOrDeepGemmExperts`,
`CutlassExpertsFp8`,
`MarlinExperts` | -| deepep_low_latency,
pplx | `DeepEPLLPrepareAndFinalize`,
`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts` | +| deepep_low_latency | `DeepEPLLPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts` | | flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` | diff --git a/docs/governance/committers.md b/docs/governance/committers.md index 2f0780a089..89aaadc2bb 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -154,7 +154,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - FlashAttention: @LucasWilkinson - FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon - Blackwell Kernels: @mgoin, @yewentao256 -- DeepEP/DeepGEMM/pplx: @mgoin, @yewentao256 +- DeepEP/DeepGEMM: @mgoin, @yewentao256 ### Integrations diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index 82fde27d71..d469e20c98 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -8,7 +8,7 @@ EP is typically coupled with Data Parallelism (DP). While DP can be used indepen Before using EP, you need to install the necessary dependencies. We are actively working on making this easier in the future: -1. **Install DeepEP and pplx-kernels**: Set up host environment following vLLM's guide for EP kernels [here](../../tools/ep_kernels). +1. **Install DeepEP**: Set up host environment following vLLM's guide for EP kernels [here](../../tools/ep_kernels). 2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation). 3. **For disaggregated serving**: Install `gdrcopy` by running the [`install_gdrcopy.sh`](../../tools/install_gdrcopy.sh) script (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/). @@ -19,7 +19,6 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to | Backend | Use Case | Features | Best For | |---------|----------|----------|----------| | `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration | -| `pplx` | Single node | Chunked prefill support, efficient intra-node communication | Single-node deployments, development | | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios | | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios | | `flashinfer_all2allv` | MNNVL systems | FlashInfer alltoallv kernels for multi-node NVLink | Systems with NVLink across nodes | @@ -71,12 +70,11 @@ For example, with `TP=2, DP=4` (8 GPUs total): The following command serves a `DeepSeek-V3-0324` model with 1-way tensor parallel, 8-way (attention) data parallel, and 8-way expert parallel. The attention weights are replicated across all GPUs, while the expert weights are split across GPUs. It will work on a H200 (or H20) node with 8 GPUs. For H100, you can try to serve a smaller model or refer to the multi-node deployment section. ```bash -# Single node EP deployment with pplx backend +# Single node EP deployment vllm serve deepseek-ai/DeepSeek-V3-0324 \ --tensor-parallel-size 1 \ # Tensor parallelism across 1 GPU --data-parallel-size 8 \ # Data parallelism across 8 processes - --enable-expert-parallel \ # Enable expert parallelism - --all2all-backend pplx # Use pplx communication backend + --enable-expert-parallel # Enable expert parallelism ``` ## Multi-Node Deployment @@ -197,7 +195,6 @@ vllm serve deepseek-ai/DeepSeek-V3-0324 \ --tensor-parallel-size 1 \ # Tensor parallelism --data-parallel-size 8 \ # Data parallelism --enable-expert-parallel \ # Enable EP - --all2all-backend pplx \ # Use pplx communication backend --enable-eplb \ # Enable load balancer --eplb-config '{"window_size":1000,"step_interval":3000,"num_redundant_experts":2,"log_balancedness":true}' ``` diff --git a/examples/online_serving/elastic_ep/serve_deepseek_v2.sh b/examples/online_serving/elastic_ep/serve_deepseek_v2.sh index b4e9220999..3ce89e1d86 100644 --- a/examples/online_serving/elastic_ep/serve_deepseek_v2.sh +++ b/examples/online_serving/elastic_ep/serve_deepseek_v2.sh @@ -64,7 +64,7 @@ vllm serve "$MODEL_NAME" \ --enforce-eager \ --enable-expert-parallel \ --enable-eplb \ - --all2all-backend pplx \ + --all2all-backend allgather_reducescatter \ --num-redundant-experts "$REDUNDANT_EXPERTS" \ --trust-remote-code \ --host "$HOST" \ diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 87cf0453be..9f67129616 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -37,7 +37,6 @@ from vllm.utils.import_utils import ( has_deep_ep, has_deep_gemm, has_mori, - has_pplx, ) from .mk_objects import ( @@ -206,10 +205,6 @@ class Config: info = expert_info(self.fused_experts_type) return info.needs_deep_gemm - def needs_pplx(self): - info = prepare_finalize_info(self.prepare_finalize_type) - return info.backend == "pplx" - def needs_deep_ep(self): info = prepare_finalize_info(self.prepare_finalize_type) return ( @@ -290,8 +285,6 @@ class Config: return False, "Needs DeepEP, but DeepEP not available." if self.needs_deep_gemm() and not has_deep_gemm(): return False, "Needs DeepGEMM, but DeepGEMM not available." - if self.needs_pplx() and not has_pplx(): # noqa: SIM103 - return False, "Needs PPLX, but PPLX not available." if self.needs_aiter() and not has_aiter(): # noqa: SIM103 return False, "Needs Aiter, but Aiter not available." if self.needs_mori() and not has_mori(): # noqa: SIM103 diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 981f993427..0ea414c3af 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -39,7 +39,6 @@ from vllm.utils.import_utils import ( has_deep_ep, has_deep_gemm, has_mori, - has_pplx, ) @@ -238,19 +237,6 @@ if has_mori(): supports_apply_weight_on_input=False, ) -if has_pplx(): - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - ) - - register_prepare_and_finalize( - PplxPrepareAndFinalize, - batched_format, - common_float_and_int_types, - blocked_quantization_support=True, - backend="pplx", - ) - if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100): from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501 FlashInferCutlassMoEPrepareAndFinalize, diff --git a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py index 3cdc7b8213..702584f9da 100644 --- a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py +++ b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py @@ -125,7 +125,7 @@ if __name__ == "__main__": description=( "Run single prepare-finalize & fused-experts combination test" "Example : python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel " # noqa: E501 - "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts" + "--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts" ) ) args = parser.parse_args() diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index ec31e66140..cd1d0a0afe 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -14,7 +14,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.config import VllmConfig, set_current_vllm_config from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe -from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_deep_gemm from vllm.utils.torch_utils import cuda_device_count_stateless, set_random_seed from vllm.v1.worker.workspace import init_workspace_manager @@ -39,12 +39,12 @@ from .modular_kernel_tools.parallel_utils import ( ) has_any_multi_gpu_package = ( - has_deep_ep() or has_deep_gemm() or has_pplx() or has_flashinfer_cutlass_fused_moe() + has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe() ) meets_multi_gpu_requirements = pytest.mark.skipif( not has_any_multi_gpu_package, - reason="Requires deep_ep or deep_gemm or pplx or flashinfer packages", + reason="Requires deep_ep or deep_gemm or flashinfer packages", ) if current_platform.is_fp8_fnuz(): @@ -341,7 +341,7 @@ if __name__ == "__main__": description=( "Run single prepare-finalize & fused-experts combination test" "Example : python3 -m tests.kernels.moe.test_modular_kernel_combinations " - "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts" + "--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts" ) ) args = parser.parse_args() diff --git a/tests/kernels/moe/test_pplx_cutlass_moe.py b/tests/kernels/moe/test_pplx_cutlass_moe.py deleted file mode 100644 index d8a6600743..0000000000 --- a/tests/kernels/moe/test_pplx_cutlass_moe.py +++ /dev/null @@ -1,365 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest -import torch - -from tests.kernels.utils import torch_experts -from vllm import _custom_ops as ops -from vllm.config import VllmConfig, set_current_vllm_config -from vllm.model_executor.layers.fused_moe import fused_topk -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEParallelConfig, - RoutingMethodType, - fp8_w8a8_moe_quant_config, -) -from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8 -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel -from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv -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 - -try: - from pplx_kernels import AllToAll - from pplx_kernels.nvshmem import ( - nvshmem_alloc_empty_unique_id, - nvshmem_finalize, - nvshmem_get_unique_id, - nvshmem_init, - ) - - has_pplx = True -except ImportError: - has_pplx = False - -requires_pplx = pytest.mark.skipif( - not has_pplx, - reason="Requires PPLX kernels", -) - -NUM_EXPERTS = [40, 64] -TOP_KS = [6, 8] - - -def rank_chunk(num, r, w): - rem = num % w - return (num // w) + (1 if r < rem else 0) - - -def chunk_by_rank(t, r, w): - num = t.shape[0] - chunk = rank_chunk(num, r, w) - rem = num % w - if rem == 0 or r < rem: - return t[(r * chunk) : (r + 1) * chunk].contiguous() - else: - long_chunks = (num // w + 1) * rem - short_chunks = (r - rem) * chunk - start = long_chunks + short_chunks - return t[start : start + chunk].contiguous() - - -def pplx_cutlass_moe( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - w1_scale: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - a1_scale: torch.Tensor, - out_dtype, - per_act_token: bool, - per_out_ch: bool, - group_name: str | None, -): - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - ) - - init_workspace_manager(torch.cuda.current_device()) - - assert torch.cuda.current_device() == pgi.local_rank - - num_tokens, hidden_dim = a.shape - intermediate_dim = w2.shape[2] - num_experts = w1.shape[0] - block_size = hidden_dim # TODO support more cases - device = pgi.device - rank = pgi.rank - world_size = pgi.world_size - rank_num_tokens = rank_chunk(num_tokens, rank, world_size) - max_num_tokens = rank_chunk(num_tokens, 0, world_size) - topk = topk_ids.shape[1] - - if block_size == hidden_dim: - scale_elems = 4 # hack to circumvent pplx data format requirements - else: - scale_elems = (hidden_dim + block_size - 1) // block_size - - args = dict( - max_num_tokens=max_num_tokens, - num_experts=num_experts, - experts_per_token=topk, - rank=rank, - world_size=world_size, - dp_size=dp_size, - hidden_dim=hidden_dim, - hidden_dim_bytes=hidden_dim, # because a.dtype.itemsize == 1 - hidden_dim_scale_bytes=scale_elems * torch.float32.itemsize, - ) - - if group_name is None: - ata = AllToAll.internode(**args) - else: - args["group_name"] = group_name - ata = AllToAll.intranode(**args) - - w1 = w1.to(device) - w2 = w2.to(device) - w1_scale = w1_scale.to(device) - w2_scale = w2_scale.to(device) - a1_scale = a1_scale.to(device) - - assert num_experts % world_size == 0 - num_local_experts = cdiv(num_experts, world_size) - num_dispatchers = pgi.world_size // dp_size - - prepare_finalize = PplxPrepareAndFinalize( - ata, - max_num_tokens=max_num_tokens, - num_local_experts=num_local_experts, - num_dispatchers=num_dispatchers, - ) - - def make_moe_config() -> FusedMoEConfig: - return FusedMoEConfig( - num_experts=num_experts, - experts_per_token=topk, - hidden_dim=hidden_dim, - intermediate_size_per_partition=intermediate_dim, - num_local_experts=num_local_experts, - num_logical_experts=num_experts, - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), - activation=MoEActivation.SILU, - in_dtype=torch.bfloat16, - device="cuda", - routing_method=RoutingMethodType.Llama4, - ) - - experts = CutlassBatchedExpertsFp8( - moe_config=make_moe_config(), - quant_config=fp8_w8a8_moe_quant_config( - per_act_token_quant=per_act_token, - per_out_ch_quant=per_out_ch, - w1_scale=chunk_by_rank(w1_scale, rank, world_size), - w2_scale=chunk_by_rank(w2_scale, rank, world_size), - a1_scale=chunk_by_rank(a1_scale, rank, world_size) - if per_act_token - else a1_scale[rank], - ), - max_num_tokens=max_num_tokens, - num_dispatchers=num_dispatchers, - ) - - fused_cutlass_experts = FusedMoEModularKernel( - prepare_finalize, - experts, - inplace=False, - ) - - a_chunk = chunk_by_rank(a, rank, world_size).to(device) - chunk_topk_weight = chunk_by_rank(topk_weights, rank, world_size).to(device) - chunk_topk_ids = ( - chunk_by_rank(topk_ids, rank, world_size).to(torch.uint32).to(device) - ) - - out = fused_cutlass_experts( - a_chunk, - chunk_by_rank(w1, rank, world_size), - chunk_by_rank(w2, rank, world_size), - chunk_topk_weight, - chunk_topk_ids, - global_num_experts=num_experts, - expert_map=None, # TODO - ) - - torch.cuda.synchronize() - - ata.destroy() - - return out[:rank_num_tokens] - - -vllm_config = VllmConfig() - - -def _pplx_moe( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - w1_scale: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - a1_scale: torch.Tensor, - out_dtype, - a_full: torch.Tensor, - w1_full: torch.Tensor, - w2_full: torch.Tensor, - per_act_token: bool, - per_out_ch: bool, - use_internode: bool, -): - try: - if use_internode: - uid = ( - nvshmem_get_unique_id() - if pgi.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - torch.distributed.broadcast(uid, src=0) - nvshmem_init(uid, pgi.rank, pgi.world_size) - else: - group_ranks = list(range(pgi.world_size)) - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo") - group_name = cpu_group.group_name - - with set_current_vllm_config(vllm_config): - torch_output = torch_experts( - a_full, w1_full, w2_full, topk_weights, topk_ids - ) - pplx_output = pplx_cutlass_moe( - pgi, - dp_size, - a, - w1, - w2, - w1_scale, - w2_scale, - topk_weights, - topk_ids, - a1_scale, - out_dtype, - per_act_token, - per_out_ch, - group_name, - ) - - torch_output = chunk_by_rank(torch_output, pgi.rank, pgi.world_size).to( - pplx_output.device - ) - - # Uncomment if more debugging is needed - # print("PPLX OUT:", pplx_output) - # print("TORCH OUT:", torch_output) - - torch.testing.assert_close(pplx_output, torch_output, atol=0.05, rtol=0) - finally: - if use_internode: - nvshmem_finalize() - - -@pytest.mark.parametrize("m", [2, 224]) -@pytest.mark.parametrize("n", [3072]) -@pytest.mark.parametrize("k", [1536]) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("per_act_token", [True, False]) -@pytest.mark.parametrize("per_out_ch", [True, False]) -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) # , [4, 2]]) -@pytest.mark.parametrize("use_internode", [False]) -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif( - (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))( - current_platform.get_device_capability() - ), - reason="Grouped gemm is not supported on this GPU type.", -) -@requires_pplx -def test_cutlass_moe_pplx( - m: int, - n: int, - k: int, - e: int, - topk: int, - per_act_token: bool, - per_out_ch: bool, - world_dp_size: tuple[int, int], - use_internode: bool, -): - set_random_seed(7) - - with set_current_vllm_config(vllm_config): - dtype = torch.half - - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10.0 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10.0 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10.0 - - n_b_scales = 2 * n if per_out_ch else 1 - k_b_scales = k if per_out_ch else 1 - - w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float8_e4m3fn) - w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn) - w1_scale = torch.empty((e, n_b_scales, 1), device="cuda", dtype=torch.float32) - w2_scale = torch.empty((e, k_b_scales, 1), device="cuda", dtype=torch.float32) - - for expert in range(e): - w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant( - w1[expert], use_per_token_if_dynamic=per_out_ch - ) - w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant( - w2[expert], use_per_token_if_dynamic=per_out_ch - ) - - w1_d = torch.empty_like(w1) - w2_d = torch.empty_like(w2) - for expert in range(e): - w1_d[expert] = (w1_q[expert].float() * w1_scale[expert]).half() - w2_d[expert] = (w2_q[expert].float() * w2_scale[expert]).half() - - score = torch.randn((m, e), device="cuda", dtype=dtype) - topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) - - world_size, dp_size = world_dp_size - a_scale1 = ( - torch.randn( - (m if per_act_token else 1, 1), device="cuda", dtype=torch.float32 - ) - / 10.0 - ) - if not per_act_token: - a_scale1 = a_scale1.repeat(world_size, 1) - - parallel_launch( - world_size, - _pplx_moe, - dp_size, - a, - w1_q, - w2_q, - w1_scale, - w2_scale, - topk_weights, - topk_ids, - a_scale1, - dtype, - a, - w1_d, - w2_d, - per_act_token, - per_out_ch, - use_internode, - ) diff --git a/tests/kernels/moe/test_pplx_moe.py b/tests/kernels/moe/test_pplx_moe.py deleted file mode 100644 index deb3b9eb4d..0000000000 --- a/tests/kernels/moe/test_pplx_moe.py +++ /dev/null @@ -1,1021 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MOE layers. - -Run `pytest tests/kernels/test_pplx_moe.py`. -""" - -import copy -import itertools -import textwrap -import traceback -from collections.abc import Callable - -import pytest -import torch - -try: - from pplx_kernels import AllToAll - from pplx_kernels.nvshmem import ( - nvshmem_alloc_empty_unique_id, - nvshmem_finalize, - nvshmem_get_unique_id, - nvshmem_init, - ) - - has_pplx = True -except ImportError: - has_pplx = False - -from tests.kernels.moe.modular_kernel_tools.parallel_utils import _set_vllm_config -from tests.kernels.moe.utils import ( - make_dummy_moe_config, - make_shared_experts, - make_test_weights, - naive_batched_moe, -) -from tests.kernels.quant_utils import dequant -from tests.kernels.utils import torch_experts -from vllm.config import VllmConfig, set_current_vllm_config -from vllm.model_executor.layers.fused_moe import fused_topk, override_config -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig -from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts -from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceDelegate, -) -from vllm.utils.math_utils import round_up -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 - -requires_pplx = pytest.mark.skipif( - not has_pplx, - reason="Requires PPLX kernels", -) - -BATCHED_MOE_MNK_FACTORS = [ - (1, 128, 128), - (33, 2048, 128), - (64, 128, 2048), - (222, 128, 128), - (222, 2048, 1024), -] - -PPLX_COMBOS = [ - # TODO(bnell): figure out why this fails, seems to be test problem - # (1, 128, 128), - (2, 128, 512), - (3, 1024, 2048), - (4, 128, 128), - (32, 1024, 512), - (45, 512, 2048), - (64, 1024, 512), - (222, 2048, 1024), - (256, 1408, 2048), -] - -NUM_EXPERTS = [8, 64] -TOP_KS = [1, 2, 6] -DTYPES = [torch.float8_e4m3fn, torch.bfloat16] - -vllm_config = VllmConfig() - - -def torch_prepare( - a: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - max_num_tokens: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - assert topk_ids.dim() == 2 - assert topk_ids.shape[0] == a.shape[0] - - num_tokens, hidden_dim = a.shape - topk = topk_ids.shape[1] - - tokens_per_expert = torch.bincount(topk_ids.view(-1), minlength=num_experts) - - assert tokens_per_expert.numel() == num_experts - - if max_num_tokens is None: - max_num_tokens = int(tokens_per_expert.max().item()) - - b_a = torch.zeros( - (num_experts, max_num_tokens, hidden_dim), dtype=a.dtype, device=a.device - ) - - token_counts = torch.zeros(num_experts, dtype=torch.int, device=a.device) - - for token in range(num_tokens): - for j in range(topk): - expert_id = topk_ids[token, j] - idx = token_counts[expert_id] - b_a[expert_id, idx : idx + 1, :] = a[token, :] - token_counts[expert_id] = token_counts[expert_id] + 1 - - return b_a, tokens_per_expert - - -def torch_finalize( - b_out: torch.Tensor, topk_weight: torch.Tensor, topk_ids: torch.Tensor -) -> torch.Tensor: - num_tokens = topk_ids.shape[0] - num_experts = b_out.shape[0] - K = b_out.shape[-1] - out = torch.zeros((num_tokens, K), dtype=b_out.dtype, device=b_out.device) - expert_counts = torch.zeros(num_experts, dtype=torch.int, device=b_out.device) - for token in range(num_tokens): - expert_ids = topk_ids[token] - for i in range(expert_ids.numel()): - expert_id = expert_ids[i] - idx = expert_counts[expert_id] - out[token, :] = ( - out[token, :] - + b_out[expert_id, idx : idx + 1, :] * topk_weight[token, i] - ) - expert_counts[expert_id] = expert_counts[expert_id] + 1 - - return out - - -def torch_batched_moe( - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weight: torch.Tensor, - topk_ids: torch.Tensor, -) -> torch.Tensor: - num_experts = w1.shape[0] - b_a, tokens_per_expert = torch_prepare(a, topk_ids, num_experts) - assert b_a.dim() == 3 - num_tokens, topk = topk_ids.shape - _, max_num_tokens, K = b_a.shape - assert num_experts == b_a.shape[0] and w2.shape[1] == K - out = torch.zeros( - (num_experts, max_num_tokens, K), dtype=b_a.dtype, device=b_a.device - ) - tmp = torch.empty( - (max_num_tokens, w1.shape[1] // 2), dtype=b_a.dtype, device=b_a.device - ) - for expert in range(num_experts): - num = tokens_per_expert[expert] - if num > 0: - torch.ops._C.silu_and_mul( - tmp[:num], b_a[expert, :num, :] @ w1[expert].transpose(0, 1) - ) - out[expert, :num, :] = tmp[:num] @ w2[expert].transpose(0, 1) - - return torch_finalize(out, topk_weight, topk_ids) - - -@pytest.mark.parametrize("m,n,k", BATCHED_MOE_MNK_FACTORS) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_fused_moe_batched_experts( - m: int, - n: int, - k: int, - e: int, - topk: int, - dtype: torch.dtype, - workspace_init, -): - set_random_seed(7) - - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 - score = torch.randn((m, e), device="cuda", dtype=dtype) - - with set_current_vllm_config(vllm_config): - topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) - baseline_output = torch_experts( - a, w1, w2, topk_weight, topk_ids - ) # only for baseline - torch_output = torch_batched_moe(a, w1, w2, topk_weight, topk_ids) - batched_output = naive_batched_moe( - a, w1, w2, topk_weight, topk_ids - ) # pick torch_experts or this - - torch.testing.assert_close(baseline_output, torch_output, atol=2e-2, rtol=0) - torch.testing.assert_close(baseline_output, batched_output, atol=2e-2, rtol=0) - - -def create_pplx_prepare_finalize( - num_tokens: int, - hidden_dim: int, - topk: int, - num_experts: int, - rank: int, - dp_size: int, - world_size: int, - in_dtype: torch.dtype, - quant_dtype: torch.dtype | None, - block_shape: list[int] | None, - per_act_token_quant: bool, - group_name: str | None, -): - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - pplx_hidden_dim_scale_bytes, - ) - - max_num_tokens = max(rank_chunk(num_tokens, 0, world_size), 1) - num_local_experts = rank_chunk(num_experts, 0, world_size) - - hidden_dim_bytes, scale_bytes = pplx_hidden_dim_scale_bytes( - max_num_tokens, - hidden_dim, - in_dtype, - quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - ) - - args = dict( - max_num_tokens=max_num_tokens, - num_experts=num_experts, - experts_per_token=topk, - rank=rank, - world_size=world_size, - dp_size=dp_size, - hidden_dim=hidden_dim, - hidden_dim_bytes=hidden_dim_bytes, - hidden_dim_scale_bytes=scale_bytes, - ) - - if group_name is None: - ata = AllToAll.internode(**args) - else: - args["group_name"] = group_name - ata = AllToAll.intranode(**args) - - prepare_finalize = PplxPrepareAndFinalize( - ata, - max_num_tokens=max_num_tokens, - num_local_experts=num_local_experts, - num_dispatchers=world_size // dp_size, - ) - - return prepare_finalize, ata - - -def rank_chunk(num: int, r: int, w: int) -> int: - rem = num % w - return (num // w) + (1 if r < rem else 0) - - -def chunk_by_rank(t: torch.Tensor, r: int, w: int) -> torch.Tensor: - chunk = rank_chunk(t.shape[0], r, w) - return t[(r * chunk) : (r + 1) * chunk] - - -def maybe_chunk_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None: - if t is not None: - return chunk_by_rank(t, r, w) - else: - return t - - -def chunk_scales_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None: - if t is not None and t.numel() > 1: - chunk = rank_chunk(t.shape[0], r, w) - return t[(r * chunk) : (r + 1) * chunk] - else: - return t - - -def chunk_scales(t: torch.Tensor | None, start: int, end: int) -> torch.Tensor | None: - if t is not None and t.numel() > 1: - return t[start:end] - else: - return t - - -def dummy_work(a: torch.Tensor) -> torch.Tensor: - return a * 1.1 - - -def pplx_prepare_finalize( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - topk_weight: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - quant_dtype: torch.dtype | None, - block_shape: list[int] | None, - per_act_token_quant: bool, - group_name: str | None, -) -> torch.Tensor: - assert torch.cuda.current_device() == pgi.local_rank - - topk = topk_ids.shape[1] - num_tokens, hidden_dim = a.shape - device = pgi.device - rank = pgi.rank - world_size = pgi.world_size - - topk_ids = topk_ids.to(dtype=torch.uint32) - - prepare_finalize, ata = create_pplx_prepare_finalize( - num_tokens, - hidden_dim, - topk, - num_experts, - rank, - dp_size, - world_size, - a.dtype, - quant_dtype, - block_shape, - per_act_token_quant, - group_name, - ) - - assert a.shape[0] == topk_ids.shape[0] - - a_chunk = chunk_by_rank(a, rank, world_size).to(device) - chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device) - chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device) - - assert a_chunk.shape[0] == chunk_topk_ids.shape[0] - - out = torch.full( - a_chunk.shape, - torch.nan, - dtype=a.dtype, - device=device, - ) - - if quant_dtype is not None and not per_act_token_quant and block_shape is None: - a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - else: - a1_scale = None - a2_scale = None - - b_a, b_a_scale, expert_num_tokens, _, _ = prepare_finalize.prepare( - a_chunk, - chunk_topk_weight, - chunk_topk_ids, - num_experts, - None, - False, - FusedMoEQuantConfig.make( - quant_dtype, - per_act_token_quant=per_act_token_quant, - per_out_ch_quant=False, - block_shape=block_shape, - a1_scale=a1_scale, - a2_scale=a2_scale, - ), - ) - - b_a = dummy_work(dequant(b_a, b_a_scale, block_shape, per_act_token_quant, a.dtype)) - - prepare_finalize.finalize( - out, - b_a, - chunk_topk_weight, - chunk_topk_ids, - False, - weight_and_reduce_impl=TopKWeightAndReduceDelegate(), - ) - - torch.cuda.synchronize() - - ata.destroy() - - num_tokens = a_chunk.shape[0] - - return out[:num_tokens] - - -def _pplx_prepare_finalize( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - score: torch.Tensor, - topk: torch.Tensor, - num_experts: int, - quant_dtype: torch.dtype | None, - block_shape: list[int] | None, - per_act_token_quant: bool, - use_internode: bool, -): - try: - if use_internode: - uid = ( - nvshmem_get_unique_id() - if pgi.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - torch.distributed.broadcast(uid, src=0) - nvshmem_init(uid, pgi.rank, pgi.world_size) - group_name = None - else: - group_ranks = list(range(pgi.world_size)) - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo") - group_name = cpu_group.group_name - - topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) - m, k = a.shape - - a_rep = torch.repeat_interleave(dummy_work(a), topk, dim=0) - - torch_output = ( - a_rep.view(m, topk, k) * topk_weight.view(m, topk, 1).to(a_rep.dtype) - ).sum(dim=1) - - pplx_output = pplx_prepare_finalize( - pgi, - dp_size, - a, - topk_weight, - topk_ids, - num_experts, - quant_dtype, - block_shape, - per_act_token_quant, - group_name, - ) - - torch_output = chunk_by_rank(torch_output, pgi.rank, pgi.world_size).to( - pgi.device - ) - - torch.testing.assert_close(pplx_output, torch_output, atol=3e-2, rtol=3e-2) - finally: - if use_internode: - nvshmem_finalize() - - -@pytest.mark.parametrize("mnk", PPLX_COMBOS) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("per_act_token_quant", [False, True]) -@pytest.mark.parametrize("block_shape", [None, [128, 128]]) -@pytest.mark.parametrize("use_internode", [False]) -@pytest.mark.optional -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_prepare_finalize_slow( - mnk: tuple[int, int, int], - e: int, - topk: int, - dtype: torch.dtype, - world_dp_size: tuple[int, int], - per_act_token_quant: bool, - block_shape: list[int] | None, - use_internode: bool, -): - if dtype == torch.float8_e4m3fn: - use_fp8_w8a8 = True - act_dtype = torch.bfloat16 - quant_dtype = dtype - else: - use_fp8_w8a8 = False - act_dtype = dtype - quant_dtype = None - - if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None): - pytest.skip("Skip quantization test for non-quantized type") - - if per_act_token_quant and block_shape is not None: - pytest.skip("Skip illegal quantization combination") - - set_random_seed(7) - m, n, k = mnk - world_size, dp_size = world_dp_size - device = "cuda" - - a = torch.randn((m, k), device=device, dtype=act_dtype) / 10 - score = torch.randn((m, e), device=device, dtype=act_dtype) - - parallel_launch( - world_size, - _pplx_prepare_finalize, - dp_size, - a, - score, - topk, - e, - quant_dtype, - block_shape, - per_act_token_quant, - use_internode, - ) - - -def pplx_moe( - group_name: str | None, - rank: int, - world_size: int, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weight: torch.Tensor, - topk_ids: torch.Tensor, - w1_scale: torch.Tensor | None = None, - w2_scale: torch.Tensor | None = None, - a1_scale: torch.Tensor | None = None, - a2_scale: torch.Tensor | None = None, - quant_dtype: torch.dtype | None = None, - per_act_token_quant=False, - block_shape: list[int] | None = None, - use_compile: bool = False, - use_cudagraphs: bool = True, - shared_experts: torch.nn.Module | None = None, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - num_tokens, hidden_dim = a.shape - num_experts = w1.shape[0] - topk = topk_ids.shape[1] - max_num_tokens = round_up(rank_chunk(a.shape[0], 0, world_size), 16) - - prepare_finalize, ata = create_pplx_prepare_finalize( - num_tokens, - hidden_dim, - topk, - num_experts, - rank, - dp_size, - world_size, - a.dtype, - quant_dtype, - block_shape, - per_act_token_quant, - group_name, - ) - - topk_ids = topk_ids.to(dtype=torch.uint32) - - # Note: workers with the same dp_rank must use the exact same inputs. - a_chunk = chunk_by_rank(a, rank, world_size) - chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size) - chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size) - - # Chunking weights like this only works for batched format - w1_chunk = chunk_by_rank(w1, rank, world_size) - w2_chunk = chunk_by_rank(w2, rank, world_size) - w1_scale_chunk = maybe_chunk_by_rank(w1_scale, rank, world_size) - w2_scale_chunk = maybe_chunk_by_rank(w2_scale, rank, world_size) - a1_scale_chunk = chunk_scales_by_rank(a1_scale, rank, world_size) - a2_scale_chunk = chunk_scales_by_rank(a2_scale, rank, world_size) - - quant_config = FusedMoEQuantConfig.make( - quant_dtype, - block_shape=block_shape, - per_act_token_quant=per_act_token_quant, - w1_scale=w1_scale_chunk, - w2_scale=w2_scale_chunk, - a1_scale=a1_scale_chunk, - a2_scale=a2_scale_chunk, - ) - - experts = BatchedTritonExperts( - max_num_tokens=max_num_tokens, - num_dispatchers=prepare_finalize.num_dispatchers(), - quant_config=quant_config, - moe_config=make_dummy_moe_config(), - ) - - fused_experts = FusedMoEModularKernel( - prepare_finalize, - experts, - shared_experts, - inplace=False, - ) - - # Note: for now use_compile will error out if the problem size is - # large enough to trigger chunking. I'm leaving the flag and - # setup code in case we are able to revisit this later. - if use_compile: - _fused_experts = torch.compile( - fused_experts, backend="inductor", fullgraph=True - ) - torch._dynamo.mark_dynamic(a_chunk, 0) - torch._dynamo.mark_dynamic(chunk_topk_weight, 0) - torch._dynamo.mark_dynamic(chunk_topk_ids, 0) - else: - _fused_experts = fused_experts - - out = _fused_experts( - a_chunk, - w1_chunk, - w2_chunk, - chunk_topk_weight, - chunk_topk_ids, - global_num_experts=num_experts, - ) - - if use_cudagraphs: - if isinstance(out, tuple): - out[0].fill_(0) - out[1].fill_(0) - else: - out.fill_(0) - stream = torch.cuda.Stream() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - out = _fused_experts( - a_chunk, - w1_chunk, - w2_chunk, - chunk_topk_weight, - chunk_topk_ids, - global_num_experts=num_experts, - ) - - torch.cuda.synchronize() - graph.replay() - - torch.cuda.synchronize() - - ata.destroy() - - return out - - -def _pplx_moe( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - score: torch.Tensor, - topk: int, - num_experts: int, - w1_s: torch.Tensor | None = None, - w2_s: torch.Tensor | None = None, - quant_dtype: torch.dtype | None = None, - per_act_token_quant: bool = False, - block_shape: list[int] | None = None, - use_internode: bool = False, - shared_experts: torch.nn.Module | None = None, -): - try: - if use_internode: - uid = ( - nvshmem_get_unique_id() - if pgi.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - torch.distributed.broadcast(uid, src=0) - nvshmem_init(uid, pgi.rank, pgi.world_size) - group_name = None - else: - group_ranks = list(range(pgi.world_size)) - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo") - group_name = cpu_group.group_name - - m, k = a.shape - e, _, n = w2.shape - - moe_config = get_default_config(m, e, n, k, topk, a.dtype, False) - - device = torch.device("cuda", pgi.rank) - rank = pgi.rank - world_size = pgi.world_size - - a = a.to(device) - w1 = w1.to(device) - w2 = w2.to(device) - w1_s = w1_s.to(device) if w1_s is not None else None - w2_s = w2_s.to(device) if w2_s is not None else None - - if quant_dtype is not None and not per_act_token_quant and block_shape is None: - a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - else: - a1_scale = None - a2_scale = None - - with set_current_vllm_config(vllm_config), override_config(moe_config): - topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) - - shared_output = shared_experts(a) if shared_experts is not None else None - - torch_output = torch_experts( - a, - w1, - w2, - topk_weight, - topk_ids, - w1_scale=w1_s, - w2_scale=w2_s, - a1_scale=a1_scale, - a2_scale=a2_scale, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - ) - - batched_output = naive_batched_moe( - a, - w1, - w2, - topk_weight, - topk_ids, - w1_scale=w1_s, - w2_scale=w2_s, - a1_scale=a1_scale, - a2_scale=a2_scale, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - ) - - pplx_outputs = pplx_moe( - group_name, - rank, - world_size, - dp_size, - a, - w1, - w2, - topk_weight, - topk_ids, - w1_scale=w1_s, - w2_scale=w2_s, - a1_scale=a1_scale, - a2_scale=a2_scale, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - shared_experts=shared_experts, - ) - - if shared_experts is None: - pplx_shared_output = None - pplx_output = pplx_outputs - assert isinstance(pplx_output, torch.Tensor) - else: - pplx_shared_output, pplx_output = pplx_outputs - - if shared_output is not None: - assert pplx_shared_output is not None - chunked_shared_output = chunk_by_rank( - shared_output, pgi.rank, pgi.world_size - ).to(pplx_shared_output.device) - else: - chunked_shared_output = None - - chunked_batch_output = chunk_by_rank( - batched_output, pgi.rank, pgi.world_size - ).to(pplx_output.device) - - torch.testing.assert_close(batched_output, torch_output, atol=3e-2, rtol=3e-2) - - torch.testing.assert_close( - pplx_output, chunked_batch_output, atol=3e-2, rtol=3e-2 - ) - - if shared_experts is not None: - assert chunked_shared_output is not None - assert pplx_shared_output is not None - torch.testing.assert_close( - pplx_shared_output, chunked_shared_output, atol=3e-2, rtol=3e-2 - ) - - finally: - if use_internode: - nvshmem_finalize() - - -@pytest.mark.parametrize("mnk", PPLX_COMBOS) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("per_act_token_quant", [False, True]) -@pytest.mark.parametrize("block_shape", [None, [128, 128]]) -@pytest.mark.parametrize("use_internode", [False]) -@pytest.mark.optional -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_moe_slow( - mnk: tuple[int, int, int], - e: int, - topk: int, - dtype: torch.dtype, - world_dp_size: tuple[int, int], - per_act_token_quant: bool, - block_shape: list[int] | None, - use_internode: bool, -): - set_random_seed(7) - m, n, k = mnk - world_size, dp_size = world_dp_size - - if dtype == torch.float8_e4m3fn: - use_fp8_w8a8 = True - quant_dtype = dtype - else: - use_fp8_w8a8 = False - quant_dtype = None - - if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None): - pytest.skip("Skip quantization test for non-quantized type") - - if per_act_token_quant and block_shape is not None: - pytest.skip("Skip illegal quantization combination") - - a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 - score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) - - (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( - e, - n, - k, - quant_dtype=quant_dtype, - block_shape=block_shape, - per_out_ch_quant=per_act_token_quant, - ) - - parallel_launch( - world_size, - _pplx_moe, - dp_size, - a, - w1, - w2, - score, - topk, - e, - w1_s, - w2_s, - quant_dtype, - per_act_token_quant, - block_shape, - use_internode, - ) - - -def _pplx_test_loop( - pgi: ProcessGroupInfo, - dp_size: int, - use_internode: bool, - use_shared_experts: bool, - make_weights: bool, - test_fn: Callable, -): - device = torch.device(f"cuda:{pgi.local_rank}") - init_workspace_manager(device) - - def format_result(msg, ex=None): - if ex is not None: - x = str(ex) - newx = x.strip(" \n\t")[:16] - if len(newx) < len(x): - newx = newx + " ..." - - prefix = "E\t" - print(f"{textwrap.indent(traceback.format_exc(), prefix)}") - print(f"FAILED {msg} - {newx}\n") - else: - print(f"PASSED {msg}") - - if use_shared_experts: - # Note: this config is only needed for the non-naive shared experts. - new_vllm_config = copy.deepcopy(vllm_config) - new_vllm_config.parallel_config.data_parallel_size = pgi.world_size - new_vllm_config.parallel_config.enable_expert_parallel = True - _set_vllm_config(new_vllm_config, pgi.world_size, pgi.rank, pgi.local_rank) - - set_random_seed(7) - combos = itertools.product( - PPLX_COMBOS, NUM_EXPERTS, TOP_KS, DTYPES, [False, True], [None, [128, 128]] - ) - exceptions = [] - count = 0 - for mnk, e, topk, dtype, per_act_token_quant, block_shape in combos: - count = count + 1 - m, n, k = mnk - - if dtype == torch.float8_e4m3fn: - use_fp8_w8a8 = True - quant_dtype = dtype - else: - use_fp8_w8a8 = False - quant_dtype = None - - test_desc = ( - f"test_pplx_moe[mnk={mnk}, e={e}, topk={topk}, " - f"dtype={dtype}, per_act_token={per_act_token_quant}, " - f"block_shape={block_shape}, use_internode={use_internode}, " - f"use_shared_experts={use_shared_experts}" - ) - - if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None): - print(f"{test_desc} - Skip quantization test for non-quantized type.") - continue - - if per_act_token_quant and block_shape is not None: - print(f"{test_desc} - Skip illegal quantization combination.") - continue - - a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 - score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) - - args = dict() - if make_weights: - (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( - e, - n, - k, - quant_dtype=quant_dtype, - block_shape=block_shape, - per_out_ch_quant=per_act_token_quant, - ) - args["w1"] = w1 - args["w2"] = w2 - args["w1_s"] = w1_s - args["w2_s"] = w2_s - - if use_shared_experts: - args["shared_experts"] = make_shared_experts( - n, - k, - in_dtype=a.dtype, - quant_dtype=quant_dtype, - ) - - try: - test_fn( - pgi=pgi, - dp_size=dp_size, - a=a, - score=score, - topk=topk, - num_experts=e, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - use_internode=use_internode, - **args, - ) - format_result(test_desc) - except Exception as ex: - format_result(test_desc, ex) - exceptions.append(ex) - - if len(exceptions) > 0: - raise RuntimeError( - f"{len(exceptions)} of {count} tests failed in child process, " - f"rank={pgi.rank}." - ) - else: - print(f"{count} of {count} tests passed in child process, rank={pgi.rank}.") - - -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("use_internode", [False]) -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_prepare_finalize( - world_dp_size: tuple[int, int], - use_internode: bool, -): - set_random_seed(7) - world_size, dp_size = world_dp_size - parallel_launch( - world_size * dp_size, - _pplx_test_loop, - dp_size, - use_internode, - False, - False, - _pplx_prepare_finalize, - ) - - -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("use_internode", [False]) -@pytest.mark.parametrize("use_shared_experts", [False, True]) -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_moe( - world_dp_size: tuple[int, int], - use_internode: bool, - use_shared_experts: bool, -): - set_random_seed(7) - world_size, dp_size = world_dp_size - parallel_launch( - world_size, - _pplx_test_loop, - dp_size, - use_internode, - use_shared_experts, - True, - _pplx_moe, - ) diff --git a/tools/ep_kernels/README.md b/tools/ep_kernels/README.md index ab0e358802..b4eabe18ca 100644 --- a/tools/ep_kernels/README.md +++ b/tools/ep_kernels/README.md @@ -4,7 +4,7 @@ Large-scale cluster-level expert parallel, as described in the [DeepSeek-V3 Tech Here we break down the requirements in 2 steps: -1. Build and install the Python libraries (both [pplx-kernels](https://github.com/ppl-ai/pplx-kernels) and [DeepEP](https://github.com/deepseek-ai/DeepEP)), including necessary dependencies like NVSHMEM. This step does not require any privileged access. Any user can do this. +1. Build and install the Python libraries ([DeepEP](https://github.com/deepseek-ai/DeepEP)), including necessary dependencies like NVSHMEM. This step does not require any privileged access. Any user can do this. 2. Configure NVIDIA driver to enable IBGDA. This step requires root access, and must be done on the host machine. Step 2 is necessary for multi-node deployment. diff --git a/tools/ep_kernels/elastic_ep/install_eep_libraries.sh b/tools/ep_kernels/elastic_ep/install_eep_libraries.sh index fe7b862159..31519c2871 100755 --- a/tools/ep_kernels/elastic_ep/install_eep_libraries.sh +++ b/tools/ep_kernels/elastic_ep/install_eep_libraries.sh @@ -76,11 +76,4 @@ popd export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem_install:$CMAKE_PREFIX_PATH -# build and install pplx, require pytorch installed -pushd "$WORKSPACE" -git clone https://github.com/ppl-ai/pplx-kernels -cd pplx-kernels -# see https://github.com/pypa/pip/issues/9955#issuecomment-838065925 -# PIP_NO_BUILD_ISOLATION=0 disables build isolation -PIP_NO_BUILD_ISOLATION=0 TORCH_CUDA_ARCH_LIST=9.0a+PTX pip install . --no-deps -v diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 148cb6e18d..3372dd10f4 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -4,12 +4,10 @@ set -ex # usage: ./install_python_libraries.sh [options] # --workspace workspace directory (default: ./ep_kernels_workspace) # --mode "install" (default) or "wheel" -# --pplx-ref pplx-kernels commit hash # --deepep-ref DeepEP commit hash # --nvshmem-ver NVSHMEM version CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} -PPLX_COMMIT_HASH=${PPLX_COMMIT_HASH:-"12cecfd"} 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} @@ -35,14 +33,6 @@ while [[ $# -gt 0 ]]; do MODE="$2" shift 2 ;; - --pplx-ref) - if [[ -z "$2" || "$2" =~ ^- ]]; then - echo "Error: --pplx-ref requires an argument." >&2 - exit 1 - fi - PPLX_COMMIT_HASH="$2" - shift 2 - ;; --deepep-ref) if [[ -z "$2" || "$2" =~ ^- ]]; then echo "Error: --deepep-ref requires an argument." >&2 @@ -188,14 +178,6 @@ do_build() { popd } -# build pplx-kernels -do_build \ - "https://github.com/ppl-ai/pplx-kernels" \ - "pplx-kernels" \ - "setup.py" \ - "$PPLX_COMMIT_HASH" \ - "" - # build DeepEP do_build \ "https://github.com/deepseek-ai/DeepEP" \ diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index e48ba6c997..37cf436203 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -988,7 +988,7 @@ def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor): return output_tensor -def get_cutlass_pplx_moe_mm_data( +def get_cutlass_batched_moe_mm_data( expert_offsets: torch.Tensor, problem_sizes1: torch.Tensor, problem_sizes2: torch.Tensor, @@ -1011,7 +1011,7 @@ def get_cutlass_pplx_moe_mm_data( multiplication in two grouped MMs used in the fused MoE operation. """ - return torch.ops._C.get_cutlass_pplx_moe_mm_data( + return torch.ops._C.get_cutlass_batched_moe_mm_data( expert_offsets, problem_sizes1, problem_sizes2, diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index d22e9a96e0..01dc61cdca 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1045,7 +1045,7 @@ class CompilationConfig: "are optimized for prefill and are incompatible with CUDA Graphs. " "In order to use CUDA Graphs for decode-optimized workloads, " "use --all2all-backend with another option, such as " - "deepep_low_latency, pplx, or allgather_reducescatter." + "deepep_low_latency or allgather_reducescatter." ) self.cudagraph_mode = CUDAGraphMode.NONE diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index cc2cfa97b5..fa4f72dccc 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -152,7 +152,6 @@ class ParallelConfig: - "naive": Naive all2all implementation using broadcasts\n - "allgather_reducescatter": All2all based on allgather and reducescatter\n - - "pplx": Use pplx kernels\n - "deepep_high_throughput": Use deepep high-throughput kernels\n - "deepep_low_latency": Use deepep low-latency kernels\n - "mori": Use mori kernels\n @@ -310,6 +309,13 @@ class ParallelConfig: f"but found: {self._api_process_rank}" ) + if self.all2all_backend == "pplx": + logger.warning( + "The 'pplx' all2all backend has been removed. " + "Falling back to 'allgather_reducescatter'." + ) + self.all2all_backend = "allgather_reducescatter" + if self.data_parallel_size_local > self.data_parallel_size: raise ValueError( f"data_parallel_size_local ({self.data_parallel_size_local}) " @@ -442,7 +448,6 @@ class ParallelConfig: # In this case, ensure the input to the experts is sequence parallel # to avoid the excess work. # - # Not needed for pplx-kernels as it can handle duplicate input tokens. @property def use_sequence_parallel_moe(self) -> bool: return ( diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 678cd45800..4acab4e3c8 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -3,14 +3,13 @@ from typing import Any import torch -import torch.distributed as dist import vllm.envs as envs from vllm.distributed import get_dp_group, get_ep_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.utils.flashinfer import has_flashinfer_all2all -from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_mori from .base_device_communicator import All2AllManagerBase, Cache @@ -235,96 +234,6 @@ class AgRsAll2AllManager(All2AllManagerBase): pass -class PPLXAll2AllManager(All2AllManagerBase): - """ - All2All communication based on PPLX kernels. - """ - - def __init__(self, cpu_group): - assert has_pplx(), ( - "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md" - " to install pplx_kernels." - ) - super().__init__(cpu_group) - - if self.internode: - # inter-node communication needs nvshmem, - # intra-node communication uses p2p mapping directly - from pplx_kernels.nvshmem import ( # type: ignore[import-not-found] - nvshmem_alloc_empty_unique_id, - nvshmem_get_unique_id, - nvshmem_init, - ) - - logger.debug( - "Initialize NVSHMEM for pplx_kernels: rank=%d, world size=%d", - self.rank, - self.world_size, - ) - uid = ( - nvshmem_get_unique_id() - if self.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - dist.broadcast( - uid, - src=dist.get_process_group_ranks(self.cpu_group)[0], - group=self.cpu_group, - ) - logger.debug("PPLX NVSHMEM UID = %s", uid) - nvshmem_init(uid, self.rank, self.world_size) - - self.handle_cache = Cache() - - def get_handle(self, kwargs): - import pplx_kernels as pplx # type: ignore[import-not-found] - - return self.handle_cache.get_or_create( - kwargs, - pplx.AllToAll.internode if self.internode else pplx.AllToAll.intranode, - ) - - def dispatch_router_logits( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - is_sequence_parallel: bool = False, - extra_tensors: list[torch.Tensor] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - raise NotImplementedError - - def dispatch( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - is_sequence_parallel: bool = False, - extra_tensors: list[torch.Tensor] | None = None, - ) -> ( - tuple[torch.Tensor, torch.Tensor, torch.Tensor] - | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] - ): - raise NotImplementedError - - def combine( - self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False - ) -> torch.Tensor: - raise NotImplementedError - - def destroy(self): - with self.handle_cache._lock: - for _, handle in self.handle_cache._cache.items(): - handle.destroy() - - if self.internode: - from pplx_kernels.nvshmem import ( - nvshmem_finalize, # type: ignore[import-not-found] - ) - - logger.debug("PPLX NVSHMEM finalize") - nvshmem_finalize() - - class DeepEPAll2AllManagerBase(All2AllManagerBase): """ All2All communication based on DeepEP High-Throughput kernels. diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 62e2b90377..dd571482f5 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -112,10 +112,6 @@ class CudaCommunicator(DeviceCommunicatorBase): from .all2all import AgRsAll2AllManager self.all2all_manager = AgRsAll2AllManager(self.cpu_group) - elif self.all2all_backend == "pplx": - from .all2all import PPLXAll2AllManager - - self.all2all_manager = PPLXAll2AllManager(self.cpu_group) elif self.all2all_backend == "deepep_high_throughput": from .all2all import DeepEPHTAll2AllManager @@ -298,7 +294,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.fi_ar_comm = None if self.all2all_manager is not None: self.all2all_manager.destroy() - self.all2all_manager = None + self.all2all_manager = None # type: ignore[assignment] def all_gatherv( self, diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 7c3701b4ea..891f19cfe2 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -159,7 +159,7 @@ class EplbModelState: NOTE: The expert_load_view now records load for all physical experts rather than just local experts. This ensures consistent load statistics - across different dispatch methods (naive all-to-all, DeepEP, pplx-kernels). + across different dispatch methods (naive all-to-all, DeepEP). The recorded load will be multiplied by dp_size when using naive all-to-all due to each DP rank contributing the same token set to the calculation. See: diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index bf8ec2dc6f..8c1bfe1c36 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any import torch @@ -24,16 +25,11 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) from vllm.platforms import current_platform -from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_mori logger = init_logger(__name__) if current_platform.is_cuda_alike(): - if has_pplx(): - from .pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - pplx_hidden_dim_scale_bytes, - ) if has_deep_ep(): from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize from .deepep_ll_prepare_finalize import ( @@ -120,51 +116,10 @@ def maybe_make_prepare_finalize( prepare_finalize: FusedMoEPrepareAndFinalize | None = None - if moe.use_pplx_kernels: - assert quant_config is not None - - hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes( - moe.max_num_tokens, - moe.hidden_dim, - moe.in_dtype, - quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - ) - - all_to_all_args = dict( - max_num_tokens=moe.max_num_tokens, - num_experts=moe.num_experts, - experts_per_token=moe.experts_per_token, # topk - rank=all2all_manager.rank, - world_size=all2all_manager.world_size, - # dp_size actually means tp_size, bug in pplx kernels - dp_size=all2all_manager.tp_group.world_size, - hidden_dim=moe.hidden_dim, - hidden_dim_bytes=hidden_dim_bytes, - hidden_dim_scale_bytes=hidden_scale_bytes, - ) - - num_dispatchers = ( - all2all_manager.world_size // all2all_manager.tp_group.world_size - ) - - # Intranode pplx a2a takes a group name while internode does not. - if not all2all_manager.internode: - all_to_all_args["group_name"] = all2all_manager.cpu_group.group_name - - handle = all2all_manager.get_handle(all_to_all_args) - - prepare_finalize = PplxPrepareAndFinalize( - handle, - max_num_tokens=moe.max_num_tokens, - num_local_experts=moe.num_local_experts, - num_dispatchers=num_dispatchers, - ) - elif moe.use_deepep_ht_kernels: + if moe.use_deepep_ht_kernels: assert moe.dp_size == all2all_manager.dp_world_size - all_to_all_args = dict() + all_to_all_args: dict[str, Any] = dict() handle = all2all_manager.get_handle(all_to_all_args) prepare_finalize = DeepEPHTPrepareAndFinalize( handle, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 87e1e244b3..33d69b57a9 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -939,10 +939,6 @@ class FusedMoEParallelConfig: def use_all2all_kernels(self): return self.dp_size > 1 and self.use_ep - @property - def use_pplx_kernels(self): - return self.use_all2all_kernels and self.all2all_backend == "pplx" - @property def use_deepep_ht_kernels(self): return ( @@ -962,7 +958,7 @@ class FusedMoEParallelConfig: @property def use_batched_activation_format(self): - return self.use_deepep_ll_kernels or self.use_pplx_kernels + return self.use_deepep_ll_kernels @property def use_naive_all2all_kernels(self): @@ -1221,10 +1217,6 @@ class FusedMoEConfig: def use_ep(self): return self.moe_parallel_config.use_ep - @property - def use_pplx_kernels(self): - return self.moe_parallel_config.use_pplx_kernels - @property def use_deepep_ht_kernels(self): return self.moe_parallel_config.use_deepep_ht_kernels diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index ae9430d29d..ac9ba56a6b 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -166,7 +166,7 @@ def run_cutlass_moe_fp8( problem_sizes1 = torch.empty((local_E, 3), dtype=torch.int32, device=device) problem_sizes2 = torch.empty((local_E, 3), dtype=torch.int32, device=device) - ops.get_cutlass_pplx_moe_mm_data( + ops.get_cutlass_batched_moe_mm_data( expert_offsets, problem_sizes1, problem_sizes2, diff --git a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py index fbd47f8c42..24ae2d3c82 100644 --- a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py @@ -493,7 +493,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): """ A reference prepare/finalize class that reorganizes the tokens into expert batched format, i.e. E x max_num_tokens x K. This is the format - that the PPLX dispatch/combine kernels use. + that the batched dispatch/combine kernels use. """ def __init__( @@ -648,7 +648,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute): """ A reference MoE expert class that operates on expert batched format, - i.e. E x max_num_tokens x K. This is the format that the pplx + i.e. E x max_num_tokens x K. This is the format that the batched dispatch/combine kernels use. """ @@ -880,7 +880,7 @@ def batched_moe_kernel_quantize_input( class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute): """ A Triton based MoE expert class that operates on expert batched format, - i.e. E x max_num_tokens x K. This is the format that the pplx + i.e. E x max_num_tokens x K. This is the format that the batched dispatch/combine kernels use. """ diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index c2c0e809d7..043b5ef266 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1172,9 +1172,9 @@ class FusedMoEModularKernel(torch.nn.Module): # This happens when none of the tokens from the all2all reach this # EP rank. Also, note that this is only relevant for CUDAGraph # incompatible all2all kernels like the DeepEP high-throughput - # kernels. CUDAGraph compatible all2all kernels like the pplx - # kernels and the DeepEP low-latency kernels are always batched - # and can never run into the tensor.numel() == 0 case. + # kernels. CUDAGraph compatible all2all kernels like the DeepEP + # low-latency kernels are always batched and can never run into + # the tensor.numel() == 0 case. if M_full == 0: assert num_chunks == 0 workspace13 = None diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index ee7db88ccb..b4f4b74ca2 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -143,10 +143,7 @@ def select_nvfp4_moe_backend( # NOTE(rob): this is kind of a hack. We need to peak into # the prepare-finalize selection to determine if we are using # the batched or standard expert format. - use_batched = ( - config.moe_parallel_config.use_deepep_ll_kernels - or config.moe_parallel_config.use_pplx_kernels - ) + use_batched = config.moe_parallel_config.use_deepep_ll_kernels activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if use_batched diff --git a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py deleted file mode 100644 index 289ac0d141..0000000000 --- a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py +++ /dev/null @@ -1,373 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import pplx_kernels as pplx -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceDelegate, -) -from vllm.model_executor.layers.fused_moe.utils import ( - _validate_scale_shape, - moe_kernel_quantize_input, -) -from vllm.utils.math_utils import cdiv, round_up - -logger = init_logger(__name__) - - -def pplx_hidden_dim_scale_bytes( - max_num_tokens: int, - hidden_dim: int, - in_dtype: torch.dtype, - quant_dtype: torch.dtype | str | None, - per_act_token_quant: bool, - block_shape: list[int] | None, -): - # All pplx byte sizes must be 16-byte aligned. - align = 16 - - # For blocked per token: set to - # cdiv(hidden_dim, block_size) * sizeof(float32) - # For per-token: set to 4 * sizeof(float32) (x4 for alignment) - if quant_dtype is not None: - assert isinstance(quant_dtype, torch.dtype) - assert quant_dtype.itemsize == 1 - hidden_dim_bytes = hidden_dim * quant_dtype.itemsize - elem_size = torch.float32.itemsize - - if per_act_token_quant: - # per-token (M x 1) - assert block_shape is None - hidden_scale_bytes = elem_size - elif block_shape is not None: - # per-group (M x K_tiles) - block_size = block_shape[1] - num_blocks = cdiv(hidden_dim, block_size) - hidden_scale_bytes = num_blocks * elem_size - else: - # per-tensor (1 x 1) - hidden_scale_bytes = elem_size - else: - hidden_dim_bytes = hidden_dim * in_dtype.itemsize - hidden_scale_bytes = 0 - - return ( - round_up(hidden_dim_bytes, align), - round_up(hidden_scale_bytes, align), - ) - - -class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): - """PPLX-based prepare and finalize for expert parallelism.""" - - def __init__( - self, - a2a: pplx.AllToAll, - max_num_tokens: int, - num_local_experts: int, - num_dispatchers: int, - ): - super().__init__() - assert max_num_tokens > 0 - assert num_local_experts > 0 - self.a2a = a2a - self.max_num_tokens = max_num_tokens - self.num_local_experts = num_local_experts - self.num_dispatchers_ = num_dispatchers - - @property - def activation_format(self) -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.BatchedExperts - - def max_num_tokens_per_rank(self) -> int | None: - return self.max_num_tokens - - def topk_indices_dtype(self) -> torch.dtype | None: - return torch.uint32 - - def num_dispatchers(self) -> int: - return self.num_dispatchers_ - - def output_is_reduced(self) -> bool: - return True - - 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, - ) -> tuple[Callable, mk.ReceiverType]: - if defer_input_quant: - raise NotImplementedError( - f"{self.__class__.__name__} does not support defer_input_quant=True. " - "Please select an MoE kernel that accepts quantized inputs." - ) - - num_tokens = a1.size(0) # M - hidden_dim = a1.size(-1) # K - - assert topk_ids.size(0) == num_tokens - # expert_map should be None because with expert map, -1 id is used for - # non-local token; this causes error when casting ids to the - # topk_indices_dtype() int32 - # - if expert_map is not None: - logger.warning_once( - "The PPLX backend does not support expert mapping. " - "The provided `expert_map` will be ignored." - ) - expert_map = None # noqa: F841 - - # Is this always going to be a1.device? - device = a1.device - - if apply_router_weight_on_input: - topk = topk_ids.size(1) - # TODO: this only works for topK=1, will need to update for topK>1 - assert topk == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - a1 = a1 * topk_weights.to(a1.dtype) - - repeat_cols = 4 - repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0) - # TODO(bnell): always pass quant_config.a1_scale? - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - (None if quant_config.per_act_token_quant else 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, - ) - - _validate_scale_shape( - a1q, a1q_scale, quant_config.per_act_token_quant, quant_config.block_shape - ) - - orig_a_scale_block_shape: int | None = None - - if a1q_scale is not None: - scalar_scales = a1q_scale.numel() == 1 - - # pplx requires 2-d scales even for scalar scales - if a1q_scale.dim() <= 1: - assert scalar_scales - a1q_scale = a1q_scale.view(1, 1) - - orig_a_scale_block_shape = a1q_scale.shape[-1] - - if not quant_config.is_block_quantized: - # TODO (bnell): use group_broadcast instead? - a1q_scale = a1q_scale.repeat(repeat_rows, repeat_cols) - - assert a1q_scale is None or a1q_scale.ndim == 2, ( - f"{0 if a1q_scale is None else (a1q_scale.ndim, a1q_scale.shape)}" - ) - - expert_num_tokens = torch.empty( - self.num_local_experts, - dtype=torch.int32, - device=device, - ) - - expert_x = torch.empty( - ( - self.num_local_experts, - self.max_num_tokens * self.num_dispatchers(), - hidden_dim, - ), - dtype=a1q.dtype, - device=device, - ) - - expert_x_scale: torch.Tensor | None = None - if a1q.dtype.itemsize == 1: - if quant_config.is_per_act_token: - # (M x 1) -> (E x M x K) - final_dim = expert_x.size(2) - elif quant_config.is_per_tensor: - # (1 x 1) -> (E x 1 x 1) - final_dim = 1 - else: - # (M x K_tiles) -> (E x M x K_tiles) - assert quant_config.block_shape is not None - num_blocks = cdiv(expert_x.size(2), quant_config.block_shape[1]) - final_dim = num_blocks - - expert_x_scale_shape = ( - self.num_local_experts, - expert_x.size(1), - round_up(final_dim, 4), # round up for alignment - ) - - expert_x_scale = torch.empty( - expert_x_scale_shape, - dtype=torch.float32, - device=expert_x.device, - ) - - # This argument is optional, defaults to indices.size(0) - # There's not much point setting this unless it is != indices.size(0) - bound_m: torch.Tensor | None = None - - self.a2a.dispatch( - out_expert_num_tokens=expert_num_tokens, - out_expert_x=expert_x, - out_expert_x_scale=expert_x_scale, - dp_x=a1q, - dp_x_scale=a1q_scale, - indices=topk_ids, - bound_m=bound_m, - do_send=True, - do_recv=False, - ) - - hook = lambda: self.a2a.dispatch( - out_expert_num_tokens=expert_num_tokens, - out_expert_x=expert_x, - out_expert_x_scale=expert_x_scale, - dp_x=a1q, - dp_x_scale=a1q_scale, - indices=topk_ids, - bound_m=bound_m, - do_send=False, - do_recv=True, - ) - - return ( - hook, - lambda: self._receiver( - expert_num_tokens, - expert_x, - expert_x_scale, - orig_a_scale_block_shape, - ), - ) - - def _receiver( - self, - expert_num_tokens: torch.Tensor, - expert_x: torch.Tensor, - expert_x_scale: torch.Tensor | None, - orig_a_scale_block_shape: int | None, - ) -> mk.PrepareResultType: - if expert_x_scale is not None: - expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape] - assert expert_x_scale.ndim == 3 - - expert_tokens_meta = mk.ExpertTokensMetadata( - expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None - ) - - return expert_x, expert_x_scale, expert_tokens_meta, None, None - - 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: - hook, receiver = self.prepare_async( - a1, - topk_weights, - topk_ids, - num_experts, - expert_map, - apply_router_weight_on_input, - quant_config, - defer_input_quant=defer_input_quant, - ) - hook() - return receiver() - - 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: - assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate), ( - "Weight application and reduction happens in the combine kernel." - ) - - # This argument is optional - # There's not much point setting this unless it is != topk_ids.size(0) - bound_m: torch.Tensor | None = None - - # TODO (bnell): fails in test_pplx_moe.py, figure out what's going on - # num_tokens = output.size(0) # M - # assert topk_ids.size(0) == num_tokens, ( - # f"{topk_ids.size(0)} == {num_tokens}") - assert topk_ids.size() == topk_weights.size(), ( - f"{topk_ids.size()} == {topk_weights.size()}" - ) - assert output.size(0) <= self.max_num_tokens, ( - f"{output.size(0)} <= {self.max_num_tokens}" - ) - assert output.size(1) == fused_expert_output.size(-1) - - # Set weights to 1 if we did them in dispatch. This is hacky. - if apply_router_weight_on_input: - topk_weights = torch.ones_like(topk_weights) - - topk_ids_u32 = topk_ids.view(dtype=torch.uint32) - - self.a2a.combine( - out_tokens=output, - indices=topk_ids_u32, - weights=topk_weights, - expert_y=fused_expert_output, - bound_m=bound_m, - do_send=True, - do_recv=False, - ) - - return lambda: self.a2a.combine( - out_tokens=output, - indices=topk_ids_u32, - weights=topk_weights, - expert_y=fused_expert_output, - bound_m=bound_m, - do_send=False, - do_recv=True, - ) - - 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: - receiver = self.finalize_async( - output, - fused_expert_output, - topk_weights, - topk_ids, - apply_router_weight_on_input, - weight_and_reduce_impl, - ) - receiver() diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 7e25c96874..9c2adf7996 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -216,8 +216,7 @@ class DefaultMoERunner(MoERunner): @property def use_dp_chunking(self) -> bool: return ( - self.moe_config.moe_parallel_config.use_pplx_kernels - or self.moe_config.moe_parallel_config.use_deepep_ll_kernels + self.moe_config.moe_parallel_config.use_deepep_ll_kernels or self.moe_config.moe_parallel_config.use_mori_kernels or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 99d4038ec3..d7b50aea2a 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -14,10 +14,11 @@ class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce): implementation does not perform weight application and reduction but cannot address the needs of all the compatible PrepareAndFinalize implementations. - For example, BatchedTritonExperts is compatible with both - PplxPrepareAndFinalize and BatchedPrepareAndFinalize. PplxPrepareAndFinalize - does the weight-application + reduction as part of the pplx combine kernel. - But the BatchedPrepareAndFinalize needs an implementation. To facilitate + For example, BatchedTritonExperts is compatible with both batched + PrepareAndFinalize implementations like DeepEPLLPrepareAndFinalize and + BatchedPrepareAndFinalize. Some PrepareAndFinalize implementations do + the weight-application + reduction as part of the combine kernel, while + BatchedPrepareAndFinalize needs an explicit implementation. To facilitate this case, the BatchedTritonExperts could use TopKWeightAndReduceDelegate so the PrepareAndFinalize implementations could choose how to weight + reduce. diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index d81f0f80d2..9318bedfff 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -798,7 +798,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # batched activation format. As self.fused_experts is not # initialized at this point, we resort to checking the MoE config # directly. - is_batched_moe = self.moe.use_pplx_kernels or self.moe.use_deepep_ll_kernels + is_batched_moe = self.moe.use_deepep_ll_kernels if is_batched_moe: num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8 else: diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 4739120d47..91e724012d 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -402,11 +402,6 @@ def _has_module(module_name: str) -> bool: return importlib.util.find_spec(module_name) is not None -def has_pplx() -> bool: - """Whether the optional `pplx_kernels` package is available.""" - return _has_module("pplx_kernels") - - def has_deep_ep() -> bool: """Whether the optional `deep_ep` package is available.""" return _has_module("deep_ep")