support qwen 3.8 flash next (#37500)

Co-authored-by: ch-wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: ispobock <26454835+ispobock@users.noreply.github.com>
Co-authored-by: JustinTong0323 <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: samuellees <26428561+samuellees@users.noreply.github.com>
Co-authored-by: YAMY1234 <74099316+YAMY1234@users.noreply.github.com>
Co-authored-by: yhyang201 <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: yizhang2077 <25844240+yizhang2077@users.noreply.github.com>
Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: Shinto C V <cshintov@gmail.com>
Co-authored-by: Julian Huang <huangzhilin.hzl@antgroup.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>

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 <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace sglang {
+
+/// \brief Round a float to the storage dtype and back (one eager aten step).
+/// The inline cvt keeps nvcc from folding the round-trip away: every eager
+/// bf16/fp16 aten op rounds its result, and this kernel must round likewise.
+template <typename T>
+SGL_DEVICE float eager_round(float x) {
+  return static_cast<float>(DTypeTrait<T>::from(x));
+}
+
+template <>
+SGL_DEVICE float eager_round<bf16_t>(float x) {
+#ifndef USE_ROCM
+  uint16_t u;
+  asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(u) : "f"(x));
+  return __bfloat162float(__ushort_as_bfloat16(u));
+#else
+  return static_cast<float>(DTypeTrait<bf16_t>::from(x));
+#endif
+}
+
+template <>
+SGL_DEVICE float eager_round<fp16_t>(float x) {
+#ifndef USE_ROCM
+  uint16_t u;
+  asm("cvt.rn.f16.f32 %0, %1;" : "=h"(u) : "f"(x));
+  return __ushort_as_half(u);
+#else
+  return static_cast<float>(DTypeTrait<fp16_t>::from(x));
+#endif
+}
+
+/**
+ * \brief Apply (M)RoPE to a normed row staged in shared memory.
+ *
+ * Reproduces the eager `get_cos_sin_with_position` + `apply_rotary_emb` chain:
+ * `axis_map[i]` selects which position axis feeds pair index i (plain RoPE
+ * uses all-zero maps; Qwen interleaved/sectioned MRoPE maps are built on the
+ * host). The cos/sin cache row is [cos(half), sin(half)] of width rotary_dim.
+ *
+ * \tparam T          Storage element type: bf16_t | fp16_t.
+ * \tparam kHeadDim   Compile-time head dimension (multiple of 32).
+ * \tparam kIsNeox    true -> NeoX pairing (d, d+half); false -> GPT-J (2i, 2i+1).
+ * \param smem_row    Normed row [kHeadDim], one warp cooperates.
+ * \param out_row     Destination row [kHeadDim].
+ * \param cos_sin_cache [num_positions, rotary_dim] fp32 cache.
+ * \param axis_map    [rotary_dim/2] int32 position-axis selector per pair.
+ * \param pos         Resolved per-axis positions for this token (>= 3 entries).
+ * \param rotary_dim  Rotated prefix length; tail dims pass through.
+ */
+template <typename T, int kHeadDim, bool kIsNeox>
+SGL_DEVICE void qsa_mrope_apply(
+    const T* __restrict__ smem_row,
+    T* __restrict__ out_row,
+    const float* __restrict__ cos_sin_cache,
+    const int32_t* __restrict__ axis_map,
+    const int64_t* pos,
+    const int32_t rotary_dim) {
+  using namespace device;
+  constexpr int kPerLane = kHeadDim / kWarpThreads;
+  using vec_t = AlignedVector<T, kPerLane>;
+  const uint32_t lane = threadIdx.x % kWarpThreads;
+  const int32_t half = rotary_dim / 2;
+
+  vec_t ov;
+#pragma unroll
+  for (int i = 0; i < kPerLane; ++i) {
+    const int32_t d = static_cast<int32_t>(lane * kPerLane) + i;
+    T o;
+    if constexpr (kIsNeox) {
+      if (d < half) {
+        const int32_t p = d + half;
+        const float* row = cos_sin_cache + pos[axis_map[d]] * rotary_dim;
+        const float c = eager_round<T>(row[d]);
+        const float s = eager_round<T>(row[half + d]);
+        const float nd = static_cast<float>(smem_row[d]);
+        const float np = static_cast<float>(smem_row[p]);
+        o = DTypeTrait<T>::from(eager_round<T>(nd * c) - eager_round<T>(np * s));
+      } else if (d < rotary_dim) {
+        const int32_t p = d - half;
+        const float* row = cos_sin_cache + pos[axis_map[p]] * rotary_dim;
+        const float c = eager_round<T>(row[p]);
+        const float s = eager_round<T>(row[half + p]);
+        const float nd = static_cast<float>(smem_row[d]);
+        const float np = static_cast<float>(smem_row[p]);
+        o = DTypeTrait<T>::from(eager_round<T>(nd * c) + eager_round<T>(np * s));
+      } else {
+        o = smem_row[d];
+      }
+    } else {
+      if (d < rotary_dim) {
+        const int32_t p = d / 2;
+        const float* row = cos_sin_cache + pos[axis_map[p]] * rotary_dim;
+        const float c = eager_round<T>(row[p]);
+        const float s = eager_round<T>(row[half + p]);
+        const int32_t q = (d % 2 == 0) ? d + 1 : d - 1;
+        const float nd = static_cast<float>(smem_row[d]);
+        const float nq = static_cast<float>(smem_row[q]);
+        const float t1 = eager_round<T>(nd * c);
+        const float t2 = eager_round<T>(nq * s);
+        o = DTypeTrait<T>::from((d % 2 == 0) ? t1 - t2 : t1 + t2);
+      } else {
+        o = smem_row[d];
+      }
+    }
+    ov[i] = o;
+  }
+  ov.store(out_row, lane);  // offset is in vector units
+}
+
+/**
+ * \brief Gemma RMSNorm of one row into shared memory (warp-cooperative).
+ *
+ * out = x * rsqrt(mean(x^2) + eps) * (1 + w), fp32 math with one rounding.
+ * The sum-of-squares uses flashinfer RMSNormKernel's exact partial layout
+ * (vec_size = 8 for 2-byte dtypes, thread t sums elements [8t, 8t+8), inactive
+ * lanes contribute 0, xor butterfly over the warp) so results stay bit-equal
+ * to the eager sgl_kernel gemma_rmsnorm this replaces.
+ */
+template <typename T, int kHeadDim>
+SGL_DEVICE void qsa_gemma_norm_row(
+    const T* __restrict__ x_row, const T* __restrict__ weight, const float eps, T* __restrict__ smem_row) {
+  using namespace device;
+  constexpr int kPerLane = kHeadDim / kWarpThreads;
+  using vec_t = AlignedVector<T, kPerLane>;
+  const uint32_t lane = threadIdx.x % kWarpThreads;
+
+  vec_t xv, wv;
+  xv.load(x_row, lane);  // offset is in vector units
+  wv.load(weight, lane);
+
+  float xf[kPerLane];
+#pragma unroll
+  for (int i = 0; i < kPerLane; ++i) {
+    xf[i] = static_cast<float>(xv[i]);
+  }
+
+  static_assert(kHeadDim % 8 == 0);
+  constexpr uint32_t kNormThreads = kHeadDim / 8;
+  float ss = 0.0f;
+  if (lane < kNormThreads) {
+    AlignedVector<T, 4> va, vb;
+    va.load(x_row, lane * 2);      // elements [8*lane, 8*lane+4)
+    vb.load(x_row, lane * 2 + 1);  // elements [8*lane+4, 8*lane+8)
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+      const float f = static_cast<float>(va[i]);
+      ss += f * f;
+    }
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+      const float f = static_cast<float>(vb[i]);
+      ss += f * f;
+    }
+  }
+  ss = warp::reduce_sum(ss);
+  const float nf = math::rsqrt(ss / kHeadDim + eps);
+
+#pragma unroll
+  for (int i = 0; i < kPerLane; ++i) {
+    const float wf = static_cast<float>(wv[i]);
+    smem_row[lane * kPerLane + i] = DTypeTrait<T>::from(xf[i] * nf * (1.0f + wf));
+  }
+  __syncwarp();
+}
+
+struct QsaIndexQPrepParams {
+  const void* qk;                 // [tokens, (num_q_heads + 1) * kHeadDim]
+  void* q_out;                    // [tokens, q_heads_padded, kHeadDim]
+  const void* weight;             // [kHeadDim]
+  const float* cos_sin_cache;     // [positions_capacity, rotary_dim]
+  const int32_t* axis_map;        // [rotary_dim / 2]
+  const int64_t* positions;       // [num_axes, tokens] (row stride may exceed tokens)
+  const int64_t* cache_loc;       // [tokens]
+  void* key_state_buffer;         // [slots, kHeadDim]
+  int64_t* rope_position_buffer;  // [slots, 3]
+  int64_t positions_stride;
+  int32_t num_axes;
+  int32_t num_q_heads;
+  int32_t q_heads_padded;
+  int32_t rotary_dim;
+  float eps;
+};
+
+/**
+ * \brief Per-token fused index-Q prep: gemma norm + MRoPE for every query
+ * head, zero-fill of padded heads, raw token-K store and RoPE-position store.
+ * One CTA (4 warps) per token; one warp per query head.
+ */
+template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
+__global__ __launch_bounds__(128) void qsa_index_q_prep_kernel(const QsaIndexQPrepParams __grid_constant__ params) {
+  using namespace device;
+  constexpr int kPerLane = kHeadDim / kWarpThreads;
+  using vec_t = AlignedVector<T, kPerLane>;
+  const uint32_t token = blockIdx.x;
+  const uint32_t warp = threadIdx.x / kWarpThreads;
+  const uint32_t lane = threadIdx.x % kWarpThreads;
+  __shared__ T smem_rows[4][kHeadDim];
+
+  device::PDLWaitPrimary<kUsePDL>();
+
+  const int64_t qk_row = static_cast<int64_t>(token) * (params.num_q_heads + 1) * kHeadDim;
+  const int64_t loc = params.cache_loc[token];
+  int64_t pos[3];
+#pragma unroll
+  for (int a = 0; a < 3; ++a) {
+    const int64_t ax = a < params.num_axes ? a : 0;
+    pos[a] = params.positions[ax * params.positions_stride + token];
+  }
+
+  for (int32_t h = static_cast<int32_t>(warp); h < params.q_heads_padded; h += 4) {
+    T* out_row = static_cast<T*>(params.q_out) + (static_cast<int64_t>(token) * params.q_heads_padded + h) * kHeadDim;
+    if (h < params.num_q_heads) {
+      const T* x_row = static_cast<const T*>(params.qk) + qk_row + h * kHeadDim;
+      qsa_gemma_norm_row<T, kHeadDim>(x_row, static_cast<const T*>(params.weight), params.eps, smem_rows[warp]);
+      qsa_mrope_apply<T, kHeadDim, kIsNeox>(
+          smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim);
+    } else {
+      vec_t zv;
+      zv.fill(DTypeTrait<T>::from(0.0f));
+      zv.store(out_row, lane);  // offset is in vector units
+    }
+  }
+
+  // Raw token-K and RoPE coordinates are stored for every token, whether or
+  // not the token completes a compression group.
+  if (warp == 0) {
+    vec_t kv;
+    kv.load(
+        static_cast<const T*>(params.qk) + qk_row + params.num_q_heads * kHeadDim,
+        lane);  // offset is in vector units
+    kv.store(static_cast<T*>(params.key_state_buffer) + loc * kHeadDim, lane);
+  }
+  if (warp == 1 && lane < 3) {
+    params.rope_position_buffer[loc * 3 + lane] = pos[lane];
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+}
+
+struct QsaIndexKCompressParams {
+  const void* key_state_buffer;         // [slots, kHeadDim]
+  const int32_t* group_locs;            // [groups, compress_ratio]
+  const int64_t* rope_position_buffer;  // [slots, 3]
+  const float* cos_sin_cache;           // [positions_capacity, rotary_dim]
+  const int32_t* axis_map;              // [rotary_dim / 2]
+  const void* weight;                   // [kHeadDim]
+  const int32_t* write_locs;            // [groups]
+  void* compressed_k_buffer;            // [compressed_slots, kHeadDim]
+  int32_t compress_ratio;
+  int32_t rotary_dim;
+  int32_t num_groups;
+  float eps;
+};
+
+/**
+ * \brief Per-group compressed-K prep: fp32 mean over the group, gemma norm,
+ * MRoPE at the group-start position, store into the compressed cache.
+ * One warp per group.
+ */
+template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
+__global__
+__launch_bounds__(128) void qsa_index_k_compress_kernel(const QsaIndexKCompressParams __grid_constant__ params) {
+  using namespace device;
+  constexpr int kPerLane = kHeadDim / kWarpThreads;
+  using vec_t = AlignedVector<T, kPerLane>;
+  const uint32_t warp = threadIdx.x / kWarpThreads;
+  const uint32_t lane = threadIdx.x % kWarpThreads;
+  const uint32_t group = blockIdx.x * 4 + warp;
+  if (group >= static_cast<uint32_t>(params.num_groups)) {
+    return;
+  }
+  __shared__ T smem_rows[4][kHeadDim];
+
+  device::PDLWaitPrimary<kUsePDL>();
+
+  const int32_t* locs = params.group_locs + group * params.compress_ratio;
+  const int32_t loc0 = locs[0];
+
+  // fp32 mean over the group, rounded to the storage dtype exactly like
+  // average_pool_qsa_keys (float().mean(dim=1).to(dtype)).
+  float mf[kPerLane];
+  {
+    float acc[kPerLane];
+    for (int32_t r = 0; r < params.compress_ratio; ++r) {
+      vec_t v;
+      v.load(
+          static_cast<const T*>(params.key_state_buffer) + static_cast<int64_t>(locs[r]) * kHeadDim,
+          lane);  // offset is in vector units
+#pragma unroll
+      for (int i = 0; i < kPerLane; ++i) {
+        const float f = static_cast<float>(v[i]);
+        acc[i] = (r == 0) ? f : acc[i] + f;
+      }
+    }
+    const float inv_ratio = 1.0f / static_cast<float>(params.compress_ratio);
+#pragma unroll
+    for (int i = 0; i < kPerLane; ++i) {
+      const T m = DTypeTrait<T>::from(acc[i] * inv_ratio);
+      mf[i] = static_cast<float>(m);
+      smem_rows[warp][lane * kPerLane + i] = m;
+    }
+    __syncwarp();
+    // Sum of squares in flashinfer RMSNormKernel's exact partial layout
+    // (thread t sums elements [8t, 8t+8), inactive lanes contribute 0), so
+    // the result stays bit-equal to the eager k_layernorm this replaces.
+    static_assert(kHeadDim % 8 == 0);
+    float ss = 0.0f;
+    if (lane < kHeadDim / 8) {
+#pragma unroll
+      for (int j = 0; j < 8; ++j) {
+        const float f = static_cast<float>(smem_rows[warp][lane * 8 + j]);
+        ss += f * f;
+      }
+    }
+    ss = warp::reduce_sum(ss);
+    const float nf = math::rsqrt(ss / kHeadDim + params.eps);
+#pragma unroll
+    for (int i = 0; i < kPerLane; ++i) {
+      const float wf = static_cast<float>(static_cast<const T*>(params.weight)[lane * kPerLane + i]);
+      smem_rows[warp][lane * kPerLane + i] = DTypeTrait<T>::from(mf[i] * nf * (1.0f + wf));
+    }
+    __syncwarp();
+  }
+
+  int64_t pos[3];
+#pragma unroll
+  for (int a = 0; a < 3; ++a) {
+    pos[a] = params.rope_position_buffer[static_cast<int64_t>(loc0) * 3 + a];
+  }
+
+  T* out_row = static_cast<T*>(params.compressed_k_buffer) + static_cast<int64_t>(params.write_locs[group]) * kHeadDim;
+  qsa_mrope_apply<T, kHeadDim, kIsNeox>(
+      smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim);
+
+  device::PDLTriggerSecondary<kUsePDL>();
+}
+
+/**
+ * \brief Validate inputs and launch `qsa_index_q_prep_kernel` (one CTA per token).
+ *
+ * \tparam T         Element type: bf16_t | fp16_t.
+ * \tparam kHeadDim  Index head dimension: 64 | 128 | 256.
+ * \tparam kIsNeox   RoPE pairing style.
+ * \tparam kUsePDL   Whether to launch with PDL enabled.
+ */
+template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
+void qsa_index_q_prep(
+    tvm::ffi::TensorView qk,
+    tvm::ffi::TensorView q_out,
+    tvm::ffi::TensorView weight,
+    tvm::ffi::TensorView cos_sin_cache,
+    tvm::ffi::TensorView axis_map,
+    tvm::ffi::TensorView positions,
+    int64_t num_axes,
+    tvm::ffi::TensorView cache_loc,
+    tvm::ffi::TensorView key_state_buffer,
+    tvm::ffi::TensorView rope_position_buffer,
+    int64_t num_q_heads,
+    int64_t rotary_dim,
+    float eps) {
+  using namespace host;
+  auto tokens = SymbolicSize{"tokens"};
+  auto device = SymbolicDevice{};
+  device.set_options<kDLCUDA>();
+  constexpr int64_t D = kHeadDim;
+
+  TensorMatcher({tokens, (num_q_heads + 1) * D}).with_dtype<T>().with_device(device).verify(qk);
+  auto heads_padded = SymbolicSize{"heads_padded"};
+  TensorMatcher({tokens, heads_padded, D}).with_dtype<T>().with_device(device).verify(q_out);
+  TensorMatcher({D}).with_dtype<T>().with_device(device).verify(weight);
+  auto cache_rows = SymbolicSize{"cos_sin_cache_rows"};
+  TensorMatcher({cache_rows, rotary_dim}).with_dtype<fp32_t>().with_device(device).verify(cos_sin_cache);
+  TensorMatcher({rotary_dim / 2}).with_dtype<int32_t>().with_device(device).verify(axis_map);
+  TensorMatcher({num_axes, tokens}).with_dtype<int64_t>().with_device(device).with_strides({-1, 1}).verify(positions);
+  TensorMatcher({tokens}).with_dtype<int64_t>().with_device(device).verify(cache_loc);
+  auto slots = SymbolicSize{"state_slots"};
+  TensorMatcher({slots, D}).with_dtype<T>().with_device(device).verify(key_state_buffer);
+  TensorMatcher({slots, 3}).with_dtype<int64_t>().with_device(device).verify(rope_position_buffer);
+
+  const int64_t num_tokens = tokens.unwrap();
+  const int64_t q_heads_padded = heads_padded.unwrap();
+  CHECK_HOST(num_tokens > 0) << "qsa_index_q_prep: no tokens";
+  CHECK_HOST(num_axes == 1 || num_axes == 3) << "qsa_index_q_prep: positions must have 1 or 3 axes, got " << num_axes;
+  CHECK_HOST(q_heads_padded >= num_q_heads)
+      << "qsa_index_q_prep: padded heads " << q_heads_padded << " < num_q_heads " << num_q_heads;
+  CHECK_HOST(rotary_dim > 0 && rotary_dim % 2 == 0 && rotary_dim <= D)
+      << "qsa_index_q_prep: invalid rotary_dim " << rotary_dim;
+
+  const auto params = QsaIndexQPrepParams{
+      .qk = qk.data_ptr(),
+      .q_out = q_out.data_ptr(),
+      .weight = weight.data_ptr(),
+      .cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
+      .axis_map = static_cast<const int32_t*>(axis_map.data_ptr()),
+      .positions = static_cast<const int64_t*>(positions.data_ptr()),
+      .cache_loc = static_cast<const int64_t*>(cache_loc.data_ptr()),
+      .key_state_buffer = key_state_buffer.data_ptr(),
+      .rope_position_buffer = static_cast<int64_t*>(rope_position_buffer.data_ptr()),
+      .positions_stride = positions.stride(0),
+      .num_axes = static_cast<int32_t>(num_axes),
+      .num_q_heads = static_cast<int32_t>(num_q_heads),
+      .q_heads_padded = static_cast<int32_t>(q_heads_padded),
+      .rotary_dim = static_cast<int32_t>(rotary_dim),
+      .eps = eps,
+  };
+  LaunchKernel(static_cast<uint32_t>(num_tokens), 128, device.unwrap())
+      .enable_pdl(kUsePDL)(qsa_index_q_prep_kernel<T, kHeadDim, kIsNeox, kUsePDL>, params);
+}
+
+/**
+ * \brief Validate inputs and launch `qsa_index_k_compress_kernel` (one warp per group).
+ */
+template <typename T, int kHeadDim, bool kIsNeox, bool kUsePDL>
+void qsa_index_k_compress(
+    tvm::ffi::TensorView key_state_buffer,
+    tvm::ffi::TensorView group_locs,
+    tvm::ffi::TensorView rope_position_buffer,
+    tvm::ffi::TensorView cos_sin_cache,
+    tvm::ffi::TensorView axis_map,
+    tvm::ffi::TensorView weight,
+    tvm::ffi::TensorView write_locs,
+    tvm::ffi::TensorView compressed_k_buffer,
+    int64_t compress_ratio,
+    int64_t rotary_dim,
+    float eps) {
+  using namespace host;
+  auto device = SymbolicDevice{};
+  device.set_options<kDLCUDA>();
+  constexpr int64_t D = kHeadDim;
+
+  auto slots = SymbolicSize{"state_slots"};
+  TensorMatcher({slots, D}).with_dtype<T>().with_device(device).verify(key_state_buffer);
+  auto groups = SymbolicSize{"groups"};
+  TensorMatcher({groups, compress_ratio}).with_dtype<int32_t>().with_device(device).verify(group_locs);
+  TensorMatcher({slots, 3}).with_dtype<int64_t>().with_device(device).verify(rope_position_buffer);
+  auto cache_rows = SymbolicSize{"cos_sin_cache_rows"};
+  TensorMatcher({cache_rows, rotary_dim}).with_dtype<fp32_t>().with_device(device).verify(cos_sin_cache);
+  TensorMatcher({rotary_dim / 2}).with_dtype<int32_t>().with_device(device).verify(axis_map);
+  TensorMatcher({D}).with_dtype<T>().with_device(device).verify(weight);
+  TensorMatcher({groups}).with_dtype<int32_t>().with_device(device).verify(write_locs);
+  auto compressed_slots = SymbolicSize{"compressed_slots"};
+  TensorMatcher({compressed_slots, D}).with_dtype<T>().with_device(device).verify(compressed_k_buffer);
+
+  const int64_t num_groups = groups.unwrap();
+  CHECK_HOST(num_groups > 0) << "qsa_index_k_compress: no groups";
+  CHECK_HOST(compress_ratio > 0 && compress_ratio <= 16)
+      << "qsa_index_k_compress: invalid compress_ratio " << compress_ratio;
+  CHECK_HOST(rotary_dim > 0 && rotary_dim % 2 == 0 && rotary_dim <= D)
+      << "qsa_index_k_compress: invalid rotary_dim " << rotary_dim;
+
+  const auto params = QsaIndexKCompressParams{
+      .key_state_buffer = key_state_buffer.data_ptr(),
+      .group_locs = static_cast<const int32_t*>(group_locs.data_ptr()),
+      .rope_position_buffer = static_cast<const int64_t*>(rope_position_buffer.data_ptr()),
+      .cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
+      .axis_map = static_cast<const int32_t*>(axis_map.data_ptr()),
+      .weight = weight.data_ptr(),
+      .write_locs = static_cast<const int32_t*>(write_locs.data_ptr()),
+      .compressed_k_buffer = compressed_k_buffer.data_ptr(),
+      .compress_ratio = static_cast<int32_t>(compress_ratio),
+      .rotary_dim = static_cast<int32_t>(rotary_dim),
+      .num_groups = static_cast<int32_t>(num_groups),
+      .eps = eps,
+  };
+  LaunchKernel(static_cast<uint32_t>(div_ceil(num_groups, 4)), 128, device.unwrap())
+      .enable_pdl(kUsePDL)(qsa_index_k_compress_kernel<T, kHeadDim, kIsNeox, kUsePDL>, params);
+}
+
+}  // namespace sglang

python/sglang/kernels/jit/csrc/elementwise/fast_topk.cuh

@@ -0,0 +1,291 @@
+// Radix-select fast top-k with the AOT fast_topk_v2 semantics: for each row b,
+// select the kTopK largest scores in [row_starts[b], row_starts[b] + lengths[b])
+// and write their indices relative to row_starts[b]; order within a row is unspecified.
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace sglang {
+
+namespace fast_topk_detail {
+
+constexpr uint32_t kThreadsPerBlock = 1024;
+// Each radix pass needs at most ~kTopK candidates in the threshold bin, so
+// 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
+constexpr size_t kSmemBytes = 8 * 1024 * sizeof(uint32_t);  // 32KB
+
+struct FastTopKParams {
+  const float* __restrict__ input;         // [B, input_stride]
+  const int32_t* __restrict__ row_starts;  // [B]
+  int32_t* __restrict__ indices;           // [B, kTopK]
+  const int32_t* __restrict__ lengths;     // [B]
+  int64_t input_stride;
+};
+
+SGL_DEVICE auto convert_to_uint8(float x) -> uint8_t {
+  const __half h = __float2half_rn(x);
+  const uint16_t bits = __half_as_ushort(h);
+  const uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
+  return static_cast<uint8_t>(key >> 8);
+}
+
+SGL_DEVICE auto convert_to_uint32(float x) -> uint32_t {
+  const uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+// When length <= kTopK, write the indices directly.
+template <int kTopK>
+SGL_DEVICE void naive_topk(const float* __restrict__ score, int32_t* __restrict__ indice, int32_t length) {
+  const auto tid = threadIdx.x;
+  for (int i = tid; i < kTopK; i += kThreadsPerBlock) {
+    indice[i] = (i < length) ? i : -1;
+  }
+}
+
+// Radix-select top-k. Assumes length > kTopK (checked by the caller).
+template <int kTopK>
+SGL_DEVICE void radix_select_topk(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
+  int topk = kTopK;
+  constexpr auto BLOCK_SIZE = kThreadsPerBlock;
+  constexpr auto RADIX = 256;
+  constexpr auto SMEM_INPUT_SIZE = kSmemBytes / (2 * sizeof(int));
+
+  alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
+  alignas(128) __shared__ int s_counter;
+  alignas(128) __shared__ int s_threshold_bin_id;
+  alignas(128) __shared__ int s_num_input[2];
+
+  auto& s_histogram = s_histogram_buf[0];
+  // allocate for two rounds
+  extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
+
+  const int tx = threadIdx.x;
+
+  // stage 1: 8bit coarse histogram
+  if (tx < RADIX + 1) s_histogram[tx] = 0;
+  __syncthreads();
+
+  for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+    const auto bin = convert_to_uint8(input[idx + row_start]);
+    ::atomicAdd(&s_histogram[bin], 1);
+  }
+  __syncthreads();
+
+  const auto run_cumsum = [&] {
+#pragma unroll 8
+    for (int i = 0; i < 8; ++i) {
+      static_assert(1 << 8 == RADIX);
+      if (tx < RADIX) {
+        const auto j = 1 << i;
+        const auto k = i & 1;
+        auto value = s_histogram_buf[k][tx];
+        if (tx < RADIX - j) {
+          value += s_histogram_buf[k][tx + j];
+        }
+        s_histogram_buf[k ^ 1][tx] = value;
+      }
+      __syncthreads();
+    }
+  };
+
+  run_cumsum();
+  if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
+    s_threshold_bin_id = tx;
+    s_num_input[0] = 0;
+    s_counter = 0;
+  }
+  __syncthreads();
+
+  const auto threshold_bin = s_threshold_bin_id;
+  topk -= s_histogram[threshold_bin + 1];
+
+  if (topk == 0) {
+    for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        index[pos] = idx;
+      }
+    }
+    __syncthreads();
+    return;
+  } else {
+    __syncthreads();
+    if (tx < RADIX + 1) {
+      s_histogram[tx] = 0;
+    }
+    __syncthreads();
+
+    for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const auto raw_input = input[idx + row_start];
+      const auto bin = static_cast<int>(convert_to_uint8(raw_input));
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        index[pos] = idx;
+      } else if (bin == threshold_bin) {
+        const auto pos = ::atomicAdd(&s_num_input[0], 1);
+        // fuse the histogram computation here
+        if (pos < int(SMEM_INPUT_SIZE)) {
+          s_input_idx[0][pos] = idx;
+          const auto bin = convert_to_uint32(raw_input);
+          const auto sub_bin = (bin >> 24) & 0xFF;
+          ::atomicAdd(&s_histogram[sub_bin], 1);
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // stage 2: refine with 8bit radix passes
+#pragma unroll 4
+  for (int round = 0; round < 4; ++round) {
+    __shared__ int s_last_remain;
+    const auto r_idx = round % 2;
+
+    // clip here to prevent overflow
+    const auto _raw_num_input = s_num_input[r_idx];
+    const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
+
+    run_cumsum();
+    if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
+      s_threshold_bin_id = tx;
+      s_num_input[r_idx ^ 1] = 0;
+      s_last_remain = topk - s_histogram[tx + 1];
+    }
+    __syncthreads();
+
+    const auto threshold_bin = s_threshold_bin_id;
+    topk -= s_histogram[threshold_bin + 1];
+
+    if (topk == 0) {
+      for (int i = tx; i < num_input; i += BLOCK_SIZE) {
+        const auto idx = s_input_idx[r_idx][i];
+        const auto offset = 24 - round * 8;
+        const auto bin = (convert_to_uint32(input[idx + row_start]) >> offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const auto pos = ::atomicAdd(&s_counter, 1);
+          index[pos] = idx;
+        }
+      }
+      __syncthreads();
+      break;
+    } else {
+      __syncthreads();
+      if (tx < RADIX + 1) {
+        s_histogram[tx] = 0;
+      }
+      __syncthreads();
+      for (int i = tx; i < num_input; i += BLOCK_SIZE) {
+        const auto idx = s_input_idx[r_idx][i];
+        const auto raw_input = input[idx + row_start];
+        const auto offset = 24 - round * 8;
+        const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const auto pos = ::atomicAdd(&s_counter, 1);
+          index[pos] = idx;
+        } else if (bin == threshold_bin) {
+          if (round == 3) {
+            const auto pos = ::atomicAdd(&s_last_remain, -1);
+            if (pos > 0) {
+              index[kTopK - pos] = idx;
+            }
+          } else {
+            const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
+            if (pos < int(SMEM_INPUT_SIZE)) {
+              // fuse the histogram computation here
+              s_input_idx[r_idx ^ 1][pos] = idx;
+              const auto bin = convert_to_uint32(raw_input);
+              const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
+              ::atomicAdd(&s_histogram[sub_bin], 1);
+            }
+          }
+        }
+      }
+      __syncthreads();
+    }
+  }
+}
+
+template <int kTopK, bool kUsePDL>
+__global__ __launch_bounds__(fast_topk_detail::kThreadsPerBlock) void fast_topk_kernel(
+    const fast_topk_detail::FastTopKParams __grid_constant__ params) {
+  using namespace fast_topk_detail;
+  device::PDLWaitPrimary<kUsePDL>();
+
+  const auto bid = static_cast<uint64_t>(blockIdx.x);
+  const auto row_start = params.row_starts == nullptr ? 0 : params.row_starts[bid];
+  const auto length = params.lengths[bid];
+  const auto indice = params.indices + bid * kTopK;
+  const auto score = params.input + bid * params.input_stride;
+  if (length <= kTopK) {
+    naive_topk<kTopK>(score, indice, length);
+  } else {
+    radix_select_topk<kTopK>(score, indice, row_start, length);
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+}
+
+}  // namespace fast_topk_detail
+
+/**
+ * \brief Per-row top-k selection over ragged rows of a fp32 score matrix.
+ *
+ * Row b selects the kTopK largest values in
+ * score[b, row_starts[b] : row_starts[b] + lengths[b]) and writes their
+ * indices (relative to row_starts[b]) into indices[b]. Unfilled slots are
+ * -1 when lengths[b] < kTopK.
+ */
+template <int kTopK, bool kUsePDL>
+struct FastTopKKernel {
+  static constexpr auto kernel = fast_topk_detail::fast_topk_kernel<kTopK, kUsePDL>;
+
+  static void
+  run(const tvm::ffi::TensorView score,
+      const tvm::ffi::TensorView row_starts,
+      const tvm::ffi::TensorView indices,
+      const tvm::ffi::TensorView lengths) {
+    using namespace host;
+    auto B = SymbolicSize{"batch"};
+    auto L = SymbolicSize{"length"};
+    auto S = SymbolicSize{"input_stride"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({B, L})  // score
+        .with_strides({S, 1})
+        .with_dtype<fp32_t>()
+        .with_device(device)
+        .verify(score);
+    TensorMatcher({B})  // row_starts
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(row_starts);
+    TensorMatcher({B, kTopK})  // indices
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(indices);
+    TensorMatcher({B})  // lengths
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(lengths);
+
+    const auto params = fast_topk_detail::FastTopKParams{
+        .input = static_cast<const float*>(score.data_ptr()),
+        .row_starts = static_cast<const int32_t*>(row_starts.data_ptr()),
+        .indices = static_cast<int32_t*>(indices.data_ptr()),
+        .lengths = static_cast<const int32_t*>(lengths.data_ptr()),
+        .input_stride = S.unwrap(),
+    };
+
+    const auto num_rows = static_cast<uint32_t>(B.unwrap());
+    LaunchKernel(num_rows, fast_topk_detail::kThreadsPerBlock, device.unwrap(), fast_topk_detail::kSmemBytes)
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace sglang

python/sglang/kernels/jit/csrc/elementwise/grouped_gemma_rmsnorm.cuh

@@ -0,0 +1,178 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace sglang {
+
+struct GroupedGemmaRMSNormParams {
+  const void* input;
+  const void* __restrict__ weight;
+  void* output;
+  uint32_t num_groups;
+  float eps;
+};
+
+/**
+ * \brief Grouped Gemma-style RMSNorm: out = x * rsqrt(mean(x^2) + eps) * (1 + w).
+ *
+ * The last dim of the input is split into `num_groups` chunks of `kGroupSize`
+ * elements. Variance is computed per (token, group) chunk, so a [M, H] input
+ * with H = num_groups * kGroupSize behaves like M * num_groups independent
+ * RMSNorm rows whose weight rows are the matching kGroupSize slice of `weight`.
+ *
+ * One CTA handles one (token, group) chunk. Since chunks are contiguous in
+ * memory, block `bid` reads/writes elements [bid * kGroupSize, (bid + 1) *
+ * kGroupSize) and uses weight slice (bid % num_groups) * kGroupSize.
+ *
+ * \tparam kGroupSize Elements per group. Must be a multiple of 512.
+ * \tparam kUsePDL    Whether to emit the PDL wait/trigger pair.
+ * \tparam Float      Element type: bf16_t | fp16_t.
+ */
+template <int64_t kGroupSize, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(kGroupSize / 16) void grouped_gemma_rmsnorm_kernel(
+    const GroupedGemmaRMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Float2 = packed_t<Float>;
+#if SGL_ARCH_BLACKWELL_OR_GREATER
+  // Blackwell: 32B vector, each thread loads/stores once
+  using Storage = AlignedVector<Float2, 8>;
+  constexpr uint32_t kNumLoads = 1;
+#else
+  // Pre-Blackwell: 16B vector, each thread loads/stores twice
+  using Storage = AlignedVector<Float2, 4>;
+  constexpr uint32_t kNumLoads = 2;
+#endif
+  constexpr uint32_t kVecLen = kNumLoads == 1 ? 8 : 4;
+  constexpr auto kNumThreads = kGroupSize / 16;
+  constexpr auto kNumWarps = kNumThreads / kWarpThreads;
+
+  const uint32_t bid = blockIdx.x;
+  const uint32_t group = bid % params.num_groups;
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  // Warp 0 writes smem[tx] for all 32 lanes in the cross-warp reduce below,
+  // so this must hold kWarpThreads entries, not kNumWarps.
+  __shared__ float smem[kWarpThreads];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto input_ptr = pointer::offset<Float>(params.input, static_cast<int64_t>(bid) * kGroupSize);
+  const auto output_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(bid) * kGroupSize);
+  const auto weight_ptr = pointer::offset<Float>(params.weight, static_cast<int64_t>(group) * kGroupSize);
+
+  Storage input_vec[kNumLoads];
+  Storage weight_vec[kNumLoads];
+#pragma unroll
+  for (uint32_t j = 0; j < kNumLoads; ++j) {
+    input_vec[j] = gmem.load(input_ptr, j);
+    weight_vec[j] = gmem.load(weight_ptr, j);
+  }
+
+  float sum_of_squares = 0.0f;
+#pragma unroll
+  for (uint32_t j = 0; j < kNumLoads; ++j) {
+#pragma unroll
+    for (uint32_t i = 0; i < kVecLen; ++i) {
+      const auto [x, y] = cast<fp32x2_t>(input_vec[j][i]);
+      sum_of_squares += x * x + y * y;
+    }
+  }
+
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  float norm_factor;
+  if constexpr (kNumWarps == 1) {
+    norm_factor = math::rsqrt(sum_of_squares / kGroupSize + params.eps);
+  } else {
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    smem[warp_id] = sum_of_squares;
+    __syncthreads();
+    if (warp_id == 0) {
+      const auto tx = threadIdx.x;
+      const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
+      sum_of_squares = warp::reduce_sum(local_sum);
+      smem[tx] = math::rsqrt(sum_of_squares / kGroupSize + params.eps);
+    }
+    __syncthreads();
+    norm_factor = smem[warp_id];
+  }
+
+#pragma unroll
+  for (uint32_t j = 0; j < kNumLoads; ++j) {
+    Storage output_vec;
+#pragma unroll
+    for (uint32_t i = 0; i < kVecLen; ++i) {
+      const auto [ix, iy] = cast<fp32x2_t>(input_vec[j][i]);
+      const auto [wx, wy] = cast<fp32x2_t>(weight_vec[j][i]);
+      output_vec[i] = cast<Float2>(fp32x2_t{ix * norm_factor * (1.0f + wx), iy * norm_factor * (1.0f + wy)});
+    }
+    gmem.store(output_ptr, output_vec, j);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kGroupSize, bool kUsePDL, typename DType>
+struct GroupedGemmaRMSNormKernel {
+  static_assert(sizeof(DType) == 2, "GroupedGemmaRMSNorm only supports 2-byte dtypes");
+  static_assert(kGroupSize % 512 == 0, "kGroupSize must be a multiple of 512");
+  static constexpr auto kernel = grouped_gemma_rmsnorm_kernel<kGroupSize, kUsePDL, DType>;
+  static constexpr auto kBlockSize = static_cast<uint32_t>(kGroupSize / 16);
+
+  /**
+   * \brief Validate tensors and launch one CTA per (token, group) chunk.
+   * \param input  [M, H] contiguous, H % kGroupSize == 0
+   * \param weight [H]
+   * \param output [M, H] contiguous, same shape/dtype/device as input
+   * \param eps    RMSNorm epsilon
+   */
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto M = SymbolicSize{"num_tokens"};
+    auto H = SymbolicSize{"hidden_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M, H})  // input
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({H})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(weight);
+    TensorMatcher({M, H})  // output
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(output);
+
+    const int64_t hidden_size = H.unwrap();
+    CHECK_HOST(hidden_size % kGroupSize == 0) << "grouped_gemma_rmsnorm: hidden_size (" << hidden_size
+                                              << ") must be divisible by group_size (" << kGroupSize << ")";
+
+    const auto params = GroupedGemmaRMSNormParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .output = output.data_ptr(),
+        .num_groups = static_cast<uint32_t>(hidden_size / kGroupSize),
+        .eps = eps,
+    };
+
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    const uint32_t num_blocks = num_tokens * params.num_groups;
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace sglang

python/sglang/kernels/jit/csrc/elementwise/hc_combine.cuh

@@ -0,0 +1,381 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace sglang {
+
+struct HcCombineParams {
+  const void* block_output;     // [M, H]
+  const void* residual;         // [M, HC * H]
+  const void* normed_residual;  // [M, HC * H]
+  const void* inject_weight;    // [HC, HC * H]
+  void* output;                 // [M, HC * H]
+};
+
+/**
+ * \brief Fused HyperConnection (gated residual) combine:
+ *
+ *   a[m, c] = 2 * sigmoid(dot(normed_residual[m, :], inject_weight[c, :]) / kHcCount)
+ *   output[m, c*H + i] = residual[m, c*H + i] + a[m, c] * block_output[m, i]
+ *
+ * One CTA handles one token row. Phase 1 computes the kHcCount gate values with a
+ * block reduction over the full HC*H row (fp32 accumulation). Phase 2 streams the
+ * HC*H output elements with vectorized 16B accesses; the block_output row is only
+ * H elements, so its re-read per branch stays in L2.
+ *
+ * \tparam kHcCount    Number of hyper-connection branches (4 in production).
+ * \tparam kHiddenSize Per-branch hidden size H. HC*H must be a multiple of
+ *                     kNumThreads * kVecLen so the row maps exactly onto the CTA.
+ * \tparam kUsePDL     Whether to emit the PDL wait/trigger pair.
+ * \tparam Float       Element type: bf16_t | fp16_t.
+ */
+template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(256) void hc_combine_kernel(const HcCombineParams __grid_constant__ params) {
+  using namespace device;
+  using Float2 = packed_t<Float>;
+  using Storage = AlignedVector<Float2, 4>;  // 8 elements, 16 bytes
+  constexpr uint32_t kVecLen = 8;
+  constexpr uint32_t kNumThreads = 256;
+  constexpr int64_t kRowSize = kHcCount * kHiddenSize;
+  constexpr uint32_t kVecsPerRow = kRowSize / kVecLen;            // 1280 for 4x2560
+  constexpr uint32_t kVecsPerThread = kVecsPerRow / kNumThreads;  // 5 for 4x2560
+  constexpr uint32_t kVecsPerBranch = kHiddenSize / kVecLen;      // 320 for 2560
+  constexpr uint32_t kNumWarps = kNumThreads / kWarpThreads;
+
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  const uint32_t m = blockIdx.x;
+
+  const auto y_ptr = pointer::offset<Float>(params.block_output, static_cast<int64_t>(m) * kHiddenSize);
+  const auto r_ptr = pointer::offset<Float>(params.residual, static_cast<int64_t>(m) * kRowSize);
+  const auto n_ptr = pointer::offset<Float>(params.normed_residual, static_cast<int64_t>(m) * kRowSize);
+  const auto w_ptr = static_cast<const Float*>(params.inject_weight);
+  const auto out_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(m) * kRowSize);
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // Phase 1: gate values a_c, accumulated in fp32 and reduced across the CTA.
+  Storage n_vec[kVecsPerThread];
+#pragma unroll
+  for (uint32_t j = 0; j < kVecsPerThread; ++j) {
+    n_vec[j] = gmem.load(n_ptr, j);
+  }
+
+  float acc[kHcCount];
+#pragma unroll
+  for (int c = 0; c < kHcCount; ++c) {
+    const auto wc_ptr = w_ptr + static_cast<int64_t>(c) * kRowSize;
+    float sum = 0.0f;
+#pragma unroll
+    for (uint32_t j = 0; j < kVecsPerThread; ++j) {
+      const Storage w_vec = gmem.load(wc_ptr, j);
+#pragma unroll
+      for (uint32_t i = 0; i < kVecLen / 2; ++i) {
+        const auto [nx, ny] = cast<fp32x2_t>(n_vec[j][i]);
+        const auto [wx, wy] = cast<fp32x2_t>(w_vec[i]);
+        sum += nx * wx + ny * wy;
+      }
+    }
+    acc[c] = warp::reduce_sum(sum);
+  }
+
+  __shared__ float smem[kHcCount][kNumWarps];
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  const uint32_t lane = threadIdx.x % kWarpThreads;
+  if (lane == 0) {
+#pragma unroll
+    for (int c = 0; c < kHcCount; ++c) {
+      smem[c][warp_id] = acc[c];
+    }
+  }
+  __syncthreads();
+  __shared__ float a_shared[kHcCount];
+  if (threadIdx.x < kHcCount) {
+    float total = 0.0f;
+#pragma unroll
+    for (uint32_t w = 0; w < kNumWarps; ++w) {
+      total += smem[threadIdx.x][w];
+    }
+    a_shared[threadIdx.x] = 2.0f / (1.0f + math::exp(-total / kHcCount));
+  }
+  __syncthreads();
+
+  // Phase 2: stream the output row. Vector `vec_idx` lies entirely inside
+  // branch `vec_idx / kVecsPerBranch` (H is a multiple of kVecLen).
+#pragma unroll
+  for (uint32_t j = 0; j < kVecsPerThread; ++j) {
+    const uint32_t vec_idx = threadIdx.x + j * kNumThreads;
+    const uint32_t branch = vec_idx / kVecsPerBranch;
+    const uint32_t col_in_branch = (vec_idx % kVecsPerBranch) * kVecLen;
+    const float a = a_shared[branch];
+
+    const Storage r_vec = gmem.load(r_ptr, j);
+    Storage y_vec;
+    y_vec.load(y_ptr, col_in_branch / kVecLen);
+    Storage out_vec;
+#pragma unroll
+    for (uint32_t i = 0; i < kVecLen / 2; ++i) {
+      const auto [rx, ry] = cast<fp32x2_t>(r_vec[i]);
+      const auto [yx, yy] = cast<fp32x2_t>(y_vec[i]);
+      out_vec[i] = cast<Float2>(fp32x2_t{rx + a * yx, ry + a * yy});
+    }
+    gmem.store(out_ptr, out_vec, j);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename DType>
+struct HcCombineKernel {
+  static_assert(sizeof(DType) == 2, "HcCombine only supports 2-byte dtypes");
+  static_assert(kHcCount > 0, "kHcCount must be positive");
+  static_assert(kHiddenSize > 0 && kHiddenSize % 8 == 0, "kHiddenSize must be a multiple of 8");
+  static_assert((kHcCount * kHiddenSize) % (256 * 8) == 0, "kHcCount * kHiddenSize must be a multiple of 2048");
+  static constexpr auto kernel = hc_combine_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
+  static constexpr uint32_t kBlockSize = 256;
+
+  /**
+   * \brief Validate tensors and launch one CTA per token row.
+   * \param block_output    [M, H] contiguous
+   * \param residual        [M, HC * H] contiguous
+   * \param normed_residual [M, HC * H] contiguous, same dtype/device as residual
+   * \param inject_weight   [HC, HC * H] contiguous, same dtype/device as residual
+   * \param output          [M, HC * H] contiguous, same shape/dtype/device as residual
+   */
+  static void
+  run(const tvm::ffi::TensorView block_output,
+      const tvm::ffi::TensorView residual,
+      const tvm::ffi::TensorView normed_residual,
+      const tvm::ffi::TensorView inject_weight,
+      const tvm::ffi::TensorView output) {
+    using namespace host;
+    auto M = SymbolicSize{"num_tokens"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M, kHiddenSize})  // block_output
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(block_output);
+    TensorMatcher({M, kHcCount * kHiddenSize})  // residual, normed_residual, output
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(residual)
+        .verify(normed_residual)
+        .verify(output);
+    TensorMatcher({kHcCount, kHcCount * kHiddenSize})  // inject_weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(inject_weight);
+
+    const auto params = HcCombineParams{
+        .block_output = block_output.data_ptr(),
+        .residual = residual.data_ptr(),
+        .normed_residual = normed_residual.data_ptr(),
+        .inject_weight = inject_weight.data_ptr(),
+        .output = output.data_ptr(),
+    };
+
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    LaunchKernel(num_tokens, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+struct HcCombineSplitParams {
+  const void* block_output;
+  const void* residual;
+  const void* normed_residual;
+  const void* inject_weight;
+  void* output;
+  float* partials;
+};
+
+namespace hc_combine_split_detail {
+
+// One CTA per warp of the original kernel: CTA s replays exactly the work of
+// threads [32s, 32s+32) of the 256-thread reference, so every float is
+// accumulated in the reference order and results stay bit-identical.
+constexpr uint32_t kSplit = 8;
+constexpr uint32_t kRefThreads = 256;
+constexpr uint32_t kGateThreads = 32;
+constexpr uint32_t kApplyThreads = 160;
+constexpr uint32_t kVecLen = 8;
+
+}  // namespace hc_combine_split_detail
+
+/**
+ * \brief Stage 1 of the split combine: partial gate dots over a K slice.
+ *
+ * Grid is (rows, kSplit) so the [HC, HC*H] inject weight is read once across
+ * the whole grid instead of once per row, and the row's traffic is spread over
+ * kSplit CTAs. Each CTA writes its own partials slot, so no atomics and no
+ * buffer clearing are needed.
+ */
+template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(hc_combine_split_detail::kGateThreads) void hc_combine_gate_kernel(
+    const HcCombineSplitParams __grid_constant__ params) {
+  using namespace device;
+  using namespace hc_combine_split_detail;
+  using Float2 = packed_t<Float>;
+  using Storage = AlignedVector<Float2, 4>;
+  constexpr uint32_t kVecLen = 8;
+  constexpr int64_t kRowSize = kHcCount * kHiddenSize;
+  constexpr uint32_t kVecsPerRow = kRowSize / kVecLen;
+  constexpr uint32_t kVecsPerThread = kVecsPerRow / kRefThreads;
+  static_assert(kVecsPerRow % kRefThreads == 0);
+  static_assert(kRefThreads / kGateThreads == kSplit);
+
+  const uint32_t m = blockIdx.x;
+  const uint32_t split = blockIdx.y / kHcCount;
+  const uint32_t c = blockIdx.y % kHcCount;
+  const uint32_t ref_tid = split * kGateThreads + threadIdx.x;
+
+  const auto n_ptr = pointer::offset<Float>(params.normed_residual, static_cast<int64_t>(m) * kRowSize);
+  const auto w_ptr = static_cast<const Float*>(params.inject_weight);
+
+  PDLWaitPrimary<kUsePDL>();
+
+  Storage n_vec[kVecsPerThread];
+#pragma unroll
+  for (uint32_t j = 0; j < kVecsPerThread; ++j) {
+    n_vec[j].load(n_ptr, ref_tid + j * kRefThreads);
+  }
+
+  {
+    const auto wc_ptr = w_ptr + static_cast<int64_t>(c) * kRowSize;
+    float sum = 0.0f;
+#pragma unroll
+    for (uint32_t j = 0; j < kVecsPerThread; ++j) {
+      Storage w_vec;
+      w_vec.load(wc_ptr, ref_tid + j * kRefThreads);
+#pragma unroll
+      for (uint32_t i = 0; i < kVecLen / 2; ++i) {
+        const auto [nx, ny] = cast<fp32x2_t>(n_vec[j][i]);
+        const auto [wx, wy] = cast<fp32x2_t>(w_vec[i]);
+        sum += nx * wx + ny * wy;
+      }
+    }
+    sum = warp::reduce_sum(sum);
+    if (threadIdx.x == 0) {
+      params.partials[(static_cast<int64_t>(m) * kSplit + split) * kHcCount + c] = sum;
+    }
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+/**
+ * \brief Stage 2: reduce the partial dots and stream the combined row.
+ *
+ * Each CTA owns one contiguous vector slice of the row. kSplit divides
+ * kHiddenSize, so a slice never straddles two branches and the gate is a
+ * per-CTA scalar.
+ */
+template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(hc_combine_split_detail::kApplyThreads) void hc_combine_apply_kernel(
+    const HcCombineSplitParams __grid_constant__ params) {
+  using namespace device;
+  using namespace hc_combine_split_detail;
+  using Float2 = packed_t<Float>;
+  using Storage = AlignedVector<Float2, 4>;
+  constexpr int64_t kRowSize = kHcCount * kHiddenSize;
+  constexpr uint32_t kVecsPerRow = kRowSize / kVecLen;
+  constexpr uint32_t kVecsPerSplit = kVecsPerRow / kSplit;
+  constexpr uint32_t kVecsPerThread = kVecsPerSplit / kApplyThreads;
+  constexpr uint32_t kVecsPerBranch = kHiddenSize / kVecLen;
+  static_assert(kVecsPerBranch % kVecsPerSplit == 0);
+
+  const uint32_t m = blockIdx.x;
+  const uint32_t split = blockIdx.y;
+  const uint32_t vec_base = split * kVecsPerSplit;
+  const uint32_t branch = vec_base / kVecsPerBranch;
+
+  const auto y_ptr = pointer::offset<Float>(params.block_output, static_cast<int64_t>(m) * kHiddenSize);
+  const auto r_ptr = pointer::offset<Float>(params.residual, static_cast<int64_t>(m) * kRowSize);
+  const auto out_ptr = pointer::offset<Float>(params.output, static_cast<int64_t>(m) * kRowSize);
+
+  PDLWaitPrimary<kUsePDL>();
+
+  float total = 0.0f;
+#pragma unroll
+  for (uint32_t s = 0; s < kSplit; ++s) {
+    total += params.partials[(static_cast<int64_t>(m) * kSplit + s) * kHcCount + branch];
+  }
+  const float a = 2.0f / (1.0f + math::exp(-total / kHcCount));
+
+#pragma unroll
+  for (uint32_t j = 0; j < kVecsPerThread; ++j) {
+    const uint32_t vec_idx = vec_base + threadIdx.x + j * kApplyThreads;
+    const uint32_t col_in_branch = (vec_idx % kVecsPerBranch);
+    Storage r_vec;
+    r_vec.load(r_ptr, vec_idx);
+    Storage y_vec;
+    y_vec.load(y_ptr, col_in_branch);
+    Storage out_vec;
+#pragma unroll
+    for (uint32_t i = 0; i < kVecLen / 2; ++i) {
+      const auto [rx, ry] = cast<fp32x2_t>(r_vec[i]);
+      const auto [yx, yy] = cast<fp32x2_t>(y_vec[i]);
+      out_vec[i] = cast<Float2>(fp32x2_t{rx + a * yx, ry + a * yy});
+    }
+    out_vec.store(out_ptr, vec_idx);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kHcCount, int64_t kHiddenSize, bool kUsePDL, typename DType>
+struct HcCombineSplitKernel {
+  static_assert(sizeof(DType) == 2, "HcCombine only supports 2-byte dtypes");
+  static constexpr auto gate_kernel = hc_combine_gate_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
+  static constexpr auto apply_kernel = hc_combine_apply_kernel<kHcCount, kHiddenSize, kUsePDL, DType>;
+
+  static void
+  run(const tvm::ffi::TensorView block_output,
+      const tvm::ffi::TensorView residual,
+      const tvm::ffi::TensorView normed_residual,
+      const tvm::ffi::TensorView inject_weight,
+      const tvm::ffi::TensorView output,
+      const tvm::ffi::TensorView partials) {
+    using namespace host;
+    using namespace hc_combine_split_detail;
+    auto M = SymbolicSize{"num_tokens"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M, kHiddenSize}).with_dtype<DType>().with_device(device).verify(block_output);
+    TensorMatcher({M, kHcCount * kHiddenSize})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(residual)
+        .verify(normed_residual)
+        .verify(output);
+    TensorMatcher({kHcCount, kHcCount * kHiddenSize}).with_dtype<DType>().with_device(device).verify(inject_weight);
+    auto part_rows = SymbolicSize{"partial_rows"};
+    TensorMatcher({part_rows, kSplit, kHcCount}).with_dtype<fp32_t>().with_device(device).verify(partials);
+
+    const auto params = HcCombineSplitParams{
+        .block_output = block_output.data_ptr(),
+        .residual = residual.data_ptr(),
+        .normed_residual = normed_residual.data_ptr(),
+        .inject_weight = inject_weight.data_ptr(),
+        .output = output.data_ptr(),
+        .partials = static_cast<float*>(partials.data_ptr()),
+    };
+
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    LaunchKernel(dim3(num_tokens, kSplit * kHcCount, 1), kGateThreads, device.unwrap())
+        .enable_pdl(kUsePDL)(gate_kernel, params);
+    LaunchKernel(dim3(num_tokens, kSplit, 1), kApplyThreads, device.unwrap()).enable_pdl(kUsePDL)(apply_kernel, params);
+  }
+};
+
+}  // namespace sglang

python/sglang/kernels/kda_kernels/README.md

@@ -20,6 +20,7 @@ load Triton, CUTLASS, or compile a JIT extension.
 | Diffusion residual-gate add | `residual_gate_add_jit.py` | [sgl-project/sglang#29361](https://github.com/sgl-project/sglang/pull/29361), merge commit `495f13fa12` |
 | LTX2 QK-norm split-RoPE | `ltx2_qknorm_split_rope_jit.py` | [sgl-project/sglang#29708](https://github.com/sgl-project/sglang/pull/29708), merge commit `fcb9f229b3` |
 | FLUX.2 FP8 producer and QKV packing fusions | `layernorm_modulate_triton.py`, `flux2_qkv_epilogue_jit.py`, `flux2_token_cat_fp8_triton.py` | [sgl-project/sglang#37162](https://github.com/sgl-project/sglang/pull/37162), merge commit `1c3ad92438` |
+| Qwen3.8 QSA packed-varlen decode on SM121 | `qwen38_qsa_sm121/` | [radixark/KDA-1.5#4](https://github.com/radixark/KDA-1.5/pull/4) at `414ce456e14a`; see the package README |
 
 For JIT kernels, the Python entry module and the corresponding source under
 `csrc/` move together. The shared `sglang.kernels.jit` loader remains build

python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/README.md

@@ -0,0 +1,51 @@
+# Qwen3.8 QSA packed-varlen decode for SM121
+
+This implementation was optimized by Codex and Kimi K3 agents through
+[KDA-1.5](https://github.com/radixark/KDA-1.5). The task and immutable real
+tensor replay were registered in [radixark/KDA-1.5 PR #4](https://github.com/radixark/KDA-1.5/pull/4)
+at commit `414ce456e14ae8546f77d9356d2c4d955c5bb7f1`. This package integrates
+winning submission `b4181149c8884ddb`; its byte-exact submitted source has SHA256
+`4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909`.
+
+The kernel is specialized for the packed QSA decode tensors captured from
+`RadixArk/Qwen3.8-Flash-Next-NVFP4` on NVIDIA GB10 (SM121):
+
+- BF16 query, key, value, and output with head dimension 256
+- one packed query row per sequence and device-side `cu_seqlens`
+- 12 query heads per KV head: TP1 uses 24Q/2KV and TP2 uses 12Q/1KV
+- all query-row counts in the validated `1 <= bs <= 128` envelope
+- `max_seqlen_k` capacity up to 2055 and captured logical selected-KV lengths
+  up to 2051 rows per sequence
+
+The implementation groups the 12 query heads that share one KV head into one
+CTA, uses BF16 tensor-core QK/PV products with FP32 online-softmax state, and
+splits long KV rows across multiple CTAs. The last arriving split performs a
+stable FP32 merge and resets its device counter in the same launch. A
+host-visible shape/topology policy selects the two measured schedules, while
+the live device `cu_seqlens_k` selects one, two, four, or eight active splits
+without a host synchronization.
+
+SM121 dispatch checks the exact Qwen3.8 contract and routes directly to this
+kernel; it is the only packed-QSA attention implementation added by this PR.
+The KDA replay passed all 15 TP1/TP2 production tensors on two independent GB10
+GPUs, and the final source passed 150,000 consecutive launches with all
+counters returning to zero.
+
+After adaptation into SGLang, the packaged kernel passed the same 15/15 replay
+with exactly one CUDA activity per row and a 2.0702x all-shape geomean over the
+generic Triton fallback (1.6951x large, 2.3653x small). On one DGX Spark running
+the full TP1 NVFP4 model with NEXTN, three-round low-concurrency serving A/B
+improved total token throughput by 4.45% at concurrency 1 and 4.00% at
+concurrency 4. A 50-example, five-shot GSM8K A/B with a 2048-token output limit
+scored 49/50 for both Triton and KDA, with the same single failed example.
+
+An additional synthetic GB10 sweep covers both TP topologies, every batch size
+from 1 through 16, and short plus saturated KV rows. All 64 cases passed; the
+maximum relative L2 against the original correct Triton implementation was
+0.002422, and speedup ranged from 1.41x to 5.09x.
+
+A follow-up extended-batch sweep covers both TP topologies, batch sizes
+17/24/32/48/64/96/128, and short, medium, plus saturated KV rows. All 42 cases
+passed with maximum relative L2 0.002410. Geomean speedup was 4.48x, the slowest
+case still improved by 1.58x, and no case regressed. The packaged scratch space
+is therefore sized for the largest tested batch, 128.

python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/__init__.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_HEAD_TOPOLOGIES = frozenset({(12, 1), (24, 2)})
+# Largest batch qualified by the extended GB10 baseline sweep.
+_MAX_BATCH = 128
+_MAX_SELECTED_KV = 2055
+_logged_fast_path = False
+
+
+def can_use_qwen38_qsa_sm121(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+) -> bool:
+    """Return whether this call matches the captured Qwen3.8 SM121 contract."""
+    if not q.is_cuda or q.ndim != 3 or q.dtype != torch.bfloat16:
+        return False
+    batch, num_q_heads, head_dim = q.shape
+    if not (0 < batch <= _MAX_BATCH) or head_dim != 256:
+        return False
+    if k.ndim != 3 or v.shape != k.shape or k.dtype != q.dtype or v.dtype != q.dtype:
+        return False
+    num_kv_heads = k.shape[1]
+    if (num_q_heads, num_kv_heads) not in _SUPPORTED_HEAD_TOPOLOGIES:
+        return False
+    if k.shape[2] != head_dim or not (0 < max_seqlen_k <= _MAX_SELECTED_KV):
+        return False
+    if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous():
+        return False
+    if q.device != k.device or q.device != v.device:
+        return False
+    if cu_seqlens_q.device != q.device or cu_seqlens_k.device != q.device:
+        return False
+    if cu_seqlens_q.dtype != torch.int32 or cu_seqlens_k.dtype != torch.int32:
+        return False
+    if cu_seqlens_q.ndim != 1 or cu_seqlens_k.ndim != 1:
+        return False
+    if not cu_seqlens_q.is_contiguous() or not cu_seqlens_k.is_contiguous():
+        return False
+    if cu_seqlens_q.numel() != batch + 1 or cu_seqlens_k.numel() != batch + 1:
+        return False
+    properties = torch.cuda.get_device_properties(q.device)
+    return (properties.major, properties.minor) == (12, 1)
+
+
+def qwen38_qsa_sm121(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """Run the KDA-generated Qwen3.8 packed QSA decode kernel."""
+    global _logged_fast_path
+    if not can_use_qwen38_qsa_sm121(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k):
+        raise ValueError("unsupported call for the KDA Qwen3.8 SM121 QSA kernel")
+
+    from .kernel import qwen38_qsa_sm121 as run_kernel
+
+    if not _logged_fast_path:
+        logger.info(
+            "Using the Codex/Kimi K3 KDA Qwen3.8 QSA kernel on SM121 "
+            "(radixark/KDA-1.5#4, submission b4181149c8884ddb)"
+        )
+        _logged_fast_path = True
+    return run_kernel(q, k, v, cu_seqlens_q, cu_seqlens_k, softmax_scale)
+
+
+__all__ = ["can_use_qwen38_qsa_sm121", "qwen38_qsa_sm121"]

python/sglang/kernels/kda_kernels/qwen38_qsa_sm121/kernel.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# KDA-1.5 submission b4181149c8884ddb (https://github.com/radixark/KDA-1.5/pull/4);
+# source SHA256 4f9977f88abfea4393a2add3a2c9255699f7e13b981dbc1a976b024b3b00e909.
+"""Shape-specialized Qwen3.8 packed QSA decode kernel for SM121."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _qsa_split_kernel(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    out_ptr,
+    cu_q_ptr,
+    cu_k_ptr,
+    partial_max_ptr,
+    partial_sum_ptr,
+    partial_acc_ptr,
+    counter_ptr,
+    softmax_scale,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_KV_HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_KV: tl.constexpr,
+    MAX_SPLITS: tl.constexpr,
+    q_stride_t: tl.constexpr,
+    q_stride_h: tl.constexpr,
+    k_stride_t: tl.constexpr,
+    k_stride_h: tl.constexpr,
+    v_stride_t: tl.constexpr,
+    v_stride_h: tl.constexpr,
+    out_stride_t: tl.constexpr,
+    out_stride_h: tl.constexpr,
+):
+    sequence_idx = tl.program_id(0)
+    split_program = tl.program_id(1)
+    kv_head_idx = split_program // MAX_SPLITS
+    split_idx = split_program - kv_head_idx * MAX_SPLITS
+    slot = sequence_idx * NUM_KV_HEADS + kv_head_idx
+    queries_per_kv = NUM_Q_HEADS // NUM_KV_HEADS
+
+    query_idx = tl.load(cu_q_ptr + sequence_idx)
+    kv_begin = tl.load(cu_k_ptr + sequence_idx)
+    kv_end = tl.load(cu_k_ptr + sequence_idx + 1)
+    kv_count = kv_end - kv_begin
+    tile_count = tl.cdiv(kv_count, BLOCK_KV)
+
+    # The split count depends only on live device metadata and the static
+    # launch geometry, so CUDA Graph replay needs no host readback.
+    batch = tl.num_programs(0)
+    n_splits = 1
+    if tile_count >= 1536 // BLOCK_KV:
+        n_splits = 2
+    elif tile_count >= 512 // BLOCK_KV and batch * NUM_KV_HEADS <= 4:
+        n_splits = 4
+    if batch == 1:
+        if tile_count >= 1024 // BLOCK_KV:
+            n_splits = 8
+        elif tile_count >= 512 // BLOCK_KV:
+            n_splits = 4
+        elif tile_count >= 256 // BLOCK_KV:
+            n_splits = 2
+    if split_idx >= n_splits:
+        return
+
+    tile_lo = (tile_count * split_idx) // n_splits
+    tile_hi = (tile_count * (split_idx + 1)) // n_splits
+    kv_start = kv_begin + tile_lo * BLOCK_KV
+    kv_stop = tl.minimum(kv_begin + tile_hi * BLOCK_KV, kv_end)
+
+    m = tl.arange(0, BLOCK_M)
+    d = tl.arange(0, HEAD_DIM)
+    q_head = kv_head_idx * queries_per_kv + m
+    q_mask = m < queries_per_kv
+    query = tl.load(
+        q_ptr + query_idx * q_stride_t + q_head[:, None] * q_stride_h + d[None, :],
+        mask=q_mask[:, None],
+        other=0.0,
+    )
+
+    n0 = tl.arange(0, BLOCK_KV)
+    k_rows = (
+        k_ptr
+        + kv_head_idx * k_stride_h
+        + kv_start.to(tl.int64) * k_stride_t
+        + n0[:, None] * k_stride_t
+    )
+    v_rows = (
+        v_ptr
+        + kv_head_idx * v_stride_h
+        + kv_start.to(tl.int64) * v_stride_t
+        + n0[:, None] * v_stride_t
+    )
+
+    running_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    running_sum = tl.zeros([BLOCK_M], tl.float32)
+    accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
+
+    kv_len = kv_stop - kv_start
+    full_end = kv_start + (kv_len // BLOCK_KV) * BLOCK_KV
+    for block_start in range(kv_start, full_end, BLOCK_KV):
+        keys = tl.load(k_rows + d[None, :])
+        scores = tl.dot(query, tl.trans(keys)) * softmax_scale
+        new_max = tl.maximum(running_max, tl.max(scores, axis=1))
+        old_scale = tl.exp(running_max - new_max)
+        probabilities = tl.exp(scores - new_max[:, None])
+        running_sum = running_sum * old_scale + tl.sum(probabilities, axis=1)
+        values = tl.load(v_rows + d[None, :])
+        accumulator = accumulator * old_scale[:, None] + tl.dot(
+            probabilities.to(tl.bfloat16), values
+        )
+        running_max = new_max
+        k_rows += BLOCK_KV * k_stride_t
+        v_rows += BLOCK_KV * v_stride_t
+
+    if full_end < kv_stop:
+        n = full_end + n0
+        n_mask = n < kv_stop
+        keys = tl.load(k_rows + d[None, :], mask=n_mask[:, None], other=0.0)
+        scores = tl.dot(query, tl.trans(keys)) * softmax_scale
+        scores = tl.where(n_mask[None, :], scores, -float("inf"))
+        new_max = tl.maximum(running_max, tl.max(scores, axis=1))
+        old_scale = tl.exp(running_max - new_max)
+        probabilities = tl.exp(scores - new_max[:, None])
+        running_sum = running_sum * old_scale + tl.sum(probabilities, axis=1)
+        values = tl.load(v_rows + d[None, :], mask=n_mask[:, None], other=0.0)
+        accumulator = accumulator * old_scale[:, None] + tl.dot(
+            probabilities.to(tl.bfloat16), values
+        )
+        running_max = new_max
+
+    if n_splits == 1:
+        output = accumulator / tl.where(running_sum > 0.0, running_sum, 1.0)[:, None]
+        tl.store(
+            out_ptr
+            + query_idx * out_stride_t
+            + q_head[:, None] * out_stride_h
+            + d[None, :],
+            output.to(out_ptr.dtype.element_ty),
+            mask=q_mask[:, None],
+        )
+        return
+
+    partial_row = (slot * MAX_SPLITS + split_idx) * BLOCK_M + m
+    tl.store(partial_max_ptr + partial_row, running_max)
+    tl.store(partial_sum_ptr + partial_row, running_sum)
+    tl.store(
+        partial_acc_ptr + partial_row[:, None] * HEAD_DIM + d[None, :],
+        accumulator,
+    )
+    tl.debug_barrier()
+    arrival = tl.atomic_add(counter_ptr + slot, 1, sem="acq_rel", scope="gpu")
+    if arrival == n_splits - 1:
+        merged_max = tl.full([BLOCK_M], -float("inf"), tl.float32)
+        for j in tl.static_range(MAX_SPLITS):
+            j_ok = (j < n_splits) & (m < BLOCK_M)
+            row = (slot * MAX_SPLITS + j) * BLOCK_M + m
+            mj = tl.load(partial_max_ptr + row, mask=j_ok, other=-float("inf"))
+            merged_max = tl.maximum(merged_max, mj)
+        merged_sum = tl.zeros([BLOCK_M], tl.float32)
+        merged_acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
+        for j in tl.static_range(MAX_SPLITS):
+            j_ok = (j < n_splits) & (m < BLOCK_M)
+            row = (slot * MAX_SPLITS + j) * BLOCK_M + m
+            mj = tl.load(partial_max_ptr + row, mask=j_ok, other=-float("inf"))
+            lj = tl.load(partial_sum_ptr + row, mask=j_ok, other=0.0)
+            weight = tl.exp(mj - merged_max)
+            merged_sum += weight * lj
+            partial = tl.load(
+                partial_acc_ptr + row[:, None] * HEAD_DIM + d[None, :],
+                mask=j_ok[:, None],
+                other=0.0,
+            )
+            merged_acc += weight[:, None] * partial
+        output = merged_acc / tl.where(merged_sum > 0.0, merged_sum, 1.0)[:, None]
+        tl.store(
+            out_ptr
+            + query_idx * out_stride_t
+            + q_head[:, None] * out_stride_h
+            + d[None, :],
+            output.to(out_ptr.dtype.element_ty),
+            mask=q_mask[:, None],
+        )
+        tl.atomic_xchg(counter_ptr + slot, 0, sem="release", scope="gpu")
+
+
+# The worst-case TP1 scratch allocation at the qualified limit is 32.3 MiB.
+_MAX_BATCH = 128
+_MAX_KV_HEADS = 2
+_BLOCK_M = 16
+_MAX_SPLITS = 8
+_MAX_SLOTS = _MAX_BATCH * _MAX_KV_HEADS
+_HEAD_DIM = 256
+_scratch: dict[int, tuple[torch.Tensor, ...]] = {}
+
+
+def _get_scratch(device: torch.device) -> tuple[torch.Tensor, ...]:
+    device_index = (
+        device.index if device.index is not None else torch.cuda.current_device()
+    )
+    scratch = _scratch.get(device_index)
+    if scratch is None:
+        rows = _MAX_SLOTS * _MAX_SPLITS * _BLOCK_M
+        scratch = (
+            torch.empty(rows, dtype=torch.float32, device=device),
+            torch.empty(rows, dtype=torch.float32, device=device),
+            torch.empty(rows * _HEAD_DIM, dtype=torch.float32, device=device),
+            torch.zeros(_MAX_SLOTS, dtype=torch.int32, device=device),
+        )
+        _scratch[device_index] = scratch
+    return scratch
+
+
+def qwen38_qsa_sm121(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """Run the KDA-generated SM121 Qwen3.8 QSA kernel."""
+    output = torch.empty_like(q)
+    partial_max, partial_sum, partial_acc, counters = _get_scratch(q.device)
+    batch, num_q_heads, head_dim = q.shape
+    num_kv_heads = k.shape[1]
+
+    # Measured schedule. The TP1 q_rows=4 shape stays on BK64 on purpose:
+    # its short and saturated rows cannot be told apart from the host.
+    use_bk32 = (num_kv_heads == 1 and batch < 12) or (num_kv_heads == 2 and batch < 4)
+    block_kv = 32 if use_bk32 else 64
+    stages = 3 if use_bk32 else 2
+    _qsa_split_kernel[(batch, num_kv_heads * _MAX_SPLITS)](
+        q,
+        k,
+        v,
+        output,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        partial_max,
+        partial_sum,
+        partial_acc,
+        counters,
+        softmax_scale,
+        NUM_Q_HEADS=num_q_heads,
+        NUM_KV_HEADS=num_kv_heads,
+        HEAD_DIM=head_dim,
+        BLOCK_M=_BLOCK_M,
+        BLOCK_KV=block_kv,
+        MAX_SPLITS=_MAX_SPLITS,
+        q_stride_t=q.stride(0),
+        q_stride_h=q.stride(1),
+        k_stride_t=k.stride(0),
+        k_stride_h=k.stride(1),
+        v_stride_t=v.stride(0),
+        v_stride_h=v.stride(1),
+        out_stride_t=output.stride(0),
+        out_stride_h=output.stride(1),
+        num_warps=4,
+        num_stages=stages,
+    )
+    return output

python/sglang/kernels/ops/attention/__init__.py

@@ -6,8 +6,21 @@ The Triton kernels migrated here live in this package
 KV-cache index/write kernels went to the ``kvcache`` group instead.
 """
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from sglang.kernels.registry import register_kernel
-from sglang.kernels.spec import KernelBackend, KernelSpec
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+)
+
+if TYPE_CHECKING:
+    import torch
 
 # (module, public_fn) migrated from layers/attention/triton_ops + model_executor.
 _TRITON_KERNELS = [
@@ -45,7 +58,78 @@ for _mod, _fn in _TRITON_KERNELS:
     )
 del _mod, _fn
 
-__all__ = []
+register_kernel(
+    KernelSpec(
+        op="attention.kda_qwen38_qsa_sm121",
+        backend=KernelBackend.TRITON,
+        target=("sglang.kernels.kda_kernels.qwen38_qsa_sm121:qwen38_qsa_sm121"),
+        capabilities=frozenset(
+            {CapabilityRequirement.cuda(min_sm=(12, 1), max_sm=(12, 1))}
+        ),
+        format_signature=FormatSignature(
+            supported_dtypes=("bfloat16",),
+            description=(
+                "Qwen3.8 packed QSA decode: D=256, 12:1 GQA, 1 <= q_rows <= 128"
+            ),
+        ),
+        description=(
+            "SM121 Qwen3.8 QSA decode optimized by Codex/Kimi K3 through KDA-1.5."
+        ),
+    )
+)
+
+
+def can_use_kda_qwen38_qsa_sm121(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+) -> bool:
+    """Check the exact E2E-qualified Qwen3.8/SM121 QSA contract."""
+    from sglang.kernels.kda_kernels.qwen38_qsa_sm121 import (
+        can_use_qwen38_qsa_sm121,
+    )
+
+    return can_use_qwen38_qsa_sm121(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k)
+
+
+def qwen38_qsa_sm121_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int = 1,
+    max_seqlen_k: int = 0,
+    softmax_scale: float = 1.0,
+    causal: bool = True,
+    **_: object,
+) -> torch.Tensor:
+    """Run the only SM121 packed-QSA kernel for its qualified contract."""
+    del causal
+    if max_seqlen_q != 1:
+        raise ValueError(f"QSA requires max_seqlen_q=1, got {max_seqlen_q}")
+    if not can_use_kda_qwen38_qsa_sm121(
+        q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_k
+    ):
+        raise ValueError(
+            "unsupported SM121 QSA call: expected BF16 D=256, 12:1 GQA, "
+            "TP1 24Q/2KV or TP2 12Q/1KV, bs<=128, and selected KV<=2055"
+        )
+    return get_kernel("attention.kda_qwen38_qsa_sm121", KernelBackend.TRITON)(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_k,
+        softmax_scale,
+    )
+
+
+__all__ = ["can_use_kda_qwen38_qsa_sm121", "qwen38_qsa_sm121_varlen"]
 
 
 # Vendored linear-attention (flash-linear-attention port) kernels relocated

python/sglang/kernels/ops/attention/qsa_indexer.py

@@ -0,0 +1,176 @@
+"""Fused QSA (Qwen4-Exp sparse attention) indexer-prep kernels.
+
+``qsa_index_q_norm_rope_store`` fuses, per token, the eager chain
+split -> GemmaRMSNorm(index q) -> MRoPE(index q) -> raw-K store ->
+RoPE-position store into one kernel launch.
+
+``qsa_index_k_compress_store`` fuses, per completed compress group, the eager
+chain gather -> fp32 mean -> GemmaRMSNorm -> MRoPE(group-start position) ->
+compressed-cache store into one kernel launch.
+
+Both kernels reproduce the eager numerics step by step (fp32 norm reduction,
+per-op rounding to the storage dtype during RoPE, fp32 group mean rounded to
+the storage dtype before the norm). Outputs are bit-comparable to the eager
+indexer path: the eager RMSNorm (flashinfer's CuTe DSL kernel) reduces sums of
+squares in an order that cannot be reproduced exactly, so a small fraction of
+rows (~1 in 30k) may flip by one bf16 ulp on a rounding boundary.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_qsa_indexer_module(
+    dtype: torch.dtype, head_dim: int, is_neox_style: bool
+) -> Module:
+    """Compile and cache the JIT QSA indexer module for one specialisation."""
+    if dtype not in (torch.bfloat16, torch.float16):
+        raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
+    if head_dim not in (64, 128, 256):
+        raise RuntimeError(
+            f"Unsupported index head_dim {head_dim}. Supported: 64, 128, 256"
+        )
+    args = make_cpp_args(dtype, head_dim, is_neox_style, is_arch_support_pdl())
+    return load_jit(
+        "qsa_indexer",
+        *args,
+        cuda_files=["attention/qsa_indexer.cuh"],
+        cuda_wrappers=[
+            ("q_prep", f"qsa_index_q_prep<{args}>"),
+            ("k_compress", f"qsa_index_k_compress<{args}>"),
+        ],
+    )
+
+
+def qsa_index_q_norm_rope_store(
+    qk: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    axis_map: torch.Tensor,
+    weight: torch.Tensor,
+    cache_loc: torch.Tensor,
+    key_state_buffer: torch.Tensor,
+    rope_position_buffer: torch.Tensor,
+    num_q_heads: int,
+    rotary_dim: int,
+    eps: float,
+    is_neox_style: bool,
+    q_heads_padded: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Per-token fused index-Q prep.
+
+    Parameters
+    ----------
+    qk              : CUDA bf16/fp16 [tokens, (num_q_heads + 1) * head_dim],
+                      contiguous fused index Q/K projection output
+    positions       : CUDA int64 [tokens] or [3, tokens] RoPE positions
+                      (a strided trailing slice is accepted)
+    cos_sin_cache   : CUDA fp32 [capacity, rotary_dim] RoPE cache
+    axis_map        : CUDA int32 [rotary_dim // 2] position-axis per pair index
+    weight          : [head_dim] gemma norm weight (kernel applies 1 + w)
+    cache_loc       : CUDA int64 [tokens] state slots of each token
+    key_state_buffer : CUDA [slots, head_dim] raw-K state buffer (written)
+    rope_position_buffer : CUDA int64 [slots, 3] position buffer (written)
+    num_q_heads     : number of index query heads
+    rotary_dim      : rotated prefix of each head row
+    eps             : RMSNorm epsilon
+    is_neox_style   : NeoX (True) or GPT-J (False) RoPE pairing
+    q_heads_padded  : output head count; heads >= num_q_heads are zero-filled
+                      (defaults to num_q_heads)
+
+    Returns
+    -------
+    CUDA tensor [tokens, q_heads_padded, head_dim]: normed + rotated index Q.
+    """
+    num_tokens = qk.shape[0]
+    head_dim = weight.shape[0]
+    if q_heads_padded is None:
+        q_heads_padded = num_q_heads
+    if positions.ndim == 1:
+        positions = positions.unsqueeze(0)
+    q_out = torch.empty(
+        (num_tokens, q_heads_padded, head_dim), dtype=qk.dtype, device=qk.device
+    )
+    module = _jit_qsa_indexer_module(qk.dtype, head_dim, is_neox_style)
+    module.q_prep(
+        qk,
+        q_out,
+        weight,
+        cos_sin_cache,
+        axis_map,
+        positions,
+        positions.shape[0],
+        cache_loc,
+        key_state_buffer,
+        rope_position_buffer,
+        num_q_heads,
+        rotary_dim,
+        eps,
+    )
+    return q_out
+
+
+def qsa_index_k_compress_store(
+    key_state_buffer: torch.Tensor,
+    group_locs: torch.Tensor,
+    rope_position_buffer: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    axis_map: torch.Tensor,
+    weight: torch.Tensor,
+    write_locs: torch.Tensor,
+    compressed_k_buffer: torch.Tensor,
+    compress_ratio: int,
+    rotary_dim: int,
+    eps: float,
+    is_neox_style: bool,
+) -> None:
+    """
+    Per-group compressed-K prep (in-place store into ``compressed_k_buffer``).
+
+    Parameters
+    ----------
+    key_state_buffer : CUDA bf16/fp16 [slots, head_dim] raw-K state buffer
+    group_locs       : CUDA int32 [groups, compress_ratio] state slots of each
+                       completed group, group-start slot in column 0 (its RoPE
+                       position rotates the group; other columns any order)
+    rope_position_buffer : CUDA int64 [slots, 3] per-slot RoPE coordinates
+    cos_sin_cache    : CUDA fp32 [capacity, rotary_dim] RoPE cache
+    axis_map         : CUDA int32 [rotary_dim // 2] position-axis per pair index
+    weight           : [head_dim] gemma norm weight (kernel applies 1 + w)
+    write_locs       : CUDA int32 [groups] compressed-cache slots to write
+    compressed_k_buffer : CUDA [compressed_slots, head_dim] (written)
+    compress_ratio   : raw keys per compressed key
+    rotary_dim       : rotated prefix of each head row
+    eps              : RMSNorm epsilon
+    is_neox_style    : NeoX (True) or GPT-J (False) RoPE pairing
+    """
+    head_dim = weight.shape[0]
+    module = _jit_qsa_indexer_module(key_state_buffer.dtype, head_dim, is_neox_style)
+    module.k_compress(
+        key_state_buffer,
+        group_locs,
+        rope_position_buffer,
+        cos_sin_cache,
+        axis_map,
+        weight,
+        write_locs,
+        compressed_k_buffer,
+        compress_ratio,
+        rotary_dim,
+        eps,
+    )

python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py

@@ -171,6 +171,8 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
     a,
     mixed_qkvz,
     mixed_ba,
+    qkvz_row_stride,
+    ba_row_stride,
     NUM_HEADS_QK: tl.constexpr,
     NUM_HEADS_V: tl.constexpr,
     HEAD_QK: tl.constexpr,
@@ -185,19 +187,19 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
     TOTAL_Q: tl.constexpr = NUM_HEADS_QK * HEAD_QK
     TOTAL_K: tl.constexpr = NUM_HEADS_QK * HEAD_QK
     TOTAL_V: tl.constexpr = NUM_HEADS_V * HEAD_V
-    TOTAL_QKVZ: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V + TOTAL_V
-    TOTAL_BA: tl.constexpr = NUM_HEADS_V * 2
 
     # ── Output dimensions ──
     QKV_DIM_T: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V
 
     # ── Read from contiguous input ──
     # q for head group i_qk: in the all_q region, offset i_qk * HEAD_QK
-    blk_q_ptr = mixed_qkvz + i_bs * TOTAL_QKVZ + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
+    blk_q_ptr = (
+        mixed_qkvz + i_bs * qkvz_row_stride + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
+    )
     # k for head group i_qk: in the all_k region
     blk_k_ptr = (
         mixed_qkvz
-        + i_bs * TOTAL_QKVZ
+        + i_bs * qkvz_row_stride
         + TOTAL_Q
         + i_qk * HEAD_QK
         + tl.arange(0, HEAD_QK)
@@ -209,7 +211,11 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
     # vector access. V_POW2 arrives as a wrapper-computed constexpr so the
     # dead branch is pruned before tl.arange validation.
     v_ld_base = (
-        mixed_qkvz + i_bs * TOTAL_QKVZ + TOTAL_Q + TOTAL_K + i_qk * V_PER_GROUP * HEAD_V
+        mixed_qkvz
+        + i_bs * qkvz_row_stride
+        + TOTAL_Q
+        + TOTAL_K
+        + i_qk * V_PER_GROUP * HEAD_V
     )
     z_ld_base = v_ld_base + TOTAL_V
 
@@ -250,12 +256,14 @@ def fused_qkvzba_split_reshape_cat_contiguous_kernel(
 
     # ── b and a from contiguous [all_b | all_a] ──
     for i in tl.static_range(V_PER_GROUP):
-        blk_b_ptr = mixed_ba + i_bs * TOTAL_BA + i_qk * V_PER_GROUP + i
+        blk_b_ptr = mixed_ba + i_bs * ba_row_stride + i_qk * V_PER_GROUP + i
         blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
         tl.store(blk_b_st_ptr, tl.load(blk_b_ptr))
 
     for i in tl.static_range(V_PER_GROUP):
-        blk_a_ptr = mixed_ba + i_bs * TOTAL_BA + NUM_HEADS_V + i_qk * V_PER_GROUP + i
+        blk_a_ptr = (
+            mixed_ba + i_bs * ba_row_stride + NUM_HEADS_V + i_qk * V_PER_GROUP + i
+        )
         blk_a_st_ptr = a + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
         tl.store(blk_a_st_ptr, tl.load(blk_a_ptr))
 
@@ -319,6 +327,8 @@ def fused_qkvzba_split_reshape_cat_contiguous(
         a,
         mixed_qkvz,
         mixed_ba,
+        mixed_qkvz.stride(0),
+        mixed_ba.stride(0),
         num_heads_qk,
         num_heads_v,
         head_qk,

python/sglang/kernels/ops/elementwise/fast_topk.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_FAST_TOPK_SUPPORTED_K = (512, 2048)
+
+
+@cache_once
+def _jit_fast_topk_module(topk: int) -> Module:
+    """Compile and cache the JIT fast top-k module for a given k."""
+    # Checks on the compile key live here, not in `fast_topk`: `cache_once`
+    # keys on `topk`, so this runs once per specialisation.
+    if topk not in _FAST_TOPK_SUPPORTED_K:
+        raise RuntimeError(
+            f"Unsupported topk {topk}. Supported: {_FAST_TOPK_SUPPORTED_K}"
+        )
+    args = make_cpp_args(topk, is_arch_support_pdl())
+    return load_jit(
+        "fast_topk",
+        *args,
+        cuda_files=["elementwise/fast_topk.cuh"],
+        cuda_wrappers=[("fast_topk", f"FastTopKKernel<{args}>::run")],
+    )
+
+
+def fast_topk(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    topk: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Per-row top-k selection over a fp32 score matrix.
+
+    Row b selects the `topk` largest values in
+    ``score[b, row_starts[b] : row_starts[b] + lengths[b]]`` and returns their
+    indices relative to ``row_starts[b]``. Slots beyond ``lengths[b]`` are -1.
+    Output order within a row is unspecified (atomic collection order).
+
+    Parameters
+    ----------
+    score      : CUDA fp32 tensor [B, L]
+    lengths    : CUDA int32 tensor [B]
+    topk       : number of indices per row; 512 or 2048
+    row_starts : optional CUDA int32 tensor [B]; defaults to zeros
+
+    Returns
+    -------
+    CUDA int32 tensor [B, topk]
+    """
+    batch = score.shape[0]
+    if row_starts is None:
+        row_starts = torch.zeros(batch, dtype=torch.int32, device=score.device)
+    indices = score.new_empty((batch, topk), dtype=torch.int32)
+
+    module = _jit_fast_topk_module(topk)
+    module.fast_topk(score, row_starts, indices, lengths)
+    return indices

python/sglang/kernels/ops/elementwise/hc_combine.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_hc_combine_module(
+    hc_count: int, hidden_size: int, dtype: torch.dtype
+) -> Module:
+    """Compile and cache the JIT HC combine module for a given shape/dtype."""
+    # Validation lives here rather than in `hc_combine`,
+    # so `cache_once` runs it once per (hc_count, hidden_size, dtype), not per call.
+    if dtype not in (torch.bfloat16, torch.float16):
+        raise RuntimeError(f"Unsupported dtype {dtype}. Supported: bfloat16, float16")
+    if hidden_size <= 0 or hidden_size % 8 != 0:
+        raise RuntimeError(
+            f"Unsupported hidden_size {hidden_size}. Must be a multiple of 8."
+        )
+    if hc_count <= 0 or (hc_count * hidden_size) % 2048 != 0:
+        raise RuntimeError(
+            f"Unsupported hc_count * hidden_size {hc_count * hidden_size}. "
+            "Must be a multiple of 2048."
+        )
+    args = make_cpp_args(hc_count, hidden_size, is_arch_support_pdl(), dtype)
+    return load_jit(
+        "hc_combine",
+        *args,
+        cuda_files=["elementwise/hc_combine.cuh"],
+        cuda_wrappers=[
+            ("hc_combine", f"HcCombineKernel<{args}>::run"),
+            ("hc_combine_split", f"HcCombineSplitKernel<{args}>::run"),
+        ],
+    )
+
+
+def hc_combine(
+    block_output: torch.Tensor,
+    residual: torch.Tensor,
+    normed_residual: torch.Tensor,
+    inject_weight: torch.Tensor,
+    hc_count: int,
+    hidden_size: int,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Fused HyperConnection (gated residual) combine.
+
+        a[m, c] = 2 * sigmoid(dot(normed_residual[m], inject_weight[c]) / hc_count)
+        out[m, c*H + i] = residual[m, c*H + i] + a[m, c] * block_output[m, i]
+
+    Mirrors ``GatedResidual._combine_compute`` in
+    ``sglang.srt.layers.hyperconnection``. All math is accumulated in fp32.
+
+    Supported dtypes: torch.bfloat16, torch.float16.
+
+    Parameters
+    ----------
+    block_output    : CUDA tensor [..., hidden_size]
+    residual        : CUDA tensor [..., hc_count * hidden_size]
+    normed_residual : CUDA tensor, same shape/dtype as residual
+    inject_weight   : CUDA tensor [hc_count, hc_count * hidden_size]
+    hc_count        : number of hyper-connection branches
+    hidden_size     : per-branch hidden size
+    out             : optional pre-allocated output tensor (same shape/dtype as residual)
+
+    Returns
+    -------
+    Combined tensor, same shape/dtype as residual.
+    """
+    y = block_output.reshape(-1, hidden_size)
+    r = residual.reshape(-1, hc_count * hidden_size)
+    n = normed_residual.reshape(-1, hc_count * hidden_size)
+    if out is None:
+        out = torch.empty_like(r)
+    else:
+        out = out.reshape(-1, hc_count * hidden_size)
+
+    module = _jit_hc_combine_module(hc_count, hidden_size, residual.dtype)
+    module.hc_combine(y, r, n, inject_weight, out)
+    return out.reshape(residual.shape)
+
+
+_SPLIT = 8
+_MAX_ROWS = 32
+_partials_cache = {}
+
+
+def _get_partials(hc_count: int, device: torch.device, rows: int) -> torch.Tensor:
+    key = (device, hc_count)
+    buf = _partials_cache.get(key)
+    if buf is None or buf.shape[0] < rows:
+        # The gate kernel writes one slot per (row, split, hc); a buffer shorter
+        # than `rows` is written past its end rather than truncated.
+        buf = torch.empty(
+            (max(rows, _MAX_ROWS), _SPLIT, hc_count),
+            dtype=torch.float32,
+            device=device,
+        )
+        _partials_cache[key] = buf
+    return buf
+
+
+def hc_combine_split(
+    block_output: torch.Tensor,
+    residual: torch.Tensor,
+    normed_residual: torch.Tensor,
+    inject_weight: torch.Tensor,
+    hc_count: int,
+    hidden_size: int,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    y = block_output.reshape(-1, hidden_size)
+    r = residual.reshape(-1, hc_count * hidden_size)
+    n = normed_residual.reshape(-1, hc_count * hidden_size)
+    if out is None:
+        out = torch.empty_like(r)
+    else:
+        out = out.reshape(-1, hc_count * hidden_size)
+    rows = r.shape[0]
+    partials = _get_partials(hc_count, r.device, rows)[:rows]
+    module = _jit_hc_combine_module(hc_count, hidden_size, residual.dtype)
+    module.hc_combine_split(y, r, n, inject_weight, out, partials)
+    return out.reshape(residual.shape)

python/sglang/kernels/ops/elementwise/hc_mix.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import torch
+
+_MAX_ROWS = 32
+_KPAD_ALIGN = 128
+
+
+def pad_lowrank(lowrank: int) -> int:
+    return (lowrank + _KPAD_ALIGN - 1) // _KPAD_ALIGN * _KPAD_ALIGN
+
+
+def hc_mix_shape_ok(hc_count: int, hidden_size: int, lowrank: int) -> bool:
+    k = hc_count * hidden_size
+    return (
+        hidden_size % 8 == 0
+        and lowrank > 0
+        and lowrank % 8 == 0
+        and k % _KPAD_ALIGN == 0
+        and k % 128 == 0
+    )
+
+
+def permute_pad_up_weight(w_up: torch.Tensor, hc_count: int) -> torch.Tensor:
+    k, lowrank = w_up.shape
+    hidden_size = k // hc_count
+    padded = torch.zeros(
+        (k, pad_lowrank(lowrank)), dtype=w_up.dtype, device=w_up.device
+    )
+    padded[:, :lowrank] = w_up.detach()
+    src = torch.arange(k, device=w_up.device)
+    perm = torch.empty(k, dtype=torch.long, device=w_up.device)
+    perm[(src % hidden_size) * hc_count + (src // hidden_size)] = src
+    return padded[perm].contiguous()
+
+
+_scratch_cache = {}
+_tactic_cache = {}
+
+
+def _get_scratch(lowrank: int, dtype: torch.dtype, device: torch.device):
+    key = (device, lowrank, dtype)
+    buf = _scratch_cache.get(key)
+    if buf is None:
+        buf = torch.zeros((_MAX_ROWS, pad_lowrank(lowrank)), dtype=dtype, device=device)
+        _scratch_cache[key] = buf
+    return buf
+
+
+def _get_tactic(m: int, n: int, k: int):
+    key = (m, n, k)
+    tactic = _tactic_cache.get(key)
+    if tactic is None:
+        from sglang.kernels.ops.gemm.dense_bf16_gemm_sm100_splitk_epilogue import (
+            SplitKTactic,
+            default_tactic,
+            validate_tactic,
+        )
+
+        if n <= 512 and k % (128 * 16) == 0:
+            tactic = SplitKTactic(mma_m=64, mma_n=8, split_k=16, ab_stages=6)
+            try:
+                validate_tactic(tactic, m, n, k)
+            except ValueError:
+                tactic = default_tactic(m, n, k)
+        else:
+            tactic = default_tactic(m, n, k)
+        _tactic_cache[key] = tactic
+    return tactic
+
+
+def hc_mix(
+    hyper_input_normed: torch.Tensor,
+    w_down: torch.Tensor,
+    w_up_permuted_padded: torch.Tensor,
+    hc_count: int,
+    hidden_size: int,
+) -> torch.Tensor:
+    from sglang.kernels.ops.gemm.dense_bf16_gemm_sm100_splitk_epilogue import (
+        run_splitk_dense_gate,
+        run_splitk_dense_silu,
+    )
+
+    rows, k = hyper_input_normed.shape
+    lowrank = w_down.shape[0]
+    kpad = w_up_permuted_padded.shape[1]
+    out = torch.empty(
+        (rows, hidden_size),
+        dtype=hyper_input_normed.dtype,
+        device=hyper_input_normed.device,
+    )
+    if rows == 0:
+        return out
+    inv_hc = 1.0 / hc_count
+    t_pad = _get_scratch(lowrank, hyper_input_normed.dtype, hyper_input_normed.device)
+    for row_start in range(0, rows, _MAX_ROWS):
+        row_end = min(row_start + _MAX_ROWS, rows)
+        m = row_end - row_start
+        chunk = hyper_input_normed[row_start:row_end]
+        run_splitk_dense_silu(
+            chunk,
+            w_down.T,
+            t_pad[:m, :lowrank],
+            True,
+            _get_tactic(m, lowrank, k),
+            inv_hc,
+        )
+        run_splitk_dense_gate(
+            t_pad[:m],
+            w_up_permuted_padded.T,
+            chunk,
+            out[row_start:row_end],
+            True,
+            _get_tactic(m, k, kpad),
+            inv_hc,
+            hc_count,
+        )
+    return out

python/sglang/kernels/ops/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)

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,
     )

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

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

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

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

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:

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

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)

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

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",

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)

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(

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",

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",

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)

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,

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)
 
 

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)

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

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,

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)

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",
+]

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",
+]

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",
+]

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,
+    )

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",
+]

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",
+]

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",
+]

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",
+]

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",
+]

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

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,
+}

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
+    }
+}

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
+    }
+}

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
+    }
+}

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

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",

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
 

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,

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

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
+        )

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

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 "

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.

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:

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

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

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

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]

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",

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,

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
 

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

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()

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,

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)

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,

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