minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 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
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
+19 -6
View File
@@ -105,7 +105,11 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
# The padded slot 0 is used for writing dummy outputs from padded tokens.
# Continuous memory improves the efficiency of Ascend`s transmission backend,
# while other backends remain unchanged.
self.k_buffer = torch.zeros(
# FIA exposes the KV cache as per-layer Python views so graph
# capture does not retain the full multi-layer tensor. HiCache's
# NPU exchange operator still requires the original contiguous
# [layer, page, token, head, dim] allocation.
self._hicache_k_buffer = torch.zeros(
(
self.layer_num,
self.size // self.page_size + 1,
@@ -116,7 +120,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
dtype=self.store_dtype,
device=self.device,
)
self.v_buffer = torch.zeros(
self._hicache_v_buffer = torch.zeros(
(
self.layer_num,
self.size // self.page_size + 1,
@@ -130,8 +134,11 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
# Keep a reference to the contiguous tensor for HiCache
# D2H/H2D transfers (transfer_kv_dim_exchange expects a
# tensor, not the per-layer list used in FIA mode below).
self.k_buffer_tensor = self.k_buffer
self.v_buffer_tensor = self.v_buffer
self.k_buffer_tensor = self._hicache_k_buffer
self.v_buffer_tensor = self._hicache_v_buffer
self.k_buffer = self._hicache_k_buffer
self.v_buffer = self._hicache_v_buffer
if self.use_fia:
# Use per-layer Python lists to avoid torch.compile capturing
@@ -139,14 +146,20 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
# Each layer view: [P*ps, 1, H, D], sharing the contiguous
# storage allocated above.
self.k_buffer = [
self.k_buffer[i].view(-1, 1, self.head_num, self.head_dim)
self._hicache_k_buffer[i].view(-1, 1, self.head_num, self.head_dim)
for i in range(self.layer_num)
]
self.v_buffer = [
self.v_buffer[i].view(-1, 1, self.head_num, self.v_head_dim)
self._hicache_v_buffer[i].view(
-1, 1, self.head_num, self.v_head_dim
)
for i in range(self.layer_num)
]
def get_hicache_transfer_buffers(self):
"""Return contiguous all-layer KV tensors for NPU HiCache IO."""
return self._hicache_k_buffer, self._hicache_v_buffer
def _init_kv_copy_and_warmup(self):
# implementation relies on self.data_strides / self.data_ptrs, which the
# NPU paged buffer layout never builds.
python/sglang/srt/mem_cache/pool_host/mamba.py
+90 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
import os
import threading
from typing import Optional
@@ -18,10 +19,13 @@ from sglang.srt.mem_cache.pool_host.common import (
ALLOC_MEMORY_FUNCS,
get_allocator_from_storage,
)
from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils import is_cuda, is_hip, is_npu
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
transfer_state_per_layer_direct_pf_lf = None
transfer_state_all_layer_direct_lf_pf = None
if _is_cuda or _is_hip:
from sgl_kernel.kvcacheio import (
transfer_kv_all_layer_direct_lf_pf,
@@ -34,10 +38,32 @@ if _is_cuda or _is_hip:
transfer_kv_mamba_lf_pf,
transfer_kv_mamba_pf_lf,
)
if _is_npu:
try:
from sgl_kernel_npu.kvcacheio import (
transfer_state_all_layer_direct_lf_pf,
transfer_state_per_layer_direct_pf_lf,
)
except ImportError:
pass
logger = logging.getLogger(__name__)
_NPU_HICACHE_MAMBA_IO_ENV = "SGLANG_NPU_HICACHE_MAMBA_IO"
_NPU_HICACHE_MAMBA_IO_MODES = {"sync", "async"}
def _npu_hicache_mamba_io_mode() -> str:
mode = os.getenv(_NPU_HICACHE_MAMBA_IO_ENV, "sync").strip().lower()
if mode not in _NPU_HICACHE_MAMBA_IO_MODES:
raise ValueError(
f"{_NPU_HICACHE_MAMBA_IO_ENV} must be one of "
f"{sorted(_NPU_HICACHE_MAMBA_IO_MODES)}, got {mode!r}."
)
return mode
class MambaPoolHost(HostKVCache):
def __init__(
self,
@@ -128,10 +154,36 @@ class MambaPoolHost(HostKVCache):
]
self.kv_buffer = self.init_kv_buffer()
self._configure_npu_mamba_io()
self._init_write_back_staging_buffers()
self.lock = threading.RLock()
self.clear()
def _configure_npu_mamba_io(self) -> None:
mode = _npu_hicache_mamba_io_mode()
if mode == "sync":
logger.info("NPU HiCache Mamba state transfer mode: sync torch fallback.")
return
required_ops = (
transfer_state_per_layer_direct_pf_lf,
transfer_state_all_layer_direct_lf_pf,
)
required_torch_ops = (
"transfer_state_per_layer_direct_pf_lf",
"transfer_state_all_layer_direct_lf_pf",
)
if any(op is None for op in required_ops) or any(
not hasattr(torch.ops.npu, op_name) for op_name in required_torch_ops
):
raise RuntimeError(
"NPU HiCache Mamba async state transfer requires "
"the per-layer PF->LF and all-layer LF->PF direct operators "
"from sgl-kernel-npu."
)
logger.info("NPU HiCache Mamba state transfer mode: native async.")
def init_kv_buffer(self):
_host_alloc = ALLOC_MEMORY_FUNCS[self.device_pool.device]
@@ -361,6 +413,22 @@ class MambaPoolHost(HostKVCache):
layer_id=layer_id,
page_size=1,
)
elif io_backend == "kernel_ascend":
if _npu_hicache_mamba_io_mode() == "async":
transfer_state_per_layer_direct_pf_lf(
src=src,
dst=dst,
src_indices=src_indices,
dst_indices=dst_indices,
layer_id=layer_id,
)
else:
host_indices = src_indices.to(dtype=torch.int64, device=src.device)
device_indices = dst_indices.to(dtype=torch.int64, device=dst.device)
values = (
src.select(1, layer_id).index_select(0, host_indices).select(1, 0)
)
dst.index_copy_(0, device_indices, values.to(device=dst.device))
else:
raise ValueError(f"Unsupported io_backend: {io_backend}")
@@ -404,6 +472,27 @@ class MambaPoolHost(HostKVCache):
dst_indices=dst_indices,
page_size=1,
)
elif io_backend == "kernel_ascend":
if _npu_hicache_mamba_io_mode() == "async":
transfer_state_all_layer_direct_lf_pf(
device_states=[src_layers],
host_states=[dst],
device_indices=src_indices,
host_indices=dst_indices,
)
else:
device_indices = src_indices.to(
dtype=torch.int64, device=src_layers.device
)
host_indices = dst_indices.to(dtype=torch.int64, device=dst.device)
values = (
src_layers.index_select(1, device_indices)
.movedim(0, 1)
.unsqueeze(2)
.contiguous()
.to(device=dst.device)
)
dst.index_copy_(0, host_indices, values)
else:
raise ValueError(f"Unsupported io_backend: {io_backend}")
python/sglang/srt/mem_cache/pool_host/mha.py
+43 -31
View File
@@ -131,7 +131,7 @@ class MHATokenToKVPoolHost(HostKVCache):
self.device_pool.device,
host_memory_registered=self.pin_memory,
)
if self.mtp_draft_device_pools:
if self.mtp_draft_device_pools and not _is_npu:
device_pools = (self.device_pool, *self.mtp_draft_device_pools)
if not _is_npu:
self.packed_device_k_data_ptrs = torch.cat(
@@ -364,22 +364,22 @@ class MHATokenToKVPoolHost(HostKVCache):
if self.layout == "page_first_direct":
# Ascend-specific: transfer KV data for all layers when layer_id == 0
if host_layer_id == 0:
device_k = getattr(
device_pool, "k_buffer_tensor", device_pool.k_buffer
)
device_v = getattr(
device_pool, "v_buffer_tensor", device_pool.v_buffer
)
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=self.k_buffer,
device_v=device_v,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
for (
device_k,
device_v,
host_k,
host_v,
) in self._npu_transfer_buffers(device_pool):
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=host_k,
device_v=device_v,
host_v=host_v,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
@@ -400,6 +400,19 @@ class MHATokenToKVPoolHost(HostKVCache):
device_pool.v_buffer,
)
def _npu_transfer_buffers(self, target_device_pool):
layer_start = 0
for pool in (target_device_pool, *self.mtp_draft_device_pools):
device_k, device_v = pool.get_hicache_transfer_buffers()
layer_end = layer_start + device_k.shape[0]
yield (
device_k,
device_v,
self.k_buffer[:, layer_start:layer_end],
self.v_buffer[:, layer_start:layer_end],
)
layer_start = layer_end
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
):
@@ -502,20 +515,19 @@ class MHATokenToKVPoolHost(HostKVCache):
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_direct":
# In FIA mode, k_buffer/v_buffer are per-layer lists;
# use the 5-D contiguous view for transfer_kv_dim_exchange.
device_k = getattr(device_pool, "k_buffer_tensor", device_pool.k_buffer)
device_v = getattr(device_pool, "v_buffer_tensor", device_pool.v_buffer)
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=self.k_buffer,
device_v=device_v,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
for device_k, device_v, host_k, host_v in self._npu_transfer_buffers(
device_pool
):
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=host_k,
device_v=device_v,
host_v=host_v,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
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()