Fix bug: Incorrect variable used in rem_total_token_offset calculatio… (#13201)

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(

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],