feat: add page-aligned index validation for CP HiCache direct transfers
Add validate_page_aligned_token_indices utility and apply it across HiCache write/load paths, NSA indexer transfers, and CUDA direct copy kernels to reject malformed (partial, misaligned, non-contiguous) page groups before they reach native transfer code. Also validate supported CP HiCache backend/layout combinations at server startup.
This commit is contained in:
@@ -40,6 +40,7 @@ from sglang.srt.layers.dp_attention import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||||
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool
|
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool
|
||||||
|
from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices
|
||||||
from sglang.srt.utils import get_device_module
|
from sglang.srt.utils import get_device_module
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -735,6 +736,16 @@ class HiCacheController:
|
|||||||
event.record()
|
event.record()
|
||||||
self.ack_load_queue.append(HiCacheAck(event, event, [node_id]))
|
self.ack_load_queue.append(HiCacheAck(event, event, [node_id]))
|
||||||
|
|
||||||
|
def _validate_cp_hicache_page_indices(
|
||||||
|
self,
|
||||||
|
host_indices: torch.Tensor,
|
||||||
|
device_indices: torch.Tensor,
|
||||||
|
) -> None:
|
||||||
|
validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
device_indices, self.page_size, "physical_device_indices"
|
||||||
|
)
|
||||||
|
|
||||||
def _write_cp(
|
def _write_cp(
|
||||||
self,
|
self,
|
||||||
device_indices: torch.Tensor,
|
device_indices: torch.Tensor,
|
||||||
@@ -762,6 +773,11 @@ class HiCacheController:
|
|||||||
host_indices = self.mem_pool_host.alloc(len(physical_device_indices))
|
host_indices = self.mem_pool_host.alloc(len(physical_device_indices))
|
||||||
if host_indices is None:
|
if host_indices is None:
|
||||||
return HiCacheWriteFailure(required_host_slots=len(physical_device_indices))
|
return HiCacheWriteFailure(required_host_slots=len(physical_device_indices))
|
||||||
|
try:
|
||||||
|
self._validate_cp_hicache_page_indices(host_indices, physical_device_indices)
|
||||||
|
except Exception:
|
||||||
|
self.mem_pool_host.free(host_indices)
|
||||||
|
raise
|
||||||
|
|
||||||
self.write_queue.append(
|
self.write_queue.append(
|
||||||
CacheOperation(host_indices, physical_device_indices, node_id, priority)
|
CacheOperation(host_indices, physical_device_indices, node_id, priority)
|
||||||
@@ -817,10 +833,17 @@ class HiCacheController:
|
|||||||
self._append_completed_load_ack(node_id)
|
self._append_completed_load_ack(node_id)
|
||||||
return device_indices
|
return device_indices
|
||||||
|
|
||||||
|
host_indices = torch.cat(host_chunks)
|
||||||
|
physical_device_indices = torch.cat(physical_chunks)
|
||||||
|
try:
|
||||||
|
self._validate_cp_hicache_page_indices(host_indices, physical_device_indices)
|
||||||
|
except Exception:
|
||||||
|
self.mem_pool_device_allocator.free(device_indices)
|
||||||
|
raise
|
||||||
self.load_queue.append(
|
self.load_queue.append(
|
||||||
CacheOperation(
|
CacheOperation(
|
||||||
torch.cat(host_chunks),
|
host_indices,
|
||||||
torch.cat(physical_chunks),
|
physical_device_indices,
|
||||||
node_id,
|
node_id,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ from sglang.srt.mem_cache.memory_pool import (
|
|||||||
MLATokenToKVPool,
|
MLATokenToKVPool,
|
||||||
NSATokenToKVPool,
|
NSATokenToKVPool,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices
|
||||||
from sglang.srt.utils import is_cuda, is_mps, is_npu, is_xpu
|
from sglang.srt.utils import is_cuda, is_mps, is_npu, is_xpu
|
||||||
|
|
||||||
_is_cuda = is_cuda()
|
_is_cuda = is_cuda()
|
||||||
@@ -1169,10 +1170,10 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
|
|||||||
def _get_indexer_page_indices(self, host_indices, device_indices):
|
def _get_indexer_page_indices(self, host_indices, device_indices):
|
||||||
if host_indices.numel() == 0:
|
if host_indices.numel() == 0:
|
||||||
return host_indices, device_indices
|
return host_indices, device_indices
|
||||||
if host_indices.numel() % self.page_size != 0:
|
validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
|
||||||
raise ValueError(
|
validate_page_aligned_token_indices(
|
||||||
"Index buffer transfer expects page-aligned indices for NSA."
|
device_indices, self.page_size, "device_indices"
|
||||||
)
|
)
|
||||||
host_page_indices = (
|
host_page_indices = (
|
||||||
host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
|
host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
|
||||||
)
|
)
|
||||||
|
|||||||
46
python/sglang/srt/mem_cache/page_index_utils.py
Normal file
46
python/sglang/srt/mem_cache/page_index_utils.py
Normal file
@@ -0,0 +1,46 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
|
||||||
|
def validate_page_aligned_token_indices(
|
||||||
|
indices: torch.Tensor,
|
||||||
|
page_size: int,
|
||||||
|
name: str,
|
||||||
|
) -> None:
|
||||||
|
"""Validate that token indices are grouped as complete contiguous pages.
|
||||||
|
|
||||||
|
HiCache page-first direct transfers treat every `page_size` entries as one
|
||||||
|
page and derive the page id from the first token. A malformed group can
|
||||||
|
make native direct-copy code address the wrong page, so validate the exact
|
||||||
|
invariant before entering transfer code.
|
||||||
|
"""
|
||||||
|
|
||||||
|
if page_size <= 0:
|
||||||
|
raise ValueError(f"page_size must be positive, got {page_size}")
|
||||||
|
if indices.dim() != 1:
|
||||||
|
raise ValueError(
|
||||||
|
f"{name} must be a 1-D tensor, got shape={tuple(indices.shape)}"
|
||||||
|
)
|
||||||
|
if indices.numel() == 0:
|
||||||
|
return
|
||||||
|
if indices.numel() % page_size != 0:
|
||||||
|
raise ValueError(
|
||||||
|
f"{name} must contain whole pages: numel={indices.numel()} page_size={page_size}"
|
||||||
|
)
|
||||||
|
|
||||||
|
page_view = indices.reshape(-1, page_size)
|
||||||
|
starts = page_view[:, 0]
|
||||||
|
if torch.any(torch.remainder(starts, page_size) != 0):
|
||||||
|
raise ValueError(
|
||||||
|
f"{name} page groups must start at page boundaries: page_size={page_size}"
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_offsets = torch.arange(
|
||||||
|
page_size, device=indices.device, dtype=indices.dtype
|
||||||
|
)
|
||||||
|
expected = starts[:, None] + expected_offsets
|
||||||
|
if not torch.equal(page_view, expected):
|
||||||
|
raise ValueError(
|
||||||
|
f"{name} page groups must be contiguous page spans: page_size={page_size}"
|
||||||
|
)
|
||||||
@@ -798,6 +798,7 @@ class ServerArgs:
|
|||||||
|
|
||||||
# Handle Hicache settings.
|
# Handle Hicache settings.
|
||||||
self._handle_hicache()
|
self._handle_hicache()
|
||||||
|
self._handle_cp_hicache_layout_validation()
|
||||||
|
|
||||||
# Handle data parallelism.
|
# Handle data parallelism.
|
||||||
self._handle_data_parallelism()
|
self._handle_data_parallelism()
|
||||||
@@ -917,6 +918,33 @@ class ServerArgs:
|
|||||||
"Disable hicache_storage_backend or disable CP shared KV."
|
"Disable hicache_storage_backend or disable CP shared KV."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _handle_cp_hicache_layout_validation(self):
|
||||||
|
if not (
|
||||||
|
self.enable_nsa_prefill_cp_shared_kv and self.enable_hierarchical_cache
|
||||||
|
):
|
||||||
|
return
|
||||||
|
|
||||||
|
supported_pairs = {
|
||||||
|
("kernel", "layer_first"),
|
||||||
|
("kernel", "page_first"),
|
||||||
|
("direct", "layer_first"),
|
||||||
|
("direct", "page_first_direct"),
|
||||||
|
}
|
||||||
|
current_pair = (self.hicache_io_backend, self.hicache_mem_layout)
|
||||||
|
if current_pair in supported_pairs:
|
||||||
|
return
|
||||||
|
|
||||||
|
raise ValueError(
|
||||||
|
"CP shared KV HiCache supports only "
|
||||||
|
"kernel/layer_first, kernel/page_first, direct/layer_first, "
|
||||||
|
"and direct/page_first_direct after HiCache normalization. "
|
||||||
|
f"Got hicache_io_backend={self.hicache_io_backend!r} "
|
||||||
|
f"hicache_mem_layout={self.hicache_mem_layout!r}. "
|
||||||
|
"page_head is MHA/Mooncake-specific, page_first_kv_split is "
|
||||||
|
"kernel_ascend-specific, and kernel_ascend CP shared KV HiCache "
|
||||||
|
"is not implemented in the CUDA NSA host path."
|
||||||
|
)
|
||||||
|
|
||||||
def _handle_deprecated_args(self):
|
def _handle_deprecated_args(self):
|
||||||
# Handle deprecated tool call parsers
|
# Handle deprecated tool call parsers
|
||||||
deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"}
|
deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"}
|
||||||
|
|||||||
@@ -686,6 +686,38 @@ inline void transfer_page_direct(
|
|||||||
/* non_blocking= */ true);
|
/* non_blocking= */ true);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
inline void validate_page_first_direct_indices(
|
||||||
|
const int64_t* indices_ptr, int64_t num_indices, int64_t page_size, const char* name) {
|
||||||
|
if (num_indices == 0) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
for (int64_t i = 0; i < num_indices; i += page_size) {
|
||||||
|
const int64_t start = indices_ptr[i];
|
||||||
|
TORCH_CHECK(
|
||||||
|
start % page_size == 0,
|
||||||
|
name,
|
||||||
|
" page groups must start at page boundaries: group=",
|
||||||
|
i / page_size,
|
||||||
|
" start=",
|
||||||
|
start,
|
||||||
|
" page_size=",
|
||||||
|
page_size);
|
||||||
|
for (int64_t offset = 1; offset < page_size; ++offset) {
|
||||||
|
TORCH_CHECK(
|
||||||
|
indices_ptr[i + offset] == start + offset,
|
||||||
|
name,
|
||||||
|
" page groups must be contiguous page spans: group=",
|
||||||
|
i / page_size,
|
||||||
|
" offset=",
|
||||||
|
offset,
|
||||||
|
" expected=",
|
||||||
|
start + offset,
|
||||||
|
" got=",
|
||||||
|
indices_ptr[i + offset]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
void transfer_kv_direct(
|
void transfer_kv_direct(
|
||||||
const std::vector<at::Tensor>& src_layers,
|
const std::vector<at::Tensor>& src_layers,
|
||||||
std::vector<at::Tensor> dst_layers,
|
std::vector<at::Tensor> dst_layers,
|
||||||
@@ -749,6 +781,8 @@ inline void transfer_kv_page_first_direct_impl(
|
|||||||
const int64_t num_pages = src_indices_cpu.size(0) / page_size;
|
const int64_t num_pages = src_indices_cpu.size(0) / page_size;
|
||||||
int64_t* src_indices_ptr = src_indices_cpu.data_ptr<int64_t>();
|
int64_t* src_indices_ptr = src_indices_cpu.data_ptr<int64_t>();
|
||||||
int64_t* dst_indices_ptr = dst_indices_cpu.data_ptr<int64_t>();
|
int64_t* dst_indices_ptr = dst_indices_cpu.data_ptr<int64_t>();
|
||||||
|
validate_page_first_direct_indices(src_indices_ptr, src_indices_cpu.numel(), page_size, "src_indices");
|
||||||
|
validate_page_first_direct_indices(dst_indices_ptr, dst_indices_cpu.numel(), page_size, "dst_indices");
|
||||||
|
|
||||||
auto fallback_to_page_copy = [&]() {
|
auto fallback_to_page_copy = [&]() {
|
||||||
if constexpr (IsLf2Pf) {
|
if constexpr (IsLf2Pf) {
|
||||||
|
|||||||
@@ -33,6 +33,15 @@ def ref_copy_with_indices_pf_direct(
|
|||||||
][layer_id].to(dst_pool.device)
|
][layer_id].to(dst_pool.device)
|
||||||
|
|
||||||
|
|
||||||
|
def make_page_indices(page_ids, page_size):
|
||||||
|
return torch.cat(
|
||||||
|
[
|
||||||
|
torch.arange(p * page_size, (p + 1) * page_size, dtype=torch.int64)
|
||||||
|
for p in page_ids
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def ref_copy_with_indices_page_head(
|
def ref_copy_with_indices_page_head(
|
||||||
src_pool,
|
src_pool,
|
||||||
dst_pool,
|
dst_pool,
|
||||||
@@ -522,6 +531,54 @@ def test_transfer_kv_pf_direct(
|
|||||||
torch.set_default_dtype(original_dtype)
|
torch.set_default_dtype(original_dtype)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(is_hip(), reason="HIP uses the fallback path for this direct validation")
|
||||||
|
def test_transfer_kv_all_layer_direct_lf_pf_rejects_misaligned_page_start():
|
||||||
|
page_size = 4
|
||||||
|
item_size = 8
|
||||||
|
num_layers = 1
|
||||||
|
src_pool = torch.randn(num_layers, 32, item_size, device="cuda")
|
||||||
|
dst_pool = torch.zeros(8, num_layers, page_size, item_size).pin_memory()
|
||||||
|
src_indices = torch.tensor([1, 2, 3, 4], dtype=torch.int64)
|
||||||
|
dst_indices = make_page_indices([0], page_size)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="page boundaries"):
|
||||||
|
transfer_kv_all_layer_direct_lf_pf(
|
||||||
|
[src_pool[0]], [dst_pool], src_indices, dst_indices, page_size
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(is_hip(), reason="HIP uses the fallback path for this direct validation")
|
||||||
|
def test_transfer_kv_all_layer_direct_lf_pf_rejects_non_contiguous_page_group():
|
||||||
|
page_size = 4
|
||||||
|
item_size = 8
|
||||||
|
num_layers = 1
|
||||||
|
src_pool = torch.randn(num_layers, 32, item_size, device="cuda")
|
||||||
|
dst_pool = torch.zeros(8, num_layers, page_size, item_size).pin_memory()
|
||||||
|
src_indices = torch.tensor([4, 5, 7, 6], dtype=torch.int64)
|
||||||
|
dst_indices = make_page_indices([0], page_size)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="contiguous page spans"):
|
||||||
|
transfer_kv_all_layer_direct_lf_pf(
|
||||||
|
[src_pool[0]], [dst_pool], src_indices, dst_indices, page_size
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(is_hip(), reason="HIP uses the fallback path for this direct validation")
|
||||||
|
def test_transfer_kv_per_layer_direct_pf_lf_rejects_misaligned_host_page_start():
|
||||||
|
page_size = 4
|
||||||
|
item_size = 8
|
||||||
|
layer_id = 0
|
||||||
|
src_pool = torch.randn(8, 1, page_size, item_size).pin_memory()
|
||||||
|
dst_pool = torch.zeros(1, 32, item_size, device="cuda")
|
||||||
|
src_indices = torch.tensor([1, 2, 3, 4], dtype=torch.int64)
|
||||||
|
dst_indices = make_page_indices([0], page_size).to("cuda")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="page boundaries"):
|
||||||
|
transfer_kv_per_layer_direct_pf_lf(
|
||||||
|
[src_pool], [dst_pool[0]], src_indices, dst_indices, layer_id, page_size
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.skipif(is_hip(), reason="HIP is not supported for this test")
|
@pytest.mark.skipif(is_hip(), reason="HIP is not supported for this test")
|
||||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||||
@pytest.mark.parametrize("num_items_to_transfer", [256, 1024])
|
@pytest.mark.parametrize("num_items_to_transfer", [256, 1024])
|
||||||
|
|||||||
@@ -195,6 +195,32 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
|
|||||||
self.assertEqual(host_pool.alloc_calls, [4])
|
self.assertEqual(host_pool.alloc_calls, [4])
|
||||||
self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
|
self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
|
||||||
|
|
||||||
|
def test_cp_write_rejects_incomplete_owned_physical_page(self):
|
||||||
|
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
|
||||||
|
controller = self.make_controller(host_pool, cp_rank=1)
|
||||||
|
logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "physical_device_indices.*whole pages"):
|
||||||
|
controller.write(logical_locs, node_id=21)
|
||||||
|
|
||||||
|
def test_cp_write_rejects_non_contiguous_owned_physical_page(self):
|
||||||
|
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
|
||||||
|
controller = self.make_controller(host_pool, cp_rank=1)
|
||||||
|
logical_locs = torch.tensor([8, 9, 11, 10], dtype=torch.int64)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
ValueError, "physical_device_indices.*contiguous page spans"
|
||||||
|
):
|
||||||
|
controller.write(logical_locs, node_id=22)
|
||||||
|
|
||||||
|
def test_cp_write_rejects_non_contiguous_host_page(self):
|
||||||
|
host_pool = FakeHostPool(torch.tensor([100, 101, 103, 102], dtype=torch.int64))
|
||||||
|
controller = self.make_controller(host_pool, cp_rank=1)
|
||||||
|
logical_locs = torch.arange(8, 12, dtype=torch.int64)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
|
||||||
|
controller.write(logical_locs, node_id=23)
|
||||||
|
|
||||||
def test_cp_write_zero_owned_returns_metadata_and_noop_ack(self):
|
def test_cp_write_zero_owned_returns_metadata_and_noop_ack(self):
|
||||||
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
|
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
|
||||||
controller = self.make_controller(host_pool, cp_rank=3)
|
controller = self.make_controller(host_pool, cp_rank=3)
|
||||||
@@ -315,6 +341,43 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
|
|||||||
self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103])
|
self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103])
|
||||||
self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23])
|
self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23])
|
||||||
|
|
||||||
|
def test_cp_load_rejects_non_contiguous_physical_device_page(self):
|
||||||
|
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
|
||||||
|
allocator = FakeAllocator(
|
||||||
|
alloc_result=torch.tensor(
|
||||||
|
[64, 65, 66, 67, 68, 69, 71, 70, 72, 73, 74, 75, 76, 77, 78, 79],
|
||||||
|
dtype=torch.int64,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
|
||||||
|
node = TreeNode()
|
||||||
|
node.host_len = 16
|
||||||
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
||||||
|
logical_len=16,
|
||||||
|
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
|
||||||
|
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
ValueError, "physical_device_indices.*contiguous page spans"
|
||||||
|
):
|
||||||
|
controller.load_cp([node], node_id=31)
|
||||||
|
|
||||||
|
def test_cp_load_rejects_non_contiguous_host_page(self):
|
||||||
|
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
|
||||||
|
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
|
||||||
|
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
|
||||||
|
node = TreeNode()
|
||||||
|
node.host_len = 16
|
||||||
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
||||||
|
logical_len=16,
|
||||||
|
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
|
||||||
|
host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
|
||||||
|
controller.load_cp([node], node_id=32)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
|
|||||||
@@ -10,11 +10,12 @@ from sglang.srt.mem_cache.memory_pool_host import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_cuda_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cuda_ci(est_time=3, suite="stage-b-test-1-gpu-small")
|
register_cuda_ci(est_time=3, suite="stage-b-test-1-gpu-small")
|
||||||
|
|
||||||
|
|
||||||
class TestNSAHiCacheTransfer(unittest.TestCase):
|
class TestNSAHiCacheTransfer(CustomTestCase):
|
||||||
def setUp(self):
|
def setUp(self):
|
||||||
if not torch.cuda.is_available():
|
if not torch.cuda.is_available():
|
||||||
self.skipTest("CUDA is required for NSA host transfer tests.")
|
self.skipTest("CUDA is required for NSA host transfer tests.")
|
||||||
@@ -36,7 +37,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
|
|||||||
]
|
]
|
||||||
return torch.cat(parts, dim=0)
|
return torch.cat(parts, dim=0)
|
||||||
|
|
||||||
def _run_device_to_host_indexer_copy(self, io_backend: str):
|
def _run_device_to_host_indexer_copy(self, io_backend: str, layout: str):
|
||||||
page_size = 1 if is_hip() else 64
|
page_size = 1 if is_hip() else 64
|
||||||
layer_num = 2
|
layer_num = 2
|
||||||
size = page_size * 4
|
size = page_size * 4
|
||||||
@@ -63,7 +64,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
|
|||||||
host_to_device_ratio=2.0,
|
host_to_device_ratio=2.0,
|
||||||
host_size=0,
|
host_size=0,
|
||||||
page_size=page_size,
|
page_size=page_size,
|
||||||
layout="layer_first",
|
layout=layout,
|
||||||
pin_memory=pin_memory,
|
pin_memory=pin_memory,
|
||||||
device="cpu",
|
device="cpu",
|
||||||
)
|
)
|
||||||
@@ -105,26 +106,88 @@ class TestNSAHiCacheTransfer(unittest.TestCase):
|
|||||||
for host_page, device_page in zip(
|
for host_page, device_page in zip(
|
||||||
host_pages.tolist(), device_pages.tolist()
|
host_pages.tolist(), device_pages.tolist()
|
||||||
):
|
):
|
||||||
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
|
if layout == "page_first_direct":
|
||||||
|
got = host_pool.index_k_with_scale_buffer[host_page, layer_id].cpu()
|
||||||
|
else:
|
||||||
|
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
|
||||||
expected = device_pool.index_k_with_scale_buffer[layer_id][
|
expected = device_pool.index_k_with_scale_buffer[layer_id][
|
||||||
device_page
|
device_page
|
||||||
].cpu()
|
].cpu()
|
||||||
self.assertTrue(torch.equal(got, expected))
|
self.assertTrue(torch.equal(got, expected))
|
||||||
host_start = host_page * page_size
|
|
||||||
device_start = device_page * page_size
|
device_start = device_page * page_size
|
||||||
got_kv = host_pool.kv_buffer[layer_id][
|
|
||||||
host_start : host_start + page_size
|
|
||||||
].cpu()
|
|
||||||
expected_kv = device_pool.kv_buffer[layer_id][
|
expected_kv = device_pool.kv_buffer[layer_id][
|
||||||
device_start : device_start + page_size
|
device_start : device_start + page_size
|
||||||
].cpu()
|
].cpu()
|
||||||
|
if layout == "page_first_direct":
|
||||||
|
got_kv = host_pool.kv_buffer[host_page, layer_id].cpu()
|
||||||
|
else:
|
||||||
|
host_start = host_page * page_size
|
||||||
|
got_kv = host_pool.kv_buffer[layer_id][
|
||||||
|
host_start : host_start + page_size
|
||||||
|
].cpu()
|
||||||
self.assertTrue(torch.equal(got_kv, expected_kv))
|
self.assertTrue(torch.equal(got_kv, expected_kv))
|
||||||
|
|
||||||
def test_device_to_host_indexer_kernel(self):
|
def test_device_to_host_indexer_kernel_layer_first(self):
|
||||||
self._run_device_to_host_indexer_copy(io_backend="kernel")
|
self._run_device_to_host_indexer_copy(
|
||||||
|
io_backend="kernel", layout="layer_first"
|
||||||
|
)
|
||||||
|
|
||||||
def test_device_to_host_indexer_direct(self):
|
def test_device_to_host_indexer_direct_layer_first(self):
|
||||||
self._run_device_to_host_indexer_copy(io_backend="direct")
|
self._run_device_to_host_indexer_copy(
|
||||||
|
io_backend="direct", layout="layer_first"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_device_to_host_indexer_direct_page_first_direct(self):
|
||||||
|
self._run_device_to_host_indexer_copy(
|
||||||
|
io_backend="direct", layout="page_first_direct"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestNSAIndexerPageIndices(CustomTestCase):
|
||||||
|
def make_host_pool_stub(self, page_size: int):
|
||||||
|
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
|
||||||
|
host_pool.page_size = page_size
|
||||||
|
return host_pool
|
||||||
|
|
||||||
|
def test_indexer_page_indices_accepts_valid_page_spans(self):
|
||||||
|
host_pool = self.make_host_pool_stub(page_size=4)
|
||||||
|
|
||||||
|
host_pages, device_pages = host_pool._get_indexer_page_indices(
|
||||||
|
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
|
||||||
|
torch.tensor([16, 17, 18, 19], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(host_pages.tolist(), [2])
|
||||||
|
self.assertEqual(device_pages.tolist(), [4])
|
||||||
|
|
||||||
|
def test_indexer_page_indices_rejects_partial_host_page(self):
|
||||||
|
host_pool = self.make_host_pool_stub(page_size=4)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "host_indices.*whole pages"):
|
||||||
|
host_pool._get_indexer_page_indices(
|
||||||
|
torch.tensor([8, 9, 10], dtype=torch.int64),
|
||||||
|
torch.tensor([16, 17, 18], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_indexer_page_indices_rejects_misaligned_device_page(self):
|
||||||
|
host_pool = self.make_host_pool_stub(page_size=4)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
ValueError, "device_indices.*start at page boundaries"
|
||||||
|
):
|
||||||
|
host_pool._get_indexer_page_indices(
|
||||||
|
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
|
||||||
|
torch.tensor([17, 18, 19, 20], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_indexer_page_indices_rejects_non_contiguous_host_page(self):
|
||||||
|
host_pool = self.make_host_pool_stub(page_size=4)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
|
||||||
|
host_pool._get_indexer_page_indices(
|
||||||
|
torch.tensor([8, 9, 11, 10], dtype=torch.int64),
|
||||||
|
torch.tensor([16, 17, 18, 19], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
64
test/registered/unit/mem_cache/test_page_index_utils.py
Normal file
64
test/registered/unit/mem_cache/test_page_index_utils.py
Normal file
@@ -0,0 +1,64 @@
|
|||||||
|
import unittest
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.mem_cache.page_index_utils import (
|
||||||
|
validate_page_aligned_token_indices,
|
||||||
|
)
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
class TestPageIndexUtils(CustomTestCase):
|
||||||
|
def test_valid_empty_indices(self):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.empty((0,), dtype=torch.int64), 4, "indices"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_valid_single_page(self):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.tensor([8, 9, 10, 11], dtype=torch.int64), 4, "indices"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_valid_multiple_pages(self):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.tensor([4, 5, 6, 7, 16, 17, 18, 19], dtype=torch.int64),
|
||||||
|
4,
|
||||||
|
"indices",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_rejects_non_1d_tensor(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "must be a 1-D tensor"):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.tensor([[0, 1], [2, 3]], dtype=torch.int64), 4, "indices"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_rejects_non_positive_page_size(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "page_size must be positive"):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.tensor([0, 1, 2, 3], dtype=torch.int64), 0, "indices"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_rejects_partial_page(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "whole pages"):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.tensor([8, 9, 10], dtype=torch.int64), 4, "indices"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_rejects_non_page_start(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "start at page boundaries"):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.tensor([9, 10, 11, 12], dtype=torch.int64), 4, "indices"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_rejects_non_contiguous_group(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "contiguous page spans"):
|
||||||
|
validate_page_aligned_token_indices(
|
||||||
|
torch.tensor([8, 9, 11, 10], dtype=torch.int64), 4, "indices"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -371,6 +371,26 @@ class TestHiCacheArgs(CustomTestCase):
|
|||||||
if expected_decode_backend is not None:
|
if expected_decode_backend is not None:
|
||||||
self.assertEqual(args.decode_attention_backend, expected_decode_backend)
|
self.assertEqual(args.decode_attention_backend, expected_decode_backend)
|
||||||
|
|
||||||
|
def _make_cp_hicache_args(self, **overrides) -> ServerArgs:
|
||||||
|
args = self._make_args(
|
||||||
|
enable_hierarchical_cache=True,
|
||||||
|
enable_nsa_prefill_context_parallel=True,
|
||||||
|
enable_nsa_prefill_cp_shared_kv=True,
|
||||||
|
nsa_prefill_cp_mode="in-seq-split",
|
||||||
|
disaggregation_mode="prefill",
|
||||||
|
page_size=64,
|
||||||
|
enable_hisparse=False,
|
||||||
|
)
|
||||||
|
for key, value in overrides.items():
|
||||||
|
setattr(args, key, value)
|
||||||
|
return args
|
||||||
|
|
||||||
|
def _normalize_and_validate_cp_hicache_args(self, **overrides) -> ServerArgs:
|
||||||
|
args = self._make_cp_hicache_args(**overrides)
|
||||||
|
args._handle_hicache()
|
||||||
|
args._handle_cp_hicache_layout_validation()
|
||||||
|
return args
|
||||||
|
|
||||||
def test_cp_shared_kv_rejects_hicache_storage_backend(self):
|
def test_cp_shared_kv_rejects_hicache_storage_backend(self):
|
||||||
with (
|
with (
|
||||||
self.assertRaisesRegex(
|
self.assertRaisesRegex(
|
||||||
@@ -468,6 +488,61 @@ class TestHiCacheArgs(CustomTestCase):
|
|||||||
|
|
||||||
self.assertEqual(args.decode_attention_backend, "triton")
|
self.assertEqual(args.decode_attention_backend, "triton")
|
||||||
|
|
||||||
|
def test_cp_hicache_accepts_supported_backend_layout_pairs(self):
|
||||||
|
cases = [
|
||||||
|
("kernel", "layer_first"),
|
||||||
|
("kernel", "page_first"),
|
||||||
|
("direct", "layer_first"),
|
||||||
|
("direct", "page_first_direct"),
|
||||||
|
]
|
||||||
|
|
||||||
|
for io_backend, mem_layout in cases:
|
||||||
|
with self.subTest(io_backend=io_backend, mem_layout=mem_layout):
|
||||||
|
args = self._normalize_and_validate_cp_hicache_args(
|
||||||
|
hicache_io_backend=io_backend,
|
||||||
|
hicache_mem_layout=mem_layout,
|
||||||
|
)
|
||||||
|
self.assertEqual(args.hicache_io_backend, io_backend)
|
||||||
|
self.assertEqual(args.hicache_mem_layout, mem_layout)
|
||||||
|
|
||||||
|
def test_cp_hicache_normalizes_supported_alias_pairs(self):
|
||||||
|
args = self._normalize_and_validate_cp_hicache_args(
|
||||||
|
hicache_io_backend="kernel",
|
||||||
|
hicache_mem_layout="page_first_direct",
|
||||||
|
)
|
||||||
|
self.assertEqual(args.hicache_io_backend, "direct")
|
||||||
|
self.assertEqual(args.hicache_mem_layout, "page_first_direct")
|
||||||
|
|
||||||
|
args = self._normalize_and_validate_cp_hicache_args(
|
||||||
|
hicache_io_backend="direct",
|
||||||
|
hicache_mem_layout="page_first",
|
||||||
|
)
|
||||||
|
self.assertEqual(args.hicache_io_backend, "direct")
|
||||||
|
self.assertEqual(args.hicache_mem_layout, "page_first_direct")
|
||||||
|
|
||||||
|
def test_cp_hicache_rejects_page_head_layout(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "CP shared KV HiCache.*page_head"):
|
||||||
|
self._normalize_and_validate_cp_hicache_args(
|
||||||
|
hicache_io_backend="kernel",
|
||||||
|
hicache_mem_layout="page_head",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_cp_hicache_rejects_cuda_kv_split_layout(self):
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
ValueError, "CP shared KV HiCache.*page_first_kv_split"
|
||||||
|
):
|
||||||
|
self._normalize_and_validate_cp_hicache_args(
|
||||||
|
hicache_io_backend="kernel",
|
||||||
|
hicache_mem_layout="page_first_kv_split",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_cp_hicache_rejects_kernel_ascend_backend(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "CP shared KV HiCache.*kernel_ascend"):
|
||||||
|
self._normalize_and_validate_cp_hicache_args(
|
||||||
|
hicache_io_backend="kernel_ascend",
|
||||||
|
hicache_mem_layout="page_first_kv_split",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user