chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Dump: add 'fh' mode (final_hidden only, fresh prefills allowed) for cache-bust vs cache-hit

Browse Source 
rid 'dumpfh-' -> dump ONLY the model's final hidden (last token, post-CP-gather = the
first-token source), allowed on FRESH prefills too -> the cache-bust-vs-cache-hit
discriminator (prefill-side vs decode-side). rid 'dump-' keeps the lean per-layer
trajectory (cache-hit only). Flush moved after the CP gather. dump_diff generalized to
compare any two phases.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-17 16:10:41 +00:00
parent d6aac79eb1
commit 268bcf78da
2 changed files with 45 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
python/sglang/srt/mem_cache/cp_hicache_trace.py
View File

@@ -196,8 +196,12 @@ def cptrace(level: int, tag: str, **fields) -> None:
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."""
"""Return (dump_dir, rid, cprank, mode) if this forward should dump, else None.
Two rid-controlled modes (gated by SGLANG_NSA_DUMP_DIR + EXTEND forward):
'dumpfh-' -> mode 'fh' : final_hidden ONLY (tiny), works for FRESH prefills too
-> the cache-BUST-vs-cache-HIT discriminator.
'dump-' -> mode 'full': lean per-layer trajectory (attn_in/out, q_all), and
only CACHE-HIT extends (skip huge fresh prefills)."""
try:
d = envs.SGLANG_NSA_DUMP_DIR.get()
except Exception:
@@ -209,20 +213,22 @@ def _dump_enabled_rid(forward_batch):
return None
rids = getattr(forward_batch, "rids", None)
rid = str(rids[0]) if rids else ""
if not rid.startswith("dump-"):
if rid.startswith("dumpfh-"):
mode = "fh"
elif rid.startswith("dump-"):
mode = "full"
# 'full' only on cache-hit extends; fresh full prefills are huge + not the repro.
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
else:
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
return d, rid, cprank, mode
def dump_tensors(forward_batch, layer_id, tag, tensors: dict, *, big_layers=None) -> None:
@@ -231,8 +237,9 @@ def dump_tensors(forward_batch, layer_id, tag, tensors: dict, *, big_layers=None
(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
info = _dump_enabled_rid(forward_batch)
if info is None or info[3] != "full":
return # per-layer trajectory only in 'full' mode (cache-hit); 'fh' skips it
try:
import torch
@@ -259,16 +266,23 @@ def dump_flush(forward_batch, *, final_hidden=None, positions=None) -> None:
info = _dump_enabled_rid(forward_batch)
if info is None:
return
d, rid, cprank = info
d, rid, cprank, mode = info
try:
import os
import torch
buf = getattr(forward_batch, "_cp_dump_buf", None)
forward_batch._cp_dump_buf = None
payload = {"layers": buf or {}, "rid": rid, "cprank": cprank}
payload = {"rid": rid, "cprank": cprank, "mode": mode}
if mode == "full":
payload["layers"] = buf or {}
if isinstance(final_hidden, torch.Tensor):
payload["final_hidden"] = final_hidden.detach().to("cpu")
# 'fh' mode: keep only the LAST token's hidden (the first-token source)
# so a fresh full-prefill dump stays tiny; 'full' keeps the extend rows.
fh = final_hidden.detach()
if mode == "fh" and fh.dim() >= 1 and fh.shape[0] > 1:
fh = fh[-1:]
payload["final_hidden"] = fh.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"):

24
python/sglang/srt/models/deepseek_v2.py
View File

@@ -2146,17 +2146,6 @@ 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
@@ -2172,6 +2161,19 @@ class DeepseekV2Model(nn.Module):
forward_batch,
torch.cuda.current_stream(),
)
# Dump flush at model end (gated by SGLANG_NSA_DUMP_DIR + rid 'dump-'/'dumpfh-').
# Placed AFTER the CP gather so final_hidden is the model's actual output (the
# first-token source); for 'fh' mode this is the cache-bust-vs-cache-hit probe.
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 len(aux_hidden_states) == 0:
return hidden_states
return hidden_states, aux_hidden_states