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

[Fix] DSV4 cached_loc invalidated when SWA mapping is rebuilt (#25889)

Browse Source 
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-20 22:38:12 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent 3a6de13cd8
commit 888a8794ef
4 changed files with 462 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+4
View File
@@ -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()
test/manual/core/test_dsv4_cached_loc_invalidation.py
+180
View File
@@ -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()
test/manual/core/test_dsv4_hicache_swa_translation_cache.py
+114
View File
@@ -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()
test/manual/core/test_dsv4_stale_loc_crash.py
+164
View File
@@ -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()