forked from Karylab-cklius/vllm
Compare commits
28
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e329dbcb0f | ||
|
|
8f3ce1492e | ||
|
|
21eaa58f71 | ||
|
|
d0359f3e04 | ||
|
|
ed0622e3a8 | ||
|
|
b5f6c5f834 | ||
|
|
bfde49e287 | ||
|
|
153ba7f0f3 | ||
|
|
87518c3027 | ||
|
|
aeee7ef939 | ||
|
|
cda19ecf4d | ||
|
|
80b18230e0 | ||
|
|
d0697cc7b6 | ||
|
|
b0755523dc | ||
|
|
993859ceb0 | ||
|
|
48a65ccb02 | ||
|
|
55842a8d69 | ||
|
|
1f45e83756 | ||
|
|
a8bffaa133 | ||
|
|
5cdddddd4a | ||
|
|
6ef1efd51f | ||
|
|
58da4ee047 | ||
|
|
1ae11e2bfc | ||
|
|
251c18d1f8 | ||
|
|
512765d52d | ||
|
|
640cc9dd7d | ||
|
|
ceade1952c | ||
|
|
747256bb5d |
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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}")
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -53,7 +53,7 @@
|
||||
"default": "cuda"
|
||||
},
|
||||
"DEEPEP_COMMIT_HASH": {
|
||||
"default": "73b6ea4"
|
||||
"default": "9249c25"
|
||||
},
|
||||
"GIT_REPO_CHECK": {
|
||||
"default": "0"
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
[
|
||||
|
||||
@@ -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)
|
||||
@@ -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}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+200
-15
@@ -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)
|
||||
|
||||
+65
-13
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -94,6 +94,10 @@ _TOOL_PARSERS_TO_REGISTER = {
|
||||
"longcat_tool_parser",
|
||||
"LongcatFlashToolParser",
|
||||
),
|
||||
"mimo": (
|
||||
"qwen3xml_tool_parser",
|
||||
"Qwen3XMLToolParser",
|
||||
),
|
||||
"minimax_m2": (
|
||||
"minimax_m2_tool_parser",
|
||||
"MinimaxM2ToolParser",
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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`
|
||||
|
||||
Reference in New Issue
Block a user