From 4fb9b5b5bad41ac7de5586ec812c421fe3aa2b8b Mon Sep 17 00:00:00 2001 From: qyb233 <1959345727@qq.com> Date: Thu, 17 Sep 2026 10:12:26 +0800 Subject: [PATCH] feat(hicache): support NPU Mamba states with FIA and async IO (#32500) --- .../hardware_backend/npu/memory_pool_npu.py | 25 ++- .../sglang/srt/mem_cache/pool_host/mamba.py | 91 ++++++++- python/sglang/srt/mem_cache/pool_host/mha.py | 74 ++++--- .../HiCache/test_npu_hicache_mamba.py | 184 +++++++++++++++++ ...test_hicache_staged_write_back_dispatch.py | 124 +++++++++++ .../mem_cache/test_npu_mamba_async_layout.py | 192 ++++++++++++++++++ 6 files changed, 652 insertions(+), 38 deletions(-) create mode 100644 test/registered/npu/basic_function/HiCache/test_npu_hicache_mamba.py create mode 100644 test/registered/unit/mem_cache/test_npu_mamba_async_layout.py diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 59b296707..4a532cd1b 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -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. diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py index 669d99662..8b8eda90e 100644 --- a/python/sglang/srt/mem_cache/pool_host/mamba.py +++ b/python/sglang/srt/mem_cache/pool_host/mamba.py @@ -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}") diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index f91e6d72a..fdac9d207 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -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: diff --git a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mamba.py b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mamba.py new file mode 100644 index 000000000..af0e8b8a2 --- /dev/null +++ b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mamba.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index 7591102e3..d0646adb8 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_npu_mamba_async_layout.py b/test/registered/unit/mem_cache/test_npu_mamba_async_layout.py new file mode 100644 index 000000000..8bec03017 --- /dev/null +++ b/test/registered/unit/mem_cache/test_npu_mamba_async_layout.py @@ -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()