diff --git a/python/sglang/srt/mem_cache/cp_hicache_trace.py b/python/sglang/srt/mem_cache/cp_hicache_trace.py index dac1ca7f2..085931115 100644 --- a/python/sglang/srt/mem_cache/cp_hicache_trace.py +++ b/python/sglang/srt/mem_cache/cp_hicache_trace.py @@ -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=) 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"): diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 1bdbed169..2ddc4863c 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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