diff --git a/sendnn_inference/v1/core/scheduler.py b/sendnn_inference/v1/core/scheduler.py index 3ec29fe23..52d7510d0 100644 --- a/sendnn_inference/v1/core/scheduler.py +++ b/sendnn_inference/v1/core/scheduler.py @@ -20,6 +20,16 @@ logger = init_logger(__name__) +# Ensure that block_size is 64 +# This ensures the rounding function is correct +assert SpyrePlatform.get_block_size() == 64 + + +def round_up_to_block_size(n: int) -> int: + # Helper function to round up to the nearest block size + # Uses bitwise alignment for better performance + return (n + 63) & ~63 + class SpyreScheduler(Scheduler): """Base class inheriting from the V1 scheduler to support static @@ -480,7 +490,6 @@ def _satisfies_last_chunk_constraints(self, request: Request) -> bool: cond3 = lambda: self.check_batch_tkv_limit_cp( request=request, new_req_tkv=new_req_tkv, - n_blocks=n_blocks, running=decoding_requests, ) @@ -504,9 +513,7 @@ def _has_scheduling_priority(self, request): num_prefills = len(self.waiting) + len(self.ongoing_prefills) return num_prefills < max_concurrent_prefills - def check_batch_tkv_limit_cp( - self, request: Request, new_req_tkv: int, n_blocks: int, running - ) -> bool: + def check_batch_tkv_limit_cp(self, request: Request, new_req_tkv: int, running) -> bool: """ Check whether adding a new sequence to the decode batch would violate Spyre's maximum batch volume constraint for chunked prefill. @@ -531,7 +538,8 @@ def check_batch_tkv_limit_cp( """ # Compute the effective token length of the new request - new_req_max_tkv = new_req_tkv + request.max_tokens - 1 + # Rounded up to the nearest block size to account for potential padding + new_req_max_tkv = round_up_to_block_size(new_req_tkv + request.max_tokens - 1) # Compute token lengths for all running requests (decode batch) decode_req_max_tkvs = [] @@ -539,9 +547,10 @@ def check_batch_tkv_limit_cp( dec_req_tkv = max(self.tkv, request.num_prompt_tokens) for req in running: n_generated_output_tokens = req.num_computed_tokens - req.num_prompt_tokens - dec_req_max_tkv = dec_req_tkv + (req.max_tokens - n_generated_output_tokens) - 1 - # Account for potential padding block - dec_req_max_tkv += self.block_size + # Rounded up to the nearest block size to account for potential padding + dec_req_max_tkv = round_up_to_block_size( + dec_req_tkv + (req.max_tokens - n_generated_output_tokens) - 1 + ) decode_req_max_tkvs.append(dec_req_max_tkv) # Sort decode requests token lengths in ascending order diff --git a/tests/v1/worker/test_scheduler_tkv_limits.py b/tests/v1/worker/test_scheduler_tkv_limits.py index 94638b1fe..7d5f2337f 100644 --- a/tests/v1/worker/test_scheduler_tkv_limits.py +++ b/tests/v1/worker/test_scheduler_tkv_limits.py @@ -80,3 +80,106 @@ def test_scheduler_tkv_limits(monkeypatch: pytest.MonkeyPatch): scheduler.update_from_output(sched_output, output) if len(scheduler.running) == 0: break + + +@pytest.mark.cpu +@pytest.mark.chunked_prefill +def test_scheduler_tkv_limits_ongoing_batch(monkeypatch: pytest.MonkeyPatch): + """ + Test that the scheduler correctly enforces the TKV limit constraint + when new requests are added during an ongoing batch. + + This test schedules a 16x8k batch that will fully fill the 128k TKV limit. + Then inject a batch of smaller requests partway through processing, + which should be able to schedule only because they are guaranteed to + finish processing just before the TKV is long enough to overrun the + limit with the larger batch size. This flexes the logic for injecting + shorter requests into a running batch, which is not tested by the + other test case in this file. + + Expected behavior (when bug is fixed): + - Test should pass without exceeding hardware constraints + + Current behavior (with bug): + - Scheduler accepts invalid batch configurations + - Test will fail with assertion errors + """ + # Setup: Use the default test model + model = REFERENCE_MODELS[InstrumentedModelRunner.DEFAULT_TEST_MODEL] + + # Build model runner with specific constraints + model_runner = InstrumentedModelRunner.build( + monkeypatch=monkeypatch, + max_num_batched_tokens=512, + max_num_seqs=32, + max_model_len=32768, + available_blocks=32768, + ) + + # Configure the TKV limit + scheduler = model_runner.scheduler + scheduler.max_batch_tkv_limit = 131072 + SpyrePlatform._max_batch_tkv_limit = 131072 + monkeypatch.setenv("VLLM_DT_MAX_BATCH_TKV_LIMIT", "131072") + + # Define prompt lengths and max tokens for requests + prompt_lengths = [1018] + [1024] * 15 + max_tokens_1 = 7168 + max_tokens_2 = 900 + + # Create and add first set of requests to the scheduler + requests = [] + for request_id, prompt_length in enumerate(prompt_lengths): + prompt = random_prompt(model=model, seed=request_id, length=prompt_length) + request = create_request_for_scheduler_test( + model=model, + request_id=request_id, + add_step=0, + max_tokens=max_tokens_1, + prompt=prompt, + use_golden_token_injection=False, + generate_hf_results=False, + ).request + requests.append(request) + scheduler.add_request(request) + + # Failure was observed in testing when first request generated 2920 tokens + target_generated_tokens = 2920 + + # Run the scheduler loop until first set of requests have generated tokens + while True: + sched_output = scheduler.schedule() + output = model_runner.execute_model(sched_output) + scheduler.update_from_output(sched_output, output) + + target_req = requests[0] + + if target_req: + generated = target_req.num_computed_tokens - target_req.num_prompt_tokens + if generated >= target_generated_tokens: + break + + # Create and add second set requests to the scheduler + for request_id, prompt_length in enumerate(prompt_lengths): + prompt = random_prompt(model=model, seed=request_id + 16, length=prompt_length) + request = create_request_for_scheduler_test( + model=model, + request_id=request_id + 16, + add_step=0, + max_tokens=max_tokens_2, + prompt=prompt, + use_golden_token_injection=False, + generate_hf_results=False, + ).request + requests.append(request) + scheduler.add_request(request) + + # Run the scheduler loop until all requests complete + # With the bug present, the scheduler will incorrectly accept a batch + # configuration that exceeds the TKV limit + while True: + sched_output = scheduler.schedule() + output = model_runner.execute_model(sched_output) + scheduler.update_from_output(sched_output, output) + if len(scheduler.running) == 0: + break