[AMD] Support DeepSeek V4 DSpark on AMD HIP platform (#30964)

python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py

@@ -110,6 +110,57 @@ def store_swa_into_unified(
     )
 
 
+@triton.jit
+def _scatter_loc_kernel(
+    kv_ptr,  # [T, D] bf16
+    loc_ptr,  # [T] int (unified row index; <0 => skip)
+    unified_ptr,  # [pages, D] bf16
+    n_rows,
+    D: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    row = tl.program_id(0)
+    if row >= n_rows:
+        return
+    loc = tl.load(loc_ptr + row).to(tl.int64)
+    if loc < 0:
+        return
+    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 scatter_bf16_into_unified(
+    *,
+    kv: torch.Tensor,  # [T, head_dim] bf16 (already norm+rope'd)
+    loc: torch.Tensor,  # [T] int32/int64 unified ring row; <0 => skip
+    unified_kv: torch.Tensor,  # [pages, head_dim] bf16
+) -> None:
+    """Scatter already-norm+rope'd bf16 K into ``unified_kv[loc]`` (skip loc < 0).
+
+    Companion to ``store_swa_into_unified`` for callers that already hold the
+    precomputed ring row index (the DSpark draft: ``get_unified_swa_loc`` for the
+    draft forward, or the commit-inject layout for target-hidden injection) and
+    need per-row commit masking expressed as ``loc == -1``.
+    """
+    n_rows, D = kv.shape
+    if n_rows == 0:
+        return
+    assert kv.is_contiguous() and kv.dtype == unified_kv.dtype
+    assert loc.is_contiguous()
+    assert unified_kv.is_contiguous()
+    _scatter_loc_kernel[(n_rows,)](
+        kv,
+        loc,
+        unified_kv,
+        n_rows,
+        D=D,
+        BLOCK_D=triton.next_power_of_2(D),
+        num_warps=8,
+    )
+
+
 # ---------------------------------------------------------------------------
 # Ragged indptr helper (shared by the decode streams + prefill builders)
 # ---------------------------------------------------------------------------

python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py

@@ -604,6 +604,37 @@ class CommitInjectLayoutResult(msgspec.Struct):
     positions: torch.Tensor
 
 
+def build_unified_commit_inject_layout(
+    *,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    block_pos_offsets: torch.Tensor,
+    commit_lens: torch.Tensor,
+    stride: int,
+    ring_stride: int,
+) -> CommitInjectLayoutResult:
+    """unified_kv counterpart of build_commit_inject_layout.
+
+    Non-unified injection translates the verify tokens' full cache locs through
+    ``full_to_swa_mapping``; under unified_kv the SWA K lives in a ring addressed
+    directly by ``state_slot * ring_stride + pos % ring_stride``, so compute the
+    ring row here instead. Uncommitted tokens (col >= commit_len) get loc = -1 and
+    are skipped by the scatter. All ops are static-shape (CUDA-graph safe).
+    """
+    bs = req_pool_indices.shape[0]
+    device = req_pool_indices.device
+    positions_2d = prefix_lens.unsqueeze(1) + block_pos_offsets[:stride]
+    positions = positions_2d.reshape(-1).to(torch.int64)
+    state_slot = (
+        req_pool_indices.to(torch.int64).view(-1, 1).expand(bs, stride).reshape(-1)
+    )
+    loc = state_slot * ring_stride + positions % ring_stride
+    col = torch.arange(stride, device=device).view(1, -1)
+    committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1)
+    swa_loc = torch.where(committed, loc, torch.full_like(loc, -1)).to(torch.int32)
+    return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions)
+
+
 class BuildCommitInjectLayout:
     @classmethod
     def execute(cls, *args, **kwargs) -> CommitInjectLayoutResult:

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

@@ -414,6 +414,7 @@ class DeepseekV4HipRadixBackend(
     # both children and leaks ROCm HSA resources (HSA_STATUS_ERROR_OUT_OF_RESOURCES).
     # TboAttnBackend reads this to skip children in the *_graph paths only.
     tbo_supports_cuda_graph = False
+    supports_ragged_verify_graph: bool = True
 
     def __init__(
         self,
@@ -456,6 +457,18 @@ class DeepseekV4HipRadixBackend(
         self.mtp_enabled = self.topk > 0
         self.speculative_num_steps = speculative_num_steps
         self.speculative_num_draft_tokens: int = get_spec().speculative_num_draft_tokens
+        self.is_dspark_draft = (
+            getattr(model_runner, "is_draft_worker", False)
+            and model_runner.spec_algorithm.is_dspark()
+        )
+        self.target_verify_num_draft_tokens = self.speculative_num_draft_tokens
+        if self.is_dspark_draft:
+            assert self.speculative_num_draft_tokens is not None
+            assert self.speculative_num_draft_tokens > 1
+            # DSpark draft workers verify gamma rows. The server arg keeps the
+            # CUDA-side convention gamma + 1, so use an explicit effective value
+            # instead of mutating speculative_num_draft_tokens in place.
+            self.target_verify_num_draft_tokens = self.speculative_num_draft_tokens - 1
         self.speculative_step_id = speculative_step_id
         self.forward_metadata: Union[
             DSV4Metadata,
@@ -532,14 +545,34 @@ class DeepseekV4HipRadixBackend(
         extend_seq_lens_cpu: List[int],
         need_compress: bool = True,
         use_prefill_cuda_graph: bool = False,
+        compress_gpu_plan: bool = False,
+        extend_start_loc: Optional[torch.Tensor] = None,
     ) -> DSV4Metadata:
-        seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually(
-            num_tokens=num_tokens,
-            seq_lens=seq_lens_cpu,
-            extend_seq_lens=extend_seq_lens_cpu,
-            req_pool_indices=req_pool_indices,
-            padded_num_tokens=out_cache_loc.shape[0],
-        )
+        if extend_start_loc is not None:
+            from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import (
+                ExpandPrefillCausally,
+            )
+
+            _expanded = ExpandPrefillCausally.execute(
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens,
+                extend_seq_lens=extend_seq_lens,
+                extend_start_loc=extend_start_loc,
+                seq_lens_cpu=None,
+                extend_seq_lens_cpu=None,
+                num_tokens=num_tokens,
+                padded_num_tokens=out_cache_loc.shape[0],
+            )
+            seq_lens_casual = _expanded.seq_lens_casual
+            req_pool_indices_repeated = _expanded.req_pool_indices_repeated
+        else:
+            seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually(
+                num_tokens=num_tokens,
+                seq_lens=seq_lens_cpu,
+                extend_seq_lens=extend_seq_lens_cpu,
+                req_pool_indices=req_pool_indices,
+                padded_num_tokens=out_cache_loc.shape[0],
+            )
         core_attn_metadata = self.make_core_attn_metadata(
             req_to_token=self.req_to_token,
             req_pool_indices_repeated=req_pool_indices_repeated,
@@ -559,6 +592,20 @@ class DeepseekV4HipRadixBackend(
         )
         if not need_compress:
             create = _create_dummy_paged_compress_data
+        elif compress_gpu_plan:
+            create = functools.partial(
+                create_paged_compressor_data,
+                is_prefill=True,
+                token_to_kv_pool=self.token_to_kv_pool,
+                req_to_token=self.req_to_token,
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens,
+                seq_lens_cpu=None,
+                extend_lens=extend_seq_lens,
+                extend_lens_cpu=None,
+                num_q_tokens=num_tokens,
+                use_prefill_cuda_graph=use_prefill_cuda_graph,
+            )
         else:
             create = functools.partial(
                 create_paged_compressor_data,
@@ -588,6 +635,7 @@ class DeepseekV4HipRadixBackend(
         extend_seq_lens: Optional[torch.Tensor] = None,
         use_prefill_cuda_graph: bool = False,
         seq_lens_cpu: Optional[List[int]] = None,
+        ragged_layout=None,
     ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]:
         # HIP path: build target-verify metadata eagerly even when
         # SGLANG_PREP_IN_CUDA_GRAPH is enabled. The raw/lazy-upgrade route can
@@ -601,6 +649,7 @@ class DeepseekV4HipRadixBackend(
             seq_lens_cpu=seq_lens_cpu,
             out_cache_loc=out_cache_loc,
             use_prefill_cuda_graph=use_prefill_cuda_graph,
+            ragged_layout=ragged_layout,
         )
 
     def init_forward_metadata_target_verify_old(
@@ -611,13 +660,38 @@ class DeepseekV4HipRadixBackend(
         seq_lens_cpu: Optional[List[int]] = None,
         out_cache_loc: Optional[torch.Tensor] = None,
         use_prefill_cuda_graph: bool = False,
+        ragged_layout=None,
     ) -> DSV4Metadata:
         batch_size = len(seq_lens)
-        seq_lens = seq_lens + self.speculative_num_draft_tokens
-        seq_lens_cpu = [x + self.speculative_num_draft_tokens for x in seq_lens_cpu]
-        extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size
-        extend_seq_lens = self._move_to_device(extend_seq_lens_cpu)
-        num_tokens = self.speculative_num_draft_tokens * batch_size
+        extend_start_loc = None
+        if ragged_layout is not None:
+            verify_lens_dev = ragged_layout.verify_lens.to(
+                device=seq_lens.device, dtype=torch.int32
+            )
+            extend_start_loc = ragged_layout.extend_start_loc.to(
+                device=seq_lens.device, dtype=torch.int32
+            )
+            extend_seq_lens = verify_lens_dev
+            seq_lens = seq_lens + verify_lens_dev.to(seq_lens.dtype)
+            # Total verify tokens to expand. For the graph path the padded layout
+            # sets total_verify_tokens == graph_num_tokens (tier); the eager path
+            # resolves a device-assembled layout whose total_verify_tokens is None,
+            # so fall back to sum(verify_lens) (== real total; padded == tier).
+            num_tokens = ragged_layout.total_verify_tokens
+            if num_tokens is None:
+                num_tokens = int(verify_lens_dev.sum().item())
+            else:
+                num_tokens = int(num_tokens)
+            extend_seq_lens_cpu = None
+            seq_lens_cpu = None
+        else:
+            seq_lens = seq_lens + self.target_verify_num_draft_tokens
+            seq_lens_cpu = [
+                x + self.target_verify_num_draft_tokens for x in seq_lens_cpu
+            ]
+            extend_seq_lens_cpu = [self.target_verify_num_draft_tokens] * batch_size
+            num_tokens = self.target_verify_num_draft_tokens * batch_size
+            extend_seq_lens = self._move_to_device(extend_seq_lens_cpu)
         if out_cache_loc is None:
             out_cache_loc = seq_lens.new_zeros(num_tokens)
         return self.init_forward_metadata_prefill(
@@ -631,6 +705,8 @@ class DeepseekV4HipRadixBackend(
             extend_seq_lens_cpu=extend_seq_lens_cpu,
             need_compress=True,
             use_prefill_cuda_graph=use_prefill_cuda_graph,
+            compress_gpu_plan=ragged_layout is not None,
+            extend_start_loc=extend_start_loc,
         )
 
     def make_forward_metadata_from_raw_verify(
@@ -640,7 +716,7 @@ class DeepseekV4HipRadixBackend(
         seq_lens = raw_metadata.seq_lens
         out_cache_loc = raw_metadata.out_cache_loc
 
-        bs, num_draft_tokens = len(seq_lens), self.speculative_num_draft_tokens
+        bs, num_draft_tokens = len(seq_lens), self.target_verify_num_draft_tokens
         seq_lens = seq_lens + num_draft_tokens
         extend_seq_lens = raw_metadata.extend_seq_lens
         if extend_seq_lens is None or extend_seq_lens.numel() != bs:
@@ -846,6 +922,8 @@ class DeepseekV4HipRadixBackend(
         chosen_max_seq_len = self.MAX_SEQ_LEN_FOR_CAPTURE
         assert actual_max_seq_len <= chosen_max_seq_len
 
+        graph_key = bs
+
         if bucket == _GraphBucket.DECODE_OR_IDLE:
             assert out_cache_loc is not None
             assert len(out_cache_loc.shape) == 1, f"{out_cache_loc.shape=}"
@@ -862,14 +940,14 @@ class DeepseekV4HipRadixBackend(
                 out_cache_loc=out_cache_loc_padded,
             )
         elif bucket == _GraphBucket.TARGET_VERIFY:
-            if resolve_ragged_verify_layout(forward_batch) is not None:
-                raise NotImplementedError(
-                    "DSV4 ragged verify is not supported on the HIP backend "
-                    "(DeepseekV4HipRadixBackend) cuda-graph path; disable "
-                    "SGLANG_RAGGED_VERIFY_MODE or use a CUDA device."
-                )
             assert out_cache_loc is not None
-            num_tokens_v = self.speculative_num_draft_tokens * bs
+            ragged_layout = resolve_ragged_verify_layout(forward_batch)
+            if ragged_layout is not None:
+                ragged_layout = ragged_layout.padded_to_bucket(padded_bs=bs)
+                num_tokens_v = ragged_layout.graph_num_tokens
+                graph_key = num_tokens_v
+            else:
+                num_tokens_v = self.target_verify_num_draft_tokens * bs
             out_cache_loc_padded = torch.nn.functional.pad(
                 out_cache_loc,
                 pad=(0, num_tokens_v - len(out_cache_loc)),
@@ -885,6 +963,7 @@ class DeepseekV4HipRadixBackend(
                 # CPU mirror already available here (== seq_lens, no D2H);
                 # pass it so target_verify skips the per-iter seq_lens.tolist() sync.
                 seq_lens_cpu=seq_lens_cpu.tolist(),
+                ragged_layout=ragged_layout,
             )
         elif bucket == _GraphBucket.DRAFT_EXTEND:
             num_tokens_per_req = self.draft_extend_num_tokens_per_req
@@ -910,7 +989,7 @@ class DeepseekV4HipRadixBackend(
             raise NotImplementedError
 
         self.replay_cuda_graph_metadata_from(
-            bs=bs, temp_metadata=temp_metadata, bucket=bucket
+            bs=graph_key, temp_metadata=temp_metadata, bucket=bucket
         )
 
         if in_capture:
@@ -955,12 +1034,7 @@ class DeepseekV4HipRadixBackend(
                 out_cache_loc=out_cache_loc,
             )
         elif forward_batch.forward_mode.is_target_verify():
-            if resolve_ragged_verify_layout(forward_batch) is not None:
-                raise NotImplementedError(
-                    "DSV4 ragged verify is not supported on the HIP backend "
-                    "(DeepseekV4HipRadixBackend); disable SGLANG_RAGGED_VERIFY_MODE "
-                    "or use a CUDA device."
-                )
+            ragged_layout = resolve_ragged_verify_layout(forward_batch)
             metadata = self.init_forward_metadata_target_verify(
                 max_seq_len=max_seq_len,
                 req_pool_indices=req_pool_indices,
@@ -970,6 +1044,7 @@ class DeepseekV4HipRadixBackend(
                 seq_lens_cpu=(
                     seq_lens_cpu.tolist() if seq_lens_cpu is not None else None
                 ),
+                ragged_layout=ragged_layout,
             )
         elif forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True):
             extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -1190,6 +1190,34 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
             page_size=self.swa_kv_pool.page_size,
         )
 
+    def set_unified_key_buffer_radix_fused_norm_rope(
+        self,
+        layer_id: int,
+        swa_loc: torch.Tensor,
+        kv: torch.Tensor,
+        kv_weight: torch.Tensor,
+        eps: float,
+        freqs_cis: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> None:
+        """unified_kv counterpart of set_swa_key_buffer_radix_fused_norm_rope.
+
+        Under unified_kv the (fp8, paged) swa_kv_pool is None -- SWA K lives in
+        the shared bf16 unified_kv ring instead. Norm+RoPE the draft KV in place
+        (the same freqs_cis path the main model uses via _compute_kv_bf16) and
+        scatter it into ``unified_kv[swa_loc]``. Rows with swa_loc < 0
+        (uncommitted verify tokens) are skipped by the scatter.
+        """
+        from sglang.kernels.ops.attention.dsv4 import fused_norm_rope_inplace
+        from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
+
+        fused_norm_rope_inplace(kv, kv_weight, eps, freqs_cis, positions)
+        runtime.scatter_bf16_into_unified(
+            kv=kv,
+            loc=swa_loc,
+            unified_kv=self.get_unified_kv(layer_id),
+        )
+
     def set_extra_key_buffer_fused(
         self,
         layer_id: int,

python/sglang/srt/models/deepseek_v4_dspark.py

@@ -9,6 +9,9 @@ import torch.nn.functional as F
 from torch import nn
 
 from sglang.kernels.ops.attention.dsv4 import fused_q_norm_rope, fused_rope_inplace
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+    is_unified_kv_triton,
+)
 from sglang.kernels.ops.speculative.dspark.dspark_draft_model import (
     BuildStepLocal,
     CommitKvProj,
@@ -135,6 +138,20 @@ class DSparkAttention(MqaAttentionBase):
         attn_backend,
         pool: DeepSeekV4TokenToKVPool,
     ) -> None:
+        if is_unified_kv_triton():
+            # unified_kv: SWA K lives in the shared bf16 ring (swa_kv_pool is
+            # None). Use the unified ring write target -- get_unified_swa_loc
+            # recomputes it from live positions for multi-step draft decode.
+            pool.set_unified_key_buffer_radix_fused_norm_rope(
+                layer_id=self.layer_id,
+                swa_loc=attn_backend.get_unified_swa_loc(forward_batch),
+                kv=kv,
+                kv_weight=self.kv_norm.weight.data,
+                eps=self.eps,
+                freqs_cis=self.freqs_cis,
+                positions=positions,
+            )
+            return
         pool.set_swa_key_buffer_radix_fused_norm_rope(
             layer_id=self.layer_id,
             swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
@@ -660,9 +677,18 @@ class DeepseekV4ForCausalLMDSpark(nn.Module):
             main_x=main_x,
             wkv_linears=[stage.self_attn.wkv for stage in self.stages],
         )
+        # Under unified_kv the swa_kv_pool is None; the caller passes a unified
+        # ring loc (state_slot * ring + pos % ring, -1 for uncommitted) so the
+        # store just needs to target the bf16 ring instead of the fp8 flashmla
+        # buffer. Same swa_loc/positions contract either way.
+        store_kv = (
+            pool.set_unified_key_buffer_radix_fused_norm_rope
+            if is_unified_kv_triton()
+            else pool.set_swa_key_buffer_radix_fused_norm_rope
+        )
         for stage, kv in zip(self.stages, kvs):
             attn = stage.self_attn
-            pool.set_swa_key_buffer_radix_fused_norm_rope(
+            store_kv(
                 layer_id=attn.layer_id,
                 swa_loc=swa_loc,
                 kv=kv,

python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py

@@ -2,9 +2,13 @@ from typing import Optional
 
 import torch
 
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+    is_unified_kv_triton,
+)
 from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func
 from sglang.kernels.ops.speculative.dspark.dspark_verify_window import (
     BuildCommitInjectLayout,
+    build_unified_commit_inject_layout,
 )
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
@@ -36,6 +40,8 @@ class TargetHiddenKvInjector:
         positions: torch.Tensor,
         cache_loc_2d: Optional[torch.Tensor] = None,
         commit_lens: Optional[torch.Tensor] = None,
+        state_slot: Optional[torch.Tensor] = None,
+        final_pos: Optional[torch.Tensor] = None,
     ) -> None:
         if target_hidden is None or target_hidden.numel() == 0:
             return
@@ -54,6 +60,14 @@ class TargetHiddenKvInjector:
             commit_lens = commit_lens.to(
                 device=device, dtype=torch.int32, non_blocking=True
             )
+        if state_slot is not None:
+            state_slot = state_slot.to(
+                device=device, dtype=torch.int64, non_blocking=True
+            )
+        if final_pos is not None:
+            final_pos = final_pos.to(
+                device=device, dtype=torch.int64, non_blocking=True
+            )
 
         pool = self.draft_model_runner.token_to_kv_pool
         if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"):
@@ -64,6 +78,8 @@ class TargetHiddenKvInjector:
                 positions=positions,
                 cache_loc_2d=cache_loc_2d,
                 commit_lens=commit_lens,
+                state_slot=state_slot,
+                final_pos=final_pos,
             )
             return
 
@@ -86,13 +102,29 @@ class TargetHiddenKvInjector:
         positions: torch.Tensor,
         cache_loc_2d: Optional[torch.Tensor],
         commit_lens: Optional[torch.Tensor],
+        state_slot: Optional[torch.Tensor] = None,
+        final_pos: Optional[torch.Tensor] = None,
     ) -> None:
-        swa_loc = pool.translate_loc_from_full_to_swa(cache_loc).to(torch.int32)
-        if commit_lens is not None and cache_loc_2d is not None:
-            bs, verify_len = cache_loc_2d.shape
-            col = torch.arange(verify_len, device=cache_loc.device).view(1, -1)
-            committed_mask = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1)
-            swa_loc = torch.where(committed_mask, swa_loc, torch.full_like(swa_loc, -1))
+        if is_unified_kv_triton():
+            swa_loc = self._unified_inject_loc(
+                pool=pool,
+                positions=positions,
+                cache_loc_2d=cache_loc_2d,
+                commit_lens=commit_lens,
+                state_slot=state_slot,
+                final_pos=final_pos,
+            )
+        else:
+            swa_loc = pool.translate_loc_from_full_to_swa(cache_loc).to(torch.int32)
+            if commit_lens is not None and cache_loc_2d is not None:
+                bs, verify_len = cache_loc_2d.shape
+                col = torch.arange(verify_len, device=cache_loc.device).view(1, -1)
+                committed_mask = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(
+                    -1
+                )
+                swa_loc = torch.where(
+                    committed_mask, swa_loc, torch.full_like(swa_loc, -1)
+                )
 
         with torch.inference_mode():
             self.draft_model.write_target_hidden_kv(
@@ -102,6 +134,43 @@ class TargetHiddenKvInjector:
                 pool=pool,
             )
 
+    def _unified_inject_loc(
+        self,
+        *,
+        pool,
+        positions: torch.Tensor,
+        cache_loc_2d: Optional[torch.Tensor],
+        commit_lens: Optional[torch.Tensor],
+        state_slot: Optional[torch.Tensor],
+        final_pos: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Ring row for target-hidden injection under unified_kv.
+
+        loc = state_slot * ring + pos % ring, with two skip (-1) rules:
+          * SWA window: only the last ``win`` tokens per req land in the ring;
+            older tokens share a ring slot (pos % ring) and would race, so drop
+            them (needed for long prefill chunks).
+          * commit gate: uncommitted verify tokens (col >= commit_len) are dropped.
+        """
+        if state_slot is None:
+            raise RuntimeError(
+                "unified_kv target-hidden injection requires state_slot "
+                "(per-token draft req_pool_indices)."
+            )
+        ring = pool.unified_swa_ring_size
+        win = pool.unified_swa_window
+        pos = positions.to(torch.int64)
+        loc = state_slot.to(torch.int64) * ring + pos % ring
+        if final_pos is not None:
+            keep = pos > (final_pos.to(torch.int64) - win)
+            loc = torch.where(keep, loc, torch.full_like(loc, -1))
+        if commit_lens is not None and cache_loc_2d is not None:
+            bs, verify_len = cache_loc_2d.shape
+            col = torch.arange(verify_len, device=positions.device).view(1, -1)
+            committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1)
+            loc = torch.where(committed, loc, torch.full_like(loc, -1))
+        return loc.to(torch.int32)
+
     def inject_ragged(
         self,
         *,
@@ -119,15 +188,25 @@ class TargetHiddenKvInjector:
         if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"):
             if hidden_strided.numel() == 0:
                 return
-            inject_layout = BuildCommitInjectLayout.execute(
-                req_pool_indices=batch.req_pool_indices,
-                req_to_token=self.model_runner.req_to_token_pool.req_to_token,
-                prefix_lens=prefix_lens,
-                block_pos_offsets=self._block_pos_offsets[:stride],
-                full_to_swa_mapping=pool.full_to_swa_index_mapping,
-                commit_lens=commit_lens,
-                stride=stride,
-            )
+            if is_unified_kv_triton():
+                inject_layout = build_unified_commit_inject_layout(
+                    req_pool_indices=batch.req_pool_indices,
+                    prefix_lens=prefix_lens,
+                    block_pos_offsets=self._block_pos_offsets[:stride],
+                    commit_lens=commit_lens,
+                    stride=stride,
+                    ring_stride=pool.unified_swa_ring_size,
+                )
+            else:
+                inject_layout = BuildCommitInjectLayout.execute(
+                    req_pool_indices=batch.req_pool_indices,
+                    req_to_token=self.model_runner.req_to_token_pool.req_to_token,
+                    prefix_lens=prefix_lens,
+                    block_pos_offsets=self._block_pos_offsets[:stride],
+                    full_to_swa_mapping=pool.full_to_swa_index_mapping,
+                    commit_lens=commit_lens,
+                    stride=stride,
+                )
             with torch.inference_mode():
                 self.draft_model.write_target_hidden_kv(
                     main_hidden=hidden.reshape(-1, hidden.shape[-1]),

python/sglang/srt/speculative/dspark_components/dspark_verify.py

@@ -5,6 +5,9 @@ from typing import Optional
 import msgspec
 import torch
 
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+    is_unified_kv_triton,
+)
 from sglang.kernels.ops.speculative.dspark.dspark_accept import (
     AcceptGreedy,
     AcceptSampling,
@@ -20,6 +23,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_verify_window import (
     BuildRaggedVerifyWindow,
     RaggedVerifyWindow,
     ScatterCompactToStrided,
+    build_unified_commit_inject_layout,
     scatter_compact_to_strided_into,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
@@ -317,12 +321,23 @@ class TargetVerifyExecutor:
         if hidden is None:
             raise RuntimeError("DSpark verify requires target hidden states, got None.")
         hidden = hidden.view(bs, self.verify_num_draft_tokens, -1)
+        state_slot = None
+        if is_unified_kv_triton():
+            # unified_kv needs the per-token draft req slot to address the SWA ring
+            # (state_slot * ring + pos % ring). Verify tokens are the latest in each
+            # req so they always fall in the window; the commit gate (via commit_lens
+            # + cache_loc_2d) drops rejected tokens, so no final_pos skip is needed.
+            vlen = verify_window.verify_cache_loc_2d.shape[1]
+            state_slot = (
+                batch.req_pool_indices[:bs].view(-1, 1).expand(bs, vlen).reshape(-1)
+            )
         self.kv_injector.inject_target_hidden(
             target_hidden=hidden.reshape(-1, hidden.shape[-1]),
             cache_loc=verify_window.verify_cache_loc,
             cache_loc_2d=verify_window.verify_cache_loc_2d,
             positions=verify_window.positions_2d.reshape(-1),
             commit_lens=commit_lens,
+            state_slot=state_slot,
         )
 
     def _run_ragged(
@@ -645,15 +660,25 @@ class DsparkVerifyEpilogue:
             torch.minimum(commit_lens, verify_lens.to(torch.int32))
             * self.inject_gate_buf
         )
-        inject_layout = BuildCommitInjectLayout.execute(
-            req_pool_indices=req_pool_indices,
-            req_to_token=ctx.resolve_req_to_token(),
-            prefix_lens=seq_lens[:bs],
-            block_pos_offsets=ctx.block_pos_offsets[: self.stride],
-            full_to_swa_mapping=pool.full_to_swa_index_mapping,
-            commit_lens=gated_commit_lens,
-            stride=self.stride,
-        )
+        if is_unified_kv_triton():
+            inject_layout = build_unified_commit_inject_layout(
+                req_pool_indices=req_pool_indices,
+                prefix_lens=seq_lens[:bs],
+                block_pos_offsets=ctx.block_pos_offsets[: self.stride],
+                commit_lens=gated_commit_lens,
+                stride=self.stride,
+                ring_stride=pool.unified_swa_ring_size,
+            )
+        else:
+            inject_layout = BuildCommitInjectLayout.execute(
+                req_pool_indices=req_pool_indices,
+                req_to_token=ctx.resolve_req_to_token(),
+                prefix_lens=seq_lens[:bs],
+                block_pos_offsets=ctx.block_pos_offsets[: self.stride],
+                full_to_swa_mapping=pool.full_to_swa_index_mapping,
+                commit_lens=gated_commit_lens,
+                stride=self.stride,
+            )
         with torch.inference_mode():
             ctx.draft_model.write_target_hidden_kv(
                 main_hidden=self.strided_hidden[: bs * self.stride],

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -5,6 +5,9 @@ from typing import Optional
 
 import torch
 
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+    is_unified_kv_triton,
+)
 from sglang.srt.configs.hybrid_arch import mambaish_config
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
@@ -449,10 +452,25 @@ class DSparkWorkerV2(BaseSpecWorker):
             ctx_lens,
             int(sum(batch.extend_lens)),
         )
+        # unified_kv injects into the SWA ring keyed by (draft req slot, position);
+        # thread the per-token state_slot + the req's final position so the
+        # injector keeps only the last SWA window (older prefill tokens share a
+        # ring slot and would race). Cheap; only consumed under unified_kv.
+        state_slot = final_pos = None
+        if is_unified_kv_triton():
+            repeats = ctx_lens.to(torch.int64)
+            state_slot = torch.repeat_interleave(
+                batch.req_pool_indices.to(device=device, dtype=torch.int64), repeats
+            )
+            final_pos = torch.repeat_interleave(
+                (draft_seq_lens + ctx_lens - 1).to(torch.int64), repeats
+            )
         self._kv_injector.inject_target_hidden(
             target_hidden=logits_output.hidden_states,
             cache_loc=batch.out_cache_loc,
             positions=positions,
+            state_slot=state_slot,
+            final_pos=final_pos,
         )
         # Avoid copying large hidden-state buffers to CPU in overlap scheduling.
         logits_output.hidden_states = None