fix(hicache/umbp): support DeepSeek-V4 hybrid HostPoolGroup (multi-po… (#30762)
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
AMD-yanfeiwang
co-authored by
Zhangheng
Zhangheng
parent
b7f87a2513
commit
a34f81251f
@@ -0,0 +1,214 @@
|
||||
"""E2E test for DeepSeek-V4 HiCache storage with the UMBP backend.
|
||||
|
||||
The first request writes the hybrid HostPoolGroup side pools to UMBP. After
|
||||
flushing the device and host radix caches, the same prompt must be restored
|
||||
from UMBP and report a storage-tier cache hit.
|
||||
|
||||
Usage:
|
||||
python3 -m pytest \
|
||||
test/registered/hicache/test_hicache_storage_umbp_backend.py -v
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import is_hip, kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_amd_ci(
|
||||
est_time=3600,
|
||||
suite="nightly-amd-8-gpu-mi35x-deepseek-v4-flash",
|
||||
nightly=True,
|
||||
)
|
||||
|
||||
DEEPSEEK_V4_FLASH_FP8_MODEL_PATH = os.environ.get(
|
||||
"DEEPSEEK_V4_FP8_MODEL_PATH", "sgl-project/DeepSeek-V4-Flash-FP8"
|
||||
)
|
||||
SERVER_LAUNCH_TIMEOUT = 3600
|
||||
PAGE_SIZE = 256
|
||||
TP_SIZE = 8
|
||||
|
||||
|
||||
@unittest.skipUnless(is_hip(), "UMBP HiCache requires ROCm.")
|
||||
@unittest.skipUnless(
|
||||
os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "unified_kv_triton")
|
||||
== "unified_kv_triton",
|
||||
"UMBP HiCache E2E only runs in the unified_kv_triton DSV4 nightly leg.",
|
||||
)
|
||||
class TestHiCacheStorageUMBPBackend(CustomTestCase):
|
||||
"""DeepSeek-V4 hybrid HostPoolGroup round trip through local UMBP L3."""
|
||||
|
||||
input_ids = list(range(4000, 5024))
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEEPSEEK_V4_FLASH_FP8_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = None
|
||||
|
||||
try:
|
||||
cls._launch_server()
|
||||
except Exception:
|
||||
cls._stop_server()
|
||||
raise
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls._stop_server()
|
||||
|
||||
@classmethod
|
||||
def _launch_server(cls):
|
||||
storage_config = {
|
||||
"dram_capacity_bytes": 1 * 1024 * 1024 * 1024,
|
||||
"ssd_enabled": True,
|
||||
"ssd_storage_dir": "/tmp/umbp_dsv4_local",
|
||||
"ssd_capacity_bytes": 20 * 1024 * 1024 * 1024,
|
||||
}
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--tp-size",
|
||||
str(TP_SIZE),
|
||||
"--attention-backend",
|
||||
"dsv4",
|
||||
"--kv-cache-dtype",
|
||||
"fp8_e4m3",
|
||||
"--page-size",
|
||||
str(PAGE_SIZE),
|
||||
"--chunked-prefill-size",
|
||||
"8192",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
"--disable-cuda-graph",
|
||||
"--disable-shared-experts-fusion",
|
||||
"--enable-hierarchical-cache",
|
||||
"--hicache-ratio",
|
||||
"2",
|
||||
"--hicache-write-policy",
|
||||
"write_through",
|
||||
"--hicache-storage-prefetch-policy",
|
||||
"wait_complete",
|
||||
"--hicache-io-backend",
|
||||
"direct",
|
||||
"--hicache-mem-layout",
|
||||
"page_first",
|
||||
"--hicache-storage-backend",
|
||||
"mori",
|
||||
"--hicache-storage-backend-extra-config",
|
||||
json.dumps(storage_config),
|
||||
"--enable-cache-report",
|
||||
"--enable-metrics",
|
||||
"--swa-full-tokens-ratio",
|
||||
"0.1",
|
||||
"--max-total-tokens",
|
||||
"20000",
|
||||
"--max-running-requests",
|
||||
"4",
|
||||
"--watchdog-timeout",
|
||||
"1200",
|
||||
]
|
||||
|
||||
env = os.environ.copy()
|
||||
# An absent master address keeps every TP rank in standalone local mode,
|
||||
# so this E2E does not require an RDMA-capable CI runner.
|
||||
env.pop("UMBP_MASTER_ADDRESS", None)
|
||||
env.update(
|
||||
{
|
||||
"SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
|
||||
"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
|
||||
"SGLANG_DSV4_FP4_EXPERTS": "0",
|
||||
"SGLANG_HACK_FLASHMLA_BACKEND": "unified_kv_triton",
|
||||
"SGLANG_USE_ROCM700A": "0",
|
||||
"AITER_BF16_FP8_MOE_BOUND": "0",
|
||||
# Correctness does not depend on pre-reserved hugepages, and
|
||||
# disabling them makes the E2E portable across MI35x runners.
|
||||
"SGLANG_HICACHE_HOST_HUGEPAGE": "0",
|
||||
"UMBP_DRAM_USE_HUGEPAGES": "0",
|
||||
}
|
||||
)
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=SERVER_LAUNCH_TIMEOUT,
|
||||
other_args=other_args,
|
||||
env=env,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _stop_server(cls):
|
||||
process = getattr(cls, "process", None)
|
||||
if process is None:
|
||||
return
|
||||
if process.poll() is None:
|
||||
# Give UMBP clients a chance to close their local tiers before the
|
||||
# process tree is force-killed.
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=60)
|
||||
except subprocess.TimeoutExpired:
|
||||
kill_process_tree(process.pid)
|
||||
cls.process = None
|
||||
|
||||
def _flush_device_and_host_cache(self):
|
||||
response = requests.post(
|
||||
self.base_url + "/flush_cache",
|
||||
params={"timeout": 60},
|
||||
timeout=90,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
def _generate(self):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"input_ids": self.input_ids,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 8,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
},
|
||||
timeout=1200,
|
||||
)
|
||||
self.assertEqual(
|
||||
response.status_code,
|
||||
200,
|
||||
f"Request failed: {response.status_code} - {response.text}",
|
||||
)
|
||||
return response.json()
|
||||
|
||||
def test_hybrid_host_pool_round_trip_from_umbp(self):
|
||||
self._flush_device_and_host_cache()
|
||||
|
||||
first = self._generate()
|
||||
self.assertEqual(first["meta_info"]["cached_tokens"], 0)
|
||||
|
||||
# Writes are asynchronous below the request path. This mirrors the
|
||||
# Mooncake E2E drain before forcing the next request to use L3.
|
||||
time.sleep(15)
|
||||
self._flush_device_and_host_cache()
|
||||
|
||||
second = self._generate()
|
||||
cached_details = second["meta_info"].get("cached_tokens_details") or {}
|
||||
storage_cached_tokens = int(cached_details.get("storage", 0))
|
||||
|
||||
self.assertGreaterEqual(
|
||||
storage_cached_tokens,
|
||||
PAGE_SIZE,
|
||||
"Expected DeepSeek-V4 side-pool KV to load from UMBP storage, "
|
||||
f"got {cached_details=}",
|
||||
)
|
||||
self.assertEqual(cached_details.get("storage_backend"), "UMBPStore")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -684,6 +684,22 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
|
||||
self.assertEqual(group.layout, "page_first")
|
||||
self.assertTrue(group.can_use_write_back_jit)
|
||||
|
||||
def test_host_pool_group_destroys_logical_anchor(self):
|
||||
logical_host_pool = LogicalHostPool(8, 2, layout="page_first")
|
||||
group = HostPoolGroup(
|
||||
[
|
||||
PoolEntry(
|
||||
name=PoolName.KV,
|
||||
host_pool=logical_host_pool,
|
||||
device_pool=None,
|
||||
layer_mapper=lambda _: 0,
|
||||
is_primary_index_anchor=True,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
self.assertIsNone(group.destroy())
|
||||
|
||||
def test_write_back_jit_hybrid_write_keeps_extra_host_indices_on_cpu(self):
|
||||
captured = {}
|
||||
|
||||
|
||||
@@ -2,25 +2,20 @@
|
||||
"""Unit tests for UMBPStore with mocked HostKVCache."""
|
||||
|
||||
import ctypes
|
||||
import importlib
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from dataclasses import dataclass
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
import mori.umbp # noqa: F401
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
from sglang.test.ci.ci_register import register_amd_ci
|
||||
|
||||
# UMBPStore wraps mori's UMBP client (AMD/ROCm only). On machines without mori
|
||||
# (e.g. NVIDIA / CPU CI) the whole TestCase is skipped instead of failing at
|
||||
# import time, so the CI runner (`python3 <file> -f`) exits cleanly.
|
||||
try:
|
||||
import mori.umbp # noqa: F401
|
||||
|
||||
HAS_MORI = True
|
||||
except ImportError:
|
||||
HAS_MORI = False
|
||||
register_amd_ci(est_time=30, suite="stage-a-test-1-gpu-small-amd")
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -90,12 +85,39 @@ class MockHostKVCache:
|
||||
return bytes(ctypes.string_at(self._buffer_ptr + v_offset, self.element_size))
|
||||
|
||||
|
||||
class MockLogicalHostPool:
|
||||
layout = "page_first"
|
||||
page_size = 1
|
||||
kv_buffer = None
|
||||
|
||||
|
||||
class MockHybridSidePool:
|
||||
page_size = 1
|
||||
|
||||
def get_page_buffer_meta(self, indices):
|
||||
return [1000 + i * 8 for i in range(len(indices))], [8] * len(indices)
|
||||
|
||||
|
||||
def import_umbp_store_module():
|
||||
"""Import UMBPStore without pulling GPU-only memory-pool dependencies."""
|
||||
module_name = "sglang.srt.mem_cache.storage.umbp.umbp_store"
|
||||
if module_name in sys.modules:
|
||||
return sys.modules[module_name]
|
||||
|
||||
fake_memory_pool_host = ModuleType("sglang.srt.mem_cache.memory_pool_host")
|
||||
fake_memory_pool_host.HostKVCache = object
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"sglang.srt.mem_cache.memory_pool_host": fake_memory_pool_host},
|
||||
):
|
||||
return importlib.import_module(module_name)
|
||||
|
||||
|
||||
def make_indices(indices):
|
||||
"""Create a list that acts like a torch.Tensor of indices."""
|
||||
return indices
|
||||
|
||||
|
||||
@unittest.skipUnless(HAS_MORI, "mori.umbp not available (AMD/ROCm only)")
|
||||
class TestUMBPStore(unittest.TestCase):
|
||||
def test_basic_set_get(self):
|
||||
from sglang.srt.mem_cache.storage.umbp.umbp_store import UMBPStore
|
||||
@@ -290,5 +312,162 @@ class TestUMBPStore(unittest.TestCase):
|
||||
store.clear()
|
||||
|
||||
|
||||
class TestUMBPStoreDefensiveSemantics(unittest.TestCase):
|
||||
@staticmethod
|
||||
def _make_v2_store():
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName
|
||||
|
||||
UMBPStore = import_umbp_store_module().UMBPStore
|
||||
store = UMBPStore.__new__(UMBPStore)
|
||||
store.client = MagicMock()
|
||||
store.client.is_distributed.return_value = False
|
||||
store.registered_pools = {}
|
||||
store._kv_anchor_is_logical = True
|
||||
store.is_mla_backend = True
|
||||
store.mla_suffix = ""
|
||||
store.mha_suffix = "0"
|
||||
store.register_mem_host_pool_v2(MockHybridSidePool(), PoolName.DEEPSEEK_V4_C4)
|
||||
return store
|
||||
|
||||
def test_constructor_preserves_logical_anchor_detection(self):
|
||||
umbp_module = import_umbp_store_module()
|
||||
|
||||
class FakeUMBPConfig:
|
||||
def __init__(self):
|
||||
self.role = None
|
||||
self.dram = SimpleNamespace(capacity_bytes=0)
|
||||
self.ssd = SimpleNamespace(
|
||||
enabled=False,
|
||||
storage_dir="/tmp",
|
||||
capacity_bytes=0,
|
||||
ssd_backend="file",
|
||||
spdk_proxy_tenant_id=0,
|
||||
spdk_proxy_tenant_quota_bytes=0,
|
||||
)
|
||||
self.distributed = None
|
||||
|
||||
@classmethod
|
||||
def from_environment(cls):
|
||||
return cls()
|
||||
|
||||
class FakeUMBPClient:
|
||||
def __init__(self, _config):
|
||||
pass
|
||||
|
||||
fake_role = SimpleNamespace(
|
||||
Standalone="standalone",
|
||||
SharedSSDLeader="leader",
|
||||
SharedSSDFollower="follower",
|
||||
)
|
||||
imported = (
|
||||
FakeUMBPClient,
|
||||
FakeUMBPConfig,
|
||||
fake_role,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
config = MockStorageConfig(
|
||||
extra_config={"dram_capacity_bytes": 1024, "ssd_enabled": False}
|
||||
)
|
||||
|
||||
with patch.object(umbp_module, "_import_umbp_client", return_value=imported):
|
||||
store = umbp_module.UMBPStore(config, MockLogicalHostPool())
|
||||
|
||||
self.assertTrue(store._kv_anchor_is_logical)
|
||||
self.assertEqual(store.batch_set_v1(["page0"], [0]), [True])
|
||||
|
||||
def test_short_batch_exists_result_fails_closed(self):
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
|
||||
store = self._make_v2_store()
|
||||
store.client.batch_exists.return_value = [True]
|
||||
transfer = PoolTransfer(
|
||||
name=PoolName.DEEPSEEK_V4_C4,
|
||||
keys=["page0", "page1"],
|
||||
host_indices=[0, 1],
|
||||
)
|
||||
|
||||
result = store.batch_exists_v2(["page0", "page1"], [transfer])
|
||||
|
||||
self.assertEqual(result.kv_hit_pages, 0)
|
||||
|
||||
def test_batch_exists_v2_narrows_queries_across_side_pools(self):
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
|
||||
store = self._make_v2_store()
|
||||
store.register_mem_host_pool_v2(MockHybridSidePool(), PoolName.DEEPSEEK_V4_C128)
|
||||
page_keys = [f"page{i}" for i in range(4)]
|
||||
store.client.batch_exists.side_effect = [
|
||||
[True, True, False, True],
|
||||
[True, False],
|
||||
]
|
||||
transfers = [
|
||||
PoolTransfer(
|
||||
name=PoolName.DEEPSEEK_V4_C4,
|
||||
keys=page_keys,
|
||||
host_indices=[0, 1, 2, 3],
|
||||
),
|
||||
PoolTransfer(
|
||||
name=PoolName.DEEPSEEK_V4_C128,
|
||||
keys=page_keys,
|
||||
host_indices=[0, 1, 2, 3],
|
||||
),
|
||||
]
|
||||
|
||||
result = store.batch_exists_v2(page_keys, transfers)
|
||||
|
||||
queried_keys = [
|
||||
invocation.args[0]
|
||||
for invocation in store.client.batch_exists.call_args_list
|
||||
]
|
||||
self.assertEqual(
|
||||
queried_keys,
|
||||
[
|
||||
[f"{key}__{PoolName.DEEPSEEK_V4_C4}" for key in page_keys],
|
||||
[f"{key}__{PoolName.DEEPSEEK_V4_C128}" for key in page_keys[:2]],
|
||||
],
|
||||
)
|
||||
self.assertEqual(result.kv_hit_pages, 1)
|
||||
self.assertEqual(
|
||||
result.extra_pool_hit_pages,
|
||||
{
|
||||
PoolName.KV: 4,
|
||||
PoolName.DEEPSEEK_V4_C4: 2,
|
||||
PoolName.DEEPSEEK_V4_C128: 1,
|
||||
},
|
||||
)
|
||||
|
||||
def test_short_batch_get_result_marks_every_page_failed(self):
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
|
||||
store = self._make_v2_store()
|
||||
store.client.batch_get_into_ptr.return_value = [True]
|
||||
transfer = PoolTransfer(
|
||||
name=PoolName.DEEPSEEK_V4_C4,
|
||||
keys=["page0", "page1"],
|
||||
host_indices=[0, 1],
|
||||
)
|
||||
|
||||
result = store.batch_get_v2([transfer])
|
||||
|
||||
self.assertEqual(result[PoolName.DEEPSEEK_V4_C4], [False, False])
|
||||
|
||||
def test_short_batch_set_result_marks_every_page_failed(self):
|
||||
from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
|
||||
|
||||
store = self._make_v2_store()
|
||||
store.client.batch_put_from_ptr.return_value = [True]
|
||||
transfer = PoolTransfer(
|
||||
name=PoolName.DEEPSEEK_V4_C4,
|
||||
keys=["page0", "page1"],
|
||||
host_indices=[0, 1],
|
||||
)
|
||||
|
||||
result = store.batch_set_v2([transfer])
|
||||
|
||||
self.assertEqual(result[PoolName.DEEPSEEK_V4_C4], [False, False])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -143,6 +143,7 @@ NIGHTLY_SUITES = {
|
||||
"nightly-amd-4-gpu",
|
||||
"nightly-amd-8-gpu",
|
||||
"nightly-amd-vlm",
|
||||
"nightly-amd-8-gpu-mi35x-deepseek-v4-flash",
|
||||
# MI35x 8-GPU suite (different model configs)
|
||||
"nightly-amd-8-gpu-mi35x",
|
||||
],
|
||||
|
||||
Reference in New Issue
Block a user