From ed42a6dcbc58616d941578ff063d54f99c007f50 Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 17 Jun 2026 14:05:54 +0000 Subject: [PATCH] Honor x-request-id as rid + env-gated NSA tensor dump for L1-vs-L2 reload diff - http_server: chat + generate endpoints read x-request-id header into rid (the Rust PD router drops the client body rid but forwards the header), so the client id reaches forward_batch.rids for exact cross-send join. - cp_hicache_trace.dump_tensors + SGLANG_NSA_DUMP_DIR: torch.save q/composed-KV/ selection/attn_out at layer 0 for rids starting 'dump-' (extend forwards), to diff L1-hit vs L2-reload by relative error (beats fp-nondeterminism that defeats binary hashes). Wired into nsa_backend flashmla_sparse path. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/entrypoints/http_server.py | 13 ++++++ python/sglang/srt/environ.py | 3 ++ .../srt/layers/attention/nsa_backend.py | 21 +++++++++ .../sglang/srt/mem_cache/cp_hicache_trace.py | 45 +++++++++++++++++++ 4 files changed, 82 insertions(+) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 6fff9e6b0..fc5242f32 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -704,6 +704,13 @@ if os.environ.get("DUMPER_SERVER_PORT") == "reuse": ) async def generate_request(obj: GenerateReqInput, request: Request): """Handle a generate request.""" + # The Rust PD router drops the client body `rid` (strict typed struct) but + # forwards the `x-request-id` header. Honor it as the request id so the + # client-set id reaches the prefill forward (forward_batch.rids) for exact + # cross-send join in debugging. The body rid was just minted to a UUID. + _xrid = request.headers.get("x-request-id") + if _xrid: + obj.rid = _xrid if obj.stream: async def stream_results() -> AsyncIterator[bytes]: @@ -1470,6 +1477,12 @@ async def openai_v1_chat_completions( request: ChatCompletionRequest, raw_request: Request ): """OpenAI-compatible chat completion endpoint.""" + # The Rust PD router drops the client body `rid` but forwards the + # `x-request-id` header; honor it so the client id reaches the prefill + # forward (forward_batch.rids) for exact cross-send join in debugging. + _xrid = raw_request.headers.get("x-request-id") + if _xrid: + request.rid = _xrid return await raw_request.app.state.openai_serving_chat.handle_request( request, raw_request ) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 161dc36ce..c178e1c57 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -274,6 +274,9 @@ class Envs: # CP HiCache round-trip KV-corruption tracing. 0=off, 1=structural lifecycle # (rid_map/backup/split/evict/reload), 2=+compose/free/ack timing. SGLANG_CP_HICACHE_KV_TRACE = EnvInt(0) + # Dir to torch.save per-stage attention tensors for L1-hit vs L2-reload diff. + # Empty=off. Only requests whose rid starts with "dump-" are dumped. + SGLANG_NSA_DUMP_DIR = EnvStr("") SGLANG_EAGLE_ACCEPT_DEBUG = EnvBool(False) SGLANG_EAGLE_ACCEPT_DEBUG_INTERVAL = EnvInt(128) SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 633009d52..5d6150f13 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -2712,6 +2712,27 @@ class NativeSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) + # EXACT L1-hit vs L2-reload dump (env+rid gated): the query, the + # composed/dequantized KV the kernel reads, the selection, and the + # output. Offline relerr per tensor localizes reload corruption past + # the fp-nondeterminism that defeats binary hashes. + try: + from sglang.srt.mem_cache.cp_hicache_trace import ( + dump_tensors as _cp_dump, + ) + + _cp_dump( + forward_batch, + layer.layer_id, + { + "q_all": q_all, + "kv_cache": kv_cache, + "page_table_1": page_table_1, + "attn_out": attn_output, + }, + ) + except Exception: + pass elif nsa_impl == "flashmla_kv": if ( self.nsa_kv_cache_store_fp8 diff --git a/python/sglang/srt/mem_cache/cp_hicache_trace.py b/python/sglang/srt/mem_cache/cp_hicache_trace.py index a65a6b191..e39fe9182 100644 --- a/python/sglang/srt/mem_cache/cp_hicache_trace.py +++ b/python/sglang/srt/mem_cache/cp_hicache_trace.py @@ -184,3 +184,48 @@ def cptrace(level: int, tag: str, **fields) -> None: logger.info("[CPTRACE %s] %s", tag, parts) except Exception: 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).""" + try: + d = envs.SGLANG_NSA_DUMP_DIR.get() + except Exception: + d = "" + if not d: + return + 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(): + 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 + os.makedirs(d, exist_ok=True) + torch.save(payload, os.path.join(d, f"{rid}_cp{cprank}_L{layer_id}.pt")) + except Exception: + pass