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:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user