chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Expose CP shared-KV batch timing evidence

Browse Source 
The bs>1 path regressed from the beginning of a run, so log throughput alone is not enough to identify whether scheduler preparation, batch splits, index/top-k, sync compose, HiCache load planning, or valid-row selection is dominating. Add a separate env-gated timing channel with rate and slow-threshold controls so production-like runs can identify the hot operation without enabling verbose structural debug logs.\n\nInstrumentation is intentionally observational: it does not change batch planning, cache residency, transfer semantics, or fast-path eligibility. The timing logs are limited and can be filtered by SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.\n\nConstraint: Need runtime evidence without starting traffic from the agent side.\nConstraint: Existing bs=1 batch-plan behavior remains expected and unchanged.\nRejected: Reuse SGLANG_CP_SHARED_KV_BS_GT1_DEBUG | existing debug logs are structural and already too noisy for timing diagnosis.\nRejected: Add unconditional timing logs | would perturb the hot path and flood long ETE runs.\nConfidence: high\nScope-risk: narrow\nDirective: Keep this timing channel observational; do not use it to gate behavior or silently change batch policy.\nTested: g0034 docker py_compile for changed Python files.\nTested: g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py — 241 passed, 5 warnings, 2 subtests passed.\nNot-tested: Live ETE timing log interpretation under user-driven traffic.\nNot-tested: Nsight correlation with timing events.
This commit is contained in:
laoyao0822
2026-06-06 14:28:26 +08:00
parent 81bfd4bfd3
commit 68defa3923
6 changed files with 417 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/environ.py
View File

@@ -205,6 +205,9 @@ class Envs:
SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)
SGLANG_CP_SHARED_KV_BS_GT1_DEBUG = EnvBool(False)
SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT = EnvInt(128)
SGLANG_CP_SHARED_KV_BS_GT1_TIMING = EnvBool(False)
SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT = EnvInt(256)
SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS = EnvFloat(0.0)
SGLANG_DEBUG_SORT_NVTX = EnvBool(False)
SGLANG_DEBUG_MOE_SORT_NVTX = EnvBool(False)
SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)

49
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -10,8 +10,10 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import (
cp_shared_kv_bs_gt1_timing_start,
get_cp_shared_kv_local_out_cache_loc,
log_cp_draft_shared_kv_debug,
log_cp_shared_kv_bs_gt1_timing,
) # noqa: F401
from sglang.srt.layers.dp_attention import get_attention_cp_group
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -3481,6 +3483,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
reduced immediately; current rows are then inserted into their padded suffix
page slots and non-current tail slack is masked from the returned locs.
"""
timing_start = cp_shared_kv_bs_gt1_timing_start()
total_slots = int(slot_remap.slot_logical_pages.numel())
if prefix_slot_spans is not None and prefix_slot_span is not None:
@@ -3618,6 +3621,29 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
merged_current_spans = (
_merge_slot_spans(current_slot_spans)
if current_slot_spans is not None
else []
)
log_cp_shared_kv_bs_gt1_timing(
"mla_partial_current_compose",
timing_start,
"cp_rank=%s layer=%s cp_size=%s total_slots=%s dense_pages=%s "
"dense_rows=%s prefix_span_pages=%s current_span_pages=%s "
"current_rows=%s materialized_by_ipc=%s kv_dtype=%s",
layout.cp_rank,
layer_id,
layout.cp_size,
total_slots,
slot_remap.dense_num_pages,
int(mixed_kv_cache.shape[0]),
sum(int(end) - int(start) for start, end in prefix_spans),
sum(int(end) - int(start) for start, end in merged_current_spans),
int(current_kv_cache.shape[0]),
materialized_by_ipc,
kv_cache.dtype,
)
return mixed_kv_cache, mixed_locs
@@ -3639,6 +3665,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
nvtx_source: str = "index.partial_current_sync",
) -> tuple[torch.Tensor, torch.Tensor]:
"""Synchronously compose prefix index materialization with current index rows."""
timing_start = cp_shared_kv_bs_gt1_timing_start()
total_slots = int(slot_remap.slot_logical_pages.numel())
if prefix_slot_spans is not None and prefix_slot_span is not None:
@@ -3756,6 +3783,28 @@ def materialize_prefix_and_reuse_current_index_page_slots(
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
merged_current_spans = (
_merge_slot_spans(current_slot_spans)
if current_slot_spans is not None
else []
)
log_cp_shared_kv_bs_gt1_timing(
"index_partial_current_compose",
timing_start,
"cp_rank=%s layer=%s cp_size=%s total_slots=%s dense_pages=%s "
"prefix_span_pages=%s current_span_pages=%s current_rows=%s "
"materialized_by_ipc=%s page_buffer_dtype=%s",
layout.cp_rank,
layer_id,
layout.cp_size,
total_slots,
int(dense_page_buffer.shape[0]),
sum(int(end) - int(start) for start, end in prefix_spans),
sum(int(end) - int(start) for start, end in merged_current_spans),
int(current_index_k.shape[0]),
materialized_by_ipc,
page_buffer.dtype,
)
return dense_page_buffer, slot_remap.dense_pages

91
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -69,6 +69,7 @@ from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.nsa.utils import (
cp_all_gather_rerange_output,
cp_shared_kv_bs_gt1_timing_start,
cp_split_and_rebuild_data,
cp_shared_kv_bs_gt1_debug_enabled,
get_cp_shared_kv_batch_plan,
@@ -77,6 +78,7 @@ from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
log_cp_shared_kv_bs_gt1_debug,
log_cp_shared_kv_bs_gt1_timing,
nsa_use_prefill_cp,
raise_cp_shared_kv_direct_write_error,
select_cp_local_valid_rows_for_cache_write,
@@ -1954,6 +1956,7 @@ class Indexer(MultiPlatformOp):
metadata: BaseIndexerMetadata,
current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> torch.Tensor:
total_timing_start = cp_shared_kv_bs_gt1_timing_start()
cp_metadata = forward_batch.nsa_cp_metadata
assert cp_metadata is not None
batch_size = int(getattr(cp_metadata, "batch_size", 1) or 1)
@@ -2005,6 +2008,7 @@ class Indexer(MultiPlatformOp):
shared_block_tables = None
current_index_kv_for_topk = current_index_kv
current_only_batch = is_current_only_extend_batch(forward_batch)
materialize_timing_start = cp_shared_kv_bs_gt1_timing_start()
if current_index_kv is not None:
current_index_kv_for_topk = None
shared_block_tables = metadata.get_page_table_64()
@@ -2025,7 +2029,21 @@ class Indexer(MultiPlatformOp):
shared_block_tables,
)
)
log_cp_shared_kv_bs_gt1_timing(
"index_topk_batch_materialize",
materialize_timing_start,
"layer=%s bs=%s current_only=%s has_current_index=%s "
"has_shared_index_buffer=%s q_tokens=%s weights_tokens=%s",
layer_id,
batch_size,
current_only_batch,
current_index_kv is not None,
shared_index_buffer is not None,
int(q_fp8.shape[0]),
int(weights.shape[0]),
)
compact_timing_start = cp_shared_kv_bs_gt1_timing_start()
cursor = 0
output_cursor = 0
cp_index: List[Tuple[int, int, int]] = []
@@ -2156,11 +2174,47 @@ class Indexer(MultiPlatformOp):
device=q_fp8.device,
)
if not compact_q_chunks:
log_cp_shared_kv_bs_gt1_timing(
"index_topk_batch_compact",
compact_timing_start,
"layer=%s bs=%s compact_rows=0 output_rows=%s cp_index=%s",
layer_id,
batch_size,
output_cursor,
len(cp_index),
)
log_cp_shared_kv_bs_gt1_timing(
"index_topk_batch_total",
total_timing_start,
"layer=%s bs=%s q_tokens=%s weights_tokens=%s "
"compact_rows=0 output_rows=%s current_only=%s",
layer_id,
batch_size,
int(q_fp8.shape[0]),
int(weights.shape[0]),
output_cursor,
current_only,
)
return result
compact_q = torch.cat(compact_q_chunks, dim=0)
compact_weights = torch.cat(compact_weight_chunks, dim=0)
compact_rows = int(compact_q.shape[0])
log_cp_shared_kv_bs_gt1_timing(
"index_topk_batch_compact",
compact_timing_start,
"layer=%s bs=%s compact_rows=%s output_rows=%s cp_index=%s "
"chunks=%s q_dtype=%s weights_dtype=%s",
layer_id,
batch_size,
compact_rows,
output_cursor,
len(cp_index),
len(compact_q_chunks),
q_fp8.dtype,
weights.dtype,
)
topk_timing_start = cp_shared_kv_bs_gt1_timing_start()
compact_topk = self._get_topk_ragged_with_cp(
forward_batch,
layer_id,
@@ -2176,18 +2230,55 @@ class Indexer(MultiPlatformOp):
actual_seq_q_tensor=None,
actual_seq_q_cu_tensor=None,
)
log_cp_shared_kv_bs_gt1_timing(
"index_topk_batch_kernel",
topk_timing_start,
"layer=%s bs=%s compact_rows=%s output_rows=%s cp_index=%s "
"current_only=%s",
layer_id,
batch_size,
compact_rows,
output_cursor,
len(cp_index),
current_only,
)
if int(compact_topk.shape[0]) != compact_rows:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_gt1_compact_topk_rows_mismatch "
f"expected={compact_rows} got={int(compact_topk.shape[0])}"
)
scatter_timing_start = cp_shared_kv_bs_gt1_timing_start()
compact_cursor = 0
for output_start, valid_q_count in compact_output_spans:
result[output_start : output_start + valid_q_count] = compact_topk[
compact_cursor : compact_cursor + valid_q_count
]
compact_cursor += valid_q_count
log_cp_shared_kv_bs_gt1_timing(
"index_topk_batch_scatter",
scatter_timing_start,
"layer=%s bs=%s compact_rows=%s output_rows=%s spans=%s",
layer_id,
batch_size,
compact_rows,
output_cursor,
len(compact_output_spans),
)
log_cp_shared_kv_bs_gt1_timing(
"index_topk_batch_total",
total_timing_start,
"layer=%s bs=%s q_tokens=%s weights_tokens=%s compact_rows=%s "
"output_rows=%s current_only=%s has_shared_index_buffer=%s",
layer_id,
batch_size,
int(q_fp8.shape[0]),
int(weights.shape[0]),
compact_rows,
output_cursor,
current_only,
shared_index_buffer is not None,
)
return result
def forward_indexer(

148
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -1,5 +1,6 @@
# temp NSA debugging environ
import logging
import time
from dataclasses import dataclass
from itertools import accumulate
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
@@ -32,6 +33,7 @@ logger = logging.getLogger(__name__)
_CP_DRAFT_SHARED_KV_DEBUG_COUNTS = {}
_CP_SHARED_KV_BS_GT1_DEBUG_COUNTS = {}
_CP_SHARED_KV_BS_GT1_TIMING_COUNTS = {}
def cp_shared_kv_bs_gt1_debug_enabled() -> bool:
@@ -56,6 +58,46 @@ def log_cp_shared_kv_bs_gt1_debug(
logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args)
def cp_shared_kv_bs_gt1_timing_enabled() -> bool:
return envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get()
def cp_shared_kv_bs_gt1_timing_start() -> Optional[float]:
if not cp_shared_kv_bs_gt1_timing_enabled():
return None
return time.perf_counter()
def log_cp_shared_kv_bs_gt1_timing(
key: str,
start_time: Optional[float],
message: str,
*args,
limit: Optional[int] = None,
slow_ms: Optional[float] = None,
) -> None:
if start_time is None or not cp_shared_kv_bs_gt1_timing_enabled():
return
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
if slow_ms is None:
slow_ms = float(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.get())
if slow_ms > 0 and elapsed_ms < slow_ms:
return
if limit is None:
limit = envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get()
limit = int(limit)
count = _CP_SHARED_KV_BS_GT1_TIMING_COUNTS.get(key, 0)
if limit > 0 and count >= limit:
return
_CP_SHARED_KV_BS_GT1_TIMING_COUNTS[key] = count + 1
logger.info(
"[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=%.3f " + message,
key,
elapsed_ms,
*args,
)
def log_cp_draft_shared_kv_debug(
key: str,
message: str,
@@ -535,6 +577,7 @@ def build_batch_page_aligned_in_seq_split_plan(
preserves phase1 narrow-output collection for bs>1 without treating the
batch as one long sequence.
"""
timing_start = cp_shared_kv_bs_gt1_timing_start()
if len(extend_lens) != len(prefix_lens):
raise ValueError(
@@ -810,6 +853,23 @@ def build_batch_page_aligned_in_seq_split_plan(
request_last_token_owner,
request_last_token_local_offset,
)
log_cp_shared_kv_bs_gt1_timing(
"batch_plan_build",
timing_start,
"cp_rank=%s cp_size=%s bs=%s page_size=%s extend_sum=%s "
"prefix_sum=%s valid_pages_sum=%s compute_pages_sum=%s "
"compute_padding=%s padding_tokens_sum=%s",
cp_rank,
cp_size,
plan.batch_size,
page_size,
sum(request_extend_lens),
sum(request_prefix_lens),
sum(request_valid_padded_pages),
sum(request_compute_padded_pages),
plan.compute_padding_enabled,
sum(request_compute_padding_tokens),
)
return plan
@@ -854,6 +914,7 @@ def split_tensor_by_cp_batch_plan(
`split_kind="compute"` materializes padded compute rows. Cache writes must
use `split_kind="valid"` so dummy compute rows never receive cache locs.
"""
timing_start = cp_shared_kv_bs_gt1_timing_start()
if mode not in ("1d", "data", "position"):
raise ValueError(f"unsupported CP batch split mode={mode!r}")
@@ -974,10 +1035,27 @@ def split_tensor_by_cp_batch_plan(
compute_padding_enabled,
static_padded_tokens,
)
log_cp_shared_kv_bs_gt1_timing(
f"split_tensor:{mode}:{split_kind}",
timing_start,
"mode=%s split_kind=%s bs=%s input_tokens=%s expected_tokens=%s "
"local_rows=%s compute_padding=%s static_padded=%s dtype=%s shape_tail=%s",
mode,
split_kind,
batch_size,
input_tokens,
expected_tokens,
int(result.shape[0]),
compute_padding_enabled,
static_padded_tokens,
tensor.dtype,
tuple(tensor.shape[1:]),
)
return result
def _get_cp_local_valid_row_indices_cache(forward_batch, plan, device: torch.device):
timing_start = cp_shared_kv_bs_gt1_timing_start()
cached = getattr(forward_batch, "cp_local_valid_row_indices_for_cache_write", None)
cached_expected_rows = getattr(
forward_batch, "cp_local_valid_compute_rows_for_cache_write", None
@@ -1046,6 +1124,16 @@ def _get_cp_local_valid_row_indices_cache(forward_batch, plan, device: torch.dev
indices = torch.empty((0,), device=device, dtype=torch.long)
forward_batch.cp_local_valid_row_indices_for_cache_write = indices
forward_batch.cp_local_valid_compute_rows_for_cache_write = local_cursor
log_cp_shared_kv_bs_gt1_timing(
"valid_row_indices_build",
timing_start,
"bs=%s compute_rows=%s valid_rows=%s device=%s compute_padding=%s",
batch_size,
local_cursor,
int(indices.numel()),
device,
bool(getattr(plan, "compute_padding_enabled", False)),
)
return indices, local_cursor
@@ -1054,6 +1142,7 @@ def select_cp_local_valid_rows_for_cache_write(
local_tensor: torch.Tensor,
) -> torch.Tensor:
"""Drop compute-padding rows before writing CP shared KV into persistent cache."""
timing_start = cp_shared_kv_bs_gt1_timing_start()
plan = get_cp_shared_kv_batch_plan(forward_batch)
if plan is None or not bool(getattr(plan, "compute_padding_enabled", False)):
@@ -1088,7 +1177,19 @@ def select_cp_local_valid_rows_for_cache_write(
return local_tensor
if indices.numel() == 0:
return local_tensor.new_empty((0, *local_tensor.shape[1:]))
return local_tensor.index_select(0, indices)
selected = local_tensor.index_select(0, indices)
log_cp_shared_kv_bs_gt1_timing(
"valid_rows_select",
timing_start,
"local_rows=%s expected_compute_rows=%s valid_rows=%s dtype=%s "
"shape_tail=%s",
local_rows,
expected_compute_rows,
int(indices.numel()),
local_tensor.dtype,
tuple(local_tensor.shape[1:]),
)
return selected
def _pad_cp_request_tensor_for_split(
@@ -1730,6 +1831,7 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
cached = getattr(forward_batch, "cp_local_out_cache_loc", None)
if cached is not None:
return cached
timing_start = cp_shared_kv_bs_gt1_timing_start()
if not getattr(forward_batch, "uses_cp_shared_kv", False):
return None
@@ -1836,6 +1938,18 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
)
if local_out_cache_loc.numel() == 0:
forward_batch.cp_local_out_cache_loc = local_out_cache_loc
log_cp_shared_kv_bs_gt1_timing(
"local_out_cache_loc",
timing_start,
"cp_rank=%s cp_size=%s batch_plan=%s split_tokens=%s "
"out_cache_tokens=%s local_tokens=0 valid_local_tokens=0 page_size=%s",
layout.cp_rank,
layout.cp_size,
batch_plan is not None,
split_tokens,
out_cache_tokens,
layout.page_size,
)
return local_out_cache_loc
if valid_locs.numel() > 0 and not torch.all(layout.owned_by_this_rank(valid_locs)):
@@ -1847,6 +1961,25 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
layout.page_size,
)
forward_batch.cp_local_out_cache_loc = local_out_cache_loc
log_cp_shared_kv_bs_gt1_timing(
"local_out_cache_loc",
timing_start,
"cp_rank=%s cp_size=%s batch_plan=%s compute_padding=%s "
"split_tokens=%s out_cache_tokens=%s local_tokens=%s "
"valid_local_tokens=%s page_size=%s forward_mode=%s",
layout.cp_rank,
layout.cp_size,
batch_plan is not None,
bool(getattr(batch_plan, "compute_padding_enabled", False))
if batch_plan is not None
else False,
split_tokens,
out_cache_tokens,
int(local_out_cache_loc.numel()),
int(valid_locs.numel()),
layout.page_size,
getattr(forward_batch, "forward_mode", None),
)
return local_out_cache_loc
@@ -1864,6 +1997,7 @@ def get_cp_shared_kv_local_physical_out_cache_loc(forward_batch: "ForwardBatch")
cached = getattr(forward_batch, "cp_local_physical_out_cache_loc", None)
if cached is not None:
return cached
timing_start = cp_shared_kv_bs_gt1_timing_start()
local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
if local_out_cache_loc is None:
@@ -1896,6 +2030,18 @@ def get_cp_shared_kv_local_physical_out_cache_loc(forward_batch: "ForwardBatch")
getattr(forward_batch, "token_to_kv_pool", None).__class__.__name__,
)
forward_batch.cp_local_physical_out_cache_loc = physical_out_cache_loc
log_cp_shared_kv_bs_gt1_timing(
"physical_out_loc",
timing_start,
"cp_rank=%s cp_size=%s page_size=%s logical_tokens=%s "
"physical_tokens=%s pool=%s",
layout.cp_rank,
layout.cp_size,
layout.page_size,
local_out_cache_loc.numel(),
physical_out_cache_loc.numel(),
getattr(forward_batch, "token_to_kv_pool", None).__class__.__name__,
)
return physical_out_cache_loc

59
python/sglang/srt/managers/cache_controller.py
View File

@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Set
import torch
from sglang.srt.environ import envs
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
if TYPE_CHECKING:
@@ -49,6 +50,32 @@ from sglang.srt.utils import get_device_module
logger = logging.getLogger(__name__)
device_module = get_device_module()
_CP_SHARED_KV_BS_GT1_CACHE_TIMING_COUNTS: Dict[str, int] = {}
def _cp_shared_kv_bs_gt1_cache_timing(
key: str,
start_time: float,
message: str,
*args,
) -> None:
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
return
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
slow_ms = float(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.get())
if slow_ms > 0 and elapsed_ms < slow_ms:
return
limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get())
count = _CP_SHARED_KV_BS_GT1_CACHE_TIMING_COUNTS.get(key, 0)
if limit > 0 and count >= limit:
return
_CP_SHARED_KV_BS_GT1_CACHE_TIMING_COUNTS[key] = count + 1
logger.info(
"[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=%.3f " + message,
key,
elapsed_ms,
*args,
)
class LayerLoadingEvent:
@@ -1474,6 +1501,14 @@ class HiCacheController:
# Fail closed: returning None lets the caller drop to cache miss
# (cold prefill). Never proceed with a non-matching owner pattern.
if device_indices is None:
_cp_shared_kv_bs_gt1_cache_timing(
"hicache_load_cp_plan",
start_time,
"node_id=%d result=alloc_none pages=%d stages_ms=%s",
node_id,
len(page_owners),
stage_durations_ms,
)
return None
padded_len_expected = sum(
@@ -1572,6 +1607,17 @@ class HiCacheController:
node_id,
)
)
_cp_shared_kv_bs_gt1_cache_timing(
"hicache_load_cp_plan",
start_time,
"node_id=%d result=zero_owned pages=%d visible_indices=%d "
"draft=%s stages_ms=%s",
node_id,
len(page_owners),
int(visible_device_indices.numel()),
self.has_draft_hicache,
[(stage, round(ms, 3)) for stage, ms in stage_durations_ms],
)
return visible_device_indices
host_indices = torch.cat(host_chunks)
@@ -1623,6 +1669,19 @@ class HiCacheController:
draft_host_indices is not None,
[(stage, round(ms, 3)) for stage, ms in stage_durations_ms],
)
_cp_shared_kv_bs_gt1_cache_timing(
"hicache_load_cp_plan",
start_time,
"node_id=%d result=queued pages=%d host_indices=%d "
"physical_indices=%d visible_indices=%d draft=%s stages_ms=%s",
node_id,
len(page_owners),
int(host_indices.numel()),
int(physical_device_indices.numel()),
int(visible_device_indices.numel()),
draft_host_indices is not None,
[(stage, round(ms, 3)) for stage, ms in stage_durations_ms],
)
return visible_device_indices
def move_indices(self, op: CacheOperation, mem_pool_host=None):

68
python/sglang/srt/managers/scheduler.py
View File

@@ -259,6 +259,7 @@ def _cp_draft_shared_kv_debug(message: str, *args) -> None:
_CP_SHARED_KV_BS_GT1_SCHED_DEBUG_COUNTS = {}
_CP_SHARED_KV_BS_GT1_SCHED_TIMING_COUNTS = {}
def _cp_shared_kv_bs_gt1_scheduler_debug(
@@ -276,6 +277,31 @@ def _cp_shared_kv_bs_gt1_scheduler_debug(
logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args)
def _cp_shared_kv_bs_gt1_scheduler_timing(
key: str,
start_time: Optional[float],
message: str,
*args,
) -> None:
if start_time is None or not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
return
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
slow_ms = float(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.get())
if slow_ms > 0 and elapsed_ms < slow_ms:
return
limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT.get())
count = _CP_SHARED_KV_BS_GT1_SCHED_TIMING_COUNTS.get(key, 0)
if limit > 0 and count >= limit:
return
_CP_SHARED_KV_BS_GT1_SCHED_TIMING_COUNTS[key] = count + 1
logger.info(
"[CP_SHARED_KV_BS_GT1_TIMING] event=%s elapsed_ms=%.3f " + message,
key,
elapsed_ms,
*args,
)
_is_npu = is_npu()
@@ -2348,6 +2374,27 @@ class Scheduler(
if dynamic_size is not None:
chunked_prefill_size = dynamic_size
cp_timing_start = (
time.perf_counter()
if (
getattr(self.server_args, "enable_nsa_prefill_cp_shared_kv", False)
and envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get()
)
else None
)
cp_timing_stage_start = cp_timing_start
cp_timing_stages = []
def record_cp_timing_stage(stage: str) -> None:
nonlocal cp_timing_stage_start
if cp_timing_stage_start is None:
return
now = time.perf_counter()
cp_timing_stages.append(
(stage, round((now - cp_timing_stage_start) * 1000.0, 3))
)
cp_timing_stage_start = now
# Prefill policy
adder = PrefillAdder(
self.page_size,
@@ -2375,6 +2422,7 @@ class Scheduler(
prefill_delayer_single_pass=prefill_delayer_single_pass,
dllm_config=self.dllm_config,
)
record_cp_timing_stage("create_adder")
if self.chunked_req is not None:
self.chunked_req.init_next_round_input()
@@ -2447,6 +2495,7 @@ class Scheduler(
else:
self.running_batch.batch_is_full = True
break
record_cp_timing_stage("scan_waiting_queue")
# Update waiting queue
can_run_list: List[Req] = adder.can_run_list
@@ -2487,15 +2536,18 @@ class Scheduler(
self.spec_algorithm,
chunked_req=self.chunked_req,
)
record_cp_timing_stage("init_new_batch")
self.max_prefill_bs = max(self.max_prefill_bs, len(can_run_list))
if self.enable_hierarchical_cache:
# todo (zhiqiang): disable cuda graph execution if hicache loading triggered
new_batch.hicache_consumer_index = (
self.tree_cache.ready_to_load_host_cache()
)
record_cp_timing_stage("hicache_ready_to_load")
try:
new_batch.prepare_for_extend()
record_cp_timing_stage("prepare_for_extend")
except KVCapacityWaitError as exc:
self._release_prefill_adder_locks(
can_run_list, skip_req=chunked_req_before_prepare
@@ -2508,6 +2560,22 @@ class Scheduler(
exc,
)
return None
_cp_shared_kv_bs_gt1_scheduler_timing(
"scheduler_prefill_prepare",
cp_timing_start,
"bs=%s extend_lens=%s prefix_lens=%s out_cache_tokens=%s "
"chunked_req=%s queue_before=%s queue_after=%s stages_ms=%s",
len(can_run_list),
list(getattr(new_batch, "extend_lens", []) or []),
list(getattr(new_batch, "prefix_lens", []) or []),
int(new_batch.out_cache_loc.numel())
if getattr(new_batch, "out_cache_loc", None) is not None
else None,
self.chunked_req is not None,
len(waiting_queue_before_prepare),
len(self.waiting_queue),
cp_timing_stages,
)
if (
getattr(self.server_args, "enable_nsa_prefill_cp_shared_kv", False)
and envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get()