dsv4.1: remaining model and runtime integration (#38798)
Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Xiaoyu Zhang <xiaoyu.zhang@radixark.ai> Co-authored-by: Yuwei An <ayw.sirius19@gmail.com> Co-authored-by: Khoa Pham <khoa.pham@radixark.ai> Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Co-authored-by: Zhichen Zeng <zczeng@uw.edu> Co-authored-by: Ke Bao <ispobaoke@gmail.com>
This commit is contained in:
Liangsheng Yin
co-authored by
BBuf
Claude Opus 5
Xiaoyu Zhang
Yuwei An
Khoa Pham
Yuhao Yang
Zhichen Zeng
Ke Bao
BBuf
Claude Opus 5
Xiaoyu Zhang
Yuwei An
Khoa Pham
Yuhao Yang
Zhichen Zeng
Ke Bao
parent
1b200ffaaa
commit
a6cf05817f
@@ -5,11 +5,17 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.attention.dsv4.kv_layout import (
|
||||
KVLayout,
|
||||
is_valid_kv_layout_pair,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
|
||||
DeepSeekV4SingleKVPool,
|
||||
DeepSeekV4TokenToKVPool,
|
||||
_CompressedPoolConfig,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
@@ -25,6 +31,8 @@ class TestDSV4CompressedPools(CustomTestCase):
|
||||
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
|
||||
pool._unified_kv = unified
|
||||
pool.uniform_fp8 = False
|
||||
pool.kv_layout = KVLayout.V4
|
||||
pool.compressed_kv_layout_option = None
|
||||
pool.compressed_pool_configs = {
|
||||
4: _CompressedPoolConfig(
|
||||
256, 64, torch.bfloat16, indexer_size=1024
|
||||
@@ -176,5 +184,151 @@ class TestDSV4CompressedPools(CustomTestCase):
|
||||
pool.get_index_k_page_size(128)
|
||||
|
||||
|
||||
HEAD_DIM = 512
|
||||
ROPE_DIM = 64
|
||||
PAGE_SIZE = 256
|
||||
FULL_SIZE = 4 * PAGE_SIZE
|
||||
|
||||
|
||||
class TestV41KVPoolLayouts(CustomTestCase):
|
||||
"""A V4.1-layout pool hands the attention kernel page-aligned buffers and
|
||||
picks the compressed layout each ratio asks for."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
override = get_context().override_server_args(page_size=PAGE_SIZE)
|
||||
override.install()
|
||||
self.addCleanup(override.restore)
|
||||
|
||||
def make_pool(self, ratios, kv_source_layers, kv_layout, compressed=None, **sizes):
|
||||
return DeepSeekV4TokenToKVPool(
|
||||
max_num_reqs=16,
|
||||
swa_size=FULL_SIZE,
|
||||
c4_size=sizes.get("c4_size", 0),
|
||||
c128_size=sizes.get("c128_size", 0),
|
||||
c4_state_pool_size=sizes.get("c4_state_pool_size", 0),
|
||||
c128_state_pool_size=sizes.get("c128_state_pool_size", 0),
|
||||
page_size=PAGE_SIZE,
|
||||
swa_page_size=PAGE_SIZE,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
c4_state_dtype=torch.float32,
|
||||
c128_state_dtype=torch.float32,
|
||||
qk_nope_head_dim=HEAD_DIM - ROPE_DIM,
|
||||
qk_rope_head_dim=ROPE_DIM,
|
||||
indexer_head_dim=128,
|
||||
layer_num=len(ratios),
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
compression_ratios=ratios,
|
||||
kv_source_layers=kv_source_layers,
|
||||
full_size=FULL_SIZE,
|
||||
kv_layout=kv_layout,
|
||||
compressed_kv_layout=compressed,
|
||||
)
|
||||
|
||||
def assert_kernel_requirements(self, pool, layout):
|
||||
"""Pages start on the kernel's alignment, and its
|
||||
(num_pages, page_size, 1, bytes_per_token) view walks one token per row."""
|
||||
for buf in pool.kv_buffer:
|
||||
self.assertEqual(buf.stride(0) % layout.page_align, 0)
|
||||
bpt = layout.bytes_per_token
|
||||
view = buf[:, : pool.page_size * bpt].view(
|
||||
buf.shape[0], pool.page_size, 1, bpt
|
||||
)
|
||||
self.assertEqual(view.stride(1), bpt)
|
||||
self.assertEqual(view.stride(0), pool.bytes_per_page_padded)
|
||||
|
||||
def test_v41_pool_buffers(self):
|
||||
for option, expect in ((None, KVLayout.V41_FP4), ("fp8", KVLayout.V41)):
|
||||
with self.subTest(compressed=option):
|
||||
pool = self.make_pool([0, 0, 2, 1, 1], [2, 3], KVLayout.V41, option)
|
||||
self.assert_kernel_requirements(pool.swa_kv_pool, KVLayout.V41)
|
||||
self.assertEqual(pool.get_swa_key_bytes_per_token(), 528)
|
||||
for ratio in (1, 2):
|
||||
layer_id = pool.sources_by_ratio[ratio][0]
|
||||
self.assertIs(pool.get_extra_key_layout(layer_id), expect)
|
||||
self.assertEqual(
|
||||
pool.get_extra_key_bytes_per_token(layer_id),
|
||||
expect.bytes_per_token,
|
||||
)
|
||||
self.assertTrue(is_valid_kv_layout_pair(pool.kv_layout, expect))
|
||||
self.assert_kernel_requirements(pool.kv_pools[ratio], expect)
|
||||
# A pool of the fp4 layout cannot be the main cache.
|
||||
with self.assertRaises(AssertionError):
|
||||
self.make_pool([0], [], KVLayout.V41_FP4)
|
||||
|
||||
def test_v41_pool_with_c4_c128(self):
|
||||
pool = self.make_pool(
|
||||
[0, 4, 128],
|
||||
[],
|
||||
KVLayout.V41,
|
||||
c4_size=PAGE_SIZE,
|
||||
c128_size=PAGE_SIZE,
|
||||
c4_state_pool_size=16,
|
||||
c128_state_pool_size=16,
|
||||
)
|
||||
for ratio in (4, 128):
|
||||
self.assertEqual(pool.kv_pools[ratio].page_size, PAGE_SIZE // ratio)
|
||||
# The 2-token c128 page is the only production page that pads.
|
||||
self.assertEqual(pool.kv_pools[128].bytes_per_page_padded, 1536)
|
||||
|
||||
|
||||
class TestPagedDSparkWithEncoderReplay(CustomTestCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
override = get_context().override_server_args(
|
||||
enable_encoder_swa_bounded_replay=True,
|
||||
speculative_algorithm="DSPARK",
|
||||
speculative_num_draft_tokens=6,
|
||||
speculative_dspark_block_size=5,
|
||||
page_size=256,
|
||||
max_running_requests=2,
|
||||
chunked_prefill_size=256,
|
||||
)
|
||||
override.install()
|
||||
self.addCleanup(override.restore)
|
||||
|
||||
def make_pool(self, *, draft):
|
||||
return DeepSeekV4TokenToKVPool(
|
||||
max_num_reqs=2,
|
||||
num_req_slots=3,
|
||||
swa_size=1024,
|
||||
c4_size=0,
|
||||
c128_size=0,
|
||||
c4_state_pool_size=0,
|
||||
c128_state_pool_size=0,
|
||||
page_size=256,
|
||||
swa_page_size=256,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
c4_state_dtype=torch.float32,
|
||||
c128_state_dtype=torch.bfloat16,
|
||||
qk_nope_head_dim=448,
|
||||
qk_rope_head_dim=64,
|
||||
indexer_head_dim=128,
|
||||
layer_num=3,
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
compression_ratios=[0, 0, 0],
|
||||
online_mtp_max_draft_tokens=6,
|
||||
full_size=2048,
|
||||
is_draft_worker=draft,
|
||||
)
|
||||
|
||||
def test_target_window_and_draft_paged_storage_share_allocator_mapping(self):
|
||||
target = self.make_pool(draft=False)
|
||||
draft = self.make_pool(draft=True)
|
||||
allocator = SWATokenToKVPoolAllocator(
|
||||
2048, 1024, 256, torch.float8_e4m3fn, "cpu", target, False
|
||||
)
|
||||
draft.register_mapping(allocator.full_to_swa_index_mapping)
|
||||
allocator.full_to_swa_index_mapping[256:512] = torch.arange(768, 1024)
|
||||
self.assertEqual(
|
||||
draft.translate_loc_from_full_to_swa(
|
||||
torch.tensor([256, 300, 511])
|
||||
).tolist(),
|
||||
[768, 812, 1023],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -48,6 +48,7 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase):
|
||||
)
|
||||
|
||||
kvcache = _mock_kvcache(DeepSeekV4TokenToKVPool)
|
||||
kvcache.swa_kv_pool = MagicMock()
|
||||
strategy = _select_strategy(kvcache, {FULL, SWA})
|
||||
self.assertIsInstance(strategy, _DeepSeekV4Strategy)
|
||||
|
||||
@@ -141,6 +142,7 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase):
|
||||
|
||||
for cls in (SWAKVPool, DeepSeekV4TokenToKVPool):
|
||||
kvcache = _mock_kvcache(cls)
|
||||
kvcache.swa_kv_pool = MagicMock()
|
||||
with self.assertRaises(AssertionError) as cm:
|
||||
_select_strategy(kvcache, {FULL})
|
||||
self.assertIn("No matching HiCache strategy", str(cm.exception))
|
||||
|
||||
Reference in New Issue
Block a user