sglang/python/sglang/srt/mem_cache/cp_hicache_trace.py

"""Env-gated structured tracing for CP HiCache KV round-trip corruption debugging.

All output is keyed by ``node_id`` (which threads prepare->backup->evict->reload)
plus a ``rid_map`` line that ties the client request id to its node_id. The
companion analyzer joins these lines by node_id and checks four invariants:
  H1 split-misattribution, H2 stale/wrong host ref on reload,
  H3 stale compose cache, H4 page-recycle-before-backup race.

Gate: SGLANG_CP_HICACHE_KV_TRACE (0=off, 1=structural lifecycle, 2=+compose/free/ack).
Off by default and cheap to leave in place: when the level is below the call's
level, nothing is computed or logged.
"""
import logging

from sglang.srt.environ import envs

logger = logging.getLogger(__name__)


def trace_level() -> int:
    try:
        return int(envs.SGLANG_CP_HICACHE_KV_TRACE.get())
    except Exception:
        return 0


def trace_enabled(level: int = 1) -> bool:
    return trace_level() >= level


def rng(t) -> str:
    """Compact summary of a 1-D index tensor/list: (cnt,min,max,first,last).

    Debug-only; .item() syncs are acceptable here. Empty -> (0)."""
    if t is None:
        return "(none)"
    try:
        n = t.numel() if hasattr(t, "numel") else len(t)
        if n == 0:
            return "(0)"
        if hasattr(t, "numel"):
            tf = t.flatten()
            return (
                f"(cnt={n},min={int(tf.min().item())},max={int(tf.max().item())},"
                f"first={int(tf[0].item())},last={int(tf[-1].item())})"
            )
        return f"(cnt={n},min={min(t)},max={max(t)},first={t[0]},last={t[-1]})"
    except Exception as e:  # never let tracing crash the hot path
        return f"(err:{e})"


def _fmt(v) -> str:
    """Render a value with NO internal spaces so the log is space-splittable
    into key=value pairs by the analyzer."""
    if isinstance(v, (list, tuple)):
        return "[" + ",".join(_fmt(x) for x in v) + "]"
    return str(v).replace(" ", "")


def khash(t) -> int:
    """Full-tensor, position-weighted int64 fingerprint of a KV slice.

    Sensitive to BOTH value and position (catches permutation/scatter, unlike a
    plain sum), covers the whole tensor (not a byte sample), and reduces to one
    int. Debug-only: the .item() sync is acceptable when tracing is on. Used to
    compare the KV bytes a node's pages hold at backup (correct, post-store, on
    the default stream) vs what comes back at reload (post-H2D); inequality =
    the round-trip delivered wrong KV (store-vs-copy race or transfer corruption)."""
    try:
        import torch

        if t is None or t.numel() == 0:
            return 0
        b = t.detach().contiguous().view(torch.uint8).reshape(-1).to(torch.int64)
        n = b.numel()
        idx = torch.arange(1, n + 1, device=b.device, dtype=torch.int64)
        return int((((b * idx).sum() ^ (b.sum() * 1000003) ^ n) & 0x7FFFFFFFFFFFFFFF).item())
    except Exception:
        return -1


def knz(t) -> int:
    """Count of nonzero bytes in a KV slice. Real KV is never all-zero, so a
    reload (or backup) slice with nz==0 = uninitialized/zero KV -- a direct,
    self-contained corruption flag (no cross-stage matching needed). Catches the
    observed '0|0|0...' garbage where the store never filled the backed-up pages."""
    try:
        import torch

        if t is None or t.numel() == 0:
            return 0
        b = t.detach().contiguous().view(torch.uint8).reshape(-1)
        return int((b != 0).sum().item())
    except Exception:
        return -1


def fwd_hash(forward_batch, layer_id, stage, t, *, level: int = 3) -> None:
    """Hash a per-layer forward tensor (attn-in/out, topk_indices, MoE-in) to
    localize where a reload forward diverges from a fresh one. Level 3 (so the
    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."""
    if trace_level() < level:
        return
    try:
        import torch

        fm = getattr(forward_batch, "forward_mode", None)
        if fm is not None and hasattr(fm, "is_extend") and not fm.is_extend():
            return
        if isinstance(t, torch.Tensor):
            h, nz, rows = khash(t), knz(t), int(t.shape[0]) if t.dim() else 0
        elif t is None:
            h, nz, rows = 0, 0, 0
        else:
            h, nz, rows = -2, -2, -2  # tuple/quant: not a plain tensor
        rids = getattr(forward_batch, "rids", None)
        lay = getattr(forward_batch, "cp_shared_kv_layout", None)
        cptrace(
            level,
            "fwd_hash",
            rid=(rids[0] if rids else "?"),
            nreq=(len(rids) if rids else 0),
            layer=layer_id,
            stage=stage,
            h=h,
            nz=nz,
            rows=rows,
            cprank=(getattr(lay, "cp_rank", -1) if lay is not None else -1),
        )
    except Exception:
        pass


def cptrace(level: int, tag: str, **fields) -> None:
    """Emit one ``[CPTRACE <tag>] k=v ...`` line if the gate >= level.

    Every value is rendered space-free so ``line.split()`` recovers the
    key=value pairs cleanly in the offline analyzer."""
    if trace_level() < level:
        return
    try:
        parts = " ".join(f"{k}={_fmt(v)}" for k, v in fields.items())
        logger.info("[CPTRACE %s] %s", tag, parts)
    except Exception:
        pass