From fe59578d1699082f4b161d937775555eb083d866 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Wed, 22 Apr 2026 09:09:50 -0700 Subject: [PATCH 1/2] [None][tests] Add TestServePrefixAwareScheduling base on LMBenchmark/synthetic-multi-round-qa This is the very first multi-round tests with meaningful prefix reuse. Aim to catch functional bugs with all config combinations of the scheduler. Caught SWA+prefix-aware and python-scheduler bugs during the test developement. Add fixes in the PR. Signed-off-by: Simeng Liu --- .../batch_manager/templatedTrie.h | 8 + .../_torch/pyexecutor/scheduler/scheduler.py | 186 +++- .../kv_cache/test_prefix_aware_scheduling.py | 799 ++++++++++++++++++ .../test_lists/test-db/l0_b200.yml | 2 + .../test_lists/test-db/l0_h100.yml | 9 + .../executor/test_kvcache_aware_router.py | 8 +- .../_torch/executor/test_py_scheduler.py | 76 ++ 7 files changed, 1058 insertions(+), 30 deletions(-) create mode 100644 tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py diff --git a/cpp/include/tensorrt_llm/batch_manager/templatedTrie.h b/cpp/include/tensorrt_llm/batch_manager/templatedTrie.h index b0e0138af1b1..e234303866a7 100644 --- a/cpp/include/tensorrt_llm/batch_manager/templatedTrie.h +++ b/cpp/include/tensorrt_llm/batch_manager/templatedTrie.h @@ -191,6 +191,14 @@ class Node auto itr = mNextNodes.find(nkey); if (itr != mNextNodes.end()) { + // Sever the back-edge on the child before erasing its forward entry. + // A shared_ptr to the child may still be held elsewhere (e.g. by + // KVCacheBlock::mLookupNode), and without this step a later + // clearValue on that child would cascade-prune up through a stale + // mPrevNode and assert because this node's mNextNodes no longer + // contains it. Passing an empty NodePtr makes the orphan behave + // as a root: its own cascade-prune stops at itself. + itr->second->setPrevNode(NodePtr{}); mNextNodes.erase(itr); auto const isRoot = mPrevNode.expired(); auto const canBeDeleted = !isRoot && mValue.empty() && mNextNodes.empty(); diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index b06c5b01f66f..d3256251c7b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -3,7 +3,7 @@ from collections import namedtuple from dataclasses import dataclass from enum import Enum -from typing import Optional, Set +from typing import Any, Callable, Optional, Set, TypeAlias, TypeVar from strenum import StrEnum @@ -15,6 +15,31 @@ from ..llm_request import LlmRequest, LlmRequestState RequestList = list[LlmRequest] +PrefixReuseSummary: TypeAlias = tb_internal.batch_manager.PrefixReuseSummary +PrefixSummaryCache: TypeAlias = dict[int, PrefixReuseSummary] +T = TypeVar("T") + + +def _call_with_optional_summary( + fn: Callable[..., T], + *args: Any, + cached_summary: Optional[PrefixReuseSummary] = None, +) -> T: + """Call ``fn(*args)`` and, if supported, pass ``cached_summary`` as a kwarg. + + The nanobind binding for ``get_remaining_blocks_to_completion`` and + ``get_needed_blocks_one_step`` accepts ``cached_summary: PrefixReuseSummary | None = None`` + so the C++ side can skip a redundant radix-tree walk. Test mocks may not + accept this kwarg yet; fall back to the positional form on TypeError to + keep the path working for both. + """ + if cached_summary is None: + return fn(*args) + try: + return fn(*args, cached_summary=cached_summary) + except TypeError: + return fn(*args) + SchedulerOutput = namedtuple( "SchedulerOutput", @@ -435,7 +460,7 @@ def schedule( # No Chunking: Schedule full context # C++ uses getNumTokens(beam=0) which is tokens.size() - numPreDecodedTokens base_tokens = req.get_num_tokens(0) - draft_tokens = req.num_draft_tokens if req.has_draft_tokens else 0 + draft_tokens = req.num_draft_tokens if req.has_draft_tokens() else 0 req_num_tokens = base_tokens + draft_tokens context_compute = _reuse_adjusted_compute( @@ -464,7 +489,7 @@ def schedule( draft_tokens = ( req.num_draft_tokens - if (req.is_last_context_chunk and req.has_draft_tokens) + if (req.is_last_context_chunk and req.has_draft_tokens()) else 0 ) # Compute cost: context compute + draft tokens @@ -788,7 +813,7 @@ def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit ) for req in requests: - if req.is_last_context_chunk and req.has_draft_tokens: + if req.is_last_context_chunk and req.has_draft_tokens(): remainder = req.context_chunk_size % unit_size remaining_space = 0 if remainder == 0 else unit_size - remainder @@ -797,7 +822,11 @@ def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit remaining_space = min(remaining_space, remaining_context_len) if capacity is not None: - remaining_space = min(remaining_space, capacity - num_ctx_tokens) + # capacity - num_ctx_tokens can go negative if the request + # batch has already overshot the compute budget, which would + # wrap draft_discard into an invalid large positive value + # below. Clamp at 0. + remaining_space = max(0, min(remaining_space, capacity - num_ctx_tokens)) num_ctx_tokens += remaining_space draft_discard = req.num_draft_tokens - remaining_space @@ -876,6 +905,11 @@ def schedule( newly_contributed_context_blocks: Set = set() newly_contributed_cross_context_blocks: Set = set() + # Summary caches are populated lazily by _beneficial_to_skip during the + # pending-loop; enough_available_blocks / decrement_reserved_blocks read + # them via .get() so only requests that actually walked the tree pay + # the cost. Mirrors C++ capacityScheduler.cpp's single-walk-per-request. + summary_by_req, cross_summary_by_req = scheduler._make_prefix_summary_caches() if not self.static_batch and skipping_is_relevant: newly_contributed_context_blocks, newly_contributed_cross_context_blocks = ( scheduler._prefill_contributed_blocks(active_requests) @@ -936,6 +970,8 @@ def schedule( req, newly_contributed_context_blocks, newly_contributed_cross_context_blocks, + summary_by_req=summary_by_req, + cross_summary_by_req=cross_summary_by_req, ) ): continue @@ -943,11 +979,25 @@ def schedule( if len(scheduled_requests) >= scheduler.max_num_requests: break + # Read cache AFTER _beneficial_to_skip so the + # now-populated PrefixReuseSummary is reused by + # enough_available_blocks / decrement_reserved_blocks. + # Mirrors the inline pattern used by + # MaxUtilizationPolicy.schedule and the C++ + # single-walk-per-request convention. + req_id = req.py_request_id + cached_summary = summary_by_req.get(req_id) + cached_cross_summary = cross_summary_by_req.get(req_id) + if req.is_context_init_state or req.is_disagg_generation_init_state: - enough_blocks = reserved_blocks.enough_available_blocks(req) + enough_blocks = reserved_blocks.enough_available_blocks( + req, cached_summary=cached_summary + ) enough_cross_blocks = True if reserved_cross_blocks is not None: - enough_cross_blocks = reserved_cross_blocks.enough_available_blocks(req) + enough_cross_blocks = reserved_cross_blocks.enough_available_blocks( + req, cached_summary=cached_cross_summary + ) if not enough_blocks or not enough_cross_blocks: break @@ -964,9 +1014,13 @@ def schedule( uniq_task_ids.add(lora_task_id) scheduled_requests.append(req) - reserved_blocks.decrement_reserved_blocks(req) + reserved_blocks.decrement_reserved_blocks( + req, cached_summary=cached_summary + ) if reserved_cross_blocks is not None: - reserved_cross_blocks.decrement_reserved_blocks(req) + reserved_cross_blocks.decrement_reserved_blocks( + req, cached_summary=cached_cross_summary + ) return scheduled_requests, [] @@ -992,6 +1046,10 @@ def schedule( seen_task_ids: set[int] = set() newly_contributed_context_blocks, _ = scheduler._prefill_contributed_blocks(active_requests) + # Summary cache populated lazily by _beneficial_to_skip; consumed by + # prepare_blocks_if_schedulable. See comment in + # GuaranteedNoEvictPolicy.schedule for rationale. + summary_by_req, _ = scheduler._make_prefix_summary_caches() def is_started_request(req: LlmRequest) -> bool: if not scheduler._can_be_scheduled(req): @@ -1020,7 +1078,10 @@ def is_started_request(req: LlmRequest) -> bool: continue if skipping_is_relevant and scheduler._beneficial_to_skip( - req, newly_contributed_context_blocks, set() + req, + newly_contributed_context_blocks, + set(), + summary_by_req=summary_by_req, ): req_it += 1 continue @@ -1032,6 +1093,7 @@ def is_started_request(req: LlmRequest) -> bool: scheduled_blocks_manager, num_scheduled_peft_pages, seen_task_ids, + cached_summary=summary_by_req.get(req.py_request_id), ) if was_scheduled: @@ -1065,11 +1127,14 @@ def _try_scheduling_request( scheduled_blocks_manager: "MaxUtilizationScheduledBlocksManager", num_scheduled_peft_pages: int, seen_task_ids: set[int], + cached_summary: Optional[PrefixReuseSummary] = None, ) -> bool: if len(scheduled_requests) >= scheduler.max_num_requests: return False - blocks_if_scheduled = scheduled_blocks_manager.prepare_blocks_if_schedulable(req) + blocks_if_scheduled = scheduled_blocks_manager.prepare_blocks_if_schedulable( + req, cached_summary=cached_summary + ) if blocks_if_scheduled is None: return False @@ -1113,22 +1178,43 @@ def __init__(self, kv_cache_manager): stats = kv_cache_manager.get_kv_cache_stats() self.available_blocks: dict[int, int] = dict(stats.num_free_blocks_per_window_size) - def decrement_reserved_blocks(self, req: LlmRequest) -> None: + def decrement_reserved_blocks( + self, req: LlmRequest, cached_summary: Optional[PrefixReuseSummary] = None + ) -> None: """ Decrement available blocks by the blocks needed to complete this request. + + If ``cached_summary`` (a ``PrefixReuseSummary``) is provided, it is forwarded to + ``get_remaining_blocks_to_completion`` so the C++ side skips re-walking the radix tree. C++ reference: scheduledBlocksManager.h:40-46 """ for window_size in self.available_blocks: - needed = self.kv_cache_manager.get_remaining_blocks_to_completion(req, window_size) + needed = _call_with_optional_summary( + self.kv_cache_manager.get_remaining_blocks_to_completion, + req, + window_size, + cached_summary=cached_summary, + ) self.available_blocks[window_size] -= needed - def enough_available_blocks(self, req: LlmRequest) -> bool: + def enough_available_blocks( + self, req: LlmRequest, cached_summary: Optional[PrefixReuseSummary] = None + ) -> bool: """ Check if there are enough available blocks for this request across all window sizes. + + If ``cached_summary`` (a ``PrefixReuseSummary``) is provided, it is forwarded to + ``get_remaining_blocks_to_completion`` so the C++ side skips re-walking the radix tree. C++ reference: scheduledBlocksManager.h:48-57 """ return all( - self.kv_cache_manager.get_remaining_blocks_to_completion(req, ws) <= avail + _call_with_optional_summary( + self.kv_cache_manager.get_remaining_blocks_to_completion, + req, + ws, + cached_summary=cached_summary, + ) + <= avail for ws, avail in self.available_blocks.items() ) @@ -1151,16 +1237,25 @@ def __init__(self, kv_cache_manager, two_steps_look_ahead: bool): window_sizes = set(kv_cache_manager.max_attention_window_vec) self.num_scheduled_blocks: dict[int, int] = {ws: 0 for ws in window_sizes} - def prepare_blocks_if_schedulable(self, req: LlmRequest) -> Optional[dict[int, int]]: + def prepare_blocks_if_schedulable( + self, req: LlmRequest, cached_summary: Optional[PrefixReuseSummary] = None + ) -> Optional[dict[int, int]]: """ Check if request can be scheduled and return new block counts if so. Returns None if request cannot fit. + + If ``cached_summary`` (a ``PrefixReuseSummary``) is provided, it is forwarded to + ``get_needed_blocks_one_step`` so the C++ side skips re-walking the radix tree. C++ reference: scheduledBlocksManager.h:80-100 """ blocks_if_scheduled = {} for window_size, num_scheduled in self.num_scheduled_blocks.items(): - required = self.kv_cache_manager.get_needed_blocks_one_step( - req, self.two_steps_look_ahead, window_size + required = _call_with_optional_summary( + self.kv_cache_manager.get_needed_blocks_one_step, + req, + self.two_steps_look_ahead, + window_size, + cached_summary=cached_summary, ) logger.debug( f"MaxUtilizationScheduler: request ID {req.request_id} " @@ -1330,11 +1425,31 @@ def _prefill_contributed_blocks(self, active_requests: RequestList) -> tuple[set return newly_contributed_context_blocks, newly_contributed_cross_context_blocks + def _make_prefix_summary_caches(self) -> tuple[PrefixSummaryCache, PrefixSummaryCache]: + """Build empty caches for `PrefixReuseSummary` keyed by `py_request_id`. + + The caches are populated **lazily** by `_beneficial_to_skip` when it + walks the radix tree for a pending request. Downstream consumers + (`enough_available_blocks`, `decrement_reserved_blocks`) reuse the + cached summary to avoid re-walking. + + Precomputing for every active request is avoided on purpose: the + ``req.get_unique_tokens(0)`` marshalling copies a potentially very + large token vector across the Python/C++ boundary and, at high + concurrency, dominated scheduler CPU time (verified via py-spy: + >80% of `_schedule`). Lazy population matches the C++ pattern in + ``capacityScheduler.cpp`` which only walks for pending-loop + candidates. + """ + return {}, {} + def _beneficial_to_skip( self, req: LlmRequest, newly_contributed_context_blocks: set, newly_contributed_cross_context_blocks: set, + summary_by_req: Optional[PrefixSummaryCache] = None, + cross_summary_by_req: Optional[PrefixSummaryCache] = None, ) -> bool: """ Check if it's beneficial to skip this request. @@ -1351,10 +1466,19 @@ def _beneficial_to_skip( ctx_new_block = None cross_new_block = None + req_id = req.py_request_id if self.kv_cache_manager is not None and self.kv_cache_manager.enable_block_reuse: - unique_tokens = req.get_unique_tokens(0) - summary = self.kv_cache_manager.analyze_prefix_reuse(unique_tokens, req) + # Use the cached summary when present; otherwise walk the tree + # once and stash the result so later helpers + # (`enough_available_blocks`, `decrement_reserved_blocks`) can + # reuse it. This mirrors C++ capacityScheduler.cpp:295-302. + summary = summary_by_req.get(req_id) if summary_by_req is not None else None + if summary is None: + unique_tokens = req.get_unique_tokens(0) + summary = self.kv_cache_manager.analyze_prefix_reuse(unique_tokens, req) + if summary_by_req is not None: + summary_by_req[req_id] = summary if summary.first_new_block is not None: if summary.first_new_block in newly_contributed_context_blocks: return True @@ -1364,15 +1488,19 @@ def _beneficial_to_skip( self.cross_kv_cache_manager is not None and self.cross_kv_cache_manager.enable_block_reuse ): - encoder_unique_tokens = req.get_encoder_unique_tokens() - if encoder_unique_tokens is not None: - summary = self.cross_kv_cache_manager.analyze_prefix_reuse( - encoder_unique_tokens, req - ) - if summary.first_new_block is not None: - if summary.first_new_block in newly_contributed_cross_context_blocks: - return True - cross_new_block = summary.first_new_block + summary = cross_summary_by_req.get(req_id) if cross_summary_by_req is not None else None + if summary is None: + encoder_unique_tokens = req.get_encoder_unique_tokens() + if encoder_unique_tokens is not None: + summary = self.cross_kv_cache_manager.analyze_prefix_reuse( + encoder_unique_tokens, req + ) + if cross_summary_by_req is not None: + cross_summary_by_req[req_id] = summary + if summary is not None and summary.first_new_block is not None: + if summary.first_new_block in newly_contributed_cross_context_blocks: + return True + cross_new_block = summary.first_new_block # Request is NOT skipped — register contributions so subsequent duplicate # requests can be deferred correctly. diff --git a/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py b/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py new file mode 100644 index 000000000000..cb125ae78a90 --- /dev/null +++ b/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py @@ -0,0 +1,799 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +r"""Integration tests for prefix-aware scheduling with trtllm-serve. + +Launches trtllm-serve with Qwen2-0.5B and runs the LMBenchmark +multi-round QA workload that originally triggered the over-admission bug: + total_num_tokens (13985) should be less than or equal to max_num_tokens (8192) + +The workload sends concurrent requests sharing a system prompt prefix +at various QPS levels, exercising the prefix-aware scheduler's reuse +estimation logic under realistic serving conditions. +""" + +import csv +import importlib.metadata +import json +import math +import os +import queue +import re +import subprocess +import sys +import threading +import time + +import pytest +import requests as req_lib +import yaml +from requests.exceptions import RequestException + +from ..common import get_free_port_in_ci +from ..conftest import llm_models_root +from ..local_venv import PythonVenvRunnerImpl +from ..trt_test_alternative import popen, print_error, print_info + +MODEL_PATH = f"{llm_models_root()}/Qwen2-0.5B" +MODEL_NAME = "Qwen2-0.5B" + +# c1e05a70 "add timeout, skip-ssl-verify, gap-between-requests, itl/throughput +# metrics to real multi-round qa (#36)" by Ziwen Ning, 2026-01-28. +LMBENCHMARK_SHA = "c1e05a708a5a1fd04a9ec09d215edbdcccdf92cb" +LMBENCHMARK_SCRIPT = "" +_LMBENCHMARK_PACKAGE = "benchmark" +_LMBENCHMARK_SCRIPT_IN_PACKAGE = "synthetic-multi-round-qa/multi-round-qa.py" +_LMBENCHMARK_INSTALL_SPEC = ( + f"{_LMBENCHMARK_PACKAGE} @ git+https://github.com/LMCache/LMBenchmark.git@{LMBENCHMARK_SHA}" +) + +# --------------------------------------------------------------------------- +# Scheduler / KV-cache config combinations +# --------------------------------------------------------------------------- +# +# Each entry is a plain dict that will be serialised to YAML and passed via +# `trtllm-serve serve --config`. All combinations test a distinct scheduling +# code path so that functional regressions in any path are caught. +# +# Dimensions covered: +# capacity_scheduler_policy : GUARANTEED_NO_EVICT (default), MAX_UTILIZATION +# enable_chunked_prefill : true, false +# use_python_scheduler : false (default), true +# max_attention_window : unset (full KV), [2048] (SWA) +# disable_overlap_scheduler : false (default), true + +SCHED_CONFIGS = [ + # ── Baseline: GUARANTEED_NO_EVICT + chunked prefill ───────────────────── + pytest.param( + { + "kv_cache_config": {"max_tokens": 200000}, + "enable_chunked_prefill": True, + "print_iter_log": True, + }, + id="guaranteed-chunked", + ), + # ── Aggressive eviction policy ─────────────────────────────────────────── + pytest.param( + { + "kv_cache_config": {"max_tokens": 200000}, + "enable_chunked_prefill": True, + "scheduler_config": {"capacity_scheduler_policy": "MAX_UTILIZATION"}, + "print_iter_log": True, + }, + id="max-util-chunked", + ), + # ── No chunked prefill (full-context scheduling) ───────────────────────── + pytest.param( + { + "kv_cache_config": {"max_tokens": 200000}, + "enable_chunked_prefill": False, + "print_iter_log": True, + }, + id="guaranteed-no-chunked", + ), + # ── Pure-Python scheduler (parity with C++ path) ───────────────────────── + pytest.param( + { + "kv_cache_config": {"max_tokens": 200000}, + "enable_chunked_prefill": True, + "scheduler_config": {"use_python_scheduler": True}, + "print_iter_log": True, + }, + id="python-scheduler", + ), + # ── Sliding-window attention (SWA, 2048-token window) ──────────────────── + # Qwen2-0.5B has 24 layers; max_attention_window=[2048] broadcasts to all. + # The 1000-token system prompt fits within the window, but older history + # rounds are evicted, exercising the SWA eviction path in the scheduler. + pytest.param( + { + "kv_cache_config": { + "max_tokens": 200000, + "max_attention_window": [2048], + }, + "enable_chunked_prefill": True, + "print_iter_log": True, + }, + id="swa-chunked", + ), + # ── Overlap scheduler disabled ──────────────────────────────────────────── + pytest.param( + { + "kv_cache_config": {"max_tokens": 200000}, + "enable_chunked_prefill": True, + "disable_overlap_scheduler": True, + "print_iter_log": True, + }, + id="no-overlap-chunked", + ), + # ── KV-cache host offload (secondary memory tier) ──────────────────────── + # Exercises the offload path: on eviction from device, blocks move to + # host memory instead of being dropped, and prefix-aware scheduling + # must account for the larger effective reuse pool. host_cache_size is + # deliberately small (1 GiB) relative to chat-history volume so that + # offload + promotion both happen during the QPS sweep. + pytest.param( + { + "kv_cache_config": { + "max_tokens": 200000, + "host_cache_size": 1024 * 1024 * 1024, + "free_gpu_memory_fraction": 0.2, + }, + "enable_chunked_prefill": True, + "print_iter_log": True, + }, + id="offload-chunked", + ), + pytest.param( + { + "kv_cache_config": { + "max_tokens": 200000, + "host_cache_size": 1024 * 1024 * 1024, + "free_gpu_memory_fraction": 0.2, + }, + # enable_chunked_prefill defaults to False here. The multi-round-QA + # workload sends prompts up to ~21k tokens (system_prompt=1000 + + # chat_history=20000); without chunking they must fit in a single + # forward pass, so max_num_tokens must exceed the worst-case prompt + # length. Default is 8192, which rejects the request at admission + # with "sum of prompt length (21044) should not exceed + # max_num_tokens (8192)". Size to 32768 to accept the workload. + "max_num_tokens": 32768, + "print_iter_log": True, + }, + id="offload-no-chunked", + ), +] + +# --------------------------------------------------------------------------- +# Error / NaN detection helpers +# --------------------------------------------------------------------------- + +_ERROR_PATTERNS = [ + "should be less than or equal to max_num_tokens", # original over-admission bug + "CUDA error", + "Traceback (most recent call last)", # uncaught Python exception + "AssertionError", +] + +_NAN_RE = re.compile(r"\bnan\b", re.IGNORECASE) + + +def _check_server_errors(server_log: str) -> str | None: + """Return the first error line found in *server_log*, or None if clean.""" + try: + with open(server_log) as f: + for line in f: + if any(p in line for p in _ERROR_PATTERNS): + return line.strip() + except OSError: + pass + return None + + +def _tail_log(log_path: str, n: int = 40) -> str: + """Return the last *n* lines of *log_path* as a single string.""" + try: + with open(log_path) as f: + lines = f.readlines() + return "".join(lines[-n:]) + except OSError: + return "(log not available)" + + +def _assert_no_server_errors(server_log: str) -> None: + """Assert the server log is error-free; include log tail on failure.""" + err = _check_server_errors(server_log) + assert err is None, ( + f"Server error detected: {err!r}\n" + f"--- last lines of {server_log} ---\n" + f"{_tail_log(server_log)}" + ) + + +# --------------------------------------------------------------------------- +# Server lifecycle helpers +# --------------------------------------------------------------------------- + + +def _wait_for_server_ready(proc, port, timeout=300, interval=2, server_log=None): + """Wait for trtllm-serve /health to return 200. + + If *server_log* is provided, the log is scanned for error patterns on each + poll iteration so that crashes during startup are caught immediately. + """ + url = f"http://127.0.0.1:{port}/health" + start = time.time() + while time.time() - start < timeout: + try: + if req_lib.get(url, timeout=interval).status_code == 200: + print_info(f"Server ready in {time.time() - start:.1f}s on port {port}") + return + except RequestException: + pass + rc = proc.poll() + if rc is not None and rc != 0: + tail = _tail_log(server_log) if server_log else "" + raise RuntimeError(f"trtllm-serve exited unexpectedly with code {rc}.\n{tail}") + if server_log: + err = _check_server_errors(server_log) + if err: + raise RuntimeError( + f"trtllm-serve reported error during startup: {err!r}\n{_tail_log(server_log)}" + ) + time.sleep(interval) + raise TimeoutError(f"trtllm-serve not ready within {timeout}s") + + +def _make_server_cmd(port: int, config_path: str) -> list[str]: + """Build the trtllm-serve command used by all tests.""" + return [ + "trtllm-serve", + "serve", + MODEL_PATH, + "--host", + "0.0.0.0", + "--port", + str(port), + "--config", + config_path, + ] + + +def _write_config(tmp_path, cfg: dict, name: str = "config.yml") -> str: + """Serialise *cfg* to YAML in *tmp_path* and return the file path.""" + path = str(tmp_path / name) + with open(path, "w") as f: + yaml.dump(cfg, f) + return path + + +# --------------------------------------------------------------------------- +# Benchmark helpers +# --------------------------------------------------------------------------- + + +def _run_lmbenchmark( + port, + output_csv, + qps=0.5, + num_users=4, + num_rounds=3, + system_prompt=500, + chat_history=2000, + answer_len=20, + duration=30, + server_log=None, +): + """Run the LMBenchmark multi-round-qa script and return the exit code. + + stdout and stderr are drained in background threads so that the main + watchdog loop can run deadline / server-log / /health checks on a fixed + cadence regardless of how chatty the benchmark is. A blocking readline + in the main loop would otherwise stall indefinitely when the benchmark + goes quiet (which is exactly what happens when the server stalls), + defeating the point of the watchdog. + """ + script_dir = os.path.dirname(LMBENCHMARK_SCRIPT) + cmd = [ + sys.executable, + LMBENCHMARK_SCRIPT, + "--num-users", + str(num_users), + "--num-rounds", + str(num_rounds), + "--qps", + str(qps), + "--shared-system-prompt", + str(system_prompt), + "--user-history-prompt", + str(chat_history), + "--answer-len", + str(answer_len), + "--model", + MODEL_NAME, + "--base-url", + f"http://localhost:{port}", + "--init-user-id", + "1", + "--output", + output_csv, + "--log-interval", + "30", + "--time", + str(duration), + ] + print_info(f"Running LMBenchmark: {' '.join(cmd)}") + + t_start = time.time() + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, cwd=script_dir + ) + + stdout_q: queue.Queue = queue.Queue() + stderr_lines: list[str] = [] + + def _drain_stdout(): + try: + for line in iter(proc.stdout.readline, ""): + stdout_q.put(line) + finally: + stdout_q.put(None) # EOF sentinel + + def _drain_stderr(): + for line in iter(proc.stderr.readline, ""): + stderr_lines.append(line) + + stdout_thread = threading.Thread(target=_drain_stdout, daemon=True) + stderr_thread = threading.Thread(target=_drain_stderr, daemon=True) + stdout_thread.start() + stderr_thread.start() + + # Grace window above the scripted duration. LMBenchmark keeps requests + # in flight after `--time` expires, and at higher QPS the backlog the + # server still has to drain scales roughly linearly with offered load. + # Observed worst-case drain on H100 PCIe + Qwen2-0.5B (5 scheduler + # variants, 3 stages each): drain ≈ 1.7·qps + 1 s. The coefficient below + # gives ≥2.8× margin over the observed worst case at every tested qps + # while still catching genuine stalls within a few minutes. + grace_s = max(45, int(qps * 5)) + deadline = time.time() + duration + grace_s + last_health_ok = time.time() + last_poll = time.time() + + def _kill(reason: str) -> int: + elapsed_at_kill = time.time() - t_start + print_error( + f"{reason} [timing] qps={qps} duration={duration}s " + f"elapsed={elapsed_at_kill:.1f}s drain={elapsed_at_kill - duration:+.1f}s " + f"grace_s={grace_s}s" + ) + proc.kill() + proc.wait() + stdout_thread.join(timeout=5) + stderr_thread.join(timeout=5) + return -1 + + def _drain_queued_stdout() -> int | None: + """Pop all queued stdout lines; return -1 if NaN seen, else None.""" + while True: + try: + line = stdout_q.get_nowait() + except queue.Empty: + return None + if line is None: + return None # EOF sentinel; stdout thread done + line = line.rstrip() + if line and _NAN_RE.search(line): + return _kill(f"NaN detected in benchmark output: {line!r}") + + try: + while proc.poll() is None: + rc = _drain_queued_stdout() + if rc is not None: + return rc + + now = time.time() + if now > deadline: + return _kill( + f"LMBenchmark exceeded deadline ({duration + grace_s}s, " + f"duration={duration}s + grace={grace_s}s at qps={qps}); " + "server may have stalled" + ) + + # Periodic checks (every 5 s). Do not gate these on stdout + # activity -- the whole point of the watchdog is to fire even + # when the benchmark is quiet. + if now - last_poll >= 5: + last_poll = now + + if server_log and _check_server_errors(server_log): + return _kill("Server error detected in log, killing benchmark") + + try: + # 5s (was 2s) — under high offered load the async /health + # endpoint can be briefly starved by the busy event loop. + resp = req_lib.get(f"http://localhost:{port}/health", timeout=5) + if resp.status_code == 200: + last_health_ok = now + except Exception: + pass + # 60s (was 30s) — extreme QPS stages need more slack before + # declaring the server stalled; genuine hangs still get caught. + if now - last_health_ok > 60: + return _kill("Server stopped responding to /health (stalled)") + + time.sleep(0.2) + + except Exception: + proc.kill() + proc.wait() + stdout_thread.join(timeout=5) + stderr_thread.join(timeout=5) + raise + + # Benchmark has exited. Wait for drain threads to finish so we see + # every remaining line, then surface any leftover output. + stdout_thread.join(timeout=5) + stderr_thread.join(timeout=5) + + while True: + try: + line = stdout_q.get_nowait() + except queue.Empty: + break + if line is None: + continue + line = line.rstrip() + if not line: + continue + if _NAN_RE.search(line): + print_error(f"NaN detected in benchmark output: {line!r}") + return -1 + print_info(f" [benchmark] {line}") + + elapsed = time.time() - t_start + drain = elapsed - duration + print_info( + f"[timing] qps={qps} duration={duration}s elapsed={elapsed:.1f}s " + f"drain={drain:+.1f}s grace_s={grace_s}s " + f"grace_used_frac={(drain / grace_s if grace_s else float('nan')):.2f}" + ) + + if proc.returncode != 0: + print_error(f"LMBenchmark exited with code {proc.returncode}") + for line in "".join(stderr_lines).strip().split("\n")[-10:]: + if line: + print_error(f" [stderr] {line}") + return proc.returncode + + +def _parse_csv_metrics(csv_path): + """Parse LMBenchmark output CSV and return TTFT metrics. + + Returns a dict with keys: + num_requests – total rows + nan_count – rows with NaN/Inf ttft (should be 0) + ttft_count – rows with a usable (parseable, finite) ttft value + ttft_avg, ttft_p50, ttft_p99 – latency percentiles (when available) + """ + if not os.path.exists(csv_path): + return None + with open(csv_path) as f: + rows = list(csv.DictReader(f)) + if not rows: + return None + + nan_count = 0 + ttfts = [] + for r in rows: + raw = r.get("ttft") + if not raw: + continue + try: + v = float(raw) + except ValueError: + nan_count += 1 + continue + if math.isnan(v) or math.isinf(v): + nan_count += 1 + else: + ttfts.append(v) + + if not ttfts: + return None + + ttfts.sort() + return { + "num_requests": len(rows), + "nan_count": nan_count, + "ttft_count": len(ttfts), + "ttft_avg": sum(ttfts) / len(ttfts), + "ttft_p50": ttfts[len(ttfts) // 2], + "ttft_p99": ttfts[int(len(ttfts) * 0.99)], + } + + +def _run_and_assert_stage( + label: str, + port: int, + output_csv: str, + server_log: str, + *, + qps: float, + num_users: int, + num_rounds: int, + system_prompt: int, + chat_history: int, + answer_len: int, + duration: int, +) -> dict: + """Run one LMBenchmark stage, assert success, and return the metrics dict.""" + rc = _run_lmbenchmark( + port=port, + output_csv=output_csv, + qps=qps, + num_users=num_users, + num_rounds=num_rounds, + system_prompt=system_prompt, + chat_history=chat_history, + answer_len=answer_len, + duration=duration, + server_log=server_log, + ) + assert rc == 0, f"{label} failed with rc={rc}. See {server_log}\n{_tail_log(server_log)}" + metrics = _parse_csv_metrics(output_csv) + assert metrics and metrics.get("num_requests", 0) > 0, ( + f"{label} produced no completed requests (csv={output_csv})" + ) + assert metrics["nan_count"] == 0, f"{label} produced {metrics['nan_count']} NaN ttft values" + assert metrics["ttft_count"] > 0, ( + f"{label} produced {metrics['num_requests']} rows but zero usable ttft " + f"values — benchmark output is broken (csv={output_csv})" + ) + return metrics + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _resolve_lmbenchmark_script() -> str: + """Return the benchmark script from the pinned installed LMBenchmark package.""" + try: + dist = importlib.metadata.distribution(_LMBENCHMARK_PACKAGE) + except importlib.metadata.PackageNotFoundError: + pytest.fail( + "LMBenchmark package is not installed after fixture setup. " + f"Expected {_LMBENCHMARK_SCRIPT_IN_PACKAGE} at commit {LMBENCHMARK_SHA}." + ) + + direct_url = dist.read_text("direct_url.json") + if direct_url is None: + pytest.fail("LMBenchmark package is missing direct_url.json provenance.") + + try: + direct_url_data = json.loads(direct_url) + except json.JSONDecodeError as exc: + pytest.fail(f"LMBenchmark direct_url.json is not valid JSON: {exc}") + + vcs_info = direct_url_data.get("vcs_info", {}) + recorded_revisions = (vcs_info.get("commit_id"), vcs_info.get("requested_revision")) + if LMBENCHMARK_SHA not in recorded_revisions: + pytest.fail( + f"LMBenchmark package revision {recorded_revisions} does not match " + f"expected {LMBENCHMARK_SHA}." + ) + + script = os.fspath(dist.locate_file(_LMBENCHMARK_SCRIPT_IN_PACKAGE)) + if not os.path.exists(script): + pytest.fail( + f"LMBenchmark package installed, but {_LMBENCHMARK_SCRIPT_IN_PACKAGE} was not found." + ) + return script + + +@pytest.fixture(scope="module") +def ensure_lmbenchmark(llm_venv: PythonVenvRunnerImpl) -> None: + """Install and resolve the pinned LMBenchmark script used by this module.""" + global LMBENCHMARK_SCRIPT + + llm_venv.run_cmd(["-m", "pip", "install", _LMBENCHMARK_INSTALL_SPEC], timeout=600) + LMBENCHMARK_SCRIPT = _resolve_lmbenchmark_script() + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestServePrefixAwareScheduling: + """E2E: trtllm-serve with shared prefixes. + + Tests prefix-aware scheduling under shared-prefix workloads that + originally triggered the total_num_tokens > max_num_tokens over-admission + bug. Includes both self-contained OpenAI-client tests and LMBenchmark- + based regression tests. + + Config variants (SCHED_CONFIGS) exercise: + * capacity_scheduler_policy: GUARANTEED_NO_EVICT, MAX_UTILIZATION + * enable_chunked_prefill: true / false + * use_python_scheduler: false / true + * max_attention_window (SWA): unset / 2048 tokens + * disable_overlap_scheduler: false / true + """ + + def test_multi_round_qa_shared_prefix_smoke(self, tmp_path, ensure_lmbenchmark): + """Pre-merge smoke: two-stage LMBenchmark run to catch the original bug. + + Uses the baseline scheduler config (GUARANTEED_NO_EVICT + chunked + prefill). Stage 1 seeds the radix tree; Stage 2 applies load at + QPS=32 — the level that originally triggered the over-admission crash. + The two stages share a single server so that accumulated radix-tree + state from Stage 1 is present when Stage 2 runs, matching the + conditions under which the original bug manifested. + + Estimated runtime: ~2.5 min (server start + 20s warmup + 45s main). + """ + baseline_cfg = { + "kv_cache_config": {"max_tokens": 200000}, + "enable_chunked_prefill": True, + "print_iter_log": True, + } + config_path = _write_config(tmp_path, baseline_cfg) + port = get_free_port_in_ci() + cmd = _make_server_cmd(port, config_path) + server_log = str(tmp_path / "server.log") + env = {**os.environ, "PYTHONUNBUFFERED": "1"} + + with open(server_log, "w") as log_f: + with popen(cmd, stderr=log_f, stdout=log_f, env=env) as proc: + _wait_for_server_ready(proc, port, server_log=server_log) + + # Stage 1: seed radix tree (mimics the low-QPS warmup that + # accumulated state before the crash in the original report). + print_info("Smoke stage 1: seeding radix tree...") + _run_and_assert_stage( + "Smoke warmup", + port, + str(tmp_path / "smoke_warmup.csv"), + server_log, + qps=4, + num_users=3, + num_rounds=3, + system_prompt=1000, + chat_history=20000, + answer_len=10, + duration=20, + ) + + # Stage 2: main load at QPS=32 (the originally failing level). + print_info("Smoke stage 2: main load at QPS=32...") + _run_and_assert_stage( + "Smoke main stage", + port, + str(tmp_path / "smoke_main.csv"), + server_log, + qps=32, + num_users=8, + num_rounds=5, + system_prompt=1000, + chat_history=20000, + answer_len=10, + duration=45, + ) + + assert proc.poll() is None, ( + f"Server exited unexpectedly. See {server_log}\n{_tail_log(server_log)}" + ) + + _assert_no_server_errors(server_log) + + @pytest.mark.parametrize( + "sched_cfg", + [ + c + for c in SCHED_CONFIGS + if c.id + in ( + "guaranteed-chunked", + "guaranteed-no-chunked", + "max-util-chunked", + "python-scheduler", + "swa-chunked", + "no-overlap-chunked", + "offload-chunked", + "offload-no-chunked", + ) + ], + ) + def test_multi_round_qa_shared_prefix(self, tmp_path, sched_cfg, ensure_lmbenchmark): + """Full QPS-sweep regression: shared prefix at escalating load. + + Launches trtllm-serve with block reuse enabled, then runs the full + QPS escalation sequence (8 -> 32) against a single + server process. The original over-admission bug manifests after + accumulated radix-tree state from earlier sweeps, so restarting + per QPS would mask the failure. + + Each stage is preceded by a warmup that seeds the radix tree. + After every stage the test asserts: + * the benchmark process exited with rc=0 + * the output CSV exists and contains completed request rows + * the CSV contains no NaN ttft values + * the server log is free of error patterns + + NOTE: This test is intended for post-merge / nightly CI due to its + ~10 min runtime per config. + """ + config_path = _write_config(tmp_path, sched_cfg) + port = get_free_port_in_ci() + cmd = _make_server_cmd(port, config_path) + server_log = str(tmp_path / "server.log") + # Parameters match the long_input_short_output workload that + # originally surfaced the bug: 1000-token shared system prompt, + # 20000-token per-user chat history, 100-token answers. + qps_values = [8, 32] + env = {**os.environ, "PYTHONUNBUFFERED": "1"} + + with open(server_log, "w") as log_f: + with popen(cmd, stderr=log_f, stdout=log_f, env=env) as proc: + _wait_for_server_ready(proc, port, server_log=server_log) + + # Warmup: seed the radix tree with the shared system prompt. + print_info("Warmup to seed the radix tree with the shared system prompt.") + _run_and_assert_stage( + "Warmup to seed the radix tree with the shared system prompt.", + port, + str(tmp_path / "warmup_1u_qps2.csv"), + server_log, + qps=2, + num_users=1, + num_rounds=2, + system_prompt=1000, + chat_history=20000, + answer_len=100, + duration=10, + ) + + for qps in qps_values: + print_info(f"Benchmark: 15 users, QPS={qps}...") + metrics = _run_and_assert_stage( + f"Benchmark at QPS={qps}", + port, + str(tmp_path / f"benchmark_15u_qps{qps}.csv"), + server_log, + qps=qps, + num_users=15, + num_rounds=3, + system_prompt=1000, + chat_history=20000, + answer_len=100, + duration=30, + ) + print_info( + f"QPS={qps}: {metrics['num_requests']} requests, " + f"ttft_p50={metrics.get('ttft_p50', float('nan')):.3f}s, " + f"ttft_p99={metrics.get('ttft_p99', float('nan')):.3f}s" + ) + + _assert_no_server_errors(server_log) + assert proc.poll() is None, ( + f"Server exited unexpectedly during QPS={qps}. " + f"See {server_log}\n{_tail_log(server_log)}" + ) + + # Final server-log sanity check across the full sweep. + _assert_no_server_errors(server_log) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index dac071fc81a2..207f1e29006f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -167,6 +167,8 @@ l0_b200: - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_long_context - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_rapid_fire - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_field_completeness + # ------------- Prefix-aware scheduling E2E tests --------------- + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke # ------------- Visual Gen tests --------------- - unittest/_torch/visual_gen/test_visual_gen_args.py - unittest/_torch/visual_gen/test_visual_gen_params.py diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 41ee7b38f71d..23059e84409b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -169,6 +169,15 @@ l0_h100: - test_e2e.py::test_openai_responses - test_e2e.py::test_openai_chat_guided_decoding[meta-llama/Llama-3.1-8B-Instruct] - test_e2e.py::test_trtllm_benchmark_serving[llama-3.1-model/Meta-Llama-3.1-8B] + # ------------- Prefix-aware scheduling E2E tests --------------- + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[guaranteed-chunked] + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[max-util-chunked] + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[python-scheduler] + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[swa-chunked] + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[no-overlap-chunked] + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[offload-chunked] + - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[offload-no-chunked] - condition: ranges: system_gpu_count: diff --git a/tests/unittest/_torch/executor/test_kvcache_aware_router.py b/tests/unittest/_torch/executor/test_kvcache_aware_router.py index 8d7dc1d2e404..ac2ed0baa67e 100644 --- a/tests/unittest/_torch/executor/test_kvcache_aware_router.py +++ b/tests/unittest/_torch/executor/test_kvcache_aware_router.py @@ -429,7 +429,12 @@ def test_empty_tokens_returns_zero(self): assert result == 0 def test_block_to_token_conversion(self): - """Verify num_blocks * tokens_per_block conversion.""" + """Verify num_blocks * tokens_per_block conversion. + + Also pins the expectation that probe_prefix_match_length calls + ``analyze_prefix_reuse`` exactly once. Guards against an accidental + regression to a separate count-path in the router. + """ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager mgr = Mock(spec=KVCacheManager) @@ -443,3 +448,4 @@ def test_block_to_token_conversion(self): result = KVCacheManager.probe_prefix_match_length(mgr, input_tokens=list(range(200))) assert result == 192 # 3 blocks * 64 tokens/block + mgr.impl.analyze_prefix_reuse.assert_called_once() diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index d6ce0505baa9..4a5163e55982 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -1708,6 +1708,82 @@ def test_draft_tokens_greater_than_chunk_size(self): assert req1.num_draft_tokens == 13 assert req2.num_draft_tokens == 5 + def test_mixed_batch_zero_draft_does_not_consume_speculative_budget(self): + """Regression: in a mixed batch, zero-draft requests must not burn the + ``remaining_space`` compute budget that a later draft-bearing request + needs. + + ``req.has_draft_tokens`` is a nanobind-bound method in + ``cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp``, so + reading it without ``()`` returned a truthy bound-method object. That + made every last-context-chunk request enter ``_fit_draft_tokens``'s + draft branch, including zero-draft ones, which then accumulated + ``remaining_space`` into ``num_ctx_tokens`` and left less room for + later speculative requests to keep their drafts. + + Setup (chunk_unit_size=16, max_num_tokens=28, three last-chunk reqs): + - req0 prompt=8 draft=0 → rs=8 on buggy path + - req1 prompt=8 draft=0 → rs=8 on buggy path + - req2 prompt=3 draft=13 → rs=13, needs budget + + Initial num_ctx_tokens sum = 8 + 8 + 3 = 19. + Fixed: req0/req1 skip draft branch; req2 keeps + min(13, 16-3=13, 28-19=9) = 9 drafts. + Buggy: req0 burns 8 (→27), req1 burns 1 (→28), req2 gets 0 drafts. + """ + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=16) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, + max_num_tokens=28, + ctx_chunk_config=config, + ) + scheduler.max_context_length = 64 + + requests = [ + make_context_request(0, prompt_len=8, draft_tokens_len=0), + make_context_request(1, prompt_len=8, draft_tokens_len=0), + make_context_request(2, prompt_len=3, draft_tokens_len=13), + ] + ctx, _ = scheduler.schedule(requests, set()) + assert len(ctx) == 3 + r0 = next(r for r in ctx if r.request_id == 0) + r1 = next(r for r in ctx if r.request_id == 1) + r2 = next(r for r in ctx if r.request_id == 2) + assert r0.num_draft_tokens == 0 + assert r1.num_draft_tokens == 0 + # Before fix this would be 0 (all drafts discarded because the earlier + # zero-draft siblings had already burned the compute budget). + assert r2.num_draft_tokens == 9, ( + f"zero-draft siblings burned speculative budget: " + f"got num_draft_tokens={r2.num_draft_tokens}, expected 9" + ) + + def test_no_draft_tokens_bypasses_fit_draft(self): + """Regression: _fit_draft_tokens must not discard when no drafts exist. + + Previously, ``req.has_draft_tokens`` (bound method, always truthy) + was accidentally evaluated without ``()``, so draft-discarding ran + for requests with zero drafts. When the capacity budget was + exhausted, ``capacity - num_ctx_tokens`` went negative and + ``draft_discard`` wrapped into a large positive value, triggering + the C++ assertion "Can't discard more draft tokens (N) than exists + (0)" in ``llmRequest.h``. Guard against that by ensuring + scheduling a batch of large no-draft requests does not raise. + """ + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=16) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, + max_num_tokens=32, + ctx_chunk_config=config, + ) + # Requests with zero draft tokens, large enough to exhaust budget. + requests = [make_context_request(i, prompt_len=24, draft_tokens_len=0) for i in range(4)] + ctx, _ = scheduler.schedule(requests, set()) + assert ctx, "expected at least one context request to be scheduled" + # All scheduled requests must still report 0 draft tokens — no discard happened. + for r in ctx: + assert r.num_draft_tokens == 0 + # ############################################################################ # From ccb8318cdcf3f93016be1ed2a78efe1451a32e06 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Thu, 30 Apr 2026 14:23:57 -0700 Subject: [PATCH 2/2] Address comments from coderabbit and Aurelien. Signed-off-by: Simeng Liu --- .../_torch/pyexecutor/scheduler/scheduler.py | 45 ++--- .../kv_cache/test_prefix_aware_scheduling.py | 182 ++++++++++++------ .../test_lists/qa/llm_function_core.txt | 1 + .../_torch/executor/test_py_scheduler.py | 28 +++ 4 files changed, 174 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index d3256251c7b3..f39064013321 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1,4 +1,5 @@ import dataclasses +import inspect from abc import ABC, abstractmethod from collections import namedtuple from dataclasses import dataclass @@ -30,15 +31,22 @@ def _call_with_optional_summary( The nanobind binding for ``get_remaining_blocks_to_completion`` and ``get_needed_blocks_one_step`` accepts ``cached_summary: PrefixReuseSummary | None = None`` so the C++ side can skip a redundant radix-tree walk. Test mocks may not - accept this kwarg yet; fall back to the positional form on TypeError to - keep the path working for both. + accept this kwarg yet; inspect the callable before passing it so real + ``TypeError`` exceptions from inside ``fn`` are not hidden. """ if cached_summary is None: return fn(*args) + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): return fn(*args, cached_summary=cached_summary) - except TypeError: + if not any( + param.name == "cached_summary" or param.kind == inspect.Parameter.VAR_KEYWORD + for param in signature.parameters.values() + ): return fn(*args) + return fn(*args, cached_summary=cached_summary) SchedulerOutput = namedtuple( @@ -821,15 +829,18 @@ def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit remaining_context_len = self.max_context_length - req.context_chunk_size remaining_space = min(remaining_space, remaining_context_len) + remaining_space = max(0, remaining_space) + kept_drafts = min(req.num_draft_tokens, remaining_space) + if capacity is not None: # capacity - num_ctx_tokens can go negative if the request - # batch has already overshot the compute budget, which would - # wrap draft_discard into an invalid large positive value - # below. Clamp at 0. - remaining_space = max(0, min(remaining_space, capacity - num_ctx_tokens)) - num_ctx_tokens += remaining_space + # batch has already overshot the compute budget. Clamp at 0 + # and only charge the draft tokens that actually remain + # attached to this request. + kept_drafts = max(0, min(kept_drafts, capacity - num_ctx_tokens)) + num_ctx_tokens += kept_drafts - draft_discard = req.num_draft_tokens - remaining_space + draft_discard = req.num_draft_tokens - kept_drafts if draft_discard > 0: logger.debug(f"Discarding {draft_discard} draft tokens") req.discard_draft_tokens(draft_discard) @@ -1426,21 +1437,7 @@ def _prefill_contributed_blocks(self, active_requests: RequestList) -> tuple[set return newly_contributed_context_blocks, newly_contributed_cross_context_blocks def _make_prefix_summary_caches(self) -> tuple[PrefixSummaryCache, PrefixSummaryCache]: - """Build empty caches for `PrefixReuseSummary` keyed by `py_request_id`. - - The caches are populated **lazily** by `_beneficial_to_skip` when it - walks the radix tree for a pending request. Downstream consumers - (`enough_available_blocks`, `decrement_reserved_blocks`) reuse the - cached summary to avoid re-walking. - - Precomputing for every active request is avoided on purpose: the - ``req.get_unique_tokens(0)`` marshalling copies a potentially very - large token vector across the Python/C++ boundary and, at high - concurrency, dominated scheduler CPU time (verified via py-spy: - >80% of `_schedule`). Lazy population matches the C++ pattern in - ``capacityScheduler.cpp`` which only walks for pending-loop - candidates. - """ + """Build empty caches for `PrefixReuseSummary` keyed by `py_request_id`.""" return {}, {} def _beneficial_to_skip( diff --git a/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py b/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py index cb125ae78a90..a087272839fd 100644 --- a/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py +++ b/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py @@ -24,16 +24,17 @@ """ import csv -import importlib.metadata import json import math import os import queue import re import subprocess -import sys import threading import time +from collections.abc import Mapping +from pathlib import Path +from typing import TypeAlias import pytest import requests as req_lib @@ -57,6 +58,7 @@ _LMBENCHMARK_INSTALL_SPEC = ( f"{_LMBENCHMARK_PACKAGE} @ git+https://github.com/LMCache/LMBenchmark.git@{LMBENCHMARK_SHA}" ) +CsvMetrics: TypeAlias = dict[str, int | float] # --------------------------------------------------------------------------- # Scheduler / KV-cache config combinations @@ -227,7 +229,13 @@ def _assert_no_server_errors(server_log: str) -> None: # --------------------------------------------------------------------------- -def _wait_for_server_ready(proc, port, timeout=300, interval=2, server_log=None): +def _wait_for_server_ready( + proc: subprocess.Popen, + port: int, + timeout: int = 300, + interval: int = 2, + server_log: str | None = None, +) -> None: """Wait for trtllm-serve /health to return 200. If *server_log* is provided, the log is scanned for error patterns on each @@ -271,7 +279,7 @@ def _make_server_cmd(port: int, config_path: str) -> list[str]: ] -def _write_config(tmp_path, cfg: dict, name: str = "config.yml") -> str: +def _write_config(tmp_path: Path, cfg: dict[str, object], name: str = "config.yml") -> str: """Serialise *cfg* to YAML in *tmp_path* and return the file path.""" path = str(tmp_path / name) with open(path, "w") as f: @@ -285,17 +293,18 @@ def _write_config(tmp_path, cfg: dict, name: str = "config.yml") -> str: def _run_lmbenchmark( - port, - output_csv, - qps=0.5, - num_users=4, - num_rounds=3, - system_prompt=500, - chat_history=2000, - answer_len=20, - duration=30, - server_log=None, -): + llm_venv: PythonVenvRunnerImpl, + port: int, + output_csv: str, + qps: float = 0.5, + num_users: int = 4, + num_rounds: int = 3, + system_prompt: int = 500, + chat_history: int = 2000, + answer_len: int = 20, + duration: int = 30, + server_log: str | None = None, +) -> int: """Run the LMBenchmark multi-round-qa script and return the exit code. stdout and stderr are drained in background threads so that the main @@ -305,9 +314,9 @@ def _run_lmbenchmark( goes quiet (which is exactly what happens when the server stalls), defeating the point of the watchdog. """ + assert LMBENCHMARK_SCRIPT, "ensure_lmbenchmark fixture must resolve LMBenchmark first" script_dir = os.path.dirname(LMBENCHMARK_SCRIPT) cmd = [ - sys.executable, LMBENCHMARK_SCRIPT, "--num-users", str(num_users), @@ -337,21 +346,30 @@ def _run_lmbenchmark( print_info(f"Running LMBenchmark: {' '.join(cmd)}") t_start = time.time() - proc = subprocess.Popen( - cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, cwd=script_dir - ) - stdout_q: queue.Queue = queue.Queue() + def _popen_lmbenchmark(call_args: list[str], env: Mapping[str, str]) -> subprocess.Popen: + return subprocess.Popen( + call_args, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + cwd=script_dir, + env=env, + ) + + proc = llm_venv.run_cmd(cmd, caller=_popen_lmbenchmark) + + stdout_q: queue.Queue[str | None] = queue.Queue() stderr_lines: list[str] = [] - def _drain_stdout(): + def _drain_stdout() -> None: try: for line in iter(proc.stdout.readline, ""): stdout_q.put(line) finally: stdout_q.put(None) # EOF sentinel - def _drain_stderr(): + def _drain_stderr() -> None: for line in iter(proc.stderr.readline, ""): stderr_lines.append(line) @@ -385,6 +403,13 @@ def _kill(reason: str) -> int: stderr_thread.join(timeout=5) return -1 + def _cleanup_if_running() -> None: + if proc.poll() is None: + proc.kill() + proc.wait() + stdout_thread.join(timeout=5) + stderr_thread.join(timeout=5) + def _drain_queued_stdout() -> int | None: """Pop all queued stdout lines; return -1 if NaN seen, else None.""" while True: @@ -398,6 +423,7 @@ def _drain_queued_stdout() -> int | None: if line and _NAN_RE.search(line): return _kill(f"NaN detected in benchmark output: {line!r}") + completed = False try: while proc.poll() is None: rc = _drain_queued_stdout() @@ -427,7 +453,7 @@ def _drain_queued_stdout() -> int | None: resp = req_lib.get(f"http://localhost:{port}/health", timeout=5) if resp.status_code == 200: last_health_ok = now - except Exception: + except RequestException: pass # 60s (was 30s) — extreme QPS stages need more slack before # declaring the server stalled; genuine hangs still get caught. @@ -435,13 +461,10 @@ def _drain_queued_stdout() -> int | None: return _kill("Server stopped responding to /health (stalled)") time.sleep(0.2) - - except Exception: - proc.kill() - proc.wait() - stdout_thread.join(timeout=5) - stderr_thread.join(timeout=5) - raise + completed = True + finally: + if not completed: + _cleanup_if_running() # Benchmark has exited. Wait for drain threads to finish so we see # every remaining line, then surface any leftover output. @@ -479,7 +502,7 @@ def _drain_queued_stdout() -> int | None: return proc.returncode -def _parse_csv_metrics(csv_path): +def _parse_csv_metrics(csv_path: str) -> CsvMetrics | None: """Parse LMBenchmark output CSV and return TTFT metrics. Returns a dict with keys: @@ -496,7 +519,7 @@ def _parse_csv_metrics(csv_path): return None nan_count = 0 - ttfts = [] + ttfts: list[float] = [] for r in rows: raw = r.get("ttft") if not raw: @@ -511,22 +534,28 @@ def _parse_csv_metrics(csv_path): else: ttfts.append(v) - if not ttfts: - return None - - ttfts.sort() - return { + metrics: CsvMetrics = { "num_requests": len(rows), "nan_count": nan_count, "ttft_count": len(ttfts), - "ttft_avg": sum(ttfts) / len(ttfts), - "ttft_p50": ttfts[len(ttfts) // 2], - "ttft_p99": ttfts[int(len(ttfts) * 0.99)], } + if not ttfts: + return metrics + + ttfts.sort() + metrics.update( + { + "ttft_avg": sum(ttfts) / len(ttfts), + "ttft_p50": ttfts[len(ttfts) // 2], + "ttft_p99": ttfts[int(len(ttfts) * 0.99)], + } + ) + return metrics def _run_and_assert_stage( label: str, + llm_venv: PythonVenvRunnerImpl, port: int, output_csv: str, server_log: str, @@ -538,9 +567,10 @@ def _run_and_assert_stage( chat_history: int, answer_len: int, duration: int, -) -> dict: +) -> CsvMetrics: """Run one LMBenchmark stage, assert success, and return the metrics dict.""" rc = _run_lmbenchmark( + llm_venv=llm_venv, port=port, output_csv=output_csv, qps=qps, @@ -570,20 +600,10 @@ def _run_and_assert_stage( # --------------------------------------------------------------------------- -def _resolve_lmbenchmark_script() -> str: - """Return the benchmark script from the pinned installed LMBenchmark package.""" - try: - dist = importlib.metadata.distribution(_LMBENCHMARK_PACKAGE) - except importlib.metadata.PackageNotFoundError: - pytest.fail( - "LMBenchmark package is not installed after fixture setup. " - f"Expected {_LMBENCHMARK_SCRIPT_IN_PACKAGE} at commit {LMBENCHMARK_SHA}." - ) - - direct_url = dist.read_text("direct_url.json") +def _validate_lmbenchmark_direct_url(direct_url: str | None) -> None: + """Verify the installed LMBenchmark package records the pinned Git revision.""" if direct_url is None: pytest.fail("LMBenchmark package is missing direct_url.json provenance.") - try: direct_url_data = json.loads(direct_url) except json.JSONDecodeError as exc: @@ -597,7 +617,41 @@ def _resolve_lmbenchmark_script() -> str: f"expected {LMBENCHMARK_SHA}." ) - script = os.fspath(dist.locate_file(_LMBENCHMARK_SCRIPT_IN_PACKAGE)) + +def _resolve_lmbenchmark_install(llm_venv: PythonVenvRunnerImpl) -> str: + """Return the benchmark script from the fixture Python environment.""" + query = f""" +import importlib.metadata +import json +import os + +package = {_LMBENCHMARK_PACKAGE!r} +script_in_package = {_LMBENCHMARK_SCRIPT_IN_PACKAGE!r} + +try: + dist = importlib.metadata.distribution(package) +except importlib.metadata.PackageNotFoundError: + raise SystemExit( + "LMBenchmark package is not installed after fixture setup. " + f"Expected {{script_in_package}}." + ) + +print( + json.dumps( + {{ + "direct_url": dist.read_text("direct_url.json"), + "script": os.fspath(dist.locate_file(script_in_package)), + }} + ) +) +""" + try: + payload = json.loads(llm_venv.run_output(query)) + except (RuntimeError, json.JSONDecodeError) as exc: + pytest.fail(f"Failed to resolve installed LMBenchmark package: {exc}") + + _validate_lmbenchmark_direct_url(payload.get("direct_url")) + script = payload["script"] if not os.path.exists(script): pytest.fail( f"LMBenchmark package installed, but {_LMBENCHMARK_SCRIPT_IN_PACKAGE} was not found." @@ -606,12 +660,13 @@ def _resolve_lmbenchmark_script() -> str: @pytest.fixture(scope="module") -def ensure_lmbenchmark(llm_venv: PythonVenvRunnerImpl) -> None: +def ensure_lmbenchmark(llm_venv: PythonVenvRunnerImpl) -> PythonVenvRunnerImpl: """Install and resolve the pinned LMBenchmark script used by this module.""" global LMBENCHMARK_SCRIPT llm_venv.run_cmd(["-m", "pip", "install", _LMBENCHMARK_INSTALL_SPEC], timeout=600) - LMBENCHMARK_SCRIPT = _resolve_lmbenchmark_script() + LMBENCHMARK_SCRIPT = _resolve_lmbenchmark_install(llm_venv) + return llm_venv # --------------------------------------------------------------------------- @@ -635,7 +690,9 @@ class TestServePrefixAwareScheduling: * disable_overlap_scheduler: false / true """ - def test_multi_round_qa_shared_prefix_smoke(self, tmp_path, ensure_lmbenchmark): + def test_multi_round_qa_shared_prefix_smoke( + self, tmp_path: Path, ensure_lmbenchmark: PythonVenvRunnerImpl + ) -> None: """Pre-merge smoke: two-stage LMBenchmark run to catch the original bug. Uses the baseline scheduler config (GUARANTEED_NO_EVICT + chunked @@ -667,6 +724,7 @@ def test_multi_round_qa_shared_prefix_smoke(self, tmp_path, ensure_lmbenchmark): print_info("Smoke stage 1: seeding radix tree...") _run_and_assert_stage( "Smoke warmup", + ensure_lmbenchmark, port, str(tmp_path / "smoke_warmup.csv"), server_log, @@ -683,6 +741,7 @@ def test_multi_round_qa_shared_prefix_smoke(self, tmp_path, ensure_lmbenchmark): print_info("Smoke stage 2: main load at QPS=32...") _run_and_assert_stage( "Smoke main stage", + ensure_lmbenchmark, port, str(tmp_path / "smoke_main.csv"), server_log, @@ -719,7 +778,12 @@ def test_multi_round_qa_shared_prefix_smoke(self, tmp_path, ensure_lmbenchmark): ) ], ) - def test_multi_round_qa_shared_prefix(self, tmp_path, sched_cfg, ensure_lmbenchmark): + def test_multi_round_qa_shared_prefix( + self, + tmp_path: Path, + sched_cfg: dict[str, object], + ensure_lmbenchmark: PythonVenvRunnerImpl, + ) -> None: """Full QPS-sweep regression: shared prefix at escalating load. Launches trtllm-serve with block reuse enabled, then runs the full @@ -756,6 +820,7 @@ def test_multi_round_qa_shared_prefix(self, tmp_path, sched_cfg, ensure_lmbenchm print_info("Warmup to seed the radix tree with the shared system prompt.") _run_and_assert_stage( "Warmup to seed the radix tree with the shared system prompt.", + ensure_lmbenchmark, port, str(tmp_path / "warmup_1u_qps2.csv"), server_log, @@ -772,6 +837,7 @@ def test_multi_round_qa_shared_prefix(self, tmp_path, sched_cfg, ensure_lmbenchm print_info(f"Benchmark: 15 users, QPS={qps}...") metrics = _run_and_assert_stage( f"Benchmark at QPS={qps}", + ensure_lmbenchmark, port, str(tmp_path / f"benchmark_15u_qps{qps}.csv"), server_log, diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index e5548a247d76..daf539eb28ce 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -875,6 +875,7 @@ disaggregated/test_workers.py::test_workers_kv_cache_aware_router_deepseek_v3_li disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_workers.py::test_workers_kv_cache_events[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-Chat-v1.0] +kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke llmapi/test_llm_examples.py::test_llmapi_server_example test_e2e.py::test_eagle3_output_repetition_4gpus[modelopt-hf-model-hub/Llama-3.3-70B-Instruct-fp8-EAGLE3-LLaMA3.3-Instruct-70B] test_e2e.py::test_eagle3_output_repetition_4gpus[Qwen3/Qwen3-30B-A3B-Qwen3/Qwen3-30B-eagle3] diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index 4a5163e55982..89cec2f84dad 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -1758,6 +1758,34 @@ def test_mixed_batch_zero_draft_does_not_consume_speculative_budget(self): f"got num_draft_tokens={r2.num_draft_tokens}, expected 9" ) + def test_short_draft_request_charges_only_kept_drafts(self) -> None: + """Regression: a request with fewer drafts than chunk remainder should + not charge the whole remainder against the shared compute budget. + """ + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=16) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, + max_num_tokens=30, + ctx_chunk_config=config, + ) + scheduler.max_context_length = 64 + + requests = [ + make_context_request(0, prompt_len=3, draft_tokens_len=1), + make_context_request(1, prompt_len=3, draft_tokens_len=13), + make_context_request(2, prompt_len=3, draft_tokens_len=13), + ] + ctx, _ = scheduler.schedule(requests, set()) + assert len(ctx) == 3 + + r0 = next(r for r in ctx if r.request_id == 0) + r1 = next(r for r in ctx if r.request_id == 1) + r2 = next(r for r in ctx if r.request_id == 2) + + assert r0.num_draft_tokens == 1 + assert r1.num_draft_tokens == 13 + assert r2.num_draft_tokens == 7 + def test_no_draft_tokens_bypasses_fit_draft(self): """Regression: _fit_draft_tokens must not discard when no drafts exist.