[Feat] QSA: Add cache sharding for Qwen3.8 hybrid models

python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh

@@ -277,11 +277,14 @@ struct QsaIndexKCompressParams {
   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]
+  const int32_t* owner_mask;            // [groups]
+  const int32_t* local_write_locs;      // [groups]
   void* compressed_k_buffer;            // [compressed_slots, kHeadDim]
   int32_t compress_ratio;
   int32_t rotary_dim;
   int32_t num_groups;
+  int32_t state_slots;
+  int32_t compressed_slots;
   float eps;
 };
 
@@ -306,7 +309,24 @@ __launch_bounds__(128) void qsa_index_k_compress_kernel(const QsaIndexKCompressP
 
   device::PDLWaitPrimary<kUsePDL>();
 
+  // Shape-derived eager launches retain padded and non-owner groups. Skip
+  // those before touching source state; their locations need not be valid.
+  if (params.owner_mask[group] == 0) {
+    device::PDLTriggerSecondary<kUsePDL>();
+    return;
+  }
   const int32_t* locs = params.group_locs + group * params.compress_ratio;
+  const int32_t write_loc = params.local_write_locs[group];
+  if (write_loc < 0 || write_loc >= params.compressed_slots) {
+    device::PDLTriggerSecondary<kUsePDL>();
+    return;
+  }
+  for (int32_t r = 0; r < params.compress_ratio; ++r) {
+    if (locs[r] < 0 || locs[r] >= params.state_slots) {
+      device::PDLTriggerSecondary<kUsePDL>();
+      return;
+    }
+  }
   const int32_t loc0 = locs[0];
 
   // fp32 mean over the group, rounded to the storage dtype exactly like
@@ -361,8 +381,7 @@ __launch_bounds__(128) void qsa_index_k_compress_kernel(const QsaIndexKCompressP
     pos[a] = params.rope_position_buffer[static_cast<int64_t>(loc0) * 3 + a];
   }
 
-  TOut* out_row =
-      static_cast<TOut*>(params.compressed_k_buffer) + static_cast<int64_t>(params.write_locs[group]) * kHeadDim;
+  TOut* out_row = static_cast<TOut*>(params.compressed_k_buffer) + static_cast<int64_t>(write_loc) * kHeadDim;
   qsa_mrope_apply<T, TOut, kHeadDim, kIsNeox>(
       smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim);
 
@@ -453,7 +472,8 @@ void qsa_index_k_compress(
     tvm::ffi::TensorView cos_sin_cache,
     tvm::ffi::TensorView axis_map,
     tvm::ffi::TensorView weight,
-    tvm::ffi::TensorView write_locs,
+    tvm::ffi::TensorView owner_mask,
+    tvm::ffi::TensorView local_write_locs,
     tvm::ffi::TensorView compressed_k_buffer,
     int64_t compress_ratio,
     int64_t rotary_dim,
@@ -472,7 +492,8 @@ void qsa_index_k_compress(
   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);
+  TensorMatcher({groups}).with_dtype<int32_t>().with_device(device).verify(owner_mask);
+  TensorMatcher({groups}).with_dtype<int32_t>().with_device(device).verify(local_write_locs);
   auto compressed_slots = SymbolicSize{"compressed_slots"};
   TensorMatcher({compressed_slots, D}).with_dtype<TOut>().with_device(device).verify(compressed_k_buffer);
 
@@ -490,11 +511,14 @@ void qsa_index_k_compress(
       .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()),
+      .owner_mask = static_cast<const int32_t*>(owner_mask.data_ptr()),
+      .local_write_locs = static_cast<const int32_t*>(local_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),
+      .state_slots = static_cast<int32_t>(slots.unwrap()),
+      .compressed_slots = static_cast<int32_t>(compressed_slots.unwrap()),
       .eps = eps,
   };
   LaunchKernel(static_cast<uint32_t>(div_ceil(num_groups, 4)), 128, device.unwrap())

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

@@ -603,6 +603,8 @@ def _dcp_lse_combine_kernel(
     out_stride_B,
     out_stride_H,
     out_stride_D,
+    out_lse_stride_B,
+    out_lse_stride_H,
     N: tl.constexpr,
     HEAD_DIM: tl.constexpr,
     IS_BASE_E: tl.constexpr,
@@ -635,7 +637,8 @@ def _dcp_lse_combine_kernel(
         )
         lse_max = tl.where(lse_i > lse_max, lse_i, lse_max)
 
-    lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
+    has_any = lse_max != -float("inf")
+    lse_max = tl.where(has_any, lse_max, 0.0)
 
     # Pass 2: accumulate weighted outputs
     weight_sum = tl.zeros([], dtype=tl.float32)
@@ -651,6 +654,7 @@ def _dcp_lse_combine_kernel(
             w = tl.exp(centered)
         else:
             w = tl.exp2(centered)
+        w = tl.where(lse_i == -float("inf"), 0.0, w)
         weight_sum += w
 
         o_offsets = (
@@ -662,7 +666,7 @@ def _dcp_lse_combine_kernel(
         partial_out = tl.load(recv_output_ptr + o_offsets).to(tl.float32)
         acc += partial_out * w
 
-    acc = acc / tl.maximum(weight_sum, 1.0)
+    acc = tl.where(has_any, acc / tl.maximum(weight_sum, 1.0), 0.0)
 
     out_offsets = (
         batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
@@ -671,10 +675,10 @@ def _dcp_lse_combine_kernel(
 
     if RETURN_LSE:
         if IS_BASE_E:
-            global_lse = tl.log(weight_sum) + lse_max
+            global_lse = tl.where(has_any, tl.log(weight_sum) + lse_max, -float("inf"))
         else:
-            global_lse = tl.log2(weight_sum) + lse_max
-        out_lse_offset = batch_idx * recv_lse_stride_B + head_idx * recv_lse_stride_H
+            global_lse = tl.where(has_any, tl.log2(weight_sum) + lse_max, -float("inf"))
+        out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H
         tl.store(out_lse_ptr + out_lse_offset, global_lse)
 
 
@@ -722,6 +726,8 @@ def dcp_lse_combine_triton(
         out.stride(0),
         out.stride(1),
         out.stride(2),
+        out_lse.stride(0) if return_lse else 0,
+        out_lse.stride(1) if return_lse else 0,
         N=N,
         HEAD_DIM=D,
         IS_BASE_E=is_lse_base_on_e,

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

@@ -24,6 +24,8 @@
 from typing import TYPE_CHECKING, Optional
 
 import torch
+import triton
+import triton.language as tl
 
 from sglang.kernels.jit.utils import (
     cache_once,
@@ -36,6 +38,49 @@
     from tvm_ffi.module import Module
 
 
+@triton.jit
+def _qsa_store_compressed_keys_kernel(
+    source,
+    cache,
+    local_locs,
+    owner_mask,
+    WIDTH: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    if tl.load(owner_mask + row):
+        slot = tl.load(local_locs + row).to(tl.int64)
+        offsets = tl.arange(0, BLOCK)
+        values = tl.load(source + row * WIDTH + offsets, mask=offsets < WIDTH)
+        tl.store(cache + slot * WIDTH + offsets, values, mask=offsets < WIDTH)
+
+
+def qsa_store_compressed_keys(
+    compressed_keys: torch.Tensor,
+    cache: torch.Tensor,
+    local_locs: torch.Tensor,
+    owner_mask: torch.Tensor,
+) -> None:
+    """Store graph fallback keys at local slots without dynamic owner filtering.
+
+    Active destinations must be unique. Masked rows may share a safe slot,
+    including slot zero even when another row owns that real local slot.
+    """
+    rows = compressed_keys.shape[0]
+    if rows == 0:
+        return
+    keys = compressed_keys.flatten(1).contiguous().to(cache.dtype)
+    width = keys.shape[1]
+    _qsa_store_compressed_keys_kernel[(rows,)](
+        keys,
+        cache,
+        local_locs,
+        owner_mask,
+        WIDTH=width,
+        BLOCK=triton.next_power_of_2(width),
+    )
+
+
 @cache_once
 def _jit_qsa_indexer_module(
     dtype: torch.dtype, out_dtype: torch.dtype, head_dim: int, is_neox_style: bool
@@ -147,7 +192,8 @@ def qsa_index_k_compress_store(
     cos_sin_cache: torch.Tensor,
     axis_map: torch.Tensor,
     weight: torch.Tensor,
-    write_locs: torch.Tensor,
+    owner_mask: torch.Tensor,
+    local_write_locs: torch.Tensor,
     compressed_k_buffer: torch.Tensor,
     compress_ratio: int,
     rotary_dim: int,
@@ -167,7 +213,9 @@ def qsa_index_k_compress_store(
     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
+    owner_mask       : CUDA int32 [groups], zero skips all reads and writes
+    local_write_locs : CUDA int32 [groups] rank-local compressed-cache slots;
+                       ignored where owner_mask is zero
     compressed_k_buffer : CUDA [compressed_slots, head_dim] (written), ring dtype
                        or fp8 e4m3
     compress_ratio   : raw keys per compressed key
@@ -186,7 +234,8 @@ def qsa_index_k_compress_store(
         cos_sin_cache,
         axis_map,
         weight,
-        write_locs,
+        owner_mask,
+        local_write_locs,
         compressed_k_buffer,
         compress_ratio,
         rotary_dim,

python/sglang/srt/arg_groups/fields/parallel.py

@@ -60,6 +60,15 @@ class Parallel(msgspec.Struct):
             aliases=["--decode-context-parallel-size"],
         ),
     ] = 1
+    qsa_cache_sharding_size: A[
+        int,
+        Arg(
+            help="The number of tensor-parallel ranks over which QSA cache, "
+            "distributed Top-K, and owner attention are sharded. GDN, PLE, and "
+            "the QSA pending ring remain replicated.",
+            aliases=["--state-cp-size"],
+        ),
+    ] = 1
     pp_size: A[
         int,
         Arg(
@@ -391,6 +400,9 @@ class Parallel(msgspec.Struct):
             "zero where decode context parallelism is off."
         )
     )
+    qsa_cache_sharding_rank = Derived(
+        doc="This process's rank in its QSA cache-sharding group."
+    )
     attn_dp_rank = Derived(
         doc=(
             "This process's index in the attention-DP group, computed from "
@@ -431,3 +443,4 @@ class Parallel(msgspec.Struct):
     attn_cp_group = Derived(doc="The attention context-parallel group.")
     shared_experts_tp_group = Derived(doc=("The shared-expert tensor-parallel group."))
     dcp_group = Derived(doc="The decode context-parallel group.")
+    qsa_cache_sharding_group = Derived(doc="The QSA cache-sharding group.")

python/sglang/srt/arg_groups/parallel_hook.py

@@ -285,6 +285,38 @@ def handle_decode_context_parallelism(server_args: Any):
                 f"got --dcp-comm-backend={cfg.dcp_comm_backend}."
             )
 
+    from sglang.srt.layers.attention.qsa.cache_sharding import (
+        validate_qsa_cache_sharding,
+    )
+
+    if cfg.qsa_cache_sharding_size == 1:
+        return
+    if parse_connector_type(cfg.model_path) == ConnectorType.INSTANCE:
+        raise ValueError("QSA cache sharding requires a local model configuration")
+    model_config = model_config_of(server_args)
+    hf_config = model_config.hf_config
+    from sglang.srt.layers.attention.qsa.config import is_qwen_qsa
+
+    validate_qsa_cache_sharding(
+        device=cfg.device,
+        is_qsa=is_qwen_qsa(hf_config),
+        tp_size=derive_attn_tp_size(
+            tp_size=cfg.tp_size,
+            attn_cp_size=resolved_view(server_args).attn_cp_size,
+            attn_dp_size=cfg.attn_dp_size,
+        ),
+        sharding_size=cfg.qsa_cache_sharding_size,
+        num_attention_heads=model_config.num_attention_heads,
+        num_key_value_heads=model_config.num_key_value_heads,
+        pp_size=cfg.pp_size,
+        attn_cp_size=cfg.attn_cp_size,
+        enable_prefill_cp=cfg.enable_prefill_cp,
+        ep_size=cfg.ep_size,
+        moe_dp_size=cfg.moe_dp_size,
+        disaggregation_mode=cfg.disaggregation_mode,
+        dcp_size=cfg.dcp_size,
+    )
+
 
 def handle_data_parallelism(server_args: Any):
     # The resets without attention DP live in the resolution pipeline