diff --git a/.gitignore b/.gitignore index 54d5170b5..ca994b078 100644 --- a/.gitignore +++ b/.gitignore @@ -283,3 +283,4 @@ tai-kernel/ # Local CP shared-KV/GSM8K debug ledger docs/advanced_features/nsa_prefill_cp_gsm8k_cachehit_temp_findings.md +docs/advanced_features/*_temp*.md diff --git a/python/sglang/srt/entrypoints/openai/serving_base.py b/python/sglang/srt/entrypoints/openai/serving_base.py index 164ccdd5a..e37ec9dc0 100644 --- a/python/sglang/srt/entrypoints/openai/serving_base.py +++ b/python/sglang/srt/entrypoints/openai/serving_base.py @@ -4,6 +4,7 @@ import json import logging import uuid from abc import ABC, abstractmethod +from http import HTTPStatus from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union import orjson @@ -37,6 +38,23 @@ class OpenAIServingBase(ABC): else None ) + @staticmethod + def _streaming_http_error_type_and_status(status_code) -> tuple[str, int]: + if isinstance(status_code, HTTPStatus): + return status_code.name, status_code.value + try: + status = HTTPStatus(int(status_code)) + return status.name, status.value + except Exception: + logger.warning( + "Invalid streaming abort status_code=%r; using INTERNAL_SERVER_ERROR", + status_code, + ) + return ( + HTTPStatus.INTERNAL_SERVER_ERROR.name, + HTTPStatus.INTERNAL_SERVER_ERROR.value, + ) + def _parse_model_parameter(self, model: str) -> Tuple[str, Optional[str]]: """Parse 'base-model:adapter-name' syntax to extract LoRA adapter. diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index cc9abc5eb..7f89a4765 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -813,10 +813,13 @@ class OpenAIServingChat(OpenAIServingBase): code = finish_reason.get( "status_code", HTTPStatus.INTERNAL_SERVER_ERROR ) + err_type, status_code = ( + self._streaming_http_error_type_and_status(code) + ) error = self.create_streaming_error_response( finish_reason.get("message", "Generation aborted."), - code.name, - code.value, + err_type, + status_code, ) yield f"data: {error}\n\n" break diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py index 0bda21907..a64b488f1 100644 --- a/python/sglang/srt/entrypoints/openai/serving_completions.py +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -274,10 +274,13 @@ class OpenAIServingCompletion(OpenAIServingBase): code = finish_reason.get( "status_code", HTTPStatus.INTERNAL_SERVER_ERROR ) + err_type, status_code = self._streaming_http_error_type_and_status( + code + ) error = self.create_streaming_error_response( finish_reason.get("message", "Generation aborted."), - code.name, - code.value, + err_type, + status_code, ) yield f"data: {error}\n\n" break diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 8184f3804..7c6e30450 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -8,7 +8,7 @@ import math import os import threading import time -from dataclasses import dataclass +from dataclasses import dataclass, field from queue import Empty from typing import TYPE_CHECKING, Dict, List, Optional, Tuple @@ -177,6 +177,9 @@ class CpHiCacheNodeMetadata: # page-aligned invariant by using logical_len as the physical length. padded_len: Optional[int] = None draft_host_indices: Optional[torch.Tensor] = None + _owner_page_counts_cache: Dict[int, Tuple[int, ...]] = field( + default_factory=dict, init=False, repr=False + ) def __post_init__(self): if self.logical_len < 0: @@ -245,6 +248,28 @@ class CpHiCacheNodeMetadata: "owned_positions must be sorted and strictly increasing" ) + def owner_page_counts(self, cp_size: int) -> Tuple[int, ...]: + cp_size = int(cp_size) + if cp_size <= 0: + raise ValueError(f"cp_size must be positive, got {cp_size}") + cached = self._owner_page_counts_cache.get(cp_size) + if cached is not None: + return cached + counts = [0 for _ in range(cp_size)] + for owner in self.page_owners.tolist(): + owner = int(owner) + if owner < 0: + continue + if owner >= cp_size: + raise RuntimeError( + "CP HiCache metadata owner exceeds CP size: " + f"owner={owner} cp_size={cp_size}" + ) + counts[owner] += 1 + result = tuple(counts) + self._owner_page_counts_cache[cp_size] = result + return result + @property def valid_len(self) -> int: return self.logical_len @@ -1112,6 +1137,10 @@ class HiRadixCache(RadixCache): def _cp_load_back_node_owner_page_counts( self, node: TreeNode, cp_size: int ) -> Tuple[int, ...]: + metadata = getattr(node, "cp_hicache", None) + if metadata is not None: + return metadata.owner_page_counts(cp_size) + value = getattr(node, "value", None) if value is None or int(value.numel()) == 0: return tuple(0 for _ in range(cp_size)) @@ -1158,18 +1187,41 @@ class HiRadixCache(RadixCache): def _plan_cp_load_back_owner_lane_evictions( self, plan: CpLoadBackPlan ) -> CpLoadBackEvictionPlan: + plan_start_time = time.perf_counter() + last_progress_time = plan_start_time deficits = [max(0, int(v)) for v in plan.deficit_by_owner] cp_size = len(deficits) planned_freed = [0 for _ in range(cp_size)] victims: List[TreeNode] = [] planned_evicted_nodes = set() candidate_nodes = set(getattr(self, "evictable_leaves", set())) + initial_candidate_count = len(candidate_nodes) + iteration = 0 while any(v > 0 for v in deficits): + iteration += 1 best_node = None best_counts = None best_score = None + scanned_candidates = 0 for node in list(candidate_nodes): + scanned_candidates += 1 + if (scanned_candidates & 1023) == 0: + now = time.perf_counter() + if now - last_progress_time >= 5.0: + logger.warning( + "[HiCache-load] slow CP owner-lane eviction planning scan: " + "iteration=%d scanned=%d candidates=%d victims=%d " + "deficits=%s planned_freed=%s elapsed_ms=%.3f", + iteration, + scanned_candidates, + len(candidate_nodes), + len(victims), + deficits, + planned_freed, + (now - plan_start_time) * 1000.0, + ) + last_progress_time = now if node in planned_evicted_nodes: continue if not self._cp_device_node_is_load_back_victim_after_plan( @@ -1223,6 +1275,22 @@ class HiRadixCache(RadixCache): break ancestor = getattr(ancestor, "parent", None) + plan_elapsed_ms = (time.perf_counter() - plan_start_time) * 1000.0 + if plan_elapsed_ms >= 1000.0 or any(v > 0 for v in deficits): + logger.warning( + "[HiCache-load] CP owner-lane eviction planning done: " + "elapsed_ms=%.3f initial_candidates=%d remaining_candidates=%d " + "iterations=%d victims=%d original_deficit_by_owner=%s " + "planned_freed_by_owner=%s remaining_deficit_by_owner=%s", + plan_elapsed_ms, + initial_candidate_count, + len(candidate_nodes), + iteration, + len(victims), + plan.deficit_by_owner, + planned_freed, + deficits, + ) return CpLoadBackEvictionPlan( victims=tuple(victims), planned_freed_by_owner=tuple(planned_freed), @@ -1232,7 +1300,20 @@ class HiRadixCache(RadixCache): def _evict_cp_load_back_owner_lanes( self, plan: CpLoadBackPlan, *, node_id: int ) -> CpLoadBackPlan: + evict_start_time = time.perf_counter() eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan) + plan_elapsed_ms = (time.perf_counter() - evict_start_time) * 1000.0 + if plan_elapsed_ms >= 1000.0: + logger.warning( + "[HiCache-load] CP owner-lane eviction plan ready: node_id=%d " + "elapsed_ms=%.3f victims=%d planned_freed_by_owner=%s " + "remaining_deficit_by_owner=%s", + node_id, + plan_elapsed_ms, + len(eviction_plan.victims), + eviction_plan.planned_freed_by_owner, + eviction_plan.remaining_deficit_by_owner, + ) if any(v > 0 for v in eviction_plan.remaining_deficit_by_owner): logger.warning( "[CP_HICACHE_FALLBACK][cp_load_back_owner_lane_eviction_insufficient] " @@ -1300,13 +1381,21 @@ class HiRadixCache(RadixCache): num_evicted += self._evict_backuped(victim) refreshed = self._refresh_cp_load_back_plan(plan) - logger.debug( - "[HiCache-load] owner-lane device eviction before CP load-back: " - "node_id=%d victims=%s num_evicted=%d required_by_owner=%s " - "before_available_by_owner=%s before_deficit_by_owner=%s " - "after_available_by_owner=%s after_deficit_by_owner=%s " - "planned_freed_by_owner=%s", + evict_elapsed_ms = (time.perf_counter() - evict_start_time) * 1000.0 + log_fn = ( + logger.warning + if evict_elapsed_ms >= 1000.0 + or any(v > 0 for v in refreshed.deficit_by_owner) + else logger.debug + ) + log_fn( + "[HiCache-load] owner-lane device eviction before CP load-back done: " + "node_id=%d elapsed_ms=%.3f victims=%s num_evicted=%d " + "required_by_owner=%s before_available_by_owner=%s " + "before_deficit_by_owner=%s after_available_by_owner=%s " + "after_deficit_by_owner=%s planned_freed_by_owner=%s", node_id, + evict_elapsed_ms, [getattr(node, "id", None) for node in eviction_plan.victims], num_evicted, plan.required_by_owner, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py index 565c658cb..bb4b471b0 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py @@ -319,6 +319,25 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): self.assertEqual(counts, (1, 1, 0, 0)) + def test_device_victim_owner_counts_use_cp_metadata_without_tensor_padding(self): + allocator = _make_allocator(page_size=4, cp_size=4) + cache = _make_cache(allocator) + node = _make_node(14, 140, [0, 0, 2], value=torch.arange(4, 16, dtype=torch.int64)) + + import sglang.srt.mem_cache.hiradix_cache as hiradix_cache + + old_pad = hiradix_cache.pad_token_locs_to_page_boundary + try: + hiradix_cache.pad_token_locs_to_page_boundary = lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("metadata-backed counts must not touch tensor padding") + ) + counts = cache._cp_load_back_node_owner_page_counts(node, cp_size=4) + finally: + hiradix_cache.pad_token_locs_to_page_boundary = old_pad + + self.assertEqual(counts, (2, 0, 1, 0)) + self.assertIs(counts, node.cp_hicache.owner_page_counts(4)) + def test_load_back_plan_fails_closed_without_cp_metadata(self): allocator = _make_allocator() cache = _make_cache(allocator)