From 9031a294b35af4233711b0bc532bb246a62a750b Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 17 Jun 2026 14:35:43 +0000 Subject: [PATCH] Complete per-forward tensor dump: all-layer trajectory + final hidden, flush per (rid,rank) dump_tensors now ACCUMULATES per-(layer,stage) into a forward buffer (every fwd_hash stage at every layer + nsa q_all/attn_out/topk all layers, kv_cache at a layer spread); dump_flush writes one file per (rid,rank) at model end with the full trajectory + final hidden + positions + seq/prefix metadata. Gated to cache-hit extends only (skip fresh full prefills) so --repeat 2 dumps just the L1-hit + L2-reload (small). dump_diff joins A(rep1) vs C by idx, relerr per token-order stage -> first divergent layer. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../srt/layers/attention/nsa_backend.py | 11 +- .../sglang/srt/mem_cache/cp_hicache_trace.py | 105 +++++++++++++----- python/sglang/srt/models/deepseek_v2.py | 11 ++ 3 files changed, 96 insertions(+), 31 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 5d6150f13..f4022eb99 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -2724,13 +2724,20 @@ class NativeSparseAttnBackend( _cp_dump( forward_batch, layer.layer_id, + "nsa", { "q_all": q_all, - "kv_cache": kv_cache, - "page_table_1": page_table_1, "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), + ) except Exception: pass elif nsa_impl == "flashmla_kv": diff --git a/python/sglang/srt/mem_cache/cp_hicache_trace.py b/python/sglang/srt/mem_cache/cp_hicache_trace.py index e39fe9182..644177317 100644 --- a/python/sglang/srt/mem_cache/cp_hicache_trace.py +++ b/python/sglang/srt/mem_cache/cp_hicache_trace.py @@ -132,6 +132,12 @@ 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 if trace_level() < level: return try: @@ -186,46 +192,87 @@ def cptrace(level: int, tag: str, **fields) -> None: pass -def dump_tensors(forward_batch, layer_id, tensors: dict, *, layers=(0,)) -> None: - """torch.save per-stage attention tensors for an EXACT L1-hit vs L2-reload diff. - - Gated by SGLANG_NSA_DUMP_DIR; fires ONLY for rids starting 'dump-' (so just the - explicitly-marked probe requests dump, not the whole flood), on the given layers - (default layer 0) and only on EXTEND forwards. Skips any tensor above a size cap - so a giant prefix can't blow up disk. Files are named {rid}_cp{rank}_L{layer}.pt - so the same content sent twice (rid 'dump-X-a' / 'dump-X-b') is joined offline by - relative error per tensor (robust to fp nondeterminism, unlike a binary hash).""" +def _dump_enabled_rid(forward_batch): + """Return (dump_dir, rid, cprank) if this forward should dump, else None. + Gated by SGLANG_NSA_DUMP_DIR + rid starting 'dump-' + EXTEND forward.""" try: d = envs.SGLANG_NSA_DUMP_DIR.get() except Exception: d = "" if not d: + return None + fm = getattr(forward_batch, "forward_mode", None) + if fm is not None and hasattr(fm, "is_extend") and not fm.is_extend(): + return None + rids = getattr(forward_batch, "rids", None) + rid = str(rids[0]) if rids else "" + if not rid.startswith("dump-"): + return None + # Only CACHE-HIT extends (have a reloaded/L1 prefix). Skip fresh full prefills + # (extend_prefix_len==0): they dump the whole sequence (huge) and aren't the + # repro. This lets `--repeat 2` dump only the L1-hit (rep1) + L2-reload (small). + epl = getattr(forward_batch, "extend_prefix_lens_cpu", None) + try: + if epl is not None and int(epl.sum().item()) <= 0: + return None + except Exception: + pass + lay = getattr(forward_batch, "cp_shared_kv_layout", None) + cprank = int(getattr(lay, "cp_rank", -1)) if lay is not None else -1 + return d, rid, cprank + + +def dump_tensors(forward_batch, layer_id, tag, tensors: dict, *, big_layers=None) -> None: + """ACCUMULATE per-(layer,stage) tensors into a per-forward buffer for a COMPLETE + L1-hit vs L2-reload diff. Written once per (rid,rank) by dump_flush() at model end + (token order, comparable; relerr offline beats fp-nondeterminism). 'big' tensors + (e.g. kv_cache, passed with big_layers=) are kept ONLY at those layers to + bound size; everything else at every layer. Gated via _dump_enabled_rid.""" + if _dump_enabled_rid(forward_batch) is None: return + try: + import torch + + buf = getattr(forward_batch, "_cp_dump_buf", None) + if buf is None: + buf = {} + forward_batch._cp_dump_buf = buf + lbuf = buf.setdefault(int(layer_id), {}) + for k, v in tensors.items(): + if big_layers is not None and int(layer_id) not in big_layers: + continue + if isinstance(v, torch.Tensor): + if v.numel() > 200_000_000: + lbuf[f"{tag}/{k}_skipped_numel"] = int(v.numel()) + continue + lbuf[f"{tag}/{k}"] = v.detach().to("cpu") + except Exception: + pass + + +def dump_flush(forward_batch, *, final_hidden=None, positions=None) -> None: + """Write the accumulated per-layer dump buffer to ONE file per (rid,rank) at model + end, plus final hidden / positions / alignment metadata, then clear the buffer.""" + info = _dump_enabled_rid(forward_batch) + if info is None: + return + d, rid, cprank = info try: import os import torch - if layer_id not in layers: - return - fm = getattr(forward_batch, "forward_mode", None) - if fm is not None and hasattr(fm, "is_extend") and not fm.is_extend(): - return - rids = getattr(forward_batch, "rids", None) - rid = str(rids[0]) if rids else "" - if not rid.startswith("dump-"): - return - lay = getattr(forward_batch, "cp_shared_kv_layout", None) - cprank = int(getattr(lay, "cp_rank", -1)) if lay is not None else -1 - payload = {"_meta": {"rid": rid, "cprank": cprank, "layer": int(layer_id)}} - for k, v in tensors.items(): + buf = getattr(forward_batch, "_cp_dump_buf", None) + forward_batch._cp_dump_buf = None + payload = {"layers": buf or {}, "rid": rid, "cprank": cprank} + if isinstance(final_hidden, torch.Tensor): + payload["final_hidden"] = final_hidden.detach().to("cpu") + if isinstance(positions, torch.Tensor): + payload["positions"] = positions.detach().to("cpu") + for name in ("seq_lens_cpu", "extend_prefix_lens_cpu", "extend_seq_lens_cpu"): + v = getattr(forward_batch, name, None) if isinstance(v, torch.Tensor): - if v.numel() > 80_000_000: # ~80M-elem cap; skip huge composed buffers - payload[k + "_skipped_numel"] = int(v.numel()) - continue - payload[k] = v.detach().to("cpu") - elif v is not None: - payload[k] = v + payload[name] = v.detach().to("cpu") os.makedirs(d, exist_ok=True) - torch.save(payload, os.path.join(d, f"{rid}_cp{cprank}_L{layer_id}.pt")) + torch.save(payload, os.path.join(d, f"{rid}_cp{cprank}.pt")) except Exception: pass diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 9e5cefaf9..1bdbed169 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -2146,6 +2146,17 @@ class DeepseekV2Model(nn.Module): else: hidden_states, _ = self.norm(hidden_states, residual) + # Complete-dump flush at model end (gated by SGLANG_NSA_DUMP_DIR + rid 'dump-'): + # write the per-layer buffer + final hidden + metadata to one file per (rid,rank). + try: + from sglang.srt.mem_cache.cp_hicache_trace import dump_flush as _cp_dump_flush + + _cp_dump_flush( + forward_batch, final_hidden=hidden_states, positions=positions + ) + except Exception: + pass + if getattr(forward_batch, "capture_draft_hidden_states", False): forward_batch.draft_hidden_states = hidden_states