CP HiCache trace: NSA indexer-K round-trip hash + forward-side attn/MoE hashes

Main KV round-trips byte-perfect yet output is garbage, so add the two
unverified components:

1. NSA INDEXER-K round-trip (level 2): hash the device index_k_with_scale_buffer
   at backup (_backup_indexer_from_device_per_layer) and reload (NSA
   load_to_device_per_layer, after _load_indexer), keyed by host-slot fingerprint
   + layer, with khash+nz. The indexer selects which tokens attention attends
   (top-k); if it corrupts on reload -> wrong selection -> garbage even with
   correct main KV.

2. FORWARD-side per-layer hashes (level 3, eager extend path only, cuda-graph
   guarded): attn-input, attn-output (pre-residual), topk_indices (the indexer's
   selection output -- direct consumer of the indexer-K), and MoE-input, in the
   DeepseekV2/GlmMoeDsa decoder layer forward. Localizes where a reload forward
   diverges: topk diverges => indexer-K cache; attn-out diverges (topk ok) =>
   main KV/page mapping; moe-in diverges (attn-out ok) => residual/MoE.

Analyzer compares indexer reload vs backup (host_fp keyed) + flags zero/degenerate
hidden states per stage. Level 3 (SGLANG_CP_HICACHE_KV_TRACE=3) captures everything.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-17 03:05:30 +00:00
parent e3cfba8441
commit a447ae8317
3 changed files with 70 additions and 0 deletions

View File

@@ -95,6 +95,44 @@ def knz(t) -> int:
return -1
def fwd_hash(forward_batch, layer_id, stage, t, *, level: int = 3) -> None:
"""Hash a per-layer forward tensor (attn-in/out, topk_indices, MoE-in) to
localize where a reload forward diverges from a fresh one. Level 3 (so the
level-2 KV run is unaffected). Guards: only the EAGER EXTEND path (the reload
case) -- never under cuda-graph decode (.item() sync would corrupt capture);
handles topk_indices=None and tuple/quant hidden_states."""
if trace_level() < level:
return
try:
import torch
fm = getattr(forward_batch, "forward_mode", None)
if fm is not None and hasattr(fm, "is_extend") and not fm.is_extend():
return
if isinstance(t, torch.Tensor):
h, nz, rows = khash(t), knz(t), int(t.shape[0]) if t.dim() else 0
elif t is None:
h, nz, rows = 0, 0, 0
else:
h, nz, rows = -2, -2, -2 # tuple/quant: not a plain tensor
rids = getattr(forward_batch, "rids", None)
lay = getattr(forward_batch, "cp_shared_kv_layout", None)
cptrace(
level,
"fwd_hash",
rid=(rids[0] if rids else "?"),
nreq=(len(rids) if rids else 0),
layer=layer_id,
stage=stage,
h=h,
nz=nz,
rows=rows,
cprank=(getattr(lay, "cp_rank", -1) if lay is not None else -1),
)
except Exception:
pass
def cptrace(level: int, tag: str, **fields) -> None:
"""Emit one ``[CPTRACE <tag>] k=v ...`` line if the gate >= level.

View File

@@ -12,6 +12,8 @@ from typing import Optional, Tuple
import psutil
import torch
from sglang.srt.mem_cache.cp_hicache_trace import cptrace, khash, knz, rng, trace_enabled
from sglang.jit_kernel.hicache import (
can_use_hicache_jit_kernel,
)
@@ -2550,6 +2552,18 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
host_page_indices, device_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
if trace_enabled(2):
_dev = device_pool.index_k_with_scale_buffer[device_layer_slot]
_rows = _dev[device_page_indices]
cptrace(
2,
"indexer_backup_hash",
layer=layer_id,
host=rng(host_indices),
kvhash=khash(_rows),
nz=knz(_rows),
npages=int(device_page_indices.numel()),
)
use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0
if use_kernel:
if self.layout == "layer_first":
@@ -2624,6 +2638,19 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
self._load_indexer_to_device_per_layer(
device_pool, host_indices, device_indices, layer_id, io_backend
)
if trace_enabled(2) and self._is_device_index_layer_active(device_pool, layer_id):
_dls = self._device_index_layer_slot(device_pool, layer_id)
_, _dpi = self._get_indexer_page_indices(host_indices, device_indices)
_rows = device_pool.index_k_with_scale_buffer[_dls][_dpi]
cptrace(
2,
"indexer_reload_hash",
layer=layer_id,
host=rng(host_indices),
kvhash=khash(_rows),
nz=knz(_rows),
npages=int(_dpi.numel()),
)
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend

View File

@@ -73,6 +73,7 @@ from sglang.srt.layers.communicator import (
enable_moe_dense_fully_dp,
get_attn_tp_context,
)
from sglang.srt.mem_cache.cp_hicache_trace import fwd_hash as _cp_fwd_hash
from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator
from sglang.srt.layers.dp_attention import (
get_attention_cp_rank,
@@ -1730,6 +1731,7 @@ class DeepseekV2DecoderLayer(nn.Module):
forward_batch,
quant_format,
)
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "attn_in", hidden_states)
previous_cp_shared_kv_num_model_layers = getattr(
forward_batch, "cp_shared_kv_num_model_layers", None
@@ -1761,10 +1763,13 @@ class DeepseekV2DecoderLayer(nn.Module):
hidden_states, topk_indices = hidden_states
else:
topk_indices = None
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "attn_out", hidden_states)
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk", topk_indices)
hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch
)
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_in", hidden_states)
should_allreduce_fusion = (
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(