From 0663ebc783e6dd2b361817620849336a7edea0dd Mon Sep 17 00:00:00 2001 From: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Date: Sat, 11 Jul 2026 11:11:06 +0800 Subject: [PATCH] [minimax-m3] Split 4/4: model + VL + glue + function-call + fp8 quant + generic infra (#28715) Co-authored-by: Xinyuan Tong Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com> --- .../csrc/gemm/per_token_group_quant_8bit.cuh | 339 ++-- python/sglang/jit_kernel/moe_fused_gate.py | 2 + .../jit_kernel/per_token_group_quant_8bit.py | 22 +- python/sglang/srt/arg_groups/overrides.py | 117 +- python/sglang/srt/configs/__init__.py | 2 + python/sglang/srt/configs/minimax_vl.py | 60 + python/sglang/srt/configs/model_config.py | 10 +- .../srt/entrypoints/openai/serving_chat.py | 7 + python/sglang/srt/entrypoints/warmup.py | 34 + python/sglang/srt/environ.py | 9 + .../srt/function_call/function_call_parser.py | 2 + python/sglang/srt/function_call/minimax_m3.py | 542 ++++++ .../layers/attention/attention_registry.py | 13 + .../attention/minimax_sparse_backend.py | 602 ++++++ .../layers/deep_gemm_wrapper/entrypoint.py | 26 + python/sglang/srt/layers/layernorm.py | 44 +- .../sglang/srt/layers/moe/ep_moe/kernels.py | 520 ++++- .../srt/layers/moe/moe_runner/deep_gemm.py | 167 +- .../layers/moe/token_dispatcher/standard.py | 9 +- python/sglang/srt/layers/moe/topk.py | 50 +- python/sglang/srt/layers/quantization/fp8.py | 296 ++- .../srt/layers/quantization/fp8_kernel.py | 308 ++- .../srt/layers/quantization/fp8_utils.py | 89 +- .../quantization/mxfp8_block_convert.py | 72 + python/sglang/srt/layers/radix_attention.py | 124 +- .../srt/managers/multi_tokenizer_mixin.py | 7 +- python/sglang/srt/managers/scheduler.py | 9 +- .../srt/model_executor/forward_batch_info.py | 4 +- python/sglang/srt/models/minimax_m3.py | 1690 +++++++++++++++++ python/sglang/srt/models/minimax_m3_vl.py | 372 ++++ python/sglang/srt/models/minimax_vl_common.py | 882 +++++++++ .../multimodal/processors/minimax_m3_vl.py | 283 +++ python/sglang/srt/parser/reasoning_parser.py | 62 +- .../sglang/srt/parser/template_detection.py | 6 + python/sglang/srt/utils/common.py | 15 + .../srt/utils/hf_transformers/common.py | 2 + .../bench_per_token_group_quant_8bit.py | 52 +- .../benchmark/bench_post_reorder_deepgemm.py | 79 + .../test_minimax_fused_qkv_index_gemm.py | 69 + .../jit/minimax/test_minimax_quant_scatter.py | 84 + .../jit/test_per_token_group_quant_8bit.py | 105 + .../jit/test_post_reorder_deepgemm.py | 90 + .../function_call/test_minimax_m3_detector.py | 539 ++++++ .../unit/parser/test_reasoning_parser.py | 90 + .../unit/parser/test_template_manager.py | 46 + 45 files changed, 7477 insertions(+), 475 deletions(-) create mode 100644 python/sglang/srt/configs/minimax_vl.py create mode 100644 python/sglang/srt/function_call/minimax_m3.py create mode 100644 python/sglang/srt/layers/attention/minimax_sparse_backend.py create mode 100644 python/sglang/srt/layers/quantization/mxfp8_block_convert.py create mode 100644 python/sglang/srt/models/minimax_m3.py create mode 100644 python/sglang/srt/models/minimax_m3_vl.py create mode 100644 python/sglang/srt/models/minimax_vl_common.py create mode 100644 python/sglang/srt/multimodal/processors/minimax_m3_vl.py create mode 100644 test/registered/jit/benchmark/bench_post_reorder_deepgemm.py create mode 100644 test/registered/jit/minimax/test_minimax_fused_qkv_index_gemm.py create mode 100644 test/registered/jit/minimax/test_minimax_quant_scatter.py create mode 100644 test/registered/jit/test_post_reorder_deepgemm.py create mode 100644 test/registered/unit/function_call/test_minimax_m3_detector.py diff --git a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh b/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh index 0afd7f412..02e721de7 100644 --- a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh +++ b/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh @@ -1,149 +1,202 @@ #include #include -#include -#include #include -#include +#include #include #include #include +#include + #include #include +#include namespace { -constexpr int kThreadsPerGroup = 16; +using deepseek_v4::fp8::cast_to_ue8m0; +using deepseek_v4::fp8::inv_scale_ue8m0; +using deepseek_v4::fp8::pack_fp8; -#ifdef USE_ROCM -// AMD implementation: HIP warps are 64-wide and require an explicit sub-group -// width, so the CUDA 32-bit-mask shuffle reduction below does not compile on -// gfx942. Delegate to the portable warp-reduce primitive, which emits -// __shfl_xor with an explicit kThreadsPerGroup sub-group width. -__device__ __forceinline__ float GroupReduceMax(float val, const int /*tid*/) { - return device::warp::reduce_max(val); -} -#else -__device__ __forceinline__ float GroupReduceMax(float val, const int tid) { - unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff; - val = fmaxf(val, __shfl_xor_sync(mask, val, 8)); - val = fmaxf(val, __shfl_xor_sync(mask, val, 4)); - val = fmaxf(val, __shfl_xor_sync(mask, val, 2)); - val = fmaxf(val, __shfl_xor_sync(mask, val, 1)); - return val; -} -#endif +// Optimized per-token-group quant to FP8-e4m3 (or int8), with optional +// column-major UE8M0 (int32-packed) scale. Memory-bound rewrite of the AOT +// `per_token_group_quant_8bit` (sgl-kernel): the previous JIT clone read the +// input TWICE and used 16 threads/group with only group/8 active. This version: +// * loads each group once into registers (single 128-bit load per thread), +// * uses exactly kGroupSize/kVec threads per group (no idle lanes), +// * sub-warp shuffle-reduces the absmax over those lanes (no shared memory), +// * launches warp-aligned, ~256-thread blocks for high occupancy / latency +// hiding (the AOT kernel's 1-warp blocks left HBM ~80% idle at prefill). +// The UE8M0 path reuses the dsv4 cast_to_ue8m0/inv_scale_ue8m0 primitives and is +// byte-identical to `sgl_per_token_group_quant_8bit_v2` (both ceil-round the +// scale and store the biased exponent byte). ROCm portability comes from the +// portable `warp::reduce_max` used directly in the kernel +// (gfx942-safe; no separate GroupReduceMax helper needed). -template -using scale_packed_t_t = std::conditional_t; +template +using scale_packed_t_t = std::conditional_t; +template +using scale_element_t_t = std::conditional_t; -template -using scale_element_t_t = std::conditional_t; +struct PerTokenGroupQuantParams { + const void* __restrict__ input; + void* __restrict__ output_q; + void* __restrict__ output_s; + int64_t num_groups; // total groups = num_tokens * num_groups_per_row + int num_groups_per_row; // hidden / group_size + int groups_per_block; // groups handled by one CTA + int scale_stride; // output_s.stride(1), in scale_packed_t elements + float eps; + float min_8bit; + float max_8bit; +}; -template -__global__ void per_token_group_quant_8bit_kernel( - const T* __restrict__ input, - DST_DTYPE* __restrict__ output_q, - scale_packed_t_t* __restrict__ output_s, - const int group_size, - const int num_groups, - const int groups_per_block, - const float eps, - const float min_8bit, - const float max_8bit, - const int num_groups_per_row = 0, - const int scale_stride = 0) { +// kGroupSize columns per group; kThreadsPerGroup threads cover one group, each +// issuing kNumVec coalesced 128-bit loads. Lane loads are interleaved +// (chunk v at element (v*kThreadsPerGroup + lane)*kVec) so consecutive lanes hit +// consecutive 16B addresses -> fully coalesced even for kNumVec > 1. +template +__global__ __launch_bounds__(256, 8) void per_token_group_quant_8bit_kernel(const PerTokenGroupQuantParams params) { using namespace device; namespace math = device::math; - (void)num_groups; + constexpr uint32_t kVec = 16u / sizeof(T); // 8 for bf16/fp16 + constexpr uint32_t kElemsPerThread = kGroupSize / kThreadsPerGroup; + constexpr uint32_t kNumVec = kElemsPerThread / kVec; // 128-bit loads per thread + static_assert(kGroupSize % (kThreadsPerGroup * kVec) == 0, "bad tiling"); + static_assert( + kThreadsPerGroup >= 1 && kThreadsPerGroup <= 32 && (kThreadsPerGroup & (kThreadsPerGroup - 1)) == 0, + "threads-per-group must be a pow2 <= 32"); - const int local_group_id = static_cast(threadIdx.x / kThreadsPerGroup); - const int lane_id = threadIdx.x % kThreadsPerGroup; + using InVec = AlignedVector; + using scale_packed_t = scale_packed_t_t; + using scale_element_t = scale_element_t_t; - const int64_t block_group_id = blockIdx.x * groups_per_block; - const int64_t global_group_id = block_group_id + local_group_id; - const int64_t block_group_offset = global_group_id * group_size; + const int local_group = threadIdx.x / kThreadsPerGroup; + const int lane = threadIdx.x % kThreadsPerGroup; + const int64_t global_group = static_cast(blockIdx.x) * params.groups_per_block + local_group; - float local_absmax = eps; - - using scale_packed_t = scale_packed_t_t; - using scale_element_t = scale_element_t_t; - static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0); - - const T* group_input = input + block_group_offset; - DST_DTYPE* group_output = static_cast(output_q) + block_group_offset; - scale_element_t* scale_output = nullptr; - - if constexpr (kIsColumnMajor) { - constexpr int kElemsPerPack = static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); - const int row_idx = global_group_id / num_groups_per_row; - const int col_idx_unpacked = global_group_id % num_groups_per_row; - const int col_idx = col_idx_unpacked / kElemsPerPack; - const int pack_idx = col_idx_unpacked % kElemsPerPack; - scale_output = reinterpret_cast(output_s) + - (col_idx * scale_stride * kElemsPerPack + row_idx * kElemsPerPack + pack_idx); - } else { - static_assert(!kScaleUE8M0); - scale_output = output_s + global_group_id; + PDLWaitPrimary(); + if (global_group >= params.num_groups) { + PDLTriggerSecondary(); + return; } - constexpr uint32_t kVecSize = 16 / sizeof(T); - using vec_t = AlignedVector; - const auto gmem_in = tile::Memory::thread(); - - const int32_t num_vec_elems = group_size / kVecSize; - - for (int32_t i = lane_id; i < num_vec_elems; i += kThreadsPerGroup) { - const vec_t input_vec = gmem_in.load(group_input, i); + const T* gin = static_cast(params.input) + global_group * kGroupSize; + DST* gout = static_cast(params.output_q) + global_group * kGroupSize; + // Load kNumVec interleaved 128-bit chunks into registers. + float vals[kElemsPerThread]; + float local_absmax = params.eps; #pragma unroll - for (uint32_t j = 0; j < kVecSize; ++j) { - const float val = static_cast(input_vec[j]); + for (uint32_t v = 0; v < kNumVec; ++v) { + InVec in_vec; + in_vec.load(gin + (v * kThreadsPerGroup + lane) * kVec, 0); +#pragma unroll + for (uint32_t j = 0; j < kVec; ++j) { + const float val = static_cast(in_vec[j]); + vals[v * kVec + j] = val; local_absmax = math::max(local_absmax, math::abs(val)); } } - - local_absmax = GroupReduceMax(local_absmax, lane_id); - - float y_s = local_absmax / max_8bit; - if constexpr (kScaleUE8M0) { - y_s = exp2f(ceilf(log2f(math::max(y_s, 1e-10f)))); + if constexpr (kThreadsPerGroup > 1) { + local_absmax = warp::reduce_max(local_absmax); } - scale_element_t y_s_quant; - if constexpr (kScaleUE8M0) { - y_s_quant = static_cast(((int)log2f(y_s)) + 127); + // Scale (byte-identical to sgl_per_token_group_quant_8bit_v2). + const float kMaxInv = 1.0f / params.max_8bit; + float inv_scale; // multiply input by this to quantize + scale_element_t scale_store; + if constexpr (kUE8M0) { + const int32_t exp = cast_to_ue8m0(local_absmax * kMaxInv); + inv_scale = inv_scale_ue8m0(exp); + scale_store = static_cast(exp); } else { - y_s_quant = y_s; + const float scale_inv = local_absmax * kMaxInv; // stored scale + inv_scale = params.max_8bit / local_absmax; // quant multiplier + scale_store = scale_inv; } - if (lane_id == 0) { - *scale_output = y_s_quant; - } - - for (int32_t i = lane_id; i < num_vec_elems; i += kThreadsPerGroup) { - const vec_t input_vec = gmem_in.load(group_input, i); - + // Quantize from registers and store kNumVec interleaved chunks. #pragma unroll - for (uint32_t j = 0; j < kVecSize; ++j) { - const float val = static_cast(input_vec[j]); - const float q_val = math::min(math::max(val / y_s, min_8bit), max_8bit); - group_output[i * kVecSize + j] = DST_DTYPE(q_val); + for (uint32_t v = 0; v < kNumVec; ++v) { + DST* o = gout + (v * kThreadsPerGroup + lane) * kVec; + if constexpr (std::is_same_v) { + AlignedVector out_vec; +#pragma unroll + for (uint32_t j = 0; j < kVec / 2; ++j) { + out_vec[j] = pack_fp8(vals[v * kVec + 2 * j] * inv_scale, vals[v * kVec + 2 * j + 1] * inv_scale); + } + out_vec.store(o, 0); + } else { + AlignedVector out_vec; +#pragma unroll + for (uint32_t j = 0; j < kVec; ++j) { + const float q = math::min(math::max(vals[v * kVec + j] * inv_scale, params.min_8bit), params.max_8bit); + out_vec[j] = static_cast(q); + } + out_vec.store(o, 0); } } + + // One scale write per group (the leading lane). + if (lane == 0) { + scale_element_t* scale_out; + if constexpr (kColMajor) { + constexpr int kPack = static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); + const int row = static_cast(global_group / params.num_groups_per_row); // token + const int col_u = static_cast(global_group % params.num_groups_per_row); // group in row + const int col = col_u / kPack; + const int pack = col_u % kPack; + scale_out = reinterpret_cast(params.output_s) + + (static_cast(col) * params.scale_stride * kPack + static_cast(row) * kPack + pack); + } else { + static_assert(!kUE8M0, "non-column-major UE8M0 is unsupported"); + scale_out = static_cast(params.output_s) + global_group; + } + *scale_out = scale_store; + } + + PDLTriggerSecondary(); } -inline int compute_groups_per_block(int64_t num_groups) { - if (num_groups % 16 == 0) return 16; - if (num_groups % 8 == 0) return 8; - if (num_groups % 4 == 0) return 4; - if (num_groups % 2 == 0) return 2; - return 1; +// Threads cooperating on one group. Heuristic: ~8 elems/thread for small groups +// (one 128-bit load) and 16 elems/thread for group=128 (two loads), capped at 8 +// threads/group to keep the sub-warp reduction (and register pressure) small. +template +constexpr int threads_per_group() { + constexpr int kVec = 16 / sizeof(T); // 8 for bf16/fp16 + int tpg = static_cast(kGroupSize) / (2 * kVec); // ~16 elems/thread (2 vecs) + if (tpg > 8) tpg = 8; + if (tpg < 1) tpg = 1; + return tpg; } -template +constexpr int kBlockThreads = 256; + +template +inline int pick_groups_per_block() { + int gpb = kBlockThreads / threads_per_group(); + if (gpb < 1) gpb = 1; + return gpb; +} + +template +void launch_quant(const PerTokenGroupQuantParams& base, int64_t num_groups, int groups_per_block, DLDevice device) { + using namespace host; + constexpr int kThreadsPerGroup = threads_per_group(); + const int num_threads = groups_per_block * kThreadsPerGroup; + const int64_t num_blocks = (num_groups + groups_per_block - 1) / groups_per_block; + PerTokenGroupQuantParams params = base; + params.groups_per_block = groups_per_block; + constexpr auto kernel = + per_token_group_quant_8bit_kernel; + LaunchKernel(static_cast(num_blocks), static_cast(num_threads), device) + .enable_pdl(kUsePDL)(kernel, params); +} + +template void per_token_group_quant_8bit( tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, @@ -154,6 +207,9 @@ void per_token_group_quant_8bit( double max_8bit, bool scale_ue8m0) { using namespace host; + static_assert( + kGroupSize == 16 || kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, + "group_size template arg must be 16/32/64/128"); auto device = SymbolicDevice{}; auto M = SymbolicSize{"num_tokens"}; @@ -163,66 +219,43 @@ void per_token_group_quant_8bit( TensorMatcher({M, K}).with_dtype().with_device(device).verify(input); TensorMatcher({M, K}).with_dtype().with_device(device).verify(output_q); - const auto num_tokens = M.unwrap(); - const auto hidden_dim = K.unwrap(); + RuntimeCheck(group_size == kGroupSize, "group_size does not match compiled template"); - const int64_t num_groups_per_row = hidden_dim / group_size; + const int64_t num_tokens = M.unwrap(); + const int64_t hidden_dim = K.unwrap(); + const int64_t num_groups_per_row = hidden_dim / kGroupSize; const int64_t num_groups = num_tokens * num_groups_per_row; + if (num_groups == 0) return; - const int groups_per_block = compute_groups_per_block(num_groups); - const int num_blocks = num_groups / groups_per_block; - const int num_threads = groups_per_block * kThreadsPerGroup; const bool is_column_major = output_s.stride(0) < output_s.stride(1); - const int scale_stride = output_s.stride(1); + const int scale_stride = static_cast(output_s.stride(1)); - const float feps = static_cast(eps); - const float fmin8 = static_cast(min_8bit); - const float fmax8 = static_cast(max_8bit); + PerTokenGroupQuantParams base{}; + base.input = input.data_ptr(); + base.output_q = output_q.data_ptr(); + base.output_s = output_s.data_ptr(); + base.num_groups = num_groups; + base.num_groups_per_row = static_cast(num_groups_per_row); + base.scale_stride = scale_stride; + base.eps = static_cast(eps); + base.min_8bit = static_cast(min_8bit); + base.max_8bit = static_cast(max_8bit); + const auto dev = input.device(); + const int gpb = pick_groups_per_block(); + + // Runtime selection between compile-time-instantiated scale-layout variants + // (the group size itself is a template arg, supplied from Python). if (is_column_major) { if (scale_ue8m0) { - LaunchKernel(num_blocks, num_threads, input.device())( - per_token_group_quant_8bit_kernel, - static_cast(input.data_ptr()), - static_cast(output_q.data_ptr()), - static_cast(output_s.data_ptr()), - static_cast(group_size), - static_cast(num_groups), - static_cast(groups_per_block), - feps, - fmin8, - fmax8, - static_cast(num_groups_per_row), - scale_stride); + launch_quant(base, num_groups, gpb, dev); } else { - LaunchKernel(num_blocks, num_threads, input.device())( - per_token_group_quant_8bit_kernel, - static_cast(input.data_ptr()), - static_cast(output_q.data_ptr()), - static_cast(output_s.data_ptr()), - static_cast(group_size), - static_cast(num_groups), - static_cast(groups_per_block), - feps, - fmin8, - fmax8, - static_cast(num_groups_per_row), - scale_stride); + launch_quant(base, num_groups, gpb, dev); } } else { - LaunchKernel(num_blocks, num_threads, input.device())( - per_token_group_quant_8bit_kernel, - static_cast(input.data_ptr()), - static_cast(output_q.data_ptr()), - static_cast(output_s.data_ptr()), - static_cast(group_size), - static_cast(num_groups), - static_cast(groups_per_block), - feps, - fmin8, - fmax8, - 0, - 0); + RuntimeCheck(!scale_ue8m0, "row-major UE8M0 unsupported"); + launch_quant(base, num_groups, gpb, dev); } } + } // namespace diff --git a/python/sglang/jit_kernel/moe_fused_gate.py b/python/sglang/jit_kernel/moe_fused_gate.py index df203bc5e..6cdfd9c3a 100644 --- a/python/sglang/jit_kernel/moe_fused_gate.py +++ b/python/sglang/jit_kernel/moe_fused_gate.py @@ -287,6 +287,8 @@ def moe_fused_gate( assert bias.ndim == 1, "bias must be 1D" assert scores.size(1) == bias.size(0), "scores and bias must have same num_experts" assert topk > num_fused_shared_experts, "topk must be > num_fused_shared_experts" + if routed_scaling_factor is None: + routed_scaling_factor = 1.0 M, N = scores.shape K = topk diff --git a/python/sglang/jit_kernel/per_token_group_quant_8bit.py b/python/sglang/jit_kernel/per_token_group_quant_8bit.py index 6df31c520..f2077c802 100644 --- a/python/sglang/jit_kernel/per_token_group_quant_8bit.py +++ b/python/sglang/jit_kernel/per_token_group_quant_8bit.py @@ -4,7 +4,12 @@ from typing import TYPE_CHECKING import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.custom_op import register_custom_op @@ -16,17 +21,22 @@ from sglang.jit_kernel.utils import CPP_DTYPE_MAP as OUTPUT_DTYPE_MAP @cache_once def _jit_per_token_group_quant_8bit_module( - dtype: torch.dtype, output_type: torch.dtype + dtype: torch.dtype, output_type: torch.dtype, group_size: int ) -> Module: - input_args = make_cpp_args(dtype) + dtype_arg = make_cpp_args(dtype) + gs_arg = make_cpp_args(group_size) + pdl_arg = make_cpp_args(is_arch_support_pdl()) out_cpp = OUTPUT_DTYPE_MAP[output_type] return load_jit( "per_token_group_quant_8bit", + *dtype_arg, + *gs_arg, + *pdl_arg, cuda_files=["gemm/per_token_group_quant_8bit.cuh"], cuda_wrappers=[ ( "per_token_group_quant_8bit", - f"per_token_group_quant_8bit<{input_args}, {out_cpp}>", + f"per_token_group_quant_8bit<{dtype_arg}, {out_cpp}, {gs_arg}, {pdl_arg}>", ) ], ) @@ -59,7 +69,9 @@ def _per_token_group_quant_8bit_custom_op( fp8_max: The maximum value of the 8-bit data type. scale_ue8m0: Whether to use UE8M0 format for scales. """ - module = _jit_per_token_group_quant_8bit_module(input.dtype, output_q.dtype) + module = _jit_per_token_group_quant_8bit_module( + input.dtype, output_q.dtype, group_size + ) module.per_token_group_quant_8bit( input, output_q, diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 2795460dc..6a8ebceb6 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -34,6 +34,7 @@ import logging from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from sglang.srt.arg_groups.arg_utils import resolvable_fields +from sglang.srt.environ import envs from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.utils.common import ( cpu_has_amx_support, @@ -45,6 +46,7 @@ from sglang.srt.utils.common import ( is_cpu, is_cuda, is_flashinfer_available, + is_gfx95_supported, is_hip, is_musa, is_npu, @@ -460,6 +462,92 @@ def _minimax_m2_overrides(server_args: Any, hf_config: Any) -> dict: return {"enable_tf32_matmul": True} +@_register_for("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration") +def _minimax_m3_overrides(server_args: Any, hf_config: Any) -> dict: + + overrides: Dict[str, Any] = {} + + quant_method = get_quantization_config(hf_config) + quant_resolved = server_args.quantization + if ( + quant_resolved is None + and not server_args._quantization_explicitly_unset + and quant_method is not None + ): + overrides["quantization"] = quant_method + quant_resolved = quant_method + + if is_hip(): + if server_args.is_attention_backend_not_set(): + overrides["attention_backend"] = "triton" + if server_args.moe_runner_backend == "auto" and quant_resolved == "mxfp8": + overrides["moe_runner_backend"] = "triton" + if not envs.USE_ROCM_AITER_ROPE_BACKEND.is_set(): + envs.USE_ROCM_AITER_ROPE_BACKEND.set("0") + aiter_fusion_resolved = server_args.enable_aiter_allreduce_fusion + if ( + server_args.ep_size > 1 + and server_args.moe_a2a_backend == "none" + and aiter_fusion_resolved + ): + logger.warning( + "Disable --enable-aiter-allreduce-fusion for MiniMax-M3 " + "standard EP on ROCm because the deferred fused all-reduce " + "corrupts sparse MoE partial outputs." + ) + overrides["enable_aiter_allreduce_fusion"] = False + aiter_fusion_resolved = False + if not aiter_fusion_resolved: + overrides["disable_custom_all_reduce"] = True + elif is_sm100_supported(): + if server_args.is_attention_backend_not_set(): + overrides["attention_backend"] = "fa4" + page_resolved = server_args.page_size + if ( + page_resolved is None + and overrides.get("attention_backend", server_args.attention_backend) + == "fa4" + ): + overrides["page_size"] = 128 + page_resolved = 128 + if server_args.moe_runner_backend == "auto" and quant_resolved == "mxfp8": + overrides["moe_runner_backend"] = "deep_gemm" + logger.info( + "MiniMax-M3 on SM100: attention_backend=" + f"{overrides.get('attention_backend', server_args.attention_backend)}, page_size={page_resolved}, " + f"moe_runner_backend={overrides.get('moe_runner_backend', server_args.moe_runner_backend)}." + ) + elif is_sm90_supported(): + if server_args.is_attention_backend_not_set(): + overrides["attention_backend"] = "fa3" + page_resolved = server_args.page_size + if ( + page_resolved is None + and overrides.get("attention_backend", server_args.attention_backend) + == "fa3" + ): + overrides["page_size"] = 128 + page_resolved = 128 + logger.info( + "MiniMax-M3 on Hopper: attention_backend=" + f"{overrides.get('attention_backend', server_args.attention_backend)}, page_size={page_resolved} " + "(MSA is SM100-only; sparse attention runs on the Triton path)." + ) + + moe_runner_resolved = overrides.get( + "moe_runner_backend", server_args.moe_runner_backend + ) + if quant_resolved is None and moe_runner_resolved in ("auto", "deep_gemm"): + if moe_runner_resolved == "deep_gemm": + logger.warning( + "MiniMax-M3: the deep_gemm MoE runner produces corrupted output " + "on bf16 full weights; overriding --moe-runner-backend to 'triton'." + ) + overrides["moe_runner_backend"] = "triton" + + return overrides + + @_register_for( "Gemma2ForCausalLM", "Gemma3ForCausalLM", @@ -525,7 +613,6 @@ def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict: # use bf16 for mxfp4 triton kernels overrides["dtype"] = "bfloat16" if server_args.moe_runner_backend == "auto": - from sglang.srt.environ import envs if is_sm100_supported() and is_mxfp4_quant_format: overrides["moe_runner_backend"] = "flashinfer_mxfp4" @@ -832,7 +919,6 @@ def _qwen3_5_hybrid_overrides(server_args: Any, hf_config: Any) -> dict: @_register_for("Qwen3VLForConditionalGeneration") def _qwen3vl_overrides(server_args: Any, hf_config: Any) -> dict: - from sglang.srt.environ import envs if ( is_hip() @@ -1279,7 +1365,6 @@ def _deepseek_spec_moe_resolution(view: Any) -> dict: backends for the DeepSeek fp4 checkpoint. Reads the mid-resolution quantization (after _deepseek_moe_quant_resolution) and the pre-a2a ep_size, exactly like the legacy in-branch writes.""" - from sglang.srt.environ import envs hf_config = view.get_model_config().hf_config model_arch = hf_config.architectures[0] @@ -1364,7 +1449,6 @@ def _deepseek_v4_sm120_moe(view: Any) -> dict: @register_post_process def _sparse_head_overlap_disable(view: Any) -> dict: - from sglang.srt.environ import envs if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set(): logger.warning( @@ -1782,7 +1866,6 @@ def _attention_backend_dual_chunk(view: Any) -> dict: def _page_size_default(view: Any) -> dict: if view.page_size is not None: return {} - from sglang.srt.environ import envs # SHUFFLE 5D vectorized KV layout (aiter backend + pa_decode_gluon) # is tuned for and prefers page_size=64 — making it the default @@ -1845,19 +1928,25 @@ def _moe_runner_backend_quant_constraints(view: Any) -> dict: "flashinfer_trtllm_routed." ) if view.quantization == "mxfp8": - if moe_runner_backend == "auto": - moe_runner_backend = "flashinfer_trtllm" - elif moe_runner_backend not in [ + is_gfx95_mxfp8 = is_hip() and is_gfx95_supported() + allowed = [ "cutlass", + "deep_gemm", "flashinfer_trtllm", "flashinfer_trtllm_routed", - ]: + ] + if is_gfx95_mxfp8: + allowed.append("triton") + mxfp8_default = "triton" if is_gfx95_mxfp8 else "flashinfer_trtllm" + if moe_runner_backend == "auto": + moe_runner_backend = mxfp8_default + elif moe_runner_backend not in allowed: logger.warning( - "mxfp8 quantization supports only cutlass, flashinfer_trtllm, " - "or flashinfer_trtllm_routed backends. " - f"Overriding {moe_runner_backend!r}." + "mxfp8 quantization supports only %s backends. " "Overriding %r.", + ", ".join(allowed), + moe_runner_backend, ) - moe_runner_backend = "flashinfer_trtllm" + moe_runner_backend = mxfp8_default if ( moe_runner_backend == "auto" and view.quantization == "modelopt_fp4" @@ -1918,7 +2007,6 @@ def _a2a_fusion_adjustments(view: Any) -> dict: def _cutlass_moe_env_override(view: Any) -> dict: - from sglang.srt.environ import envs if envs.SGLANG_CUTLASS_MOE.get(): logger.warning( @@ -1940,7 +2028,6 @@ _A2A_EP_SPANNING_BACKENDS = frozenset( @register_post_process def _a2a_backend_overrides(view: Any) -> dict: - from sglang.srt.environ import envs moe_a2a_backend = view.moe_a2a_backend if view.enable_deepep_waterfill and moe_a2a_backend != "deepep": diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index aa05b73be..5d5d14d68 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -24,6 +24,7 @@ from sglang.srt.configs.lfm2_vl import Lfm2VlConfig from sglang.srt.configs.locate_anything import LocateAnythingConfig from sglang.srt.configs.longcat_flash import LongcatFlashConfig from sglang.srt.configs.minicpmv4_6 import MiniCPMV4_6Config, MiniCPMV4_6VisionConfig +from sglang.srt.configs.minimax_vl import MiniMaxM3VLConfig from sglang.srt.configs.nano_nemotron_vl import ( NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, @@ -82,6 +83,7 @@ __all__ = [ "JetNemotronConfig", "JetVLMConfig", "Step3p5Config", + "MiniMaxM3VLConfig", "Step3p7Config", "Qwen3ASRConfig", "UnlimitedVLConfig", diff --git a/python/sglang/srt/configs/minimax_vl.py b/python/sglang/srt/configs/minimax_vl.py new file mode 100644 index 000000000..37c9a1c8b --- /dev/null +++ b/python/sglang/srt/configs/minimax_vl.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +from transformers.configuration_utils import PretrainedConfig +from transformers.models.auto import CONFIG_MAPPING + + +def _coerce_sub_config( + sub_config: Optional[dict], default_model_type: str +) -> Optional[PretrainedConfig]: + """Convert a config dict to a ``PretrainedConfig``. + + Unknown ``model_type`` (e.g. M3's ``minimax_m2``, absent from + ``CONFIG_MAPPING``) falls back to ``PretrainedConfig`` so dict keys + still become real attributes. + """ + if not isinstance(sub_config, dict): + return sub_config + model_type = sub_config.get("model_type", default_model_type) + cls = CONFIG_MAPPING.get(model_type, PretrainedConfig) + return cls(**sub_config) + + +class MiniMaxVLBaseConfig(PretrainedConfig): + def __init__( + self, + vision_config: Optional[dict] = None, + text_config: Optional[dict] = None, + image_token_index: int = 200025, + video_token_index: int = 200026, + image_seq_length: int = 576, + process_image_mode: str = "dynamic_res", + projector_hidden_act: str = "gelu", + multimodal_projector_bias: bool = True, + vision_feature_layer: int = -1, + vision_feature_select_strategy: str = "full", + img_token_compression_config: Optional[dict] = None, + image_grid_pinpoints: Optional[str] = None, + **kwargs, + ): + self.vision_config = _coerce_sub_config(vision_config, "clip_vision_model") + self.text_config = _coerce_sub_config(text_config, "mixtral") + + self.image_token_index = image_token_index + self.video_token_index = video_token_index + self.image_seq_length = image_seq_length + self.process_image_mode = process_image_mode + self.projector_hidden_act = projector_hidden_act + self.multimodal_projector_bias = multimodal_projector_bias + self.vision_feature_layer = vision_feature_layer + self.vision_feature_select_strategy = vision_feature_select_strategy + self.img_token_compression_config = img_token_compression_config or {} + self.image_grid_pinpoints = image_grid_pinpoints + + super().__init__(**kwargs) + + +class MiniMaxM3VLConfig(MiniMaxVLBaseConfig): + model_type = "minimax_m3_vl" diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index e9b50cc26..465553fbb 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1726,6 +1726,8 @@ piecewise_cuda_graph_disabled_model_archs = [ # cleanly (vision encoder runs eagerly outside the graph via general_mm_embed_routine). multimodal_piecewise_cuda_graph_supported_model_archs = [ "Cohere2VisionForConditionalGeneration", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", ] if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): @@ -1786,14 +1788,6 @@ def is_piecewise_cuda_graph_disabled_model(model_architectures: List[str]): ) -# Multimodal archs whose LM-decoder prefill is validated under piecewise CUDA -# graph (capture wraps only the decoder; the image encoder runs eager). -multimodal_piecewise_cuda_graph_supported_archs = [ - "MiniMaxM3SparseForCausalLM", - "MiniMaxM3SparseForConditionalGeneration", -] - - def is_multimodal_piecewise_cuda_graph_supported(model_architectures: List[str]): """Whether a multimodal arch may keep prefill piecewise CUDA graph enabled.""" return any( diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index bd0607682..b14be748e 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -1791,6 +1791,13 @@ class OpenAIServingChat(OpenAIServingBase): if not self.reasoning_parser: return False + if self.reasoning_parser == "minimax-m3": + # M3 template prefills for thinking_mode=enabled, so it never + # appears in output and reasoning must be forced. Mirrors reasoning_parser.py. + return (request.chat_template_kwargs or {}).get( + "thinking_mode" + ) == "enabled" + if self.reasoning_parser == "hunyuan": # Hy3-preview template emits no when reasoning_effort is # "no_think" / "none" / unset; forcing reasoning would route all diff --git a/python/sglang/srt/entrypoints/warmup.py b/python/sglang/srt/entrypoints/warmup.py index 485d2dbf8..2608621f6 100644 --- a/python/sglang/srt/entrypoints/warmup.py +++ b/python/sglang/srt/entrypoints/warmup.py @@ -125,3 +125,37 @@ async def voice_chat(disaggregation_mode: str, tokenizer_manager: TokenizerManag generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST await tokenizer_manager.generate_request(generate_req_input, None).__anext__() + + +@warmup("prefill_shapes") +async def prefill_shapes(disaggregation_mode: str, tokenizer_manager: TokenizerManager): + """Warmup Triton kernels across a wide range of prefill seq_lens (up to 32K). + + Uses power-of-2 sizes plus intermediate points to cover the shape space + that fused_moe, attention extend, and other Triton kernels may encounter. + """ + page_size = 64 + sizes = set() + base = 64 + while base <= 32768: + sizes.add(base) + mid = base * 3 // 2 + mid = (mid + page_size - 1) // page_size * page_size + if mid <= 32768: + sizes.add(mid) + base *= 2 + sizes = sorted(sizes) + + for size in tqdm.tqdm(sizes, desc="Warmup prefill shapes (up to 32K)"): + generate_req_input = GenerateReqInput( + input_ids=(np.random.randint(2**16, size=[size])).tolist(), + sampling_params={ + "max_new_tokens": 1, + "temperature": 0.0, + }, + ) + if disaggregation_mode != "null": + generate_req_input.bootstrap_room = 0 + generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST + + await tokenizer_manager.generate_request(generate_req_input, None).__anext__() diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2b3ad5cc5..78cd951b1 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -520,6 +520,7 @@ class Envs: SGLANG_AITER_KV_CACHE_LAYOUT = EnvStr("nhd") SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False) SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False) + USE_ROCM_AITER_ROPE_BACKEND = EnvStr("0") SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(4096) # Enable dual-stream MoE (shared experts vs routed experts) on the # ROCm/AITER path. Requires GPU_MAX_HW_QUEUES>=5 to avoid HW-queue serialization. @@ -907,6 +908,14 @@ class Envs: SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK = EnvBool(False) SGLANG_OPT_USE_TOPK_V2 = EnvBool(True) + # Reroutes the generic fp8 per-token-group quant (every model, not just MiniMax) + # to the V1 JIT kernel. Off by default; V1 is byte-identical to V2. + SGLANG_OPT_USE_JIT_PER_TOKEN_GROUP_QUANT = EnvBool(False) + SGLANG_OPT_USE_BF16_ROUTER_GEMM = EnvBool(True) + SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE = EnvBool(False) + SGLANG_DISABLE_MSA = EnvBool(False) + SGLANG_OPT_USE_MSA_DECODE_UNDER_GRAPH = EnvBool(False) + # MiniMax-M3 sparse decode indexer: single JIT radix-select kernel replaces the 2-stage split-K Triton topk. SGLANG_OPT_USE_MINIMAX_DECODE_TOPK_RADIX = EnvBool(True) diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index b934eb5d1..62ac02715 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -33,6 +33,7 @@ from sglang.srt.function_call.llama32_detector import Llama32Detector from sglang.srt.function_call.mimo_detector import MiMoDetector from sglang.srt.function_call.minicpm5_detector import MiniCPM5Detector from sglang.srt.function_call.minimax_m2 import MinimaxM2Detector +from sglang.srt.function_call.minimax_m3 import MinimaxM3Detector from sglang.srt.function_call.mistral_detector import MistralDetector from sglang.srt.function_call.poolside_v1_detector import PoolsideV1Detector from sglang.srt.function_call.pythonic_detector import PythonicDetector @@ -82,6 +83,7 @@ class FunctionCallParser: "step3": Step3Detector, "step3p5": Qwen3CoderDetector, "minimax-m2": MinimaxM2Detector, + "minimax-m3": MinimaxM3Detector, "trinity": TrinityDetector, "interns1": InternlmDetector, "hermes": HermesDetector, diff --git a/python/sglang/srt/function_call/minimax_m3.py b/python/sglang/srt/function_call/minimax_m3.py new file mode 100644 index 000000000..062fc3ec6 --- /dev/null +++ b/python/sglang/srt/function_call/minimax_m3.py @@ -0,0 +1,542 @@ +import json +import logging +import re +from typing import Any, Dict, List, Optional, Tuple + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + ToolCallItem, + _GetInfoFunc, +) + +logger = logging.getLogger(__name__) + + +MINIMAX_NS_TOKEN = "]<]minimax[>[" +STRING_TYPES = {"string", "str"} +INTEGER_TYPES = {"integer", "int"} +NUMBER_TYPES = {"number", "float"} +BOOLEAN_TYPES = {"boolean", "bool"} +SCALAR_TYPES = STRING_TYPES | INTEGER_TYPES | NUMBER_TYPES | BOOLEAN_TYPES | {"null"} +CONTAINER_TYPES = {"object", "array"} +# Exact-case only; "none"/"nil" are valid string enum values, never null-coerced. +NULL_STRINGS = {"null", "Null", "NULL"} + + +class MinimaxM3Detector(BaseFormatDetector): + TOOL_CALL_START = MINIMAX_NS_TOKEN + "" + TOOL_CALL_END = MINIMAX_NS_TOKEN + "" + + INVOKE_PREFIX = MINIMAX_NS_TOKEN + '" + PARAM_START_PREFIX = MINIMAX_NS_TOKEN + "<" + TAG_SPACING_CHARS = " \t\r\n" + TAG_SPACING_RE = re.compile(re.escape(MINIMAX_NS_TOKEN) + r"[ \t\r\n]+(?=<)") + + def __init__(self): + super().__init__() + self._in_tool_call = False + self._current_function_name: Optional[str] = None + self._current_function_schema: Optional[Dict[str, Any]] = None + self._current_param_name: Optional[str] = None + self._current_param_schema: Optional[Dict[str, Any]] = None + self._current_param_buffer = "" + self._current_param_is_complex = False + self._current_string_started = False + self._is_first_param = True + + @classmethod + def _normalize_tag_spacing(cls, text: str) -> str: + return cls.TAG_SPACING_RE.sub(MINIMAX_NS_TOKEN, text) + + @classmethod + def _flushable_prefix_length(cls, text: str, token: str) -> int: + for length in range(min(len(token) - 1, len(text)), 0, -1): + if token.startswith(text[-length:]): + return len(text) - length + + namespace_start = text.rfind(MINIMAX_NS_TOKEN) + if namespace_start == -1: + return len(text) + + suffix = text[namespace_start + len(MINIMAX_NS_TOKEN) :] + if suffix and all(ch in cls.TAG_SPACING_CHARS for ch in suffix): + return namespace_start + return len(text) + + def has_tool_call(self, text: str) -> bool: + return self.TOOL_CALL_START in self._normalize_tag_spacing(text) + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + original_text = text + try: + text = self._normalize_tag_spacing(text) + normal_text, calls = self._extract(text, tools) + return StreamingParseResult(normal_text=normal_text, calls=calls) + except Exception as exc: + logger.warning( + "invalid MiniMax M3 tool call returned as content: %s", + exc, + exc_info=True, + ) + return StreamingParseResult(normal_text=original_text, calls=[]) + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError + + def _extract(self, text: str, tools: List[Tool]) -> Tuple[str, List[ToolCallItem]]: + normal_parts: List[str] = [] + calls: List[ToolCallItem] = [] + cursor = 0 + while True: + start = text.find(self.TOOL_CALL_START, cursor) + if start == -1: + normal_parts.append(text[cursor:]) + break + + normal_parts.append(text[cursor:start]) + + end = text.find(self.TOOL_CALL_END, start) + if end == -1: + normal_parts.append(text[start:]) + break + + block = text[start + len(self.TOOL_CALL_START) : end] + cursor = end + len(self.TOOL_CALL_END) + calls.extend(self._parse_block(block, tools)) + + return "".join(normal_parts), calls + + def _parse_block(self, block: str, tools: List[Tool]) -> List[ToolCallItem]: + results: List[ToolCallItem] = [] + cursor = 0 + while True: + start = block.find(self.INVOKE_PREFIX, cursor) + end = block.find(self.INVOKE_SUFFIX, start) + if start == -1 or end == -1: + break + + invoke_str = block[start:end] + cursor = end + len(self.INVOKE_SUFFIX) + + name_end = invoke_str.find('">', len(self.INVOKE_PREFIX)) + if name_end == -1: + continue + + func_name = invoke_str[len(self.INVOKE_PREFIX) : name_end] + body = invoke_str[name_end + len('">') :] + params = self._parse_parameter( + body, self._get_function_parameters_schema(func_name, tools) + ) + action = {"name": func_name, "arguments": params} + try: + parsed_calls = self.parse_base_json(action, tools) + for call in parsed_calls: + call.tool_index = len(results) + results.append(call) + except Exception: + logger.warning("invalid tool call for %s dropped", func_name) + return results + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + self._buffer += new_text + self._buffer = self._normalize_tag_spacing(self._buffer) + normal_text = "" + calls: List[ToolCallItem] = [] + + while True: + if not self._in_tool_call: + start = self._buffer.find(self.TOOL_CALL_START) + if start == -1: + flush_length = self._flushable_prefix_length( + self._buffer, self.TOOL_CALL_START + ) + normal_text += self._buffer[:flush_length] + self._buffer = self._buffer[flush_length:] + break + + normal_text += self._buffer[:start] + self._buffer = self._buffer[start + len(self.TOOL_CALL_START) :] + self._in_tool_call = True + self._current_function_name = None + continue + + if self._current_function_name is None: + if self._consume_tool_call_end(): + continue + if not self._consume_invoke_start(tools, calls): + break + continue + + if self._current_param_name is None: + if self._consume_invoke_end(calls): + continue + if not self._consume_param_start(calls): + break + continue + + if self._current_param_is_complex: + if not self._consume_complex_param(calls): + break + elif not self._consume_scalar_param(calls): + break + + return StreamingParseResult(normal_text=normal_text, calls=calls) + + def _consume_tool_call_end(self) -> bool: + start = self._buffer.find(self.TOOL_CALL_END) + invoke_start = self._buffer.find(self.INVOKE_PREFIX) + if start == -1 or (invoke_start != -1 and invoke_start < start): + return False + + self._buffer = self._buffer[start + len(self.TOOL_CALL_END) :] + self._in_tool_call = False + return True + + def _consume_invoke_start( + self, tools: List[Tool], calls: List[ToolCallItem] + ) -> bool: + start = self._buffer.find(self.INVOKE_PREFIX) + if start == -1: + return False + + name_start = start + len(self.INVOKE_PREFIX) + name_end = self._buffer.find('">', name_start) + if name_end == -1: + return False + + function_name = self._buffer[name_start:name_end] + self._buffer = self._buffer[name_end + len('">') :] + self._current_function_name = function_name + self._current_function_schema = self._get_function_parameters_schema( + function_name, tools + ) + self._is_first_param = True + + if self.current_tool_id == -1: + self.current_tool_id = 0 + self._append_stream_call(calls, "", name=function_name) + self._append_stream_call(calls, "{") + return True + + def _consume_invoke_end(self, calls: List[ToolCallItem]) -> bool: + start = self._buffer.find(self.INVOKE_SUFFIX) + param_start = self._buffer.find(self.PARAM_START_PREFIX) + if start == -1 or (param_start != -1 and param_start < start): + return False + + self._buffer = self._buffer[start + len(self.INVOKE_SUFFIX) :] + self._append_stream_call(calls, "}") + self.current_tool_id += 1 + self._current_function_name = None + self._current_function_schema = None + return True + + def _consume_param_start(self, calls: List[ToolCallItem]) -> bool: + start = self._buffer.find(self.PARAM_START_PREFIX) + if start == -1: + return False + + gt = self._buffer.find(">", start + len(self.PARAM_START_PREFIX)) + if gt == -1: + return False + + tag = self._buffer[start + len(self.PARAM_START_PREFIX) : gt].strip() + self._buffer = self._buffer[gt + 1 :] + self._current_param_name = tag + self._current_param_schema = self._get_child_schema( + self._current_function_schema, tag + ) + self._current_param_buffer = "" + self._current_param_is_complex = self._schema_has_type( + self._current_param_schema, ("object", "array") + ) or self._buffer.startswith(self.PARAM_START_PREFIX) + self._current_string_started = False + + prefix = "{}{}: ".format( + "" if self._is_first_param else ", ", + json.dumps(tag, ensure_ascii=False), + ) + self._append_stream_call(calls, prefix) + self._is_first_param = False + return True + + def _consume_complex_param(self, calls: List[ToolCallItem]) -> bool: + end_token = self._parameter_end_token(self._current_param_name) + end = self._buffer.find(end_token) + if end == -1: + flush_length = self._flushable_prefix_length(self._buffer, end_token) + self._current_param_buffer += self._buffer[:flush_length] + self._buffer = self._buffer[flush_length:] + return False + + self._current_param_buffer += self._buffer[:end] + value = self._parse_parameter( + self._current_param_buffer, self._current_param_schema + ) + self._append_stream_call(calls, json.dumps(value, ensure_ascii=False)) + self._buffer = self._buffer[end + len(end_token) :] + self._clear_current_param() + return True + + def _consume_scalar_param(self, calls: List[ToolCallItem]) -> bool: + end_token = self._parameter_end_token(self._current_param_name) + end = self._buffer.find(end_token) + if end == -1: + flush_length = self._flushable_prefix_length(self._buffer, end_token) + text = self._buffer[:flush_length] + self._buffer = self._buffer[flush_length:] + self._stream_scalar_text(text, calls) + return False + + self._stream_scalar_text(self._buffer[:end], calls) + if self._schema_has_type(self._current_param_schema, tuple(STRING_TYPES)): + if not self._current_string_started: + self._append_stream_call(calls, '"') + self._append_stream_call(calls, '"') + else: + value = self._convert_leaf_value( + self._current_param_buffer, self._current_param_schema + ) + self._append_stream_call(calls, json.dumps(value, ensure_ascii=False)) + + self._buffer = self._buffer[end + len(end_token) :] + self._clear_current_param() + return True + + def _stream_scalar_text(self, text: str, calls: List[ToolCallItem]) -> None: + if not text: + return + + if self._schema_has_type(self._current_param_schema, tuple(STRING_TYPES)): + escaped = json.dumps(text, ensure_ascii=False)[1:-1] + if self._current_string_started: + self._append_stream_call(calls, escaped) + else: + self._append_stream_call(calls, '"' + escaped) + self._current_string_started = True + else: + self._current_param_buffer += text + + def _clear_current_param(self) -> None: + self._current_param_name = None + self._current_param_schema = None + self._current_param_buffer = "" + self._current_param_is_complex = False + self._current_string_started = False + + def _append_stream_call( + self, calls: List[ToolCallItem], parameters: str, *, name: Optional[str] = None + ) -> None: + if ( + name is None + and calls + and calls[-1].tool_index == self.current_tool_id + and calls[-1].name is None + ): + calls[-1].parameters += parameters + else: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=name, + parameters=parameters, + ) + ) + + def _parameter_end_token(self, tag: Optional[str]) -> str: + return MINIMAX_NS_TOKEN + f"" + + def _get_function_parameters_schema( + self, function_name: str, tools: List[Tool] + ) -> Optional[Dict[str, Any]]: + for tool in tools: + if tool.function.name == function_name: + parameters = tool.function.parameters + if isinstance(parameters, dict): + return parameters + break + return None + + def _get_child_schema( + self, parent_schema: Any, child_tag: str, parent_value: Any = None + ) -> Optional[Dict]: + if not isinstance(parent_schema, dict): + return None + + if self._schema_has_type(parent_schema, ("array",)) and child_tag == "item": + return self._get_array_item_schema(parent_schema, parent_value) + + properties = parent_schema.get("properties") + if isinstance(properties, dict) and child_tag in properties: + child_schema = properties[child_tag] + return child_schema if isinstance(child_schema, dict) else None + + additional_properties = parent_schema.get("additionalProperties") + if isinstance(additional_properties, dict): + return additional_properties + + return None + + def _get_array_item_schema( + self, array_schema: Dict[str, Any], array_value: Any + ) -> Optional[Dict]: + item_index = len(array_value) if isinstance(array_value, list) else 0 + + prefix_items = array_schema.get("prefixItems") + if isinstance(prefix_items, list) and item_index < len(prefix_items): + item_schema = prefix_items[item_index] + return item_schema if isinstance(item_schema, dict) else None + + additional_items = array_schema.get("additionalItems") + if isinstance(additional_items, dict): + return additional_items + + items = array_schema.get("items") + if isinstance(items, dict): + return items + + if isinstance(prefix_items, list) and prefix_items: + item_schema = prefix_items[-1] + return item_schema if isinstance(item_schema, dict) else None + + return None + + def _schema_types(self, schema: Any) -> List[str]: + if not isinstance(schema, dict): + return [] + + schema_type = schema.get("type") + if isinstance(schema_type, str): + return [schema_type.lower()] + if isinstance(schema_type, list): + return [t.lower() for t in schema_type if isinstance(t, str)] + return [] + + def _schema_has_type(self, schema: Any, schema_types: Tuple[str, ...]) -> bool: + return any(t in self._schema_types(schema) for t in schema_types) + + def _is_scalar_schema(self, schema: Any) -> bool: + schema_types = set(self._schema_types(schema)) + return bool(schema_types & SCALAR_TYPES) and not schema_types & CONTAINER_TYPES + + def _new_container_for_schema(self, schema: Any) -> Any: + if self._schema_has_type(schema, ("array",)): + return [] + return {} + + def _convert_leaf_value(self, value: str, schema: Any) -> Any: + schema_types = set(self._schema_types(schema)) + null_permitted = "null" in schema_types + + # Return verbatim; streaming emits strings literally, so null-coercing + # here would diverge the non-streaming path from streaming. + if schema_types & STRING_TYPES and not null_permitted: + return value + + if null_permitted: + if value in NULL_STRINGS: + return None + if value == "": + return None + + lower_value = value.lower().strip() + + if schema_types & INTEGER_TYPES: + try: + return int(value) + except (TypeError, ValueError): + pass + + if schema_types & NUMBER_TYPES: + try: + parsed = float(value) + return parsed if parsed != int(parsed) else int(parsed) + except (TypeError, ValueError): + pass + + if schema_types & BOOLEAN_TYPES: + if lower_value in ("true", "1", "yes", "on"): + return True + if lower_value in ("false", "0", "no", "off"): + return False + + return value + + def _assign_child( + self, parent_value: Any, child_tag: str, child_value: Any + ) -> None: + if isinstance(parent_value, list): + parent_value.append(child_value) + return + + if child_tag in parent_value: + existing_value = parent_value[child_tag] + if isinstance(existing_value, list): + existing_value.append(child_value) + else: + parent_value[child_tag] = [existing_value, child_value] + return + + parent_value[child_tag] = child_value + + def _parse_parameter(self, body: str, parameters_schema: Optional[Dict]) -> dict: + if self._schema_has_type( + parameters_schema, ("array",) + ) and self._body_starts_with_item(body): + root: Any = [] + else: + root = {} + stack: List[Dict[str, Any]] = [ + {"tag": "", "schema": parameters_schema, "value": root} + ] + + for chunk in body.split(MINIMAX_NS_TOKEN): + chunk = chunk.strip() + if not chunk: + continue + + if chunk.startswith("", 2) + tag = chunk[2:gt].strip() if gt != -1 else chunk[2:].strip() + if len(stack) == 1: + raise ValueError(f"unexpected closing tag: {tag}") + if stack[-1]["tag"] != tag: + raise ValueError( + f"mismatched closing tag: expected {stack[-1]['tag']}, got {tag}" + ) + + frame = stack.pop() + self._assign_child(stack[-1]["value"], frame["tag"], frame["value"]) + continue + + if chunk.startswith("<"): + gt = chunk.index(">") + tag = chunk[1:gt].strip() + text = chunk[gt + 1 :] + parent_frame = stack[-1] + child_schema = self._get_child_schema( + parent_frame["schema"], tag, parent_frame["value"] + ) + + if text or self._is_scalar_schema(child_schema): + value = self._convert_leaf_value(text, child_schema) + else: + value = self._new_container_for_schema(child_schema) + stack.append({"tag": tag, "schema": child_schema, "value": value}) + + return root + + def _body_starts_with_item(self, body: str) -> bool: + for chunk in body.split(MINIMAX_NS_TOKEN): + chunk = chunk.strip() + if chunk: + return chunk.startswith("") + return False diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 76f855fd5..99e430809 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -258,6 +258,19 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac runner.hybrid_gdn_config is not None and runner.use_mla_backend ), "hybrid_gdn can only be used with non-MLA models." + from sglang.srt.configs.model_config import is_minimax_sparse + + if is_minimax_sparse(runner.model_config.hf_config): + from sglang.srt.layers.attention.minimax_sparse_backend import ( + MiniMaxHybridAttnBackend, + MiniMaxSparseAttnBackend, + ) + + sparse_backend = MiniMaxSparseAttnBackend(runner) + return MiniMaxHybridAttnBackend( + full_attn_backend, sparse_backend, sparse_backend.sparse_layer_ids + ) + if cfg := runner.mambaish_config: from sglang.srt.layers.attention.fla.utils import check_environments from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py new file mode 100644 index 000000000..d234aa088 --- /dev/null +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -0,0 +1,602 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.configs.model_config import ( + get_minimax_sparse_attention_config, + get_minimax_sparse_disable_value_layer_ids, + get_minimax_sparse_layer_ids, + get_minimax_sparse_score_type, +) +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.minimax_sparse_ops.minimax_sparse import ( + minimax_sparse_decode, + minimax_sparse_prefill, +) +from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + +logger = logging.getLogger(__name__) + + +class MiniMaxSparseAttnBackend(AttentionBackend): + def __init__(self, runner: ModelRunner): + assert isinstance(runner.token_to_kv_pool, MiniMaxSparseKVPool) + self.kv_pool = runner.token_to_kv_pool + self.req_to_token = runner.req_to_token_pool.req_to_token + self.max_context_len = int(runner.model_config.context_len) + + hf_config = runner.model_config.hf_config + sparse_cfg = get_minimax_sparse_attention_config(hf_config) + self.idx_head_dim = sparse_cfg["sparse_index_dim"] + self.dense_layer_ids, self.sparse_layer_ids = get_minimax_sparse_layer_ids( + sparse_cfg + ) + self.disable_value_layer_ids: set[int] = set( + get_minimax_sparse_disable_value_layer_ids(sparse_cfg) + ) + self.score_type: str = get_minimax_sparse_score_type(sparse_cfg) + + # Plain Python int so it is safe inside CUDA graphs (no .item() at graph time). + self._max_seqlen_q: int = 1 + self._max_seqlen_k: int = 1 + + self.block_size_q = 1 + self.block_size_k = sparse_cfg["sparse_block_size"] + if "sparse_init_block" in sparse_cfg: + self.init_blocks = sparse_cfg["sparse_init_block"] + else: + init_tokens = sparse_cfg["sparse_init_tokens"] + self.init_blocks = ( + init_tokens + self.block_size_k - 1 + ) // self.block_size_k + if "sparse_local_block" in sparse_cfg: + self.local_blocks = sparse_cfg["sparse_local_block"] + else: + local_tokens = sparse_cfg["sparse_local_tokens"] + self.local_blocks = ( + local_tokens + self.block_size_k - 1 + ) // self.block_size_k + 1 + self.topk_blocks = sparse_cfg["sparse_topk_blocks"] + + # MSA (fmha_sm100) is SM100-only; fall back to the Triton sparse path when + # the kernel is unavailable or its constraints don't hold. + from sglang.srt.environ import envs + from sglang.srt.layers.attention.minimax_sparse_ops.msa import ( + msa_available, + ) + + # MSA (fmha_sm100) is bf16/fp16-only; an fp8 main KV cache must stay on the + # Triton sparse path (it dequants fp8 on load). + _main_kv_is_fp8 = self.kv_pool.main_pool.dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + self.use_msa = ( + not envs.SGLANG_DISABLE_MSA.get() + and msa_available() + and self.block_size_k == 128 + and self.kv_pool.page_size == self.block_size_k + and self.topk_blocks in (4, 8, 16, 32) + and not _main_kv_is_fp8 + ) + if ( + not self.use_msa + and not envs.SGLANG_DISABLE_MSA.get() + and msa_available() + and self.block_size_k == 128 + and self.kv_pool.page_size != self.block_size_k + ): + logger.warning( + "MiniMax-M3 MSA decode disabled: page_size=%d != sparse block size " + "%d. Pass --page-size 128 (with an attention backend that allows it, " + "e.g. fa4) to enable the faster MSA kernel; falling back to the " + "Triton sparse path.", + self.kv_pool.page_size, + self.block_size_k, + ) + self._msa_dec_meta = None + if self.use_msa: + from sglang.srt.runtime_context import get_parallel + + self.num_q_heads = ( + runner.model_config.num_attention_heads // get_parallel().attn_tp_size + ) + self.num_kv_heads = self.kv_pool.main_pool.head_num + self._msa_nb_max = ( + self.max_context_len + self.block_size_k - 1 + ) // self.block_size_k + self._msa_cg: dict[int, tuple] = {} + + self.page_size = self.kv_pool.page_size + self.use_dense_sparse_decode = ( + envs.SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE.get() + and self.block_size_k % self.page_size == 0 + ) + # MSA fmha_sm100 decode is NOT cuda-graph-safe: captured/replayed it returns + # wrong results (~14% GSM8K loss on B200). Gate capture via cuda_graph_config, + # not legacy disable_* flags — they disagree under config-native flags and would + # capture the unsafe MSA decode kernel. + from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + Phase, + check_cuda_graph_backend, + ) + + _sa = getattr(runner, "server_args", None) + _decode_cuda_graph = not check_cuda_graph_backend( + Phase.DECODE, Backend.DISABLED + ) + self._use_msa_decode = self.use_msa and ( + not _decode_cuda_graph or envs.SGLANG_OPT_USE_MSA_DECODE_UNDER_GRAPH.get() + ) + + # MSA + spec decode + cuda graph crashes mid-capture: TARGET_VERIFY batches + # route to forward_extend, dereferencing absent extend metadata. Fail at startup. + if ( + self.use_msa + and _decode_cuda_graph + and getattr(_sa, "speculative_algorithm", None) is not None + ): + raise NotImplementedError( + "MiniMax-M3 MSA attention does not support speculative decoding under " + "CUDA graph. Use --disable-cuda-graph, set SGLANG_DISABLE_MSA=1, or " + "disable speculative decoding." + ) + self._msa_owns_decode = self._use_msa_decode and not ( + self.use_dense_sparse_decode and self.kv_pool.main_pool.head_num == 1 + ) + self.dense_backend: Optional[AttentionBackend] = None + + logger.info( + f"[MiniMaxSparse] Backend initialized " + f"(score_type={self.score_type!r}, " + f"main_attn={'MSA' if self.use_msa else 'triton'}, " + f"disable_value_layers={sorted(self.disable_value_layer_ids)})" + ) + + def init_forward_metadata_out_graph( + self, forward_batch: ForwardBatch, in_capture: bool = False + ): + # cuda-graph replay views are a SimpleNamespace without extend_seq_lens_cpu, + # and TARGET_VERIFY sets it to None despite is_extend() — getattr covers both. + self._msa_dec_meta = None + extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None) + if extend_lens is not None: + self._max_seqlen_q = int(max(extend_lens)) + else: + self._max_seqlen_q = 1 + if in_capture and forward_batch.forward_mode.is_decode_or_idle(): + self._max_seqlen_k = self.max_context_len + else: + self._max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item()) + + # Build plan + page table eager (outside capture) so captured forward_decode + # runs only device-side ops; host-side code can't be captured. + if self._msa_owns_decode and forward_batch.forward_mode.is_decode_or_idle(): + self._prepare_msa_decode_meta(forward_batch) + + def _prepare_msa_decode_meta(self, forward_batch: ForwardBatch): + from sglang.srt.layers.attention.minimax_sparse_ops.msa import ( + build_msa_decode_cg_plan, + update_msa_decode_cg_meta, + ) + + bs = forward_batch.seq_lens.shape[0] + if bs == 0: + return + entry = self._msa_cg.get(bs) + if entry is None: + device = forward_batch.seq_lens.device + plan = build_msa_decode_cg_plan( + self.num_q_heads, + self.num_kv_heads, + self.block_size_k, + self.topk_blocks, + bs, + device=device, + ) + kv_indices_buf = torch.zeros( + bs * self._msa_nb_max, dtype=torch.int32, device=device + ) + entry = (plan, kv_indices_buf) + self._msa_cg[bs] = entry + plan, kv_indices_buf = entry + update_msa_decode_cg_meta( + plan, + kv_indices_buf, + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + self.block_size_k, + self.topk_blocks, + self.num_q_heads, + self.num_kv_heads, + ) + self._msa_dec_meta = (kv_indices_buf, plan) + + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): + pass + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + pass + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + @staticmethod + def _is_sparse_kv_cached_by_fusion( + forward_batch: ForwardBatch, layer_id: int + ) -> bool: + layer_ids = forward_batch.minimax_m3_precached_sparse_layers + return layer_ids is not None and layer_id in layer_ids + + def forward( + self, + q, + k, + v, + layer, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + if forward_batch.forward_mode.is_idle(): + idx_q = kwargs.get("idx_q") + num_idx_heads = idx_q.shape[1] + disable_value = layer.layer_id in self.disable_value_layer_ids + idx_out: Optional[torch.Tensor] = ( + None + if disable_value + else q.new_zeros(q.shape[0], num_idx_heads * self.idx_head_dim) + ) + out = q.new_zeros(q.shape[0], layer.tp_q_head_num * layer.v_head_dim) + return idx_out, out + else: + return super().forward( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer, + forward_batch: ForwardBatch, + save_kv_cache=True, + *, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + idx_v: Optional[torch.Tensor], + ): + disable_value = layer.layer_id in self.disable_value_layer_ids + kv_cached_by_fusion = self._is_sparse_kv_cached_by_fusion( + forward_batch, layer.layer_id + ) + if not kv_cached_by_fusion: + self.kv_pool.set_fused_kv_index_buffer( + layer, + forward_batch.out_cache_loc, + k, + v, + idx_k, + None if disable_value else idx_v, + ) + k_cache, v_cache = self.kv_pool.get_kv_buffer(layer.layer_id) + if disable_value: + idx_k_cache = self.kv_pool.get_index_k_buffer(layer.layer_id) + idx_v_cache = None + else: + idx_k_cache, idx_v_cache = self.kv_pool.get_index_kv_buffer(layer.layer_id) + + cu_seqlens = torch.cat( + [ + torch.zeros( + 1, dtype=torch.int32, device=forward_batch.extend_seq_lens.device + ), + forward_batch.extend_seq_lens.to(torch.int32).cumsum(0).to(torch.int32), + ] + ) + seq_lens = forward_batch.seq_lens.to(torch.int32) + if forward_batch.extend_prefix_lens is not None: + prefix_lens = forward_batch.extend_prefix_lens.to(torch.int32) + else: + prefix_lens = torch.zeros_like(seq_lens) + + # DP attention pads q beyond the real token count for collective alignment; + # trim to actual tokens so the sparse kernel sees consistent shapes. + if forward_batch.extend_seq_lens_cpu is not None: + actual_num_tokens = int(sum(forward_batch.extend_seq_lens_cpu)) + else: + actual_num_tokens = int(cu_seqlens[-1].item()) + original_num_tokens = q.shape[0] + if actual_num_tokens < original_num_tokens: + q = q[:actual_num_tokens] + idx_q = idx_q[:actual_num_tokens] + + idx_o, o = minimax_sparse_prefill( + q, + k_cache, + v_cache, + None, + idx_q, + idx_k_cache, + idx_v_cache, + None, + self.req_to_token, + forward_batch.req_pool_indices, + cu_seqlens, + seq_lens, + prefix_lens, + self._max_seqlen_q, + self._max_seqlen_k, + self.block_size_q, + self.block_size_k, + self.topk_blocks, + self.init_blocks, + self.local_blocks, + score_type=self.score_type, + disable_index_value=disable_value, + use_msa=self.use_msa, + seqlens_cpu=forward_batch.extend_seq_lens_cpu, + ) + + if actual_num_tokens < original_num_tokens: + pad_len = original_num_tokens - actual_num_tokens + o = torch.cat([o, o.new_zeros(pad_len, *o.shape[1:])], dim=0) + if idx_o is not None: + idx_o = torch.cat( + [idx_o, idx_o.new_zeros(pad_len, *idx_o.shape[1:])], dim=0 + ) + + return ( + ( + None + if idx_o is None + else idx_o.reshape(original_num_tokens, -1).contiguous() + ), + o.reshape(original_num_tokens, -1).contiguous(), + ) + + def _dense_sparse_main_decode( + self, + q: torch.Tensor, + page_table: torch.Tensor, + real_seq_lens: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + layer, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend + + if isinstance(self.dense_backend, TRTLLMHAAttnBackend): + import flashinfer + + ps = self.page_size + nkv = 1 + head_dim = q.size(-1) + # [max_slots, nkv, D] -> [num_pages, page_size, nkv, D] + # -> [num_pages, nkv, page_size, D] (HND, trtllm default) + kc = k_cache.view(-1, ps, nkv, head_dim).permute(0, 2, 1, 3) + vc = v_cache.view(-1, ps, nkv, head_dim).permute(0, 2, 1, 3) + return flashinfer.decode.trtllm_batch_decode_with_kv_cache( # type: ignore + query=q.contiguous(), + kv_cache=(kc, vc), + workspace_buffer=self.dense_backend.workspace_buffer, + block_tables=page_table, + seq_lens=real_seq_lens, + max_seq_len=self.topk_blocks * self.block_size_k, + bmm1_scale=layer.scaling, + bmm2_scale=1.0, + ) + raise NotImplementedError( + "dense sparse decode currently supports trtllm_mha only (fa3 is TODO)" + ) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + *, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + idx_v: Optional[torch.Tensor], + **kwargs, + ): + assert len(kwargs) == 0 + disable_value = layer.layer_id in self.disable_value_layer_ids + self.kv_pool.set_fused_kv_index_buffer( + layer, + forward_batch.out_cache_loc, + k, + v, + idx_k, + None if disable_value else idx_v, + ) + k_cache, v_cache = self.kv_pool.get_kv_buffer(layer.layer_id) + if disable_value: + idx_k_cache = self.kv_pool.get_index_k_buffer(layer.layer_id) + idx_v_cache = None + else: + idx_k_cache, idx_v_cache = self.kv_pool.get_index_kv_buffer(layer.layer_id) + + attn_fn = None + if self.use_dense_sparse_decode and k_cache.shape[1] == 1: + + def attn_fn(main_q, page_table, real_seq_lens): + return self._dense_sparse_main_decode( + main_q, + page_table, + real_seq_lens, + k_cache, + v_cache, + layer, + forward_batch, + ) + + msa_kv_indices = msa_plan = None + if self._use_msa_decode and attn_fn is None: + if self._msa_dec_meta is not None: + msa_kv_indices, msa_plan = self._msa_dec_meta + elif q.shape[0] > 0: + # Rebuilding the plan inline would run host-side code inside + # CUDA-graph capture; fail loudly instead. + raise RuntimeError( + "MSA decode metadata missing: init_forward_metadata_out_graph " + "did not prepare the plan for this forward (gate mismatch)." + ) + + idx_o, o = minimax_sparse_decode( + q, + None, + k_cache, + v_cache, + idx_q, + None, + idx_k_cache, + idx_v_cache, + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + self._max_seqlen_k, + 1, + self.block_size_k, + self.topk_blocks, + self.init_blocks, + self.local_blocks, + score_type=self.score_type, + disable_index_value=disable_value, + dense_main_attn_fn=attn_fn, + page_size=self.page_size, + use_msa=self._use_msa_decode, + msa_kv_indices=msa_kv_indices, + msa_plan=msa_plan, + ) + return ( + None if idx_o is None else idx_o.reshape(q.shape[0], -1).contiguous(), + o.reshape(q.shape[0], -1).contiguous(), + ) + + +class MiniMaxHybridAttnBackend(AttentionBackend): + def __init__( + self, + dense_backend: AttentionBackend, + sparse_backend: MiniMaxSparseAttnBackend, + sparse_layer_ids: list[int], + ): + self.dense = dense_backend + self.sparse = sparse_backend + self.sparse_layer_ids = sparse_layer_ids + self.sparse.dense_backend = dense_backend + + def init_forward_metadata(self, forward_batch: ForwardBatch): + self.sparse.init_forward_metadata(forward_batch) + self.dense.init_forward_metadata(forward_batch) + + def init_forward_metadata_out_graph( + self, forward_batch: ForwardBatch, in_capture: bool = False + ): + self.sparse.init_forward_metadata_out_graph(forward_batch, in_capture) + self.dense.init_forward_metadata_out_graph(forward_batch, in_capture) + + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): + self.sparse.init_forward_metadata_in_graph(forward_batch) + self.dense.init_forward_metadata_in_graph(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + self.dense.init_cuda_graph_state(max_bs, max_num_tokens) + self.sparse.init_cuda_graph_state(max_bs, max_num_tokens) + + def get_cuda_graph_seq_len_fill_value(self): + return self.sparse.get_cuda_graph_seq_len_fill_value() + + def forward( + self, + q, + k, + v, + layer, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + if layer.layer_id in self.sparse_layer_ids: + return self.sparse.forward( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + # DP attention pads q to an even length but flashinfer builds qo_indptr from + # extend_seq_lens, so padded q.shape[0] != qo_indptr[-1] and paged-prefill + # raises. Trim q and re-pad output; k/v stay untrimmed so KV-cache writes + # align with out_cache_loc. + mode = forward_batch.forward_mode + if mode.is_extend() and forward_batch.extend_seq_lens_cpu is not None: + actual_num_tokens = int(sum(forward_batch.extend_seq_lens_cpu)) + original_num_tokens = q.shape[0] + if actual_num_tokens < original_num_tokens: + o = self.dense.forward( + q[:actual_num_tokens], + k, + v, + layer, + forward_batch, + save_kv_cache, + **kwargs, + ) + pad_len = original_num_tokens - actual_num_tokens + return torch.cat([o, o.new_zeros(pad_len, *o.shape[1:])], dim=0) + + return self.dense.forward( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_extend( + self, + q, + k, + v, + layer, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + if layer.layer_id in self.sparse_layer_ids: + return self.sparse.forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + else: + return self.dense.forward_extend( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + + def forward_decode( + self, + q, + k, + v, + layer, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + **kwargs, + ): + if layer.layer_id in self.sparse_layer_ids: + return self.sparse.forward_decode( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) + else: + return self.dense.forward_decode( + q, k, v, layer, forward_batch, save_kv_cache, **kwargs + ) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 003a716a7..6276dc976 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -171,6 +171,32 @@ def gemm_nt_f8f8bf16( ) +def gemm_nt_mxfp8_f8f8bf16( + lhs: Tuple[torch.Tensor, torch.Tensor], + rhs: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, +): + m, k = lhs[0].shape + n, _ = rhs[0].shape + num_groups = 1 + kernel_type = compile_utils.DeepGemmKernelType.GEMM_NT_F8F8BF16 + + _sanity_check_input(lhs) + _sanity_check_input(rhs) + + disable_cast = lhs[1].dtype == torch.int and rhs[1].dtype == torch.int + + with compile_utils.deep_gemm_execution_hook(m, n, k, num_groups, kernel_type): + deep_gemm.fp8_fp4_gemm_nt( + lhs, + rhs, + out, + recipe_a=(1, 32), + recipe_b=(1, 32), + disable_ue8m0_cast=disable_cast, + ) + + def gemm_nt_bf16bf16f32( lhs: torch.Tensor, rhs: torch.Tensor, diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index fd1bf385f..2c95c01ec 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -93,6 +93,7 @@ if _is_cuda or _is_xpu or _is_musa: ) _has_aiter_layer_norm = False _has_vllm_rms_norm = False +_has_rocm_triton_gemma_rms_norm = False if _use_aiter: from aiter import layernorm2d_fwd as layer_norm from aiter import rmsnorm2d_fwd as rms_norm @@ -109,6 +110,19 @@ elif _is_hip: # Fallback: vllm not available, will use forward_native _has_vllm_rms_norm = False +if _is_hip: + try: + from sglang.jit_kernel.minimax_m3.rmsnorm import ( + gemma_fused_add_rmsnorm as rocm_triton_gemma_fused_add_rmsnorm, + ) + from sglang.jit_kernel.minimax_m3.rmsnorm import ( + gemma_rmsnorm as rocm_triton_gemma_rmsnorm, + ) + + _has_rocm_triton_gemma_rms_norm = True + except ImportError: + _has_rocm_triton_gemma_rms_norm = False + if _is_cuda: # HF-semantics RMSNorm kernel (JIT-compiled). Used when `cast_x_before_out_mul=True` # (the transformers backend path) to produce outputs that are numerically identical @@ -357,6 +371,10 @@ class RMSNorm(MultiPlatformOp): if residual is not None: return x, residual return x + if self.weight.data.dtype != x.dtype: + # AITER's ROCm rmsnorm2d_fwd requires weight/activation dtypes to match; + # FP32 weight + BF16 activation yields finite-but-corrupted output on gfx950. + return self.forward_native(x, residual, post_residual_addition) # Aiter's RMSNorm kernels expect 2D contiguous inputs. Keep the # already-safe layout as a zero-copy path, and only normalize strided or # higher-rank views such as Q/K slices from packed QKV projections. @@ -749,24 +767,24 @@ class GemmaRMSNorm(MultiPlatformOp): residual: Optional[torch.Tensor] = None, post_residual_addition: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if _use_aiter and _has_rocm_triton_gemma_rms_norm: + if residual is not None: + if post_residual_addition is not None: + residual = residual + post_residual_addition + return rocm_triton_gemma_fused_add_rmsnorm( + x, residual, self.weight.data, self.variance_epsilon + ) + return rocm_triton_gemma_rmsnorm(x, self.weight.data, self.variance_epsilon) + if not _has_vllm_rms_norm: return self.forward_native(x, residual, post_residual_addition) - w = self.gemma_weight if _use_aiter: - # aiter API: rms_norm(input, weight, eps) -> output - # fused_add_rms_norm(output, input, residual, residual_out, weight, eps) - if residual is not None: - output = torch.empty_like(x) - residual_out = torch.empty_like(x) - if post_residual_addition is not None: - residual = residual + post_residual_addition - fused_add_rms_norm( - output, x, residual, residual_out, w, self.variance_epsilon - ) - return output, residual_out - return rms_norm(x, w, self.variance_epsilon) + # AITER's ROCm rmsnorm2d_fwd has the same dtype requirement here; + # keep Gemma RMSNorm on native torch math for correctness. + return self.forward_native(x, residual, post_residual_addition) else: + w = self.gemma_weight # vllm API: rms_norm(out, input, weight, eps) -> None (in-place) # fused_add_rms_norm(out, input, residual_out, residual, weight, eps) if not x.is_contiguous(): diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py index 086fa0db2..4226ebb48 100644 --- a/python/sglang/srt/layers/moe/ep_moe/kernels.py +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -1,4 +1,5 @@ import logging +from typing import Tuple import torch import triton @@ -181,23 +182,6 @@ def deepep_run_moe_deep_preprocess(topk_ids: torch.Tensor, num_experts: int): return reorder_topk_ids, src2dst, seg_indptr -@triton.jit -def compute_seg_indptr_triton_kernel(reorder_topk_ids, seg_indptr, num_toks): - expert_id_minus_1 = tl.program_id(0) - 1 - low = 0 - high = num_toks - 1 - target_location = -1 - while low <= high: - mid = (low + high) // 2 - - if tl.load(reorder_topk_ids + mid) > expert_id_minus_1: - high = mid - 1 - else: - low = mid + 1 - target_location = mid - tl.store(seg_indptr + expert_id_minus_1 + 1, target_location + 1) - - def cutlass_w4_run_moe_ep_preproess(topk_ids: torch.Tensor): _, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) @@ -303,9 +287,12 @@ def _silu_and_mul_post_quant_kernel( size_n, fp8_max, fp8_min, + QUANT_GROUP_SIZE: tl.constexpr, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, SCALE_UE8M0: tl.constexpr, + GEMM1_ALPHA: tl.constexpr, + GEMM1_CLAMP_LIMIT: tl.constexpr, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) @@ -320,13 +307,15 @@ def _silu_and_mul_post_quant_kernel( stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64) stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64) + N_GROUPS: tl.constexpr = BLOCK_N // QUANT_GROUP_SIZE + offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d - output_scale_offs = ( + scale_base = ( output_scale_ptr + expert_id * stride_output_scale_0 - + hidden_dim_block_index * stride_output_scale_2 + + hidden_dim_block_index * N_GROUPS * stride_output_scale_2 ) for token_index in tl.range( @@ -342,23 +331,35 @@ def _silu_and_mul_post_quant_kernel( mask=offs_in_d < size_n, other=0.0, ) - gate = gate / (1 + tl.exp(-gate)) - gate = gate.to(input_ptr.dtype.element_ty) - gate_up = up * gate - _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) - output_s = _absmax / fp8_max + if GEMM1_ALPHA > 0: + gate = tl.minimum(gate, GEMM1_CLAMP_LIMIT) + up = tl.clamp(up, -GEMM1_CLAMP_LIMIT, GEMM1_CLAMP_LIMIT) + gate_up = gate * tl.sigmoid(gate * GEMM1_ALPHA) * (up + 1) + else: + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + + gate_up_2d = tl.reshape(gate_up, (N_GROUPS, QUANT_GROUP_SIZE)) + group_absmax = tl.max(tl.abs(gate_up_2d), axis=1) + group_absmax = tl.maximum(group_absmax, 1e-10) + + output_s = group_absmax / fp8_max if SCALE_UE8M0: - output_s = tl.exp2(tl.ceil(tl.log2(tl.abs(output_s)))) - output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to( - output_ptr.dtype.element_ty - ) + output_s = tl.exp2(tl.ceil(tl.log2(output_s))) + + inv_s = tl.reshape(1.0 / output_s, (N_GROUPS, 1)) + output_q_2d = tl.clamp(gate_up_2d * inv_s, fp8_min, fp8_max) + output_q = tl.reshape(output_q_2d, (BLOCK_N,)).to(output_ptr.dtype.element_ty) + tl.store( output_ptr_offs + token_index * stride_output_1, output_q, mask=offs_in_d < size_n, ) + scale_offs = scale_base + token_index * stride_output_scale_1 tl.store( - output_scale_offs + token_index * stride_output_scale_1, + scale_offs + tl.arange(0, N_GROUPS) * stride_output_scale_2, output_s, ) @@ -370,6 +371,8 @@ def silu_and_mul_masked_post_quant_fwd( quant_group_size: int, masked_m: torch.Tensor, scale_ue8m0: bool = False, + gemm1_alpha: float = 0.0, + gemm1_clamp_limit: float = 0.0, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -396,11 +399,18 @@ def silu_and_mul_masked_post_quant_fwd( else: BLOCK_NUM_PER_EXPERT = 32 - BLOCK_N = quant_group_size + groups_total = size_n // quant_group_size + gpb = 4 + while gpb > 1: + block_n = quant_group_size * gpb + if (block_n & (block_n - 1) == 0) and (groups_total % gpb == 0): + break + gpb //= 2 + BLOCK_N = quant_group_size * gpb + num_warps = 1 NUM_STAGES = 6 hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N) - assert BLOCK_N % quant_group_size == 0 grid = ( hidden_dim_split_block_num, @@ -423,10 +433,155 @@ def silu_and_mul_masked_post_quant_fwd( size_n, fp8_max, fp8_min, + QUANT_GROUP_SIZE=quant_group_size, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, num_warps=num_warps, SCALE_UE8M0=scale_ue8m0, + GEMM1_ALPHA=gemm1_alpha if gemm1_alpha is not None else 0.0, + GEMM1_CLAMP_LIMIT=gemm1_clamp_limit if gemm1_clamp_limit is not None else 0.0, + ) + return + + +@triton.jit +def _silu_and_mul_post_quant_packed_kernel( + input_ptr, + stride_input_0, + stride_input_1, + stride_input_2, + output_ptr, + stride_output_0, + stride_output_1, + stride_output_2, + scale_ptr, # int32 [E, G//4, m_max] (MN-major: packed scale stored token-minor) + stride_scale_e, + stride_scale_g4, + stride_scale_m, + masked_m_ptr, + num_experts, + size_n, + fp8_max, + fp8_min, + QUANT_GROUP_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, # == 4 * QUANT_GROUP_SIZE (one packed int32 per block) + GEMM1_ALPHA: tl.constexpr, + GEMM1_CLAMP_LIMIT: tl.constexpr, + E_PADDED: tl.constexpr, +): + # Flat work dim rides axis 0 (x): num_real_tokens*topk can exceed the 65535 grid.y/z limit. + work_id = tl.program_id(0) + hidden_dim_block_index = tl.program_id(1) + + e_off = tl.arange(0, E_PADDED) + mm = tl.load(masked_m_ptr + e_off, mask=e_off < num_experts, other=0) + incl = tl.cumsum(mm) + total = tl.sum(mm) + if work_id >= total: + return + excl = incl - mm # first global slot of each expert + owner = (excl <= work_id) & (work_id < incl) + expert_id = tl.sum(tl.where(owner, e_off, 0)) + token_index = work_id - tl.sum(tl.where(owner, excl, 0)) + + stride_input_0 = tl.cast(stride_input_0, tl.int64) + stride_input_1 = tl.cast(stride_input_1, tl.int64) + stride_output_0 = tl.cast(stride_output_0, tl.int64) + stride_output_1 = tl.cast(stride_output_1, tl.int64) + + N_GROUPS: tl.constexpr = BLOCK_N // QUANT_GROUP_SIZE + + offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) + mask_d = offs_in_d < size_n + in_base = input_ptr + expert_id * stride_input_0 + token_index * stride_input_1 + out_base = output_ptr + expert_id * stride_output_0 + token_index * stride_output_1 + + gate = tl.load(in_base + offs_in_d, mask=mask_d, other=0.0).to(tl.float32) + up = tl.load(in_base + offs_in_d + size_n, mask=mask_d, other=0.0) + if GEMM1_ALPHA > 0: + gate = tl.minimum(gate, GEMM1_CLAMP_LIMIT) + up = tl.clamp(up, -GEMM1_CLAMP_LIMIT, GEMM1_CLAMP_LIMIT) + gate_up = gate * tl.sigmoid(gate * GEMM1_ALPHA) * (up + 1) + else: + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + + gate_up_2d = tl.reshape(gate_up, (N_GROUPS, QUANT_GROUP_SIZE)) + group_absmax = tl.max(tl.abs(gate_up_2d), axis=1) + group_absmax = tl.maximum(group_absmax, 1e-10) + output_s = group_absmax / fp8_max + # UE8M0: round to a power of two so the (>>23)&0xFF exponent extraction below is valid. + output_s = tl.exp2(tl.ceil(tl.log2(output_s))) + + inv_s = tl.reshape(1.0 / output_s, (N_GROUPS, 1)) + output_q_2d = tl.clamp(gate_up_2d * inv_s, fp8_min, fp8_max) + output_q = tl.reshape(output_q_2d, (BLOCK_N,)).to(output_ptr.dtype.element_ty) + tl.store(out_base + offs_in_d, output_q, mask=mask_d) + + # Pack 4 UE8M0 exponent bytes little-endian into one int32, matching deep_gemm's + # get_mn_major_tma_aligned_packed_ue8m0_tensor (fp32>>23 -> uint8 -> 4-group int32 view). + s_bits = output_s.to(tl.int32, bitcast=True) + expo = (s_bits >> 23) & 0xFF + shifts = tl.arange(0, N_GROUPS) * 8 + packed = tl.sum(expo << shifts) + scale_off = ( + expert_id * stride_scale_e + + hidden_dim_block_index * stride_scale_g4 + + token_index * stride_scale_m + ) + tl.store(scale_ptr + scale_off, packed) + + +def silu_and_mul_masked_post_quant_packed_fwd( + input: torch.Tensor, + output: torch.Tensor, + output_scale_packed: torch.Tensor, + quant_group_size: int, + masked_m: torch.Tensor, + num_real_tokens: int, + topk: int, + gemm1_alpha: float = 0.0, + gemm1_clamp_limit: float = 0.0, +): + assert input.is_contiguous() + assert output.dtype == torch.float8_e4m3fn + assert output.is_contiguous() + assert input.dim() == 3 and input.shape[-1] % 2 == 0 + + E, m_max, _ = input.shape + size_n = input.shape[-1] // 2 + assert size_n % quant_group_size == 0 + G = size_n // quant_group_size + assert G % 4 == 0, "packed UE8M0 path requires num_groups % 4 == 0" + BLOCK_N = quant_group_size * 4 + assert size_n % BLOCK_N == 0, "packed UE8M0 path requires size_n % (4*group) == 0" + hidden_dim_split = size_n // BLOCK_N + assert tuple(output_scale_packed.shape) == (E, hidden_dim_split, m_max) + assert output_scale_packed.dtype == torch.int32 + + finfo = torch.finfo(torch.float8_e4m3fn) + fp8_max = finfo.max + + grid = (num_real_tokens * topk, hidden_dim_split) + _silu_and_mul_post_quant_packed_kernel[grid]( + input, + *input.stride(), + output, + *output.stride(), + output_scale_packed, + *output_scale_packed.stride(), + masked_m, + E, + size_n, + fp8_max, + -fp8_max, + QUANT_GROUP_SIZE=quant_group_size, + BLOCK_N=BLOCK_N, + GEMM1_ALPHA=gemm1_alpha if gemm1_alpha is not None else 0.0, + GEMM1_CLAMP_LIMIT=gemm1_clamp_limit if gemm1_clamp_limit is not None else 0.0, + E_PADDED=triton.next_power_of_2(E), + num_warps=1, ) return @@ -677,6 +832,83 @@ def post_reorder_for_cutlass_moe( ) +@triton.jit +def post_reorder_deepgemm_triton_kernel( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + topk, + num_tokens, + hidden_size, + routed_scaling_factor: float, + BLOCK_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + """`expert_id >= 0` includes the shared expert at num_experts (padding=-1); don't + switch to the cutlass `!= num_local_experts` gate. routed_scaling_factor is folded into the store. + """ + OutDtype = output_ptr.dtype.element_ty + + offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE) + mask = offset < hidden_size + + down_output_ptr_offs = down_output_ptr + offset + output_ptr_offs = output_ptr + offset + + start_src_idx = tl.program_id(0) + step = tl.num_programs(0) + + for src_idx_int32 in tl.range( + start_src_idx, num_tokens, step, num_stages=NUM_STAGES + ): + src_idx = src_idx_int32.to(tl.int64) + token_src2dst_ptr = src2dst_ptr + src_idx * topk + token_topk_ids_ptr = topk_ids_ptr + src_idx * topk + token_topk_weights_ptr = topk_weights_ptr + src_idx * topk + + sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for idx in range(topk): + expert_id = tl.load(token_topk_ids_ptr + idx) + if expert_id >= 0: + dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64) + weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32) + load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size + in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32) + sum_vec += in_data * weight_scale + sum_vec *= routed_scaling_factor + store_ptr_offs = output_ptr_offs + src_idx * hidden_size + tl.store(store_ptr_offs, sum_vec.to(OutDtype), mask=mask) + + +def post_reorder_deepgemm( + down_output, + output, + src2dst, + topk_ids, + topk_weights, + topk, + num_tokens, + hidden_size, + routed_scaling_factor: float, +): + grid, block_dim = _get_launch_config_2d(down_output.device, num_tokens, hidden_size) + post_reorder_deepgemm_triton_kernel[grid]( + down_output, + output, + src2dst, + topk_ids, + topk_weights, + topk, + num_tokens, + hidden_size, + float(routed_scaling_factor), + BLOCK_SIZE=block_dim, + NUM_STAGES=3, + ) + + @triton.jit def post_reorder_triton_kernel( down_output_ptr, @@ -846,9 +1078,10 @@ def ep_scatter( m_indices: torch.Tensor, output_index: torch.Tensor, scale_ue8m0: bool = False, + quant_block_size: int = 128, ): BLOCK_E = 128 # token num of per expert is aligned to 128 - BLOCK_D = 128 # block size of quantization + BLOCK_D = quant_block_size # block size of quantization num_warps = 8 num_experts = num_recv_tokens_per_expert.shape[0] hidden_size = recv_x.shape[1] @@ -1100,32 +1333,71 @@ def tma_align_input_scale(input_scale: torch.Tensor): @triton.jit -def compute_masked_m_triton_kernel(seg_indptr, masked_m): - expert_id = tl.program_id(0) - start = tl.load(seg_indptr + expert_id) - end = tl.load(seg_indptr + expert_id + 1) - tl.store(masked_m + expert_id, (end - start)) - - -@triton.jit -def deepgemm_compute_src2dst_triton_kernel( - topk_ids, - reorder_ids, - seg_indptr, - src2dst, +def fused_moe_dispatch_index_triton_kernel( + topk_ids_ptr, # flat (num_toks,) int32; -1 = padding (drives the `expert >= 0` gate) + src2dst_ptr, + masked_m_ptr, m_max, num_toks, + num_experts, BLOCK_SIZE: tl.constexpr, + ZERO_INIT: tl.constexpr, ): - pid = tl.program_id(axis=0) - dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = dst_id < num_toks - src_id = tl.load(reorder_ids + dst_id, mask=mask) - expert_id = tl.load(topk_ids + src_id, mask=(src_id < num_toks)) - expert_dst_start = tl.load(seg_indptr + expert_id, mask=(expert_id >= 0)) - expert_dst_offset = dst_id - expert_dst_start - dst_id = expert_id * m_max + expert_dst_offset - tl.store(src2dst + src_id, dst_id, mask=mask) + # Each token picks top_k distinct experts, so per-expert count <= num_tokens < m_max; + # dst = expert*m_max + offset never spills into the next expert's region. + pid = tl.program_id(0) + if ZERO_INIT: + # Zero the cursor in-kernel and barrier before any atomic_add (single-block path). + e_off = tl.arange(0, BLOCK_SIZE) + tl.store( + masked_m_ptr + e_off, + tl.zeros((BLOCK_SIZE,), dtype=tl.int32), + mask=e_off < num_experts, + ) + tl.debug_barrier() + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < num_toks + expert = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) + valid = mask & (expert >= 0) + # Clamp masked lanes to bin 0 so the masked atomic's pointer stays in-bounds. + expert_safe = tl.where(valid, expert, 0) + offset = tl.atomic_add(masked_m_ptr + expert_safe, 1, mask=valid) + dst = expert_safe * m_max + offset + tl.store(src2dst_ptr + offs, dst, mask=valid) + + +def fused_moe_dispatch_index( + topk_ids: torch.Tensor, + num_local_experts: int, + m_max: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + num_toks = topk_ids.numel() + src2dst = torch.empty(num_toks, device=topk_ids.device, dtype=torch.int32) + # masked_m doubles as the atomic cursor and the final per-expert count; must be zeroed before any atomic_add. + single_block = max(num_toks, num_local_experts) <= 1024 + if single_block: + BLOCK_SIZE = triton.next_power_of_2(max(num_toks, num_local_experts)) + masked_m = torch.empty( + num_local_experts, device=topk_ids.device, dtype=torch.int32 + ) + grid = (1,) + else: + BLOCK_SIZE = 256 + masked_m = torch.zeros( + num_local_experts, device=topk_ids.device, dtype=torch.int32 + ) + grid = (triton.cdiv(num_toks, BLOCK_SIZE),) + fused_moe_dispatch_index_triton_kernel[grid]( + topk_ids.view(-1), + src2dst, + masked_m, + m_max, + num_toks, + num_local_experts, + BLOCK_SIZE=BLOCK_SIZE, + ZERO_INIT=single_block, + ) + return masked_m, src2dst @triton.jit @@ -1139,8 +1411,12 @@ def fill_gateup_input_triton_kernel( topk, hidden_size, scale_size, + m_max, + scale_row_stride, + scale_col_stride, BLOCK_SIZE: tl.constexpr, - IS_FP8: tl.constexpr, + IS_FP8: tl.constexpr = True, + SCALE_MN_MAJOR: tl.constexpr = False, ): src_idx_int32 = tl.program_id(0) @@ -1149,7 +1425,7 @@ def fill_gateup_input_triton_kernel( topk_ids_ptr = topk_ids_ptr + src_idx * topk src_ptr = input_ptr + src_idx * hidden_size if IS_FP8: - scale_src_ptr = scale_ptr + src_idx * scale_size + scale_src_ptr = scale_ptr + src_idx * scale_row_stride vec = tl.arange(0, BLOCK_SIZE) for idx in range(topk): @@ -1165,12 +1441,28 @@ def fill_gateup_input_triton_kernel( tl.store(dst_ptr + offset, in_data, mask=mask) if IS_FP8: - scale_dst_ptr = gateup_input_scale_ptr + dst_idx * scale_size - for start_offset in tl.range(0, scale_size, BLOCK_SIZE): - offset = start_offset + vec - mask = offset < scale_size - in_scale = tl.load(scale_src_ptr + offset, mask=mask) - tl.store(scale_dst_ptr + offset, in_scale, mask=mask) + if SCALE_MN_MAJOR: + expert = dst_idx // m_max + m = dst_idx % m_max + scale_dst_ptr = ( + gateup_input_scale_ptr + expert * scale_size * m_max + m + ) + for start_offset in tl.range(0, scale_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < scale_size + in_scale = tl.load( + scale_src_ptr + offset * scale_col_stride, mask=mask + ) + tl.store(scale_dst_ptr + offset * m_max, in_scale, mask=mask) + else: + scale_dst_ptr = gateup_input_scale_ptr + dst_idx * scale_size + for start_offset in tl.range(0, scale_size, BLOCK_SIZE): + offset = start_offset + vec + mask = offset < scale_size + in_scale = tl.load( + scale_src_ptr + offset * scale_col_stride, mask=mask + ) + tl.store(scale_dst_ptr + offset, in_scale, mask=mask) def moe_ep_deepgemm_preprocess( @@ -1180,71 +1472,91 @@ def moe_ep_deepgemm_preprocess( top_k: int, block_shape, output_dtype: torch.dtype = torch.float8_e4m3fn, -): - reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) - seg_indptr = torch.zeros( - num_local_experts + 1, device=topk_ids.device, dtype=torch.int64 - ) - src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32) - masked_m = torch.empty(num_local_experts, device=topk_ids.device, dtype=torch.int32) - - compute_seg_indptr_triton_kernel[(num_local_experts + 1,)]( - reorder_topk_ids, seg_indptr, topk_ids.numel() - ) - - grid = lambda meta: (triton.cdiv(topk_ids.numel(), meta["BLOCK_SIZE"]),) - compute_masked_m_triton_kernel[(num_local_experts,)](seg_indptr, masked_m) - + use_mxfp8: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m}) https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/jit_kernels/m_grouped_gemm.py#L165 m_max = (hidden_states.size(0) // 256 + 1) * 256 expected_m = (topk_ids.numel() - 1) // num_local_experts + 1 + + masked_m, src2dst = fused_moe_dispatch_index(topk_ids, num_local_experts, m_max) + gateup_input = torch.empty( (num_local_experts, m_max, hidden_states.size(1)), device=hidden_states.device, dtype=output_dtype, ) - deepgemm_compute_src2dst_triton_kernel[grid]( - topk_ids, - reorder_ids, - seg_indptr, - src2dst, - m_max, - topk_ids.numel(), - BLOCK_SIZE=256, - ) - if block_shape is None: block_shape = [128, 128] assert len(block_shape) == 2 block_n, block_k = block_shape[0], block_shape[1] is_fp8 = output_dtype == torch.float8_e4m3fn - if is_fp8: - # TODO: fuse this with the preprocess - hidden_states, scale = per_token_group_quant_fp8(hidden_states, block_k) + if is_fp8 and use_mxfp8: + from sglang.jit_kernel.minimax_quant_ue8m0 import ( + per_token_quant_fp8_ue8m0_scatter, + ) + num_groups = hidden_states.size(1) // block_k + gateup_input_scale = torch.empty( + (gateup_input.size(0), num_groups // 4, m_max), + device=hidden_states.device, + dtype=torch.int32, + ) + per_token_quant_fp8_ue8m0_scatter( + hidden_states, + gateup_input, + gateup_input_scale, + src2dst, + topk_ids, + top_k, + m_max, + group_size=block_k, + ) + gateup_input_scale = gateup_input_scale.transpose(1, 2) + elif is_fp8: + hidden_states, scale = per_token_group_quant_fp8(hidden_states, block_k) gateup_input_scale = torch.empty( (gateup_input.size(0), gateup_input.size(1), scale.size(1)), device=hidden_states.device, dtype=scale.dtype, ) + fill_gateup_input_triton_kernel[(hidden_states.shape[0],)]( + hidden_states, + scale, + gateup_input, + gateup_input_scale, + src2dst, + topk_ids, + top_k, + hidden_states.size(1), + scale.size(1), + m_max, + scale.stride(0), + scale.stride(1), + BLOCK_SIZE=1024, + IS_FP8=True, + SCALE_MN_MAJOR=False, + ) else: scale = None gateup_input_scale = None - - fill_gateup_input_triton_kernel[(hidden_states.shape[0],)]( - hidden_states, - scale, - gateup_input, - gateup_input_scale, - src2dst, - topk_ids, - top_k, - hidden_states.size(1), - scale.size(1) if is_fp8 else 0, - BLOCK_SIZE=1024, - IS_FP8=is_fp8, - ) + fill_gateup_input_triton_kernel[(hidden_states.shape[0],)]( + hidden_states, + scale, + gateup_input, + gateup_input_scale, + src2dst, + topk_ids, + top_k, + hidden_states.size(1), + 0, + m_max, + 0, + 0, + BLOCK_SIZE=1024, + IS_FP8=False, + SCALE_MN_MAJOR=False, + ) return ( masked_m, diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 4c3fda886..36d388421 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -119,6 +119,17 @@ class DeepGemmMoeQuantInfo(MoeQuantInfo): block_shape: Optional[List[int]] = None # DSV4 mxfp4 layout flag; selects recipe_a=(1,128)/recipe_b=(1,32) downstream. is_fp4_experts: bool = False + use_mxfp8: bool = False + + def __post_init__(self): + if self.use_mxfp8: + assert self.block_shape == [ + 1, + 32, + ], f"MXFP8 requires block_shape [1, 32], got {self.block_shape}" + assert ( + deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + ), "MXFP8 requires DEEPGEMM_SCALE_UE8M0=True" class DeepGemmRunnerCore(MoeRunnerCore): @@ -377,12 +388,31 @@ class DeepGemmRunnerCore(MoeRunnerCore): w13_scale = quant_info.w13_scale w2_scale = quant_info.w2_scale - recipe_a, recipe_b = ( - ((1, 128), (1, 32)) if quant_info.is_fp4_experts else (None, None) - ) - hidden_states_device = running_state["hidden_states_device"] + use_mxfp8 = quant_info.use_mxfp8 + scale_block_size = quant_info.block_shape[1] if quant_info.block_shape else 128 + + if use_mxfp8: + recipe_b = tuple(quant_info.block_shape) + # gran_k is set by the dispatch path (standard=block_shape[1], DeepEP-LL=128), + # not inferable from K; inferring it silently mis-reads the activation scale. + gran_k_act = running_state.get( + "mxfp8_act_gran_k", quant_info.block_shape[1] + ) + _, _, k_for_recipe = hidden_states.shape + act_sf_last = hidden_states_scale.shape[-1] + assert ceil_div(k_for_recipe, gran_k_act * 4) == act_sf_last, ( + f"MXFP8 gateup scale mismatch: gran_k={gran_k_act}, K={k_for_recipe}, " + f"act_sf_last={act_sf_last}, expected " + f"{ceil_div(k_for_recipe, gran_k_act * 4)}" + ) + recipe_a = (quant_info.block_shape[0], gran_k_act) + elif quant_info.is_fp4_experts: + recipe_a, recipe_b = (1, 128), (1, 32) + else: + recipe_a, recipe_b = None, None + # GroupGemm-0 if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: if hidden_states_scale.dtype != torch.int: @@ -436,24 +466,58 @@ class DeepGemmRunnerCore(MoeRunnerCore): gateup_output, swiglu_limit=self.swiglu_limit ) gateup_output = einops.rearrange( - gateup_output, "(grp tok) hidden -> grp tok hidden", grp=num_groups + gateup_output, + "(grp tok) hidden -> grp tok hidden", + grp=num_groups, ) - # Act + # Act. + topk_ids_rs = running_state.get("topk_ids") + num_real_tokens = ( + topk_ids_rs.shape[0] + if ( + use_mxfp8 + and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + and topk_ids_rs is not None + and "src2dst" in running_state + ) + else None + ) down_input, down_input_scale = _varlen_deep_gemm_silu_mul_quant( gateup_output, masked_m, - group_size=128, + group_size=scale_block_size, topk=self.config.top_k, swiglu_limit=swiglu_limit_arg, swizzle=self.use_swizzle, + gemm1_alpha=self.config.gemm1_alpha, + gemm1_clamp_limit=self.config.gemm1_clamp_limit, + num_real_tokens=num_real_tokens, ) del gateup_output + # Down activation is quantised locally at scale_block_size (never DeepEP-LL), + # so its gran_k differs from gateup recipe_a. + recipe_a_down = recipe_a + if use_mxfp8: + recipe_a_down = (quant_info.block_shape[0], scale_block_size) + # GroupGemm-1 n = w2_weight.shape[1] - if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: + if ( + use_mxfp8 + and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + and down_input_scale.dtype != torch.int32 + ): + import deep_gemm.utils.layout + + down_input_scale = ( + deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor( + down_input_scale + ) + ) + elif deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: down_input_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor( down_input_scale ) @@ -481,7 +545,7 @@ class DeepGemmRunnerCore(MoeRunnerCore): down_output, masked_m, expected_m, - recipe_a=recipe_a, + recipe_a=recipe_a_down, recipe_b=recipe_b, **gemm_overlap_args_dict, ) @@ -600,6 +664,7 @@ def pre_permute_standard_to_deep_gemm( runner_config.top_k, quant_info.block_shape, output_dtype=output_dtype, + use_mxfp8=quant_info.use_mxfp8, ) ) @@ -611,6 +676,9 @@ def pre_permute_standard_to_deep_gemm( running_state["hidden_states_dtype"] = hidden_states_dtype running_state["hidden_states_device"] = hidden_states_device running_state["src2dst"] = src2dst + running_state["mxfp8_act_gran_k"] = ( + quant_info.block_shape[1] if quant_info.block_shape else 128 + ) return DeepGemmRunnerInput( hidden_states=hidden_states, @@ -628,7 +696,7 @@ def post_permute_deep_gemm_to_standard( runner_config: MoeRunnerConfig, running_state: dict, ) -> StandardCombineInput: - from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel + from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_deepgemm from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput hidden_states_shape = running_state["hidden_states_shape"] @@ -641,22 +709,23 @@ def post_permute_deep_gemm_to_standard( output = torch.empty( hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device ) - post_reorder_triton_kernel[(hidden_states_shape[0],)]( + post_reorder_deepgemm( runner_output.hidden_states, output, src2dst, topk_ids, topk_weights, runner_config.top_k, + hidden_states_shape[0], hidden_states_shape[1], - BLOCK_SIZE=512, + ( + runner_config.routed_scaling_factor + if runner_config.routed_scaling_factor is not None + else 1.0 + ), ) - dispose_tensor(runner_output.hidden_states) - if runner_config.routed_scaling_factor is not None: - output *= runner_config.routed_scaling_factor - return StandardCombineInput( hidden_states=output, ) @@ -678,6 +747,8 @@ def pre_permute_deepep_ll_to_deep_gemm( running_state["hidden_states_shape"] = hidden_states.shape running_state["hidden_states_dtype"] = hidden_states.dtype running_state["hidden_states_device"] = hidden_states.device + # DeepEP-LL FP8 dispatch quantises activations at a fixed 128 block, not the checkpoint block_shape. + running_state["mxfp8_act_gran_k"] = 128 return DeepGemmRunnerInput( hidden_states=hidden_states, @@ -833,6 +904,9 @@ def _varlen_deep_gemm_silu_mul_quant( topk: int, swiglu_limit: Optional[float] = None, swizzle: bool = False, + gemm1_alpha: Optional[float] = None, + gemm1_clamp_limit: Optional[float] = None, + num_real_tokens: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_post_quant_fwd from sglang.srt.layers.quantization.fp8_kernel import ( @@ -840,6 +914,9 @@ def _varlen_deep_gemm_silu_mul_quant( ) if _MASKED_GEMM_FAST_ACT: + assert ( + gemm1_alpha is None + ), "gemm1_alpha is not supported with SGLANG_MASKED_GEMM_FAST_ACT" assert not swizzle, ( "SGLANG_OPT_FIX_MEGA_MOE_MEMORY is incompatible with " "SGLANG_MASKED_GEMM_FAST_ACT (swizzled layout only supported by JIT act)" @@ -865,6 +942,42 @@ def _varlen_deep_gemm_silu_mul_quant( D = D_2 // 2 del D_2 G = D // group_size + + # Fused UE8M0 pack needs 4 groups per packed int32 (the G%4 and D guards below). + if ( + gemm1_alpha is not None + and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + and num_real_tokens is not None + and G % 4 == 0 + and D % (group_size * 4) == 0 + ): + from sglang.srt.layers.moe.ep_moe.kernels import ( + silu_and_mul_masked_post_quant_packed_fwd, + ) + + assert ( + swiglu_limit is None + ), "swiglu_limit and gemm1_alpha are mutually exclusive" + assert not swizzle, "swizzle is not supported with gemm1_alpha" + down_input = torch.empty( + (E, N, D), device=hidden_states_device, dtype=torch.float8_e4m3fn + ) + down_input_scale_packed = torch.empty( + (E, G // 4, N), device=hidden_states_device, dtype=torch.int32 + ) + silu_and_mul_masked_post_quant_packed_fwd( + gateup_output, + down_input, + down_input_scale_packed, + group_size, + masked_m, + num_real_tokens=num_real_tokens, + topk=topk, + gemm1_alpha=gemm1_alpha, + gemm1_clamp_limit=gemm1_clamp_limit or 0.0, + ) + return down_input, down_input_scale_packed.transpose(-1, -2) + down_input = torch.empty( (E, N, D), device=hidden_states_device, @@ -874,6 +987,8 @@ def _varlen_deep_gemm_silu_mul_quant( use_jit_ep_activation = envs.SGLANG_OPT_USE_JIT_EP_ACTIVATION.get() if N % 4 != 0 or G % 4 != 0 or D // 8 < E: use_jit_ep_activation = False + if gemm1_alpha is not None: + use_jit_ep_activation = False if use_jit_ep_activation: packed_ue8m0 = deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 @@ -897,12 +1012,18 @@ def _varlen_deep_gemm_silu_mul_quant( if packed_ue8m0: down_input_scale = down_input_scale.transpose(-1, -2) else: - assert ( - swiglu_limit is None - ), "swiglu_limit (DeepSeek V4) requires SGLANG_OPT_USE_JIT_EP_ACTIVATION=True" - assert ( - not swizzle - ), "SGLANG_OPT_FIX_MEGA_MOE_MEMORY requires SGLANG_OPT_USE_JIT_EP_ACTIVATION=True" + if gemm1_alpha is not None: + assert ( + swiglu_limit is None + ), "swiglu_limit and gemm1_alpha are mutually exclusive" + assert not swizzle, "swizzle is not supported with gemm1_alpha" + else: + assert ( + swiglu_limit is None + ), "swiglu_limit (DeepSeek V4) requires SGLANG_OPT_USE_JIT_EP_ACTIVATION=True" + assert ( + not swizzle + ), "SGLANG_OPT_FIX_MEGA_MOE_MEMORY requires SGLANG_OPT_USE_JIT_EP_ACTIVATION=True" down_input_scale = torch.empty( (E, N, G), device=hidden_states_device, @@ -915,6 +1036,8 @@ def _varlen_deep_gemm_silu_mul_quant( group_size, masked_m, scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + gemm1_alpha=gemm1_alpha or 0.0, + gemm1_clamp_limit=gemm1_clamp_limit or 0.0, ) return down_input, down_input_scale diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 650da2797..244f19e11 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -25,6 +25,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import ( ) from sglang.srt.layers.moe.topk import StandardTopKOutput, TopKOutput, TopKOutputChecker from sglang.srt.layers.moe.utils import ( + get_moe_a2a_backend, get_moe_runner_backend, should_use_flashinfer_cutlass_moe_fp4_allgather, ) @@ -84,7 +85,6 @@ assert isinstance(StandardCombineInput, CombineInput) class StandardDispatcher(BaseDispatcher): - def __init__(self, moe_runner_config: MoeRunnerConfig): super().__init__() self.moe_ep_size = get_parallel().moe_ep_size @@ -92,6 +92,11 @@ class StandardDispatcher(BaseDispatcher): self.enable_flashinfer_cutlass_moe = backend.is_flashinfer_cutlass() self.enable_flashinfer_mxfp4_moe = backend.is_flashinfer_mxfp4() self.enable_flashinfer_trtllm_routed_moe = backend.is_flashinfer_trtllm_routed() + # AITER fast paths can be on while the MoE runner stays Triton; only the + # AITER runner keeps global expert IDs, so Triton must remap to local range. + self.use_aiter_moe_runner = backend.is_aiter() or ( + backend.is_auto() and _use_aiter and get_moe_a2a_backend().supports_aiter() + ) # Skip local expert mapping when the backend handles EP with global expert IDs: # - cutlass / cutedsl / trtllm_routed handle EP internally # - mxfp4 dispatcher mapping is already global @@ -195,7 +200,7 @@ class StandardDispatcher(BaseDispatcher): ) if self.local_expert_mapping is not None and not self.skip_local_expert_mapping: - if _use_aiter: + if self.use_aiter_moe_runner: self.expert_mask_gpu = ( ( (self.local_expert_mapping >= 0) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 824d63463..a57fd612e 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -182,7 +182,7 @@ if _is_cuda or _is_hip or _is_xpu: from sglang.kernels.ops.moe import topk_softmax try: - from sgl_kernel import topk_sigmoid + from sglang.jit_kernel.moe_topk_sigmoid import topk_sigmoid except ImportError: pass if _use_aiter: @@ -748,6 +748,9 @@ def fused_topk( renormalize: bool, correction_bias: Optional[torch.Tensor] = None, scoring_func: str = "softmax", + routed_scaling_factor: Optional[float] = None, + apply_routed_scaling_factor_on_output: Optional[bool] = False, + num_fused_shared_experts: int = 0, packed_out: Optional[torch.Tensor] = None, num_token_non_padded: Optional[torch.Tensor] = None, ): @@ -825,6 +828,10 @@ def fused_topk( topk_group=1, need_renorm=renormalize, ) + if apply_routed_scaling_factor_on_output: + topk_weights *= ( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ) elif _is_cuda and envs.SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK.get(): # Unified Triton router (subsumes the AOT topk_sigmoid CUDA kernel). from sglang.jit_kernel.moe_fused_gate import ( @@ -846,14 +853,30 @@ def fused_topk( topk, scoring_func="sigmoid", renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, ) else: + if num_fused_shared_experts > 1: + raise ValueError( + "sigmoid topk supports at most one fused shared expert" + ) + scale = ( + routed_scaling_factor + if ( + apply_routed_scaling_factor_on_output + and routed_scaling_factor is not None + ) + else 1.0 + ) topk_sigmoid( topk_weights, topk_ids, gating_output, renormalize, correction_bias, + scale, + num_fused_shared_experts, ) else: raise ValueError(f"Invalid scoring function: {scoring_func}") @@ -873,10 +896,17 @@ def grouped_topk_gpu( num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, + scoring_func: str = "softmax", ): assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" - scores = torch.softmax(gating_output, dim=-1) + if scoring_func == "softmax": + scores = torch.softmax(gating_output, dim=-1) + elif scoring_func == "sigmoid": + scores = gating_output.sigmoid() + else: + raise ValueError(f"Unsupported scoring function: {scoring_func}") + num_token = scores.shape[0] num_experts = scores.shape[1] group_scores = ( @@ -2001,6 +2031,7 @@ def select_experts( num_fused_shared_experts=num_fused_shared_experts, routed_scaling_factor=routed_scaling_factor, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + scoring_func=scoring_func, ) else: topk_weights, topk_ids = biased_grouped_topk( @@ -2030,14 +2061,16 @@ def select_experts( scoring_func=scoring_func, ) elif custom_routing_function is None: - if scoring_func != "sqrtsoftplus": + if scoring_func not in ("sqrtsoftplus", "sigmoid"): assert not apply_routed_scaling_factor_on_output, "Not implemented" - if scoring_func == "sqrtsoftplus": + # Keep sigmoid flag-off byte-identical: only use the JIT gate when the flag is on. + use_jit_fused_gate = envs.SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK.get() + if scoring_func == "sqrtsoftplus" or ( + scoring_func == "sigmoid" and use_jit_fused_gate + ): _biased_topk = ( - biased_topk_jit_kernel_impl - if envs.SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK.get() - else biased_topk_impl + biased_topk_jit_kernel_impl if use_jit_fused_gate else biased_topk_impl ) topk_weights, topk_ids = _biased_topk( @@ -2110,6 +2143,9 @@ def select_experts( renormalize=renormalize, correction_bias=correction_bias, scoring_func=scoring_func, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, **_fused_topk_kwargs, ) else: diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 151f80e72..34d9d5539 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -94,6 +94,7 @@ from sglang.srt.utils import ( is_sm100_supported, is_sm120_supported, log_info_on_rank0, + mxfp8_block_convert_required, print_warning_once, set_weight_attrs, use_intel_amx_backend, @@ -113,6 +114,10 @@ _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() _is_fp8_fnuz = is_fp8_fnuz() +_is_gfx95_supported = is_gfx95_supported() +# gfx942 (MI300) has no MX matmul HW; MXFP8 checkpoints are converted to +# block-fp8 [128,128] at load and run through the native block-fp8 kernels. +_mxfp8_to_block_fp8_required = mxfp8_block_convert_required() _use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip _use_aiter = envs.SGLANG_USE_AITER.get() and _is_hip _is_shuffle_moe_mxfp4 = is_gfx95_supported() @@ -279,6 +284,10 @@ class Fp8Config(QuantizationConfig): return 0 # NPU bypasses CUDA capability checks if _is_musa: return 31 + if self.use_mxfp8 and _is_hip and _is_gfx95_supported: + return 95 + if self.use_mxfp8 and _mxfp8_to_block_fp8_required: + return 94 return 100 if self.use_mxfp8 else 80 @@ -307,10 +316,13 @@ class Fp8Config(QuantizationConfig): normalized.append(f"model.{base}") ignored_layers = normalized weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) - if use_mxfp8 and weight_block_size is not None: - logger.warning( - "MXFP8 ignoring incoming weight_block_size in config.json; it is fixed to [1, 32]." - ) + if use_mxfp8: + # MXFP8 (OCP) spec fixes block size to [1, 32]; ckpt field is metadata only. + if weight_block_size is not None and weight_block_size != [1, 32]: + logger.warning( + "MXFP8 overriding weight_block_size=%s from config.json -> [1, 32].", + weight_block_size, + ) weight_block_size = [1, 32] return cls( is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, @@ -429,9 +441,11 @@ class Fp8LinearMethod(LinearMethodBase): self.block_quant = ( self.use_mxfp8 or self.quant_config.weight_block_size is not None ) + self.convert_mxfp8_to_block = self.use_mxfp8 and _mxfp8_to_block_fp8_required + self.weight_block_size = self.quant_config.weight_block_size self.w8a8_block_fp8_linear = None self.w8a8_mxfp8_linear = None - if self.use_mxfp8: + if self.use_mxfp8 and not self.convert_mxfp8_to_block: self.w8a8_mxfp8_linear = dispatch_w8a8_mxfp8_linear() else: self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear() @@ -581,6 +595,31 @@ class Fp8LinearMethod(LinearMethodBase): layer.register_parameter("input_scale", None) def process_weights_after_loading_block_quant(self, layer: Module) -> None: + if self.convert_mxfp8_to_block: + from sglang.srt.layers.quantization.mxfp8_block_convert import ( + convert_mxfp8_weight_to_block_fp8, + ) + + qweight, scale = convert_mxfp8_weight_to_block_fp8( + layer.weight.data, layer.weight_scale_inv.data, block=128 + ) + layer.weight = Parameter(qweight, requires_grad=False) + layer.weight_scale_inv = Parameter(scale, requires_grad=False) + self.use_mxfp8 = False + self.convert_mxfp8_to_block = False + self.weight_block_size = [128, 128] + elif self.use_mxfp8: + # MXFP8 (e4m3fn + UE8M0) must NOT be fnuz-normalized; check before + # the fnuz branch since is_fp8_fnuz() is also True on gfx942. + if not self.is_checkpoint_fp8_serialized: + self._quantize_mxfp8_weights(layer) + return + # MXFP8 scales are stored as UE8M0 uint8; no requantization here. + # Keep parameter object to preserve weight_loader attrs for hot reload. + layer.weight_scale_inv.requires_grad_(False) + layer.weight_scale_inv.format_ue8m0 = True + self._process_mxfp8_linear_weight_scale(layer) + return # If ROCm, normalize the weights and scales to e4m3fnuz if _is_fp8_fnuz: # activation_scheme: dynamic @@ -599,16 +638,6 @@ class Fp8LinearMethod(LinearMethodBase): layer.weight_scale_inv.data, requires_grad=False ) return - elif self.use_mxfp8: - if not self.is_checkpoint_fp8_serialized: - self._quantize_mxfp8_weights(layer) - return - # MXFP8 scales are stored as UE8M0 uint8; no requantization here. - # Keep parameter object to preserve weight_loader attrs for hot reload. - layer.weight_scale_inv.requires_grad_(False) - layer.weight_scale_inv.format_ue8m0 = True - self._process_mxfp8_linear_weight_scale(layer) - return else: # Requantize block scales to UE8M0 when DeepGEMM is the active runner. use_deepgemm_runner = ( @@ -697,8 +726,32 @@ class Fp8LinearMethod(LinearMethodBase): "weight_scale_inv_swizzled", block_scale_interleave(scale_u8.contiguous()).contiguous(), ) + elif get_fp8_gemm_runner_backend().is_deep_gemm(): + from sglang.srt.layers.deep_gemm_wrapper.configurer import ( + DEEPGEMM_SCALE_UE8M0, + ) + + n, k = layer.weight.shape + scale_u8 = layer.weight_scale_inv.data + scale_fp32 = ( + (scale_u8.contiguous().view(-1).to(torch.int32) << 23) + .view(torch.float32) + .view(n, k // 32) + ) + if DEEPGEMM_SCALE_UE8M0: + # Pre-packed; GEMM must be called with disable_ue8m0_cast=True. + import deep_gemm.utils.layout + + scale_packed = ( + deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor( + scale_fp32 + ) + ) + else: + scale_packed = scale_fp32 + copy_or_rebind_param(layer, "weight_scale_inv_deepgemm", scale_packed) else: - # Triton path consumes canonical 2D UE8M0 scales directly. + # Triton path consumes canonical 2D UE8M0 uint8 scales directly. return def _quantize_mxfp8_weights(self, layer: Module) -> None: @@ -854,6 +907,27 @@ class Fp8LinearMethod(LinearMethodBase): weight_scale = layer.weight_scale_inv_swizzled elif backend.is_flashinfer_trtllm(): weight_scale = layer.weight_scale_inv_shuffled + elif get_fp8_gemm_runner_backend().is_deep_gemm(): + weight_scale = getattr( + layer, "weight_scale_inv_deepgemm", layer.weight_scale_inv + ) + if isinstance(x, tuple): + return self.w8a8_mxfp8_linear( + input=x[0], + weight=layer.weight, + weight_scale=weight_scale, + input_scale=x[1], + bias=bias, + weight_scale_fallback=layer.weight_scale_inv, + ) + return self.w8a8_mxfp8_linear( + input=x, + weight=layer.weight, + weight_scale=weight_scale, + input_scale=None, + bias=bias, + weight_scale_fallback=layer.weight_scale_inv, + ) else: weight_scale = layer.weight_scale_inv if isinstance(x, tuple): @@ -878,7 +952,7 @@ class Fp8LinearMethod(LinearMethodBase): x, layer.weight, layer.weight_scale_inv, - self.quant_config.weight_block_size, + self.weight_block_size, bias, x.dtype, True, # is_vnni @@ -888,7 +962,7 @@ class Fp8LinearMethod(LinearMethodBase): return self.w8a8_block_fp8_linear( input=x[0], weight=layer.weight, - block_size=self.quant_config.weight_block_size, + block_size=self.weight_block_size, weight_scale=layer.weight_scale_inv, input_scale=x[1], bias=bias, @@ -897,7 +971,7 @@ class Fp8LinearMethod(LinearMethodBase): return self.w8a8_block_fp8_linear( input=x, weight=layer.weight, - block_size=self.quant_config.weight_block_size, + block_size=self.weight_block_size, weight_scale=layer.weight_scale_inv, input_scale=None, bias=bias, @@ -933,6 +1007,8 @@ class Fp8MoEMethod(FusedMoEMethodBase): self.block_quant = ( self.use_mxfp8 or self.quant_config.weight_block_size is not None ) + self.convert_mxfp8_to_block = self.use_mxfp8 and _mxfp8_to_block_fp8_required + self.weight_block_size = self.quant_config.weight_block_size self.is_fp4_expert = self.quant_config.is_fp4_experts self.dequant_fp4_to_fp8 = self.quant_config.dequant_fp4_to_fp8 self.with_bias = False @@ -1340,6 +1416,52 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer.w2_weight.is_shuffled = is_shuffled return + if self.convert_mxfp8_to_block: + # Only aiter-shuffle when the MoE runner is aiter; the triton runner + # consumes un-shuffled weights (shuffling the wrong runner corrupts output). + self._convert_mxfp8_moe_to_block_fp8(layer) + self.use_mxfp8 = False + self.convert_mxfp8_to_block = False + self.weight_block_size = [128, 128] + if _is_fp8_fnuz: + w13_weight, w13_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.w13_weight, + weight_scale=layer.w13_weight_scale_inv, + input_scale=None, + ) + w2_weight, w2_weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.w2_weight, + weight_scale=layer.w2_weight_scale_inv, + input_scale=None, + ) + layer.w13_weight = Parameter(w13_weight, requires_grad=False) + layer.w13_weight_scale_inv = Parameter( + w13_weight_scale, requires_grad=False + ) + layer.w2_weight = Parameter(w2_weight, requires_grad=False) + layer.w2_weight_scale_inv = Parameter( + w2_weight_scale, requires_grad=False + ) + layer.w13_input_scale = None + layer.w2_input_scale = None + runner_is_aiter = ( + getattr(self, "runner", None) is not None + and self.runner.runner_backend.is_aiter() + ) + if _use_aiter and runner_is_aiter: + layer.w13_weight.data = shuffle_weight( + layer.w13_weight.contiguous(), (16, 16) + ) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight.contiguous(), (16, 16) + ) + return + elif self.use_mxfp8: + self._process_mxfp8_moe_weights( + layer, quantize=not self.quant_config.is_checkpoint_fp8_serialized + ) + return + # If ROCm, normalize the weights and scales to e4m3fnuz if _is_fp8_fnuz: # activation_scheme: dynamic @@ -1384,10 +1506,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): _is_cpu_amx_available ), "Fp8MoEMethod on CPU requires that CPU has AMX support" _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) - elif self.use_mxfp8: - self._process_mxfp8_moe_weights( - layer, quantize=not self.quant_config.is_checkpoint_fp8_serialized - ) else: # For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0 from sglang.srt.layers import deep_gemm_wrapper @@ -1472,10 +1590,41 @@ class Fp8MoEMethod(FusedMoEMethodBase): use_deepgemm_runner=True, ) + def _convert_mxfp8_moe_to_block_fp8(self, layer: Module) -> None: + from sglang.srt.layers.quantization.mxfp8_block_convert import ( + convert_mxfp8_weight_to_block_fp8, + ) + + def convert(w, s): + E, N, K = w.shape + qw = torch.empty_like(w) + sn = (N + 127) // 128 + sk = (K + 127) // 128 + scale = torch.empty((E, sn, sk), dtype=torch.float32, device=w.device) + for e in range(E): + qe, se = convert_mxfp8_weight_to_block_fp8(w[e], s[e], block=128) + qw[e] = qe + scale[e] = se + return qw, scale + + w13_q, w13_s = convert(layer.w13_weight.data, layer.w13_weight_scale_inv.data) + w2_q, w2_s = convert(layer.w2_weight.data, layer.w2_weight_scale_inv.data) + layer.w13_weight = Parameter(w13_q, requires_grad=False) + layer.w2_weight = Parameter(w2_q, requires_grad=False) + layer.w13_weight_scale_inv = Parameter(w13_s, requires_grad=False) + layer.w2_weight_scale_inv = Parameter(w2_s, requires_grad=False) + layer.w13_input_scale = None + layer.w2_input_scale = None + def _process_mxfp8_moe_weights(self, layer: Module, quantize: bool = True) -> None: - if not (_is_cuda and is_sm100_supported()): - raise RuntimeError("MXFP8 MoE quantization requires SM100.") + if not ( + (_is_cuda and is_sm100_supported()) or (_is_hip and _is_gfx95_supported) + ): + raise RuntimeError( + "MXFP8 MoE quantization requires SM100 or ROCm gfx95 " + "(gfx942 converts MXFP8 to block-fp8 at load instead)." + ) def _quantize_and_swizzle_with_cutlass_es_kernel(weight: torch.Tensor): from sgl_kernel import es_sm100_mxfp8_blockscaled_grouped_quant @@ -1545,8 +1694,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): scale = scale.view(num_experts, aligned_m, k // 32) num_warps = 8 scale = _swizzle_mxfp8_sf(scale, num_warps) - scale = scale.data.view(num_experts, aligned_m, k // 32) - return scale + # convert_layout may pad for alignment; we can't view back to the + # unpadded shape, so return the (possibly padded) swizzled tensor. + return scale.data def _quantize_and_swizzle_with_triton_kernel(weight: torch.Tensor): @@ -1573,14 +1723,81 @@ class Fp8MoEMethod(FusedMoEMethodBase): ) return qweight.view_as(weight), scale_u8 + from sglang.srt.layers.quantization.mxfp8_block_convert import ( + _ue8m0_to_fp32, + ) + + def _quantize_for_deepgemm(weight: torch.Tensor): + weight = weight.contiguous() + num_experts, m, k = weight.shape + assert k % 32 == 0, f"{k=} must be divisible by 32 for MXFP8" + + weight_flat = weight.view(-1, k).contiguous() + qweight, scale_u8 = mxfp8_group_quantize(weight_flat) + qweight = qweight.view_as(weight) + scale_fp32 = _ue8m0_to_fp32(scale_u8).view(num_experts, m, k // 32) + scale_packed = _pack_moe_scale_for_deepgemm(scale_fp32) + return qweight, scale_packed + + def _pack_moe_scale_for_deepgemm(scale_fp32: torch.Tensor) -> torch.Tensor: + """Blackwell: int32 MN-major TMA-packed. Hopper returns fp32 (FP4 API converts).""" + from sglang.srt.layers.deep_gemm_wrapper.configurer import ( + DEEPGEMM_SCALE_UE8M0, + ) + + if DEEPGEMM_SCALE_UE8M0: + import deep_gemm.utils.layout + + return ( + deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor( + scale_fp32 + ) + ) + return scale_fp32 + + def _convert_ue8m0_scales_for_deepgemm( + scale_u8: torch.Tensor, shape: tuple + ) -> torch.Tensor: + num_experts, m, k_groups = shape[0], shape[1], scale_u8.shape[-1] + scale_fp32 = _ue8m0_to_fp32(scale_u8.contiguous().view(-1)).view( + num_experts, m, k_groups + ) + return _pack_moe_scale_for_deepgemm(scale_fp32) + if quantize: - if get_moe_runner_backend().is_cutlass(): + if _is_hip: + w13_q, w13_s_u8 = mxfp8_group_quantize( + layer.w13_weight.data.contiguous().view( + -1, layer.w13_weight.data.shape[-1] + ) + ) + w2_q, w2_s_u8 = mxfp8_group_quantize( + layer.w2_weight.data.contiguous().view( + -1, layer.w2_weight.data.shape[-1] + ) + ) + w13_q = w13_q.view_as(layer.w13_weight.data) + w2_q = w2_q.view_as(layer.w2_weight.data) + w13_s = w13_s_u8.view( + layer.w13_weight.data.shape[0], + layer.w13_weight.data.shape[1], + layer.w13_weight.data.shape[2] // 32, + ) + w2_s = w2_s_u8.view( + layer.w2_weight.data.shape[0], + layer.w2_weight.data.shape[1], + layer.w2_weight.data.shape[2] // 32, + ) + elif get_moe_runner_backend().is_cutlass(): w13_q, w13_s = _quantize_and_swizzle_with_cutlass_es_kernel( layer.w13_weight.data ) w2_q, w2_s = _quantize_and_swizzle_with_cutlass_es_kernel( layer.w2_weight.data ) + elif get_moe_runner_backend().is_deep_gemm(): + w13_q, w13_s = _quantize_for_deepgemm(layer.w13_weight.data) + w2_q, w2_s = _quantize_for_deepgemm(layer.w2_weight.data) elif ( get_moe_runner_backend().is_flashinfer_trtllm() or get_moe_runner_backend().is_flashinfer_trtllm_routed() @@ -1598,7 +1815,12 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer.w2_weight.data ) else: - if ( + if _is_hip: + w13_q = layer.w13_weight.data + w2_q = layer.w2_weight.data + w13_s = layer.w13_weight_scale_inv.data + w2_s = layer.w2_weight_scale_inv.data + elif ( get_moe_runner_backend().is_flashinfer_trtllm() or get_moe_runner_backend().is_flashinfer_trtllm_routed() ): @@ -1606,6 +1828,15 @@ class Fp8MoEMethod(FusedMoEMethodBase): w2_q = layer.w2_weight.data w13_s = layer.w13_weight_scale_inv.data w2_s = layer.w2_weight_scale_inv.data + elif get_moe_runner_backend().is_deep_gemm(): + w13_q = layer.w13_weight.data + w2_q = layer.w2_weight.data + w13_s = _convert_ue8m0_scales_for_deepgemm( + layer.w13_weight_scale_inv.data, layer.w13_weight.data.shape + ) + w2_s = _convert_ue8m0_scales_for_deepgemm( + layer.w2_weight_scale_inv.data, layer.w2_weight.data.shape + ) else: w13_q = layer.w13_weight.data w2_q = layer.w2_weight.data @@ -1887,12 +2118,14 @@ class Fp8MoEMethod(FusedMoEMethodBase): pass def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: + use_rocm_mxfp8 = self.use_mxfp8 and _is_hip and _is_gfx95_supported return TritonMoeQuantInfo( w13_weight=layer.w13_weight, w2_weight=layer.w2_weight, b13=getattr(layer, "w13_weight_bias", None), b2=getattr(layer, "w2_weight_bias", None), - use_fp8_w8a8=True, + use_mxfp8=use_rocm_mxfp8, + use_fp8_w8a8=not use_rocm_mxfp8, w13_scale=( layer.w13_weight_scale_inv if self.block_quant @@ -1903,7 +2136,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): ), a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, - block_shape=self.quant_config.weight_block_size, + block_shape=self.weight_block_size, ) def apply( @@ -2070,6 +2303,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): w2_scale=w2_scale, block_shape=block_shape, is_fp4_experts=self.is_fp4_expert, + use_mxfp8=self.use_mxfp8, ) elif ( self.runner.runner_backend.is_flashinfer_trtllm() diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 0f2b7eef9..51ef12b38 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -29,6 +29,7 @@ except: pass from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.utils import ( ceil_align, @@ -146,6 +147,17 @@ def deep_gemm_fp8_fp8_bf16_nt( deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C) +@register_custom_op(mutates_args=["C"]) +def deep_gemm_mxfp8_fp8_bf16_nt( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, +) -> None: + deep_gemm_wrapper.gemm_nt_mxfp8_f8f8bf16((A, As), (B, Bs), C) + + @triton.jit def _per_token_group_quant_8bit( # Pointers to inputs and output @@ -334,7 +346,6 @@ def _per_token_group_quant_8bit_raw( if scale_ue8m0: from deep_gemm import transform_sf_into_required_layout - assert group_size == 128 x_s = transform_sf_into_required_layout( x_s, num_groups=None, @@ -404,7 +415,6 @@ def _per_token_group_quant_8bit_fuse_silu_and_mul( scale_ue8m0=scale_ue8m0, ) - assert group_size == 128 output_scale = transform_sf_into_required_layout( output_scale_for_kernel, num_groups=output.shape[0], @@ -465,7 +475,7 @@ def create_per_token_group_quant_fp8_output_scale( if scale_ue8m0: if column_major_scales and scale_tma_aligned: *x_batch, x_q_mn, x_q_k = x_shape - x_s_mn, x_s_k = x_q_mn, x_q_k // 128 + x_s_mn, x_s_k = x_q_mn, x_q_k // group_size aligned_mn = ceil_align(x_s_mn, 4) aligned_k = ceil_align(x_s_k, 4) # TODO(FIXME): Fix cuda kernel and recover here to empty. @@ -510,6 +520,105 @@ def create_per_token_group_quant_fp8_output_scale( ) +_V2_KERNEL_SUPPORTED_GROUP_SIZES = (16, 32, 64, 128) + + +def _run_per_token_group_quant_8bit_kernel( + x: torch.Tensor, + x_q: torch.Tensor, + x_s: torch.Tensor, + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + *, + column_major_scales: bool, + scale_ue8m0: bool, + fuse_silu_and_mul: bool, + masked_m: Optional[torch.Tensor], + enable_v2: Optional[bool], +) -> None: + # V1 JIT (.cuh) is byte-identical to V2 but CUDA-only and col-major-UE8M0-only; + # gate it to opt-in plain-2D non-MUSA calls, else fall back to V2 / AOT v1. + if enable_v2 is None: + enable_v2 = group_size in _V2_KERNEL_SUPPORTED_GROUP_SIZES or _is_musa + + use_jit_per_token_group_v1_quant = ( + envs.SGLANG_OPT_USE_JIT_PER_TOKEN_GROUP_QUANT.get() + and enable_v2 + and not _is_musa + and not fuse_silu_and_mul + and masked_m is None + and x.dim() == 2 + and group_size in _V2_KERNEL_SUPPORTED_GROUP_SIZES + and not (scale_ue8m0 and not column_major_scales) + ) + + if use_jit_per_token_group_v1_quant: + sgl_per_token_group_quant_8bit_jit( + input=x, + output_q=x_q, + output_s=x_s, + group_size=group_size, + eps=eps, + fp8_min=fp8_min, + fp8_max=fp8_max, + scale_ue8m0=scale_ue8m0, + ) + return + + if not enable_sgl_per_token_group_quant_8bit: + assert not enable_v2 + sgl_per_token_group_quant_fp8( + x, x_q, x_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0 + ) + return + + if enable_v2 and _is_musa: + # JIT v2 .cuh is CUDA-only (no MUSA fallback); AOT v2 carries the USE_MUSA path. + sgl_per_token_group_quant_8bit( + x, + x_q, + x_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + fuse_silu_and_mul, + masked_m, + enable_v2=True, + ) + elif enable_v2: + sgl_per_token_group_quant_8bit_jit_v2( + x, + x_q, + x_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0=scale_ue8m0, + fuse_silu_and_mul=fuse_silu_and_mul, + masked_m=masked_m, + ) + else: + # JIT kernels static_assert on group_size in {16,32,64,128}; keep AOT v1 otherwise. + sgl_per_token_group_quant_8bit( + x, + x_q, + x_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + fuse_silu_and_mul, + masked_m, + enable_v2=enable_v2, + ) + + def sglang_per_token_group_quant_fp8( x: torch.Tensor, group_size: int, @@ -538,60 +647,21 @@ def sglang_per_token_group_quant_fp8( scale_ue8m0=scale_ue8m0, ) - # Enable v2 kernel by default on supported group sizes - _V2_KERNEL_SUPPORTED_GROUP_SIZES = [16, 32, 64, 128] - if enable_v2 is None: - enable_v2 = group_size in _V2_KERNEL_SUPPORTED_GROUP_SIZES or _is_musa - if x.shape[0] > 0: - # Temporary - if enable_sgl_per_token_group_quant_8bit: - if enable_v2 and _is_musa: - # The JIT v2 .cuh uses CUDA-only inline PTX (ld/st.global.v4) and - # has no MUSA fallback, so keep MUSA on the AOT v2 op, which - # carries the USE_MUSA vector load/store fallbacks. - sgl_per_token_group_quant_8bit( - x, - x_q, - x_s, - group_size, - eps, - fp8_min, - fp8_max, - scale_ue8m0, - fuse_silu_and_mul, - masked_m, - enable_v2=True, - ) - elif enable_v2: - sgl_per_token_group_quant_8bit_jit_v2( - x, - x_q, - x_s, - group_size, - eps, - fp8_min, - fp8_max, - scale_ue8m0=scale_ue8m0, - fuse_silu_and_mul=fuse_silu_and_mul, - masked_m=masked_m, - ) - else: - sgl_per_token_group_quant_8bit_jit( - input=x, - output_q=x_q, - output_s=x_s, - group_size=group_size, - eps=eps, - fp8_min=fp8_min, - fp8_max=fp8_max, - scale_ue8m0=scale_ue8m0, - ) - else: - assert not enable_v2 - sgl_per_token_group_quant_fp8( - x, x_q, x_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0 - ) + _run_per_token_group_quant_8bit_kernel( + x, + x_q, + x_s, + group_size, + eps, + fp8_min, + fp8_max, + column_major_scales=column_major_scales, + scale_ue8m0=scale_ue8m0, + fuse_silu_and_mul=fuse_silu_and_mul, + masked_m=masked_m, + enable_v2=enable_v2, + ) return x_q, x_s @@ -609,7 +679,8 @@ def sglang_per_token_group_quant_fp8_row_padded( and scales_a). Allocating the quant outputs with rows already aligned to ``row_alignment`` makes the wrapper's pad_tensor() short-circuit (pad_rows == 0), removing 2x fill + 2x cat kernels per GEMM. Rows in [m, m_pad) are - uninitialized garbage; the caller must slice the GEMM output back to m. + zero-filled to match the legacy pad_tensor contract so the padded GEMM is + bit-exact; the caller still slices the GEMM output back to m. """ assert x.dim() == 2, "row-padded quant expects a 2D input" assert ( @@ -633,19 +704,37 @@ def sglang_per_token_group_quant_fp8_row_padded( (k // group_size, m_pad), device=x.device, dtype=torch.float32 ).transpose(0, 1) if m > 0: - sgl_per_token_group_quant_8bit( - x, - x_q[:m], - x_s[:m], - group_size, - eps, - fp8_min, - fp8_max, - False, # scale_ue8m0 - False, # fuse_silu_and_mul - None, # masked_m - enable_v2=True, - ) + # V1 JIT (.cuh) is CUDA-only; MUSA must stay on the AOT v2 op below. + if envs.SGLANG_OPT_USE_JIT_PER_TOKEN_GROUP_QUANT.get() and not _is_musa: + sgl_per_token_group_quant_8bit_jit( + input=x, + output_q=x_q[:m], + output_s=x_s[:m], + group_size=group_size, + eps=eps, + fp8_min=fp8_min, + fp8_max=fp8_max, + scale_ue8m0=False, + ) + else: + sgl_per_token_group_quant_8bit( + x, + x_q[:m], + x_s[:m], + group_size, + eps, + fp8_min, + fp8_max, + False, # scale_ue8m0 + False, # fuse_silu_and_mul + None, # masked_m + enable_v2=True, + ) + if m_pad != m: + # Tail rows feed the cutlass GEMM's padded region; zero them so the padded + # GEMM stays bit-exact with the legacy pad_tensor path (torch.empty is garbage). + x_q[m:].zero_() + x_s[m:].zero_() return x_q, x_s @@ -1294,6 +1383,25 @@ def w8a8_block_fp8_matmul_deepgemm( return C +def w8a8_mxfp8_matmul_deepgemm( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + assert A.is_contiguous() and B.is_contiguous() + assert output_dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + + M = A.numel() // A.shape[-1] + N, K = B.shape + C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype) + + deep_gemm_mxfp8_fp8_bf16_nt(A, As, B, Bs, C) + + return C + + def w8a8_block_fp8_matmul_triton( A: torch.Tensor, B: torch.Tensor, @@ -1537,6 +1645,66 @@ def mxfp8_block_scaled_matmul_triton( return output +@triton.jit +def _pack_mxfp8_scales_kernel( + scale_ptr, + out_ptr, + M: tl.constexpr, + K_GROUPS: tl.constexpr, + SCALE_K: tl.constexpr, + TOTAL: tl.constexpr, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < TOTAL + + idx256 = offs % 256 + tmp = offs // 256 + two = tmp % 2 + tmp = tmp // 2 + scale_k = tmp % SCALE_K + scale_m = tmp // SCALE_K + + within = two * 256 + idx256 + row_inner_32 = within // 16 + rem = within - row_inner_32 * 16 + row_outer_4 = rem // 4 + k_inner_4 = rem - row_outer_4 * 4 + + row = scale_m * 128 + row_outer_4 * 32 + row_inner_32 + col = scale_k * 4 + k_inner_4 + value = tl.load(scale_ptr + row * K_GROUPS + col, mask & (row < M), other=127) + tl.store(out_ptr + offs, value, mask) + + +def pack_mxfp8_scales_triton(scale_u8: torch.Tensor) -> torch.Tensor: + assert scale_u8.dim() == 2, f"Expected 2D scale tensor, got {scale_u8.dim()}D" + scale_u8 = scale_u8.contiguous() + m, k_groups = scale_u8.shape + assert ( + k_groups % 4 == 0 + ), f"{k_groups=} must be divisible by 4 (K must be multiple of 128)" + + scale_m = triton.cdiv(m, 128) + scale_k = k_groups // 4 + out = torch.empty( + (1, scale_m, scale_k, 2, 256), dtype=scale_u8.dtype, device=scale_u8.device + ) + total = out.numel() + block = 1024 + grid = (triton.cdiv(total, block),) + _pack_mxfp8_scales_kernel[grid]( + scale_u8, + out, + m, + k_groups, + scale_k, + total, + BLOCK=block, + ) + return out + + @triton.jit def _per_tensor_quant_mla_fp8_stage1( x_ptr, diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index b96c90b97..3e6ee25d2 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -25,6 +25,7 @@ from sglang.srt.layers.quantization.fp8_kernel import ( fp8_min, is_fp8_fnuz, mxfp8_block_scaled_matmul_triton, + pack_mxfp8_scales_triton, per_token_group_quant_fp8, scaled_fp8_quant, sglang_per_token_quant_fp8, @@ -449,11 +450,67 @@ def dispatch_w8a8_block_fp8_linear() -> Callable: def dispatch_w8a8_mxfp8_linear() -> Callable: backend = get_fp8_gemm_runner_backend() - if backend.is_flashinfer_cutlass() or backend.is_flashinfer_trtllm(): + if backend.is_deep_gemm(): + return _deepgemm_w8a8_mxfp8_linear_with_fallback + elif backend.is_flashinfer_cutlass() or backend.is_flashinfer_trtllm(): return flashinfer_mxfp8_blockscaled_linear + elif backend.is_triton(): + return triton_mxfp8_blockscaled_linear + elif _is_hip and _is_gfx95_supported: + from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import ( + dot_scaled_mxfp8_blockscaled_linear, + ) + + return dot_scaled_mxfp8_blockscaled_linear return triton_mxfp8_blockscaled_linear +def _deepgemm_w8a8_mxfp8_linear_with_fallback( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + weight_scale_fallback: Optional[torch.Tensor] = None, +) -> torch.Tensor: + from sglang.srt.layers.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + w8a8_mxfp8_matmul_deepgemm, + ) + + assert input_scale is None + output_dtype = input.dtype + + shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 + dtype_supported = output_dtype == torch.bfloat16 + + if not (shape_supported and dtype_supported): + return triton_mxfp8_blockscaled_linear( + input, weight, weight_scale_fallback, input_scale, bias + ) + + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = sglang_per_token_group_quant_fp8( + input_2d, + 32, + column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + + # weight_scale format is set per-backend by _process_mxfp8_linear_weight_scale + # (int32 packed TMA-aligned on Blackwell, float32 on Hopper); NOT uint8 — Triton form is routed to the fallback above. + + output = w8a8_mxfp8_matmul_deepgemm( + q_input, weight, x_scale, weight_scale, output_dtype=output_dtype + ) + if bias is not None: + output += bias + return output.to(dtype=output_dtype).view(*output_shape) + + def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: """Dispatch based on explicitly selected backend.""" if backend.is_flashinfer_trtllm(): @@ -968,6 +1025,14 @@ def mxfp8_group_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: def _pack_mxfp8_scales(scale_u8: torch.Tensor) -> torch.Tensor: + if ( + _is_hip + and _is_gfx95_supported + and scale_u8.is_cuda + and scale_u8.shape[0] % 128 != 0 + ): + return pack_mxfp8_scales_triton(scale_u8) + # Pack (M, K//32) UE8M0 scales into the layout expected by tl.dot_scaled. assert scale_u8.dim() == 2, f"Expected 2D scale tensor, got {scale_u8.dim()}D" scale_u8 = scale_u8.contiguous() @@ -1035,8 +1100,13 @@ def _raw_triton_mxfp8_blockscaled_linear( bias: Optional[torch.Tensor] = None, output_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: - if not (_is_cuda and (_is_sm100_supported or _is_sm120_supported)): - raise RuntimeError("MXFP8 dense linear requires Blackwell GPUs (SM100/SM120).") + if not ( + (_is_cuda and (_is_sm100_supported or _is_sm120_supported)) + or (_is_hip and _is_gfx95_supported) + ): + raise RuntimeError( + "MXFP8 dense linear requires Blackwell GPUs (SM100/SM120) or ROCm gfx95." + ) input_2d = input.view(-1, input.shape[-1]).contiguous() output_shape = [*input.shape[:-1], weight.shape[0]] @@ -1052,6 +1122,13 @@ def _raw_triton_mxfp8_blockscaled_linear( assert n % block_n == 0, f"{n=} must be divisible by {block_n}" assert weight.dtype == torch.float8_e4m3fn, "MXFP8 weight must be FP8 E4M3." assert weight_scale.dtype == torch.uint8, "MXFP8 weight_scale must be UE8M0 uint8." + assert weight_scale.dim() in ( + 2, + 5, + ), ( + "MXFP8 weight_scale must be canonical 2D or packed 5D, " + f"got {weight_scale.dim()}D." + ) if input_scale is None: q_input, x_scale_u8 = mxfp8_group_quantize(input_2d) @@ -1085,7 +1162,11 @@ def _raw_triton_mxfp8_blockscaled_linear( x_scale_u8 = torch.cat([x_scale_u8, pad_scale], dim=0) a_scale_packed = _pack_mxfp8_scales(x_scale_u8) - b_scale_packed = _pack_mxfp8_scales(weight_scale) + b_scale_packed = ( + weight_scale.contiguous() + if weight_scale.dim() == 5 + else _pack_mxfp8_scales(weight_scale) + ) num_stages = 1 if _is_sm120_supported else (4 if _is_sm100_supported else 1) output = triton_mxfp8_block_scaled_matmul( diff --git a/python/sglang/srt/layers/quantization/mxfp8_block_convert.py b/python/sglang/srt/layers/quantization/mxfp8_block_convert.py new file mode 100644 index 000000000..47e84dabf --- /dev/null +++ b/python/sglang/srt/layers/quantization/mxfp8_block_convert.py @@ -0,0 +1,72 @@ +"""Convert MXFP8 weights to block-fp8 [128,128] for AMD gfx942 (CDNA3 / MI300). + +gfx942 has no hardware MX-scaled matmul: Triton's ``tl.dot_scaled`` fails to +lower and the gfx950 ``mfma_scale`` intrinsics are unavailable. So MXFP8 +checkpoints (e4m3fn weights + 1x32 UE8M0 scales) are converted at load time to +block-wise FP8 [128,128] (e4m3fn + fp32 scales), which runs through SGLang's +native DeepSeek-V3 block-fp8 kernels (aiter / triton). The conversion is: + + bf16 = e4m3.to(f32) * exp2(ue8m0_scale.to(f32) - 127.0) # dequant 1x32 + block-fp8 = per-128x128-block quantize(bf16) # requant 128x128 +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +MXFP8_BLOCK_SIZE = 32 + + +def _ue8m0_to_fp32(scale_u8: torch.Tensor) -> torch.Tensor: + """UE8M0 uint8 (biased exponent, bias 127) -> fp32 multiplier 2^(v-127).""" + return (scale_u8.to(torch.int32) << 23).view(torch.float32) + + +def dequant_mxfp8_2d_to_bf16( + weight: torch.Tensor, scale_u8: torch.Tensor +) -> torch.Tensor: + """Dequant a 2D MXFP8 tensor (e4m3fn + 1x32 UE8M0 scales) to bf16. + + weight: [N, K] float8_e4m3fn; scale_u8: [N, K//32] uint8. + """ + n, k = weight.shape + descale = _ue8m0_to_fp32(scale_u8).unsqueeze(-1) # [N, K//32, 1] + deq = weight.to(torch.float32).view(n, k // MXFP8_BLOCK_SIZE, MXFP8_BLOCK_SIZE) + return (deq * descale).view(n, k).to(torch.bfloat16) + + +def bf16_to_block_fp8_128( + weight: torch.Tensor, block: int = 128 +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D bf16/fp32 weight to block-wise FP8 (e4m3fn) + fp32 scales. + + Returns (qweight [N,K] e4m3fn, scale [ceil(N/block), ceil(K/block)] fp32). + Mirrors the DeepSeek-V3 block-fp8 contract (divide by e4m3fn max 448). + The downstream gfx942 path normalizes e4m3fn -> e4m3fnuz separately. + """ + n, k = weight.shape + pn = ((n + block - 1) // block) * block + pk = ((k + block - 1) // block) * block + xp = torch.zeros((pn, pk), dtype=torch.float32, device=weight.device) + xp[:n, :k] = weight.to(torch.float32) + xv = xp.view(pn // block, block, pk // block, block) + amax = xv.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-4) + sf = amax / 448.0 + xq = (xv / sf).to(torch.float8_e4m3fn) + qweight = xq.view(pn, pk)[:n, :k].contiguous() + scale = sf.view(pn // block, pk // block).contiguous() + return qweight, scale + + +def convert_mxfp8_weight_to_block_fp8( + weight: torch.Tensor, scale_u8: torch.Tensor, block: int = 128 +) -> Tuple[torch.Tensor, torch.Tensor]: + """MXFP8 (e4m3fn + 1x32 UE8M0) -> block-fp8 [block,block] (e4m3fn + fp32). + + Used on gfx942 to run MXFP8 checkpoints through the fast native block-fp8 + kernels. + """ + bf16 = dequant_mxfp8_2d_to_bf16(weight, scale_u8) + return bf16_to_block_fp8_128(bf16, block=block) diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 7eaf30533..854a9eef4 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -35,6 +35,22 @@ from sglang.srt.utils.custom_op import register_custom_op _is_hip = is_hip() + +def _zero_padded_pcg_tail(buf: torch.Tensor, context) -> None: + """Zero the padded tail ``buf`` leaves as torch.empty garbage under PCG + replay, so NaN/Inf cannot reach residual / MoE routing / allreduce.""" + pcg_static_tokens = context.num_tokens + actual_tokens = context.raw_num_tokens + if ( + pcg_static_tokens is not None + and actual_tokens is not None + and pcg_static_tokens > actual_tokens + ): + first_dim = buf.shape[0] + elems_per_token = buf.numel() // first_dim + buf.view(first_dim, elems_per_token)[actual_tokens:].zero_() + + if TYPE_CHECKING: from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -128,6 +144,31 @@ class RadixAttention(nn.Module): forward_batch.forward_mode.is_extend() and get_tc_piecewise_forward_context() is not None ): + if kwargs.get("idx_q") is not None: + if is_in_breakable_cuda_graph(): + return get_attn_backend().forward( + q, k, v, self, forward_batch, save_kv_cache, **kwargs + ) + idx_q = kwargs["idx_q"] + idx_k = kwargs["idx_k"] + idx_v = kwargs.get("idx_v") + attn_out = q.new_empty( + (q.shape[0], self.tp_q_head_num * self.v_head_dim) + ) + idx_out = q.new_empty((q.shape[0], idx_q.shape[1] * idx_q.shape[2])) + unified_sparse_attention_with_output( + q, + k, + v, + attn_out, + idx_out, + idx_q, + idx_k, + save_kv_cache, + self.layer_id, + idx_v=idx_v, + ) + return idx_out, attn_out if self.qk_head_dim != self.v_head_dim: output = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim)) else: @@ -232,24 +273,71 @@ def unified_attention_with_output( if ret.data_ptr() != output.data_ptr(): output[:real_num_tokens].view(ret.shape).copy_(ret) - if _is_hip: - # During PCG replay on AMD, varlen attention kernels only fill positions - # 0..actual_tokens-1 and leave padded positions with uninitialized - # garbage from torch.empty. Zero these so garbage (NaN/Inf) does not - # propagate through residual connections, MoE routing, and allreduce. - # Use context.raw_num_tokens (pre-padding count from PCG runner) - # instead of forward_batch.extend_num_tokens, because - # extend_num_tokens is None for TARGET_VERIFY (EAGLE) batches. - pcg_static_tokens = context.num_tokens - actual_tokens = context.raw_num_tokens - if ( - pcg_static_tokens is not None - and actual_tokens is not None - and pcg_static_tokens > actual_tokens - ): - first_dim = output.shape[0] - elems_per_token = output.numel() // first_dim - output.view(first_dim, elems_per_token)[actual_tokens:].zero_() + # During PCG replay the attention backend writes only the narrowed + # real-token slice (output[:real_num_tokens]) and leaves padded positions + # as uninitialized torch.empty garbage. Zero them so garbage (NaN/Inf) does + # not propagate through residual connections, MoE routing, and allreduce. + # This affects every backend that varlen-writes under PCG, not just ROCm. + # Use context.raw_num_tokens (pre-padding count from PCG runner) instead of + # forward_batch.extend_num_tokens, which is None for TARGET_VERIFY batches. + _zero_padded_pcg_tail(output, context) + return + + +@register_custom_op(mutates_args=["attn_out", "idx_out"]) +@register_split_op() +def unified_sparse_attention_with_output( + query: torch.Tensor, + key: Optional[torch.Tensor], + value: Optional[torch.Tensor], + attn_out: torch.Tensor, + idx_out: torch.Tensor, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + save_kv_cache: bool, + layer_id: int, + *, + idx_v: Optional[torch.Tensor] = None, +) -> None: + context = get_tc_piecewise_forward_context() + forward_batch = context.forward_batch + attention_layer = context.attention_layers[layer_id] + real_num_tokens = forward_batch.num_token_non_padded_cpu + + query = query[:real_num_tokens] + if key is not None: + key = key[:real_num_tokens] + if value is not None: + value = value[:real_num_tokens] + idx_q = idx_q[:real_num_tokens] + idx_k = idx_k[:real_num_tokens] + if idx_v is not None: + idx_v = idx_v[:real_num_tokens] + + original_out_cache_loc = forward_batch.out_cache_loc + forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens] + + ret_idx, ret_out = get_attn_backend().forward( + query, + key, + value, + attention_layer, + forward_batch, + save_kv_cache, + idx_q=idx_q, + idx_k=idx_k, + idx_v=idx_v, + ) + forward_batch.out_cache_loc = original_out_cache_loc + + attn_out[:real_num_tokens].view(ret_out.shape).copy_(ret_out) + # disable_value layers return ret_idx=None; the guard keeps idx_out's + # untouched real-token slice safe (model returns before index_o_proj). + if ret_idx is not None: + idx_out[:real_num_tokens].view(ret_idx.shape).copy_(ret_idx) + + for buf in (attn_out, idx_out): + _zero_padded_pcg_tail(buf, context) return diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index f05bba746..99e809a40 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -599,6 +599,9 @@ class TokenizerWorker(TokenizerManager): port_args: PortArgs, ): setproctitle.setproctitle(f"sglang::tokenizer_worker:{os.getpid()}") + import torch + + torch.set_num_threads(1) # prevent init prefill bootstrapserver again disaggregation_mode = server_args.disaggregation_mode server_args.override( @@ -711,9 +714,7 @@ async def print_exception_wrapper(func): def get_main_process_id() -> int: - """ - Get the main process ID. - """ + """Get the main process ID.""" return multiprocessing.current_process()._parent_pid diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 00f97dcaa..92a405dd8 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -39,7 +39,7 @@ from torch.cuda import Stream as CudaStream from torch.distributed import barrier from sglang.jit_kernel.ngram_embedding import update_token_table -from sglang.srt.configs.model_config import ModelConfig, ModelImpl +from sglang.srt.configs.model_config import ModelConfig, ModelImpl, is_minimax_sparse from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix from sglang.srt.disaggregation.decode import ( @@ -1139,8 +1139,11 @@ class Scheduler( if ( self.disaggregation_mode == DisaggregationMode.DECODE - ): # *2 for the headroom. - buffer_size = (self.req_to_token_pool.size) * 2 + ): # *8 headroom for MiniMax-M3; *2 for other models. + buffer_multiplier = ( + 8 if is_minimax_sparse(self.model_config.hf_config) else 2 + ) + buffer_size = (self.req_to_token_pool.size) * buffer_multiplier self.req_to_metadata_buffer_idx_allocator = ReqToMetadataIdxAllocator( buffer_size ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 9b29865b2..6f97b707b 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -32,7 +32,7 @@ import warnings from dataclasses import dataclass from enum import IntEnum, auto from functools import total_ordering -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union import torch @@ -443,6 +443,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # the carried topk lives on spec_info (see EagleDraftInput.dsa_topk_indices). reuse_dsa_topk_indices: Optional[bool] = False + minimax_m3_precached_sparse_layers: Optional[Set[int]] = None + # === Forward-derived (built in init_new on the forward stream; FB-owned) === # Position information positions: torch.Tensor = None diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py new file mode 100644 index 000000000..49241cfa0 --- /dev/null +++ b/python/sglang/srt/models/minimax_m3.py @@ -0,0 +1,1690 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from DeepSeek and Mixtral implementation +"""Inference-only MiniMax M3 model compatible with HuggingFace weights.""" + +import logging +from contextlib import nullcontext +from typing import Iterable, List, Optional, Set, Tuple, Union + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo +from sglang.srt.configs.model_config import ( + get_minimax_sparse_disable_value_layer_ids, + get_minimax_sparse_layer_ids, +) +from sglang.srt.distributed import ( + get_pp_group, + tensor_model_parallel_all_reduce, +) +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + ScatterMode, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.moe.utils import get_moe_a2a_backend +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.utils.common import get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + Phase, + check_cuda_graph_backend, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.forward_context import ( + get_forward_context, + has_forward_context, +) +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.models.minimax_m2 import MiniMaxM2RMSNormTP +from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.utils import ( + add_prefix, + get_device_sm, + is_cuda, + is_hip, + log_info_on_rank0, + make_layers, +) +from sglang.srt.utils.hf_transformers_utils import get_rope_config + +_is_cuda = is_cuda() +_is_hip = is_hip() +_device_sm = get_device_sm() + +_FP8_KV_DTYPES = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float8_e4m3fnuz, +) + +# rotary_dim required by the fused qknorm+rope JIT kernel: rotary_dim/2 must +# equal the CUDA warp size (32) so each warp norms+ropes one head in one pass. +_M3_FUSED_QKNORM_ROPE_ROTARY_DIM = 64 + +_has_rocm_qk_norm_rope = False +if _is_hip: + try: + from sglang.jit_kernel.minimax_m3.qk_norm_rope import ( + qk_gemma_rmsnorm_rope, + sparse_qk_index_gemma_rmsnorm_rope, + sparse_qk_index_gemma_rmsnorm_rope_cache, + ) + + _has_rocm_qk_norm_rope = True + except ImportError: + _has_rocm_qk_norm_rope = False + +logger = logging.getLogger(__name__) + + +class MultiHeadRMSNorm(nn.Module): + def __init__( + self, + num_heads: int, + head_dim: int, + eps: float = 1e-6, + apply_layernorm_1p: bool = False, + ) -> None: + super().__init__() + self.tp_world = get_parallel().attn_tp_size + self.tp_rank = get_parallel().attn_tp_rank + self.num_heads = num_heads + self.num_heads_per_tp = num_heads // self.tp_world + self.head_dim = head_dim + self.weight = nn.Parameter( + torch.ones(self.num_heads_per_tp, self.head_dim, dtype=torch.float32) + ) + self.weight.weight_loader = self.weight_loader + self.apply_layernorm_1p = apply_layernorm_1p + self.variance_epsilon = eps + + @staticmethod + def weight_loader( + param: nn.Parameter, + loaded_weight: torch.Tensor, + ) -> None: + tp_world = get_parallel().attn_tp_size + tp_rank = get_parallel().attn_tp_rank + + shard_size = loaded_weight.shape[0] // tp_world + shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) + param.data.copy_(loaded_weight[shard].reshape_as(param)) + + def forward( + self, + x: torch.Tensor, + ) -> torch.Tensor: + orig_dtype = x.dtype + x = x.view(-1, self.num_heads_per_tp, self.head_dim).to(torch.float32) + variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32) + x = x * torch.rsqrt(variance + self.variance_epsilon) + if self.apply_layernorm_1p: + x = x * (self.weight + 1)[None, ...] + else: + x = x * self.weight[None, ...] + x = x.view(-1, self.num_heads_per_tp * self.head_dim) + return x.to(orig_dtype) + + +class _FusedQKVIndexProj(nn.Module): + def __init__( + self, + quant_method, + weight: torch.Tensor, + weight_scale_inv: Optional[torch.Tensor], + input_size_per_partition: int, + logical_widths: List[int], + orig_dtype: torch.dtype, + ) -> None: + super().__init__() + # Named ``_qm`` (not ``quant_method``) so the loader's post-process loop + # skips this module; the backend scale layout is derived once below. + self._qm = quant_method + self.register_parameter("weight", nn.Parameter(weight, requires_grad=False)) + self.input_size_per_partition = input_size_per_partition + self.output_size_per_partition = weight.shape[0] + self.logical_widths = logical_widths + self.orig_dtype = orig_dtype + self.input_scale = None + if weight_scale_inv is not None: + self.register_parameter( + "weight_scale_inv", nn.Parameter(weight_scale_inv, requires_grad=False) + ) + self.weight_scale_inv.format_ue8m0 = True + # Must derive the backend scale layout here: the loader skips this + # module (see ``_qm``), so it won't run process_weights_after_loading. + quant_method._process_mxfp8_linear_weight_scale(self) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self._qm.apply(self, x, None) + + +def build_minimax_fused_qkv_index(model: nn.Module) -> None: + for module in model.modules(): + if isinstance(module, MiniMaxM3Attention): + module.maybe_build_fused_qkv_index() + + +class MiniMaxM3MLP(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + reduce_results: bool = True, + intermediate_size: int = None, + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + hidden_size = config.hidden_size + hidden_act = config.hidden_act + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "swigluoai": + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( + swiglu_no_interleaved_with_alpha_and_limit, + ) + + self.act_fn = lambda x: swiglu_no_interleaved_with_alpha_and_limit( + x, config.swiglu_alpha, config.swiglu_limit + ) + else: + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + + def forward( + self, + x, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj( + x, + skip_all_reduce=should_allreduce_fusion or use_reduce_scatter, + ) + return x + + +class MiniMaxM3MoE(nn.Module): + """MiniMax MoE implementation using DeepEP for Expert Parallel support.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_parallel().tp_size + self.n_shared_experts = getattr(config, "n_shared_experts", None) + self.num_fused_shared_experts = ( + 0 + if get_server_args().disable_shared_experts_fusion + else config.n_shared_experts + ) + + 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.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 + + self.experts = get_moe_impl_class(quant_config)( + num_experts=config.num_local_experts + + self.num_fused_shared_experts + + get_server_args().ep_num_redundant_experts, + num_fused_shared_experts=self.num_fused_shared_experts, + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + layer_id=layer_id, + quant_config=quant_config, + activation="silu", + is_gated=True, + gemm1_alpha=config.swiglu_alpha, + gemm1_clamp_limit=config.swiglu_limit, + prefix=add_prefix("experts", prefix), + gate_up_interleaved=False, + ) + self.topk = TopK( + top_k=config.num_experts_per_tok + self.num_fused_shared_experts, + renormalize=True, + layer_id=layer_id, + scoring_func=config.scoring_func, + correction_bias=self.e_score_correction_bias, + num_fused_shared_experts=self.num_fused_shared_experts, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scaling_factor_on_output=True, + ) + + if self.n_shared_experts is not None and self.num_fused_shared_experts == 0: + intermediate_size = config.intermediate_size * self.n_shared_experts + # DeepEP all-gathers (not all-reduces) the layer output, so a TP-sharded + # shared MLP would leave an unreduced partial; replicate (tp_size=1), like GLM4 / DSV2. + shared_experts_tp1 = get_moe_a2a_backend().is_deepep() + self.shared_experts = MiniMaxM3MLP( + config=config, + quant_config=quant_config, + prefix=add_prefix("shared_experts", prefix), + reduce_results=False, + intermediate_size=intermediate_size, + **(dict(tp_rank=0, tp_size=1) if shared_experts_tp1 else {}), + ) + else: + self.shared_experts = None + + self.bf16_router_gemm = envs.SGLANG_OPT_USE_BF16_ROUTER_GEMM.get() + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_local_experts, + bias=False, + params_dtype=torch.bfloat16 if self.bf16_router_gemm else torch.float32, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + + self.layer_id = layer_id + + if get_moe_a2a_backend().is_deepep(): + self.ep_size = get_parallel().moe_ep_size + self.top_k = config.num_experts_per_tok + + @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, + forward_batch: ForwardBatch, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if get_moe_a2a_backend().is_deepep(): + return self.forward_deepep(hidden_states, forward_batch) + else: + return self.forward_normal( + hidden_states, should_allreduce_fusion, use_reduce_scatter + ) + + def forward_normal( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if hidden_states.shape[0] > 0: + shared_output = self._forward_shared_experts(hidden_states) + router_logits = self._compute_router_logits(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + else: + shared_output = None + topk_output = self.topk.empty_topk_output(hidden_states.device) + + final_hidden_states = self.experts(hidden_states, topk_output) + + if shared_output is not None: + final_hidden_states = final_hidden_states + shared_output + if self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states + + def forward_deepep( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + shared_output = None + if hidden_states.shape[0] > 0: + shared_output = self._forward_shared_experts(hidden_states) + router_logits = self._compute_router_logits(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=forward_batch.num_token_non_padded, + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ), + ) + else: + topk_output = self.topk.empty_topk_output(hidden_states.device) + + # DeepEP returns the complete per-token routed result (no TP all-reduce here); + # shared experts are replicated (tp_size=1), so both add directly. + final_hidden_states = self.experts(hidden_states, topk_output) + + if shared_output is not None: + final_hidden_states = final_hidden_states + shared_output + + return final_hidden_states + + def _compute_router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.bf16_router_gemm: + return torch.mm( + hidden_states, self.gate.weight.t(), out_dtype=torch.float32 + ) + router_logits, _ = self.gate(hidden_states.to(torch.float32)) + return router_logits + + def _forward_shared_experts(self, hidden_states: torch.Tensor): + if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0): + return self.shared_experts(hidden_states) + else: + return None + + +class MiniMaxM3Attention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + is_sparse_attention_layer: bool = False, + disable_index_value: bool = False, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.is_sparse_attention_layer = is_sparse_attention_layer + self.disable_index_value = is_sparse_attention_layer and disable_index_value + + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_size = attn_tp_size + self.attn_tp_rank = attn_tp_rank + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size + self.total_num_kv_heads = config.num_key_value_heads + + if self.total_num_kv_heads >= attn_tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % attn_tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + 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.rope_theta, self.rope_scaling = get_rope_config(config) + self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.rotary_dim = getattr(config, "rotary_dim", self.head_dim) + + self.qk_norm_type = getattr(config, "qk_norm_type", "per_layer") + self.use_gemma_norm = getattr(config, "use_gemma_norm", False) + + if self.is_sparse_attention_layer: + assert self.qk_norm_type == "per_head", ( + f"sparse attention only supports qk_norm_type='per_head', " + f"got {self.qk_norm_type!r}" + ) + sparse_cfg = config.sparse_attention_config + self.total_idx_heads = sparse_cfg["sparse_num_index_heads"] + self.idx_head_dim = sparse_cfg["sparse_index_dim"] + # idx_replica_size ranks share one idx head; pre-divide idx_o on the activation + # (not the weight) so the TP all-reduce sums right and stays FP8-quant-safe. + if self.total_idx_heads >= attn_tp_size: + assert self.total_idx_heads % attn_tp_size == 0 + else: + assert attn_tp_size % self.total_idx_heads == 0 + self.idx_head_tp_size = min(attn_tp_size, self.total_idx_heads) + self.idx_replica_size = attn_tp_size // self.idx_head_tp_size + self.idx_head_rank = attn_tp_rank // self.idx_replica_size + self.num_idx_heads = self.total_idx_heads // self.idx_head_tp_size + + 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, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + reduce_results=False, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=self.max_position_embeddings, + base=self.rope_theta, + rope_scaling=self.rope_scaling, + ) + + if self.is_sparse_attention_layer: + self.index_qkv_proj = QKVParallelLinear( + self.hidden_size, + self.idx_head_dim, + self.total_idx_heads, + total_num_kv_heads=1, + bias=False, + quant_config=quant_config, + v_head_size=(0 if self.disable_index_value else self.idx_head_dim), + tp_rank=self.idx_head_rank, + tp_size=self.idx_head_tp_size, + prefix=add_prefix("index_qkv_proj", prefix), + ) + + if self.disable_index_value: + self.index_o_proj = None + else: + self.index_o_proj = RowParallelLinear( + self.total_idx_heads * self.idx_head_dim, + self.hidden_size, + bias=False, + input_is_parallel=True, + reduce_results=False, + quant_config=quant_config, + prefix=add_prefix("index_o_proj", prefix), + tp_rank=self.idx_head_rank, + tp_size=self.idx_head_tp_size, + ) + self.index_rotary_emb = self.rotary_emb + + if self.qk_norm_type == "per_layer": + if attn_tp_size > 1: + self.q_norm = MiniMaxM2RMSNormTP( + self.total_num_heads * self.head_dim, eps=config.rms_norm_eps + ) + self.k_norm = MiniMaxM2RMSNormTP( + self.total_num_kv_heads * self.head_dim, eps=config.rms_norm_eps + ) + else: + self.q_norm = RMSNorm( + self.total_num_heads * self.head_dim, eps=config.rms_norm_eps + ) + self.k_norm = RMSNorm( + self.total_num_kv_heads * self.head_dim, eps=config.rms_norm_eps + ) + elif self.qk_norm_type == "per_head": + if self.use_gemma_norm: + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + else: + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + if self.is_sparse_attention_layer: + if self.use_gemma_norm: + self.index_q_norm = GemmaRMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + self.index_k_norm = GemmaRMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + else: + self.index_q_norm = RMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + self.index_k_norm = RMSNorm( + self.idx_head_dim, eps=config.rms_norm_eps + ) + elif self.qk_norm_type == "multi_head": + self.q_norm = MultiHeadRMSNorm( + self.total_num_heads, + self.head_dim, + eps=config.rms_norm_eps, + apply_layernorm_1p=self.use_gemma_norm, + ) + self.k_norm = MultiHeadRMSNorm( + self.total_num_kv_heads, + self.head_dim, + eps=config.rms_norm_eps, + apply_layernorm_1p=self.use_gemma_norm, + ) + else: + raise ValueError(f"Invalid qk_norm_type: {self.qk_norm_type}") + + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + self._use_fused_qknorm_rope = ( + _is_cuda + and self.qk_norm_type == "per_head" + and self.use_gemma_norm + and self.head_dim == 128 + and self.rotary_dim == _M3_FUSED_QKNORM_ROPE_ROTARY_DIM + and getattr(self.rotary_emb, "is_neox_style", False) + ) + + self._fuse_qkv_index_enabled = self.is_sparse_attention_layer and ( + _is_cuda or _is_hip + ) + self._fused_qkv_index = None + self._fused_main_size = self.q_size + 2 * self.kv_size + + self._combined_qknorm_ok = ( + self.is_sparse_attention_layer + and self._use_fused_qknorm_rope + and self.idx_head_dim == 128 + and self.rotary_emb.cos_sin_cache.dtype == torch.float32 + ) + if self.is_sparse_attention_layer: + off_iq = self.num_heads + 2 * self.num_kv_heads + off_ik = off_iq + self.num_idx_heads + self._qknorm_group_meta = ( + (0, self.num_heads), + (self.num_heads, self.num_kv_heads), + (off_iq, self.num_idx_heads), + (off_ik, 1), + ) + + self._can_use_rocm_qk_norm_rope_static = ( + _has_rocm_qk_norm_rope + and self.qk_norm_type == "per_head" + and self.use_gemma_norm + and self.q_norm.variance_epsilon == self.k_norm.variance_epsilon + and hasattr(self.rotary_emb, "cos_sin_cache") + and self.rotary_emb.rotary_dim == self.rotary_dim + and self.rotary_dim <= self.head_dim + ) + self._can_use_rocm_index_qk_norm_rope_static = ( + self.is_sparse_attention_layer + and _has_rocm_qk_norm_rope + and self.use_gemma_norm + and self.index_q_norm.variance_epsilon == self.index_k_norm.variance_epsilon + and hasattr(self.index_rotary_emb, "cos_sin_cache") + and self.index_rotary_emb.rotary_dim == self.rotary_dim + and self.rotary_dim <= self.idx_head_dim + ) + self._can_use_rocm_sparse_qk_index_norm_rope_static = ( + self.is_sparse_attention_layer + and self._can_use_rocm_qk_norm_rope_static + and self._can_use_rocm_index_qk_norm_rope_static + and self.idx_head_dim == self.head_dim + and self.index_q_norm.variance_epsilon == self.q_norm.variance_epsilon + and self.index_k_norm.variance_epsilon == self.q_norm.variance_epsilon + and self.index_rotary_emb is self.rotary_emb + ) + + def _can_use_rocm_qk_norm_rope( + self, positions: torch.Tensor, q: torch.Tensor, k: torch.Tensor + ) -> bool: + return ( + self._can_use_rocm_qk_norm_rope_static + and positions.dim() == 1 + and q.dim() == 2 + and k.dim() == 2 + and q.dtype in (torch.bfloat16, torch.float16) + and k.dtype == q.dtype + ) + + def _qk_norm_rope( + self, positions: torch.Tensor, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self._can_use_rocm_qk_norm_rope(positions, q, k): + return qk_gemma_rmsnorm_rope( + q, + k, + self.q_norm.weight.data, + self.k_norm.weight.data, + positions, + self.rotary_emb.cos_sin_cache, + self.q_norm.variance_epsilon, + self.head_dim, + self.rotary_dim, + self.rotary_emb.is_neox_style, + ) + q, k = self._qk_norm(q, k) + return self.rotary_emb(positions, q, k) + + def _qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self.qk_norm_type == "per_layer": + if self.attn_tp_size > 1: + q, k = MiniMaxM2RMSNormTP.forward_qk( + self.q_norm, self.k_norm, q.contiguous(), k.contiguous() + ) + else: + q = self.q_norm(q.contiguous()) + k = self.k_norm(k.contiguous()) + elif self.qk_norm_type == "per_head": + q_shape = q.shape + k_shape = k.shape + q = q.reshape(-1, self.head_dim).contiguous() + k = k.reshape(-1, self.head_dim).contiguous() + q = self.q_norm(q).reshape(q_shape) + k = self.k_norm(k).reshape(k_shape) + elif self.qk_norm_type == "multi_head": + q = self.q_norm(q.contiguous()) + k = self.k_norm(k.contiguous()) + else: + raise ValueError(f"Invalid qk_norm_type: {self.qk_norm_type}") + return q, k + + def _index_qk_norm_rope( + self, positions: torch.Tensor, idx_q: torch.Tensor, idx_k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if ( + self._can_use_rocm_index_qk_norm_rope_static + and positions.dim() == 1 + and idx_q.dim() == 2 + and idx_k.dim() == 2 + and idx_q.dtype in (torch.bfloat16, torch.float16) + and idx_k.dtype == idx_q.dtype + ): + return qk_gemma_rmsnorm_rope( + idx_q, + idx_k, + self.index_q_norm.weight.data, + self.index_k_norm.weight.data, + positions, + self.index_rotary_emb.cos_sin_cache, + self.index_q_norm.variance_epsilon, + self.idx_head_dim, + self.rotary_dim, + self.index_rotary_emb.is_neox_style, + ) + idx_q, idx_k = self._index_qk_norm(idx_q, idx_k) + return self.index_rotary_emb(positions, idx_q, idx_k) + + def _index_qk_norm( + self, idx_q: torch.Tensor, idx_k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + idx_q_shape = idx_q.shape + idx_k_shape = idx_k.shape + idx_q = idx_q.reshape(-1, self.idx_head_dim) + idx_k = idx_k.reshape(-1, self.idx_head_dim) + idx_q = self.index_q_norm(idx_q).reshape(idx_q_shape) + idx_k = self.index_k_norm(idx_k).reshape(idx_k_shape) + return idx_q, idx_k + + def _split_index_qkv( + self, idx_qkv: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + q_size = self.num_idx_heads * self.idx_head_dim + if self.disable_index_value: + idx_q, idx_k = idx_qkv.split([q_size, self.idx_head_dim], dim=-1) + idx_v = None + else: + idx_q, idx_k, idx_v = idx_qkv.split( + [q_size, self.idx_head_dim, self.idx_head_dim], dim=-1 + ) + return idx_q, idx_k, idx_v + + def maybe_build_fused_qkv_index(self) -> None: + if not self._fuse_qkv_index_enabled or self._fused_qkv_index is not None: + return + from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod + + qp, ip = self.qkv_proj, self.index_qkv_proj + qm = qp.quant_method + if type(ip.quant_method) is not type(qm): + return + + # gfx942 converts MXFP8->block-fp8 in process_weights_after_loading; the + # fused module skips that pass, so keep two separate (converted) GEMMs. + if getattr(qm, "convert_mxfp8_to_block", False): + return + + is_unquant = isinstance(qm, UnquantizedLinearMethod) + use_mxfp8 = getattr(qm, "use_mxfp8", False) and hasattr(qp, "weight_scale_inv") + if not (is_unquant or use_mxfp8): + return + + weight = torch.cat([qp.weight.data, ip.weight.data], dim=0).contiguous() + if is_unquant: + scale = None + else: + scale = torch.cat( + [qp.weight_scale_inv.data, ip.weight_scale_inv.data], dim=0 + ).contiguous() + + holder = _FusedQKVIndexProj( + qm, + weight, + scale, + getattr(qp, "input_size_per_partition", qp.input_size), + [qp.output_size_per_partition, ip.output_size_per_partition], + getattr(qp, "orig_dtype", qp.params_dtype), + ) + self.add_module("fused_qkv_index_proj", holder) + self._fused_qkv_index = holder + + # Free the dead originals and drop their quant_method so the loader's + # post-process loop ignores them (see ``_qm``). + for m in (qp, ip): + m.quant_method = None + for attr in ("weight", "weight_scale_inv"): + p = getattr(m, attr, None) + if isinstance(p, nn.Parameter): + p.data = torch.empty(0, dtype=p.dtype, device=p.data.device) + + def _qknorm_groups(self): + weights = ( + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + ) + return [ + (w, off, cnt) for w, (off, cnt) in zip(weights, self._qknorm_group_meta) + ] + + def _can_use_rocm_sparse_qk_index_norm_rope( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + ) -> bool: + return ( + self._can_use_rocm_sparse_qk_index_norm_rope_static + and positions.dim() == 1 + and q.dim() == 2 + and k.dim() == 2 + and idx_q.dim() == 2 + and idx_k.dim() == 2 + and q.dtype in (torch.bfloat16, torch.float16) + and k.dtype == q.dtype + and idx_q.dtype == q.dtype + and idx_k.dtype == q.dtype + ) + + def _sparse_qk_index_norm_rope( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if self._can_use_rocm_sparse_qk_index_norm_rope(positions, q, k, idx_q, idx_k): + return sparse_qk_index_gemma_rmsnorm_rope( + q, + k, + idx_q, + idx_k, + self.q_norm.weight.data, + self.k_norm.weight.data, + self.index_q_norm.weight.data, + self.index_k_norm.weight.data, + positions, + self.rotary_emb.cos_sin_cache, + self.q_norm.variance_epsilon, + self.head_dim, + self.rotary_dim, + self.rotary_emb.is_neox_style, + ) + q, k = self._qk_norm_rope(positions, q, k) + idx_q, idx_k = self._index_qk_norm_rope(positions, idx_q, idx_k) + return q, k, idx_q, idx_k + + @staticmethod + def _mark_sparse_kv_cached_by_fusion( + forward_batch: ForwardBatch, layer_id: int + ) -> None: + layer_ids = forward_batch.minimax_m3_precached_sparse_layers + if layer_ids is None: + layer_ids = set() + forward_batch.minimax_m3_precached_sparse_layers = layer_ids + layer_ids.add(layer_id) + + @staticmethod + def _get_sparse_kv_pool(): + if not has_forward_context(): + return None + attn_backend = get_forward_context().attn_backend + sparse_backend = getattr(attn_backend, "sparse", None) + return getattr(sparse_backend, "kv_pool", None) + + def _sparse_qk_index_norm_rope_cache( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + idx_q: torch.Tensor, + idx_k: torch.Tensor, + idx_v: Optional[torch.Tensor], + forward_batch: ForwardBatch, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + kv_pool = self._get_sparse_kv_pool() + # The fused kernel writes normed bf16 K/V straight into the paged cache, so an + # fp8 main K/V cache (--kv-cache-dtype fp8_*) can't use it; fall back to norm+rope. + main_kv_is_fp8 = kv_pool is not None and kv_pool.dtype in _FP8_KV_DTYPES + can_use_cache_fusion = ( + not main_kv_is_fp8 + and idx_v is None + and self._can_use_rocm_sparse_qk_index_norm_rope( + positions, q, k, idx_q, idx_k + ) + and getattr(forward_batch, "out_cache_loc", None) is not None + and v.dim() == 2 + and v.dtype == q.dtype + and v.shape == k.shape + ) + if can_use_cache_fusion and kv_pool is not None: + layer_id = self.attn.layer_id + k_cache, v_cache = kv_pool.get_kv_buffer(layer_id) + idx_k_cache = kv_pool.get_index_k_buffer(layer_id) + q, k, idx_q, idx_k = sparse_qk_index_gemma_rmsnorm_rope_cache( + q, + k, + v, + idx_q, + idx_k, + k_cache, + v_cache, + idx_k_cache, + forward_batch.out_cache_loc, + self.q_norm.weight.data, + self.k_norm.weight.data, + self.index_q_norm.weight.data, + self.index_k_norm.weight.data, + positions, + self.rotary_emb.cos_sin_cache, + self.q_norm.variance_epsilon, + self.head_dim, + self.rotary_dim, + self.rotary_emb.is_neox_style, + ) + self._mark_sparse_kv_cached_by_fusion(forward_batch, layer_id) + return q, k, idx_q, idx_k + return self._sparse_qk_index_norm_rope(positions, q, k, idx_q, idx_k) + + def forward_prepare( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ): + fused_out = None + if self._fused_qkv_index is not None: + fused_out = self.fused_qkv_index_proj(hidden_states) + qkv = fused_out[:, : self._fused_main_size] + + if self._combined_qknorm_ok: + from sglang.jit_kernel.minimax_qknorm_rope import ( + minimax_qknorm_rope_grouped, + ) + + minimax_qknorm_rope_grouped( + fused_out, + self._qknorm_groups(), + self.rotary_emb.cos_sin_cache, + positions, + self.q_norm.variance_epsilon, + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + idx_qkv = fused_out[:, self._fused_main_size :] + idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv) + inner_state = (q, k, v, idx_q, idx_k, idx_v, forward_batch) + return None, forward_batch, inner_state + else: + qkv, _ = self.qkv_proj(hidden_states) + + if self._use_fused_qknorm_rope: + from sglang.jit_kernel.minimax_qknorm_rope import minimax_qknorm_rope + + minimax_qknorm_rope( + qkv, + self.q_norm.weight, + self.k_norm.weight, + self.rotary_emb.cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + self.num_kv_heads, + self.q_norm.variance_epsilon, + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + main_qk_already_normed = True + else: + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + main_qk_already_normed = False + + if self.is_sparse_attention_layer: + if fused_out is not None: + idx_qkv = fused_out[:, self._fused_main_size :] + else: + idx_qkv, _ = self.index_qkv_proj(hidden_states) + + if main_qk_already_normed: + use_fused_index_norm_rope = ( + self._use_fused_qknorm_rope + and self.idx_head_dim == 128 + and self.index_rotary_emb.cos_sin_cache.dtype == torch.float32 + ) + if use_fused_index_norm_rope: + from sglang.jit_kernel.minimax_qknorm_rope import ( + minimax_qknorm_rope, + ) + + minimax_qknorm_rope( + idx_qkv, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.index_rotary_emb.cos_sin_cache, + positions, + self.num_idx_heads, + 1, + 0 if self.disable_index_value else 1, + self.index_q_norm.variance_epsilon, + ) + idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv) + else: + idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv) + idx_q, idx_k = self._index_qk_norm_rope(positions, idx_q, idx_k) + else: + idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv) + q, k, idx_q, idx_k = self._sparse_qk_index_norm_rope_cache( + positions, q, k, v, idx_q, idx_k, idx_v, forward_batch + ) + + inner_state = (q, k, v, idx_q, idx_k, idx_v, forward_batch) + else: + if not main_qk_already_normed: + q, k = self._qk_norm_rope(positions, q, k) + inner_state = (q, k, v, forward_batch) + return None, forward_batch, inner_state + + def forward_core(self, intermediate_state): + _, _, inner_state = intermediate_state + + if self.is_sparse_attention_layer: + q, k, v, idx_q, idx_k, idx_v, forward_batch = inner_state + q = q.view(q.shape[0], self.num_heads, self.head_dim) + k = k.view(k.shape[0], self.num_kv_heads, self.head_dim) + v = v.view(v.shape[0], self.num_kv_heads, self.head_dim) + idx_q = idx_q.reshape(idx_q.shape[0], self.num_idx_heads, self.idx_head_dim) + idx_k = idx_k.reshape(idx_k.shape[0], 1, self.idx_head_dim) + if idx_v is not None: + idx_v = idx_v.reshape(idx_v.shape[0], 1, self.idx_head_dim) + idx_o, attn_output = self.attn( + q, k, v, forward_batch, idx_q=idx_q, idx_k=idx_k, idx_v=idx_v + ) + output, _ = self.o_proj(attn_output) + if self.disable_index_value: + return output + # idx_replica_size ranks produce identical idx_o; pre-divide idx_o (not the + # o_proj weight) so the TP all-reduce sums right and stays FP8-quant-safe. + if self.idx_replica_size > 1: + idx_o = idx_o / self.idx_replica_size + idx_output, _ = self.index_o_proj(idx_o) + return output + idx_output + + q, k, v, forward_batch = inner_state + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + s = self.forward_prepare( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + return self.forward_core(s) + + +class MiniMaxM3DecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.layer_id = layer_id + + sparse_attention_config = getattr(config, "sparse_attention_config", None) + if sparse_attention_config is not None: + _, sparse_layer_ids = get_minimax_sparse_layer_ids(sparse_attention_config) + is_sparse_attention_layer = layer_id in sparse_layer_ids + disable_value_layer_ids = set( + get_minimax_sparse_disable_value_layer_ids(sparse_attention_config) + ) + disable_index_value = layer_id in disable_value_layer_ids + else: + is_sparse_attention_layer = False + disable_index_value = False + + self.self_attn = MiniMaxM3Attention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + is_sparse_attention_layer=is_sparse_attention_layer, + disable_index_value=disable_index_value, + ) + + moe_layer_freq = getattr(config, "moe_layer_freq", None) + # Means "MLP is a sparse MoE", not attention sparsity. Kept as ``is_layer_sparse`` + # because LayerCommunicator / LayerScatterModes / other models read this attr. + self.is_layer_sparse = ( + moe_layer_freq[layer_id] != 0 if moe_layer_freq is not None else True + ) + + if self.is_layer_sparse: + self.mlp = MiniMaxM3MoE( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = MiniMaxM3MLP( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + intermediate_size=config.dense_intermediate_size, + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + + self.use_gemma_norm = getattr(config, "use_gemma_norm", False) + if self.use_gemma_norm: + self.input_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + else: + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def _is_layer_sparse(lid): + if moe_layer_freq is None: + return True + if lid < 0 or lid >= config.num_hidden_layers: + return True + return moe_layer_freq[lid] != 0 + + is_previous_layer_sparse = _is_layer_sparse(layer_id - 1) + is_next_layer_sparse = _is_layer_sparse(layer_id + 1) + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + hidden_states, residual = ( + self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( + hidden_states, + residual, + forward_batch, + captured_last_layer_outputs=captured_last_layer_outputs, + **kwargs, + ) + ) + + if hidden_states.shape[0] != 0: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + if self.is_layer_sparse and get_parallel().tp_size > 1: + # Sparse MoE outputs are TP-partial; deferring their all-reduce into the next + # layer's fusion re-triggers the M3 no-EOS runaway. Force immediate all-reduce. + should_allreduce_fusion = False + + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + if self.is_layer_sparse or hidden_states.shape[0] != 0: + hidden_states = self.mlp( + hidden_states, + should_allreduce_fusion, + use_reduce_scatter, + ) + + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + else: + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + +class MiniMaxM3Model(nn.Module): + """MiniMax Model implementation.""" + + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.padding_idx = getattr(config, "pad_token_id", 0) + self.vocab_size = config.vocab_size + self.pp_group = get_pp_group() + self.use_gemma_norm = getattr(config, "use_gemma_norm", False) + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + use_attn_tp_group=is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + def layer_fn(idx, prefix: str) -> nn.Module: + return MiniMaxM3DecoderLayer( + config=config, + layer_id=idx, + quant_config=quant_config, + prefix=prefix, + ) + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + layer_fn, + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + if self.use_gemma_norm: + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + + self.layers_to_capture = [] + + def get_input_embeddings(self) -> torch.Tensor: + return self.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors, Tuple[torch.Tensor, list[torch.Tensor]]]: + if self.pp_group.is_first_rank: + if input_embeds is None: + embeds = self.get_input_embeddings() + hidden_states = embeds(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + + aux_hidden_states = [] + if forward_batch.can_run_tbo: + hidden_states, residual = model_forward_maybe_tbo( + layers=self.layers, + enable_tbo=True, + input_data_scatter_mode=ScatterMode.model_input_output(), + positions=positions, + forward_batch=forward_batch, + hidden_states=hidden_states, + residual=residual, + ) + else: + for i in range(self.start_layer, self.end_layer): + # NOTE: torch dynamo does not support graph break in context manager + ctx = ( + nullcontext() + if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) + else get_global_expert_distribution_recorder().with_current_layer(i) + ) + with ctx: + layer = self.layers[i] + hidden_states, residual = layer( + positions=positions, + forward_batch=forward_batch, + hidden_states=hidden_states, + residual=residual, + captured_last_layer_outputs=( + aux_hidden_states + if getattr(layer, "_is_layer_to_capture", False) + else None + ), + ) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + if hidden_states.shape[0] != 0: + if residual is not None: + hidden_states, _ = self.norm(hidden_states, residual) + else: + hidden_states = self.norm(hidden_states) + + if len(aux_hidden_states) == 0: + return hidden_states + return hidden_states, aux_hidden_states + + +class MiniMaxM3SparseForCausalLM(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.quant_config = quant_config + self.pp_group = get_pp_group() + + self.num_fused_shared_experts = 0 + self.determine_num_fused_shared_experts() + + self.model = MiniMaxM3Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + + if self.pp_group.is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_server_args().enable_dp_lm_head, + ) + + self.logits_processor = LogitsProcessor(config) + else: + self.lm_head = PPMissingLayer() + + self.capture_aux_hidden_states = False + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def determine_num_fused_shared_experts(self): + if get_server_args().disable_shared_experts_fusion: + return + + disable_reason = None + if not getattr(self.config, "n_shared_experts", None): + disable_reason = "No shared experts are defined in the config." + elif not _is_cuda: + disable_reason = "Shared experts fusion currently requires CUDA devices." + elif _is_cuda and (_device_sm is not None) and (_device_sm < 80): + disable_reason = "Shared experts fusion requires SM80 or newer GPUs." + elif get_parallel().moe_ep_size > 1: + disable_reason = "Shared experts fusion is not supported together with expert parallelism yet." + elif get_moe_a2a_backend().is_deepep(): + disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled." + + if disable_reason is not None: + from sglang.srt.arg_groups.overrides import declare_load_time_override + + declare_load_time_override( + "MiniMaxM3ForCausalLM.determine_num_fused_shared_experts", + {"disable_shared_experts_fusion": True}, + ) + log_info_on_rank0( + logger, + f"{disable_reason} Shared experts fusion optimization is disabled.", + ) + return + + self.num_fused_shared_experts = self.config.n_shared_experts + assert ( + self.num_fused_shared_experts == 1 + ), "Only 1 fused shared expert is supported for MiniMax-M3" + log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") + + def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None): + if not self.pp_group.is_last_rank: + return + + self.capture_aux_hidden_states = True + if layer_ids is None: + num_layers = self.config.num_hidden_layers + self.model.layers_to_capture = [ + 2, + num_layers // 2, + num_layers - 3, + ] + else: + self.model.layers_to_capture = [val + 1 for val in layer_ids] + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors + ) + + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states + ) + else: + return hidden_states + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + """Load model weights with proper mapping for MiniMax architecture.""" + + stacked_params_mapping = [ + # Leading "." on ".qkv_proj" prevents it from falsely matching the sparse + # index_q/k/v_proj weights (remapped separately below). + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + # Value-disabled layers have no ".index_v_proj" weight, so that entry never matches. + if getattr(self.config, "sparse_attention_config", None) is not None: + stacked_params_mapping += [ + (".index_qkv_proj", ".index_q_proj", "q"), + (".index_qkv_proj", ".index_k_proj", "k"), + (".index_qkv_proj", ".index_v_proj", "v"), + ] + + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_local_experts + self.num_fused_shared_experts, + ) + + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if layer_id is not None and ( + layer_id < self.model.start_layer or layer_id >= self.model.end_layer + ): + continue + + name = name.replace(".block_sparse_moe", ".mlp") + + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.num_local_experts}", + ) + name = name.replace("gate_proj", "w1") + name = name.replace("down_proj", "w2") + name = name.replace("up_proj", "w3") + + if "rotary_emb.inv_freq" in name: + continue + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Must skip experts before the name.replace below, else gate_proj -> + # gate_up_proj -> gate_gate_up_proj double-remap breaks load. + if "mlp.experts." in name: + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + 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: + is_expert_weight = False + + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + + is_expert_weight = True + + 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=shard_id, + expert_id=expert_id, + ) + break + else: + if is_expert_weight: + continue + + if name.endswith(".bias") and name not in params_dict: + continue + + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + if name in params_dict: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + try: + weight_loader(param, loaded_weight) + except Exception as e: + logger.warning(f"Error loading weight {name}: {e}") + continue + else: + logger.warning(f"Parameter {name} not found in params_dict") + loaded_params.add(name) + + # Run before the loader's process pass: the raw fp8 weight + uint8 scale are + # final here (mxfp8 post-process only derives the packed scale, not these). + build_minimax_fused_qkv_index(self) + return loaded_params + + @classmethod + def get_model_config_for_expert_location(cls, config): + from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation + + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_local_experts, + num_groups=None, + ) + + +def get_spec_layer_idx_from_weight_name( + config: PretrainedConfig, weight_name: str +) -> Optional[int]: + # M3 checkpoints emit MTP weights as model.mtp.layers.{i}.*; skip them here + # (no NextN module is built for the main model). + if hasattr(config, "num_mtp_modules") and (config.num_mtp_modules > 0): + for i in range(config.num_mtp_modules): + if weight_name.startswith(f"model.mtp.layers.{i}."): + return config.num_hidden_layers + i + return None + + +EntryClass = [MiniMaxM3SparseForCausalLM] diff --git a/python/sglang/srt/models/minimax_m3_vl.py b/python/sglang/srt/models/minimax_m3_vl.py new file mode 100644 index 000000000..c3d563962 --- /dev/null +++ b/python/sglang/srt/models/minimax_m3_vl.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import Iterable, List, Optional, Tuple + +import torch +import torch.nn as nn + +from sglang.srt.distributed import ( + get_pp_group, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.utils import get_moe_a2a_backend +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.utils.common import get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import ( + MultimodalDataItem, + MultimodalInputs, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.models.minimax_m3 import ( + MiniMaxM3Model, + MiniMaxM3SparseForCausalLM, + build_minimax_fused_qkv_index, + get_spec_layer_idx_from_weight_name, +) +from sglang.srt.models.minimax_vl_common import ( + CLIPVisionConfig, + MiniMaxVLVisionModel, + get_image_feature, + get_video_feature, + load_vision_weight, + merge_vit_qkv_weights, +) +from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.utils import add_prefix, get_device_sm, is_cuda, log_info_on_rank0 +from sglang.srt.utils.hf_transformers_utils import get_rope_config + +logger = logging.getLogger(__name__) + + +_is_cuda = is_cuda() +_device_sm = get_device_sm() + + +class MiniMaxM3SparseForConditionalGeneration(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.pp_group = get_pp_group() + + self.use_data_parallel = get_server_args().mm_enable_dp_encoder + + self.num_fused_shared_experts = 0 + self._determine_num_fused_shared_experts() + + vision_config_raw = config.vision_config + assert vision_config_raw is not None, "vision_config is required" + if hasattr(vision_config_raw, "to_dict"): + vision_config_dict = vision_config_raw.to_dict() + else: + vision_config_dict = vision_config_raw + vision_config = CLIPVisionConfig.from_dict(vision_config_dict) + self.vision_config = vision_config + + text_hidden_size = getattr(config.text_config, "hidden_size", None) + assert text_hidden_size is not None, "text_hidden_size is required" + projector_hidden_size = getattr(config, "projector_hidden_size", None) + + # Vision model skips quantization: CLIP dimensions (head_dim=80) are not + # compatible with MXFP8 kernel alignment requirements (128). + self.vision_tower = MiniMaxVLVisionModel( + config=vision_config, + text_hidden_size=text_hidden_size, + projector_hidden_size=projector_hidden_size, + quant_config=None, + prefix=add_prefix("vision_tower", prefix), + multimodal_projector_bias=getattr( + config, "multimodal_projector_bias", True + ), + patch_merge_bias=getattr(config, "patch_merge_bias", True), + ) + + text_config = config.text_config + self.model = MiniMaxM3Model( + config=text_config, + quant_config=quant_config, + prefix=add_prefix("language_model.model", prefix), + ) + + if self.pp_group.is_last_rank: + self.lm_head = ParallelLMHead( + text_config.vocab_size, + text_config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("language_model.lm_head", prefix), + use_attn_tp_group=get_server_args().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + + _, text_rope_scaling = get_rope_config(text_config) + self.is_mrope_enabled = ( + text_rope_scaling is not None and "mrope_section" in text_rope_scaling + ) + + self.logits_processor = LogitsProcessor(text_config) + + def _determine_num_fused_shared_experts(self) -> None: + text_config = self.config.text_config + server_args = get_server_args() + if server_args.disable_shared_experts_fusion: + return + + disable_reason = None + if not getattr(text_config, "n_shared_experts", None): + disable_reason = "No shared experts are defined in the config." + elif not _is_cuda: + disable_reason = "Shared experts fusion currently requires CUDA devices." + elif (_device_sm is not None) and (_device_sm < 80): + disable_reason = "Shared experts fusion requires SM80 or newer GPUs." + elif get_parallel().moe_ep_size > 1: + disable_reason = ( + "Shared experts fusion is not supported together with expert " + "parallelism yet." + ) + elif get_moe_a2a_backend().is_deepep(): + disable_reason = ( + "Shared experts fusion is not supported when Deepep MoE backend " + "is enabled." + ) + + if disable_reason is not None: + from sglang.srt.arg_groups.overrides import declare_load_time_override + + declare_load_time_override( + "MiniMaxM3VLForCausalLM._determine_num_fused_shared_experts", + {"disable_shared_experts_fusion": True}, + ) + log_info_on_rank0( + logger, + f"{disable_reason} Shared experts fusion optimization is disabled.", + ) + return + + self.num_fused_shared_experts = text_config.n_shared_experts + assert ( + self.num_fused_shared_experts == 1 + ), "Only 1 fused shared expert is supported" + log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") + + @classmethod + def get_model_config_for_expert_location(cls, config): + # EP asserts if this hook is absent on the top-level arch; VL nests the + # LM config under text_config, so delegate there (fall back to config). + text_config = getattr(config, "text_config", None) or config + return MiniMaxM3SparseForCausalLM.get_model_config_for_expert_location( + text_config + ) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + return MultiModalityDataPaddingPatternMultimodalTokens().pad_input_tokens( + input_ids, mm_inputs + ) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + return get_image_feature(self.vision_tower, items, self.use_data_parallel) + + def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + return get_video_feature(self.vision_tower, items, self.use_data_parallel) + + def get_input_embeddings(self): + return self.model.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = False, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ): + if self.is_mrope_enabled: + positions = forward_batch.mrope_positions + + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + multimodal_model=self, + positions=positions, + pp_proxy_tensors=pp_proxy_tensors, + ) + + if self.pp_group.is_last_rank and not get_embedding: + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + ) + return hidden_states + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + + # ``.qkv_proj`` (with the leading dot) prevents matching e.g. + # ``index_q_proj`` in the sparse-attention branch. + llm_stacked_params_mapping = [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + if ( + getattr(self.config.text_config, "sparse_attention_config", None) + is not None + ): + llm_stacked_params_mapping += [ + (".index_qkv_proj", ".index_q_proj", "q"), + (".index_qkv_proj", ".index_k_proj", "k"), + (".index_qkv_proj", ".index_v_proj", "v"), + ] + + num_experts = getattr(self.config.text_config, "num_local_experts", 0) + expert_params_mapping = ( + FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=num_experts + self.num_fused_shared_experts, + ) + if num_experts > 0 + else [] + ) + + params_dict = dict(self.named_parameters()) + vit_qkv_weights: dict = {} + vit_qkv_biases: dict = {} + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + if name.startswith("language_model."): + self._load_llm_weight( + name[len("language_model.") :], + loaded_weight, + params_dict, + llm_stacked_params_mapping, + expert_params_mapping, + ) + continue + + load_vision_weight( + name, loaded_weight, params_dict, vit_qkv_weights, vit_qkv_biases + ) + + merge_vit_qkv_weights(vit_qkv_weights, vit_qkv_biases, params_dict) + + build_minimax_fused_qkv_index(self) + + def _load_llm_weight( + self, + name: str, + loaded_weight: torch.Tensor, + params_dict: dict, + llm_stacked_params_mapping: list, + expert_params_mapping: list, + ) -> None: + if "block_sparse_moe" in name: + name = name.replace("block_sparse_moe", "mlp") + + layer_id = get_layer_id(name) + if layer_id is not None and ( + layer_id < self.model.start_layer or layer_id >= self.model.end_layer + ): + return + + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.text_config.num_local_experts}", + ) + name = name.replace("gate_proj", "w1") + name = name.replace("down_proj", "w2") + name = name.replace("up_proj", "w3") + + if ( + get_spec_layer_idx_from_weight_name(self.config.text_config, name) + is not None + ): + return + + for param_name, weight_name, shard_id in llm_stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts." in name: + continue + new_name = name.replace(weight_name, param_name) + if new_name.endswith(".bias") and new_name not in params_dict: + continue + if new_name not in params_dict: + continue + param = params_dict[new_name] + param.weight_loader(param, loaded_weight, shard_id) + return + + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + is_expert_weight = True + new_name = name.replace(weight_name, param_name) + if new_name not in params_dict: + continue + param = params_dict[new_name] + param.weight_loader( + param, + loaded_weight, + new_name, + shard_id=shard_id, + expert_id=expert_id, + ) + return + if is_expert_weight: + return + + if name.endswith(".bias") and name not in params_dict: + return + remapped = maybe_remap_kv_scale_name(name, params_dict) + if remapped is None: + return + if remapped not in params_dict: + logger.warning(f"Parameter {remapped} not found in params_dict") + return + param = params_dict[remapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + try: + weight_loader(param, loaded_weight) + except Exception as e: + logger.warning(f"Error loading weight {remapped}: {e}") + + +EntryClass = [MiniMaxM3SparseForConditionalGeneration] diff --git a/python/sglang/srt/models/minimax_vl_common.py b/python/sglang/srt/models/minimax_vl_common.py new file mode 100644 index 000000000..a86831a6f --- /dev/null +++ b/python/sglang/srt/models/minimax_vl_common.py @@ -0,0 +1,882 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from dataclasses import dataclass, fields +from typing import List, Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn + +from sglang.srt.layers.activation import get_act_fn +from sglang.srt.layers.attention.vision import ( + BATCH_BUCKETS, + FLASHINFER_MAX_SEQLEN_BUCKETS, + FLASHINFER_WORKSPACE_SIZE_BYTES, + VisionAttention, +) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.rotary_embedding.utils import rotate_half +from sglang.srt.managers.schedule_batch import MultimodalDataItem +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model +from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.utils import add_prefix, get_compiler_backend, round_up + +logger = logging.getLogger(__name__) + + +@dataclass +class CLIPVisionConfig: + hidden_size: int + intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_channels: int + image_size: int + patch_size: int + hidden_act: str + layer_norm_eps: float + img_token_compression_config: dict + position_embedding_type: str + rope_mode: str + rope_theta: float + vision_segment_max_frames: Optional[int] + + @classmethod + def from_dict(cls, d: dict) -> "CLIPVisionConfig": + valid_keys = {f.name for f in fields(cls)} + filtered = {k: v for k, v in d.items() if k in valid_keys} + if "rope_theta" not in filtered and isinstance(d.get("rope_parameters"), dict): + rope_theta = d["rope_parameters"].get("rope_theta") + if rope_theta is not None: + filtered["rope_theta"] = rope_theta + return cls(**filtered) + + +class MiniMaxVLMultiModalProjector(nn.Module): + def __init__( + self, + vision_hidden_size: int, + text_hidden_size: int, + projector_hidden_act: str, + multimodal_projector_bias: bool, + projector_hidden_size: Optional[int] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_data_parallel: bool = False, + ): + super().__init__() + + mid_size = ( + projector_hidden_size + if projector_hidden_size is not None + else text_hidden_size + ) + + tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size + tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank + + self.linear_1 = ColumnParallelLinear( + vision_hidden_size, + mid_size, + bias=multimodal_projector_bias, + quant_config=quant_config, + prefix=f"{prefix}.linear_1", + tp_size=tp_size, + tp_rank=tp_rank, + ) + assert ( + projector_hidden_act == "gelu" + ), f"Only gelu activation is supported, got {projector_hidden_act}" + self.act = get_act_fn(projector_hidden_act) + self.linear_2 = RowParallelLinear( + mid_size, + text_hidden_size, + bias=multimodal_projector_bias, + quant_config=quant_config, + prefix=f"{prefix}.linear_2", + tp_size=tp_size, + tp_rank=tp_rank, + use_dp_attention_reduce=is_dp_attention_enabled(), + ) + + def forward(self, image_features: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.linear_1(image_features) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.linear_2(hidden_states) + return hidden_states + + +class MiniMaxVLPatchMerger(nn.Module): + def __init__( + self, + spatial_merge_size: int, + text_hidden_size: int, + projector_hidden_act: str, + patch_merge_bias: bool, + projector_hidden_size: Optional[int] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_data_parallel: bool = False, + ): + super().__init__() + self.spatial_merge_size = spatial_merge_size + + mid_size = ( + projector_hidden_size + if projector_hidden_size is not None + else text_hidden_size + ) + + tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size + tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank + + self.linear_1 = ColumnParallelLinear( + text_hidden_size * spatial_merge_size**2, + mid_size, + bias=patch_merge_bias, + quant_config=quant_config, + prefix=f"{prefix}.linear_1", + tp_size=tp_size, + tp_rank=tp_rank, + ) + assert ( + projector_hidden_act == "gelu" + ), f"Only gelu activation is supported, got {projector_hidden_act}" + self.act = get_act_fn(projector_hidden_act) + self.linear_2 = RowParallelLinear( + mid_size, + text_hidden_size, + bias=patch_merge_bias, + quant_config=quant_config, + prefix=f"{prefix}.linear_2", + tp_size=tp_size, + tp_rank=tp_rank, + use_dp_attention_reduce=is_dp_attention_enabled(), + ) + + def forward(self, image_features: torch.Tensor) -> torch.Tensor: + image_features = image_features.reshape( + image_features.shape[0] // (self.spatial_merge_size**2), -1 + ) + hidden_states, _ = self.linear_1(image_features) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.linear_2(hidden_states) + return hidden_states + + +def _prepare_rotary_cos_sin( + freqs: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + cos = freqs.cos().repeat(1, 2).unsqueeze(-2).float() + sin = freqs.sin().repeat(1, 2).unsqueeze(-2).float() + return cos, sin + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def _minimax_rope_applier( + q: torch.Tensor, + k: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + x_shape=None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """3D RoPE uses rope_dim=60 < head_dim=64; trailing dims pass through unrotated.""" + cos, sin = position_embeddings + rot_dim = cos.shape[-1] + + q_rot = q[..., :rot_dim].float() + q_pass = q[..., rot_dim:] + k_rot = k[..., :rot_dim].float() + k_pass = k[..., rot_dim:] + + q_rot = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_rot = (k_rot * cos) + (rotate_half(k_rot) * sin) + + q = torch.cat((q_rot.to(q_pass.dtype), q_pass), dim=-1) + k = torch.cat((k_rot.to(k_pass.dtype), k_pass), dim=-1) + return q, k + + +class CLIPVisionEmbeddings(nn.Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + + self.config = config + self.embed_dim = config.hidden_size + self.patch_size = config.patch_size + self.input_num_channels = config.num_channels + + self.temporal_patch_size = config.img_token_compression_config.get( + "temporal_patch_size", 2 + ) + + self.patch_embedding = nn.Conv3d( + in_channels=self.input_num_channels, + out_channels=self.embed_dim, + kernel_size=(self.temporal_patch_size, self.patch_size, self.patch_size), + stride=(self.temporal_patch_size, self.patch_size, self.patch_size), + bias=False, + ) + + def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: + if self.patch_embedding.weight.dtype != pixel_values.dtype: + self.patch_embedding = self.patch_embedding.to(pixel_values.dtype) + + assert ( + pixel_values.dim() == 2 + ), f"pixel_values must be 2D, got {pixel_values.dim()}D" + pixel_values = pixel_values.reshape( + pixel_values.shape[0], + self.input_num_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + patch_embeds = self.patch_embedding(pixel_values) + patch_embeds = patch_embeds.reshape(patch_embeds.shape[0], -1) + return patch_embeds + + +class CLIPEncoderLayer(nn.Module): + def __init__( + self, + config: CLIPVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_data_parallel: bool = False, + workspace_buffer: Optional[torch.Tensor] = None, + ) -> None: + super().__init__() + + self.embed_dim = config.hidden_size + self.use_data_parallel = use_data_parallel + tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size + tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank + + self.self_attn = VisionAttention( + embed_dim=config.hidden_size, + num_heads=config.num_attention_heads, + projection_size=config.hidden_size, + use_qkv_parallel=True, + flatten_batch=True, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + use_data_parallel=use_data_parallel, + use_dp_attention_reduce=is_dp_attention_enabled(), + customized_position_embedding_applier=_minimax_rope_applier, + workspace_buffer=workspace_buffer, + ) + + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.fc1 = ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp.fc1", + tp_size=tp_size, + tp_rank=tp_rank, + ) + hidden_act = getattr(config, "hidden_act", "gelu") + assert ( + hidden_act == "gelu" + ), f"Only gelu activation is supported, got {hidden_act}" + self.act = get_act_fn(hidden_act) + self.fc2 = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp.fc2", + tp_size=tp_size, + tp_rank=tp_rank, + use_dp_attention_reduce=is_dp_attention_enabled(), + ) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seq_len: torch.Tensor, + rotary_pos_emb: torch.Tensor, + max_seqlen: Optional[int] = None, + sequence_lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.self_attn( + x=hidden_states, + cu_seqlens=cu_seq_len, + position_embeddings=rotary_pos_emb, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states, _ = self.fc1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.fc2(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class CLIPEncoder(nn.Module): + def __init__( + self, + config: CLIPVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_data_parallel: bool = False, + workspace_buffer: Optional[torch.Tensor] = None, + ) -> None: + super().__init__() + + self.config = config + self.use_data_parallel = use_data_parallel + + self.layers = nn.ModuleList( + [ + CLIPEncoderLayer( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.layers.{layer_idx}", + use_data_parallel=use_data_parallel, + workspace_buffer=workspace_buffer, + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + + def forward( + self, + inputs_embeds, + cu_seq_len: torch.Tensor, + rotary_pos_emb: torch.Tensor, + max_seqlen: Optional[int] = None, + sequence_lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + hidden_states = inputs_embeds + cos_sin = _prepare_rotary_cos_sin(rotary_pos_emb) + + for encoder_layer in self.layers: + hidden_states = encoder_layer( + hidden_states, + cu_seq_len, + cos_sin, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + + return hidden_states + + +class MiniMaxVLVisionTransformer(nn.Module): + def __init__( + self, + config: CLIPVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + *, + require_post_norm: Optional[bool] = None, + prefix: str = "", + use_data_parallel: bool = False, + ) -> None: + super().__init__() + + self.config = config + self.use_data_parallel = use_data_parallel + embed_dim = config.hidden_size + + self.temporal_patch_size = config.img_token_compression_config.get( + "temporal_patch_size", 2 + ) + self.spatial_merge_size = config.img_token_compression_config.get( + "spatial_merge_size", 2 + ) + + self.embeddings = CLIPVisionEmbeddings(config) + # NOTE: Typo "layrnorm" matches the original transformers code and the + # weight names used in the published checkpoints; do not "fix" it. + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + workspace_buffer: Optional[torch.Tensor] = None + if ( + get_server_args().mm_attention_backend == "flashinfer_cudnn" + and torch.cuda.is_available() + ): + workspace_buffer = torch.empty( + FLASHINFER_WORKSPACE_SIZE_BYTES, + dtype=torch.uint8, + device=torch.device("cuda", torch.cuda.current_device()), + ) + + self.encoder = CLIPEncoder( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.encoder", + use_data_parallel=use_data_parallel, + workspace_buffer=workspace_buffer, + ) + + assert ( + self.config.position_embedding_type == "rope" + ), "Only rope position embedding is supported" + assert self.config.rope_mode == "3d", "Only 3D RoPE is supported" + rope_theta = getattr(config, "rope_theta") + assert rope_theta is not None, "rope_theta must be set" + self.vision_segment_max_frames = getattr(config, "vision_segment_max_frames") + + head_dim = embed_dim // config.num_attention_heads + rope_dims = 2 * (head_dim // 2) + + 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)) + + 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) + + num_hidden_layers = config.num_hidden_layers + if len(self.encoder.layers) > config.num_hidden_layers: + raise ValueError( + f"The original encoder only has {num_hidden_layers} " + f"layers, but you requested {len(self.encoder.layers)} layers." + ) + + if require_post_norm is None: + require_post_norm = len(self.encoder.layers) == num_hidden_layers + self.post_layernorm = ( + nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + if require_post_norm + else None + ) + + def _get_3d_rope_embed( + self, grid_t: int, grid_h: int, grid_w: int, spatial_merge_size: int + ) -> torch.Tensor: + 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) + ) + hpos_ids = hpos_ids.reshape( + grid_h // spatial_merge_size, + spatial_merge_size, + grid_w // spatial_merge_size, + spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3) + hpos_ids = hpos_ids.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) + ) + wpos_ids = wpos_ids.reshape( + grid_h // spatial_merge_size, + spatial_merge_size, + grid_w // spatial_merge_size, + spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3) + wpos_ids = wpos_ids.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) + freqs_h = torch.outer(seq_hw, self.inv_freq_h) + freqs_w = torch.outer(seq_hw, self.inv_freq_w) + + emb_t = freqs_t[tpos_ids] + emb_h = freqs_h[hpos_ids] + emb_w = freqs_w[wpos_ids] + + return torch.cat([emb_t, emb_h, emb_w], dim=-1) + + def _get_rope_embed_3d(self, grid_thw, spatial_merge_size: int) -> torch.Tensor: + all_rope_embeds = [ + self._get_3d_rope_embed(grid_t, grid_h, grid_w, spatial_merge_size) + for grid_t, grid_h, grid_w in grid_thw + ] + return torch.cat(all_rope_embeds, dim=0) + + def _apply_max_frames_limit( + self, origin_grid_thw: list[list[int]] + ) -> List[List[int]]: + if self.vision_segment_max_frames is None: + return origin_grid_thw + max_frames = self.vision_segment_max_frames + ret_grid_thw = [] + for grid_t, grid_h, grid_w in origin_grid_thw: + if grid_t <= max_frames: + ret_grid_thw.append([grid_t, grid_h, grid_w]) + else: + for i in range(0, grid_t, max_frames): + sub_grid_t = min(max_frames, grid_t - i) + ret_grid_thw.append([sub_grid_t, grid_h, grid_w]) + return ret_grid_thw + + def _compute_cu_seq_len( + self, + grid_thw: list[list[int]], + device: torch.device, + ) -> torch.Tensor: + grid_thw = self._apply_max_frames_limit(grid_thw) + cu_seq_len = [0] + for grid_t, grid_h, grid_w in grid_thw: + cu_seq_len.append(grid_t * grid_h * grid_w) + cu_seq_len = torch.tensor(cu_seq_len, device=device).to(torch.int32) + cu_seq_len = torch.cumsum(cu_seq_len, dim=0).to(torch.int32) + return cu_seq_len + + @staticmethod + def _bucket_flashinfer_batch_size(batch_size: int) -> int: + return next( + (b for b in BATCH_BUCKETS if b >= batch_size), + round_up(batch_size, BATCH_BUCKETS[0]), + ) + + @staticmethod + def _bucket_flashinfer_max_seqlen(real_max_seqlen: int) -> int: + if real_max_seqlen <= 0: + return FLASHINFER_MAX_SEQLEN_BUCKETS[0] + return next( + (s for s in FLASHINFER_MAX_SEQLEN_BUCKETS if s >= real_max_seqlen), + round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]), + ) + + @classmethod + def _compute_flashinfer_sequence_lengths_padded( + cls, + token_cu_seqlens: np.ndarray, + ) -> np.ndarray: + assert token_cu_seqlens.ndim == 1 and token_cu_seqlens.size >= 2 + B = int(token_cu_seqlens.size - 1) + seq_lens = (token_cu_seqlens[1:] - token_cu_seqlens[:-1]).astype(np.int32) + B_padded = cls._bucket_flashinfer_batch_size(B) + if B_padded != B: + pad = np.zeros((B_padded - B,), dtype=np.int32) + seq_lens = np.concatenate([seq_lens, pad], axis=0) + return seq_lens + + @classmethod + def _compute_flashinfer_batch_offsets_packed( + cls, + token_cu_seqlens: np.ndarray, + *, + elem_per_token: int, + ) -> np.ndarray: + assert token_cu_seqlens.ndim == 1 and token_cu_seqlens.size >= 2 + B = int(token_cu_seqlens.size - 1) + B_padded = cls._bucket_flashinfer_batch_size(B) + token_indptr = token_cu_seqlens.astype(np.int64, copy=False) + if B_padded != B: + pad = np.full((B_padded - B,), token_indptr[-1], dtype=token_indptr.dtype) + token_indptr = np.concatenate([token_indptr, pad], axis=0) + elem_indptr = (token_indptr * int(elem_per_token)).astype(np.int32) + return np.concatenate([elem_indptr, elem_indptr, elem_indptr], axis=0) + + def _build_flashinfer_cudnn_inputs( + self, + cu_seq_len: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, int]: + device = cu_seq_len.device + token_cu_seqlens_np = cu_seq_len.detach().cpu().numpy().astype(np.int32) + + real_seq_lens = token_cu_seqlens_np[1:] - token_cu_seqlens_np[:-1] + max_seqlen = self._bucket_flashinfer_max_seqlen( + int(real_seq_lens.max()) if real_seq_lens.size > 0 else 0 + ) + + seq_lens_padded = self._compute_flashinfer_sequence_lengths_padded( + token_cu_seqlens_np + ) + + attn_tp_size = 1 if self.use_data_parallel else get_parallel().attn_tp_size + elem_per_token = self.config.hidden_size // attn_tp_size + + offsets_packed = self._compute_flashinfer_batch_offsets_packed( + token_cu_seqlens_np, + elem_per_token=elem_per_token, + ) + + sequence_lengths = ( + torch.from_numpy(seq_lens_padded) + .to(device=device, dtype=torch.int32, non_blocking=True) + .view(-1, 1, 1, 1) + ) + cu_seqlens_packed = torch.from_numpy(offsets_packed).to( + device=device, dtype=torch.int32, non_blocking=True + ) + return cu_seqlens_packed, sequence_lengths, int(max_seqlen) + + def forward( + self, + pixel_values: torch.Tensor, + grid_thw: list[list[int]], + ) -> torch.Tensor: + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + assert pixel_values.dtype == torch.bfloat16, "pixel_values must be bfloat16" + + hidden_states = self.embeddings(pixel_values) + hidden_states = self.pre_layrnorm(hidden_states) + + grid_thw = self._apply_max_frames_limit(grid_thw) + + cu_seq_len = self._compute_cu_seq_len(grid_thw, hidden_states.device) + rotary_pos_emb = self._get_rope_embed_3d(grid_thw, self.spatial_merge_size) + + assert ( + rotary_pos_emb.device == hidden_states.device + ), "rotary_pos_emb and hidden_states must be on the same device" + + max_seqlen: Optional[int] = None + sequence_lengths: Optional[torch.Tensor] = None + encoder_cu_seq_len = cu_seq_len + if get_server_args().mm_attention_backend == "flashinfer_cudnn": + ( + encoder_cu_seq_len, + sequence_lengths, + max_seqlen, + ) = self._build_flashinfer_cudnn_inputs(cu_seq_len) + + return self.encoder( + inputs_embeds=hidden_states, + cu_seq_len=encoder_cu_seq_len, + rotary_pos_emb=rotary_pos_emb, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + + +class MiniMaxVLVisionModel(nn.Module): + def __init__( + self, + config: CLIPVisionConfig, + text_hidden_size: int, + projector_hidden_size: Optional[int] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + multimodal_projector_bias: bool = True, + patch_merge_bias: bool = True, + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + + self.use_data_parallel = get_server_args().mm_enable_dp_encoder + self.vision_config = config + + self.vision_model = MiniMaxVLVisionTransformer( + config=config, + quant_config=quant_config, + prefix=add_prefix("vision_model", prefix), + use_data_parallel=self.use_data_parallel, + ) + + self.multi_modal_projector = MiniMaxVLMultiModalProjector( + vision_hidden_size=config.hidden_size, + text_hidden_size=text_hidden_size, + projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"), + multimodal_projector_bias=multimodal_projector_bias, + projector_hidden_size=projector_hidden_size, + quant_config=quant_config, + prefix=add_prefix("multi_modal_projector", prefix), + use_data_parallel=self.use_data_parallel, + ) + + spatial_merge_size = config.img_token_compression_config.get( + "spatial_merge_size", 2 + ) + self.spatial_merge_size = spatial_merge_size + self.patch_merge_mlp = MiniMaxVLPatchMerger( + spatial_merge_size=spatial_merge_size, + text_hidden_size=text_hidden_size, + projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"), + patch_merge_bias=patch_merge_bias, + projector_hidden_size=projector_hidden_size, + quant_config=quant_config, + prefix=add_prefix("patch_merge_mlp", prefix), + use_data_parallel=self.use_data_parallel, + ) + self.dtype = self.vision_model.embeddings.patch_embedding.weight.dtype + # Required by run_dp_sharded_mrope_vision_model when input is empty. + self.out_hidden_size = text_hidden_size + + def forward( + self, + pixel_values: torch.Tensor, + grid_thw: list[list[int]], + ) -> torch.Tensor: + hidden_states = self.vision_model(pixel_values=pixel_values, grid_thw=grid_thw) + if hidden_states.dim() == 3: + hidden_states = hidden_states.squeeze(0) + hidden_states = self.multi_modal_projector(hidden_states) + hidden_states = self.patch_merge_mlp(hidden_states) + return hidden_states + + +def _run_vision_tower( + vision_tower: MiniMaxVLVisionModel, + pixel_values: torch.Tensor, + grid_thw: list[list[int]], + use_data_parallel: bool, +) -> torch.Tensor: + if use_data_parallel: + return run_dp_sharded_mrope_vision_model( + vision_tower, + pixel_values, + grid_thw, + rope_type="rope_3d", + ) + return vision_tower(pixel_values, grid_thw=grid_thw) + + +def get_image_feature( + vision_tower: MiniMaxVLVisionModel, + items: List[MultimodalDataItem], + use_data_parallel: bool, +) -> torch.Tensor: + pixel_values = torch.cat([item.feature for item in items], dim=0).type( + vision_tower.dtype + ) + image_grid_thw: list[list[int]] = [] + for item in items: + image_grid_thw.extend(item.image_grid_thw.tolist()) + return _run_vision_tower( + vision_tower, pixel_values, image_grid_thw, use_data_parallel + ) + + +def get_video_feature( + vision_tower: MiniMaxVLVisionModel, + items: List[MultimodalDataItem], + use_data_parallel: bool, +) -> torch.Tensor: + pixel_values = torch.cat([item.feature for item in items], dim=0).type( + vision_tower.dtype + ) + video_grid_thw: list[list[int]] = [] + for item in items: + video_grid_thw.extend(item.video_grid_thw.tolist()) + assert pixel_values.dim() == 2, pixel_values.dim() + return _run_vision_tower( + vision_tower, pixel_values, video_grid_thw, use_data_parallel + ) + + +def _parse_vit_layer_idx(name: str) -> Optional[int]: + parts = name.split(".") + for i, p in enumerate(parts): + if p == "layers" and i + 1 < len(parts): + return int(parts[i + 1]) + return None + + +def load_vision_weight( + name: str, + loaded_weight: torch.Tensor, + params_dict: dict, + vit_qkv_weights: dict, + vit_qkv_biases: dict, +) -> None: + if ( + "self_attn.q_proj" in name + or "self_attn.k_proj" in name + or "self_attn.v_proj" in name + ): + if name.endswith(".weight"): + target = vit_qkv_weights + elif name.endswith(".bias"): + target = vit_qkv_biases + else: + return + layer_idx = _parse_vit_layer_idx(name) + if layer_idx is None: + return + qkv_type = "q" if "q_proj" in name else ("k" if "k_proj" in name else "v") + target.setdefault(layer_idx, {})[qkv_type] = loaded_weight + return + + param_name = name + if "vision_tower.vision_model." in param_name: + param_name = param_name.replace(".mlp.fc1.", ".fc1.") + param_name = param_name.replace(".mlp.fc2.", ".fc2.") + param_name = param_name.replace(".self_attn.out_proj.", ".self_attn.proj.") + if name.startswith("patch_merge_mlp.") or name.startswith("multi_modal_projector."): + param_name = "vision_tower." + param_name + + if param_name in params_dict: + param = params_dict[param_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +def merge_vit_qkv_weights( + vit_qkv_weights: dict, + vit_qkv_biases: dict, + params_dict: dict, +) -> None: + for layer_idx, qkv_dict in vit_qkv_weights.items(): + if {"q", "k", "v"} <= qkv_dict.keys(): + merged = torch.cat([qkv_dict["q"], qkv_dict["k"], qkv_dict["v"]], dim=0) + param_name = ( + f"vision_tower.vision_model.encoder.layers.{layer_idx}" + ".self_attn.qkv_proj.weight" + ) + if param_name in params_dict: + param = params_dict[param_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, merged) + + for layer_idx, qkv_dict in vit_qkv_biases.items(): + if {"q", "k", "v"} <= qkv_dict.keys(): + merged = torch.cat([qkv_dict["q"], qkv_dict["k"], qkv_dict["v"]], dim=0) + param_name = ( + f"vision_tower.vision_model.encoder.layers.{layer_idx}" + ".self_attn.qkv_proj.bias" + ) + if param_name in params_dict: + param = params_dict[param_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, merged) diff --git a/python/sglang/srt/multimodal/processors/minimax_m3_vl.py b/python/sglang/srt/multimodal/processors/minimax_m3_vl.py new file mode 100644 index 000000000..b24ec2bc8 --- /dev/null +++ b/python/sglang/srt/multimodal/processors/minimax_m3_vl.py @@ -0,0 +1,283 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +"""HF processor classes live in sglang.srt.configs.minimax_vl_processor to avoid circular imports with model classes.""" + +import math +import re +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torchvision +from torchvision.transforms import InterpolationMode + +from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput +from sglang.srt.models.minimax_m3_vl import MiniMaxM3SparseForConditionalGeneration +from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor, + MultimodalSpecialTokens, +) +from sglang.srt.utils import round_up + + +def get_hw_multiple_of( + image_size: Tuple[int, int], + multiple: int, + max_size: Union[None, int, Tuple[int, int]] = None, +) -> Tuple[int, int]: + w, h = image_size + + if isinstance(max_size, int): + ratio = 1.0 + max_dim = max(w, h) + if max_dim > max_size: + ratio = max_size / max_dim + new_w = round_up(round(w * ratio), multiple) + new_h = round_up(round(h * ratio), multiple) + return new_w, new_h + + new_w = round_up(w, multiple) + new_h = round_up(h, multiple) + + if max_size is not None: + assert isinstance(max_size, (list, tuple)) and len(max_size) == 2 + max_w, max_h = max_size + assert max_w % multiple == 0 and max_h % multiple == 0 + + if new_w > max_w or new_h > max_h: + new_w_ = min((new_w * max_w) // new_w, (new_w * max_h) // new_h) + new_h_ = min((new_h * max_w) // new_w, (new_h * max_h) // new_h) + new_w = new_w_ + new_h = new_h_ + + new_w = ( + new_w + if new_w % multiple == 0 + else new_w + (multiple - new_w % multiple) + ) + new_h = ( + new_h + if new_h % multiple == 0 + else new_h + (multiple - new_h % multiple) + ) + + assert new_w % multiple == 0 and new_h % multiple == 0 + assert new_w <= max_w and new_h <= max_h + + return new_w, new_h + + +def _compute_sampled_frame_indices( + total_frames: int, + video_fps: float, + fps: float, + max_frames: Optional[int] = None, +) -> List[int]: + """Frame indices must match SFT extract_frame.py constant-mode sampling (>=1/fps apart, always keep last) or eval diverges from training.""" + 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] + if max_frames is not None and len(indices) > max_frames > 0: + last = indices[-1] + if max_frames == 1: + # max_frames == 1 would divide by (max_frames - 1) == 0 below; keep only the last frame. + indices = [last] + else: + step = len(indices) / (max_frames - 1) + indices = [indices[int(i * step)] for i in range(max_frames - 1)] + indices.append(last) + return indices + + +async def get_video_tensor( + vr, + image_factor: int, + max_size: Tuple[int, int], + fps: Optional[float] = None, + frame_max_size: Optional[int] = None, + max_frames: Optional[int] = None, +) -> Tuple[torch.Tensor, dict]: + if fps is None: + fps = 1.0 + if frame_max_size is None: + frame_max_size = max_size[0] + if fps <= 0: + raise ValueError(f"video fps must be > 0, got {fps}") + + if isinstance(vr, torch.Tensor): + video_tchw = vr + _, _, height, width = video_tchw.shape + resized_width, resized_height = get_hw_multiple_of( + (width, height), image_factor, max_size + ) + resized = torchvision.transforms.functional.resize( + video_tchw, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + ) + return resized, { + "total_num_frames": resized.shape[0], + "fps": None, + "frames_indices": None, + } + + total_frames = len(vr) + video_fps = vr.avg_fps + if video_fps <= 0 or total_frames <= 0: + raise ValueError( + f"Invalid video metadata: fps={video_fps}, frames={total_frames}" + ) + indices = _compute_sampled_frame_indices(total_frames, video_fps, fps, max_frames) + video_tchw = vr.get_frames_as_tensor(indices) + video_tchw = video_tchw.permute(0, 3, 1, 2).float() + + _, _, height, width = video_tchw.shape + resized_width, resized_height = get_hw_multiple_of( + (width, height), image_factor, frame_max_size + ) + resized = torchvision.transforms.functional.resize( + video_tchw, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + ) + return resized, { + "total_num_frames": total_frames, + "fps": video_fps, + "frames_indices": indices, + } + + +class MiniMaxM3VLProcessor(BaseMultimodalProcessor): + models = [ + MiniMaxM3SparseForConditionalGeneration, + ] + + gpu_image_decode = False + + # M3's tokenizer has no pad_token. + tokenizer_padding = False + + IMAGE_TOKEN = "]<]image[>[" + VIDEO_TOKEN = "]<]video[>[" + IMAGE_START_TOKEN = "]<]start of image[>[" + IMAGE_END_TOKEN = "]<]end of image[>[" + + @staticmethod + def _token_id(tokenizer, token): + token_id = tokenizer.convert_tokens_to_ids(token) + assert token_id is not None, f"token id for {token!r} not found" + return token_id + + @property + def spatial_merge_size(self): + return self._processor.image_processor.merge_size + + def _video_resize_config(self): + video_processor = self._processor.video_processor + image_factor = video_processor.patch_size * video_processor.merge_size + # Newer M3 video processors expose max_pixels (area) instead of max_size; derive an equivalent (max_w, max_h) cap. + max_size = getattr(video_processor, "max_size", None) + if max_size is None: + max_pixels = getattr(video_processor, "max_pixels", None) + if max_pixels is not None: + side = int(math.isqrt(int(max_pixels))) + side -= side % image_factor + max_size = (side, side) + else: + max_size = video_processor._max_size_from_size(video_processor.size) + assert max_size is not None, "video processor max_size is required" + return image_factor, max_size + + def __init__(self, hf_config, server_args, _processor, *args, **kwargs): + super().__init__(hf_config, server_args, _processor, *args, **kwargs) + + tokenizer = _processor.tokenizer + assert tokenizer is not None, "tokenizer is required" + + self.IM_TOKEN_ID = self._token_id(tokenizer, self.IMAGE_TOKEN) + self.VIDEO_TOKEN_ID = self._token_id(tokenizer, self.VIDEO_TOKEN) + self.IM_START_TOKEN_ID = self._token_id(tokenizer, self.IMAGE_START_TOKEN) + self.IM_END_TOKEN_ID = self._token_id(tokenizer, self.IMAGE_END_TOKEN) + self.video_fps = self.video_config.pop("fps", None) + self.video_frame_max_size = self.video_config.pop("frame_max_size", None) + self.video_max_frames = self.video_config.pop("max_frames", None) + + self.mm_tokens = MultimodalSpecialTokens( + image_token=self.IMAGE_TOKEN, + image_token_id=self.IM_TOKEN_ID, + image_token_regex=re.compile( + r"|<\|image\|>|<\|image_pad\|>|\]\<\]image\[\>\[" + ), + video_token=self.VIDEO_TOKEN, + video_token_id=self.VIDEO_TOKEN_ID, + video_token_regex=re.compile(r"``, so a non-thinking reply may open with one stray closer; drop it unless thinking. + """ + + def __init__( + self, + stream_reasoning: bool = True, + force_reasoning: bool = False, + continue_final_message: bool = False, + previous_content: str = "", + force_nonempty_content: bool = False, + ): + super().__init__( + "", + "", + force_reasoning=force_reasoning, + stream_reasoning=stream_reasoning, + continue_final_message=continue_final_message, + previous_content=previous_content, + ) + self._lead_buffer = "" + self._checked_leading_close = False + self._force_nonempty_content = force_nonempty_content + + def detect_and_parse(self, text: str) -> StreamingParseResult: + if not self._in_reasoning and text.lstrip().startswith(self.think_end_token): + text = text.lstrip()[len(self.think_end_token) :] + ret = super().detect_and_parse(text) + if self._force_nonempty_content and not ret.normal_text: + ret.normal_text, ret.reasoning_text = ret.reasoning_text, ret.normal_text + return ret + + def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: + # ```` is a single token, so a stray leading closer arrives whole. + if not self._checked_leading_close and not self._in_reasoning: + self._lead_buffer += new_text + stripped = self._lead_buffer.lstrip() + if not stripped: + return StreamingParseResult() + self._checked_leading_close = True + if stripped.startswith(self.think_end_token): + new_text = stripped[len(self.think_end_token) :] + else: + new_text = self._lead_buffer + self._lead_buffer = "" + if not new_text: + return StreamingParseResult() + return super().parse_streaming_increment(new_text) + + class MistralDetector(BaseReasoningFormatDetector): """ Detector for Mistral models with reasoning (e.g., Mistral-Small-4-119B-2603). @@ -1089,6 +1142,7 @@ class ReasoningParser: "qwen3-thinking": Qwen3Detector, "minimax": Qwen3Detector, "minimax-append-think": MiniMaxAppendThinkDetector, + "minimax-m3": MiniMaxM3Detector, "step3": DeepSeekR1Detector, "step3p5": DeepSeekR1Detector, "mistral": MistralDetector, @@ -1113,6 +1167,8 @@ class ReasoningParser: if not detector_class: raise ValueError(f"Unsupported model type: {model_type}") + chat_template_kwargs = getattr(request, "chat_template_kwargs", None) or {} + # Special cases where we override force_reasoning if model_type.lower() in { "qwen3-thinking", @@ -1121,6 +1177,11 @@ class ReasoningParser: }: force_reasoning = True + # M3 consumes the start tag only for thinking_mode=enabled + # (absent from output → must force); mirror serving_chat's M3 branch. + if model_type.lower() == "minimax-m3" and force_reasoning is None: + force_reasoning = chat_template_kwargs.get("thinking_mode") == "enabled" + # Only pass force_reasoning if explicitly set, let detectors use their defaults kwargs = {"stream_reasoning": stream_reasoning} if force_reasoning is not None: @@ -1135,7 +1196,6 @@ class ReasoningParser: kwargs["continue_final_message"] = True kwargs["previous_content"] = request.messages[-1].content - chat_template_kwargs = getattr(request, "chat_template_kwargs", None) or {} if chat_template_kwargs.get("force_nonempty_content") is True: kwargs["force_nonempty_content"] = True diff --git a/python/sglang/srt/parser/template_detection.py b/python/sglang/srt/parser/template_detection.py index 24a1a1b50..489fa820a 100644 --- a/python/sglang/srt/parser/template_detection.py +++ b/python/sglang/srt/parser/template_detection.py @@ -293,6 +293,10 @@ def _is_minimax(ctx): return ctx.has_text("") +def _is_minimax_m3(ctx): + return ctx.has_text("") or ctx.has_text("]<]minimax[>[") + + def _is_minicpm5(ctx): if ctx.has_vocab(" torch.Tensor: @@ -171,6 +171,7 @@ def _make_sglang_bench_fn( group_size: int, dst_dtype: torch.dtype, flags: dict, + provider: str = "sglang", ): """ Adapter that pre-allocates output tensors and returns a zero-arg callable @@ -213,17 +214,37 @@ def _make_sglang_bench_fn( scale_ue8m0=scale_ue8m0, ) - def _run(): - sglang_per_token_group_quant_8bit( - input=x_input, - output_q=output_q, - output_s=output_s, - group_size=group_size, - eps=1e-10, - fp8_min=fp8_min, - fp8_max=fp8_max, - scale_ue8m0=scale_ue8m0, - ) + if provider == "aot_v2": + from sgl_kernel import sgl_per_token_group_quant_8bit as aot_quant + + def _run(): + aot_quant( + x_input, + output_q, + output_s, + group_size, + 1e-10, + fp8_min, + fp8_max, + scale_ue8m0, + False, # fuse_silu_and_mul (already applied to x_input) + None, # masked_m (flattened to 2D) + enable_v2=True, + ) + + else: + + def _run(): + sglang_per_token_group_quant_8bit( + input=x_input, + output_q=output_q, + output_s=output_s, + group_size=group_size, + eps=1e-10, + fp8_min=fp8_min, + fp8_max=fp8_max, + scale_ue8m0=scale_ue8m0, + ) return _run @@ -270,13 +291,14 @@ def benchmark( dst_dtype=dst_dtype, **{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]}, ) - elif provider == "sglang": + elif provider in ("sglang", "aot_v2"): kernel_names = "per_token_group_quant_8bit_kernel" bench_fn = _make_sglang_bench_fn( x=x, group_size=group_size, dst_dtype=dst_dtype, flags=flags, + provider=provider, ) else: raise ValueError(f"Unknown provider: {provider}") diff --git a/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py b/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py new file mode 100644 index 000000000..874ab4b41 --- /dev/null +++ b/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py @@ -0,0 +1,79 @@ +import torch + +from sglang.jit_kernel.benchmark import marker +from sglang.srt.layers.moe.ep_moe.kernels import ( + post_reorder_deepgemm, + post_reorder_triton_kernel, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +HIDDEN = 6144 +NUM_EXPERTS = 129 +TOP_K = 5 +RSF = 2.0 + + +def _build(num_tokens): + m_max = (num_tokens // 256 + 1) * 256 + down_output = torch.randn( + NUM_EXPERTS * m_max, HIDDEN, dtype=torch.bfloat16, device="cuda" + ) + topk_ids = torch.randint( + 0, NUM_EXPERTS, (num_tokens, TOP_K), dtype=torch.int32, device="cuda" + ) + src2dst = ( + topk_ids.long() * m_max + + torch.randint(0, m_max, (num_tokens, TOP_K), device="cuda") + ).to(torch.int32) + topk_weights = torch.rand(num_tokens, TOP_K, dtype=torch.float32, device="cuda") + return down_output, src2dst, topk_ids, topk_weights + + +def _new(down_output, src2dst, topk_ids, topk_weights): + num_tokens = topk_ids.shape[0] + out = torch.empty(num_tokens, HIDDEN, dtype=torch.bfloat16, device="cuda") + post_reorder_deepgemm( + down_output, + out, + src2dst, + topk_ids, + topk_weights, + TOP_K, + num_tokens, + HIDDEN, + RSF, + ) + return out + + +def _old(down_output, src2dst, topk_ids, topk_weights): + num_tokens = topk_ids.shape[0] + out = torch.empty(num_tokens, HIDDEN, dtype=torch.bfloat16, device="cuda") + post_reorder_triton_kernel[(num_tokens,)]( + down_output, out, src2dst, topk_ids, topk_weights, TOP_K, HIDDEN, BLOCK_SIZE=512 + ) + out *= RSF + return out + + +FN_MAP = {"fused": _new, "legacy": _old} + + +@marker.parametrize("num_tokens", [1, 8, 64, 256, 1024, 4096, 16384], [64, 4096]) +@marker.benchmark("impl", ["fused", "legacy"]) +def benchmark(num_tokens: int, impl: str): + down_output, src2dst, topk_ids, topk_weights = _build(num_tokens) + return marker.do_bench( + FN_MAP[impl], + input_args=(down_output, src2dst, topk_ids, topk_weights), + graph_clone_args=(0,), + memory_args=None, + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/jit/minimax/test_minimax_fused_qkv_index_gemm.py b/test/registered/jit/minimax/test_minimax_fused_qkv_index_gemm.py new file mode 100644 index 000000000..4c79af7f4 --- /dev/null +++ b/test/registered/jit/minimax/test_minimax_fused_qkv_index_gemm.py @@ -0,0 +1,69 @@ +import pytest +import torch + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +dev = "cuda" + + +def _pack_weight_scale(scale_u8: torch.Tensor) -> torch.Tensor: + from sglang.srt.layers.deep_gemm_wrapper.configurer import DEEPGEMM_SCALE_UE8M0 + + n, kk = scale_u8.shape + scale_fp32 = ( + (scale_u8.contiguous().view(-1).to(torch.int32) << 23) + .view(torch.float32) + .view(n, kk) + ) + if DEEPGEMM_SCALE_UE8M0: + import deep_gemm.utils.layout + + return deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor( + scale_fp32 + ) + return scale_fp32 + + +@pytest.mark.parametrize("T", [1, 7, 64, 256, 1024]) +@pytest.mark.parametrize("N1,N2,K", [(1280, 256, 6144), (1024, 128, 2048)]) +def test_fused_equals_separate(T, N1, N2, K): + from sglang.srt.layers.quantization.fp8_utils import ( + _deepgemm_w8a8_mxfp8_linear_with_fallback as mxfp8_linear, + ) + + torch.manual_seed(T * 7 + N1 + K) + G = K // 32 + + def rand_w(n): + return (torch.randn(n, K, device=dev) * 0.2).to(torch.float8_e4m3fn) + + def rand_s(n): + # mid-range UE8M0 exponents (~2^-3..2^3) to avoid inf/zero blowups. + return torch.randint(124, 131, (n, G), device=dev, dtype=torch.uint8) + + w1, w2 = rand_w(N1), rand_w(N2) + s1, s2 = rand_s(N1), rand_s(N2) + x = torch.randn(T, K, dtype=torch.bfloat16, device=dev) + + s1p, s2p = _pack_weight_scale(s1), _pack_weight_scale(s2) + out1 = mxfp8_linear(x, w1, s1p, weight_scale_fallback=s1) + out2 = mxfp8_linear(x, w2, s2p, weight_scale_fallback=s2) + ref = torch.cat([out1, out2], dim=-1) + + w = torch.cat([w1, w2], dim=0).contiguous() + s = torch.cat([s1, s2], dim=0).contiguous() + sp = _pack_weight_scale(s) + fused = mxfp8_linear(x, w, sp, weight_scale_fallback=s) + + assert fused.shape == ref.shape + assert torch.equal( + fused, ref + ), f"max abs diff {(fused.float() - ref.float()).abs().max().item()}" + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/jit/minimax/test_minimax_quant_scatter.py b/test/registered/jit/minimax/test_minimax_quant_scatter.py new file mode 100644 index 000000000..3e4ea6b07 --- /dev/null +++ b/test/registered/jit/minimax/test_minimax_quant_scatter.py @@ -0,0 +1,84 @@ +import random +import sys + +import pytest +import torch + +from sglang.jit_kernel.minimax_quant_ue8m0 import ( + per_token_quant_fp8_ue8m0, + per_token_quant_fp8_ue8m0_scatter, +) +from sglang.srt.layers.moe.ep_moe.kernels import fill_gateup_input_triton_kernel +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +dev = "cuda" + + +@pytest.mark.parametrize("num_tokens", [1, 7, 64, 256]) +@pytest.mark.parametrize("topk", [4, 5, 8]) +@pytest.mark.parametrize("hidden,group", [(6144, 32), (2048, 32), (4096, 128)]) +def test_quant_scatter_matches_quant_plus_fill(num_tokens, topk, hidden, group): + arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device()) + if arch_major <= 9: + pytest.skip("UE8M0 fusion is Blackwell-only") + + E = 129 # 128 routed + 1 fused shared + G4 = (hidden // group) // 4 + m_max = (num_tokens // 256 + 1) * 256 + torch.manual_seed(num_tokens * 91 + topk + hidden) + random.seed(num_tokens) + + x = (torch.randn(num_tokens, hidden, device=dev, dtype=torch.bfloat16)) * 4.0 + + tids = torch.empty(num_tokens, topk, dtype=torch.int32, device=dev) + tids_cpu = torch.empty(num_tokens, topk, dtype=torch.int32) + s2d = [0] * (num_tokens * topk) + cur = [0] * E + for t in range(num_tokens): + for j, e in enumerate(random.sample(range(E), topk)): + tids_cpu[t, j] = e + s2d[t * topk + j] = e * m_max + cur[e] + cur[e] += 1 + tids.copy_(tids_cpu) + s2d = torch.tensor(s2d, dtype=torch.int32, device=dev) + + x_q, x_sf = per_token_quant_fp8_ue8m0(x, group) + gi_ref = torch.zeros(E, m_max, hidden, device=dev, dtype=torch.float8_e4m3fn) + gs_ref = torch.zeros(E, G4, m_max, device=dev, dtype=torch.int32) + fill_gateup_input_triton_kernel[(num_tokens,)]( + x_q, + x_sf, + gi_ref, + gs_ref, + s2d, + tids, + topk, + hidden, + G4, + m_max, + x_sf.stride(0), + x_sf.stride(1), + BLOCK_SIZE=1024, + SCALE_MN_MAJOR=True, + ) + + gi_new = torch.zeros(E, m_max, hidden, device=dev, dtype=torch.float8_e4m3fn) + gs_new = torch.zeros(E, G4, m_max, device=dev, dtype=torch.int32) + per_token_quant_fp8_ue8m0_scatter(x, gi_new, gs_new, s2d, tids, topk, m_max, group) + + for t in range(num_tokens): + for j in range(topk): + e = int(tids_cpu[t, j]) + m = int(s2d[t * topk + j]) % m_max + assert torch.equal( + gi_new[e, m].view(torch.uint8), gi_ref[e, m].view(torch.uint8) + ), f"fp8 mismatch token={t} slot={j} expert={e}" + assert torch.equal( + gs_new[e, :, m], gs_ref[e, :, m] + ), f"scale mismatch token={t} slot={j} expert={e}" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-x"])) diff --git a/test/registered/jit/test_per_token_group_quant_8bit.py b/test/registered/jit/test_per_token_group_quant_8bit.py index 25a139212..6f5073960 100644 --- a/test/registered/jit/test_per_token_group_quant_8bit.py +++ b/test/registered/jit/test_per_token_group_quant_8bit.py @@ -17,11 +17,17 @@ register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) if not torch.cuda.is_available(): pytest.skip("CUDA required", allow_module_level=True) +from sgl_kernel import ( # noqa: E402 + sgl_per_token_group_quant_8bit as aot_per_token_group_quant_8bit, +) from sgl_kernel.test_utils import ( # noqa: E402 assert_all_close_or_tiny_diff, create_per_token_group_quant_test_data, ) +from sglang.jit_kernel.per_token_group_quant_8bit import ( # noqa: E402 + per_token_group_quant_8bit as jit_per_token_group_quant_8bit, +) from sglang.srt.layers.quantization.fp8_kernel import ( # noqa: E402 create_per_token_group_quant_fp8_output_scale, ) @@ -224,5 +230,104 @@ def test_per_token_group_quant_with_column_major( raise +LAYOUTS = [ + (False, False, False), + (True, False, False), + (True, True, False), + (True, True, True), +] + +CONFIGS = list( + itertools.product( + [1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192], + [512, 1536, 2048, 4096, 6144, 7168, 16384], + [16, 32, 64, 128], + LAYOUTS, + [fp8_type_], + ) +) + + +@pytest.mark.parametrize( + "num_tokens, hidden_dim, group_size, layout, dst_dtype", CONFIGS +) +def test_jit_matches_aot_v2_byte_identical( + num_tokens, hidden_dim, group_size, layout, dst_dtype +): + column_major_scales, scale_tma_aligned, scale_ue8m0 = layout + + arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device()) + if scale_ue8m0 and arch_major <= 9: + pytest.skip("UE8M0 fusion is Blackwell-only") + if hidden_dim % group_size != 0: + pytest.skip("hidden_dim must be divisible by group_size") + + torch.manual_seed(num_tokens * 131 + hidden_dim + group_size) + x = (torch.randn(num_tokens, hidden_dim, device="cuda", dtype=torch.bfloat16)) * 3.0 + + fp8_max = torch.finfo(dst_dtype).max + fp8_min = -fp8_max + + def _alloc(): + q = torch.empty_like(x, dtype=dst_dtype) + s = create_per_token_group_quant_fp8_output_scale( + x_shape=x.shape, + device=x.device, + group_size=group_size, + column_major_scales=column_major_scales, + scale_tma_aligned=scale_tma_aligned, + scale_ue8m0=scale_ue8m0, + ) + return q, s + + q_aot, s_aot = _alloc() + aot_per_token_group_quant_8bit( + x, + q_aot, + s_aot, + group_size, + 1e-10, + fp8_min, + fp8_max, + scale_ue8m0, + False, + None, + enable_v2=True, + ) + + q_jit, s_jit = _alloc() + jit_per_token_group_quant_8bit( + x, q_jit, s_jit, group_size, 1e-10, fp8_min, fp8_max, scale_ue8m0=scale_ue8m0 + ) + + # AOT v2 uses -use_fast_math reciprocal; this JIT uses precise division, so an + # exact fp8 midpoint can round to an adjacent code (1-ULP, JIT more accurate). + qj = q_jit.view(torch.uint8) + qa = q_aot.view(torch.uint8) + if not torch.equal(qj, qa): + mism = qj != qa + bj = qj[mism].to(torch.int16) + ba = qa[mism].to(torch.int16) + same_sign = (bj & 0x80) == (ba & 0x80) + one_ulp = (bj - ba).abs() == 1 + assert bool( + (same_sign & one_ulp).all() + ), f"q mismatch > 1 fp8 ULP {num_tokens=} {hidden_dim=} {group_size=} {layout=}" + assert mism.float().mean() < 0.01, ( + f"too many fp8 ties ({int(mism.sum())}/{mism.numel()}) " + f"{num_tokens=} {hidden_dim=} {group_size=} {layout=}" + ) + + if scale_ue8m0: + assert torch.equal( + s_jit[:num_tokens].reshape(num_tokens, -1).view(torch.int32), + s_aot[:num_tokens].reshape(num_tokens, -1).view(torch.int32), + ), f"ue8m0 scale mismatch {num_tokens=} {hidden_dim=} {group_size=}" + else: + assert torch.equal( + s_jit[:num_tokens].float(), s_aot[:num_tokens].float() + ), f"float scale mismatch {num_tokens=} {hidden_dim=} {group_size=}" + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/jit/test_post_reorder_deepgemm.py b/test/registered/jit/test_post_reorder_deepgemm.py new file mode 100644 index 000000000..181e315e4 --- /dev/null +++ b/test/registered/jit/test_post_reorder_deepgemm.py @@ -0,0 +1,90 @@ +import pytest +import torch + +from sglang.srt.layers.moe.ep_moe.kernels import ( + post_reorder_deepgemm, + post_reorder_triton_kernel, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +dev = "cuda" + + +def _build( + num_tokens, + hidden, + num_routed=128, + top_k_routed=4, + with_shared=True, + pad_frac=0.0, + seed=0, +): + num_experts = num_routed + (1 if with_shared else 0) + top_k = top_k_routed + (1 if with_shared else 0) + m_max = (num_tokens // 256 + 1) * 256 + g = torch.Generator(device="cpu").manual_seed(seed) + down_output = torch.randn( + num_experts * m_max, hidden, dtype=torch.bfloat16, device=dev + ) + topk_ids = torch.full((num_tokens, top_k), -1, dtype=torch.int32, device=dev) + src2dst = torch.full((num_tokens, top_k), -1, dtype=torch.int32, device=dev) + topk_weights = torch.rand(num_tokens, top_k, dtype=torch.float32, device=dev) + counts = [0] * num_experts + for t in range(num_tokens): + experts = torch.randperm(num_routed, generator=g)[:top_k_routed].tolist() + if with_shared: + experts = experts + [num_routed] + for slot, e in enumerate(experts): + if ( + pad_frac > 0 + and slot < top_k - 1 + and torch.rand(1, generator=g).item() < pad_frac + ): + continue + src2dst[t, slot] = e * m_max + counts[e] + counts[e] += 1 + topk_ids[t, slot] = e + return down_output, src2dst, topk_ids, topk_weights, top_k + + +def _ref(down_output, src2dst, topk_ids, topk_weights, top_k, hidden, rsf): + num_tokens = topk_ids.shape[0] + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=dev) + do = down_output.float() + for slot in range(top_k): + valid = topk_ids[:, slot] >= 0 + dst = src2dst[:, slot].long().clamp_min(0) + out[valid] += (do[dst] * topk_weights[:, slot, None])[valid] + return out * rsf + + +@pytest.mark.parametrize("num_tokens", [1, 8, 128, 1024, 4096]) +@pytest.mark.parametrize("pad_frac", [0.0, 0.3]) +@pytest.mark.parametrize("rsf", [1.0, 2.0]) +def test_post_reorder_deepgemm(num_tokens, pad_frac, rsf): + hidden = 6144 + do, s2d, tids, tw, tk = _build( + num_tokens, hidden, pad_frac=pad_frac, seed=num_tokens + ) + ref = _ref(do, s2d, tids, tw, tk, hidden, rsf) + + new = torch.empty(num_tokens, hidden, dtype=torch.bfloat16, device=dev) + post_reorder_deepgemm(do, new, s2d, tids, tw, tk, num_tokens, hidden, rsf) + + old = torch.empty(num_tokens, hidden, dtype=torch.bfloat16, device=dev) + post_reorder_triton_kernel[(num_tokens,)]( + do, old, s2d, tids, tw, tk, hidden, BLOCK_SIZE=512 + ) + old *= rsf + + torch.testing.assert_close(new.float(), ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(new.float(), old.float(), rtol=5e-2, atol=0.5) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/unit/function_call/test_minimax_m3_detector.py b/test/registered/unit/function_call/test_minimax_m3_detector.py new file mode 100644 index 000000000..8ffe51399 --- /dev/null +++ b/test/registered/unit/function_call/test_minimax_m3_detector.py @@ -0,0 +1,539 @@ +import json +import unittest + +from sglang.srt.entrypoints.openai.protocol import Function, Tool +from sglang.srt.function_call.minimax_m3 import MINIMAX_NS_TOKEN, MinimaxM3Detector +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=7, suite="base-a-test-cpu") + +NS = MINIMAX_NS_TOKEN + + +def _make_tools(): + return [ + Tool( + type="function", + function=Function( + name="get_current_date", + description="Get the current date", + parameters={}, + ), + ), + Tool( + type="function", + function=Function( + name="get_weather", + description="Get weather information", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + }, + "required": ["city"], + }, + ), + ), + Tool( + type="function", + function=Function( + name="search", + description="Search the web", + parameters={ + "type": "object", + "properties": { + "query": {"type": "string"}, + "count": {"type": "integer"}, + "ratio": {"type": "number"}, + "verbose": {"type": "boolean"}, + }, + "required": ["query"], + }, + ), + ), + Tool( + type="function", + function=Function( + name="create_event", + description="Create a calendar event", + parameters={ + "type": "object", + "properties": { + "title": {"type": "string"}, + "location": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "zip": {"type": "integer"}, + }, + }, + "tags": { + "type": "array", + "items": {"type": "string"}, + }, + }, + }, + ), + ), + Tool( + type="function", + function=Function( + name="add_note", + description="Add a free-form note", + parameters={ + "type": "object", + "properties": {"note": {"type": "string"}}, + }, + ), + ), + Tool( + type="function", + function=Function( + name="configure", + description="Configure runtime options", + parameters={ + "type": "object", + "properties": { + "mode": { + "type": "string", + "enum": ["none", "low", "high"], + }, + "optional": {"type": ["string", "null"]}, + }, + }, + ), + ), + ] + + +def _wire(*lines): + return "".join(NS + line for line in lines) + + +def _segments(*lines): + return [NS + line for line in lines] + + +def _collect_streamed_tool_calls(all_calls): + tools = {} + for c in all_calls: + idx = c.tool_index + if idx not in tools: + tools[idx] = {"name": c.name or "", "parameters": c.parameters or ""} + else: + if c.name: + tools[idx]["name"] += c.name + if c.parameters: + tools[idx]["parameters"] += c.parameters + return [tools[i] for i in sorted(tools.keys())] + + +def _stream_segments(segments, tools): + detector = MinimaxM3Detector() + all_calls = [] + for seg in _segments(*segments): + all_calls.extend(detector.parse_streaming_increment(seg, tools).calls) + collected = _collect_streamed_tool_calls(all_calls) + return [{"name": c["name"], "args": json.loads(c["parameters"])} for c in collected] + + +def _parse_segments(segments, tools): + detector = MinimaxM3Detector() + result = detector.detect_and_parse(_wire(*segments), tools) + return [ + {"name": c.name, "args": json.loads(c.parameters)} for c in result.calls + ], result.normal_text + + +class TestMinimaxM3HasToolCall(CustomTestCase): + def setUp(self): + self.detector = MinimaxM3Detector() + + def test_has_tool_call_true(self): + text = _wire( + "", + '', + "", + "", + ) + self.assertTrue(self.detector.has_tool_call(text)) + + def test_has_tool_call_false(self): + self.assertFalse( + self.detector.has_tool_call("The weather in Beijing is sunny.") + ) + + +class TestMinimaxM3DetectAndParse(CustomTestCase): + def setUp(self): + self.tools = _make_tools() + + def test_no_tool_call(self): + text = "This is a plain response with no tool call." + calls, normal = _parse_segments_text(text, self.tools) + self.assertEqual(len(calls), 0) + self.assertEqual(normal, text) + + def test_single_tool_call(self): + segments = ( + "", + '', + "Beijing", + "", + "", + "", + ) + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["name"], "get_weather") + self.assertEqual(calls[0]["args"], {"city": "Beijing"}) + + def test_zero_arg_tool_call(self): + segments = ( + "", + '', + "", + "", + ) + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["name"], "get_current_date") + self.assertEqual(calls[0]["args"], {}) + + def test_multiple_tool_calls_separate_blocks(self): + segments = ( + "", + '', + "Beijing", + "", + "", + "", + "", + '', + "Tokyo", + "", + "", + "", + ) + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(len(calls), 2) + self.assertEqual(calls[0]["args"]["city"], "Beijing") + self.assertEqual(calls[1]["args"]["city"], "Tokyo") + + def test_multiple_invokes_one_block(self): + segments = ( + "", + '', + "Beijing", + "", + "", + '', + "Tokyo", + "", + "", + "", + ) + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(len(calls), 2) + self.assertEqual(calls[0]["args"]["city"], "Beijing") + self.assertEqual(calls[1]["args"]["city"], "Tokyo") + + def test_content_before_tool_call(self): + segments = ( + "", + '', + "Paris", + "", + "", + "", + ) + detector = MinimaxM3Detector() + text = "Let me check the weather." + _wire(*segments) + result = detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.normal_text, "Let me check the weather.") + + def test_typed_scalars(self): + segments = ( + "", + '', + "pizza", + "", + "7", + "", + "0.5", + "", + "true", + "", + "", + "", + ) + calls, _ = _parse_segments(segments, self.tools) + args = calls[0]["args"] + self.assertEqual(args["query"], "pizza") + self.assertEqual(args["count"], 7) + self.assertIsInstance(args["count"], int) + self.assertAlmostEqual(args["ratio"], 0.5) + self.assertIs(args["verbose"], True) + + def test_nested_object_and_array(self): + segments = ( + "", + '', + "Standup", + "", + "", + "NYC", + "", + "10001", + "", + "", + "", + "red", + "", + "blue", + "", + "", + "", + "", + ) + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(len(calls), 1) + self.assertEqual( + calls[0]["args"], + { + "title": "Standup", + "location": {"city": "NYC", "zip": 10001}, + "tags": ["red", "blue"], + }, + ) + + def test_string_with_special_characters(self): + value = 'a "quoted" word with \\ backslash and\nnewline' + segments = ( + "", + '', + "" + value, + "", + "", + "", + ) + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(calls[0]["args"], {"note": value}) + + +class TestMinimaxM3NoneNullRegression(CustomTestCase): + def setUp(self): + self.tools = _make_tools() + + def _configure_segments(self, param, value): + return ( + "", + '', + "<{}>{}".format(param, value), + "".format(param), + "", + "", + ) + + def test_plain_string_none_not_coerced(self): + for value in ("none", "nil", "null"): + with self.subTest(value=value): + segments = self._configure_segments("mode", value) + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(calls[0]["args"], {"mode": value}) + self.assertIsInstance(calls[0]["args"]["mode"], str) + + def test_plain_string_none_not_coerced_streaming(self): + for value in ("none", "nil", "null"): + with self.subTest(value=value): + segments = self._configure_segments("mode", value) + calls = _stream_segments(segments, self.tools) + self.assertEqual(calls[0]["args"], {"mode": value}) + self.assertIsInstance(calls[0]["args"]["mode"], str) + + def test_streaming_and_non_streaming_agree_for_string(self): + for value in ("none", "nil", "null"): + with self.subTest(value=value): + segments = self._configure_segments("mode", value) + non_stream, _ = _parse_segments(segments, self.tools) + stream = _stream_segments(segments, self.tools) + self.assertEqual(non_stream, stream) + + def test_nullable_param_null_becomes_none(self): + segments = self._configure_segments("optional", "null") + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(calls[0]["args"], {"optional": None}) + + def test_nullable_param_none_stays_string(self): + segments = self._configure_segments("optional", "none") + calls, _ = _parse_segments(segments, self.tools) + self.assertEqual(calls[0]["args"], {"optional": "none"}) + + +class TestMinimaxM3Streaming(CustomTestCase): + def setUp(self): + self.tools = _make_tools() + + def test_normal_text_only(self): + detector = MinimaxM3Detector() + result = detector.parse_streaming_increment("Hello there.", self.tools) + self.assertEqual(result.normal_text, "Hello there.") + self.assertEqual(len(result.calls), 0) + + def test_single_tool_call_chunked(self): + segments = ( + "", + '', + "Beijing", + "", + "", + "", + ) + calls = _stream_segments(segments, self.tools) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["name"], "get_weather") + self.assertEqual(calls[0]["args"], {"city": "Beijing"}) + + def test_streaming_matches_non_streaming(self): + cases = { + "weather": ( + "", + '', + "Beijing", + "", + "", + "", + ), + "typed": ( + "", + '', + "pizza", + "", + "7", + "", + "0.5", + "", + "true", + "", + "", + "", + ), + "nested": ( + "", + '', + "Standup", + "", + "", + "NYC", + "", + "10001", + "", + "", + "", + "red", + "", + "blue", + "", + "", + "", + "", + ), + "special": ( + "", + '', + "" + 'a "q" and \\ back\nslash', + "", + "", + "", + ), + "multi_invoke": ( + "", + '', + "Beijing", + "", + "", + '', + "Tokyo", + "", + "", + "", + ), + } + for name, segments in cases.items(): + with self.subTest(case=name): + non_stream, _ = _parse_segments(segments, self.tools) + stream = _stream_segments(segments, self.tools) + self.assertEqual(non_stream, stream) + + def test_streaming_sequential_tool_index(self): + segments = ( + "", + '', + "Beijing", + "", + "", + '', + "Tokyo", + "", + "", + "", + ) + detector = MinimaxM3Detector() + all_calls = [] + for seg in _segments(*segments): + all_calls.extend(detector.parse_streaming_increment(seg, self.tools).calls) + self.assertEqual(sorted({c.tool_index for c in all_calls}), [0, 1]) + + +class TestMinimaxM3Malformed(CustomTestCase): + def setUp(self): + self.tools = _make_tools() + + def test_truncated_no_closing_tags(self): + text = _wire( + "", + '', + "Beijing", + ) + detector = MinimaxM3Detector() + result = detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 0) + self.assertEqual(result.normal_text, text) + + def test_mismatched_closing_tag(self): + text = _wire( + "", + '', + "Beijing", + "", + "", + "", + ) + detector = MinimaxM3Detector() + result = detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 0) + self.assertEqual(result.normal_text, text) + + def test_truncated_streaming_does_not_crash(self): + segments = ( + "", + '', + "Beijing", + ) + detector = MinimaxM3Detector() + for seg in _segments(*segments): + detector.parse_streaming_increment(seg, self.tools) + + +def _parse_segments_text(text, tools): + detector = MinimaxM3Detector() + result = detector.detect_and_parse(text, tools) + return [ + {"name": c.name, "args": json.loads(c.parameters)} for c in result.calls + ], result.normal_text + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/parser/test_reasoning_parser.py b/test/registered/unit/parser/test_reasoning_parser.py index 8bc59efe4..ef0d26b8a 100644 --- a/test/registered/unit/parser/test_reasoning_parser.py +++ b/test/registered/unit/parser/test_reasoning_parser.py @@ -762,6 +762,96 @@ class TestMiniMaxAppendThinkDetector(CustomTestCase): self.assertEqual(result.normal_text, "Second") +class TestMiniMaxM3Detector(CustomTestCase): + """Test cases for MiniMaxM3Detector multi-turn stray-closer handling.""" + + def _detector(self, force_reasoning=False): + from sglang.srt.parser.reasoning_parser import MiniMaxM3Detector + + return MiniMaxM3Detector(force_reasoning=force_reasoning) + + def test_drops_leading_stray_close_non_stream(self): + """Non-thinking multi-turn reply opening with a stray .""" + result = self._detector().detect_and_parse("The answer is 42.") + self.assertEqual(result.normal_text, "The answer is 42.") + self.assertEqual(result.reasoning_text or "", "") + + def test_drops_leading_stray_close_with_whitespace(self): + result = self._detector().detect_and_parse("\nHello") + self.assertEqual(result.normal_text, "Hello") + + def test_plain_reply_untouched(self): + result = self._detector().detect_and_parse("Just a normal answer.") + self.assertEqual(result.normal_text, "Just a normal answer.") + + def test_real_reasoning_block_non_stream(self): + result = self._detector().detect_and_parse( + "reasoning herefinal" + ) + self.assertEqual(result.reasoning_text, "reasoning here") + self.assertEqual(result.normal_text, "final") + + def test_thinking_mode_close_not_dropped(self): + """force_reasoning=True means the closer ends reasoning, not a stray drop.""" + result = self._detector(force_reasoning=True).detect_and_parse( + "reasoninganswer" + ) + self.assertEqual(result.reasoning_text, "reasoning") + self.assertEqual(result.normal_text, "answer") + + def test_drops_leading_stray_close_stream_single_token(self): + detector = self._detector() + first = detector.parse_streaming_increment("") + self.assertEqual(first.normal_text or "", "") + second = detector.parse_streaming_increment("The answer.") + self.assertEqual(second.normal_text, "The answer.") + + def test_drops_leading_stray_close_stream_after_whitespace(self): + """A leading whitespace token before the atomic is buffered.""" + detector = self._detector() + self.assertEqual(detector.parse_streaming_increment(" ").normal_text or "", "") + self.assertEqual( + detector.parse_streaming_increment("").normal_text or "", "" + ) + self.assertEqual( + detector.parse_streaming_increment("Hello").normal_text, "Hello" + ) + + def test_plain_reply_stream_untouched(self): + detector = self._detector() + out = detector.parse_streaming_increment("Hello") + self.assertEqual(out.normal_text, "Hello") + + def test_minimax_m3_model_type(self): + from sglang.srt.parser.reasoning_parser import MiniMaxM3Detector + + parser = ReasoningParser("minimax-m3") + self.assertIsInstance(parser.detector, MiniMaxM3Detector) + + def test_force_nonempty_content_via_chat_template_kwargs(self): + """force_nonempty_content must reach the M3 detector without a TypeError.""" + from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionMessageUserParam, + ChatCompletionRequest, + ) + + request = ChatCompletionRequest( + model="test", + messages=[ChatCompletionMessageUserParam(role="user", content="Hi")], + chat_template_kwargs={"force_nonempty_content": True}, + ) + parser = ReasoningParser("minimax-m3", request=request) + self.assertTrue(parser.detector._force_nonempty_content) + + def test_force_nonempty_content_swaps_when_no_content(self): + from sglang.srt.parser.reasoning_parser import MiniMaxM3Detector + + detector = MiniMaxM3Detector(force_reasoning=True, force_nonempty_content=True) + result = detector.detect_and_parse("only reasoning, no closer") + self.assertEqual(result.normal_text, "only reasoning, no closer") + self.assertEqual(result.reasoning_text or "", "") + + class TestReasoningParserAdvanced(CustomTestCase): """Additional tests for ReasoningParser init edge cases.""" diff --git a/test/registered/unit/parser/test_template_manager.py b/test/registered/unit/parser/test_template_manager.py index bc78175a3..85fb3a75b 100644 --- a/test/registered/unit/parser/test_template_manager.py +++ b/test/registered/unit/parser/test_template_manager.py @@ -130,6 +130,27 @@ class TestTemplateManagerReasoningDetection(unittest.TestCase): self.assertIsNone(config) self.assertEqual(parser, "minimax") + MINIMAX_M3_TEMPLATE = ( + "{%- set ns_token = ']<]minimax[>[' -%}\n" + "{%- set toolcall_begin_token = ns_token ~ '' -%}\n" + "\n" + ) + + def test_minimax_m3_detected_via_mm_think_signature(self): + _, config, parser = self._detect( + self.MINIMAX_M3_TEMPLATE, ["", ""] + ) + + self.assertIsNone(config) + self.assertEqual(parser, "minimax-m3") + + def test_minimax_m2_not_misclassified_as_m3(self): + template = """ + {%- set toolcall_begin_token = '' -%} + """ + _, _, parser = self._detect(template, [""]) + self.assertEqual(parser, "minimax") + class TestTemplateDetectionRuleMatrix(unittest.TestCase): """Table-driven tests for REASONING_PARSER_RULES and REASONING_MODE_RULES.""" @@ -376,6 +397,13 @@ class TestToolCallParserDetection(unittest.TestCase): ("gpt_oss", "<|channel|>analysis<|message|>", [], "gpt-oss"), ("gemma4", "<|channel>content", [], "gemma4"), ("minimax_maps_to_m2", "", [], "minimax-m2"), + ( + "minimax_m3_ns_token_only", + "{%- set ns_token = ']<]minimax[>[' -%}\n" + "{%- set toolcall_begin_token = ns_token ~ '' -%}", + [], + "minimax-m3", + ), ( "deepseekv3", "{% if not thinking is defined %}{% set thinking = false %}{% endif %}", @@ -583,6 +611,24 @@ class TestToolCallParserDetection(unittest.TestCase): self.assertLess(minicpm5_idx, rule_names.index("mimo")) self.assertLess(minicpm5_idx, rule_names.index("qwen")) + def test_minimax_m3_rule_precedes_m2_in_both_registries(self): + for rules in (REASONING_PARSER_RULES, TOOL_CALL_PARSER_RULES): + names = [rule.name for rule in rules] + self.assertLess(names.index("minimax_m3"), names.index("minimax")) + + def test_minicpm5_not_misclassified_as_qwen(self): + template = ( + "{% set enable_thinking = enable_thinking if enable_thinking is defined else true %}" + '\n' + '\n{{ param.value }}' + "\n" + ) + force, config = detect_reasoning_pattern(template) + result = detect_tool_call_parser( + template, _DummyTokenizer(["