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

feat(hicache): support NPU Mamba states with FIA and async IO (#32500)

Browse Source
This commit is contained in:
qyb233
2026-09-17 10:12:26 +08:00
committed by GitHub
parent 4c85172f3a
commit 4fb9b5b5ba
6 changed files with 652 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/npu/basic_function/HiCache/test_npu_hicache_mamba.py
+184
View File
@@ -0,0 +1,184 @@
"""NPU HiCache L3 coverage for hybrid Mamba models."""
import os
import shutil
import tempfile
import time
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_npu_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_npu_ci(est_time=400, suite="nightly-1-npu-a3", nightly=True)
TEST_MODEL_MATRIX = {
"/root/.cache/modelscope/hub/models/Qwen/Qwen3.5-0.8B": {
"target_token_id": 1000,
},
}
class TestNPUMambaHiCache(CustomTestCase):
"""Exercise Mamba state write-back and L3 load-back on NPU."""
# A reusable Mamba checkpoint must be strictly inside the prompt. Ending the
# prompt exactly at 1024 does not make that boundary reusable by a second
# identical request. 3200 leaves the 3072 checkpoint inside the prompt and
# also exercises the configured 1024-token chunked-prefill path.
prompt_tokens = 3200
@classmethod
def setUpClass(cls):
cls.models = TEST_MODEL_MATRIX.keys()
cls.base_url = DEFAULT_URL_FOR_TEST
cls.common_args = [
"--tp-size",
"1",
"--mem-fraction-static",
"0.20",
"--chunked-prefill-size",
"1024",
"--page-size",
"64",
"--enable-hierarchical-cache",
"--enable-cache-report",
"--hicache-size",
"1",
"--hicache-storage-backend",
"file",
"--hicache-storage-prefetch-policy",
"wait_complete",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"kernel",
"--hicache-mem-layout",
"page_first_direct",
]
def _generate(self, token_id: int) -> dict:
response = requests.post(
f"{self.base_url}/generate",
json={
"input_ids": [token_id] * self.prompt_tokens,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 16,
"ignore_eos": True,
},
},
timeout=120,
)
self.assertEqual(response.status_code, 200, response.text)
return response.json()
@staticmethod
def _storage_files(storage_dir: str) -> set[str]:
return {
os.path.relpath(os.path.join(root, name), storage_dir)
for root, _, names in os.walk(storage_dir)
for name in names
if os.path.getsize(os.path.join(root, name)) > 0
}
def _wait_for_mamba_storage_file(
self, storage_dir: str, files_before: set[str], timeout: float = 30
) -> set[str]:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
files_after = self._storage_files(storage_dir)
new_files = files_after - files_before
if any("mamba" in name.lower() for name in new_files):
return files_after
time.sleep(0.1)
self.fail("Timed out waiting for a new Mamba HiCache storage file.")
def test_mamba_state_restores_from_l3(self):
for model in self.models:
with self.subTest(model=model):
storage_dir = tempfile.mkdtemp(prefix="npu_mamba_hicache_")
process = None
try:
process = popen_launch_server(
model,
self.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=self.common_args,
env={
**os.environ,
"CUDA_VISIBLE_DEVICES": "0",
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": storage_dir,
},
)
target_token_id = TEST_MODEL_MATRIX[model]["target_token_id"]
storage_files_before = self._storage_files(storage_dir)
first = self._generate(target_token_id)
first_meta = first["meta_info"]
self.assertEqual(int(first_meta.get("cached_tokens", 0)), 0)
# The first cache hit reaches write_through's hit-count
# threshold and starts the asynchronous L1 -> L2 -> L3 copy.
warm = self._generate(target_token_id)
warm_details = warm["meta_info"].get("cached_tokens_details") or {}
self.assertGreater(
int(warm_details.get("device", 0) or 0),
0,
f"Expected an L1 device hit; got meta_info={warm['meta_info']}",
)
self.assertEqual(
int(warm_details.get("host", 0) or 0),
0,
f"Expected no L2 contribution: {warm_details}",
)
self.assertEqual(
int(warm_details.get("storage", 0) or 0),
0,
f"Expected no L3 contribution: {warm_details}",
)
# Do not infer persistence from the request response: wait
# until the Mamba sidecar has physically reached the file
# backend before clearing the in-memory cache levels.
storage_files_after = self._wait_for_mamba_storage_file(
storage_dir, storage_files_before
)
self.assertTrue(storage_files_after - storage_files_before)
# Remove L1/L2 residency without clearing the file backend.
flush = requests.post(
f"{self.base_url}/flush_cache",
params={"timeout": 30},
timeout=40,
)
flush.raise_for_status()
restored = self._generate(target_token_id)
restored_meta = restored["meta_info"]
details = restored_meta.get("cached_tokens_details") or {}
self.assertGreater(
int(details.get("storage", 0) or 0),
0,
f"Expected an L3 storage hit; got meta_info={restored_meta}",
)
self.assertEqual(
restored["text"],
first["text"],
"Mamba L3 restore changed deterministic generation output.",
)
finally:
if process:
kill_process_tree(process.pid)
shutil.rmtree(storage_dir, ignore_errors=True)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py
+124
View File
@@ -585,6 +585,57 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
torch.equal(host.v_buffer[host_indices, layer_id], expected_v[layer_id])
)
def test_npu_mha_transfer_uses_contiguous_hicache_backing(self):
host = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
host.layout = "page_first_direct"
host.page_size = 2
host.kv_buffer = torch.empty(2, 2, 2, 2, 1, 1)
device_k = torch.empty(2, 3, 2, 1, 1)
device_v = torch.empty_like(device_k)
device_pool = SimpleNamespace(
# FIA exposes lists here; these must not be sent to the operator.
k_buffer=[device_k[layer].reshape(-1, 1, 1, 1) for layer in range(2)],
v_buffer=[device_v[layer].reshape(-1, 1, 1, 1) for layer in range(2)],
get_hicache_transfer_buffers=mock.Mock(return_value=(device_k, device_v)),
)
host_indices = _indices(0, 2)
device_indices = _indices(2, 4)
directions = SimpleNamespace(H2D="H2D", D2H="D2H")
with (
mock.patch(
f"{MHA_POOL_HOST_MODULE}.TransferDirection",
directions,
create=True,
),
mock.patch(
f"{MHA_POOL_HOST_MODULE}.transfer_kv_dim_exchange",
create=True,
) as transfer,
):
host.backup_from_device_all_layer(
device_pool,
host_indices,
device_indices,
io_backend="kernel_ascend",
)
host.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=0,
io_backend="kernel_ascend",
)
self.assertEqual(device_pool.get_hicache_transfer_buffers.call_count, 2)
self.assertEqual(transfer.call_count, 2)
for call in transfer.call_args_list:
self.assertIs(call.kwargs["device_k"], device_k)
self.assertIs(call.kwargs["device_v"], device_v)
self.assertEqual(transfer.call_args_list[0].kwargs["direction"], "D2H")
self.assertEqual(transfer.call_args_list[1].kwargs["direction"], "H2D")
def test_mla_backup_then_load_roundtrip_uses_staged(self):
layer_num = 2
kv_cache_dim = 5
@@ -750,6 +801,79 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
)
)
def test_mamba_kernel_npu_backup_then_load_roundtrip(self):
num_layers = 2
host_indices = torch.tensor([1, 3], dtype=torch.int64)
device_indices = torch.tensor([2, 5], dtype=torch.int64)
temporal = torch.arange(num_layers * 8 * 3, dtype=torch.float32).reshape(
num_layers, 8, 3
)
conv = (
torch.arange(num_layers * 8 * 2, dtype=torch.float32).reshape(
num_layers, 8, 2
)
/ 8
).to(torch.bfloat16)
device_pool = SimpleNamespace(
mamba_cache=SimpleNamespace(temporal=temporal.clone(), conv=[conv.clone()])
)
expected_temporal = device_pool.mamba_cache.temporal[:, device_indices].clone()
expected_conv = device_pool.mamba_cache.conv[0][:, device_indices].clone()
host = MambaPoolHost.__new__(MambaPoolHost)
host.layout = "page_first_direct"
host.num_mamba_layers = num_layers
host.temporal_state_elem_size = 3
host.temporal_buffer = torch.zeros(8, num_layers, 1, 3, dtype=torch.float32)
host.conv_state_shapes = [(2,)]
host.conv_buffer = [torch.zeros(8, num_layers, 1, 2, dtype=torch.bfloat16)]
host.temporal_staging_buffer = None
host.conv_staging_buffers = [None]
host._temporal_can_use_jit = False
host._conv_can_use_jit = [False]
host.temporal_device_ptrs = torch.empty(0, dtype=torch.uint64)
host.conv_device_ptrs = [torch.empty(0, dtype=torch.uint64)]
host.backup_from_device_all_layer(
device_pool,
host_indices,
device_indices,
io_backend="kernel_ascend",
)
device_pool.mamba_cache.temporal.zero_()
device_pool.mamba_cache.conv[0].zero_()
for layer_id in range(num_layers):
host.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id,
io_backend="kernel_ascend",
)
self.assertTrue(
torch.equal(
device_pool.mamba_cache.temporal[:, device_indices], expected_temporal
)
)
self.assertTrue(
torch.equal(
device_pool.mamba_cache.conv[0][:, device_indices], expected_conv
)
)
self.assertTrue(
torch.equal(
host.temporal_buffer[host_indices].squeeze(2).transpose(0, 1),
expected_temporal,
)
)
self.assertTrue(
torch.equal(
host.conv_buffer[0][host_indices].squeeze(2).transpose(0, 1),
expected_conv,
)
)
def test_deepseek_v4_paged_pool_backup_then_load_roundtrip_uses_staged(self):
layer_num = 2
slot_page_size = 2
test/registered/unit/mem_cache/test_npu_mamba_async_layout.py
+192
View File
@@ -0,0 +1,192 @@
import os
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import torch
import sglang.srt.mem_cache.pool_host.mamba as mamba_pool_host
from sglang.srt.mem_cache.pool_host.mamba import (
MambaPoolHost,
_npu_hicache_mamba_io_mode,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestNPUMambaAsyncConfig(unittest.TestCase):
def test_accepts_explicit_sync_and_async_modes(self):
for mode in ("sync", "async"):
with (
self.subTest(mode=mode),
patch.dict(os.environ, {"SGLANG_NPU_HICACHE_MAMBA_IO": mode}),
):
self.assertEqual(_npu_hicache_mamba_io_mode(), mode)
def test_rejects_auto_mode(self):
with (
patch.dict(os.environ, {"SGLANG_NPU_HICACHE_MAMBA_IO": "auto"}),
self.assertRaisesRegex(ValueError, "must be one of"),
):
_npu_hicache_mamba_io_mode()
@patch.object(mamba_pool_host, "_is_npu", True)
@patch.object(
mamba_pool_host,
"_npu_hicache_mamba_io_mode",
return_value="async",
)
@patch.object(mamba_pool_host, "transfer_state_per_layer_direct_pf_lf", None)
@patch.object(mamba_pool_host, "transfer_state_all_layer_direct_lf_pf", None)
def test_async_requires_native_operator(self, _mock_mode):
pool = MambaPoolHost.__new__(MambaPoolHost)
with self.assertRaisesRegex(
RuntimeError, "per-layer PF->LF and all-layer LF->PF"
):
pool._configure_npu_mamba_io()
def test_async_h2d_is_dispatched_per_component_from_copy_helper(self):
pool = MambaPoolHost.__new__(MambaPoolHost)
host = torch.empty(4, 3, 1, 2, 4)
device_layers = torch.empty(3, 5, 2, 4)
host_indices = torch.tensor([1, 3])
device_indices = torch.tensor([0, 4])
transfer_op = MagicMock()
with (
patch.object(
mamba_pool_host,
"transfer_state_per_layer_direct_pf_lf",
transfer_op,
),
patch.object(
mamba_pool_host,
"_npu_hicache_mamba_io_mode",
return_value="async",
),
):
pool._copy_tensor_pf_lf(
src=host,
dst=device_layers[2],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=2,
num_layers=3,
io_backend="kernel_ascend",
)
transfer_op.assert_called_once()
call = transfer_op.call_args.kwargs
self.assertIs(call["src"], host)
self.assertEqual(call["dst"].data_ptr(), device_layers[2].data_ptr())
self.assertIs(call["src_indices"], host_indices)
self.assertIs(call["dst_indices"], device_indices)
self.assertEqual(call["layer_id"], 2)
def test_async_d2h_is_dispatched_per_component_from_copy_helper(self):
pool = MambaPoolHost.__new__(MambaPoolHost)
device_layers = torch.empty(3, 5, 2, 4)
host = torch.empty(4, 3, 1, 2, 4)
device_indices = torch.tensor([0, 4])
host_indices = torch.tensor([1, 3])
transfer_op = MagicMock()
with (
patch.object(
mamba_pool_host,
"transfer_state_all_layer_direct_lf_pf",
transfer_op,
),
patch.object(
mamba_pool_host,
"_npu_hicache_mamba_io_mode",
return_value="async",
),
):
pool._copy_tensor_all_layers_lf_pf(
src_layers=device_layers,
dst=host,
src_indices=device_indices,
dst_indices=host_indices,
num_layers=3,
io_backend="kernel_ascend",
src_ptrs=torch.empty(0),
)
transfer_op.assert_called_once_with(
device_states=[device_layers],
host_states=[host],
device_indices=device_indices,
host_indices=host_indices,
)
def test_sync_h2d_torch_fallback_is_preserved(self):
pool = MambaPoolHost.__new__(MambaPoolHost)
host = torch.arange(4 * 3 * 1 * 2, dtype=torch.float32).reshape(4, 3, 1, 2)
device_layers = torch.zeros(3, 5, 2)
host_indices = torch.tensor([1, 3])
device_indices = torch.tensor([0, 4])
with patch.object(
mamba_pool_host,
"_npu_hicache_mamba_io_mode",
return_value="sync",
):
pool._copy_tensor_pf_lf(
src=host,
dst=device_layers[2],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=2,
num_layers=3,
io_backend="kernel_ascend",
)
torch.testing.assert_close(
device_layers[2, device_indices], host[host_indices, 2, 0]
)
def test_copy_helpers_keep_static_signatures(self):
self.assertIsInstance(
MambaPoolHost.__dict__["_copy_tensor_pf_lf"], staticmethod
)
self.assertIsInstance(
MambaPoolHost.__dict__["_copy_tensor_all_layers_lf_pf"], staticmethod
)
def test_conv_only_load_skips_empty_temporal_component(self):
pool = MambaPoolHost.__new__(MambaPoolHost)
pool.layout = "page_first_direct"
pool.temporal_state_elem_size = 0
pool.num_mamba_layers = 3
pool.temporal_buffer = torch.empty(4, 3, 1, 0)
pool.conv_state_shapes = [torch.Size([2, 4])]
pool.conv_buffer = [torch.empty(4, 3, 1, 2, 4)]
pool._copy_tensor_pf_lf = MagicMock()
device_pool = SimpleNamespace(
mamba_cache=SimpleNamespace(
temporal=torch.empty(3, 4, 0),
conv=[torch.empty(3, 4, 2, 4)],
)
)
pool.load_to_device_per_layer(
device_pool=device_pool,
host_indices=torch.tensor([1]),
device_indices=torch.tensor([2]),
layer_id=1,
io_backend="kernel_ascend",
)
pool._copy_tensor_pf_lf.assert_called_once()
call = pool._copy_tensor_pf_lf.call_args.kwargs
self.assertIs(call["src"], pool.conv_buffer[0])
self.assertEqual(
call["dst"].data_ptr(), device_pool.mamba_cache.conv[0][1].data_ptr()
)
if __name__ == "__main__":
unittest.main()