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

[AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated (#38192)

Browse Source 
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
yuttian1
2026-09-07 13:13:04 -07:00
committed by GitHub
co-authored by hnyls2002 Liangsheng Yin
parent 6287ebf43a
commit 570087ceda
20 changed files with 763 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/attention/test_c4_v2.py
+90 -1
View File
@@ -7,7 +7,11 @@ import pytest
import torch
import triton
from sglang.kernels.ops.attention.dsv4 import compress_forward
from sglang.kernels.ops.attention.dsv4 import (
CompressorDecodePlan,
CompressorPrefillPlan,
compress_forward,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kernels.deepseek_v4.common import (
@@ -122,6 +126,91 @@ def _make_inputs(
# -----------------------------------------------------------------------------
@pytest.mark.parametrize("ring_size", [8, 16])
@pytest.mark.parametrize(
("gpu_inputs", "use_cuda_graph"),
[(False, False), (True, False), (True, True)],
)
def test_unified_request_ring_plans_ignore_full_to_state(
ring_size: int, gpu_inputs: bool, use_cuda_graph: bool
) -> None:
"""C4 plans must address state by request slot, not the full-cache map."""
device = torch.device(get_device())
req_pool_indices = torch.tensor([2, 5], dtype=torch.int64, device=device)
req_to_token = torch.zeros((6, 16), dtype=torch.int32, device=device)
full_to_state = torch.zeros(1, dtype=torch.int64, device=device)
seq_lens = torch.tensor([8, 12], dtype=torch.int64)
extend_lens = torch.tensor([4, 4], dtype=torch.int64)
if gpu_inputs:
seq_lens = seq_lens.to(device)
extend_lens = extend_lens.to(device)
prefill = CompressorPrefillPlan.generate(
compress_ratio=RATIO,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
extend_lens=extend_lens,
req_to_token=req_to_token,
full_to_state=full_to_state,
swa_page_size=256,
ring_size=ring_size,
num_q_tokens=8,
use_cuda_graph=use_cuda_graph,
use_req_ring=True,
)
plan_c = prefill.plan_c.view(torch.int32).reshape(-1, 4).cpu()
plan_w = prefill.plan_w.view(torch.int32).reshape(-1, 2).cpu()
valid_c = plan_c[plan_c[:, 2] >= 0]
got_reads = {
int(row[1].item()) & 0xFFFF: (int(row[2].item()), int(row[3].item()))
for row in valid_c
}
expected_reads = {
3: (
(2 * ring_size + 3 % ring_size) // RATIO,
(2 * ring_size + 7 % ring_size) // RATIO,
),
7: (
(5 * ring_size + 7 % ring_size) // RATIO,
(5 * ring_size + 11 % ring_size) // RATIO,
),
}
assert got_reads == expected_reads
valid_w = plan_w[plan_w[:, 1] >= 0]
got_writes = {int(row[0].item()): int(row[1].item()) for row in valid_w}
expected_writes = {
**{j: 2 * ring_size + (4 + j) % ring_size for j in range(4)},
**{4 + j: 5 * ring_size + (8 + j) % ring_size for j in range(4)},
}
assert got_writes == expected_writes
assert {got_writes[j] for j in range(4)}.isdisjoint(
{got_writes[j] for j in range(4, 8)}
)
decode = CompressorDecodePlan.generate(
compress_ratio=RATIO,
req_pool_indices=req_pool_indices,
req_to_token=req_to_token,
full_to_state=full_to_state,
seq_lens=torch.tensor([8, 12], dtype=torch.int64, device=device),
swa_page_size=256,
ring_size=ring_size,
use_req_ring=True,
)
got_decode = decode.plan_d.view(torch.int32).reshape(-1, 4).cpu()
expected_decode = torch.tensor(
[
[8, 2 * ring_size + 7 % ring_size, *expected_reads[3]],
[12, 5 * ring_size + 11 % ring_size, *expected_reads[7]],
],
dtype=torch.int32,
)
assert torch.equal(got_decode, expected_decode)
assert got_decode[0, 1] != got_decode[1, 1]
@pytest.mark.parametrize("mode", ["legacy", "paged"])
@pytest.mark.parametrize("seq_len", [4, 8, 32, 256, 1024])
def test_prefill_no_context(mode: str, seq_len: int) -> None:
test/registered/unit/mem_cache/test_dsv4_c4_state_lifecycle.py
+162
View File
@@ -0,0 +1,162 @@
"""CPU/mock tests for unified DSV4 C4 request-state lifecycle."""
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
from sglang.srt.disaggregation.decode import DecodeReqToTokenPool
from sglang.srt.mem_cache.allocation import alloc_req_slots
from sglang.srt.mem_cache.deepseek_v4_compress_state import KVAndScore
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _request(req_pool_idx=None, *, reused=False):
return SimpleNamespace(
kv=SimpleNamespace(
req_pool_idx=req_pool_idx,
kv_committed_len=1 if reused else 0,
kv_allocated_len=1 if reused else 0,
holds_kv=reused,
),
inflight_middle_chunks=1 if reused else 0,
)
def _mark_reused(req):
req.kv.kv_committed_len = 1
req.kv.kv_allocated_len = 1
req.kv.holds_kv = True
req.inflight_middle_chunks = 1
def _c4_pool(rows: int, width: int, ring_size: int):
return SimpleNamespace(
ratio=4,
ring_size=ring_size,
kv_score_buffer=KVAndScore(torch.full((rows, width), 7.0)),
)
def _token_pool(unified: bool, ring_size: int = 8):
logical_rows = 4 * ring_size
physical_rows = logical_rows + ring_size + 4
attn = _c4_pool(physical_rows, width=12, ring_size=ring_size)
indexer = _c4_pool(physical_rows, width=8, ring_size=ring_size)
c128 = SimpleNamespace(
ratio=128,
ring_size=128,
kv_score_buffer=KVAndScore(torch.full((physical_rows, 8), 9.0)),
)
token_pool = object.__new__(DeepSeekV4TokenToKVPool)
token_pool._unified_kv = unified
token_pool.compress_state_pools = [attn, c128]
token_pool.indexer_compress_state_pools = [indexer, None]
token_pool.get_ring_size = MagicMock(return_value=ring_size)
return token_pool, attn, indexer, c128, logical_rows
class TestUnifiedC4StateLifecycle(unittest.TestCase):
def test_pool_size_is_exact_request_ring_product(self):
configurator = object.__new__(DSV4PoolConfigurator)
configurator.disaggregation_mode = "decode"
configurator.disaggregation_decode_extra_slots = 3
configurator.c4_ring_size = 16
self.assertEqual(configurator._unified_c4_state_pool_size(10), 14 * 16)
def test_clear_resets_only_selected_request_rings(self):
ring_size = 8
token_pool, attn, indexer, c128, logical_rows = _token_pool(
unified=True, ring_size=ring_size
)
token_pool.clear_c4_req_states([1, 3])
selected = torch.tensor(list(range(8, 16)) + list(range(24, 32)))
untouched = torch.tensor(list(range(0, 8)) + list(range(16, 24)))
for pool in (attn, indexer):
state = pool.kv_score_buffer.kv_score
half = state.shape[-1] // 2
self.assertTrue(
torch.equal(
state[selected, :half], torch.zeros_like(state[selected, :half])
)
)
self.assertTrue(torch.isneginf(state[selected, half:]).all())
self.assertTrue((state[untouched] == 7).all())
self.assertTrue((state[logical_rows:] == 7).all())
self.assertTrue((c128.kv_score_buffer.kv_score == 9).all())
def test_clear_is_noop_off_the_unified_path(self):
"""The non-unified (fp8) pool addresses C4 state by SWA page, so a
req-slot reset must not touch it."""
token_pool, attn, indexer, _, _ = _token_pool(unified=False)
token_pool.clear_c4_req_states([1, 3])
for pool in (attn, indexer):
self.assertTrue((pool.kv_score_buffer.kv_score == 7).all())
def test_req_pool_hook_fires_for_new_slots_only(self):
req_pool = ReqToTokenPool(3, 16, "cpu", enable_memory_saver=False)
hook = MagicMock()
req_pool.register_on_alloc_rows(hook)
reused = _request()
# First admission: a brand-new slot, so its C4 ring must be cleared.
(reused_idx,) = alloc_req_slots(req_pool, [reused], None)
hook.assert_called_once_with([reused_idx])
# Chunked continuation reuses the same slot -- clearing it here would
# wipe the state captured by the previous chunk.
hook.reset_mock()
_mark_reused(reused)
self.assertEqual(alloc_req_slots(req_pool, [reused], None), [reused_idx])
hook.assert_not_called()
# Mixed batch: only the newly allocated slot is reported.
fresh = _request()
indices = alloc_req_slots(req_pool, [reused, fresh], None)
self.assertEqual(indices[0], reused_idx)
self.assertNotEqual(indices[1], reused_idx)
hook.assert_called_once_with([indices[1]])
def test_decode_req_pool_hook_fires_for_new_slots_only(self):
"""PD decode pre-allocates through DecodeReqToTokenPool, which has its
own alloc; it must report fresh rows the same way."""
req_pool = DecodeReqToTokenPool(
2, 16, "cpu", enable_memory_saver=False, pre_alloc_size=2
)
hook = MagicMock()
req_pool.register_on_alloc_rows(hook)
first = _request()
(first_idx,) = req_pool.alloc([first])
hook.assert_called_once_with([first_idx])
hook.reset_mock()
_mark_reused(first)
second = _request()
indices = req_pool.alloc([first, second])
self.assertEqual(indices[0], first_idx)
hook.assert_called_once_with([indices[1]])
hook.reset_mock()
self.assertEqual(req_pool.alloc([first]), [first_idx])
hook.assert_not_called()
def test_req_pool_without_hook_is_unchanged(self):
req_pool = ReqToTokenPool(2, 16, "cpu", enable_memory_saver=False)
(idx,) = alloc_req_slots(req_pool, [_request()], None)
self.assertGreater(idx, 0)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py
+2
View File
@@ -36,6 +36,8 @@ def _make_self(*, page_size: int, full_available: int, swa_available: int):
return SimpleNamespace(
page_size=page_size,
# alloc_extend reads _swa_req_ring; pin the paged-SWA path.
_swa_req_ring=False,
full_attn_allocator=SimpleNamespace(
available_size=lambda: full_available,
alloc_extend=MagicMock(return_value=full_indices),
test/registered/unit/model_executor/test_pool_configurator.py
+73 -2
View File
@@ -1044,9 +1044,9 @@ class TestSWAPoolFloor(CustomTestCase):
)
self.assertEqual(config.swa_max_total_num_tokens, 3072)
def _dsv4_sizes(self, max_tokens, page_size):
def _dsv4_sizes(self, max_tokens, page_size, unified=False):
"""Exercise the DSV4 size arithmetic without a full V4 model fixture:
_compute_dsv4_sizes reads only these five attributes."""
_compute_dsv4_sizes reads only these six attributes."""
from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator
cfg = object.__new__(DSV4PoolConfigurator)
@@ -1055,6 +1055,7 @@ class TestSWAPoolFloor(CustomTestCase):
cfg.swa_page_size = 128
cfg.c4_ring_size = 8
cfg.c4_shrink_factor = 1
cfg._unified = unified
return cfg._compute_dsv4_sizes(max_tokens, page_size)
def test_dsv4_rejects_single_page_pool(self):
@@ -1068,6 +1069,76 @@ class TestSWAPoolFloor(CustomTestCase):
sizes = self._dsv4_sizes(max_tokens=32768, page_size=256)
self.assertEqual(sizes.full_max_total_num_tokens, 32768)
self.assertEqual(sizes.swa_max_total_num_tokens, 3072)
# Non-unified: the c4 state pool scales with the paged SWA pool.
self.assertEqual(sizes.c4_state_pool_size, 3072 // 128 * 8)
def test_dsv4_token_cap_never_grows_total_footprint(self):
"""Regression: the token-cap path subtracts no fixed-pool bias, so
capping the budget-derived token count must still shrink the total."""
cfg = self._dsv4_configurator_for_budget()
page_size = 128
budget = 256 * (1 << 30)
base = cfg.calculate_pool_sizes(budget, page_size)
base_bytes = self._dsv4_total_bytes(cfg, base.max_total_num_tokens)
self.assertLessEqual(base_bytes, budget)
for numerator in (999, 900, 500, 100, 1):
capped_tokens = (
base.max_total_num_tokens * numerator // 1000 // page_size * page_size
)
if capped_tokens <= 0:
continue
capped = cfg.calculate_pool_sizes_from_max_tokens(capped_tokens, page_size)
capped_bytes = self._dsv4_total_bytes(cfg, capped.max_total_num_tokens)
with self.subTest(numerator=numerator):
self.assertLessEqual(capped_bytes, base_bytes)
# White-box 671B-class shape: the byte arithmetic runs without a model fixture.
def _dsv4_configurator_for_budget(self):
from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator
cfg = object.__new__(DSV4PoolConfigurator)
cfg.qk_nope_head_dim, cfg.qk_rope_head_dim = 128, 64
cfg.attn_head_dim = 192
cfg.indexer_head_dim = 128
cfg.num_layers_total = 61
cfg.num_layers_ca4 = 61
cfg.num_layers_ca128 = 61
cfg.c4_ring_size = 8
cfg.c128_ring_size = 128
cfg._swa_ring_size = 128
cfg._spec_infl = 1.0
cfg.context_len = 65536
cfg.bytes_per_full_token = 576.0
cfg.requested_max_running_requests_per_worker = None
cfg.swa_ratio = 0.1
cfg.sliding_window_size = 4096
cfg.swa_page_size = 128
cfg.c4_shrink_factor = 1
cfg.online_c128_mtp_max_draft_tokens = 0
cfg.disaggregation_mode = None
cfg.disaggregation_decode_extra_slots = 0
cfg._unified = True
return cfg
# Token pool plus the three request-scoped fixed pools, sized from the
# concurrency resolve_max_num_reqs derives from this token count.
def _dsv4_total_bytes(self, cfg, tokens):
estimated = max(min(int(tokens / cfg.context_len * 512), 4096), 2048)
max_running_requests = min(estimated, tokens // 2)
return int(
tokens * cfg.bytes_per_full_token
+ cfg._fixed_swa_bytes(max_running_requests)
+ cfg._fixed_c4_state_bytes(max_running_requests)
+ cfg._get_c128_state_fixed_bytes(max_running_requests)
)
def test_dsv4_unified_c4_state_not_token_scaled(self):
# Unified-KV sizes the c4 state ring from max_running_requests in
# finalize_with_max_running_requests, so it must not scale here.
sizes = self._dsv4_sizes(max_tokens=32768, page_size=256, unified=True)
self.assertEqual(sizes.full_max_total_num_tokens, 32768)
self.assertEqual(sizes.swa_max_total_num_tokens, 3072)
self.assertEqual(sizes.c4_state_pool_size, 0)
if __name__ == "__main__":