diff --git a/python/sglang/srt/mem_cache/cp_hicache_trace.py b/python/sglang/srt/mem_cache/cp_hicache_trace.py index e588c0a1e..f944fbc96 100644 --- a/python/sglang/srt/mem_cache/cp_hicache_trace.py +++ b/python/sglang/srt/mem_cache/cp_hicache_trace.py @@ -95,6 +95,44 @@ def knz(t) -> int: 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 ] k=v ...`` line if the gate >= level. diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 91339f571..8d0844b1e 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -12,6 +12,8 @@ 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, ) @@ -2550,6 +2552,18 @@ 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": @@ -2624,6 +2638,19 @@ 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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 01609d1d5..006a669c8 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -73,6 +73,7 @@ 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, @@ -1730,6 +1731,7 @@ 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 @@ -1761,10 +1763,13 @@ 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(