Lean tensor dump: residual trajectory only (attn_in/attn_out + q_all), drop kv/MoE

The full-stage+kv dump (1.7GB/file, synchronous in-forward) hung/crashed a scheduler
child. Trim to the residual-stream trajectory (attn_in/attn_out per layer) + q_all +
final_hidden (~tens of MB/forward) — enough to localize the first divergent layer (Stage 1).
Drop dense-space kv_cache/topk (not comparable anyway) and the MoE substages. Dumps go to
local NVMe (/ssd) to avoid beegfs stall.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-17 15:09:39 +00:00
parent 9031a294b3
commit d6aac79eb1
2 changed files with 10 additions and 18 deletions

View File

@@ -2725,18 +2725,7 @@ class NativeSparseAttnBackend(
forward_batch,
layer.layer_id,
"nsa",
{
"q_all": q_all,
"attn_out": attn_output,
"topk_indices": page_table_1,
},
)
_cp_dump(
forward_batch,
layer.layer_id,
"nsa_kv",
{"kv_cache": kv_cache},
big_layers=(0, 19, 39, 59, 77),
{"q_all": q_all},
)
except Exception:
pass

View File

@@ -132,12 +132,15 @@ def fwd_hash(forward_batch, layer_id, stage, t, *, level: int = 3) -> None:
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."""
# Complete-dump accumulation runs regardless of trace level (gated internally
# by SGLANG_NSA_DUMP_DIR + rid 'dump-'): captures every stage at every layer.
try:
dump_tensors(forward_batch, layer_id, stage, {"v": t})
except Exception:
pass
# LEAN dump (gated internally by SGLANG_NSA_DUMP_DIR + rid 'dump-'): only the
# residual-stream trajectory (attn_in/attn_out per layer) — enough to localize
# the first divergent layer. (The full-stage/kv dump was too heavy and hung the
# server.) Stage 2 drills into the localized layer.
if stage in ("attn_in", "attn_out"):
try:
dump_tensors(forward_batch, layer_id, stage, {"v": t})
except Exception:
pass
if trace_level() < level:
return
try: