[AMD][DSV4] feat: enable DSpark with fp8 unified_kv on gfx950 (#38901)

Co-authored-by: HAI <hixiao@gmail.com>

test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py

@@ -244,6 +244,23 @@ class TestUnifiedFp8SwaScatter(CustomTestCase):
         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)

test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py

@@ -1,5 +1,7 @@
 import contextlib
 import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
 
 import torch
 
@@ -9,7 +11,10 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
     DeepSeekV4TokenToKVPool,
     DeepSeekV4UnifiedKVPool,
     dsv4_unified_row_bytes,
+    resolve_unified_kv_fp8,
 )
+from sglang.srt.mem_cache.kv_cache_configurator import unified_fp8_for_dsv4_pool
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -237,5 +242,213 @@ class TestDSV4UnifiedRegionBuffers(CustomTestCase):
             self.assertAlmostEqual((nope + rope) / whole, 0.625)
 
 
+class TestResolveUnifiedKvFp8(CustomTestCase):
+    def test_override_false_wins_over_env(self):
+        env_mod = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate"
+        with patch(f"{env_mod}.is_unified_kv_fp8", return_value=True):
+            self.assertFalse(resolve_unified_kv_fp8(False))
+            self.assertTrue(resolve_unified_kv_fp8(True))
+            self.assertTrue(resolve_unified_kv_fp8(None))
+
+    def test_none_follows_env_off(self):
+        env_mod = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate"
+        with patch(f"{env_mod}.is_unified_kv_fp8", return_value=False):
+            self.assertFalse(resolve_unified_kv_fp8(None))
+            self.assertFalse(resolve_unified_kv_fp8(False))
+            self.assertTrue(resolve_unified_kv_fp8(True))
+
+
+class TestDsv4PoolFp8Gate(CustomTestCase):
+    _ENV = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate"
+
+    def _layout(self, *, is_draft_worker, algo, env_on=True):
+        with patch(f"{self._ENV}.is_unified_kv_fp8", return_value=env_on):
+            return unified_fp8_for_dsv4_pool(
+                is_draft_worker=is_draft_worker, spec_algorithm=algo
+            )
+
+    def test_dspark_draft_stays_bf16_when_env_on(self):
+        self.assertFalse(
+            self._layout(
+                is_draft_worker=True, algo=SpeculativeAlgorithm.DSPARK, env_on=True
+            )
+        )
+
+    def test_eagle_draft_stays_two_pool_when_env_on(self):
+        self.assertTrue(
+            self._layout(
+                is_draft_worker=True, algo=SpeculativeAlgorithm.EAGLE, env_on=True
+            )
+        )
+
+    def test_target_stays_two_pool_under_dspark_and_eagle(self):
+        for algo in (SpeculativeAlgorithm.DSPARK, SpeculativeAlgorithm.EAGLE):
+            with self.subTest(algo=algo):
+                self.assertTrue(
+                    self._layout(is_draft_worker=False, algo=algo, env_on=True)
+                )
+
+    def test_env_off_is_bf16_for_every_worker(self):
+        for draft, algo in (
+            (True, SpeculativeAlgorithm.DSPARK),
+            (True, SpeculativeAlgorithm.EAGLE),
+            (False, SpeculativeAlgorithm.DSPARK),
+        ):
+            with self.subTest(draft=draft, algo=algo):
+                self.assertFalse(
+                    self._layout(is_draft_worker=draft, algo=algo, env_on=False)
+                )
+
+
+class TestBuildDsv4KvPoolPassesGate(CustomTestCase):
+    class _RecPool:
+        last = None
+
+        def __init__(self, **kwargs):
+            type(self).last = kwargs
+            self._unified_kv = False
+            self._unified_kv_fp8 = kwargs.get("unified_fp8")
+
+    def _kvc(self, *, is_draft_worker, spec_algorithm):
+        from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
+
+        kvc = object.__new__(KVCacheConfigurator)
+        kvc.is_draft_worker = is_draft_worker
+        kvc.spec_algorithm = spec_algorithm
+        kvc.layer_info = SimpleNamespace(
+            num_effective_layers=1, start_layer=0, end_layer=1
+        )
+        kvc.model_config = SimpleNamespace(
+            compress_ratios=[0],
+            window_size=256,
+            qk_nope_head_dim=NOPE_DIM,
+            qk_rope_head_dim=ROPE_DIM,
+            index_head_dim=128,
+            hf_config=SimpleNamespace(kv_source_layer_ids=[]),
+        )
+        kvc.kv_cache_dtype = torch.bfloat16
+        kvc.device = "cpu"
+        return kvc
+
+    def _build(self, *, is_draft_worker, spec_algorithm):
+        kvc = self._kvc(is_draft_worker=is_draft_worker, spec_algorithm=spec_algorithm)
+        sched = MagicMock()
+        sched.page_size = 256
+        exec_cfg = MagicMock()
+        exec_cfg.features.enable_memory_saver = False
+        mem = MagicMock()
+        mem.enable_hisparse = False
+        par = MagicMock()
+        par.attn_dcp_size = 1
+        req = SimpleNamespace(req_to_token=torch.zeros(4, 1))
+        rec = self._RecPool
+        rec.last = None
+        env = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate"
+        with (
+            patch(
+                "sglang.srt.mem_cache.kv_cache_configurator.DeepSeekV4TokenToKVPool",
+                rec,
+            ),
+            patch(
+                "sglang.srt.mem_cache.kv_cache_configurator.get_schedule",
+                return_value=sched,
+            ),
+            patch(
+                "sglang.srt.mem_cache.kv_cache_configurator.get_exec",
+                return_value=exec_cfg,
+            ),
+            patch(
+                "sglang.srt.mem_cache.kv_cache_configurator.get_memory",
+                return_value=mem,
+            ),
+            patch(
+                "sglang.srt.mem_cache.kv_cache_configurator.get_parallel",
+                return_value=par,
+            ),
+            patch(
+                "sglang.srt.mem_cache.kv_cache_configurator.max_speculative_num_draft_tokens",
+                return_value=0,
+            ),
+            patch(f"{env}.is_unified_kv_fp8", return_value=True),
+        ):
+            kvc._build_dsv4_kv_pool(
+                max_running_requests=2,
+                full_max_total_num_tokens=256,
+                swa_max_total_num_tokens=256,
+                c4_max_total_num_tokens=0,
+                c128_max_total_num_tokens=1,
+                c4_state_pool_size=0,
+                c128_state_pool_size=0,
+                c4_state_dtype=None,
+                c128_state_dtype=None,
+                req_to_token_pool=req,
+            )
+        return rec.last
+
+    def test_dspark_draft_ctor_gets_unified_fp8_false(self):
+        kw = self._build(
+            is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.DSPARK
+        )
+        self.assertIsNotNone(kw)
+        self.assertFalse(kw["unified_fp8"])
+
+    def test_eagle_draft_ctor_gets_unified_fp8_true(self):
+        kw = self._build(
+            is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.EAGLE
+        )
+        self.assertIsNotNone(kw)
+        self.assertTrue(kw["unified_fp8"])
+
+    def test_target_ctor_gets_unified_fp8_true(self):
+        kw = self._build(
+            is_draft_worker=False, spec_algorithm=SpeculativeAlgorithm.DSPARK
+        )
+        self.assertIsNotNone(kw)
+        self.assertTrue(kw["unified_fp8"])
+
+
+class TestUnifiedKvPoolFollowsCtorFp8(CustomTestCase):
+    def _pool(self, fp8):
+        return DeepSeekV4UnifiedKVPool(
+            stage_ratios=[0],
+            num_slots=2,
+            num_blocks=1,
+            page_size=256,
+            qk_nope_head_dim=NOPE_DIM,
+            qk_rope_head_dim=ROPE_DIM,
+            device="cpu",
+            memory_saver_adapter=_StubMemorySaver(),
+            custom_mem_pool=None,
+            swa_ring_size=8,
+            fp8=fp8,
+        )
+
+    def test_dspark_draft_layout_has_no_rope_pool(self):
+        env = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate"
+        with patch(f"{env}.is_unified_kv_fp8", return_value=True):
+            fp8 = unified_fp8_for_dsv4_pool(
+                is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.DSPARK
+            )
+        self.assertFalse(fp8)
+        pool = self._pool(fp8)
+        buf = pool.kv_buffer[0]
+        self.assertEqual(buf.dtype, torch.bfloat16)
+        self.assertEqual(buf.shape[1], NOPE_DIM + ROPE_DIM)
+        self.assertIsNone(pool.kv_buffer_rope[0])
+
+    def test_eagle_draft_layout_has_rope_pool(self):
+        env = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate"
+        with patch(f"{env}.is_unified_kv_fp8", return_value=True):
+            fp8 = unified_fp8_for_dsv4_pool(
+                is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.EAGLE
+            )
+        self.assertTrue(fp8)
+        pool = self._pool(fp8)
+        buf, rope = pool.kv_buffer[0], pool.kv_buffer_rope[0]
+        self.assertEqual(buf.dtype, torch.float8_e4m3fn)
+        self.assertEqual(rope.dtype, torch.bfloat16)
+        self.assertEqual(rope.shape[1], ROPE_DIM)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_pool_configurator.py

@@ -1223,6 +1223,7 @@ class TestSWAPoolFloor(CustomTestCase):
         cfg.request_window_bytes = 0
         cfg.bytes_per_swa_token = 0.0
         cfg._unified_fp8 = False
+        cfg._dspark_draft_on_bf16 = False
         # object.__new__ skips __init__; bf16 unified row is 2B * latent
         cfg._unified_row_bytes = cfg.attn_head_dim * 2
         return cfg
@@ -1291,6 +1292,38 @@ class TestSWAPoolFloor(CustomTestCase):
         self.assertEqual(sizes.swa_max_total_num_tokens, 3072)
         self.assertEqual(sizes.c4_state_pool_size, 0)
 
+    def test_dsv4_fp8_dspark_swa_ring_includes_bf16_draft(self):
+        """Target fp8 ring + one-layer bf16 draft ring, not (T+1)/T * fp8."""
+        from sglang.srt.mem_cache.deepseek_v4_memory_pool import dsv4_unified_row_bytes
+
+        cfg = self._dsv4_configurator_for_budget()
+        cfg._unified_fp8 = True
+        cfg._unified_row_bytes = dsv4_unified_row_bytes(448, 64, fp8=True)
+        cfg.qk_nope_head_dim, cfg.qk_rope_head_dim = 448, 64
+        cfg._dspark_draft_on_bf16 = True
+        cfg._spec_infl = (cfg.num_layers_total + 1) / cfg.num_layers_total
+        mrr = 32
+        slots = cfg._get_num_req_slots(mrr)
+        got = cfg._fixed_swa_bytes(mrr)
+        target = (
+            slots * cfg._swa_ring_size * cfg._unified_row_bytes * cfg.num_layers_total
+        )
+        draft = slots * cfg._swa_ring_size * dsv4_unified_row_bytes(448, 64, False)
+        self.assertEqual(got, target + draft)
+        mtp_formula = int(target * cfg._spec_infl)
+        self.assertGreater(got, mtp_formula)
+
+    def test_dsv4_fp8_mtp_swa_ring_keeps_spec_inflation(self):
+        cfg = self._dsv4_configurator_for_budget()
+        cfg._unified_fp8 = True
+        cfg._unified_row_bytes = 640
+        cfg._dspark_draft_on_bf16 = False
+        cfg._spec_infl = (cfg.num_layers_total + 1) / cfg.num_layers_total
+        mrr = 32
+        slots = cfg._get_num_req_slots(mrr)
+        target = slots * cfg._swa_ring_size * 640 * cfg.num_layers_total
+        self.assertEqual(cfg._fixed_swa_bytes(mrr), int(target * cfg._spec_infl))
+
 
 if __name__ == "__main__":
     unittest.main()