[FEAT] Capture fixed-depth speculative decode with CUDA Graphs

tests/test_speculative.py

@@ -302,7 +302,6 @@ def test_probability_filters_match_expected_and_keep_boundary_ties():
     [
         dict(cache_config=CacheConfig(layout="shared")),
         dict(execution_config=ExecutionConfig(async_scheduling=True)),
-        dict(execution_config=ExecutionConfig(cuda_graphs=True)),
         dict(scheduler_config=SchedulerConfig(enable_preemption=True)),
         dict(scheduler_config=SchedulerConfig(mode="no_refill")),
     ],
@@ -375,3 +374,38 @@ def run(speculative, sampling):
     greedy = [SamplingParams(max_tokens=n, ignore_eos=True) for n in [8, 5, 3]]
     a, b = run(True, greedy), run(False, greedy)
     assert [o.token_ids for o in a] == [o.token_ids for o in b]
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("max_graph_rows", [2, 8])
+def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows):
+    m = model(dtype=torch.bfloat16).to("cuda")
+    prompts = [[2, 3, 4], [7, 8]]
+    params = SamplingParams(max_tokens=9, ignore_eos=True)
+
+    def run(graphs):
+        llm = LLM(
+            m,
+            speculative_config=SpeculativeConfig(3),
+            cache_config=CacheConfig(128, 16),
+            attention_backend="triton",
+            execution_config=ExecutionConfig(
+                cuda_graphs=graphs, cuda_graph_max_batch_size=max_graph_rows
+            ),
+        )
+        outputs = [llm.generate(prompts, params) for _ in range(2)]
+        assert llm.engine.cache_manager.num_used_blocks == 0
+        return llm.engine.speculative_runner, [
+            [(out.token_ids, out.exit_depths) for out in batch] for batch in outputs
+        ]
+
+    _, eager = run(False)
+    runner, graphed = run(True)
+    assert graphed == eager
+    assert runner.graphs.captures > 0
+    assert runner.graphs.replays > runner.graphs.captures
+    assert runner.coda_graphs.captures > 0
+    if max_graph_rows == 2:
+        assert runner.graphs.fallbacks > 0
+    else:
+        assert any(rows > 2 for rows, _, _, _ in runner.graphs.entries)

vllm_rlt/engine/llm_engine.py

@@ -41,10 +41,8 @@ def __init__(
                 raise ValueError("speculative target_loops must equal the model full depth")
             if self.exit_config.mode != "ouro":
                 raise ValueError("speculative decoding requires fixed-depth ouro exit mode")
-            if self.execution_config.async_scheduling or self.execution_config.cuda_graphs:
-                raise ValueError(
-                    "speculative decoding currently requires synchronous eager execution"
-                )
+            if self.execution_config.async_scheduling:
+                raise ValueError("speculative decoding currently requires synchronous execution")
             if scheduler_config.enable_preemption or scheduler_config.mode != "refill":
                 raise ValueError(
                     "speculative decoding requires refill scheduling without preemption"
@@ -98,7 +96,7 @@ def __init__(
             raise ValueError("prefill_uva requires CUDA FA4 with last_exited KV")
         self.scheduler = Scheduler(scheduler_config, self.cache_manager, speculative_config)
         self.speculative_runner = (
-            SpeculativeRunner(model, self.cache_manager, speculative_config)
+            SpeculativeRunner(model, self.cache_manager, speculative_config, self.execution_config)
             if speculative_config is not None
             else None
         )

vllm_rlt/worker/cuda_graph.py

@@ -7,6 +7,19 @@
 import torch
 
 
+def _capture(hidden, stream, pool, run):
+    """Warm and capture one stable-input graph on its own stream."""
+    torch.cuda.synchronize(hidden.device)
+    with torch.cuda.stream(stream):
+        for _ in range(2):
+            run()
+    stream.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, pool=pool, stream=stream):
+        output = run()
+    return graph, output
+
+
 class _DeviceCache:
     """Capture only GPU cache operations, never Python allocation/written state."""
 
@@ -20,6 +33,16 @@ def _write_prepared(self, layer, batch, k, v):
 
     def _attend_prepared(self, layer, batch, q):
         cache = self.cache
+        if batch.cu_seqlens_q is not None:
+            return cache.attention.prefill(
+                q,
+                cache.key_cache[:, layer],
+                cache.value_cache[:, layer],
+                batch.block_tables,
+                batch.context_lengths,
+                batch.cu_seqlens_q,
+                batch.max_seqlen_q,
+            )
         return cache.attention(
             q,
             cache.key_cache[:, layer],
@@ -42,6 +65,8 @@ def __init__(self, model, cache, execution, compute_gate):
     @torch.inference_mode()
     def run(self, hidden, batch):
         count = len(batch.rows)
+        packed = batch.cu_seqlens_q is not None
+        key = (count, len(batch.context_lengths), batch.max_seqlen_q if packed else 1, packed)
         if len(batch.position_ids) != count:
             batch = replace(
                 batch,
@@ -52,9 +77,8 @@ def run(self, hidden, batch):
         if (
             count > self.execution.cuda_graph_max_batch_size
             or not count
-            or batch.cu_seqlens_q is not None
             or (
-                count not in self.entries
+                key not in self.entries
                 and len(self.entries) >= self.execution.cuda_graph_max_graphs
             )
         ):
@@ -64,51 +88,67 @@ def run(self, hidden, batch):
             )
         cache = self.cache
         cache._require_live_batch(batch)
-        # Decode writes exactly one position per request/depth. Validate history
-        # before launch; capture/replay must not freeze or repeat host mutations.
+        # A speculative verification batch may write consecutive positions at
+        # one depth. Its first position needs committed history; later rows are
+        # written together before attention reads them.
+        frontier = {}
         for allocation, depth, pos in batch.rows:
-            for layer in range(cache.num_layers):
-                cache._require_prefix(allocation, layer, depth, pos)
+            row_key = (id(allocation), depth)
+            previous = frontier.get(row_key)
+            if previous is None:
+                for layer in range(cache.num_layers):
+                    cache._require_prefix(allocation, layer, depth, pos)
+            elif pos != previous + 1:
+                self.fallbacks += 1
+                return self.model.recurrent_prepared(
+                    hidden, batch, cache, compute_gate=self.compute_gate
+                )
+            frontier[row_key] = pos
         stream = torch.cuda.current_stream(cache.device)
         if self.last_event is not None:
             stream.wait_event(self.last_event)
-        entry = self.entries.get(count)
+        entry = self.entries.get(key)
         if entry is None:
             width = math.ceil(self.model.config.max_position_embeddings / cache.block_size)
+            tables = len(batch.context_lengths)
             entry = SimpleNamespace(
                 hidden=torch.empty_like(hidden[:count]),
                 metadata=SimpleNamespace(
                     position_ids=torch.empty(count, device=cache.device, dtype=torch.long),
                     write_blocks=torch.empty(count, device=cache.device, dtype=torch.long),
                     write_offsets=torch.empty(count, device=cache.device, dtype=torch.long),
                     block_tables=torch.zeros(
-                        (count, width), device=cache.device, dtype=torch.int32
+                        (tables, width), device=cache.device, dtype=torch.int32
                     ),
-                    context_lengths=torch.empty(count, device=cache.device, dtype=torch.int32),
+                    context_lengths=torch.empty(tables, device=cache.device, dtype=torch.int32),
+                    cu_seqlens_q=(
+                        torch.empty(tables + 1, device=cache.device, dtype=torch.int32)
+                        if packed
+                        else None
+                    ),
+                    max_seqlen_q=batch.max_seqlen_q,
                 ),
             )
         entry.hidden.copy_(hidden[:count])
-        for name in ("position_ids", "write_blocks", "write_offsets", "context_lengths"):
-            getattr(entry.metadata, name).copy_(getattr(batch, name)[:count])
+        entry.metadata.position_ids.copy_(batch.position_ids[:count])
+        entry.metadata.write_blocks.copy_(batch.write_blocks[:count])
+        entry.metadata.write_offsets.copy_(batch.write_offsets[:count])
+        entry.metadata.context_lengths.copy_(batch.context_lengths)
+        if packed:
+            entry.metadata.cu_seqlens_q.copy_(batch.cu_seqlens_q)
         width = batch.block_tables.shape[1]
-        entry.metadata.block_tables[:, :width].copy_(batch.block_tables[:count])
-        if count not in self.entries:
-            # First-use capture is intentionally outside steady-state timing.
-            # Drain outstanding streams: CUDA capture cannot race other launches.
-            torch.cuda.synchronize(cache.device)
+        entry.metadata.block_tables[:, :width].copy_(batch.block_tables)
+        if key not in self.entries:
             proxy = _DeviceCache(cache)
-            with torch.cuda.stream(self.capture_stream):
-                for _ in range(2):
-                    self.model.recurrent_prepared(
-                        entry.hidden, entry.metadata, proxy, compute_gate=self.compute_gate
-                    )
-            self.capture_stream.synchronize()
-            entry.graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(entry.graph, pool=self.pool, stream=self.capture_stream):
-                entry.output = self.model.recurrent_prepared(
+            entry.graph, entry.output = _capture(
+                entry.hidden,
+                self.capture_stream,
+                self.pool,
+                lambda: self.model.recurrent_prepared(
                     entry.hidden, entry.metadata, proxy, compute_gate=self.compute_gate
-                )
-            self.entries[count] = entry
+                ),
+            )
+            self.entries[key] = entry
             self.captures += 1
         entry.graph.replay()
         # Outputs must outlive reuse of graph-private buffers and the shared pool.
@@ -120,3 +160,36 @@ def run(self, hidden, batch):
                 allocation.written[cache._plane(depth)][layer].add(pos)
         self.replays += 1
         return output
+
+
+class CodaGraphs:
+    """Capture fixed-row LM heads while keeping sampling and stop decisions eager."""
+
+    def __init__(self, model, execution):
+        self.model, self.execution = model, execution
+        self.entries = {}
+        self.pool = torch.cuda.graph_pool_handle()
+        self.stream = torch.cuda.Stream(device=next(model.parameters()).device)
+        self.captures = self.replays = self.fallbacks = 0
+
+    @torch.inference_mode()
+    def run(self, hidden):
+        count = len(hidden)
+        if count > self.execution.cuda_graph_max_batch_size or (
+            count not in self.entries and len(self.entries) >= self.execution.cuda_graph_max_graphs
+        ):
+            self.fallbacks += 1
+            return self.model.coda(hidden)
+        entry = self.entries.get(count)
+        if entry is None:
+            entry = SimpleNamespace(hidden=torch.empty_like(hidden))
+        entry.hidden.copy_(hidden)
+        if count not in self.entries:
+            entry.graph, entry.output = _capture(
+                entry.hidden, self.stream, self.pool, lambda: self.model.coda(entry.hidden)
+            )
+            self.entries[count] = entry
+            self.captures += 1
+        entry.graph.replay()
+        self.replays += 1
+        return entry.output.clone()

vllm_rlt/worker/speculative.py

@@ -8,6 +8,7 @@
 
 import torch
 
+from vllm_rlt.worker.cuda_graph import CodaGraphs, RecurrentGraphs
 from vllm_rlt.worker.sampling import (
     draw,
     generator_for,
@@ -33,21 +34,35 @@ class SpeculativeStats:
 
 
 class SpeculativeRunner:
-    def __init__(self, model, cache, config):
+    def __init__(self, model, cache, config, execution):
         self.model, self.cache, self.config = model, cache, config
         self.device = next(model.parameters()).device
         self.stats = SpeculativeStats()
+        self.graphs = (
+            RecurrentGraphs(model, cache, execution, False) if execution.cuda_graphs else None
+        )
+        self.coda_graphs = CodaGraphs(model, execution) if execution.cuda_graphs else None
 
     def _core(self, hidden, ids, positions, depth, *, packed=False):
         batch = self.cache._prepare_batch(
             ids,
             [depth] * len(ids),
             positions,
-            packed_prefill=packed and getattr(self.cache.attention, "generation", None) == 4,
+            packed_prefill=packed and self.cache.attention_info.get("generation") == 4,
         )
-        hidden, _ = self.model.recurrent_prepared(hidden, batch, self.cache, compute_gate=False)
+        if self.graphs is not None:
+            hidden, _ = self.graphs.run(hidden, batch)
+        else:
+            hidden, _ = self.model.recurrent_prepared(hidden, batch, self.cache, compute_gate=False)
         return hidden
 
+    def _coda(self, hidden):
+        return (
+            self.coda_graphs.run(hidden)
+            if self.coda_graphs is not None
+            else self.model.coda(hidden)
+        )
+
     @torch.inference_mode()
     def execute(self, batch):
         items = batch.items
@@ -74,7 +89,7 @@ def execute(self, batch):
                     drafting.append((row, i))
             if not drafting:
                 continue
-            logits = self.model.coda(torch.stack([hidden[row] for row, _ in drafting]))
+            logits = self._coda(torch.stack([hidden[row] for row, _ in drafting]))
             for row, (_, i) in enumerate(drafting):
                 request = items[i].request
                 if request.sampling_params.temperature == 0:
@@ -92,7 +107,7 @@ def execute(self, batch):
         hidden = torch.stack([h for request_states in states for h in request_states])
         for depth in range(self.config.draft_loops, self.config.target_loops):
             hidden = self._core(hidden, ids, positions, depth, packed=True)
-        logits = self.model.coda(hidden)
+        logits = self._coda(hidden)
         results, start = [], 0
         for i, item in enumerate(items):
             rows = logits[start : start + item.token_count]