CP HiCache: multi-slab host KV cache to respect cudaHostRegister per-call ceiling
A single cudaHostRegister over a >~1 TB host buffer fails with cudaErrorMemoryAllocation on B300 (hard per-call ceiling: 512 GiB OK, 1024 GiB FAIL), crashing the hicache_size=1600 (~1.5 TB CP shared-L2 slab) prefill at startup. The registration cannot simply be chunked: a memcpy (cudaMemcpyBatchAsync, the CP-L2 H2D/D2H transfer) fails with cudaErrorInvalidValue when its host range straddles a registration boundary (verified empirically on b300-049). Fix: physically split the host cache into multiple page-aligned slabs, each <= a safe single-registration size (default 480 GiB, env SGLANG_CP_HICACHE_MAX_SLAB_GB), reusing the existing SharedHostTensorGroupAllocator + per-slab transfer splitting (_host_transfer_segments). Each slab is one whole registration and no transfer crosses a boundary; small configs (hicache_size<=400) stay single-slab unchanged. - memory_pool_host.py: add cp_hicache_max_single_register_bytes() + the fail-loud _check_single_cuda_host_register_size guard; revert the (transfer-unsafe) registration chunking back to one cudaHostRegister per buffer/slab. - hiradix_cache.py: _cp_shared_l2_slab_pages_by_payload auto-caps each payload's slab <= the ceiling so large caches auto-split. - cp_l3_slab_accessor.py: CpSharedL2SlabAccessor is now slab-count-aware (CpL3SlabSpan + per-slab dispatch via global_base_page; per-slab layer stride); _cp_l3_slab_spans rewires _maybe_init_cp_l3 off the single-slab assumption. L3 disk slabs / slot pool / LMDB index / GC are content-addressed and unchanged. Tests: multi-slab accessor incl. a non-circular torch.frombuffer-layout check; a real allocate_group + _cp_l3_slab_spans roundtrip; slab-cap auto-split; L3 store cross-slab spill/reload. Reviewed by 3 adversarial agents, no correctness bugs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -7,6 +7,8 @@ import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
_MEM = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache"
|
||||
|
||||
@@ -52,10 +54,28 @@ def _make_slab():
|
||||
return mm, lo
|
||||
|
||||
|
||||
def _make_multislab(page_counts, *, fill=True):
|
||||
"""Build len(page_counts) slabs tiling [0, sum(page_counts)) in layer_page_first
|
||||
layout (per-slab layer stride = num_pages * SLICE), returning CpL3SlabSpan list."""
|
||||
spans = []
|
||||
base = 0
|
||||
for num in page_counts:
|
||||
mm = mmap.mmap(-1, N_LAYERS * num * SLICE)
|
||||
if fill:
|
||||
per_slab_layer_stride = num * SLICE
|
||||
for layer in range(N_LAYERS):
|
||||
for local in range(num):
|
||||
off = layer * per_slab_layer_stride + local * SLICE
|
||||
mm[off:off + SLICE] = _slice_pattern(layer, base + local)
|
||||
spans.append(a.CpL3SlabSpan(base, num, mm))
|
||||
base += num
|
||||
return spans
|
||||
|
||||
|
||||
class TestAccessor(unittest.TestCase):
|
||||
def test_gather_concatenates_layers_in_order(self):
|
||||
mm, lo = _make_slab()
|
||||
acc = a.CpSharedL2SlabAccessor(mm, lo)
|
||||
acc = a.CpSharedL2SlabAccessor.from_single_mmap(mm, lo)
|
||||
self.assertEqual(acc.page_blob_bytes, N_LAYERS * SLICE)
|
||||
for page in range(PAGE_NUM):
|
||||
expect = b"".join(_slice_pattern(layer, page) for layer in range(N_LAYERS))
|
||||
@@ -63,7 +83,7 @@ class TestAccessor(unittest.TestCase):
|
||||
|
||||
def test_gather_into_with_offset(self):
|
||||
mm, lo = _make_slab()
|
||||
acc = a.CpSharedL2SlabAccessor(mm, lo)
|
||||
acc = a.CpSharedL2SlabAccessor.from_single_mmap(mm, lo)
|
||||
dst = bytearray(64 + lo.page_blob_bytes)
|
||||
n = acc.gather_into(1, dst, 64)
|
||||
self.assertEqual(n, lo.page_blob_bytes)
|
||||
@@ -73,9 +93,9 @@ class TestAccessor(unittest.TestCase):
|
||||
|
||||
def test_gather_scatter_roundtrip(self):
|
||||
src_mm, lo = _make_slab()
|
||||
acc_src = a.CpSharedL2SlabAccessor(src_mm, lo)
|
||||
acc_src = a.CpSharedL2SlabAccessor.from_single_mmap(src_mm, lo)
|
||||
dst_mm = mmap.mmap(-1, lo.total_bytes)
|
||||
acc_dst = a.CpSharedL2SlabAccessor(dst_mm, lo)
|
||||
acc_dst = a.CpSharedL2SlabAccessor.from_single_mmap(dst_mm, lo)
|
||||
for page in range(PAGE_NUM):
|
||||
blob = acc_src.gather(page)
|
||||
acc_dst.scatter_from(page, blob, 0)
|
||||
@@ -86,16 +106,16 @@ class TestAccessor(unittest.TestCase):
|
||||
|
||||
def test_scatter_from_offset(self):
|
||||
src_mm, lo = _make_slab()
|
||||
acc_src = a.CpSharedL2SlabAccessor(src_mm, lo)
|
||||
acc_src = a.CpSharedL2SlabAccessor.from_single_mmap(src_mm, lo)
|
||||
dst_mm = mmap.mmap(-1, lo.total_bytes)
|
||||
acc_dst = a.CpSharedL2SlabAccessor(dst_mm, lo)
|
||||
acc_dst = a.CpSharedL2SlabAccessor.from_single_mmap(dst_mm, lo)
|
||||
framed = bytearray(64) + bytearray(acc_src.gather(2)) # simulate header+payload slot buffer
|
||||
acc_dst.scatter_from(2, framed, 64)
|
||||
self.assertEqual(acc_dst.gather(2), acc_src.gather(2))
|
||||
|
||||
def test_bounds_fail_loud(self):
|
||||
mm, lo = _make_slab()
|
||||
acc = a.CpSharedL2SlabAccessor(mm, lo)
|
||||
acc = a.CpSharedL2SlabAccessor.from_single_mmap(mm, lo)
|
||||
with self.assertRaises(ValueError):
|
||||
acc.gather(PAGE_NUM) # page out of range
|
||||
with self.assertRaises(ValueError):
|
||||
@@ -107,7 +127,78 @@ class TestAccessor(unittest.TestCase):
|
||||
lo = a.CpL3SlabLayout(n_layers=N_LAYERS, page_num=PAGE_NUM, slice_bytes=SLICE)
|
||||
small = mmap.mmap(-1, lo.total_bytes - 16)
|
||||
with self.assertRaises(ValueError):
|
||||
a.CpSharedL2SlabAccessor(small, lo)
|
||||
a.CpSharedL2SlabAccessor.from_single_mmap(small, lo)
|
||||
|
||||
# --- multi-slab (host cache split across N physical slabs) ---
|
||||
|
||||
def test_multislab_gather_dispatches_per_slab(self):
|
||||
spans = _make_multislab([2, 2]) # pages 0-1 in slab0, pages 2-3 in slab1 (global_base 2)
|
||||
acc = a.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE)
|
||||
self.assertEqual(acc.total_pages, PAGE_NUM)
|
||||
for page in range(PAGE_NUM):
|
||||
expect = b"".join(_slice_pattern(layer, page) for layer in range(N_LAYERS))
|
||||
self.assertEqual(acc.gather(page), expect) # slab1 pages addressed via global_base + per-slab stride
|
||||
|
||||
def test_multislab_roundtrip_uneven_split(self):
|
||||
src = _make_multislab([1, 3]) # slab1 base=1, per-slab layer strides differ (1*SLICE vs 3*SLICE)
|
||||
dst = _make_multislab([1, 3], fill=False)
|
||||
acc_src = a.CpSharedL2SlabAccessor(src, n_layers=N_LAYERS, slice_bytes=SLICE)
|
||||
acc_dst = a.CpSharedL2SlabAccessor(dst, n_layers=N_LAYERS, slice_bytes=SLICE)
|
||||
for page in range(PAGE_NUM):
|
||||
acc_dst.scatter_from(page, acc_src.gather(page), 0)
|
||||
for page in range(PAGE_NUM):
|
||||
self.assertEqual(acc_dst.gather(page), acc_src.gather(page))
|
||||
|
||||
def test_multislab_matches_single_slab_bytes(self):
|
||||
# A page in slab1 must gather identically whether the payload is one slab or two.
|
||||
single_mm, lo = _make_slab()
|
||||
acc_single = a.CpSharedL2SlabAccessor.from_single_mmap(single_mm, lo)
|
||||
acc_multi = a.CpSharedL2SlabAccessor(
|
||||
_make_multislab([2, 2]), n_layers=N_LAYERS, slice_bytes=SLICE
|
||||
)
|
||||
for page in range(PAGE_NUM):
|
||||
self.assertEqual(acc_multi.gather(page), acc_single.gather(page))
|
||||
|
||||
def test_multislab_noncontiguous_spans_fail(self):
|
||||
mm0 = mmap.mmap(-1, N_LAYERS * 2 * SLICE)
|
||||
mm1 = mmap.mmap(-1, N_LAYERS * 2 * SLICE)
|
||||
with self.assertRaises(ValueError): # gap: slab1 base 3 != 0+2
|
||||
a.CpSharedL2SlabAccessor(
|
||||
[a.CpL3SlabSpan(0, 2, mm0), a.CpL3SlabSpan(3, 2, mm1)],
|
||||
n_layers=N_LAYERS,
|
||||
slice_bytes=SLICE,
|
||||
)
|
||||
|
||||
def test_multislab_slab_mmap_too_small_fails(self):
|
||||
spans = _make_multislab([2, 2])
|
||||
spans[1] = a.CpL3SlabSpan(2, 2, mmap.mmap(-1, N_LAYERS * 2 * SLICE - 8))
|
||||
with self.assertRaises(ValueError):
|
||||
a.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE)
|
||||
|
||||
def test_accessor_stride_matches_real_torch_frombuffer_view(self):
|
||||
"""Non-circular layout check: write each slab's pages through a real
|
||||
torch.frombuffer(mmap).view(n_layers, num_pages, slice) -- exactly how the
|
||||
production per-slab tensor is C-contiguously laid out (page_dim sized to the
|
||||
SLAB's num_pages) -- and read them back through the accessor. A bug where
|
||||
the accessor used the TOTAL page_num (instead of per-slab num_pages) for the
|
||||
layer stride would read the wrong bytes here and fail."""
|
||||
page_counts = [2, 3] # uneven -> per-slab layer strides differ (2*SLICE vs 3*SLICE)
|
||||
spans, base = [], 0
|
||||
for num in page_counts:
|
||||
mm = mmap.mmap(-1, N_LAYERS * num * SLICE)
|
||||
# Write via the production-style C-contiguous view, NOT the accessor's formula.
|
||||
t = torch.frombuffer(mm, dtype=torch.uint8).view(N_LAYERS, num, SLICE)
|
||||
for layer in range(N_LAYERS):
|
||||
for local in range(num):
|
||||
t[layer, local] = torch.frombuffer(
|
||||
bytearray(_slice_pattern(layer, base + local)), dtype=torch.uint8
|
||||
)
|
||||
spans.append(a.CpL3SlabSpan(base, num, mm))
|
||||
base += num
|
||||
acc = a.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE)
|
||||
for page in range(sum(page_counts)):
|
||||
expect = b"".join(_slice_pattern(layer, page) for layer in range(N_LAYERS))
|
||||
self.assertEqual(acc.gather(page), expect)
|
||||
|
||||
def test_factory_mla_dims(self):
|
||||
lo = a.CpL3SlabLayout.for_mla(layer_num=78, page_num=100, page_size=64, kv_cache_dim=576, itemsize=1)
|
||||
|
||||
@@ -65,7 +65,24 @@ def _make_slab_and_accessor():
|
||||
for page in range(PAGE_NUM):
|
||||
off = layer * lo.layer_stride_bytes + page * SLICE
|
||||
mm[off:off + SLICE] = _pattern(layer, page)
|
||||
return mm, acc_mod.CpSharedL2SlabAccessor(mm, lo)
|
||||
return mm, acc_mod.CpSharedL2SlabAccessor.from_single_mmap(mm, lo)
|
||||
|
||||
|
||||
def _make_multislab_accessor(page_counts):
|
||||
"""Build a multi-slab target_kv accessor tiling [0, sum(page_counts)) -- each slab
|
||||
its own mmap with per-slab layer stride. Returns ([(base,num,mm)...], accessor)."""
|
||||
spans, slabs, base = [], [], 0
|
||||
for num in page_counts:
|
||||
mm = mmap.mmap(-1, N_LAYERS * num * SLICE)
|
||||
stride = num * SLICE
|
||||
for layer in range(N_LAYERS):
|
||||
for local in range(num):
|
||||
off = layer * stride + local * SLICE
|
||||
mm[off:off + SLICE] = _pattern(layer, base + local)
|
||||
spans.append(acc_mod.CpL3SlabSpan(base, num, mm))
|
||||
slabs.append((base, num, mm))
|
||||
base += num
|
||||
return slabs, acc_mod.CpSharedL2SlabAccessor(spans, n_layers=N_LAYERS, slice_bytes=SLICE)
|
||||
|
||||
|
||||
def _wait_ack(qsize_fn, n=1, timeout=5.0):
|
||||
@@ -301,5 +318,53 @@ class TestCpL3Store(unittest.TestCase):
|
||||
self.assertFalse(self.store.has_inflight())
|
||||
|
||||
|
||||
class TestCpL3StoreMultiSlab(unittest.TestCase):
|
||||
"""The store must work when the host cache is split across multiple physical slabs:
|
||||
the accessor dispatches each global page to its owning slab (per-slab layer stride)."""
|
||||
|
||||
def setUp(self):
|
||||
self._td = tempfile.TemporaryDirectory()
|
||||
self.slabs, self.acc = _make_multislab_accessor([4, 4]) # pages 0-3 | 4-7
|
||||
cfg = cfg_mod.CpL3Config.from_dict({
|
||||
"backend": "posix",
|
||||
"require_plp": False,
|
||||
"index_map_gb": 0.05,
|
||||
"disks": [{"path": os.path.join(self._td.name, "disk0"), "budget_gb": 0.02}],
|
||||
})
|
||||
self.store = store_mod.CpL3Store.from_config(
|
||||
cfg, cp_rank=0, cp_size=1, accessors={"target_kv": self.acc})
|
||||
self.store.connect(cfg)
|
||||
|
||||
def tearDown(self):
|
||||
self.store.close()
|
||||
self._td.cleanup()
|
||||
|
||||
def _zero_page(self, global_page):
|
||||
for base, num, mm in self.slabs:
|
||||
if base <= global_page < base + num:
|
||||
local, stride = global_page - base, num * SLICE
|
||||
for layer in range(N_LAYERS):
|
||||
off = layer * stride + local * SLICE
|
||||
mm[off:off + SLICE] = bytes(SLICE)
|
||||
return
|
||||
|
||||
def test_spill_reload_roundtrip_page_in_second_slab(self):
|
||||
page = 5 # slab1: global_base 4, local 1
|
||||
orig = self.acc.gather(page)
|
||||
pages = {"target_kv": [(page, _h(page))]}
|
||||
op_id = self.store.submit_spill("obj-ms", pages, last_access=1)
|
||||
self.assertTrue(_wait_ack(self.store.ack_durable_qsize, 1))
|
||||
self.store.drain_gather_acks(self.store.ack_gather_qsize())
|
||||
self.assertEqual(
|
||||
self.store.drain_durable_acks(self.store.ack_durable_qsize()), [(op_id, True)]
|
||||
)
|
||||
self._zero_page(page)
|
||||
self.assertNotEqual(self.acc.gather(page), orig)
|
||||
self.store.submit_reload("obj-ms", pages)
|
||||
self.assertTrue(_wait_ack(self.store.ack_reload_qsize, 1))
|
||||
self.store.drain_reload_acks(self.store.ack_reload_qsize())
|
||||
self.assertEqual(self.acc.gather(page), orig) # byte-exact reload into slab1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
159
test/registered/unit/mem_cache/test_cp_shared_l2_slab_cap.py
Normal file
159
test/registered/unit/mem_cache/test_cp_shared_l2_slab_cap.py
Normal file
@@ -0,0 +1,159 @@
|
||||
"""Unit tests for the CP shared-L2 per-payload slab-size cap.
|
||||
|
||||
A single cudaHostRegister cannot exceed a per-call ceiling (and cannot be
|
||||
chunked), so large host caches are split into multiple slabs each <= the
|
||||
ceiling. `_cp_shared_l2_slab_pages_by_payload` derives the per-payload slab page
|
||||
count; this verifies the auto-cap, env override, explicit-size capping, and the
|
||||
one-page-too-big fail-loud.
|
||||
"""
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.cp_l3_slab_accessor import CpSharedL2SlabAccessor
|
||||
from sglang.srt.mem_cache.cp_shared_l2_pool import (
|
||||
PAYLOAD_TARGET_KV,
|
||||
CpSharedL2SlabInfo,
|
||||
build_cp_shared_l2_slabs_by_payload,
|
||||
)
|
||||
from sglang.srt.mem_cache.hiradix_cache import (
|
||||
HiRadixCache,
|
||||
_cp_shared_l2_slab_pages_by_payload,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool_host import (
|
||||
SharedHostTensorAllocator,
|
||||
SharedHostTensorGroupAllocator,
|
||||
cp_hicache_max_single_register_bytes,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# hiradix_cache imports sgl_kernel (needs libcuda to import); no GPU is used.
|
||||
register_cuda_ci(est_time=1, suite="stage-b-test-1-gpu-small")
|
||||
|
||||
GiB = 1024**3
|
||||
|
||||
|
||||
def _args(slab_size_gb=0):
|
||||
return SimpleNamespace(cp_shared_l2_slab_size_gb=slab_size_gb)
|
||||
|
||||
|
||||
def _slab_pages(server_args, bytes_per_page_by_payload):
|
||||
return _cp_shared_l2_slab_pages_by_payload(
|
||||
server_args=server_args,
|
||||
page_size=64,
|
||||
bytes_per_page_by_payload=bytes_per_page_by_payload,
|
||||
)
|
||||
|
||||
|
||||
class TestCpSharedL2SlabCap(CustomTestCase):
|
||||
def test_auto_caps_each_payload_at_ceiling(self):
|
||||
ceiling = cp_hicache_max_single_register_bytes()
|
||||
out = _slab_pages(_args(0), {"target_kv": 10 * GiB, "index_k": 4 * GiB})
|
||||
self.assertEqual(out["target_kv"], ceiling // (10 * GiB))
|
||||
self.assertEqual(out["index_k"], ceiling // (4 * GiB))
|
||||
# Each slab's registration bytes stay <= the ceiling.
|
||||
self.assertLessEqual(out["target_kv"] * 10 * GiB, ceiling)
|
||||
|
||||
def test_env_override_changes_cap(self):
|
||||
with patch.dict(os.environ, {"SGLANG_CP_HICACHE_MAX_SLAB_GB": "256"}):
|
||||
out = _slab_pages(_args(0), {"target_kv": 8 * GiB})
|
||||
self.assertEqual(out["target_kv"], (256 * GiB) // (8 * GiB)) # 32
|
||||
|
||||
def test_explicit_slab_size_capped_at_ceiling(self):
|
||||
ceiling = cp_hicache_max_single_register_bytes()
|
||||
# 4000 GB >> ceiling -> capped down to the ceiling.
|
||||
out = _slab_pages(_args(4000), {"target_kv": 4 * GiB})
|
||||
self.assertEqual(out["target_kv"], ceiling // (4 * GiB))
|
||||
|
||||
def test_explicit_small_slab_size_respected(self):
|
||||
# 100 GB < ceiling (~515 GB) -> honoured as-is.
|
||||
out = _slab_pages(_args(100), {"target_kv": 1 * GiB})
|
||||
self.assertEqual(out["target_kv"], (100 * 1000**3) // (1 * GiB))
|
||||
|
||||
def test_page_larger_than_budget_fails_loud(self):
|
||||
with patch.dict(os.environ, {"SGLANG_CP_HICACHE_MAX_SLAB_GB": "1"}): # 1 GiB budget
|
||||
with self.assertRaises(ValueError):
|
||||
_slab_pages(_args(0), {"target_kv": 2 * GiB}) # one page > budget
|
||||
|
||||
def test_auto_cap_yields_multiple_slabs_for_large_payload(self):
|
||||
"""B3: a realistic payload + the auto-cap must build >1 page-aligned slab,
|
||||
each <= ceiling, tiling [0, total) contiguously (so the group allocator path
|
||||
and per-slab transfer splitting actually engage)."""
|
||||
ceiling = cp_hicache_max_single_register_bytes()
|
||||
bpp = 1 * GiB # 1 GiB per logical page
|
||||
cap_pages = ceiling // bpp
|
||||
total_pages = cap_pages * 3 + 5 # spans ~3.x slabs
|
||||
slab_pages = _slab_pages(_args(0), {PAYLOAD_TARGET_KV: bpp}) # auto-cap
|
||||
slabs = build_cp_shared_l2_slabs_by_payload(
|
||||
{PAYLOAD_TARGET_KV: total_pages}, slab_pages
|
||||
)[PAYLOAD_TARGET_KV]
|
||||
self.assertGreater(len(slabs), 1)
|
||||
base = 0
|
||||
for s in slabs:
|
||||
self.assertLessEqual(s.num_pages * bpp, ceiling) # each slab <= one registration
|
||||
self.assertEqual(s.global_base_page, base) # contiguous tiling
|
||||
base += s.num_pages
|
||||
self.assertEqual(base, total_pages) # exact
|
||||
|
||||
|
||||
class TestCpL3SlabSpansWiring(CustomTestCase):
|
||||
"""B2: _cp_l3_slab_spans (the production wiring) + the accessor read against a
|
||||
REAL SharedHostTensorGroupAllocator allocate_group tensor (not a hand-filled toy)."""
|
||||
|
||||
def test_group_spans_and_accessor_roundtrip(self):
|
||||
n_layers, trailing = 3, 2 # int16 -> slice_bytes = trailing * 2 = 4
|
||||
slice_bytes = trailing * 2
|
||||
infos = (
|
||||
CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=3),
|
||||
CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=3, num_pages=2),
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
allocator = SharedHostTensorGroupAllocator(
|
||||
slabs=infos, directory=tmp, name="target-kv", creator_rank=0,
|
||||
validate_production=False,
|
||||
)
|
||||
group = allocator.allocate_group(
|
||||
(n_layers, 5, trailing), dtype=torch.int16, device="cpu", page_dim=1
|
||||
)
|
||||
# Write a distinct value per (layer, global_page) via the REAL per-slab tensors.
|
||||
for view in group.slab_views:
|
||||
b = int(view.slab_info.global_base_page)
|
||||
for layer in range(n_layers):
|
||||
for local in range(int(view.slab_info.num_pages)):
|
||||
view.tensor[layer, local] = layer * 100 + (b + local)
|
||||
# Production wiring derives spans from the live allocator:
|
||||
spans = HiRadixCache._cp_l3_slab_spans(allocator, 5, "target_kv")
|
||||
self.assertEqual(
|
||||
[(s.global_base_page, s.num_pages) for s in spans], [(0, 3), (3, 2)]
|
||||
)
|
||||
acc = CpSharedL2SlabAccessor(spans, n_layers=n_layers, slice_bytes=slice_bytes)
|
||||
for gp in range(5):
|
||||
vals = torch.frombuffer(
|
||||
bytearray(acc.gather(gp)), dtype=torch.int16
|
||||
).view(n_layers, trailing)
|
||||
for layer in range(n_layers):
|
||||
self.assertTrue(bool((vals[layer] == layer * 100 + gp).all()))
|
||||
group.release_tensor_view_for_close()
|
||||
group.close()
|
||||
|
||||
def test_single_slab_spans_branch(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
alloc = SharedHostTensorAllocator(
|
||||
directory=tmp, name="single", creator_rank=0, validate_production=False
|
||||
)
|
||||
alloc.allocate((3, 4, 2), dtype=torch.int16, device="cpu")
|
||||
spans = HiRadixCache._cp_l3_slab_spans(alloc, 4, "target_kv")
|
||||
self.assertEqual(len(spans), 1)
|
||||
self.assertEqual((spans[0].global_base_page, spans[0].num_pages), (0, 4))
|
||||
self.assertIs(spans[0].mmap, alloc.mapping.mmap)
|
||||
alloc.release_tensor_view_for_close()
|
||||
alloc.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user