From eae3837f4b0bc2ec574057acfd75c752681f52f2 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Tue, 3 Mar 2026 02:04:53 -0800 Subject: [PATCH 1/7] init tests Signed-off-by: Lanyu Liao --- .../_torch/executor/test_py_scheduler.py | 2241 +++++++++++++++++ 1 file changed, 2241 insertions(+) create mode 100644 tests/unittest/_torch/executor/test_py_scheduler.py diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py new file mode 100644 index 000000000000..633e54a6d26b --- /dev/null +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -0,0 +1,2241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. +""" +Unit tests for Python scheduler implementations (PyMicroBatchScheduler, +PyCapacityScheduler, SimpleUnifiedScheduler). + +These tests validate the pure-Python scheduler logic using mock objects, +without requiring GPU or real KV cache managers. They are aligned with +the C++ scheduler unit tests in: + - cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp + - cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp +""" + +from dataclasses import dataclass, field +from enum import IntEnum +from typing import List, Optional + +import pytest + +# --------------------------------------------------------------------------- +# Mock LlmRequestState — mirrors the C++ enum values exactly. +# We define our own so these tests don't depend on the C++ bindings at import +# time, making them runnable in environments where bindings aren't built. +# --------------------------------------------------------------------------- + + +# Try to import the real scheduler classes. If the bindings are not available +# (e.g. no CUDA environment), the entire module is skipped. +try: + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ( + ChunkingPolicy, + ContextChunkingConfig, + PyCapacityScheduler, + PyMicroBatchScheduler, + SimpleUnifiedScheduler, + ) + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy + + SCHEDULER_AVAILABLE = True +except ImportError: + SCHEDULER_AVAILABLE = False + + +# MockLlmRequestState uses the REAL enum values from C++ bindings. +# These values are NOT sequential (e.g., CONTEXT_INIT=10, not 2). +# We derive them from the real LlmRequestState when available. +if SCHEDULER_AVAILABLE: + + class MockLlmRequestState(IntEnum): + UNKNOWN = LlmRequestState.UNKNOWN.value + ENCODER_INIT = LlmRequestState.ENCODER_INIT.value + CONTEXT_INIT = LlmRequestState.CONTEXT_INIT.value + GENERATION_IN_PROGRESS = LlmRequestState.GENERATION_IN_PROGRESS.value + GENERATION_TO_COMPLETE = LlmRequestState.GENERATION_TO_COMPLETE.value + GENERATION_COMPLETE = LlmRequestState.GENERATION_COMPLETE.value + DISAGG_GENERATION_INIT = LlmRequestState.DISAGG_GENERATION_INIT.value + +else: + + class MockLlmRequestState(IntEnum): + UNKNOWN = 0 + ENCODER_INIT = 1 + CONTEXT_INIT = 10 + GENERATION_IN_PROGRESS = 13 + GENERATION_TO_COMPLETE = 14 + GENERATION_COMPLETE = 20 + DISAGG_GENERATION_INIT = 8 + +pytestmark = pytest.mark.skipif( + not SCHEDULER_AVAILABLE, + reason="tensorrt_llm scheduler not importable (missing bindings or CUDA)", +) + +# ============================================================================ +# Mock infrastructure +# ============================================================================ + + +@dataclass +class MockLlmRequest: + """ + Minimal mock of LlmRequest that exposes exactly the attributes/methods + the Python scheduler accesses. Attribute names match the real binding API. + """ + + request_id: int + _state: int = MockLlmRequestState.CONTEXT_INIT + prompt_len: int = 10 + _context_current_position: int = 0 + _context_chunk_size: int = 0 + beam_width: int = 1 + _draft_tokens_len: int = 0 + encoder_output_len: int = 0 + lora_task_id: Optional[int] = None + py_request_id: int = 0 # Used by MaxUtilization's scheduling_remove_sequence + _input_tokens: Optional[List[int]] = None + + def __post_init__(self): + self.py_request_id = self.request_id + + # --- State properties --- + @property + def state(self): + return LlmRequestState(self._state) if SCHEDULER_AVAILABLE else self._state + + @state.setter + def state(self, value): + self._state = int(value) + + @property + def state_value(self) -> int: + return self._state + + @property + def is_context_init_state(self) -> bool: + return self._state == MockLlmRequestState.CONTEXT_INIT + + @property + def is_generation_in_progress_state(self) -> bool: + return self._state == MockLlmRequestState.GENERATION_IN_PROGRESS + + @property + def is_encoder_init_state(self) -> bool: + return self._state == MockLlmRequestState.ENCODER_INIT + + @property + def is_disagg_generation_init_state(self) -> bool: + return self._state == MockLlmRequestState.DISAGG_GENERATION_INIT + + @property + def is_generation_to_complete_state(self) -> bool: + return self._state == MockLlmRequestState.GENERATION_TO_COMPLETE + + # --- Context chunking properties --- + @property + def context_remaining_length(self) -> int: + return self.prompt_len - self._context_current_position + + @property + def context_chunk_size(self) -> int: + return self._context_chunk_size + + @context_chunk_size.setter + def context_chunk_size(self, value: int): + # Clamp to context_remaining_length, matching C++ behavior + self._context_chunk_size = min(value, self.context_remaining_length) + + @property + def is_first_context_chunk(self) -> bool: + return self._context_current_position == 0 + + @property + def is_last_context_chunk(self) -> bool: + return self._context_chunk_size >= self.context_remaining_length + + # --- Token / draft properties --- + def get_num_tokens(self, beam: int) -> int: + """For context requests, returns prompt_len (matching C++ getNumTokens).""" + return self.prompt_len + + @property + def num_draft_tokens(self) -> int: + return self._draft_tokens_len + + @property + def has_draft_tokens(self) -> bool: + return self._draft_tokens_len > 0 + + @property + def draft_tokens(self) -> list: + return [0] * self._draft_tokens_len + + def get_beam_width_by_iter(self, for_next_iteration: bool = False) -> int: + return self.beam_width + + def discard_draft_tokens(self, n: int): + self._draft_tokens_len = max(0, self._draft_tokens_len - n) + + # --- Methods used by capacity scheduler --- + def get_unique_tokens(self, beam: int): + return self._input_tokens + + def get_encoder_unique_tokens(self): + return None + + def move_to_next_context_chunk(self): + """Simulate moving to the next context chunk.""" + self._context_current_position += self._context_chunk_size + self._context_chunk_size = 0 + + +# Helper to create requests in specific states +def make_context_request( + request_id: int, + prompt_len: int = 10, + beam_width: int = 1, + draft_tokens_len: int = 0, + context_position: int = 0, +) -> MockLlmRequest: + return MockLlmRequest( + request_id=request_id, + _state=MockLlmRequestState.CONTEXT_INIT, + prompt_len=prompt_len, + _context_current_position=context_position, + beam_width=beam_width, + _draft_tokens_len=draft_tokens_len, + ) + + +def make_generation_request( + request_id: int, + beam_width: int = 1, + draft_tokens_len: int = 0, +) -> MockLlmRequest: + return MockLlmRequest( + request_id=request_id, + _state=MockLlmRequestState.GENERATION_IN_PROGRESS, + beam_width=beam_width, + _draft_tokens_len=draft_tokens_len, + ) + + +def make_encoder_request( + request_id: int, + encoder_output_len: int = 10, +) -> MockLlmRequest: + return MockLlmRequest( + request_id=request_id, + _state=MockLlmRequestState.ENCODER_INIT, + encoder_output_len=encoder_output_len, + ) + + +def make_disagg_gen_init_request(request_id: int) -> MockLlmRequest: + return MockLlmRequest( + request_id=request_id, + _state=MockLlmRequestState.DISAGG_GENERATION_INIT, + ) + + +def make_completed_request(request_id: int) -> MockLlmRequest: + return MockLlmRequest( + request_id=request_id, + _state=MockLlmRequestState.GENERATION_COMPLETE, + ) + + +# ============================================================================ +# Mock KV Cache Manager +# ============================================================================ + + +@dataclass +class MockKVCacheStats: + num_free_blocks_per_window_size: dict = field(default_factory=lambda: {128: 100}) + + +class MockKVCacheManager: + """Mock KV cache manager for capacity scheduler tests.""" + + def __init__( + self, + num_free_blocks: int = 100, + window_sizes: Optional[list] = None, + blocks_per_request: int = 5, + is_variable_window: bool = False, + enable_block_reuse: bool = False, + ): + self._window_sizes = window_sizes or [128] + self._num_free_blocks = num_free_blocks + self._blocks_per_request = blocks_per_request + self.is_variable_window = is_variable_window + self.enable_block_reuse = enable_block_reuse + self.max_attention_window_vec = self._window_sizes + self._scheduling_started = False + + def get_kv_cache_stats(self) -> MockKVCacheStats: + return MockKVCacheStats( + num_free_blocks_per_window_size={ + ws: self._num_free_blocks for ws in self._window_sizes + } + ) + + def get_remaining_blocks_to_completion(self, req, window_size: int) -> int: + return self._blocks_per_request + + def get_needed_blocks_one_step( + self, req, two_step_lookahead: bool, window_size: int + ) -> int: + return self._blocks_per_request + + def scheduling_has_free_blocks(self, total: int, window_size: int) -> bool: + return total <= self._num_free_blocks + + def start_scheduling(self): + self._scheduling_started = True + + def scheduling_remove_sequence(self, req_id: int): + pass + + def find_new_context_block(self, unique_tokens, req): + return None + + def get_max_resource_count(self) -> int: + return self._num_free_blocks + + def get_needed_resource_to_completion(self, req) -> int: + return self._blocks_per_request + + +# ============================================================================ +# Mock PEFT Cache Manager +# ============================================================================ + + +class MockPeftCacheManager: + def __init__(self, max_pages: int = 100, pages_per_request: int = 10): + self.max_device_pages = max_pages + self._pages_per_request = pages_per_request + + def determine_num_pages(self, req) -> int: + return self._pages_per_request + + +# ############################################################################ +# +# Phase 2: PyMicroBatchScheduler Tests +# +# ############################################################################ + + +class TestPyMicroBatchSchedulerBasic: + """ + Tests for PyMicroBatchScheduler — single-step scheduling decisions. + Aligned with C++ MicroBatchSchedulerTest in microBatchSchedulerTest.cpp. + """ + + def test_simple_context_only(self): + """All requests are context requests, batch size allows 2.""" + scheduler = PyMicroBatchScheduler(max_batch_size=2, max_num_tokens=None) + requests = [ + make_context_request(0, prompt_len=10), + make_context_request(1, prompt_len=10), + make_context_request(2, prompt_len=10), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + assert len(gen) == 0 + assert ctx[0].request_id == 0 + assert ctx[1].request_id == 1 + + def test_simple_generation_only(self): + """All requests are generation requests, batch size allows 2.""" + scheduler = PyMicroBatchScheduler(max_batch_size=2, max_num_tokens=None) + requests = [ + make_generation_request(0), + make_generation_request(1), + make_generation_request(2), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 0 + assert len(gen) == 2 + assert gen[0].request_id == 0 + assert gen[1].request_id == 1 + + def test_context_generation_overlap(self): + """ + Mixed batch: context + generation requests. + C++ ref: SimpleWithOverlap + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=None) + requests = [ + make_context_request(0, prompt_len=10), + make_generation_request(1), + make_context_request(2, prompt_len=10), + make_generation_request(3), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + assert len(gen) == 2 + assert {r.request_id for r in ctx} == {0, 2} + assert {r.request_id for r in gen} == {1, 3} + + def test_max_num_tokens_limits_context(self): + """ + max_num_tokens limits how many context tokens can be scheduled. + C++ ref: SimpleNoOverlapMaxNumTokens + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=15) + # Each context request has 10 tokens. Two would be 20 > 15. + requests = [ + make_context_request(0, prompt_len=10), + make_context_request(1, prompt_len=10), + ] + ctx, gen = scheduler.schedule(requests, set()) + # Only 1 fits within token budget + assert len(ctx) == 1 + assert ctx[0].request_id == 0 + + def test_max_num_tokens_allows_gen_after_context(self): + """ + After scheduling a context request, generation requests still fit + if their token count (beam_width) fits in remaining budget. + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=12) + requests = [ + make_context_request(0, prompt_len=10), + make_generation_request(1, beam_width=1), + make_generation_request(2, beam_width=1), + ] + ctx, gen = scheduler.schedule(requests, set()) + # context: 10 tokens, gen1: 1 token, gen2: 1 token => total 12 + assert len(ctx) == 1 + assert len(gen) == 2 + + def test_max_batch_size_limits_total(self): + """Batch size limits total (context + generation).""" + scheduler = PyMicroBatchScheduler(max_batch_size=2, max_num_tokens=None) + requests = [ + make_context_request(0, prompt_len=5), + make_generation_request(1), + make_generation_request(2), + ] + ctx, gen = scheduler.schedule(requests, set()) + # batch_size=2: should schedule context_0 + gen_1 + assert len(ctx) + len(gen) == 2 + + def test_beam_width_1(self): + """ + Generation requests with beam_width=1 each cost 1 token. + C++ ref: SimpleMaxNumTokensBW1 + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=12) + requests = [ + make_context_request(0, prompt_len=10, beam_width=1), + make_generation_request(1, beam_width=1), + make_generation_request(2, beam_width=1), + make_generation_request(3, beam_width=1), + ] + ctx, gen = scheduler.schedule(requests, set()) + # context: 10, gen: 1+1 = 12 total. Can't fit gen_3 (would be 13). + assert len(ctx) == 1 + assert len(gen) == 2 + + def test_beam_width_4(self): + """ + Generation requests with beam_width=4 each cost 4 tokens. + C++ ref: SimpleMaxNumTokensBW4 + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=15) + requests = [ + make_context_request(0, prompt_len=10, beam_width=4), + make_generation_request(1, beam_width=4), + make_generation_request(2, beam_width=4), + ] + ctx, gen = scheduler.schedule(requests, set()) + # context: 10, gen1: 4 = 14. gen2: +4 = 18 > 15. + assert len(ctx) == 1 + assert len(gen) == 1 + + def test_beam_width_mismatch_skipped(self): + """ + Generation requests with different beam widths are skipped. + C++ ensures all gen requests in a batch have same beam_width. + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=None) + requests = [ + make_generation_request(0, beam_width=1), + make_generation_request(1, beam_width=4), + make_generation_request(2, beam_width=1), + ] + ctx, gen = scheduler.schedule(requests, set()) + # gen_0 sets beam_width=1, gen_1 is skipped (beam_width=4), gen_2 fits + assert len(gen) == 2 + assert gen[0].request_id == 0 + assert gen[1].request_id == 2 + + def test_draft_tokens_count_toward_budget(self): + """ + Draft tokens are added to the token count for both context and gen. + C++ ref: DraftTokensMaxNumTokens + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=15) + # Context request: 10 prompt + 3 draft = 13 tokens + requests = [ + make_context_request(0, prompt_len=10, draft_tokens_len=3), + make_generation_request(1, draft_tokens_len=2), + ] + ctx, gen = scheduler.schedule(requests, set()) + # context: 10+3=13, gen: 1+2=3, total=16 > 15 => only context fits + assert len(ctx) == 1 + assert len(gen) == 0 + + def test_gen_draft_tokens(self): + """ + Generation with draft tokens: cost = beam_width + num_draft_tokens. + C++ ref: GenDraftTokensMaxNumTokens + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=10) + requests = [ + make_generation_request(0, beam_width=1, draft_tokens_len=3), + make_generation_request(1, beam_width=1, draft_tokens_len=3), + make_generation_request(2, beam_width=1, draft_tokens_len=3), + ] + ctx, gen = scheduler.schedule(requests, set()) + # Each gen costs 1+3=4. Two fit (8), three don't (12 > 10). + assert len(gen) == 2 + + def test_inflight_requests_excluded(self): + """Requests already in flight are skipped.""" + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=None) + requests = [ + make_context_request(0, prompt_len=10), + make_context_request(1, prompt_len=10), + make_generation_request(2), + ] + ctx, gen = scheduler.schedule(requests, {0, 2}) + # Only request 1 is not in flight + assert len(ctx) == 1 + assert ctx[0].request_id == 1 + assert len(gen) == 0 + + def test_completed_requests_filtered(self): + """Requests in GENERATION_COMPLETE state are filtered out.""" + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=None) + requests = [ + make_context_request(0, prompt_len=10), + make_completed_request(1), + make_generation_request(2), + ] + ctx, gen = scheduler.schedule(requests, set()) + # Completed request 1 is filtered by state gating + assert len(ctx) == 1 + assert len(gen) == 1 + + def test_encoder_request(self): + """Encoder requests use encoder_output_len as token count.""" + scheduler = PyMicroBatchScheduler( + max_batch_size=4, + max_num_tokens=20, + no_schedule_until_state=LlmRequestState.ENCODER_INIT, + ) + requests = [ + make_encoder_request(0, encoder_output_len=15), + make_context_request(1, prompt_len=10), + ] + ctx, gen = scheduler.schedule(requests, set()) + # Encoder: 15 tokens, context: 10. Total 25 > 20. + assert len(ctx) == 1 + assert ctx[0].request_id == 0 + + +# ############################################################################ +# +# Phase 3: Context Chunking Tests +# +# ############################################################################ + + +class TestPyMicroBatchSchedulerChunking: + """ + Tests for context chunking logic in PyMicroBatchScheduler. + Aligned with C++ ContextChunkingTest in microBatchSchedulerTest.cpp. + """ + + # --- EQUAL_PROGRESS policy --- + + def test_equal_progress_basic(self): + """ + Two context requests split equally across token budget. + C++ ref: ContextChunkingTest with EQUAL_PROGRESS + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=10, ctx_chunk_config=config + ) + requests = [ + make_context_request(0, prompt_len=20), + make_context_request(1, prompt_len=20), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + # Each should get ~5 tokens (equal progress, unit=5, total=10) + total_chunk = sum(r.context_chunk_size for r in ctx) + assert total_chunk <= 10 + + def test_equal_progress_uneven_remaining(self): + """ + One request has less remaining context than the other. + Equal progress should still give fair allocation. + After chunking, sort puts not-last-chunk requests first. + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=15, ctx_chunk_config=config + ) + requests = [ + make_context_request(0, prompt_len=3), # Only 3 tokens remaining + make_context_request(1, prompt_len=20), # Lots remaining + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + # Look up by request_id since sort reorders (not-last-chunk first) + req0 = next(r for r in ctx if r.request_id == 0) + req1 = next(r for r in ctx if r.request_id == 1) + # Request 0 should get at most 3 (all it has) + assert req0.context_chunk_size <= 3 + # Total should fit within budget + assert req0.context_chunk_size + req1.context_chunk_size <= 15 + + def test_fcfs_basic(self): + """ + FIRST_COME_FIRST_SERVED: first request gets as much as possible. + """ + config = ContextChunkingConfig( + ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5 + ) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=12, ctx_chunk_config=config + ) + requests = [ + make_context_request(0, prompt_len=20), + make_context_request(1, prompt_len=20), + ] + ctx, gen = scheduler.schedule(requests, set()) + # FCFS: request 0 gets up to budget, request 1 gets remainder + assert len(ctx) >= 1 + # First request should get more tokens + assert ctx[0].context_chunk_size >= ctx[-1].context_chunk_size or len(ctx) == 1 + + def test_fcfs_fills_first_request(self): + """FCFS fills the first request completely if budget allows. + After chunking, sort puts not-last-chunk requests first.""" + config = ContextChunkingConfig( + ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5 + ) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=25, ctx_chunk_config=config + ) + requests = [ + make_context_request(0, prompt_len=10), + make_context_request(1, prompt_len=20), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + # Look up by request_id since sort reorders (not-last-chunk first) + req0 = next(r for r in ctx if r.request_id == 0) + req1 = next(r for r in ctx if r.request_id == 1) + # First request should get all 10 (full context) + assert req0.context_chunk_size == 10 + # Second gets remainder up to its need + assert req1.context_chunk_size <= 15 + + def test_chunk_with_generation(self): + """ + Chunked context + generation in the same batch. + Generation tokens reduce the available budget for context chunks. + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=15, ctx_chunk_config=config + ) + requests = [ + make_generation_request(0), # costs 1 token + make_context_request(1, prompt_len=20), + make_context_request(2, prompt_len=20), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(gen) == 1 + # Remaining budget for context: 15 - 1 = 14 + total_ctx_tokens = sum(r.context_chunk_size for r in ctx) + assert total_ctx_tokens <= 14 + + def test_chunk_size_zero_not_scheduled(self): + """ + If a request gets chunk_size=0 after chunking, it's not added to + the scheduled context requests. + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=2, max_num_tokens=5, ctx_chunk_config=config + ) + requests = [ + make_context_request(0, prompt_len=20), + make_context_request(1, prompt_len=20), + ] + ctx, gen = scheduler.schedule(requests, set()) + # With budget 5, at most one request gets chunk_size=5, the other might get 0 + for r in ctx: + assert r.context_chunk_size > 0 + + def test_chunking_with_max_context_length(self): + """ + max_context_length (same as max_num_tokens) limits individual chunk size. + C++ ref: SimpleNoOverlapMaxContextLength + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=12, ctx_chunk_config=config + ) + requests = [ + make_context_request(0, prompt_len=20), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 1 + # max_context_length = max_num_tokens = 12, so chunk <= 12 + assert ctx[0].context_chunk_size <= 12 + + def test_continued_chunking(self): + """ + A request that has already processed part of its context + (context_position > 0) continues from where it left off. + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=10, ctx_chunk_config=config + ) + req = make_context_request(0, prompt_len=20, context_position=10) + # remaining = 20 - 10 = 10 + ctx, gen = scheduler.schedule([req], set()) + assert len(ctx) == 1 + assert ctx[0].context_chunk_size <= 10 # remaining context + + def test_last_chunk_allows_draft_tokens(self): + """ + On the last chunk, draft tokens are included if they fit within + the remaining space in the chunk unit. + C++ ref: DraftTokensNoDiscard + """ + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=10) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=20, ctx_chunk_config=config + ) + # prompt_len=8, so chunk_size will be 8. Unit=10, remainder=2. + # Draft tokens=2 fits in remainder. + req = make_context_request(0, prompt_len=8, draft_tokens_len=2) + ctx, gen = scheduler.schedule([req], set()) + assert len(ctx) == 1 + assert req.is_last_context_chunk + + def test_draft_tokens_discarded_when_no_space(self): + """ + Draft tokens that don't fit in chunk unit remainder are discarded. + C++ ref: DraftTokensDiscard + """ + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=20, ctx_chunk_config=config + ) + # prompt_len=5, chunk_size=5, unit=5, remainder=0. Draft=3 won't fit. + req = make_context_request(0, prompt_len=5, draft_tokens_len=3) + ctx, gen = scheduler.schedule([req], set()) + assert len(ctx) == 1 + + def test_chunked_context_draft_tokens_max_num_tokens(self): + """ + Chunked context + draft tokens: maxNumTokens limits total draft budget. + C++ ref: ChunkedContextDraftTokensMaxNumTokens + maxNumTokens=8192, maxBatchSize=64, chunkUnitSize=64, FCFS, + promptLen=2041, draftLen=8, 4 requests. + 2041 = 31*64 + 57, so remainder in unit = 64-57 = 7. + Each request's draft reduced from 8 to 7. + """ + config = ContextChunkingConfig( + ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=64 + ) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, + max_num_tokens=8192, + ctx_chunk_config=config, + ) + requests = [ + make_context_request(i, prompt_len=2041, draft_tokens_len=8) + for i in range(4) + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 4 + for req in ctx: + assert req.num_draft_tokens == 7 + + def test_chunked_context_draft_tokens_max_context_length(self): + """ + Chunked context + draft tokens: maxContextLength limits individual draft. + C++ ref: ChunkedContextDraftTokensMaxContextLength + maxContextLength=10, maxNumTokens=8192, maxBatchSize=64, + chunkUnitSize=64, FCFS, promptLen=6, draftLen=5, 2 requests. + chunk_size=6, unit=64, remainder=6, space_in_unit=58. + But maxContextLength-chunk_size = 10-6 = 4, so remaining_space=4. + Draft reduced from 5 to 4. + """ + config = ContextChunkingConfig( + ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=64 + ) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, + max_num_tokens=8192, + ctx_chunk_config=config, + ) + scheduler.max_context_length = 10 + requests = [ + make_context_request(0, prompt_len=6, draft_tokens_len=5), + make_context_request(1, prompt_len=6, draft_tokens_len=5), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + for req in ctx: + assert req.num_draft_tokens == 4 + + def test_no_chunking_context_fits(self): + """Without chunking, context is scheduled in full if it fits.""" + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=20, ctx_chunk_config=None + ) + req = make_context_request(0, prompt_len=15) + ctx, gen = scheduler.schedule([req], set()) + assert len(ctx) == 1 + + def test_no_chunking_context_exceeds_budget(self): + """Without chunking, cumulative tokens exceeding budget stops scheduling. + Each individual request must fit within max_context_length (== max_num_tokens), + but the cumulative token count is checked against the budget. The first request + that would push the total over the limit breaks the loop.""" + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=10, ctx_chunk_config=None + ) + requests = [ + make_context_request(0, prompt_len=8), + make_context_request(1, prompt_len=8), + ] + ctx, gen = scheduler.schedule(requests, set()) + # First request (8) fits (8 <= 10). Second (8+8=16 > 10) breaks the loop. + assert len(ctx) == 1 + assert ctx[0].request_id == 0 + + def test_sort_by_lora_task_id(self): + """ + Requests are sorted by lora_task_id for performance. + C++ ref: sortRequests in inflightBatchingUtils.cpp + """ + scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=None) + r0 = make_generation_request(0) + r0.lora_task_id = 5 + r1 = make_generation_request(1) + r1.lora_task_id = None + r2 = make_generation_request(2) + r2.lora_task_id = 3 + ctx, gen = scheduler.schedule([r0, r1, r2], set()) + # None < any value, so order should be: r1(None), r2(3), r0(5) + assert gen[0].request_id == 1 + assert gen[1].request_id == 2 + assert gen[2].request_id == 0 + + +# ############################################################################ +# +# Phase 3b: Direct Context Chunking Tests (mirrors C++ ContextChunkingTest) +# +# ############################################################################ + + +def _run_context_chunking_test( + context_lengths: List[int], + chunk_unit_size: int, + ep_positions: List[List[int]], + fcfs_positions: List[List[int]], + ctx_tokens_capacity: Optional[int] = None, + max_context_length: Optional[int] = None, + draft_lengths: Optional[List[int]] = None, + ep_draft_lengths: Optional[List[List[int]]] = None, + fcfs_draft_lengths: Optional[List[List[int]]] = None, +): + """ + Helper that mirrors the C++ ContextChunkingTest fixture. + + For each policy (EQUAL_PROGRESS and FCFS), it: + 1. Creates MockLlmRequests with given context_lengths and optional draft_lengths. + 2. Creates a PyMicroBatchScheduler with the right ContextChunkingConfig. + 3. If max_context_length is set, overrides scheduler.max_context_length. + 4. For each iteration (each element in positions list): + a. Filters requests where context_remaining_length > 0. + b. Calls scheduler._set_ctx_requests_chunk_size(active_reqs, ctx_tokens_capacity). + c. For each active req, calls req.move_to_next_context_chunk(). + d. Verifies _context_current_position matches expected positions for ALL requests. + 5. After all iterations, verifies final draft_lengths if specified. + """ + policies_and_data = [ + (ChunkingPolicy.EQUAL_PROGRESS, ep_positions, ep_draft_lengths), + (ChunkingPolicy.FIRST_COME_FIRST_SERVED, fcfs_positions, fcfs_draft_lengths), + ] + + for policy, positions_list, final_draft_lens in policies_and_data: + # Create fresh requests for each policy + requests = [] + for i, ctx_len in enumerate(context_lengths): + dl = draft_lengths[i] if draft_lengths else 0 + req = MockLlmRequest( + request_id=i, + _state=MockLlmRequestState.CONTEXT_INIT, + prompt_len=ctx_len, + _context_current_position=0, + _context_chunk_size=0, + _draft_tokens_len=dl, + ) + requests.append(req) + + # Create scheduler + config = ContextChunkingConfig(policy, chunk_unit_size=chunk_unit_size) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, + max_num_tokens=1000, # large enough not to limit + ctx_chunk_config=config, + ) + + if max_context_length is not None: + scheduler.max_context_length = max_context_length + + # Run iterations + for iteration_idx, expected_positions in enumerate(positions_list): + # Filter active requests (those with remaining context) + active_reqs = [r for r in requests if r.context_remaining_length > 0] + + scheduler._set_ctx_requests_chunk_size(active_reqs, ctx_tokens_capacity) + + # Move each active request to next chunk + for req in active_reqs: + req.move_to_next_context_chunk() + + # Verify positions for ALL requests (including completed ones) + for req_idx, req in enumerate(requests): + assert req._context_current_position == expected_positions[req_idx], ( + f"Policy {policy.name}, iteration {iteration_idx}, " + f"request {req_idx}: expected position " + f"{expected_positions[req_idx]}, got " + f"{req._context_current_position}" + ) + + # Verify final draft lengths if specified + if final_draft_lens is not None: + for req_idx, req in enumerate(requests): + assert req._draft_tokens_len == final_draft_lens[req_idx], ( + f"Policy {policy.name}, request {req_idx}: " + f"expected draft_len {final_draft_lens[req_idx]}, " + f"got {req._draft_tokens_len}" + ) + + +class TestContextChunkingDirect: + """ + Direct tests for _set_ctx_requests_chunk_size logic, mirroring + C++ ContextChunkingTest in microBatchSchedulerTest.cpp. + Each test calls _run_context_chunking_test with parameters matching + exactly the C++ test cases. + """ + + def test_no_limit(self): + """C++ ref: NoLimit""" + _run_context_chunking_test([25, 25], 20, [[25, 25]], [[25, 25]]) + + def test_context_length_never_satisfied(self): + """C++ ref: ContextLengthNeverSatisfied""" + _run_context_chunking_test( + [25, 25], 100, [[0, 0]], [[0, 0]], max_context_length=20 + ) + + def test_chunk_longer_than_context(self): + """C++ ref: ChunkLongerThanContext""" + _run_context_chunking_test( + [25, 25], 30, [[25, 25]], [[25, 25]], max_context_length=25 + ) + + def test_context_length_satisfied(self): + """C++ ref: ContextLengthSatisfied""" + _run_context_chunking_test( + [10, 25], + 10, + [[10, 20], [10, 25]], + [[10, 20], [10, 25]], + max_context_length=20, + ) + + def test_token_capacity_smaller_than_context(self): + """C++ ref: TokenCapacitySmallerThanContext""" + _run_context_chunking_test( + [25, 25], + 20, + [[20, 0], [25, 0], [25, 20], [25, 25]], + [[20, 0], [25, 0], [25, 20], [25, 25]], + ctx_tokens_capacity=20, + ) + + def test_token_capacity_smaller_than_chunk_unit(self): + """C++ ref: TokenCapacitySmallerThanChunkUnit""" + _run_context_chunking_test( + [25, 25], 20, [[0, 0]], [[0, 0]], ctx_tokens_capacity=10 + ) + + def test_scheduling_order(self): + """C++ ref: SchedulingOrder""" + _run_context_chunking_test( + [25, 25], + 5, + [[15, 15], [25, 25]], + [[25, 5], [25, 25]], + ctx_tokens_capacity=30, + ) + + def test_completion_order(self): + """C++ ref: CompletionOrder""" + _run_context_chunking_test( + [25, 15], + 5, + [[15, 15], [25, 15]], + [[25, 5], [25, 15]], + ctx_tokens_capacity=30, + ) + + def test_long_first_short_later(self): + """C++ ref: LongFirstShortLater""" + _run_context_chunking_test( + [25, 15], + 5, + [[10, 10], [20, 15], [25, 15]], + [[10, 10], [20, 15], [25, 15]], + ctx_tokens_capacity=30, + max_context_length=10, + ) + + def test_front_priority(self): + """C++ ref: FrontPriority""" + _run_context_chunking_test( + [25, 25], + 5, + [[10, 5], [20, 10], [25, 20], [25, 25]], + [[15, 0], [25, 5], [25, 20], [25, 25]], + ctx_tokens_capacity=15, + ) + + def test_draft_tokens_discard(self): + """C++ ref: DraftTokensDiscard""" + _run_context_chunking_test( + [27, 27], + 5, + [[15, 15], [27, 27]], + [[27, 0], [27, 27]], + ctx_tokens_capacity=30, + draft_lengths=[5, 5], + ep_draft_lengths=[3, 3], + fcfs_draft_lengths=[3, 3], + ) + + def test_draft_tokens_discard2(self): + """C++ ref: DraftTokensDiscard2""" + _run_context_chunking_test( + [17, 17], + 5, + [[15, 15], [17, 17]], + [[17, 10], [17, 17]], + ctx_tokens_capacity=30, + draft_lengths=[5, 5], + ep_draft_lengths=[3, 3], + fcfs_draft_lengths=[3, 3], + ) + + def test_draft_tokens_discard3(self): + """C++ ref: DraftTokensDiscard3""" + _run_context_chunking_test( + [27, 27], + 5, + [[10, 10], [20, 20], [27, 27]], + [[20, 0], [27, 10], [27, 27]], + ctx_tokens_capacity=20, + draft_lengths=[5, 5], + ep_draft_lengths=[3, 3], + fcfs_draft_lengths=[3, 3], + ) + + def test_draft_tokens_discard_due_to_token_capacity(self): + """C++ ref: DraftTokensDiscardDueToTokenCapacity""" + _run_context_chunking_test( + [23, 17], + 5, + [[10, 10], [23, 17]], + [[20, 0], [23, 17]], + ctx_tokens_capacity=20, + draft_lengths=[5, 5], + ep_draft_lengths=[0, 0], + fcfs_draft_lengths=[0, 0], + ) + + def test_draft_tokens_discard_due_to_max_context_length(self): + """C++ ref: DraftTokensDiscardDueToMaxContextLength""" + _run_context_chunking_test( + [6, 6], + 5, + [[6, 6]], + [[6, 6]], + ctx_tokens_capacity=30, + max_context_length=10, + draft_lengths=[5, 5], + ep_draft_lengths=[4, 4], + fcfs_draft_lengths=[4, 4], + ) + + def test_draft_tokens_discard_all(self): + """C++ ref: DraftTokensDiscardAll""" + _run_context_chunking_test( + [25, 25], + 5, + [[25, 25]], + [[25, 25]], + ctx_tokens_capacity=50, + draft_lengths=[5, 5], + ep_draft_lengths=[0, 0], + fcfs_draft_lengths=[0, 0], + ) + + def test_draft_tokens_discard_all2(self): + """C++ ref: DraftTokensDiscardAll2""" + _run_context_chunking_test( + [25, 25], + 5, + [[15, 10], [25, 25]], + [[25, 0], [25, 25]], + ctx_tokens_capacity=25, + draft_lengths=[5, 5], + ep_draft_lengths=[0, 0], + fcfs_draft_lengths=[0, 0], + ) + + def test_draft_tokens_no_discard(self): + """C++ ref: DraftTokensNoDiscard""" + _run_context_chunking_test( + [25, 25], + 10, + [[20, 10], [25, 25]], + [[25, 0], [25, 25]], + ctx_tokens_capacity=30, + draft_lengths=[5, 5], + ep_draft_lengths=[5, 5], + fcfs_draft_lengths=[5, 5], + ) + + def test_draft_tokens_no_chunking_discard_all(self): + """C++ ref: DraftTokensNoChunkingDiscardAll""" + _run_context_chunking_test( + [4128], + 64, + [[4128]], + [[4128]], + max_context_length=4128, + draft_lengths=[3], + ep_draft_lengths=[0], + fcfs_draft_lengths=[0], + ) + + def test_draft_tokens_no_chunking_discard_some(self): + """C++ ref: DraftTokensNoChunkingDiscardSome""" + _run_context_chunking_test( + [4127], + 64, + [[4127]], + [[4127]], + max_context_length=4128, + draft_lengths=[3], + ep_draft_lengths=[1], + fcfs_draft_lengths=[1], + ) + + def test_draft_tokens_no_chunking_discard_none(self): + """C++ ref: DraftTokensNoChunkingDiscardNone""" + _run_context_chunking_test( + [4125], + 64, + [[4125]], + [[4125]], + max_context_length=4128, + draft_lengths=[3], + ep_draft_lengths=[3], + fcfs_draft_lengths=[3], + ) + + +class TestDraftTokensGreaterThanChunkSize: + """ + Tests that when draft tokens > chunk unit, they get properly trimmed. + C++ ref: DraftTokensGreaterThanChunkSize in microBatchSchedulerTest.cpp + """ + + def test_draft_tokens_greater_than_chunk_size(self): + """ + C++ ref: DraftTokensGreaterThanChunkSize + maxNumTokens=40, maxBatchSize=64, chunkUnitSize=16, FCFS policy, + maxContextLength=64. 3 requests: promptLen=3, draftLen=17. + + After scheduling, expected: + - All 3 scheduled + - Request 0: draftTokens = 13 (unit=16, context=3, remainder=13) + - Request 1: draftTokens = 13 + - Request 2: draftTokens = 5 (remaining budget) + """ + config = ContextChunkingConfig( + ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=16 + ) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, + max_num_tokens=40, + ctx_chunk_config=config, + ) + scheduler.max_context_length = 64 + + requests = [ + make_context_request(0, prompt_len=3, draft_tokens_len=17), + make_context_request(1, prompt_len=3, draft_tokens_len=17), + make_context_request(2, prompt_len=3, draft_tokens_len=17), + ] + + ctx, gen = scheduler.schedule(requests, set()) + + assert len(ctx) == 3 + req0 = next(r for r in ctx if r.request_id == 0) + req1 = next(r for r in ctx if r.request_id == 1) + req2 = next(r for r in ctx if r.request_id == 2) + + assert req0.num_draft_tokens == 13 + assert req1.num_draft_tokens == 13 + assert req2.num_draft_tokens == 5 + + +# ############################################################################ +# +# Phase 4: PyCapacityScheduler Tests +# +# ############################################################################ + + +class TestPyCapacitySchedulerMaxRequests: + """Tests for MaxRequestsPolicy — no KV cache, simple count limit.""" + + def test_basic_scheduling(self): + """Schedule up to max_num_requests.""" + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=None, # triggers MaxRequestsPolicy + ) + requests = [ + make_context_request(0), + make_context_request(1), + make_context_request(2), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) == 2 + assert len(paused) == 0 + + def test_mixed_states(self): + """Only schedulable states are included.""" + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=None, + ) + requests = [ + make_context_request(0), + make_generation_request(1), + make_completed_request(2), # should be filtered + make_context_request(3), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + # Completed request filtered out + assert len(fitting) == 3 + assert all(r.request_id != 2 for r in fitting) + + def test_empty_requests(self): + """No requests to schedule.""" + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=None, + ) + fitting, disagg, paused = scheduler.schedule_request([]) + assert len(fitting) == 0 + + +class TestPyCapacitySchedulerGuaranteedNoEvict: + """ + Tests for GuaranteedNoEvictPolicy. + C++ ref: capacitySchedulerTest.cpp GuaranteedCompletion tests + """ + + def test_requests_fit(self): + """ + Both requests fit: free blocks (100) > 2 * blocks_per_request (5). + C++ ref: SimpleShouldFit + """ + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0), + make_context_request(1), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) == 2 + + def test_not_enough_blocks(self): + """ + Not enough blocks for the second request. + C++ ref: SimpleDoesntFitGuaranteedCompletion + """ + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0), + make_context_request(1), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + # First request takes 5 blocks, leaving 2 < 5 for second + assert len(fitting) == 1 + assert fitting[0].request_id == 0 + + def test_generation_scheduled_first(self): + """ + In-progress generation requests are scheduled before context requests. + They consume blocks from the pool. + """ + kv = MockKVCacheManager(num_free_blocks=12, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_generation_request(0), # takes 5 blocks, leaves 7 + make_context_request(1), # takes 5 blocks, leaves 2 + make_context_request(2), # needs 5, only 2 left + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) == 2 + assert {r.request_id for r in fitting} == {0, 1} + + def test_max_num_requests_honored(self): + """max_num_requests is respected even if blocks are available.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0), + make_context_request(1), + make_context_request(2), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) == 2 + + +class TestPyCapacitySchedulerMaxUtilization: + """ + Tests for MaxUtilizationPolicy. + C++ ref: capacitySchedulerTest.cpp MaxUtilization tests + """ + + def test_fits(self): + """Requests fit within free blocks.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + requests = [ + make_context_request(0), + make_generation_request(1), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) == 2 + assert len(paused) == 0 + + def test_doesnt_fit_pauses(self): + """ + When a new request can't fit, MaxUtilization tries to pause + already-started requests to make room. + C++ ref: SimpleDoesntFitMaxUtilization + """ + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + # gen_0 is started (in progress), context_1 is new + requests = [ + make_generation_request(0), + make_context_request(1), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + # gen_0 takes 5 blocks, context_1 needs 5 but only 2 left. + # MaxUtilization pauses gen_0 to make room for context_1. + assert len(fitting) + len(paused) >= 1 + + def test_no_requests_to_pause(self): + """If no started requests to pause, scheduling stops.""" + kv = MockKVCacheManager(num_free_blocks=3, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + requests = [ + make_context_request(0), + make_context_request(1), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + # Neither fits (3 < 5), no started requests to pause + assert len(fitting) == 0 + assert len(paused) == 0 + + +class TestPyCapacitySchedulerStaticBatch: + """ + Tests for STATIC_BATCH policy. + C++ ref: SimpleFitsStaticBatch + """ + + def test_schedules_when_no_active_generation(self): + """Static batch: schedule all context requests when idle.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.STATIC_BATCH, + ) + requests = [ + make_context_request(0), + make_context_request(1), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) == 2 + + def test_no_new_context_when_generation_active(self): + """Static batch: no new context when generation is in progress.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.STATIC_BATCH, + ) + requests = [ + make_generation_request(0), + make_context_request(1), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + # Static batch: gen_0 is active, so only gen_0 scheduled (no new context) + assert len(fitting) == 1 + assert fitting[0].request_id == 0 + + +class TestPyCapacitySchedulerDisagg: + """Tests for disaggregated generation init handling.""" + + def test_disagg_gen_init_classified_separately(self): + """Disagg gen init requests are separated in output. + Uses GUARANTEED_NO_EVICT policy with a KV cache manager, since + MaxRequestsPolicy (used when kv_cache_manager=None) does not + handle disagg_generation_init state.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0), + make_disagg_gen_init_request(1), + make_generation_request(2), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert any(r.request_id == 0 for r in fitting) + assert any(r.request_id == 2 for r in fitting) + assert len(disagg) == 1 + assert disagg[0].request_id == 1 + + def test_disagg_gen_init_bypasses_state_gating(self): + """ + Disagg gen init requests bypass normal state gating + (no_schedule_until / no_schedule_after). + """ + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_disagg_gen_init_request(0), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(disagg) == 1 + + +class TestPyCapacitySchedulerStateGating: + """Tests for state-based filtering.""" + + def test_unknown_state_filtered(self): + """UNKNOWN state is before no_schedule_until, filtered out.""" + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=None, + ) + req = MockLlmRequest(request_id=0, _state=MockLlmRequestState.UNKNOWN) + fitting, disagg, paused = scheduler.schedule_request([req]) + assert len(fitting) == 0 + + def test_generation_complete_filtered(self): + """GENERATION_COMPLETE is at no_schedule_after, filtered out.""" + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=None, + ) + fitting, disagg, paused = scheduler.schedule_request( + [make_completed_request(0)] + ) + assert len(fitting) == 0 + + def test_generation_to_complete_included(self): + """GENERATION_TO_COMPLETE is within schedulable range and handled + like GENERATION_IN_PROGRESS by all capacity scheduler policies.""" + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=None, + ) + req = MockLlmRequest( + request_id=0, + _state=MockLlmRequestState.GENERATION_TO_COMPLETE, + ) + fitting, disagg, paused = scheduler.schedule_request([req]) + # GENERATION_TO_COMPLETE passes state gating and is handled by MaxRequestsPolicy + assert len(fitting) == 1 + + +# ############################################################################ +# +# Phase 5: PyCapacityScheduler Advanced Tests +# +# ############################################################################ + + +class TestPyCapacitySchedulerLora: + """Tests for LoRA/PEFT integration in capacity scheduling.""" + + def test_lora_fits(self): + """ + LoRA requests fit within PEFT cache. + C++ ref: SimpleLoraFitsDuplicateTask + """ + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + peft = MockPeftCacheManager(max_pages=100, pages_per_request=10) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + peft_cache_manager=peft, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0) + r0.lora_task_id = 1 + r1 = make_context_request(1) + r1.lora_task_id = 1 # same task — no extra pages needed + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) == 2 + + def test_lora_doesnt_fit(self): + """ + LoRA requests exceed PEFT cache. + C++ ref: SimpleLoraDoesntFitDuplicateTask + """ + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + peft = MockPeftCacheManager(max_pages=15, pages_per_request=10) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + peft_cache_manager=peft, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0) + r0.lora_task_id = 1 + r1 = make_context_request(1) + r1.lora_task_id = 2 # different task — needs 10 more pages + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + # First task: 10 pages, second task: 10 pages, total 20 > 15 + assert len(fitting) == 1 + + +# ############################################################################ +# +# Phase 6: SimpleUnifiedScheduler Integration Tests +# +# ############################################################################ + + +class TestSimpleUnifiedScheduler: + """ + Tests for the two-stage scheduling pipeline: + PyCapacityScheduler → PyMicroBatchScheduler + """ + + def test_capacity_then_microbatch(self): + """Capacity filters, then microbatch selects within token budget. + max_batch_size is used as max_num_requests for capacity scheduler, + so it must be large enough for all requests to pass capacity.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = SimpleUnifiedScheduler( + max_batch_size=4, + max_num_tokens=15, + kv_cache_manager=kv, + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0, prompt_len=10), + make_context_request(1, prompt_len=10), + make_generation_request(2), + ] + output = scheduler.schedule_request(requests, set()) + # Capacity: all 3 fit (plenty of blocks, max_num_requests=4) + # Microbatch: gen_2 (1) + context_0 (10) = 11 <= 15, context_1 (10) would be 21 > 15 + assert output.num_fitting_requests == 3 + assert len(output.context_requests) + len(output.generation_requests) <= 2 + + def test_can_schedule_dry_run(self): + """can_schedule() checks capacity without side effects.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = SimpleUnifiedScheduler( + max_batch_size=4, + max_num_tokens=100, + kv_cache_manager=kv, + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0), + make_generation_request(1), + ] + assert scheduler.can_schedule(requests) is True + + def test_can_schedule_returns_false(self): + """can_schedule() returns False when capacity is insufficient.""" + kv = MockKVCacheManager(num_free_blocks=3, blocks_per_request=5) + scheduler = SimpleUnifiedScheduler( + max_batch_size=4, + max_num_tokens=100, + kv_cache_manager=kv, + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0), + make_context_request(1), + ] + # First takes 5 blocks > 3 free + assert scheduler.can_schedule(requests) is False + + def test_full_pipeline_output_structure(self): + """Verify SchedulerOutput has all expected fields.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = SimpleUnifiedScheduler( + max_batch_size=4, + max_num_tokens=100, + kv_cache_manager=kv, + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_context_request(0, prompt_len=10), + make_generation_request(1), + ] + output = scheduler.schedule_request(requests, set()) + assert hasattr(output, "context_requests") + assert hasattr(output, "generation_requests") + assert hasattr(output, "paused_requests") + assert hasattr(output, "fitting_disagg_gen_init_requests") + assert hasattr(output, "num_fitting_requests") + assert len(output.context_requests) == 1 + assert len(output.generation_requests) == 1 + assert output.context_requests[0].request_id == 0 + assert output.generation_requests[0].request_id == 1 + + def test_paused_requests_propagated(self): + """Paused requests from capacity scheduler appear in output.""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = SimpleUnifiedScheduler( + max_batch_size=4, + max_num_tokens=100, + kv_cache_manager=kv, + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + # With MAX_UTILIZATION and plenty of resources, nothing should be paused + requests = [ + make_context_request(0, prompt_len=10), + make_generation_request(1), + ] + output = scheduler.schedule_request(requests, set()) + assert isinstance(output.paused_requests, list) + + +# ############################################################################ +# +# Phase 7: Additional PyCapacityScheduler Tests +# +# ############################################################################ + + +class TestPyCapacitySchedulerCrossKVCache: + """ + Tests for cross-attention KV cache scheduling. + C++ ref: capacitySchedulerTest.cpp cross KV cache tests. + """ + + def test_should_fit_with_cross_blocks(self): + """C++ ref: SimpleShouldFitWithCrossBlocks""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + cross_kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=2) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + cross_kv_cache_manager=cross_kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0, prompt_len=10) + r0.encoder_output_len = 10 + r1 = make_context_request(1, prompt_len=10) + r1.encoder_output_len = 10 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) == 2 + + def test_doesnt_fit_with_cross_blocks(self): + """C++ ref: SimpleDoesntFitWithCrossBlocks - cross kv cache too small for both""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + cross_kv = MockKVCacheManager(num_free_blocks=1, blocks_per_request=1) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + cross_kv_cache_manager=cross_kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0, prompt_len=10) + r0.encoder_output_len = 10 + r1 = make_context_request(1, prompt_len=10) + r1.encoder_output_len = 10 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) == 1 + + +class TestPyCapacitySchedulerPriority: + """ + Tests for priority-related scheduling behavior. + C++ ref: capacitySchedulerTest.cpp priority tests. + """ + + def test_requests_sorted_by_priorities(self): + """C++ ref: RequestsSortedByPriorities. + Python scheduler doesn't have insertRequestInOrder, but this tests + that requests are processed in the order they're given.""" + scheduler = PyCapacityScheduler(max_num_requests=12, kv_cache_manager=None) + # Create 12 requests - all should be scheduled in order + requests = [make_context_request(i) for i in range(12)] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) == 12 + assert [r.request_id for r in fitting] == list(range(12)) + + def test_doesnt_fit_priorities(self): + """C++ ref: SimpleDoesntFitPriorities - MAX_UTILIZATION with limited blocks""" + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + requests = [make_context_request(0), make_context_request(1)] + fitting, disagg, paused = scheduler.schedule_request(requests) + # Only first request fits (5 blocks), second needs 5 but only 2 left + assert len(fitting) >= 1 + + +class TestPyCapacitySchedulerChunked: + """ + Tests for chunked context request capacity scheduling. + C++ ref: capacitySchedulerTest.cpp chunked tests. + """ + + def test_should_fit_in_chunk(self): + """C++ ref: SimpleShouldFitInChunk - chunked context requests fit""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0, prompt_len=50, context_position=0) + r0._context_chunk_size = 20 + r1 = make_context_request(1, prompt_len=50, context_position=0) + r1._context_chunk_size = 20 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) == 2 + + def test_doesnt_fit_guaranteed_completion_in_chunk(self): + """C++ ref: SimpleDoesntFitGuaranteedCompletionInChunk""" + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0, prompt_len=30) + r0._context_chunk_size = 20 + r1 = make_context_request(1, prompt_len=30) + r1._context_chunk_size = 20 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) == 1 + + def test_doesnt_fit_max_utilization_in_chunk(self): + """C++ ref: SimpleDoesntFitMaxUtilizationInChunk""" + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + r0 = make_context_request(0, prompt_len=30) + r0._context_chunk_size = 20 + r1 = make_context_request(1, prompt_len=30) + r1._context_chunk_size = 20 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + # At least one should be scheduled + assert len(fitting) >= 1 + + def test_doesnt_fit_max_utilization_in_chunked_cache(self): + """C++ ref: SimpleDoesntFitMaxUtilizationInChunkedCache - only 1 fits due to limited blocks""" + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=7) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + r0 = make_context_request(0, prompt_len=70) + r0._context_chunk_size = 40 + r1 = make_context_request(1, prompt_len=70) + r1._context_chunk_size = 40 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) == 1 + + def test_doesnt_fit_max_utilization_draft_tokens(self): + """C++ ref: SimpleDoesntFitMaxUtilizationDraftTokens""" + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + r0 = make_context_request(0, prompt_len=10, draft_tokens_len=5) + r1 = make_context_request(1, prompt_len=10, draft_tokens_len=10) + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + +class TestPyCapacitySchedulerDynamicAddition: + """ + Tests for dynamic request addition during scheduling. + C++ ref: capacitySchedulerTest.cpp dynamic addition tests. + """ + + def test_adding_new_requests_max_utilization(self): + """C++ ref: SimpleDoesntFitAddingNewRequestsMaxUtilization + Initial state: 2 gen requests in progress, 2 new context requests. + With limited blocks, not all can run.""" + kv = MockKVCacheManager(num_free_blocks=10, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + requests = [ + make_generation_request(0), + make_generation_request(1), + make_context_request(2), + make_context_request(3), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) + len(paused) >= 2 + + def test_adding_new_requests_guaranteed_completion(self): + """C++ ref: SimpleDoesntFitAddingNewRequestsGuaranteedCompletion""" + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + requests = [ + make_generation_request(0), + make_context_request(1), + make_context_request(2), + make_context_request(3), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + # gen_0 takes 5 blocks (7-5=2 left), none of the context requests fit + assert len(fitting) >= 1 + assert fitting[0].request_id == 0 + + def test_adding_new_requests_guaranteed_completion_in_chunk(self): + """C++ ref: SimpleDoesntFitAddingNewRequestsGuaranteedCompletionInChunk""" + kv = MockKVCacheManager(num_free_blocks=7, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_generation_request(0) + r1 = make_context_request(1, prompt_len=30) + r1._context_chunk_size = 20 + r2 = make_context_request(2, prompt_len=30) + r2._context_chunk_size = 20 + fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2]) + assert len(fitting) >= 1 + + def test_surpass_max_num_requests_with_priorities(self): + """C++ ref: SimpleSurpassMaxNumRequestsWithPriorities""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=2) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + # In C++, high-priority requests preempt lower-priority ones. + # In Python, maxNumRequests limits total scheduled. + requests = [ + make_context_request(0), + make_context_request(1), + make_context_request(2), + make_context_request(3), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) <= 2 + + def test_adding_new_requests_max_utilization_priorities(self): + """C++ ref: SimpleDoesntFitAddingNewRequestsMaxUtilizationPriorities""" + kv = MockKVCacheManager(num_free_blocks=10, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + requests = [ + make_generation_request(0), + make_generation_request(1), + make_context_request(2), + make_context_request(3), + ] + fitting, disagg, paused = scheduler.schedule_request(requests) + assert len(fitting) + len(paused) >= 2 + + +class TestPyCapacitySchedulerKVCacheReuse: + """ + Tests for KV cache reuse-aware scheduling. + C++ ref: DelayDuplicate*, ReuseAware*, MaxUtilizationReuse*, NoReuse* tests + in capacitySchedulerTest.cpp. + """ + + def test_delay_duplicate_request(self): + """C++ ref: DelayDuplicateRequest - identical requests delayed for reuse""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=3, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + tokens = list(range(21)) + r0 = make_context_request(0, prompt_len=21) + r0._input_tokens = tokens + r1 = make_context_request(1, prompt_len=21) + r1._input_tokens = tokens + r2 = make_context_request(2, prompt_len=21) + r2._input_tokens = tokens + fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2]) + # With reuse enabled, beneficial_to_skip may delay r1 and r2 + assert len(fitting) >= 1 + + def test_delay_duplicate_request_chunked(self): + """C++ ref: DelayDuplicateRequestChunked""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=5, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + tokens = list(range(50)) + r0 = make_context_request(0, prompt_len=50) + r0._input_tokens = tokens + r0._context_chunk_size = 20 + r1 = make_context_request(1, prompt_len=50) + r1._input_tokens = tokens + r1._context_chunk_size = 20 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + def test_delay_five_requests_complicated(self): + """C++ ref: DelayFiveRequestsComplicated""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=5, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0, prompt_len=11) + r0._input_tokens = [x + 1 for x in range(11)] # unique + r1 = make_context_request(1, prompt_len=21) + r1._input_tokens = [x + 2 for x in range(21)] # unique + r2 = make_context_request(2, prompt_len=11) + r2._input_tokens = list(range(11)) + r3 = make_context_request(3, prompt_len=21) + r3._input_tokens = list(range(21)) # first 11 shared + r4 = make_context_request(4, prompt_len=31) + r4._input_tokens = list(range(31)) # first 11 shared + fitting, disagg, paused = scheduler.schedule_request( + [r0, r1, r2, r3, r4] + ) + assert len(fitting) >= 1 + + def test_reuse_aware_allows_more_requests(self): + """C++ ref: ReuseAwareSchedulingAllowsMoreRequestsWithSharedPrefix""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + tokens = list(range(20)) + r0 = make_context_request(0, prompt_len=20) + r0._input_tokens = tokens + r1 = make_context_request(1, prompt_len=20) + r1._input_tokens = tokens + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + def test_reuse_aware_partial_prefix_match(self): + """C++ ref: ReuseAwareSchedulingWithPartialPrefixMatch""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=3, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + tokens0 = list(range(30)) + tokens1 = list(range(20)) + [999] * 10 + r0 = make_context_request(0, prompt_len=30) + r0._input_tokens = tokens0 + r1 = make_context_request(1, prompt_len=30) + r1._input_tokens = tokens1 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + def test_no_reuse_with_different_prompts(self): + """C++ ref: NoReuseWithDifferentPrompts""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + ) + r0 = make_context_request(0, prompt_len=20) + r0._input_tokens = [100] * 20 + r1 = make_context_request(1, prompt_len=20) + r1._input_tokens = [200] * 20 + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + # Different prompts: no reuse, both scheduled + assert len(fitting) == 2 + + def test_reuse_aware_max_utilization(self): + """C++ ref: ReuseAwareSchedulingMaxUtilizationPolicy""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + tokens = list(range(20)) + r0 = make_context_request(0, prompt_len=20) + r0._input_tokens = tokens + r1 = make_context_request(1, prompt_len=20) + r1._input_tokens = tokens + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + def test_max_utilization_reuse_reduces_needed_blocks(self): + """C++ ref: MaxUtilizationReuseReducesNeededBlocksOneStep""" + kv = MockKVCacheManager( + num_free_blocks=6, blocks_per_request=3, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=3, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + tokens = list(range(30)) + r0 = make_context_request(0, prompt_len=30) + r0._input_tokens = tokens + r1 = make_context_request(1, prompt_len=30) + r1._input_tokens = tokens + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + def test_max_utilization_under_memory_pressure_with_reuse(self): + """C++ ref: MaxUtilizationUnderMemoryPressureWithReuse""" + kv = MockKVCacheManager( + num_free_blocks=5, blocks_per_request=2, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + tokens = list(range(20)) + r0 = make_context_request(0, prompt_len=20) + r0._input_tokens = tokens + r1 = make_context_request(1, prompt_len=20) + r1._input_tokens = tokens + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + def test_max_utilization_incremental_reuse(self): + """C++ ref: MaxUtilizationMultipleRequestsIncrementalReuse""" + kv = MockKVCacheManager( + num_free_blocks=15, blocks_per_request=2, enable_block_reuse=True + ) + scheduler = PyCapacityScheduler( + max_num_requests=4, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + tokens = list(range(20)) + r0 = make_context_request(0, prompt_len=20) + r0._input_tokens = tokens + r1 = make_context_request(1, prompt_len=20) + r1._input_tokens = tokens + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(fitting) >= 1 + + def test_no_reuse_when_disabled(self): + """C++ ref: MaxUtilizationNoReuseWhenDisabled""" + kv = MockKVCacheManager( + num_free_blocks=100, blocks_per_request=2, enable_block_reuse=False + ) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + tokens = list(range(20)) + r0 = make_context_request(0, prompt_len=20) + r0._input_tokens = tokens + r1 = make_context_request(1, prompt_len=20) + r1._input_tokens = tokens + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + # No reuse: both scheduled together + assert len(fitting) == 2 + + +class TestPyCapacitySchedulerDisaggAdvanced: + """ + Advanced tests for disaggregated generation init scheduling. + C++ ref: capacitySchedulerTest.cpp disagg tests. + """ + + def test_disagg_gen_init_max_utilization(self): + """C++ ref: DisaggGenInitMaxUtilization""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + r0 = make_disagg_gen_init_request(0) + r1 = make_disagg_gen_init_request(1) + fitting, disagg, paused = scheduler.schedule_request([r0, r1]) + assert len(disagg) == 2 + + +class TestPyCapacitySchedulerStaticBatchAdvanced: + """ + Advanced tests for static batch scheduling. + C++ ref: capacitySchedulerTest.cpp static batch tests. + """ + + def test_static_batch_fits(self): + """C++ ref: SimpleFitsStaticBatch - third request waits for batch to complete""" + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5) + scheduler = PyCapacityScheduler( + max_num_requests=2, + kv_cache_manager=kv, + scheduler_policy=CapacitySchedulerPolicy.STATIC_BATCH, + ) + # When gen is active, no new context allowed + r0 = make_generation_request(0) + r1 = make_context_request(1) + r2 = make_context_request(2) + fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2]) + assert len(fitting) == 1 + assert fitting[0].request_id == 0 From c36502e90f7b14dec383a355beadc9c511a20d48 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Tue, 3 Mar 2026 18:14:49 -0800 Subject: [PATCH 2/7] add IT Signed-off-by: Lanyu Liao --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 4 + .../defs/accuracy/test_llm_api_pytorch.py | 56 ++++- .../test_lists/test-db/l0_b200.yml | 8 + .../test_lists/test-db/l0_dgx_b200.yml | 4 + .../test-db/l0_gb200_multi_gpus.yml | 4 + .../_torch/executor/test_overlap_scheduler.py | 15 +- .../_torch/executor/test_py_scheduler.py | 214 +++++++++++------- 8 files changed, 215 insertions(+), 92 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 722b4a54de33..3d5e954c0a66 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1103,7 +1103,7 @@ def create_py_executor_instance( if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: scheduler_capacity += 1 - use_python_scheduler = os.getenv("TLLM_USE_PYTHON_SCHEDULER", "0") == "1" + use_python_scheduler = scheduler_config.use_python_scheduler if scheduler_config is not None else False if use_python_scheduler and not isinstance(kv_cache_manager, KVCacheManagerV2): scheduler = SimpleUnifiedScheduler( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 68cd9e0ef3ca..9e4c49af31b8 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1626,6 +1626,10 @@ class SchedulerConfig(StrictBaseModel, PybindMirror): default=WaitingQueuePolicy.FCFS, description="The waiting queue scheduling policy") + use_python_scheduler: bool = Field( + default=False, + description="Use pure-Python scheduler instead of C++ scheduler.") + def _to_pybind(self): return _SchedulerConfig( capacity_scheduler_policy=self.capacity_scheduler_policy._to_pybind( diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index e1870edcc2a5..547c111727d1 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -57,7 +57,8 @@ def patched_start_mpi_pool(self): AutoDecodingConfig, CudaGraphConfig, DeepSeekSparseAttentionConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, RocketSparseAttentionConfig, - SamplingParams, SkipSoftmaxAttentionConfig, TorchCompileConfig) + SamplingParams, SchedulerConfig, SkipSoftmaxAttentionConfig, + TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -1432,6 +1433,34 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) + @parametrize_with_ids("mtp_nextn", [0, 2]) + @parametrize_with_ids("attention_dp,cuda_graph,overlap_scheduler", + [(False, False, False), (True, True, True)]) + @parametrize_with_ids("enable_chunked_prefill", [False, True]) + def test_bfloat16_python_scheduler(self, mtp_nextn, attention_dp, + cuda_graph, overlap_scheduler, + enable_chunked_prefill): + scheduler_config = SchedulerConfig(use_python_scheduler=True) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) + pytorch_config = dict( + disable_overlap_scheduler=not overlap_scheduler, + cuda_graph_config=CudaGraphConfig() if cuda_graph else None, + ) + mtp_config = None + if mtp_nextn > 0: + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) + with LLM(self.MODEL_PATH, + kv_cache_config=kv_cache_config, + scheduler_config=scheduler_config, + enable_chunked_prefill=enable_chunked_prefill, + max_num_tokens=256 if enable_chunked_prefill else 8192, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) def test_bfloat16_2_model_mtp(self): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.3) @@ -1494,6 +1523,31 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip_less_device(4) + @parametrize_with_ids("mtp_nextn", + [0, pytest.param(2, marks=skip_pre_hopper)]) + @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(4, 1, 1), (4, 1, 4)], + ids=["tp4", "ep4"]) + def test_bfloat16_4gpus_python_scheduler(self, tp_size, pp_size, ep_size, + mtp_nextn): + scheduler_config = SchedulerConfig(use_python_scheduler=True) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) + pytorch_config = dict(cuda_graph_config=CudaGraphConfig(), ) + mtp_config = None + if mtp_nextn > 0: + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) + with LLM(self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + scheduler_config=scheduler_config, + **pytorch_config, + enable_attention_dp=True, + speculative_config=mtp_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @skip_pre_hopper @parametrize_with_ids("torch_compile", [False, True]) @parametrize_with_ids("fp8kv,attention_dp,cuda_graph,overlap_scheduler", diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f9c4b13fd806..64e622b50241 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -24,6 +24,14 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=True] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 7220d25dccb7..f7ef10964ca6 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -50,6 +50,10 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_nixl[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[tp4-mtp_nextn=0] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[tp4-mtp_nextn=2] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[ep4-mtp_nextn=0] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[ep4-mtp_nextn=2] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] TIMEOUT (60) # ------------- VisualGen multi-GPU tests --------------- - unittest/_torch/visual_gen/multi_gpu diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 9fb818aec1a4..38f24a571ec8 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -25,6 +25,10 @@ l0_gb200_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[tp4-mtp_nextn=0] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[tp4-mtp_nextn=2] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[ep4-mtp_nextn=0] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[ep4-mtp_nextn=2] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] diff --git a/tests/unittest/_torch/executor/test_overlap_scheduler.py b/tests/unittest/_torch/executor/test_overlap_scheduler.py index e9d21dfe066e..8cd239052bb5 100644 --- a/tests/unittest/_torch/executor/test_overlap_scheduler.py +++ b/tests/unittest/_torch/executor/test_overlap_scheduler.py @@ -7,6 +7,7 @@ from tensorrt_llm import LLM, SamplingParams from tensorrt_llm.llmapi import CudaGraphConfig from tensorrt_llm.llmapi import KvCacheConfig as TRT_KvCacheConfig +from tensorrt_llm.llmapi.llm_args import SchedulerConfig # A test case of mmlu_llama from lm_eval @@ -24,7 +25,7 @@ def model_path(): def create_llm(model_dir, disable_overlap_scheduler, sampler_type, - env_overrides=None): + scheduler_config=None): """Create LLM with specific overlap scheduler setting""" pytorch_config = dict(disable_overlap_scheduler=disable_overlap_scheduler, sampler_type=sampler_type) @@ -41,7 +42,7 @@ def create_llm(model_dir, kv_cache_config=trt_kv_cache_config, max_num_tokens= 128, # Only one request longer than max_num_tokens is required to test chunked prefill - env_overrides=env_overrides, + scheduler_config=scheduler_config, ) @@ -52,10 +53,8 @@ def create_llm(model_dir, @pytest.mark.mpi_ray_parity def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, use_python_scheduler): - # Use env_overrides to pass env var to MPI subprocess - env_overrides = { - "TLLM_USE_PYTHON_SCHEDULER": "1" - } if use_python_scheduler else {} + scheduler_config = SchedulerConfig( + use_python_scheduler=True) if use_python_scheduler else None # Test configuration prompts = test_case["prompts"] @@ -75,7 +74,7 @@ def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, with create_llm(model_path, disable_overlap_scheduler=False, sampler_type=sampler_type, - env_overrides=env_overrides) as llm: + scheduler_config=scheduler_config) as llm: outputs_with_overlap = llm.generate(prompts, sampling_params=sampling_config, use_tqdm=True) @@ -87,7 +86,7 @@ def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, with create_llm(model_path, disable_overlap_scheduler=True, sampler_type=sampler_type, - env_overrides=env_overrides) as llm: + scheduler_config=scheduler_config) as llm: outputs_without_overlap = llm.generate(prompts, sampling_params=sampling_config, use_tqdm=True) diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index 633e54a6d26b..df1771e9ff57 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -39,6 +39,7 @@ # Try to import the real scheduler classes. If the bindings are not available # (e.g. no CUDA environment), the entire module is skipped. try: + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ( ChunkingPolicy, ContextChunkingConfig, @@ -46,7 +47,6 @@ PyMicroBatchScheduler, SimpleUnifiedScheduler, ) - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy SCHEDULER_AVAILABLE = True @@ -79,6 +79,7 @@ class MockLlmRequestState(IntEnum): GENERATION_COMPLETE = 20 DISAGG_GENERATION_INIT = 8 + pytestmark = pytest.mark.skipif( not SCHEDULER_AVAILABLE, reason="tensorrt_llm scheduler not importable (missing bindings or CUDA)", @@ -289,17 +290,13 @@ def __init__( def get_kv_cache_stats(self) -> MockKVCacheStats: return MockKVCacheStats( - num_free_blocks_per_window_size={ - ws: self._num_free_blocks for ws in self._window_sizes - } + num_free_blocks_per_window_size={ws: self._num_free_blocks for ws in self._window_sizes} ) def get_remaining_blocks_to_completion(self, req, window_size: int) -> int: return self._blocks_per_request - def get_needed_blocks_one_step( - self, req, two_step_lookahead: bool, window_size: int - ) -> int: + def get_needed_blocks_one_step(self, req, two_step_lookahead: bool, window_size: int) -> int: return self._blocks_per_request def scheduling_has_free_blocks(self, total: int, window_size: int) -> bool: @@ -562,6 +559,104 @@ def test_encoder_request(self): assert len(ctx) == 1 assert ctx[0].request_id == 0 + def test_simple_no_overlap(self): + """ + With max_batch_size=2 and 4 context requests, only 2 are scheduled. + After transitioning to generation, 2 gen requests fill the batch. + C++ ref: SimpleNoOverlap (multi-iteration; here we test single-step + scheduling decisions that compose the same behavior). + """ + scheduler = PyMicroBatchScheduler(max_batch_size=2, max_num_tokens=None) + + # Step 1: 4 context requests, only 2 fit + requests = [ + make_context_request(0, prompt_len=10), + make_context_request(1, prompt_len=10), + make_context_request(2, prompt_len=10), + make_context_request(3, prompt_len=10), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + assert len(gen) == 0 + assert ctx[0].request_id == 0 + assert ctx[1].request_id == 1 + + # Step 2: first 2 become generation, remaining 2 still context + # Generation fills the batch, context requests wait + requests = [ + make_generation_request(0), + make_generation_request(1), + make_context_request(2, prompt_len=10), + make_context_request(3, prompt_len=10), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(gen) == 2 + assert gen[0].request_id == 0 + assert gen[1].request_id == 1 + # Context requests 2,3 don't fit due to batch_size=2 + assert len(ctx) == 0 + + # Step 3: first 2 complete, now context 2,3 get scheduled + requests = [ + make_context_request(2, prompt_len=10), + make_context_request(3, prompt_len=10), + ] + ctx, gen = scheduler.schedule(requests, set()) + assert len(ctx) == 2 + assert ctx[0].request_id == 2 + assert ctx[1].request_id == 3 + + def test_simple_no_overlap_max_num_tokens(self): + """ + Context chunking with token budget enforcement across multiple steps. + C++ ref: SimpleNoOverlapMaxNumTokens + Req 0, 1: promptLen=12, maxNewTokens=5, maxNumTokens=7, chunkUnitSize=5 + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=2, max_num_tokens=7, ctx_chunk_config=config + ) + + # Step 1 (it=0): Only req0 gets a chunk of 5, req1 doesn't fit + # C++: Req 0: (0,1,2,3,4), Req 1: () + r0 = make_context_request(0, prompt_len=12) + r1 = make_context_request(1, prompt_len=12) + ctx, gen = scheduler.schedule([r0, r1], set()) + assert len(ctx) >= 1 + # First request gets a chunk within budget + req0 = next(r for r in ctx if r.request_id == 0) + assert req0.context_chunk_size <= 7 + total_tokens = sum(r.context_chunk_size for r in ctx) + assert total_tokens <= 7 + + def test_simple_no_overlap_max_context_length(self): + """ + Context chunking with max_context_length limiting chunk sizes. + C++ ref: SimpleNoOverlapMaxContextLength + Requests with promptLen=10 and 17, maxContextLength=12, chunkUnitSize=5. + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=2, max_num_tokens=None, ctx_chunk_config=config + ) + # Override max_context_length (in C++ this is a separate constructor arg) + scheduler.max_context_length = 12 + + # Two requests with promptLen=10 fit within maxContextLength=12 + r0 = make_context_request(0, prompt_len=10) + r1 = make_context_request(1, prompt_len=10) + ctx, gen = scheduler.schedule([r0, r1], set()) + assert len(ctx) == 2 + # Each chunk should be at most max_context_length + for r in ctx: + assert r.context_chunk_size <= 12 + + # Request with promptLen=17 needs chunking (17 > 12) + r3 = make_context_request(3, prompt_len=17) + ctx2, gen2 = scheduler.schedule([r3], set()) + assert len(ctx2) == 1 + assert ctx2[0].context_chunk_size <= 12 + # ############################################################################ # @@ -608,7 +703,7 @@ def test_equal_progress_uneven_remaining(self): max_batch_size=4, max_num_tokens=15, ctx_chunk_config=config ) requests = [ - make_context_request(0, prompt_len=3), # Only 3 tokens remaining + make_context_request(0, prompt_len=3), # Only 3 tokens remaining make_context_request(1, prompt_len=20), # Lots remaining ] ctx, gen = scheduler.schedule(requests, set()) @@ -625,9 +720,7 @@ def test_fcfs_basic(self): """ FIRST_COME_FIRST_SERVED: first request gets as much as possible. """ - config = ContextChunkingConfig( - ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5 - ) + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5) scheduler = PyMicroBatchScheduler( max_batch_size=4, max_num_tokens=12, ctx_chunk_config=config ) @@ -644,9 +737,7 @@ def test_fcfs_basic(self): def test_fcfs_fills_first_request(self): """FCFS fills the first request completely if budget allows. After chunking, sort puts not-last-chunk requests first.""" - config = ContextChunkingConfig( - ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5 - ) + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5) scheduler = PyMicroBatchScheduler( max_batch_size=4, max_num_tokens=25, ctx_chunk_config=config ) @@ -774,18 +865,13 @@ def test_chunked_context_draft_tokens_max_num_tokens(self): 2041 = 31*64 + 57, so remainder in unit = 64-57 = 7. Each request's draft reduced from 8 to 7. """ - config = ContextChunkingConfig( - ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=64 - ) + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=64) scheduler = PyMicroBatchScheduler( max_batch_size=64, max_num_tokens=8192, ctx_chunk_config=config, ) - requests = [ - make_context_request(i, prompt_len=2041, draft_tokens_len=8) - for i in range(4) - ] + requests = [make_context_request(i, prompt_len=2041, draft_tokens_len=8) for i in range(4)] ctx, gen = scheduler.schedule(requests, set()) assert len(ctx) == 4 for req in ctx: @@ -801,9 +887,7 @@ def test_chunked_context_draft_tokens_max_context_length(self): But maxContextLength-chunk_size = 10-6 = 4, so remaining_space=4. Draft reduced from 5 to 4. """ - config = ContextChunkingConfig( - ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=64 - ) + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=64) scheduler = PyMicroBatchScheduler( max_batch_size=64, max_num_tokens=8192, @@ -971,15 +1055,11 @@ def test_no_limit(self): def test_context_length_never_satisfied(self): """C++ ref: ContextLengthNeverSatisfied""" - _run_context_chunking_test( - [25, 25], 100, [[0, 0]], [[0, 0]], max_context_length=20 - ) + _run_context_chunking_test([25, 25], 100, [[0, 0]], [[0, 0]], max_context_length=20) def test_chunk_longer_than_context(self): """C++ ref: ChunkLongerThanContext""" - _run_context_chunking_test( - [25, 25], 30, [[25, 25]], [[25, 25]], max_context_length=25 - ) + _run_context_chunking_test([25, 25], 30, [[25, 25]], [[25, 25]], max_context_length=25) def test_context_length_satisfied(self): """C++ ref: ContextLengthSatisfied""" @@ -1003,9 +1083,7 @@ def test_token_capacity_smaller_than_context(self): def test_token_capacity_smaller_than_chunk_unit(self): """C++ ref: TokenCapacitySmallerThanChunkUnit""" - _run_context_chunking_test( - [25, 25], 20, [[0, 0]], [[0, 0]], ctx_tokens_capacity=10 - ) + _run_context_chunking_test([25, 25], 20, [[0, 0]], [[0, 0]], ctx_tokens_capacity=10) def test_scheduling_order(self): """C++ ref: SchedulingOrder""" @@ -1211,9 +1289,7 @@ def test_draft_tokens_greater_than_chunk_size(self): - Request 1: draftTokens = 13 - Request 2: draftTokens = 5 (remaining budget) """ - config = ContextChunkingConfig( - ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=16 - ) + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=16) scheduler = PyMicroBatchScheduler( max_batch_size=64, max_num_tokens=40, @@ -1348,8 +1424,8 @@ def test_generation_scheduled_first(self): ) requests = [ make_generation_request(0), # takes 5 blocks, leaves 7 - make_context_request(1), # takes 5 blocks, leaves 2 - make_context_request(2), # needs 5, only 2 left + make_context_request(1), # takes 5 blocks, leaves 2 + make_context_request(2), # needs 5, only 2 left ] fitting, disagg, paused = scheduler.schedule_request(requests) assert len(fitting) == 2 @@ -1535,14 +1611,13 @@ def test_generation_complete_filtered(self): max_num_requests=4, kv_cache_manager=None, ) - fitting, disagg, paused = scheduler.schedule_request( - [make_completed_request(0)] - ) + fitting, disagg, paused = scheduler.schedule_request([make_completed_request(0)]) assert len(fitting) == 0 - def test_generation_to_complete_included(self): - """GENERATION_TO_COMPLETE is within schedulable range and handled - like GENERATION_IN_PROGRESS by all capacity scheduler policies.""" + def test_generation_to_complete_filtered(self): + """GENERATION_TO_COMPLETE prohibits re-scheduling per C++ design. + C++ ref: capacityScheduler.cpp — this state is not included in + scheduling conditions for either MaxRequests or GuaranteedNoEvict.""" scheduler = PyCapacityScheduler( max_num_requests=4, kv_cache_manager=None, @@ -1552,8 +1627,7 @@ def test_generation_to_complete_included(self): _state=MockLlmRequestState.GENERATION_TO_COMPLETE, ) fitting, disagg, paused = scheduler.schedule_request([req]) - # GENERATION_TO_COMPLETE passes state gating and is handled by MaxRequestsPolicy - assert len(fitting) == 1 + assert len(fitting) == 0 # ############################################################################ @@ -1987,9 +2061,7 @@ class TestPyCapacitySchedulerKVCacheReuse: def test_delay_duplicate_request(self): """C++ ref: DelayDuplicateRequest - identical requests delayed for reuse""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=3, kv_cache_manager=kv, @@ -2008,9 +2080,7 @@ def test_delay_duplicate_request(self): def test_delay_duplicate_request_chunked(self): """C++ ref: DelayDuplicateRequestChunked""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=5, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=5, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=2, kv_cache_manager=kv, @@ -2028,9 +2098,7 @@ def test_delay_duplicate_request_chunked(self): def test_delay_five_requests_complicated(self): """C++ ref: DelayFiveRequestsComplicated""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=5, kv_cache_manager=kv, @@ -2046,16 +2114,12 @@ def test_delay_five_requests_complicated(self): r3._input_tokens = list(range(21)) # first 11 shared r4 = make_context_request(4, prompt_len=31) r4._input_tokens = list(range(31)) # first 11 shared - fitting, disagg, paused = scheduler.schedule_request( - [r0, r1, r2, r3, r4] - ) + fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2, r3, r4]) assert len(fitting) >= 1 def test_reuse_aware_allows_more_requests(self): """C++ ref: ReuseAwareSchedulingAllowsMoreRequestsWithSharedPrefix""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=4, kv_cache_manager=kv, @@ -2071,9 +2135,7 @@ def test_reuse_aware_allows_more_requests(self): def test_reuse_aware_partial_prefix_match(self): """C++ ref: ReuseAwareSchedulingWithPartialPrefixMatch""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=3, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=3, kv_cache_manager=kv, @@ -2090,9 +2152,7 @@ def test_reuse_aware_partial_prefix_match(self): def test_no_reuse_with_different_prompts(self): """C++ ref: NoReuseWithDifferentPrompts""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=2, kv_cache_manager=kv, @@ -2108,9 +2168,7 @@ def test_no_reuse_with_different_prompts(self): def test_reuse_aware_max_utilization(self): """C++ ref: ReuseAwareSchedulingMaxUtilizationPolicy""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=2, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=4, kv_cache_manager=kv, @@ -2126,9 +2184,7 @@ def test_reuse_aware_max_utilization(self): def test_max_utilization_reuse_reduces_needed_blocks(self): """C++ ref: MaxUtilizationReuseReducesNeededBlocksOneStep""" - kv = MockKVCacheManager( - num_free_blocks=6, blocks_per_request=3, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=6, blocks_per_request=3, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=3, kv_cache_manager=kv, @@ -2144,9 +2200,7 @@ def test_max_utilization_reuse_reduces_needed_blocks(self): def test_max_utilization_under_memory_pressure_with_reuse(self): """C++ ref: MaxUtilizationUnderMemoryPressureWithReuse""" - kv = MockKVCacheManager( - num_free_blocks=5, blocks_per_request=2, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=5, blocks_per_request=2, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=2, kv_cache_manager=kv, @@ -2162,9 +2216,7 @@ def test_max_utilization_under_memory_pressure_with_reuse(self): def test_max_utilization_incremental_reuse(self): """C++ ref: MaxUtilizationMultipleRequestsIncrementalReuse""" - kv = MockKVCacheManager( - num_free_blocks=15, blocks_per_request=2, enable_block_reuse=True - ) + kv = MockKVCacheManager(num_free_blocks=15, blocks_per_request=2, enable_block_reuse=True) scheduler = PyCapacityScheduler( max_num_requests=4, kv_cache_manager=kv, @@ -2180,9 +2232,7 @@ def test_max_utilization_incremental_reuse(self): def test_no_reuse_when_disabled(self): """C++ ref: MaxUtilizationNoReuseWhenDisabled""" - kv = MockKVCacheManager( - num_free_blocks=100, blocks_per_request=2, enable_block_reuse=False - ) + kv = MockKVCacheManager(num_free_blocks=100, blocks_per_request=2, enable_block_reuse=False) scheduler = PyCapacityScheduler( max_num_requests=2, kv_cache_manager=kv, From bd48da521691f50ed7fa77e3fff7dd286c916672 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Tue, 3 Mar 2026 19:34:36 -0800 Subject: [PATCH 3/7] fix tests Signed-off-by: Lanyu Liao --- .../nanobind/batch_manager/bindings.cpp | 1 + .../_torch/pyexecutor/scheduler/scheduler.py | 3 +- .../_torch/executor/test_py_scheduler.py | 399 +++++------------- 3 files changed, 117 insertions(+), 286 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 044c9461dd2c..5e87afc8439c 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -153,6 +153,7 @@ void initBindings(nb::module_& m) nb::arg("num_tokens_per_iteration"), nb::arg("model_config")) .def_prop_ro("orig_prompt_len", &GenLlmReq::getOrigPromptLen) .def("has_draft_tokens", &GenLlmReq::hasDraftTokens) + .def("discard_draft_tokens", &GenLlmReq::discardDraftTokens, nb::arg("num_tokens_to_discard")) .def("move_to_next_context_chunk", &GenLlmReq::moveToNextContextChunk) .def_prop_ro("is_last_context_chunk", &GenLlmReq::isLastContextChunk) .def_prop_ro("is_first_context_chunk", &GenLlmReq::isFirstContextChunk) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 52564264bdb3..e6b4e375a39c 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -655,8 +655,7 @@ def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit draft_discard = req.num_draft_tokens - remaining_space if draft_discard > 0: logger.debug(f"Discarding {draft_discard} draft tokens") - if hasattr(req, "discard_draft_tokens"): - req.discard_draft_tokens(draft_discard) + req.discard_draft_tokens(draft_discard) class SchedulerPolicyBase(ABC): diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index df1771e9ff57..4e9ff3732506 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -16,246 +16,113 @@ Unit tests for Python scheduler implementations (PyMicroBatchScheduler, PyCapacityScheduler, SimpleUnifiedScheduler). -These tests validate the pure-Python scheduler logic using mock objects, -without requiring GPU or real KV cache managers. They are aligned with -the C++ scheduler unit tests in: +These tests validate the pure-Python scheduler logic using real LlmRequest +objects (from C++ bindings) and mock KV cache managers, without requiring +GPU. They are aligned with the C++ scheduler unit tests in: - cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp - cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp """ from dataclasses import dataclass, field -from enum import IntEnum from typing import List, Optional import pytest -# --------------------------------------------------------------------------- -# Mock LlmRequestState — mirrors the C++ enum values exactly. -# We define our own so these tests don't depend on the C++ bindings at import -# time, making them runnable in environments where bindings aren't built. -# --------------------------------------------------------------------------- - - -# Try to import the real scheduler classes. If the bindings are not available -# (e.g. no CUDA environment), the entire module is skipped. -try: - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState - from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ( - ChunkingPolicy, - ContextChunkingConfig, - PyCapacityScheduler, - PyMicroBatchScheduler, - SimpleUnifiedScheduler, - ) - from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy - - SCHEDULER_AVAILABLE = True -except ImportError: - SCHEDULER_AVAILABLE = False - - -# MockLlmRequestState uses the REAL enum values from C++ bindings. -# These values are NOT sequential (e.g., CONTEXT_INIT=10, not 2). -# We derive them from the real LlmRequestState when available. -if SCHEDULER_AVAILABLE: - - class MockLlmRequestState(IntEnum): - UNKNOWN = LlmRequestState.UNKNOWN.value - ENCODER_INIT = LlmRequestState.ENCODER_INIT.value - CONTEXT_INIT = LlmRequestState.CONTEXT_INIT.value - GENERATION_IN_PROGRESS = LlmRequestState.GENERATION_IN_PROGRESS.value - GENERATION_TO_COMPLETE = LlmRequestState.GENERATION_TO_COMPLETE.value - GENERATION_COMPLETE = LlmRequestState.GENERATION_COMPLETE.value - DISAGG_GENERATION_INIT = LlmRequestState.DISAGG_GENERATION_INIT.value - -else: - - class MockLlmRequestState(IntEnum): - UNKNOWN = 0 - ENCODER_INIT = 1 - CONTEXT_INIT = 10 - GENERATION_IN_PROGRESS = 13 - GENERATION_TO_COMPLETE = 14 - GENERATION_COMPLETE = 20 - DISAGG_GENERATION_INIT = 8 - - -pytestmark = pytest.mark.skipif( - not SCHEDULER_AVAILABLE, - reason="tensorrt_llm scheduler not importable (missing bindings or CUDA)", +from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest, + LlmRequestState, + SamplingConfig) +from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ( + ChunkingPolicy, + ContextChunkingConfig, + PyCapacityScheduler, + PyMicroBatchScheduler, + SimpleUnifiedScheduler, ) +from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy -# ============================================================================ -# Mock infrastructure -# ============================================================================ - - -@dataclass -class MockLlmRequest: - """ - Minimal mock of LlmRequest that exposes exactly the attributes/methods - the Python scheduler accesses. Attribute names match the real binding API. - """ - - request_id: int - _state: int = MockLlmRequestState.CONTEXT_INIT - prompt_len: int = 10 - _context_current_position: int = 0 - _context_chunk_size: int = 0 - beam_width: int = 1 - _draft_tokens_len: int = 0 - encoder_output_len: int = 0 - lora_task_id: Optional[int] = None - py_request_id: int = 0 # Used by MaxUtilization's scheduling_remove_sequence - _input_tokens: Optional[List[int]] = None - - def __post_init__(self): - self.py_request_id = self.request_id - - # --- State properties --- - @property - def state(self): - return LlmRequestState(self._state) if SCHEDULER_AVAILABLE else self._state - - @state.setter - def state(self, value): - self._state = int(value) - - @property - def state_value(self) -> int: - return self._state - - @property - def is_context_init_state(self) -> bool: - return self._state == MockLlmRequestState.CONTEXT_INIT - - @property - def is_generation_in_progress_state(self) -> bool: - return self._state == MockLlmRequestState.GENERATION_IN_PROGRESS - - @property - def is_encoder_init_state(self) -> bool: - return self._state == MockLlmRequestState.ENCODER_INIT - - @property - def is_disagg_generation_init_state(self) -> bool: - return self._state == MockLlmRequestState.DISAGG_GENERATION_INIT - - @property - def is_generation_to_complete_state(self) -> bool: - return self._state == MockLlmRequestState.GENERATION_TO_COMPLETE - - # --- Context chunking properties --- - @property - def context_remaining_length(self) -> int: - return self.prompt_len - self._context_current_position - - @property - def context_chunk_size(self) -> int: - return self._context_chunk_size - - @context_chunk_size.setter - def context_chunk_size(self, value: int): - # Clamp to context_remaining_length, matching C++ behavior - self._context_chunk_size = min(value, self.context_remaining_length) - - @property - def is_first_context_chunk(self) -> bool: - return self._context_current_position == 0 - - @property - def is_last_context_chunk(self) -> bool: - return self._context_chunk_size >= self.context_remaining_length - - # --- Token / draft properties --- - def get_num_tokens(self, beam: int) -> int: - """For context requests, returns prompt_len (matching C++ getNumTokens).""" - return self.prompt_len - - @property - def num_draft_tokens(self) -> int: - return self._draft_tokens_len - - @property - def has_draft_tokens(self) -> bool: - return self._draft_tokens_len > 0 - - @property - def draft_tokens(self) -> list: - return [0] * self._draft_tokens_len - - def get_beam_width_by_iter(self, for_next_iteration: bool = False) -> int: - return self.beam_width - - def discard_draft_tokens(self, n: int): - self._draft_tokens_len = max(0, self._draft_tokens_len - n) - - # --- Methods used by capacity scheduler --- - def get_unique_tokens(self, beam: int): - return self._input_tokens - - def get_encoder_unique_tokens(self): - return None - def move_to_next_context_chunk(self): - """Simulate moving to the next context chunk.""" - self._context_current_position += self._context_chunk_size - self._context_chunk_size = 0 +def _make_request( + request_id: int, + prompt_len: int = 10, + beam_width: int = 1, + draft_tokens_len: int = 0, + encoder_output_len: int = 0, + lora_task_id: Optional[int] = None, + state: LlmRequestState = LlmRequestState.CONTEXT_INIT, + input_tokens: Optional[List[int]] = None, +) -> LlmRequest: + tokens = input_tokens if input_tokens is not None else list( + range(prompt_len)) + draft = list(range(draft_tokens_len)) if draft_tokens_len > 0 else None + req = LlmRequest( + request_id=request_id, + max_new_tokens=10, + input_tokens=tokens, + sampling_config=SamplingConfig(beam_width), + is_streaming=False, + draft_tokens=draft, + lora_task_id=lora_task_id, + encoder_output_len=encoder_output_len if encoder_output_len > 0 else None, + ) + req.state = state + return req -# Helper to create requests in specific states def make_context_request( request_id: int, prompt_len: int = 10, beam_width: int = 1, draft_tokens_len: int = 0, context_position: int = 0, -) -> MockLlmRequest: - return MockLlmRequest( +) -> LlmRequest: + req = _make_request( request_id=request_id, - _state=MockLlmRequestState.CONTEXT_INIT, prompt_len=prompt_len, - _context_current_position=context_position, beam_width=beam_width, - _draft_tokens_len=draft_tokens_len, + draft_tokens_len=draft_tokens_len, + state=LlmRequestState.CONTEXT_INIT, ) + if context_position > 0: + req.context_chunk_size = context_position + req.move_to_next_context_chunk() + return req def make_generation_request( request_id: int, beam_width: int = 1, draft_tokens_len: int = 0, -) -> MockLlmRequest: - return MockLlmRequest( +) -> LlmRequest: + return _make_request( request_id=request_id, - _state=MockLlmRequestState.GENERATION_IN_PROGRESS, beam_width=beam_width, - _draft_tokens_len=draft_tokens_len, + draft_tokens_len=draft_tokens_len, + state=LlmRequestState.GENERATION_IN_PROGRESS, ) def make_encoder_request( request_id: int, encoder_output_len: int = 10, -) -> MockLlmRequest: - return MockLlmRequest( +) -> LlmRequest: + return _make_request( request_id=request_id, - _state=MockLlmRequestState.ENCODER_INIT, encoder_output_len=encoder_output_len, + state=LlmRequestState.ENCODER_INIT, ) -def make_disagg_gen_init_request(request_id: int) -> MockLlmRequest: - return MockLlmRequest( +def make_disagg_gen_init_request(request_id: int) -> LlmRequest: + return _make_request( request_id=request_id, - _state=MockLlmRequestState.DISAGG_GENERATION_INIT, + state=LlmRequestState.DISAGG_GENERATION_INIT, ) -def make_completed_request(request_id: int) -> MockLlmRequest: - return MockLlmRequest( +def make_completed_request(request_id: int) -> LlmRequest: + return _make_request( request_id=request_id, - _state=MockLlmRequestState.GENERATION_COMPLETE, + state=LlmRequestState.GENERATION_COMPLETE, ) @@ -935,12 +802,9 @@ def test_sort_by_lora_task_id(self): C++ ref: sortRequests in inflightBatchingUtils.cpp """ scheduler = PyMicroBatchScheduler(max_batch_size=4, max_num_tokens=None) - r0 = make_generation_request(0) - r0.lora_task_id = 5 - r1 = make_generation_request(1) - r1.lora_task_id = None - r2 = make_generation_request(2) - r2.lora_task_id = 3 + r0 = _make_request(0, state=LlmRequestState.GENERATION_IN_PROGRESS, lora_task_id=5) + r1 = _make_request(1, state=LlmRequestState.GENERATION_IN_PROGRESS) + r2 = _make_request(2, state=LlmRequestState.GENERATION_IN_PROGRESS, lora_task_id=3) ctx, gen = scheduler.schedule([r0, r1, r2], set()) # None < any value, so order should be: r1(None), r2(3), r0(5) assert gen[0].request_id == 1 @@ -970,14 +834,14 @@ def _run_context_chunking_test( Helper that mirrors the C++ ContextChunkingTest fixture. For each policy (EQUAL_PROGRESS and FCFS), it: - 1. Creates MockLlmRequests with given context_lengths and optional draft_lengths. + 1. Creates LlmRequests with given context_lengths and optional draft_lengths. 2. Creates a PyMicroBatchScheduler with the right ContextChunkingConfig. 3. If max_context_length is set, overrides scheduler.max_context_length. 4. For each iteration (each element in positions list): a. Filters requests where context_remaining_length > 0. b. Calls scheduler._set_ctx_requests_chunk_size(active_reqs, ctx_tokens_capacity). c. For each active req, calls req.move_to_next_context_chunk(). - d. Verifies _context_current_position matches expected positions for ALL requests. + d. Verifies context position matches expected positions for ALL requests. 5. After all iterations, verifies final draft_lengths if specified. """ policies_and_data = [ @@ -990,13 +854,10 @@ def _run_context_chunking_test( requests = [] for i, ctx_len in enumerate(context_lengths): dl = draft_lengths[i] if draft_lengths else 0 - req = MockLlmRequest( + req = make_context_request( request_id=i, - _state=MockLlmRequestState.CONTEXT_INIT, prompt_len=ctx_len, - _context_current_position=0, - _context_chunk_size=0, - _draft_tokens_len=dl, + draft_tokens_len=dl, ) requests.append(req) @@ -1024,20 +885,20 @@ def _run_context_chunking_test( # Verify positions for ALL requests (including completed ones) for req_idx, req in enumerate(requests): - assert req._context_current_position == expected_positions[req_idx], ( + assert req.context_current_position == expected_positions[req_idx], ( f"Policy {policy.name}, iteration {iteration_idx}, " f"request {req_idx}: expected position " f"{expected_positions[req_idx]}, got " - f"{req._context_current_position}" + f"{req.context_current_position}" ) # Verify final draft lengths if specified if final_draft_lens is not None: for req_idx, req in enumerate(requests): - assert req._draft_tokens_len == final_draft_lens[req_idx], ( + assert req.num_draft_tokens == final_draft_lens[req_idx], ( f"Policy {policy.name}, request {req_idx}: " f"expected draft_len {final_draft_lens[req_idx]}, " - f"got {req._draft_tokens_len}" + f"got {req.num_draft_tokens}" ) @@ -1601,7 +1462,7 @@ def test_unknown_state_filtered(self): max_num_requests=4, kv_cache_manager=None, ) - req = MockLlmRequest(request_id=0, _state=MockLlmRequestState.UNKNOWN) + req = _make_request(request_id=0, state=LlmRequestState.UNKNOWN) fitting, disagg, paused = scheduler.schedule_request([req]) assert len(fitting) == 0 @@ -1622,9 +1483,9 @@ def test_generation_to_complete_filtered(self): max_num_requests=4, kv_cache_manager=None, ) - req = MockLlmRequest( + req = _make_request( request_id=0, - _state=MockLlmRequestState.GENERATION_TO_COMPLETE, + state=LlmRequestState.GENERATION_TO_COMPLETE, ) fitting, disagg, paused = scheduler.schedule_request([req]) assert len(fitting) == 0 @@ -1653,10 +1514,8 @@ def test_lora_fits(self): peft_cache_manager=peft, scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) - r0 = make_context_request(0) - r0.lora_task_id = 1 - r1 = make_context_request(1) - r1.lora_task_id = 1 # same task — no extra pages needed + r0 = _make_request(0, lora_task_id=1) + r1 = _make_request(1, lora_task_id=1) # same task — no extra pages needed fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) == 2 @@ -1673,10 +1532,8 @@ def test_lora_doesnt_fit(self): peft_cache_manager=peft, scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) - r0 = make_context_request(0) - r0.lora_task_id = 1 - r1 = make_context_request(1) - r1.lora_task_id = 2 # different task — needs 10 more pages + r0 = _make_request(0, lora_task_id=1) + r1 = _make_request(1, lora_task_id=2) # different task — needs 10 more pages fitting, disagg, paused = scheduler.schedule_request([r0, r1]) # First task: 10 pages, second task: 10 pages, total 20 > 15 assert len(fitting) == 1 @@ -1889,9 +1746,9 @@ def test_should_fit_in_chunk(self): scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) r0 = make_context_request(0, prompt_len=50, context_position=0) - r0._context_chunk_size = 20 + r0.context_chunk_size = 20 r1 = make_context_request(1, prompt_len=50, context_position=0) - r1._context_chunk_size = 20 + r1.context_chunk_size = 20 fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) == 2 @@ -1904,9 +1761,9 @@ def test_doesnt_fit_guaranteed_completion_in_chunk(self): scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) r0 = make_context_request(0, prompt_len=30) - r0._context_chunk_size = 20 + r0.context_chunk_size = 20 r1 = make_context_request(1, prompt_len=30) - r1._context_chunk_size = 20 + r1.context_chunk_size = 20 fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) == 1 @@ -1919,9 +1776,9 @@ def test_doesnt_fit_max_utilization_in_chunk(self): scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, ) r0 = make_context_request(0, prompt_len=30) - r0._context_chunk_size = 20 + r0.context_chunk_size = 20 r1 = make_context_request(1, prompt_len=30) - r1._context_chunk_size = 20 + r1.context_chunk_size = 20 fitting, disagg, paused = scheduler.schedule_request([r0, r1]) # At least one should be scheduled assert len(fitting) >= 1 @@ -1935,9 +1792,9 @@ def test_doesnt_fit_max_utilization_in_chunked_cache(self): scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, ) r0 = make_context_request(0, prompt_len=70) - r0._context_chunk_size = 40 + r0.context_chunk_size = 40 r1 = make_context_request(1, prompt_len=70) - r1._context_chunk_size = 40 + r1.context_chunk_size = 40 fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) == 1 @@ -2009,9 +1866,9 @@ def test_adding_new_requests_guaranteed_completion_in_chunk(self): ) r0 = make_generation_request(0) r1 = make_context_request(1, prompt_len=30) - r1._context_chunk_size = 20 + r1.context_chunk_size = 20 r2 = make_context_request(2, prompt_len=30) - r2._context_chunk_size = 20 + r2.context_chunk_size = 20 fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2]) assert len(fitting) >= 1 @@ -2068,12 +1925,9 @@ def test_delay_duplicate_request(self): scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) tokens = list(range(21)) - r0 = make_context_request(0, prompt_len=21) - r0._input_tokens = tokens - r1 = make_context_request(1, prompt_len=21) - r1._input_tokens = tokens - r2 = make_context_request(2, prompt_len=21) - r2._input_tokens = tokens + r0 = _make_request(0, prompt_len=21, input_tokens=tokens) + r1 = _make_request(1, prompt_len=21, input_tokens=tokens) + r2 = _make_request(2, prompt_len=21, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2]) # With reuse enabled, beneficial_to_skip may delay r1 and r2 assert len(fitting) >= 1 @@ -2087,12 +1941,10 @@ def test_delay_duplicate_request_chunked(self): scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) tokens = list(range(50)) - r0 = make_context_request(0, prompt_len=50) - r0._input_tokens = tokens - r0._context_chunk_size = 20 - r1 = make_context_request(1, prompt_len=50) - r1._input_tokens = tokens - r1._context_chunk_size = 20 + r0 = _make_request(0, prompt_len=50, input_tokens=tokens) + r0.context_chunk_size = 20 + r1 = _make_request(1, prompt_len=50, input_tokens=tokens) + r1.context_chunk_size = 20 fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) >= 1 @@ -2104,16 +1956,11 @@ def test_delay_five_requests_complicated(self): kv_cache_manager=kv, scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) - r0 = make_context_request(0, prompt_len=11) - r0._input_tokens = [x + 1 for x in range(11)] # unique - r1 = make_context_request(1, prompt_len=21) - r1._input_tokens = [x + 2 for x in range(21)] # unique - r2 = make_context_request(2, prompt_len=11) - r2._input_tokens = list(range(11)) - r3 = make_context_request(3, prompt_len=21) - r3._input_tokens = list(range(21)) # first 11 shared - r4 = make_context_request(4, prompt_len=31) - r4._input_tokens = list(range(31)) # first 11 shared + r0 = _make_request(0, prompt_len=11, input_tokens=[x + 1 for x in range(11)]) + r1 = _make_request(1, prompt_len=21, input_tokens=[x + 2 for x in range(21)]) + r2 = _make_request(2, prompt_len=11, input_tokens=list(range(11))) + r3 = _make_request(3, prompt_len=21, input_tokens=list(range(21))) + r4 = _make_request(4, prompt_len=31, input_tokens=list(range(31))) fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2, r3, r4]) assert len(fitting) >= 1 @@ -2126,10 +1973,8 @@ def test_reuse_aware_allows_more_requests(self): scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) tokens = list(range(20)) - r0 = make_context_request(0, prompt_len=20) - r0._input_tokens = tokens - r1 = make_context_request(1, prompt_len=20) - r1._input_tokens = tokens + r0 = _make_request(0, prompt_len=20, input_tokens=tokens) + r1 = _make_request(1, prompt_len=20, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) >= 1 @@ -2143,10 +1988,8 @@ def test_reuse_aware_partial_prefix_match(self): ) tokens0 = list(range(30)) tokens1 = list(range(20)) + [999] * 10 - r0 = make_context_request(0, prompt_len=30) - r0._input_tokens = tokens0 - r1 = make_context_request(1, prompt_len=30) - r1._input_tokens = tokens1 + r0 = _make_request(0, prompt_len=30, input_tokens=tokens0) + r1 = _make_request(1, prompt_len=30, input_tokens=tokens1) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) >= 1 @@ -2158,10 +2001,8 @@ def test_no_reuse_with_different_prompts(self): kv_cache_manager=kv, scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, ) - r0 = make_context_request(0, prompt_len=20) - r0._input_tokens = [100] * 20 - r1 = make_context_request(1, prompt_len=20) - r1._input_tokens = [200] * 20 + r0 = _make_request(0, prompt_len=20, input_tokens=[100] * 20) + r1 = _make_request(1, prompt_len=20, input_tokens=[200] * 20) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) # Different prompts: no reuse, both scheduled assert len(fitting) == 2 @@ -2175,10 +2016,8 @@ def test_reuse_aware_max_utilization(self): scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, ) tokens = list(range(20)) - r0 = make_context_request(0, prompt_len=20) - r0._input_tokens = tokens - r1 = make_context_request(1, prompt_len=20) - r1._input_tokens = tokens + r0 = _make_request(0, prompt_len=20, input_tokens=tokens) + r1 = _make_request(1, prompt_len=20, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) >= 1 @@ -2191,10 +2030,8 @@ def test_max_utilization_reuse_reduces_needed_blocks(self): scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, ) tokens = list(range(30)) - r0 = make_context_request(0, prompt_len=30) - r0._input_tokens = tokens - r1 = make_context_request(1, prompt_len=30) - r1._input_tokens = tokens + r0 = _make_request(0, prompt_len=30, input_tokens=tokens) + r1 = _make_request(1, prompt_len=30, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) >= 1 @@ -2207,10 +2044,8 @@ def test_max_utilization_under_memory_pressure_with_reuse(self): scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, ) tokens = list(range(20)) - r0 = make_context_request(0, prompt_len=20) - r0._input_tokens = tokens - r1 = make_context_request(1, prompt_len=20) - r1._input_tokens = tokens + r0 = _make_request(0, prompt_len=20, input_tokens=tokens) + r1 = _make_request(1, prompt_len=20, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) >= 1 @@ -2223,10 +2058,8 @@ def test_max_utilization_incremental_reuse(self): scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, ) tokens = list(range(20)) - r0 = make_context_request(0, prompt_len=20) - r0._input_tokens = tokens - r1 = make_context_request(1, prompt_len=20) - r1._input_tokens = tokens + r0 = _make_request(0, prompt_len=20, input_tokens=tokens) + r1 = _make_request(1, prompt_len=20, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) assert len(fitting) >= 1 @@ -2239,10 +2072,8 @@ def test_no_reuse_when_disabled(self): scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, ) tokens = list(range(20)) - r0 = make_context_request(0, prompt_len=20) - r0._input_tokens = tokens - r1 = make_context_request(1, prompt_len=20) - r1._input_tokens = tokens + r0 = _make_request(0, prompt_len=20, input_tokens=tokens) + r1 = _make_request(1, prompt_len=20, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) # No reuse: both scheduled together assert len(fitting) == 2 From 4077869f4ece4ac372338b5d6c1362f6671a9b10 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Wed, 4 Mar 2026 04:32:34 -0800 Subject: [PATCH 4/7] fix Signed-off-by: Lanyu Liao --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index e69843c36380..ab4dd176d623 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1474,10 +1474,10 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, task.evaluate(llm) @pytest.mark.skip_less_device_memory(60000) - @parametrize_with_ids("mtp_nextn", [0, 2]) + @parametrize_with_ids("enable_chunked_prefill", [False, True]) @parametrize_with_ids("attention_dp,cuda_graph,overlap_scheduler", [(False, False, False), (True, True, True)]) - @parametrize_with_ids("enable_chunked_prefill", [False, True]) + @parametrize_with_ids("mtp_nextn", [0, 2]) def test_bfloat16_python_scheduler(self, mtp_nextn, attention_dp, cuda_graph, overlap_scheduler, enable_chunked_prefill): From e6b904659175f0bcd5def21898d0a3bb462af756 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Wed, 4 Mar 2026 06:18:30 -0800 Subject: [PATCH 5/7] add IT Signed-off-by: Lanyu Liao Made-with: Cursor --- tensorrt_llm/serve/openai_server.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index f52fe64afe2c..296426f965eb 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1656,7 +1656,8 @@ async def openai_video_generation_sync(self, actual_output_path = MediaStorage.save_video( video=output.video, - output_path=self.media_storage_path / f"{video_id}{resolved_ext}", + output_path=self.media_storage_path / + f"{video_id}{resolved_ext}", audio=output.audio, frame_rate=request.fps or params.frame_rate, format=resolved_fmt, @@ -1836,7 +1837,8 @@ async def _generate_video_background( actual_output_path = MediaStorage.save_video( video=output.video, - output_path=self.media_storage_path / f"{video_id}{resolved_ext}", + output_path=self.media_storage_path / + f"{video_id}{resolved_ext}", audio=output.audio, frame_rate=request.fps or params.frame_rate, format=resolved_fmt, From 7112110f45b7912e7d2741faa69256329f2b39b7 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Wed, 4 Mar 2026 18:10:55 -0800 Subject: [PATCH 6/7] fix Signed-off-by: Lanyu Liao --- tests/unittest/_torch/executor/test_overlap_scheduler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/executor/test_overlap_scheduler.py b/tests/unittest/_torch/executor/test_overlap_scheduler.py index 8cd239052bb5..306c5e45101b 100644 --- a/tests/unittest/_torch/executor/test_overlap_scheduler.py +++ b/tests/unittest/_torch/executor/test_overlap_scheduler.py @@ -27,6 +27,8 @@ def create_llm(model_dir, sampler_type, scheduler_config=None): """Create LLM with specific overlap scheduler setting""" + if scheduler_config is None: + scheduler_config = SchedulerConfig() pytorch_config = dict(disable_overlap_scheduler=disable_overlap_scheduler, sampler_type=sampler_type) @@ -54,7 +56,7 @@ def create_llm(model_dir, def test_overlap_scheduler_consistency(model_path, test_case, sampler_type, use_python_scheduler): scheduler_config = SchedulerConfig( - use_python_scheduler=True) if use_python_scheduler else None + use_python_scheduler=use_python_scheduler) # Test configuration prompts = test_case["prompts"] From 35e8348c23500f374811109f21ebbac70f397862 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Thu, 5 Mar 2026 20:59:13 -0800 Subject: [PATCH 7/7] add new tests in llm_function_core.txt Signed-off-by: Lanyu Liao --- .../integration/test_lists/qa/llm_function_core.txt | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 132ae0b3b6be..58a11c67f63d 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -83,6 +83,14 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0- accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-enable_chunked_prefill=False] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=False] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-enable_chunked_prefill=True] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales[mtp=disable-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales[mtp=disable-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] @@ -92,6 +100,10 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUT accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=0.75-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=0.75-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[tp4-mtp_nextn=0] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[tp4-mtp_nextn=2] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[ep4-mtp_nextn=0] +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[ep4-mtp_nextn=2] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus_static_eplb accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=0-moe_backend=WIDEEP] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=2-moe_backend=WIDEEP]