@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4a0f7888a3
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 04:45 UTC

[Core] feat: Implement Priority Scheduling in V1 Engine (#19057)

Signed-off-by: amit <amit.man@gmail.com> Co-authored-by: Roger Wang <Rogerw0108@gmail.com>

amit committed Jun 23, 2025 at 06:18 UTC 4a0f7888a3f0de1d460d90f43b8f3e1826b68dbb
7 files changed +896 -30
docs/usage/v1_guide.md
+12
@@ -45,6 +45,18 @@ For each item, our progress towards V1 support falls into one of the following s
45 - **🟠 Delayed**: Temporarily dropped in V1 but planned to be re-introduced later.
46 - **🔴 Deprecated**: Not planned for V1 unless there is strong demand.
47
48 +!!! note
49 + vLLM V1’s unified scheduler treats both prompt and output tokens the same
50 + way by using a simple dictionary (e.g., `{request_id: num_tokens}`) to dynamically
51 + allocate a fixed token budget per request, enabling features like chunked prefills,
52 + prefix caching, and speculative decoding without a strict separation between prefill
53 + and decode phases.
54 +
55 +The V1 scheduler supports multiple scheduling policies, including First-Come,
56 +First-Served (FCFS) and priority-based scheduling (where requests are processed
57 +based on assigned priority, with FCFS as a tie-breaker), configurable via the
58 +`--scheduling-policy` argument.
59 +
60 ### Hardware
61
62 | Hardware | Status |
tests/v1/core/test_scheduler.py
+589 -1
@@ -1150,7 +1150,6 @@ def test_kv_connector_handles_preemption():
1150 assert len(scheduler.running) == 1
1151 _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
1152 assert len(scheduler.running) == 0
1153 - assert len(scheduler.waiting) == 1
1153 # All memory should be freed since nothing is running.
1154 assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() \
1155 == NUM_BLOCKS - 1
@@ -1265,3 +1264,592 @@ def test_memory_leak():
1264
1265 # Confirm no memory leak.
1266 assert_scheduler_empty(scheduler)
1267 +
1268 +
1269 +def create_scheduler_with_priority(
1270 + model: str = "facebook/opt-125m",
1271 + max_num_seqs: int = 16,
1272 + max_num_batched_tokens: int = 8192,
1273 + enable_prefix_caching: Optional[bool] = None,
1274 + long_prefill_token_threshold: int = 0,
1275 + disable_chunked_mm_input: bool = False,
1276 + use_kv_connector: bool = False,
1277 + num_blocks: int = 10000,
1278 + block_size: int = 16,
1279 + max_model_len: Optional[int] = None,
1280 + num_speculative_tokens: Optional[int] = None,
1281 +) -> Scheduler:
1282 + '''Create scheduler with priority policy enabled.
1283 +
1284 + Args:
1285 + model: model under test
1286 + max_num_seqs: max sequences to schedule
1287 + max_num_batch_tokens: max num tokens to batch
1288 + enable_prefix_caching: optionally force APC config
1289 + (True/False) or use default
1290 + (None)
1291 +
1292 + Returns:
1293 + {class}`Scheduler` instance with priority scheduling
1294 + '''
1295 + if max_model_len is None:
1296 + max_model_len = max_num_batched_tokens
1297 + scheduler_config = SchedulerConfig(
1298 + max_num_seqs=max_num_seqs,
1299 + max_num_batched_tokens=max_num_batched_tokens,
1300 + max_model_len=max_model_len,
1301 + long_prefill_token_threshold=long_prefill_token_threshold,
1302 + disable_chunked_mm_input=disable_chunked_mm_input,
1303 + enable_chunked_prefill=True,
1304 + policy="priority", # Enable priority scheduling
1305 + )
1306 + model_config = ModelConfig(
1307 + model=model,
1308 + task="auto",
1309 + tokenizer=model,
1310 + tokenizer_mode="auto",
1311 + trust_remote_code=True,
1312 + dtype="float16",
1313 + seed=42,
1314 + )
1315 + # Cache config, optionally force APC
1316 + kwargs_cache = ({} if enable_prefix_caching is None else {
1317 + 'enable_prefix_caching': enable_prefix_caching
1318 + })
1319 + cache_config = CacheConfig(
1320 + block_size=block_size,
1321 + gpu_memory_utilization=0.9,
1322 + swap_space=0,
1323 + cache_dtype="auto",
1324 + **kwargs_cache,
1325 + )
1326 + kv_transfer_config = KVTransferConfig(
1327 + kv_connector="SharedStorageConnector",
1328 + kv_role="kv_both",
1329 + kv_connector_extra_config={"shared_storage_path": "local_storage"},
1330 + ) if use_kv_connector else None
1331 +
1332 + speculative_config: Optional[SpeculativeConfig] = None
1333 + if num_speculative_tokens is not None:
1334 + speculative_config = SpeculativeConfig(
1335 + model="ngram", num_speculative_tokens=num_speculative_tokens)
1336 +
1337 + vllm_config = VllmConfig(
1338 + scheduler_config=scheduler_config,
1339 + model_config=model_config,
1340 + cache_config=cache_config,
1341 + kv_transfer_config=kv_transfer_config,
1342 + speculative_config=speculative_config,
1343 + )
1344 + kv_cache_config = KVCacheConfig(
1345 + num_blocks=num_blocks, # A large number of blocks to hold all requests
1346 + kv_cache_tensors=[],
1347 + kv_cache_groups=[
1348 + KVCacheGroupSpec(['layer'],
1349 + FullAttentionSpec(block_size, 1, 1, torch.float32,
1350 + False))
1351 + ],
1352 + )
1353 + cache_config.num_gpu_blocks = num_blocks
1354 + return Scheduler(
1355 + vllm_config=vllm_config,
1356 + kv_cache_config=kv_cache_config,
1357 + log_stats=True,
1358 + structured_output_manager=StructuredOutputManager(vllm_config),
1359 + )
1360 +
1361 +
1362 +def create_requests_with_priority(
1363 + num_requests: int,
1364 + priorities: list[int],
1365 + arrival_times: Optional[list[float]] = None,
1366 + num_tokens: int = 10,
1367 + mm_positions: Optional[list[PlaceholderRange]] = None,
1368 + max_tokens: int = 16,
1369 + stop_token_ids: Optional[list[int]] = None,
1370 + prompt_logprobs: Optional[int] = None):
1371 + """Create requests with specified priorities and arrival times."""
1372 + assert len(priorities) == num_requests
1373 + if arrival_times is not None:
1374 + assert len(arrival_times) == num_requests
1375 + else:
1376 + arrival_times = [float(i) for i in range(num_requests)]
1377 +
1378 + sampling_params = SamplingParams(ignore_eos=False,
1379 + max_tokens=max_tokens,
1380 + stop_token_ids=stop_token_ids,
1381 + prompt_logprobs=prompt_logprobs)
1382 + requests = []
1383 + for i in range(num_requests):
1384 + if mm_positions is not None:
1385 + mm_position = mm_positions[i]
1386 + mm_inputs = [MultiModalKwargs({})] * len(mm_position)
1387 + else:
1388 + mm_position = None
1389 + mm_inputs = None
1390 + request = Request(
1391 + request_id=f"{i}",
1392 + prompt_token_ids=[i] * num_tokens,
1393 + sampling_params=sampling_params,
1394 + pooling_params=None,
1395 + multi_modal_inputs=mm_inputs,
1396 + multi_modal_placeholders=mm_position,
1397 + multi_modal_hashes=None,
1398 + eos_token_id=EOS_TOKEN_ID,
1399 + arrival_time=arrival_times[i],
1400 + priority=priorities[i],
1401 + )
1402 + requests.append(request)
1403 + return requests
1404 +
1405 +
1406 +def test_priority_scheduling_basic_ordering():
1407 + """Test that requests are scheduled in priority order
1408 + (lower value = higher priority)."""
1409 + scheduler = create_scheduler_with_priority()
1410 +
1411 + # Create requests with different priorities
1412 + # Priority 0 (highest), 1, 2 (lowest)
1413 + priorities = [2, 0, 1] # Add in non-priority order
1414 + arrival_times = [1.0, 2.0, 3.0] # All different arrival times
1415 + requests = create_requests_with_priority(num_requests=3,
1416 + priorities=priorities,
1417 + arrival_times=arrival_times)
1418 +
1419 + # Add requests in non-priority order
1420 + for request in requests:
1421 + scheduler.add_request(request)
1422 +
1423 + # Schedule and verify priority order
1424 + output = scheduler.schedule()
1425 +
1426 + # Should schedule all requests since they fit in budget
1427 + assert len(output.scheduled_new_reqs) == 3
1428 +
1429 + # Verify they are scheduled in priority order:
1430 + # req_1 (priority 0), req_2 (priority 1), req_0 (priority 2)
1431 + scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
1432 + assert scheduled_req_ids == ["1", "2", "0"]
1433 +
1434 +
1435 +def test_priority_scheduling_arrival_time_tiebreaker():
1436 + """Test that arrival time is used
1437 + as tiebreaker when priorities are equal."""
1438 + scheduler = create_scheduler_with_priority()
1439 +
1440 + # Create requests with same priority but different arrival times
1441 + priorities = [1, 1, 1] # All same priority
1442 + arrival_times = [3.0, 1.0, 2.0] # Different arrival times
1443 + requests = create_requests_with_priority(num_requests=3,
1444 + priorities=priorities,
1445 + arrival_times=arrival_times)
1446 +
1447 + # Add requests in non-arrival order
1448 + for request in requests:
1449 + scheduler.add_request(request)
1450 +
1451 + # Schedule and verify arrival time order
1452 + output = scheduler.schedule()
1453 +
1454 + # Should schedule all requests since they fit in budget
1455 + assert len(output.scheduled_new_reqs) == 3
1456 +
1457 + # Verify they are scheduled in arrival time order:
1458 + # req_1 (1.0), req_2 (2.0), req_0 (3.0)
1459 + scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
1460 + assert scheduled_req_ids == ["1", "2", "0"]
1461 +
1462 +
1463 +def test_priority_scheduling_mixed_priority_and_arrival():
1464 + """Test priority scheduling with mixed priorities and arrival times."""
1465 + scheduler = create_scheduler_with_priority()
1466 +
1467 + # Create requests with mixed priorities and arrival times
1468 + priorities = [2, 1, 1, 0] # Mixed priorities
1469 + arrival_times = [1.0, 3.0, 2.0, 4.0] # Mixed arrival times
1470 + requests = create_requests_with_priority(num_requests=4,
1471 + priorities=priorities,
1472 + arrival_times=arrival_times)
1473 +
1474 + # Add requests
1475 + for request in requests:
1476 + scheduler.add_request(request)
1477 +
1478 + # Schedule and verify order
1479 + output = scheduler.schedule()
1480 +
1481 + # Should schedule all requests since they fit in budget
1482 + assert len(output.scheduled_new_reqs) == 4
1483 +
1484 + # Expected order:
1485 + # 1. req_3 (priority 0, arrival 4.0)
1486 + # 2. req_2 (priority 1, arrival 2.0) - earlier arrival than req_1
1487 + # 3. req_1 (priority 1, arrival 3.0)
1488 + # 4. req_0 (priority 2, arrival 1.0)
1489 + scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
1490 + assert scheduled_req_ids == ["3", "2", "1", "0"]
1491 +
1492 +
1493 +def test_priority_scheduling_preemption():
1494 + """Test that priority scheduling preempts
1495 + lower priority requests when memory is constrained."""
1496 + # Create scheduler with very limited memory to force preemption
1497 + scheduler = create_scheduler_with_priority(
1498 + max_num_seqs=3, # Allow multiple requests
1499 + max_num_batched_tokens=200,
1500 + num_blocks=6, # Very limited blocks to force memory pressure
1501 + block_size=16, # Standard block size
1502 + )
1503 +
1504 + # Create initial low-priority requests that will consume most memory
1505 + low_priority_requests = create_requests_with_priority(
1506 + num_requests=2,
1507 + priorities=[5, 5], # Low priority
1508 + arrival_times=[1.0, 2.0],
1509 + num_tokens=30 # Large enough to consume significant memory
1510 + )
1511 +
1512 + # Add and schedule low priority requests
1513 + for request in low_priority_requests:
1514 + scheduler.add_request(request)
1515 +
1516 + output = scheduler.schedule()
1517 + assert len(output.scheduled_new_reqs) == 2
1518 +
1519 + # Simulate model execution to move requests to running state
1520 + model_output = ModelRunnerOutput(
1521 + req_ids=[req.request_id for req in low_priority_requests],
1522 + req_id_to_index={
1523 + req.request_id: i
1524 + for i, req in enumerate(low_priority_requests)
1525 + },
1526 + sampled_token_ids=[[100] for _ in low_priority_requests],
1527 + spec_token_ids=None,
1528 + logprobs=None,
1529 + prompt_logprobs_dict={},
1530 + pooler_output=[],
1531 + )
1532 + scheduler.update_from_output(output, model_output)
1533 +
1534 + # Verify both requests are running
1535 + assert len(scheduler.running) == 2
1536 +
1537 + # Now add a high-priority request that requires memory allocation
1538 + # This should trigger preemption due to memory constraints
1539 + high_priority_request = create_requests_with_priority(
1540 + num_requests=1,
1541 + priorities=[0], # High priority
1542 + arrival_times=[3.0],
1543 + num_tokens=30 # Large enough to require significant memory
1544 + )[0]
1545 +
1546 + scheduler.add_request(high_priority_request)
1547 +
1548 + # Schedule again - this should trigger
1549 + # preemption when trying to allocate memory
1550 + output = scheduler.schedule()
1551 +
1552 + # Due to the scheduler's design, if preemption happens
1553 + # during running request scheduling,
1554 + # waiting requests won't be scheduled in the same step
1555 + # Let's check if preemption occurred by looking at the waiting queue
1556 +
1557 + # If preemption happened, we should see requests in the
1558 + # waiting queue
1559 + if len(scheduler.waiting) > 1: # high priority + preempted request
1560 + # Preemption occurred - verify the high priority request
1561 + # gets scheduled next
1562 + output2 = scheduler.schedule()
1563 + assert len(output2.scheduled_new_reqs) == 1
1564 + # High priority request
1565 + assert output2.scheduled_new_reqs[0].req_id == "0"
1566 + else:
1567 + # No preemption needed - all requests fit
1568 + # This is also valid behavior if memory allows
1569 + assert len(output.scheduled_new_reqs) == 1
1570 + # High priority request
1571 + assert output.scheduled_new_reqs[0].req_id == "0"
1572 +
1573 +
1574 +def test_priority_scheduling_no_preemption_when_space_available():
1575 + """Test that preemption doesn't happen
1576 + when there's space for new requests."""
1577 + scheduler = create_scheduler_with_priority(
1578 + max_num_seqs=3, # Allow 3 concurrent requests
1579 + max_num_batched_tokens=200, # Sufficient token budget
1580 + )
1581 +
1582 + # Add two low-priority running requests
1583 + low_priority_requests = create_requests_with_priority(
1584 + num_requests=2,
1585 + priorities=[5, 5],
1586 + arrival_times=[1.0, 2.0],
1587 + num_tokens=30)
1588 +
1589 + for request in low_priority_requests:
1590 + scheduler.add_request(request)
1591 +
1592 + output = scheduler.schedule()
1593 + model_output = ModelRunnerOutput(
1594 + req_ids=[req.request_id for req in low_priority_requests],
1595 + req_id_to_index={
1596 + req.request_id: i
1597 + for i, req in enumerate(low_priority_requests)
1598 + },
1599 + sampled_token_ids=[[100] for _ in low_priority_requests],
1600 + spec_token_ids=None,
1601 + logprobs=None,
1602 + prompt_logprobs_dict={},
1603 + pooler_output=[],
1604 + )
1605 + scheduler.update_from_output(output, model_output)
1606 +
1607 + # Add high-priority request
1608 + high_priority_request = create_requests_with_priority(num_requests=1,
1609 + priorities=[0],
1610 + arrival_times=[3.0],
1611 + num_tokens=30)[0]
1612 +
1613 + scheduler.add_request(high_priority_request)
1614 +
1615 + # Schedule - should not preempt since there's space
1616 + output = scheduler.schedule()
1617 +
1618 + # Should schedule the new request without preemption
1619 + assert len(output.scheduled_new_reqs) == 1
1620 + assert len(scheduler.running) == 3 # All three requests running
1621 + assert len(scheduler.waiting) == 0 # No requests waiting
1622 +
1623 +
1624 +def test_priority_scheduling_preemption_victim_selection():
1625 + """Test that the correct victim is selected for
1626 + preemption based on priority and arrival time."""
1627 + # This test verifies the priority-based victim selection logic
1628 + # by checking the waiting queue order after adding requests with different
1629 + # priorities
1630 + scheduler = create_scheduler_with_priority(
1631 + max_num_seqs=1, # Force sequential processing to test priority order
1632 + )
1633 +
1634 + # Create requests with different priorities
1635 + requests = create_requests_with_priority(
1636 + num_requests=3,
1637 + priorities=[3, 2, 0], # Different priorities: low, medium, high
1638 + arrival_times=[1.0, 2.0, 3.0],
1639 + num_tokens=10)
1640 +
1641 + # Add all requests
1642 + for request in requests:
1643 + scheduler.add_request(request)
1644 +
1645 + # Schedule - should only schedule the highest priority request
1646 + # (req_2, priority 0)
1647 + output = scheduler.schedule()
1648 + assert len(output.scheduled_new_reqs) == 1
1649 + assert output.scheduled_new_reqs[0].req_id == "2" # Highest priority
1650 +
1651 + # Verify the waiting queue has the remaining requests in priority order
1652 + assert len(scheduler.waiting) == 2
1653 +
1654 + # Extract waiting requests and verify priority order
1655 + waiting_requests = list(scheduler.waiting)
1656 +
1657 + waiting_priorities = [req.priority for req in waiting_requests]
1658 + waiting_req_ids = [req.request_id for req in waiting_requests]
1659 +
1660 + # Should be req_1 (priority 2) then req_0 (priority 3)
1661 + assert waiting_priorities == [2, 3]
1662 + assert waiting_req_ids == ["1", "0"]
1663 +
1664 +
1665 +def test_priority_scheduling_equal_priority_preemption():
1666 + """Test arrival time tiebreaker when requests have equal priority."""
1667 + # This test verifies that arrival time is used as a tiebreaker for equal
1668 + # priorities
1669 + scheduler = create_scheduler_with_priority(
1670 + max_num_seqs=1, # Force sequential processing
1671 + )
1672 +
1673 + # Create requests with same priority but different arrival times
1674 + requests = create_requests_with_priority(
1675 + num_requests=3,
1676 + priorities=[2, 2, 2], # Same priority
1677 + arrival_times=[3.0, 1.0, 2.0], # Different arrival times
1678 + num_tokens=10)
1679 +
1680 + # Add all requests
1681 + for request in requests:
1682 + scheduler.add_request(request)
1683 +
1684 + # Schedule - should schedule the request with earliest arrival time
1685 + output = scheduler.schedule()
1686 + assert len(output.scheduled_new_reqs) == 1
1687 + assert output.scheduled_new_reqs[0].req_id == "1" # Earliest arrival (1.0)
1688 +
1689 + # Verify the waiting queue has remaining requests in arrival time order
1690 + assert len(scheduler.waiting) == 2
1691 +
1692 + # Extract waiting requests and verify arrival time order
1693 + waiting_requests = list(scheduler.waiting)
1694 +
1695 + waiting_arrival_times = [req.arrival_time for req in waiting_requests]
1696 + waiting_req_ids = [req.request_id for req in waiting_requests]
1697 +
1698 + # Should be req_2 (arrival 2.0) then req_0 (arrival 3.0)
1699 + assert waiting_arrival_times == [2.0, 3.0]
1700 + assert waiting_req_ids == ["2", "0"]
1701 +
1702 +
1703 +def test_priority_scheduling_waiting_queue_order():
1704 + """Test that the waiting queue maintains priority order."""
1705 + scheduler = create_scheduler_with_priority(
1706 + max_num_seqs=1, # Only one request can run at a time
1707 + )
1708 +
1709 + # Create multiple requests with different priorities
1710 + requests = create_requests_with_priority(
1711 + num_requests=4,
1712 + priorities=[3, 1, 2, 0], # Mixed priorities
1713 + arrival_times=[1.0, 2.0, 3.0, 4.0],
1714 + num_tokens=10)
1715 +
1716 + # Add all requests
1717 + for request in requests:
1718 + scheduler.add_request(request)
1719 +
1720 + # Schedule - should only schedule the highest priority request
1721 + # (req_3, priority 0)
1722 + output = scheduler.schedule()
1723 + assert len(output.scheduled_new_reqs) == 1
1724 + assert output.scheduled_new_reqs[0].req_id == "3"
1725 +
1726 + # Verify waiting queue has remaining requests in priority order
1727 + assert len(scheduler.waiting) == 3
1728 +
1729 + # Extract requests from waiting queue
1730 + # (it's a heap, so we need to pop to see order)
1731 + waiting_requests = list(scheduler.waiting)
1732 +
1733 + waiting_priorities = [req.priority for req in waiting_requests]
1734 + waiting_req_ids = [req.request_id for req in waiting_requests]
1735 +
1736 + # Should be ordered by priority: req_1 (1), req_2 (2), req_0 (3)
1737 + assert waiting_req_ids == ["1", "2", "0"]
1738 + assert waiting_priorities == [1, 2, 3]
1739 +
1740 +
1741 +def test_priority_scheduling_fcfs_fallback():
1742 + """Test that FCFS behavior is maintained when all
1743 + requests have same priority."""
1744 + scheduler = create_scheduler_with_priority()
1745 +
1746 + # Create requests with same priority but different arrival times
1747 + priorities = [1, 1, 1, 1] # All same priority
1748 + arrival_times = [4.0, 1.0, 3.0, 2.0] # Different arrival times
1749 + requests = create_requests_with_priority(num_requests=4,
1750 + priorities=priorities,
1751 + arrival_times=arrival_times)
1752 +
1753 + # Add requests
1754 + for request in requests:
1755 + scheduler.add_request(request)
1756 +
1757 + # Schedule
1758 + output = scheduler.schedule()
1759 +
1760 + # Should schedule all requests in arrival time order
1761 + assert len(output.scheduled_new_reqs) == 4
1762 + scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
1763 +
1764 + # Expected order by arrival time:
1765 + # req_1 (1.0), req_3 (2.0), req_2 (3.0), req_0 (4.0)
1766 + assert scheduled_req_ids == ["1", "3", "2", "0"]
1767 +
1768 +
1769 +def test_priority_scheduling_with_limited_slots():
1770 + """Test priority scheduling when max_num_seqs limits concurrent requests."""
1771 + scheduler = create_scheduler_with_priority(
1772 + max_num_seqs=2, # Only allow 2 concurrent requests
1773 + max_num_batched_tokens=1000, # Plenty of token budget
1774 + )
1775 +
1776 + # Create requests with different priorities
1777 + requests = create_requests_with_priority(
1778 + num_requests=4,
1779 + priorities=[3, 1, 2, 0], # Mixed priorities
1780 + arrival_times=[1.0, 2.0, 3.0, 4.0],
1781 + num_tokens=10)
1782 +
1783 + # Add all requests
1784 + for request in requests:
1785 + scheduler.add_request(request)
1786 +
1787 + # Schedule - should only schedule the 2 highest priority requests
1788 + output = scheduler.schedule()
1789 + assert len(output.scheduled_new_reqs) == 2
1790 +
1791 + # Should schedule req_3 (priority 0) and req_1 (priority 1)
1792 + scheduled_req_ids = [req.req_id for req in output.scheduled_new_reqs]
1793 + assert "3" in scheduled_req_ids # Priority 0
1794 + assert "1" in scheduled_req_ids # Priority 1
1795 +
1796 + # Remaining requests should be in waiting queue in priority order
1797 + assert len(scheduler.waiting) == 2
1798 +
1799 + # Extract waiting requests and verify order
1800 + waiting_requests = list(scheduler.waiting)
1801 + waiting_priorities = [req.priority for req in waiting_requests]
1802 + waiting_req_ids = [req.request_id for req in waiting_requests]
1803 +
1804 + # Should be req_2 (priority 2) then req_0 (priority 3)
1805 + assert waiting_priorities == [2, 3]
1806 + assert waiting_req_ids == ["2", "0"]
1807 +
1808 +
1809 +def test_priority_scheduling_heap_property():
1810 + """Test that the waiting queue maintains heap
1811 + property for priority scheduling."""
1812 + scheduler = create_scheduler_with_priority(
1813 + max_num_seqs=1, # Only one request can run at a time
1814 + )
1815 +
1816 + # Add requests in random priority order
1817 + priorities = [5, 1, 8, 3, 2, 7, 4, 6]
1818 + arrival_times = [float(i) for i in range(len(priorities))]
1819 + requests = create_requests_with_priority(num_requests=len(priorities),
1820 + priorities=priorities,
1821 + arrival_times=arrival_times,
1822 + num_tokens=10)
1823 +
1824 + # Add all requests
1825 + for request in requests:
1826 + scheduler.add_request(request)
1827 +
1828 + # Schedule one request at a time and verify priority order
1829 + scheduled_priorities = []
1830 +
1831 + while scheduler.waiting:
1832 + output = scheduler.schedule()
1833 + if output.scheduled_new_reqs:
1834 + req = output.scheduled_new_reqs[0]
1835 + scheduled_priorities.append(requests[int(req.req_id)].priority)
1836 +
1837 + # Simulate completion to make room for next request
1838 + model_output = ModelRunnerOutput(
1839 + req_ids=[req.req_id],
1840 + req_id_to_index={req.req_id: 0},
1841 + sampled_token_ids=[[100]],
1842 + spec_token_ids=None,
1843 + logprobs=None,
1844 + prompt_logprobs_dict={},
1845 + pooler_output=[],
1846 + )
1847 + scheduler.update_from_output(output, model_output)
1848 +
1849 + # Finish the request to make room for the next one
1850 + scheduler.finish_requests(req.req_id,
1851 + RequestStatus.FINISHED_STOPPED)
1852 +
1853 + # Verify requests were scheduled in priority order (lowest value first)
1854 + expected_priorities = sorted(priorities)
1855 + assert scheduled_priorities == expected_priorities
vllm/v1/core/sched/request_queue.py new
+224
@@ -0,0 +1,224 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +import heapq
7 +from abc import ABC, abstractmethod
8 +from collections import deque
9 +from collections.abc import Iterable, Iterator
10 +from enum import Enum
11 +
12 +from vllm.v1.request import Request
13 +
14 +
15 +class SchedulingPolicy(Enum):
16 + """Enum for scheduling policies."""
17 + FCFS = "fcfs"
18 + PRIORITY = "priority"
19 +
20 +
21 +class RequestQueue(ABC):
22 + """Abstract base class for request queues."""
23 +
24 + @abstractmethod
25 + def add_request(self, request: Request) -> None:
26 + """Add a request to the queue according to the policy."""
27 + pass
28 +
29 + @abstractmethod
30 + def pop_request(self) -> Request:
31 + """Pop a request from the queue according to the policy."""
32 + pass
33 +
34 + @abstractmethod
35 + def peek_request(self) -> Request:
36 + """Peek at the request at the front of the queue without removing it."""
37 + pass
38 +
39 + @abstractmethod
40 + def prepend_request(self, request: Request) -> None:
41 + """Prepend a request to the front of the queue."""
42 + pass
43 +
44 + @abstractmethod
45 + def prepend_requests(self, requests: RequestQueue) -> None:
46 + """Prepend all requests from another queue to the front of this
47 + queue."""
48 + pass
49 +
50 + @abstractmethod
51 + def remove_request(self, request: Request) -> None:
52 + """Remove a specific request from the queue."""
53 + pass
54 +
55 + @abstractmethod
56 + def remove_requests(self, requests: Iterable[Request]) -> None:
57 + """Remove multiple specific requests from the queue."""
58 + pass
59 +
60 + @abstractmethod
61 + def __bool__(self) -> bool:
62 + """Check if queue has any requests."""
63 + pass
64 +
65 + @abstractmethod
66 + def __len__(self) -> int:
67 + """Get number of requests in queue."""
68 + pass
69 +
70 + @abstractmethod
71 + def __iter__(self) -> Iterator[Request]:
72 + """Iterate over the queue according to the policy."""
73 + pass
74 +
75 + @abstractmethod
76 + def __reversed__(self) -> Iterator[Request]:
77 + """Iterate over the queue in reverse order."""
78 + pass
79 +
80 +
81 +class FCFSRequestQueue(deque[Request], RequestQueue):
82 + """A first-come-first-served queue that supports deque operations."""
83 +
84 + def add_request(self, request: Request) -> None:
85 + """Add a request to the queue according to FCFS policy."""
86 + self.append(request)
87 +
88 + def pop_request(self) -> Request:
89 + """Pop a request from the queue according to FCFS policy."""
90 + return self.popleft()
91 +
92 + def peek_request(self) -> Request:
93 + """Peek at the next request in the queue without removing it."""
94 + if not self:
95 + raise IndexError("peek from an empty queue")
96 + return self[0]
97 +
98 + def prepend_request(self, request: Request) -> None:
99 + """Prepend a request to the front of the queue."""
100 + self.appendleft(request)
101 +
102 + def prepend_requests(self, requests: RequestQueue) -> None:
103 + """Prepend all requests from another queue to the front of this
104 + queue."""
105 + self.extendleft(reversed(requests))
106 +
107 + def remove_request(self, request: Request) -> None:
108 + """Remove a specific request from the queue."""
109 + self.remove(request)
110 +
111 + def remove_requests(self, requests: Iterable[Request]) -> None:
112 + """Remove multiple specific requests from the queue."""
113 + requests_to_remove = set(requests)
114 + filtered_requests = [
115 + req for req in self if req not in requests_to_remove
116 + ]
117 + # deque does not support in-place filtering, so we need to clear
118 + # and extend
119 + self.clear()
120 + self.extend(filtered_requests)
121 +
122 + def __bool__(self) -> bool:
123 + """Check if queue has any requests."""
124 + return len(self) > 0
125 +
126 + def __len__(self) -> int:
127 + """Get number of requests in queue."""
128 + return super().__len__()
129 +
130 + def __iter__(self) -> Iterator[Request]:
131 + """Iterate over the queue according to FCFS policy."""
132 + return super().__iter__()
133 +
134 + def __reversed__(self) -> Iterator[Request]:
135 + """Iterate over the queue in reverse order."""
136 + return super().__reversed__()
137 +
138 +
139 +class PriorityRequestQueue(RequestQueue):
140 + """
141 + A priority queue that supports heap operations.
142 +
143 + Requests with a smaller value of `priority` are processed first.
144 + If multiple requests have the same priority, the one with the earlier
145 + `arrival_time` is processed first.
146 + """
147 +
148 + def __init__(self) -> None:
149 + self._heap: list[tuple[int, float, Request]] = []
150 +
151 + def add_request(self, request: Request) -> None:
152 + """Add a request to the queue according to priority policy."""
153 + heapq.heappush(self._heap,
154 + (request.priority, request.arrival_time, request))
155 +
156 + def pop_request(self) -> Request:
157 + """Pop a request from the queue according to priority policy."""
158 + if not self._heap:
159 + raise IndexError("pop from empty heap")
160 + _, _, request = heapq.heappop(self._heap)
161 + return request
162 +
163 + def peek_request(self) -> Request:
164 + """Peek at the next request in the queue without removing it."""
165 + if not self._heap:
166 + raise IndexError("peek from empty heap")
167 + _, _, request = self._heap[0]
168 + return request
169 +
170 + def prepend_request(self, request: Request) -> None:
171 + """Add a request to the queue according to priority policy.
172 +
173 + Note: In a priority queue, there is no concept of prepending to the
174 + front. Requests are ordered by (priority, arrival_time)."""
175 + self.add_request(request)
176 +
177 + def prepend_requests(self, requests: RequestQueue) -> None:
178 + """Add all requests from another queue according to priority policy.
179 +
180 + Note: In a priority queue, there is no concept of prepending to the
181 + front. Requests are ordered by (priority, arrival_time)."""
182 + for request in requests:
183 + self.add_request(request)
184 +
185 + def remove_request(self, request: Request) -> None:
186 + """Remove a specific request from the queue."""
187 + self._heap = [(p, t, r) for p, t, r in self._heap if r != request]
188 + heapq.heapify(self._heap)
189 +
190 + def remove_requests(self, requests: Iterable[Request]) -> None:
191 + """Remove multiple specific requests from the queue."""
192 + requests_to_remove = set(requests)
193 + self._heap = [(p, t, r) for p, t, r in self._heap
194 + if r not in requests_to_remove]
195 + heapq.heapify(self._heap)
196 +
197 + def __bool__(self) -> bool:
198 + """Check if queue has any requests."""
199 + return bool(self._heap)
200 +
201 + def __len__(self) -> int:
202 + """Get number of requests in queue."""
203 + return len(self._heap)
204 +
205 + def __iter__(self) -> Iterator[Request]:
206 + """Iterate over the queue according to priority policy."""
207 + heap_copy = self._heap[:]
208 + while heap_copy:
209 + _, _, request = heapq.heappop(heap_copy)
210 + yield request
211 +
212 + def __reversed__(self) -> Iterator[Request]:
213 + """Iterate over the queue in reverse priority order."""
214 + return reversed(list(self))
215 +
216 +
217 +def create_request_queue(policy: SchedulingPolicy) -> RequestQueue:
218 + """Create request queue based on scheduling policy."""
219 + if policy == SchedulingPolicy.PRIORITY:
220 + return PriorityRequestQueue()
221 + elif policy == SchedulingPolicy.FCFS:
222 + return FCFSRequestQueue()
223 + else:
224 + raise ValueError(f"Unknown scheduling policy: {policy}")
vllm/v1/core/sched/scheduler.py
+61 -27
@@ -22,6 +22,8 @@ from vllm.v1.core.kv_cache_manager import KVCacheManager
22 from vllm.v1.core.sched.interface import SchedulerInterface
23 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
24 SchedulerOutput)
25 +from vllm.v1.core.sched.request_queue import (SchedulingPolicy,
26 + create_request_queue)
27 from vllm.v1.core.sched.utils import check_stop
28 from vllm.v1.engine import (EngineCoreEventType, EngineCoreOutput,
29 EngineCoreOutputs)
@@ -94,8 +96,16 @@ class Scheduler(SchedulerInterface):
96
97 # req_id -> Request
98 self.requests: dict[str, Request] = {}
99 + # Scheduling policy
100 + if self.scheduler_config.policy == "priority":
101 + self.policy = SchedulingPolicy.PRIORITY
102 + elif self.scheduler_config.policy == "fcfs":
103 + self.policy = SchedulingPolicy.FCFS
104 + else:
105 + raise ValueError(
106 + f"Unknown scheduling policy: {self.scheduler_config.policy}")
107 # Priority queues for requests.
98 - self.waiting: deque[Request] = deque()
108 + self.waiting = create_request_queue(self.policy)
109 self.running: list[Request] = []
110
111 # The request IDs that are finished in between the previous and the
@@ -247,7 +257,15 @@ class Scheduler(SchedulerInterface):
257 if new_blocks is None:
258 # The request cannot be scheduled.
259 # Preempt the lowest-priority request.
250 - preempted_req = self.running.pop()
260 + if self.policy == SchedulingPolicy.PRIORITY:
261 + preempted_req = max(
262 + self.running,
263 + key=lambda r: (r.priority, r.arrival_time),
264 + )
265 + self.running.remove(preempted_req)
266 + else:
267 + preempted_req = self.running.pop()
268 +
269 self.kv_cache_manager.free(preempted_req)
270 preempted_req.status = RequestStatus.PREEMPTED
271 preempted_req.num_computed_tokens = 0
@@ -255,7 +273,7 @@ class Scheduler(SchedulerInterface):
273 preempted_req.record_event(
274 EngineCoreEventType.PREEMPTED, scheduled_timestamp)
275
258 - self.waiting.appendleft(preempted_req)
276 + self.waiting.prepend_request(preempted_req)
277 preempted_reqs.append(preempted_req)
278 if preempted_req == request:
279 # No more request to preempt.
@@ -311,9 +329,9 @@ class Scheduler(SchedulerInterface):
329 if req.lora_request and req.lora_request.lora_int_id > 0)
330 assert len(scheduled_loras) <= self.lora_config.max_loras
331
314 - # Use a temporary deque to collect requests that need to be skipped
315 - # and put back at the head of the waiting queue later
316 - skipped_waiting_requests: deque[Request] = deque()
332 + # Use a temporary RequestQueue to collect requests that need to be
333 + # skipped and put back at the head of the waiting queue later
334 + skipped_waiting_requests = create_request_queue(self.policy)
335
336 # Next, schedule the WAITING requests.
337 if not preempted_reqs:
@@ -321,7 +339,7 @@ class Scheduler(SchedulerInterface):
339 if len(self.running) == self.max_num_running_reqs:
340 break
341
324 - request = self.waiting[0]
342 + request = self.waiting.peek_request()
343
344 # KVTransfer: skip request if still waiting for remote kvs.
345 if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
@@ -332,8 +350,8 @@ class Scheduler(SchedulerInterface):
350 logger.debug(
351 "%s is still in WAITING_FOR_REMOTE_KVS state.",
352 request.request_id)
335 - self.waiting.popleft()
336 - skipped_waiting_requests.appendleft(request)
353 + self.waiting.pop_request()
354 + skipped_waiting_requests.prepend_request(request)
355 continue
356
357 # Skip request if the structured output request is still waiting
@@ -343,19 +361,18 @@ class Scheduler(SchedulerInterface):
361 if structured_output_req and structured_output_req.grammar:
362 request.status = RequestStatus.WAITING
363 else:
346 - self.waiting.popleft()
347 - skipped_waiting_requests.appendleft(request)
364 + self.waiting.pop_request()
365 + skipped_waiting_requests.prepend_request(request)
366 continue
367
368 # Check that adding the request still respects the max_loras
369 # constraint.
352 - if self.lora_config and request.lora_request and (
353 - len(scheduled_loras) == self.lora_config.max_loras
354 - and request.lora_request.lora_int_id
355 - not in scheduled_loras):
370 + if (self.lora_config and request.lora_request and
371 + (len(scheduled_loras) == self.lora_config.max_loras and
372 + request.lora_request.lora_int_id not in scheduled_loras)):
373 # Scheduling would exceed max_loras, skip.
357 - self.waiting.popleft()
358 - skipped_waiting_requests.appendleft(request)
374 + self.waiting.pop_request()
375 + skipped_waiting_requests.prepend_request(request)
376 continue
377
378 num_external_computed_tokens = 0
@@ -407,8 +424,8 @@ class Scheduler(SchedulerInterface):
424 # pooling requests to be chunked
425 if not self.scheduler_config.chunked_prefill_enabled and \
426 num_new_tokens > token_budget:
410 - self.waiting.popleft()
411 - skipped_waiting_requests.appendleft(request)
427 + self.waiting.pop_request()
428 + skipped_waiting_requests.prepend_request(request)
429 continue
430
431 num_new_tokens = min(num_new_tokens, token_budget)
@@ -448,17 +465,19 @@ class Scheduler(SchedulerInterface):
465 num_external_computed_tokens,
466 )
467
451 - self.waiting.popleft()
468 + # Request was already popped from self.waiting
469 + # unless it was re-added above due to new_blocks being None.
470 + request = self.waiting.pop_request()
471 if load_kv_async:
472 # If loading async, allocate memory and put request
473 # into the WAITING_FOR_REMOTE_KV state.
455 - skipped_waiting_requests.appendleft(request)
474 + skipped_waiting_requests.prepend_request(request)
475 request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
476 continue
477
478 if request.use_structured_output:
460 - structured_output_request_ids[
461 - request.request_id] = req_index
479 + structured_output_request_ids[request.request_id] = (
480 + req_index)
481 req_index += 1
482 self.running.append(request)
483 if self.log_stats:
@@ -494,7 +513,7 @@ class Scheduler(SchedulerInterface):
513
514 # Put back any skipped requests at the head of the waiting queue
515 if skipped_waiting_requests:
497 - self.waiting.extendleft(skipped_waiting_requests)
516 + self.waiting.prepend_requests(skipped_waiting_requests)
517
518 # Check if the scheduling constraints are satisfied.
519 total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
@@ -896,7 +915,7 @@ class Scheduler(SchedulerInterface):
915 return len(self.running), len(self.waiting)
916
917 def add_request(self, request: Request) -> None:
899 - self.waiting.append(request)
918 + self.waiting.add_request(request)
919 self.requests[request.request_id] = request
920 if self.log_stats:
921 request.record_event(EngineCoreEventType.QUEUED)
@@ -917,16 +936,31 @@ class Scheduler(SchedulerInterface):
936 else:
937 request_ids = set(request_ids)
938
939 + running_requests_to_remove = []
940 + waiting_requests_to_remove = []
941 + valid_requests = []
942 +
943 + # First pass: collect requests to remove from queues
944 for req_id in request_ids:
945 request = self.requests.get(req_id)
946 if request is None:
947 # Invalid request ID.
948 continue
949
950 + valid_requests.append(request)
951 if request.status == RequestStatus.RUNNING:
927 - self.running.remove(request)
952 + running_requests_to_remove.append(request)
953 else:
929 - self.waiting.remove(request)
954 + waiting_requests_to_remove.append(request)
955 +
956 + # Remove all requests from queues at once for better efficiency
957 + for request in running_requests_to_remove:
958 + self.running.remove(request)
959 + if waiting_requests_to_remove:
960 + self.waiting.remove_requests(waiting_requests_to_remove)
961 +
962 + # Second pass: set status and free requests
963 + for request in valid_requests:
964 request.status = finished_status
965 self._free_request(request)
966
vllm/v1/engine/__init__.py
+1
@@ -68,6 +68,7 @@ class EngineCoreRequest(
68 # belong to, to cover a race condition where the request is sent before
69 # a wave finished notification is received.
70 current_wave: int = 0
71 + priority: int = 0
72
73
74 class EngineCoreEventType(enum.IntEnum):
vllm/v1/engine/processor.py
+1 -2
@@ -219,8 +219,6 @@ class Processor:
219 # TODO(woosuk): Support encoder-decoder models.
220 self._validate_lora(lora_request)
221 self._validate_params(params, lora_request)
222 - if priority != 0:
223 - raise ValueError("V1 does not support priority yet.")
222 if trace_headers is not None:
223 raise ValueError("V1 does not support tracing yet.")
224 if prompt_adapter_request is not None:
@@ -340,6 +338,7 @@ class Processor:
338 arrival_time=arrival_time,
339 lora_request=lora_request,
340 cache_salt=decoder_inputs.get("cache_salt"),
341 + priority=priority,
342 data_parallel_rank=data_parallel_rank,
343 )
344
vllm/v1/request.py
+8
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import enum
5 +import time
6 from typing import TYPE_CHECKING, Any, Optional, Union
7
8 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
@@ -30,18 +31,23 @@ class Request:
31 pooling_params: Optional[PoolingParams],
32 eos_token_id: Optional[int],
33 client_index: int = 0,
34 + arrival_time: Optional[float] = None,
35 lora_request: Optional["LoRARequest"] = None,
36 structured_output_request: Optional["StructuredOutputRequest"] = None,
37 cache_salt: Optional[str] = None,
38 + priority: int = 0,
39 ) -> None:
40 self.request_id = request_id
41 self.client_index = client_index
42 + self.priority = priority
43 self.sampling_params = sampling_params
44 self.pooling_params = pooling_params
45 # Because of LoRA, the eos token id can be different for each request.
46 self.eos_token_id = eos_token_id
47 self.lora_request = lora_request
48 self.structured_output_request = structured_output_request
49 + self.arrival_time = arrival_time if arrival_time is not None else \
50 + time.time()
51
52 self.status = RequestStatus.WAITING
53 if sampling_params and sampling_params.guided_decoding is not None:
@@ -118,11 +124,13 @@ class Request:
124 sampling_params=request.sampling_params,
125 pooling_params=request.pooling_params,
126 eos_token_id=request.eos_token_id,
127 + arrival_time=request.arrival_time,
128 lora_request=request.lora_request,
129 structured_output_request=StructuredOutputRequest(
130 sampling_params=request.sampling_params) \
131 if request.sampling_params else None,
132 cache_salt=request.cache_salt,
133 + priority=request.priority,
134 )
135
136 def append_output_token_ids(