From 50d1c2fc5a24ebd2e9adbbb18047601e7e0238dd Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 17 Jun 2026 04:03:47 +0000 Subject: [PATCH] CP HiCache trace: compose gather-output + span-owner + descriptor + MoE-stage hashes Both cached KV components round-trip byte-perfect, so the bug is downstream in the cross-rank COMPOSE/GATHER on a reload forward. Add level-3 hashes targeting the active suspect = materialize_prefix_and_reuse_current_kv_page_slots (reloaded prefix gathered via modulo-owner IPC descriptors + fresh current via page_inverse staging; the abutting prefix|current boundary is the suspect): - gather_out: the live composed dense KV the attention actually reads (mixed_locs gather), with prefix_pages/current_pages; nz==0 = zero/uninitialized composed KV (the observed 0|0|0), self-contained. - span_owner: per prefix/current span -> owner ranks (modulo) + physical pages + logical pages + per-span hash/nz, to verify the prefix span maps to modulo owners and catch a boundary conflation/off-by-one. - ipc_desc: owner/src/logical-page ranges in build_cp_shared_kv_ipc_page_descriptors (the modulo-owner prefix gather chokepoint). - MoE stages (forward_deepep, via fwd_hash, valid-row local tensors only, never the a2a-permuted intermediate): moe_postsel, router_logits, topk_ids/topk_w, experts_out, moe_out_compact, moe_out_restored -> brackets select/router/topk/ dispatch-combine/shared-add/row-restore. (The symm ComposePlan is dormant/unwired, so excluded.) All level 3, eager-extend only, try/except-guarded, helpers khash/knz/rng. Analyzer flags zero composed KV + zero spans + MoE-stage zeros. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../attention/nsa/cp_shared_kv_runtime.py | 59 +++++++++++++++++++ python/sglang/srt/models/deepseek_v2.py | 7 +++ 2 files changed, 66 insertions(+) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 44f8eb9ae..64f639b75 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -10,6 +10,8 @@ import torch from sglang.srt.environ import envs from sglang.srt.mem_cache.cp_hicache_trace import cptrace as _cptrace +from sglang.srt.mem_cache.cp_hicache_trace import khash as _khash +from sglang.srt.mem_cache.cp_hicache_trace import knz as _knz from sglang.srt.mem_cache.cp_hicache_trace import rng as _cprng from sglang.srt.mem_cache.cp_hicache_trace import trace_enabled as _cptrace_enabled from sglang.srt.layers.attention.nsa.utils import ( @@ -2240,6 +2242,21 @@ def build_cp_shared_kv_ipc_page_descriptors( invalid_value = torch.full_like(owner_ranks, -1) owner_ranks = torch.where(invalid, invalid_value, owner_ranks) src_page_indices = torch.where(invalid, invalid_value, src_page_indices) + if _cptrace_enabled(3): + try: + _valid = ~invalid + _cptrace( + 3, + "ipc_desc", + cprank=layout.cp_rank, + slots=int(logical_pages.numel()), + valid=int(_valid.sum().item()), + owners=_cprng(owner_ranks[_valid]), + src=_cprng(src_page_indices[_valid]), + lpg=_cprng(logical_pages[_valid]), + ) + except Exception: + pass return owner_ranks.contiguous(), src_page_indices.contiguous() @@ -5675,6 +5692,48 @@ def materialize_prefix_and_reuse_current_kv_page_slots( materialized_by_ipc, kv_cache.dtype, ) + if _cptrace_enabled(3): + try: + _live = mixed_locs.reshape(-1) + _live = _live[_live >= 0] + _gath = ( + mixed_kv_cache.index_select(0, _live.to(torch.long)) + if _live.numel() + else mixed_kv_cache[:0] + ) + _cptrace( + 3, + "gather_out", + cprank=layout.cp_rank, + layer=layer_id, + total_slots=int(total_slots), + live_rows=int(_live.numel()), + prefix_pages=sum(int(e) - int(s) for s, e in prefix_spans), + current_pages=sum(int(e) - int(s) for s, e in merged_current_spans), + h=_khash(_gath), + nz=_knz(_gath), + ) + _slp = slot_remap.slot_logical_pages.reshape(-1) + for _kind, _spans in (("prefix", prefix_spans), ("current", merged_current_spans)): + for _s, _e in _spans: + _sl = slot_range_to_token_slice(page_size, int(_s), int(_e)) + _lp = _slp[int(_s) : int(_e)].to(torch.long) + _cptrace( + 3, + "span_owner", + cprank=layout.cp_rank, + layer=layer_id, + kind=_kind, + s=int(_s), + e=int(_e), + owners=_cprng(layout.owner_for_logical_pages(_lp)), + phys=_cprng(layout.logical_pages_to_physical(_lp)), + lpg=_cprng(_lp), + h=_khash(mixed_kv_cache[_sl]), + nz=_knz(mixed_kv_cache[_sl]), + ) + except Exception: + pass return mixed_kv_cache, mixed_locs diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 006a669c8..9e5cefaf9 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -775,6 +775,7 @@ class DeepseekV2MoE(nn.Module): hidden_states, ) + _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_postsel", hidden_states) shared_output = None sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn sbo_overlap_dispatch_flag = ( @@ -787,6 +788,7 @@ class DeepseekV2MoE(nn.Module): if hidden_states.shape[0] > 0: # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, forward_batch=forward_batch) + _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "router_logits", router_logits) if not sbo_enabled_flag: if self.alt_stream is not None: self.alt_stream.wait_stream(torch.cuda.current_stream()) @@ -804,6 +806,8 @@ class DeepseekV2MoE(nn.Module): layer_id=self.layer_id, ), ) + _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk_ids", getattr(topk_output, "topk_ids", None)) + _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk_w", getattr(topk_output, "topk_weights", None)) else: topk_output = self.topk.empty_topk_output(hidden_states.device) @@ -956,6 +960,7 @@ class DeepseekV2MoE(nn.Module): hidden_states=hidden_states, topk_output=topk_output, ) + _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "experts_out", final_hidden_states) if ( hidden_states.shape[0] > 0 @@ -979,12 +984,14 @@ class DeepseekV2MoE(nn.Module): ): final_hidden_states *= self.routed_scaling_factor + _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_out_compact", final_hidden_states) if local_compute_hidden_states is not None: final_hidden_states = restore_cp_local_valid_rows_for_moe( forward_batch, final_hidden_states, local_compute_hidden_states, ) + _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_out_restored", final_hidden_states) return final_hidden_states