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

@@ -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(