iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model] Add MiniMax M3 support (#45381)

Browse Source 
Signed-off-by: youkaichao <youkaichao@gmail.com>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: functionstackx <47992694+functionstackx@users.noreply.github.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
youkaichao
2026-06-16 01:01:25 +08:00
committed by GitHub
co-authored by OpenAI Codex Isotr0py Thien Tran Bugen Zhao Jee Jee Li Roger Wang functionstackx Yongye Zhu Jee Jee Li
parent a3195fab7b
commit 0a1c5034f5
108 changed files with 14734 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3
View File
@@ -15,6 +15,9 @@ vllm/third_party/flashmla/flash_mla_interface.py
# DeepGEMM vendored package built from source
vllm/third_party/deep_gemm/
# fmha_sm100 vendored package built from source
vllm/third_party/fmha_sm100/
# triton jit
.triton
CMakeLists.txt
+2
View File
@@ -440,6 +440,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
"csrc/libtorch_stable/pos_encoding_kernels.cu"
"csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
"csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
"csrc/libtorch_stable/layernorm_kernels.cu"
"csrc/libtorch_stable/layernorm_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
@@ -1398,6 +1399,7 @@ endif()
# For CUDA we also build and ship some external projects.
if (VLLM_GPU_LANG STREQUAL "CUDA")
include(cmake/external_projects/deepgemm.cmake)
include(cmake/external_projects/fmha_sm100.cmake)
include(cmake/external_projects/flashmla.cmake)
include(cmake/external_projects/qutlass.cmake)
cmake/external_projects/fmha_sm100.cmake
+50
View File
@@ -0,0 +1,50 @@
include(FetchContent)
# If FMHA_SM100_SRC_DIR is set, fmha_sm100 is installed from that directory
# instead of downloading. This is useful for local MSA development.
if(DEFINED ENV{FMHA_SM100_SRC_DIR})
set(FMHA_SM100_SRC_DIR $ENV{FMHA_SM100_SRC_DIR})
endif()
if(FMHA_SM100_SRC_DIR)
FetchContent_Declare(
fmha_sm100
SOURCE_DIR ${FMHA_SM100_SRC_DIR}
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
)
else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
)
endif()
FetchContent_GetProperties(fmha_sm100)
if(NOT fmha_sm100_POPULATED)
FetchContent_Populate(fmha_sm100)
endif()
message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}")
add_custom_target(fmha_sm100)
install(FILES
"${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/__init__.py"
"${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/sparse.py"
DESTINATION vllm/third_party/fmha_sm100
COMPONENT fmha_sm100)
install(DIRECTORY "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/cute/"
DESTINATION vllm/third_party/fmha_sm100/cute
COMPONENT fmha_sm100
FILES_MATCHING
REGEX "/__pycache__(/.*)?$" EXCLUDE
REGEX ".*\\.pyc$" EXCLUDE
PATTERN "example.py" EXCLUDE
PATTERN "test_*.py" EXCLUDE
PATTERN "*.py"
PATTERN "build_k2q_csr.cu")
csrc/libtorch_stable/activation_kernels.cu
+77 -41
View File
@@ -10,11 +10,20 @@
namespace vllm {
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
// `alpha` and `beta` are applied to opposite operands:
// - alpha lives INSIDE the activation (the activated half): the gated
// activation computes act_half * sigmoid(alpha * act_half).
// - beta is added to the OTHER (non-activated) half before the multiply.
// So the result is always ACT(act_half, alpha) * (other_half + beta).
// Which half is which depends on `act_first` (see below). Defaults
// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
bool act_first, bool HAS_CLAMP>
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
const scalar_t& y,
const float limit) {
const float limit,
const float alpha,
const float beta) {
if constexpr (act_first) {
scalar_t gate = x;
scalar_t up = y;
@@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
gate = (scalar_t)fminf((float)gate, limit);
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
}
return ACT_FN(gate) * up;
// act_first: gate is the activated half -> alpha applies to gate;
// beta is added to up (the non-activated half).
return (scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta));
} else {
scalar_t gate = x;
scalar_t up = y;
@@ -30,55 +41,68 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
up = (scalar_t)fminf((float)up, limit);
}
return gate * ACT_FN(up);
// !act_first: up is the activated half -> alpha applies to up;
// beta is added to gate (the non-activated half).
return (scalar_t)(((float)gate + beta) * ACT_FN(up, alpha));
}
}
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
template <typename packed_t,
packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
bool act_first, bool HAS_CLAMP>
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
const packed_t& y,
const float limit) {
const float limit,
const float alpha,
const float beta) {
if constexpr (act_first) {
packed_t gate = x;
packed_t up = y;
float2 u = cast_to_float2(up);
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fminf(g.x, limit);
g.y = fminf(g.y, limit);
u.x = fmaxf(fminf(u.x, limit), -limit);
u.y = fmaxf(fminf(u.y, limit), -limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(PACKED_ACT_FN(gate), up);
// act_first: gate is the activated half -> alpha applies to gate;
// beta is added to up (the non-activated half).
float2 activated = cast_to_float2(PACKED_ACT_FN(gate, alpha));
activated.x *= u.x + beta;
activated.y *= u.y + beta;
return cast_to_packed<packed_t>(activated);
} else {
packed_t gate = x;
packed_t up = y;
float2 g = cast_to_float2(gate);
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fmaxf(fminf(g.x, limit), -limit);
g.y = fmaxf(fminf(g.y, limit), -limit);
u.x = fminf(u.x, limit);
u.y = fminf(u.y, limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(gate, PACKED_ACT_FN(up));
// !act_first: up is the activated half -> alpha applies to up;
// beta is added to gate (the non-activated half).
float2 activated = cast_to_float2(PACKED_ACT_FN(up, alpha));
activated.x *= g.x + beta;
activated.y *= g.y + beta;
return cast_to_packed<packed_t>(activated);
}
}
// Activation and gating kernel template.
template <typename scalar_t, typename packed_t,
scalar_t (*ACT_FN)(const scalar_t&),
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
scalar_t (*ACT_FN)(const scalar_t&, const float),
packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
__global__ void act_and_mul_kernel(
scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., 2, d]
const int d, const float limit) {
const int d, const float limit, const float alpha, const float beta) {
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
const scalar_t* y_ptr = x_ptr + d;
scalar_t* out_ptr = out + blockIdx.x * d;
@@ -105,7 +129,7 @@ __global__ void act_and_mul_kernel(
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
x.elts[j] =
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
x.elts[j], y.elts[j], limit);
x.elts[j], y.elts[j], limit, alpha, beta);
}
if constexpr (use_256b) {
st256(x, &out_vec[i]);
@@ -118,29 +142,34 @@ __global__ void act_and_mul_kernel(
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
out_ptr[idx] =
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
x, y, limit, alpha, beta);
}
}
}
// Gated activations take an `alpha` argument that scales the sigmoid input
// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
// non-default alpha. Activations that do not use alpha simply ignore it.
template <typename T>
__device__ __forceinline__ T silu_kernel(const T& x) {
// x * sigmoid(x)
return (T)(((float)x) / (1.0f + expf((float)-x)));
__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
// x * sigmoid(alpha * x)
return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
}
template <typename packed_t>
__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
// x * sigmoid(x)
__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
const float alpha) {
// x * sigmoid(alpha * x)
float2 fval = cast_to_float2(val);
fval.x = fval.x / (1.0f + expf(-fval.x));
fval.y = fval.y / (1.0f + expf(-fval.y));
fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
return cast_to_packed<packed_t>(fval);
}
template <typename T>
__device__ __forceinline__ T gelu_kernel(const T& x) {
__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'none' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -150,7 +179,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
}
template <typename packed_t>
__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'none' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -162,7 +192,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
}
template <typename T>
__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
__device__ __forceinline__ T gelu_tanh_kernel(const T& x,
const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'tanh' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -176,7 +207,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
template <typename packed_t>
__device__ __forceinline__ packed_t
packed_gelu_tanh_kernel(const packed_t& val) {
packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'tanh' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -202,7 +233,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
// clamped (max only) and up input is clamped (both sides) before the
// activation function is applied.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
HAS_CLAMP, LIMIT) \
HAS_CLAMP, LIMIT, ALPHA, BETA) \
auto dtype = input.scalar_type(); \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
@@ -230,7 +261,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<scalar_t>(), \
input.const_data_ptr<scalar_t>(), d, LIMIT); \
input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
}); \
} else { \
VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
@@ -240,7 +271,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<scalar_t>(), \
input.const_data_ptr<scalar_t>(), d, LIMIT); \
input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
}); \
} \
} else { \
@@ -252,7 +283,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
d, LIMIT); \
d, LIMIT, ALPHA, BETA); \
}); \
}
@@ -260,14 +291,18 @@ void silu_and_mul(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, false, 0.0f);
true, false, 0.0f, 1.0f, 0.0f);
}
void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input, // [..., 2 * d]
double limit) {
double limit, double alpha, double beta) {
// out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
// * (up.clamp(+-limit) + beta)
// alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, true, (float)limit);
true, true, (float)limit, (float)alpha,
(float)beta);
}
void mul_and_silu(torch::stable::Tensor& out, // [..., d]
@@ -276,21 +311,22 @@ void mul_and_silu(torch::stable::Tensor& out, // [..., d]
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
// applies the silu to the latter half of the input.
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
false, false, 0.0f);
false, false, 0.0f, 1.0f, 0.0f);
}
void gelu_and_mul(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
true, false, 0.0f);
true, false, 0.0f, 1.0f, 0.0f);
}
void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
vllm::packed_gelu_tanh_kernel, true, false,
0.0f, 1.0f, 0.0f);
}
namespace vllm {
csrc/libtorch_stable/fp32_router_gemm.cu
+42 -39
View File
@@ -175,49 +175,52 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
}
// ---------------------------------------------------------------------------
// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
// Explicit instantiations: M=1..32, for both input types, for the supported
// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3].
// ---------------------------------------------------------------------------
#define INSTANTIATE(T, M) \
template void invokeFp32RouterGemm<T, M, 256, 3072>( \
float*, T const*, float const*, cudaStream_t);
#define INSTANTIATE(T, M, E, H) \
template void invokeFp32RouterGemm<T, M, E, H>(float*, T const*, \
float const*, cudaStream_t);
#define INSTANTIATE_ALL(T) \
INSTANTIATE(T, 1) \
INSTANTIATE(T, 2) \
INSTANTIATE(T, 3) \
INSTANTIATE(T, 4) \
INSTANTIATE(T, 5) \
INSTANTIATE(T, 6) \
INSTANTIATE(T, 7) \
INSTANTIATE(T, 8) \
INSTANTIATE(T, 9) \
INSTANTIATE(T, 10) \
INSTANTIATE(T, 11) \
INSTANTIATE(T, 12) \
INSTANTIATE(T, 13) \
INSTANTIATE(T, 14) \
INSTANTIATE(T, 15) \
INSTANTIATE(T, 16) \
INSTANTIATE(T, 17) \
INSTANTIATE(T, 18) \
INSTANTIATE(T, 19) \
INSTANTIATE(T, 20) \
INSTANTIATE(T, 21) \
INSTANTIATE(T, 22) \
INSTANTIATE(T, 23) \
INSTANTIATE(T, 24) \
INSTANTIATE(T, 25) \
INSTANTIATE(T, 26) \
INSTANTIATE(T, 27) \
INSTANTIATE(T, 28) \
INSTANTIATE(T, 29) \
INSTANTIATE(T, 30) \
INSTANTIATE(T, 31) \
INSTANTIATE(T, 32)
#define INSTANTIATE_ALL(T, E, H) \
INSTANTIATE(T, 1, E, H) \
INSTANTIATE(T, 2, E, H) \
INSTANTIATE(T, 3, E, H) \
INSTANTIATE(T, 4, E, H) \
INSTANTIATE(T, 5, E, H) \
INSTANTIATE(T, 6, E, H) \
INSTANTIATE(T, 7, E, H) \
INSTANTIATE(T, 8, E, H) \
INSTANTIATE(T, 9, E, H) \
INSTANTIATE(T, 10, E, H) \
INSTANTIATE(T, 11, E, H) \
INSTANTIATE(T, 12, E, H) \
INSTANTIATE(T, 13, E, H) \
INSTANTIATE(T, 14, E, H) \
INSTANTIATE(T, 15, E, H) \
INSTANTIATE(T, 16, E, H) \
INSTANTIATE(T, 17, E, H) \
INSTANTIATE(T, 18, E, H) \
INSTANTIATE(T, 19, E, H) \
INSTANTIATE(T, 20, E, H) \
INSTANTIATE(T, 21, E, H) \
INSTANTIATE(T, 22, E, H) \
INSTANTIATE(T, 23, E, H) \
INSTANTIATE(T, 24, E, H) \
INSTANTIATE(T, 25, E, H) \
INSTANTIATE(T, 26, E, H) \
INSTANTIATE(T, 27, E, H) \
INSTANTIATE(T, 28, E, H) \
INSTANTIATE(T, 29, E, H) \
INSTANTIATE(T, 30, E, H) \
INSTANTIATE(T, 31, E, H) \
INSTANTIATE(T, 32, E, H)
INSTANTIATE_ALL(float)
INSTANTIATE_ALL(__nv_bfloat16)
INSTANTIATE_ALL(float, 256, 3072)
INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
INSTANTIATE_ALL(float, 128, 6144)
INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
#undef INSTANTIATE_ALL
#undef INSTANTIATE
csrc/libtorch_stable/fp32_router_gemm_entry.cu
+42 -18
View File
@@ -22,36 +22,42 @@ inline int getSMVersion() {
} // namespace
static constexpr int FP32_NUM_EXPERTS = 256;
static constexpr int FP32_HIDDEN_DIM = 3072;
static constexpr int FP32_MAX_TOKENS = 32;
// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
return (hidden_dim == 3072 && num_experts == 256) ||
(hidden_dim == 6144 && num_experts == 128);
}
// Forward declarations — 4 template params must match fp32_router_gemm.cu
template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
void invokeFp32RouterGemm(float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream);
// LoopUnroller templated on InputT
template <typename InputT, int kBegin, int kEnd>
// LoopUnroller templated on InputT, kNumExperts and kHiddenDim
template <typename InputT, int kNumExperts, int kHiddenDim, int kBegin,
int kEnd>
struct Fp32LoopUnroller {
static void unroll(int num_tokens, float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream) {
if (num_tokens == kBegin) {
invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
invokeFp32RouterGemm<InputT, kBegin, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
} else {
Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
mat_a, mat_b, stream);
Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kBegin + 1,
kEnd>::unroll(num_tokens, output, mat_a, mat_b, stream);
}
}
};
template <typename InputT, int kEnd>
struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
template <typename InputT, int kNumExperts, int kHiddenDim, int kEnd>
struct Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kEnd, kEnd> {
static void unroll(int num_tokens, float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream) {
if (num_tokens == kEnd) {
invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
invokeFp32RouterGemm<InputT, kEnd, kNumExperts, kHiddenDim>(
output, mat_a, mat_b, stream);
} else {
throw std::invalid_argument(
@@ -60,6 +66,23 @@ struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
}
};
// Dispatch over the supported (num_experts, hidden_dim) pairs.
template <typename InputT>
void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
float* output, InputT const* mat_a,
float const* mat_b, cudaStream_t stream) {
if (num_experts == 256 && hidden_dim == 3072) {
Fp32LoopUnroller<InputT, 256, 3072, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, stream);
} else if (num_experts == 128 && hidden_dim == 6144) {
Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, output, mat_a, mat_b, stream);
} else {
throw std::invalid_argument(
"fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
}
}
void fp32_router_gemm(
torch::stable::Tensor& output, // [num_tokens, num_experts]
torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim]
@@ -85,10 +108,10 @@ void fp32_router_gemm(
STD_TORCH_CHECK(
mat_a.size(1) == mat_b.size(1),
"fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
"fp32_router_gemm: expected hidden_dim=3072");
STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
"fp32_router_gemm: expected num_experts=256");
STD_TORCH_CHECK(
fp32_router_gemm_supported(hidden_dim, num_experts),
"fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
"(3072, 256) and (6144, 128)");
STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
"fp32_router_gemm: num_tokens must be in [0, 32]");
STD_TORCH_CHECK(
@@ -113,12 +136,13 @@ void fp32_router_gemm(
if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
auto const* mat_a_ptr =
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
dispatchFp32RouterGemm<__nv_bfloat16>(num_experts, hidden_dim, num_tokens,
out_ptr, mat_a_ptr, mat_b_ptr,
stream);
} else {
auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
dispatchFp32RouterGemm<float>(num_experts, hidden_dim, num_tokens, out_ptr,
mat_a_ptr, mat_b_ptr, stream);
}
}
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+635
View File
@@ -0,0 +1,635 @@
/*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* Horizontally-fused MiniMax-M3 attention pre-processing kernel.
*
* Replaces the per-token Python sequence in
* ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
*
* q = q_norm(q); k = k_norm(k); q, k = rotary_emb(pos, q, k)
* index_q = index_q_norm(index_q); index_k = index_k_norm(index_k)
* index_q, index_k = rotary_emb(pos, index_q, index_k)
* _insert_kv(k, v, index_k)
*
* All branches share head_dim=128 and the *same* partial-NeoX RoPE table
* (``rotary_dim`` rotated, the trailing dims pass through). The four norms
* are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
* independent weights.
*
* Everything lives in a single fused ``qkv`` tensor. The sparse layer's
* fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
*
* [ q | k | v | index_q | index_k ] (the "5 results")
*
* while the dense layer emits just ``[ q | k | v ]``. The kernel reads the
* index branch straight out of that packed row -- no separate index tensors.
*
* One kernel, one grid; each warp owns one (token, head-slot) pair. Slot
* enumeration per token:
* [0, nq) Q heads -> norm(q_w) + RoPE, write
* qkv [nq, nq+nkv) K heads -> norm(k_w) + RoPE, write
* qkv
* (+ insert into key cache)
* [nq+nkv, nq+2*nkv) V heads -> insert into value cache
* IQ heads (niq) -> norm(iq_w) + RoPE, write iq
* IK (1) -> norm(ik_w) + RoPE
* (+ insert into index cache)
*
* The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
* fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
*
* Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
* template bools (3 instantiations: dense <false,false>, sparse-profiling
* <true,false>, sparse-serving <true,true>), so the index slots, the V slots
* and the cache inserts fold away entirely on paths that don't use them. The
* dense layer passes no caches/index: norm+RoPE happens in place and the
* generic ``Attention`` layer owns the cache write.
*
* Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
* ``qkv`` tensor. Caches (bf16) are scatter-written by slot.
*/
#include <cmath>
#include <cuda_runtime.h>
#include <type_traits>
#include "torch_utils.h"
#include "../cuda_compat.h"
#include "../type_convert.cuh"
#include "dispatch_utils.h"
#ifndef FINAL_MASK
#ifdef USE_ROCM
#define FINAL_MASK 0xffffffffffffffffULL
#else
#define FINAL_MASK 0xffffffffu
#endif
#endif
namespace vllm {
namespace minimax_m3_fused_ops {
namespace {
inline int getSMVersion() {
auto* props = get_device_prop();
return props->major * 10 + props->minor;
}
} // namespace
// ────────────────────────────────────────────────────────────────────────────
// Constants (hard-coded for MiniMax-M3-preview).
// ────────────────────────────────────────────────────────────────────────────
constexpr int kHeadDim = 128;
constexpr int kNumLanes = 32;
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 4
// ────────────────────────────────────────────────────────────────────────────
// Helpers
// ────────────────────────────────────────────────────────────────────────────
__device__ __forceinline__ float warpReduceSum(float val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
}
return val;
}
// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``), rounded
// back to scalar_t like the materialized unfused norm output, followed by
// partial NeoX RoPE on the leading ``rotary_dim`` dims. Each lane owns
// ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
template <typename scalar_t>
__device__ __forceinline__ void normAndRope(
float (&elems)[kElemsPerLane], int const laneId, float const eps,
scalar_t const* __restrict__ weight, // [kHeadDim] or nullptr (no norm)
bool const do_rope, int const rotary_dim,
scalar_t const* __restrict__ cos_ptr, // cos_sin_cache + pos*rotary_dim
bool const apply_norm) {
// ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
if (apply_norm) {
float sumsq = 0.0f;
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
sumsq = warpReduceSum(sumsq);
float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
int const dim = laneId * kElemsPerLane + i;
float const w = 1.0f + static_cast<float>(weight[dim]);
elems[i] = elems[i] * rms_rcp * w;
}
}
// ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
// half = rotary_dim/2. Pair (i, i+half) for i in [0, half). Lane L owns
// dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
// first half (own=x[i]) or second half (own=x[i+half]); its partner lives
// ``half/4`` lanes away (XOR with that distance).
if (do_rope) {
int const half = rotary_dim / 2;
int const dim0 = laneId * kElemsPerLane;
bool const in_rope = dim0 < rotary_dim;
int const lane_xor = half / kElemsPerLane; // partner-lane distance
float partner[kElemsPerLane];
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
}
if (in_rope) {
bool const first_half = dim0 < half;
int const i_base = first_half ? dim0 : (dim0 - half); // cos/sin index
scalar_t const* sin_ptr = cos_ptr + half;
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
float const c = static_cast<float>(cos_ptr[i_base + i]);
float const s = static_cast<float>(sin_ptr[i_base + i]);
if (first_half) {
elems[i] = elems[i] * c - partner[i] * s;
} else {
elems[i] = elems[i] * c + partner[i] * s;
}
}
}
}
}
// Load 4 contiguous bf16 -> 4 fp32 registers.
template <typename scalar_t>
__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
float (&elems)[kElemsPerLane]) {
using Converter = vllm::_typeConvert<scalar_t>;
uint2 v = *reinterpret_cast<uint2 const*>(src);
auto const* p =
reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
#pragma unroll
for (int i = 0; i < kElemsPerLane / 2; i++) {
float2 f2 = Converter::convert(p[i]);
elems[2 * i] = f2.x;
elems[2 * i + 1] = f2.y;
}
}
// Store 4 fp32 registers -> 4 contiguous bf16.
template <typename scalar_t>
__device__ __forceinline__ void storeElems(
scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
using Converter = vllm::_typeConvert<scalar_t>;
uint2 v;
auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
#pragma unroll
for (int i = 0; i < kElemsPerLane / 2; i++) {
p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
}
*reinterpret_cast<uint2*>(dst) = v;
}
// ────────────────────────────────────────────────────────────────────────────
// Kernel
// ────────────────────────────────────────────────────────────────────────────
// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
// Each warp = one (token, slot).
//
// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
// branch decisions that distinguish the dense layer from the sparse layer
// (index slots, KV/index inserts, V slots) fold away per instantiation.
// Three instantiations are built: dense <false,false>, sparse-profiling
// <true,false> and sparse-serving <true,true>. Slots per token:
// Q : nq (always — norm+RoPE)
// K : nkv (always — norm+RoPE; +K-cache insert)
// V : nkv only if kInsertKV (V-cache insert; no warps in dense)
// IQ: niq only if kIsSparse (norm+RoPE)
// IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert)
template <typename scalar_t, bool kIsSparse, bool kInsertKV>
__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
scalar_t* __restrict__ index_q_out, // [N, niq*128] contiguous, or nullptr
scalar_t const* __restrict__ q_norm_w,
scalar_t const* __restrict__ k_norm_w,
scalar_t const* __restrict__ iq_norm_w,
scalar_t const* __restrict__ ik_norm_w,
scalar_t const* __restrict__ cos_sin_cache, // [max_pos, rotary_dim]
int64_t const* __restrict__ positions, // [N] i64
int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr
int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
scalar_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
float const eps, int const rotary_dim, int const num_tokens, int const nq,
int const nkv, int const niq, int const block_size,
// kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
// The head_dim (last) dim is always innermost-contiguous (stride 1), so the
// NHD/HND layout choice is fully captured by these four strides: NHD keeps
// s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
int64_t const kv_s_head) {
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
// _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
// runs with bf16/fp16 inputs in practice. Discard the bf16 body there.
if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
return;
} else {
#endif
int const warpsPerBlock = blockDim.x / 32;
int const laneId = threadIdx.x % 32;
int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
// Slot layout (compile-time gated: dense has neither V nor index slots).
int const v_slots = kInsertKV ? nkv : 0;
int const idx_slots = kIsSparse ? niq + 1 : 0;
int const slots_per_token = nq + nkv + v_slots + idx_slots;
int const tokenIdx = globalWarpIdx / slots_per_token;
int const slot = globalWarpIdx % slots_per_token;
if (tokenIdx >= num_tokens) return;
// Slot boundaries.
int const k_begin = nq;
int const v_begin = nq + nkv; // valid only when kInsertKV
int const iq_begin = nq + nkv + v_slots; // index block start
int const ik_slot = iq_begin + niq; // valid only when kIsSparse
bool const isQ = slot < k_begin;
bool const isK = slot >= k_begin && slot < v_begin;
bool isV = false;
if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
bool isIQ = false, isIK = false;
if constexpr (kIsSparse) {
isIQ = slot >= iq_begin && slot < ik_slot;
isIK = slot == ik_slot;
}
int const dim_base = laneId * kElemsPerLane;
// Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
// layer additionally packs [index_q (niq heads) | index_k (1 head)].
int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
// ── Resolve source pointer + per-branch parameters. ────────────────────
scalar_t* row_ptr = nullptr; // in-place output location
scalar_t const* norm_w = nullptr; // nullptr -> skip norm (V)
bool do_rope = true;
int head = 0; // kv head index for inserts
if (isQ) {
row_ptr =
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
norm_w = q_norm_w;
} else if (isK) {
head = slot - k_begin;
row_ptr =
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
norm_w = k_norm_w;
} else if (isV) {
// qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
// correct in-tensor offset.
head = slot - v_begin;
row_ptr =
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
norm_w = nullptr; // V: no norm, no rope
do_rope = false;
} else if (isIQ) {
// index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
int const ih = slot - iq_begin;
row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
(nq + 2 * nkv + ih) * kHeadDim;
norm_w = iq_norm_w;
} else { // isIK -- single shared index key at (nq+2*nkv+niq)*128.
row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
(nq + 2 * nkv + niq) * kHeadDim;
norm_w = ik_norm_w;
}
// Store destination. Q and index_q are gathered into dedicated contiguous
// output buffers (when provided) so the downstream SM100 sparse kernel's
// flat TMA descriptor can address them as [tokens*heads, head_dim]; this
// folds the de-interleaving into the store the kernel already does, instead
// of a separate q.contiguous() copy. Everything else stays in place.
scalar_t* store_ptr = row_ptr;
if (isQ && q_out != nullptr) {
store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
slot * kHeadDim;
} else if (isIQ && index_q_out != nullptr) {
store_ptr = index_q_out +
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
(slot - iq_begin) * kHeadDim;
}
// PDL: wait for the predecessor kernel (the qkv-projection GEMM that
// produces ``qkv``) to finish before touching any global memory. No-op
// when PDL is not enabled on the launch. The CUDA runtime wrapper emits
// the griddepcontrol.wait PTX with the required memory clobber internally.
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
// ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
float elems[kElemsPerLane];
loadElems<scalar_t>(row_ptr + dim_base, elems);
if (!isV) {
int64_t const pos = positions[tokenIdx];
scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
cos_ptr, /*apply_norm=*/norm_w != nullptr);
storeElems<scalar_t>(store_ptr + dim_base, elems);
}
// ── Cache inserts (sparse serving only). ───────────────────────────────
if constexpr (kInsertKV) {
// Guard (not early-return) so every thread reaches the PDL trigger below.
int64_t const sm = (isK || isV)
? slot_mapping[tokenIdx]
: (isIK ? index_slot_mapping[tokenIdx] : -1);
if (sm >= 0) { // skip padded / unscheduled tokens
if (isIK) {
scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
storeElems<scalar_t>(dst, elems);
} else if (isK || isV) {
// kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
// Paging is logical (block = sm/block_size, token = sm%block_size);
// the physical NHD/HND layout is honoured via the passed strides.
int64_t const b = sm / block_size;
int64_t const t = sm % block_size;
int const kv = isK ? 0 : 1;
int64_t const off =
b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
storeElems<scalar_t>(kv_cache + off + dim_base, elems);
}
}
}
// PDL: signal that this kernel is done so a dependent successor may launch
// early. No-op when PDL is not enabled on the launch.
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
#endif
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
}
#endif
}
// ────────────────────────────────────────────────────────────────────────────
// Launch wrapper
// ────────────────────────────────────────────────────────────────────────────
template <typename scalar_t>
void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
scalar_t const* q_norm_w, scalar_t const* k_norm_w,
scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
scalar_t const* cos_sin_cache,
int64_t const* positions, int64_t const* slot_mapping,
int64_t const* index_slot_mapping, scalar_t* kv_cache,
scalar_t* index_cache, float const eps,
int const rotary_dim, int const num_tokens,
int const nq, int const nkv, int const niq,
int const block_size, int64_t const kv_s_block,
int64_t const kv_s_kv, int64_t const kv_s_token,
int64_t const kv_s_head, bool const has_index,
bool const insert_kv, cudaStream_t stream) {
// Slot count must match the kernel's compile-time gating.
int const v_slots = insert_kv ? nkv : 0;
int const idx_slots = has_index ? niq + 1 : 0;
int const slots_per_token = nq + nkv + v_slots + idx_slots;
constexpr int kBlockSize = 256;
constexpr int kWarpsPerBlock = kBlockSize / 32;
int64_t const total_warps =
static_cast<int64_t>(num_tokens) * slots_per_token;
int const grid =
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
if (grid == 0) return;
#ifndef USE_ROCM
// PDL: enable programmatic stream serialization whenever the hardware
// supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable, so
// leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
static int const sm_version = getSMVersion();
cudaLaunchConfig_t config;
config.gridDim = dim3(grid);
config.blockDim = dim3(kBlockSize);
config.dynamicSmemBytes = 0;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.attrs = attrs;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
#define LAUNCH(IS_SPARSE, INSERT) \
cudaLaunchKernelEx( \
&config, \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>, \
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \
cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, \
kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
#else
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
// clang-format off
#define LAUNCH(IS_SPARSE, INSERT) \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT> \
<<<grid, kBlockSize, 0, stream>>>( \
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \
ik_norm_w, cos_sin_cache, positions, slot_mapping, \
index_slot_mapping, kv_cache, index_cache, eps, rotary_dim, \
num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
kv_s_token, kv_s_head)
// clang-format on
#endif
if (has_index) {
if (insert_kv) {
LAUNCH(true, true); // sparse serving
} else {
LAUNCH(true, false); // sparse profiling
}
} else {
// Dense layer: never has an index branch and never inserts here (the
// generic Attention layer owns the KV insert).
LAUNCH(false, false);
}
#undef LAUNCH
}
} // namespace minimax_m3_fused_ops
} // namespace vllm
// ────────────────────────────────────────────────────────────────────────────
// Torch op wrapper
// ────────────────────────────────────────────────────────────────────────────
void fused_minimax_m3_qknorm_rope_kv_insert(
torch::stable::Tensor& qkv, // [N, qkv_row] (packs index if sparse)
torch::stable::Tensor const& q_norm_weight, // [128]
torch::stable::Tensor const& k_norm_weight, // [128]
torch::stable::Tensor const& cos_sin_cache, // [max_pos, rotary_dim]
torch::stable::Tensor const& positions, // [N] i64
int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
std::optional<torch::stable::Tensor> index_q_norm_weight, // [128]
std::optional<torch::stable::Tensor> index_k_norm_weight, // [128]
int64_t num_index_heads, // niq; 0 => dense
std::optional<torch::stable::Tensor> slot_mapping, // [N] i64
std::optional<torch::stable::Tensor> index_slot_mapping, // [N] i64
std::optional<torch::stable::Tensor> kv_cache, // [nb,2,bs,nkv,128]
std::optional<torch::stable::Tensor> index_cache, // [nb,bs,128]
int64_t block_size,
std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
std::optional<torch::stable::Tensor>
index_q_out) { // [N, niq*128] contiguous
STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
"qkv must be contiguous CUDA");
STD_TORCH_CHECK(
positions.is_cuda() &&
positions.scalar_type() == torch::headeronly::ScalarType::Long,
"positions must be int64 CUDA");
STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
"cos_sin_cache must be contiguous CUDA");
STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
"cos_sin_cache dtype must match qkv");
STD_TORCH_CHECK(
cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
"cos_sin_cache shape [max_pos, rotary_dim]");
STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
k_norm_weight.scalar_type() == qkv.scalar_type(),
"q/k norm weight dtype must match qkv");
STD_TORCH_CHECK(
q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
"q/k norm weight must have 128 elements");
STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
"rotary_dim must be a positive multiple of 8 and <= 128");
int const num_tokens = static_cast<int>(qkv.size(0));
int const nq = static_cast<int>(num_heads);
int const nkv = static_cast<int>(num_kv_heads);
int const niq = static_cast<int>(num_index_heads);
// The sparse layer packs the index branch ([index_q (niq heads) | index_k
// (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
bool const has_index = niq > 0;
bool const insert_kv = kv_cache.has_value();
int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
int const expected_row =
(nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
STD_TORCH_CHECK(qkv.size(1) == expected_row,
"qkv last dim must be (num_heads + 2*num_kv_heads"
" + num_index_heads + 1) * 128 for sparse, "
"(num_heads + 2*num_kv_heads) * 128 for dense");
// Only the sparse layer inserts here (dense lets the generic Attention layer
// own the KV write); there is no dense+insert kernel instantiation.
STD_TORCH_CHECK(
!insert_kv || has_index,
"insert mode (kv_cache) requires the index branch (sparse layer)");
if (has_index) {
STD_TORCH_CHECK(
index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
"index branch requires both index norm weights");
STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
index_k_norm_weight->scalar_type() == qkv.scalar_type(),
"index norm weights dtype must match qkv");
STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
index_k_norm_weight->numel() == kHeadDim,
"index norm weights must have 128 elements");
}
// kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
// off the tensor so the kernel honours whatever physical layout the attention
// backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
// op argument is needed -- the strides ride along with the tensor itself.
int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
torch::stable::Tensor const* effective_index_slot_mapping = nullptr;
if (insert_kv) {
STD_TORCH_CHECK(
slot_mapping.has_value() && slot_mapping->is_cuda() &&
slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
"insert mode requires int64 CUDA slot_mapping");
STD_TORCH_CHECK(
!index_slot_mapping.has_value() ||
(index_slot_mapping->is_cuda() &&
index_slot_mapping->scalar_type() ==
torch::headeronly::ScalarType::Long &&
index_slot_mapping->numel() == slot_mapping->numel()),
"index_slot_mapping must be int64 CUDA with slot_mapping length");
STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
"kv_cache dtype must match qkv (bf16 cache only)");
STD_TORCH_CHECK(index_cache.has_value() &&
index_cache->scalar_type() == qkv.scalar_type(),
"insert mode requires matching index_cache");
STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
"kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
"head_dim (stride(4)==1)");
kv_s_block = kv_cache->stride(0);
kv_s_kv = kv_cache->stride(1);
kv_s_token = kv_cache->stride(2);
kv_s_head = kv_cache->stride(3);
effective_index_slot_mapping = index_slot_mapping.has_value()
? &index_slot_mapping.value()
: &slot_mapping.value();
}
// Optional contiguous gather targets: when given, the normed/roped q (and
// index_q) are written here instead of in place, so callers avoid a separate
// .contiguous() copy. index_q_out only makes sense on the sparse path.
if (q_out.has_value()) {
STD_TORCH_CHECK(
q_out->is_cuda() && q_out->is_contiguous() &&
q_out->scalar_type() == qkv.scalar_type(),
"q_out must be a contiguous CUDA tensor matching qkv dtype");
STD_TORCH_CHECK(
q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
"q_out must have num_tokens * num_heads * 128 elements");
}
if (index_q_out.has_value()) {
STD_TORCH_CHECK(
has_index,
"index_q_out requires the index branch (num_index_heads > 0)");
STD_TORCH_CHECK(
index_q_out->is_cuda() && index_q_out->is_contiguous() &&
index_q_out->scalar_type() == qkv.scalar_type(),
"index_q_out must be a contiguous CUDA tensor matching qkv dtype");
STD_TORCH_CHECK(index_q_out->numel() ==
static_cast<int64_t>(num_tokens) * niq * kHeadDim,
"index_q_out must have num_tokens * num_index_heads * 128 "
"elements");
}
const torch::stable::accelerator::DeviceGuard device_guard(
qkv.get_device_index());
auto stream = get_current_cuda_stream(qkv.get_device_index());
VLLM_STABLE_DISPATCH_HALF_TYPES(
qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
using st = scalar_t;
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
reinterpret_cast<st*>(qkv.data_ptr()),
q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr())
: nullptr,
index_q_out.has_value()
? reinterpret_cast<st*>(index_q_out->data_ptr())
: nullptr,
reinterpret_cast<st const*>(q_norm_weight.data_ptr()),
reinterpret_cast<st const*>(k_norm_weight.data_ptr()),
has_index
? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())
: nullptr,
has_index
? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())
: nullptr,
reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),
reinterpret_cast<int64_t const*>(positions.data_ptr()),
insert_kv
? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())
: nullptr,
insert_kv ? reinterpret_cast<int64_t const*>(
effective_index_slot_mapping->data_ptr())
: nullptr,
insert_kv ? reinterpret_cast<st*>(kv_cache->data_ptr()) : nullptr,
(insert_kv && has_index)
? reinterpret_cast<st*>(index_cache->data_ptr())
: nullptr,
static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
nq, nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv,
kv_s_token, kv_s_head, has_index, insert_kv, stream);
});
}
csrc/libtorch_stable/ops.h
+20 -1
View File
@@ -281,6 +281,24 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
int64_t const nranks, double const eps);
#endif
// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
// the index branch and scatters k/v/index_k into their paged caches.
void fused_minimax_m3_qknorm_rope_kv_insert(
torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
torch::stable::Tensor const& k_norm_weight,
torch::stable::Tensor const& cos_sin_cache,
torch::stable::Tensor const& positions, int64_t num_heads,
int64_t num_kv_heads, int64_t rotary_dim, double eps,
std::optional<torch::stable::Tensor> index_q_norm_weight,
std::optional<torch::stable::Tensor> index_k_norm_weight,
int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
std::optional<torch::stable::Tensor> index_slot_mapping,
std::optional<torch::stable::Tensor> kv_cache,
std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
std::optional<torch::stable::Tensor> q_out,
std::optional<torch::stable::Tensor> index_q_out);
// Sampler kernels (shared CUDA/ROCm)
void apply_repetition_penalties_(
torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -346,7 +364,8 @@ void free_shared_buffer(int64_t buffer);
// Activation kernels (shared CUDA/ROCm)
void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void silu_and_mul_clamp(torch::stable::Tensor& out,
torch::stable::Tensor& input, double limit);
torch::stable::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_tanh_and_mul(torch::stable::Tensor& out,
csrc/libtorch_stable/torch_bindings.cpp
+18 -2
View File
@@ -461,6 +461,18 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"float eps) -> (Tensor, Tensor)");
#endif
// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
ops.def(
"fused_minimax_m3_qknorm_rope_kv_insert("
"Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
"Tensor cos_sin_cache, Tensor positions, int num_heads, "
"int num_kv_heads, int rotary_dim, float eps, "
"Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
"int num_index_heads, "
"Tensor? slot_mapping, Tensor? index_slot_mapping, "
"Tensor!? kv_cache, Tensor!? index_cache, "
"int block_size, Tensor!? q_out, Tensor!? index_q_out) -> ()");
// Apply repetition penalties to logits in-place.
ops.def(
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -488,9 +500,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
// SwiGLU activation with input clamping.
// alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
// the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
ops.def(
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
"-> ()");
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
"float alpha=1.0, float beta=0.0) -> ()");
// Activation function used in GeGLU with `none` approximation.
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -679,6 +693,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
#endif
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
// Sampler kernels (shared CUDA/ROCm)
ops.impl("apply_repetition_penalties_",
csrc/ops.h
+2 -1
View File
@@ -50,7 +50,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
torch::Tensor& scale);
docs/design/attention_backends.md
+14 -2
View File
@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -188,6 +188,18 @@ Priority is **1 = highest** (tried first).
>
> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
## MiniMax M3 Sparse Attention Backends
Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
layers. It is wired in directly by the model and is not part of the
automatic priority lists above. A lightning indexer scores KV blocks, the
top-k blocks (plus fixed init/local blocks) are selected, and attention
attends only to those blocks; index keys live in a separate side cache.
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
## MLA (Multi-head Latent Attention) Backends
MLA uses separate backends for prefill and decode phases.
pyproject.toml
+2
View File
@@ -162,6 +162,8 @@ dout = "dout"
Pn = "Pn"
arange = "arange"
thw = "thw"
# temporal position ids (parallels hpos/wpos in vision RoPE)
tpos = "tpos"
subtile = "subtile"
HSA = "HSA"
setp = "setp"
requirements/common.txt
+1
View File
@@ -29,6 +29,7 @@ xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine ==
typing_extensions >= 4.10
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
partial-json-parser # used for parsing partial JSON outputs
jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
pyzmq >= 25.0.0
msgspec
mistral_common[image] >= 1.11.3
requirements/test/cuda.txt
+1
View File
@@ -360,6 +360,7 @@ jsonpointer==3.0.0
# via jsonschema
jsonschema==4.23.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# ray
requirements/test/rocm.txt
+2
View File
@@ -439,6 +439,8 @@ jsonpointer==3.1.0
# via jsonschema
jsonschema==4.26.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
# hypothesis-jsonschema
# mcp
# mistral-common
requirements/test/xpu.txt
+1
View File
@@ -229,6 +229,7 @@ jsonlines==4.0.0
# via lm-eval
jsonschema==4.26.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# schemathesis
rust/src/chat/src/lib.rs
+2 -2
View File
@@ -277,7 +277,7 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
}
#[test]
@@ -288,6 +288,6 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
}
}
rust/src/chat/src/parser/reasoning/mod.rs
+7 -3
View File
@@ -5,9 +5,9 @@ use std::sync::LazyLock;
pub use vllm_reasoning_parser::{
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
Step3p5ReasoningParser,
KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser, Step3p5ReasoningParser,
};
use vllm_tokenizer::DynTokenizer;
@@ -24,6 +24,7 @@ pub mod names {
pub const KIMI: &str = "kimi";
pub const KIMI_K2: &str = "kimi_k2";
pub const MINIMAX_M2: &str = "minimax_m2";
pub const MINIMAX_M3: &str = "minimax_m3";
pub const NEMOTRON_V3: &str = "nemotron_v3";
pub const QWEN3: &str = "qwen3";
pub const SEED_OSS: &str = "seed_oss";
@@ -62,6 +63,7 @@ impl ReasoningParserFactory {
.register_parser::<KimiReasoningParser>(names::KIMI)
.register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
.register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
.register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
.register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
.register_parser::<Qwen3ReasoningParser>(names::QWEN3)
.register_parser::<SeedOssReasoningParser>(names::SEED_OSS)
@@ -90,6 +92,8 @@ impl ReasoningParserFactory {
.register_pattern("step3", names::STEP3)
.register_pattern("seed-oss", names::SEED_OSS)
.register_pattern("seedoss", names::SEED_OSS)
.register_pattern("minimax-m3", names::MINIMAX_M3)
.register_pattern("mm-m3", names::MINIMAX_M3)
.register_pattern("minimax", names::MINIMAX_M2)
.register_pattern("mm-m2", names::MINIMAX_M2)
.register_pattern("cohere", names::COHERE_CMD)
rust/src/chat/src/parser/reasoning/tests.rs
+15
View File
@@ -34,10 +34,12 @@ fn factory_contains_and_lists_registered_parsers() {
assert!(factory.contains(names::DEEPSEEK_V4));
assert!(factory.contains(names::SEED_OSS));
assert!(factory.contains(names::STEP3P5));
assert!(factory.contains(names::MINIMAX_M3));
assert!(factory.list().contains(&names::QWEN3.to_string()));
assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
assert!(factory.list().contains(&names::SEED_OSS.to_string()));
assert!(factory.list().contains(&names::STEP3P5.to_string()));
assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
}
#[test]
@@ -88,6 +90,19 @@ fn factory_routes_seed_oss_models() {
);
}
#[test]
fn factory_resolves_minimax_m3_before_generic_minimax() {
let factory = ReasoningParserFactory::new();
assert_eq!(
factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"),
Some(names::MINIMAX_M3)
);
assert_eq!(
factory.resolve_name_for_model("mm-m3"),
Some(names::MINIMAX_M3)
);
}
#[test]
fn factory_rejects_unknown_parser_names() {
let tokenizer = Arc::new(FakeTokenizer);
rust/src/chat/src/parser/tool/mod.rs
+7 -2
View File
@@ -6,8 +6,9 @@ pub use vllm_tool_parser::{
DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser,
HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser,
MinimaxM2ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser,
Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput,
MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser,
Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
ToolParserOutput,
};
use crate::parser::ParserFactory;
@@ -32,6 +33,7 @@ pub mod names {
pub const LLAMA3_JSON: &str = "llama3_json";
pub const LLAMA4_JSON: &str = "llama4_json";
pub const MINIMAX_M2: &str = "minimax_m2";
pub const MINIMAX_M3: &str = "minimax_m3";
pub const MISTRAL: &str = "mistral";
pub const PHI4_MINI_JSON: &str = "phi4_mini_json";
pub const QWEN3_CODER: &str = "qwen3_coder";
@@ -73,6 +75,7 @@ impl ToolParserFactory {
.register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
.register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
.register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
.register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
.register_parser::<MistralToolParser>(names::MISTRAL)
.register_parser::<Phi4MiniJsonToolParser>(names::PHI4_MINI_JSON)
.register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
@@ -111,6 +114,8 @@ impl ToolParserFactory {
.register_pattern("gemma-4", names::GEMMA4)
.register_pattern("granite-4", names::GRANITE4)
.register_pattern("kimi-k2", names::KIMI_K2)
.register_pattern("minimax-m3", names::MINIMAX_M3)
.register_pattern("mm-m3", names::MINIMAX_M3)
.register_pattern("minimax", names::MINIMAX_M2)
.register_pattern("mm-m2", names::MINIMAX_M2);
rust/src/chat/src/parser/tool/tests.rs
+8
View File
@@ -157,6 +157,14 @@ fn factory_new_resolves_default_patterns() {
factory.resolve_name_for_model("tencent/Hy3-preview"),
Some(names::HY_V3)
);
assert_eq!(
factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
Some(names::MINIMAX_M3)
);
assert_eq!(
factory.resolve_name_for_model("org/mm-m3-base"),
Some(names::MINIMAX_M3)
);
assert_eq!(
factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
Some(names::MINIMAX_M2)
rust/src/reasoning-parser/src/lib.rs
+2
View File
@@ -19,6 +19,7 @@ mod deepseek_r1;
mod delimited;
mod gemma4;
mod kimi;
mod minimax_m3;
mod qwen3;
mod seed_oss;
mod step3p5;
@@ -31,6 +32,7 @@ pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
pub(crate) use self::delimited::DelimitedReasoningParser;
pub use self::gemma4::Gemma4ReasoningParser;
pub use self::kimi::KimiReasoningParser;
pub use self::minimax_m3::MiniMaxM3ReasoningParser;
pub use self::qwen3::Qwen3ReasoningParser;
pub use self::seed_oss::SeedOssReasoningParser;
pub use self::step3p5::Step3p5ReasoningParser;
rust/src/reasoning-parser/src/minimax_m3.rs
+98
View File
@@ -0,0 +1,98 @@
use vllm_tokenizer::DynTokenizer;
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
const M3_THINK_START: &str = "<mm:think>";
const M3_THINK_END: &str = "</mm:think>";
/// Reasoning parser for MiniMax M3 style outputs.
///
/// MiniMax M3 uses `<mm:think>...</mm:think>` delimiters. Its chat template may
/// prefill either delimiter depending on the requested thinking mode, so the
/// shared delimited parser derives the starting state from the rendered prompt.
pub struct MiniMaxM3ReasoningParser {
inner: DelimitedReasoningParser,
/// True until the first response text is classified. Only this position may
/// drop a stray `</mm:think>` emitted at the start of a response.
at_response_start: bool,
/// Holds an initial suffix like `</mm` while it may still complete into the
/// leading closer on a later chunk.
leading_end_buffer: String,
}
impl MiniMaxM3ReasoningParser {
/// Create a MiniMax M3 parser backed by the shared delimited state machine.
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
Ok(Self {
inner: DelimitedReasoningParser::new(tokenizer, M3_THINK_START, M3_THINK_END, false)?,
at_response_start: true,
leading_end_buffer: String::new(),
})
}
/// Drop a response-leading `</mm:think>` while preserving later unmatched
/// closers as ordinary content.
fn push_inner(&mut self, delta: &str) -> ReasoningDelta {
if self.at_response_start && !self.inner.in_reasoning() {
self.leading_end_buffer.push_str(delta);
let buffered = std::mem::take(&mut self.leading_end_buffer);
if buffered.is_empty() {
return ReasoningDelta::default();
}
if let Some(rest) = buffered.strip_prefix(M3_THINK_END) {
self.at_response_start = false;
return self.inner.push(rest);
}
if M3_THINK_END.starts_with(buffered.as_str()) {
self.leading_end_buffer = buffered;
return ReasoningDelta::default();
}
self.at_response_start = false;
return self.inner.push(&buffered);
}
self.inner.push(delta)
}
}
fn append_delta(target: &mut ReasoningDelta, delta: ReasoningDelta) {
if let Some(reasoning) = delta.reasoning {
target.push_reasoning(&reasoning);
}
if let Some(content) = delta.content {
target.push_content(&content);
}
}
impl ReasoningParser for MiniMaxM3ReasoningParser {
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tokenizer)?))
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.inner.initialize(prompt_token_ids);
self.at_response_start = true;
self.leading_end_buffer.clear();
Ok(())
}
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
Ok(self.push_inner(delta))
}
fn finish(&mut self) -> Result<ReasoningDelta> {
let mut delta = ReasoningDelta::default();
if !self.leading_end_buffer.is_empty() {
let pending = std::mem::take(&mut self.leading_end_buffer);
self.at_response_start = false;
append_delta(&mut delta, self.inner.push(&pending));
}
append_delta(&mut delta, self.inner.finish());
Ok(delta)
}
}
rust/src/reasoning-parser/src/tests.rs
+67 -1
View File
@@ -3,7 +3,8 @@ use std::sync::Arc;
use vllm_tokenizer::Tokenizer;
use super::{
DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
Qwen3ReasoningParser, ReasoningParser,
};
pub(crate) struct FakeTokenizer;
@@ -32,6 +33,8 @@ impl Tokenizer for FakeTokenizer {
"<|END_THINKING|>" => Some(4),
"◁think▷" => Some(5),
"◁/think▷" => Some(6),
"<mm:think>" => Some(8),
"</mm:think>" => Some(9),
"<seed:think>" => Some(10),
"</seed:think>" => Some(11),
_ => None,
@@ -161,3 +164,66 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_handles_explicit_think_delimiters() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_drops_leading_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("</mm:think>answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_preserves_non_leading_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("XXX</mm:think>YYY").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("XXX</mm:think>YYY"));
}
#[test]
fn minimax_m3_drops_split_leading_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
assert!(parser.push("</mm").unwrap().is_empty());
let delta = parser.push(":think>answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_uses_prompt_prefilled_start_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[8]).unwrap();
let delta = parser.push("reason</mm:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_uses_prompt_prefilled_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[9]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
rust/src/tool-parser/python/src/lib.rs
+2
View File
@@ -38,6 +38,8 @@ macro_rules! tool_parser_factory {
// Export a tool parser to Python by registering it here.
tool_parser_factory! {
MinimaxM3ToolParser,
// Below are the parsers just for testing purposes on Python side.
DeepSeekV4ToolParser,
KimiK2ToolParser,
rust/src/tool-parser/src/lib.rs
+2
View File
@@ -10,6 +10,7 @@ mod hy_v3;
mod json;
mod kimi_k2;
mod minimax_m2;
mod minimax_m3;
mod parameters;
mod qwen_coder;
#[cfg(any(test, feature = "test-util"))]
@@ -30,6 +31,7 @@ pub use json::{
};
pub use kimi_k2::KimiK2ToolParser;
pub use minimax_m2::MinimaxM2ToolParser;
pub use minimax_m3::MinimaxM3ToolParser;
pub use qwen_coder::Qwen3CoderToolParser;
use serde::{Deserialize, Serialize};
use serde_json::Value;
rust/src/tool-parser/src/minimax_m3.rs
+885
View File
@@ -0,0 +1,885 @@
use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};
use winnow::combinator::{alt, delimited, seq};
use winnow::error::{ContextError, ErrMode};
use winnow::prelude::*;
use winnow::stream::Partial;
use winnow::token::{literal, rest, take_until};
use super::parameters::{ParamElement, ParamInput, ToolSchemas};
use super::utils::{parse_buffered_event, safe_text_len};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::Tool;
const NAMESPACE: &str = "]<]minimax[>[";
const TOOL_CALL_START: &str = "]<]minimax[>[<tool_call>";
const TOOL_CALL_END: &str = "]<]minimax[>[</tool_call>";
const INVOKE_START: &str = "]<]minimax[>[<invoke";
const INVOKE_END: &str = "]<]minimax[>[</invoke>";
const ELEMENT_START: &str = "]<]minimax[>[<";
const ELEMENT_END_START: &str = "]<]minimax[>[</";
const MIXED_TEXT_FIELD: &str = "$text";
type MinimaxM3Input<'i> = Partial<&'i str>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MinimaxM3Mode {
Text,
ToolBlock,
Done,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum MinimaxM3Event {
Text {
len: usize,
},
ToolBlockStart,
Invoke {
name: String,
params: Vec<(String, ParamInput)>,
},
ToolBlockEnd,
IgnoredRest,
}
/// Tool parser for MiniMax M3 namespace-delimited XML-style tool calls.
///
/// Example tool call content with recursive parameters:
///
/// ```text
/// ]<]minimax[>[<tool_call>
/// ]<]minimax[>[<invoke name="create_order">
/// ]<]minimax[>[<user_id>42]<]minimax[>[</user_id>
/// ]<]minimax[>[<shipping>
/// ]<]minimax[>[<city>Singapore]<]minimax[>[</city>
/// ]<]minimax[>[<zip>018956]<]minimax[>[</zip>
/// ]<]minimax[>[</shipping>
/// ]<]minimax[>[<items>
/// ]<]minimax[>[<item>
/// ]<]minimax[>[<sku>book-001]<]minimax[>[</sku>
/// ]<]minimax[>[<qty>2]<]minimax[>[</qty>
/// ]<]minimax[>[</item>
/// ]<]minimax[>[</items>
/// ]<]minimax[>[</invoke>
/// ]<]minimax[>[</tool_call>
/// ```
///
/// With a schema where `shipping` is an object and `items` is an array of
/// objects, recursive parameter conversion produces:
///
/// ```json
/// {
/// "user_id": 42,
/// "shipping": {
/// "city": "Singapore",
/// "zip": 18956
/// },
/// "items": [
/// {
/// "sku": "book-001",
/// "qty": 2
/// }
/// ]
/// }
/// ```
///
/// MiniMax M3 emits the namespace marker `]<]minimax[>[` before each structural
/// tag. Arguments are emitted only after a full `<invoke>` block is parsed.
pub struct MinimaxM3ToolParser {
buffer: String,
mode: MinimaxM3Mode,
emitted_tool_count: usize,
tool_parameters: ToolSchemas,
}
impl MinimaxM3ToolParser {
/// Create a MiniMax M3 tool parser.
pub fn new(tools: &[Tool]) -> Self {
Self {
buffer: String::new(),
mode: MinimaxM3Mode::Text,
emitted_tool_count: 0,
tool_parameters: ToolSchemas::from_tools(tools),
}
}
/// Apply one parsed MiniMax M3 event to parser state and output.
fn apply_event(&mut self, event: MinimaxM3Event, output: &mut ToolParserOutput) -> Result<()> {
match event {
MinimaxM3Event::Text { len: consumed_len } => {
output.normal_text.push_str(&self.buffer[..consumed_len]);
}
MinimaxM3Event::ToolBlockStart => self.mode = MinimaxM3Mode::ToolBlock,
MinimaxM3Event::Invoke { name, params } => {
let arguments = self.tool_parameters.convert_params_with_schema(&name, params);
let arguments = serde_json::to_string(&arguments)
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
output.calls.push(ToolCallDelta {
tool_index: self.emitted_tool_count,
name: Some(name),
arguments,
});
self.emitted_tool_count += 1;
}
MinimaxM3Event::ToolBlockEnd => self.mode = MinimaxM3Mode::Done,
MinimaxM3Event::IgnoredRest => {}
}
Ok(())
}
}
impl ToolParser for MinimaxM3ToolParser {
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tools)))
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
self.buffer.push_str(chunk);
while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
parse_next_minimax_m3_event(input, self.mode)
})? {
self.apply_event(event, output)?;
self.buffer.drain(..consumed_len);
}
Ok(())
}
fn finish(&mut self) -> Result<ToolParserOutput> {
let mut output = ToolParserOutput::default();
match self.mode {
MinimaxM3Mode::Text => {
output.normal_text.push_str(&self.buffer);
}
MinimaxM3Mode::ToolBlock => {
if !self.buffer.trim_start().is_empty() {
return Err(parsing_failed!("incomplete MiniMax M3 tool call"));
}
}
MinimaxM3Mode::Done => {}
}
let _ = self.reset();
Ok(output)
}
fn reset(&mut self) -> String {
self.mode = MinimaxM3Mode::Text;
self.emitted_tool_count = 0;
std::mem::take(&mut self.buffer)
}
}
/// Parse a MiniMax M3 event for the current parser mode.
fn parse_next_minimax_m3_event(
input: &mut MinimaxM3Input<'_>,
mode: MinimaxM3Mode,
) -> ModalResult<MinimaxM3Event> {
match mode {
MinimaxM3Mode::Text => parse_text_event(input),
MinimaxM3Mode::ToolBlock => parse_tool_block_event(input),
MinimaxM3Mode::Done => ignored_rest_event(input),
}
}
/// Parse a text-mode MiniMax M3 event.
fn parse_text_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
alt((tool_block_start_event, safe_text_event)).parse_next(input)
}
/// Parse a MiniMax M3 tool-block start marker.
fn tool_block_start_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
literal(TOOL_CALL_START).value(MinimaxM3Event::ToolBlockStart).parse_next(input)
}
/// Parse a safe text run before the next MiniMax M3 marker.
fn safe_text_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
safe_text_len(input, TOOL_CALL_START).map(|len| MinimaxM3Event::Text { len })
}
/// Parse one event inside a MiniMax M3 tool block.
fn parse_tool_block_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
alt((tool_block_end_event, invoke_event)).parse_next(input)
}
/// Parse a MiniMax M3 tool-block end marker.
fn tool_block_end_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
(ws0, literal(TOOL_CALL_END))
.value(MinimaxM3Event::ToolBlockEnd)
.parse_next(input)
}
/// Parse a complete MiniMax M3 invoke block.
fn invoke_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
let (name, body) = seq!(
_: ws0,
_: literal(INVOKE_START),
_: (ws1, literal("name=")),
partial_attr_value,
_: literal(">"),
take_until(0.., INVOKE_END),
_: literal(INVOKE_END),
)
.parse_next(input)?;
let params = parse_invoke_params(body)?;
Ok(MinimaxM3Event::Invoke {
name: name.trim().to_string(),
params,
})
}
/// Parse all parameter elements inside a complete MiniMax M3 invoke body.
fn parse_invoke_params(invoke_body: &str) -> ModalResult<Vec<(String, ParamInput)>> {
let mut input = invoke_body;
let mut elements = Vec::new();
loop {
let _ = ws0.parse_next(&mut input)?;
if input.is_empty() {
break;
}
if input.starts_with(ELEMENT_START) {
elements.push(parameter_element(&mut input)?);
continue;
}
if input.starts_with(NAMESPACE) {
return malformed();
}
// Be tolerant: ordinary text at an invokeparameter boundary ends this invoke.
// Keep parsed parameters and drop the remaining invoke body.
break;
}
Ok(elements.into_iter().map(|element| (element.name, element.value)).collect())
}
/// Parse a MiniMax M3 parameter element.
fn parameter_element(input: &mut &str) -> ModalResult<ParamElement> {
let name = open_element_tag(input)?.to_string();
let value = element_body(input, &name)?;
close_element_tag(input, &name)?;
Ok(ParamElement { name, value })
}
/// Parse a MiniMax M3 opening element tag.
fn open_element_tag<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
let name = seq!(
_: literal(ELEMENT_START),
take_until(1.., ">"),
_: literal(">"),
)
.parse_next(input)?;
let name = name.0;
if name.starts_with('/') || name.trim().is_empty() {
return malformed();
}
Ok(name)
}
/// Parse a MiniMax M3 closing element tag.
fn close_element_tag(input: &mut &str, name: &str) -> ModalResult<()> {
literal(ELEMENT_END_START).void().parse_next(input)?;
literal(name).void().parse_next(input)?;
literal(">").void().parse_next(input)
}
/// Parse the body of one MiniMax M3 element.
fn element_body(input: &mut &str, closing_name: &str) -> ModalResult<ParamInput> {
let close_tag = format!("{ELEMENT_END_START}{closing_name}>");
let mut text = String::new();
let mut elements = Vec::new();
loop {
text.push_str(text_until_namespace(input)?);
if input.starts_with(&close_tag) {
// Close tag reached, end of element body.
break;
}
if input.starts_with(ELEMENT_START) {
// Child element start reached, parse child element recursively.
elements.push(parameter_element(input)?);
continue;
}
if input.starts_with(NAMESPACE) {
// Unexpected namespace marker.
return malformed();
}
}
if elements.is_empty() {
Ok(ParamInput::Text(text))
} else {
if !text.trim().is_empty() {
push_mixed_text_element(&mut elements, text);
}
Ok(ParamInput::Elements(elements))
}
}
/// Parse text until the next MiniMax M3 namespace marker.
fn text_until_namespace<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
take_until(0.., NAMESPACE).parse_next(input)
}
/// Preserve mixed text content under a reserved object field.
///
/// By default, the field name is `$text`, but if that collides with an existing
/// child element name, prepend `$` until there is no collision.
fn push_mixed_text_element(elements: &mut Vec<ParamElement>, text: String) {
let mut name = MIXED_TEXT_FIELD.to_string();
while elements.iter().any(|element| element.name == name) {
name.insert(0, '$');
}
elements.push(ParamElement {
name,
value: ParamInput::Text(text),
});
}
/// Parse a quoted or unquoted XML attribute value from partial streaming input.
fn partial_attr_value<'i>(input: &mut MinimaxM3Input<'i>) -> ModalResult<&'i str> {
alt((
delimited(literal("\""), take_until(1.., "\""), literal("\"")),
delimited(literal("'"), take_until(1.., "'"), literal("'")),
take_until(1.., ">"),
))
.parse_next(input)
}
/// Parse ignored rest after the MiniMax M3 tool block ends.
fn ignored_rest_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
rest.value(MinimaxM3Event::IgnoredRest).parse_next(input)
}
fn malformed<T>() -> ModalResult<T> {
Err(ErrMode::Cut(ContextError::new()))
}
#[cfg(test)]
mod tests {
use expect_test::expect;
use serde_json::{Value, json};
use thiserror_ext::AsReport;
use super::{
ELEMENT_END_START, ELEMENT_START, INVOKE_END, INVOKE_START, MinimaxM3ToolParser,
TOOL_CALL_END, TOOL_CALL_START, ToolParser,
};
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
use crate::{Tool, ToolParserTestExt as _};
fn element(name: &str, body: &str) -> String {
format!("{ELEMENT_START}{name}>{body}{ELEMENT_END_START}{name}>")
}
fn invoke(function_name: &str, body: &str) -> String {
format!("{INVOKE_START} name=\"{function_name}\">{body}{INVOKE_END}")
}
fn build_tool_block(invokes: &[(&str, String)]) -> String {
let invokes = invokes
.iter()
.map(|(function_name, body)| invoke(function_name, body))
.collect::<Vec<_>>()
.join("\n");
format!("{TOOL_CALL_START}\n{invokes}\n{TOOL_CALL_END}")
}
fn m3_test_tools() -> Vec<Tool> {
let mut tools = test_tools();
tools.push(Tool {
name: "create_order".to_string(),
description: None,
parameters: json!({
"type": "object",
"properties": {
"user_id": { "type": "integer" },
"urgent": { "type": "boolean" },
"note": { "type": "string" },
"shipping": {
"type": "object",
"properties": {
"city": { "type": "string" },
"zip": { "type": "integer" }
}
},
"items": {
"type": "array",
"items": {
"type": "object",
"properties": {
"sku": { "type": "string" },
"qty": { "type": "integer" }
}
}
},
"metadata": {
"type": "object",
"additionalProperties": { "type": "integer" }
},
"duplicate_demo": {
"type": "object",
"properties": {
"tag": { "type": "string" }
}
},
"schema_mismatch_array": {
"type": "array",
"items": { "type": "integer" }
}
}
}),
strict: None,
});
tools
}
fn order_arguments() -> String {
let shipping = element(
"shipping",
&format!(
"{}{}",
element("city", "Singapore"),
element("zip", "018956")
),
);
let first_item = element(
"item",
&format!("{}{}", element("sku", "book-001"), element("qty", "2")),
);
let second_item = element(
"item",
&format!("{}{}", element("sku", "pen-007"), element("qty", "5")),
);
let items = element("items", &format!("{first_item}{second_item}"));
let metadata = element(
"metadata",
&format!("{}{}", element("score", "42"), element("rank", "7")),
);
let duplicate_demo = element(
"duplicate_demo",
&format!("{}{}", element("tag", "a"), element("tag", "b")),
);
let schema_mismatch_array = element(
"schema_mismatch_array",
&format!("{}{}", element("x", "1"), element("x", "2")),
);
[
element("user_id", "42"),
element("urgent", "true"),
element("note", "Please leave at front desk."),
shipping,
items,
metadata,
duplicate_demo,
schema_mismatch_array,
element(
"unknown_struct",
&format!("{}{}", element("a", "1"), element("a", "2")),
),
]
.join("")
}
#[test]
fn minimax_m3_parse_complete_without_tool_call_keeps_text() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser.parse_complete("Hello, world!").unwrap();
assert_eq!(output.normal_text, "Hello, world!");
assert!(output.calls.is_empty());
}
#[test]
fn minimax_m3_parse_complete_extracts_single_tool_call() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_complete(&build_tool_block(&[(
"get_weather",
format!("{}{}", element("city", "Seattle"), element("days", "5")),
)]))
.unwrap();
assert!(output.normal_text.is_empty());
assert_eq!(output.calls.len(), 1);
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({ "city": "Seattle", "days": 5 })
);
}
#[test]
fn minimax_m3_parse_complete_preserves_prefix_and_ignores_trailing_text() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = format!(
"Let me check. {} This trailing text is ignored.",
build_tool_block(&[("get_weather", element("city", "Seattle"))])
);
let output = parser.parse_complete(&output).unwrap();
assert_eq!(output.normal_text, "Let me check. ");
assert_eq!(output.calls.len(), 1);
}
#[test]
fn minimax_m3_parse_complete_extracts_multiple_invokes() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_complete(&build_tool_block(&[
("get_weather", element("city", "Seattle")),
("get_weather", element("city", "NYC")),
]))
.unwrap();
assert_eq!(output.calls.len(), 2);
assert_eq!(output.calls[0].tool_index, 0);
assert_eq!(output.calls[1].tool_index, 1);
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({ "city": "Seattle" })
);
assert_eq!(
serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
json!({ "city": "NYC" })
);
}
#[test]
fn minimax_m3_invoke_body_junk_drops_rest_of_invoke() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_complete(&build_tool_block(&[(
"get_weather",
[
element("city", "Seattle"),
"I need to use the city above.".to_string(),
element("days", "5"),
]
.join(""),
)]))
.unwrap();
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({ "city": "Seattle" })
);
}
#[test]
fn minimax_m3_parse_complete_converts_schema_types() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_complete(&build_tool_block(&[(
"convert",
[
element("whole", "5.0"),
element("flag", "true"),
element("payload", r#"{"nested":true}"#),
element("items", "[1,2]"),
element("empty", "42"),
]
.join(""),
)]))
.unwrap();
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({
"whole": 5.0,
"flag": true,
"payload": { "nested": true },
"items": [1, 2],
"empty": "42",
})
);
}
#[test]
fn minimax_m3_parse_complete_converts_nested_arguments() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_complete(&build_tool_block(&[("create_order", order_arguments())]))
.unwrap();
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({
"user_id": 42,
"urgent": true,
"note": "Please leave at front desk.",
"shipping": {
"city": "Singapore",
"zip": 18956
},
"items": [
{
"sku": "book-001",
"qty": 2
},
{
"sku": "pen-007",
"qty": 5
}
],
"metadata": {
"score": 42,
"rank": 7
},
"duplicate_demo": {
"tag": ["a", "b"]
},
"schema_mismatch_array": [1, 2],
"unknown_struct": {
"a": ["1", "2"]
}
})
);
}
#[test]
fn minimax_m3_parse_complete_handles_multiline_leaf_parameters() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_complete(&build_tool_block(&[(
"calculate_area",
[
element("shape", "\nrectangle\n"),
element("dimensions", r#"{"width":10,"height":20}"#),
element("precision", "2"),
]
.join(""),
)]))
.unwrap();
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({
"shape": "\nrectangle\n",
"dimensions": { "width": 10, "height": 20 },
"precision": 2,
})
);
}
#[test]
fn minimax_m3_streaming_extracts_single_tool_call() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = collect_stream(
&mut parser,
&[
TOOL_CALL_START,
&invoke("get_weather", &element("city", "Seattle")),
TOOL_CALL_END,
],
);
assert!(output.normal_text.is_empty());
assert_eq!(output.calls.len(), 1);
assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({ "city": "Seattle" })
);
}
#[test]
fn minimax_m3_streaming_preserves_prefix_text() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = collect_stream(
&mut parser,
&[
"Let me check. ",
TOOL_CALL_START,
&invoke("get_weather", &element("city", "Seattle")),
TOOL_CALL_END,
],
);
assert_eq!(output.normal_text, "Let me check. ");
assert_eq!(output.calls.len(), 1);
}
#[test]
fn minimax_m3_streaming_without_tool_call_emits_text_incrementally() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
assert_eq!(output.normal_text, "Hello, world!");
assert!(output.calls.is_empty());
}
#[test]
fn minimax_m3_streaming_handles_marker_split_across_chunks() {
let text = build_tool_block(&[("get_weather", element("city", "Seattle"))]);
let chunks = split_by_chars(&text, 3);
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = collect_stream(&mut parser, &chunks);
assert_eq!(output.calls.len(), 1);
assert!(output.normal_text.is_empty());
}
#[test]
fn minimax_m3_streaming_extracts_multiple_invokes_in_order() {
let text = build_tool_block(&[
("get_weather", element("city", "Seattle")),
("get_weather", element("city", "NYC")),
]);
let chunks = split_by_chars(&text, 7);
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = collect_stream(&mut parser, &chunks);
assert_eq!(output.calls.len(), 2);
assert_eq!(output.calls[0].tool_index, 0);
assert_eq!(output.calls[1].tool_index, 1);
}
#[test]
fn minimax_m3_streaming_does_not_emit_incomplete_tool_call() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_chunk(&format!(
"{TOOL_CALL_START}{INVOKE_START} name=\"get_weather\">"
))
.unwrap();
assert!(output.normal_text.is_empty());
assert!(output.calls.is_empty());
}
#[test]
fn minimax_m3_streaming_ignores_text_after_tool_block() {
let text = format!(
"{} ignored",
build_tool_block(&[("get_weather", element("city", "Seattle"))])
);
let chunks = split_by_chars(&text, 5);
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = collect_stream(&mut parser, &chunks);
assert!(output.normal_text.is_empty());
assert_eq!(output.calls.len(), 1);
}
#[test]
fn minimax_m3_finish_fails_incomplete_tool_call() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
parser
.parse_chunk(&format!(
"{TOOL_CALL_START}{INVOKE_START} name=\"get_weather\">"
))
.unwrap();
assert!(parser.finish().is_err());
}
#[test]
fn minimax_m3_finish_recovers_after_bare_tool_block_start() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
parser.parse_chunk(TOOL_CALL_START).unwrap();
let output = parser.finish().unwrap();
assert!(output.normal_text.is_empty());
assert!(output.calls.is_empty());
}
#[test]
fn minimax_m3_finish_recovers_completed_invoke_with_whitespace_tail() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let output = parser
.parse_complete(&format!(
"{}\n{}\n \n",
TOOL_CALL_START,
invoke("get_weather", &element("city", "Seattle"))
))
.unwrap();
assert_eq!(output.calls.len(), 1);
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({ "city": "Seattle" })
);
}
#[test]
fn minimax_m3_finish_fails_partial_outer_end_marker() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
parser
.parse_chunk(&format!(
"{}\n{}\n{}",
TOOL_CALL_START,
invoke("get_weather", &element("city", "Seattle")),
&TOOL_CALL_END[..3]
))
.unwrap();
assert!(parser.finish().is_err());
}
#[test]
fn minimax_m3_malformed_tool_call_fails_fast() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let error = parser
.parse_chunk(&format!(
"{TOOL_CALL_START}{ELEMENT_START}bad>{TOOL_CALL_END}"
))
.unwrap_err();
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
}
#[test]
fn minimax_m3_mixed_content_is_preserved_as_text_field() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let body = element(
"payload",
&format!("text before {} text after", element("child", "value")),
);
let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({
"payload": {
"child": "value",
"$text": "text before text after"
}
})
);
}
#[test]
fn minimax_m3_mixed_text_field_avoids_child_name_collision() {
let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
let body = element(
"payload",
&format!(
"text{}{}",
element("$text", "child text"),
element("child", "value")
),
);
let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
assert_eq!(
serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
json!({
"payload": {
"$text": "child text",
"$$text": "text",
"child": "value"
}
})
);
}
}
setup.py
+19
View File
@@ -432,6 +432,19 @@ class cmake_build_ext(build_ext):
dirs_exist_ok=True,
)
# copy vendored fmha_sm100 package from build_lib to source tree
# for editable installs
fmha_sm100_build = os.path.join(
self.build_lib, "vllm", "third_party", "fmha_sm100"
)
if os.path.exists(fmha_sm100_build):
print(f"Copying {fmha_sm100_build} to vllm/third_party/fmha_sm100")
shutil.copytree(
fmha_sm100_build,
"vllm/third_party/fmha_sm100",
dirs_exist_ok=True,
)
class precompiled_build_ext(build_ext):
"""Disables extension building when using precompiled binaries."""
@@ -787,6 +800,7 @@ class precompiled_wheel_utils:
)
# DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.)
deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*")
fmha_sm100_regex = re.compile(r"vllm/third_party/fmha_sm100/.*")
file_members = []
for member in wheel.filelist:
if member.filename in exact_members:
@@ -812,6 +826,7 @@ class precompiled_wheel_utils:
or triton_kernels_regex.match(member.filename)
or flashmla_regex.match(member.filename)
or deep_gemm_regex.match(member.filename)
or fmha_sm100_regex.match(member.filename)
):
file_members.append(member)
@@ -1120,6 +1135,8 @@ if _is_cuda():
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
# Optional since it won't build on unsupported architectures
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
# fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
if _is_cpu():
import platform
@@ -1150,6 +1167,8 @@ package_data = {
"third_party/deep_gemm/include/**/*.cuh",
"third_party/deep_gemm/include/**/*.h",
"third_party/deep_gemm/include/**/*.hpp",
# fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake)
"third_party/fmha_sm100/cute/src/sm100/build_k2q_csr/build_k2q_csr.cu",
]
}
tests/kernels/attention/test_minimax_m3.py
+854
View File
@@ -0,0 +1,854 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Correctness tests for MiniMax M3 sparse prefill attention kernels."""
import pytest
import torch
from vllm import _custom_ops as ops
from vllm.models.minimax_m3.common.indexer import (
MiniMaxM3IndexerBackend,
)
from vllm.models.minimax_m3.common.ops.index_topk import (
minimax_m3_index_decode,
minimax_m3_index_score,
minimax_m3_index_topk,
)
from vllm.models.minimax_m3.common.ops.sparse_attn import (
minimax_m3_sparse_attn,
minimax_m3_sparse_attn_decode,
)
from vllm.models.minimax_m3.common.sparse_attention import (
MiniMaxM3SparseBackend,
)
from vllm.platforms import current_platform
from vllm.v1.attention.backends.utils import set_kv_cache_layout
from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
from vllm.v1.worker.utils import AttentionGroup
if not (current_platform.is_cuda() or current_platform.is_rocm()):
pytest.skip(
"MiniMax M3 attention kernels require CUDA or ROCm.",
allow_module_level=True,
)
@pytest.fixture
def kv_layout(request):
"""Set the global KV cache layout for one test and restore it after."""
set_kv_cache_layout(request.param)
try:
yield request.param
finally:
set_kv_cache_layout(None)
def _stride_order_for(backend: type[MiniMaxM3SparseBackend], ndim: int) -> tuple:
"""Mirror the allocator's stride-order resolution (identity fallback)."""
try:
stride_order = backend.get_kv_cache_stride_order()
assert len(stride_order) == ndim
except (AttributeError, NotImplementedError):
stride_order = tuple(range(ndim))
return stride_order
def _allocate_main_kv_via_contract(
num_pages: int, device: torch.device | str = "cuda"
) -> torch.Tensor:
"""Build the main KV cache exactly as the production allocator does for the
currently active layout: allocate the physical (permuted) tensor, then
expose the inverse-permuted logical-NHD view the backend sees."""
logical_shape = MiniMaxM3SparseBackend.get_kv_cache_shape(
num_pages, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
)
stride_order = _stride_order_for(MiniMaxM3SparseBackend, len(logical_shape))
physical_shape = tuple(logical_shape[i] for i in stride_order)
inv_order = [stride_order.index(i) for i in range(len(stride_order))]
raw = torch.randn(physical_shape, device=device, dtype=DTYPE)
return raw.permute(*inv_order)
NUM_Q_HEADS = 32
NUM_KV_HEADS = 2
HEAD_DIM = 128
BLOCK_SIZE = 128
DTYPE = torch.bfloat16
SM_SCALE = HEAD_DIM**-0.5
TOPK = 16
# Index top-k kernels.
def _reference_index_topk(
idx_q: torch.Tensor,
index_kv_cache: torch.Tensor,
block_table: torch.Tensor,
q_lens: torch.Tensor,
seq_lens: torch.Tensor,
prefix_lens: torch.Tensor,
topk: int,
init_blocks: int,
local_blocks: int,
sm_scale: float,
) -> torch.Tensor:
total_q, num_idx_heads, _ = idx_q.shape
out = torch.full(
(num_idx_heads, total_q, topk), -1, device=idx_q.device, dtype=torch.int32
)
q_start = 0
for req_id, (q_len, seq_len, prefix_len) in enumerate(
zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
):
q_end = q_start + q_len
q = idx_q[q_start:q_end]
num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
pages = block_table[req_id, :num_blocks]
k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale
q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
k_pos = torch.arange(k.shape[0], device=idx_q.device)
score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf"))
score = score.reshape(num_idx_heads, q_len, num_blocks, BLOCK_SIZE)
score_tensor = score.max(dim=3).values
valid_blocks = (q_pos + BLOCK_SIZE) // BLOCK_SIZE
for local_q, num_valid_blocks in enumerate(valid_blocks.tolist()):
end = min(init_blocks, num_valid_blocks)
score_tensor[:, local_q, :end] = 1e30
start = max(0, num_valid_blocks - local_blocks)
score_tensor[:, local_q, start:num_valid_blocks] = 1e29
k = min(topk, num_valid_blocks)
topk_idx = score_tensor[:, local_q].topk(k, dim=1).indices
out[:, q_start + local_q, :k] = topk_idx
q_start = q_end
return out
def _assert_topk_indices_equal_unordered(
actual: torch.Tensor,
expected: torch.Tensor,
) -> None:
"""Compare selected sparse blocks without requiring a deterministic order."""
assert actual.shape == expected.shape
actual_flat = actual.cpu().reshape(-1, actual.shape[-1]).tolist()
expected_flat = expected.cpu().reshape(-1, expected.shape[-1]).tolist()
for actual_row, expected_row in zip(actual_flat, expected_flat):
assert set(actual_row) == set(expected_row)
def test_prefill_index_topk_correctness():
topk = 6
init_blocks = 0
local_blocks = 1
num_idx_heads = 2
head_dim = 16
q_lens = torch.tensor((4, 3), device="cuda", dtype=torch.int32)
prefix_lens = torch.tensor((0, 1024), device="cuda", dtype=torch.int32)
seq_lens = prefix_lens + q_lens
batch = q_lens.numel()
max_seq_len = seq_lens.max().item()
max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_pages = batch * max_blocks
cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
cu_seqlens[1:] = q_lens.cumsum(0)
block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
batch, max_blocks
)
idx_q = torch.ones(q_lens.sum().item(), num_idx_heads, head_dim, device="cuda")
index_kv_cache = torch.empty(num_pages, BLOCK_SIZE, head_dim, device="cuda")
for req_id in range(batch):
for block_id in range(max_blocks):
page = block_table[req_id, block_id]
index_kv_cache[page].fill_(block_id + 1)
score = minimax_m3_index_score(
idx_q,
index_kv_cache,
block_table,
cu_seqlens,
seq_lens,
prefix_lens,
max_query_len=q_lens.max().item(),
max_seq_len=max_seq_len,
num_kv_heads=num_idx_heads,
sm_scale=head_dim**-0.5,
)
actual = minimax_m3_index_topk(
score,
cu_seqlens,
prefix_lens,
max_query_len=q_lens.max().item(),
topk=topk,
init_blocks=init_blocks,
local_blocks=local_blocks,
)
expected = _reference_index_topk(
idx_q,
index_kv_cache,
block_table,
q_lens,
seq_lens,
prefix_lens,
topk,
init_blocks,
local_blocks,
head_dim**-0.5,
)
_assert_topk_indices_equal_unordered(actual, expected)
@pytest.mark.parametrize("decode_query_len", [1, 4])
@pytest.mark.parametrize("num_padded_reqs", [0, 2])
def test_decode_index_topk_correctness(
decode_query_len: int,
num_padded_reqs: int,
):
topk = 6
init_blocks = 0
local_blocks = 1
num_idx_heads = 2
head_dim = 16
active_seq_lens = torch.tensor((7, 129, 1025), device="cuda", dtype=torch.int32)
q_lens = torch.full_like(active_seq_lens, decode_query_len)
prefix_lens = active_seq_lens - decode_query_len
active_batch = active_seq_lens.numel()
batch = active_batch + num_padded_reqs
seq_lens = torch.cat(
[
active_seq_lens,
torch.zeros(num_padded_reqs, device="cuda", dtype=torch.int32),
]
)
max_seq_len = active_seq_lens.max().item()
max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_pages = active_batch * max_blocks
active_block_table = torch.randperm(
num_pages, device="cuda", dtype=torch.int32
).reshape(active_batch, max_blocks)
block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
block_table[:active_batch] = active_block_table
idx_q = torch.randn(
batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
)
index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda")
actual = minimax_m3_index_decode(
idx_q,
index_kv_cache,
block_table,
seq_lens,
max_seq_len=max_seq_len,
topk=topk,
init_blocks=init_blocks,
local_blocks=local_blocks,
num_kv_heads=num_idx_heads,
sm_scale=head_dim**-0.5,
decode_query_len=decode_query_len,
)
expected = torch.full_like(actual, -1)
active_tokens = active_batch * decode_query_len
expected[:, :active_tokens] = _reference_index_topk(
idx_q[:active_tokens],
index_kv_cache,
block_table[:active_batch],
q_lens,
active_seq_lens,
prefix_lens,
topk,
init_blocks,
local_blocks,
head_dim**-0.5,
)
_assert_topk_indices_equal_unordered(actual, expected)
# Sparse attention kernels.
def _reference_sparse_attn(
q: torch.Tensor,
kv_cache: torch.Tensor,
topk_idx: torch.Tensor,
block_table: torch.Tensor,
q_lens: torch.Tensor,
seq_lens: torch.Tensor,
prefix_lens: torch.Tensor,
) -> torch.Tensor:
out = torch.empty_like(q, dtype=torch.float32)
gqa_group_size = NUM_Q_HEADS // NUM_KV_HEADS
q_start = 0
for req_id, (q_len, seq_len, prefix_len) in enumerate(
zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
):
q_end = q_start + q_len
q_req = q[q_start:q_end]
positions = torch.arange(seq_len, device="cuda")
pages = block_table[req_id, positions // BLOCK_SIZE]
rows = positions % BLOCK_SIZE
k_req = kv_cache[pages, 0, rows]
v_req = kv_cache[pages, 1, rows].float()
q_pos = prefix_len + torch.arange(q_len, device="cuda")
key_blocks = positions // BLOCK_SIZE
causal_mask = positions.unsqueeze(0) <= q_pos.unsqueeze(1)
for kv_head in range(NUM_KV_HEADS):
selected = topk_idx[kv_head, q_start:q_end]
selected_mask = (key_blocks[None, :, None] == selected[:, None, :]).any(-1)
mask = causal_mask & selected_mask
head_start = kv_head * gqa_group_size
head_end = head_start + gqa_group_size
q_heads = q_req[:, head_start:head_end].transpose(0, 1)
k_head = k_req[:, kv_head].T.expand(gqa_group_size, -1, -1)
scores = torch.bmm(q_heads, k_head, out_dtype=torch.float32)
scores = scores.transpose(0, 1) * SM_SCALE
probs = torch.softmax(
scores.masked_fill(~mask[:, None, :], -float("inf")), -1
)
out[q_start:q_end, head_start:head_end] = torch.einsum(
"qhk,kd->qhd", probs, v_req[:, kv_head]
)
q_start += q_len
return out.to(q.dtype)
@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
@pytest.mark.parametrize(
("q_lens", "kv_lens"),
[
((129, 257), (129, 257)),
((65, 129, 257), (129, 257, 385)),
],
)
def test_prefill_sparse_attention_correctness(
kv_layout: str,
q_lens: tuple[int, ...],
kv_lens: tuple[int, ...],
):
assert len(q_lens) == len(kv_lens)
assert all(kv_len >= q_len for q_len, kv_len in zip(q_lens, kv_lens))
# Build paged-KV metadata, including a non-identity page order.
batch = len(q_lens)
pages_per_req = [(kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE for kv_len in kv_lens]
max_blocks = max(pages_per_req)
num_pages = sum(pages_per_req)
physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
base_page = 0
for req_id, num_req_pages in enumerate(pages_per_req):
block_table[req_id, :num_req_pages] = physical_pages[
base_page : base_page + num_req_pages
]
base_page += num_req_pages
q_lens_t = torch.tensor(q_lens, device="cuda", dtype=torch.int32)
seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32)
prefix_lens = seq_lens - q_lens_t
cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
cu_seqlens[1:] = q_lens_t.cumsum(0)
total_q = sum(q_lens)
max_seqlen_q = max(q_lens)
q_shape = (total_q, NUM_Q_HEADS, HEAD_DIM)
q = torch.randn(q_shape, device="cuda", dtype=DTYPE)
# Allocate the main KV cache through the backend layout contract so the
# physical storage matches the active layout (contiguous NHD or strided
# HND), while the kernels and reference see the logical-NHD view.
kv_cache = _allocate_main_kv_via_contract(num_pages)
# Build sparse block indices with the same contract as the real M3 indexer:
# one forced local block, then score-selected older causal blocks.
topk_shape = (NUM_KV_HEADS, total_q, TOPK)
topk_idx = torch.full(topk_shape, -1, device="cuda", dtype=torch.int32)
q_start = 0
for q_len, prefix_len in zip(q_lens_t.tolist(), prefix_lens.tolist()):
for local_q in range(q_len):
current_block = (prefix_len + local_q) // BLOCK_SIZE
older_blocks = torch.randperm(
current_block, device="cuda", dtype=torch.int32
)
selected = torch.cat(
[
torch.tensor([current_block], device="cuda", dtype=torch.int32),
older_blocks[: TOPK - 1],
]
)
topk_idx[:, q_start + local_q, : selected.numel()] = selected
q_start += q_len
actual = torch.empty_like(q)
minimax_m3_sparse_attn(
q,
kv_cache,
topk_idx,
block_table,
cu_seqlens,
seq_lens,
prefix_lens,
max_seqlen_q,
NUM_KV_HEADS,
SM_SCALE,
actual,
)
expected = _reference_sparse_attn(
q,
kv_cache,
topk_idx,
block_table,
q_lens_t,
seq_lens,
prefix_lens,
)
torch.accelerator.synchronize()
error = (actual.float() - expected.float()).abs()
assert error.mean().item() < 2.5e-4
assert error.max().item() < 1.7e-2
def test_main_backend_layout_contract():
"""The main sparse backend exposes the logical-NHD shape and the
flash_attn-style stride order for each layout."""
nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
assert logical == (nb, 2, bs, h, d)
# The old HND-ordered shape is no longer the logical shape.
assert logical != (nb, 2, h, bs, d)
try:
set_kv_cache_layout("HND")
assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 3, 2, 4)
set_kv_cache_layout("NHD")
assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3, 4)
finally:
set_kv_cache_layout(None)
for layout in ("NHD", "HND"):
try:
set_kv_cache_layout(layout)
order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
finally:
set_kv_cache_layout(None)
# Valid permutation: no duplicates, covers every axis.
assert set(order) == set(range(len(order)))
# M3 has no cross-layer KV blocks.
with pytest.raises(NotImplementedError):
MiniMaxM3SparseBackend.get_kv_cache_stride_order(
include_num_layers_dimension=True
)
def test_main_backend_unknown_layout_raises(monkeypatch):
"""An unrecognized layout (injected past env-var validation) is rejected."""
import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
monkeypatch.setattr(sparse_attn_mod, "get_kv_cache_layout", lambda: "BOGUS")
with pytest.raises(ValueError, match="Unknown cache layout format"):
MiniMaxM3SparseBackend.get_kv_cache_stride_order()
def test_indexer_backend_stride_order_is_identity():
"""The 3-dim indexer cache must not inherit the parent's 5-element stride
order; it overrides to the 3-element identity so the allocator keeps the
contiguous layout."""
assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2)
# Cross-layer (per-layer-stacked) KV blocks are not supported.
with pytest.raises(NotImplementedError):
MiniMaxM3IndexerBackend.get_kv_cache_stride_order(
include_num_layers_dimension=True
)
# The stride order matches the 3-dim indexer shape rank.
indexer_shape = MiniMaxM3IndexerBackend.get_kv_cache_shape(
5, BLOCK_SIZE, 1, HEAD_DIM
)
assert len(indexer_shape) == 3
assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
def test_hnd_allocation_is_byte_identical_to_transpose():
"""Under HND the backend-visible logical view is byte-identical to the
pre-change allocate-HND-then-transpose(2, 3) workaround."""
nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
try:
set_kv_cache_layout("HND")
stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
finally:
set_kv_cache_layout(None)
physical_shape = tuple(logical[i] for i in stride_order)
# The physical (permuted) shape equals the old hardcoded HND shape.
assert physical_shape == (nb, 2, h, bs, d)
inv_order = [stride_order.index(i) for i in range(len(stride_order))]
raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
view = raw.permute(*inv_order)
expected = raw.view((nb, 2, h, bs, d)).transpose(2, 3)
assert view.shape == expected.shape
assert view.stride() == expected.stride()
assert view.storage_offset() == expected.storage_offset()
# Negative: the identity (wrong) stride order under HND does not reproduce
# the transpose view.
wrong_view = raw.view(logical)
assert wrong_view.stride() != expected.stride()
def test_main_cache_is_block_first_and_unpadded():
"""The allocator's contiguous-view branch (not the padded-strided branch)
is used for the main GQA cache: its spec is unpadded and the physical
layout keeps num_blocks as the first dimension under both layouts."""
from vllm.v1.kv_cache_interface import FullAttentionSpec
spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=NUM_KV_HEADS,
head_size=HEAD_DIM,
head_size_v=HEAD_DIM,
dtype=DTYPE,
)
# Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
assert spec.page_size_padded is None
logical = MiniMaxM3SparseBackend.get_kv_cache_shape(
4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
)
for layout in ("NHD", "HND"):
try:
set_kv_cache_layout(layout)
order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
finally:
set_kv_cache_layout(None)
inv_order = [order.index(i) for i in range(len(order))]
# Physical first dim is num_blocks (block-first); required by the
# padded-strided branch's block-first assumption if it were ever taken.
assert inv_order[0] == 0
assert logical[order[0]] == logical[0]
def _build_decode_inputs(
seq_lens_list: tuple[int, ...],
decode_query_len: int = 1,
num_padded_reqs: int = 0,
):
"""Shared decode setup: uniform query tokens per request, a non-identity
block table, and topk indices selecting the current block plus older causal
blocks for each query token."""
active_batch = len(seq_lens_list)
batch = active_batch + num_padded_reqs
pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list]
max_blocks = max(pages_per_req)
num_pages = sum(pages_per_req)
physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
base_page = 0
for req_id, num_req_pages in enumerate(pages_per_req):
block_table[req_id, :num_req_pages] = physical_pages[
base_page : base_page + num_req_pages
]
base_page += num_req_pages
seq_lens = torch.tensor(
(*seq_lens_list, *([0] * num_padded_reqs)),
device="cuda",
dtype=torch.int32,
)
q = torch.randn(
batch * decode_query_len, NUM_Q_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE
)
topk_idx = torch.full(
(NUM_KV_HEADS, batch * decode_query_len, TOPK),
-1,
device="cuda",
dtype=torch.int32,
)
token_id = 0
for req_id, seq_len in enumerate(seq_lens_list):
for local_q in range(decode_query_len):
query_pos = seq_len - decode_query_len + local_q
current_block = query_pos // BLOCK_SIZE
older_blocks = torch.randperm(
current_block, device="cuda", dtype=torch.int32
)
selected = torch.cat(
[
torch.tensor([current_block], device="cuda", dtype=torch.int32),
older_blocks[: TOPK - 1],
]
)
topk_idx[:, token_id, : selected.numel()] = selected
token_id += 1
return q, block_table, seq_lens, topk_idx, num_pages
@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
@pytest.mark.parametrize(
"seq_lens_list",
[(130, 257), (129, 200, 384)],
)
@pytest.mark.parametrize("decode_query_len", [1, 4])
@pytest.mark.parametrize("num_padded_reqs", [0, 2])
def test_decode_sparse_attention_correctness(
kv_layout: str,
seq_lens_list: tuple[int, ...],
decode_query_len: int,
num_padded_reqs: int,
):
"""Decode (split-K) parity under both layouts: this is the only coverage of
the decode-site cache feed, and the strided HND case fails if the kernel
ignores the cache strides."""
torch.manual_seed(0)
q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(
seq_lens_list, decode_query_len, num_padded_reqs
)
kv_cache = _allocate_main_kv_via_contract(num_pages)
actual = torch.empty_like(q)
minimax_m3_sparse_attn_decode(
q,
kv_cache,
topk_idx,
block_table,
seq_lens,
NUM_KV_HEADS,
SM_SCALE,
actual,
decode_query_len,
)
# Reuse the prefill reference: decode is a uniform query chunk ending at
# seq_len - 1 for each request.
active_batch = len(seq_lens_list)
active_tokens = active_batch * decode_query_len
q_lens_t = torch.full(
(len(seq_lens_list),), decode_query_len, device="cuda", dtype=torch.int32
)
active_seq_lens = seq_lens[:active_batch]
prefix_lens = active_seq_lens - q_lens_t
expected = _reference_sparse_attn(
q[:active_tokens],
kv_cache,
topk_idx[:, :active_tokens],
block_table[:active_batch],
q_lens_t,
active_seq_lens,
prefix_lens,
)
torch.accelerator.synchronize()
error = (actual[:active_tokens].float() - expected.float()).abs()
assert error.mean().item() < 2.5e-4
assert error.max().item() < 1.7e-2
def test_decode_wrong_layout_breaks_parity():
"""Negative (AC-3/AC-5): consuming the physical HND buffer as if it were
already contiguous-NHD (i.e. skipping the allocator's inverse permute)
reorders the K/V content, so the decode output no longer matches the
reference computed on the correct logical view. The mislabeled tensor keeps
the same shape as the correct view, so the kernel stays in bounds."""
torch.manual_seed(0)
seq_lens_list = (130, 257)
q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
# Physical HND storage [blocks, 2, heads, block, dim].
phys = torch.randn(
(num_pages, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM), device="cuda", dtype=DTYPE
)
# Correct logical-NHD view (strided) vs. the same bytes mislabeled as a
# contiguous-NHD cache — same shape, different content mapping.
correct = phys.permute(0, 1, 3, 2, 4)
wrong = phys.reshape(num_pages, 2, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM)
q_lens_t = torch.ones(len(seq_lens_list), device="cuda", dtype=torch.int32)
prefix_lens = seq_lens - q_lens_t
expected = _reference_sparse_attn(
q, correct, topk_idx, block_table, q_lens_t, seq_lens, prefix_lens
)
actual = torch.empty_like(q)
minimax_m3_sparse_attn_decode(
q, wrong, topk_idx, block_table, seq_lens, NUM_KV_HEADS, SM_SCALE, actual, 1
)
torch.accelerator.synchronize()
assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2
def _make_attn_group(backend, spec):
return AttentionGroup(
backend=backend,
layer_names=["main"],
kv_cache_spec=spec,
kv_cache_group_id=0,
)
def test_main_cache_byte_identical_through_production_allocator():
"""AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main
`FullAttentionSpec` under HND and assert the backend-visible view has the
same shape, stride, and storage offset as the pre-change
allocate-HND-then-transpose path; the indexer `MLAAttentionSpec` allocates
through the same path to its 3-dim shape."""
nb = 4
spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=NUM_KV_HEADS,
head_size=HEAD_DIM,
head_size_v=HEAD_DIM,
dtype=DTYPE,
)
raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
group = _make_attn_group(MiniMaxM3SparseBackend, spec)
try:
set_kv_cache_layout("HND")
kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {})
finally:
set_kv_cache_layout(None)
view = kv_caches["main"]
oracle = raw.view(DTYPE).view((nb, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM))
oracle = oracle.transpose(2, 3)
assert tuple(view.shape) == tuple(oracle.shape)
assert view.stride() == oracle.stride()
assert view.storage_offset() == oracle.storage_offset()
# Indexer cache allocates through the same path under both layouts.
ispec = MLAAttentionSpec(
block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
)
for layout in ("NHD", "HND"):
iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
igroup = AttentionGroup(
backend=MiniMaxM3IndexerBackend,
layer_names=["idx"],
kv_cache_spec=ispec,
kv_cache_group_id=0,
)
try:
set_kv_cache_layout(layout)
iout = _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
finally:
set_kv_cache_layout(None)
assert tuple(iout["idx"].shape) == (nb, BLOCK_SIZE, HEAD_DIM)
def test_indexer_inherited_stride_order_trips_allocator_assert():
"""AC-4 negative: without the indexer override, the inherited 5-element
stride order trips the allocator's `len(stride_order) == len(shape)` assert
for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the
allocator's `(AttributeError, NotImplementedError)` fallback."""
class _BrokenIndexerBackend(MiniMaxM3IndexerBackend):
# Simulate inheriting the parent's 5-element stride order.
get_kv_cache_stride_order = staticmethod(
MiniMaxM3SparseBackend.get_kv_cache_stride_order
)
nb = 4
ispec = MLAAttentionSpec(
block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
)
iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
igroup = AttentionGroup(
backend=_BrokenIndexerBackend,
layer_names=["idx"],
kv_cache_spec=ispec,
kv_cache_group_id=0,
)
try:
set_kv_cache_layout("HND")
with pytest.raises(AssertionError):
_reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
finally:
set_kv_cache_layout(None)
def test_padded_main_cache_is_flagged():
"""AC-2.1 negative: the M3 main cache relies on the allocator's
contiguous-view branch (`page_size_padded is None`). A spec that sets
`page_size_padded` is explicitly flagged rather than silently wrong-strided."""
def _require_unpadded_block_first(spec, stride_order):
inv_order = [stride_order.index(i) for i in range(len(stride_order))]
assert spec.page_size_padded is None, (
"main GQA cache must be unpadded to use the contiguous-view "
"allocator branch"
)
assert inv_order[0] == 0, "main GQA cache must remain block-first"
try:
set_kv_cache_layout("HND")
stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
finally:
set_kv_cache_layout(None)
good = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=NUM_KV_HEADS,
head_size=HEAD_DIM,
head_size_v=HEAD_DIM,
dtype=DTYPE,
)
_require_unpadded_block_first(good, stride_order) # passes
padded = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=NUM_KV_HEADS,
head_size=HEAD_DIM,
head_size_v=HEAD_DIM,
dtype=DTYPE,
page_size_padded=good.page_size_bytes + 128,
)
with pytest.raises(AssertionError):
_require_unpadded_block_first(padded, stride_order)
@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
def test_reshape_and_cache_flash_write_persists(kv_layout: str):
"""AC-5 write path: the `reshape_and_cache_flash` write site now consumes
`self.kv_cache.unbind(1)` directly. Writing through those views must persist
into the bound storage (read back through an independent logical view) under
both layouts — a `.contiguous()` copy of the unbind slice would leave the
bound storage unchanged."""
torch.manual_seed(0)
num_pages = 4
kv_cache = _allocate_main_kv_via_contract(num_pages)
with torch.no_grad():
kv_cache.zero_()
# Exactly the production write-site code under test.
key_cache, value_cache = kv_cache.unbind(1)
num_tokens = 12
slot_mapping = torch.randperm(num_pages * BLOCK_SIZE, device="cuda")[
:num_tokens
].to(torch.int64)
key = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
value = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
scale = torch.ones((), device="cuda")
ops.reshape_and_cache_flash(
key, value, key_cache, value_cache, slot_mapping, "auto", scale, scale
)
torch.accelerator.synchronize()
# Read back through the independent logical view; proves the writes landed
# in the engine-bound storage, not a detached copy.
for t in range(num_tokens):
slot = int(slot_mapping[t].item())
blk, intra = divmod(slot, BLOCK_SIZE)
torch.testing.assert_close(kv_cache[blk, 0, intra], key[t])
torch.testing.assert_close(kv_cache[blk, 1, intra], value[t])
tests/kernels/core/test_fused_allreduce_gemma_rms_norm.py
+109
View File
@@ -0,0 +1,109 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the manual AllReduce + GemmaRMSNorm fusion used by MiniMax M3.
``fused_allreduce_gemma_rms_norm`` must match the unfused model path, i.e.
``GemmaRMSNorm(all_reduce(partial), residual)``, both on the flashinfer fast
path (TP>1 with flashinfer + NVSwitch) and on the eager fallback (TP==1, or when
flashinfer is unavailable / the GPU has no NVSwitch).
"""
import pytest
import torch
from torch.multiprocessing import spawn
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
fused_allreduce_gemma_rms_norm,
)
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
from vllm.platforms import current_platform
from vllm.utils.network_utils import get_open_port
from vllm.utils.torch_utils import set_random_seed
@ensure_current_vllm_config()
def _worker_fused_ar_norm(
local_rank,
world_size,
port,
num_tokens,
hidden_size,
dtype,
seed,
eps,
):
"""Per-rank worker: compare the fused helper vs all_reduce + GemmaRMSNorm."""
device = torch.device(f"cuda:{local_rank}")
torch.accelerator.set_device_index(device)
init_test_distributed_environment(
world_size, 1, local_rank, port, local_rank=local_rank
)
# Norm weights are identical across ranks (replicated GemmaRMSNorm).
set_random_seed(seed)
norm = GemmaRMSNorm(hidden_size, eps=eps).cuda().to(dtype)
with torch.no_grad():
norm.weight.normal_(mean=0.0, std=0.1)
# Residual is shared across ranks; the partial o_proj output differs per rank
# (each rank holds a partial sum that all_reduce combines).
torch.manual_seed(seed + 7)
residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
torch.manual_seed(seed + 1000 + local_rank)
partial = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
# Reference: the unfused model path.
reduced = tensor_model_parallel_all_reduce(partial.clone())
ref_out, ref_res = norm(reduced, residual.clone())
# Fused helper (flashinfer fast path when available, else fallback).
out, res = fused_allreduce_gemma_rms_norm(partial.clone(), residual.clone(), norm)
torch.accelerator.synchronize()
torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(res, ref_res, atol=2e-2, rtol=2e-2)
cleanup_dist_env_and_memory()
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason="CUDA required",
)
# world_size=1 exercises the TP==1 identity branch on a single GPU; >1 exercises
# the all_reduce + GemmaRMSNorm equivalence (flashinfer kernel or fallback).
@pytest.mark.parametrize("world_size", [1, 2, 4])
@pytest.mark.parametrize("num_tokens", [1, 128, 333])
@pytest.mark.parametrize("hidden_size", [2048, 4096])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("eps", [1e-6])
@pytest.mark.parametrize("seed", [42])
def test_fused_allreduce_gemma_rms_norm(
world_size,
num_tokens,
hidden_size,
dtype,
eps,
seed,
):
num_gpus = current_platform.device_count()
if num_gpus < world_size:
pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
port = str(get_open_port())
spawn(
_worker_fused_ar_norm,
args=(
world_size,
port,
num_tokens,
hidden_size,
dtype,
seed,
eps,
),
nprocs=world_size,
join=True,
)
tests/kernels/test_fp32_router_gemm.py
+21 -15
View File
@@ -1,6 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for fp32_router_gemm kernel: activation×weight→fp32, H=3072, E=256.
"""Tests for fp32_router_gemm kernel: activation×weight→fp32.
Supported (hidden_size, num_experts) pairs:
(3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
Correctness baseline: torch.matmul in float64.
"""
@@ -10,8 +13,8 @@ import torch
from vllm._custom_ops import fp32_router_gemm
NUM_EXPERTS = 256
HIDDEN_DIM = 3072
# (hidden_size, num_experts)
SHAPES = [(3072, 256), (6144, 128)]
# Absolute tolerance for fp32 kernel vs float64 reference
ATOL_FP32 = 2e-4
ATOL_BF16 = 2e-2 # bf16 activation has lower precision
@@ -30,49 +33,52 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
return torch.nn.functional.linear(mat_a.float(), mat_b.float())
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
def test_fp32_activation(num_tokens: int):
def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
"""fp32 activation → fp32 output should match reference closely."""
_requires_sm90()
torch.manual_seed(42)
device = torch.device("cuda")
mat_a = torch.randn(num_tokens, HIDDEN_DIM, dtype=torch.float32, device=device)
mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
mat_a = torch.randn(num_tokens, hidden_dim, dtype=torch.float32, device=device)
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
out = fp32_router_gemm(mat_a, mat_b)
ref = _ref(mat_a, mat_b)
assert out.shape == (num_tokens, NUM_EXPERTS)
assert out.shape == (num_tokens, num_experts)
assert out.dtype == torch.float32
torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0)
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
def test_bf16_activation(num_tokens: int):
def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
"""bf16 activation → fp32 output should match reference within bf16 error."""
_requires_sm90()
torch.manual_seed(42)
device = torch.device("cuda")
mat_a_bf16 = torch.randn(
num_tokens, HIDDEN_DIM, dtype=torch.bfloat16, device=device
num_tokens, hidden_dim, dtype=torch.bfloat16, device=device
)
mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
out = fp32_router_gemm(mat_a_bf16, mat_b)
ref = _ref(mat_a_bf16, mat_b).to(device)
assert out.shape == (num_tokens, NUM_EXPERTS)
assert out.shape == (num_tokens, num_experts)
assert out.dtype == torch.float32
torch.testing.assert_close(out, ref, atol=ATOL_BF16, rtol=0)
def test_output_shape_and_dtype():
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
def test_output_shape_and_dtype(hidden_dim: int, num_experts: int):
"""Basic shape and dtype checks."""
_requires_sm90()
device = torch.device("cuda")
mat_a = torch.randn(4, HIDDEN_DIM, dtype=torch.float32, device=device)
mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
mat_a = torch.randn(4, hidden_dim, dtype=torch.float32, device=device)
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
out = fp32_router_gemm(mat_a, mat_b)
assert out.shape == (4, NUM_EXPERTS)
assert out.shape == (4, num_experts)
assert out.dtype == torch.float32
assert out.device.type == "cuda"
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+244
View File
@@ -0,0 +1,244 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit test for the horizontally-fused MiniMax-M3 attention pre-processing
kernel:
fused_minimax_m3_qknorm_rope_kv_insert
- q / k / index_q / index_k: Gemma RMSNorm + partial NeoX RoPE (in place)
- sparse (insert) mode: scatter k/v into the paged bf16 KV cache and the
index key into the index cache by its own slot mapping.
Reference: PyTorch Gemma RMSNorm with the same dtype materialization boundary
as the unfused path, followed by vLLM CUDA rotary_embedding-style NeoX RoPE.
"""
import pytest
import torch
import vllm._custom_ops as ops
HEAD_DIM = 128
ROTARY_DIM = 64
def _op_available() -> bool:
return hasattr(torch.ops._C, "fused_minimax_m3_qknorm_rope_kv_insert")
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available() or not _op_available(),
reason="CUDA not available or fused MiniMax-M3 op not built in",
)
def make_cos_sin_cache(max_pos, rotary_dim, base, dtype, device):
inv_freq = 1.0 / (
base
** (
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device)
/ rotary_dim
)
)
t = torch.arange(max_pos, dtype=torch.float32, device=device)
freqs = torch.einsum("i,j->ij", t, inv_freq) # [max_pos, rotary_dim/2]
cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1) # [max_pos, rotary_dim]
return cache.to(dtype)
def gemma_rmsnorm(x, weight, eps):
"""x: [..., 128]; weight: [128]. Returns original dtype."""
xf = x.float()
var = xf.pow(2).mean(dim=-1, keepdim=True)
out = xf * torch.rsqrt(var + eps)
out = out * (1.0 + weight.float())
return out.to(x.dtype)
def apply_rope_neox_partial(x, positions, cos_sin_cache, rotary_dim):
"""NeoX-style RoPE on the leading rotary_dim dims; rest pass through.
x: [num_tokens, num_heads, head_dim]
cos_sin_cache: [max_pos, rotary_dim] (cos||sin), read as float (matches the
kernel, which loads the bf16 cache and converts to fp32).
"""
half = rotary_dim // 2
cs = cos_sin_cache[positions].float() # [num_tokens, rotary_dim]
cos = cs[..., :half].unsqueeze(1) # [nt, 1, half]
sin = cs[..., half:].unsqueeze(1)
rot = x[..., :rotary_dim].float()
x1 = rot[..., :half]
x2 = rot[..., half:]
o1 = x1 * cos - x2 * sin
o2 = x2 * cos + x1 * sin
out = x.clone()
out[..., :half] = o1
out[..., half:rotary_dim] = o2
return out.to(x.dtype)
def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
"""[nt, nheads, 128] -> Gemma norm + neox partial rope."""
normed = gemma_rmsnorm(x, weight, eps)
roped = apply_rope_neox_partial(normed, positions, cos_sin_cache, ROTARY_DIM)
return roped
# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
@pytest.mark.parametrize("num_heads,num_kv_heads", [(8, 2), (16, 4), (64, 4)])
def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads):
torch.manual_seed(0)
device, dtype, eps = "cuda", torch.bfloat16, 1e-6
base, max_pos = 5_000_000.0, 4096
q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
positions = torch.randint(
0, max_pos, (num_tokens,), dtype=torch.int64, device=device
)
qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
qkv = torch.randn(num_tokens, qsz + 2 * kvsz, dtype=dtype, device=device)
qkv_orig = qkv.clone()
ops.fused_minimax_m3_qknorm_rope_kv_insert(
qkv, q_w, k_w, cos_sin, positions, num_heads, num_kv_heads, ROTARY_DIM, eps
)
q_out, k_out, v_out = qkv.split([qsz, kvsz, kvsz], dim=-1)
q_in, k_in, v_in = qkv_orig.split([qsz, kvsz, kvsz], dim=-1)
q_ref = norm_rope_ref(
q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps
).view(num_tokens, qsz)
k_ref = norm_rope_ref(
k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
k_w,
positions,
cos_sin,
eps,
).view(num_tokens, kvsz)
torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
# V is untouched.
torch.testing.assert_close(v_out, v_in, rtol=0, atol=0)
# ── Test 2: sparse mode (full: index branch + cache inserts) ─────────────────
@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
@pytest.mark.parametrize("block_size", [16, 64])
def test_sparse_full(num_tokens, block_size):
torch.manual_seed(1)
device, dtype, eps = "cuda", torch.bfloat16, 1e-6
base, max_pos = 5_000_000.0, 4096
num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
positions = torch.randint(
0, max_pos, (num_tokens,), dtype=torch.int64, device=device
)
qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
# Single fused tensor packing [q | k | v | index_q | index_k].
qkv = torch.randn(
num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
)
qkv_orig = qkv.clone()
splits = [qsz, kvsz, kvsz, iqsz, iksz]
num_blocks = (num_tokens + block_size - 1) // block_size + 1
kv_cache = torch.zeros(
num_blocks, 2, block_size, num_kv_heads, HEAD_DIM, dtype=dtype, device=device
)
index_cache = torch.zeros(
num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
)
slot_mapping = torch.randperm(
num_blocks * block_size, dtype=torch.int64, device=device
)[:num_tokens]
index_slot_mapping = torch.roll(slot_mapping, shifts=1)
# Contiguous gather targets: the kernel writes the normed/roped q and
# index_q here (de-interleaved from the packed qkv); k/v/index_k stay in
# place inside qkv and are scatter-inserted into the caches.
q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
index_q = torch.empty(num_tokens, iqsz, dtype=dtype, device=device)
ops.fused_minimax_m3_qknorm_rope_kv_insert(
qkv,
q_w,
k_w,
cos_sin,
positions,
num_heads,
num_kv_heads,
ROTARY_DIM,
eps,
iq_w,
ik_w,
num_idx_heads,
slot_mapping,
index_slot_mapping,
kv_cache,
index_cache,
block_size,
q_out,
index_q,
)
# ── norm+rope parity. q/index_q land in their gather buffers; k/index_k are
# rewritten in place inside qkv. ──
_, k_out, _, _, index_k = qkv.split(splits, dim=-1)
q_in, k_in, v_in, iq_orig, ik_orig = qkv_orig.split(splits, dim=-1)
q_ref = norm_rope_ref(
q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps
).view(num_tokens, qsz)
k_ref = norm_rope_ref(
k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
k_w,
positions,
cos_sin,
eps,
).view(num_tokens, kvsz)
iq_ref = norm_rope_ref(
iq_orig.view(num_tokens, num_idx_heads, HEAD_DIM),
iq_w,
positions,
cos_sin,
eps,
).view(num_tokens, num_idx_heads * HEAD_DIM)
ik_ref = norm_rope_ref(
ik_orig.view(num_tokens, 1, HEAD_DIM), ik_w, positions, cos_sin, eps
).view(num_tokens, HEAD_DIM)
torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2)
torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2)
# ── Cache inserts. ──
# Main cache layout is [num_blocks, 2, block_size, num_kv_heads, head_dim]
# (the K/V axis sits *before* block_size); index cache is [nb, bs, head_dim].
idx_flat = index_cache.view(num_blocks * block_size, HEAD_DIM)
k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM) # v is raw (no norm/rope)
for t in range(num_tokens):
s = slot_mapping[t].item()
b, pos = s // block_size, s % block_size
torch.testing.assert_close(
kv_cache[b, 0, pos], k_ref_h[t], rtol=1e-2, atol=1e-2
)
torch.testing.assert_close(kv_cache[b, 1, pos], v_ref_h[t], rtol=0, atol=0)
index_s = index_slot_mapping[t].item()
torch.testing.assert_close(idx_flat[index_s], ik_ref[t], rtol=1e-2, atol=1e-2)
tests/kernels/test_minimax_m3_amd_ops.py
+337
View File
@@ -0,0 +1,337 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reference-vs-optimized unit tests for the MiniMax-M3 AMD/ROCm fused kernels.
Each optimized kernel added for the ROCm port has a slow PyTorch reference; the
tests assert the two agree within tolerance:
* Gemma RMSNorm (plain + fused-add-residual) -> fp32 PyTorch normalize
* SwiGLU-OAI (split layout) -> fp32 PyTorch elementwise
* Fused MXFP8 activation quant (Triton) -> _mxfp8_e4m3_quantize_torch
* Native MXFP8 linear (dot_scaled) -> dequant-to-bf16 @ matmul
* Native MXFP8 MoE (dot_scaled grouped GEMM) -> dequant-to-bf16 MoE math
The native MXFP8 GEMMs also guard the ``dot_scaled`` rhs-scale orientation: the
scale is loaded ``[N, K//32]`` and passed WITHOUT transpose; a stray ``.T``
makes the shape ``[K//32, N]`` and Triton raises before producing output, so any
regression there fails these tests loudly.
Hardware scope: the whole module is ROCm-only (these are the AMD path; NVIDIA
uses the FlashInfer kernels). The norm/activation/quant kernels run on any ROCm
arch; the native MXFP8 ``dot_scaled`` linear/MoE tests are additionally gated to
CDNA4 gfx95x (``@requires_gfx950``) since gfx942 uses the BF16 emulation path.
Run: pytest tests/kernels/test_minimax_m3_amd_ops.py -v
"""
import pytest
import torch
from vllm.platforms import current_platform
if not current_platform.is_rocm():
pytest.skip("MiniMax-M3 AMD fused ops require ROCm.", allow_module_level=True)
if not torch.cuda.is_available():
pytest.skip("Requires a GPU.", allow_module_level=True)
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( # noqa: E402
_mxfp8_e4m3_quantize_torch,
_mxfp8_e4m3_quantize_triton,
dequant_mxfp8_to_bf16,
)
from vllm.models.minimax_m3.amd.ops import ( # noqa: E402
gemma_fused_add_rmsnorm,
gemma_rmsnorm,
swiglu_oai_split,
)
from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402
DEVICE = "cuda"
EPS = 1e-6
def _gcn_arch() -> str:
try:
return torch.cuda.get_device_properties(0).gcnArchName
except Exception: # pragma: no cover - no device / non-AMD
return ""
# The pure-Triton norm/activation/quant kernels run on any ROCm arch (CDNA3
# gfx942 and CDNA4 gfx950). The native MXFP8 ``dot_scaled`` GEMMs (linear + MoE)
# use CDNA4 hardware microscaling and are gated to gfx95x in the source
# (``RocmDotScaledMxfp8LinearKernel.is_supported``; the MoE oracle routes gfx942
# to the BF16 emulation path instead) — so those tests are gfx950-only.
requires_gfx950 = pytest.mark.skipif(
"gfx95" not in _gcn_arch(),
reason="native MXFP8 dot_scaled is a CDNA4 (gfx95x) feature; "
"gfx942 uses the BF16 emulation path instead.",
)
def _relerr(a: torch.Tensor, b: torch.Tensor) -> float:
a = a.float()
b = b.float()
return ((a - b).norm() / (b.norm() + 1e-8)).item()
# --------------------------------------------------------------------------- #
# Gemma RMSNorm
# --------------------------------------------------------------------------- #
def _ref_gemma_rmsnorm(x, w, eps, residual=None):
orig_dtype = x.dtype
xf = x.float()
res_out = None
if residual is not None:
xf = xf + residual.float()
res_out = xf.to(orig_dtype)
xf = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps)
xf = xf * (1.0 + w.float())
out = xf.to(orig_dtype)
return out if residual is None else (out, res_out)
@pytest.mark.parametrize("shape", [(1, 4096), (37, 6144), (128, 2048)])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("seed", [0, 1234])
@torch.inference_mode()
def test_gemma_rmsnorm(shape, dtype, seed):
torch.manual_seed(seed)
x = torch.randn(*shape, device=DEVICE, dtype=dtype)
w = torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1
got = gemma_rmsnorm(x, w, EPS)
ref = _ref_gemma_rmsnorm(x, w, EPS)
assert got.shape == x.shape
assert _relerr(got, ref) < 5e-3
@pytest.mark.parametrize("shape", [(1, 6144), (64, 4096)])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@torch.inference_mode()
def test_gemma_fused_add_rmsnorm(shape, dtype):
torch.manual_seed(0)
x = torch.randn(*shape, device=DEVICE, dtype=dtype)
res = torch.randn(*shape, device=DEVICE, dtype=dtype)
w = torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1
got_out, got_res = gemma_fused_add_rmsnorm(x, res, w, EPS)
ref_out, ref_res = _ref_gemma_rmsnorm(x, w, EPS, residual=res)
assert _relerr(got_out, ref_out) < 5e-3
# residual_out is the pre-norm sum (x + res): bit-for-bit identical cast.
assert torch.equal(got_res, ref_res)
@torch.inference_mode()
def test_gemma_rmsnorm_per_head_strided():
"""q_norm/k_norm normalize a non-contiguous ``qkv.split`` slice over head_dim."""
torch.manual_seed(0)
T, H, D, kv = 7, 48, 128, 8
total = (H + 2 * kv) * D
qkv = torch.randn(T, total, device=DEVICE, dtype=torch.bfloat16)
q = qkv[..., : H * D] # non-contiguous view (row stride == total)
q_by_head = q.view(T, H, D)
assert not q_by_head.is_contiguous()
w = torch.randn(D, device=DEVICE, dtype=torch.bfloat16) * 0.1
got = gemma_rmsnorm(q_by_head, w, EPS)
ref = _ref_gemma_rmsnorm(q_by_head, w, EPS)
assert got.shape == q_by_head.shape
assert _relerr(got, ref) < 5e-3
def test_num_warps_monotonic():
assert _num_warps(128) <= _num_warps(2048) <= _num_warps(8192)
# --------------------------------------------------------------------------- #
# SwiGLU-OAI (split layout)
# --------------------------------------------------------------------------- #
def _ref_swiglu(gate_up, alpha, beta, limit):
d = gate_up.shape[-1] // 2
gate = gate_up[..., :d].float()
up = gate_up[..., d:].float()
if limit is not None:
gate = gate.clamp(max=limit)
up = up.clamp(min=-limit, max=limit)
return (gate * torch.sigmoid(alpha * gate) * (up + beta)).to(gate_up.dtype)
@pytest.mark.parametrize("m,inter", [(1, 768), (64, 1536), (128, 1024)])
@pytest.mark.parametrize("limit", [7.0, None])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@torch.inference_mode()
def test_swiglu_oai_split(m, inter, limit, dtype):
torch.manual_seed(0)
gate_up = torch.randn(m, 2 * inter, device=DEVICE, dtype=dtype)
got = swiglu_oai_split(gate_up, alpha=1.702, beta=1.0, limit=limit)
ref = _ref_swiglu(gate_up, 1.702, 1.0, limit)
assert got.shape == (m, inter)
assert _relerr(got, ref) < 5e-3
# --------------------------------------------------------------------------- #
# Fused MXFP8 activation quant (Triton vs torch reference)
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("shape", [(64, 4096), (1, 6144), (333, 2048)])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@torch.inference_mode()
def test_mxfp8_quant_triton_matches_torch(shape, dtype):
torch.manual_seed(0)
x = torch.randn(*shape, device=DEVICE, dtype=dtype)
xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False)
xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
# E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at
# most a 1-step difference at exact powers of two.
assert (s_k.int() - s_t.int()).abs().max().item() <= 1
# Dequantized values agree to fp8 granularity.
deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
deq_k = dequant_mxfp8_to_bf16(xq_k, s_k)
assert _relerr(deq_k, deq_t) < 1e-2
# --------------------------------------------------------------------------- #
# Native MXFP8 linear (dot_scaled) vs dequant-to-bf16 matmul
# --------------------------------------------------------------------------- #
@requires_gfx950
@pytest.mark.parametrize("m,n,k", [(64, 256, 128), (37, 512, 256), (1, 6144, 4096)])
@torch.inference_mode()
def test_mxfp8_native_linear(m, n, k):
from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
_mxfp8_dot_scaled_linear,
)
torch.manual_seed(0)
w_bf16 = torch.randn(n, k, device=DEVICE, dtype=torch.bfloat16) * 0.1
w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) * 0.5
got = _mxfp8_dot_scaled_linear(x, w_fp8, w_scale)
# Reference: consume the SAME quantized weights (isolates activation-quant
# noise) -> dequant to bf16, plain matmul.
w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
ref = torch.nn.functional.linear(x, w_deq).to(x.dtype)
assert got.shape == (m, n)
# Only the activation is re-quantized inside the kernel -> small MX noise.
assert _relerr(got, ref) < 5e-2
# --------------------------------------------------------------------------- #
# Native MXFP8 MoE (dot_scaled grouped GEMM) vs dequant-to-bf16 MoE math
# --------------------------------------------------------------------------- #
def _ref_moe(x, w13, w2, topk_weights, topk_ids, alpha, beta, limit):
T, H = x.shape
inter = w2.shape[-1]
top_k = topk_ids.shape[1]
out = torch.zeros(T, H, device=x.device, dtype=torch.float32)
for t in range(T):
for j in range(top_k):
e = int(topk_ids[t, j].item())
g1 = x[t].float() @ w13[e].float().T # [2I]
gate = g1[:inter]
up = g1[inter:]
if limit is not None:
gate = gate.clamp(max=limit)
up = up.clamp(min=-limit, max=limit)
act = gate * torch.sigmoid(alpha * gate) * (up + beta)
g2 = act @ w2[e].float().T # [H]
out[t] += topk_weights[t, j].float() * g2
return out.to(x.dtype)
@requires_gfx950
@pytest.mark.parametrize(
"T,H,inter,E,top_k", [(8, 256, 512, 8, 2), (1, 512, 256, 16, 4)]
)
@torch.inference_mode()
def test_mxfp8_native_moe(T, H, inter, E, top_k):
from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
fused_moe_mxfp8_native,
)
torch.manual_seed(0)
alpha, beta, limit = 1.702, 1.0, 7.0
w13_bf16 = torch.randn(E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
w2_bf16 = torch.randn(E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(
w13_bf16, is_sf_swizzled_layout=False
)
w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(w2_bf16, is_sf_swizzled_layout=False)
x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
topk_weights = topk_weights.to(torch.float32)
topk_ids = topk_ids.to(torch.int32)
got = fused_moe_mxfp8_native(
x,
w13_fp8,
w13_scale,
w2_fp8,
w2_scale,
topk_weights,
topk_ids,
alpha=alpha,
beta=beta,
limit=limit,
global_num_experts=E,
expert_map=None,
)
# Reference consumes the dequantized weights (same bits the kernel reads).
w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids, alpha, beta, limit)
assert got.shape == (T, H)
assert _relerr(got, ref) < 5e-2
# --------------------------------------------------------------------------- #
# MXFP8 linear emulation: BF16-at-load (default) vs per-step dequant + switch
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("shape", [(512, 2048), (1, 6144)])
@pytest.mark.parametrize("act_dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("dequant_at_load", [True, False])
@torch.inference_mode()
def test_mxfp8_linear_emulation_bf16_at_load(
shape, act_dtype, dequant_at_load, monkeypatch
):
"""EmulationMxfp8LinearKernel load-time BF16 dequant (default) and the
``VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0`` per-step fallback must produce the
same result; the dtype-match (BF16/FP16 activations) must also hold."""
from vllm.model_executor.kernels.linear.mxfp8.emulation import (
EmulationMxfp8LinearKernel,
)
from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
Mxfp8LinearLayerConfig,
)
monkeypatch.setenv(
"VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD", "1" if dequant_at_load else "0"
)
N, K = shape
torch.manual_seed(0)
w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16)
w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
assert w_scale.shape == (N, K // 32)
# Reference: dequant once, plain linear in the activation dtype.
w_ref = dequant_mxfp8_to_bf16(w_fp8, w_scale).to(act_dtype)
x = torch.randn(7, K, device=DEVICE, dtype=act_dtype)
out_ref = torch.nn.functional.linear(x, w_ref)
layer = torch.nn.Module()
layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
kernel = EmulationMxfp8LinearKernel(Mxfp8LinearLayerConfig())
kernel.process_weights_after_loading(layer)
if dequant_at_load:
# weights converted to BF16 at load (>= 2-byte)
assert layer.weight.element_size() >= 2
else:
# opt-out: weights stay 1-byte MXFP8, dequant happens per-step
assert layer.weight.element_size() == 1
out = kernel.apply_weights(layer, x)
assert out.dtype == act_dtype # dtype-match preserved (no tl.dot/F.linear crash)
assert _relerr(out.float(), out_ref.float()) < 2e-2
tests/models/multimodal/processing/test_minimax_m3.py
+138
View File
@@ -0,0 +1,138 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for MiniMax-M3 VL ``max_long_side_pixel`` resize support.
These exercise the vendored processor directly (no checkpoint / GPU needed), so
they validate the long-side resize spec and the resulting prompt-token counts
deterministically.
"""
import pytest
import torch
from vllm.transformers_utils.processors.minimax_m3 import (
IMAGE_MAX_TOTAL_PIXELS,
MIN_SHORT_SIDE_PIXEL,
VIDEO_MAX_TOTAL_PIXELS,
MiniMaxM3VLImageProcessor,
MiniMaxM3VLVideoProcessor,
smart_resize,
)
# Long sides are multiples of patch_size*merge_size (28) so the rounding is
# exact and the expected token counts are unambiguous.
LONG_SIDES = [252, 504, 1008]
MERGE2 = 2**2 # merge_size ** 2
def _image_tokens(grid_thw) -> int:
g = list(grid_thw)
return int(g[0] * g[1] * g[2]) // MERGE2
# --------------------------------------------------------------------------- #
# smart_resize: the long-side spec (a) shrink / (b) enlarge / (c) hard cap
# --------------------------------------------------------------------------- #
def test_smart_resize_long_side_shrink():
# (a) long side exceeds the cap -> shrink so the long side equals the cap.
h, w = smart_resize(
2048, 1024, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9
)
assert max(h, w) == 1008
assert (h, w) == (1008, 504) # aspect ratio preserved
def test_smart_resize_short_side_enlarge():
# (b) long side within the cap but short side below the floor -> enlarge so
# the short side reaches min_short_side_pixel.
h, w = smart_resize(
200, 40, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9
)
assert min(h, w) == MIN_SHORT_SIDE_PIXEL # 112
def test_smart_resize_total_pixels_raises():
# (c) still over the area cap after resizing -> raise instead of inferring.
with pytest.raises(ValueError, match="max_total_pixels"):
smart_resize(
5000,
5000,
factor=28,
max_long_side_pixel=4000,
max_total_pixels=IMAGE_MAX_TOTAL_PIXELS,
)
def test_smart_resize_backward_compatible_area_bound():
# Without max_long_side_pixel the original Qwen-style area bound is used.
assert smart_resize(2048, 2048, factor=28, max_pixels=451584) == (672, 672)
# --------------------------------------------------------------------------- #
# Image processor: monotonic prompt-token counts for 252 < 504 < 1008
# --------------------------------------------------------------------------- #
def test_image_tokens_increase_with_max_long_side_pixel():
proc = MiniMaxM3VLImageProcessor()
counts = []
for long_side in LONG_SIDES:
patches = proc.get_number_of_image_patches(
2048, 2048, images_kwargs={"max_long_side_pixel": long_side}
)
counts.append(patches // MERGE2)
assert counts == [81, 324, 1296]
assert counts[0] < counts[1] < counts[2]
def test_image_processor_defaults_match_spec():
proc = MiniMaxM3VLImageProcessor()
assert proc.max_long_side_pixel is None # opt-in
assert proc.min_short_side_pixel == MIN_SHORT_SIDE_PIXEL
assert proc.max_total_pixels == IMAGE_MAX_TOTAL_PIXELS
def test_image_preprocess_pipeline_monotonic():
proc = MiniMaxM3VLImageProcessor()
image = torch.randint(0, 255, (3, 2048, 2048), dtype=torch.uint8)
counts = []
for long_side in LONG_SIDES:
out = proc.preprocess(
[image],
do_resize=True,
max_long_side_pixel=long_side,
return_tensors="pt",
)
counts.append(_image_tokens(out["image_grid_thw"][0]))
assert counts == [81, 324, 1296]
# --------------------------------------------------------------------------- #
# Video processor: same monotonic behavior + volumetric (w*h*frames) cap
# --------------------------------------------------------------------------- #
def test_video_tokens_increase_with_max_long_side_pixel():
proc = MiniMaxM3VLVideoProcessor()
assert proc.max_total_pixels == VIDEO_MAX_TOTAL_PIXELS
video = torch.randint(0, 255, (4, 3, 2048, 2048), dtype=torch.uint8)
counts = []
for long_side in LONG_SIDES:
out = proc.preprocess(
videos=[video],
do_resize=True,
max_long_side_pixel=long_side,
return_tensors="pt",
)
counts.append(_image_tokens(out["video_grid_thw"][0]))
assert counts[0] < counts[1] < counts[2]
def test_video_volumetric_cap_raises():
proc = MiniMaxM3VLVideoProcessor()
# 400 frames at a 1008-long-side square: 1008*1008*400 >> 301,056,000.
video = torch.randint(0, 255, (400, 3, 2048, 2048), dtype=torch.uint8)
with pytest.raises(ValueError, match="max_total_pixels"):
proc.preprocess(
videos=[video],
do_resize=True,
max_long_side_pixel=1008,
return_tensors="pt",
)
tests/models/registry.py
+15
View File
@@ -420,6 +420,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"MiniMaxAI/MiniMax-M2",
trust_remote_code=True,
),
"MiniMaxM3SparseForCausalLM": _HfExamplesInfo(
"MiniMaxAI/MiniMax-M3",
trust_remote_code=True,
is_available_online=False,
),
"Ministral3ForCausalLM": _HfExamplesInfo("mistralai/Ministral-3-3B-Instruct-2512"),
"MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
"MistralLarge3ForCausalLM": _HfExamplesInfo(
@@ -1099,6 +1104,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"MiniMaxAI/MiniMax-VL-01",
trust_remote_code=True,
),
"MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo(
"MiniMaxAI/MiniMax-M3",
trust_remote_code=True,
is_available_online=False,
),
"Mistral3ForConditionalGeneration": _HfExamplesInfo(
"mistralai/Mistral-Small-3.1-24B-Instruct-2503",
extras={"fp8": "nm-testing/Mistral-Small-3.1-24B-Instruct-2503-FP8-dynamic"},
@@ -1601,6 +1611,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
speculative_model="XiaomiMiMo/MiMo-V2.5-Omni",
is_available_online=False,
),
"MiniMaxM3MTP": _HfExamplesInfo(
"MiniMaxAI/MiniMax-M3",
trust_remote_code=True,
is_available_online=False,
),
"NemotronHMTPModel": _HfExamplesInfo(
"nvidia/Nemotron-Super-Placeholder",
speculative_model="nvidia/Nemotron-Super-Placeholder",
tests/reasoning/test_minimax_m3_reasoning_parser.py
+320
View File
@@ -0,0 +1,320 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import string
from collections.abc import Sequence
import pytest
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.reasoning import ReasoningParserManager
from vllm.reasoning.minimax_m3_reasoning_parser import MiniMaxM3ReasoningParser
pytestmark = pytest.mark.skip_global_cleanup
class MiniMaxM3Tokenizer:
"""Small tokenizer with MiniMax M3 reasoning tags as special tokens."""
special_tokens = ("<mm:think>", "</mm:think>")
def __init__(self):
self._token_to_id: dict[str, int] = {}
self._id_to_token: dict[int, str] = {}
for token in self.special_tokens:
self._add_token(token)
for char in string.printable:
self._add_token(char)
def _add_token(self, token: str) -> int:
token_id = self._token_to_id.get(token)
if token_id is None:
token_id = len(self._token_to_id) + 1
self._token_to_id[token] = token_id
self._id_to_token[token_id] = token
return token_id
def get_vocab(self) -> dict[str, int]:
return dict(self._token_to_id)
def encode(
self,
text: str,
truncation: bool | None = None,
max_length: int | None = None,
add_special_tokens: bool = True,
) -> list[int]:
return [self._add_token(token) for token in self.tokenize(text)]
def decode(
self, ids: Sequence[int] | int, skip_special_tokens: bool = False
) -> str:
if isinstance(ids, int):
ids = [ids]
return "".join(self._id_to_token[token_id] for token_id in ids)
def tokenize(self, text: str) -> list[str]:
tokens: list[str] = []
pos = 0
while pos < len(text):
for special_token in self.special_tokens:
if text.startswith(special_token, pos):
tokens.append(special_token)
pos += len(special_token)
break
else:
tokens.append(text[pos])
pos += 1
return tokens
def convert_ids_to_tokens(
self,
ids: Sequence[int],
skip_special_tokens: bool = False,
) -> list[str]:
return [self._id_to_token[token_id] for token_id in ids]
def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:
if isinstance(tokens, str):
return self._add_token(tokens)
return [self._add_token(token) for token in tokens]
def convert_tokens_to_string(self, tokens: list[str]) -> str:
return "".join(tokens)
def make_parser(
chat_template_kwargs: dict[str, str] | None = None,
) -> tuple[MiniMaxM3ReasoningParser, MiniMaxM3Tokenizer]:
tokenizer = MiniMaxM3Tokenizer()
return (
MiniMaxM3ReasoningParser(tokenizer, chat_template_kwargs=chat_template_kwargs),
tokenizer,
)
def run_streaming(
parser: MiniMaxM3ReasoningParser,
tokenizer: MiniMaxM3Tokenizer,
chunks: list[str],
) -> tuple[str | None, str | None, list[bool]]:
previous_text = ""
previous_token_ids: list[int] = []
reasoning_parts: list[str] = []
content_parts: list[str] = []
reasoning_end_states: list[bool] = []
for chunk in chunks:
delta_token_ids = tokenizer.encode(chunk, add_special_tokens=False)
current_text = previous_text + chunk
current_token_ids = previous_token_ids + delta_token_ids
delta = parser.extract_reasoning_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=chunk,
previous_token_ids=previous_token_ids,
current_token_ids=current_token_ids,
delta_token_ids=delta_token_ids,
)
reasoning_end_states.append(
parser.is_reasoning_end_streaming(current_token_ids, delta_token_ids)
)
if delta is not None:
if delta.reasoning is not None:
reasoning_parts.append(delta.reasoning)
if delta.content is not None:
content_parts.append(delta.content)
previous_text = current_text
previous_token_ids = current_token_ids
return (
"".join(reasoning_parts) or None,
"".join(content_parts) or None,
reasoning_end_states,
)
def test_parser_registration():
parser_cls = ReasoningParserManager.get_reasoning_parser("minimax_m3")
assert parser_cls is MiniMaxM3ReasoningParser
def test_nonstreaming_extracts_explicit_reasoning_block():
parser, _ = make_parser()
request = ChatCompletionRequest(messages=[], model="test-model")
reasoning, content = parser.extract_reasoning(
"<mm:think>plan</mm:think>answer", request
)
assert reasoning == "plan"
assert content == "answer"
def test_nonstreaming_without_start_tag_is_content():
parser, _ = make_parser()
request = ChatCompletionRequest(messages=[], model="test-model")
reasoning, content = parser.extract_reasoning("plain answer", request)
assert reasoning is None
assert content == "plain answer"
def test_nonstreaming_drops_leading_end_tag():
parser, _ = make_parser()
request = ChatCompletionRequest(messages=[], model="test-model")
reasoning, content = parser.extract_reasoning("</mm:think>answer", request)
assert reasoning is None
assert content == "answer"
def test_nonstreaming_non_leading_end_tag_is_content():
parser, _ = make_parser()
request = ChatCompletionRequest(messages=[], model="test-model")
reasoning, content = parser.extract_reasoning("XXX</mm:think>YYY", request)
assert reasoning is None
assert content == "XXX</mm:think>YYY"
def test_nonstreaming_enabled_mode_starts_in_reasoning():
parser, _ = make_parser(chat_template_kwargs={"thinking_mode": "enabled"})
request = ChatCompletionRequest(messages=[], model="test-model")
reasoning, content = parser.extract_reasoning("plan</mm:think>answer", request)
assert reasoning == "plan"
assert content == "answer"
def test_nonstreaming_open_reasoning_block():
parser, _ = make_parser()
request = ChatCompletionRequest(messages=[], model="test-model")
reasoning, content = parser.extract_reasoning("<mm:think>still thinking", request)
assert reasoning == "still thinking"
assert content is None
def test_streaming_reasoning_tags_are_not_returned():
parser, tokenizer = make_parser()
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["<mm:think>", "plan", "</mm:think>", "answer"],
)
assert reasoning == "plan"
assert content == "answer"
assert end_states == [False, False, True, True]
def test_streaming_boundary_can_emit_reasoning_and_content():
parser, tokenizer = make_parser()
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["<mm:think>plan</mm:think>answer"],
)
assert reasoning == "plan"
assert content == "answer"
assert end_states == [True]
def test_streaming_drops_leading_end_tag():
parser, tokenizer = make_parser()
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["</mm:think>", "answer"],
)
assert reasoning is None
assert content == "answer"
assert end_states == [True, True]
def test_streaming_non_leading_end_tag_is_content():
parser, tokenizer = make_parser()
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["XXX</mm:think>YYY"],
)
assert reasoning is None
assert content == "XXX</mm:think>YYY"
assert end_states == [True]
def test_streaming_enabled_mode_starts_in_reasoning():
parser, tokenizer = make_parser(chat_template_kwargs={"thinking_mode": "enabled"})
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["plan", "</mm:think>", "answer"],
)
assert reasoning == "plan"
assert content == "answer"
assert end_states == [False, True, True]
def test_streaming_plain_content_ends_reasoning_phase():
parser, tokenizer = make_parser()
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["plain ", "answer"],
)
assert reasoning is None
assert content == "plain answer"
assert end_states == [True, True]
def test_token_id_helpers():
parser, tokenizer = make_parser()
output_ids = tokenizer.encode(
"<mm:think>abc</mm:think>def", add_special_tokens=False
)
open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False)
content_ids = tokenizer.encode("plain", add_special_tokens=False)
assert parser.is_reasoning_end(output_ids)
assert not parser.is_reasoning_end(open_reasoning_ids)
assert not parser.is_reasoning_end(content_ids)
assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
assert parser.extract_content_ids(open_reasoning_ids) == []
assert parser.extract_content_ids(content_ids) == content_ids
assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
def test_token_id_helpers_enabled_mode():
parser, tokenizer = make_parser(chat_template_kwargs={"thinking_mode": "enabled"})
output_ids = tokenizer.encode("abc</mm:think>def", add_special_tokens=False)
open_reasoning_ids = tokenizer.encode("abc", add_special_tokens=False)
assert parser.is_reasoning_end(output_ids)
assert not parser.is_reasoning_end(open_reasoning_ids)
assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
assert parser.extract_content_ids(open_reasoning_ids) == []
assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
assert parser.count_reasoning_tokens(open_reasoning_ids) == len(
tokenizer.encode("abc")
)
tests/tool_parsers/test_minimax_m3_tool_parser.py
+261
View File
@@ -0,0 +1,261 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from typing import Any
import pytest
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionToolsParam,
FunctionDefinition,
)
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.tool_parsers import ToolParserManager
from vllm.tool_parsers.minimax_m3_tool_parser import MinimaxM3ToolParser
pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
NS = "]<]minimax[>["
EOS_ID = 99
class FakeTokenizer:
"""Minimal fake tokenizer for unit tests."""
def __init__(self):
self.model_tokenizer = True
self.vocab: dict[str, int] = {}
def get_vocab(self) -> dict[str, int]:
return self.vocab
def sample_tools() -> list[ChatCompletionToolsParam]:
return [
ChatCompletionToolsParam(
function=FunctionDefinition(
name="create_order",
parameters={
"type": "object",
"properties": {
"user_id": {"type": "integer"},
"urgent": {"type": "boolean"},
"note": {"type": "string"},
"shipping": {
"type": "object",
"properties": {
"city": {"type": "string"},
"zip": {"type": "integer"},
},
},
"items": {
"type": "array",
"items": {
"type": "object",
"properties": {
"sku": {"type": "string"},
"qty": {"type": "integer"},
},
},
},
"metadata": {
"type": "object",
"additionalProperties": {"type": "string"},
},
"duplicate_demo": {"type": "object"},
},
},
),
)
]
@pytest.fixture
def parser() -> MinimaxM3ToolParser:
return MinimaxM3ToolParser(FakeTokenizer(), tools=sample_tools())
def build_order_call() -> str:
return (
f"{NS}<tool_call>\n"
f'{NS}<invoke name="create_order">'
f"{NS}<user_id>42{NS}</user_id>"
f"{NS}<urgent>true{NS}</urgent>"
f"{NS}<note>Please leave at front desk.{NS}</note>"
f"{NS}<shipping>"
f"{NS}<city>Singapore{NS}</city>"
f"{NS}<zip>018956{NS}</zip>"
f"{NS}</shipping>"
f"{NS}<items>"
f"{NS}<item>{NS}<sku>book-001{NS}</sku>{NS}<qty>2{NS}</qty>{NS}</item>"
f"{NS}<item>{NS}<sku>pen-007{NS}</sku>{NS}<qty>5{NS}</qty>{NS}</item>"
f"{NS}</items>"
f"{NS}<metadata>"
f"{NS}<source>mobile{NS}</source>"
f"{NS}<campaign>may-launch{NS}</campaign>"
f"{NS}</metadata>"
f"{NS}<duplicate_demo>"
f"{NS}<tag>a{NS}</tag>"
f"{NS}<tag>b{NS}</tag>"
f"{NS}</duplicate_demo>"
f"{NS}</invoke>\n"
f"{NS}</tool_call>"
)
def build_order_invocation(user_id: int) -> str:
return (
f'{NS}<invoke name="create_order">'
f"{NS}<user_id>{user_id}{NS}</user_id>"
f"{NS}</invoke>"
)
def build_multiple_order_call() -> str:
return (
f"{NS}<tool_call>\n"
f"{build_order_invocation(1)}\n"
f"{build_order_invocation(2)}\n"
f"{NS}</tool_call>"
)
def _feed(
parser: MinimaxM3ToolParser, chunks: list[str | tuple[str, list[int]]]
) -> list[DeltaMessage]:
previous = ""
results: list[DeltaMessage] = []
for chunk in chunks:
if isinstance(chunk, tuple):
delta, delta_ids = chunk
else:
delta = chunk
delta_ids = []
current = previous + delta
result = parser.extract_tool_calls_streaming(
previous_text=previous,
current_text=current,
delta_text=delta,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=delta_ids,
request=None,
)
if result is not None:
results.append(result)
previous = current
return results
def _collect_content(results: list[DeltaMessage]) -> str:
return "".join(result.content for result in results if result.content)
def _collect_tool_calls(results: list[DeltaMessage]) -> dict[int, dict[str, Any]]:
tool_calls: dict[int, dict[str, Any]] = {}
for result in results:
for tool_call in result.tool_calls or []:
tool_calls.setdefault(
tool_call.index,
{"id": None, "name": "", "arguments": ""},
)
if tool_call.id:
tool_calls[tool_call.index]["id"] = tool_call.id
if tool_call.function:
if tool_call.function.name:
tool_calls[tool_call.index]["name"] += tool_call.function.name
if tool_call.function.arguments:
tool_calls[tool_call.index]["arguments"] += (
tool_call.function.arguments
)
return tool_calls
def test_minimax_m3_parser_registered():
assert ToolParserManager.get_tool_parser("minimax_m3") is MinimaxM3ToolParser
def test_non_streaming_nested_tool_call(parser):
result = parser.extract_tool_calls(
"I will create it.\n" + build_order_call(),
request=None,
)
assert result.tools_called
assert result.content == "I will create it.\n"
assert len(result.tool_calls) == 1
tool_call = result.tool_calls[0]
assert tool_call.function.name == "create_order"
assert json.loads(tool_call.function.arguments) == {
"user_id": 42,
"urgent": True,
"note": "Please leave at front desk.",
"shipping": {"city": "Singapore", "zip": 18956},
"items": [
{"sku": "book-001", "qty": 2},
{"sku": "pen-007", "qty": 5},
],
"metadata": {
"source": "mobile",
"campaign": "may-launch",
},
"duplicate_demo": {"tag": ["a", "b"]},
}
def test_non_streaming_without_tool_call_keeps_content(parser):
result = parser.extract_tool_calls("plain response", request=None)
assert not result.tools_called
assert result.tool_calls == []
assert result.content == "plain response"
def test_non_streaming_multiple_tool_calls(parser):
result = parser.extract_tool_calls(build_multiple_order_call(), request=None)
assert result.tools_called
assert result.content is None
assert [tool_call.function.name for tool_call in result.tool_calls] == [
"create_order",
"create_order",
]
assert [
json.loads(tool_call.function.arguments)["user_id"]
for tool_call in result.tool_calls
] == [1, 2]
def test_streaming_without_tool_call_emits_text(parser):
results = _feed(parser, ["plain ", "response"])
assert _collect_content(results) == "plain response"
assert _collect_tool_calls(results) == {}
def test_streaming_nested_tool_call(parser):
tool_call_text = build_order_call()
results = _feed(
parser,
[
"I will create it.\n",
tool_call_text[:5],
tool_call_text[5:17],
tool_call_text[17:120],
tool_call_text[120:],
("", [EOS_ID]),
],
)
assert _collect_content(results) == "I will create it.\n"
tool_calls = _collect_tool_calls(results)
assert len(tool_calls) == 1
assert tool_calls[0]["name"] == "create_order"
assert tool_calls[0]["id"] is not None
assert json.loads(tool_calls[0]["arguments"]) == json.loads(
parser.streamed_args_for_tool[0]
)
assert json.loads(parser.prev_tool_call_arr[0]["arguments"])["items"][1]["qty"] == 5
assert results[-1].content is None
tools/pre_commit/generate_attention_backend_docs.py
+36 -4
View File
@@ -30,6 +30,7 @@ REPO_ROOT = Path(__file__).parent.parent.parent
RELEVANT_PATTERNS = [
"vllm/v1/attention/backends/*.py",
"vllm/v1/attention/backends/**/*.py",
"vllm/models/minimax_m3/common/sparse_attention.py",
"vllm/model_executor/layers/attention/mla_attention.py",
"vllm/platforms/cuda.py",
"tools/pre_commit/generate_attention_backend_docs.py",
@@ -1633,6 +1634,24 @@ def generate_mla_section(
return "\n".join(lines)
def generate_minimax_section(backends: list[dict[str, Any]]) -> str:
"""Generate the MiniMax M3 sparse attention section."""
lines = [
"## MiniMax M3 Sparse Attention Backends",
"",
'Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")',
"layers. It is wired in directly by the model and is not part of the",
"automatic priority lists above. A lightning indexer scores KV blocks, the",
"top-k blocks (plus fixed init/local blocks) are selected, and attention",
"attends only to those blocks; index keys live in a separate side cache.",
"",
]
columns = _build_columns(is_mla=False, has_versions=False)
lines.extend(_render_table(columns, backends))
lines.append("")
return "\n".join(lines)
# ---------------------------------------------------------------------------
# Top-level orchestration
# ---------------------------------------------------------------------------
@@ -1669,15 +1688,24 @@ def generate_docs() -> str:
if fi_features:
all_backends = _expand_flashinfer_variants(all_backends, fi_features)
# DeepSeek V4 (*_DSV4) decode backends get their own subsection rather than
# mixing into the main MLA / standard tables (the ROCm V4 backend isn't
# flagged is_mla by the AST heuristic, so filter purely on the name).
# DeepSeek V4 (*_DSV4) decode backends and MiniMax M3 sparse backends each
# get their own subsection rather than mixing into the main MLA / standard
# tables (the ROCm V4 backend isn't flagged is_mla by the AST heuristic, so
# filter purely on the name).
def _is_v4(b: dict[str, Any]) -> bool:
return b["name"].endswith("_DSV4")
def _is_minimax(b: dict[str, Any]) -> bool:
return not b["is_mla"] and not _is_v4(b) and b["name"].startswith("MINIMAX")
v4_decode_backends = [b for b in all_backends if _is_v4(b)]
minimax_backends = [b for b in all_backends if _is_minimax(b)]
mla_backends = [b for b in all_backends if b["is_mla"] and not _is_v4(b)]
non_mla_backends = [b for b in all_backends if not b["is_mla"] and not _is_v4(b)]
non_mla_backends = [
b
for b in all_backends
if not b["is_mla"] and not _is_v4(b) and not _is_minimax(b)
]
# Generate documentation
script_path = "tools/pre_commit/generate_attention_backend_docs.py"
@@ -1726,6 +1754,10 @@ def generate_docs() -> str:
if footnotes:
doc_lines.append("\n>\n".join(footnotes) + "\n")
# Add MiniMax M3 sparse section (separate category after standard GQA)
if minimax_backends:
doc_lines.append(generate_minimax_section(minimax_backends))
# Add MLA section with prefill and decode backends
doc_lines.append(
generate_mla_section(mla_prefill_backends, mla_backends, v4_decode_backends)
vllm/_custom_ops.py
+64
View File
@@ -2614,6 +2614,70 @@ def reshape_and_cache_flash(
)
def fused_minimax_m3_qknorm_rope_kv_insert(
qkv: torch.Tensor,
q_norm_weight: torch.Tensor,
k_norm_weight: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
num_heads: int,
num_kv_heads: int,
rotary_dim: int,
eps: float,
index_q_norm_weight: torch.Tensor | None = None,
index_k_norm_weight: torch.Tensor | None = None,
num_index_heads: int = 0,
slot_mapping: torch.Tensor | None = None,
index_slot_mapping: torch.Tensor | None = None,
kv_cache: torch.Tensor | None = None,
index_cache: torch.Tensor | None = None,
block_size: int = 0,
q_out: torch.Tensor | None = None,
index_q_out: torch.Tensor | None = None,
) -> None:
"""Fused MiniMax-M3 attention pre-processing (in-place).
Applies Gemma RMSNorm + partial NeoX RoPE to ``qkv`` in place. ``qkv`` is a
single fused tensor:
- dense layer (``num_index_heads == 0``): ``[q | k | v]``;
- sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q |
index_k]`` — the index branch is read straight out of ``qkv``.
When ``kv_cache`` is given (sparse serving), also scatter-inserts the
normed/roped k & v into the paged bf16 KV cache by ``slot_mapping`` and the
index key into ``index_cache`` by ``index_slot_mapping``. If
``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches.
If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N,
niq*128]``) are given, the normed/roped q / index_q are written there
instead of in place — folding the de-interleave into this kernel's store so
callers skip a separate ``.contiguous()`` copy before the SM100 sparse
attention's flat TMA descriptor.
"""
torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert(
qkv,
q_norm_weight,
k_norm_weight,
cos_sin_cache,
positions,
num_heads,
num_kv_heads,
rotary_dim,
eps,
index_q_norm_weight,
index_k_norm_weight,
num_index_heads,
slot_mapping,
index_slot_mapping,
kv_cache,
index_cache,
block_size,
q_out,
index_q_out,
)
def concat_and_cache_mla(
kv_c: torch.Tensor,
k_pe: torch.Tensor,
vllm/config/attention.py
+6
View File
@@ -9,6 +9,8 @@ from vllm.config.utils import config
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
from vllm.v1.attention.backends.registry import AttentionBackendEnum
IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"]
@config
class AttentionConfig:
@@ -50,6 +52,10 @@ class AttentionConfig:
use_fp4_indexer_cache: bool = False
"""If set, use fp4 indexer cache for dsv32 family model (not support yet)"""
indexer_kv_dtype: IndexerKVDType = "bf16"
"""Data type for the sparse-attention indexer K cache. Quantized formats
(fp8, mxfp4, nvfp4) require indexer kernel support in the backend."""
use_non_causal: bool = False
"""Whether to use non-causal (bidirectional) attention."""
vllm/config/speculative.py
+31
View File
@@ -45,6 +45,7 @@ MTPModelTypes = Literal[
"qwen3_next_mtp",
"qwen3_5_mtp",
"longcat_flash_mtp",
"minimax_m3_mtp",
"mtp",
"pangu_ultra_moe_mtp",
"step3p5_mtp",
@@ -528,6 +529,36 @@ class SpeculativeConfig:
text_config.num_kv_shared_layers = 0
hf_config.update({"n_predict": 1, "architectures": ["Gemma4MTPModel"]})
if (
hf_config.model_type == "minimax_m3_vl"
or initial_architecture == "MiniMaxM3SparseForConditionalGeneration"
):
# MTP modules live on the language model of this VL checkpoint, so
# promote text_config before rewriting it into an MTP config.
quantization_config = getattr(hf_config, "quantization_config", None)
hf_config = getattr(hf_config, "text_config", hf_config)
if (
quantization_config is not None
and getattr(hf_config, "quantization_config", None) is None
):
hf_config.update({"quantization_config": quantization_config})
hf_config.model_type = "minimax_m3_mtp"
n_predict = getattr(hf_config, "num_mtp_modules", 1)
hf_config.update(
{"n_predict": n_predict, "architectures": ["MiniMaxM3MTP"]}
)
elif (
hf_config.model_type == "minimax_m3_mtp"
or initial_architecture == "MiniMaxM3MTP"
):
# Standalone MTP checkpoints already use a flat MTP config with no
# VL wrapper / text_config to promote, so just normalize the
# architecture and derive n_predict from num_mtp_modules.
n_predict = getattr(hf_config, "num_mtp_modules", 1)
hf_config.update(
{"n_predict": n_predict, "architectures": ["MiniMaxM3MTP"]}
)
return hf_config
def __post_init__(self):
vllm/config/vllm.py
+9 -3
View File
@@ -1101,20 +1101,26 @@ class VllmConfig:
)
self.compilation_config.mode = CompilationMode.NONE
# DeepSeek V4's model classes don't carry @support_torch_compile —
# For model classes don't carry @support_torch_compile —
# the breakable cudagraph is the supported PIECEWISE path. Auto-enable
# it unless the user has explicitly opted out via the env var.
if (
self.model_config is not None
and "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
and any(
a in ("DeepseekV4ForCausalLM", "DeepSeekV4MTPModel")
a
in (
"DeepseekV4ForCausalLM",
"DeepSeekV4MTPModel",
"MiniMaxM3SparseForCausalLM",
"MiniMaxM3SparseForConditionalGeneration",
)
for a in self.model_config.architectures
)
):
os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
logger.info_once(
"Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1 for DeepSeek V4. "
"Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
"Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
)
vllm/envs.py
+10
View File
@@ -113,6 +113,7 @@ if TYPE_CHECKING:
VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True
VLLM_DISABLE_PYNCCL: bool = False
VLLM_USE_OINK_OPS: bool = False
VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True
VLLM_ROCM_USE_AITER: bool = False
VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
VLLM_ROCM_USE_AITER_LINEAR: bool = True
@@ -1092,6 +1093,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
),
# Disable aiter ops unless specifically enabled.
# Acts as a parent switch to enable the rest of the other operations.
# On hardware without a native MXFP8 kernel (e.g. ROCm gfx942 / MI300), the
# MXFP8 emulation path dequantizes weights MXFP8->BF16 once at load time and
# runs as a BF16 checkpoint (no per-step dequant). Set to 0 to fall back to
# per-step dequant: keeps the 1-byte MXFP8 weights (~half the weight memory)
# at the cost of dequantizing every forward step (much slower). Default on.
"VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD": lambda: (
os.getenv("VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD", "True").lower()
in ("true", "1")
),
"VLLM_ROCM_USE_AITER": lambda: (
os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1")
),
vllm/model_executor/kernels/linear/__init__.py
+6
View File
@@ -93,6 +93,9 @@ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
from vllm.model_executor.kernels.linear.mxfp8.marlin import (
MarlinMxfp8LinearKernel,
)
from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
RocmDotScaledMxfp8LinearKernel,
)
from vllm.model_executor.kernels.linear.mxfp8.xpu import (
XPUMxFp8LinearKernel,
)
@@ -383,6 +386,9 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
EmulationMxfp8LinearKernel,
],
PlatformEnum.ROCM: [
# Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and
# falls through to BF16 emulation (hipBLASLt) elsewhere / on regression.
RocmDotScaledMxfp8LinearKernel,
EmulationMxfp8LinearKernel,
],
PlatformEnum.XPU: [
vllm/model_executor/kernels/linear/mxfp8/emulation.py
+25 -1
View File
@@ -33,6 +33,17 @@ class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
# Dequantize MXFP8 -> BF16 ONCE here, at load time, so apply_weights runs
# a plain BF16 linear with no per-step dequant -- i.e. run as if from a
# BF16 checkpoint. The 1-byte MXFP8 weight is replaced by BF16 (2x its
# size, but linear weights are small vs the MoE experts); the tiny E8M0
# scale is kept for the dtype/ndim asserts but is otherwise unused.
# Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the MXFP8
# weight and dequant per-step in apply_weights instead.
import vllm.envs as envs
if envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD:
weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
layer.weight = Parameter(weight.contiguous(), requires_grad=False)
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
@@ -42,6 +53,17 @@ class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
weight = layer.weight
# Load-time dequant path: weights are already BF16/FP16 (>= 2-byte), so
# run a plain linear -- no per-step dequant. (MXFP8 weights are 1-byte.)
if weight.element_size() >= 2:
# F.linear requires x and weight share a dtype; .to() is a no-op when
# they already match (e.g. both BF16).
output = torch.nn.functional.linear(x, weight.to(x.dtype), bias)
return output.to(x.dtype)
# Fallback: weights still in MXFP8 -- dequant on the fly (other archs /
# if a future caller skips the load-time conversion above).
weight_scale = layer.weight_scale
if weight_scale.dtype != MXFP8_SCALE_DTYPE:
raise ValueError(
@@ -55,6 +77,8 @@ class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
f"Ensure process_weights_after_loading was called."
)
weight_bf16 = dequant_mxfp8_to_bf16(layer.weight, weight_scale)
# Cast to x's dtype: dequant yields BF16, but F.linear needs both operands
# to match (e.g. an FP16 model). No-op when x is already BF16.
weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale).to(x.dtype)
output = torch.nn.functional.linear(x, weight_bf16, bias)
return output.to(x.dtype)
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
-10
View File
@@ -56,8 +56,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
input_shape = x.shape
input_2d = x.view(-1, K)
M_orig = input_2d.shape[0]
min_dim = 128
assert min_dim <= K, (
@@ -72,11 +70,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
f"out_features is too small for mm_mxfp8."
)
M_padded = ((M_orig + min_dim - 1) // min_dim) * min_dim
if M_padded != M_orig:
pad_rows = M_padded - M_orig
input_2d = torch.nn.functional.pad(input_2d, (0, 0, 0, pad_rows))
input_mxfp8, input_scale = mxfp8_e4m3_quantize(
input_2d, is_sf_swizzled_layout=True
)
@@ -93,9 +86,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
backend="cutlass",
)
if M_padded != M_orig:
output = output[:M_orig, :]
if bias is not None:
output = output + bias
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+171
View File
@@ -0,0 +1,171 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``.
Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
matrix cores. Falls back (via the kernel selector) to the BF16
``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
``K % 128 != 0``.
"""
import torch
from torch.nn.parameter import Parameter
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
MXFP8_BLOCK_SIZE,
MXFP8_SCALE_DTYPE,
dequant_mxfp8_to_bf16,
mxfp8_e4m3_quantize,
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
@triton.jit
def _mxfp8_linear_kernel(
x_ptr,
xs_ptr,
w_ptr,
ws_ptr,
out_ptr,
M,
N,
K,
stride_xm,
stride_xk,
stride_xsm,
stride_xsk,
stride_wn,
stride_wk,
stride_wsn,
stride_wsk,
stride_om,
stride_on,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
offs_sk = tl.arange(0, BLOCK_K // 32)
m_mask = offs_m < M
n_mask = offs_n < N
x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk
xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk
w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for _ in range(0, tl.cdiv(K, BLOCK_K)):
x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3")
x_ptrs += BLOCK_K * stride_xk
w_ptrs += BLOCK_K * stride_wk
xs_ptrs += (BLOCK_K // 32) * stride_xsk
ws_ptrs += (BLOCK_K // 32) * stride_wsk
o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
tl.store(
o_ptrs, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :]
)
def _mxfp8_dot_scaled_linear(
x: torch.Tensor, # [M, K] bf16/fp16
w: torch.Tensor, # [N, K] fp8 e4m3
w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0)
) -> torch.Tensor:
M, K = x.shape
N = w.shape[0]
x_q, x_scale = mxfp8_e4m3_quantize(x)
out = torch.empty((M, N), dtype=x.dtype, device=x.device)
BLOCK_M, BLOCK_N, BLOCK_K = 64, 128, 128
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
_mxfp8_linear_kernel[grid](
x_q,
x_scale,
w,
w_scale,
out,
M,
N,
K,
x_q.stride(0),
x_q.stride(1),
x_scale.stride(0),
x_scale.stride(1),
w.stride(0),
w.stride(1),
w_scale.stride(0),
w_scale.stride(1),
out.stride(0),
out.stride(1),
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
BLOCK_K=BLOCK_K,
num_warps=8,
)
return out
class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
"""Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_rocm():
return False, "not ROCm"
# supports_mx() == gfx95x (CDNA4 native microscaling hardware). On other
# archs dot_scaled would upcast to BF16, so the kernel selector falls
# through to the BF16 emulation (hipBLASLt) path instead.
if not current_platform.supports_mx():
return False, "native MX requires CDNA4 (gfx95x)"
return True, None
@classmethod
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight = layer.weight.data # [N, K] fp8
N, K = weight.shape
scale_k = K // MXFP8_BLOCK_SIZE
weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
layer.weight = Parameter(weight.contiguous(), requires_grad=False)
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
raise ValueError(
f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got "
f"{layer.weight_scale.dtype}."
)
out_shape = (*x.shape[:-1], layer.weight.shape[0])
x2d = x.reshape(-1, x.shape[-1])
if x2d.shape[-1] % 128 == 0:
out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
else:
# dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
out = out.reshape(out_shape)
if bias is not None:
out = out + bias
return out
vllm/model_executor/layers/activation.py
+22 -5
View File
@@ -158,17 +158,28 @@ class SiluAndMulWithClamp(CustomOp):
Computes:
gate = clamp(x[..., :d], max=swiglu_limit)
up = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
out = silu(gate) * up
where d = x.shape[-1] // 2.
out = gate * sigmoid(alpha * gate) * (up + beta)
where d = x.shape[-1] // 2. The defaults alpha=1.0, beta=0.0 reduce this to
``silu(gate) * up``; SwiGLU-OAI style models pass alpha (sigmoid scale) and
beta=1.0 (up bias).
Shapes:
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
return: (num_tokens, d) or (batch_size, seq_len, d)
"""
def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
def __init__(
self,
swiglu_limit: float,
alpha: float = 1.0,
beta: float = 0.0,
*,
compile_native: bool = True,
):
super().__init__(compile_native=compile_native)
self.swiglu_limit = float(swiglu_limit)
self.alpha = float(alpha)
self.beta = float(beta)
if current_platform.is_rocm() or current_platform.is_xpu():
self._forward_method = self.forward_native
elif current_platform.is_cuda_alike():
@@ -180,18 +191,24 @@ class SiluAndMulWithClamp(CustomOp):
d = x.shape[-1] // 2
gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
return F.silu(gate) * up
return gate * torch.sigmoid(self.alpha * gate) * (up + self.beta)
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
d = x.shape[-1] // 2
output_shape = x.shape[:-1] + (d,)
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
self.op(out, x, self.swiglu_limit)
self.op(out, x, self.swiglu_limit, self.alpha, self.beta)
return out
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
return self.forward_native(x)
def extra_repr(self) -> str:
return (
f"swiglu_limit={self.swiglu_limit!r}, "
f"alpha={self.alpha!r}, beta={self.beta!r}"
)
# --8<-- [start:mul_and_silu]
@CustomOp.register("mul_and_silu")
vllm/model_executor/layers/attention/attention.py
+2
View File
@@ -7,6 +7,7 @@ import torch
import torch.nn as nn
import vllm.envs as envs
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
from vllm.config import CacheConfig, get_current_vllm_config
from vllm.config.vllm import VllmConfig
from vllm.forward_context import ForwardContext, get_forward_context
@@ -730,6 +731,7 @@ direct_register_custom_op(
)
@eager_break_during_capture
@maybe_transfer_kv_layer
def unified_attention_with_output(
query: torch.Tensor,
vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
+143
View File
@@ -0,0 +1,143 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Manual fusion of tensor-parallel all-reduce with the following GemmaRMSNorm.
Under tensor parallelism a ``RowParallelLinear`` (e.g. attention ``o_proj``)
produces a per-rank partial sum that is all-reduced, and the result is then fed
into a ``GemmaRMSNorm`` that adds the residual and normalizes. flashinfer ships a
kernel that fuses all-reduce + residual-add + RMSNorm into a single launch; this
helper drives it directly (no torch.compile pass) for models that run eager.
Scope: attention output only, no quantization. When the flashinfer fast path is
not applicable (TP==1, flashinfer/NVSwitch unavailable, unsupported dtype, or an
oversize batch) it falls back to ``all_reduce`` + ``GemmaRMSNorm``, which is
numerically identical to the unfused model path.
"""
import torch
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
from vllm.distributed.parallel_state import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
get_tp_group,
)
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
MiB = 1024 * 1024
# flashinfer fused all-reduce + RMSNorm is wired as a registered custom op in
# allreduce_rms_fusion; both that op and the workspace helpers only exist when
# flashinfer.comm.allreduce_fusion is importable.
try:
from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
flashinfer_trtllm_fused_allreduce_norm,
)
from vllm.distributed.device_communicators.flashinfer_all_reduce import (
flashinfer_comm,
get_fi_ar_workspace,
)
_AR_RESIDUAL_RMS_NORM = (
flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm
if flashinfer_comm is not None
else None
)
except ImportError:
flashinfer_trtllm_fused_allreduce_norm = None # type: ignore[assignment]
get_fi_ar_workspace = None # type: ignore[assignment]
_AR_RESIDUAL_RMS_NORM = None
_FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16)
def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
"""Workspace token budget for flashinfer fused all-reduce, or None if the
current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
from vllm.config.compilation import PassConfig
max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
if not max_size_mb:
return None
element_size = torch.tensor([], dtype=dtype).element_size()
return int(max_size_mb * MiB) // (hidden_size * element_size)
def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool, int]:
"""Whether the flashinfer fused path applies; returns (ok, max_token_num)."""
if (
flashinfer_trtllm_fused_allreduce_norm is None
or get_fi_ar_workspace is None
or _AR_RESIDUAL_RMS_NORM is None
):
return False, 0
if (
not hidden_states.is_cuda
or hidden_states.dim() != 2
or not hidden_states.is_contiguous()
or hidden_states.dtype not in _FI_SUPPORTED_DTYPES
):
return False, 0
num_tokens, hidden_size = hidden_states.shape
max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype)
if max_token_num is None or num_tokens > max_token_num:
return False, 0
# Lazily create / fetch the (globally cached) workspace; returns None on
# GPUs without NVSwitch, in which case we fall back gracefully.
workspace = get_fi_ar_workspace(
world_size=tp_size,
rank=get_tensor_model_parallel_rank(),
max_token_num=max_token_num,
hidden_dim=hidden_size,
dtype=hidden_states.dtype,
group=get_tp_group().device_group,
)
if workspace is None:
return False, 0
return True, max_token_num
def fused_allreduce_gemma_rms_norm(
hidden_states: torch.Tensor,
residual: torch.Tensor,
norm: GemmaRMSNorm,
) -> tuple[torch.Tensor, torch.Tensor]:
"""All-reduce ``hidden_states`` + add ``residual`` + GemmaRMSNorm, fused.
``hidden_states`` is the per-rank *partial* (un-reduced) output of a
row-parallel linear; ``norm`` is the GemmaRMSNorm applied right after.
Returns ``(normed_output, new_residual)``, equivalent to
``norm(all_reduce(hidden_states), residual)``.
"""
tp_size = get_tensor_model_parallel_world_size()
if tp_size == 1:
# No all-reduce needed; identical to the unfused path.
return norm(hidden_states, residual)
ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
if ok:
norm_out = torch.empty_like(hidden_states)
# With norm_out provided, the kernel writes the new residual
# (all_reduce(hidden_states) + residual) into the hidden_states buffer
# and the normalized result into norm_out, leaving `residual` untouched.
flashinfer_trtllm_fused_allreduce_norm(
allreduce_in=hidden_states,
residual=residual,
rms_gamma=norm.weight,
rms_eps=norm.variance_epsilon,
world_size=tp_size,
weight_bias=1.0, # GemmaRMSNorm-style
launch_with_pdl=True,
fp32_acc=True,
max_token_num=max_token_num,
pattern_code=_AR_RESIDUAL_RMS_NORM,
norm_out=norm_out,
)
return norm_out, hidden_states
# Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model).
reduced = tensor_model_parallel_all_reduce(hidden_states)
return norm(reduced, residual)
vllm/model_executor/layers/fused_moe/activation.py
+25 -2
View File
@@ -17,7 +17,12 @@ class MoEActivation(Enum):
GELU = "gelu"
GELU_TANH = "gelu_tanh"
RELU2 = "relu2"
# SWIGLUOAI expects gate/up *interleaved* in w13 ([gate0, up0, gate1, ...]),
# as in gpt-oss checkpoints. SWIGLUOAI_UNINTERLEAVE has identical math but
# expects the *packed* layout ([all gates; all ups]), as produced by a
# MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3).
SWIGLUOAI = "swigluoai"
SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave"
SWIGLUSTEP = "swiglustep"
# Non-gated activations (no mul with gate) expect input of shape [..., d]
@@ -73,6 +78,7 @@ _CUSTOM_OP_NAMES: dict[MoEActivation, str] = {
MoEActivation.GELU: "gelu_and_mul",
MoEActivation.GELU_TANH: "gelu_tanh_and_mul",
MoEActivation.SWIGLUOAI: "swigluoai_and_mul",
MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp",
MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",
MoEActivation.RELU2: "relu2",
MoEActivation.SILU_NO_MUL: "silu_and_mul",
@@ -105,8 +111,17 @@ def apply_moe_activation(
activation: MoEActivation,
output: torch.Tensor,
input: torch.Tensor,
*,
clamp_limit: float | None = None,
alpha: float = 1.0,
beta: float = 0.0,
) -> torch.Tensor:
"""Apply MoE activation function."""
"""Apply MoE activation function.
``clamp_limit``/``alpha``/``beta`` (from the quant config) drive the clamped
SwiGLU kernels: ``SILU`` + ``clamp_limit`` and ``SWIGLUOAI_UNINTERLEAVE`` both
map to ``silu_and_mul_with_clamp``. Other activations ignore them.
"""
assert input.dim() == 2, "Input must be 2D"
assert output.dim() == 2, "Output must be 2D"
if activation.is_gated:
@@ -122,13 +137,21 @@ def apply_moe_activation(
# Activations with gated multiplication (gate × activation(up))
if activation == MoEActivation.SILU:
torch.ops._C.silu_and_mul(output, input)
if clamp_limit is not None:
# Fused silu(clamp(gate)) * clamp(up); equivalent to swiglu_limit_func.
torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, 1.0, 0.0)
else:
torch.ops._C.silu_and_mul(output, input)
elif activation == MoEActivation.GELU:
torch.ops._C.gelu_and_mul(output, input)
elif activation == MoEActivation.GELU_TANH:
torch.ops._C.gelu_tanh_and_mul(output, input)
elif activation == MoEActivation.SWIGLUOAI:
torch.ops._C.swigluoai_and_mul(output, input)
elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
# SwiGLU-OAI on packed w13 (gate = first half, up = second half).
assert clamp_limit is not None, "SWIGLUOAI_UNINTERLEAVE requires clamp_limit"
torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, alpha, beta)
elif activation == MoEActivation.SWIGLUSTEP:
from vllm.model_executor.layers.activation import swiglustep_and_mul_triton
vllm/model_executor/layers/fused_moe/config.py
+17
View File
@@ -900,6 +900,9 @@ def fp8_w8a16_moe_quant_config(
w1_bias: torch.Tensor | None = None,
w2_bias: torch.Tensor | None = None,
block_shape: list[int] | None = None,
gemm1_alpha: float | None = None,
gemm1_beta: float | None = None,
gemm1_clamp_limit: float | None = None,
) -> FusedMoEQuantConfig:
"""
Construct a quant config for 16-bit float activations and fp8 weights.
@@ -925,6 +928,9 @@ def fp8_w8a16_moe_quant_config(
None,
w2_bias,
),
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=gemm1_clamp_limit,
)
@@ -979,15 +985,24 @@ def int4_w4afp8_moe_quant_config(
def biased_moe_quant_config(
w1_bias: torch.Tensor | None,
w2_bias: torch.Tensor | None,
gemm1_alpha: float | None = None,
gemm1_beta: float | None = None,
gemm1_clamp_limit: float | None = None,
) -> FusedMoEQuantConfig:
"""
Construct a quant config for unquantized activations with biases.
gemm1_alpha/gemm1_beta/gemm1_clamp_limit carry the SwiGLU gate params
through to the fused activation kernel (e.g. swigluoai_uninterleave).
"""
return FusedMoEQuantConfig(
_a1=FusedMoEQuantDesc(),
_a2=FusedMoEQuantDesc(),
_w1=FusedMoEQuantDesc(bias=w1_bias),
_w2=FusedMoEQuantDesc(bias=w2_bias),
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=gemm1_clamp_limit,
)
@@ -1268,6 +1283,8 @@ class FusedMoEConfig:
# are filtered out by `FusedMoEExperts.is_supported_config` so the oracle
# cannot silently select one and drop the clamp.
swiglu_limit: float | None = None
swiglu_alpha: float | None = None
swiglu_beta: float | None = None
max_capture_size: int = 0
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+86 -14
View File
@@ -130,6 +130,9 @@ def _fwd_kernel_ep_scatter_2(
HIDDEN_SIZE_PAD: tl.constexpr,
SCALE_HIDDEN_SIZE: tl.constexpr,
SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
PACK_UE8M0: tl.constexpr,
SCALE_PACKED_SIZE: tl.constexpr,
SCALE_PACKED_SIZE_PAD: tl.constexpr,
):
start_token_id = tl.program_id(0)
grid_num = tl.num_programs(0)
@@ -137,16 +140,47 @@ def _fwd_kernel_ep_scatter_2(
offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
mask = offset_in < HIDDEN_SIZE
offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
mask_s = offset_in_s < SCALE_HIDDEN_SIZE
output_tensor_stride0 = output_tensor_stride0.to(tl.int64)
if PACK_UE8M0:
# One int32 per 4 consecutive 32-wide UE8M0 groups, stored MN-major.
offs_pk = tl.arange(0, SCALE_PACKED_SIZE_PAD)
mask_pk = offs_pk < SCALE_PACKED_SIZE
else:
offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
mask_s = offset_in_s < SCALE_HIDDEN_SIZE
for token_id in range(start_token_id, total_token_num, grid_num):
to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
to_copy_s = tl.load(
recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s
)
if PACK_UE8M0:
# Pack 4 UE8M0 bytes into one int32 (byte j = group 4*pk+j).
base_s = recv_x_scale + token_id * recv_x_scale_stride0
g0, g1 = offs_pk * 4, offs_pk * 4 + 1
g2, g3 = offs_pk * 4 + 2, offs_pk * 4 + 3
b0 = tl.load(
base_s + g0 * recv_x_scale_stride1, mask=g0 < SCALE_HIDDEN_SIZE
)
b1 = tl.load(
base_s + g1 * recv_x_scale_stride1, mask=g1 < SCALE_HIDDEN_SIZE
)
b2 = tl.load(
base_s + g2 * recv_x_scale_stride1, mask=g2 < SCALE_HIDDEN_SIZE
)
b3 = tl.load(
base_s + g3 * recv_x_scale_stride1, mask=g3 < SCALE_HIDDEN_SIZE
)
packed_s = (
b0.to(tl.int32)
| (b1.to(tl.int32) << 8)
| (b2.to(tl.int32) << 16)
| (b3.to(tl.int32) << 24)
)
else:
to_copy_s = tl.load(
recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s,
mask=mask_s,
)
for topk_index in tl.range(0, topk_num, 1, num_stages=4):
expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
@@ -164,11 +198,21 @@ def _fwd_kernel_ep_scatter_2(
output_tensor_ptr = (
output_tensor + dest_token_index_i64 * output_tensor_stride0
)
tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
output_tensor_scale_ptr = (
output_tensor_scale + dest_token_index * output_tensor_scale_stride0
)
tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s)
if PACK_UE8M0:
tl.store(
output_tensor_scale_ptr + offs_pk * output_tensor_scale_stride1,
packed_s,
mask=mask_pk,
)
else:
tl.store(
output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s
)
@torch.no_grad()
@@ -183,9 +227,11 @@ def ep_scatter(
output_tensor_scale: torch.Tensor,
m_indices: torch.Tensor,
output_index: torch.Tensor,
block_size: int = 128,
pack_ue8m0: bool = False,
):
BLOCK_E = 128 # token num of per expert is aligned to 128
BLOCK_D = 128 # block size of quantization
BLOCK_D = block_size # block size of activation-scale quantization
num_warps = 8
num_experts = num_recv_tokens_per_expert.shape[0]
hidden_size = recv_x.shape[1]
@@ -195,6 +241,10 @@ def ep_scatter(
assert m_indices.shape[0] % BLOCK_E == 0
assert expert_start_loc.shape[0] == num_experts
# pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
scale_hidden_size = hidden_size // BLOCK_D
scale_packed_size = (scale_hidden_size + 3) // 4 if pack_ue8m0 else 1
_fwd_kernel_ep_scatter_1[(grid,)](
num_recv_tokens_per_expert,
expert_start_loc,
@@ -234,8 +284,11 @@ def ep_scatter(
num_warps=num_warps,
HIDDEN_SIZE=hidden_size,
HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D,
SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D),
SCALE_HIDDEN_SIZE=scale_hidden_size,
SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size),
PACK_UE8M0=pack_ue8m0,
SCALE_PACKED_SIZE=scale_packed_size,
SCALE_PACKED_SIZE_PAD=triton.next_power_of_2(scale_packed_size),
)
return
@@ -352,6 +405,7 @@ def deepgemm_moe_permute(
expert_map: torch.Tensor | None,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
aq_out: torch.Tensor | None = None,
block_size: int | None = None,
):
assert aq.ndim == 2
assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
@@ -359,6 +413,10 @@ def deepgemm_moe_permute(
device = aq.device
block_m, block_k = get_mk_alignment_for_contiguous_layout()
# The activation-scale group size may differ from the M/K tile alignment
# (e.g. MXFP8 uses a 32-element scale group while block_k stays 128).
if block_size is not None:
block_k = block_size
M_sum = compute_aligned_M(
M=topk_ids.size(0),
@@ -376,9 +434,21 @@ def deepgemm_moe_permute(
if aq_out is None:
aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype)
aq_scale_out = torch.empty(
(M_sum, H // block_k), device=device, dtype=torch.float32
)
# uint8 UE8M0 (MXFP8) -> scatter packs into DeepGEMM's int32 MN-major
# TMA-aligned layout; float32 (FP8/FP4) scattered row-major as-is.
pack_ue8m0 = aq_scale.dtype == torch.uint8
sf_k = H // block_k
if pack_ue8m0:
packed_sf_k = (sf_k + 3) // 4
tma_aligned_mn = round_up(M_sum, 4)
aq_scale_out = torch.empty_strided(
(M_sum, packed_sf_k),
(1, tma_aligned_mn),
device=device,
dtype=torch.int32,
)
else:
aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32)
# DeepGEMM uses negative values in m_indices (here expert_ids) to mark
# completely invalid / padded blocks that should be skipped. We always
@@ -412,6 +482,8 @@ def deepgemm_moe_permute(
output_tensor_scale=aq_scale_out,
m_indices=expert_ids,
output_index=inv_perm,
block_size=block_k,
pack_ue8m0=pack_ue8m0,
)
return aq_out, aq_scale_out, expert_ids, inv_perm
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+68 -14
View File
@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8Dynamic128Sym,
kFp8Static128BlockSym,
kMxfp4Static,
kMxfp8Dynamic,
kMxfp8Static,
)
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import (
DeepGemmQuantScaleFMT,
get_mk_alignment_for_contiguous_layout,
@@ -123,12 +126,26 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig):
super().__init__(moe_config=moe_config, quant_config=quant_config)
assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()
assert quant_config.quant_dtype == torch.float8_e4m3fn
# MXFP8: FP8 e4m3 values + UE8M0 1x32 block scales (Blackwell). Reuses
# the same grouped GEMM (aliased to fp8_fp4) with recipe (1, 32).
self.mxfp8 = quant_config.block_shape == [1, 32]
if self.mxfp8:
assert quant_config.quant_dtype == "mxfp8"
else:
assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()
assert quant_config.quant_dtype == torch.float8_e4m3fn
assert not quant_config.per_act_token_quant
assert not quant_config.per_out_ch_quant
self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
# Gated-activation params: silu == swigluoai with alpha=1, beta=0.
# FP8 (silu) configs leave these None, reproducing plain silu.
self.gemm1_alpha = (
quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
)
self.gemm1_beta = (
quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
)
@staticmethod
def activation_format() -> mk.FusedMoEActivationFormat:
@@ -147,14 +164,25 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
weight_key: QuantKey | None,
activation_key: QuantKey | None,
) -> bool:
SUPPORTED_W_A = [
(kFp8Static128BlockSym, kFp8Dynamic128Sym),
]
return (weight_key, activation_key) in SUPPORTED_W_A
if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):
return True
# MXFP8 1x32 uses the fp8_fp4 grouped GEMM with recipe (1, 32) — only
# available on Blackwell (SM100).
if (weight_key, activation_key) == (kMxfp8Static, kMxfp8Dynamic):
return current_platform.is_device_capability_family(100)
return False
@staticmethod
def _supports_activation(activation: MoEActivation) -> bool:
return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]
# silu/swigluoai go through the fused alpha/beta kernel; swiglustep
# uses the unfused activation path. The fused kernel reads packed w13
# (gate = first half, up = second half), so it implements the
# *uninterleaved* SwiGLU-OAI variant.
return activation in [
MoEActivation.SILU,
MoEActivation.SWIGLUSTEP,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
]
@staticmethod
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -179,7 +207,9 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
activation: MoEActivation,
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
assert self.block_shape is not None
block_m = self.block_shape[0]
# Use the contiguous-layout M alignment (matches apply()); block_shape[0]
# is the quant block (1 for MXFP8) and would under-size the workspace.
block_m = get_mk_alignment_for_contiguous_layout()[0]
M_sum = compute_aligned_M(
M, topk, local_num_experts, block_m, expert_tokens_meta
)
@@ -201,14 +231,24 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
M_sum, N = input.size()
activation_out_dim = self.adjust_N_for_activation(N, activation)
# 1. DeepGemm UE8M0: fused SiLU+mul+clamp+quant+pack
# silu and swigluoai are both expressible by the fused gated kernel via
# (alpha, beta): silu uses alpha=1, beta=0; swigluoai uses config values.
# The fused kernel reads packed w13, hence SWIGLUOAI_UNINTERLEAVE.
fused_gated = activation in (
MoEActivation.SILU,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
)
# 1. DeepGemm UE8M0: fused gate+mul+clamp+quant+pack
if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
if activation == MoEActivation.SILU:
if fused_gated:
return fused_silu_mul_fp8_quant_packed(
input=input,
output_q=output,
group_size=block_k,
clamp_limit=self.gemm1_clamp_limit,
alpha=self.gemm1_alpha,
beta=self.gemm1_beta,
)
act_out = torch.empty(
(M_sum, activation_out_dim), dtype=input.dtype, device=input.device
@@ -221,14 +261,17 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
)
return a2q, a2q_scale
# 2. Hopper / non‑E8M0: prefer the fused SiLU+mul+quant kernel
if activation == MoEActivation.SILU:
# 2. Hopper / non‑E8M0: prefer the fused gate+mul+quant kernel
if fused_gated:
use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
return silu_mul_per_token_group_quant_fp8_colmajor(
input=input,
output=output,
use_ue8m0=use_ue8m0,
clamp_limit=self.gemm1_clamp_limit,
group_size=block_k,
alpha=self.gemm1_alpha,
beta=self.gemm1_beta,
)
# 3. fallback path for non-SiLU activations in non‑UE8M0 cases.
@@ -292,12 +335,23 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
expert_map=expert_map,
expert_tokens_meta=expert_tokens_meta,
aq_out=a1q_perm,
# MXFP8 uses a 32-element activation-scale group (block_shape[1]);
# FP8-block keeps the default (128) alignment.
block_size=self.block_shape[1] if self.mxfp8 else None,
)
assert a1q.size(0) == M_sum
# MXFP8 (1x32) drives the fp8_fp4-aliased grouped GEMM with recipe
# (1, 32); the FP8 block path keeps the default (128) recipe.
gemm_kwargs = (
{"recipe_a": (1, self.block_shape[1]), "recipe_b": (1, self.block_shape[1])}
if self.mxfp8
else {}
)
mm1_out = _resize_cache(workspace2, (M_sum, N))
m_grouped_fp8_gemm_nt_contiguous(
(a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids
(a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids, **gemm_kwargs
)
activation_out_dim = self.adjust_N_for_activation(N, activation)
@@ -310,7 +364,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
mm2_out = _resize_cache(workspace2, (M_sum, K))
m_grouped_fp8_gemm_nt_contiguous(
(a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids
(a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids, **gemm_kwargs
)
if apply_router_weight_on_input:
vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+5 -1
View File
@@ -801,7 +801,11 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
return TopKWeightAndReduceDelegate()
def activation(
self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
self,
activation: MoEActivation,
output: torch.Tensor,
input: torch.Tensor,
**kwargs,
) -> None:
gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+1
View File
@@ -787,6 +787,7 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
activation: MoEActivation,
output: torch.Tensor,
input: torch.Tensor,
**kwargs,
) -> None:
quant_config = self.quant_config or FUSED_MOE_UNQUANTIZED_CONFIG
if activation == MoEActivation.SWIGLUOAI:
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+51 -23
View File
@@ -28,10 +28,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
TopKWeightAndReduceDelegate,
TopKWeightAndReduceNoOP,
)
from vllm.model_executor.layers.fused_moe.utils import (
_resize_cache,
swiglu_limit_func,
)
from vllm.model_executor.layers.fused_moe.utils import _resize_cache
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
get_marlin_input_dtype,
marlin_make_workspace_new,
@@ -74,9 +71,7 @@ def _fused_marlin_moe(
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor,
activation: MoEActivation = MoEActivation.SILU,
activation_func: Callable[
[MoEActivation, torch.Tensor, torch.Tensor], None
] = apply_moe_activation,
activation_func: Callable[..., None] = apply_moe_activation,
input_global_scale1: torch.Tensor | None = None,
input_global_scale2: torch.Tensor | None = None,
global_scale1: torch.Tensor | None = None,
@@ -94,6 +89,8 @@ def _fused_marlin_moe(
input_dtype: torch.dtype | None = None,
is_k_full: bool = True,
clamp_limit: float | None = None,
gemm1_alpha: float = 1.0,
gemm1_beta: float = 0.0,
) -> torch.Tensor:
assert hidden_states.ndim == 2
M, K = hidden_states.size()
@@ -161,18 +158,16 @@ def _fused_marlin_moe(
use_fp32_reduce=True,
is_zp_float=False,
)
if clamp_limit is not None and activation == MoEActivation.SILU:
swiglu_limit_func(
intermediate_cache2,
intermediate_cache1.view(-1, w13_num_shards * N),
clamp_limit,
)
else:
activation_func(
activation,
intermediate_cache2,
intermediate_cache1.view(-1, w13_num_shards * N),
)
# apply_moe_activation fuses the clamp/gate params: SILU + clamp_limit and
# SWIGLUOAI_UNINTERLEAVE both map to the silu_and_mul_with_clamp kernel.
activation_func(
activation,
intermediate_cache2,
intermediate_cache1.view(-1, w13_num_shards * N),
clamp_limit=clamp_limit,
alpha=gemm1_alpha,
beta=gemm1_beta,
)
if output is None:
output = intermediate_cache3
@@ -238,9 +233,7 @@ def fused_marlin_moe(
apply_router_weight_on_input: bool = False,
global_num_experts: int = -1,
activation: MoEActivation = MoEActivation.SILU,
activation_func: Callable[
[MoEActivation, torch.Tensor, torch.Tensor], None
] = apply_moe_activation,
activation_func: Callable[..., None] = apply_moe_activation,
moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
expert_map: torch.Tensor | None = None,
input_global_scale1: torch.Tensor | None = None,
@@ -260,6 +253,8 @@ def fused_marlin_moe(
output: torch.Tensor | None = None,
input_dtype: torch.dtype | None = None,
clamp_limit: float | None = None,
gemm1_alpha: float = 1.0,
gemm1_beta: float = 0.0,
) -> torch.Tensor:
"""
This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -373,6 +368,8 @@ def fused_marlin_moe(
input_dtype=input_dtype,
is_k_full=is_k_full,
clamp_limit=clamp_limit,
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
).view(-1, topk, K)
if output is None:
@@ -415,6 +412,8 @@ def batched_fused_marlin_moe(
output: torch.Tensor | None = None,
input_dtype: torch.dtype | None = None,
clamp_limit: float | None = None,
gemm1_alpha: float = 1.0,
gemm1_beta: float = 0.0,
) -> torch.Tensor:
"""
This function massages the inputs so the batched hidden_states can be
@@ -544,6 +543,8 @@ def batched_fused_marlin_moe(
input_dtype=input_dtype,
is_k_full=is_k_full,
clamp_limit=clamp_limit,
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
)
output = output.view(B, BATCH_TOKENS_MAX, K)
@@ -579,6 +580,15 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
self.is_k_full = is_k_full
self.input_dtype = get_marlin_input_dtype()
self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
# Gated-activation params (used by SWIGLUOAI_UNINTERLEAVE on packed w13).
# silu == swigluoai with alpha=1, beta=0; configs that don't set these
# (plain silu) fall back to the silu identity.
self.gemm1_alpha = (
quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
)
self.gemm1_beta = (
quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
)
super().__init__(
moe_config=moe_config,
@@ -627,6 +637,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
MoEActivation.GELU,
MoEActivation.GELU_TANH,
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
MoEActivation.SWIGLUSTEP,
MoEActivation.SILU_NO_MUL,
MoEActivation.GELU_NO_MUL,
@@ -787,6 +798,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
is_k_full=self.is_k_full,
input_dtype=self.input_dtype,
clamp_limit=self.gemm1_clamp_limit,
gemm1_alpha=self.gemm1_alpha,
gemm1_beta=self.gemm1_beta,
)
return
@@ -805,6 +818,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
act_enum: MoEActivation,
act_output: torch.Tensor,
act_input: torch.Tensor,
*,
clamp_limit: float | None = None,
alpha: float = 1.0,
beta: float = 0.0,
) -> None:
# act_input = intermediate_cache1 (M*topk, 2N for gated)
# act_output = intermediate_cache2 (M*topk, N)
@@ -834,7 +851,14 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
"tlm": token_lora_mapping,
}
)
self.activation(act_enum, act_output, act_input)
self.activation(
act_enum,
act_output,
act_input,
clamp_limit=clamp_limit,
alpha=alpha,
beta=beta,
)
lora_state["cache2"] = act_output
def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None:
@@ -888,6 +912,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
is_k_full=self.is_k_full,
input_dtype=self.input_dtype,
clamp_limit=self.gemm1_clamp_limit,
gemm1_alpha=self.gemm1_alpha,
gemm1_beta=self.gemm1_beta,
)
def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
@@ -996,4 +1022,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
input_dtype=self.input_dtype,
is_k_full=self.is_k_full,
clamp_limit=self.gemm1_clamp_limit,
gemm1_alpha=self.gemm1_alpha,
gemm1_beta=self.gemm1_beta,
)
vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py
+176
View File
@@ -0,0 +1,176 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MXFP8 (1x32 block, E8M0 scale) MoE experts on Triton.
``Mxfp8TritonExpertsBase`` stashes E8M0 weight scales for checkpoint layout.
``Mxfp8EmulationTritonExperts`` dequantizes to BF16 and runs ``TritonExperts``
for devices without a native MXFP8 MoE kernel (e.g. ROCm gfx942 / MI300).
"""
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEQuantConfig,
)
from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
dequant_mxfp8_to_bf16,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
kMxfp8Dynamic,
kMxfp8Static,
)
logger = init_logger(__name__)
class Mxfp8TritonExpertsBase(TritonExperts):
"""Shared MXFP8 MoE setup: stash E8M0 scales, clear scales on ``quant_config``."""
def __init__(
self,
moe_config: FusedMoEConfig,
quant_config: FusedMoEQuantConfig,
):
super().__init__(moe_config, quant_config)
self.w1_scale_val = self.quant_config.w1_scale
self.w2_scale_val = self.quant_config.w2_scale
self.quant_config._w1.scale = None
self.quant_config._w2.scale = None
@staticmethod
def _supports_quant_scheme(
weight_key: QuantKey | None,
activation_key: QuantKey | None,
) -> bool:
return (weight_key, activation_key) == (kMxfp8Static, kMxfp8Dynamic)
@staticmethod
def _supports_activation(activation) -> bool:
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
return True
return TritonExperts._supports_activation(activation)
class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase):
"""Dequantize MXFP8 weights to BF16 on the fly and run ``TritonExperts``."""
def __init__(
self,
moe_config: FusedMoEConfig,
quant_config: FusedMoEQuantConfig,
):
super().__init__(moe_config, quant_config)
logger.warning_once(
"Using Mxfp8EmulationTritonExperts MoE backend. Weights are "
"dequantized to BF16 on the fly; this is slower than a native "
"MXFP8 MoE kernel and is intended for devices without one."
)
@property
def quant_dtype(self) -> torch.dtype | str | None:
# BF16 fallback: do not MXFP8-quantize activations in ``TritonExperts``.
return None
@property
def block_shape(self) -> list[int] | None:
return None
@property
def expects_unquantized_inputs(self) -> bool:
return True
@staticmethod
def _supports_current_device() -> bool:
return True
def activation(
self,
activation,
output: torch.Tensor,
input: torch.Tensor,
**kwargs,
):
"""Apply GEMM1 activation with quant-config alpha/beta/clamp."""
from vllm.model_executor.layers.fused_moe.activation import (
MoEActivation,
apply_moe_activation,
)
if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
limit = self.quant_config.gemm1_clamp_limit
if limit is None:
raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit")
alpha = self.quant_config.gemm1_alpha
alpha = 1.702 if alpha is None else float(alpha)
beta = self.quant_config.gemm1_beta
beta = 1.0 if beta is None else float(beta)
apply_moe_activation(
activation,
output,
input,
clamp_limit=float(limit),
alpha=alpha,
beta=beta,
)
return
super().activation(activation, output, input)
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,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
a2_scale: torch.Tensor | None,
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
):
# If the weights were already dequantized to BF16 at load time
# (process_weights_after_loading on devices without a native MXFP8 MoE
# kernel), use them directly -- no per-step dequant. MXFP8 weights are
# 1-byte FP8 (element_size 1); BF16/FP16 are >= 2 bytes.
if w1.element_size() >= 2:
# tl.dot requires w and activations share a dtype; .to() is a no-op
# when they already match (e.g. both BF16).
w1_bf16 = w1.to(hidden_states.dtype)
w2_bf16 = w2.to(hidden_states.dtype)
else:
w1_bf16 = dequant_mxfp8_to_bf16(w1, self.w1_scale_val).to(
hidden_states.dtype
)
w2_bf16 = dequant_mxfp8_to_bf16(w2, self.w2_scale_val).to(
hidden_states.dtype
)
super().apply(
output=output,
hidden_states=hidden_states,
w1=w1_bf16,
w2=w2_bf16,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=activation,
global_num_experts=global_num_experts,
expert_map=expert_map,
a1q_scale=None,
a2_scale=None,
workspace13=workspace13,
workspace2=workspace2,
expert_tokens_meta=expert_tokens_meta,
apply_router_weight_on_input=apply_router_weight_on_input,
)
vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py
+326
View File
@@ -0,0 +1,326 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Native MXFP8 (1x32 block, E8M0 scale) MoE for AMD CDNA4 (gfx950) via Triton
``tl.dot_scaled`` (hardware microscaling matmul).
The expert GEMMs consume the FP8 E4M3 weights and their E8M0 block scales
directly (no dequant-to-BF16), and activations are MXFP8-quantized per token.
On CDNA4 ``dot_scaled`` maps to the native MX matrix-core ops; on other archs
Triton upcasts to BF16 (so this stays correct, just not faster) — but the
oracle only selects this path on gfx950 and routes everything else to the
BF16 ``Mxfp8EmulationTritonExperts`` fallback.
Structure mirrors vLLM's ``fused_moe_kernel``: tokens are sorted by expert
(``moe_align_block_size``); each program computes a ``[BLOCK_M, BLOCK_N]`` tile
for one expert, accumulating over K with ``dot_scaled``. SwiGLU-OAI activation
and the top-k weighted reduction run in PyTorch between/after the two GEMMs.
"""
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.experts.mxfp8_emulation_moe import (
Mxfp8TritonExpertsBase,
)
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
moe_align_block_size,
)
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
mxfp8_e4m3_quantize,
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
logger = init_logger(__name__)
@triton.jit
def _mxfp8_grouped_gemm_kernel(
a_ptr,
a_scale_ptr,
b_ptr,
b_scale_ptr,
c_ptr,
topk_weights_ptr,
sorted_token_ids_ptr,
expert_ids_ptr,
num_tokens_post_padded_ptr,
N,
K,
num_valid_tokens,
top_k,
stride_am,
stride_ak,
stride_asm,
stride_ask,
stride_be,
stride_bn,
stride_bk,
stride_bse,
stride_bsn,
stride_bsk,
stride_cm,
stride_cn,
A_DIV: tl.constexpr,
MUL_WEIGHT: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
num_post = tl.load(num_tokens_post_padded_ptr)
if pid_m * BLOCK_M >= num_post:
return
offs_tid = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_token = tl.load(sorted_token_ids_ptr + offs_tid).to(tl.int64)
token_mask = offs_token < num_valid_tokens
off_e = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
offs_sk = tl.arange(0, BLOCK_K // 32)
a_row = offs_token // A_DIV
a_ptrs = a_ptr + a_row[:, None] * stride_am + offs_k[None, :] * stride_ak
as_ptrs = a_scale_ptr + a_row[:, None] * stride_asm + offs_sk[None, :] * stride_ask
b_ptrs = (
b_ptr
+ off_e * stride_be
+ offs_n[:, None] * stride_bn
+ offs_k[None, :] * stride_bk
)
bs_ptrs = (
b_scale_ptr
+ off_e * stride_bse
+ offs_n[:, None] * stride_bsn
+ offs_sk[None, :] * stride_bsk
)
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
n_mask = offs_n < N
for _ in range(0, tl.cdiv(K, BLOCK_K)):
a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
b = tl.load(b_ptrs, mask=n_mask[:, None], other=0.0)
asc = tl.load(as_ptrs, mask=token_mask[:, None], other=0)
bsc = tl.load(bs_ptrs, mask=n_mask[:, None], other=0)
acc += tl.dot_scaled(a, asc, "e4m3", b.T, bsc, "e4m3")
a_ptrs += BLOCK_K * stride_ak
b_ptrs += BLOCK_K * stride_bk
as_ptrs += (BLOCK_K // 32) * stride_ask
bs_ptrs += (BLOCK_K // 32) * stride_bsk
if MUL_WEIGHT:
w = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
acc = acc * w[:, None]
c_ptrs = c_ptr + offs_token[:, None] * stride_cm + offs_n[None, :] * stride_cn
tl.store(
c_ptrs,
acc.to(c_ptr.dtype.element_ty),
mask=token_mask[:, None] & n_mask[None, :],
)
def _grouped_gemm_mxfp8(
a_q: torch.Tensor, # [M, K] fp8 e4m3
a_scale: torch.Tensor, # [M, K//32] uint8 (E8M0)
w: torch.Tensor, # [E, N, K] fp8 e4m3
w_scale: torch.Tensor, # [E, N, K//32] uint8 (E8M0)
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor,
num_valid_tokens: int,
top_k: int,
block_m: int,
out_dtype: torch.dtype,
a_div: int,
mul_weight_by: torch.Tensor | None = None,
expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
M_routed = num_valid_tokens
E, N, K = w.shape
assert K % 128 == 0, f"MXFP8 native MoE requires K%128==0, got K={K}"
# Under expert parallelism (expert_map set) tokens routed to non-local
# experts are dropped from sorted_token_ids, so their output rows are never
# written — zero them so the downstream reduction ignores their garbage.
alloc = torch.zeros if expert_map is not None else torch.empty
out = alloc((M_routed, N), dtype=out_dtype, device=a_q.device)
BLOCK_N = 128
BLOCK_K = 128
grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, BLOCK_N))
_mxfp8_grouped_gemm_kernel[grid](
a_q,
a_scale,
w,
w_scale,
out,
mul_weight_by if mul_weight_by is not None else a_q,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
N,
K,
num_valid_tokens,
top_k,
a_q.stride(0),
a_q.stride(1),
a_scale.stride(0),
a_scale.stride(1),
w.stride(0),
w.stride(1),
w.stride(2),
w_scale.stride(0),
w_scale.stride(1),
w_scale.stride(2),
out.stride(0),
out.stride(1),
A_DIV=a_div,
MUL_WEIGHT=mul_weight_by is not None,
BLOCK_M=block_m,
BLOCK_N=BLOCK_N,
BLOCK_K=BLOCK_K,
num_warps=8,
)
return out
def fused_moe_mxfp8_native(
hidden_states: torch.Tensor, # [T, H] bf16
w13: torch.Tensor, # [E, 2I, H] fp8
w13_scale: torch.Tensor, # [E, 2I, H//32] uint8
w2: torch.Tensor, # [E, H, I] fp8
w2_scale: torch.Tensor, # [E, H, I//32] uint8
topk_weights: torch.Tensor, # [T, top_k]
topk_ids: torch.Tensor, # [T, top_k] (global expert ids)
*,
alpha: float,
beta: float,
limit: float | None,
global_num_experts: int,
expert_map: torch.Tensor | None,
) -> torch.Tensor:
T, H = hidden_states.shape
top_k = topk_ids.shape[1]
M = T * top_k
block_m = 64
sorted_ids, expert_ids, num_post = moe_align_block_size(
topk_ids,
block_m,
global_num_experts,
expert_map,
ignore_invalid_experts=expert_map is not None,
)
# GEMM1: x (mxfp8) @ w13^T -> [M, 2I]
a_q, a_s = mxfp8_e4m3_quantize(hidden_states)
g1 = _grouped_gemm_mxfp8(
a_q,
a_s,
w13,
w13_scale,
sorted_ids,
expert_ids,
num_post,
M,
top_k,
block_m,
hidden_states.dtype,
a_div=top_k,
expert_map=expert_map,
) # [M, 2I]
# SwiGLU-OAI (split layout: gate=g1[:, :I], up=g1[:, I:]) FUSED with the
# GEMM2 MXFP8 activation-quant in one fp32 Triton pass — no bf16 ``act``
# round-trip to HBM. Bit-exact vs the unfused swiglu+quant chain on measured
# MoE shapes, and ~1.2-1.9x faster on that step in isolation. (Not the #22
# ``silu_and_mul_with_clamp`` op: it rounds intermediates to bf16, rel ~3e-3.)
# Lazy import: the amd.ops package pulls in the minimax_m3 platform dispatch,
# only resolvable after the model module finishes loading.
from vllm.models.minimax_m3.amd.ops import swiglu_oai_quantize_mxfp8
# GEMM2: act (mxfp8) @ w2^T -> [M, H], weighted by topk_weights, then reduce.
act_q, act_s = swiglu_oai_quantize_mxfp8(g1, alpha=alpha, beta=beta, limit=limit)
g2 = _grouped_gemm_mxfp8(
act_q,
act_s,
w2,
w2_scale,
sorted_ids,
expert_ids,
num_post,
M,
top_k,
block_m,
torch.float32,
a_div=1,
mul_weight_by=topk_weights.reshape(-1).to(torch.float32),
expert_map=expert_map,
) # [M, H] == [T*top_k, H]
return g2.view(T, top_k, H).sum(dim=1).to(hidden_states.dtype)
class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
"""Native MXFP8 MoE (CDNA4 ``dot_scaled``) on gfx950."""
@property
def quant_dtype(self) -> torch.dtype | str | None:
return self.quant_config.quant_dtype
@property
def block_shape(self) -> list[int] | None:
return self.quant_config.block_shape
@property
def expects_unquantized_inputs(self) -> bool:
# Activations are MXFP8-quantized inside ``fused_moe_mxfp8_native``.
return True
@staticmethod
def _supports_current_device() -> bool:
return current_platform.is_rocm() and current_platform.supports_mx()
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,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
a2_scale: torch.Tensor | None,
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
):
alpha = self.quant_config.gemm1_alpha
alpha = 1.702 if alpha is None else float(alpha)
beta = self.quant_config.gemm1_beta
beta = 1.0 if beta is None else float(beta)
limit = self.quant_config.gemm1_clamp_limit
limit = None if limit is None else float(limit)
out = fused_moe_mxfp8_native(
hidden_states,
w1,
self.w1_scale_val,
w2,
self.w2_scale_val,
topk_weights,
topk_ids,
alpha=alpha,
beta=beta,
limit=limit,
global_num_experts=global_num_experts,
expert_map=expert_map,
)
output.copy_(out)
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+32 -2
View File
@@ -64,6 +64,15 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
self.quantization_emulation = False
super().__init__(moe_config, quant_config)
self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
# Gated-activation params: silu == swigluoai with alpha=1, beta=0.
self.gemm1_alpha = (
quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
)
self.gemm1_beta = (
quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
)
@staticmethod
def activation_format() -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
@@ -107,6 +116,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
MoEActivation.GELU,
MoEActivation.GELU_TANH,
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
MoEActivation.SWIGLUSTEP,
MoEActivation.SILU_NO_MUL,
MoEActivation.GELU_NO_MUL,
@@ -129,14 +139,34 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
return TopKWeightAndReduceNoOP()
def activation(
self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
self,
activation: MoEActivation,
output: torch.Tensor,
input: torch.Tensor,
**kwargs,
) -> None:
gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
swiglu_limit_func(output, input, float(gemm1_clamp_limit))
return
super().activation(activation, output, input)
# SWIGLUOAI_UNINTERLEAVE routes to the silu_and_mul_with_clamp kernel and
# needs the clamped-SwiGLU params (gemm1_clamp_limit/alpha/beta read from
# the quant config in __init__) forwarded; without a clamp_limit it
# asserts. Other activations ignore alpha/beta/clamp_limit.
if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
assert gemm1_clamp_limit is not None, (
"SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit"
)
super().activation(
activation,
output,
input,
clamp_limit=gemm1_clamp_limit,
alpha=self.gemm1_alpha,
beta=self.gemm1_beta,
)
def workspace_shapes(
self,
vllm/model_executor/layers/fused_moe/layer.py
+6
View File
@@ -120,6 +120,8 @@ def FusedMoE(
scoring_func: str = "softmax",
routed_scaling_factor: float = 1.0,
swiglu_limit: float | None = None,
swiglu_alpha: float | None = None,
swiglu_beta: float | None = None,
e_score_correction_bias: torch.Tensor | None = None,
apply_router_weight_on_input: bool = False,
activation: str = "silu",
@@ -322,6 +324,8 @@ def FusedMoE(
device=vllm_config.device_config.device,
routing_method=router.routing_method_type, # Not ideal
swiglu_limit=swiglu_limit,
swiglu_alpha=swiglu_alpha,
swiglu_beta=swiglu_beta,
max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
)
@@ -353,6 +357,8 @@ def FusedMoE(
if not apply_routed_scale_to_output
else 1.0,
swiglu_limit=swiglu_limit,
swiglu_alpha=swiglu_alpha,
swiglu_beta=swiglu_beta,
# TODO get from router? needs to be truncated?
e_score_correction_bias=e_score_correction_bias,
apply_router_weight_on_input=apply_router_weight_on_input,
vllm/model_executor/layers/fused_moe/modular_kernel.py
+11 -2
View File
@@ -880,9 +880,18 @@ class FusedMoEExpertsModular(FusedMoEExperts):
return N if not activation.is_gated else N // 2
def activation(
self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
self,
activation: MoEActivation,
output: torch.Tensor,
input: torch.Tensor,
*,
clamp_limit: float | None = None,
alpha: float = 1.0,
beta: float = 0.0,
) -> None:
apply_moe_activation(activation, output, input)
apply_moe_activation(
activation, output, input, clamp_limit=clamp_limit, alpha=alpha, beta=beta
)
@abstractmethod
def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+21 -4
View File
@@ -52,6 +52,13 @@ class Fp8MoeBackend(Enum):
BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
XPU = "XPU"
CPU = "CPU"
# Dequantize-to-BF16 emulation for MXFP8 on devices without a native
# MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here.
EMULATION = "EMULATION"
# MXFP8 MoE via a Triton ``dot_scaled`` kernel that lowers to CDNA4
# (gfx950) native MX matrix-core ops. Weights stay in MXFP8 (no load-time
# format conversion); the FP8 values + E8M0 scales are consumed directly.
NATIVE_MXFP8 = "NATIVE_MXFP8"
def _get_priority_backends(
@@ -463,6 +470,10 @@ def convert_to_fp8_moe_kernel_format(
Fp8MoeBackend.VLLM_CUTLASS,
Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
Fp8MoeBackend.XPU,
# EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes
# the MXFP8 weights as-is — neither needs a load-time layout change.
Fp8MoeBackend.EMULATION,
Fp8MoeBackend.NATIVE_MXFP8,
]:
raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}")
@@ -481,6 +492,8 @@ def make_fp8_moe_quant_config(
per_act_token_quant: bool = False,
per_out_ch_quant: bool = False,
swiglu_limit: float | None = None,
gemm1_alpha: float | None = None,
gemm1_beta: float | None = None,
) -> FusedMoEQuantConfig:
"""
Create FusedMoEQuantConfig for the specified FP8 Backend.
@@ -503,6 +516,9 @@ def make_fp8_moe_quant_config(
w1_bias=w1_bias,
w2_bias=w2_bias,
block_shape=block_shape,
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=swiglu_limit,
)
# Flashinfer CUTLASS per-tensor uses single dq scale
@@ -522,10 +538,9 @@ def make_fp8_moe_quant_config(
g2_alphas=(w2_scale * a2_scale).squeeze(),
gemm1_clamp_limit=swiglu_limit,
)
# MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to
# _mxfp8_e4m3_quantize rather than standard FP8 block quantization.
# Non-swizzled layout is required since the TRTLLM kernel expects
# scales in (num_tokens, hidden_dim // 32) format.
# MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are
# the non-swizzled (num_tokens, hidden_dim // 32) uint8 UE8M0 layout for all
# backends; the DeepGEMM expert permute repacks them for the grouped GEMM.
if block_shape == [1, 32]:
return FusedMoEQuantConfig.make(
"mxfp8",
@@ -537,6 +552,8 @@ def make_fp8_moe_quant_config(
a2_scale=a2_scale,
block_shape=block_shape,
is_scale_swizzled=False,
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=swiglu_limit,
)
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+49 -1
View File
@@ -12,22 +12,43 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kMxfp8Dynamic,
kMxfp8Static,
)
from vllm.platforms import current_platform
logger = init_logger(__name__)
_SUPPORTED_BACKENDS = (
Fp8MoeBackend.FLASHINFER_TRTLLM,
Fp8MoeBackend.DEEPGEMM,
Fp8MoeBackend.MARLIN,
Fp8MoeBackend.XPU,
)
_BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
"flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM,
"deep_gemm": Fp8MoeBackend.DEEPGEMM,
"marlin": Fp8MoeBackend.MARLIN,
"xpu": Fp8MoeBackend.XPU,
}
def _mxfp8_backend_to_kernel_cls(
backend: Fp8MoeBackend,
) -> list[type[mk.FusedMoEExperts]]:
"""Resolve the MXFP8 expert classes for a backend.
DeepGEMM resolves directly to ``DeepGemmExperts`` (not the
``TritonOrDeepGemmExperts`` wrapper, whose Triton fallback cannot handle the
MXFP8 1x32 scheme); all other backends defer to the FP8 resolver.
"""
if backend == Fp8MoeBackend.DEEPGEMM:
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
DeepGemmExperts,
)
return [DeepGemmExperts]
return backend_to_kernel_cls(backend)
def _select_kernel_cls(
backend: Fp8MoeBackend,
config: FusedMoEConfig,
@@ -39,7 +60,7 @@ def _select_kernel_cls(
else mk.FusedMoEActivationFormat.Standard
)
last_reason: str | None = None
for cls in backend_to_kernel_cls(backend):
for cls in _mxfp8_backend_to_kernel_cls(backend):
supported, reason = cls.is_supported_config(
cls,
config,
@@ -55,6 +76,29 @@ def _select_kernel_cls(
)
def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
"""ROCm fallback when vendor MXFP8 backends are unavailable."""
if current_platform.supports_mx():
from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
Mxfp8NativeTritonExperts,
)
logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.")
return Fp8MoeBackend.NATIVE_MXFP8, Mxfp8NativeTritonExperts
from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
Mxfp8EmulationTritonExperts,
)
logger.info_once(
"No native MXFP8 MoE backend available on this device; "
"MXFP8 weights will be dequantized to BF16 once at load time and the "
"MoE will run in BF16 (no per-step dequant)."
)
return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts
def select_mxfp8_moe_backend(
config: FusedMoEConfig,
) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
@@ -88,4 +132,8 @@ def select_mxfp8_moe_backend(
logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
return backend, experts_cls
# simplify the logic for rocm, refactor later when more backends are supported
if current_platform.is_rocm():
return _select_rocm_mxfp8_backend()
raise ValueError("No MXFP8 MoE backends available.")
vllm/model_executor/layers/fused_moe/routed_experts.py
+4
View File
@@ -72,6 +72,8 @@ class RoutedExperts(PluggableLayer):
scoring_func: str = "softmax",
routed_scaling_factor: float = 1.0,
swiglu_limit: float | None = None,
swiglu_alpha: float | None = None,
swiglu_beta: float | None = None,
e_score_correction_bias: torch.Tensor | None = None,
apply_router_weight_on_input: bool = False,
):
@@ -103,6 +105,8 @@ class RoutedExperts(PluggableLayer):
self.scoring_func = scoring_func
self.routed_scaling_factor = routed_scaling_factor
self.swiglu_limit = swiglu_limit
self.swiglu_alpha = swiglu_alpha
self.swiglu_beta = swiglu_beta
self.e_score_correction_bias = e_score_correction_bias
self.apply_router_weight_on_input = apply_router_weight_on_input
# End random parameters
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+4 -5
View File
@@ -29,9 +29,9 @@ class GateLinear(ReplicatedLinear):
DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
# Dimensions supported by the fp32 specialized kernel
FP32_SUPPORTED_NUM_EXPERTS = [256]
FP32_SUPPORTED_HIDDEN_SIZES = [3072]
# (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
# (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}
FP32_MAX_TOKENS = 32
def __init__(
@@ -82,8 +82,7 @@ class GateLinear(ReplicatedLinear):
and self.weight.dtype == torch.float32
and current_platform.is_cuda()
and (is_hopper or is_blackwell)
and output_size in self.FP32_SUPPORTED_NUM_EXPERTS
and input_size in self.FP32_SUPPORTED_HIDDEN_SIZES
and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
)
# cuBLAS bf16→fp32 eligibility
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+20 -8
View File
@@ -11,7 +11,6 @@ import vllm.envs as envs
from vllm.logger import init_logger
from vllm.model_executor.custom_op import CustomOp
from vllm.model_executor.layers.fused_moe.config import (
FUSED_MOE_UNQUANTIZED_CONFIG,
FusedMoEConfig,
FusedMoEQuantConfig,
biased_moe_quant_config,
@@ -184,11 +183,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
if not is_weight_update:
# Setup moe kernel only on the first call. For the unquantized
# method, moe_quant_config is either the constant
# FUSED_MOE_UNQUANTIZED_CONFIG or biased_moe_quant_config(...)
# which references layer.w{13,2}_bias; since weight updates
# mutate those bias tensors in place, the kernel does not need
# to be re-built.
# method, moe_quant_config carries no quantized scales -- only
# optional w{13,2}_bias references and SwiGLU gate params. Since
# weight updates mutate those bias tensors in place, the kernel
# does not need to be re-built.
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
assert self.moe_quant_config is not None
assert self.experts_cls is not None
@@ -272,13 +270,27 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
)
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
# SwiGLU/swigluoai gate params live on the layer; plumb them into the
# quant config so the fused activation (e.g. swigluoai_uninterleave on
# MiniMax-M3) receives gemm1_clamp_limit/alpha/beta.
gemm1_alpha = getattr(layer, "swiglu_alpha", None)
gemm1_beta = getattr(layer, "swiglu_beta", None)
gemm1_clamp_limit = getattr(layer, "swiglu_limit", None)
if self.moe.has_bias:
return biased_moe_quant_config(
layer.w13_bias,
layer.w2_bias,
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=gemm1_clamp_limit,
)
else:
return FUSED_MOE_UNQUANTIZED_CONFIG
return FusedMoEQuantConfig.make(
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=gemm1_clamp_limit,
)
def apply(
self,
vllm/model_executor/layers/fused_moe/utils.py
+2
View File
@@ -313,6 +313,8 @@ def moe_kernel_quantize_input(
"moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
"quantization emulation. Please open an issue."
)
# Non-swizzled (M, K/32) uint8 UE8M0 scales; deepgemm_moe_permute packs
# them for DeepGEMM, TRTLLM takes them as-is.
return _mxfp8_e4m3_quantize(
A,
A_scale,
vllm/model_executor/layers/linear.py
+186
View File
@@ -1105,6 +1105,7 @@ class QKVParallelLinear(ColumnParallelLinear):
shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
shard_size = self._get_shard_size_mapping(loaded_shard_id)
assert shard_offset is not None and shard_size is not None
if isinstance(param, BlockQuantScaleParameter):
weight_block_size = getattr(self, "weight_block_size", None)
@@ -1302,6 +1303,191 @@ class QKVParallelLinear(ColumnParallelLinear):
param_data.copy_(loaded_weight)
class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):
"""QKV projection fused with a lightning-indexer's index_q/index_k.
NOTE: MiniMax-M3-specific. This is tailored to the M3 sparse-attention
layers (it assumes the indexer's head count equals the KV head count and
shares the main head_dim); it is not a general-purpose linear layer. It
lives here only to sit alongside QKVParallelLinear, whose sharding /
weight-loading machinery it reuses.
A single column-parallel GEMM emits, per rank::
[q | k | v | index_q | index_k]
``index_q`` must have the same head count as the KV heads
(``total_num_index_heads == total_num_kv_heads``) and ``index_head_size ==
head_size``, so it shards exactly like K/V -- including the KV-head
*replication* path when ``tp_size > total_num_kv_heads`` (this is what makes
a TP size greater than the KV-head count work). ``index_k`` is a single
shared head, replicated to every rank.
"""
def __init__(
self,
hidden_size: int,
head_size: int,
total_num_heads: int,
total_num_kv_heads: int,
total_num_index_heads: int,
index_head_size: int,
bias: bool = False,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
# index_q rides the KV-head sharding/replication path, so its head count
# must match the KV heads.
assert total_num_index_heads == total_num_kv_heads, (
"MinimaxM3QKVParallelLinearWithIndexer requires "
"total_num_index_heads == total_num_kv_heads"
)
self.hidden_size = hidden_size
self.head_size = head_size
self.v_head_size = head_size
self.total_num_heads = total_num_heads
self.total_num_kv_heads = total_num_kv_heads
self.total_num_index_heads = total_num_index_heads
self.index_head_size = index_head_size
tp_size = get_tensor_model_parallel_world_size()
self.num_heads = divide(self.total_num_heads, tp_size)
if tp_size >= self.total_num_kv_heads:
self.num_kv_heads = 1
self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
else:
self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
self.num_kv_head_replicas = 1
# index_q shards identically to the KV heads.
self.num_index_heads = self.num_kv_heads
# Global per-group sizes (replicated groups counted x tp_size, matching
# the QKVParallelLinear convention). index_k is a single replicated head.
q = self.num_heads * self.head_size
kv = self.num_kv_heads * self.head_size
iq = self.num_index_heads * self.index_head_size
ik = self.index_head_size
self.output_sizes = [
q * tp_size, # q
kv * tp_size, # k
kv * tp_size, # v
iq * tp_size, # index_q
ik * tp_size, # index_k (replicated)
]
# Skip QKVParallelLinear.__init__ (3-group layout); build the 5-group
# column-parallel weight directly.
ColumnParallelLinear.__init__(
self,
input_size=self.hidden_size,
output_size=sum(self.output_sizes),
bias=bias,
gather_output=False,
quant_config=quant_config,
prefix=prefix,
)
def validate_shard_id(self, loaded_shard_id: str | None) -> None:
if loaded_shard_id is None:
return
if loaded_shard_id not in ("q", "k", "v", "index_q", "index_k"):
raise ValueError(
"Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of "
"'q', 'k', 'v', 'index_q', 'index_k'; got "
f"{loaded_shard_id}."
)
def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None:
h = self.head_size
nq, nkv, nidx = self.num_heads, self.num_kv_heads, self.num_index_heads
return {
"q": 0,
"k": nq * h,
"v": (nq + nkv) * h,
"index_q": (nq + 2 * nkv) * h,
"index_k": (nq + 2 * nkv + nidx) * h,
}.get(loaded_shard_id)
def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None:
h = self.head_size
return {
"q": self.num_heads * h,
"k": self.num_kv_heads * h,
"v": self.num_kv_heads * h,
"index_q": self.num_index_heads * h,
"index_k": self.index_head_size,
}.get(loaded_shard_id)
def weight_loader_v2(
self,
param: BasevLLMParameter,
loaded_weight: torch.Tensor,
loaded_shard_id: str | None = None,
) -> None:
self.validate_shard_id(loaded_shard_id)
# Index checkpoints are never pre-fused on disk; a shard id is always given.
assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k")
shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
shard_size = self._get_shard_size_mapping(loaded_shard_id)
assert shard_offset is not None and shard_size is not None
if isinstance(param, BlockQuantScaleParameter):
weight_block_size = getattr(self, "weight_block_size", None)
shard_size, shard_offset = adjust_block_scale_shard(
weight_block_size, shard_size, shard_offset
)
# index_k is fully replicated: num_heads == tp_size makes
# load_qkv_weight pick shard_id_int == 0 on every rank. q/k/v/index_q ride
# the KV-head replication factor.
num_heads = (
self.tp_size if loaded_shard_id == "index_k" else self.num_kv_head_replicas
)
param.load_qkv_weight(
loaded_weight=loaded_weight,
num_heads=num_heads,
shard_id=loaded_shard_id,
shard_offset=shard_offset,
shard_size=shard_size,
tp_rank=self.tp_rank,
)
def weight_loader(
self,
param: Parameter,
loaded_weight: torch.Tensor,
loaded_shard_id: str | None = None,
) -> None:
# Unquantized (bf16) path. MXFP8 checkpoints use weight_loader_v2; this
# keeps an unquantized load correct too.
self.validate_shard_id(loaded_shard_id)
assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k")
output_dim = getattr(param, "output_dim", None)
assert output_dim is not None
shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
shard_size = self._get_shard_size_mapping(loaded_shard_id)
assert shard_offset is not None and shard_size is not None
if isinstance(param, BlockQuantScaleParameter):
weight_block_size = getattr(self, "weight_block_size", None)
shard_size, shard_offset = adjust_block_scale_shard(
weight_block_size, shard_size, shard_offset
)
param_data = param.data.narrow(output_dim, shard_offset, shard_size)
if loaded_shard_id == "q":
shard_rank = self.tp_rank
elif loaded_shard_id == "index_k":
shard_rank = 0 # replicated to every rank
else:
shard_rank = self.tp_rank // self.num_kv_head_replicas
loaded_weight = loaded_weight.narrow(
output_dim, shard_rank * shard_size, shard_size
)
assert param_data.shape == loaded_weight.shape
param_data.copy_(loaded_weight)
# --8<-- [start:row_parallel_linear]
@PluggableLayer.register("row_parallel_linear")
class RowParallelLinear(LinearBase):
vllm/model_executor/layers/quantization/__init__.py
+9 -8
View File
@@ -163,17 +163,18 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
"deepseek_v4_fp8": DeepseekV4FP8Config,
"humming": HummingConfig,
"online": OnlineQuantizationConfig,
# MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the
# ModelOpt MXFP8 config (same format). The "mxfp8" online shorthand
# below only applies to the `--quantization mxfp8` CLI path.
"mxfp8": ModelOptMxFp8Config,
}
# Register online shorthands as quantization methods so the user can
# specify "LLM(..., quantization='fp8_per_tensor')" as shorthand for
# creating a more complicated online quant config object.
# Register online shorthands (e.g. "fp8_per_tensor") as quant methods.
# setdefault so a shorthand that is also a checkpoint method (e.g. "mxfp8")
# keeps its checkpoint config; the shorthand still works via the
# `--quantization` CLI path in `resolve_quantization_config`.
for shorthand in _ONLINE_SHORTHANDS:
assert shorthand not in method_to_config, (
f"Online quant shorthand {shorthand!r} conflicts with an "
f"existing quantization method"
)
method_to_config[shorthand] = OnlineQuantizationConfig
method_to_config.setdefault(shorthand, OnlineQuantizationConfig)
# Update the `method_to_config` with customized quantization methods.
method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
vllm/model_executor/layers/quantization/modelopt.py
+77 -1
View File
@@ -2,11 +2,12 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from fnmatch import fnmatch
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, cast
import torch
from torch.nn.parameter import Parameter
import vllm.envs as envs
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.config import get_current_vllm_config
from vllm.logger import init_logger
@@ -27,6 +28,7 @@ from vllm.model_executor.layers.fused_moe import (
SharedExperts,
)
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
Fp8MoeBackend,
convert_to_fp8_moe_kernel_format,
make_fp8_moe_kernel,
make_fp8_moe_quant_config,
@@ -1720,6 +1722,22 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
return "modelopt_mxfp8"
return None
@classmethod
def from_config(cls, config: dict[str, Any]) -> "ModelOptMxFp8Config":
# MiniMax-style checkpoints tag `quant_method: "mxfp8"` + `ignored_layers`
# (same on-disk format as ModelOpt MXFP8); normalize to the ModelOpt
# schema and reuse the shared parser.
if "quantization" not in config and not config.get("quant_algo"):
config = {
"quant_method": "modelopt",
"quantization": {
"quant_algo": "MXFP8",
"kv_cache_quant_algo": config.get("kv_cache_quant_algo"),
"exclude_modules": config.get("ignored_layers", []) or [],
},
}
return cast("ModelOptMxFp8Config", super().from_config(config))
@classmethod
def _from_config(
cls,
@@ -1823,6 +1841,12 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
layer.register_parameter("weight_scale", weight_scale)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# Idempotent: the emulation kernel may dequant the weight to BF16 at load
# time (>=2-byte). If already converted, there is nothing left to do --
# avoid re-running the MXFP8-only validation/conversion below.
if layer.weight.element_size() >= 2:
return
# Validate weight tensor
if layer.weight.ndim != 2:
raise ValueError(
@@ -2065,6 +2089,44 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
torch.stack(w2_scale_shuffled).contiguous(),
)
def _dequant_mxfp8_weights_to_bf16(self, layer: RoutedExperts) -> None:
"""One-time MXFP8->BF16 weight dequant for the emulation path.
On devices without a native MXFP8 MoE kernel (e.g. gfx942 / MI300),
``Mxfp8EmulationTritonExperts`` otherwise dequantizes every expert
weight to BF16 on *every* forward step -- the dominant cost (conc1
~1.3 tok/s). Doing the dequant once here and replacing the MXFP8
parameters with BF16 makes the MoE run exactly like a plain BF16
checkpoint (full precision, no per-step dequant); SwiGLU-OAI is still
applied by the experts' ``activation()`` override. The MXFP8 weights
are freed by ``replace_parameter`` (BF16 is 2x their size; the small
E8M0 scale tensors are left in place, unused).
"""
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
dequant_mxfp8_to_bf16,
)
target_dtype = getattr(layer, "orig_dtype", torch.bfloat16)
num_experts = layer.w13_weight.shape[0]
# dequant_mxfp8_to_bf16 handles arbitrary leading dims (*x.shape[:-1]),
# so dequant the whole [E, N, K] weight in one vectorized call.
w13_bf16 = dequant_mxfp8_to_bf16(layer.w13_weight, layer.w13_weight_scale).to(
target_dtype
)
w2_bf16 = dequant_mxfp8_to_bf16(layer.w2_weight, layer.w2_weight_scale).to(
target_dtype
)
replace_parameter(layer, "w13_weight", w13_bf16)
replace_parameter(layer, "w2_weight", w2_bf16)
logger.info_once(
"MXFP8->BF16 load-time dequant complete (%d experts/layer); MoE "
"now runs in BF16 with no per-step dequant.",
num_experts,
)
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
# TODO(bnell): why is this required only for mxfp8?
if getattr(layer, "_already_called_process_weights_after_loading", False):
@@ -2102,6 +2164,17 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
routing_tables=layer._expert_routing_tables(),
)
# No native MXFP8 MoE kernel on this device (e.g. gfx942): the emulation
# experts would dequant MXFP8->BF16 every forward step. Convert the
# weights to BF16 once, here, so the MoE runs like a BF16 checkpoint.
# Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the 1-byte
# MXFP8 weights and dequant per-step (~half the memory, much slower).
if (
self.mxfp8_backend == Fp8MoeBackend.EMULATION
and envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD
):
self._dequant_mxfp8_weights_to_bf16(layer)
def maybe_make_prepare_finalize(
self,
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
@@ -2131,6 +2204,9 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
a1_scale=None,
a2_scale=None,
block_shape=self.weight_block_size,
swiglu_limit=getattr(layer, "swiglu_limit", None),
gemm1_alpha=getattr(layer, "swiglu_alpha", None),
gemm1_beta=getattr(layer, "swiglu_beta", None),
)
def apply_monolithic(
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+128 -70
View File
@@ -159,75 +159,104 @@ def _silu_mul_quant_fp8_packed_kernel(
output_q_stride_m,
output_scale_stride_k,
clamp_limit,
alpha,
beta,
N: tl.constexpr,
NUM_GROUPS: tl.constexpr,
GROUPS_PER_ROW: tl.constexpr,
PACKS_PER_ROW: tl.constexpr,
fp8_min: tl.constexpr,
fp8_max: tl.constexpr,
GROUP_SIZE: tl.constexpr,
PACKS_PER_CTA: tl.constexpr,
BLOCK_M: tl.constexpr,
HAS_CLAMP: tl.constexpr,
):
N_2: tl.constexpr = N // 2
GROUPS_PER_PACK: tl.constexpr = 4
hidden_size: tl.constexpr = N // 2
pid_pack = tl.program_id(0)
pid_m = tl.program_id(1)
m_offset = pid_m.to(tl.int64) * BLOCK_M
pack_tile = tl.program_id(0)
row_start = tl.program_id(1).to(tl.int64) * BLOCK_M
row_step = tl.num_programs(1).to(tl.int64) * BLOCK_M
if m_offset >= M:
return
groups_per_cta: tl.constexpr = PACKS_PER_CTA * GROUPS_PER_PACK
elems_per_cta: tl.constexpr = groups_per_cta * GROUP_SIZE
col_start = pack_tile * elems_per_cta
col_offsets = tl.arange(0, elems_per_cta)
row_offsets = tl.arange(0, BLOCK_M)
pack_offsets = tl.arange(0, PACKS_PER_CTA)
offs_m = tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, GROUP_SIZE)
row_mask = (m_offset + offs_m) < M
col_mask = (col_start + col_offsets) < (GROUPS_PER_ROW * GROUP_SIZE)
base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m
base_out_offset = (m_offset + offs_m[:, None]) * output_q_stride_m
# persistent with grid_m-stride loop
while row_start < M:
rows = row_start + row_offsets
row_mask = rows < M
input_row_start = rows[:, None] * input_stride_m
output_row_start = rows[:, None] * output_q_stride_m
packed_scale = tl.zeros((BLOCK_M,), dtype=tl.int32)
gate_flat = tl.load(
input_ptr + input_row_start + col_start + col_offsets[None, :],
mask=row_mask[:, None] & col_mask[None, :],
other=0.0,
)
up_flat = tl.load(
input_ptr
+ input_row_start
+ hidden_size
+ col_start
+ col_offsets[None, :],
mask=row_mask[:, None] & col_mask[None, :],
other=0.0,
)
for pack_idx in tl.static_range(4):
group_id = pid_pack * 4 + pack_idx
gate = tl.reshape(gate_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(
tl.float32
)
up = tl.reshape(up_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(tl.float32)
if group_id < NUM_GROUPS:
n_offset = group_id * GROUP_SIZE
if HAS_CLAMP:
gate = tl.minimum(gate, clamp_limit)
up = tl.clamp(up, -clamp_limit, clamp_limit)
act_ptrs = input_ptr + base_row_offset + n_offset + offs_n[None, :]
act_in = tl.load(act_ptrs, mask=row_mask[:, None], other=0.0)
# Unified gated activation: silu == swigluoai with alpha=1, beta=0.
# glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
glu = gate / (1.0 + tl.exp(-gate * alpha))
y = (up + beta) * glu
# Round through bf16 to match unfused precision path
y = y.to(tl.bfloat16).to(tl.float32)
mul_ptrs = act_ptrs + N_2
mul_in = tl.load(mul_ptrs, mask=row_mask[:, None], other=0.0)
absmax = tl.max(tl.abs(y), axis=2)
scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
exponent = tl.ceil(tl.log2(scale_raw))
scale = tl.math.exp2(exponent)
act_f32 = act_in.to(tl.float32)
mul_f32 = mul_in.to(tl.float32)
y_q = tl.clamp(y / scale[:, :, None], fp8_min, fp8_max)
if HAS_CLAMP:
act_f32 = tl.minimum(act_f32, clamp_limit)
mul_f32 = tl.clamp(mul_f32, -clamp_limit, clamp_limit)
y_q_flat = tl.reshape(y_q, (BLOCK_M, elems_per_cta))
tl.store(
output_q_ptr + output_row_start + col_start + col_offsets[None, :],
y_q_flat.to(output_q_ptr.dtype.element_ty),
mask=row_mask[:, None] & col_mask[None, :],
)
y = (act_f32 / (1.0 + tl.exp(-act_f32))) * mul_f32
# Round through bf16 to match unfused precision path
y = y.to(tl.bfloat16).to(tl.float32)
scale_byte = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
scale_bytes = tl.reshape(scale_byte, (BLOCK_M, PACKS_PER_CTA, GROUPS_PER_PACK))
shifts = tl.arange(0, GROUPS_PER_PACK) * 8
packed_scale = tl.sum(scale_bytes << shifts[None, None, :], axis=2)
absmax = tl.max(tl.abs(y), axis=1)
scale_pack = pack_tile * PACKS_PER_CTA + pack_offsets
scale_ptrs = (
output_scale_ptr
+ scale_pack[None, :] * output_scale_stride_k
+ rows[:, None]
)
tl.store(
scale_ptrs,
packed_scale,
mask=row_mask[:, None] & (scale_pack[None, :] < PACKS_PER_ROW),
)
scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
exponent = tl.ceil(tl.log2(scale_raw))
scale = tl.math.exp2(exponent)
y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max)
out_q_ptrs = output_q_ptr + base_out_offset + n_offset + offs_n[None, :]
tl.store(
out_q_ptrs,
y_q.to(output_q_ptr.dtype.element_ty),
mask=row_mask[:, None],
)
exponent_biased = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
packed_scale = packed_scale | (exponent_biased << (pack_idx * 8))
scale_ptrs = output_scale_ptr + pid_pack * output_scale_stride_k + m_offset + offs_m
tl.store(scale_ptrs, packed_scale, mask=row_mask)
row_start += row_step
def silu_mul_quant_fp8_packed_triton(
@@ -235,37 +264,48 @@ def silu_mul_quant_fp8_packed_triton(
group_size: int = 128,
output_q: torch.Tensor | None = None,
clamp_limit: float | None = None,
alpha: float = 1.0,
beta: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor]:
assert input.dim() == 2
assert input.is_contiguous()
M, N = input.shape
N_2 = N // 2
hidden_size = N // 2
assert N_2 % group_size == 0
assert hidden_size % group_size == 0
fp8_dtype = torch.float8_e4m3fn
finfo = torch.finfo(fp8_dtype)
fp8_min, fp8_max = finfo.min, finfo.max
num_groups_per_row = N_2 // group_size
num_packed_groups = (num_groups_per_row + 3) // 4
tma_aligned_M = ((M + 3) // 4) * 4
groups_per_row = hidden_size // group_size
groups_per_pack = 4 # pack 4 UE8M0 scales to a single INT32
packs_per_row = triton.cdiv(groups_per_row, groups_per_pack)
if output_q is None:
output_q = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device)
output_q = torch.empty((M, hidden_size), dtype=fp8_dtype, device=input.device)
aligned_m = triton.cdiv(M, 4) * 4
output_scale_packed = torch.empty(
(num_packed_groups, tma_aligned_M),
(packs_per_row, aligned_m),
dtype=torch.int32,
device=input.device,
).T[:M, :]
BLOCK_M = 8
grid = (num_packed_groups, (M + BLOCK_M - 1) // BLOCK_M)
num_warps = max(4, group_size // 32)
# Tuned for group_size=32 (MXFP8) and group_size=128 (DeepSeek-V4)
num_warps = 4
num_stages = 2
if group_size < 128:
BM = 1
packs_per_cta = 8
else:
BM = 1 if M < 512 else 4
packs_per_cta = 2 if M < 512 else 1
grid_n = triton.cdiv(packs_per_row, packs_per_cta)
grid_m = min(triton.cdiv(M, BM), 4096)
grid = (grid_n, grid_m)
has_clamp = clamp_limit is not None
_silu_mul_quant_fp8_packed_kernel[grid](
@@ -277,12 +317,16 @@ def silu_mul_quant_fp8_packed_triton(
output_q.stride(0),
output_scale_packed.stride(1),
clamp_limit if has_clamp else 0.0,
alpha,
beta,
N=N,
NUM_GROUPS=num_groups_per_row,
GROUPS_PER_ROW=groups_per_row,
PACKS_PER_ROW=packs_per_row,
fp8_min=fp8_min,
fp8_max=fp8_max,
GROUP_SIZE=group_size,
BLOCK_M=BLOCK_M,
PACKS_PER_CTA=packs_per_cta,
BLOCK_M=BM,
HAS_CLAMP=has_clamp,
num_warps=num_warps,
num_stages=num_stages,
@@ -303,6 +347,8 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
# Information for float8
eps,
clamp_limit,
alpha,
beta,
fp8_min: tl.constexpr,
fp8_max: tl.constexpr,
use_ue8m0: tl.constexpr,
@@ -348,10 +394,14 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to(
y_ptr.dtype.element_ty
)
# Unified gated activation: silu == swigluoai with alpha=1, beta=0.
# glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
# Keep glu/up at input precision (narrow before the mul) so the alpha=1,
# beta=0 defaults match the C++ silu_and_mul path bit-for-bit.
act_in = act_in.to(tl.float32)
one_f32 = tl.cast(1, tl.float32)
silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty)
y = (silu_out * mul_in).to(tl.float32)
glu = (act_in / (1.0 + tl.exp(-act_in * alpha))).to(y_ptr.dtype.element_ty)
up = (mul_in.to(tl.float32) + beta).to(y_ptr.dtype.element_ty)
y = (glu * up).to(tl.float32)
# quant
_absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps)
@@ -379,11 +429,15 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
use_ue8m0: bool | None = None,
eps: float = 1e-10,
clamp_limit: float | None = None,
group_size: int = 128,
alpha: float = 1.0,
beta: float = 0.0,
):
"""
silu+mul + block-fp8 quant with group size 128.
Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate
(silu: alpha=1, beta=0; swigluoai: alpha, beta from config).
"""
GROUP_SIZE = 128
GROUP_SIZE = group_size
assert input.ndim == 2
if output is not None:
assert output.ndim == 2
@@ -431,6 +485,8 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
output_scales.stride(-1),
eps,
clamp_limit if has_clamp else 0.0,
alpha,
beta,
fp8_min,
fp8_max,
use_ue8m0,
@@ -1015,9 +1071,10 @@ def deepgemm_post_process_fp8_weight_block(
f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
)
if ws.dtype == torch.float8_e8m0fnu:
# Scales already in E8M0 from checkpoint — upcast to fp32
# and skip requantization (weights already have power-of-two scales).
if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
# Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
# bits as uint8 for MXFP8) — upcast to fp32 and skip requantization
# (weights already have power-of-two scales).
ws = _upcast_e8m0_to_fp32(ws)
else:
assert ws.dtype == torch.float32, (
@@ -1057,7 +1114,8 @@ def deepgemm_post_process_fp8_weight_block(
ws = ws.unsqueeze(0)
# From https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/utils/layout.hpp#L46
recipe = (1, 128, 128)
# (1, block_n, block_k): (1, 128, 128) for FP8 block, (1, 1, 32) for MXFP8.
recipe = (1, quant_block_shape[0], quant_block_shape[1])
# Ref : https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/gemm.hpp
# DeepGemm uses the `transform_sf_into_required_layout` function to
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+97
View File
@@ -84,6 +84,92 @@ def _mxfp8_e4m3_quantize_torch(
return x_fp8, scales_uint8
def _mxfp8_quant_triton_kernel():
"""Lazily-built Triton kernel: per-32-block E8M0 scale + FP8-E4M3 quant.
Fuses what ``_mxfp8_e4m3_quantize_torch`` does in several elementwise passes
into one launch. Each program handles ``[BLOCK_M, 32]`` (one MX block).
"""
from vllm.triton_utils import tl, triton
@triton.jit
def _kernel(
x_ptr,
xq_ptr,
s_ptr,
M,
K,
sxm,
sxk,
sqm,
sqk,
ssm,
ssk,
BLOCK_M: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_b = tl.program_id(1) # which 32-element block along K
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_k = pid_b * 32 + tl.arange(0, 32)
m_mask = offs_m < M
x = tl.load(
x_ptr + offs_m[:, None] * sxm + offs_k[None, :] * sxk,
mask=m_mask[:, None],
other=0.0,
).to(tl.float32)
amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-30) # [BLOCK_M]
sb = tl.floor(tl.log2(amax)) + 127.0
sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
descale = tl.exp2(sb - 127.0)
xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
tl.store(
xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk,
xq,
mask=m_mask[:, None],
)
tl.store(s_ptr + offs_m * ssm + pid_b * ssk, sb.to(tl.uint8), mask=m_mask)
return _kernel
_MXFP8_QUANT_KERNEL = None
def _mxfp8_e4m3_quantize_triton(
x: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Fused 2D MXFP8 quant (non-swizzled, row-major [M, K//32] scales)."""
from vllm.triton_utils import triton
global _MXFP8_QUANT_KERNEL
if _MXFP8_QUANT_KERNEL is None:
_MXFP8_QUANT_KERNEL = _mxfp8_quant_triton_kernel()
M, K = x.shape
x = x.contiguous()
xq = torch.empty((M, K), dtype=MXFP8_VALUE_DTYPE, device=x.device)
scales = torch.empty(
(M, K // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device
)
BLOCK_M = 64
grid = (triton.cdiv(M, BLOCK_M), K // MXFP8_BLOCK_SIZE)
_MXFP8_QUANT_KERNEL[grid](
x,
xq,
scales,
M,
K,
x.stride(0),
x.stride(1),
xq.stride(0),
xq.stride(1),
scales.stride(0),
scales.stride(1),
BLOCK_M=BLOCK_M,
)
return xq, scales
def _mxfp8_e4m3_quantize_impl(
x: torch.Tensor,
is_sf_swizzled_layout: bool = False,
@@ -103,6 +189,17 @@ def _mxfp8_e4m3_quantize_impl(
x_scales = x_scales.view(x.size(0), -1)
return x_q, x_scales
# ROCm: a single fused Triton kernel beats the multi-pass torch path for the
# common 2D, non-swizzled activation-quant case (used by the native MX
# linear/MoE). Falls back to torch otherwise (3D weights, swizzled layout).
if (
current_platform.is_rocm()
and not is_sf_swizzled_layout
and x.ndim == 2
and x.shape[-1] % MXFP8_BLOCK_SIZE == 0
):
return _mxfp8_e4m3_quantize_triton(x)
return _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout)
vllm/model_executor/models/registry.py
+9
View File
@@ -164,6 +164,10 @@ _TEXT_GENERATION_MODELS = {
"MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
"MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
"MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
"MiniMaxM3SparseForCausalLM": (
"vllm.models.minimax_m3",
"MiniMaxM3SparseForCausalLM",
),
"Ministral3ForCausalLM": ("mistral", "MistralForCausalLM"),
"MistralForCausalLM": ("mistral", "MistralForCausalLM"),
"MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
@@ -483,6 +487,10 @@ _MULTIMODAL_MODELS = {
"MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),
"MiDashengLMModel": ("midashenglm", "MiDashengLMModel"),
"MiMoV2OmniForCausalLM": ("mimo_v2_omni", "MiMoV2OmniForCausalLM"),
"MiniMaxM3SparseForConditionalGeneration": (
"vllm.models.minimax_m3",
"MiniMaxM3SparseForConditionalGeneration",
),
"MiniMaxVL01ForConditionalGeneration": (
"minimax_vl_01",
"MiniMaxVL01ForConditionalGeneration",
@@ -620,6 +628,7 @@ _SPECULATIVE_DECODING_MODELS = {
"EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
"DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
"MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
"Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"),
"ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
"ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
vllm/model_executor/warmup/kernel_warmup.py
+6
View File
@@ -53,6 +53,10 @@ def _resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
def kernel_warmup(worker: "Worker"):
from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
minimax_m3_msa_warmup,
)
# Deep GEMM warmup
do_deep_gemm_warmup = (
envs.VLLM_USE_DEEP_GEMM
@@ -64,6 +68,8 @@ def kernel_warmup(worker: "Worker"):
max_tokens = worker.scheduler_config.max_num_batched_tokens
deep_gemm_warmup(model, max_tokens)
minimax_m3_msa_warmup(worker)
enable_flashinfer_autotune = (
worker.vllm_config.kernel_config.enable_flashinfer_autotune
)
vllm/model_executor/warmup/minimax_m3_msa_warmup.py
+43
View File
@@ -0,0 +1,43 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import TYPE_CHECKING
from vllm.logger import init_logger
from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
from vllm.platforms import current_platform
from vllm.tracing import instrument
if TYPE_CHECKING:
from vllm.v1.worker.gpu_worker import Worker
logger = init_logger(__name__)
@instrument(span_name="MiniMax M3 MSA warmup")
def minimax_m3_msa_warmup(worker: "Worker") -> None:
sparse_module = next(
(
module
for module in worker.get_model().modules()
if isinstance(module, MiniMaxM3SparseAttention)
),
None,
)
if sparse_module is None:
return
if not (
current_platform.is_cuda() and current_platform.is_device_capability_family(100)
):
return
logger.info("Warming up MiniMax M3 MSA kernels.")
# Cover sparse prefill through the normal model path.
worker.model_runner._dummy_run(
num_tokens=16,
skip_eplb=True,
is_profile=True,
force_attention=True,
create_mixed_batch=True,
)
vllm/models/minimax_m3/__init__.py
+33
View File
@@ -0,0 +1,33 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MiniMax M3 model — hardware-isolated entry point.
The implementation lives under ``nvidia/`` and ``amd/``; this module picks the
right one for the current platform and re-exports the public classes used by
the model registry. (Mirrors ``vllm.models.deepseek_v4``.)
"""
from typing import TYPE_CHECKING
from vllm.platforms import current_platform
# The NVIDIA branch is the static default that type-checkers see; the ROCm
# branch overrides it at runtime (kept type-compatible via type: ignore).
if TYPE_CHECKING or not current_platform.is_rocm():
from .nvidia.model import (
MiniMaxM3SparseForCausalLM,
MiniMaxM3SparseForConditionalGeneration,
)
from .nvidia.mtp import MiniMaxM3MTP
else:
from .amd.model import ( # type: ignore[assignment]
MiniMaxM3SparseForCausalLM,
MiniMaxM3SparseForConditionalGeneration,
)
from .amd.mtp import MiniMaxM3MTP # type: ignore[assignment]
__all__ = [
"MiniMaxM3MTP",
"MiniMaxM3SparseForCausalLM",
"MiniMaxM3SparseForConditionalGeneration",
]
vllm/models/minimax_m3/amd/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/minimax_m3/amd/model.py
+1216
View File
File diff suppressed because it is too large Load Diff
vllm/models/minimax_m3/amd/mtp.py
+330
View File
@@ -0,0 +1,330 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MiniMax M3 MTP (multi-token prediction) draft model -- ROCm/AMD variant.
Byte-identical to ``nvidia/mtp.py`` except this file lives under ``amd/`` so its
``from .model import ...`` resolves to the self-contained AMD model (native Gemma
RMSNorm, native MXFP8 MoE, Triton sparse attention). The MTP logic is
platform-agnostic. (Mirrors ``vllm.models.deepseek_v4.amd.mtp``.)
TODO(future, separate diff): since this is byte-identical to ``nvidia/mtp.py``,
both copies could be consolidated into a single ``common/mtp.py`` that dispatches
its model import (``..amd.model`` vs ``..nvidia.model``) via
``current_platform.is_rocm()`` -- the same dispatch ``minimax_m3/__init__.py``
uses. This was prototyped and VERIFIED working (``MiniMaxM3MTP`` resolves through
``common.mtp`` to the AMD decoder layer / RMSNorm on ROCm), but it deletes the
upstream ``nvidia/mtp.py`` and touches the NVIDIA load path, so it is deferred to
a dedicated refactor diff to keep this AMD-enablement change NVIDIA-untouched.
"""
from collections.abc import Iterable
import regex as re
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.model_executor.layers.fused_moe import (
fused_moe_make_expert_params_mapping,
)
from vllm.model_executor.layers.linear import (
ReplicatedLinear,
)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader,
maybe_remap_kv_scale_name,
)
from vllm.model_executor.models.utils import (
maybe_prefix,
)
from vllm.sequence import IntermediateTensors
from .model import (
MiniMAXGemmaRMSNorm,
MiniMaxM3DecoderLayer,
)
class MiniMaxM3MultiTokenPredictorLayer(nn.Module):
def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
super().__init__()
assert vllm_config.speculative_config is not None
config = vllm_config.speculative_config.draft_model_config.hf_config
cache_config = vllm_config.cache_config
quant_config = vllm_config.quant_config
self.enorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.hnorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.eh_proj = ReplicatedLinear(
config.hidden_size * 2,
config.hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.eh_proj",
)
self.transformer_layer = MiniMaxM3DecoderLayer(
config=config,
prefix=prefix,
cache_config=cache_config,
quant_config=quant_config,
force_sparse_attn=True,
force_moe=True,
)
self.final_layernorm = MiniMAXGemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
previous_hidden_states: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
spec_step_index: int = 0,
) -> torch.Tensor:
assert inputs_embeds is not None
# Mask out inputs at position 0, as not needed by MTP.
inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
# Combine the normalized token embeddings with the normalized
# previous hidden states.
inputs_embeds = self.enorm(inputs_embeds)
previous_hidden_states = self.hnorm(previous_hidden_states)
hidden_states, _ = self.eh_proj(
torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
)
# Apply transformer layer.
hidden_states, residual = self.transformer_layer(
positions=positions,
hidden_states=hidden_states,
residual=None,
)
hidden_states += residual
return hidden_states
class MiniMaxM3MultiTokenPredictor(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
assert vllm_config.speculative_config is not None
# Use the draft (MTP) config, not the target model's. This is flat for a
# standalone checkpoint, and the promoted text_config for a bundled one.
config = vllm_config.speculative_config.draft_model_config.hf_config
self.num_mtp_layers = config.num_mtp_modules
self.layers = torch.nn.ModuleDict(
{
str(idx): MiniMaxM3MultiTokenPredictorLayer(
vllm_config, f"{prefix}.layers.{idx}"
)
for idx in range(self.num_mtp_layers)
}
)
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
prefix=maybe_prefix(prefix, "embed_tokens"),
)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
previous_hidden_states: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> torch.Tensor:
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
current_step_idx = spec_step_idx % self.num_mtp_layers
return self.layers[str(current_step_idx)](
input_ids,
positions,
previous_hidden_states,
inputs_embeds,
current_step_idx,
)
class MiniMaxM3MTP(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
assert vllm_config.speculative_config is not None
self.config = vllm_config.speculative_config.draft_model_config.hf_config
self.quant_config = vllm_config.quant_config
self.model = MiniMaxM3MultiTokenPredictor(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
self.lm_head = ParallelLMHead(
self.config.vocab_size,
self.config.hidden_size,
quant_config=self.quant_config,
prefix=maybe_prefix(prefix, "lm_head"),
)
self.logits_processor = LogitsProcessor(self.config.vocab_size)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
hidden_states: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> torch.Tensor:
return self.model(
input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
)
def compute_logits(
self,
hidden_states: torch.Tensor,
spec_step_idx: int = 0,
) -> torch.Tensor | None:
current_step_idx = spec_step_idx % self.model.num_mtp_layers
mtp_layer = self.model.layers[str(current_step_idx)]
return self.logits_processor(
self.lm_head, mtp_layer.final_layernorm(hidden_states)
)
def _get_mtp_layer_idx_from_weight_name(self, name: str) -> int | None:
"""Return the MTP layer index in *.mtp.layers.{idx}.*, else None."""
match = re.search(r"\.mtp\.layers\.(\d+)\.", name)
return int(match.group(1)) if match else None
def _map_checkpoint_name(self, name: str) -> str | None:
"""Map a full checkpoint key to this MTP module's parameter name.
The MTP module only owns the *.mtp.layers.* weights plus the token
embedding and LM head, which the checkpoint shares with the main model.
Everything else belongs to other modules and is ignored here by returning
None.
"""
# In the bundled checkpoint, the MTP weights are prefixed with
# "language_model". The standalone MTP checkpoint has no such prefix.
# Strip it if present.
name = name.removeprefix("language_model.")
if name == "model.embed_tokens.weight":
return "model.embed_tokens.weight"
if name == "lm_head.weight":
return "lm_head.weight"
if "model.mtp.layers" in name:
if "weight_scale_inv" in name:
# The checkpoint stores block scales as "weight_scale_inv".
# The ModelOpt MXFP8 layers expose them as "weight_scale".
name = name.replace("weight_scale_inv", "weight_scale")
# Strip "mtp" from prefix.
return name.replace(".mtp.", ".")
return None
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Map q/k/v projections to qkv_proj, and gate/up projections to gate_up_proj.
stacked_params_mapping: list[tuple[str, str, int | str]] = [
# (param_name, shard_name, shard_id)
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".qkv_proj", ".index_q_proj", "index_q"),
(".qkv_proj", ".index_k_proj", "index_k"),
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
]
# Map expert weights w1/w2/w3 to gate/down/up.
# (param_name, weight_name, expert_id, shard_id)
expert_params_mapping = fused_moe_make_expert_params_mapping(
self,
ckpt_gate_proj_name="w1",
ckpt_down_proj_name="w2",
ckpt_up_proj_name="w3",
num_experts=self.config.num_local_experts,
)
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
loaded_mtp_layers: set[int] = set()
for name, loaded_weight in weights:
mtp_layer = self._get_mtp_layer_idx_from_weight_name(name)
mapped_name = self._map_checkpoint_name(name)
if mapped_name is None:
# This weight does not belong to the MTP module, so skip it.
continue
name = mapped_name
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
# Routed experts (w1/w2/w3) are handled below. Don't let the
# stacked mapping rewrite them.
if ("block_sparse_moe.experts." in name) and name not in params_dict:
continue
name = name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
for (
param_name,
weight_name,
expert_id,
expert_shard_id,
) in expert_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(
param,
loaded_weight,
name,
shard_id=expert_shard_id,
expert_id=expert_id,
)
break
else:
remapped_name = maybe_remap_kv_scale_name(name, params_dict)
if remapped_name is None or remapped_name not in params_dict:
continue
name = remapped_name
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
if mtp_layer is not None:
loaded_mtp_layers.add(mtp_layer)
# Validate that weights were loaded for each MTP layer.
for layer_idx in range(self.model.num_mtp_layers):
if layer_idx not in loaded_mtp_layers:
raise ValueError(
f"Failed to load MTP layer {layer_idx} weights from checkpoint."
)
return loaded_params
vllm/models/minimax_m3/amd/ops/__init__.py
+24
View File
@@ -0,0 +1,24 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""AMD/ROCm fused Triton ops for MiniMax-M3.
These replace per-element PyTorch fallbacks (FlashInfer / fused HIP kernels are
unavailable on ROCm) with single-pass Triton kernels to cut launch overhead and
intermediate-tensor traffic during decode.
"""
from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import (
gemma_fused_add_rmsnorm,
gemma_rmsnorm,
)
from vllm.models.minimax_m3.amd.ops.swiglu_oai import (
swiglu_oai_quantize_mxfp8,
swiglu_oai_split,
)
__all__ = [
"gemma_rmsnorm",
"gemma_fused_add_rmsnorm",
"swiglu_oai_split",
"swiglu_oai_quantize_mxfp8",
]
vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py
+155
View File
@@ -0,0 +1,155 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Fused Gemma-style RMSNorm for AMD ROCm via Triton.
Gemma RMSNorm = normalize(x) * (1 + weight), computed in fp32. FlashInfer's
``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` CUDA kernels are unavailable on
ROCm, so the AMD path previously used a ~8-op PyTorch sequence (float cast, add,
pow, mean, rsqrt, two muls, cast) — each a separate kernel launch materializing
fp32 intermediates. These kernels collapse that into a single pass per row.
Two entry points:
* ``gemma_rmsnorm(x, w, eps)`` -> normalized tensor
* ``gemma_fused_add_rmsnorm(x, res, w, eps)`` -> (normalized, x + res)
Both normalize over the last dim and broadcast ``weight`` (shape [N]) over it,
so they serve both the full-hidden norms (input/post-attn/final) and the
per-head q_norm/k_norm (N == head_dim). Inputs may be non-contiguous views
(e.g. ``qkv.split`` slices); strides are passed through and outputs are written
contiguous.
"""
import torch
from vllm.triton_utils import tl, triton
@triton.jit
def _gemma_rmsnorm_kernel(
x_ptr,
w_ptr,
out_ptr,
n_cols,
stride_row,
stride_col,
eps,
BLOCK_N: tl.constexpr,
):
row = tl.program_id(0)
cols = tl.arange(0, BLOCK_N)
mask = cols < n_cols
x = tl.load(x_ptr + row * stride_row + cols * stride_col, mask=mask, other=0.0).to(
tl.float32
)
var = tl.sum(x * x, axis=0) / n_cols
rstd = 1.0 / tl.sqrt(var + eps)
w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
out = x * rstd * (1.0 + w)
tl.store(
out_ptr + row * n_cols + cols,
out.to(out_ptr.dtype.element_ty),
mask=mask,
)
@triton.jit
def _gemma_fused_add_rmsnorm_kernel(
x_ptr,
res_ptr,
w_ptr,
out_ptr,
res_out_ptr,
n_cols,
stride_xrow,
stride_xcol,
stride_rrow,
stride_rcol,
eps,
BLOCK_N: tl.constexpr,
):
row = tl.program_id(0)
cols = tl.arange(0, BLOCK_N)
mask = cols < n_cols
x = tl.load(
x_ptr + row * stride_xrow + cols * stride_xcol, mask=mask, other=0.0
).to(tl.float32)
r = tl.load(
res_ptr + row * stride_rrow + cols * stride_rcol, mask=mask, other=0.0
).to(tl.float32)
s = x + r
# residual_out is the pre-norm sum (consumed by the next layer's add).
tl.store(
res_out_ptr + row * n_cols + cols,
s.to(res_out_ptr.dtype.element_ty),
mask=mask,
)
var = tl.sum(s * s, axis=0) / n_cols
rstd = 1.0 / tl.sqrt(var + eps)
w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
out = s * rstd * (1.0 + w)
tl.store(
out_ptr + row * n_cols + cols,
out.to(out_ptr.dtype.element_ty),
mask=mask,
)
def _num_warps(block_n: int) -> int:
if block_n >= 4096:
return 16
if block_n >= 1024:
return 8
return 4
def gemma_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
orig_shape = x.shape
n = orig_shape[-1]
x2 = x.reshape(-1, n)
m = x2.shape[0]
out = torch.empty((m, n), dtype=x.dtype, device=x.device)
block_n = triton.next_power_of_2(n)
_gemma_rmsnorm_kernel[(m,)](
x2,
weight,
out,
n,
x2.stride(0),
x2.stride(1),
eps,
BLOCK_N=block_n,
num_warps=_num_warps(block_n),
)
return out.reshape(orig_shape)
def gemma_fused_add_rmsnorm(
x: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
orig_shape = x.shape
n = orig_shape[-1]
x2 = x.reshape(-1, n)
r2 = residual.reshape(-1, n)
m = x2.shape[0]
out = torch.empty((m, n), dtype=x.dtype, device=x.device)
res_out = torch.empty((m, n), dtype=x.dtype, device=x.device)
block_n = triton.next_power_of_2(n)
_gemma_fused_add_rmsnorm_kernel[(m,)](
x2,
r2,
weight,
out,
res_out,
n,
x2.stride(0),
x2.stride(1),
r2.stride(0),
r2.stride(1),
eps,
BLOCK_N=block_n,
num_warps=_num_warps(block_n),
)
return out.reshape(orig_shape), res_out.reshape(orig_shape)
vllm/models/minimax_m3/amd/ops/swiglu_oai.py
+221
View File
@@ -0,0 +1,221 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Fused SwiGLU-OAI activation (split layout) for AMD ROCm via Triton.
SwiGLU-OAI on a ``[*, 2I]`` split-layout input (gate = first half, up = second
half):
gate = clamp(gate, max=limit)
up = clamp(up, -limit, +limit)
out = gate * sigmoid(alpha * gate) * (up + beta)
On ROCm the dense MLP and the native MXFP8 MoE (between its two GEMMs) fell back
to a chain of elementwise PyTorch ops with fp32 intermediates: vLLM's shared
``SiluAndMulWithClamp`` blanket-routes ROCm to ``forward_native``, and the MoE
applies the activation inline in PyTorch. This Triton kernel collapses that into
a single pass producing the ``[*, I]`` output directly, and computes in fp32
(rel ~1e-6 vs reference).
Note: the vectorized ``torch.ops._C.silu_and_mul_with_clamp`` op IS built on
ROCm and is ~1.2-2.2x faster in isolation, but the win is launch overhead that
HIP graphs already eliminate — measured end-to-end throughput is identical
(within noise), so we keep the fp32-accurate Triton kernel.
"""
import torch
from vllm.triton_utils import tl, triton
@triton.jit
def _swiglu_oai_kernel(
g_ptr,
out_ptr,
n_inter,
stride_gm,
stride_gn,
stride_om,
stride_on,
alpha,
beta,
limit,
HAS_LIMIT: tl.constexpr,
BLOCK_I: tl.constexpr,
):
row = tl.program_id(0)
pid_i = tl.program_id(1)
cols = pid_i * BLOCK_I + tl.arange(0, BLOCK_I)
mask = cols < n_inter
gate = tl.load(g_ptr + row * stride_gm + cols * stride_gn, mask=mask, other=0.0).to(
tl.float32
)
up = tl.load(
g_ptr + row * stride_gm + (n_inter + cols) * stride_gn,
mask=mask,
other=0.0,
).to(tl.float32)
if HAS_LIMIT:
gate = tl.minimum(gate, limit)
up = tl.minimum(tl.maximum(up, -limit), limit)
out = gate * tl.sigmoid(alpha * gate) * (up + beta)
tl.store(
out_ptr + row * stride_om + cols * stride_on,
out.to(out_ptr.dtype.element_ty),
mask=mask,
)
@triton.jit
def _swiglu_oai_quant_kernel(
g_ptr,
aq_ptr,
as_ptr,
M,
n_inter,
stride_gm,
stride_gn,
stride_qm,
stride_qn,
stride_sm,
stride_sk,
alpha,
beta,
limit,
HAS_LIMIT: tl.constexpr,
BLOCK_M: tl.constexpr,
):
"""SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0)
quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one
MX block): it reads the matching gate/up columns from ``g1`` (``[M, 2I]``),
computes the SwiGLU in fp32, then derives the block E8M0 scale and emits the
FP8 values + scale in a single pass — no bf16 ``act`` round-trip to HBM.
"""
pid_m = tl.program_id(0)
pid_b = tl.program_id(1) # which 32-element block along I
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_c = pid_b * 32 + tl.arange(0, 32)
m_mask = offs_m < M
gate = tl.load(
g_ptr + offs_m[:, None] * stride_gm + offs_c[None, :] * stride_gn,
mask=m_mask[:, None],
other=0.0,
).to(tl.float32)
up = tl.load(
g_ptr + offs_m[:, None] * stride_gm + (n_inter + offs_c)[None, :] * stride_gn,
mask=m_mask[:, None],
other=0.0,
).to(tl.float32)
if HAS_LIMIT:
gate = tl.minimum(gate, limit)
up = tl.minimum(tl.maximum(up, -limit), limit)
act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32
amax = tl.maximum(tl.max(tl.abs(act), axis=1), 1e-30) # [BLOCK_M]
sb = tl.minimum(tl.maximum(tl.floor(tl.log2(amax)) + 127.0, 0.0), 254.0)
descale = tl.exp2(sb - 127.0)
aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty)
tl.store(
aq_ptr + offs_m[:, None] * stride_qm + offs_c[None, :] * stride_qn,
aq,
mask=m_mask[:, None],
)
tl.store(
as_ptr + offs_m * stride_sm + pid_b * stride_sk, sb.to(tl.uint8), mask=m_mask
)
def swiglu_oai_quantize_mxfp8(
gate_up: torch.Tensor,
alpha: float,
beta: float,
limit: float | None,
block_m: int = 64,
) -> tuple[torch.Tensor, torch.Tensor]:
"""SwiGLU-OAI on split-layout ``[M, 2I]`` fused with MXFP8 activation-quant.
Returns ``(act_q [M, I] float8_e4m3fn, act_scale [M, I//32] uint8 E8M0)``,
identical to ``mxfp8_e4m3_quantize(swiglu_oai_split(gate_up))`` but in a
single Triton pass (no bf16 intermediate). Used between the two GEMMs of the
native MXFP8 MoE. Numerically equivalent to the unfused chain (bit-exact on
measured MoE shapes); marginally more accurate (fp32 act, no bf16 round-trip).
"""
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
MXFP8_BLOCK_SIZE,
MXFP8_SCALE_DTYPE,
MXFP8_VALUE_DTYPE,
)
two_i = gate_up.shape[-1]
n_inter = two_i // 2
assert n_inter % MXFP8_BLOCK_SIZE == 0, (
f"fused swiglu+quant needs I % {MXFP8_BLOCK_SIZE} == 0, got I={n_inter}"
)
g1 = gate_up.reshape(-1, two_i).contiguous()
M = g1.shape[0]
aq = torch.empty((M, n_inter), dtype=MXFP8_VALUE_DTYPE, device=g1.device)
asc = torch.empty(
(M, n_inter // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=g1.device
)
grid = (triton.cdiv(M, block_m), n_inter // MXFP8_BLOCK_SIZE)
_swiglu_oai_quant_kernel[grid](
g1,
aq,
asc,
M,
n_inter,
g1.stride(0),
g1.stride(1),
aq.stride(0),
aq.stride(1),
asc.stride(0),
asc.stride(1),
float(alpha),
float(beta),
0.0 if limit is None else float(limit),
HAS_LIMIT=limit is not None,
BLOCK_M=block_m,
num_warps=4,
)
return aq, asc
def swiglu_oai_split(
gate_up: torch.Tensor,
alpha: float,
beta: float,
limit: float | None,
out_dtype: torch.dtype | None = None,
) -> torch.Tensor:
"""SwiGLU-OAI on a split-layout ``[*, 2I]`` tensor -> ``[*, I]``."""
orig_shape = gate_up.shape
two_i = orig_shape[-1]
n_inter = two_i // 2
x2 = gate_up.reshape(-1, two_i)
m = x2.shape[0]
dt = out_dtype if out_dtype is not None else gate_up.dtype
out = torch.empty((m, n_inter), dtype=dt, device=gate_up.device)
# Tile tuned on gfx950. The SwiGLU intermediate is sharded across tensor
# parallel ranks (per-rank n_inter = I / tp: dense I=12288, MoE I=3072), and
# a 512-wide tile (4 warps, ~2 elems/lane) only helps once the per-rank slice
# is large enough to be bandwidth-bound — at TP=1 prefill that is ~1.25-1.35x
# faster than 256. For small sharded slices (high TP) the kernel is launch-
# bound (~12us) and a wide tile can slightly regress, so fall back to 256.
# Decode is launch-bound at every TP. num_warps=8 underfills this tile, so it
# is pinned to 4.
block_i = 512 if n_inter >= 2048 else 256
grid = (m, triton.cdiv(n_inter, block_i))
_swiglu_oai_kernel[grid](
x2,
out,
n_inter,
x2.stride(0),
x2.stride(1),
out.stride(0),
out.stride(1),
float(alpha),
float(beta),
0.0 if limit is None else float(limit),
HAS_LIMIT=limit is not None,
BLOCK_I=block_i,
num_warps=4,
)
return out.reshape(*orig_shape[:-1], n_inter)
vllm/models/minimax_m3/common/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/minimax_m3/common/indexer.py
+512
View File
@@ -0,0 +1,512 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MiniMax M3 lightning indexer: side cache, metadata, and impl.
The indexer scores KV blocks with the index heads and selects the top-k blocks
(plus fixed init/local blocks) that the main block-sparse attention
(``sparse_attention.py``) then attends to. It owns its own side cache
(``MiniMaxM3IndexerCache``, one index-key vector per token), metadata, and
metadata builder, mirroring how DeepSeek V4 keeps the indexer separate from the
main attention.
``MiniMaxM3Indexer`` is the ``nn.Module`` the attention layer holds (like
``DeepseekV4Indexer``); it picks a kernel impl in ``__init__`` (via
``select_indexer_impl_cls``) and delegates ``forward`` to it.
"""
from dataclasses import dataclass
from typing import ClassVar
import torch
from torch import nn
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
from vllm.config.attention import IndexerKVDType
from vllm.config.cache import CacheDType
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.models.minimax_m3.common.ops.index_topk import (
minimax_m3_index_decode,
minimax_m3_index_score,
minimax_m3_index_topk,
)
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionMetadata,
AttentionMetadataBuilder,
CommonAttentionMetadata,
MultipleOf,
)
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheSpec,
MLAAttentionSpec,
)
class MiniMaxM3IndexerBackend(AttentionBackend):
"""Indexer side-cache backend (key-only)."""
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16]
# bf16 today; mirrors the main backend to keep spec validation permissive.
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"bfloat16",
"fp8",
"fp8_e4m3",
"fp8_e5m2",
]
@staticmethod
def get_name() -> str:
return "MINIMAX_M3_SPARSE_INDEXER"
@staticmethod
def get_impl_cls() -> type["MiniMaxM3IndexerImpl"]:
# Concrete impl chosen by select_indexer_impl_cls; base for introspection.
return MiniMaxM3IndexerImpl
@staticmethod
def get_builder_cls() -> type["MiniMaxM3IndexerMetadataBuilder"]:
return MiniMaxM3IndexerTritonMetadataBuilder
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return [128]
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
return [128]
@classmethod
def is_sparse(cls) -> bool:
return True
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
return (num_blocks, block_size, head_size)
@staticmethod
def get_kv_cache_stride_order(
include_num_layers_dimension: bool = False,
) -> tuple[int, ...]:
if include_num_layers_dimension:
# M3 does not use cross-layer (per-layer-stacked) KV blocks.
raise NotImplementedError
return (0, 1, 2)
class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
"""Side KV cache for the indexer's per-token index keys (key-only).
Registers itself in the static forward context so the KV-cache manager
allocates it (like ``DeepseekV32IndexerCache``).
"""
def __init__(
self,
head_dim: int,
prefix: str,
cache_config: CacheConfig | None = None,
indexer_kv_dtype: IndexerKVDType = "bf16",
backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend,
) -> None:
super().__init__()
if indexer_kv_dtype != "bf16":
raise NotImplementedError(
f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported yet "
"for the MiniMax M3 indexer cache (only 'bf16')."
)
self.kv_cache = torch.tensor([])
self.head_dim = head_dim
self.indexer_kv_dtype = indexer_kv_dtype
# Storage dtype for the side cache (bf16 today; quantized layouts later).
self.dtype = torch.bfloat16
self.prefix = prefix
self.cache_config = cache_config
# Impl-chosen backend -> each impl gets its own builder (get_attn_backend).
self.backend_cls = backend_cls
compilation_config = get_current_vllm_config().compilation_config
if prefix in compilation_config.static_forward_context:
raise ValueError(f"Duplicate layer name: {prefix}")
compilation_config.static_forward_context[prefix] = self
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
# Key-only: MLAAttentionSpec budgets one vector/token (not 2x for K+V).
return MLAAttentionSpec(
block_size=vllm_config.cache_config.block_size,
num_kv_heads=1,
head_size=self.head_dim,
dtype=self.dtype,
)
def forward(self) -> None: ...
def get_attn_backend(self) -> type[AttentionBackend]:
return self.backend_cls
@dataclass
class MiniMaxM3IndexerPrefillMetadata:
"""Per-prefill index-scoring state."""
cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0
seq_lens: torch.Tensor # [num_prefills] int32, total KV lengths
context_lens: torch.Tensor # [num_prefills] int32 (cached/context tokens)
block_table: torch.Tensor
max_query_len: int
max_seq_len: int
@dataclass
class MiniMaxM3IndexerDecodeMetadata:
"""Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform
per-request query length (1, or 1 + num_speculative_tokens)."""
seq_lens: torch.Tensor # [num_decodes] int32
block_table: torch.Tensor
max_seq_len: int
decode_query_len: int
@dataclass
class MiniMaxM3IndexerMetadata(AttentionMetadata):
"""Indexer metadata, split into prefill and decode sub-metadata."""
seq_lens: torch.Tensor
max_seq_len: int
slot_mapping: torch.Tensor
num_actual_tokens: int # total query tokens (decode-first batch)
# Split counts; identical to the main metadata's (same reorder threshold).
num_decodes: int
num_decode_tokens: int
num_prefills: int
num_prefill_tokens: int
prefill: MiniMaxM3IndexerPrefillMetadata | None = None
decode: MiniMaxM3IndexerDecodeMetadata | None = None
class MiniMaxM3IndexerMetadataBuilder(
AttentionMetadataBuilder[MiniMaxM3IndexerMetadata]
):
"""Abstract base: shared setup only. The Triton and MSA builders are
parallel subclasses that each own their full ``build`` (no shared code)."""
# Full cudagraphs for uniform decode batches (incl. spec-decode verify).
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
# Raised to 1 + num_speculative_tokens by _init_reorder_batch_threshold when
# spec decode is on; matches the main builder so the splits agree.
reorder_batch_threshold: int = 1
def __init__(
self,
kv_cache_spec: AttentionSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
) -> None:
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
hf_config = vllm_config.model_config.hf_config
text_config = getattr(hf_config, "text_config", hf_config)
sparse_cfg = text_config.sparse_attention_config
# Index-query head count from model config (cache spec has 1 vec/token).
total_index_heads = sparse_cfg["sparse_num_index_heads"]
tp_size = get_tensor_model_parallel_world_size()
if total_index_heads >= tp_size:
assert total_index_heads % tp_size == 0
else:
assert tp_size % total_index_heads == 0
self.num_index_heads = max(1, total_index_heads // tp_size)
self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
# Stable context-length buffer for decode cudagraph replays.
self.context_len_buffer = torch.empty(
vllm_config.scheduler_config.max_num_batched_tokens,
dtype=torch.int32,
device=device,
)
class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
"""Triton indexer metadata: no SM100 fmha_sm100 plan."""
def build(
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
) -> MiniMaxM3IndexerMetadata:
num_reqs = common_attn_metadata.num_reqs
num_tokens = common_attn_metadata.num_actual_tokens
query_start_loc = common_attn_metadata.query_start_loc
seq_lens = common_attn_metadata.seq_lens
block_table = common_attn_metadata.block_table_tensor
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
split_decodes_and_prefills(
common_attn_metadata,
decode_threshold=self.reorder_batch_threshold,
require_uniform=True,
)
)
assert num_decodes + num_prefills == num_reqs
assert num_decode_tokens + num_prefill_tokens == num_tokens
# Decode-first batch: context lengths into the stable cudagraph buffer.
context_lens = self.context_len_buffer[:num_reqs]
context_lens.copy_(
common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
)
prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None
if num_prefills > 0:
prefill_metadata = MiniMaxM3IndexerPrefillMetadata(
cu_seqlens_q=(query_start_loc[num_decodes:] - num_decode_tokens).to(
torch.int32
),
seq_lens=seq_lens[num_decodes:],
context_lens=context_lens[num_decodes:],
block_table=block_table[num_decodes:],
max_query_len=common_attn_metadata.max_query_len,
max_seq_len=common_attn_metadata.max_seq_len,
)
decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None
if num_decodes > 0:
qsl_cpu = common_attn_metadata.query_start_loc_cpu
query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
decode_query_len = int(query_lens_cpu[0].item())
assert decode_query_len > 0
assert torch.all(
(query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
)
assert num_decode_tokens == num_decodes * decode_query_len
decode_metadata = MiniMaxM3IndexerDecodeMetadata(
seq_lens=seq_lens[:num_decodes],
block_table=block_table[:num_decodes],
max_seq_len=common_attn_metadata.max_seq_len,
decode_query_len=decode_query_len,
)
return MiniMaxM3IndexerMetadata(
seq_lens=seq_lens,
max_seq_len=common_attn_metadata.max_seq_len,
slot_mapping=common_attn_metadata.slot_mapping,
num_actual_tokens=num_tokens,
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
prefill=prefill_metadata,
decode=decode_metadata,
)
class MiniMaxM3IndexerImpl(nn.Module):
"""Abstract base for the indexer kernel impls.
Each impl owns its side cache and reports its backend via
``indexer_backend_cls`` (so each gets its own builder). The Triton and MSA
subclasses each own a full ``forward`` returning ``(decode_topk,
prefill_topk)`` -- no shared forward code.
"""
# Set by each impl so the side cache reports the matching backend + builder.
indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerBackend
def __init__(
self,
*,
num_kv_heads: int,
scale: float,
topk_blocks: int,
sparse_block_size: int,
num_index_heads: int,
index_head_dim: int,
prefix: str,
init_blocks: int = 0,
local_blocks: int = 0,
score_type: str = "max",
cache_config: CacheConfig | None = None,
indexer_kv_dtype: IndexerKVDType = "bf16",
) -> None:
super().__init__()
self.num_kv_heads = num_kv_heads
self.scale = scale
self.topk_blocks = topk_blocks
self.block_size = sparse_block_size
self.init_blocks = init_blocks
self.local_blocks = local_blocks
self.score_type = score_type
self.num_index_heads = num_index_heads
self.index_head_dim = index_head_dim
self.indexer_kv_dtype = indexer_kv_dtype
# Owns the side cache (registers itself in the static forward context).
self.index_cache = MiniMaxM3IndexerCache(
head_dim=index_head_dim,
prefix=f"{prefix}.index_cache",
cache_config=cache_config,
indexer_kv_dtype=indexer_kv_dtype,
backend_cls=type(self).indexer_backend_cls,
)
def forward(
self,
index_query: torch.Tensor,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
"""Return ``(decode_topk, prefill_topk)``; implemented per kernel impl."""
raise NotImplementedError
class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
"""Triton indexer score + top-k for both prefill and decode."""
def forward(
self,
index_query: torch.Tensor,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
attn_metadata = get_forward_context().attn_metadata
if not isinstance(attn_metadata, dict):
return None, None # profiling run; caches unbound
index_md = attn_metadata[self.index_cache.prefix]
assert isinstance(index_md, MiniMaxM3IndexerMetadata)
num_tokens = index_md.num_actual_tokens
nd = index_md.num_decode_tokens
iq = index_query[:num_tokens].view(
-1, self.num_index_heads, self.index_head_dim
)
kv = self.index_cache.kv_cache
decode_topk: torch.Tensor | None = None
prefill_topk: torch.Tensor | None = None
if index_md.num_decodes > 0:
d = index_md.decode
assert d is not None
decode_topk = minimax_m3_index_decode(
iq[:nd],
kv,
d.block_table,
d.seq_lens,
d.max_seq_len,
self.topk_blocks,
self.init_blocks,
self.local_blocks,
self.num_kv_heads,
self.scale,
d.decode_query_len,
)
if index_md.num_prefills > 0:
p = index_md.prefill
assert p is not None
score = minimax_m3_index_score(
iq[nd:],
kv,
p.block_table,
p.cu_seqlens_q,
p.seq_lens,
p.context_lens,
p.max_query_len,
p.max_seq_len,
self.num_kv_heads,
self.scale,
)
prefill_topk = minimax_m3_index_topk(
score,
p.cu_seqlens_q,
p.context_lens,
p.max_query_len,
self.topk_blocks,
self.init_blocks,
self.local_blocks,
)
return decode_topk, prefill_topk
def select_indexer_impl_cls(
*,
indexer_kv_dtype: IndexerKVDType = "bf16",
) -> type[MiniMaxM3IndexerImpl]:
"""Pick the indexer impl off the index-cache dtype.
The SM100 MSA indexer score path is disabled for now; use the local Triton
indexer. If re-enabled, add a NVIDIA-specific ``MiniMaxM3IndexerImpl`` here.
"""
if indexer_kv_dtype in ("mxfp4", "nvfp4"):
raise NotImplementedError(
f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
"CuteDSL indexer impl."
)
if indexer_kv_dtype != "bf16":
raise NotImplementedError(
f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
"Triton indexer impl."
)
return MiniMaxM3IndexerTritonImpl
class MiniMaxM3Indexer(nn.Module):
"""Indexer module held by the attention layer (like ``DeepseekV4Indexer``).
Picks the kernel impl in ``__init__`` (``select_indexer_impl_cls``) and
delegates ``forward``; exposes the impl's side cache via ``index_cache``.
"""
def __init__(
self,
*,
num_kv_heads: int,
scale: float,
topk_blocks: int,
sparse_block_size: int,
num_index_heads: int,
index_head_dim: int,
prefix: str,
init_blocks: int = 0,
local_blocks: int = 0,
score_type: str = "max",
cache_config: CacheConfig | None = None,
indexer_kv_dtype: IndexerKVDType = "bf16",
) -> None:
super().__init__()
impl_cls = select_indexer_impl_cls(
indexer_kv_dtype=indexer_kv_dtype,
)
self.impl = impl_cls(
num_kv_heads=num_kv_heads,
scale=scale,
topk_blocks=topk_blocks,
sparse_block_size=sparse_block_size,
num_index_heads=num_index_heads,
index_head_dim=index_head_dim,
prefix=prefix,
init_blocks=init_blocks,
local_blocks=local_blocks,
score_type=score_type,
cache_config=cache_config,
indexer_kv_dtype=indexer_kv_dtype,
)
@property
def index_cache(self) -> MiniMaxM3IndexerCache:
return self.impl.index_cache
@property
def num_index_heads(self) -> int:
return self.impl.num_index_heads
def forward(
self,
index_query: torch.Tensor,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
return self.impl(index_query)
vllm/models/minimax_m3/common/mm_preprocess.py
+514
View File
@@ -0,0 +1,514 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
from collections.abc import Mapping, Sequence
from typing import cast
import torch
from transformers import BatchFeature
from transformers.video_utils import VideoMetadata
from vllm.config.multimodal import (
BaseDummyOptions,
ImageDummyOptions,
VideoDummyOptions,
)
from vllm.inputs import MultiModalDataDict
from vllm.multimodal.inputs import (
MultiModalFieldConfig,
MultiModalKwargsItems,
)
from vllm.multimodal.parse import (
ImageSize,
MultiModalDataItems,
MultiModalDataParser,
)
from vllm.multimodal.processing import (
BaseDummyInputsBuilder,
BaseMultiModalProcessor,
BaseProcessingInfo,
PromptReplacement,
PromptUpdate,
PromptUpdateDetails,
)
from vllm.multimodal.video import (
VIDEO_LOADER_REGISTRY,
VideoBackend,
VideoSourceMetadata,
VideoTargetMetadata,
)
from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3Config
from vllm.transformers_utils.processors.minimax_m3 import (
MIN_SHORT_SIDE_PIXEL,
MiniMaxM3VLImageProcessor,
MiniMaxM3VLVideoProcessor,
MiniMaxVLProcessor,
smart_resize,
)
# Upper bound on the number of frames used to build the dummy video during
# memory profiling. Sized to the worst-case video the processor accepts:
# ``max_total_pixels // max_pixels_per_frame`` = 301,056,000 // 602,112 = 500
# frames, each at the video processor's per-frame ``max_pixels`` (768 * 28 * 28
# = 602,112). This reaches the true worst-case ~192,000 vision tokens, but only
# because the dummy video is sized via ``get_video_size_with_most_features()``
# (the video ``max_pixels`` bound), not the smaller image bound. Without a cap,
# ``_get_max_video_frames(seq_len)`` with M3's large ``max_model_len`` yields
# ~1400 frames, producing a multi-GB dummy tensor that overflows the
# multimodal encoder cache.
_MAX_FRAMES_PER_VIDEO = 500
class MiniMaxM3VLProcessingInfo(BaseProcessingInfo):
IMAGE_TOKEN = "]<]image[>["
VIDEO_TOKEN = "]<]video[>["
VISION_START_TOKEN = "]<]start of image[>["
VISION_END_TOKEN = "]<]end of image[>["
def get_hf_config(self) -> MiniMaxM3Config:
return self.ctx.get_hf_config(MiniMaxM3Config)
def get_hf_processor(self, **kwargs: object) -> MiniMaxVLProcessor:
# The released checkpoint only ships the processor as remote code
# (via ``auto_map``). Construct the vendored processor directly so the
# model loads without ``--trust-remote-code``.
return self.ctx.get_hf_processor(MiniMaxVLProcessor, **kwargs)
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
return {"image": None, "video": None}
def get_mm_max_tokens_per_item(
self,
seq_len: int,
mm_counts: Mapping[str, int],
) -> Mapping[str, int]:
return {
"image": self.get_max_image_tokens(),
"video": self.get_max_video_tokens(seq_len, mm_counts),
}
def get_image_processor(self, **kwargs: object) -> MiniMaxM3VLImageProcessor:
return self.get_hf_processor(**kwargs).image_processor
def get_video_processor(self, **kwargs: object) -> MiniMaxM3VLVideoProcessor:
return self.get_hf_processor(**kwargs).video_processor
def _get_vision_info(
self,
*,
image_width: int,
image_height: int,
num_frames: int,
image_processor,
) -> tuple[ImageSize, int]:
"""Compute resized image size and number of vision tokens.
Mirrors the processor's Qwen-style ``smart_resize`` (area bound by
``max_pixels``) so token counts match the actual processor output.
"""
patch_size: int = image_processor.patch_size
merge_size: int = image_processor.merge_size
temporal_patch_size: int = image_processor.temporal_patch_size
factor = patch_size * merge_size
max_pixels: int = image_processor.max_pixels
# Long-side resize spec (opt-in). ``image_processor`` is the *video*
# processor when counting video tokens, so read the bounds off it.
max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None)
min_short_side_pixel = getattr(
image_processor, "min_short_side_pixel", MIN_SHORT_SIDE_PIXEL
)
new_h, new_w = smart_resize(
image_height,
image_width,
factor=factor,
max_pixels=max_pixels,
max_long_side_pixel=max_long_side_pixel,
min_short_side_pixel=min_short_side_pixel,
# Token counting must not raise; the volumetric/area cap is enforced
# in the processor's _preprocess on the real inputs.
max_total_pixels=None,
)
grid_h = new_h // patch_size
grid_w = new_w // patch_size
# Pad frames to be divisible by temporal_patch_size
padded_frames = num_frames + (-num_frames % temporal_patch_size)
grid_t = max(padded_frames // temporal_patch_size, 1)
num_tokens = grid_t * grid_h * grid_w // (merge_size**2)
return ImageSize(width=new_w, height=new_h), num_tokens
def get_num_image_tokens(
self,
*,
image_width: int,
image_height: int,
image_processor,
mm_kwargs: Mapping[str, object],
) -> int:
_, n = self._get_vision_info(
image_width=image_width,
image_height=image_height,
num_frames=1,
image_processor=image_processor,
)
return n
def get_num_video_tokens(
self,
*,
image_width: int,
image_height: int,
num_frames: int,
image_processor,
mm_kwargs: Mapping[str, object],
) -> int:
_, n = self._get_vision_info(
image_width=image_width,
image_height=image_height,
num_frames=num_frames,
image_processor=image_processor,
)
return n
def get_image_size_with_most_features(self) -> ImageSize:
# Largest square (a multiple of patch_size*merge_size) whose area is
# within the image processor's bound — this yields the most vision
# tokens for one image. With the long-side spec the square side is
# capped by ``max_long_side_pixel`` (and the fixed ``max_total_pixels``);
# otherwise it is bound by the ``max_pixels`` area.
image_processor = self.get_image_processor()
factor = image_processor.patch_size * image_processor.merge_size
max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None)
if max_long_side_pixel is not None:
side_px = min(
max_long_side_pixel,
math.isqrt(image_processor.max_total_pixels),
)
else:
side_px = math.isqrt(image_processor.max_pixels)
side = max(factor, (side_px // factor) * factor)
return ImageSize(width=side, height=side)
def get_video_size_with_most_features(self) -> ImageSize:
# Per-frame size that yields the most vision tokens, bound by the
# *video* processor's ``max_pixels`` (which differs from the image
# bound). Token count depends only on area, so maximize the area
# achievable with both sides a multiple of patch_size*merge_size rather
# than picking the largest square — a square (e.g. 756x756 for M3's
# 602,112 bound) leaves area on the table, undercounting frames.
video_processor = self.get_video_processor()
factor = video_processor.patch_size * video_processor.merge_size
per_frame_pixels = video_processor.max_pixels
max_long_side_pixel = getattr(video_processor, "max_long_side_pixel", None)
if max_long_side_pixel is not None:
# Long-side spec: a frame's worst case is a square capped by
# ``max_long_side_pixel`` (per-frame area, not the volumetric cap).
per_frame_pixels = min(per_frame_pixels, max_long_side_pixel**2)
units = per_frame_pixels // (factor * factor) # h_u * w_u
h_u = math.isqrt(units)
while units % h_u:
h_u -= 1
return ImageSize(width=(units // h_u) * factor, height=h_u * factor)
def get_max_image_tokens(self) -> int:
image_processor = self.get_image_processor()
size = self.get_image_size_with_most_features()
return self.get_num_image_tokens(
image_width=size.width,
image_height=size.height,
image_processor=image_processor,
mm_kwargs={},
)
def _get_max_video_frames(self, max_tokens: int) -> int:
video_processor = self.get_video_processor()
size = self.get_video_size_with_most_features()
num_frames = 1
while True:
next_n = self.get_num_video_tokens(
image_width=size.width,
image_height=size.height,
num_frames=num_frames + 1,
image_processor=video_processor,
mm_kwargs={},
)
if next_n > max_tokens:
break
num_frames += 1
return num_frames
def get_num_frames_with_most_features(
self,
seq_len: int,
mm_counts: Mapping[str, int],
max_frames_per_video: int = _MAX_FRAMES_PER_VIDEO,
) -> int:
max_videos = mm_counts.get("video", 0)
max_total_frames = self._get_max_video_frames(seq_len)
max_frames_per_video = min(
max_total_frames // max(max_videos, 1), max_frames_per_video
)
return max(max_frames_per_video, 1)
def get_max_video_tokens(
self,
seq_len: int,
mm_counts: Mapping[str, int],
) -> int:
video_processor = self.get_video_processor()
size = self.get_video_size_with_most_features()
return self.get_num_video_tokens(
image_width=size.width,
image_height=size.height,
num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
image_processor=video_processor,
mm_kwargs={},
)
class MiniMaxM3VLDummyInputsBuilder(BaseDummyInputsBuilder[MiniMaxM3VLProcessingInfo]):
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
num_images = mm_counts.get("image", 0)
num_videos = mm_counts.get("video", 0)
image_token: str = self.info.IMAGE_TOKEN
video_token: str = self.info.VIDEO_TOKEN
return image_token * num_images + video_token * num_videos
def get_dummy_mm_data(
self,
seq_len: int,
mm_counts: Mapping[str, int],
mm_options: Mapping[str, BaseDummyOptions],
) -> MultiModalDataDict:
size = self.info.get_image_size_with_most_features()
video_size = self.info.get_video_size_with_most_features()
num_frames = self.info.get_num_frames_with_most_features(seq_len, mm_counts)
return {
"image": self._get_dummy_images(
width=size.width,
height=size.height,
num_images=mm_counts.get("image", 0),
overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
),
"video": self._get_dummy_videos(
width=video_size.width,
height=video_size.height,
num_frames=num_frames,
num_videos=mm_counts.get("video", 0),
overrides=cast(VideoDummyOptions | None, mm_options.get("video")),
),
}
class MiniMaxM3VLMultiModalProcessor(
BaseMultiModalProcessor[MiniMaxM3VLProcessingInfo]
):
def _get_data_parser(self) -> MultiModalDataParser:
# Request video metadata (fps + sampled frame indices) so the HF
# processor can emit per-frame ``]<]X.X seconds[>[`` timestamp markers,
# matching MiniMax's reference video token stream. ``_get_prompt_updates``
# reconstructs the same markers from the metadata to keep the prompt
# replacement aligned with the processor output.
return MultiModalDataParser(video_needs_metadata=True)
def _call_hf_processor(
self,
prompt: str,
mm_data: Mapping[str, object],
mm_kwargs: Mapping[str, object],
tok_kwargs: Mapping[str, object],
) -> BatchFeature:
mm_data = dict(mm_data)
# With ``video_needs_metadata=True`` each video arrives as a
# ``(frames, metadata)`` tuple. Split the frames back out and forward the
# metadata as ``VideoMetadata`` so the processor emits timestamps.
videos = cast(list | None, mm_data.get("videos"))
video_metadata: list[VideoMetadata] | None = None
if videos:
frames_only = []
video_metadata = []
for item in videos:
if isinstance(item, tuple) and len(item) == 2:
frames, meta = item
else:
frames, meta = item, {}
frames_only.append(frames)
meta = {
k: v for k, v in (meta or {}).items() if k != "do_sample_frames"
}
# VideoMetadata requires total_num_frames; derive it for
# dummy/profiling videos whose metadata omits it. fps and
# frames_indices default to None there → no timestamps, which
# stays consistent with _get_prompt_updates.
meta.setdefault("total_num_frames", len(frames))
video_metadata.append(VideoMetadata(**meta))
mm_data["videos"] = frames_only
# Override the video processor's default do_resize=False (set for a
# pre-resized pipeline) to True for vLLM's raw-frame inputs.
merged = dict(do_resize=True, **mm_kwargs, **tok_kwargs)
data = dict(text=prompt, **mm_data)
if video_metadata is not None:
data["video_metadata"] = video_metadata
return self.info.ctx.call_hf_processor(
self.info.get_hf_processor(**mm_kwargs),
data,
merged,
)
def _get_mm_fields_config(
self,
hf_inputs: BatchFeature,
hf_processor_mm_kwargs: Mapping[str, object],
) -> Mapping[str, MultiModalFieldConfig]:
image_grid_thw = hf_inputs.get("image_grid_thw")
video_grid_thw = hf_inputs.get("video_grid_thw")
# Total patches per item (grid_t * grid_h * grid_w)
image_grid_sizes = (
image_grid_thw.prod(-1)
if image_grid_thw is not None
else torch.empty(0, dtype=torch.long)
)
video_grid_sizes = (
video_grid_thw.prod(-1)
if video_grid_thw is not None
else torch.empty(0, dtype=torch.long)
)
return {
"pixel_values": MultiModalFieldConfig.flat_from_sizes(
"image", image_grid_sizes
),
"image_grid_thw": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
"pixel_values_videos": MultiModalFieldConfig.flat_from_sizes(
"video", video_grid_sizes
),
"video_grid_thw": MultiModalFieldConfig.batched("video", keep_on_cpu=True),
}
def _get_prompt_updates(
self,
mm_items: MultiModalDataItems,
hf_processor_mm_kwargs: Mapping[str, object],
out_mm_kwargs: MultiModalKwargsItems,
) -> Sequence[PromptUpdate]:
hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
tokenizer = self.info.get_tokenizer()
vocab = tokenizer.get_vocab()
image_token_id: int = vocab[self.info.IMAGE_TOKEN]
video_token_id: int = vocab[self.info.VIDEO_TOKEN]
start_token_id: int = vocab[self.info.VISION_START_TOKEN]
end_token_id: int = vocab[self.info.VISION_END_TOKEN]
merge_length: int = hf_processor.image_processor.merge_size**2
def get_image_replacement(item_idx: int):
grid_thw: torch.Tensor = out_mm_kwargs["image"][item_idx][
"image_grid_thw"
].data
# grid_thw shape: (3,) = [1, grid_h, grid_w]
N = int(grid_thw.prod().item()) // merge_length
full = [start_token_id] + [image_token_id] * N + [end_token_id]
return PromptUpdateDetails.select_token_id(full, image_token_id)
# Per-video metadata (fps + sampled frame indices) is carried on the
# parsed video items; used to reproduce the HF processor's timestamps.
video_items = mm_items.get("video")
video_metadata = getattr(video_items, "metadata", None)
temporal_patch_size: int = hf_processor.video_processor.temporal_patch_size
def get_video_replacement(item_idx: int):
grid_thw: torch.Tensor = out_mm_kwargs["video"][item_idx][
"video_grid_thw"
].data
# grid_thw shape: (3,) = [grid_t, grid_h, grid_w]
# HF model uses VIDEO_TOKEN (not IMAGE_TOKEN) for video frame content:
# processing_minimax.py L245: replace(placeholder, self.VIDEO_TOKEN)
T = int(grid_thw[0].item())
M = int(grid_thw[1].item() * grid_thw[2].item()) // merge_length
# Reproduce the HF processor's per-frame timestamp markers
# (processing_minimax.py: ts = frames_indices[frame*tps] / fps,
# rendered as "]<]X.X seconds[>["). Falls back to no timestamps when
# metadata is unavailable (keeping the replacement aligned with the
# processor output in both cases).
meta = (
video_metadata[item_idx]
if video_metadata is not None and item_idx < len(video_metadata)
else None
)
fps = meta.get("fps") if meta else None
frames_indices = meta.get("frames_indices") if meta else None
full: list[int] = []
for frame_idx in range(T):
if fps is not None and frames_indices is not None:
idx = min(frame_idx * temporal_patch_size, len(frames_indices) - 1)
ts = frames_indices[idx] / fps
full += tokenizer.encode(
f"]<]{ts:.1f} seconds[>[", add_special_tokens=False
)
full += [start_token_id] + [video_token_id] * M + [end_token_id]
return PromptUpdateDetails.select_token_id(full, video_token_id)
return [
PromptReplacement(
modality="image",
target=[image_token_id],
replacement=get_image_replacement,
),
PromptReplacement(
modality="video",
target=[video_token_id],
replacement=get_video_replacement,
),
]
# TODO(Isotr0py): Tie with MinimaxVideoProcessor
# after https://github.com/vllm-project/vllm/pull/44126
@VIDEO_LOADER_REGISTRY.register("minimax_m3_vl")
class MiniMaxM3VideoBackend(VideoBackend):
@classmethod
def compute_frames_index_to_sample(
cls,
source: VideoSourceMetadata,
target: VideoTargetMetadata,
**kwargs,
) -> list[int]:
total_frames = source.total_frames_num
video_fps = source.original_fps
fps = target.fps
if total_frames <= 0 or video_fps <= 0 or fps <= 0:
return [0] if total_frames > 0 else []
read_time_interval = 1.0 / fps
eps = 1e-4
indices: list[int] = []
prev_kept_ts = -float("inf")
while True:
if not indices:
target_frame = 0
else:
target_ts = prev_kept_ts + read_time_interval - eps
target_frame = math.ceil(target_ts * video_fps)
target_frame = max(target_frame, indices[-1] + 1)
if target_frame >= total_frames:
break
indices.append(target_frame)
prev_kept_ts = target_frame / video_fps
last_frame_idx = total_frames - 1
last_ts = last_frame_idx / video_fps
if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps:
indices.append(last_frame_idx)
if not indices:
indices = [0]
return indices
vllm/models/minimax_m3/common/ops/__init__.py
+18
View File
@@ -0,0 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Cross-platform (Triton) kernels for MiniMax M3 sparse attention."""
from .index_topk import (
minimax_m3_index_decode,
minimax_m3_index_score,
minimax_m3_index_topk,
)
from .sparse_attn import minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode
__all__ = [
"minimax_m3_index_decode",
"minimax_m3_index_score",
"minimax_m3_index_topk",
"minimax_m3_sparse_attn",
"minimax_m3_sparse_attn_decode",
]
vllm/models/minimax_m3/common/ops/index_topk.py
+898
View File
@@ -0,0 +1,898 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k.
Index queries score each 128-token block of index keys (max over the block),
then the top-k blocks (plus forced init/local blocks) are selected per query
token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the
sparse block size (128), so one sparse block maps to exactly one page.
Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head).
Only the paths MiniMax M3 uses are implemented: score_type="max", index value
disabled (score-only indexer), single shared index head. The selected block ids
feed the block-sparse attention kernels in ``sparse_attn``.
"""
import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import round_up
# One sparse block == one KV page.
SPARSE_BLOCK_SIZE = 128
# ---------------------------------------------------------------------------
# Bitonic top-k helpers (layout-agnostic).
# ---------------------------------------------------------------------------
@triton.jit
def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr):
n_outer: tl.constexpr = x.numel >> n_dims
shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)]
y = tl.reshape(x, shape)
mask = tl.arange(0, 2)[None, :, None]
left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype)
right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype)
left = tl.reshape(left, x.shape)
right = tl.reshape(right, x.shape)
y_idx = tl.reshape(ids, shape)
left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape)
right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape)
left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype)
right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype)
idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True)
ileft = left.to(idtype, bitcast=True)
iright = right.to(idtype, bitcast=True)
ix = x.to(idtype, bitcast=True)
cond = (left > right) != flip
ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix))
new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids))
return ret.to(x.dtype, bitcast=True), new_ids
@triton.jit
def _bitonic_merge(
x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr
):
n_outer: tl.constexpr = x.numel >> n_dims
tl.static_assert(stage <= n_dims)
if order == 2:
shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage]
flip = tl.reshape(
tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape
)
else:
flip = order
for i in tl.static_range(stage):
x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims)
return x, ids
# ---------------------------------------------------------------------------
# Index block-score kernel (paged). score[h, token, block] = max over the
# 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so
# each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1).
# ---------------------------------------------------------------------------
# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
# might lose pointer alignment, which trigger Triton recompiles. we don't actually
# need pointer alignment for those tensors anyway because we do scalar load.
@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
def _index_block_score_kernel(
q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
score_ptr, # [num_idx_heads, total_q, max_block]
block_table_ptr, # [num_reqs, max_blocks]
cu_seqlens, # [batch+1] query start offsets
seq_lens, # [batch] total K length
prefix_lens, # [batch] context length before this chunk's queries
num_idx_heads,
head_dim: tl.constexpr,
sm_scale,
stride_q_n,
stride_q_h,
stride_q_d,
stride_ik_blk,
stride_ik_pos,
stride_ik_d,
stride_s_h,
stride_s_n,
stride_s_k,
stride_bt_b,
BLOCK_SIZE_Q: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
):
sm_scale_log2e = sm_scale * 1.4426950409
pid_q = tl.program_id(0)
pid_bh = tl.program_id(1)
pid_b = pid_bh // num_idx_heads
pid_h = pid_bh % num_idx_heads
seq_start = tl.load(cu_seqlens + pid_b)
q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start
seq_len = tl.load(seq_lens + pid_b)
prefix_len = tl.load(prefix_lens + pid_b)
if BLOCK_SIZE_Q * pid_q >= q_len:
return
q_ptrs = tl.make_block_ptr(
base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h,
shape=(q_len, head_dim),
strides=(stride_q_n, stride_q_d),
offsets=(pid_q * BLOCK_SIZE_Q, 0),
block_shape=(BLOCK_SIZE_Q, head_dim),
order=(1, 0),
)
q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero")
q_start = prefix_len + pid_q * BLOCK_SIZE_Q
off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len
off_k = tl.arange(0, BLOCK_SIZE_K)
off_d = tl.arange(0, head_dim)
# Block table row for this request.
bt_row = block_table_ptr + pid_b * stride_bt_b
# Causal window: only blocks up to the last query token's position.
hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
for i in tl.range(0, hi, BLOCK_SIZE_K):
blk = i // BLOCK_SIZE_K
page = tl.load(bt_row + blk).to(tl.int64)
pos = i + off_k
# index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed)
# we don't need masked load for K, because KV cache ensures
# allocation is multiple of BLOCK_SIZE_K.
# for tokens beyond seqlen, they will be masked in qk later.
k = tl.load(
ik_cache_ptr
+ page * stride_ik_blk
+ off_k[None, :] * stride_ik_pos
+ off_d[:, None] * stride_ik_d,
)
qk = tl.dot(q, k) * sm_scale_log2e
# apply causal mask as needed
if q_start < i + BLOCK_SIZE_K:
qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf"))
# one sparse block per K-tile -> max over the 128 positions
score = tl.max(qk, axis=1) # [BLOCK_SIZE_Q]
s_ptrs = (
score_ptr
+ pid_h * stride_s_h
+ (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q))
* stride_s_n
+ blk * stride_s_k
)
q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len
tl.store(s_ptrs, score, mask=q_store_mask)
# ---------------------------------------------------------------------------
# Top-k selection over per-token block scores (layout-agnostic). block_size_q
# is 1 for M3, so top-k is computed per query token.
# ---------------------------------------------------------------------------
# since prefill metadata is sliced from mixed batch metadata, prefix_lens
# might lose pointer alignment, which trigger Triton recompiles. we don't actually
# need pointer alignment for those tensors anyway because we do scalar load.
@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
@triton.autotune(
configs=[
triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
],
key=["BLOCK_SIZE_T"],
)
@triton.jit(do_not_specialize_on_alignment=["prefix_lens"])
def _topk_index_kernel(
s_ptr, # [num_heads, total_q, max_block]
ti_ptr, # [num_heads, total_q, topk]
sample_interval: tl.constexpr, # block_size_q (1 for M3)
block_size: tl.constexpr, # sparse block size (128)
cu_seqlens,
cu_seqblocks_q,
prefix_lens,
topk,
init_blocks: tl.constexpr,
local_blocks: tl.constexpr,
stride_s_h,
stride_s_n,
stride_s_k,
stride_ti_h,
stride_ti_n,
stride_ti_t,
BLOCK_SIZE_K: tl.constexpr,
BLOCK_SIZE_T: tl.constexpr,
MASK_INIT: tl.constexpr,
MASK_LOCAL: tl.constexpr,
):
tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T)
pid_q = tl.program_id(0)
pid_b = tl.program_id(1)
pid_h = tl.program_id(2)
seq_start = tl.load(cu_seqlens + pid_b)
block_start = tl.load(cu_seqblocks_q + pid_b)
block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start
prefix_len = tl.load(prefix_lens + pid_b)
if pid_q >= block_num:
return
off_k = tl.arange(0, BLOCK_SIZE_K)
off_t = tl.arange(0, BLOCK_SIZE_T)
s_ptrs = (
s_ptr
+ (seq_start + pid_q * sample_interval) * stride_s_n
+ pid_h * stride_s_h
+ off_k * stride_s_k
)
topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size
for i in tl.range(0, valid_blocks, BLOCK_SIZE_K):
causal_mask = i + off_k < valid_blocks
local_mask = i + off_k >= max(0, valid_blocks - local_blocks)
init_mask = i + off_k < init_blocks
score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32)
score = tl.where(score != score, -1e30, score)
s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
if MASK_INIT:
score = tl.where(causal_mask & init_mask, score - 1e29, score)
else:
score = tl.where(causal_mask & init_mask, 1e30, score)
if MASK_LOCAL:
score = tl.where(causal_mask & local_mask, score - 1e28, score)
else:
score = tl.where(causal_mask & local_mask, 1e29, score)
topk_score, last_topk_score = score, topk_score
topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx)
n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
for j in tl.static_range(1, n_dims):
topk_score, topk_idx = _bitonic_merge(
topk_score, topk_idx.to(tl.int32), j, 2, n_dims
)
if i != 0:
topk_score, topk_idx = _bitonic_merge(
topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
)
topk_score_new = last_topk_score * left_half_mask + topk_score * (
1 - left_half_mask
)
topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
1 - left_half_mask
)
topk_score, topk_idx = _bitonic_merge(
topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
)
else:
topk_score, topk_idx = _bitonic_merge(
topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
)
topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
topk_idx = tl.sum(
topk_mask[:, None]
* tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
axis=0,
)
ti_ptrs = (
ti_ptr
+ (block_start + pid_q) * stride_ti_n
+ pid_h * stride_ti_h
+ off_t * stride_ti_t
)
store_mask = off_t < topk
valid_mask = off_t < valid_blocks
topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1)
tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask)
# ---------------------------------------------------------------------------
# Decode index-score kernel (split-K over seq blocks). Decode batches are
# flattened request-major, with a runtime query length used to map each query
# token back to its request metadata. Chunk counts depend only on shape
# constants so the grid is fixed within a cuda graph. Base-2 (exp2/log2)
# softmax matches prefill.
# ---------------------------------------------------------------------------
@triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"])
def _decode_index_score_kernel(
q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
score_ptr, # [num_idx_heads, total_q, max_block]
block_table_ptr, # [num_reqs, max_blocks]
seq_lens, # [num_reqs]
num_idx_heads: tl.constexpr,
head_dim: tl.constexpr,
init_blocks,
local_blocks,
sm_scale,
decode_query_len,
stride_q_n,
stride_q_h,
stride_q_d,
stride_ik_blk,
stride_ik_pos,
stride_ik_d,
stride_s_h,
stride_s_n,
stride_s_k,
stride_bt_b,
BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
num_kv_chunks,
USE_PDL: tl.constexpr,
):
sm_scale_log2e = sm_scale * 1.4426950409
pid_b = tl.program_id(0) # flattened query-token id
pid_c = tl.program_id(1)
req_id = pid_b // decode_query_len
q_offset = pid_b - req_id * decode_query_len
if USE_PDL:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
seq_len = tl.load(seq_lens + req_id)
query_pos = seq_len - decode_query_len + q_offset
# Full-CG padding uses zero-length request rows. Clamp to an empty
# attention range instead of letting padded rows produce negative lengths.
kv_len = tl.maximum(query_pos + 1, 0)
num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
# block-aligned fixed-count split: grid independent of seq_len (cuda graph).
chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks
chunk_start_block = pid_c * chunk_size_blocks
chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks)
if chunk_start_block >= chunk_end_block:
return
off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block
off_d = tl.arange(0, head_dim)
bt_row = block_table_ptr + req_id * stride_bt_b
# Force-select init (1e30) and local (1e29, higher priority) blocks.
local_start = tl.maximum(0, num_blocks - local_blocks)
# query vectors across all heads
q = tl.load(
q_ptr
+ pid_b * stride_q_n
+ tl.arange(0, num_idx_heads) * stride_q_h
+ off_d[:, None] * stride_q_d,
) # [D,H]
for blk in tl.range(chunk_start_block, chunk_end_block):
page = tl.load(bt_row + blk).to(tl.int64)
pos = blk * BLOCK_SIZE_K + off_k
pos_mask = pos < kv_len
# we don't need masked load for K, because KV cache ensures
# allocation is multiple of BLOCK_SIZE_K.
# for tokens beyond seqlen, they will be masked in qk later.
k = tl.load(
ik_cache_ptr
+ page * stride_ik_blk
+ off_k[:, None] * stride_ik_pos
+ off_d * stride_ik_d,
) # [N,D]
kq = tl.dot(k, q) * sm_scale_log2e # [N,H]
kq = tl.where(pos_mask[:, None], kq, float("-inf"))
score = tl.max(kq, axis=0) # [H]
is_init = blk < init_blocks
is_local = (blk >= local_start) & (blk < num_blocks)
score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score))
tl.store(
score_ptr
+ tl.arange(0, num_idx_heads) * stride_s_h
+ pid_b * stride_s_n
+ blk * stride_s_k,
score,
)
# ---------------------------------------------------------------------------
# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local
# blocks are already encoded in the scores.
# ---------------------------------------------------------------------------
@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
@triton.autotune(
configs=[
triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3),
triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
],
key=["topk"],
)
@triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"])
def _topk_index_partial_kernel(
s_ptr, # score: [num_idx_heads, total_q, max_block]
ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape
seq_lens, # [num_reqs]
block_size: tl.constexpr, # sparse block size (128)
topk: tl.constexpr,
chunk_blocks, # how many score-blocks each chunk owns
decode_query_len,
stride_s_h,
stride_s_b,
stride_s_k,
stride_ts_c,
stride_ts_h,
stride_ts_b,
stride_ts_t,
stride_ti_c,
stride_ti_h,
stride_ti_b,
stride_ti_t,
BLOCK_SIZE_K: tl.constexpr,
BLOCK_SIZE_T: tl.constexpr,
USE_PDL: tl.constexpr,
):
tl.static_assert(topk < BLOCK_SIZE_K)
pid_b = tl.program_id(0) # flattened query-token id
pid_h = tl.program_id(1)
pid_chunk = tl.program_id(2)
req_id = pid_b // decode_query_len
q_offset = pid_b - req_id * decode_query_len
if USE_PDL:
tl.extra.cuda.gdc_wait()
seq_len = tl.load(seq_lens + req_id)
query_pos = seq_len - decode_query_len + q_offset
# Full-CG padding uses zero-length request rows. Clamp to an empty
# attention range instead of letting padded rows produce negative lengths.
kv_len = tl.maximum(query_pos + 1, 0)
num_blocks = (kv_len + block_size - 1) // block_size
# Slice this chunk owns within [0, num_blocks).
chunk_start = pid_chunk * chunk_blocks
chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks)
chunk_actual = tl.maximum(chunk_end - chunk_start, 0)
off_k = tl.arange(0, BLOCK_SIZE_K)
off_t = tl.arange(0, BLOCK_SIZE_T)
s_ptrs = (
s_ptr
+ pid_b * stride_s_b
+ pid_h * stride_s_h
+ (chunk_start + off_k) * stride_s_k
)
topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
# Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty
# chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0.
for i in tl.range(0, chunk_actual, BLOCK_SIZE_K):
mask = off_k < chunk_actual - i
score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32)
score = tl.where(score != score, -1e30, score)
s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
topk_score, last_topk_score = score, topk_score
topk_idx, last_topk_idx = (
tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global
topk_idx,
)
n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
for j in tl.static_range(1, n_dims):
topk_score, topk_idx = _bitonic_merge(
topk_score, topk_idx.to(tl.int32), j, 2, n_dims
)
if i != 0:
topk_score, topk_idx = _bitonic_merge(
topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
)
topk_score_new = last_topk_score * left_half_mask + topk_score * (
1 - left_half_mask
)
topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
1 - left_half_mask
)
topk_score, topk_idx = _bitonic_merge(
topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
)
else:
topk_score, topk_idx = _bitonic_merge(
topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
# Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort).
topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
final_score = tl.sum(
topk_mask_extract[:, None]
* tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
axis=0,
)
final_idx = tl.sum(
topk_mask_extract[:, None]
* tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
axis=0,
)
# Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0
# sentinels and lose to real scores in the merge stage.
ts_ptrs = (
ts_partial_ptr
+ pid_chunk * stride_ts_c
+ pid_b * stride_ts_b
+ pid_h * stride_ts_h
+ off_t * stride_ts_t
)
ti_ptrs = (
ti_partial_ptr
+ pid_chunk * stride_ti_c
+ pid_b * stride_ti_b
+ pid_h * stride_ti_h
+ off_t * stride_ti_t
)
tl.store(ts_ptrs, final_score)
tl.store(ti_ptrs, final_idx)
@triton.heuristics(
{
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]),
"BLOCK_SIZE_K": lambda args: triton.next_power_of_2(
args["num_topk_chunks"] * triton.next_power_of_2(args["topk"])
),
}
)
@triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"])
def _topk_index_merge_kernel(
ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape
ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk]
seq_lens, # [num_reqs]
block_size: tl.constexpr, # sparse block size (128)
topk: tl.constexpr,
decode_query_len,
stride_ts_c,
stride_ts_h,
stride_ts_b,
stride_ts_t,
stride_ti_c,
stride_ti_h,
stride_ti_b,
stride_ti_t,
stride_tif_h,
stride_tif_b,
stride_tif_t,
num_topk_chunks,
BLOCK_SIZE_K: tl.constexpr,
BLOCK_SIZE_T: tl.constexpr,
USE_PDL: tl.constexpr,
):
pid_b = tl.program_id(0) # flattened query-token id
pid_h = tl.program_id(1)
req_id = pid_b // decode_query_len
q_offset = pid_b - req_id * decode_query_len
if USE_PDL:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
seq_len = tl.load(seq_lens + req_id)
query_pos = seq_len - decode_query_len + q_offset
# Full-CG padding uses zero-length request rows. Clamp to an empty
# attention range instead of letting padded rows produce negative lengths.
kv_len = tl.maximum(query_pos + 1, 0)
num_blocks = (kv_len + block_size - 1) // block_size
# Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K.
# Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T,
# in_chunk = p % BLOCK_SIZE_T.
off = tl.arange(0, BLOCK_SIZE_K)
chunk_idx = off // BLOCK_SIZE_T
in_chunk_idx = off % BLOCK_SIZE_T
valid = chunk_idx < num_topk_chunks
score_offset = (
chunk_idx * stride_ts_c
+ pid_h * stride_ts_h
+ pid_b * stride_ts_b
+ in_chunk_idx * stride_ts_t
)
idx_offset = (
chunk_idx * stride_ti_c
+ pid_h * stride_ti_h
+ pid_b * stride_ti_b
+ in_chunk_idx * stride_ti_t
)
score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to(
tl.float32
)
score = tl.where(score != score, -1e30, score)
idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32)
# Full bitonic descending sort of BLOCK_SIZE_K items.
n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
for j in tl.static_range(1, n_dims):
score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims)
score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims)
# Extract first BLOCK_SIZE_T positions — these are the global top-K.
extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
topk_idx_final = tl.sum(
extract_mask[:, None]
* tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
axis=0,
)
off_t = tl.arange(0, BLOCK_SIZE_T)
tif_ptrs = (
ti_final_ptr
+ pid_h * stride_tif_h
+ pid_b * stride_tif_b
+ off_t * stride_tif_t
)
store_mask = off_t < topk
topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1)
tl.store(
tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask
)
# ---------------------------------------------------------------------------
# Python wrappers
# ---------------------------------------------------------------------------
@torch.no_grad()
def minimax_m3_index_score(
idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
block_table: torch.Tensor, # [batch, max_blocks]
cu_seqlens_q: torch.Tensor, # [batch+1] int32
seq_lens: torch.Tensor, # [batch] int32
prefix_lens: torch.Tensor, # [batch] int32
max_query_len: int,
max_seq_len: int,
num_kv_heads: int,
sm_scale: float,
) -> torch.Tensor:
"""Compute per-token index scores for each visible sparse block.
Returns score [num_kv_heads, total_q, max_block], where each score is the
max over a 128-token index-K block. M3 has num_idx_heads == num_kv_heads.
"""
total_q, num_idx_heads, head_dim = idx_q.shape
assert num_idx_heads == num_kv_heads, (
"M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
)
batch = cu_seqlens_q.shape[0] - 1
max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
# Keep score strides 16-divisible to avoid Triton recompiles.
score_block_stride = round_up(max_block, 16)
score = torch.empty(
(num_idx_heads, total_q, score_block_stride),
dtype=torch.float32,
device=idx_q.device,
)
BLOCK_SIZE_Q = 64
grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads)
_index_block_score_kernel[grid_score](
idx_q,
index_kv_cache,
score,
block_table,
cu_seqlens_q,
seq_lens,
prefix_lens,
num_idx_heads,
head_dim,
sm_scale,
idx_q.stride(0),
idx_q.stride(1),
idx_q.stride(2),
index_kv_cache.stride(0),
index_kv_cache.stride(1),
index_kv_cache.stride(2),
score.stride(0),
score.stride(1),
score.stride(2),
block_table.stride(0),
BLOCK_SIZE_Q=BLOCK_SIZE_Q,
BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
)
return score
@torch.no_grad()
def minimax_m3_index_topk(
score: torch.Tensor, # [num_idx_heads, total_q, max_block]
cu_seqlens_q: torch.Tensor, # [batch+1] int32
prefix_lens: torch.Tensor, # [batch] int32
max_query_len: int,
topk: int,
init_blocks: int,
local_blocks: int,
) -> torch.Tensor:
"""Select index top-k from a precomputed score tensor."""
num_idx_heads = score.shape[0]
batch = cu_seqlens_q.shape[0] - 1
total_q = score.shape[1]
topk_idx = torch.empty(
(num_idx_heads, total_q, topk),
dtype=torch.int32,
device=score.device,
)
# block_size_q == 1 -> query blocks coincide with query tokens.
grid_topk = (max_query_len, batch, num_idx_heads)
_topk_index_kernel[grid_topk](
score,
topk_idx,
1, # sample_interval (block_size_q)
SPARSE_BLOCK_SIZE,
cu_seqlens_q,
cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
prefix_lens,
topk,
init_blocks,
local_blocks,
score.stride(0),
score.stride(1),
score.stride(2),
topk_idx.stride(0),
topk_idx.stride(1),
topk_idx.stride(2),
MASK_INIT=False,
MASK_LOCAL=False,
)
return topk_idx
@torch.no_grad()
def minimax_m3_index_decode(
idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
block_table: torch.Tensor, # [num_reqs, max_blocks]
seq_lens: torch.Tensor, # [num_reqs] int32
max_seq_len: int,
topk: int,
init_blocks: int,
local_blocks: int,
num_kv_heads: int,
sm_scale: float,
decode_query_len: int,
) -> torch.Tensor:
"""Decode index block-score + top-k, both split-K (cudagraph-safe).
Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
"""
total_q, num_idx_heads, head_dim = idx_q.shape
assert num_idx_heads == num_kv_heads, (
"M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
)
assert total_q == seq_lens.shape[0] * decode_query_len
batch = total_q
max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
use_pdl = current_platform.is_arch_support_pdl()
# `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
# SM9+); this ROCm Triton rejects it even when False ("Keyword argument
# launch_pdl was specified but unrecognised"). Only pass it when PDL is
# actually supported -- on ROCm use_pdl is always False, so it's omitted.
pdl_launch = {"launch_pdl": True} if use_pdl else {}
# Keep score strides 16-divisible to avoid Triton recompiles.
score_block_stride = round_up(max_block, 16)
score = torch.empty(
(num_idx_heads, total_q, score_block_stride),
dtype=torch.float32,
device=idx_q.device,
)
# split-K over seq blocks; chunk count depends only on shape constants so
# the grid is fixed within a cuda graph.
TARGET_GRID = 4096
MAX_NUM_KV_CHUNKS = 256
target = max(
1, min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, batch * num_idx_heads))
)
num_kv_chunks = 1 << (target.bit_length() - 1)
grid_score = (batch, num_kv_chunks)
_decode_index_score_kernel[grid_score](
idx_q,
index_kv_cache,
score,
block_table,
seq_lens,
num_idx_heads,
head_dim,
init_blocks,
local_blocks,
sm_scale,
decode_query_len,
idx_q.stride(0),
idx_q.stride(1),
idx_q.stride(2),
index_kv_cache.stride(0),
index_kv_cache.stride(1),
index_kv_cache.stride(2),
score.stride(0),
score.stride(1),
score.stride(2),
block_table.stride(0),
BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
num_kv_chunks=num_kv_chunks,
USE_PDL=use_pdl,
**pdl_launch,
)
topk_idx = torch.empty(
(num_idx_heads, total_q, topk),
dtype=torch.int32,
device=idx_q.device,
)
# Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
# pow2(num_topk_chunks * pow2(topk)) candidates.
TOPK_TARGET_GRID = 64
MAX_NUM_TOPK_CHUNKS = 16
topk_target = max(
1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))
)
num_topk_chunks = 1 << (topk_target.bit_length() - 1)
block_size_t = triton.next_power_of_2(topk)
chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks
topk_score_partial = torch.empty(
num_topk_chunks,
num_idx_heads,
batch,
block_size_t,
dtype=torch.float32,
device=idx_q.device,
)
topk_idx_partial = torch.empty(
num_topk_chunks,
num_idx_heads,
batch,
block_size_t,
dtype=torch.int32,
device=idx_q.device,
)
_topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)](
score,
topk_score_partial,
topk_idx_partial,
seq_lens,
SPARSE_BLOCK_SIZE,
topk,
chunk_blocks,
decode_query_len,
score.stride(0),
score.stride(1),
score.stride(2),
topk_score_partial.stride(0),
topk_score_partial.stride(1),
topk_score_partial.stride(2),
topk_score_partial.stride(3),
topk_idx_partial.stride(0),
topk_idx_partial.stride(1),
topk_idx_partial.stride(2),
topk_idx_partial.stride(3),
USE_PDL=use_pdl,
**pdl_launch,
)
_topk_index_merge_kernel[(batch, num_idx_heads)](
topk_score_partial,
topk_idx_partial,
topk_idx,
seq_lens,
SPARSE_BLOCK_SIZE,
topk,
decode_query_len,
topk_score_partial.stride(0),
topk_score_partial.stride(1),
topk_score_partial.stride(2),
topk_score_partial.stride(3),
topk_idx_partial.stride(0),
topk_idx_partial.stride(1),
topk_idx_partial.stride(2),
topk_idx_partial.stride(3),
topk_idx.stride(0),
topk_idx.stride(1),
topk_idx.stride(2),
num_topk_chunks=num_topk_chunks,
USE_PDL=use_pdl,
**pdl_launch,
)
return topk_idx
vllm/models/minimax_m3/common/ops/sparse_attn.py
+593
View File
@@ -0,0 +1,593 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Triton kernels for MiniMax M3 block-sparse GQA attention.
The main heads attend only to the blocks selected by the lightning indexer (see
``index_topk``). Adapted to vLLM's paged KV cache: the KV page size is forced to
equal the sparse block size (128), so one selected block maps to exactly one
page.
Main K/V cache layout (vLLM):
``(num_blocks, 2, 128, num_kv_heads, head_dim)`` K=[:,0] V=[:,1]
Only the paths MiniMax M3 uses are implemented: no attention sink, base-2
(exp2/log2) softmax. The decode kernels use split-K (flash-decoding) over the
selected blocks with a separate merge step, since one query token per request
leaves the prefill kernels (which parallelize over the query dim) idle.
"""
import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
# One sparse block == one KV page.
SPARSE_BLOCK_SIZE = 128
_SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None
def _sparse_attn_num_stages_kwarg() -> dict:
"""Triton ``num_stages`` override for the sparse-attn GEMM kernels.
Forced only where required: CDNA3 (gfx942) caps LDS at
64 KB, and the default 2-stage pipeline double-buffers the 128x128 K/V tiles
to ~66 KB ("out of resource: shared memory"), so pin gfx942 to a single
stage (~32 KB, which fits). Everywhere else (NVIDIA, CDNA4 gfx950) return an
empty kwarg and let Triton keep its own default -- don't second-guess it.
Cached: the arch is fixed per process.
"""
global _SPARSE_ATTN_NUM_STAGES_KWARG
if _SPARSE_ATTN_NUM_STAGES_KWARG is None:
kwarg: dict = {}
if current_platform.is_rocm():
from vllm.platforms.rocm import on_gfx942
if on_gfx942():
kwarg = {"num_stages": 1}
_SPARSE_ATTN_NUM_STAGES_KWARG = kwarg
return _SPARSE_ATTN_NUM_STAGES_KWARG
# ---------------------------------------------------------------------------
# GQA block-sparse attention (paged). Main heads attend only to the selected
# blocks. BLOCK_SIZE_K == 128 so each selected block is one page.
# ---------------------------------------------------------------------------
# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
# might lose pointer alignment, which trigger Triton recompiles. we don't actually
# need pointer alignment for those tensors anyway because we do scalar load.
@triton.heuristics(
{
"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
"BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
"BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]
* triton.next_power_of_2(args["gqa_group_size"]),
}
)
@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
def _gqa_sparse_fwd_kernel(
q_ptr, # [total_q, num_heads, head_dim]
kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
o_ptr, # [total_q, num_heads, head_dim]
block_table_ptr, # [num_reqs, max_blocks]
cu_seqlens_q,
cu_seqblocks_q,
seq_lens,
prefix_lens,
num_kv_heads,
gqa_group_size,
head_dim,
max_topk,
num_q_loop,
sm_scale,
stride_qn,
stride_qh,
stride_qd,
stride_kv_blk,
stride_kv_kv,
stride_kv_pos,
stride_kv_h,
stride_kv_d,
stride_th,
stride_tn,
stride_tk,
stride_on,
stride_oh,
stride_od,
stride_bt_b,
BLOCK_SIZE_Q: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
BLOCK_SIZE_D: tl.constexpr,
BLOCK_SIZE_H: tl.constexpr,
BLOCK_SIZE_T: tl.constexpr,
BLOCK_SIZE_QH: tl.constexpr,
USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
):
sm_scale_log2e = sm_scale * 1.4426950409
pid_q = tl.program_id(0)
pid_kh = tl.program_id(1)
pid_b = tl.program_id(2)
pid_h = pid_kh * gqa_group_size
q_start = tl.load(cu_seqlens_q + pid_b)
q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start
q_block_start = tl.load(cu_seqblocks_q + pid_b)
q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start
seq_len = tl.load(seq_lens + pid_b)
prefix_len = tl.load(prefix_lens + pid_b)
if pid_q * num_q_loop >= q_block_len:
return
real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop)
bt_row = block_table_ptr + pid_b * stride_bt_b
off_n = tl.arange(0, BLOCK_SIZE_K)
off_d = tl.arange(0, BLOCK_SIZE_D)
d_mask = off_d < head_dim
for j in range(real_q_loop):
pid_q_j = pid_q * num_q_loop + j
t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
off_t = tl.arange(0, BLOCK_SIZE_T)
topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)
real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
q_ptrs = tl.make_block_ptr(
base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
shape=(q_len, gqa_group_size, head_dim),
strides=(stride_qn, stride_qh, stride_qd),
offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(2, 1, 0),
)
q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
off_q = (
tl.arange(0, BLOCK_SIZE_Q)[:, None]
+ pid_q_j * BLOCK_SIZE_Q
+ prefix_len
- tl.arange(0, BLOCK_SIZE_K)[None, :]
)
m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32)
q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D)
for _ in range(real_topk):
blk = tl.load(t_ptr_j).to(tl.int32)
t_ptr_j = t_ptr_j + stride_tk
c = blk * BLOCK_SIZE_K
page = tl.load(bt_row + blk).to(tl.int64)
pos = c + off_n
pos_mask = pos < seq_len
k = tl.load(
kv_cache_ptr
+ page * stride_kv_blk
+ 0 * stride_kv_kv
+ off_n[None, :] * stride_kv_pos
+ pid_kh * stride_kv_h
+ off_d[:, None] * stride_kv_d,
mask=d_mask[:, None] & pos_mask[None, :],
other=0.0,
)
if USE_FP8:
k = k.to(q.dtype)
qk = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
# causal: q_abs_pos - k_off >= block_start (c)
qk += tl.where(off_q[:, None, :] >= c, 0, float("-inf"))
qk = tl.reshape(qk, BLOCK_SIZE_QH, BLOCK_SIZE_K)
qk += tl.dot(q, k) * sm_scale_log2e
qk += tl.where(pos_mask[None, :], 0, float("-inf"))
m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
p = tl.exp2(qk - m_ij[:, None])
l_ij = tl.sum(p, axis=1)
acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
v = tl.load(
kv_cache_ptr
+ page * stride_kv_blk
+ 1 * stride_kv_kv
+ off_n[:, None] * stride_kv_pos
+ pid_kh * stride_kv_h
+ off_d[None, :] * stride_kv_d,
mask=pos_mask[:, None] & d_mask[None, :],
other=0.0,
)
if USE_FP8:
v = v.to(q.dtype)
acc_o += tl.dot(p.to(v.dtype), v)
m_i = m_ij
lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
o_ptrs = tl.make_block_ptr(
base=o_ptr + q_start * stride_on + pid_h * stride_oh,
shape=(q_len, gqa_group_size, head_dim),
strides=(stride_on, stride_oh, stride_od),
offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(2, 1, 0),
)
tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
# ---------------------------------------------------------------------------
# Decode kernels (split-K). Decode batches are flattened request-major, with a
# runtime query length used to map each query token back to its request metadata.
# This parallelizes over the selected top-k blocks, producing partials that the
# merge kernel combines (flash-decoding). All chunk counts depend only on shape
# constants so the grid is fixed within a cuda graph. Base-2 (exp2/log2)
# softmax matches the prefill kernel.
# ---------------------------------------------------------------------------
@triton.heuristics(
{
"BLOCK_SIZE_H": lambda args: max(
16, triton.next_power_of_2(args["gqa_group_size"])
),
"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
}
)
@triton.jit(do_not_specialize=["decode_query_len"])
def _gqa_sparse_decode_kernel(
q_ptr, # [total_q, num_heads, head_dim]
kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
o_ptr, # partial out: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
lse_ptr, # partial lse (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
block_table_ptr, # [num_reqs, max_blocks]
seq_lens, # [num_reqs]
total_q,
gqa_group_size,
head_dim,
max_topk,
sm_scale,
decode_query_len,
stride_qn,
stride_qh,
stride_qd,
stride_kv_blk,
stride_kv_kv,
stride_kv_pos,
stride_kv_h,
stride_kv_d,
stride_th,
stride_tn,
stride_tk,
stride_o_c,
stride_o_b,
stride_o_h,
stride_o_d,
stride_l_c,
stride_l_b,
stride_l_h,
stride_bt_b,
BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
NUM_TOPK_CHUNKS: tl.constexpr,
BLOCK_SIZE_H: tl.constexpr,
BLOCK_SIZE_D: tl.constexpr,
BLOCK_SIZE_T: tl.constexpr,
USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
USE_PDL: tl.constexpr,
):
sm_scale_log2e = sm_scale * 1.4426950409
# split-K over the topk dimension: pid(0) folds (query-token, chunk).
pid_bc, pid_kh = tl.program_id(0), tl.program_id(1)
pid_b = pid_bc % total_q
pid_c = pid_bc // total_q
req_id = pid_b // decode_query_len
q_offset = pid_b - req_id * decode_query_len
pid_h = pid_kh * gqa_group_size
chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS
chunk_start_topk = pid_c * chunk_size_topk
chunk_end_compiletime = chunk_start_topk + chunk_size_topk
if USE_PDL:
tl.extra.cuda.gdc_wait()
seq_len = tl.load(seq_lens + req_id)
query_pos = seq_len - decode_query_len + q_offset
# Full-CG padding uses zero-length request rows. Clamp to an empty
# attention range instead of letting padded rows produce negative lengths.
kv_len = tl.maximum(query_pos + 1, 0)
# number of valid (non-padded) selected blocks for this query token
off_t = tl.arange(0, BLOCK_SIZE_T)
idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn
topk_idx = tl.load(idx_base + off_t * stride_tk, mask=off_t < max_topk, other=-1)
real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
chunk_end_topk = tl.minimum(chunk_end_compiletime, real_topk)
off_n = tl.arange(0, BLOCK_SIZE_K)
off_d = tl.arange(0, BLOCK_SIZE_D)
d_mask = off_d < head_dim
bt_row = block_table_ptr + req_id * stride_bt_b
m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32)
q_ptrs = tl.make_block_ptr(
base=q_ptr + pid_b * stride_qn + pid_h * stride_qh,
shape=(gqa_group_size, head_dim),
strides=(stride_qh, stride_qd),
offsets=(0, 0),
block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(1, 0),
)
q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero")
cur_idx_ptr = idx_base + chunk_start_topk * stride_tk
for _ in tl.range(chunk_start_topk, chunk_end_topk):
blk = tl.load(cur_idx_ptr).to(tl.int32)
cur_idx_ptr = cur_idx_ptr + stride_tk
c = blk * BLOCK_SIZE_K
page = tl.load(bt_row + blk).to(tl.int64)
pos = c + off_n
pos_mask = pos < kv_len
k = tl.load(
kv_cache_ptr
+ page * stride_kv_blk
+ 0 * stride_kv_kv
+ off_n[None, :] * stride_kv_pos
+ pid_kh * stride_kv_h
+ off_d[:, None] * stride_kv_d,
mask=d_mask[:, None] & pos_mask[None, :],
other=0.0,
)
if USE_FP8:
k = k.to(q.dtype)
qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
qk += tl.where(pos_mask[None, :], 0, float("-inf"))
qk += tl.dot(q, k) * sm_scale_log2e
m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
p = tl.exp2(qk - m_ij[:, None])
l_ij = tl.sum(p, axis=1)
acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
v = tl.load(
kv_cache_ptr
+ page * stride_kv_blk
+ 1 * stride_kv_kv
+ off_n[:, None] * stride_kv_pos
+ pid_kh * stride_kv_h
+ off_d[None, :] * stride_kv_d,
mask=pos_mask[:, None] & d_mask[None, :],
other=0.0,
)
if USE_FP8:
v = v.to(q.dtype)
acc_o += tl.dot(p.to(v.dtype), v)
m_i = m_ij
lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
# Empty chunks for active rows must store zero output; otherwise the merge
# can hit 0 * NaN. All-empty padded rows may still produce NaNs in merge.
scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), tl.zeros_like(lse_i))
acc_o = acc_o * scale[:, None]
o_ptrs = tl.make_block_ptr(
base=o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h,
shape=(gqa_group_size, head_dim),
strides=(stride_o_h, stride_o_d),
offsets=(0, 0),
block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(1, 0),
)
tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1))
lse_ptrs = tl.make_block_ptr(
base=lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h,
shape=(gqa_group_size,),
strides=(stride_l_h,),
offsets=(0,),
block_shape=(BLOCK_SIZE_H,),
order=(0,),
)
tl.store(lse_ptrs, lse_i.to(lse_ptr.dtype.element_ty), boundary_check=(0,))
@triton.heuristics(
{"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}
)
@triton.jit
def _merge_topk_attn_out_kernel(
o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
out_ptr, # merged out: [total_q, num_heads, head_dim]
head_dim,
stride_o_c,
stride_o_b,
stride_o_h,
stride_o_d,
stride_l_c,
stride_l_b,
stride_l_h,
stride_out_n,
stride_out_h,
stride_out_d,
NUM_TOPK_CHUNKS: tl.constexpr,
BLOCK_SIZE_D: tl.constexpr,
USE_PDL: tl.constexpr,
):
pid_b, pid_h = tl.program_id(0), tl.program_id(1)
# NOTE: assume seq_lens is safe to load before gdc_wait()
if USE_PDL:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
off_c = tl.arange(0, NUM_TOPK_CHUNKS)
off_d = tl.arange(0, BLOCK_SIZE_D)
o_ptrs = tl.make_block_ptr(
base=o_ptr + pid_b * stride_o_b + pid_h * stride_o_h,
shape=(NUM_TOPK_CHUNKS, head_dim),
strides=(stride_o_c, stride_o_d),
offsets=(0, 0),
block_shape=(NUM_TOPK_CHUNKS, BLOCK_SIZE_D),
order=(1, 0),
)
lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c
o = tl.load(o_ptrs, boundary_check=(0, 1), padding_option="zero")
lse = tl.load(lse_ptrs) # empty chunks contribute -inf -> weight 0
lse_max = tl.max(lse, axis=0)
weights = tl.exp2(lse - lse_max)
weights = weights / tl.sum(weights, axis=0)
o_merged = tl.sum(o * weights[:, None], axis=0)
out_ptrs = (
out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d
)
tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim)
# ---------------------------------------------------------------------------
# Python wrappers
# ---------------------------------------------------------------------------
@torch.no_grad()
def minimax_m3_sparse_attn(
q: torch.Tensor, # [total_q, num_heads, head_dim]
kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
block_table: torch.Tensor, # [batch, max_blocks]
cu_seqlens_q: torch.Tensor, # [batch+1] int32
seq_lens: torch.Tensor, # [batch] int32
prefix_lens: torch.Tensor, # [batch] int32
max_query_len: int,
num_kv_heads: int,
sm_scale: float,
output: torch.Tensor, # [total_q, num_heads, head_dim]
) -> None:
"""GQA block-sparse attention over the selected blocks. block_size_q == 1."""
total_q, num_heads, head_dim = q.shape
batch = cu_seqlens_q.shape[0] - 1
topk = topk_idx.shape[-1]
gqa_group_size = num_heads // num_kv_heads
use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
grid = (max_query_len, num_kv_heads, batch)
_gqa_sparse_fwd_kernel[grid](
q,
kv_cache,
topk_idx,
output,
block_table,
cu_seqlens_q,
cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
seq_lens,
prefix_lens,
num_kv_heads,
gqa_group_size,
head_dim,
topk,
1, # num_q_loop
sm_scale,
q.stride(0),
q.stride(1),
q.stride(2),
kv_cache.stride(0),
kv_cache.stride(1),
kv_cache.stride(2),
kv_cache.stride(3),
kv_cache.stride(4),
topk_idx.stride(0),
topk_idx.stride(1),
topk_idx.stride(2),
output.stride(0),
output.stride(1),
output.stride(2),
block_table.stride(0),
BLOCK_SIZE_Q=1,
BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
USE_FP8=use_fp8,
**_sparse_attn_num_stages_kwarg(),
)
@torch.no_grad()
def minimax_m3_sparse_attn_decode(
q: torch.Tensor, # [total_q, num_heads, head_dim]
kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
block_table: torch.Tensor, # [num_reqs, max_blocks]
seq_lens: torch.Tensor, # [num_reqs] int32
num_kv_heads: int,
sm_scale: float,
output: torch.Tensor, # [total_q, num_heads, head_dim]
decode_query_len: int,
) -> None:
"""GQA block-sparse attention for decode (split-K over the top-k blocks)."""
total_q, num_heads, head_dim = q.shape
assert total_q == seq_lens.shape[0] * decode_query_len
max_topk = topk_idx.shape[-1]
gqa_group_size = num_heads // num_kv_heads
use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
use_pdl = current_platform.is_arch_support_pdl()
# `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
# SM9+); this ROCm Triton rejects it even when False ("Keyword argument
# launch_pdl was specified but unrecognised"). Only pass it when PDL is
# actually supported -- on ROCm use_pdl is always False, so it's omitted.
pdl_launch = {"launch_pdl": True} if use_pdl else {}
# split-K over the selected blocks; chunk count is shape-constant (cuda graph).
TARGET_GRID = 256
target = max(1, min(max_topk, TARGET_GRID // max(1, total_q * num_kv_heads)))
num_topk_chunks = 1 << (target.bit_length() - 1)
o_partial = torch.empty(
num_topk_chunks, total_q, num_heads, head_dim, dtype=q.dtype, device=q.device
)
lse_partial = torch.empty(
num_topk_chunks, total_q, num_heads, dtype=torch.float32, device=q.device
)
grid = (total_q * num_topk_chunks, num_kv_heads)
_gqa_sparse_decode_kernel[grid](
q,
kv_cache,
topk_idx,
o_partial,
lse_partial,
block_table,
seq_lens,
total_q,
gqa_group_size,
head_dim,
max_topk,
sm_scale,
decode_query_len,
q.stride(0),
q.stride(1),
q.stride(2),
kv_cache.stride(0),
kv_cache.stride(1),
kv_cache.stride(2),
kv_cache.stride(3),
kv_cache.stride(4),
topk_idx.stride(0),
topk_idx.stride(1),
topk_idx.stride(2),
o_partial.stride(0),
o_partial.stride(1),
o_partial.stride(2),
o_partial.stride(3),
lse_partial.stride(0),
lse_partial.stride(1),
lse_partial.stride(2),
block_table.stride(0),
BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
NUM_TOPK_CHUNKS=num_topk_chunks,
USE_FP8=use_fp8,
USE_PDL=use_pdl,
**_sparse_attn_num_stages_kwarg(),
**pdl_launch,
)
merge_grid = (total_q, num_heads)
_merge_topk_attn_out_kernel[merge_grid](
o_partial,
lse_partial,
output,
head_dim,
o_partial.stride(0),
o_partial.stride(1),
o_partial.stride(2),
o_partial.stride(3),
lse_partial.stride(0),
lse_partial.stride(1),
lse_partial.stride(2),
output.stride(0),
output.stride(1),
output.stride(2),
NUM_TOPK_CHUNKS=num_topk_chunks,
USE_PDL=use_pdl,
**pdl_launch,
)
vllm/models/minimax_m3/common/sparse_attention.py
+398
View File
@@ -0,0 +1,398 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Main block-sparse GQA attention for MiniMax M3 sparse layers.
The lightning indexer (``indexer.py``) selects the top-k KV blocks; this module
holds the main attention that attends only to those blocks: the paged K/V cache
backend, its metadata + builder, and the impl that consumes the indexer's
``topk_idx``. The Triton attend kernel lives here; the SM100 (MSA)
``build_k2q_csr`` + ``sparse_atten_func`` attend lives in
``nvidia/sparse_attention_msa.py``.
``MiniMaxM3SparseBackend`` and ``MiniMaxM3SparseMetadata`` are referenced by the
attention-backend registry (by dotted path) and by spec-decode, so they must
keep these names and stay in this module.
"""
from dataclasses import dataclass
from typing import ClassVar
import torch
from vllm.config import VllmConfig
from vllm.config.cache import CacheDType
from vllm.forward_context import get_forward_context
from vllm.models.minimax_m3.common.ops.sparse_attn import (
SPARSE_BLOCK_SIZE,
minimax_m3_sparse_attn,
minimax_m3_sparse_attn_decode,
)
from vllm.platforms import current_platform
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionImplBase,
AttentionLayer,
AttentionMetadata,
AttentionMetadataBuilder,
CommonAttentionMetadata,
MultipleOf,
)
from vllm.v1.attention.backends.utils import (
get_kv_cache_layout,
split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
class MiniMaxM3SparseBackend(AttentionBackend):
"""Block-sparse GQA backend for MiniMax M3 sparse attention layers."""
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16]
# bf16 or fp8 (e4m3/e5m2): the Triton kernels dequant fp8 before the dots.
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"bfloat16",
"fp8",
"fp8_e4m3",
"fp8_e5m2",
]
@staticmethod
def get_name() -> str:
return "MINIMAX_M3_SPARSE"
@staticmethod
def get_impl_cls() -> type["MiniMaxM3SparseImpl"]:
# Concrete impl chosen by select_main_impl_cls; base for introspection.
return MiniMaxM3SparseImpl
@staticmethod
def get_builder_cls() -> type["MiniMaxM3SparseMetadataBuilder"]:
return MiniMaxM3SparseMetadataBuilder
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return [128]
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
# Page size == sparse block size (one sparse block per KV page).
return [128]
@classmethod
def is_sparse(cls) -> bool:
return True
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
return (num_blocks, 2, block_size, num_kv_heads, head_size)
@staticmethod
def get_kv_cache_stride_order(
include_num_layers_dimension: bool = False,
) -> tuple[int, ...]:
# Permutation from get_kv_cache_shape to the actual memory layout.
if include_num_layers_dimension:
raise NotImplementedError # no cross-layer KV blocks in M3
cache_layout = get_kv_cache_layout()
if cache_layout == "NHD":
stride_order = (0, 1, 2, 3, 4)
elif cache_layout == "HND":
stride_order = (0, 1, 3, 2, 4)
else:
raise ValueError(f"Unknown cache layout format {cache_layout}.")
return stride_order
@dataclass
class MiniMaxM3SparsePrefillMetadata:
"""Per-prefill state; ``cu_seqlens_k``/``total_kv_blocks`` feed the MSA CSR."""
cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0
cu_seqlens_k: torch.Tensor # [num_prefills + 1] int32, cumulative KV lengths
seq_lens: torch.Tensor # [num_prefills] int32, total KV lengths
context_lens: torch.Tensor # [num_prefills] int32 (cached/context tokens)
block_table: torch.Tensor
max_query_len: int
max_seq_len: int
total_kv_blocks: int
@dataclass
class MiniMaxM3SparseDecodeMetadata:
"""Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform
per-request query length (1, or 1 + num_speculative_tokens)."""
seq_lens: torch.Tensor # [num_decodes] int32
block_table: torch.Tensor
decode_query_len: int
@dataclass
class MiniMaxM3SparseMetadata(AttentionMetadata):
"""Sparse-attention metadata, split into prefill and decode sub-metadata."""
seq_lens: torch.Tensor
max_seq_len: int
slot_mapping: torch.Tensor
num_actual_tokens: int # total query tokens (decode-first batch)
# Split counts (batch reordered decode-first).
num_decodes: int
num_decode_tokens: int
num_prefills: int
num_prefill_tokens: int
prefill: MiniMaxM3SparsePrefillMetadata | None = None
decode: MiniMaxM3SparseDecodeMetadata | None = None
class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMetadata]):
# Full cudagraphs for uniform decode batches, incl. spec-decode verify
# batches with >1 query token/request.
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
# Raised to 1 + num_speculative_tokens by _init_reorder_batch_threshold when
# spec decode is on; must match the indexer builder so the splits agree.
reorder_batch_threshold: int = 1
def __init__(
self,
kv_cache_spec: AttentionSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
) -> None:
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
# Stable context-length buffer for decode cudagraph replays.
self.context_len_buffer = torch.empty(
vllm_config.scheduler_config.max_num_batched_tokens,
dtype=torch.int32,
device=device,
)
def build(
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
) -> MiniMaxM3SparseMetadata:
num_reqs = common_attn_metadata.num_reqs
num_tokens = common_attn_metadata.num_actual_tokens
query_start_loc = common_attn_metadata.query_start_loc
seq_lens = common_attn_metadata.seq_lens
block_table = common_attn_metadata.block_table_tensor
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
split_decodes_and_prefills(
common_attn_metadata,
decode_threshold=self.reorder_batch_threshold,
require_uniform=True,
)
)
assert num_decodes + num_prefills == num_reqs
assert num_decode_tokens + num_prefill_tokens == num_tokens
# Decode-first batch: context lengths into the stable cudagraph buffer.
context_lens = self.context_len_buffer[:num_reqs]
context_lens.copy_(
common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
)
prefill_metadata: MiniMaxM3SparsePrefillMetadata | None = None
if num_prefills > 0:
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
assert seq_lens_cpu is not None
prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]
prefill_total_kv_blocks = (
((prefill_seq_lens_cpu + SPARSE_BLOCK_SIZE - 1) // SPARSE_BLOCK_SIZE)
.sum()
.item()
)
prefill_kv_lens = seq_lens[num_decodes:]
prefill_cu_seqlens_k = torch.empty(
num_prefills + 1, dtype=torch.int32, device=seq_lens.device
)
prefill_cu_seqlens_k[0] = 0
torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:])
prefill_metadata = MiniMaxM3SparsePrefillMetadata(
cu_seqlens_q=(query_start_loc[num_decodes:] - num_decode_tokens).to(
torch.int32
),
cu_seqlens_k=prefill_cu_seqlens_k,
seq_lens=prefill_kv_lens,
context_lens=context_lens[num_decodes:],
block_table=block_table[num_decodes:],
max_query_len=common_attn_metadata.max_query_len,
max_seq_len=common_attn_metadata.max_seq_len,
total_kv_blocks=prefill_total_kv_blocks,
)
decode_metadata: MiniMaxM3SparseDecodeMetadata | None = None
if num_decodes > 0:
qsl_cpu = common_attn_metadata.query_start_loc_cpu
query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
decode_query_len = int(query_lens_cpu[0].item())
assert decode_query_len > 0
assert torch.all(
(query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
)
assert num_decode_tokens == num_decodes * decode_query_len
decode_metadata = MiniMaxM3SparseDecodeMetadata(
seq_lens=seq_lens[:num_decodes],
block_table=block_table[:num_decodes],
decode_query_len=decode_query_len,
)
return MiniMaxM3SparseMetadata(
seq_lens=seq_lens,
max_seq_len=common_attn_metadata.max_seq_len,
slot_mapping=common_attn_metadata.slot_mapping,
num_actual_tokens=num_tokens,
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
prefill=prefill_metadata,
decode=decode_metadata,
)
class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
"""Abstract base for block-sparse GQA over the indexer-selected blocks.
Inherits ``AttentionImplBase`` for a custom forward signature (the layer
pre-inserts K/V and runs the indexer, so forward takes the queries +
``topk_idx``). The Triton and MSA subclasses each own a full ``forward`` --
no shared forward code.
"""
def __init__(
self,
num_heads: int,
head_size: int,
scale: float,
num_kv_heads: int | None = None,
kv_cache_dtype: str = "auto",
*,
topk_blocks: int,
sparse_block_size: int,
) -> None:
self.num_heads = num_heads
self.head_size = head_size
self.scale = scale
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
self.kv_cache_dtype = kv_cache_dtype
self.use_fp8_kv = is_quantized_kv_cache(kv_cache_dtype)
self.kv_cache_fp8_dtype = (
torch.float8_e5m2 if "e5m2" in kv_cache_dtype else torch.float8_e4m3fn
)
# Sparse selection parameters (block_size == page size == SPARSE_BLOCK_SIZE).
self.topk_blocks = topk_blocks
self.block_size = sparse_block_size
def forward(
self,
layer: AttentionLayer,
query: torch.Tensor,
kv_cache: torch.Tensor,
topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
output: torch.Tensor,
) -> torch.Tensor:
"""Attend the queries to the indexer-selected blocks. Per kernel."""
raise NotImplementedError
class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
"""Triton block-sparse attend (``minimax_m3_sparse_attn``) + Triton decode."""
def forward(
self,
layer: AttentionLayer,
query: torch.Tensor,
kv_cache: torch.Tensor,
topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
output: torch.Tensor,
) -> torch.Tensor:
attn_metadata = get_forward_context().attn_metadata
if not isinstance(attn_metadata, dict):
return output # profiling run; caches unbound
main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined]
assert isinstance(main_md, MiniMaxM3SparseMetadata)
decode_topk, prefill_topk = topk_idx
nd = main_md.num_decode_tokens
num_tokens = main_md.num_actual_tokens
hd = self.head_size
q = query[:num_tokens].view(-1, self.num_heads, hd)
out = output[:num_tokens].view(-1, self.num_heads, hd)
kv_cache = (
kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
)
# Decode [:nd]: split-K over the selected blocks (request-major chunks).
if main_md.num_decodes > 0:
d = main_md.decode
assert d is not None and decode_topk is not None
minimax_m3_sparse_attn_decode(
q[:nd],
kv_cache,
decode_topk,
d.block_table,
d.seq_lens,
self.num_kv_heads,
self.scale,
out[:nd],
d.decode_query_len,
)
# Prefill [nd:]: cu_seqlens_q already rebased to 0.
if main_md.num_prefills > 0:
p = main_md.prefill
assert p is not None and prefill_topk is not None
minimax_m3_sparse_attn(
q[nd:],
kv_cache,
prefill_topk,
p.block_table,
p.cu_seqlens_q,
p.seq_lens,
p.context_lens,
p.max_query_len,
self.num_kv_heads,
self.scale,
out[nd:],
)
return output
def select_main_impl_cls(
*,
topk_blocks: int,
kv_cache_dtype: str,
) -> type[MiniMaxM3SparseImpl]:
"""Pick the main attend impl off the main KV-cache dtype.
bf16 on Blackwell (SM100) uses the MSA attend; fp8 or non-Blackwell falls
back to Triton. The MSA module is imported lazily so AMD/non-SM100 never
import fmha_sm100.
"""
if (
current_platform.is_cuda()
and current_platform.is_device_capability_family(100)
and topk_blocks in (4, 8, 16, 32)
and not is_quantized_kv_cache(kv_cache_dtype)
):
from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
MiniMaxM3SparseMSAImpl,
)
return MiniMaxM3SparseMSAImpl
return MiniMaxM3SparseTritonImpl
vllm/models/minimax_m3/common/vision_tower.py
+765
View File
@@ -0,0 +1,765 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
import numpy as np
import torch
import torch.nn as nn
from einops import rearrange
from transformers import PretrainedConfig
from vllm.distributed import parallel_state
from vllm.distributed import utils as dist_utils
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
)
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.utils import maybe_prefix
from vllm.model_executor.models.vision import (
get_vit_attn_backend,
is_vit_use_data_parallel,
)
from vllm.platforms import current_platform
# ROCm caps a kernel-launch gridDim.y at 65536. The HIP flash-attn Triton
# rotary kernel launches grid.y = cdiv(seqlen, BLOCK_M), so it fails with
# hipErrorInvalidValue once cdiv(seqlen, BLOCK_M) > 65536. Used below to decide
# when RoPE must be applied per video segment instead of in one launch.
_HIP_MAX_GRID_DIM_Y = 65536
class MiniMaxVLPatchEmbed(nn.Module):
"""Conv3d-based patch embedding.
Takes flat tokens of shape (N, C * temporal_patch_size * patch_size²)
and projects each to a hidden-size embedding.
"""
def __init__(self, config: PretrainedConfig) -> None:
super().__init__()
compression = config.img_token_compression_config
temporal_patch_size = compression.get("temporal_patch_size", 2)
patch_size = config.patch_size
num_channels = config.num_channels
self.patch_size = patch_size
self.temporal_patch_size = temporal_patch_size
self.num_channels = num_channels
self.hidden_size = config.hidden_size
self.patch_embedding = nn.Conv3d(
in_channels=num_channels,
out_channels=config.hidden_size,
kernel_size=(temporal_patch_size, patch_size, patch_size),
stride=(temporal_patch_size, patch_size, patch_size),
bias=False,
)
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
# pixel_values: (N, C * temporal_patch_size * patch_size²)
if self.patch_embedding.weight.dtype != pixel_values.dtype:
self.patch_embedding = self.patch_embedding.to(pixel_values.dtype)
x = pixel_values.reshape(
pixel_values.shape[0],
self.num_channels,
self.temporal_patch_size,
self.patch_size,
self.patch_size,
)
return self.patch_embedding(x).reshape(x.shape[0], -1)
class MiniMaxVLAttention(nn.Module):
"""Multi-head attention with MiniMax's partial 3D RoPE.
Partial means only the first ``rot_dim`` (< head_dim) dimensions of
Q and K are rotated; the remaining dims are passed through unchanged.
"""
def __init__(
self,
embed_dim: int,
num_heads: int,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
use_data_parallel = is_vit_use_data_parallel()
self.tp_size = (
1
if use_data_parallel
else parallel_state.get_tensor_model_parallel_world_size()
)
self.head_dim = embed_dim // num_heads
self.num_heads_per_partition = dist_utils.divide(num_heads, self.tp_size)
self.qkv_proj = QKVParallelLinear(
hidden_size=embed_dim,
head_size=self.head_dim,
total_num_heads=num_heads,
total_num_kv_heads=num_heads,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
disable_tp=use_data_parallel,
)
self.out_proj = RowParallelLinear(
input_size=embed_dim,
output_size=embed_dim,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.out_proj",
disable_tp=use_data_parallel,
)
self.attn = MMEncoderAttention(
num_heads=self.num_heads_per_partition,
head_size=self.head_dim,
prefix=f"{prefix}.attn",
)
# ApplyRotaryEmb handles the internal cos/sin repeat and partial
# rotation (ro_dim = half_rot_dim * 2 < head_dim for MiniMax).
# enable_fp32_compute=True runs the rotation in fp32 (q/k upcast,
# fp32 cos/sin), matching the reference ``_minimax_rope_applier``.
self.apply_rotary_emb = ApplyRotaryEmb(
enforce_enable=True, enable_fp32_compute=True
)
def _apply_rotary_emb(
self,
qk_reshaped: torch.Tensor,
rotary_cos: torch.Tensor,
rotary_sin: torch.Tensor,
seq_len: int,
rotary_segment_lengths: list[int] | None,
) -> torch.Tensor:
# Default fast path (all NVIDIA inputs, and ROCm short clips/images):
# a single rotary kernel launch. ``rotary_segment_lengths`` is only
# populated on ROCm (see ``MiniMaxVLVisionTransformer.forward``), so
# the per-segment path below is ROCm-only and never touches the
# NVIDIA/CUDA code path.
if not current_platform.is_rocm() or rotary_segment_lengths is None:
return self.apply_rotary_emb(qk_reshaped, rotary_cos, rotary_sin)
# ROCm only: the HIP flash-attn Triton rotary kernel fails with
# hipErrorInvalidValue once grid.y = cdiv(seqlen, BLOCK_M) exceeds
# _HIP_MAX_GRID_DIM_Y (65536). BLOCK_M is 8 for rotary_dim <= 128
# (MiniMax-M3 vision: rotary_dim=78), giving a hard limit of
# 65536 * BLOCK_M tokens — measured exactly as 524288 OK / 524289 fail.
# Only long videos cross it; since vision_segment_max_frames caps each
# segment at a few frames (<< limit), applying RoPE per segment keeps
# every sub-call in range. Splitting on segment boundaries is
# mathematically exact because rotary_cos/sin are precomputed per token.
# Images and short clips stay on the single-kernel fast path above.
rotary_dim = rotary_cos.shape[-1] * 2
block_m = 8 if rotary_dim <= 128 else 4
hip_rotary_max_seqlen = _HIP_MAX_GRID_DIM_Y * block_m
if seq_len <= hip_rotary_max_seqlen or len(rotary_segment_lengths) <= 1:
return self.apply_rotary_emb(qk_reshaped, rotary_cos, rotary_sin)
qk_segments = qk_reshaped.split(rotary_segment_lengths, dim=1)
cos_segments = rotary_cos.split(rotary_segment_lengths, dim=0)
sin_segments = rotary_sin.split(rotary_segment_lengths, dim=0)
return torch.cat(
[
self.apply_rotary_emb(qk_s, cos_s, sin_s)
for qk_s, cos_s, sin_s in zip(qk_segments, cos_segments, sin_segments)
],
dim=1,
)
def forward(
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
rotary_cos: torch.Tensor,
rotary_sin: torch.Tensor,
max_seqlen: torch.Tensor,
rotary_segment_lengths: list[int] | None = None,
sequence_lengths: torch.Tensor | None = None,
) -> torch.Tensor:
# x: (N, 1, embed_dim) [seq=N, batch=1, chan=embed_dim]
x_qkv, _ = self.qkv_proj(x) # (N, 1, 3 * heads_per_part * head_dim)
seq_len, batch_size, _ = x_qkv.shape
# Rearrange to (b=1, N, 3, heads, head_dim) — same as Qwen2_5_VisionAttention
qkv = rearrange(
x_qkv,
"s b (three head d) -> b s three head d",
three=3,
head=self.num_heads_per_partition,
)
qk, v = qkv[:, :, :2], qkv[:, :, 2] # (b,N,2,h,d) and (b,N,h,d)
# Stack q/k → (2*b, N, heads, head_dim) for joint RoPE application.
# rotary_cos/sin: (N, half_rot_dim) — ApplyRotaryEmb expands internally
# and rotates only the first 2*half_rot_dim dims, passing the rest through.
qk_reshaped = rearrange(qk, "b s two h d -> (two b) s h d", two=2).contiguous()
qk_rotated = self._apply_rotary_emb(
qk_reshaped, rotary_cos, rotary_sin, seq_len, rotary_segment_lengths
)
qk_rotated = qk_rotated.view(
2, batch_size, seq_len, self.num_heads_per_partition, self.head_dim
)
q, k = qk_rotated.unbind(dim=0) # each (b=1, N, heads, head_dim)
# Flash attention → (b, N, heads, head_dim)
context = self.attn(
query=q,
key=k,
value=v,
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
sequence_lengths=sequence_lengths,
)
# Back to (N, 1, embed_dim)
context = rearrange(context, "b s h d -> s b (h d)", b=batch_size)
output, _ = self.out_proj(context)
return output
class MiniMaxVLEncoderLayer(nn.Module):
"""Single CLIP-style transformer block."""
def __init__(
self,
config: PretrainedConfig,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
embed_dim = config.hidden_size
self.layer_norm1 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
self.self_attn = MiniMaxVLAttention(
embed_dim=embed_dim,
num_heads=config.num_attention_heads,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
)
self.layer_norm2 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
use_data_parallel = is_vit_use_data_parallel()
self.fc1 = ColumnParallelLinear(
config.hidden_size,
config.intermediate_size,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.fc1",
disable_tp=use_data_parallel,
)
self.act = get_act_fn(getattr(config, "hidden_act", "gelu"))
self.fc2 = RowParallelLinear(
config.intermediate_size,
config.hidden_size,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.fc2",
disable_tp=use_data_parallel,
)
def forward(
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
rotary_cos: torch.Tensor,
rotary_sin: torch.Tensor,
max_seqlen: torch.Tensor,
rotary_segment_lengths: list[int] | None = None,
sequence_lengths: torch.Tensor | None = None,
) -> torch.Tensor:
# x: (N, 1, hidden_size)
x = x + self.self_attn(
self.layer_norm1(x),
cu_seqlens,
rotary_cos,
rotary_sin,
max_seqlen,
rotary_segment_lengths,
sequence_lengths,
)
residual = x
x, _ = self.fc1(self.layer_norm2(x))
x = self.act(x)
x, _ = self.fc2(x)
return residual + x
class MiniMaxVLEncoder(nn.Module):
def __init__(
self,
config: PretrainedConfig,
num_hidden_layers_override: int | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
n = (
config.num_hidden_layers
if num_hidden_layers_override is None
else num_hidden_layers_override
)
self.layers = nn.ModuleList(
[
MiniMaxVLEncoderLayer(
config=config,
quant_config=quant_config,
prefix=f"{prefix}.layers.{i}",
)
for i in range(n)
]
)
def forward(
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
rotary_cos: torch.Tensor,
rotary_sin: torch.Tensor,
max_seqlen: torch.Tensor,
rotary_segment_lengths: list[int] | None = None,
sequence_lengths: torch.Tensor | None = None,
) -> torch.Tensor:
for layer in self.layers:
x = layer(
x,
cu_seqlens,
rotary_cos,
rotary_sin,
max_seqlen,
rotary_segment_lengths,
sequence_lengths,
)
return x
class MiniMaxVLVisionTransformer(nn.Module):
"""CLIP-based ViT with 3D RoPE (t/h/w decomposed).
Faithfully mirrors the reference ``MiniMaxVLVisionTransformer``.
FLASHINFER backend is not supported; standard flash-attn is used.
"""
def __init__(
self,
config: PretrainedConfig,
num_hidden_layers_override: int | None = None,
require_post_norm: bool | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
compression = config.img_token_compression_config
self.spatial_merge_size: int = compression.get("spatial_merge_size", 2)
self.temporal_patch_size: int = compression.get("temporal_patch_size", 2)
self.vision_segment_max_frames: int | None = getattr(
config, "vision_segment_max_frames", None
)
self.use_data_parallel = is_vit_use_data_parallel()
embed_dim = config.hidden_size
head_dim = embed_dim // config.num_attention_heads
# Backend selection + sharding info for building encoder metadata.
# Defaults to FLASH_ATTN on SM80+; --mm-encoder-attn-backend FLASHINFER
# selects the cuDNN ViT prefill path.
self.hidden_size = embed_dim
self.tp_size = (
1
if self.use_data_parallel
else parallel_state.get_tensor_model_parallel_world_size()
)
self.attn_backend = get_vit_attn_backend(
head_size=head_dim, dtype=torch.get_default_dtype()
)
rope_dims = 2 * (head_dim // 2)
# Split rope dims evenly across t/h/w (same formula as the reference)
self.t_dim = int(2 * ((rope_dims // 3) // 2))
self.h_dim = int(2 * ((rope_dims // 3) // 2))
self.w_dim = int(2 * ((rope_dims // 3) // 2))
# rot_dim = t_dim + h_dim + w_dim (may be < head_dim)
rope_theta: float = getattr(config, "rope_theta", 10000.0)
inv_freq_t = 1.0 / (
rope_theta
** (torch.arange(0, self.t_dim, 2, dtype=torch.float32) / self.t_dim)
)
inv_freq_h = 1.0 / (
rope_theta
** (torch.arange(0, self.h_dim, 2, dtype=torch.float32) / self.h_dim)
)
inv_freq_w = 1.0 / (
rope_theta
** (torch.arange(0, self.w_dim, 2, dtype=torch.float32) / self.w_dim)
)
self.register_buffer("inv_freq_t", inv_freq_t, persistent=False)
self.register_buffer("inv_freq_h", inv_freq_h, persistent=False)
self.register_buffer("inv_freq_w", inv_freq_w, persistent=False)
self.embeddings = MiniMaxVLPatchEmbed(config)
self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
n_layers = config.num_hidden_layers
if num_hidden_layers_override is None:
num_hidden_layers_override = n_layers
self.encoder = MiniMaxVLEncoder(
config=config,
num_hidden_layers_override=num_hidden_layers_override,
quant_config=quant_config,
prefix=f"{prefix}.encoder",
)
if require_post_norm is None:
require_post_norm = num_hidden_layers_override == n_layers
self.post_layernorm = (
nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
if require_post_norm
else None
)
# out_hidden_size needed by run_dp_sharded_mrope_vision_model
self.out_hidden_size = embed_dim
# ── RoPE helpers ─────────────────────────────────────────────────────
def _get_3d_rope_embed(
self, grid_t: int, grid_h: int, grid_w: int, spatial_merge_size: int
) -> torch.Tensor:
"""Compute 3D RoPE frequencies for a single (T, H, W) grid.
Returns (T*H*W, half_rot_dim) on the same device as inv_freq buffers.
Mirrors the reference ``_get_3d_rope_embed`` exactly.
"""
tokens_per_frame = grid_h * grid_w
tpos_ids = (
torch.arange(grid_t, device=self.inv_freq_t.device)
.unsqueeze(1)
.expand(-1, tokens_per_frame)
.flatten()
)
hpos_ids = (
torch.arange(grid_h, device=self.inv_freq_h.device)
.unsqueeze(1)
.expand(-1, grid_w)
.reshape(
grid_h // spatial_merge_size,
spatial_merge_size,
grid_w // spatial_merge_size,
spatial_merge_size,
)
.permute(0, 2, 1, 3)
.unsqueeze(0)
.expand(grid_t, -1, -1, -1, -1)
.flatten()
)
wpos_ids = (
torch.arange(grid_w, device=self.inv_freq_w.device)
.unsqueeze(0)
.expand(grid_h, -1)
.reshape(
grid_h // spatial_merge_size,
spatial_merge_size,
grid_w // spatial_merge_size,
spatial_merge_size,
)
.permute(0, 2, 1, 3)
.unsqueeze(0)
.expand(grid_t, -1, -1, -1, -1)
.flatten()
)
max_t = max(grid_t, 1)
max_hw = max(grid_h, grid_w)
seq_t = torch.arange(
max_t, device=self.inv_freq_t.device, dtype=self.inv_freq_t.dtype
)
seq_hw = torch.arange(
max_hw, device=self.inv_freq_h.device, dtype=self.inv_freq_h.dtype
)
freqs_t = torch.outer(seq_t, self.inv_freq_t) # (max_t, t_dim/2)
freqs_h = torch.outer(seq_hw, self.inv_freq_h) # (max_hw, h_dim/2)
freqs_w = torch.outer(seq_hw, self.inv_freq_w) # (max_hw, w_dim/2)
return torch.cat(
[freqs_t[tpos_ids], freqs_h[hpos_ids], freqs_w[wpos_ids]], dim=-1
) # (T*H*W, half_rot_dim)
def _get_rope_embed_3d(
self, grid_thw: list[list[int]], spatial_merge_size: int
) -> torch.Tensor:
embeds = [
self._get_3d_rope_embed(t, h, w, spatial_merge_size) for t, h, w in grid_thw
]
return torch.cat(embeds, dim=0) # (total_N, half_rot_dim)
# ── Frame-limit helper (mirrors the reference) ───────────────────────
def _apply_max_frames_limit(self, grid_thw: list[list[int]]) -> list[list[int]]:
if self.vision_segment_max_frames is None:
return grid_thw
max_f = self.vision_segment_max_frames
out: list[list[int]] = []
for t, h, w in grid_thw:
if t <= max_f:
out.append([t, h, w])
else:
for i in range(0, t, max_f):
out.append([min(max_f, t - i), h, w])
return out
# ── Forward ──────────────────────────────────────────────────────────
def forward(
self,
pixel_values: torch.Tensor,
grid_thw: list[list[int]],
) -> torch.Tensor:
# pixel_values: (total_N, C * temporal_patch_size * patch_size²)
# Output: (total_N, hidden_size)
hidden = self.embeddings(pixel_values) # (total_N, hidden_size)
hidden = self.pre_layrnorm(hidden)
limited = self._apply_max_frames_limit(grid_thw)
# Token-level cumulative sequence lengths (one segment per limited grid).
lens = [t * h * w for t, h, w in limited]
cu_seqlens_np = np.zeros(len(lens) + 1, dtype=np.int32)
np.cumsum(np.array(lens, dtype=np.int32), out=cu_seqlens_np[1:])
# Backend-specific encoder metadata. For FLASH_ATTN this returns the raw
# token cu_seqlens, the max segment length, and sequence_lengths=None;
# for FLASHINFER (cuDNN) it repacks cu_seqlens into element-offset
# indptrs, buckets max_seqlen, and builds padded per-sequence lengths.
sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
self.attn_backend, cu_seqlens_np, hidden.device
)
max_seqlen = torch.tensor(
MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens_np),
dtype=torch.int32,
)
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
self.attn_backend,
cu_seqlens_np,
self.hidden_size,
self.tp_size,
hidden.device,
)
# 3D RoPE: (total_N, half_rot_dim); ApplyRotaryEmb expands internally
freqs = self._get_rope_embed_3d(limited, self.spatial_merge_size)
freqs = freqs.to(device=hidden.device)
# Keep cos/sin in fp32; ApplyRotaryEmb(enable_fp32_compute=True) runs the
# rotation in fp32 to match the reference precision.
rotary_cos, rotary_sin = freqs.cos(), freqs.sin()
# Encoder expects (N, 1, hidden_size) — add batch dim
hidden = hidden.unsqueeze(1)
# On ROCm, the flash_attn Triton rotary kernel can fail with
# hipErrorInvalidValue when seqlen is very large, e.g. 192k video
# tokens; pass per-segment lengths so RoPE can be applied in chunks.
# On other platforms leave it None -> single-kernel fast path, so the
# NVIDIA/CUDA code path is unchanged.
rotary_segment_lengths = lens if current_platform.is_rocm() else None
hidden = self.encoder(
hidden,
cu_seqlens,
rotary_cos,
rotary_sin,
max_seqlen,
rotary_segment_lengths,
sequence_lengths=sequence_lengths,
)
hidden = hidden.squeeze(1) # back to (total_N, hidden_size)
if self.post_layernorm is not None:
hidden = self.post_layernorm(hidden)
return hidden
class MiniMaxVLMultiModalProjector(nn.Module):
"""Two-layer MLP projector: vision_hidden → text_hidden."""
def __init__(
self,
vision_hidden_size: int,
text_hidden_size: int,
projector_hidden_size: int | None,
multimodal_projector_bias: bool,
projector_hidden_act: str = "gelu",
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
mid = projector_hidden_size if projector_hidden_size else text_hidden_size
use_dp = is_vit_use_data_parallel()
self.linear_1 = ColumnParallelLinear(
vision_hidden_size,
mid,
bias=multimodal_projector_bias,
quant_config=quant_config,
prefix=f"{prefix}.linear_1",
disable_tp=use_dp,
)
self.act = get_act_fn(projector_hidden_act)
self.linear_2 = RowParallelLinear(
mid,
text_hidden_size,
bias=multimodal_projector_bias,
quant_config=quant_config,
prefix=f"{prefix}.linear_2",
disable_tp=use_dp,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x, _ = self.linear_1(x)
x = self.act(x)
x, _ = self.linear_2(x)
return x
class MiniMaxVLPatchMerger(nn.Module):
def __init__(
self,
spatial_merge_size: int,
text_hidden_size: int,
projector_hidden_size: int | None,
patch_merge_bias: bool,
projector_hidden_act: str = "gelu",
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
self.spatial_merge_size = spatial_merge_size
mid = projector_hidden_size if projector_hidden_size else text_hidden_size
merge_in = text_hidden_size * spatial_merge_size**2
use_dp = is_vit_use_data_parallel()
self.linear_1 = ColumnParallelLinear(
merge_in,
mid,
bias=patch_merge_bias,
quant_config=quant_config,
prefix=f"{prefix}.linear_1",
disable_tp=use_dp,
)
self.act = get_act_fn(projector_hidden_act)
self.linear_2 = RowParallelLinear(
mid,
text_hidden_size,
bias=patch_merge_bias,
quant_config=quant_config,
prefix=f"{prefix}.linear_2",
disable_tp=use_dp,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (N, text_hidden_size) → (N // merge_size², text_hidden_size)
x = x.reshape(x.shape[0] // (self.spatial_merge_size**2), -1)
x, _ = self.linear_1(x)
x = self.act(x)
x, _ = self.linear_2(x)
return x
class MiniMaxVLVisionModel(nn.Module):
"""Full vision model: ViT → projector → patch merger."""
def __init__(
self,
config: PretrainedConfig,
text_hidden_size: int,
projector_hidden_size: int | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
compression = config.img_token_compression_config
spatial_merge_size: int = compression.get("spatial_merge_size", 2)
self.spatial_merge_size = spatial_merge_size
self.use_data_parallel = is_vit_use_data_parallel()
# The released checkpoint ships no ``post_layernorm`` weights and
# uses ``vision_feature_layer=-1`` with ``vision_feature_select_strategy
# ="full"``, i.e. the raw last encoder hidden state (CLIP's
# ``last_hidden_state`` is taken before the post layernorm). Applying an
# untrained post layernorm here would corrupt the visual features.
self.vision_model = MiniMaxVLVisionTransformer(
config=config,
require_post_norm=False,
quant_config=quant_config,
prefix=maybe_prefix(prefix, "vision_model"),
)
self.multi_modal_projector = MiniMaxVLMultiModalProjector(
vision_hidden_size=config.hidden_size,
text_hidden_size=text_hidden_size,
projector_hidden_size=projector_hidden_size,
multimodal_projector_bias=getattr(
config, "multimodal_projector_bias", True
),
projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
quant_config=quant_config,
prefix=maybe_prefix(prefix, "multi_modal_projector"),
)
self.patch_merge_mlp = MiniMaxVLPatchMerger(
spatial_merge_size=spatial_merge_size,
text_hidden_size=text_hidden_size,
projector_hidden_size=projector_hidden_size,
patch_merge_bias=getattr(config, "patch_merge_bias", True),
projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
quant_config=quant_config,
prefix=maybe_prefix(prefix, "patch_merge_mlp"),
)
self.dtype = self.vision_model.embeddings.patch_embedding.weight.dtype
self.out_hidden_size = text_hidden_size
def forward(
self,
pixel_values: torch.Tensor,
grid_thw: list[list[int]],
) -> torch.Tensor:
hidden = self.vision_model(pixel_values=pixel_values, grid_thw=grid_thw)
if hidden.dim() == 3:
hidden = hidden.squeeze(0)
hidden = self.multi_modal_projector(hidden)
hidden = self.patch_merge_mlp(hidden)
return hidden
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("qkv_proj.", "q_proj.", "q"),
("qkv_proj.", "k_proj.", "k"),
("qkv_proj.", "v_proj.", "v"),
]
params_dict = dict(self.named_parameters(remove_duplicate=False))
loaded_params: set[str] = set()
for name, loaded_weight in weights:
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params
vllm/models/minimax_m3/nvidia/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/minimax_m3/nvidia/model.py
+1177
View File
File diff suppressed because it is too large Load Diff
vllm/models/minimax_m3/nvidia/mtp.py
+312
View File
@@ -0,0 +1,312 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
import regex as re
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.model_executor.layers.fused_moe import (
fused_moe_make_expert_params_mapping,
)
from vllm.model_executor.layers.linear import (
ReplicatedLinear,
)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader,
maybe_remap_kv_scale_name,
)
from vllm.model_executor.models.utils import (
maybe_prefix,
)
from vllm.sequence import IntermediateTensors
from .model import (
MiniMAXGemmaRMSNorm,
MiniMaxM3DecoderLayer,
)
class MiniMaxM3MultiTokenPredictorLayer(nn.Module):
def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
super().__init__()
assert vllm_config.speculative_config is not None
config = vllm_config.speculative_config.draft_model_config.hf_config
quant_config = vllm_config.quant_config
self.enorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.hnorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.eh_proj = ReplicatedLinear(
config.hidden_size * 2,
config.hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.eh_proj",
)
self.transformer_layer = MiniMaxM3DecoderLayer(
vllm_config=vllm_config,
prefix=prefix,
force_sparse_attn=True,
force_moe=True,
is_mtp_block=True,
)
self.final_layernorm = MiniMAXGemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
previous_hidden_states: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
spec_step_index: int = 0,
) -> torch.Tensor:
assert inputs_embeds is not None
# Mask out inputs at position 0, as not needed by MTP.
inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
# Combine the normalized token embeddings with the normalized
# previous hidden states.
inputs_embeds = self.enorm(inputs_embeds)
previous_hidden_states = self.hnorm(previous_hidden_states)
hidden_states, _ = self.eh_proj(
torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
)
# Apply transformer layer.
hidden_states, residual = self.transformer_layer(
positions=positions,
hidden_states=hidden_states,
residual=None,
)
hidden_states += residual
return hidden_states
class MiniMaxM3MultiTokenPredictor(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
assert vllm_config.speculative_config is not None
# Use the draft (MTP) config, not the target model's. This is flat for a
# standalone checkpoint, and the promoted text_config for a bundled one.
config = vllm_config.speculative_config.draft_model_config.hf_config
self.num_mtp_layers = config.num_mtp_modules
self.layers = torch.nn.ModuleDict(
{
str(idx): MiniMaxM3MultiTokenPredictorLayer(
vllm_config, f"{prefix}.layers.{idx}"
)
for idx in range(self.num_mtp_layers)
}
)
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
prefix=maybe_prefix(prefix, "embed_tokens"),
)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
previous_hidden_states: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> torch.Tensor:
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
current_step_idx = spec_step_idx % self.num_mtp_layers
return self.layers[str(current_step_idx)](
input_ids,
positions,
previous_hidden_states,
inputs_embeds,
current_step_idx,
)
class MiniMaxM3MTP(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
assert vllm_config.speculative_config is not None
self.config = vllm_config.speculative_config.draft_model_config.hf_config
self.quant_config = vllm_config.quant_config
self.model = MiniMaxM3MultiTokenPredictor(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
self.lm_head = ParallelLMHead(
self.config.vocab_size,
self.config.hidden_size,
quant_config=self.quant_config,
prefix=maybe_prefix(prefix, "lm_head"),
)
self.logits_processor = LogitsProcessor(self.config.vocab_size)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
hidden_states: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
spec_step_idx: int = 0,
) -> torch.Tensor:
return self.model(
input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
)
def compute_logits(
self,
hidden_states: torch.Tensor,
spec_step_idx: int = 0,
) -> torch.Tensor | None:
current_step_idx = spec_step_idx % self.model.num_mtp_layers
mtp_layer = self.model.layers[str(current_step_idx)]
return self.logits_processor(
self.lm_head, mtp_layer.final_layernorm(hidden_states)
)
def _get_mtp_layer_idx_from_weight_name(self, name: str) -> int | None:
"""Return the MTP layer index in *.mtp.layers.{idx}.*, else None."""
match = re.search(r"\.mtp\.layers\.(\d+)\.", name)
return int(match.group(1)) if match else None
def _map_checkpoint_name(self, name: str) -> str | None:
"""Map a full checkpoint key to this MTP module's parameter name.
The MTP module only owns the *.mtp.layers.* weights plus the token
embedding and LM head, which the checkpoint shares with the main model.
Everything else belongs to other modules and is ignored here by returning
None.
"""
# In the bundled checkpoint, the MTP weights are prefixed with
# "language_model". The standalone MTP checkpoint has no such prefix.
# Strip it if present.
name = name.removeprefix("language_model.")
if name == "model.embed_tokens.weight":
return "model.embed_tokens.weight"
if name == "lm_head.weight":
return "lm_head.weight"
if "model.mtp.layers" in name:
if "weight_scale_inv" in name:
# The checkpoint stores block scales as "weight_scale_inv".
# The ModelOpt MXFP8 layers expose them as "weight_scale".
name = name.replace("weight_scale_inv", "weight_scale")
# Strip "mtp" from prefix.
return name.replace(".mtp.", ".")
return None
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Map q/k/v projections to qkv_proj, and gate/up projections to gate_up_proj.
stacked_params_mapping: list[tuple[str, str, int | str]] = [
# (param_name, shard_name, shard_id)
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".qkv_proj", ".index_q_proj", "index_q"),
(".qkv_proj", ".index_k_proj", "index_k"),
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
]
# Map expert weights w1/w2/w3 to gate/down/up.
# (param_name, weight_name, expert_id, shard_id)
expert_params_mapping = fused_moe_make_expert_params_mapping(
self,
ckpt_gate_proj_name="w1",
ckpt_down_proj_name="w2",
ckpt_up_proj_name="w3",
num_experts=self.config.num_local_experts,
)
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
loaded_mtp_layers: set[int] = set()
for name, loaded_weight in weights:
mtp_layer = self._get_mtp_layer_idx_from_weight_name(name)
mapped_name = self._map_checkpoint_name(name)
if mapped_name is None:
# This weight does not belong to the MTP module, so skip it.
continue
name = mapped_name
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
# Routed experts (w1/w2/w3) are handled below. Don't let the
# stacked mapping rewrite them.
if ("block_sparse_moe.experts." in name) and name not in params_dict:
continue
name = name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
for (
param_name,
weight_name,
expert_id,
expert_shard_id,
) in expert_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(
param,
loaded_weight,
name,
shard_id=expert_shard_id,
expert_id=expert_id,
)
break
else:
remapped_name = maybe_remap_kv_scale_name(name, params_dict)
if remapped_name is None or remapped_name not in params_dict:
continue
name = remapped_name
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
if mtp_layer is not None:
loaded_mtp_layers.add(mtp_layer)
# Validate that weights were loaded for each MTP layer.
for layer_idx in range(self.model.num_mtp_layers):
if layer_idx not in loaded_mtp_layers:
raise ValueError(
f"Failed to load MTP layer {layer_idx} weights from checkpoint."
)
return loaded_params
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+110
View File
@@ -0,0 +1,110 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MSA (SM100/Blackwell) block-sparse attend for MiniMax M3.
Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``);
decode falls back to the Triton split-K kernel (no MSA decode yet). ``fmha_sm100``
imports are function-local, so this module is import-safe on AMD/non-SM100.
"""
import torch
from vllm.forward_context import get_forward_context
from vllm.models.minimax_m3.common.ops.sparse_attn import (
SPARSE_BLOCK_SIZE,
minimax_m3_sparse_attn_decode,
)
from vllm.models.minimax_m3.common.sparse_attention import (
MiniMaxM3SparseImpl,
MiniMaxM3SparseMetadata,
)
from vllm.v1.attention.backend import AttentionLayer
class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
"""MSA block-sparse attend (``fmha_sm100``); Triton split-K decode."""
def forward(
self,
layer: AttentionLayer,
query: torch.Tensor,
kv_cache: torch.Tensor,
topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
output: torch.Tensor,
) -> torch.Tensor:
attn_metadata = get_forward_context().attn_metadata
if not isinstance(attn_metadata, dict):
return output # profiling run; caches unbound
main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined]
assert isinstance(main_md, MiniMaxM3SparseMetadata)
decode_topk, prefill_topk = topk_idx
nd = main_md.num_decode_tokens
num_tokens = main_md.num_actual_tokens
hd = self.head_size
q = query[:num_tokens].view(-1, self.num_heads, hd)
out = output[:num_tokens].view(-1, self.num_heads, hd)
kv_cache = (
kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
)
# Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
if main_md.num_decodes > 0:
d = main_md.decode
assert d is not None and decode_topk is not None
minimax_m3_sparse_attn_decode(
q[:nd],
kv_cache,
decode_topk,
d.block_table,
d.seq_lens,
self.num_kv_heads,
self.scale,
out[:nd],
d.decode_query_len,
)
# Prefill [nd:]: MSA sparse FMHA over the selected blocks.
if main_md.num_prefills > 0:
from vllm.third_party.fmha_sm100.sparse import (
build_k2q_csr,
sparse_atten_func,
)
p = main_md.prefill
assert p is not None and prefill_topk is not None
qp = q[nd:]
k_cache = kv_cache[:, 0].transpose(1, 2)
v_cache = kv_cache[:, 1].transpose(1, 2)
k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr(
prefill_topk,
p.cu_seqlens_q,
p.cu_seqlens_k,
SPARSE_BLOCK_SIZE,
total_k=0,
max_seqlen_k=p.max_seq_len,
max_seqlen_q=p.max_query_len,
total_rows=p.total_kv_blocks,
qhead_per_kv=qp.shape[1] // self.num_kv_heads,
return_schedule=True,
)
sparse_atten_func(
qp,
k_cache,
v_cache,
k2q_row_ptr,
k2q_q_indices,
topK=self.topk_blocks,
blk_kv=SPARSE_BLOCK_SIZE,
causal=True,
softmax_scale=self.scale,
cu_seqlens_q=p.cu_seqlens_q,
cu_seqlens_k=p.cu_seqlens_k,
max_seqlen_q=p.max_query_len,
max_seqlen_k=p.max_seq_len,
page_table=p.block_table,
seqused_k=p.seq_lens,
schedule=schedule,
out=out[nd:],
)
return output
vllm/reasoning/__init__.py
+4
View File
@@ -92,6 +92,10 @@ _REASONING_PARSERS_TO_REGISTER = {
"minimax_m2_reasoning_parser",
"MiniMaxM2AppendThinkReasoningParser",
),
"minimax_m3": (
"minimax_m3_reasoning_parser",
"MiniMaxM3ReasoningParser",
),
"mistral": (
"mistral_reasoning_parser",
"MistralReasoningParser",
vllm/reasoning/minimax_m3_reasoning_parser.py
+171
View File
@@ -0,0 +1,171 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable, Sequence
from typing import TYPE_CHECKING
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser):
"""Reasoning parser for MiniMax M3 explicit thinking blocks.
MiniMax M3 emits reasoning as:
<mm:think>reasoning text</mm:think>assistant content
The M3 tokenizer exposes both markers as complete vocabulary tokens. The
chat template may also prefill the start marker when
``thinking_mode="enabled"``, so generated text can begin directly inside a
reasoning block without emitting ``<mm:think>`` again.
"""
@property
def start_token(self) -> str:
return "<mm:think>"
@property
def end_token(self) -> str:
return "</mm:think>"
def __init__(self, tokenizer, *args, **kwargs):
super().__init__(tokenizer, *args, **kwargs)
chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled"
self._at_response_start = True
def extract_reasoning(
self,
model_output: str,
request: "ChatCompletionRequest | ResponsesRequest",
) -> tuple[str | None, str | None]:
# MiniMax M3 can start a response with a stray closer. Drop that first
# token only; later unmatched closers stay visible as content.
if not self._initial_in_reasoning and model_output.startswith(self.end_token):
content = model_output[len(self.end_token) :]
return None, content or None
if self._initial_in_reasoning and self.start_token not in model_output:
reasoning, end, content = model_output.partition(self.end_token)
if not end:
return model_output, None
return reasoning, content or None
if self.start_token not in model_output:
return None, model_output
content_before, _, after_start = model_output.partition(self.start_token)
reasoning, end, content_after = after_start.partition(self.end_token)
if not end:
return reasoning, content_before or None
return reasoning, (content_before + content_after) or None
def is_reasoning_end_streaming(
self, input_ids: Sequence[int], delta_ids: Iterable[int]
) -> bool:
delta_ids = tuple(delta_ids)
if self.end_token_id in delta_ids:
return True
if self.end_token_id in input_ids:
return True
if self._initial_in_reasoning:
return False
if self.start_token_id not in input_ids:
return bool(input_ids)
return False
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
if self.end_token_id in input_ids:
end_index = len(input_ids) - 1 - input_ids[::-1].index(self.end_token_id)
return input_ids[end_index + 1 :]
if self._initial_in_reasoning and self.start_token_id not in input_ids:
return []
if self.start_token_id not in input_ids:
return input_ids
return []
def extract_reasoning_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
) -> DeltaMessage | None:
if not delta_text:
return None
if self._at_response_start and not self._initial_in_reasoning:
# Apply the leading-closer tolerance once. Later unmatched closers
# stay visible as content.
self._at_response_start = False
if delta_text.startswith(self.end_token):
delta_text = delta_text[len(self.end_token) :]
if not delta_text:
return None
if delta_token_ids and delta_token_ids[0] == self.end_token_id:
delta_token_ids = delta_token_ids[1:]
if self.end_token_id in previous_token_ids:
return DeltaMessage(content=delta_text)
if (
self._initial_in_reasoning
and self.start_token_id not in previous_token_ids
and self.start_token_id not in delta_token_ids
):
if self.end_token_id in delta_token_ids:
reasoning, _, content = delta_text.partition(self.end_token)
return DeltaMessage(
reasoning=reasoning or None,
content=content or None,
)
return DeltaMessage(reasoning=delta_text)
if (
self.start_token_id not in previous_token_ids
and self.start_token_id not in delta_token_ids
):
return DeltaMessage(content=delta_text)
if self.end_token_id in delta_token_ids:
reasoning_text, _, content = delta_text.partition(self.end_token)
if self.start_token_id in delta_token_ids:
_, _, reasoning_text = reasoning_text.partition(self.start_token)
return DeltaMessage(
reasoning=reasoning_text or None,
content=content or None,
)
if self.start_token_id in delta_token_ids:
_, _, reasoning = delta_text.partition(self.start_token)
return DeltaMessage(reasoning=reasoning) if reasoning else None
return DeltaMessage(reasoning=delta_text)
def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
if not self._initial_in_reasoning:
return super().count_reasoning_tokens(token_ids)
count = 0
depth = 1
for token_id in token_ids:
if token_id == self.start_token_id:
depth += 1
continue
if token_id == self.end_token_id:
if depth > 0:
depth -= 1
continue
if depth > 0:
count += 1
return count
vllm/tool_parsers/__init__.py
+4
View File
@@ -126,6 +126,10 @@ _TOOL_PARSERS_TO_REGISTER = {
"minimax_m2_tool_parser",
"MinimaxM2ToolParser",
),
"minimax_m3": (
"minimax_m3_tool_parser",
"MinimaxM3ToolParser",
),
"minimax": (
"minimax_tool_parser",
"MinimaxToolParser",
vllm/tool_parsers/minimax_m3_tool_parser.py
+19
View File
@@ -0,0 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.tool_parsers.rust_tool_parser import RustToolParser
class MinimaxM3ToolParser(RustToolParser):
"""Adapter from the Rust MiniMax M3 parser to vLLM ToolParser.
The real M3 grammar lives in the Rust tool-parser crate. This class only
configures the generic Rust bridge with the MiniMax M3 parser name.
M3 is not M2 with renamed tags: it prefixes each structural tag with the
MiniMax namespace marker, allows multiple ``<invoke>`` tags in one wrapper,
and represents nested arguments with parameter-name XML tags.
"""
rust_parser_name = "MinimaxM3ToolParser"
tool_call_start_token = "]<]minimax[>[<tool_call>"
vllm/transformers_utils/config.py
+2
View File
@@ -103,6 +103,8 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
medusa="MedusaConfig",
mellum="MellumConfig",
midashenglm="MiDashengLMConfig",
minimax_m3_vl="MiniMaxM3Config",
minimax_m3_mtp="MiniMaxM3MTPConfig",
moondream3="Moondream3Config",
eagle="EAGLEConfig",
speculators="SpeculatorsConfig",

Some files were not shown because too many files have changed in this diff Show More