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) <noreply@anthropic.com>

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":

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=<spread>) 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

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