[AMD] Add unified kv attention support in dpsk-v4 (#27380)

Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>

python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh

@@ -368,7 +368,7 @@ INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNor
 // Each thread loads kVecSize=2 BF16, so 256 threads cover the full 512 elems.
 // Cache layout: 584 bytes/token = 448 fp8 nope + 64 (=32 bf16x2) rope + 8 scale.
 // ----------------------------------------------------------------------------
-template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
+template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL, bool kBf16Store = false>
 FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) {
   using namespace device;
   using enum ForwardMode;
@@ -379,7 +379,10 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
   // Last warp owns the rope tail. The remaining 7 warps each emit one
   // 64-element fp8 group (own UE8M0 scale).
   constexpr uint32_t kRopeWarp = kNumWarps - 1;
-  constexpr int64_t kPageBytes = host::div_ceil(584ll << kPageBits, 576) * 576;
+  // kBf16Store: write the whole head_dim as plain BF16 (no fp8 / no scale) into a
+  // [num_slots, head_dim] bf16 cache (page_size==1) at row out_loc
+  constexpr int64_t kPageBytes =
+      kBf16Store ? ((kHeadDim * 2ll) << kPageBits) : host::div_ceil(584ll << kPageBits, 576) * 576;
   static_assert(kHeadDim == kBlockSize * kVecSize);
   static_assert(kRopeDim == kWarpThreads * kVecSize);
   static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
@@ -450,12 +453,23 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
   const int32_t page = out_loc >> kPageBits;
   const int32_t offset = out_loc & ((1 << kPageBits) - 1);
   const auto page_ptr = params.kvcache + page * kPageBytes;
-  const auto value_ptr = page_ptr + offset * 576;
+  const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576);
 
   PDLTriggerSecondary<kUsePDL>();
 
   // part 2: rope on the rope warp (BF16 store), or per-warp FP8 quant + store.
-  if (warp_id == kRopeWarp) {
+  if constexpr (kBf16Store) {
+    Float2 d = data;
+    if (warp_id == kRopeWarp) {
+      const auto x_real = data[0];
+      const auto x_imag = data[1];
+      const auto freq_real = freq[0];
+      const auto freq_imag = freq[1];
+      d[0] = x_real * freq_real - x_imag * freq_imag;
+      d[1] = x_real * freq_imag + x_imag * freq_real;
+    }
+    reinterpret_cast<bf16x2_t*>(value_ptr)[tx] = cast<bf16x2_t>(fp32x2_t{d[0], d[1]});
+  } else if (warp_id == kRopeWarp) {
     // Each rope-warp lane owns exactly one (real, imag) pair within the rope
     // tail. Apply rotation, downcast to BF16, write to the slot's rope region.
     const auto x_real = data[0];
@@ -485,13 +499,15 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
   }
 }
 
-template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL>
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL, bool kBf16Store = false>
 struct FusedNormRopeKernel {
   static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
   static constexpr bool kIsIndexer = (kHeadDim == 128);
+  static_assert(!(kIsIndexer && kBf16Store), "bf16 store only for flashmla head_dim=512");
   static constexpr int64_t kIndexerBytes = 132 * kPageSize;
   static constexpr int64_t kFlashMLABytes = host::div_ceil(584 * kPageSize, 576) * 576;
-  static constexpr int64_t kPageBytes = kIsIndexer ? kIndexerBytes : kFlashMLABytes;
+  static constexpr int64_t kBf16Bytes = kHeadDim * 2 * kPageSize;  // plain bf16 cache
+  static constexpr int64_t kPageBytes = kBf16Store ? kBf16Bytes : (kIsIndexer ? kIndexerBytes : kFlashMLABytes);
 
   /// TODO: Let's fix the config for now.
   static_assert(kRopeDim == 64 && (kHeadDim == 128 || kHeadDim == 512));
@@ -502,7 +518,7 @@ struct FusedNormRopeKernel {
     if constexpr (kIsIndexer) {
       return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL>;
     } else {
-      return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL>;
+      return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store>;
     }
   }
 

python/sglang/jit_kernel/dsv4/compress.py

@@ -23,8 +23,11 @@ def _jit_compress_norm_rope_module(
     head_dim: int,
     rope_dim: int,
     page_size: int,
+    bf16_store: bool = False,
 ) -> Module:
-    args = make_cpp_args(dtype, head_dim, rope_dim, page_size, is_arch_support_pdl())
+    args = make_cpp_args(
+        dtype, head_dim, rope_dim, page_size, is_arch_support_pdl(), bf16_store
+    )
     cuda_wrappers = [("forward", f"FusedNormRopeKernel<{args}>::forward")]
     if head_dim == 128:
         cuda_wrappers.append(
@@ -339,12 +342,13 @@ def compress_norm_rope_store(
     kvcache: torch.Tensor,
     page_size: int,
     use_fp4: bool = False,
+    bf16_store: bool = False,
 ) -> None:
     if use_fp4:
         assert kv.shape[-1] == 128
     freq_cis = torch.view_as_real(freq_cis).flatten(-2)
     module = _jit_compress_norm_rope_module(
-        kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size
+        kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size, bf16_store
     )
     fn = module.forward_fp4 if use_fp4 else module.forward
     fn(

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -238,6 +238,7 @@ class DSV4AttnMetadata:
             self.c4_topk_lengths_clamp1,
             self.c128_out_loc,
             _,
+            _,
             self.c128_topk_lengths_clamp1,
             self.c128_page_indices,
         ) = _init_compression_metadata_triton(

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -113,11 +113,28 @@ class DSV4AttnMetadata:
     c4_topk_lengths_raw: Optional[torch.Tensor] = None
     c4_topk_lengths_clamp1: Optional[torch.Tensor] = None
     c4_sparse_topk_lengths: torch.Tensor = field(init=False)
+    c4_sparse_topk_lengths_raw: torch.Tensor = field(init=False)
     c4_sparse_page_indices: torch.Tensor = field(init=False)
+    c4_sparse_raw_indices: Optional[torch.Tensor] = field(init=False, default=None)
 
     c128_out_loc: Optional[torch.Tensor] = None
     c128_page_indices: Optional[torch.Tensor] = None
     c128_topk_lengths_clamp1: Optional[torch.Tensor] = None
+    c128_topk_lengths_raw: Optional[torch.Tensor] = None
+
+    # unified_kv: per-forward prebuilt ragged decode index
+    unified_swa_indices: Optional[torch.Tensor] = None
+    unified_swa_indptr: Optional[torch.Tensor] = None
+    unified_hca_indices: Optional[torch.Tensor] = None
+    unified_hca_indptr: Optional[torch.Tensor] = None
+    unified_csa_indices: Optional[torch.Tensor] = None
+    unified_csa_indptr: Optional[torch.Tensor] = None
+
+    # unified_kv: per-forward prefill/extend per-token mapping
+    unified_pf_state_slot: Optional[torch.Tensor] = None
+    unified_pf_chunk_start: Optional[torch.Tensor] = None
+    unified_pf_cu_q: Optional[torch.Tensor] = None
+    unified_pf_final_pos: Optional[torch.Tensor] = None
 
     c1_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False)
     c4_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False)
@@ -157,10 +174,23 @@ class DSV4AttnMetadata:
                 "swa_topk_lengths",
                 "c128_page_indices",
                 "c128_topk_lengths_clamp1",
+                "c128_topk_lengths_raw",
                 "c4_topk_lengths_raw",
                 "c4_topk_lengths_clamp1",
                 "c4_sparse_topk_lengths",
+                "c4_sparse_topk_lengths_raw",
                 "c4_sparse_page_indices",
+                "c4_sparse_raw_indices",
+                "unified_swa_indices",
+                "unified_swa_indptr",
+                "unified_hca_indices",
+                "unified_hca_indptr",
+                "unified_csa_indices",
+                "unified_csa_indptr",
+                "unified_pf_state_slot",
+                "unified_pf_chunk_start",
+                "unified_pf_cu_q",
+                "unified_pf_final_pos",
             ],
             assign_fields=[
                 # Recomputed by the recorded init_forward_metadata_in_graph op
@@ -185,6 +215,7 @@ class DSV4AttnMetadata:
             self.c4_topk_lengths_clamp1,
             self.c128_out_loc,
             _,
+            self.c128_topk_lengths_raw,
             self.c128_topk_lengths_clamp1,
             self.c128_page_indices,
         ) = _init_compression_metadata_triton(
@@ -209,6 +240,7 @@ class DSV4AttnMetadata:
         "c4_topk_lengths_clamp1",
         "c128_page_indices",
         "c128_topk_lengths_clamp1",
+        "c128_topk_lengths_raw",
     ]
     _CP_GLOBAL_FIELDS = [
         "raw_out_loc",
@@ -260,6 +292,10 @@ class DSV4AttnMetadata:
         self.c4_sparse_topk_lengths = torch.clamp(
             self.c4_topk_lengths_clamp1, max=self.c4_sparse_topk
         )
+        assert self.c4_topk_lengths_raw is not None
+        self.c4_sparse_topk_lengths_raw = torch.clamp(
+            self.c4_topk_lengths_raw, max=self.c4_sparse_topk
+        )
         self.c4_sparse_page_indices = torch.full(
             (self.c4_topk_lengths_clamp1.size(0), self.c4_sparse_topk),
             -1,
@@ -427,6 +463,7 @@ class DeepseekV4HipRadixBackend(
             out_loc=out_cache_loc,
             need_compress=True,
         )
+        self._attach_unified_kv_decode_streams(core_attn_metadata, req_pool_indices)
 
         indexer_metadata = self.init_forward_metadata_indexer(core_attn_metadata)
 
@@ -475,6 +512,9 @@ class DeepseekV4HipRadixBackend(
             need_compress=need_compress,
             is_prefill=True,
         )
+        self._attach_unified_kv_prefill_meta(
+            core_attn_metadata, req_pool_indices, seq_lens, extend_seq_lens
+        )
         indexer_metadata = (
             self.init_forward_metadata_indexer(core_attn_metadata)
             if need_compress
@@ -621,6 +661,7 @@ class DeepseekV4HipRadixBackend(
             out_loc=out_cache_loc,
             need_compress=True,
         )
+        self._attach_unified_kv_decode_streams(core_attn_metadata, req_pool_indices)
         indexer_metadata = self.init_forward_metadata_indexer(core_attn_metadata)
 
         create = functools.partial(
@@ -954,6 +995,191 @@ class DeepseekV4HipRadixBackend(
         if current_raw is not None:
             self.forward_metadata = current_raw
 
+    def _attach_unified_kv_decode_streams(
+        self, core: "DSV4AttnMetadata", req_pool_indices: torch.Tensor
+    ) -> None:
+        """build the ragged decode index streams once per forward"""
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_triton,
+        )
+
+        if not is_unified_kv_triton():
+            return
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime
+
+        pool = self.token_to_kv_pool
+        N = core.positions_casual.shape[0]
+        (
+            core.unified_swa_indices,
+            core.unified_swa_indptr,
+            core.unified_hca_indices,
+            core.unified_hca_indptr,
+            core.unified_csa_indices,
+            core.unified_csa_indptr,
+        ) = runtime.build_decode_streams(
+            state_slot=req_pool_indices[:N],
+            positions=core.positions_casual,
+            swa_len=core.swa_topk_lengths,
+            hca_len=core.c128_topk_lengths_raw,
+            csa_len=core.c4_sparse_topk_lengths_raw,
+            hca_page_indices=core.c128_page_indices,
+            csa_width=core.c4_sparse_page_indices.shape[1],
+            win=pool.unified_swa_window,
+            ring_stride=pool.unified_swa_ring_size,
+            swa_pages=pool.unified_swa_pages,
+        )
+
+    def _attach_unified_kv_prefill_meta(
+        self,
+        core: "DSV4AttnMetadata",
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+    ) -> None:
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_triton,
+        )
+
+        if not is_unified_kv_triton():
+            return
+        device = req_pool_indices.device
+        bs = req_pool_indices.shape[0]
+        seq_lens = seq_lens.to(torch.int64)
+        extend_seq_lens = extend_seq_lens.to(torch.int64)
+        # token -> req index (length L = sum(extend_seq_lens))
+        bid = torch.repeat_interleave(
+            torch.arange(bs, device=device, dtype=torch.int64), extend_seq_lens
+        )
+        core.unified_pf_state_slot = req_pool_indices[bid]
+        core.unified_pf_chunk_start = (seq_lens - extend_seq_lens)[bid]
+        cu_q_per_req = torch.cumsum(extend_seq_lens, dim=0) - extend_seq_lens
+        core.unified_pf_cu_q = cu_q_per_req[bid]
+        core.unified_pf_final_pos = (seq_lens - 1)[bid]
+
+    def _forward_unified_kv(
+        self,
+        *,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        compress_ratio: Literal[0, 4, 128],
+        attn_sink: torch.Tensor,
+        core_attn_metadata: "DSV4AttnMetadata",
+        save_kv_cache: bool = True,
+    ) -> torch.Tensor:
+        """unified_kv paged-attention path over the bf16 unified_kv"""
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime
+
+        pool = self.token_to_kv_pool
+        layer_id = layer.layer_id
+        unified = pool.get_unified_kv(layer_id)
+        win = pool.unified_swa_window
+        ring_stride = pool.unified_swa_ring_size
+        swa_pages = pool.unified_swa_pages
+
+        if q.ndim == 4:
+            q = q.squeeze(1)
+        device = q.device
+        positions = forward_batch.positions.to(torch.int64)
+        T = q.shape[0]
+        positions = positions[:T]
+
+        c128_pi = getattr(core_attn_metadata, "c128_page_indices", None)
+        c4_pi = getattr(core_attn_metadata, "c4_sparse_page_indices", None)
+
+        # decode
+        is_decode = forward_batch.forward_mode.is_decode_or_idle()
+        if is_decode:
+            state_slot = forward_batch.req_pool_indices[:T]
+            if save_kv_cache:
+                runtime.store_swa_into_unified(
+                    kv=kv,
+                    state_slot=state_slot,
+                    positions=positions,
+                    unified_kv=unified,
+                    win=win,
+                    ring_stride=ring_stride,
+                    final_pos=positions,
+                )
+            if compress_ratio == 0:
+                kv_indices = core_attn_metadata.unified_swa_indices
+                kv_indptr = core_attn_metadata.unified_swa_indptr
+            elif compress_ratio == 128:
+                kv_indices = core_attn_metadata.unified_hca_indices
+                kv_indptr = core_attn_metadata.unified_hca_indptr
+            elif compress_ratio == 4:
+                kv_indices = core_attn_metadata.unified_csa_indices
+                kv_indptr = core_attn_metadata.unified_csa_indptr
+                runtime.fill_compress_tail(
+                    indices=kv_indices,
+                    indptr=kv_indptr,
+                    prefix_len=core_attn_metadata.swa_topk_lengths[:T],
+                    page_indices=c4_pi[:T],
+                    valid_len=core_attn_metadata.c4_sparse_topk_lengths_raw[:T],
+                    swa_pages=swa_pages,
+                )
+            else:
+                raise ValueError(f"bad compress_ratio {compress_ratio}")
+            return runtime.decode(
+                q=q,
+                unified_kv=unified,
+                kv_indices=kv_indices,
+                kv_indptr=kv_indptr,
+                attn_sink=attn_sink,
+                softmax_scale=self.softmax_scale,
+            )
+
+        # prefill / extend
+        state_slot = core_attn_metadata.unified_pf_state_slot
+        chunk_start = core_attn_metadata.unified_pf_chunk_start
+        cu_q = core_attn_metadata.unified_pf_cu_q
+        final_pos = core_attn_metadata.unified_pf_final_pos
+
+        kpre_i, kpre_p, kext_i, kext_p = runtime.build_prefill_indices(
+            compress_ratio=compress_ratio,
+            state_slot=state_slot,
+            positions=positions,
+            chunk_start=chunk_start,
+            cu_q=cu_q,
+            win=win,
+            ring_stride=ring_stride,
+            swa_pages=swa_pages,
+            c128_page_indices=c128_pi,
+            c4_sparse_page_indices=c4_pi,
+        )
+
+        if kpre_p.shape[0] < T + 1:
+            pad = T + 1 - kpre_p.shape[0]
+            kpre_p = torch.cat([kpre_p, kpre_p[-1:].expand(pad)])
+            kext_p = torch.cat([kext_p, kext_p[-1:].expand(pad)])
+        o = runtime.prefill(
+            q=q,
+            unified_kv=unified,
+            kv_indices_prefix=kpre_i,
+            kv_indptr_prefix=kpre_p,
+            kv_extend=kv,
+            kv_indices_extend=kext_i,
+            kv_indptr_extend=kext_p,
+            attn_sink=attn_sink,
+            softmax_scale=self.softmax_scale,
+        )
+
+        # write this chunk's SWA K into the ring for future chunks / decode
+        # only the final-window tokens per request
+        if save_kv_cache:
+            n_real = state_slot.shape[0]
+            runtime.store_swa_into_unified(
+                kv=kv[:n_real],
+                state_slot=state_slot,
+                positions=positions[:n_real],
+                unified_kv=unified,
+                win=win,
+                ring_stride=ring_stride,
+                final_pos=final_pos,
+            )
+        return o
+
     def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
         """Resolve the SWA KV-store write target for the current forward.
 
@@ -1022,6 +1248,22 @@ class DeepseekV4HipRadixBackend(
         token_to_kv_pool = self.token_to_kv_pool
         assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
 
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_triton,
+        )
+
+        if is_unified_kv_triton():
+            return self._forward_unified_kv(
+                q=q,
+                kv=swa_k,
+                layer=layer,
+                forward_batch=forward_batch,
+                compress_ratio=compress_ratio,
+                attn_sink=attn_sink,
+                core_attn_metadata=core_attn_metadata,
+                save_kv_cache=save_kv_cache,
+            )
+
         if isinstance(core_attn_metadata, DSV4AttnMetadata):
             if save_kv_cache:
                 self.store_cache(layer_id, swa_k, forward_batch)
@@ -1218,7 +1460,9 @@ class DeepseekV4HipRadixBackend(
             core_attn_metadata.init_flashmla_related()
         else:
             core_attn_metadata.c4_sparse_topk_lengths = None
+            core_attn_metadata.c4_sparse_topk_lengths_raw = None
             core_attn_metadata.c4_sparse_page_indices = None
+            core_attn_metadata.c4_sparse_raw_indices = None
             core_attn_metadata.c1_flashmla_metadata = _create_flashmla_metadata()
             core_attn_metadata.c4_flashmla_metadata = None
             core_attn_metadata.c128_flashmla_metadata = None

python/sglang/srt/layers/attention/dsv4/compressor.py

@@ -406,9 +406,9 @@ class Compressor(nn.Module):
         return ret
 
     def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
-        if _tgemm is not None:
+        if _tgemm is not None and not envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
-            # linear_bf16_fp32 uses tgemm.mm + .float(); skip the .float() cast
+            # v1 compress goes through fused_compress_triton, which promotes
-            # because downstream Triton kernels promote bf16→fp32 internally.
+            # bf16->fp32 internally, so skip the .float() cast.
             kv_score = _tgemm.mm(x, self.wkv_gate.weight, otype=x.dtype)
         else:
             kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)

python/sglang/srt/layers/attention/dsv4/compressor_v2.py

@@ -428,6 +428,7 @@ class CompressorBackendMixin:
         page_size: int,
         out_loc: torch.Tensor,
         use_fp4_indexer: bool = False,
+        bf16_store: bool = False,
     ) -> None:
         assert compress_ratio == 4 or compress_ratio == 128
         assert rotate == is_indexer == (head_dim == 128)
@@ -468,6 +469,7 @@ class CompressorBackendMixin:
             kvcache=kv_cache,
             page_size=page_size,
             use_fp4=use_fp4_indexer,
+            bf16_store=bf16_store,
         )
 
     def forward_unified(
@@ -485,6 +487,10 @@ class CompressorBackendMixin:
         kv_score_input = compressor.compute_kv_score(x, forward_batch)
 
         state_pool = compressor.get_state_pool(self)
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_triton,
+        )
+
         if _is_hip and not envs.SGLANG_OPT_USE_JIT_NORM.get():
             self._forward_unified_hip(
                 token_to_kv_pool=token_to_kv_pool,
@@ -498,9 +504,15 @@ class CompressorBackendMixin:
             use_fp4_indexer = (
                 compressor.is_in_indexer and self.enable_deepseek_v4_fp4_indexer
             )
+            bf16_store = False
             if compressor.is_in_indexer:
                 kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id)
                 page_size = token_to_kv_pool.get_index_k_page_size()
+            elif is_unified_kv_triton():
+                kv_cache = token_to_kv_pool.get_unified_kv(layer_id)
+                page_size = 1
+                out_loc = out_loc + token_to_kv_pool.unified_swa_pages
+                bf16_store = True
             else:
                 _, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
                 assert compress_kv_pool is not None
@@ -527,6 +539,7 @@ class CompressorBackendMixin:
                 page_size=page_size,
                 out_loc=out_loc,
                 use_fp4_indexer=use_fp4_indexer,
+                bf16_store=bf16_store,
             )
 
     def _forward_unified_hip(

python/sglang/srt/layers/attention/dsv4/metadata_kernel.py

@@ -17,6 +17,7 @@ def _init_compressed_attn_metadata_kernel(
     c4_seq_lens_clamp1_ptr,
     c128_out_loc_ptr,
     c128_positions_ptr,
+    c128_seq_lens_raw_ptr,
     c128_seq_lens_clamp1_ptr,
     c128_page_indices_ptr,
     bs,
@@ -54,6 +55,7 @@ def _init_compressed_attn_metadata_kernel(
 
     tl.store(c128_out_loc_ptr + batch_id, c128_out_loc)
     tl.store(c128_positions_ptr + batch_id, c128_positions)
+    tl.store(c128_seq_lens_raw_ptr + batch_id, c128_seq_lens_raw)
     tl.store(c128_seq_lens_clamp1_ptr + batch_id, c128_seq_lens_clamp1)
 
     if COMPUTE_PAGE_INDICES:
@@ -99,6 +101,7 @@ def _init_compressed_attn_metadata_triton(
     torch.Tensor,
     torch.Tensor,
     torch.Tensor,
+    torch.Tensor,
     Optional[torch.Tensor],
 ]:
     bs = seq_lens.shape[0]
@@ -111,6 +114,7 @@ def _init_compressed_attn_metadata_triton(
 
     c128_out_loc = torch.empty(bs, dtype=torch.int32, device=device)
     c128_positions = torch.empty(bs, dtype=torch.int32, device=device)
+    c128_seq_lens_raw = torch.empty(bs, dtype=torch.int32, device=device)
     c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device)
 
     if compute_page_indices:
@@ -146,6 +150,7 @@ def _init_compressed_attn_metadata_triton(
         c4_seq_lens_clamp1,
         c128_out_loc,
         c128_positions,
+        c128_seq_lens_raw,
         c128_seq_lens_clamp1,
         (
             c128_page_indices
@@ -168,6 +173,7 @@ def _init_compressed_attn_metadata_triton(
         c4_seq_lens_clamp1,
         c128_out_loc,
         c128_positions,
+        c128_seq_lens_raw,
         c128_seq_lens_clamp1,
         c128_page_indices,
     )
@@ -188,6 +194,7 @@ def init_compression_metadata(
     torch.Tensor,
     torch.Tensor,
     torch.Tensor,
+    torch.Tensor,
     Optional[torch.Tensor],
 ]:
     return _init_compressed_attn_metadata_triton(

python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+import functools
+import os
+
+
+@functools.lru_cache(maxsize=1)
+def is_unified_kv_triton() -> bool:
+    return os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "") == "unified_kv_triton"

python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py

@@ -0,0 +1,881 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Sparse decode attention over a unified KV pool with per-token paged indices.
+
+Designed for V4 decode + CUDAGraph: replaces the per-fwd `kv_flat_sa`
+materialization (whose shape depends on `n_committed_per_seq` → varies per
+fwd → blocks CG capture) with a single unified KV pool indexed via paged
+indices, mirroring `aiter.mla.mla_decode_fwd`'s API style.
+
+Caller contract:
+  unified_kv:       [total_pages, D] BF16  (page_size=1)
+    Conceptually merges the SWA ring buffer and the compressor paged cache
+    of a single V4 layer. Slots in `[0, swa_pages)` reference SWA entries
+    (state_slot * win + ring); slots in `[swa_pages, ...)` reference
+    compressed-K entries (block_id * K_PER_BLOCK + slot_in_block).
+  kv_indices: [total_indices] int32 — per-token slot lists, flat.
+    Per-token entries live in
+    `kv_indices[kv_indptr[t] : kv_indptr[t+1]]`.
+    **All entries MUST be valid slot ids in [0, unified_kv.shape[0]).**
+    The production decode index builder (``write_v4_paged_decode_indices``)
+    emits ragged-packed indices with no sentinels; CG-padded tokens get
+    a zero-length slice via ``indptr[t+1] == indptr[t]``. The kernel no
+    longer carries a per-iter ``slot >= 0`` sentinel check.
+  kv_indptr:  [N+1] int32 — true prefix sum (variable per-token len).
+  attn_sink:        [H] per-head learnable softmax-denom bias (V4 specific).
+  softmax_scale:    float.
+
+Returns:
+  out: [N, H, D] same dtype as q.
+
+Numerics: online-softmax in log2 domain (exp2 with qk_scale = softmax_scale *
+LOG2E), with attention sink folded as a virtual K. Bit-close to the PyTorch
+reference (``_sparse_attn_ragged_torch``) within fp32-accumulation tolerance.
+
+Architecture:
+  - Small T (T * ceil(H/block_h) < ~1.7×CU): split-K + reduce-kernel path.
+    Split kernel emits (m, l, acc) fp32 partials; reduce combines splits and
+    folds attn_sink.
+  - Large T: single-pass FUSED kernel that does softmax + sink + write in one
+    shot (no partial-buffer alloc, no second kernel launch).
+
+CUDAGraph-safe: kv_splits and tile config depend only on capture-time
+shapes (``T``, ``H``); the kernels' early-return / segment-mask logic is
+driven by ``kv_indptr`` values, which are runtime data but don't affect
+the captured launch sequence.
+"""
+
+from __future__ import annotations
+
+import functools
+
+import torch
+import triton
+import triton.language as tl
+from aiter.ops.triton.utils.device_info import get_num_sms
+
+LOG2E = 1.4426950408889634  # log2(e); folded into qk_scale so softmax can use exp2.
+_MAX_KV_SPLITS = 64  # Hard cap on kv_splits (see _kv_splits_heuristic).
+
+# FP8 KV cache (1xGROUP_SIZE block-scale quantization).
+#
+# Storage: unified_kv[total_pages, D] in e4m3fnuz + kv_scales[total_pages,
+# D // GROUP_SIZE] in fp32. Per-slot, D is split into NUM_GROUPS chunks of
+# GROUP_SIZE elements; each chunk shares one fp32 scale.
+# Dequant in-kernel: kv_bf16 = kv_fp8.to(fp32) * scale[d // GROUP_SIZE], cast
+# back to q.dtype before the second dot.
+#
+# GROUP_SIZE=64 matches the user's 1x64 quant spec. V4-Pro: D=512 → 8 scales
+# per slot, 4 bytes each → +6.25% storage on top of the fp8 pool (vs the
+# halving from bf16→fp8 = 2× saving — net ~46% read bandwidth reduction).
+_FP8_GROUP_SIZE = 64
+_FP8_DTYPE = torch.float8_e4m3fnuz
+
+
+@functools.lru_cache(maxsize=1)
+def _cu_count() -> int:
+    """Compute-unit count of the active GPU, queried once via aiter.
+
+    Wrapped in ``lru_cache`` so the first decode call pays the device-property
+    lookup and all subsequent calls hit the cache — important inside a hot
+    decode loop and CUDAGraph capture (no data-dependent host work).
+    """
+    return get_num_sms()
+
+
+# ---------------------------------------------------------------------------
+# Heuristics — pure-Python, capture-time deterministic.
+# ---------------------------------------------------------------------------
+
+
+def _kernel_config(block_h: int) -> tuple[int, int, int]:
+    """Pick (BLOCK_K, num_warps, num_stages) without autotune.
+
+    Depends ONLY on ``block_h`` (a function of H, capture-time shape) so the
+    config is identical between CUDAGraph capture and replay regardless of
+    per-token K. In production ``kv_indices.shape[0]`` is a padded bucket
+    whose value is unrelated to the true per-token kv_len, so any heuristic
+    that reads it would mis-tune at capture time.
+
+    Derived from autotune statistics over ~150 shapes on MI355:
+      - BLOCK_K=16 dominated (~78% of best configs for D=512); D=512 has
+        enough load width that wider K tiles spill regs more than they buy.
+      - num_warps:
+          block_h ≤ 32 → 4 warps; block_h ≥ 64 → 8 warps.
+        Pre-v11 the threshold was 16, which gave H=32 nw=8 — that was a
+        regression worth +30% geomean (max +57%) on H=32 across all T,
+        from a Phase-1 fine-tune sweep on top of v05+v09. With block_h=32
+        the MFMA tile is 32×16×D; 4 waves (256 threads) is the sweet
+        spot for AMD wave64 register budget. 8 warps over-distribute and
+        leave each warp with too little MFMA to amortize pipeline fill.
+        block_h=64 still wants 8 warps (one wave per row tile is too
+        little ILP).
+      - num_stages=2 is the safe default — deeper pipelining (3) helps only
+        when the K loop has many iterations; we cannot know per-token K at
+        capture time, so the conservative pick avoids regressing short-K.
+    """
+    block_k = 16
+    num_warps = 4 if block_h <= 32 else 8
+    num_stages = 2
+    return block_k, num_warps, num_stages
+
+
+def _prev_pow2(n: int) -> int:
+    """Largest power of two ≤ ``n``. For n < 1 returns 1."""
+    if n < 1:
+        return 1
+    return 1 << (n.bit_length() - 1)
+
+
+def _kv_splits_heuristic(
+    T: int,
+    H: int,
+    block_h: int,
+    num_cu: int | None = None,
+    target_wg_per_cu: float = 2.0,
+    max_kv_splits: int = _MAX_KV_SPLITS,
+) -> int:
+    """Pick KV_SPLITS to fill the GPU. CUDAGraph-safe: depends ONLY on
+    capture-time scalars (``T``, ``H``, ``block_h``). Never reads any
+    tensor value or shape — production callers must not assume kv_indices
+    layout encodes per-token kv_len.
+
+    Split-K trades a single-kernel pass for (split, reduce) two-kernel +
+    partial-buffer allocation. The trade is worth it when the base grid
+    ``T * ceil(H/block_h)`` underfills the device.
+
+      base_ctas  = T * ceil(H / block_h)
+      target_wg  = target_wg_per_cu * num_cu     (≈ 1.7x to hide load-imbalance)
+      if base_ctas >= target_wg:  splits = 1     (grid already saturates GPU)
+      else:                       splits = prev_pow2(min(target_wg/base_ctas,
+                                                          max_kv_splits))
+
+    ``max_kv_splits`` (default 64) caps the number of split-kernel CTAs per
+    token. Higher values would buy more parallelism for bs=1 long-ctx, but
+    when per-token K is short most splits fall-through and the launch
+    overhead dominates. 64 is the sweet spot for MI300/MI355.
+
+    Rounded DOWN to a power of two — rounding up over-splits when
+    splits_to_fill isn't already pow2 (e.g. T=2 → 258 → 512 doubles the wg
+    count past target, halves per-split work → 4× slowdown on bs=2 ctx=16384).
+    """
+    if num_cu is None:
+        num_cu = _cu_count()
+    target_wg = max(1, int(target_wg_per_cu * num_cu))
+    head_blocks = max(1, (H + block_h - 1) // block_h)
+    base_ctas = max(1, T * head_blocks)
+    if base_ctas >= target_wg:
+        return 1
+
+    splits_to_fill = max(1, target_wg // base_ctas)
+    return _prev_pow2(min(splits_to_fill, max_kv_splits))
+
+
+# ---------------------------------------------------------------------------
+# Kernels.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _paged_decode_fused_kernel(
+    q_ptr,  # [N, H, D]
+    unified_kv_ptr,  # [total_pages, D] bf16/fp16, or fp8 when QUANT_KV
+    kv_scales_ptr,  # [total_pages, NUM_GROUPS] fp32 when QUANT_KV (dummy otherwise)
+    kv_indices_ptr,  # [total_indices] int32
+    kv_indptr_ptr,  # [N+1] int32
+    attn_sink_ptr,  # [H]
+    out_ptr,  # [N, H, D]
+    q_stride_t,
+    q_stride_h,
+    q_stride_d,
+    kv_stride_n,
+    kv_stride_d,
+    ks_stride_n,  # row stride of kv_scales (groups are contiguous, stride=1)
+    out_stride_t,
+    out_stride_h,
+    out_stride_d,
+    qk_scale,  # = softmax_scale * LOG2E
+    log2e,  # = LOG2E, to lift natural-log sink into log2 domain
+    H: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    QUANT_KV: tl.constexpr,  # True → dequant fp8 KV via kv_scales
+    GROUP_SIZE: tl.constexpr,  # scale block width along D (e.g. 64)
+    NUM_GROUPS: tl.constexpr,  # D // GROUP_SIZE (constexpr; D % GROUP_SIZE == 0)
+):
+    """Single-pass online-softmax with sink folded inline — fast path for
+    cases where ``kv_splits = 1`` (base grid already saturates the GPU). Skips
+    the partial-buffer alloc + reduce-kernel launch that the 2-kernel
+    split-K path needs.
+
+    Grid: ``(N, ceil(H / BLOCK_H))``. One CTA owns one token and one head-tile.
+    """
+    t = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    d_offs = tl.arange(0, BLOCK_D)
+    h_mask = h_offs < H
+    d_mask = d_offs < D
+
+    q = tl.load(
+        q_ptr
+        + t * q_stride_t
+        + h_offs[:, None] * q_stride_h
+        + d_offs[None, :] * q_stride_d,
+        mask=h_mask[:, None] & d_mask[None, :],
+        other=0.0,
+    )
+
+    kv_start = tl.load(kv_indptr_ptr + t)
+    kv_end = tl.load(kv_indptr_ptr + t + 1)
+    kv_len = kv_end - kv_start
+    num_tiles = tl.cdiv(kv_len, BLOCK_K)
+
+    neg_large = -3.4028234663852886e38
+    m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
+    l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
+
+    k_offs = tl.arange(0, BLOCK_K)
+    if QUANT_KV:
+        # Compile-time per-D-element group index: d_offs // GROUP_SIZE has
+        # NUM_GROUPS distinct values; redundant scale loads at the same
+        # address are coalesced through L1, no per-element scalar issue.
+        g_idx_per_d = d_offs // GROUP_SIZE
+    # num_stages=3 on the inner K loop overrides the launch-time default
+    # (2) for this loop only. Deeper SW pipeline keeps 2 in-flight KV
+    # gathers (vs 1 with stages=2) while the current MFMA runs — better
+    # hides AMD HBM gather latency on D=512. Cost: ~1 extra KV tile
+    # (BLOCK_K*BLOCK_D*2 = 16KB at bf16) staged in regs/LDS per CTA.
+    for j in tl.range(0, num_tiles, num_stages=3):
+        k_start = j * BLOCK_K
+        k_pos = k_start + k_offs
+        valid = k_pos < kv_len  # in_range; no sentinel check (see contract)
+        slot = tl.load(
+            kv_indices_ptr + kv_start + k_pos,
+            mask=valid,
+            other=0,  # any in-bounds slot; the read is masked out below
+        )
+
+        kv_raw = tl.load(
+            unified_kv_ptr
+            + slot[:, None] * kv_stride_n
+            + d_offs[None, :] * kv_stride_d,
+            mask=valid[:, None] & d_mask[None, :],
+            other=0.0,
+        )
+        if QUANT_KV:
+            # 1xGROUP_SIZE block-scale dequant via direct broadcast load —
+            # avoids the explicit reshape + 3D intermediate that pinned
+            # too many bf16 tiles in flight at once. The masked load with
+            # d_offs // GROUP_SIZE as the column index produces a virtual
+            # [BLOCK_K, BLOCK_D] scales tile but in IR is a coalesced
+            # NUM_GROUPS-wide load per row.
+            scales_full = tl.load(
+                kv_scales_ptr + slot[:, None] * ks_stride_n + g_idx_per_d[None, :],
+                mask=valid[:, None] & d_mask[None, :],
+                other=0.0,
+            ).to(q.dtype)
+            kv = kv_raw.to(q.dtype) * scales_full
+        else:
+            kv = kv_raw
+
+        scores = tl.dot(q, tl.trans(kv)) * qk_scale
+        # K: drop h_mask from the per-iter where. In V4-Pro every realistic
+        # (H, BLOCK_H) pair has H % BLOCK_H == 0 → h_mask is statically
+        # all-True, but the runtime ``h_offs < H`` compare prevents Triton
+        # from constant-folding. Masking only on ``valid`` is sufficient:
+        # ``neg_large`` on invalid k positions makes exp2(scores - m) ≈ 0
+        # in the subsequent dot, and the masked store at the end gates
+        # invalid h rows from polluting the output.
+        scores = tl.where(valid[None, :], scores, neg_large)
+
+        m_block = tl.max(scores, axis=1)
+        m_new = tl.maximum(m_i, m_block)
+        alpha = tl.exp2(m_i - m_new)
+        p = tl.exp2(scores - m_new[:, None])
+        l_new = l_i * alpha + tl.sum(p, axis=1)
+
+        acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
+        m_i = m_new
+        l_i = l_new
+
+    # Fold attn_sink as a virtual K of weight 1. sink is a natural-log bias;
+    # multiply by log2e so it lives in the same log2 domain as our online m_i.
+    sink_raw = tl.load(attn_sink_ptr + h_offs, mask=h_mask, other=neg_large).to(
+        tl.float32
+    )
+    sink = sink_raw * log2e
+    m_final = tl.maximum(m_i, sink)
+    alpha_kv = tl.exp2(m_i - m_final)
+    alpha_sink = tl.exp2(sink - m_final)
+    l_final = l_i * alpha_kv + alpha_sink
+
+    denom = tl.maximum(l_final, 1.0e-30)
+    out = tl.where(
+        l_final[:, None] > 0.0, (acc * alpha_kv[:, None]) / denom[:, None], 0.0
+    )
+    tl.store(
+        out_ptr
+        + t * out_stride_t
+        + h_offs[:, None] * out_stride_h
+        + d_offs[None, :] * out_stride_d,
+        out.to(out_ptr.dtype.element_ty),
+        mask=h_mask[:, None] & d_mask[None, :],
+    )
+
+
+@triton.jit
+def _paged_decode_split_kernel(
+    q_ptr,  # [N, H, D]
+    unified_kv_ptr,  # [total_pages, D] bf16/fp16, or fp8 when QUANT_KV
+    kv_scales_ptr,  # [total_pages, NUM_GROUPS] fp32 when QUANT_KV (dummy otherwise)
+    kv_indices_ptr,  # [total_indices] int32
+    kv_indptr_ptr,  # [N+1] int32
+    m_partial_ptr,  # [N, KV_SPLITS, H_padded] fp32
+    l_partial_ptr,  # [N, KV_SPLITS, H_padded] fp32
+    acc_partial_ptr,  # [N, KV_SPLITS, H_padded, D] fp32
+    q_stride_t,
+    q_stride_h,
+    q_stride_d,
+    kv_stride_n,
+    kv_stride_d,
+    ks_stride_n,  # row stride of kv_scales (groups are contiguous, stride=1)
+    mp_stride_t,
+    mp_stride_k,
+    mp_stride_h,
+    lp_stride_t,
+    lp_stride_k,
+    lp_stride_h,
+    ap_stride_t,
+    ap_stride_k,
+    ap_stride_h,
+    ap_stride_d,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    KV_SPLITS: tl.constexpr,
+    qk_scale,  # = softmax_scale * LOG2E
+    BLOCK_H: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    QUANT_KV: tl.constexpr,  # True → dequant fp8 KV via kv_scales
+    GROUP_SIZE: tl.constexpr,  # scale block width along D (e.g. 64)
+    NUM_GROUPS: tl.constexpr,  # D // GROUP_SIZE
+):
+    """3D split-K + exp2-softmax sparse paged-decode. Grid: (N, ceil(H/BLOCK_H), KV_SPLITS).
+
+    Emits pre-sink (m, l, acc) partials in fp32. The reduce kernel folds
+    ``attn_sink`` and combines splits.
+    """
+    t = tl.program_id(0)
+    pid_h = tl.program_id(1)
+    pid_k = tl.program_id(2)
+
+    h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    d_offs = tl.arange(0, BLOCK_D)
+    h_mask = h_offs < H
+    d_mask = d_offs < D
+
+    q = tl.load(
+        q_ptr
+        + t * q_stride_t
+        + h_offs[:, None] * q_stride_h
+        + d_offs[None, :] * q_stride_d,
+        mask=h_mask[:, None] & d_mask[None, :],
+        other=0.0,
+    )
+
+    kv_start = tl.load(kv_indptr_ptr + t)
+    kv_end = tl.load(kv_indptr_ptr + t + 1)
+    kv_len = kv_end - kv_start
+
+    # tiles_per_segment pattern from aiter's kernel_unified_attention_3d:
+    # KV_SPLITS is constexpr; tiles_per_segment derived at runtime. Splits
+    # whose tile range is past kv_len early-return WITHOUT writing — the
+    # reduce kernel uses the same constexpr BLOCK_K/KV_SPLITS to compute
+    # ``act_num_segments = cdiv(kv_len, tiles_per_segment * BLOCK_K)`` and
+    # masks unwritten slots out of its load. This saves a per-empty-split
+    # write of BLOCK_H*BLOCK_D fp32 (16 KB at BLOCK_H=16, BLOCK_D=512) which
+    # dominated short-K + many-splits cases.
+    tiles_per_segment = tl.cdiv(kv_len, KV_SPLITS * BLOCK_K)
+    if pid_k * tiles_per_segment * BLOCK_K >= kv_len:
+        return
+    num_tiles = tl.cdiv(kv_len, BLOCK_K)
+    tile_start = pid_k * tiles_per_segment
+    tile_end = tl.minimum((pid_k + 1) * tiles_per_segment, num_tiles)
+
+    neg_large = -3.4028234663852886e38
+    m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
+    l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
+
+    k_offs = tl.arange(0, BLOCK_K)
+    if QUANT_KV:
+        g_idx_per_d = d_offs // GROUP_SIZE
+    # num_stages=3 (see fused kernel comment for rationale).
+    for j in tl.range(tile_start, tile_end, num_stages=3):
+        k_start = j * BLOCK_K
+        k_pos = k_start + k_offs
+        valid = k_pos < kv_len  # in_range; no sentinel check (see contract)
+        slot = tl.load(
+            kv_indices_ptr + kv_start + k_pos,
+            mask=valid,
+            other=0,  # any in-bounds slot; masked out below
+        )
+
+        kv_raw = tl.load(
+            unified_kv_ptr
+            + slot[:, None] * kv_stride_n
+            + d_offs[None, :] * kv_stride_d,
+            mask=valid[:, None] & d_mask[None, :],
+            other=0.0,
+        )
+        if QUANT_KV:
+            scales_full = tl.load(
+                kv_scales_ptr + slot[:, None] * ks_stride_n + g_idx_per_d[None, :],
+                mask=valid[:, None] & d_mask[None, :],
+                other=0.0,
+            ).to(q.dtype)
+            kv = kv_raw.to(q.dtype) * scales_full
+        else:
+            kv = kv_raw
+
+        scores = tl.dot(q, tl.trans(kv)) * qk_scale
+        # K (same as fused kernel): drop h_mask from per-iter where.
+        scores = tl.where(valid[None, :], scores, neg_large)
+
+        m_block = tl.max(scores, axis=1)
+        m_new = tl.maximum(m_i, m_block)
+        alpha = tl.exp2(m_i - m_new)
+        p = tl.exp2(scores - m_new[:, None])
+        l_new = l_i * alpha + tl.sum(p, axis=1)
+
+        acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
+        m_i = m_new
+        l_i = l_new
+
+    m_base = t * mp_stride_t + pid_k * mp_stride_k
+    tl.store(m_partial_ptr + m_base + h_offs * mp_stride_h, m_i, mask=h_mask)
+    l_base = t * lp_stride_t + pid_k * lp_stride_k
+    tl.store(l_partial_ptr + l_base + h_offs * lp_stride_h, l_i, mask=h_mask)
+    a_base = t * ap_stride_t + pid_k * ap_stride_k
+    tl.store(
+        acc_partial_ptr
+        + a_base
+        + h_offs[:, None] * ap_stride_h
+        + d_offs[None, :] * ap_stride_d,
+        acc,
+        mask=h_mask[:, None] & d_mask[None, :],
+    )
+
+
+@triton.jit
+def _paged_decode_reduce_kernel(
+    m_partial_ptr,  # [N, KV_SPLITS, H_padded] fp32
+    l_partial_ptr,  # [N, KV_SPLITS, H_padded] fp32
+    acc_partial_ptr,  # [N, KV_SPLITS, H_padded, D] fp32
+    attn_sink_ptr,  # [H]
+    kv_indptr_ptr,  # [N+1] int32
+    out_ptr,  # [N, H, D]
+    mp_stride_t,
+    mp_stride_k,
+    mp_stride_h,
+    lp_stride_t,
+    lp_stride_k,
+    lp_stride_h,
+    ap_stride_t,
+    ap_stride_k,
+    ap_stride_h,
+    ap_stride_d,
+    out_stride_t,
+    out_stride_h,
+    out_stride_d,
+    log2e,  # = LOG2E, used to convert natural-log sink → log2 domain
+    H: tl.constexpr,
+    D: tl.constexpr,
+    KV_SPLITS: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    D_CHUNK: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """2D-tile reduce: combine KV_SPLITS partials, fold attn_sink, write
+    final output. Grid: ``(T, H, ceil(D / D_CHUNK))`` — one CTA owns one
+    (token, single-head, D-chunk) tuple.
+
+    Rewrite of the prior 3D-load reduce inspired by:
+      - aiter ``_fwd_kernel_stage2`` (mla_decode_rope.py): scalar control
+        flow + one D-tile per CTA, online accumulation across splits.
+      - AKO4X ``hybrid_2d_reduce`` (B200 reference, +3.37x): merged 2D
+        tile load ``[KV_SPLITS, D_CHUNK]`` replaces strided 3D access.
+
+    Why this wins on MI355 split path (T ≤ ~256, kv_splits > 1):
+
+    1. **2D tile fits registers**. The old 3D load tile
+       ``[KV_SPLITS=64, BLOCK_H=1, BLOCK_D=512]`` = 32K fp32 = 128 KB —
+       didn't fit in registers, spilled to LDS, and used a strided 3D
+       address compute. New ``[KV_SPLITS, D_CHUNK]`` = 64×64×4 = 16 KB
+       fits one wave's VGPR with headroom.
+    2. **D-chunked grid widens occupancy**. Old grid was ``(T, H)`` —
+       e.g. T=1 H=16 → 16 CTAs into 256 CUs (6% occupancy). New grid
+       ``(T, H, D/D_CHUNK)`` → 16 × (512/64) = 128 CTAs at T=1 H=16
+       (50% occupancy). Reduce was the latency bottleneck at small T.
+    3. **Scalar control flow for sink fold**. Sink computation needs
+       (m_max, l_combined) which only depend on (t, h) — same value
+       across all dc programs for the same (t, h). They recompute it
+       redundantly but it's a tiny scalar reduce; cheaper than passing
+       through LDS.
+    """
+    t = tl.program_id(0)
+    h = tl.program_id(1)
+    dc = tl.program_id(2)
+
+    d_offs = dc * D_CHUNK + tl.arange(0, D_CHUNK)
+    k_offs = tl.arange(0, KV_SPLITS)
+    d_mask = d_offs < D
+
+    neg_large = -3.4028234663852886e38
+
+    kv_start = tl.load(kv_indptr_ptr + t)
+    kv_end = tl.load(kv_indptr_ptr + t + 1)
+    kv_len = kv_end - kv_start
+    # CTA-level early return for empty tokens (CUDAGraph padding, or any
+    # caller-supplied zero-length slice). Split kernel skipped these without
+    # writing partials → partial buffers hold garbage; the segm_mask path
+    # below would still mask it correctly but consumes the masked-load BW
+    # and the sink-fold arithmetic. Skipping the whole CTA also halves the
+    # reduce-kernel cost on mixed-kv batches with many padded tokens.
+    if kv_len == 0:
+        out_off = t * out_stride_t + h * out_stride_h + d_offs * out_stride_d
+        tl.store(
+            out_ptr + out_off,
+            tl.zeros([D_CHUNK], dtype=out_ptr.dtype.element_ty),
+            mask=d_mask,
+        )
+        return
+    tiles_per_segment = tl.cdiv(kv_len, KV_SPLITS * BLOCK_K)
+    act_num_segments = tl.cdiv(kv_len, tl.maximum(tiles_per_segment, 1) * BLOCK_K)
+    segm_mask = k_offs < act_num_segments
+
+    # 1D loads for (m, l) along splits — single head h.
+    m_p = tl.load(
+        m_partial_ptr + t * mp_stride_t + k_offs * mp_stride_k + h * mp_stride_h,
+        mask=segm_mask,
+        other=neg_large,
+    )  # [KV_SPLITS]
+    l_p = tl.load(
+        l_partial_ptr + t * lp_stride_t + k_offs * lp_stride_k + h * lp_stride_h,
+        mask=segm_mask,
+        other=0.0,
+    )  # [KV_SPLITS]
+
+    # 2D-tile load for acc partials — the key change vs old 3D-strided load.
+    a_p = tl.load(
+        acc_partial_ptr
+        + t * ap_stride_t
+        + k_offs[:, None] * ap_stride_k
+        + h * ap_stride_h
+        + d_offs[None, :] * ap_stride_d,
+        mask=segm_mask[:, None] & d_mask[None, :],
+        other=0.0,
+    )  # [KV_SPLITS, D_CHUNK]
+
+    # Combine across splits.
+    m_max = tl.max(m_p, axis=0)  # scalar
+    alpha_split = tl.exp2(m_p - m_max)  # [KV_SPLITS]
+    l_combined = tl.sum(l_p * alpha_split, axis=0)  # scalar
+    acc_combined = tl.sum(a_p * alpha_split[:, None], axis=0)  # [D_CHUNK]
+
+    # Fold attn_sink (recomputed across dc — scalar work, negligible).
+    sink_raw = tl.load(attn_sink_ptr + h).to(tl.float32)
+    sink = sink_raw * log2e
+    m_final = tl.maximum(m_max, sink)
+    alpha_kv = tl.exp2(m_max - m_final)
+    alpha_sink = tl.exp2(sink - m_final)
+    l_final = l_combined * alpha_kv + alpha_sink
+
+    denom = tl.maximum(l_final, 1.0e-30)
+    # Direct divide (acc*alpha_kv)/denom, matching the single-CTA reference.
+    # The prior `acc * (alpha_kv/denom)` precomputed a reciprocal-scaled scalar
+    # (~1 extra ulp per element). Under split-K that per-kv_splits rounding
+    # diverges across batch shapes (T) and, via MTP greedy spec-acceptance,
+    # flips tokens — breaking MTP losslessness. The D-chunked multi-CTA layout
+    # (perf) is untouched; only the final normalize arithmetic changes.
+    acc_final = acc_combined * alpha_kv
+    out = tl.where(l_final > 0.0, acc_final / denom, 0.0)
+
+    tl.store(
+        out_ptr + t * out_stride_t + h * out_stride_h + d_offs * out_stride_d,
+        out.to(out_ptr.dtype.element_ty),
+        mask=d_mask,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Wrapper.
+# ---------------------------------------------------------------------------
+
+
+def _sparse_attn_v4_paged_decode_triton(
+    q: torch.Tensor,
+    unified_kv: torch.Tensor,
+    kv_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+    kv_scales: torch.Tensor | None = None,
+    block_h: int | None = None,
+    kv_splits: int | None = None,
+    block_k: int | None = None,
+) -> torch.Tensor:
+    """V4 sparse decode Triton implementation: split-K with FUSED fast path,
+    exp2 softmax, CG-safe heuristic. ``block_h`` and ``kv_splits`` are
+    escape hatches for benchmarks; production callers pass neither.
+
+    When ``kv_scales`` is provided, ``unified_kv`` must be e4m3fnuz and
+    ``kv_scales`` must be ``[total_pages, D // GROUP_SIZE]`` fp32 — 1xGROUP_SIZE
+    block-scale quantization. Dequant happens in-kernel; the dot still runs
+    in q.dtype.
+    """
+    if not q.is_cuda:
+        raise RuntimeError(
+            "Triton sparse_attn_v4_paged_decode requires CUDA/HIP tensors"
+        )
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        raise RuntimeError(
+            f"sparse_attn_v4_paged_decode expects fp16/bf16 q, got {q.dtype}"
+        )
+
+    quant_kv = kv_scales is not None
+    if quant_kv:
+        if unified_kv.dtype != _FP8_DTYPE:
+            raise RuntimeError(
+                f"kv_scales supplied but unified_kv is {unified_kv.dtype}, "
+                f"expected {_FP8_DTYPE}"
+            )
+        if kv_scales.dtype != torch.float32:
+            raise RuntimeError(f"kv_scales must be fp32, got {kv_scales.dtype}")
+        D_check = unified_kv.shape[-1]
+        if D_check % _FP8_GROUP_SIZE != 0:
+            raise RuntimeError(
+                f"D={D_check} must be divisible by GROUP_SIZE={_FP8_GROUP_SIZE}"
+            )
+        expected_g = D_check // _FP8_GROUP_SIZE
+        if kv_scales.shape != (unified_kv.shape[0], expected_g):
+            raise RuntimeError(
+                f"kv_scales shape {tuple(kv_scales.shape)} does not match "
+                f"expected ({unified_kv.shape[0]}, {expected_g})"
+            )
+        if kv_scales.stride(-1) != 1:
+            kv_scales = kv_scales.contiguous()
+    else:
+        if unified_kv.dtype != q.dtype:
+            raise RuntimeError(
+                f"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}"
+            )
+
+    T, H, D = q.shape
+    out = torch.empty_like(q)
+
+    if block_h is None:
+        block_h = triton.next_power_of_2(min(H, 64))
+    else:
+        block_h = triton.next_power_of_2(block_h)
+    block_h = max(block_h, 16)  # AMD MFMA min tile
+
+    n_head_blocks = (H + block_h - 1) // block_h
+    h_padded = n_head_blocks * block_h
+    block_d = triton.next_power_of_2(D)
+
+    if kv_splits is None:
+        kv_splits = _kv_splits_heuristic(T, H, block_h)
+
+    qk_scale = float(softmax_scale) * LOG2E
+    _bk, num_warps, num_stages = _kernel_config(block_h)
+    if block_k is None:
+        # fp8 dequant inflates per-tile ALU work ~4×; a wider K tile amortizes
+        # the per-tile dequant cost (scale load + cast + multiply) over more
+        # MFMA work. Empirically BLOCK_K=32 wins ~20% over BLOCK_K=16 on fp8
+        # (bs=512 ctx=4096: 3000µs → 2300µs) without hurting bf16.
+        block_k = 32 if quant_kv else _bk
+
+    # Kernel reads (kv_scales_ptr, ks_stride_n) only when QUANT_KV — supply a
+    # dummy 1-element fp32 tensor on the bf16 path so the launch signature
+    # stays uniform (avoids a separate JIT specialization per call).
+    if quant_kv:
+        kv_scales_arg = kv_scales
+        ks_stride_n_arg = kv_scales.stride(0)
+        num_groups_arg = D // _FP8_GROUP_SIZE
+    else:
+        kv_scales_arg = q.new_empty(1, dtype=torch.float32)
+        ks_stride_n_arg = 1
+        # NUM_GROUPS still needs a constexpr value (unused at compile time
+        # because the QUANT_KV=False branch elides the dequant code).
+        num_groups_arg = 1
+
+    # Fast path: when the base grid (T * n_head_blocks) already saturates the
+    # GPU, kv_splits=1 and a single-pass fused kernel beats split+reduce by
+    # skipping the partial-buffer alloc and the second kernel launch.
+    if kv_splits == 1:
+        grid_fused = (T, n_head_blocks)
+        _paged_decode_fused_kernel[grid_fused](
+            q,
+            unified_kv,
+            kv_scales_arg,
+            kv_indices,
+            kv_indptr,
+            attn_sink,
+            out,
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            unified_kv.stride(0),
+            unified_kv.stride(1),
+            ks_stride_n_arg,
+            out.stride(0),
+            out.stride(1),
+            out.stride(2),
+            qk_scale,
+            LOG2E,
+            H,
+            D,
+            BLOCK_H=block_h,
+            BLOCK_D=block_d,
+            BLOCK_K=block_k,
+            QUANT_KV=quant_kv,
+            GROUP_SIZE=_FP8_GROUP_SIZE,
+            NUM_GROUPS=num_groups_arg,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
+        return out
+
+    # Split-K path: split kernel writes (m, l, acc) partials in log2 domain;
+    # reduce kernel combines them, folds attn_sink, writes final output.
+    # Empty splits early-return without writing — reduce masks them out using
+    # the same constexpr BLOCK_K to derive ``act_num_segments``.
+    m_partial = torch.empty(
+        (T, kv_splits, h_padded), dtype=torch.float32, device=q.device
+    )
+    l_partial = torch.empty_like(m_partial)
+    acc_partial = torch.empty(
+        (T, kv_splits, h_padded, D), dtype=torch.float32, device=q.device
+    )
+
+    grid_split = (T, n_head_blocks, kv_splits)
+    _paged_decode_split_kernel[grid_split](
+        q,
+        unified_kv,
+        kv_scales_arg,
+        kv_indices,
+        kv_indptr,
+        m_partial,
+        l_partial,
+        acc_partial,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        unified_kv.stride(0),
+        unified_kv.stride(1),
+        ks_stride_n_arg,
+        m_partial.stride(0),
+        m_partial.stride(1),
+        m_partial.stride(2),
+        l_partial.stride(0),
+        l_partial.stride(1),
+        l_partial.stride(2),
+        acc_partial.stride(0),
+        acc_partial.stride(1),
+        acc_partial.stride(2),
+        acc_partial.stride(3),
+        H,
+        D,
+        kv_splits,
+        qk_scale,
+        BLOCK_H=block_h,
+        BLOCK_D=block_d,
+        BLOCK_K=block_k,
+        QUANT_KV=quant_kv,
+        GROUP_SIZE=_FP8_GROUP_SIZE,
+        NUM_GROUPS=num_groups_arg,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+
+    # 2D-tile reduce: grid = (T, H, ceil(D/D_CHUNK)). One CTA per
+    # (token, single-head, D-chunk). Adaptive D_CHUNK based on whether the
+    # base reduce grid (T*H) already saturates the GPU:
+    #   - large T*H (≥ 2*num_CU=512 on MI355): the grid is already at or
+    #     above the launch target; further D-splitting just over-fragments
+    #     (T=128 H=128 with D_CHUNK=64 → 131K CTAs → 2× slowdown). Use
+    #     D_CHUNK = block_d (1 d-block per CTA, grid = T*H).
+    #   - small T*H: D-split widens the grid to fill CUs at small-T
+    #     latency-critical decode. The 2D tile [KV_SPLITS, D_CHUNK] should
+    #     stay ≤ 16 KB fp32 to fit registers.
+    base_grid_t_h = T * H
+    target_reduce_wg = 2 * _cu_count()
+    if base_grid_t_h >= target_reduce_wg:
+        d_chunk = block_d
+    else:
+        d_chunks_needed = max(1, target_reduce_wg // base_grid_t_h)
+        d_chunks_needed = min(d_chunks_needed, block_d // 32)
+        d_chunk = max(32, triton.next_power_of_2(block_d // d_chunks_needed))
+    grid_reduce = (T, H, (D + d_chunk - 1) // d_chunk)
+    _paged_decode_reduce_kernel[grid_reduce](
+        m_partial,
+        l_partial,
+        acc_partial,
+        attn_sink,
+        kv_indptr,
+        out,
+        m_partial.stride(0),
+        m_partial.stride(1),
+        m_partial.stride(2),
+        l_partial.stride(0),
+        l_partial.stride(1),
+        l_partial.stride(2),
+        acc_partial.stride(0),
+        acc_partial.stride(1),
+        acc_partial.stride(2),
+        acc_partial.stride(3),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        LOG2E,
+        H,
+        D,
+        kv_splits,
+        BLOCK_D=block_d,
+        D_CHUNK=d_chunk,
+        BLOCK_K=block_k,
+        num_warps=4,
+    )
+    return out
+
+
+def sparse_attn_v4_paged_decode(
+    q: torch.Tensor,
+    unified_kv: torch.Tensor,
+    kv_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+    kv_scales: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """V4 decode sparse attention over a unified KV pool with paged indices.
+
+    When ``kv_scales`` is provided, ``unified_kv`` must be fp8 (e4m3fnuz) and
+    will be dequantized in-kernel using 1xGROUP_SIZE (default 64) block scales.
+    """
+    return _sparse_attn_v4_paged_decode_triton(
+        q,
+        unified_kv,
+        kv_indices,
+        kv_indptr,
+        attn_sink,
+        softmax_scale,
+        kv_scales=kv_scales,
+    )

python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""V4 paged-decode index scatter — single Triton kernel writes SWA window-
+prefix paged offsets into the three ragged-packed destination buffers
+(`kv_indices_swa` / `kv_indices_csa` / `kv_indices_hca`).
+
+Replaces the prior chain (numpy `_build_window_topk_np` + `index_copy_`):
+
+    window_topk_np = _build_window_topk_np(positions, win, ring_stride)  # [T, win]
+    swa_paged_2d = torch.where(window_topk >= 0, slot * ring_stride + topk, -1)
+    swa_paged_flat = swa_paged_2d.reshape(-1)
+    swa_indices_gpu[:T*win].copy_(swa_paged_flat)
+    csa_indices_gpu.index_copy_(0, csa_win_pos, swa_paged_flat)
+    hca_indices_gpu.index_copy_(0, hca_win_pos, swa_paged_flat)
+
+Two simplifications vs the prior implementation (see plan
+`sequential-noodling-turing.md` for details):
+
+1. The ring-index formula `ring = (pos - win + 1 + w) % ring_stride` is computed
+   inline inside the kernel from `positions[t]`. The `[T, win]`
+   `window_topk` intermediate buffer (mnbt·win·4 = 4 MB at typical config)
+   is gone; no separate CPU build + H2D copy.
+2. The destination layout is now ragged-packed (same as prefill): each
+   token's SWA prefix segment has length `n = min(positions[t]+1, win)`
+   (NOT a fixed `win` padded with `-1` sentinels). The caller's
+   `swa_indptr` / `csa_indptr` / `hca_indptr` reflect this ragged sizing.
+
+Bytewise correctness: for tokens with `position >= win-1` (all `n == win`),
+the output is identical to the prior implementation. For shorter
+positions, the prior layout wrote `(win - n)` leading `-1` entries that
+the sparse-attention kernel masked out; the new layout omits those slots
+entirely, saving sparse-attn loop iterations.
+
+Caller contract:
+- Grid = T (one program per token).
+- `batch_id_per_token[:T]` may carry `-1` sentinels in the CG-padded tail —
+  kernel checks and bails (matches `_attach_v4_per_fwd_meta` convention).
+- `swa_indptr` / `csa_indptr` / `hca_indptr` must reflect the ragged-packed
+  sizing: per-token slot count = `min(positions[t]+1, win) + n_compress[t]`
+  where `n_compress[t]` is 0 for SWA, `min(n_committed_csa, index_topk)`
+  for CSA, `n_committed_hca` for HCA.
+- `swa_indices` / `csa_indices` / `hca_indices` capacity ≥ corresponding
+  indptr[T]; this kernel only writes the SWA-prefix segment
+  `[indptr[t], indptr[t] + n)` per token. The compress-tail is filled
+  elsewhere (HCA: numpy fill in caller, CSA: `csa_translate_pack` per layer).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _v4_paged_decode_indices_kernel(
+    state_slot_per_seq_ptr,  # [bs] int32
+    batch_id_per_token_ptr,  # [T+pad] int — sentinel -1 in pad tail
+    positions_ptr,  # [T+pad] int — global token position
+    swa_indptr_ptr,  # [T+1] int32 — ragged SWA-prefix cumsum
+    csa_indptr_ptr,  # [T+1] int32 — ragged (SWA + CSA topk)
+    hca_indptr_ptr,  # [T+1] int32 — ragged (SWA + HCA committed)
+    swa_indices_ptr,  # [swa_total] int32, output
+    csa_indices_ptr,  # [csa_total] int32, output (writes SWA-prefix segment only)
+    hca_indices_ptr,  # [hca_total] int32, output (writes SWA-prefix segment only)
+    ring_stride,  # win_with_spec — stride into unified_kv SWA region (paper §3.6.1)
+    win: tl.constexpr,  # window_size — max SWA prefix slots
+    BLOCK_N: tl.constexpr,  # next_pow2(win)
+):
+    """One program per token. Writes `n = min(positions[t]+1, win)` paged
+    offsets to the SWA prefix segment of each of SWA/CSA/HCA index buffers.
+
+    For token `t`:
+        bid = batch_id_per_token[t]                  # bail if -1 (CG pad)
+        slot = state_slot_per_seq[bid]
+        pos = positions[t]
+        n = min(pos + 1, win)
+        # Old -1 sentinels were at the leading `win - n` cols; reparameterize
+        # to skip them: i in [0, n) → abs_pos = pos - n + 1 + i ∈ [0, pos].
+        for i in range(n):
+            abs_pos = pos - n + 1 + i
+            ring = abs_pos % ring_stride
+            paged = slot * ring_stride + ring
+            swa_indices[swa_indptr[t] + i] = paged
+            csa_indices[csa_indptr[t] + i] = paged
+            hca_indices[hca_indptr[t] + i] = paged
+    """
+    t = tl.program_id(0)
+    bid = tl.load(batch_id_per_token_ptr + t)
+    if bid < 0:
+        return  # CG-padded sentinel — leave outputs untouched
+
+    slot = tl.load(state_slot_per_seq_ptr + bid)
+    pos = tl.load(positions_ptr + t)
+    # `n` = actual valid SWA prefix count. Cast to match `win` (compile-time
+    # int) — pos is i32/i64 from positions buffer.
+    n = tl.minimum(pos + 1, win)
+    swa_base = tl.load(swa_indptr_ptr + t)
+    csa_base = tl.load(csa_indptr_ptr + t)
+    hca_base = tl.load(hca_indptr_ptr + t)
+
+    i = tl.arange(0, BLOCK_N)
+    mask = i < n
+    abs_pos = pos - n + 1 + i  # ∈ [0, pos] for valid i
+    ring_idx = abs_pos % ring_stride
+    paged = slot * ring_stride + ring_idx
+
+    tl.store(swa_indices_ptr + swa_base + i, paged, mask=mask)
+    tl.store(csa_indices_ptr + csa_base + i, paged, mask=mask)
+    tl.store(hca_indices_ptr + hca_base + i, paged, mask=mask)
+
+
+def write_v4_paged_decode_indices(
+    *,
+    state_slot_per_seq: torch.Tensor,
+    batch_id_per_token: torch.Tensor,
+    positions: torch.Tensor,
+    swa_indptr: torch.Tensor,
+    csa_indptr: torch.Tensor,
+    hca_indptr: torch.Tensor,
+    swa_indices: torch.Tensor,
+    csa_indices: torch.Tensor,
+    hca_indices: torch.Tensor,
+    T: int,
+    win: int,
+    ring_stride: int,
+) -> None:
+    """In-place fill SWA / CSA / HCA window-prefix offsets via a single
+    Triton kernel. Replaces the prior `_build_window_topk_np` (CPU O(T·win))
+    + `index_copy_` chain. All inputs are persistent forward_vars buffers —
+    no allocator churn.
+
+    Args (all GPU tensors except T/win/ring_stride):
+      state_slot_per_seq:  [bs]   int32 — per-seq state cache slot.
+      batch_id_per_token:  [>=T]  int   — token→seq map; -1 sentinel skipped.
+      positions:           [>=T]  int   — global token position
+                                   (forward_vars["positions"]); used to derive
+                                   `n = min(pos+1, win)` per token + the ring
+                                   index `(pos - n + 1 + i) % ring_stride`.
+      swa_indptr:          [>=T+1] int32 — ragged SWA-prefix cumsum, where
+                                   `swa_indptr[t+1] - swa_indptr[t] =
+                                    min(positions[t]+1, win)`.
+      csa_indptr:          [>=T+1] int32 — ragged CSA buffer indptr (SWA
+                                   prefix + CSA topk per token).
+      hca_indptr:          [>=T+1] int32 — ragged HCA buffer indptr (SWA
+                                   prefix + HCA committed per token).
+      swa_indices:         [>=swa_indptr[T]] int32 OUT — fully written by
+                                   this kernel (no other source).
+      csa_indices:         [>=csa_indptr[T]] int32 OUT — window-prefix
+                                   `[csa_indptr[t], +n)` written here; CSA
+                                   topk tail filled per-layer by
+                                   `csa_translate_pack`.
+      hca_indices:         [>=hca_indptr[T]] int32 OUT — same semantics; HCA
+                                   compress tail filled in the caller via
+                                   numpy fill.
+      T:                   int — number of real tokens (grid size).
+      win:                 int — SWA window size (typically 128 for V4-Pro).
+      ring_stride:                  int — `win_with_spec = window_size + max_spec_steps`,
+                                 stride into unified_kv SWA region per slot
+                                 AND modulo for ring-index wrap.
+    """
+    if T == 0:
+        return
+    assert state_slot_per_seq.dim() == 1
+    assert batch_id_per_token.dim() == 1 and batch_id_per_token.shape[0] >= T
+    assert positions.dim() == 1 and positions.shape[0] >= T
+    assert swa_indptr.dim() == 1 and swa_indptr.shape[0] >= T + 1
+    assert csa_indptr.dim() == 1 and csa_indptr.shape[0] >= T + 1
+    assert hca_indptr.dim() == 1 and hca_indptr.shape[0] >= T + 1
+    assert swa_indices.dim() == 1
+    assert csa_indices.dim() == 1
+    assert hca_indices.dim() == 1
+
+    BLOCK_N = triton.next_power_of_2(win)
+    _v4_paged_decode_indices_kernel[(T,)](
+        state_slot_per_seq,
+        batch_id_per_token,
+        positions,
+        swa_indptr,
+        csa_indptr,
+        hca_indptr,
+        swa_indices,
+        csa_indices,
+        hca_indices,
+        ring_stride,
+        win=win,
+        BLOCK_N=BLOCK_N,
+    )

python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py

@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Sparse prefill attention with two KV sources: paged `unified_kv` (history)
+and per-fwd flat `kv` (current chunk's input).
+
+Designed for V4 prefill: indexes the two KV sources directly without
+materialising a per-fwd `kv_flat_sa` packed tensor.
+
+Caller contract:
+  unified_kv:        [total_pages, D] BF16 — prefix source. Same buffer as
+    decode kernel: SWA ring slots in `[0, swa_pages)`, compress pages in
+    `[swa_pages, total_pages)`. For prefill, prefix indices select
+    (a) prior-chunk SWA history, (b) CSA topk, (c) HCA all-committed.
+  kv_indices_prefix: [total_prefix_indices] int32 — flat per-token slot
+    lists. Per-token entries live in
+    `kv_indices_prefix[kv_indptr_prefix[t] : kv_indptr_prefix[t+1]]`.
+    `-1` entries are skipped (sentinel).
+  kv_indptr_prefix:  [N+1] int32 — true prefix sum (variable per-token len).
+
+  kv:                [total_tokens, D] BF16 — extend source = current
+    fwd's just-computed K (NOT yet written to swa_kv ring). Layout matches
+    `swa_write` input.
+  kv_indices_extend: [total_extend_indices] int32 — flat per-token row idx
+    lists into `kv`. Per-token entries live in
+    `kv_indices_extend[kv_indptr_extend[t] : kv_indptr_extend[t+1]]`.
+    `-1` entries are skipped (rare for extend; usually all valid).
+  kv_indptr_extend:  [N+1] int32 — true prefix sum.
+
+  attn_sink:         [H] per-head learnable softmax-denom bias (V4 specific).
+  softmax_scale:     float.
+
+Per-token K loop iterates two regions sequentially, sharing the online
+softmax accumulator (m_i, l_i, acc) across regions. Order of regions does
+not affect correctness (online softmax is order-invariant).
+
+Returns:
+  out: [N, H, D] same dtype as q.
+
+Numerics: identical online-softmax + sink finalization to
+`sparse_attn_v4_paged_decode` — bit-exact when the extend region is empty
+(then equivalent to a decode call with the same prefix indices).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+# OPUS gfx950 paged-prefill kernel is preferred when importable; otherwise fall
+# back to the Triton implementation below.
+try:
+    from aiter.ops.pa_sparse_prefill_opus import pa_sparse_prefill_opus
+
+    _HAS_OPUS = True
+except ImportError:
+    pa_sparse_prefill_opus = None
+    _HAS_OPUS = False
+
+
+@triton.jit
+def _sparse_attn_v4_paged_prefill_kernel(
+    q_ptr,  # [N, H, D]
+    unified_kv_ptr,  # [total_pages, D]   — prefix source
+    kv_indices_prefix_ptr,  # [total_prefix_indices] int32
+    kv_indptr_prefix_ptr,  # [N+1] int32
+    kv_ptr,  # [total_tokens, D]    — extend source
+    kv_indices_extend_ptr,  # [total_extend_indices] int32
+    kv_indptr_extend_ptr,  # [N+1] int32
+    attn_sink_ptr,  # [H]
+    out_ptr,  # [N, H, D]
+    q_stride_t: tl.constexpr,
+    q_stride_h: tl.constexpr,
+    q_stride_d: tl.constexpr,
+    pkv_stride_n: tl.constexpr,  # unified_kv stride 0 (= D usually)
+    pkv_stride_d: tl.constexpr,  # unified_kv stride 1 (= 1 usually)
+    ekv_stride_n: tl.constexpr,  # kv stride 0
+    ekv_stride_d: tl.constexpr,  # kv stride 1
+    out_stride_t: tl.constexpr,
+    out_stride_h: tl.constexpr,
+    out_stride_d: tl.constexpr,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    softmax_scale: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    t = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+    d_offs = tl.arange(0, BLOCK_D)
+    h_mask = h_offs < H
+    d_mask = d_offs < D
+
+    q = tl.load(
+        q_ptr
+        + t * q_stride_t
+        + h_offs[:, None] * q_stride_h
+        + d_offs[None, :] * q_stride_d,
+        mask=h_mask[:, None] & d_mask[None, :],
+        other=0.0,
+    )
+
+    neg_large = -3.4028234663852886e38
+    m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
+    l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
+
+    k_offs = tl.arange(0, BLOCK_K)
+
+    # ===== Region 1: prefix from unified_kv =====
+    p_start = tl.load(kv_indptr_prefix_ptr + t)
+    p_end = tl.load(kv_indptr_prefix_ptr + t + 1)
+    p_len = p_end - p_start
+
+    for k_start in tl.range(0, p_len, BLOCK_K):
+        k_pos = k_start + k_offs
+        in_range = k_pos < p_len
+        slot = tl.load(
+            kv_indices_prefix_ptr + p_start + k_pos,
+            mask=in_range,
+            other=-1,
+        )
+        valid = in_range & (slot >= 0)
+        slot_clamped = tl.maximum(slot, 0)
+
+        kv = tl.load(
+            unified_kv_ptr
+            + slot_clamped[:, None] * pkv_stride_n
+            + d_offs[None, :] * pkv_stride_d,
+            mask=valid[:, None] & d_mask[None, :],
+            other=0.0,
+        )
+
+        scores = tl.dot(q, tl.trans(kv)) * softmax_scale
+        scores = tl.where(h_mask[:, None] & valid[None, :], scores, neg_large)
+
+        m_block = tl.max(scores, axis=1)
+        m_new = tl.maximum(m_i, m_block)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(scores - m_new[:, None])
+        p = tl.where(h_mask[:, None] & valid[None, :], p, 0.0)
+        l_new = l_i * alpha + tl.sum(p, axis=1)
+
+        acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
+        m_i = m_new
+        l_i = l_new
+
+    # ===== Region 2: extend from kv (per-fwd flat) =====
+    e_start = tl.load(kv_indptr_extend_ptr + t)
+    e_end = tl.load(kv_indptr_extend_ptr + t + 1)
+    e_len = e_end - e_start
+
+    for k_start in tl.range(0, e_len, BLOCK_K):
+        k_pos = k_start + k_offs
+        in_range = k_pos < e_len
+        slot = tl.load(
+            kv_indices_extend_ptr + e_start + k_pos,
+            mask=in_range,
+            other=-1,
+        )
+        valid = in_range & (slot >= 0)
+        slot_clamped = tl.maximum(slot, 0)
+
+        kv = tl.load(
+            kv_ptr
+            + slot_clamped[:, None] * ekv_stride_n
+            + d_offs[None, :] * ekv_stride_d,
+            mask=valid[:, None] & d_mask[None, :],
+            other=0.0,
+        )
+
+        scores = tl.dot(q, tl.trans(kv)) * softmax_scale
+        scores = tl.where(h_mask[:, None] & valid[None, :], scores, neg_large)
+
+        m_block = tl.max(scores, axis=1)
+        m_new = tl.maximum(m_i, m_block)
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(scores - m_new[:, None])
+        p = tl.where(h_mask[:, None] & valid[None, :], p, 0.0)
+        l_new = l_i * alpha + tl.sum(p, axis=1)
+
+        acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
+        m_i = m_new
+        l_i = l_new
+
+    # ===== Sink finalization =====
+    # Online softmax + sink integration: sink is a virtual extra K with V=0,
+    # contributing only to the denominator. After main loops, (m_i, l_i, acc)
+    # are in m_i frame; sink may shift max to m_final = max(m_i, sink), so
+    # rescale BOTH l_i (for denom) AND acc (for numerator) by alpha to switch
+    # to m_final frame. The sink itself adds exp(sink - m_final) to l_final
+    # but contributes 0 to acc since V_sink = 0.
+    sink = tl.load(attn_sink_ptr + h_offs, mask=h_mask, other=neg_large).to(tl.float32)
+    m_final = tl.maximum(m_i, sink)
+    alpha = tl.exp(m_i - m_final)
+    l_final = l_i * alpha + tl.exp(sink - m_final)
+
+    denom = tl.maximum(l_final, 1.0e-30)
+    out = tl.where(l_final[:, None] > 0.0, (acc * alpha[:, None]) / denom[:, None], 0.0)
+    tl.store(
+        out_ptr
+        + t * out_stride_t
+        + h_offs[:, None] * out_stride_h
+        + d_offs[None, :] * out_stride_d,
+        out,
+        mask=h_mask[:, None] & d_mask[None, :],
+    )
+
+
+def _sparse_attn_v4_paged_prefill_triton(
+    q: torch.Tensor,
+    unified_kv: torch.Tensor,
+    kv_indices_prefix: torch.Tensor,
+    kv_indptr_prefix: torch.Tensor,
+    kv: torch.Tensor,
+    kv_indices_extend: torch.Tensor,
+    kv_indptr_extend: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+) -> torch.Tensor:
+    if not q.is_cuda:
+        raise RuntimeError(
+            "Triton sparse_attn_v4_paged_prefill requires CUDA/HIP tensors"
+        )
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        raise RuntimeError(
+            f"sparse_attn_v4_paged_prefill expects fp16/bf16 q, got {q.dtype}"
+        )
+    if unified_kv.dtype != q.dtype:
+        raise RuntimeError(
+            f"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}"
+        )
+    if kv.dtype != q.dtype:
+        raise RuntimeError(f"kv dtype mismatch: kv={kv.dtype}, q={q.dtype}")
+    if unified_kv.size(-1) != kv.size(-1):
+        raise RuntimeError(
+            f"head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={kv.size(-1)}"
+        )
+
+    T, H, D = q.shape
+    out = torch.empty_like(q)
+    kv_indices_prefix = kv_indices_prefix.to(torch.int32).contiguous()
+    kv_indptr_prefix = kv_indptr_prefix.to(torch.int32).contiguous()
+    kv_indices_extend = kv_indices_extend.to(torch.int32).contiguous()
+    kv_indptr_extend = kv_indptr_extend.to(torch.int32).contiguous()
+
+    block_h = 16  # AMD MFMA min tile
+    block_d = triton.next_power_of_2(D)
+    block_k = 16 if D >= 256 else 32
+    _sparse_attn_v4_paged_prefill_kernel[(T, triton.cdiv(H, block_h))](
+        q,
+        unified_kv,
+        kv_indices_prefix,
+        kv_indptr_prefix,
+        kv,
+        kv_indices_extend,
+        kv_indptr_extend,
+        attn_sink,
+        out,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        unified_kv.stride(0),
+        unified_kv.stride(1),
+        kv.stride(0),
+        kv.stride(1),
+        out.stride(0),
+        out.stride(1),
+        out.stride(2),
+        H,
+        D,
+        float(softmax_scale),
+        BLOCK_H=block_h,
+        BLOCK_D=block_d,
+        BLOCK_K=block_k,
+        num_warps=8,
+    )
+    return out
+
+
+def sparse_attn_v4_paged_prefill(
+    q: torch.Tensor,
+    unified_kv: torch.Tensor,
+    kv_indices_prefix: torch.Tensor,
+    kv_indptr_prefix: torch.Tensor,
+    kv: torch.Tensor,
+    kv_indices_extend: torch.Tensor,
+    kv_indptr_extend: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+) -> torch.Tensor:
+    """V4 prefill sparse attention over two KV sources (paged unified_kv +
+    flat per-fwd kv).
+
+    Args:
+      q:                 [T, H, D] BF16/FP16 — query.
+      unified_kv:        [total_pages, D] BF16/FP16 — prefix source (paged).
+      kv_indices_prefix: [total_prefix] int32 — flat per-token slot lists into
+        unified_kv. -1 sentinels skipped.
+      kv_indptr_prefix:  [T+1] int32 — true prefix sum.
+      kv:                [total_tokens, D] BF16/FP16 — extend source (this
+        fwd's input K, NOT yet in swa_kv ring).
+      kv_indices_extend: [total_extend] int32 — flat per-token row idx lists
+        into kv. -1 sentinels skipped.
+      kv_indptr_extend:  [T+1] int32 — true prefix sum.
+      attn_sink:         [H] — per-head softmax-denom bias.
+      softmax_scale:     float.
+
+    Returns:
+      out: [T, H, D] same dtype as q.
+    """
+    if _HAS_OPUS:
+        # OPUS contract differs from the Triton kernel in two ways the Triton
+        # path tolerates implicitly:
+        #  - it requires a FULLY-contiguous q (it only asserts stride(2)==1 but
+        #    indexes assuming [T,H,D] contiguous); a non-contiguous head stride
+        #    silently reads wrong/out-of-bounds addresses. q from the model is
+        #    often a view, so force contiguity.
+        #  - it requires ``attn_sink.size(0) == H``; attn_sink is the full
+        #    per-head Parameter, so slice to the H query heads.
+        q = q.contiguous()
+        H = q.shape[1]
+        if attn_sink.shape[0] != H:
+            attn_sink = attn_sink[:H].contiguous()
+        if (
+            kv.stride(0) != unified_kv.stride(0)
+            and kv.shape[0] == 1
+            and kv.stride(1) == 1
+        ):
+            kv = kv.as_strided(kv.shape, (kv.shape[1], 1))
+        return pa_sparse_prefill_opus(
+            q,
+            unified_kv,
+            kv_indices_prefix,
+            kv_indptr_prefix,
+            kv,
+            kv_indices_extend,
+            kv_indptr_extend,
+            attn_sink,
+            softmax_scale,
+        )
+    return _sparse_attn_v4_paged_prefill_triton(
+        q,
+        unified_kv,
+        kv_indices_prefix,
+        kv_indptr_prefix,
+        kv,
+        kv_indices_extend,
+        kv_indptr_extend,
+        attn_sink,
+        softmax_scale,
+    )

python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py

@@ -0,0 +1,440 @@
+"""Runtime glue for the unified_kv backend.
+
+Builds unified_kv-style flat ``kv_indices`` / ``kv_indptr`` from SGLang's already-computed
+DSV4 metadata, scatters SWA K into the bf16 ``unified_kv`` ring, and dispatches the
+vendored paged decode/prefill kernels.
+
+unified_kv[L] layout (page_size 1, bf16, row-major):
+  - rows ``[0, swa_pages)``    = SWA ring (``state_slot * win + pos % win``);
+  - rows ``[swa_pages, ...)``  = compressed K (``swa_pages + page_index``), where
+    SGLang metadata already encodes the compressed slot id:
+      HCA (ratio 128): ``c128_page_indices``      (== phys_block, k_per_block=1)
+      CSA (ratio   4): ``c4_sparse_page_indices``  (== phys_block*32 + slot)
+
+Index layout: RAGGED-PACKED. Each token's segment is tightly packed
+(``kv_indptr`` is a true prefix sum of per-token valid lengths) so the
+attention K-loop scans only real entries. The backing buffer is still
+allocated at the fixed worst-case capacity ``N * (win + Wc)`` so its shape is
+static across CUDA-graph replay; only ``kv_indptr`` values (and the written
+prefix) vary per forward. Compressed valid entries are front-packed in the
+``*_page_indices`` rows (the same contract the non-unified_kv flashmla path relies
+on via ``topk_length``); the per-token compressed count is recovered from the
+``kv_indptr`` delta inside the kernel, so no extra length tensor is threaded.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_decode import (
+    sparse_attn_v4_paged_decode,
+)
+from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_decode_indices import (
+    write_v4_paged_decode_indices,
+)
+from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_prefill import (
+    sparse_attn_v4_paged_prefill,
+)
+
+
+# ---------------------------------------------------------------------------
+# SWA ring scatter
+# ---------------------------------------------------------------------------
+@triton.jit
+def _swa_scatter_kernel(
+    kv_ptr,  # [T, D] bf16
+    state_slot_ptr,  # [T] int
+    positions_ptr,  # [T] int
+    final_pos_ptr,  # [T] int
+    unified_ptr,  # [pages, D] bf16
+    n_rows,
+    ring_stride,  # SWA ring per-slot stride
+    win: tl.constexpr,
+    D: tl.constexpr,
+    HAS_FINAL: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    row = tl.program_id(0)
+    if row >= n_rows:
+        return
+    pos = tl.load(positions_ptr + row)
+    if HAS_FINAL:
+        fp = tl.load(final_pos_ptr + row)
+        if pos <= fp - win:
+            return
+    s = tl.load(state_slot_ptr + row)
+    loc = s * ring_stride + (pos % ring_stride)
+    offs = tl.arange(0, BLOCK_D)
+    mask = offs < D
+    vals = tl.load(kv_ptr + row * D + offs, mask=mask, other=0.0)
+    tl.store(unified_ptr + loc * D + offs, vals, mask=mask)
+
+
+def store_swa_into_unified(
+    *,
+    kv: torch.Tensor,  # [T, head_dim] bf16
+    state_slot: torch.Tensor,  # [T] int
+    positions: torch.Tensor,  # [T] int
+    unified_kv: torch.Tensor,  # [pages, head_dim] bf16
+    win: int,  # SWA attention window length
+    ring_stride: int,  # SWA ring stride
+    final_pos: Optional[torch.Tensor] = None,  # [T] req's last position
+) -> None:
+    n_rows, D = kv.shape
+    if n_rows == 0:
+        return
+
+    has_final = final_pos is not None
+    fp_arg = final_pos if has_final else positions
+    assert kv.is_contiguous() and kv.dtype == unified_kv.dtype
+    assert state_slot.is_contiguous() and positions.is_contiguous()
+    assert fp_arg.is_contiguous()
+    _swa_scatter_kernel[(n_rows,)](
+        kv,
+        state_slot,
+        positions,
+        fp_arg,
+        unified_kv,
+        n_rows,
+        ring_stride,
+        win=win,
+        D=D,
+        HAS_FINAL=has_final,
+        BLOCK_D=triton.next_power_of_2(D),
+        num_warps=8,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Ragged indptr helper (shared by the decode streams + prefill builders)
+# ---------------------------------------------------------------------------
+def _lengths_to_indptr(lengths: torch.Tensor) -> torch.Tensor:
+    """[N] int32 per-token lengths -> [N+1] int32 indptr"""
+    return F.pad(torch.cumsum(lengths, dim=0, dtype=torch.int32), (1, 0))
+
+
+def decode(
+    *,
+    q: torch.Tensor,  # [T, H, D] (local heads)
+    unified_kv: torch.Tensor,  # [pages, D] bf16
+    kv_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    attn_sink: torch.Tensor,  # [H] fp32
+    softmax_scale: float,
+) -> torch.Tensor:
+    return sparse_attn_v4_paged_decode(
+        q, unified_kv, kv_indices, kv_indptr, attn_sink, softmax_scale
+    )
+
+
+@triton.jit
+def _fill_compress_tail_kernel(
+    indices_ptr,  # [*] int32 (out)
+    indptr_ptr,  # [N+1] int32
+    prefix_len_ptr,  # [N] int
+    page_idx_ptr,  # [N, Wc] int
+    valid_len_ptr,  # [N] int
+    swa_pages,
+    Wc: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Per token: write valid_len compressed slots (swa_pages+page_idx, -1 for empty) into the stream tail at indptr[t]+prefix_len[t]."""
+    t = tl.program_id(0)
+    cbase = tl.load(indptr_ptr + t) + tl.load(prefix_len_ptr + t).to(tl.int32)
+    nc = tl.load(valid_len_ptr + t).to(tl.int32)
+    for off in tl.range(0, Wc, BLOCK):
+        j = off + tl.arange(0, BLOCK)
+        m = j < nc
+        j_clamped = tl.minimum(j, Wc - 1)
+        pi = tl.load(page_idx_ptr + t * Wc + j_clamped, mask=m, other=-1).to(tl.int32)
+        slot = tl.where(pi >= 0, pi + swa_pages, -1)
+        tl.store(indices_ptr + cbase + j, slot, mask=m)
+
+
+def fill_compress_tail(
+    *,
+    indices: torch.Tensor,
+    indptr: torch.Tensor,
+    prefix_len: torch.Tensor,
+    page_indices: torch.Tensor,  # [N, Wc] int32
+    valid_len: torch.Tensor,
+    swa_pages: int,
+) -> None:
+    N, Wc = page_indices.shape
+    if N == 0:
+        return
+    assert prefix_len.is_contiguous() and page_indices.is_contiguous()
+    assert valid_len.is_contiguous()
+    _fill_compress_tail_kernel[(N,)](
+        indices,
+        indptr,
+        prefix_len,
+        page_indices,
+        valid_len,
+        swa_pages,
+        Wc=Wc,
+        BLOCK=min(1024, triton.next_power_of_2(max(Wc, 1))),
+        num_warps=4,
+    )
+
+
+def build_decode_streams(
+    *,
+    state_slot: torch.Tensor,  # [N] int
+    positions: torch.Tensor,  # [N] int
+    swa_len: torch.Tensor,  # [N] int
+    hca_len: torch.Tensor,  # [N] int
+    csa_len: torch.Tensor,  # [N] int
+    hca_page_indices: torch.Tensor,  # [N, hca_width] int32
+    csa_width: int,
+    win: int,  # SWA attention window length
+    ring_stride: int,  # SWA ring per-slot stride
+    swa_pages: int,
+) -> Tuple[
+    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+]:
+    device = state_slot.device
+    N = state_slot.shape[0]
+    assert state_slot.is_contiguous() and positions.is_contiguous()
+    state_slot = state_slot.to(torch.int32)
+    positions = positions.to(torch.int32)
+    hca_width = hca_page_indices.shape[1]
+
+    swa_p = _lengths_to_indptr(swa_len)
+    hca_p = _lengths_to_indptr(swa_len + hca_len)
+    csa_p = _lengths_to_indptr(swa_len + csa_len)
+
+    swa_i = torch.empty(N * win, dtype=torch.int32, device=device)
+    hca_i = torch.empty(N * (win + hca_width), dtype=torch.int32, device=device)
+    csa_i = torch.empty(N * (win + csa_width), dtype=torch.int32, device=device)
+
+    if N > 0:
+        batch_id = torch.arange(N, dtype=torch.int32, device=device)
+        write_v4_paged_decode_indices(
+            state_slot_per_seq=state_slot,
+            batch_id_per_token=batch_id,
+            positions=positions,
+            swa_indptr=swa_p,
+            csa_indptr=csa_p,
+            hca_indptr=hca_p,
+            swa_indices=swa_i,
+            csa_indices=csa_i,
+            hca_indices=hca_i,
+            T=N,
+            win=win,
+            ring_stride=ring_stride,
+        )
+        fill_compress_tail(
+            indices=hca_i,
+            indptr=hca_p,
+            prefix_len=swa_len,
+            page_indices=hca_page_indices[:N],
+            valid_len=hca_len,
+            swa_pages=swa_pages,
+        )
+    return swa_i, swa_p, hca_i, hca_p, csa_i, csa_p
+
+
+# ---------------------------------------------------------------------------
+# Prefill index builder (ragged-packed: paged prefix + flat extend)
+# ---------------------------------------------------------------------------
+@triton.jit
+def _prefill_lengths_kernel(
+    positions_ptr,  # [T] int
+    chunk_start_ptr,  # [T] int
+    page_idx_ptr,  # [T, Wc] int (front-packed, -1 padded)
+    prefix_len_ptr,  # [T] int32 out
+    extend_len_ptr,  # [T] int32 out
+    win: tl.constexpr,
+    Wc: tl.constexpr,
+    HAS_COMPRESS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Per token: write extend/prefix segment lengths"""
+    t = tl.program_id(0)
+    pos = tl.load(positions_ptr + t).to(tl.int32)
+    cstart = tl.load(chunk_start_ptr + t).to(tl.int32)
+    tpic = pos - cstart
+    swa_low = tl.maximum(pos - win + 1, 0)
+    extend_count = tl.minimum(tpic + 1, win)
+    prefix_swa_count = tl.minimum(tl.maximum(cstart - swa_low, 0), win)
+    tl.store(extend_len_ptr + t, extend_count)
+    if HAS_COMPRESS:
+        nc = 0
+        for off in tl.range(0, Wc, BLOCK):
+            j = off + tl.arange(0, BLOCK)
+            m = j < Wc
+            j_clamped = tl.minimum(j, Wc - 1)
+            pi = tl.load(page_idx_ptr + t * Wc + j_clamped, mask=m, other=-1)
+            nc += tl.sum(tl.where(m & (pi >= 0), 1, 0))
+        tl.store(prefix_len_ptr + t, prefix_swa_count + nc)
+    else:
+        tl.store(prefix_len_ptr + t, prefix_swa_count)
+
+
+@triton.jit
+def _build_prefill_indices_kernel(
+    positions_ptr,  # [T] int
+    chunk_start_ptr,  # [T] int
+    cu_q_ptr,  # [T] int
+    state_slot_ptr,  # [T] int
+    page_idx_ptr,  # [T, Wc] int (front-packed, -1 padded)
+    pre_indptr_ptr,  # [T+1] int32 (prefix stream ragged indptr)
+    ext_indptr_ptr,  # [T+1] int32 (extend stream ragged indptr)
+    pre_out_ptr,
+    ext_out_ptr,
+    swa_pages,
+    ring_stride,  # SWA ring per-slot stride
+    win: tl.constexpr,
+    Wc: tl.constexpr,
+    HAS_COMPRESS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Per token: write extend rows + prefix (SWA ring slots ++ swa_pages+compressed slots) as two ragged segments"""
+    t = tl.program_id(0)
+    pos = tl.load(positions_ptr + t).to(tl.int32)
+    cstart = tl.load(chunk_start_ptr + t).to(tl.int32)
+    cuq = tl.load(cu_q_ptr + t).to(tl.int32)
+    s = tl.load(state_slot_ptr + t).to(tl.int32)
+
+    tpic = pos - cstart
+    swa_low = tl.maximum(pos - win + 1, 0)
+    extend_count = tl.minimum(tpic + 1, win)
+    prefix_swa_count = tl.minimum(tl.maximum(cstart - swa_low, 0), win)
+
+    ebase = tl.load(ext_indptr_ptr + t)
+    pbase = tl.load(pre_indptr_ptr + t)
+
+    # ---- extend: rows into the current-chunk kv tensor ----
+    ext_start = cuq + tpic - extend_count + 1
+    for off in tl.range(0, win, BLOCK):
+        k = off + tl.arange(0, BLOCK)
+        m = k < extend_count
+        tl.store(ext_out_ptr + ebase + k, ext_start + k, mask=m)
+
+    # ---- prefix SWA: prior-chunk ring slots (stride = ring_stride) ----
+    for off in tl.range(0, win, BLOCK):
+        k = off + tl.arange(0, BLOCK)
+        m = k < prefix_swa_count
+        gp = swa_low + k
+        tl.store(pre_out_ptr + pbase + k, s * ring_stride + (gp % ring_stride), mask=m)
+
+    # ---- prefix compressed: swa_pages + front-packed page index ----
+    if HAS_COMPRESS:
+        nc = tl.load(pre_indptr_ptr + t + 1) - pbase - prefix_swa_count
+        cbase = pbase + prefix_swa_count
+        for off in tl.range(0, Wc, BLOCK):
+            j = off + tl.arange(0, BLOCK)
+            m = j < nc
+            j_clamped = tl.minimum(j, Wc - 1)
+            pi = tl.load(page_idx_ptr + t * Wc + j_clamped, mask=m, other=0).to(
+                tl.int32
+            )
+            tl.store(pre_out_ptr + cbase + j, pi + swa_pages, mask=m)
+
+
+def build_prefill_indices(
+    *,
+    compress_ratio: int,
+    state_slot: torch.Tensor,  # [T] int (per token)
+    positions: torch.Tensor,  # [T] int (per token absolute position)
+    chunk_start: torch.Tensor,  # [T] int (absolute start of this chunk for token's seq)
+    cu_q: torch.Tensor,  # [T] int (row in extend `kv` of the seq's first chunk token)
+    win: int,  # SWA attention window length
+    ring_stride: int,  # SWA ring per-slot stride / modulo (win_with_spec)
+    swa_pages: int,
+    c128_page_indices: Optional[torch.Tensor],
+    c4_sparse_page_indices: Optional[torch.Tensor],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Build ragged prefill indices: prefix (SWA ring + swa_pages + compressed) into unified_kv + extend into current-chunk kv; returns (prefix_indices, prefix_indptr, extend_indices, extend_indptr)."""
+    device = state_slot.device
+    T = state_slot.shape[0]
+    assert positions.is_contiguous() and chunk_start.is_contiguous()
+    assert cu_q.is_contiguous() and state_slot.is_contiguous()
+
+    if compress_ratio == 0:
+        page_idx = None
+    elif compress_ratio == 128:
+        assert c128_page_indices is not None
+        page_idx = c128_page_indices[:T]
+    elif compress_ratio == 4:
+        assert c4_sparse_page_indices is not None
+        page_idx = c4_sparse_page_indices[:T]
+    else:
+        raise ValueError(f"bad compress_ratio {compress_ratio}")
+
+    has_compress = page_idx is not None
+    if has_compress:
+        assert page_idx.is_contiguous()
+    Wc = page_idx.shape[1] if has_compress else 0
+
+    block = min(1024, triton.next_power_of_2(max(win, Wc, 1)))
+    prefix_len = torch.empty(T, dtype=torch.int32, device=device)
+    extend_len = torch.empty(T, dtype=torch.int32, device=device)
+    _prefill_lengths_kernel[(T,)](
+        positions,
+        chunk_start,
+        page_idx if has_compress else positions,  # dummy ptr when no compress
+        prefix_len,
+        extend_len,
+        win=win,
+        Wc=Wc if has_compress else 1,
+        HAS_COMPRESS=has_compress,
+        BLOCK=block,
+        num_warps=4,
+    )
+    kv_indptr_prefix = _lengths_to_indptr(prefix_len)
+    kv_indptr_extend = _lengths_to_indptr(extend_len)
+
+    kv_indices_prefix = torch.empty(T * (win + Wc), dtype=torch.int32, device=device)
+    kv_indices_extend = torch.empty(T * win, dtype=torch.int32, device=device)
+
+    _build_prefill_indices_kernel[(T,)](
+        positions,
+        chunk_start,
+        cu_q,
+        state_slot,
+        page_idx if has_compress else state_slot,  # dummy ptr when no compress
+        kv_indptr_prefix,
+        kv_indptr_extend,
+        kv_indices_prefix,
+        kv_indices_extend,
+        swa_pages,
+        ring_stride,
+        win=win,
+        Wc=Wc if has_compress else 1,
+        HAS_COMPRESS=has_compress,
+        BLOCK=block,
+        num_warps=4,
+    )
+    return kv_indices_prefix, kv_indptr_prefix, kv_indices_extend, kv_indptr_extend
+
+
+def prefill(
+    *,
+    q: torch.Tensor,  # [T, H, D]
+    unified_kv: torch.Tensor,  # [pages, D]
+    kv_indices_prefix: torch.Tensor,
+    kv_indptr_prefix: torch.Tensor,
+    kv_extend: torch.Tensor,  # [T, D] current-chunk K (bf16, norm+rope'd)
+    kv_indices_extend: torch.Tensor,
+    kv_indptr_extend: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+) -> torch.Tensor:
+    return sparse_attn_v4_paged_prefill(
+        q,
+        unified_kv,
+        kv_indices_prefix,
+        kv_indptr_prefix,
+        kv_extend,
+        kv_indices_extend,
+        kv_indptr_extend,
+        attn_sink,
+        softmax_scale,
+    )

python/sglang/srt/layers/fused_qk_norm_rope_store.py

@@ -95,6 +95,7 @@ def _fused_qk_norm_rope_store_kernel(
     FP8_MAX: tl.constexpr,
     BYTES_PER_TOKEN: tl.constexpr,
     SWA_PAGE_SIZE: tl.constexpr,
+    BF16_STORE: tl.constexpr,
 ):
     pid_m = tl.program_id(0).to(tl.int64)
     pid_h = tl.program_id(1).to(tl.int64)
@@ -206,7 +207,24 @@ def _fused_qk_norm_rope_store_kernel(
     VALUE_STRIDE: tl.constexpr = DIM_NOPE + ROPE_DIM * 2
     SCALE_BYTES: tl.constexpr = NUM_NOPE_TILES + 1
 
-    if HAS_SWA_STORE:
+    if HAS_SWA_STORE and BF16_STORE:
+        # unified_kv unified_kv: write the whole head_dim as plain bf16 into a
+        # [num_slots, head_dim] bf16 cache at row=loc (no fp8 / no scale).
+        loc = tl.load(swa_loc_ptr + src_id, mask=src_mask, other=0)
+        row_base = loc.to(tl.int64)[:, None] * swa_cache_stride_page
+        # nope
+        tl.store(
+            swa_cache_ptr + row_base + offs_d_full[None, :],
+            kv_normed.to(swa_cache_ptr.dtype.element_ty),
+            mask=src_mask[:, None] & nope_d_mask[None, :],
+        )
+        # pe
+        tl.store(
+            swa_cache_ptr + row_base + (NOPE_DIM + d_pe_offs[None, :]),
+            kv_pe.to(swa_cache_ptr.dtype.element_ty),
+            mask=src_mask[:, None],
+        )
+    elif HAS_SWA_STORE:
         loc = tl.load(swa_loc_ptr + src_id, mask=src_mask, other=0)
         page_id = loc // SWA_PAGE_SIZE
         page_off = loc % SWA_PAGE_SIZE
@@ -288,15 +306,17 @@ def fused_qk_norm_rope_swa_store(
     swa_page_size: int = 128,
     q_out: Optional[torch.Tensor] = None,
     dtype: torch.dtype = torch.bfloat16,
+    bf16_store: bool = False,
 ) -> torch.Tensor:
-    """Fused Q norm + KV norm + RoPE + optional FP8 paged SWA store.
+    """Fused Q norm + KV norm + RoPE + optional SWA store.
 
     Args:
         q: [M, N] or [splitk, M, N] where N = num_local_heads * head_dim
         kv: [M, head_dim=512] mutated in-place (norm + RoPE)
-        swa_cache: paged SWA KV pool buffer [num_pages, bytes_per_page] uint8
+        swa_cache: paged SWA KV pool buffer [num_pages, bytes_per_page] uint8 OR a plain [num_slots, head_dim] bf16 cache
         swa_loc: [M] int32 pre-translated paged indices
         swa_page_size: tokens per SWA page (default 128)
+        bf16_store: write the whole head_dim as plain bf16 at swa_cache[swa_loc]
     """
     head_dim = kv.shape[1]
 
@@ -375,6 +395,7 @@ def fused_qk_norm_rope_swa_store(
         FP8_MAX=fp8_info.max,
         BYTES_PER_TOKEN=bytes_per_token,
         SWA_PAGE_SIZE=swa_page_size,
+        BF16_STORE=bf16_store,
         num_warps=num_warps,
     )
     return q_out

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -374,6 +374,65 @@ class DeepSeekV4LayerItem(NamedTuple):
     compress_kv_pool: Optional[DeepSeekV4SingleKVPool] = None
 
 
+class DeepSeekV4UnifiedKVPool:
+    """
+    Layout:
+    unified_kv[L]: ``[swa_pages + compress_pages, head_dim]`` bf16
+    - rows ``[0, swa_pages)``   = SWA ring (``req_pool_indices * swa_window + pos % swa_window``)
+    - rows ``[swa_pages, ...)`` = compressed (``swa_pages + page_index``)
+    """
+
+    K_PER_BLOCK = {0: 0, 4: 32, 128: 1}
+
+    def __init__(
+        self,
+        *,
+        stage_ratios: List[int],
+        num_slots: int,
+        num_blocks: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        device: str,
+        memory_saver_adapter,
+        custom_mem_pool,
+        swa_ring_size: int,
+    ):
+        self.swa_ring_size = swa_ring_size
+        self.head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.num_slots = num_slots
+        self.swa_pages = num_slots * self.swa_ring_size
+        self.num_blocks = num_blocks
+        self.k_per_block = dict(self.K_PER_BLOCK)
+
+        bufs = []
+        with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            with (
+                torch.cuda.use_mem_pool(custom_mem_pool)
+                if custom_mem_pool
+                else nullcontext()
+            ):
+                for ratio in stage_ratios:
+                    compress_pages = self.num_blocks * self.k_per_block[ratio]
+                    bufs.append(
+                        torch.zeros(
+                            self.swa_pages + compress_pages,
+                            self.head_dim,
+                            dtype=torch.bfloat16,
+                            device=device,
+                        )
+                    )
+        self.kv_buffer = bufs
+
+    def get_unified_kv(self, local_layer_id: int) -> torch.Tensor:
+        return self.kv_buffer[local_layer_id]
+
+    def get_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
+        data_ptrs = [b.data_ptr() for b in self.kv_buffer]
+        data_lens = [b.nbytes for b in self.kv_buffer]
+        item_lens = [b[0].nbytes for b in self.kv_buffer]
+        return data_ptrs, data_lens, item_lens
+
+
 class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
 
     def __init__(
@@ -395,6 +454,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         device: str,
         enable_memory_saver: bool,
         compression_ratios: List[int],
+        sliding_window: int = 128,
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
         enable_hisparse: bool = False,
@@ -441,6 +501,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         stage_ratios = compression_ratios[self._stage_start : self._stage_end]
 
         assert page_size % swa_page_size == 0
+        self.sliding_window = sliding_window
 
         self.swa_size = swa_size
         self.swa_window_size = swa_page_size
@@ -455,41 +516,75 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         c128_layer_num = sum(1 for r in stage_ratios if r == 128)
         c4_page_size = page_size // 4
         c128_page_size = page_size // 128
-        self.swa_kv_pool = DeepSeekV4SingleKVPool(
+
-            swa_size,
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
-            swa_page_size,
+            is_unified_kv_triton,
-            dtype,
-            qk_nope_head_dim,
-            qk_rope_head_dim,
-            layer_num,
-            device,
-            enable_memory_saver,
         )
 
-        c4_kv_pool_type = DeepSeekV4SingleKVPool
+        self._unified_kv = is_unified_kv_triton()
-        if enable_hisparse:
-            c4_kv_pool_type = HiSparseC4DevicePool
-        self.c4_kv_pool = c4_kv_pool_type(
-            c4_size,
-            c4_page_size,
-            dtype,
-            qk_nope_head_dim,
-            qk_rope_head_dim,
-            c4_layer_num,
-            device,
-            enable_memory_saver,
-        )
 
-        self.c128_kv_pool = DeepSeekV4SingleKVPool(
+        if self._unified_kv:
-            c128_size,
+            self.swa_kv_pool = None
-            c128_page_size,
+            self.c4_kv_pool = None
-            dtype,
+            self.c128_kv_pool = None
-            qk_nope_head_dim,
+            server_args = get_global_server_args()
-            qk_rope_head_dim,
+            spec_extra = (
-            c128_layer_num,
+                (server_args.speculative_num_draft_tokens - 1)
-            device,
+                if server_args.speculative_algorithm is not None
-            enable_memory_saver,
+                else 0
-        )
+            )
+            self.unified_kv_pool = DeepSeekV4UnifiedKVPool(
+                stage_ratios=stage_ratios,
+                num_slots=self.max_num_reqs + 1,
+                num_blocks=self.c128_size,
+                qk_nope_head_dim=qk_nope_head_dim,
+                qk_rope_head_dim=qk_rope_head_dim,
+                device=device,
+                memory_saver_adapter=self.memory_saver_adapter,
+                custom_mem_pool=self.custom_mem_pool,
+                swa_ring_size=self.sliding_window + spec_extra,
+            )
+
+            self.unified_swa_window = self.sliding_window
+            self.unified_swa_ring_size = self.sliding_window + spec_extra
+            self.unified_swa_pages = self.unified_kv_pool.swa_pages
+        else:
+            self.unified_kv_pool = None
+            self.swa_kv_pool = DeepSeekV4SingleKVPool(
+                swa_size,
+                swa_page_size,
+                dtype,
+                qk_nope_head_dim,
+                qk_rope_head_dim,
+                layer_num,
+                device,
+                enable_memory_saver,
+            )
+
+            c4_kv_pool_type = DeepSeekV4SingleKVPool
+            if enable_hisparse:
+                c4_kv_pool_type = HiSparseC4DevicePool
+            self.c4_kv_pool = c4_kv_pool_type(
+                c4_size,
+                c4_page_size,
+                dtype,
+                qk_nope_head_dim,
+                qk_rope_head_dim,
+                c4_layer_num,
+                device,
+                enable_memory_saver,
+            )
+
+            self.c128_kv_pool = DeepSeekV4SingleKVPool(
+                c128_size,
+                c128_page_size,
+                dtype,
+                qk_nope_head_dim,
+                qk_rope_head_dim,
+                c128_layer_num,
+                device,
+                enable_memory_saver,
+            )
 
         indexer_size = (
             self.c4_logical_size
@@ -513,6 +608,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         else:
             self._init_paged_compress_states(enable_memory_saver)
 
+    def get_unified_kv(self, layer_id: int) -> torch.Tensor:
+        return self.unified_kv_pool.get_unified_kv(layer_id - self._stage_start)
+
     def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
         self.full_to_swa_index_mapping = full_to_swa_index_mapping
 
@@ -530,11 +628,19 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         data_lens: List[int] = []
         item_lens: List[int] = []
 
-        for bufs in [
+        if self._unified_kv:
-            self.c4_kv_pool.kv_buffer,
+            buf_groups = [
-            self.c4_indexer_kv_pool.index_k_with_scale_buffer,
+                self.unified_kv_pool.kv_buffer,
-            self.c128_kv_pool.kv_buffer,
+                self.c4_indexer_kv_pool.index_k_with_scale_buffer,
-        ]:
+            ]
+        else:
+            buf_groups = [
+                self.c4_kv_pool.kv_buffer,
+                self.c4_indexer_kv_pool.index_k_with_scale_buffer,
+                self.c128_kv_pool.kv_buffer,
+            ]
+
+        for bufs in buf_groups:
             for buf in bufs:
                 assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
                 data_ptrs.append(buf.data_ptr())
@@ -548,11 +654,12 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         data_lens: List[int] = []
         item_lens: List[int] = []
 
-        for buf in self.swa_kv_pool.kv_buffer:
+        if not self._unified_kv:
-            assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
+            for buf in self.swa_kv_pool.kv_buffer:
-            data_ptrs.append(buf.data_ptr())
+                assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
-            data_lens.append(buf.nbytes)
+                data_ptrs.append(buf.data_ptr())
-            item_lens.append(buf[0].nbytes)
+                data_lens.append(buf.nbytes)
+                item_lens.append(buf[0].nbytes)
 
         for pools in [
             self.compress_state_pools,

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -408,6 +408,7 @@ class ModelRunnerKVCacheMixin:
                 c128_state_pool_size=self.c128_state_pool_size,
                 page_size=self.page_size,
                 swa_page_size=swa_page_size,
+                sliding_window=self.model_config.window_size,
                 dtype=self.kv_cache_dtype,
                 state_dtype=self.state_dtype,
                 qk_nope_head_dim=self.model_config.qk_nope_head_dim,

python/sglang/srt/models/deepseek_v4.py

@@ -785,8 +785,17 @@ class MQALayer(nn.Module):
         use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
         kv: Optional[torch.Tensor]
 
-        if self.use_fused_qk_norm_rope:
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_triton,
+        )
 
+        unified = is_unified_kv_triton()
+        is_decode = forward_batch.forward_mode.is_decode_or_idle()
+        do_fused_store = (unified and is_decode) or (
+            not unified and self.use_fused_qk_norm_rope
+        )
+
+        if do_fused_store:
             if _is_gfx95_supported:
                 q_for_wqb, q_lora = _fused_rmsnorm_fp8_quant(
                     q_lora,
@@ -804,15 +813,33 @@ class MQALayer(nn.Module):
                 else self.wkv(x_linear)[0]
             )
 
+            token_to_kv_pool = get_token_to_kv_pool()
+            if unified:
+                swa_ring_size = token_to_kv_pool.unified_swa_ring_size
+                swa_cache = token_to_kv_pool.get_unified_kv(self.layer_id)
+                # ring slot = req_slot * ring + pos % ring, per token.
+                # positions is per-token; req_pool_indices is per-req.
+                req_slot = forward_batch.req_pool_indices.to(torch.int64)
+                if req_slot.shape[0] != positions.shape[0]:
+                    req_slot = req_slot.repeat_interleave(
+                        positions.shape[0] // req_slot.shape[0]
+                    )
+                swa_loc = (
+                    req_slot * swa_ring_size + positions.to(torch.int64) % swa_ring_size
+                ).to(torch.int32)
+                swa_page_size, bf16_store = 1, True
+            else:
+                swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
+                swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
+                swa_page_size, bf16_store = (
+                    token_to_kv_pool.swa_kv_pool.page_size,
+                    False,
+                )
+
             from sglang.srt.layers.fused_qk_norm_rope_store import (
                 fused_qk_norm_rope_swa_store,
             )
 
-            token_to_kv_pool = get_token_to_kv_pool()
-            swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
-            swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
-            swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
-
             q = fused_qk_norm_rope_swa_store(
                 q=q,
                 kv=kv,
@@ -829,9 +856,11 @@ class MQALayer(nn.Module):
                 swa_page_size=swa_page_size,
                 q_out=q_out,
                 dtype=x.dtype,
+                bf16_store=bf16_store,
             )
+            kv = None
 
-            if use_cp:
+            if not unified and use_cp:
                 # DSA CP: keep bf16 kv around for the cross-rank all-gather, then
                 # write to the FlashMLA cache after gather.
                 kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a)
@@ -844,7 +873,11 @@ class MQALayer(nn.Module):
         else:
             q_lora = self.q_norm(q_lora)
             q = self._compute_q_b(q_lora, positions, q_out)
-            if use_cp:
+            if unified:
+                # unified_kv prefill: keep bf16 kv; the backend writes
+                # the ring AFTER attention (2-source path).
+                kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
+            elif use_cp:
                 # NSA CP: keep bf16 kv around for the cross-rank all-gather, then
                 # write to the FlashMLA cache after gather.
                 kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
@@ -956,34 +989,50 @@ class MQALayer(nn.Module):
         # tell the backend to skip its own store_cache. When `kv is None`
         # (no DSA-CP), pass `q` as a sentinel for the `k is v` assert; the
         # attention path doesn't read it once `save_kv_cache=False`.
-        attn_q = q_padded if q_padded is not None else q
         attn_k = kv if kv is not None else q
-        save_kv_cache = False
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
-        if forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph():
+            is_unified_kv_triton,
-            o = attn_q.new_empty(
+        )
-                (*attn_q.shape[:-1], self.attn_mqa.v_head_dim),
+
-            )
+        if is_unified_kv_triton():
-            bcg_deepseek_v4_attention_with_output(
-                attn_q,
-                attn_k,
-                o,
-                self.attn_mqa.layer_id,
-                self.compress_ratio,
-                self.attn_sink,
-                save_kv_cache,
-            )
-        else:
             o = attn_backend.forward(
-                q=attn_q,
+                q=q_out if q_out is not None else q,
                 k=attn_k,
                 v=attn_k,
                 layer=self.attn_mqa,
                 forward_batch=forward_batch,
                 compress_ratio=self.compress_ratio,
                 attn_sink=self.attn_sink,
-                save_kv_cache=save_kv_cache,
+                save_kv_cache=kv is not None,
             )
-        o = o[:, tp_slice, :]
+        else:
+            attn_q = q_padded if q_padded is not None else q
+            save_kv_cache = False
+            if forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph():
+                o = attn_q.new_empty(
+                    (*attn_q.shape[:-1], self.attn_mqa.v_head_dim),
+                )
+                bcg_deepseek_v4_attention_with_output(
+                    attn_q,
+                    attn_k,
+                    o,
+                    self.attn_mqa.layer_id,
+                    self.compress_ratio,
+                    self.attn_sink,
+                    save_kv_cache,
+                )
+            else:
+                o = attn_backend.forward(
+                    q=attn_q,
+                    k=attn_k,
+                    v=attn_k,
+                    layer=self.attn_mqa,
+                    forward_batch=forward_batch,
+                    compress_ratio=self.compress_ratio,
+                    attn_sink=self.attn_sink,
+                    save_kv_cache=save_kv_cache,
+                )
+            o = o[:, tp_slice, :]
         fused_rope_inplace(
             o[..., -self.qk_rope_head_dim :],
             None,