diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 166ace228..8a5a6247d 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -492,6 +492,10 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): self.full_to_swa_index_mapping = full_to_swa_index_mapping + self.cached_loc = None # mapping replaced; discard any cached translation + + def invalidate_loc_cache(self) -> None: + self.cached_loc = None def get_ring_size(self, compress_ratio: int) -> int: server_args = get_global_server_args() diff --git a/test/manual/core/test_dsv4_cached_loc_invalidation.py b/test/manual/core/test_dsv4_cached_loc_invalidation.py new file mode 100644 index 000000000..ff7314fc4 --- /dev/null +++ b/test/manual/core/test_dsv4_cached_loc_invalidation.py @@ -0,0 +1,180 @@ +"""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() diff --git a/test/manual/core/test_dsv4_hicache_swa_translation_cache.py b/test/manual/core/test_dsv4_hicache_swa_translation_cache.py new file mode 100644 index 000000000..5a83c2e5b --- /dev/null +++ b/test/manual/core/test_dsv4_hicache_swa_translation_cache.py @@ -0,0 +1,114 @@ +"""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 import UnifiedRadixTreeTestMixin +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.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() diff --git a/test/manual/core/test_dsv4_stale_loc_crash.py b/test/manual/core/test_dsv4_stale_loc_crash.py new file mode 100644 index 000000000..3b31ace55 --- /dev/null +++ b/test/manual/core/test_dsv4_stale_loc_crash.py @@ -0,0 +1,164 @@ +"""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()