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

Avoid remap-cache aliasing across tensor views

Browse Source 
The CP shared-KV remap caches may receive tensor views that share the same storage but represent different logical page rows.  Keying only by storage pointer and shape can reuse a stale remap for another view, corrupting cache-hit materialization.  Use the actual tensor data pointer plus stride, storage offset, and version so different views and mutations are not collapsed into one cache entry.

Constraint: CP shared-KV bs>1 cache-hit paths reuse small tensor views over shared backing tensors.

Rejected: Clear the remap cache on every call | would avoid aliasing but add avoidable hot-path churn and hide the identity bug.

Confidence: high

Scope-risk: narrow

Directive: Do not reduce the remap cache key back to storage pointer only; same-storage views are semantically distinct here.

Tested: g0034 cjy-glm5-new PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_token_slot_remap_cache_distinguishes_same_storage_views test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_paged_slot_remap_cache_distinguishes_same_storage_views test/registered/unit/speculative/test_eagle_worker_v2_cp_hidden.py

Not-tested: local pytest for CP runtime import is blocked by missing starlette in the local base environment.
(cherry picked from commit 7360ef13565dbf21428e3b121112135a2955f913)
This commit is contained in:
laoyao0822
2026-06-22 01:05:41 +08:00
committed by leavelet
parent c1637cfb68
commit 1d7149df1c
2 changed files with 79 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -241,12 +241,21 @@ class SharedPagedBufferSlotRemap:
)
def _tensor_identity_key(tensor: torch.Tensor) -> tuple[int, tuple[int, ...], str, str]:
def _tensor_identity_key(
tensor: torch.Tensor,
) -> tuple[int, tuple[int, ...], tuple[int, ...], int, str, str, int | None]:
try:
version = int(tensor._version)
except Exception:
version = None
return (
int(tensor.untyped_storage().data_ptr()),
int(tensor.data_ptr()),
tuple(int(dim) for dim in tensor.shape),
tuple(int(stride) for stride in tensor.stride()),
int(tensor.storage_offset()),
str(tensor.dtype),
str(tensor.device),
version,
)