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

support qwen 3.8 flash next (#37500)

Browse Source 
Co-authored-by: ch-wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: ispobock <26454835+ispobock@users.noreply.github.com>
Co-authored-by: JustinTong0323 <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: samuellees <26428561+samuellees@users.noreply.github.com>
Co-authored-by: YAMY1234 <74099316+YAMY1234@users.noreply.github.com>
Co-authored-by: yhyang201 <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: yizhang2077 <25844240+yizhang2077@users.noreply.github.com>
Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: Shinto C V <cshintov@gmail.com>
Co-authored-by: Julian Huang <huangzhilin.hzl@antgroup.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>
This commit is contained in:
Qiaolin Yu
2026-09-08 13:56:21 -07:00
committed by GitHub
co-authored by ch-wan ispobock JustinTong0323 samuellees YAMY1234 yhyang201 yizhang2077 zijiexia Shinto C V Julian Huang Xiaoyu Zhang yhyang201
parent afe90a8bc9
commit 52fecfdf09
91 changed files with 16418 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/fused_moe_triton/common_utils.py
+7
View File
@@ -91,6 +91,13 @@ def get_model_config(
E = config.num_experts // ep_size
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
elif architecture in [
"Qwen4ExpForCausalLM",
"Qwen4ExpForConditionalGeneration",
]:
E = config.num_experts // ep_size
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
elif architecture in [
"DeepseekV2ForCausalLM",
"DeepseekV3ForCausalLM",
python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh
+486
View File
@@ -0,0 +1,486 @@
// Fused QSA (Qwen4-Exp sparse attention) indexer-prep kernels.
// Outputs are bit-identical to the eager bf16/fp16 aten chain, mirrored step by step.
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <tvm/ffi/container/tensor.h>
#include <cstdint>
namespace sglang {
/// \brief Round a float to the storage dtype and back (one eager aten step).
/// The inline cvt keeps nvcc from folding the round-trip away: every eager
/// bf16/fp16 aten op rounds its result, and this kernel must round likewise.
template <typename T>
SGL_DEVICE float eager_round(float x) {
return static_cast<float>(DTypeTrait<T>::from(x));
}
template <>
SGL_DEVICE float eager_round<bf16_t>(float x) {
#ifndef USE_ROCM
uint16_t u;
asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(u) : "f"(x));
return __bfloat162float(__ushort_as_bfloat16(u));
#else
return static_cast<float>(DTypeTrait<bf16_t>::from(x));
#endif
}
template <>
SGL_DEVICE float eager_round<fp16_t>(float x) {
#ifndef USE_ROCM
uint16_t u;
asm("cvt.rn.f16.f32 %0, %1;" : "=h"(u) : "f"(x));
return __ushort_as_half(u);
#else
return static_cast<float>(DTypeTrait<fp16_t>::from(x));
#endif
}
/**
* \brief Apply (M)RoPE to a normed row staged in shared memory.
*
* Reproduces the eager `get_cos_sin_with_position` + `apply_rotary_emb` chain:
* `axis_map[i]` selects which position axis feeds pair index i (plain RoPE
* uses all-zero maps; Qwen interleaved/sectioned MRoPE maps are built on the
* host). The cos/sin cache row is [cos(half), sin(half)] of width rotary_dim.
*
* \tparam T Storage element type: bf16_t | fp16_t.
* \tparam kHeadDim Compile-time head dimension (multiple of 32).
* \tparam kIsNeox true -> NeoX pairing (d, d+half); false -> GPT-J (2i, 2i+1).
* \param smem_row Normed row [kHeadDim], one warp cooperates.
* \param out_row Destination row [kHeadDim].
* \param cos_sin_cache [num_positions, rotary_dim] fp32 cache.
* \param axis_map [rotary_dim/2] int32 position-axis selector per pair.
* \param pos Resolved per-axis positions for this token (>= 3 entries).
* \param rotary_dim Rotated prefix length; tail dims pass through.
*/
template <typename T, int kHeadDim, bool kIsNeox>
SGL_DEVICE void qsa_mrope_apply(
const T* __restrict__ smem_row,
T* __restrict__ out_row,
const float* __restrict__ cos_sin_cache,
const int32_t* __restrict__ axis_map,
const int64_t* pos,
const int32_t rotary_dim) {
using namespace device;
constexpr int kPerLane = kHeadDim / kWarpThreads;
using vec_t = AlignedVector<T, kPerLane>;
const uint32_t lane = threadIdx.x % kWarpThreads;
const int32_t half = rotary_dim / 2;
vec_t ov;
#pragma unroll
for (int i = 0; i < kPerLane; ++i) {
const int32_t d = static_cast<int32_t>(lane * kPerLane) + i;
T o;
if constexpr (kIsNeox) {
if (d < half) {
const int32_t p = d + half;
const float* row = cos_sin_cache + pos[axis_map[d]] * rotary_dim;
const float c = eager_round<T>(row[d]);
const float s = eager_round<T>(row[half + d]);
const float nd = static_cast<float>(smem_row[d]);
const float np = static_cast<float>(smem_row[p]);
o = DTypeTrait<T>::from(eager_round<T>(nd * c) - eager_round<T>(np * s));
} else if (d < rotary_dim) {
const int32_t p = d - half;
const float* row = cos_sin_cache + pos[axis_map[p]] * rotary_dim;
const float c = eager_round<T>(row[p]);
const float s = eager_round<T>(row[half + p]);
const float nd = static_cast<float>(smem_row[d]);
const float np = static_cast<float>(smem_row[p]);
o = DTypeTrait<T>::from(eager_round<T>(nd * c) + eager_round<T>(np * s));
} else {
o = smem_row[d];
}
} else {
if (d < rotary_dim) {
const int32_t p = d / 2;
const float* row = cos_sin_cache + pos[axis_map[p]] * rotary_dim;
const float c = eager_round<T>(row[p]);
const float s = eager_round<T>(row[half + p]);
const int32_t q = (d % 2 == 0) ? d + 1 : d - 1;
const float nd = static_cast<float>(smem_row[d]);
const float nq = static_cast<float>(smem_row[q]);
const float t1 = eager_round<T>(nd * c);
const float t2 = eager_round<T>(nq * s);
o = DTypeTrait<T>::from((d % 2 == 0) ? t1 - t2 : t1 + t2);
} else {
o = smem_row[d];
}
}
ov[i] = o;
}
ov.store(out_row, lane); // offset is in vector units
}
/**
* \brief Gemma RMSNorm of one row into shared memory (warp-cooperative).
*
* out = x * rsqrt(mean(x^2) + eps) * (1 + w), fp32 math with one rounding.
* The sum-of-squares uses flashinfer RMSNormKernel's exact partial layout
* (vec_size = 8 for 2-byte dtypes, thread t sums elements [8t, 8t+8), inactive
* lanes contribute 0, xor butterfly over the warp) so results stay bit-equal
* to the eager sgl_kernel gemma_rmsnorm this replaces.
*/
template <typename T, int kHeadDim>
SGL_DEVICE void qsa_gemma_norm_row(
const T* __restrict__ x_row, const T* __restrict__ weight, const float eps, T* __restrict__ smem_row) {
using namespace device;
constexpr int kPerLane = kHeadDim / kWarpThreads;
using vec_t = AlignedVector<T, kPerLane>;
const uint32_t lane = threadIdx.x % kWarpThreads;
vec_t xv, wv;
xv.load(x_row, lane); // offset is in vector units
wv.load(weight, lane);
float xf[kPerLane];
#pragma unroll
for (int i = 0; i < kPerLane; ++i) {
xf[i] = static_cast<float>(xv[i]);
}
static_assert(kHeadDim % 8 == 0);
constexpr uint32_t kNormThreads = kHeadDim / 8;
float ss = 0.0f;
if (lane < kNormThreads) {
AlignedVector<T, 4> va, vb;
va.load(x_row, lane * 2); // elements [8*lane, 8*lane+4)
vb.load(x_row, lane * 2 + 1); // elements [8*lane+4, 8*lane+8)
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float f = static_cast<float>(va[i]);
ss += f * f;
}
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float f = static_cast<float>(vb[i]);
ss += f * f;
}
}
ss = warp::reduce_sum(ss);
const float nf = math::rsqrt(ss / kHeadDim + eps);
#pragma unroll
for (int i = 0; i < kPerLane; ++i) {
const float wf = static_cast<float>(wv[i]);
smem_row[lane * kPerLane + i] = DTypeTrait<T>::from(xf[i] * nf * (1.0f + wf));
}
__syncwarp();
}
struct QsaIndexQPrepParams {
const void* qk; // [tokens, (num_q_heads + 1) * kHeadDim]
void* q_out; // [tokens, q_heads_padded, kHeadDim]
const void* weight; // [kHeadDim]
const float* cos_sin_cache; // [positions_capacity, rotary_dim]
const int32_t* axis_map; // [rotary_dim / 2]
const int64_t* positions; // [num_axes, tokens] (row stride may exceed tokens)
const int64_t* cache_loc; // [tokens]
void* key_state_buffer; // [slots, kHeadDim]
int64_t* rope_position_buffer; // [slots, 3]
int64_t positions_stride;
int32_t num_axes;
int32_t num_q_heads;
int32_t q_heads_padded;
int32_t rotary_dim;
float eps;
};
/**
* \brief Per-token fused index-Q prep: gemma norm + MRoPE for every query
* head, zero-fill of padded heads, raw token-K store and RoPE-position store.
* One CTA (4 warps) per token; one warp per query head.
*/
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
__global__ __launch_bounds__(128) void qsa_index_q_prep_kernel(const QsaIndexQPrepParams __grid_constant__ params) {
using namespace device;
constexpr int kPerLane = kHeadDim / kWarpThreads;
using vec_t = AlignedVector<T, kPerLane>;
const uint32_t token = blockIdx.x;
const uint32_t warp = threadIdx.x / kWarpThreads;
const uint32_t lane = threadIdx.x % kWarpThreads;
__shared__ T smem_rows[4][kHeadDim];
device::PDLWaitPrimary<kUsePDL>();
const int64_t qk_row = static_cast<int64_t>(token) * (params.num_q_heads + 1) * kHeadDim;
const int64_t loc = params.cache_loc[token];
int64_t pos[3];
#pragma unroll
for (int a = 0; a < 3; ++a) {
const int64_t ax = a < params.num_axes ? a : 0;
pos[a] = params.positions[ax * params.positions_stride + token];
}
for (int32_t h = static_cast<int32_t>(warp); h < params.q_heads_padded; h += 4) {
T* out_row = static_cast<T*>(params.q_out) + (static_cast<int64_t>(token) * params.q_heads_padded + h) * kHeadDim;
if (h < params.num_q_heads) {
const T* x_row = static_cast<const T*>(params.qk) + qk_row + h * kHeadDim;
qsa_gemma_norm_row<T, kHeadDim>(x_row, static_cast<const T*>(params.weight), params.eps, smem_rows[warp]);
qsa_mrope_apply<T, kHeadDim, kIsNeox>(
smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim);
} else {
vec_t zv;
zv.fill(DTypeTrait<T>::from(0.0f));
zv.store(out_row, lane); // offset is in vector units
}
}
// Raw token-K and RoPE coordinates are stored for every token, whether or
// not the token completes a compression group.
if (warp == 0) {
vec_t kv;
kv.load(
static_cast<const T*>(params.qk) + qk_row + params.num_q_heads * kHeadDim,
lane); // offset is in vector units
kv.store(static_cast<T*>(params.key_state_buffer) + loc * kHeadDim, lane);
}
if (warp == 1 && lane < 3) {
params.rope_position_buffer[loc * 3 + lane] = pos[lane];
}
device::PDLTriggerSecondary<kUsePDL>();
}
struct QsaIndexKCompressParams {
const void* key_state_buffer; // [slots, kHeadDim]
const int32_t* group_locs; // [groups, compress_ratio]
const int64_t* rope_position_buffer; // [slots, 3]
const float* cos_sin_cache; // [positions_capacity, rotary_dim]
const int32_t* axis_map; // [rotary_dim / 2]
const void* weight; // [kHeadDim]
const int32_t* write_locs; // [groups]
void* compressed_k_buffer; // [compressed_slots, kHeadDim]
int32_t compress_ratio;
int32_t rotary_dim;
int32_t num_groups;
float eps;
};
/**
* \brief Per-group compressed-K prep: fp32 mean over the group, gemma norm,
* MRoPE at the group-start position, store into the compressed cache.
* One warp per group.
*/
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
__global__
__launch_bounds__(128) void qsa_index_k_compress_kernel(const QsaIndexKCompressParams __grid_constant__ params) {
using namespace device;
constexpr int kPerLane = kHeadDim / kWarpThreads;
using vec_t = AlignedVector<T, kPerLane>;
const uint32_t warp = threadIdx.x / kWarpThreads;
const uint32_t lane = threadIdx.x % kWarpThreads;
const uint32_t group = blockIdx.x * 4 + warp;
if (group >= static_cast<uint32_t>(params.num_groups)) {
return;
}
__shared__ T smem_rows[4][kHeadDim];
device::PDLWaitPrimary<kUsePDL>();
const int32_t* locs = params.group_locs + group * params.compress_ratio;
const int32_t loc0 = locs[0];
// fp32 mean over the group, rounded to the storage dtype exactly like
// average_pool_qsa_keys (float().mean(dim=1).to(dtype)).
float mf[kPerLane];
{
float acc[kPerLane];
for (int32_t r = 0; r < params.compress_ratio; ++r) {
vec_t v;
v.load(
static_cast<const T*>(params.key_state_buffer) + static_cast<int64_t>(locs[r]) * kHeadDim,
lane); // offset is in vector units
#pragma unroll
for (int i = 0; i < kPerLane; ++i) {
const float f = static_cast<float>(v[i]);
acc[i] = (r == 0) ? f : acc[i] + f;
}
}
const float inv_ratio = 1.0f / static_cast<float>(params.compress_ratio);
#pragma unroll
for (int i = 0; i < kPerLane; ++i) {
const T m = DTypeTrait<T>::from(acc[i] * inv_ratio);
mf[i] = static_cast<float>(m);
smem_rows[warp][lane * kPerLane + i] = m;
}
__syncwarp();
// Sum of squares in flashinfer RMSNormKernel's exact partial layout
// (thread t sums elements [8t, 8t+8), inactive lanes contribute 0), so
// the result stays bit-equal to the eager k_layernorm this replaces.
static_assert(kHeadDim % 8 == 0);
float ss = 0.0f;
if (lane < kHeadDim / 8) {
#pragma unroll
for (int j = 0; j < 8; ++j) {
const float f = static_cast<float>(smem_rows[warp][lane * 8 + j]);
ss += f * f;
}
}
ss = warp::reduce_sum(ss);
const float nf = math::rsqrt(ss / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kPerLane; ++i) {
const float wf = static_cast<float>(static_cast<const T*>(params.weight)[lane * kPerLane + i]);
smem_rows[warp][lane * kPerLane + i] = DTypeTrait<T>::from(mf[i] * nf * (1.0f + wf));
}
__syncwarp();
}
int64_t pos[3];
#pragma unroll
for (int a = 0; a < 3; ++a) {
pos[a] = params.rope_position_buffer[static_cast<int64_t>(loc0) * 3 + a];
}
T* out_row = static_cast<T*>(params.compressed_k_buffer) + static_cast<int64_t>(params.write_locs[group]) * kHeadDim;
qsa_mrope_apply<T, kHeadDim, kIsNeox>(
smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim);
device::PDLTriggerSecondary<kUsePDL>();
}
/**
* \brief Validate inputs and launch `qsa_index_q_prep_kernel` (one CTA per token).
*
* \tparam T Element type: bf16_t | fp16_t.
* \tparam kHeadDim Index head dimension: 64 | 128 | 256.
* \tparam kIsNeox RoPE pairing style.
* \tparam kUsePDL Whether to launch with PDL enabled.
*/
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
void qsa_index_q_prep(
tvm::ffi::TensorView qk,
tvm::ffi::TensorView q_out,
tvm::ffi::TensorView weight,
tvm::ffi::TensorView cos_sin_cache,
tvm::ffi::TensorView axis_map,
tvm::ffi::TensorView positions,
int64_t num_axes,
tvm::ffi::TensorView cache_loc,
tvm::ffi::TensorView key_state_buffer,
tvm::ffi::TensorView rope_position_buffer,
int64_t num_q_heads,
int64_t rotary_dim,
float eps) {
using namespace host;
auto tokens = SymbolicSize{"tokens"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
constexpr int64_t D = kHeadDim;
TensorMatcher({tokens, (num_q_heads + 1) * D}).with_dtype<T>().with_device(device).verify(qk);
auto heads_padded = SymbolicSize{"heads_padded"};
TensorMatcher({tokens, heads_padded, D}).with_dtype<T>().with_device(device).verify(q_out);
TensorMatcher({D}).with_dtype<T>().with_device(device).verify(weight);
auto cache_rows = SymbolicSize{"cos_sin_cache_rows"};
TensorMatcher({cache_rows, rotary_dim}).with_dtype<fp32_t>().with_device(device).verify(cos_sin_cache);
TensorMatcher({rotary_dim / 2}).with_dtype<int32_t>().with_device(device).verify(axis_map);
TensorMatcher({num_axes, tokens}).with_dtype<int64_t>().with_device(device).with_strides({-1, 1}).verify(positions);
TensorMatcher({tokens}).with_dtype<int64_t>().with_device(device).verify(cache_loc);
auto slots = SymbolicSize{"state_slots"};
TensorMatcher({slots, D}).with_dtype<T>().with_device(device).verify(key_state_buffer);
TensorMatcher({slots, 3}).with_dtype<int64_t>().with_device(device).verify(rope_position_buffer);
const int64_t num_tokens = tokens.unwrap();
const int64_t q_heads_padded = heads_padded.unwrap();
CHECK_HOST(num_tokens > 0) << "qsa_index_q_prep: no tokens";
CHECK_HOST(num_axes == 1 || num_axes == 3) << "qsa_index_q_prep: positions must have 1 or 3 axes, got " << num_axes;
CHECK_HOST(q_heads_padded >= num_q_heads)
<< "qsa_index_q_prep: padded heads " << q_heads_padded << " < num_q_heads " << num_q_heads;
CHECK_HOST(rotary_dim > 0 && rotary_dim % 2 == 0 && rotary_dim <= D)
<< "qsa_index_q_prep: invalid rotary_dim " << rotary_dim;
const auto params = QsaIndexQPrepParams{
.qk = qk.data_ptr(),
.q_out = q_out.data_ptr(),
.weight = weight.data_ptr(),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.axis_map = static_cast<const int32_t*>(axis_map.data_ptr()),
.positions = static_cast<const int64_t*>(positions.data_ptr()),
.cache_loc = static_cast<const int64_t*>(cache_loc.data_ptr()),
.key_state_buffer = key_state_buffer.data_ptr(),
.rope_position_buffer = static_cast<int64_t*>(rope_position_buffer.data_ptr()),
.positions_stride = positions.stride(0),
.num_axes = static_cast<int32_t>(num_axes),
.num_q_heads = static_cast<int32_t>(num_q_heads),
.q_heads_padded = static_cast<int32_t>(q_heads_padded),
.rotary_dim = static_cast<int32_t>(rotary_dim),
.eps = eps,
};
LaunchKernel(static_cast<uint32_t>(num_tokens), 128, device.unwrap())
.enable_pdl(kUsePDL)(qsa_index_q_prep_kernel<T, kHeadDim, kIsNeox, kUsePDL>, params);
}
/**
* \brief Validate inputs and launch `qsa_index_k_compress_kernel` (one warp per group).
*/
template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
void qsa_index_k_compress(
tvm::ffi::TensorView key_state_buffer,
tvm::ffi::TensorView group_locs,
tvm::ffi::TensorView rope_position_buffer,
tvm::ffi::TensorView cos_sin_cache,
tvm::ffi::TensorView axis_map,
tvm::ffi::TensorView weight,
tvm::ffi::TensorView write_locs,
tvm::ffi::TensorView compressed_k_buffer,
int64_t compress_ratio,
int64_t rotary_dim,
float eps) {
using namespace host;
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
constexpr int64_t D = kHeadDim;
auto slots = SymbolicSize{"state_slots"};
TensorMatcher({slots, D}).with_dtype<T>().with_device(device).verify(key_state_buffer);
auto groups = SymbolicSize{"groups"};
TensorMatcher({groups, compress_ratio}).with_dtype<int32_t>().with_device(device).verify(group_locs);
TensorMatcher({slots, 3}).with_dtype<int64_t>().with_device(device).verify(rope_position_buffer);
auto cache_rows = SymbolicSize{"cos_sin_cache_rows"};
TensorMatcher({cache_rows, rotary_dim}).with_dtype<fp32_t>().with_device(device).verify(cos_sin_cache);
TensorMatcher({rotary_dim / 2}).with_dtype<int32_t>().with_device(device).verify(axis_map);
TensorMatcher({D}).with_dtype<T>().with_device(device).verify(weight);
TensorMatcher({groups}).with_dtype<int32_t>().with_device(device).verify(write_locs);
auto compressed_slots = SymbolicSize{"compressed_slots"};
TensorMatcher({compressed_slots, D}).with_dtype<T>().with_device(device).verify(compressed_k_buffer);
const int64_t num_groups = groups.unwrap();
CHECK_HOST(num_groups > 0) << "qsa_index_k_compress: no groups";
CHECK_HOST(compress_ratio > 0 && compress_ratio <= 16)
<< "qsa_index_k_compress: invalid compress_ratio " << compress_ratio;
CHECK_HOST(rotary_dim > 0 && rotary_dim % 2 == 0 && rotary_dim <= D)
<< "qsa_index_k_compress: invalid rotary_dim " << rotary_dim;
const auto params = QsaIndexKCompressParams{
.key_state_buffer = key_state_buffer.data_ptr(),
.group_locs = static_cast<const int32_t*>(group_locs.data_ptr()),
.rope_position_buffer = static_cast<const int64_t*>(rope_position_buffer.data_ptr()),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.axis_map = static_cast<const int32_t*>(axis_map.data_ptr()),
.weight = weight.data_ptr(),
.write_locs = static_cast<const int32_t*>(write_locs.data_ptr()),
.compressed_k_buffer = compressed_k_buffer.data_ptr(),
.compress_ratio = static_cast<int32_t>(compress_ratio),
.rotary_dim = static_cast<int32_t>(rotary_dim),
.num_groups = static_cast<int32_t>(num_groups),
.eps = eps,
};
LaunchKernel(static_cast<uint32_t>(div_ceil(num_groups, 4)), 128, device.unwrap())
.enable_pdl(kUsePDL)(qsa_index_k_compress_kernel<T, kHeadDim, kIsNeox, kUsePDL>, params);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/fast_topk.cuh
+291
View File
@@ -0,0 +1,291 @@
// Radix-select fast top-k with the AOT fast_topk_v2 semantics: for each row b,
// select the kTopK largest scores in [row_starts[b], row_starts[b] + lengths[b])
// and write their indices relative to row_starts[b]; order within a row is unspecified.
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <tvm/ffi/container/tensor.h>
namespace sglang {
namespace fast_topk_detail {
constexpr uint32_t kThreadsPerBlock = 1024;
// Each radix pass needs at most ~kTopK candidates in the threshold bin, so
// 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
constexpr size_t kSmemBytes = 8 * 1024 * sizeof(uint32_t); // 32KB
struct FastTopKParams {
const float* __restrict__ input; // [B, input_stride]
const int32_t* __restrict__ row_starts; // [B]
int32_t* __restrict__ indices; // [B, kTopK]
const int32_t* __restrict__ lengths; // [B]
int64_t input_stride;
};
SGL_DEVICE auto convert_to_uint8(float x) -> uint8_t {
const __half h = __float2half_rn(x);
const uint16_t bits = __half_as_ushort(h);
const uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
return static_cast<uint8_t>(key >> 8);
}
SGL_DEVICE auto convert_to_uint32(float x) -> uint32_t {
const uint32_t bits = __float_as_uint(x);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
// When length <= kTopK, write the indices directly.
template <int kTopK>
SGL_DEVICE void naive_topk(const float* __restrict__ score, int32_t* __restrict__ indice, int32_t length) {
const auto tid = threadIdx.x;
for (int i = tid; i < kTopK; i += kThreadsPerBlock) {
indice[i] = (i < length) ? i : -1;
}
}
// Radix-select top-k. Assumes length > kTopK (checked by the caller).
template <int kTopK>
SGL_DEVICE void radix_select_topk(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
int topk = kTopK;
constexpr auto BLOCK_SIZE = kThreadsPerBlock;
constexpr auto RADIX = 256;
constexpr auto SMEM_INPUT_SIZE = kSmemBytes / (2 * sizeof(int));
alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
alignas(128) __shared__ int s_counter;
alignas(128) __shared__ int s_threshold_bin_id;
alignas(128) __shared__ int s_num_input[2];
auto& s_histogram = s_histogram_buf[0];
// allocate for two rounds
extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
const int tx = threadIdx.x;
// stage 1: 8bit coarse histogram
if (tx < RADIX + 1) s_histogram[tx] = 0;
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = convert_to_uint8(input[idx + row_start]);
::atomicAdd(&s_histogram[bin], 1);
}
__syncthreads();
const auto run_cumsum = [&] {
#pragma unroll 8
for (int i = 0; i < 8; ++i) {
static_assert(1 << 8 == RADIX);
if (tx < RADIX) {
const auto j = 1 << i;
const auto k = i & 1;
auto value = s_histogram_buf[k][tx];
if (tx < RADIX - j) {
value += s_histogram_buf[k][tx + j];
}
s_histogram_buf[k ^ 1][tx] = value;
}
__syncthreads();
}
};
run_cumsum();
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
s_threshold_bin_id = tx;
s_num_input[0] = 0;
s_counter = 0;
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
topk -= s_histogram[threshold_bin + 1];
if (topk == 0) {
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
}
}
__syncthreads();
return;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto raw_input = input[idx + row_start];
const auto bin = static_cast<int>(convert_to_uint8(raw_input));
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
} else if (bin == threshold_bin) {
const auto pos = ::atomicAdd(&s_num_input[0], 1);
// fuse the histogram computation here
if (pos < int(SMEM_INPUT_SIZE)) {
s_input_idx[0][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> 24) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
__syncthreads();
}
// stage 2: refine with 8bit radix passes
#pragma unroll 4
for (int round = 0; round < 4; ++round) {
__shared__ int s_last_remain;
const auto r_idx = round % 2;
// clip here to prevent overflow
const auto _raw_num_input = s_num_input[r_idx];
const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
run_cumsum();
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
s_threshold_bin_id = tx;
s_num_input[r_idx ^ 1] = 0;
s_last_remain = topk - s_histogram[tx + 1];
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
topk -= s_histogram[threshold_bin + 1];
if (topk == 0) {
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(input[idx + row_start]) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
}
}
__syncthreads();
break;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__syncthreads();
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto raw_input = input[idx + row_start];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
} else if (bin == threshold_bin) {
if (round == 3) {
const auto pos = ::atomicAdd(&s_last_remain, -1);
if (pos > 0) {
index[kTopK - pos] = idx;
}
} else {
const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
if (pos < int(SMEM_INPUT_SIZE)) {
// fuse the histogram computation here
s_input_idx[r_idx ^ 1][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
}
__syncthreads();
}
}
}
template <int kTopK, bool kUsePDL>
__global__ __launch_bounds__(fast_topk_detail::kThreadsPerBlock) void fast_topk_kernel(
const fast_topk_detail::FastTopKParams __grid_constant__ params) {
using namespace fast_topk_detail;
device::PDLWaitPrimary<kUsePDL>();
const auto bid = static_cast<uint64_t>(blockIdx.x);
const auto row_start = params.row_starts == nullptr ? 0 : params.row_starts[bid];
const auto length = params.lengths[bid];
const auto indice = params.indices + bid * kTopK;
const auto score = params.input + bid * params.input_stride;
if (length <= kTopK) {
naive_topk<kTopK>(score, indice, length);
} else {
radix_select_topk<kTopK>(score, indice, row_start, length);
}
device::PDLTriggerSecondary<kUsePDL>();
}
} // namespace fast_topk_detail
/**
* \brief Per-row top-k selection over ragged rows of a fp32 score matrix.
*
* Row b selects the kTopK largest values in
* score[b, row_starts[b] : row_starts[b] + lengths[b]) and writes their
* indices (relative to row_starts[b]) into indices[b]. Unfilled slots are
* -1 when lengths[b] < kTopK.
*/
template <int kTopK, bool kUsePDL>
struct FastTopKKernel {
static constexpr auto kernel = fast_topk_detail::fast_topk_kernel<kTopK, kUsePDL>;
static void
run(const tvm::ffi::TensorView score,
const tvm::ffi::TensorView row_starts,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView lengths) {
using namespace host;
auto B = SymbolicSize{"batch"};
auto L = SymbolicSize{"length"};
auto S = SymbolicSize{"input_stride"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({B, L}) // score
.with_strides({S, 1})
.with_dtype<fp32_t>()
.with_device(device)
.verify(score);
TensorMatcher({B}) // row_starts
.with_dtype<int32_t>()
.with_device(device)
.verify(row_starts);
TensorMatcher({B, kTopK}) // indices
.with_dtype<int32_t>()
.with_device(device)
.verify(indices);
TensorMatcher({B}) // lengths
.with_dtype<int32_t>()
.with_device(device)
.verify(lengths);
const auto params = fast_topk_detail::FastTopKParams{
.input = static_cast<const float*>(score.data_ptr()),
.row_starts = static_cast<const int32_t*>(row_starts.data_ptr()),
.indices = static_cast<int32_t*>(indices.data_ptr()),
.lengths = static_cast<const int32_t*>(lengths.data_ptr()),
.input_stride = S.unwrap(),
};
const auto num_rows = static_cast<uint32_t>(B.unwrap());
LaunchKernel(num_rows, fast_topk_detail::kThreadsPerBlock, device.unwrap(), fast_topk_detail::kSmemBytes)
.enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/grouped_gemma_rmsnorm.cuh
+178
View File
@@ -0,0 +1,178 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#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 <tvm/ffi/container/tensor.h>
namespace sglang {
struct GroupedGemmaRMSNormParams {
const void* input;
const void* __restrict__ weight;
void* output;
uint32_t num_groups;
float eps;
};
/**
* \brief Grouped Gemma-style RMSNorm: out = x * rsqrt(mean(x^2) + eps) * (1 + w).
*
* The last dim of the input is split into `num_groups` chunks of `kGroupSize`
* elements. Variance is computed per (token, group) chunk, so a [M, H] input
* with H = num_groups * kGroupSize behaves like M * num_groups independent
* RMSNorm rows whose weight rows are the matching kGroupSize slice of `weight`.
*
* One CTA handles one (token, group) chunk. Since chunks are contiguous in
* memory, block `bid` reads/writes elements [bid * kGroupSize, (bid + 1) *
* kGroupSize) and uses weight slice (bid % num_groups) * kGroupSize.
*
* \tparam kGroupSize Elements per group. Must be a multiple of 512.
* \tparam kUsePDL Whether to emit the PDL wait/trigger pair.
* \tparam Float Element type: bf16_t | fp16_t.
*/
template <int64_t kGroupSize, bool kUsePDL, typename Float>
__global__ __launch_bounds__(kGroupSize / 16) void grouped_gemma_rmsnorm_kernel(
const GroupedGemmaRMSNormParams __grid_constant__ params) {
using namespace device;
using Float2 = packed_t<Float>;
#if SGL_ARCH_BLACKWELL_OR_GREATER
// Blackwell: 32B vector, each thread loads/stores once
using Storage = AlignedVector<Float2, 8>;
constexpr uint32_t kNumLoads = 1;
#else
// Pre-Blackwell: 16B vector, each thread loads/stores twice
using Storage = AlignedVector<Float2, 4>;
constexpr uint32_t kNumLoads = 2;
#endif
constexpr uint32_t kVecLen = kNumLoads == 1 ? 8 : 4;
constexpr auto kNumThreads = kGroupSize / 16;
constexpr auto kNumWarps = kNumThreads / kWarpThreads;
const uint32_t bid = blockIdx.x;
const uint32_t group = bid % params.num_groups;
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
// Warp 0 writes smem[tx] for all 32 lanes in the cross-warp reduce below,
// so this must hold kWarpThreads entries, not kNumWarps.
__shared__ float smem[kWarpThreads];
PDLWaitPrimary<kUsePDL>();
const auto input_ptr = pointer::offset<Float>(params.input, static_cast<int64_t>(bid) * kGroupSize);
const auto output_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(bid) * kGroupSize);
const auto weight_ptr = pointer::offset<Float>(params.weight, static_cast<int64_t>(group) * kGroupSize);
Storage input_vec[kNumLoads];
Storage weight_vec[kNumLoads];
#pragma unroll
for (uint32_t j = 0; j < kNumLoads; ++j) {
input_vec[j] = gmem.load(input_ptr, j);
weight_vec[j] = gmem.load(weight_ptr, j);
}
float sum_of_squares = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < kNumLoads; ++j) {
#pragma unroll
for (uint32_t i = 0; i < kVecLen; ++i) {
const auto [x, y] = cast<fp32x2_t>(input_vec[j][i]);
sum_of_squares += x * x + y * y;
}
}
sum_of_squares = warp::reduce_sum(sum_of_squares);
float norm_factor;
if constexpr (kNumWarps == 1) {
norm_factor = math::rsqrt(sum_of_squares / kGroupSize + params.eps);
} else {
const auto warp_id = threadIdx.x / kWarpThreads;
smem[warp_id] = sum_of_squares;
__syncthreads();
if (warp_id == 0) {
const auto tx = threadIdx.x;
const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
sum_of_squares = warp::reduce_sum(local_sum);
smem[tx] = math::rsqrt(sum_of_squares / kGroupSize + params.eps);
}
__syncthreads();
norm_factor = smem[warp_id];
}
#pragma unroll
for (uint32_t j = 0; j < kNumLoads; ++j) {
Storage output_vec;
#pragma unroll
for (uint32_t i = 0; i < kVecLen; ++i) {
const auto [ix, iy] = cast<fp32x2_t>(input_vec[j][i]);
const auto [wx, wy] = cast<fp32x2_t>(weight_vec[j][i]);
output_vec[i] = cast<Float2>(fp32x2_t{ix * norm_factor * (1.0f + wx), iy * norm_factor * (1.0f + wy)});
}
gmem.store(output_ptr, output_vec, j);
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kGroupSize, bool kUsePDL, typename DType>
struct GroupedGemmaRMSNormKernel {
static_assert(sizeof(DType) == 2, "GroupedGemmaRMSNorm only supports 2-byte dtypes");
static_assert(kGroupSize % 512 == 0, "kGroupSize must be a multiple of 512");
static constexpr auto kernel = grouped_gemma_rmsnorm_kernel<kGroupSize, kUsePDL, DType>;
static constexpr auto kBlockSize = static_cast<uint32_t>(kGroupSize / 16);
/**
* \brief Validate tensors and launch one CTA per (token, group) chunk.
* \param input [M, H] contiguous, H % kGroupSize == 0
* \param weight [H]
* \param output [M, H] contiguous, same shape/dtype/device as input
* \param eps RMSNorm epsilon
*/
static void
run(const tvm::ffi::TensorView input,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView output,
float eps) {
using namespace host;
auto M = SymbolicSize{"num_tokens"};
auto H = SymbolicSize{"hidden_size"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({M, H}) // input
.with_dtype<DType>()
.with_device(device)
.verify(input);
TensorMatcher({H}) // weight
.with_dtype<DType>()
.with_device(device)
.verify(weight);
TensorMatcher({M, H}) // output
.with_dtype<DType>()
.with_device(device)
.verify(output);
const int64_t hidden_size = H.unwrap();
CHECK_HOST(hidden_size % kGroupSize == 0) << "grouped_gemma_rmsnorm: hidden_size (" << hidden_size
<< ") must be divisible by group_size (" << kGroupSize << ")";
const auto params = GroupedGemmaRMSNormParams{
.input = input.data_ptr(),
.weight = weight.data_ptr(),
.output = output.data_ptr(),
.num_groups = static_cast<uint32_t>(hidden_size / kGroupSize),
.eps = eps,
};
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
const uint32_t num_blocks = num_tokens * params.num_groups;
LaunchKernel(num_blocks, kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/hc_combine.cuh
+381
View File
@@ -0,0 +1,381 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#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 <tvm/ffi/container/tensor.h>
namespace sglang {
struct HcCombineParams {
const void* block_output; // [M, H]
const void* residual; // [M, HC * H]
const void* normed_residual; // [M, HC * H]
const void* inject_weight; // [HC, HC * H]
void* output; // [M, HC * H]
};
/**
* \brief Fused HyperConnection (gated residual) combine:
*
* a[m, c] = 2 * sigmoid(dot(normed_residual[m, :], inject_weight[c, :]) / kHcCount)
* output[m, c*H + i] = residual[m, c*H + i] + a[m, c] * block_output[m, i]
*
* One CTA handles one token row. Phase 1 computes the kHcCount gate values with a
* block reduction over the full HC*H row (fp32 accumulation). Phase 2 streams the
* HC*H output elements with vectorized 16B accesses; the block_output row is only
* H elements, so its re-read per branch stays in L2.
*
* \tparam kHcCount Number of hyper-connection branches (4 in production).
* \tparam kHiddenSize Per-branch hidden size H. HC*H must be a multiple of
* kNumThreads * kVecLen so the row maps exactly onto the CTA.
* \tparam kUsePDL Whether to emit the PDL wait/trigger pair.
* \tparam Float Element type: bf16_t | fp16_t.
*/
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
__global__ __launch_bounds__(256) void hc_combine_kernel(const HcCombineParams __grid_constant__ params) {
using namespace device;
using Float2 = packed_t<Float>;
using Storage = AlignedVector<Float2, 4>; // 8 elements, 16 bytes
constexpr uint32_t kVecLen = 8;
constexpr uint32_t kNumThreads = 256;
constexpr int64_t kRowSize = kHcCount * kHiddenSize;
constexpr uint32_t kVecsPerRow = kRowSize / kVecLen; // 1280 for 4x2560
constexpr uint32_t kVecsPerThread = kVecsPerRow / kNumThreads; // 5 for 4x2560
constexpr uint32_t kVecsPerBranch = kHiddenSize / kVecLen; // 320 for 2560
constexpr uint32_t kNumWarps = kNumThreads / kWarpThreads;
const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
const uint32_t m = blockIdx.x;
const auto y_ptr = pointer::offset<Float>(params.block_output, static_cast<int64_t>(m) * kHiddenSize);
const auto r_ptr = pointer::offset<Float>(params.residual, static_cast<int64_t>(m) * kRowSize);
const auto n_ptr = pointer::offset<Float>(params.normed_residual, static_cast<int64_t>(m) * kRowSize);
const auto w_ptr = static_cast<const Float*>(params.inject_weight);
const auto out_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(m) * kRowSize);
PDLWaitPrimary<kUsePDL>();
// Phase 1: gate values a_c, accumulated in fp32 and reduced across the CTA.
Storage n_vec[kVecsPerThread];
#pragma unroll
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
n_vec[j] = gmem.load(n_ptr, j);
}
float acc[kHcCount];
#pragma unroll
for (int c = 0; c < kHcCount; ++c) {
const auto wc_ptr = w_ptr + static_cast<int64_t>(c) * kRowSize;
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
const Storage w_vec = gmem.load(wc_ptr, j);
#pragma unroll
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
const auto [nx, ny] = cast<fp32x2_t>(n_vec[j][i]);
const auto [wx, wy] = cast<fp32x2_t>(w_vec[i]);
sum += nx * wx + ny * wy;
}
}
acc[c] = warp::reduce_sum(sum);
}
__shared__ float smem[kHcCount][kNumWarps];
const uint32_t warp_id = threadIdx.x / kWarpThreads;
const uint32_t lane = threadIdx.x % kWarpThreads;
if (lane == 0) {
#pragma unroll
for (int c = 0; c < kHcCount; ++c) {
smem[c][warp_id] = acc[c];
}
}
__syncthreads();
__shared__ float a_shared[kHcCount];
if (threadIdx.x < kHcCount) {
float total = 0.0f;
#pragma unroll
for (uint32_t w = 0; w < kNumWarps; ++w) {
total += smem[threadIdx.x][w];
}
a_shared[threadIdx.x] = 2.0f / (1.0f + math::exp(-total / kHcCount));
}
__syncthreads();
// Phase 2: stream the output row. Vector `vec_idx` lies entirely inside
// branch `vec_idx / kVecsPerBranch` (H is a multiple of kVecLen).
#pragma unroll
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
const uint32_t vec_idx = threadIdx.x + j * kNumThreads;
const uint32_t branch = vec_idx / kVecsPerBranch;
const uint32_t col_in_branch = (vec_idx % kVecsPerBranch) * kVecLen;
const float a = a_shared[branch];
const Storage r_vec = gmem.load(r_ptr, j);
Storage y_vec;
y_vec.load(y_ptr, col_in_branch / kVecLen);
Storage out_vec;
#pragma unroll
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
const auto [rx, ry] = cast<fp32x2_t>(r_vec[i]);
const auto [yx, yy] = cast<fp32x2_t>(y_vec[i]);
out_vec[i] = cast<Float2>(fp32x2_t{rx + a * yx, ry + a * yy});
}
gmem.store(out_ptr, out_vec, j);
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename DType>
struct HcCombineKernel {
static_assert(sizeof(DType) == 2, "HcCombine only supports 2-byte dtypes");
static_assert(kHcCount > 0, "kHcCount must be positive");
static_assert(kHiddenSize > 0 && kHiddenSize % 8 == 0, "kHiddenSize must be a multiple of 8");
static_assert((kHcCount * kHiddenSize) % (256 * 8) == 0, "kHcCount * kHiddenSize must be a multiple of 2048");
static constexpr auto kernel = hc_combine_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
static constexpr uint32_t kBlockSize = 256;
/**
* \brief Validate tensors and launch one CTA per token row.
* \param block_output [M, H] contiguous
* \param residual [M, HC * H] contiguous
* \param normed_residual [M, HC * H] contiguous, same dtype/device as residual
* \param inject_weight [HC, HC * H] contiguous, same dtype/device as residual
* \param output [M, HC * H] contiguous, same shape/dtype/device as residual
*/
static void
run(const tvm::ffi::TensorView block_output,
const tvm::ffi::TensorView residual,
const tvm::ffi::TensorView normed_residual,
const tvm::ffi::TensorView inject_weight,
const tvm::ffi::TensorView output) {
using namespace host;
auto M = SymbolicSize{"num_tokens"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({M, kHiddenSize}) // block_output
.with_dtype<DType>()
.with_device(device)
.verify(block_output);
TensorMatcher({M, kHcCount * kHiddenSize}) // residual, normed_residual, output
.with_dtype<DType>()
.with_device(device)
.verify(residual)
.verify(normed_residual)
.verify(output);
TensorMatcher({kHcCount, kHcCount * kHiddenSize}) // inject_weight
.with_dtype<DType>()
.with_device(device)
.verify(inject_weight);
const auto params = HcCombineParams{
.block_output = block_output.data_ptr(),
.residual = residual.data_ptr(),
.normed_residual = normed_residual.data_ptr(),
.inject_weight = inject_weight.data_ptr(),
.output = output.data_ptr(),
};
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
LaunchKernel(num_tokens, kBlockSize, device.unwrap()) //
.enable_pdl(kUsePDL)(kernel, params);
}
};
struct HcCombineSplitParams {
const void* block_output;
const void* residual;
const void* normed_residual;
const void* inject_weight;
void* output;
float* partials;
};
namespace hc_combine_split_detail {
// One CTA per warp of the original kernel: CTA s replays exactly the work of
// threads [32s, 32s+32) of the 256-thread reference, so every float is
// accumulated in the reference order and results stay bit-identical.
constexpr uint32_t kSplit = 8;
constexpr uint32_t kRefThreads = 256;
constexpr uint32_t kGateThreads = 32;
constexpr uint32_t kApplyThreads = 160;
constexpr uint32_t kVecLen = 8;
} // namespace hc_combine_split_detail
/**
* \brief Stage 1 of the split combine: partial gate dots over a K slice.
*
* Grid is (rows, kSplit) so the [HC, HC*H] inject weight is read once across
* the whole grid instead of once per row, and the row's traffic is spread over
* kSplit CTAs. Each CTA writes its own partials slot, so no atomics and no
* buffer clearing are needed.
*/
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
__global__ __launch_bounds__(hc_combine_split_detail::kGateThreads) void hc_combine_gate_kernel(
const HcCombineSplitParams __grid_constant__ params) {
using namespace device;
using namespace hc_combine_split_detail;
using Float2 = packed_t<Float>;
using Storage = AlignedVector<Float2, 4>;
constexpr uint32_t kVecLen = 8;
constexpr int64_t kRowSize = kHcCount * kHiddenSize;
constexpr uint32_t kVecsPerRow = kRowSize / kVecLen;
constexpr uint32_t kVecsPerThread = kVecsPerRow / kRefThreads;
static_assert(kVecsPerRow % kRefThreads == 0);
static_assert(kRefThreads / kGateThreads == kSplit);
const uint32_t m = blockIdx.x;
const uint32_t split = blockIdx.y / kHcCount;
const uint32_t c = blockIdx.y % kHcCount;
const uint32_t ref_tid = split * kGateThreads + threadIdx.x;
const auto n_ptr = pointer::offset<Float>(params.normed_residual, static_cast<int64_t>(m) * kRowSize);
const auto w_ptr = static_cast<const Float*>(params.inject_weight);
PDLWaitPrimary<kUsePDL>();
Storage n_vec[kVecsPerThread];
#pragma unroll
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
n_vec[j].load(n_ptr, ref_tid + j * kRefThreads);
}
{
const auto wc_ptr = w_ptr + static_cast<int64_t>(c) * kRowSize;
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
Storage w_vec;
w_vec.load(wc_ptr, ref_tid + j * kRefThreads);
#pragma unroll
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
const auto [nx, ny] = cast<fp32x2_t>(n_vec[j][i]);
const auto [wx, wy] = cast<fp32x2_t>(w_vec[i]);
sum += nx * wx + ny * wy;
}
}
sum = warp::reduce_sum(sum);
if (threadIdx.x == 0) {
params.partials[(static_cast<int64_t>(m) * kSplit + split) * kHcCount + c] = sum;
}
}
PDLTriggerSecondary<kUsePDL>();
}
/**
* \brief Stage 2: reduce the partial dots and stream the combined row.
*
* Each CTA owns one contiguous vector slice of the row. kSplit divides
* kHiddenSize, so a slice never straddles two branches and the gate is a
* per-CTA scalar.
*/
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
__global__ __launch_bounds__(hc_combine_split_detail::kApplyThreads) void hc_combine_apply_kernel(
const HcCombineSplitParams __grid_constant__ params) {
using namespace device;
using namespace hc_combine_split_detail;
using Float2 = packed_t<Float>;
using Storage = AlignedVector<Float2, 4>;
constexpr int64_t kRowSize = kHcCount * kHiddenSize;
constexpr uint32_t kVecsPerRow = kRowSize / kVecLen;
constexpr uint32_t kVecsPerSplit = kVecsPerRow / kSplit;
constexpr uint32_t kVecsPerThread = kVecsPerSplit / kApplyThreads;
constexpr uint32_t kVecsPerBranch = kHiddenSize / kVecLen;
static_assert(kVecsPerBranch % kVecsPerSplit == 0);
const uint32_t m = blockIdx.x;
const uint32_t split = blockIdx.y;
const uint32_t vec_base = split * kVecsPerSplit;
const uint32_t branch = vec_base / kVecsPerBranch;
const auto y_ptr = pointer::offset<Float>(params.block_output, static_cast<int64_t>(m) * kHiddenSize);
const auto r_ptr = pointer::offset<Float>(params.residual, static_cast<int64_t>(m) * kRowSize);
const auto out_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(m) * kRowSize);
PDLWaitPrimary<kUsePDL>();
float total = 0.0f;
#pragma unroll
for (uint32_t s = 0; s < kSplit; ++s) {
total += params.partials[(static_cast<int64_t>(m) * kSplit + s) * kHcCount + branch];
}
const float a = 2.0f / (1.0f + math::exp(-total / kHcCount));
#pragma unroll
for (uint32_t j = 0; j < kVecsPerThread; ++j) {
const uint32_t vec_idx = vec_base + threadIdx.x + j * kApplyThreads;
const uint32_t col_in_branch = (vec_idx % kVecsPerBranch);
Storage r_vec;
r_vec.load(r_ptr, vec_idx);
Storage y_vec;
y_vec.load(y_ptr, col_in_branch);
Storage out_vec;
#pragma unroll
for (uint32_t i = 0; i < kVecLen / 2; ++i) {
const auto [rx, ry] = cast<fp32x2_t>(r_vec[i]);
const auto [yx, yy] = cast<fp32x2_t>(y_vec[i]);
out_vec[i] = cast<Float2>(fp32x2_t{rx + a * yx, ry + a * yy});
}
out_vec.store(out_ptr, vec_idx);
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename DType>
struct HcCombineSplitKernel {
static_assert(sizeof(DType) == 2, "HcCombine only supports 2-byte dtypes");
static constexpr auto gate_kernel = hc_combine_gate_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
static constexpr auto apply_kernel = hc_combine_apply_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
static void
run(const tvm::ffi::TensorView block_output,
const tvm::ffi::TensorView residual,
const tvm::ffi::TensorView normed_residual,
const tvm::ffi::TensorView inject_weight,
const tvm::ffi::TensorView output,
const tvm::ffi::TensorView partials) {
using namespace host;
using namespace hc_combine_split_detail;
auto M = SymbolicSize{"num_tokens"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({M, kHiddenSize}).with_dtype<DType>().with_device(device).verify(block_output);
TensorMatcher({M, kHcCount * kHiddenSize})
.with_dtype<DType>()
.with_device(device)
.verify(residual)
.verify(normed_residual)
.verify(output);
TensorMatcher({kHcCount, kHcCount * kHiddenSize}).with_dtype<DType>().with_device(device).verify(inject_weight);
auto part_rows = SymbolicSize{"partial_rows"};
TensorMatcher({part_rows, kSplit, kHcCount}).with_dtype<fp32_t>().with_device(device).verify(partials);
const auto params = HcCombineSplitParams{
.block_output = block_output.data_ptr(),
.residual = residual.data_ptr(),
.normed_residual = normed_residual.data_ptr(),
.inject_weight = inject_weight.data_ptr(),
.output = output.data_ptr(),
.partials = static_cast<float*>(partials.data_ptr()),
};
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
LaunchKernel(dim3(num_tokens, kSplit * kHcCount, 1), kGateThreads, device.unwrap())
.enable_pdl(kUsePDL)(gate_kernel, params);
LaunchKernel(dim3(num_tokens, kSplit, 1), kApplyThreads, device.unwrap()).enable_pdl(kUsePDL)(apply_kernel, params);
}
};
} // namespace sglang
python/sglang/kernels/kda_kernels/README.md
+1
View File
@@ -20,6 +20,7 @@ load Triton, CUTLASS, or compile a JIT extension.
| Diffusion residual-gate add | `residual_gate_add_jit.py` | [sgl-project/sglang#29361](https://github.com/sgl-project/sglang/pull/29361), merge commit `495f13fa12` |
| LTX2 QK-norm split-RoPE | `ltx2_qknorm_split_rope_jit.py` | [sgl-project/sglang#29708](https://github.com/sgl-project/sglang/pull/29708), merge commit `fcb9f229b3` |
| FLUX.2 FP8 producer and QKV packing fusions | `layernorm_modulate_triton.py`, `flux2_qkv_epilogue_jit.py`, `flux2_token_cat_fp8_triton.py` | [sgl-project/sglang#37162](https://github.com/sgl-project/sglang/pull/37162), merge commit `1c3ad92438` |
| Qwen3.8 QSA packed-varlen decode on SM121 | `qwen38_qsa_sm121/` | [radixark/KDA-1.5#4](https://github.com/radixark/KDA-1.5/pull/4) at `414ce456e14a`; see the package README |
For JIT kernels, the Python entry module and the corresponding source under
`csrc/` move together. The shared `sglang.kernels.jit` loader remains build
python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/README.md
+51
View File
@@ -0,0 +1,51 @@
# Qwen3.8 QSA packed-varlen decode for SM121
This implementation was optimized by Codex and Kimi K3 agents through
[KDA-1.5](https://github.com/radixark/KDA-1.5). The task and immutable real
tensor replay were registered in [radixark/KDA-1.5 PR #4](https://github.com/radixark/KDA-1.5/pull/4)
at commit `414ce456e14ae8546f77d9356d2c4d955c5bb7f1`. This package integrates
winning submission `b4181149c8884ddb`; its byte-exact submitted source has SHA256
`4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909`.
The kernel is specialized for the packed QSA decode tensors captured from
`RadixArk/Qwen3.8-Flash-Next-NVFP4` on NVIDIA GB10 (SM121):
- BF16 query, key, value, and output with head dimension 256
- one packed query row per sequence and device-side `cu_seqlens`
- 12 query heads per KV head: TP1 uses 24Q/2KV and TP2 uses 12Q/1KV
- all query-row counts in the validated `1 <= bs <= 128` envelope
- `max_seqlen_k` capacity up to 2055 and captured logical selected-KV lengths
up to 2051 rows per sequence
The implementation groups the 12 query heads that share one KV head into one
CTA, uses BF16 tensor-core QK/PV products with FP32 online-softmax state, and
splits long KV rows across multiple CTAs. The last arriving split performs a
stable FP32 merge and resets its device counter in the same launch. A
host-visible shape/topology policy selects the two measured schedules, while
the live device `cu_seqlens_k` selects one, two, four, or eight active splits
without a host synchronization.
SM121 dispatch checks the exact Qwen3.8 contract and routes directly to this
kernel; it is the only packed-QSA attention implementation added by this PR.
The KDA replay passed all 15 TP1/TP2 production tensors on two independent GB10
GPUs, and the final source passed 150,000 consecutive launches with all
counters returning to zero.
After adaptation into SGLang, the packaged kernel passed the same 15/15 replay
with exactly one CUDA activity per row and a 2.0702x all-shape geomean over the
generic Triton fallback (1.6951x large, 2.3653x small). On one DGX Spark running
the full TP1 NVFP4 model with NEXTN, three-round low-concurrency serving A/B
improved total token throughput by 4.45% at concurrency 1 and 4.00% at
concurrency 4. A 50-example, five-shot GSM8K A/B with a 2048-token output limit
scored 49/50 for both Triton and KDA, with the same single failed example.
An additional synthetic GB10 sweep covers both TP topologies, every batch size
from 1 through 16, and short plus saturated KV rows. All 64 cases passed; the
maximum relative L2 against the original correct Triton implementation was
0.002422, and speedup ranged from 1.41x to 5.09x.
A follow-up extended-batch sweep covers both TP topologies, batch sizes
17/24/32/48/64/96/128, and short, medium, plus saturated KV rows. All 42 cases
passed with maximum relative L2 0.002410. Geomean speedup was 4.48x, the slowest
case still improved by 1.58x, and no case regressed. The packaged scratch space
is therefore sized for the largest tested batch, 128.
python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/__init__.py
+83
View File
@@ -0,0 +1,83 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import logging
import torch
logger = logging.getLogger(__name__)
_SUPPORTED_HEAD_TOPOLOGIES = frozenset({(12, 1), (24, 2)})
# Largest batch qualified by the extended GB10 baseline sweep.
_MAX_BATCH = 128
_MAX_SELECTED_KV = 2055
_logged_fast_path = False
def can_use_qwen38_qsa_sm121(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_k: int,
) -> bool:
"""Return whether this call matches the captured Qwen3.8 SM121 contract."""
if not q.is_cuda or q.ndim != 3 or q.dtype != torch.bfloat16:
return False
batch, num_q_heads, head_dim = q.shape
if not (0 < batch <= _MAX_BATCH) or head_dim != 256:
return False
if k.ndim != 3 or v.shape != k.shape or k.dtype != q.dtype or v.dtype != q.dtype:
return False
num_kv_heads = k.shape[1]
if (num_q_heads, num_kv_heads) not in _SUPPORTED_HEAD_TOPOLOGIES:
return False
if k.shape[2] != head_dim or not (0 < max_seqlen_k <= _MAX_SELECTED_KV):
return False
if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous():
return False
if q.device != k.device or q.device != v.device:
return False
if cu_seqlens_q.device != q.device or cu_seqlens_k.device != q.device:
return False
if cu_seqlens_q.dtype != torch.int32 or cu_seqlens_k.dtype != torch.int32:
return False
if cu_seqlens_q.ndim != 1 or cu_seqlens_k.ndim != 1:
return False
if not cu_seqlens_q.is_contiguous() or not cu_seqlens_k.is_contiguous():
return False
if cu_seqlens_q.numel() != batch + 1 or cu_seqlens_k.numel() != batch + 1:
return False
properties = torch.cuda.get_device_properties(q.device)
return (properties.major, properties.minor) == (12, 1)
def qwen38_qsa_sm121(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_k: int,
softmax_scale: float,
) -> torch.Tensor:
"""Run the KDA-generated Qwen3.8 packed QSA decode kernel."""
global _logged_fast_path
if not can_use_qwen38_qsa_sm121(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k):
raise ValueError("unsupported call for the KDA Qwen3.8 SM121 QSA kernel")
from .kernel import qwen38_qsa_sm121 as run_kernel
if not _logged_fast_path:
logger.info(
"Using the Codex/Kimi K3 KDA Qwen3.8 QSA kernel on SM121 "
"(radixark/KDA-1.5#4, submission b4181149c8884ddb)"
)
_logged_fast_path = True
return run_kernel(q, k, v, cu_seqlens_q, cu_seqlens_k, softmax_scale)
__all__ = ["can_use_qwen38_qsa_sm121", "qwen38_qsa_sm121"]
python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/kernel.py
+269
View File
@@ -0,0 +1,269 @@
# SPDX-License-Identifier: Apache-2.0
# KDA-1.5 submission b4181149c8884ddb (https://github.com/radixark/KDA-1.5/pull/4);
# source SHA256 4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909.
"""Shape-specialized Qwen3.8 packed QSA decode kernel for SM121."""
from __future__ import annotations
import torch
import triton
import triton.language as tl
@triton.jit
def _qsa_split_kernel(
q_ptr,
k_ptr,
v_ptr,
out_ptr,
cu_q_ptr,
cu_k_ptr,
partial_max_ptr,
partial_sum_ptr,
partial_acc_ptr,
counter_ptr,
softmax_scale,
NUM_Q_HEADS: tl.constexpr,
NUM_KV_HEADS: tl.constexpr,
HEAD_DIM: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_KV: tl.constexpr,
MAX_SPLITS: tl.constexpr,
q_stride_t: tl.constexpr,
q_stride_h: tl.constexpr,
k_stride_t: tl.constexpr,
k_stride_h: tl.constexpr,
v_stride_t: tl.constexpr,
v_stride_h: tl.constexpr,
out_stride_t: tl.constexpr,
out_stride_h: tl.constexpr,
):
sequence_idx = tl.program_id(0)
split_program = tl.program_id(1)
kv_head_idx = split_program // MAX_SPLITS
split_idx = split_program - kv_head_idx * MAX_SPLITS
slot = sequence_idx * NUM_KV_HEADS + kv_head_idx
queries_per_kv = NUM_Q_HEADS // NUM_KV_HEADS
query_idx = tl.load(cu_q_ptr + sequence_idx)
kv_begin = tl.load(cu_k_ptr + sequence_idx)
kv_end = tl.load(cu_k_ptr + sequence_idx + 1)
kv_count = kv_end - kv_begin
tile_count = tl.cdiv(kv_count, BLOCK_KV)
# The split count depends only on live device metadata and the static
# launch geometry, so CUDA Graph replay needs no host readback.
batch = tl.num_programs(0)
n_splits = 1
if tile_count >= 1536 // BLOCK_KV:
n_splits = 2
elif tile_count >= 512 // BLOCK_KV and batch * NUM_KV_HEADS <= 4:
n_splits = 4
if batch == 1:
if tile_count >= 1024 // BLOCK_KV:
n_splits = 8
elif tile_count >= 512 // BLOCK_KV:
n_splits = 4
elif tile_count >= 256 // BLOCK_KV:
n_splits = 2
if split_idx >= n_splits:
return
tile_lo = (tile_count * split_idx) // n_splits
tile_hi = (tile_count * (split_idx + 1)) // n_splits
kv_start = kv_begin + tile_lo * BLOCK_KV
kv_stop = tl.minimum(kv_begin + tile_hi * BLOCK_KV, kv_end)
m = tl.arange(0, BLOCK_M)
d = tl.arange(0, HEAD_DIM)
q_head = kv_head_idx * queries_per_kv + m
q_mask = m < queries_per_kv
query = tl.load(
q_ptr + query_idx * q_stride_t + q_head[:, None] * q_stride_h + d[None, :],
mask=q_mask[:, None],
other=0.0,
)
n0 = tl.arange(0, BLOCK_KV)
k_rows = (
k_ptr
+ kv_head_idx * k_stride_h
+ kv_start.to(tl.int64) * k_stride_t
+ n0[:, None] * k_stride_t
)
v_rows = (
v_ptr
+ kv_head_idx * v_stride_h
+ kv_start.to(tl.int64) * v_stride_t
+ n0[:, None] * v_stride_t
)
running_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
running_sum = tl.zeros([BLOCK_M], tl.float32)
accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
kv_len = kv_stop - kv_start
full_end = kv_start + (kv_len // BLOCK_KV) * BLOCK_KV
for block_start in range(kv_start, full_end, BLOCK_KV):
keys = tl.load(k_rows + d[None, :])
scores = tl.dot(query, tl.trans(keys)) * softmax_scale
new_max = tl.maximum(running_max, tl.max(scores, axis=1))
old_scale = tl.exp(running_max - new_max)
probabilities = tl.exp(scores - new_max[:, None])
running_sum = running_sum * old_scale + tl.sum(probabilities, axis=1)
values = tl.load(v_rows + d[None, :])
accumulator = accumulator * old_scale[:, None] + tl.dot(
probabilities.to(tl.bfloat16), values
)
running_max = new_max
k_rows += BLOCK_KV * k_stride_t
v_rows += BLOCK_KV * v_stride_t
if full_end < kv_stop:
n = full_end + n0
n_mask = n < kv_stop
keys = tl.load(k_rows + d[None, :], mask=n_mask[:, None], other=0.0)
scores = tl.dot(query, tl.trans(keys)) * softmax_scale
scores = tl.where(n_mask[None, :], scores, -float("inf"))
new_max = tl.maximum(running_max, tl.max(scores, axis=1))
old_scale = tl.exp(running_max - new_max)
probabilities = tl.exp(scores - new_max[:, None])
running_sum = running_sum * old_scale + tl.sum(probabilities, axis=1)
values = tl.load(v_rows + d[None, :], mask=n_mask[:, None], other=0.0)
accumulator = accumulator * old_scale[:, None] + tl.dot(
probabilities.to(tl.bfloat16), values
)
running_max = new_max
if n_splits == 1:
output = accumulator / tl.where(running_sum > 0.0, running_sum, 1.0)[:, None]
tl.store(
out_ptr
+ query_idx * out_stride_t
+ q_head[:, None] * out_stride_h
+ d[None, :],
output.to(out_ptr.dtype.element_ty),
mask=q_mask[:, None],
)
return
partial_row = (slot * MAX_SPLITS + split_idx) * BLOCK_M + m
tl.store(partial_max_ptr + partial_row, running_max)
tl.store(partial_sum_ptr + partial_row, running_sum)
tl.store(
partial_acc_ptr + partial_row[:, None] * HEAD_DIM + d[None, :],
accumulator,
)
tl.debug_barrier()
arrival = tl.atomic_add(counter_ptr + slot, 1, sem="acq_rel", scope="gpu")
if arrival == n_splits - 1:
merged_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
for j in tl.static_range(MAX_SPLITS):
j_ok = (j < n_splits) & (m < BLOCK_M)
row = (slot * MAX_SPLITS + j) * BLOCK_M + m
mj = tl.load(partial_max_ptr + row, mask=j_ok, other=-float("inf"))
merged_max = tl.maximum(merged_max, mj)
merged_sum = tl.zeros([BLOCK_M], tl.float32)
merged_acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
for j in tl.static_range(MAX_SPLITS):
j_ok = (j < n_splits) & (m < BLOCK_M)
row = (slot * MAX_SPLITS + j) * BLOCK_M + m
mj = tl.load(partial_max_ptr + row, mask=j_ok, other=-float("inf"))
lj = tl.load(partial_sum_ptr + row, mask=j_ok, other=0.0)
weight = tl.exp(mj - merged_max)
merged_sum += weight * lj
partial = tl.load(
partial_acc_ptr + row[:, None] * HEAD_DIM + d[None, :],
mask=j_ok[:, None],
other=0.0,
)
merged_acc += weight[:, None] * partial
output = merged_acc / tl.where(merged_sum > 0.0, merged_sum, 1.0)[:, None]
tl.store(
out_ptr
+ query_idx * out_stride_t
+ q_head[:, None] * out_stride_h
+ d[None, :],
output.to(out_ptr.dtype.element_ty),
mask=q_mask[:, None],
)
tl.atomic_xchg(counter_ptr + slot, 0, sem="release", scope="gpu")
# The worst-case TP1 scratch allocation at the qualified limit is 32.3 MiB.
_MAX_BATCH = 128
_MAX_KV_HEADS = 2
_BLOCK_M = 16
_MAX_SPLITS = 8
_MAX_SLOTS = _MAX_BATCH * _MAX_KV_HEADS
_HEAD_DIM = 256
_scratch: dict[int, tuple[torch.Tensor, ...]] = {}
def _get_scratch(device: torch.device) -> tuple[torch.Tensor, ...]:
device_index = (
device.index if device.index is not None else torch.cuda.current_device()
)
scratch = _scratch.get(device_index)
if scratch is None:
rows = _MAX_SLOTS * _MAX_SPLITS * _BLOCK_M
scratch = (
torch.empty(rows, dtype=torch.float32, device=device),
torch.empty(rows, dtype=torch.float32, device=device),
torch.empty(rows * _HEAD_DIM, dtype=torch.float32, device=device),
torch.zeros(_MAX_SLOTS, dtype=torch.int32, device=device),
)
_scratch[device_index] = scratch
return scratch
def qwen38_qsa_sm121(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
softmax_scale: float,
) -> torch.Tensor:
"""Run the KDA-generated SM121 Qwen3.8 QSA kernel."""
output = torch.empty_like(q)
partial_max, partial_sum, partial_acc, counters = _get_scratch(q.device)
batch, num_q_heads, head_dim = q.shape
num_kv_heads = k.shape[1]
# Measured schedule. The TP1 q_rows=4 shape stays on BK64 on purpose:
# its short and saturated rows cannot be told apart from the host.
use_bk32 = (num_kv_heads == 1 and batch < 12) or (num_kv_heads == 2 and batch < 4)
block_kv = 32 if use_bk32 else 64
stages = 3 if use_bk32 else 2
_qsa_split_kernel[(batch, num_kv_heads * _MAX_SPLITS)](
q,
k,
v,
output,
cu_seqlens_q,
cu_seqlens_k,
partial_max,
partial_sum,
partial_acc,
counters,
softmax_scale,
NUM_Q_HEADS=num_q_heads,
NUM_KV_HEADS=num_kv_heads,
HEAD_DIM=head_dim,
BLOCK_M=_BLOCK_M,
BLOCK_KV=block_kv,
MAX_SPLITS=_MAX_SPLITS,
q_stride_t=q.stride(0),
q_stride_h=q.stride(1),
k_stride_t=k.stride(0),
k_stride_h=k.stride(1),
v_stride_t=v.stride(0),
v_stride_h=v.stride(1),
out_stride_t=output.stride(0),
out_stride_h=output.stride(1),
num_warps=4,
num_stages=stages,
)
return output
python/sglang/kernels/ops/attention/__init__.py
+86 -2
View File
@@ -6,8 +6,21 @@ The Triton kernels migrated here live in this package
KV-cache index/write kernels went to the ``kvcache`` group instead.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from sglang.kernels.registry import register_kernel
from sglang.kernels.spec import KernelBackend, KernelSpec
from sglang.kernels.selector import get_kernel
from sglang.kernels.spec import (
CapabilityRequirement,
FormatSignature,
KernelBackend,
KernelSpec,
)
if TYPE_CHECKING:
import torch
# (module, public_fn) migrated from layers/attention/triton_ops + model_executor.
_TRITON_KERNELS = [
@@ -45,7 +58,78 @@ for _mod, _fn in _TRITON_KERNELS:
)
del _mod, _fn
__all__ = []
register_kernel(
KernelSpec(
op="attention.kda_qwen38_qsa_sm121",
backend=KernelBackend.TRITON,
target=("sglang.kernels.kda_kernels.qwen38_qsa_sm121:qwen38_qsa_sm121"),
capabilities=frozenset(
{CapabilityRequirement.cuda(min_sm=(12, 1), max_sm=(12, 1))}
),
format_signature=FormatSignature(
supported_dtypes=("bfloat16",),
description=(
"Qwen3.8 packed QSA decode: D=256, 12:1 GQA, 1 <= q_rows <= 128"
),
),
description=(
"SM121 Qwen3.8 QSA decode optimized by Codex/Kimi K3 through KDA-1.5."
),
)
)
def can_use_kda_qwen38_qsa_sm121(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_k: int,
) -> bool:
"""Check the exact E2E-qualified Qwen3.8/SM121 QSA contract."""
from sglang.kernels.kda_kernels.qwen38_qsa_sm121 import (
can_use_qwen38_qsa_sm121,
)
return can_use_qwen38_qsa_sm121(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k)
def qwen38_qsa_sm121_varlen(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_q: int = 1,
max_seqlen_k: int = 0,
softmax_scale: float = 1.0,
causal: bool = True,
**_: object,
) -> torch.Tensor:
"""Run the only SM121 packed-QSA kernel for its qualified contract."""
del causal
if max_seqlen_q != 1:
raise ValueError(f"QSA requires max_seqlen_q=1, got {max_seqlen_q}")
if not can_use_kda_qwen38_qsa_sm121(
q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k
):
raise ValueError(
"unsupported SM121 QSA call: expected BF16 D=256, 12:1 GQA, "
"TP1 24Q/2KV or TP2 12Q/1KV, bs<=128, and selected KV<=2055"
)
return get_kernel("attention.kda_qwen38_qsa_sm121", KernelBackend.TRITON)(
q,
k,
v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_k,
softmax_scale,
)
__all__ = ["can_use_kda_qwen38_qsa_sm121", "qwen38_qsa_sm121_varlen"]
# Vendored linear-attention (flash-linear-attention port) kernels relocated
python/sglang/kernels/ops/attention/qsa_indexer.py
+176
View File
@@ -0,0 +1,176 @@
"""Fused QSA (Qwen4-Exp sparse attention) indexer-prep kernels.
``qsa_index_q_norm_rope_store`` fuses, per token, the eager chain
split -> GemmaRMSNorm(index q) -> MRoPE(index q) -> raw-K store ->
RoPE-position store into one kernel launch.
``qsa_index_k_compress_store`` fuses, per completed compress group, the eager
chain gather -> fp32 mean -> GemmaRMSNorm -> MRoPE(group-start position) ->
compressed-cache store into one kernel launch.
Both kernels reproduce the eager numerics step by step (fp32 norm reduction,
per-op rounding to the storage dtype during RoPE, fp32 group mean rounded to
the storage dtype before the norm). Outputs are bit-comparable to the eager
indexer path: the eager RMSNorm (flashinfer's CuTe DSL kernel) reduces sums of
squares in an order that cannot be reproduced exactly, so a small fraction of
rows (~1 in 30k) may flip by one bf16 ulp on a rounding boundary.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_qsa_indexer_module(
dtype: torch.dtype, head_dim: int, is_neox_style: bool
) -> Module:
"""Compile and cache the JIT QSA indexer module for one specialisation."""
if dtype not in (torch.bfloat16, torch.float16):
raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
if head_dim not in (64, 128, 256):
raise RuntimeError(
f"Unsupported index head_dim {head_dim}. Supported: 64, 128, 256"
)
args = make_cpp_args(dtype, head_dim, is_neox_style, is_arch_support_pdl())
return load_jit(
"qsa_indexer",
*args,
cuda_files=["attention/qsa_indexer.cuh"],
cuda_wrappers=[
("q_prep", f"qsa_index_q_prep<{args}>"),
("k_compress", f"qsa_index_k_compress<{args}>"),
],
)
def qsa_index_q_norm_rope_store(
qk: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
axis_map: torch.Tensor,
weight: torch.Tensor,
cache_loc: torch.Tensor,
key_state_buffer: torch.Tensor,
rope_position_buffer: torch.Tensor,
num_q_heads: int,
rotary_dim: int,
eps: float,
is_neox_style: bool,
q_heads_padded: Optional[int] = None,
) -> torch.Tensor:
"""
Per-token fused index-Q prep.
Parameters
----------
qk : CUDA bf16/fp16 [tokens, (num_q_heads + 1) * head_dim],
contiguous fused index Q/K projection output
positions : CUDA int64 [tokens] or [3, tokens] RoPE positions
(a strided trailing slice is accepted)
cos_sin_cache : CUDA fp32 [capacity, rotary_dim] RoPE cache
axis_map : CUDA int32 [rotary_dim // 2] position-axis per pair index
weight : [head_dim] gemma norm weight (kernel applies 1 + w)
cache_loc : CUDA int64 [tokens] state slots of each token
key_state_buffer : CUDA [slots, head_dim] raw-K state buffer (written)
rope_position_buffer : CUDA int64 [slots, 3] position buffer (written)
num_q_heads : number of index query heads
rotary_dim : rotated prefix of each head row
eps : RMSNorm epsilon
is_neox_style : NeoX (True) or GPT-J (False) RoPE pairing
q_heads_padded : output head count; heads >= num_q_heads are zero-filled
(defaults to num_q_heads)
Returns
-------
CUDA tensor [tokens, q_heads_padded, head_dim]: normed + rotated index Q.
"""
num_tokens = qk.shape[0]
head_dim = weight.shape[0]
if q_heads_padded is None:
q_heads_padded = num_q_heads
if positions.ndim == 1:
positions = positions.unsqueeze(0)
q_out = torch.empty(
(num_tokens, q_heads_padded, head_dim), dtype=qk.dtype, device=qk.device
)
module = _jit_qsa_indexer_module(qk.dtype, head_dim, is_neox_style)
module.q_prep(
qk,
q_out,
weight,
cos_sin_cache,
axis_map,
positions,
positions.shape[0],
cache_loc,
key_state_buffer,
rope_position_buffer,
num_q_heads,
rotary_dim,
eps,
)
return q_out
def qsa_index_k_compress_store(
key_state_buffer: torch.Tensor,
group_locs: torch.Tensor,
rope_position_buffer: torch.Tensor,
cos_sin_cache: torch.Tensor,
axis_map: torch.Tensor,
weight: torch.Tensor,
write_locs: torch.Tensor,
compressed_k_buffer: torch.Tensor,
compress_ratio: int,
rotary_dim: int,
eps: float,
is_neox_style: bool,
) -> None:
"""
Per-group compressed-K prep (in-place store into ``compressed_k_buffer``).
Parameters
----------
key_state_buffer : CUDA bf16/fp16 [slots, head_dim] raw-K state buffer
group_locs : CUDA int32 [groups, compress_ratio] state slots of each
completed group, group-start slot in column 0 (its RoPE
position rotates the group; other columns any order)
rope_position_buffer : CUDA int64 [slots, 3] per-slot RoPE coordinates
cos_sin_cache : CUDA fp32 [capacity, rotary_dim] RoPE cache
axis_map : CUDA int32 [rotary_dim // 2] position-axis per pair index
weight : [head_dim] gemma norm weight (kernel applies 1 + w)
write_locs : CUDA int32 [groups] compressed-cache slots to write
compressed_k_buffer : CUDA [compressed_slots, head_dim] (written)
compress_ratio : raw keys per compressed key
rotary_dim : rotated prefix of each head row
eps : RMSNorm epsilon
is_neox_style : NeoX (True) or GPT-J (False) RoPE pairing
"""
head_dim = weight.shape[0]
module = _jit_qsa_indexer_module(key_state_buffer.dtype, head_dim, is_neox_style)
module.k_compress(
key_state_buffer,
group_locs,
rope_position_buffer,
cos_sin_cache,
axis_map,
weight,
write_locs,
compressed_k_buffer,
compress_ratio,
rotary_dim,
eps,
)
python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py
+17 -7
View File
@@ -171,6 +171,8 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
a,
mixed_qkvz,
mixed_ba,
qkvz_row_stride,
ba_row_stride,
NUM_HEADS_QK: tl.constexpr,
NUM_HEADS_V: tl.constexpr,
HEAD_QK: tl.constexpr,
@@ -185,19 +187,19 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
TOTAL_Q: tl.constexpr = NUM_HEADS_QK * HEAD_QK
TOTAL_K: tl.constexpr = NUM_HEADS_QK * HEAD_QK
TOTAL_V: tl.constexpr = NUM_HEADS_V * HEAD_V
TOTAL_QKVZ: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V + TOTAL_V
TOTAL_BA: tl.constexpr = NUM_HEADS_V * 2
# ── Output dimensions ──
QKV_DIM_T: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V
# ── Read from contiguous input ──
# q for head group i_qk: in the all_q region, offset i_qk * HEAD_QK
blk_q_ptr = mixed_qkvz + i_bs * TOTAL_QKVZ + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
blk_q_ptr = (
mixed_qkvz + i_bs * qkvz_row_stride + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
)
# k for head group i_qk: in the all_k region
blk_k_ptr = (
mixed_qkvz
+ i_bs * TOTAL_QKVZ
+ i_bs * qkvz_row_stride
+ TOTAL_Q
+ i_qk * HEAD_QK
+ tl.arange(0, HEAD_QK)
@@ -209,7 +211,11 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
# vector access. V_POW2 arrives as a wrapper-computed constexpr so the
# dead branch is pruned before tl.arange validation.
v_ld_base = (
mixed_qkvz + i_bs * TOTAL_QKVZ + TOTAL_Q + TOTAL_K + i_qk * V_PER_GROUP * HEAD_V
mixed_qkvz
+ i_bs * qkvz_row_stride
+ TOTAL_Q
+ TOTAL_K
+ i_qk * V_PER_GROUP * HEAD_V
)
z_ld_base = v_ld_base + TOTAL_V
@@ -250,12 +256,14 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
# ── b and a from contiguous [all_b | all_a] ──
for i in tl.static_range(V_PER_GROUP):
blk_b_ptr = mixed_ba + i_bs * TOTAL_BA + i_qk * V_PER_GROUP + i
blk_b_ptr = mixed_ba + i_bs * ba_row_stride + i_qk * V_PER_GROUP + i
blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
tl.store(blk_b_st_ptr, tl.load(blk_b_ptr))
for i in tl.static_range(V_PER_GROUP):
blk_a_ptr = mixed_ba + i_bs * TOTAL_BA + NUM_HEADS_V + i_qk * V_PER_GROUP + i
blk_a_ptr = (
mixed_ba + i_bs * ba_row_stride + NUM_HEADS_V + i_qk * V_PER_GROUP + i
)
blk_a_st_ptr = a + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
tl.store(blk_a_st_ptr, tl.load(blk_a_ptr))
@@ -319,6 +327,8 @@ def fused_qkvzba_split_reshape_cat_contiguous(
a,
mixed_qkvz,
mixed_ba,
mixed_qkvz.stride(0),
mixed_ba.stride(0),
num_heads_qk,
num_heads_v,
head_qk,
python/sglang/kernels/ops/elementwise/fast_topk.py
+70
View File
@@ -0,0 +1,70 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
if TYPE_CHECKING:
from tvm_ffi.module import Module
_FAST_TOPK_SUPPORTED_K = (512, 2048)
@cache_once
def _jit_fast_topk_module(topk: int) -> Module:
"""Compile and cache the JIT fast top-k module for a given k."""
# Checks on the compile key live here, not in `fast_topk`: `cache_once`
# keys on `topk`, so this runs once per specialisation.
if topk not in _FAST_TOPK_SUPPORTED_K:
raise RuntimeError(
f"Unsupported topk {topk}. Supported: {_FAST_TOPK_SUPPORTED_K}"
)
args = make_cpp_args(topk, is_arch_support_pdl())
return load_jit(
"fast_topk",
*args,
cuda_files=["elementwise/fast_topk.cuh"],
cuda_wrappers=[("fast_topk", f"FastTopKKernel<{args}>::run")],
)
def fast_topk(
score: torch.Tensor,
lengths: torch.Tensor,
topk: int,
row_starts: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Per-row top-k selection over a fp32 score matrix.
Row b selects the `topk` largest values in
``score[b, row_starts[b] : row_starts[b] + lengths[b]]`` and returns their
indices relative to ``row_starts[b]``. Slots beyond ``lengths[b]`` are -1.
Output order within a row is unspecified (atomic collection order).
Parameters
----------
score : CUDA fp32 tensor [B, L]
lengths : CUDA int32 tensor [B]
topk : number of indices per row; 512 or 2048
row_starts : optional CUDA int32 tensor [B]; defaults to zeros
Returns
-------
CUDA int32 tensor [B, topk]
"""
batch = score.shape[0]
if row_starts is None:
row_starts = torch.zeros(batch, dtype=torch.int32, device=score.device)
indices = score.new_empty((batch, topk), dtype=torch.int32)
module = _jit_fast_topk_module(topk)
module.fast_topk(score, row_starts, indices, lengths)
return indices
python/sglang/kernels/ops/elementwise/hc_combine.py
+135
View File
@@ -0,0 +1,135 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_hc_combine_module(
hc_count: int, hidden_size: int, dtype: torch.dtype
) -> Module:
"""Compile and cache the JIT HC combine module for a given shape/dtype."""
# Validation lives here rather than in `hc_combine`,
# so `cache_once` runs it once per (hc_count, hidden_size, dtype), not per call.
if dtype not in (torch.bfloat16, torch.float16):
raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
if hidden_size <= 0 or hidden_size % 8 != 0:
raise RuntimeError(
f"Unsupported hidden_size {hidden_size}. Must be a multiple of 8."
)
if hc_count <= 0 or (hc_count * hidden_size) % 2048 != 0:
raise RuntimeError(
f"Unsupported hc_count * hidden_size {hc_count * hidden_size}. "
"Must be a multiple of 2048."
)
args = make_cpp_args(hc_count, hidden_size, is_arch_support_pdl(), dtype)
return load_jit(
"hc_combine",
*args,
cuda_files=["elementwise/hc_combine.cuh"],
cuda_wrappers=[
("hc_combine", f"HcCombineKernel<{args}>::run"),
("hc_combine_split", f"HcCombineSplitKernel<{args}>::run"),
],
)
def hc_combine(
block_output: torch.Tensor,
residual: torch.Tensor,
normed_residual: torch.Tensor,
inject_weight: torch.Tensor,
hc_count: int,
hidden_size: int,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Fused HyperConnection (gated residual) combine.
a[m, c] = 2 * sigmoid(dot(normed_residual[m], inject_weight[c]) / hc_count)
out[m, c*H + i] = residual[m, c*H + i] + a[m, c] * block_output[m, i]
Mirrors ``GatedResidual._combine_compute`` in
``sglang.srt.layers.hyperconnection``. All math is accumulated in fp32.
Supported dtypes: torch.bfloat16, torch.float16.
Parameters
----------
block_output : CUDA tensor [..., hidden_size]
residual : CUDA tensor [..., hc_count * hidden_size]
normed_residual : CUDA tensor, same shape/dtype as residual
inject_weight : CUDA tensor [hc_count, hc_count * hidden_size]
hc_count : number of hyper-connection branches
hidden_size : per-branch hidden size
out : optional pre-allocated output tensor (same shape/dtype as residual)
Returns
-------
Combined tensor, same shape/dtype as residual.
"""
y = block_output.reshape(-1, hidden_size)
r = residual.reshape(-1, hc_count * hidden_size)
n = normed_residual.reshape(-1, hc_count * hidden_size)
if out is None:
out = torch.empty_like(r)
else:
out = out.reshape(-1, hc_count * hidden_size)
module = _jit_hc_combine_module(hc_count, hidden_size, residual.dtype)
module.hc_combine(y, r, n, inject_weight, out)
return out.reshape(residual.shape)
_SPLIT = 8
_MAX_ROWS = 32
_partials_cache = {}
def _get_partials(hc_count: int, device: torch.device, rows: int) -> torch.Tensor:
key = (device, hc_count)
buf = _partials_cache.get(key)
if buf is None or buf.shape[0] < rows:
# The gate kernel writes one slot per (row, split, hc); a buffer shorter
# than `rows` is written past its end rather than truncated.
buf = torch.empty(
(max(rows, _MAX_ROWS), _SPLIT, hc_count),
dtype=torch.float32,
device=device,
)
_partials_cache[key] = buf
return buf
def hc_combine_split(
block_output: torch.Tensor,
residual: torch.Tensor,
normed_residual: torch.Tensor,
inject_weight: torch.Tensor,
hc_count: int,
hidden_size: int,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
y = block_output.reshape(-1, hidden_size)
r = residual.reshape(-1, hc_count * hidden_size)
n = normed_residual.reshape(-1, hc_count * hidden_size)
if out is None:
out = torch.empty_like(r)
else:
out = out.reshape(-1, hc_count * hidden_size)
rows = r.shape[0]
partials = _get_partials(hc_count, r.device, rows)[:rows]
module = _jit_hc_combine_module(hc_count, hidden_size, residual.dtype)
module.hc_combine_split(y, r, n, inject_weight, out, partials)
return out.reshape(residual.shape)
python/sglang/kernels/ops/elementwise/hc_mix.py
+118
View File
@@ -0,0 +1,118 @@
from __future__ import annotations
import torch
_MAX_ROWS = 32
_KPAD_ALIGN = 128
def pad_lowrank(lowrank: int) -> int:
return (lowrank + _KPAD_ALIGN - 1) // _KPAD_ALIGN * _KPAD_ALIGN
def hc_mix_shape_ok(hc_count: int, hidden_size: int, lowrank: int) -> bool:
k = hc_count * hidden_size
return (
hidden_size % 8 == 0
and lowrank > 0
and lowrank % 8 == 0
and k % _KPAD_ALIGN == 0
and k % 128 == 0
)
def permute_pad_up_weight(w_up: torch.Tensor, hc_count: int) -> torch.Tensor:
k, lowrank = w_up.shape
hidden_size = k // hc_count
padded = torch.zeros(
(k, pad_lowrank(lowrank)), dtype=w_up.dtype, device=w_up.device
)
padded[:, :lowrank] = w_up.detach()
src = torch.arange(k, device=w_up.device)
perm = torch.empty(k, dtype=torch.long, device=w_up.device)
perm[(src % hidden_size) * hc_count + (src // hidden_size)] = src
return padded[perm].contiguous()
_scratch_cache = {}
_tactic_cache = {}
def _get_scratch(lowrank: int, dtype: torch.dtype, device: torch.device):
key = (device, lowrank, dtype)
buf = _scratch_cache.get(key)
if buf is None:
buf = torch.zeros((_MAX_ROWS, pad_lowrank(lowrank)), dtype=dtype, device=device)
_scratch_cache[key] = buf
return buf
def _get_tactic(m: int, n: int, k: int):
key = (m, n, k)
tactic = _tactic_cache.get(key)
if tactic is None:
from sglang.kernels.ops.gemm.dense_bf16_gemm_sm100_splitk_epilogue import (
SplitKTactic,
default_tactic,
validate_tactic,
)
if n <= 512 and k % (128 * 16) == 0:
tactic = SplitKTactic(mma_m=64, mma_n=8, split_k=16, ab_stages=6)
try:
validate_tactic(tactic, m, n, k)
except ValueError:
tactic = default_tactic(m, n, k)
else:
tactic = default_tactic(m, n, k)
_tactic_cache[key] = tactic
return tactic
def hc_mix(
hyper_input_normed: torch.Tensor,
w_down: torch.Tensor,
w_up_permuted_padded: torch.Tensor,
hc_count: int,
hidden_size: int,
) -> torch.Tensor:
from sglang.kernels.ops.gemm.dense_bf16_gemm_sm100_splitk_epilogue import (
run_splitk_dense_gate,
run_splitk_dense_silu,
)
rows, k = hyper_input_normed.shape
lowrank = w_down.shape[0]
kpad = w_up_permuted_padded.shape[1]
out = torch.empty(
(rows, hidden_size),
dtype=hyper_input_normed.dtype,
device=hyper_input_normed.device,
)
if rows == 0:
return out
inv_hc = 1.0 / hc_count
t_pad = _get_scratch(lowrank, hyper_input_normed.dtype, hyper_input_normed.device)
for row_start in range(0, rows, _MAX_ROWS):
row_end = min(row_start + _MAX_ROWS, rows)
m = row_end - row_start
chunk = hyper_input_normed[row_start:row_end]
run_splitk_dense_silu(
chunk,
w_down.T,
t_pad[:m, :lowrank],
True,
_get_tactic(m, lowrank, k),
inv_hc,
)
run_splitk_dense_gate(
t_pad[:m],
w_up_permuted_padded.T,
chunk,
out[row_start:row_end],
True,
_get_tactic(m, k, kpad),
inv_hc,
hc_count,
)
return out
python/sglang/kernels/ops/gemm/dense_bf16_gemm_sm100_splitk_epilogue.py
+1293
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/ops/layernorm/grouped_gemma_rmsnorm.py
+77
View File
@@ -0,0 +1,77 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
if TYPE_CHECKING:
from tvm_ffi.module import Module
@cache_once
def _jit_grouped_gemma_rmsnorm_module(group_size: int, dtype: torch.dtype) -> Module:
"""Compile and cache the JIT grouped Gemma RMSNorm module."""
# Validation sits under `cache_once`: once per (group_size, dtype) specialisation,
# not once per call.
if dtype not in (torch.bfloat16, torch.float16):
raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
if group_size <= 0 or group_size % 512 != 0:
raise RuntimeError(
f"Unsupported group_size {group_size}. Must be a multiple of 512."
)
args = make_cpp_args(group_size, is_arch_support_pdl(), dtype)
return load_jit(
"grouped_gemma_rmsnorm",
*args,
cuda_files=["elementwise/grouped_gemma_rmsnorm.cuh"],
cuda_wrappers=[
("grouped_gemma_rmsnorm", f"GroupedGemmaRMSNormKernel<{args}>::run")
],
)
def grouped_gemma_rmsnorm(
input: torch.Tensor,
weight: torch.Tensor,
group_size: int,
eps: float = 1e-6,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Grouped Gemma-style RMSNorm: out = x * rsqrt(mean(x^2) + eps) * (1 + weight).
The last dimension is split into groups of `group_size` elements; variance
is computed per group. With group_size == input.size(-1) this reduces to a
plain Gemma RMSNorm.
Supported dtypes: torch.bfloat16, torch.float16.
Parameters
----------
input : CUDA tensor [..., hidden_size], hidden_size % group_size == 0
weight : CUDA tensor [hidden_size]
group_size : elements per variance group (multiple of 512)
eps : RMSNorm epsilon
out : optional pre-allocated output tensor (same shape/dtype as input)
Returns
-------
Normalized tensor, same shape/dtype as input.
"""
hidden_size = input.size(-1)
x = input.reshape(-1, hidden_size)
if out is None:
out = torch.empty_like(x)
else:
out = out.reshape(-1, hidden_size)
module = _jit_grouped_gemma_rmsnorm_module(group_size, input.dtype)
module.grouped_gemma_rmsnorm(x, weight, out, eps)
return out.reshape(input.shape)
python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py
+8 -4
View File
@@ -740,13 +740,15 @@ def _fused_commit_track_indices_kernel(
last_correct_out_ptr,
track_steps_out_ptr,
dtn,
accept_stride,
interval,
HAS_TRACK: tl.constexpr,
):
b = tl.program_id(0).to(tl.int64)
al = tl.load(accept_lens_ptr + b).to(tl.int64)
row = b * accept_stride
base = b * dtn
last = tl.load(accept_index_ptr + base + al - 1).to(tl.int64) - base
last = tl.load(accept_index_ptr + row + al - 1).to(tl.int64) - base
tl.store(last_correct_out_ptr + b, last)
if HAS_TRACK:
pre = tl.load(seq_lens_ptr + b).to(tl.int64)
@@ -755,7 +757,7 @@ def _fused_commit_track_indices_kernel(
tp = (post // interval) * interval
ti = tp - pre - 1
ti = tl.where(ti < 0, 0, ti)
cand = tl.load(accept_index_ptr + base + ti).to(tl.int64) - base
cand = tl.load(accept_index_ptr + row + ti).to(tl.int64) - base
tl.store(track_steps_out_ptr + b, tl.where(cross, cand, -1))
@@ -766,8 +768,9 @@ def fused_commit_track_indices(
draft_token_num: int,
mamba_track_interval: int,
):
"""Single-launch replacement for the eager index math in
``commit_mamba_states_after_verify`` (index ranges, gathers, floordiv chain)."""
"""Single-launch replacement for the eager verify-commit index math;
accept_index is [bs, tree_depth] but its values index bs * draft_token_num rows."""
accept_index = accept_index.contiguous()
bs = accept_lens.shape[0]
last_correct_step_indices = torch.empty(
bs, dtype=torch.int64, device=accept_lens.device
@@ -785,6 +788,7 @@ def fused_commit_track_indices(
last_correct_step_indices,
mamba_steps_to_track,
draft_token_num,
accept_index.shape[1],
mamba_track_interval,
HAS_TRACK=has_track,
)
python/sglang/kernels/ops/qwen4_ple.py
+309
View File
@@ -0,0 +1,309 @@
"""Bitwise-exact fused kernels for decode-sized Qwen4 PLE paths."""
from __future__ import annotations
import torch
import triton
import triton.language as tl
_QWEN4_NGRAM_SIZE = 3
_QWEN4_HEADS_PER_NGRAM = 8
_QWEN4_NGRAM_HEADS = 16
_QWEN4_HC_COUNT = 4
_QWEN4_HIDDEN_SIZE = 2560
_QWEN4_MAX_SHORT_CONV_STATE_LEN = 16
@triton.jit
def _round_bf16_to_fp32(value):
"""RNE-round fp32 to BF16 precision while retaining an fp32 register."""
bits = value.to(tl.int32, bitcast=True)
rounding_bias = 0x7FFF + ((bits >> 16) & 1)
rounded_bits = (bits + rounding_bias) & -65536
return rounded_bits.to(tl.float32, bitcast=True)
@triton.jit
def _qwen4_ngram_hash_kernel(
contexts_ptr,
multipliers_ptr,
vocab_sizes_ptr,
offsets_ptr,
output_ptr,
num_outputs,
eos_token_id,
NGRAM_SIZE: tl.constexpr,
HEADS_PER_NGRAM: tl.constexpr,
NGRAM_HEADS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
output_idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = output_idx < num_outputs
token_idx = output_idx // NGRAM_HEADS
head_idx = output_idx % NGRAM_HEADS
context_base = token_idx * NGRAM_SIZE
token_0 = tl.load(contexts_ptr + context_base, mask=mask, other=0)
token_1 = tl.load(contexts_ptr + context_base + 1, mask=mask, other=0)
token_2 = tl.load(contexts_ptr + context_base + 2, mask=mask, other=0)
multiplier_0 = tl.load(multipliers_ptr)
multiplier_1 = tl.load(multipliers_ptr + 1)
multiplier_2 = tl.load(multipliers_ptr + 2)
# Only the final position of each three-token context is materialized.
previous_1 = tl.where(token_1 == eos_token_id, eos_token_id, token_1)
previous_2 = tl.where(
(token_0 == eos_token_id) | (token_1 == eos_token_id),
eos_token_id,
token_0,
)
mixed = (token_2 * multiplier_0) ^ (previous_1 * multiplier_1)
mixed_3 = mixed ^ (previous_2 * multiplier_2)
mixed = tl.where(head_idx < HEADS_PER_NGRAM, mixed, mixed_3)
vocab_size = tl.load(vocab_sizes_ptr + head_idx, mask=mask, other=1)
offset = tl.load(offsets_ptr + head_idx, mask=mask, other=0)
tl.store(output_ptr + output_idx, mixed % vocab_size + offset, mask=mask)
def can_fuse_qwen4_ngram_hash(
contexts: torch.Tensor,
multipliers: torch.Tensor,
vocab_sizes: torch.Tensor,
offsets: torch.Tensor,
) -> bool:
"""Return whether inputs match the fixed Qwen4 PLE hash contract."""
return (
contexts.is_cuda
and contexts.dtype == torch.long
and contexts.dim() == 2
and contexts.shape[1] == _QWEN4_NGRAM_SIZE
and contexts.is_contiguous()
and multipliers.is_cuda
and multipliers.dtype == torch.long
and multipliers.numel() == _QWEN4_NGRAM_SIZE
and vocab_sizes.is_cuda
and vocab_sizes.dtype == torch.long
and vocab_sizes.numel() == _QWEN4_NGRAM_HEADS
and offsets.is_cuda
and offsets.dtype == torch.long
and offsets.numel() == _QWEN4_NGRAM_HEADS
)
def fused_qwen4_ngram_hash(
contexts: torch.Tensor,
multipliers: torch.Tensor,
vocab_sizes: torch.Tensor,
offsets: torch.Tensor,
eos_token_id: int,
) -> torch.Tensor:
"""Return the 16 Qwen4 PLE N-gram IDs in one kernel launch."""
if not can_fuse_qwen4_ngram_hash(contexts, multipliers, vocab_sizes, offsets):
raise ValueError("unsupported input for fused Qwen4 PLE N-gram hash")
output = torch.empty(
(contexts.shape[0], _QWEN4_NGRAM_HEADS),
dtype=torch.long,
device=contexts.device,
)
num_outputs = output.numel()
if num_outputs:
block_size = 256
_qwen4_ngram_hash_kernel[(triton.cdiv(num_outputs, block_size),)](
contexts,
multipliers,
vocab_sizes,
offsets,
output,
num_outputs,
eos_token_id,
NGRAM_SIZE=_QWEN4_NGRAM_SIZE,
HEADS_PER_NGRAM=_QWEN4_HEADS_PER_NGRAM,
NGRAM_HEADS=_QWEN4_NGRAM_HEADS,
BLOCK_SIZE=block_size,
num_warps=4,
)
return output
@triton.jit
def _qwen4_gate_value_kernel(
gate_ptr,
value_ptr,
output_ptr,
num_tokens,
HC_COUNT: tl.constexpr,
HIDDEN_SIZE: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
token_group = tl.program_id(0)
token = token_group // HC_COUNT
hidden = tl.arange(0, BLOCK_SIZE)
mask = (token < num_tokens) & (hidden < HIDDEN_SIZE)
# `gate` arrives already rounded to bf16 by the eager multiply/reduce/divide;
# every remaining eager bf16 rounding boundary is reproduced below.
gate = tl.load(gate_ptr + token_group).to(tl.float32)
magnitude = tl.maximum(tl.abs(gate), 1.0e-6)
root = _round_bf16_to_fp32(tl.sqrt(magnitude))
sign = tl.where(gate > 0.0, 1.0, tl.where(gate < 0.0, -1.0, 0.0))
transformed = _round_bf16_to_fp32(root * sign)
activated = _round_bf16_to_fp32(tl.sigmoid(transformed))
value = tl.load(value_ptr + token * HIDDEN_SIZE + hidden, mask=mask, other=0.0).to(
tl.float32
)
output_offset = token_group * HIDDEN_SIZE + hidden
tl.store(output_ptr + output_offset, activated * value, mask=mask)
def can_fuse_qwen4_gate_value(gate: torch.Tensor, value: torch.Tensor) -> bool:
"""Return whether inputs match Qwen4's fixed BF16 gate/value contract."""
return (
gate.is_cuda
and gate.dtype == torch.bfloat16
and gate.dim() == 3
and gate.shape[1:] == (_QWEN4_HC_COUNT, 1)
and gate.is_contiguous()
and value.is_cuda
and value.dtype == gate.dtype
and value.shape == (gate.shape[0], _QWEN4_HIDDEN_SIZE)
and value.is_contiguous()
)
def fused_qwen4_gate_value(gate: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
"""Apply Qwen4's post-reduction gate and value broadcast in one kernel."""
if not can_fuse_qwen4_gate_value(gate, value):
raise ValueError("unsupported input for fused Qwen4 PLE gate/value")
output = torch.empty(
(gate.shape[0], _QWEN4_HC_COUNT, _QWEN4_HIDDEN_SIZE),
dtype=value.dtype,
device=value.device,
)
if gate.shape[0]:
_qwen4_gate_value_kernel[(gate.shape[0] * _QWEN4_HC_COUNT,)](
gate,
value,
output,
gate.shape[0],
HC_COUNT=_QWEN4_HC_COUNT,
HIDDEN_SIZE=_QWEN4_HIDDEN_SIZE,
BLOCK_SIZE=4096,
num_warps=8,
)
return output
@triton.jit
def _qwen4_short_conv_state_kernel(
state_ptr,
state_indices_ptr,
x_ptr,
conv_input_ptr,
num_tokens,
CHANNELS: tl.constexpr,
STATE_LEN: tl.constexpr,
BLOCK_CHANNELS: tl.constexpr,
BLOCK_STATE_LEN: tl.constexpr,
):
token = tl.program_id(0)
channel = tl.program_id(1) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS)[:, None]
state_col = tl.arange(0, BLOCK_STATE_LEN)[None, :]
channel_mask = (token < num_tokens) & (channel < CHANNELS)
state_mask = channel_mask & (state_col < STATE_LEN)
state_index = tl.load(state_indices_ptr + token, mask=token < num_tokens, other=0)
state_base = state_index * CHANNELS * STATE_LEN
state_offset = state_base + channel * STATE_LEN + state_col
output_base = token * CHANNELS * (STATE_LEN + 1)
output_offset = output_base + channel * (STATE_LEN + 1) + state_col
# Materialize every old state value before advancing the in-place cache,
# so the convolution input equals the native index_select + cat result.
old_state = tl.load(state_ptr + state_offset, mask=state_mask, other=0.0)
tl.store(conv_input_ptr + output_offset, old_state, mask=state_mask)
x = tl.load(x_ptr + token * CHANNELS + channel, mask=channel_mask, other=0.0)
tl.store(
conv_input_ptr + output_base + channel * (STATE_LEN + 1) + STATE_LEN,
x,
mask=channel_mask,
)
tl.debug_barrier()
# Slot 0 is the CUDA-graph padding slot and may appear in several rows;
# its post-step value is unobservable, so skip it to avoid duplicate writers.
update_mask = state_mask & (state_col < STATE_LEN - 1) & (state_index != 0)
next_value = tl.load(
conv_input_ptr + output_offset + 1, mask=update_mask, other=0.0
)
tl.store(state_ptr + state_offset, next_value, mask=update_mask)
tl.store(
state_ptr + state_base + channel * STATE_LEN + STATE_LEN - 1,
x,
mask=channel_mask & (state_index != 0),
)
def can_fuse_qwen4_short_conv_state(
state: torch.Tensor,
state_indices: torch.Tensor,
x: torch.Tensor,
) -> bool:
"""Return whether decode state movement can use the exact fused kernel."""
return (
state.is_cuda
and state.dtype in (torch.bfloat16, torch.float16)
and state.dim() == 3
and state.is_contiguous()
and 0 < state.shape[2] <= _QWEN4_MAX_SHORT_CONV_STATE_LEN
and state_indices.is_cuda
and state_indices.dtype == torch.long
and state_indices.dim() == 1
and state_indices.is_contiguous()
and x.is_cuda
and x.dtype == state.dtype
and x.dim() == 2
and x.is_contiguous()
and x.shape == (state_indices.shape[0], state.shape[1])
)
def fused_qwen4_short_conv_state(
state: torch.Tensor,
state_indices: torch.Tensor,
x: torch.Tensor,
) -> torch.Tensor:
"""Build ``[selected state, x]`` and advance real decode slots in one launch."""
if not can_fuse_qwen4_short_conv_state(state, state_indices, x):
raise ValueError("unsupported input for fused Qwen4 short-conv state")
state_len = state.shape[2]
conv_input = torch.empty(
(x.shape[0], x.shape[1], state_len + 1),
dtype=x.dtype,
device=x.device,
)
if x.shape[0]:
block_channels = 128
block_state_len = triton.next_power_of_2(state_len)
_qwen4_short_conv_state_kernel[
(x.shape[0], triton.cdiv(x.shape[1], block_channels))
](
state,
state_indices,
x,
conv_input,
x.shape[0],
CHANNELS=state.shape[1],
STATE_LEN=state_len,
BLOCK_CHANNELS=block_channels,
BLOCK_STATE_LEN=block_state_len,
num_warps=8,
)
return conv_input
python/sglang/srt/arg_groups/choices.py
+1
View File
@@ -78,6 +78,7 @@ ATTENTION_BACKEND_CHOICES = [
"flex_attention",
"dsa",
"nsa", # Deprecated alias for "dsa"
"qsa",
"dsv4",
"compressed", # Deprecated alias for "dsv4"
# NVIDIA specific
python/sglang/srt/arg_groups/fields/exec_.py
+10
View File
@@ -847,6 +847,16 @@ class ExecOffload:
"Steps to prefetch in offloading.",
] = 1
offload_mode: A[str, "Mode of offloading."] = "cpu"
ple_offload_embedding: A[
Optional[bool],
Arg(
help="Offload Qwen4 PLE n-gram embedding weights to CPU pinned "
"memory. Enabled by default for BF16 Qwen4-Exp on CUDA; use "
"--no-ple-offload-embedding to disable.",
action=argparse.BooleanOptionalAction,
resolvable=True,
),
] = None
@dataclasses.dataclass
python/sglang/srt/arg_groups/memory_hook.py
+13
View File
@@ -24,6 +24,19 @@ logger = logging.getLogger(__name__)
_DEFAULT_PP_PREFILL_CUDA_GRAPH_MAX_TOKENS = 8192
def handle_offload_compatibility(server_args: Any) -> None:
"""Flag-only check; re-run after the model overrides fill in the PLE default."""
cfg = resolving_view(server_args)
if cfg.ple_offload_embedding and (
cfg.cpu_offload_gb > 0 or cfg.offload_group_size > 0
):
raise ValueError(
"--ple-offload-embedding cannot be combined with "
"--cpu-offload-gb or --offload-group-size: generic layer offload "
"would stage the pinned PLE embedding back to the device."
)
def handle_gpu_memory_settings(server_args: Any, gpu_mem):
"""
Configure GPU memory-dependent settings including
python/sglang/srt/arg_groups/model_hook.py
+1
View File
@@ -592,6 +592,7 @@ def handle_model_specific_adjustments(server_args: Any):
"Qwen3_5MoeForConditionalGeneration",
"InternS2PreviewForConditionalGeneration",
"Qwen3_5ForConditionalGeneration",
"Qwen4ExpForConditionalGeneration",
]:
# The quantization/moe_runner_backend resolution moved to the
# override registry (arg_groups/overrides.py:
python/sglang/srt/arg_groups/model_overrides/__init__.py
+1
View File
@@ -35,3 +35,4 @@ from sglang.srt.arg_groups.model_overrides import olmo2 # noqa: F401
from sglang.srt.arg_groups.model_overrides import qwen3_5 # noqa: F401
from sglang.srt.arg_groups.model_overrides import qwen3_moe # noqa: F401
from sglang.srt.arg_groups.model_overrides import qwen3_vl # noqa: F401
from sglang.srt.arg_groups.model_overrides import qwen4_exp # noqa: F401
python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py
+2 -1
View File
@@ -1,6 +1,6 @@
"""Config-time override declarations for qwen3_moe.
Architectures: InternS2PreviewForConditionalGeneration, Qwen3MoeForCausalLM, Qwen3NextForCausalLM, Qwen3VLMoeForConditionalGeneration, Qwen3_5ForConditionalGeneration, Qwen3_5MoeForConditionalGeneration.
Architectures: InternS2PreviewForConditionalGeneration, Qwen3MoeForCausalLM, Qwen3NextForCausalLM, Qwen3VLMoeForConditionalGeneration, Qwen3_5ForConditionalGeneration, Qwen3_5MoeForConditionalGeneration, Qwen4ExpForConditionalGeneration.
"""
import logging
@@ -23,6 +23,7 @@ logger = logging.getLogger(__name__)
"Qwen3_5MoeForConditionalGeneration",
"InternS2PreviewForConditionalGeneration",
"Qwen3_5ForConditionalGeneration",
"Qwen4ExpForConditionalGeneration",
)
def _qwen3_moe_family_overrides(server_args: Any, hf_config: Any) -> dict:
cfg = resolving_view(server_args)
python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py
+82
View File
@@ -0,0 +1,82 @@
"""Config-time override declarations for qwen4_exp.
Architectures: Qwen4ExpForConditionalGeneration.
"""
import logging
from typing import Any, Dict
from sglang.srt.arg_groups.model_override_base import (
_register_for,
get_default_attn_backend,
mamba_extra_buffer_of,
model_config_of,
resolved_view,
resolving_view,
use_mla_backend,
)
from sglang.srt.runtime_context import get_platform
logger = logging.getLogger(__name__)
@_register_for("Qwen4ExpForConditionalGeneration")
def _qwen4_exp_overrides(server_args: Any, hf_config: Any) -> dict:
"""Compressed QSA must own ``page_size`` here,
so the qwen3_5 hybrid attention-shape policy is restated rather than shared.
page_size=64 needs page-aligned full-KV allocation (slots are full_slot // ratio),
which MambaRadixCache allows only with mamba extra-buffer or --disable-radix-cache.
"""
cfg = resolving_view(server_args)
if cfg.disaggregation_mode != "null":
raise ValueError("Qwen4-Exp does not support PD disaggregation yet")
if cfg.enable_unified_memory:
raise ValueError("Qwen4-Exp does not support --enable-unified-memory yet")
overrides: Dict[str, Any] = {}
if cfg.ple_offload_embedding is None:
import torch
overrides["ple_offload_embedding"] = (
get_platform().is_cuda
and model_config_of(server_args).dtype == torch.bfloat16
)
text_config = getattr(hf_config, "text_config", hf_config)
if (
getattr(text_config, "num_experts", None) is not None
and cfg.moe_dense_tp_size == 1
):
overrides["moe_dense_tp_size"] = None
if get_platform().is_sm100 and cfg.attention_backend is None:
sm100_default_attn_backend = "triton"
default_attn_backend = get_default_attn_backend(
server_args,
use_mla_backend=use_mla_backend(server_args),
model_config=model_config_of(server_args),
)
if default_attn_backend == "trtllm_mha" and not (
not mamba_extra_buffer_of(resolved_view(server_args))
and not cfg.disable_radix_cache
and cfg.speculative_algorithm is None
):
sm100_default_attn_backend = "trtllm_mha"
overrides["attention_backend"] = sm100_default_attn_backend
overrides["page_size"] = 64 if sm100_default_attn_backend == "trtllm_mha" else 1
from sglang.srt.layers.attention.qsa.config import (
QSA_VARIANT_COMPRESSED,
parse_qsa_profile,
)
profile = parse_qsa_profile(hf_config)
if profile is not None and profile.variant == QSA_VARIANT_COMPRESSED:
# Compressed slot = full_slot // ratio; all backends need page-aligned pages.
# mamba_radix_cache_strategy resolves later, so do not gate on it.
overrides["page_size"] = 64
logger.info(
"Setting page size to 64 for compressed QSA "
"(full//ratio compressed addressing)."
)
return overrides
python/sglang/srt/arg_groups/overrides.py
+3
View File
@@ -519,6 +519,7 @@ _MAMBA_RADIX_CACHE_ARCHS = frozenset(
# Qwen3.8-2.4T-A95B ships as Qwen3_5MoeForCausalLM.
"Qwen3_5MoeForCausalLM",
"Qwen3_5ForCausalLM",
"Qwen4ExpForConditionalGeneration",
"MiniCPMV4_6ForConditionalGeneration",
"NemotronHForCausalLM",
"NemotronHPuzzleForCausalLM",
@@ -544,6 +545,7 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset(
"Qwen3_5MoeForCausalLM",
"Qwen3_5ForCausalLM",
"Qwen3NextForCausalLM",
"Qwen4ExpForConditionalGeneration",
"InternS2PreviewForConditionalGeneration",
"MiniCPMV4_6ForConditionalGeneration",
"BailingMoeV2_5ForCausalLM",
@@ -1032,6 +1034,7 @@ _FLASHINFER_ALLREDUCE_FUSION_ARCHS = frozenset(
"Qwen3MoeForCausalLM",
"Qwen3VLMoeForConditionalGeneration",
"Qwen3NextForCausalLM",
"Qwen4ExpForConditionalGeneration",
"KimiK25ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
"InternS2PreviewForConditionalGeneration",
python/sglang/srt/arg_groups/pipeline.py
+5
View File
@@ -91,6 +91,9 @@ def run_resolution_pipeline(server_args: Any) -> None:
)
handle_hicache_ratio_default(server_args)
from sglang.srt.arg_groups.memory_hook import handle_offload_compatibility
handle_offload_compatibility(server_args)
from sglang.srt.arg_groups.validation_hook import (
validate_experimental_sgl_marlin,
validate_prefill_decode_interval,
@@ -231,6 +234,8 @@ def run_resolution_pipeline(server_args: Any) -> None:
)
handle_model_specific_adjustments(server_args)
# After the model overrides: Qwen4-Exp declares the PLE offload default there.
handle_offload_compatibility(server_args)
# Set kernel backends.
run_post_process_pass(server_args, _sampling_backend_default)
python/sglang/srt/arg_groups/speculative_hook.py
+2
View File
@@ -863,6 +863,8 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
"PixtralForConditionalGeneration",
"HYV3ForCausalLM",
"HYV4ForCausalLM",
# Qwen4-Exp ships its NEXTN draft layer inside the target checkpoint.
"Qwen4ExpForConditionalGeneration",
]:
if cfg.speculative_draft_model_path is None:
declare_resolution(
python/sglang/srt/configs/__init__.py
+3
View File
@@ -68,6 +68,7 @@ from sglang.srt.configs.qwen3_5 import (
)
from sglang.srt.configs.qwen3_asr import Qwen3ASRConfig
from sglang.srt.configs.qwen3_next import Qwen3NextConfig
from sglang.srt.configs.qwen4_exp import Qwen4ExpConfig, Qwen4ExpTextConfig
from sglang.srt.configs.spark2_5 import Spark2_5Config
from sglang.srt.configs.step3_vl import (
Step3TextConfig,
@@ -110,6 +111,8 @@ __all__ = [
"KimiK25Config",
"LagunaConfig",
"Qwen3NextConfig",
"Qwen4ExpConfig",
"Qwen4ExpTextConfig",
"Qwen3_5Config",
"Qwen3_5MoeConfig",
"Qwen3_5TextConfig",
python/sglang/srt/configs/model_config.py
+29 -2
View File
@@ -212,12 +212,19 @@ def is_deepseek_v4(config) -> bool:
)
def is_qwen4_exp(config) -> bool:
return _hf_arch(config) in (
"Qwen4ExpForConditionalGeneration",
"Qwen4ExpForCausalLMMTP",
)
def resolve_spec_hidden_size(
hf_config, hidden_size: int, hc_mult: int
) -> tuple[int, Optional[int]]:
# Only DSV4 carries the hc-flattened stream across the target→draft
# DSV4 and Qwen4-Exp carry the hc-flattened stream across the target->draft
# boundary; other hc models (hy_v4) collapse to hidden_size first.
if hc_mult <= 1 or not is_deepseek_v4(hf_config):
if hc_mult <= 1 or not (is_deepseek_v4(hf_config) or is_qwen4_exp(hf_config)):
return hidden_size, None
hc_hidden_size = hidden_size * hc_mult
return hc_hidden_size, hc_hidden_size
@@ -836,6 +843,23 @@ class ModelConfig:
self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP"
self.hf_config.num_nextn_predict_layers = 1
if (
is_draft_model
and self.hf_config.architectures[0] == "Qwen4ExpForConditionalGeneration"
):
# The target's ModelConfig shares this hf_config object; deep-copy
# before the MTP rewrites below so the target keeps its full depth.
self.hf_config = copy.deepcopy(self.hf_config)
self.hf_text_config = get_hf_text_config(self.hf_config)
self.hf_config.architectures[0] = "Qwen4ExpForCausalLMMTP"
text_config = self.hf_text_config
text_config.num_nextn_predict_layers = 1
# layers_block_type follows layer_types, not num_hidden_layers,
# so both must shrink for the draft's full_attention_layer_ids to be [0].
text_config.num_hidden_layers = 1
text_config.layer_types = ["full_attention"]
text_config.full_attention_interval = 1
if is_draft_model and self.hf_config.architectures[0] == "Qwen3MoeForCausalLM":
self.hf_config.architectures[0] = "Qwen3MoeForCausalLMMTP"
self.hf_config.num_nextn_predict_layers = 1
@@ -2056,6 +2080,7 @@ multimodal_model_archs = [
"Qwen3VLMoeForConditionalGeneration",
"Qwen3_5ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
"Qwen4ExpForConditionalGeneration",
"InternS2PreviewForConditionalGeneration",
"InternS2MobiusForConditionalGeneration",
"Qwen3ASRForConditionalGeneration",
@@ -2118,6 +2143,8 @@ multimodal_breakable_cuda_graph_supported_model_archs = [
"PaddleOCRVLForConditionalGeneration",
"Qwen3_5ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
# Qwen4-Exp is intentionally absent: QSA builds host-side sparse metadata
# per forward and cannot serve the breakable prefill capture.
"MuseGlimmerForConditionalGeneration",
"KimiK3ForConditionalGeneration",
"KimiK25ForConditionalGeneration",
python/sglang/srt/configs/qwen4_exp.py
+165
View File
@@ -0,0 +1,165 @@
from transformers import PretrainedConfig
from sglang.srt.configs.qwen3_next import Qwen3NextConfig
from sglang.srt.configs.qwen3_vl import Qwen3VLVisionConfig
class Qwen4ExpVisionConfig(Qwen3VLVisionConfig):
model_type = "qwen4_exp"
base_config_key = "vision_config"
def __init__(self, **kwargs):
super().__init__(**kwargs)
class Qwen4ExpTextConfig(Qwen3NextConfig):
model_type = "qwen4_exp_text"
base_config_key = "text_config"
keys_to_ignore_at_inference = ["past_key_values"]
# ModelConfig sizes the speculative hidden width off the DSV4 mHC name.
attribute_map = {"hc_mult": "hc_count"}
def __init__(
self,
hc_count=4,
hc_lowrank=320,
ple_layer_ids=None,
ple_embed_dim=None,
ple_conv_kernel_size=4,
ngram_size=3,
heads_per_ngram=8,
ngram_vocab_size_base=20000000,
make_ngram_vocab_size_divisible_by=128,
ple_offload_embedding=False,
ple_embedding_dtype=None,
index_share_for_mtp_iteration=True,
rope_parameters=None,
layer_types=None,
**kwargs,
):
if hc_count <= 1:
raise ValueError(f"Qwen4-Exp requires hc_count > 1, got {hc_count}.")
# Newer checkpoints spell rope_scaling/rope_theta as rope_parameters;
# Qwen3NextConfig.__init__ only reads the old names.
if rope_parameters is not None:
if kwargs.get("rope_scaling") is None:
kwargs["rope_scaling"] = rope_parameters
if kwargs.get("rope_theta") is None and "rope_theta" in rope_parameters:
kwargs["rope_theta"] = rope_parameters["rope_theta"]
if (
kwargs.get("partial_rotary_factor") is None
and "partial_rotary_factor" in rope_parameters
):
kwargs["partial_rotary_factor"] = rope_parameters[
"partial_rotary_factor"
]
super().__init__(**kwargs)
if self.rope_scaling is None:
self.rope_scaling = rope_parameters or {}
self.rope_parameters = rope_parameters or self.rope_scaling
self.hc_count = hc_count
self.hc_lowrank = hc_lowrank
self.layer_types = layer_types
self.ple_layer_ids = ple_layer_ids or []
self.ple_embed_dim = ple_embed_dim or self.hidden_size
self.ple_conv_kernel_size = ple_conv_kernel_size
self.ngram_size = ngram_size
self.heads_per_ngram = heads_per_ngram
self.ngram_vocab_size_base = ngram_vocab_size_base
self.make_ngram_vocab_size_divisible_by = make_ngram_vocab_size_divisible_by
self.ple_offload_embedding = ple_offload_embedding
# "float8_e4m3fn" keeps fp8 PLE tables fp8-resident; text_config-scoped.
self.ple_embedding_dtype = ple_embedding_dtype
# Draft decode steps reuse the draft-extend indexer top-k (IndexShare).
self.index_share_for_mtp_iteration = index_share_for_mtp_iteration
@property
def layers_block_type(self):
if self.layer_types is not None:
return [
(
"attention"
if layer_type in ("full_attention", "qwen_sparse_attention")
else layer_type
)
for layer_type in self.layer_types
]
return super().layers_block_type
@property
def short_conv_layer_ids(self):
if not self.ple_layer_ids:
return []
return sorted({int(layer_id) - 1 for layer_id in self.ple_layer_ids})
@property
def short_conv_state_shape(self):
if not self.short_conv_layer_ids:
return None
ple_state_len = (self.ple_conv_kernel_size - 1) * self.ngram_size
ple_channels = self.hidden_size * self.hc_count
return ple_channels, ple_state_len
@property
def ngram_context_len(self):
if not self.ple_layer_ids:
return 0
return max(int(self.ngram_size) - 1, 0)
class Qwen4ExpConfig(PretrainedConfig):
model_type = "qwen4_exp"
sub_configs = {
"vision_config": Qwen4ExpVisionConfig,
"text_config": Qwen4ExpTextConfig,
}
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
text_config=None,
vision_config=None,
image_token_id=248056,
video_token_id=248057,
vision_start_token_id=248053,
vision_end_token_id=248054,
tie_word_embeddings=False,
rope_parameters=None,
**kwargs,
):
# The nested text config is authoritative; old exports also copied this
# value to the top level.
if text_config is not None:
kwargs.pop("split_ngram_parts", None)
# Backward compatibility: older Qwen4-Exp checkpoints were text-only
# and stored text attributes at the top level.
text_kwargs = (
dict(kwargs)
if text_config is None
and "hidden_size" in kwargs
and "num_hidden_layers" in kwargs
else {}
)
if isinstance(vision_config, dict):
self.vision_config = self.sub_configs["vision_config"](**vision_config)
elif vision_config is None:
self.vision_config = self.sub_configs["vision_config"]()
else:
self.vision_config = vision_config
if isinstance(text_config, dict):
self.text_config = self.sub_configs["text_config"](**text_config)
elif text_config is None:
self.text_config = self.sub_configs["text_config"](**text_kwargs)
else:
self.text_config = text_config
self.image_token_id = image_token_id
self.video_token_id = video_token_id
self.vision_start_token_id = vision_start_token_id
self.vision_end_token_id = vision_end_token_id
self.rope_parameters = rope_parameters or getattr(
self.text_config, "rope_parameters", {}
)
super().__init__(**kwargs, tie_word_embeddings=tie_word_embeddings)
python/sglang/srt/configs/update_config.py
+1
View File
@@ -308,6 +308,7 @@ def adjust_config_with_unaligned_cpu_tp(
[model_config.hf_config, "vision_config", "qwen3_vl", "num_heads"],
[model_config.hf_config, "vision_config", "qwen3_5_moe", "num_heads"],
[model_config.hf_config, "vision_config", "qwen3_5", "num_heads"],
[model_config.hf_config, "vision_config", "qwen4_exp", "num_heads"],
[model_config.hf_config, "vision_config", "mllama", "attention_heads"],
[
model_config.hf_config,
python/sglang/srt/disaggregation/utils.py
+5 -1
View File
@@ -19,7 +19,7 @@ import numpy as np
import torch
import torch.distributed as dist
from sglang.srt.configs.model_config import get_dsa_mtp_topk_width
from sglang.srt.configs.model_config import get_dsa_mtp_topk_width, is_deepseek_dsa
from sglang.srt.disaggregation.base import KVPoll
from sglang.srt.environ import envs
from sglang.srt.runtime_context import (
@@ -71,6 +71,10 @@ def get_dsa_seed_metadata_dim(hf_config) -> int:
"""Return the model-defined PD seed width, independent of local spec mode."""
if not getattr(hf_config, "index_share_for_mtp_iteration", False):
return 0
# QSA models reuse the same flag for their draft-side index sharing but
# carry no DSA seed metadata over PD.
if not is_deepseek_dsa(hf_config):
return 0
return get_dsa_mtp_topk_width(hf_config)
python/sglang/srt/environ.py
+6
View File
@@ -298,6 +298,12 @@ class Envs:
# keeping access relatively ordered.
SGLANG_SORT_WEIGHT_FILES = EnvInt(0)
SGLANG_DISABLED_MODEL_ARCHS = EnvTuple(tuple())
# Shard the Qwen4-Exp PLE n-gram embedding within each attention-TP group
# instead of gathering DP tokens for a global-TP lookup.
SGLANG_USE_ATTN_TP_NGRAM = EnvBool(False)
# Bitwise-exact, shape-guarded Qwen4 PLE decode fusion. Unsupported inputs
# and phases fall back to the original implementation.
SGLANG_ENABLE_QWEN4_PLE_FUSION = EnvBool(True)
SGLANG_PREFETCH_BLOCK_SIZE_MB = EnvInt(16)
SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION = EnvBool(False)
SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False)
python/sglang/srt/layers/attention/attention_registry.py
+18
View File
@@ -143,6 +143,15 @@ def create_dsa_backend(runner):
return DeepseekSparseAttnBackend(runner)
@register_attention_backend("qsa")
def create_qsa_backend(runner):
from sglang.srt.layers.attention.qwen_sparse_attn_backend import (
QwenSparseAttnBackend,
)
return QwenSparseAttnBackend(runner)
@register_attention_backend("nsa")
def _create_nsa_compat(runner):
warnings.warn(
@@ -454,6 +463,15 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
)
logger.info(f"Using hybrid linear attention backend for hybrid GDN models.")
linear_attn_backend = GDNAttnBackend(runner)
from sglang.srt.layers.attention.qsa.config import is_qwen_qsa
if is_qwen_qsa(runner.model_config.hf_config):
from sglang.srt.layers.attention.qwen_sparse_attn_backend import (
QwenSparseAttnBackend,
)
logger.info("Using QSA for sparse full-attention layers.")
full_attn_backend = QwenSparseAttnBackend(runner)
elif mamba2_config(runner.model_config) is not None:
from sglang.srt.configs.lfm2 import Lfm2Config
from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+105 -10
View File
@@ -11,6 +11,7 @@ from sglang.kernels.ops.mamba.mamba_state_indices_triton import (
)
from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
fused_conv_window_scatter_with_mask,
fused_mamba_state_scatter_with_mask,
scatter_mamba_states_after_mtp_verify,
track_mamba_states_all_layers,
track_mamba_states_if_needed,
@@ -32,12 +33,7 @@ from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import (
get_exec,
get_memory,
get_spec,
mamba_cache_chunk_size,
)
from sglang.srt.runtime_context import get_exec, get_memory, get_spec
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.spec_info import SpecInput
@@ -339,11 +335,13 @@ class MambaAttnBackendBase(AttentionBackend):
the last complete chunk boundary, mamba_track_mask rows only)."""
conv_state_len = self.conv_states_shape[-1]
lens_to_track = (
forward_batch.mamba_track_seqlens - forward_batch.extend_prefix_lens
# Shared with the Qwen4-Exp PLE side states so the boundary can never
# drift between them.
aligned_len = forward_batch.mamba_track_aligned_lens()
assert aligned_len is not None, (
"conv-state tracking requires mamba_track_seqlens and extend_prefix_lens; "
"this path should only run when the track mask is set on an extend batch"
)
chunk_size = mamba_cache_chunk_size()
aligned_len = (lens_to_track // chunk_size) * chunk_size
start_indices = query_start_loc[:-1] + aligned_len - conv_state_len
start_indices = start_indices[forward_batch.mamba_track_mask]
@@ -1109,6 +1107,11 @@ class HybridLinearAttnBackend(AttentionBackend):
for attn_backend in self.attn_backend_list:
attn_backend.init_forward_metadata_in_graph(forward_batch)
def get_indexer_metadata(self, layer_id: int, forward_batch: ForwardBatch):
if layer_id in self.full_attn_layers:
return self.full_attn_backend.get_indexer_metadata(layer_id, forward_batch)
return None
def on_after_cuda_graph_warmup(self):
for attn_backend in self.attn_backend_list:
attn_backend.on_after_cuda_graph_warmup()
@@ -1382,6 +1385,98 @@ class HybridLinearAttnBackend(AttentionBackend):
mamba_steps_to_track,
)
self._update_ple_state_after_mtp_verify(
state_indices_tensor,
last_correct_step_indices,
mamba_track_indices,
mamba_steps_to_track,
)
@staticmethod
def _scatter_speculative_state_with_mask(
dst: torch.Tensor,
src: torch.Tensor,
dst_indices_raw: torch.Tensor,
step_indices_raw: torch.Tensor,
):
if dst is None or src is None or step_indices_raw.numel() == 0:
return
if dst.is_cuda and src.is_cuda:
fused_mamba_state_scatter_with_mask(
dst, src, dst_indices_raw, step_indices_raw
)
return
device = dst.device
dst_indices = dst_indices_raw.to(device=device, dtype=torch.long)
steps = step_indices_raw.to(device=device, dtype=torch.long)
src_indices = torch.arange(steps.shape[0], device=device, dtype=torch.long)
valid = (
(steps >= 0)
& (steps < src.shape[2])
& (dst_indices >= 0)
& (dst_indices < dst.shape[1])
& (src_indices < src.shape[1])
)
valid_indices = valid.nonzero(as_tuple=True)[0]
if valid_indices.numel() == 0:
return
dst[:, dst_indices[valid_indices]] = src[
:, src_indices[valid_indices], steps[valid_indices]
]
def _update_ple_state_after_mtp_verify(
self,
state_indices_tensor: torch.Tensor,
last_correct_step_indices: torch.Tensor,
mamba_track_indices: Optional[torch.Tensor],
mamba_steps_to_track: Optional[torch.Tensor],
):
"""Roll the accepted per-step PLE side states into their main slots."""
req_to_token_pool = self.linear_attn_backend.req_to_token_pool
if mamba_track_indices is not None:
assert mamba_steps_to_track is not None
state_pairs = []
short_conv_pool = req_to_token_pool.short_conv_pool
if (
short_conv_pool.conv_state is not None
and short_conv_pool.intermediate_conv_state is not None
):
state_pairs.append(
(
short_conv_pool.conv_state,
short_conv_pool.intermediate_conv_state,
)
)
ngram_pool = req_to_token_pool.ngram_pool
if (
ngram_pool.context is not None
and ngram_pool.intermediate_context is not None
):
state_pairs.append(
(
ngram_pool.context.unsqueeze(0),
ngram_pool.intermediate_context.unsqueeze(0),
)
)
for state, intermediate_state in state_pairs:
self._scatter_speculative_state_with_mask(
state,
intermediate_state,
state_indices_tensor,
last_correct_step_indices,
)
if mamba_track_indices is not None:
self._scatter_speculative_state_with_mask(
state,
intermediate_state,
mamba_track_indices,
mamba_steps_to_track,
)
class ShortConvHybridAttnBackend(HybridLinearAttnBackend):
"""HybridLinearAttnBackend variant for short-conv hybrid models (ZAYA1 CCA,
python/sglang/srt/layers/attention/qsa/__init__.py
+40
View File
@@ -0,0 +1,40 @@
"""Simple QSA operators for Qwen4-Exp.
The package intentionally avoids eager imports so reference tensor helpers can
be used without constructing the full SGLang runtime.
"""
__all__ = [
"QSAIndexer",
"QSAIndexerMetadata",
"QSAProfile",
"QwenDSAIndexer",
"build_qsa_indexer",
"get_qsa_indexer_metadata",
"is_qwen_qsa",
"parse_qsa_profile",
]
def __getattr__(name):
if name == "QSAIndexer":
from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer
return QSAIndexer
if name == "QwenDSAIndexer":
from sglang.srt.layers.attention.qsa.dsa_indexer import QwenDSAIndexer
return QwenDSAIndexer
if name == "QSAIndexerMetadata":
from sglang.srt.layers.attention.qsa.metadata import QSAIndexerMetadata
return QSAIndexerMetadata
if name in {"QSAProfile", "is_qwen_qsa", "parse_qsa_profile"}:
from sglang.srt.layers.attention.qsa import config as qsa_config
return getattr(qsa_config, name)
if name in {"build_qsa_indexer", "get_qsa_indexer_metadata"}:
from sglang.srt.layers.attention.qsa import glue as qsa_glue
return getattr(qsa_glue, name)
raise AttributeError(name)
python/sglang/srt/layers/attention/qsa/config.py
+177
View File
@@ -0,0 +1,177 @@
"""Shared QSA profile parsing across model variants.
``QSAProfile`` normalizes each model family's HF-config indexer schema,
so backends, draft utilities and model glue branch on a stable variant name,
not on raw config keys. ``compressed`` is Qwen4-Exp block compression;
``tokenwise`` is qsa_0511 / Qwen3.5-DSA per-token indexing.
DeepSeek NSA configs also expose ``index_topk``,
so the tokenwise schema is additionally gated on a Qwen ``model_type``.
"""
from __future__ import annotations
from typing import Optional
import msgspec
# QSA variant names.
QSA_VARIANT_COMPRESSED = "compressed"
QSA_VARIANT_TOKENWISE = "tokenwise"
# Rotary layouts the indexer can consume.
QSA_ROPE_MROPE = "mrope"
QSA_ROPE_PLAIN = "plain"
_COMPRESSED_FIELDS = (
"indexer_n_heads",
"indexer_kv_heads",
"indexer_head_dim",
"indexer_budget",
"indexer_compress_ratio",
)
_TOKENWISE_FIELDS = (
"index_topk",
"index_n_heads",
"index_kv_heads",
"index_head_dim",
)
# fast_topk_v2 only supports these compressed block top-k widths.
_COMPRESSED_BLOCK_TOPK = frozenset({512, 2048})
# fast_topk_v2 only supports a 2048-wide tokenwise top-k.
_TOKENWISE_BUDGET = 2048
class QSAProfile(msgspec.Struct, frozen=True):
"""Normalized sparse-attention indexer description for one model."""
variant: str # QSA_VARIANT_COMPRESSED | QSA_VARIANT_TOKENWISE
n_heads: int # index query heads
kv_heads: int # index key/value heads
head_dim: int # per-head index dimension
budget: int # tokens selected per query row
compress_ratio: int # 1 for tokenwise variants
rope_mode: str # rotary layout the indexer expects
@property
def block_topk(self) -> int:
"""Compressed blocks selected per query row (== budget for tokenwise)."""
return self.budget // self.compress_ratio
def _text_config(config):
return getattr(config, "text_config", config)
def _is_qwen_family(config) -> bool:
model_type = str(getattr(config, "model_type", "") or "")
return model_type.startswith("qwen")
def _require_fields(config, fields) -> dict:
missing = [name for name in fields if getattr(config, name, None) is None]
if missing:
raise ValueError(f"QSA config is missing required fields: {missing}")
return {name: int(getattr(config, name)) for name in fields}
def _parse_compressed(text_config) -> QSAProfile:
values = _require_fields(text_config, _COMPRESSED_FIELDS)
if any(value <= 0 for value in values.values()):
raise ValueError(f"QSA config values must be positive: {values}")
if values["indexer_kv_heads"] != 1:
raise ValueError("the QSA MQA operators require indexer_kv_heads=1")
ratio = values["indexer_compress_ratio"]
budget = values["indexer_budget"]
if ratio < 2:
# Padding rows carry logical length 1, which must never reach a
# compression boundary; ratio >= 2 guarantees that.
raise ValueError(f"QSA requires indexer_compress_ratio >= 2, got {ratio}")
if budget % ratio != 0:
raise ValueError(
"indexer_budget must be divisible by indexer_compress_ratio, got "
f"{budget} / {ratio}"
)
if budget // ratio not in _COMPRESSED_BLOCK_TOPK:
raise ValueError(
"fast_topk_v2 requires indexer_budget / indexer_compress_ratio "
f"to be one of {sorted(_COMPRESSED_BLOCK_TOPK)}, got {budget // ratio}"
)
return QSAProfile(
variant=QSA_VARIANT_COMPRESSED,
n_heads=values["indexer_n_heads"],
kv_heads=values["indexer_kv_heads"],
head_dim=values["indexer_head_dim"],
budget=budget,
compress_ratio=ratio,
# The compressed indexer consumes the Qwen4-Exp layer's own (m)rope.
rope_mode=QSA_ROPE_MROPE,
)
def _parse_tokenwise(text_config) -> QSAProfile:
values = _require_fields(text_config, _TOKENWISE_FIELDS)
if any(value <= 0 for value in values.values()):
raise ValueError(f"QSA config values must be positive: {values}")
if values["index_topk"] != _TOKENWISE_BUDGET:
raise ValueError(
f"fast_topk_v2 only supports index_topk = {_TOKENWISE_BUDGET}, "
f"got {values['index_topk']}"
)
if values["index_kv_heads"] != 1:
raise ValueError(
f"QSA tokenwise index requires index_kv_heads = 1 (MQA), "
f"got {values['index_kv_heads']}"
)
return QSAProfile(
variant=QSA_VARIANT_TOKENWISE,
n_heads=values["index_n_heads"],
kv_heads=values["index_kv_heads"],
head_dim=values["index_head_dim"],
budget=values["index_topk"],
compress_ratio=1,
# The tokenwise indexer owns plain per-token rotary positions.
rope_mode=QSA_ROPE_PLAIN,
)
def parse_qsa_profile(config) -> Optional[QSAProfile]:
"""QSA profile of config, None if absent; malformed schemas raise ValueError."""
if config is None:
return None
text_config = _text_config(config)
if text_config is None:
return None
has_compressed = getattr(text_config, "indexer_n_heads", None) is not None
has_tokenwise = getattr(
text_config, "index_topk", None
) is not None and _is_qwen_family(text_config)
if has_compressed and has_tokenwise:
raise ValueError(
"Ambiguous QSA config: both compressed (indexer_*) and tokenwise "
"(index_*) indexer fields are set"
)
if has_compressed:
return _parse_compressed(text_config)
if has_tokenwise:
return _parse_tokenwise(text_config)
return None
def is_qwen_qsa(config) -> bool:
"""Return whether the config describes a supported Qwen QSA variant."""
return parse_qsa_profile(config) is not None
__all__ = [
"QSAProfile",
"QSA_ROPE_MROPE",
"QSA_ROPE_PLAIN",
"QSA_VARIANT_COMPRESSED",
"QSA_VARIANT_TOKENWISE",
"is_qwen_qsa",
"parse_qsa_profile",
]
python/sglang/srt/layers/attention/qsa/dsa_indexer.py
+337
View File
@@ -0,0 +1,337 @@
"""Tokenwise (per-token) QSA indexer for Qwen3Next-DSA models.
A tokenwise profile has ``compress_ratio = 1`` and ``block_topk = budget = 2048``;
it never consumes the compressed-only MQA inputs.
Only the BF16 torch reference path is implemented;
requesting the FP8 or TileLang fast paths fails loudly.
"""
from __future__ import annotations
import logging
import torch
from sglang.srt.layers.attention.qsa.config import (
QSA_VARIANT_TOKENWISE,
parse_qsa_profile,
)
from sglang.srt.layers.attention.qsa.kernel import qsa_fast_topk
from sglang.srt.layers.attention.qsa.qsa_indexer import _qsa_prefill_row_chunk_size
from sglang.srt.layers.layernorm import GemmaRMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils import MultiPlatformOp
logger = logging.getLogger(__name__)
def torch_dsa_weighted_mqa_logits(
q: torch.Tensor,
w: torch.Tensor,
k: torch.Tensor,
score_scale: float,
) -> torch.Tensor:
"""Lightning-Index scoring reference: ReLU dot-product weighted per head."""
if k.ndim == 4:
if k.shape[2] != 1 or k.shape[0] != q.shape[0]:
raise ValueError(
"tokenwise MQA requires per-row k [rows, keys, 1, head_dim], "
f"got {k.shape}"
)
scores = torch.relu(torch.einsum("mhd,mkhd->mkh", q.float(), k.float()))
else:
if k.ndim != 3 or k.shape[1] != 1:
raise ValueError(
f"tokenwise MQA requires k [keys, 1, head_dim], got {k.shape}"
)
scores = torch.relu(torch.einsum("mhd,khd->mkh", q.float(), k.float()))
return (scores * w.float().unsqueeze(1)).sum(dim=-1) / score_scale
class QwenDSAIndexer(MultiPlatformOp):
"""Tokenwise Lightning Indexer with the compressed ``QSAIndexer`` forward contract;
returns per-row logical token indices consumed as ``topk_indices``."""
def __init__(
self,
config,
layer_id: int,
quant_config=None,
prefix: str = "",
page_size: int = 64,
max_model_len=None,
) -> None:
super().__init__()
profile = parse_qsa_profile(config)
if profile is None or profile.variant != QSA_VARIANT_TOKENWISE:
raise ValueError(
"QwenDSAIndexer requires a tokenwise QSA config (index_topk/), "
f"got profile={profile}"
)
if page_size != 64:
# The paged index-K layout and every fast path assume 64-token
# pages, matching qsa_0511.
raise ValueError(f"tokenwise QSA requires page_size = 64, got {page_size}")
self.qsa_profile = profile
self.layer_id = int(layer_id)
self.index_n_heads = profile.n_heads
self.index_kv_heads = profile.kv_heads
self.index_head_dim = profile.head_dim
self.token_topk = profile.budget
self.score_scale = float(profile.head_dim) ** 0.5
self.page_size = page_size
self.max_model_len = max_model_len
# Fused Q/K/W projection. Output layout:
# q_raw: [M, index_n_heads * index_head_dim]
# k_raw: [M, index_kv_heads * index_head_dim]
# w: [M, index_n_heads] per-head scalar weight
self.index_q_dim = self.index_n_heads * self.index_head_dim
self.index_k_dim = self.index_kv_heads * self.index_head_dim
self.index_w_dim = self.index_n_heads
self.index_qkw_proj = ReplicatedLinear(
config.hidden_size,
self.index_q_dim + self.index_k_dim + self.index_w_dim,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.index_qkw_proj" if prefix else "index_qkw_proj",
)
self.index_q_layernorm = GemmaRMSNorm(
self.index_head_dim, eps=getattr(config, "rms_norm_eps", 1e-6)
)
self.index_k_layernorm = GemmaRMSNorm(
self.index_head_dim, eps=getattr(config, "rms_norm_eps", 1e-6)
)
# The indexer keeps its own RoPE instance shaped for index_head_dim;
# its rotary width follows the main attention's partial_rotary_factor.
rope_scaling = getattr(config, "rope_scaling", None)
if rope_scaling is None:
rope_scaling = getattr(config, "rope_parameters", None)
rope_theta = getattr(config, "rope_theta", 10000)
if isinstance(rope_scaling, dict) and "rope_theta" in rope_scaling:
rope_theta = rope_scaling["rope_theta"]
main_head_dim = getattr(config, "head_dim", None)
if main_head_dim is None:
main_head_dim = getattr(config, "hidden_size") // getattr(
config, "num_attention_heads"
)
partial_rotary_factor = getattr(config, "partial_rotary_factor", None)
if partial_rotary_factor is None and isinstance(rope_scaling, dict):
partial_rotary_factor = rope_scaling.get("partial_rotary_factor")
if partial_rotary_factor is None:
partial_rotary_factor = 1.0
indexer_rotary_dim = min(
self.index_head_dim, int(main_head_dim * float(partial_rotary_factor))
)
if indexer_rotary_dim <= 0 or indexer_rotary_dim % 2 != 0:
raise ValueError(
"tokenwise QSA indexer requires a positive even rotary dim, got "
f"{indexer_rotary_dim=} from {main_head_dim=} and "
f"{partial_rotary_factor=}"
)
self.rotary_emb = get_rope_wrapper(
head_size=self.index_head_dim,
rotary_dim=indexer_rotary_dim,
max_position=getattr(config, "max_position_embeddings", 8192),
base=rope_theta,
rope_scaling=rope_scaling if isinstance(rope_scaling, dict) else None,
is_neox_style=True,
dtype=torch.get_default_dtype(),
)
def project_qkw(self, hidden_states: torch.Tensor, positions: torch.Tensor):
"""Fused Q/K/W projection, per-head RMS norm and indexer RoPE."""
qkw, _ = self.index_qkw_proj(hidden_states)
q_raw, k_raw, w = torch.split(
qkw, [self.index_q_dim, self.index_k_dim, self.index_w_dim], dim=-1
)
q = self.index_q_layernorm(q_raw.reshape(-1, self.index_head_dim)).reshape(
-1, self.index_n_heads, self.index_head_dim
)
k = self.index_k_layernorm(k_raw.reshape(-1, self.index_head_dim)).reshape(
-1, self.index_kv_heads, self.index_head_dim
)
q, k = self.rotary_emb(positions, q, k)
return q, w, k
def forward_cuda(
self,
hidden_states: torch.Tensor,
positions: torch.Tensor,
forward_batch,
indexer_metadata,
) -> torch.Tensor:
forward_mode = forward_batch.forward_mode
is_target_verify = getattr(forward_mode, "is_target_verify", lambda: False)()
is_draft_extend = getattr(forward_mode, "is_draft_extend", lambda **_: False)(
include_v2=True
)
is_paged = forward_mode.is_decode() or is_target_verify or is_draft_extend
if is_paged:
# Paged rows take their causal length from the paged metadata,
# not the model's RoPE coordinate, as in the compressed QSAIndexer.
logical_positions = indexer_metadata.get_seqlens_expanded() - 1
else:
logical_positions = getattr(forward_batch, "positions", None)
if logical_positions is None:
logical_positions = positions[0] if positions.ndim == 2 else positions
logical_positions = logical_positions.flatten()
# DP padding adds token rows that belong to no request;
# token_to_batch_idx is the source of truth for semantic rows.
num_valid_tokens = indexer_metadata.get_token_to_batch_idx().numel()
if logical_positions.numel() < num_valid_tokens:
raise ValueError(
"tokenwise QSA logical positions are shorter than the request "
f"mapping: positions={logical_positions.numel()}, "
f"mapping={num_valid_tokens}"
)
if hidden_states.shape[0] < num_valid_tokens:
raise ValueError(
"tokenwise QSA hidden states are shorter than the request "
f"mapping: hidden={hidden_states.shape[0]}, "
f"mapping={num_valid_tokens}"
)
position_tokens = (
positions.shape[-1] if positions.ndim == 2 else positions.numel()
)
if position_tokens < num_valid_tokens:
raise ValueError(
"tokenwise QSA RoPE positions are shorter than the request "
f"mapping: positions={position_tokens}, "
f"mapping={num_valid_tokens}"
)
logical_positions = logical_positions[:num_valid_tokens]
hidden_states = hidden_states[:num_valid_tokens]
positions = (
positions[:, :num_valid_tokens]
if positions.ndim == 2
else positions[:num_valid_tokens]
)
if num_valid_tokens == 0:
return torch.empty(
(0, self.token_topk),
dtype=torch.int32,
device=hidden_states.device,
)
q, w, k = self.project_qkw(hidden_states, positions)
pool = indexer_metadata.token_to_kv_pool
out_cache_loc = getattr(indexer_metadata, "out_cache_loc", None)
if out_cache_loc is None:
out_cache_loc = forward_batch.out_cache_loc
pool.set_dsa_index_k_buffer(self.layer_id, out_cache_loc[:num_valid_tokens], k)
if is_paged:
return self._select_paged(q, w, indexer_metadata)
return self._select_prefill(q, w, logical_positions, indexer_metadata)
def _select_paged(
self,
q: torch.Tensor,
w: torch.Tensor,
indexer_metadata,
) -> torch.Tensor:
"""Per-query-row top-k over ``[0, row_len)`` for paged modes."""
pool = indexer_metadata.token_to_kv_pool
index_k = pool.get_dsa_index_k_buffer(self.layer_id)
sequence_lengths = indexer_metadata.sequence_lengths.to(torch.int32)
table = indexer_metadata.token_slot_table
rows, max_len = table.shape
if rows != indexer_metadata.token_to_batch_idx.numel():
raise ValueError(
"tokenwise QSA paged modes need one slot-table row per query "
f"row: table_rows={rows}, "
f"mapping={indexer_metadata.token_to_batch_idx.numel()}"
)
output = torch.full(
(rows, self.token_topk), -1, dtype=torch.int32, device=q.device
)
row_chunk = _qsa_prefill_row_chunk_size(rows, max_len, self.index_n_heads)
table_long = table.long()
for row_start in range(0, rows, row_chunk):
row_end = min(row_start + row_chunk, rows)
# Table columns at/after each row's length hold stale slots; the
# gathers stay in range and fast_topk masks them out by length.
k_chunk = index_k.index_select(0, table_long[row_start:row_end].reshape(-1))
k_chunk = k_chunk.reshape(row_end - row_start, max_len, 1, -1)
logits = torch_dsa_weighted_mqa_logits(
q[row_start:row_end],
w[row_start:row_end],
k_chunk,
self.score_scale,
)
lengths = sequence_lengths[row_start:row_end]
selected = qsa_fast_topk(
logits,
torch.zeros_like(lengths),
lengths.clamp(min=0, max=max_len),
topk=self.token_topk,
)
output[row_start:row_end].copy_(selected)
return output
def _select_prefill(
self,
q: torch.Tensor,
w: torch.Tensor,
logical_positions: torch.Tensor,
indexer_metadata,
) -> torch.Tensor:
"""Packed per-sequence top-k with causal windows for extend modes."""
pool = indexer_metadata.token_to_kv_pool
index_k = pool.get_dsa_index_k_buffer(self.layer_id)
sequence_lengths = indexer_metadata.sequence_lengths.to(torch.int32)
table = indexer_metadata.token_slot_table
query_sequence_ids = indexer_metadata.token_to_batch_idx.long()
row_ends_all = (logical_positions.to(torch.int32) + 1).clamp(
min=0, max=table.shape[1]
)
rows = q.shape[0]
output = torch.full(
(rows, self.token_topk), -1, dtype=torch.int32, device=q.device
)
for sequence_id in range(sequence_lengths.numel()):
seq_len = int(sequence_lengths[sequence_id].item())
row_mask = query_sequence_ids == sequence_id
if seq_len <= 0 or not bool(row_mask.any()):
continue
row_indices = row_mask.nonzero(as_tuple=True)[0]
slots = table[sequence_id, :seq_len].long()
k_seq = index_k.index_select(0, slots)
row_chunk = _qsa_prefill_row_chunk_size(
row_indices.numel(), seq_len, self.index_n_heads
)
for chunk_start in range(0, row_indices.numel(), row_chunk):
chunk_rows = row_indices[chunk_start : chunk_start + row_chunk]
row_ends = row_ends_all.index_select(0, chunk_rows)
logits = torch_dsa_weighted_mqa_logits(
q.index_select(0, chunk_rows),
w.index_select(0, chunk_rows),
k_seq,
self.score_scale,
)
selected = qsa_fast_topk(
logits,
torch.zeros_like(row_ends),
row_ends,
topk=self.token_topk,
)
# Tensor indexing returns a copy on read; use index_put style
# assignment or the selection would never reach `output`.
output[chunk_rows] = selected
return output
__all__ = [
"QwenDSAIndexer",
"torch_dsa_weighted_mqa_logits",
]
python/sglang/srt/layers/attention/qsa/glue.py
+86
View File
@@ -0,0 +1,86 @@
"""Assembly helpers shared by models that carry a QSA indexer."""
from __future__ import annotations
from sglang.srt.layers.attention.qsa.config import (
QSA_VARIANT_COMPRESSED,
parse_qsa_profile,
)
def build_qsa_indexer(
config,
*,
layer_id: int,
quant_config=None,
prefix: str = "",
rotary_emb=None,
):
profile = parse_qsa_profile(config)
if profile is None:
raise ValueError(
"build_qsa_indexer requires a config with a QSA indexer schema"
)
if profile.variant == QSA_VARIANT_COMPRESSED:
# The compressed indexer reuses the layer's own Qwen4-Exp RoPE
# (mrope); there is intentionally no plain-rope path for it here.
from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer
return QSAIndexer(
config=config,
layer_id=layer_id,
quant_config=quant_config,
prefix=prefix,
rotary_emb=rotary_emb,
)
# Tokenwise (Qwen3Next-DSA): the Lightning Indexer owns its plain
# per-token RoPE; a shared layer rotary is neither needed nor accepted.
from sglang.srt.layers.attention.qsa.dsa_indexer import QwenDSAIndexer
return QwenDSAIndexer(
config=config,
layer_id=layer_id,
quant_config=quant_config,
prefix=prefix,
)
def resolve_qsa_sparse_backend(attn_backend):
"""Backend owning the QSA MTP sparse-selection hooks;
a hybrid wrapper keeps them on its full-attention side.
``set_mtp_shared_sparse_indices`` is the probe, the one hook all owners define."""
if hasattr(attn_backend, "set_mtp_shared_sparse_indices"):
return attn_backend
full_attn_backend = getattr(attn_backend, "full_attn_backend", None)
if full_attn_backend is not None and hasattr(
full_attn_backend, "set_mtp_shared_sparse_indices"
):
return full_attn_backend
return attn_backend
def get_qsa_indexer_metadata(attn_backend, layer_id: int, forward_batch):
"""Fetch indexer metadata from a (possibly hybrid-wrapped) backend."""
metadata = None
get_metadata = getattr(attn_backend, "get_indexer_metadata", None)
if get_metadata is not None:
metadata = get_metadata(layer_id, forward_batch)
if metadata is None:
full_attn_backend = getattr(attn_backend, "full_attn_backend", None)
if full_attn_backend is not None and full_attn_backend is not attn_backend:
get_metadata = getattr(full_attn_backend, "get_indexer_metadata", None)
if get_metadata is not None:
metadata = get_metadata(layer_id, forward_batch)
if metadata is None:
raise RuntimeError("QSA backend did not provide indexer metadata")
return metadata
__all__ = [
"build_qsa_indexer",
"get_qsa_indexer_metadata",
"resolve_qsa_sparse_backend",
]
python/sglang/srt/layers/attention/qsa/graph_metadata.py
+219
View File
@@ -0,0 +1,219 @@
"""GPU builders for QSA CUDA-graph replay metadata.
Compressed addressing is pure arithmetic over the page-aligned full-KV cache:
a group's compressed slot is any of its raw slots floor-divided by the compress ratio,
so per-row graph buffers are rebuilt from request lengths and ``req_to_token`` alone;
accept-dependent speculative lengths never need the host.
Both kernels run once eagerly at capture warmup,
then are recorded into the main CUDA graph through ``init_forward_metadata_in_graph``;
all inputs are stable-address runner buffers.
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl
@triton.jit
def _qsa_graph_layout_kernel(
# Request-level inputs.
seq_lens_ptr, # [bs] base lengths
req_pool_ptr, # [bs] request pool slots
extend_lens_ptr, # [bs] per-request extend lengths (draft extend)
# Row layout buffers (graph persistent state).
row_seq_lens_ptr,
row_prefix_lens_ptr,
row_req_pool_ptr,
bs,
num_tokens,
num_padding,
extend_len, # uniform extend length (target verify); 0 -> extend_lens_ptr
MODE: tl.constexpr, # 0 = decode, 1 = target verify, 2 = draft extend
):
pid = tl.program_id(0)
if MODE == 0:
if pid < bs:
real = pid < bs - num_padding
seq_len = tl.load(seq_lens_ptr + pid).to(tl.int32)
req = tl.load(req_pool_ptr + pid).to(tl.int64)
# Padding rows alias request slot 0: it is never allocated, so
# its pending-ring rows are the inert dump for their state
# stores, and its req_to_token row reads stay in-bounds.
req = tl.where(real, req, 0)
seq_len = tl.where(real, seq_len, 1)
prefix = tl.maximum(seq_len - 1, 0)
tl.store(row_seq_lens_ptr + pid, seq_len)
tl.store(row_req_pool_ptr + pid, req.to(tl.int32))
tl.store(row_prefix_lens_ptr + pid, prefix)
return
real_reqs = bs - num_padding
if pid == bs:
# Tail program: dummy rows for the static capacity past the real layout.
if MODE == 1:
row_start = real_reqs * extend_len
else:
row_start = 0
for j in range(real_reqs):
row_start += tl.load(extend_lens_ptr + j)
for row in range(row_start, num_tokens):
tl.store(row_seq_lens_ptr + row, 1)
tl.store(row_prefix_lens_ptr + row, 0)
# Request slot 0 is never allocated: inert for ring stores and
# in-bounds for every row read.
tl.store(row_req_pool_ptr + row, 0)
return
if MODE == 1:
eff = tl.where(pid < real_reqs, extend_len, 0)
offset = tl.minimum(pid, real_reqs) * extend_len
else:
eff = 0
offset = 0
for j in range(bs):
e_j = tl.where(j < real_reqs, tl.load(extend_lens_ptr + j), 0)
offset += tl.where(j < pid, e_j, 0)
eff = tl.where(j == pid, e_j, eff)
base = tl.load(seq_lens_ptr + pid).to(tl.int32)
req = tl.load(req_pool_ptr + pid).to(tl.int64)
if MODE == 1:
prefix = base
limit = base + eff
else:
prefix = tl.maximum(base - eff, 0)
limit = base
for j in range(eff):
row = offset + j
seq_len = tl.minimum(prefix + 1 + j, limit)
tl.store(row_seq_lens_ptr + row, seq_len)
tl.store(row_prefix_lens_ptr + row, prefix)
tl.store(row_req_pool_ptr + row, req.to(tl.int32))
@triton.jit
def _qsa_graph_row_metadata_kernel(
# Row layout buffers (filled by the layout kernel).
row_seq_lens_ptr,
row_req_pool_ptr,
# Graph output buffers.
compressed_lens_ptr,
write_locs_ptr,
page_table_ptr,
logical_positions_ptr,
state_slots_ptr,
ring_locs_ptr,
# Pool state.
req_to_token_ptr,
req_to_token_row_stride,
max_pages,
RATIO: tl.constexpr,
FULL_PAGE: tl.constexpr, # full-KV tokens per page
PAGE_BLOCK: tl.constexpr,
):
row = tl.program_id(0)
seq_len = tl.load(row_seq_lens_ptr + row).to(tl.int32)
req = tl.load(row_req_pool_ptr + row).to(tl.int64)
token_row = req * req_to_token_row_stride
current = tl.maximum(seq_len - 1, 0)
last_loc = tl.load(req_to_token_ptr + token_row + current).to(tl.int32)
compressed = seq_len // RATIO
tl.store(compressed_lens_ptr + row, compressed)
# The page-aligned allocator keeps each compression group inside one page,
# so last_loc // RATIO is the group's compressed slot; slot 0 is the padding slot.
boundary = (seq_len > 0) & (seq_len % RATIO == 0)
write_loc = tl.where(boundary, last_loc // RATIO, 0)
tl.store(write_locs_ptr + row, write_loc)
tl.store(logical_positions_ptr + row, current)
tl.store(state_slots_ptr + row, req * RATIO + (current % RATIO).to(tl.int64))
ring_base = row.to(tl.int64) * RATIO
for k in tl.static_range(RATIO):
member = tl.maximum(current - (RATIO - 1 - k), 0)
slot = req * RATIO + (member % RATIO).to(tl.int64)
tl.store(ring_locs_ptr + ring_base + k, slot.to(tl.int32))
# Page-table entries are the request's FULL-KV page ids, read from the
# page-aligned req_to_token row; the scoring kernels turn them into
# compressed slots as page_id * (FULL_PAGE // RATIO) + block_in_page.
table_row = page_table_ptr + row.to(tl.int64) * max_pages
offs = tl.arange(0, PAGE_BLOCK)
row_width_pages = req_to_token_row_stride // FULL_PAGE
for p0 in range(0, max_pages, PAGE_BLOCK):
idx = p0 + offs
valid = idx < tl.minimum(max_pages, row_width_pages)
loc = tl.load(
req_to_token_ptr + token_row + idx * FULL_PAGE, mask=valid, other=0
)
tl.store(table_row + idx, tl.maximum(loc // FULL_PAGE, 0), mask=valid)
def supports_graph_metadata_kernels(pool, device) -> bool:
"""Whether the CUDA fast path can serve this pool/device pair."""
from sglang.srt.mem_cache.qsa_kv_pool import QSATokenToKVPool
return torch.device(device).type == "cuda" and isinstance(pool, QSATokenToKVPool)
def launch_graph_metadata(
*,
mode,
bs,
num_rows,
seq_lens,
req_pool_indices,
extend_lens,
extend_len,
num_padding,
metadata,
req_to_token,
pool,
) -> None:
indexer = metadata.indexer_metadata
max_pages = indexer.graph_compressed_page_table.shape[1]
row_seq_lens = metadata.sequence_lengths
row_req_pool = metadata.row_req_pool_indices
row_prefix_lens = indexer.graph_prefix_lengths
_qsa_graph_layout_kernel[(bs + 1,)](
seq_lens,
req_pool_indices,
(
extend_lens
if extend_lens is not None
else row_seq_lens # unused dummy pointer
),
row_seq_lens,
row_prefix_lens,
row_req_pool,
bs,
num_rows,
num_padding,
extend_len,
MODE=mode,
num_warps=1,
)
_qsa_graph_row_metadata_kernel[(num_rows,)](
row_seq_lens,
row_req_pool,
indexer.graph_compressed_lengths,
indexer.graph_write_locs,
indexer.graph_compressed_page_table,
indexer.decode_logical_positions,
indexer.pending_ring_slots,
indexer.graph_ring_group_locs,
req_to_token,
req_to_token.stride(0),
max_pages,
RATIO=indexer.compress_ratio,
FULL_PAGE=pool.qsa_compressed_page_size * indexer.compress_ratio,
PAGE_BLOCK=128,
num_warps=1,
)
python/sglang/srt/layers/attention/qsa/kernel.py
+332
View File
@@ -0,0 +1,332 @@
"""CUDA kernels and tensor transforms for simple QSA."""
from __future__ import annotations
from typing import Optional
import torch
import triton
import triton.language as tl
def average_pool_qsa_keys(key_groups: torch.Tensor) -> torch.Tensor:
"""FP32-average complete key groups shaped ``[groups, ratio, kv_heads, dim]``."""
if key_groups.ndim != 4:
raise ValueError(
"QSA key groups must be [groups, ratio, kv_heads, head_dim], "
f"got {key_groups.shape}"
)
return key_groups.float().mean(dim=1).to(key_groups.dtype)
def qsa_fast_topk(
logits: torch.Tensor,
row_starts: torch.Tensor,
row_ends: torch.Tensor,
topk: int,
) -> torch.Tensor:
"""Select compressed blocks, with a compatibility fallback for top-k 512."""
lengths = (row_ends - row_starts).to(device=logits.device, dtype=torch.int32)
starts = row_starts.to(device=logits.device, dtype=torch.int32)
if logits.is_cuda:
if topk == 512:
# Prefer the JIT kernel: it ships with the sglang python package,
# so top-k 512 works regardless of the installed sgl_kernel version.
from sglang.kernels.ops.elementwise.fast_topk import fast_topk
return fast_topk(logits, lengths, topk=512, row_starts=starts)
from sgl_kernel import top_k as top_k_module
supported_topk = getattr(top_k_module, "_FAST_TOPK_SUPPORTED_K", (2048,))
if topk in supported_topk:
return top_k_module.fast_topk_v2(
logits, lengths, topk=topk, row_starts=starts
)
raise ValueError(
f"QSA top-k {topk} is unsupported by sgl_kernel; "
f"supported values are {supported_topk}"
)
# CPU/reference path mirrors the CUDA operator's fixed-width, relative output.
output = torch.full(
(logits.shape[0], topk),
-1,
dtype=torch.int32,
device=logits.device,
)
for row in range(logits.shape[0]):
start = int(starts[row])
length = int(lengths[row])
width = min(length, topk)
if width:
output[row, :width] = torch.topk(
logits[row, start : start + length], width
).indices.to(torch.int32)
return output
def torch_expand_qsa_block_indices(
block_indices: torch.Tensor,
query_positions: torch.Tensor,
sequence_lengths: torch.Tensor,
compress_ratio: int,
token_topk: int,
) -> torch.Tensor:
"""Expand compressed block indices into fixed-width logical token indices."""
block_topk = (token_topk + compress_ratio - 1) // compress_ratio
final_topk = token_topk + compress_ratio - 1
if block_indices.ndim != 2 or block_indices.shape[1] != block_topk:
raise ValueError(
f"expected block indices [M, {block_topk}], got "
f"{tuple(block_indices.shape)}"
)
rows = block_indices.shape[0]
if query_positions.numel() != rows or sequence_lengths.numel() != rows:
raise ValueError("query positions and sequence lengths must match top-k rows")
device = block_indices.device
blocks = block_indices.long()
offsets = torch.arange(compress_ratio, device=device, dtype=torch.long)
expanded = blocks.unsqueeze(-1) * compress_ratio + offsets
expanded = torch.where(
blocks.unsqueeze(-1) >= 0, expanded, torch.full_like(expanded, -1)
).reshape(rows, block_topk * compress_ratio)
expanded = expanded[:, :token_topk]
query_positions = query_positions.to(device=device, dtype=torch.long)
sequence_lengths = sequence_lengths.to(device=device, dtype=torch.long)
expanded = torch.where(
(expanded >= 0) & (expanded < sequence_lengths.unsqueeze(1)),
expanded,
torch.full_like(expanded, -1),
)
tail_offsets = torch.arange(compress_ratio - 1, device=device, dtype=torch.long)
visible_tokens = query_positions + 1
tail_start = (
torch.div(visible_tokens, compress_ratio, rounding_mode="floor")
* compress_ratio
)
tail_count = visible_tokens - tail_start
tail = tail_start.unsqueeze(1) + tail_offsets.unsqueeze(0)
tail_valid = (tail_offsets.unsqueeze(0) < tail_count.unsqueeze(1)) & (
tail < sequence_lengths.unsqueeze(1)
)
tail = torch.where(tail_valid, tail, torch.full_like(tail, -1))
result = torch.cat([expanded, tail], dim=1)
# Keep all valid entries contiguous. This is required by the FA2 packing path.
order = torch.arange(final_topk, device=device).unsqueeze(0).expand(rows, -1)
sort_key = torch.where(result >= 0, order, order + final_topk)
return result.gather(1, torch.argsort(sort_key, dim=1, stable=True)).to(torch.int32)
@triton.jit
def _expand_qsa_block_indices_kernel(
block_indices,
query_positions,
sequence_lengths,
output,
block_stride: tl.constexpr,
output_stride: tl.constexpr,
BLOCK_TOPK: tl.constexpr,
COMPRESS_RATIO: tl.constexpr,
TOKEN_TOPK: tl.constexpr,
FINAL_TOPK: tl.constexpr,
OUTPUT_BLOCK_SIZE: tl.constexpr,
):
row = tl.program_id(0)
cols = tl.arange(0, OUTPUT_BLOCK_SIZE)
sequence_length = tl.load(sequence_lengths + row)
source_block_cols = cols // COMPRESS_RATIO
offsets = cols % COMPRESS_RATIO
blocks = tl.load(
block_indices + row * block_stride + source_block_cols,
mask=(cols < TOKEN_TOPK) & (source_block_cols < BLOCK_TOPK),
other=-1,
)
expanded = blocks * COMPRESS_RATIO + offsets
expanded_valid = (
(cols < TOKEN_TOPK)
& (blocks >= 0)
& (expanded >= 0)
& (expanded < sequence_length)
)
valid_block_count = tl.sum(
(
(cols < BLOCK_TOPK)
& (
tl.load(
block_indices + row * block_stride + cols,
mask=cols < BLOCK_TOPK,
other=-1,
)
>= 0
)
).to(tl.int32),
axis=0,
)
valid_token_count = tl.minimum(valid_block_count * COMPRESS_RATIO, TOKEN_TOPK)
query_position = tl.load(query_positions + row)
visible_tokens = query_position + 1
tail_start = (visible_tokens // COMPRESS_RATIO) * COMPRESS_RATIO
tail_offset = cols - valid_token_count
tail_count = visible_tokens - tail_start
tail = tail_start + tail_offset
tail_valid = (
(tail_offset >= 0)
& (tail_offset < COMPRESS_RATIO - 1)
& (tail_offset < tail_count)
& (tail < sequence_length)
)
result = tl.where(
expanded_valid & (cols < valid_token_count),
expanded,
tl.where(tail_valid, tail, -1),
)
tl.store(
output + row * output_stride + cols,
result,
mask=cols < FINAL_TOPK,
)
def triton_expand_qsa_block_indices(
block_indices: torch.Tensor,
query_positions: torch.Tensor,
sequence_lengths: torch.Tensor,
compress_ratio: int,
token_topk: int,
) -> torch.Tensor:
"""CUDA fast path for fast_topk_v2 output (valid blocks precede -1 padding)."""
rows, block_topk = block_indices.shape
final_topk = token_topk + compress_ratio - 1
output = torch.empty(
(rows, final_topk), dtype=torch.int32, device=block_indices.device
)
if rows == 0:
return output
_expand_qsa_block_indices_kernel[(rows,)](
block_indices,
query_positions,
sequence_lengths,
output,
block_indices.stride(0),
output.stride(0),
BLOCK_TOPK=block_topk,
COMPRESS_RATIO=compress_ratio,
TOKEN_TOPK=token_topk,
FINAL_TOPK=final_topk,
OUTPUT_BLOCK_SIZE=triton.next_power_of_2(final_topk),
num_warps=8,
)
return output
def expand_qsa_block_indices(
block_indices: torch.Tensor,
query_positions: torch.Tensor,
sequence_lengths: torch.Tensor,
compress_ratio: int,
token_topk: int,
) -> torch.Tensor:
"""Expand compressed blocks with Triton on CUDA and Torch elsewhere."""
block_topk = (token_topk + compress_ratio - 1) // compress_ratio
if block_indices.ndim != 2 or block_indices.shape[1] != block_topk:
raise ValueError(
f"expected block indices [M, {block_topk}], got "
f"{tuple(block_indices.shape)}"
)
rows = block_indices.shape[0]
if query_positions.numel() != rows or sequence_lengths.numel() != rows:
raise ValueError("query positions and sequence lengths must match top-k rows")
if block_indices.is_cuda:
# The Triton kernel loads positions/lengths as scalars, so any integer
# dtype works; skip the int64 conversion copies.
return triton_expand_qsa_block_indices(
block_indices.contiguous(),
query_positions.to(device=block_indices.device).contiguous(),
sequence_lengths.to(device=block_indices.device).contiguous(),
compress_ratio,
token_topk,
)
return torch_expand_qsa_block_indices(
block_indices,
query_positions,
sequence_lengths,
compress_ratio,
token_topk,
)
def qsa_sparse_attention(
q: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
token_slots: torch.Tensor,
softmax_scale: Optional[float] = None,
) -> torch.Tensor:
"""Torch reference for sparse GQA over physical token slots."""
if q.ndim != 3 or k_cache.ndim != 3 or v_cache.ndim != 3:
raise ValueError("q, k_cache and v_cache must be rank-3 tensors")
if token_slots.ndim != 2 or token_slots.shape[0] != q.shape[0]:
raise ValueError(
"token slots must be [query_tokens, selected_tokens], got "
f"{token_slots.shape}"
)
if q.shape[-1] != k_cache.shape[-1] or q.shape[-1] != v_cache.shape[-1]:
raise ValueError("Q/K/V head dimensions must match")
if q.shape[1] % k_cache.shape[1] != 0:
raise ValueError("query heads must be divisible by KV heads")
return qsa_sparse_attention_reference(
q, k_cache, v_cache, token_slots, softmax_scale
)
def qsa_sparse_attention_reference(
q: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
token_slots: torch.Tensor,
softmax_scale: Optional[float] = None,
) -> torch.Tensor:
"""Device-agnostic sparse GQA reference."""
scale = softmax_scale or q.shape[-1] ** -0.5
outputs = []
repeats = q.shape[1] // k_cache.shape[1]
for row in range(q.shape[0]):
valid = token_slots[row] >= 0
slots = token_slots[row, valid].long()
if slots.numel() == 0:
outputs.append(torch.zeros_like(q[row]))
continue
keys = k_cache.index_select(0, slots).repeat_interleave(repeats, dim=1)
values = v_cache.index_select(0, slots).repeat_interleave(repeats, dim=1)
scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float()) * scale
probabilities = torch.softmax(scores, dim=-1)
outputs.append(
torch.einsum("hk,khd->hd", probabilities, values.float()).to(q.dtype)
)
return torch.stack(outputs)
__all__ = [
"average_pool_qsa_keys",
"expand_qsa_block_indices",
"torch_expand_qsa_block_indices",
"triton_expand_qsa_block_indices",
"qsa_fast_topk",
"qsa_sparse_attention",
"qsa_sparse_attention_reference",
]
python/sglang/srt/layers/attention/qsa/metadata.py
+313
View File
@@ -0,0 +1,313 @@
"""Metadata owned by the simple QSA implementation.
QSA intentionally does not inherit the NSA metadata abstraction. This module
contains only fields and transforms consumed by the indexer.
"""
from __future__ import annotations
from typing import Optional, Tuple
import msgspec
import torch
from sglang.srt.layers.attention.qsa.kernel import qsa_fast_topk
def build_qsa_row_ranges(
sequence_lengths: torch.Tensor,
query_positions: torch.Tensor,
query_sequence_ids: torch.Tensor,
compress_ratio: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build packed compressed-key ranges for prefill scoring."""
sequence_lengths = sequence_lengths.to(dtype=torch.int32)
compressed_lengths = torch.div(
sequence_lengths, compress_ratio, rounding_mode="floor"
)
compressed_cu_seqlens = torch.nn.functional.pad(
compressed_lengths.cumsum(0), (1, 0)
).to(torch.int32)
query_sequence_ids = query_sequence_ids.to(
device=sequence_lengths.device, dtype=torch.long
)
row_starts = compressed_cu_seqlens.index_select(0, query_sequence_ids)
visible_blocks = torch.div(
query_positions.to(device=sequence_lengths.device, dtype=torch.int32) + 1,
compress_ratio,
rounding_mode="floor",
)
max_blocks = compressed_lengths.index_select(0, query_sequence_ids)
row_ends = row_starts + torch.minimum(visible_blocks, max_blocks)
return row_starts, row_ends, compressed_cu_seqlens
class QSAIndexerMetadata(msgspec.Struct, frozen=True):
"""All per-forward metadata consumed specifically by ``QSAIndexer``.
Row layout contract:
- ``sequence_lengths``/``token_slot_table`` carry one row per *sequence*
for extend modes and one row per *query token* for the paged modes
(decode, target_verify, draft_extend).
- ``token_to_batch_idx`` maps every query/token row handled by the indexer
onto a row of ``sequence_lengths``/``token_slot_table``; DP attention
token padding adds physical rows beyond this mapping, never inside it.
- For the paged modes the mapping is the identity
(``arange(num_query_rows)``), so page-table/MQA inputs built per
``sequence_lengths`` row line up with per-query sparse-attention rows.
"""
sequence_lengths: torch.Tensor
token_to_batch_idx: torch.Tensor
token_slot_table: torch.Tensor
out_cache_loc: torch.Tensor
token_to_kv_pool: object
compress_ratio: int
block_topk: int
req_pool_indices: Optional[torch.Tensor] = None
# Parallel per-group arrays for the groups compressed this forward:
# slot, sequence-local group-end position, and owning metadata row.
# The first member's token row in this forward's packed tensors is extend only,
# where group-aligned chunks keep every member in-chunk; None on paged forwards.
write_locs: Optional[torch.Tensor] = None
compress_group_positions: Optional[torch.Tensor] = None
compress_sequence_ids: Optional[torch.Tensor] = None
compress_member_rows: Optional[torch.Tensor] = None
is_cuda_graph: bool = False
graph_write_locs: Optional[torch.Tensor] = None
graph_compressed_page_table: Optional[torch.Tensor] = None
graph_compressed_lengths: Optional[torch.Tensor] = None
graph_prefix_lengths: Optional[torch.Tensor] = None
decode_page_table: Optional[torch.Tensor] = None
decode_lengths: Optional[torch.Tensor] = None
decode_logical_positions: Optional[torch.Tensor] = None
pending_ring_slots: Optional[torch.Tensor] = None
compress_group_ring_locs: Optional[torch.Tensor] = None
extend_rope_matrix: Optional[torch.Tensor] = None
graph_ring_group_locs: Optional[torch.Tensor] = None
def get_seqlens_int32(self) -> torch.Tensor:
return self.sequence_lengths.to(torch.int32)
def get_token_slot_table(self) -> torch.Tensor:
return self.token_slot_table
def get_seqlens_expanded(self) -> torch.Tensor:
return self.get_seqlens_int32().index_select(
0, self.get_token_to_batch_idx().long()
)
def get_token_to_batch_idx(self) -> torch.Tensor:
return self.token_to_batch_idx
def topk_transform(
self,
logits: torch.Tensor,
topk: int,
row_starts: Optional[torch.Tensor] = None,
row_ends: Optional[torch.Tensor] = None,
**kwargs,
) -> torch.Tensor:
if topk != self.block_topk:
raise ValueError(
f"QSA compressed top-k must be {self.block_topk}, got {topk}"
)
if row_starts is None or row_ends is None:
raise ValueError("QSA top-k transform requires row_starts and row_ends")
return qsa_fast_topk(logits, row_starts, row_ends, topk=self.block_topk)
def get_prefill_mqa_inputs(
self,
layer_id: int,
positions: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Gather packed compressed K and ragged ranges for prefill MQA."""
pool = self.token_to_kv_pool
ratio = self.compress_ratio
compressed_buffer = pool.get_qsa_compressed_k_buffer(layer_id)
parts = []
sequence_lengths = self.sequence_lengths.to(torch.int32)
sequence_lengths_list = sequence_lengths.tolist()
for sequence_id in range(len(sequence_lengths_list)):
complete_blocks = int(sequence_lengths_list[sequence_id]) // ratio
if complete_blocks == 0:
continue
# compressed slot = first raw slot // ratio; the allocator is page-aligned,
# so each group is contiguous in one page (see QSATokenToKVPool).
compressed_locs = (
self.token_slot_table[
sequence_id, : complete_blocks * ratio : ratio
].long()
// ratio
)
parts.append(compressed_buffer.index_select(0, compressed_locs))
compressed_keys = (
torch.cat(parts, dim=0)
if parts
else compressed_buffer.new_empty(
(0, pool.qsa_index_kv_heads, pool.qsa_index_head_dim)
)
)
num_valid_tokens = self.token_to_batch_idx.numel()
if positions.numel() < num_valid_tokens:
raise ValueError(
"QSA prefill positions are shorter than the request mapping: "
f"positions={positions.numel()}, mapping={num_valid_tokens}"
)
positions = positions[:num_valid_tokens]
row_starts, row_ends, _ = build_qsa_row_ranges(
sequence_lengths,
positions.to(sequence_lengths.device),
self.token_to_batch_idx.to(sequence_lengths.device),
self.compress_ratio,
)
return compressed_keys, row_starts, row_ends, sequence_lengths
def get_decode_mqa_inputs(
self, layer_id: int
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
"""Paged compressed-K cache inputs for decode MQA, one row per query row."""
pool = self.token_to_kv_pool
num_rows = self.sequence_lengths.numel()
if self.token_slot_table.shape[0] != num_rows:
raise ValueError(
"QSA decode page-table rows must match the per-query sequence "
f"lengths: table_rows={self.token_slot_table.shape[0]}, "
f"rows={num_rows}"
)
compressed_cache = pool.get_qsa_compressed_k_buffer(layer_id).reshape(
-1,
pool.qsa_compressed_page_size,
pool.qsa_index_kv_heads,
pool.qsa_index_head_dim,
)
if self.is_cuda_graph:
if (
self.graph_compressed_page_table is None
or self.graph_compressed_lengths is None
):
raise RuntimeError("QSA CUDA graph decode metadata is incomplete")
return (
compressed_cache,
self.graph_compressed_page_table,
self.graph_compressed_lengths,
self.graph_compressed_page_table.shape[1]
* pool.qsa_compressed_page_size,
)
if self.decode_page_table is not None and self.decode_lengths is not None:
return (
compressed_cache,
self.decode_page_table,
self.decode_lengths,
self.decode_page_table.shape[1] * pool.qsa_compressed_page_size,
)
compressed_page_table, compressed_lengths = compressed_decode_view(
compressed_page_size=pool.qsa_compressed_page_size,
compress_ratio=self.compress_ratio,
sequence_lengths=self.sequence_lengths,
token_slot_table=self.token_slot_table,
)
return (
compressed_cache,
compressed_page_table,
compressed_lengths,
compressed_page_table.shape[1] * pool.qsa_compressed_page_size,
)
def build_pending_ring_slots(
*,
token_to_batch_idx: torch.Tensor,
req_pool_indices: torch.Tensor,
sequence_lengths: torch.Tensor,
logical_positions: torch.Tensor,
compress_ratio: int,
is_extend: bool,
) -> torch.Tensor:
"""Pending-ring slot ``req_pool_idx * ratio + position % ratio`` per token.
On extend, tokens before the pending tail dump into rows [0, ratio),
which no request owns (request slot 0 is never allocated); CUDA-graph safe."""
rows = token_to_batch_idx.long()[: logical_positions.numel()]
requests = req_pool_indices.long()[rows]
positions = logical_positions.long()
slots = requests * compress_ratio + positions % compress_ratio
if is_extend:
lengths = sequence_lengths.long()[rows]
pending = positions >= (lengths // compress_ratio) * compress_ratio
slots = torch.where(pending, slots, positions % compress_ratio)
return slots
def build_group_ring_slots(
*,
req_pool_indices: torch.Tensor,
group_end_positions: torch.Tensor,
sequence_ids: torch.Tensor,
compress_ratio: int,
) -> torch.Tensor:
"""Ring slots of a planned group's members, oldest first."""
requests = req_pool_indices.long()[sequence_ids]
offsets = torch.arange(
compress_ratio - 1,
-1,
-1,
device=group_end_positions.device,
dtype=torch.long,
)
positions = (group_end_positions[:, None] - offsets[None, :]).clamp_min(0)
return requests[:, None] * compress_ratio + positions % compress_ratio
def build_rope_position_matrix(
rope_positions: torch.Tensor, num_tokens: int
) -> torch.Tensor:
"""This forward's RoPE coordinates as the [tokens, 3] layout the fused
compress kernel reads."""
if rope_positions.ndim == 1:
return (
rope_positions[:num_tokens].long().unsqueeze(1).expand(-1, 3)
).contiguous()
return rope_positions[:, :num_tokens].long().transpose(0, 1).contiguous()
def compressed_decode_view(
*,
compressed_page_size: int,
compress_ratio: int,
sequence_lengths: torch.Tensor,
token_slot_table: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compressed page table and lengths for decode MQA.
Page-table entries are full-KV page ids read off the page-aligned
token-slot rows; the scoring kernel converts them to compressed
slots as page_id * compressed_page_size + block_in_page. Entries
past a row's compressed length are stale-but-unread (bounded by
compressed_lengths); clamp keeps them non-negative.
"""
full_page = compressed_page_size * compress_ratio
compressed_lengths = torch.div(
sequence_lengths.to(torch.int32),
compress_ratio,
rounding_mode="floor",
)
compressed_page_table = (
(token_slot_table[:, ::full_page].long() // full_page)
.clamp_min(0)
.to(torch.int32)
)
return compressed_page_table, compressed_lengths
__all__ = [
"QSAIndexerMetadata",
"build_qsa_row_ranges",
"build_pending_ring_slots",
"build_group_ring_slots",
"build_rope_position_matrix",
"compressed_decode_view",
]
python/sglang/srt/layers/attention/qsa/mqa.py
+415
View File
@@ -0,0 +1,415 @@
"""Weight-free TileLang MQA operators for the simple QSA indexer;
the torch implementations are the fallback and the reference."""
import math
from typing import Optional
import torch
try:
import flashinfer.comm # noqa: F401
except ImportError:
pass
try:
import tilelang
from tilelang import language as T
HAS_TILELANG = True
except ImportError:
tilelang = None
T = None
HAS_TILELANG = False
def _validate_q(q: torch.Tensor) -> None:
if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0:
raise ValueError(f"QSA requires q [tokens, heads, head_dim], got {q.shape}")
def _validate_k(k: torch.Tensor) -> None:
if k.ndim != 3 or k.shape[1] != 1 or k.shape[2] <= 0:
raise ValueError(f"QSA MQA requires k [tokens, 1, head_dim], got {k.shape}")
def torch_qsa_mqa_prefill(
q: torch.Tensor,
k: torch.Tensor,
row_starts: torch.Tensor,
row_ends: torch.Tensor,
score_scale: Optional[float] = None,
) -> torch.Tensor:
"""Torch reference for packed, variable-length prefill MQA."""
_validate_q(q)
_validate_k(k)
if q.shape[-1] != k.shape[-1]:
raise ValueError("QSA query and key head dimensions must match")
scores = torch.einsum("mhd,nd->mnh", q.float(), k[:, 0].float())
logits = torch.relu(scores).sum(dim=-1) / (score_scale or math.sqrt(q.shape[-1]))
columns = torch.arange(k.shape[0], device=q.device).unsqueeze(0)
valid = (columns >= row_starts.to(q.device).reshape(-1, 1)) & (
columns < row_ends.to(q.device).reshape(-1, 1)
)
return logits.masked_fill(~valid, -float("inf"))
def _validate_decode_inputs(
q: torch.Tensor,
k_cache: torch.Tensor,
page_table: torch.Tensor,
context_lens: torch.Tensor,
) -> None:
_validate_q(q)
if k_cache.ndim != 4 or k_cache.shape[2] != 1:
raise ValueError(
"QSA decode cache must be [pages, page_size, 1, head_dim], "
f"got {tuple(k_cache.shape)}"
)
if k_cache.shape[-1] != q.shape[-1]:
raise ValueError("QSA query and key head dimensions must match")
if page_table.ndim != 2 or page_table.shape[0] != q.shape[0]:
raise ValueError("QSA decode page table must have one row per query")
if context_lens.numel() != q.shape[0]:
raise ValueError("QSA decode context lengths must have one entry per query")
def torch_qsa_mqa_decode(
q: torch.Tensor,
k_cache: torch.Tensor,
page_table: torch.Tensor,
context_lens: torch.Tensor,
max_model_len: int,
score_scale: Optional[float] = None,
) -> torch.Tensor:
"""Torch reference for variable-length paged decode MQA."""
_validate_decode_inputs(q, k_cache, page_table, context_lens)
batch = q.shape[0]
page_size = k_cache.shape[1]
total = page_table.shape[1] * page_size
gathered = k_cache[page_table.long().clamp_min(0).reshape(-1), :, 0].reshape(
batch, total, q.shape[-1]
)
scores = torch.einsum("bhd,bnd->bnh", q.float(), gathered.float())
scores = torch.relu(scores).sum(dim=-1) / (score_scale or math.sqrt(q.shape[-1]))
positions = torch.arange(total, device=q.device).unsqueeze(0)
scores.masked_fill_(
positions >= context_lens.to(q.device).reshape(-1, 1), -float("inf")
)
logits = torch.full(
(batch, max_model_len), -float("inf"), dtype=torch.float32, device=q.device
)
copy_len = min(total, max_model_len)
if copy_len:
logits[:, :copy_len] = scores[:, :copy_len]
return logits
if HAS_TILELANG:
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
}
)
def _tilelang_qsa_mqa_prefill_kernel(
heads: int,
head_dim: int,
block_n: int = 64,
block_q: int = 32,
num_stages: int = 3,
threads: int = 512,
):
rows = T.dynamic("rows")
keys = T.dynamic("keys")
@T.prim_func
def kernel(
Q: T.Tensor([rows * heads, head_dim], T.bfloat16), # type: ignore
K: T.Tensor([keys, head_dim], T.bfloat16), # type: ignore
Logits: T.Tensor([rows, keys], T.float32), # type: ignore
Starts: T.Tensor([rows], T.int32), # type: ignore
Ends: T.Tensor([rows], T.int32), # type: ignore
):
with T.Kernel(T.ceildiv(rows, block_q), threads=threads) as bx:
q_shared = T.alloc_shared([block_q * heads, head_dim], T.bfloat16)
k_shared = T.alloc_shared([block_n, head_dim], T.bfloat16)
scores = T.alloc_fragment([block_n, block_q * heads], T.float32)
scores_3d = T.reshape(scores, (block_n, block_q, heads))
reduced = T.alloc_fragment([block_n, block_q], T.float32)
row_base = bx * block_q
start_min = T.alloc_var(T.int32)
end_max = T.alloc_var(T.int32)
start_min = 2147483647
end_max = -2147483648
for qi in T.serial(block_q):
start_min = T.min(start_min, T.min(Starts[row_base + qi], keys))
end_max = T.max(end_max, T.min(Ends[row_base + qi], keys))
T.copy(Q[row_base * heads, 0], q_shared)
for ni in T.Pipelined(
T.ceildiv(end_max - start_min, block_n), num_stages=num_stages
):
T.copy(K[start_min + ni * block_n, 0], k_shared)
T.gemm(
k_shared,
q_shared,
scores,
transpose_B=True,
clear_accum=True,
policy=T.GemmWarpPolicy.FullCol,
)
for n, qi, head in T.Parallel(block_n, block_q, heads):
scores_3d[n, qi, head] = T.max(scores_3d[n, qi, head], 0.0)
T.reduce_sum(scores_3d, reduced, dim=-1, clear=True)
for qi, n in T.Parallel(block_q, block_n):
Logits[row_base + qi, start_min + ni * block_n + n] = reduced[
n, qi
]
return kernel
@tilelang.jit
def _tilelang_qsa_mqa_mask_kernel(threads: int = 512, block_k: int = 4096):
rows = T.dynamic("rows")
keys = T.dynamic("keys")
@T.prim_func
def kernel(
Logits: T.Tensor([rows, keys], T.float32), # type: ignore
Starts: T.Tensor([rows], T.int32), # type: ignore
Ends: T.Tensor([rows], T.int32), # type: ignore
):
with T.Kernel(rows, threads=threads) as bx:
tx = T.thread_binding(0, threads, thread="threadIdx.x")
for block in T.Pipelined(T.ceildiv(keys, block_k)):
for item in T.serial(block_k // threads):
column = block * block_k + item * threads + tx
if column < Starts[bx] or column >= Ends[bx]:
Logits[bx, column] = -T.infinity(T.float32)
return kernel
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
}
)
def _tilelang_qsa_mqa_decode_kernel(
heads: int,
head_dim: int,
page_size: int = 64,
groups_per_cta: int = 1,
num_stages: int = 3,
threads: int = 128,
):
# The MMA layout needs 64 GEMM rows; a compressed page has full_page // ratio,
# so pages are packed as sub-pages of one 64-row tile.
GROUP = 64
assert GROUP % page_size == 0, page_size
sub_pages = GROUP // page_size
batch = T.dynamic("batch")
pages = T.dynamic("pages")
max_pages = T.dynamic("max_pages")
max_model_len = T.dynamic("max_model_len")
@T.prim_func
def kernel(
Q: T.Tensor([batch, 1, heads, head_dim], T.bfloat16), # type: ignore
KCache: T.Tensor([pages, page_size, 1, head_dim], T.bfloat16), # type: ignore
PageTable: T.Tensor([batch, max_pages], T.int32), # type: ignore
ContextLens: T.Tensor([batch], T.int32), # type: ignore
Logits: T.Tensor([batch, max_model_len], T.float32), # type: ignore
Scale: T.float32,
):
with T.Kernel(
batch,
T.ceildiv(T.ceildiv(max_pages, sub_pages), groups_per_cta),
threads=threads,
) as (bx, group_block):
q_shared = T.alloc_shared([heads, head_dim], T.bfloat16)
k_shared = T.alloc_shared([GROUP, head_dim], T.bfloat16)
scores = T.alloc_fragment([GROUP, heads], T.float32)
reduced = T.alloc_fragment([GROUP], T.float32)
T.copy(Q[bx, 0, :, :], q_shared)
context_len = ContextLens[bx]
for gi in T.Pipelined(groups_per_cta, num_stages=num_stages):
group = group_block * groups_per_cta + gi
if group * GROUP < context_len:
# TileLang's pipeliner rejects dynamic loops around smem copies;
# unroll at the Python level instead.
for sp in range(sub_pages):
if (group * sub_pages + sp) * page_size < context_len:
T.copy(
KCache[
PageTable[bx, group * sub_pages + sp],
:,
0,
:,
],
k_shared[sp * page_size : (sp + 1) * page_size, :],
)
T.gemm(
k_shared,
q_shared,
scores,
transpose_B=True,
clear_accum=True,
policy=T.GemmWarpPolicy.FullCol,
)
for token, head in T.Parallel(GROUP, heads):
scores[token, head] = T.max(scores[token, head], 0.0)
T.reduce_sum(scores, reduced, dim=1, clear=True)
for token in T.Parallel(GROUP):
position = group * GROUP + token
if position < context_len:
Logits[bx, position] = reduced[token] / Scale
return kernel
def tilelang_qsa_mqa_prefill(
q: torch.Tensor,
k: torch.Tensor,
row_starts: torch.Tensor,
row_ends: torch.Tensor,
score_scale: Optional[float] = None,
) -> torch.Tensor:
"""Validated TileLang packed prefill kernel with weights removed."""
if not HAS_TILELANG:
raise RuntimeError("TileLang is unavailable")
_validate_q(q)
_validate_k(k)
rows, keys = q.shape[0], k.shape[0]
if not rows or not keys:
logits = torch.zeros((rows, keys), dtype=torch.float32, device=q.device)
return logits.masked_fill_(
torch.ones_like(logits, dtype=torch.bool), -float("inf")
)
heads, head_dim = q.shape[1:]
block_q = max(1, 128 // heads)
padding = (-rows) % block_q
padded_rows = rows + padding
# A torch.cat of the padding rows would copy the whole [rows, keys] fp32 matrix,
# doubling the dominant prefill buffer; allocate pre-padded instead.
logits = torch.zeros((padded_rows, keys), dtype=torch.float32, device=q.device)
q_padded = q.to(torch.bfloat16).contiguous()
starts = row_starts.to(device=q.device, dtype=torch.int32).contiguous()
ends = row_ends.to(device=q.device, dtype=torch.int32).contiguous()
if padding:
q_padded = torch.cat([q_padded, q_padded.new_zeros(padding, heads, head_dim)])
starts = torch.cat([starts, starts[-1:].expand(padding)])
ends = torch.cat([ends, ends[-1:].expand(padding)])
_tilelang_qsa_mqa_prefill_kernel(heads=heads, head_dim=head_dim, block_q=block_q)(
q_padded.reshape(-1, head_dim),
k[:, 0].to(torch.bfloat16).contiguous(),
logits,
starts,
ends,
)
# A leading-dimension slice that retains every column is already
# contiguous, so do not copy this large matrix again when removing padding.
logits = logits[:rows]
logits.div_(score_scale or math.sqrt(head_dim))
_tilelang_qsa_mqa_mask_kernel()(logits, starts[:rows], ends[:rows])
return logits
def tilelang_qsa_mqa_decode(
q: torch.Tensor,
k_cache: torch.Tensor,
page_table: torch.Tensor,
context_lens: torch.Tensor,
max_model_len: int,
score_scale: Optional[float] = None,
) -> torch.Tensor:
"""Validated TileLang paged decode kernel with weights removed."""
if not HAS_TILELANG:
raise RuntimeError("TileLang is unavailable")
_validate_decode_inputs(q, k_cache, page_table, context_lens)
page_size = int(k_cache.shape[1])
if page_size < 8 or 64 % page_size != 0:
raise ValueError(
"TileLang QSA decode requires a compressed page size of "
f"8/16/32/64 (64-row GEMM sub-page packing), got {page_size}"
)
logits = torch.full(
(q.shape[0], max_model_len),
-float("inf"),
dtype=torch.float32,
device=q.device,
)
if not q.shape[0] or not max_model_len:
return logits
# The validated MMA layout requires N (the Q-head dimension) to be a
# multiple of eight. Zero-padding preserves the weight-free head sum.
query_heads, head_dim = q.shape[1:]
kernel_heads = max(8, ((query_heads + 7) // 8) * 8)
q_kernel = q.to(torch.bfloat16)
if kernel_heads != query_heads:
q_kernel = torch.cat(
[
q_kernel,
q_kernel.new_zeros(q.shape[0], kernel_heads - query_heads, head_dim),
],
dim=1,
)
_tilelang_qsa_mqa_decode_kernel(
heads=kernel_heads, head_dim=head_dim, page_size=page_size
)(
q_kernel.unsqueeze(1).contiguous(),
k_cache.to(torch.bfloat16).contiguous(),
page_table.to(device=q.device, dtype=torch.int32).contiguous(),
context_lens.to(device=q.device, dtype=torch.int32).contiguous(),
logits,
float(score_scale or math.sqrt(head_dim)),
)
return logits
def qsa_mqa_prefill(
q: torch.Tensor,
k: torch.Tensor,
row_starts: torch.Tensor,
row_ends: torch.Tensor,
score_scale: Optional[float] = None,
) -> torch.Tensor:
if q.is_cuda and HAS_TILELANG:
return tilelang_qsa_mqa_prefill(q, k, row_starts, row_ends, score_scale)
return torch_qsa_mqa_prefill(q, k, row_starts, row_ends, score_scale)
def qsa_mqa_decode(
q: torch.Tensor,
k_cache: torch.Tensor,
page_table: torch.Tensor,
context_lens: torch.Tensor,
max_model_len: int,
score_scale: Optional[float] = None,
) -> torch.Tensor:
if q.is_cuda and HAS_TILELANG:
return tilelang_qsa_mqa_decode(
q, k_cache, page_table, context_lens, max_model_len, score_scale
)
return torch_qsa_mqa_decode(
q, k_cache, page_table, context_lens, max_model_len, score_scale
)
__all__ = [
"HAS_TILELANG",
"qsa_mqa_decode",
"qsa_mqa_prefill",
"tilelang_qsa_mqa_decode",
"tilelang_qsa_mqa_prefill",
"torch_qsa_mqa_decode",
"torch_qsa_mqa_prefill",
]
python/sglang/srt/layers/attention/qsa/qsa_indexer.py
+633
View File
@@ -0,0 +1,633 @@
"""QSA indexer for Qwen4-Exp checkpoints."""
from __future__ import annotations
from typing import Tuple
import torch
from sglang.srt.layers.attention.qsa.kernel import (
average_pool_qsa_keys,
expand_qsa_block_indices,
qsa_fast_topk,
)
from sglang.srt.layers.attention.qsa.metadata import (
build_group_ring_slots,
build_pending_ring_slots,
build_rope_position_matrix,
)
from sglang.srt.layers.attention.qsa.mqa import qsa_mqa_decode, qsa_mqa_prefill
from sglang.srt.layers.layernorm import GemmaRMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.runner import get_is_capture_mode
# Cap on the fp32 [query_rows, compressed_keys] prefill logits workspace;
# top-k is per row, so tiling rows does not change the selection.
_QSA_PREFILL_LOGITS_BUDGET_BYTES = 128 * 1024 * 1024
def _qsa_prefill_row_chunk_size(rows: int, keys: int, heads: int) -> int:
if rows <= 0 or keys <= 0:
return max(rows, 1)
block_q = max(1, 128 // heads)
bytes_per_row = keys * torch.float32.itemsize
max_padded_rows = max(block_q, _QSA_PREFILL_LOGITS_BUDGET_BYTES // bytes_per_row)
max_padded_rows = max(block_q, max_padded_rows // block_q * block_q)
return min(rows, max_padded_rows)
class QSAIndexer(MultiPlatformOp):
"""Config-driven fused-QK, weight-free sparse-attention indexer."""
def __init__(
self,
config,
layer_id: int,
quant_config=None,
prefix: str = "",
rotary_emb=None,
) -> None:
self._validate_config(config)
super().__init__()
self.layer_id = int(layer_id)
self.index_n_heads = int(config.indexer_n_heads)
self.index_kv_heads = int(config.indexer_kv_heads)
self.index_head_dim = int(config.indexer_head_dim)
self.token_topk = int(config.indexer_budget)
self.compress_ratio = int(config.indexer_compress_ratio)
self.block_topk = self.token_topk // self.compress_ratio
if rotary_emb is None:
raise ValueError("QSAIndexer must reuse its Qwen4-Exp attention RoPE")
self.rotary_emb = rotary_emb
if not 0 < self.rotary_emb.rotary_dim <= self.index_head_dim:
raise ValueError(
"Qwen4-Exp attention RoPE rotary_dim must fit the QSA index head: "
f"{self.rotary_emb.rotary_dim=} {self.index_head_dim=}"
)
self.index_qk_proj = ReplicatedLinear(
config.hidden_size,
(self.index_n_heads + self.index_kv_heads) * self.index_head_dim,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj",
)
self.q_layernorm = GemmaRMSNorm(
self.index_head_dim, eps=getattr(config, "rms_norm_eps", 1e-6)
)
self.k_layernorm = GemmaRMSNorm(
self.index_head_dim, eps=getattr(config, "rms_norm_eps", 1e-6)
)
self._rope_axis_map_cache = None
@staticmethod
def _validate_config(config) -> None:
names = (
"indexer_n_heads",
"indexer_kv_heads",
"indexer_head_dim",
"indexer_budget",
"indexer_compress_ratio",
)
missing = [name for name in names if getattr(config, name, None) is None]
if missing:
raise ValueError(f"QSA config is missing required fields: {missing}")
values = {name: int(getattr(config, name)) for name in names}
if any(value <= 0 for value in values.values()):
raise ValueError(f"QSA config values must be positive: {values}")
if values["indexer_compress_ratio"] < 2:
# DP token-padding rows carry logical length 1, which must never
# reach a compression boundary; ratio >= 2 guarantees that.
raise ValueError(
"QSA requires indexer_compress_ratio >= 2, got "
f"{values['indexer_compress_ratio']}"
)
if values["indexer_kv_heads"] != 1:
raise ValueError("the QSA MQA operators require indexer_kv_heads=1")
if values["indexer_budget"] % values["indexer_compress_ratio"] != 0:
raise ValueError(
"indexer_budget must be divisible by indexer_compress_ratio"
)
block_topk = values["indexer_budget"] // values["indexer_compress_ratio"]
if block_topk not in (512, 2048):
raise ValueError(
"fast_topk_v2 requires indexer_budget / indexer_compress_ratio "
f"to be 512 or 2048, got {block_topk}"
)
def _use_fused_prep(self, tensor: torch.Tensor) -> bool:
"""Whether the fused indexer-prep kernels support this configuration."""
return (
tensor.is_cuda
and tensor.dtype in (torch.bfloat16, torch.float16)
and self.index_head_dim in (64, 128, 256)
and self.rotary_emb.rotary_dim % 2 == 0
and not getattr(self.rotary_emb, "mrope_interleaved_glm", False)
and len(getattr(self.rotary_emb, "mrope_section", None) or ()) in (0, 3)
and getattr(self.rotary_emb, "cos_sin_cache", None) is not None
and self.rotary_emb.cos_sin_cache.is_cuda
and self.rotary_emb.cos_sin_cache.dtype == torch.float32
)
def _rope_axis_map(self, device) -> torch.Tensor:
"""axis_map[i] is the MRoPE position axis whose cos/sin rotary pair i reads."""
cache = self._rope_axis_map_cache
if cache is not None and cache.device == device:
return cache
half = self.rotary_emb.rotary_dim // 2
section = getattr(self.rotary_emb, "mrope_section", None) or None
axis_map = torch.zeros(half, dtype=torch.int32)
if section is not None:
s0, s1, s2 = (int(v) for v in section)
if getattr(self.rotary_emb, "mrope_interleaved", False):
pair = torch.arange(half, dtype=torch.int32)
axis_map[((pair % 3) == 1) & (pair < s1 * 3)] = 1
axis_map[((pair % 3) == 2) & (pair < s2 * 3)] = 2
else:
axis_map[s0 : s0 + s1] = 1
axis_map[s0 + s1 :] = 2
self._rope_axis_map_cache = axis_map.to(device)
return self._rope_axis_map_cache
def project_qk(
self,
hidden_states: torch.Tensor,
positions: torch.Tensor,
pool=None,
cache_loc: torch.Tensor | None = None,
q_heads_padded: int | None = None,
) -> Tuple[torch.Tensor, torch.Tensor, bool]:
qk, _ = self.index_qk_proj(hidden_states)
token_k = qk[:, self.index_n_heads * self.index_head_dim :].reshape(
-1, self.index_kv_heads, self.index_head_dim
)
if (
pool is not None
and cache_loc is not None
and qk.shape[0] > 0
and self._use_fused_prep(qk)
):
from sglang.kernels.ops.attention.qsa_indexer import (
qsa_index_q_norm_rope_store,
)
if not get_is_capture_mode() and hasattr(
self.rotary_emb, "_ensure_cos_sin_cache_length"
):
self.rotary_emb._ensure_cos_sin_cache_length(
int(positions.max().item())
)
key_state_buffer = pool.get_qsa_key_state_buffer(self.layer_id)
q = qsa_index_q_norm_rope_store(
qk,
positions.long(),
self.rotary_emb.cos_sin_cache,
self._rope_axis_map(qk.device),
self.q_layernorm.weight.data,
cache_loc[: qk.shape[0]].long(),
key_state_buffer.view(key_state_buffer.shape[0], -1),
pool.qsa_rope_position_buffer,
self.index_n_heads,
self.rotary_emb.rotary_dim,
self.q_layernorm.variance_epsilon,
self.rotary_emb.is_neox_style,
q_heads_padded=q_heads_padded,
)
return q, token_k, True
q_raw = qk[:, : self.index_n_heads * self.index_head_dim]
q = self.q_layernorm(q_raw.reshape(-1, self.index_head_dim)).reshape(
-1, self.index_n_heads, self.index_head_dim
)
q = self.apply_rope(positions, q)
return q, token_k, False
def normalize_compressed_keys(
self, compressed_keys: torch.Tensor, block_positions: torch.Tensor
) -> torch.Tensor:
normalized = self.k_layernorm(
compressed_keys.reshape(-1, self.index_head_dim)
).reshape(-1, self.index_kv_heads, self.index_head_dim)
return self.apply_rope(block_positions, normalized)
def _use_fused_compress(self, pool) -> bool:
return getattr(
pool, "qsa_rope_position_buffer", None
) is not None and self._use_fused_prep(
pool.get_qsa_key_state_buffer(self.layer_id)
)
def _fused_compress_store(
self,
pool,
group_locs: torch.Tensor,
write_locs: torch.Tensor,
source_keys: torch.Tensor | None = None,
source_rope: torch.Tensor | None = None,
) -> None:
"""Fused mean -> gemma norm -> MRoPE -> compressed-cache store;
a None source_keys/source_rope reads the members from the pending ring."""
from sglang.kernels.ops.attention.qsa_indexer import (
qsa_index_k_compress_store,
)
if source_keys is None:
source_keys = pool.get_qsa_key_state_buffer(self.layer_id)
if source_rope is None:
source_rope = pool.qsa_rope_position_buffer
compressed_buffer = pool.get_qsa_compressed_k_buffer(self.layer_id)
qsa_index_k_compress_store(
source_keys.reshape(source_keys.shape[0], -1)
.contiguous()
.to(pool.index_state_dtype),
group_locs.to(torch.int32),
source_rope,
self.rotary_emb.cos_sin_cache,
self._rope_axis_map(source_keys.device),
self.k_layernorm.weight.data,
write_locs.to(torch.int32),
compressed_buffer.view(compressed_buffer.shape[0], -1),
self.compress_ratio,
self.rotary_emb.rotary_dim,
self.k_layernorm.variance_epsilon,
self.rotary_emb.is_neox_style,
)
def _pending_ring_slots(
self, metadata, logical_positions: torch.Tensor, is_extend: bool
) -> torch.Tensor:
return build_pending_ring_slots(
token_to_batch_idx=metadata.token_to_batch_idx,
req_pool_indices=metadata.req_pool_indices,
sequence_lengths=metadata.sequence_lengths,
logical_positions=logical_positions,
compress_ratio=self.compress_ratio,
is_extend=is_extend,
)
def _group_ring_slots(
self, metadata, group_end_positions: torch.Tensor, sequence_ids: torch.Tensor
) -> torch.Tensor:
return build_group_ring_slots(
req_pool_indices=metadata.req_pool_indices,
group_end_positions=group_end_positions,
sequence_ids=sequence_ids,
compress_ratio=self.compress_ratio,
)
def update_key_state_and_compress(
self,
token_k: torch.Tensor,
logical_positions: torch.Tensor,
rope_positions: torch.Tensor,
metadata,
state_slots: torch.Tensor | None = None,
state_stored: bool = False,
) -> None:
"""Store the pending-group ring and compress each completed group."""
pool = metadata.token_to_kv_pool
is_extend = metadata.compress_member_rows is not None
if not state_stored:
if state_slots is None:
state_slots = self._pending_ring_slots(
metadata, logical_positions, is_extend
)
pool.set_qsa_key_state_buffer(
self.layer_id, state_slots[: token_k.shape[0]], token_k
)
pool.set_qsa_rope_position_buffer(
state_slots[: token_k.shape[0]], rope_positions
)
if metadata.is_cuda_graph:
self._compress_decode_cuda_graph(metadata)
return
if metadata.write_locs is None:
raise RuntimeError(
"QSA metadata is missing the precomputed write plan; the "
"sparse-attention backend derives it from the batch lengths"
)
if metadata.write_locs.numel() == 0:
return
group_end_positions = metadata.compress_group_positions.long()
compressed_locs = metadata.write_locs
if is_extend:
# Extend chunks are group-aligned; each planned group lies in this forward,
# so read its members from the packed chunk tensors.
member_rows = metadata.compress_member_rows.long()
group_locs = member_rows[:, None] + torch.arange(
self.compress_ratio, device=member_rows.device, dtype=torch.long
)
source_keys = token_k
source_rope = metadata.extend_rope_matrix
if source_rope is None:
source_rope = build_rope_position_matrix(
rope_positions, token_k.shape[0]
)
else:
# Paged eager rows (speculative fallback) complete at most one
# group each; its members are exactly the pending ring window.
group_locs = metadata.compress_group_ring_locs
if group_locs is None:
group_locs = self._group_ring_slots(
metadata,
group_end_positions,
metadata.compress_sequence_ids.long(),
)
source_keys = pool.get_qsa_key_state_buffer(self.layer_id)
source_rope = pool.qsa_rope_position_buffer
if self._use_fused_compress(pool):
self._fused_compress_store(
pool,
group_locs,
compressed_locs,
source_keys=source_keys,
source_rope=source_rope,
)
return
key_groups = source_keys[group_locs]
pooled = average_pool_qsa_keys(key_groups)
compressed_rope_positions = self._rope_from_matrix(
source_rope[group_locs[:, 0]]
)
normalized = self.normalize_compressed_keys(pooled, compressed_rope_positions)
pool.set_qsa_compressed_k_buffer(self.layer_id, compressed_locs, normalized)
def _compress_decode_cuda_graph(self, metadata) -> None:
"""Fixed-shape graph-replay compression; non-boundary rows write slot 0."""
if metadata.graph_write_locs is None or metadata.graph_ring_group_locs is None:
raise RuntimeError("QSA CUDA graph compression metadata is incomplete")
pool = metadata.token_to_kv_pool
group_locs = metadata.graph_ring_group_locs
if self._use_fused_compress(pool):
self._fused_compress_store(
pool,
group_locs,
metadata.graph_write_locs,
)
return
key_groups = pool.get_qsa_key_state_buffer(self.layer_id)[group_locs]
compressed = average_pool_qsa_keys(key_groups)
compressed_rope_positions = self._rope_from_matrix(
pool.qsa_rope_position_buffer[group_locs[:, 0]]
)
compressed = self.normalize_compressed_keys(
compressed, compressed_rope_positions
)
pool.set_qsa_compressed_k_buffer(
self.layer_id, metadata.graph_write_locs, compressed.contiguous()
)
def _rope_from_matrix(self, positions: torch.Tensor) -> torch.Tensor:
"""[n, 3] slot coordinates -> the layout apply_rope expects."""
positions = positions.transpose(0, 1)
if not getattr(self.rotary_emb, "mrope_section", None):
return positions[0]
return positions
def apply_rope(self, positions: torch.Tensor, tensor: torch.Tensor) -> torch.Tensor:
if tensor.numel() == 0:
return tensor
positions = positions.long()
num_positions = (
positions.shape[-1] if positions.ndim == 2 else positions.numel()
)
if num_positions != tensor.shape[0]:
raise ValueError("QSA RoPE positions must match the token dimension")
if not get_is_capture_mode() and hasattr(
self.rotary_emb, "_ensure_cos_sin_cache_length"
):
self.rotary_emb._ensure_cos_sin_cache_length(int(positions.max().item()))
# position_cos/position_sin repeat cos/sin to the full rotary width;
# apply_rotary_emb consumes one half.
self.rotary_emb.get_cos_sin_with_position(positions)
rotary_dim = self.rotary_emb.rotary_dim
half_rotary_dim = rotary_dim // 2
cos = self.rotary_emb.position_cos.reshape(num_positions, -1)[
:, :half_rotary_dim
]
sin = self.rotary_emb.position_sin.reshape(num_positions, -1)[
:, :half_rotary_dim
]
rotated = apply_rotary_emb(
tensor[..., :rotary_dim],
cos,
sin,
self.rotary_emb.is_neox_style,
)
return torch.cat([rotated, tensor[..., rotary_dim:]], dim=-1)
def select_prefill_tokens(
self,
q: torch.Tensor,
compressed_keys: torch.Tensor,
row_starts: torch.Tensor,
row_ends: torch.Tensor,
query_positions: torch.Tensor,
sequence_lengths_for_rows: torch.Tensor,
) -> torch.Tensor:
rows = q.shape[0]
output = torch.empty(
(rows, self.token_topk + self.compress_ratio - 1),
dtype=torch.int32,
device=q.device,
)
if rows == 0:
return output
row_chunk_size = _qsa_prefill_row_chunk_size(
rows, compressed_keys.shape[0], q.shape[1]
)
for row_start in range(0, rows, row_chunk_size):
row_end = min(row_start + row_chunk_size, rows)
chunk_slice = slice(row_start, row_end)
if compressed_keys.shape[0] == 0:
block_indices = torch.full(
(row_end - row_start, self.block_topk),
-1,
dtype=torch.int32,
device=q.device,
)
logits = None
else:
logits = qsa_mqa_prefill(
q[chunk_slice],
compressed_keys,
row_starts[chunk_slice],
row_ends[chunk_slice],
)
block_indices = qsa_fast_topk(
logits,
row_starts[chunk_slice],
row_ends[chunk_slice],
topk=self.block_topk,
)
selected = expand_qsa_block_indices(
block_indices,
query_positions[chunk_slice],
sequence_lengths_for_rows[chunk_slice],
compress_ratio=self.compress_ratio,
token_topk=self.token_topk,
)
output[chunk_slice].copy_(selected)
del logits, block_indices, selected
return output
def select_decode_tokens(
self,
q: torch.Tensor,
compressed_cache: torch.Tensor,
compressed_page_table: torch.Tensor,
compressed_lengths: torch.Tensor,
max_model_len: int,
query_positions: torch.Tensor,
sequence_lengths: torch.Tensor,
) -> torch.Tensor:
logits = qsa_mqa_decode(
q,
compressed_cache,
compressed_page_table,
compressed_lengths,
max_model_len,
)
if logits.is_cuda and self.block_topk == 512:
# Decode rows start at zero, so compressed lengths double as row lengths;
# skip the generic zero-fill + subtract.
from sglang.kernels.ops.elementwise.fast_topk import fast_topk
block_indices = fast_topk(
logits,
compressed_lengths.to(torch.int32),
topk=self.block_topk,
row_starts=None,
)
else:
row_starts = torch.zeros_like(compressed_lengths, dtype=torch.int32)
block_indices = qsa_fast_topk(
logits, row_starts, compressed_lengths, topk=self.block_topk
)
return expand_qsa_block_indices(
block_indices,
query_positions,
sequence_lengths,
compress_ratio=self.compress_ratio,
token_topk=self.token_topk,
)
def forward_cuda(
self,
hidden_states: torch.Tensor,
positions: torch.Tensor,
forward_batch,
indexer_metadata,
) -> torch.Tensor:
forward_mode = forward_batch.forward_mode
is_target_verify = getattr(forward_mode, "is_target_verify", lambda: False)()
is_draft_extend = getattr(forward_mode, "is_draft_extend_v2", lambda: False)()
if forward_mode.is_decode() or is_target_verify or is_draft_extend:
# EAGLE/MTP may advance the model's RoPE coordinate independently
# from the physical paged-KV position. Compression and sparse
# selection must use the latter; otherwise a draft step can index
# token_slot_table[:, seq_len] one past the valid range.
logical_positions = indexer_metadata.decode_logical_positions
if logical_positions is None:
logical_positions = indexer_metadata.get_seqlens_expanded() - 1
else:
logical_positions = getattr(forward_batch, "positions", None)
if logical_positions is None:
logical_positions = positions[0] if positions.ndim == 2 else positions
logical_positions = logical_positions.flatten()
# DP MAX_LEN padding adds token rows without assigning them to a
# request. token_to_batch_idx is the source of truth for semantic rows.
num_valid_tokens = indexer_metadata.get_token_to_batch_idx().numel()
if logical_positions.numel() < num_valid_tokens:
raise ValueError(
"QSA logical positions are shorter than the request mapping: "
f"positions={logical_positions.numel()}, mapping={num_valid_tokens}"
)
if hidden_states.shape[0] < num_valid_tokens:
raise ValueError(
"QSA hidden states are shorter than the request mapping: "
f"hidden={hidden_states.shape[0]}, mapping={num_valid_tokens}"
)
position_tokens = (
positions.shape[-1] if positions.ndim == 2 else positions.numel()
)
if position_tokens < num_valid_tokens:
raise ValueError(
"QSA RoPE positions are shorter than the request mapping: "
f"positions={position_tokens}, mapping={num_valid_tokens}"
)
logical_positions = logical_positions[:num_valid_tokens]
hidden_states = hidden_states[:num_valid_tokens]
positions = (
positions[:, :num_valid_tokens]
if positions.ndim == 2
else positions[:num_valid_tokens]
)
state_slots = indexer_metadata.pending_ring_slots
if state_slots is None:
state_slots = self._pending_ring_slots(
indexer_metadata,
logical_positions,
indexer_metadata.compress_member_rows is not None,
)
q, token_k, state_stored = self.project_qk(
hidden_states,
positions,
pool=indexer_metadata.token_to_kv_pool,
cache_loc=state_slots,
q_heads_padded=(
# The tilelang decode MQA kernel needs query heads in multiples of 8.
((self.index_n_heads + 7) // 8) * 8
if (forward_mode.is_decode() or is_target_verify or is_draft_extend)
else None
),
)
self.update_key_state_and_compress(
token_k,
logical_positions,
positions,
indexer_metadata,
state_slots=state_slots,
state_stored=state_stored,
)
if forward_mode.is_decode() or is_target_verify or is_draft_extend:
compressed_cache, page_table, compressed_lengths, max_model_len = (
indexer_metadata.get_decode_mqa_inputs(self.layer_id)
)
return self.select_decode_tokens(
q,
compressed_cache,
page_table,
compressed_lengths,
max_model_len,
logical_positions,
indexer_metadata.get_seqlens_int32(),
)
compressed_keys, row_starts, row_ends, sequence_lengths = (
indexer_metadata.get_prefill_mqa_inputs(self.layer_id, logical_positions)
)
query_sequence_ids = indexer_metadata.get_token_to_batch_idx()
row_sequence_lengths = sequence_lengths.index_select(
0, query_sequence_ids.long()
)
return self.select_prefill_tokens(
q,
compressed_keys,
row_starts,
row_ends,
logical_positions,
row_sequence_lengths,
)
__all__ = [
"QSAIndexer",
]
python/sglang/srt/layers/attention/qsa/sparse_attn.py
+453
View File
@@ -0,0 +1,453 @@
"""Validated sparse GQA operators migrated from the QSA reference branch."""
from typing import Optional
import torch
import triton
import triton.language as tl
_H20_CONFIGS = [
(32, (32, 8, 2)),
(64, (64, 8, 2)),
(1024, (32, 4, 2)),
(float("inf"), (16, 1, 2)),
]
_L20_CONFIGS = [
(32, (32, 8, 2)),
(64, (64, 8, 2)),
(128, (64, 4, 2)),
(512, (32, 4, 2)),
(float("inf"), (16, 1, 2)),
]
def _get_best_config(total_q: int):
table = _H20_CONFIGS if "H20" in torch.cuda.get_device_name(0) else _L20_CONFIGS
return next(cfg for limit, cfg in table if total_q <= limit)
@triton.jit
def _sparse_gqa_prefill(
q,
k,
v,
out,
indices,
cu_seqlens,
scale,
topk,
sq_m: tl.constexpr,
sq_h: tl.constexpr,
sq_d: tl.constexpr,
sk_n: tl.constexpr,
sk_h: tl.constexpr,
sk_d: tl.constexpr,
sv_n: tl.constexpr,
sv_h: tl.constexpr,
sv_d: tl.constexpr,
so_m: tl.constexpr,
so_h: tl.constexpr,
so_d: tl.constexpr,
si_m: tl.constexpr,
si_g: tl.constexpr,
si_n: tl.constexpr,
NUM_KV_HEADS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
HEAD_DIM: tl.constexpr,
):
batch_group = tl.program_id(1)
group = batch_group % NUM_KV_HEADS
batch = batch_group // NUM_KV_HEADS
seq_start = tl.load(cu_seqlens + batch).to(tl.int64)
seq_end = tl.load(cu_seqlens + batch + 1).to(tl.int64)
query_relative = tl.program_id(0).to(tl.int64)
query = seq_start + query_relative
if query >= seq_end:
return
row_topk = tl.minimum(topk, query_relative + 1)
row_limit = tl.minimum(topk, ((row_topk + BLOCK_N - 1) // BLOCK_N) * BLOCK_N)
offs_h = tl.arange(0, BLOCK_M)
offs_d = tl.arange(0, HEAD_DIM)
head_start = group * GROUP_SIZE
q_values = tl.load(
q
+ query * sq_m
+ (head_start + offs_h[:, None]) * sq_h
+ offs_d[None, :] * sq_d,
mask=(offs_h < GROUP_SIZE)[:, None],
other=0.0,
)
q_values = (q_values * scale * 1.4426950408).to(q_values.dtype)
k_base = k + seq_start * sk_n + group * sk_h
v_base = v + seq_start * sv_n + group * sv_h
idx_row = indices + query * si_m + group * si_g
max_value = tl.full([BLOCK_M], -float("inf"), tl.float32)
normalizer = tl.zeros([BLOCK_M], tl.float32)
accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
offs_n = tl.arange(0, BLOCK_N)
for start in range(0, row_limit, BLOCK_N):
current = start + offs_n
token = tl.load(idx_row + current * si_n, mask=current < topk, other=-1)
valid = token >= 0
keys = tl.load(
k_base + token[None, :] * sk_n + offs_d[:, None] * sk_d,
mask=valid[None, :],
other=0.0,
)
values = tl.load(
v_base + token[:, None] * sv_n + offs_d[None, :] * sv_d,
mask=valid[:, None],
other=0.0,
)
scores = tl.where(valid[None, :], tl.dot(q_values, keys), -float("inf"))
next_max = tl.maximum(max_value, tl.max(scores, 1))
alpha = tl.math.exp2(max_value - next_max)
probabilities = tl.math.exp2(scores - next_max[:, None])
accumulator = tl.dot(
probabilities.to(values.dtype), values, accumulator * alpha[:, None]
)
normalizer = normalizer * alpha + tl.sum(probabilities, 1)
max_value = next_max
output = accumulator / normalizer[:, None]
tl.store(
out
+ query * so_m
+ (head_start + offs_h[:, None]) * so_h
+ offs_d[None, :] * so_d,
output,
mask=(offs_h < GROUP_SIZE)[:, None],
)
def sparse_gqa_fwd_interface_triton(q, k, v, max_seqlen_k, indices, cu_seqlens, scale):
total_q, num_q_heads, head_dim = q.shape
num_kv_heads = k.shape[1]
group_size = num_q_heads // num_kv_heads
block_m = max(16, triton.next_power_of_2(group_size))
block_n, warps, stages = _get_best_config(total_q)
out = torch.empty_like(q)
_sparse_gqa_prefill[(max_seqlen_k, (cu_seqlens.shape[0] - 1) * num_kv_heads)](
q,
k,
v,
out,
indices,
cu_seqlens,
scale,
indices.shape[-1],
q.stride(0),
q.stride(1),
q.stride(2),
k.stride(0),
k.stride(1),
k.stride(2),
v.stride(0),
v.stride(1),
v.stride(2),
out.stride(0),
out.stride(1),
out.stride(2),
indices.stride(0),
indices.stride(1) if indices.ndim == 3 else 0,
indices.stride(2) if indices.ndim == 3 else indices.stride(1),
NUM_KV_HEADS=num_kv_heads,
GROUP_SIZE=group_size,
BLOCK_M=block_m,
BLOCK_N=block_n,
HEAD_DIM=head_dim,
num_warps=warps,
num_stages=stages,
)
return out
@triton.jit
def _sparse_gqa_chunk_prefill(
q,
k,
v,
out,
indices,
cu_q,
cu_k,
kv_lens,
scale,
topk,
sq_m: tl.constexpr,
sq_h: tl.constexpr,
sq_d: tl.constexpr,
sk_n: tl.constexpr,
sk_h: tl.constexpr,
sk_d: tl.constexpr,
sv_n: tl.constexpr,
sv_h: tl.constexpr,
sv_d: tl.constexpr,
so_m: tl.constexpr,
so_h: tl.constexpr,
so_d: tl.constexpr,
si_m: tl.constexpr,
si_g: tl.constexpr,
si_n: tl.constexpr,
NUM_KV_HEADS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
HEAD_DIM: tl.constexpr,
):
query_relative = tl.program_id(0).to(tl.int64)
batch_group = tl.program_id(1)
group = batch_group % NUM_KV_HEADS
batch = batch_group // NUM_KV_HEADS
q_start = tl.load(cu_q + batch)
q_end = tl.load(cu_q + batch + 1)
query = (q_start + query_relative).to(tl.int64)
if query >= q_end:
return
k_start = tl.load(cu_k + batch).to(tl.int64)
kv_len = tl.load(kv_lens + batch).to(tl.int64)
visible = query_relative + kv_len - (q_end - q_start) + 1
row_topk = tl.minimum(topk, visible)
row_limit = tl.minimum(topk, ((row_topk + BLOCK_N - 1) // BLOCK_N) * BLOCK_N)
offs_h = tl.arange(0, BLOCK_M)
offs_d = tl.arange(0, HEAD_DIM)
q_values = tl.load(
q
+ query * sq_m
+ (group * GROUP_SIZE + offs_h[:, None]) * sq_h
+ offs_d[None, :] * sq_d,
mask=(offs_h < GROUP_SIZE)[:, None],
other=0.0,
)
q_values = (q_values * scale * 1.4426950408).to(q_values.dtype)
k_base = k + k_start * sk_n + group * sk_h
v_base = v + k_start * sv_n + group * sv_h
idx_row = indices + query * si_m + group * si_g
max_value = tl.full([BLOCK_M], -float("inf"), tl.float32)
normalizer = tl.zeros([BLOCK_M], tl.float32)
accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
offs_n = tl.arange(0, BLOCK_N)
for start in range(0, row_limit, BLOCK_N):
current = start + offs_n
token = tl.load(idx_row + current * si_n, mask=current < topk, other=-1)
valid = token >= 0
keys = tl.load(
k_base + token[None, :] * sk_n + offs_d[:, None] * sk_d,
mask=valid[None, :],
other=0.0,
)
values = tl.load(
v_base + token[:, None] * sv_n + offs_d[None, :] * sv_d,
mask=valid[:, None],
other=0.0,
)
scores = tl.where(valid[None, :], tl.dot(q_values, keys), -float("inf"))
next_max = tl.maximum(max_value, tl.max(scores, 1))
alpha = tl.math.exp2(max_value - next_max)
probabilities = tl.math.exp2(scores - next_max[:, None])
accumulator = tl.dot(
probabilities.to(values.dtype), values, accumulator * alpha[:, None]
)
normalizer = normalizer * alpha + tl.sum(probabilities, 1)
max_value = next_max
output = accumulator / normalizer[:, None]
tl.store(
out
+ query * so_m
+ (group * GROUP_SIZE + offs_h[:, None]) * so_h
+ offs_d[None, :] * so_d,
output,
mask=(offs_h < GROUP_SIZE)[:, None],
)
def sparse_gqa_fwd_interface_triton_ck(q, k, v, indices, cu_q, cu_k, kv_lens, scale):
k, v = k.contiguous(), v.contiguous()
total_q, num_q_heads, head_dim = q.shape
num_kv_heads = k.shape[1]
group_size = num_q_heads // num_kv_heads
max_q = int((cu_q[1:] - cu_q[:-1]).max().item())
block_m = max(16, triton.next_power_of_2(group_size))
block_n, warps, stages = _get_best_config(total_q)
out = torch.empty_like(q)
_sparse_gqa_chunk_prefill[(max_q, (cu_q.shape[0] - 1) * num_kv_heads)](
q,
k,
v,
out,
indices,
cu_q,
cu_k,
kv_lens,
scale,
indices.shape[-1],
q.stride(0),
q.stride(1),
q.stride(2),
k.stride(0),
k.stride(1),
k.stride(2),
v.stride(0),
v.stride(1),
v.stride(2),
out.stride(0),
out.stride(1),
out.stride(2),
indices.stride(0),
indices.stride(1) if indices.ndim == 3 else 0,
indices.stride(2) if indices.ndim == 3 else indices.stride(1),
NUM_KV_HEADS=num_kv_heads,
GROUP_SIZE=group_size,
BLOCK_M=block_m,
BLOCK_N=block_n,
HEAD_DIM=head_dim,
num_warps=warps,
num_stages=stages,
)
return out
@triton.jit
def _fa2_valid_counts(
seq_lens,
indices,
counts,
topk: tl.constexpr,
stride_i: tl.constexpr,
BLOCK_TOPK: tl.constexpr,
):
row = tl.program_id(0)
cols = tl.arange(0, BLOCK_TOPK)
length = tl.load(seq_lens + row)
positions = tl.load(
indices + row * stride_i + cols,
mask=cols < topk,
other=-1,
)
valid = (positions >= 0) & (positions < length)
tl.store(counts + row, tl.sum(valid.to(tl.int32), axis=0))
@triton.jit
def _fa2_prefix_sum(counts, cu_k, batch, BLOCK_B: tl.constexpr):
rows = tl.arange(0, BLOCK_B)
valid_rows = rows < batch
row_counts = tl.load(counts + rows, mask=valid_rows, other=0)
tl.store(cu_k, 0)
tl.store(cu_k + rows + 1, tl.cumsum(row_counts, 0), mask=valid_rows)
def qwen_sparse_fa2_cu_seqlens_triton(
seq_lens, indices, counts, cu_k, batch, topk, block_b: Optional[int] = None
):
block_b = block_b or triton.next_power_of_2(batch)
# One request per program: Triton caps a tile at 1M elements,
# which [next_pow2(topk), next_pow2(batch)] exceeds at topk=2051, batch=512.
_fa2_valid_counts[(batch,)](
seq_lens,
indices,
counts,
topk,
indices.stride(0),
BLOCK_TOPK=triton.next_power_of_2(topk),
num_warps=8,
)
# Prefix sum is only over the batch dimension and remains a small 1-D
# tensor, including during CUDA graph capture.
_fa2_prefix_sum[(1,)](
counts,
cu_k,
batch,
BLOCK_B=block_b,
num_warps=8,
)
@triton.jit
def _compact_kv(
k,
v,
req_to_token,
req_indices,
indices,
seq_lens,
cu_k,
out_k,
out_v,
topk: tl.constexpr,
heads: tl.constexpr,
dim: tl.constexpr,
req_stride: tl.constexpr,
idx_stride: tl.constexpr,
BLOCK_TOPK: tl.constexpr,
BLOCK_D: tl.constexpr,
):
batch, head, block = tl.program_id(0), tl.program_id(1), tl.program_id(2)
cols = block * BLOCK_TOPK + tl.arange(0, BLOCK_TOPK)
dims = tl.arange(0, BLOCK_D)
length = tl.load(seq_lens + batch)
req = tl.load(req_indices + batch)
pack_start = tl.load(cu_k + batch)
valid_count = tl.load(cu_k + batch + 1) - pack_start
positions = tl.load(indices + batch * idx_stride + cols, mask=cols < topk, other=-1)
valid = (cols < valid_count) & (positions >= 0) & (positions < length)
slots = tl.load(
req_to_token + req * req_stride + tl.where(valid, positions, 0),
mask=valid,
other=0,
)
src = slots[:, None] * heads * dim + head * dim + dims[None, :]
dst = (pack_start + cols)[:, None] * heads * dim + head * dim + dims[None, :]
mask = valid[:, None] & (dims[None, :] < dim)
tl.store(out_k + dst, tl.load(k + src, mask=mask, other=0.0), mask=mask)
tl.store(out_v + dst, tl.load(v + src, mask=mask, other=0.0), mask=mask)
def qwen_sparse_valid_counts_triton(seq_lens, indices, counts, batch, topk):
"""Valid-count pass alone, without the packed cu_seqlens prefix sum."""
_fa2_valid_counts[(batch,)](
seq_lens,
indices,
counts,
topk,
indices.stride(0),
BLOCK_TOPK=triton.next_power_of_2(topk),
num_warps=8,
)
def qwen_sparse_kv_extraction_compact_triton(
k, v, req_to_token, req_indices, indices, seq_lens, cu_k, out_k, out_v, batch, topk
):
_, heads, dim = k.shape
block_topk = 16
_compact_kv[(batch, heads, triton.cdiv(topk, block_topk))](
k,
v,
req_to_token,
req_indices,
indices,
seq_lens,
cu_k,
out_k,
out_v,
topk,
heads,
dim,
req_to_token.stride(0),
indices.stride(0),
BLOCK_TOPK=block_topk,
BLOCK_D=triton.next_power_of_2(dim),
num_warps=8,
)
__all__ = [
"qwen_sparse_fa2_cu_seqlens_triton",
"qwen_sparse_valid_counts_triton",
"qwen_sparse_kv_extraction_compact_triton",
"sparse_gqa_fwd_interface_triton",
"sparse_gqa_fwd_interface_triton_ck",
]
python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py
+1764
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/hc_mix_triton.py
+219
View File
@@ -0,0 +1,219 @@
"""Fused HC low-rank mix for decode-size batches.
One persistent kernel replaces the five-kernel `GatedResidual._mix_compute` chain.
One CTA per SM keeps every CTA resident, so the software grid barrier cannot deadlock;
the last CTA to finish resets the barrier counters,
so a captured CUDA graph replays with them in their initial state.
Row counts beyond ``_FUSED_MIX_MAX_ROWS`` stay on the torch.compile path.
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl
_FUSED_MIX_MAX_ROWS = 16
@triton.jit
def _grid_barrier(counter_ptr, num_ctas):
tl.atomic_add(counter_ptr, 1, sem="acq_rel", scope="gpu")
while tl.atomic_add(counter_ptr, 0, sem="acq_rel", scope="gpu") < num_ctas:
pass
@triton.jit
def _hc_mix_persistent_kernel(
x_ptr,
w_down_ptr,
w_up_ptr,
t_raw_ptr,
out_ptr,
counters_ptr,
K,
LOWRANK,
HS,
num_rows,
num_ctas,
inv_hc,
ROWS: tl.constexpr,
HC: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
BLOCK_J: tl.constexpr,
BLOCK_R: tl.constexpr,
):
pid = tl.program_id(0)
offs_m = tl.arange(0, ROWS)
mask_m = offs_m < num_rows
zero_span = ROWS * LOWRANK
offs_z = tl.arange(0, 256)
for z0 in range(pid * 256, zero_span, num_ctas * 256):
idx = z0 + offs_z
tl.store(t_raw_ptr + idx, 0.0, mask=idx < zero_span)
_grid_barrier(counters_ptr + 0, num_ctas)
offs_k = tl.arange(0, BLOCK_K)
offs_n = tl.arange(0, BLOCK_N)
n_blocks = tl.cdiv(LOWRANK, BLOCK_N)
k_chunks = tl.cdiv(K, BLOCK_K)
for tile in range(pid, n_blocks * k_chunks, num_ctas):
nb = tile % n_blocks
kc = tile // n_blocks
n = nb * BLOCK_N + offs_n
k = kc * BLOCK_K + offs_k
mask_n = n < LOWRANK
xt = tl.load(
x_ptr + offs_m[:, None] * K + k[None, :],
mask=mask_m[:, None],
other=0.0,
)
w = tl.load(
w_down_ptr + n[:, None] * K + k[None, :],
mask=mask_n[:, None],
other=0.0,
)
acc = tl.dot(xt, tl.trans(w))
tl.atomic_add(
t_raw_ptr + offs_m[:, None] * LOWRANK + n[None, :],
acc,
mask=mask_n[None, :],
sem="relaxed",
scope="gpu",
)
_grid_barrier(counters_ptr + 1, num_ctas)
offs_j = tl.arange(0, BLOCK_J)
offs_r = tl.arange(0, BLOCK_R)
offs_g = tl.arange(0, HC)
j_blocks = tl.cdiv(HS, BLOCK_J)
for jb in range(pid, j_blocks, num_ctas):
j = jb * BLOCK_J + offs_j
mask_j = j < HS
gj = offs_g[:, None] * HS + j[None, :]
gj_flat = tl.reshape(gj, (HC * BLOCK_J,))
mask_gj = tl.reshape(
tl.broadcast_to(mask_j[None, :], (HC, BLOCK_J)), (HC * BLOCK_J,)
)
acc = tl.zeros((ROWS, HC * BLOCK_J), dtype=tl.float32)
for r0 in range(0, LOWRANK, BLOCK_R):
r = r0 + offs_r
mask_r = r < LOWRANK
a = tl.load(
t_raw_ptr + offs_m[:, None] * LOWRANK + r[None, :],
mask=mask_r[None, :],
other=0.0,
)
a = a * inv_hc
t = (a * tl.sigmoid(a)).to(x_ptr.dtype.element_ty)
w = tl.load(
w_up_ptr + gj_flat[:, None] * LOWRANK + r[None, :],
mask=mask_gj[:, None] & mask_r[None, :],
other=0.0,
)
acc = tl.dot(t, tl.trans(w), acc)
gate = tl.sigmoid(tl.reshape(acc, (ROWS, HC, BLOCK_J)))
xg = tl.load(
x_ptr
+ offs_m[:, None, None] * (HC * HS)
+ offs_g[None, :, None] * HS
+ j[None, None, :],
mask=mask_m[:, None, None] & mask_j[None, None, :],
other=0.0,
).to(tl.float32)
out = tl.sum(gate * xg, axis=1) * inv_hc
tl.store(
out_ptr + offs_m[:, None] * HS + j[None, :],
out.to(out_ptr.dtype.element_ty),
mask=mask_m[:, None] & mask_j[None, :],
)
ticket = tl.atomic_add(counters_ptr + 2, 1, sem="acq_rel", scope="gpu")
if ticket == num_ctas - 1:
tl.store(counters_ptr + 0, 0)
tl.store(counters_ptr + 1, 0)
tl.store(counters_ptr + 2, 0)
_counters_cache = {}
def _get_counters(device: torch.device) -> torch.Tensor:
buf = _counters_cache.get(device)
if buf is None:
buf = torch.zeros(3, dtype=torch.int32, device=device)
_counters_cache[device] = buf
return buf
def _deterministic_inference() -> bool:
from sglang.srt.runtime_context import get_exec
try:
exec_cfg = get_exec()
except ValueError:
return False
return bool(exec_cfg.deterministic.enable_deterministic_inference)
def fused_hc_mix_supported(
hyper_input_normed: torch.Tensor, w_down: torch.Tensor, w_up: torch.Tensor
) -> bool:
# The persistent kernel accumulates the down projection with
# device-scope atomics, so summation order varies across replays.
if _deterministic_inference():
return False
return (
hyper_input_normed.is_cuda
and hyper_input_normed.dtype in (torch.bfloat16, torch.float16)
and w_down.dtype == hyper_input_normed.dtype
and w_up.dtype == hyper_input_normed.dtype
and hyper_input_normed.shape[0] <= _FUSED_MIX_MAX_ROWS
and hyper_input_normed.dim() == 2
and hyper_input_normed.shape[1] % 2048 == 0
and hyper_input_normed.is_contiguous()
and w_down.is_contiguous()
and w_up.is_contiguous()
)
def fused_hc_mix(
hyper_input_normed: torch.Tensor,
w_down: torch.Tensor,
w_up: torch.Tensor,
hc: int,
hs: int,
) -> torch.Tensor:
rows, k = hyper_input_normed.shape
lowrank = w_down.shape[0]
rows_pad = 16
device = hyper_input_normed.device
num_ctas = torch.cuda.get_device_properties(device).multi_processor_count
t_raw = torch.empty((rows_pad, lowrank), dtype=torch.float32, device=device)
out = torch.empty((rows, hs), dtype=hyper_input_normed.dtype, device=device)
if rows == 0:
return out
_hc_mix_persistent_kernel[(num_ctas,)](
hyper_input_normed,
w_down,
w_up,
t_raw,
out,
_get_counters(device),
k,
lowrank,
hs,
rows,
num_ctas,
1.0 / hc,
ROWS=rows_pad,
HC=hc,
BLOCK_N=32,
BLOCK_K=256,
BLOCK_J=32,
BLOCK_R=64,
num_warps=8,
)
return out
python/sglang/srt/layers/hyperconnection.py
+333
View File
@@ -0,0 +1,333 @@
from typing import Optional
import msgspec
import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.srt.layers.hc_mix_triton import fused_hc_mix, fused_hc_mix_supported
class HyperConnectionConfig(msgspec.Struct, frozen=True):
hc_count: int = 4
hidden_size: int = 64
params_dtype: torch.dtype = torch.bfloat16
mtp_hc: bool = False
hc_lowrank: int = 16
rms_norm_eps: float = 1e-6
hc_per_branch_norm: bool = False
class GroupedGemmaRMSNorm(nn.Module):
def __init__(
self, hidden_size: int, eps: float = 1e-6, group_size: Optional[int] = None
):
super().__init__()
if group_size is not None and hidden_size % group_size != 0:
raise ValueError(
f"hidden_size ({hidden_size}) must be divisible by group_size ({group_size})"
)
self.weight = nn.Parameter(torch.zeros(hidden_size))
self.variance_epsilon = eps
self.group_size = group_size
self.weight.weight_loader = self._weight_loader
# The JIT kernel requires group_size to be a multiple of 512; this is
# init-static, so resolve it once here (device/dtype stay per-call).
effective_group_size = group_size if group_size is not None else hidden_size
self._jit_group_size = (
effective_group_size if effective_group_size % 512 == 0 else None
)
def _weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
assert param.size() == loaded_weight.size()
param.data.copy_(loaded_weight)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if (
self._jit_group_size is not None
and x.is_cuda
and x.dtype in (torch.bfloat16, torch.float16)
):
from sglang.kernels.ops.layernorm.grouped_gemma_rmsnorm import (
grouped_gemma_rmsnorm,
)
return grouped_gemma_rmsnorm(
x, self.weight, self._jit_group_size, self.variance_epsilon
)
input_dtype = x.dtype
x_float = x.float()
if self.group_size is None:
variance = x_float.pow(2).mean(dim=-1, keepdim=True)
x_norm = x_float * torch.rsqrt(variance + self.variance_epsilon)
else:
x_grouped = x_float.reshape(
*x_float.shape[:-1],
x_float.shape[-1] // self.group_size,
self.group_size,
)
variance = x_grouped.pow(2).mean(dim=-1, keepdim=True)
x_norm = (
x_grouped * torch.rsqrt(variance + self.variance_epsilon)
).flatten(-2)
return (x_norm * (1.0 + self.weight.float())).to(input_dtype)
class HyperConnectionBase(nn.Module):
def __init__(
self,
config: HyperConnectionConfig,
use_mix: bool = True,
use_combine: bool = True,
role: Optional[str] = None,
):
super().__init__()
self.config = config
self.hc_count = config.hc_count
if config.mtp_hc and role is not None and "mtp" in role:
self.hc_count = self.hc_count + 1
self.hidden_size = config.hidden_size
self.params_dtype = config.params_dtype
def mix(self, hyper_input: torch.Tensor):
assert hyper_input.shape[-1] == self.hc_count * self.hidden_size
mixed_input = hyper_input.view(
*hyper_input.shape[:-1], self.hc_count, self.hidden_size
).mean(dim=-2)
return mixed_input, hyper_input
def combine(
self, block_output: torch.Tensor, residual: torch.Tensor
) -> torch.Tensor:
assert residual.shape[-1] == self.hc_count * self.hidden_size
assert block_output.shape[-1] == self.hidden_size
residual_reshaped = residual.view(
*residual.shape[:-1], self.hc_count, self.hidden_size
)
combined_output = residual_reshaped + block_output.unsqueeze(-2)
combined_output = combined_output.view(
*residual.shape[:-1], self.hc_count * self.hidden_size
)
return combined_output
class GatedResidual(HyperConnectionBase):
def __init__(
self,
config: HyperConnectionConfig,
use_mix: bool = True,
use_combine: bool = True,
role: Optional[str] = None,
):
super().__init__(config, use_mix, use_combine, role)
norm_dim = (
self.config.hidden_size * self.hc_count
if self.config.hc_per_branch_norm
else self.config.hidden_size
)
norm_group_size = (
self.config.hidden_size if self.config.hc_per_branch_norm else None
)
self.hc_norm = GroupedGemmaRMSNorm(
norm_dim, eps=self.config.rms_norm_eps, group_size=norm_group_size
)
if use_mix:
self.input_mix_weight_down = nn.Linear(
self.hidden_size * self.hc_count,
self.config.hc_lowrank,
bias=False,
device=torch.cuda.current_device(),
dtype=config.params_dtype,
)
self.input_mix_weight_up = nn.Linear(
self.config.hc_lowrank,
self.hc_count * self.hidden_size,
bias=False,
device=torch.cuda.current_device(),
dtype=config.params_dtype,
)
lowrank = self.config.hc_lowrank
self._jit_mix_ok = (
torch.cuda.is_available()
# The CuTe split-K pair is tcgen05 (sm_100 family) only.
and torch.cuda.get_device_capability()[0] == 10
and (self.hc_count * self.hidden_size) % 2048 == 0
and self.hidden_size % 8 == 0
and lowrank > 0
and lowrank % 8 == 0
)
self._mix_up_weight_padded = None
if use_combine:
self.block_inject_weight = nn.Linear(
self.hidden_size * self.hc_count,
self.hc_count,
bias=False,
device=torch.cuda.current_device(),
dtype=config.params_dtype,
)
# hc_combine rejects other shapes; device and dtype are checked per call.
self._jit_combine_ok = (
self.hidden_size % 8 == 0
and (self.hc_count * self.hidden_size) % 2048 == 0
)
vecs = self.hc_count * self.hidden_size // 8
self._split_combine_ok = (
self._jit_combine_ok
and vecs % (8 * 160) == 0
and (self.hidden_size // 8) % (vecs // 8) == 0
)
def _mix_compute(
hyper_input_normed: torch.Tensor,
input_mix_weight_down: torch.Tensor,
input_mix_weight_up: torch.Tensor,
hc: int,
hs: int,
) -> torch.Tensor:
input_mix_weight = F.silu(
F.linear(hyper_input_normed, input_mix_weight_down) / hc
)
input_mix_weight = F.linear(input_mix_weight, input_mix_weight_up)
input_mix_weight = torch.sigmoid(input_mix_weight)
input_mix_weight = input_mix_weight.unflatten(-1, (hc, hs))
output = (
input_mix_weight * hyper_input_normed.unflatten(-1, (hc, hs))
).mean(dim=-2)
return output
def _combine_compute(
block_output: torch.Tensor,
residual: torch.Tensor,
normed_residual: torch.Tensor,
block_inject_weight: torch.Tensor,
hc: int,
hs: int,
) -> torch.Tensor:
R = residual.unflatten(-1, (hc, hs))
block_inject_weight_out = 2 * torch.sigmoid(
F.linear(normed_residual, block_inject_weight) / hc
)
injection = block_output.unsqueeze(-2) * block_inject_weight_out.unsqueeze(
-1
)
return (R + injection).flatten(-2)
self._mix_compute = torch.compile(_mix_compute)
self._combine_compute = torch.compile(_combine_compute)
def mix(self, hyper_input: torch.Tensor):
assert hyper_input.shape[-1] == self.hc_count * self.hidden_size
if hyper_input.shape[0] == 0:
mixed_input = hyper_input.new_empty(
(*hyper_input.shape[:-1], self.hidden_size), dtype=self.params_dtype
)
return mixed_input, (hyper_input, hyper_input)
if self.config.hc_per_branch_norm:
hyper_input_normed = self.hc_norm(hyper_input)
else:
hyper_input_normed = self.hc_norm(
hyper_input.unflatten(-1, (self.hc_count, self.hidden_size))
).flatten(-2)
if (
self._jit_mix_ok
and hyper_input_normed.is_cuda
and hyper_input_normed.dtype in (torch.bfloat16, torch.float16)
and hyper_input_normed.shape[0] <= 24
):
from sglang.kernels.ops.elementwise.hc_mix import (
hc_mix,
permute_pad_up_weight,
)
if self._mix_up_weight_padded is None:
self._mix_up_weight_padded = permute_pad_up_weight(
self.input_mix_weight_up.weight, self.hc_count
)
mixed_input = hc_mix(
hyper_input_normed,
self.input_mix_weight_down.weight.data,
self._mix_up_weight_padded,
self.hc_count,
self.hidden_size,
).to(self.params_dtype)
elif fused_hc_mix_supported(
hyper_input_normed,
self.input_mix_weight_down.weight,
self.input_mix_weight_up.weight,
):
mixed_input = fused_hc_mix(
hyper_input_normed,
self.input_mix_weight_down.weight,
self.input_mix_weight_up.weight,
self.hc_count,
self.hidden_size,
).to(self.params_dtype)
else:
mixed_input = self._mix_compute(
hyper_input_normed,
self.input_mix_weight_down.weight,
self.input_mix_weight_up.weight,
self.hc_count,
self.hidden_size,
).to(self.params_dtype)
return mixed_input, (hyper_input, hyper_input_normed)
def combine(self, block_output: torch.Tensor, residuals) -> torch.Tensor:
hyper_input, hyper_input_normed = residuals
assert hyper_input.shape[-1] == self.hc_count * self.hidden_size
assert block_output.shape[-1] == self.hidden_size
if block_output.shape[0] == 0:
return hyper_input.to(self.params_dtype)
if (
self._jit_combine_ok
and block_output.is_cuda
and block_output.dtype in (torch.bfloat16, torch.float16)
and hyper_input.dtype == block_output.dtype
and hyper_input_normed.dtype == block_output.dtype
and self.block_inject_weight.weight.dtype == block_output.dtype
):
if self._split_combine_ok and block_output.shape[0] <= 32:
from sglang.kernels.ops.elementwise.hc_combine import (
hc_combine_split,
)
return hc_combine_split(
block_output,
hyper_input,
hyper_input_normed,
self.block_inject_weight.weight.data,
self.hc_count,
self.hidden_size,
)
from sglang.kernels.ops.elementwise.hc_combine import hc_combine
return hc_combine(
block_output,
hyper_input,
hyper_input_normed,
self.block_inject_weight.weight,
self.hc_count,
self.hidden_size,
)
updated_residuals = self._combine_compute(
block_output,
hyper_input,
hyper_input_normed,
self.block_inject_weight.weight,
self.hc_count,
self.hidden_size,
).to(self.params_dtype)
return updated_residuals
HYPERCONNECTION_CLASS_DICT = {
"hyperconnection_average": HyperConnectionBase,
"gated_residual_simple": GatedResidual,
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=160,device_name=NVIDIA_H200.json
+74
View File
@@ -0,0 +1,74 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"1024": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 3
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 3
},
"8192": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 2
}
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=320,device_name=NVIDIA_H200.json
+74
View File
@@ -0,0 +1,74 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 3
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 3
},
"1024": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 3
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 3
},
"8192": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 3
}
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=80,device_name=NVIDIA_H200.json
+74
View File
@@ -0,0 +1,74 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"1024": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 2
},
"8192": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 8,
"num_warps": 4,
"num_stages": 2
}
}
python/sglang/srt/layers/quantization/unquant.py
+87 -4
View File
@@ -132,6 +132,48 @@ _BF16_SPLITK_TUNED_TACTICS = {
(16, 2560, 8192): (64, 16, 2, 11),
(24, 2560, 8192): (64, 32, 2, 9),
(32, 2560, 8192): (64, 32, 2, 9),
# Qwen4-Exp TP4 decode shapes, measured on B300 (sm103) under CUDA graph replay;
# unlisted (m, n, k) keep the CuTe DSL/cuBLAS path.
(1, 320, 2560): (64, 8, 4, 11),
(1, 512, 2560): (64, 8, 4, 11),
(1, 640, 2560): (64, 8, 4, 10),
(1, 2560, 1536): (64, 8, 2, 6),
(1, 2560, 2560): (64, 8, 2, 6),
(1, 3584, 2560): (64, 8, 2, 6),
(1, 4096, 2560): (64, 8, 2, 6),
(1, 4120, 2560): (64, 8, 2, 6),
(2, 320, 2560): (64, 8, 4, 10),
(2, 512, 2560): (64, 8, 4, 11),
(2, 640, 2560): (64, 8, 4, 11),
(2, 2560, 1536): (64, 8, 2, 6),
(2, 2560, 2560): (64, 8, 2, 6),
(2, 3584, 2560): (64, 8, 2, 6),
(2, 4096, 2560): (64, 8, 2, 6),
(2, 4120, 2560): (64, 8, 2, 6),
(3, 320, 2560): (64, 8, 4, 10),
(3, 512, 2560): (64, 8, 4, 10),
(3, 640, 2560): (64, 8, 4, 10),
(3, 2560, 1536): (64, 8, 2, 6),
(3, 2560, 2560): (64, 8, 2, 6),
(3, 3584, 2560): (64, 8, 2, 6),
(3, 4096, 2560): (64, 8, 2, 6),
(3, 4120, 2560): (64, 8, 2, 6),
(4, 320, 2560): (64, 8, 4, 12),
(4, 512, 2560): (64, 8, 4, 10),
(4, 640, 2560): (64, 8, 4, 11),
(4, 2560, 1536): (64, 8, 2, 6),
(4, 2560, 2560): (64, 8, 2, 6),
(4, 3584, 2560): (64, 8, 2, 6),
(4, 4096, 2560): (64, 8, 2, 6),
(4, 4120, 2560): (64, 8, 2, 6),
(8, 320, 2560): (64, 8, 4, 11),
(8, 512, 2560): (64, 8, 4, 10),
(8, 640, 2560): (64, 8, 4, 11),
(8, 2560, 1536): (64, 8, 2, 10),
(8, 2560, 2560): (64, 8, 2, 6),
(8, 3584, 2560): (64, 8, 2, 6),
(8, 4096, 2560): (64, 8, 2, 6),
(8, 4120, 2560): (64, 8, 2, 6),
}
@@ -139,6 +181,21 @@ def use_bf16_splitk_gemm(m: int, n: int, k: int) -> bool:
return (m, n, k) in _BF16_SPLITK_TUNED_TACTICS
def precompile_splitk_tactics() -> bool:
"""JIT-compile every tuned tactic through the real dispatch,
so CUDA graph capture never hits a cold kernel."""
if not _enable_bf16_splitk_gemm:
return False
device = torch.cuda.current_device()
for m, n, k in _BF16_SPLITK_TUNED_TACTICS:
x = torch.zeros(m, k, dtype=torch.bfloat16, device=device)
weight = torch.zeros(n, k, dtype=torch.bfloat16, device=device)
out = torch.empty(m, n, dtype=torch.bfloat16, device=device)
_bf16_splitk_gemm_out(x, weight, None, out)
torch.cuda.synchronize()
return True
def should_enable_bf16_splitk_gemm(backend: Bf16GemmBackend) -> bool:
"""Return whether the optional Split-K path should be initialized."""
return backend.is_cutedsl() and envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.get()
@@ -223,11 +280,12 @@ def _bf16_gemm_dispatch_fake(
return x.new_empty((*x.shape[:-1], weight.shape[0]))
def _bf16_splitk_gemm(
x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
def _bf16_splitk_gemm_out(
x_2d: torch.Tensor,
weight: torch.Tensor,
bias: Optional[torch.Tensor],
out: torch.Tensor,
) -> torch.Tensor:
x_2d = x.view(-1, x.shape[-1])
out = torch.empty((x_2d.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)
m, n, k = x_2d.shape[0], weight.shape[0], weight.shape[1]
if bias is None and _prefer_direct(m, n, k):
tactic = _direct_default_tactic(m, n, k)
@@ -242,6 +300,15 @@ def _bf16_splitk_gemm(
True,
tactic,
)
return out
def _bf16_splitk_gemm(
x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
) -> torch.Tensor:
x_2d = x.view(-1, x.shape[-1])
out = torch.empty((x_2d.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)
_bf16_splitk_gemm_out(x_2d, weight, bias, out)
return out.view(*x.shape[:-1], weight.shape[0])
@@ -459,6 +526,22 @@ class UnquantizedLinearMethod(LinearMethodBase):
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Run an inference-only BF16 linear into caller-owned storage."""
if (
_enable_bf16_splitk_gemm
and bias is None
and x.is_cuda
and x.ndim == 2
and x.dtype == torch.bfloat16
and layer.weight.dtype == torch.bfloat16
and output.dtype == torch.bfloat16
and output.is_contiguous()
and output.shape == (x.shape[0], layer.weight.shape[0])
and not layer.weight.requires_grad
and use_bf16_splitk_gemm(
x.shape[0], layer.weight.shape[0], layer.weight.shape[1]
)
):
return _bf16_splitk_gemm_out(x, layer.weight, None, output)
if (
get_bf16_gemm_backend().is_cutedsl()
and x.is_cuda
python/sglang/srt/layers/rotary_embedding/mrope_rope_index.py
+2
View File
@@ -74,6 +74,7 @@ def get_rope_index(
model_type.startswith("qwen3_vl")
or model_type.startswith("qwen3_vl_moe")
or model_type.startswith("qwen3_5")
or model_type == "qwen4_exp"
or model_type.startswith("interns2_mobius")
or model_type.startswith("cosmos3_omni")
or model_type.startswith("cosmos3_edge")
@@ -162,6 +163,7 @@ def get_rope_index(
"qwen3_vl_moe",
"qwen3_5",
"qwen3_5_moe",
"qwen4_exp",
"intern_s2_preview",
"interns2_mobius",
"cosmos3_omni",
python/sglang/srt/layers/vocab_parallel_embedding.py
+12 -2
View File
@@ -234,6 +234,7 @@ class VocabParallelEmbedding(torch.nn.Module):
embedding_dim: int,
*,
params_dtype: Optional[torch.dtype] = None,
output_dtype: Optional[torch.dtype] = None,
org_num_embeddings: Optional[int] = None,
padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
quant_config: Optional[QuantizationConfig] = None,
@@ -244,6 +245,7 @@ class VocabParallelEmbedding(torch.nn.Module):
):
super().__init__()
self.quant_config = quant_config
self.output_dtype = output_dtype
self.enable_tp = enable_tp
self.use_attn_tp_group = use_attn_tp_group
@@ -539,10 +541,13 @@ class VocabParallelEmbedding(torch.nn.Module):
)
if self.tp_size == 1:
with symm_alloc:
return self.quant_method.embedding(self, input_.long())
output_parallel = self.quant_method.embedding(self, input_.long())
if self.output_dtype is not None:
output_parallel = output_parallel.to(self.output_dtype)
return output_parallel
if self._use_triton_embedding(input_):
with symm_alloc:
return fused_vocab_parallel_embedding(
output_parallel = fused_vocab_parallel_embedding(
input_,
self.weight,
self.shard_indices.org_vocab_start_index,
@@ -551,6 +556,9 @@ class VocabParallelEmbedding(torch.nn.Module):
self.shard_indices.added_vocab_start_index,
self.shard_indices.added_vocab_end_index,
)
if self.output_dtype is not None:
output_parallel = output_parallel.to(self.output_dtype)
return output_parallel
# Map out-of-shard ids to index 0, gather, then zero those rows.
masked_input, input_mask = get_masked_input_and_mask(
input_,
@@ -562,6 +570,8 @@ class VocabParallelEmbedding(torch.nn.Module):
)
with symm_alloc:
output_parallel = self.quant_method.embedding(self, masked_input.long())
if self.output_dtype is not None:
output_parallel = output_parallel.to(self.output_dtype)
output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
return output_parallel
python/sglang/srt/mem_cache/kv_cache_configurator.py
+45 -2
View File
@@ -1106,6 +1106,20 @@ class KVCacheConfigurator:
"--enable-linear-replayssm-spec with DSPARK/DFLASH requires a KDA "
"(kimi_linear) model; got a non-KDA model."
)
from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig
ple_kwargs = {}
if isinstance(self.mambaish_config, Qwen4ExpTextConfig):
ple_kwargs = dict(
short_conv_layer_ids=[
i
for i in self.mambaish_config.short_conv_layer_ids
if self.layer_info.start_layer <= i < self.layer_info.end_layer
],
short_conv_state_shape=self.mambaish_config.short_conv_state_shape,
ngram_context_len=self.mambaish_config.ngram_context_len,
ngram_eos_token_id=int(self.mambaish_config.eos_token_id),
)
req_to_token_pool = HybridReqToTokenPool(
size=max_num_reqs,
mamba_size=get_schedule().max_mamba_cache_size,
@@ -1117,6 +1131,7 @@ class KVCacheConfigurator:
mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(),
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
**ple_kwargs,
# A PD prefill server never runs TARGET_VERIFY, so skip the
# verify-only per-draft-token state snapshots (see the draft-head
# case above: None => the pool skips SpeculativeState).
@@ -1840,7 +1855,36 @@ class KVCacheConfigurator:
if self.kv_cache_dtype_str == "mxfp8" and not self.use_mla_backend
else mha_pool_class
)
token_to_kv_pool = HybridLinearKVPool(
from sglang.srt.layers.attention.qsa.config import (
QSA_VARIANT_TOKENWISE,
parse_qsa_profile,
)
from sglang.srt.mem_cache.qsa_kv_pool import (
QSATokenToKVPool,
QwenDSATokenToKVPool,
)
qsa_profile = parse_qsa_profile(self.model_config.hf_config)
if qsa_profile is None:
pool_class = HybridLinearKVPool
extra_args["use_mla"] = self.use_mla_backend
elif qsa_profile.variant == QSA_VARIANT_TOKENWISE:
pool_class = QwenDSATokenToKVPool
extra_args.update(
qsa_index_kv_heads=qsa_profile.kv_heads,
qsa_index_head_dim=qsa_profile.head_dim,
qsa_token_budget=qsa_profile.budget,
)
else:
pool_class = QSATokenToKVPool
extra_args.update(
qsa_index_kv_heads=qsa_profile.kv_heads,
qsa_index_head_dim=qsa_profile.head_dim,
qsa_compress_ratio=qsa_profile.compress_ratio,
qsa_token_topk=qsa_profile.budget,
num_request_slots=req_to_token_pool.req_to_token.shape[0],
)
token_to_kv_pool = pool_class(
page_size=self.pool_page_size,
size=max_total_num_tokens,
dtype=self.kv_cache_dtype,
@@ -1854,7 +1898,6 @@ class KVCacheConfigurator:
mamba_pool=req_to_token_pool.mamba_pool,
enable_memory_saver=get_exec().features.enable_memory_saver,
enable_kv_cache_copy=(get_spec().speculative_algorithm is not None),
use_mla=self.use_mla_backend,
start_layer=self.layer_info.start_layer,
full_kv_pool_class=full_pool_class,
quant_method=quant_method,
python/sglang/srt/mem_cache/memory_pool.py
+100
View File
@@ -383,6 +383,15 @@ class MambaPool:
# Upstream states use (dim, K-1); subclasses may preserve another layout.
conv_window_axis = -1
# Slot-lifecycle side states (see ple_state_pool.SlotIndexedState);
# class-level default because UnifiedMambaPool skips MambaPool.__init__.
_slot_siblings: Tuple = ()
def register_slot_state(self, state) -> None:
"""Attach a state that rides along on clear / copy / host round-trip,
so a slot never changes owner with a stale sibling row attached."""
self._slot_siblings = [*self._slot_siblings, state]
@dataclass(frozen=True, kw_only=True)
class State:
conv: List[torch.Tensor]
@@ -977,6 +986,8 @@ class MambaPool:
def clear_slots(self, indices: torch.Tensor):
"""Zero out mamba state at the given pool indices. Must run on forward stream."""
for sibling in self._slot_siblings:
sibling.reset_slots(indices)
if self._should_fuse_slot_ops():
from sglang.srt.mem_cache.mamba_slot_fused import fused_clear_conv_slots
@@ -1047,6 +1058,8 @@ class MambaPool:
]
if self.replayssm_write_pos is not None:
self.replayssm_write_pos[dst_indices] = 0
for sibling in self._slot_siblings:
sibling.copy_slots(src_indices, dst_indices)
def get_cpu_copy(self, indices):
current_platform.synchronize()
@@ -1057,10 +1070,19 @@ class MambaPool:
temporal_cpu = self.mamba_cache.temporal[:, indices].to(
"cpu", non_blocking=True
)
siblings_cpu = [s.get_cpu_slots(indices) for s in self._slot_siblings]
current_platform.synchronize()
if self._slot_siblings:
return conv_cpu, temporal_cpu, siblings_cpu
return conv_cpu, temporal_cpu
def load_cpu_copy(self, mamba_cache_cpu, indices):
# The trailing element exists exactly when this instance registered siblings:
# the pool that saved the copy is the pool that loads it.
siblings_cpu = None
if self._slot_siblings:
siblings_cpu = mamba_cache_cpu[-1]
mamba_cache_cpu = mamba_cache_cpu[:-1]
# Accept historical 3-tuples, but request-keyed replay scratch is not
# restored with a physical checkpoint slot.
if len(mamba_cache_cpu) == 3:
@@ -1073,6 +1095,9 @@ class MambaPool:
self.mamba_cache.temporal[:, indices] = temporal_cpu.to(
self.mamba_cache.temporal.device, non_blocking=True
)
if siblings_cpu is not None:
for sibling, data in zip(self._slot_siblings, siblings_cpu):
sibling.load_cpu_slots(data, indices)
current_platform.synchronize()
_NON_TRANSFER_STATE_FIELDS = frozenset(
@@ -1202,6 +1227,10 @@ class HybridReqToTokenPool(ReqToTokenPool):
linear_replayssm_cache_len: int = 16,
mamba_envelope_layout: bool = False,
enable_linear_replayssm_spec: bool = False,
short_conv_layer_ids: Optional[List[int]] = None,
short_conv_state_shape: Optional[Tuple[int, int]] = None,
ngram_context_len: int = 0,
ngram_eos_token_id: int = 0,
):
super().__init__(
size=size,
@@ -1216,6 +1245,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
self.enable_memory_saver = enable_memory_saver
self.start_layer = start_layer if start_layer is not None else 0
self.layer_transfer_counter = None
self.ple_window_cache = None
self._init_mamba_pool(
mamba_size=mamba_size,
mamba_spec_state_size=mamba_spec_state_size,
@@ -1229,6 +1259,10 @@ class HybridReqToTokenPool(ReqToTokenPool):
linear_replayssm_cache_len=linear_replayssm_cache_len,
mamba_envelope_layout=mamba_envelope_layout,
enable_linear_replayssm_spec=enable_linear_replayssm_spec,
short_conv_layer_ids=short_conv_layer_ids,
short_conv_state_shape=short_conv_state_shape,
ngram_context_len=ngram_context_len,
ngram_eos_token_id=ngram_eos_token_id,
)
def _init_mamba_pool(
@@ -1245,6 +1279,10 @@ class HybridReqToTokenPool(ReqToTokenPool):
linear_replayssm_cache_len: int = 16,
mamba_envelope_layout: bool = False,
enable_linear_replayssm_spec: bool = False,
short_conv_layer_ids: Optional[List[int]] = None,
short_conv_state_shape: Optional[Tuple[int, int]] = None,
ngram_context_len: int = 0,
ngram_eos_token_id: int = 0,
):
self.mamba_pool = self.mamba_pool_cls(
size=mamba_size,
@@ -1266,6 +1304,36 @@ class HybridReqToTokenPool(ReqToTokenPool):
)
self.mamba_map = {layer_id: i for i, layer_id in enumerate(mamba_layer_ids)}
# Qwen4-Exp PLE side states; built disabled rather than None without a config,
# so every hybrid model has both attributes.
from sglang.srt.mem_cache.ple_state_pool import NGramPool, ShortConvPool
self.short_conv_pool = ShortConvPool(
size=mamba_size,
spec_state_size=mamba_spec_state_size,
state_shape=short_conv_state_shape,
layer_ids=short_conv_layer_ids or [],
dtype=cache_params.dtype.conv,
device=device,
enable_memory_saver=self.enable_memory_saver,
speculative_num_draft_tokens=speculative_num_draft_tokens,
)
self.ngram_pool = NGramPool(
size=mamba_size,
spec_state_size=mamba_spec_state_size,
context_len=ngram_context_len,
eos_token_id=ngram_eos_token_id,
device=device,
enable_memory_saver=self.enable_memory_saver,
speculative_num_draft_tokens=speculative_num_draft_tokens,
)
# Disabled pools stay off the sibling list so the host-offload payload
# keeps its legacy shape for every non-PLE hybrid model.
if self.short_conv_pool.enabled:
self.mamba_pool.register_slot_state(self.short_conv_pool)
if self.ngram_pool.enabled:
self.mamba_pool.register_slot_state(self.ngram_pool)
# Optional int8 checkpoint pool: the radix caches states here (int8) instead
# of holding them in the active bf16 pool -> ~2x cached-prefix capacity at
# fixed memory. Strategy-agnostic (no_buffer / extra_buffer / spec).
@@ -1279,6 +1347,15 @@ class HybridReqToTokenPool(ReqToTokenPool):
mamba_layer_ids=mamba_layer_ids,
device=device,
)
if self.mamba_ckpt_pool is not None and (
self.short_conv_pool.enabled or self.ngram_pool.enabled
):
# The int8 checkpoint pool frees the bf16 slot after donating its state,
# taking the bf16-slot-indexed PLE side states with it.
raise ValueError(
"--enable-int8-mamba-checkpoint is incompatible with Qwen4-Exp "
"PLE side states"
)
self.device = device
req_pool_size = self.req_to_token.shape[0]
@@ -1438,6 +1515,27 @@ class HybridReqToTokenPool(ReqToTokenPool):
def mamba2_layer_cache(self, layer_id: int):
return self.mamba_pool.mamba2_layer_cache(self.mamba2_layer_index(layer_id))
def short_conv_layer_cache(self, layer_id: int) -> torch.Tensor:
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
return self.short_conv_pool.layer_cache(layer_id)
def short_conv_layer_intermediate_cache(
self, layer_id: int
) -> Optional[torch.Tensor]:
return self.short_conv_pool.layer_intermediate_cache(layer_id)
def get_ngram_context(self, ngram_indices: torch.Tensor) -> torch.Tensor:
return self.ngram_pool.get_context(ngram_indices)
def set_ngram_context(
self, ngram_indices: torch.Tensor, context: torch.Tensor
) -> None:
self.ngram_pool.set_context(ngram_indices, context)
def set_ngram_intermediate_context(self, context: torch.Tensor) -> None:
self.ngram_pool.set_intermediate_context(context)
def copy_mamba_state(
self, src_index: torch.Tensor, dst_index: torch.Tensor
) -> None:
@@ -1607,6 +1705,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
logger.info("Reset HybridReqToTokenPool")
super().clear()
self.mamba_allocator.clear()
self.short_conv_pool.clear()
self.ngram_pool.clear()
# The int8 checkpoint pool holds radix-cached states in its own slots; a
# flush/reset drops the radix tree, so its slots must be released too,
# otherwise the (now unreferenced) slots leak and break the int8-pool
python/sglang/srt/mem_cache/ple_state_pool.py
+218
View File
@@ -0,0 +1,218 @@
"""Per-request PLE side states (short-conv window, N-gram context),
addressed by the request's MambaPool slot and registered on that pool,
so every slot lifecycle event carries them along."""
from __future__ import annotations
from contextlib import nullcontext
from typing import Any, List, Optional, Protocol, Tuple
import torch
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
class SlotIndexedState(Protocol):
"""Per-request state addressed by MambaPool slot index,
carried through every slot handoff;
every method must be a no-op when the backing tensor is None (a disabled pool)."""
def reset_slots(self, indices: torch.Tensor) -> None: ...
def copy_slots(self, src_index: torch.Tensor, dst_index: torch.Tensor) -> None: ...
def get_cpu_slots(self, indices: torch.Tensor) -> Any: ...
def load_cpu_slots(self, data: Any, indices: torch.Tensor) -> None: ...
class ShortConvPool:
def __init__(
self,
*,
size: int,
state_shape: Optional[Tuple[int, int]],
layer_ids: List[int],
dtype: torch.dtype,
device: str,
spec_state_size: int = 0,
enable_memory_saver: bool = False,
speculative_num_draft_tokens: Optional[int] = None,
):
self.size = size
self.device = device
self.layer_map = {layer_id: i for i, layer_id in enumerate(layer_ids)}
self.conv_state = None
self.intermediate_conv_state = None
if not layer_ids or state_shape is None:
return
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=enable_memory_saver
)
self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
maybe_init_custom_mem_pool(device=self.device)
)
with (
self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE),
(
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.enable_custom_mem_pool
else nullcontext()
),
):
self.conv_state = torch.zeros(
size=(len(layer_ids), size + 1) + state_shape,
dtype=dtype,
device=device,
)
if speculative_num_draft_tokens is not None:
self.intermediate_conv_state = torch.zeros(
size=(
len(layer_ids),
spec_state_size + 1,
speculative_num_draft_tokens,
)
+ state_shape,
dtype=dtype,
device=device,
)
@property
def enabled(self) -> bool:
return self.conv_state is not None
def layer_cache(self, layer_id: int) -> torch.Tensor:
assert self.conv_state is not None
assert layer_id in self.layer_map
return self.conv_state[self.layer_map[layer_id]]
def layer_intermediate_cache(self, layer_id: int) -> Optional[torch.Tensor]:
if self.intermediate_conv_state is None:
return None
assert layer_id in self.layer_map
return self.intermediate_conv_state[self.layer_map[layer_id]]
def clear(self):
if self.conv_state is not None:
self.conv_state.zero_()
# SlotIndexedState: slot is dim 1, behind the layer dim.
def reset_slots(self, indices: torch.Tensor) -> None:
if self.conv_state is not None and indices.numel() > 0:
self.conv_state[:, indices] = 0
def copy_slots(self, src_index: torch.Tensor, dst_index: torch.Tensor) -> None:
if self.conv_state is not None:
self.conv_state[:, dst_index] = self.conv_state[:, src_index]
def get_cpu_slots(self, indices: torch.Tensor) -> Any:
if self.conv_state is None:
return None
return self.conv_state[:, indices].to("cpu", non_blocking=True)
def load_cpu_slots(self, data: Any, indices: torch.Tensor) -> None:
if self.conv_state is None or data is None:
return
self.conv_state[:, indices] = data.to(self.conv_state.device, non_blocking=True)
class NGramPool:
def __init__(
self,
*,
size: int,
context_len: int,
eos_token_id: int,
device: str,
spec_state_size: int = 0,
enable_memory_saver: bool = False,
speculative_num_draft_tokens: Optional[int] = None,
):
self.size = size
self.context_len = context_len
self.eos_token_id = eos_token_id
self.device = device
self.context = None
self.intermediate_context = None
if context_len <= 0:
return
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=enable_memory_saver
)
self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
maybe_init_custom_mem_pool(device=self.device)
)
with (
self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE),
(
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.enable_custom_mem_pool
else nullcontext()
),
):
self.context = torch.full(
(size + 1, context_len),
eos_token_id,
dtype=torch.long,
device=device,
)
if speculative_num_draft_tokens is not None:
self.intermediate_context = torch.full(
(spec_state_size + 1, speculative_num_draft_tokens, context_len),
eos_token_id,
dtype=torch.long,
device=device,
)
@property
def enabled(self) -> bool:
return self.context is not None
def get_context(self, indices: torch.Tensor) -> torch.Tensor:
assert self.context is not None
return self.context.index_select(0, indices.to(dtype=torch.long))
def set_context(self, indices: torch.Tensor, context: torch.Tensor):
if self.context is not None and indices.numel() > 0:
self.context[indices.to(dtype=torch.long)] = context.to(
device=self.context.device, dtype=self.context.dtype
)
def set_intermediate_context(self, context: torch.Tensor):
if self.intermediate_context is not None and context.numel() > 0:
self.intermediate_context[: context.shape[0], : context.shape[1]].copy_(
context.to(device=self.context.device, dtype=self.context.dtype)
)
def clear(self):
if self.context is not None:
self.context.fill_(self.eos_token_id)
# SlotIndexedState: slot is dim 0, no layer dim.
def reset_slots(self, indices: torch.Tensor) -> None:
if self.context is not None and indices.numel() > 0:
self.context[indices.to(dtype=torch.long)] = self.eos_token_id
def copy_slots(self, src_index: torch.Tensor, dst_index: torch.Tensor) -> None:
if self.context is not None:
src = src_index.to(dtype=torch.long)
dst = dst_index.to(dtype=torch.long)
self.context[dst] = self.context[src]
def get_cpu_slots(self, indices: torch.Tensor) -> Any:
if self.context is None:
return None
return self.context[indices.to(dtype=torch.long)].to("cpu", non_blocking=True)
def load_cpu_slots(self, data: Any, indices: torch.Tensor) -> None:
if self.context is None or data is None:
return
self.context[indices.to(dtype=torch.long)] = data.to(
self.context.device, non_blocking=True
)
python/sglang/srt/mem_cache/qsa_kv_pool.py
+312
View File
@@ -0,0 +1,312 @@
"""KV pools carrying the QSA sparse-attention indexer caches."""
from __future__ import annotations
from typing import List, Optional
import torch
from sglang.srt.mem_cache.memory_pool import GB, HybridLinearKVPool, MambaPool
def _index_k_bytes(*, kv_heads: int, head_dim: int, dtype: torch.dtype) -> int:
return kv_heads * head_dim * dtype.itemsize
class QSATokenToKVPool(HybridLinearKVPool):
"""Hybrid KV pool with the minimal BF16 state required by simple QSA."""
# Full-KV pages are a multiple of the compress ratio, so no group straddles pages;
# ``compressed_slot = full_slot // ratio`` needs no ownership bookkeeping;
# lifecycle rides the full-KV allocator and radix tree.
# Full slot 0 is the reserved padding slot; compressed slot 0 is the inert dump.
index_state_dtype = torch.bfloat16
@classmethod
def qsa_bytes_per_token(
cls, *, kv_heads: int, head_dim: int, compress_ratio: int, num_layers: int
) -> int:
"""Per-token QSA index-cache cost: compressed keys only;
the per-request pending ring is budgeted with the other per-request buffers."""
index_k_bytes = _index_k_bytes(
kv_heads=kv_heads, head_dim=head_dim, dtype=cls.index_state_dtype
)
return index_k_bytes // compress_ratio * num_layers
def __init__(
self,
*,
size: int,
dtype: torch.dtype,
page_size: int,
head_num: int,
head_dim: int,
full_attention_layer_ids: List[int],
device: str,
mamba_pool: MambaPool,
qsa_index_kv_heads: int,
qsa_index_head_dim: int,
qsa_compress_ratio: int,
qsa_token_topk: int,
num_request_slots: int,
enable_memory_saver: bool = False,
enable_kv_cache_copy: bool = False,
start_layer: Optional[int] = None,
full_kv_pool_class: Optional[type] = None,
quant_method=None,
post_capture_active: bool = False,
):
if page_size <= 1 or page_size % qsa_compress_ratio != 0:
raise ValueError(
"compressed QSA requires a paged full-KV cache with the page "
"a multiple of the compress ratio (compressed slots are "
f"full_slot // ratio): page_size={page_size}, "
f"ratio={qsa_compress_ratio}. With MambaRadixCache this "
"needs the mamba extra-buffer strategy or "
"--disable-radix-cache (see the Qwen4-Exp arg overrides)."
)
# super().__init__ computes mem_usage via the overridden get_kv_size_bytes,
# so the QSA buffers get placeholders first; mem_usage is recomputed last.
self.qsa_key_state_buffer_pool = []
self.qsa_compressed_k_buffer_pool = []
self.qsa_rope_position_buffer = torch.empty(0)
super().__init__(
size=size,
dtype=dtype,
page_size=page_size,
head_num=head_num,
head_dim=head_dim,
full_attention_layer_ids=full_attention_layer_ids,
device=device,
mamba_pool=mamba_pool,
enable_memory_saver=enable_memory_saver,
enable_kv_cache_copy=enable_kv_cache_copy,
use_mla=False,
start_layer=start_layer,
full_kv_pool_class=full_kv_pool_class,
quant_method=quant_method,
post_capture_active=post_capture_active,
)
if (
min(
qsa_index_kv_heads,
qsa_index_head_dim,
qsa_compress_ratio,
qsa_token_topk,
)
<= 0
):
raise ValueError("QSA cache configuration values must be positive")
if qsa_token_topk % qsa_compress_ratio != 0:
raise ValueError("qsa_token_topk must be divisible by qsa_compress_ratio")
self.qsa_compress_ratio = int(qsa_compress_ratio)
self.qsa_index_head_dim = int(qsa_index_head_dim)
self.qsa_index_kv_heads = int(qsa_index_kv_heads)
self.qsa_token_topk = int(qsa_token_topk)
self.qsa_block_topk = self.qsa_token_topk // self.qsa_compress_ratio
state_size = size + page_size
# Compressed slots mirror the full-KV slot space 1:ratio; the "page"
# seen by the scoring kernels is one full-KV page's worth of groups.
self.qsa_compressed_page_size = page_size // self.qsa_compress_ratio
self.qsa_compressed_capacity = -(state_size // -self.qsa_compress_ratio)
# Pre-compression index-K state is a per-request ring, not a per-token cache:
# only the pending group's ``ratio`` members must survive a forward,
# addressed as ``req_pool_idx * ratio + position % ratio``.
# Request slot 0 is never allocated, so rows [0, ratio) are the inert dump.
if num_request_slots <= 0:
raise ValueError(
f"QSA pending ring needs request slots, got {num_request_slots}"
)
self.qsa_num_request_slots = int(num_request_slots)
ring_slots = self.qsa_num_request_slots * self.qsa_compress_ratio
self.qsa_key_state_buffer_pool = [
torch.zeros(
(ring_slots, self.qsa_index_kv_heads, self.qsa_index_head_dim),
dtype=self.index_state_dtype,
device=device,
)
for _ in full_attention_layer_ids
]
# Layer-independent MRoPE coordinate of every pending key;
# the compress kernel rotates the pooled key at the group's real start position.
self.qsa_rope_position_buffer = torch.zeros(
(ring_slots, 3), dtype=torch.int64, device=device
)
# One contiguous allocation behind per-layer views: every layer's
# compressed pages are addressable from a single base pointer.
self.qsa_compressed_flat = torch.zeros(
(
len(full_attention_layer_ids),
self.qsa_compressed_capacity
* self.qsa_index_kv_heads
* self.qsa_index_head_dim,
),
dtype=self.index_state_dtype,
device=device,
)
self.qsa_compressed_k_buffer_pool = [
self.qsa_compressed_flat[layer_offset].view(
self.qsa_compressed_capacity,
self.qsa_index_kv_heads,
self.qsa_index_head_dim,
)
for layer_offset in range(len(full_attention_layer_ids))
]
k_size, v_size = self.get_kv_size_bytes()
self.mem_usage = (k_size + v_size) / GB
def get_qsa_key_state_buffer(self, layer_id: int) -> torch.Tensor:
return self.qsa_key_state_buffer_pool[
self._transfer_full_attention_id(layer_id)
]
def set_qsa_key_state_buffer(
self, layer_id: int, loc: torch.Tensor, token_k: torch.Tensor
) -> None:
buffer = self.get_qsa_key_state_buffer(layer_id)
buffer[loc.long()] = token_k.to(buffer.dtype)
def set_qsa_rope_position_buffer(
self, loc: torch.Tensor, positions: torch.Tensor
) -> None:
positions = positions.long()
if positions.ndim == 1:
positions = positions.unsqueeze(0).expand(3, -1)
if positions.ndim != 2 or positions.shape[0] != 3:
raise ValueError(
f"QSA RoPE positions must be [tokens] or [3, tokens], got {positions.shape}"
)
self.qsa_rope_position_buffer[loc.long()] = positions.transpose(0, 1)
def get_qsa_rope_position_buffer(self, loc: torch.Tensor) -> torch.Tensor:
return self.qsa_rope_position_buffer[loc.long()]
def get_qsa_compressed_k_buffer(self, layer_id: int) -> torch.Tensor:
return self.qsa_compressed_k_buffer_pool[
self._transfer_full_attention_id(layer_id)
]
def set_qsa_compressed_k_buffer(
self, layer_id: int, loc: torch.Tensor, compressed_k: torch.Tensor
) -> None:
buffer = self.get_qsa_compressed_k_buffer(layer_id)
buffer[loc.long()] = compressed_k.to(buffer.dtype)
def get_kv_size_bytes(self):
k_size, v_size = super().get_kv_size_bytes()
qsa_k_size = (
sum(
tensor.numel() * tensor.element_size()
for tensor in self.qsa_key_state_buffer_pool
)
+ sum(
tensor.numel() * tensor.element_size()
for tensor in self.qsa_compressed_k_buffer_pool
)
+ self.qsa_rope_position_buffer.numel() * 8
)
return k_size + qsa_k_size, v_size
class QwenDSATokenToKVPool(HybridLinearKVPool):
"""Hybrid KV pool carrying the per-token index-K cache of tokenwise QSA:
a ``[size + page_size, index_kv_heads, index_head_dim]`` BF16 buffer per DSA layer,
addressed by raw KV slots; the FP8 deep_gemm layout is deliberately absent."""
index_state_dtype = torch.bfloat16
@classmethod
def qsa_bytes_per_token(
cls, *, kv_heads: int, head_dim: int, num_layers: int
) -> int:
return (
_index_k_bytes(
kv_heads=kv_heads, head_dim=head_dim, dtype=cls.index_state_dtype
)
* num_layers
)
def __init__(
self,
*,
size: int,
dtype: torch.dtype,
page_size: int,
head_num: int,
head_dim: int,
full_attention_layer_ids: List[int],
device: str,
mamba_pool: MambaPool,
qsa_index_kv_heads: int,
qsa_index_head_dim: int,
qsa_token_budget: int,
enable_memory_saver: bool = False,
enable_kv_cache_copy: bool = False,
start_layer: Optional[int] = None,
full_kv_pool_class: Optional[type] = None,
quant_method=None,
post_capture_active: bool = False,
):
if page_size != 64:
raise ValueError(
"tokenwise QSA requires KV-cache page_size 64 for its paged "
f"indexer buffer, got {page_size}"
)
self.dsa_index_k_buffer_pool = []
super().__init__(
size=size,
dtype=dtype,
page_size=page_size,
head_num=head_num,
head_dim=head_dim,
full_attention_layer_ids=full_attention_layer_ids,
device=device,
mamba_pool=mamba_pool,
enable_memory_saver=enable_memory_saver,
enable_kv_cache_copy=enable_kv_cache_copy,
use_mla=False,
start_layer=start_layer,
full_kv_pool_class=full_kv_pool_class,
quant_method=quant_method,
post_capture_active=post_capture_active,
)
if qsa_index_kv_heads != 1:
raise ValueError(
f"tokenwise QSA requires index_kv_heads = 1 (MQA), got "
f"{qsa_index_kv_heads}"
)
if min(qsa_index_kv_heads, qsa_index_head_dim, qsa_token_budget) <= 0:
raise ValueError("QSA cache configuration values must be positive")
self.qsa_compress_ratio = 1
self.qsa_index_kv_heads = int(qsa_index_kv_heads)
self.qsa_index_head_dim = int(qsa_index_head_dim)
self.qsa_token_topk = int(qsa_token_budget)
self.qsa_block_topk = int(qsa_token_budget)
state_size = size + page_size
self.dsa_index_k_buffer_pool = [
torch.zeros(
(state_size, self.qsa_index_kv_heads, self.qsa_index_head_dim),
dtype=self.index_state_dtype,
device=device,
)
for _ in full_attention_layer_ids
]
k_size, v_size = self.get_kv_size_bytes()
self.mem_usage = (k_size + v_size) / GB
def set_dsa_index_k_buffer(
self, layer_id: int, loc: torch.Tensor, index_k: torch.Tensor
) -> None:
buffer = self.get_dsa_index_k_buffer(layer_id)
buffer[loc.long()] = index_k.to(buffer.dtype)
def get_dsa_index_k_buffer(self, layer_id: int) -> torch.Tensor:
return self.dsa_index_k_buffer_pool[self._transfer_full_attention_id(layer_id)]
def get_kv_size_bytes(self):
k_size, v_size = super().get_kv_size_bytes()
dsa_k_size = sum(
tensor.numel() * tensor.element_size()
for tensor in self.dsa_index_k_buffer_pool
)
return k_size + dsa_k_size, v_size
python/sglang/srt/mem_cache/unified_memory_pool.py
+24
View File
@@ -1063,11 +1063,35 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
enable_linear_replayssm: bool = False,
linear_replayssm_cache_len: int = 16,
enable_linear_replayssm_spec: bool = False,
short_conv_layer_ids: Optional[List[int]] = None,
short_conv_state_shape=None,
ngram_context_len: int = 0,
ngram_eos_token_id: int = 0,
):
# mamba_envelope_layout / speculative_eagle_topk / enable_linear_replayssm /
# linear_replayssm_cache_len / enable_linear_replayssm_spec: accepted to match
# the parent signature but NOT forwarded — the shared pool's conv/temporal
# state are fixed-shape views (replayssm/spec are gated off under unified).
if short_conv_layer_ids or ngram_context_len:
raise ValueError(
"Qwen4-Exp PLE side states are not supported with "
"--enable-unified-memory"
)
from sglang.srt.mem_cache.ple_state_pool import NGramPool, ShortConvPool
self.short_conv_pool = ShortConvPool(
size=0,
state_shape=None,
layer_ids=[],
dtype=torch.bfloat16,
device=device,
)
self.ngram_pool = NGramPool(
size=0,
context_len=0,
eos_token_id=0,
device=device,
)
assert mamba_size == self._shared_mamba_size, (
f"UnifiedHybridReqToTokenPool._init_mamba_pool: mamba_size={mamba_size} "
f"!= unified_buffer.max_slots({self._mamba_sub_pool_name!r}) - 1 "
python/sglang/srt/model_executor/forward_batch_info.py
+21
View File
@@ -55,6 +55,7 @@ from sglang.srt.runtime_context import (
get_exec,
get_lora,
get_parallel,
mamba_cache_chunk_size,
)
from sglang.srt.speculative.spec_info import SpecInputType
from sglang.srt.utils import (
@@ -1059,6 +1060,26 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
sharded=sharded,
)
def mamba_track_aligned_lens(self) -> Optional[torch.Tensor]:
"""Tokens of this extend chunk covered by the tracked mamba state,
floored to the mamba_cache_chunk_size boundary the scheduler snapshots at;
the +1 that _force_track_h adds cancels under the floor.
Sole home of this math: every side state snapshotting alongside mamba calls it.
None means tracking is skipped for this forward:
no mask, or a prefill CUDA-graph replay without mamba_track_seqlens.
Masked-off rows hold garbage.
"""
if (
self.mamba_track_mask is None
or self.mamba_track_seqlens is None
or self.extend_prefix_lens is None
):
return None
chunk_size = mamba_cache_chunk_size()
lens_to_track = self.mamba_track_seqlens - self.extend_prefix_lens
return (lens_to_track // chunk_size) * chunk_size
def merge_mm_inputs(self) -> Optional[MultimodalInputs]:
"""
Merge all multimodal inputs in the batch into a single MultiModalInputs object.
python/sglang/srt/model_executor/model_runner_components/load_model_utils.py
+17
View File
@@ -265,6 +265,17 @@ def load_model_with_memory_saver(
# Remove monkey_patch when linear.py quant remove dependencies with vllm
monkey_patch_vllm_parallel_state()
if not is_draft_worker:
architectures = model_config.hf_config.architectures or []
is_qwen4_exp = "Qwen4ExpForConditionalGeneration" in architectures
ple_offload_embedding = get_exec().offload.ple_offload_embedding
if ple_offload_embedding and not is_qwen4_exp:
raise ValueError(
"--ple-offload-embedding only supports Qwen4ExpForConditionalGeneration"
)
if is_qwen4_exp:
model_config.hf_text_config.ple_offload_embedding = ple_offload_embedding
enable_cpu_backup = get_exec().features.enable_weights_cpu_backup or (
is_draft_worker and get_exec().features.enable_draft_weights_cpu_backup
)
@@ -312,6 +323,12 @@ def load_model_with_memory_saver(
remote_instance_weight_info = (
loader.remote_instance_transfer_engine_weight_info
)
if (
not is_draft_worker
and get_exec().offload.ple_offload_embedding
and device == "cuda"
):
current_platform.empty_cache()
# Cache needs to be cleared after loading model weights (in the loader.load_model function).
# To avoid conflict with memory_saver_adapter.region, empty_cache operation is now moved here.
if _is_npu:
python/sglang/srt/model_executor/pool_configurator.py
+33
View File
@@ -406,8 +406,40 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
n * (model_config.head_dim + model_config.v_head_dim) * num_layers
) // scale_block_size
cell_size += self._compute_qsa_cell_size(
hf_config=model_config.hf_config, num_layers=num_layers
)
return cell_size
@staticmethod
def _compute_qsa_cell_size(*, hf_config, num_layers: int) -> int:
from sglang.srt.layers.attention.qsa.config import (
QSA_VARIANT_COMPRESSED,
parse_qsa_profile,
)
from sglang.srt.mem_cache.qsa_kv_pool import (
QSATokenToKVPool,
QwenDSATokenToKVPool,
)
if num_layers == 0:
return 0
qsa_profile = parse_qsa_profile(hf_config)
if qsa_profile is None:
return 0
if qsa_profile.variant == QSA_VARIANT_COMPRESSED:
return QSATokenToKVPool.qsa_bytes_per_token(
kv_heads=qsa_profile.kv_heads,
head_dim=qsa_profile.head_dim,
compress_ratio=qsa_profile.compress_ratio,
num_layers=num_layers,
)
return QwenDSATokenToKVPool.qsa_bytes_per_token(
kv_heads=qsa_profile.kv_heads,
head_dim=qsa_profile.head_dim,
num_layers=num_layers,
)
def _compute_dsa_indexer_cell_size(
self,
*,
@@ -482,6 +514,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
def calculate_pool_sizes(
self, available_bytes: int, page_size: int
) -> MemoryPoolConfig:
available_bytes = max(available_bytes, 0)
max_total_num_tokens = (
available_bytes // self._cell_size
if self._cell_size
python/sglang/srt/models/qwen3_5.py
+105 -40
View File
@@ -68,6 +68,10 @@ from sglang.srt.layers.parameter import (
PerTensorScaleParameter,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.unquant import (
UnquantizedLinearMethod,
bf16_gemm_dispatch,
)
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.layers.rotary_embedding import get_rope
@@ -100,6 +104,7 @@ from sglang.srt.models.utils import (
from sglang.srt.runtime_context import (
get_exec,
get_forward,
get_lora,
get_parallel,
get_stream,
)
@@ -128,6 +133,7 @@ _is_npu = is_npu()
_is_cpu = is_cpu()
_is_gfx95 = is_gfx95_supported()
_is_hip = is_hip()
_QWEN3_5_MOE_TEXT_MODEL_TYPES = ("qwen3_5_moe_text", "qwen4_exp_text")
_is_xpu = is_xpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_hip_use_alt_stream = get_bool_env_var("SGLANG_ALT_STREAM") and _is_hip
@@ -387,6 +393,8 @@ class Qwen3_5GatedDeltaNet(nn.Module):
# `weight_scale_inv` / `weight_scale` / `input_scale` if present.
self._bind_packed_weight_loaders(self.in_proj_qkvz)
self._bind_packed_weight_loaders(self.in_proj_ba)
self._fused_in_proj_weight: Optional[torch.Tensor] = None
self._fused_in_proj_qkvz_width = 0
self._fused_input_proj_cpu_enabled = LazyValue(
lambda: (
_is_cpu
@@ -645,6 +653,32 @@ class Qwen3_5GatedDeltaNet(nn.Module):
return query, key, value, z, b, a
def finalize_fused_in_proj(self) -> None:
"""Stack in_proj_qkvz + in_proj_ba into one GEMM weight;
the module weights become row views of it,
so weight reload and dtype checks still see them."""
if not _is_cuda or self._fused_in_proj_weight is not None:
return
if get_lora().enable_lora or get_lora().lora_paths:
# LoRA wraps the individual Linear modules; the fused GEMM would
# bypass their adapters.
return
qkvz, ba = self.in_proj_qkvz, self.in_proj_ba
if not (
isinstance(qkvz.quant_method, UnquantizedLinearMethod)
and isinstance(ba.quant_method, UnquantizedLinearMethod)
and qkvz.weight.dtype == torch.bfloat16
and ba.weight.dtype == torch.bfloat16
and qkvz.bias is None
and ba.bias is None
):
return
fused = torch.cat([qkvz.weight.data, ba.weight.data], dim=0).contiguous()
self._fused_in_proj_qkvz_width = qkvz.weight.shape[0]
qkvz.weight.data = fused[: self._fused_in_proj_qkvz_width]
ba.weight.data = fused[self._fused_in_proj_qkvz_width :]
self._fused_in_proj_weight = fused
def _forward_input_proj(self, hidden_states: torch.Tensor):
# AMD/aiter fused AR+RMSNorm+per-group-quant path ships a
# ``(bf16, fp8, scale)`` 3-tuple so the FP8 ``in_proj_qkvz`` can
@@ -654,6 +688,21 @@ class Qwen3_5GatedDeltaNet(nn.Module):
if _use_aiter and isinstance(hidden_states, tuple):
return self._forward_input_proj_fused_quant_amd(hidden_states)
if (
self._fused_in_proj_weight is not None
and hidden_states.dtype == torch.bfloat16
# Measured on cuBLAS above ~1k rows:
# the merged (m, 4120) GEMM is ~10% slower than the two separate GEMMs.
and hidden_states.shape[0] <= 1024
):
fused_out = bf16_gemm_dispatch(
hidden_states, self._fused_in_proj_weight, None
)
return (
fused_out[:, : self._fused_in_proj_qkvz_width],
fused_out[:, self._fused_in_proj_qkvz_width :],
)
if (
_is_cpu
or _is_npu
@@ -899,7 +948,7 @@ class Qwen3_5LinearDecoderLayer(nn.Module):
# NOTE: Determine the MLP type based on the model type
# Qwen3.5 use all layers for MLP / Qwen3.5-MoE use sparse MoE blocks
if config.model_type == "qwen3_5_moe_text":
if config.model_type in _QWEN3_5_MOE_TEXT_MODEL_TYPES:
self.mlp = Qwen2MoeSparseMoeBlock(
layer_id=layer_id,
config=config,
@@ -1151,7 +1200,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
is_layer_sparse = False
is_previous_layer_sparse = False
is_next_layer_sparse = False
elif config.model_type == "qwen3_5_moe_text":
elif config.model_type in _QWEN3_5_MOE_TEXT_MODEL_TYPES:
self.mlp = Qwen2MoeSparseMoeBlock(
layer_id=layer_id,
config=config,
@@ -1343,6 +1392,37 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
)
return q, k, v, gate
def _prepare_qkv_gate(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
if _is_cuda and self.attn_output_gate:
return self.forward_prepare_cuda_fused(
positions=positions,
hidden_states=hidden_states,
)
if (_is_hip or _is_xpu or _is_cpu) and self.attn_output_gate:
return self.forward_prepare_fused_gate(
positions=positions,
hidden_states=hidden_states,
)
if (
not _is_npu
or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
or not self.attn_output_gate
):
return self.forward_prepare_native(
positions=positions,
hidden_states=hidden_states,
)
return self.forward_prepare_npu(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
)
def self_attention(
self,
positions: torch.Tensor,
@@ -1350,31 +1430,11 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
forward_batch: ForwardBatch,
) -> torch.Tensor:
"""Full attention forward pass."""
if _is_cuda and self.attn_output_gate:
q, k, v, gate = self.forward_prepare_cuda_fused(
positions=positions,
hidden_states=hidden_states,
)
elif (_is_hip or _is_xpu or _is_cpu) and self.attn_output_gate:
q, k, v, gate = self.forward_prepare_fused_gate(
positions=positions,
hidden_states=hidden_states,
)
elif (
not _is_npu
or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
or not self.attn_output_gate
):
q, k, v, gate = self.forward_prepare_native(
positions=positions,
hidden_states=hidden_states,
)
else:
q, k, v, gate = self.forward_prepare_npu(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
)
q, k, v, gate = self._prepare_qkv_gate(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
)
attn_output = self.attn(q, k, v, forward_batch)
@@ -1490,6 +1550,8 @@ QWEN3_5_KV_SCALE_MAPPER = WeightsMapper(
class Qwen3_5ForCausalLM(nn.Module):
"""Qwen3.5 Model with support for dense variant."""
decoder_layer_types = ALL_DECODER_LAYER_TYPES
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
@@ -1530,14 +1592,14 @@ class Qwen3_5ForCausalLM(nn.Module):
elif module_name == "gate_up_proj":
# MoE: shared expert uses shared_expert_intermediate_size
# Dense: regular MLP uses intermediate_size
is_moe = "moe" in getattr(config, "model_type", "")
is_moe = config.model_type in _QWEN3_5_MOE_TEXT_MODEL_TYPES
if is_moe:
inter = config.shared_expert_intermediate_size
else:
inter = config.intermediate_size
return config.hidden_size, inter * 2
elif module_name == "down_proj":
is_moe = "moe" in getattr(config, "model_type", "")
is_moe = config.model_type in _QWEN3_5_MOE_TEXT_MODEL_TYPES
if is_moe:
inter = config.shared_expert_intermediate_size
else:
@@ -1571,20 +1633,12 @@ class Qwen3_5ForCausalLM(nn.Module):
alt_stream = get_stream("alt") if _is_cuda or _hip_use_alt_stream else None
# Embedding layer
if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
org_num_embeddings=config.vocab_size,
enable_tp=not is_dp_attention_enabled(),
)
else:
self.embed_tokens = PPMissingLayer()
self.embed_tokens = self._build_embed_tokens(config)
# Decoder layers
def get_layer(idx: int, prefix: str):
layer_type = config.layers_block_type[idx]
layer_class = ALL_DECODER_LAYER_TYPES[layer_type]
layer_class = self.decoder_layer_types[layer_type]
if layer_type == "attention":
prefix = add_prefix("self_attn", prefix)
else:
@@ -1656,6 +1710,17 @@ class Qwen3_5ForCausalLM(nn.Module):
self.layers_to_capture = []
def _build_embed_tokens(self, config: Qwen3_5TextConfig) -> nn.Module:
"""Embedding sharding hook for models reusing this backbone."""
if not self.pp_group.is_first_rank:
return PPMissingLayer()
return VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
org_num_embeddings=config.vocab_size,
enable_tp=not is_dp_attention_enabled(),
)
def get_input_embeddings(self):
return self.embed_tokens
@@ -2678,7 +2743,7 @@ def _qwen3_5_shared_experts_fusion_disable_reason(hf_config, quant_config):
if not _is_hip:
return None
text_config = getattr(hf_config, "text_config", hf_config)
if getattr(text_config, "model_type", None) != "qwen3_5_moe_text":
if text_config.model_type not in _QWEN3_5_MOE_TEXT_MODEL_TYPES:
return None
if can_fuse_shared_expert(text_config, quant_config):
return None
python/sglang/srt/models/qwen3_vl.py
+7 -1
View File
@@ -1367,7 +1367,13 @@ class Qwen3VLForConditionalGeneration(nn.Module):
config.vision_config.deepstack_visual_indexes
)
self.num_deepstack_embeddings = len(self.deepstack_visual_indexes)
self.use_deepstack = {Modality.IMAGE: True, Modality.VIDEO: True}
# Only enable deepstack when the checkpoint declares deepstack
# capture layers (Qwen4-Exp ships an empty list).
self.use_deepstack = (
{Modality.IMAGE: True, Modality.VIDEO: True}
if self.num_deepstack_embeddings > 0
else {}
)
else:
self.deepstack_visual_indexes = []
self.num_deepstack_embeddings = 0
python/sglang/srt/models/qwen4_exp.py
+2131
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/models/qwen4_exp_mtp.py
+230
View File
@@ -0,0 +1,230 @@
"""Inference-only Qwen4-Exp MTP speculative decoding."""
import copy
import logging
from contextlib import ExitStack
from typing import Optional
import torch
from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.layernorm import GemmaRMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.qwen3_5_mtp import Qwen3_5ForCausalLMMTP, _mtp_quant_config
from sglang.srt.models.qwen4_exp import Qwen4ExpModel
from sglang.srt.runtime_context import get_model, get_parallel
from sglang.srt.utils import add_prefix, is_npu
logger = logging.getLogger(__name__)
class Qwen4ExpForCausalLMMTP(Qwen3_5ForCausalLMMTP):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
nn.Module.__init__(self)
self.is_multimodal = hasattr(config, "text_config")
if self.is_multimodal:
config = config.text_config
# Deepcopy so MTP-only mutations below don't leak into the main model.
config = copy.deepcopy(config)
config.num_hidden_layers = 1
config.layer_types = ["full_attention"]
config.full_attention_interval = 1
config.ple_layer_ids = []
quant_config = _mtp_quant_config(quant_config)
self.config = config
self.tp_size = get_parallel().tp_size
self.quant_config = quant_config
self.pp_group = get_pp_group()
self.hidden_size = config.hidden_size
self.hc_count = config.hc_count
self._mtp_input_fusion = self._init_mtp_input_fusion(config)
self.model = Qwen4ExpModel(
config,
quant_config,
prefix=add_prefix("mtp", prefix),
is_nextn=True,
)
self.lm_head = ParallelLMHead(
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_parallel().enable_dp_lm_head,
)
self.logits_processor = LogitsProcessor(config)
def _init_pre_fc_norms(self, config: PretrainedConfig) -> None:
self.pre_fc_norm_embedding = GemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
hidden_norm_size = (
self.hc_count * config.hidden_size
if self.hc_count > 1
else config.hidden_size
)
self.pre_fc_norm_hidden = GemmaRMSNorm(
hidden_norm_size, eps=config.rms_norm_eps
)
def _init_linear_projections(self, config: PretrainedConfig) -> None:
self.fc_embedding = nn.Linear(
config.hidden_size, config.hidden_size, bias=False
)
self.fc_hidden = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
def _init_standard_fusion(self, config: PretrainedConfig):
self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
self._init_pre_fc_norms(config)
return self._fuse_standard
def _init_mtp_input_fusion(self, config: PretrainedConfig):
if self.hc_count <= 1:
return self._init_standard_fusion(config)
self._init_linear_projections(config)
self._init_pre_fc_norms(config)
return self._fuse_residual_linear_shared
def _fuse_residual_linear_shared(
self, input_embeds: torch.Tensor, hidden_states: torch.Tensor
) -> torch.Tensor:
input_embeds = self.fc_embedding(self.pre_fc_norm_embedding(input_embeds))
orig_shape = hidden_states.shape
hidden_states = self.pre_fc_norm_hidden(hidden_states)
decoder_view = hidden_states.view(
*hidden_states.shape[:-1], self.hc_count, self.hidden_size
)
encoder_inputs = self.fc_hidden(decoder_view)
return (input_embeds.unsqueeze(-2) + encoder_inputs).view(orig_shape)
def _fuse_standard(
self, input_embeds: torch.Tensor, hidden_states: torch.Tensor
) -> torch.Tensor:
input_embeds = self.pre_fc_norm_embedding(input_embeds)
hidden_states = self.pre_fc_norm_hidden(hidden_states)
return self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
def _npu_quant_context(self):
exit_stack = ExitStack()
if (
is_npu()
and self.quant_config is None
and get_model().quantization is not None
):
exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
exit_stack.enter_context(
envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
)
return exit_stack
def _prepare_input_embeds(
self,
input_ids: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: Optional[torch.Tensor],
) -> torch.Tensor:
assert input_embeds is None
input_embeds = forward_batch.mm_input_embeds
if (
forward_batch.forward_mode.is_extend()
and forward_batch.contains_mm_inputs()
and not forward_batch.forward_mode.is_draft_extend_v2()
):
assert input_embeds is not None
last_indices = (
forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1
).long()
input_embeds[last_indices] = self.model.embed_tokens(
input_ids[last_indices]
)
if input_embeds is None:
input_embeds = self.model.embed_tokens(input_ids)
return input_embeds
def _set_hc_logits_hidden_states(
self,
logits_output,
hc_hidden_states: Optional[torch.Tensor],
forward_batch: ForwardBatch,
) -> None:
if hc_hidden_states is None:
return
# The EAGLE v2 future map holds one hidden state per request;
# reduce a token-shaped draft-extend HC tensor to its last token per request.
if forward_batch.forward_mode.is_draft_extend_v2():
# Mirror LogitsProcessor: the graph path selects rows via
# spec_info.select_index and the worker no longer re-indexes them.
select_index = forward_batch.spec_info.select_index
if select_index is not None:
hc_hidden_states = hc_hidden_states[select_index]
elif (
forward_batch.extend_seq_lens is not None
and hc_hidden_states.shape[0] != forward_batch.extend_seq_lens.shape[0]
):
last_index = (
torch.cumsum(forward_batch.extend_seq_lens.to(torch.int64), dim=0) - 1
)
hc_hidden_states = hc_hidden_states[last_index]
assert hc_hidden_states.shape[-1] == self.hc_count * self.hidden_size
logits_output.hidden_states = hc_hidden_states
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: Optional[torch.Tensor] = None,
**kwargs,
):
with self._npu_quant_context():
input_embeds = self._prepare_input_embeds(
input_ids, forward_batch, input_embeds
)
hidden_states = forward_batch.spec_info.hidden_states
if not forward_batch.forward_mode.is_idle():
hidden_states = self._mtp_input_fusion(input_embeds, hidden_states)
with get_global_expert_distribution_recorder().disable_this_region():
model_output = self.model(
input_ids,
positions,
forward_batch,
hidden_states,
)
hc_hidden_states = None
if isinstance(model_output, tuple):
hidden_states, hc_hidden_states = model_output
else:
hidden_states = model_output
logits_output = self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch
)
self._set_hc_logits_hidden_states(
logits_output, hc_hidden_states, forward_batch
)
return logits_output
EntryClass = [Qwen4ExpForCausalLMMTP]
python/sglang/srt/multimodal/processors/qwen_vl.py
+6
View File
@@ -32,6 +32,7 @@ from sglang.srt.models.qwen3_5_mtp import Qwen3_5ForCausalLMMTP
from sglang.srt.models.qwen3_omni_moe import Qwen3OmniMoeForConditionalGeneration
from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
from sglang.srt.models.qwen3_vl_moe import Qwen3VLMoeForConditionalGeneration
from sglang.srt.models.qwen4_exp import Qwen4ExpForConditionalGeneration
from sglang.srt.multimodal.processors.base_processor import (
BaseMultimodalProcessor as SGLangBaseProcessor,
)
@@ -301,6 +302,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
InternS2MobiusForConditionalGeneration,
Qwen3OmniMoeForConditionalGeneration,
Cosmos3ForConditionalGeneration,
Qwen4ExpForConditionalGeneration,
]
def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
@@ -312,6 +314,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
"qwen3_vl_moe",
"qwen3_5",
"qwen3_5_moe",
"qwen4_exp",
"intern_s2_preview",
"interns2_mobius",
):
@@ -522,6 +525,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
"qwen3_vl_moe",
"qwen3_5",
"qwen3_5_moe",
"qwen4_exp",
"intern_s2_preview",
"interns2_mobius",
"cosmos3_omni",
@@ -659,6 +663,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
"qwen3_vl_moe",
"qwen3_5",
"qwen3_5_moe",
"qwen4_exp",
"intern_s2_preview",
"cosmos3_omni",
]
@@ -768,6 +773,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
"qwen3_vl_moe",
"qwen3_5",
"qwen3_5_moe",
"qwen4_exp",
"intern_s2_preview",
"interns2_mobius",
"cosmos3_omni",
python/sglang/srt/speculative/draft_utils.py
+44
View File
@@ -1,3 +1,6 @@
from typing import Optional
from sglang.srt.layers.attention.qsa.config import QSA_VARIANT_COMPRESSED, QSAProfile
from sglang.srt.runtime_context import attention_backends, get_spec
from sglang.srt.utils.common import (
cpu_has_amx_support,
@@ -31,11 +34,13 @@ class DraftBackendFactory:
topk: int,
speculative_num_steps: int,
seed_dsa_topk_from_draft_extend: bool = False,
qsa_profile: Optional[QSAProfile] = None,
):
self.draft_model_runner = draft_model_runner
self.topk = topk
self.speculative_num_steps = speculative_num_steps
self.seed_dsa_topk_from_draft_extend = seed_dsa_topk_from_draft_extend
self.qsa_profile = qsa_profile
# The draft runner's own backend, not the process-wide config.
self.draft_attn_backend = draft_model_runner.draft_attention_backend
@@ -81,6 +86,9 @@ class DraftBackendFactory:
if self.speculative_num_steps <= 1:
return None
if self.qsa_profile is not None:
return self._create_qwen_qsa_decode_backend()
# Returns a per-step CONTAINER, not an AttentionBackend, so
# attn_backend_wrapper_for_draft_extend cannot give it a conv sidecar.
_assert_draft_needs_no_conv_sidecar(self.draft_model_runner)
@@ -112,6 +120,9 @@ class DraftBackendFactory:
)
def create_draft_extend_backend(self):
if self.qsa_profile is not None:
return self._create_qwen_qsa_draft_extend_backend()
backend_map = {
"flashinfer": self._create_flashinfer_prefill_backend,
"triton": self._create_triton_prefill_backend,
@@ -156,6 +167,39 @@ class DraftBackendFactory:
wrapped.decode_attention_backend_str = backend.decode_attention_backend_str
return wrapped
@staticmethod
def _stamp_qsa(backend) -> None:
backend.prefill_attention_backend_str = "qsa"
backend.decode_attention_backend_str = "qsa"
def _create_qwen_qsa_draft_extend_backend(self):
if self.qsa_profile.variant != QSA_VARIANT_COMPRESSED:
# Tokenwise QSA has no graph-stable indexer metadata: draft extend
# stays eager instead of falling back to a dense backend.
return None
from sglang.srt.layers.attention.qwen_sparse_attn_backend import (
QwenSparseAttnBackend,
)
# The draft is full-attention only: give it a QSA backend of its own
# instead of the hybrid wrapper whose linear side has no draft layers.
backend = QwenSparseAttnBackend(self.draft_model_runner)
self._stamp_qsa(backend)
return backend
def _create_qwen_qsa_decode_backend(self):
from sglang.srt.layers.attention.qwen_sparse_attn_backend import (
QwenSparseMultiStepDraftBackend,
)
backend = QwenSparseMultiStepDraftBackend(
self.draft_model_runner, self.topk, self.speculative_num_steps
)
self._stamp_qsa(backend)
for child in backend.attn_backends:
self._stamp_qsa(child)
return backend
def _create_dsa_decode_backend(self):
from sglang.srt.layers.attention.dsa_backend import (
DeepseekSparseAttnMultiStepBackend,
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+4 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass
from types import SimpleNamespace
from typing import TYPE_CHECKING, Callable, Optional
import torch
@@ -667,7 +668,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
# Prepare per-step draft attention metadata (kv_indptr / kv_indices for
# each speculative step). The glue-graph optimisation is not applied
# here — see __init__ comment for why.
self.draft_attn_backend.init_forward_metadata_out_graph(forward_batch)
self.draft_attn_backend.init_forward_metadata_out_graph(
SimpleNamespace(**vars(forward_batch), num_padding=bs - raw_bs)
)
self.raw_bs = raw_bs
self.bs = bs
python/sglang/srt/speculative/eagle_worker_v2.py
+71
View File
@@ -21,6 +21,12 @@ from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGra
from sglang.srt.kv_canary.runner.canary_manager import context_tuple
from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
from sglang.srt.layers.attention.qsa.config import parse_qsa_profile
from sglang.srt.layers.attention.qwen_sparse_attn_backend import (
QSAMTPSharedSparseIndices,
QwenSparseAttnBackend,
QwenSparseMultiStepDraftBackend,
)
from sglang.srt.layers.attention.tokenspeed_mla_backend import TokenspeedMLABackend
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
@@ -135,6 +141,19 @@ _is_xpu = is_xpu()
logger = logging.getLogger(__name__)
def _qsa_index_share_requested(hf_config) -> bool:
"""--json-model-override-args writes top-level hf_config attributes, while
checkpoint configs carry the flag on the nested text_config; read both."""
text_config = getattr(hf_config, "text_config", hf_config)
return bool(
getattr(
text_config,
"index_share_for_mtp_iteration",
getattr(hf_config, "index_share_for_mtp_iteration", False),
)
)
class EagleDraftWorker(EagleDraftWorkerBase):
def __init__(
self,
@@ -351,6 +370,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self.topk,
self.speculative_num_steps,
seed_dsa_topk_from_draft_extend=self.seed_dsa_topk_from_draft_extend,
qsa_profile=parse_qsa_profile(self.draft_runner.model_config.hf_config),
)
# Initialize decode attention backend
@@ -364,8 +384,58 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self.draft_runner.draft_attn_backend = self.draft_attn_backend
if self.draft_extend_attn_backend is not None:
self.draft_runner.attn_backend = self.draft_extend_attn_backend
self._configure_qsa_mtp_index_share()
self.tree_mask_mode = default_tree_mask_mode()
def _configure_qsa_mtp_index_share(self) -> None:
"""Reuse the draft-extend QSA selection across the MTP decode steps;
chain speculation only: with topk > 1 decode rows are not request-major."""
from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer
hf_config = self.draft_runner.model_config.hf_config
if (
not _qsa_index_share_requested(hf_config)
or self.topk != 1
or self.speculative_num_steps <= 1
or not isinstance(self.draft_attn_backend, QwenSparseMultiStepDraftBackend)
or not isinstance(self.draft_extend_attn_backend, QwenSparseAttnBackend)
):
return
if get_spec().speculative_adaptive:
# Adaptive speculation switches SpecRuntimeState between the draft-extend
# capture and the decode lookup; per-state index buffers would not match.
logger.warning(
"index_share_for_mtp_iteration is disabled under adaptive "
"speculative decoding"
)
return
layer_ids = sorted(
{
module.layer_id
for module in self.draft_runner.model.modules()
if isinstance(module, QSAIndexer)
}
)
if not layer_ids:
return
pool = self.draft_runner.token_to_kv_pool
# The expansion emits token_topk + ratio - 1 columns (top-k blocks
# plus the uncompressed tail of the capture position).
expanded_width = pool.qsa_token_topk + pool.qsa_compress_ratio - 1
state = QSAMTPSharedSparseIndices(
layer_ids=layer_ids,
num_requests=self.draft_runner.req_to_token_pool.req_to_token.shape[0],
token_topk=expanded_width,
tail_width=get_spec().speculative_num_steps + 1,
device=self.draft_runner.device,
)
for backend in (self.draft_attn_backend, self.draft_extend_attn_backend):
backend.set_mtp_shared_sparse_indices(state)
logger.info(
"QSA MTP index sharing enabled: draft decode steps reuse the "
f"draft-extend selection for layers {layer_ids}"
)
def _capture_cuda_graphs(self):
"""Capture the draft worker's own cuda graphs (decode + draft-extend)."""
self.cuda_graph_runner = None
@@ -450,6 +520,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
TRTLLMHAAttnBackend,
TokenspeedMLABackend,
FlashInferAttnBackend,
QwenSparseAttnBackend,
]
if _is_cuda or _is_musa:
# DSA is CUDA-only; import lazily so non-CUDA builds don't pull in
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+13
View File
@@ -26,6 +26,7 @@ from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.multi_layer_eagle_draft_extend_npu_graph_runner import (
MultiLayerEagleMultiStepDraftExtendNpuGraphRunner,
)
from sglang.srt.layers.attention.qsa.config import parse_qsa_profile
from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
speculative_moe_backend_context,
@@ -368,6 +369,17 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
self.draft_runner_list[i].model.set_embed_and_head(embed, head)
def init_attention_backend(self):
from sglang.srt.speculative.eagle_worker_v2 import (
_qsa_index_share_requested,
)
hf_config = self.draft_runner_list[0].model_config.hf_config
if _qsa_index_share_requested(hf_config):
logger.warning(
"index_share_for_mtp_iteration is not supported with "
"multi-layer EAGLE; the draft indexer runs every step"
)
qsa_profile = parse_qsa_profile(hf_config)
# Create attn backends
self.draft_extend_attn_backend_list = []
for step in range(self.speculative_num_steps):
@@ -375,6 +387,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
self.draft_runner_list[step],
self.topk,
self.speculative_num_steps,
qsa_profile=qsa_profile,
)
self.draft_extend_attn_backend_list.append(
draft_backend_factory.create_draft_extend_backend()
python/sglang/srt/speculative/spec_utils.py
+17
View File
@@ -805,6 +805,23 @@ def _verify_commit_step_indices(
mamba-track interval-crossing step (-1 = no crossing; None when tracking
is off)."""
bs = accept_lens.shape[0]
if accept_index.is_cuda:
from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
fused_commit_track_indices,
)
track_grid = (
mamba_track_grid(batch.tree_cache.page_size)
if batch.mamba_track_indices is not None
else 0
)
return fused_commit_track_indices(
accept_index,
accept_lens,
batch.seq_lens if track_grid > 0 else None,
draft_token_num,
track_grid,
)
accept_indices_offset = torch.arange(
0,
bs * draft_token_num,
python/sglang/srt/utils/common.py
+6
View File
@@ -319,6 +319,12 @@ is_sm90_supported = lru_cache(maxsize=1)(
)
# RTX Blackwell. Unlike is_sm120_supported(), this excludes SM121/GB10.
@lru_cache(maxsize=1)
def is_sm120() -> bool:
return is_cuda() and torch.cuda.get_device_capability() == (12, 0)
# GB10 (DGX Spark and OEM equivalents). Not expressible via
# _check_cuda_device_version, which only matches on the major.
@lru_cache(maxsize=1)
python/sglang/srt/utils/hf_transformers/common.py
+4
View File
@@ -76,6 +76,8 @@ from sglang.srt.configs import (
Qwen3_5MoeTextConfig,
Qwen3_5TextConfig,
Qwen3NextConfig,
Qwen4ExpConfig,
Qwen4ExpTextConfig,
Spark2_5Config,
Step3p5Config,
Step3p7Config,
@@ -126,6 +128,8 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
Glm5NextTextConfig,
KimiLinearConfig,
Qwen3NextConfig,
Qwen4ExpConfig,
Qwen4ExpTextConfig,
FalconH1Config,
GraniteMoeHybridConfig,
HYV4Config,
python/sglang/test/run_eval.py
+2
View File
@@ -299,6 +299,8 @@ def run_eval(args):
return _run_sgl_eval("mmmu_pro_vision", args)
elif args.eval_name == "aime25":
return _run_sgl_eval("aime25", args)
elif args.eval_name == "aime26":
return _run_sgl_eval("aime26", args)
elif args.eval_name == "gsm8k":
if getattr(args, "api", None) == "sgl_eval":
# Only the nightly correctness eval opts into sgl-eval (zero-shot
test/registered/e2e/models/test_qwen4_exp_models.py
+90
View File
@@ -0,0 +1,90 @@
"""Qwen3.8-Flash-Next (Qwen4-Exp) E2E on B200; the plain-serving case is kept:
MTP's verify widths never exercise the QSA sparse-decode path."""
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
try_cached_model,
)
register_cuda_ci(est_time=1500, stage="base-c", runner_config="4-gpu-b200")
MODEL = "RadixArk/Qwen3.8-Flash-Next-NVFP4"
SERVER_LAUNCH_TIMEOUT = 3600
GSM8K_SCORE_THRESHOLD = 0.94
BASE_ARGS = [
"--tp-size",
"4",
"--mem-fraction-static",
"0.85",
"--chunked-prefill-size",
"8192",
"--linear-attn-prefill-backend",
"flashinfer",
"--linear-attn-decode-backend",
"flashinfer",
"--mamba-ssm-dtype",
"bfloat16",
"--reasoning-parser",
"qwen3-thinking",
]
class _Qwen4ExpServer:
speculative_args: list[str] = []
model = try_cached_model(MODEL)
base_url = DEFAULT_URL_FOR_TEST
gsm8k_backend = "sgl_eval"
gsm8k_thinking = True
gsm8k_num_examples = 200
gsm8k_num_threads = 32
gsm8k_max_tokens = 16384
gsm8k_score_threshold = GSM8K_SCORE_THRESHOLD
@classmethod
def setUpClass(cls):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=BASE_ARGS + cls.speculative_args,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
class TestQwen4ExpBase(_Qwen4ExpServer, GSM8KMixin, CustomTestCase):
"""Normal autoregressive serving."""
class TestQwen4ExpMTP(_Qwen4ExpServer, GSM8KMixin, CustomTestCase):
"""NEXTN MTP serving (3 steps, topk 1, 4 draft tokens)."""
# GSM8K accept length measured at 3.02-3.03 (max 4.0 with 3 steps);
# 2.9 leaves noise margin while still failing on a real drop.
gsm8k_accept_length_thres = 2.9
speculative_args = [
"--speculative-algorithm",
"NEXTN",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
]
if __name__ == "__main__":
unittest.main()
test/registered/kernel/embeddings/test_qwen4_ple_offload.py
+195
View File
@@ -0,0 +1,195 @@
from types import SimpleNamespace
import pytest
import torch
from torch import nn
from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbeddingShardIndices,
)
from sglang.srt.models import qwen4_exp as qwen4_exp_module
from sglang.srt.models.qwen4_exp import (
Qwen4ExpPinnedHostEmbedding,
Qwen4ExpPLELayer,
)
from sglang.srt.utils import set_weight_attrs
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available(), reason="CUDA is required for this test."
)
def _make_source_embedding(
*,
dtype=torch.bfloat16,
embedding_dim=7,
vocab_start=0,
vocab_end=8,
org_vocab_size=8,
tp_size=1,
num_added_embeddings=0,
):
local_rows = vocab_end - vocab_start
weight = nn.Parameter(
torch.empty((local_rows, embedding_dim), dtype=dtype, device="cuda"),
requires_grad=False,
)
set_weight_attrs(
weight,
{
"input_dim": 1,
"output_dim": 0,
"weight_loader": lambda *_args, **_kwargs: None,
},
)
shard_indices = VocabParallelEmbeddingShardIndices(
padded_org_vocab_start_index=vocab_start,
padded_org_vocab_end_index=vocab_end,
padded_added_vocab_start_index=org_vocab_size,
padded_added_vocab_end_index=org_vocab_size,
org_vocab_start_index=vocab_start,
org_vocab_end_index=vocab_end,
added_vocab_start_index=org_vocab_size,
added_vocab_end_index=org_vocab_size,
)
return SimpleNamespace(
weight=weight,
quant_config=None,
enable_tp=True,
use_attn_tp_group=False,
tp_size=tp_size,
num_embeddings=org_vocab_size + num_added_embeddings,
org_vocab_size=org_vocab_size,
padding_size=1,
num_added_embeddings=num_added_embeddings,
use_presharded_weights=False,
org_vocab_size_padded=org_vocab_size,
num_embeddings_padded=org_vocab_size + num_added_embeddings,
shard_indices=shard_indices,
embedding_dim=embedding_dim,
weight_scale=None,
quant_method=UnquantizedEmbeddingMethod(),
num_embeddings_per_partition=local_rows,
num_org_embeddings_per_partition=local_rows,
num_added_embeddings_per_partition=0,
)
def _load_rows(offloaded, rows):
pointer = offloaded.weight.data_ptr()
offloaded.weight_loader(offloaded.weight, rows)
assert offloaded.weight.data_ptr() == pointer
assert offloaded.weight.is_pinned()
assert offloaded.weight.weight_loader.__self__ is offloaded
assert offloaded.quant_method is None
@pytest.mark.parametrize("input_dtype", [torch.int32, torch.int64])
@pytest.mark.parametrize("embedding_dim", [7, 64, 257])
def test_qwen4_ple_pinned_gather_tp1(input_dtype, embedding_dim):
source = _make_source_embedding(embedding_dim=embedding_dim)
offloaded = Qwen4ExpPinnedHostEmbedding(source)
rows = torch.arange(8 * embedding_dim, dtype=torch.bfloat16, device="cuda").reshape(
8, embedding_dim
)
_load_rows(offloaded, rows)
ids = torch.tensor([[0, 7, 3], [4, 1, 6]], dtype=input_dtype, device="cuda")
expected = rows.index_select(0, ids.long().flatten()).reshape(
*ids.shape, embedding_dim
)
actual = offloaded(ids)
assert actual.shape == expected.shape
assert actual.is_contiguous()
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_qwen4_ple_pinned_gather_shard_boundaries_and_out_buffer():
embedding_dim = 13
source = _make_source_embedding(
embedding_dim=embedding_dim,
vocab_start=4,
vocab_end=8,
org_vocab_size=8,
tp_size=2,
)
offloaded = Qwen4ExpPinnedHostEmbedding(source)
rows = torch.arange(8 * embedding_dim, dtype=torch.bfloat16, device="cuda").reshape(
8, embedding_dim
)
_load_rows(offloaded, rows)
ids = torch.tensor([[-1, 3, 4], [7, 8, 100]], device="cuda")
output = torch.full(
(*ids.shape, embedding_dim),
torch.nan,
dtype=torch.bfloat16,
device="cuda",
)
actual = offloaded.gather(ids, out=output)
expected = torch.zeros_like(output)
expected[0, 2] = rows[4]
expected[1, 0] = rows[7]
assert actual.data_ptr() == output.data_ptr()
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_qwen4_ple_pinned_gather_empty_input():
offloaded = Qwen4ExpPinnedHostEmbedding(_make_source_embedding())
_load_rows(offloaded, torch.zeros((8, 7), dtype=torch.bfloat16, device="cuda"))
ids = torch.empty((0, 3), dtype=torch.int64, device="cuda")
actual = offloaded.gather(ids)
assert actual.shape == (0, 3, 7)
assert actual.numel() == 0
def test_qwen4_ple_pinned_embedding_rejects_unsupported_weights():
with pytest.raises(TypeError, match="requires bfloat16"):
Qwen4ExpPinnedHostEmbedding(_make_source_embedding(dtype=torch.float16))
with pytest.raises(NotImplementedError, match="added vocabulary"):
Qwen4ExpPinnedHostEmbedding(_make_source_embedding(num_added_embeddings=1))
def test_qwen4_ple_prefetch_buffer_lifecycle(monkeypatch):
layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
nn.Module.__init__(layer)
layer.ple_embed_dim = 7
layer.ple_embedding = SimpleNamespace(
ngram_embedding=Qwen4ExpPinnedHostEmbedding(
_make_source_embedding(embedding_dim=layer.ple_embed_dim)
)
)
layer._graph_prefetch_buffers = {}
layer._eager_prefetch_buffer = None
lookup_ids = torch.empty((0,), dtype=torch.int64, device="cuda")
monkeypatch.setattr(qwen4_exp_module, "get_is_capture_mode", lambda: False)
eager_large = layer._get_prefetch_buffer(8, lookup_ids)
eager_small = layer._get_prefetch_buffer(3, lookup_ids)
assert eager_small.data_ptr() == eager_large.data_ptr()
assert layer._eager_prefetch_buffer.shape == (8, layer.ple_embed_dim)
eager_grown = layer._get_prefetch_buffer(12, lookup_ids)
eager_grown_small = layer._get_prefetch_buffer(4, lookup_ids)
assert eager_grown_small.data_ptr() == eager_grown.data_ptr()
assert layer._eager_prefetch_buffer.shape == (12, layer.ple_embed_dim)
monkeypatch.setattr(qwen4_exp_module, "get_is_capture_mode", lambda: True)
graph_three = layer._get_prefetch_buffer(3, lookup_ids)
graph_five = layer._get_prefetch_buffer(5, lookup_ids)
graph_three_reused = layer._get_prefetch_buffer(3, lookup_ids)
assert graph_three_reused.data_ptr() == graph_three.data_ptr()
assert graph_five.data_ptr() != graph_three.data_ptr()
assert set(layer._graph_prefetch_buffers) == {3, 5}
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernel/hyperconnection/test_hc_mix_triton.py
+84
View File
@@ -0,0 +1,84 @@
import sys
import pytest
import torch
import torch.nn.functional as F
from sglang.srt.layers.hc_mix_triton import (
_FUSED_MIX_MAX_ROWS,
fused_hc_mix,
fused_hc_mix_supported,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
HC_COUNT = 4
HIDDEN_SIZE = 2560
LOWRANK = 320
def _reference_mix(
hyper_input_normed: torch.Tensor,
w_down: torch.Tensor,
w_up: torch.Tensor,
hc: int,
hs: int,
compute_dtype: torch.dtype = torch.float64,
) -> torch.Tensor:
"""Mirrors GatedResidual._mix_compute in hyperconnection.py."""
x = hyper_input_normed.to(compute_dtype)
t = F.silu(F.linear(x, w_down.to(compute_dtype)) / hc)
u = torch.sigmoid(F.linear(t, w_up.to(compute_dtype)))
return (u.unflatten(-1, (hc, hs)) * x.unflatten(-1, (hc, hs))).mean(dim=-2)
def _make_inputs(num_tokens: int, dtype: torch.dtype):
torch.manual_seed(0)
x = torch.randn(num_tokens, HC_COUNT * HIDDEN_SIZE, dtype=dtype, device="cuda")
w_down = (
torch.randn(LOWRANK, HC_COUNT * HIDDEN_SIZE, dtype=dtype, device="cuda") * 0.02
)
w_up = (
torch.randn(HC_COUNT * HIDDEN_SIZE, LOWRANK, dtype=dtype, device="cuda") * 0.02
)
return x, w_down, w_up
_TOLERANCES = {
torch.bfloat16: dict(rtol=1e-2, atol=5e-3),
torch.float16: dict(rtol=2e-3, atol=1e-3),
}
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("num_tokens", [1, 4, 7, _FUSED_MIX_MAX_ROWS])
def test_fused_hc_mix_matches_reference(dtype, num_tokens):
x, w_down, w_up = _make_inputs(num_tokens, dtype)
assert fused_hc_mix_supported(x, w_down, w_up)
out = fused_hc_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
ref = _reference_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
torch.testing.assert_close(out.to(torch.float64), ref, **_TOLERANCES[dtype])
def test_fused_hc_mix_no_less_accurate_than_eager():
"""The fused kernel (fp32 accumulation throughout) must not be farther
from the fp64 reference than the eager bf16 chain it replaces."""
x, w_down, w_up = _make_inputs(8, torch.bfloat16)
ref = _reference_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
fused = fused_hc_mix(x, w_down, w_up, HC_COUNT, HIDDEN_SIZE)
eager = _reference_mix(
x, w_down, w_up, HC_COUNT, HIDDEN_SIZE, compute_dtype=torch.bfloat16
)
fused_err = (fused.to(torch.float64) - ref).abs().max()
eager_err = (eager.to(torch.float64) - ref).abs().max()
assert fused_err <= eager_err * 1.5 + 1e-6
def test_fused_hc_mix_gate_rejects_prefill_rows():
x, w_down, w_up = _make_inputs(_FUSED_MIX_MAX_ROWS + 1, torch.bfloat16)
assert not fused_hc_mix_supported(x, w_down, w_up)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernel/jit/test_fast_topk.py
+141
View File
@@ -0,0 +1,141 @@
import pytest
import torch
from sglang.kernels.ops.elementwise.fast_topk import fast_topk
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
def _check_topk_values(score, lengths, indices, topk, row_starts):
"""fast_topk leaves order and tie-breaking unspecified,
so compare the sorted top-k values rather than index sets."""
for b in range(score.shape[0]):
start = int(row_starts[b]) if row_starts is not None else 0
length = int(lengths[b])
section = score[b, start : start + length]
row = indices[b]
if length <= topk:
# naive path: identity indices, then -1 fill
assert torch.equal(
row[:length].cpu(), torch.arange(length, dtype=torch.int32)
)
assert (row[length:] == -1).all()
continue
assert (row >= 0).all(), "long rows must fill every slot"
picked = section[row.long()]
expected = torch.topk(section, topk).values
assert torch.equal(
picked.sort(descending=True).values, expected.sort(descending=True).values
), f"row {b}: top-{topk} value multiset mismatch"
@pytest.mark.parametrize("topk", [512, 2048])
@pytest.mark.parametrize(
"batch,length",
[
(1, 4096),
(7, 3000),
(33, 32768),
(128, 2050),
],
)
def test_fast_topk_long_rows(topk, batch, length):
torch.manual_seed(0)
score = torch.randn(batch, length, dtype=torch.float32, device="cuda")
lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
indices = fast_topk(score, lengths, topk)
_check_topk_values(score, lengths, indices, topk, None)
@pytest.mark.parametrize("topk", [512, 2048])
def test_fast_topk_short_and_mixed_rows(topk):
torch.manual_seed(0)
max_len = topk + 128
batch = 8
score = torch.randn(batch, max_len, dtype=torch.float32, device="cuda")
# rows shorter than k (naive path), exactly k, and longer than k
lens = [1, topk // 3, topk - 1, topk, topk + 1, topk + 7, 17, max_len]
lengths = torch.tensor(lens[:batch], dtype=torch.int32, device="cuda")
indices = fast_topk(score, lengths, topk)
_check_topk_values(score, lengths, indices, topk, None)
@pytest.mark.parametrize("topk", [512, 2048])
def test_fast_topk_ragged_with_row_starts(topk):
torch.manual_seed(0)
batch, width = 16, 8192
score = torch.randn(batch, width, dtype=torch.float32, device="cuda")
row_starts = torch.randint(0, 2048, (batch,), dtype=torch.int32, device="cuda")
lengths = torch.randint(1, 2048, (batch,), dtype=torch.int32, device="cuda")
lengths = torch.minimum(lengths, width - row_starts).to(torch.int32)
# ensure some rows are longer than k
lengths[0] = min(width - int(row_starts[0]), topk + 100)
indices = fast_topk(score, lengths, topk, row_starts=row_starts)
_check_topk_values(score, lengths, indices, topk, row_starts)
@pytest.mark.parametrize("topk", [512, 2048])
def test_fast_topk_row_stride(topk):
torch.manual_seed(0)
batch, length = 8, 4096
base = torch.randn(batch, 2 * length, dtype=torch.float32, device="cuda")
score = base[:, :length] # stride(0) == 2*length, stride(1) == 1
lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
indices = fast_topk(score, lengths, topk)
_check_topk_values(score, lengths, indices, topk, None)
@pytest.mark.parametrize("topk", [512, 2048])
@pytest.mark.parametrize(
"fill",
[
"binary", # only 0s and 1s: extreme duplication at the threshold bin
"few_levels", # a handful of distinct levels incl. negatives
"constant", # whole rows of one value
],
)
def test_fast_topk_duplicate_heavy(topk, fill):
torch.manual_seed(0)
batch, length = 16, 8192
if fill == "binary":
score = torch.randint(0, 2, (batch, length), dtype=torch.float32, device="cuda")
elif fill == "few_levels":
levels = torch.tensor([-5.0, -1.0, 0.0, 0.5, 2.0], device="cuda")
score = levels[torch.randint(0, 5, (batch, length), device="cuda")]
else:
score = torch.full((batch, length), 3.25, dtype=torch.float32, device="cuda")
lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
indices = fast_topk(score, lengths, topk)
_check_topk_values(score, lengths, indices, topk, None)
@pytest.mark.parametrize("topk", [512, 2048])
def test_fast_topk_negative_and_zero(topk):
torch.manual_seed(0)
batch, length = 8, 16384
score = torch.randn(batch, length, dtype=torch.float32, device="cuda") * 100
score[:, : length // 3] = 0.0 # long zero prefix
score[:, length // 3 : length // 2] = -1e30 # very negative block
lengths = torch.full((batch,), length, dtype=torch.int32, device="cuda")
indices = fast_topk(score, lengths, topk)
_check_topk_values(score, lengths, indices, topk, None)
def test_fast_topk_unsupported_k():
score = torch.randn(2, 4096, dtype=torch.float32, device="cuda")
lengths = torch.full((2,), 4096, dtype=torch.int32, device="cuda")
with pytest.raises(RuntimeError, match="topk"):
fast_topk(score, lengths, 1024)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernel/jit/test_grouped_gemma_rmsnorm.py
+87
View File
@@ -0,0 +1,87 @@
import pytest
import torch
from sglang.kernels.ops.layernorm.grouped_gemma_rmsnorm import grouped_gemma_rmsnorm
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
def _reference_grouped_gemma_rmsnorm(
x: torch.Tensor,
weight: torch.Tensor,
group_size: int,
eps: float,
compute_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""Mirrors GroupedGemmaRMSNorm.forward (hyperconnection.py); keep them in sync."""
x_float = x.to(compute_dtype)
hidden = x_float.shape[-1]
x_grouped = x_float.reshape(*x_float.shape[:-1], hidden // group_size, group_size)
variance = x_grouped.pow(2).mean(dim=-1, keepdim=True)
x_norm = (x_grouped * torch.rsqrt(variance + eps)).flatten(-2)
return x_norm * (1.0 + weight.to(compute_dtype))
# Tolerances are at the output-dtype quantization floor, measured against the
# fp64 reference on 4xB300 (sm103): bf16 max rel err 3.9e-3 (1 ulp), fp16
# 4.9e-4 (0.5 ulp). The kernel computes in fp32 like the eager reference.
_TOLERANCES = {
torch.bfloat16: dict(rtol=5e-3, atol=5e-3),
torch.float16: dict(rtol=1e-3, atol=1e-3),
}
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize(
"num_tokens,hidden_size,group_size",
[
(1, 10240, 2560), # production shape (HC 4 x 2560)
(7, 10240, 2560),
(128, 10240, 2560),
(33, 1024, 512),
(5, 512, 512), # single group == plain gemma rmsnorm
(1024, 2048, 1024),
],
)
@pytest.mark.parametrize("eps", [1e-6, 1e-5])
def test_grouped_gemma_rmsnorm_correctness(
dtype, num_tokens, hidden_size, group_size, eps
):
torch.manual_seed(0)
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda")
weight = torch.randn(hidden_size, dtype=dtype, device="cuda") * 0.2
out = grouped_gemma_rmsnorm(x, weight, group_size, eps)
expected = _reference_grouped_gemma_rmsnorm(
x, weight, group_size, eps, compute_dtype=torch.float64
).to(dtype)
torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
def test_grouped_gemma_rmsnorm_out_param():
x = torch.randn(64, 10240, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(10240, dtype=torch.bfloat16, device="cuda") * 0.2
out = torch.empty_like(x)
result = grouped_gemma_rmsnorm(x, weight, 2560, 1e-6, out=out)
expected = _reference_grouped_gemma_rmsnorm(
x, weight, 2560, 1e-6, compute_dtype=torch.float64
).to(x.dtype)
assert result.data_ptr() == out.data_ptr()
torch.testing.assert_close(result, expected, **_TOLERANCES[x.dtype])
def test_grouped_gemma_rmsnorm_bad_group_size():
x = torch.randn(4, 10240, dtype=torch.bfloat16, device="cuda")
weight = torch.zeros(10240, dtype=torch.bfloat16, device="cuda")
with pytest.raises(RuntimeError, match="group_size"):
grouped_gemma_rmsnorm(x, weight, 1000, 1e-6)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernel/jit/test_hc_combine.py
+161
View File
@@ -0,0 +1,161 @@
import pytest
import torch
import torch.nn.functional as F
from sglang.kernels.ops.elementwise.hc_combine import hc_combine
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
HC_COUNT = 4
HIDDEN_SIZE = 2560
def _reference_hc_combine(
block_output: torch.Tensor,
residual: torch.Tensor,
normed_residual: torch.Tensor,
inject_weight: torch.Tensor,
hc: int,
hs: int,
compute_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""Eager reference mirroring ``GatedResidual._combine_compute``;
``compute_dtype=torch.float64`` is the near-exact reference."""
R = residual.to(compute_dtype).unflatten(-1, (hc, hs))
gates = 2 * torch.sigmoid(
F.linear(normed_residual.to(compute_dtype), inject_weight.to(compute_dtype))
/ hc
)
injection = block_output.to(compute_dtype).unsqueeze(-2) * gates.unsqueeze(-1)
return (R + injection).flatten(-2)
def _make_inputs(
num_tokens: int, dtype: torch.dtype, hc: int = HC_COUNT, hs: int = HIDDEN_SIZE
):
torch.manual_seed(0)
block_output = torch.randn(num_tokens, hs, dtype=dtype, device="cuda")
residual = torch.randn(num_tokens, hc * hs, dtype=dtype, device="cuda")
normed_residual = torch.randn(num_tokens, hc * hs, dtype=dtype, device="cuda")
inject_weight = torch.randn(hc, hc * hs, dtype=dtype, device="cuda") * 0.02
return block_output, residual, normed_residual, inject_weight
# Worst case over M in {1, 7, 128, 8192} against the fp64 reference on B300 (sm103):
# bf16 max rel err 7.8e-3 (1 ulp at a binade edge), fp16 below 1e-3 (1 ulp = 9.8e-4);
# the residual is fp32 reordering flipping the final rounding.
_TOLERANCES = {
torch.bfloat16: dict(rtol=1e-2, atol=5e-3),
torch.float16: dict(rtol=1e-3, atol=1e-3),
}
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("num_tokens", [1, 7, 128, 8192])
def test_hc_combine_correctness(dtype, num_tokens):
block_output, residual, normed_residual, inject_weight = _make_inputs(
num_tokens, dtype
)
out = hc_combine(
block_output,
residual,
normed_residual,
inject_weight,
HC_COUNT,
HIDDEN_SIZE,
)
expected = _reference_hc_combine(
block_output,
residual,
normed_residual,
inject_weight,
HC_COUNT,
HIDDEN_SIZE,
compute_dtype=torch.float64,
).to(dtype)
torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_hc_combine_out_param(dtype):
block_output, residual, normed_residual, inject_weight = _make_inputs(64, dtype)
out = torch.empty_like(residual)
result = hc_combine(
block_output,
residual,
normed_residual,
inject_weight,
HC_COUNT,
HIDDEN_SIZE,
out=out,
)
expected = _reference_hc_combine(
block_output,
residual,
normed_residual,
inject_weight,
HC_COUNT,
HIDDEN_SIZE,
compute_dtype=torch.float64,
).to(dtype)
assert result.data_ptr() == out.data_ptr()
torch.testing.assert_close(result, expected, **_TOLERANCES[dtype])
def test_hc_combine_3d_input():
dtype = torch.bfloat16
block_output, residual, normed_residual, inject_weight = _make_inputs(32, dtype)
block_output = block_output.reshape(4, 8, HIDDEN_SIZE)
residual = residual.reshape(4, 8, HC_COUNT * HIDDEN_SIZE)
normed_residual = normed_residual.reshape(4, 8, HC_COUNT * HIDDEN_SIZE)
out = hc_combine(
block_output,
residual,
normed_residual,
inject_weight,
HC_COUNT,
HIDDEN_SIZE,
)
expected = _reference_hc_combine(
block_output,
residual,
normed_residual,
inject_weight,
HC_COUNT,
HIDDEN_SIZE,
compute_dtype=torch.float64,
).to(dtype)
assert out.shape == residual.shape
torch.testing.assert_close(out, expected, **_TOLERANCES[dtype])
def test_hc_combine_bad_hidden_size():
dtype = torch.bfloat16
block_output, residual, normed_residual, inject_weight = _make_inputs(
4,
dtype,
hc=4,
hs=1000, # 4 * 1000 = 4000, not a multiple of 2048
)
with pytest.raises(RuntimeError, match="2048"):
hc_combine(
block_output,
residual,
normed_residual,
inject_weight,
4,
1000,
)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernel/qsa/test_qsa.py
+1528
View File
File diff suppressed because it is too large Load Diff
test/registered/kernel/qsa/test_qsa_indexer.py
+273
View File
@@ -0,0 +1,273 @@
"""Fused QSA indexer-prep kernels must match the eager indexer path bit-for-bit,
up to rare last-ulp RMSNorm flips (see assert_bit_comparable)."""
from types import SimpleNamespace
import pytest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
from sglang.srt.layers.attention.qsa.kernel import (
average_pool_qsa_keys,
expand_qsa_block_indices,
torch_expand_qsa_block_indices,
)
from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer
from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
# MRotaryEmbedding reads the exec config bag at init; publish a minimal
# process context for the bare pytest process.
from sglang.srt.runtime_context import publish
from sglang.srt.server_args import ServerArgs
publish(ServerArgs(model_path="dummy"), role="test")
HEAD_DIM = 128
NUM_Q_HEADS = 4
RATIO = 4
HIDDEN = 2560
EPS = 1e-6
def _make_config():
return SimpleNamespace(
indexer_n_heads=NUM_Q_HEADS,
indexer_kv_heads=1,
indexer_head_dim=HEAD_DIM,
indexer_budget=2048,
indexer_compress_ratio=RATIO,
hidden_size=HIDDEN,
rms_norm_eps=EPS,
)
def _make_rotary(mrope_section, mrope_interleaved, device, dtype=torch.bfloat16):
return MRotaryEmbedding(
head_size=HEAD_DIM,
rotary_dim=HEAD_DIM,
max_position_embeddings=32768,
base=1000000,
is_neox_style=True,
dtype=dtype,
mrope_section=mrope_section,
mrope_interleaved=mrope_interleaved,
)
def _make_indexer(rotary, device, dtype=torch.bfloat16):
# Build under the model dtype like ModelRunner does; device-only .to()
# afterwards so the fp32 cos_sin_cache buffer keeps its dtype.
prev_dtype = torch.get_default_dtype()
torch.set_default_dtype(dtype)
try:
indexer = QSAIndexer(
_make_config(), layer_id=0, quant_config=None, rotary_emb=rotary
)
indexer.to(device=device)
finally:
torch.set_default_dtype(prev_dtype)
with torch.no_grad():
out_features = (NUM_Q_HEADS + 1) * HEAD_DIM
indexer.index_qk_proj.weight.data.copy_(
torch.randn(out_features, HIDDEN, device=device, dtype=dtype) * 0.02
)
for norm in (indexer.q_layernorm, indexer.k_layernorm):
w = torch.randn(HEAD_DIM, device=device, dtype=dtype) * 0.1
norm._weight_loader(norm.weight, w)
return indexer
class FakePool:
"""Minimal stand-in for the QSA KV pool buffers used by the indexer."""
index_state_dtype = torch.bfloat16
def __init__(self, num_slots, num_compressed, device, dtype=torch.bfloat16):
self.key_state = torch.zeros(num_slots, 1, HEAD_DIM, dtype=dtype, device=device)
self.qsa_rope_position_buffer = torch.zeros(
num_slots, 3, dtype=torch.int64, device=device
)
self.compressed = torch.zeros(
num_compressed, 1, HEAD_DIM, dtype=dtype, device=device
)
def get_qsa_key_state_buffer(self, layer_id):
return self.key_state
def set_qsa_key_state_buffer(self, layer_id, loc, token_k):
self.key_state[loc.long()] = token_k.to(self.key_state.dtype)
def set_qsa_rope_position_buffer(self, loc, positions):
positions = positions.long()
if positions.ndim == 1:
positions = positions.unsqueeze(0).expand(3, -1)
self.qsa_rope_position_buffer[loc.long()] = positions.transpose(0, 1)
def get_qsa_rope_position_buffer(self, loc):
return self.qsa_rope_position_buffer[loc.long()]
def get_qsa_compressed_k_buffer(self, layer_id):
return self.compressed
def set_qsa_compressed_k_buffer(self, layer_id, loc, compressed_k):
self.compressed[loc.long()] = compressed_k.to(self.compressed.dtype)
def assert_bit_comparable(actual, expected, max_frac=1e-5, max_abs=0.02):
"""Eager RMSNorm (flashinfer CuTe DSL) reduces in an unreproducible order,
so ~1 row in 30k flips by 1-2 bf16 ulp; max_frac and max_abs bound that."""
diff = (actual.float() - expected.float()).abs()
mismatches = int((diff > 0).sum())
allowed = max(16, int(max_frac * actual.numel()))
assert mismatches <= allowed, f"{mismatches} mismatched elements"
if mismatches:
peak = diff.max().item()
assert peak <= max_abs, f"largest deviation {peak} exceeds {max_abs}"
def _eager_compress_reference(indexer, pool, group_locs, write_locs):
"""The pre-fusion compression chain, via the indexer's own helpers."""
key_groups = pool.get_qsa_key_state_buffer(0)[group_locs.long()]
pooled = average_pool_qsa_keys(key_groups)
rope_positions = indexer._rope_from_matrix(
pool.get_qsa_rope_position_buffer(group_locs[:, 0])
)
normalized = indexer.normalize_compressed_keys(pooled, rope_positions)
pool.set_qsa_compressed_k_buffer(0, write_locs, normalized)
@pytest.mark.parametrize("num_groups", [1, 5, 2000])
@pytest.mark.parametrize(
"mrope_section, mrope_interleaved",
[([24, 20, 20], True), ([24, 20, 20], False), (None, False)],
)
def test_fused_compress_matches_eager(num_groups, mrope_section, mrope_interleaved):
device = torch.device("cuda")
dtype = torch.bfloat16
torch.manual_seed(num_groups)
rotary = _make_rotary(mrope_section, mrope_interleaved, device, dtype)
indexer = _make_indexer(rotary, device, dtype)
pool_ref = FakePool(8192, 4096, device, dtype)
pool_new = FakePool(8192, 4096, device, dtype)
pool_new.key_state.copy_(
pool_ref.key_state.copy_(
torch.randn(8192, 1, HEAD_DIM, device=device, dtype=dtype)
)
)
positions = torch.randint(0, 30000, (8192, 3), device=device)
pool_new.qsa_rope_position_buffer.copy_(positions)
pool_ref.qsa_rope_position_buffer.copy_(positions)
# Random groups; slot 0 doubles as the CUDA-graph dummy write target, so
# allow repeats there too.
group_locs = torch.randint(0, 8192, (num_groups, RATIO), device=device).to(
torch.int32
)
write_locs = torch.randperm(4096, device=device)[:num_groups].to(torch.int32)
_eager_compress_reference(indexer, pool_ref, group_locs, write_locs)
indexer._fused_compress_store(pool_new, group_locs, write_locs)
assert_bit_comparable(pool_new.compressed, pool_ref.compressed)
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
def test_expand_block_indices_int_inputs(dtype):
device = torch.device("cuda")
torch.manual_seed(0)
rows, block_topk, token_topk, ratio = 37, 512, 2048, 4
query_positions = torch.randint(0, 8000, (rows,), dtype=dtype, device=device)
sequence_lengths = (
query_positions + torch.randint(1, 9, (rows,), dtype=dtype, device=device)
).to(dtype)
# Production contract: top-k only selects blocks inside [0, seq_len//4),
# so no selected block ever masks out against sequence_lengths.
counts = torch.randint(0, block_topk + 1, (rows,))
block_indices = torch.full((rows, block_topk), -1, dtype=torch.int32)
seq_lens_host = sequence_lengths.cpu()
for r in range(rows):
limit = max(int(seq_lens_host[r]) // ratio, 1)
count = min(int(counts[r]), limit)
if count:
block_indices[r, :count] = torch.randperm(limit)[:count].to(torch.int32)
block_indices = block_indices.to(device)
out = expand_qsa_block_indices(
block_indices, query_positions, sequence_lengths, ratio, token_topk
)
ref = torch_expand_qsa_block_indices(
block_indices.cpu(),
query_positions.cpu(),
sequence_lengths.cpu(),
ratio,
token_topk,
)
assert torch.equal(out.cpu(), ref)
def test_decode_selection_equivalent():
"""Last-ulp norm flips must not change the selected blocks:
scores are fp32 sums of 128-dim dots, so a 1-ulp flip only matters on exact ties."""
from sglang.srt.layers.attention.qsa.kernel import qsa_fast_topk
from sglang.srt.layers.attention.qsa.mqa import torch_qsa_mqa_decode
device = torch.device("cuda")
dtype = torch.bfloat16
torch.manual_seed(7)
rotary = _make_rotary([24, 20, 20], True, device, dtype)
indexer = _make_indexer(rotary, device, dtype)
batch, max_pages, page_size = 4, 32, 64
max_model_len = max_pages * page_size
hidden = torch.randn(batch, HIDDEN, device=device, dtype=dtype)
positions = (
torch.arange(8000, 8000 + batch, device=device)
.unsqueeze(0)
.expand(3, -1)
.contiguous()
)
qk, _ = indexer.index_qk_proj(hidden)
# Eager index q.
q_ref = indexer.q_layernorm(qk[:, : NUM_Q_HEADS * HEAD_DIM].reshape(-1, HEAD_DIM))
q_ref = q_ref.reshape(batch, NUM_Q_HEADS, HEAD_DIM)
q_ref = indexer.apply_rope(positions, q_ref)
# Fused index q.
pool = FakePool(64, 4096, device, dtype)
cache_loc = torch.arange(1, batch + 1, device=device)
q_new, _, stored = indexer.project_qk(
hidden, positions, pool=pool, cache_loc=cache_loc
)
assert stored
compressed_cache = torch.randn(
64, page_size, 1, HEAD_DIM, device=device, dtype=dtype
)
page_table = torch.arange(max_pages, dtype=torch.int32, device=device).repeat(
batch, 1
)
context_lens = torch.full((batch,), 1500, dtype=torch.int32, device=device)
def select(q):
logits = torch_qsa_mqa_decode(
q, compressed_cache, page_table, context_lens, max_model_len
)
row_starts = torch.zeros_like(context_lens)
return qsa_fast_topk(logits, row_starts, context_lens, topk=512)
idx_ref = select(q_ref)
idx_new = select(q_new[:, :NUM_Q_HEADS].contiguous())
for row in range(batch):
ref_set = set(idx_ref[row][idx_ref[row] >= 0].tolist())
new_set = set(idx_new[row][idx_new[row] >= 0].tolist())
assert ref_set == new_set, f"row {row}: selection mismatch"
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernel/speculative/test_verify_commit_triton.py
+81
View File
@@ -0,0 +1,81 @@
import sys
import pytest
import torch
from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
fused_commit_track_indices,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
def _reference(accept_index, accept_lens, seq_lens, draft_token_num, track_interval):
"""Mirrors the eager branch of spec_utils._verify_commit_step_indices."""
bs = accept_lens.shape[0]
offset = torch.arange(
0,
bs * draft_token_num,
step=draft_token_num,
dtype=accept_lens.dtype,
device=accept_lens.device,
)
req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device)
last = accept_index[req_idx, (accept_lens - 1).to(torch.int64)] - offset
if track_interval <= 0:
return last, None
pre = seq_lens
post = seq_lens + accept_lens
mask = pre // track_interval != post // track_interval
point = post // track_interval * track_interval
ith = torch.clamp(point - pre - 1, min=0).to(torch.int64)
cand = accept_index[req_idx, ith] - offset
track = torch.where(mask, cand, torch.full_like(cand, -1))
return last, track
@pytest.mark.parametrize("bs", [1, 3, 48, 257])
@pytest.mark.parametrize("track_interval", [0, 64])
@pytest.mark.parametrize("tree_depth", [4, 3])
def test_verify_commit_steps_matches_eager(bs, track_interval, tree_depth):
"""The fused kernel must match eager on both outputs near tracking boundaries
and when accept_index rows (max_tree_depth) are narrower than draft_token_num."""
if not torch.cuda.is_available():
pytest.skip("needs CUDA")
torch.manual_seed(bs + track_interval + tree_depth)
device = "cuda"
draft_token_num = 4
accept_lens = torch.randint(
1, tree_depth + 1, (bs,), device=device, dtype=torch.int32
)
tree_nodes = torch.argsort(torch.rand(bs, draft_token_num, device=device), dim=1)[
:, :tree_depth
]
accept_index = (
torch.arange(bs, device=device, dtype=torch.int64).unsqueeze(1)
* draft_token_num
+ tree_nodes
).to(torch.int32)
# Cluster seq lens around tracking boundaries to exercise the crossing.
seq_lens = torch.randint(60, 70, (bs,), device=device, dtype=torch.int64)
exp_last, exp_track = _reference(
accept_index, accept_lens, seq_lens, draft_token_num, track_interval
)
got_last, got_track = fused_commit_track_indices(
accept_index,
accept_lens,
seq_lens if track_interval > 0 else None,
draft_token_num,
track_interval,
)
assert torch.equal(got_last, exp_last)
if track_interval > 0:
assert torch.equal(got_track, exp_track)
else:
assert got_track is None
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/unit/configs/test_model_config.py
+23
View File
@@ -7,7 +7,9 @@ from sglang.srt.configs.model_config import (
ModelConfig,
get_hybrid_layer_ids,
is_embedding_gemma,
resolve_spec_hidden_size,
)
from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -69,6 +71,27 @@ class TestDraftModelConfig(CustomTestCase):
self.assertEqual(config.hf_config.num_nextn_predict_layers, 1)
self.assertEqual(config.hf_text_config.num_nextn_predict_layers, 1)
def test_qwen4_exp_spec_hidden_size_keeps_hc_width(self):
"""Qwen4-Exp's MTP draft consumes the hc-flattened target stream,
so spec_hidden_size must stay hidden_size * hc_mult; hy_v4 collapses first."""
hidden_size, hc_mult = 2560, 4
self.assertEqual(Qwen4ExpTextConfig(hc_count=hc_mult).hc_mult, hc_mult)
for arch in ("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLMMTP"):
hf_config = SimpleNamespace(architectures=[arch])
self.assertEqual(
resolve_spec_hidden_size(
hf_config=hf_config, hidden_size=hidden_size, hc_mult=hc_mult
),
(hidden_size * hc_mult, hidden_size * hc_mult),
)
hy_v4 = SimpleNamespace(architectures=["HYV4ForCausalLM"])
self.assertEqual(
resolve_spec_hidden_size(
hf_config=hy_v4, hidden_size=hidden_size, hc_mult=hc_mult
),
(hidden_size, None),
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_mamba_unittest.py
+151
View File
@@ -631,6 +631,157 @@ class TestMamba(unittest.TestCase):
return tree, allocator, req_to_token_pool, make_dummy_req
# Qwen4-Exp's PLE N-gram window is 2 wide (ngram_size=3) and its "no history"
# sentinel is the eos id; pick a recognisable one for the tests.
NGRAM_CONTEXT_LEN = 2
NGRAM_EOS = 248044
def _setup_pool_with_ngram(self, ngram_context_len: int = NGRAM_CONTEXT_LEN):
server_args = ServerArgs(model_path="dummy", page_size=1)
server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE
set_global_server_args_for_scheduler(server_args)
with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"):
shape = Mamba2StateShape.create(
tp_world_size=1,
intermediate_size=4096,
n_groups=16,
num_heads=32,
head_dim=128,
state_size=128,
conv_kernel=4,
)
cache_params = Mamba2CacheParams(shape=shape, layers=[0])
return HybridReqToTokenPool(
size=10,
mamba_size=20,
mamba_spec_state_size=10,
max_context_len=128,
device=get_device(),
enable_memory_saver=False,
cache_params=cache_params,
mamba_layer_ids=[0],
enable_mamba_extra_buffer=False,
speculative_num_draft_tokens=3,
ngram_context_len=ngram_context_len,
ngram_eos_token_id=self.NGRAM_EOS,
)
# Slot-sibling parity: each test below pins one way a mamba slot changes owner.
def test_slot_siblings_registered(self):
"""Enabled PLE side states register on the pool that owns the slots;
disabled ones stay off so the host-offload payload keeps its legacy shape."""
_, _, base_pool, _ = self._setup_tree_and_allocator()
# The default hybrid setup has no PLE config: no siblings ride along.
self.assertEqual(len(base_pool.mamba_pool._slot_siblings), 0)
pool = self._setup_pool_with_ngram()
self.assertEqual(len(pool.mamba_pool._slot_siblings), 1)
def test_ngram_clear_slots_resets_window(self):
"""A recycled slot must not carry its previous owner's N-gram window;
the sibling reset must ride the same deferred ``clear_slots`` call."""
pool = self._setup_pool_with_ngram()
mamba_pool = pool.mamba_pool
ngram = pool.ngram_pool
victim = pool.mamba_allocator.alloc(1)
ngram.context[victim.long()] = 777 # poison, as a real request's history
mamba_pool.clear_slots(victim)
self.assertTrue(
torch.all(ngram.context[victim.long()] == self.NGRAM_EOS),
f"clear_slots left a dirty N-gram row: {ngram.context[victim.long()]}",
)
def test_ngram_copy_from_copies_window(self):
"""copy_from carries the window, so radix cow gets the cached prefix's state."""
pool = self._setup_pool_with_ngram()
mamba_pool = pool.mamba_pool
ngram = pool.ngram_pool
src = pool.mamba_allocator.alloc(1)
dst = pool.mamba_allocator.alloc(1)
window = torch.tensor(
[[55, 66]], dtype=ngram.context.dtype, device=ngram.context.device
)
ngram.context[src.long()] = window
mamba_pool.copy_from(src, dst)
self.assertTrue(
torch.equal(ngram.context[dst.long()], window),
f"copy_from lost the N-gram window: got {ngram.context[dst.long()]}",
)
def test_ngram_cpu_offload_roundtrip(self):
"""The window survives a host offload round-trip along with mamba state."""
pool = self._setup_pool_with_ngram()
mamba_pool = pool.mamba_pool
ngram = pool.ngram_pool
indices = pool.mamba_allocator.alloc(2)
window = torch.tensor(
[[11, 12], [13, 14]],
dtype=ngram.context.dtype,
device=ngram.context.device,
)
ngram.context[indices.long()] = window
saved = mamba_pool.get_cpu_copy(indices)
ngram.context[indices.long()] = self.NGRAM_EOS # simulate slot reuse
mamba_pool.load_cpu_copy(saved, indices)
self.assertTrue(
torch.equal(ngram.context[indices.long()], window),
f"offload round-trip lost the window: got {ngram.context[indices.long()]}",
)
def test_ngram_pool_absent_keeps_legacy_offload_shape(self):
"""Disabled pool stays inert: legacy 2-tuple offload payload, no sibling."""
pool = self._setup_pool_with_ngram(ngram_context_len=0)
self.assertIsNone(pool.ngram_pool.context)
self.assertEqual(len(pool.mamba_pool._slot_siblings), 0)
src = pool.mamba_allocator.alloc(1)
payload = pool.mamba_pool.get_cpu_copy(src)
self.assertEqual(len(payload), 2)
pool.mamba_pool.load_cpu_copy(payload, src)
def test_mamba_track_aligned_lens_math(self):
"""Floor division must swallow the scheduler's `aligned + 1` (_force_track_h),
or the PLE side states snapshot one token past the mamba state."""
from types import SimpleNamespace
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def aligned_for(chunk_size, track_seqlens, prefix_lens):
server_args = ServerArgs(model_path="dummy", page_size=1)
server_args._mamba_cache_chunk_size = chunk_size
set_global_server_args_for_scheduler(server_args)
fake = SimpleNamespace(
mamba_track_mask=torch.tensor([True] * len(track_seqlens)),
mamba_track_seqlens=torch.tensor(track_seqlens, dtype=torch.int64),
extend_prefix_lens=torch.tensor(prefix_lens, dtype=torch.int64),
)
return ForwardBatch.mamba_track_aligned_lens(fake).tolist()
# normal: track_seqlens = prefix + extend_input_len
self.assertEqual(
aligned_for(64, [100 + 64, 100 + 100, 100 + 127], [100, 100, 100]),
[64, 64, 64],
)
# _force_track_h with chunk > 64: track_seqlens = aligned + 1
self.assertEqual(aligned_for(128, [100 + 128 + 1], [100]), [128])
self.assertEqual(aligned_for(128, [100 + 256 + 1], [100]), [256])
# branching point inside the chunk, also handed over as +1
self.assertEqual(aligned_for(64, [100 + 64 + 1], [100]), [64])
# a masked-off row carries -1 and must come out non-positive, so the
# caller's clamp(min=0) routes it harmlessly
self.assertLessEqual(aligned_for(64, [-1], [100])[0], 0)
# restore the chunk size the rest of the suite expects
server_args = ServerArgs(model_path="dummy", page_size=1)
server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE
set_global_server_args_for_scheduler(server_args)
def test_mamba_pool_cpu_offload(self):
"""MambaPool.get_cpu_copy / load_cpu_copy round-trips conv and temporal state."""
_, _, req_to_token_pool, _ = self._setup_tree_and_allocator()
test/registered/unit/server_args/test_server_args.py
+17
View File
@@ -105,6 +105,23 @@ _mock_device.start()
class TestPrepareServerArgs(CustomTestCase):
def test_ple_embedding_offload_rejects_generic_weight_offload(self):
for generic_offload in (
{"cpu_offload_gb": 1},
{"offload_group_size": 1},
):
with (
self.subTest(generic_offload=generic_offload),
self.assertRaisesRegex(
ValueError, "ple-offload-embedding cannot be combined"
),
):
ServerArgs(
model_path="dummy",
ple_offload_embedding=True,
**generic_offload,
).resolve_once()
def test_weight_cache_daemon_allows_static_eplb(self):
args = ServerArgs(
model_path="dummy",
test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py
+8 -2
View File
@@ -252,7 +252,10 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase):
existing_backend = object()
decode_backend = object()
worker.server_args = _fake_server_args()
worker.draft_runner = SimpleNamespace(attn_backend=existing_backend)
worker.draft_runner = SimpleNamespace(
attn_backend=existing_backend,
model_config=SimpleNamespace(hf_config=SimpleNamespace()),
)
worker.topk = 1
worker.speculative_num_steps = 2
worker.seed_dsa_topk_from_draft_extend = False
@@ -274,7 +277,10 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase):
decode_backend = object()
draft_extend_backend = object()
worker.server_args = _fake_server_args()
worker.draft_runner = SimpleNamespace(attn_backend=existing_backend)
worker.draft_runner = SimpleNamespace(
attn_backend=existing_backend,
model_config=SimpleNamespace(hf_config=SimpleNamespace()),
)
worker.topk = 1
worker.speculative_num_steps = 2
worker.seed_dsa_topk_from_draft_extend = True
test/registered/unit/test_model_overrides.py
+37
View File
@@ -88,6 +88,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
"uses_mamba_radix_cache",
"mamba_radix_cache_strategy",
"mamba_full_memory_ratio",
"ple_offload_embedding",
"speculative_moe_runner_backend",
"speculative_moe_a2a_backend",
"disable_shared_experts_fusion",
@@ -605,12 +606,48 @@ class TestGoldenModelOverrides(_IsolatedPublish):
def test_control_arch_keeps_pristine_dtype(self):
sa = self._construct("LlamaForCausalLM", "llama")
self.assertEqual(self._resolved(sa, "dtype"), "auto")
self.assertIsNone(self._resolved(sa, "ple_offload_embedding"))
declared = {f for _s, d in sa._resolved_overrides for f in d}
self.assertNotIn("dtype", declared) # no arch declaration for Llama
# publish still projects the whitelisted leaf with the pristine
# value: readers only ever read flags.
self.assertEqual((self._publish(sa), self._leaf("dtype"))[1], "auto")
def test_qwen4_rejects_pd_and_unified_memory(self):
qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
for kwargs, message in (
({"disaggregation_mode": "prefill"}, "PD disaggregation"),
({"disaggregation_mode": "decode"}, "PD disaggregation"),
({"enable_unified_memory": True}, "enable-unified-memory"),
):
with self.subTest(**kwargs):
with self.assertRaisesRegex(ValueError, message):
self._construct(*qwen4, **kwargs)
def test_qwen4_ple_offload_default(self):
qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
with override_platform(is_cuda=True):
for kwargs, expected in (
({}, True),
({"dtype": "float16"}, False),
({"ple_offload_embedding": False}, False),
({"ple_offload_embedding": False, "cpu_offload_gb": 1}, False),
):
with self.subTest(kwargs=kwargs):
self.assertEqual(
self._resolved(
self._construct(*qwen4, **kwargs),
"ple_offload_embedding",
),
expected,
)
with self.assertRaisesRegex(ValueError, "cannot be combined"):
self._construct(*qwen4, cpu_offload_gb=1)
with override_platform(is_cuda=False, is_hip=True):
self.assertFalse(
self._resolved(self._construct(*qwen4), "ple_offload_embedding")
)
def test_minimax_m2_enables_tf32_matmul(self):
sa = self._construct("MiniMaxM2ForCausalLM", "llama")
self.assertTrue(self._resolved(sa, "enable_tf32_matmul"))