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/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()