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:
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user