[Metrics] Discount queued prefill load by recent cache hits when waiting-queue matching is off (#35248)
This commit is contained in:
@@ -522,6 +522,9 @@ class Envs:
|
|||||||
SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION = EnvInt(4096)
|
SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION = EnvInt(4096)
|
||||||
SGLANG_MAX_NEW_TOKENS_LIMIT = EnvInt(None)
|
SGLANG_MAX_NEW_TOKENS_LIMIT = EnvInt(None)
|
||||||
SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
|
SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
|
||||||
|
# Window for the token-weighted recent cache-hit rate used to estimate
|
||||||
|
# waiting-queue prefill load.
|
||||||
|
SGLANG_CACHE_HIT_RATE_WINDOW_SECONDS = EnvFloat(15.0)
|
||||||
SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES = EnvInt(None)
|
SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES = EnvInt(None)
|
||||||
SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK = EnvFloat(None)
|
SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK = EnvFloat(None)
|
||||||
SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1)
|
SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1)
|
||||||
|
|||||||
@@ -293,6 +293,13 @@ class SchedulePolicy:
|
|||||||
return CacheAgnosticPolicy.FCFS
|
return CacheAgnosticPolicy.FCFS
|
||||||
return self.policy
|
return self.policy
|
||||||
|
|
||||||
|
def waiting_queue_prefix_matched(self, waiting_queue: List[Req]) -> bool:
|
||||||
|
policy = self._determine_active_policy(waiting_queue)
|
||||||
|
return (
|
||||||
|
isinstance(policy, CacheAwarePolicy)
|
||||||
|
or self.tree_cache.supports_fast_match_prefix()
|
||||||
|
)
|
||||||
|
|
||||||
def _validate_and_adjust_policy(
|
def _validate_and_adjust_policy(
|
||||||
self, policy: str, tree_cache: BasePrefixCache
|
self, policy: str, tree_cache: BasePrefixCache
|
||||||
) -> Policy:
|
) -> Policy:
|
||||||
|
|||||||
@@ -2123,6 +2123,10 @@ class Scheduler(
|
|||||||
spec_algorithm=self.spec_algorithm,
|
spec_algorithm=self.spec_algorithm,
|
||||||
get_running_batch=lambda: self.running_batch,
|
get_running_batch=lambda: self.running_batch,
|
||||||
get_waiting_queue=lambda: self.waiting_queue,
|
get_waiting_queue=lambda: self.waiting_queue,
|
||||||
|
waiting_queue_prefix_matched=lambda: self.policy.waiting_queue_prefix_matched(
|
||||||
|
self.waiting_queue
|
||||||
|
),
|
||||||
|
get_recent_cache_hit_rate=lambda: self.metrics_reporter.recent_cache_hit_rate,
|
||||||
get_stats=lambda: self.metrics_reporter.stats,
|
get_stats=lambda: self.metrics_reporter.stats,
|
||||||
get_chunked_req=lambda: self.chunked_req,
|
get_chunked_req=lambda: self.chunked_req,
|
||||||
get_disagg_prefill_bootstrap_queue=lambda: self.disagg_prefill_bootstrap_queue,
|
get_disagg_prefill_bootstrap_queue=lambda: self.disagg_prefill_bootstrap_queue,
|
||||||
|
|||||||
@@ -43,6 +43,8 @@ class SchedulerLoadInquirer:
|
|||||||
spec_algorithm: SpeculativeAlgorithm
|
spec_algorithm: SpeculativeAlgorithm
|
||||||
get_running_batch: Callable
|
get_running_batch: Callable
|
||||||
get_waiting_queue: Callable
|
get_waiting_queue: Callable
|
||||||
|
waiting_queue_prefix_matched: Callable
|
||||||
|
get_recent_cache_hit_rate: Callable
|
||||||
get_stats: Callable
|
get_stats: Callable
|
||||||
get_chunked_req: Callable
|
get_chunked_req: Callable
|
||||||
get_disagg_prefill_bootstrap_queue: Callable
|
get_disagg_prefill_bootstrap_queue: Callable
|
||||||
@@ -76,13 +78,17 @@ class SchedulerLoadInquirer:
|
|||||||
return num_pending_tokens
|
return num_pending_tokens
|
||||||
|
|
||||||
def get_num_waiting_uncached_tokens(self) -> int:
|
def get_num_waiting_uncached_tokens(self) -> int:
|
||||||
"""Get uncached input tokens waiting for prefill compute."""
|
"""Estimate input tokens waiting for prefill compute."""
|
||||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||||
return 0
|
return 0
|
||||||
|
waiting_queue_prefix_matched = self.waiting_queue_prefix_matched()
|
||||||
|
cache_miss_rate = 1.0 - self.get_recent_cache_hit_rate()
|
||||||
num_tokens = 0
|
num_tokens = 0
|
||||||
for req in self.get_waiting_queue():
|
for req in self.get_waiting_queue():
|
||||||
# if match-in-waiting-queue disabled, this metric returns seq_lens
|
if waiting_queue_prefix_matched:
|
||||||
num_tokens += max(0, req.seqlen - req.num_matched_prefix_tokens)
|
num_tokens += max(0, req.seqlen - req.num_matched_prefix_tokens)
|
||||||
|
else:
|
||||||
|
num_tokens += int(req.seqlen * cache_miss_rate)
|
||||||
cr = self.get_chunked_req()
|
cr = self.get_chunked_req()
|
||||||
if cr is not None:
|
if cr is not None:
|
||||||
num_tokens += max(0, cr.seqlen - len(cr.prefix_indices))
|
num_tokens += max(0, cr.seqlen - len(cr.prefix_indices))
|
||||||
|
|||||||
@@ -39,6 +39,28 @@ logger = logging.getLogger(__name__)
|
|||||||
RECORD_STEP_TIME = envs.SGLANG_RECORD_STEP_TIME.get()
|
RECORD_STEP_TIME = envs.SGLANG_RECORD_STEP_TIME.get()
|
||||||
LOG_FORWARD_ITERS = envs.SGLANG_LOG_FORWARD_ITERS.get()
|
LOG_FORWARD_ITERS = envs.SGLANG_LOG_FORWARD_ITERS.get()
|
||||||
ENABLE_METRICS_DEVICE_TIMER = envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.get()
|
ENABLE_METRICS_DEVICE_TIMER = envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.get()
|
||||||
|
CACHE_HIT_RATE_WINDOW_SECONDS = envs.SGLANG_CACHE_HIT_RATE_WINDOW_SECONDS.get()
|
||||||
|
|
||||||
|
|
||||||
|
class _CacheHitRateWindow:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.samples = deque()
|
||||||
|
self.hit_tokens = 0
|
||||||
|
self.total_tokens = 0
|
||||||
|
|
||||||
|
def add(self, hit_tokens: int, total_tokens: int, now: float) -> float:
|
||||||
|
if total_tokens > 0:
|
||||||
|
self.samples.append((now, hit_tokens, total_tokens))
|
||||||
|
self.hit_tokens += hit_tokens
|
||||||
|
self.total_tokens += total_tokens
|
||||||
|
|
||||||
|
cutoff = now - CACHE_HIT_RATE_WINDOW_SECONDS
|
||||||
|
while self.samples and self.samples[0][0] <= cutoff:
|
||||||
|
_, expired_hit_tokens, expired_total_tokens = self.samples.popleft()
|
||||||
|
self.hit_tokens -= expired_hit_tokens
|
||||||
|
self.total_tokens -= expired_total_tokens
|
||||||
|
|
||||||
|
return self.hit_tokens / self.total_tokens if self.total_tokens > 0 else 0.0
|
||||||
|
|
||||||
|
|
||||||
def _decode_total_seq_lens(batch: ScheduleBatch) -> int:
|
def _decode_total_seq_lens(batch: ScheduleBatch) -> int:
|
||||||
@@ -118,6 +140,10 @@ class SchedulerMetricsReporter:
|
|||||||
self._eplb_balancedness_history = [
|
self._eplb_balancedness_history = [
|
||||||
deque(maxlen=window_size) for window_size in EPLB_BALANCEDNESS_WINDOW_SIZES
|
deque(maxlen=window_size) for window_size in EPLB_BALANCEDNESS_WINDOW_SIZES
|
||||||
]
|
]
|
||||||
|
self.cache_hit_rate_window = _CacheHitRateWindow()
|
||||||
|
# Windowed rate for waiting-queue load estimation only; the exported
|
||||||
|
# cache_hit_rate stats keep their per-report semantics.
|
||||||
|
self.recent_cache_hit_rate = 0.0
|
||||||
|
|
||||||
def _init_metrics(
|
def _init_metrics(
|
||||||
self,
|
self,
|
||||||
@@ -654,6 +680,11 @@ class SchedulerMetricsReporter:
|
|||||||
cache_hit_rate = (
|
cache_hit_rate = (
|
||||||
effective_hit_tokens / total_tokens if total_tokens > 0 else 0.0
|
effective_hit_tokens / total_tokens if total_tokens > 0 else 0.0
|
||||||
)
|
)
|
||||||
|
self.recent_cache_hit_rate = self.cache_hit_rate_window.add(
|
||||||
|
effective_hit_tokens,
|
||||||
|
total_tokens,
|
||||||
|
now,
|
||||||
|
)
|
||||||
self.metrics_collector.increment_effective_prefill_tokens(
|
self.metrics_collector.increment_effective_prefill_tokens(
|
||||||
input_tokens=effective_input_tokens,
|
input_tokens=effective_input_tokens,
|
||||||
device_hit_tokens=prefill_stats.log_device_hit_tokens,
|
device_hit_tokens=prefill_stats.log_device_hit_tokens,
|
||||||
|
|||||||
@@ -0,0 +1,70 @@
|
|||||||
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||||
|
from sglang.srt.managers.schedule_policy import (
|
||||||
|
CacheAgnosticPolicy,
|
||||||
|
CacheAwarePolicy,
|
||||||
|
SchedulePolicy,
|
||||||
|
)
|
||||||
|
from sglang.srt.managers.scheduler_components.load_inquirer import (
|
||||||
|
SchedulerLoadInquirer,
|
||||||
|
)
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
class TestSchedulePolicyWaitingQueueMatching(unittest.TestCase):
|
||||||
|
def make_policy(self, policy, supports_fast_match_prefix):
|
||||||
|
schedule_policy = object.__new__(SchedulePolicy)
|
||||||
|
schedule_policy.policy = policy
|
||||||
|
schedule_policy.tree_cache = SimpleNamespace(
|
||||||
|
supports_fast_match_prefix=lambda: supports_fast_match_prefix
|
||||||
|
)
|
||||||
|
return schedule_policy
|
||||||
|
|
||||||
|
def test_cache_agnostic_policy_requires_fast_matching(self):
|
||||||
|
policy = self.make_policy(CacheAgnosticPolicy.FCFS, False)
|
||||||
|
self.assertFalse(policy.waiting_queue_prefix_matched([]))
|
||||||
|
|
||||||
|
policy.tree_cache = SimpleNamespace(supports_fast_match_prefix=lambda: True)
|
||||||
|
self.assertTrue(policy.waiting_queue_prefix_matched([]))
|
||||||
|
|
||||||
|
def test_lpm_queue_limit_can_disable_matching(self):
|
||||||
|
policy = self.make_policy(CacheAwarePolicy.LPM, False)
|
||||||
|
self.assertTrue(policy.waiting_queue_prefix_matched([None] * 128))
|
||||||
|
self.assertFalse(policy.waiting_queue_prefix_matched([None] * 129))
|
||||||
|
|
||||||
|
|
||||||
|
class TestSchedulerLoadInquirer(unittest.TestCase):
|
||||||
|
def make_inquirer(self, waiting_queue_prefix_matched):
|
||||||
|
waiting_req = SimpleNamespace(seqlen=100, num_matched_prefix_tokens=20)
|
||||||
|
chunked_req = SimpleNamespace(seqlen=50, prefix_indices=range(10))
|
||||||
|
return SimpleNamespace(
|
||||||
|
disaggregation_mode=DisaggregationMode.NULL,
|
||||||
|
get_waiting_queue=lambda: [waiting_req],
|
||||||
|
waiting_queue_prefix_matched=lambda: waiting_queue_prefix_matched,
|
||||||
|
get_chunked_req=lambda: chunked_req,
|
||||||
|
get_recent_cache_hit_rate=lambda: 0.75,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_waiting_tokens_are_estimated_when_prefix_matching_is_skipped(self):
|
||||||
|
inquirer = self.make_inquirer(waiting_queue_prefix_matched=False)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
SchedulerLoadInquirer.get_num_waiting_uncached_tokens(inquirer),
|
||||||
|
65,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_waiting_tokens_use_exact_match_when_prefix_matching_is_done(self):
|
||||||
|
inquirer = self.make_inquirer(waiting_queue_prefix_matched=True)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
SchedulerLoadInquirer.get_num_waiting_uncached_tokens(inquirer),
|
||||||
|
120,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -13,6 +13,7 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState
|
|||||||
from sglang.srt.managers.scheduler_components.metrics_reporter import (
|
from sglang.srt.managers.scheduler_components.metrics_reporter import (
|
||||||
PrefillStats,
|
PrefillStats,
|
||||||
SchedulerMetricsReporter,
|
SchedulerMetricsReporter,
|
||||||
|
_CacheHitRateWindow,
|
||||||
)
|
)
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
@@ -136,6 +137,12 @@ class TestForwardPassMetrics(unittest.TestCase):
|
|||||||
self.reporter = _make_reporter(self, self.scheduler)
|
self.reporter = _make_reporter(self, self.scheduler)
|
||||||
self.scheduler.enable_fpm = True
|
self.scheduler.enable_fpm = True
|
||||||
|
|
||||||
|
def test_cache_hit_rate_window_keeps_last_15s_of_tokens(self):
|
||||||
|
window = _CacheHitRateWindow()
|
||||||
|
self.assertEqual(window.add(hit_tokens=20, total_tokens=100, now=0.0), 0.2)
|
||||||
|
self.assertEqual(window.add(hit_tokens=80, total_tokens=100, now=10.0), 0.5)
|
||||||
|
self.assertEqual(window.add(hit_tokens=90, total_tokens=100, now=15.0), 0.85)
|
||||||
|
|
||||||
def _make_batch(self, **overrides):
|
def _make_batch(self, **overrides):
|
||||||
defaults = dict(
|
defaults = dict(
|
||||||
forward_mode=_FakeForwardMode(),
|
forward_mode=_FakeForwardMode(),
|
||||||
|
|||||||
Reference in New Issue
Block a user