minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Unify full→SWA index translation in init_forward_metadata; drop pool caches (#27091)

Browse Source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-06-03 16:12:27 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 8980eb82de
commit c9ca56da8c
29 changed files with 274 additions and 814 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/core/test_dsv4_cached_loc_invalidation.py
-180
View File
@@ -1,180 +0,0 @@
"""Regression test for PR #25889: DeepSeekV4TokenToKVPool.register_mapping()
must clear cached_loc.
Bug scenario (pre-fix):
During a forward pass, the first SWA layer (layer_id == start_layer) computes
and caches `cached_loc` via translate_loc_from_full_to_swa(). If HiCache then
loads back SWA KV from host, it calls register_mapping() to install the new
full->swa index mapping. Before the fix, register_mapping() only stored the
new tensor but did NOT clear cached_loc. Subsequent SWA layers (layer_id >
start_layer) saw `cached_loc is not None` and returned the stale translation,
silently writing KV to wrong SWA pool slots.
Fix (PR #25889): add `self.cached_loc = None` in register_mapping().
Test structure:
- test_stale_without_fix: shows the stale-return bug using a replica
of the pre-fix logic
- test_correct_with_fix: verifies the fixed logic returns fresh values
- test_actual_pool_register_mapping: exercises the real production method
directly (bypassing the full __init__)
Run with:
python -m pytest test/manual/core/test_dsv4_cached_loc_invalidation.py -v
"""
import unittest
import torch
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.utils import get_device
from sglang.test.test_utils import CustomTestCase
# ---------------------------------------------------------------------------
# Minimal stub that replicates DeepSeekV4TokenToKVPool's caching pattern.
# Used to demonstrate both sides of the bug without constructing the full pool.
# ---------------------------------------------------------------------------
class _DSV4CacheStub:
"""Stripped-down replica of the caching logic in DeepSeekV4TokenToKVPool."""
start_layer = 3
def __init__(self, device):
self.device = device
self.cached_loc = None
self.full_to_swa_index_mapping = None
# --- pre-fix version ---
def register_mapping_buggy(self, mapping: torch.Tensor) -> None:
self.full_to_swa_index_mapping = mapping
# BUG: cached_loc not cleared → stale on next mid-forward call
# --- post-fix version (PR #25889) ---
def register_mapping_fixed(self, mapping: torch.Tensor) -> None:
self.full_to_swa_index_mapping = mapping
self.cached_loc = None # THE FIX
def _translate(self, raw_loc: torch.Tensor) -> torch.Tensor:
return self.full_to_swa_index_mapping[raw_loc]
def get_swa_loc(self, layer_id: int, raw_loc: torch.Tensor) -> torch.Tensor:
"""Exact replica of set_swa_key_buffer_radix_fused caching branch."""
if layer_id == self.start_layer or self.cached_loc is None:
self.cached_loc = self._translate(raw_loc)
return self.cached_loc
def _make_mapping(indices, values, size=32, device="cpu"):
m = torch.zeros(size, dtype=torch.int64, device=device)
m[indices] = torch.tensor(values, dtype=torch.int64, device=device)
return m
class TestDSV4CachedLocBugAndFix(CustomTestCase):
"""Shows the pre-fix bug and verifies the post-fix behaviour."""
def setUp(self):
self.device = get_device()
self.raw_loc = torch.tensor([0, 1, 2, 3], device=self.device)
self.mapping_v1 = _make_mapping(
[0, 1, 2, 3], [10, 11, 12, 13], device=self.device
)
self.mapping_v2 = _make_mapping(
[0, 1, 2, 3], [20, 21, 22, 23], device=self.device
)
def test_stale_without_fix(self):
"""Without the fix, register_mapping() mid-forward leaves a stale cached_loc.
Sequence:
1. start_layer: cached_loc = translate(mapping_v1) = [10..13]
2. HiCache load-back: register_mapping(mapping_v2) ← buggy, no clear
3. start_layer+1: layer_id != start_layer, cached_loc is not None
→ returns stale [10..13], NOT the correct [20..23]
"""
stub = _DSV4CacheStub(self.device)
stub.register_mapping_buggy(self.mapping_v1)
# Forward pass — first SWA layer primes the cache.
loc = stub.get_swa_loc(stub.start_layer, self.raw_loc)
self.assertEqual(loc.tolist(), [10, 11, 12, 13])
# HiCache load-back installs new mapping (buggy path).
stub.register_mapping_buggy(self.mapping_v2)
self.assertIsNotNone(stub.cached_loc, "Bug: cached_loc not cleared")
# Next SWA layer — should use mapping_v2 but returns mapping_v1.
loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc)
self.assertEqual(
loc_next.tolist(),
[10, 11, 12, 13],
"Confirms the bug: stale cached_loc [10..13] returned instead of [20..23]",
)
def test_correct_with_fix(self):
"""With the fix, register_mapping() clears cached_loc; next layer recomputes.
Same sequence as test_stale_without_fix but using the fixed register_mapping.
"""
stub = _DSV4CacheStub(self.device)
stub.register_mapping_fixed(self.mapping_v1)
# Forward pass — first SWA layer primes the cache.
loc = stub.get_swa_loc(stub.start_layer, self.raw_loc)
self.assertEqual(loc.tolist(), [10, 11, 12, 13])
# HiCache load-back installs new mapping (fixed path).
stub.register_mapping_fixed(self.mapping_v2)
self.assertIsNone(
stub.cached_loc, "Fix: cached_loc cleared by register_mapping"
)
# Next SWA layer — recomputes with mapping_v2.
loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc)
self.assertEqual(
loc_next.tolist(),
[20, 21, 22, 23],
"Fix works: fresh translation [20..23] from new mapping",
)
class TestDSV4ActualPoolRegisterMapping(CustomTestCase):
"""Exercises the production DeepSeekV4TokenToKVPool.register_mapping() directly.
Uses __new__ to bypass the complex __init__ (which needs full GPU pool setup)
and tests only the register_mapping / cached_loc contract.
"""
def test_register_mapping_clears_cached_loc(self):
device = get_device()
# Bypass full __init__ — only the fields register_mapping touches matter.
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
old_loc = torch.tensor([10, 11, 12, 13], device=device)
pool.cached_loc = old_loc
pool.full_to_swa_index_mapping = None
new_mapping = torch.arange(64, dtype=torch.int64, device=device)
pool.register_mapping(new_mapping)
self.assertIsNone(pool.cached_loc, "register_mapping must clear cached_loc")
self.assertIs(pool.full_to_swa_index_mapping, new_mapping)
def test_register_mapping_clears_none_cached_loc(self):
"""Idempotent when cached_loc is already None."""
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
pool.cached_loc = None
pool.full_to_swa_index_mapping = None
mapping = torch.arange(16, dtype=torch.int64, device=get_device())
pool.register_mapping(mapping)
self.assertIsNone(pool.cached_loc)
self.assertIs(pool.full_to_swa_index_mapping, mapping)
if __name__ == "__main__":
unittest.main()
test/manual/core/test_dsv4_hicache_swa_translation_cache.py
-114
View File
@@ -1,114 +0,0 @@
"""E2E regression for PR #25889: DSV4 cached_loc stale after HiCache load-back.
Bug (pre-fix):
DeepSeekV4TokenToKVPool.register_mapping() replaces full_to_swa_index_mapping
but does NOT clear self.cached_loc. When SGLANG_OPT_CACHE_SWA_TRANSLATION=True,
set_swa_key_buffer_radix_fused() caches the full→SWA translation across SWA
layers. After a HiCache commit/load-back that calls register_mapping() with a
new mapping, subsequent SWA layers reuse the stale cached_loc and write KV to
wrong SWA pool slots — producing divergent logprobs.
Fix (PR #25889):
register_mapping() sets self.cached_loc = None so the first SWA layer in the
next forward pass recomputes the translation from the fresh mapping.
Test strategy:
Subclass the existing DSV4 Flash HiCache KL suite and activate the SWA
translation cache via SGLANG_OPT_CACHE_SWA_TRANSLATION=1. The mixin tests
(test_multiturn_logprobs_match, test_multiturn_prefill_cache_hit_branching,
test_multiturn_decode_cache_hit_branching) compare logprobs from cold and
warm radix-cache hits. Without the fix the stale translation corrupts SWA KV
data and the KL divergence exceeds the threshold; with the fix it stays within.
"""
import unittest
from test_unified_radix_cache_kl_hicache import (
DSV4_FLASH_LAUNCH_TIMEOUT,
DSV4_FLASH_MODEL,
_assert_dsv4_decode_cached_tokens,
)
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
from sglang.test.kl_multiturn_utils import get_input_ids
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestDSV4HiCacheSWATranslationCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""DSV4 Flash FP8 + HiCache + SWA translation cache enabled.
Identical server config to TestUnifiedDeepSeekV4FlashHiCache but with
SGLANG_OPT_CACHE_SWA_TRANSLATION=1 to activate the cached_loc path.
Without PR #25889 the KL tests fail; with the fix they pass.
"""
kl_threshold = 0.005
sampling_temperature = 0
decode_hit_request_batch_size = 3
decode_hit_inter_batch_delay_s = 0.5
decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens)
gsm8k_threshold = 0.90
num_gsm8k_questions = 100
@unittest.skipIf(True, "Covered by test_multiturn_prefill_cache_hit_branching.")
def test_multiturn_logprobs_match(self):
pass
@classmethod
def setUpClass(cls):
cls.model = DSV4_FLASH_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DSV4_FLASH_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp-size",
"4",
"--attention-backend",
"compressed",
"--page-size",
"256",
"--chunked-prefill-size",
"8192",
"--mem-fraction-static",
"0.9",
"--disable-shared-experts-fusion",
"--enable-hierarchical-cache",
"--hicache-ratio",
"4",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"direct",
"--hicache-mem-layout",
"page_first_direct",
"--swa-full-tokens-ratio",
"0.25",
"--max-total-tokens",
"20000",
"--max-running-requests",
"4",
],
env={
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
# Activate the SWA translation cache — the flag that exposes the bug.
"SGLANG_OPT_CACHE_SWA_TRANSLATION": "1",
},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/manual/core/test_dsv4_stale_loc_crash.py
-164
View File
@@ -1,164 +0,0 @@
"""Crash regression for PR #25889: stale cached_loc after register_mapping().
Bug:
DeepSeekV4TokenToKVPool caches the full→SWA index translation in
self.cached_loc (when SGLANG_OPT_CACHE_SWA_TRANSLATION=True).
register_mapping() was replacing full_to_swa_index_mapping without
clearing cached_loc, so a subsequent set_swa_key_buffer_radix_fused
call would use stale SWA indices and, if those indices exceed the
current pool size, raise a RuntimeError (OOB tensor access).
Crash scenario reproduced here:
Pass 1 (large SWA pool, size=8): first SWA layer primes
cached_loc = [4, 5, 6, 7].
register_mapping() called with new mapping (pre-fix: cache not cleared).
Pass 2 (smaller SWA pool, size=4): same SWA layer finds cached_loc is
not None → uses stale [4, 5, 6, 7] → OOB write on size-4 pool →
RuntimeError.
Fix: register_mapping() sets self.cached_loc = None so the next call
recomputes with the fresh mapping.
Run with:
python -m pytest test/registered/unit/mem_cache/test_dsv4_stale_loc_crash.py -v
"""
import unittest
import torch
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
_NUM_HEADS = 2
_HEAD_DIM = 8
_SWA_LARGE = 8 # size-8 pool used during pass 1
_SWA_SMALL = 4 # size-4 pool used during pass 2 (simulates post-HiCache loadback)
class _SWAPoolMock:
"""Minimal SWA pool mock whose set_key_buffer_fused does a real tensor write.
A write with OOB swa_loc raises RuntimeError, reproducing the crash.
"""
def __init__(self, size: int):
self.buf = torch.zeros(size, _NUM_HEADS, _HEAD_DIM)
def set_key_buffer_fused(
self, local_layer_id: int, swa_loc: torch.Tensor, cache_k: torch.Tensor
) -> None:
n = swa_loc.numel()
self.buf[swa_loc.long()] = cache_k[:n].reshape(n, _NUM_HEADS, _HEAD_DIM)
def _build_pool(
mapping: torch.Tensor,
swa_pool: _SWAPoolMock,
start_layer: int = 0,
) -> DeepSeekV4TokenToKVPool:
"""Create a minimal DeepSeekV4TokenToKVPool via __new__, bypassing __init__."""
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
pool.cached_loc = None
pool._should_cache_swa = True
pool.start_layer = start_layer
pool.full_to_swa_index_mapping = mapping
pool.swa_kv_pool = swa_pool
# _swa_local_layer_id: map global SWA layer id → local index 0 for this test.
pool._swa_local_layer_id = lambda lid: 0
return pool
def _mapping(indices: list, size: int = 16) -> torch.Tensor:
m = torch.full((size,), -1, dtype=torch.int64)
for i, v in enumerate(indices):
m[i] = v
return m
class TestDSV4StaleLocCrash(CustomTestCase):
"""
Two paired tests that together constitute the crash regression for #25889.
test_crash_without_fix: reproduces the RuntimeError that occurs when
register_mapping() does NOT clear cached_loc.
test_fix_prevents_crash: same sequence with the fixed register_mapping()
→ no error, correct SWA slots written.
"""
def setUp(self):
# raw_loc: full-pool indices for 4 tokens.
self.raw_loc = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
# cache_k: synthetic key data for 4 tokens.
self.cache_k = torch.ones(4, _NUM_HEADS, _HEAD_DIM)
# SWA layer id > start_layer (0), so cached_loc is only reset by the
# "is None" branch, NOT by the "layer_id == start_layer" branch.
self.swa_layer_id = 1
# mapping_v1: raw_loc [0,1,2,3] → SWA slots [4,5,6,7] (valid for size-8 pool).
self.mapping_v1 = _mapping([4, 5, 6, 7])
# mapping_v2: same raw_loc → SWA slots [0,1,2,3] (valid for size-4 pool).
self.mapping_v2 = _mapping([0, 1, 2, 3])
def test_crash_without_fix(self):
"""Without the fix, stale cached_loc [4,5,6,7] causes RuntimeError on
the size-4 pool used in pass 2."""
large_pool = _SWAPoolMock(_SWA_LARGE)
pool = _build_pool(self.mapping_v1, large_pool, start_layer=0)
# Pass 1: swa_layer_id=1, cached_loc is None → compute and cache [4,5,6,7].
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7], "Pass 1: cache primed")
# PRE-FIX register_mapping: replace mapping WITHOUT clearing cached_loc.
pool.full_to_swa_index_mapping = self.mapping_v2
# Pass 2 on a smaller pool (size=4). cached_loc is still [4,5,6,7].
# OOB write → RuntimeError.
pool.swa_kv_pool = _SWAPoolMock(_SWA_SMALL)
with self.assertRaises((RuntimeError, IndexError)):
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
def test_fix_prevents_crash(self):
"""With the fix, register_mapping() clears cached_loc. Pass 2 recomputes
[0,1,2,3] from mapping_v2 and writes to the correct size-4 pool slots."""
large_pool = _SWAPoolMock(_SWA_LARGE)
pool = _build_pool(self.mapping_v1, large_pool, start_layer=0)
# Pass 1: prime cache → cached_loc = [4,5,6,7].
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7])
# FIXED register_mapping: clears cached_loc.
pool.register_mapping(self.mapping_v2)
self.assertIsNone(
pool.cached_loc, "Fix: cached_loc cleared by register_mapping"
)
# Pass 2 on the smaller pool: cached_loc is None → recompute with mapping_v2.
small_pool = _SWAPoolMock(_SWA_SMALL)
pool.swa_kv_pool = small_pool
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
self.assertEqual(
pool.cached_loc.tolist(), [0, 1, 2, 3], "Fresh indices after fix"
)
# Verify data landed in the correct SWA slots [0-3], not the stale [4-7].
self.assertTrue(
small_pool.buf[0:4].abs().sum().item() > 0,
"Correct SWA slots 0-3 received data",
)
if __name__ == "__main__":
unittest.main()
test/manual/core/test_swa_loc_translation_cache.py
-233
View File
@@ -1,233 +0,0 @@
"""Manual tests for SWAKVPool.translate_loc_from_full_to_swa cache behaviour.
These tests cover three properties introduced by PR #25824:
1. Cache key uses data_ptr() — correctly distinguishes views at different
offsets within the same storage (untyped_storage().data_ptr() would not).
2. Allocator mutations invalidate the cache — alloc/free/clear/
set_full_to_swa_mapping each call invalidate_loc_cache() so the next
translation sees the fresh mapping.
3. BaseSWAKVPool.invalidate_loc_cache is a no-op default — subclasses that
don't cache (e.g. DSV4) can be called safely without AttributeError.
Run with:
python -m pytest test/manual/core/test_swa_loc_translation_cache.py -v
"""
import unittest
import torch
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.utils import get_device
from sglang.test.test_utils import CustomTestCase
def _build_pool(
kv_size: int = 32,
kv_size_swa: int = 32,
page_size: int = 1,
):
device = get_device()
num_layers = 8
full_layer_ids = [0, 4]
swa_layer_ids = [i for i in range(num_layers) if i not in set(full_layer_ids)]
pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=torch.bfloat16,
head_num=4,
head_dim=64,
swa_attention_layer_ids=swa_layer_ids,
full_attention_layer_ids=full_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=torch.bfloat16,
device=device,
kvcache=pool,
need_sort=False,
)
return pool, allocator, device
class TestCacheKeyDataPtr(CustomTestCase):
"""Cache key uses data_ptr(), which encodes the storage offset."""
def test_same_offset_view_is_cache_hit(self):
"""Two different Python objects pointing to the same base are a hit."""
pool, allocator, device = _build_pool()
loc = allocator.alloc(4)
self.assertIsNotNone(loc)
# Create two slice objects at offset 0 — same data_ptr, same numel.
view_a = loc[:4]
view_b = loc[:4]
self.assertIsNot(view_a, view_b) # different Python objects
self.assertEqual(view_a.data_ptr(), view_b.data_ptr())
result_a = pool.translate_loc_from_full_to_swa(view_a)
result_b = pool.translate_loc_from_full_to_swa(view_b)
# Both should return the identical tensor (cache hit).
self.assertIs(result_a, result_b)
def test_different_offset_view_is_cache_miss(self):
"""Views at different offsets produce different data_ptr → cache miss."""
pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
loc = allocator.alloc(10)
self.assertIsNotNone(loc)
self.assertGreaterEqual(loc.numel(), 10)
view_lo = loc[0:5]
view_hi = loc[5:10]
self.assertEqual(view_lo.numel(), view_hi.numel()) # same numel
# Different data_ptr (different storage offset).
self.assertNotEqual(view_lo.data_ptr(), view_hi.data_ptr())
# Prime the cache with view_lo.
result_lo = pool.translate_loc_from_full_to_swa(view_lo)
# view_hi should be a cache miss and produce a distinct translation.
result_hi = pool.translate_loc_from_full_to_swa(view_hi)
# They should NOT be the same object (different cache entries).
self.assertIsNot(result_lo, result_hi)
# And the content must differ (different full indices → different swa).
self.assertFalse(torch.equal(result_lo, result_hi))
def test_storage_base_ptr_would_collide(self):
"""Demonstrate that untyped_storage().data_ptr() WOULD collide for the
two views above — confirming data_ptr() is the right key."""
t = torch.arange(20, device=get_device())
a, b = t[0:10], t[5:15]
# Same storage base — old key would collide.
self.assertEqual(a.untyped_storage().data_ptr(), b.untyped_storage().data_ptr())
self.assertEqual(a.numel(), b.numel())
# But data_ptr differs — new key is safe.
self.assertNotEqual(a.data_ptr(), b.data_ptr())
class TestAllocatorMutationInvalidation(CustomTestCase):
"""Each allocator method that writes the mapping calls invalidate_loc_cache."""
def _prime_and_check_invalidation(self, pool, allocator, mutate_fn):
"""Helper: prime cache, mutate, assert fresh translation."""
loc = allocator.alloc(4)
self.assertIsNotNone(loc)
# Prime the cache.
first = pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
# Mutate — should invalidate.
mutate_fn(allocator, loc)
# Cache must be cleared after mutation.
self.assertIsNone(pool._cached_loc_key)
self.assertIsNone(pool._cached_swa_loc)
def test_alloc_invalidates(self):
pool, allocator, _ = _build_pool()
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
# Another alloc should invalidate.
allocator.alloc(4)
self.assertIsNone(pool._cached_loc_key)
def test_free_swa_invalidates(self):
pool, allocator, _ = _build_pool()
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
allocator.free_swa(loc)
self.assertIsNone(pool._cached_loc_key)
def test_clear_invalidates(self):
pool, allocator, _ = _build_pool()
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
allocator.clear()
self.assertIsNone(pool._cached_loc_key)
def test_set_full_to_swa_mapping_invalidates(self):
"""HiCache load-back path: set_full_to_swa_mapping must invalidate."""
pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
# Simulate HiCache rebuild with new swa indices.
new_swa = torch.arange(4, dtype=torch.int64, device=device)
allocator.set_full_to_swa_mapping(loc, new_swa)
self.assertIsNone(pool._cached_loc_key)
# Translation after rebuild should reflect the new mapping.
result = pool.translate_loc_from_full_to_swa(loc)
self.assertEqual(result.tolist(), new_swa.tolist())
class TestBaseClassNoOp(CustomTestCase):
"""BaseSWAKVPool.invalidate_loc_cache is a no-op default — must not raise."""
def test_noop_does_not_raise(self):
# BaseSWAKVPool is abstract; instantiate via SWAKVPool which inherits.
pool, _, _ = _build_pool()
# Calling on the concrete class uses the override — that's fine.
pool.invalidate_loc_cache() # must not raise
pool.invalidate_loc_cache() # idempotent
def test_base_class_noop_directly(self):
"""Call the base-class method directly to verify it's a true no-op."""
pool, _, _ = _build_pool()
# Prime the cache first.
loc = pool.full_to_swa_index_mapping # any tensor
pool._cached_loc_key = ("dummy", 1)
pool._cached_swa_loc = torch.zeros(1)
# Call the BASE class method directly — should not clear the cache
# (it's a no-op; the concrete override is what clears).
BaseSWAKVPool.invalidate_loc_cache(pool)
# base no-op: cache untouched
self.assertIsNotNone(pool._cached_loc_key)
class TestExplicitInvalidationCycle(CustomTestCase):
"""Simulates the per-forward-pass invalidation done by model_runner."""
def test_fresh_translation_after_explicit_invalidation(self):
"""After invalidate_loc_cache(), a new alloc produces the right mapping."""
pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
# First "forward pass": alloc 4 tokens, translate.
loc1 = allocator.alloc(4)
trans1 = pool.translate_loc_from_full_to_swa(loc1).clone()
# Simulate start of next forward pass: model_runner calls invalidate.
pool.invalidate_loc_cache()
self.assertIsNone(pool._cached_loc_key)
# Alloc 4 more (mapping changes), translate loc1 again.
loc2 = allocator.alloc(4)
# Alloc already invalidated; translate loc1 with fresh mapping.
trans1_after = pool.translate_loc_from_full_to_swa(loc1)
# loc1's SWA mapping hasn't changed (same full→swa assignment),
# so result should be equal — but it must have been recomputed
# (cache key was None before this call).
self.assertEqual(trans1.tolist(), trans1_after.tolist())
# loc2 should have different translation than loc1.
trans2 = pool.translate_loc_from_full_to_swa(loc2)
# They have different indices, so translation differs.
self.assertFalse(torch.equal(trans1_after, trans2))
if __name__ == "__main__":
unittest.main()