CP HiCache trace: add KV-value round-trip fingerprint (backup vs reload)
The index trace proved per-node addressing is self-consistent, so the corruption must be in the KV VALUES, not the bookkeeping. Add a full-tensor position-weighted fingerprint (khash) of the actual device KV at two points, both keyed by (node_id, layer_id) in owned-position order so they compare directly: - backup_kv_hash: device KV at backup, taken on the DEFAULT stream (outside the write_stream block) so it is the correct POST-STORE value, not the racy write_stream copy's view. - reload_kv_hash: device KV after the H2D load (on load_stream, in-order after the copy). If a reload hash matches no backup hash for that node+layer, the round-trip delivered wrong KV -- catching BOTH a per-layer store-vs-copy ordering race (backup-correct != reloaded-stale) AND transfer corruption, in one run. Since the compose/attention is identical fresh-vs-reload and stateless-per-forward, all-ranks byte-correct shards => correct output, so this per-rank value check is complete for the round-trip. Gated at SGLANG_CP_HICACHE_KV_TRACE=2. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Set
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import cptrace, rng, trace_enabled
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import cptrace, khash, rng, trace_enabled
|
||||
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -1489,6 +1489,23 @@ class HiCacheController:
|
||||
state.draft_host_indices, self.write_stream
|
||||
)
|
||||
|
||||
# Value fingerprint of the EXACT device KV being backed up, taken on the
|
||||
# DEFAULT stream (outside the write_stream block) so it reflects the
|
||||
# correct post-store value -- NOT the racy write_stream copy's view. If a
|
||||
# per-layer store/copy ordering race exists, the host copy will diverge
|
||||
# from THIS hash, and the reload hash (post-H2D) will not match it.
|
||||
if trace_enabled(2):
|
||||
for state in target_transfer_states:
|
||||
kv = self.mem_pool_device.get_key_buffer(layer_id)
|
||||
cptrace(
|
||||
2,
|
||||
"backup_kv_hash",
|
||||
node_id=state.reservation.node_id,
|
||||
layer=layer_id,
|
||||
kvhash=khash(kv[state.physical_device_indices]),
|
||||
nrows=int(state.physical_device_indices.numel()),
|
||||
)
|
||||
|
||||
for state in final_states:
|
||||
self._append_layer_write_ack(state)
|
||||
|
||||
@@ -1951,6 +1968,21 @@ class HiCacheController:
|
||||
self.io_backend,
|
||||
)
|
||||
submitted_target = True
|
||||
# Value fingerprint of what actually landed on device
|
||||
# after H2D (on load_stream, in-order after the load).
|
||||
# Compare to backup_kv_hash(node,layer): inequality =
|
||||
# round-trip delivered wrong KV. node_ids is the merged
|
||||
# op's nodes (single node in the c=1 rehit pass).
|
||||
if trace_enabled(2):
|
||||
kv = self.mem_pool_device.get_key_buffer(i)
|
||||
cptrace(
|
||||
2,
|
||||
"reload_kv_hash",
|
||||
node_ids=node_ids,
|
||||
layer=i,
|
||||
kvhash=khash(kv[device_indices]),
|
||||
nrows=int(device_indices.numel()),
|
||||
)
|
||||
producer_event.complete(i)
|
||||
elif op is None and i < self.layer_num:
|
||||
producer_event.complete(i)
|
||||
|
||||
@@ -57,6 +57,28 @@ def _fmt(v) -> str:
|
||||
return str(v).replace(" ", "")
|
||||
|
||||
|
||||
def khash(t) -> int:
|
||||
"""Full-tensor, position-weighted int64 fingerprint of a KV slice.
|
||||
|
||||
Sensitive to BOTH value and position (catches permutation/scatter, unlike a
|
||||
plain sum), covers the whole tensor (not a byte sample), and reduces to one
|
||||
int. Debug-only: the .item() sync is acceptable when tracing is on. Used to
|
||||
compare the KV bytes a node's pages hold at backup (correct, post-store, on
|
||||
the default stream) vs what comes back at reload (post-H2D); inequality =
|
||||
the round-trip delivered wrong KV (store-vs-copy race or transfer corruption)."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if t is None or t.numel() == 0:
|
||||
return 0
|
||||
b = t.detach().contiguous().view(torch.uint8).reshape(-1).to(torch.int64)
|
||||
n = b.numel()
|
||||
idx = torch.arange(1, n + 1, device=b.device, dtype=torch.int64)
|
||||
return int((((b * idx).sum() ^ (b.sum() * 1000003) ^ n) & 0x7FFFFFFFFFFFFFFF).item())
|
||||
except Exception:
|
||||
return -1
|
||||
|
||||
|
||||
def cptrace(level: int, tag: str, **fields) -> None:
|
||||
"""Emit one ``[CPTRACE <tag>] k=v ...`` line if the gate >= level.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user