diff --git a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py index 039af217d5a6..79ae94ab03dd 100644 --- a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py +++ b/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) # --------------------------------------------------------------------------- diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py index f6c6aa6cac60..af6ffc86d051 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index a6adfc929cee..f73b79aaacf9 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/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 @@ def __init__( 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 @@ def init_forward_metadata_prefill( 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 @@ def init_forward_metadata_prefill( ) 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 @@ def init_forward_metadata_target_verify( 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 @@ def init_forward_metadata_target_verify( 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 @@ def init_forward_metadata_target_verify_old( 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 @@ def init_forward_metadata_target_verify_old( 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 @@ def make_forward_metadata_from_raw_verify( 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 @@ def init_forward_metadata_out_graph( 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 @@ def init_forward_metadata_out_graph( 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 @@ def init_forward_metadata_out_graph( # 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 @@ def init_forward_metadata_out_graph( 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 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: 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 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: 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 diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 062e60bc9c05..f8ac0aadc833 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -1192,6 +1192,34 @@ def set_swa_key_buffer_radix_fused_norm_rope( 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, diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 2b01ec2e2a08..1af8df913230 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -9,6 +9,9 @@ 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 @@ def _store_block_kv( 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 @@ def write_target_hidden_kv( 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, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py index d1e44145a360..baaf55eb7e5f 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -2,9 +2,13 @@ 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 @@ def inject_target_hidden( 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 @@ def inject_target_hidden( 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 @@ def inject_target_hidden( 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 @@ def _inject_mla( 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 @@ def _inject_mla( 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 @@ def inject_ragged( 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]), diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 4d5478cd53c7..e95beca3c97a 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -5,6 +5,9 @@ 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 @@ BuildRaggedVerifyWindow, RaggedVerifyWindow, ScatterCompactToStrided, + build_unified_commit_inject_layout, scatter_compact_to_strided_into, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput @@ -309,12 +313,23 @@ def commit_hidden( 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( @@ -637,15 +652,25 @@ def _commit_inject( 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], diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index ed64e21f37c8..fa74bd1510fe 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -5,6 +5,9 @@ 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 @@ def _forward_prefill( 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 diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py new file mode 100644 index 000000000000..944f7493ed0d --- /dev/null +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py @@ -0,0 +1,201 @@ +"""MI35x DeepSeek-V4-Pro-DSpark unified_kv GSM8K accuracy test (8-GPU). + +Runs the production AMD DSpark static configuration with the HIP dsv4 backend and +SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton. The test uses the full GSM8K set +to catch regressions in unified-KV target-hidden injection, verify metadata, and +DSpark acceptance. + +Registry: nightly-amd-8-gpu-mi35x-deepseek-v4-pro-dspark suite +""" + +import os +import unittest +from types import SimpleNamespace + +import requests +import torch + +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime +from sglang.kernels.ops.speculative.dspark import dspark_verify_window +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + +register_amd_ci( + est_time=7200, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro-dspark", nightly=True +) + +DEEPSEEK_V4_DSPARK_MODEL_PATH = os.environ.get( + "DEEPSEEK_V4_DSPARK_MODEL_PATH", "deepseek-ai/DeepSeek-V4-Pro-DSpark" +) +SERVER_LAUNCH_TIMEOUT = 5400 +FLASHMLA_BACKEND = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "unified_kv_triton") +GSM8K_ACCURACY_THRESHOLD = 0.92 +AVG_SPEC_ACCEPT_LENGTH_THRESHOLD = 3.0 +DEVICE = torch.device("cuda") + +COMMON_ENV_VARS = { + "SGLANG_DEFAULT_THINKING": "1", + "SGLANG_DSV4_REASONING_EFFORT": "max", + "SGLANG_USE_ROCM700A": "0", + "SGLANG_HACK_FLASHMLA_BACKEND": FLASHMLA_BACKEND, + "AITER_BF16_FP8_MOE_BOUND": "0", +} + +DSPARK_ENV_VARS = { + "SGLANG_RAGGED_VERIFY_MODE": "static", +} + +# FP4 variant (matches test_deepseek_v4_pro_fp4.py; V4-Pro also auto-detects it). +FP4_ENV_VARS = { + "SGLANG_DSV4_FP4_EXPERTS": "true", +} + + +class TestDSparkUnifiedKVKernelsAMD(CustomTestCase): + def test_build_unified_commit_inject_layout(self): + stride, ring_stride = 7, 128 + req_pool_indices = torch.tensor([3, 0, 5, 1], device=DEVICE, dtype=torch.int32) + prefix_lens = torch.tensor( + [10, 127, 128, 255], device=DEVICE, dtype=torch.int64 + ) + block_pos_offsets = torch.arange(stride, device=DEVICE, dtype=torch.int64) + commit_lens = torch.tensor([0, 3, stride, 5], device=DEVICE, dtype=torch.int32) + + got = dspark_verify_window.build_unified_commit_inject_layout( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + block_pos_offsets=block_pos_offsets, + commit_lens=commit_lens, + stride=stride, + ring_stride=ring_stride, + ) + + positions_2d = prefix_lens.view(-1, 1) + block_pos_offsets[:stride] + loc_2d = req_pool_indices.to(torch.int64).view(-1, 1) * ring_stride + loc_2d = loc_2d + positions_2d % ring_stride + col = torch.arange(stride, device=DEVICE).view(1, -1) + committed = col < commit_lens.to(torch.long).view(-1, 1) + ref_loc = torch.where(committed, loc_2d, torch.full_like(loc_2d, -1)).to( + torch.int32 + ) + + self.assertTrue(torch.equal(got.positions, positions_2d.reshape(-1))) + self.assertTrue(torch.equal(got.swa_loc, ref_loc.reshape(-1))) + + def test_scatter_bf16_into_unified(self): + torch.manual_seed(20) + n_rows, dim, n_pages = 8, 16, 32 + kv = torch.randn(n_rows, dim, device=DEVICE).to(torch.bfloat16).contiguous() + loc = torch.tensor( + [3, -1, 5, 7, 0, -1, 9, 11], device=DEVICE, dtype=torch.int32 + ) + unified = torch.zeros(n_pages, dim, device=DEVICE, dtype=torch.bfloat16) + expected = unified.clone() + keep = loc >= 0 + expected[loc[keep].long()] = kv[keep] + + runtime.scatter_bf16_into_unified(kv=kv, loc=loc, unified_kv=unified) + self.assertTrue(torch.equal(unified, expected)) + + with self.assertRaises(AssertionError): + runtime.scatter_bf16_into_unified(kv=kv, loc=loc, unified_kv=unified.t()) + + +class TestDeepseekV4DSparkUnifiedKVGSM8K(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEEPSEEK_V4_DSPARK_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + env = os.environ.copy() + env.update(COMMON_ENV_VARS) + env.update(DSPARK_ENV_VARS) + env.update(FP4_ENV_VARS) + other_args = [ + "--trust-remote-code", + "--tp", + "8", + "--dp", + "8", + "--enable-dp-attention", + "--enable-dp-lm-head", + "--enable-prefill-delayer", + "--disable-radix-cache", + "--attention-backend", + "dsv4", + "--page-size", + "256", + "--mem-fraction-static", + "0.9", + "--swa-full-tokens-ratio", + "0.15", + "--disable-shared-experts-fusion", + "--tool-call-parser", + "deepseekv4", + "--reasoning-parser", + "deepseek-v4", + "--kv-cache-dtype", + "fp8_e4m3", + "--chunked-prefill-size", + "65536", + "--cuda-graph-max-bs", + "512", + "--max-running-requests", + "512", + "--speculative-algorithm", + "DSPARK", + "--speculative-dspark-block-size", + "5", + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + env=env, + ) + + @classmethod + def tearDownClass(cls): + if getattr(cls, "process", None) is not None: + kill_process_tree(cls.process.pid) + + def test_full_gsm8k_unified_kv_dspark_static(self): + requests.get(self.base_url + "/flush_cache") + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=1319, + parallel=512, + max_new_tokens=512, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"{metrics=}") + + server_info = requests.get(self.base_url + "/server_info") + avg_spec_accept_length = server_info.json()["internal_states"][0][ + "avg_spec_accept_length" + ] + print(f"{avg_spec_accept_length=}") + + if is_in_ci(): + write_github_step_summary( + "### test_gsm8k (deepseek-v4-pro-dspark unified_kv static MI35x)\n" + f"accuracy={metrics['accuracy']:.3f}\n" + f"avg_spec_accept_length={avg_spec_accept_length:.2f}\n" + ) + self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD) + self.assertGreater(avg_spec_accept_length, AVG_SPEC_ACCEPT_LENGTH_THRESHOLD) + + +if __name__ == "__main__": + unittest.main()