Tkv limit bug

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,17 +538,19 @@ 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 = []
         # Decide new tkv based on max of current tkv or new request prompt tokens
         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

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