diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index d9af0c883..fff4cc98b 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -704,7 +704,7 @@ class PrefillAdder: for i, running_req in enumerate(self.running_batch.reqs): if running_req in preemptible_reqs: self.rem_total_token_offset -= ( - self._get_running_request_total_token_offset(req) + self._get_running_request_total_token_offset(running_req) ) release_counter += 1 self.running_batch.release_req( diff --git a/test/srt/test_priority_scheduling.py b/test/srt/test_priority_scheduling.py index ef79188f8..923bcb2d3 100644 --- a/test/srt/test_priority_scheduling.py +++ b/test/srt/test_priority_scheduling.py @@ -288,6 +288,86 @@ class TestPrioritySchedulingMultipleRunningRequests(CustomTestCase): _verify_genereate_responses(responses, expected_status_and_error_messages, []) + def test_priority_scheduling_preemption_token_offset_calculation(self): + """ + Verify correct token offset calculation during preemption. + + This test specifically targets the bug where rem_total_token_offset was incorrectly + calculated using the incoming request's tokens instead of the preempted request's tokens + (related to issue #13111 and PR #13201). + + THE BUG: + In schedule_policy.py line 700, the code was using: + self.rem_total_token_offset -= self._get_running_request_total_token_offset(req) + Instead of: + self.rem_total_token_offset -= self._get_running_request_total_token_offset(running_req) + + WHY THIS TEST CATCHES THE BUG: + - Request 1 (preempted): 8000 tokens - This is what SHOULD be freed + - Request 3 (incoming): 1000 tokens - This is what WAS freed (bug) + - Token difference: 8000 - 1000 = 7000 tokens incorrectly accounted + + With the bug, the system thinks it only freed 1000 tokens instead of 8000 tokens. + This causes incorrect memory accounting and can lead to: + 1. Scheduler believes less memory is available than actually is + 2. Subsequent requests (like Request 4) may fail to schedule or cause issues + 3. Memory calculations become increasingly inaccurate with each preemption + + The test creates a scenario where: + 1. A low-priority request with many tokens (8000) starts running + 2. A high-priority request with few tokens (1000) arrives and triggers preemption + 3. The system must correctly free 8000 tokens from the preempted request + 4. Additional requests can be scheduled only if tokens were correctly freed + 5. Execution order validates priority-based scheduling works correctly + + The large token difference (8x) makes the bug's impact obvious and testable. + """ + responses = asyncio.run( + send_concurrent_generate_requests_with_custom_params( + self.base_url, + [ + { + "priority": 0, + "sampling_params": {"max_new_tokens": 8000}, + }, # Low priority, large token count - will be preempted + { + "priority": 1, + "sampling_params": {"max_new_tokens": 5000}, + }, # Medium priority, medium token count - queued initially + { + "priority": 100, + "sampling_params": {"max_new_tokens": 1000}, + }, # High priority, small token count - triggers preemption + { + "priority": 50, + "sampling_params": {"max_new_tokens": 2000}, + }, # Should be schedulable after correct token accounting + ], + ) + ) + + # All requests should complete successfully + # The key is that the fourth request should be schedulable because + # the system correctly freed tokens from the first (preempted) request + expected_status_and_error_messages = [ + (200, None), + (200, None), + (200, None), + (200, None), + ] + + e2e_latencies = [] + _verify_genereate_responses( + responses, expected_status_and_error_messages, e2e_latencies + ) + + # Verify execution order: high priority requests finish before low priority ones + # Request 3 (priority 100) should finish first + # Request 4 (priority 50) should finish second + # Request 2 (priority 1) should finish third + # Request 1 (priority 0) should finish last (after being preempted) + assert e2e_latencies[2] < e2e_latencies[3] < e2e_latencies[1] < e2e_latencies[0] + def _verify_genereate_responses( responses: Tuple[int, Any, float],