diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 7e0ac848d..d1053d5b0 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index d1c9d583b..13adfd7ae 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 47373ffc6..8eca28c58 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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( diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index d7363afa0..93eed48c8 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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 diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 55cc30565..2923f1717 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 9d2e66b3d..cf945ffdb 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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()