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) <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user