@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7ad894c86a
Fork 0

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

[Bugfix] Prevent cuMemcpyBatchAsync segfault with MTP and KV offloading (#44784)

Signed-off-by: joshua <joshua.abraham@multicorewareinc.com> Co-authored-by: joshua <joshua.abraham@multicorewareinc.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

joshua abraham committed Jun 16, 2026 at 13:28 UTC 7ad894c86a2f3615fe72d739c25567803b5924ec
2 files changed +744 -5
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+682
@@ -14,6 +14,7 @@ from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
14 from vllm.distributed.kv_events import BlockRemoved, BlockStored
15 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
16 OffloadingConnectorScheduler,
17 + RequestOffloadState,
18 )
19 from vllm.v1.core.kv_cache_utils import BlockHash
20 from vllm.v1.kv_cache_interface import (
@@ -28,6 +29,7 @@ from vllm.v1.kv_offload.base import (
29 ReqContext,
30 RequestOffloadingContext,
31 get_offload_block_hash,
32 + make_offload_key,
33 )
34 from vllm.v1.request import RequestStatus
35
@@ -1432,3 +1434,683 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
1434
1435 # The external lookup must have been completely skipped.
1436 runner.manager.lookup.assert_not_called()
1437 +
1438 +
1439 +# ---------------------------------------------------------------------------
1440 +# Eagle/MTP test class
1441 +# ---------------------------------------------------------------------------
1442 +
1443 +
1444 +class TestEagle:
1445 + """Tests for Eagle/MTP speculative decoding support in the offloading
1446 + connector scheduler — both _lookup() unit tests and integration tests."""
1447 +
1448 + # -------------------------------------------------------------------
1449 + # Helpers
1450 + # -------------------------------------------------------------------
1451 +
1452 + @staticmethod
1453 + def _group_keys(group_idx: int, int_hashes: list[int]) -> list:
1454 + return [make_offload_key(str(h).encode(), group_idx) for h in int_hashes]
1455 +
1456 + @staticmethod
1457 + def _make_req_status(
1458 + scheduler: OffloadingConnectorScheduler,
1459 + *,
1460 + num_tokens: int,
1461 + num_computed_tokens: int = 0,
1462 + offload_keys_per_group: list[list[int]],
1463 + ) -> RequestOffloadState:
1464 + """Build RequestOffloadState with synthetic offload keys."""
1465 + req = MagicMock()
1466 + req.request_id = "test-req"
1467 + req.num_tokens = num_tokens
1468 + req.kv_transfer_params = None
1469 +
1470 + state = RequestOffloadState(
1471 + config=scheduler.config,
1472 + req=req,
1473 + req_context=ReqContext(req_id="test-req"),
1474 + offloading_context=RequestOffloadingContext(
1475 + policy=OffloadPolicy.BLOCK_LEVEL
1476 + ),
1477 + num_locally_computed_tokens=num_computed_tokens,
1478 + )
1479 + for idx, (gs, hashes) in enumerate(
1480 + zip(state.group_states, offload_keys_per_group)
1481 + ):
1482 + gs.offload_keys = TestEagle._group_keys(
1483 + scheduler.config.kv_group_configs[idx].group_idx, hashes
1484 + )
1485 + return state
1486 +
1487 + # -------------------------------------------------------------------
1488 + # Lookup unit tests: call _lookup() directly via request_runner
1489 + # -------------------------------------------------------------------
1490 +
1491 + def test_full_attn_lookup_pops_one_block(self, request_runner):
1492 + """Full-attn eagle group with 3 blocks all hit → pop to 2 blocks."""
1493 + block_size = 4
1494 + groups = [
1495 + KVCacheGroupSpec(
1496 + ["layer0"],
1497 + FullAttentionSpec(
1498 + block_size=block_size,
1499 + num_kv_heads=1,
1500 + head_size=1,
1501 + dtype=torch.float32,
1502 + ),
1503 + is_eagle_group=True,
1504 + ),
1505 + ]
1506 + runner = request_runner(
1507 + block_size=block_size,
1508 + num_gpu_blocks=100,
1509 + async_scheduling=False,
1510 + kv_cache_groups=groups,
1511 + )
1512 + runner.manager.lookup.side_effect = lambda key, req_context: (
1513 + int(get_offload_block_hash(key).decode()) in {1, 2, 3}
1514 + )
1515 + sched = runner.connector_scheduler
1516 + req_status = self._make_req_status(
1517 + sched, num_tokens=12, offload_keys_per_group=[[1, 2, 3]]
1518 + )
1519 + # 3 hits, pop to 2 → 2 * block_size = 8 tokens loadable
1520 + assert sched._lookup(req_status) == 8
1521 +
1522 + def test_full_attn_lookup_single_block_returns_zero(self, request_runner):
1523 + """Full-attn eagle group with 1 block hit → pop to 0 → returns 0."""
1524 + block_size = 4
1525 + groups = [
1526 + KVCacheGroupSpec(
1527 + ["layer0"],
1528 + FullAttentionSpec(
1529 + block_size=block_size,
1530 + num_kv_heads=1,
1531 + head_size=1,
1532 + dtype=torch.float32,
1533 + ),
1534 + is_eagle_group=True,
1535 + ),
1536 + ]
1537 + runner = request_runner(
1538 + block_size=block_size,
1539 + num_gpu_blocks=100,
1540 + async_scheduling=False,
1541 + kv_cache_groups=groups,
1542 + )
1543 + runner.manager.lookup.side_effect = lambda key, req_context: (
1544 + int(get_offload_block_hash(key).decode()) in {1}
1545 + )
1546 + sched = runner.connector_scheduler
1547 + req_status = self._make_req_status(
1548 + sched, num_tokens=4, offload_keys_per_group=[[1]]
1549 + )
1550 + # 1 hit, pop to 0 → new_num_hit_tokens < block_size → return 0
1551 + assert sched._lookup(req_status) == 0
1552 +
1553 + def test_full_attn_lookup_no_hits_returns_zero(self, request_runner):
1554 + """Full-attn eagle group with 0 hits returns 0 before pop."""
1555 + block_size = 4
1556 + groups = [
1557 + KVCacheGroupSpec(
1558 + ["layer0"],
1559 + FullAttentionSpec(
1560 + block_size=block_size,
1561 + num_kv_heads=1,
1562 + head_size=1,
1563 + dtype=torch.float32,
1564 + ),
1565 + is_eagle_group=True,
1566 + ),
1567 + ]
1568 + runner = request_runner(
1569 + block_size=block_size,
1570 + num_gpu_blocks=100,
1571 + async_scheduling=False,
1572 + kv_cache_groups=groups,
1573 + )
1574 + runner.manager.lookup.return_value = False
1575 + sched = runner.connector_scheduler
1576 + req_status = self._make_req_status(
1577 + sched, num_tokens=8, offload_keys_per_group=[[1, 2]]
1578 + )
1579 + assert sched._lookup(req_status) == 0
1580 +
1581 + def test_sw_lookup_inflates_query_max(self, request_runner):
1582 + """SW eagle group inflates query_max so _sliding_window_lookup gets
1583 + one extra key beyond what max_hit_size_tokens alone would yield.
1584 +
1585 + With block_size=4, W=2, eagle, num_tokens=13, 4 keys all hitting:
1586 + - max_hit = 13-1 = 12 (SW reduction)
1587 + - Without inflation: num_blocks = cdiv(12,4) = 3 → only 3 keys
1588 + - With inflation: query_max = min(12+4, 4*4=16) = 16,
1589 + num_blocks = cdiv(16,4) = 4 → 4 keys passed to SW
1590 + - SW finds window of 3 (required=W+1=3) at idx 1 → returns 4
1591 + - Pop: 4-1=3 → max_hit = min(12, 12) = 12. Result: 12.
1592 + """
1593 + block_size = 4
1594 + sw_blocks = 2
1595 + groups = [
1596 + KVCacheGroupSpec(
1597 + ["layer0"],
1598 + SlidingWindowSpec(
1599 + block_size=block_size,
1600 + num_kv_heads=1,
1601 + head_size=1,
1602 + dtype=torch.float32,
1603 + sliding_window=sw_blocks * block_size,
1604 + ),
1605 + is_eagle_group=True,
1606 + ),
1607 + ]
1608 + runner = request_runner(
1609 + block_size=block_size,
1610 + num_gpu_blocks=100,
1611 + async_scheduling=False,
1612 + kv_cache_groups=groups,
1613 + )
1614 + runner.manager.lookup.side_effect = lambda key, req_context: (
1615 + int(get_offload_block_hash(key).decode()) in {1, 2, 3, 4}
1616 + )
1617 + sched = runner.connector_scheduler
1618 +
1619 + captured_keys: list = []
1620 + orig_sw_lookup = type(sched)._sliding_window_lookup
1621 +
1622 + def capturing_sw_lookup(self_arg, keys, window, req_context):
1623 + captured_keys.append(list(keys))
1624 + return orig_sw_lookup(self_arg, keys, window, req_context)
1625 +
1626 + sched._sliding_window_lookup = lambda keys, window, req_ctx: (
1627 + capturing_sw_lookup(sched, keys, window, req_ctx)
1628 + )
1629 +
1630 + req_status = self._make_req_status(
1631 + sched, num_tokens=13, offload_keys_per_group=[[1, 2, 3, 4]]
1632 + )
1633 + result = sched._lookup(req_status)
1634 + assert len(captured_keys) == 1
1635 + # Inflation bumped from 3 keys (cdiv(12,4)) to 4 keys (cdiv(16,4))
1636 + assert len(captured_keys[0]) == 4
1637 + # SW finds window of 3 → returns 4, pop to 3 → 3*4=12
1638 + assert result == 12
1639 +
1640 + def test_sw_lookup_requires_extra_window_block(self, request_runner):
1641 + """SW eagle with W=2 and only 2 keys (both hit) uses prefix fallback.
1642 +
1643 + Since required_window = W+1 = 3 but only 2 keys are available
1644 + (inflation is capped by len(offload_keys)), _sliding_window_lookup
1645 + can never find a window of 3. It falls back to prefix count (2).
1646 + Pop: 2-1=1 → max_hit = 4. Result: 4 tokens (degraded from full hit).
1647 + """
1648 + block_size = 4
1649 + sw_blocks = 2
1650 + groups = [
1651 + KVCacheGroupSpec(
1652 + ["layer0"],
1653 + SlidingWindowSpec(
1654 + block_size=block_size,
1655 + num_kv_heads=1,
1656 + head_size=1,
1657 + dtype=torch.float32,
1658 + sliding_window=sw_blocks * block_size,
1659 + ),
1660 + is_eagle_group=True,
1661 + ),
1662 + ]
1663 + runner = request_runner(
1664 + block_size=block_size,
1665 + num_gpu_blocks=100,
1666 + async_scheduling=False,
1667 + kv_cache_groups=groups,
1668 + )
1669 + runner.manager.lookup.side_effect = lambda key, req_context: (
1670 + int(get_offload_block_hash(key).decode()) in {1, 2}
1671 + )
1672 + sched = runner.connector_scheduler
1673 + req_status = self._make_req_status(
1674 + sched, num_tokens=9, offload_keys_per_group=[[1, 2]]
1675 + )
1676 + # Prefix fallback returns 2, pop to 1 → 1*4 = 4 tokens
1677 + assert sched._lookup(req_status) == 4
1678 +
1679 + def test_sw_lookup_w_plus_one_hits_returns_w_blocks(self, request_runner):
1680 + """SW eagle with W=2, 3 contiguous hits → pop to 2 → returns 2*bs."""
1681 + block_size = 4
1682 + sw_blocks = 2
1683 + groups = [
1684 + KVCacheGroupSpec(
1685 + ["layer0"],
1686 + SlidingWindowSpec(
1687 + block_size=block_size,
1688 + num_kv_heads=1,
1689 + head_size=1,
1690 + dtype=torch.float32,
1691 + sliding_window=sw_blocks * block_size,
1692 + ),
1693 + is_eagle_group=True,
1694 + ),
1695 + ]
1696 + runner = request_runner(
1697 + block_size=block_size,
1698 + num_gpu_blocks=100,
1699 + async_scheduling=False,
1700 + kv_cache_groups=groups,
1701 + )
1702 + runner.manager.lookup.side_effect = lambda key, req_context: (
1703 + int(get_offload_block_hash(key).decode()) in {1, 2, 3}
1704 + )
1705 + sched = runner.connector_scheduler
1706 + # num_tokens=13 → max_hit=13-1=12, query_max=min(12+4,12)=12
1707 + # num_blocks=cdiv(12,4)=3, keys=[1,2,3], required_window=3
1708 + # SW finds window of 3, pop to 2 → 2*4=8
1709 + req_status = self._make_req_status(
1710 + sched, num_tokens=13, offload_keys_per_group=[[1, 2, 3]]
1711 + )
1712 + assert sched._lookup(req_status) == 8
1713 +
1714 + def test_eagle_verified_prevents_double_pop(self, request_runner):
1715 + """Once an eagle group has popped, it doesn't pop again on re-iteration.
1716 +
1717 + Setup: group 0 = non-eagle full-attn (3 blocks), group 1 = eagle
1718 + full-attn (3 blocks). Both see all hits. Eagle pops to 2 and tightens
1719 + max_hit to 8. Group 0 re-runs (convergence) but since eagle_verified
1720 + contains group 1, it won't pop again — result stays at 8 tokens.
1721 + """
1722 + block_size = 4
1723 + groups = [
1724 + KVCacheGroupSpec(
1725 + ["layer0"],
1726 + FullAttentionSpec(
1727 + block_size=block_size,
1728 + num_kv_heads=1,
1729 + head_size=1,
1730 + dtype=torch.float32,
1731 + ),
1732 + is_eagle_group=False,
1733 + ),
1734 + KVCacheGroupSpec(
1735 + ["layer1"],
1736 + FullAttentionSpec(
1737 + block_size=block_size,
1738 + num_kv_heads=2,
1739 + head_size=1,
1740 + dtype=torch.float32,
1741 + ),
1742 + is_eagle_group=True,
1743 + ),
1744 + ]
1745 + runner = request_runner(
1746 + block_size=block_size,
1747 + num_gpu_blocks=100,
1748 + async_scheduling=False,
1749 + kv_cache_groups=groups,
1750 + )
1751 + runner.manager.lookup.side_effect = lambda key, req_context: (
1752 + int(get_offload_block_hash(key).decode()) in {1, 2, 3}
1753 + )
1754 + sched = runner.connector_scheduler
1755 + req_status = self._make_req_status(
1756 + sched,
1757 + num_tokens=12,
1758 + offload_keys_per_group=[[1, 2, 3], [1, 2, 3]],
1759 + )
1760 + # Group 0: prefix finds 3 → max_hit=12, num_hit=12
1761 + # Group 1 (eagle): prefix finds 3, pop to 2 → max_hit=8, num_hit=8
1762 + # num_hit(8) < prev num_hit(12) AND group IS eagle → no clear
1763 + # No re-iteration triggered (eagle shrink doesn't trigger re-loop)
1764 + # Final: 8 tokens
1765 + assert sched._lookup(req_status) == 8
1766 +
1767 + def test_non_eagle_tighten_clears_eagle_verified(self, request_runner):
1768 + """Non-eagle group tightening clears eagle_verified → eagle re-pops.
1769 +
1770 + Groups: 0=non-eagle full-attn, 1=eagle full-attn.
1771 + Group 0 has only 1 hit (out of 3 keys) → max_hit tightens to 4.
1772 + This clears eagle_verified. Group 1 runs with max_hit=4 → only 1
1773 + key queried, 1 hit, pop to 0 → returns 0.
1774 + """
1775 + block_size = 4
1776 + groups = [
1777 + KVCacheGroupSpec(
1778 + ["layer0"],
1779 + FullAttentionSpec(
1780 + block_size=block_size,
1781 + num_kv_heads=1,
1782 + head_size=1,
1783 + dtype=torch.float32,
1784 + ),
1785 + is_eagle_group=False,
1786 + ),
1787 + KVCacheGroupSpec(
1788 + ["layer1"],
1789 + FullAttentionSpec(
1790 + block_size=block_size,
1791 + num_kv_heads=2,
1792 + head_size=1,
1793 + dtype=torch.float32,
1794 + ),
1795 + is_eagle_group=True,
1796 + ),
1797 + ]
1798 + runner = request_runner(
1799 + block_size=block_size,
1800 + num_gpu_blocks=100,
1801 + async_scheduling=False,
1802 + kv_cache_groups=groups,
1803 + )
1804 + # Group 0 keys [10,11,12]: only 10 hits.
1805 + # Group 1 keys [1,2,3]: all hit.
1806 + runner.manager.lookup.side_effect = lambda key, req_context: (
1807 + int(get_offload_block_hash(key).decode()) in {10, 1, 2, 3}
1808 + )
1809 + sched = runner.connector_scheduler
1810 + req_status = self._make_req_status(
1811 + sched,
1812 + num_tokens=12,
1813 + offload_keys_per_group=[[10, 11, 12], [1, 2, 3]],
1814 + )
1815 + # Group 0 (non-eagle FA): prefix finds 1 hit → max_hit=4, num_hit=4
1816 + # Group 1 (eagle FA): max_hit=4 → num_blocks=1, keys=[1].
1817 + # Finds 1 hit, pop to 0 → new_num_hit = 0 < block_size → return 0
1818 + assert sched._lookup(req_status) == 0
1819 +
1820 + def test_eagle_verified_survives_eagle_tighten(self, request_runner):
1821 + """Eagle group tightening does NOT clear eagle_verified.
1822 +
1823 + Groups: 0=non-eagle full-attn, 1=eagle full-attn.
1824 + Group 0 finds 3 hits (max_hit=12). Group 1 finds 3 hits, pops to 2
1825 + (max_hit=8). Since group 1 IS eagle, eagle_verified is NOT cleared.
1826 + Result: 8 tokens (eagle only pops once).
1827 + """
1828 + block_size = 4
1829 + groups = [
1830 + KVCacheGroupSpec(
1831 + ["layer0"],
1832 + FullAttentionSpec(
1833 + block_size=block_size,
1834 + num_kv_heads=1,
1835 + head_size=1,
1836 + dtype=torch.float32,
1837 + ),
1838 + is_eagle_group=False,
1839 + ),
1840 + KVCacheGroupSpec(
1841 + ["layer1"],
1842 + FullAttentionSpec(
1843 + block_size=block_size,
1844 + num_kv_heads=2,
1845 + head_size=1,
1846 + dtype=torch.float32,
1847 + ),
1848 + is_eagle_group=True,
1849 + ),
1850 + ]
1851 + runner = request_runner(
1852 + block_size=block_size,
1853 + num_gpu_blocks=100,
1854 + async_scheduling=False,
1855 + kv_cache_groups=groups,
1856 + )
1857 + runner.manager.lookup.side_effect = lambda key, req_context: (
1858 + int(get_offload_block_hash(key).decode()) in {1, 2, 3}
1859 + )
1860 + sched = runner.connector_scheduler
1861 + req_status = self._make_req_status(
1862 + sched,
1863 + num_tokens=12,
1864 + offload_keys_per_group=[[1, 2, 3], [1, 2, 3]],
1865 + )
1866 + # Group 0: 3 hits → max_hit=12, num_hit=12
1867 + # Group 1 (eagle): 3 hits, pop to 2 → max_hit=8, num_hit=8
1868 + # Tightened but IS eagle → no clear. No re-iteration.
1869 + assert sched._lookup(req_status) == 8
1870 +
1871 + # -------------------------------------------------------------------
1872 + # Integration tests: store and load via request_runner
1873 + # -------------------------------------------------------------------
1874 +
1875 + @pytest.mark.parametrize("async_scheduling", [True, False])
1876 + def test_full_attn_store_excludes_trailing_block(
1877 + self, request_runner, async_scheduling: bool
1878 + ):
1879 + """Eagle full-attention group stores all blocks except the trailing
1880 + one.
1881 +
1882 + Setup: 2 groups — group 0 is normal full-attention, group 1 is
1883 + eagle full-attention. With a 3-block prompt, group 1 should store
1884 + only blocks 0 and 1, skipping block 2 (the volatile tail).
1885 + """
1886 + block_size = 4
1887 + block_size_factor = 1
1888 + offloaded_block_size = block_size * block_size_factor
1889 + num_gpu_blocks = 100
1890 +
1891 + kv_cache_groups = [
1892 + KVCacheGroupSpec(
1893 + ["layer0"],
1894 + FullAttentionSpec(
1895 + block_size=block_size,
1896 + num_kv_heads=1,
1897 + head_size=1,
1898 + dtype=torch.float32,
1899 + ),
1900 + ),
1901 + KVCacheGroupSpec(
1902 + ["layer1"],
1903 + FullAttentionSpec(
1904 + block_size=block_size,
1905 + num_kv_heads=2,
1906 + head_size=1,
1907 + dtype=torch.float32,
1908 + ),
1909 + is_eagle_group=True,
1910 + ),
1911 + ]
1912 +
1913 + runner = request_runner(
1914 + block_size=block_size,
1915 + num_gpu_blocks=num_gpu_blocks,
1916 + async_scheduling=async_scheduling,
1917 + kv_cache_groups=kv_cache_groups,
1918 + block_size_factor=block_size_factor,
1919 + )
1920 +
1921 + kv_group_configs = runner.connector_scheduler.config.kv_group_configs
1922 + assert len(kv_group_configs) == 2
1923 + assert not kv_group_configs[0].is_eagle_group
1924 + assert kv_group_configs[1].is_eagle_group
1925 +
1926 + runner.new_request(token_ids=[0] * offloaded_block_size * 3)
1927 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1928 + generate_store_output(keys)
1929 + )
1930 + runner.run(
1931 + decoded_tokens=[EOS_TOKEN_ID],
1932 + expected_stored=(
1933 + (0, 0),
1934 + (0, 1),
1935 + (0, 2),
1936 + (1, 0),
1937 + (1, 1),
1938 + ),
1939 + expected_flushed=(
1940 + (0, 0),
1941 + (0, 1),
1942 + (0, 2),
1943 + (1, 0),
1944 + (1, 1),
1945 + )
1946 + if not async_scheduling
1947 + else (),
1948 + )
1949 +
1950 + @pytest.mark.parametrize("async_scheduling", [True, False])
1951 + def test_sw_store_excludes_trailing_block(
1952 + self, request_runner, async_scheduling: bool
1953 + ):
1954 + """Eagle sliding-window group stores all blocks except the trailing
1955 + one."""
1956 + block_size = 4
1957 + sliding_window = 8
1958 + num_gpu_blocks = 100
1959 +
1960 + kv_cache_groups = [
1961 + KVCacheGroupSpec(
1962 + ["layer0"],
1963 + SlidingWindowSpec(
1964 + block_size=block_size,
1965 + num_kv_heads=1,
1966 + head_size=1,
1967 + dtype=torch.float32,
1968 + sliding_window=sliding_window,
1969 + ),
1970 + is_eagle_group=True,
1971 + ),
1972 + ]
1973 +
1974 + runner = request_runner(
1975 + block_size=block_size,
1976 + num_gpu_blocks=num_gpu_blocks,
1977 + async_scheduling=async_scheduling,
1978 + kv_cache_groups=kv_cache_groups,
1979 + )
1980 +
1981 + kv_group_configs = runner.connector_scheduler.config.kv_group_configs
1982 + assert len(kv_group_configs) == 1
1983 + assert kv_group_configs[0].is_eagle_group
1984 + assert kv_group_configs[0].sliding_window_size_in_blocks == 2
1985 +
1986 + runner.new_request(token_ids=[0] * block_size * 3)
1987 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1988 + generate_store_output(keys)
1989 + )
1990 + runner.run(
1991 + decoded_tokens=[EOS_TOKEN_ID],
1992 + expected_stored=((0, 0), (0, 1)),
1993 + expected_flushed=((0, 0), (0, 1)) if not async_scheduling else (),
1994 + )
1995 +
1996 + @pytest.mark.parametrize("async_scheduling", [True, False])
1997 + def test_single_block_nothing_stored(self, request_runner, async_scheduling: bool):
1998 + """An eagle group with only one block stores nothing: that block is
1999 + the tail."""
2000 + block_size = 4
2001 + block_size_factor = 1
2002 + offloaded_block_size = block_size * block_size_factor
2003 + num_gpu_blocks = 100
2004 +
2005 + kv_cache_groups = [
2006 + KVCacheGroupSpec(
2007 + ["layer0"],
2008 + FullAttentionSpec(
2009 + block_size=block_size,
2010 + num_kv_heads=1,
2011 + head_size=1,
2012 + dtype=torch.float32,
2013 + ),
2014 + is_eagle_group=True,
2015 + ),
2016 + ]
2017 +
2018 + runner = request_runner(
2019 + block_size=block_size,
2020 + num_gpu_blocks=num_gpu_blocks,
2021 + async_scheduling=async_scheduling,
2022 + kv_cache_groups=kv_cache_groups,
2023 + block_size_factor=block_size_factor,
2024 + )
2025 +
2026 + runner.new_request(token_ids=[0] * offloaded_block_size)
2027 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2028 + generate_store_output(keys)
2029 + )
2030 + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=())
2031 + runner.manager.prepare_store.assert_not_called()
2032 +
2033 + @pytest.mark.parametrize("async_scheduling", [True, False])
2034 + def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool):
2035 + """Eagle group constrains load: convergence tightens both groups.
2036 +
2037 + Store 3 offloaded blocks per group (eagle group skips tail → stores
2038 + 2). Then a new request loads from CPU. The eagle group's post-pop hit
2039 + (2) does not tighten below group 0's hit (3), so both groups load
2040 + normally.
2041 + """
2042 + block_size = 4
2043 + block_size_factor = 1
2044 + offloaded_block_size = block_size * block_size_factor
2045 + num_gpu_blocks = 100
2046 +
2047 + kv_cache_groups = [
2048 + KVCacheGroupSpec(
2049 + ["layer0"],
2050 + FullAttentionSpec(
2051 + block_size=block_size,
2052 + num_kv_heads=1,
2053 + head_size=1,
2054 + dtype=torch.float32,
2055 + ),
2056 + ),
2057 + KVCacheGroupSpec(
2058 + ["layer1"],
2059 + FullAttentionSpec(
2060 + block_size=block_size,
2061 + num_kv_heads=2,
2062 + head_size=1,
2063 + dtype=torch.float32,
2064 + ),
2065 + is_eagle_group=True,
2066 + ),
2067 + ]
2068 +
2069 + runner = request_runner(
2070 + block_size=block_size,
2071 + num_gpu_blocks=num_gpu_blocks,
2072 + async_scheduling=async_scheduling,
2073 + kv_cache_groups=kv_cache_groups,
2074 + block_size_factor=block_size_factor,
2075 + )
2076 +
2077 + runner.new_request(token_ids=[0] * offloaded_block_size * 3)
2078 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2079 + generate_store_output(keys)
2080 + )
2081 + runner.run(
2082 + decoded_tokens=[EOS_TOKEN_ID],
2083 + expected_stored=(
2084 + (0, 0),
2085 + (0, 1),
2086 + (0, 2),
2087 + (1, 0),
2088 + (1, 1),
2089 + ),
2090 + expected_flushed=(
2091 + (0, 0),
2092 + (0, 1),
2093 + (0, 2),
2094 + (1, 0),
2095 + (1, 1),
2096 + )
2097 + if not async_scheduling
2098 + else (),
2099 + )
2100 +
2101 + runner.scheduler.reset_prefix_cache()
2102 +
2103 + runner.new_request(token_ids=[0] * offloaded_block_size * 3 + [1])
2104 + runner.manager.lookup.return_value = True
2105 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2106 + generate_store_output([])
2107 + )
2108 + runner.run(
2109 + decoded_tokens=[EOS_TOKEN_ID],
2110 + expected_loaded=(
2111 + (0, 0),
2112 + (0, 1),
2113 + (1, 0),
2114 + (1, 1),
2115 + ),
2116 + )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+62 -5
@@ -77,6 +77,10 @@ class GroupOffloadConfig(NamedTuple):
77 # than the MLA full-attention group).
78 # None for full-attention groups or when the optimization doesn't apply.
79 alignment_block_count: int | None = None
80 + # True for EAGLE/MTP draft-model attention groups. The trailing block
81 + # of these groups is volatile and lacks a stable hash, so it must
82 + # be excluded from store and load scheduling.
83 + is_eagle_group: bool = False
84
85
86 def get_sliding_window_size_in_blocks(
@@ -155,6 +159,27 @@ class SchedulerOffloadConfig(NamedTuple):
159 return None
160 return per_segment
161
162 + eagle_groups = {
163 + idx
164 + for idx, g in enumerate(spec.kv_cache_config.kv_cache_groups)
165 + if g.is_eagle_group
166 + }
167 +
168 + use_eagle = (
169 + spec.vllm_config.speculative_config is not None
170 + and spec.vllm_config.speculative_config.use_eagle()
171 + )
172 + if use_eagle and not eagle_groups:
173 + eagle_groups = set(range(len(spec.kv_cache_config.kv_cache_groups)))
174 +
175 + if eagle_groups:
176 + logger.info(
177 + "KV offloading: EAGLE/MTP draft attention groups %s "
178 + "detected. The trailing block of these groups will be "
179 + "excluded from offloading due to volatility.",
180 + sorted(eagle_groups),
181 + )
182 +
183 return cls(
184 num_workers=spec.vllm_config.parallel_config.world_size,
185 kv_group_configs=tuple(
@@ -175,6 +200,7 @@ class SchedulerOffloadConfig(NamedTuple):
200 alignment_block_count=_alignment_block_count(
201 gpu_block_size * spec.block_size_factor, sw
202 ),
203 + is_eagle_group=idx in eagle_groups,
204 )
205 for idx, gpu_block_size in enumerate(spec.gpu_block_size)
206 ),
@@ -436,6 +462,11 @@ class OffloadingConnectorScheduler:
462 num_hit_tokens: int = 0
463 defer_lookup = False
464 lookup_groups = self._lookup_groups
465 +
466 + # Tracks which eagle groups have already popped their volatile trailing block
467 + # in the current convergence iteration. Reset when a non-eagle group
468 + # tightens the hit boundary, requiring a fresh pop.
469 + eagle_verified: set[int] = set()
470 while lookup_groups:
471 looked_up_sliding_window: bool = False
472 groups_iter = iter(lookup_groups)
@@ -453,6 +484,10 @@ class OffloadingConnectorScheduler:
484 >= req_status.req.num_tokens // offloaded_block_size
485 )
486
487 + is_eagle_unverified = (
488 + group_config.is_eagle_group and group_idx not in eagle_verified
489 + )
490 +
491 # Constrain to block-aligned boundary for this group
492 max_hit_size_tokens = min(
493 max_hit_size_tokens, len(offload_keys) * offloaded_block_size
@@ -461,14 +496,24 @@ class OffloadingConnectorScheduler:
496 # we can only load less than a block, better skip
497 return 0
498
499 + sliding_window_size_in_blocks = (
500 + group_config.sliding_window_size_in_blocks
501 + )
502 +
503 + # For eagle groups, query one extra block that will be popped.
504 + # We only need to increase the query size for sliding window groups.
505 + query_max = max_hit_size_tokens
506 + if is_eagle_unverified and sliding_window_size_in_blocks is not None:
507 + query_max = min(
508 + max_hit_size_tokens + offloaded_block_size,
509 + len(offload_keys) * offloaded_block_size,
510 + )
511 +
512 num_blocks = min(
465 - cdiv(max_hit_size_tokens, offloaded_block_size), len(offload_keys)
513 + cdiv(query_max, offloaded_block_size), len(offload_keys)
514 )
515 start_block_idx = num_computed_tokens // offloaded_block_size
516 offload_keys = offload_keys[start_block_idx:num_blocks]
469 - sliding_window_size_in_blocks = (
470 - group_config.sliding_window_size_in_blocks
471 - )
517
518 # end index (in the sliced offload_keys) up to which we
519 # have backend-confirmed hits
@@ -478,9 +523,12 @@ class OffloadingConnectorScheduler:
523 offload_keys, req_status.req_context
524 )
525 else:
526 + required_window = sliding_window_size_in_blocks
527 + if is_eagle_unverified:
528 + required_window += 1
529 num_hit_blocks = self._sliding_window_lookup(
530 offload_keys,
483 - sliding_window_size_in_blocks,
531 + required_window,
532 req_status.req_context,
533 )
534 if num_hit_blocks == 0:
@@ -489,6 +537,10 @@ class OffloadingConnectorScheduler:
537 if num_hit_blocks is None:
538 defer_lookup = True
539 else:
540 + if is_eagle_unverified:
541 + num_hit_blocks -= 1
542 + eagle_verified.add(group_idx)
543 +
544 max_hit_size_tokens = min(
545 max_hit_size_tokens,
546 offloaded_block_size * (start_block_idx + num_hit_blocks),
@@ -500,6 +552,8 @@ class OffloadingConnectorScheduler:
552 return 0
553
554 if new_num_hit_tokens < num_hit_tokens:
555 + if not group_config.is_eagle_group:
556 + eagle_verified.clear()
557 if defer_lookup:
558 # make another iteration on all groups to check
559 # if we still need to defer lookup
@@ -791,6 +845,9 @@ class OffloadingConnectorScheduler:
845 self.config.kv_group_configs, req_status.group_states
846 ):
847 num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
848 + if group_config.is_eagle_group:
849 + num_blocks = max(0, num_blocks - 1)
850 +
851 start_block_idx = group_state.next_stored_block_idx
852 if num_blocks <= start_block_idx:
853 continue