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