diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index afc8fc49a2a..e1ecfeb8415 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -70,7 +70,7 @@ steps: - tests/kernels/moe/test_batched_deepgemm.py - tests/kernels/attention/test_deepgemm_attention.py commands: - - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm + - pytest -v -s kernels/quantization/test_block_fp8.py - pytest -v -s kernels/moe/test_deepgemm.py - pytest -v -s kernels/moe/test_batched_deepgemm.py - pytest -v -s kernels/attention/test_deepgemm_attention.py @@ -155,5 +155,14 @@ steps: commands: - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py - pytest -v -s kernels/moe/test_deepep_moe.py - - pytest -v -s kernels/moe/test_pplx_cutlass_moe.py - # - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main + +- label: Kernels Fp4 MoE Test (B200) + timeout_in_minutes: 60 + device: b200 + num_devices: 1 + optional: true + commands: + - pytest -v -s kernels/moe/test_cutedsl_moe.py + - pytest -v -s kernels/moe/test_flashinfer_moe.py + - pytest -v -s kernels/moe/test_nvfp4_moe.py + - pytest -v -s kernels/moe/test_ocp_mx_moe.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index adf50a185e5..653d6c42e9a 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,7 +2,7 @@ # for more info about CODEOWNERS file # This lists cover the "core" components of vLLM that require careful review -/vllm/compilation @zou3519 @youkaichao @ProExpertProg +/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni @@ -54,11 +54,14 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett /vllm/v1/kv_cache_interface.py @heheda12345 /vllm/v1/kv_offload @ApostaC @orozery -/vllm/v1/worker/gpu/kv_connector.py @orozery +/vllm/v1/engine @njhill +/vllm/v1/executor @njhill +/vllm/v1/worker @njhill /vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche # Model runner V2 -/vllm/v1/worker/gpu @WoosukKwon +/vllm/v1/worker/gpu @WoosukKwon @njhill +/vllm/v1/worker/gpu/kv_connector.py @orozery # Test ownership /.buildkite/lm-eval-harness @mgoin diff --git a/.github/mergify.yml b/.github/mergify.yml index 080767ca721..9c53342d173 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -259,8 +259,7 @@ pull_request_rules: - files=benchmarks/run_structured_output_benchmark.sh - files=docs/features/structured_outputs.md - files=examples/offline_inference/structured_outputs.py - - files=examples/online_serving/openai_chat_completion_structured_outputs.py - - files=examples/online_serving/openai_chat_completion_structured_outputs_with_reasoning.py + - files=examples/online_serving/structured_outputs/structured_outputs.py - files~=^tests/v1/structured_output/ - files=tests/v1/entrypoints/llm/test_struct_output_generate.py - files~=^vllm/v1/structured_output/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 55127a514f1..479d6db1eeb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -725,7 +725,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # CUTLASS MoE kernels # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and ONLY works - # on Hopper). get_cutlass_(pplx_)moe_mm_data should only be compiled + # on Hopper). get_cutlass_(batched_)moe_mm_data should only be compiled # if it's possible to compile MoE kernels that use its output. cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) @@ -971,7 +971,8 @@ set(VLLM_MOE_EXT_SRC if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu" - "csrc/moe/grouped_topk_kernels.cu") + "csrc/moe/grouped_topk_kernels.cu" + "csrc/moe/router_gemm.cu") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h index b71db356944..d8d962887da 100644 --- a/csrc/moe/moe_ops.h +++ b/csrc/moe/moe_ops.h @@ -58,6 +58,10 @@ void shuffle_rows(const torch::Tensor& input_tensor, torch::Tensor& output_tensor); #ifndef USE_ROCM +// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16) +torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input, + torch::Tensor const& weight); + // DeepSeek V3 optimized router GEMM kernel for SM90+ // Computes output = mat_a @ mat_b.T where: // mat_a: [num_tokens, hidden_dim] in bf16 diff --git a/csrc/moe/router_gemm.cu b/csrc/moe/router_gemm.cu new file mode 100644 index 00000000000..a939f8846ff --- /dev/null +++ b/csrc/moe/router_gemm.cu @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// bf16 x bf16 -> fp32 router GEMM via cuBLAS. +// Uses CUBLAS_COMPUTE_32F so bf16 operands accumulate into fp32, +// matching TRT-LLM's cuBLAS fallback behaviour in dsv3RouterGemmOp. + +#include +#include +#include + +// cuBLAS column-major math for row-major PyTorch tensors: +// weight[N,K]_row lda=K -> cuBLAS sees (K,N) col-major; CUBLAS_OP_T -> +// (N,K) input[M,K]_row ldb=K -> cuBLAS sees (K,M) col-major; CUBLAS_OP_N +// -> (K,M) out[M,N]_row ldc=N -> cuBLAS sees (N,M) col-major (written as +// output^T) +// cuBLAS: C(N,M) = weight(N,K) @ input(K,M) => C^T = output[M,N] +// params: m=N, n=M, k=K, lda=K (weight), ldb=K (input), ldc=N (output) + +torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input, + torch::Tensor const& weight) { + TORCH_CHECK(input.dtype() == torch::kBFloat16, + "router_gemm_bf16_fp32: input must be bfloat16"); + TORCH_CHECK(weight.dtype() == torch::kBFloat16, + "router_gemm_bf16_fp32: weight must be bfloat16"); + TORCH_CHECK(input.dim() == 2 && weight.dim() == 2, + "router_gemm_bf16_fp32: input and weight must be 2-D"); + TORCH_CHECK(input.size(1) == weight.size(1), + "router_gemm_bf16_fp32: inner dimensions must match"); + + int64_t const M = input.size(0); + int64_t const N = weight.size(0); + int64_t const K = input.size(1); + + auto out = torch::empty({M, N}, input.options().dtype(torch::kFloat32)); + + cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(); + TORCH_CUDABLAS_CHECK( + cublasSetStream(handle, at::cuda::getCurrentCUDAStream())); + + float const alpha = 1.0f; + float const beta = 0.0f; + + TORCH_CUDABLAS_CHECK(cublasGemmEx( + handle, CUBLAS_OP_T, CUBLAS_OP_N, static_cast(N), + static_cast(M), static_cast(K), &alpha, weight.data_ptr(), + CUDA_R_16BF, static_cast(K), input.data_ptr(), CUDA_R_16BF, + static_cast(K), &beta, out.data_ptr(), CUDA_R_32F, + static_cast(N), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + + return out; +} diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp index 43859945145..7b627a6f876 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/moe/torch_bindings.cpp @@ -125,6 +125,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { "Tensor)"); m.impl("grouped_topk", torch::kCUDA, &grouped_topk); + // cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16) + m.def("router_gemm_bf16_fp32(Tensor input, Tensor weight) -> Tensor"); + m.impl("router_gemm_bf16_fp32", torch::kCUDA, &router_gemm_bf16_fp32); + // DeepSeek V3 optimized router GEMM for SM90+ m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); // conditionally compiled so impl registration is in source file diff --git a/csrc/ops.h b/csrc/ops.h index 5e2b475fa8c..690342b37ca 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -269,13 +269,13 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2, const int64_t n, const int64_t k, const bool swap_ab); -void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, const int64_t n, - const int64_t k); +void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets, + torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, + const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, + const int64_t padded_m, const int64_t n, + const int64_t k); void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a, torch::Tensor const& b, diff --git a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu index eae500cb632..41cf170a243 100644 --- a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu @@ -263,12 +263,10 @@ void get_cutlass_moe_mm_data_caller( } template -__global__ void compute_pplx_data(int32_t* expert_offsets, - int32_t* problem_sizes1, - int32_t* problem_sizes2, - const int32_t* __restrict__ expert_num_tokens, - const int padded_m, const int n, - const int k) { +__global__ void compute_batched_moe_data( + int32_t* expert_offsets, int32_t* problem_sizes1, int32_t* problem_sizes2, + const int32_t* __restrict__ expert_num_tokens, const int padded_m, + const int n, const int k) { int expert_idx = threadIdx.x; expert_offsets[expert_idx] = expert_idx * padded_m; @@ -289,24 +287,22 @@ __global__ void compute_pplx_data(int32_t* expert_offsets, } } -void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, - const int64_t n, const int64_t k) { +void get_cutlass_batched_moe_mm_data_caller( + torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k) { auto stream = at::cuda::getCurrentCUDAStream(expert_offsets.device().index()); if (num_local_experts * padded_m > SWAP_AB_THRESHOLD) { - compute_pplx_data<<<1, num_local_experts, 0, stream>>>( + compute_batched_moe_data<<<1, num_local_experts, 0, stream>>>( static_cast(expert_offsets.data_ptr()), static_cast(problem_sizes1.data_ptr()), static_cast(problem_sizes2.data_ptr()), static_cast(expert_num_tokens.data_ptr()), padded_m, n, k); } else { - compute_pplx_data<<<1, num_local_experts, 0, stream>>>( + compute_batched_moe_data<<<1, num_local_experts, 0, stream>>>( static_cast(expert_offsets.data_ptr()), static_cast(problem_sizes1.data_ptr()), static_cast(problem_sizes2.data_ptr()), diff --git a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 82ccc19608c..d6e82f1db9f 100644 --- a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -82,13 +82,11 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller( torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2, const int64_t n, const int64_t k, const bool swap_ab); -void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, - const int64_t n, const int64_t k); +void get_cutlass_batched_moe_mm_data_caller( + torch::Tensor& expert_offsets, torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k); #endif void cutlass_scaled_mm_azp_sm75(torch::Tensor& c, torch::Tensor const& a, @@ -319,29 +317,30 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( version_num, ". Required capability: 90, 100, or 120"); } -void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets, - torch::Tensor& problem_sizes1, - torch::Tensor& problem_sizes2, - const torch::Tensor& expert_num_tokens, - const int64_t num_local_experts, - const int64_t padded_m, const int64_t n, - const int64_t k) { +void get_cutlass_batched_moe_mm_data(torch::Tensor& expert_offsets, + torch::Tensor& problem_sizes1, + torch::Tensor& problem_sizes2, + const torch::Tensor& expert_num_tokens, + const int64_t num_local_experts, + const int64_t padded_m, const int64_t n, + const int64_t k) { // This function currently gets compiled only if we have a valid cutlass moe // mm to run it for. int32_t version_num = get_sm_version_num(); #if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120) - get_cutlass_pplx_moe_mm_data_caller(expert_offsets, problem_sizes1, - problem_sizes2, expert_num_tokens, - num_local_experts, padded_m, n, k); + get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1, + problem_sizes2, expert_num_tokens, + num_local_experts, padded_m, n, k); return; #endif - TORCH_CHECK_NOT_IMPLEMENTED( - false, - "No compiled get_cutlass_pplx_moe_mm_data: no cutlass_scaled_mm kernel " - "for CUDA device capability: ", - version_num, ". Required capability: 90, 100, or 120"); + TORCH_CHECK_NOT_IMPLEMENTED(false, + "No compiled get_cutlass_batched_moe_mm_data: no " + "cutlass_scaled_mm kernel " + "for CUDA device capability: ", + version_num, + ". Required capability: 90, 100, or 120"); } void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 39b6bc98a84..8be30b20910 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -489,19 +489,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { &get_cutlass_moe_mm_problem_sizes_from_expert_offsets); // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM and PPLX. It takes expert_num_tokens and non_zero_expert_idxs + // GEMM in batched expert format. It takes expert_num_tokens // as an input, and computes expert_offsets (token start indices of each // expert). In addition to this, it computes problem sizes for each expert's // multiplication used by the two mms called from fused MoE operation. ops.def( - "get_cutlass_pplx_moe_mm_data(Tensor! expert_offsets, " + "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " " Tensor! problem_sizes1, " " Tensor! problem_sizes2, " " Tensor expert_num_tokens, " " int num_local_experts, int padded_m, " " int n, int k) -> ()"); - ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA, - &get_cutlass_pplx_moe_mm_data); + ops.impl("get_cutlass_batched_moe_mm_data", torch::kCUDA, + &get_cutlass_batched_moe_mm_data); // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) ops.def( diff --git a/docker/Dockerfile b/docker/Dockerfile index 717f27b6b23..495a480b758 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -308,7 +308,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ #################### CSRC BUILD IMAGE #################### #################### EXTENSIONS BUILD IMAGE #################### -# Build DeepGEMM, pplx-kernels, DeepEP - runs in PARALLEL with csrc-build +# Build DeepGEMM, DeepEP - runs in PARALLEL with csrc-build # This stage is independent and doesn't affect csrc cache FROM base AS extensions-build ARG CUDA_VERSION @@ -335,10 +335,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Ensure the wheel dir exists so COPY won't fail when DeepGEMM is skipped RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped -# Build pplx-kernels and DeepEP wheels +# Build DeepEP wheels COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management -ARG PPLX_COMMIT_HASH=12cecfd ARG DEEPEP_COMMIT_HASH=73b6ea4 ARG NVSHMEM_VER RUN --mount=type=cache,target=/root/.cache/uv \ @@ -347,7 +346,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ /tmp/install_python_libraries.sh \ --workspace /tmp/ep_kernels_workspace \ --mode wheel \ - ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \ ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} \ ${NVSHMEM_VER:+--nvshmem-ver "$NVSHMEM_VER"} && \ find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete @@ -676,7 +674,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH -# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage +# Install EP kernels wheels (DeepEP) that have been built in the `build` stage RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \ --mount=type=cache,target=/root/.cache/uv \ uv pip install --system ep_kernels/dist/*.whl --verbose \ diff --git a/docker/versions.json b/docker/versions.json index 24f4b6e7d1b..fa090c10c44 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -52,9 +52,6 @@ "DEEPGEMM_GIT_REF": { "default": "477618cd51baffca09c4b0b87e97c03fe827ef03" }, - "PPLX_COMMIT_HASH": { - "default": "12cecfd" - }, "DEEPEP_COMMIT_HASH": { "default": "73b6ea4" }, diff --git a/docs/design/fused_moe_modular_kernel.md b/docs/design/fused_moe_modular_kernel.md index 975df8ba29d..9db356cdf53 100644 --- a/docs/design/fused_moe_modular_kernel.md +++ b/docs/design/fused_moe_modular_kernel.md @@ -15,7 +15,7 @@ Based on the format of the input activations, FusedMoE implementations are broad The input activation format completely depends on the All2All Dispatch being used. * In the Contiguous variant, the All2All Dispatch returns the activations as a contiguous tensor of shape (M, K) along with TopK Ids and TopK weights of shape (M, num_topk). Look at `DeepEPHTPrepareAndFinalize` for an example. -* In the Batched variant, the All2All Dispatch returns the activations as a tensor of shape (num_experts, max_tokens, K). Here, the activations/tokens that subscribe to the same expert are batched together. Note that not all entries of the tensor are valid. The activations tensor is typically accompanied by an `expert_num_tokens` tensor of size `num_experts`, where `expert_num_tokens[i]` indicates the number of valid tokens that subscribe to the ith expert. Look at `PplxPrepareAndFinalize` or `DeepEPLLPrepareAndFinalize` for an example. +* In the Batched variant, the All2All Dispatch returns the activations as a tensor of shape (num_experts, max_tokens, K). Here, the activations/tokens that subscribe to the same expert are batched together. Note that not all entries of the tensor are valid. The activations tensor is typically accompanied by an `expert_num_tokens` tensor of size `num_experts`, where `expert_num_tokens[i]` indicates the number of valid tokens that subscribe to the ith expert. Look at `DeepEPLLPrepareAndFinalize` for an example. The FusedMoE operation is generally made of multiple operations, in both the Contiguous and Batched variants, as described in the diagrams below @@ -132,7 +132,6 @@ class FusedMoEModularKernel: Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & Combine implementation / kernel. For example, -* PplxPrepareAndFinalize type is backed by Pplx All2All kernels, * DeepEPHTPrepareAndFinalize type is backed by DeepEP High-Throughput All2All kernels, and * DeepEPLLPrepareAndFinalize type is backed by DeepEP Low-Latency All2All kernels. @@ -229,7 +228,7 @@ Doing this will add the new implementation to the test suite. ### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` Compatibility The unit test file [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script. -Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` +Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts` As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` compatibility. When invoked with incompatible types, the script will error. @@ -238,7 +237,7 @@ with incompatible types, the script will error. Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py) The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible `FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` types. -Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts` +Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts` ## FusedMoEPrepareAndFinalize Implementations diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 37cc61d4626..a977ce9b9bb 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -656,7 +656,7 @@ vLLM has support for OpenTelemetry tracing: - Added by and reinstated by - Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces` - [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/) -- [User-facing docs](../examples/online_serving/opentelemetry.md) +- [User-facing docs](../../examples/online_serving/opentelemetry/README.md) - [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f) - [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 04ceeede3eb..ac5acb66bdb 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -33,7 +33,6 @@ th { | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass | |---------|--------------------|--------------|---------------|-------|-----------------------|-----------| | naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] | -| pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] | | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] | | flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] | @@ -68,7 +67,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes. There are a number of MoE experts kernel implementations for different quantization types and architectures. Most follow the general API of the base Triton [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts] function. Many have modular kernel adapters, so they can be used with compatible all2all backends. This table lists each experts kernel and its particular properties. -Each kernel must be provided with one of the supported input activation formats. Some flavors of kernels support both standard and batched formats through different entry points, e.g. `TritonExperts` and `BatchedTritonExperts`. Batched format kernels are currently only needed for matching with certain all2all backends, e.g. `pplx` and `DeepEPLLPrepareAndFinalize`. +Each kernel must be provided with one of the supported input activation formats. Some flavors of kernels support both standard and batched formats through different entry points, e.g. `TritonExperts` and `BatchedTritonExperts`. Batched format kernels are currently only needed for matching with certain all2all backends, e.g. `DeepEPLLPrepareAndFinalize`. Similar to the backend kernels, each experts kernel only supports certain quantization formats. For non-modular experts, the activations will be in the original type and quantized internally by the kernel. Modular experts will expect the activations to already be in the quantized format. Both types of experts will yield outputs in the original activation type. @@ -110,5 +109,5 @@ The following table shows "families" of modular kernels that are intended to wor | backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses | |---------|-----------------------------------------|----------------------------------------------| | deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,
`TritonExperts`,
`TritonOrDeepGemmExperts`,
`CutlassExpertsFp8`,
`MarlinExperts` | -| deepep_low_latency,
pplx | `DeepEPLLPrepareAndFinalize`,
`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts` | +| deepep_low_latency | `DeepEPLLPrepareAndFinalize` | `BatchedDeepGemmExperts`,
`BatchedTritonExperts`,
`CutlassBatchedExpertsFp8`,
`BatchedMarlinExperts` | | flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` | diff --git a/docs/governance/committers.md b/docs/governance/committers.md index 2f0780a0897..df874418f1c 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -55,6 +55,7 @@ Sorted alphabetically by GitHub handle: - [@ywang96](https://github.com/ywang96): Multimodality, benchmarks - [@zhuohan123](https://github.com/zhuohan123): Project lead, RL integration, numerics - [@zou3519](https://github.com/zou3519): Compilation +- [@BoyuanFeng](https://github.com/BoyuanFeng): Compilation, CUDAGraph ### Emeritus Committers @@ -113,7 +114,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Multi-modal Input Processing: Components that load and process image/video/audio data into feature tensors - @DarkLight1337, @ywang96, @Isotr0py - torch compile: The torch.compile integration in vLLM, custom passes & transformations - - @ProExpertProg, @zou3519, @youkaichao + - @ProExpertProg, @zou3519, @youkaichao, @BoyuanFeng - State space models: The state space models implementation in vLLM - @tdoublep, @tlrmchlsmth - Reasoning and tool calling parsers @@ -154,7 +155,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - FlashAttention: @LucasWilkinson - FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon - Blackwell Kernels: @mgoin, @yewentao256 -- DeepEP/DeepGEMM/pplx: @mgoin, @yewentao256 +- DeepEP/DeepGEMM: @mgoin, @yewentao256 ### Integrations diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 5f821ef7a4b..eca66041da0 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -369,6 +369,7 @@ th { | `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ | | `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ | | `ArcticForCausalLM` | Arctic | `Snowflake/snowflake-arctic-base`, `Snowflake/snowflake-arctic-instruct`, etc. | | ✅︎ | +| `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ | | `BaiChuanForCausalLM` | Baichuan2, Baichuan | `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc. | ✅︎ | ✅︎ | | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ | diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index 82fde27d71f..d469e20c986 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -8,7 +8,7 @@ EP is typically coupled with Data Parallelism (DP). While DP can be used indepen Before using EP, you need to install the necessary dependencies. We are actively working on making this easier in the future: -1. **Install DeepEP and pplx-kernels**: Set up host environment following vLLM's guide for EP kernels [here](../../tools/ep_kernels). +1. **Install DeepEP**: Set up host environment following vLLM's guide for EP kernels [here](../../tools/ep_kernels). 2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation). 3. **For disaggregated serving**: Install `gdrcopy` by running the [`install_gdrcopy.sh`](../../tools/install_gdrcopy.sh) script (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/). @@ -19,7 +19,6 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to | Backend | Use Case | Features | Best For | |---------|----------|----------|----------| | `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration | -| `pplx` | Single node | Chunked prefill support, efficient intra-node communication | Single-node deployments, development | | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios | | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios | | `flashinfer_all2allv` | MNNVL systems | FlashInfer alltoallv kernels for multi-node NVLink | Systems with NVLink across nodes | @@ -71,12 +70,11 @@ For example, with `TP=2, DP=4` (8 GPUs total): The following command serves a `DeepSeek-V3-0324` model with 1-way tensor parallel, 8-way (attention) data parallel, and 8-way expert parallel. The attention weights are replicated across all GPUs, while the expert weights are split across GPUs. It will work on a H200 (or H20) node with 8 GPUs. For H100, you can try to serve a smaller model or refer to the multi-node deployment section. ```bash -# Single node EP deployment with pplx backend +# Single node EP deployment vllm serve deepseek-ai/DeepSeek-V3-0324 \ --tensor-parallel-size 1 \ # Tensor parallelism across 1 GPU --data-parallel-size 8 \ # Data parallelism across 8 processes - --enable-expert-parallel \ # Enable expert parallelism - --all2all-backend pplx # Use pplx communication backend + --enable-expert-parallel # Enable expert parallelism ``` ## Multi-Node Deployment @@ -197,7 +195,6 @@ vllm serve deepseek-ai/DeepSeek-V3-0324 \ --tensor-parallel-size 1 \ # Tensor parallelism --data-parallel-size 8 \ # Data parallelism --enable-expert-parallel \ # Enable EP - --all2all-backend pplx \ # Use pplx communication backend --enable-eplb \ # Enable load balancer --eplb-config '{"window_size":1000,"step_interval":3000,"num_redundant_experts":2,"log_balancedness":true}' ``` diff --git a/docs/serving/openai_compatible_server.md b/docs/serving/openai_compatible_server.md index 97ed7d45fbb..1053b614eb5 100644 --- a/docs/serving/openai_compatible_server.md +++ b/docs/serving/openai_compatible_server.md @@ -84,7 +84,7 @@ In order for the language model to support chat protocol, vLLM requires the mode a chat template in its tokenizer configuration. The chat template is a Jinja2 template that specifies how roles, messages, and other chat-specific tokens are encoded in the input. -An example chat template for `NousResearch/Meta-Llama-3-8B-Instruct` can be found [here](https://github.com/meta-llama/llama3?tab=readme-ov-file#instruction-tuned-models) +An example chat template for `NousResearch/Meta-Llama-3-8B-Instruct` can be found [here](https://llama.com/docs/model-cards-and-prompt-formats/meta-llama-3/#prompt-template-for-meta-llama-3) Some models do not provide a chat template even though they are instruction/chat fine-tuned. For those models, you can manually specify their chat template in the `--chat-template` parameter with the file path to the chat diff --git a/examples/online_serving/elastic_ep/serve_deepseek_v2.sh b/examples/online_serving/elastic_ep/serve_deepseek_v2.sh index b4e92209996..3ce89e1d86f 100644 --- a/examples/online_serving/elastic_ep/serve_deepseek_v2.sh +++ b/examples/online_serving/elastic_ep/serve_deepseek_v2.sh @@ -64,7 +64,7 @@ vllm serve "$MODEL_NAME" \ --enforce-eager \ --enable-expert-parallel \ --enable-eplb \ - --all2all-backend pplx \ + --all2all-backend allgather_reducescatter \ --num-redundant-experts "$REDUNDANT_EXPERTS" \ --trust-remote-code \ --host "$HOST" \ diff --git a/tests/distributed/test_mq_connect_ip.py b/tests/distributed/test_mq_connect_ip.py new file mode 100644 index 00000000000..4b0cdda3ad9 --- /dev/null +++ b/tests/distributed/test_mq_connect_ip.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Test that MessageQueue uses the local node's IP for binding, +not a remote master_addr. This validates the fix for cross-node +data-parallel where each DP group leader must bind to its own IP. + +The bug: multiproc_executor used `parallel_config.master_addr` as +`connect_ip` for every DP group's MessageQueue. For DP groups whose +leader is NOT on the master node, binding to master_addr fails with +"Cannot assign requested address". + +The fix: use `get_ip()` (local node IP) instead of `master_addr`. +""" + +import pytest +import zmq + +from vllm.distributed.device_communicators.shm_broadcast import MessageQueue +from vllm.utils.network_utils import get_ip + + +def test_mq_bind_with_local_ip(): + """MessageQueue with remote readers should successfully bind + when connect_ip is the local node's IP.""" + # n_reader=2, n_local_reader=1 means 1 remote reader, + # which triggers the remote ZMQ socket bind. + mq = MessageQueue( + n_reader=2, + n_local_reader=1, + connect_ip=get_ip(), + ) + handle = mq.export_handle() + assert handle.remote_subscribe_addr is not None + # The bound address should contain our local IP + local_ip = get_ip() + assert ( + local_ip in handle.remote_subscribe_addr + or f"[{local_ip}]" in handle.remote_subscribe_addr + ) + del mq + + +def test_mq_bind_with_non_local_ip_fails(): + """MessageQueue should fail to bind when connect_ip is a + non-local IP address (simulating the bug where master_addr + from a different node was used).""" + # Use a non-local IP that we definitely can't bind to. + # 198.51.100.1 is from TEST-NET-2 (RFC 5737), never locally assigned. + non_local_ip = "198.51.100.1" + with pytest.raises(zmq.error.ZMQError, match="Cannot assign requested address"): + MessageQueue( + n_reader=2, + n_local_reader=1, + connect_ip=non_local_ip, + ) + + +def test_mq_bind_defaults_to_local_ip(): + """When connect_ip is None, MessageQueue should auto-detect + the local IP and bind successfully.""" + mq = MessageQueue( + n_reader=2, + n_local_reader=1, + connect_ip=None, # should fallback to get_ip() + ) + handle = mq.export_handle() + assert handle.remote_subscribe_addr is not None + del mq + + +if __name__ == "__main__": + test_mq_bind_with_local_ip() + print("PASSED: test_mq_bind_with_local_ip") + test_mq_bind_with_non_local_ip_fails() + print("PASSED: test_mq_bind_with_non_local_ip_fails") + test_mq_bind_defaults_to_local_ip() + print("PASSED: test_mq_bind_defaults_to_local_ip") + print("\nAll tests passed!") diff --git a/tests/entrypoints/openai/responses/test_simple.py b/tests/entrypoints/openai/responses/test_simple.py index b67f0d34115..bbf3cc80ad4 100644 --- a/tests/entrypoints/openai/responses/test_simple.py +++ b/tests/entrypoints/openai/responses/test_simple.py @@ -6,6 +6,7 @@ import pytest_asyncio from openai import OpenAI from ....utils import RemoteOpenAIServer +from .conftest import validate_streaming_event_stack MODEL_NAME = "Qwen/Qwen3-8B" @@ -219,3 +220,23 @@ async def test_extra_sampling_params(client: OpenAI, model_name: str): assert response.status in ["completed", "incomplete"] assert len(response.output) > 0 assert response.output[0].content[0].text # Has text output + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_streaming_types( + pairs_of_event_types: dict[str, str], client: OpenAI, model_name: str +): + stream = await client.responses.create( + model=model_name, + input="tell me a story about a cat in 20 words", + reasoning={"effort": "low"}, + tools=[], + stream=True, + background=False, + ) + events = [] + async for event in stream: + events.append(event) + + validate_streaming_event_stack(events, pairs_of_event_types) diff --git a/tests/entrypoints/openai/test_completion_error.py b/tests/entrypoints/openai/test_completion_error.py index e48cc32e540..1e7a3d0a8c3 100644 --- a/tests/entrypoints/openai/test_completion_error.py +++ b/tests/entrypoints/openai/test_completion_error.py @@ -221,6 +221,19 @@ async def test_completion_error_stream(): assert chunks[-1] == "data: [DONE]\n\n" +def test_json_schema_response_format_missing_schema(): + """When response_format type is 'json_schema' but the json_schema field + is not provided, request construction should raise a validation error + so the API returns 400 instead of 500.""" + with pytest.raises(Exception, match="json_schema.*must be provided"): + CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + response_format={"type": "json_schema"}, + ) + + def test_negative_prompt_token_ids_nested(): """Negative token IDs in prompt (nested list) should raise validation error.""" with pytest.raises(Exception, match="greater than or equal to 0"): diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index 570bf7fc865..9a084769762 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -84,6 +84,209 @@ def ref_paged_attn( return torch.cat(outputs, dim=0) +def _make_paged_kv_metadata( + kv_lens: list[int], + block_size: int, + num_blocks: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build paged-KV metadata tensors for fast_plan_decode tests. + + Returns: + kv_indptr – CPU int32, shape [num_seqs + 1] + kv_indices – CUDA int32, shape [total_blocks] + kv_last_page_lens – CPU int32, shape [num_seqs] + block_tables – CUDA int32, shape [num_seqs, max_blocks_per_seq] + """ + num_seqs = len(kv_lens) + max_blocks = (max(kv_lens) + block_size - 1) // block_size + block_tables = torch.randint( + 0, num_blocks, (num_seqs, max_blocks), dtype=torch.int32, device="cuda" + ) + + indptr_list = [0] + indices_list: list[int] = [] + last_lens_list: list[int] = [] + for i, seq_len in enumerate(kv_lens): + n = (seq_len + block_size - 1) // block_size + indices_list.extend(block_tables[i, :n].cpu().tolist()) + indptr_list.append(indptr_list[-1] + n) + last_lens_list.append(seq_len % block_size or block_size) + + return ( + torch.tensor(indptr_list, dtype=torch.int32, device="cpu"), + torch.tensor(indices_list, dtype=torch.int32, device="cuda"), + torch.tensor(last_lens_list, dtype=torch.int32, device="cpu"), + block_tables, + ) + + +def _make_cg_decode_wrapper( + num_seqs: int, + kv_indices_buffer: torch.Tensor, + workspace_buffer: torch.Tensor, + use_tensor_cores: bool = True, +) -> "flashinfer.BatchDecodeWithPagedKVCacheWrapper": + """Create a cudagraph-enabled BatchDecodeWithPagedKVCacheWrapper. + + *kv_indices_buffer* is shared with the caller so that fast_plan_decode + can avoid the device-to-device index copy on subsequent (cudagraph) calls. + """ + return flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, + "NHD", + use_cuda_graph=True, + paged_kv_indptr_buffer=torch.zeros( + num_seqs + 1, dtype=torch.int32, device="cuda" + ), + paged_kv_indices_buffer=kv_indices_buffer, + paged_kv_last_page_len_buffer=torch.zeros( + num_seqs, dtype=torch.int32, device="cuda" + ), + use_tensor_cores=use_tensor_cores, + ) + + +def test_fast_decode_plan_importable() -> None: + """fast_decode_plan must be importable from flashinfer.decode. + + This is a forward-compatibility smoke test: if FlashInfer reorganises its + public API the import will fail before any other test does. + """ + from flashinfer.decode import fast_decode_plan # noqa: F401 + + assert callable(fast_decode_plan) + + +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode +def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None: + """On the first call fast_plan_decode must route through self.plan() and + flip vllm_first_call to False on the wrapper object.""" + from unittest.mock import patch + + from vllm.v1.attention.backends.flashinfer import fast_plan_decode + + torch.set_default_device("cuda") + set_random_seed(0) + + kv_lens = [128, 64] + block_size = 16 + num_seqs = len(kv_lens) + num_query_heads, num_kv_heads = 8, 2 + head_size = 128 + + kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata( + kv_lens, block_size, NUM_BLOCKS + ) + + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8) + wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices.clone(), workspace) + + assert getattr(wrapper, "vllm_first_call", True) is True + + with patch.object(wrapper, "plan", wraps=wrapper.plan) as mock_plan: + fast_plan_decode( + wrapper, + indptr_cpu=kv_indptr, + indices=kv_indices, + last_page_len_cpu=kv_last_page_lens, + num_qo_heads=num_query_heads, + num_kv_heads=num_kv_heads, + head_dim=head_size, + page_size=block_size, + q_data_type=dtype, + kv_data_type=dtype, + ) + mock_plan.assert_called_once() + + assert wrapper.vllm_first_call is False, ( + "vllm_first_call should be False after the first fast_plan_decode call" + ) + + +@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]]) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode +def test_fast_plan_decode_matches_full_plan( + kv_lens: list[int], + num_heads: tuple[int, int], + head_size: int, + block_size: int, + dtype: torch.dtype, +) -> None: + """fast_plan_decode's cudagraph path (delegating to FlashInfer's + fast_decode_plan) must produce attention output numerically identical to + a standard plan() call. + + Both the warmup call (self.plan) and the subsequent fast call + (fast_decode_plan) are verified against the same reference. + """ + from vllm.v1.attention.backends.flashinfer import fast_plan_decode + + torch.set_default_device("cuda") + set_random_seed(0) + num_seqs = len(kv_lens) + num_query_heads, num_kv_heads = num_heads + + query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype) + key_value_cache = torch.randn( + NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype + ) + + kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata( + kv_lens, block_size, NUM_BLOCKS + ) + + # Reference output via the standard plan() + workspace_ref = torch.empty(128 * 1024 * 1024, dtype=torch.int8) + ref_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_ref, "NHD", use_tensor_cores=True + ) + ref_wrapper.plan( + kv_indptr, + kv_indices, + kv_last_page_lens, + num_query_heads, + num_kv_heads, + head_size, + block_size, + "NONE", + q_data_type=dtype, + kv_data_type=dtype, + ) + ref_output = ref_wrapper.run(query, key_value_cache) + + # CUDAGraph wrapper exercised through fast_plan_decode + kv_indices_buf = kv_indices.clone() + workspace_cg = torch.empty(128 * 1024 * 1024, dtype=torch.int8) + cg_wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices_buf, workspace_cg) + + plan_kwargs: dict = dict( + indptr_cpu=kv_indptr, + indices=kv_indices_buf, + last_page_len_cpu=kv_last_page_lens, + num_qo_heads=num_query_heads, + num_kv_heads=num_kv_heads, + head_dim=head_size, + page_size=block_size, + q_data_type=dtype, + kv_data_type=dtype, + ) + + # First call – warmup path (routes through self.plan) + fast_plan_decode(cg_wrapper, **plan_kwargs) + warmup_output = cg_wrapper.run(query, key_value_cache) + torch.testing.assert_close(warmup_output, ref_output, atol=1e-2, rtol=1e-2) + + # Second call – fast path (routes through fast_decode_plan from FlashInfer) + fast_plan_decode(cg_wrapper, **plan_kwargs) + fast_output = cg_wrapper.run(query, key_value_cache) + torch.testing.assert_close(fast_output, ref_output, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]]) @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES) diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py index d76c57f9eb1..bc99ed57637 100644 --- a/tests/kernels/attention/test_mha_attn.py +++ b/tests/kernels/attention/test_mha_attn.py @@ -9,9 +9,12 @@ Test: import itertools from unittest.mock import patch +import numpy as np import pytest import torch +from vllm.config import get_current_vllm_config +from vllm.config.multimodal import MultiModalConfig from vllm.model_executor.layers.attention import MMEncoderAttention from vllm.platforms import current_platform from vllm.platforms.cpu import CpuPlatform @@ -224,3 +227,110 @@ def test_mha_attn_varlen_forward( ref_output.append(output_i) ref_output = torch.cat(ref_output, dim=1) torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) + + +@pytest.mark.parametrize("var_seq_len", VAR_SEQ_LENS) +@pytest.mark.parametrize( + "dtype", + [torch.bfloat16, torch.half], +) +@pytest.mark.parametrize("device", CUDA_DEVICES) +def test_mha_attn_varlen_forward_flashinfer( + default_vllm_config, + var_seq_len: list[int], + dtype: torch.dtype, + device: str, +): + """Test MMEncoderAttention varlen forward with FLASHINFER backend (head_size=72). + + Exercises the path that uses --mm-encoder-attn-backend=FLASHINFER with + recomputed cu_seqlens, max_seqlen, and sequence_lengths as in qwen3_vl + vision encoder. + """ + pytest.importorskip("flashinfer") + + num_heads = 16 + head_size = 72 + set_random_seed(0) + torch.set_default_device(device) + torch.set_default_dtype(dtype) + + # Override vllm config so get_vit_attn_backend returns FLASHINFER (simulates + # --mm-encoder-attn-backend=FLASHINFER). + vllm_config = get_current_vllm_config() + old_model_config = getattr(vllm_config, "model_config", None) + minimal_model_config = type( + "MinimalModelConfig", + (), + { + "multimodal_config": MultiModalConfig( + mm_encoder_attn_backend=AttentionBackendEnum.FLASHINFER + ), + }, + )() + vllm_config.model_config = minimal_model_config + try: + total_len = sum(var_seq_len) + # Stride of second dim = 3 * num_heads * head_size (same as qwen2_5_vl + # after qkv rearrange and unbind: qkv shape (b, s, 3, head, head_dim)). + qkv = torch.randn(1, total_len, 3, num_heads, head_size) + q, k, v = qkv.unbind(dim=2) + + cu_seqlens_np = np.array( + [0] + list(itertools.accumulate(var_seq_len)), dtype=np.int32 + ) + hidden_size = num_heads * head_size + tp_size = 1 + + sequence_lengths_np = MMEncoderAttention.maybe_compute_sequence_lengths( + AttentionBackendEnum.FLASHINFER, cu_seqlens_np + ) + sequence_lengths = torch.from_numpy(sequence_lengths_np).to( + device, dtype=torch.int32, non_blocking=True + ) + + max_seqlen_val = MMEncoderAttention.compute_max_seqlen( + AttentionBackendEnum.FLASHINFER, cu_seqlens_np + ) + max_seqlen = torch.tensor(max_seqlen_val, device=device, dtype=torch.int32) + + cu_seqlens_np = MMEncoderAttention.maybe_recompute_cu_seqlens( + AttentionBackendEnum.FLASHINFER, + cu_seqlens_np, + hidden_size, + tp_size, + ) + cu_seqlens = torch.from_numpy(cu_seqlens_np).to( + device, dtype=torch.int32, non_blocking=True + ) + + scale = 1.0 / head_size**0.5 + attn = MMEncoderAttention( + num_heads, + head_size, + scale=scale, + num_kv_heads=num_heads, + ) + assert attn.attn_backend == AttentionBackendEnum.FLASHINFER + + output = attn( + q, + k, + v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + + ref_output = [] + for q_i, k_i, v_i in zip( + torch.split(q, var_seq_len, dim=1), + torch.split(k, var_seq_len, dim=1), + torch.split(v, var_seq_len, dim=1), + ): + output_i = ref_attention(q_i, k_i, v_i, scale=scale) + ref_output.append(output_i) + ref_output = torch.cat(ref_output, dim=1) + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) + finally: + vllm_config.model_config = old_model_config diff --git a/tests/kernels/helion/test_pattern_matching.py b/tests/kernels/helion/test_pattern_matching.py new file mode 100644 index 00000000000..1cab249a18c --- /dev/null +++ b/tests/kernels/helion/test_pattern_matching.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test make_fx tracing and inductor pattern matching with HelionKernelWrapper.""" + +import contextlib +from unittest.mock import Mock, patch + +import pytest +import torch + +from vllm.utils.import_utils import has_helion + +if not has_helion(): + pytest.skip( + "Helion is not installed. Install with: pip install vllm[helion]", + allow_module_level=True, + ) + +import helion +import helion.language as hl +from helion._compat import requires_torch_version + +if not requires_torch_version("2.11"): + pytest.skip( + "HigherOrderOp requires PyTorch >= 2.11", + allow_module_level=True, + ) + +from helion._compiler._dynamo.higher_order_ops import ( + helion_kernel_side_table, + helion_kernel_wrapper_mutation, +) +from torch._inductor.pattern_matcher import ( + PatternMatcherPass, + fwd_only, + register_replacement, + select_decomp_table, +) +from torch.fx.experimental.proxy_tensor import make_fx + +from vllm.kernels.helion.config_manager import ConfigManager +from vllm.kernels.helion.register import HelionKernelWrapper + + +@contextlib.contextmanager +def _helion_mock_context(): + configs = { + "default": helion.Config(block_sizes=[64], num_warps=2, num_stages=2), + } + mock_config_manager = Mock(spec=ConfigManager) + mock_config_manager.get_platform_configs = Mock(return_value=configs) + + with ( + patch( + "vllm.kernels.helion.config_manager.ConfigManager.get_instance", + return_value=mock_config_manager, + ), + patch( + "vllm.kernels.helion.utils.get_canonical_gpu_name", + return_value="nvidia_h200", + ), + ): + yield + + +class TestMakeFxHop: + def setup_method(self): + helion_kernel_side_table.reset_table() + + def test_make_fx_symbolic(self): + def raw_add_scale( + x: torch.Tensor, y: torch.Tensor, scale: float + ) -> tuple[torch.Tensor, int, torch.Tensor]: + out_x = torch.empty_like(x) + out_y = torch.empty_like(x) + for tile in hl.tile(x.size()): + out_x[tile] = x[tile] + y[tile] * scale + out_y[tile] = out_x[tile] * 2.0 + return out_x, 42, out_y + + input_x = torch.randn(7, 13) + input_y = torch.randn(7, 13) + scale = 0.5 + + with _helion_mock_context(): + wrapper = HelionKernelWrapper( + raw_kernel_func=raw_add_scale, + op_name="test_make_fx", + fake_impl=lambda *a, **kw: None, + ) + wrapper.register_config_picker(lambda args, keys: "default") + + def fn(x, y): + return wrapper(x, y, scale) + + gm = make_fx(fn, tracing_mode="symbolic")(input_x, input_y) + + hop_nodes = [ + n + for n in gm.graph.nodes + if n.op == "call_function" and n.target is helion_kernel_wrapper_mutation + ] + assert len(hop_nodes) == 1 + node = hop_nodes[0] + + assert node.kwargs["constant_args"]["scale"] == scale + assert set(node.kwargs["tensor_args"]) == {"x", "y"} + + specs = node.kwargs["output_spec"]["leaf_specs"] + tensor_specs = [s for s in specs if s["type"] == "tensor"] + scalar_specs = [s for s in specs if s["type"] == "scalar"] + assert len(tensor_specs) == 2 + assert len(scalar_specs) == 1 + + for spec in tensor_specs: + assert spec["dtype"] == input_x.dtype + + assert scalar_specs[0]["scalar_value"] == 42 + + for val in node.meta["val"]: + assert all(isinstance(s, torch.SymInt) for s in val.shape) + + # Both out_x and out_y are empty_like(x), so output shapes == input shape + input_node = next(n for n in gm.graph.nodes if n.op == "placeholder") + input_shape = input_node.meta["val"].shape + for val in node.meta["val"]: + assert len(val.shape) == len(input_shape) + for out_s, in_s in zip(val.shape, input_shape): + assert out_s == in_s + + def test_pattern_matcher_replaces_with_helion_hop(self): + def raw_silu_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + M, N = x.size() + out = torch.empty_like(x) + for tile_m, tile_n in hl.tile([M, N]): + out[tile_m, tile_n] = ( + torch.nn.functional.silu(x[tile_m, tile_n]) * y[tile_m, tile_n] + ) + return out + + with _helion_mock_context(): + wrapper = HelionKernelWrapper( + raw_kernel_func=raw_silu_mul, + op_name="test_pm_silu_mul", + fake_impl=lambda *a, **kw: None, + ) + wrapper.register_config_picker(lambda args, keys: "default") + + def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(x) * y + + def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return wrapper(x, y) + + inputs = [torch.randn(8, 16), torch.randn(8, 16)] + + pm_pass = PatternMatcherPass(pass_name="test_helion_replacement") + register_replacement(pattern, replacement, inputs, fwd_only, pm_pass) + + def model(x, y): + return torch.nn.functional.silu(x) * y + + decompositions = select_decomp_table() + input_x = torch.randn(8, 16) + input_y = torch.randn(8, 16) + gm = make_fx(model, decompositions, tracing_mode="symbolic")( + input_x, input_y + ) + + def count_hop_nodes(graph): + return sum( + 1 + for n in graph.nodes + if n.op == "call_function" + and n.target is helion_kernel_wrapper_mutation + ) + + assert count_hop_nodes(gm.graph) == 0 + + match_count = pm_pass.apply(gm.graph) + gm.graph.lint() + gm.recompile() + + assert match_count == 1 + assert count_hop_nodes(gm.graph) == 1 + + hop_node = next( + n + for n in gm.graph.nodes + if n.op == "call_function" + and n.target is helion_kernel_wrapper_mutation + ) + + # raw_silu_mul returns empty_like(x), so output shape == input shape + for val in hop_node.meta["val"]: + assert all(isinstance(s, torch.SymInt) for s in val.shape) + + input_node = next(n for n in gm.graph.nodes if n.op == "placeholder") + input_shape = input_node.meta["val"].shape + output_shape = hop_node.meta["val"][0].shape + assert len(output_shape) == len(input_shape) + for out_s, in_s in zip(output_shape, input_shape): + assert out_s == in_s diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py index 02b05be74d1..bee72d58a06 100644 --- a/tests/kernels/helion/test_register.py +++ b/tests/kernels/helion/test_register.py @@ -4,8 +4,7 @@ Unit tests for Helion kernel registration. Tests ConfiguredHelionKernel, HelionKernelWrapper, and PresetConfigSearch -including config picker registration, custom autotuner integration, and -PyTorch op registration. +including config picker registration and custom autotuner integration. """ from unittest.mock import Mock, patch @@ -25,6 +24,7 @@ import helion from vllm.kernels.helion.config_manager import ConfigManager from vllm.kernels.helion.register import ( + _HOP_AVAILABLE, ConfiguredHelionKernel, HelionKernelWrapper, get_kernel_by_name, @@ -451,8 +451,10 @@ class TestHelionKernelWrapper: ): wrapper.get_configured_op() - def test_get_configured_op_returns_cached_op(self, sample_kernel, sample_configs): - """Test get_configured_op returns cached op when already registered.""" + def test_get_configured_op_returns_cached_kernel( + self, sample_kernel, sample_configs + ): + """Test get_configured_op returns cached ConfiguredHelionKernel.""" def fake_impl(*args, **kwargs): return torch.zeros_like(args[0]) @@ -470,6 +472,46 @@ class TestHelionKernelWrapper: mock_config_manager = Mock(spec=ConfigManager) mock_config_manager.get_platform_configs = Mock(return_value=sample_configs) + with ( + patch( + "vllm.kernels.helion.config_manager.ConfigManager.get_instance", + return_value=mock_config_manager, + ), + patch( + "vllm.kernels.helion.utils.get_canonical_gpu_name", + return_value="nvidia_h200", + ), + patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel, + ): + mock_decorated = Mock() + mock_kernel.return_value = Mock(return_value=mock_decorated) + + result1 = wrapper.get_configured_op() + result2 = wrapper.get_configured_op() + assert result1 is result2 + + @pytest.mark.skipif( + _HOP_AVAILABLE, reason="CustomOp path not used when HOP available" + ) + def test_get_or_register_custom_op_returns_cached_op( + self, sample_kernel, sample_configs + ): + def fake_impl(*args, **kwargs): + return torch.zeros_like(args[0]) + + def default_picker(args, config_keys): + return "default" + + wrapper = HelionKernelWrapper( + raw_kernel_func=sample_kernel, + op_name="test_kernel", + fake_impl=fake_impl, + ) + wrapper._config_picker = default_picker + + mock_config_manager = Mock(spec=ConfigManager) + mock_config_manager.get_platform_configs = Mock(return_value=sample_configs) + existing_op = Mock() mock_namespace = Mock() mock_namespace.test_kernel = existing_op @@ -488,12 +530,15 @@ class TestHelionKernelWrapper: ): mock_decorated = Mock() mock_kernel.return_value = Mock(return_value=mock_decorated) - result = wrapper.get_configured_op() + result = wrapper._get_or_register_custom_op() assert result is existing_op - def test_get_configured_op_registers_new_op(self, sample_kernel, sample_configs): - """Test get_configured_op creates and registers new op.""" - + @pytest.mark.skipif( + _HOP_AVAILABLE, reason="CustomOp path not used when HOP available" + ) + def test_get_or_register_custom_op_registers_new_op( + self, sample_kernel, sample_configs + ): def fake_impl(*args, **kwargs): return torch.zeros_like(args[0]) @@ -542,11 +587,10 @@ class TestHelionKernelWrapper: ): mock_decorated = Mock() mock_kernel.return_value = Mock(return_value=mock_decorated) - result = wrapper.get_configured_op() + result = wrapper._get_or_register_custom_op() mock_register.assert_called_once() assert result is new_op - # Check that op_func is the decorated kernel, not ConfiguredHelionKernel assert mock_register.call_args[1]["op_func"] is mock_decorated diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 87cf0453bea..9f67129616f 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -37,7 +37,6 @@ from vllm.utils.import_utils import ( has_deep_ep, has_deep_gemm, has_mori, - has_pplx, ) from .mk_objects import ( @@ -206,10 +205,6 @@ class Config: info = expert_info(self.fused_experts_type) return info.needs_deep_gemm - def needs_pplx(self): - info = prepare_finalize_info(self.prepare_finalize_type) - return info.backend == "pplx" - def needs_deep_ep(self): info = prepare_finalize_info(self.prepare_finalize_type) return ( @@ -290,8 +285,6 @@ class Config: return False, "Needs DeepEP, but DeepEP not available." if self.needs_deep_gemm() and not has_deep_gemm(): return False, "Needs DeepGEMM, but DeepGEMM not available." - if self.needs_pplx() and not has_pplx(): # noqa: SIM103 - return False, "Needs PPLX, but PPLX not available." if self.needs_aiter() and not has_aiter(): # noqa: SIM103 return False, "Needs Aiter, but Aiter not available." if self.needs_mori() and not has_mori(): # noqa: SIM103 diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 981f993427b..0ea414c3af4 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -39,7 +39,6 @@ from vllm.utils.import_utils import ( has_deep_ep, has_deep_gemm, has_mori, - has_pplx, ) @@ -238,19 +237,6 @@ if has_mori(): supports_apply_weight_on_input=False, ) -if has_pplx(): - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - ) - - register_prepare_and_finalize( - PplxPrepareAndFinalize, - batched_format, - common_float_and_int_types, - blocked_quantization_support=True, - backend="pplx", - ) - if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100): from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501 FlashInferCutlassMoEPrepareAndFinalize, diff --git a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py index 3cdc7b82130..702584f9da5 100644 --- a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py +++ b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py @@ -125,7 +125,7 @@ if __name__ == "__main__": description=( "Run single prepare-finalize & fused-experts combination test" "Example : python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel " # noqa: E501 - "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts" + "--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts" ) ) args = parser.parse_args() diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index ec31e66140a..cd1d0a0afe9 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -14,7 +14,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.config import VllmConfig, set_current_vllm_config from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe -from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_deep_gemm from vllm.utils.torch_utils import cuda_device_count_stateless, set_random_seed from vllm.v1.worker.workspace import init_workspace_manager @@ -39,12 +39,12 @@ from .modular_kernel_tools.parallel_utils import ( ) has_any_multi_gpu_package = ( - has_deep_ep() or has_deep_gemm() or has_pplx() or has_flashinfer_cutlass_fused_moe() + has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe() ) meets_multi_gpu_requirements = pytest.mark.skipif( not has_any_multi_gpu_package, - reason="Requires deep_ep or deep_gemm or pplx or flashinfer packages", + reason="Requires deep_ep or deep_gemm or flashinfer packages", ) if current_platform.is_fp8_fnuz(): @@ -341,7 +341,7 @@ if __name__ == "__main__": description=( "Run single prepare-finalize & fused-experts combination test" "Example : python3 -m tests.kernels.moe.test_modular_kernel_combinations " - "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts" + "--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts" ) ) args = parser.parse_args() diff --git a/tests/kernels/moe/test_pplx_cutlass_moe.py b/tests/kernels/moe/test_pplx_cutlass_moe.py deleted file mode 100644 index d8a6600743e..00000000000 --- a/tests/kernels/moe/test_pplx_cutlass_moe.py +++ /dev/null @@ -1,365 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest -import torch - -from tests.kernels.utils import torch_experts -from vllm import _custom_ops as ops -from vllm.config import VllmConfig, set_current_vllm_config -from vllm.model_executor.layers.fused_moe import fused_topk -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEParallelConfig, - RoutingMethodType, - fp8_w8a8_moe_quant_config, -) -from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8 -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel -from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import set_random_seed -from vllm.v1.worker.workspace import init_workspace_manager - -from ...utils import multi_gpu_test -from .parallel_utils import ProcessGroupInfo, parallel_launch - -try: - from pplx_kernels import AllToAll - from pplx_kernels.nvshmem import ( - nvshmem_alloc_empty_unique_id, - nvshmem_finalize, - nvshmem_get_unique_id, - nvshmem_init, - ) - - has_pplx = True -except ImportError: - has_pplx = False - -requires_pplx = pytest.mark.skipif( - not has_pplx, - reason="Requires PPLX kernels", -) - -NUM_EXPERTS = [40, 64] -TOP_KS = [6, 8] - - -def rank_chunk(num, r, w): - rem = num % w - return (num // w) + (1 if r < rem else 0) - - -def chunk_by_rank(t, r, w): - num = t.shape[0] - chunk = rank_chunk(num, r, w) - rem = num % w - if rem == 0 or r < rem: - return t[(r * chunk) : (r + 1) * chunk].contiguous() - else: - long_chunks = (num // w + 1) * rem - short_chunks = (r - rem) * chunk - start = long_chunks + short_chunks - return t[start : start + chunk].contiguous() - - -def pplx_cutlass_moe( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - w1_scale: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - a1_scale: torch.Tensor, - out_dtype, - per_act_token: bool, - per_out_ch: bool, - group_name: str | None, -): - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - ) - - init_workspace_manager(torch.cuda.current_device()) - - assert torch.cuda.current_device() == pgi.local_rank - - num_tokens, hidden_dim = a.shape - intermediate_dim = w2.shape[2] - num_experts = w1.shape[0] - block_size = hidden_dim # TODO support more cases - device = pgi.device - rank = pgi.rank - world_size = pgi.world_size - rank_num_tokens = rank_chunk(num_tokens, rank, world_size) - max_num_tokens = rank_chunk(num_tokens, 0, world_size) - topk = topk_ids.shape[1] - - if block_size == hidden_dim: - scale_elems = 4 # hack to circumvent pplx data format requirements - else: - scale_elems = (hidden_dim + block_size - 1) // block_size - - args = dict( - max_num_tokens=max_num_tokens, - num_experts=num_experts, - experts_per_token=topk, - rank=rank, - world_size=world_size, - dp_size=dp_size, - hidden_dim=hidden_dim, - hidden_dim_bytes=hidden_dim, # because a.dtype.itemsize == 1 - hidden_dim_scale_bytes=scale_elems * torch.float32.itemsize, - ) - - if group_name is None: - ata = AllToAll.internode(**args) - else: - args["group_name"] = group_name - ata = AllToAll.intranode(**args) - - w1 = w1.to(device) - w2 = w2.to(device) - w1_scale = w1_scale.to(device) - w2_scale = w2_scale.to(device) - a1_scale = a1_scale.to(device) - - assert num_experts % world_size == 0 - num_local_experts = cdiv(num_experts, world_size) - num_dispatchers = pgi.world_size // dp_size - - prepare_finalize = PplxPrepareAndFinalize( - ata, - max_num_tokens=max_num_tokens, - num_local_experts=num_local_experts, - num_dispatchers=num_dispatchers, - ) - - def make_moe_config() -> FusedMoEConfig: - return FusedMoEConfig( - num_experts=num_experts, - experts_per_token=topk, - hidden_dim=hidden_dim, - intermediate_size_per_partition=intermediate_dim, - num_local_experts=num_local_experts, - num_logical_experts=num_experts, - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), - activation=MoEActivation.SILU, - in_dtype=torch.bfloat16, - device="cuda", - routing_method=RoutingMethodType.Llama4, - ) - - experts = CutlassBatchedExpertsFp8( - moe_config=make_moe_config(), - quant_config=fp8_w8a8_moe_quant_config( - per_act_token_quant=per_act_token, - per_out_ch_quant=per_out_ch, - w1_scale=chunk_by_rank(w1_scale, rank, world_size), - w2_scale=chunk_by_rank(w2_scale, rank, world_size), - a1_scale=chunk_by_rank(a1_scale, rank, world_size) - if per_act_token - else a1_scale[rank], - ), - max_num_tokens=max_num_tokens, - num_dispatchers=num_dispatchers, - ) - - fused_cutlass_experts = FusedMoEModularKernel( - prepare_finalize, - experts, - inplace=False, - ) - - a_chunk = chunk_by_rank(a, rank, world_size).to(device) - chunk_topk_weight = chunk_by_rank(topk_weights, rank, world_size).to(device) - chunk_topk_ids = ( - chunk_by_rank(topk_ids, rank, world_size).to(torch.uint32).to(device) - ) - - out = fused_cutlass_experts( - a_chunk, - chunk_by_rank(w1, rank, world_size), - chunk_by_rank(w2, rank, world_size), - chunk_topk_weight, - chunk_topk_ids, - global_num_experts=num_experts, - expert_map=None, # TODO - ) - - torch.cuda.synchronize() - - ata.destroy() - - return out[:rank_num_tokens] - - -vllm_config = VllmConfig() - - -def _pplx_moe( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - w1_scale: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - a1_scale: torch.Tensor, - out_dtype, - a_full: torch.Tensor, - w1_full: torch.Tensor, - w2_full: torch.Tensor, - per_act_token: bool, - per_out_ch: bool, - use_internode: bool, -): - try: - if use_internode: - uid = ( - nvshmem_get_unique_id() - if pgi.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - torch.distributed.broadcast(uid, src=0) - nvshmem_init(uid, pgi.rank, pgi.world_size) - else: - group_ranks = list(range(pgi.world_size)) - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo") - group_name = cpu_group.group_name - - with set_current_vllm_config(vllm_config): - torch_output = torch_experts( - a_full, w1_full, w2_full, topk_weights, topk_ids - ) - pplx_output = pplx_cutlass_moe( - pgi, - dp_size, - a, - w1, - w2, - w1_scale, - w2_scale, - topk_weights, - topk_ids, - a1_scale, - out_dtype, - per_act_token, - per_out_ch, - group_name, - ) - - torch_output = chunk_by_rank(torch_output, pgi.rank, pgi.world_size).to( - pplx_output.device - ) - - # Uncomment if more debugging is needed - # print("PPLX OUT:", pplx_output) - # print("TORCH OUT:", torch_output) - - torch.testing.assert_close(pplx_output, torch_output, atol=0.05, rtol=0) - finally: - if use_internode: - nvshmem_finalize() - - -@pytest.mark.parametrize("m", [2, 224]) -@pytest.mark.parametrize("n", [3072]) -@pytest.mark.parametrize("k", [1536]) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("per_act_token", [True, False]) -@pytest.mark.parametrize("per_out_ch", [True, False]) -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) # , [4, 2]]) -@pytest.mark.parametrize("use_internode", [False]) -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif( - (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))( - current_platform.get_device_capability() - ), - reason="Grouped gemm is not supported on this GPU type.", -) -@requires_pplx -def test_cutlass_moe_pplx( - m: int, - n: int, - k: int, - e: int, - topk: int, - per_act_token: bool, - per_out_ch: bool, - world_dp_size: tuple[int, int], - use_internode: bool, -): - set_random_seed(7) - - with set_current_vllm_config(vllm_config): - dtype = torch.half - - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10.0 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10.0 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10.0 - - n_b_scales = 2 * n if per_out_ch else 1 - k_b_scales = k if per_out_ch else 1 - - w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float8_e4m3fn) - w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn) - w1_scale = torch.empty((e, n_b_scales, 1), device="cuda", dtype=torch.float32) - w2_scale = torch.empty((e, k_b_scales, 1), device="cuda", dtype=torch.float32) - - for expert in range(e): - w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant( - w1[expert], use_per_token_if_dynamic=per_out_ch - ) - w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant( - w2[expert], use_per_token_if_dynamic=per_out_ch - ) - - w1_d = torch.empty_like(w1) - w2_d = torch.empty_like(w2) - for expert in range(e): - w1_d[expert] = (w1_q[expert].float() * w1_scale[expert]).half() - w2_d[expert] = (w2_q[expert].float() * w2_scale[expert]).half() - - score = torch.randn((m, e), device="cuda", dtype=dtype) - topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) - - world_size, dp_size = world_dp_size - a_scale1 = ( - torch.randn( - (m if per_act_token else 1, 1), device="cuda", dtype=torch.float32 - ) - / 10.0 - ) - if not per_act_token: - a_scale1 = a_scale1.repeat(world_size, 1) - - parallel_launch( - world_size, - _pplx_moe, - dp_size, - a, - w1_q, - w2_q, - w1_scale, - w2_scale, - topk_weights, - topk_ids, - a_scale1, - dtype, - a, - w1_d, - w2_d, - per_act_token, - per_out_ch, - use_internode, - ) diff --git a/tests/kernels/moe/test_pplx_moe.py b/tests/kernels/moe/test_pplx_moe.py deleted file mode 100644 index deb3b9eb4d7..00000000000 --- a/tests/kernels/moe/test_pplx_moe.py +++ /dev/null @@ -1,1021 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the MOE layers. - -Run `pytest tests/kernels/test_pplx_moe.py`. -""" - -import copy -import itertools -import textwrap -import traceback -from collections.abc import Callable - -import pytest -import torch - -try: - from pplx_kernels import AllToAll - from pplx_kernels.nvshmem import ( - nvshmem_alloc_empty_unique_id, - nvshmem_finalize, - nvshmem_get_unique_id, - nvshmem_init, - ) - - has_pplx = True -except ImportError: - has_pplx = False - -from tests.kernels.moe.modular_kernel_tools.parallel_utils import _set_vllm_config -from tests.kernels.moe.utils import ( - make_dummy_moe_config, - make_shared_experts, - make_test_weights, - naive_batched_moe, -) -from tests.kernels.quant_utils import dequant -from tests.kernels.utils import torch_experts -from vllm.config import VllmConfig, set_current_vllm_config -from vllm.model_executor.layers.fused_moe import fused_topk, override_config -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig -from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts -from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceDelegate, -) -from vllm.utils.math_utils import round_up -from vllm.utils.torch_utils import set_random_seed -from vllm.v1.worker.workspace import init_workspace_manager - -from ...utils import multi_gpu_test -from .parallel_utils import ProcessGroupInfo, parallel_launch - -requires_pplx = pytest.mark.skipif( - not has_pplx, - reason="Requires PPLX kernels", -) - -BATCHED_MOE_MNK_FACTORS = [ - (1, 128, 128), - (33, 2048, 128), - (64, 128, 2048), - (222, 128, 128), - (222, 2048, 1024), -] - -PPLX_COMBOS = [ - # TODO(bnell): figure out why this fails, seems to be test problem - # (1, 128, 128), - (2, 128, 512), - (3, 1024, 2048), - (4, 128, 128), - (32, 1024, 512), - (45, 512, 2048), - (64, 1024, 512), - (222, 2048, 1024), - (256, 1408, 2048), -] - -NUM_EXPERTS = [8, 64] -TOP_KS = [1, 2, 6] -DTYPES = [torch.float8_e4m3fn, torch.bfloat16] - -vllm_config = VllmConfig() - - -def torch_prepare( - a: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - max_num_tokens: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - assert topk_ids.dim() == 2 - assert topk_ids.shape[0] == a.shape[0] - - num_tokens, hidden_dim = a.shape - topk = topk_ids.shape[1] - - tokens_per_expert = torch.bincount(topk_ids.view(-1), minlength=num_experts) - - assert tokens_per_expert.numel() == num_experts - - if max_num_tokens is None: - max_num_tokens = int(tokens_per_expert.max().item()) - - b_a = torch.zeros( - (num_experts, max_num_tokens, hidden_dim), dtype=a.dtype, device=a.device - ) - - token_counts = torch.zeros(num_experts, dtype=torch.int, device=a.device) - - for token in range(num_tokens): - for j in range(topk): - expert_id = topk_ids[token, j] - idx = token_counts[expert_id] - b_a[expert_id, idx : idx + 1, :] = a[token, :] - token_counts[expert_id] = token_counts[expert_id] + 1 - - return b_a, tokens_per_expert - - -def torch_finalize( - b_out: torch.Tensor, topk_weight: torch.Tensor, topk_ids: torch.Tensor -) -> torch.Tensor: - num_tokens = topk_ids.shape[0] - num_experts = b_out.shape[0] - K = b_out.shape[-1] - out = torch.zeros((num_tokens, K), dtype=b_out.dtype, device=b_out.device) - expert_counts = torch.zeros(num_experts, dtype=torch.int, device=b_out.device) - for token in range(num_tokens): - expert_ids = topk_ids[token] - for i in range(expert_ids.numel()): - expert_id = expert_ids[i] - idx = expert_counts[expert_id] - out[token, :] = ( - out[token, :] - + b_out[expert_id, idx : idx + 1, :] * topk_weight[token, i] - ) - expert_counts[expert_id] = expert_counts[expert_id] + 1 - - return out - - -def torch_batched_moe( - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weight: torch.Tensor, - topk_ids: torch.Tensor, -) -> torch.Tensor: - num_experts = w1.shape[0] - b_a, tokens_per_expert = torch_prepare(a, topk_ids, num_experts) - assert b_a.dim() == 3 - num_tokens, topk = topk_ids.shape - _, max_num_tokens, K = b_a.shape - assert num_experts == b_a.shape[0] and w2.shape[1] == K - out = torch.zeros( - (num_experts, max_num_tokens, K), dtype=b_a.dtype, device=b_a.device - ) - tmp = torch.empty( - (max_num_tokens, w1.shape[1] // 2), dtype=b_a.dtype, device=b_a.device - ) - for expert in range(num_experts): - num = tokens_per_expert[expert] - if num > 0: - torch.ops._C.silu_and_mul( - tmp[:num], b_a[expert, :num, :] @ w1[expert].transpose(0, 1) - ) - out[expert, :num, :] = tmp[:num] @ w2[expert].transpose(0, 1) - - return torch_finalize(out, topk_weight, topk_ids) - - -@pytest.mark.parametrize("m,n,k", BATCHED_MOE_MNK_FACTORS) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_fused_moe_batched_experts( - m: int, - n: int, - k: int, - e: int, - topk: int, - dtype: torch.dtype, - workspace_init, -): - set_random_seed(7) - - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 - score = torch.randn((m, e), device="cuda", dtype=dtype) - - with set_current_vllm_config(vllm_config): - topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) - baseline_output = torch_experts( - a, w1, w2, topk_weight, topk_ids - ) # only for baseline - torch_output = torch_batched_moe(a, w1, w2, topk_weight, topk_ids) - batched_output = naive_batched_moe( - a, w1, w2, topk_weight, topk_ids - ) # pick torch_experts or this - - torch.testing.assert_close(baseline_output, torch_output, atol=2e-2, rtol=0) - torch.testing.assert_close(baseline_output, batched_output, atol=2e-2, rtol=0) - - -def create_pplx_prepare_finalize( - num_tokens: int, - hidden_dim: int, - topk: int, - num_experts: int, - rank: int, - dp_size: int, - world_size: int, - in_dtype: torch.dtype, - quant_dtype: torch.dtype | None, - block_shape: list[int] | None, - per_act_token_quant: bool, - group_name: str | None, -): - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - pplx_hidden_dim_scale_bytes, - ) - - max_num_tokens = max(rank_chunk(num_tokens, 0, world_size), 1) - num_local_experts = rank_chunk(num_experts, 0, world_size) - - hidden_dim_bytes, scale_bytes = pplx_hidden_dim_scale_bytes( - max_num_tokens, - hidden_dim, - in_dtype, - quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - ) - - args = dict( - max_num_tokens=max_num_tokens, - num_experts=num_experts, - experts_per_token=topk, - rank=rank, - world_size=world_size, - dp_size=dp_size, - hidden_dim=hidden_dim, - hidden_dim_bytes=hidden_dim_bytes, - hidden_dim_scale_bytes=scale_bytes, - ) - - if group_name is None: - ata = AllToAll.internode(**args) - else: - args["group_name"] = group_name - ata = AllToAll.intranode(**args) - - prepare_finalize = PplxPrepareAndFinalize( - ata, - max_num_tokens=max_num_tokens, - num_local_experts=num_local_experts, - num_dispatchers=world_size // dp_size, - ) - - return prepare_finalize, ata - - -def rank_chunk(num: int, r: int, w: int) -> int: - rem = num % w - return (num // w) + (1 if r < rem else 0) - - -def chunk_by_rank(t: torch.Tensor, r: int, w: int) -> torch.Tensor: - chunk = rank_chunk(t.shape[0], r, w) - return t[(r * chunk) : (r + 1) * chunk] - - -def maybe_chunk_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None: - if t is not None: - return chunk_by_rank(t, r, w) - else: - return t - - -def chunk_scales_by_rank(t: torch.Tensor | None, r: int, w: int) -> torch.Tensor | None: - if t is not None and t.numel() > 1: - chunk = rank_chunk(t.shape[0], r, w) - return t[(r * chunk) : (r + 1) * chunk] - else: - return t - - -def chunk_scales(t: torch.Tensor | None, start: int, end: int) -> torch.Tensor | None: - if t is not None and t.numel() > 1: - return t[start:end] - else: - return t - - -def dummy_work(a: torch.Tensor) -> torch.Tensor: - return a * 1.1 - - -def pplx_prepare_finalize( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - topk_weight: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - quant_dtype: torch.dtype | None, - block_shape: list[int] | None, - per_act_token_quant: bool, - group_name: str | None, -) -> torch.Tensor: - assert torch.cuda.current_device() == pgi.local_rank - - topk = topk_ids.shape[1] - num_tokens, hidden_dim = a.shape - device = pgi.device - rank = pgi.rank - world_size = pgi.world_size - - topk_ids = topk_ids.to(dtype=torch.uint32) - - prepare_finalize, ata = create_pplx_prepare_finalize( - num_tokens, - hidden_dim, - topk, - num_experts, - rank, - dp_size, - world_size, - a.dtype, - quant_dtype, - block_shape, - per_act_token_quant, - group_name, - ) - - assert a.shape[0] == topk_ids.shape[0] - - a_chunk = chunk_by_rank(a, rank, world_size).to(device) - chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device) - chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device) - - assert a_chunk.shape[0] == chunk_topk_ids.shape[0] - - out = torch.full( - a_chunk.shape, - torch.nan, - dtype=a.dtype, - device=device, - ) - - if quant_dtype is not None and not per_act_token_quant and block_shape is None: - a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - else: - a1_scale = None - a2_scale = None - - b_a, b_a_scale, expert_num_tokens, _, _ = prepare_finalize.prepare( - a_chunk, - chunk_topk_weight, - chunk_topk_ids, - num_experts, - None, - False, - FusedMoEQuantConfig.make( - quant_dtype, - per_act_token_quant=per_act_token_quant, - per_out_ch_quant=False, - block_shape=block_shape, - a1_scale=a1_scale, - a2_scale=a2_scale, - ), - ) - - b_a = dummy_work(dequant(b_a, b_a_scale, block_shape, per_act_token_quant, a.dtype)) - - prepare_finalize.finalize( - out, - b_a, - chunk_topk_weight, - chunk_topk_ids, - False, - weight_and_reduce_impl=TopKWeightAndReduceDelegate(), - ) - - torch.cuda.synchronize() - - ata.destroy() - - num_tokens = a_chunk.shape[0] - - return out[:num_tokens] - - -def _pplx_prepare_finalize( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - score: torch.Tensor, - topk: torch.Tensor, - num_experts: int, - quant_dtype: torch.dtype | None, - block_shape: list[int] | None, - per_act_token_quant: bool, - use_internode: bool, -): - try: - if use_internode: - uid = ( - nvshmem_get_unique_id() - if pgi.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - torch.distributed.broadcast(uid, src=0) - nvshmem_init(uid, pgi.rank, pgi.world_size) - group_name = None - else: - group_ranks = list(range(pgi.world_size)) - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo") - group_name = cpu_group.group_name - - topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) - m, k = a.shape - - a_rep = torch.repeat_interleave(dummy_work(a), topk, dim=0) - - torch_output = ( - a_rep.view(m, topk, k) * topk_weight.view(m, topk, 1).to(a_rep.dtype) - ).sum(dim=1) - - pplx_output = pplx_prepare_finalize( - pgi, - dp_size, - a, - topk_weight, - topk_ids, - num_experts, - quant_dtype, - block_shape, - per_act_token_quant, - group_name, - ) - - torch_output = chunk_by_rank(torch_output, pgi.rank, pgi.world_size).to( - pgi.device - ) - - torch.testing.assert_close(pplx_output, torch_output, atol=3e-2, rtol=3e-2) - finally: - if use_internode: - nvshmem_finalize() - - -@pytest.mark.parametrize("mnk", PPLX_COMBOS) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("per_act_token_quant", [False, True]) -@pytest.mark.parametrize("block_shape", [None, [128, 128]]) -@pytest.mark.parametrize("use_internode", [False]) -@pytest.mark.optional -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_prepare_finalize_slow( - mnk: tuple[int, int, int], - e: int, - topk: int, - dtype: torch.dtype, - world_dp_size: tuple[int, int], - per_act_token_quant: bool, - block_shape: list[int] | None, - use_internode: bool, -): - if dtype == torch.float8_e4m3fn: - use_fp8_w8a8 = True - act_dtype = torch.bfloat16 - quant_dtype = dtype - else: - use_fp8_w8a8 = False - act_dtype = dtype - quant_dtype = None - - if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None): - pytest.skip("Skip quantization test for non-quantized type") - - if per_act_token_quant and block_shape is not None: - pytest.skip("Skip illegal quantization combination") - - set_random_seed(7) - m, n, k = mnk - world_size, dp_size = world_dp_size - device = "cuda" - - a = torch.randn((m, k), device=device, dtype=act_dtype) / 10 - score = torch.randn((m, e), device=device, dtype=act_dtype) - - parallel_launch( - world_size, - _pplx_prepare_finalize, - dp_size, - a, - score, - topk, - e, - quant_dtype, - block_shape, - per_act_token_quant, - use_internode, - ) - - -def pplx_moe( - group_name: str | None, - rank: int, - world_size: int, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weight: torch.Tensor, - topk_ids: torch.Tensor, - w1_scale: torch.Tensor | None = None, - w2_scale: torch.Tensor | None = None, - a1_scale: torch.Tensor | None = None, - a2_scale: torch.Tensor | None = None, - quant_dtype: torch.dtype | None = None, - per_act_token_quant=False, - block_shape: list[int] | None = None, - use_compile: bool = False, - use_cudagraphs: bool = True, - shared_experts: torch.nn.Module | None = None, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - num_tokens, hidden_dim = a.shape - num_experts = w1.shape[0] - topk = topk_ids.shape[1] - max_num_tokens = round_up(rank_chunk(a.shape[0], 0, world_size), 16) - - prepare_finalize, ata = create_pplx_prepare_finalize( - num_tokens, - hidden_dim, - topk, - num_experts, - rank, - dp_size, - world_size, - a.dtype, - quant_dtype, - block_shape, - per_act_token_quant, - group_name, - ) - - topk_ids = topk_ids.to(dtype=torch.uint32) - - # Note: workers with the same dp_rank must use the exact same inputs. - a_chunk = chunk_by_rank(a, rank, world_size) - chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size) - chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size) - - # Chunking weights like this only works for batched format - w1_chunk = chunk_by_rank(w1, rank, world_size) - w2_chunk = chunk_by_rank(w2, rank, world_size) - w1_scale_chunk = maybe_chunk_by_rank(w1_scale, rank, world_size) - w2_scale_chunk = maybe_chunk_by_rank(w2_scale, rank, world_size) - a1_scale_chunk = chunk_scales_by_rank(a1_scale, rank, world_size) - a2_scale_chunk = chunk_scales_by_rank(a2_scale, rank, world_size) - - quant_config = FusedMoEQuantConfig.make( - quant_dtype, - block_shape=block_shape, - per_act_token_quant=per_act_token_quant, - w1_scale=w1_scale_chunk, - w2_scale=w2_scale_chunk, - a1_scale=a1_scale_chunk, - a2_scale=a2_scale_chunk, - ) - - experts = BatchedTritonExperts( - max_num_tokens=max_num_tokens, - num_dispatchers=prepare_finalize.num_dispatchers(), - quant_config=quant_config, - moe_config=make_dummy_moe_config(), - ) - - fused_experts = FusedMoEModularKernel( - prepare_finalize, - experts, - shared_experts, - inplace=False, - ) - - # Note: for now use_compile will error out if the problem size is - # large enough to trigger chunking. I'm leaving the flag and - # setup code in case we are able to revisit this later. - if use_compile: - _fused_experts = torch.compile( - fused_experts, backend="inductor", fullgraph=True - ) - torch._dynamo.mark_dynamic(a_chunk, 0) - torch._dynamo.mark_dynamic(chunk_topk_weight, 0) - torch._dynamo.mark_dynamic(chunk_topk_ids, 0) - else: - _fused_experts = fused_experts - - out = _fused_experts( - a_chunk, - w1_chunk, - w2_chunk, - chunk_topk_weight, - chunk_topk_ids, - global_num_experts=num_experts, - ) - - if use_cudagraphs: - if isinstance(out, tuple): - out[0].fill_(0) - out[1].fill_(0) - else: - out.fill_(0) - stream = torch.cuda.Stream() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - out = _fused_experts( - a_chunk, - w1_chunk, - w2_chunk, - chunk_topk_weight, - chunk_topk_ids, - global_num_experts=num_experts, - ) - - torch.cuda.synchronize() - graph.replay() - - torch.cuda.synchronize() - - ata.destroy() - - return out - - -def _pplx_moe( - pgi: ProcessGroupInfo, - dp_size: int, - a: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - score: torch.Tensor, - topk: int, - num_experts: int, - w1_s: torch.Tensor | None = None, - w2_s: torch.Tensor | None = None, - quant_dtype: torch.dtype | None = None, - per_act_token_quant: bool = False, - block_shape: list[int] | None = None, - use_internode: bool = False, - shared_experts: torch.nn.Module | None = None, -): - try: - if use_internode: - uid = ( - nvshmem_get_unique_id() - if pgi.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - torch.distributed.broadcast(uid, src=0) - nvshmem_init(uid, pgi.rank, pgi.world_size) - group_name = None - else: - group_ranks = list(range(pgi.world_size)) - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo") - group_name = cpu_group.group_name - - m, k = a.shape - e, _, n = w2.shape - - moe_config = get_default_config(m, e, n, k, topk, a.dtype, False) - - device = torch.device("cuda", pgi.rank) - rank = pgi.rank - world_size = pgi.world_size - - a = a.to(device) - w1 = w1.to(device) - w2 = w2.to(device) - w1_s = w1_s.to(device) if w1_s is not None else None - w2_s = w2_s.to(device) if w2_s is not None else None - - if quant_dtype is not None and not per_act_token_quant and block_shape is None: - a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32) - else: - a1_scale = None - a2_scale = None - - with set_current_vllm_config(vllm_config), override_config(moe_config): - topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) - - shared_output = shared_experts(a) if shared_experts is not None else None - - torch_output = torch_experts( - a, - w1, - w2, - topk_weight, - topk_ids, - w1_scale=w1_s, - w2_scale=w2_s, - a1_scale=a1_scale, - a2_scale=a2_scale, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - ) - - batched_output = naive_batched_moe( - a, - w1, - w2, - topk_weight, - topk_ids, - w1_scale=w1_s, - w2_scale=w2_s, - a1_scale=a1_scale, - a2_scale=a2_scale, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - ) - - pplx_outputs = pplx_moe( - group_name, - rank, - world_size, - dp_size, - a, - w1, - w2, - topk_weight, - topk_ids, - w1_scale=w1_s, - w2_scale=w2_s, - a1_scale=a1_scale, - a2_scale=a2_scale, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - shared_experts=shared_experts, - ) - - if shared_experts is None: - pplx_shared_output = None - pplx_output = pplx_outputs - assert isinstance(pplx_output, torch.Tensor) - else: - pplx_shared_output, pplx_output = pplx_outputs - - if shared_output is not None: - assert pplx_shared_output is not None - chunked_shared_output = chunk_by_rank( - shared_output, pgi.rank, pgi.world_size - ).to(pplx_shared_output.device) - else: - chunked_shared_output = None - - chunked_batch_output = chunk_by_rank( - batched_output, pgi.rank, pgi.world_size - ).to(pplx_output.device) - - torch.testing.assert_close(batched_output, torch_output, atol=3e-2, rtol=3e-2) - - torch.testing.assert_close( - pplx_output, chunked_batch_output, atol=3e-2, rtol=3e-2 - ) - - if shared_experts is not None: - assert chunked_shared_output is not None - assert pplx_shared_output is not None - torch.testing.assert_close( - pplx_shared_output, chunked_shared_output, atol=3e-2, rtol=3e-2 - ) - - finally: - if use_internode: - nvshmem_finalize() - - -@pytest.mark.parametrize("mnk", PPLX_COMBOS) -@pytest.mark.parametrize("e", NUM_EXPERTS) -@pytest.mark.parametrize("topk", TOP_KS) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("per_act_token_quant", [False, True]) -@pytest.mark.parametrize("block_shape", [None, [128, 128]]) -@pytest.mark.parametrize("use_internode", [False]) -@pytest.mark.optional -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_moe_slow( - mnk: tuple[int, int, int], - e: int, - topk: int, - dtype: torch.dtype, - world_dp_size: tuple[int, int], - per_act_token_quant: bool, - block_shape: list[int] | None, - use_internode: bool, -): - set_random_seed(7) - m, n, k = mnk - world_size, dp_size = world_dp_size - - if dtype == torch.float8_e4m3fn: - use_fp8_w8a8 = True - quant_dtype = dtype - else: - use_fp8_w8a8 = False - quant_dtype = None - - if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None): - pytest.skip("Skip quantization test for non-quantized type") - - if per_act_token_quant and block_shape is not None: - pytest.skip("Skip illegal quantization combination") - - a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 - score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) - - (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( - e, - n, - k, - quant_dtype=quant_dtype, - block_shape=block_shape, - per_out_ch_quant=per_act_token_quant, - ) - - parallel_launch( - world_size, - _pplx_moe, - dp_size, - a, - w1, - w2, - score, - topk, - e, - w1_s, - w2_s, - quant_dtype, - per_act_token_quant, - block_shape, - use_internode, - ) - - -def _pplx_test_loop( - pgi: ProcessGroupInfo, - dp_size: int, - use_internode: bool, - use_shared_experts: bool, - make_weights: bool, - test_fn: Callable, -): - device = torch.device(f"cuda:{pgi.local_rank}") - init_workspace_manager(device) - - def format_result(msg, ex=None): - if ex is not None: - x = str(ex) - newx = x.strip(" \n\t")[:16] - if len(newx) < len(x): - newx = newx + " ..." - - prefix = "E\t" - print(f"{textwrap.indent(traceback.format_exc(), prefix)}") - print(f"FAILED {msg} - {newx}\n") - else: - print(f"PASSED {msg}") - - if use_shared_experts: - # Note: this config is only needed for the non-naive shared experts. - new_vllm_config = copy.deepcopy(vllm_config) - new_vllm_config.parallel_config.data_parallel_size = pgi.world_size - new_vllm_config.parallel_config.enable_expert_parallel = True - _set_vllm_config(new_vllm_config, pgi.world_size, pgi.rank, pgi.local_rank) - - set_random_seed(7) - combos = itertools.product( - PPLX_COMBOS, NUM_EXPERTS, TOP_KS, DTYPES, [False, True], [None, [128, 128]] - ) - exceptions = [] - count = 0 - for mnk, e, topk, dtype, per_act_token_quant, block_shape in combos: - count = count + 1 - m, n, k = mnk - - if dtype == torch.float8_e4m3fn: - use_fp8_w8a8 = True - quant_dtype = dtype - else: - use_fp8_w8a8 = False - quant_dtype = None - - test_desc = ( - f"test_pplx_moe[mnk={mnk}, e={e}, topk={topk}, " - f"dtype={dtype}, per_act_token={per_act_token_quant}, " - f"block_shape={block_shape}, use_internode={use_internode}, " - f"use_shared_experts={use_shared_experts}" - ) - - if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None): - print(f"{test_desc} - Skip quantization test for non-quantized type.") - continue - - if per_act_token_quant and block_shape is not None: - print(f"{test_desc} - Skip illegal quantization combination.") - continue - - a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 - score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) - - args = dict() - if make_weights: - (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( - e, - n, - k, - quant_dtype=quant_dtype, - block_shape=block_shape, - per_out_ch_quant=per_act_token_quant, - ) - args["w1"] = w1 - args["w2"] = w2 - args["w1_s"] = w1_s - args["w2_s"] = w2_s - - if use_shared_experts: - args["shared_experts"] = make_shared_experts( - n, - k, - in_dtype=a.dtype, - quant_dtype=quant_dtype, - ) - - try: - test_fn( - pgi=pgi, - dp_size=dp_size, - a=a, - score=score, - topk=topk, - num_experts=e, - quant_dtype=quant_dtype, - per_act_token_quant=per_act_token_quant, - block_shape=block_shape, - use_internode=use_internode, - **args, - ) - format_result(test_desc) - except Exception as ex: - format_result(test_desc, ex) - exceptions.append(ex) - - if len(exceptions) > 0: - raise RuntimeError( - f"{len(exceptions)} of {count} tests failed in child process, " - f"rank={pgi.rank}." - ) - else: - print(f"{count} of {count} tests passed in child process, rank={pgi.rank}.") - - -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("use_internode", [False]) -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_prepare_finalize( - world_dp_size: tuple[int, int], - use_internode: bool, -): - set_random_seed(7) - world_size, dp_size = world_dp_size - parallel_launch( - world_size * dp_size, - _pplx_test_loop, - dp_size, - use_internode, - False, - False, - _pplx_prepare_finalize, - ) - - -@pytest.mark.parametrize("world_dp_size", [[2, 1]]) -@pytest.mark.parametrize("use_internode", [False]) -@pytest.mark.parametrize("use_shared_experts", [False, True]) -@requires_pplx -@multi_gpu_test(num_gpus=2) -def test_pplx_moe( - world_dp_size: tuple[int, int], - use_internode: bool, - use_shared_experts: bool, -): - set_random_seed(7) - world_size, dp_size = world_dp_size - parallel_launch( - world_size, - _pplx_test_loop, - dp_size, - use_internode, - use_shared_experts, - True, - _pplx_moe, - ) diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index 2c54267ef90..936516576ce 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -37,13 +37,15 @@ vllm_config = VllmConfig() # Test configurations DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32] +# Quantization test configs NUM_TOKENS = [7, 2050] D = [512, 4096, 5120, 13824] GROUP_SIZE = [64, 128, 512] COLUMN_MAJOR_SCALES = [True, False] TMA_ALIGNED_SCALES = [True, False] -M = [1, 7, 8, 83, 84, 4096] -N = [128, 512, 7168, 7748, 13824] +# Matmul test configs +M = [1, 7, 8, 83, 4096] +N = [128, 512, 576, 7168, 13824] K = [256, 3884, 4096, 13824, 16384] # Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8 # and its hidden size is 7168. @@ -162,8 +164,6 @@ def test_w8a8_block_fp8_cutlass_matmul(): k_tiles = (K + block_k - 1) // block_k Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale - # Hopper requires row-major format for scales - Bs_cutlass = Bs.T.contiguous() if current_platform.is_device_capability(90) else Bs A_fp8, As = per_token_group_quant_fp8( A_fp32, block_size[1], column_major_scales=False @@ -174,9 +174,7 @@ def test_w8a8_block_fp8_cutlass_matmul(): ) ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) - out = cutlass_scaled_mm( - A_fp8_cutlass, B_fp8, As_cutlass, Bs_cutlass, block_size, out_dtype - ) + out = cutlass_scaled_mm(A_fp8_cutlass, B_fp8, As_cutlass, Bs, block_size, out_dtype) rel_diff = torch.mean( torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) diff --git a/tests/lora/test_moe_lora_align_sum.py b/tests/lora/test_moe_lora_align_sum.py index 3a17f3eba6e..bb46b4d8680 100644 --- a/tests/lora/test_moe_lora_align_sum.py +++ b/tests/lora/test_moe_lora_align_sum.py @@ -47,6 +47,8 @@ def test_moe_lora_align_block_size( # compute paddings max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) + if topk_ids.numel() < num_experts: + max_num_tokens_padded = topk_ids.numel() * block_size max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size) # init output tensors diff --git a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py index df5b077ce9d..5001b98b6d2 100644 --- a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py +++ b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py @@ -116,6 +116,32 @@ class TestCheckInterleavedAudioVideo: is_video, is_audio, is_video.sum().item(), is_audio.sum().item() ) + def test_batched_non_interleaved_no_false_positive(self): + """ + Regression test for https://github.com/vllm-project/vllm/issues/35394. + + 5 identical non-interleaved mixed-modality requests batched together: + each has [audio][image][video] in separate blocks with text between them. + Across the batch, audio from request N falls between video blocks of + request N and request N+1, causing the global ranges to overlap. + check_interleaved_audio_video must return False (not a false positive). + """ + # Build one request: [text][audio*5][text][image*4][text][video*6][text] + single_ids, _ = make_token_seq(5, 4, 6) + # Batch 5 identical requests (separated by text tokens to simulate padding) + sep = torch.tensor([TEXT_TOKEN_ID] * 3) + batched_ids = torch.cat([single_ids, sep] * 5) + is_multimodal = ( + (batched_ids == AUDIO_TOKEN_ID) + | (batched_ids == IMAGE_TOKEN_ID) + | (batched_ids == VIDEO_TOKEN_ID) + ) + is_video = is_multimodal & (batched_ids == VIDEO_TOKEN_ID) + is_audio = is_multimodal & (batched_ids == AUDIO_TOKEN_ID) + assert not check_interleaved_audio_video( + is_video, is_audio, is_video.sum().item(), is_audio.sum().item() + ), "Batched non-interleaved requests should not be detected as interleaved" + # --------------------------------------------------------------------------- # Tests for embed_input_ids via a minimal mock diff --git a/tests/models/registry.py b/tests/models/registry.py index 0978c93dac5..c8e47ad502f 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -194,6 +194,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "ArcticForCausalLM": _HfExamplesInfo( "Snowflake/snowflake-arctic-instruct", trust_remote_code=True ), + "AXK1ForCausalLM": _HfExamplesInfo("skt/A.X-K1", trust_remote_code=True), "BaiChuanForCausalLM": _HfExamplesInfo( "baichuan-inc/Baichuan-7B", trust_remote_code=True ), diff --git a/tests/reasoning/test_qwen3_reasoning_parser.py b/tests/reasoning/test_qwen3_reasoning_parser.py index db2bc16ffbb..411c7ba485a 100644 --- a/tests/reasoning/test_qwen3_reasoning_parser.py +++ b/tests/reasoning/test_qwen3_reasoning_parser.py @@ -9,6 +9,7 @@ from tests.reasoning.utils import ( run_reasoning_extraction, run_reasoning_extraction_streaming, ) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.reasoning import ReasoningParser, ReasoningParserManager parser_name = "qwen3" @@ -58,12 +59,14 @@ WITH_THINK_STREAM = { "content": "This is the rest", } -# --- No think tokens at all (thinking disabled) --- +# --- No think tokens at all (thinking enabled, truncated) --- +# With thinking enabled (default), no think tokens means the output was +# truncated before could be generated. All output is reasoning. WITHOUT_THINK = { "output": "This is the rest", - "reasoning": None, - "content": "This is the rest", + "reasoning": "This is the rest", + "content": None, } # In streaming, the parser cannot distinguish "thinking disabled" from # "reasoning in progress" when no think tokens have appeared yet. @@ -87,10 +90,12 @@ MULTILINE_REASONING = { "reasoning": "This is a reasoning\nsection", "content": "This is the rest\nThat", } +# Truncated output: present but no (thinking enabled). +# Everything is reasoning because the output was cut off mid-thought. ONLY_OPEN_TAG = { "output": "This is a reasoning section", - "reasoning": None, - "content": "This is a reasoning section", + "reasoning": "This is a reasoning section", + "content": None, } ONLY_OPEN_TAG_STREAM = { @@ -99,6 +104,20 @@ ONLY_OPEN_TAG_STREAM = { "content": None, } +# Truncated output without prefix (Qwen3.5 style where +# is in the prompt). No means truncation — all is reasoning. +TRUNCATED_NO_START_TOKEN = { + "output": "This is a reasoning section", + "reasoning": "This is a reasoning section", + "content": None, +} + +TRUNCATED_NO_START_TOKEN_STREAM = { + "output": "This is a reasoning section", + "reasoning": "This is a reasoning section", + "content": None, +} + TEST_CASES = [ pytest.param( False, @@ -170,6 +189,16 @@ TEST_CASES = [ ONLY_OPEN_TAG_STREAM, id="only_open_tag_stream", ), + pytest.param( + False, + TRUNCATED_NO_START_TOKEN, + id="truncated_no_start_token", + ), + pytest.param( + True, + TRUNCATED_NO_START_TOKEN_STREAM, + id="truncated_no_start_token_stream", + ), ] @@ -249,3 +278,46 @@ def test_reasoning_streaming_multi_token_deltas( assert reconstructor.reasoning == expected_reasoning assert (reconstructor.other_content or None) == expected_content + + +# --- Tests for enable_thinking=False (thinking explicitly disabled) --- + + +THINKING_DISABLED_CASES = [ + pytest.param( + "This is plain content", + None, + "This is plain content", + id="thinking_disabled_plain_content", + ), + pytest.param( + "Some output without think tokens", + None, + "Some output without think tokens", + id="thinking_disabled_no_think_tokens", + ), +] + + +@pytest.mark.parametrize( + "output, expected_reasoning, expected_content", THINKING_DISABLED_CASES +) +def test_reasoning_thinking_disabled( + output: str, + expected_reasoning: str | None, + expected_content: str | None, + qwen3_tokenizer, +): + """When enable_thinking=False, output without is all content.""" + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( + qwen3_tokenizer, + chat_template_kwargs={"enable_thinking": False}, + ) + + reasoning, content = parser.extract_reasoning( + model_output=output, + request=ChatCompletionRequest(messages=[], model="test-model"), + ) + + assert reasoning == expected_reasoning + assert content == expected_content diff --git a/tests/v1/e2e/test_mamba_prefix_cache.py b/tests/v1/e2e/test_mamba_prefix_cache.py index 38cfdcdb3ec..5aa72ccb3b0 100644 --- a/tests/v1/e2e/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/test_mamba_prefix_cache.py @@ -325,6 +325,7 @@ def get_fake_process_mamba_fn( requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + copy_bufs: mamba_utils.MambaCopyBuffers, ): nonlocal copy_info copy_info = None @@ -337,6 +338,7 @@ def get_fake_process_mamba_fn( requests, forward_context, mamba_state_copy_funcs, + copy_bufs, ) if cur_step_action is not None: check_copy_info( @@ -355,6 +357,7 @@ def get_fake_process_mamba_fn( mamba_state_idx: dict[str, int], forward_context: dict[str, Any], mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + copy_bufs: mamba_utils.MambaCopyBuffers, ): nonlocal copy_info copy_info = None @@ -366,6 +369,7 @@ def get_fake_process_mamba_fn( mamba_state_idx, forward_context, mamba_state_copy_funcs, + copy_bufs, ) if cur_step_action is not None: check_copy_info( @@ -376,19 +380,15 @@ def get_fake_process_mamba_fn( ) return ret - def fake_copy_fn( - src_state_list: list[int], - dest_state_list: list[int], - num_elements_list: list[int], - ): + def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers): nonlocal copy_info assert copy_info is None + n = copy_bufs.offset + src_state_list = copy_bufs.src_ptrs.cpu[:n].tolist() + dest_state_list = copy_bufs.dst_ptrs.cpu[:n].tolist() + num_elements_list = copy_bufs.sizes.cpu[:n].tolist() copy_info = (src_state_list, dest_state_list, num_elements_list) - return original_copy_fn( - src_state_list, - dest_state_list, - num_elements_list, - ) + return original_copy_fn(copy_bufs) return fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index cb38aa70d3f..93e6822e639 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -38,7 +38,7 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.worker.gpu_input_batch import InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner -from vllm.v1.worker.utils import AttentionGroup +from vllm.v1.worker.utils import AttentionGroup, select_common_block_size BLOCK_SIZE = 16 NUM_BLOCKS = 10 @@ -209,7 +209,7 @@ def test_select_common_block_size_prefers_manager_block_size(): AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0), ] - selected_size = GPUModelRunner.select_common_block_size(128, attn_groups) + selected_size = select_common_block_size(128, attn_groups) assert selected_size == 128 @@ -221,7 +221,7 @@ def test_select_common_block_size_uses_largest_shared_int(): AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0), ] - selected_size = GPUModelRunner.select_common_block_size(256, attn_groups) + selected_size = select_common_block_size(256, attn_groups) assert selected_size == 64 @@ -234,7 +234,7 @@ def test_select_common_block_size_no_valid_option(): ] with pytest.raises(ValueError): - GPUModelRunner.select_common_block_size(48, attn_groups) + select_common_block_size(48, attn_groups) def test_update_states_new_request(model_runner, dist_init): diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 38eb250fb17..df3b7de9b4c 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -62,6 +62,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): {}, {}, (), + MagicMock(), ) assert mamba_state_idx == {"keep": 99} diff --git a/tools/ep_kernels/README.md b/tools/ep_kernels/README.md index ab0e358802b..b4eabe18ca1 100644 --- a/tools/ep_kernels/README.md +++ b/tools/ep_kernels/README.md @@ -4,7 +4,7 @@ Large-scale cluster-level expert parallel, as described in the [DeepSeek-V3 Tech Here we break down the requirements in 2 steps: -1. Build and install the Python libraries (both [pplx-kernels](https://github.com/ppl-ai/pplx-kernels) and [DeepEP](https://github.com/deepseek-ai/DeepEP)), including necessary dependencies like NVSHMEM. This step does not require any privileged access. Any user can do this. +1. Build and install the Python libraries ([DeepEP](https://github.com/deepseek-ai/DeepEP)), including necessary dependencies like NVSHMEM. This step does not require any privileged access. Any user can do this. 2. Configure NVIDIA driver to enable IBGDA. This step requires root access, and must be done on the host machine. Step 2 is necessary for multi-node deployment. diff --git a/tools/ep_kernels/elastic_ep/install_eep_libraries.sh b/tools/ep_kernels/elastic_ep/install_eep_libraries.sh index fe7b862159d..31519c28716 100755 --- a/tools/ep_kernels/elastic_ep/install_eep_libraries.sh +++ b/tools/ep_kernels/elastic_ep/install_eep_libraries.sh @@ -76,11 +76,4 @@ popd export CMAKE_PREFIX_PATH=$WORKSPACE/nvshmem_install:$CMAKE_PREFIX_PATH -# build and install pplx, require pytorch installed -pushd "$WORKSPACE" -git clone https://github.com/ppl-ai/pplx-kernels -cd pplx-kernels -# see https://github.com/pypa/pip/issues/9955#issuecomment-838065925 -# PIP_NO_BUILD_ISOLATION=0 disables build isolation -PIP_NO_BUILD_ISOLATION=0 TORCH_CUDA_ARCH_LIST=9.0a+PTX pip install . --no-deps -v diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 148cb6e18d8..3372dd10f4d 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -4,12 +4,10 @@ set -ex # usage: ./install_python_libraries.sh [options] # --workspace workspace directory (default: ./ep_kernels_workspace) # --mode "install" (default) or "wheel" -# --pplx-ref pplx-kernels commit hash # --deepep-ref DeepEP commit hash # --nvshmem-ver NVSHMEM version CUDA_HOME=${CUDA_HOME:-/usr/local/cuda} -PPLX_COMMIT_HASH=${PPLX_COMMIT_HASH:-"12cecfd"} DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"} NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"} # Default supports both CUDA 12 and 13 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace} @@ -35,14 +33,6 @@ while [[ $# -gt 0 ]]; do MODE="$2" shift 2 ;; - --pplx-ref) - if [[ -z "$2" || "$2" =~ ^- ]]; then - echo "Error: --pplx-ref requires an argument." >&2 - exit 1 - fi - PPLX_COMMIT_HASH="$2" - shift 2 - ;; --deepep-ref) if [[ -z "$2" || "$2" =~ ^- ]]; then echo "Error: --deepep-ref requires an argument." >&2 @@ -188,14 +178,6 @@ do_build() { popd } -# build pplx-kernels -do_build \ - "https://github.com/ppl-ai/pplx-kernels" \ - "pplx-kernels" \ - "setup.py" \ - "$PPLX_COMMIT_HASH" \ - "" - # build DeepEP do_build \ "https://github.com/deepseek-ai/DeepEP" \ diff --git a/tools/profiler/print_layerwise_table.py b/tools/profiler/print_layerwise_table.py index d7a24a59859..06a8c58537b 100644 --- a/tools/profiler/print_layerwise_table.py +++ b/tools/profiler/print_layerwise_table.py @@ -33,7 +33,10 @@ if __name__ == "__main__": "--json-trace", type=str, required=True, - help="json trace file output by examples/offline_inference/profiling.py", + help=( + "JSON trace file generated by scripts that use " + "vllm.profiler.layerwise_profile" + ), ) parser.add_argument( "--phase", diff --git a/tools/profiler/visualize_layerwise_profile.py b/tools/profiler/visualize_layerwise_profile.py index ed4bf0beb71..83b8b3a7520 100644 --- a/tools/profiler/visualize_layerwise_profile.py +++ b/tools/profiler/visualize_layerwise_profile.py @@ -564,8 +564,10 @@ if __name__ == "__main__": "--json-trace", type=str, required=True, - help="json trace file output by \ - examples/offline_inference/profiling.py", + help=( + "JSON trace file generated by scripts that use " + "vllm.profiler.layerwise_profile" + ), ) parser.add_argument( "--output-directory", type=str, required=False, help="Directory to output plots" diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index e48ba6c997e..69f080ae20d 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -988,7 +988,7 @@ def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor): return output_tensor -def get_cutlass_pplx_moe_mm_data( +def get_cutlass_batched_moe_mm_data( expert_offsets: torch.Tensor, problem_sizes1: torch.Tensor, problem_sizes2: torch.Tensor, @@ -1011,7 +1011,7 @@ def get_cutlass_pplx_moe_mm_data( multiplication in two grouped MMs used in the fused MoE operation. """ - return torch.ops._C.get_cutlass_pplx_moe_mm_data( + return torch.ops._C.get_cutlass_batched_moe_mm_data( expert_offsets, problem_sizes1, problem_sizes2, @@ -2190,6 +2190,23 @@ def moe_wna16_gemm( ) +def router_gemm_bf16_fp32(input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """bf16 x bf16 -> fp32 GEMM via cuBLAS. weight shape: (N, K).""" + return torch.ops._moe_C.router_gemm_bf16_fp32(input, weight) + + +if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "router_gemm_bf16_fp32"): + + @register_fake("_moe_C::router_gemm_bf16_fp32") + def router_gemm_bf16_fp32_fake( + input: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + return torch.empty( + input.shape[0], weight.shape[0], dtype=torch.float32, device=input.device + ) + + def dsv3_router_gemm( hidden_states: torch.Tensor, router_weight: torch.Tensor, diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py index 55a5a2e5df4..a9b64adcb3f 100644 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ b/vllm/compilation/passes/fusion/collective_fusion.py @@ -53,7 +53,7 @@ class GEMMReduceScatterPattern(BasePattern): gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter( mul, mm_weight, - "avg", + "sum", scatter_dim=0, group_name=self.tp.device_group.group_name, ) @@ -150,7 +150,7 @@ class ScaledMMReduceScatterPattern(BasePattern): mat2, scale_a, scale_b, - "avg", + "sum", scatter_dim, # orig_scatter_dim scatter_dim, # scatter_dim_after_maybe_reshape self.tp.device_group.group_name, @@ -285,7 +285,7 @@ class CutlassScaledMMReduceScatterPattern(BasePattern): mat2, scale_a, scale_b, - "avg", + "sum", scatter_dim, # orig_scatter_dim scatter_dim, # scatter_dim_after_maybe_reshape self.tp.device_group.group_name, diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index d22e9a96e0f..01dc61cdcad 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1045,7 +1045,7 @@ class CompilationConfig: "are optimized for prefill and are incompatible with CUDA Graphs. " "In order to use CUDA Graphs for decode-optimized workloads, " "use --all2all-backend with another option, such as " - "deepep_low_latency, pplx, or allgather_reducescatter." + "deepep_low_latency or allgather_reducescatter." ) self.cudagraph_mode = CUDAGraphMode.NONE diff --git a/vllm/config/model.py b/vllm/config/model.py index 5fb81ee424e..012b2b1c9eb 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -883,6 +883,7 @@ class ModelConfig: "modelopt", "modelopt_fp4", "modelopt_mxfp8", + "modelopt_mixed", "petit_nvfp4", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index cc2cfa97b50..fa4f72dccca 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -152,7 +152,6 @@ class ParallelConfig: - "naive": Naive all2all implementation using broadcasts\n - "allgather_reducescatter": All2all based on allgather and reducescatter\n - - "pplx": Use pplx kernels\n - "deepep_high_throughput": Use deepep high-throughput kernels\n - "deepep_low_latency": Use deepep low-latency kernels\n - "mori": Use mori kernels\n @@ -310,6 +309,13 @@ class ParallelConfig: f"but found: {self._api_process_rank}" ) + if self.all2all_backend == "pplx": + logger.warning( + "The 'pplx' all2all backend has been removed. " + "Falling back to 'allgather_reducescatter'." + ) + self.all2all_backend = "allgather_reducescatter" + if self.data_parallel_size_local > self.data_parallel_size: raise ValueError( f"data_parallel_size_local ({self.data_parallel_size_local}) " @@ -442,7 +448,6 @@ class ParallelConfig: # In this case, ensure the input to the experts is sequence parallel # to avoid the excess work. # - # Not needed for pplx-kernels as it can handle duplicate input tokens. @property def use_sequence_parallel_moe(self) -> bool: return ( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 127c16ac769..7f7b21316c7 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -126,6 +126,9 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool: # tp-dp combination broken: # https://github.com/vllm-project/vllm/issues/34458 and cfg.parallel_config.data_parallel_size == 1 + # tp-pp combination broken: + # https://github.com/vllm-project/vllm/issues/35426 + and cfg.parallel_config.pipeline_parallel_size == 1 ) @@ -857,7 +860,7 @@ class VllmConfig: self.compilation_config.pass_config.fuse_gemm_comms = False else: # Compute SP threshold early; disable if None (model too - # small) before +rms_norm gets forced into custom_ops. + # small for SP to be beneficial). pass_config = self.compilation_config.pass_config if pass_config.sp_min_token_num is None: from vllm.compilation.passes.fusion.sequence_parallelism import ( @@ -880,14 +883,6 @@ class VllmConfig: self.compilation_config.pass_config.enable_sp = False self.compilation_config.pass_config.fuse_gemm_comms = False - if self.compilation_config.pass_config.enable_sp: - if "-rms_norm" in self.compilation_config.custom_ops: - logger.warning( - "RMS norm force disabled, sequence parallelism might break" - ) - else: - self.compilation_config.custom_ops.append("+rms_norm") - if self.compilation_config.fast_moe_cold_start is None: # resolve default behavior: try to be as safe as possible # this config is unsafe if any spec decoding draft model has a MOE. diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 678cd45800d..4acab4e3c88 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -3,14 +3,13 @@ from typing import Any import torch -import torch.distributed as dist import vllm.envs as envs from vllm.distributed import get_dp_group, get_ep_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.utils.flashinfer import has_flashinfer_all2all -from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_mori from .base_device_communicator import All2AllManagerBase, Cache @@ -235,96 +234,6 @@ class AgRsAll2AllManager(All2AllManagerBase): pass -class PPLXAll2AllManager(All2AllManagerBase): - """ - All2All communication based on PPLX kernels. - """ - - def __init__(self, cpu_group): - assert has_pplx(), ( - "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md" - " to install pplx_kernels." - ) - super().__init__(cpu_group) - - if self.internode: - # inter-node communication needs nvshmem, - # intra-node communication uses p2p mapping directly - from pplx_kernels.nvshmem import ( # type: ignore[import-not-found] - nvshmem_alloc_empty_unique_id, - nvshmem_get_unique_id, - nvshmem_init, - ) - - logger.debug( - "Initialize NVSHMEM for pplx_kernels: rank=%d, world size=%d", - self.rank, - self.world_size, - ) - uid = ( - nvshmem_get_unique_id() - if self.rank == 0 - else nvshmem_alloc_empty_unique_id() - ) - dist.broadcast( - uid, - src=dist.get_process_group_ranks(self.cpu_group)[0], - group=self.cpu_group, - ) - logger.debug("PPLX NVSHMEM UID = %s", uid) - nvshmem_init(uid, self.rank, self.world_size) - - self.handle_cache = Cache() - - def get_handle(self, kwargs): - import pplx_kernels as pplx # type: ignore[import-not-found] - - return self.handle_cache.get_or_create( - kwargs, - pplx.AllToAll.internode if self.internode else pplx.AllToAll.intranode, - ) - - def dispatch_router_logits( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - is_sequence_parallel: bool = False, - extra_tensors: list[torch.Tensor] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - raise NotImplementedError - - def dispatch( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - is_sequence_parallel: bool = False, - extra_tensors: list[torch.Tensor] | None = None, - ) -> ( - tuple[torch.Tensor, torch.Tensor, torch.Tensor] - | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] - ): - raise NotImplementedError - - def combine( - self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False - ) -> torch.Tensor: - raise NotImplementedError - - def destroy(self): - with self.handle_cache._lock: - for _, handle in self.handle_cache._cache.items(): - handle.destroy() - - if self.internode: - from pplx_kernels.nvshmem import ( - nvshmem_finalize, # type: ignore[import-not-found] - ) - - logger.debug("PPLX NVSHMEM finalize") - nvshmem_finalize() - - class DeepEPAll2AllManagerBase(All2AllManagerBase): """ All2All communication based on DeepEP High-Throughput kernels. diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index ff2d7436b27..3c347ef756d 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -27,6 +27,7 @@ from vllm.utils.torch_utils import cuda_device_count_stateless logger = init_logger(__name__) +KiB = 1024 MiB = 1024 * 1024 # Max size for each world size in case symmetric memory is available # For different SM architectures @@ -60,17 +61,44 @@ SYMM_MEM_ALL_REDUCE_MAX_SIZES = { }, } +# NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks. +# PyNCCL-symm outperforms custom_AR for small and large tensor sizes, +# while custom_AR wins for mid-range sizes. +# +# Benchmark results (8 GPUs): +# 2K - 16K: PyNCCL-symm wins (1.35x - 1.48x faster) +# 32K - 64K: custom_AR wins +# 128K - 1G: PyNCCL-symm wins (1.12x - 6.14x faster) +# +# Benchmark results (4 GPUs): +# 2K - 16K: PyNCCL-symm wins (1.21x - 1.30x faster) +# 32K - 256K: custom_AR wins (1.07x - 1.35x faster) +# 512K - 1G: PyNCCL-symm wins (1.10x - 2.32x faster) +# +# The config defines ranges where custom_AR is preferred (symm_mem disabled). NCCL_SYMM_MEM_ALL_REDUCE_CONFIG: dict[str, Any] = { "min_world_size": 4, - "thresholds": { - 4: 2 * MiB, # 2 MB - 8: 1 * MiB, # 1 MB + # Ranges where custom_AR outperforms NCCL symm_mem: (lower_bound, upper_bound) + # NCCL symm_mem will NOT be used for sizes in range: lower < size < upper + "custom_ar_preferred_ranges": { + 4: (16 * KiB, 512 * KiB), # custom_AR wins for 32K-256K + 8: (16 * KiB, 128 * KiB), # custom_AR wins for 32K-64K }, "always_use_above_world_size": 8, # Always use symm mem for world_size > 8 } def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) -> bool: + """ + Determine if NCCL symmetric memory allreduce should be used. + + Based on H100 and GB200 benchmarks, NCCL symm_mem is preferred for: + - Small tensors (≤16K): Lower latency than custom_AR + - Large tensors (≥128K for 8 GPUs, ≥512K for 4 GPUs): Better bandwidth + + Custom_AR is preferred for mid-range sizes where its P2P approach + has lower overhead than the symm_mem copy-in/copy-out pattern. + """ from vllm.distributed.device_communicators.pynccl_allocator import ( is_symmetric_memory_enabled, ) @@ -80,11 +108,20 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) if not is_symmetric_memory_enabled(): return False + if world_size < NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["min_world_size"]: return False - threshold = NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["thresholds"].get(world_size) - if threshold is not None and input_tensor.nbytes >= threshold: - return True + + tensor_size = input_tensor.nbytes + custom_ar_range = NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"].get( + world_size + ) + + if custom_ar_range is not None: + lower_bound, upper_bound = custom_ar_range + # Use symm_mem for small sizes (≤ lower_bound) and large sizes (≥ upper_bound) + # Use custom_AR (not symm_mem) for mid-range sizes + return tensor_size <= lower_bound or tensor_size >= upper_bound return world_size > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"] diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 62e2b90377f..dd571482f5c 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -112,10 +112,6 @@ class CudaCommunicator(DeviceCommunicatorBase): from .all2all import AgRsAll2AllManager self.all2all_manager = AgRsAll2AllManager(self.cpu_group) - elif self.all2all_backend == "pplx": - from .all2all import PPLXAll2AllManager - - self.all2all_manager = PPLXAll2AllManager(self.cpu_group) elif self.all2all_backend == "deepep_high_throughput": from .all2all import DeepEPHTAll2AllManager @@ -298,7 +294,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.fi_ar_comm = None if self.all2all_manager is not None: self.all2all_manager.destroy() - self.all2all_manager = None + self.all2all_manager = None # type: ignore[assignment] def all_gatherv( self, diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 7c3701b4ea9..891f19cfe25 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -159,7 +159,7 @@ class EplbModelState: NOTE: The expert_load_view now records load for all physical experts rather than just local experts. This ensures consistent load statistics - across different dispatch methods (naive all-to-all, DeepEP, pplx-kernels). + across different dispatch methods (naive all-to-all, DeepEP). The recorded load will be multiplied by dp_size when using naive all-to-all due to each DP rank contributing the same token set to the calculation. See: diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 226dd6c1a2b..02e6e0d036a 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -259,7 +259,7 @@ class CompletionRequest(OpenAIBaseModel): structured_outputs_kwargs["json"] = json_schema.json_schema elif response_format.type == "structural_tag": structural_tag = response_format - assert structural_tag is not None and isinstance( + assert isinstance( structural_tag, ( LegacyStructuralTagResponseFormat, @@ -313,6 +313,34 @@ class CompletionRequest(OpenAIBaseModel): skip_clone=True, # Created fresh per request, safe to skip clone ) + @model_validator(mode="before") + @classmethod + def validate_response_format(cls, data): + response_format = data.get("response_format") + if response_format is None: + return data + + rf_type = ( + response_format.get("type") + if isinstance(response_format, dict) + else getattr(response_format, "type", None) + ) + + if rf_type == "json_schema": + json_schema = ( + response_format.get("json_schema") + if isinstance(response_format, dict) + else getattr(response_format, "json_schema", None) + ) + if json_schema is None: + raise VLLMValidationError( + "When response_format type is 'json_schema', the " + "'json_schema' field must be provided.", + parameter="response_format", + ) + + return data + @model_validator(mode="before") @classmethod def check_structured_outputs_count(cls, data): diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index b9d526e25de..3cfb6fffc3e 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -85,6 +85,8 @@ from vllm.entrypoints.openai.responses.protocol import ( ResponseCreatedEvent, ResponseInProgressEvent, ResponseInputOutputMessage, + ResponseReasoningPartAddedEvent, + ResponseReasoningPartDoneEvent, ResponsesRequest, ResponsesResponse, ResponseUsage, @@ -1339,6 +1341,19 @@ class OpenAIServingResponses(OpenAIServing): ), ) ) + yield _increment_sequence_number_and_return( + ResponseReasoningPartAddedEvent( + type="response.reasoning_part.added", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + content_index=current_content_index, + part=ResponseReasoningTextContent( + text="", + type="reasoning_text", + ), + ) + ) else: yield _increment_sequence_number_and_return( ResponseOutputItemAddedEvent( @@ -1354,22 +1369,21 @@ class OpenAIServingResponses(OpenAIServing): ), ) ) - yield _increment_sequence_number_and_return( - ResponseContentPartAddedEvent( - type="response.content_part.added", - sequence_number=-1, - output_index=current_output_index, - item_id=current_item_id, - content_index=current_content_index, - part=ResponseOutputText( - type="output_text", - text="", - annotations=[], - logprobs=[], - ), + yield _increment_sequence_number_and_return( + ResponseContentPartAddedEvent( + type="response.content_part.added", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + content_index=current_content_index, + part=ResponseOutputText( + type="output_text", + text="", + annotations=[], + logprobs=[], + ), + ) ) - ) - current_content_index += 1 first_delta_sent = True # todo(kebe7jun) tool call support @@ -1397,6 +1411,19 @@ class OpenAIServingResponses(OpenAIServing): text=reason_content, ) ) + yield _increment_sequence_number_and_return( + ResponseReasoningPartDoneEvent( + type="response.reasoning_part.done", + sequence_number=-1, + item_id=current_item_id, + output_index=current_output_index, + content_index=current_content_index, + part=ResponseReasoningTextContent( + text=reason_content, + type="reasoning_text", + ), + ) + ) current_content_index = 0 reasoning_item = ResponseReasoningItem( type="reasoning", @@ -1418,6 +1445,8 @@ class OpenAIServingResponses(OpenAIServing): item=reasoning_item, ) ) + current_output_index += 1 + current_item_id = str(uuid.uuid4()) yield _increment_sequence_number_and_return( ResponseOutputItemAddedEvent( type="response.output_item.added", @@ -1432,8 +1461,6 @@ class OpenAIServingResponses(OpenAIServing): ), ) ) - current_output_index += 1 - current_item_id = str(uuid.uuid4()) yield _increment_sequence_number_and_return( ResponseContentPartAddedEvent( type="response.content_part.added", @@ -1449,7 +1476,6 @@ class OpenAIServingResponses(OpenAIServing): ), ) ) - current_content_index += 1 # reset previous delta messages previous_delta_messages = [] @@ -1485,7 +1511,6 @@ class OpenAIServingResponses(OpenAIServing): ), ) ) - current_content_index += 1 previous_delta_messages.append(delta_message) if previous_delta_messages: @@ -1505,7 +1530,19 @@ class OpenAIServingResponses(OpenAIServing): text=reason_content, ) ) - current_content_index += 1 + yield _increment_sequence_number_and_return( + ResponseReasoningPartDoneEvent( + type="response.reasoning_part.done", + sequence_number=-1, + item_id=current_item_id, + output_index=current_output_index, + content_index=current_content_index, + part=ResponseReasoningTextContent( + text=reason_content, + type="reasoning_text", + ), + ) + ) reasoning_item = ResponseReasoningItem( type="reasoning", content=[ @@ -1543,7 +1580,6 @@ class OpenAIServingResponses(OpenAIServing): item_id=current_item_id, ) ) - current_content_index += 1 part = ResponseOutputText( text=final_content, type="output_text", @@ -1559,7 +1595,6 @@ class OpenAIServingResponses(OpenAIServing): part=part, ) ) - current_content_index += 1 item = ResponseOutputMessage( type="message", role="assistant", diff --git a/vllm/envs.py b/vllm/envs.py index d560cfc7753..07d9f81eac9 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1752,10 +1752,7 @@ def compile_factors() -> dict[str, object]: "VLLM_ENABLE_V1_MULTIPROCESSING", "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE", "VLLM_CPU_KVCACHE_SPACE", - "VLLM_CPU_OMP_THREADS_BIND", - "VLLM_CPU_NUM_OF_RESERVED_CPU", "VLLM_CPU_MOE_PREPACK", - "VLLM_CPU_SGL_KERNEL", "VLLM_TEST_FORCE_LOAD_FORMAT", "VLLM_ENABLE_CUDA_COMPATIBILITY", "VLLM_CUDA_COMPATIBILITY_PATH", diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py index 3114631ddec..cd0ef83fc0a 100644 --- a/vllm/kernels/helion/register.py +++ b/vllm/kernels/helion/register.py @@ -31,8 +31,8 @@ by key matches the config returned by the autotuner. Key Classes ----------- -- HelionKernelWrapper: Wraps raw kernel + config_picker, creates configured ops -- ConfiguredHelionKernel: Platform-specific kernel registered as PyTorch custom op +- HelionKernelWrapper: Wraps raw kernel + config_picker, creates configured kernels +- ConfiguredHelionKernel: Platform-specific kernel with pre-tuned configs - PresetConfigSearch: Custom autotuner that returns pre-tuned configs """ @@ -53,10 +53,27 @@ if not has_helion(): ) import helion +from helion._compat import requires_torch_version from helion.autotuner.base_search import BaseAutotuner from helion.runtime.config import Config from helion.runtime.settings import default_autotuner_fn +# TODO(gmagogsfm): Remove CustomOp fallback path (_get_or_register_custom_op, +# vllm_helion_lib, direct_register_custom_op) once vLLM requires PyTorch >= 2.11. +_HOP_AVAILABLE = requires_torch_version("2.11") + +if _HOP_AVAILABLE: + import torch.utils._pytree as pytree + from helion._compiler._dynamo.higher_order_ops import ( + helion_kernel_side_table, + helion_kernel_wrapper_mutation, + ) + from helion._compiler._dynamo.variables import infer_output_spec + from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + get_proxy_mode, + ) + logger = init_logger(__name__) vllm_helion_lib = Library("vllm_helion", "FRAGMENT") # noqa @@ -233,7 +250,7 @@ class ConfiguredHelionKernel: class HelionKernelWrapper: - """Wrapper for Helion kernels that creates config-specific PyTorch custom ops.""" + """Wrapper for Helion kernels with pre-tuned config selection and HOP support.""" def __init__( self, @@ -252,11 +269,86 @@ class HelionKernelWrapper: self._config_picker: ( Callable[[tuple[Any, ...], list[str]], str | None] | None ) = None + self._configured_kernel: ConfiguredHelionKernel | None = None self._input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None def __call__(self, *args, **kwargs): - configured_op = self.get_configured_op() - return configured_op(*args, **kwargs) + # CustomOp fallback: register as torch custom op for torch.compile + # compatibility on older PyTorch lacking HOP/EffectType support + if not _HOP_AVAILABLE: + custom_op = self._get_or_register_custom_op() + return custom_op(*args, **kwargs) + # HOP tracing: record HigherOrderOp in the FX graph + if get_proxy_mode() is not None: + return self._call_via_hop(args, kwargs) + # Eager: run the configured kernel directly + return self.get_configured_op()(*args, **kwargs) + + def _call_via_hop( + self, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + kernel = self.get_configured_op()._decorated_kernel + kernel_idx = helion_kernel_side_table.add_kernel(kernel) + + constant_args, tensor_args = self._partition_args(kernel, args, kwargs) + + all_named = {**constant_args, **tensor_args} + full_args = tuple( + all_named.get(n, p.default) + for n, p in kernel.signature.parameters.items() # type: ignore[attr-defined] + if n in all_named or p.default is not p.empty + ) + + with disable_proxy_modes_tracing(): + output_spec = infer_output_spec(kernel, full_args) + + hop_result = helion_kernel_wrapper_mutation( + kernel_idx=kernel_idx, + constant_args=constant_args, + tensor_args=tensor_args, + output_spec=output_spec, + ) + + tree_spec_str = output_spec.get("tree_spec_str") + if tree_spec_str is None: + return None + tree_spec = pytree.treespec_loads(tree_spec_str) + + hop_iter = iter(hop_result) + reconstructed = [] + for spec in output_spec["leaf_specs"]: + is_constant_scalar = spec["type"] == "scalar" and not isinstance( + spec.get("scalar_value"), torch.SymInt + ) + if is_constant_scalar: + reconstructed.append(spec["scalar_value"]) + else: + reconstructed.append(next(hop_iter)) + return pytree.tree_unflatten(reconstructed, tree_spec) + + @staticmethod + def _partition_args( + kernel: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> tuple[dict[str, Any], dict[str, Any]]: + constant_args: dict[str, Any] = {} + tensor_args: dict[str, Any] = {} + params = list(kernel.signature.parameters.keys()) + for i, val in enumerate(args): + name = params[i] + if isinstance(val, torch.Tensor): + tensor_args[name] = val + else: + constant_args[name] = val + for name, val in kwargs.items(): + if isinstance(val, torch.Tensor): + tensor_args[name] = val + else: + constant_args[name] = val + return constant_args, tensor_args def register_config_picker( self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None] @@ -309,29 +401,32 @@ class HelionKernelWrapper: ) return autotune_kernel.autotune(inputs) - def get_configured_op(self) -> Any: + def get_configured_op(self) -> ConfiguredHelionKernel: assert self._config_picker is not None, ( f"No config picker registered for kernel '{self.op_name}'. " f"Use @{self.op_name}.register_config_picker to register one." ) + if self._configured_kernel is None: + self._configured_kernel = ConfiguredHelionKernel( + op_name=self.op_name, + config_picker=self._config_picker, + raw_kernel_func=self.raw_kernel_func, + helion_settings=self.helion_settings, + ) + + return self._configured_kernel + + def _get_or_register_custom_op(self) -> Any: if hasattr(torch.ops.vllm_helion, self.op_name): - logger.debug("Op vllm_helion::%s already registered", self.op_name) return getattr(torch.ops.vllm_helion, self.op_name) - configured_kernel = ConfiguredHelionKernel( - op_name=self.op_name, - config_picker=self._config_picker, - raw_kernel_func=self.raw_kernel_func, - helion_settings=self.helion_settings, - ) + configured_kernel = self.get_configured_op() logger.info("Registering op: vllm_helion::%s", self.op_name) direct_register_custom_op( op_name=self.op_name, - op_func=configured_kernel._decorated_kernel, # Register decorated kernel - # TODO(gmagogsfm): Implement automatic mutation/aliasing detection - # for Helion kernels. + op_func=configured_kernel._decorated_kernel, mutates_args=None, fake_impl=self._fake_impl, target_lib=vllm_helion_lib, diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index b75d297ba5c..5f2604892ce 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -351,6 +351,8 @@ class PunicaWrapperGPU(PunicaWrapperBase): max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) if pad_sorted_ids: max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) + if topk_ids.numel() < num_experts: + max_num_tokens_padded = topk_ids.numel() * block_size sorted_ids = torch.empty( (max_loras * max_num_tokens_padded,), dtype=torch.int32, diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index e59806abb04..d89366bbdfe 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -2,21 +2,93 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import torch from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.models.vision import get_vit_attn_backend +from vllm.utils.math_utils import round_up from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.vit_attn_wrappers import ( vit_flash_attn_wrapper, + vit_flashinfer_wrapper, vit_torch_sdpa_wrapper, vit_triton_attn_wrapper, ) logger = init_logger(__name__) +# Batch buckets for cuDNN graph caching. +# Graphs use batch size and max sequence length as cache key. +# This avoids creating a new graph for each unique set of +# batch size and max sequence length at runtime. +# From the cuDNN team's performance measurements, there +# is no significant kernel performance difference between padding +# to a smaller batch size/seq length and padding to larger +# ones. The bucketing here is solely used to avoid memory +# operation overhead, which won't be needed if we have CUDA +# graph support in the future. +# TODO: Remove buckets after issue #34763 +# (cuda graph support) is addressed. +FLASHINFER_BATCH_BUCKETS = [8, 16, 32, 64] +FLASHINFER_MAX_SEQLEN_BUCKETS = [ + 1 * 1024, + 2 * 1024, + 4 * 1024, + 8 * 1024, + 16 * 1024, + 32 * 1024, + 64 * 1024, + 128 * 1024, +] + +# Workspace buffer for FlashInfer CuDNN backend +FLASHINFER_CUDNN_WORKSPACE_SIZE_BYTES = 128 * 1024 * 1024 +_flashinfer_workspace_buffer: torch.Tensor | None = None + + +def _get_flashinfer_workspace_buffer() -> torch.Tensor: + global _flashinfer_workspace_buffer + if _flashinfer_workspace_buffer is None: + _flashinfer_workspace_buffer = torch.zeros( + FLASHINFER_CUDNN_WORKSPACE_SIZE_BYTES, + dtype=torch.uint8, + device="cuda", + ) + return _flashinfer_workspace_buffer + + +def add_padding_to_seqlens( + seq: np.ndarray, + batch_size: int, + padding_value: int, +) -> np.ndarray: + batch_size_padded = next( + (b for b in FLASHINFER_BATCH_BUCKETS if b >= batch_size), + round_up(batch_size, FLASHINFER_BATCH_BUCKETS[0]), + ) + if batch_size_padded == batch_size: + return seq + return np.concatenate( + [ + seq, + np.full((batch_size_padded - batch_size,), padding_value, dtype=seq.dtype), + ] + ) + + +def bucket_flashinfer_max_seqlen( + real_max_seqlen: int, +) -> int: + if real_max_seqlen <= 0: + return FLASHINFER_MAX_SEQLEN_BUCKETS[0] + return next( + (s for s in FLASHINFER_MAX_SEQLEN_BUCKETS if s >= real_max_seqlen), + round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]), + ) + # --8<-- [start:mm_encoder_attn] @CustomOp.register("mm_encoder_attn") @@ -24,6 +96,67 @@ class MMEncoderAttention(CustomOp): """Multi-headed attention without any cache, used for multimodal encoder.""" # --8<-- [end:mm_encoder_attn] + @classmethod + def compute_max_seqlen( + cls, + attn_backend: AttentionBackendEnum, + cu_seqlens: np.ndarray, + ) -> int: + max_seqlen = 0 + if ( + attn_backend + in ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.ROCM_AITER_FA, + AttentionBackendEnum.TRITON_ATTN, + AttentionBackendEnum.FLASHINFER, + ) + and len(cu_seqlens) >= 2 + ): + max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max()) + if attn_backend == AttentionBackendEnum.FLASHINFER: + max_seqlen = bucket_flashinfer_max_seqlen(max_seqlen) + return max_seqlen + + @classmethod + def maybe_compute_sequence_lengths( + cls, + attn_backend: AttentionBackendEnum, + cu_seqlens: np.ndarray, + ) -> np.ndarray | None: + if attn_backend != AttentionBackendEnum.FLASHINFER: + return None + sequence_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + sequence_lengths = add_padding_to_seqlens( + sequence_lengths, len(sequence_lengths), 0 + ) + return sequence_lengths + + @classmethod + def maybe_recompute_cu_seqlens( + cls, + attn_backend: AttentionBackendEnum, + cu_seqlens: np.ndarray, + hidden_size: int, + tp_size: int, + ) -> np.ndarray: + if attn_backend != AttentionBackendEnum.FLASHINFER: + return cu_seqlens + + batch_size = len(cu_seqlens) - 1 + scale = hidden_size // tp_size + cu_seqlens = cu_seqlens * scale + + cu_seqlens_qko = cu_seqlens + cu_seqlens_v = cu_seqlens * 3 + + cu_seqlens_qko = add_padding_to_seqlens( + cu_seqlens_qko, batch_size, cu_seqlens_qko[-1] + ) + cu_seqlens_v = add_padding_to_seqlens( + cu_seqlens_v, batch_size, cu_seqlens_v[-1] + ) + return np.concatenate([cu_seqlens_qko, cu_seqlens_v]) def __init__( self, @@ -46,10 +179,9 @@ class MMEncoderAttention(CustomOp): self.num_heads = num_heads self.head_size = head_size - self.scale = scale + self.scale = 1.0 / (head_size**0.5) if scale is None else scale self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads self.layer_name = prefix - assert self.num_heads % self.num_kv_heads == 0, ( f"num_heads ({self.num_heads}) is not " f"divisible by num_kv_heads ({self.num_kv_heads})" @@ -75,6 +207,9 @@ class MMEncoderAttention(CustomOp): get_flash_attn_version() if self.is_flash_attn_backend else None ) + if self.attn_backend == AttentionBackendEnum.FLASHINFER: + _get_flashinfer_workspace_buffer() + logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.") @classmethod @@ -201,6 +336,27 @@ class MMEncoderAttention(CustomOp): output = output.reshape(bsz, q_len, -1) return output + def _forward_flashinfer( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, + sequence_lengths: torch.Tensor + | None = None, # Only used for FlashInfer CuDNN backend + ) -> torch.Tensor: + return vit_flashinfer_wrapper( + q=query, + k=key, + v=value, + scale=self.scale, + workspace_buffer=_get_flashinfer_workspace_buffer(), + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + def forward_native( self, query: torch.Tensor, @@ -208,6 +364,8 @@ class MMEncoderAttention(CustomOp): value: torch.Tensor, cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention + sequence_lengths: torch.Tensor + | None = None, # Only used for FlashInfer CuDNN backend ) -> torch.Tensor: return self._forward_sdpa(query, key, value, cu_seqlens) @@ -218,11 +376,17 @@ class MMEncoderAttention(CustomOp): value: torch.Tensor, cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention + sequence_lengths: torch.Tensor + | None = None, # Only used for FlashInfer CuDNN backend ) -> torch.Tensor: if self.is_flash_attn_backend: return self._forward_fa(query, key, value, cu_seqlens, max_seqlen) elif self.attn_backend == AttentionBackendEnum.TRITON_ATTN: return self._forward_triton(query, key, value, cu_seqlens, max_seqlen) + elif self.attn_backend == AttentionBackendEnum.FLASHINFER: + return self._forward_flashinfer( + query, key, value, cu_seqlens, max_seqlen, sequence_lengths + ) elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA: return self._forward_sdpa(query, key, value, cu_seqlens) else: @@ -238,6 +402,8 @@ class MMEncoderAttention(CustomOp): value: torch.Tensor, cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention + sequence_lengths: torch.Tensor + | None = None, # Only used for FlashInfer CuDNN backend ) -> torch.Tensor: return self._forward_sdpa(query, key, value, cu_seqlens) @@ -248,6 +414,8 @@ class MMEncoderAttention(CustomOp): value: torch.Tensor, cu_seqlens: torch.Tensor | None = None, max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention + sequence_lengths: torch.Tensor + | None = None, # Only used for FlashInfer CuDNN backend ) -> torch.Tensor: if self.attn_backend == AttentionBackendEnum.FLASH_ATTN: return self._forward_fa(query, key, value, cu_seqlens, max_seqlen) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index c6cb31b629a..be901bd2449 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import ( from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, ) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, @@ -64,6 +65,7 @@ __all__ = [ "FusedMoEPermuteExpertsUnpermute", "FusedMoEActivationFormat", "FusedMoEPrepareAndFinalize", + "GateLinear", "RoutingMethodType", "SharedFusedMoE", "ZeroExpertFusedMoE", diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index bf8ec2dc6f2..8c1bfe1c367 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any import torch @@ -24,16 +25,11 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import ( MoEPrepareAndFinalizeNoEP, ) from vllm.platforms import current_platform -from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx +from vllm.utils.import_utils import has_deep_ep, has_mori logger = init_logger(__name__) if current_platform.is_cuda_alike(): - if has_pplx(): - from .pplx_prepare_finalize import ( - PplxPrepareAndFinalize, - pplx_hidden_dim_scale_bytes, - ) if has_deep_ep(): from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize from .deepep_ll_prepare_finalize import ( @@ -120,51 +116,10 @@ def maybe_make_prepare_finalize( prepare_finalize: FusedMoEPrepareAndFinalize | None = None - if moe.use_pplx_kernels: - assert quant_config is not None - - hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes( - moe.max_num_tokens, - moe.hidden_dim, - moe.in_dtype, - quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - ) - - all_to_all_args = dict( - max_num_tokens=moe.max_num_tokens, - num_experts=moe.num_experts, - experts_per_token=moe.experts_per_token, # topk - rank=all2all_manager.rank, - world_size=all2all_manager.world_size, - # dp_size actually means tp_size, bug in pplx kernels - dp_size=all2all_manager.tp_group.world_size, - hidden_dim=moe.hidden_dim, - hidden_dim_bytes=hidden_dim_bytes, - hidden_dim_scale_bytes=hidden_scale_bytes, - ) - - num_dispatchers = ( - all2all_manager.world_size // all2all_manager.tp_group.world_size - ) - - # Intranode pplx a2a takes a group name while internode does not. - if not all2all_manager.internode: - all_to_all_args["group_name"] = all2all_manager.cpu_group.group_name - - handle = all2all_manager.get_handle(all_to_all_args) - - prepare_finalize = PplxPrepareAndFinalize( - handle, - max_num_tokens=moe.max_num_tokens, - num_local_experts=moe.num_local_experts, - num_dispatchers=num_dispatchers, - ) - elif moe.use_deepep_ht_kernels: + if moe.use_deepep_ht_kernels: assert moe.dp_size == all2all_manager.dp_world_size - all_to_all_args = dict() + all_to_all_args: dict[str, Any] = dict() handle = all2all_manager.get_handle(all_to_all_args) prepare_finalize = DeepEPHTPrepareAndFinalize( handle, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 87e1e244b3a..33d69b57a93 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -939,10 +939,6 @@ class FusedMoEParallelConfig: def use_all2all_kernels(self): return self.dp_size > 1 and self.use_ep - @property - def use_pplx_kernels(self): - return self.use_all2all_kernels and self.all2all_backend == "pplx" - @property def use_deepep_ht_kernels(self): return ( @@ -962,7 +958,7 @@ class FusedMoEParallelConfig: @property def use_batched_activation_format(self): - return self.use_deepep_ll_kernels or self.use_pplx_kernels + return self.use_deepep_ll_kernels @property def use_naive_all2all_kernels(self): @@ -1221,10 +1217,6 @@ class FusedMoEConfig: def use_ep(self): return self.moe_parallel_config.use_ep - @property - def use_pplx_kernels(self): - return self.moe_parallel_config.use_pplx_kernels - @property def use_deepep_ht_kernels(self): return self.moe_parallel_config.use_deepep_ht_kernels diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=96,device_name=NVIDIA_H200,dtype=fp8_w8a8.json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=96,device_name=NVIDIA_H200,dtype=fp8_w8a8.json new file mode 100644 index 00000000000..620fe9365aa --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=96,device_name=NVIDIA_H200,dtype=fp8_w8a8.json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=96,device_name=NVIDIA_H200.json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=96,device_name=NVIDIA_H200.json new file mode 100644 index 00000000000..fc7dda8a784 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=96,device_name=NVIDIA_H200.json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index ae9430d29d9..ac9ba56a6b7 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -166,7 +166,7 @@ def run_cutlass_moe_fp8( problem_sizes1 = torch.empty((local_E, 3), dtype=torch.int32, device=device) problem_sizes2 = torch.empty((local_E, 3), dtype=torch.int32, device=device) - ops.get_cutlass_pplx_moe_mm_data( + ops.get_cutlass_batched_moe_mm_data( expert_offsets, problem_sizes1, problem_sizes2, diff --git a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py index fbd47f8c423..24ae2d3c82c 100644 --- a/vllm/model_executor/layers/fused_moe/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_batched_moe.py @@ -493,7 +493,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): """ A reference prepare/finalize class that reorganizes the tokens into expert batched format, i.e. E x max_num_tokens x K. This is the format - that the PPLX dispatch/combine kernels use. + that the batched dispatch/combine kernels use. """ def __init__( @@ -648,7 +648,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute): """ A reference MoE expert class that operates on expert batched format, - i.e. E x max_num_tokens x K. This is the format that the pplx + i.e. E x max_num_tokens x K. This is the format that the batched dispatch/combine kernels use. """ @@ -880,7 +880,7 @@ def batched_moe_kernel_quantize_input( class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute): """ A Triton based MoE expert class that operates on expert batched format, - i.e. E x max_num_tokens x K. This is the format that the pplx + i.e. E x max_num_tokens x K. This is the format that the batched dispatch/combine kernels use. """ diff --git a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py index 5617156bf2f..2fcb7f19378 100644 --- a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py @@ -6,6 +6,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -178,7 +179,40 @@ def triton_kernel_moe_forward( apply_router_weight_on_input: bool = False, global_num_experts: int = -1, expert_map: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, ) -> torch.Tensor: + if ( + quant_config is not None + and quant_config.use_mxfp4_w4a8 + and rocm_aiter_ops.is_enabled() + ): + from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, topk, sm_first=not renormalize + ) + return triton_kernel_fused_mxfp4_w4a8_experts( + None, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + activation=activation.value, + quant_config=quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + expert_map=expert_map, + unpadded_N_w1=unpadded_N_w1, + unpadded_K_w1=unpadded_K_w1, + unpadded_N_w2=unpadded_N_w2, + unpadded_K_w2=unpadded_K_w2, + ) + if expert_map is not None: # With expert parallelism, legacy_routing produces routing data # using global expert IDs which don't correspond to local weight @@ -210,6 +244,9 @@ def triton_kernel_moe_forward( effective_global_num_experts = global_num_experts output = torch.empty_like(hidden_states) + effective_quant_config = ( + quant_config if quant_config is not None else FUSED_MOE_UNQUANTIZED_CONFIG + ) return triton_kernel_fused_experts( output, @@ -221,7 +258,7 @@ def triton_kernel_moe_forward( scatter_idx, topk=topk, activation=activation, - quant_config=quant_config, + quant_config=effective_quant_config, apply_router_weight_on_input=apply_router_weight_on_input, global_num_experts=effective_global_num_experts, expert_map=effective_expert_map, @@ -252,8 +289,7 @@ def triton_kernel_fused_experts( assert activation == MoEActivation.SWIGLUOAI, ( "Only SWIGLUOAI activation is supported" ) - if quant_config is None: - quant_config = FUSED_MOE_UNQUANTIZED_CONFIG + assert quant_config is not None # type check, uint8 means mxfp4 assert hidden_states.dtype == torch.bfloat16 @@ -330,6 +366,98 @@ def triton_kernel_fused_experts( return output_tensor +# This is a triton implementation of the fused_experts function +def triton_kernel_fused_mxfp4_w4a8_experts( + output_tensor: torch.Tensor, + hidden_states: torch.Tensor, + w1, # Tensor or triton_kernels.Tensor + w2, # Tensor or triton_kernels.Tensor + routing_data, # RoutingData + gather_indx, # GatherIndx + scatter_indx, # ScatterIndx + activation: str = "silu", + quant_config: FusedMoEQuantConfig | None = None, + swiglu_alpha: float = 1.702, + swiglu_limit: float = 7.0, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + a1q_scale: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, +) -> torch.Tensor: + assert quant_config is not None + # type check, uint8 means mxfp4 + assert hidden_states.dtype == torch.bfloat16 + assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32 + assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32 + + # Shape check, only check non-mxfp4 + assert hidden_states.shape[-1] == w1.shape[-2] + assert w2.shape[-1] == w1.shape[1] + + E, _, N = w1.shape + + if global_num_experts == -1: + global_num_experts = E + + gammas = routing_data.gate_scal if routing_data else None + + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + + assert quant_config.w1_precision is not None, ( + "w1_precision in quant config can't be None" + ) + assert quant_config.w2_precision is not None, ( + "w2_precision in quant config can't be None" + ) + + hidden_states = downcast_to_static_fp8( + hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale + ) + + intermediate_cache1 = moe_gemm_a8w4( + hidden_states, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + quant_config.w1_precision.flex_ctx.lhs_data.scale, + quant_config.w2_precision.flex_ctx.lhs_data.scale, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale="CDNA4_SCALE", + out_dtype=torch.float8_e4m3fn, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + + intermediate_cache3 = moe_gemm_a8w4( + intermediate_cache1, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + quant_config.w2_precision.flex_ctx.lhs_data.scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale="CDNA4_SCALE", + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + + return intermediate_cache3 + + def make_routing_data( topk_ids: torch.Tensor, topk_weights: torch.Tensor, @@ -520,6 +648,9 @@ class OAITritonExperts(BaseOAITritonExperts): expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): + if self.quant_config is None: + self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG + if expert_map is not None: topk_ids = expert_map[topk_ids] diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 679b79ce971..a7dee700415 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -525,16 +525,18 @@ class FusedMoE(CustomOp): # Round up hidden size before creating moe_config. # This way moe_config is created with the correct hidden_size from the start. + unpadded_hidden_size = hidden_size + self.model_type = ( + self.vllm_config.model_config.hf_config.model_type + if self.vllm_config.model_config is not None + else None + ) hidden_size = maybe_roundup_hidden_size( hidden_size=hidden_size, act_dtype=moe_in_dtype, moe_parallel_config=self.moe_parallel_config, is_lora_enabled=vllm_config.lora_config is not None, - model_type=( - self.vllm_config.model_config.hf_config.model_type - if self.vllm_config.model_config is not None - else None - ), + model_type=self.model_type, is_mxfp4_quant=( quant_config is not None and quant_config.is_mxfp4_quant(prefix, self) ), @@ -610,6 +612,7 @@ class FusedMoE(CustomOp): moe_quant_params = { "num_experts": self.local_num_experts, "hidden_size": hidden_size, + "unpadded_hidden_size": unpadded_hidden_size, "intermediate_size_per_partition": self.intermediate_size_per_partition, "params_dtype": params_dtype, "weight_loader": self.weight_loader, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index c2c0e809d70..043b5ef2669 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1172,9 +1172,9 @@ class FusedMoEModularKernel(torch.nn.Module): # This happens when none of the tokens from the all2all reach this # EP rank. Also, note that this is only relevant for CUDAGraph # incompatible all2all kernels like the DeepEP high-throughput - # kernels. CUDAGraph compatible all2all kernels like the pplx - # kernels and the DeepEP low-latency kernels are always batched - # and can never run into the tensor.numel() == 0 case. + # kernels. CUDAGraph compatible all2all kernels like the DeepEP + # low-latency kernels are always batched and can never run into + # the tensor.numel() == 0 case. if M_full == 0: assert num_chunks == 0 workspace13 = None diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index ee7db88ccbd..b4f4b74ca2d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -143,10 +143,7 @@ def select_nvfp4_moe_backend( # NOTE(rob): this is kind of a hack. We need to peak into # the prepare-finalize selection to determine if we are using # the batched or standard expert format. - use_batched = ( - config.moe_parallel_config.use_deepep_ll_kernels - or config.moe_parallel_config.use_pplx_kernels - ) + use_batched = config.moe_parallel_config.use_deepep_ll_kernels activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if use_batched diff --git a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py deleted file mode 100644 index 289ac0d1413..00000000000 --- a/vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py +++ /dev/null @@ -1,373 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import pplx_kernels as pplx -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceDelegate, -) -from vllm.model_executor.layers.fused_moe.utils import ( - _validate_scale_shape, - moe_kernel_quantize_input, -) -from vllm.utils.math_utils import cdiv, round_up - -logger = init_logger(__name__) - - -def pplx_hidden_dim_scale_bytes( - max_num_tokens: int, - hidden_dim: int, - in_dtype: torch.dtype, - quant_dtype: torch.dtype | str | None, - per_act_token_quant: bool, - block_shape: list[int] | None, -): - # All pplx byte sizes must be 16-byte aligned. - align = 16 - - # For blocked per token: set to - # cdiv(hidden_dim, block_size) * sizeof(float32) - # For per-token: set to 4 * sizeof(float32) (x4 for alignment) - if quant_dtype is not None: - assert isinstance(quant_dtype, torch.dtype) - assert quant_dtype.itemsize == 1 - hidden_dim_bytes = hidden_dim * quant_dtype.itemsize - elem_size = torch.float32.itemsize - - if per_act_token_quant: - # per-token (M x 1) - assert block_shape is None - hidden_scale_bytes = elem_size - elif block_shape is not None: - # per-group (M x K_tiles) - block_size = block_shape[1] - num_blocks = cdiv(hidden_dim, block_size) - hidden_scale_bytes = num_blocks * elem_size - else: - # per-tensor (1 x 1) - hidden_scale_bytes = elem_size - else: - hidden_dim_bytes = hidden_dim * in_dtype.itemsize - hidden_scale_bytes = 0 - - return ( - round_up(hidden_dim_bytes, align), - round_up(hidden_scale_bytes, align), - ) - - -class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): - """PPLX-based prepare and finalize for expert parallelism.""" - - def __init__( - self, - a2a: pplx.AllToAll, - max_num_tokens: int, - num_local_experts: int, - num_dispatchers: int, - ): - super().__init__() - assert max_num_tokens > 0 - assert num_local_experts > 0 - self.a2a = a2a - self.max_num_tokens = max_num_tokens - self.num_local_experts = num_local_experts - self.num_dispatchers_ = num_dispatchers - - @property - def activation_format(self) -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.BatchedExperts - - def max_num_tokens_per_rank(self) -> int | None: - return self.max_num_tokens - - def topk_indices_dtype(self) -> torch.dtype | None: - return torch.uint32 - - def num_dispatchers(self) -> int: - return self.num_dispatchers_ - - def output_is_reduced(self) -> bool: - return True - - def supports_async(self) -> bool: - return True - - def prepare_async( - self, - a1: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> tuple[Callable, mk.ReceiverType]: - if defer_input_quant: - raise NotImplementedError( - f"{self.__class__.__name__} does not support defer_input_quant=True. " - "Please select an MoE kernel that accepts quantized inputs." - ) - - num_tokens = a1.size(0) # M - hidden_dim = a1.size(-1) # K - - assert topk_ids.size(0) == num_tokens - # expert_map should be None because with expert map, -1 id is used for - # non-local token; this causes error when casting ids to the - # topk_indices_dtype() int32 - # - if expert_map is not None: - logger.warning_once( - "The PPLX backend does not support expert mapping. " - "The provided `expert_map` will be ignored." - ) - expert_map = None # noqa: F841 - - # Is this always going to be a1.device? - device = a1.device - - if apply_router_weight_on_input: - topk = topk_ids.size(1) - # TODO: this only works for topK=1, will need to update for topK>1 - assert topk == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - a1 = a1 * topk_weights.to(a1.dtype) - - repeat_cols = 4 - repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0) - # TODO(bnell): always pass quant_config.a1_scale? - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - (None if quant_config.per_act_token_quant else quant_config.a1_scale), - quant_dtype=quant_config.quant_dtype, - per_act_token_quant=quant_config.per_act_token_quant, - block_shape=quant_config.block_shape, - ) - - _validate_scale_shape( - a1q, a1q_scale, quant_config.per_act_token_quant, quant_config.block_shape - ) - - orig_a_scale_block_shape: int | None = None - - if a1q_scale is not None: - scalar_scales = a1q_scale.numel() == 1 - - # pplx requires 2-d scales even for scalar scales - if a1q_scale.dim() <= 1: - assert scalar_scales - a1q_scale = a1q_scale.view(1, 1) - - orig_a_scale_block_shape = a1q_scale.shape[-1] - - if not quant_config.is_block_quantized: - # TODO (bnell): use group_broadcast instead? - a1q_scale = a1q_scale.repeat(repeat_rows, repeat_cols) - - assert a1q_scale is None or a1q_scale.ndim == 2, ( - f"{0 if a1q_scale is None else (a1q_scale.ndim, a1q_scale.shape)}" - ) - - expert_num_tokens = torch.empty( - self.num_local_experts, - dtype=torch.int32, - device=device, - ) - - expert_x = torch.empty( - ( - self.num_local_experts, - self.max_num_tokens * self.num_dispatchers(), - hidden_dim, - ), - dtype=a1q.dtype, - device=device, - ) - - expert_x_scale: torch.Tensor | None = None - if a1q.dtype.itemsize == 1: - if quant_config.is_per_act_token: - # (M x 1) -> (E x M x K) - final_dim = expert_x.size(2) - elif quant_config.is_per_tensor: - # (1 x 1) -> (E x 1 x 1) - final_dim = 1 - else: - # (M x K_tiles) -> (E x M x K_tiles) - assert quant_config.block_shape is not None - num_blocks = cdiv(expert_x.size(2), quant_config.block_shape[1]) - final_dim = num_blocks - - expert_x_scale_shape = ( - self.num_local_experts, - expert_x.size(1), - round_up(final_dim, 4), # round up for alignment - ) - - expert_x_scale = torch.empty( - expert_x_scale_shape, - dtype=torch.float32, - device=expert_x.device, - ) - - # This argument is optional, defaults to indices.size(0) - # There's not much point setting this unless it is != indices.size(0) - bound_m: torch.Tensor | None = None - - self.a2a.dispatch( - out_expert_num_tokens=expert_num_tokens, - out_expert_x=expert_x, - out_expert_x_scale=expert_x_scale, - dp_x=a1q, - dp_x_scale=a1q_scale, - indices=topk_ids, - bound_m=bound_m, - do_send=True, - do_recv=False, - ) - - hook = lambda: self.a2a.dispatch( - out_expert_num_tokens=expert_num_tokens, - out_expert_x=expert_x, - out_expert_x_scale=expert_x_scale, - dp_x=a1q, - dp_x_scale=a1q_scale, - indices=topk_ids, - bound_m=bound_m, - do_send=False, - do_recv=True, - ) - - return ( - hook, - lambda: self._receiver( - expert_num_tokens, - expert_x, - expert_x_scale, - orig_a_scale_block_shape, - ), - ) - - def _receiver( - self, - expert_num_tokens: torch.Tensor, - expert_x: torch.Tensor, - expert_x_scale: torch.Tensor | None, - orig_a_scale_block_shape: int | None, - ) -> mk.PrepareResultType: - if expert_x_scale is not None: - expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape] - assert expert_x_scale.ndim == 3 - - expert_tokens_meta = mk.ExpertTokensMetadata( - expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None - ) - - return expert_x, expert_x_scale, expert_tokens_meta, None, None - - def prepare( - self, - a1: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None, - apply_router_weight_on_input: bool, - quant_config: FusedMoEQuantConfig, - defer_input_quant: bool = False, - ) -> mk.PrepareResultType: - hook, receiver = self.prepare_async( - a1, - topk_weights, - topk_ids, - num_experts, - expert_map, - apply_router_weight_on_input, - quant_config, - defer_input_quant=defer_input_quant, - ) - hook() - return receiver() - - def finalize_async( - self, - output: torch.Tensor, - fused_expert_output: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - apply_router_weight_on_input: bool, - weight_and_reduce_impl: mk.TopKWeightAndReduce, - ) -> Callable: - assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate), ( - "Weight application and reduction happens in the combine kernel." - ) - - # This argument is optional - # There's not much point setting this unless it is != topk_ids.size(0) - bound_m: torch.Tensor | None = None - - # TODO (bnell): fails in test_pplx_moe.py, figure out what's going on - # num_tokens = output.size(0) # M - # assert topk_ids.size(0) == num_tokens, ( - # f"{topk_ids.size(0)} == {num_tokens}") - assert topk_ids.size() == topk_weights.size(), ( - f"{topk_ids.size()} == {topk_weights.size()}" - ) - assert output.size(0) <= self.max_num_tokens, ( - f"{output.size(0)} <= {self.max_num_tokens}" - ) - assert output.size(1) == fused_expert_output.size(-1) - - # Set weights to 1 if we did them in dispatch. This is hacky. - if apply_router_weight_on_input: - topk_weights = torch.ones_like(topk_weights) - - topk_ids_u32 = topk_ids.view(dtype=torch.uint32) - - self.a2a.combine( - out_tokens=output, - indices=topk_ids_u32, - weights=topk_weights, - expert_y=fused_expert_output, - bound_m=bound_m, - do_send=True, - do_recv=False, - ) - - return lambda: self.a2a.combine( - out_tokens=output, - indices=topk_ids_u32, - weights=topk_weights, - expert_y=fused_expert_output, - bound_m=bound_m, - do_send=False, - do_recv=True, - ) - - def finalize( - self, - output: torch.Tensor, - fused_expert_output: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - apply_router_weight_on_input: bool, - weight_and_reduce_impl: mk.TopKWeightAndReduce, - ) -> None: - receiver = self.finalize_async( - output, - fused_expert_output, - topk_weights, - topk_ids, - apply_router_weight_on_input, - weight_and_reduce_impl, - ) - receiver() diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py new file mode 100644 index 00000000000..77d8e756026 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +from torch.nn.parameter import Parameter + +from vllm.model_executor.custom_op import PluggableLayer +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.platforms import current_platform + + +@PluggableLayer.register("gate_linear") +class GateLinear(ReplicatedLinear): + """MoE gate linear layer with three-tier GEMM dispatch: + + 1. DSV3 specialized kernel (SM90+, batch<=16, supported dims) + 2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype) + 3. F.linear via ReplicatedLinear (ultimate fallback) + + The ``out_dtype`` attribute is mutable and can be set after init + (e.g. when the required dtype depends on the expert quantization + method which is only known later). + """ + + # Dimensions supported by the DSV3 specialized kernel + DSV3_SUPPORTED_NUM_EXPERTS = [256, 384] + DSV3_SUPPORTED_HIDDEN_SIZES = [7168] + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = False, + out_dtype: torch.dtype | None = None, + params_dtype: torch.dtype | None = None, + force_fp32_compute: bool = False, + prefix: str = "", + ): + is_hopper_or_blackwell = current_platform.is_device_capability( + (9, 0) + ) or current_platform.is_device_capability_family(100) + can_use_specialized_kernels = ( + current_platform.is_cuda() and is_hopper_or_blackwell and not bias + ) + + # If fp32 compute is required and no specialized kernel is available, + # store weights in fp32 so Tier 3 computes in fp32 natively. + if force_fp32_compute and not can_use_specialized_kernels: + params_dtype = torch.float32 + + super().__init__( + input_size, + output_size, + bias=bias, + params_dtype=params_dtype, + quant_config=None, + prefix=prefix, + ) + self.out_dtype = out_dtype + + # DSV3 specialized kernel eligibility (SM90+, exact dims) + self.allow_specialized_router_gemm = can_use_specialized_kernels + self.allow_dsv3_router_gemm = ( + self.allow_specialized_router_gemm + and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS + and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES + ) + + # cuBLAS bf16→fp32 eligibility + self.allow_cublas_router_gemm = ( + self.allow_specialized_router_gemm + and self.weight.dtype == torch.bfloat16 + and self.out_dtype == torch.float32 + ) + + def set_out_dtype(self, out_dtype: torch.dtype) -> None: + """Set output dtype for the router logits after init. + + Useful when the required dtype depends on the expert quantization + method which is only known after the gate is constructed. + """ + if self.out_dtype is not None: + raise ValueError("out_dtype has already been set") + self.out_dtype = out_dtype + + if ( + not self.allow_cublas_router_gemm + and self.allow_specialized_router_gemm + and out_dtype == torch.float32 + ): + self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16 + + def forward( + self, x: torch.Tensor + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: + import vllm._custom_ops as ops + + # Tier 1: DSV3 specialized kernel + if self.allow_dsv3_router_gemm and x.shape[0] <= 16: + output = ops.dsv3_router_gemm( + hidden_states=x, + router_weight=self.weight, + output_dtype=self.out_dtype, + ) + return output, None + + # Tier 2: cuBLAS bf16→fp32 + if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: + output = ops.router_gemm_bf16_fp32(x, self.weight) + return output, None + + # Tier 3: F.linear (ReplicatedLinear) + if self.out_dtype is not None and x.dtype != self.weight.dtype: + x = x.to(self.weight.dtype) + output, output_bias = super().forward(x) + if self.out_dtype is not None and output.dtype != self.out_dtype: + output = output.to(self.out_dtype) + return output, output_bias diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 7e25c968740..9c2adf7996e 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -216,8 +216,7 @@ class DefaultMoERunner(MoERunner): @property def use_dp_chunking(self) -> bool: return ( - self.moe_config.moe_parallel_config.use_pplx_kernels - or self.moe_config.moe_parallel_config.use_deepep_ll_kernels + self.moe_config.moe_parallel_config.use_deepep_ll_kernels or self.moe_config.moe_parallel_config.use_mori_kernels or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 99d4038ec38..d7b50aea2ad 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -14,10 +14,11 @@ class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce): implementation does not perform weight application and reduction but cannot address the needs of all the compatible PrepareAndFinalize implementations. - For example, BatchedTritonExperts is compatible with both - PplxPrepareAndFinalize and BatchedPrepareAndFinalize. PplxPrepareAndFinalize - does the weight-application + reduction as part of the pplx combine kernel. - But the BatchedPrepareAndFinalize needs an implementation. To facilitate + For example, BatchedTritonExperts is compatible with both batched + PrepareAndFinalize implementations like DeepEPLLPrepareAndFinalize and + BatchedPrepareAndFinalize. Some PrepareAndFinalize implementations do + the weight-application + reduction as part of the combine kernel, while + BatchedPrepareAndFinalize needs an explicit implementation. To facilitate this case, the BatchedTritonExperts could use TopKWeightAndReduceDelegate so the PrepareAndFinalize implementations could choose how to weight + reduce. diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 17b90c97012..72f42de06ee 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -510,6 +510,7 @@ class RMSNormGated(CustomOp): norm_before_gate: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None, + activation: str = "swish", ): """Initialize RMSNormGated. @@ -524,10 +525,12 @@ class RMSNormGated(CustomOp): If False and z is provided: out = norm(x * silu(z)) device: Device to create parameters on dtype: Data type for parameters + activation: Activation function name for gating """ factory_kwargs = {"device": device, "dtype": dtype} super().__init__() self.eps = eps + self.activation = activation self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) self.register_parameter("bias", None) self.group_size = group_size @@ -577,7 +580,7 @@ class RMSNormGated(CustomOp): if z is not None and self.norm_before_gate: out = out * F.silu(z) - return out + return out.to(x.dtype) def forward_cuda( self, x: torch.Tensor, z: torch.Tensor | None = None diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 09e67f562d0..2fb54e7751a 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -18,6 +18,7 @@ QuantizationMethods = Literal[ "modelopt", "modelopt_fp4", "modelopt_mxfp8", + "modelopt_mixed", "gguf", "gptq_marlin", "awq_marlin", @@ -120,7 +121,12 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: from .gptq import GPTQConfig from .gptq_marlin import GPTQMarlinConfig from .inc import INCConfig - from .modelopt import ModelOptFp8Config, ModelOptMxFp8Config, ModelOptNvFp4Config + from .modelopt import ( + ModelOptFp8Config, + ModelOptMixedPrecisionConfig, + ModelOptMxFp8Config, + ModelOptNvFp4Config, + ) from .moe_wna16 import MoeWNA16Config from .mxfp4 import Mxfp4Config from .petit import PetitNvFp4Config @@ -135,6 +141,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "modelopt": ModelOptFp8Config, "modelopt_fp4": ModelOptNvFp4Config, "modelopt_mxfp8": ModelOptMxFp8Config, + "modelopt_mixed": ModelOptMixedPrecisionConfig, "gguf": GGUFConfig, "gptq_marlin": GPTQMarlinConfig, "awq_marlin": AWQMarlinConfig, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 4c059da4170..999bb632504 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -12,7 +12,7 @@ from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_fp8_linear_kernel, ) -from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention import Attention, MLAAttention from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -114,6 +114,8 @@ QUANT_ALGOS = [ "NVFP4", # MXFP8 "MXFP8", + # MIXED_PRECISION, + "MIXED_PRECISION", ] KV_CACHE_QUANT_ALGOS = ["FP8"] @@ -181,7 +183,7 @@ class ModelOptQuantConfigBase(QuantizationConfig): self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": # handle kv-cache first so we can focus only on weight quantization thereafter - if isinstance(layer, Attention): + if isinstance(layer, (Attention, MLAAttention)): return self.KVCacheMethodCls(self) # handle exclusion @@ -235,6 +237,26 @@ class ModelOptQuantConfigBase(QuantizationConfig): self.exclude_modules = hf_to_vllm_mapper.apply_list(new_exclude_modules) + @staticmethod + def _extract_modelopt_quant_algo( + hf_quant_cfg: dict[str, Any] | None, + ) -> str | None: + """Extract upper-cased quant_algo from a modelopt config. + + Returns the quant_algo string (upper-cased), or None if the config + is not a modelopt config. + """ + if hf_quant_cfg is None: + return None + if hf_quant_cfg.get("quant_method", "").lower() != "modelopt": + return None + if "quantization" in hf_quant_cfg: + quant_config = hf_quant_cfg["quantization"] + if isinstance(quant_config, dict): + return str(quant_config.get("quant_algo", "")).upper() + return None + return str(hf_quant_cfg.get("quant_algo", "")).upper() + @staticmethod def get_config_filenames() -> list[str]: return ["hf_quant_config.json"] @@ -272,10 +294,20 @@ class ModelOptQuantConfigBase(QuantizationConfig): # "exclude_modules" is the key in the legacy hf_quant_config.json exclude_modules = quant_config.get("exclude_modules", []) else: - # Compressed-tensors style format: + # Compressed-tensors style format (config.json quantization_config): # {"quant_algo": "...", "quant_method": "modelopt"} quant_method = config.get("quant_algo") - kv_cache_quant_method = config.get("kv_cache_quant_algo") + + # "kv_cache_scheme" (a dict) instead of "kv_cache_quant_algo" (a string). + kv_cache_scheme = config.get("kv_cache_scheme") + if isinstance(kv_cache_scheme, dict) and ( + kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 8 + ): + kv_cache_quant_method = "FP8" + else: + kv_cache_quant_method = None + # "ignore" is the key in config.json exclude_modules = config.get("ignore", []) group_size_raw = config.get("group_size") @@ -379,32 +411,9 @@ class ModelOptFp8Config(ModelOptQuantConfigBase): def override_quantization_method( cls, hf_quant_cfg, user_quant ) -> QuantizationMethods | None: - """Detect if this ModelOpt config should be used based on - quantization config.""" - - if hf_quant_cfg is None: - return None - - # Use the community standard 'quant_method' - quant_method = hf_quant_cfg.get("quant_method", "").lower() - - # Only proceed if the method is explicitly "modelopt" - if quant_method != "modelopt": - return None - - # Look for ModelOpt-specific config structure - if "quantization" in hf_quant_cfg: - quant_config = hf_quant_cfg["quantization"] - if isinstance(quant_config, dict): - quant_algo = str(quant_config.get("quant_algo", "")) - if quant_algo.upper() == "FP8": - return "modelopt" - else: - # Check for compressed-tensors style config with specific quant_algo - quant_algo = str(hf_quant_cfg.get("quant_algo", "")) - if quant_algo.upper() == "FP8": - return "modelopt" - + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) + if algo is not None and algo == "FP8": + return "modelopt" return None @classmethod @@ -1031,32 +1040,9 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase): def override_quantization_method( cls, hf_quant_cfg, user_quant ) -> QuantizationMethods | None: - """Detect if this ModelOpt FP4 config should be used based on - quantization config.""" - if hf_quant_cfg is None: - return None - - # Use the community standard 'quant_method' - quant_method = hf_quant_cfg.get("quant_method", "").lower() - - # Only proceed if the method is explicitly "modelopt" - if quant_method != "modelopt": - return None - - # Look for ModelOpt-specific config structure - if "quantization" in hf_quant_cfg: - quant_config = hf_quant_cfg["quantization"] - if isinstance(quant_config, dict): - quant_algo = quant_config.get("quant_algo", "") - if "NVFP4" in quant_algo: - return "modelopt_fp4" - else: - # Check for compressed-tensors style config with specific - # quant_algo field - quant_algo = hf_quant_cfg.get("quant_algo", "") - if isinstance(quant_algo, str) and "FP4" in quant_algo.upper(): - return "modelopt_fp4" - + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) + if algo is not None and ("NVFP4" in algo or "FP4" in algo): + return "modelopt_fp4" return None @classmethod @@ -1619,31 +1605,9 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase): def override_quantization_method( cls, hf_quant_cfg, user_quant ) -> QuantizationMethods | None: - """Detect if this ModelOpt MXFP8 config should be used based on - quantization config.""" - if hf_quant_cfg is None: - return None - - # Use the community standard 'quant_method' - quant_method = hf_quant_cfg.get("quant_method", "").lower() - - # Only proceed if the method is explicitly "modelopt" - if quant_method != "modelopt": - return None - - # Look for ModelOpt-specific config structure - if "quantization" in hf_quant_cfg: - quant_config = hf_quant_cfg["quantization"] - if isinstance(quant_config, dict): - quant_algo = str(quant_config.get("quant_algo", "")).upper() - if "MXFP8" in quant_algo: - return "modelopt_mxfp8" - else: - # Check for compressed-tensors style config with specific quant_algo - quant_algo = str(hf_quant_cfg.get("quant_algo", "")).upper() - if "MXFP8" in quant_algo: - return "modelopt_mxfp8" - + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) + if algo is not None and "MXFP8" in algo: + return "modelopt_mxfp8" return None @classmethod @@ -1841,3 +1805,188 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase): # Register the method classes for ModelOptMxFp8Config ModelOptMxFp8Config.LinearMethodCls = ModelOptMxFp8LinearMethod ModelOptMxFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod + + +class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): + """Config class for ModelOpt MIXED_PRECISION. + + Supports checkpoints where different layers use different quantization + algorithms (e.g., FP8 for dense layers and NVFP4 for MoE experts). + The per-layer algorithm is specified in the ``quantized_layers`` dict + inside ``config.json``'s ``quantization_config`` (preferred) or the + legacy ``hf_quant_config.json``. + """ + + def __init__( + self, + kv_cache_quant_method: str | None, + exclude_modules: list[str], + quantized_layers: dict[str, dict[str, Any]], + fp8_config: ModelOptFp8Config, + nvfp4_config: ModelOptNvFp4Config, + ) -> None: + super().__init__(exclude_modules) + self.kv_cache_quant_method = kv_cache_quant_method + self.quantized_layers = quantized_layers + self.fp8_config = fp8_config + self.nvfp4_config = nvfp4_config + + def get_name(self) -> QuantizationMethods: + return "modelopt_mixed" + + def get_supported_act_dtypes(self) -> list[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 89 + + @classmethod + def override_quantization_method( + cls, hf_quant_cfg, user_quant + ) -> QuantizationMethods | None: + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) + if algo is not None and algo == "MIXED_PRECISION": + return "modelopt_mixed" + return None + + @classmethod + def _from_config( + cls, + *, + quant_method: str, + kv_cache_quant_method: str | None, + exclude_modules: list[str], + original_config: dict[str, Any], + group_size: int | None, + **kwargs: Any, + ) -> "ModelOptMixedPrecisionConfig": + if "quantization" in original_config: + quantized_layers = original_config["quantization"].get( + "quantized_layers", {} + ) + else: + quantized_layers = original_config.get("quantized_layers", {}) + + if not quantized_layers: + raise ValueError( + "MIXED_PRECISION quant_algo requires a non-empty " + "'quantized_layers' mapping in the quantization config." + ) + + # Determine group_size from the first NVFP4 entry if not provided. + if group_size is None: + for layer_info in quantized_layers.values(): + if layer_info.get("quant_algo", "").upper() == "NVFP4": + group_size = layer_info.get("group_size", 16) + break + if group_size is None: + group_size = 16 + + fp8_config = ModelOptFp8Config( + quant_method="FP8", + is_checkpoint_fp8_serialized=True, + kv_cache_quant_method=kv_cache_quant_method, + exclude_modules=[], + ) + nvfp4_config = ModelOptNvFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=kv_cache_quant_method, + exclude_modules=[], + group_size=group_size, + ) + + return cls( + kv_cache_quant_method=kv_cache_quant_method, + exclude_modules=exclude_modules, + quantized_layers=quantized_layers, + fp8_config=fp8_config, + nvfp4_config=nvfp4_config, + ) + + def _resolve_quant_algo(self, prefix: str) -> str | None: + """Look up the quant_algo for a vLLM-side layer prefix. + + Tries three strategies in order: + 1. Direct lookup in ``quantized_layers``. + 2. Packed/fused-layer lookup (unfuse via ``packed_modules_mapping``). + 3. Prefix-based lookup for FusedMoE (any child key starts with + ``prefix + "."``). + + Returns the upper-cased quant_algo string, or *None* if the prefix + is not found. + """ + # 1. Direct lookup + if prefix in self.quantized_layers: + return self.quantized_layers[prefix]["quant_algo"].upper() + + # 2. Packed / fused layer lookup + proj_name = prefix.rsplit(".", 1)[-1] + if self.packed_modules_mapping and proj_name in self.packed_modules_mapping: + algos: set[str] = set() + base = prefix.rsplit(".", 1)[0] + for shard_name in self.packed_modules_mapping[proj_name]: + shard_prefix = f"{base}.{shard_name}" + if shard_prefix in self.quantized_layers: + algos.add(self.quantized_layers[shard_prefix]["quant_algo"].upper()) + if len(algos) == 1: + return algos.pop() + if len(algos) > 1: + raise ValueError( + f"Mixed quant_algo within fused layer {prefix}: " + f"{algos}. All shards must use the same quantization." + ) + + # 3. Prefix-based lookup (for FusedMoE / parent modules) + prefix_dot = prefix + "." + for key, info in self.quantized_layers.items(): + if key.startswith(prefix_dot): + return info["quant_algo"].upper() + + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> "QuantizeMethodBase | None": + """Return quantize-method based on layer.""" + # KV-cache quantization + if isinstance(layer, Attention): + if self.kv_cache_quant_method: + return ModelOptFp8KVCacheMethod(self) + return None + + # Excluded layers + if self.is_layer_excluded(prefix): + if isinstance(layer, LinearBase): + return UnquantizedLinearMethod() + return None + + quant_algo = self._resolve_quant_algo(prefix) + + if isinstance(layer, LinearBase): + if quant_algo == "FP8": + return ModelOptFp8LinearMethod(self.fp8_config) + if quant_algo == "NVFP4": + return ModelOptNvFp4LinearMethod(self.nvfp4_config) + # Layer not in quantized_layers — leave unquantized + return UnquantizedLinearMethod() + + if isinstance(layer, FusedMoE): + if quant_algo == "FP8": + return ModelOptFp8MoEMethod( + quant_config=self.fp8_config, + moe_config=layer.moe_config, + ) + if quant_algo == "NVFP4": + return ModelOptNvFp4FusedMoE( + quant_config=self.nvfp4_config, + moe_config=layer.moe_config, + ) + return None + + return None + + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): + super().apply_vllm_mapper(hf_to_vllm_mapper) + if self.quantized_layers: + self.quantized_layers = hf_to_vllm_mapper.apply_dict(self.quantized_layers) diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index d81f0f80d2e..9318bedfff1 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -798,7 +798,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # batched activation format. As self.fused_experts is not # initialized at this point, we resort to checking the MoE config # directly. - is_batched_moe = self.moe.use_pplx_kernels or self.moe.use_deepep_ll_kernels + is_batched_moe = self.moe.use_deepep_ll_kernels if is_batched_moe: num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8 else: diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 8394857cf9c..b2abbce1aa1 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -5,8 +5,8 @@ from typing import Any import torch -import vllm.envs as envs from vllm import _custom_ops as ops +from vllm import envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import get_current_vllm_config from vllm.logger import init_logger @@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.mxfp4 import ( from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( prepare_fp8_moe_layer_for_marlin, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, OCP_MX_Scheme, @@ -49,7 +50,11 @@ from vllm.utils.math_utils import round_up logger = init_logger(__name__) -__all__ = ["QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod", "QuarkOCP_MX_MoEMethod"] +__all__ = [ + "QuarkMoEMethod", + "QuarkOCP_MX_MoEMethod", + "QuarkOCP_MX_MoEMethod_OSS", +] class QuarkMoEMethod(FusedMoEMethodBase): @@ -71,14 +76,30 @@ class QuarkMoEMethod(FusedMoEMethodBase): "output_tensors and bias " "quantized are not supported" ) + weight_config = layer_quant_config.get("weight") input_config = layer_quant_config.get("input_tensors") + if quant_config._is_fp8_w4a8(weight_config, input_config): return QuarkW4A8Fp8MoEMethod(weight_config, input_config, module.moe_config) elif quant_config._is_fp8_w8a8(weight_config, input_config): return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config) elif quant_config._is_w_ocp_mx_a_x(weight_config, input_config): - return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config) + emulate = not current_platform.supports_mx() or not ( + rocm_aiter_ops.is_fused_moe_enabled() + ) + if ( + input_config.get("dtype") == "fp8_e4m3" + and not input_config.get("is_dynamic") + and not emulate + ): + return QuarkOCP_MX_MoEMethod_OSS( + weight_config, input_config, module.moe_config + ) + else: + return QuarkOCP_MX_MoEMethod( + weight_config, input_config, module.moe_config + ) else: raise RuntimeError("Unsupported FusedMoe scheme") @@ -706,13 +727,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): get_current_vllm_config().model_config.hf_config, "model_type", None ) - self._emulate = ( + self.emulate = ( not current_platform.supports_mx() or not self.ocp_mx_scheme.startswith("w_mxfp4") ) and (self.mxfp4_backend is None or not self.use_rocm_aiter_moe) - self.emulate = True if self.model_type == "gpt_oss" else self._emulate - if self.emulate: logger.warning_once( f"The current mode (supports_mx={current_platform.supports_mx()}, " @@ -753,6 +772,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): ) params_dtype = torch.uint8 + self.intermediate_size_per_partition = intermediate_size_per_partition if self.model_type == "gpt_oss": if current_platform.is_rocm(): intermediate_size_per_partition_after_pad = round_up( @@ -765,6 +785,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): else: intermediate_size_per_partition_after_pad = intermediate_size_per_partition + self.unpadded_hidden_size = extra_weight_attrs.get( + "unpadded_hidden_size", hidden_size + ) + # WEIGHTS w13_weight = torch.nn.Parameter( torch.empty( @@ -991,30 +1015,20 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if not self.emulate: - if ( - self.model_type == "gpt_oss" - and self.mxfp4_backend == Mxfp4Backend.TRITON - ): - raise NotImplementedError( - "Triton kernel implemented fused MoE for GPT_OSS model " - "in Quark(MoE) format is not integrated or provided yet." - ) + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( + rocm_aiter_fused_experts, + ) - else: - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( - rocm_aiter_fused_experts, - ) - - return rocm_aiter_fused_experts( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=layer.activation, - quant_config=self.moe_quant_config, - expert_map=layer.expert_map, - ) + return rocm_aiter_fused_experts( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + quant_config=self.moe_quant_config, + expert_map=layer.expert_map, + ) else: from vllm.model_executor.layers.fused_moe import fused_experts @@ -1031,3 +1045,133 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): expert_map=layer.expert_map, quant_config=self.moe_quant_config, ) + + +class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod): + def __init__( + self, + weight_config: dict[str, Any], + input_config: dict[str, Any], + moe: FusedMoEConfig, + ): + super().__init__(weight_config, input_config, moe) + + def process_weights_after_loading(self, layer): + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig + + w13_bias = layer.w13_bias.to(torch.float32) + w2_bias = layer.w2_bias.to(torch.float32) + + layer.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False) + layer.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False) + + # FIXME warp need to be adjusted based on batch size + # only apply to batched mode + if self.moe.use_ep: + num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8 + else: + num_warps = 8 + + w13_weight, w13_flex, w13_scale = _swizzle_mxfp4( + layer.w13_weight, layer.w13_weight_scale, num_warps + ) + w2_weight, w2_flex, w2_scale = _swizzle_mxfp4( + layer.w2_weight, layer.w2_weight_scale, num_warps + ) + + self.w13_weight_triton_tensor = w13_weight + self.w2_weight_triton_tensor = w2_weight + + # need to delete the original weights to save memory on single GPU + del layer.w13_weight + del layer.w2_weight + layer.w13_weight = None + layer.w2_weight = None + torch.cuda.empty_cache() + + if self.static_input_scales: + if layer.w13_input_scale is None or layer.w2_input_scale is None: + raise ValueError( + "QuantConfig has static quantization, but found " + "activation scales are None." + ) + if not all_close_1d(layer.w13_input_scale) or not all_close_1d( + layer.w2_input_scale + ): + logger.warning_once( + "Found input_scales that are not equal for " + "fp8 MoE layer. Using the maximum across experts " + "for each layer." + ) + + layer.w13_input_scale = torch.nn.Parameter( + layer.w13_input_scale.max().to(torch.float32), requires_grad=False + ) + layer.w2_input_scale = torch.nn.Parameter( + layer.w2_input_scale.max().to(torch.float32), requires_grad=False + ) + + from triton_kernels.numerics import InFlexData + + lhs_data13 = InFlexData(scale=layer.w13_input_scale) + lhs_data2 = InFlexData(scale=layer.w2_input_scale) + + self.w13_precision_config = PrecisionConfig( + weight_scale=w13_scale, + flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13), + ) + + self.w2_precision_config = PrecisionConfig( + weight_scale=w2_scale, + flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2), + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + return mxfp4_w4a8_moe_quant_config( + w1_scale=self.w13_precision_config, + w2_scale=self.w2_precision_config, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + w1_bias=layer.w13_bias, + w2_bias=layer.w2_bias, + block_shape=None, + ) + + @property + def is_monolithic(self) -> bool: + return True + + def apply_monolithic( + self, + layer: torch.nn.Module, + x: torch.Tensor, + router_logits: torch.Tensor, + expert_map: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if layer.enable_eplb: + raise NotImplementedError( + "EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet." + ) + + from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501 + triton_kernel_moe_forward, + ) + + return triton_kernel_moe_forward( + hidden_states=x, + w1=self.w13_weight_triton_tensor, + w2=self.w2_weight_triton_tensor, + gating_output=router_logits, + topk=layer.top_k, + renormalize=layer.renormalize, + global_num_experts=layer.global_num_experts, + expert_map=expert_map, + quant_config=self.moe_quant_config, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + unpadded_N_w1=self.intermediate_size_per_partition * 2, + unpadded_K_w1=self.unpadded_hidden_size, + unpadded_N_w2=self.unpadded_hidden_size, + unpadded_K_w2=self.intermediate_size_per_partition, + ) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 44dcd076e0b..24b2f61b867 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -287,7 +287,17 @@ def get_quant_config( ) if hf_quant_config is not None: - return quant_cls.from_config(hf_quant_config) + # For modelopt_mixed, config.json's quantization_config may or may + # not contain the per-layer quantized_layers map. Newer checkpoints + # embed it directly; older ones keep it only in hf_quant_config.json. + # If it is missing, fall through to the file-based loading path. + if ( + model_config.quantization == "modelopt_mixed" + and "quantized_layers" not in hf_quant_config + ): + pass # fall through to file-based loading below + else: + return quant_cls.from_config(hf_quant_config) # if hf_quant_config is None, we will try to get config from # hf_overrides @@ -365,8 +375,8 @@ def get_quant_config( if model_config.quantization == "bitsandbytes": config["adapter_name_or_path"] = model_config.model - elif model_config.quantization == "modelopt": - if config["producer"]["name"] == "modelopt": + elif model_config.quantization in ("modelopt", "modelopt_mixed"): + if config.get("producer", {}).get("name") == "modelopt": return quant_cls.from_config(config) else: raise ValueError( diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py new file mode 100644 index 00000000000..f5ed4400fb6 --- /dev/null +++ b/vllm/model_executor/models/AXK1.py @@ -0,0 +1,1168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py +# Copyright 2023 The vLLM team. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only A.X K1 model.""" + +import typing +from collections.abc import Callable, Iterable +from itertools import islice + +import torch +from torch import nn + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, ParallelConfig, VllmConfig +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import SharedFusedMoE +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.deepseek_v2 import ( + DeepseekAttention, + DeepseekV2MLP, + yarn_get_mscale, +) +from vllm.model_executor.models.utils import sequence_parallel_chunk +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.AXK1 import AXK1Config + +from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP +from .utils import ( + PPMissingLayer, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) + +logger = init_logger(__name__) + + +class AXK1MLP(DeepseekV2MLP): + pass + + +class AXK1MoE(nn.Module): + def __init__( + self, + config: AXK1Config, + parallel_config: ParallelConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + + self.routed_scaling_factor = config.routed_scaling_factor + + self.ep_group = get_ep_group().device_group + self.ep_rank = get_ep_group().rank_in_group + self.ep_size = self.ep_group.size() + self.n_routed_experts: int = config.n_routed_experts + self.n_shared_experts: int = config.n_shared_experts + + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.n_routed_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + if config.topk_method == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32) + ) + else: + self.gate.e_score_correction_bias = None + + # Load balancing settings. + eplb_config = parallel_config.eplb_config + self.enable_eplb = parallel_config.enable_eplb + + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_logical_experts = self.n_routed_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() + self.is_fusion_moe_shared_experts_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) + if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled: + self.shared_experts = None + else: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + + self.shared_experts = AXK1MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + is_sequence_parallel=self.is_sequence_parallel, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + reduce_results=False, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + use_grouped_topk=True, + num_expert_group=config.n_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.scoring_func, + # we do scaling outside, set factor to 1.0 to avoid double mul + # aiter applies routed_scaling_factor internally + routed_scaling_factor=1.0 + if not self.is_rocm_aiter_moe_enabled + else self.routed_scaling_factor, + e_score_correction_bias=self.gate.e_score_correction_bias, + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + is_sequence_parallel=self.is_sequence_parallel, + n_shared_experts=config.n_shared_experts + if self.is_fusion_moe_shared_experts_enabled + else None, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # Chunk the hidden states so they aren't replicated across TP ranks. + # This avoids duplicate computation in self.experts. + # TODO: We can replace the all_reduce at the end of attn with a + # reduce_scatter instead of chunking here. + if self.is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + + if self.experts.is_internal_router: + # In this case, the gate/router runs inside the FusedMoE class + fused_moe_out = self.experts( + hidden_states=hidden_states, router_logits=hidden_states + ) + else: + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + fused_moe_out = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + + shared_output, final_hidden_states = fused_moe_out + if self.shared_experts is None: + assert shared_output is None + + # Fix FP16 overflow + # See AXK1DecoderLayer for more details. + if hidden_states.dtype != torch.float16: + if not self.is_rocm_aiter_moe_enabled: + final_hidden_states *= self.routed_scaling_factor + elif self.shared_experts is not None: + assert shared_output is not None + shared_output *= 1.0 / self.routed_scaling_factor + + if self.shared_experts is not None: + assert shared_output is not None + final_hidden_states += shared_output + + if self.is_sequence_parallel: + final_hidden_states = tensor_model_parallel_all_gather( + final_hidden_states, 0 + ) + final_hidden_states = final_hidden_states[:num_tokens] + elif self.tp_size > 1: + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( + final_hidden_states + ) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +def _get_llama_4_scaling( + original_max_position_embeddings: int, scaling_beta: float, positions: torch.Tensor +) -> torch.Tensor: + scaling = 1 + scaling_beta * torch.log( + 1 + torch.floor(positions / original_max_position_embeddings) + ) + # Broadcast over num_heads and head_dim + return scaling[..., None, None] + + +class AXK1Attention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: AXK1Config, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + max_position_embeddings: int = 8192, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + topk_indices_buffer: torch.Tensor | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_local_heads = num_heads // tp_size + self.scaling = self.qk_head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + assert topk_indices_buffer is None, ( + "topk_indices_buffer is not \ + supported for AXK1Attention" + ) + + if self.q_lora_rank is not None: + self.q_a_proj = ReplicatedLinear( + self.hidden_size, + self.q_lora_rank, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_a_proj", + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + # O projection. + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + + self.rotary_emb = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=False, + ) + + if config.rope_parameters["rope_type"] == "deepseek_yarn": + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False) + scaling_factor = config.rope_parameters["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + + self.attn = Attention( + self.num_local_heads, + self.qk_head_dim, + self.scaling, + num_kv_heads=self.num_local_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ) -> torch.Tensor: + if self.q_lora_rank is not None: + q = self.q_a_proj(hidden_states)[0] + q = self.q_a_layernorm(q) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + latent_cache = latent_cache.unsqueeze(1) + kv_a = self.kv_a_layernorm(kv_a) + kv = self.kv_b_proj(kv_a)[0] + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_pe = latent_cache[:, :, self.kv_lora_rank :] + + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + + q[..., self.qk_nope_head_dim :] = q_pe + k = torch.empty_like(q) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + # Apply llama 4 scaling if provided + if llama_4_scaling is not None: + q *= llama_4_scaling + + # padding value to qk_head_dim for alignment + v = torch.nn.functional.pad( + v, [0, self.qk_head_dim - self.v_head_dim], value=0 + ).view(-1, self.num_local_heads * self.qk_head_dim) + attn_output = self.attn(q, k, v) + attn_output = attn_output.view(-1, self.num_local_heads, self.qk_head_dim)[ + ..., : self.v_head_dim + ].reshape(-1, self.num_local_heads * self.v_head_dim) + output, _ = self.o_proj(attn_output) + return output + + +class AXK1MLAAttention(nn.Module): + """ + Main reference: DeepseekV2 paper, and FlashInfer Implementation + (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551). + + For more info see MLACommonImpl in: + vllm/v1/attention/backends/mla/utils.py + """ + + def __init__( + self, + vllm_config: VllmConfig, + config: AXK1Config, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + max_position_embeddings: int = 8192, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + topk_indices_buffer: torch.Tensor | None = None, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_local_heads = num_heads // tp_size + + self.scaling = self.qk_head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + + if self.q_lora_rank is not None: + self.fused_qkv_a_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + disable_tp=True, + ) + else: + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + + if self.q_lora_rank is not None: + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + + self.rotary_emb = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=False, + ) + + if config.rope_parameters["rope_type"] == "deepseek_yarn": + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False) + scaling_factor = config.rope_parameters["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + + mla_modules = MLAModules( + kv_a_layernorm=self.kv_a_layernorm, + kv_b_proj=self.kv_b_proj, + rotary_emb=self.rotary_emb, + o_proj=self.o_proj, + fused_qkv_a_proj=self.fused_qkv_a_proj + if self.q_lora_rank is not None + else None, + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa + if self.q_lora_rank is None + else None, + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None, + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None, + q_proj=self.q_proj if self.q_lora_rank is None else None, + indexer=None, + indexer_rotary_emb=None, + is_sparse=False, + topk_indices_buffer=topk_indices_buffer, + ) + + self.mla_attn = MultiHeadLatentAttentionWrapper( + self.hidden_size, + self.num_local_heads, + self.scaling, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.q_lora_rank, + self.kv_lora_rank, + mla_modules, + cache_config, + quant_config, + prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ) -> torch.Tensor: + return self.mla_attn(positions, hidden_states, llama_4_scaling) + + +class AXK1DecoderLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + config: AXK1Config | None = None, + ) -> None: + super().__init__() + + if config is None: + config = vllm_config.model_config.hf_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + self.config = config + + self.hidden_size = config.hidden_size + max_position_embeddings = config.max_position_embeddings + # DecoderLayers are created with `make_layers` which passes the prefix + # with the layer's index. + layer_idx = int(prefix.split(sep=".")[-1]) + self.layer_idx = layer_idx + + # verify MLA attention specific fields + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + kv_lora_rank = config.kv_lora_rank + use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)) + self.use_mha = use_mha + + if use_mha: + attn_cls = DeepseekAttention + elif model_config.use_mla: + attn_cls = AXK1MLAAttention + else: + attn_cls = AXK1Attention + self.self_attn = attn_cls( + vllm_config=vllm_config, + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=kv_lora_rank, + max_position_embeddings=max_position_embeddings, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + topk_indices_buffer=None, + ) + + self.is_layer_sparse = self._is_layer_sparse() + if self.is_layer_sparse: + self.mlp = AXK1MoE( + config=config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = AXK1MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_mlp_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.routed_scaling_factor = config.routed_scaling_factor + + def _is_layer_sparse(self) -> bool: + return ( + self.config.n_routed_experts is not None + and self.layer_idx >= self.config.first_k_dense_replace + and self.layer_idx % self.config.moe_layer_freq == 0 + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + llama_4_scaling: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states.clone() + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + attn_kwargs = { + "positions": positions, + "hidden_states": hidden_states, + } + if not self.use_mha: + attn_kwargs["llama_4_scaling"] = llama_4_scaling + hidden_states = self.self_attn(**attn_kwargs) + + if ( + not isinstance(self.self_attn, DeepseekAttention) + and hidden_states.dtype == torch.float16 + ): + # Fix FP16 overflow + # We scale both hidden_states and residual before + # rmsnorm, and rmsnorm result would not affect by scale. + hidden_states *= 1.0 / self.routed_scaling_factor + if self.layer_idx == 0: + # The residual is shared by all layers, we only scale it on + # first layer. + residual *= 1.0 / self.routed_scaling_factor + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + + if self.is_layer_sparse: + hidden_states = self.post_mlp_layernorm(hidden_states) + + if isinstance(self.mlp, AXK1MLP) and hidden_states.dtype == torch.float16: + # Fix FP16 overflow + # Scaling the AXK1MLP output, it is the input of + # input_layernorm of next decoder layer. + # The scaling of AXK1MOE output would be done in the forward + # of AXK1MOE + hidden_states *= 1.0 / self.routed_scaling_factor + + return hidden_states, residual + + +@support_torch_compile +class AXK1Model(nn.Module): + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config: AXK1Config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.device = current_platform.device_type + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: AXK1DecoderLayer(vllm_config, prefix), + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + # Compute llama 4 scaling once per forward pass if enabled + llama_4_scaling_config = getattr(self.config, "llama_4_scaling", None) + llama_4_scaling: torch.Tensor | None + if llama_4_scaling_config is not None: + llama_4_scaling = _get_llama_4_scaling( + original_max_position_embeddings=llama_4_scaling_config[ + "original_max_position_embeddings" + ], + scaling_beta=llama_4_scaling_config["beta"], + positions=positions, + ) + else: + llama_4_scaling = None + + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual = layer( + positions, hidden_states, residual, llama_4_scaling + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class AXK1MixtureOfExperts(MixtureOfExperts): + moe_mlp_layers: list[AXK1MoE] + """ + List of MoE MLP layers in the model. + """ + + def extract_moe_parameters(self, example_moe: AXK1MoE | None): + if example_moe is None: + self.num_moe_layers = 0 + self.num_expert_groups = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_shared_experts = 0 + self.num_redundant_experts = 0 + logger.warning("AXK1: No AXK1MoE layer found in model.layers.") + else: + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_shared_experts = example_moe.n_shared_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for moe in self.moe_mlp_layers: + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + +class AXK1ForCausalLM( + nn.Module, SupportsPP, AXK1MixtureOfExperts, SupportsLoRA, SupportsEagle +): + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + } + model_cls = AXK1Model + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: AXK1Config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + self.use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)) + + if self.use_mha: + self.packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"] + + # `packed_modules_mapping` needs to be modified before + # initializing AXK1Model, as it is passed inplace to + # quantization config init and may be used to select the + # quant_method for relevant layers during initialization. + self.fuse_qkv_a_proj = config.q_lora_rank is not None + if self.fuse_qkv_a_proj: + self.packed_modules_mapping["fused_qkv_a_proj"] = [ + "q_a_proj", + "kv_a_proj_with_mqa", + ] + + self.model = self.model_cls( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + # Set MoE hyperparameters + self.num_moe_layers = ( + self.config.num_hidden_layers - self.config.first_k_dense_replace + ) + self.set_moe_parameters() + + def set_moe_parameters(self): + self.expert_weights = [] + + self.num_expert_groups = getattr(self.config, "n_group", 1) + + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + + assert isinstance(layer, AXK1DecoderLayer) + if isinstance(layer.mlp, AXK1MoE): + # Pick last one layer since the first ones may be dense layers. + example_moe = layer.mlp + self.moe_mlp_layers.append(layer.mlp) + self.moe_layers.append(layer.mlp.experts) + + self.extract_moe_parameters(example_moe) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return SharedFusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts, + num_redundant_experts=0, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + rocm_aiter_moe_shared_expert_enabled = ( + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ) + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + mla_params_mapping = [ + ("fused_qkv_a_proj", "q_a_proj", 0), + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), + ] + mha_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + if self.use_mha: + stacked_params_mapping.extend(mha_params_mapping) + else: + stacked_params_mapping.extend(mla_params_mapping) + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + + ( + self.config.n_shared_experts + if rocm_aiter_moe_shared_expert_enabled + else 0 + ), + num_redundant_experts=self.num_redundant_experts, + ) + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + + is_fusion_moe_shared_experts_layer = ( + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) + ) + + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + if is_fusion_moe_shared_experts_layer: + continue + name_mapped = name.replace(weight_name, param_name) + + # QKV fusion is optional, fall back to normal + # weight loading if it's not enabled + # if go with fusion option, then update name + if ( + param_name == "fused_qkv_a_proj" + ) and name_mapped not in params_dict: + continue + else: + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + + # Special handling: when AITER fusion_shared_experts is enabled, + # checkpoints may provide a single widened shared_experts tensor + # without explicit expert indices + # (e.g. ...mlp.shared_experts.gate_proj.weight). + # For models with multiple shared experts, split that tensor + # evenly into per-shared-expert slices and load them into + # appended expert slots mlp.experts.{n_routed_experts + j}.* + # accordingly. + num_chunks = 1 + if is_fusion_moe_shared_experts_layer: + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1 + # Determine split axis based on op type + # gate/up: ColumnParallel → split along dim 0 + # down: RowParallel → split along dim 1 + split_dim = ( + 1 + if ("down_proj.weight" in name and loaded_weight.ndim > 1) + else 0 + ) + total = loaded_weight.shape[split_dim] + assert total % num_chunks == 0, ( + f"Shared expert weight dim {total} " + f"not divisible by num_chunks {num_chunks}" + ) + chunk_size = total // num_chunks + + for j in range(num_chunks): + chunk_name = name + weight_to_load = loaded_weight + + if is_fusion_moe_shared_experts_layer: + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size) + if loaded_weight.ndim == 1: + weight_to_load = loaded_weight[chunk_slice] + elif split_dim == 0: + weight_to_load = loaded_weight[chunk_slice, :] + else: + weight_to_load = loaded_weight[:, chunk_slice] + # Synthesize an expert-style name so expert mapping + # can route it + chunk_name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts + j}", + ) + + # Use expert_params_mapping to locate the destination + # param and delegate to its expert-aware weight_loader + # with expert_id. + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in chunk_name: + continue + + # Anyway, this is an expert weight and should not be + # attempted to load as other weights later + is_expert_weight = True + + # Do not modify `name` since the loop may continue here + # Instead, create a new variable + name_mapped = chunk_name.replace(weight_name, param_name) + + if is_pp_missing_parameter(name_mapped, self): + continue + + param = params_dict[name_mapped] + # We should ask the weight loader to return success or + # not here since otherwise we may skip experts with + # other available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + weight_to_load, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + if not is_fusion_moe_shared_experts_layer: + name = name_mapped + else: + loaded_params.add(name_mapped) + break + else: + if is_expert_weight: + # We've checked that this is an expert weight + # However it's not mapped locally to this rank + # So we simply skip it + continue + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + if not is_fusion_moe_shared_experts_layer: + loaded_params.add(name) + + return loaded_params + + +def get_spec_layer_idx_from_weight_name( + config: AXK1Config, weight_name: str +) -> int | None: + if config.num_nextn_predict_layers and config.num_nextn_predict_layers > 0: + layer_idx = config.num_hidden_layers + for i in range(config.num_nextn_predict_layers): + if weight_name.startswith(f"model.layers.{layer_idx + i}."): + return layer_idx + i + return None diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index ea0f118a0d2..2ec219d4023 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -213,7 +213,7 @@ class NomicBertModelConfig(VerifyAndUpdateConfig): "Nomic context extension is disabled. " "Changing max_model_len from %s to %s. " "To enable context extension, see: " - "https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/context_extension.html", + "https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/context_extension.py", max_model_len_before, model_config.max_model_len, ) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 768f4e20b34..c3e1ddb7dbb 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -47,7 +47,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.fused_moe import SharedFusedMoE +from vllm.model_executor.layers.fused_moe import GateLinear, SharedFusedMoE from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -221,73 +221,6 @@ class DeepseekV2MLP(nn.Module): return x -class DeepSeekV2Gate(ReplicatedLinear): - def __init__( - self, - hidden_size: int, - n_experts: int, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - assert quant_config is None - super().__init__( - hidden_size, - n_experts, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.gate", - ) - - # Unquantized only, will be called "weight". - assert hasattr(self, "weight") - is_hopper_or_blackwell = current_platform.is_device_capability( - (9, 0) - ) or current_platform.is_device_capability_family(100) - SUPPORTED_NUM_EXPERTS = [256, 384] - SUPPORTED_HIDDEN_SIZES = [7168] - - self.allow_dsv3_router_gemm = ( - current_platform.is_cuda() - and is_hopper_or_blackwell - and n_experts in SUPPORTED_NUM_EXPERTS - and hidden_size in SUPPORTED_HIDDEN_SIZES - ) - - self._out_dtype: torch.dtype | None = None - - def set_out_dtype(self, out_dtype: torch.dtype) -> None: - """ - Set out dtype for the router logits. This is needed after - __init__, b/c we need to check if the trtllm kernel is - selected before we decide between bf16 and fp32. - """ - - if self._out_dtype is not None: - raise ValueError("out_dtype has already been set") - else: - self._out_dtype = out_dtype - - @property - def out_dtype(self) -> torch.dtype: - if self._out_dtype is None: - raise ValueError("out_dtype has not been set yet") - return self._out_dtype - - def forward( - self, - x: torch.Tensor, - ) -> tuple[torch.Tensor, None]: - """ - Use specialized GEMM for low batch size for DSV3 and KIMI. - """ - if self.allow_dsv3_router_gemm and x.shape[0] <= 16: - return ops.dsv3_router_gemm( - hidden_states=x, router_weight=self.weight, output_dtype=self.out_dtype - ), None - else: - return super().forward(x) - - class DeepseekV2MoE(nn.Module): def __init__( self, @@ -316,10 +249,9 @@ class DeepseekV2MoE(nn.Module): "Only silu is supported for now." ) - self.gate = DeepSeekV2Gate( + self.gate = GateLinear( config.hidden_size, config.n_routed_experts, - quant_config=None, prefix=f"{prefix}.gate", ) if getattr(config, "topk_method", None) == "noaux_tc": diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 46cf7fe9782..51b36b1cae3 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -44,6 +44,7 @@ from vllm.model_executor.models.internvl import ( ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM +from vllm.model_executor.models.parakeet import ParakeetExtractor, ProjectedParakeet from vllm.model_executor.models.radio import RadioModel, calc_seq_lens from vllm.model_executor.models.utils import ( init_vllm_registered_model, @@ -55,12 +56,14 @@ from vllm.multimodal.evs import ( compute_retention_mask, ) from vllm.multimodal.inputs import ( + AudioItem, MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems, VideoItem, ) from vllm.multimodal.parse import ( + AudioProcessorItems, ImageEmbeddingItems, ImageProcessorItems, ImageSize, @@ -91,9 +94,29 @@ Image.MAX_IMAGE_PIXELS = None # Disable the limit entirely # Alternative: Set a specific higher limit # Image.MAX_IMAGE_PIXELS = 300000000 # ~300M pixels + +class NanoNemotronVLAudioFeatureInputs(TensorSchema): + """ + Dimensions: + - b: Number of audio clips + - t: Audio feature length + - f: Feature size (mel bins) + """ + + type: Literal["audio_features"] = "audio_features" + input_audio_features: Annotated[torch.Tensor, TensorShape("b", "t", "f")] + feature_attention_mask: Annotated[torch.Tensor, TensorShape("b", "t")] + audio_feature_lengths: Annotated[torch.Tensor, TensorShape("b")] + + +MAX_AUDIO_LEN_S = 10 * 60 # 10 minutes + IMG_START = "" IMG_END = "" IMG_CONTEXT = "" +AUDIO_START = "" +AUDIO_END = "" +AUDIO_CONTEXT = "" # Profiling # MAX_FRAMES = 16 @@ -820,6 +843,11 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): self.video_token = video_token self.video_pruning_rate = video_pruning_rate + self.audio_extractor: ParakeetExtractor | None = None + raw_sound_config = getattr(config, "sound_config", None) + if raw_sound_config is not None: + self.audio_extractor = ParakeetExtractor(raw_sound_config) + # Pre-tokenize special tokens for video processing # to avoid repeated tokenization self._img_start_token_ids = tokenizer.encode( @@ -952,11 +980,53 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): text = [t.replace("