diff --git a/.gitignore b/.gitignore index 8dde75e43e4b..c70200ed0914 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 8405958a4198..1259ec0c1bf7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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) diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake new file mode 100644 index 000000000000..15610552f236 --- /dev/null +++ b/cmake/external_projects/fmha_sm100.cmake @@ -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") diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index cdab456348e2..e1dc01346055 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/libtorch_stable/activation_kernels.cu @@ -10,11 +10,20 @@ namespace vllm { -template __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 __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(g); - up = cast_to_packed(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(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(g); up = cast_to_packed(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(activated); } } // Activation and gating kernel template. template + 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( - 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(x, y, limit); + out_ptr[idx] = compute( + 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 -__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 -__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(fval); } template -__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 -__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 -__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 __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::Type>, \ ACT_FIRST, true, HAS_CLAMP, true><<>>( \ out.mutable_data_ptr(), \ - input.const_data_ptr(), d, LIMIT); \ + input.const_data_ptr(), 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::Type>, \ ACT_FIRST, true, HAS_CLAMP, false><<>>( \ out.mutable_data_ptr(), \ - input.const_data_ptr(), d, LIMIT); \ + input.const_data_ptr(), d, LIMIT, ALPHA, BETA); \ }); \ } \ } else { \ @@ -252,7 +283,7 @@ packed_gelu_tanh_kernel(const packed_t& val) { PACKED_KERNEL::Type>, \ ACT_FIRST, false, HAS_CLAMP><<>>( \ out.mutable_data_ptr(), input.const_data_ptr(), \ - 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 { diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/libtorch_stable/fp32_router_gemm.cu index 04397e0893c2..80374d66a025 100644 --- a/csrc/libtorch_stable/fp32_router_gemm.cu +++ b/csrc/libtorch_stable/fp32_router_gemm.cu @@ -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( \ - 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) - -INSTANTIATE_ALL(float) -INSTANTIATE_ALL(__nv_bfloat16) +#define INSTANTIATE(T, M, E, H) \ + template void invokeFp32RouterGemm(float*, T const*, \ + float const*, cudaStream_t); + +#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, 256, 3072) +INSTANTIATE_ALL(__nv_bfloat16, 256, 3072) +INSTANTIATE_ALL(float, 128, 6144) +INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) #undef INSTANTIATE_ALL #undef INSTANTIATE diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/libtorch_stable/fp32_router_gemm_entry.cu index 4baa740de930..b4bc0a11d200 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/libtorch_stable/fp32_router_gemm_entry.cu @@ -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 void invokeFp32RouterGemm(float* output, InputT const* mat_a, float const* mat_b, cudaStream_t stream); -// LoopUnroller templated on InputT -template +// LoopUnroller templated on InputT, kNumExperts and kHiddenDim +template 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( + invokeFp32RouterGemm( output, mat_a, mat_b, stream); } else { - Fp32LoopUnroller::unroll(num_tokens, output, - mat_a, mat_b, stream); + Fp32LoopUnroller::unroll(num_tokens, output, mat_a, mat_b, stream); } } }; -template -struct Fp32LoopUnroller { +template +struct Fp32LoopUnroller { static void unroll(int num_tokens, float* output, InputT const* mat_a, float const* mat_b, cudaStream_t stream) { if (num_tokens == kEnd) { - invokeFp32RouterGemm( + invokeFp32RouterGemm( output, mat_a, mat_b, stream); } else { throw std::invalid_argument( @@ -60,6 +66,23 @@ struct Fp32LoopUnroller { } }; +// Dispatch over the supported (num_experts, hidden_dim) pairs. +template +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::unroll( + num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 128 && hidden_dim == 6144) { + Fp32LoopUnroller::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(mat_a.data_ptr()); - Fp32LoopUnroller::unroll( - num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream); + dispatchFp32RouterGemm(num_experts, hidden_dim, num_tokens, out_ptr, + mat_a_ptr, mat_b_ptr, stream); } } diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu new file mode 100644 index 000000000000..5dd610f28788 --- /dev/null +++ b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -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 , sparse-profiling + * , sparse-serving ), 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 +#include +#include + +#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 +__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(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + int const dim = laneId * kElemsPerLane + i; + float const w = 1.0f + static_cast(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(cos_ptr[i_base + i]); + float const s = static_cast(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 +__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src, + float (&elems)[kElemsPerLane]) { + using Converter = vllm::_typeConvert; + uint2 v = *reinterpret_cast(src); + auto const* p = + reinterpret_cast(&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 +__device__ __forceinline__ void storeElems( + scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { + using Converter = vllm::_typeConvert; + uint2 v; + auto* p = reinterpret_cast(&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(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 , sparse-profiling +// and sparse-serving . 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 +__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 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) { + 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(tokenIdx) * qkv_row + slot * kHeadDim; + norm_w = q_norm_w; + } else if (isK) { + head = slot - k_begin; + row_ptr = + qkv + static_cast(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(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(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(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(tokenIdx) * nq * kHeadDim + + slot * kHeadDim; + } else if (isIQ && index_q_out != nullptr) { + store_ptr = index_q_out + + static_cast(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(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(elems, laneId, eps, norm_w, do_rope, rotary_dim, + cos_ptr, /*apply_norm=*/norm_w != nullptr); + storeElems(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(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(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 +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(num_tokens) * slots_per_token; + int const grid = + static_cast((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, \ + 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 \ + <<>>( \ + 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 index_q_norm_weight, // [128] + std::optional index_k_norm_weight, // [128] + int64_t num_index_heads, // niq; 0 => dense + std::optional slot_mapping, // [N] i64 + std::optional index_slot_mapping, // [N] i64 + std::optional kv_cache, // [nb,2,bs,nkv,128] + std::optional index_cache, // [nb,bs,128] + int64_t block_size, + std::optional q_out, // [N, nq*128] contiguous + std::optional + 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(qkv.size(0)); + int const nq = static_cast(num_heads); + int const nkv = static_cast(num_kv_heads); + int const niq = static_cast(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(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(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( + reinterpret_cast(qkv.data_ptr()), + q_out.has_value() ? reinterpret_cast(q_out->data_ptr()) + : nullptr, + index_q_out.has_value() + ? reinterpret_cast(index_q_out->data_ptr()) + : nullptr, + reinterpret_cast(q_norm_weight.data_ptr()), + reinterpret_cast(k_norm_weight.data_ptr()), + has_index + ? reinterpret_cast(index_q_norm_weight->data_ptr()) + : nullptr, + has_index + ? reinterpret_cast(index_k_norm_weight->data_ptr()) + : nullptr, + reinterpret_cast(cos_sin_cache.data_ptr()), + reinterpret_cast(positions.data_ptr()), + insert_kv + ? reinterpret_cast(slot_mapping->data_ptr()) + : nullptr, + insert_kv ? reinterpret_cast( + effective_index_slot_mapping->data_ptr()) + : nullptr, + insert_kv ? reinterpret_cast(kv_cache->data_ptr()) : nullptr, + (insert_kv && has_index) + ? reinterpret_cast(index_cache->data_ptr()) + : nullptr, + static_cast(eps), static_cast(rotary_dim), num_tokens, + nq, nkv, niq, static_cast(block_size), kv_s_block, kv_s_kv, + kv_s_token, kv_s_head, has_index, insert_kv, stream); + }); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 05e55e7198ca..d5144d76818a 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -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 index_q_norm_weight, + std::optional index_k_norm_weight, + int64_t num_index_heads, std::optional slot_mapping, + std::optional index_slot_mapping, + std::optional kv_cache, + std::optional index_cache, int64_t block_size, + std::optional q_out, + std::optional 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, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index b1c166b1d3aa..7d9a39a7a4b7 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -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_", diff --git a/csrc/ops.h b/csrc/ops.h index e39bae08f198..b909c5711d49 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -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); diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index a585cd77ffb6..fcf05cf68594 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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. diff --git a/pyproject.toml b/pyproject.toml index c782cc326bc1..031f8d1a0a28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/requirements/common.txt b/requirements/common.txt index ea53b8d25dd0..fde1ba4f0c9e 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -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 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index a3e1466c763a..c6d9ed24adbf 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -360,6 +360,7 @@ jsonpointer==3.0.0 # via jsonschema jsonschema==4.23.0 # via + # -c requirements/common.txt # hypothesis-jsonschema # mistral-common # ray diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 7488490ff000..879a32864442 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -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 diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 6d5435462ff0..820ce27bc3dd 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -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 diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 1148560787b3..e66db04c22ef 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -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()); } } diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index aa4d45964387..7de8a9d5fa12 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -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::(names::KIMI) .register_parser::(names::KIMI_K2) .register_parser::(names::MINIMAX_M2) + .register_parser::(names::MINIMAX_M3) .register_parser::(names::NEMOTRON_V3) .register_parser::(names::QWEN3) .register_parser::(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) diff --git a/rust/src/chat/src/parser/reasoning/tests.rs b/rust/src/chat/src/parser/reasoning/tests.rs index 803926d16bad..58d987770c65 100644 --- a/rust/src/chat/src/parser/reasoning/tests.rs +++ b/rust/src/chat/src/parser/reasoning/tests.rs @@ -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); diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 960d1d62af47..7561aa071ac8 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -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::(names::LLAMA3_JSON) .register_parser::(names::LLAMA4_JSON) .register_parser::(names::MINIMAX_M2) + .register_parser::(names::MINIMAX_M3) .register_parser::(names::MISTRAL) .register_parser::(names::PHI4_MINI_JSON) .register_parser::(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); diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index 5a2778157b9e..c40500adc74a 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -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) diff --git a/rust/src/reasoning-parser/src/lib.rs b/rust/src/reasoning-parser/src/lib.rs index f8f8d7c8726d..1f71e14cef7b 100644 --- a/rust/src/reasoning-parser/src/lib.rs +++ b/rust/src/reasoning-parser/src/lib.rs @@ -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; diff --git a/rust/src/reasoning-parser/src/minimax_m3.rs b/rust/src/reasoning-parser/src/minimax_m3.rs new file mode 100644 index 000000000000..69d4e416dfa8 --- /dev/null +++ b/rust/src/reasoning-parser/src/minimax_m3.rs @@ -0,0 +1,98 @@ +use vllm_tokenizer::DynTokenizer; + +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result}; + +const M3_THINK_START: &str = ""; +const M3_THINK_END: &str = ""; + +/// Reasoning parser for MiniMax M3 style outputs. +/// +/// MiniMax M3 uses `...` 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 `` emitted at the start of a response. + at_response_start: bool, + /// Holds an initial suffix like ` Result { + 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 `` 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> + 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 { + Ok(self.push_inner(delta)) + } + + fn finish(&mut self) -> Result { + 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) + } +} diff --git a/rust/src/reasoning-parser/src/tests.rs b/rust/src/reasoning-parser/src/tests.rs index 7e33e0cfc1b9..22c026d35818 100644 --- a/rust/src/reasoning-parser/src/tests.rs +++ b/rust/src/reasoning-parser/src/tests.rs @@ -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), + "" => Some(8), + "" => Some(9), "" => Some(10), "" => 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("reasonanswer").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("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("XXXYYY").unwrap(); + assert_eq!(delta.reasoning, None); + assert_eq!(delta.content.as_deref(), Some("XXXYYY")); +} + +#[test] +fn minimax_m3_drops_split_leading_end_marker() { + let tokenizer = Arc::new(FakeTokenizer); + let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap(); + + assert!(parser.push("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("reasonanswer").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")); +} diff --git a/rust/src/tool-parser/python/src/lib.rs b/rust/src/tool-parser/python/src/lib.rs index 81aed04b1cc4..e5ae0fa7b69a 100644 --- a/rust/src/tool-parser/python/src/lib.rs +++ b/rust/src/tool-parser/python/src/lib.rs @@ -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, diff --git a/rust/src/tool-parser/src/lib.rs b/rust/src/tool-parser/src/lib.rs index f611cbb7d1aa..b5f0b80d0452 100644 --- a/rust/src/tool-parser/src/lib.rs +++ b/rust/src/tool-parser/src/lib.rs @@ -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; diff --git a/rust/src/tool-parser/src/minimax_m3.rs b/rust/src/tool-parser/src/minimax_m3.rs new file mode 100644 index 000000000000..ac79ec1a577c --- /dev/null +++ b/rust/src/tool-parser/src/minimax_m3.rs @@ -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[>["; +const TOOL_CALL_END: &str = "]<]minimax[>["; +const INVOKE_START: &str = "]<]minimax[>[ = 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[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[42]<]minimax[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[Singapore]<]minimax[>[ +/// ]<]minimax[>[018956]<]minimax[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[book-001]<]minimax[>[ +/// ]<]minimax[>[2]<]minimax[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[ +/// ]<]minimax[>[ +/// ``` +/// +/// 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 `` 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> + 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 { + 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 { + 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 { + 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 { + 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 { + 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 { + 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 { + (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 { + 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> { + 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 { + 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 { + 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, 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 { + rest.value(MinimaxM3Event::IgnoredRest).parse_next(input) +} + +fn malformed() -> ModalResult { + 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::>() + .join("\n"); + format!("{TOOL_CALL_START}\n{invokes}\n{TOOL_CALL_END}") + } + + fn m3_test_tools() -> Vec { + 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::(&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::(&output.calls[0].arguments).unwrap(), + json!({ "city": "Seattle" }) + ); + assert_eq!( + serde_json::from_str::(&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::(&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::(&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::(&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::(&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::(&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::(&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::(&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::(&output.calls[0].arguments).unwrap(), + json!({ + "payload": { + "$text": "child text", + "$$text": "text", + "child": "value" + } + }) + ); + } +} diff --git a/setup.py b/setup.py index 8ef2d5eec32e..99bf8d91b50f 100644 --- a/setup.py +++ b/setup.py @@ -432,6 +432,19 @@ def run(self): 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 @@ def extract_precompiled_and_patch_package( ) # 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 @@ def extract_precompiled_and_patch_package( 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 @@ def _read_requirements(filename: str) -> list[str]: # 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 @@ def _read_requirements(filename: str) -> list[str]: "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", ] } diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py new file mode 100644 index 000000000000..32b4bc97ede7 --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3.py @@ -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]) diff --git a/tests/kernels/core/test_fused_allreduce_gemma_rms_norm.py b/tests/kernels/core/test_fused_allreduce_gemma_rms_norm.py new file mode 100644 index 000000000000..cb936ce33ad1 --- /dev/null +++ b/tests/kernels/core/test_fused_allreduce_gemma_rms_norm.py @@ -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, + ) diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index f855eb7aa171..0673a438c546 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -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 @@ 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" diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py new file mode 100644 index 000000000000..3268d125bb2c --- /dev/null +++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py @@ -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) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py new file mode 100644 index 000000000000..9a14edc42713 --- /dev/null +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -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 diff --git a/tests/models/multimodal/processing/test_minimax_m3.py b/tests/models/multimodal/processing/test_minimax_m3.py new file mode 100644 index 000000000000..04d6aa4778c3 --- /dev/null +++ b/tests/models/multimodal/processing/test_minimax_m3.py @@ -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", + ) diff --git a/tests/models/registry.py b/tests/models/registry.py index f5431e799e90..931dbfb61eb9 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -420,6 +420,11 @@ def check_available_online( "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 @@ def check_available_online( "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 @@ def check_available_online( 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", diff --git a/tests/reasoning/test_minimax_m3_reasoning_parser.py b/tests/reasoning/test_minimax_m3_reasoning_parser.py new file mode 100644 index 000000000000..e2cd14562c04 --- /dev/null +++ b/tests/reasoning/test_minimax_m3_reasoning_parser.py @@ -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 = ("", "") + + 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( + "plananswer", 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("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("XXXYYY", request) + + assert reasoning is None + assert content == "XXXYYY" + + +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("plananswer", 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("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, + ["", "plan", "", "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, + ["plananswer"], + ) + + 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, + ["", "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, + ["XXXYYY"], + ) + + assert reasoning is None + assert content == "XXXYYY" + 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", "", "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( + "abcdef", add_special_tokens=False + ) + open_reasoning_ids = tokenizer.encode("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("abcdef", 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") + ) diff --git a/tests/tool_parsers/test_minimax_m3_tool_parser.py b/tests/tool_parsers/test_minimax_m3_tool_parser.py new file mode 100644 index 000000000000..fd1acabde2e5 --- /dev/null +++ b/tests/tool_parsers/test_minimax_m3_tool_parser.py @@ -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}\n" + f'{NS}' + f"{NS}42{NS}" + f"{NS}true{NS}" + f"{NS}Please leave at front desk.{NS}" + f"{NS}" + f"{NS}Singapore{NS}" + f"{NS}018956{NS}" + f"{NS}" + f"{NS}" + f"{NS}{NS}book-001{NS}{NS}2{NS}{NS}" + f"{NS}{NS}pen-007{NS}{NS}5{NS}{NS}" + f"{NS}" + f"{NS}" + f"{NS}mobile{NS}" + f"{NS}may-launch{NS}" + f"{NS}" + f"{NS}" + f"{NS}a{NS}" + f"{NS}b{NS}" + f"{NS}" + f"{NS}\n" + f"{NS}" + ) + + +def build_order_invocation(user_id: int) -> str: + return ( + f'{NS}' + f"{NS}{user_id}{NS}" + f"{NS}" + ) + + +def build_multiple_order_call() -> str: + return ( + f"{NS}\n" + f"{build_order_invocation(1)}\n" + f"{build_order_invocation(2)}\n" + f"{NS}" + ) + + +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 diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 1f7150ce6a74..7bc7f1de4b34 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -30,6 +30,7 @@ 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 _is_v4(b: dict[str, Any]) -> bool: 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) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 38fcca66dc05..3878f3038bd2 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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, diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 52ce9f102a6c..48db183d5a3a 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -9,6 +9,8 @@ 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.""" diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index eba8653d63b6..de505e122cf6 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -45,6 +45,7 @@ "qwen3_next_mtp", "qwen3_5_mtp", "longcat_flash_mtp", + "minimax_m3_mtp", "mtp", "pangu_ultra_moe_mtp", "step3p5_mtp", @@ -528,6 +529,36 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: 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): diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ca2244a7324e..a3bfa56f579e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1101,20 +1101,26 @@ def __post_init__(self): ) 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." ) diff --git a/vllm/envs.py b/vllm/envs.py index 8b5544fd0aa4..a44ca3487462 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -113,6 +113,7 @@ 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 @@ def _resolve_rust_frontend_path() -> str | None: ), # 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") ), diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index f9d2d9970de2..919d71fb8e8b 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -93,6 +93,9 @@ 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 @@ def _filter_kernels_by_backend( 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: [ diff --git a/vllm/model_executor/kernels/linear/mxfp8/emulation.py b/vllm/model_executor/kernels/linear/mxfp8/emulation.py index a7cc29be758b..79b2fba3889b 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/emulation.py +++ b/vllm/model_executor/kernels/linear/mxfp8/emulation.py @@ -33,6 +33,17 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: 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 @@ def apply_weights( 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 @@ def apply_weights( 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) diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py index 336da511ad8b..8188fd596096 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -56,8 +56,6 @@ def apply_weights( 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 @@ def apply_weights( 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 @@ def apply_weights( backend="cutlass", ) - if M_padded != M_orig: - output = output[:M_orig, :] - if bias is not None: output = output + bias diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py new file mode 100644 index 000000000000..abc98df80abd --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -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 diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index ddad6801adc4..80bf251b2d88 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -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 @@ def forward_native(self, x: torch.Tensor) -> torch.Tensor: 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") diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 2e17a55ce7c4..5974e09624de 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -7,6 +7,7 @@ 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 @@ def unified_kv_cache_update_fake( ) +@eager_break_during_capture @maybe_transfer_kv_layer def unified_attention_with_output( query: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py new file mode 100644 index 000000000000..e49e135b26a9 --- /dev/null +++ b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py @@ -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) diff --git a/vllm/model_executor/layers/fused_moe/activation.py b/vllm/model_executor/layers/fused_moe/activation.py index b2e67e6220a9..2d8d46cacb7f 100644 --- a/vllm/model_executor/layers/fused_moe/activation.py +++ b/vllm/model_executor/layers/fused_moe/activation.py @@ -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 @@ def from_str(cls, s: str) -> "MoEActivation": 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 diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 1b063559b8d7..0755699d1a45 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -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 diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index df69fa328ca7..c74cb2d9a7b6 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -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 diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 3b354dd3ef13..5681d12554f3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -33,7 +33,10 @@ 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 @@ def _supports_quant_scheme( 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 @@ def workspace_shapes( 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 @@ def _act_mul_quant( 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 @@ def _act_mul_quant( ) 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 @@ def apply( 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 @@ def apply( 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: diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 1f5724ac39cc..21bda8e173fd 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -801,7 +801,11 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: 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: diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 03bf925fbd9b..a7f31afc5ef8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -787,6 +787,7 @@ def activation( activation: MoEActivation, output: torch.Tensor, input: torch.Tensor, + **kwargs, ) -> None: quant_config = self.quant_config or FUSED_MOE_UNQUANTIZED_CONFIG if activation == MoEActivation.SWIGLUOAI: diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 64c68018f365..867f71b9bf64 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -28,10 +28,7 @@ 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 @@ def __init__( 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 @@ def _supports_activation(activation: MoEActivation) -> bool: MoEActivation.GELU, MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.SWIGLUSTEP, MoEActivation.SILU_NO_MUL, MoEActivation.GELU_NO_MUL, @@ -787,6 +798,8 @@ def apply( 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 @@ def activation_with_lora( 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 @@ def activation_with_lora( "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 @@ def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None: 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 @@ def apply( 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, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py new file mode 100644 index 000000000000..71dd7634a697 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py @@ -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, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py new file mode 100644 index 000000000000..33851fdc8622 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -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) diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 25dd0584de0c..d81458b37514 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -64,6 +64,15 @@ def __init__( 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 @@ def _supports_activation(activation: MoEActivation) -> bool: MoEActivation.GELU, MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.SWIGLUSTEP, MoEActivation.SILU_NO_MUL, MoEActivation.GELU_NO_MUL, @@ -129,14 +139,34 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: 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, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 15806ca4f890..225484385865 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -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, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index d3176668016f..e80224be70fa 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -880,9 +880,18 @@ def adjust_N_for_activation(N: int, activation: MoEActivation) -> int: 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: diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 3a65e7360f06..acbf2cb46ad4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -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, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 64e6cb93fa80..d0d7c76481b0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -12,22 +12,43 @@ 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.") diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 9a75d6a3f1a0..669d1d376902 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -72,6 +72,8 @@ def __init__( 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 @@ def __init__( 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 diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 5867ce3e9a57..f230b4d57909 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -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 @@ def __init__( 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 diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index e980700d3ea6..bd4393be5e7c 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -11,7 +11,6 @@ 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 @@ def _setup_kernel( 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 @@ def process_weights_after_loading(self, layer: "RoutedExperts") -> None: ) 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, diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index cb2cd5e94a5b..b8c84ad2af25 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -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, diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index f7f9fe4c3dbe..9ee3a231b91f 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -1105,6 +1105,7 @@ def weight_loader_v2( 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 @@ def weight_loader( 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): diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index b0a245bb6037..f47dcae310a4 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -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) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 1d6264f77602..2bdb26e1a8dd 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -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 @@ 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 @@ def override_quantization_method( 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 @@ def create_weights( 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 @@ def _shuffle_weights_for_trtllm(self, layer: torch.nn.Module) -> None: 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 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: 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 @@ def get_fused_moe_quant_config( 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( diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 71442fb1add2..66a9aa86bde3 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -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 - - pid_pack = tl.program_id(0) - pid_m = tl.program_id(1) - m_offset = pid_m.to(tl.int64) * BLOCK_M - - if m_offset >= M: - return - - offs_m = tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, GROUP_SIZE) - row_mask = (m_offset + offs_m) < M - - base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m - base_out_offset = (m_offset + offs_m[:, None]) * output_q_stride_m - - packed_scale = tl.zeros((BLOCK_M,), dtype=tl.int32) - - for pack_idx in tl.static_range(4): - group_id = pid_pack * 4 + pack_idx - - if group_id < NUM_GROUPS: - n_offset = group_id * GROUP_SIZE - - 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) - - mul_ptrs = act_ptrs + N_2 - mul_in = tl.load(mul_ptrs, mask=row_mask[:, None], other=0.0) - - act_f32 = act_in.to(tl.float32) - mul_f32 = mul_in.to(tl.float32) - - if HAS_CLAMP: - act_f32 = tl.minimum(act_f32, clamp_limit) - mul_f32 = tl.clamp(mul_f32, -clamp_limit, clamp_limit) - - 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) - - absmax = tl.max(tl.abs(y), axis=1) - - scale_raw = tl.maximum(absmax / fp8_max, 1e-10) - exponent = tl.ceil(tl.log2(scale_raw)) - scale = tl.math.exp2(exponent) + GROUPS_PER_PACK: tl.constexpr = 4 + hidden_size: tl.constexpr = N // 2 + + 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 + + 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) + + col_mask = (col_start + col_offsets) < (GROUPS_PER_ROW * GROUP_SIZE) + + # 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 + + 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, + ) - y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max) + 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 HAS_CLAMP: + gate = tl.minimum(gate, clamp_limit) + up = tl.clamp(up, -clamp_limit, clamp_limit) + + # 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) + + 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) + + y_q = tl.clamp(y / scale[:, :, None], fp8_min, fp8_max) + + 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, :], + ) - 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], - ) + 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) - exponent_biased = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32) - packed_scale = packed_scale | (exponent_biased << (pack_idx * 8)) + 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_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 diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index a12918225348..e6063b463284 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -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) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 6c197ad3c59b..ecd31f2dc01a 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -164,6 +164,10 @@ "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 @@ "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 @@ "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"), diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c3725064a6d9..61d2376abb8c 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -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 ) diff --git a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py new file mode 100644 index 000000000000..18bf14249111 --- /dev/null +++ b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py @@ -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, + ) diff --git a/vllm/models/minimax_m3/__init__.py b/vllm/models/minimax_m3/__init__.py new file mode 100644 index 000000000000..f9ddb2a9d218 --- /dev/null +++ b/vllm/models/minimax_m3/__init__.py @@ -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", +] diff --git a/vllm/models/minimax_m3/amd/__init__.py b/vllm/models/minimax_m3/amd/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py new file mode 100644 index 000000000000..b80d3b8b3b8c --- /dev/null +++ b/vllm/models/minimax_m3/amd/model.py @@ -0,0 +1,1216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only MiniMax M3 (text backbone) model — AMD ROCm implementation. + +Self-contained per-platform impl (mirrors ``deepseek_v4/amd``). It is identical +to ``../nvidia/model.py`` except for RMS normalization: FlashInfer's Gemma +RMSNorm kernels are CUDA-only, so ``MiniMAXGemmaRMSNorm`` here uses a native +(FlashInfer-free) implementation. + +The MiniMax-M3-preview config selects a single set of branches: + * qk_norm_type == "per_head" + * hidden_act == "swigluoai" + * use_gemma_norm == True -> Gemma-style RMSNorm everywhere + * attention_output_gate == False + * scoring_func == "sigmoid" with a routing-bias correction term + * sparse_attention_config present -> a subset of layers run the extra + "index" attention branch. +""" + +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import PretrainedConfig + +from vllm import _custom_ops as ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, +) +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + fused_allreduce_gemma_rms_norm, +) +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + GateLinear, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + MinimaxM3QKVParallelLinearWithIndexer, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model +from vllm.models.minimax_m3.amd.ops import ( + gemma_fused_add_rmsnorm, + gemma_rmsnorm, + swiglu_oai_split, +) +from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer +from vllm.models.minimax_m3.common.mm_preprocess import ( + MiniMaxM3VLDummyInputsBuilder, + MiniMaxM3VLMultiModalProcessor, + MiniMaxM3VLProcessingInfo, +) +from vllm.models.minimax_m3.common.sparse_attention import ( + MiniMaxM3SparseBackend, + MiniMaxM3SparseImpl, + select_main_impl_cls, +) +from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + get_kv_quant_mode, +) + + +def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]: + """Layer ids whose attention runs the extra sparse "index" branch.""" + cfg = getattr(config, "sparse_attention_config", None) + if not cfg: + return set() + freq = cfg.get("sparse_attention_freq") + if freq is None: + return set() + return {i for i, f in enumerate(freq) if f != 0} + + +def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool: + """Whether this layer's MLP is a sparse MoE block (vs a dense MLP).""" + moe_layer_freq = getattr(config, "moe_layer_freq", None) + if moe_layer_freq is None: + return True + return moe_layer_freq[layer_id] != 0 + + +def _build_rotary_emb(config: PretrainedConfig, head_dim: int): + """Build the (partial NeoX) RoPE, honoring an optional ``rope_scaling`` config. + + Without scaling the cos/sin cache is sized to ``max_position_embeddings`` + (524288 native); a request whose positions exceed that reads the cache out of + bounds and the worker hard-crashes (no Python traceback). When ``rope_scaling`` + is set (e.g. YaRN ``factor: 2`` to reach 1M), thread it into ``get_rope`` so the + proper scaled embedding is built and its cache covers + ``original_max_position_embeddings * factor`` positions. Default behavior + (no scaling) is unchanged. Shared by the dense and sparse attention layers, and + the index branch reuses the returned module. + + Note: for the VL checkpoint, set ``rope_scaling`` on the *text* config + (``--hf-overrides '{"text_config":{"rope_scaling":{...}}}'``) -- that is the + config the decoder reads here; a top-level override does not reach it. + """ + rope_parameters = { + "rope_theta": config.rope_theta, + "partial_rotary_factor": config.partial_rotary_factor, + } + max_position = config.max_position_embeddings + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling: + rope_parameters.update(rope_scaling) + # HF uses "rope_type" (older configs: "type"); get_rope reads "rope_type". + if "rope_type" not in rope_parameters and "type" in rope_scaling: + rope_parameters["rope_type"] = rope_scaling["type"] + rope_parameters.setdefault( + "original_max_position_embeddings", config.max_position_embeddings + ) + factor = float(rope_scaling.get("factor", 1.0)) + # Cover the extended range (informational for get_rope's default branch; + # the YaRN embedding sizes its own cache from original * factor). + max_position = int(rope_parameters["original_max_position_embeddings"] * factor) + return get_rope( + head_dim, + max_position=max_position, + rope_parameters=rope_parameters, + ) + + +class MiniMAXGemmaRMSNorm(nn.Module): + """Gemma-style RMS normalization (native ROCm implementation). + + Normalizes in fp32 and scales by ``(1 + weight)`` — numerically equivalent + to the FlashInfer ``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` kernels + used in the NVIDIA path, which are unavailable on ROCm. When ``residual`` is + given, the fused add + norm returns the updated ``(normed, residual)`` pair. + + The fp32 normalize + scale + (optional) residual-add run in a single fused + Triton pass (``amd.ops.gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm``) instead + of a chain of elementwise PyTorch kernels. + """ + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + def forward( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is None: + return gemma_rmsnorm(x, self.weight, self.variance_epsilon) + return gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon) + + +class MiniMaxM3MLP(nn.Module): + """Dense SwiGLU-OAI MLP (used by the leading dense layers).""" + + def __init__( + self, + config: PretrainedConfig, + intermediate_size: int, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if config.hidden_act != "swigluoai": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only swigluoai is supported." + ) + # gate * sigmoid(alpha * gate) * (up + beta), with both halves clamped. + # Kept as our fp32 Triton kernel (not the #22 SWIGLUOAI_UNINTERLEAVE op + # ``silu_and_mul_with_clamp``): that op IS built on ROCm but rounds + # intermediates to bf16 (rel ~3e-3 vs our fp32 ~1e-6), which costs gsm8k + # accuracy since this activation feeds the MXFP8 quant + MoE. + self.swiglu_alpha = config.swiglu_alpha + self.swiglu_beta = config.swiglu_beta + self.swiglu_limit = config.swiglu_limit + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = swiglu_oai_split( + gate_up, + alpha=self.swiglu_alpha, + beta=self.swiglu_beta, + limit=self.swiglu_limit, + ) + x, _ = self.down_proj(x) + return x + + +class MiniMaxM3MoE(nn.Module): + """Sigmoid-routed MoE block with a routing-bias correction and a shared + expert.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + if self.tp_size > config.num_local_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_local_experts}." + ) + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.n_shared_experts = getattr(config, "n_shared_experts", None) + + # Sigmoid routing uses a per-expert score-correction bias for selection. + self.use_routing_bias = getattr(config, "use_routing_bias", False) + if self.use_routing_bias: + self.e_score_correction_bias = nn.Parameter( + torch.empty(config.num_local_experts, dtype=torch.float32) + ) + self.e_score_correction_bias.weight_loader = ( + MiniMaxM3MoE.ebias_weight_loader + ) + else: + self.e_score_correction_bias = None + + # Router weights are stored in fp32; GateLinear upcasts the bf16 + # activations and computes the gate in fp32 (fp32 router logits). + self.gate = GateLinear( + config.hidden_size, + config.num_local_experts, + bias=False, + params_dtype=torch.float32, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.shared_experts: MiniMaxM3MLP | None = None + if self.n_shared_experts: + self.shared_experts = MiniMaxM3MLP( + config=config, + intermediate_size=config.intermediate_size * self.n_shared_experts, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + + self.experts = FusedMoE( + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + scoring_func=config.scoring_func, + e_score_correction_bias=self.e_score_correction_bias, + renormalize=True, + activation="swigluoai_uninterleave", + swiglu_limit=config.swiglu_limit, + swiglu_alpha=config.swiglu_alpha, + swiglu_beta=config.swiglu_beta, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, + router_logits_dtype=self.gate.out_dtype, + shared_experts=self.shared_experts, + quant_config=quant_config, + prefix=f"{prefix}.experts", + ) + + @staticmethod + def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None: + assert param.size() == loaded_weight.size() + param.data.copy_(loaded_weight.to(torch.float32)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # router_logits: (num_tokens, n_experts); GateLinear casts to fp32. + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class MiniMaxM3Attention(nn.Module): + """Dense attention with per-head QK norm and partial RoPE.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + cache_config: CacheConfig | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + # reduce_results=False: the attention all-reduce is fused with the + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer + # via fused_allreduce_gemma_rms_norm. + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + reduce_results=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True). + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. Honors + # config.rope_scaling (e.g. YaRN) so long-context positions are covered. + self.rotary_emb = _build_rotary_emb(config, self.head_dim) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place (dense + # mode: no index branch, no KV-cache insert). Matches nvidia/model.py and + # replaces the unfused split -> q_norm/k_norm -> rotary_emb chain; verified + # bit-equivalent on ROCm (q/k rel ~2e-3 bf16 noise, v untouched). + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + self.rotary_emb.cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + self.rotary_emb.rotary_dim, + self.q_norm.variance_epsilon, + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): + """Block-sparse attention layer with the lightning-indexer branch. + + This is a merged attention layer: it owns the projections (qkv + index + q/k), per-head QK norms and RoPE, *and* the attention-backend wiring that a + generic ``Attention`` layer would normally provide — it binds the + ``MiniMaxM3SparseBackend`` + main impl, registers the main paged K/V cache, + and owns the lightning indexer (``MiniMaxM3Indexer``), which holds the + index-key side cache. + + The index branch (index_{q,k}_proj + index_{q,k}_norm) feeds the sparse + top-k block selection. M3 always disables the index value/output + projections (``sparse_disable_index_value`` set for every sparse layer), so + ``index_{v,o}_proj`` are never created. + """ + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + cache_config: CacheConfig | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + # Sparse "index" branch dims. index_q has the same head count as the KV + # heads (sparse_num_index_heads == num_key_value_heads), so it shards + # identically -- including replication when tp_size > num_key_value_heads. + sparse_cfg = config.sparse_attention_config + self.total_idx_heads = sparse_cfg["sparse_num_index_heads"] + self.num_idx_heads = self.num_kv_heads + self.idx_head_dim = sparse_cfg["sparse_index_dim"] + self.index_q_size = self.num_idx_heads * self.idx_head_dim + + # Single fused projection: q, k, v, index_q, index_k in one GEMM. + self.qkv_proj = MinimaxM3QKVParallelLinearWithIndexer( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + self.total_idx_heads, + self.idx_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + # reduce_results=False: the attention all-reduce is fused with the + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer + # via fused_allreduce_gemma_rms_norm. + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + reduce_results=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True). + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. Honors + # config.rope_scaling (e.g. YaRN) so long-context positions are covered. + self.rotary_emb = _build_rotary_emb(config, self.head_dim) + + self.index_q_norm = MiniMAXGemmaRMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + self.index_k_norm = MiniMAXGemmaRMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + self.index_rotary_emb = self.rotary_emb + + # Attention-backend wiring. + vllm_config = get_current_vllm_config() + self.layer_name = f"{prefix}.attn" + self.kv_cache_dtype = ( + cache_config.cache_dtype if cache_config is not None else "auto" + ) + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) + # fp8 main-K/V cache: the fused qknorm+rope+kv-insert op is bf16-cache-only + # (asserts kv_cache dtype == qkv), so on the fp8 path we run it in + # norm+rope-only mode and write the cache via the fp8-capable + # reshape_and_cache_flash in _insert_kv. (index cache stays bf16.) + self._fp8_kv = "fp8" in self.kv_cache_dtype + + self.attn_backend = MiniMaxM3SparseBackend + # Indexer and main attention are separate impls. On ROCm the SM100 gate + # is always False, so both pick Triton and the index cache stays bf16. + # impl is AttentionImplBase (broader than AttentionLayerBase's annotation). + self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment] + topk_blocks=sparse_cfg["sparse_topk_blocks"], + kv_cache_dtype=self.kv_cache_dtype, + )( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + kv_cache_dtype=self.kv_cache_dtype, + topk_blocks=sparse_cfg["sparse_topk_blocks"], + sparse_block_size=sparse_cfg["sparse_block_size"], + ) + # Self-contained nn.Module: owns its side cache, selects its impl in init + # (Triton on ROCm, where the SM100 gate is always False). + self.indexer = MiniMaxM3Indexer( + num_kv_heads=self.num_kv_heads, + scale=self.scaling, + topk_blocks=sparse_cfg["sparse_topk_blocks"], + sparse_block_size=sparse_cfg["sparse_block_size"], + num_index_heads=self.num_idx_heads, + index_head_dim=self.idx_head_dim, + prefix=self.layer_name, + init_blocks=sparse_cfg.get("sparse_init_block", 0), + local_blocks=sparse_cfg.get("sparse_local_block", 0), + score_type=sparse_cfg.get("sparse_score_type", "max"), + cache_config=cache_config, + ) + + # Register the main K/V cache so the KV-cache manager allocates it. + compilation_config = vllm_config.compilation_config + if self.layer_name in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {self.layer_name}") + compilation_config.static_forward_context[self.layer_name] = self + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache + + def get_attn_backend(self) -> type[MiniMaxM3SparseBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + # Main GQA K/V cache. Block size may change after load, refresh it. + return FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + head_size_v=self.head_dim, + dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + ) + + def _insert_kv( + self, + key: torch.Tensor, + value: torch.Tensor, + index_key: torch.Tensor, + main_slot_mapping: torch.Tensor, + index_slot_mapping: torch.Tensor, + ) -> None: + """Write main K/V (fp8-quantizing) and index-K into their paged caches. + + Used only on the fp8-KV path: the fused #20 op is bf16-cache-only, so it + runs in norm+rope-only mode and the (already normed/roped) k/v/index_k are + written here via ``reshape_and_cache_flash`` (which honors kv_cache_dtype, + unit scale -- matching the fp8 read path added in #33). Mirrors the + pre-#20 unfused insert. The index cache stays bf16 (no quant). + """ + key_cache, value_cache = self.kv_cache.unbind(1) + scale = torch.ones((), device=key.device) + ops.reshape_and_cache_flash( + key.view(-1, self.num_kv_heads, self.head_dim), + value.view(-1, self.num_kv_heads, self.head_dim), + key_cache, + value_cache, + main_slot_mapping, + self.kv_cache_dtype, + scale, + scale, + ) + idx_cache = self.indexer.index_cache.kv_cache.view(-1, self.idx_head_dim) + idx_cache[index_slot_mapping] = index_key.to(idx_cache.dtype) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # Single fused projection emitting [q | k | v | index_q | index_k]. + qkv, _ = self.qkv_proj(hidden_states) + + # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on the + # main (q/k) and index (index_q/index_k) branches, all read straight out + # of the single fused ``qkv`` tensor. Once the paged caches are bound the + # kernel also inserts k/v and the index key into them (each with its own + # slot_mapping); the memory-profiling run (caches unbound, no slot_mapping) + # short-circuits to zeros below. Replaces the + # q_norm/k_norm/rotary_emb/index_*_norm/index_rotary_emb/_insert_kv chain. + # (#20 fused_minimax_m3_qknorm_rope_kv_insert; HIP/CDNA path. The main and + # index slot mappings are read from the forward context's slot_mapping + # dict, matching the breakable-cudagraph path -- see nvidia/model.py.) + cos_sin_cache = self.rotary_emb.cos_sin_cache + rotary_dim = self.rotary_emb.rotary_dim + eps = self.q_norm.variance_epsilon + num_tokens = qkv.shape[0] + + fwd_slot_mapping = get_forward_context().slot_mapping + if ( + not isinstance(fwd_slot_mapping, dict) + or self.layer_name not in fwd_slot_mapping + ): + # Memory-profiling run: caches not yet bound, slot_mapping is empty. + return qkv.new_zeros((num_tokens, self.hidden_size)) + + main_slot_mapping = fwd_slot_mapping[self.layer_name] + index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] + q = qkv.new_empty((num_tokens, self.q_size)) + index_q = qkv.new_empty((num_tokens, self.index_q_size)) + # On the fp8-KV path the fused op cannot write the (fp8) cache, so pass + # kv_cache/index_cache = None -> insert_kv=False (norm+rope only): it still + # de-interleaves q/index_q and rewrites the normed/roped k & index_k in + # place in qkv, leaving v raw (correct -- v is never normed/roped). We then + # write the cache via _insert_kv below. + insert_via_fused = not self._fp8_kv + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + rotary_dim, + eps, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.num_idx_heads, + main_slot_mapping, + index_slot_mapping, + self.kv_cache if insert_via_fused else None, + self.indexer.index_cache.kv_cache if insert_via_fused else None, + self.kv_cache.size(2), # paged-cache block size + q, + index_q, + ) + if not insert_via_fused: + # Extract the normed/roped k, raw v, normed/roped index_k from qkv + # ([q | k | v | index_q | index_k], all head_dim=128) and fp8-insert. + kv = self.num_kv_heads * self.head_dim + # These are strided views into qkv (row stride = full qkv width), but + # their last dim is contiguous, so `_insert_kv`'s `.view(-1, nkv, + # head_dim)` works on them and `reshape_and_cache_flash` honors the + # input stride -- no `.contiguous()` needed (verified bit-identical; + # avoids a [N, kv] copy per step on the fp8-KV path). + k = qkv[:, self.q_size : self.q_size + kv] + v = qkv[:, self.q_size + kv : self.q_size + 2 * kv] + ik0 = self.q_size + 2 * kv + self.index_q_size + index_k = qkv[:, ik0 : ik0 + self.num_idx_heads * self.idx_head_dim] + self._insert_kv(k, v, index_k, main_slot_mapping, index_slot_mapping) + + output = torch.empty_like(q) + attn_output = self._run_attention(q, index_q, output) + output, _ = self.o_proj(attn_output) + return output + + @eager_break_during_capture + def _run_attention( + self, + query: torch.Tensor, + index_query: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + # Single eager break around both: their split-K kernels read per-request + # metadata and can't be captured into a cudagraph. + topk_idx = self.indexer(index_query) + return self.impl.forward(self, query, self.kv_cache, topk_idx, output) + + +class MiniMaxM3DecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + prefix: str, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + force_sparse_attn: bool = False, + force_moe: bool = False, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + # DecoderLayers are created with `make_layers` which passes the prefix + # with the layer's index. + layer_id = int(prefix.split(sep=".")[-1]) + self.layer_id = layer_id + + is_sparse_attention_layer = ( + force_sparse_attn or layer_id in _sparse_attention_layer_ids(config) + ) + + if is_sparse_attention_layer: + self.self_attn = MiniMaxM3SparseAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + cache_config=cache_config, + ) + else: + self.self_attn = MiniMaxM3Attention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + cache_config=cache_config, + ) + + # Dense layers store the FFN under `mlp`; MoE layers under + # `block_sparse_moe` -- matching the checkpoint's naming. + self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id) + if self.is_moe_layer: + self.block_sparse_moe = MiniMaxM3MoE( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + ) + else: + self.mlp = MiniMaxM3MLP( + config=config, + intermediate_size=config.dense_intermediate_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm. + self.input_layernorm = MiniMAXGemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = MiniMAXGemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + hidden_states, residual = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + ffn = self.block_sparse_moe if self.is_moe_layer else self.mlp + hidden_states = ffn(hidden_states) + return hidden_states, residual + + +class MiniMaxM3Model(nn.Module, EagleModelMixin): + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_text_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.config = config + + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: MiniMaxM3DecoderLayer( + config, + prefix, + cache_config=cache_config, + quant_config=quant_config, + ), + prefix=f"{prefix}.layers", + ) + + self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + + # EAGLE3 is not yet compatible with pipeline parallel + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): + hidden_states, residual = layer(positions, hidden_states, residual) + self._maybe_add_hidden_state( + aux_hidden_states, idx + 1, hidden_states, residual + ) + + hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # Checkpoint experts use w1=gate, w2=down, w3=up. + return 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, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # q/k/v_proj -> fused qkv_proj; gate_proj/up_proj -> fused gate_up_proj + # (dense MLP and shared expert). On sparse layers the indexer + # index_q/index_k_proj fold into the same fused qkv_proj + # (MinimaxM3QKVParallelLinearWithIndexer); these entries simply never match on + # dense layers, whose checkpoints have no index_*_proj weights. Leading + # dots keep `q_proj`/`k_proj` from matching `index_q_proj`/`index_k_proj` + # (preceded by `_`, not `.`). + 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), + ] + + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = self.get_expert_mapping() + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + # The MTP module is not modeled yet. + if "mtp." in name: + continue + + # The checkpoint stores block scales as ``weight_scale_inv``; the + # ModelOpt MXFP8 layers expose them as ``weight_scale``. + if "weight_scale_inv" in name: + name = name.replace("weight_scale_inv", "weight_scale") + + 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.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + 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 is_pp_missing_parameter(name, self): + continue + 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: + if name.endswith(".bias") and name not in params_dict: + continue + remapped = maybe_remap_kv_scale_name(name, params_dict) + if remapped is None: + continue + name = remapped + if is_pp_missing_parameter(name, self): + continue + # Modules not modeled yet (e.g. attention) are skipped until + # they are ported. + if name not in params_dict: + continue + 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 + + +class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3): + """MiniMax M3 (sparse/dense backbone) for causal language modeling.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + self.model = MiniMaxM3Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(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, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + +# TODO(refactor): this VL wrapper is platform-agnostic and byte-identical to the +# NVIDIA copy — it only orchestrates the shared vision tower + the per-platform +# language model (resolved via ``init_vllm_registered_model``). Hoist it into +# ``common/`` to drop the amd/nvidia duplication once the split stabilizes. +@MULTIMODAL_REGISTRY.register_processor( + MiniMaxM3VLMultiModalProcessor, + info=MiniMaxM3VLProcessingInfo, + dummy_inputs=MiniMaxM3VLDummyInputsBuilder, +) +class MiniMaxM3SparseForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsEagle3 +): + """Top-level (VL) entry point for MiniMax M3. + + Owns the shared MiniMax-M3 vision tower on ROCm and delegates text + generation to the AMD language-model path. + """ + + # The vision tower runs replicated per rank under ``--mm-encoder-tp-mode + # data``; ``run_dp_sharded_mrope_vision_model`` shards the work across + # ranks (see ``_process_image_input`` / ``_process_video_input``). + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "multi_modal_projector.": "vision_tower.multi_modal_projector.", + "patch_merge_mlp.": "vision_tower.patch_merge_mlp.", + }, + orig_to_new_substr={ + ".mlp.fc1.": ".fc1.", + ".mlp.fc2.": ".fc2.", + }, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return MiniMaxM3VLProcessingInfo.IMAGE_TOKEN + if modality == "video": + return MiniMaxM3VLProcessingInfo.VIDEO_TOKEN + raise ValueError(f"Unsupported modality: {modality!r}") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.config = config + self.quant_config = vllm_config.quant_config + self.multimodal_config = vllm_config.model_config.multimodal_config + assert self.multimodal_config is not None + self.use_data_parallel = self.multimodal_config.mm_encoder_tp_mode == "data" + + text_hidden_size = getattr(config.text_config, "hidden_size", None) + assert text_hidden_size is not None, "text_config.hidden_size is required" + projector_hidden_size = getattr(config, "projector_hidden_size", None) + + with self._mark_tower_model(vllm_config, {"image", "video"}): + vision_config = config.vision_config + self.vision_tower = MiniMaxVLVisionModel( + config=PretrainedConfig.from_dict(vision_config), + text_hidden_size=text_hidden_size, + projector_hidden_size=projector_hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["MiniMaxM3SparseForCausalLM"], + ) + + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: + pixel_values = kwargs.pop("pixel_values", None) + image_grid_thw = kwargs.pop("image_grid_thw", None) + if pixel_values is None: + return None + return {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw} + + def _parse_and_validate_video_input(self, **kwargs: object) -> dict | None: + pixel_values_videos = kwargs.pop("pixel_values_videos", None) + video_grid_thw = kwargs.pop("video_grid_thw", None) + if pixel_values_videos is None: + return None + return { + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": video_grid_thw, + } + + def _process_image_input(self, image_input: dict) -> tuple[torch.Tensor, ...]: + pixel_values: torch.Tensor = image_input["pixel_values"].type( + self.vision_tower.dtype + ) + grid_thw: torch.Tensor = image_input["image_grid_thw"] + assert grid_thw.ndim == 2 + + if self.use_data_parallel: + # Already returns a per-item tuple of embeddings. + return run_dp_sharded_mrope_vision_model( + self.vision_tower, + pixel_values, + grid_thw.tolist(), + rope_type="rope_3d", + ) + + image_embeds = self.vision_tower( + pixel_values=pixel_values, + grid_thw=grid_thw.tolist(), + ) + + # Split the concatenated output into one tensor per image item. + merge_size = self.vision_tower.spatial_merge_size + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist() + return image_embeds.split(sizes) + + def _process_video_input(self, video_input: dict) -> tuple[torch.Tensor, ...]: + pixel_values: torch.Tensor = video_input["pixel_values_videos"].type( + self.vision_tower.dtype + ) + grid_thw: torch.Tensor = video_input["video_grid_thw"] + assert grid_thw.ndim == 2 + + if self.use_data_parallel: + # Already returns a per-item tuple of embeddings. + return run_dp_sharded_mrope_vision_model( + self.vision_tower, + pixel_values, + grid_thw.tolist(), + rope_type="rope_3d", + ) + + video_embeds = self.vision_tower( + pixel_values=pixel_values, + grid_thw=grid_thw.tolist(), + ) + + # Split the concatenated output into one tensor per video item. + merge_size = self.vision_tower.spatial_merge_size + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist() + return video_embeds.split(sizes) + + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> dict[str, dict]: + mm_input_by_modality: dict[str, dict] = {} + for input_key in kwargs: + if input_key == "pixel_values" and "image" not in mm_input_by_modality: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is not None: + mm_input_by_modality["image"] = image_input + if ( + input_key == "pixel_values_videos" + and "video" not in mm_input_by_modality + ): + video_input = self._parse_and_validate_video_input(**kwargs) + if video_input is not None: + mm_input_by_modality["video"] = video_input + return mm_input_by_modality + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) + if not mm_input_by_modality: + return [] + + multimodal_embeddings: list[torch.Tensor] = [] + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "image": + image_embeddings = self._process_image_input(multimodal_input) + multimodal_embeddings.extend(image_embeddings) + if modality == "video": + video_embeddings = self._process_video_input(multimodal_input) + multimodal_embeddings.extend(video_embeddings) + + return tuple(multimodal_embeddings) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.language_model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.language_model.get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/models/minimax_m3/amd/mtp.py b/vllm/models/minimax_m3/amd/mtp.py new file mode 100644 index 000000000000..f62face1d2e1 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mtp.py @@ -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 diff --git a/vllm/models/minimax_m3/amd/ops/__init__.py b/vllm/models/minimax_m3/amd/ops/__init__.py new file mode 100644 index 000000000000..22d96f9de979 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/__init__.py @@ -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", +] diff --git a/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py b/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py new file mode 100644 index 000000000000..cb74877f6824 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py @@ -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) diff --git a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py new file mode 100644 index 000000000000..836649b725b4 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py @@ -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) diff --git a/vllm/models/minimax_m3/common/__init__.py b/vllm/models/minimax_m3/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py new file mode 100644 index 000000000000..e43ad60914f3 --- /dev/null +++ b/vllm/models/minimax_m3/common/indexer.py @@ -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) diff --git a/vllm/models/minimax_m3/common/mm_preprocess.py b/vllm/models/minimax_m3/common/mm_preprocess.py new file mode 100644 index 000000000000..208adfffea51 --- /dev/null +++ b/vllm/models/minimax_m3/common/mm_preprocess.py @@ -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 diff --git a/vllm/models/minimax_m3/common/ops/__init__.py b/vllm/models/minimax_m3/common/ops/__init__.py new file mode 100644 index 000000000000..b3a7c2d9f6ec --- /dev/null +++ b/vllm/models/minimax_m3/common/ops/__init__.py @@ -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", +] diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py new file mode 100644 index 000000000000..c32ff38d998f --- /dev/null +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -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 diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py new file mode 100644 index 000000000000..40287e166b2a --- /dev/null +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -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, + ) diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py new file mode 100644 index 000000000000..8cca0e8e2998 --- /dev/null +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -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 diff --git a/vllm/models/minimax_m3/common/vision_tower.py b/vllm/models/minimax_m3/common/vision_tower.py new file mode 100644 index 000000000000..23b8b3ed3197 --- /dev/null +++ b/vllm/models/minimax_m3/common/vision_tower.py @@ -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 diff --git a/vllm/models/minimax_m3/nvidia/__init__.py b/vllm/models/minimax_m3/nvidia/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py new file mode 100644 index 000000000000..e2cd62704fda --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -0,0 +1,1177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only MiniMax M3 (text backbone) model. + +The MiniMax-M3-preview config selects a single set of branches: + * qk_norm_type == "per_head" + * hidden_act == "swigluoai" + * use_gemma_norm == True -> Gemma-style RMSNorm everywhere + * attention_output_gate == False + * scoring_func == "sigmoid" with a routing-bias correction term + * sparse_attention_config present -> a subset of layers run the extra + "index" attention branch. +""" + +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import PretrainedConfig + +from vllm import _custom_ops as ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.activation import SiluAndMulWithClamp +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + fused_allreduce_gemma_rms_norm, +) +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + GateLinear, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + MinimaxM3QKVParallelLinearWithIndexer, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model +from vllm.models.minimax_m3.common.indexer import ( + MiniMaxM3Indexer, + MiniMaxM3IndexerMetadata, +) +from vllm.models.minimax_m3.common.mm_preprocess import ( + MiniMaxM3VLDummyInputsBuilder, + MiniMaxM3VLMultiModalProcessor, + MiniMaxM3VLProcessingInfo, +) +from vllm.models.minimax_m3.common.sparse_attention import ( + MiniMaxM3SparseBackend, + MiniMaxM3SparseImpl, + MiniMaxM3SparseMetadata, + select_main_impl_cls, +) +from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + get_kv_quant_mode, +) + + +def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]: + """Layer ids whose attention runs the extra sparse "index" branch.""" + cfg = getattr(config, "sparse_attention_config", None) + if not cfg: + return set() + freq = cfg.get("sparse_attention_freq") + if freq is None: + return set() + return {i for i, f in enumerate(freq) if f != 0} + + +def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool: + """Whether this layer's MLP is a sparse MoE block (vs a dense MLP).""" + moe_layer_freq = getattr(config, "moe_layer_freq", None) + if moe_layer_freq is None: + return True + return moe_layer_freq[layer_id] != 0 + + +class MiniMAXGemmaRMSNorm(nn.Module): + """Gemma-style RMS normalization backed by FlashInfer kernels. + + When ``residual`` is given, the fused add + norm runs in place and the + updated ``(x, residual)`` pair is returned. + """ + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + def forward( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + from flashinfer.norm import gemma_fused_add_rmsnorm, gemma_rmsnorm + + if residual is None: + return gemma_rmsnorm(x, self.weight, self.variance_epsilon) + + # gemma_fused_add_rmsnorm mutates x and residual in place. + gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon) + return x, residual + + +class MiniMaxM3MLP(nn.Module): + """Dense SwiGLU-OAI MLP (used by the leading dense layers).""" + + def __init__( + self, + config: PretrainedConfig, + intermediate_size: int, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if config.hidden_act != "swigluoai": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only swigluoai is supported." + ) + # gate * sigmoid(alpha * gate) * (up + beta), with both halves clamped. + self.act_fn = SiluAndMulWithClamp( + swiglu_limit=config.swiglu_limit, + alpha=config.swiglu_alpha, + beta=config.swiglu_beta, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MiniMaxM3MoE(nn.Module): + """Sigmoid-routed MoE block with a routing-bias correction and a shared + expert.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + if self.tp_size > config.num_local_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_local_experts}." + ) + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.n_shared_experts = getattr(config, "n_shared_experts", None) + + # Sigmoid routing uses a per-expert score-correction bias for selection. + self.use_routing_bias = getattr(config, "use_routing_bias", False) + if self.use_routing_bias: + self.e_score_correction_bias = nn.Parameter( + torch.empty(config.num_local_experts, dtype=torch.float32) + ) + self.e_score_correction_bias.weight_loader = ( + MiniMaxM3MoE.ebias_weight_loader + ) + else: + self.e_score_correction_bias = None + + # Router weights are stored in fp32; GateLinear upcasts the bf16 + # activations and computes the gate in fp32 (fp32 router logits). + self.gate = GateLinear( + config.hidden_size, + config.num_local_experts, + bias=False, + params_dtype=torch.float32, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.shared_experts: MiniMaxM3MLP | None = None + if self.n_shared_experts: + self.shared_experts = MiniMaxM3MLP( + config=config, + intermediate_size=config.intermediate_size * self.n_shared_experts, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + + self.experts = FusedMoE( + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + scoring_func=config.scoring_func, + e_score_correction_bias=self.e_score_correction_bias, + renormalize=True, + # w13 (gate_up_proj) is loaded packed via MergedColumnParallelLinear + # ([all gates; all ups]), so use the uninterleaved SwiGLU-OAI variant + # rather than the interleaved gpt-oss layout. + activation="swigluoai_uninterleave", + swiglu_limit=config.swiglu_limit, + swiglu_alpha=config.swiglu_alpha, + swiglu_beta=config.swiglu_beta, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=True, + router_logits_dtype=self.gate.out_dtype, + shared_experts=self.shared_experts, + quant_config=quant_config, + prefix=f"{prefix}.experts", + ) + + @staticmethod + def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None: + assert param.size() == loaded_weight.size() + param.data.copy_(loaded_weight.to(torch.float32)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # router_logits: (num_tokens, n_experts); GateLinear casts to fp32. + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class MiniMaxM3Attention(nn.Module): + """Dense attention with per-head QK norm and partial RoPE.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + cache_config: CacheConfig | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + # reduce_results=False: the attention all-reduce is fused with the + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer + # via fused_allreduce_gemma_rms_norm. + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + reduce_results=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True). + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters={ + "rope_theta": config.rope_theta, + "partial_rotary_factor": config.partial_rotary_factor, + }, + ) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place. + + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + self.rotary_emb.cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + self.rotary_emb.rotary_dim, + self.q_norm.variance_epsilon, + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): + """Block-sparse attention layer with the lightning-indexer branch. + + This is a merged attention layer: it owns the projections (qkv + index + q/k), per-head QK norms and RoPE, *and* the attention-backend wiring that a + generic ``Attention`` layer would normally provide — it binds the + ``MiniMaxM3SparseBackend`` + main impl, registers the main paged K/V cache, + and owns the lightning indexer (``MiniMaxM3Indexer``), which holds the + index-key side cache. + + The index branch (index_{q,k}_proj + index_{q,k}_norm) feeds the sparse + top-k block selection. M3 always disables the index value/output + projections (``sparse_disable_index_value`` set for every sparse layer), so + ``index_{v,o}_proj`` are never created. + """ + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + cache_config: CacheConfig | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + # Sparse "index" branch dims. index_q has the same head count as the KV + # heads (sparse_num_index_heads == num_key_value_heads), so it shards + # identically -- including replication when tp_size > num_key_value_heads. + sparse_cfg = config.sparse_attention_config + self.total_idx_heads = sparse_cfg["sparse_num_index_heads"] + self.num_idx_heads = self.num_kv_heads + self.idx_head_dim = sparse_cfg["sparse_index_dim"] + self.index_q_size = self.num_idx_heads * self.idx_head_dim + + # Single fused projection: q, k, v, index_q, index_k in one GEMM. + self.qkv_proj = MinimaxM3QKVParallelLinearWithIndexer( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + self.total_idx_heads, + self.idx_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + # reduce_results=False: the attention all-reduce is fused with the + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer + # via fused_allreduce_gemma_rms_norm. + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + reduce_results=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True). + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters={ + "rope_theta": config.rope_theta, + "partial_rotary_factor": config.partial_rotary_factor, + }, + ) + + self.index_q_norm = MiniMAXGemmaRMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + self.index_k_norm = MiniMAXGemmaRMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + self.index_rotary_emb = self.rotary_emb + + # Attention-backend wiring. + vllm_config = get_current_vllm_config() + self.layer_name = f"{prefix}.attn" + self.kv_cache_dtype = ( + cache_config.cache_dtype if cache_config is not None else "auto" + ) + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) + # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main + # cache (--attention-config '{"indexer_kv_dtype": ...}'). + self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype + + self.attn_backend = MiniMaxM3SparseBackend + # Indexer (top-k selection) and main attention are separate impls, each + # picking Triton vs MSA off its cache dtype. impl is AttentionImplBase + # (broader than the AttentionImpl that AttentionLayerBase annotates). + self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment] + topk_blocks=sparse_cfg["sparse_topk_blocks"], + kv_cache_dtype=self.kv_cache_dtype, + )( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + kv_cache_dtype=self.kv_cache_dtype, + topk_blocks=sparse_cfg["sparse_topk_blocks"], + sparse_block_size=sparse_cfg["sparse_block_size"], + ) + # Self-contained nn.Module: owns its side cache, selects its impl in init. + self.indexer = MiniMaxM3Indexer( + num_kv_heads=self.num_kv_heads, + scale=self.scaling, + topk_blocks=sparse_cfg["sparse_topk_blocks"], + sparse_block_size=sparse_cfg["sparse_block_size"], + num_index_heads=self.num_idx_heads, + index_head_dim=self.idx_head_dim, + prefix=self.layer_name, + init_blocks=sparse_cfg.get("sparse_init_block", 0), + local_blocks=sparse_cfg.get("sparse_local_block", 0), + score_type=sparse_cfg.get("sparse_score_type", "max"), + cache_config=cache_config, + indexer_kv_dtype=self.indexer_kv_dtype, + ) + + # Register the main K/V cache so the KV-cache manager allocates it. + compilation_config = vllm_config.compilation_config + if self.layer_name in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {self.layer_name}") + compilation_config.static_forward_context[self.layer_name] = self + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache + + def get_attn_backend(self) -> type[MiniMaxM3SparseBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + # Main GQA K/V cache. Block size may change after load, refresh it. + return FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + head_size_v=self.head_dim, + dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + ) + + def _insert_kv( + self, key: torch.Tensor, value: torch.Tensor, index_key: torch.Tensor + ) -> None: + """Write main K/V and index-K into their paged caches. + + No-op during the profiling run, where caches are not yet bound and + ``attn_metadata`` is None. + """ + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return + main_meta = attn_metadata[self.layer_name] + index_meta = attn_metadata[self.indexer.index_cache.prefix] + assert isinstance(main_meta, MiniMaxM3SparseMetadata) + assert isinstance(index_meta, MiniMaxM3IndexerMetadata) + + # Identity scale: unused for the bf16 cache, required arg of the op. + key_cache, value_cache = self.kv_cache.unbind(1) + scale = torch.ones((), device=key.device) + ops.reshape_and_cache_flash( + key.view(-1, self.num_kv_heads, self.head_dim), + value.view(-1, self.num_kv_heads, self.head_dim), + key_cache, + value_cache, + main_meta.slot_mapping, + self.kv_cache_dtype, + scale, + scale, + ) + + # Index-key cache: single vector per token, scatter by slot. + idx_cache = self.indexer.index_cache.kv_cache.view(-1, self.idx_head_dim) + idx_cache[index_meta.slot_mapping] = index_key.to(idx_cache.dtype) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # Single fused projection emitting [q | k | v | index_q | index_k]. + qkv, _ = self.qkv_proj(hidden_states) + + # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on the + # main (q/k) and index (index_q/index_k) branches, all read straight out + # of the single fused ``qkv`` tensor (the "5 results"). Once the paged + # caches are bound the kernel also inserts k/v and the index key into + # them; the initial memory-profiling run (caches unbound, no slot_mapping) + # short-circuits to zeros below. Replaces the + # q_norm/k_norm/rotary_emb/index_*_norm/index_rotary_emb/_insert_kv + # sequence. k/v and index_k are rewritten in place inside qkv (and + # scatter-inserted into the caches); q and index_q are de-interleaved + # straight into the dedicated contiguous ``q``/``index_q`` buffers below. + + cos_sin_cache = self.rotary_emb.cos_sin_cache + rotary_dim = self.rotary_emb.rotary_dim + eps = self.q_norm.variance_epsilon + num_tokens = qkv.shape[0] + + fwd_slot_mapping = get_forward_context().slot_mapping + if ( + not isinstance(fwd_slot_mapping, dict) + or self.layer_name not in fwd_slot_mapping + ): + # Memory-profiling run: caches not yet bound, slot_mapping is empty. + return qkv.new_zeros((num_tokens, self.hidden_size)) + + main_slot_mapping = fwd_slot_mapping[self.layer_name] + index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] + q = qkv.new_empty((num_tokens, self.q_size)) + index_q = qkv.new_empty((num_tokens, self.index_q_size)) + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + rotary_dim, + eps, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.num_idx_heads, + main_slot_mapping, + index_slot_mapping, + self.kv_cache, + self.indexer.index_cache.kv_cache, + self.kv_cache.size(2), # paged-cache block size + q, + index_q, + ) + + output = torch.empty_like(q) + attn_output = self._run_attention(q, index_q, output) + output, _ = self.o_proj(attn_output) + return output + + @eager_break_during_capture + def _run_attention( + self, + query: torch.Tensor, + index_query: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + # Single eager break around both: their split-K kernels read per-request + # metadata and can't be captured into a cudagraph. + topk_idx = self.indexer(index_query) + return self.impl.forward(self, query, self.kv_cache, topk_idx, output) + + +class MiniMaxM3DecoderLayer(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str, + force_sparse_attn: bool = False, + force_moe: bool = False, + is_mtp_block: bool = False, + ) -> None: + super().__init__() + if is_mtp_block: + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + else: + config = vllm_config.model_config.hf_text_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.hidden_size = config.hidden_size + # DecoderLayers are created with `make_layers` which passes the prefix + # with the layer's index. + layer_id = int(prefix.split(sep=".")[-1]) + self.layer_id = layer_id + + # Complete the preceding dense MLP's deferred all-reduce + # (reduce_results=False), fused into this layer's input_layernorm. + # Disable this fusion when PP is set + self.fuse_input_allreduce = ( + layer_id > 0 + and not _is_moe_layer(config, layer_id - 1) + and vllm_config.parallel_config.pipeline_parallel_size == 1 + ) + + is_sparse_attention_layer = ( + force_sparse_attn or layer_id in _sparse_attention_layer_ids(config) + ) + + if is_sparse_attention_layer: + self.self_attn = MiniMaxM3SparseAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + cache_config=cache_config, + ) + else: + self.self_attn = MiniMaxM3Attention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + cache_config=cache_config, + ) + + # Dense layers store the FFN under `mlp`; MoE layers under + # `block_sparse_moe` -- matching the checkpoint's naming. + self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id) + if self.is_moe_layer: + self.block_sparse_moe = MiniMaxM3MoE( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + ) + else: + self.mlp = MiniMaxM3MLP( + config=config, + intermediate_size=config.dense_intermediate_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + reduce_results=vllm_config.parallel_config.pipeline_parallel_size > 1, + ) + + # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm. + self.input_layernorm = MiniMAXGemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = MiniMAXGemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.fuse_input_allreduce and residual is not None: + hidden_states, residual = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.input_layernorm + ) + else: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + hidden_states, residual = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + ffn = self.block_sparse_moe if self.is_moe_layer else self.mlp + hidden_states = ffn(hidden_states) + return hidden_states, residual + + +class MiniMaxM3Model(nn.Module, EagleModelMixin): + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + self.config = config + + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: MiniMaxM3DecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + ), + prefix=f"{prefix}.layers", + ) + + self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + + # EAGLE3 is not yet compatible with pipeline parallel + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): + hidden_states, residual = layer(positions, hidden_states, residual) + self._maybe_add_hidden_state( + aux_hidden_states, idx + 1, hidden_states, residual + ) + + hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # Checkpoint experts use w1=gate, w2=down, w3=up. + return 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, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # q/k/v_proj -> fused qkv_proj; gate_proj/up_proj -> fused gate_up_proj + # (dense MLP and shared expert). On sparse layers the indexer + # index_q/index_k_proj fold into the same fused qkv_proj + # (MinimaxM3QKVParallelLinearWithIndexer); these entries simply never match on + # dense layers, whose checkpoints have no index_*_proj weights. Leading + # dots keep `q_proj`/`k_proj` from matching `index_q_proj`/`index_k_proj` + # (preceded by `_`, not `.`). + 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), + ] + + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = self.get_expert_mapping() + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + # The MTP module is not modeled yet. + if "mtp." in name: + continue + + # The checkpoint stores block scales as ``weight_scale_inv``; the + # ModelOpt MXFP8 layers expose them as ``weight_scale``. + if "weight_scale_inv" in name: + name = name.replace("weight_scale_inv", "weight_scale") + + 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.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + 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 is_pp_missing_parameter(name, self): + continue + 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: + if name.endswith(".bias") and name not in params_dict: + continue + remapped = maybe_remap_kv_scale_name(name, params_dict) + if remapped is None: + continue + name = remapped + if is_pp_missing_parameter(name, self): + continue + # Modules not modeled yet (e.g. attention) are skipped until + # they are ported. + if name not in params_dict: + continue + 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 + + +class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3): + """MiniMax M3 (sparse/dense backbone) for causal language modeling.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + self.model = MiniMaxM3Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(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, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + +@MULTIMODAL_REGISTRY.register_processor( + MiniMaxM3VLMultiModalProcessor, + info=MiniMaxM3VLProcessingInfo, + dummy_inputs=MiniMaxM3VLDummyInputsBuilder, +) +class MiniMaxM3SparseForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsEagle3 +): + """Top-level (VL) entry point for MiniMax M3. + + The vision tower is not modeled yet; this wrapper routes the text + backbone by constructing ``MiniMaxM3SparseForCausalLM`` from the nested + ``text_config`` and delegating generation to it. + """ + + # The vision tower runs replicated per rank under ``--mm-encoder-tp-mode + # data``; ``run_dp_sharded_mrope_vision_model`` shards the work across + # ranks (see ``_process_image_input`` / ``_process_video_input``). + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "multi_modal_projector.": "vision_tower.multi_modal_projector.", + "patch_merge_mlp.": "vision_tower.patch_merge_mlp.", + }, + orig_to_new_substr={ + ".mlp.fc1.": ".fc1.", + ".mlp.fc2.": ".fc2.", + }, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return MiniMaxM3VLProcessingInfo.IMAGE_TOKEN + if modality == "video": + return MiniMaxM3VLProcessingInfo.VIDEO_TOKEN + raise ValueError(f"Unsupported modality: {modality!r}") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.config = config + self.quant_config = vllm_config.quant_config + self.multimodal_config = vllm_config.model_config.multimodal_config + assert self.multimodal_config is not None + self.use_data_parallel = self.multimodal_config.mm_encoder_tp_mode == "data" + + text_hidden_size = getattr(config.text_config, "hidden_size", None) + assert text_hidden_size is not None, "text_config.hidden_size is required" + projector_hidden_size = getattr(config, "projector_hidden_size", None) + + with self._mark_tower_model(vllm_config, {"image", "video"}): + vision_config = config.vision_config + self.vision_tower = MiniMaxVLVisionModel( + config=PretrainedConfig.from_dict(vision_config), + text_hidden_size=text_hidden_size, + projector_hidden_size=projector_hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["MiniMaxM3SparseForCausalLM"], + ) + + # Expose language model / lm_head for EAGLE3 spec decode. + @property + def model(self) -> nn.Module: + return self.language_model.model + + @property + def lm_head(self) -> nn.Module: + return self.language_model.lm_head + + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: + pixel_values = kwargs.pop("pixel_values", None) + image_grid_thw = kwargs.pop("image_grid_thw", None) + if pixel_values is None: + return None + return {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw} + + def _parse_and_validate_video_input(self, **kwargs: object) -> dict | None: + pixel_values_videos = kwargs.pop("pixel_values_videos", None) + video_grid_thw = kwargs.pop("video_grid_thw", None) + if pixel_values_videos is None: + return None + return { + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": video_grid_thw, + } + + def _process_image_input(self, image_input: dict) -> tuple[torch.Tensor, ...]: + pixel_values: torch.Tensor = image_input["pixel_values"].type( + self.vision_tower.dtype + ) + grid_thw: torch.Tensor = image_input["image_grid_thw"] + assert grid_thw.ndim == 2 + + if self.use_data_parallel: + # Already returns a per-item tuple of embeddings. + return run_dp_sharded_mrope_vision_model( + self.vision_tower, + pixel_values, + grid_thw.tolist(), + rope_type="rope_3d", + ) + + image_embeds = self.vision_tower( + pixel_values=pixel_values, + grid_thw=grid_thw.tolist(), + ) + + # Split the concatenated output into one tensor per image item. + merge_size = self.vision_tower.spatial_merge_size + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist() + return image_embeds.split(sizes) + + def _process_video_input(self, video_input: dict) -> tuple[torch.Tensor, ...]: + pixel_values: torch.Tensor = video_input["pixel_values_videos"].type( + self.vision_tower.dtype + ) + grid_thw: torch.Tensor = video_input["video_grid_thw"] + assert grid_thw.ndim == 2 + + if self.use_data_parallel: + # Already returns a per-item tuple of embeddings. + return run_dp_sharded_mrope_vision_model( + self.vision_tower, + pixel_values, + grid_thw.tolist(), + rope_type="rope_3d", + ) + + video_embeds = self.vision_tower( + pixel_values=pixel_values, + grid_thw=grid_thw.tolist(), + ) + + # Split the concatenated output into one tensor per video item. + merge_size = self.vision_tower.spatial_merge_size + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist() + return video_embeds.split(sizes) + + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> dict[str, dict]: + mm_input_by_modality: dict[str, dict] = {} + for input_key in kwargs: + if input_key == "pixel_values" and "image" not in mm_input_by_modality: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is not None: + mm_input_by_modality["image"] = image_input + if ( + input_key == "pixel_values_videos" + and "video" not in mm_input_by_modality + ): + video_input = self._parse_and_validate_video_input(**kwargs) + if video_input is not None: + mm_input_by_modality["video"] = video_input + return mm_input_by_modality + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) + if not mm_input_by_modality: + return [] + + multimodal_embeddings: list[torch.Tensor] = [] + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "image": + image_embeddings = self._process_image_input(multimodal_input) + multimodal_embeddings.extend(image_embeddings) + if modality == "video": + video_embeddings = self._process_video_input(multimodal_input) + multimodal_embeddings.extend(video_embeddings) + + return tuple(multimodal_embeddings) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.language_model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.language_model.get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/models/minimax_m3/nvidia/mtp.py b/vllm/models/minimax_m3/nvidia/mtp.py new file mode 100644 index 000000000000..e2c7f8821d96 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/mtp.py @@ -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 diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py new file mode 100644 index 000000000000..6ab59f8c4b57 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -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 diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 5d301b8201ed..bb3b67524722 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -92,6 +92,10 @@ "minimax_m2_reasoning_parser", "MiniMaxM2AppendThinkReasoningParser", ), + "minimax_m3": ( + "minimax_m3_reasoning_parser", + "MiniMaxM3ReasoningParser", + ), "mistral": ( "mistral_reasoning_parser", "MistralReasoningParser", diff --git a/vllm/reasoning/minimax_m3_reasoning_parser.py b/vllm/reasoning/minimax_m3_reasoning_parser.py new file mode 100644 index 000000000000..ec75ce78bfbb --- /dev/null +++ b/vllm/reasoning/minimax_m3_reasoning_parser.py @@ -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: + + reasoning textassistant 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 ```` again. + """ + + @property + def start_token(self) -> str: + return "" + + @property + def end_token(self) -> str: + return "" + + 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 diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index a6a931d5b2c7..6a70510e6ff7 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -126,6 +126,10 @@ "minimax_m2_tool_parser", "MinimaxM2ToolParser", ), + "minimax_m3": ( + "minimax_m3_tool_parser", + "MinimaxM3ToolParser", + ), "minimax": ( "minimax_tool_parser", "MinimaxToolParser", diff --git a/vllm/tool_parsers/minimax_m3_tool_parser.py b/vllm/tool_parsers/minimax_m3_tool_parser.py new file mode 100644 index 000000000000..a8628448c448 --- /dev/null +++ b/vllm/tool_parsers/minimax_m3_tool_parser.py @@ -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 ```` tags in one wrapper, + and represents nested arguments with parameter-name XML tags. + """ + + rust_parser_name = "MinimaxM3ToolParser" + tool_call_start_token = "]<]minimax[>[" diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 04a296551ddd..21b5e7494d7b 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -103,6 +103,8 @@ def __getitem__(self, key): medusa="MedusaConfig", mellum="MellumConfig", midashenglm="MiDashengLMConfig", + minimax_m3_vl="MiniMaxM3Config", + minimax_m3_mtp="MiniMaxM3MTPConfig", moondream3="Moondream3Config", eagle="EAGLEConfig", speculators="SpeculatorsConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index e91f89b2d09a..021eb2ea419e 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -53,6 +53,9 @@ "MedusaConfig": "vllm.transformers_utils.configs.medusa", "MellumConfig": "vllm.transformers_utils.configs.mellum", "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm", + "MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3", + "MiniMaxM3MTPConfig": "vllm.transformers_utils.configs.minimax_m3", + "MiniMaxM3TextConfig": "vllm.transformers_utils.configs.minimax_m3", "MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator", "Moondream3Config": "vllm.transformers_utils.configs.moondream3", "Moondream3TextConfig": "vllm.transformers_utils.configs.moondream3", @@ -124,6 +127,9 @@ "MedusaConfig", "MellumConfig", "MiDashengLMConfig", + "MiniMaxM3Config", + "MiniMaxM3MTPConfig", + "MiniMaxM3TextConfig", "MLPSpeculatorConfig", "Moondream3Config", "Moondream3TextConfig", diff --git a/vllm/transformers_utils/configs/minimax_m3.py b/vllm/transformers_utils/configs/minimax_m3.py new file mode 100644 index 000000000000..c340dda85a6c --- /dev/null +++ b/vllm/transformers_utils/configs/minimax_m3.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +from transformers import PretrainedConfig + + +class MiniMaxM3TextConfig(PretrainedConfig): + """Config for the MiniMax M3 text backbone (MiniMaxM3SparseForCausalLM). + + Defaults mirror the ``text_config`` of the MiniMax-M3-preview checkpoint. + """ + + model_type = "minimax_m3_text" + architectures = ["MiniMaxM3SparseForCausalLM"] + + def __init__( + self, + vocab_size: int = 200064, + hidden_size: int = 6144, + intermediate_size: int = 3072, + dense_intermediate_size: int = 12288, + shared_intermediate_size: int = 3072, + num_hidden_layers: int = 60, + num_attention_heads: int = 64, + num_key_value_heads: int = 4, + head_dim: int = 128, + max_position_embeddings: int = 524288, + rms_norm_eps: float = 1e-6, + use_gemma_norm: bool = True, + attention_output_gate: bool = False, + rope_theta: float = 5000000, + rotary_dim: int = 64, + partial_rotary_factor: float = 0.5, + hidden_act: str = "swigluoai", + swiglu_alpha: float = 1.702, + # SwiGLU-OAI uses the (up + 1) bias, i.e. beta=1.0 (matches the + # reference: gate * sigmoid(gate * alpha) * (up + 1)). The checkpoint + # config omits swiglu_beta, so this default must stay 1.0. + swiglu_beta: float = 1.0, + swiglu_limit: float = 7.0, + use_qk_norm: bool = True, + qk_norm_type: str = "per_head", + num_local_experts: int = 128, + num_experts_per_tok: int = 4, + n_shared_experts: int = 1, + scoring_func: str = "sigmoid", + use_routing_bias: bool = True, + routed_scaling_factor: float = 2.0, + num_mtp_modules: int = 1, + moe_layer_freq: list[int] | None = None, + sparse_attention_config: dict[str, Any] | None = None, + tie_word_embeddings: bool = False, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.dense_intermediate_size = dense_intermediate_size + self.shared_intermediate_size = shared_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.use_gemma_norm = use_gemma_norm + self.attention_output_gate = attention_output_gate + self.rope_theta = rope_theta + self.rotary_dim = rotary_dim + self.partial_rotary_factor = partial_rotary_factor + self.hidden_act = hidden_act + self.swiglu_alpha = swiglu_alpha + self.swiglu_beta = swiglu_beta + self.swiglu_limit = swiglu_limit + self.use_qk_norm = use_qk_norm + self.qk_norm_type = qk_norm_type + self.num_local_experts = num_local_experts + self.num_experts_per_tok = num_experts_per_tok + self.n_shared_experts = n_shared_experts + self.scoring_func = scoring_func + self.use_routing_bias = use_routing_bias + self.routed_scaling_factor = routed_scaling_factor + self.num_mtp_modules = num_mtp_modules + # First 3 layers are dense; the remaining 57 are sparse MoE. + self.moe_layer_freq = ( + moe_layer_freq if moe_layer_freq is not None else [0] * 3 + [1] * 57 + ) + self.sparse_attention_config = ( + sparse_attention_config + if sparse_attention_config is not None + else { + "use_sparse_attention": True, + "sparse_index_dim": 128, + "sparse_num_index_heads": 4, + "sparse_topk_blocks": 16, + "sparse_block_size": 128, + "sparse_disable_index_value": [0] * 3 + [1] * 57, + "sparse_score_type": "max", + "sparse_init_block": 0, + "sparse_local_block": 1, + "sparse_attention_freq": [0] * 3 + [1] * 57, + } + ) + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + +class MiniMaxM3MTPConfig(MiniMaxM3TextConfig): + """Config for a standalone MiniMax M3 MTP (multi-token prediction) head. + + The MTP transformer layer is structurally a single MiniMax M3 decoder + layer, so this reuses the text backbone schema. Standalone MTP checkpoints + use ``model_type='minimax_m3_mtp'`` and a single hidden layer. + """ + + model_type = "minimax_m3_mtp" + architectures = ["MiniMaxM3MTP"] + + def __init__(self, num_hidden_layers: int = 1, **kwargs): + super().__init__(num_hidden_layers=num_hidden_layers, **kwargs) + + +class MiniMaxM3Config(PretrainedConfig): + """Top-level MiniMax M3 (VL) config. + + Holds the text backbone as ``text_config`` so that + ``config.get_text_config()`` extracts the MiniMaxM3SparseForCausalLM + backbone. Vision components are kept as a raw dict passthrough and are + not modeled here. + """ + + model_type = "minimax_m3_vl" + + def __init__( + self, + text_config: dict | MiniMaxM3TextConfig | None = None, + vision_config: dict | None = None, + **kwargs, + ): + if text_config is None: + text_config = MiniMaxM3TextConfig() + elif isinstance(text_config, dict): + text_config = MiniMaxM3TextConfig(**text_config) + self.text_config = text_config + self.vision_config = vision_config + + self.hidden_size = text_config.hidden_size + + super().__init__(**kwargs) diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index a64be9618923..e4ece0a41972 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -31,6 +31,9 @@ "MiMoOmniProcessor", "MiniCPMOProcessor", "MiniCPMVProcessor", + "MiniMaxM3VLImageProcessor", + "MiniMaxM3VLVideoProcessor", + "MiniMaxVLProcessor", "MistralCommonPixtralProcessor", "MistralCommonVoxtralProcessor", "NanoNemotronVLProcessor", @@ -64,6 +67,9 @@ "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni", "MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo", "MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv", + "MiniMaxM3VLImageProcessor": "vllm.transformers_utils.processors.minimax_m3", + "MiniMaxM3VLVideoProcessor": "vllm.transformers_utils.processors.minimax_m3", + "MiniMaxVLProcessor": "vllm.transformers_utils.processors.minimax_m3", "MistralCommonPixtralProcessor": "vllm.transformers_utils.processors.pixtral", "MistralCommonVoxtralProcessor": "vllm.transformers_utils.processors.voxtral", "Moondream3Processor": "vllm.transformers_utils.processors.moondream3", diff --git a/vllm/transformers_utils/processors/minimax_m3.py b/vllm/transformers_utils/processors/minimax_m3.py new file mode 100644 index 000000000000..13dbce5368f8 --- /dev/null +++ b/vllm/transformers_utils/processors/minimax_m3.py @@ -0,0 +1,736 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax M3 VL HuggingFace-compatible Processor / ImageProcessor / +VideoProcessor, vendored into vLLM so the model loads without +``--trust-remote-code`` (the released checkpoint only ships these classes as +remote code via ``auto_map``). + +Adapted verbatim from the ``MiniMaxAI/Minimax-M3-preview`` repository files +``image_processor.py``, ``video_processor.py`` and ``processing_minimax.py`` +(revision ``db01c0fe``). Both image and video processors use Qwen-style +``smart_resize`` (bound by total pixels). The original async frame-sampling +helpers are intentionally omitted: vLLM performs its own frame loading and +feeds decoded frames to the processor. +""" + +import math + +import regex as re +import torch +from torchvision.transforms import InterpolationMode +from transformers import AutoTokenizer, BatchFeature +from transformers.image_processing_utils_fast import ( + BaseImageProcessorFast, + group_images_by_shape, + reorder_images, +) +from transformers.image_utils import PILImageResampling, SizeDict +from transformers.processing_utils import ( + ImagesKwargs, + ProcessingKwargs, + ProcessorMixin, + Unpack, + VideosKwargs, +) +from transformers.utils import TensorType +from transformers.video_processing_utils import BaseVideoProcessor +from transformers.video_utils import group_videos_by_shape, reorder_videos + +# Maximum allowed aspect ratio before smart_resize rejects the input. +MAX_RATIO = 200 + +# Fixed (non-configurable) bounds for the long-side resize logic, per the +# MiniMax-M3 size spec. ``min_short_side_pixel`` is the floor the short edge is +# enlarged to; ``*_MAX_TOTAL_PIXELS`` is the hard area cap that, once exceeded, +# aborts processing instead of downscaling. +MIN_SHORT_SIDE_PIXEL = 112 +IMAGE_MAX_TOTAL_PIXELS = 12_845_056 # 3584 ** 2 (width * height) +VIDEO_MAX_TOTAL_PIXELS = 301_056_000 # width * height * frames + + +def round_by_factor(number: int | float, factor: int) -> int: + return round(number / factor) * factor + + +def ceil_by_factor(number: int | float, factor: int) -> int: + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int | float, factor: int) -> int: + return math.floor(number / factor) * factor + + +def _smart_resize_by_long_side( + height: int, + width: int, + factor: int, + max_long_side_pixel: int, + min_short_side_pixel: int, + max_total_pixels: int | None, +) -> tuple[int, int]: + """Long-side based resize (MiniMax-M3 size spec). + + (a) if the long side exceeds ``max_long_side_pixel`` → shrink so the long + side equals ``max_long_side_pixel``; + (b) else if the short side is below ``min_short_side_pixel`` → enlarge so the + short side equals ``min_short_side_pixel``; + (c) if the resulting area still exceeds ``max_total_pixels`` → raise. + + (a) and (b) are mutually exclusive (they branch on the *original* long side). + Both sides are then rounded to a multiple of ``factor``. For videos the + ``max_total_pixels`` cap is volumetric (width * height * frames) and is + enforced by the caller, so pass ``max_total_pixels=None`` here. + """ + long_side = max(height, width) + short_side = min(height, width) + + scaled_height: float = height + scaled_width: float = width + if long_side > max_long_side_pixel: + beta = max_long_side_pixel / long_side + scaled_height = height * beta + scaled_width = width * beta + elif short_side < min_short_side_pixel: + beta = min_short_side_pixel / short_side + scaled_height = height * beta + scaled_width = width * beta + + h_bar = max(factor, round_by_factor(scaled_height, factor)) + w_bar = max(factor, round_by_factor(scaled_width, factor)) + + if max_total_pixels is not None and h_bar * w_bar > max_total_pixels: + raise ValueError( + f"image area {h_bar * w_bar} exceeds max_total_pixels " + f"{max_total_pixels} after resizing" + ) + return h_bar, w_bar + + +def smart_resize( + height: int, + width: int, + factor: int = 28, + min_pixels: int = 4 * 28 * 28, + max_pixels: int = 451584, + max_long_side_pixel: int | None = None, + min_short_side_pixel: int = MIN_SHORT_SIDE_PIXEL, + max_total_pixels: int | None = None, +) -> tuple[int, int]: + """Rescale (height, width) so each side is a multiple of ``factor``. + + When ``max_long_side_pixel`` is set, use the MiniMax-M3 long-side resize + spec (see :func:`_smart_resize_by_long_side`). Otherwise fall back to the + Qwen-VL area bound, keeping the total area within ``[min_pixels, max_pixels]``. + """ + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, " + f"got {max(height, width) / min(height, width)}" + ) + if max_long_side_pixel is not None: + return _smart_resize_by_long_side( + height, + width, + factor=factor, + max_long_side_pixel=max_long_side_pixel, + min_short_side_pixel=min_short_side_pixel, + max_total_pixels=max_total_pixels, + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +class MiniMaxM3VLImageProcessorKwargs(ImagesKwargs, total=False): # type: ignore[call-arg] + patch_size: int + temporal_patch_size: int + merge_size: int + max_pixels: int + max_long_side_pixel: int + + +class MiniMaxM3VLImageProcessor(BaseImageProcessorFast): + do_resize = True + resample = PILImageResampling.BICUBIC + # required by base-class validation, not used as the resize bound + size = {"height": 672, "width": 672} + default_to_square = False + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + image_mean = [0.48145466, 0.4578275, 0.40821073] + image_std = [0.26862954, 0.26130258, 0.27577711] + do_convert_rgb = True + patch_size = 14 + temporal_patch_size = 2 + merge_size = 2 + max_pixels = 451584 # 672 * 672 + # Long-side resize spec (opt-in via ``max_long_side_pixel``). The latter two + # are fixed per the spec and are not exposed as configurable kwargs. + max_long_side_pixel = None + min_short_side_pixel = MIN_SHORT_SIDE_PIXEL + max_total_pixels = IMAGE_MAX_TOTAL_PIXELS + valid_kwargs = MiniMaxM3VLImageProcessorKwargs + model_input_names = ["pixel_values", "image_grid_thw"] + + def __init__(self, **kwargs: Unpack[MiniMaxM3VLImageProcessorKwargs]): + super().__init__(**kwargs) + + def preprocess( + self, images, **kwargs: Unpack[MiniMaxM3VLImageProcessorKwargs] + ) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _preprocess( + self, + images: list[torch.Tensor], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: "float | list[float] | None", + image_std: "float | list[float] | None", + patch_size: int, + temporal_patch_size: int, + merge_size: int, + max_pixels: int, + max_long_side_pixel: "int | None", + disable_grouping: "bool | None", + return_tensors: "str | TensorType | None", + **kwargs, + ) -> BatchFeature: + grouped_images, grouped_images_index = group_images_by_shape( + images, disable_grouping=disable_grouping + ) + resized_images_grouped = {} + factor = patch_size * merge_size + for shape, stacked_images in grouped_images.items(): + height, width = stacked_images.shape[-2:] + if do_resize: + resized_height, resized_width = smart_resize( + height, + width, + factor=factor, + max_pixels=max_pixels, + max_long_side_pixel=max_long_side_pixel, + min_short_side_pixel=self.min_short_side_pixel, + max_total_pixels=self.max_total_pixels, + ) + stacked_images = self.resize( + stacked_images, + size=SizeDict(height=resized_height, width=resized_width), + resample=resample, + ) + resized_images_grouped[shape] = stacked_images + + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + processed_grids = {} + + for shape, stacked_images in grouped_images.items(): + resized_height, resized_width = stacked_images.shape[-2:] + + patches = self.rescale_and_normalize( + stacked_images, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + if patches.ndim == 4: + patches = patches.unsqueeze(1) + + if patches.shape[1] % temporal_patch_size != 0: + repeats = patches[:, -1:].repeat( + 1, + temporal_patch_size - (patches.shape[1] % temporal_patch_size), + 1, + 1, + 1, + ) + patches = torch.cat([patches, repeats], dim=1) + + batch_size, grid_t, channel = patches.shape[:3] + grid_t = grid_t // temporal_patch_size + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + patches = patches.view( + batch_size, + grid_t, + temporal_patch_size, + channel, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9) + + flatten_patches = patches.reshape( + batch_size, + grid_t * grid_h * grid_w, + channel * temporal_patch_size * patch_size * patch_size, + ) + + processed_images_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_images = reorder_images( + processed_images_grouped, grouped_images_index + ) + processed_grids = reorder_images(processed_grids, grouped_images_index) + + pixel_values = torch.cat(processed_images, dim=0) + image_grid_thw = torch.tensor(processed_grids, dtype=torch.long) + + return BatchFeature( + data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + tensor_type=return_tensors, + ) + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None): + images_kwargs = images_kwargs or {} + patch_size = images_kwargs.get("patch_size", self.patch_size) + merge_size = images_kwargs.get("merge_size", self.merge_size) + max_pixels = images_kwargs.get("max_pixels", self.max_pixels) + max_long_side_pixel = images_kwargs.get( + "max_long_side_pixel", self.max_long_side_pixel + ) + + resized_height, resized_width = smart_resize( + height, + width, + factor=patch_size * merge_size, + max_pixels=max_pixels, + max_long_side_pixel=max_long_side_pixel, + min_short_side_pixel=self.min_short_side_pixel, + max_total_pixels=self.max_total_pixels, + ) + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + return grid_h * grid_w + + +class MiniMaxM3VLVideoProcessorKwargs(VideosKwargs, total=False): # type: ignore[call-arg] + patch_size: int + temporal_patch_size: int + merge_size: int + min_pixels: int + max_pixels: int + max_long_side_pixel: int + total_pixels: int + min_frames: int + max_frames: int + fps: "float | int" + + +class MiniMaxM3VLVideoProcessor(BaseVideoProcessor): + do_resize = True + resample = PILImageResampling.BICUBIC + size = {"height": 672, "width": 672} + default_to_square = False + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + image_mean = [0.48145466, 0.4578275, 0.40821073] + image_std = [0.26862954, 0.26130258, 0.27577711] + do_convert_rgb = True + do_sample_frames = False + patch_size = 14 + temporal_patch_size = 2 + merge_size = 2 + min_pixels = 4 * 28 * 28 + max_pixels = 768 * 28 * 28 # 602,112 + total_pixels = int(64000 * 28 * 28 * 0.9) # ~45M, ~64k tokens budget + # Long-side resize spec (opt-in via ``max_long_side_pixel``). The video + # ``max_total_pixels`` cap is volumetric (width * height * frames) and is + # enforced in ``_preprocess`` once the frame count is known. + max_long_side_pixel = None + min_short_side_pixel = MIN_SHORT_SIDE_PIXEL + max_total_pixels = VIDEO_MAX_TOTAL_PIXELS + fps = 1.0 + min_frames = 4 + max_frames = 768 + valid_kwargs = MiniMaxM3VLVideoProcessorKwargs + model_input_names = ["pixel_values_videos", "video_grid_thw"] + + def __init__(self, **kwargs: Unpack[MiniMaxM3VLVideoProcessorKwargs]): + super().__init__(**kwargs) + + def _preprocess( + self, + videos: list[torch.Tensor], + do_convert_rgb: bool, + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: "float | list[float] | None", + image_std: "float | list[float] | None", + patch_size: int, + temporal_patch_size: int, + merge_size: int, + min_pixels: int, + max_pixels: int, + max_long_side_pixel: "int | None" = None, + return_tensors: "str | TensorType | None" = None, + **kwargs, + ) -> BatchFeature: + grouped_videos, grouped_videos_index = group_videos_by_shape(videos) + resized_videos_grouped = {} + factor = patch_size * merge_size + for shape, stacked_videos in grouped_videos.items(): + batch_size, num_frames, channels, height, width = stacked_videos.shape + resized_height, resized_width = height, width + if do_resize: + resized_height, resized_width = smart_resize( + height, + width, + factor=factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + max_long_side_pixel=max_long_side_pixel, + min_short_side_pixel=self.min_short_side_pixel, + # Per-frame raise disabled; the video cap is volumetric and + # is enforced below once num_frames is known. + max_total_pixels=None, + ) + if ( + max_long_side_pixel is not None + and resized_height * resized_width * num_frames + > self.max_total_pixels + ): + raise ValueError( + f"video area {resized_height * resized_width * num_frames} " + f"(width * height * frames) exceeds max_total_pixels " + f"{self.max_total_pixels} after resizing" + ) + stacked_videos = stacked_videos.view( + batch_size * num_frames, channels, height, width + ) + stacked_videos = self.resize( + stacked_videos, + size=SizeDict(height=resized_height, width=resized_width), + resample=resample, + ) + stacked_videos = stacked_videos.view( + batch_size, + num_frames, + channels, + resized_height, + resized_width, + ) + resized_videos_grouped[shape] = stacked_videos + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) + + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) + processed_videos_grouped = {} + processed_grids = {} + for shape, stacked_videos in grouped_videos.items(): + resized_height, resized_width = stacked_videos.shape[-2:] + patches = self.rescale_and_normalize( + stacked_videos, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + + if pad := -patches.shape[1] % temporal_patch_size: + repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1) + patches = torch.cat([patches, repeats], dim=1) + + batch_size, grid_t, channels = patches.shape[:3] + grid_t = grid_t // temporal_patch_size + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + patches = patches.view( + batch_size, + grid_t, + temporal_patch_size, + channels, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9) + flatten_patches = patches.reshape( + batch_size, + grid_t * grid_h * grid_w, + channels * temporal_patch_size * patch_size * patch_size, + ) + + processed_videos_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_videos = reorder_videos( + processed_videos_grouped, grouped_videos_index + ) + processed_grids = reorder_videos(processed_grids, grouped_videos_index) + pixel_values_videos = torch.cat(processed_videos, dim=0) + video_grid_thw = torch.tensor(processed_grids, dtype=torch.long) + + return BatchFeature( + data={ + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": video_grid_thw, + }, + tensor_type=return_tensors, + ) + + +class MiniMaxVLProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg] + _defaults = { + "videos_kwargs": { + "do_resize": False, + "return_metadata": True, + }, + } + + +class MiniMaxVLProcessor(ProcessorMixin): + IMAGE_TOKEN = "]<]image[>[" + VIDEO_TOKEN = "]<]video[>[" + VISION_START_TOKEN = "]<]start of image[>[" + VISION_END_TOKEN = "]<]end of image[>[" + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + # Bypass ProcessorMixin's dynamic module lookup, which breaks in + # transformers >= 5.9 when image_processor_class is a string: the + # register() API now stores classes as {"pil": cls} dicts in + # _extra_content, but get_possibly_dynamic_module() still calls + # .__name__ on the raw value, crashing with AttributeError on dicts. + tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, **kwargs + ) + image_processor = MiniMaxM3VLImageProcessor.from_pretrained( + pretrained_model_name_or_path, **kwargs + ) + video_processor = MiniMaxM3VLVideoProcessor.from_pretrained( + pretrained_model_name_or_path, **kwargs + ) + return cls( + image_processor=image_processor, + tokenizer=tokenizer, + video_processor=video_processor, + ) + + def __init__( + self, image_processor=None, tokenizer=None, video_processor=None, **kwargs + ): + self.image_token_id = tokenizer.convert_tokens_to_ids(self.IMAGE_TOKEN) + self.video_token_id = tokenizer.convert_tokens_to_ids(self.VIDEO_TOKEN) + super().__init__(image_processor, tokenizer, video_processor) + # Video expansion also uses image start/end tokens. Separate video + # start/end tokens exist in the tokenizer, but the original MiniMax + # serving path did not use them; keep that behavior for compatibility. + self.vision_start_token_id = tokenizer.convert_tokens_to_ids( + self.VISION_START_TOKEN + ) + self.vision_end_token_id = tokenizer.convert_tokens_to_ids( + self.VISION_END_TOKEN + ) + + def _prune_video_tokens( + self, + input_text: str, + video_segments: list[int], + video_token: str, + ) -> str: + """Prune video tokens by temporal_patch_size (e.g., 2:1). + + Expects the prompt to carry exactly sum(video_segments) video tokens + — i.e. one token per *sampled* frame — then drops tokens. + """ + # If no videos or temporal_patch_size <= 1, no pruning needed + if not video_segments or self.video_processor.temporal_patch_size <= 1: + return input_text + + # Split while keeping delimiters + special_tokens = [video_token] + pattern = "|".join(map(re.escape, special_tokens)) + parts = re.split(f"({pattern})", input_text) + + def is_timestamp(text: str) -> bool: + """Check if text ends with timestamp format like ']<]0.0 seconds[>['""" + return ( + text.endswith("seconds[>[") + or text.endswith("seconds[>[ ") + or text.endswith("seconds [>[") + or text.endswith("seconds [>[ ") + ) + + def extract_timestamp(text: str) -> str: + """Extract timestamp text from the end, starting from ']<]'""" + start_index = text.rfind("]<]") + if start_index == -1: + raise ValueError(f"Failed to extract timestamp: {text}") + return text[start_index:] + + # Build new text with pruned video tokens + final_parts = [] + current_seg_idx = 0 # Which video segment we're in + frame_in_seg = 0 # Frame index within current segment + last_timestamp_len = 0 # Length of timestamp to potentially remove + + for part in parts: + if part == video_token: + if current_seg_idx < len(video_segments): + if frame_in_seg % self.video_processor.temporal_patch_size == 0: + # Keep this video token + final_parts.append(part) + frame_in_seg += 1 + if frame_in_seg >= video_segments[current_seg_idx]: + current_seg_idx += 1 + frame_in_seg = 0 + last_timestamp_len = 0 + else: + # Skip this video token + frame_in_seg += 1 + if frame_in_seg >= video_segments[current_seg_idx]: + current_seg_idx += 1 + frame_in_seg = 0 + # Remove the timestamp that was already appended + if last_timestamp_len > 0: + assert len(final_parts) > 0 + final_parts[-1] = final_parts[-1][:-last_timestamp_len] + last_timestamp_len = 0 + else: + # No more video segments, keep as is + final_parts.append(part) + last_timestamp_len = 0 + else: + # Text part + final_parts.append(part) + # Check if this text ends with a timestamp + if is_timestamp(part): + last_timestamp_len = len(extract_timestamp(part)) + else: + last_timestamp_len = 0 + + return "".join(final_parts) + + def __call__( + self, + images=None, + text=None, + videos=None, + **kwargs: Unpack[MiniMaxVLProcessorKwargs], + ) -> BatchFeature: + output_kwargs = self._merge_kwargs( + MiniMaxVLProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if images is not None: + images_kwargs = output_kwargs["images_kwargs"] + image_inputs = self.image_processor(images=images, **images_kwargs) + image_grid_thw = image_inputs["image_grid_thw"] + else: + image_inputs = {} + image_grid_thw = None + + if videos is not None: + videos_kwargs = output_kwargs["videos_kwargs"] + video_inputs = self.video_processor(videos=videos, **videos_kwargs) + video_grid_thw = video_inputs["video_grid_thw"] + if not kwargs.get("return_metadata"): + video_metadata = video_inputs.pop("video_metadata") + else: + video_metadata = video_inputs["video_metadata"] + else: + video_inputs = {} + video_grid_thw = None + + if not isinstance(text, list): + text = [text] + text = text.copy() + + # Expand image tokens + if image_grid_thw is not None: + merge_length = self.image_processor.merge_size**2 + placeholder = "]<]placeholder[>[" + index = 0 + for i in range(len(text)): + while self.IMAGE_TOKEN in text[i]: + num_tokens = image_grid_thw[index].prod() // merge_length + text[i] = text[i].replace( + self.IMAGE_TOKEN, + self.VISION_START_TOKEN + + placeholder * num_tokens + + self.VISION_END_TOKEN, + 1, + ) + index += 1 + text[i] = text[i].replace(placeholder, self.IMAGE_TOKEN) + + # Expand video tokens + if video_grid_thw is not None: + merge_length = self.image_processor.merge_size**2 + placeholder = "]<]placeholder[>[" + index = 0 + for i in range(len(text)): + while self.VIDEO_TOKEN in text[i]: + metadata = video_metadata[index] + grid_t = video_grid_thw[index][0] + frame_seqlen = video_grid_thw[index][1:].prod() // merge_length + + video_placeholder = "" + for frame_idx in range(grid_t): + if ( + metadata.fps is not None + and metadata.frames_indices is not None + ): + ts = ( + metadata.frames_indices[ + min( + frame_idx + * self.video_processor.temporal_patch_size, + len(metadata.frames_indices) - 1, + ) + ] + / metadata.fps + ) + video_placeholder += f"]<]{ts:.1f} seconds[>[" + video_placeholder += ( + self.VISION_START_TOKEN + + placeholder * frame_seqlen + + self.VISION_END_TOKEN + ) + + text[i] = text[i].replace(self.VIDEO_TOKEN, video_placeholder, 1) + index += 1 + text[i] = text[i].replace(placeholder, self.VIDEO_TOKEN) + + # Tokenize + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + + return BatchFeature( + data={**text_inputs, **image_inputs, **video_inputs}, + tensor_type=return_tensors, + ) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 73e1cce56d56..486aa7e4054b 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -336,9 +336,24 @@ class FlashInferBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - # Note: Not sure for all platforms, but on Blackwell, - # only support a page size of 16, 32, 64. - return [16, 32, 64] + # Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA + # on Blackwell); advertise them only when usable so selection never + # picks a large kernel block we cannot serve. + use_large_pages = False + vllm_config = get_current_vllm_config_or_none() + if vllm_config is not None and vllm_config.model_config is not None: + pc = vllm_config.parallel_config + mc = vllm_config.model_config + num_qo_heads = mc.get_num_attention_heads(pc) + num_kv_heads = mc.get_num_kv_heads(pc) + use_large_pages = ( + num_kv_heads > 0 + and num_qo_heads // num_kv_heads > 1 + and can_use_trtllm_attention(num_qo_heads, num_kv_heads) + ) + if not use_large_pages: + return [16, 32, 64] + return [16, 32, 64, 128, 256, 512, 1024] @staticmethod def get_name() -> str: @@ -647,6 +662,12 @@ def __init__( # if TRTLLM attention kernel is not used when building attn metadata can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads) + # Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by + # get_supported_kernel_block_sizes). + assert self.page_size <= 64 or ( + can_use_trtllm and self.num_qo_heads // self.num_kv_heads > 1 + ), f"Unexpected FlashInfer page size {self.page_size} without trtllm-gen GQA" + if ( can_use_trtllm and not vllm_config.attention_config.disable_flashinfer_q_quantization @@ -917,6 +938,10 @@ def build( # - Decode (FI native or TRTLLM) use_cascade = common_prefix_len > 0 uses_spec_reorder = self.reorder_batch_threshold > 1 + # Page sizes >= 128 must use trtllm-gen; force it for prefill too. + prefill_force_trtllm = ( + True if page_size >= 128 else self.attention_config.use_trtllm_attention + ) prefill_use_trtllm = use_trtllm_attention( self.num_qo_heads, self.num_kv_heads, @@ -926,7 +951,7 @@ def build( self.cache_dtype, self.q_data_type, is_prefill=True, - force_use_trtllm=self.attention_config.use_trtllm_attention, + force_use_trtllm=prefill_force_trtllm, has_sinks=self.has_sinks, has_spec=uses_spec_reorder, ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 2cd2bb5b9860..bdaa752a6032 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -91,6 +91,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend" ) FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" + MINIMAX_M3_SPARSE = ( + "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend" + ) NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend" FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" ROCM_AITER_UNIFIED_ATTN = ( diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index e11798ce6b00..d4f2c1007b09 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.config import ( CUDAGraphMode, VllmConfig, @@ -250,6 +251,13 @@ def __init__( DeepseekV4ROCMAiterMLASparseMetadata, DeepseekV4ROCMAiterSparseSWAMetadata, ) + + # MiniMax-M3 sparse (lightning-indexer) attention. The multi-step + # drafting machinery is shared code at num_speculative_tokens>1. + # this just opts the metadata into the ROCm allowlist. + from vllm.models.minimax_m3.common.sparse_attention import ( + MiniMaxM3SparseMetadata, + ) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerMetadata, ) @@ -265,6 +273,7 @@ def __init__( DeepseekV4ROCMAiterMLASparseMetadata, DeepseekV4ROCMAiterSparseSWAMetadata, DeepseekV32IndexerMetadata, + MiniMaxM3SparseMetadata, ] # ROCM_AITER_FA is an optional backend # We check is_enabled() here to avoid importing the backend module during @@ -457,8 +466,11 @@ def propose( batch_size = common_attn_metadata.batch_size() if self.method in ("eagle3", "dflash"): + model = self.model + if isinstance(model, BreakableCUDAGraphWrapper): + model = model.unwrap() assert isinstance( - self.model, + model, ( Eagle3LlamaForCausalLM, Eagle3DeepseekV2ForCausalLM, diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 87a2aac9d4ca..d9c041ba0b89 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -322,7 +322,7 @@ def __getitem__(self, idx: int) -> "BlockTable": return self.block_tables[idx] -@triton.jit +@triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) def _compute_slot_mapping_kernel( num_tokens, max_num_tokens,