Compare commits

...
Author SHA1 Message Date
Roger WangandOpenAI Codex e329dbcb0f [Build] Update DeepEP pin for low-latency fabric support
Bump the repo's DeepEP pin to the revision that adds the use_fabric low-latency buffer option required for cross-node MNNVL initialization.

This keeps the runtime DeepEP low-latency fix deployable instead of relying on a local DeepEP build newer than the repo default.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-04-18 18:06:19 -07:00
Roger WangandOpenAI Codex 8f3ce1492e [Bugfix] Harden DeepEP low-latency cross-node NVFP4
Pass use_fabric when DeepEP low-latency MNNVL is enabled so cross-node buffer initialization uses the fabric API, and gracefully fall back when the local DeepEP build does not expose low_latency_dispatch(use_nvfp4=...).

This complements the existing NVFP4 all2all backend selection fix for issue #37931.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-04-18 18:05:06 -07:00
Roger WangandOpenAI Codex 21eaa58f71 [Bugfix] Fix NVFP4 CuteDSL batched all2all selection
Promote --moe-backend=flashinfer_cutedsl to the batched NVFP4 backend when the selected all2all path requires batched expert activations, and make flashinfer_nvlink_one_sided prepare/finalize speak the batched expert contract.

Also add regression coverage for the backend selection and one-sided regroup/reduce helpers.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-04-18 17:47:35 -07:00
d0359f3e04 [Bugfix] Guard mxfp4_experts_quant bindings on ENABLE_NVFP4_SM100 (#40191)
Signed-off-by: ultranationalism <www913363043@gmail.com>
Signed-off-by: mgoin <mike.goin12@gmail.com>
Co-authored-by: mgoin <mike.goin12@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-18 13:58:46 -07:00
Dan AlistarhandGitHub ed0622e3a8 [Attention] TurboQuant: remove redundant random signs, add prior art attribution (#40194)
Signed-off-by: Dan Alistarh <d.alistarh@gmail.com>
2026-04-18 14:31:59 -04:00
Yusuf MohammadandGitHub b5f6c5f834 Added general ND x ND matmul and unit test for it (#39909)
Signed-off-by: Yusuf <yusufmohammad@live.com>
2026-04-18 10:05:21 -04:00
Jee Jee LiandGitHub bfde49e287 [DOC] Add fuse_minimax_qk_norm (#39782)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
2026-04-18 00:41:37 -07:00
153ba7f0f3 [Refactor] Drop direct dependency on librosa (#39079)
Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-04-18 06:55:38 +00:00
Chinmay-Kulkarni-AMDandGitHub 87518c3027 [ZenCPU] AMD Zen CPU Backend with supported dtypes via zentorch weekly (#39967)
Signed-off-by: Chinmay Kulkarni <Chinmay.Kulkarni@amd.com>
2026-04-18 06:22:37 +00:00
Rishapveer SinghandGitHub aeee7ef939 [Bugfix] Fix k_proj's bias for GLM-ASR (#40160)
Signed-off-by: Rishapveer Singh <singhrishapveer@gmail.com>
2026-04-17 22:34:33 -07:00
z1yingandGitHub cda19ecf4d [Doc] Fix outdated source reference comment in anthropic/serving.py (#40189)
Signed-off-by: z1ying <tzzying@outlook.com>
2026-04-17 22:31:13 -07:00
80b18230e0 [Frontend] Add multimodal support to /inference/v1/generate endpoint (#38405)
Signed-off-by: Nithin Chalapathi <nithin.ch10@gmail.com>
Signed-off-by: Nithin Chalapathi <nithinc@berkeley.edu>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-04-17 20:31:56 -07:00
z1yingandGitHub d0697cc7b6 [Doc] Add Realtime Transcription section to supported_models.md (#39845)
Signed-off-by: Ziying Tao <tzzying@outlook.com>
2026-04-18 03:26:14 +00:00
milesialGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
b0755523dc [Core] Reduce mm scheduler, get_num_embed overhead (#40143)
Signed-off-by: milesial <milesial@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-04-18 11:25:49 +08:00
Chaojun ZhangGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
993859ceb0 [XPU] fix all_reduce all-zero accuracy issue under torch.compile (#39844)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-04-18 02:33:07 +00:00
Michael GoinandGitHub 48a65ccb02 [CI] Speed up test_fused_marlin_moe (#40178)
Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: Michael Goin <mgoin64@gmail.com>
2026-04-17 19:26:21 -07:00
55842a8d69 [XPU]fake impl for xpu fp8_gemm (#39984)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-04-18 08:53:56 +08:00
Michael GoinandGitHub 1f45e83756 Remove outdated tests test_mixtral_moe and test_duplicated_ignored_sequence_group (#40175)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-04-17 16:49:43 -07:00
Michael GoinandGitHub a8bffaa133 [Kernel] Add MXFP4 W4A4 CUTLASS MoE kernel for SM100 (#37463)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-04-17 16:42:32 -07:00
5cdddddd4a [Kernel] [Helion] Force disable HOP path due to performance regression (#40171)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
Co-authored-by: Claude Sonnet 4 <noreply@anthropic.com>
2026-04-17 17:36:49 -04:00
aditi-amdandGitHub 6ef1efd51f [ROCm] Fix TurboQuant on ROCm: backend routing, flash-attn compat, int64 overflow (#39953)
Signed-off-by: aditi <aditi.rana@amd.com>
2026-04-17 13:08:37 -07:00
Ryan RockandGitHub 58da4ee047 [AMD][CI] Update DeepEP branch (#38396)
Signed-off-by: Ryan Rock <ryan.rock@amd.com>
2026-04-17 14:30:20 -05:00
Andreas KaratzasandGitHub 1ae11e2bfc [ROCm][CI] Build fastsafetensors from source so it links against libamdhip64 (#39978)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-04-17 14:30:08 -05:00
251c18d1f8 skip fp8e4b15 on xpu (#39957)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-04-17 16:55:08 +00:00
512765d52d [Misc][UX] Map mimo reasoning and tooling parsers (#40089)
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-04-17 16:49:21 +00:00
allgatherGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
640cc9dd7d feat: Add LoRA support for Gemma4ForConditionalGeneration (#39291)
Signed-off-by: allgather <all2allops@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-04-17 09:39:19 -07:00
ceade1952c [BugFix] Support custom tool parsers when tool_choice is required and named function (#39870)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Co-authored-by: sfeng33 <4florafeng@gmail.com>
2026-04-17 16:38:10 +00:00
747256bb5d [Bugfix][Core] Fix stuck chunked pipeline parallelism with async scheduling (#38726)
Signed-off-by: Jing Wang <jingwang96@qq.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-04-17 16:02:50 +00:00
81 changed files with 3263 additions and 564 deletions
+1 -1
View File
@@ -2613,6 +2613,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
@@ -3601,7 +3602,6 @@ steps:
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
- pytest -v -s tests/distributed/test_context_parallel.py
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
+1
View File
@@ -141,6 +141,7 @@ steps:
- pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
- pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
- pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
- pytest -v -s tests/kernels/moe/test_flashinfer.py
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
+3 -1
View File
@@ -952,7 +952,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
+9
View File
@@ -134,4 +134,13 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
torch::stable::Tensor& input,
torch::stable::Tensor& input_global_scale);
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
const torch::stable::Tensor& a,
const torch::stable::Tensor& b,
const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets);
#endif
@@ -0,0 +1,468 @@
/*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
* Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
*/
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include <cutlass/arch/arch.h>
#include "cutlass_extensions/common.hpp"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/group_array_problem_shape.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/util/packed_stride.hpp"
#include <cassert>
using namespace cute;
// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
template <typename ElementAB, typename ElementC, typename ElementSF,
typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
__global__ void __mxfp4_get_group_gemm_starts(
ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
ElementAB* a_base_as_int, ElementAB* b_base_as_int,
ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
const int64_t a_stride_val, const int64_t b_stride_val,
const int64_t c_stride_val, const int K, const int N) {
int64_t expert_id = threadIdx.x;
if (expert_id >= gridDim.x * blockDim.x) {
return;
}
int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
int64_t group_size = 32;
int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
"unexpected problem sizes");
int64_t half_k = static_cast<int64_t>(k / 2);
int64_t group_k = static_cast<int64_t>(k / group_size);
// Shape of A as uint8/byte = [M, K // 2]
a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
// Shape of B as uint8/byte = [E, N, K // 2]
b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
// Shape of C = [M, N]
out_offsets[expert_id] = out_base_as_int + expert_offset * n;
// Shape of a_scale = [sum(sf_sizes), K // group_size]
a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
0 &&
"TMA requires 128-byte alignment");
// Shape of B scale = [E, N, K // group_size]
b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
0 &&
"TMA requires 128-byte alignment");
// Initialize strides
a_strides[expert_id] = a_stride_val;
b_strides[expert_id] = b_stride_val;
c_strides[expert_id] = c_stride_val;
LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
*layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
*layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
}
#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \
TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
LayoutSFB, ScaleConfig) \
else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
__mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
LayoutSFB, ScaleConfig> \
<<<1, num_experts, 0, stream>>>( \
static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \
static_cast<C_TYPE**>(out_starts.data_ptr()), \
static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \
static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \
static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \
static_cast<C_TYPE*>(out_tensors.data_ptr()), \
static_cast<SF_TYPE*>(a_scales.data_ptr()), \
static_cast<SF_TYPE*>(b_scales.data_ptr()), \
static_cast<int32_t*>(expert_offsets.data_ptr()), \
static_cast<int32_t*>(sf_offsets.data_ptr()), \
static_cast<int32_t*>(problem_sizes.data_ptr()), \
static_cast<int64_t*>(a_strides.data_ptr()), \
static_cast<int64_t*>(b_strides.data_ptr()), \
static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val, \
b_stride_val, c_stride_val, K, N); \
}
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
void mxfp4_run_get_group_gemm_starts(
const torch::stable::Tensor& a_starts,
const torch::stable::Tensor& b_starts,
const torch::stable::Tensor& out_starts,
const torch::stable::Tensor& a_scales_starts,
const torch::stable::Tensor& b_scales_starts,
const torch::stable::Tensor& layout_sfa,
const torch::stable::Tensor& layout_sfb,
const torch::stable::Tensor& a_strides,
const torch::stable::Tensor& b_strides,
const torch::stable::Tensor& c_strides, int64_t a_stride_val,
int64_t b_stride_val, int64_t c_stride_val,
torch::stable::Tensor const& a_tensors,
torch::stable::Tensor const& b_tensors,
torch::stable::Tensor const& out_tensors,
torch::stable::Tensor const& a_scales,
torch::stable::Tensor const& b_scales,
torch::stable::Tensor const& expert_offsets,
torch::stable::Tensor const& sf_offsets,
torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
int num_experts = (int)expert_offsets.size(0);
auto stream = get_current_cuda_stream(a_tensors.get_device_index());
STD_TORCH_CHECK(out_tensors.size(1) == N,
"Output tensor shape doesn't match expected shape");
STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
"b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
" dimension must match");
if (false) {
}
// MXFP4 uses E8M0 (float_ue8m0_t) scale factors
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
torch::headeronly::ScalarType::BFloat16,
cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
ScaleConfig)
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
torch::headeronly::ScalarType::Half, half,
LayoutSFA, LayoutSFB, ScaleConfig)
else {
STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
}
}
template <typename OutType>
void run_mxfp4_blockwise_scaled_group_mm_sm100(
torch::stable::Tensor& output, const torch::stable::Tensor& a,
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
using ProblemShape =
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
using ElementType = cutlass::float_e2m1_t;
using ElementSFType = cutlass::float_ue8m0_t;
using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementC = OutType;
using ElementD = ElementC;
using ElementAccumulator = float;
// Layout definitions
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::RowMajor;
using LayoutD = LayoutC;
static constexpr int AlignmentA = 32;
static constexpr int AlignmentB = 32;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
// Architecture definitions
using ArchTag = cutlass::arch::Sm100;
using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
using StageCountType = cutlass::gemm::collective::StageCountAuto;
using ClusterShape = Shape<_1, _1, _1>;
struct MMA1SMConfig {
using MmaTileShape = Shape<_128, _128, _128>;
using KernelSchedule =
cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
};
using CollectiveEpilogue =
typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
ClusterShape, Shape<_128, _64>, ElementAccumulator,
ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
LayoutC*, AlignmentD,
typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
using CollectiveMainloop =
typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
ElementB, LayoutB*, AlignmentB, ElementAccumulator,
typename MMA1SMConfig::MmaTileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
sizeof(typename CollectiveEpilogue::SharedStorage))>,
typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
using GemmKernel =
cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
CollectiveEpilogue>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using StrideA = typename Gemm::GemmKernel::InternalStrideA;
using StrideB = typename Gemm::GemmKernel::InternalStrideB;
using StrideC = typename Gemm::GemmKernel::InternalStrideC;
using StrideD = typename Gemm::GemmKernel::InternalStrideD;
using LayoutSFA =
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
using LayoutSFB =
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
using ScaleConfig =
typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
int num_experts = static_cast<int>(expert_offsets.size(0));
torch::stable::Tensor a_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor out_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor a_scales_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_scales_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor layout_sfa = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor layout_sfb = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor a_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor c_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
// Create an instance of the GEMM
Gemm gemm_op;
UnderlyingProblemShape* problem_sizes_as_shapes =
static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
// Set the Scheduler info
cutlass::KernelHardwareInfo hw_info;
using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
PersistentTileSchedulerSm100GroupParams<
typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
scheduler.raster_order = RasterOrderOptions::AlongM;
hw_info.device_id = a.get_device_index();
static std::unordered_map<int, int> cached_sm_counts;
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
cached_sm_counts[hw_info.device_id] =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
hw_info.device_id);
}
hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
// Mainloop Arguments
typename GemmKernel::MainloopArguments mainloop_args{
static_cast<const ElementType**>(a_ptrs.data_ptr()),
static_cast<StrideA*>(a_strides1.data_ptr()),
static_cast<const ElementType**>(b_ptrs.data_ptr()),
static_cast<StrideB*>(b_strides1.data_ptr()),
static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
// Epilogue Arguments
typename GemmKernel::EpilogueArguments epilogue_args{
{}, // epilogue.thread
nullptr,
static_cast<StrideC*>(c_strides1.data_ptr()),
static_cast<ElementD**>(out_ptrs.data_ptr()),
static_cast<StrideC*>(c_strides1.data_ptr())};
auto& fusion_args = epilogue_args.thread;
// Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
fusion_args.alpha_ptr = nullptr;
fusion_args.beta_ptr = nullptr;
fusion_args.alpha = 1.0f;
fusion_args.alpha_ptr_array = nullptr;
fusion_args.dAlpha = {_0{}, _0{}, 0};
fusion_args.beta = 0.0f;
fusion_args.beta_ptr_array = nullptr;
fusion_args.dBeta = {_0{}, _0{}, 0};
// Gemm Arguments
typename GemmKernel::Arguments args{
cutlass::gemm::GemmUniversalMode::kGrouped,
{num_experts, problem_sizes_as_shapes, nullptr},
mainloop_args,
epilogue_args,
hw_info,
scheduler};
size_t workspace_size = Gemm::get_workspace_size(args);
auto workspace =
torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
std::nullopt, a.device());
const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
auto can_implement_status = gemm_op.can_implement(args);
STD_TORCH_CHECK(
can_implement_status == cutlass::Status::kSuccess,
"Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
// Run the GEMM
auto status = gemm_op.initialize(args, workspace.data_ptr());
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
"Failed to initialize MXFP4 GEMM: status=", (int)status,
" workspace_size=", workspace_size,
" num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
status = gemm_op.run(args, workspace.data_ptr(), stream);
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
"Failed to run MXFP4 GEMM");
}
template <typename OutType>
void run_mxfp4_blockwise_scaled_group_mm(
torch::stable::Tensor& output, const torch::stable::Tensor& a,
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
int32_t version_num = get_sm_version_num();
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
if (version_num >= 100 && version_num < 120) {
run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
output, a, b, a_blockscale, b_blockscales, problem_sizes,
expert_offsets, sf_offsets, M, N, K);
return;
}
#endif
STD_TORCH_CHECK_NOT_IMPLEMENTED(
false,
"No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
version_num, ". Required capability: 100");
}
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
// E8M0 scale factors stored as uint8
constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
#endif
#define CHECK_TYPE(x, st, m) \
STD_TORCH_CHECK(x.scalar_type() == st, \
": Inconsistency of torch::stable::Tensor type:", m)
#define CHECK_TH_CUDA(x, m) \
STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
#define CHECK_CONTIGUOUS(x, m) \
STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
#define CHECK_INPUT(x, st, m) \
CHECK_TH_CUDA(x, m); \
CHECK_CONTIGUOUS(x, m); \
CHECK_TYPE(x, st, m)
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
const torch::stable::Tensor& a,
const torch::stable::Tensor& b,
const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets) {
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
// Input validation
CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
// MXFP4 uses E8M0 scale factors (stored as uint8)
CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
STD_TORCH_CHECK(
a_blockscale.dim() == 2,
"expected a_blockscale to be of shape [num_experts, rounded_m,"
" k // group_size], observed rank: ",
a_blockscale.dim())
STD_TORCH_CHECK(b_blockscales.dim() == 3,
"expected b_blockscale to be of shape: "
" [num_experts, n, k // group_size], observed rank: ",
b_blockscales.dim())
STD_TORCH_CHECK(problem_sizes.dim() == 2,
"problem_sizes must be a 2D tensor");
STD_TORCH_CHECK(problem_sizes.size(1) == 3,
"problem_sizes must have the shape (num_experts, 3)");
STD_TORCH_CHECK(
problem_sizes.size(0) == expert_offsets.size(0),
"Number of experts in problem_sizes must match expert_offsets");
STD_TORCH_CHECK(
problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
"problem_sizes must be int32.");
int M = static_cast<int>(a.size(0));
int N = static_cast<int>(b.size(1));
int E = static_cast<int>(b.size(0));
int K = static_cast<int>(2 * b.size(2));
if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
output, a, b, a_blockscale, b_blockscales, problem_sizes,
expert_offsets, sf_offsets, M, N, K);
} else {
run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
output, a, b, a_blockscale, b_blockscales, problem_sizes,
expert_offsets, sf_offsets, M, N, K);
}
#else
STD_TORCH_CHECK_NOT_IMPLEMENTED(
false,
"No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
"SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
#endif
}
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
}
@@ -0,0 +1,432 @@
/*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* MXFP4 activation quantization kernel for MoE experts.
* Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
* over 32-element groups.
*
* Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
* - Block size 32 (2 threads per SF in PACK16 mode)
* - E8M0 (power-of-two) scale factors
* - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
*/
// MXFP4 requires PACK16 mode (16 elements per thread) so that
// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
// Must be defined before any header that (transitively) includes
// nvfp4_utils.cuh.
#define NVFP4_ENABLE_ELTS16 1
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <cuda_fp8.h>
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include "cuda_utils.h"
#include "nvfp4_utils.cuh"
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
#include "launch_bounds_utils.h"
namespace vllm {
// MXFP4 block size constants
static constexpr int MXFP4_SF_VEC_SIZE = 32;
// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
static constexpr int MXFP4_NUM_THREADS_PER_SF =
MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
// MXFP4 quantization kernel for experts.
// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
// SiLU(gate)*up before quantization.
template <class Type, bool FUSE_SILU_MUL = false,
bool SMALL_NUM_EXPERTS = false>
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
fp4_packed_t* out, uint32_t* SFout,
uint32_t* input_offset_by_experts,
uint32_t* output_scale_offset_by_experts,
int n_experts, bool low_latency) {
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
// MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
int32_t const numKTiles = (numCols + 127) / 128;
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
globalIdx += gridDim.x * blockDim.x) {
int rowIdx = globalIdx / colsPerRow;
int colIdx = globalIdx % colsPerRow;
int rowIdx_in_expert = 0;
int expert_idx = 0;
if constexpr (SMALL_NUM_EXPERTS) {
for (int i = 0; i < n_experts; i++) {
uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
if (rowIdx >= current_offset && rowIdx < next_offset) {
rowIdx_in_expert = rowIdx - current_offset;
expert_idx = i;
break;
}
}
} else {
uint32_t local_offsets[17];
for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
*reinterpret_cast<int4*>(local_offsets) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start]));
*reinterpret_cast<int4*>(local_offsets + 4) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start + 4]));
*reinterpret_cast<int4*>(local_offsets + 8) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start + 8]));
*reinterpret_cast<int4*>(local_offsets + 12) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start + 12]));
local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
#pragma unroll
for (int i = 0; i < 16; i++) {
if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
rowIdx_in_expert = rowIdx - local_offsets[i];
expert_idx = chunk_start + i;
break;
}
}
}
}
// Load input and optionally apply fused SiLU+Mul
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
PackedVec quant_input;
if constexpr (FUSE_SILU_MUL) {
PackedVec in_vec_up =
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
quant_input = compute_silu_mul(in_vec, in_vec_up);
} else {
quant_input = in_vec;
}
// In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
int64_t outOffset = rowIdx * colsPerRow + colIdx;
auto& out_pos = out[outOffset];
uint32_t* SFout_in_expert =
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
// Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
auto sf_out =
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
// Block E8M0 scales only; no extra tensor-level scale in this path
constexpr float SFScaleVal = 1.0f;
// UE8M0_SF=true for MXFP4 E8M0 scale factors
out_pos =
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
quant_input, SFScaleVal, sf_out);
}
}
// Large M_topk variant using shared memory for expert offsets
template <class Type, bool FUSE_SILU_MUL = false,
bool SMALL_NUM_EXPERTS = false>
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
fp4_packed_t* out, uint32_t* SFout,
uint32_t* input_offset_by_experts,
uint32_t* output_scale_offset_by_experts,
int n_experts) {
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
// MXFP4: numKTiles = ceil(numCols / 128)
int32_t const numKTiles = (numCols + 127) / 128;
extern __shared__ uint32_t shared_input_offsets[];
if constexpr (SMALL_NUM_EXPERTS) {
for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
shared_input_offsets[i] = input_offset_by_experts[i];
}
} else {
for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
*reinterpret_cast<int4*>(&shared_input_offsets[i]) =
*reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
}
if (threadIdx.x == 0) {
shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
}
}
__syncthreads();
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
globalIdx += gridDim.x * blockDim.x) {
int rowIdx = globalIdx / colsPerRow;
int colIdx = globalIdx % colsPerRow;
int rowIdx_in_expert = 0;
int expert_idx = 0;
// Binary search through experts using shared memory
int left = 0, right = n_experts - 1;
while (left <= right) {
int mid = (left + right) / 2;
uint32_t mid_offset = shared_input_offsets[mid];
uint32_t next_offset = shared_input_offsets[mid + 1];
if (rowIdx >= mid_offset && rowIdx < next_offset) {
rowIdx_in_expert = rowIdx - mid_offset;
expert_idx = mid;
break;
} else if (rowIdx < mid_offset) {
right = mid - 1;
} else {
left = mid + 1;
}
}
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
PackedVec quant_input;
if constexpr (FUSE_SILU_MUL) {
PackedVec in_vec_up =
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
quant_input = compute_silu_mul(in_vec, in_vec_up);
} else {
quant_input = in_vec;
}
int64_t outOffset = rowIdx * colsPerRow + colIdx;
auto& out_pos = out[outOffset];
// MXFP4 has no global scale - only block-level E8M0 scale factors
constexpr float SFScaleVal = 1.0f;
uint32_t* SFout_in_expert =
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
auto sf_out =
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
out_pos =
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
quant_input, SFScaleVal, sf_out);
}
}
template <typename T, bool FUSE_SILU_MUL = false>
void mxfp4_quant_impl(void* output, void* output_scale, void* input,
void* input_offset_by_experts,
void* output_scale_offset_by_experts, int m_topk, int k,
int n_experts, cudaStream_t stream) {
int multiProcessorCount =
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
int const workSizePerRow = k / ELTS_PER_THREAD;
int const totalWorkSize = m_topk * workSizePerRow;
dim3 block(std::min(workSizePerRow, 512));
int const numBlocksPerSM =
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
multiProcessorCount * numBlocksPerSM));
while (grid.x <= multiProcessorCount && block.x > 64) {
grid.x *= 2;
block.x = (block.x + 1) / 2;
}
int const blockRepeat =
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
if (blockRepeat > 1) {
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
if (n_experts >= 4) {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
<<<grid, block, shared_mem_size, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts);
} else {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
<<<grid, block, shared_mem_size, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts);
}
} else {
if (n_experts >= 16) {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
<<<grid, block, 0, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts, /* bool low_latency */ true);
} else {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts, /* bool low_latency */ true);
}
}
}
} // namespace vllm
/*Quantization entry for mxfp4 experts quantization*/
#define CHECK_TH_CUDA(x, m) \
STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x, m) \
STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
#define CHECK_INPUT(x, m) \
CHECK_TH_CUDA(x, m); \
CHECK_CONTIGUOUS(x, m);
constexpr auto HALF = torch::headeronly::ScalarType::Half;
constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
constexpr auto INT = torch::headeronly::ScalarType::Int;
constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
static constexpr int MXFP4_BLOCK_SIZE = 32;
static void validate_mxfp4_experts_quant_inputs(
torch::stable::Tensor const& output,
torch::stable::Tensor const& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts, int64_t m_topk, int64_t k) {
CHECK_INPUT(output, "output");
CHECK_INPUT(output_scale, "output_scale");
CHECK_INPUT(input, "input");
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
STD_TORCH_CHECK(output.dim() == 2);
STD_TORCH_CHECK(output_scale.dim() == 2);
STD_TORCH_CHECK(input.dim() == 2);
STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
// output is uint8 (two mxfp4 values packed into one uint8)
// output_scale is int32 (four E8M0 values packed into one int32)
STD_TORCH_CHECK(output.scalar_type() == UINT8);
STD_TORCH_CHECK(output_scale.scalar_type() == INT);
STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
STD_TORCH_CHECK(output.size(0) == m_topk);
STD_TORCH_CHECK(output.size(1) == k / 2);
int scales_k = k / MXFP4_BLOCK_SIZE;
// K-dimension scale columns padded to a multiple of 4 for swizzle layout
int padded_k = (scales_k + (4 - 1)) / 4 * 4;
// 4 = 4 E8M0 values packed into one int32
STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
}
void mxfp4_experts_quant(
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts) {
auto m_topk = input.size(0);
auto k = input.size(1);
validate_mxfp4_experts_quant_inputs(
output, output_scale, input, input_offset_by_experts,
output_scale_offset_by_experts, n_experts, m_topk, k);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
stream);
});
}
void silu_and_mul_mxfp4_experts_quant(
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts) {
auto m_topk = input.size(0);
auto k_times_2 = input.size(1);
STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
auto k = k_times_2 / 2;
validate_mxfp4_experts_quant_inputs(
output, output_scale, input, input_offset_by_experts,
output_scale_offset_by_experts, n_experts, m_topk, k);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
stream);
});
}
// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied
// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to
// .cpp files and cannot gate the registration from there.
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
m.impl("silu_and_mul_mxfp4_experts_quant",
TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
}
+21 -3
View File
@@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
" Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
// cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
ops.def(
"cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
" Tensor a_blockscale, Tensor b_blockscales,"
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
// Compute NVFP4 block quantized tensor.
ops.def(
"scaled_fp4_quant(Tensor input,"
@@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts) -> ()");
// Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
ops.def(
"mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
"Tensor input, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
// Fused SiLU+Mul+MXFP4 experts quantization.
ops.def(
"silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
"output_scale,"
"Tensor input, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
// Fused SiLU+Mul+NVFP4 quantization.
ops.def(
"silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
@@ -233,9 +252,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("silu_and_mul_scaled_fp4_experts_quant",
TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
// W4A8 ops: impl registrations are in the source files
// (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
// mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
// W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
#endif
}
+1 -1
View File
@@ -330,7 +330,7 @@ WORKDIR /workspace
# 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 DEEPEP_COMMIT_HASH=73b6ea4
ARG DEEPEP_COMMIT_HASH=9249c25
ARG NVSHMEM_VER
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p /tmp/ep_kernels_workspace/dist && \
+7 -32
View File
@@ -192,9 +192,10 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
FROM base AS build_deep
ARG ROCSHMEM_BRANCH="ba0bf0f3"
ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
ARG DEEPEP_BRANCH="e84464ec"
ARG DEEPEP_BRANCH="5d90af8b"
ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
ARG DEEPEP_NIC="cx7"
ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
ENV ROCSHMEM_DIR=/opt/rocshmem
RUN git clone ${ROCSHMEM_REPO} \
@@ -202,13 +203,11 @@ RUN git clone ${ROCSHMEM_REPO} \
&& git checkout ${ROCSHMEM_BRANCH} \
&& mkdir -p projects/rocshmem/build \
&& cd projects/rocshmem/build \
&& cmake .. \
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-DROCM_PATH=/opt/rocm \
-DCMAKE_POSITION_INDEPENDENT_CODE=ON \
-DUSE_EXTERNAL_MPI=OFF \
&& make -j \
&& make install
&& bash ../scripts/build_configs/all_backends \
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-DROCM_PATH=/opt/rocm \
-DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-DUSE_EXTERNAL_MPI=OFF
# Build DeepEP wheel.
# DeepEP looks for rocshmem at ROCSHMEM_DIR.
@@ -262,30 +261,6 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
&& echo "Detected vLLM version: ${VLLM_VERSION}" \
&& echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
# Fail if git-based package dependencies are found in requirements files
# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
RUN echo "Checking for git-based packages in requirements files..." \
&& echo "Checking common.txt for git-based packages:" \
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
echo "ERROR: Git-based packages found in common.txt:"; \
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
echo "Please publish these packages to PyPI instead of using git dependencies."; \
exit 1; \
else \
echo " ✓ No git-based packages found in common.txt"; \
fi \
&& echo "Checking rocm.txt for git-based packages:" \
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
echo "ERROR: Git-based packages found in rocm.txt:"; \
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
echo "Please publish these packages to PyPI instead of using git dependencies."; \
exit 1; \
else \
echo " ✓ No git-based packages found in rocm.txt"; \
fi \
&& echo "All requirements files are clean - no git-based packages found"
# Pin vLLM dependencies to exact versions of custom ROCm wheels
# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
+1 -1
View File
@@ -53,7 +53,7 @@
"default": "cuda"
},
"DEEPEP_COMMIT_HASH": {
"default": "73b6ea4"
"default": "9249c25"
},
"GIT_REPO_CHECK": {
"default": "0"
+3 -3
View File
@@ -193,7 +193,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:
The API server takes care of basic audio I/O and optional chunking before building prompts:
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `librosa`.
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
- Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
@@ -206,8 +206,8 @@ Relevant server logic:
async def _preprocess_speech_to_text(...):
language = self.model_cls.validate_language(request.language)
...
y, sr = librosa.load(bytes_, sr=self.asr_config.sample_rate)
duration = librosa.get_duration(y=y, sr=sr)
y, sr = load_audio(bytes_, sr=self.asr_config.sample_rate)
duration = get_audio_duration(y=y, sr=sr)
do_split_audio = (self.asr_config.allow_audio_chunking
and duration > self.asr_config.max_audio_clip_s)
chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
+34
View File
@@ -21,6 +21,7 @@ or just on the low or high end.
| Fusion | `PassConfig` flag | Fused operations | Default at | E2E Speedup | Fullgraph | `num_tokens` |
| ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
| [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low |
| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm) | `fuse_minimax_qk_norm` | Q/K variance all-reduce → Q/K RMSNorm | Off by default | 2-3% | No | Low |
| [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always |
| [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always |
| [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low |
@@ -40,6 +41,7 @@ The table below lists the quantization schemes supported by each fusion on each
| Fusion | SM100 (Blackwell) | SM90 (Hopper) | SM89 (Ada) | SM80 (Ampere) | ROCm |
| ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
@@ -54,6 +56,9 @@ The table below lists the quantization schemes supported by each fusion on each
fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
for per-backend details.
\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
`enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
@@ -184,6 +189,35 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
- Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
### MiniMax QK Norm (`fuse_minimax_qk_norm`)
!!! info
This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
optimization level.
**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
per-token Q/K variances before applying RMS normalization to Q and K.
This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
Example:
```bash
vllm serve MiniMaxAI/MiniMax-M2.5 \
--tensor-parallel-size 4 \
--compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
```
**Code locations.**
- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
### Sequence Parallelism (`enable_sp`)
**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
+3 -3
View File
@@ -300,12 +300,12 @@ Full example: [examples/offline_inference/audio_language.py](../../examples/offl
Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech.
```python
import librosa
from vllm import LLM, SamplingParams
from vllm.multimodal.audio import split_audio
from vllm.multimodal.media.audio import load_audio
# Load long audio file
audio, sr = librosa.load("long_audio.wav", sr=16000)
audio, sr = load_audio("long_audio.wav", sr=16000)
# Split into chunks at low-energy (quiet) regions
chunks = split_audio(
@@ -832,7 +832,7 @@ Then, you can use the OpenAI client as follows:
base_url=openai_api_base,
)
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
audio_url = AudioAsset("winning_call").url
audio_base64 = encode_base64_content_from_url(audio_url)
+18
View File
@@ -682,6 +682,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
!!! note
`VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.
#### Realtime Transcription
Speech models that support streaming transcription via the
[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api)
WebSocket endpoint.
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | |
| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | |
!!! note
`VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`.
`Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`.
You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'`
when serving.
## Pooling Models
See [this page](pooling_models/README.md) for more information on how to use pooling models.
+1 -1
View File
@@ -60,7 +60,7 @@ We currently support the following OpenAI APIs:
- [Translation API](#translations-api) (`/v1/audio/translations`)
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
- [Realtime API](#realtime-api) (`/v1/realtime`)
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription).
In addition, we have the following custom APIs:
@@ -0,0 +1,117 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Disaggregated multimodal serving: render → generate round-trip.
Demonstrates the two-phase disaggregated flow:
1. /v1/chat/completions/render preprocesses a multimodal chat request
into token IDs and serialized tensor features.
2. /inference/v1/generate runs inference on the preprocessed tokens.
The render response is passed *directly* to generate with only
``sampling_params`` added, showing that the two endpoints compose with
zero client-side transformation.
Launch the server first:
vllm serve Qwen/Qwen3-VL-2B-Instruct \
--dtype bfloat16 --max-model-len 4096 --enforce-eager
Then run this script:
python example_mm_serve.py
"""
import io
import pybase64 as base64
import requests
from PIL import Image
from transformers import AutoTokenizer
BASE_URL = "http://localhost:8000"
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
def make_data_url(image: Image.Image) -> str:
"""Encode a PIL image as a base64 data URL."""
buf = io.BytesIO()
image.save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode()
return f"data:image/png;base64,{b64}"
def main():
# -- Step 1: Create a test image (solid red) -------------------------
image = Image.new("RGB", (224, 224), color=(255, 0, 0))
data_url = make_data_url(image)
print("Created 224x224 red test image")
# -- Step 2: Render (preprocess) -------------------------------------
render_payload = {
"model": MODEL_NAME,
"messages": [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": data_url}},
{
"type": "text",
"text": "What color is this image? Answer in one word.",
},
],
}
],
}
print("\n--- Render ---")
render_resp = requests.post(
f"{BASE_URL}/v1/chat/completions/render", json=render_payload
)
render_resp.raise_for_status()
render_data = render_resp.json()
print(f"Response keys: {list(render_data.keys())}")
print(f"Number of token_ids: {len(render_data['token_ids'])}")
features = render_data.get("features")
if features and features.get("kwargs_data"):
print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
for modality, items in features["kwargs_data"].items():
print(
f" {modality}: {len(items)} item(s), "
f"first item type: {type(items[0])} length: {len(items[0])}"
if items
else "First item: (empty)"
)
else:
print("WARNING: no kwargs_data in render response")
# -- Step 3: Generate (inference) ------------------------------------
# Pass the render output directly — only add sampling_params.
generate_payload = render_data
generate_payload["sampling_params"] = {
"max_tokens": 20,
"temperature": 0.0,
}
print("\n--- Generate ---")
gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
gen_resp.raise_for_status()
gen_data = gen_resp.json()
# -- Step 4: Decode & print ------------------------------------------
output_ids = gen_data["choices"][0]["token_ids"]
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
text = tokenizer.decode(output_ids, skip_special_tokens=True)
print(f"Output token count: {len(output_ids)}")
print(f"Generated text: {text!r}")
if "red" in text.lower():
print("\nModel correctly identified the red image.")
else:
print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
if __name__ == "__main__":
main()
@@ -267,7 +267,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
{
"type": "input_audio",
"input_audio": {
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
"data": audio_base64,
"format": "wav",
},
@@ -292,7 +292,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
{
"type": "audio_url",
"audio_url": {
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
"url": audio_url
},
},
@@ -316,7 +316,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
{
"type": "audio_url",
"audio_url": {
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
"url": f"data:audio/ogg;base64,{audio_base64}"
},
},
@@ -12,7 +12,6 @@ model, for example:
Requirements:
- vllm with audio support
- websockets
- librosa
- numpy
The script:
@@ -26,12 +25,12 @@ import argparse
import asyncio
import json
import librosa
import numpy as np
import pybase64 as base64
import websockets
from vllm.assets.audio import AudioAsset
from vllm.multimodal.media.audio import load_audio
def audio_to_pcm16_base64(audio_path: str) -> str:
@@ -39,7 +38,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str:
Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
"""
# Load audio and resample to 16kHz mono
audio, _ = librosa.load(audio_path, sr=16000, mono=True)
audio, _ = load_audio(audio_path, sr=16000, mono=True)
# Convert to PCM16
pcm16 = (audio * 32767).astype(np.int16)
# Encode as base64
+1 -1
View File
@@ -22,4 +22,4 @@ timm>=1.0.17
# To be consistent with test_quark.py
amd-quark>=0.8.99
# Required for faster safetensors model loading
fastsafetensors >= 0.2.2
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2
+1 -1
View File
@@ -55,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
numba==0.61.2 # Required for N-gram speculative decoding
numpy
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
instanttensor>=0.1.5
pydantic>=2.12 # 2.11 leads to error on python 3.13
decord==0.6.0
+1 -1
View File
@@ -277,7 +277,7 @@ fastar==0.10.0
# via fastapi-cloud-cli
fastparquet==2026.3.0
# via genai-perf
fastsafetensors==0.2.2
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
# via
# -c requirements/rocm.txt
# -r requirements/test/rocm.in
+3 -1
View File
@@ -1085,7 +1085,9 @@ setup(
install_requires=get_requirements(),
extras_require={
# AMD Zen CPU optimizations via zentorch
"zen": ["zentorch"],
"zen": [
"zentorch-weekly==5.2.1.dev20260408"
], # Zentorch has weekly releases. This pulls the known-good version.
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
@@ -13,7 +13,6 @@ import io
import time
from statistics import mean, median
import librosa
import pytest
import soundfile
import torch
@@ -21,6 +20,7 @@ from datasets import load_dataset
from evaluate import load
from transformers.models.whisper.english_normalizer import EnglishTextNormalizer
from vllm.multimodal.audio import get_audio_duration
from vllm.tokenizers import get_tokenizer
from ....models.registry import HF_EXAMPLE_MODELS
@@ -84,7 +84,7 @@ async def process_dataset(model, client, data, concurrent_request):
trust_remote_code=model_info.trust_remote_code,
)
# Warmup call as the first `librosa.load` server-side is quite slow.
# Warmup call as the first `load_audio` server-side is quite slow.
audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"]
_ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "")
@@ -118,7 +118,7 @@ def print_performance_metrics(results, total_time):
def add_duration(sample):
y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000
sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000
return sample
@@ -5,7 +5,6 @@ import asyncio
import json
import warnings
import librosa
import numpy as np
import pybase64 as base64
import pytest
@@ -14,6 +13,7 @@ import websockets
from tests.entrypoints.openai.conftest import add_attention_backend
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
from vllm.assets.audio import AudioAsset
from vllm.multimodal.media.audio import load_audio
# Increase engine iteration timeout for ROCm where first-use JIT compilation
# can exceed the default 60s, causing a silent deadlock in feed_tokens.
@@ -56,7 +56,7 @@ async def send_event(ws, event: dict) -> None:
def mary_had_lamb_audio_chunks() -> list[str]:
"""Audio split into ~1 second chunks for streaming."""
path = AudioAsset("mary_had_lamb").get_local_path()
audio, _ = librosa.load(str(path), sr=16000, mono=True)
audio, _ = load_audio(str(path), sr=16000, mono=True)
# Split into ~0.1 second chunks (1600 samples at 16kHz)
chunk_size = 1600
@@ -6,7 +6,6 @@ import asyncio
import io
import json
import librosa
import numpy as np
import openai
import pytest
@@ -14,6 +13,7 @@ import pytest_asyncio
import soundfile as sf
from tests.utils import RemoteOpenAIServer
from vllm.multimodal.media.audio import load_audio
from vllm.platforms import current_platform
MODEL_NAME = "openai/whisper-large-v3-turbo"
@@ -134,7 +134,7 @@ async def test_bad_requests(mary_had_lamb, whisper_client):
@pytest.mark.asyncio
async def test_long_audio_request(mary_had_lamb, whisper_client):
mary_had_lamb.seek(0)
audio, sr = librosa.load(mary_had_lamb)
audio, sr = load_audio(mary_had_lamb)
# Add small silence after each audio for repeatability in the split process
audio = np.pad(audio, (0, 1600))
repeated_audio = np.tile(audio, 10)
@@ -7,7 +7,6 @@ import io
import json
import httpx
import librosa
import numpy as np
import openai
import pytest
@@ -17,6 +16,7 @@ import soundfile as sf
from tests.entrypoints.openai.conftest import add_attention_backend
from tests.utils import RemoteOpenAIServer
from vllm.logger import init_logger
from vllm.multimodal.media.audio import load_audio
logger = init_logger(__name__)
@@ -264,7 +264,7 @@ async def test_long_audio_request(foscolo, client_and_model):
if model_name == "google/gemma-3n-E2B-it":
pytest.skip("Gemma3n does not support long audio requests")
foscolo.seek(0)
audio, sr = librosa.load(foscolo)
audio, sr = load_audio(foscolo)
repeated_audio = np.tile(audio, 2)
# Repeated audio to buffer
buffer = io.BytesIO()
+111
View File
@@ -0,0 +1,111 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Roundtrip tests for multimodal serde used by the disagg generate endpoint."""
import torch
from vllm.entrypoints.serve.disagg.mm_serde import (
decode_mm_kwargs_item,
encode_mm_kwargs_item,
)
from vllm.entrypoints.serve.disagg.protocol import (
MultiModalFeatures,
PlaceholderRangeInfo,
)
from vllm.multimodal.inputs import (
MultiModalBatchedField,
MultiModalFieldElem,
MultiModalFlatField,
MultiModalKwargsItem,
MultiModalSharedField,
)
def test_mm_kwargs_item_roundtrip():
"""Full roundtrip test with all three field types and multiple dtypes."""
e1 = MultiModalFieldElem(
data=torch.zeros(1000, dtype=torch.bfloat16),
field=MultiModalBatchedField(),
)
e2 = MultiModalFieldElem(
data=torch.ones(100, dtype=torch.int32),
field=MultiModalSharedField(batch_size=4),
)
e3 = MultiModalFieldElem(
data=torch.randn(20, dtype=torch.float32),
field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0),
)
item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3})
encoded = encode_mm_kwargs_item(item)
# Encoded result is a base64 string
assert isinstance(encoded, str)
decoded = decode_mm_kwargs_item(encoded)
assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"}
assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data)
assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data)
assert torch.equal(item["embeds"].data, decoded["embeds"].data)
assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField)
assert isinstance(decoded["grid_thw"].field, MultiModalSharedField)
assert isinstance(decoded["embeds"].field, MultiModalFlatField)
def test_mm_kwargs_item_none_data():
"""Roundtrip with None data field."""
elem = MultiModalFieldElem(
data=None,
field=MultiModalSharedField(batch_size=2),
)
item = MultiModalKwargsItem({"empty": elem})
encoded = encode_mm_kwargs_item(item)
decoded = decode_mm_kwargs_item(encoded)
assert decoded["empty"].data is None
assert isinstance(decoded["empty"].field, MultiModalSharedField)
def test_mm_kwargs_item_nested_tensors():
"""Roundtrip with nested tensor data."""
nested = [torch.randn(3, 4), torch.randn(5, 4)]
elem = MultiModalFieldElem(
data=nested,
field=MultiModalBatchedField(),
)
item = MultiModalKwargsItem({"nested": elem})
encoded = encode_mm_kwargs_item(item)
decoded = decode_mm_kwargs_item(encoded)
decoded_data = decoded["nested"].data
assert len(decoded_data) == 2
assert torch.equal(nested[0], decoded_data[0])
assert torch.equal(nested[1], decoded_data[1])
def test_mm_features_with_kwargs_data():
"""Test that MultiModalFeatures can carry serialized tensor data."""
elem = MultiModalFieldElem(
data=torch.randn(5, 3, dtype=torch.float32),
field=MultiModalBatchedField(),
)
item = MultiModalKwargsItem({"pixel_values": elem})
encoded = encode_mm_kwargs_item(item)
features = MultiModalFeatures(
mm_hashes={"image": ["abc123"]},
mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]},
kwargs_data={"image": [encoded]},
)
# JSON roundtrip
json_str = features.model_dump_json()
features2 = MultiModalFeatures.model_validate_json(json_str)
assert features2.mm_hashes == {"image": ["abc123"]}
assert features2.kwargs_data is not None
assert len(features2.kwargs_data["image"]) == 1
decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
assert torch.equal(elem.data, decoded["pixel_values"].data)
@@ -0,0 +1,158 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for multimodal features through the /inference/v1/generate endpoint.
Mirrors test_serving_tokens.py but exercises the multimodal piping
using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render ->
/generate -> /detokenize path. Intentionally avoids running the HF
processor in the pytest parent process to keep os.fork() in sibling
tests (e.g. test_weight_transfer_llm.py) deadlock-free.
"""
import os
import httpx
import pytest
import pytest_asyncio
from PIL import Image
from tests.utils import RemoteOpenAIServer
from vllm.multimodal.utils import encode_image_url
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
GEN_ENDPOINT = "/inference/v1/generate"
RENDER_ENDPOINT = "/v1/chat/completions/render"
DETOKENIZE_ENDPOINT = "/detokenize"
@pytest.fixture(scope="module")
def test_image():
return Image.new("RGB", (224, 224), color=(255, 0, 0))
@pytest.fixture(scope="module")
def server():
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"4096",
"--enforce-eager",
"--no-enable-prefix-caching",
]
envs = os.environ.copy()
envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client(server: RemoteOpenAIServer):
transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
async with httpx.AsyncClient(
transport=transport,
base_url=server.url_root,
timeout=600,
headers=headers,
) as c:
yield c
@pytest.mark.asyncio
async def test_render_to_generate_roundtrip(client, test_image):
"""End-to-end: render a multimodal chat -> feed into generate -> decode.
All preprocessing and detokenization happens in the server subprocess;
the pytest parent never imports transformers or touches torch tensors.
"""
data_url = encode_image_url(test_image, format="PNG")
render_payload = {
"model": MODEL_NAME,
"messages": [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": data_url}},
{
"type": "text",
"text": "What color is this image? Answer in one word.",
},
],
}
],
}
render_resp = await client.post(RENDER_ENDPOINT, json=render_payload)
render_resp.raise_for_status()
render_data = render_resp.json()
# Validate render output structure: keys exist and values are non-empty
# and well-typed.
assert "token_ids" in render_data
assert isinstance(render_data["token_ids"], list)
assert len(render_data["token_ids"]) > 0
assert all(isinstance(t, int) for t in render_data["token_ids"])
assert "features" in render_data
features = render_data["features"]
assert features is not None
assert isinstance(features, dict)
assert "mm_hashes" in features
assert "image" in features["mm_hashes"]
image_hashes = features["mm_hashes"]["image"]
assert isinstance(image_hashes, list)
assert len(image_hashes) > 0
assert all(isinstance(h, str) and h for h in image_hashes)
assert "mm_placeholders" in features
assert "image" in features["mm_placeholders"]
image_placeholders = features["mm_placeholders"]["image"]
assert isinstance(image_placeholders, list)
assert len(image_placeholders) > 0
for p in image_placeholders:
assert isinstance(p.get("offset"), int)
assert isinstance(p.get("length"), int)
assert p["length"] > 0
assert "kwargs_data" in features
assert "image" in features["kwargs_data"]
assert len(features["kwargs_data"]["image"]) > 0
# Build generate request from render output
generate_payload = render_data
generate_payload["sampling_params"] = {
"max_tokens": 10,
"temperature": 0.0,
}
gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload)
gen_resp.raise_for_status()
gen_data = gen_resp.json()
assert "choices" in gen_data
assert isinstance(gen_data["choices"], list)
assert len(gen_data["choices"]) >= 1
choice = gen_data["choices"][0]
assert "token_ids" in choice
assert isinstance(choice["token_ids"], list)
assert len(choice["token_ids"]) > 0
assert all(isinstance(t, int) for t in choice["token_ids"])
detok_resp = await client.post(
DETOKENIZE_ENDPOINT,
json={"model": MODEL_NAME, "tokens": choice["token_ids"]},
)
detok_resp.raise_for_status()
detok_data = detok_resp.json()
assert "prompt" in detok_data
text = detok_data["prompt"]
assert isinstance(text, str)
assert len(text) > 0
assert "red" in text.lower(), (
f"Expected model to identify the red image, got: {text!r}"
)
+45 -166
View File
@@ -14,8 +14,6 @@ import pytest
import torch
from torch.nn import Parameter
from torch.nn import functional as F
from transformers import MixtralConfig
from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
import vllm.model_executor.layers.fused_moe # noqa
from tests.kernels.moe.utils import (
@@ -24,10 +22,7 @@ from tests.kernels.moe.utils import (
modular_triton_fused_moe,
)
from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
from vllm._aiter_ops import rocm_aiter_ops
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.distributed.parallel_state import init_distributed_environment
from vllm.forward_context import get_forward_context, set_forward_context
from vllm.model_executor.layers.fused_moe import (
MoEActivation,
fused_topk,
@@ -56,12 +51,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
marlin_quantize,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
from vllm.model_executor.models.mixtral import MixtralMoE
from vllm.platforms import current_platform
from vllm.scalar_type import ScalarType, scalar_types
from vllm.utils.math_utils import next_power_of_2
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.worker.workspace import init_workspace_manager
def iterative_moe(
@@ -150,12 +143,14 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
{
"a_type": [scalar_types.bfloat16],
"b_type": scalar_types.float4_e2m1f,
"c_type": [scalar_types.bfloat16],
"group_blocks": [2],
},
# MXFP8
{
"a_type": [scalar_types.bfloat16],
"b_type": scalar_types.float8_e4m3fn,
"c_type": [scalar_types.bfloat16],
"group_blocks": [2],
},
# AWQ-INT4 with INT8 activation
@@ -681,154 +676,35 @@ def test_fused_moe_wn16(
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("padding", [True, False])
@pytest.mark.parametrize(
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
)
@torch.inference_mode()
def test_mixtral_moe(
default_vllm_config,
dist_init,
dtype: torch.dtype,
padding: bool,
use_rocm_aiter: bool,
monkeypatch,
):
"""Make sure our Mixtral MoE implementation agrees with the one from
huggingface."""
# Explicitly set AITER env var based on test parameter to ensure
# consistent behavior regardless of external environment
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
rocm_aiter_ops.refresh_env_variables()
if use_rocm_aiter and dtype == torch.float32:
pytest.skip("AITER ROCm test skip for float32")
monkeypatch.setenv("RANK", "0")
monkeypatch.setenv("LOCAL_RANK", "0")
monkeypatch.setenv("WORLD_SIZE", "1")
monkeypatch.setenv("MASTER_ADDR", "localhost")
monkeypatch.setenv("MASTER_PORT", "12345")
init_distributed_environment()
init_workspace_manager(torch.accelerator.current_device_index())
# Instantiate our and huggingface's MoE blocks
vllm_config.compilation_config.static_forward_context = dict()
with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
config = MixtralConfig()
hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
vllm_moe = MixtralMoE(
num_experts=config.num_local_experts,
top_k=config.num_experts_per_tok,
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
params_dtype=dtype,
tp_size=1,
dp_size=1,
).cuda()
# Load the weights
vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
if isinstance(hf_moe.experts, torch.nn.ModuleList):
# Transformers v4
for i in range(config.num_local_experts):
weights = (
hf_moe.experts[i].w1.weight.data,
hf_moe.experts[i].w3.weight.data,
)
vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
else:
# Transformers v5
vllm_moe.experts.w13_weight.data[:] = hf_moe.experts.gate_up_proj.data
vllm_moe.experts.w2_weight.data[:] = hf_moe.experts.down_proj.data
# TODO: remove this line after https://github.com/huggingface/transformers/pull/43622
hf_moe.experts.config._experts_implementation = "eager"
# Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda")
# vLLM uses 1D query [num_tokens, hidden_dim]
vllm_inputs = hf_inputs.flatten(0, 1)
# Pad the weight if moe padding is enabled
if padding:
vllm_moe.experts.w13_weight = Parameter(
F.pad(vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[
..., 0:-128
],
requires_grad=False,
)
vllm_moe.experts.w2_weight = Parameter(
F.pad(vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],
requires_grad=False,
)
torch.accelerator.synchronize()
torch.accelerator.empty_cache()
# FIXME (zyongye) fix this after we move self.kernel
# assignment in FusedMoE.__init__
vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)
# need to override the forward context for unittests, otherwise it assumes
# we're running the model forward pass (the model specified in vllm_config)
get_forward_context().all_moe_layers = None
# Run forward passes for both MoE blocks
hf_states = hf_moe.forward(hf_inputs)
if isinstance(hf_states, tuple):
# Transformers v4
hf_states = hf_states[0]
vllm_states = vllm_moe.forward(vllm_inputs)
mixtral_moe_tol = {
torch.float32: 1e-3,
torch.float16: 1e-3,
torch.bfloat16: 1e-2,
}
if use_rocm_aiter:
# The values of rtol and atol are set based on the tests in ROCM AITER package.
# https://github.com/ROCm/aiter/blob/dfed377f4be7da96ca2d75ac0761f569676f7240/op_tests/test_moe.py#L174
torch.testing.assert_close(
hf_states.flatten(0, 1), vllm_states, rtol=0.01, atol=100
)
else:
torch.testing.assert_close(
hf_states.flatten(0, 1),
vllm_states,
rtol=mixtral_moe_tol[dtype],
atol=mixtral_moe_tol[dtype],
)
MARLIN_MOE_SCENARIOS = [
# (m, n, k, e, topk, ep_size, act_order, is_k_full)
# No act_order: is_k_full=True matches usual case (marlin_is_k_full).
# N>=256 required for Marlin kernel thread config for MXFP8.
# Single token, small matrices
(1, 128, 256, 5, 2, 1, False, True),
# Single token, large matrices
(1, 1024, 2048, 5, 2, 1, False, True),
# Unaligned m, small matrices
(133, 256, 256, 5, 2, 1, False, True),
# Unaligned m, large matrices
(133, 1024, 2048, 12, 3, 1, False, True),
# Aligned batch, small matrices
(128, 256, 256, 5, 2, 1, False, True),
# Aligned batch, large matrices
(128, 1024, 2048, 12, 3, 1, False, True),
# Expert parallelism
(64, 1024, 2048, 12, 3, 4, False, True),
# Act order with is_k_full=True (no tensor parallelism)
(1, 1024, 2048, 5, 2, 1, True, True),
# Act order with is_k_full=False (tensor parallelism)
(133, 256, 256, 5, 2, 1, True, False),
]
def marlin_moe_generate_valid_test_cases():
import itertools
m_list = [1, 123, 666]
n_list = [128, 1024]
k_list = [256, 2048]
e_list = [5, 12]
topk_list = [2, 3]
ep_size_list = [1, 4]
act_order_list = [True, False]
is_k_full_list = [True, False]
all_combinations = itertools.product(
MOE_MARLIN_QUANT_TEST_CONFIGS,
m_list,
n_list,
k_list,
e_list,
topk_list,
ep_size_list,
act_order_list,
is_k_full_list,
)
def is_invalid(
def is_valid(
a_type,
b_type,
c_type,
@@ -845,39 +721,42 @@ def marlin_moe_generate_valid_test_cases():
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
if group_size > 0 and k % group_size != 0:
return False
if act_order and group_size in [-1, k, n]:
return False
if group_size in [k, n]:
return False
if not act_order and is_k_full:
if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
return False
return a_type.size_bits < 16 or a_type is c_type
cases = []
for case in all_combinations:
quant_test_config, m, n, k, _, _, _, act_order, *_ = case
if act_order and not quant_test_config.get("support_act_order", False):
continue
for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS:
f16_types = [scalar_types.float16]
inner_combinations = itertools.product(
quant_test_config.get("a_type", f16_types),
[quant_test_config["b_type"]],
quant_test_config.get("c_type", f16_types),
quant_test_config["group_blocks"],
inner_combinations = list(
itertools.product(
quant_test_config.get("a_type", f16_types),
[quant_test_config["b_type"]],
quant_test_config.get("c_type", f16_types),
quant_test_config["group_blocks"],
)
)
supports_act_order = quant_test_config.get("support_act_order", False)
for sub_case in inner_combinations:
if (
sub_case[0] == scalar_types.float8_e4m3fn
and current_platform.get_device_capability() not in [89, 120]
):
continue
args = sub_case + (m, n, k) + case[4:]
if is_invalid(*args):
cases.append(args)
for scenario in MARLIN_MOE_SCENARIOS:
m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
if act_order and not supports_act_order:
continue
args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
if is_valid(*args):
cases.append(args)
return cases
+248
View File
@@ -0,0 +1,248 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels."""
import random
import pytest
import torch
from tests.kernels.utils import torch_moe_single
from vllm import _custom_ops as ops
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
random.seed(42)
set_random_seed(42)
MXFP4_BLOCK_SIZE = 32
def align(val: int, alignment: int = 128) -> int:
return int((val + alignment - 1) // alignment * alignment)
def calc_diff(x, y):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
def is_sm100_supported() -> bool:
return current_platform.is_cuda() and current_platform.is_device_capability_family(
100
)
def compute_ref_output(
input_tensor: torch.Tensor,
weight_list: list[torch.Tensor],
expert_offsets: list[int],
expert_offset: int,
num_experts: int,
) -> torch.Tensor:
"""Reference output using torch_moe_single with top-1 routing."""
score = torch.full(
(expert_offset, num_experts),
-1e9,
device=input_tensor.device,
dtype=torch.float32,
)
for g in range(num_experts):
start = expert_offsets[g]
end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
score[start:end, g] = 0.0
return torch_moe_single(
input_tensor, torch.stack(weight_list, dim=0), score, topk=1
)
@pytest.mark.skipif(
not is_sm100_supported(),
reason="cutlass_mxfp4_group_mm requires CUDA SM100",
)
@pytest.mark.parametrize("num_experts", [8, 16, 32])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16])
def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
"""
Test the MXFP4 grouped GEMM kernel by:
1. Creating random per-expert inputs and weights
2. Quantizing both to MXFP4 using the CUDA kernel
3. Running the CUTLASS grouped GEMM
4. Comparing against BF16 reference
"""
device = "cuda"
alignment = 128
# N and K must be multiples of 128 for clean swizzle layout
n_g = random.randint(1, 16) * alignment
k_g = random.randint(1, 16) * alignment
expert_offset = 0
expert_offsets_input = []
problem_sizes = []
input_list = []
weight_list = []
for g in range(num_experts):
m_g = random.randint(1, 256)
expert_offsets_input.append(expert_offset)
expert_offset += m_g
problem_sizes.append([m_g, n_g, k_g])
input_list.append(
torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype)
)
weight_list.append(
torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype)
)
input_tensor = torch.concat(input_list, dim=0) # [M_total, K]
# --- Quantize INPUTS via mxfp4_experts_quant ---
input_bs_offsets = []
tot = 0
for g in range(num_experts):
input_bs_offsets.append(tot)
tot += align(problem_sizes[g][0], 128)
input_bs_offsets.append(tot)
_inp_expert_offsets = torch.tensor(
expert_offsets_input + [expert_offset], device=device, dtype=torch.int32
)
_inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32)
input_quant, input_sf = ops.mxfp4_experts_quant(
input_tensor,
_inp_expert_offsets,
_inp_bs_offsets,
num_experts,
topk=1,
)
# --- Quantize WEIGHTS via mxfp4_experts_quant ---
# Treat each expert's N weight rows as an "expert" with N tokens
weight_tensor = torch.concat(weight_list, dim=0) # [E*N, K]
weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
# N is always multiple of 128, so blockscale offsets are clean
weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
_wt_expert_offsets = torch.tensor(
weight_expert_offsets, device=device, dtype=torch.int32
)
_wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32)
weight_quant, weight_sf = ops.mxfp4_experts_quant(
weight_tensor,
_wt_expert_offsets,
_wt_bs_offsets,
num_experts,
topk=1,
)
# Reshape weight quantized data to [E, N, K//2]
weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2)
# Reshape weight scale factors to [E, N, K//32]
# The quant kernel produces uint8 SF buffer. Each row has K//32 SFs.
scales_per_row = k_g // MXFP4_BLOCK_SIZE
weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row]
weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row)
# Output
output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
_problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32)
_expert_offsets = torch.tensor(
expert_offsets_input, device=device, dtype=torch.int32
)
_input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32)
# Run the MXFP4 grouped GEMM
ops.cutlass_mxfp4_moe_mm(
output,
input_quant,
weight_quant,
input_sf,
weight_sf_3d,
_problem_sizes,
_expert_offsets,
_input_bs,
)
# Reference: BF16 matmul
ref_output = compute_ref_output(
input_tensor=input_tensor,
weight_list=weight_list,
expert_offsets=expert_offsets_input,
expert_offset=expert_offset,
num_experts=num_experts,
)
# Compare per-expert
for g in range(num_experts):
start = expert_offsets_input[g]
end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset
if start == end:
continue
baseline = ref_output[start:end]
actual = output[start:end]
diff = calc_diff(actual, baseline)
print(
f"m_g={end - start} n_g={n_g} k_g={k_g} "
f"num_experts={num_experts}, "
f"out_dtype={out_dtype}, diff={diff:.5f}"
)
# FP4 quantization is very lossy (~4 bits precision)
# Comparing quantized vs full-precision gives cosine diff of 0.05-0.15
assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold"
@pytest.mark.skipif(
not is_sm100_supported(),
reason="mxfp4_experts_quant requires CUDA SM100",
)
def test_mxfp4_experts_quant_basic():
"""
Basic smoke test for the MXFP4 experts quantization kernel.
"""
device = "cuda"
num_experts = 4
k = 256
tokens_per_expert = 16
total_tokens = tokens_per_expert * num_experts
input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5
expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)]
blockscale_offsets = [
align(i * tokens_per_expert, 128) for i in range(num_experts + 1)
]
_expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32)
_blockscale_offsets = torch.tensor(
blockscale_offsets, device=device, dtype=torch.int32
)
output, output_sf = ops.mxfp4_experts_quant(
input_tensor,
_expert_offsets,
_blockscale_offsets,
num_experts,
topk=1,
)
assert output.shape == (total_tokens, k // 2)
assert output.dtype == torch.uint8
assert output_sf.dtype == torch.uint8
assert output.any(), "Quantized output is all zeros"
print(
f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}"
)
print("PASSED")
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
@@ -0,0 +1,202 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import patch
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.moe.utils import make_dummy_moe_config
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
NvFp4MoeBackend,
select_nvfp4_moe_backend,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import ( # noqa: E501
_group_rank_batched_inputs_by_local_expert,
_reduce_local_expert_outputs_to_rank_batched_payload,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
kNvfp4Dynamic,
kNvfp4Static,
)
class _StandardNvFp4Kernel:
@staticmethod
def is_supported_config(
_cls,
moe_config,
weight_key,
activation_key,
activation_format,
):
if activation_format == mk.FusedMoEActivationFormat.Standard:
return True, None
return False, f"{activation_format.value} activation format"
class _BatchedNvFp4Kernel:
@staticmethod
def is_supported_config(
_cls,
moe_config,
weight_key,
activation_key,
activation_format,
):
if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
return True, None
return False, f"{activation_format.value} activation format"
class _UnsupportedNvFp4Kernel:
@staticmethod
def is_supported_config(
_cls,
moe_config,
weight_key,
activation_key,
activation_format,
):
return False, "unsupported"
def _make_nvfp4_config(all2all_backend: str):
moe_config = make_dummy_moe_config(num_experts=8, hidden_dim=16)
moe_config.moe_backend = "flashinfer_cutedsl"
moe_config.moe_parallel_config.dp_size = 2
moe_config.moe_parallel_config.use_ep = True
moe_config.moe_parallel_config.all2all_backend = all2all_backend
return moe_config
def _fake_backend_to_kernel_cls(backend: NvFp4MoeBackend):
if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
return [_StandardNvFp4Kernel]
if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED:
return [_BatchedNvFp4Kernel]
return [_UnsupportedNvFp4Kernel]
@patch(
"vllm.model_executor.layers.fused_moe.oracle.nvfp4.backend_to_kernel_cls",
side_effect=_fake_backend_to_kernel_cls,
)
def test_select_nvfp4_backend_uses_standard_cutedsl_for_standard_all2all(
mock_backend_to_kernel_cls,
):
moe_config = _make_nvfp4_config("allgather_reducescatter")
backend, experts_cls = select_nvfp4_moe_backend(
moe_config,
kNvfp4Static,
kNvfp4Dynamic,
)
assert backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
assert experts_cls is _StandardNvFp4Kernel
@patch(
"vllm.model_executor.layers.fused_moe.oracle.nvfp4.backend_to_kernel_cls",
side_effect=_fake_backend_to_kernel_cls,
)
def test_select_nvfp4_backend_promotes_cutedsl_for_batched_all2all(
mock_backend_to_kernel_cls,
):
moe_config = _make_nvfp4_config("flashinfer_nvlink_one_sided")
backend, experts_cls = select_nvfp4_moe_backend(
moe_config,
kNvfp4Static,
kNvfp4Dynamic,
)
assert moe_config.moe_parallel_config.use_batched_activation_format
assert backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
assert experts_cls is _BatchedNvFp4Kernel
def test_flashinfer_one_sided_rank_batched_regroup_and_reduce():
hidden_states = torch.tensor(
[
[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0], [0.0, 0.0]],
[[40.0, 41.0], [50.0, 51.0], [0.0, 0.0], [0.0, 0.0]],
]
)
hidden_scales = torch.tensor(
[
[[1.0], [2.0], [3.0], [0.0]],
[[4.0], [5.0], [0.0], [0.0]],
]
)
topk_ids = torch.tensor(
[
[[2, 0], [3, 2], [1, 0], [0, 0]],
[[3, 1], [2, 3], [0, 0], [0, 0]],
],
dtype=torch.int32,
)
topk_weights = torch.tensor(
[
[[0.70, 0.30], [0.40, 0.60], [0.50, 0.50], [0.0, 0.0]],
[[0.80, 0.20], [0.25, 0.75], [0.0, 0.0], [0.0, 0.0]],
],
dtype=torch.float32,
)
(
batched_hidden_states,
batched_hidden_scales,
expert_tokens_meta,
dispatch_metadata,
) = _group_rank_batched_inputs_by_local_expert(
hidden_states,
hidden_scales,
topk_ids,
topk_weights,
num_local_experts=2,
first_local_expert=2,
num_dispatchers=2,
runtime_max_tokens_per_rank=4,
source_num_tokens=[3, 2],
)
assert expert_tokens_meta.expert_num_tokens.tolist() == [3, 3]
torch.testing.assert_close(
batched_hidden_states[0, :3],
torch.tensor([[10.0, 11.0], [20.0, 21.0], [50.0, 51.0]]),
)
torch.testing.assert_close(
batched_hidden_states[1, :3],
torch.tensor([[20.0, 21.0], [40.0, 41.0], [50.0, 51.0]]),
)
assert batched_hidden_scales is not None
torch.testing.assert_close(
batched_hidden_scales[0, :3],
torch.tensor([[1.0], [2.0], [5.0]]),
)
torch.testing.assert_close(
batched_hidden_scales[1, :3],
torch.tensor([[2.0], [4.0], [5.0]]),
)
fused_expert_output = torch.tensor(
[
[[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [0.0, 0.0]],
[[10.0, 10.0], [20.0, 20.0], [30.0, 30.0], [0.0, 0.0]],
]
)
combine_payload = _reduce_local_expert_outputs_to_rank_batched_payload(
fused_expert_output,
dispatch_metadata,
num_dispatchers=2,
runtime_max_tokens_per_rank=4,
apply_router_weight_on_input=False,
)
expected = torch.zeros((2, 4, 2))
expected[0, 0] = torch.tensor([0.70, 0.70])
expected[0, 1] = torch.tensor([5.20, 5.20])
expected[1, 0] = torch.tensor([16.0, 16.0])
expected[1, 1] = torch.tensor([23.25, 23.25])
torch.testing.assert_close(combine_payload, expected)
@@ -4,7 +4,6 @@
import os
from collections.abc import Sequence
import librosa
import pytest
import regex as re
from huggingface_hub import snapshot_download
@@ -14,6 +13,7 @@ from vllm.assets.image import ImageAsset
from vllm.logprobs import SampleLogprobs
from vllm.lora.request import LoRARequest
from vllm.multimodal.image import convert_image_mode, rescale_image_size
from vllm.multimodal.media.audio import load_audio
from ....conftest import (
IMAGE_ASSETS,
@@ -290,7 +290,7 @@ def test_vision_speech_models(
num_logprobs: int,
) -> None:
# use the example speech question so that the model outputs are reasonable
audio = librosa.load(speech_question, sr=None)
audio = load_audio(speech_question, sr=None)
image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
inputs_vision_speech = [
@@ -4,11 +4,11 @@
from collections.abc import Sequence
from typing import Any
import librosa
import pytest
from transformers import AutoModelForSpeechSeq2Seq
from vllm.assets.audio import AudioAsset
from vllm.multimodal.audio import AudioResampler
from vllm.platforms import current_platform
from ....conftest import HfRunner, PromptAudioInput, VllmRunner
@@ -93,13 +93,12 @@ def run_test(
def resampled_assets() -> list[tuple[Any, int]]:
audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
sampled_assets = []
resampler = AudioResampler(target_sr=WHISPER_SAMPLE_RATE)
for asset in audio_assets:
audio, orig_sr = asset.audio_and_sample_rate
# Resample to Whisper's expected sample rate (16kHz)
if orig_sr != WHISPER_SAMPLE_RATE:
audio = librosa.resample(
audio, orig_sr=orig_sr, target_sr=WHISPER_SAMPLE_RATE
)
audio = resampler.resample(audio, orig_sr=orig_sr)
sampled_assets.append(
(audio, WHISPER_SAMPLE_RATE),
)
+2 -2
View File
@@ -3,12 +3,12 @@
from pathlib import Path
from unittest.mock import patch
import librosa
import numpy as np
import pybase64 as base64
import pytest
from vllm.multimodal.media import AudioMediaIO
from vllm.multimodal.media.audio import load_audio
from ...conftest import AudioTestAssets
@@ -73,6 +73,6 @@ def test_audio_media_io_from_video(video_assets):
video_path = video_assets[0].video_path
with open(video_path, "rb") as f:
audio, sr = audio_io.load_bytes(f.read())
audio_ref, sr_ref = librosa.load(video_path, sr=None)
audio_ref, sr_ref = load_audio(video_path, sr=None)
assert sr == sr_ref
np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
+3 -6
View File
@@ -26,11 +26,8 @@ def test_placeholder_range_get_num_embeds(is_embed, expected):
"is_embed,expected",
[
(None, None),
(
torch.tensor([False, True, False, True, True]),
torch.tensor([0, 1, 1, 2, 3]),
),
(torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
(torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]),
(torch.tensor([True, True, True]), [1, 2, 3]),
],
)
def test_placeholder_range_embeds_cumsum(is_embed, expected):
@@ -41,6 +38,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected):
assert pr.embeds_cumsum is None
return
assert torch.equal(pr.embeds_cumsum, expected)
assert pr.embeds_cumsum == expected
# cached_property should return the same object on repeated access
assert pr.embeds_cumsum is pr.embeds_cumsum
+28 -54
View File
@@ -18,9 +18,7 @@ from vllm.model_executor.layers.quantization.turboquant.config import (
TQ_PRESETS,
TurboQuantConfig,
)
from vllm.model_executor.layers.quantization.turboquant.quantizer import (
generate_wht_signs,
)
from vllm.platforms import current_platform
from vllm.utils.math_utils import next_power_of_2
# ============================================================================
@@ -345,7 +343,8 @@ class TestLloydMax:
# Rotation matrix tests (GPU required)
# ============================================================================
CUDA_AVAILABLE = torch.cuda.is_available()
GPGPU_AVAILABLE = torch.cuda.is_available() or torch.xpu.is_available()
DEVICE_TYPE = current_platform.device_type
def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
@@ -360,16 +359,16 @@ def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Te
return Q.to(device)
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
class TestRotationMatrix:
"""Tests for the QR-based rotation (standalone benchmarks only)."""
@pytest.mark.parametrize("dim", [64, 96, 128, 256])
def test_rotation_matrix_shape_and_orthogonal(self, dim):
Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
Pi = generate_rotation_matrix(dim, seed=42, device=DEVICE_TYPE)
assert Pi.shape == (dim, dim)
eye = Pi @ Pi.T
assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
f"Pi not orthogonal for dim={dim}"
)
@@ -385,13 +384,13 @@ class TestRotationMatrix:
def test_rotation_matrix_det_is_pm1(self):
"""Orthogonal matrix determinant must be +1 or -1."""
Pi = generate_rotation_matrix(128, seed=42, device="cuda")
Pi = generate_rotation_matrix(128, seed=42, device=DEVICE_TYPE)
det = torch.linalg.det(Pi)
assert abs(abs(det.item()) - 1.0) < 1e-4
# ============================================================================
# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard)
# Hadamard rotation tests (serving path: _build_hadamard)
# ============================================================================
@@ -403,58 +402,34 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
return (H / math.sqrt(d)).to(torch.device(device))
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
class TestWHTRotation:
"""Tests for the WHT rotation actually used in serving."""
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
class TestHadamardRotation:
"""Tests for the Hadamard rotation used in serving."""
@pytest.mark.parametrize("dim", [64, 128, 256])
def test_wht_orthonormal(self, dim):
"""signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
signs = generate_wht_signs(dim, seed=42, device="cuda")
H = _build_hadamard(dim, "cuda")
PiT = (signs.unsqueeze(1) * H).contiguous()
eye = PiT @ PiT.T
assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
f"WHT rotation not orthonormal for dim={dim}"
def test_hadamard_orthonormal(self, dim):
"""H must be orthonormal: H @ H^T = I."""
H = _build_hadamard(dim, DEVICE_TYPE)
eye = H @ H.T
assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
f"Hadamard not orthonormal for dim={dim}"
)
@pytest.mark.parametrize("dim", [64, 128, 256])
def test_wht_self_inverse(self, dim):
"""PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
signs = generate_wht_signs(dim, seed=42, device="cuda")
H = _build_hadamard(dim, "cuda")
PiT = (signs.unsqueeze(1) * H).contiguous()
Pi = PiT.T.contiguous()
# Pi @ PiT should be identity (rotation then inverse)
result = Pi @ PiT
assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
f"WHT rotation not self-inverse for dim={dim}"
def test_hadamard_symmetric(self, dim):
"""Sylvester Hadamard must be symmetric: H = H^T."""
H = _build_hadamard(dim, DEVICE_TYPE)
assert torch.allclose(H, H.T, atol=1e-6), (
f"Hadamard not symmetric for dim={dim}"
)
def test_wht_signs_deterministic(self):
"""Same seed must produce identical signs."""
s1 = generate_wht_signs(128, seed=42)
s2 = generate_wht_signs(128, seed=42)
assert torch.equal(s1, s2)
def test_wht_signs_different_seeds(self):
"""Different seeds must produce different signs."""
s1 = generate_wht_signs(128, seed=42)
s2 = generate_wht_signs(128, seed=99)
assert not torch.equal(s1, s2)
def test_wht_signs_are_pm1(self):
"""All sign values must be exactly +1 or -1."""
signs = generate_wht_signs(128, seed=42)
assert torch.all(signs.abs() == 1.0)
# ============================================================================
# Store → Decode round-trip test (GPU + Triton required)
# ============================================================================
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
class TestStoreDecodeRoundTrip:
"""End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
@@ -487,13 +462,12 @@ class TestStoreDecodeRoundTrip:
block_size = 16
num_blocks = 1
device = torch.device("cuda")
device = torch.device(DEVICE_TYPE)
# Generate rotation
signs = generate_wht_signs(D, seed=42, device=device)
H = _build_hadamard(D, "cuda")
PiT = (signs.unsqueeze(1) * H).contiguous().float()
Pi = PiT.T.contiguous()
# Pure Hadamard rotation (symmetric: H = H^T, so Pi = PiT = H)
H = _build_hadamard(D, DEVICE_TYPE)
PiT = H
Pi = H
# Generate centroids
centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
-16
View File
@@ -17,22 +17,6 @@ from vllm import LLM, SamplingParams
from vllm.platforms import current_platform
@pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len")
def test_duplicated_ignored_sequence_group():
"""https://github.com/vllm-project/vllm/issues/1655"""
sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=256)
llm = LLM(
model="distilbert/distilgpt2",
max_num_batched_tokens=4096,
tensor_parallel_size=1,
)
prompts = ["This is a short prompt", "This is a very long prompt " * 1000]
outputs = llm.generate(prompts, sampling_params=sampling_params)
assert len(prompts) == len(outputs)
@pytest.mark.parametrize(
"model",
[
+41 -2
View File
@@ -7,17 +7,39 @@ from tests.models.utils import check_embeddings_close
from vllm.utils.serial_utils import (
EMBED_DTYPES,
ENDIANNESS,
MM_METADATA_DTYPES,
EmbedDType,
Endianness,
MmMetadataDType,
binary2tensor,
tensor2binary,
)
FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys())
INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys())
def _build_integer_tensor(
embed_dtype: MmMetadataDType, shape: tuple[int, ...]
) -> torch.Tensor:
torch_dtype = MM_METADATA_DTYPES[embed_dtype].torch_dtype
if torch_dtype is torch.bool:
return torch.randint(0, 2, shape, dtype=torch.int32).to(torch.bool)
if torch_dtype is torch.uint8:
return torch.randint(0, 256, shape, dtype=torch.uint8)
if torch_dtype is torch.int32:
return torch.randint(-(2**20), 2**20, shape, dtype=torch.int32)
if torch_dtype is torch.int64:
return torch.randint(-(2**62), 2**62, shape, dtype=torch.int64)
raise AssertionError(f"Unsupported non-floating embed dtype: {embed_dtype}")
@pytest.mark.parametrize("endianness", ENDIANNESS)
@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys())
@pytest.mark.parametrize("embed_dtype", FLOAT_EMBED_DTYPES)
@torch.inference_mode()
def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
def test_encode_and_decode_floats(embed_dtype: EmbedDType, endianness: Endianness):
for i in range(10):
tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
shape = tensor.shape
@@ -40,3 +62,20 @@ def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
name_1="new",
tol=1e-2,
)
@pytest.mark.parametrize("endianness", ENDIANNESS)
@pytest.mark.parametrize("embed_dtype", INTEGER_EMBED_DTYPES)
@torch.inference_mode()
def test_encode_and_decode_integers(
embed_dtype: MmMetadataDType, endianness: Endianness
):
shape = (2, 3, 5, 7, 11, 13)
for i in range(10):
tensor = _build_integer_tensor(embed_dtype, shape)
binary = tensor2binary(tensor, embed_dtype, endianness)
new_tensor = binary2tensor(binary, shape, embed_dtype, endianness)
assert new_tensor.dtype == MM_METADATA_DTYPES[embed_dtype].torch_dtype
torch.testing.assert_close(tensor, new_tensor, atol=0, rtol=0)
@@ -0,0 +1,105 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Test batch-invariant matmul against torch.matmul for various shape combinations.
Tests correctness (matches torch.matmul) and batch invariance (result for one
item doesn't change based on other items in the batch).
"""
import pytest
import torch
from utils import skip_unsupported
from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant
from vllm.platforms import current_platform
DEVICE_TYPE = current_platform.device_type
@skip_unsupported
@pytest.mark.parametrize(
"a_shape,b_shape",
[
# 2D x 2D
((32, 64), (64, 16)),
# 2D x 3D
((64, 16), (4, 16, 32)),
# 3D x 2D
((4, 32, 64), (64, 16)),
# 4D x 2D
((1, 4, 32, 64), (64, 16)),
# 3D x 3D
((4, 32, 64), (4, 64, 16)),
# 3D x 4D
((2, 32, 64), (1, 2, 64, 16)),
# 4D x 3D (Gemma4 pattern)
((1, 2, 32, 64), (2, 64, 16)),
# 4D x 4D
((1, 2, 32, 64), (4, 2, 64, 16)),
# 2D x 4D
((32, 64), (1, 2, 64, 16)),
# 2D x 5D
((32, 64), (1, 2, 2, 64, 16)),
# 5D x 2D
((1, 2, 2, 32, 64), (64, 16)),
# 5D x 5D
((1, 2, 4, 32, 64), (1, 2, 4, 64, 16)),
],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_matmul_correctness(a_shape, b_shape, dtype):
"""
Compare matmul_batch_invariant against torch.matmul for various shapes.
"""
device = torch.device(DEVICE_TYPE)
torch.manual_seed(42)
a = torch.rand(a_shape, dtype=dtype, device=device)
b = torch.rand(b_shape, dtype=dtype, device=device)
# Standard implementation (CUDA ops)
standard_output = torch.matmul(a, b)
# Batch-invariant implementation (Triton)
triton_output = matmul_batch_invariant(a, b)
# Compare outputs
# Use looser tolerance for bfloat16 due to its lower precision
if dtype == torch.bfloat16:
rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16
else:
rtol, atol = 1e-2, 1e-2 # 1% for float16/float32
torch.testing.assert_close(
triton_output,
standard_output,
rtol=rtol,
atol=atol,
msg=f"matmul mismatch for a ndim={a.ndim}, b ndim={b.ndim},",
)
@skip_unsupported
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_matmul_batch_invariance(dtype):
"""
Verify that the result for one item is bitwise identical regardless
of what other items are in the batch.
"""
device = torch.device(DEVICE_TYPE)
torch.manual_seed(42)
a_single = torch.rand((1, 64, 32), dtype=dtype, device=device)
b = torch.rand((32, 128), dtype=dtype, device=device)
standard_output = matmul_batch_invariant(a_single, b)
a_batch = torch.rand((8, 64, 32), dtype=dtype, device=device)
a_batch[3] = a_single[0]
batch_output = matmul_batch_invariant(a_batch, b)
batch_output_a = batch_output[3]
assert torch.equal(standard_output[0], batch_output_a)
+1 -1
View File
@@ -8,7 +8,7 @@ set -ex
# --nvshmem-ver <ver> NVSHMEM version
CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"9249c25"}
NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"} # Default supports both CUDA 12 and 13
WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
MODE=${MODE:-install}
+135
View File
@@ -1150,6 +1150,38 @@ def cutlass_fp4_moe_mm(
)
def cutlass_mxfp4_moe_mm(
out_tensors: torch.Tensor,
a_tensors: torch.Tensor,
b_tensors: torch.Tensor,
a_scales: torch.Tensor,
b_scales: torch.Tensor,
problem_sizes: torch.Tensor,
expert_offsets: torch.Tensor,
sf_offsets: torch.Tensor,
):
"""
An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
- a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
- a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
- Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
- expert_offsets/sf_offsets: expert boundary indices
- problem_sizes: (num_experts, 3) with (M, N, K) per expert
"""
return torch.ops._C.cutlass_mxfp4_group_mm(
out_tensors,
a_tensors,
b_tensors,
a_scales,
b_scales,
problem_sizes,
expert_offsets,
sf_offsets,
)
def mxfp8_experts_quant(
input_tensor: torch.Tensor,
problem_sizes: torch.Tensor,
@@ -1848,6 +1880,109 @@ def silu_and_mul_scaled_fp4_experts_quant(
return output, output_scales
def mxfp4_experts_quant(
input_tensor: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
n_experts: int,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Quantize input tensor to MXFP4 for packed MoE inputs.
Uses 32-element blocks with E8M0 (power-of-two) scale factors.
MXFP4 has no global scale - only block-level E8M0 scale factors.
Args:
input_tensor: [m_topk, k] BF16/FP16 activations
expert_offsets: [n_experts+1] token boundaries per expert
blockscale_offsets: [n_experts+1] SF row boundaries per expert
n_experts: number of experts
topk: number of top-k experts
Returns:
output: [m_topk, k//2] packed E2M1 values (uint8)
output_scales: E8M0 blockscales in swizzled layout (uint8 view)
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k = input_tensor.shape
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
f"{MAX_TOKENS_PER_EXPERT})"
f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
)
scales_k = k // 32
padded_k = (scales_k + (4 - 1)) // 4
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.mxfp4_experts_quant(
output,
output_scales,
input_tensor,
expert_offsets,
blockscale_offsets,
n_experts,
)
# E8M0 SFs are stored as uint8
output_scales = output_scales.view(torch.uint8)
return output, output_scales
def silu_and_mul_mxfp4_experts_quant(
input_tensor: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
n_experts: int,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
MXFP4 has no global scale - only block-level E8M0 scale factors.
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k_times_2 = input_tensor.shape
assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
k = k_times_2 // 2
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
scales_k = k // 32
padded_k = (scales_k + (4 - 1)) // 4
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.silu_and_mul_mxfp4_experts_quant(
output,
output_scales,
input_tensor,
expert_offsets,
blockscale_offsets,
n_experts,
)
output_scales = output_scales.view(torch.uint8)
return output, output_scales
# fp8
def scaled_fp8_quant(
input: torch.Tensor,
+17
View File
@@ -22,6 +22,23 @@ else:
except ImportError:
from torch.library import impl_abstract as register_fake
if hasattr(torch.ops._xpu_C, "fp8_gemm"):
@register_fake("_xpu_C::fp8_gemm")
def _fp8_gemm_fake(
q_input: torch.Tensor,
q_weight: torch.Tensor,
out_dtype: torch.dtype,
input_scales: torch.Tensor,
weight_scale: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
input_2d = q_input.view(-1, q_input.shape[-1])
M = input_2d.size(0)
N = q_weight.size(1)
return torch.empty((M, N), dtype=out_dtype, device=q_input.device)
if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"):
@register_fake("_xpu_C::fp8_gemm_w8a16")
@@ -414,9 +414,11 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
num_qps_per_rank=num_qps_per_rank,
)
if not current_platform.is_rocm():
use_mnnvl = envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL
kwargs.update(
allow_nvlink_for_low_latency_mode=True,
allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
allow_mnnvl=use_mnnvl,
use_fabric=use_mnnvl,
explicitly_destroy=True,
)
return kwargs
@@ -47,9 +47,10 @@ class XpuCommunicator(DeviceCommunicatorBase):
self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
logger.info("Using AgRs manager on XPU device.")
def all_reduce(self, input_) -> torch.Tensor:
dist.all_reduce(input_, group=self.device_group)
return input_
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
output = input_.clone() if torch.compiler.is_compiling() else input_
dist.all_reduce(output, group=self.device_group)
return output
def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
world_size = self.world_size
+1 -1
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from
# https://github.com/vllm/vllm/entrypoints/openai/serving_chat.py
# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/chat_completion/serving.py
"""Anthropic Messages API serving handler"""
@@ -557,6 +557,20 @@ class OpenAIServingChat(OpenAIServing):
and self._should_stream_with_auto_tool_parsing(request)
)
# Determine whether required/named tool_choice should fall back to
# the auto tool_parser path instead of the standard JSON-based parsing.
# This happens when the parser declares supports_required_and_named=False
# (e.g. GLM models that output XML instead of JSON).
tool_choice_uses_parser = (
self.tool_parser is not None
and not self.tool_parser.supports_required_and_named
and request.tools
and (
request.tool_choice == "required"
or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
)
)
all_previous_token_ids: list[list[int]] | None
function_name_returned = [False] * num_choices
if self.tool_call_id_type == "kimi_k2":
@@ -569,7 +583,12 @@ class OpenAIServingChat(OpenAIServing):
# Only one of these will be used, thus previous_texts and
# all_previous_token_ids will not be used twice in the same iteration.
if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
if (
is_mistral_grammar_path
or tool_choice_auto
or tool_choice_uses_parser
or reasoning_parser
):
# These are only required in "auto" tool choice case
all_previous_token_ids = [[] for _ in range(num_choices)]
reasoning_end_arr = [False] * num_choices
@@ -764,7 +783,12 @@ class OpenAIServingChat(OpenAIServing):
delta_message: DeltaMessage | None
# just update previous_texts and previous_token_ids
if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
if (
is_mistral_grammar_path
or tool_choice_auto
or tool_choice_uses_parser
or reasoning_parser
):
assert previous_texts is not None
assert all_previous_token_ids is not None
previous_text = previous_texts[i]
@@ -813,7 +837,9 @@ class OpenAIServingChat(OpenAIServing):
if result.tools_called:
tools_streamed[i] = True
# handle streaming deltas for tools with named tool_choice
elif tool_choice_function_name:
# Skip when tool_choice_uses_parser so it falls through
# to the auto tool_parser branches below.
elif tool_choice_function_name and not tool_choice_uses_parser:
# When encountering think end id in prompt_token_ids
# i.e {"enable_thinking": False},
# check BEFORE calling the parser to avoid a spurious
@@ -851,7 +877,6 @@ class OpenAIServingChat(OpenAIServing):
):
reasoning_end_arr[i] = True
if delta_message and delta_message.content:
# This need to be added to next `delta_text`
current_text = delta_message.content
delta_message.content = None
else:
@@ -896,7 +921,12 @@ class OpenAIServingChat(OpenAIServing):
)
tools_streamed[i] = True
elif request.tool_choice == "required":
# Skip when tool_choice_uses_parser so it falls through
# to the auto tool_parser branches below.
elif (
request.tool_choice == "required"
and not tool_choice_uses_parser
):
assert previous_texts is not None
previous_text = previous_texts[i]
current_text = previous_text + delta_text
@@ -966,7 +996,10 @@ class OpenAIServingChat(OpenAIServing):
# update the previous values for the next iteration
if (
is_mistral_grammar_path or tool_choice_auto or reasoning_parser
is_mistral_grammar_path
or tool_choice_auto
or tool_choice_uses_parser
or reasoning_parser
) and not self.use_harmony:
assert previous_texts is not None
assert all_previous_token_ids is not None
+26 -5
View File
@@ -627,7 +627,7 @@ class OpenAIServing:
and isinstance(request.tool_choice, ToolChoiceFunction)
):
assert content is not None
# Forced Function Call
# Forced Function Call (Responses API)
function_calls.append(
FunctionCall(name=request.tool_choice.name, arguments=content)
)
@@ -636,14 +636,20 @@ class OpenAIServing:
not use_mistral_tool_parser
and request.tool_choice
and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
):
# Named function with standard JSON-based parsing
assert content is not None
# Forced Function Call
function_calls.append(
FunctionCall(name=request.tool_choice.function.name, arguments=content)
)
content = None # Clear content since tool is called.
elif not use_mistral_tool_parser and request.tool_choice == "required":
elif (
not use_mistral_tool_parser
and request.tool_choice == "required"
and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
):
# "required" with standard JSON-based parsing
tool_calls = []
with contextlib.suppress(ValidationError):
content = content or ""
@@ -662,15 +668,30 @@ class OpenAIServing:
use_mistral_tool_parser
or (
enable_auto_tools
and (request.tool_choice == "auto" or request.tool_choice is None)
and (
request.tool_choice == "auto"
or request.tool_choice is None
or (
not tool_parser_cls.supports_required_and_named
and request.tools
and (
request.tool_choice == "required"
or isinstance(
request.tool_choice,
ChatCompletionNamedToolChoiceParam,
)
)
)
)
)
):
# Automatic Tool Call Parsing (also used as fallback for
# required/named when supports_required_and_named=False)
if tokenizer is None:
raise ValueError(
"Tokenizer not available when `skip_tokenizer_init=True`"
)
# Automatic Tool Call Parsing
try:
tool_parser = tool_parser_cls(tokenizer, request.tools)
except RuntimeError as e:
+27
View File
@@ -0,0 +1,27 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Encode/decode utilities for multimodal tensors and field metadata
over JSON/HTTP, used by the disaggregated generate endpoint."""
from __future__ import annotations
import pybase64
from vllm.multimodal.inputs import MultiModalKwargsItem
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
_encoder = MsgpackEncoder(size_threshold=2**62) # force all tensors inline
_decoder = MsgpackDecoder(t=MultiModalKwargsItem)
def encode_mm_kwargs_item(item: MultiModalKwargsItem) -> str:
"""Serialize a MultiModalKwargsItem to a base64 string."""
bufs = _encoder.encode(item)
assert len(bufs) == 1, "All tensors should be inline"
return pybase64.b64encode(bufs[0]).decode("ascii")
def decode_mm_kwargs_item(data: str) -> MultiModalKwargsItem:
"""Deserialize a base64 string back to a MultiModalKwargsItem."""
raw = pybase64.b64decode(data)
return _decoder.decode(raw)
+9 -8
View File
@@ -35,14 +35,6 @@ class MultiModalFeatures(BaseModel):
Carries hashes (for cache lookup / identification) and placeholder
positions so the downstream `/generate` service knows *where* in
the token sequence each multimodal item lives.
Note:
Phase 1 metadata only.
Phase 2 should add `mm_kwargs` (processed tensor data) using a
binary transport so the ``/generate` side can skip re-processing.
The `/generate` endpoint must also be updated to inject these
features into `EngineInput` before passing to
`InputProcessor.process_inputs`.
"""
mm_hashes: dict[str, list[str]]
@@ -51,6 +43,15 @@ class MultiModalFeatures(BaseModel):
mm_placeholders: dict[str, list[PlaceholderRangeInfo]]
"""Per-modality placeholder ranges in the token sequence."""
kwargs_data: dict[str, list[str | None]] | None = None
"""Per-modality serialized tensor data.
Each value is a list parallel to ``mm_hashes[modality]``. A ``str``
entry is a base64-encoded ``MultiModalKwargsItem``; ``None`` means
the item should be resolved from cache. The entire field is
``None`` for metadata-only (cache-hit) responses.
"""
class GenerateRequest(BaseModel):
request_id: str = Field(
+43 -5
View File
@@ -25,6 +25,7 @@ from vllm.entrypoints.openai.engine.protocol import (
)
from vllm.entrypoints.openai.engine.serving import OpenAIServing, clamp_prompt_logprobs
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item
from vllm.entrypoints.serve.disagg.protocol import (
GenerateRequest,
GenerateResponse,
@@ -34,8 +35,14 @@ from vllm.entrypoints.serve.disagg.protocol import (
)
from vllm.entrypoints.serve.render.serving import OpenAIServingRender
from vllm.entrypoints.utils import should_include_usage
from vllm.inputs import EngineInput, mm_input
from vllm.logger import init_logger
from vllm.logprobs import Logprob
from vllm.multimodal.inputs import (
MultiModalKwargsItem,
MultiModalKwargsItems,
PlaceholderRange,
)
from vllm.outputs import RequestOutput
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.utils.collection_utils import as_list
@@ -103,11 +110,42 @@ class ServingTokens(OpenAIServing):
if raw_request:
raw_request.state.request_metadata = request_metadata
(engine_input,) = await self.openai_serving_render.preprocess_completion(
request,
prompt_input=request.token_ids,
prompt_embeds=None,
)
engine_input: EngineInput
if features := request.features:
# Convert PlaceholderRangeInfo → PlaceholderRange per modality.
mm_placeholders: dict[str, list[PlaceholderRange]] = {
modality: [
PlaceholderRange(offset=p.offset, length=p.length) for p in ranges
]
for modality, ranges in features.mm_placeholders.items()
}
# Deserialize tensor data when present; None → cache hit.
mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {}
if features.kwargs_data is not None:
for modality, items in features.kwargs_data.items():
mm_kwargs[modality] = [
decode_mm_kwargs_item(item) if item is not None else None
for item in items
]
else:
for modality, hashes in features.mm_hashes.items():
mm_kwargs[modality] = [None] * len(hashes)
engine_input = mm_input(
prompt_token_ids=request.token_ids,
mm_kwargs=MultiModalKwargsItems(mm_kwargs),
mm_hashes=features.mm_hashes,
mm_placeholders=mm_placeholders,
cache_salt=request.cache_salt,
)
else:
(engine_input,) = await self.openai_serving_render.preprocess_completion(
request,
prompt_input=request.token_ids,
prompt_embeds=None,
skip_mm_cache=True,
)
# Schedule the request and get the result generator.
result_generator: AsyncGenerator[RequestOutput, None] | None = None
+20 -4
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
from http import HTTPStatus
from typing import Any
from typing import Any, cast
from openai_harmony import Message as OpenAIMessage
@@ -25,6 +25,7 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
render_for_completion,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.entrypoints.serve.disagg.mm_serde import encode_mm_kwargs_item
from vllm.entrypoints.serve.disagg.protocol import (
GenerateRequest,
MultiModalFeatures,
@@ -37,6 +38,7 @@ from vllm.entrypoints.utils import (
from vllm.inputs import (
EngineInput,
MultiModalHashes,
MultiModalInput,
MultiModalPlaceholders,
PromptType,
SingletonPrompt,
@@ -251,6 +253,7 @@ class OpenAIServingRender:
default_template_kwargs=self.default_chat_template_kwargs,
tool_dicts=tool_dicts,
tool_parser=tool_parser,
skip_mm_cache=True,
reasoning_parser=self.reasoning_parser,
)
else:
@@ -342,6 +345,7 @@ class OpenAIServingRender:
request,
prompt_input=request.prompt,
prompt_embeds=request.prompt_embeds,
skip_mm_cache=True,
)
return engine_inputs
@@ -357,9 +361,10 @@ class OpenAIServingRender:
if engine_input.get("type") != "multimodal":
return None
# At this point engine_input is a MultiModalInputs TypedDict.
mm_hashes: MultiModalHashes = engine_input["mm_hashes"] # type: ignore[typeddict-item]
raw_placeholders: MultiModalPlaceholders = engine_input["mm_placeholders"] # type: ignore[typeddict-item]
# At this point engine_input is a MultiModalInput TypedDict.
mm_engine_input = cast(MultiModalInput, engine_input)
mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
mm_placeholders = {
modality: [
@@ -368,9 +373,20 @@ class OpenAIServingRender:
for modality, ranges in raw_placeholders.items()
}
# Serialize tensor data per modality.
kwargs_data: dict[str, list[str | None]] | None = None
if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
kwargs_data = {}
for modality, items in raw_mm_kwargs.items():
kwargs_data[modality] = [
encode_mm_kwargs_item(item) if item is not None else None
for item in items
]
return MultiModalFeatures(
mm_hashes=mm_hashes,
mm_placeholders=mm_placeholders,
kwargs_data=kwargs_data,
)
def _make_request_with_harmony(
+3 -2
View File
@@ -53,14 +53,15 @@ 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")
# FIXME(gmagogsfm): Re-enable HOP path once performance regression is fixed.
# _HOP_AVAILABLE = requires_torch_version("2.11")
_HOP_AVAILABLE = False
if _HOP_AVAILABLE:
from helion._compat import supports_torch_compile_fusion
@@ -406,33 +406,16 @@ class Attention(nn.Module, AttentionLayerBase):
def _init_turboquant_buffers(
self, cache_dtype: str, head_size: int, prefix: str
) -> None:
"""Initialize TurboQuant rotation/projection matrices and centroids."""
"""Initialize TurboQuant centroids for Lloyd-Max quantization."""
from vllm.model_executor.layers.quantization.turboquant.centroids import (
get_centroids,
)
from vllm.model_executor.layers.quantization.turboquant.config import (
TurboQuantConfig,
)
from vllm.model_executor.layers.quantization.turboquant.quantizer import (
generate_wht_signs,
)
tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size)
# Each layer needs a unique rotation matrix so quantization errors
# don't correlate across layers. Stride must exceed max head_dim to
# ensure non-overlapping RNG streams between adjacent layers.
_TQ_LAYER_SEED_STRIDE = 1337
from vllm.model_executor.models.utils import extract_layer_index
layer_idx = extract_layer_index(prefix)
seed = tq_config.seed + layer_idx * _TQ_LAYER_SEED_STRIDE
self.register_buffer(
"_tq_signs",
generate_wht_signs(head_size, seed=seed),
)
self.register_buffer(
"_tq_centroids",
get_centroids(head_size, tq_config.centroid_bits),
+24 -31
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
import os
from collections.abc import Callable
from typing import Any
@@ -611,51 +612,43 @@ def matmul_batch_invariant(a, b, *, out=None):
out.copy_(result)
return out
return result
elif a.ndim == 3 and b.ndim == 3:
# Handle batched case like bmm
return bmm_batch_invariant(a, b, out=out)
elif a.ndim == 3 and b.ndim == 2:
# Handle 3D x 2D: common for linear layers
# (batch, seq, hidden) @ (hidden, out) -> (batch, seq, out)
# Reshape to 2D, do mm, reshape back
batch, seq, hidden = a.shape
elif b.ndim == 2:
# Handle ND x 2D: Common for linear layers
# (..., batch, seq, hidden) @ (hidden, out) -> (..., batch, seq, out)
batch_dims = a.shape[:-1]
hidden = a.shape[-1]
out_dim = b.shape[-1]
a_2d = a.reshape(-1, hidden)
result_2d = matmul_persistent(a_2d, b)
result = result_2d.reshape(batch, seq, -1)
result = result_2d.reshape(batch_dims + (out_dim,))
if out is not None:
out.copy_(result)
return out
return result
elif a.ndim == 2 and b.ndim == 3:
# Handle 2D x 3D: (M, K) @ (B, K, N) -> (B, M, N)
# By broadcasting `a` to 3D, we can reuse the batched matrix
# multiplication logic.
a_expanded = a.unsqueeze(0).expand(b.shape[0], -1, -1)
return bmm_batch_invariant(a_expanded, b, out=out)
elif a.ndim == 4 and b.ndim == 4:
# Handle 4D attention tensors: [batch, heads, seq, dim]
# Reshape to 3D, process, reshape back
batch, heads, seq_a, dim_a = a.shape
_, _, dim_b, seq_b = b.shape
# Reshape to [batch*heads, seq_a, dim_a]
a_3d = a.reshape(batch * heads, seq_a, dim_a)
b_3d = b.reshape(batch * heads, dim_b, seq_b)
elif a.ndim >= 2 and b.ndim >= 3:
# Generic handler for 2D x ND and ND x ND (except 1D)
# Broadcast dims to ensure both matrices have the same shape
# If 2D x ND, then unsqueeze to add a dim to a
if a.ndim == 2:
a = a.unsqueeze(0)
broadcast_shape = torch.broadcast_shapes(a.shape[:-2], b.shape[:-2])
a = a.expand(broadcast_shape + a.shape[-2:])
b = b.expand(broadcast_shape + b.shape[-2:])
batch_dim = math.prod(broadcast_shape)
# Reuse broadcast shape to get all dims except mm dims
a_3d = a.reshape(batch_dim, a.shape[-2], a.shape[-1])
b_3d = b.reshape(batch_dim, b.shape[-2], b.shape[-1])
# Do batched matmul
result_3d = bmm_batch_invariant(a_3d, b_3d)
# Reshape back to [batch, heads, seq_a, seq_b]
result = result_3d.reshape(batch, heads, seq_a, seq_b)
# Reshape back to [broadcast_shape, seq_a, seq_b]
result = result_3d.reshape(broadcast_shape + (a.shape[-2], b.shape[-1]))
if out is not None:
out.copy_(result)
return out
return result
else:
raise ValueError(
f"matmul_batch_invariant currently only supports 2D x 2D, 3D x 3D, "
f"3D x 2D, 2D x 3D, and 4D x 4D, "
f"matmul_batch_invariant requires both inputs be at least 2D "
f"got shapes {a.shape} and {b.shape}"
)
+24 -1
View File
@@ -762,6 +762,25 @@ def nvfp4_moe_quant_config(
)
def mxfp4_moe_quant_config(
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
) -> FusedMoEQuantConfig:
"""
Construct a quant config for MXFP4 x MXFP4 MoE.
MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
separate alphas / global activation scales in this config.
"""
return FusedMoEQuantConfig.make(
"mxfp4",
w1_scale=w1_scale,
w2_scale=w2_scale,
per_act_token_quant=False,
per_out_ch_quant=False,
block_shape=None,
)
def nvfp4_w4a16_moe_quant_config(
g1_alphas: torch.Tensor,
g2_alphas: torch.Tensor,
@@ -971,7 +990,11 @@ class FusedMoEParallelConfig:
@property
def use_batched_activation_format(self):
return self.use_deepep_ll_kernels
return (
self.use_deepep_ll_kernels
or self.use_fi_nvl_one_sided_kernels
or self.use_nixl_ep_kernels
)
@property
def use_ag_rs_all2all_kernels(self):
@@ -36,6 +36,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8DynamicTokenSym,
kFp8StaticChannelSym,
kFp8StaticTensorSym,
kMxfp4Dynamic,
kMxfp4Static,
kNvfp4Dynamic,
kNvfp4Static,
)
@@ -795,6 +797,299 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
)
def run_cutlass_moe_mxfp4(
output: torch.Tensor,
a: torch.Tensor,
w1_fp4: torch.Tensor,
w1_blockscale: torch.Tensor,
w2_fp4: torch.Tensor,
w2_blockscale: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
activation: MoEActivation,
workspace13: torch.Tensor,
workspace2: torch.Tensor,
m: int,
n: int,
k: int,
e: int,
device: torch.device,
apply_router_weight_on_input: bool = False,
) -> None:
"""MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM."""
is_gated = activation.is_gated
w1_n = n * 2 if is_gated else n
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8"
assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8"
assert (
w1_fp4.ndim == 3
and w2_fp4.ndim == 3
and w1_blockscale.ndim == 3
and w2_blockscale.ndim == 3
), "All Weights must be of rank 3 for cutlass_moe_mxfp4"
m_a, k_a = a.shape
e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape
e_w2, k_w2, half_n_w2 = w2_fp4.shape
assert e_w1 == e_w2 and e_w1 == e
assert k_a == half_k_w1 * 2 and k == k_w2
assert w1_n_actual == w1_n and half_n_w2 * 2 == n
assert m == m_a
assert 2 * half_k_w1 == k_w2
assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
assert topk_weights.size(0) == m and topk_ids.size(0) == m
topk = topk_ids.size(1)
out_dtype = a.dtype
num_topk = topk_ids.size(1)
expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device)
problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device)
a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
if apply_router_weight_on_input:
assert num_topk == 1, (
"apply_router_weight_on_input is only implemented for topk=1"
)
a.mul_(topk_weights.to(out_dtype))
ops.get_cutlass_moe_mm_data(
topk_ids,
expert_offsets,
problem_sizes1,
problem_sizes2,
a_map,
c_map,
e,
n,
k,
blockscale_offsets,
is_gated=is_gated,
)
a = ops.shuffle_rows(a, a_map)
rep_a_fp4, rep_a_blockscale = ops.mxfp4_experts_quant(
a,
expert_offsets,
blockscale_offsets,
e,
num_topk,
)
c1 = _resize_cache(workspace13, (m * topk, w1_n))
c2 = _resize_cache(workspace2, (m * topk, n))
c3 = _resize_cache(workspace13, (m * topk, k))
ops.cutlass_mxfp4_moe_mm(
c1,
rep_a_fp4,
w1_fp4,
rep_a_blockscale,
w1_blockscale,
problem_sizes1,
expert_offsets[:-1],
blockscale_offsets[:-1],
)
del rep_a_fp4, rep_a_blockscale
if activation == MoEActivation.SILU:
int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant(
c1, expert_offsets, blockscale_offsets, e, num_topk
)
else:
apply_moe_activation(activation, c2, c1)
int_fp4, int_blockscale = ops.mxfp4_experts_quant(
c2, expert_offsets, blockscale_offsets, e, num_topk
)
ops.cutlass_mxfp4_moe_mm(
c3,
int_fp4,
w2_fp4,
int_blockscale,
w2_blockscale,
problem_sizes2,
expert_offsets[:-1],
blockscale_offsets[:-1],
)
del int_fp4, int_blockscale
c3 = ops.shuffle_rows(c3, c_map)
assert output.dtype == out_dtype
if not apply_router_weight_on_input:
output.copy_(
(
c3.view(m, num_topk, k)
* topk_weights.view(m, num_topk, 1).to(out_dtype)
).sum(dim=1),
non_blocking=True,
)
else:
output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True)
return
def swizzle_mxfp4_scales(
scales: torch.Tensor,
N: int,
K: int,
) -> torch.Tensor:
"""Swizzle flat [N, K//32] E8M0 scales to CUTLASS tiled layout.
CUTLASS expects MX scale factors in a tiled layout:
[numMTiles, numKTiles, 32, 4, 4]
where numMTiles = ceil(N/128), numKTiles = ceil(K/128),
and the inner dimensions correspond to the swizzle pattern:
mTileIdx = mIdx / 128
outerMIdx = mIdx % 32
innerMIdx = (mIdx / 32) % 4
kTileIdx = kIdx / 4
innerKIdx = kIdx % 4
with kIdx = col_in_scale_space (i.e., index into K//32).
"""
assert scales.dtype == torch.uint8
num_scale_cols = K // 32 # number of E8M0 scale values per row
num_m_tiles = (N + 127) // 128
num_k_tiles = (num_scale_cols + 3) // 4
# Pad N to multiple of 128 and scale_cols to multiple of 4
padded_N = num_m_tiles * 128
padded_scale_cols = num_k_tiles * 4
# Start with flat scales, pad if needed
padded = torch.zeros(
padded_N, padded_scale_cols, dtype=torch.uint8, device=scales.device
)
padded[:N, :num_scale_cols] = scales
# Reshape to tile structure:
# [numMTiles, 4, 32, numKTiles, 4]
# mTileIdx, innerMIdx, outerMIdx, kTileIdx, innerKIdx
tiled = padded.reshape(num_m_tiles, 4, 32, num_k_tiles, 4)
# Permute to [numMTiles, numKTiles, 32, 4, 4]
# (outerMIdx, innerMIdx, innerKIdx)
tiled = tiled.permute(0, 3, 2, 1, 4).contiguous()
return tiled.reshape(-1)
class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
"""CUTLASS MXFP4 x MXFP4 fused MoE expert implementation."""
@property
def expects_unquantized_inputs(self) -> bool:
return True
@staticmethod
def _supports_current_device() -> bool:
p = current_platform
return p.is_cuda() and p.is_device_capability_family(100)
@staticmethod
def _supports_no_act_and_mul() -> bool:
return True
@staticmethod
def _supports_quant_scheme(
weight_key: QuantKey | None,
activation_key: QuantKey | None,
) -> bool:
return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic)
@staticmethod
def _supports_activation(activation: MoEActivation) -> bool:
return activation in [
MoEActivation.SILU,
MoEActivation.GELU,
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUSTEP,
MoEActivation.SILU_NO_MUL,
MoEActivation.GELU_NO_MUL,
MoEActivation.RELU2_NO_MUL,
]
@staticmethod
def _supports_parallel_config(
moe_parallel_config: FusedMoEParallelConfig,
) -> bool:
return moe_parallel_config.ep_size == 1
@staticmethod
def activation_format() -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
def supports_expert_map(self) -> bool:
return False
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
return TopKWeightAndReduceNoOP()
def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
return act_dtype
def workspace_shapes(
self,
M: int,
N: int,
K: int,
topk: int,
global_num_experts: int,
local_num_experts: int,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
activation: MoEActivation,
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
workspace1 = (M * topk, max(2 * N, K))
workspace2 = (M * topk, N)
output = (M, K)
return (workspace1, workspace2, output)
def apply(
self,
output: torch.Tensor,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
activation: MoEActivation,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
a2_scale: torch.Tensor | None,
workspace13: torch.Tensor | None,
workspace2: torch.Tensor | None,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
):
e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
n = w2.shape[2] * 2
run_cutlass_moe_mxfp4(
output=output,
a=hidden_states,
w1_fp4=w1,
w1_blockscale=self.w1_scale,
w2_fp4=w2,
w2_blockscale=self.w2_scale,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=activation,
workspace13=workspace13,
workspace2=workspace2,
m=m,
n=n,
k=k,
e=e,
device=hidden_states.device,
apply_router_weight_on_input=apply_router_weight_on_input,
)
# W4A8
def run_cutlass_moe_w4a8_fp8(
output: torch.Tensor,
@@ -4,7 +4,6 @@
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm import envs
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 (
@@ -126,7 +125,7 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
# We use global_num_experts due to how moe_align_block_size handles
# expert_maps.
K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
K_dim = self.moe_config.hidden_dim if self.moe_config.hidden_dim == K * 2 else K
output_shape = (local_num_experts, M, K_dim)
workspace2 = (local_num_experts, M, N)
workspace1 = output_shape
@@ -163,13 +162,12 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
assert self.w1_scale.ndim == 3
assert self.w2_scale.ndim == 3
input_global_scale = (
None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
use_prequantized_inputs = (
a1q_scale is not None and hidden_states.dtype == torch.uint8
)
input_global_scale = None if use_prequantized_inputs else self.a1_gscale
flashinfer_hidden_states = (
(hidden_states, a1q_scale)
if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
else hidden_states
(hidden_states, a1q_scale) if use_prequantized_inputs else hidden_states
)
flashinfer_cutedsl_moe_masked(
hidden_states=flashinfer_hidden_states,
@@ -162,10 +162,9 @@ 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
activation_format = (
mk.FusedMoEActivationFormat.BatchedExperts
if use_batched
if config.moe_parallel_config.use_batched_activation_format
else mk.FusedMoEActivationFormat.Standard
)
@@ -205,16 +204,22 @@ def select_nvfp4_moe_backend(
raise ValueError(_make_log_unsupported(backend, reason))
def _resolve_requested_backend(
backend: NvFp4MoeBackend,
) -> NvFp4MoeBackend:
if (
activation_format == mk.FusedMoEActivationFormat.BatchedExperts
and backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
):
return NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
return backend
# Handle explicit moe_backend from user.
runner_backend = config.moe_backend
if runner_backend != "auto":
requested_backend = map_nvfp4_backend(runner_backend)
# For batched activation format, use batched variant if available.
if (
activation_format == mk.FusedMoEActivationFormat.BatchedExperts
and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
):
requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
requested_backend = _resolve_requested_backend(
map_nvfp4_backend(runner_backend)
)
return _return_or_raise(
requested_backend, config, weight_key, activation_key, activation_format
)
@@ -227,7 +232,9 @@ def select_nvfp4_moe_backend(
elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):
# If user is explicit about backend, validate it.
backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()]
backend = _resolve_requested_backend(
fi_2_vllm_backend_map[get_flashinfer_moe_backend()]
)
return _return_or_raise(
backend, config, weight_key, activation_key, activation_format
)
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import inspect
from collections.abc import Callable
import deep_ep
@@ -120,6 +121,20 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
# time. This setting is handled by post_init_setup.
self.use_ue8m0_dispatch = False
# Check if DeepEP supports use_nvfp4 in low_latency_dispatch.
# This requires the hybrid-ep branch of DeepEP.
self.has_nvfp4_support = (
"use_nvfp4" in inspect.signature(buffer.low_latency_dispatch).parameters
)
if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH and not self.has_nvfp4_support:
logger.warning_once(
"VLLM_DEEPEPLL_NVFP4_DISPATCH=1 but DeepEP does not support "
"use_nvfp4 in low_latency_dispatch. Falling back to FP8/BF16 "
"dispatch. Install DeepEP from the hybrid-ep branch for "
"NvFP4 dispatch support: "
"https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep"
)
def post_init_setup(self, fused_experts: mk.FusedMoEExperts):
if not fused_experts.supports_packed_ue8m0_act_scales():
# Early exit.
@@ -183,27 +198,29 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
assert isinstance(x, (torch.Tensor, tuple))
q_dtype = quant_config.quant_dtype
if q_dtype == "nvfp4" and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH:
nvfp4_quant_path = (
q_dtype == "nvfp4"
and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
and self.has_nvfp4_support
)
if nvfp4_quant_path:
logger.info_once(
"Since VLLM_DEEPEPLL_NVFP4_DISPATCH==1, make sure "
"using the hybrid-ep branch of DeepEP"
"(https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep)"
"Quantization is fused with DeepEP nvfp4 dispatch (hybrid-ep branch)"
)
assert isinstance(x, tuple)
x_scales = x[1]
x = x[0].permute(2, 0, 1)
num_experts, max_tokens, hidden_dim_by_2 = x.shape
hidden_dim = hidden_dim_by_2 * 2
logger.info_once(
"Quantization is fused with DeepEP nvfp4 dispatch for "
"FlashInfer CUTEDSL as VLLM_DEEPEPLL_NVFP4_DISPATCH==1"
)
else:
if q_dtype == "nvfp4":
q_dtype = None
logger.info_once(
"Using DeepEP bfloat16 dispatch for FlashInfer CUTEDSL as "
"VLLM_DEEPEPLL_NVFP4_DISPATCH==0"
"Using DeepEP bfloat16 dispatch for FlashInfer CUTEDSL "
"(nvfp4 dispatch %s)",
"not supported by this DeepEP build"
if not self.has_nvfp4_support
else "disabled",
)
assert isinstance(x, torch.Tensor)
num_experts, max_tokens, hidden_dim = x.size()
@@ -260,7 +277,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
use_nvfp4 = False
nvfp4_dispatch = (
quant_config.quant_dtype == "nvfp4" and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
quant_config.quant_dtype == "nvfp4"
and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
and self.has_nvfp4_support
)
if nvfp4_dispatch:
use_nvfp4 = True
@@ -1,5 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
from dataclasses import dataclass
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -14,6 +17,156 @@ def get_local_sizes():
return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
@dataclass
class _OneSidedDispatchMetadata:
combine_indices: list[torch.Tensor]
combine_weights: list[torch.Tensor]
def _normalize_source_num_tokens(
source_num_tokens: Sequence[int] | None,
num_dispatchers: int,
runtime_max_tokens_per_rank: int,
) -> list[int]:
if source_num_tokens is None:
return [runtime_max_tokens_per_rank] * num_dispatchers
normalized = [int(x) for x in source_num_tokens[:num_dispatchers]]
if len(normalized) < num_dispatchers:
normalized.extend(
[runtime_max_tokens_per_rank] * (num_dispatchers - len(normalized))
)
return [min(max(x, 0), runtime_max_tokens_per_rank) for x in normalized]
def _group_rank_batched_inputs_by_local_expert(
hidden_states: torch.Tensor,
hidden_scales: torch.Tensor | None,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
*,
num_local_experts: int,
first_local_expert: int,
num_dispatchers: int,
runtime_max_tokens_per_rank: int,
source_num_tokens: Sequence[int] | None,
) -> tuple[
torch.Tensor,
torch.Tensor | None,
mk.ExpertTokensMetadata,
_OneSidedDispatchMetadata,
]:
hidden_dim = hidden_states.shape[-1]
batched_hidden_states = hidden_states.new_empty(
(num_local_experts, runtime_max_tokens_per_rank, hidden_dim)
)
batched_hidden_scales = (
None
if hidden_scales is None
else hidden_scales.new_empty(
(num_local_experts, runtime_max_tokens_per_rank, hidden_scales.shape[-1])
)
)
tokens_per_expert = torch.zeros(
num_local_experts, dtype=torch.int32, device=hidden_states.device
)
combine_indices: list[torch.Tensor] = []
combine_weights: list[torch.Tensor] = []
valid_source_num_tokens = _normalize_source_num_tokens(
source_num_tokens, num_dispatchers, runtime_max_tokens_per_rank
)
for local_expert in range(num_local_experts):
global_expert = first_local_expert + local_expert
expert_indices: list[torch.Tensor] = []
expert_weights: list[torch.Tensor] = []
cursor = 0
for dispatcher, num_tokens in enumerate(valid_source_num_tokens):
if num_tokens == 0:
continue
token_idx, topk_slot_idx = torch.where(
topk_ids[dispatcher, :num_tokens] == global_expert
)
rows = token_idx.numel()
if rows == 0:
continue
batched_hidden_states[local_expert, cursor : cursor + rows] = hidden_states[
dispatcher, token_idx
]
if batched_hidden_scales is not None:
assert hidden_scales is not None
batched_hidden_scales[local_expert, cursor : cursor + rows] = (
hidden_scales[dispatcher, token_idx]
)
expert_indices.append(
dispatcher * runtime_max_tokens_per_rank + token_idx.to(torch.int64)
)
expert_weights.append(
topk_weights[dispatcher, token_idx, topk_slot_idx].contiguous()
)
cursor += rows
tokens_per_expert[local_expert] = cursor
combine_indices.append(
torch.cat(expert_indices)
if expert_indices
else torch.empty(0, dtype=torch.int64, device=hidden_states.device)
)
combine_weights.append(
torch.cat(expert_weights)
if expert_weights
else torch.empty(0, dtype=topk_weights.dtype, device=topk_weights.device)
)
expert_tokens_meta = mk.ExpertTokensMetadata(
expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None
)
dispatch_metadata = _OneSidedDispatchMetadata(
combine_indices=combine_indices,
combine_weights=combine_weights,
)
return (
batched_hidden_states,
batched_hidden_scales,
expert_tokens_meta,
dispatch_metadata,
)
def _reduce_local_expert_outputs_to_rank_batched_payload(
fused_expert_output: torch.Tensor,
dispatch_metadata: _OneSidedDispatchMetadata,
*,
num_dispatchers: int,
runtime_max_tokens_per_rank: int,
apply_router_weight_on_input: bool,
) -> torch.Tensor:
hidden_dim = fused_expert_output.shape[-1]
combine_payload = fused_expert_output.new_zeros(
(num_dispatchers, runtime_max_tokens_per_rank, hidden_dim)
)
flat_payload = combine_payload.view(-1, hidden_dim)
for local_expert, linear_indices in enumerate(dispatch_metadata.combine_indices):
rows = linear_indices.numel()
if rows == 0:
continue
expert_output = fused_expert_output[local_expert, :rows]
if not apply_router_weight_on_input:
expert_output = expert_output * dispatch_metadata.combine_weights[
local_expert
].to(expert_output.dtype).unsqueeze(-1)
flat_payload.index_add_(0, linear_indices, expert_output)
return combine_payload
class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
"""FlashInfer implementation using the Moe AlltoAll kernel."""
@@ -31,8 +184,14 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
self.num_experts = num_experts
self.hidden_size = hidden_size
self.num_dispatchers_ = num_dispatchers
self.all2all_manager = get_ep_group().device_communicator.all2all_manager
assert self.num_experts % self.num_dispatchers_ == 0, (
"flashinfer_nvlink_one_sided requires evenly sharded local experts."
)
self.num_local_experts = self.num_experts // self.num_dispatchers_
self.first_local_expert = self.all2all_manager.rank * self.num_local_experts
self.dispatch_metadata: _OneSidedDispatchMetadata | None = None
self.all2all_manager.initialize(
max_num_tokens=self.max_num_tokens,
top_k=self.top_k,
@@ -42,10 +201,10 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
@property
def activation_format(self) -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
return mk.FusedMoEActivationFormat.BatchedExperts
def max_num_tokens_per_rank(self) -> int | None:
return None
return self.max_num_tokens
def num_dispatchers(self) -> int:
return self.num_dispatchers_
@@ -67,6 +226,10 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
quant_config: FusedMoEQuantConfig,
defer_input_quant: bool = False,
) -> mk.PrepareResultType:
if defer_input_quant:
raise NotImplementedError(
f"{self.__class__.__name__} does not support defer_input_quant=True."
)
if apply_router_weight_on_input:
topk = topk_ids.size(1)
assert topk == 1, (
@@ -104,7 +267,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
)
if a1q_scale is not None:
a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
# Apply scale interleaving only for CUTLASS (not TRT-LLM)
# Swizzle after dispatch when the selected MoE kernel expects it.
if (
quant_config.quant_dtype == "nvfp4"
and quant_config.is_nvfp4_scale_swizzled
@@ -112,15 +275,32 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
a1q_scale_recv = a1q_scale_recv.view(-1, self.hidden_size // 16)
a1q_scale_recv = a1q_scale_recv.view(
self.num_dispatchers_,
self.runtime_max_tokens_per_rank,
self.hidden_size // 16,
)
else:
a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads
a1q_scale_recv = None
a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1])
topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1])
topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1])
(
a1q_recv,
a1q_scale_recv,
expert_tokens_meta,
self.dispatch_metadata,
) = _group_rank_batched_inputs_by_local_expert(
a1q_recv,
a1q_scale_recv,
topk_ids_recv,
topk_weights_recv,
num_local_experts=self.num_local_experts,
first_local_expert=self.first_local_expert,
num_dispatchers=self.num_dispatchers_,
runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
source_num_tokens=global_num_tokens_cpu,
)
return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv
return a1q_recv, a1q_scale_recv, expert_tokens_meta, None, None
def finalize(
self,
@@ -132,15 +312,20 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
weight_and_reduce_impl: mk.TopKWeightAndReduce,
) -> None:
assert self.all2all_manager.moe_alltoall is not None
ep_size = self.all2all_manager.world_size
hidden_size = fused_expert_output.shape[-1]
fused_expert_output = fused_expert_output.view(
ep_size, self.runtime_max_tokens_per_rank, hidden_size
assert self.dispatch_metadata is not None, (
"flashinfer_nvlink_one_sided finalize called before prepare"
)
combine_payload = _reduce_local_expert_outputs_to_rank_batched_payload(
fused_expert_output,
self.dispatch_metadata,
num_dispatchers=self.num_dispatchers_,
runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
apply_router_weight_on_input=apply_router_weight_on_input,
)
combined_output = self.all2all_manager.moe_alltoall.combine(
payload=fused_expert_output,
payload=combine_payload,
runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
)
self.dispatch_metadata = None
output.copy_(combined_output)
@@ -4,6 +4,7 @@
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 import (
FusedMoE,
@@ -11,6 +12,10 @@ from vllm.model_executor.layers.fused_moe import (
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
mxfp4_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
CutlassExpertsMxfp4,
)
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
MarlinExperts,
@@ -36,7 +41,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
super().__init__(moe)
self.group_size = 32
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
self.experts_cls = MarlinExperts
self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
self.experts_cls: type[mk.FusedMoEExperts]
if self.use_cutlass_mxfp4:
logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE", scope="local")
self.experts_cls = CutlassExpertsMxfp4
else:
logger.info_once("Using MarlinExperts for MXFP4 MoE", scope="local")
self.experts_cls = MarlinExperts
def create_weights(
self,
@@ -109,11 +121,19 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
def get_fused_moe_quant_config(
self, layer: torch.nn.Module
) -> FusedMoEQuantConfig | None:
return make_mxfp4_moe_quant_config(
mxfp4_backend=self.mxfp4_backend,
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
)
if self.use_cutlass_mxfp4:
# W4A4: both weights and activations quantized to MXFP4
return mxfp4_moe_quant_config(
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
)
else:
# W4A16: weight-only via Marlin
return make_mxfp4_moe_quant_config(
mxfp4_backend=self.mxfp4_backend,
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
)
def process_weights_after_loading(self, layer: FusedMoE) -> None:
layer.w13_weight = torch.nn.Parameter(
@@ -126,13 +146,45 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
)
delattr(layer, "w2_weight_packed")
logger.warning_once(
"Your GPU does not have native support for FP4 computation but "
"FP4 quantization is being used. Weight-only FP4 compression "
"will be used leveraging the Marlin kernel. This may degrade "
"performance for compute-heavy workloads."
)
prepare_moe_fp4_layer_for_marlin(layer)
if self.use_cutlass_mxfp4:
# Swizzle weight scales from flat checkpoint layout [E, N, K//32]
# to CUTLASS tiled layout [E, numMTiles*numKTiles*512].
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
swizzle_mxfp4_scales,
)
E = layer.w13_weight_scale.shape[0]
w13_N = layer.w13_weight_scale.shape[1]
w13_scale_K = layer.w13_weight_scale.shape[2]
w13_K = w13_scale_K * 32
w2_M = layer.w2_weight_scale.shape[1]
w2_scale_N = layer.w2_weight_scale.shape[2]
w2_N = w2_scale_N * 32
swizzled_w13 = []
swizzled_w2 = []
for e_idx in range(E):
s13 = layer.w13_weight_scale[e_idx]
sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
s2 = layer.w2_weight_scale[e_idx]
sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
layer.w13_weight_scale = torch.nn.Parameter(
torch.stack(swizzled_w13), requires_grad=False
)
layer.w2_weight_scale = torch.nn.Parameter(
torch.stack(swizzled_w2), requires_grad=False
)
else:
logger.warning_once(
"Your GPU does not have native support for FP4 computation "
"but FP4 quantization is being used. Weight-only FP4 "
"compression will be used leveraging the Marlin kernel. "
"This may degrade performance for compute-heavy workloads."
)
prepare_moe_fp4_layer_for_marlin(layer)
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
if self.moe_quant_config is not None:
@@ -1,12 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""TurboQuant: Near-optimal KV-cache quantization for vLLM.
"""TurboQuant: KV-cache quantization for vLLM.
PolarQuant compression: random rotation + per-coordinate Lloyd-Max
scalar quantization for keys, uniform quantization for values.
Hadamard rotation + per-coordinate Lloyd-Max scalar quantization for
keys, uniform quantization for values.
Reference: "TurboQuant: Online Vector Quantization with Near-optimal
Distortion Rate" (ICLR 2026), Zandieh et al.
The technique implemented here consists of the scalar case of the HIGGS
quantization method (Malinovskii et al., "Pushing the Limits of Large
Language Model Quantization via the Linearity Theorem", NAACL 2025;
preprint arXiv:2411.17525): rotation + optimized grid + optional
re-normalization, applied to KV cache compression. A first application
of this approach to KV-cache compression is in "Cache Me If You Must:
Adaptive Key-Value Quantization for Large Language Models" (Shutova
et al., ICML 2025; preprint arXiv:2501.19392). Both these references
pre-date the TurboQuant paper (Zandieh et al., ICLR 2026).
"""
from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig
@@ -36,10 +36,22 @@ TQ_PRESETS: dict[str, dict] = {
class TurboQuantConfig:
"""Configuration for TurboQuant KV-cache quantization.
Uses PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys
and uniform quantization for values. QJL is intentionally omitted
community consensus (5+ independent groups) found it hurts attention
quality by amplifying variance through softmax.
Applies Hadamard rotation followed by per-coordinate Lloyd-Max scalar
quantization for keys, and uniform quantization for values.
Historical note: this is the scalar case of the HIGGS quantization
method (Malinovskii et al., "Pushing the Limits of Large Language Model
Quantization via the Linearity Theorem", NAACL 2025; preprint
arXiv:2411.17525): rotation + optimized grid + optional re-normalization,
applied to KV cache compression. A first application of this approach to
KV-cache compression is in "Cache Me If You Must: Adaptive Key-Value
Quantization for Large Language Models" (Shutova et al., ICML 2025;
preprint arXiv:2501.19392). Both these references pre-date the
TurboQuant paper.
QJL is intentionally omitted community consensus (5+ independent
groups) found it hurts attention quality by amplifying variance through
softmax.
Named presets (use via --kv-cache-dtype):
turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL
@@ -53,8 +65,6 @@ class TurboQuantConfig:
rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys.
value_quant_bits: Bits per value dimension for uniform quantization.
3 = 8 levels, 4 = 16 levels (default).
seed: Base seed for deterministic random matrix generation.
Actual seed per layer = seed + layer_idx * 1337.
norm_correction: Re-normalize centroid vectors to unit norm before
inverse rotation during dequant. Fixes quantization-induced norm
distortion, improving PPL by ~0.8% at 4-bit.
@@ -63,7 +73,7 @@ class TurboQuantConfig:
head_dim: int = 128
key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys
value_quant_bits: int = 4 # 3-4 = uniform quantized values
seed: int = 42
seed: int = 42 # kept for backward compatibility; no longer used internally
norm_correction: bool = False
@property
@@ -2,23 +2,5 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""TurboQuant quantizer utilities.
Serving path uses generate_wht_signs() for WHT rotation sign buffers.
Triton kernels handle all quantization, packing, and dequantization on GPU.
"""
import torch
_CPU = torch.device("cpu")
def generate_wht_signs(d: int, seed: int, device: torch.device = _CPU) -> torch.Tensor:
"""Generate deterministic random ±1 signs for WHT rotation.
Used with Walsh-Hadamard Transform for per-layer rotation randomization.
Same seed derivation as QR (per-layer via seed + layer_idx * stride).
"""
gen = torch.Generator(device="cpu")
gen.manual_seed(seed)
bits = torch.randint(0, 2, (d,), generator=gen, device="cpu")
signs = bits.float() * 2 - 1
return signs.to(device)
+10 -2
View File
@@ -67,6 +67,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
from .interfaces import (
MultiModalEmbeddings,
SupportsEagle3,
SupportsLoRA,
SupportsMultiModal,
SupportsPP,
)
@@ -880,6 +881,7 @@ class Gemma4ForConditionalGeneration(
nn.Module,
SupportsMultiModal,
SupportsPP,
SupportsLoRA,
SupportsEagle3,
):
packed_modules_mapping = {
@@ -1358,10 +1360,16 @@ class Gemma4ForConditionalGeneration(
def get_mm_mapping(self) -> MultiModelKeys:
"""Get the module prefix mapping for multimodal models."""
connectors = ["embed_vision"]
tower_models = ["vision_tower"]
if self.audio_tower is not None:
connectors.append("embed_audio")
tower_models.append("audio_tower")
return MultiModelKeys.from_string_field(
language_model="language_model",
connector=["embed_vision", "embed_audio"],
tower_model=["vision_tower", "audio_tower"],
connector=connectors,
tower_model=tower_models,
)
@classmethod
+3 -1
View File
@@ -66,7 +66,7 @@ from .interfaces import (
SupportsTranscription,
)
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
from .whisper import ISO639_1_SUPPORTED_LANGS
from .whisper import ISO639_1_SUPPORTED_LANGS, _create_fake_bias_for_k_proj
class GlmAsrEncoderRotaryEmbedding(nn.Module):
@@ -499,6 +499,8 @@ class GlmAsrEncoder(nn.Module):
"""Custom weight loading to handle q_proj/k_proj/v_proj -> qkv_proj mapping."""
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
+6 -7
View File
@@ -145,14 +145,15 @@ class PlaceholderRange:
"""
@cached_property
def embeds_cumsum(self) -> torch.Tensor | None:
return None if self.is_embed is None else self.is_embed.cumsum(dim=0)
def embeds_cumsum(self) -> list[int] | None:
# python list so python indexing avoids torch C++ overhead/conversions/deallocs
return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist()
def get_num_embeds(self) -> int:
if self.embeds_cumsum is None:
return self.length
return int(self.embeds_cumsum[-1])
return self.embeds_cumsum[-1] if self.embeds_cumsum else 0
def get_embeds_indices_in_range(
self, start_idx: int, end_idx: int
@@ -170,10 +171,8 @@ class PlaceholderRange:
if self.embeds_cumsum is None:
return start_idx, end_idx
embeds_start_idx = (
int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0
)
embeds_end_idx = int(self.embeds_cumsum[end_idx - 1])
embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0
embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0
return embeds_start_idx, embeds_end_idx
+3 -3
View File
@@ -29,9 +29,9 @@ except ImportError:
soundfile = PlaceholderModule("soundfile") # type: ignore[assignment]
# Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`, soundfile
# being librosa's main backend. Used to validate if an audio loading error is due to a
# server error vs a client error (invalid audio file).
# Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`,
# soundfile being the main audio loading backend. Used to validate if an audio
# loading error is due to a server error vs a client error (invalid audio file).
# 0 = sf_error(NULL) race condition: when multiple threads fail sf_open_virtual
# concurrently, one thread may clear the global error before another reads it,
# producing code=0 ("Garbled error message from libsndfile" in soundfile).
+1
View File
@@ -382,6 +382,7 @@ def _get_backend_priorities(
if is_aiter_found_and_supported():
backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
backends.append(AttentionBackendEnum.TRITON_ATTN)
backends.append(AttentionBackendEnum.TURBOQUANT)
return backends
+8
View File
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.logger import init_logger
from vllm.platforms.cpu import CpuPlatform
@@ -22,3 +24,9 @@ class ZenCpuPlatform(CpuPlatform):
def is_zen_cpu(self) -> bool:
# is_cpu() also returns True for this platform (inherited from CpuPlatform).
return True
# Currently, AMD CPUs do not support float16 compute.
# Hence explicitly return bfloat16 and float32.
@property
def supported_dtypes(self) -> list[torch.dtype]:
return [torch.bfloat16, torch.float32]
+4
View File
@@ -60,6 +60,10 @@ _REASONING_PARSERS_TO_REGISTER = {
"kimi_k2_reasoning_parser",
"KimiK2ReasoningParser",
),
"mimo": (
"qwen3_reasoning_parser",
"Qwen3ReasoningParser",
),
"minimax_m2": (
"minimax_m2_reasoning_parser",
"MiniMaxM2ReasoningParser",
+4
View File
@@ -94,6 +94,10 @@ _TOOL_PARSERS_TO_REGISTER = {
"longcat_tool_parser",
"LongcatFlashToolParser",
),
"mimo": (
"qwen3xml_tool_parser",
"Qwen3XMLToolParser",
),
"minimax_m2": (
"minimax_m2_tool_parser",
"MinimaxM2ToolParser",
+11
View File
@@ -44,6 +44,17 @@ class ToolParser:
derived classes.
"""
# When True (default), the serving layer uses the standard JSON-based
# parsing for tool_choice="required" and named function tool_choice,
# which works for models where guided decoding produces well-formed
# JSON output (e.g. Hermes).
# Subclasses set False when the standard parsing does not work for
# their model's output format (e.g. GLM models that use XML). When
# False, the serving layer falls back to the tool_parser's
# extract_tool_calls / extract_tool_calls_streaming methods for
# required/named tool_choice, treating them the same as "auto".
supports_required_and_named: bool = True
def __init__(
self,
tokenizer: TokenizerLike,
@@ -23,6 +23,8 @@ logger = init_logger(__name__)
class Glm47MoeModelToolParser(Glm4MoeModelToolParser):
supports_required_and_named = False
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
super().__init__(tokenizer, tools)
# GLM-4.7 format: <tool_call>func_name[<arg_key>...]*</tool_call>
+22 -1
View File
@@ -20,6 +20,7 @@ import regex as re
from vllm.entrypoints.chat_utils import make_tool_call_id
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionNamedToolChoiceParam,
ChatCompletionRequest,
)
from vllm.entrypoints.openai.engine.protocol import (
@@ -50,6 +51,8 @@ class Glm4MoeModelToolParser(ToolParser):
call, and diffs against what was previously sent to emit only new content.
"""
supports_required_and_named = False
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
super().__init__(tokenizer, tools)
# Stateful streaming fields
@@ -156,7 +159,25 @@ class Glm4MoeModelToolParser(ToolParser):
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
"""Adjust request parameters for tool call token handling."""
"""Adjust request parameters for tool call token handling.
For required/named tool_choice, skip setting structured_outputs
because GLM models output tool calls in XML format (per chat
template). Guided decoding would force JSON output, conflicting
with the XML format and causing parsing failures.
"""
if request.tools:
tc = request.tool_choice
if tc == "required" or isinstance(tc, ChatCompletionNamedToolChoiceParam):
# Do NOT call super().adjust_request() for required/named,
# because it would set structured_outputs and force JSON
# output via guided decoding. GLM models use XML tool-call
# syntax (defined in the chat template), so guided decoding
# must be skipped to let the model output XML freely.
# The tool_parser handles extraction from XML output.
if request.tool_choice != "none":
request.skip_special_tokens = False
return request
request = super().adjust_request(request)
if request.tools and request.tool_choice != "none":
# Ensure tool call tokens (<tool_call>, </tool_call>) are not skipped
@@ -4,11 +4,11 @@ import logging
import math
import random
import librosa
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torchaudio.functional import melscale_fbanks
from transformers import AutoFeatureExtractor, AutoProcessor, BatchFeature
from transformers.feature_extraction_sequence_utils import (
SequenceFeatureExtractor,
@@ -129,17 +129,15 @@ class FilterbankFeatures(nn.Module):
self.pad_min_duration = 0.0
self.pad_direction = "both"
filterbanks = torch.tensor(
librosa.filters.mel(
sr=sample_rate,
n_fft=self.n_fft,
n_mels=nfilt,
fmin=lowfreq,
fmax=highfreq,
norm=mel_norm,
),
dtype=torch.float,
).unsqueeze(0)
filterbanks = melscale_fbanks(
n_freqs=self.n_fft // 2 + 1,
f_min=lowfreq,
f_max=highfreq,
n_mels=nfilt,
sample_rate=sample_rate,
norm=mel_norm,
mel_scale="slaney",
).T.unsqueeze(0)
self.register_buffer("fb", filterbanks)
# Calculate maximum sequence length
+16 -6
View File
@@ -27,6 +27,7 @@ class DTypeInfo:
EmbedDType = Literal["float32", "float16", "bfloat16", "fp8_e4m3", "fp8_e5m2"]
MmMetadataDType = Literal["int32", "int64", "uint8", "bool"]
Endianness = Literal["native", "big", "little"]
EncodingFormat = Literal["float", "base64", "bytes", "bytes_only"]
@@ -42,6 +43,15 @@ EMBED_DTYPES: Mapping[EmbedDType, DTypeInfo] = {
"fp8_e4m3": DTypeInfo(torch.float8_e4m3fn, torch.uint8, np.uint8),
"fp8_e5m2": DTypeInfo(torch.float8_e5m2, torch.uint8, np.uint8),
}
MM_METADATA_DTYPES: Mapping[MmMetadataDType, DTypeInfo] = {
"int32": DTypeInfo(torch.int32, torch.int32, np.int32),
"int64": DTypeInfo(torch.int64, torch.int64, np.int64),
"uint8": DTypeInfo(torch.uint8, torch.uint8, np.uint8),
"bool": DTypeInfo(torch.bool, torch.uint8, np.uint8),
}
_ALL_SERIAL_DTYPES: Mapping[str, DTypeInfo] = {
k: v for d in (EMBED_DTYPES, MM_METADATA_DTYPES) for k, v in d.items()
}
ENDIANNESS: tuple[Endianness, ...] = get_args(Endianness)
@@ -56,14 +66,14 @@ def tensor2base64(x: torch.Tensor) -> str:
def tensor2binary(
tensor: torch.Tensor,
embed_dtype: EmbedDType,
embed_dtype: "EmbedDType | MmMetadataDType",
endianness: Endianness,
) -> bytes:
assert isinstance(tensor, torch.Tensor)
assert embed_dtype in EMBED_DTYPES
assert embed_dtype in _ALL_SERIAL_DTYPES
assert endianness in ENDIANNESS
dtype_info = EMBED_DTYPES[embed_dtype]
dtype_info = _ALL_SERIAL_DTYPES[embed_dtype]
np_array = (
tensor.to(dtype_info.torch_dtype)
@@ -82,13 +92,13 @@ def tensor2binary(
def binary2tensor(
binary: bytes,
shape: tuple[int, ...],
embed_dtype: EmbedDType,
embed_dtype: "EmbedDType | MmMetadataDType",
endianness: Endianness,
) -> torch.Tensor:
assert embed_dtype in EMBED_DTYPES
assert embed_dtype in _ALL_SERIAL_DTYPES
assert endianness in ENDIANNESS
dtype_info = EMBED_DTYPES[embed_dtype]
dtype_info = _ALL_SERIAL_DTYPES[embed_dtype]
np_array = np.frombuffer(binary, dtype=dtype_info.numpy_view_dtype).reshape(shape)
+12 -24
View File
@@ -279,29 +279,25 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
)
def _ensure_on_device(self, layer, device):
"""One-time derivation of TQ buffers (rotation matrices, midpoints).
"""One-time derivation of TQ buffers (rotation matrix, midpoints).
Registered buffers (_tq_signs, _tq_centroids) are already on the
correct device via register_buffer + model.to(device).
The Hadamard rotation is shared across all layers: random sign
flips do not improve Lloyd-Max quantization quality because the
quantizer is symmetric around zero (sign-flipping a coordinate
maps it to the mirror centroid with identical distortion).
"""
if not hasattr(layer, "_tq_cached"):
D = layer._tq_signs.shape[0]
signs = layer._tq_signs.to(device=device, dtype=torch.float32)
D = self.head_size
# WHT rotation: orthonormal + self-inverse, enabling future
# Pure Hadamard: orthonormal + symmetric (H = H^T), enabling
# in-kernel butterfly fusion and trivial inverse for continuation.
H = _build_hadamard(D, str(device))
layer._tq_PiT = (signs.unsqueeze(1) * H).contiguous()
layer._tq_Pi = layer._tq_PiT.T.contiguous()
layer._tq_PiT = H
layer._tq_Pi = H
c = layer._tq_centroids.to(device=device, dtype=torch.float32)
# Precompute midpoints for threshold-based quantization
c_sorted, _ = c.sort()
layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2
# Decode buffers (_tq_mid_o_buf, _tq_output_buf, _tq_lse_buf)
# are pre-allocated via register_buffer in Attention.__init__
# and moved to GPU by model.to(device) — no allocation needed
# here. The memory profiler sees them before KV cache sizing.
layer._tq_cached = True
def do_kv_cache_update(
@@ -507,8 +503,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
# max_query_len == max_seq_len means no request has prior cached KV.
# Both are Python ints — no GPU sync.
if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len:
output = torch.empty(N, Hq, D, device=query.device, dtype=query.dtype)
flash_attn_varlen_func(
return flash_attn_varlen_func(
q=query,
k=key,
v=value,
@@ -518,9 +513,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
max_seqlen_k=attn_metadata.max_query_len,
softmax_scale=self.scale,
causal=True,
out=output,
)
return output
# Continuation or no flash_attn: per-request attention.
# For continuation chunks (seq_len > q_len), we must attend to
@@ -557,10 +550,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
if q_len == seq_len:
# First-chunk prefill: all K/V are in the current batch.
if _HAS_FLASH_ATTN:
out = torch.empty_like(q_seq)
_cu_2[1] = q_len
cu = _cu_2
flash_attn_varlen_func(
out = flash_attn_varlen_func(
q=q_seq,
k=k_seq,
v=v_seq,
@@ -570,7 +562,6 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
max_seqlen_k=q_len,
softmax_scale=self.scale,
causal=True,
out=out,
)
else:
q_t = q_seq.transpose(0, 1).contiguous()
@@ -733,10 +724,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
# Attention: q_len queries attending to seq_len K/V with causal mask
if _HAS_FLASH_ATTN:
output = torch.empty(q_len, Hq, D, device=device, dtype=query.dtype)
cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32)
cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32)
flash_attn_varlen_func(
return flash_attn_varlen_func(
q=query,
k=k_full,
v=v_full,
@@ -746,9 +736,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
max_seqlen_k=seq_len,
softmax_scale=self.scale,
causal=True,
out=output,
)
return output
else:
# SDPA fallback: expand KV for GQA, build causal mask
q_t = query.transpose(0, 1).unsqueeze(0) # (1, Hq, q_len, D)
@@ -13,6 +13,7 @@ from typing import Any
import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.v1.attention.ops.triton_decode_attention import (
_fwd_kernel_stage2,
@@ -22,10 +23,15 @@ _FP8_E4B15: dict[int, int] = {}
def _use_fp8_e4b15(device: int = 0) -> int:
"""Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0."""
"""Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0.
On non-CUDA platforms (e.g. XPU), always returns 0 (use e4nv format).
"""
if device not in _FP8_E4B15:
cap = torch.cuda.get_device_capability(device)
_FP8_E4B15[device] = 1 if cap < (8, 9) else 0
if current_platform.is_cuda_alike():
cap = torch.cuda.get_device_capability(device)
_FP8_E4B15[device] = 1 if cap < (8, 9) else 0
else:
_FP8_E4B15[device] = 0
return _FP8_E4B15[device]
@@ -137,12 +143,12 @@ def _tq_decode_stage1(
Block_table_ptr + bt_base + page_idx,
mask=kv_mask,
other=0,
)
).to(tl.int64)
slot_bases = (
block_nums * stride_cache_block
+ page_off * stride_cache_pos
+ kv_head * stride_cache_head
+ page_off.to(tl.int64) * stride_cache_pos
+ tl.cast(kv_head, tl.int64) * stride_cache_head
)
# ============================================================
@@ -350,11 +356,11 @@ def _tq_full_dequant_kv(
page_idx = pos // BLOCK_SIZE
page_off = pos % BLOCK_SIZE
block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx)
block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64)
slot_base = (
block_num * stride_cache_block
+ page_off * stride_cache_pos
+ hid * stride_cache_head
+ tl.cast(page_off, tl.int64) * stride_cache_pos
+ tl.cast(hid, tl.int64) * stride_cache_head
)
d_offs = tl.arange(0, BLOCK_D)
@@ -174,10 +174,13 @@ def _tq_fused_store_fp8(
slot = tl.load(Slot_mapping_ptr + token_idx)
if slot < 0:
return
blk = slot // BLOCK_SIZE
off = slot % BLOCK_SIZE
blk = (slot // BLOCK_SIZE).to(tl.int64)
off = (slot % BLOCK_SIZE).to(tl.int64)
head_idx_i64 = tl.cast(head_idx, tl.int64)
slot_base = (
blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
blk * stride_cache_block
+ off * stride_cache_pos
+ head_idx_i64 * stride_cache_head
)
base = pid * D
@@ -259,10 +262,13 @@ def _tq_fused_store_mse(
slot = tl.load(Slot_mapping_ptr + token_idx)
if slot < 0:
return
blk = slot // BLOCK_SIZE
off = slot % BLOCK_SIZE
blk = (slot // BLOCK_SIZE).to(tl.int64)
off = (slot % BLOCK_SIZE).to(tl.int64)
head_idx_i64 = tl.cast(head_idx, tl.int64)
slot_base = (
blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
blk * stride_cache_block
+ off * stride_cache_pos
+ head_idx_i64 * stride_cache_head
)
base = pid * D
+18 -4
View File
@@ -3758,6 +3758,15 @@ class GPUModelRunner(
return slot_mappings_by_gid, slot_mappings_by_layer
def _is_all_reqs_chunked_prefill(self) -> bool:
"""Check if all scheduled requests are marked to discard sampled tokens.
This is true when `discard_request_mask` is set for every scheduled
request (e.g., for chunked prefill requests that are not the last
prefill chunk)."""
num_reqs = self.input_batch.num_reqs
return bool(self.discard_request_mask.np[:num_reqs].all())
@torch.inference_mode()
def execute_model(
self,
@@ -4361,9 +4370,12 @@ class GPUModelRunner(
assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, (
"PP+async expects sampled_token_ids to have shape [num_reqs, 1]"
)
torch.distributed.broadcast(
sampled_token_ids, src=pp.rank, group=pp.device_group
)
# Skip for chunked prefill: sampled tokens are dummy
# and will be discarded, no need to broadcast.
if not self._is_all_reqs_chunked_prefill():
torch.distributed.broadcast(
sampled_token_ids, src=pp.rank, group=pp.device_group
)
def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None:
"""Receive sampled token ids broadcast from last PP stage"""
@@ -4372,7 +4384,9 @@ class GPUModelRunner(
num_reqs = self.input_batch.num_reqs
# `prev_sampled_token_ids` is expected to have shape [num_reqs, 1].
recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device)
torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
# skip for chunked prefill.
if not self._is_all_reqs_chunked_prefill():
torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
self.input_batch.prev_sampled_token_ids = recv
# construct `prev_req_id_to_index` here so `_prepare_input_ids`