[FEAT] [PERF] Remove per-candidate and metadata host syncs from speculative decoding

tests/test_prepared_kv.py

@@ -2,6 +2,7 @@
 
 import math
 import weakref
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -273,10 +274,47 @@ def forbidden(*args, **kwargs):
     finally:
         for handle in handles:
             handle.remove()
-    assert metadata_counts == [5] * 4
+    # Two staged host tensors (int64 addresses, int32 tables/lengths) per core traversal.
+    assert metadata_counts == [2] * 4
     assert len(layer_ids) == 96
     for traversal in range(4):
         group = layer_ids[traversal * 24 : (traversal + 1) * 24]
         assert len(set(group)) == 1
         assert group[0][1] == rotary_ids[traversal]
     assert all(reference() is None for reference in preparations)
+
+
+@pytest.mark.parametrize("packed", [False, True])
+def test_depth_batches_share_depth_independent_metadata(packed):
+    cache = make_cache(num_blocks=64)
+    if packed:
+        cache.attention = SimpleNamespace(generation=4)
+    assert cache.allocate("a", 7) and cache.allocate("b", 5)
+    ids = ["a"] * 3 + ["b"] * 2
+    positions = [4, 5, 6, 2, 3]
+    depth_sets = [[depth] * len(ids) for depth in range(4)]
+    batches = cache._prepare_batches(ids, depth_sets, positions, packed_prefill=packed)
+    for depths, batch in zip(depth_sets, batches):
+        single = cache._prepare_batch(ids, depths, positions, packed_prefill=packed)
+        assert batch.rows == single.rows and batch.allocations == single.allocations
+        assert batch.max_seqlen_q == single.max_seqlen_q
+        for name in (
+            "position_ids",
+            "write_blocks",
+            "write_offsets",
+            "block_tables",
+            "context_lengths",
+            "cu_seqlens_q",
+        ):
+            ours, reference = getattr(batch, name), getattr(single, name)
+            assert (ours is None) == (reference is None)
+            if ours is not None:
+                assert torch.equal(ours, reference), name
+    # Each depth addresses its own plane; depth-independent metadata is one copy.
+    assert len({tuple(batch.write_blocks.tolist()) for batch in batches}) == 4
+    for name in ("position_ids", "write_offsets", "context_lengths"):
+        assert len({getattr(batch, name).data_ptr() for batch in batches}) == 1
+    with pytest.raises(ValueError, match="depth"):
+        cache._prepare_batches(ids, [[0] * 5, [4] * 5], positions)
+    with pytest.raises(ValueError, match="equal lengths"):
+        cache._prepare_batches(ids, [[0] * 5, [1] * 4], positions)

tests/test_speculative.py

@@ -1,5 +1,6 @@
 """Fixed-depth oracle, rollback, lifecycle and distribution tests for speculation."""
 
+import warnings
 from collections import Counter
 
 import pytest
@@ -19,6 +20,7 @@
 from vllm_rlt.models.reference import dense_reference
 from vllm_rlt.request import Request, Stage
 from vllm_rlt.worker.sampling import probabilities, rejection_sample
+from vllm_rlt.worker.speculative import greedy_accept
 
 
 def model(seed=123, dtype=torch.float32):
@@ -192,6 +194,17 @@ def controlled(h):
     assert final.token_ids[len(out.token_ids) :] == expected.token_ids
 
 
+@pytest.mark.parametrize("k", [0, 1, 2, 4, 8])
+def test_greedy_accept_commits_target_prefix_through_first_mismatch(k):
+    candidates = list(range(100, 100 + k))
+    for rejected in range(k + 1):
+        # Rows before `rejected` agree; row `rejected` is the correction, or the bonus.
+        targets = candidates[:rejected] + [7] + list(range(200, 200 + k - rejected))
+        tokens, accepted = greedy_accept(candidates, targets)
+        assert accepted == rejected
+        assert tokens == candidates[:rejected] + [7]
+
+
 def test_eos_in_accepted_prefix_never_delivers_following_tokens(monkeypatch):
     m = model()
     e = engine(m)
@@ -376,6 +389,32 @@ def run(speculative, sampling):
     assert [o.token_ids for o in a] == [o.token_ids for o in b]
 
 
+@pytest.mark.gpu
+def test_greedy_round_reads_back_once_with_ragged_requests():
+    # Draft IDs stay on the device; every speculative round, including ragged and
+    # budget-shortened ones, synchronizes with the host exactly once.
+    torch.manual_seed(123)
+    m = OuroForCausalLM(OuroConfig.tiny(hidden_size=256, head_dim=64)).to(
+        device="cuda", dtype=torch.bfloat16
+    )
+    e = engine(m, k=3, cache_config=CacheConfig(128, 16), attention_backend="triton")
+    for rid, prompt, n in [("a", [2, 3, 4], 9), ("b", [7, 8], 6), ("c", [9], 4)]:
+        e.add_request(rid, prompt, SamplingParams(max_tokens=n, ignore_eos=True))
+    rounds = 0
+    while e.has_unfinished_requests():
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            torch.cuda.set_sync_debug_mode("warn")
+            try:
+                e.step()
+            finally:
+                torch.cuda.set_sync_debug_mode("default")
+        if e.last_schedule.stage == Stage.SPECULATIVE:
+            rounds += 1
+            assert sum("synchroniz" in str(w.message) for w in caught) == 1
+    assert rounds > 1
+
+
 @pytest.mark.gpu
 @pytest.mark.parametrize("max_graph_rows", [2, 8])
 def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows):

vllm_rlt/core/kv_cache_manager.py

@@ -425,77 +425,138 @@ def _prepare_batch(
         packed_prefill: bool = False,
     ) -> _PreparedKVBatch:
         """Build layer-independent addresses once; do not initialize any KV slot."""
-        rows = tuple(self._validate_rows(request_ids, depths, positions))
-        addresses = [
-            (
-                allocation.block_tables[self._plane(depth)][position // self.block_size],
-                position % self.block_size,
+        return self._prepare_batches(
+            request_ids,
+            (depths,),
+            positions,
+            for_write=for_write,
+            packed_prefill=packed_prefill,
+        )[0]
+
+    def _prepare_batches(
+        self,
+        request_ids: Sequence[str],
+        depth_sets: Sequence[Sequence[int]],
+        positions: Sequence[int] | torch.Tensor,
+        *,
+        for_write: bool = True,
+        packed_prefill: bool = False,
+    ) -> tuple[_PreparedKVBatch, ...]:
+        """Prepare one batch per depth assignment of the same request/position rows.
+
+        Positions and write offsets do not depend on depth; they are built and
+        copied once and shared as read-only views, as are context lengths and
+        query boundaries when the per-depth grouping is identical. Block addresses
+        and page tables are built per depth assignment. No KV slot is initialized:
+        ownership and written-prefix checks still run when each batch executes.
+        """
+        if not depth_sets:
+            raise ValueError("at least one depth assignment is required")
+        if packed_prefill and (
+            self.layout != "last_exited" or getattr(self.attention, "generation", None) != 4
+        ):
+            raise ValueError("packed prefill requires LAST_EXITED and FlashAttention-4")
+        first = tuple(self._validate_rows(request_ids, depth_sets[0], positions))
+        row_sets = [first]
+        for depths in depth_sets[1:]:
+            if len(depths) != len(first):
+                raise ValueError("request_ids, depths, and positions must have equal lengths")
+            for depth in depths:
+                self._validate_depth(depth)
+            row_sets.append(
+                tuple(
+                    (allocation, int(depth), position)
+                    for (allocation, _, position), depth in zip(first, depths)
+                )
             )
-            for allocation, depth, position in rows
-        ]
-        if for_write and len(set(addresses)) != len(addresses):
-            raise ValueError(
-                "a write batch cannot contain duplicate request/depth/position addresses"
+        n = len(first)
+        position_list = [position for _, _, position in first]
+        offsets = [position % self.block_size for position in position_list]
+        width = max((position // self.block_size + 1 for position in position_list), default=0)
+        # Stage all metadata in two pinned host tensors and copy them without
+        # blocking: a pageable H2D copy would wait for all queued GPU work. The
+        # caching host allocator keeps each staging block alive until its copy ends.
+        wide = position_list + offsets
+        narrow = []
+        shared = {}
+        plans = []
+        for rows in row_sets:
+            blocks = [
+                allocation.block_tables[self._plane(depth)][position // self.block_size]
+                for allocation, depth, position in rows
+            ]
+            if for_write and len(set(zip(blocks, offsets))) != n:
+                raise ValueError(
+                    "a write batch cannot contain duplicate request/depth/position addresses"
+                )
+            table_rows, cumulative, max_query = (
+                self._group_packed_rows(rows) if packed_prefill else (rows, None, 1)
             )
-        width = max((position // self.block_size + 1 for _, _, position in rows), default=0)
-        table_rows = rows
-        cumulative = None
-        max_query = 1
-        if packed_prefill:
-            if self.layout != "last_exited" or getattr(self.attention, "generation", None) != 4:
-                raise ValueError("packed prefill requires LAST_EXITED and FlashAttention-4")
-            # Group only consecutive positions of one request/depth. The last
-            # position supplies the causal key length for the whole query chunk.
-            ends, cumulative, seen = [], [0], set()
-            for index, (allocation, depth, position) in enumerate(rows):
-                key = (id(allocation), depth)
-                if index and key == (id(rows[index - 1][0]), rows[index - 1][1]):
-                    if position != rows[index - 1][2] + 1:
-                        raise ValueError("packed prefill positions must be contiguous")
-                    ends[-1] = (allocation, depth, position)
-                    cumulative[-1] = index + 1
-                else:
-                    if key in seen:
-                        raise ValueError("packed prefill request/depth must form one sequence")
-                    seen.add(key)
-                    ends.append((allocation, depth, position))
-                    cumulative.append(index + 1)
-            table_rows = ends
-            max_query = max((b - a for a, b in zip(cumulative, cumulative[1:])), default=1)
-        tables = []
-        for allocation, depth, _ in table_rows:
-            table = allocation.block_tables[self._plane(depth)][:width]
-            tables.append(list(table) + [-1] * (width - len(table)))
-        allocations = dict(zip(request_ids, (allocation for allocation, _, _ in rows)))
-        return _PreparedKVBatch(
-            owner=self,
-            rows=rows,
-            allocations=tuple(allocations.items()),
-            position_ids=torch.tensor(
-                [position for _, _, position in rows], device=self.device, dtype=torch.long
-            ),
-            write_blocks=torch.tensor(
-                [block for block, _ in addresses], device=self.device, dtype=torch.long
-            ),
-            write_offsets=torch.tensor(
-                [offset for _, offset in addresses], device=self.device, dtype=torch.long
-            ),
-            block_tables=torch.tensor(tables, device=self.device, dtype=torch.int32).reshape(
-                len(table_rows), width
-            ),
-            context_lengths=torch.tensor(
-                [position + 1 for _, _, position in table_rows],
-                device=self.device,
-                dtype=torch.int32,
-            ),
-            writable=for_write,
-            cu_seqlens_q=(
-                torch.tensor(cumulative, device=self.device, dtype=torch.int32)
-                if cumulative is not None
-                else None
-            ),
-            max_seqlen_q=max_query,
-        )
+            block_start = len(wide)
+            wide += blocks
+            table_start = len(narrow)
+            for allocation, depth, _ in table_rows:
+                table = allocation.block_tables[self._plane(depth)][:width]
+                narrow += table
+                narrow += [-1] * (width - len(table))
+            lengths = [position + 1 for _, _, position in table_rows]
+            key = (tuple(lengths), None if cumulative is None else tuple(cumulative))
+            if key not in shared:
+                shared[key] = len(narrow)
+                narrow += lengths + (cumulative or [])
+            packed = cumulative is not None
+            plans.append(
+                (rows, block_start, table_start, len(table_rows), shared[key], packed, max_query)
+            )
+        wide = self._stage(wide, torch.long)
+        narrow = self._stage(narrow, torch.int32)
+        allocations = tuple(dict(zip(request_ids, (a for a, _, _ in first))).items())
+        batches = []
+        for rows, block_start, table_start, t, lengths_start, packed, max_query in plans:
+            batches.append(
+                _PreparedKVBatch(
+                    owner=self,
+                    rows=rows,
+                    allocations=allocations,
+                    position_ids=wide[:n],
+                    write_blocks=wide[block_start : block_start + n],
+                    write_offsets=wide[n : 2 * n],
+                    block_tables=narrow[table_start : table_start + t * width].reshape(t, width),
+                    context_lengths=narrow[lengths_start : lengths_start + t],
+                    writable=for_write,
+                    cu_seqlens_q=(
+                        narrow[lengths_start + t : lengths_start + 2 * t + 1] if packed else None
+                    ),
+                    max_seqlen_q=max_query,
+                )
+            )
+        return tuple(batches)
+
+    def _group_packed_rows(self, rows):
+        """Group consecutive positions of one request/depth into packed query chunks.
+
+        The last position supplies the causal key length for the whole chunk.
+        """
+        ends, cumulative, seen = [], [0], set()
+        for index, (allocation, depth, position) in enumerate(rows):
+            key = (id(allocation), depth)
+            if index and key == (id(rows[index - 1][0]), rows[index - 1][1]):
+                if position != rows[index - 1][2] + 1:
+                    raise ValueError("packed prefill positions must be contiguous")
+                ends[-1] = (allocation, depth, position)
+                cumulative[-1] = index + 1
+            else:
+                if key in seen:
+                    raise ValueError("packed prefill request/depth must form one sequence")
+                seen.add(key)
+                ends.append((allocation, depth, position))
+                cumulative.append(index + 1)
+        max_query = max((b - a for a, b in zip(cumulative, cumulative[1:])), default=1)
+        return ends, cumulative, max_query
+
+    def _stage(self, values, dtype):
+        host = torch.tensor(values, dtype=dtype, pin_memory=self.device.type == "cuda")
+        return host.to(self.device, non_blocking=True)
 
     def _require_live_batch(self, batch: _PreparedKVBatch) -> None:
         if batch.owner is not self: