minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[minimax-m3] Split 4/4: model + VL + glue + function-call + fp8 quant + generic infra (#28715)

Browse Source 
Co-authored-by: Xinyuan Tong <xinyuan-tong@users.noreply.github.com>
Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com>
This commit is contained in:
Xinyuan Tong
2026-07-11 11:11:06 +08:00
committed by GitHub
co-authored by Xinyuan Tong zijiexia
parent e3ceccf781
commit 0663ebc783
45 changed files with 7477 additions and 475 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh
+186 -153
View File
@@ -1,149 +1,202 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/atomic.cuh>
#include <sgl_kernel/cta.cuh>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <cstddef>
#include <cstdint>
#include <type_traits>
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<kThreadsPerGroup>(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<kThreadsPerGroup>` used directly in the kernel
// (gfx942-safe; no separate GroupReduceMax helper needed).
template <bool kScaleUE8M0>
using scale_packed_t_t = std::conditional_t<kScaleUE8M0, uint32_t, float>;
template <bool kUE8M0>
using scale_packed_t_t = std::conditional_t<kUE8M0, uint32_t, float>;
template <bool kUE8M0>
using scale_element_t_t = std::conditional_t<kUE8M0, uint8_t, float>;
template <bool kScaleUE8M0>
using scale_element_t_t = std::conditional_t<kScaleUE8M0, uint8_t, float>;
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 <typename T, typename DST_DTYPE, bool kIsColumnMajor, bool kScaleUE8M0>
__global__ void per_token_group_quant_8bit_kernel(
const T* __restrict__ input,
DST_DTYPE* __restrict__ output_q,
scale_packed_t_t<kScaleUE8M0>* __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 <typename T, typename DST, int64_t kGroupSize, int kThreadsPerGroup, bool kColMajor, bool kUE8M0, bool kUsePDL>
__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<int>(threadIdx.x / kThreadsPerGroup);
const int lane_id = threadIdx.x % kThreadsPerGroup;
using InVec = AlignedVector<T, kVec>;
using scale_packed_t = scale_packed_t_t<kUE8M0>;
using scale_element_t = scale_element_t_t<kUE8M0>;
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<int64_t>(blockIdx.x) * params.groups_per_block + local_group;
float local_absmax = eps;
using scale_packed_t = scale_packed_t_t<kScaleUE8M0>;
using scale_element_t = scale_element_t_t<kScaleUE8M0>;
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<DST_DTYPE*>(output_q) + block_group_offset;
scale_element_t* scale_output = nullptr;
if constexpr (kIsColumnMajor) {
constexpr int kElemsPerPack = static_cast<int>(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<scale_element_t*>(output_s) +
(col_idx * scale_stride * kElemsPerPack + row_idx * kElemsPerPack + pack_idx);
} else {
static_assert(!kScaleUE8M0);
scale_output = output_s + global_group_id;
PDLWaitPrimary<kUsePDL>();
if (global_group >= params.num_groups) {
PDLTriggerSecondary<kUsePDL>();
return;
}
constexpr uint32_t kVecSize = 16 / sizeof(T);
using vec_t = AlignedVector<T, kVecSize>;
const auto gmem_in = tile::Memory<vec_t>::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<const T*>(params.input) + global_group * kGroupSize;
DST* gout = static_cast<DST*>(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<float>(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<float>(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<kThreadsPerGroup>(local_absmax);
}
scale_element_t y_s_quant;
if constexpr (kScaleUE8M0) {
y_s_quant = static_cast<uint8_t>(((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<uint8_t>(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<float>(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<DST, fp8_e4m3_t>) {
AlignedVector<fp8x2_e4m3_t, kVec / 2> 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<DST, kVec> 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<DST>(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<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
const int row = static_cast<int>(global_group / params.num_groups_per_row); // token
const int col_u = static_cast<int>(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<scale_element_t*>(params.output_s) +
(static_cast<int64_t>(col) * params.scale_stride * kPack + static_cast<int64_t>(row) * kPack + pack);
} else {
static_assert(!kUE8M0, "non-column-major UE8M0 is unsupported");
scale_out = static_cast<scale_element_t*>(params.output_s) + global_group;
}
*scale_out = scale_store;
}
PDLTriggerSecondary<kUsePDL>();
}
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 <int64_t kGroupSize, typename T>
constexpr int threads_per_group() {
constexpr int kVec = 16 / sizeof(T); // 8 for bf16/fp16
int tpg = static_cast<int>(kGroupSize) / (2 * kVec); // ~16 elems/thread (2 vecs)
if (tpg > 8) tpg = 8;
if (tpg < 1) tpg = 1;
return tpg;
}
template <typename DType, typename OutType>
constexpr int kBlockThreads = 256;
template <int64_t kGroupSize, typename T>
inline int pick_groups_per_block() {
int gpb = kBlockThreads / threads_per_group<kGroupSize, T>();
if (gpb < 1) gpb = 1;
return gpb;
}
template <typename T, typename DST, int64_t kGroupSize, bool kColMajor, bool kUE8M0, bool kUsePDL>
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<kGroupSize, T>();
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<T, DST, kGroupSize, kThreadsPerGroup, kColMajor, kUE8M0, kUsePDL>;
LaunchKernel(static_cast<uint32_t>(num_blocks), static_cast<uint32_t>(num_threads), device)
.enable_pdl(kUsePDL)(kernel, params);
}
template <typename DType, typename OutType, int64_t kGroupSize, bool kUsePDL>
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<DType>().with_device(device).verify(input);
TensorMatcher({M, K}).with_dtype<OutType>().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<int>(output_s.stride(1));
const float feps = static_cast<float>(eps);
const float fmin8 = static_cast<float>(min_8bit);
const float fmax8 = static_cast<float>(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<int>(num_groups_per_row);
base.scale_stride = scale_stride;
base.eps = static_cast<float>(eps);
base.min_8bit = static_cast<float>(min_8bit);
base.max_8bit = static_cast<float>(max_8bit);
const auto dev = input.device();
const int gpb = pick_groups_per_block<kGroupSize, DType>();
// 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<DType, OutType, true, true>,
static_cast<const DType*>(input.data_ptr()),
static_cast<OutType*>(output_q.data_ptr()),
static_cast<uint32_t*>(output_s.data_ptr()),
static_cast<int>(group_size),
static_cast<int>(num_groups),
static_cast<int>(groups_per_block),
feps,
fmin8,
fmax8,
static_cast<int>(num_groups_per_row),
scale_stride);
launch_quant<DType, OutType, kGroupSize, true, true, kUsePDL>(base, num_groups, gpb, dev);
} else {
LaunchKernel(num_blocks, num_threads, input.device())(
per_token_group_quant_8bit_kernel<DType, OutType, true, false>,
static_cast<const DType*>(input.data_ptr()),
static_cast<OutType*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
static_cast<int>(group_size),
static_cast<int>(num_groups),
static_cast<int>(groups_per_block),
feps,
fmin8,
fmax8,
static_cast<int>(num_groups_per_row),
scale_stride);
launch_quant<DType, OutType, kGroupSize, true, false, kUsePDL>(base, num_groups, gpb, dev);
}
} else {
LaunchKernel(num_blocks, num_threads, input.device())(
per_token_group_quant_8bit_kernel<DType, OutType, false, false>,
static_cast<const DType*>(input.data_ptr()),
static_cast<OutType*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
static_cast<int>(group_size),
static_cast<int>(num_groups),
static_cast<int>(groups_per_block),
feps,
fmin8,
fmax8,
0,
0);
RuntimeCheck(!scale_ue8m0, "row-major UE8M0 unsupported");
launch_quant<DType, OutType, kGroupSize, false, false, kUsePDL>(base, num_groups, gpb, dev);
}
}
} // namespace
python/sglang/jit_kernel/moe_fused_gate.py
+2
View File
@@ -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
python/sglang/jit_kernel/per_token_group_quant_8bit.py
+17 -5
View File
@@ -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,
python/sglang/srt/arg_groups/overrides.py
+102 -15
View File
@@ -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":
python/sglang/srt/configs/__init__.py
+2
View File
@@ -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",
python/sglang/srt/configs/minimax_vl.py
+60
View File
@@ -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"
python/sglang/srt/configs/model_config.py
+2 -8
View File
@@ -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(
python/sglang/srt/entrypoints/openai/serving_chat.py
+7
View File
@@ -1791,6 +1791,13 @@ class OpenAIServingChat(OpenAIServingBase):
if not self.reasoning_parser:
return False
if self.reasoning_parser == "minimax-m3":
# M3 template prefills <mm:think> 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 <think> when reasoning_effort is
# "no_think" / "none" / unset; forcing reasoning would route all
python/sglang/srt/entrypoints/warmup.py
+34
View File
@@ -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__()
python/sglang/srt/environ.py
+9
View File
@@ -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)
python/sglang/srt/function_call/function_call_parser.py
+2
View File
@@ -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,
python/sglang/srt/function_call/minimax_m3.py
+542
View File
@@ -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>"
TOOL_CALL_END = MINIMAX_NS_TOKEN + "</tool_call>"
INVOKE_PREFIX = MINIMAX_NS_TOKEN + '<invoke name="'
INVOKE_SUFFIX = MINIMAX_NS_TOKEN + "</invoke>"
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"</{tag}>"
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("</"):
gt = chunk.find(">", 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("<item>")
return False
python/sglang/srt/layers/attention/attention_registry.py
+13
View File
@@ -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
python/sglang/srt/layers/attention/minimax_sparse_backend.py
+602
View File
@@ -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
)
python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
+26
View File
@@ -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,
python/sglang/srt/layers/layernorm.py
+31 -13
View File
@@ -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():
python/sglang/srt/layers/moe/ep_moe/kernels.py
+416 -104
View File
@@ -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,
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+145 -22
View File
@@ -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
python/sglang/srt/layers/moe/token_dispatcher/standard.py
+7 -2
View File
@@ -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)
python/sglang/srt/layers/moe/topk.py
+43 -7
View File
@@ -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:
python/sglang/srt/layers/quantization/fp8.py
+265 -31
View File
@@ -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()
python/sglang/srt/layers/quantization/fp8_kernel.py
+238 -70
View File
@@ -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,
python/sglang/srt/layers/quantization/fp8_utils.py
+85 -4
View File
@@ -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(
python/sglang/srt/layers/quantization/mxfp8_block_convert.py
+72
View File
@@ -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)
python/sglang/srt/layers/radix_attention.py
+106 -18
View File
@@ -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
python/sglang/srt/managers/multi_tokenizer_mixin.py
+4 -3
View File
@@ -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
python/sglang/srt/managers/scheduler.py
+6 -3
View File
@@ -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
)
python/sglang/srt/model_executor/forward_batch_info.py
+3 -1
View File
@@ -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
python/sglang/srt/models/minimax_m3.py
+1690
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/models/minimax_m3_vl.py
+372
View File
@@ -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]
python/sglang/srt/models/minimax_vl_common.py
+882
View File
@@ -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)
python/sglang/srt/multimodal/processors/minimax_m3_vl.py
+283
View File
@@ -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\|>|<\|image_pad\|>|\]\<\]image\[\>\["
),
video_token=self.VIDEO_TOKEN,
video_token_id=self.VIDEO_TOKEN_ID,
video_token_regex=re.compile(r"<video>|<\|video\|>|\]\<\]video\[\>\["),
).build(_processor)
async def process_mm_data_async(
self,
image_data: Optional[List],
audio_data: Optional[List],
input_text: str,
request_obj,
**kwargs,
) -> Dict:
base_output = await self.load_mm_data(
prompt=input_text,
image_data=image_data,
video_data=request_obj.video_data,
multimodal_tokens=self.mm_tokens,
)
video_metadata = None
if base_output.videos:
image_factor, max_size = self._video_resize_config()
videos_processed = [
await get_video_tensor(
video,
image_factor=image_factor,
max_size=max_size,
fps=self.video_fps,
frame_max_size=self.video_frame_max_size,
max_frames=self.video_max_frames,
)
for video in base_output.videos
]
base_output.videos, video_metadata = map(list, zip(*videos_processed))
mm_items, input_ids, ret = self.process_and_combine_mm_data(
base_output=base_output,
mm_tokens=self.mm_tokens,
video_metadata=video_metadata,
)
return MultimodalProcessorOutput(
input_ids=input_ids.tolist() if hasattr(input_ids, "tolist") else input_ids,
mm_items=mm_items,
im_start_id=self.IM_START_TOKEN_ID,
im_end_id=self.IM_END_TOKEN_ID,
im_token_id=self.IM_TOKEN_ID,
video_token_id=self.VIDEO_TOKEN_ID,
)
python/sglang/srt/parser/reasoning_parser.py
+61 -1
View File
@@ -502,6 +502,59 @@ class Nemotron3Detector(BaseReasoningFormatDetector):
return ret
class MiniMaxM3Detector(BaseReasoningFormatDetector):
"""MiniMax-M3 detector. Format: (<mm:think>)*(.*)</mm:think>.
In multi-turn chats M3 prefixes earlier non-thinking turns with a bare
``</mm:think>``, 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__(
"<mm:think>",
"</mm:think>",
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:
# ``</mm:think>`` 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 <mm:think> 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
python/sglang/srt/parser/template_detection.py
+6
View File
@@ -293,6 +293,10 @@ def _is_minimax(ctx):
return ctx.has_text("<minimax:tool_call>")
def _is_minimax_m3(ctx):
return ctx.has_text("<mm:think>") or ctx.has_text("]<]minimax[>[")
def _is_minicpm5(ctx):
if ctx.has_vocab("<function") and ctx.has_vocab("<param"):
return True
@@ -360,6 +364,7 @@ REASONING_PARSER_RULES = (
DetectionRule(name="hunyuan", value="hunyuan", predicate=_is_hunyuan),
DetectionRule(name="poolside_v1", value="poolside_v1", predicate=_is_poolside_v1),
DetectionRule(name="mimo", value="mimo", predicate=_is_mimo),
DetectionRule(name="minimax_m3", value="minimax-m3", predicate=_is_minimax_m3),
DetectionRule(name="minimax", value="minimax", predicate=_is_minimax),
DetectionRule(name="step3p5", value="step3p5", predicate=_is_step3p5),
DetectionRule(name="step3", value="step3", predicate=_is_step3),
@@ -385,6 +390,7 @@ TOOL_CALL_PARSER_RULES = (
DetectionRule(name="gemma4", value="gemma4", predicate=_is_gemma4),
DetectionRule(name="gpt_oss", value="gpt-oss", predicate=_is_gpt_oss),
DetectionRule(name="kimi_k2", value="kimi_k2", predicate=_is_kimi_k2),
DetectionRule(name="minimax_m3", value="minimax-m3", predicate=_is_minimax_m3),
DetectionRule(name="minimax", value="minimax-m2", predicate=_is_minimax),
DetectionRule(name="interns1", value="interns1", predicate=_is_interns1),
DetectionRule(name="mistral", value="mistral", predicate=_is_mistral),
python/sglang/srt/utils/common.py
+15
View File
@@ -4012,6 +4012,21 @@ def parse_module_path(module_path, function_name, create_dummy):
return final_module, None
@lru_cache(maxsize=1)
def mxfp8_block_convert_required():
"""Whether MXFP8 weights must be converted to block-fp8 [128,128] at load.
gfx942 (CDNA3) has no hardware MX-scaled matmul: ``tl.dot_scaled`` fails to
lower and the gfx950 ``mfma_scale`` intrinsics are unavailable. So MXFP8
checkpoints there are converted to block-fp8 [128,128] at load and run
through the native block-fp8 kernels. gfx95 keeps its native MX path (this
returns False there).
"""
if not torch.version.hip:
return False
return is_gfx942_supported() and not is_gfx95_supported()
# LoRA-related constants and utilities
SUPPORTED_LORA_TARGET_MODULES = [
"q_proj",
python/sglang/srt/utils/hf_transformers/common.py
+2
View File
@@ -43,6 +43,7 @@ from sglang.srt.configs import (
LongcatFlashConfig,
MiniCPMV4_6Config,
MiniCPMV4_6VisionConfig,
MiniMaxM3VLConfig,
MultiModalityConfig,
NemotronH_Nano_Omni_Reasoning_V3_Config,
NemotronH_Nano_VL_V2_Config,
@@ -112,6 +113,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
Step3p7Config,
MiniCPMV4_6Config,
MiniCPMV4_6VisionConfig,
MiniMaxM3VLConfig,
]
}