[Feat] DCP + HiCache L2 Support (ported from kimi-k3) (#33112)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuwei An
co-authored by
Claude Opus 5
Claude Opus 5
parent
f8e62a9224
commit
1a3bea77f2
@@ -0,0 +1,112 @@
|
||||
"""HiCache L2 under decode context parallelism (DCP) + UnifiedRadixCache.
|
||||
|
||||
Under DCP the radix layer allocates widened logical indices while each rank's
|
||||
buffers hold only its 1/dcp_size shard, so a missing translation makes cache
|
||||
hits return another rank's KV. The KL cases catch that as a large divergence.
|
||||
|
||||
Blackwell-only: the MLA DCP decode path needs ``tokenspeed_mla`` (SM100/12x).
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
|
||||
from sglang.test.kl_multiturn_utils import (
|
||||
get_input_ids,
|
||||
make_mamba_decode_assert,
|
||||
make_mamba_prefill_assert,
|
||||
)
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=1500, stage="extra-b", runner_config="4-gpu-b200")
|
||||
|
||||
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
|
||||
DCP_SIZE = 4
|
||||
PAGE_SIZE = 64
|
||||
WIDENED_PAGE = PAGE_SIZE * DCP_SIZE
|
||||
MAX_MAMBA_CACHE_SIZE = 256
|
||||
# Bound the host pools directly. Sizing them off the device pool would need
|
||||
# --max-total-tokens, which triggers out-of-range KV writes under DCP.
|
||||
HICACHE_SIZE_GB = 10
|
||||
|
||||
|
||||
class TestUnifiedKimiLinearDcpHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
|
||||
"""Kimi Linear + DCP4 + HiCache L2 + UnifiedRadixCache."""
|
||||
|
||||
kl_threshold = 0.01
|
||||
gsm8k_threshold = 0.85
|
||||
mmlu_threshold = 0.4
|
||||
prefill_cache_assert = staticmethod(
|
||||
make_mamba_prefill_assert(chunk_size=WIDENED_PAGE)
|
||||
)
|
||||
decode_cache_assert = staticmethod(
|
||||
make_mamba_decode_assert(track_interval=WIDENED_PAGE)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = KIMI_LINEAR_MODEL
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp-size",
|
||||
"4",
|
||||
"--dcp-size",
|
||||
str(DCP_SIZE),
|
||||
"--page-size",
|
||||
str(PAGE_SIZE),
|
||||
"--attention-backend",
|
||||
"tokenspeed_mla",
|
||||
"--kv-cache-dtype",
|
||||
"fp8_e4m3",
|
||||
"--dcp-comm-backend",
|
||||
"a2a",
|
||||
"--dcp-replicate-q-proj",
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--random-seed",
|
||||
"0",
|
||||
"--cuda-graph-max-bs-decode",
|
||||
"64",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
"--mem-fraction-static",
|
||||
"0.80",
|
||||
"--enable-hierarchical-cache",
|
||||
"--hicache-size",
|
||||
str(HICACHE_SIZE_GB),
|
||||
"--hicache-write-policy",
|
||||
"write_through",
|
||||
"--max-running-requests",
|
||||
"64",
|
||||
"--max-mamba-cache-size",
|
||||
str(MAX_MAMBA_CACHE_SIZE),
|
||||
"--enable-metrics",
|
||||
],
|
||||
env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
|
||||
)
|
||||
cls.input_ids = get_input_ids(cls.model, num_samples=18, trust_remote_code=True)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.process.terminate()
|
||||
try:
|
||||
cls.process.wait(timeout=60)
|
||||
except subprocess.TimeoutExpired:
|
||||
pass
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,210 @@
|
||||
"""HiCache under decode context parallelism (DCP): host-pool index math.
|
||||
|
||||
Under DCP the radix/controller layer works in a widened logical index space
|
||||
(page_size * dcp_size wide pages, dcp_size * physical capacity), while each
|
||||
rank's device and host buffers only materialize the owned 1/dcp_size token
|
||||
shard (owner rule: index % dcp_size == dcp_rank, physical row = index //
|
||||
dcp_size — the same rule the device-side KV write and page-table kernels
|
||||
use). These tests cover the translation helper, the logical/physical host
|
||||
pool sizing, and that the transfer entry points hand *physical* rows to the
|
||||
kernels.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
DCP_SIZE = 8
|
||||
PHYSICAL_PAGE = 64
|
||||
WIDENED_PAGE = PHYSICAL_PAGE * DCP_SIZE
|
||||
|
||||
|
||||
def _fake_mla_device_pool(size: int = 1024) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
size=size,
|
||||
store_dtype=torch.float16,
|
||||
kv_lora_rank=8,
|
||||
qk_rope_head_dim=4,
|
||||
layer_num=2,
|
||||
start_layer=0,
|
||||
end_layer=1,
|
||||
device="cpu",
|
||||
layers_to_capture=None,
|
||||
layer_shard_enabled=False,
|
||||
)
|
||||
|
||||
|
||||
def _make_host_pool(dcp_rank: int, device_size: int = 1024) -> MLATokenToKVPoolHost:
|
||||
return MLATokenToKVPoolHost(
|
||||
_fake_mla_device_pool(device_size),
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=WIDENED_PAGE,
|
||||
layout="layer_first",
|
||||
pin_memory=False,
|
||||
device="cpu",
|
||||
dcp_size=DCP_SIZE,
|
||||
dcp_rank=dcp_rank,
|
||||
)
|
||||
|
||||
|
||||
class TestDcpKernelIndices(CustomTestCase):
|
||||
def _bare_pool(self, dcp_size: int, dcp_rank: int) -> MLATokenToKVPoolHost:
|
||||
pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
|
||||
pool.dcp_size = dcp_size
|
||||
pool.dcp_rank = dcp_rank
|
||||
return pool
|
||||
|
||||
def test_identity_without_dcp(self):
|
||||
pool = self._bare_pool(1, 0)
|
||||
indices = torch.arange(37)
|
||||
self.assertIs(pool.dcp_kernel_indices(indices), indices)
|
||||
|
||||
def test_aligned_page_translates_to_full_physical_page(self):
|
||||
# One widened page starting at logical 512 covers physical rows
|
||||
# 64..127 on every rank.
|
||||
indices = torch.arange(WIDENED_PAGE, 2 * WIDENED_PAGE)
|
||||
for rank in range(DCP_SIZE):
|
||||
pool = self._bare_pool(DCP_SIZE, rank)
|
||||
out = pool.dcp_kernel_indices(indices)
|
||||
torch.testing.assert_close(
|
||||
out, torch.arange(PHYSICAL_PAGE, 2 * PHYSICAL_PAGE)
|
||||
)
|
||||
|
||||
def test_matches_owner_rule_on_merged_unordered_pages(self):
|
||||
# Concatenation of non-adjacent widened pages in arbitrary order, as
|
||||
# produced by merged CacheOperations after allocator churn.
|
||||
pages = [3, 0, 5]
|
||||
indices = torch.cat(
|
||||
[torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE) for p in pages]
|
||||
)
|
||||
for rank in range(DCP_SIZE):
|
||||
pool = self._bare_pool(DCP_SIZE, rank)
|
||||
out = pool.dcp_kernel_indices(indices)
|
||||
expected = (
|
||||
indices[indices % DCP_SIZE == rank] // DCP_SIZE
|
||||
) # owner rule, same as filter_dcp_local_kv_indices
|
||||
torch.testing.assert_close(out, expected)
|
||||
self.assertEqual(out.numel() * DCP_SIZE, indices.numel())
|
||||
|
||||
def test_ragged_run_is_rejected(self):
|
||||
pool = self._bare_pool(DCP_SIZE, 0)
|
||||
with self.assertRaises(AssertionError):
|
||||
pool.dcp_kernel_indices(torch.arange(WIDENED_PAGE + 1))
|
||||
|
||||
def test_positional_residue_pairing_survives_host_sort(self):
|
||||
# move_indices (direct/layer_first) sorts host indices and permutes
|
||||
# device indices to match. Independent residue filtering of both
|
||||
# tensors must keep the same token positions on every rank.
|
||||
g = torch.Generator().manual_seed(0)
|
||||
host_pages = [7, 2]
|
||||
device_pages = [1, 4]
|
||||
host = torch.cat(
|
||||
[torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE) for p in host_pages]
|
||||
)
|
||||
device = torch.cat(
|
||||
[
|
||||
torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE)
|
||||
for p in device_pages
|
||||
]
|
||||
)
|
||||
# token identity: position i pairs host[i] <-> device[i]
|
||||
perm = torch.randperm(host.numel(), generator=g)
|
||||
# sort host as move_indices does, permuting device alongside
|
||||
host_sorted, order = host[perm].sort()
|
||||
device_matched = device[perm][order]
|
||||
for rank in range(DCP_SIZE):
|
||||
pool = self._bare_pool(DCP_SIZE, rank)
|
||||
host_mask = host_sorted % DCP_SIZE == rank
|
||||
device_mask = device_matched % DCP_SIZE == rank
|
||||
# same positions selected on both sides -> pairing preserved
|
||||
torch.testing.assert_close(host_mask, device_mask)
|
||||
self.assertEqual(
|
||||
pool.dcp_kernel_indices(host_sorted).numel(),
|
||||
host.numel() // DCP_SIZE,
|
||||
)
|
||||
|
||||
|
||||
class TestHostPoolSizingUnderDcp(CustomTestCase):
|
||||
def test_logical_and_physical_sizing(self):
|
||||
pool = _make_host_pool(dcp_rank=3)
|
||||
# kernel-facing page is physical
|
||||
self.assertEqual(pool.page_size, PHYSICAL_PAGE)
|
||||
self.assertEqual(pool.logical_page_size, WIDENED_PAGE)
|
||||
# physical rows = ratio * device physical size, page aligned
|
||||
self.assertEqual(pool.size, pool.page_num * PHYSICAL_PAGE)
|
||||
self.assertEqual(pool.logical_size, pool.size * DCP_SIZE)
|
||||
# buffers materialize physical rows only
|
||||
self.assertEqual(pool.kv_buffer.shape[1], pool.size)
|
||||
# allocator surface is logical
|
||||
self.assertEqual(pool.free_slots.numel(), pool.logical_size)
|
||||
self.assertEqual(pool.mem_state.numel(), pool.logical_size)
|
||||
|
||||
def test_alloc_is_widened_page_granular(self):
|
||||
pool = _make_host_pool(dcp_rank=0)
|
||||
out = pool.alloc(WIDENED_PAGE)
|
||||
self.assertEqual(out.numel(), WIDENED_PAGE)
|
||||
with self.assertRaises(AssertionError):
|
||||
pool.alloc(PHYSICAL_PAGE) # not a multiple of the widened page
|
||||
|
||||
def test_non_dcp_pool_unchanged(self):
|
||||
pool = MLATokenToKVPoolHost(
|
||||
_fake_mla_device_pool(),
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=PHYSICAL_PAGE,
|
||||
layout="layer_first",
|
||||
pin_memory=False,
|
||||
device="cpu",
|
||||
)
|
||||
self.assertEqual(pool.page_size, PHYSICAL_PAGE)
|
||||
self.assertEqual(pool.logical_size, pool.size)
|
||||
self.assertEqual(pool.logical_page_size, PHYSICAL_PAGE)
|
||||
|
||||
|
||||
class TestTransferEntryPointsTranslate(CustomTestCase):
|
||||
def _run_backup(self, pool, host_indices, device_indices):
|
||||
device_pool = SimpleNamespace(
|
||||
data_ptrs=torch.zeros(2, dtype=torch.uint64),
|
||||
kv_buffer=[torch.zeros(1)] * 2,
|
||||
)
|
||||
# create=True: mla.py imports the kernel only under `if _is_cuda or
|
||||
# _is_hip`, so the name is absent on the CPU runner this test targets.
|
||||
with mock.patch(
|
||||
"sglang.srt.mem_cache.pool_host.mla.transfer_kv_all_layer_mla",
|
||||
create=True,
|
||||
) as kernel:
|
||||
pool.can_use_jit = False
|
||||
pool.can_use_write_back_jit = False
|
||||
with mock.patch.object(
|
||||
MLATokenToKVPoolHost, "_is_device_layer_sharded", return_value=False
|
||||
):
|
||||
pool.backup_from_device_all_layer(
|
||||
device_pool, host_indices, device_indices, io_backend="kernel"
|
||||
)
|
||||
return kernel.call_args.kwargs
|
||||
|
||||
def test_backup_receives_physical_rows(self):
|
||||
pool = _make_host_pool(dcp_rank=5)
|
||||
logical = torch.arange(2 * WIDENED_PAGE)
|
||||
kwargs = self._run_backup(pool, logical, logical.clone())
|
||||
expected = torch.arange(2 * PHYSICAL_PAGE)
|
||||
torch.testing.assert_close(kwargs["src_indices"], expected)
|
||||
torch.testing.assert_close(kwargs["dst_indices"], expected)
|
||||
|
||||
def test_l3_data_page_is_guarded(self):
|
||||
pool = _make_host_pool(dcp_rank=0)
|
||||
with self.assertRaises(AssertionError):
|
||||
pool.get_data_page(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user