Remove CP HiCache debug instrumentation
The CP shared-KV cache-hit corruption bug is fixed and verified end-to-end, so drop the per-layer/per-request trace probes added during the investigation. Each probe was a Python call plus an UNCACHED os.getenv (EnvField.get() re-reads os.environ on every call) in the hot path even when disabled -- ~1000 per forward from the 11 per-layer deepseek probes alone -- plus a latent risk that an accidentally-set SGLANG_NSA_DUMP_DIR would dump tensors in-forward and tank throughput. Delete cp_hicache_trace.py and every reference to it: - deepseek_v2: 11 per-layer fwd_hash probes + the final dump-flush block - nsa_backend: the per-compose NSA tensor-dump block - cp_shared_kv_runtime: 5 imports + the if _cptrace_enabled(...) compose blocks - cache_controller / hiradix_cache / allocator / memory_pool_host: the cptrace/khash/knz/rng round-trip + lifecycle hashes - environ: the SGLANG_CP_HICACHE_KV_TRACE and SGLANG_NSA_DUMP_DIR flags Kept: SGLANG_DEBUG_CP_SHARED_KV / SGLANG_CP_TRANSFER_LOG transfer logging and the x-request-id passthrough (independent of the trace module). Verified: imports clean, 235 mem_cache + reasoning unit tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
@@ -271,16 +271,10 @@ class Envs:
|
||||
SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS = EnvInt(-1)
|
||||
SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False)
|
||||
SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False)
|
||||
# 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)
|
||||
# Logging-only: emit the CP shared-KV sender/worker transfer-partition dumps
|
||||
# (main-KV pages/positions vs NSA-state pages/positions) WITHOUT the side effects
|
||||
# of SGLANG_DEBUG_CP_SHARED_KV (which disables tai materialize -> fail-fast).
|
||||
SGLANG_CP_TRANSFER_LOG = EnvBool(False)
|
||||
# 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)
|
||||
|
||||
@@ -9,11 +9,6 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import cptrace as _cptrace
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import khash as _khash
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import knz as _knz
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import rng as _cprng
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import trace_enabled as _cptrace_enabled
|
||||
from sglang.srt.layers.attention.nsa.utils import (
|
||||
cp_shared_kv_bs_gt1_timing_start,
|
||||
get_cp_shared_kv_local_out_cache_loc,
|
||||
@@ -2076,22 +2071,6 @@ def fill_current_index_page_slots(
|
||||
safe_req = torch.clamp(req_long, min=0, max=max(batch_rows - 1, 0))
|
||||
dense_pages = page_inverse[safe_req, safe_pages.to(torch.long)].to(torch.long)
|
||||
valid_rows = valid_pages & (dense_pages > 0)
|
||||
if _cptrace_enabled(2):
|
||||
# H3: in-range logical locs that fail to map to a real dense page
|
||||
# (dense<=0) right after a reload = stale/wrong page_inverse.
|
||||
_in_range = int(valid_pages.sum().item())
|
||||
_mapped = int(valid_rows.sum().item())
|
||||
_cptrace(
|
||||
2,
|
||||
"remap",
|
||||
total=int(current_locs.numel()),
|
||||
in_range=_in_range,
|
||||
mapped=_mapped,
|
||||
unmapped=_in_range - _mapped,
|
||||
batch_rows=int(batch_rows),
|
||||
capacity=int(capacity),
|
||||
dense=_cprng(dense_pages),
|
||||
)
|
||||
if not torch.any(valid_rows):
|
||||
return dense_page_buffer
|
||||
|
||||
@@ -2242,21 +2221,6 @@ def build_cp_shared_kv_ipc_page_descriptors(
|
||||
invalid_value = torch.full_like(owner_ranks, -1)
|
||||
owner_ranks = torch.where(invalid, invalid_value, owner_ranks)
|
||||
src_page_indices = torch.where(invalid, invalid_value, src_page_indices)
|
||||
if _cptrace_enabled(3):
|
||||
try:
|
||||
_valid = ~invalid
|
||||
_cptrace(
|
||||
3,
|
||||
"ipc_desc",
|
||||
cprank=layout.cp_rank,
|
||||
slots=int(logical_pages.numel()),
|
||||
valid=int(_valid.sum().item()),
|
||||
owners=_cprng(owner_ranks[_valid]),
|
||||
src=_cprng(src_page_indices[_valid]),
|
||||
lpg=_cprng(logical_pages[_valid]),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return owner_ranks.contiguous(), src_page_indices.contiguous()
|
||||
|
||||
|
||||
@@ -4462,28 +4426,10 @@ def get_cp_shared_kv_flattened_request_row_ids(
|
||||
cached = getattr(forward_batch, "cp_flattened_row_ids", None)
|
||||
cached_key = getattr(forward_batch, "cp_flattened_row_ids_key", None)
|
||||
if cached is not None and cached_key == key and cached.device == device:
|
||||
if _cptrace_enabled(2):
|
||||
_cptrace(
|
||||
2,
|
||||
"compose",
|
||||
fb=id(forward_batch),
|
||||
key=key,
|
||||
hit=True,
|
||||
row_ids=_cprng(cached),
|
||||
)
|
||||
return cached
|
||||
row_ids = build_flattened_request_row_ids(seq_lens_cpu, device=device)
|
||||
forward_batch.cp_flattened_row_ids = row_ids
|
||||
forward_batch.cp_flattened_row_ids_key = key
|
||||
if _cptrace_enabled(2):
|
||||
_cptrace(
|
||||
2,
|
||||
"compose",
|
||||
fb=id(forward_batch),
|
||||
key=key,
|
||||
hit=False,
|
||||
row_ids=_cprng(row_ids),
|
||||
)
|
||||
return row_ids
|
||||
|
||||
|
||||
@@ -5662,48 +5608,6 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
|
||||
materialized_by_ipc,
|
||||
kv_cache.dtype,
|
||||
)
|
||||
if _cptrace_enabled(3):
|
||||
try:
|
||||
_live = mixed_locs.reshape(-1)
|
||||
_live = _live[_live >= 0]
|
||||
_gath = (
|
||||
mixed_kv_cache.index_select(0, _live.to(torch.long))
|
||||
if _live.numel()
|
||||
else mixed_kv_cache[:0]
|
||||
)
|
||||
_cptrace(
|
||||
3,
|
||||
"gather_out",
|
||||
cprank=layout.cp_rank,
|
||||
layer=layer_id,
|
||||
total_slots=int(total_slots),
|
||||
live_rows=int(_live.numel()),
|
||||
prefix_pages=sum(int(e) - int(s) for s, e in prefix_spans),
|
||||
current_pages=sum(int(e) - int(s) for s, e in merged_current_spans),
|
||||
h=_khash(_gath),
|
||||
nz=_knz(_gath),
|
||||
)
|
||||
_slp = slot_remap.slot_logical_pages.reshape(-1)
|
||||
for _kind, _spans in (("prefix", prefix_spans), ("current", merged_current_spans)):
|
||||
for _s, _e in _spans:
|
||||
_sl = slot_range_to_token_slice(page_size, int(_s), int(_e))
|
||||
_lp = _slp[int(_s) : int(_e)].to(torch.long)
|
||||
_cptrace(
|
||||
3,
|
||||
"span_owner",
|
||||
cprank=layout.cp_rank,
|
||||
layer=layer_id,
|
||||
kind=_kind,
|
||||
s=int(_s),
|
||||
e=int(_e),
|
||||
owners=_cprng(layout.owner_for_logical_pages(_lp)),
|
||||
phys=_cprng(layout.logical_pages_to_physical(_lp)),
|
||||
lpg=_cprng(_lp),
|
||||
h=_khash(mixed_kv_cache[_sl]),
|
||||
nz=_knz(mixed_kv_cache[_sl]),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return mixed_kv_cache, mixed_locs
|
||||
|
||||
|
||||
|
||||
@@ -2696,23 +2696,6 @@ 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,
|
||||
"nsa",
|
||||
{"q_all": q_all},
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
elif nsa_impl == "flashmla_kv":
|
||||
if (
|
||||
self.nsa_kv_cache_store_fp8
|
||||
|
||||
@@ -23,7 +23,6 @@ from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Set
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import cptrace, khash, knz, rng, trace_enabled
|
||||
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -1142,22 +1141,6 @@ class HiCacheController:
|
||||
self.draft_mem_pool_host.free(draft_host_indices)
|
||||
raise
|
||||
|
||||
if trace_enabled(1):
|
||||
cptrace(
|
||||
1,
|
||||
"backup_reserve",
|
||||
node_id=node_id,
|
||||
logical_len=logical_len,
|
||||
padded_len=padded_len,
|
||||
owned=owned_positions.numel(),
|
||||
host_idx=rng(host_indices),
|
||||
phys=rng(physical_device_indices),
|
||||
page_owners=page_owners.numel()
|
||||
if page_owners is not None
|
||||
else 0,
|
||||
draft=self.has_draft_hicache,
|
||||
draft_host_idx=rng(draft_host_indices),
|
||||
)
|
||||
return HiCacheWriteReservation(
|
||||
metadata=CpHiCacheNodeMetadata(
|
||||
logical_len=logical_len,
|
||||
@@ -1418,21 +1401,6 @@ class HiCacheController:
|
||||
self.io_backend,
|
||||
)
|
||||
grouped_tensors.extend([target_host_indices, target_device_indices])
|
||||
if trace_enabled(1) and layer_id == 0:
|
||||
cptrace(
|
||||
1,
|
||||
"backup_d2h",
|
||||
node_ids=[s.reservation.node_id for s in target_transfer_states],
|
||||
layer=layer_id,
|
||||
host_idx=rng(target_host_indices),
|
||||
phys=rng(target_device_indices),
|
||||
len_match=(
|
||||
target_host_indices.numel()
|
||||
== target_device_indices.numel()
|
||||
),
|
||||
draft=False,
|
||||
)
|
||||
|
||||
if draft_transfer_states:
|
||||
draft_host_indices = self._concat_layer_write_tensors(
|
||||
[state.draft_host_indices for state in draft_transfer_states]
|
||||
@@ -1451,21 +1419,6 @@ class HiCacheController:
|
||||
self.io_backend,
|
||||
)
|
||||
grouped_tensors.extend([draft_host_indices, draft_device_indices])
|
||||
if trace_enabled(1) and layer_id == 0:
|
||||
cptrace(
|
||||
1,
|
||||
"backup_d2h",
|
||||
node_ids=[s.reservation.node_id for s in draft_transfer_states],
|
||||
layer=layer_id,
|
||||
host_idx=rng(draft_host_indices),
|
||||
phys=rng(draft_device_indices),
|
||||
len_match=(
|
||||
draft_host_indices.numel()
|
||||
== draft_device_indices.numel()
|
||||
),
|
||||
draft=True,
|
||||
)
|
||||
|
||||
for tensor in grouped_tensors:
|
||||
self._record_tensor_on_stream(tensor, self.write_stream)
|
||||
|
||||
@@ -1494,21 +1447,6 @@ class HiCacheController:
|
||||
# correct post-store value -- NOT the racy write_stream copy's view. If a
|
||||
# per-layer store/copy ordering race exists, the host copy will diverge
|
||||
# from THIS hash, and the reload hash (post-H2D) will not match it.
|
||||
if trace_enabled(2):
|
||||
for state in target_transfer_states:
|
||||
kv = self.mem_pool_device.get_key_buffer(layer_id)
|
||||
rows = kv[state.physical_device_indices]
|
||||
cptrace(
|
||||
2,
|
||||
"backup_kv_hash",
|
||||
node_id=state.reservation.node_id,
|
||||
layer=layer_id,
|
||||
host=rng(state.host_indices),
|
||||
kvhash=khash(rows),
|
||||
nz=knz(rows),
|
||||
nrows=int(state.physical_device_indices.numel()),
|
||||
)
|
||||
|
||||
for state in final_states:
|
||||
self._append_layer_write_ack(state)
|
||||
|
||||
@@ -1705,26 +1643,6 @@ class HiCacheController:
|
||||
node_device_indices = device_indices[offset : offset + padded_len]
|
||||
offset += padded_len
|
||||
visible_chunks.append(node_device_indices[:valid_len])
|
||||
if trace_enabled(1):
|
||||
# CROSS-RANK invariant: the visible logical-loc sequence for this
|
||||
# node MUST be byte-identical on every CP rank — the cross-rank
|
||||
# gather reads token j from its owner rank at physical(loc_j), so
|
||||
# if rank A and rank R disagree on loc_j, A reads R's slot for a
|
||||
# DIFFERENT token (right bytes, wrong place). Emitted for EVERY
|
||||
# node (incl. zero-owned, before the continue below) so the
|
||||
# analyzer can diff this hash across ranks by node_id. Divergence
|
||||
# here is produced by wall-clock-LRU device eviction giving ranks
|
||||
# different free buckets in alloc_pages_with_owners.
|
||||
_vis = node_device_indices[:valid_len]
|
||||
cptrace(
|
||||
1,
|
||||
"visible_locs_hash",
|
||||
node_id=getattr(node, "id", -1),
|
||||
cprank=self.cp_shared_kv_layout.cp_rank,
|
||||
n=int(_vis.numel()),
|
||||
locs_hash=khash(_vis),
|
||||
locs=rng(_vis),
|
||||
)
|
||||
if self.has_draft_hicache:
|
||||
draft_host_indices = getattr(
|
||||
meta, "draft_host_indices", None
|
||||
@@ -1760,48 +1678,6 @@ class HiCacheController:
|
||||
host_chunks.append(meta.host_indices)
|
||||
if draft_host_indices is not None:
|
||||
draft_host_chunks.append(draft_host_indices)
|
||||
if trace_enabled(1):
|
||||
cptrace(
|
||||
1,
|
||||
"reload_node",
|
||||
node_id=getattr(node, "id", -1),
|
||||
valid_len=valid_len,
|
||||
padded_len=padded_len,
|
||||
owned=owned_positions.numel(),
|
||||
host_idx=rng(meta.host_indices),
|
||||
sel_logical=rng(selected_logical_locs),
|
||||
phys=rng(physical_chunks[-1]),
|
||||
draft_host_idx=rng(draft_host_indices),
|
||||
)
|
||||
# MAKE-OR-BREAK reload invariant (the code above SKIPS this for
|
||||
# perf): every page this rank reloads into must be owned by THIS
|
||||
# rank, i.e. owner_for_logical_pages((loc)//page_size) == cp_rank.
|
||||
# A nonzero bad_owner = alloc_pages_with_owners/padding produced a
|
||||
# per-page owner permutation -> this rank loads its host bytes into
|
||||
# pages owned by another rank -> attention reads garbage (the
|
||||
# aggregate byte-hash can't see this).
|
||||
try:
|
||||
_lay = self.cp_shared_kv_layout
|
||||
_ps = int(getattr(_lay, "page_size", getattr(self, "page_size", 0)))
|
||||
if _ps > 0:
|
||||
_pages = selected_logical_locs // _ps
|
||||
_own = _lay.owner_for_logical_pages(_pages)
|
||||
_bad = int((_own != _lay.cp_rank).sum().item())
|
||||
cptrace(
|
||||
1,
|
||||
"owner_check",
|
||||
node_id=getattr(node, "id", -1),
|
||||
cprank=_lay.cp_rank,
|
||||
owned=int(owned_positions.numel()),
|
||||
bad_owner=_bad,
|
||||
valid_len=valid_len,
|
||||
padded_len=padded_len,
|
||||
tail_pad=padded_len - valid_len,
|
||||
pages=rng(_pages),
|
||||
owners=rng(_own),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
record_stage("build_chunks")
|
||||
|
||||
visible_device_indices = (
|
||||
@@ -2025,19 +1901,6 @@ class HiCacheController:
|
||||
# Compare to backup_kv_hash(node,layer): inequality =
|
||||
# round-trip delivered wrong KV. node_ids is the merged
|
||||
# op's nodes (single node in the c=1 rehit pass).
|
||||
if trace_enabled(2):
|
||||
kv = self.mem_pool_device.get_key_buffer(i)
|
||||
rows = kv[device_indices]
|
||||
cptrace(
|
||||
2,
|
||||
"reload_kv_hash",
|
||||
node_ids=node_ids,
|
||||
layer=i,
|
||||
host=rng(host_indices),
|
||||
kvhash=khash(rows),
|
||||
nz=knz(rows),
|
||||
nrows=int(device_indices.numel()),
|
||||
)
|
||||
producer_event.complete(i)
|
||||
elif op is None and i < self.layer_num:
|
||||
producer_event.complete(i)
|
||||
|
||||
@@ -27,7 +27,6 @@ import torch
|
||||
import triton
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import cptrace, rng, trace_enabled
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2
|
||||
@@ -848,14 +847,6 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
if free_index.numel() == 0:
|
||||
return
|
||||
|
||||
if trace_enabled(2):
|
||||
cptrace(
|
||||
2,
|
||||
"dev_free",
|
||||
tokens=rng(free_index),
|
||||
pages=rng(torch.unique(free_index // self.page_size)),
|
||||
in_free_group=not self.is_not_in_free_group,
|
||||
)
|
||||
if self.is_not_in_free_group:
|
||||
free_page_indices = torch.unique(free_index // self.page_size)
|
||||
self._append_pages_to_owner_buckets(
|
||||
@@ -1093,20 +1084,6 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
|
||||
page_size, dtype=torch.int64, device=self.device
|
||||
).unsqueeze(0)
|
||||
out_indices = (base + offsets).reshape(-1)
|
||||
if trace_enabled(1) and any(c > 0 for c in selected_release_counts):
|
||||
# Tapped the DEFERRED-FREE (release) bucket because a lane's free
|
||||
# bucket was exhausted — only under eviction pressure. A release page
|
||||
# may still be read by an in-flight kernel (deferred free exists to
|
||||
# prevent exactly this), so handing it to an H2D reload here is a
|
||||
# use-after-free / aliasing risk the byte-round-trip hash can't see.
|
||||
cptrace(
|
||||
1,
|
||||
"release_draw",
|
||||
cprank=self.cp_rank,
|
||||
pages=len(page_compute_owners),
|
||||
free_by_owner=selected_free_counts,
|
||||
release_by_owner=selected_release_counts,
|
||||
)
|
||||
self._consume_owner_bucket_prefix(
|
||||
release=False, counts_by_owner=selected_free_counts
|
||||
)
|
||||
|
||||
@@ -1,295 +0,0 @@
|
||||
"""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 _content_key(forward_batch) -> int:
|
||||
"""Stable per-request content fingerprint so the SAME request forwarded from
|
||||
an L1-hit (known-good) and an L2-reload (suspect) can be JOINED across the
|
||||
log WITHOUT rid (the Rust PD gateway strips the client rid -> server mints a
|
||||
UUID). Derived from the extend input-ids + total seq length, which are
|
||||
identical for the same content regardless of where the prefix KV came from.
|
||||
Cached on the ForwardBatch (one compute, shared by every layer/stage).
|
||||
Only meaningful for a SINGLE-request (bs=1) forward; for bs>1 it mixes
|
||||
requests and simply won't join (harmless)."""
|
||||
ck = getattr(forward_batch, "_cp_content_key", None)
|
||||
if ck is not None:
|
||||
return ck
|
||||
ck = -1
|
||||
try:
|
||||
import torch
|
||||
|
||||
ids = getattr(forward_batch, "input_ids", None)
|
||||
sl = getattr(forward_batch, "seq_lens", None)
|
||||
h = khash(ids) if isinstance(ids, torch.Tensor) else 0
|
||||
slsum = int(sl.sum().item()) if isinstance(sl, torch.Tensor) else 0
|
||||
nreq = len(getattr(forward_batch, "rids", []) or [])
|
||||
ck = ((h ^ (slsum * 1000003) ^ (nreq * 998244353)) & 0x7FFFFFFFFFFFFFFF)
|
||||
except Exception:
|
||||
ck = -1
|
||||
try:
|
||||
forward_batch._cp_content_key = ck
|
||||
except Exception:
|
||||
pass
|
||||
return ck
|
||||
|
||||
|
||||
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."""
|
||||
# LEAN dump (gated internally by SGLANG_NSA_DUMP_DIR + rid 'dump-'): only the
|
||||
# residual-stream trajectory (attn_in/attn_out per layer) — enough to localize
|
||||
# the first divergent layer. (The full-stage/kv dump was too heavy and hung the
|
||||
# server.) Stage 2 drills into the localized layer.
|
||||
if stage in ("attn_in", "attn_out"):
|
||||
try:
|
||||
dump_tensors(forward_batch, layer_id, stage, {"v": t})
|
||||
except Exception:
|
||||
pass
|
||||
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
|
||||
# Single-request (bs<=1) forwards only: the content-key join (L1-hit vs
|
||||
# L2-reload of the SAME content) is meaningful only at bs==1 (bs>1 mixes
|
||||
# requests -> ck is a blend), and this also skips the c=24 flood-evict
|
||||
# forwards so the level-3 log stays small and focused on the comparison.
|
||||
_rids0 = getattr(forward_batch, "rids", None)
|
||||
if _rids0 is not None and len(_rids0) > 1:
|
||||
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),
|
||||
ck=_content_key(forward_batch),
|
||||
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
|
||||
|
||||
|
||||
def _dump_enabled_rid(forward_batch):
|
||||
"""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:
|
||||
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 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
|
||||
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, mode
|
||||
|
||||
|
||||
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."""
|
||||
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
|
||||
|
||||
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, mode = info
|
||||
try:
|
||||
import os
|
||||
import torch
|
||||
|
||||
buf = getattr(forward_batch, "_cp_dump_buf", None)
|
||||
forward_batch._cp_dump_buf = None
|
||||
payload = {"rid": rid, "cprank": cprank, "mode": mode}
|
||||
if mode == "full":
|
||||
payload["layers"] = buf or {}
|
||||
if isinstance(final_hidden, torch.Tensor):
|
||||
# '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"):
|
||||
v = getattr(forward_batch, name, None)
|
||||
if isinstance(v, torch.Tensor):
|
||||
payload[name] = v.detach().to("cpu")
|
||||
os.makedirs(d, exist_ok=True)
|
||||
torch.save(payload, os.path.join(d, f"{rid}_cp{cprank}.pt"))
|
||||
except Exception:
|
||||
pass
|
||||
@@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import cptrace, rng, trace_enabled
|
||||
from sglang.srt.managers.cache_controller import (
|
||||
HiCacheController,
|
||||
HiCacheWriteFailure,
|
||||
@@ -1310,25 +1309,6 @@ class HiRadixCache(RadixCache):
|
||||
) -> CpLoadBackPlan:
|
||||
evict_start_time = time.perf_counter()
|
||||
eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan)
|
||||
if trace_enabled(1):
|
||||
# ROOT-CAUSE signal: device load-back eviction victim selection is
|
||||
# keyed on wall-clock last_access_time (per-rank monotonic clock, NOT
|
||||
# replicated across CP ranks — unlike host eviction which uses the
|
||||
# deterministic (priority, node.id)). So ranks can pick DIFFERENT
|
||||
# victims -> free DIFFERENT pages -> diverge alloc_pages_with_owners
|
||||
# -> req_to_token disagrees across ranks. Log the victim id set per
|
||||
# rank for this load op; the analyzer diffs it across ranks. A
|
||||
# cross-rank mismatch confirms the divergence trigger.
|
||||
cptrace(
|
||||
1,
|
||||
"victim_set",
|
||||
node_id=node_id,
|
||||
cprank=self._cp_hicache_cp_rank(),
|
||||
nvictims=len(eviction_plan.victims),
|
||||
victims=sorted(
|
||||
int(getattr(v, "id", -1)) for v in eviction_plan.victims
|
||||
),
|
||||
)
|
||||
plan_elapsed_ms = (time.perf_counter() - evict_start_time) * 1000.0
|
||||
if plan_elapsed_ms >= 1000.0:
|
||||
logger.warning(
|
||||
@@ -2705,21 +2685,6 @@ class HiRadixCache(RadixCache):
|
||||
len(kv_indices),
|
||||
result.metadata.owned_positions.numel(),
|
||||
)
|
||||
if trace_enabled(1):
|
||||
md = result.metadata
|
||||
cptrace(
|
||||
1,
|
||||
"rid_map",
|
||||
rid=getattr(req, "rid", "<unknown>"),
|
||||
node_id=node_id,
|
||||
logical_len=len(kv_indices),
|
||||
owned=md.owned_positions.numel(),
|
||||
host_idx=rng(getattr(md, "host_indices", None)),
|
||||
page_owners=getattr(md, "page_owners", None).numel()
|
||||
if getattr(md, "page_owners", None) is not None
|
||||
else 0,
|
||||
)
|
||||
|
||||
def prepare_write_backup_for_req(self, req) -> None:
|
||||
if self.disable or not self._uses_cp_hicache:
|
||||
return
|
||||
@@ -3034,13 +2999,6 @@ class HiRadixCache(RadixCache):
|
||||
)
|
||||
|
||||
finish_count = int(queue_size.item())
|
||||
if trace_enabled(2) and finish_count > 0:
|
||||
finished_ids = [
|
||||
aid
|
||||
for entry in self.cache_controller.ack_write_queue[:finish_count]
|
||||
for aid in entry[2]
|
||||
]
|
||||
cptrace(2, "write_ack", finished=finish_count, node_ids=finished_ids)
|
||||
logger.debug(
|
||||
"[HiCache-write] writing_check: ongoing=%d ack_queue=%d local_finished=%d sync_finished=%d tp_size=%d",
|
||||
len(self.ongoing_write_through),
|
||||
@@ -3369,8 +3327,6 @@ class HiRadixCache(RadixCache):
|
||||
|
||||
def _evict_backuped(self, node: TreeNode):
|
||||
# GPU -> CPU demotion: no BlockRemoved since block is still reachable via load_back
|
||||
_trace_wp = self._node_host_write_pending(node) if trace_enabled(1) else None
|
||||
_trace_dev = rng(node.value) if trace_enabled(1) else None
|
||||
device_resident_len = self._node_device_resident_len(node)
|
||||
freed_len = self.cache_controller.evict_device(node.value)
|
||||
assert freed_len > 0
|
||||
@@ -3383,17 +3339,6 @@ class HiRadixCache(RadixCache):
|
||||
node.lock_ref,
|
||||
self._node_backuped(node),
|
||||
)
|
||||
if trace_enabled(1):
|
||||
cptrace(
|
||||
1,
|
||||
"evict",
|
||||
node_id=node.id,
|
||||
freed_len=freed_len,
|
||||
device_resident=device_resident_len,
|
||||
backed=self._node_backuped(node),
|
||||
write_pending=_trace_wp,
|
||||
dev=_trace_dev,
|
||||
)
|
||||
node.value = None
|
||||
self._update_leaf_status(node)
|
||||
self._update_host_leaf_status(node)
|
||||
@@ -3712,15 +3657,6 @@ class HiRadixCache(RadixCache):
|
||||
for loaded_node in nodes_to_load:
|
||||
host_len = self._node_host_len(loaded_node)
|
||||
loaded_node.value = device_indices[offset : offset + host_len].clone()
|
||||
if trace_enabled(1):
|
||||
cptrace(
|
||||
1,
|
||||
"reload_assign",
|
||||
node_id=loaded_node.id,
|
||||
host_len=host_len,
|
||||
offset=offset,
|
||||
value=rng(loaded_node.value),
|
||||
)
|
||||
offset += host_len
|
||||
metadata = getattr(loaded_node, "cp_hicache", None)
|
||||
physical_loaded_len += int(
|
||||
@@ -4310,43 +4246,11 @@ class HiRadixCache(RadixCache):
|
||||
child.value = child.value[split_len:].clone()
|
||||
if self._uses_cp_hicache:
|
||||
if self._node_backuped(child):
|
||||
if trace_enabled(1):
|
||||
_orig = child.cp_hicache
|
||||
_orig_owned = _orig.owned_positions.numel()
|
||||
_orig_host = rng(getattr(_orig, "host_indices", None))
|
||||
_orig_po = (
|
||||
_orig.page_owners.numel()
|
||||
if getattr(_orig, "page_owners", None) is not None
|
||||
else 0
|
||||
)
|
||||
new_node.cp_hicache, child.cp_hicache = child.cp_hicache.split(
|
||||
split_len
|
||||
)
|
||||
new_node.host_len = split_len
|
||||
child.host_len = child.host_len - split_len
|
||||
if trace_enabled(1):
|
||||
pf, sf = new_node.cp_hicache, child.cp_hicache
|
||||
cptrace(
|
||||
1,
|
||||
"split",
|
||||
orig_id=child.id,
|
||||
prefix_id=new_node.id,
|
||||
suffix_id=child.id,
|
||||
split_len=split_len,
|
||||
orig_owned=_orig_owned,
|
||||
orig_host=_orig_host,
|
||||
orig_po=_orig_po,
|
||||
pre_owned=pf.owned_positions.numel(),
|
||||
pre_host=rng(getattr(pf, "host_indices", None)),
|
||||
pre_po=pf.page_owners.numel()
|
||||
if getattr(pf, "page_owners", None) is not None
|
||||
else 0,
|
||||
suf_owned=sf.owned_positions.numel(),
|
||||
suf_host=rng(getattr(sf, "host_indices", None)),
|
||||
suf_po=sf.page_owners.numel()
|
||||
if getattr(sf, "page_owners", None) is not None
|
||||
else 0,
|
||||
)
|
||||
elif child.backuped:
|
||||
new_node.host_value = child.host_value[:split_len].clone()
|
||||
child.host_value = child.host_value[split_len:].clone()
|
||||
|
||||
@@ -12,7 +12,6 @@ from typing import Optional, Tuple
|
||||
import psutil
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import cptrace, khash, knz, rng, trace_enabled
|
||||
|
||||
from sglang.jit_kernel.hicache import (
|
||||
can_use_hicache_jit_kernel,
|
||||
@@ -2552,18 +2551,6 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
|
||||
host_page_indices, device_page_indices = self._get_indexer_page_indices(
|
||||
host_indices, device_indices
|
||||
)
|
||||
if trace_enabled(2):
|
||||
_dev = device_pool.index_k_with_scale_buffer[device_layer_slot]
|
||||
_rows = _dev[device_page_indices]
|
||||
cptrace(
|
||||
2,
|
||||
"indexer_backup_hash",
|
||||
layer=layer_id,
|
||||
host=rng(host_indices),
|
||||
kvhash=khash(_rows),
|
||||
nz=knz(_rows),
|
||||
npages=int(device_page_indices.numel()),
|
||||
)
|
||||
use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0
|
||||
if use_kernel:
|
||||
if self.layout == "layer_first":
|
||||
@@ -2638,20 +2625,6 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
|
||||
self._load_indexer_to_device_per_layer(
|
||||
device_pool, host_indices, device_indices, layer_id, io_backend
|
||||
)
|
||||
if trace_enabled(2) and self._is_device_index_layer_active(device_pool, layer_id):
|
||||
_dls = self._device_index_layer_slot(device_pool, layer_id)
|
||||
_, _dpi = self._get_indexer_page_indices(host_indices, device_indices)
|
||||
_rows = device_pool.index_k_with_scale_buffer[_dls][_dpi]
|
||||
cptrace(
|
||||
2,
|
||||
"indexer_reload_hash",
|
||||
layer=layer_id,
|
||||
host=rng(host_indices),
|
||||
kvhash=khash(_rows),
|
||||
nz=knz(_rows),
|
||||
npages=int(_dpi.numel()),
|
||||
)
|
||||
|
||||
def backup_from_device_all_layer(
|
||||
self, device_pool, host_indices, device_indices, io_backend
|
||||
):
|
||||
|
||||
@@ -73,7 +73,6 @@ from sglang.srt.layers.communicator import (
|
||||
enable_moe_dense_fully_dp,
|
||||
get_attn_tp_context,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_hicache_trace import fwd_hash as _cp_fwd_hash
|
||||
from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator
|
||||
from sglang.srt.layers.dp_attention import (
|
||||
get_attention_cp_rank,
|
||||
@@ -775,7 +774,6 @@ class DeepseekV2MoE(nn.Module):
|
||||
hidden_states,
|
||||
)
|
||||
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_postsel", hidden_states)
|
||||
shared_output = None
|
||||
sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
|
||||
sbo_overlap_dispatch_flag = (
|
||||
@@ -788,7 +786,6 @@ class DeepseekV2MoE(nn.Module):
|
||||
if hidden_states.shape[0] > 0:
|
||||
# router_logits: (num_tokens, n_experts)
|
||||
router_logits = self.gate(hidden_states, forward_batch=forward_batch)
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "router_logits", router_logits)
|
||||
if not sbo_enabled_flag:
|
||||
if self.alt_stream is not None:
|
||||
self.alt_stream.wait_stream(torch.cuda.current_stream())
|
||||
@@ -806,8 +803,6 @@ class DeepseekV2MoE(nn.Module):
|
||||
layer_id=self.layer_id,
|
||||
),
|
||||
)
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk_ids", getattr(topk_output, "topk_ids", None))
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk_w", getattr(topk_output, "topk_weights", None))
|
||||
else:
|
||||
topk_output = self.topk.empty_topk_output(hidden_states.device)
|
||||
|
||||
@@ -960,7 +955,6 @@ class DeepseekV2MoE(nn.Module):
|
||||
hidden_states=hidden_states,
|
||||
topk_output=topk_output,
|
||||
)
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "experts_out", final_hidden_states)
|
||||
|
||||
if (
|
||||
hidden_states.shape[0] > 0
|
||||
@@ -984,14 +978,12 @@ class DeepseekV2MoE(nn.Module):
|
||||
):
|
||||
final_hidden_states *= self.routed_scaling_factor
|
||||
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_out_compact", final_hidden_states)
|
||||
if local_compute_hidden_states is not None:
|
||||
final_hidden_states = restore_cp_local_valid_rows_for_moe(
|
||||
forward_batch,
|
||||
final_hidden_states,
|
||||
local_compute_hidden_states,
|
||||
)
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_out_restored", final_hidden_states)
|
||||
|
||||
return final_hidden_states
|
||||
|
||||
@@ -1738,7 +1730,6 @@ class DeepseekV2DecoderLayer(nn.Module):
|
||||
forward_batch,
|
||||
quant_format,
|
||||
)
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "attn_in", hidden_states)
|
||||
|
||||
previous_cp_shared_kv_num_model_layers = getattr(
|
||||
forward_batch, "cp_shared_kv_num_model_layers", None
|
||||
@@ -1770,13 +1761,10 @@ class DeepseekV2DecoderLayer(nn.Module):
|
||||
hidden_states, topk_indices = hidden_states
|
||||
else:
|
||||
topk_indices = None
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "attn_out", hidden_states)
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk", topk_indices)
|
||||
|
||||
hidden_states, residual = self.layer_communicator.prepare_mlp(
|
||||
hidden_states, residual, forward_batch
|
||||
)
|
||||
_cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_in", hidden_states)
|
||||
|
||||
should_allreduce_fusion = (
|
||||
self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
|
||||
@@ -2162,18 +2150,6 @@ class DeepseekV2Model(nn.Module):
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user