[AMD][DSV4] feat: enable DSpark with fp8 unified_kv on gfx950

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -68,6 +68,19 @@ def get_swa_ring_size(sliding_window: int, is_speculative: bool = False) -> int:
     return sliding_window + spec_extra
 
 
+def resolve_unified_kv_fp8(unified_fp8: Optional[bool] = None) -> bool:
+    """Per-pool fp8 layout. None follows SGLANG_DSV4_UNIFIED_KV_FP8.
+
+    A caller may pass False so this pool keeps the bf16 ring while the env
+    stays on (target fused-Q still keys off the global switch).
+    """
+    from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+        is_unified_kv_fp8,
+    )
+
+    return is_unified_kv_fp8() if unified_fp8 is None else bool(unified_fp8)
+
+
 class DeepSeekV4SingleKVPool(KVCache):
     def __init__(
         self,
@@ -679,6 +692,7 @@ def __init__(
         enable_hisparse: bool = False,
         online_mtp_max_draft_tokens: int = 0,
         num_req_slots: Optional[int] = None,
+        unified_fp8: Optional[bool] = None,
     ):
         super().__init__(
             swa_size,
@@ -692,13 +706,6 @@ def __init__(
         )
         c4_logical_size = c128_size * 32
 
-        logger.info(
-            "Initialize DeepSeekV4TokenToKVPool with "
-            f"{max_num_reqs=} {swa_size=} {c4_size=} "
-            f"{c4_logical_size=} {c128_size=} "
-            f"{c4_state_pool_size=} {c128_state_pool_size=}"
-        )
-
         self.max_num_reqs = max_num_reqs
         # PD preallocation can exceed max_num_reqs;
         # the SWA ring must cover every addressable req_pool_idx.
@@ -714,7 +721,19 @@ def __init__(
 
         # Resolve the unified-kv gate before any sizing so the two cannot drift.
         self._unified_kv = is_unified_kv_triton()
-        self._unified_kv_fp8 = is_unified_kv_fp8()
+        self._unified_kv_fp8 = resolve_unified_kv_fp8(unified_fp8)
+        logger.info(
+            "Initialize DeepSeekV4TokenToKVPool with "
+            f"{max_num_reqs=} {swa_size=} {c4_size=} "
+            f"{c4_logical_size=} {c128_size=} "
+            f"{c4_state_pool_size=} {c128_state_pool_size=} "
+            f"unified={self._unified_kv} unified_fp8={self._unified_kv_fp8}"
+        )
+        if is_unified_kv_fp8() and not self._unified_kv_fp8:
+            logger.info(
+                "SGLANG_DSV4_UNIFIED_KV_FP8 is on; this pool stays bf16 "
+                "(unified_fp8=False)"
+            )
         # Uniform 512-dim e4m3 layout for the trtllm attention backend
         self.uniform_fp8 = (
             not self._unified_kv

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -131,6 +131,16 @@ def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]:
 _is_npu = is_npu()
 
 
+def unified_fp8_for_dsv4_pool(*, is_draft_worker: bool, spec_algorithm) -> bool:
+    """Per-pool fp8 layout. DSpark draft writers scatter bf16, so that pool
+    stays a bf16 ring; MTP/EAGLE NextN follows the env."""
+    from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+        is_unified_kv_fp8,
+    )
+
+    return is_unified_kv_fp8() and not (is_draft_worker and spec_algorithm.is_dspark())
+
+
 def _should_enable_lazy_compaction() -> bool:
     """Lazy compaction default — ON unless
     `SGLANG_DISABLE_LAZY_COMPACTION=1` (escape hatch for A/B / rollback).
@@ -1377,6 +1387,11 @@ def _build_dsv4_kv_pool(
         else:
             pool_cls = DeepSeekV4TokenToKVPool
 
+        unified_fp8 = unified_fp8_for_dsv4_pool(
+            is_draft_worker=self.is_draft_worker,
+            spec_algorithm=self.spec_algorithm,
+        )
+
         token_to_kv_pool = pool_cls(
             max_num_reqs=max_running_requests,
             # SWA ring is indexed by req_pool_idx; PD decode inflates req_to_token
@@ -1404,6 +1419,7 @@ def _build_dsv4_kv_pool(
             end_layer=self.layer_info.end_layer,
             enable_hisparse=get_memory().enable_hisparse,
             online_mtp_max_draft_tokens=(max_speculative_num_draft_tokens() or 0),
+            unified_fp8=unified_fp8,
         )
         if not self.is_draft_worker and token_to_kv_pool._unified_kv:
             # The draft pool has no C4 layers and shares this req pool, so only

python/sglang/srt/model_executor/pool_configurator.py

@@ -934,6 +934,9 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
     is the request-scoped fixed pools that do not scale with full_token.
     """
 
+    # object.__new__ stubs (SWA floor tests) skip __init__
+    _dspark_draft_on_bf16 = False
+
     def __init__(self, kvc: KVCacheConfigurator):
         self.kv_cache_dtype_str = kvc.kv_cache_dtype_str
         cfg = kvc.model_config
@@ -1006,6 +1009,11 @@ def __init__(self, kvc: KVCacheConfigurator):
 
         self._unified = is_unified_kv_triton()
         self._unified_fp8 = is_unified_kv_fp8()
+        # DSpark draft still allocates a bf16 ring; target fp8 * (T+1)/T would
+        # under-count that ring (640 vs 1024). MTP keeps the old inflation.
+        self._dspark_draft_on_bf16 = bool(
+            self._unified_fp8 and kvc.spec_algorithm.is_dspark()
+        )
         self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
         # Row width across both pools: 1024 B bf16, 640 B fp8. Read from the pool
         # module so sizing can't drift from the allocation.
@@ -1244,20 +1252,32 @@ def _resolve_max_running_requests_per_worker(self, available_bytes: int) -> int:
 
     def _fixed_swa_bytes(self, max_running_requests: int) -> int:
         """Unified_kv SWA is a fixed per-request ring, sized by concurrency
-        (num_req_slots) rather than by full_token. Return its byte footprint
-        across all full layers, inflated for the draft worker the same way as the
-        per-token coeff. Returns 0 on the non-unified path (where SWA is already
-        accounted per-token)."""
+        (num_req_slots) rather than by full_token. MTP inflates the target ring
+        by _spec_infl; DSpark+fp8 adds a bf16 draft ring instead (640 vs 1024).
+        Returns 0 on the non-unified path (SWA already counted per-token)."""
         if not self._unified:
             return 0
         num_req_slots = self._get_num_req_slots(max_running_requests)
-        ring_bytes = (
+        target_ring = (
             num_req_slots
             * self._swa_ring_size
             * self._unified_row_bytes
             * self.num_layers_total
         )
-        return int(ring_bytes * self._spec_infl)
+        if self._dspark_draft_on_bf16:
+            from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+                dsv4_unified_row_bytes,
+            )
+
+            draft_row = dsv4_unified_row_bytes(
+                self.qk_nope_head_dim, self.qk_rope_head_dim, fp8=False
+            )
+            # 1 layer is what the shipped DSpark drafts allocate. A multi-stage
+            # draft would under-count by ~9 MB/layer (128-wide window, ~65 req
+            # slots), which the (T+1)/T on bytes_per_full_token already covers.
+            draft_ring = num_req_slots * self._swa_ring_size * draft_row
+            return int(target_ring + draft_ring)
+        return int(target_ring * self._spec_infl)
 
     def _to_config(self, sizes: _DSV4PoolSizes) -> MemoryPoolConfig:
         full = sizes.full_max_total_num_tokens
@@ -1326,6 +1346,7 @@ def calculate_pool_sizes(
         logger.info(
             f"DSV4 memory calculation: unified={self._unified}, "
             f"unified_fp8={self._unified_fp8}, "
+            f"dspark_draft_bf16={self._dspark_draft_on_bf16}, "
             f"bytes_per_full_token={self.bytes_per_full_token:.2f}, "
             f"available_bytes={available_bytes / (1 << 30):.2f} GB, "
             f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, "

test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py

@@ -244,6 +244,23 @@ def test_dtype_mismatch_rejected(self):
         with self.assertRaises(AssertionError):
             _store(kv, pool_nope, self.state_slot, self.positions)
 
+    def test_scatter_bf16_still_writes_a_bf16_draft_pool(self):
+        """Cut A: DSpark draft stays on the bf16 ring; scatter_bf16 must not
+        dtype-assert against that pool even when the fp8 env is on."""
+        kv = torch.randn(
+            self.n_rows, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16
+        ).contiguous()
+        loc = (
+            self.state_slot.long() * RING_STRIDE + self.positions.long() % RING_STRIDE
+        ).to(torch.int32)
+        loc[1] = -1
+        pool = torch.zeros(N_PAGES, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16)
+        expected = pool.clone()
+        keep = loc >= 0
+        expected[loc[keep].long()] = kv[keep]
+        runtime.scatter_bf16_into_unified(kv=kv, loc=loc, unified_kv=pool)
+        self.assertTrue(torch.equal(pool, expected))
+
     def test_empty_batch_is_a_noop(self):
         empty_slot = torch.zeros(0, device=DEVICE, dtype=torch.int32)
         kv_nope, _ = _packed_nope(0)