diff --git a/benchmark/kernels/fused_moe_triton/common_utils.py b/benchmark/kernels/fused_moe_triton/common_utils.py index e509834a4..c08a6176a 100644 --- a/benchmark/kernels/fused_moe_triton/common_utils.py +++ b/benchmark/kernels/fused_moe_triton/common_utils.py @@ -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", diff --git a/python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh b/python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh new file mode 100644 index 000000000..01efe40af --- /dev/null +++ b/python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh @@ -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 +#include + +#include +#include +#include +#include +#include + +#include + +#include + +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 +SGL_DEVICE float eager_round(float x) { + return static_cast(DTypeTrait::from(x)); +} + +template <> +SGL_DEVICE float eager_round(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(DTypeTrait::from(x)); +#endif +} + +template <> +SGL_DEVICE float eager_round(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(DTypeTrait::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 +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; + 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(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(row[d]); + const float s = eager_round(row[half + d]); + const float nd = static_cast(smem_row[d]); + const float np = static_cast(smem_row[p]); + o = DTypeTrait::from(eager_round(nd * c) - eager_round(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(row[p]); + const float s = eager_round(row[half + p]); + const float nd = static_cast(smem_row[d]); + const float np = static_cast(smem_row[p]); + o = DTypeTrait::from(eager_round(nd * c) + eager_round(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(row[p]); + const float s = eager_round(row[half + p]); + const int32_t q = (d % 2 == 0) ? d + 1 : d - 1; + const float nd = static_cast(smem_row[d]); + const float nq = static_cast(smem_row[q]); + const float t1 = eager_round(nd * c); + const float t2 = eager_round(nq * s); + o = DTypeTrait::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 +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; + 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(xv[i]); + } + + static_assert(kHeadDim % 8 == 0); + constexpr uint32_t kNormThreads = kHeadDim / 8; + float ss = 0.0f; + if (lane < kNormThreads) { + AlignedVector 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(va[i]); + ss += f * f; + } +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float f = static_cast(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(wv[i]); + smem_row[lane * kPerLane + i] = DTypeTrait::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 +__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; + 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(); + + const int64_t qk_row = static_cast(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(warp); h < params.q_heads_padded; h += 4) { + T* out_row = static_cast(params.q_out) + (static_cast(token) * params.q_heads_padded + h) * kHeadDim; + if (h < params.num_q_heads) { + const T* x_row = static_cast(params.qk) + qk_row + h * kHeadDim; + qsa_gemma_norm_row(x_row, static_cast(params.weight), params.eps, smem_rows[warp]); + qsa_mrope_apply( + smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim); + } else { + vec_t zv; + zv.fill(DTypeTrait::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(params.qk) + qk_row + params.num_q_heads * kHeadDim, + lane); // offset is in vector units + kv.store(static_cast(params.key_state_buffer) + loc * kHeadDim, lane); + } + if (warp == 1 && lane < 3) { + params.rope_position_buffer[loc * 3 + lane] = pos[lane]; + } + + device::PDLTriggerSecondary(); +} + +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 +__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; + 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(params.num_groups)) { + return; + } + __shared__ T smem_rows[4][kHeadDim]; + + device::PDLWaitPrimary(); + + 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(params.key_state_buffer) + static_cast(locs[r]) * kHeadDim, + lane); // offset is in vector units +#pragma unroll + for (int i = 0; i < kPerLane; ++i) { + const float f = static_cast(v[i]); + acc[i] = (r == 0) ? f : acc[i] + f; + } + } + const float inv_ratio = 1.0f / static_cast(params.compress_ratio); +#pragma unroll + for (int i = 0; i < kPerLane; ++i) { + const T m = DTypeTrait::from(acc[i] * inv_ratio); + mf[i] = static_cast(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(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(static_cast(params.weight)[lane * kPerLane + i]); + smem_rows[warp][lane * kPerLane + i] = DTypeTrait::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(loc0) * 3 + a]; + } + + T* out_row = static_cast(params.compressed_k_buffer) + static_cast(params.write_locs[group]) * kHeadDim; + qsa_mrope_apply( + smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim); + + device::PDLTriggerSecondary(); +} + +/** + * \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 +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(); + constexpr int64_t D = kHeadDim; + + TensorMatcher({tokens, (num_q_heads + 1) * D}).with_dtype().with_device(device).verify(qk); + auto heads_padded = SymbolicSize{"heads_padded"}; + TensorMatcher({tokens, heads_padded, D}).with_dtype().with_device(device).verify(q_out); + TensorMatcher({D}).with_dtype().with_device(device).verify(weight); + auto cache_rows = SymbolicSize{"cos_sin_cache_rows"}; + TensorMatcher({cache_rows, rotary_dim}).with_dtype().with_device(device).verify(cos_sin_cache); + TensorMatcher({rotary_dim / 2}).with_dtype().with_device(device).verify(axis_map); + TensorMatcher({num_axes, tokens}).with_dtype().with_device(device).with_strides({-1, 1}).verify(positions); + TensorMatcher({tokens}).with_dtype().with_device(device).verify(cache_loc); + auto slots = SymbolicSize{"state_slots"}; + TensorMatcher({slots, D}).with_dtype().with_device(device).verify(key_state_buffer); + TensorMatcher({slots, 3}).with_dtype().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(cos_sin_cache.data_ptr()), + .axis_map = static_cast(axis_map.data_ptr()), + .positions = static_cast(positions.data_ptr()), + .cache_loc = static_cast(cache_loc.data_ptr()), + .key_state_buffer = key_state_buffer.data_ptr(), + .rope_position_buffer = static_cast(rope_position_buffer.data_ptr()), + .positions_stride = positions.stride(0), + .num_axes = static_cast(num_axes), + .num_q_heads = static_cast(num_q_heads), + .q_heads_padded = static_cast(q_heads_padded), + .rotary_dim = static_cast(rotary_dim), + .eps = eps, + }; + LaunchKernel(static_cast(num_tokens), 128, device.unwrap()) + .enable_pdl(kUsePDL)(qsa_index_q_prep_kernel, params); +} + +/** + * \brief Validate inputs and launch `qsa_index_k_compress_kernel` (one warp per group). + */ +template +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(); + constexpr int64_t D = kHeadDim; + + auto slots = SymbolicSize{"state_slots"}; + TensorMatcher({slots, D}).with_dtype().with_device(device).verify(key_state_buffer); + auto groups = SymbolicSize{"groups"}; + TensorMatcher({groups, compress_ratio}).with_dtype().with_device(device).verify(group_locs); + TensorMatcher({slots, 3}).with_dtype().with_device(device).verify(rope_position_buffer); + auto cache_rows = SymbolicSize{"cos_sin_cache_rows"}; + TensorMatcher({cache_rows, rotary_dim}).with_dtype().with_device(device).verify(cos_sin_cache); + TensorMatcher({rotary_dim / 2}).with_dtype().with_device(device).verify(axis_map); + TensorMatcher({D}).with_dtype().with_device(device).verify(weight); + TensorMatcher({groups}).with_dtype().with_device(device).verify(write_locs); + auto compressed_slots = SymbolicSize{"compressed_slots"}; + TensorMatcher({compressed_slots, D}).with_dtype().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(group_locs.data_ptr()), + .rope_position_buffer = static_cast(rope_position_buffer.data_ptr()), + .cos_sin_cache = static_cast(cos_sin_cache.data_ptr()), + .axis_map = static_cast(axis_map.data_ptr()), + .weight = weight.data_ptr(), + .write_locs = static_cast(write_locs.data_ptr()), + .compressed_k_buffer = compressed_k_buffer.data_ptr(), + .compress_ratio = static_cast(compress_ratio), + .rotary_dim = static_cast(rotary_dim), + .num_groups = static_cast(num_groups), + .eps = eps, + }; + LaunchKernel(static_cast(div_ceil(num_groups, 4)), 128, device.unwrap()) + .enable_pdl(kUsePDL)(qsa_index_k_compress_kernel, params); +} + +} // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/elementwise/fast_topk.cuh b/python/sglang/kernels/jit/csrc/elementwise/fast_topk.cuh new file mode 100644 index 000000000..4b015b2f0 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/elementwise/fast_topk.cuh @@ -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 +#include + +#include + +#include + +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(~bits) : static_cast(bits | 0x8000); + return static_cast(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 +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 +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(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(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 +__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(); + + const auto bid = static_cast(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(score, indice, length); + } else { + radix_select_topk(score, indice, row_start, length); + } + + device::PDLTriggerSecondary(); +} + +} // 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 +struct FastTopKKernel { + static constexpr auto kernel = fast_topk_detail::fast_topk_kernel; + + 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(); + + TensorMatcher({B, L}) // score + .with_strides({S, 1}) + .with_dtype() + .with_device(device) + .verify(score); + TensorMatcher({B}) // row_starts + .with_dtype() + .with_device(device) + .verify(row_starts); + TensorMatcher({B, kTopK}) // indices + .with_dtype() + .with_device(device) + .verify(indices); + TensorMatcher({B}) // lengths + .with_dtype() + .with_device(device) + .verify(lengths); + + const auto params = fast_topk_detail::FastTopKParams{ + .input = static_cast(score.data_ptr()), + .row_starts = static_cast(row_starts.data_ptr()), + .indices = static_cast(indices.data_ptr()), + .lengths = static_cast(lengths.data_ptr()), + .input_stride = S.unwrap(), + }; + + const auto num_rows = static_cast(B.unwrap()); + LaunchKernel(num_rows, fast_topk_detail::kThreadsPerBlock, device.unwrap(), fast_topk_detail::kSmemBytes) + .enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/elementwise/grouped_gemma_rmsnorm.cuh b/python/sglang/kernels/jit/csrc/elementwise/grouped_gemma_rmsnorm.cuh new file mode 100644 index 000000000..abdde55a6 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/elementwise/grouped_gemma_rmsnorm.cuh @@ -0,0 +1,178 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +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 +__global__ __launch_bounds__(kGroupSize / 16) void grouped_gemma_rmsnorm_kernel( + const GroupedGemmaRMSNormParams __grid_constant__ params) { + using namespace device; + using Float2 = packed_t; +#if SGL_ARCH_BLACKWELL_OR_GREATER + // Blackwell: 32B vector, each thread loads/stores once + using Storage = AlignedVector; + constexpr uint32_t kNumLoads = 1; +#else + // Pre-Blackwell: 16B vector, each thread loads/stores twice + using Storage = AlignedVector; + 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::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(); + + const auto input_ptr = pointer::offset(params.input, static_cast(bid) * kGroupSize); + const auto output_ptr = pointer::offset(params.output, static_cast(bid) * kGroupSize); + const auto weight_ptr = pointer::offset(params.weight, static_cast(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(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(input_vec[j][i]); + const auto [wx, wy] = cast(weight_vec[j][i]); + output_vec[i] = cast(fp32x2_t{ix * norm_factor * (1.0f + wx), iy * norm_factor * (1.0f + wy)}); + } + gmem.store(output_ptr, output_vec, j); + } + + PDLTriggerSecondary(); +} + +template +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; + static constexpr auto kBlockSize = static_cast(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(); + + TensorMatcher({M, H}) // input + .with_dtype() + .with_device(device) + .verify(input); + TensorMatcher({H}) // weight + .with_dtype() + .with_device(device) + .verify(weight); + TensorMatcher({M, H}) // output + .with_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(hidden_size / kGroupSize), + .eps = eps, + }; + + const auto num_tokens = static_cast(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 diff --git a/python/sglang/kernels/jit/csrc/elementwise/hc_combine.cuh b/python/sglang/kernels/jit/csrc/elementwise/hc_combine.cuh new file mode 100644 index 000000000..9d9c5fc32 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/elementwise/hc_combine.cuh @@ -0,0 +1,381 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +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 +__global__ __launch_bounds__(256) void hc_combine_kernel(const HcCombineParams __grid_constant__ params) { + using namespace device; + using Float2 = packed_t; + using Storage = AlignedVector; // 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::cta(kNumThreads); + const uint32_t m = blockIdx.x; + + const auto y_ptr = pointer::offset(params.block_output, static_cast(m) * kHiddenSize); + const auto r_ptr = pointer::offset(params.residual, static_cast(m) * kRowSize); + const auto n_ptr = pointer::offset(params.normed_residual, static_cast(m) * kRowSize); + const auto w_ptr = static_cast(params.inject_weight); + const auto out_ptr = pointer::offset(params.output, static_cast(m) * kRowSize); + + PDLWaitPrimary(); + + // 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(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(n_vec[j][i]); + const auto [wx, wy] = cast(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(r_vec[i]); + const auto [yx, yy] = cast(y_vec[i]); + out_vec[i] = cast(fp32x2_t{rx + a * yx, ry + a * yy}); + } + gmem.store(out_ptr, out_vec, j); + } + + PDLTriggerSecondary(); +} + +template +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; + 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(); + + TensorMatcher({M, kHiddenSize}) // block_output + .with_dtype() + .with_device(device) + .verify(block_output); + TensorMatcher({M, kHcCount * kHiddenSize}) // residual, normed_residual, output + .with_dtype() + .with_device(device) + .verify(residual) + .verify(normed_residual) + .verify(output); + TensorMatcher({kHcCount, kHcCount * kHiddenSize}) // inject_weight + .with_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(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 +__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; + using Storage = AlignedVector; + 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(params.normed_residual, static_cast(m) * kRowSize); + const auto w_ptr = static_cast(params.inject_weight); + + PDLWaitPrimary(); + + 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(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(n_vec[j][i]); + const auto [wx, wy] = cast(w_vec[i]); + sum += nx * wx + ny * wy; + } + } + sum = warp::reduce_sum(sum); + if (threadIdx.x == 0) { + params.partials[(static_cast(m) * kSplit + split) * kHcCount + c] = sum; + } + } + + PDLTriggerSecondary(); +} + +/** + * \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 +__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; + using Storage = AlignedVector; + 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(params.block_output, static_cast(m) * kHiddenSize); + const auto r_ptr = pointer::offset(params.residual, static_cast(m) * kRowSize); + const auto out_ptr = pointer::offset(params.output, static_cast(m) * kRowSize); + + PDLWaitPrimary(); + + float total = 0.0f; +#pragma unroll + for (uint32_t s = 0; s < kSplit; ++s) { + total += params.partials[(static_cast(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(r_vec[i]); + const auto [yx, yy] = cast(y_vec[i]); + out_vec[i] = cast(fp32x2_t{rx + a * yx, ry + a * yy}); + } + out_vec.store(out_ptr, vec_idx); + } + + PDLTriggerSecondary(); +} + +template +struct HcCombineSplitKernel { + static_assert(sizeof(DType) == 2, "HcCombine only supports 2-byte dtypes"); + static constexpr auto gate_kernel = hc_combine_gate_kernel; + static constexpr auto apply_kernel = hc_combine_apply_kernel; + + 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(); + + TensorMatcher({M, kHiddenSize}).with_dtype().with_device(device).verify(block_output); + TensorMatcher({M, kHcCount * kHiddenSize}) + .with_dtype() + .with_device(device) + .verify(residual) + .verify(normed_residual) + .verify(output); + TensorMatcher({kHcCount, kHcCount * kHiddenSize}).with_dtype().with_device(device).verify(inject_weight); + auto part_rows = SymbolicSize{"partial_rows"}; + TensorMatcher({part_rows, kSplit, kHcCount}).with_dtype().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(partials.data_ptr()), + }; + + const auto num_tokens = static_cast(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 diff --git a/python/sglang/kernels/kda_kernels/README.md b/python/sglang/kernels/kda_kernels/README.md index b1b892272..dc4d451ad 100644 --- a/python/sglang/kernels/kda_kernels/README.md +++ b/python/sglang/kernels/kda_kernels/README.md @@ -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 diff --git a/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/README.md b/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/README.md new file mode 100644 index 000000000..2659868b8 --- /dev/null +++ b/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/README.md @@ -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. diff --git a/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/__init__.py b/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/__init__.py new file mode 100644 index 000000000..5295972cd --- /dev/null +++ b/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/__init__.py @@ -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"] diff --git a/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/kernel.py b/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/kernel.py new file mode 100644 index 000000000..fa9aa3d4e --- /dev/null +++ b/python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/kernel.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index 6350324fc..42e51effc 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/python/sglang/kernels/ops/attention/__init__.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/qsa_indexer.py b/python/sglang/kernels/ops/attention/qsa_indexer.py new file mode 100644 index 000000000..7e7f71f5f --- /dev/null +++ b/python/sglang/kernels/ops/attention/qsa_indexer.py @@ -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, + ) diff --git a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py index 2c6096531..e1f8fb7d4 100644 --- a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py +++ b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py @@ -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, diff --git a/python/sglang/kernels/ops/elementwise/fast_topk.py b/python/sglang/kernels/ops/elementwise/fast_topk.py new file mode 100644 index 000000000..7e8757fc7 --- /dev/null +++ b/python/sglang/kernels/ops/elementwise/fast_topk.py @@ -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 diff --git a/python/sglang/kernels/ops/elementwise/hc_combine.py b/python/sglang/kernels/ops/elementwise/hc_combine.py new file mode 100644 index 000000000..ca86184c0 --- /dev/null +++ b/python/sglang/kernels/ops/elementwise/hc_combine.py @@ -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) diff --git a/python/sglang/kernels/ops/elementwise/hc_mix.py b/python/sglang/kernels/ops/elementwise/hc_mix.py new file mode 100644 index 000000000..3919b33e0 --- /dev/null +++ b/python/sglang/kernels/ops/elementwise/hc_mix.py @@ -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 diff --git a/python/sglang/kernels/ops/gemm/dense_bf16_gemm_sm100_splitk_epilogue.py b/python/sglang/kernels/ops/gemm/dense_bf16_gemm_sm100_splitk_epilogue.py new file mode 100644 index 000000000..015b94285 --- /dev/null +++ b/python/sglang/kernels/ops/gemm/dense_bf16_gemm_sm100_splitk_epilogue.py @@ -0,0 +1,1293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# Vendored from flashinfer-ai/flashinfer PR #4266 at 629147317d4149a12e53bcef27808bac380c283f. +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.cute.math as cute_math +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = 227 * 1024 + +#: K extent of one CTA tile. +_CTA_K = 128 + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4, 8, 16) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 + +#: A/B pipeline stage bounds. +_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 + + +@dataclasses.dataclass(frozen=True, slots=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], " + f"got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across " + f"split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + base = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(_MIN_AB_STAGES, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", + "run_splitk_dense_silu", + "run_splitk_dense_gate", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: cute.Pointer, + peer_cta_rank_in_cluster: Int32, + *, + loc=None, + ip=None, +) -> Int32: + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: cute.Pointer, + mbar_ptr: cute.Pointer, + peer_cta_rank_in_cluster: Int32, + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank( + smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip + ).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank( + mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip + ).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 + + +def _sigmoid_f32(v): + return cute_math.rcp(cute_math.exp(v * -1.0) + 1.0) + + +#: Epilogue modes; "none" preserves the vendored store path byte-for-byte. +_EPILOGUE_MODES = ("none", "silu", "gate") + +#: Named barrier for the gate epilogue's SMEM staging round-trip. +_GATE_BARRIER_ID = 7 + +#: Alignment of the gate epilogue SMEM tile. +_GATE_TILE_ALIGN_BYTES = 16 + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + epilogue_mode: str = "none", + epilogue_scale: float = 1.0, + epilogue_group: int = 1, + ) -> None: + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + self.epilogue_mode = epilogue_mode + self.epilogue_scale = epilogue_scale + self.epilogue_group = epilogue_group + + if epilogue_mode not in _EPILOGUE_MODES: + raise ValueError(f"unsupported epilogue_mode={epilogue_mode}") + if epilogue_mode == "gate": + if tactic.split_k != 1: + raise ValueError("gate epilogue requires split_k=1") + if has_bias: + raise ValueError("gate epilogue does not support bias") + if epilogue_group < 2 or tactic.mma_m % epilogue_group: + raise ValueError( + f"gate epilogue_group={epilogue_group} must divide " + f"mma_m={tactic.mma_m}" + ) + gate_out_elems = (tactic.mma_m // epilogue_group) * tactic.mma_n + if gate_out_elems % 128: + raise ValueError( + f"gate tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{gate_out_elems} outputs; must be a multiple of 128" + ) + gate_smem = ( + _align_up(_smem_bytes(tactic, tactic.ab_stages), _GATE_TILE_ALIGN_BYTES) + + tactic.mma_m * tactic.mma_n * _FP32_BYTES + ) + if gate_smem > _SMEM_CAPACITY_BYTES: + raise ValueError( + f"gate epilogue needs {gate_smem} B of shared memory; " + f"only {_SMEM_CAPACITY_BYTES} B available" + ) + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = ( + (tactic.split_k - 1) * self.epilog_threads * values_per_thread + ) + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + x: cute.Tensor, + out: cute.Tensor, + stream: _cuda.CUstream, + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel(a, b, c, bias, x, out).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + mX: cute.Tensor, # Gate activation; dead unless epilogue_mode="gate" + mOut: cute.Tensor, # Gate output; dead unless epilogue_mode="gate" + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc( + tiled_mma, self.mma_tiler_mn, acc_stage=1 + ) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if cutlass.const_expr(self.epilogue_mode == "gate"): + gate_tile = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n)), + alignment=_GATE_TILE_ALIGN_BYTES, + ) + else: + gate_tile = mailbox + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + gate_tile, + mX, + mOut, + bidx, + n_idx, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + gate_tile, + mX: cute.Tensor, + mOut: cute.Tensor, + bidx: cutlass.Int32, + n_idx: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr( + self.epilog_threads * values_per_thread + ) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + + epi_tid * values_per_thread + + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + + if cutlass.const_expr(self.epilogue_mode == "silu"): + scaled = rAcc.load() * self.epilogue_scale + rAcc.store(scaled * _sigmoid_f32(scaled)) + + if cutlass.const_expr(self.epilogue_mode == "gate"): + group = self.epilogue_group + rSig = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rSig.store(_sigmoid_f32(rAcc.load())) + sGate_epi = cute.flat_divide(gate_tile, epi_tile) + cute_ext.partition_and_copy(thr_t2r, rSig, sGate_epi[None, None, 0, 0]) + cute.arch.barrier( + barrier_id=_GATE_BARRIER_ID, + number_of_threads=self.epilog_threads, + ) + + rows = cute.size(mOut, mode=[0]) + hs = cute.size(mOut, mode=[1]) + j_per_tile = self.cta_m // group + total_out = j_per_tile * self.cta_n + for it in cutlass.range_constexpr(total_out // self.epilog_threads): + elem = it * self.epilog_threads + epi_tid + j_local = elem % j_per_tile + m_local = elem // j_per_tile + m_global = n_idx * self.cta_n + m_local + j_global = bidx * j_per_tile + j_local + if m_global < rows: + gated = cutlass.Float32(0.0) + for g in cutlass.range_constexpr(group): + sig = gate_tile[j_local * group + g, m_local] + xv = mX[m_global, g * hs + j_global].to(cutlass.Float32) + gated = gated + sig * xv + mOut[m_global, j_global] = (gated * self.epilogue_scale).to( + c_dtype + ) + else: + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + +import torch as _torch + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + c, + c, + stream, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + c_swapped = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c_swapped, + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + c_swapped, + c_swapped, + stream, + ) + + +@cute.experimental.jit +def _bmm_gate( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + x: cute.Tensor, + out: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + x, + out, + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [ + leading_dim + ] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, + epilogue_mode: str = "none", +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic( + _make_layout_tensor(shape, dtype, leading_dim), leading_dim + ) + for shape, leading_dim in zip( + ((batch, n, k), (batch, k, m), (batch, n, m)), + (a_leading, b_leading, c_leading), + strict=True, + ) + ) + if epilogue_mode == "gate": + return ( + *tensors, + _from_dlpack_dynamic(_torch.empty((n, m), dtype=dtype, device="cuda"), 1), + _from_dlpack_dynamic(_torch.empty((n, m), dtype=dtype, device="cuda"), 1), + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple( + _detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap) + ) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip( + (a_swap, b_swap, c_swap), leading_dims, strict=True + ) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided( + size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0) + ), + 1, + 2, + ), + leading_dims, + ) + + +# Tactic hashes all compile-time tile, split, and stage fields. +_SPLITK_COMPILE_CACHE: dict = {} + + +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], + epilogue_mode: str = "none", + epilogue_scale: float = 1.0, + epilogue_group: int = 1, +): + key = ( + dtype, + tactic, + use_pdl, + has_bias, + epilogue_mode, + epilogue_scale, + epilogue_group, + *leading_dims, + ) + cached = _SPLITK_COMPILE_CACHE.get(key) + if cached is not None: + return cached + + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError( + f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}" + ) + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + epilogue_mode=epilogue_mode, + epilogue_scale=epilogue_scale, + epilogue_group=epilogue_group, + ) + compile_tensors = _make_compile_repr_tensors( + dtype, has_bias, *leading_dims, epilogue_mode=epilogue_mode + ) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if epilogue_mode == "gate": + compiled = cute_ext.compile(_bmm_gate, kernel, *compile_tensors, stream) + elif has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + _SPLITK_COMPILE_CACHE[key] = compiled + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any( + tensor.dtype != a.dtype for tensor in tensors + ): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + + def _is_dense_2d(tensor: _torch.Tensor) -> bool: + rows, cols = tensor.shape + return (tensor.stride(1) == 1 and tensor.stride(0) >= cols) or ( + tensor.stride(0) == 1 and tensor.stride(1) >= rows + ) + + if any(not _is_dense_2d(tensor) for tensor in (a, b, out)): + raise ValueError( + "a, b, and out must be row-major or column-major matrices " + "(padded leading strides allowed)" + ) + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError( + f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}" + ) + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and ( + bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous() + ): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +def run_splitk_dense_silu( + a, + b, + out, + pdl: bool, + tactic: SplitKTactic, + scale: float, +): + """Run ``bf16(silu(scale * (A[M,K] @ B[K,N])))``.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=False, + leading_dims=cute_tensors[4], + epilogue_mode="silu", + epilogue_scale=scale, + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*cute_tensors[:3], stream) + return out + + +def run_splitk_dense_gate( + a, + b, + x, + out, + pdl: bool, + tactic: SplitKTactic, + scale: float, + group: int, +): + """Fused ``logits = A @ B`` (B pre-permuted so ``n = j*group + g``) with + ``out[m, j] = bf16(scale * sum_g sigmoid(logits[m, j*group+g]) * + fp32(x[m, g*hs+j]))``; the logits matrix is never materialized and ``x`` + doubles as its shape anchor.""" + m, n, k = _validate_runtime_tensors(a, b, None, x) + validate_tactic(tactic, m, n, k) + if n % group: + raise ValueError(f"N={n} must be divisible by group={group}") + if n % tactic.mma_m: + raise ValueError( + f"gate epilogue requires N={n} divisible by mma_m={tactic.mma_m}" + ) + hs = n // group + if ( + out.ndim != 2 + or out.shape != (m, hs) + or out.dtype != a.dtype + or out.device != a.device + or out.stride(1) != 1 + or out.stride(0) < hs + or out.data_ptr() % 32 + or x.stride(1) != 1 + ): + raise ValueError( + f"gate out must be a 32-byte-aligned row-major {(m, hs)} tensor " + f"matching a, and x must be row-major" + ) + cute_tensors = _to_cute_swap(a, b, x, None) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=False, + leading_dims=cute_tensors[4], + epilogue_mode="gate", + epilogue_scale=scale, + epilogue_group=group, + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled( + *cute_tensors[:3], + _from_dlpack_dynamic(x, 1), + _from_dlpack_dynamic(out, 1), + stream, + ) + return out diff --git a/python/sglang/kernels/ops/layernorm/grouped_gemma_rmsnorm.py b/python/sglang/kernels/ops/layernorm/grouped_gemma_rmsnorm.py new file mode 100644 index 000000000..c05adb707 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/grouped_gemma_rmsnorm.py @@ -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) diff --git a/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py b/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py index 9ff741a81..29be4ccc5 100644 --- a/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py +++ b/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py @@ -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, ) diff --git a/python/sglang/kernels/ops/qwen4_ple.py b/python/sglang/kernels/ops/qwen4_ple.py new file mode 100644 index 000000000..f91d9f2e9 --- /dev/null +++ b/python/sglang/kernels/ops/qwen4_ple.py @@ -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 diff --git a/python/sglang/srt/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index c735df017..954bc39d4 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -78,6 +78,7 @@ ATTENTION_BACKEND_CHOICES = [ "flex_attention", "dsa", "nsa", # Deprecated alias for "dsa" + "qsa", "dsv4", "compressed", # Deprecated alias for "dsv4" # NVIDIA specific diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index f80d877c9..289a08ef1 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -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 diff --git a/python/sglang/srt/arg_groups/memory_hook.py b/python/sglang/srt/arg_groups/memory_hook.py index 437a4d8bb..267738f41 100644 --- a/python/sglang/srt/arg_groups/memory_hook.py +++ b/python/sglang/srt/arg_groups/memory_hook.py @@ -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 diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index e1678f3bf..14f8e910b 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -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: diff --git a/python/sglang/srt/arg_groups/model_overrides/__init__.py b/python/sglang/srt/arg_groups/model_overrides/__init__.py index 88dcb05e6..2d6240e9f 100644 --- a/python/sglang/srt/arg_groups/model_overrides/__init__.py +++ b/python/sglang/srt/arg_groups/model_overrides/__init__.py @@ -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 diff --git a/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py b/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py index 2eeb402bd..68840f8e2 100644 --- a/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py +++ b/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py @@ -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) diff --git a/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py b/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py new file mode 100644 index 000000000..ae8e2af72 --- /dev/null +++ b/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py @@ -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 diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index bbfaf34ea..486e56786 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -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", diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 3cff4c0f9..28440f244 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/python/sglang/srt/arg_groups/pipeline.py @@ -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) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 9866a4d69..b2f39e61d 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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( diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 093f29125..7278a8d08 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -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", diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index ff3b6d340..b589e250e 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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", diff --git a/python/sglang/srt/configs/qwen4_exp.py b/python/sglang/srt/configs/qwen4_exp.py new file mode 100644 index 000000000..a161face8 --- /dev/null +++ b/python/sglang/srt/configs/qwen4_exp.py @@ -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) diff --git a/python/sglang/srt/configs/update_config.py b/python/sglang/srt/configs/update_config.py index fa2c803b2..918ee9ab9 100644 --- a/python/sglang/srt/configs/update_config.py +++ b/python/sglang/srt/configs/update_config.py @@ -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, diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 84b58da33..9c4acca59 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -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) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 576772665..6f64ca795 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 20a7e2003..4466b55c3 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -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 diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 673e1cb3f..5610a251d 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/qsa/__init__.py b/python/sglang/srt/layers/attention/qsa/__init__.py new file mode 100644 index 000000000..d7a9e3b12 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/__init__.py @@ -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) diff --git a/python/sglang/srt/layers/attention/qsa/config.py b/python/sglang/srt/layers/attention/qsa/config.py new file mode 100644 index 000000000..4e0c0c954 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/config.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qsa/dsa_indexer.py b/python/sglang/srt/layers/attention/qsa/dsa_indexer.py new file mode 100644 index 000000000..88f84cb86 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/dsa_indexer.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qsa/glue.py b/python/sglang/srt/layers/attention/qsa/glue.py new file mode 100644 index 000000000..326895770 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/glue.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qsa/graph_metadata.py b/python/sglang/srt/layers/attention/qsa/graph_metadata.py new file mode 100644 index 000000000..dc91f87e0 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/graph_metadata.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/qsa/kernel.py b/python/sglang/srt/layers/attention/qsa/kernel.py new file mode 100644 index 000000000..56070c496 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/kernel.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qsa/metadata.py b/python/sglang/srt/layers/attention/qsa/metadata.py new file mode 100644 index 000000000..a9e3270f4 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/metadata.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qsa/mqa.py b/python/sglang/srt/layers/attention/qsa/mqa.py new file mode 100644 index 000000000..6446ed145 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/mqa.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py new file mode 100644 index 000000000..4de4c6a8d --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qsa/sparse_attn.py b/python/sglang/srt/layers/attention/qsa/sparse_attn.py new file mode 100644 index 000000000..e2d714815 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/sparse_attn.py @@ -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", +] diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py new file mode 100644 index 000000000..4f8eaefef --- /dev/null +++ b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py @@ -0,0 +1,1764 @@ +"""Sparse-attention backend for Qwen4-Exp models with an indexer. + +The backend is installed as the full-attention side of Qwen4-Exp's hybrid backend; +linear-attention layers continue to use GDN. +""" + +from __future__ import annotations + +import logging +import math +from copy import copy +from functools import lru_cache +from typing import Dict, Optional, Tuple + +import msgspec +import torch +import torch.nn.functional as F + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.qsa.config import ( + QSA_VARIANT_COMPRESSED, + is_qwen_qsa, + parse_qsa_profile, +) +from sglang.srt.layers.attention.qsa.kernel import qsa_sparse_attention +from sglang.srt.layers.attention.qsa.metadata import ( + QSAIndexerMetadata, + build_group_ring_slots, + build_pending_ring_slots, + build_rope_position_matrix, + compressed_decode_view, +) +from sglang.srt.layers.attention.qsa.sparse_attn import ( + qwen_sparse_fa2_cu_seqlens_triton, + qwen_sparse_kv_extraction_compact_triton, + qwen_sparse_valid_counts_triton, + sparse_gqa_fwd_interface_triton, + sparse_gqa_fwd_interface_triton_ck, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode + +logger = logging.getLogger(__name__) + + +_TRTLLM_SPARSE_PAGE_SIZE = 64 + + +@lru_cache(maxsize=1) +def _resolve_trtllm_sparse_decode(): + """FlashInfer paged decode for the post-gather sparse attention; + the FA4 varlen fallback runs a prefill-shaped kernel at decode row counts.""" + from sglang.srt.utils import is_sm100_supported, is_sm120 + + # This path is numerically validated on SM100 and SM120. Do not widen it + # to every SM12x device: it silently corrupts long-context decode on + # SM121/GB10. + if not (is_sm100_supported() or is_sm120()): + return None + try: + from flashinfer.decode import trtllm_batch_decode_with_kv_cache + except ImportError: + return None + return trtllm_batch_decode_with_kv_cache + + +@lru_cache(maxsize=1) +def _resolve_flash_attn_varlen_func(): + from sglang.srt.utils import is_sm121 + + if is_sm121(): + from sglang.kernels.ops.attention import ( + qwen38_qsa_sm121_varlen, + ) + + return qwen38_qsa_sm121_varlen + try: + from flash_attn import flash_attn_varlen_func + + return flash_attn_varlen_func + except ImportError: + pass + try: + from flash_attn.cute.interface import flash_attn_varlen_func as cute_varlen_func + + def flash_attn_varlen_func(*args, **kwargs): + output = cute_varlen_func(*args, **kwargs) + # The cute interface returns (out, lse); lse is None here. + return output[0] if isinstance(output, tuple) else output + + return flash_attn_varlen_func + except ImportError as exc: + raise ImportError( + "QSA decode requires flash_attn (FA2) or flash-attn-4 " + "(FA4 cute) for its packed varlen fallback." + ) from exc + + +class QwenSparseAttnMetadata(msgspec.Struct, frozen=True): + """Per-forward metadata consumed by core sparse attention.""" + + sequence_lengths: torch.Tensor + token_to_batch_idx: torch.Tensor + token_slot_table: torch.Tensor + indexer_metadata: QSAIndexerMetadata + row_req_pool_indices: Optional[torch.Tensor] = None + is_cuda_graph: bool = False + fa2_valid_counts: Optional[torch.Tensor] = None + fa2_cu_seqlens_k: Optional[torch.Tensor] = None + fa2_cu_seqlens_q: Optional[torch.Tensor] = None + + +class QSAMTPSharedSparseIndices: + """Draft-extend top-k reused by one MTP iteration's decode steps, + causally valid because a draft moves at most speculative_num_steps positions. + Row ``num_requests`` is the trash row for padded/degenerate requests.""" + + def __init__( + self, *, layer_ids, num_requests, token_topk, tail_width, device + ) -> None: + self.layer_slots = {int(l): i for i, l in enumerate(sorted(layer_ids))} + self.tail_width = tail_width + self.trash_row = num_requests + # Logical index 0 keeps never-captured rows (graph warmup dummies) + # attending exactly the first token instead of an empty/invalid set. + self.indices = torch.zeros( + (len(self.layer_slots), num_requests + 1, token_topk + tail_width), + dtype=torch.int32, + device=device, + ) + self.captured_len = torch.ones( + (len(self.layer_slots), num_requests + 1), + dtype=torch.int32, + device=device, + ) + self._tail_offsets = torch.arange(tail_width, device=device) + + def capture( + self, + topk_indices: torch.Tensor, + req_pool_indices: torch.Tensor, + captured_lens: torch.Tensor, + layer_id: int, + ) -> None: + slot = self.layer_slots[int(layer_id)] + rows = req_pool_indices.to(torch.long) + self.indices[slot, :, : topk_indices.shape[1]].index_copy_( + 0, rows, topk_indices.to(self.indices.dtype) + ) + self.captured_len[slot].index_copy_(0, rows, captured_lens.to(torch.int32)) + + def lookup( + self, + req_pool_indices: torch.Tensor, + current_positions: torch.Tensor, + layer_id: int, + ) -> torch.Tensor: + """Frozen selection plus a tail of positions drafted since the capture; + tail slots past current_position hold -1, which downstream drops.""" + slot = self.layer_slots[int(layer_id)] + rows = req_pool_indices.to(torch.long) + out = self.indices[slot, rows] + base = self.captured_len[slot, rows].to(torch.int64) + tail = base.unsqueeze(1) + self._tail_offsets.unsqueeze(0) + valid = tail <= current_positions.to(torch.int64).unsqueeze(1) + out[:, out.shape[1] - self.tail_width :] = torch.where(valid, tail, -1).to( + out.dtype + ) + return out + + +class QwenSparseAttnBackend(AttentionBackend): + """QSA backend using trtllm-gen decode with a packed FA2/FA4 fallback.""" + + # Every seq_lens_cpu read here has a device fallback (one readback); + # the graphed decode path never reads it, so opting out is safe. + needs_cpu_seq_lens: bool = False + + def __init__(self, runner=None) -> None: + self.runner = runner + self.token_to_kv_pool = getattr(runner, "token_to_kv_pool", None) + self.device = getattr(runner, "device", None) + model_config = getattr(runner, "model_config", None) + config = getattr(model_config, "hf_text_config", None) + if config is None: + config = getattr(model_config, "hf_config", None) + # Compressed (Qwen4-Exp) and tokenwise (Qwen3Next-DSA) QSA share this backend. + self.qsa_profile = parse_qsa_profile(config) + self.max_context_len = int(getattr(model_config, "context_len", 0)) + self.compress_ratio = ( + self.qsa_profile.compress_ratio + if self.qsa_profile is not None + else int(getattr(config, "indexer_compress_ratio", 4)) + ) + req_pool = getattr(runner, "req_to_token_pool", None) + self.req_to_token = getattr(req_pool, "req_to_token", None) + self.req_to_token_pool = req_pool + self.forward_metadata: Optional[QwenSparseAttnMetadata] = None + self._cuda_graph_metadata: Dict[ + Tuple[ForwardMode, int], QwenSparseAttnMetadata + ] = {} + self._cuda_graph_max_tokens = 0 + self._fa2_scratch: Dict[ + Tuple[int, int, torch.dtype, torch.device], + Tuple[torch.Tensor, torch.Tensor], + ] = {} + self._graph_seq_lens = None + self._graph_token_to_batch = None + self._graph_cu_seqlens_q = None + self._graph_fa2_valid_counts = None + self._graph_fa2_cu_seqlens_k = None + self._graph_write_locs = None + self._graph_compressed_page_table = None + self._graph_compressed_lengths = None + self._graph_prefix_lengths = None + self._graph_dummy_token_slot_table = None + self._graph_dummy_out_cache_loc = None + self._graph_row_req_pool_indices = None + self._trtllm_sparse_tables = {} + self._mtp_shared_sparse_indices = None + self._trtllm_workspace = None + self._graph_extend_lens = None + self._graph_extend_lens_pin = None + + @staticmethod + def _is_speculative_paged_mode(forward_mode) -> bool: + if forward_mode is None: + return False + return forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2() + + def _require_chain_speculation(self, forward_mode, spec_info) -> None: + if forward_mode is None or not forward_mode.is_target_verify(): + return + if int(getattr(spec_info, "topk", 1) or 1) != 1: + raise NotImplementedError( + "Qwen QSA target verification supports only speculative_eagle_topk=1" + ) + draft_tokens = int(getattr(spec_info, "draft_token_num", 0) or 0) + if draft_tokens > self.compress_ratio: + # The pending-group ring keys state by position % ratio; a verify + # window wider than the ratio would collide within one forward. + raise NotImplementedError( + "Qwen QSA requires speculative_num_draft_tokens <= the QSA " + f"compress ratio ({self.compress_ratio}): the pending " + f"index-key ring holds one group; got {draft_tokens}" + ) + + @staticmethod + def _speculative_max_row_length(forward_batch, sequence_lengths) -> int: + """Sync-free host-side over-bound for the token-slot gather width: + request length plus the draft window covers every speculative row.""" + seq_lens_cpu = forward_batch.seq_lens_cpu + if seq_lens_cpu is None or seq_lens_cpu.numel() == 0: + logger.warning_once( + "QSA speculative metadata without CPU request lengths: the " + "token-slot bound reads them back once per forward" + ) + return max(1, int(sequence_lengths.max())) + spec_info = forward_batch.spec_info + draft_window = int(spec_info.draft_token_num) if spec_info is not None else 0 + return max(1, int(seq_lens_cpu.max()) + draft_window) + + @staticmethod + def _speculative_row_to_request(forward_batch, num_rows: int) -> torch.Tensor: + batch_size = int(forward_batch.req_pool_indices.numel()) + if batch_size == 0: + if num_rows == 0: + return torch.zeros( + 0, + dtype=torch.long, + device=forward_batch.req_pool_indices.device, + ) + raise ValueError( + "QSA speculative query rows cannot be mapped to an empty batch: " + f"rows={num_rows}" + ) + extend_seq_lens = getattr(forward_batch, "extend_seq_lens", None) + if extend_seq_lens is not None: + # DP token padding rows belong to no request; alias them to request row 0, + # as the CUDA-graph layout does, so gathers stay in bounds. + repeats = extend_seq_lens[:batch_size].to(dtype=torch.long) + real_rows = int(repeats.sum().item()) + if real_rows > num_rows: + raise ValueError( + "QSA speculative query rows are fewer than the extend " + f"mapping: rows={num_rows}, mapped={real_rows}" + ) + row_to_request = torch.repeat_interleave( + torch.arange( + batch_size, + dtype=torch.long, + device=forward_batch.req_pool_indices.device, + ), + repeats, + ) + padding = num_rows - real_rows + if padding: + row_to_request = torch.cat( + [ + row_to_request, + row_to_request.new_zeros(padding), + ] + ) + return row_to_request + if num_rows % batch_size != 0: + raise ValueError( + "QSA speculative query rows cannot be mapped to requests: " + f"rows={num_rows}, batch={batch_size}" + ) + return torch.arange( + batch_size, + dtype=torch.long, + device=forward_batch.req_pool_indices.device, + ).repeat_interleave(num_rows // batch_size) + + @staticmethod + def _as_cpu_int_tensor(values, size: int) -> torch.Tensor: + if isinstance(values, torch.Tensor): + return values[:size].detach().cpu().to(torch.int32) + return torch.tensor(values[:size], dtype=torch.int32) + + @classmethod + def _graph_speculative_layout( + cls, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens_cpu, + forward_mode, + spec_info, + num_padding: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + base_lengths = cls._as_cpu_int_tensor(seq_lens_cpu, bs) + if forward_mode.is_target_verify(): + extend_len = int(spec_info.draft_token_num) + extend_lengths = torch.full((bs,), extend_len, dtype=torch.int32) + else: + extend_lengths_cpu = getattr(spec_info, "extend_seq_lens_cpu", None) + if extend_lengths_cpu is not None: + extend_lengths = cls._as_cpu_int_tensor(extend_lengths_cpu, bs) + if extend_lengths.numel() < bs: + extend_lengths = torch.cat( + [ + extend_lengths, + torch.zeros(bs - extend_lengths.numel(), dtype=torch.int32), + ] + ) + else: + extend_lengths = torch.full( + (bs,), num_tokens // max(bs, 1), dtype=torch.int32 + ) + num_padding = max(0, min(int(num_padding), bs)) + if num_padding: + extend_lengths = extend_lengths.clone() + extend_lengths[bs - num_padding :] = 0 + effective_lengths = ( + base_lengths + extend_lengths + if forward_mode.is_target_verify() + else base_lengths + ) + + row_lengths = [] + row_prefix_lengths = [] + for seq_len, extend_len in zip( + effective_lengths.tolist(), extend_lengths.tolist() + ): + prefix_len = max(int(seq_len) - int(extend_len), 0) + row_lengths.append( + torch.arange( + prefix_len + 1, + prefix_len + int(extend_len) + 1, + dtype=torch.int32, + ).clamp_max(int(seq_len)) + ) + row_prefix_lengths.append( + torch.full((int(extend_len),), prefix_len, dtype=torch.int32) + ) + row_lengths = ( + torch.cat(row_lengths) if row_lengths else torch.empty(0, dtype=torch.int32) + ) + row_prefix_lengths = ( + torch.cat(row_prefix_lengths) + if row_prefix_lengths + else torch.empty(0, dtype=torch.int32) + ) + actual_rows = row_lengths.numel() + if actual_rows > num_tokens: + raise ValueError( + "QSA CUDA graph speculative layout has inconsistent token count: " + f"capacity={num_tokens}, actual={actual_rows}" + ) + repeats = extend_lengths.to(device=req_pool_indices.device, dtype=torch.long) + row_req_pool_indices = torch.repeat_interleave(req_pool_indices[:bs], repeats) + # A draft-extend replay may hold fewer accepted rows than the captured shape; + # the runner packs real rows first, so give the tail inert metadata. + padding = num_tokens - actual_rows + if padding: + row_lengths = torch.cat( + [row_lengths, torch.ones(padding, dtype=torch.int32)] + ) + row_prefix_lengths = torch.cat( + [row_prefix_lengths, torch.zeros(padding, dtype=torch.int32)] + ) + dummy_req = ( + req_pool_indices[0] + if req_pool_indices.numel() + else torch.zeros((), dtype=torch.int32, device=req_pool_indices.device) + ) + row_req_pool_indices = torch.cat( + [ + row_req_pool_indices, + dummy_req.to(row_req_pool_indices.dtype).expand(padding), + ] + ) + return row_lengths, row_req_pool_indices, row_prefix_lengths + + def _empty_metadata(self, forward_batch) -> QwenSparseAttnMetadata: + """Well-formed zero-row metadata for IDLE/zero-request forwards.""" + + device = forward_batch.seq_lens.device + sequence_lengths = torch.zeros(0, dtype=torch.int32, device=device) + token_to_batch_idx = torch.zeros(0, dtype=torch.int32, device=device) + token_slot_table = torch.zeros((0, 1), dtype=torch.int32, device=device) + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + if out_cache_loc is None: + out_cache_loc = torch.zeros(0, dtype=torch.int64, device=device) + indexer_metadata = QSAIndexerMetadata( + sequence_lengths=sequence_lengths, + token_to_batch_idx=token_to_batch_idx, + token_slot_table=token_slot_table, + out_cache_loc=out_cache_loc, + token_to_kv_pool=self.token_to_kv_pool, + compress_ratio=self.token_to_kv_pool.qsa_compress_ratio, + block_topk=self.token_to_kv_pool.qsa_block_topk, + ) + return QwenSparseAttnMetadata( + sequence_lengths=sequence_lengths, + token_to_batch_idx=token_to_batch_idx, + token_slot_table=token_slot_table, + indexer_metadata=indexer_metadata, + row_req_pool_indices=torch.zeros(0, dtype=torch.int32, device=device), + ) + + @staticmethod + def _qsa_write_plan( + *, + token_slot_table, + start_blocks, + end_blocks, + capacity, + compress_ratio, + row_token_starts=None, + prefix_lens=None, + ): + """Compact per-row block ranges into ``capacity`` write entries, + a shape-derived bound (no sync); padding writes the inert reserved slot 0.""" + device = token_slot_table.device + # The table width is a host-side bound; + # assert on device so a short table fails loudly without a sync. + torch._assert_async( + (end_blocks * compress_ratio <= token_slot_table.shape[1]).all() + ) + counts = (end_blocks - start_blocks).clamp_min(0) + ends = torch.cumsum(counts, 0) + starts = ends - counts + entries = torch.arange(capacity, dtype=torch.long, device=device) + # Which row owns each compacted entry, and its ordinal within the row. + rows = torch.searchsorted(ends, entries, right=True) + valid = entries < ends[-1] if counts.numel() else entries < 0 + rows = torch.where(valid, rows.clamp_max(max(counts.numel() - 1, 0)), 0) + blocks = torch.where(valid, start_blocks[rows] + entries - starts[rows], 0) + write_locs = torch.where( + valid, + token_slot_table[rows, blocks * compress_ratio].long() // compress_ratio, + torch.zeros_like(blocks), + ).to(torch.int32) + group_end_positions = blocks * compress_ratio + (compress_ratio - 1) + member_rows = None + if row_token_starts is not None: + # member_rows index this forward's packed token rows; + # the chunk is group-aligned, so a group starts at block * ratio - prefix. + member_rows = torch.where( + valid, + row_token_starts[rows] + blocks * compress_ratio - prefix_lens[rows], + torch.zeros_like(blocks), + ) + return write_locs, group_end_positions, rows, member_rows + + def _qsa_build_write_plan( + self, + *, + forward_batch, + speculative_paged: bool, + token_slot_table, + sequence_lengths, + ): + ratio = self.token_to_kv_pool.qsa_compress_ratio + lengths = sequence_lengths.long() + end_blocks = lengths // ratio + if forward_batch.forward_mode.is_decode() or speculative_paged: + # A paged row compresses exactly the block its length completes; + # its members come from the per-request pending ring. + start_blocks = end_blocks - (lengths % ratio == 0).long() + return self._qsa_write_plan( + token_slot_table=token_slot_table, + start_blocks=start_blocks, + end_blocks=end_blocks, + capacity=int(lengths.numel()), + compress_ratio=ratio, + ) + extend_lens = forward_batch.extend_seq_lens + if extend_lens is None: + raise ValueError("QSA extend write plan requires extend_seq_lens") + extend_lens = extend_lens.long()[: lengths.numel()] + prefix_lens = (lengths - extend_lens).clamp_min(0) + # Prefix sharing is page-granular and the page is a ratio + # multiple, so a matched prefix always covers whole groups. A + # misaligned prefix would leave a shared group half-written. + torch._assert_async((prefix_lens % ratio == 0).all()) + # Each row spans at most ceil(extend_len / ratio) blocks, so the + # token count and row count bound the plan without a sync. + capacity = int(forward_batch.input_ids.numel()) // ratio + int(lengths.numel()) + row_token_starts = torch.cumsum(extend_lens, 0) - extend_lens + return self._qsa_write_plan( + token_slot_table=token_slot_table, + start_blocks=prefix_lens // ratio, + end_blocks=end_blocks, + capacity=capacity, + compress_ratio=ratio, + row_token_starts=row_token_starts, + prefix_lens=prefix_lens, + ) + + def _metadata_from_forward_batch(self, forward_batch) -> QwenSparseAttnMetadata: + self._require_chain_speculation( + forward_batch.forward_mode, + getattr(forward_batch, "spec_info", None), + ) + if self.device is None: + self.device = forward_batch.seq_lens.device + if not self.max_context_len: + self.max_context_len = self.req_to_token.shape[1] + if forward_batch.forward_mode.is_idle() or forward_batch.seq_lens.numel() == 0: + # DP idle forwards reach metadata init even though layers skip attention; + # so do the zero-row DECODE steps the MTP wrapper makes of them. + return self._empty_metadata(forward_batch) + original_mode = getattr(forward_batch, "_original_forward_mode", None) + if original_mode is not None and self._is_speculative_paged_mode(original_mode): + # DP MAX_LEN pseudo-extend sets extend_seq_lens == 1 per request, + # losing the draft fan-out this mapping needs. + raise ValueError( + "QSA cannot build metadata for DP MAX_LEN pseudo-extend of " + f"speculative mode {original_mode}: token rows would be " + "mis-mapped to requests" + ) + speculative_paged = self._is_speculative_paged_mode(forward_batch.forward_mode) + if speculative_paged: + logical_positions = forward_batch.positions + if logical_positions.ndim == 2: + logical_positions = logical_positions[0] + logical_positions = logical_positions.flatten() + sequence_lengths = (logical_positions + 1).to(torch.int32) + row_to_request = self._speculative_row_to_request( + forward_batch, logical_positions.numel() + ) + row_req_pool_indices = forward_batch.req_pool_indices.index_select( + 0, row_to_request + ) + max_length = self._speculative_max_row_length( + forward_batch, sequence_lengths + ) + token_slot_table = self.req_to_token[ + row_req_pool_indices.long(), :max_length + ].to(torch.int32) + token_to_batch_idx = torch.arange( + sequence_lengths.numel(), + device=sequence_lengths.device, + dtype=torch.int32, + ) + else: + sequence_lengths = forward_batch.seq_lens.to(torch.int32) + batch_size = sequence_lengths.numel() + if forward_batch.seq_lens_cpu is not None: + max_length = int(forward_batch.seq_lens_cpu[:batch_size].max()) + else: + max_length = int(sequence_lengths.max()) + row_req_pool_indices = forward_batch.req_pool_indices[:batch_size] + token_slot_table = self.req_to_token[ + row_req_pool_indices.long(), :max_length + ].to(torch.int32) + positions = forward_batch.positions + num_position_tokens = ( + positions.shape[-1] if positions.ndim == 2 else positions.numel() + ) + if forward_batch.forward_mode.is_decode(): + if num_position_tokens != batch_size: + raise ValueError( + "QSA decode requires exactly one query row per request: " + f"rows={num_position_tokens}, batch={batch_size}" + ) + token_to_batch_idx = torch.arange( + batch_size, + device=sequence_lengths.device, + dtype=torch.int32, + ) + else: + extend_seq_lens = forward_batch.extend_seq_lens + if extend_seq_lens is None: + raise ValueError("QSA extend metadata requires extend_seq_lens") + token_to_batch_idx = torch.repeat_interleave( + torch.arange( + batch_size, + device=sequence_lengths.device, + dtype=torch.int32, + ), + extend_seq_lens.to(torch.long), + ) + # More mapped rows than physical would index past them; + # fewer is fine (DP MAX_LEN padding is trimmed downstream). + if token_to_batch_idx.numel() > num_position_tokens: + raise ValueError( + "QSA extend request mapping exceeds token rows: " + f"mapping={token_to_batch_idx.numel()}, " + f"positions={num_position_tokens}" + ) + write_locs = None + group_positions = None + group_sequence_ids = None + group_member_rows = None + decode_page_table = None + decode_lengths = None + decode_logical_positions = None + pending_ring_slots = None + compress_group_ring_locs = None + extend_rope_matrix = None + if ( + self.qsa_profile is None + or self.qsa_profile.variant == QSA_VARIANT_COMPRESSED + ): + write_locs, group_positions, group_sequence_ids, group_member_rows = ( + self._qsa_build_write_plan( + forward_batch=forward_batch, + speculative_paged=speculative_paged, + token_slot_table=token_slot_table, + sequence_lengths=sequence_lengths, + ) + ) + decode_like = speculative_paged or forward_batch.forward_mode.is_decode() + if decode_like: + decode_logical_positions = ( + logical_positions.to(torch.int32) + if speculative_paged + else sequence_lengths - 1 + ) + ring_logical_positions = decode_logical_positions + else: + extend_positions = forward_batch.positions + if extend_positions.ndim == 2: + extend_positions = extend_positions[0] + ring_logical_positions = extend_positions.flatten()[ + : token_to_batch_idx.numel() + ] + if not self.should_reuse_mtp_sparse_indices(forward_batch): + if decode_like: + pool = self.token_to_kv_pool + decode_page_table, decode_lengths = compressed_decode_view( + compressed_page_size=pool.qsa_compressed_page_size, + compress_ratio=pool.qsa_compress_ratio, + sequence_lengths=sequence_lengths, + token_slot_table=token_slot_table, + ) + pending_ring_slots = build_pending_ring_slots( + token_to_batch_idx=token_to_batch_idx, + req_pool_indices=row_req_pool_indices, + sequence_lengths=sequence_lengths, + logical_positions=ring_logical_positions, + compress_ratio=self.compress_ratio, + is_extend=group_member_rows is not None, + ) + if write_locs.numel(): + if group_member_rows is not None: + rope_source = ( + forward_batch.mrope_positions + if forward_batch.mrope_positions is not None + else forward_batch.positions + ) + extend_rope_matrix = build_rope_position_matrix( + rope_source, token_to_batch_idx.numel() + ) + else: + compress_group_ring_locs = build_group_ring_slots( + req_pool_indices=row_req_pool_indices, + group_end_positions=group_positions.long(), + sequence_ids=group_sequence_ids.long(), + compress_ratio=self.compress_ratio, + ) + indexer_metadata = QSAIndexerMetadata( + sequence_lengths=sequence_lengths, + token_to_batch_idx=token_to_batch_idx, + token_slot_table=token_slot_table, + out_cache_loc=forward_batch.out_cache_loc, + token_to_kv_pool=self.token_to_kv_pool, + compress_ratio=self.token_to_kv_pool.qsa_compress_ratio, + block_topk=self.token_to_kv_pool.qsa_block_topk, + req_pool_indices=row_req_pool_indices, + write_locs=write_locs, + compress_group_positions=group_positions, + compress_sequence_ids=group_sequence_ids, + compress_member_rows=group_member_rows, + decode_page_table=decode_page_table, + decode_lengths=decode_lengths, + decode_logical_positions=decode_logical_positions, + pending_ring_slots=pending_ring_slots, + compress_group_ring_locs=compress_group_ring_locs, + extend_rope_matrix=extend_rope_matrix, + ) + return QwenSparseAttnMetadata( + sequence_lengths=sequence_lengths, + token_to_batch_idx=token_to_batch_idx, + token_slot_table=token_slot_table, + indexer_metadata=indexer_metadata, + row_req_pool_indices=row_req_pool_indices, + ) + + def init_forward_metadata(self, forward_batch): + if forward_batch.forward_mode.is_idle(): + self.forward_metadata = None + return + self.forward_metadata = self._metadata_from_forward_batch(forward_batch) + + def init_forward_metadata_out_graph( + self, + forward_batch, + in_capture: bool = False, + ): + if in_capture: + num_tokens = ( + forward_batch.input_ids.shape[0] + if self._is_speculative_paged_mode(forward_batch.forward_mode) + else forward_batch.batch_size + ) + self._capture_cuda_graph_metadata( + bs=forward_batch.batch_size, + num_tokens=num_tokens, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + forward_mode=forward_batch.forward_mode, + spec_info=forward_batch.spec_info, + ) + else: + num_padding = getattr(forward_batch, "num_padding", None) + self._replay_cuda_graph_metadata( + bs=forward_batch.batch_size, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + forward_mode=forward_batch.forward_mode, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + num_padding=num_padding if num_padding is not None else 0, + ) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: + if self.device is None: + raise RuntimeError( + "QSA backend requires a ModelRunner to initialize CUDA graph state" + ) + self._cuda_graph_max_tokens = max_num_tokens + max_blocks = math.ceil(self.max_context_len / self.compress_ratio) + max_pages = max( + 1, + math.ceil(max_blocks / self.token_to_kv_pool.qsa_compressed_page_size), + ) + self._graph_seq_lens = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_token_to_batch = torch.arange( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_cu_seqlens_q = torch.arange( + max_num_tokens + 1, dtype=torch.int32, device=self.device + ) + self._graph_fa2_valid_counts = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_fa2_cu_seqlens_k = torch.zeros( + max_num_tokens + 1, dtype=torch.int32, device=self.device + ) + self._graph_write_locs = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_compressed_page_table = torch.zeros( + (max_num_tokens, max_pages), dtype=torch.int32, device=self.device + ) + self._graph_compressed_lengths = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_prefix_lengths = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_dummy_token_slot_table = torch.zeros( + (max_num_tokens, 1), dtype=torch.int32, device=self.device + ) + self._graph_dummy_out_cache_loc = torch.zeros( + max_num_tokens, dtype=torch.int64, device=self.device + ) + self._graph_row_req_pool_indices = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_logical_positions = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) + self._graph_state_slots = torch.zeros( + max_num_tokens, dtype=torch.int64, device=self.device + ) + self._graph_ring_group_locs = torch.zeros( + (max_num_tokens, self.compress_ratio), + dtype=torch.int32, + device=self.device, + ) + # Two pinned halves: when a replay refills the pinned buffer, + # the previous async HtoD copy from it may still be queued. + self._graph_extend_lens = torch.zeros( + max_bs, dtype=torch.int32, device=self.device + ) + self._graph_extend_lens_pin = [ + torch.zeros(max_bs, dtype=torch.int32, pin_memory=True) for _ in range(2) + ] + self._extend_lens_pin_idx = 0 + + def _require_compressed_cuda_graph_support(self) -> None: + if ( + self.qsa_profile is not None + and self.qsa_profile.variant != QSA_VARIANT_COMPRESSED + ): + raise NotImplementedError( + "QSA tokenwise CUDA-graph execution requires graph-stable " + "indexer metadata, which is not available in this tree yet; " + "run tokenwise QSA with --disable-cuda-graph" + ) + + def _capture_cuda_graph_metadata( + self, + *, + bs, + num_tokens, + req_pool_indices, + seq_lens, + forward_mode, + spec_info, + ) -> None: + self._require_compressed_cuda_graph_support() + self._require_chain_speculation(forward_mode, spec_info) + if self.token_to_kv_pool is None: + self.token_to_kv_pool = getattr(self.runner, "token_to_kv_pool", None) + if self.req_to_token is None: + req_pool = getattr(self.runner, "req_to_token_pool", None) + self.req_to_token = getattr(req_pool, "req_to_token", None) + self.req_to_token_pool = req_pool + if self._graph_seq_lens is None or self.token_to_kv_pool is None: + raise RuntimeError("QSA CUDA graph state is not initialized") + pool = self.token_to_kv_pool + speculative_paged = self._is_speculative_paged_mode(forward_mode) + metadata_rows = num_tokens if speculative_paged else bs + if speculative_paged: + row_lengths, row_req_pool_indices, row_prefix_lengths = ( + self._graph_speculative_layout( + bs, + num_tokens, + req_pool_indices, + seq_lens[:bs].detach().cpu(), + forward_mode, + spec_info, + num_padding=bs, + ) + ) + self._graph_prefix_lengths[:metadata_rows].copy_( + row_prefix_lengths.to(device=self.device) + ) + self._graph_seq_lens[:metadata_rows].copy_( + row_lengths.to(device=self.device) + ) + self._graph_row_req_pool_indices[:metadata_rows].copy_( + row_req_pool_indices.to(torch.int32) + ) + else: + self._graph_seq_lens[:bs].copy_(seq_lens[:bs].to(torch.int32)) + self._graph_row_req_pool_indices[:bs].copy_( + req_pool_indices[:bs].to(torch.int32) + ) + self._graph_prefix_lengths[:bs].copy_( + (seq_lens[:bs] - 1).clamp_min(0).to(torch.int32) + ) + indexer_metadata = QSAIndexerMetadata( + sequence_lengths=self._graph_seq_lens[:metadata_rows], + token_to_batch_idx=self._graph_token_to_batch[:metadata_rows], + token_slot_table=self._graph_dummy_token_slot_table[:metadata_rows], + out_cache_loc=self._graph_dummy_out_cache_loc[:metadata_rows], + token_to_kv_pool=pool, + compress_ratio=pool.qsa_compress_ratio, + block_topk=pool.qsa_block_topk, + req_pool_indices=self._graph_row_req_pool_indices[:metadata_rows], + is_cuda_graph=True, + graph_write_locs=self._graph_write_locs[:metadata_rows], + graph_compressed_page_table=self._graph_compressed_page_table[ + :metadata_rows + ], + graph_compressed_lengths=self._graph_compressed_lengths[:metadata_rows], + graph_prefix_lengths=self._graph_prefix_lengths[:metadata_rows], + decode_logical_positions=self._graph_logical_positions[:metadata_rows], + pending_ring_slots=self._graph_state_slots[:metadata_rows], + graph_ring_group_locs=self._graph_ring_group_locs[:metadata_rows], + ) + metadata = QwenSparseAttnMetadata( + sequence_lengths=self._graph_seq_lens[:metadata_rows], + token_to_batch_idx=self._graph_token_to_batch[:metadata_rows], + token_slot_table=self._graph_dummy_token_slot_table[:metadata_rows], + indexer_metadata=indexer_metadata, + row_req_pool_indices=self._graph_row_req_pool_indices[:metadata_rows], + is_cuda_graph=True, + fa2_valid_counts=self._graph_fa2_valid_counts[:metadata_rows], + fa2_cu_seqlens_k=self._graph_fa2_cu_seqlens_k[: metadata_rows + 1], + fa2_cu_seqlens_q=self._graph_cu_seqlens_q[: metadata_rows + 1], + ) + self._cuda_graph_metadata[(forward_mode, bs)] = metadata + self.forward_metadata = metadata + + def _replay_cuda_graph_metadata( + self, + *, + bs, + req_pool_indices, + seq_lens, + forward_mode, + spec_info, + seq_lens_cpu, + num_padding: int = 0, + ) -> None: + self._require_chain_speculation(forward_mode, spec_info) + metadata = self._cuda_graph_metadata[(forward_mode, bs)] + # Compressed addressing is arithmetic over req_to_token; pure reads below. + if self._can_replay_with_gpu_kernels(metadata, seq_lens): + self._replay_cuda_graph_metadata_gpu( + metadata, + bs=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + forward_mode=forward_mode, + spec_info=spec_info, + seq_lens_cpu=seq_lens_cpu, + num_padding=num_padding, + ) + self.forward_metadata = metadata + return + if seq_lens_cpu is None: + # Host-fallback staging only (non-CUDA pools): an explicit + # slow-path readback, never the serving path. + seq_lens_cpu = seq_lens[:bs].cpu() + if self._is_speculative_paged_mode(forward_mode): + row_lengths, row_req_pool_indices, row_prefix_lengths = ( + self._graph_speculative_layout( + bs, + metadata.sequence_lengths.numel(), + req_pool_indices, + seq_lens_cpu, + forward_mode, + spec_info, + num_padding=num_padding or 0, + ) + ) + metadata.sequence_lengths.copy_(row_lengths.to(device=self.device)) + metadata.row_req_pool_indices.copy_(row_req_pool_indices.to(torch.int32)) + metadata.indexer_metadata.graph_prefix_lengths.copy_( + row_prefix_lengths.to(device=self.device) + ) + else: + metadata.sequence_lengths.copy_(seq_lens[:bs].to(torch.int32)) + metadata.row_req_pool_indices.copy_(req_pool_indices[:bs].to(torch.int32)) + metadata.indexer_metadata.graph_prefix_lengths.copy_( + (seq_lens[:bs] - 1).clamp_min(0).to(torch.int32) + ) + self._update_qsa_cuda_graph_metadata( + metadata.indexer_metadata, metadata.row_req_pool_indices + ) + self.forward_metadata = metadata + + def _can_replay_with_gpu_kernels(self, metadata, seq_lens) -> bool: + if self.req_to_token is None: + return False + from sglang.srt.layers.attention.qsa.graph_metadata import ( + supports_graph_metadata_kernels, + ) + + return seq_lens.is_cuda and supports_graph_metadata_kernels( + metadata.indexer_metadata.token_to_kv_pool, seq_lens.device + ) + + def _stage_extend_lens(self, spec_info, bs: int, num_tokens: int): + + extend_lens = getattr(spec_info, "extend_seq_lens_tensor", None) + if extend_lens is not None and extend_lens.numel() >= bs: + return extend_lens[:bs] + pin = self._graph_extend_lens_pin[self._extend_lens_pin_idx] + self._extend_lens_pin_idx = 1 - self._extend_lens_pin_idx + values = getattr(spec_info, "extend_seq_lens_cpu", None) + if values is None: + # Mirrors the legacy fallback: an even per-request share of the + # captured token capacity. + pin[:bs].fill_(num_tokens // max(bs, 1)) + elif isinstance(values, torch.Tensor): + pin[:bs].copy_(values[:bs].to(torch.int32)) + else: + pin[:bs].copy_(torch.tensor(values[:bs], dtype=torch.int32)) + staged = self._graph_extend_lens + staged[:bs].copy_(pin[:bs], non_blocking=True) + return staged[:bs] + + def _replay_cuda_graph_metadata_gpu( + self, + metadata, + *, + bs, + req_pool_indices, + seq_lens, + forward_mode, + spec_info, + seq_lens_cpu, + num_padding: int = 0, + ) -> None: + """Sync-free replay refresh: nothing here may read back to the host, + since launch_graph_metadata rebuilds every per-row graph buffer on device.""" + + from sglang.srt.layers.attention.qsa.graph_metadata import ( + launch_graph_metadata, + ) + + indexer = metadata.indexer_metadata + pool = indexer.token_to_kv_pool + speculative = self._is_speculative_paged_mode(forward_mode) + num_rows = metadata.sequence_lengths.numel() if speculative else bs + num_padding = max(0, min(int(num_padding or 0), bs)) + if bs <= 0 or num_rows <= 0: + return + + if speculative: + if forward_mode.is_target_verify(): + mode = 1 + extend_lens = None + extend_len = int(spec_info.draft_token_num) + else: + mode = 2 + extend_len = 0 + extend_lens = self._stage_extend_lens(spec_info, bs, num_rows) + else: + mode = 0 + extend_lens = None + extend_len = 0 + launch_graph_metadata( + mode=mode, + bs=bs, + num_rows=num_rows, + seq_lens=seq_lens, + req_pool_indices=req_pool_indices, + extend_lens=extend_lens, + extend_len=extend_len, + num_padding=num_padding, + metadata=metadata, + req_to_token=self.req_to_token, + pool=pool, + ) + + def _update_qsa_cuda_graph_metadata( + self, metadata: QSAIndexerMetadata, req_pool_indices: torch.Tensor + ) -> None: + """Replay refresh for pools/devices without the graph_metadata kernels.""" + if self.req_to_token is None: + raise RuntimeError("QSA req_to_token table is not initialized") + pool = metadata.token_to_kv_pool + lengths = metadata.sequence_lengths.long() + req_indices = req_pool_indices.long() + ratio = metadata.compress_ratio + full_page = pool.qsa_compressed_page_size * ratio + + current_positions = (lengths - 1).clamp_min(0) + compressed_lengths = torch.div(lengths, ratio, rounding_mode="floor") + metadata.graph_compressed_lengths.copy_(compressed_lengths.to(torch.int32)) + + # Boundary rows write their group's slot (last raw slot // ratio); + # every other row keeps the inert reserved slot 0. + boundary = (lengths % ratio == 0) & (lengths > 0) + last_locs = self.req_to_token[req_indices, current_positions].long() + write_locs = torch.where( + boundary, + last_locs // ratio, + torch.zeros_like(lengths), + ) + metadata.graph_write_locs.copy_(write_locs.to(torch.int32)) + + metadata.decode_logical_positions.copy_(current_positions.to(torch.int32)) + metadata.pending_ring_slots.copy_( + build_pending_ring_slots( + token_to_batch_idx=metadata.token_to_batch_idx, + req_pool_indices=req_pool_indices, + sequence_lengths=metadata.sequence_lengths, + logical_positions=current_positions, + compress_ratio=ratio, + is_extend=False, + ) + ) + metadata.graph_ring_group_locs.copy_( + build_group_ring_slots( + req_pool_indices=req_pool_indices, + group_end_positions=current_positions, + sequence_ids=metadata.token_to_batch_idx.long(), + compress_ratio=ratio, + ).to(torch.int32) + ) + + # Page-table entries are full-KV page ids from the token-slot rows. + page_table = metadata.graph_compressed_page_table + max_pages = page_table.shape[1] + row_width_pages = self.req_to_token.shape[1] // full_page + num_pages = min(max_pages, row_width_pages) + table = ( + self.req_to_token[req_indices, : num_pages * full_page : full_page].long() + // full_page + ).clamp_min(0) + page_table[:, :num_pages].copy_(table.to(torch.int32)) + + def get_indexer_metadata(self, layer_id: int, forward_batch): + if self.forward_metadata is None: + self.init_forward_metadata(forward_batch) + assert self.forward_metadata is not None + return self.forward_metadata.indexer_metadata + + def set_mtp_shared_sparse_indices(self, state) -> None: + self._mtp_shared_sparse_indices = state + + def should_reuse_mtp_sparse_indices(self, forward_batch) -> bool: + """Draft decode steps reuse the draft-extend selection.""" + return ( + self._mtp_shared_sparse_indices is not None + and forward_batch.forward_mode.is_decode() + ) + + def should_capture_mtp_sparse_indices(self, forward_batch) -> bool: + """Capture on both draft-extend flavors (a DP MAX_LEN rewrite is neither): + DRAFT_EXTEND_V2, and the draft runner's plain post-prefill EXTEND.""" + if self._mtp_shared_sparse_indices is None: + return False + if forward_batch._original_forward_mode is not None: + return False + mode = forward_batch.forward_mode + return mode.is_draft_extend_v2() or mode.is_extend_without_speculative() + + def capture_mtp_sparse_indices( + self, topk_indices: torch.Tensor, forward_batch, layer_id: int, metadata=None + ) -> None: + """Store each request's final accepted row as the iteration seed.""" + if topk_indices.shape[0] == 0: + return + if metadata is None: + metadata = self.get_indexer_metadata(layer_id, forward_batch) + if metadata.is_cuda_graph or forward_batch.forward_mode.is_draft_extend_v2(): + self._capture_mtp_sparse_indices_from_extend_lens( + topk_indices, forward_batch, metadata, layer_id + ) + return + if metadata.req_pool_indices is None or metadata.req_pool_indices.numel() == 0: + return + state = self._mtp_shared_sparse_indices + row_to_req = metadata.get_token_to_batch_idx().long() + row_req_pool_indices = metadata.req_pool_indices[ + row_to_req[: topk_indices.shape[0]] + ] + is_last = torch.ones_like(row_req_pool_indices, dtype=torch.bool) + if row_req_pool_indices.numel() > 1: + is_last[:-1] = row_req_pool_indices[:-1] != row_req_pool_indices[1:] + anchor_rows = is_last.nonzero().flatten() + req_rows = row_req_pool_indices[anchor_rows] + captured_lens = metadata.get_seqlens_expanded()[anchor_rows] + state.capture(topk_indices[anchor_rows], req_rows, captured_lens, layer_id) + + @staticmethod + def _capture_extend_seq_lens(forward_batch) -> torch.Tensor: + """Same resolution order as _stage_extend_lens, + so the capture and in-graph layout read the same replay-refreshed buffer.""" + spec_info = getattr(forward_batch, "spec_info", None) + extend_seq_lens = getattr(spec_info, "extend_seq_lens_tensor", None) + if extend_seq_lens is None: + extend_seq_lens = getattr(forward_batch, "extend_seq_lens", None) + if extend_seq_lens is None: + raise RuntimeError( + "QSA draft-extend capture requires GPU extend lengths " + "(spec_info.extend_seq_lens_tensor or " + "forward_batch.extend_seq_lens)" + ) + return extend_seq_lens + + def _capture_mtp_sparse_indices_from_extend_lens( + self, topk_indices, forward_batch, metadata, layer_id: int + ) -> None: + """DRAFT_EXTEND_V2 packs each request as [front rows][draft-window rows], + so its last accepted row is block_end - (window - front - num_accept); + rows with no real request (zero extend, padding slot 0) go to the trash row.""" + + state = self._mtp_shared_sparse_indices + bs = int(forward_batch.batch_size) + if bs == 0 or forward_batch.req_pool_indices.numel() < bs: + return + request_extend_lens = self._capture_extend_seq_lens(forward_batch)[:bs].long() + block_ends = request_extend_lens.cumsum(0) - 1 + spec_info = forward_batch.spec_info + num_accept_tokens = getattr(spec_info, "num_accept_tokens", None) + if num_accept_tokens is not None and num_accept_tokens.numel() >= bs: + # Only EagleDraftExtendInput carries accept counts; + # prefill-side EXTEND has no draft window, so its rows end at the seed. + front_tokens = int(spec_info.num_front_tokens) + anchor_rows = block_ends - ( + request_extend_lens - front_tokens - num_accept_tokens[:bs].long() + ) + else: + anchor_rows = block_ends + anchor_rows = anchor_rows.clamp(min=0, max=topk_indices.shape[0] - 1) + req_pool_rows = forward_batch.req_pool_indices[:bs].long() + req_rows = torch.where( + (request_extend_lens > 0) & (req_pool_rows > 0), + req_pool_rows, + state.trash_row, + ) + captured_lens = metadata.get_seqlens_expanded()[anchor_rows] + state.capture( + topk_indices.index_select(0, anchor_rows), + req_rows, + captured_lens, + layer_id, + ) + + def lookup_mtp_sparse_indices(self, forward_batch, layer_id: int) -> torch.Tensor: + metadata = self.get_indexer_metadata(layer_id, forward_batch) + logical_positions = metadata.decode_logical_positions + if logical_positions is None: + logical_positions = metadata.get_seqlens_expanded() - 1 + return self._mtp_shared_sparse_indices.lookup( + forward_batch.req_pool_indices, + logical_positions, + layer_id, + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + def _resolve_metadata(self, forward_batch) -> QwenSparseAttnMetadata: + if self.forward_metadata is None: + self.init_forward_metadata(forward_batch) + assert self.forward_metadata is not None + return self.forward_metadata + + @staticmethod + def _logical_to_physical( + logical_indices: torch.Tensor, metadata: QwenSparseAttnMetadata + ) -> torch.Tensor: + sequence_ids = metadata.token_to_batch_idx.long() + if sequence_ids.numel() != logical_indices.shape[0]: + raise ValueError("QSA top-k rows do not match query rows") + row_lengths = metadata.sequence_lengths.to(torch.int32).index_select( + 0, sequence_ids + ) + valid = (logical_indices >= 0) & (logical_indices < row_lengths.unsqueeze(1)) + safe = logical_indices.clamp( + min=0, max=metadata.token_slot_table.shape[1] - 1 + ).long() + slots = metadata.token_slot_table[sequence_ids[:, None], safe] + return torch.where(valid, slots, torch.full_like(slots, -1)).to(torch.int32) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer, + forward_batch, + save_kv_cache: bool = True, + topk_indices: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + if topk_indices is None: + raise ValueError("QSA sparse attention requires topk_indices") + if save_kv_cache: + self.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) + num_output_rows = q.shape[0] + num_valid_rows = topk_indices.shape[0] + if num_valid_rows > num_output_rows: + raise ValueError( + "QSA top-k rows exceed query rows: " + f"topk={num_valid_rows}, query={num_output_rows}" + ) + # DP attention may pad q beyond the indexer's query rows. + # The kernels see only valid rows; the output is zero-padded to q's row count. + q = q[:num_valid_rows] + if self._is_speculative_paged_mode(forward_batch.forward_mode): + output = self._forward_paged_attention( + q, layer, forward_batch, topk_indices + ) + return self._pad_extend_output(output, num_output_rows) + if not q.is_cuda: + metadata = self._resolve_metadata(forward_batch) + slots = self._logical_to_physical(topk_indices, metadata) + pool = self.token_to_kv_pool + output = qsa_sparse_attention( + q, + pool.get_key_buffer(layer.layer_id), + pool.get_value_buffer(layer.layer_id), + slots, + layer.scaling, + ) + return self._pad_extend_output(output, num_output_rows) + + topk_indices = topk_indices.to(torch.int32).contiguous() + extend_lens = [int(x) for x in forward_batch.extend_seq_lens_cpu] + sequence_lens = [int(x) for x in forward_batch.seq_lens_cpu] + prefix_lens = [ + sequence_lens[i] - extend_lens[i] for i in range(len(extend_lens)) + ] + cu_seqlens_q = F.pad( + forward_batch.extend_seq_lens.to(q.device, dtype=torch.int32).cumsum(0), + (1, 0), + ).contiguous() + if not any(prefix_lens): + output = sparse_gqa_fwd_interface_triton( + q.contiguous(), + k[:num_valid_rows].contiguous(), + v[:num_valid_rows].contiguous(), + max(sequence_lens, default=1), + topk_indices, + cu_seqlens_q, + layer.scaling, + ) + return self._pad_extend_output(output, num_output_rows) + + # The validated chunk-prefill kernel consumes tightly packed full-context + # K/V. Current-chunk K/V has already been committed to the cache above. + pool = self.token_to_kv_pool + k_buffer = pool.get_key_buffer(layer.layer_id) + v_buffer = pool.get_value_buffer(layer.layer_id) + req_to_token = self.req_to_token_pool.req_to_token + req_indices = forward_batch.req_pool_indices.tolist() + k_parts = [ + k_buffer.index_select( + 0, req_to_token[req_indices[i], : sequence_lens[i]].long() + ) + for i in range(len(sequence_lens)) + ] + v_parts = [ + v_buffer.index_select( + 0, req_to_token[req_indices[i], : sequence_lens[i]].long() + ) + for i in range(len(sequence_lens)) + ] + sequence_lens_tensor = torch.tensor( + sequence_lens, dtype=torch.int32, device=q.device + ) + cu_seqlens_k = F.pad(sequence_lens_tensor.cumsum(0), (1, 0)).contiguous() + output = sparse_gqa_fwd_interface_triton_ck( + q.contiguous(), + torch.cat(k_parts), + torch.cat(v_parts), + topk_indices, + cu_seqlens_q, + cu_seqlens_k, + sequence_lens_tensor, + layer.scaling, + ) + return self._pad_extend_output(output, num_output_rows) + + @staticmethod + def _pad_extend_output(output: torch.Tensor, num_rows: int) -> torch.Tensor: + output = output.reshape(output.shape[0], -1) + if output.shape[0] == num_rows: + return output + if output.shape[0] > num_rows: + raise ValueError( + "QSA attention output rows exceed padded query rows: " + f"output={output.shape[0]}, query={num_rows}" + ) + padded = output.new_zeros((num_rows, output.shape[1])) + padded[: output.shape[0]].copy_(output) + return padded + + def _get_fa2_scratch( + self, + capacity: int, + num_kv_heads: int, + head_dim: int, + dtype: torch.dtype, + device: torch.device, + ) -> Tuple[torch.Tensor, torch.Tensor]: + key = (num_kv_heads, head_dim, dtype, device) + buffers = self._fa2_scratch.get(key) + if buffers is None or buffers[0].shape[0] < capacity: + shape = (capacity, num_kv_heads, head_dim) + buffers = ( + torch.empty(shape, dtype=dtype, device=device), + torch.empty(shape, dtype=dtype, device=device), + ) + self._fa2_scratch[key] = buffers + return buffers[0][:capacity], buffers[1][:capacity] + + def _get_trtllm_sparse_tables(self, batch, pages_per_row, page, device): + key = (batch, pages_per_row, device) + cached = self._trtllm_sparse_tables.get(key) + if cached is None: + stride = pages_per_row * page + cu = torch.arange(batch + 1, dtype=torch.int32, device=device) * stride + block_tables = ( + torch.arange(batch, dtype=torch.int32, device=device)[:, None] + * pages_per_row + + torch.arange(pages_per_row, dtype=torch.int32, device=device)[None, :] + ).contiguous() + cached = (cu, block_tables) + self._trtllm_sparse_tables[key] = cached + return cached + + def _forward_trtllm_sparse( + self, + q: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, + layer, + forward_batch, + metadata, + topk_indices: torch.Tensor, + trtllm_decode, + ) -> torch.Tensor: + """Pack selected KV at page-aligned row strides for FlashInfer's paged decode, + driven by a static arange block table and the per-row valid counts.""" + batch, topk = topk_indices.shape + page = _TRTLLM_SPARSE_PAGE_SIZE + pages_per_row = (topk + page - 1) // page + stride = pages_per_row * page + device = q.device + sequence_lens = metadata.sequence_lengths + if metadata.is_cuda_graph: + valid_counts = metadata.fa2_valid_counts + if valid_counts is None: + raise RuntimeError("QSA CUDA graph metadata is incomplete") + else: + valid_counts = torch.empty(batch, dtype=torch.int32, device=device) + qwen_sparse_valid_counts_triton( + sequence_lens, topk_indices, valid_counts, batch, topk + ) + cu_strided, block_tables = self._get_trtllm_sparse_tables( + batch, pages_per_row, page, device + ) + capacity_rows = self._cuda_graph_max_tokens if metadata.is_cuda_graph else batch + packed_k, packed_v = self._get_fa2_scratch( + max(capacity_rows, batch) * stride, + k_buffer.shape[1], + k_buffer.shape[2], + k_buffer.dtype, + k_buffer.device, + ) + qwen_sparse_kv_extraction_compact_triton( + k_buffer, + v_buffer, + self.req_to_token_pool.req_to_token, + ( + metadata.row_req_pool_indices + if metadata.row_req_pool_indices is not None + else forward_batch.req_pool_indices + ), + topk_indices, + sequence_lens, + cu_strided, + packed_k, + packed_v, + batch, + topk, + ) + num_kv_heads = k_buffer.shape[1] + head_dim = k_buffer.shape[2] + kc = ( + packed_k[: batch * stride] + .view(-1, page, num_kv_heads, head_dim) + .permute(0, 2, 1, 3) + ) + vc = ( + packed_v[: batch * stride] + .view(-1, page, num_kv_heads, head_dim) + .permute(0, 2, 1, 3) + ) + if self._trtllm_workspace is None: + self._trtllm_workspace = torch.zeros( + 128 * 1024 * 1024, dtype=torch.uint8, device=device + ) + output = trtllm_decode( + query=q.contiguous(), + kv_cache=(kc, vc), + workspace_buffer=self._trtllm_workspace, + block_tables=block_tables, + seq_lens=valid_counts, + max_seq_len=stride, + bmm1_scale=layer.scaling, + bmm2_scale=1.0, + ) + return output.reshape(q.shape[0], -1) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer, + forward_batch, + save_kv_cache: bool = True, + topk_indices: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + if topk_indices is None: + raise ValueError("QSA sparse attention requires topk_indices") + if save_kv_cache: + self.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) + return self._forward_paged_attention(q, layer, forward_batch, topk_indices) + + def _forward_paged_attention( + self, + q: torch.Tensor, + layer, + forward_batch, + topk_indices: torch.Tensor, + ) -> torch.Tensor: + pool = self.token_to_kv_pool + k_buffer = pool.get_key_buffer(layer.layer_id) + v_buffer = pool.get_value_buffer(layer.layer_id) + if not q.is_cuda: + metadata = self._resolve_metadata(forward_batch) + slots = self._logical_to_physical(topk_indices, metadata) + output = qsa_sparse_attention(q, k_buffer, v_buffer, slots, layer.scaling) + return output.reshape(q.shape[0], -1) + + metadata = self._resolve_metadata(forward_batch) + topk_indices = topk_indices.to(torch.int32).contiguous() + trtllm_decode = _resolve_trtllm_sparse_decode() + if trtllm_decode is not None: + return self._forward_trtllm_sparse( + q, + k_buffer, + v_buffer, + layer, + forward_batch, + metadata, + topk_indices, + trtllm_decode, + ) + + flash_attn_varlen_func = _resolve_flash_attn_varlen_func() + batch, topk = topk_indices.shape + sequence_lens = metadata.sequence_lengths + if metadata.is_cuda_graph: + valid_counts = metadata.fa2_valid_counts + cu_seqlens_k = metadata.fa2_cu_seqlens_k + cu_seqlens_q = metadata.fa2_cu_seqlens_q + if valid_counts is None or cu_seqlens_k is None or cu_seqlens_q is None: + raise RuntimeError("QSA CUDA graph FA2 metadata is incomplete") + else: + valid_counts = torch.empty(batch, dtype=torch.int32, device=q.device) + cu_seqlens_k = torch.empty(batch + 1, dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.arange(batch + 1, dtype=torch.int32, device=q.device) + qwen_sparse_fa2_cu_seqlens_triton( + sequence_lens, + topk_indices, + valid_counts, + cu_seqlens_k, + batch, + topk, + ) + scratch_capacity = ( + self._cuda_graph_max_tokens * topk + if metadata.is_cuda_graph + else batch * topk + ) + packed_k, packed_v = self._get_fa2_scratch( + scratch_capacity, + k_buffer.shape[1], + k_buffer.shape[2], + k_buffer.dtype, + k_buffer.device, + ) + qwen_sparse_kv_extraction_compact_triton( + k_buffer, + v_buffer, + self.req_to_token_pool.req_to_token, + ( + metadata.row_req_pool_indices + if metadata.row_req_pool_indices is not None + else forward_batch.req_pool_indices + ), + topk_indices, + sequence_lens, + cu_seqlens_k, + packed_k, + packed_v, + batch, + topk, + ) + output = flash_attn_varlen_func( + q=q, + k=packed_k, + v=packed_v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=1, + max_seqlen_k=topk, + softmax_scale=layer.scaling, + causal=True, + ) + return output.reshape(q.shape[0], -1) + + +class QwenSparseMultiStepDraftBackend: + """Per-step QSA metadata for consecutive Qwen4-Exp MTP draft decoding.""" + + needs_cpu_seq_lens: bool = False + + def __init__(self, model_runner, topk: int, speculative_num_steps: int): + if topk != 1: + raise NotImplementedError( + "Qwen4-Exp QSA MTP currently supports speculative_eagle_topk=1" + ) + self.model_runner = model_runner + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.attn_backends = [ + QwenSparseAttnBackend(model_runner) + for _ in range(speculative_num_steps - 1) + ] + + @staticmethod + def _as_cpu_lengths(seq_lens_cpu, seq_lens: torch.Tensor) -> torch.Tensor: + if seq_lens_cpu is None: + return seq_lens.detach().cpu().to(torch.int32) + if isinstance(seq_lens_cpu, torch.Tensor): + return seq_lens_cpu.to(dtype=torch.int32) + return torch.tensor(seq_lens_cpu, dtype=torch.int32) + + def _step_out_cache_loc(self, forward_batch, step: int): + """Mirror EagleDraftWorker.draft_forward's per-step out_cache_loc slice; + a mismatch makes every QSA step write into the first ``batch_size`` slots.""" + + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + steps = self.speculative_num_steps + batch_rows = int( + getattr( + forward_batch, + "batch_size", + forward_batch.seq_lens.numel(), + ) + ) + if out_cache_loc is None or steps <= 1: + return out_cache_loc + if out_cache_loc.numel() != batch_rows * self.topk * steps: + # Idle/zero-row and non-draft batches do not use the interleaved layout. + return out_cache_loc + # Same expression chain as EagleDraftWorker.draft_forward (views only). + return ( + out_cache_loc.reshape(batch_rows, self.topk, steps) + .permute(2, 0, 1) + .reshape(steps, -1)[step] + ) + + def _make_step_forward_batch(self, forward_batch, step: int, num_padding: int = 0): + step_forward_batch = copy(forward_batch) + step_forward_batch.forward_mode = ForwardMode.DECODE + step_forward_batch.seq_lens = (forward_batch.seq_lens + step + 1).to( + torch.int32 + ) + if forward_batch.seq_lens_cpu is None: + # GPU-only serving: downstream metadata paths derive host bounds + # from batch shape; do not force a per-step D2H here. + step_forward_batch.seq_lens_cpu = None + else: + step_forward_batch.seq_lens_cpu = self._as_cpu_lengths( + forward_batch.seq_lens_cpu, forward_batch.seq_lens + ) + (step + 1) + num_padding = max( + 0, + min(int(num_padding), int(step_forward_batch.seq_lens.numel())), + ) + if num_padding: + step_forward_batch.seq_lens = step_forward_batch.seq_lens.clone() + step_forward_batch.seq_lens[-num_padding:] = 1 + if step_forward_batch.seq_lens_cpu is not None: + step_forward_batch.seq_lens_cpu = ( + step_forward_batch.seq_lens_cpu.clone() + ) + step_forward_batch.seq_lens_cpu[-num_padding:] = 1 + step_forward_batch.batch_size = int(step_forward_batch.seq_lens.numel()) + step_forward_batch.out_cache_loc = self._step_out_cache_loc(forward_batch, step) + return step_forward_batch + + def set_mtp_shared_sparse_indices(self, state) -> None: + for backend in self.attn_backends: + backend.set_mtp_shared_sparse_indices(state) + + def init_forward_metadata(self, forward_batch): + for step, backend in enumerate(self.attn_backends): + backend.init_forward_metadata( + self._make_step_forward_batch(forward_batch, step) + ) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for backend in self.attn_backends: + backend.init_cuda_graph_state(max_bs, max_num_tokens) + + def init_forward_metadata_out_graph(self, forward_batch, in_capture: bool = False): + if in_capture: + for step, backend in enumerate(self.attn_backends): + # Warmup must not write via shared req_pool_idx 0; + # pad every capture row to stay below the first compression boundary. + step_batch = self._make_step_forward_batch( + forward_batch, + step, + num_padding=forward_batch.batch_size, + ) + backend._capture_cuda_graph_metadata( + bs=step_batch.batch_size, + num_tokens=step_batch.batch_size, + req_pool_indices=step_batch.req_pool_indices, + seq_lens=step_batch.seq_lens, + forward_mode=ForwardMode.DECODE, + spec_info=step_batch.spec_info, + ) + return + + num_padding = getattr(forward_batch, "num_padding", None) + num_padding = num_padding if num_padding is not None else 0 + for step, backend in enumerate(self.attn_backends): + step_batch = self._make_step_forward_batch( + forward_batch, step, num_padding=num_padding + ) + backend._replay_cuda_graph_metadata( + bs=step_batch.batch_size, + req_pool_indices=step_batch.req_pool_indices, + seq_lens=step_batch.seq_lens, + forward_mode=ForwardMode.DECODE, + spec_info=step_batch.spec_info, + seq_lens_cpu=step_batch.seq_lens_cpu, + num_padding=num_padding, + ) + + def init_forward_metadata_in_graph(self, forward_batch) -> None: + pass + + +__all__ = [ + "is_qwen_qsa", + "QwenSparseAttnMetadata", + "QwenSparseAttnBackend", + "QwenSparseMultiStepDraftBackend", +] diff --git a/python/sglang/srt/layers/hc_mix_triton.py b/python/sglang/srt/layers/hc_mix_triton.py new file mode 100644 index 000000000..325369d62 --- /dev/null +++ b/python/sglang/srt/layers/hc_mix_triton.py @@ -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 diff --git a/python/sglang/srt/layers/hyperconnection.py b/python/sglang/srt/layers/hyperconnection.py new file mode 100644 index 000000000..61516924d --- /dev/null +++ b/python/sglang/srt/layers/hyperconnection.py @@ -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, +} diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=160,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=160,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..666162be0 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=160,device_name=NVIDIA_H200.json @@ -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 + } +} diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=320,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=320,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..5676817c1 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=320,device_name=NVIDIA_H200.json @@ -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 + } +} diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=80,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=80,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..ef2d06cbd --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=512,N=80,device_name=NVIDIA_H200.json @@ -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 + } +} diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 1eeca21f8..c0dc373a4 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -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 diff --git a/python/sglang/srt/layers/rotary_embedding/mrope_rope_index.py b/python/sglang/srt/layers/rotary_embedding/mrope_rope_index.py index e074bd84a..f0e17f3a2 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope_rope_index.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope_rope_index.py @@ -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", diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index 74dd32e82..761f8e554 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -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 diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index fffed010b..705709a05 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -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, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 1ec550735..5133a2276 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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 diff --git a/python/sglang/srt/mem_cache/ple_state_pool.py b/python/sglang/srt/mem_cache/ple_state_pool.py new file mode 100644 index 000000000..a4242636d --- /dev/null +++ b/python/sglang/srt/mem_cache/ple_state_pool.py @@ -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 + ) diff --git a/python/sglang/srt/mem_cache/qsa_kv_pool.py b/python/sglang/srt/mem_cache/qsa_kv_pool.py new file mode 100644 index 000000000..f8c1c0e8e --- /dev/null +++ b/python/sglang/srt/mem_cache/qsa_kv_pool.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 5699a0f99..8810e5fba 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -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 " diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 62b9c054e..a9a2aa8df 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -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. diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index 54850895e..bbee9773b 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -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: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 473f099e2..4e70ad77f 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2de7eeee9..d3691ea5e 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -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 diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index f282eedbe..92bd58687 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -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 diff --git a/python/sglang/srt/models/qwen4_exp.py b/python/sglang/srt/models/qwen4_exp.py new file mode 100644 index 000000000..32a51803e --- /dev/null +++ b/python/sglang/srt/models/qwen4_exp.py @@ -0,0 +1,2131 @@ +"""Inference-only Qwen4-Exp (text + VL) on the Qwen3.5 backbone.""" + +import math +from contextlib import nullcontext +from typing import Any, Iterable, Optional, Set, Tuple + +import msgspec +import sympy +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from torch import nn + +from sglang.kernels.ops.elementwise.elementwise import fused_sigmoid_mul +from sglang.srt.configs.qwen4_exp import Qwen4ExpConfig, Qwen4ExpTextConfig +from sglang.srt.distributed import get_tp_group, tensor_model_parallel_all_reduce +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.layers.communicator import get_attn_tp_context +from sglang.srt.layers.dp_attention import ( + attn_tp_all_gather, + attn_tp_all_reduce, + dp_gather_replicate, + dp_scatter, + get_attention_dp_size, + get_dp_global_num_tokens, + get_global_dp_buffer, + get_local_dp_buffer, + is_allocation_symmetric, + is_dp_attention_enabled, +) +from sglang.srt.layers.hyperconnection import ( + GatedResidual, + HyperConnectionConfig, +) +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod +from sglang.srt.layers.utils import get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.forward_context import ( + get_attn_backend, + get_req_to_token_pool, +) +from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.qwen3_5 import ( + Qwen3_5AttentionDecoderLayer, + Qwen3_5ForCausalLM, + Qwen3_5GatedDeltaNet, + Qwen3_5LinearDecoderLayer, +) +from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils import logger + +# Decode/verify-sized batches only: at prefill sizes both chains are compute +# bound and serializing them on one stream is faster than contending. +_QSA_INDEXER_OVERLAP_TOKEN_THRESHOLD = 1024 + + +def _get_ple_forward_mode(forward_batch: ForwardBatch) -> ForwardMode: + if forward_batch._original_forward_mode is not None: + return forward_batch._original_forward_mode + return forward_batch.forward_mode + + +def _get_processed_token_count( + forward_batch: ForwardBatch, physical_tokens: int +) -> int: + processed_tokens = forward_batch.global_num_token_non_padded_cpu + if processed_tokens is None and forward_batch.extend_seq_lens_cpu is not None: + processed_tokens = sum(forward_batch.extend_seq_lens_cpu) + if processed_tokens is None: + return physical_tokens + processed_tokens = int(processed_tokens) + if not 0 <= processed_tokens <= physical_tokens: + raise RuntimeError( + f"invalid PLE token counts: {processed_tokens=}, {physical_tokens=}" + ) + return processed_tokens + + +class _PLEBatch(msgspec.Struct, frozen=True): + mode: ForwardMode + use_decode_fast_path: bool + physical_tokens: int + processed_tokens: int + lengths: torch.Tensor + row_width: int + req_indices: torch.Tensor + token_offsets: torch.Tensor + valid_tokens: torch.Tensor + state_indices: torch.Tensor + ngram_context: Optional[torch.Tensor] + ngram_eos_token_id: Optional[int] + + +def _prepare_ple_batch( + input_ids: torch.Tensor, + forward_batch: ForwardBatch, + *, + ngram_size: Optional[int], + ngram_eos_token_id: Optional[int], +) -> Optional[_PLEBatch]: + """Prepare the token layout and the shared N-gram history once per forward.""" + + if forward_batch.tbo_parent_token_range is not None: + raise NotImplementedError("Qwen4 PLE is not compatible with two-batch overlap") + spec_algorithm = forward_batch.spec_algorithm + if spec_algorithm is not None and spec_algorithm.is_ngram(): + raise NotImplementedError("Qwen4 PLE does not support NGRAM speculation") + if ( + forward_batch.spec_info is not None + and getattr(forward_batch.spec_info, "topk", 1) != 1 + ): + raise NotImplementedError("Qwen4 PLE speculative decoding supports only topk=1") + + mode = _get_ple_forward_mode(forward_batch) + get_req_to_token_pool().ple_window_cache = None + if mode.is_idle(): + return None + use_decode_fast_path = ( + envs.SGLANG_ENABLE_QWEN4_PLE_FUSION.get() and mode.is_decode() + ) + + if input_ids.dim() > 1: + input_ids = input_ids.reshape(-1) + physical_tokens = input_ids.shape[0] + processed_tokens = _get_processed_token_count(forward_batch, physical_tokens) + tokens = input_ids[:processed_tokens] + positions = torch.arange(processed_tokens, device=tokens.device, dtype=torch.long) + + if mode.is_target_verify(): + assert forward_batch.spec_info is not None + row_width = int(forward_batch.spec_info.draft_token_num) + if row_width <= 0 or processed_tokens % row_width != 0: + raise RuntimeError( + "target verify rows must contain complete draft strides: " + f"{processed_tokens=} {row_width=}" + ) + sequence_count = processed_tokens // row_width + # Eager verify rows can be shorter than the row stride; + # ignore the synthetic one-token lengths from DP attention's verify-as-EXTEND. + lengths = ( + forward_batch.extend_seq_lens[:sequence_count].long() + if forward_batch.forward_mode.is_target_verify() + and forward_batch.extend_seq_lens is not None + else torch.full( + (sequence_count,), + row_width, + dtype=torch.long, + device=tokens.device, + ) + ) + if lengths.shape[0] != sequence_count: + raise RuntimeError( + "target verify length metadata does not match its fixed rows: " + f"{lengths.shape[0]=} {sequence_count=}" + ) + req_indices = torch.div(positions, row_width, rounding_mode="floor") + token_offsets = positions - req_indices * row_width + elif mode.is_decode(): + lengths = torch.ones(processed_tokens, dtype=torch.long, device=tokens.device) + row_width = 1 + req_indices = positions + token_offsets = torch.zeros_like(positions) + else: + if forward_batch.extend_seq_lens is None: + raise RuntimeError(f"PLE requires sequence lengths in {mode!r}") + lengths = forward_batch.extend_seq_lens.long() + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + row_width = ( + max(extend_seq_lens_cpu, default=0) + if extend_seq_lens_cpu is not None + else processed_tokens + ) + query_start_loc = torch.cat( + [lengths.new_zeros(1), torch.cumsum(lengths, dim=0)] + ) + sequence_count = lengths.shape[0] + req_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 + if processed_tokens: + req_indices = req_indices.clamp(min=0, max=sequence_count - 1) + token_offsets = positions - query_start_loc.index_select(0, req_indices) + + sequence_count = lengths.shape[0] + if use_decode_fast_path: + # One token per decode row: every offset is valid, no index-select needed. + valid_tokens = torch.ones( + processed_tokens, dtype=torch.bool, device=tokens.device + ) + else: + valid_tokens = token_offsets < lengths.index_select(0, req_indices) + + state_indices = ( + get_req_to_token_pool() + .get_mamba_indices(forward_batch.req_pool_indices[:sequence_count]) + .long() + ) + + # CUDA graph padding uses request slot 0, which may belong to a real request. + # Map padded sequences to the state pools' reserved dummy slot instead. + out_cache_loc = forward_batch.out_cache_loc + if use_decode_fast_path: + if out_cache_loc is not None: + state_indices = torch.where( + out_cache_loc[:sequence_count].ne(0), + state_indices, + torch.zeros_like(state_indices), + ) + else: + valid = lengths.ne(0) + if out_cache_loc is not None and mode.is_decode(): + valid = valid & out_cache_loc[:sequence_count].ne(0) + elif out_cache_loc is not None and mode.is_target_verify(): + valid = valid & out_cache_loc[:processed_tokens].reshape( + sequence_count, row_width + ).ne(0).any(dim=1) + state_indices = torch.where( + valid, state_indices, torch.zeros_like(state_indices) + ) + + ngram_context = None + if ngram_size is not None: + assert ngram_eos_token_id is not None + if use_decode_fast_path: + # One token per decode row: + # this view is the padded tensor the general path would materialize. + padded = tokens.unsqueeze(1) + else: + padded = tokens.new_full((sequence_count, row_width), ngram_eos_token_id) + if processed_tokens: + padded[req_indices, token_offsets] = torch.where( + valid_tokens, + tokens, + tokens.new_full((), ngram_eos_token_id), + ) + history = get_req_to_token_pool().get_ngram_context(state_indices) + if history.shape[1] != ngram_size - 1: + raise RuntimeError( + "Qwen4 PLE N-gram cache has the wrong context width: " + f"{history.shape[1]=} {ngram_size=}" + ) + ngram_context = torch.cat([history, padded], dim=1) + + return _PLEBatch( + mode=mode, + use_decode_fast_path=use_decode_fast_path, + physical_tokens=physical_tokens, + processed_tokens=processed_tokens, + lengths=lengths, + row_width=row_width, + req_indices=req_indices, + token_offsets=token_offsets, + valid_tokens=valid_tokens, + state_indices=state_indices, + ngram_context=ngram_context, + ngram_eos_token_id=ngram_eos_token_id, + ) + + +def _commit_ple_batch(batch: Optional[_PLEBatch], forward_batch: ForwardBatch) -> None: + """Commit the shared N-gram history after every PLE layer consumed it.""" + + if batch is None or batch.ngram_context is None or not batch.processed_tokens: + return + + pool = get_req_to_token_pool() + context = batch.ngram_context + context_len = context.shape[1] - batch.row_width + if batch.mode.is_target_verify(): + step_contexts = context.unfold(1, context_len, 1)[:, 1:] + valid_steps = batch.valid_tokens.reshape( + batch.lengths.shape[0], batch.row_width + ) + pool.set_ngram_intermediate_context( + torch.where( + valid_steps.unsqueeze(-1), + step_contexts, + torch.full_like(step_contexts, batch.ngram_eos_token_id), + ) + ) + return + + if batch.use_decode_fast_path: + # Decode advances every two-token history by exactly one column. Slicing + # preserves the int64 values while avoiding arange + gather launches. + next_context = context[:, batch.row_width :] + pool.set_ngram_context(batch.state_indices, next_context) + track = _ple_track_targets(forward_batch, batch) + if track is not None: + track_indices, _ = track + pool.set_ngram_context(track_indices, next_context) + return + + context_cols = torch.arange(context_len, device=context.device, dtype=torch.long) + next_context = context.gather( + 1, batch.lengths.unsqueeze(1) + context_cols.unsqueeze(0) + ) + pool.set_ngram_context(batch.state_indices, next_context) + + track = _ple_track_targets(forward_batch, batch) + if track is not None: + track_indices, track_offsets = track + pool.set_ngram_context( + track_indices, + context.gather(1, track_offsets.unsqueeze(1) + context_cols.unsqueeze(0)), + ) + + +def _ple_track_targets( + forward_batch: ForwardBatch, batch: _PLEBatch +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Destination slots and gather offsets for the extra-buffer track snapshot. + + With extra_buffer the radix tree caches the ping-pong track slot, + not the working slot. + Both PLE side states are laid out [incoming_state | chunk tokens], + so the boundary value is the state's own gather at a smaller offset; + callers differ only in tensor rank, hence offsets rather than a gather. + Masked-off rows route to reserved slot 0, so the shape stays graph-capturable. + None when tracking is inactive or its metadata is absent. + """ + track_indices = forward_batch.mamba_track_indices + track_mask = forward_batch.mamba_track_mask + if track_indices is None or track_mask is None: + return None + + rows = batch.lengths.shape[0] + track_indices = track_indices[:rows] + dst = torch.where(track_mask[:rows], track_indices, torch.zeros_like(track_indices)) + + if batch.mode.is_decode(): + # One token per step, so the boundary offset is the current one. Decode never + # carries mamba_track_seqlens, so this path must not consult it. + return dst, batch.lengths + + aligned = forward_batch.mamba_track_aligned_lens() + if aligned is None: + return None + + return dst, aligned[:rows].clamp(min=0).minimum(batch.lengths) + + +def _pad_token_rows(x: torch.Tensor, total_tokens: int) -> torch.Tensor: + if x.shape[0] == total_tokens: + return x + out = x.new_zeros((total_tokens, *x.shape[1:])) + out[: x.shape[0]] = x + return out + + +def _use_attn_tp_ngram() -> bool: + return is_dp_attention_enabled() and envs.SGLANG_USE_ATTN_TP_NGRAM.get() + + +class Qwen4ExpPLEGroupedNorm(nn.Module): + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + group_size: Optional[int] = None, + ) -> 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.eps = eps + self.group_size = group_size + self.weight = nn.Parameter(torch.zeros(hidden_size)) + # 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 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.eps) + compute_dtype = x.dtype + x_float = x.float() + if self.group_size is None: + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + else: + group_shape = x_float.shape[:-1] + (-1, self.group_size) + variance = x_float.reshape(group_shape).pow(2).mean(dim=-1, keepdim=True) + variance = variance.expand(group_shape).reshape_as(x_float) + x_norm = x_float * torch.rsqrt(variance + self.eps) + weight = self.weight.float() + 1.0 + return (x_norm * weight).to(compute_dtype) + + +class Qwen4ExpNGramEmbedding(nn.Module): + _MASK64 = (1 << 64) - 1 + _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 + _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 + _SPLITMIX_M2 = 0x94D049BB133111EB + _PRIME_1 = 10007 + + def __init__( + self, + config: Qwen4ExpTextConfig, + embedding_dim: int, + ple_layer_index: int = 0, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.config = config + self.ngram_embed_dim = int(embedding_dim) + self.ngram_size = int(config.ngram_size) + self.heads_per_ngram = int(config.heads_per_ngram) + self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram + self.ple_layer_index = int(ple_layer_index) + self.unigram_vocab_size = int(config.vocab_size) + if self.ngram_size < 2: + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") + if self.heads_per_ngram <= 0: + raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") + if self.ngram_embed_dim % self.ngram_heads != 0: + raise ValueError( + "ple_embed_dim must be divisible by total ngram heads: " + f"{self.ngram_embed_dim} % {self.ngram_heads} != 0" + ) + self.ngram_vocab_size_base = int(config.ngram_vocab_size_base) + if self.ngram_vocab_size_base <= 0: + raise ValueError("ngram_vocab_size_base must be > 0") + self.make_ngram_vocab_size_divisible_by = int( + config.make_ngram_vocab_size_divisible_by + ) + self.head_dim_per_ngram = self.ngram_embed_dim // self.ngram_heads + self.eos_token_id = int(config.eos_token_id) + self.enable_ple_fusion = envs.SGLANG_ENABLE_QWEN4_PLE_FUSION.get() + + self.register_buffer( + "layer_multipliers", + self._build_layer_multipliers(self.ngram_size), + persistent=True, + ) + head_vocab_sizes, head_offsets, total_vocab_size = ( + self._build_head_vocab_and_offsets() + ) + self.register_buffer( + "ngram_heads_vocab_sizes", + torch.tensor(head_vocab_sizes, dtype=torch.long), + persistent=True, + ) + self.register_buffer( + "ngram_heads_offsets", + torch.tensor(head_offsets, dtype=torch.long), + persistent=True, + ) + padded_vocab_size = ( + (total_vocab_size + self.make_ngram_vocab_size_divisible_by - 1) + // self.make_ngram_vocab_size_divisible_by + ) * self.make_ngram_vocab_size_divisible_by + self.use_attn_tp_ngram = _use_attn_tp_ngram() + self.gather_dp_tokens = ( + is_dp_attention_enabled() + and get_attention_dp_size() > 1 + and not self.use_attn_tp_ngram + ) + self.ngram_embedding = VocabParallelEmbedding( + padded_vocab_size, + self.head_dim_per_ngram, + params_dtype=( + torch.float8_e4m3fn + if (quant_config is not None and quant_config.get_name() == "fp8") + or getattr(config, "ple_embedding_dtype", None) == "float8_e4m3fn" + else torch.bfloat16 + ), + output_dtype=torch.bfloat16, + use_attn_tp_group=self.use_attn_tp_ngram, + ) + self.ngram_embedding.register_buffer( + "weight_scale", torch.ones(1, dtype=torch.bfloat16), persistent=True + ) + + @classmethod + def _splitmix64(cls, x: int) -> int: + x = (x + cls._SPLITMIX_GAMMA) & cls._MASK64 + x = ((x ^ (x >> 30)) * cls._SPLITMIX_M1) & cls._MASK64 + x = ((x ^ (x >> 27)) * cls._SPLITMIX_M2) & cls._MASK64 + return (x ^ (x >> 31)) & cls._MASK64 + + def _build_layer_multipliers(self, size: int) -> torch.Tensor: + seed = int(getattr(self.config, "seed", 1234)) + max_long = (1 << 63) - 1 + m_max = max_long // max(self.unigram_vocab_size, 1) + half_bound = max(1, m_max // 2) + values = [] + base_seed = seed + self._PRIME_1 * self.ple_layer_index + for idx in range(size): + x0 = (base_seed + self._SPLITMIX_GAMMA * (idx + 1)) & self._MASK64 + mixed = self._splitmix64(x0) + values.append(int(2 * (mixed % half_bound) + 1)) + return torch.tensor(values, dtype=torch.long) + + @staticmethod + def _find_nth_prime_after(start: int, n: int) -> int: + prime = int(start) + for _ in range(n): + prime = int(sympy.nextprime(prime)) + return prime + + def _build_head_vocab_and_offsets(self): + sizes = [] + offsets = [] + total = 0 + for head_idx in range(self.ngram_heads): + global_head_idx = self.ple_layer_index * self.ngram_heads + head_idx + size = self._find_nth_prime_after( + self.ngram_vocab_size_base - 1, global_head_idx + 1 + ) + sizes.append(size) + offsets.append(total) + total += size + return sizes, offsets, total + + def _embed_ngram_ids( + self, + ngram_ids: torch.Tensor, + forward_batch: ForwardBatch, + physical_tokens: int, + ) -> torch.Tensor: + lookup_ids, semantic_tokens = self._prepare_embedding_lookup( + ngram_ids, forward_batch, physical_tokens + ) + embeddings = self.ngram_embedding(lookup_ids) + embeddings = embeddings * self.ngram_embedding.weight_scale + return self._finish_embedding_lookup( + embeddings, semantic_tokens, forward_batch, physical_tokens + ) + + def _prepare_embedding_lookup( + self, + ngram_ids: torch.Tensor, + forward_batch: ForwardBatch, + physical_tokens: int, + ) -> Tuple[torch.Tensor, int]: + semantic_tokens = ngram_ids.shape[0] + if not self.gather_dp_tokens: + return ngram_ids, semantic_tokens + + padded_ngram_ids = _pad_token_rows(ngram_ids, physical_tokens) + global_tokens = forward_batch.global_dp_buffer_len + if global_tokens is None: + raise RuntimeError( + "global-TP Qwen4 N-gram lookup under DP attention requires a " + "DP token layout; set SGLANG_USE_ATTN_TP_NGRAM=1 to shard the " + "table within each attention-TP group" + ) + + global_ngram_ids = ngram_ids.new_empty((global_tokens, *ngram_ids.shape[1:])) + dp_gather_replicate( + global_ngram_ids, padded_ngram_ids.contiguous(), forward_batch + ) + return global_ngram_ids, semantic_tokens + + def _finish_embedding_lookup( + self, + embeddings: torch.Tensor, + semantic_tokens: int, + forward_batch: ForwardBatch, + physical_tokens: int, + ) -> torch.Tensor: + if not self.gather_dp_tokens: + return embeddings + local_embeddings = embeddings.new_empty( + (physical_tokens, *embeddings.shape[1:]) + ) + dp_scatter(local_embeddings, embeddings.contiguous(), forward_batch) + return local_embeddings[:semantic_tokens] + + def _hash_contexts( + self, contexts: torch.Tensor, *, decode_sized: bool = False + ) -> torch.Tensor: + contexts = contexts.to(torch.long) + if self.enable_ple_fusion and decode_sized: + from sglang.kernels.ops.qwen4_ple import ( + can_fuse_qwen4_ngram_hash, + fused_qwen4_ngram_hash, + ) + + if can_fuse_qwen4_ngram_hash( + contexts, + self.layer_multipliers, + self.ngram_heads_vocab_sizes, + self.ngram_heads_offsets, + ): + return fused_qwen4_ngram_hash( + contexts, + self.layer_multipliers, + self.ngram_heads_vocab_sizes, + self.ngram_heads_offsets, + self.eos_token_id, + ) + + pool = get_req_to_token_pool() + cached = pool.ple_window_cache + if cached is not None and cached[1] is contexts and cached[2] is not None: + shifted_tokens = cached[2] + assert len(shifted_tokens) == self.ngram_size + else: + shifted_tokens = [contexts] + for shift in range(1, self.ngram_size): + shifted_tokens.append(self._shift_right_ignore_eos(contexts, shift)) + if cached is not None and cached[1] is contexts: + pool.ple_window_cache = (cached[0], contexts, shifted_tokens) + + blocks = [] + for ngram in range(2, self.ngram_size + 1): + ngram_idx = ngram - 2 + start_idx = ngram_idx * self.heads_per_ngram + end_idx = start_idx + self.heads_per_ngram + mix = shifted_tokens[0] * self.layer_multipliers[0] + for pos in range(1, ngram): + mix = torch.bitwise_xor( + mix, shifted_tokens[pos] * self.layer_multipliers[pos] + ) + head_vocab_sizes = self.ngram_heads_vocab_sizes[start_idx:end_idx] + head_offsets = self.ngram_heads_offsets[start_idx:end_idx] + ngram_ids = torch.remainder( + mix[:, -1:].unsqueeze(-1), head_vocab_sizes.view(1, 1, -1) + ) + ngram_ids = ngram_ids + head_offsets.view(1, 1, -1) + blocks.append(ngram_ids[:, 0]) + return torch.cat(blocks, dim=-1) + + def _shift_right_ignore_eos(self, tensor: torch.Tensor, n: int) -> torch.Tensor: + if n == 0: + return tensor + batch_size, seq_len = tensor.shape + idx = torch.arange(seq_len, device=tensor.device, dtype=torch.long) + eos_mask = tensor == self.eos_token_id + eos_pos = torch.where(eos_mask, idx, -1) + prev_eos_inclusive = torch.cummax(eos_pos, dim=1).values + prev_eos = torch.cat( + [eos_pos.new_full((batch_size, 1), -1), prev_eos_inclusive[:, :-1]], + dim=1, + ) + segment_start = prev_eos + 1 + pos_in_segment = idx.unsqueeze(0) - segment_start + src_idx = idx - n + gather_idx = torch.clamp(src_idx, min=0).unsqueeze(0).expand(batch_size, -1) + shifted = tensor.gather(dim=1, index=gather_idx) + valid_mask = (pos_in_segment >= n) & (src_idx.unsqueeze(0) >= 0) + return torch.where(valid_mask, shifted, tensor.new_full((), self.eos_token_id)) + + def forward_idle(self, forward_batch: ForwardBatch) -> None: + if not self.gather_dp_tokens: + return + input_ids = forward_batch.input_ids.reshape(-1) + dummy_ids = input_ids.new_zeros((input_ids.shape[0], self.ngram_heads)) + self._embed_ngram_ids(dummy_ids, forward_batch, input_ids.shape[0]) + + def forward( + self, + batch: _PLEBatch, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + ngram_ids = self.compute_ngram_ids(batch) + embeddings = self._embed_ngram_ids( + ngram_ids, forward_batch, batch.physical_tokens + ) + return embeddings.flatten(start_dim=-2) + + def compute_ngram_ids(self, batch: _PLEBatch) -> torch.Tensor: + assert batch.ngram_context is not None + pool = get_req_to_token_pool() + cached = pool.ple_window_cache + if cached is not None and cached[0] is batch: + contexts = cached[1] + else: + if batch.use_decode_fast_path: + contexts = batch.ngram_context + else: + contexts = batch.ngram_context.unfold(1, self.ngram_size, 1)[ + batch.req_indices, batch.token_offsets + ] + contexts = contexts.to(torch.long) + pool.ple_window_cache = (batch, contexts, None) + return self._hash_contexts( + contexts, + decode_sized=batch.mode.is_decode() or batch.mode.is_target_verify(), + ) + + +@triton.jit +def _gather_ple_embedding_from_pinned_kernel( + weight_ptr, + ids_ptr, + output_ptr, + embedding_dim, + tp_vocab_start, + tp_vocab_end, + is_fp8: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row_id = tl.program_id(0) + global_idx = tl.load(ids_ptr + row_id) + in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end) + local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0) + offsets = tl.arange(0, BLOCK_D) + mask = offsets < embedding_dim + if is_fp8: + weight_ptr = weight_ptr.to(tl.int64).to(tl.pointer_type(tl.float8e4nv)) + else: + weight_ptr = weight_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16)) + values = tl.load( + weight_ptr + local_idx * embedding_dim + offsets, + mask=mask, + other=0.0, + ).to(tl.bfloat16) + tl.store( + output_ptr + row_id * embedding_dim + offsets, + tl.where(in_range, values, 0.0), + mask=mask, + ) + + +class Qwen4ExpPinnedHostEmbedding(VocabParallelEmbedding): + """PLE table read directly from pinned host memory. + + The table stays in its checkpoint storage dtype (fp8 with a per-tensor + weight_scale for fp8 checkpoints, bf16 otherwise); gathers emit bf16. + """ + + _COPIED_ATTRIBUTES = ( + "quant_config", + "enable_tp", + "use_attn_tp_group", + "tp_size", + "num_embeddings", + "org_vocab_size", + "padding_size", + "num_added_embeddings", + "use_presharded_weights", + "org_vocab_size_padded", + "num_embeddings_padded", + "shard_indices", + "embedding_dim", + "num_embeddings_per_partition", + "num_org_embeddings_per_partition", + "num_added_embeddings_per_partition", + ) + + def __init__(self, embedding: VocabParallelEmbedding) -> None: + nn.Module.__init__(self) + if not isinstance(embedding.quant_method, UnquantizedEmbeddingMethod): + raise NotImplementedError( + "PLE embedding offload requires an unquantized embedding table" + ) + if embedding.weight.dtype not in (torch.bfloat16, torch.float8_e4m3fn): + raise TypeError( + "PLE embedding offload requires bfloat16 or fp8 weights, got " + f"{embedding.weight.dtype}" + ) + if embedding.num_added_embeddings: + raise NotImplementedError( + "PLE embedding offload does not support added vocabulary rows" + ) + for name in self._COPIED_ATTRIBUTES: + setattr(self, name, getattr(embedding, name)) + # The unquantized CUDA post-load hook is a no-op. Exclude this CPU-only + # table so the generic loader does not stage it back to GPU unnecessarily. + self.quant_method = None + + source_weight = embedding.weight + cpu_weight = nn.Parameter( + torch.empty( + source_weight.shape, + dtype=source_weight.dtype, + device="cpu", + pin_memory=True, + ), + requires_grad=False, + ) + for name, value in vars(source_weight).items(): + setattr(cpu_weight, name, value) + cpu_weight.weight_loader = self.weight_loader + self.register_parameter("weight", cpu_weight) + # The scale is tiny; keep it with the model instead of offloading it + # with the table. + self.register_buffer("weight_scale", embedding.weight_scale, persistent=True) + del embedding.weight + self._block_d = triton.next_power_of_2(self.embedding_dim) + + def allocate_output( + self, shape: Tuple[int, ...], device: torch.device + ) -> torch.Tensor: + allocation_context = nullcontext() + if self.tp_size > 1: + allocation_context = use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ) + with allocation_context, torch.inference_mode(False): + # The gather kernel emits bf16 rows regardless of the table dtype. + return torch.empty(shape, dtype=torch.bfloat16, device=device) + + def gather( + self, input_ids: torch.Tensor, out: Optional[torch.Tensor] = None + ) -> torch.Tensor: + expected_shape = (*input_ids.shape, self.embedding_dim) + if out is None: + output = self.allocate_output(expected_shape, input_ids.device) + else: + if tuple(out.shape) != expected_shape: + raise ValueError( + f"invalid PLE prefetch output shape: {tuple(out.shape)} != " + f"{expected_shape}" + ) + if out.dtype != torch.bfloat16 or out.device != input_ids.device: + raise ValueError( + "PLE prefetch output must be bfloat16 on the id device" + ) + output = out + + flat_ids = input_ids.reshape(-1).long() + if flat_ids.numel(): + _gather_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)]( + self.weight.data_ptr(), + flat_ids, + output, + embedding_dim=self.embedding_dim, + tp_vocab_start=self.shard_indices.org_vocab_start_index, + tp_vocab_end=self.shard_indices.org_vocab_end_index, + is_fp8=self.weight.dtype == torch.float8_e4m3fn, + BLOCK_D=self._block_d, + ) + return output + + def reduce(self, output: torch.Tensor) -> torch.Tensor: + if self.tp_size > 1 and not get_attn_tp_context().input_scattered: + if self.use_attn_tp_group: + return attn_tp_all_reduce(output) + return tensor_model_parallel_all_reduce(output) + return output + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.reduce(self.gather(input_ids)) + + +class Qwen4ExpPLELayer(nn.Module): + def __init__( + self, + config: Qwen4ExpTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + layer_id: Optional[int] = None, + ple_layer_index: int = 0, + ) -> None: + super().__init__() + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.ple_embed_dim = config.ple_embed_dim + self.conv_kernel_size = config.ple_conv_kernel_size + self.hc_count = config.hc_count + self.hc_hidden_size = self.hidden_size * self.hc_count + self.ple_embedding = Qwen4ExpNGramEmbedding( + config, + self.ple_embed_dim, + ple_layer_index=ple_layer_index, + quant_config=quant_config, + ) + if config.ple_offload_embedding: + self.ple_embedding.ngram_embedding = Qwen4ExpPinnedHostEmbedding( + self.ple_embedding.ngram_embedding + ) + self.short_conv_dilation = self.ple_embedding.ngram_size + self.short_conv_state_len = ( + self.conv_kernel_size - 1 + ) * self.short_conv_dilation + self.conv_channels = self.hc_hidden_size + self.key_proj = ReplicatedLinear( + self.ple_embed_dim, + self.conv_channels, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.key_proj", + ) + self.value_proj = ReplicatedLinear( + self.ple_embed_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.value_proj", + ) + norm_hidden = self.hc_hidden_size + norm_group = self.hidden_size + self.norm_key = Qwen4ExpPLEGroupedNorm( + norm_hidden, + eps=config.rms_norm_eps, + group_size=norm_group, + ) + self.norm_query = Qwen4ExpPLEGroupedNorm( + norm_hidden, + eps=config.rms_norm_eps, + group_size=norm_group, + ) + self.norm_conv = Qwen4ExpPLEGroupedNorm( + norm_hidden, + eps=config.rms_norm_eps, + group_size=norm_group, + ) + self.conv1d = nn.Conv1d( + in_channels=self.conv_channels, + out_channels=self.conv_channels, + kernel_size=self.conv_kernel_size, + groups=self.conv_channels, + padding=(self.conv_kernel_size - 1) * self.short_conv_dilation, + dilation=self.short_conv_dilation, + bias=False, + ) + nn.init.zeros_(self.conv1d.weight) + self._prefetch_stream = ( + torch.cuda.Stream() if config.ple_offload_embedding else None + ) + self._graph_prefetch_buffers = {} + self._eager_prefetch_buffer = None + self._prefetch_state = None + + def _apply_ple_norm(self, norm: nn.Module, x: torch.Tensor) -> torch.Tensor: + y = norm(x.flatten(-2, -1)) + return y.unflatten(-1, (self.hc_count, self.hidden_size)) + + def _short_conv( + self, + x: torch.Tensor, + forward_batch: ForwardBatch, + batch: _PLEBatch, + ) -> torch.Tensor: + if x.shape[0] == 0: + return x + pool = get_req_to_token_pool() + conv_state = pool.short_conv_layer_cache(self.layer_id) + + if batch.use_decode_fast_path: + # With row_width=1 the padded/transpose path is x.unsqueeze(-1), + # and each state boundary is a one-column shift; conv and SiLU stay native. + from sglang.kernels.ops.qwen4_ple import ( + can_fuse_qwen4_short_conv_state, + fused_qwen4_short_conv_state, + ) + + fused_state = can_fuse_qwen4_short_conv_state( + conv_state, batch.state_indices, x + ) + if fused_state: + conv_input = fused_qwen4_short_conv_state( + conv_state, batch.state_indices, x + ) + else: + state = conv_state.index_select(0, batch.state_indices).to( + dtype=x.dtype + ) + conv_input = torch.cat([state, x.unsqueeze(-1)], dim=-1) + conv_output = F.conv1d( + conv_input, + self.conv1d.weight.to(dtype=x.dtype), + bias=None, + dilation=self.short_conv_dilation, + groups=self.conv_channels, + ).squeeze(-1) + next_state = conv_input[:, :, batch.row_width :] + if not fused_state: + conv_state[batch.state_indices] = next_state.to(dtype=conv_state.dtype) + + track = _ple_track_targets(forward_batch, batch) + if track is not None: + track_indices, _ = track + conv_state[track_indices] = next_state.to(dtype=conv_state.dtype) + return F.silu(conv_output) + + state = conv_state.index_select(0, batch.state_indices).to(dtype=x.dtype) + padded_seq = x.new_zeros( + (batch.lengths.shape[0], batch.row_width, self.conv_channels) + ) + padded_seq[batch.req_indices, batch.token_offsets] = x + conv_input = torch.cat([state, padded_seq.transpose(1, 2)], dim=-1) + conv_output = F.conv1d( + conv_input, + self.conv1d.weight.to(dtype=x.dtype), + bias=None, + dilation=self.short_conv_dilation, + groups=self.conv_channels, + ).transpose(1, 2) + + if batch.mode.is_target_verify(): + intermediate_cache = pool.short_conv_layer_intermediate_cache(self.layer_id) + if intermediate_cache is not None: + if self.short_conv_state_len: + intermediate_state = ( + conv_input.unfold(2, self.short_conv_state_len, 1)[ + :, :, 1 : batch.row_width + 1 + ] + .permute(0, 2, 1, 3) + .contiguous() + ) + else: + intermediate_state = x.new_empty( + ( + batch.lengths.shape[0], + batch.row_width, + self.conv_channels, + 0, + ) + ) + valid_steps = batch.valid_tokens.reshape( + batch.lengths.shape[0], batch.row_width, 1, 1 + ) + intermediate_state = torch.where( + valid_steps, + intermediate_state, + torch.zeros_like(intermediate_state), + ) + intermediate_cache[: batch.lengths.shape[0], : batch.row_width].copy_( + intermediate_state.to(dtype=intermediate_cache.dtype) + ) + else: + state_cols = torch.arange( + self.short_conv_state_len, device=x.device, dtype=torch.long + ) + + def _gather_at(offsets: torch.Tensor) -> torch.Tensor: + return conv_input.gather( + 2, + (offsets.unsqueeze(1) + state_cols.unsqueeze(0)) + .unsqueeze(1) + .expand(-1, self.conv_channels, -1), + ) + + next_state = _gather_at(batch.lengths) + conv_state[batch.state_indices] = next_state.to(dtype=conv_state.dtype) + + # Same boundary mamba uses, into the slot the radix tree reads. + track = _ple_track_targets(forward_batch, batch) + if track is not None: + track_indices, track_offsets = track + conv_state[track_indices] = _gather_at(track_offsets).to( + dtype=conv_state.dtype + ) + + return F.silu(conv_output[batch.req_indices, batch.token_offsets]) + + def forward_idle(self, forward_batch: ForwardBatch) -> None: + if self._prefetch_state is not None: + self._consume_prefetched_embeddings(forward_batch) + else: + self.ple_embedding.forward_idle(forward_batch) + + def _allocate_prefetch_buffer( + self, lookup_tokens: int, lookup_ids: torch.Tensor + ) -> torch.Tensor: + offloaded_embedding = self.ple_embedding.ngram_embedding + return offloaded_embedding.allocate_output( + (lookup_tokens, self.ple_embed_dim), lookup_ids.device + ) + + def _get_prefetch_buffer( + self, lookup_tokens: int, lookup_ids: torch.Tensor + ) -> torch.Tensor: + if get_is_capture_mode(): + buffer = self._graph_prefetch_buffers.get(lookup_tokens) + if buffer is None: + buffer = self._allocate_prefetch_buffer(lookup_tokens, lookup_ids) + self._graph_prefetch_buffers[lookup_tokens] = buffer + return buffer + + buffer = self._eager_prefetch_buffer + if buffer is None or buffer.shape[0] < lookup_tokens: + buffer = self._allocate_prefetch_buffer(lookup_tokens, lookup_ids) + self._eager_prefetch_buffer = buffer + return buffer[:lookup_tokens] + + def start_prefetch( + self, + batch: Optional[_PLEBatch], + forward_batch: ForwardBatch, + ) -> None: + """Gather PLE rows via UVA while the preceding decoder layer runs.""" + if self._prefetch_stream is None: + return + if self._prefetch_state is not None: + raise RuntimeError("PLE prefetch state was not consumed before reuse") + if batch is None: + if not self.ple_embedding.gather_dp_tokens: + return + physical_tokens = forward_batch.input_ids.numel() + ngram_ids = forward_batch.input_ids.new_zeros( + (physical_tokens, self.ple_embedding.ngram_heads) + ) + else: + physical_tokens = batch.physical_tokens + ngram_ids = self.ple_embedding.compute_ngram_ids(batch) + + lookup_ids, semantic_tokens = self.ple_embedding._prepare_embedding_lookup( + ngram_ids, forward_batch, physical_tokens + ) + lookup_tokens = lookup_ids.shape[0] + if lookup_tokens == 0: + return + prefetched = self._get_prefetch_buffer(lookup_tokens, lookup_ids) + output_view = prefetched.view(lookup_tokens, self.ple_embedding.ngram_heads, -1) + offloaded_embedding = self.ple_embedding.ngram_embedding + + stream = self._prefetch_stream + stream.wait_stream(torch.cuda.current_stream()) + lookup_ids.record_stream(stream) + with torch.cuda.stream(stream): + offloaded_embedding.gather(lookup_ids, out=output_view) + self._prefetch_state = prefetched, semantic_tokens, physical_tokens + + def _consume_prefetched_embeddings( + self, forward_batch: ForwardBatch + ) -> torch.Tensor: + if self._prefetch_state is None: + raise RuntimeError("PLE prefetch state is missing") + embeddings, semantic_tokens, physical_tokens = self._prefetch_state + torch.cuda.current_stream().wait_stream(self._prefetch_stream) + embeddings = self.ple_embedding.ngram_embedding.reduce(embeddings) + embeddings = embeddings * self.ple_embedding.ngram_embedding.weight_scale + embeddings = self.ple_embedding._finish_embedding_lookup( + embeddings, + semantic_tokens, + forward_batch, + physical_tokens, + ) + self._prefetch_state = None + return embeddings + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + batch: _PLEBatch, + ) -> torch.Tensor: + hidden_states = hidden_states[: batch.processed_tokens] + if self._prefetch_state is not None: + embeddings = self._consume_prefetched_embeddings(forward_batch) + else: + embeddings = self.ple_embedding(batch, forward_batch) + key, _ = self.key_proj(embeddings) + value, _ = self.value_proj(embeddings) + token_count = hidden_states.shape[0] + hidden_size = self.hidden_size + hc_count = self.hc_count + if hidden_states.shape[-1] != hc_count * hidden_size: + raise RuntimeError( + "PLE hidden size does not match its hyper-connection layout: " + f"expected {hc_count * hidden_size}, got {hidden_states.shape[-1]}" + ) + key = key.reshape(token_count, hc_count, hidden_size) + query = hidden_states.reshape(token_count, hc_count, hidden_size) + key_normed = self._apply_ple_norm(self.norm_key, key) + query_normed = self._apply_ple_norm(self.norm_query, query) + gate = (key_normed * query_normed).sum(dim=-1, keepdim=True) + gate = gate / math.sqrt(hidden_size) + fused_gate_value = False + if batch.use_decode_fast_path: + from sglang.kernels.ops.qwen4_ple import ( + can_fuse_qwen4_gate_value, + fused_qwen4_gate_value, + ) + + fused_gate_value = can_fuse_qwen4_gate_value(gate, value) + if fused_gate_value: + gated_value = fused_qwen4_gate_value(gate, value) + else: + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gate = torch.sigmoid(gate) + gated_value = gate * value.unsqueeze(-2) + gated_value_normed = self._apply_ple_norm(self.norm_conv, gated_value) + gated_value = gated_value.flatten(-2) + gated_value_normed = gated_value_normed.flatten(-2) + conv_output = self._short_conv( + gated_value_normed, + forward_batch, + batch, + ) + output = gated_value + conv_output + if not batch.use_decode_fast_path: + output = torch.where( + batch.valid_tokens.unsqueeze(-1), + output, + torch.zeros_like(output), + ) + return _pad_token_rows(output, batch.physical_tokens) + + +class Qwen4ExpLayerExtensionMixin: + def _init_qwen4_exp_layer_extensions( + self, + config: Qwen4ExpTextConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + self.hc_count = config.hc_count + self.hidden_size = config.hidden_size + self.ple = None + + for attr_name in ( + "input_layernorm", + "post_attention_layernorm", + "layer_communicator", + ): + if hasattr(self, attr_name): + delattr(self, attr_name) + + if (layer_id + 1) in config.ple_layer_ids: + ple_layer_ids_sorted = sorted(set(config.ple_layer_ids)) + ple_layer_index = { + abs_id: index for index, abs_id in enumerate(ple_layer_ids_sorted) + }[layer_id + 1] + # Strip the block-type segment like the dense mlp (PLE is attn's sibling); + # else the quant prefix misses the ckpt skip-list -> NaN. + ple_prefix = prefix.replace(".linear_attn", "").replace(".self_attn", "") + self.ple = Qwen4ExpPLELayer( + config, + quant_config=quant_config, + prefix=f"{ple_prefix}.ple" if ple_prefix else "ple", + layer_id=layer_id, + ple_layer_index=ple_layer_index, + ) + + hc_config = HyperConnectionConfig( + hc_count=self.hc_count, + hidden_size=self.hidden_size, + params_dtype=torch.bfloat16, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.attn_hyper_connection = GatedResidual( + hc_config, + use_mix=True, + use_combine=True, + ) + self.mlp_hyper_connection = GatedResidual( + hc_config, + use_mix=True, + use_combine=True, + ) + + def _prepare_qwen4_exp_attn( + self, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + forward_batch: ForwardBatch, + *, + ple_batch: Optional[_PLEBatch], + ): + hc_dim = self.hc_count * self.hidden_size + if hidden_states.shape[-1] != hc_dim: + assert hidden_states.shape[-1] == self.hidden_size + hidden_states = torch.cat( + [hidden_states for _ in range(self.hc_count)], dim=-1 + ) + + if self.ple is not None: + if ple_batch is None: + if not _get_ple_forward_mode(forward_batch).is_idle(): + raise RuntimeError( + "non-idle Qwen4 PLE forward is missing its batch" + ) + self.ple.forward_idle(forward_batch) + else: + ple_query = ( + hidden_states if residual is None else hidden_states + residual + ) + hidden_states = hidden_states + self.ple( + ple_query, forward_batch, ple_batch + ) + + hidden_states, residual = self.attn_hyper_connection.mix(hidden_states) + return hidden_states, residual + + def _prepare_qwen4_exp_mlp( + self, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + forward_batch: ForwardBatch, + ): + if not forward_batch.forward_mode.is_idle(): + hidden_states = attn_tp_all_reduce(hidden_states) + hidden_states = self.attn_hyper_connection.combine(hidden_states, residual) + hidden_states, residual = self.mlp_hyper_connection.mix(hidden_states) + return hidden_states, residual + + def _qwen4_exp_use_dp_moe_gather(self) -> bool: + return get_attention_dp_size() > 1 and get_moe_a2a_backend().is_none() + + def _qwen4_exp_use_attn_tp_a2a_scatter(self) -> bool: + return get_parallel().attn_tp_size > 1 and not get_moe_a2a_backend().is_none() + + def _run_qwen4_exp_mlp( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + if not self.config.num_experts: + return self.mlp(hidden_states) + + use_dp_moe_gather = self._qwen4_exp_use_dp_moe_gather() + use_attn_tp_a2a_scatter = self._qwen4_exp_use_attn_tp_a2a_scatter() + + if use_dp_moe_gather: + hidden_states, local_hidden_states = ( + get_global_dp_buffer(get_tp_group()), + hidden_states, + ) + dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) + elif hidden_states.shape[0] == 0 and get_moe_a2a_backend().is_none(): + # Only safe to short-circuit an empty batch when the MoE holds no collective; + # under deepep an idle DP rank must still join dispatch/combine or peers hang. + return hidden_states + + attn_tp_chunks = None + if use_attn_tp_a2a_scatter: + attn_tp_size = get_parallel().attn_tp_size + attn_tp_chunks = list(hidden_states.tensor_split(attn_tp_size)) + hidden_states = attn_tp_chunks[get_parallel().attn_tp_rank].contiguous() + + hidden_states = self.mlp(hidden_states, forward_batch) + + if use_dp_moe_gather: + hidden_states, global_hidden_states = ( + get_local_dp_buffer(get_tp_group()), + hidden_states, + ) + if should_use_dp_reduce_scatterv(): + get_tp_group().reduce_scatterv( + global_hidden_states, + output=hidden_states, + sizes=get_dp_global_num_tokens(), + ) + else: + dp_scatter(hidden_states, global_hidden_states, forward_batch) + elif use_attn_tp_a2a_scatter: + assert attn_tp_chunks is not None + gathered = [torch.empty_like(t) for t in attn_tp_chunks] + attn_tp_all_gather(gathered, hidden_states.contiguous()) + hidden_states = torch.cat(gathered) + + return hidden_states + + def _postprocess_qwen4_exp_layer( + self, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + forward_batch: ForwardBatch, + ): + hidden_states = self.mlp_hyper_connection.combine(hidden_states, residual) + return hidden_states, None + + +class Qwen4ExpLinearDecoderLayer( + Qwen4ExpLayerExtensionMixin, Qwen3_5LinearDecoderLayer +): + def __init__( + self, + config: Qwen4ExpTextConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, + ) -> None: + super().__init__(config, layer_id, quant_config, prefix, alt_stream, is_nextn) + self._init_qwen4_exp_layer_extensions(config, layer_id, quant_config, prefix) + + def forward( + self, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + **kwargs, + ): + forward_batch = kwargs.get("forward_batch", None) + + hidden_states, residual = self._prepare_qwen4_exp_attn( + hidden_states, + residual, + forward_batch, + ple_batch=kwargs.get("ple_batch"), + ) + + if not forward_batch.forward_mode.is_idle(): + hidden_states = self.linear_attn(hidden_states, forward_batch) + + hidden_states, residual = self._prepare_qwen4_exp_mlp( + hidden_states, residual, forward_batch + ) + hidden_states = self._run_qwen4_exp_mlp(hidden_states, forward_batch) + return self._postprocess_qwen4_exp_layer(hidden_states, residual, forward_batch) + + +class Qwen4ExpAttentionDecoderLayer( + Qwen4ExpLayerExtensionMixin, Qwen3_5AttentionDecoderLayer +): + def __init__( + self, + config: Qwen4ExpTextConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, + ) -> None: + config.attn_output_gate = True + super().__init__(config, layer_id, quant_config, prefix, alt_stream, is_nextn) + from sglang.srt.layers.attention.qsa.config import is_qwen_qsa + from sglang.srt.layers.attention.qsa.glue import build_qsa_indexer + + self.is_qsa = is_qwen_qsa(config) + if self.is_qsa: + self.indexer = build_qsa_indexer( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.indexer" if prefix else "indexer", + rotary_emb=self.rotary_emb, + ) + self._init_qwen4_exp_layer_extensions(config, layer_id, quant_config, prefix) + + def _compute_qsa_topk_indices( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + from sglang.srt.layers.attention.qsa.glue import ( + get_qsa_indexer_metadata, + resolve_qsa_sparse_backend, + ) + + backend = get_attn_backend() + sparse_backend = resolve_qsa_sparse_backend(backend) + should_reuse = getattr(sparse_backend, "should_reuse_mtp_sparse_indices", None) + if should_reuse is not None and should_reuse(forward_batch): + # MTP decode steps reuse the draft-extend's target-aligned + # selection; the indexer never runs inside the decode graph. + return sparse_backend.lookup_mtp_sparse_indices( + forward_batch, self.layer_id + ) + indexer_metadata = get_qsa_indexer_metadata( + backend, self.layer_id, forward_batch + ) + topk_indices = self.indexer( + hidden_states, + positions, + forward_batch, + indexer_metadata, + ) + should_capture = getattr( + sparse_backend, "should_capture_mtp_sparse_indices", None + ) + if should_capture is not None and should_capture(forward_batch): + sparse_backend.capture_mtp_sparse_indices( + topk_indices, forward_batch, self.layer_id, metadata=indexer_metadata + ) + return topk_indices + + def self_attention( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + overlap_indexer = ( + self.is_qsa + and self.alt_stream is not None + and get_is_capture_mode() + and hidden_states.shape[0] < _QSA_INDEXER_OVERLAP_TOKEN_THRESHOLD + ) + attention_kwargs = {} + if overlap_indexer: + # Safe to overlap: the indexer reads only hidden_states/positions, + # and writes QSA-private pool buffers. + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + with torch.cuda.stream(self.alt_stream): + topk_indices = self._compute_qsa_topk_indices( + hidden_states, positions, forward_batch + ) + + q, k, v, gate = self._prepare_qkv_gate( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + if overlap_indexer: + current_stream.wait_stream(self.alt_stream) + # Allocated on alt_stream, consumed by attention on the current + # stream; tell the caching allocator before alt_stream is reused. + topk_indices.record_stream(current_stream) + attention_kwargs["topk_indices"] = topk_indices + elif self.is_qsa: + attention_kwargs["topk_indices"] = self._compute_qsa_topk_indices( + hidden_states, positions, forward_batch + ) + + attn_output = self.attn(q, k, v, forward_batch, **attention_kwargs) + if gate is not None: + if attn_output.is_cuda: + # The strided 3D gate view feeds the kernel directly, so the + # gate reshape copy disappears along with the sigmoid + mul. + attn_output = fused_sigmoid_mul(attn_output, gate, inplace=True) + else: + gate = gate.reshape(gate.shape[0], -1) if gate.ndim == 3 else gate + attn_output = attn_output * torch.sigmoid(gate) + output, _ = self.o_proj(attn_output) + return output + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + forward_batch: ForwardBatch, + **kwargs: Any, + ): + hidden_states, residual = self._prepare_qwen4_exp_attn( + hidden_states, + residual, + forward_batch, + ple_batch=kwargs.get("ple_batch"), + ) + + if not forward_batch.forward_mode.is_idle(): + hidden_states = self.self_attention( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self._prepare_qwen4_exp_mlp( + hidden_states, residual, forward_batch + ) + hidden_states = self._run_qwen4_exp_mlp(hidden_states, forward_batch) + return self._postprocess_qwen4_exp_layer(hidden_states, residual, forward_batch) + + +ALL_DECODER_LAYER_TYPES = { + "attention": Qwen4ExpAttentionDecoderLayer, + "full_attention": Qwen4ExpAttentionDecoderLayer, + "linear_attention": Qwen4ExpLinearDecoderLayer, +} + + +class Qwen4ExpModel(Qwen3_5ForCausalLM): + decoder_layer_types = ALL_DECODER_LAYER_TYPES + + def _build_embed_tokens(self, config: Qwen4ExpTextConfig) -> nn.Module: + return VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + use_attn_tp_group=is_dp_attention_enabled(), + ) + + def __init__( + self, + config: Qwen4ExpTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + is_nextn: bool = False, + ) -> None: + super().__init__(config, quant_config, prefix, is_nextn) + self.hc_count = config.hc_count + self.hidden_size = config.hidden_size + self.has_ple = bool(config.ple_layer_ids) + self.ple_ngram_size = int(config.ngram_size) if self.has_ple else None + self.ple_ngram_eos_token_id = ( + int(config.eos_token_id) if self.ple_ngram_size is not None else None + ) + if hasattr(self, "norm"): + delattr(self, "norm") + hc_config = HyperConnectionConfig( + hc_count=self.hc_count, + hidden_size=self.hidden_size, + params_dtype=torch.bfloat16, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.hyper_connection_mixer = GatedResidual(hc_config, use_combine=False) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(input_ids) + + ple_batch = ( + _prepare_ple_batch( + input_ids, + forward_batch, + ngram_size=self.ple_ngram_size, + ngram_eos_token_id=self.ple_ngram_eos_token_id, + ) + if self.has_ple + else None + ) + residual = None + aux_hidden_states = [] + for i in range(self.start_layer, self.end_layer): + layer = self.layers[i] + if i + 1 < self.end_layer: + next_ple = getattr(self.layers[i + 1], "ple", None) + if next_ple is not None: + next_ple.start_prefetch(ple_batch, forward_batch) + with get_global_expert_distribution_recorder().with_current_layer(i): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + ple_batch=ple_batch, + captured_last_layer_outputs=( + aux_hidden_states + if getattr(layer, "_is_layer_to_capture", False) + else None + ), + ) + + _commit_ple_batch(ple_batch, forward_batch) + + hc_hidden_states = hidden_states + hidden_states, _ = self.hyper_connection_mixer.mix(hidden_states) + if not forward_batch.forward_mode.is_idle(): + return hidden_states, hc_hidden_states + + if len(aux_hidden_states) == 0: + return hidden_states + return hidden_states, aux_hidden_states + + +class Qwen4ExpVLModel(Qwen4ExpModel): + def __init__( + self, + config: Qwen4ExpTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config=config, quant_config=quant_config, prefix=prefix) + self.last_hc_hidden_states = None + + def get_input_embeddings(self) -> nn.Module: + return self.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[Any] = None, + input_deepstack_embeds: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self.last_hc_hidden_states = None + # mm routine passes input_ids=None; PLE needs the real ids. + if input_ids is None: + input_ids = forward_batch.input_ids + model_output = super().forward( + input_ids=input_ids, + positions=positions, + forward_batch=forward_batch, + inputs_embeds=input_embeds, + ) + if isinstance(model_output, tuple): + hidden_states, self.last_hc_hidden_states = model_output + return hidden_states + return model_output + + +class Qwen4ExpForConditionalGeneration(Qwen3VLForConditionalGeneration): + packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping + hf_to_sglang_mapper = None + + @staticmethod + def shared_experts_fusion_disable_reason(hf_config, quant_config): + return Qwen4ExpVLModel.shared_experts_fusion_disable_reason( + hf_config, quant_config + ) + + def __init__( + self, + config: Qwen4ExpConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + language_model_cls=Qwen4ExpVLModel, + ) -> None: + super().__init__(config, quant_config, prefix, language_model_cls) + rope_config = getattr(self.config, "rope_parameters", None) or getattr( + self.config, "rope_scaling", {} + ) + self.is_mrope_enabled = ( + "mrope_section" in rope_config and not self.language_model_only + ) + self.deepstack_visual_indexes = ( + self.visual.deepstack_visual_indexes if self.visual is not None else [] + ) + + @torch.no_grad() + def forward(self, *args, **kwargs): + output = super().forward(*args, **kwargs) + hc_hidden_states = self.model.last_hc_hidden_states + if hc_hidden_states is not None and isinstance(output, LogitsProcessorOutput): + output.hidden_states = hc_hidden_states + return output + + def _load_qwen4_exp_ple_buffer( + self, + name: str, + loaded_weight: torch.Tensor, + buffers: dict, + loaded_buffers: Set[str], + ) -> bool: + if ".ple.ple_embedding." not in name: + return False + buffer_name = name.rsplit(".", 1)[-1] + if buffer_name.startswith("hashstats_"): + return True + if buffer_name == "token_lookup": + return True + if buffer_name not in { + "layer_multipliers", + "ngram_heads_offsets", + "ngram_heads_vocab_sizes", + "weight_scale", + }: + return False + buffer = buffers.get(name) + if buffer is None: + return False + if buffer.shape != loaded_weight.shape: + raise ValueError( + f"Shape mismatch for {name}: expected {tuple(buffer.shape)}, " + f"got {tuple(loaded_weight.shape)}" + ) + buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) + loaded_buffers.add(name) + return True + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + # Checkpoints use the qwen3.5 head-first in_proj layout, + # matching Qwen3_5GatedDeltaNet's forward, not qwen3-next's group-first. + ("in_proj_qkvz.", "in_proj_qkv.", (0, 1, 2)), + ("in_proj_qkvz.", "in_proj_z.", 3), + ("in_proj_ba.", "in_proj_b.", 0), + ("in_proj_ba.", "in_proj_a.", 1), + ] + + num_experts = getattr(self.config, "num_experts", None) + expert_params_mapping = ( + FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=num_experts, + ) + if num_experts is not None + else [] + ) + fused_expert_params_mapping = [ + ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), + ("experts.w2_weight", "experts.down_proj", 0, "w2"), + ] + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale_inv", + "_weight_scale_inv", + ".input_scale_inv", + "_input_scale_inv", + "_weight_scale", + "_input_scale", + ) + + def load_fused_expert_weights( + name: str, + params_dict: dict, + loaded_weight: torch.Tensor, + shard_id: str, + num_experts: int, + ) -> bool: + if name not in params_dict: + return False + param = params_dict[name] + weight_loader = param.weight_loader + for expert_id in range(num_experts): + weight_loader( + param, + loaded_weight[expert_id], + name, + shard_id, + expert_id, + ) + return True + + def copy_ple_rows_to_tp_embedding( + emb, loaded_weight: torch.Tensor, row_start: int, row_end: int + ) -> None: + tp_start = emb.shard_indices.org_vocab_start_index + tp_end = emb.shard_indices.org_vocab_end_index + ov_start = max(row_start, tp_start) + ov_end = min(row_end, tp_end) + if ov_start < ov_end: + local_start = ov_start - tp_start + src_start = ov_start - row_start + n_rows = ov_end - ov_start + emb.weight.data[local_start : local_start + n_rows].copy_( + loaded_weight[src_start : src_start + n_rows].to( + device=emb.weight.device, dtype=emb.weight.dtype + ) + ) + + def load_qwen4_exp_ple_shard(name: str, loaded_weight: torch.Tensor) -> bool: + if ".ngram_embedding.shard_" not in name: + return False + import re + + match = re.search(r"\.ngram_embedding\.shard_(\d+)\.weight$", name) + if not match: + return False + shard_idx = int(match.group(1)) + mod_prefix = name[: name.index(".ngram_embedding.shard_")] + ple_mod = ple_modules.get(mod_prefix) + if ple_mod is None: + return False + emb = ple_mod.ngram_embedding + if ( + loaded_weight.dtype == torch.float8_e4m3fn + and emb.weight.dtype != torch.float8_e4m3fn + ): + if isinstance(emb, Qwen4ExpPinnedHostEmbedding): + # offload gathers from pinned host memory; a swapped-in + # pageable tensor would fault in the Triton kernel. + raise ValueError( + "fp8 PLE auto-switch is unsupported with " + "ple_offload_embedding; set " + 'text_config.ple_embedding_dtype="float8_e4m3fn" instead' + ) + logger.info( + "PLE embedding switched to fp8 storage: %s (%s)", + mod_prefix, + tuple(emb.weight.data.shape), + ) + old_weight_data = emb.weight.data + # StartupWeightLoadManager enforces tensor identity/dtype; this + # swap breaks that contract if the model is ever enrolled. + emb.weight = torch.nn.Parameter( + torch.empty_like(old_weight_data, dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + del old_weight_data + # params_dict was snapshotted before the loop; drop the stale + # entry or it pins the old bf16 storage until load end. + params_dict.pop(f"{mod_prefix}.ngram_embedding.weight", None) + torch.cuda.empty_cache() + if ( + emb.weight.dtype == torch.float8_e4m3fn + and loaded_weight.dtype != torch.float8_e4m3fn + ): + if not getattr(load_qwen4_exp_ple_shard, "_warned_downcast", False): + load_qwen4_exp_ple_shard._warned_downcast = True + logger.warning( + "PLE checkpoint shards are %s but the embedding storage " + "is fp8 (ple_embedding_dtype / fp8 quant config); " + "downcasting is lossy", + loaded_weight.dtype, + ) + shard_size = ( + emb.org_vocab_size + ple_num_sync_shards - 1 + ) // ple_num_sync_shards + shard_start = shard_idx * shard_size + actual_rows = loaded_weight.shape[0] + shard_end = shard_start + actual_rows + copy_ple_rows_to_tp_embedding(emb, loaded_weight, shard_start, shard_end) + loaded_shard_params.add(f"{mod_prefix}.ngram_embedding.weight") + return True + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + buffers = dict(self.named_buffers()) + + ple_modules = { + mod_name: mod + for mod_name, mod in self.named_modules() + if isinstance(mod, Qwen4ExpNGramEmbedding) + } + text_config = getattr(self.config, "text_config", self.config) + ple_num_sync_shards = int( + getattr( + text_config, + "split_ngram_parts", + getattr(self.config, "split_ngram_parts", 512), + ) + ) + loaded_params: Set[str] = set() + loaded_buffers: Set[str] = set() + loaded_shard_params: Set[str] = set() + skipped_visual_count = 0 + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "mtp" in name: + continue + if "visual" in name and self.language_model_only: + skipped_visual_count += 1 + continue + if "language_model" in name: + name = name.replace("model.language_model.", "model.") + if ".self_attn." in name: + name = name.replace(".self_attn", "") + if name.endswith(".k_proj.k_scale"): + name = name.replace(".k_proj.k_scale", ".attn.k_scale") + elif name.endswith(".v_proj.v_scale"): + name = name.replace(".v_proj.v_scale", ".attn.v_scale") + + if self._load_qwen4_exp_ple_buffer( + name, loaded_weight, buffers, loaded_buffers + ): + continue + if load_qwen4_exp_ple_shard(name, loaded_weight): + continue + if ".ple.ple_embedding.ngram_embedding." in name and name.endswith( + ".weight" + ): + raise ValueError( + f"unsupported PLE weight layout (expected shard_N shards): {name}" + ) + + if ( + self.config.tie_word_embeddings + and self.pp_group.is_last_rank + and "model.embed_tokens.weight" in name + and "lm_head.weight" in params_dict + ): + lm_head_param = params_dict["lm_head.weight"] + weight_loader = getattr( + lm_head_param, "weight_loader", default_weight_loader + ) + weight_loader(lm_head_param, loaded_weight) + + layer_id = get_layer_id(name) + if layer_id is not None and ( + layer_id < self.start_layer or layer_id >= self.end_layer + ): + continue + + is_fused_expert = ( + "experts.gate_up_proj" in name or "experts.down_proj" in name + ) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "visual" in name or "mlp.experts" in name: + continue + mapped_name = name.replace(weight_name, param_name) + if ( + mapped_name.endswith(ignore_suffixes) + and mapped_name not in params_dict + ): + continue + if mapped_name not in params_dict: + continue + param = params_dict[mapped_name] + param.weight_loader(param, loaded_weight, shard_id) + name = mapped_name + break + else: + is_expert_weight = False + current_expert_params_mapping = ( + fused_expert_params_mapping + if is_fused_expert + else expert_params_mapping + ) + for mapping in current_expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + if "visual" in name or self.config.encoder_only: + continue + is_expert_weight = True + mapped_name = name.replace(weight_name, param_name) + if is_fused_expert: + if "experts.gate_up_proj" in name: + gate_weight, up_weight = loaded_weight.chunk(2, dim=-2) + if not load_fused_expert_weights( + mapped_name, + params_dict, + gate_weight, + "w1", + num_experts, + ): + raise KeyError(f"Parameter {mapped_name} not found") + if not load_fused_expert_weights( + mapped_name, + params_dict, + up_weight, + "w3", + num_experts, + ): + raise KeyError(f"Parameter {mapped_name} not found") + else: + if not load_fused_expert_weights( + mapped_name, + params_dict, + loaded_weight, + shard_id, + num_experts, + ): + raise KeyError(f"Parameter {mapped_name} not found") + else: + if ( + mapped_name.endswith(ignore_suffixes) + and mapped_name not in params_dict + ): + continue + param = params_dict[mapped_name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + mapped_name, + shard_id=shard_id, + expert_id=expert_id, + ) + name = mapped_name + break + else: + if is_expert_weight: + continue + if "visual" in name: + name = name.replace("attn.qkv.", "attn.qkv_proj.") + name = name.replace("model.visual.", "visual.") + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name.endswith("_scale") and name not in params_dict: + assert abs(loaded_weight.item() - 1.0) < 1e-6, ( + f"Expected 1.0, got {loaded_weight.item()} in skipped {name}" + ) + continue + if name in params_dict: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning( + "Parameter %s not found while loading Qwen4-Exp VL weights", + name, + ) + continue + loaded_params.add(name) + + loaded_params.update(loaded_buffers) + loaded_params.update(loaded_shard_params) + + if skipped_visual_count > 0: + logger.info( + f"[language_model_only] Qwen4 load_weights: skipped " + f"{skipped_visual_count} visual weights" + ) + + for module in self.modules(): + if isinstance(module, Qwen3_5GatedDeltaNet): + module.finalize_fused_in_proj() + + return loaded_params + + def precompile_kernels_after_loading(self) -> None: + from sglang.srt.layers.quantization.unquant import precompile_splitk_tactics + + if precompile_splitk_tactics(): + logger.info("Precompiled BF16 split-K GEMM tactics for Qwen4-Exp") + + @classmethod + def get_model_config_for_expert_location(cls, config): + text_config = getattr(config, "text_config", config) + if getattr(text_config, "num_experts", None) is None: + return None + return ModelConfigForExpertLocation( + num_layers=text_config.num_hidden_layers, + num_logical_experts=text_config.num_experts, + num_groups=None, + ) + + +EntryClass = [Qwen4ExpForConditionalGeneration] diff --git a/python/sglang/srt/models/qwen4_exp_mtp.py b/python/sglang/srt/models/qwen4_exp_mtp.py new file mode 100644 index 000000000..31a873014 --- /dev/null +++ b/python/sglang/srt/models/qwen4_exp_mtp.py @@ -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] diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py index 44795b027..eba723827 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py @@ -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", diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index 13de33ad6..f93f2f648 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index f6efe923f..cb603f1ff 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 192d57d82..539b96eb0 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 472537b7e..71405f244 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -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() diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index be1419c66..89bc4e2d9 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -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, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 9eb70a498..bbffe4795 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -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) diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 6b3eea06d..d61d2407c 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -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, diff --git a/python/sglang/test/run_eval.py b/python/sglang/test/run_eval.py index cf7d9a0f7..7b6a4117d 100644 --- a/python/sglang/test/run_eval.py +++ b/python/sglang/test/run_eval.py @@ -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 diff --git a/test/registered/e2e/models/test_qwen4_exp_models.py b/test/registered/e2e/models/test_qwen4_exp_models.py new file mode 100644 index 000000000..3627ee569 --- /dev/null +++ b/test/registered/e2e/models/test_qwen4_exp_models.py @@ -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() diff --git a/test/registered/kernel/embeddings/test_qwen4_ple_offload.py b/test/registered/kernel/embeddings/test_qwen4_ple_offload.py new file mode 100644 index 000000000..015d5fdd4 --- /dev/null +++ b/test/registered/kernel/embeddings/test_qwen4_ple_offload.py @@ -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"])) diff --git a/test/registered/kernel/hyperconnection/test_hc_mix_triton.py b/test/registered/kernel/hyperconnection/test_hc_mix_triton.py new file mode 100644 index 000000000..f9b8159af --- /dev/null +++ b/test/registered/kernel/hyperconnection/test_hc_mix_triton.py @@ -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"])) diff --git a/test/registered/kernel/jit/test_fast_topk.py b/test/registered/kernel/jit/test_fast_topk.py new file mode 100644 index 000000000..c203acb49 --- /dev/null +++ b/test/registered/kernel/jit/test_fast_topk.py @@ -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"])) diff --git a/test/registered/kernel/jit/test_grouped_gemma_rmsnorm.py b/test/registered/kernel/jit/test_grouped_gemma_rmsnorm.py new file mode 100644 index 000000000..a6562cee2 --- /dev/null +++ b/test/registered/kernel/jit/test_grouped_gemma_rmsnorm.py @@ -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"])) diff --git a/test/registered/kernel/jit/test_hc_combine.py b/test/registered/kernel/jit/test_hc_combine.py new file mode 100644 index 000000000..10bc98113 --- /dev/null +++ b/test/registered/kernel/jit/test_hc_combine.py @@ -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"])) diff --git a/test/registered/kernel/qsa/test_qsa.py b/test/registered/kernel/qsa/test_qsa.py new file mode 100644 index 000000000..dbd3cd237 --- /dev/null +++ b/test/registered/kernel/qsa/test_qsa.py @@ -0,0 +1,1528 @@ +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from sglang.kernels.ops.attention import qwen38_qsa_sm121_varlen +from sglang.srt.configs.qwen4_exp import Qwen4ExpConfig +from sglang.srt.layers.attention import qwen_sparse_attn_backend as qsa_backend_module +from sglang.srt.layers.attention.qsa import dsa_indexer as dsa_indexer_module +from sglang.srt.layers.attention.qsa import qsa_indexer as qsa_indexer_module +from sglang.srt.layers.attention.qsa.kernel import ( + expand_qsa_block_indices, + qsa_fast_topk, + qsa_sparse_attention, + torch_expand_qsa_block_indices, + triton_expand_qsa_block_indices, +) +from sglang.srt.layers.attention.qsa.metadata import ( + QSAIndexerMetadata, + build_qsa_row_ranges, +) +from sglang.srt.layers.attention.qsa.mqa import ( + qsa_mqa_decode, + qsa_mqa_prefill, +) +from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer +from sglang.srt.layers.attention.qsa.sparse_attn import ( + qwen_sparse_fa2_cu_seqlens_triton, + qwen_sparse_kv_extraction_compact_triton, +) +from sglang.srt.layers.attention.qwen_sparse_attn_backend import ( + QwenSparseAttnBackend, + QwenSparseMultiStepDraftBackend, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +COMPRESS_RATIO = 4 +TOKEN_TOPK = 2048 +BLOCK_TOPK = TOKEN_TOPK // COMPRESS_RATIO +FINAL_TOPK = TOKEN_TOPK + COMPRESS_RATIO - 1 + + +@pytest.mark.parametrize( + ("capability", "expected"), + [((12, 0), True), ((12, 1), False), ((10, 0), False)], +) +def test_is_sm120_matches_exact_capability(monkeypatch, capability, expected): + from sglang.srt.utils import common + + common.is_sm120.cache_clear() + monkeypatch.setattr(common, "is_cuda", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability) + + try: + assert common.is_sm120() is expected + finally: + common.is_sm120.cache_clear() + + +@pytest.mark.parametrize( + ("capability", "expected"), + [((12, 1), True), ((12, 0), False), ((10, 0), False)], +) +def test_is_sm121_matches_exact_capability(monkeypatch, capability, expected): + from sglang.srt.utils import common + + common.is_sm121.cache_clear() + monkeypatch.setattr(common, "is_cuda", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability) + + try: + assert common.is_sm121() is expected + finally: + common.is_sm121.cache_clear() + + +@pytest.mark.parametrize( + ("sm100", "sm120", "expected_enabled"), + [(False, True, True), (True, False, True), (False, False, False)], + ids=["sm120", "sm100", "other-sm12x"], +) +def test_qsa_trtllm_sparse_decode_arch_gate( + monkeypatch, sm100, sm120, expected_enabled +): + resolver = qsa_backend_module._resolve_trtllm_sparse_decode + resolver.cache_clear() + + trtllm_decode_func = object() + flashinfer_decode = ModuleType("flashinfer.decode") + flashinfer_decode.trtllm_batch_decode_with_kv_cache = trtllm_decode_func + + monkeypatch.setattr("sglang.srt.utils.is_sm100_supported", lambda: sm100) + monkeypatch.setattr("sglang.srt.utils.is_sm120", lambda: sm120) + monkeypatch.setitem(sys.modules, flashinfer_decode.__name__, flashinfer_decode) + + try: + expected = trtllm_decode_func if expected_enabled else None + assert resolver() is expected + finally: + resolver.cache_clear() + + +def test_qsa_sm121_resolves_kda_varlen_kernel(monkeypatch): + resolver = qsa_backend_module._resolve_flash_attn_varlen_func + resolver.cache_clear() + monkeypatch.setattr("sglang.srt.utils.is_sm121", lambda: True) + + try: + assert resolver() is qwen38_qsa_sm121_varlen + finally: + resolver.cache_clear() + + +def test_qsa_sm121_compaction_and_attention_match_sparse_reference(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 1): + pytest.skip("SM121-only kernel") + torch.manual_seed(2028) + device = torch.device("cuda") + batch, topk = 3, FINAL_TOPK + num_q_heads, num_kv_heads, head_dim = 24, 2, 256 + sequence_lengths = torch.tensor([17, 1050, 4096], dtype=torch.int32) + valid_counts_cpu = [17, 911, topk] + max_sequence_length = int(sequence_lengths.max()) + + req_to_token = torch.arange( + batch * max_sequence_length, dtype=torch.int32, device=device + ).reshape(batch, max_sequence_length) + indices = torch.full((batch, topk), -1, dtype=torch.int32, device=device) + for row, valid_count in enumerate(valid_counts_cpu): + indices[row, :valid_count] = torch.randperm( + int(sequence_lengths[row]), device=device, dtype=torch.int64 + )[:valid_count].to(torch.int32) + + slots = torch.full_like(indices, -1) + for row, valid_count in enumerate(valid_counts_cpu): + slots[row, :valid_count] = req_to_token[row, indices[row, :valid_count].long()] + + pool_size = batch * max_sequence_length + q = torch.randn(batch, num_q_heads, head_dim, dtype=torch.bfloat16, device=device) + k_cache = torch.randn( + pool_size, num_kv_heads, head_dim, dtype=torch.bfloat16, device=device + ) + v_cache = torch.randn_like(k_cache) + valid_counts = torch.empty(batch, dtype=torch.int32, device=device) + cu_k = torch.empty(batch + 1, dtype=torch.int32, device=device) + cu_q = torch.arange(batch + 1, dtype=torch.int32, device=device) + qwen_sparse_fa2_cu_seqlens_triton( + sequence_lengths.to(device), + indices, + valid_counts, + cu_k, + batch, + topk, + ) + packed_k = torch.empty( + batch * topk, + num_kv_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + packed_v = torch.empty_like(packed_k) + qwen_sparse_kv_extraction_compact_triton( + k_cache, + v_cache, + req_to_token, + torch.arange(batch, dtype=torch.int32, device=device), + indices, + sequence_lengths.to(device), + cu_k, + packed_k, + packed_v, + batch, + topk, + ) + + scale = head_dim**-0.5 + actual = qwen38_qsa_sm121_varlen( + q, + packed_k, + packed_v, + cu_q, + cu_k, + max_seqlen_k=topk, + softmax_scale=scale, + ) + expected = qsa_sparse_attention(q, k_cache, v_cache, slots, scale) + assert valid_counts.tolist() == valid_counts_cpu + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +def _compressed_config_namespace(**overrides): + fields = dict( + model_type="qwen4_exp", + indexer_n_heads=8, + indexer_kv_heads=1, + indexer_head_dim=128, + indexer_budget=TOKEN_TOPK, + indexer_compress_ratio=COMPRESS_RATIO, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _tokenwise_config_namespace(**overrides): + fields = dict( + model_type="qwen3_5", + index_topk=2048, + index_n_heads=64, + index_kv_heads=1, + index_head_dim=128, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test_qsa_profile_parses_compressed_qwen4_exp_schema(): + from sglang.srt.layers.attention.qsa.config import ( + QSA_ROPE_MROPE, + QSA_VARIANT_COMPRESSED, + is_qwen_qsa, + parse_qsa_profile, + ) + + wrapped = Qwen4ExpConfig( + text_config={ + "hc_count": 4, + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 1, + "indexer_n_heads": 8, + "indexer_kv_heads": 1, + "indexer_head_dim": 128, + "indexer_budget": TOKEN_TOPK, + "indexer_compress_ratio": COMPRESS_RATIO, + }, + ) + for config in (wrapped, _compressed_config_namespace()): + profile = parse_qsa_profile(config) + assert profile.variant == QSA_VARIANT_COMPRESSED + assert profile.n_heads == 8 + assert profile.kv_heads == 1 + assert profile.head_dim == 128 + assert profile.budget == TOKEN_TOPK + assert profile.compress_ratio == COMPRESS_RATIO + assert profile.block_topk == BLOCK_TOPK + assert profile.rope_mode == QSA_ROPE_MROPE + assert is_qwen_qsa(config) + # The legacy backend module keeps re-exporting the shared detector. + assert qsa_backend_module.is_qwen_qsa is is_qwen_qsa + assert not is_qwen_qsa(SimpleNamespace()) + assert parse_qsa_profile(None) is None + + +def test_qsa_profile_rejects_malformed_compressed_schema(): + from sglang.srt.layers.attention.qsa.config import parse_qsa_profile + + bad_configs = { + "missing": _compressed_config_namespace(indexer_budget=None), + "ratio_one": _compressed_config_namespace(indexer_compress_ratio=1), + "indivisible": _compressed_config_namespace(indexer_budget=TOKEN_TOPK - 2), + "bad_block_topk": _compressed_config_namespace( + indexer_budget=1024, indexer_compress_ratio=4 + ), + "kv_heads": _compressed_config_namespace(indexer_kv_heads=2), + } + for name, config in bad_configs.items(): + try: + parse_qsa_profile(config) + except ValueError: + continue + raise AssertionError(f"{name} compressed config must be rejected") + + +def test_qsa_glue_builds_indexer_per_variant(monkeypatch): + from sglang.srt.layers.attention.qsa.glue import build_qsa_indexer + + recorded = {} + + class _FakeIndexer: + def __init__( + self, config, layer_id, quant_config=None, prefix="", rotary_emb=None + ): + recorded.update( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=prefix, + rotary_emb=rotary_emb, + ) + + monkeypatch.setattr(qsa_indexer_module, "QSAIndexer", _FakeIndexer) + rotary = object() + config = _compressed_config_namespace() + indexer = build_qsa_indexer( + config, layer_id=7, quant_config="qc", prefix="p", rotary_emb=rotary + ) + assert isinstance(indexer, _FakeIndexer) + assert recorded == dict( + config=config, layer_id=7, quant_config="qc", prefix="p", rotary_emb=rotary + ) + + # Tokenwise configs build the Lightning Indexer through the same glue. + class _FakeDSAIndexer: + def __init__(self, config, layer_id, quant_config=None, prefix="", **kw): + recorded.update( + dsa_config=config, + dsa_layer_id=layer_id, + dsa_quant_config=quant_config, + dsa_prefix=prefix, + ) + + monkeypatch.setattr(dsa_indexer_module, "QwenDSAIndexer", _FakeDSAIndexer) + dsa_indexer = build_qsa_indexer( + _tokenwise_config_namespace(), layer_id=2, prefix="q" + ) + assert isinstance(dsa_indexer, _FakeDSAIndexer) + assert recorded["dsa_layer_id"] == 2 + assert recorded["dsa_prefix"] == "q" + + try: + build_qsa_indexer(SimpleNamespace(), layer_id=0, rotary_emb=rotary) + except ValueError as exc: + assert "QSA indexer schema" in str(exc) + else: + raise AssertionError("non-QSA configs must be rejected") + + +def test_qsa_glue_fetches_indexer_metadata_without_model_unwrap(): + from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( + HybridLinearAttnBackend, + ) + from sglang.srt.layers.attention.qsa.glue import get_qsa_indexer_metadata + + full = SimpleNamespace( + get_indexer_metadata=lambda layer_id, forward_batch: f"meta-{layer_id}", + token_to_kv_pool=None, + req_to_token_pool=None, + kv_index_translator=None, + needs_cpu_seq_lens=True, + ) + hybrid = HybridLinearAttnBackend( + full_attn_backend=full, + linear_attn_backend=SimpleNamespace(needs_cpu_seq_lens=True), + full_attn_layers=[3], + ) + # Full-attention layers fetch through the hybrid wrapper directly. + assert get_qsa_indexer_metadata(hybrid, 3, object()) == "meta-3" + # A backend that provides no indexer metadata anywhere must surface an + # explicit error instead of silently running without sparse selection. + empty = SimpleNamespace(get_indexer_metadata=lambda layer_id, batch: None) + try: + get_qsa_indexer_metadata(empty, 3, object()) + except RuntimeError as exc: + assert "indexer metadata" in str(exc) + else: + raise AssertionError("QSA must fail when no indexer metadata exists") + + +def test_qsa_draft_extend_backend_decision_follows_profile(): + from sglang.srt.layers.attention.qsa.config import parse_qsa_profile + from sglang.srt.speculative.draft_utils import DraftBackendFactory + + def factory(config): + runner = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), + draft_attention_backend=None, + ) + return DraftBackendFactory( + draft_model_runner=runner, + topk=1, + speculative_num_steps=3, + qsa_profile=parse_qsa_profile(config), + ) + + compressed = factory(_compressed_config_namespace()) + backend = compressed.create_draft_extend_backend() + assert isinstance(backend, QwenSparseAttnBackend) + assert backend.runner is compressed.draft_model_runner + assert backend.decode_attention_backend_str == "qsa" + # Tokenwise profiles stay eager (no graph-stable indexer metadata); they + # must never silently fall back to a dense backend either. + assert factory(_tokenwise_config_namespace()).create_draft_extend_backend() is None + + +def _make_mtp_draft_batch(steps: int, seq_lens=(8, 16), loc_base: int = 40): + bs = len(seq_lens) + pool = _FakeQSAPool(capacity=bs * steps + 256) + req_to_token = torch.stack( + [torch.arange(i * 64, (i + 1) * 64, dtype=torch.int32) for i in range(bs)] + ) + runner = SimpleNamespace( + device="cpu", + token_to_kv_pool=pool, + req_to_token_pool=SimpleNamespace(req_to_token=req_to_token), + model_config=SimpleNamespace( + context_len=64, + hf_config=SimpleNamespace(indexer_compress_ratio=COMPRESS_RATIO), + ), + ) + backend = QwenSparseMultiStepDraftBackend( + runner, topk=1, speculative_num_steps=steps + ) + full_out_cache_loc = loc_base + torch.arange(bs * steps, dtype=torch.int32) + forward_batch = SimpleNamespace( + token_to_kv_pool=pool, + req_to_token_pool=runner.req_to_token_pool, + batch_size=bs, + req_pool_indices=torch.arange(bs, dtype=torch.int32), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + seq_lens_cpu=torch.tensor(seq_lens, dtype=torch.int32), + positions=torch.tensor([length - 1 for length in seq_lens], dtype=torch.int64), + out_cache_loc=full_out_cache_loc, + forward_mode=ForwardMode.DECODE, + ) + return backend, forward_batch, pool + + +def test_qsa_cuda_graph_pads_dynamic_draft_extend_rows(): + row_lengths, row_req_pool_indices, row_prefix_lengths = ( + QwenSparseAttnBackend._graph_speculative_layout( + bs=2, + num_tokens=8, + req_pool_indices=torch.tensor([3, 7], dtype=torch.int32), + seq_lens_cpu=torch.tensor([11, 21], dtype=torch.int32), + forward_mode=ForwardMode.DRAFT_EXTEND_V2, + spec_info=SimpleNamespace(extend_seq_lens_cpu=[1, 2]), + ) + ) + assert row_lengths.tolist() == [11, 20, 21, 1, 1, 1, 1, 1] + assert row_req_pool_indices.tolist() == [3, 7, 7, 3, 3, 3, 3, 3] + assert row_prefix_lengths.tolist() == [10, 19, 19, 0, 0, 0, 0, 0] + + +def test_qsa_cuda_graph_target_verify_ignores_capture_bucket_requests(): + row_lengths, row_req_pool_indices, row_prefix_lengths = ( + QwenSparseAttnBackend._graph_speculative_layout( + bs=2, + num_tokens=8, + req_pool_indices=torch.tensor([3, 0], dtype=torch.int32), + seq_lens_cpu=torch.tensor([9, 1], dtype=torch.int32), + forward_mode=ForwardMode.TARGET_VERIFY, + spec_info=SimpleNamespace(draft_token_num=4), + num_padding=1, + ) + ) + assert row_lengths.tolist() == [10, 11, 12, 13, 1, 1, 1, 1] + assert row_req_pool_indices.tolist() == [3] * 8 + assert row_prefix_lengths.tolist() == [9, 9, 9, 9, 0, 0, 0, 0] + + +def test_qsa_target_verify_rejects_branching_speculation(): + backend = QwenSparseAttnBackend.__new__(QwenSparseAttnBackend) + backend.compress_ratio = 4 + try: + backend._require_chain_speculation( + ForwardMode.TARGET_VERIFY, SimpleNamespace(topk=2) + ) + except NotImplementedError as exc: + assert "topk=1" in str(exc) + else: + raise AssertionError("QSA target verification must reject tree branches") + # The pending-group ring keys state by position % ratio: a verify window + # wider than the ratio would collide within one forward. + try: + backend._require_chain_speculation( + ForwardMode.TARGET_VERIFY, SimpleNamespace(topk=1, draft_token_num=5) + ) + except NotImplementedError as exc: + assert "compress ratio" in str(exc) + else: + raise AssertionError("QSA must reject draft windows wider than the ratio") + backend._require_chain_speculation( + ForwardMode.TARGET_VERIFY, SimpleNamespace(topk=1, draft_token_num=4) + ) + + +def test_qsa_mtp_cuda_graph_padding_stays_below_compression_boundary(): + backend, forward_batch, _ = _make_mtp_draft_batch(steps=4, seq_lens=(3, 1)) + + class Recorder: + def __init__(self): + self.capture_lengths = None + self.replay_lengths = None + + def _capture_cuda_graph_metadata(self, **kwargs): + self.capture_lengths = kwargs["seq_lens"].clone() + + def _replay_cuda_graph_metadata(self, *, bs, seq_lens, **kwargs): + self.replay_lengths = seq_lens[:bs].clone() + + recorders = [Recorder() for _ in backend.attn_backends] + backend.attn_backends = recorders + forward_batch.spec_info = SimpleNamespace() + + backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) + forward_batch.num_padding = 1 + backend.init_forward_metadata_out_graph(forward_batch) + + assert [r.capture_lengths.tolist() for r in recorders] == [[1, 1]] * 3 + assert [r.replay_lengths.tolist() for r in recorders] == [ + [4, 1], + [5, 1], + [6, 1], + ] + + +def test_qsa_indexer_ignores_dp_attention_token_padding(): + calls = {} + + def run(num_rows): + mapping = torch.zeros(15, dtype=torch.int32) + metadata = SimpleNamespace( + token_to_kv_pool=None, + compress_member_rows=None, + decode_logical_positions=None, + pending_ring_slots=None, + get_token_to_batch_idx=lambda: mapping, + get_prefill_mqa_inputs=lambda layer_id, logical_positions: ( + torch.empty(0, 1, 128), + torch.zeros(15, dtype=torch.int32), + torch.zeros(15, dtype=torch.int32), + torch.tensor([15], dtype=torch.int32), + ), + ) + indexer = SimpleNamespace( + layer_id=0, + project_qk=lambda hidden, rope_positions, **kwargs: ( + hidden.reshape(-1, 1, 2), + hidden.reshape(-1, 1, 2), + False, + ), + _pending_ring_slots=lambda metadata, logical_positions, is_extend: ( + torch.zeros(logical_positions.numel(), dtype=torch.long) + ), + update_key_state_and_compress=lambda token_k, logical, rope, meta, state_slots=None, state_stored=False: ( + calls.update( + token_rows=token_k.shape[0], + logical_rows=logical.numel(), + rope_rows=rope.numel(), + ) + ), + select_prefill_tokens=lambda q, keys, starts, ends, logical, lengths: ( + q.squeeze(1) + ), + ) + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + positions=torch.cat( + [torch.arange(15), torch.zeros(num_rows - 15, dtype=torch.long)] + ), + out_cache_loc=torch.arange(num_rows, dtype=torch.int32), + ) + hidden = torch.arange(num_rows * 2, dtype=torch.float32).reshape(num_rows, 2) + return QSAIndexer.forward_cuda( + indexer, hidden, forward_batch.positions, forward_batch, metadata + ) + + unpadded = run(15) + padded = run(16) + torch.testing.assert_close(padded, unpadded) + assert padded.shape[0] == 15 + assert calls == {"token_rows": 15, "logical_rows": 15, "rope_rows": 15} + + +def test_qsa_cuda_extend_ignores_dp_attention_padding(monkeypatch): + if not torch.cuda.is_available(): + return + + kernel_shapes = [] + + def fake_sparse_gqa(q, k, v, max_seqlen_k, indices, cu_seqlens, scale): + kernel_shapes.append((q.shape[0], k.shape[0], v.shape[0], indices.shape[0])) + return q + 1 + + monkeypatch.setattr( + qsa_backend_module, "sparse_gqa_fwd_interface_triton", fake_sparse_gqa + ) + backend = QwenSparseAttnBackend.__new__(QwenSparseAttnBackend) + + class Pool: + def set_kv_buffer(self, layer, loc, k, v): + pass + + pool = Pool() + backend.token_to_kv_pool = pool + layer = SimpleNamespace(tp_q_head_num=1, head_dim=2, layer_id=0, scaling=1.0) + topk = torch.zeros(15, 3, dtype=torch.int32, device="cuda") + + def run(num_rows): + values = torch.arange(num_rows * 2, dtype=torch.float32, device="cuda").reshape( + num_rows, 2 + ) + batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + token_to_kv_pool=pool, + out_cache_loc=torch.arange(num_rows, dtype=torch.int32, device="cuda"), + extend_seq_lens=torch.tensor([15], dtype=torch.int32, device="cuda"), + extend_seq_lens_cpu=[15], + seq_lens_cpu=[15], + ) + return backend.forward_extend( + values, values.clone(), values.clone(), layer, batch, topk_indices=topk + ) + + unpadded = run(15) + padded = run(16) + torch.testing.assert_close(padded[:15], unpadded) + torch.testing.assert_close(padded[15], torch.zeros(2, device="cuda")) + assert kernel_shapes == [(15, 15, 15, 15), (15, 15, 15, 15)] + + +def _make_paged_extend_backend(): + metadata = SimpleNamespace( + token_to_batch_idx=torch.zeros(3, dtype=torch.int32), + sequence_lengths=torch.tensor([8], dtype=torch.int32), + token_slot_table=torch.arange(16, dtype=torch.int32).reshape(1, 16), + ) + backend = QwenSparseAttnBackend.__new__(QwenSparseAttnBackend) + backend.forward_metadata = metadata + + class Pool: + def set_kv_buffer(self, layer, loc, k, v): + pass + + def get_key_buffer(self, layer_id): + return torch.zeros(16, 1, 2) + + def get_value_buffer(self, layer_id): + return torch.zeros(16, 1, 2) + + layer = SimpleNamespace(tp_q_head_num=1, head_dim=2, layer_id=0, scaling=1.0) + pool = Pool() + backend.token_to_kv_pool = pool + return backend, pool, layer + + +def test_qsa_paged_extend_trims_padding_rows_and_restores_output(monkeypatch): + kernel_rows = [] + + def fake_sparse_attention(q, k, v, slots, scale): + kernel_rows.append((q.shape[0], slots.shape[0])) + return q + 1 + + monkeypatch.setattr( + qsa_backend_module, "qsa_sparse_attention", fake_sparse_attention + ) + backend, pool, layer = _make_paged_extend_backend() + topk = torch.zeros(3, 2, dtype=torch.int32) + + def run(num_rows, mode): + values = torch.arange(num_rows * 2, dtype=torch.float32).reshape(num_rows, 2) + batch = SimpleNamespace( + forward_mode=mode, + token_to_kv_pool=pool, + out_cache_loc=torch.arange(num_rows, dtype=torch.int32), + ) + return backend.forward_extend( + values, values.clone(), values.clone(), layer, batch, topk_indices=topk + ) + + # DP attention pads the physical q rows past the semantic draft rows in + # every speculative paged mode; padding must round-trip through the kernel. + for mode in (ForwardMode.TARGET_VERIFY, ForwardMode.DRAFT_EXTEND_V2): + unpadded = run(3, mode) + padded = run(5, mode) + torch.testing.assert_close(padded[:3], unpadded) + torch.testing.assert_close(padded[3:], torch.zeros(2, 2)) + assert padded.shape == (5, 2) + # The reference path (and any kernel behind it) only ever saw valid rows. + assert kernel_rows == [(3, 3)] * 4 + + # More semantic rows than physical q rows is a bug and must not silently + # truncate the top-k table. + short_q = torch.zeros(2, 2, dtype=torch.float32) + batch = SimpleNamespace( + forward_mode=ForwardMode.TARGET_VERIFY, + token_to_kv_pool=pool, + out_cache_loc=torch.arange(2, dtype=torch.int32), + ) + try: + backend.forward_extend( + short_q, short_q.clone(), short_q.clone(), layer, batch, topk_indices=topk + ) + except ValueError as exc: + assert "top-k rows exceed query rows" in str(exc) + else: + raise AssertionError("QSA must reject top-k rows beyond the query rows") + + +def test_qsa_indexer_rejects_shorter_source_than_request_mapping(): + metadata = SimpleNamespace( + get_token_to_batch_idx=lambda: torch.zeros(16, dtype=torch.int32) + ) + indexer = SimpleNamespace() + batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND, positions=torch.arange(15)) + try: + QSAIndexer.forward_cuda( + indexer, torch.zeros(15, 2), batch.positions, batch, metadata + ) + except ValueError as exc: + assert "logical positions are shorter" in str(exc) + else: + raise AssertionError("QSA must reject a request mapping longer than positions") + + +def _make_qsa_runner_and_pool(num_reqs=4): + pool = _FakeQSAPool(capacity=256) + req_to_token = torch.stack( + [torch.arange(i * 64, (i + 1) * 64, dtype=torch.int32) for i in range(num_reqs)] + ) + runner = SimpleNamespace( + device="cpu", + token_to_kv_pool=pool, + req_to_token_pool=SimpleNamespace(req_to_token=req_to_token), + model_config=SimpleNamespace( + context_len=64, + hf_config=SimpleNamespace(indexer_compress_ratio=COMPRESS_RATIO), + ), + ) + return runner, pool, runner.req_to_token_pool + + +def test_qsa_idle_metadata_builds_empty_rows(): + runner, pool, req_pool = _make_qsa_runner_and_pool() + idle_batch = SimpleNamespace( + token_to_kv_pool=pool, + req_to_token_pool=req_pool, + req_pool_indices=torch.empty(0, dtype=torch.int32), + seq_lens=torch.empty(0, dtype=torch.int32), + seq_lens_cpu=torch.empty(0, dtype=torch.int64), + positions=torch.empty(0, dtype=torch.int64), + out_cache_loc=torch.empty(0, dtype=torch.int32), + # The MTP wrapper forwards idle batches as zero-row DECODE steps. + forward_mode=ForwardMode.DECODE, + ) + backend = QwenSparseAttnBackend(runner) + backend.init_forward_metadata(idle_batch) + metadata = backend.forward_metadata + assert metadata.sequence_lengths.numel() == 0 + assert metadata.token_to_batch_idx.numel() == 0 + assert metadata.row_req_pool_indices.numel() == 0 + assert metadata.token_slot_table.shape[0] == 0 + assert metadata.indexer_metadata.out_cache_loc.numel() == 0 + + # The MTP multi-step wrapper forwards idle batches as zero-row DECODE + # steps; those must not fall into the empty extend/max() path either. + draft = QwenSparseMultiStepDraftBackend(runner, topk=1, speculative_num_steps=2) + draft.init_forward_metadata(idle_batch) + step_metadata = draft.attn_backends[0].forward_metadata + assert step_metadata.sequence_lengths.numel() == 0 + assert step_metadata.row_req_pool_indices.numel() == 0 + # Per-step out_cache_loc slicing must stay empty without allocating rows. + for attn_backend in draft.attn_backends: + assert attn_backend.forward_metadata.indexer_metadata.out_cache_loc.numel() == 0 + + +def test_qsa_decode_requires_one_query_row_per_request(): + runner, pool, req_pool = _make_qsa_runner_and_pool() + backend = QwenSparseAttnBackend(runner) + forward_batch = SimpleNamespace( + token_to_kv_pool=pool, + req_to_token_pool=req_pool, + req_pool_indices=torch.tensor([1, 3], dtype=torch.int32), + seq_lens=torch.tensor([8, 16], dtype=torch.int32), + seq_lens_cpu=torch.tensor([8, 16], dtype=torch.int32), + # 2 request rows but 3 query rows: this layout is ambiguous under DP. + positions=torch.tensor([7, 15, 0], dtype=torch.int64), + out_cache_loc=torch.arange(3, dtype=torch.int32), + forward_mode=ForwardMode.DECODE, + ) + try: + backend.init_forward_metadata(forward_batch) + except ValueError as exc: + assert "exactly one query row per request" in str(exc) + else: + raise AssertionError("QSA decode must reject multi-row requests") + + +def test_qsa_extend_rope_matrix_uses_mrope_coordinates(): + """extend_rope_matrix must be built from forward_batch.mrope_positions ([3, N]) + when set, so prefill- and decode-compressed keys RoPE with the same coordinates.""" + runner, pool, req_pool = _make_qsa_runner_and_pool() + backend = QwenSparseAttnBackend(runner) + num_tokens = 8 + flat = torch.arange(num_tokens, dtype=torch.int64) + mrope = torch.stack([flat, flat + 100, flat + 200]) + forward_batch = SimpleNamespace( + token_to_kv_pool=pool, + req_to_token_pool=req_pool, + req_pool_indices=torch.tensor([1], dtype=torch.int32), + seq_lens=torch.tensor([num_tokens], dtype=torch.int32), + seq_lens_cpu=torch.tensor([num_tokens], dtype=torch.int32), + forward_mode=ForwardMode.EXTEND, + extend_seq_lens=torch.tensor([num_tokens], dtype=torch.int32), + extend_prefix_lens=torch.zeros(1, dtype=torch.int32), + positions=flat, + mrope_positions=mrope, + input_ids=torch.zeros(num_tokens, dtype=torch.int32), + out_cache_loc=torch.arange(num_tokens, dtype=torch.int32), + _original_forward_mode=None, + ) + backend.init_forward_metadata(forward_batch) + got = backend.forward_metadata.indexer_metadata.extend_rope_matrix + assert got is not None + assert torch.equal(got, mrope.transpose(0, 1)) + + +def test_qsa_speculative_pseudo_extend_is_rejected(): + runner, pool, req_pool = _make_qsa_runner_and_pool() + backend = QwenSparseAttnBackend(runner) + for original_mode in (ForwardMode.TARGET_VERIFY, ForwardMode.DRAFT_EXTEND_V2): + forward_batch = SimpleNamespace( + token_to_kv_pool=pool, + req_to_token_pool=req_pool, + req_pool_indices=torch.tensor([1], dtype=torch.int32), + seq_lens=torch.tensor([10], dtype=torch.int32), + seq_lens_cpu=torch.tensor([10], dtype=torch.int32), + # DP MAX_LEN pseudo-extend rewrites the mode and loses the + # per-request draft fan-out of the original speculative mode. + forward_mode=ForwardMode.EXTEND, + extend_seq_lens=torch.ones(1, dtype=torch.int32), + positions=torch.tensor([9], dtype=torch.int64), + out_cache_loc=torch.arange(1, dtype=torch.int32), + _original_forward_mode=original_mode, + ) + try: + backend.init_forward_metadata(forward_batch) + except ValueError as exc: + assert "pseudo-extend" in str(exc) + else: + raise AssertionError(f"QSA must reject pseudo-extend of {original_mode}") + + +class _FakeQSAPool: + qsa_index_kv_heads = 1 + qsa_index_head_dim = 128 + qsa_compressed_page_size = 64 + qsa_compress_ratio = COMPRESS_RATIO + qsa_token_topk = TOKEN_TOPK + qsa_block_topk = BLOCK_TOPK + + def __init__(self, capacity=32): + self.token_k = torch.zeros(capacity, 1, 128, dtype=torch.bfloat16) + self.compressed_k = torch.zeros(3 * 64, 1, 128, dtype=torch.bfloat16) + self.mapping = torch.full((capacity,), -1, dtype=torch.int32) + self.qsa_index_loc_map = self.mapping + self.qsa_page_loc_map = torch.full((capacity,), -1, dtype=torch.int32) + self.qsa_free_index_slots = torch.empty(0, dtype=torch.int32) + self.rope_positions = torch.zeros(capacity, 3, dtype=torch.int64) + self.qsa_page_table = None + self.qsa_page_owned = None + self.next_page = 1 + self.bulk_alloc_calls = 0 + + def set_qsa_key_state_buffer(self, layer_id, loc, token_k): + self.token_k[loc.long()] = token_k + + def get_qsa_key_state_buffer(self, layer_id): + return self.token_k + + def clear_qsa_key_state_buffer(self, layer_id, loc): + self.token_k[loc.long()] = 0 + + def set_qsa_rope_position_buffer(self, loc, positions): + positions = positions.long() + if positions.ndim == 1: + positions = positions.unsqueeze(0).expand(3, -1) + self.rope_positions[loc.long()] = positions.transpose(0, 1) + + def get_qsa_rope_position_buffer(self, loc): + return self.rope_positions[loc.long()] + + def get_or_alloc_qsa_index_locs(self, full_locs, page_anchor_locs, page_offsets): + result = torch.empty_like(full_locs, dtype=torch.int32) + for index, (full_loc, anchor, offset) in enumerate( + zip(full_locs.long(), page_anchor_locs.long(), page_offsets.int()) + ): + if self.mapping[anchor] < 0: + self.mapping[anchor] = self.next_page * 64 + self.next_page += 1 + result[index] = self.mapping[anchor] + offset + self.mapping[full_loc] = result[index] + return result + + def alloc_qsa_compressed_pages(self, num_pages): + self.bulk_alloc_calls += 1 + pages = torch.arange( + self.next_page, self.next_page + num_pages, dtype=torch.int32 + ) + self.next_page += num_pages + return pages + + def copy_qsa_compressed_page_prefixes( + self, source_pages, destination_pages, copy_blocks + ): + for source, destination, blocks in zip( + source_pages.tolist(), destination_pages.tolist(), copy_blocks.tolist() + ): + src = source * 64 + dst = destination * 64 + self.compressed_k[dst : dst + blocks] = self.compressed_k[ + src : src + blocks + ] + + def alloc_qsa_compressed_page(self, source_page=-1, copy_blocks=0): + page = int(self.alloc_qsa_compressed_pages(1)[0]) + if source_page >= 0 and copy_blocks > 0: + self.copy_qsa_compressed_page_prefixes( + torch.tensor([source_page]), + torch.tensor([page]), + torch.tensor([copy_blocks]), + ) + return page + + def set_qsa_index_locs(self, full_locs, page, page_offsets): + result = page * 64 + page_offsets.int() + self.mapping[full_locs.long()] = result + return result + + def get_qsa_index_locs(self, full_locs): + result = self.mapping[full_locs.long()] + assert torch.all(result >= 0) + return result + + def get_qsa_compressed_page_locs(self, page_anchor_locs): + return self.get_qsa_index_locs(page_anchor_locs) // 64 + + def set_qsa_compressed_k_buffer(self, layer_id, loc, compressed_k): + self.compressed_k[loc.long()] = compressed_k + + def get_qsa_compressed_k_buffer(self, layer_id): + return self.compressed_k + + +class _DispatchIndexer: + layer_id = 3 + index_n_heads = 4 + compress_ratio = 4 + _pending_ring_slots = QSAIndexer._pending_ring_slots + + def __init__(self): + self.selected = None + self.logical_positions = None + + @staticmethod + def project_qk(hidden_states, positions, **kwargs): + rows = hidden_states.shape[0] + return torch.zeros(rows, 4, 128), torch.zeros(rows, 1, 128), False + + def select_prefill_tokens(self, *args): + self.selected = "prefill" + return torch.tensor([1]) + + def select_decode_tokens(self, *args): + self.selected = "decode" + return torch.tensor([2]) + + def update_key_state_and_compress( + self, token_k, logical_positions, rope_positions, metadata, **kwargs + ): + self.logical_positions = logical_positions.clone() + + +class _DispatchMetadata: + token_to_kv_pool = None + out_cache_loc = None + compress_member_rows = None + decode_logical_positions = None + pending_ring_slots = None + # Consumed by the real _pending_ring_slots helper the dispatch indexer + # borrows: one token row owned by request slot 1. + token_to_batch_idx = torch.zeros(2, dtype=torch.int32) + req_pool_indices = torch.ones(2, dtype=torch.int32) + sequence_lengths = torch.ones(2, dtype=torch.int32) + + @staticmethod + def get_decode_mqa_inputs(layer_id): + return ( + torch.zeros(2, 1, 1, 128), + torch.zeros(1, 1, dtype=torch.int32), + torch.ones(1, dtype=torch.int32), + 1, + ) + + @staticmethod + def get_prefill_mqa_inputs(layer_id, positions): + return ( + torch.zeros(1, 1, 128), + torch.zeros(1, dtype=torch.int32), + torch.ones(1, dtype=torch.int32), + torch.ones(1, dtype=torch.int32), + ) + + @staticmethod + def get_token_to_batch_idx(): + return torch.zeros(1, dtype=torch.int32) + + @staticmethod + def get_seqlens_int32(): + return torch.ones(1, dtype=torch.int32) + + @staticmethod + def get_seqlens_expanded(): + return torch.tensor([7], dtype=torch.int32) + + +class _ForwardMode: + def __init__(self, decode): + self.decode = decode + + def is_decode(self): + return self.decode + + +def test_qsa_row_ranges_do_not_cross_sequences(): + sequence_lengths = torch.tensor([10, 7], dtype=torch.int32) + query_positions = torch.tensor([8, 9, 4, 6], dtype=torch.int32) + query_sequence_ids = torch.tensor([0, 0, 1, 1], dtype=torch.int32) + starts, ends, compressed_cu = build_qsa_row_ranges( + sequence_lengths, query_positions, query_sequence_ids, COMPRESS_RATIO + ) + assert compressed_cu.tolist() == [0, 2, 3] + assert starts.tolist() == [0, 0, 2, 2] + assert ends.tolist() == [2, 2, 3, 3] + + +def test_qsa_weight_free_mqa_logits_matches_explicit_formula(): + torch.manual_seed(1) + q = torch.randn(3, 4, 128, dtype=torch.bfloat16) + k = torch.randn(5, 1, 128, dtype=torch.bfloat16) + starts = torch.tensor([0, 1, 3], dtype=torch.int32) + ends = torch.tensor([2, 4, 5], dtype=torch.int32) + actual = qsa_mqa_prefill(q, k, starts, ends) + broadcast_k = k.expand(-1, 4, -1) + expected = torch.einsum("mhd,nhd->mnh", q.float(), broadcast_k.float()) + expected = torch.relu(expected).sum(-1) / (128**0.5) + columns = torch.arange(5).unsqueeze(0) + expected.masked_fill_( + (columns < starts[:, None]) | (columns >= ends[:, None]), -float("inf") + ) + torch.testing.assert_close(actual, expected) + + +def test_qsa_prefill_selection_microchunks_rows(monkeypatch): + rows, keys, heads, head_dim = 65, 64, 4, 8 + token_topk, compress_ratio = 8, 4 + indexer = SimpleNamespace( + token_topk=token_topk, + compress_ratio=compress_ratio, + block_topk=token_topk // compress_ratio, + ) + torch.manual_seed(2) + q = torch.randn(rows, heads, head_dim, dtype=torch.bfloat16) + k = torch.randn(keys, 1, head_dim, dtype=torch.bfloat16) + starts = torch.zeros(rows, dtype=torch.int32) + ends = torch.full((rows,), keys, dtype=torch.int32) + positions = torch.full((rows,), keys * compress_ratio - 1, dtype=torch.long) + sequence_lengths = torch.full((rows,), keys * compress_ratio, dtype=torch.int32) + + monkeypatch.setattr( + qsa_indexer_module, + "_QSA_PREFILL_LOGITS_BUDGET_BYTES", + 32 * keys * torch.float32.itemsize, + ) + assert qsa_indexer_module._qsa_prefill_row_chunk_size(rows, keys, heads) == 32 + actual = QSAIndexer.select_prefill_tokens( + indexer, q, k, starts, ends, positions, sequence_lengths + ) + + logits = qsa_mqa_prefill(q, k, starts, ends) + blocks = qsa_fast_topk(logits, starts, ends, topk=indexer.block_topk) + expected = expand_qsa_block_indices( + blocks, + positions, + sequence_lengths, + compress_ratio=compress_ratio, + token_topk=token_topk, + ) + torch.testing.assert_close(actual, expected) + + +def test_qsa_forward_cuda_dispatches_prefill_and_decode_mqa(): + indexer = _DispatchIndexer() + metadata = _DispatchMetadata() + inputs = torch.zeros(1, 16) + positions = torch.zeros(1, dtype=torch.int32) + + prefill_result = QSAIndexer.forward_cuda( + indexer, + inputs, + positions, + SimpleNamespace(forward_mode=_ForwardMode(False)), + metadata, + ) + assert indexer.selected == "prefill" + assert indexer.logical_positions.tolist() == [0] + assert prefill_result.item() == 1 + + positions.fill_(99) + decode_result = QSAIndexer.forward_cuda( + indexer, + inputs, + positions, + SimpleNamespace(forward_mode=_ForwardMode(True)), + metadata, + ) + assert indexer.selected == "decode" + assert indexer.logical_positions.tolist() == [6] + assert decode_result.item() == 2 + + +def test_qsa_block_expansion_adds_only_incomplete_tail(): + blocks = torch.full((2, BLOCK_TOPK), -1, dtype=torch.int32) + blocks[0, 0] = 0 + blocks[1, :2] = torch.tensor([1, 0]) + result = expand_qsa_block_indices( + blocks, + query_positions=torch.tensor([5, 10]), + sequence_lengths=torch.tensor([6, 11]), + compress_ratio=COMPRESS_RATIO, + token_topk=TOKEN_TOPK, + ) + assert result.shape == (2, FINAL_TOPK) + assert result[0, :6].tolist() == [0, 1, 2, 3, 4, 5] + assert sorted(result[1, :11].tolist()) == list(range(11)) + assert torch.all(result[0, 6:] == -1) + assert torch.all(result[1, 11:] == -1) + + +def test_qsa_triton_block_expansion_matches_torch_reference(): + if not torch.cuda.is_available(): + return + torch.manual_seed(9) + # Regression coverage for the largest default CUDA graph batch. The old + # 2-D Triton expansion exceeded the per-program element limit at bs=512. + rows = 512 + blocks = torch.full((rows, BLOCK_TOPK), -1, dtype=torch.int32) + query_positions = torch.randint(0, 3000, (rows,), dtype=torch.int64) + sequence_lengths = query_positions + 1 + complete_blocks = (query_positions + 1) // COMPRESS_RATIO + valid_counts = torch.randint(0, BLOCK_TOPK + 1, (rows,)) + for row in range(rows): + width = min(int(valid_counts[row]), int(complete_blocks[row])) + if width: + blocks[row, :width] = torch.randperm(int(complete_blocks[row]))[:width].to( + torch.int32 + ) + + expected = torch_expand_qsa_block_indices( + blocks, + query_positions, + sequence_lengths, + COMPRESS_RATIO, + TOKEN_TOPK, + ) + actual = triton_expand_qsa_block_indices( + blocks.cuda(), + query_positions.cuda(), + sequence_lengths.cuda(), + COMPRESS_RATIO, + TOKEN_TOPK, + ).cpu() + assert torch.equal(actual, expected) + + +def test_qsa_decode_mqa_reads_paged_cache(): + torch.manual_seed(7) + q = torch.randn(2, 4, 128, dtype=torch.bfloat16) + cache = torch.randn(8, 64, 1, 128, dtype=torch.bfloat16) + page_table = torch.tensor([[3, 1, 5], [4, 2, 0]], dtype=torch.int32) + context_lens = torch.tensor([130, 67], dtype=torch.int32) + actual = qsa_mqa_decode( + q, + cache, + page_table, + context_lens, + max_model_len=192, + ) + + gathered = cache[page_table.long(), :, 0].reshape(2, 192, 128) + expected = torch.einsum("bhd,bnd->bnh", q.float(), gathered.float()) + expected = torch.relu(expected).sum(-1) / (128**0.5) + positions = torch.arange(192).unsqueeze(0) + expected.masked_fill_(positions >= context_lens[:, None], -float("inf")) + torch.testing.assert_close(actual, expected) + + +def test_qsa_mtp_step_out_cache_loc_matches_draft_forward_layout(): + """EagleDraftWorker.draft_forward gives each MTP draft step an out_cache_loc slice; + the step's metadata must reference that slice, not the first batch_size slots. + """ + from sglang.srt.layers.attention.qwen_sparse_attn_backend import ( + QwenSparseMultiStepDraftBackend, + ) + + backend = QwenSparseMultiStepDraftBackend.__new__(QwenSparseMultiStepDraftBackend) + backend.topk, backend.speculative_num_steps = 1, 3 + bs, topk, steps = 4, 1, 3 + flat = torch.arange(bs * topk * steps, dtype=torch.int64) + fb = SimpleNamespace(out_cache_loc=flat, batch_size=bs, seq_lens=torch.ones(bs)) + # Reference: the exact draft_forward expression chain. + reference = flat.reshape(bs, topk, steps).permute(2, 0, 1).reshape(steps, -1) + for step in range(steps): + got = backend._step_out_cache_loc(fb, step) + assert torch.equal(got, reference[step]), step + # Steps must address disjoint slots -- the guarded regression is every + # step landing on the first bs slots. + step_slices = [backend._step_out_cache_loc(fb, i) for i in range(steps)] + all_slots = torch.cat(step_slices) + assert all_slots.unique().numel() == bs * topk * steps + # Non-draft layouts pass through untouched instead of inventing slots. + odd = SimpleNamespace( + out_cache_loc=torch.arange(5), batch_size=bs, seq_lens=torch.ones(bs) + ) + assert torch.equal(backend._step_out_cache_loc(odd, 1), odd.out_cache_loc) + backend.speculative_num_steps = 1 + assert torch.equal(backend._step_out_cache_loc(fb, 0), flat) + + +def test_qsa_graph_metadata_kernels_match_legacy_host_path(): + """For decode rows and target-verify fan-out (boundary and non-boundary), + replay kernels and the host refresh must build identical graph buffers.""" + from sglang.srt.layers.attention.qsa.graph_metadata import launch_graph_metadata + + device = "cuda" + ratio, full_page = 4, 64 + + class _Pool: + qsa_compress_ratio = ratio + # compressed slots per full-KV page (page 64 tokens / ratio 4) + qsa_compressed_page_size = full_page // ratio + qsa_block_topk = 512 + + def run_case(mode, bs, num_rows, seq_lens_list, extend_len, extend_lens=None): + pool = _Pool() + backend = QwenSparseAttnBackend.__new__(QwenSparseAttnBackend) + # A page-aligned token table: request r's token i lives in full page + # (r * 32 + i // 64) at offset i % 64, mirroring the paged allocator. + rows = torch.arange(8, dtype=torch.int32, device=device)[:, None] + cols = torch.arange(2048, dtype=torch.int32, device=device)[None, :] + backend.req_to_token = ( + (rows * 32 + cols // full_page) * full_page + cols % full_page + ).contiguous() + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device) + req_pool_indices = torch.arange(bs, dtype=torch.int32, device=device) + + def make_metadata(): + indexer = QSAIndexerMetadata( + sequence_lengths=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + token_to_batch_idx=torch.arange( + num_rows, dtype=torch.int32, device=device + ), + token_slot_table=torch.zeros( + (num_rows, 1), dtype=torch.int32, device=device + ), + out_cache_loc=torch.zeros(num_rows, dtype=torch.int64, device=device), + token_to_kv_pool=pool, + compress_ratio=ratio, + block_topk=512, + req_pool_indices=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + is_cuda_graph=True, + graph_write_locs=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + graph_compressed_page_table=torch.zeros( + (num_rows, 32), dtype=torch.int32, device=device + ), + graph_compressed_lengths=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + graph_prefix_lengths=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + decode_logical_positions=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + pending_ring_slots=torch.zeros( + num_rows, dtype=torch.int64, device=device + ), + graph_ring_group_locs=torch.zeros( + (num_rows, ratio), dtype=torch.int32, device=device + ), + ) + return qsa_backend_module.QwenSparseAttnMetadata( + sequence_lengths=indexer.sequence_lengths, + token_to_batch_idx=indexer.token_to_batch_idx, + token_slot_table=indexer.token_slot_table, + indexer_metadata=indexer, + row_req_pool_indices=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + is_cuda_graph=True, + ) + + # Path 1: recorded kernels. + kernel_metadata = make_metadata() + launch_graph_metadata( + mode=mode, + bs=bs, + num_rows=num_rows, + seq_lens=seq_lens, + req_pool_indices=req_pool_indices, + extend_lens=( + None + if extend_lens is None + else torch.tensor(extend_lens, dtype=torch.int32, device=device) + ), + extend_len=extend_len, + num_padding=0, + metadata=kernel_metadata, + req_to_token=backend.req_to_token, + pool=pool, + ) + # Path 2: legacy host refresh over the layout the kernels produced. + host_metadata = make_metadata() + host_metadata.sequence_lengths.copy_(kernel_metadata.sequence_lengths) + host_metadata.row_req_pool_indices.copy_(kernel_metadata.row_req_pool_indices) + backend._update_qsa_cuda_graph_metadata( + host_metadata.indexer_metadata, host_metadata.row_req_pool_indices + ) + for field in ( + "graph_write_locs", + "graph_compressed_page_table", + "graph_compressed_lengths", + "decode_logical_positions", + "pending_ring_slots", + "graph_ring_group_locs", + ): + kernel_buf = getattr(kernel_metadata.indexer_metadata, field) + host_buf = getattr(host_metadata.indexer_metadata, field) + assert torch.equal(kernel_buf, host_buf), (mode, field) + + # Decode: lengths straddling boundaries (256 is a boundary, others not). + run_case(mode=0, bs=4, num_rows=4, seq_lens_list=[255, 256, 257, 512], extend_len=0) + # Target verify: 2 requests x 3 draft rows, one request crossing a boundary. + run_case(mode=1, bs=2, num_rows=6, seq_lens_list=[254, 300], extend_len=3) + # Draft extend: ragged per-request rows plus a dummy-tail row. + run_case( + mode=2, + bs=2, + num_rows=6, + seq_lens_list=[256, 303], + extend_len=0, + extend_lens=[3, 2], + ) + + +def _qsa_expected_graph_layout( + *, mode, bs, num_rows, seq_lens, req_pool, extend_lens, extend_len, num_padding +): + """Independent host oracle for the graph row layout and padded dummy tail.""" + real_reqs = bs - num_padding + row_lens, row_prefix, row_reqs = [], [], [] + for pid in range(bs): + base = seq_lens[pid] + if mode == 0: + row_lens.append(base) + row_prefix.append(max(base - 1, 0)) + row_reqs.append(req_pool[pid]) + continue + eff = ( + (extend_len if pid < real_reqs else 0) + if mode == 1 + else (extend_lens[pid] if pid < real_reqs else 0) + ) + prefix, limit = (base, base + eff) if mode == 1 else (max(base - eff, 0), base) + for j in range(eff): + row_lens.append(min(prefix + 1 + j, limit)) + row_prefix.append(prefix) + row_reqs.append(req_pool[pid]) + # Dummy tail: length 1, prefix 0, aliased to request slot 0 (never + # allocated, so its pending-ring rows are the inert store dump). + while len(row_lens) < num_rows: + row_lens.append(1) + row_prefix.append(0) + row_reqs.append(0) + return row_lens, row_prefix, row_reqs + + +def test_qsa_graph_layout_covers_speculative_rows_and_padded_tail(): + """The layout kernel must rebuild speculative row fan-out and the padded dummy tail; + the row-metadata kernel must derive compressed slots from those rows.""" + from sglang.srt.layers.attention.qsa.graph_metadata import launch_graph_metadata + + device = "cuda" + ratio, full_page = 4, 64 + + class _Pool: + qsa_compress_ratio = ratio + qsa_compressed_page_size = full_page // ratio + qsa_block_topk = 512 + + def run_case(mode, bs, num_rows, seq_lens, extend_lens, extend_len, num_padding): + pool = _Pool() + backend = QwenSparseAttnBackend.__new__(QwenSparseAttnBackend) + rows = torch.arange(8, dtype=torch.int32, device=device)[:, None] + cols = torch.arange(4096, dtype=torch.int32, device=device)[None, :] + backend.req_to_token = ( + (rows * 64 + cols // full_page) * full_page + cols % full_page + ).contiguous() + req_pool = list(range(bs)) + max_pages = 64 + indexer = QSAIndexerMetadata( + sequence_lengths=torch.zeros(num_rows, dtype=torch.int32, device=device), + token_to_batch_idx=torch.arange(num_rows, dtype=torch.int32, device=device), + token_slot_table=torch.zeros( + (num_rows, 1), dtype=torch.int32, device=device + ), + out_cache_loc=torch.zeros(num_rows, dtype=torch.int64, device=device), + token_to_kv_pool=pool, + compress_ratio=ratio, + block_topk=512, + req_pool_indices=torch.zeros(num_rows, dtype=torch.int32, device=device), + is_cuda_graph=True, + graph_write_locs=torch.zeros(num_rows, dtype=torch.int32, device=device), + graph_compressed_page_table=torch.zeros( + (num_rows, max_pages), dtype=torch.int32, device=device + ), + graph_compressed_lengths=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + graph_prefix_lengths=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + decode_logical_positions=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + pending_ring_slots=torch.zeros(num_rows, dtype=torch.int64, device=device), + graph_ring_group_locs=torch.zeros( + (num_rows, ratio), dtype=torch.int32, device=device + ), + ) + metadata = qsa_backend_module.QwenSparseAttnMetadata( + sequence_lengths=indexer.sequence_lengths, + token_to_batch_idx=indexer.token_to_batch_idx, + token_slot_table=indexer.token_slot_table, + indexer_metadata=indexer, + row_req_pool_indices=torch.zeros( + num_rows, dtype=torch.int32, device=device + ), + is_cuda_graph=True, + ) + launch_graph_metadata( + mode=mode, + bs=bs, + num_rows=num_rows, + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), + req_pool_indices=torch.tensor(req_pool, dtype=torch.int32, device=device), + extend_lens=( + None + if extend_lens is None + else torch.tensor(extend_lens, dtype=torch.int32, device=device) + ), + extend_len=extend_len, + num_padding=num_padding, + metadata=metadata, + req_to_token=backend.req_to_token, + pool=pool, + ) + torch.cuda.synchronize() + + exp_lens, exp_prefix, exp_reqs = _qsa_expected_graph_layout( + mode=mode, + bs=bs, + num_rows=num_rows, + seq_lens=seq_lens, + req_pool=req_pool, + extend_lens=extend_lens, + extend_len=extend_len, + num_padding=num_padding, + ) + assert metadata.sequence_lengths.tolist() == exp_lens, (mode, "lengths") + assert indexer.graph_prefix_lengths.tolist() == exp_prefix, (mode, "prefix") + assert metadata.row_req_pool_indices.tolist() == exp_reqs, (mode, "req rows") + + r2t = backend.req_to_token.cpu() + for row, (length, req) in enumerate(zip(exp_lens, exp_reqs)): + assert int(indexer.graph_compressed_lengths[row]) == length // ratio + if length > 0 and length % ratio == 0: + expect = int(r2t[req, length - 1]) // ratio + else: + expect = 0 # non-boundary rows keep the inert reserved slot + assert int(indexer.graph_write_locs[row]) == expect, (row, length) + + # Target verify: uniform 4-token window, 2 padded request slots. + run_case( + mode=1, + bs=6, + num_rows=32, + seq_lens=[254, 255, 256, 300, 7, 7], + extend_lens=None, + extend_len=4, + num_padding=2, + ) + # Draft extend: per-request extend lengths (accept-count dependent), padded. + run_case( + mode=2, + bs=5, + num_rows=24, + seq_lens=[260, 512, 257, 9, 9], + extend_lens=[3, 1, 4, 0, 0], + extend_len=0, + num_padding=2, + ) + # Decode with a padded tail (dummy rows alias request slot 0). + run_case( + mode=0, + bs=4, + num_rows=4, + seq_lens=[256, 1024, 1025, 4096], + extend_lens=None, + extend_len=0, + num_padding=0, + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernel/qsa/test_qsa_indexer.py b/test/registered/kernel/qsa/test_qsa_indexer.py new file mode 100644 index 000000000..9664b83f9 --- /dev/null +++ b/test/registered/kernel/qsa/test_qsa_indexer.py @@ -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"])) diff --git a/test/registered/kernel/speculative/test_verify_commit_triton.py b/test/registered/kernel/speculative/test_verify_commit_triton.py new file mode 100644 index 000000000..fed30ec8f --- /dev/null +++ b/test/registered/kernel/speculative/test_verify_commit_triton.py @@ -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"])) diff --git a/test/registered/unit/configs/test_model_config.py b/test/registered/unit/configs/test_model_config.py index 6ef9629a4..bdc43661f 100644 --- a/test/registered/unit/configs/test_model_config.py +++ b/test/registered/unit/configs/test_model_config.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index 8404ece04..7bfa12dd4 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/test/registered/unit/mem_cache/test_mamba_unittest.py @@ -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() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index aa2230eff..e8db519fb 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -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", diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py index a35ce3ee1..038fbf884 100644 --- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py +++ b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py @@ -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 diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 7618dc2dc..722be661b 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -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"))