[unified-memory] Hierarchical cache for every unified pool shape (#37507)
This commit is contained in:
Cheng Wan
@@ -459,29 +459,11 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
|
||||
self._run_pause_on_decode_running_batch("retract", weight_update=True)
|
||||
)
|
||||
|
||||
async def _get_decode_num_running_reqs(self, session):
|
||||
"""Query current decode running_batch size from /v1/loads."""
|
||||
async with session.get(
|
||||
self.decode_url + "/v1/loads?include=core",
|
||||
timeout=aiohttp.ClientTimeout(total=5),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
body = await resp.json()
|
||||
return sum(load["num_running_reqs"] for load in body["loads"])
|
||||
|
||||
async def _wait_for_decode_running_batch(self, session, timeout):
|
||||
deadline = asyncio.get_running_loop().time() + timeout
|
||||
while asyncio.get_running_loop().time() < deadline:
|
||||
if await self._get_decode_num_running_reqs(session) > 0:
|
||||
return
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
self.fail("Timed out waiting for decode running_batch to become non-empty")
|
||||
|
||||
async def _run_pause_on_decode_running_batch(self, mode, weight_update=False):
|
||||
num_requests = 2
|
||||
max_new_tokens = 512
|
||||
prompt = "Write a detailed numbered explanation of distributed inference. " * 12
|
||||
decode_started = [asyncio.Event() for _ in range(num_requests)]
|
||||
|
||||
async def _post(session, url, json_data, timeout=30):
|
||||
async with session.post(
|
||||
@@ -493,20 +475,37 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
|
||||
return await resp.json()
|
||||
|
||||
async def _generate(session, request_id):
|
||||
return await _post(
|
||||
session,
|
||||
async with session.post(
|
||||
self.lb_url + "/generate",
|
||||
{
|
||||
json={
|
||||
"text": f"Request {request_id}: {prompt}",
|
||||
"background": True,
|
||||
"stream": True,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"ignore_eos": True,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
},
|
||||
},
|
||||
timeout=180,
|
||||
)
|
||||
timeout=aiohttp.ClientTimeout(total=180),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
response = None
|
||||
async for line in resp.content:
|
||||
line = line.strip()
|
||||
if not line.startswith(b"data: "):
|
||||
continue
|
||||
data = line[len(b"data: ") :]
|
||||
if data == b"[DONE]":
|
||||
break
|
||||
response = json.loads(data)
|
||||
self.assertNotIn("error", response)
|
||||
# Prefill produces the first token. A later token proves this
|
||||
# request has reached running_batch on the decode worker.
|
||||
if response["meta_info"]["completion_tokens"] > 1:
|
||||
decode_started[request_id].set()
|
||||
self.assertIsNotNone(response, "Generation stream returned no output")
|
||||
return response
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [
|
||||
@@ -515,12 +514,17 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
|
||||
decode_paused = False
|
||||
|
||||
try:
|
||||
await self._wait_for_decode_running_batch(session, timeout=30)
|
||||
await asyncio.sleep(0.1)
|
||||
# /v1/loads can still report a previous batch. Wait for every
|
||||
# current request to decode so none can arrive in the prealloc
|
||||
# queue after the pause and prevent the weight-update flush.
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(*(event.wait() for event in decode_started)),
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
any(not task.done() for task in tasks),
|
||||
"All requests finished before decode retract pause was issued.",
|
||||
all(not task.done() for task in tasks),
|
||||
"A request finished before decode retract pause was issued.",
|
||||
)
|
||||
|
||||
await _post(
|
||||
@@ -580,6 +584,9 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
|
||||
for response in responses:
|
||||
self.assertIn("text", response)
|
||||
self.assertGreater(len(response["text"]), 0)
|
||||
self.assertEqual(
|
||||
response["meta_info"]["completion_tokens"], max_new_tokens
|
||||
)
|
||||
|
||||
self.assertGreater(
|
||||
sum(
|
||||
|
||||
@@ -0,0 +1,240 @@
|
||||
"""Compare unified-memory HiCache reloads against a resident-cache reference.
|
||||
|
||||
Evict a target prefix with distinct filler requests, require a host hit on
|
||||
reload, and compare generated text and output logprobs. Both servers use the
|
||||
same unified-memory configuration to keep attention reduction order comparable.
|
||||
Covers GDN, SWA, tri-pool, and MLA layouts.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
|
||||
|
||||
_COMMON_ARGS = [
|
||||
"--trust-remote-code",
|
||||
"--enable-unified-memory",
|
||||
"--enable-cache-report",
|
||||
"--max-running-requests",
|
||||
"1",
|
||||
"--context-length",
|
||||
"4096",
|
||||
]
|
||||
|
||||
# Distinct filler prefixes must evict the target from device memory.
|
||||
_SMALL_POOL = ["--max-total-tokens", "8192"]
|
||||
|
||||
_PREFIX = (
|
||||
"The following is a detailed technical description of a distributed inference "
|
||||
"system with paged attention, radix prefix caching and hierarchical offload. "
|
||||
) * 90
|
||||
_TARGET = _PREFIX + " Question one:"
|
||||
_CONTINUATION = _TARGET + " Explain how it works."
|
||||
|
||||
|
||||
def _generate(base_url, text, max_new_tokens=32, logprobs=True):
|
||||
payload = {
|
||||
"text": text,
|
||||
"sampling_params": {"temperature": 0.0, "max_new_tokens": max_new_tokens},
|
||||
}
|
||||
if logprobs:
|
||||
payload["return_logprob"] = True
|
||||
# Output logprobs suffice; asking for prompt logprobs from zero
|
||||
# caps the reusable prefix at zero and bypasses HiCache entirely.
|
||||
payload["logprob_start_len"] = -1
|
||||
resp = requests.post(f"{base_url}/generate", json=payload, timeout=600)
|
||||
assert resp.status_code == 200, resp.text
|
||||
data = resp.json()
|
||||
lp = (
|
||||
[t[0] for t in data["meta_info"]["output_token_logprobs"]] if logprobs else None
|
||||
)
|
||||
return data["text"], lp, data["meta_info"]
|
||||
|
||||
|
||||
class UnifiedMemoryHiCacheBase(CustomTestCase):
|
||||
"""Compare identical unified-memory configurations with and without HiCache."""
|
||||
|
||||
model: str = ""
|
||||
extra_args: list = []
|
||||
server_env: dict = {}
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls is UnifiedMemoryHiCacheBase:
|
||||
raise unittest.SkipTest("base class")
|
||||
base_args = _COMMON_ARGS + cls.extra_args
|
||||
cls.hicache_url = "http://127.0.0.1:8157"
|
||||
cls.reference_url = "http://127.0.0.1:8158"
|
||||
env = {**os.environ, **cls.server_env} if cls.server_env else None
|
||||
hicache_args = ["--enable-hierarchical-cache"]
|
||||
if "--hicache-size" not in base_args:
|
||||
hicache_args += ["--hicache-ratio", "4"]
|
||||
cls.process_hicache = popen_launch_server(
|
||||
cls.model,
|
||||
cls.hicache_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=base_args + hicache_args,
|
||||
env=env,
|
||||
)
|
||||
cls.addClassCleanup(kill_process_tree, cls.process_hicache.pid)
|
||||
cls.process_reference = popen_launch_server(
|
||||
cls.model,
|
||||
cls.reference_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=base_args + ["--base-gpu-id", "1"],
|
||||
env=env,
|
||||
)
|
||||
cls.addClassCleanup(kill_process_tree, cls.process_reference.pid)
|
||||
|
||||
def _force_host_round_trip(self):
|
||||
"""Evict the target off the device so the next hit must come from L2."""
|
||||
for i in range(8):
|
||||
_generate(
|
||||
self.hicache_url,
|
||||
f"Document {i}. "
|
||||
+ (f"Unique filler {i} about an unrelated subject. " * 300),
|
||||
max_new_tokens=8,
|
||||
logprobs=False,
|
||||
)
|
||||
|
||||
def _flush_both(self):
|
||||
"""Reset cache state to match prefill boundaries and reduction order."""
|
||||
for url in (self.hicache_url, self.reference_url):
|
||||
requests.post(f"{url}/flush_cache", timeout=180)
|
||||
time.sleep(3)
|
||||
|
||||
def test_load_back_matches_no_hicache(self):
|
||||
"""Host reloads preserve generated text and logprobs within tolerance."""
|
||||
self._flush_both()
|
||||
cold_text, cold_lp, _ = _generate(self.hicache_url, _TARGET)
|
||||
ref_cold_text, ref_cold_lp, _ = _generate(self.reference_url, _TARGET)
|
||||
self._force_host_round_trip()
|
||||
# Extend the prefix so both servers compute new KV rows. Repeating it
|
||||
# would let only the resident reference reuse its original final-token KV.
|
||||
warm_text, warm_lp, warm_meta = _generate(self.hicache_url, _CONTINUATION)
|
||||
ref_text, ref_lp, _ = _generate(self.reference_url, _CONTINUATION)
|
||||
|
||||
self.assertGreater(
|
||||
(warm_meta.get("cached_tokens_details") or {}).get("host", 0),
|
||||
0,
|
||||
msg=f"Target did not reload from host: {warm_meta}",
|
||||
)
|
||||
self.assertEqual(cold_text, ref_cold_text)
|
||||
self.assertEqual(warm_text, ref_text)
|
||||
# Match the reference's prefill boundary in each comparison: cold
|
||||
# against cold, and an L2 prefix hit against a resident prefix hit.
|
||||
for label, lp, reference in (
|
||||
("cold", cold_lp, ref_cold_lp),
|
||||
("after-L2-reload", warm_lp, ref_lp),
|
||||
):
|
||||
self.assertEqual(len(lp), len(reference))
|
||||
delta = max(abs(a - b) for a, b in zip(lp, reference))
|
||||
self.assertAlmostEqual(
|
||||
delta,
|
||||
0.0,
|
||||
places=5,
|
||||
msg=f"{label} diverged from the no-HiCache reference by {delta}",
|
||||
)
|
||||
|
||||
def test_server_survives_the_round_trip(self):
|
||||
"""Cache churn must leave both schedulers healthy."""
|
||||
self._force_host_round_trip()
|
||||
for url in (self.hicache_url, self.reference_url):
|
||||
resp = requests.get(f"{url}/health", timeout=30)
|
||||
self.assertEqual(resp.status_code, 200)
|
||||
|
||||
|
||||
class TestUnifiedMemoryHiCacheGDN(UnifiedMemoryHiCacheBase):
|
||||
"""MHA full attention with envelope-strided gated-delta-net state."""
|
||||
|
||||
model = "yujiepan/qwen3.5-tiny-random"
|
||||
extra_args = _SMALL_POOL + [
|
||||
"--linear-attn-backend",
|
||||
"triton",
|
||||
"--mamba-backend",
|
||||
"triton",
|
||||
"--max-mamba-cache-size",
|
||||
"8",
|
||||
"--mem-fraction-static",
|
||||
"0.6",
|
||||
]
|
||||
|
||||
|
||||
class TestUnifiedMemoryHiCacheSWA(UnifiedMemoryHiCacheBase):
|
||||
"""Hybrid SWA reloads bind pages to the full-attention pool's virtual IDs."""
|
||||
|
||||
model = "yujiepan/gemma-4e-tiny-random"
|
||||
extra_args = _SMALL_POOL + [
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
]
|
||||
|
||||
|
||||
class TestUnifiedMemoryHiCacheTriPool(UnifiedMemoryHiCacheBase):
|
||||
"""Full attention, sliding-window attention, and ShortConv state together."""
|
||||
|
||||
# The test revision is the reduced checkpoint used by Inkling CI.
|
||||
model = "thinkingmachines/Inkling"
|
||||
server_env = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}
|
||||
extra_args = _SMALL_POOL + [
|
||||
"--revision",
|
||||
"test",
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--page-size",
|
||||
"128",
|
||||
"--mamba-radix-cache-strategy",
|
||||
"extra_buffer",
|
||||
"--swa-full-tokens-ratio",
|
||||
"0.8",
|
||||
"--max-mamba-cache-size",
|
||||
"8",
|
||||
"--mamba-full-memory-ratio",
|
||||
"0.1",
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
# Bound total host memory across all three component pools.
|
||||
"--hicache-size",
|
||||
"8",
|
||||
]
|
||||
|
||||
|
||||
class TestUnifiedMemoryHiCacheMLA(UnifiedMemoryHiCacheBase):
|
||||
"""MLA full attention with KDA state and MLA-specific transfer pointers."""
|
||||
|
||||
model = "yujiepan/kimi-linear-tiny-random"
|
||||
extra_args = _SMALL_POOL + [
|
||||
"--max-mamba-cache-size",
|
||||
"8",
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
"--linear-attn-backend",
|
||||
"triton",
|
||||
"--mamba-backend",
|
||||
"triton",
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--cuda-graph-backend-decode",
|
||||
"disabled",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -146,15 +146,19 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
|
||||
"""
|
||||
|
||||
class _Peer:
|
||||
def __init__(self, gate):
|
||||
def __init__(self, gate, host_gate=None):
|
||||
self.lazy_compaction = True
|
||||
self._free_phys_pages = [0, 1, 2, 3] # only len() is read
|
||||
self.entry_bytes_per_page = 512
|
||||
self.disagg_move_gate = gate
|
||||
self.host_transfer_move_gate = host_gate
|
||||
|
||||
def _is_frontier_transparent(self):
|
||||
return False
|
||||
|
||||
# Exercise the production predicate when checking each gate.
|
||||
moves_blocked = MultiEndedAllocator.moves_blocked
|
||||
|
||||
class _Owner:
|
||||
"""Stands in for a grow-up END pool: the credit walks the chain from
|
||||
`_growth_side_neighbor()`, so the stub must expose what that walk reads,
|
||||
@@ -167,8 +171,8 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
|
||||
|
||||
_growth_side_neighbor = MultiEndedAllocator._growth_side_neighbor
|
||||
|
||||
def _credit(self, gate):
|
||||
peer = self._Peer(gate)
|
||||
def _credit(self, gate, host_gate=None):
|
||||
peer = self._Peer(gate, host_gate)
|
||||
owner = self._Owner(peer)
|
||||
return MultiEndedAllocator._peer_drainable_hole_bytes(owner)
|
||||
|
||||
@@ -179,6 +183,10 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
|
||||
self.assertEqual(self._credit(gate=lambda: True), 4 * 512)
|
||||
# Gate closed: an urgent flush would move nothing, so credit nothing.
|
||||
self.assertEqual(self._credit(gate=lambda: False), 0)
|
||||
# Either the RDMA gate or the HiCache gate can block compaction.
|
||||
self.assertEqual(self._credit(gate=None, host_gate=lambda: True), 4 * 512)
|
||||
self.assertEqual(self._credit(gate=None, host_gate=lambda: False), 0)
|
||||
self.assertEqual(self._credit(gate=lambda: True, host_gate=lambda: False), 0)
|
||||
|
||||
|
||||
class TestMoveGateRejectsNonPdNode(CustomTestCase):
|
||||
@@ -232,9 +240,7 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
|
||||
# REACHES rather than on what the stub was given.
|
||||
_MEMBER_ATTRS = ("full_attn_allocator", "swa_attn_allocator", "mamba_allocator")
|
||||
|
||||
# The members each composite's gate must reach. The tri-pool row is the one
|
||||
# that matters: it inherits the setter, so an enumeration written inside
|
||||
# that setter would silently leave the third member ungated.
|
||||
# Inherited gate setters must cover every member, including tri-pool Mamba.
|
||||
_EXPECTED_COVERAGE = {
|
||||
"UnifiedMambaTokenToKVPoolAllocator": {
|
||||
"full_attn_allocator",
|
||||
@@ -269,25 +275,26 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
|
||||
def gate() -> bool:
|
||||
return True
|
||||
|
||||
alloc.set_disagg_move_gate(gate)
|
||||
if slot == "disagg_move_gate":
|
||||
alloc.set_disagg_move_gate(gate)
|
||||
else:
|
||||
alloc.set_host_transfer_move_gate(gate)
|
||||
return {
|
||||
attr
|
||||
for attr in self._MEMBER_ATTRS
|
||||
if getattr(getattr(alloc, attr), slot) is gate
|
||||
}
|
||||
|
||||
def test_the_gate_reaches_every_member(self):
|
||||
"""A gate that reaches only some members is not a weaker gate, it is no
|
||||
gate: the ungated end relocates its own pages under the very transfer
|
||||
the gate was installed for.
|
||||
"""
|
||||
def test_every_gate_reaches_every_member(self):
|
||||
"""Both transfer gates must protect every sub-pool from relocation."""
|
||||
for name, expected in self._EXPECTED_COVERAGE.items():
|
||||
with self.subTest(composite=name):
|
||||
self.assertEqual(
|
||||
self._members_reached(name, "disagg_move_gate"),
|
||||
expected,
|
||||
f"{name}.disagg_move_gate does not cover every member",
|
||||
)
|
||||
for slot in ("disagg_move_gate", "host_transfer_move_gate"):
|
||||
with self.subTest(composite=name, slot=slot):
|
||||
self.assertEqual(
|
||||
self._members_reached(name, slot),
|
||||
expected,
|
||||
f"{name}.{slot} does not cover every member",
|
||||
)
|
||||
|
||||
def test_gate_setters_do_not_enumerate_members_themselves(self):
|
||||
"""The structural half of the rule above: a setter that names its
|
||||
@@ -298,10 +305,13 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
|
||||
|
||||
for name in self._EXPECTED_COVERAGE:
|
||||
cls = self._allocator_class(name)
|
||||
with self.subTest(composite=name):
|
||||
body = inspect.getsource(cls.set_disagg_move_gate)
|
||||
self.assertIn("install_move_gate", body)
|
||||
self.assertNotIn("_move_gate = ", body)
|
||||
for setter in ("set_disagg_move_gate", "set_host_transfer_move_gate"):
|
||||
if setter not in vars(cls):
|
||||
continue # inherited, and the inherited one is checked above
|
||||
with self.subTest(composite=name, setter=setter):
|
||||
body = inspect.getsource(getattr(cls, setter))
|
||||
self.assertIn("install_move_gate", body)
|
||||
self.assertNotIn("_move_gate = ", body)
|
||||
|
||||
def test_swa_composite_translates_the_swa_side_separately(self):
|
||||
"""The SWA sub-pool runs its OWN compaction, so a full-side physical id
|
||||
|
||||
@@ -31,6 +31,8 @@ WIDENED_PAGE = PHYSICAL_PAGE * DCP_SIZE
|
||||
def _fake_mla_device_pool(size: int = 1024) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
size=size,
|
||||
# Match KVCache's default: this static pool's size already counts tokens.
|
||||
host_capacity_tokens=None,
|
||||
store_dtype=torch.float16,
|
||||
kv_lora_rank=8,
|
||||
qk_rope_head_dim=4,
|
||||
|
||||
@@ -27,7 +27,7 @@ import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
from test_multi_ended_allocator import _FakeUnifiedSWAKVPool
|
||||
from test_multi_ended_allocator import _FakeKVCache, _FakeUnifiedSWAKVPool
|
||||
|
||||
from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
|
||||
UnifiedSWATokenToKVPoolAllocator,
|
||||
@@ -673,7 +673,10 @@ class TestWriteLoc(CustomTestCase):
|
||||
)
|
||||
allocator = UnifiedMambaTokenToKVPoolAllocator(
|
||||
unified_buffer=pool,
|
||||
kvcache=SimpleNamespace(full_kv_pool=None, mamba_pool=None),
|
||||
kvcache=SimpleNamespace(
|
||||
full_kv_pool=_FakeKVCache(pool.max_slots("full")),
|
||||
mamba_pool=_FakeKVCache(pool.max_slots("mamba")),
|
||||
),
|
||||
device="cpu",
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,345 @@
|
||||
"""Host reloads must preserve ID domains, allocation ownership, and stream order."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.dcp.layout import maybe_dcp_kernel_indices
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
||||
HybridCacheController,
|
||||
)
|
||||
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import _split_hicache_size
|
||||
from sglang.srt.mem_cache.l2_transfer import L2Transfer, L2TransferEngine
|
||||
from sglang.srt.mem_cache.pool_host.group import PoolEntry
|
||||
from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
|
||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestHiCacheIndexDomains(unittest.TestCase):
|
||||
def test_dcp_translation_uses_local_virtual_ids(self):
|
||||
# Two non-adjacent pages, relocated to different physical pages.
|
||||
page_size, dcp_size = 4, 2
|
||||
logical = torch.cat((torch.arange(8, 16), torch.arange(24, 32)))
|
||||
v2p = torch.tensor([0, 5, 4, 2])
|
||||
|
||||
def translate(ids):
|
||||
return v2p[ids // page_size] * (page_size * 3) + ids % page_size
|
||||
|
||||
transfer = L2Transfer(
|
||||
SimpleNamespace(dcp_size=dcp_size),
|
||||
SimpleNamespace(host_transfer_translate=translate),
|
||||
logical.clone(),
|
||||
logical,
|
||||
)
|
||||
resolved = L2TransferEngine._resolve_device_indices(transfer)
|
||||
for rank in range(dcp_size):
|
||||
expected = translate(maybe_dcp_kernel_indices(logical, dcp_size, rank))
|
||||
torch.testing.assert_close(
|
||||
maybe_dcp_kernel_indices(resolved, dcp_size, rank), expected
|
||||
)
|
||||
|
||||
def test_fixed_size_uses_host_capacity_for_shared_buffers(self):
|
||||
pools = [
|
||||
SimpleNamespace(host_capacity_bytes=n, get_kv_size_bytes=lambda: (0, 0))
|
||||
for n in (600, 300, 100)
|
||||
]
|
||||
self.assertEqual(_split_hicache_size(10, tuple(pools)), (6, 3, 1))
|
||||
|
||||
|
||||
class TestHostGateCapacity(unittest.TestCase):
|
||||
def test_schedulable_memo_tracks_gate_without_allocator_mutation(self):
|
||||
from test_unified_capacity_memo import _build
|
||||
|
||||
inst, allocator, kvcache = _build(lazy=True)
|
||||
ids = inst._alloc(allocator, kvcache, 8)
|
||||
allocator.free_swa(ids[2:6])
|
||||
full = allocator.full_attn_allocator
|
||||
state = {"open": True}
|
||||
allocator.set_host_transfer_move_gate(lambda: state["open"])
|
||||
epoch = full._chain_capacity_epoch()
|
||||
before = full.schedulable_available_size()
|
||||
state["open"] = False
|
||||
self.assertEqual(full._chain_capacity_epoch(), epoch)
|
||||
self.assertEqual(full.schedulable_available_size(), full.available_size())
|
||||
self.assertGreater(before, full.schedulable_available_size())
|
||||
self.assertFalse(allocator._compaction_allowed())
|
||||
self.assertEqual(allocator.verify_byte_accounting(), [])
|
||||
state["open"] = True
|
||||
self.assertEqual(full.schedulable_available_size(), before)
|
||||
|
||||
|
||||
class TestSwaLoadAllocation(unittest.TestCase):
|
||||
def _controller(self, bind, free=None, evict=None):
|
||||
controller = object.__new__(HybridCacheController)
|
||||
entry = PoolEntry(
|
||||
name=PoolName.SWA,
|
||||
host_pool=SimpleNamespace(),
|
||||
device_pool=SimpleNamespace(),
|
||||
layer_mapper=lambda i: i,
|
||||
device_indices_from_anchor_fn=bind,
|
||||
device_free_fn=free or Mock(),
|
||||
device_evict_fn=evict,
|
||||
)
|
||||
controller.mem_pool_host = SimpleNamespace(entry_map={PoolName.SWA: entry})
|
||||
return controller
|
||||
|
||||
def _transfer(self, parts, count):
|
||||
transfer = PoolTransfer(name=PoolName.SWA, host_indices=torch.arange(count))
|
||||
transfer.anchor_index_parts = parts
|
||||
return transfer
|
||||
|
||||
def test_swa_only_load_uses_resident_full_ids(self):
|
||||
bind = Mock(side_effect=lambda x: x + 100)
|
||||
controller = self._controller(bind)
|
||||
transfer = self._transfer([torch.tensor([13, 14])], 2)
|
||||
result = controller._resolve_device_transfers(
|
||||
[transfer], torch.empty(0, dtype=torch.int64)
|
||||
)
|
||||
self.assertIsNotNone(result)
|
||||
torch.testing.assert_close(transfer.device_indices, torch.tensor([113, 114]))
|
||||
|
||||
def test_mixed_load_skips_resident_swa_nodes(self):
|
||||
bind = Mock(side_effect=lambda x: x + 100)
|
||||
controller = self._controller(bind)
|
||||
transfer = self._transfer([torch.tensor([13, 14]), slice(2, 4)], 4)
|
||||
controller._resolve_device_transfers(
|
||||
[transfer], torch.tensor([20, 21, 22, 23, 24, 25])
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
transfer.device_indices, torch.tensor([113, 114, 122, 123])
|
||||
)
|
||||
|
||||
def test_binding_retries_after_eviction(self):
|
||||
bind = Mock(side_effect=[None, torch.tensor([41, 42])])
|
||||
evict = Mock()
|
||||
controller = self._controller(bind, evict=evict)
|
||||
transfer = self._transfer([slice(0, 2)], 2)
|
||||
self.assertIsNotNone(
|
||||
controller._resolve_device_transfers([transfer], torch.tensor([11, 12]))
|
||||
)
|
||||
evict.assert_called_once_with(2)
|
||||
self.assertEqual(bind.call_count, 2)
|
||||
|
||||
def test_rollback_releases_swa_binding_by_virtual_ids(self):
|
||||
free = Mock()
|
||||
controller = self._controller(lambda x: x + 100, free=free)
|
||||
transfer = self._transfer([torch.tensor([13, 14])], 2)
|
||||
# A missing sidecar source fails after SWA has bound its pages.
|
||||
sidecar = PoolTransfer(name=PoolName.MAMBA, indices_from_pool=PoolName.INDEXER)
|
||||
result = controller._resolve_device_transfers(
|
||||
[transfer, sidecar], torch.empty(0, dtype=torch.int64)
|
||||
)
|
||||
self.assertIsNone(result)
|
||||
torch.testing.assert_close(free.call_args.args[0], torch.tensor([13, 14]))
|
||||
self.assertIsNone(transfer.device_indices)
|
||||
|
||||
def test_independent_allocations_precede_kernel_id_resolution(self):
|
||||
events = []
|
||||
controller = self._controller(lambda x: events.append("bind") or x + 100)
|
||||
controller.mem_pool_host.entry_map[PoolName.MAMBA] = PoolEntry(
|
||||
name=PoolName.MAMBA,
|
||||
host_pool=SimpleNamespace(),
|
||||
device_pool=SimpleNamespace(),
|
||||
layer_mapper=lambda i: i,
|
||||
device_alloc_fn=lambda n: events.append("mamba") or torch.arange(n),
|
||||
device_free_fn=Mock(),
|
||||
)
|
||||
swa = self._transfer([slice(0, 2)], 2)
|
||||
mamba = PoolTransfer(name=PoolName.MAMBA, host_indices=torch.arange(1))
|
||||
self.assertIsNotNone(
|
||||
controller._resolve_device_transfers([swa, mamba], torch.tensor([10, 11]))
|
||||
)
|
||||
self.assertEqual(events, ["mamba", "bind"])
|
||||
|
||||
def test_tree_spec_preserves_node_correspondence(self):
|
||||
kv = PoolTransfer(
|
||||
name=PoolName.KV, host_indices=torch.arange(4), nodes_to_load=[2, 3]
|
||||
)
|
||||
swa = PoolTransfer(
|
||||
name=PoolName.SWA, host_indices=torch.arange(4), nodes_to_load=[1, 3]
|
||||
)
|
||||
nodes = {
|
||||
i: SimpleNamespace(
|
||||
id=i,
|
||||
key=[i, i],
|
||||
load_back_pending_id=None,
|
||||
component_data={
|
||||
ComponentType.FULL: SimpleNamespace(
|
||||
value=torch.tensor([10, 11]) if i == 1 else None
|
||||
)
|
||||
},
|
||||
)
|
||||
for i in (1, 2, 3)
|
||||
}
|
||||
full_component = SimpleNamespace(
|
||||
component_type=ComponentType.FULL,
|
||||
build_hicache_transfers=lambda *a, **k: [kv],
|
||||
)
|
||||
swa_component = SimpleNamespace(
|
||||
component_type=ComponentType.SWA,
|
||||
build_hicache_transfers=lambda *a, **k: [swa],
|
||||
)
|
||||
core = SimpleNamespace(
|
||||
node_by_id=nodes.__getitem__,
|
||||
components=[full_component, swa_component],
|
||||
components_by_type={
|
||||
ComponentType.FULL: full_component,
|
||||
ComponentType.SWA: swa_component,
|
||||
},
|
||||
)
|
||||
UnifiedTreeCore.build_load_back_spec(core, 3)
|
||||
parts = swa.anchor_index_parts
|
||||
torch.testing.assert_close(parts[0], torch.tensor([10, 11]))
|
||||
self.assertEqual(parts[1], slice(2, 4))
|
||||
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
|
||||
class TestTransferStreamOrdering(unittest.TestCase):
|
||||
def test_load_translation_follows_supplied_start_event(self):
|
||||
# The start event precedes translation; transfer must wait for both.
|
||||
engine = L2TransferEngine("kernel")
|
||||
ids = torch.tensor([2, 4, 7], device="cuda")
|
||||
output = torch.full_like(ids, -1)
|
||||
resolved = torch.full_like(ids, -2)
|
||||
torch.cuda.synchronize()
|
||||
start = torch.cuda.Event()
|
||||
start.record()
|
||||
|
||||
def translate(indices):
|
||||
self.assertEqual(torch.cuda.current_stream(), engine.host_to_device_stream)
|
||||
torch.cuda._sleep(20_000_000)
|
||||
resolved.copy_(indices + 100)
|
||||
return resolved
|
||||
|
||||
host = SimpleNamespace(
|
||||
layer_num=1,
|
||||
load_to_device_per_layer=lambda pool, h, d, layer, backend, **kw: (
|
||||
output.copy_(d)
|
||||
),
|
||||
)
|
||||
transfer = L2Transfer(
|
||||
host,
|
||||
SimpleNamespace(host_transfer_translate=translate),
|
||||
torch.arange(3),
|
||||
ids,
|
||||
)
|
||||
completion = engine.submit_host_to_device(
|
||||
[transfer], transfer_layer_id_max=1, start_event=start
|
||||
)
|
||||
completion.finish_event.synchronize()
|
||||
torch.testing.assert_close(output.cpu(), torch.tensor([102, 104, 107]))
|
||||
|
||||
|
||||
class TestSwaBackupAfterCompaction(unittest.TestCase):
|
||||
def test_backup_resolves_current_binding(self):
|
||||
from sglang.srt.mem_cache.unified_cache.components.base import (
|
||||
CacheTransferPhase,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_cache.components.swa import SWAComponent
|
||||
|
||||
node = SimpleNamespace(
|
||||
component_data={
|
||||
ComponentType.FULL: SimpleNamespace(value=torch.tensor([3, 7])),
|
||||
ComponentType.SWA: SimpleNamespace(value=torch.tensor([103, 107])),
|
||||
},
|
||||
id=1,
|
||||
)
|
||||
component = SimpleNamespace(
|
||||
component_type=ComponentType.SWA,
|
||||
tree_core=SimpleNamespace(has_swa_host_pool=True),
|
||||
_collect_unbacked_swa_nodes=lambda n: [n],
|
||||
_unified_allocator=lambda: object(),
|
||||
_translate_full_to_swa=lambda x: x + 200,
|
||||
)
|
||||
transfers = SWAComponent.build_hicache_transfers(
|
||||
component, node, CacheTransferPhase.BACKUP_HOST
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
transfers[0].device_indices, torch.tensor([203, 207])
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
|
||||
class TestDirectBackendTranslation(unittest.TestCase):
|
||||
def test_cpu_indices_translate_on_device_and_return_to_cpu(self):
|
||||
mapping = torch.tensor([0, 9, 6], device="cuda")
|
||||
|
||||
def translate(ids):
|
||||
self.assertTrue(ids.is_cuda)
|
||||
return mapping[ids]
|
||||
|
||||
transfer = L2Transfer(
|
||||
SimpleNamespace(),
|
||||
SimpleNamespace(device="cuda", host_transfer_translate=translate),
|
||||
torch.tensor([0, 1]),
|
||||
torch.tensor([2, 1]),
|
||||
)
|
||||
result = L2TransferEngine._resolve_device_indices(transfer)
|
||||
self.assertEqual(result.device.type, "cpu")
|
||||
torch.testing.assert_close(result, torch.tensor([6, 9]))
|
||||
|
||||
|
||||
class TestTriPoolAssembly(unittest.TestCase):
|
||||
def test_swa_allocation_and_rollback_match_pool_id_ownership(self):
|
||||
from sglang.srt.mem_cache.hybrid_cache import hybrid_pool_assembler as assembler
|
||||
|
||||
for unified in (False, True):
|
||||
with self.subTest(unified=unified):
|
||||
swa_allocator = SimpleNamespace(alloc=Mock(), free=Mock())
|
||||
composite = SimpleNamespace(
|
||||
swa_attn_allocator=swa_allocator,
|
||||
bind_swa_for_loaded_rows=Mock(),
|
||||
free_swa=Mock(),
|
||||
)
|
||||
params = SimpleNamespace(
|
||||
token_to_kv_pool_allocator=composite,
|
||||
req_to_token_pool=SimpleNamespace(
|
||||
mamba_allocator=SimpleNamespace(alloc=Mock(), free=Mock())
|
||||
),
|
||||
)
|
||||
memory = MagicMock(enable_unified_memory=unified, hicache_size=0)
|
||||
host = MagicMock()
|
||||
with (
|
||||
patch.object(assembler, "get_memory", return_value=memory),
|
||||
patch.object(
|
||||
assembler, "_get_allocator_type", return_value="default"
|
||||
),
|
||||
patch.object(assembler, "build_kv_host_pool", return_value=host),
|
||||
patch.object(assembler, "MambaPoolHost", return_value=host),
|
||||
patch.object(assembler, "HybridCacheController"),
|
||||
):
|
||||
group, _ = assembler.build_hybrid_mamba_swa_stack(
|
||||
params=params,
|
||||
full_kv_pool=object(),
|
||||
swa_kv_pool=object(),
|
||||
mamba_pool=object(),
|
||||
full_layer_mapping={0: 0},
|
||||
swa_layer_mapping={1: 0},
|
||||
mamba_layer_mapping={2: 0},
|
||||
page_size=1,
|
||||
tp_group=None,
|
||||
load_cache_event=None,
|
||||
storage_backend=None,
|
||||
)
|
||||
entry = group.entry_map[PoolName.SWA]
|
||||
if unified:
|
||||
self.assertIs(
|
||||
entry.device_indices_from_anchor_fn,
|
||||
composite.bind_swa_for_loaded_rows,
|
||||
)
|
||||
self.assertIs(entry.device_free_fn, composite.free_swa)
|
||||
self.assertIsNone(entry.device_alloc_fn)
|
||||
else:
|
||||
self.assertIs(entry.device_alloc_fn, swa_allocator.alloc)
|
||||
self.assertIs(entry.device_free_fn, swa_allocator.free)
|
||||
self.assertIsNone(entry.device_indices_from_anchor_fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Cover strided Mamba state staging and shared allocator wiring for HiCache."""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=20, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestStridedStateDetection(CustomTestCase):
|
||||
def test_contiguous_slots_are_not_strided(self):
|
||||
for shape in ((8, 4), (8, 4, 3), (1, 5)):
|
||||
with self.subTest(shape=shape):
|
||||
self.assertFalse(
|
||||
MambaPoolHost._slots_are_strided(torch.zeros(shape)),
|
||||
"a contiguous slot array must take the direct path",
|
||||
)
|
||||
|
||||
def test_envelope_strided_slots_are_detected(self):
|
||||
num_slots, per_slot, envelope = 6, 4, 10
|
||||
raw = torch.zeros(num_slots * envelope)
|
||||
view = torch.as_strided(raw, size=(num_slots, per_slot), stride=(envelope, 1))
|
||||
self.assertTrue(MambaPoolHost._slots_are_strided(view))
|
||||
|
||||
def test_empty_tensor_is_not_strided(self):
|
||||
self.assertFalse(MambaPoolHost._slots_are_strided(torch.zeros((0, 4))))
|
||||
|
||||
def test_staging_round_trip_preserves_slot_contents(self):
|
||||
"""Gather and scatter preserve selected slots without changing their neighbors."""
|
||||
num_slots, per_slot, envelope = 6, 4, 10
|
||||
raw = torch.arange(num_slots * envelope, dtype=torch.float32)
|
||||
view = torch.as_strided(raw, size=(num_slots, per_slot), stride=(envelope, 1))
|
||||
indices = torch.tensor([4, 1, 3])
|
||||
|
||||
staged = view.index_select(0, indices)
|
||||
self.assertTrue(staged.is_contiguous())
|
||||
for row, slot in enumerate(indices.tolist()):
|
||||
self.assertTrue(torch.equal(staged[row], view[slot]))
|
||||
|
||||
dst = torch.zeros_like(raw)
|
||||
dst_view = torch.as_strided(
|
||||
dst, size=(num_slots, per_slot), stride=(envelope, 1)
|
||||
)
|
||||
dst_view.index_copy_(0, indices, staged)
|
||||
for slot in indices.tolist():
|
||||
self.assertTrue(torch.equal(dst_view[slot], view[slot]))
|
||||
# A partial transfer must leave unselected slots untouched.
|
||||
for slot in set(range(num_slots)) - set(indices.tolist()):
|
||||
self.assertTrue(torch.all(dst_view[slot] == 0))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
class TestMambaSlotWiringIsShared(CustomTestCase):
|
||||
"""Both Mamba factories must install slot allocation and transfer translation.
|
||||
|
||||
Inspect the AST so this check does not need GPU-backed pools.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _assignments_to(attr: str):
|
||||
"""Find functions assigning the attribute, excluding None initialization."""
|
||||
import ast
|
||||
import inspect
|
||||
|
||||
from sglang.srt.mem_cache import unified_memory_pool
|
||||
|
||||
tree = ast.parse(inspect.getsource(unified_memory_pool))
|
||||
found = []
|
||||
for fn in ast.walk(tree):
|
||||
if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
for node in ast.walk(fn):
|
||||
if not isinstance(node, ast.Assign):
|
||||
continue
|
||||
if isinstance(node.value, ast.Constant) and node.value.value is None:
|
||||
continue
|
||||
for tgt in node.targets:
|
||||
if isinstance(tgt, ast.Attribute) and tgt.attr == attr:
|
||||
found.append(fn.name)
|
||||
return found
|
||||
|
||||
def test_only_the_shared_hook_wraps_the_slot_allocator(self):
|
||||
self.assertEqual(
|
||||
sorted(set(self._assignments_to("mamba_allocator"))),
|
||||
["_wire_mamba_slot_allocator"],
|
||||
"a factory wraps the mamba end itself; it will miss "
|
||||
"host_transfer_translate exactly as the tri-pool factory did",
|
||||
)
|
||||
|
||||
def test_every_unified_mamba_factory_calls_the_shared_hook(self):
|
||||
import ast
|
||||
import inspect
|
||||
|
||||
from sglang.srt.mem_cache import unified_memory_pool
|
||||
|
||||
tree = ast.parse(inspect.getsource(unified_memory_pool))
|
||||
# Factories whose composite holds a mamba end.
|
||||
expected = {"init_unified_mamba_pools", "init_unified_mamba_swa_pools"}
|
||||
callers = {
|
||||
fn.name
|
||||
for fn in ast.walk(tree)
|
||||
if isinstance(fn, ast.FunctionDef)
|
||||
for node in ast.walk(fn)
|
||||
if isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Name)
|
||||
and node.func.id == "_wire_mamba_slot_allocator"
|
||||
}
|
||||
self.assertEqual(expected, expected & callers, f"missing: {expected - callers}")
|
||||
Reference in New Issue
Block a user