minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
amd-danli103
2026-09-20 01:16:39 -07:00
committed by GitHub
co-authored by HAI
parent a8a4d86be9
commit e54009240a
6 changed files with 333 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+27 -8
View File
@@ -78,6 +78,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)
def _num_dsv4_physical_kv_pages(
size: int, physical_page_size: int, logical_page_size: int
) -> int:
@@ -867,6 +880,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
enable_hisparse: bool = False,
online_mtp_max_draft_tokens: int = 0,
num_req_slots: Optional[int] = None,
unified_fp8: Optional[bool] = None,
kv_source_layers: Sequence[int] = (),
full_size: Optional[int] = None,
is_draft_worker: bool = False,
@@ -893,13 +907,6 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
self.compressed_kv_layout_option = compressed_kv_layout
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.
@@ -915,7 +922,19 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
# 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
+16
View File
@@ -133,6 +133,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).
@@ -1365,6 +1375,11 @@ class KVCacheConfigurator:
kv_layout=kv_layout, compressed_kv_layout=compressed_kv_layout
)
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
@@ -1392,6 +1407,7 @@ class KVCacheConfigurator:
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,
kv_source_layers=kv_source_layers,
full_size=full_max_total_num_tokens,
**({"is_draft_worker": self.is_draft_worker} if not _is_npu else {}),
python/sglang/srt/model_executor/pool_configurator.py
+27 -6
View File
@@ -968,6 +968,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
@@ -986,6 +989,11 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
# Resolve the unified-kv gate before any sizing so the two cannot drift.
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()
)
# Row width across both unified pools: 1024 B bf16, 640 B fp8.
self._unified_row_bytes = dsv4_unified_row_bytes(
self.qk_nope_head_dim, self.qk_rope_head_dim, self._unified_fp8
@@ -1395,20 +1403,32 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
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
@@ -1486,6 +1506,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
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
+17
View File
@@ -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
+213
View File
@@ -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
+33
View File
@@ -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()