A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 16:30 UTC
[WideEP] Remove pplx all2all backend (#33724)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Tyler Michael Smith committed
Feb 26, 2026 at 17:30 UTC
eb19955c37089056883831838dc155340ae67edd
39 files changed
+107
-2069
.buildkite/test_areas/kernels.yaml
+11
-2
@@ -155,5 +155,14 @@ steps:
155
commands:
156
- pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py
157
- pytest -v -s kernels/moe/test_deepep_moe.py
158
- - pytest -v -s kernels/moe/test_pplx_cutlass_moe.py
159
- # - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main
158
+
159
+- label: Kernels Fp4 MoE Test (B200)
160
+ timeout_in_minutes: 60
161
+ device: b200
162
+ num_devices: 1
163
+ optional: true
164
+ commands:
165
+ - pytest -v -s kernels/moe/test_cutedsl_moe.py
166
+ - pytest -v -s kernels/moe/test_flashinfer_moe.py
167
+ - pytest -v -s kernels/moe/test_nvfp4_moe.py
168
+ - pytest -v -s kernels/moe/test_ocp_mx_moe.py
CMakeLists.txt
+1
-1
@@ -725,7 +725,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
725
# CUTLASS MoE kernels
726
727
# The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works
728
- # on Hopper). get_cutlass_(pplx_)moe_mm_data should only be compiled
728
+ # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled
729
# if it's possible to compile MoE kernels that use its output.
730
cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}")
731
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
csrc/ops.h
+7
-7
@@ -269,13 +269,13 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
269
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
270
const int64_t n, const int64_t k, const bool swap_ab);
271
272
-void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
273
- torch::Tensor& problem_sizes1,
274
- torch::Tensor& problem_sizes2,
275
- const torch::Tensor& expert_num_tokens,
276
- const int64_t num_local_experts,
277
- const int64_t padded_m, const int64_t n,
278
- const int64_t k);
272
+void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets,
273
+ torch::Tensor& problem_sizes1,
274
+ torch::Tensor& problem_sizes2,
275
+ const torch::Tensor& expert_num_tokens,
276
+ const int64_t num_local_experts,
277
+ const int64_t padded_m, const int64_t n,
278
+ const int64_t k);
279
280
void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
281
torch::Tensor const& b,
csrc/quantization/w8a8/cutlass/moe/moe_data.cu
+11
-15
@@ -263,12 +263,10 @@ void get_cutlass_moe_mm_data_caller(
263
}
264
265
template <bool SWAP_AB>
266
-__global__ void compute_pplx_data(int32_t* expert_offsets,
267
- int32_t* problem_sizes1,
268
- int32_t* problem_sizes2,
269
- const int32_t* __restrict__ expert_num_tokens,
270
- const int padded_m, const int n,
271
- const int k) {
266
+__global__ void compute_batched_moe_data(
267
+ int32_t* expert_offsets, int32_t* problem_sizes1, int32_t* problem_sizes2,
268
+ const int32_t* __restrict__ expert_num_tokens, const int padded_m,
269
+ const int n, const int k) {
270
int expert_idx = threadIdx.x;
271
expert_offsets[expert_idx] = expert_idx * padded_m;
272
@@ -289,24 +287,22 @@ __global__ void compute_pplx_data(int32_t* expert_offsets,
287
}
288
}
289
292
-void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,
293
- torch::Tensor& problem_sizes1,
294
- torch::Tensor& problem_sizes2,
295
- const torch::Tensor& expert_num_tokens,
296
- const int64_t num_local_experts,
297
- const int64_t padded_m,
298
- const int64_t n, const int64_t k) {
290
+void get_cutlass_batched_moe_mm_data_caller(
291
+ torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1,
292
+ torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens,
293
+ const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
294
+ const int64_t k) {
295
auto stream = at::cuda::getCurrentCUDAStream(expert_offsets.device().index());
296
297
if (num_local_experts * padded_m > SWAP_AB_THRESHOLD) {
302
- compute_pplx_data<false><<<1, num_local_experts, 0, stream>>>(
298
+ compute_batched_moe_data<false><<<1, num_local_experts, 0, stream>>>(
299
static_cast<int32_t*>(expert_offsets.data_ptr()),
300
static_cast<int32_t*>(problem_sizes1.data_ptr()),
301
static_cast<int32_t*>(problem_sizes2.data_ptr()),
302
static_cast<const int32_t*>(expert_num_tokens.data_ptr()), padded_m, n,
303
k);
304
} else {
309
- compute_pplx_data<true><<<1, num_local_experts, 0, stream>>>(
305
+ compute_batched_moe_data<true><<<1, num_local_experts, 0, stream>>>(
306
static_cast<int32_t*>(expert_offsets.data_ptr()),
307
static_cast<int32_t*>(problem_sizes1.data_ptr()),
308
static_cast<int32_t*>(problem_sizes2.data_ptr()),
csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu
+21
-22
@@ -82,13 +82,11 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
82
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
83
const int64_t n, const int64_t k, const bool swap_ab);
84
85
-void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,
86
- torch::Tensor& problem_sizes1,
87
- torch::Tensor& problem_sizes2,
88
- const torch::Tensor& expert_num_tokens,
89
- const int64_t num_local_experts,
90
- const int64_t padded_m,
91
- const int64_t n, const int64_t k);
85
+void get_cutlass_batched_moe_mm_data_caller(
86
+ torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1,
87
+ torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens,
88
+ const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
89
+ const int64_t k);
90
#endif
91
92
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(
317
version_num, ". Required capability: 90, 100, or 120");
318
}
319
322
-void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
323
- torch::Tensor& problem_sizes1,
324
- torch::Tensor& problem_sizes2,
325
- const torch::Tensor& expert_num_tokens,
326
- const int64_t num_local_experts,
327
- const int64_t padded_m, const int64_t n,
328
- const int64_t k) {
320
+void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets,
321
+ torch::Tensor& problem_sizes1,
322
+ torch::Tensor& problem_sizes2,
323
+ const torch::Tensor& expert_num_tokens,
324
+ const int64_t num_local_experts,
325
+ const int64_t padded_m, const int64_t n,
326
+ const int64_t k) {
327
// This function currently gets compiled only if we have a valid cutlass moe
328
// mm to run it for.
329
int32_t version_num = get_sm_version_num();
330
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
331
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
332
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
335
- get_cutlass_pplx_moe_mm_data_caller(expert_offsets, problem_sizes1,
336
- problem_sizes2, expert_num_tokens,
337
- num_local_experts, padded_m, n, k);
333
+ get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
334
+ problem_sizes2, expert_num_tokens,
335
+ num_local_experts, padded_m, n, k);
336
return;
337
#endif
340
- TORCH_CHECK_NOT_IMPLEMENTED(
341
- false,
342
- "No compiled get_cutlass_pplx_moe_mm_data: no cutlass_scaled_mm kernel "
343
- "for CUDA device capability: ",
344
- version_num, ". Required capability: 90, 100, or 120");
338
+ TORCH_CHECK_NOT_IMPLEMENTED(false,
339
+ "No compiled get_cutlass_batched_moe_mm_data: no "
340
+ "cutlass_scaled_mm kernel "
341
+ "for CUDA device capability: ",
342
+ version_num,
343
+ ". Required capability: 90, 100, or 120");
344
}
345
346
void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
csrc/torch_bindings.cpp
+4
-4
@@ -489,19 +489,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
489
&get_cutlass_moe_mm_problem_sizes_from_expert_offsets);
490
491
// A function that computes data required to run fused MoE with w8a8 grouped
492
- // GEMM and PPLX. It takes expert_num_tokens and non_zero_expert_idxs
492
+ // GEMM in batched expert format. It takes expert_num_tokens
493
// as an input, and computes expert_offsets (token start indices of each
494
// expert). In addition to this, it computes problem sizes for each expert's
495
// multiplication used by the two mms called from fused MoE operation.
496
ops.def(
497
- "get_cutlass_pplx_moe_mm_data(Tensor! expert_offsets, "
497
+ "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, "
498
" Tensor! problem_sizes1, "
499
" Tensor! problem_sizes2, "
500
" Tensor expert_num_tokens, "
501
" int num_local_experts, int padded_m, "
502
" int n, int k) -> ()");
503
- ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA,
504
- &get_cutlass_pplx_moe_mm_data);
503
+ ops.impl("get_cutlass_batched_moe_mm_data", torch::kCUDA,
504
+ &get_cutlass_batched_moe_mm_data);
505
506
// Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3)
507
ops.def(
docker/Dockerfile
+3
-5
@@ -308,7 +308,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
308
#################### CSRC BUILD IMAGE ####################
309
310
#################### EXTENSIONS BUILD IMAGE ####################
311
-# Build DeepGEMM, pplx-kernels, DeepEP - runs in PARALLEL with csrc-build
311
+# Build DeepGEMM, DeepEP - runs in PARALLEL with csrc-build
312
# This stage is independent and doesn't affect csrc cache
313
FROM base AS extensions-build
314
ARG CUDA_VERSION
@@ -335,10 +335,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
335
# Ensure the wheel dir exists so COPY won't fail when DeepGEMM is skipped
336
RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
337
338
-# Build pplx-kernels and DeepEP wheels
338
+# Build DeepEP wheels
339
COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
340
# Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
341
-ARG PPLX_COMMIT_HASH=12cecfd
341
ARG DEEPEP_COMMIT_HASH=73b6ea4
342
ARG NVSHMEM_VER
343
RUN --mount=type=cache,target=/root/.cache/uv \
@@ -347,7 +346,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
346
/tmp/install_python_libraries.sh \
347
--workspace /tmp/ep_kernels_workspace \
348
--mode wheel \
350
- ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \
349
${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} \
350
${NVSHMEM_VER:+--nvshmem-ver "$NVSHMEM_VER"} && \
351
find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
@@ -676,7 +674,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
674
# Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH
675
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
676
679
-# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage
677
+# Install EP kernels wheels (DeepEP) that have been built in the `build` stage
678
RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
679
--mount=type=cache,target=/root/.cache/uv \
680
uv pip install --system ep_kernels/dist/*.whl --verbose \
docker/versions.json
-3
@@ -52,9 +52,6 @@
52
"DEEPGEMM_GIT_REF": {
53
"default": "477618cd51baffca09c4b0b87e97c03fe827ef03"
54
},
55
- "PPLX_COMMIT_HASH": {
56
- "default": "12cecfd"
57
- },
55
"DEEPEP_COMMIT_HASH": {
56
"default": "73b6ea4"
57
},
docs/design/fused_moe_modular_kernel.md
+3
-4
@@ -15,7 +15,7 @@ Based on the format of the input activations, FusedMoE implementations are broad
15
The input activation format completely depends on the All2All Dispatch being used.
16
17
* 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.
18
-* 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.
18
+* 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.
19
20
The FusedMoE operation is generally made of multiple operations, in both the Contiguous and Batched variants, as described in the diagrams below
21
@@ -132,7 +132,6 @@ class FusedMoEModularKernel:
132
133
Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & Combine implementation / kernel. For example,
134
135
-* PplxPrepareAndFinalize type is backed by Pplx All2All kernels,
135
* DeepEPHTPrepareAndFinalize type is backed by DeepEP High-Throughput All2All kernels, and
136
* DeepEPLLPrepareAndFinalize type is backed by DeepEP Low-Latency All2All kernels.
137
@@ -229,7 +228,7 @@ Doing this will add the new implementation to the test suite.
228
### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` Compatibility
229
230
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.
232
-Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts`
231
+Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts`
232
As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` compatibility. When invoked
233
with incompatible types, the script will error.
234
@@ -238,7 +237,7 @@ with incompatible types, the script will error.
237
Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py)
238
The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible
239
`FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` types.
241
-Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts`
240
+Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts`
241
242
## FusedMoEPrepareAndFinalize Implementations
243
docs/design/moe_kernel_features.md
+2
-3
@@ -33,7 +33,6 @@ th {
33
| Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
34
|---------|--------------------|--------------|---------------|-------|-----------------------|-----------|
35
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
36
-| pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
36
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
37
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
38
| 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.
67
68
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.
69
71
-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`.
70
+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`.
71
72
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.
73
@@ -110,5 +109,5 @@ The following table shows "families" of modular kernels that are intended to wor
109
| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |
110
|---------|-----------------------------------------|----------------------------------------------|
111
| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`, </br>`MarlinExperts` |
113
-| deepep_low_latency,</br>pplx | `DeepEPLLPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`CutlassBatchedExpertsFp8`,</br>`BatchedMarlinExperts` |
112
+| deepep_low_latency | `DeepEPLLPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`CutlassBatchedExpertsFp8`,</br>`BatchedMarlinExperts` |
113
| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` |
docs/governance/committers.md
+1
-1
@@ -154,7 +154,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
154
- FlashAttention: @LucasWilkinson
155
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
156
- Blackwell Kernels: @mgoin, @yewentao256
157
-- DeepEP/DeepGEMM/pplx: @mgoin, @yewentao256
157
+- DeepEP/DeepGEMM: @mgoin, @yewentao256
158
159
### Integrations
160
docs/serving/expert_parallel_deployment.md
+3
-6
@@ -8,7 +8,7 @@ EP is typically coupled with Data Parallelism (DP). While DP can be used indepen
8
9
Before using EP, you need to install the necessary dependencies. We are actively working on making this easier in the future:
10
11
-1. **Install DeepEP and pplx-kernels**: Set up host environment following vLLM's guide for EP kernels [here](../../tools/ep_kernels).
11
+1. **Install DeepEP**: Set up host environment following vLLM's guide for EP kernels [here](../../tools/ep_kernels).
12
2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation).
13
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/).
14
@@ -19,7 +19,6 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
19
| Backend | Use Case | Features | Best For |
20
|---------|----------|----------|----------|
21
| `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration |
22
-| `pplx` | Single node | Chunked prefill support, efficient intra-node communication | Single-node deployments, development |
22
| `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
23
| `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
24
| `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):
70
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.
71
72
```bash
74
-# Single node EP deployment with pplx backend
73
+# Single node EP deployment
74
vllm serve deepseek-ai/DeepSeek-V3-0324 \
75
--tensor-parallel-size 1 \ # Tensor parallelism across 1 GPU
76
--data-parallel-size 8 \ # Data parallelism across 8 processes
78
- --enable-expert-parallel \ # Enable expert parallelism
79
- --all2all-backend pplx # Use pplx communication backend
77
+ --enable-expert-parallel # Enable expert parallelism
78
```
79
80
## Multi-Node Deployment
@@ -197,7 +195,6 @@ vllm serve deepseek-ai/DeepSeek-V3-0324 \
195
--tensor-parallel-size 1 \ # Tensor parallelism
196
--data-parallel-size 8 \ # Data parallelism
197
--enable-expert-parallel \ # Enable EP
200
- --all2all-backend pplx \ # Use pplx communication backend
198
--enable-eplb \ # Enable load balancer
199
--eplb-config '{"window_size":1000,"step_interval":3000,"num_redundant_experts":2,"log_balancedness":true}'
200
```
examples/online_serving/elastic_ep/serve_deepseek_v2.sh
+1
-1
@@ -64,7 +64,7 @@ vllm serve "$MODEL_NAME" \
64
--enforce-eager \
65
--enable-expert-parallel \
66
--enable-eplb \
67
- --all2all-backend pplx \
67
+ --all2all-backend allgather_reducescatter \
68
--num-redundant-experts "$REDUNDANT_EXPERTS" \
69
--trust-remote-code \
70
--host "$HOST" \
tests/kernels/moe/modular_kernel_tools/common.py
-7
@@ -37,7 +37,6 @@ from vllm.utils.import_utils import (
37
has_deep_ep,
38
has_deep_gemm,
39
has_mori,
40
- has_pplx,
40
)
41
42
from .mk_objects import (
@@ -206,10 +205,6 @@ class Config:
205
info = expert_info(self.fused_experts_type)
206
return info.needs_deep_gemm
207
209
- def needs_pplx(self):
210
- info = prepare_finalize_info(self.prepare_finalize_type)
211
- return info.backend == "pplx"
212
-
208
def needs_deep_ep(self):
209
info = prepare_finalize_info(self.prepare_finalize_type)
210
return (
@@ -290,8 +285,6 @@ class Config:
285
return False, "Needs DeepEP, but DeepEP not available."
286
if self.needs_deep_gemm() and not has_deep_gemm():
287
return False, "Needs DeepGEMM, but DeepGEMM not available."
293
- if self.needs_pplx() and not has_pplx(): # noqa: SIM103
294
- return False, "Needs PPLX, but PPLX not available."
288
if self.needs_aiter() and not has_aiter(): # noqa: SIM103
289
return False, "Needs Aiter, but Aiter not available."
290
if self.needs_mori() and not has_mori(): # noqa: SIM103
tests/kernels/moe/modular_kernel_tools/mk_objects.py
-14
@@ -39,7 +39,6 @@ from vllm.utils.import_utils import (
39
has_deep_ep,
40
has_deep_gemm,
41
has_mori,
42
- has_pplx,
42
)
43
44
@@ -238,19 +237,6 @@ if has_mori():
237
supports_apply_weight_on_input=False,
238
)
239
241
-if has_pplx():
242
- from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
243
- PplxPrepareAndFinalize,
244
- )
245
-
246
- register_prepare_and_finalize(
247
- PplxPrepareAndFinalize,
248
- batched_format,
249
- common_float_and_int_types,
250
- blocked_quantization_support=True,
251
- backend="pplx",
252
- )
253
-
240
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
241
from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501
242
FlashInferCutlassMoEPrepareAndFinalize,
tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py
+1
-1
@@ -125,7 +125,7 @@ if __name__ == "__main__":
125
description=(
126
"Run single prepare-finalize & fused-experts combination test"
127
"Example : python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel " # noqa: E501
128
- "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts"
128
+ "--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts"
129
)
130
)
131
args = parser.parse_args()
tests/kernels/moe/test_modular_kernel_combinations.py
+4
-4
@@ -14,7 +14,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14
from vllm.config import VllmConfig, set_current_vllm_config
15
from vllm.platforms import current_platform
16
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
17
-from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx
17
+from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
18
from vllm.utils.torch_utils import cuda_device_count_stateless, set_random_seed
19
from vllm.v1.worker.workspace import init_workspace_manager
20
@@ -39,12 +39,12 @@ from .modular_kernel_tools.parallel_utils import (
39
)
40
41
has_any_multi_gpu_package = (
42
- has_deep_ep() or has_deep_gemm() or has_pplx() or has_flashinfer_cutlass_fused_moe()
42
+ has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe()
43
)
44
45
meets_multi_gpu_requirements = pytest.mark.skipif(
46
not has_any_multi_gpu_package,
47
- reason="Requires deep_ep or deep_gemm or pplx or flashinfer packages",
47
+ reason="Requires deep_ep or deep_gemm or flashinfer packages",
48
)
49
50
if current_platform.is_fp8_fnuz():
@@ -341,7 +341,7 @@ if __name__ == "__main__":
341
description=(
342
"Run single prepare-finalize & fused-experts combination test"
343
"Example : python3 -m tests.kernels.moe.test_modular_kernel_combinations "
344
- "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts"
344
+ "--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts"
345
)
346
)
347
args = parser.parse_args()
tests/kernels/moe/test_pplx_cutlass_moe.py
deleted
-365
@@ -1,365 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-
5
-import pytest
6
-import torch
7
-
8
-from tests.kernels.utils import torch_experts
9
-from vllm import _custom_ops as ops
10
-from vllm.config import VllmConfig, set_current_vllm_config
11
-from vllm.model_executor.layers.fused_moe import fused_topk
12
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
13
-from vllm.model_executor.layers.fused_moe.config import (
14
- FusedMoEConfig,
15
- FusedMoEParallelConfig,
16
- RoutingMethodType,
17
- fp8_w8a8_moe_quant_config,
18
-)
19
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8
20
-from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
21
-from vllm.platforms import current_platform
22
-from vllm.utils.math_utils import cdiv
23
-from vllm.utils.torch_utils import set_random_seed
24
-from vllm.v1.worker.workspace import init_workspace_manager
25
-
26
-from ...utils import multi_gpu_test
27
-from .parallel_utils import ProcessGroupInfo, parallel_launch
28
-
29
-try:
30
- from pplx_kernels import AllToAll
31
- from pplx_kernels.nvshmem import (
32
- nvshmem_alloc_empty_unique_id,
33
- nvshmem_finalize,
34
- nvshmem_get_unique_id,
35
- nvshmem_init,
36
- )
37
-
38
- has_pplx = True
39
-except ImportError:
40
- has_pplx = False
41
-
42
-requires_pplx = pytest.mark.skipif(
43
- not has_pplx,
44
- reason="Requires PPLX kernels",
45
-)
46
-
47
-NUM_EXPERTS = [40, 64]
48
-TOP_KS = [6, 8]
49
-
50
-
51
-def rank_chunk(num, r, w):
52
- rem = num % w
53
- return (num // w) + (1 if r < rem else 0)
54
-
55
-
56
-def chunk_by_rank(t, r, w):
57
- num = t.shape[0]
58
- chunk = rank_chunk(num, r, w)
59
- rem = num % w
60
- if rem == 0 or r < rem:
61
- return t[(r * chunk) : (r + 1) * chunk].contiguous()
62
- else:
63
- long_chunks = (num // w + 1) * rem
64
- short_chunks = (r - rem) * chunk
65
- start = long_chunks + short_chunks
66
- return t[start : start + chunk].contiguous()
67
-
68
-
69
-def pplx_cutlass_moe(
70
- pgi: ProcessGroupInfo,
71
- dp_size: int,
72
- a: torch.Tensor,
73
- w1: torch.Tensor,
74
- w2: torch.Tensor,
75
- w1_scale: torch.Tensor,
76
- w2_scale: torch.Tensor,
77
- topk_weights: torch.Tensor,
78
- topk_ids: torch.Tensor,
79
- a1_scale: torch.Tensor,
80
- out_dtype,
81
- per_act_token: bool,
82
- per_out_ch: bool,
83
- group_name: str | None,
84
-):
85
- from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
86
- PplxPrepareAndFinalize,
87
- )
88
-
89
- init_workspace_manager(torch.cuda.current_device())
90
-
91
- assert torch.cuda.current_device() == pgi.local_rank
92
-
93
- num_tokens, hidden_dim = a.shape
94
- intermediate_dim = w2.shape[2]
95
- num_experts = w1.shape[0]
96
- block_size = hidden_dim # TODO support more cases
97
- device = pgi.device
98
- rank = pgi.rank
99
- world_size = pgi.world_size
100
- rank_num_tokens = rank_chunk(num_tokens, rank, world_size)
101
- max_num_tokens = rank_chunk(num_tokens, 0, world_size)
102
- topk = topk_ids.shape[1]
103
-
104
- if block_size == hidden_dim:
105
- scale_elems = 4 # hack to circumvent pplx data format requirements
106
- else:
107
- scale_elems = (hidden_dim + block_size - 1) // block_size
108
-
109
- args = dict(
110
- max_num_tokens=max_num_tokens,
111
- num_experts=num_experts,
112
- experts_per_token=topk,
113
- rank=rank,
114
- world_size=world_size,
115
- dp_size=dp_size,
116
- hidden_dim=hidden_dim,
117
- hidden_dim_bytes=hidden_dim, # because a.dtype.itemsize == 1
118
- hidden_dim_scale_bytes=scale_elems * torch.float32.itemsize,
119
- )
120
-
121
- if group_name is None:
122
- ata = AllToAll.internode(**args)
123
- else:
124
- args["group_name"] = group_name
125
- ata = AllToAll.intranode(**args)
126
-
127
- w1 = w1.to(device)
128
- w2 = w2.to(device)
129
- w1_scale = w1_scale.to(device)
130
- w2_scale = w2_scale.to(device)
131
- a1_scale = a1_scale.to(device)
132
-
133
- assert num_experts % world_size == 0
134
- num_local_experts = cdiv(num_experts, world_size)
135
- num_dispatchers = pgi.world_size // dp_size
136
-
137
- prepare_finalize = PplxPrepareAndFinalize(
138
- ata,
139
- max_num_tokens=max_num_tokens,
140
- num_local_experts=num_local_experts,
141
- num_dispatchers=num_dispatchers,
142
- )
143
-
144
- def make_moe_config() -> FusedMoEConfig:
145
- return FusedMoEConfig(
146
- num_experts=num_experts,
147
- experts_per_token=topk,
148
- hidden_dim=hidden_dim,
149
- intermediate_size_per_partition=intermediate_dim,
150
- num_local_experts=num_local_experts,
151
- num_logical_experts=num_experts,
152
- moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
153
- activation=MoEActivation.SILU,
154
- in_dtype=torch.bfloat16,
155
- device="cuda",
156
- routing_method=RoutingMethodType.Llama4,
157
- )
158
-
159
- experts = CutlassBatchedExpertsFp8(
160
- moe_config=make_moe_config(),
161
- quant_config=fp8_w8a8_moe_quant_config(
162
- per_act_token_quant=per_act_token,
163
- per_out_ch_quant=per_out_ch,
164
- w1_scale=chunk_by_rank(w1_scale, rank, world_size),
165
- w2_scale=chunk_by_rank(w2_scale, rank, world_size),
166
- a1_scale=chunk_by_rank(a1_scale, rank, world_size)
167
- if per_act_token
168
- else a1_scale[rank],
169
- ),
170
- max_num_tokens=max_num_tokens,
171
- num_dispatchers=num_dispatchers,
172
- )
173
-
174
- fused_cutlass_experts = FusedMoEModularKernel(
175
- prepare_finalize,
176
- experts,
177
- inplace=False,
178
- )
179
-
180
- a_chunk = chunk_by_rank(a, rank, world_size).to(device)
181
- chunk_topk_weight = chunk_by_rank(topk_weights, rank, world_size).to(device)
182
- chunk_topk_ids = (
183
- chunk_by_rank(topk_ids, rank, world_size).to(torch.uint32).to(device)
184
- )
185
-
186
- out = fused_cutlass_experts(
187
- a_chunk,
188
- chunk_by_rank(w1, rank, world_size),
189
- chunk_by_rank(w2, rank, world_size),
190
- chunk_topk_weight,
191
- chunk_topk_ids,
192
- global_num_experts=num_experts,
193
- expert_map=None, # TODO
194
- )
195
-
196
- torch.cuda.synchronize()
197
-
198
- ata.destroy()
199
-
200
- return out[:rank_num_tokens]
201
-
202
-
203
-vllm_config = VllmConfig()
204
-
205
-
206
-def _pplx_moe(
207
- pgi: ProcessGroupInfo,
208
- dp_size: int,
209
- a: torch.Tensor,
210
- w1: torch.Tensor,
211
- w2: torch.Tensor,
212
- w1_scale: torch.Tensor,
213
- w2_scale: torch.Tensor,
214
- topk_weights: torch.Tensor,
215
- topk_ids: torch.Tensor,
216
- a1_scale: torch.Tensor,
217
- out_dtype,
218
- a_full: torch.Tensor,
219
- w1_full: torch.Tensor,
220
- w2_full: torch.Tensor,
221
- per_act_token: bool,
222
- per_out_ch: bool,
223
- use_internode: bool,
224
-):
225
- try:
226
- if use_internode:
227
- uid = (
228
- nvshmem_get_unique_id()
229
- if pgi.rank == 0
230
- else nvshmem_alloc_empty_unique_id()
231
- )
232
- torch.distributed.broadcast(uid, src=0)
233
- nvshmem_init(uid, pgi.rank, pgi.world_size)
234
- else:
235
- group_ranks = list(range(pgi.world_size))
236
- cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
237
- group_name = cpu_group.group_name
238
-
239
- with set_current_vllm_config(vllm_config):
240
- torch_output = torch_experts(
241
- a_full, w1_full, w2_full, topk_weights, topk_ids
242
- )
243
- pplx_output = pplx_cutlass_moe(
244
- pgi,
245
- dp_size,
246
- a,
247
- w1,
248
- w2,
249
- w1_scale,
250
- w2_scale,
251
- topk_weights,
252
- topk_ids,
253
- a1_scale,
254
- out_dtype,
255
- per_act_token,
256
- per_out_ch,
257
- group_name,
258
- )
259
-
260
- torch_output = chunk_by_rank(torch_output, pgi.rank, pgi.world_size).to(
261
- pplx_output.device
262
- )
263
-
264
- # Uncomment if more debugging is needed
265
- # print("PPLX OUT:", pplx_output)
266
- # print("TORCH OUT:", torch_output)
267
-
268
- torch.testing.assert_close(pplx_output, torch_output, atol=0.05, rtol=0)
269
- finally:
270
- if use_internode:
271
- nvshmem_finalize()
272
-
273
-
274
-@pytest.mark.parametrize("m", [2, 224])
275
-@pytest.mark.parametrize("n", [3072])
276
-@pytest.mark.parametrize("k", [1536])
277
-@pytest.mark.parametrize("e", NUM_EXPERTS)
278
-@pytest.mark.parametrize("topk", TOP_KS)
279
-@pytest.mark.parametrize("per_act_token", [True, False])
280
-@pytest.mark.parametrize("per_out_ch", [True, False])
281
-@pytest.mark.parametrize("world_dp_size", [[2, 1]]) # , [4, 2]])
282
-@pytest.mark.parametrize("use_internode", [False])
283
-@multi_gpu_test(num_gpus=2)
284
-@pytest.mark.skipif(
285
- (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
286
- current_platform.get_device_capability()
287
- ),
288
- reason="Grouped gemm is not supported on this GPU type.",
289
-)
290
-@requires_pplx
291
-def test_cutlass_moe_pplx(
292
- m: int,
293
- n: int,
294
- k: int,
295
- e: int,
296
- topk: int,
297
- per_act_token: bool,
298
- per_out_ch: bool,
299
- world_dp_size: tuple[int, int],
300
- use_internode: bool,
301
-):
302
- set_random_seed(7)
303
-
304
- with set_current_vllm_config(vllm_config):
305
- dtype = torch.half
306
-
307
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10.0
308
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10.0
309
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10.0
310
-
311
- n_b_scales = 2 * n if per_out_ch else 1
312
- k_b_scales = k if per_out_ch else 1
313
-
314
- w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float8_e4m3fn)
315
- w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn)
316
- w1_scale = torch.empty((e, n_b_scales, 1), device="cuda", dtype=torch.float32)
317
- w2_scale = torch.empty((e, k_b_scales, 1), device="cuda", dtype=torch.float32)
318
-
319
- for expert in range(e):
320
- w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
321
- w1[expert], use_per_token_if_dynamic=per_out_ch
322
- )
323
- w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
324
- w2[expert], use_per_token_if_dynamic=per_out_ch
325
- )
326
-
327
- w1_d = torch.empty_like(w1)
328
- w2_d = torch.empty_like(w2)
329
- for expert in range(e):
330
- w1_d[expert] = (w1_q[expert].float() * w1_scale[expert]).half()
331
- w2_d[expert] = (w2_q[expert].float() * w2_scale[expert]).half()
332
-
333
- score = torch.randn((m, e), device="cuda", dtype=dtype)
334
- topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
335
-
336
- world_size, dp_size = world_dp_size
337
- a_scale1 = (
338
- torch.randn(
339
- (m if per_act_token else 1, 1), device="cuda", dtype=torch.float32
340
- )
341
- / 10.0
342
- )
343
- if not per_act_token:
344
- a_scale1 = a_scale1.repeat(world_size, 1)
345
-
346
- parallel_launch(
347
- world_size,
348
- _pplx_moe,
349
- dp_size,
350
- a,
351
- w1_q,
352
- w2_q,
353
- w1_scale,
354
- w2_scale,
355
- topk_weights,
356
- topk_ids,
357
- a_scale1,
358
- dtype,
359
- a,
360
- w1_d,
361
- w2_d,
362
- per_act_token,
363
- per_out_ch,
364
- use_internode,
365
- )
tests/kernels/moe/test_pplx_moe.py
deleted
-1021
@@ -1,1021 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for the MOE layers.
4
-
5
-Run `pytest tests/kernels/test_pplx_moe.py`.
6
-"""
7
-
8
-import copy
9
-import itertools
10
-import textwrap
11
-import traceback
12
-from collections.abc import Callable
13
-
14
-import pytest
15
-import torch
16
-
17
-try:
18
- from pplx_kernels import AllToAll
19
- from pplx_kernels.nvshmem import (
20
- nvshmem_alloc_empty_unique_id,
21
- nvshmem_finalize,
22
- nvshmem_get_unique_id,
23
- nvshmem_init,
24
- )
25
-
26
- has_pplx = True
27
-except ImportError:
28
- has_pplx = False
29
-
30
-from tests.kernels.moe.modular_kernel_tools.parallel_utils import _set_vllm_config
31
-from tests.kernels.moe.utils import (
32
- make_dummy_moe_config,
33
- make_shared_experts,
34
- make_test_weights,
35
- naive_batched_moe,
36
-)
37
-from tests.kernels.quant_utils import dequant
38
-from tests.kernels.utils import torch_experts
39
-from vllm.config import VllmConfig, set_current_vllm_config
40
-from vllm.model_executor.layers.fused_moe import fused_topk, override_config
41
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
42
-from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts
43
-from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config
44
-from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
45
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
46
- TopKWeightAndReduceDelegate,
47
-)
48
-from vllm.utils.math_utils import round_up
49
-from vllm.utils.torch_utils import set_random_seed
50
-from vllm.v1.worker.workspace import init_workspace_manager
51
-
52
-from ...utils import multi_gpu_test
53
-from .parallel_utils import ProcessGroupInfo, parallel_launch
54
-
55
-requires_pplx = pytest.mark.skipif(
56
- not has_pplx,
57
- reason="Requires PPLX kernels",
58
-)
59
-
60
-BATCHED_MOE_MNK_FACTORS = [
61
- (1, 128, 128),
62
- (33, 2048, 128),
63
- (64, 128, 2048),
64
- (222, 128, 128),
65
- (222, 2048, 1024),
66
-]
67
-
68
-PPLX_COMBOS = [
69
- # TODO(bnell): figure out why this fails, seems to be test problem
70
- # (1, 128, 128),
71
- (2, 128, 512),
72
- (3, 1024, 2048),
73
- (4, 128, 128),
74
- (32, 1024, 512),
75
- (45, 512, 2048),
76
- (64, 1024, 512),
77
- (222, 2048, 1024),
78
- (256, 1408, 2048),
79
-]
80
-
81
-NUM_EXPERTS = [8, 64]
82
-TOP_KS = [1, 2, 6]
83
-DTYPES = [torch.float8_e4m3fn, torch.bfloat16]
84
-
85
-vllm_config = VllmConfig()
86
-
87
-
88
-def torch_prepare(
89
- a: torch.Tensor,
90
- topk_ids: torch.Tensor,
91
- num_experts: int,
92
- max_num_tokens: int | None = None,
93
-) -> tuple[torch.Tensor, torch.Tensor]:
94
- assert topk_ids.dim() == 2
95
- assert topk_ids.shape[0] == a.shape[0]
96
-
97
- num_tokens, hidden_dim = a.shape
98
- topk = topk_ids.shape[1]
99
-
100
- tokens_per_expert = torch.bincount(topk_ids.view(-1), minlength=num_experts)
101
-
102
- assert tokens_per_expert.numel() == num_experts
103
-
104
- if max_num_tokens is None:
105
- max_num_tokens = int(tokens_per_expert.max().item())
106
-
107
- b_a = torch.zeros(
108
- (num_experts, max_num_tokens, hidden_dim), dtype=a.dtype, device=a.device
109
- )
110
-
111
- token_counts = torch.zeros(num_experts, dtype=torch.int, device=a.device)
112
-
113
- for token in range(num_tokens):
114
- for j in range(topk):
115
- expert_id = topk_ids[token, j]
116
- idx = token_counts[expert_id]
117
- b_a[expert_id, idx : idx + 1, :] = a[token, :]
118
- token_counts[expert_id] = token_counts[expert_id] + 1
119
-
120
- return b_a, tokens_per_expert
121
-
122
-
123
-def torch_finalize(
124
- b_out: torch.Tensor, topk_weight: torch.Tensor, topk_ids: torch.Tensor
125
-) -> torch.Tensor:
126
- num_tokens = topk_ids.shape[0]
127
- num_experts = b_out.shape[0]
128
- K = b_out.shape[-1]
129
- out = torch.zeros((num_tokens, K), dtype=b_out.dtype, device=b_out.device)
130
- expert_counts = torch.zeros(num_experts, dtype=torch.int, device=b_out.device)
131
- for token in range(num_tokens):
132
- expert_ids = topk_ids[token]
133
- for i in range(expert_ids.numel()):
134
- expert_id = expert_ids[i]
135
- idx = expert_counts[expert_id]
136
- out[token, :] = (
137
- out[token, :]
138
- + b_out[expert_id, idx : idx + 1, :] * topk_weight[token, i]
139
- )
140
- expert_counts[expert_id] = expert_counts[expert_id] + 1
141
-
142
- return out
143
-
144
-
145
-def torch_batched_moe(
146
- a: torch.Tensor,
147
- w1: torch.Tensor,
148
- w2: torch.Tensor,
149
- topk_weight: torch.Tensor,
150
- topk_ids: torch.Tensor,
151
-) -> torch.Tensor:
152
- num_experts = w1.shape[0]
153
- b_a, tokens_per_expert = torch_prepare(a, topk_ids, num_experts)
154
- assert b_a.dim() == 3
155
- num_tokens, topk = topk_ids.shape
156
- _, max_num_tokens, K = b_a.shape
157
- assert num_experts == b_a.shape[0] and w2.shape[1] == K
158
- out = torch.zeros(
159
- (num_experts, max_num_tokens, K), dtype=b_a.dtype, device=b_a.device
160
- )
161
- tmp = torch.empty(
162
- (max_num_tokens, w1.shape[1] // 2), dtype=b_a.dtype, device=b_a.device
163
- )
164
- for expert in range(num_experts):
165
- num = tokens_per_expert[expert]
166
- if num > 0:
167
- torch.ops._C.silu_and_mul(
168
- tmp[:num], b_a[expert, :num, :] @ w1[expert].transpose(0, 1)
169
- )
170
- out[expert, :num, :] = tmp[:num] @ w2[expert].transpose(0, 1)
171
-
172
- return torch_finalize(out, topk_weight, topk_ids)
173
-
174
-
175
-@pytest.mark.parametrize("m,n,k", BATCHED_MOE_MNK_FACTORS)
176
-@pytest.mark.parametrize("e", NUM_EXPERTS)
177
-@pytest.mark.parametrize("topk", TOP_KS)
178
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
179
-def test_fused_moe_batched_experts(
180
- m: int,
181
- n: int,
182
- k: int,
183
- e: int,
184
- topk: int,
185
- dtype: torch.dtype,
186
- workspace_init,
187
-):
188
- set_random_seed(7)
189
-
190
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
191
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
192
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
193
- score = torch.randn((m, e), device="cuda", dtype=dtype)
194
-
195
- with set_current_vllm_config(vllm_config):
196
- topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
197
- baseline_output = torch_experts(
198
- a, w1, w2, topk_weight, topk_ids
199
- ) # only for baseline
200
- torch_output = torch_batched_moe(a, w1, w2, topk_weight, topk_ids)
201
- batched_output = naive_batched_moe(
202
- a, w1, w2, topk_weight, topk_ids
203
- ) # pick torch_experts or this
204
-
205
- torch.testing.assert_close(baseline_output, torch_output, atol=2e-2, rtol=0)
206
- torch.testing.assert_close(baseline_output, batched_output, atol=2e-2, rtol=0)
207
-
208
-
209
-def create_pplx_prepare_finalize(
210
- num_tokens: int,
211
- hidden_dim: int,
212
- topk: int,
213
- num_experts: int,
214
- rank: int,
215
- dp_size: int,
216
- world_size: int,
217
- in_dtype: torch.dtype,
218
- quant_dtype: torch.dtype | None,
219
- block_shape: list[int] | None,
220
- per_act_token_quant: bool,
221
- group_name: str | None,
222
-):
223
- from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
224
- PplxPrepareAndFinalize,
225
- pplx_hidden_dim_scale_bytes,
226
- )
227
-
228
- max_num_tokens = max(rank_chunk(num_tokens, 0, world_size), 1)
229
- num_local_experts = rank_chunk(num_experts, 0, world_size)
230
-
231
- hidden_dim_bytes, scale_bytes = pplx_hidden_dim_scale_bytes(
232
- max_num_tokens,
233
- hidden_dim,
234
- in_dtype,
235
- quant_dtype,
236
- per_act_token_quant=per_act_token_quant,
237
- block_shape=block_shape,
238
- )
239
-
240
- args = dict(
241
- max_num_tokens=max_num_tokens,
242
- num_experts=num_experts,
243
- experts_per_token=topk,
244
- rank=rank,
245
- world_size=world_size,
246
- dp_size=dp_size,
247
- hidden_dim=hidden_dim,
248
- hidden_dim_bytes=hidden_dim_bytes,
249
- hidden_dim_scale_bytes=scale_bytes,
250
- )
251
-
252
- if group_name is None:
253
- ata = AllToAll.internode(**args)
254
- else:
255
- args["group_name"] = group_name
256
- ata = AllToAll.intranode(**args)
257
-
258
- prepare_finalize = PplxPrepareAndFinalize(
259
- ata,
260
- max_num_tokens=max_num_tokens,
261
- num_local_experts=num_local_experts,
262
- num_dispatchers=world_size // dp_size,
263
- )
264
-
265
- return prepare_finalize, ata
266
-
267
-
268
-def rank_chunk(num: int, r: int, w: int) -> int:
269
- rem = num % w
270
- return (num // w) + (1 if r < rem else 0)
271
-
272
-
273
-def chunk_by_rank(t: torch.Tensor, r: int, w: int) -> torch.Tensor:
274
- chunk = rank_chunk(t.shape[0], r, w)
275
- return t[(r * chunk) : (r + 1) * chunk]
276
-
277
-
278
-def maybe_chunk_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None:
279
- if t is not None:
280
- return chunk_by_rank(t, r, w)
281
- else:
282
- return t
283
-
284
-
285
-def chunk_scales_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None:
286
- if t is not None and t.numel() > 1:
287
- chunk = rank_chunk(t.shape[0], r, w)
288
- return t[(r * chunk) : (r + 1) * chunk]
289
- else:
290
- return t
291
-
292
-
293
-def chunk_scales(t: torch.Tensor | None, start: int, end: int) -> torch.Tensor | None:
294
- if t is not None and t.numel() > 1:
295
- return t[start:end]
296
- else:
297
- return t
298
-
299
-
300
-def dummy_work(a: torch.Tensor) -> torch.Tensor:
301
- return a * 1.1
302
-
303
-
304
-def pplx_prepare_finalize(
305
- pgi: ProcessGroupInfo,
306
- dp_size: int,
307
- a: torch.Tensor,
308
- topk_weight: torch.Tensor,
309
- topk_ids: torch.Tensor,
310
- num_experts: int,
311
- quant_dtype: torch.dtype | None,
312
- block_shape: list[int] | None,
313
- per_act_token_quant: bool,
314
- group_name: str | None,
315
-) -> torch.Tensor:
316
- assert torch.cuda.current_device() == pgi.local_rank
317
-
318
- topk = topk_ids.shape[1]
319
- num_tokens, hidden_dim = a.shape
320
- device = pgi.device
321
- rank = pgi.rank
322
- world_size = pgi.world_size
323
-
324
- topk_ids = topk_ids.to(dtype=torch.uint32)
325
-
326
- prepare_finalize, ata = create_pplx_prepare_finalize(
327
- num_tokens,
328
- hidden_dim,
329
- topk,
330
- num_experts,
331
- rank,
332
- dp_size,
333
- world_size,
334
- a.dtype,
335
- quant_dtype,
336
- block_shape,
337
- per_act_token_quant,
338
- group_name,
339
- )
340
-
341
- assert a.shape[0] == topk_ids.shape[0]
342
-
343
- a_chunk = chunk_by_rank(a, rank, world_size).to(device)
344
- chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device)
345
- chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device)
346
-
347
- assert a_chunk.shape[0] == chunk_topk_ids.shape[0]
348
-
349
- out = torch.full(
350
- a_chunk.shape,
351
- torch.nan,
352
- dtype=a.dtype,
353
- device=device,
354
- )
355
-
356
- if quant_dtype is not None and not per_act_token_quant and block_shape is None:
357
- a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
358
- a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
359
- else:
360
- a1_scale = None
361
- a2_scale = None
362
-
363
- b_a, b_a_scale, expert_num_tokens, _, _ = prepare_finalize.prepare(
364
- a_chunk,
365
- chunk_topk_weight,
366
- chunk_topk_ids,
367
- num_experts,
368
- None,
369
- False,
370
- FusedMoEQuantConfig.make(
371
- quant_dtype,
372
- per_act_token_quant=per_act_token_quant,
373
- per_out_ch_quant=False,
374
- block_shape=block_shape,
375
- a1_scale=a1_scale,
376
- a2_scale=a2_scale,
377
- ),
378
- )
379
-
380
- b_a = dummy_work(dequant(b_a, b_a_scale, block_shape, per_act_token_quant, a.dtype))
381
-
382
- prepare_finalize.finalize(
383
- out,
384
- b_a,
385
- chunk_topk_weight,
386
- chunk_topk_ids,
387
- False,
388
- weight_and_reduce_impl=TopKWeightAndReduceDelegate(),
389
- )
390
-
391
- torch.cuda.synchronize()
392
-
393
- ata.destroy()
394
-
395
- num_tokens = a_chunk.shape[0]
396
-
397
- return out[:num_tokens]
398
-
399
-
400
-def _pplx_prepare_finalize(
401
- pgi: ProcessGroupInfo,
402
- dp_size: int,
403
- a: torch.Tensor,
404
- score: torch.Tensor,
405
- topk: torch.Tensor,
406
- num_experts: int,
407
- quant_dtype: torch.dtype | None,
408
- block_shape: list[int] | None,
409
- per_act_token_quant: bool,
410
- use_internode: bool,
411
-):
412
- try:
413
- if use_internode:
414
- uid = (
415
- nvshmem_get_unique_id()
416
- if pgi.rank == 0
417
- else nvshmem_alloc_empty_unique_id()
418
- )
419
- torch.distributed.broadcast(uid, src=0)
420
- nvshmem_init(uid, pgi.rank, pgi.world_size)
421
- group_name = None
422
- else:
423
- group_ranks = list(range(pgi.world_size))
424
- cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
425
- group_name = cpu_group.group_name
426
-
427
- topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
428
- m, k = a.shape
429
-
430
- a_rep = torch.repeat_interleave(dummy_work(a), topk, dim=0)
431
-
432
- torch_output = (
433
- a_rep.view(m, topk, k) * topk_weight.view(m, topk, 1).to(a_rep.dtype)
434
- ).sum(dim=1)
435
-
436
- pplx_output = pplx_prepare_finalize(
437
- pgi,
438
- dp_size,
439
- a,
440
- topk_weight,
441
- topk_ids,
442
- num_experts,
443
- quant_dtype,
444
- block_shape,
445
- per_act_token_quant,
446
- group_name,
447
- )
448
-
449
- torch_output = chunk_by_rank(torch_output, pgi.rank, pgi.world_size).to(
450
- pgi.device
451
- )
452
-
453
- torch.testing.assert_close(pplx_output, torch_output, atol=3e-2, rtol=3e-2)
454
- finally:
455
- if use_internode:
456
- nvshmem_finalize()
457
-
458
-
459
-@pytest.mark.parametrize("mnk", PPLX_COMBOS)
460
-@pytest.mark.parametrize("e", NUM_EXPERTS)
461
-@pytest.mark.parametrize("topk", TOP_KS)
462
-@pytest.mark.parametrize("dtype", DTYPES)
463
-@pytest.mark.parametrize("world_dp_size", [[2, 1]])
464
-@pytest.mark.parametrize("per_act_token_quant", [False, True])
465
-@pytest.mark.parametrize("block_shape", [None, [128, 128]])
466
-@pytest.mark.parametrize("use_internode", [False])
467
-@pytest.mark.optional
468
-@requires_pplx
469
-@multi_gpu_test(num_gpus=2)
470
-def test_pplx_prepare_finalize_slow(
471
- mnk: tuple[int, int, int],
472
- e: int,
473
- topk: int,
474
- dtype: torch.dtype,
475
- world_dp_size: tuple[int, int],
476
- per_act_token_quant: bool,
477
- block_shape: list[int] | None,
478
- use_internode: bool,
479
-):
480
- if dtype == torch.float8_e4m3fn:
481
- use_fp8_w8a8 = True
482
- act_dtype = torch.bfloat16
483
- quant_dtype = dtype
484
- else:
485
- use_fp8_w8a8 = False
486
- act_dtype = dtype
487
- quant_dtype = None
488
-
489
- if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
490
- pytest.skip("Skip quantization test for non-quantized type")
491
-
492
- if per_act_token_quant and block_shape is not None:
493
- pytest.skip("Skip illegal quantization combination")
494
-
495
- set_random_seed(7)
496
- m, n, k = mnk
497
- world_size, dp_size = world_dp_size
498
- device = "cuda"
499
-
500
- a = torch.randn((m, k), device=device, dtype=act_dtype) / 10
501
- score = torch.randn((m, e), device=device, dtype=act_dtype)
502
-
503
- parallel_launch(
504
- world_size,
505
- _pplx_prepare_finalize,
506
- dp_size,
507
- a,
508
- score,
509
- topk,
510
- e,
511
- quant_dtype,
512
- block_shape,
513
- per_act_token_quant,
514
- use_internode,
515
- )
516
-
517
-
518
-def pplx_moe(
519
- group_name: str | None,
520
- rank: int,
521
- world_size: int,
522
- dp_size: int,
523
- a: torch.Tensor,
524
- w1: torch.Tensor,
525
- w2: torch.Tensor,
526
- topk_weight: torch.Tensor,
527
- topk_ids: torch.Tensor,
528
- w1_scale: torch.Tensor | None = None,
529
- w2_scale: torch.Tensor | None = None,
530
- a1_scale: torch.Tensor | None = None,
531
- a2_scale: torch.Tensor | None = None,
532
- quant_dtype: torch.dtype | None = None,
533
- per_act_token_quant=False,
534
- block_shape: list[int] | None = None,
535
- use_compile: bool = False,
536
- use_cudagraphs: bool = True,
537
- shared_experts: torch.nn.Module | None = None,
538
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
539
- num_tokens, hidden_dim = a.shape
540
- num_experts = w1.shape[0]
541
- topk = topk_ids.shape[1]
542
- max_num_tokens = round_up(rank_chunk(a.shape[0], 0, world_size), 16)
543
-
544
- prepare_finalize, ata = create_pplx_prepare_finalize(
545
- num_tokens,
546
- hidden_dim,
547
- topk,
548
- num_experts,
549
- rank,
550
- dp_size,
551
- world_size,
552
- a.dtype,
553
- quant_dtype,
554
- block_shape,
555
- per_act_token_quant,
556
- group_name,
557
- )
558
-
559
- topk_ids = topk_ids.to(dtype=torch.uint32)
560
-
561
- # Note: workers with the same dp_rank must use the exact same inputs.
562
- a_chunk = chunk_by_rank(a, rank, world_size)
563
- chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size)
564
- chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size)
565
-
566
- # Chunking weights like this only works for batched format
567
- w1_chunk = chunk_by_rank(w1, rank, world_size)
568
- w2_chunk = chunk_by_rank(w2, rank, world_size)
569
- w1_scale_chunk = maybe_chunk_by_rank(w1_scale, rank, world_size)
570
- w2_scale_chunk = maybe_chunk_by_rank(w2_scale, rank, world_size)
571
- a1_scale_chunk = chunk_scales_by_rank(a1_scale, rank, world_size)
572
- a2_scale_chunk = chunk_scales_by_rank(a2_scale, rank, world_size)
573
-
574
- quant_config = FusedMoEQuantConfig.make(
575
- quant_dtype,
576
- block_shape=block_shape,
577
- per_act_token_quant=per_act_token_quant,
578
- w1_scale=w1_scale_chunk,
579
- w2_scale=w2_scale_chunk,
580
- a1_scale=a1_scale_chunk,
581
- a2_scale=a2_scale_chunk,
582
- )
583
-
584
- experts = BatchedTritonExperts(
585
- max_num_tokens=max_num_tokens,
586
- num_dispatchers=prepare_finalize.num_dispatchers(),
587
- quant_config=quant_config,
588
- moe_config=make_dummy_moe_config(),
589
- )
590
-
591
- fused_experts = FusedMoEModularKernel(
592
- prepare_finalize,
593
- experts,
594
- shared_experts,
595
- inplace=False,
596
- )
597
-
598
- # Note: for now use_compile will error out if the problem size is
599
- # large enough to trigger chunking. I'm leaving the flag and
600
- # setup code in case we are able to revisit this later.
601
- if use_compile:
602
- _fused_experts = torch.compile(
603
- fused_experts, backend="inductor", fullgraph=True
604
- )
605
- torch._dynamo.mark_dynamic(a_chunk, 0)
606
- torch._dynamo.mark_dynamic(chunk_topk_weight, 0)
607
- torch._dynamo.mark_dynamic(chunk_topk_ids, 0)
608
- else:
609
- _fused_experts = fused_experts
610
-
611
- out = _fused_experts(
612
- a_chunk,
613
- w1_chunk,
614
- w2_chunk,
615
- chunk_topk_weight,
616
- chunk_topk_ids,
617
- global_num_experts=num_experts,
618
- )
619
-
620
- if use_cudagraphs:
621
- if isinstance(out, tuple):
622
- out[0].fill_(0)
623
- out[1].fill_(0)
624
- else:
625
- out.fill_(0)
626
- stream = torch.cuda.Stream()
627
- graph = torch.cuda.CUDAGraph()
628
- with torch.cuda.graph(graph, stream=stream):
629
- out = _fused_experts(
630
- a_chunk,
631
- w1_chunk,
632
- w2_chunk,
633
- chunk_topk_weight,
634
- chunk_topk_ids,
635
- global_num_experts=num_experts,
636
- )
637
-
638
- torch.cuda.synchronize()
639
- graph.replay()
640
-
641
- torch.cuda.synchronize()
642
-
643
- ata.destroy()
644
-
645
- return out
646
-
647
-
648
-def _pplx_moe(
649
- pgi: ProcessGroupInfo,
650
- dp_size: int,
651
- a: torch.Tensor,
652
- w1: torch.Tensor,
653
- w2: torch.Tensor,
654
- score: torch.Tensor,
655
- topk: int,
656
- num_experts: int,
657
- w1_s: torch.Tensor | None = None,
658
- w2_s: torch.Tensor | None = None,
659
- quant_dtype: torch.dtype | None = None,
660
- per_act_token_quant: bool = False,
661
- block_shape: list[int] | None = None,
662
- use_internode: bool = False,
663
- shared_experts: torch.nn.Module | None = None,
664
-):
665
- try:
666
- if use_internode:
667
- uid = (
668
- nvshmem_get_unique_id()
669
- if pgi.rank == 0
670
- else nvshmem_alloc_empty_unique_id()
671
- )
672
- torch.distributed.broadcast(uid, src=0)
673
- nvshmem_init(uid, pgi.rank, pgi.world_size)
674
- group_name = None
675
- else:
676
- group_ranks = list(range(pgi.world_size))
677
- cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
678
- group_name = cpu_group.group_name
679
-
680
- m, k = a.shape
681
- e, _, n = w2.shape
682
-
683
- moe_config = get_default_config(m, e, n, k, topk, a.dtype, False)
684
-
685
- device = torch.device("cuda", pgi.rank)
686
- rank = pgi.rank
687
- world_size = pgi.world_size
688
-
689
- a = a.to(device)
690
- w1 = w1.to(device)
691
- w2 = w2.to(device)
692
- w1_s = w1_s.to(device) if w1_s is not None else None
693
- w2_s = w2_s.to(device) if w2_s is not None else None
694
-
695
- if quant_dtype is not None and not per_act_token_quant and block_shape is None:
696
- a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
697
- a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
698
- else:
699
- a1_scale = None
700
- a2_scale = None
701
-
702
- with set_current_vllm_config(vllm_config), override_config(moe_config):
703
- topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
704
-
705
- shared_output = shared_experts(a) if shared_experts is not None else None
706
-
707
- torch_output = torch_experts(
708
- a,
709
- w1,
710
- w2,
711
- topk_weight,
712
- topk_ids,
713
- w1_scale=w1_s,
714
- w2_scale=w2_s,
715
- a1_scale=a1_scale,
716
- a2_scale=a2_scale,
717
- quant_dtype=quant_dtype,
718
- per_act_token_quant=per_act_token_quant,
719
- block_shape=block_shape,
720
- )
721
-
722
- batched_output = naive_batched_moe(
723
- a,
724
- w1,
725
- w2,
726
- topk_weight,
727
- topk_ids,
728
- w1_scale=w1_s,
729
- w2_scale=w2_s,
730
- a1_scale=a1_scale,
731
- a2_scale=a2_scale,
732
- quant_dtype=quant_dtype,
733
- per_act_token_quant=per_act_token_quant,
734
- block_shape=block_shape,
735
- )
736
-
737
- pplx_outputs = pplx_moe(
738
- group_name,
739
- rank,
740
- world_size,
741
- dp_size,
742
- a,
743
- w1,
744
- w2,
745
- topk_weight,
746
- topk_ids,
747
- w1_scale=w1_s,
748
- w2_scale=w2_s,
749
- a1_scale=a1_scale,
750
- a2_scale=a2_scale,
751
- quant_dtype=quant_dtype,
752
- per_act_token_quant=per_act_token_quant,
753
- block_shape=block_shape,
754
- shared_experts=shared_experts,
755
- )
756
-
757
- if shared_experts is None:
758
- pplx_shared_output = None
759
- pplx_output = pplx_outputs
760
- assert isinstance(pplx_output, torch.Tensor)
761
- else:
762
- pplx_shared_output, pplx_output = pplx_outputs
763
-
764
- if shared_output is not None:
765
- assert pplx_shared_output is not None
766
- chunked_shared_output = chunk_by_rank(
767
- shared_output, pgi.rank, pgi.world_size
768
- ).to(pplx_shared_output.device)
769
- else:
770
- chunked_shared_output = None
771
-
772
- chunked_batch_output = chunk_by_rank(
773
- batched_output, pgi.rank, pgi.world_size
774
- ).to(pplx_output.device)
775
-
776
- torch.testing.assert_close(batched_output, torch_output, atol=3e-2, rtol=3e-2)
777
-
778
- torch.testing.assert_close(
779
- pplx_output, chunked_batch_output, atol=3e-2, rtol=3e-2
780
- )
781
-
782
- if shared_experts is not None:
783
- assert chunked_shared_output is not None
784
- assert pplx_shared_output is not None
785
- torch.testing.assert_close(
786
- pplx_shared_output, chunked_shared_output, atol=3e-2, rtol=3e-2
787
- )
788
-
789
- finally:
790
- if use_internode:
791
- nvshmem_finalize()
792
-
793
-
794
-@pytest.mark.parametrize("mnk", PPLX_COMBOS)
795
-@pytest.mark.parametrize("e", NUM_EXPERTS)
796
-@pytest.mark.parametrize("topk", TOP_KS)
797
-@pytest.mark.parametrize("dtype", DTYPES)
798
-@pytest.mark.parametrize("world_dp_size", [[2, 1]])
799
-@pytest.mark.parametrize("per_act_token_quant", [False, True])
800
-@pytest.mark.parametrize("block_shape", [None, [128, 128]])
801
-@pytest.mark.parametrize("use_internode", [False])
802
-@pytest.mark.optional
803
-@requires_pplx
804
-@multi_gpu_test(num_gpus=2)
805
-def test_pplx_moe_slow(
806
- mnk: tuple[int, int, int],
807
- e: int,
808
- topk: int,
809
- dtype: torch.dtype,
810
- world_dp_size: tuple[int, int],
811
- per_act_token_quant: bool,
812
- block_shape: list[int] | None,
813
- use_internode: bool,
814
-):
815
- set_random_seed(7)
816
- m, n, k = mnk
817
- world_size, dp_size = world_dp_size
818
-
819
- if dtype == torch.float8_e4m3fn:
820
- use_fp8_w8a8 = True
821
- quant_dtype = dtype
822
- else:
823
- use_fp8_w8a8 = False
824
- quant_dtype = None
825
-
826
- if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
827
- pytest.skip("Skip quantization test for non-quantized type")
828
-
829
- if per_act_token_quant and block_shape is not None:
830
- pytest.skip("Skip illegal quantization combination")
831
-
832
- a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
833
- score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
834
-
835
- (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
836
- e,
837
- n,
838
- k,
839
- quant_dtype=quant_dtype,
840
- block_shape=block_shape,
841
- per_out_ch_quant=per_act_token_quant,
842
- )
843
-
844
- parallel_launch(
845
- world_size,
846
- _pplx_moe,
847
- dp_size,
848
- a,
849
- w1,
850
- w2,
851
- score,
852
- topk,
853
- e,
854
- w1_s,
855
- w2_s,
856
- quant_dtype,
857
- per_act_token_quant,
858
- block_shape,
859
- use_internode,
860
- )
861
-
862
-
863
-def _pplx_test_loop(
864
- pgi: ProcessGroupInfo,
865
- dp_size: int,
866
- use_internode: bool,
867
- use_shared_experts: bool,
868
- make_weights: bool,
869
- test_fn: Callable,
870
-):
871
- device = torch.device(f"cuda:{pgi.local_rank}")
872
- init_workspace_manager(device)
873
-
874
- def format_result(msg, ex=None):
875
- if ex is not None:
876
- x = str(ex)
877
- newx = x.strip(" \n\t")[:16]
878
- if len(newx) < len(x):
879
- newx = newx + " ..."
880
-
881
- prefix = "E\t"
882
- print(f"{textwrap.indent(traceback.format_exc(), prefix)}")
883
- print(f"FAILED {msg} - {newx}\n")
884
- else:
885
- print(f"PASSED {msg}")
886
-
887
- if use_shared_experts:
888
- # Note: this config is only needed for the non-naive shared experts.
889
- new_vllm_config = copy.deepcopy(vllm_config)
890
- new_vllm_config.parallel_config.data_parallel_size = pgi.world_size
891
- new_vllm_config.parallel_config.enable_expert_parallel = True
892
- _set_vllm_config(new_vllm_config, pgi.world_size, pgi.rank, pgi.local_rank)
893
-
894
- set_random_seed(7)
895
- combos = itertools.product(
896
- PPLX_COMBOS, NUM_EXPERTS, TOP_KS, DTYPES, [False, True], [None, [128, 128]]
897
- )
898
- exceptions = []
899
- count = 0
900
- for mnk, e, topk, dtype, per_act_token_quant, block_shape in combos:
901
- count = count + 1
902
- m, n, k = mnk
903
-
904
- if dtype == torch.float8_e4m3fn:
905
- use_fp8_w8a8 = True
906
- quant_dtype = dtype
907
- else:
908
- use_fp8_w8a8 = False
909
- quant_dtype = None
910
-
911
- test_desc = (
912
- f"test_pplx_moe[mnk={mnk}, e={e}, topk={topk}, "
913
- f"dtype={dtype}, per_act_token={per_act_token_quant}, "
914
- f"block_shape={block_shape}, use_internode={use_internode}, "
915
- f"use_shared_experts={use_shared_experts}"
916
- )
917
-
918
- if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
919
- print(f"{test_desc} - Skip quantization test for non-quantized type.")
920
- continue
921
-
922
- if per_act_token_quant and block_shape is not None:
923
- print(f"{test_desc} - Skip illegal quantization combination.")
924
- continue
925
-
926
- a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
927
- score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
928
-
929
- args = dict()
930
- if make_weights:
931
- (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
932
- e,
933
- n,
934
- k,
935
- quant_dtype=quant_dtype,
936
- block_shape=block_shape,
937
- per_out_ch_quant=per_act_token_quant,
938
- )
939
- args["w1"] = w1
940
- args["w2"] = w2
941
- args["w1_s"] = w1_s
942
- args["w2_s"] = w2_s
943
-
944
- if use_shared_experts:
945
- args["shared_experts"] = make_shared_experts(
946
- n,
947
- k,
948
- in_dtype=a.dtype,
949
- quant_dtype=quant_dtype,
950
- )
951
-
952
- try:
953
- test_fn(
954
- pgi=pgi,
955
- dp_size=dp_size,
956
- a=a,
957
- score=score,
958
- topk=topk,
959
- num_experts=e,
960
- quant_dtype=quant_dtype,
961
- per_act_token_quant=per_act_token_quant,
962
- block_shape=block_shape,
963
- use_internode=use_internode,
964
- **args,
965
- )
966
- format_result(test_desc)
967
- except Exception as ex:
968
- format_result(test_desc, ex)
969
- exceptions.append(ex)
970
-
971
- if len(exceptions) > 0:
972
- raise RuntimeError(
973
- f"{len(exceptions)} of {count} tests failed in child process, "
974
- f"rank={pgi.rank}."
975
- )
976
- else:
977
- print(f"{count} of {count} tests passed in child process, rank={pgi.rank}.")
978
-
979
-
980
-@pytest.mark.parametrize("world_dp_size", [[2, 1]])
981
-@pytest.mark.parametrize("use_internode", [False])
982
-@requires_pplx
983
-@multi_gpu_test(num_gpus=2)
984
-def test_pplx_prepare_finalize(
985
- world_dp_size: tuple[int, int],
986
- use_internode: bool,
987
-):
988
- set_random_seed(7)
989
- world_size, dp_size = world_dp_size
990
- parallel_launch(
991
- world_size * dp_size,
992
- _pplx_test_loop,
993
- dp_size,
994
- use_internode,
995
- False,
996
- False,
997
- _pplx_prepare_finalize,
998
- )
999
-
1000
-
1001
-@pytest.mark.parametrize("world_dp_size", [[2, 1]])
1002
-@pytest.mark.parametrize("use_internode", [False])
1003
-@pytest.mark.parametrize("use_shared_experts", [False, True])
1004
-@requires_pplx
1005
-@multi_gpu_test(num_gpus=2)
1006
-def test_pplx_moe(
1007
- world_dp_size: tuple[int, int],
1008
- use_internode: bool,
1009
- use_shared_experts: bool,
1010
-):
1011
- set_random_seed(7)
1012
- world_size, dp_size = world_dp_size
1013
- parallel_launch(
1014
- world_size,
1015
- _pplx_test_loop,
1016
- dp_size,
1017
- use_internode,
1018
- use_shared_experts,
1019
- True,
1020
- _pplx_moe,
1021
- )
tools/ep_kernels/README.md
+1
-1
@@ -4,7 +4,7 @@ Large-scale cluster-level expert parallel, as described in the [DeepSeek-V3 Tech
4
5
Here we break down the requirements in 2 steps:
6
7
-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.
7
+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.
8
2. Configure NVIDIA driver to enable IBGDA. This step requires root access, and must be done on the host machine.
9
10
Step 2 is necessary for multi-node deployment.
tools/ep_kernels/elastic_ep/install_eep_libraries.sh
-7
@@ -76,11 +76,4 @@ popd
76
77
export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem_install:$CMAKE_PREFIX_PATH
78
79
-# build and install pplx, require pytorch installed
80
-pushd "$WORKSPACE"
81
-git clone https://github.com/ppl-ai/pplx-kernels
82
-cd pplx-kernels
83
-# see https://github.com/pypa/pip/issues/9955#issuecomment-838065925
84
-# PIP_NO_BUILD_ISOLATION=0 disables build isolation
85
-PIP_NO_BUILD_ISOLATION=0 TORCH_CUDA_ARCH_LIST=9.0a+PTX pip install . --no-deps -v
79
tools/ep_kernels/install_python_libraries.sh
-18
@@ -4,12 +4,10 @@ set -ex
4
# usage: ./install_python_libraries.sh [options]
5
# --workspace <dir> workspace directory (default: ./ep_kernels_workspace)
6
# --mode <mode> "install" (default) or "wheel"
7
-# --pplx-ref <commit> pplx-kernels commit hash
7
# --deepep-ref <commit> DeepEP commit hash
8
# --nvshmem-ver <ver> NVSHMEM version
9
10
CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
12
-PPLX_COMMIT_HASH=${PPLX_COMMIT_HASH:-"12cecfd"}
11
DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
12
NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"} # Default supports both CUDA 12 and 13
13
WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
@@ -35,14 +33,6 @@ while [[ $# -gt 0 ]]; do
33
MODE="$2"
34
shift 2
35
;;
38
- --pplx-ref)
39
- if [[ -z "$2" || "$2" =~ ^- ]]; then
40
- echo "Error: --pplx-ref requires an argument." >&2
41
- exit 1
42
- fi
43
- PPLX_COMMIT_HASH="$2"
44
- shift 2
45
- ;;
36
--deepep-ref)
37
if [[ -z "$2" || "$2" =~ ^- ]]; then
38
echo "Error: --deepep-ref requires an argument." >&2
@@ -188,14 +178,6 @@ do_build() {
178
popd
179
}
180
191
-# build pplx-kernels
192
-do_build \
193
- "https://github.com/ppl-ai/pplx-kernels" \
194
- "pplx-kernels" \
195
- "setup.py" \
196
- "$PPLX_COMMIT_HASH" \
197
- ""
198
-
181
# build DeepEP
182
do_build \
183
"https://github.com/deepseek-ai/DeepEP" \
vllm/_custom_ops.py
+2
-2
@@ -988,7 +988,7 @@ def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor):
988
return output_tensor
989
990
991
-def get_cutlass_pplx_moe_mm_data(
991
+def get_cutlass_batched_moe_mm_data(
992
expert_offsets: torch.Tensor,
993
problem_sizes1: torch.Tensor,
994
problem_sizes2: torch.Tensor,
@@ -1011,7 +1011,7 @@ def get_cutlass_pplx_moe_mm_data(
1011
multiplication in two grouped MMs used in
1012
the fused MoE operation.
1013
"""
1014
- return torch.ops._C.get_cutlass_pplx_moe_mm_data(
1014
+ return torch.ops._C.get_cutlass_batched_moe_mm_data(
1015
expert_offsets,
1016
problem_sizes1,
1017
problem_sizes2,
vllm/config/compilation.py
+1
-1
@@ -1045,7 +1045,7 @@ class CompilationConfig:
1045
"are optimized for prefill and are incompatible with CUDA Graphs. "
1046
"In order to use CUDA Graphs for decode-optimized workloads, "
1047
"use --all2all-backend with another option, such as "
1048
- "deepep_low_latency, pplx, or allgather_reducescatter."
1048
+ "deepep_low_latency or allgather_reducescatter."
1049
)
1050
self.cudagraph_mode = CUDAGraphMode.NONE
1051
vllm/config/parallel.py
+7
-2
@@ -152,7 +152,6 @@ class ParallelConfig:
152
153
- "naive": Naive all2all implementation using broadcasts\n
154
- "allgather_reducescatter": All2all based on allgather and reducescatter\n
155
- - "pplx": Use pplx kernels\n
155
- "deepep_high_throughput": Use deepep high-throughput kernels\n
156
- "deepep_low_latency": Use deepep low-latency kernels\n
157
- "mori": Use mori kernels\n
@@ -310,6 +309,13 @@ class ParallelConfig:
309
f"but found: {self._api_process_rank}"
310
)
311
312
+ if self.all2all_backend == "pplx":
313
+ logger.warning(
314
+ "The 'pplx' all2all backend has been removed. "
315
+ "Falling back to 'allgather_reducescatter'."
316
+ )
317
+ self.all2all_backend = "allgather_reducescatter"
318
+
319
if self.data_parallel_size_local > self.data_parallel_size:
320
raise ValueError(
321
f"data_parallel_size_local ({self.data_parallel_size_local}) "
@@ -442,7 +448,6 @@ class ParallelConfig:
448
# In this case, ensure the input to the experts is sequence parallel
449
# to avoid the excess work.
450
#
445
- # Not needed for pplx-kernels as it can handle duplicate input tokens.
451
@property
452
def use_sequence_parallel_moe(self) -> bool:
453
return (
vllm/distributed/device_communicators/all2all.py
+1
-92
@@ -3,14 +3,13 @@
3
from typing import Any
4
5
import torch
6
-import torch.distributed as dist
6
7
import vllm.envs as envs
8
from vllm.distributed import get_dp_group, get_ep_group
9
from vllm.forward_context import get_forward_context
10
from vllm.logger import init_logger
11
from vllm.utils.flashinfer import has_flashinfer_all2all
13
-from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx
12
+from vllm.utils.import_utils import has_deep_ep, has_mori
13
14
from .base_device_communicator import All2AllManagerBase, Cache
15
@@ -235,96 +234,6 @@ class AgRsAll2AllManager(All2AllManagerBase):
234
pass
235
236
238
-class PPLXAll2AllManager(All2AllManagerBase):
239
- """
240
- All2All communication based on PPLX kernels.
241
- """
242
-
243
- def __init__(self, cpu_group):
244
- assert has_pplx(), (
245
- "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md"
246
- " to install pplx_kernels."
247
- )
248
- super().__init__(cpu_group)
249
-
250
- if self.internode:
251
- # inter-node communication needs nvshmem,
252
- # intra-node communication uses p2p mapping directly
253
- from pplx_kernels.nvshmem import ( # type: ignore[import-not-found]
254
- nvshmem_alloc_empty_unique_id,
255
- nvshmem_get_unique_id,
256
- nvshmem_init,
257
- )
258
-
259
- logger.debug(
260
- "Initialize NVSHMEM for pplx_kernels: rank=%d, world size=%d",
261
- self.rank,
262
- self.world_size,
263
- )
264
- uid = (
265
- nvshmem_get_unique_id()
266
- if self.rank == 0
267
- else nvshmem_alloc_empty_unique_id()
268
- )
269
- dist.broadcast(
270
- uid,
271
- src=dist.get_process_group_ranks(self.cpu_group)[0],
272
- group=self.cpu_group,
273
- )
274
- logger.debug("PPLX NVSHMEM UID = %s", uid)
275
- nvshmem_init(uid, self.rank, self.world_size)
276
-
277
- self.handle_cache = Cache()
278
-
279
- def get_handle(self, kwargs):
280
- import pplx_kernels as pplx # type: ignore[import-not-found]
281
-
282
- return self.handle_cache.get_or_create(
283
- kwargs,
284
- pplx.AllToAll.internode if self.internode else pplx.AllToAll.intranode,
285
- )
286
-
287
- def dispatch_router_logits(
288
- self,
289
- hidden_states: torch.Tensor,
290
- router_logits: torch.Tensor,
291
- is_sequence_parallel: bool = False,
292
- extra_tensors: list[torch.Tensor] | None = None,
293
- ) -> tuple[torch.Tensor, torch.Tensor]:
294
- raise NotImplementedError
295
-
296
- def dispatch(
297
- self,
298
- hidden_states: torch.Tensor,
299
- topk_weights: torch.Tensor,
300
- topk_ids: torch.Tensor,
301
- is_sequence_parallel: bool = False,
302
- extra_tensors: list[torch.Tensor] | None = None,
303
- ) -> (
304
- tuple[torch.Tensor, torch.Tensor, torch.Tensor]
305
- | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
306
- ):
307
- raise NotImplementedError
308
-
309
- def combine(
310
- self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
311
- ) -> torch.Tensor:
312
- raise NotImplementedError
313
-
314
- def destroy(self):
315
- with self.handle_cache._lock:
316
- for _, handle in self.handle_cache._cache.items():
317
- handle.destroy()
318
-
319
- if self.internode:
320
- from pplx_kernels.nvshmem import (
321
- nvshmem_finalize, # type: ignore[import-not-found]
322
- )
323
-
324
- logger.debug("PPLX NVSHMEM finalize")
325
- nvshmem_finalize()
326
-
327
-
237
class DeepEPAll2AllManagerBase(All2AllManagerBase):
238
"""
239
All2All communication based on DeepEP High-Throughput kernels.
vllm/distributed/device_communicators/cuda_communicator.py
+1
-5
@@ -112,10 +112,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
112
from .all2all import AgRsAll2AllManager
113
114
self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
115
- elif self.all2all_backend == "pplx":
116
- from .all2all import PPLXAll2AllManager
117
-
118
- self.all2all_manager = PPLXAll2AllManager(self.cpu_group)
115
elif self.all2all_backend == "deepep_high_throughput":
116
from .all2all import DeepEPHTAll2AllManager
117
@@ -298,7 +294,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
294
self.fi_ar_comm = None
295
if self.all2all_manager is not None:
296
self.all2all_manager.destroy()
301
- self.all2all_manager = None
297
+ self.all2all_manager = None # type: ignore[assignment]
298
299
def all_gatherv(
300
self,
vllm/distributed/eplb/eplb_state.py
+1
-1
@@ -159,7 +159,7 @@ class EplbModelState:
159
160
NOTE: The expert_load_view now records load for all physical experts
161
rather than just local experts. This ensures consistent load statistics
162
- across different dispatch methods (naive all-to-all, DeepEP, pplx-kernels).
162
+ across different dispatch methods (naive all-to-all, DeepEP).
163
The recorded load will be multiplied by dp_size when using naive all-to-all
164
due to each DP rank contributing the same token set to the calculation.
165
See:
vllm/model_executor/layers/fused_moe/all2all_utils.py
+4
-49
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from typing import Any
5
6
import torch
7
@@ -24,16 +25,11 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
25
MoEPrepareAndFinalizeNoEP,
26
)
27
from vllm.platforms import current_platform
27
-from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx
28
+from vllm.utils.import_utils import has_deep_ep, has_mori
29
30
logger = init_logger(__name__)
31
32
if current_platform.is_cuda_alike():
32
- if has_pplx():
33
- from .pplx_prepare_finalize import (
34
- PplxPrepareAndFinalize,
35
- pplx_hidden_dim_scale_bytes,
36
- )
33
if has_deep_ep():
34
from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
35
from .deepep_ll_prepare_finalize import (
@@ -120,51 +116,10 @@ def maybe_make_prepare_finalize(
116
117
prepare_finalize: FusedMoEPrepareAndFinalize | None = None
118
123
- if moe.use_pplx_kernels:
124
- assert quant_config is not None
125
-
126
- hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
127
- moe.max_num_tokens,
128
- moe.hidden_dim,
129
- moe.in_dtype,
130
- quant_config.quant_dtype,
131
- per_act_token_quant=quant_config.per_act_token_quant,
132
- block_shape=quant_config.block_shape,
133
- )
134
-
135
- all_to_all_args = dict(
136
- max_num_tokens=moe.max_num_tokens,
137
- num_experts=moe.num_experts,
138
- experts_per_token=moe.experts_per_token, # topk
139
- rank=all2all_manager.rank,
140
- world_size=all2all_manager.world_size,
141
- # dp_size actually means tp_size, bug in pplx kernels
142
- dp_size=all2all_manager.tp_group.world_size,
143
- hidden_dim=moe.hidden_dim,
144
- hidden_dim_bytes=hidden_dim_bytes,
145
- hidden_dim_scale_bytes=hidden_scale_bytes,
146
- )
147
-
148
- num_dispatchers = (
149
- all2all_manager.world_size // all2all_manager.tp_group.world_size
150
- )
151
-
152
- # Intranode pplx a2a takes a group name while internode does not.
153
- if not all2all_manager.internode:
154
- all_to_all_args["group_name"] = all2all_manager.cpu_group.group_name
155
-
156
- handle = all2all_manager.get_handle(all_to_all_args)
157
-
158
- prepare_finalize = PplxPrepareAndFinalize(
159
- handle,
160
- max_num_tokens=moe.max_num_tokens,
161
- num_local_experts=moe.num_local_experts,
162
- num_dispatchers=num_dispatchers,
163
- )
164
- elif moe.use_deepep_ht_kernels:
119
+ if moe.use_deepep_ht_kernels:
120
assert moe.dp_size == all2all_manager.dp_world_size
121
167
- all_to_all_args = dict()
122
+ all_to_all_args: dict[str, Any] = dict()
123
handle = all2all_manager.get_handle(all_to_all_args)
124
prepare_finalize = DeepEPHTPrepareAndFinalize(
125
handle,
vllm/model_executor/layers/fused_moe/config.py
+1
-9
@@ -939,10 +939,6 @@ class FusedMoEParallelConfig:
939
def use_all2all_kernels(self):
940
return self.dp_size > 1 and self.use_ep
941
942
- @property
943
- def use_pplx_kernels(self):
944
- return self.use_all2all_kernels and self.all2all_backend == "pplx"
945
-
942
@property
943
def use_deepep_ht_kernels(self):
944
return (
@@ -962,7 +958,7 @@ class FusedMoEParallelConfig:
958
959
@property
960
def use_batched_activation_format(self):
965
- return self.use_deepep_ll_kernels or self.use_pplx_kernels
961
+ return self.use_deepep_ll_kernels
962
963
@property
964
def use_naive_all2all_kernels(self):
@@ -1221,10 +1217,6 @@ class FusedMoEConfig:
1217
def use_ep(self):
1218
return self.moe_parallel_config.use_ep
1219
1224
- @property
1225
- def use_pplx_kernels(self):
1226
- return self.moe_parallel_config.use_pplx_kernels
1227
-
1220
@property
1221
def use_deepep_ht_kernels(self):
1222
return self.moe_parallel_config.use_deepep_ht_kernels
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+1
-1
@@ -166,7 +166,7 @@ def run_cutlass_moe_fp8(
166
problem_sizes1 = torch.empty((local_E, 3), dtype=torch.int32, device=device)
167
problem_sizes2 = torch.empty((local_E, 3), dtype=torch.int32, device=device)
168
169
- ops.get_cutlass_pplx_moe_mm_data(
169
+ ops.get_cutlass_batched_moe_mm_data(
170
expert_offsets,
171
problem_sizes1,
172
problem_sizes2,
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+3
-3
@@ -493,7 +493,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
493
"""
494
A reference prepare/finalize class that reorganizes the tokens into
495
expert batched format, i.e. E x max_num_tokens x K. This is the format
496
- that the PPLX dispatch/combine kernels use.
496
+ that the batched dispatch/combine kernels use.
497
"""
498
499
def __init__(
@@ -648,7 +648,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
648
class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
649
"""
650
A reference MoE expert class that operates on expert batched format,
651
- i.e. E x max_num_tokens x K. This is the format that the pplx
651
+ i.e. E x max_num_tokens x K. This is the format that the batched
652
dispatch/combine kernels use.
653
"""
654
@@ -880,7 +880,7 @@ def batched_moe_kernel_quantize_input(
880
class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
881
"""
882
A Triton based MoE expert class that operates on expert batched format,
883
- i.e. E x max_num_tokens x K. This is the format that the pplx
883
+ i.e. E x max_num_tokens x K. This is the format that the batched
884
dispatch/combine kernels use.
885
"""
886
vllm/model_executor/layers/fused_moe/modular_kernel.py
+3
-3
@@ -1172,9 +1172,9 @@ class FusedMoEModularKernel(torch.nn.Module):
1172
# This happens when none of the tokens from the all2all reach this
1173
# EP rank. Also, note that this is only relevant for CUDAGraph
1174
# incompatible all2all kernels like the DeepEP high-throughput
1175
- # kernels. CUDAGraph compatible all2all kernels like the pplx
1176
- # kernels and the DeepEP low-latency kernels are always batched
1177
- # and can never run into the tensor.numel() == 0 case.
1175
+ # kernels. CUDAGraph compatible all2all kernels like the DeepEP
1176
+ # low-latency kernels are always batched and can never run into
1177
+ # the tensor.numel() == 0 case.
1178
if M_full == 0:
1179
assert num_chunks == 0
1180
workspace13 = None
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+1
-4
@@ -143,10 +143,7 @@ def select_nvfp4_moe_backend(
143
# NOTE(rob): this is kind of a hack. We need to peak into
144
# the prepare-finalize selection to determine if we are using
145
# the batched or standard expert format.
146
- use_batched = (
147
- config.moe_parallel_config.use_deepep_ll_kernels
148
- or config.moe_parallel_config.use_pplx_kernels
149
- )
146
+ use_batched = config.moe_parallel_config.use_deepep_ll_kernels
147
activation_format = (
148
mk.FusedMoEActivationFormat.BatchedExperts
149
if use_batched
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
deleted
-373
@@ -1,373 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Callable
4
-
5
-import pplx_kernels as pplx
6
-import torch
7
-
8
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
-from vllm.logger import init_logger
10
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
12
- TopKWeightAndReduceDelegate,
13
-)
14
-from vllm.model_executor.layers.fused_moe.utils import (
15
- _validate_scale_shape,
16
- moe_kernel_quantize_input,
17
-)
18
-from vllm.utils.math_utils import cdiv, round_up
19
-
20
-logger = init_logger(__name__)
21
-
22
-
23
-def pplx_hidden_dim_scale_bytes(
24
- max_num_tokens: int,
25
- hidden_dim: int,
26
- in_dtype: torch.dtype,
27
- quant_dtype: torch.dtype | str | None,
28
- per_act_token_quant: bool,
29
- block_shape: list[int] | None,
30
-):
31
- # All pplx byte sizes must be 16-byte aligned.
32
- align = 16
33
-
34
- # For blocked per token: set to
35
- # cdiv(hidden_dim, block_size) * sizeof(float32)
36
- # For per-token: set to 4 * sizeof(float32) (x4 for alignment)
37
- if quant_dtype is not None:
38
- assert isinstance(quant_dtype, torch.dtype)
39
- assert quant_dtype.itemsize == 1
40
- hidden_dim_bytes = hidden_dim * quant_dtype.itemsize
41
- elem_size = torch.float32.itemsize
42
-
43
- if per_act_token_quant:
44
- # per-token (M x 1)
45
- assert block_shape is None
46
- hidden_scale_bytes = elem_size
47
- elif block_shape is not None:
48
- # per-group (M x K_tiles)
49
- block_size = block_shape[1]
50
- num_blocks = cdiv(hidden_dim, block_size)
51
- hidden_scale_bytes = num_blocks * elem_size
52
- else:
53
- # per-tensor (1 x 1)
54
- hidden_scale_bytes = elem_size
55
- else:
56
- hidden_dim_bytes = hidden_dim * in_dtype.itemsize
57
- hidden_scale_bytes = 0
58
-
59
- return (
60
- round_up(hidden_dim_bytes, align),
61
- round_up(hidden_scale_bytes, align),
62
- )
63
-
64
-
65
-class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
66
- """PPLX-based prepare and finalize for expert parallelism."""
67
-
68
- def __init__(
69
- self,
70
- a2a: pplx.AllToAll,
71
- max_num_tokens: int,
72
- num_local_experts: int,
73
- num_dispatchers: int,
74
- ):
75
- super().__init__()
76
- assert max_num_tokens > 0
77
- assert num_local_experts > 0
78
- self.a2a = a2a
79
- self.max_num_tokens = max_num_tokens
80
- self.num_local_experts = num_local_experts
81
- self.num_dispatchers_ = num_dispatchers
82
-
83
- @property
84
- def activation_format(self) -> mk.FusedMoEActivationFormat:
85
- return mk.FusedMoEActivationFormat.BatchedExperts
86
-
87
- def max_num_tokens_per_rank(self) -> int | None:
88
- return self.max_num_tokens
89
-
90
- def topk_indices_dtype(self) -> torch.dtype | None:
91
- return torch.uint32
92
-
93
- def num_dispatchers(self) -> int:
94
- return self.num_dispatchers_
95
-
96
- def output_is_reduced(self) -> bool:
97
- return True
98
-
99
- def supports_async(self) -> bool:
100
- return True
101
-
102
- def prepare_async(
103
- self,
104
- a1: torch.Tensor,
105
- topk_weights: torch.Tensor,
106
- topk_ids: torch.Tensor,
107
- num_experts: int,
108
- expert_map: torch.Tensor | None,
109
- apply_router_weight_on_input: bool,
110
- quant_config: FusedMoEQuantConfig,
111
- defer_input_quant: bool = False,
112
- ) -> tuple[Callable, mk.ReceiverType]:
113
- if defer_input_quant:
114
- raise NotImplementedError(
115
- f"{self.__class__.__name__} does not support defer_input_quant=True. "
116
- "Please select an MoE kernel that accepts quantized inputs."
117
- )
118
-
119
- num_tokens = a1.size(0) # M
120
- hidden_dim = a1.size(-1) # K
121
-
122
- assert topk_ids.size(0) == num_tokens
123
- # expert_map should be None because with expert map, -1 id is used for
124
- # non-local token; this causes error when casting ids to the
125
- # topk_indices_dtype() int32
126
- #
127
- if expert_map is not None:
128
- logger.warning_once(
129
- "The PPLX backend does not support expert mapping. "
130
- "The provided `expert_map` will be ignored."
131
- )
132
- expert_map = None # noqa: F841
133
-
134
- # Is this always going to be a1.device?
135
- device = a1.device
136
-
137
- if apply_router_weight_on_input:
138
- topk = topk_ids.size(1)
139
- # TODO: this only works for topK=1, will need to update for topK>1
140
- assert topk == 1, (
141
- "apply_router_weight_on_input is only implemented for topk=1"
142
- )
143
- a1 = a1 * topk_weights.to(a1.dtype)
144
-
145
- repeat_cols = 4
146
- repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0)
147
- # TODO(bnell): always pass quant_config.a1_scale?
148
- a1q, a1q_scale = moe_kernel_quantize_input(
149
- a1,
150
- (None if quant_config.per_act_token_quant else quant_config.a1_scale),
151
- quant_dtype=quant_config.quant_dtype,
152
- per_act_token_quant=quant_config.per_act_token_quant,
153
- block_shape=quant_config.block_shape,
154
- )
155
-
156
- _validate_scale_shape(
157
- a1q, a1q_scale, quant_config.per_act_token_quant, quant_config.block_shape
158
- )
159
-
160
- orig_a_scale_block_shape: int | None = None
161
-
162
- if a1q_scale is not None:
163
- scalar_scales = a1q_scale.numel() == 1
164
-
165
- # pplx requires 2-d scales even for scalar scales
166
- if a1q_scale.dim() <= 1:
167
- assert scalar_scales
168
- a1q_scale = a1q_scale.view(1, 1)
169
-
170
- orig_a_scale_block_shape = a1q_scale.shape[-1]
171
-
172
- if not quant_config.is_block_quantized:
173
- # TODO (bnell): use group_broadcast instead?
174
- a1q_scale = a1q_scale.repeat(repeat_rows, repeat_cols)
175
-
176
- assert a1q_scale is None or a1q_scale.ndim == 2, (
177
- f"{0 if a1q_scale is None else (a1q_scale.ndim, a1q_scale.shape)}"
178
- )
179
-
180
- expert_num_tokens = torch.empty(
181
- self.num_local_experts,
182
- dtype=torch.int32,
183
- device=device,
184
- )
185
-
186
- expert_x = torch.empty(
187
- (
188
- self.num_local_experts,
189
- self.max_num_tokens * self.num_dispatchers(),
190
- hidden_dim,
191
- ),
192
- dtype=a1q.dtype,
193
- device=device,
194
- )
195
-
196
- expert_x_scale: torch.Tensor | None = None
197
- if a1q.dtype.itemsize == 1:
198
- if quant_config.is_per_act_token:
199
- # (M x 1) -> (E x M x K)
200
- final_dim = expert_x.size(2)
201
- elif quant_config.is_per_tensor:
202
- # (1 x 1) -> (E x 1 x 1)
203
- final_dim = 1
204
- else:
205
- # (M x K_tiles) -> (E x M x K_tiles)
206
- assert quant_config.block_shape is not None
207
- num_blocks = cdiv(expert_x.size(2), quant_config.block_shape[1])
208
- final_dim = num_blocks
209
-
210
- expert_x_scale_shape = (
211
- self.num_local_experts,
212
- expert_x.size(1),
213
- round_up(final_dim, 4), # round up for alignment
214
- )
215
-
216
- expert_x_scale = torch.empty(
217
- expert_x_scale_shape,
218
- dtype=torch.float32,
219
- device=expert_x.device,
220
- )
221
-
222
- # This argument is optional, defaults to indices.size(0)
223
- # There's not much point setting this unless it is != indices.size(0)
224
- bound_m: torch.Tensor | None = None
225
-
226
- self.a2a.dispatch(
227
- out_expert_num_tokens=expert_num_tokens,
228
- out_expert_x=expert_x,
229
- out_expert_x_scale=expert_x_scale,
230
- dp_x=a1q,
231
- dp_x_scale=a1q_scale,
232
- indices=topk_ids,
233
- bound_m=bound_m,
234
- do_send=True,
235
- do_recv=False,
236
- )
237
-
238
- hook = lambda: self.a2a.dispatch(
239
- out_expert_num_tokens=expert_num_tokens,
240
- out_expert_x=expert_x,
241
- out_expert_x_scale=expert_x_scale,
242
- dp_x=a1q,
243
- dp_x_scale=a1q_scale,
244
- indices=topk_ids,
245
- bound_m=bound_m,
246
- do_send=False,
247
- do_recv=True,
248
- )
249
-
250
- return (
251
- hook,
252
- lambda: self._receiver(
253
- expert_num_tokens,
254
- expert_x,
255
- expert_x_scale,
256
- orig_a_scale_block_shape,
257
- ),
258
- )
259
-
260
- def _receiver(
261
- self,
262
- expert_num_tokens: torch.Tensor,
263
- expert_x: torch.Tensor,
264
- expert_x_scale: torch.Tensor | None,
265
- orig_a_scale_block_shape: int | None,
266
- ) -> mk.PrepareResultType:
267
- if expert_x_scale is not None:
268
- expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape]
269
- assert expert_x_scale.ndim == 3
270
-
271
- expert_tokens_meta = mk.ExpertTokensMetadata(
272
- expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None
273
- )
274
-
275
- return expert_x, expert_x_scale, expert_tokens_meta, None, None
276
-
277
- def prepare(
278
- self,
279
- a1: torch.Tensor,
280
- topk_weights: torch.Tensor,
281
- topk_ids: torch.Tensor,
282
- num_experts: int,
283
- expert_map: torch.Tensor | None,
284
- apply_router_weight_on_input: bool,
285
- quant_config: FusedMoEQuantConfig,
286
- defer_input_quant: bool = False,
287
- ) -> mk.PrepareResultType:
288
- hook, receiver = self.prepare_async(
289
- a1,
290
- topk_weights,
291
- topk_ids,
292
- num_experts,
293
- expert_map,
294
- apply_router_weight_on_input,
295
- quant_config,
296
- defer_input_quant=defer_input_quant,
297
- )
298
- hook()
299
- return receiver()
300
-
301
- def finalize_async(
302
- self,
303
- output: torch.Tensor,
304
- fused_expert_output: torch.Tensor,
305
- topk_weights: torch.Tensor,
306
- topk_ids: torch.Tensor,
307
- apply_router_weight_on_input: bool,
308
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
309
- ) -> Callable:
310
- assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate), (
311
- "Weight application and reduction happens in the combine kernel."
312
- )
313
-
314
- # This argument is optional
315
- # There's not much point setting this unless it is != topk_ids.size(0)
316
- bound_m: torch.Tensor | None = None
317
-
318
- # TODO (bnell): fails in test_pplx_moe.py, figure out what's going on
319
- # num_tokens = output.size(0) # M
320
- # assert topk_ids.size(0) == num_tokens, (
321
- # f"{topk_ids.size(0)} == {num_tokens}")
322
- assert topk_ids.size() == topk_weights.size(), (
323
- f"{topk_ids.size()} == {topk_weights.size()}"
324
- )
325
- assert output.size(0) <= self.max_num_tokens, (
326
- f"{output.size(0)} <= {self.max_num_tokens}"
327
- )
328
- assert output.size(1) == fused_expert_output.size(-1)
329
-
330
- # Set weights to 1 if we did them in dispatch. This is hacky.
331
- if apply_router_weight_on_input:
332
- topk_weights = torch.ones_like(topk_weights)
333
-
334
- topk_ids_u32 = topk_ids.view(dtype=torch.uint32)
335
-
336
- self.a2a.combine(
337
- out_tokens=output,
338
- indices=topk_ids_u32,
339
- weights=topk_weights,
340
- expert_y=fused_expert_output,
341
- bound_m=bound_m,
342
- do_send=True,
343
- do_recv=False,
344
- )
345
-
346
- return lambda: self.a2a.combine(
347
- out_tokens=output,
348
- indices=topk_ids_u32,
349
- weights=topk_weights,
350
- expert_y=fused_expert_output,
351
- bound_m=bound_m,
352
- do_send=False,
353
- do_recv=True,
354
- )
355
-
356
- def finalize(
357
- self,
358
- output: torch.Tensor,
359
- fused_expert_output: torch.Tensor,
360
- topk_weights: torch.Tensor,
361
- topk_ids: torch.Tensor,
362
- apply_router_weight_on_input: bool,
363
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
364
- ) -> None:
365
- receiver = self.finalize_async(
366
- output,
367
- fused_expert_output,
368
- topk_weights,
369
- topk_ids,
370
- apply_router_weight_on_input,
371
- weight_and_reduce_impl,
372
- )
373
- receiver()
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+1
-2
@@ -216,8 +216,7 @@ class DefaultMoERunner(MoERunner):
216
@property
217
def use_dp_chunking(self) -> bool:
218
return (
219
- self.moe_config.moe_parallel_config.use_pplx_kernels
220
- or self.moe_config.moe_parallel_config.use_deepep_ll_kernels
219
+ self.moe_config.moe_parallel_config.use_deepep_ll_kernels
220
or self.moe_config.moe_parallel_config.use_mori_kernels
221
or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels
222
) and envs.VLLM_ENABLE_MOE_DP_CHUNK
vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py
+5
-4
@@ -14,10 +14,11 @@ class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce):
14
implementation does not perform weight application and reduction
15
but cannot address the needs of all the compatible PrepareAndFinalize
16
implementations.
17
- For example, BatchedTritonExperts is compatible with both
18
- PplxPrepareAndFinalize and BatchedPrepareAndFinalize. PplxPrepareAndFinalize
19
- does the weight-application + reduction as part of the pplx combine kernel.
20
- But the BatchedPrepareAndFinalize needs an implementation. To facilitate
17
+ For example, BatchedTritonExperts is compatible with both batched
18
+ PrepareAndFinalize implementations like DeepEPLLPrepareAndFinalize and
19
+ BatchedPrepareAndFinalize. Some PrepareAndFinalize implementations do
20
+ the weight-application + reduction as part of the combine kernel, while
21
+ BatchedPrepareAndFinalize needs an explicit implementation. To facilitate
22
this case, the BatchedTritonExperts could use TopKWeightAndReduceDelegate
23
so the PrepareAndFinalize implementations could choose how to
24
weight + reduce.
vllm/model_executor/layers/quantization/mxfp4.py
+1
-1
@@ -798,7 +798,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
798
# batched activation format. As self.fused_experts is not
799
# initialized at this point, we resort to checking the MoE config
800
# directly.
801
- is_batched_moe = self.moe.use_pplx_kernels or self.moe.use_deepep_ll_kernels
801
+ is_batched_moe = self.moe.use_deepep_ll_kernels
802
if is_batched_moe:
803
num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
804
else:
vllm/utils/import_utils.py
-5
@@ -402,11 +402,6 @@ def _has_module(module_name: str) -> bool:
402
return importlib.util.find_spec(module_name) is not None
403
404
405
-def has_pplx() -> bool:
406
- """Whether the optional `pplx_kernels` package is available."""
407
- return _has_module("pplx_kernels")
408
-
409
-
405
def has_deep_ep() -> bool:
406
"""Whether the optional `deep_ep` package is available."""
407
return _has_module("deep_ep")