From aaa31de8a847ef234126c697d6a2e15e13decad0 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Fri, 9 Oct 2026 13:53:52 +0000
Subject: [PATCH 01/20] feat: add opt-in loop-level speculative pipeline
---
benchmarks/loop_speculative.py | 180 ++++++++++++++++++++++
docs/README.md | 1 +
docs/loop_speculative.md | 140 +++++++++++++++++
tests/test_loop_speculative.py | 185 +++++++++++++++++++++++
tests/test_speculative.py | 87 +++++++----
vllm_rlt/config.py | 3 +
vllm_rlt/entrypoints/runtime_args.py | 8 +
vllm_rlt/worker/speculative.py | 217 ++++++++++++++++++++++++++-
8 files changed, 788 insertions(+), 33 deletions(-)
create mode 100644 benchmarks/loop_speculative.py
create mode 100644 docs/loop_speculative.md
create mode 100644 tests/test_loop_speculative.py
diff --git a/benchmarks/loop_speculative.py b/benchmarks/loop_speculative.py
new file mode 100644
index 00000000..993e5423
--- /dev/null
+++ b/benchmarks/loop_speculative.py
@@ -0,0 +1,180 @@
+"""Matched synchronous baseline/pipeline feasibility check on a local Ouro model.
+
+Exclude prefill, interleave trials, and retain raw timings and generated tokens.
+This small experiment is not the full RFC #43 accuracy/performance acceptance matrix.
+"""
+
+import argparse
+import hashlib
+import json
+import statistics
+import subprocess
+import time
+from dataclasses import asdict
+from pathlib import Path
+
+import torch
+
+from vllm_rlt import (
+ CacheConfig,
+ ExecutionConfig,
+ SamplingParams,
+ SchedulerConfig,
+ SpeculativeConfig,
+)
+from vllm_rlt.engine.llm_engine import LLMEngine
+from vllm_rlt.models import OuroForCausalLM, load_tokenizer
+from vllm_rlt.request import Stage
+from vllm_rlt.worker.speculative import SpeculativeStats
+
+
+@torch.inference_mode()
+def trial(engine, prompts, new_tokens):
+ params = SamplingParams(max_tokens=new_tokens + 1, ignore_eos=True)
+ for i, prompt in enumerate(prompts):
+ engine.add_request(str(i), prompt, params)
+ # All prompts fit one prefill batch, so no decode can precede this barrier.
+ while any(not r.generated_token_ids for r in engine.scheduler.requests.values()):
+ engine.step()
+ assert engine.last_schedule.stage != Stage.SPECULATIVE
+ engine.speculative_runner.stats = SpeculativeStats()
+ torch.cuda.synchronize()
+ torch.cuda.reset_peak_memory_stats()
+ start = time.perf_counter()
+ outputs = {}
+ while engine.has_unfinished_requests():
+ for output in engine.step():
+ if output.finished:
+ outputs[output.request_id] = output.token_ids
+ torch.cuda.synchronize()
+ elapsed = time.perf_counter() - start
+ committed = len(prompts) * new_tokens
+ stats = asdict(engine.speculative_runner.stats)
+ assert stats["committed_tokens"] == committed
+ return dict(
+ seconds=elapsed,
+ committed_tokens=committed,
+ tokens_per_second=committed / elapsed,
+ peak_allocated_bytes=torch.cuda.max_memory_allocated(),
+ stats=stats,
+ tokens=[outputs[str(i)] for i in range(len(prompts))],
+ )
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--model", type=Path, required=True, help="Existing local checkpoint")
+ parser.add_argument("--prompts", type=Path, help="JSON list of prompt strings")
+ parser.add_argument("--output", type=Path, required=True)
+ parser.add_argument("--backend", choices=["triton", "flash_attn_4"], default="triton")
+ parser.add_argument("--cuda-graphs", action="store_true")
+ parser.add_argument("--k", type=int, default=4)
+ parser.add_argument("--draft-loops", type=int, default=2)
+ parser.add_argument("--concurrency", type=int, default=1)
+ parser.add_argument("--new-tokens", type=int, default=32)
+ parser.add_argument("--repeats", type=int, default=5)
+ args = parser.parse_args()
+ if not args.model.is_dir():
+ parser.error("--model must be a complete local checkpoint directory")
+ if min(args.concurrency, args.new_tokens, args.repeats, args.k) < 1:
+ parser.error("concurrency, new-tokens, repeats and k must be positive")
+ torch.set_num_threads(1)
+ texts = (
+ json.loads(args.prompts.read_text())
+ if args.prompts
+ else [
+ "Explain why the sky appears blue during the day.",
+ "Write a Python function that returns the Fibonacci sequence.",
+ ]
+ )
+ if not texts or not all(isinstance(t, str) and t for t in texts):
+ parser.error("prompts must be a nonempty JSON list of nonempty strings")
+ tokenizer = load_tokenizer(str(args.model))
+ prompts = [tokenizer.encode(texts[i % len(texts)]) for i in range(args.concurrency)]
+ model = OuroForCausalLM.from_pretrained(args.model, device="cuda", dtype=torch.bfloat16)
+ blocks = model.config.total_ut_steps * sum(
+ (len(p) + args.new_tokens + 15) // 16 for p in prompts
+ )
+ scheduler = SchedulerConfig(
+ max_num_seqs=args.concurrency,
+ max_num_batched_tokens=max(sum(map(len, prompts)), args.concurrency * (args.k + 1)),
+ prefill_chunk_size=max(map(len, prompts)),
+ )
+ execution = ExecutionConfig(cuda_graphs=args.cuda_graphs)
+ engines = {
+ name: LLMEngine(
+ model,
+ cache_config=CacheConfig(num_blocks=blocks, block_size=16),
+ scheduler_config=scheduler,
+ execution_config=execution,
+ attention_backend=args.backend,
+ speculative_config=SpeculativeConfig(
+ args.k, args.draft_loops, model.config.total_ut_steps, enabled
+ ),
+ )
+ for name, enabled in [("baseline", False), ("pipeline", True)]
+ }
+ raw = {name: [] for name in engines}
+ try:
+ for engine in engines.values():
+ trial(engine, prompts, args.new_tokens)
+ for repeat in range(args.repeats):
+ order = list(engines) if repeat % 2 == 0 else list(reversed(engines))
+ for name in order:
+ result = trial(engines[name], prompts, args.new_tokens)
+ raw[name].append(result)
+ print(
+ f"{name} repeat={repeat}: {result['tokens_per_second']:.2f} tokens/s",
+ flush=True,
+ )
+ finally:
+ for engine in engines.values():
+ engine.close()
+ summary = {
+ name: dict(
+ median_seconds=statistics.median(r["seconds"] for r in trials),
+ min_seconds=min(r["seconds"] for r in trials),
+ max_seconds=max(r["seconds"] for r in trials),
+ )
+ for name, trials in raw.items()
+ }
+ agreement = []
+ for baseline, pipeline in zip(raw["baseline"], raw["pipeline"]):
+ agreement.append(
+ [
+ next((i for i, (a, b) in enumerate(zip(left, right)) if a != b), None)
+ for left, right in zip(baseline["tokens"], pipeline["tokens"])
+ ]
+ )
+ report = dict(
+ arguments={k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()},
+ git_head=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(),
+ tracked_diff_sha256=hashlib.sha256(
+ subprocess.check_output(["git", "diff", "HEAD"])
+ ).hexdigest(),
+ torch=torch.__version__,
+ cuda=torch.version.cuda,
+ gpu=torch.cuda.get_device_name(),
+ dtype="bfloat16",
+ checkpoint_config_sha256=hashlib.sha256(
+ (args.model / "config.json").read_bytes()
+ ).hexdigest(),
+ prompt_token_ids=prompts,
+ scheduler=asdict(scheduler),
+ execution=asdict(execution),
+ summary=summary,
+ speedup=summary["baseline"]["median_seconds"] / summary["pipeline"]["median_seconds"],
+ first_divergent_output_positions=agreement,
+ trials=raw,
+ limitations="Small greedy feasibility check; not BF16 acceptance validation. "
+ "Peak memory includes both resident engines. Checkpoint config hash does not hash weights.",
+ )
+ args.output.parent.mkdir(parents=True, exist_ok=True)
+ args.output.write_text(json.dumps(report, indent=2) + "\n")
+ print(
+ json.dumps(dict(summary=summary, speedup=report["speedup"], divergence=agreement), indent=2)
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/docs/README.md b/docs/README.md
index 02848af0..ef5f60a6 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -13,6 +13,7 @@ an environment to receiving the first generated response. The
| [KV layout examples](kv_layout_computation.md) | SHARED and LAST_EXITED semantics and worked attention examples |
| [KV cache walkthrough](kv_cache_walkthrough.md) | Allocation state, contracts, cross-module dependencies, and the M4 migration plan |
| [Runtime configuration](cdb_runtime.md) | Exit policies, KV layouts, execution options, and CUDA Graphs |
+| [Loop-level speculation](loop_speculative.md) | Opt-in mixed-depth execution, KV dependencies, and feasibility measurements |
| [Asynchronous scheduling](https://github.com/hsliuustc0106/vllm-rlt/pull/30) | CPU/GPU pipelining and single-stream or multi-stream execution |
| [FlashAttention](https://github.com/hsliuustc0106/vllm-rlt/pull/30) | FA2/FA3/FA4 installation, hardware selection, and constraints |
| [Cache and scheduling features](https://github.com/hsliuustc0106/vllm-rlt/pull/31) | Prefix reuse, incremental KV, priorities, and preemption |
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
new file mode 100644
index 00000000..9faaa0cf
--- /dev/null
+++ b/docs/loop_speculative.md
@@ -0,0 +1,140 @@
+# Loop-level speculative execution
+
+This opt-in synchronous prototype implements the mixed-depth batching portion of
+[RFC #43](https://github.com/ThinkFlowLab/vllm-rlt/issues/43): advance shallow work
+for later token positions in the same recurrent call as deep work for earlier
+positions. It retains the existing finite speculative window and atomic round
+commit. The default remains the original block verifier.
+
+## Enable
+
+```python
+SpeculativeConfig(
+ num_speculative_tokens=4,
+ draft_loops=2,
+ target_loops=4,
+ loop_level_pipelining=True,
+)
+```
+
+The offline and serving CLI accept `--speculative-tokens 4 --loop-level-pipelining`.
+Existing restrictions apply: `last_exited`, full fixed target depth, synchronous
+refill scheduling, and no preemption or PD. Greedy and temperature/top-k/top-p
+sampling use the same acceptance implementation as the original runner.
+
+## Schedule and state
+
+For K=2, d=2 and D=4, A is the last committed input, B and C are proposals.
+Numbers below are zero-based depths; columns are successive recurrent calls.
+
+| Input position | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 |
+| --- | --- | --- | --- | --- | --- | --- | --- | --- |
+| A | 0 | 1 | 2 | 3 | | | | |
+| B | | | 0 | 1 | 2 | 3 | | |
+| C | | | | | 0 | 1 | 2 | 3 |
+
+After A completes depth 1, its shallow head proposes B. Call 3 batches A at
+depth 2 and B at depth 0. After B completes depth 1, it proposes C while A's
+full-depth head produces the distribution that verifies B. These head rows are
+batched too. C's full-depth distribution supplies the bonus token.
+
+Each branch stores its item index, position offset, hidden state and completed
+loop count. A child starts d calls behind its parent. Therefore, at every layer,
+the same-depth KV for earlier positions is already available before the child's
+attention. `_prepare_batch` supplies a separate depth plane and causal context
+length for every row. The pipeline uses single-query paged attention, not packed
+multi-query prefill. No shallow KV is copied into uncomputed deep planes.
+
+Each position/depth is computed exactly once. Proposals remain on the device and
+are read back with greedy target IDs once per round. Sampling retains the actual
+filtered q distribution and draws all drafts before acceptance, preserving the
+existing per-request RNG ordering. Different GEMM shapes can still change BF16
+results; same-seed agreement is not a general sampling correctness guarantee.
+
+Completed target logits are reordered into item/position order for the common
+verifier. `SpeculativeResult` and engine commit stay unchanged: keep KV through
+the initial input and accepted candidates; invalidate every depth's rejected
+suffix; leave correction/bonus unforwarded. Engine EOS/length handling and
+exception cleanup remain responsible for the request lifecycle.
+
+## Scope and benefit assessment
+
+This version does **not** implement streaming acceptance, early rejection pruning,
+cross-round hidden-state retention, or the complete continuous LoopSpec algorithm.
+It runs all K+1 positions to target depth even when an earlier proposal will be
+rejected. No new CUDA streams or asynchronous scheduler are introduced.
+
+For a uniform window, both the original block verifier and this pipeline require
+`d*K + D` recurrent calls and `(K+1)*D` recurrent rows. Pipelining changes batch
+shapes and head batching, not total recurrent work or that critical path. Its
+cost also includes extra per-step metadata preparation and loss of the original
+verifier's packed FA4 attention. Speedup must be measured; mixing depths alone
+does not guarantee an improvement. Reducing round bubbles or wasted suffix work
+would require a follow-up design for acceptance and state lifetime.
+
+## Validation
+
+```bash
+OMP_NUM_THREADS=1 python -m pytest -q
+python -m ruff check .
+python -m ruff format --check vllm_rlt tests benchmarks
+
+# On an allocated CUDA device with Triton:
+python -m pytest -q tests/test_speculative.py --run-gpu -m gpu -k 'not flash_attn_4'
+```
+
+Tests exercise mixed depths including non-divisible d/D, actual hidden/logit/KV
+agreement against a serial reference, every rejection boundary and bonus,
+greedy/sampling batches, K=1/2/4/8, shortened windows, EOS, prefix reuse,
+incremental allocation, cancellation and partial execution failure. CUDA coverage
+checks eager/graph parity and one greedy readback per round. The inherited GPU
+readback test also replaces its obsolete `OuroConfig.tiny()` call with the current
+test helper.
+
+For a bounded real-checkpoint feasibility run:
+
+```bash
+python -m benchmarks.loop_speculative \
+ --model /path/to/local/Ouro-1.4B \
+ --cuda-graphs --k 4 --concurrency 1 --new-tokens 32 --repeats 5 \
+ --output /tmp/loop-speculative.json
+```
+
+The harness warms both engines, completes matched prefill and first-token
+generation before timing, ignores EOS to commit equal work, synchronizes timing
+boundaries, alternates trial order, and saves raw timings, outputs, first
+divergences and runtime metadata. `--prompts` accepts a JSON list of strings;
+`--backend flash_attn_4` requires compatible hardware. Peak memory includes both
+resident engines; the recorded checkpoint config hash is not a weight hash.
+This small run does not replace RFC #43's fixed-input real-model diagnostics,
+sampling distribution checks, long-context/serving matrix or BF16 acceptance
+criteria.
+
+### Local feasibility results (2026-10-07)
+
+Base commit: `ecb1f8b505b7e831815b40aec3b4598619cca23a`, plus this working-tree
+implementation. Local Ouro-1.4B, A100-SXM4-40GB, PyTorch 2.14.0+cu130,
+BF16/Triton, CUDA Graphs, d=2/D=4, 32 new tokens after prefill/first-token,
+five alternating measured repetitions after warmup. GPU 0 was idle before the
+experiment. These are the harness's two short default prompts, cycled to fill
+the batch; C=1 uses only the first. No workload generalization is implied.
+
+| Concurrency | K | Baseline tokens/s | Pipeline tokens/s | Throughput change |
+| --- | --- | --- | --- | --- |
+| 1 | 1 | 53.74 | 53.13 | -1.1% |
+| 1 | 2 | 57.18 | 55.90 | -2.2% |
+| 1 | 4 | 58.59 | 56.92 | -2.8% |
+| 1 | 8 | 44.92 | 43.17 | -3.9% |
+| 8 | 1 | 383.68 | 380.84 | -0.7% |
+| 8 | 2 | 409.14 | 402.28 | -1.7% |
+| 8 | 4 | 420.43 | 409.88 | -2.5% |
+| 8 | 8 | 324.42 | 316.53 | -2.4% |
+
+Throughput is fixed committed tokens divided by median elapsed decode time.
+Every measured sequence agreed with the matching original speculative run.
+This does not prove general BF16 equivalence or agreement with native decoding.
+FA4, long contexts, native real-model baselines and serving latency were not
+evaluated. On this bounded workload the prototype is correct in the exercised
+cases but **does not improve throughput**; it remains opt-in for experimentation.
+Raw JSON from this local run is retained under `artifacts/loop-speculative/`
+(ignored by git), including per-trial timings, outputs and first divergences.
diff --git a/tests/test_loop_speculative.py b/tests/test_loop_speculative.py
new file mode 100644
index 00000000..c31287a7
--- /dev/null
+++ b/tests/test_loop_speculative.py
@@ -0,0 +1,185 @@
+"""Mixed-depth execution, round equivalence and pipeline rollback contracts."""
+
+import argparse
+from collections import Counter
+
+import pytest
+import torch
+
+from tests.test_speculative import drain, engine, model
+from vllm_rlt import LLM, CacheConfig, SamplingParams, SpeculativeConfig
+from vllm_rlt.core.scheduler import ScheduledItem, SchedulerOutput
+from vllm_rlt.entrypoints.runtime_args import add_runtime_args, runtime_configs
+from vllm_rlt.request import Request, Stage
+
+
+@pytest.mark.parametrize("draft_depth", [1, 2, 3])
+@pytest.mark.parametrize("sampling", [False, True])
+def test_pipeline_round_matches_baseline_and_batches_mixed_depths(draft_depth, sampling):
+ m = model()
+ snapshots = []
+ for pipelined in (False, True):
+ e = engine(m)
+ runner = e.speculative_runner
+ runner.config = SpeculativeConfig(4, draft_depth, 4, pipelined)
+ cache = e.cache_manager
+ items = []
+ for i, count in enumerate((5, 3, 1)):
+ rid = str(i)
+ prefix = [2, 3 + i]
+ cache.allocate(rid, 16)
+ for pos, token in enumerate(prefix):
+ h = m.prelude(torch.tensor([token]))
+ for depth in range(4):
+ h, _ = m.recurrent(h, [rid], [depth], [pos], cache)
+ params = SamplingParams(
+ temperature=0.8 if sampling and i != 1 else 0,
+ top_k=7,
+ top_p=0.8,
+ seed=42,
+ ignore_eos=True,
+ )
+ request = Request(rid, prefix, params, generated_token_ids=[7 + i])
+ items.append(ScheduledItem(request, len(prefix), count))
+ calls, work, heads = [], Counter(), []
+ core, verify = runner._core, runner._verify
+
+ def record(hidden, ids, positions, depth, **kwargs):
+ depths = [row[1] for row in kwargs["batch"].rows]
+ calls.append(list(zip(ids, positions, depths)))
+ work.update(zip(ids, positions, depths))
+ return core(hidden, ids, positions, depth, **kwargs)
+
+ def record_verify(items, drafted, proposals, logits):
+ heads.append(logits.clone())
+ return verify(items, drafted, proposals, logits)
+
+ runner._core, runner._verify = record, record_verify
+ results = runner.execute(SchedulerOutput(Stage.SPECULATIVE, items))
+ expected_work = Counter(
+ (item.request.request_id, pos, depth)
+ for item in items
+ for pos in range(item.token_start, item.token_start + item.token_count)
+ for depth in range(4)
+ )
+ assert work == expected_work # Each position/depth is computed once.
+ if pipelined:
+ assert any(len({d for _, _, d in call}) > 1 for call in calls)
+ # History at a given depth must finish before its later position.
+ completed = set()
+ for call in calls:
+ for rid, pos, depth in call:
+ assert pos == 2 or (rid, pos - 1, depth) in completed
+ completed.update(call)
+ assert len(calls) == draft_depth * 4 + 4
+ kv = [
+ tensor.clone()
+ for item in items
+ for depth in range(4)
+ for layer in range(m.config.num_hidden_layers)
+ for tensor in cache.read(
+ layer, item.request.request_id, depth, item.token_start + item.token_count
+ )
+ ]
+ snapshots.append((results, heads[0], kv, runner.stats))
+ baseline, pipeline = snapshots
+ assert baseline[0] == pipeline[0]
+ torch.testing.assert_close(baseline[1], pipeline[1], atol=4e-5, rtol=4e-5)
+ for a, b in zip(baseline[2], pipeline[2]):
+ torch.testing.assert_close(a, b, atol=4e-5, rtol=4e-5)
+ assert baseline[3] == pipeline[3]
+
+
+@pytest.mark.parametrize("rejected", [0, 1, 2, 3])
+@pytest.mark.parametrize("temperature", [0, 1])
+def test_pipeline_correction_bonus_and_next_round(rejected, temperature, monkeypatch):
+ m = model()
+ e = engine(m, pipelined=True, cache_config=CacheConfig(128, 2))
+ e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=12, ignore_eos=True))
+ while not e.step():
+ pass
+ request = e.scheduler.requests["r"]
+ request.sampling_params = SamplingParams(
+ max_tokens=12, ignore_eos=True, temperature=temperature, seed=42
+ )
+ start = request.position
+ runner = e.speculative_runner
+ core, coda = runner._core_mixed, runner._coda
+ head_rows = []
+
+ def record(hidden, ids, positions, depths):
+ nonlocal head_rows
+ output = core(hidden, ids, positions, depths)
+ # Force proposal/target distributions by semantic position and depth,
+ # independent of which other rows share this batched LM-head call.
+ head_rows = [
+ (p - start, False) for p, d in zip(positions, depths) if d == 1 and p - start < 3
+ ]
+ head_rows += [(p - start, True) for p, d in zip(positions, depths) if d == 3]
+ return output
+
+ def controlled(hidden):
+ assert len(hidden) == len(head_rows)
+ logits = torch.full((len(hidden), m.config.vocab_size), -torch.inf)
+ for row, (offset, target) in enumerate(head_rows):
+ token = 20 + offset if target and offset >= rejected else 10 + offset
+ logits[row, token] = 0
+ return logits
+
+ monkeypatch.setattr(runner, "_core_mixed", record)
+ monkeypatch.setattr(runner, "_coda", controlled)
+ output = e.step()[0]
+ count = rejected + 1
+ assert output.token_ids[-count:] == list(range(10, 10 + rejected)) + [20 + rejected]
+ assert runner.stats.accepted_tokens == rejected
+ for plane in e.cache_manager._get_allocation("r").written:
+ assert all(w.prefix == start + count and not w.pending for w in plane)
+ monkeypatch.setattr(runner, "_core_mixed", core)
+ monkeypatch.setattr(runner, "_coda", coda)
+ # Replay a common committed prefix with native greedy to expose stale KV.
+ request.sampling_params = SamplingParams(max_tokens=12, ignore_eos=True)
+ expected = LLM(m).generate(
+ [output.prompt_token_ids + output.token_ids],
+ SamplingParams(max_tokens=12 - len(output.token_ids), ignore_eos=True),
+ )[0]
+ assert drain(e)["r"].token_ids[len(output.token_ids) :] == expected.token_ids
+ assert e.cache_manager.num_used_blocks == 0
+
+
+def test_pipeline_config_and_cli():
+ assert not SpeculativeConfig(3).loop_level_pipelining
+ for invalid in (None, 0, 1, "true"):
+ with pytest.raises(ValueError, match="loop_level_pipelining"):
+ SpeculativeConfig(3, loop_level_pipelining=invalid)
+ parser = argparse.ArgumentParser()
+ add_runtime_args(parser)
+ parser.set_defaults(num_blocks=128, block_size=16)
+ args = parser.parse_args(["--speculative-tokens", "3", "--loop-level-pipelining"])
+ assert runtime_configs(args)["speculative_config"].loop_level_pipelining
+ args = parser.parse_args(["--loop-level-pipelining"])
+ with pytest.raises(ValueError, match="requires --speculative-tokens"):
+ runtime_configs(args)
+
+
+def test_pipeline_partial_execution_failure_releases_kv(monkeypatch):
+ e = engine(model(), pipelined=True)
+ e.add_request("r", [2, 3], SamplingParams(max_tokens=12, ignore_eos=True))
+ while not e.step():
+ pass
+ runner = e.speculative_runner
+ core = runner._core_mixed
+ calls = 0
+
+ def fail_after_write(*args):
+ nonlocal calls
+ result = core(*args)
+ calls += 1
+ if calls == 3:
+ raise RuntimeError("partial pipeline failure")
+ return result
+
+ monkeypatch.setattr(runner, "_core_mixed", fail_after_write)
+ with pytest.raises(RuntimeError, match="partial pipeline failure"):
+ e.step()
+ assert not e.has_unfinished_requests()
+ assert e.cache_manager.num_used_blocks == 0
diff --git a/tests/test_speculative.py b/tests/test_speculative.py
index b7fb44ab..bdcd70ee 100644
--- a/tests/test_speculative.py
+++ b/tests/test_speculative.py
@@ -29,8 +29,10 @@ def model(seed=123, dtype=torch.float32):
return OuroForCausalLM(tiny_ouro_config()).to(dtype=dtype)
-def engine(m, k=3, **kwargs):
- return LLMEngine(m, speculative_config=SpeculativeConfig(k), **kwargs)
+def engine(m, k=3, *, pipelined=False, **kwargs):
+ return LLMEngine(
+ m, speculative_config=SpeculativeConfig(k, loop_level_pipelining=pipelined), **kwargs
+ )
def drain(e):
@@ -61,14 +63,15 @@ def test_speculative_prefill_routes_by_declared_capabilities(generation, support
@pytest.mark.parametrize("k", [1, 2, 4, 8])
@pytest.mark.parametrize("budget", [1, 3, 32])
-def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget, pipelined):
m = model()
prompts = [[2], [3, 4, 5, 6, 7], [9, 3]]
params = [SamplingParams(max_tokens=n, ignore_eos=True) for n in [1, 13, 8]]
expected = LLM(m).generate(prompts, params)
llm = LLM(
m,
- speculative_config=SpeculativeConfig(k),
+ speculative_config=SpeculativeConfig(k, loop_level_pipelining=pipelined),
cache_config=CacheConfig(256, 2),
scheduler_config=SchedulerConfig(max_num_batched_tokens=budget, prefill_chunk_size=2),
)
@@ -87,7 +90,8 @@ def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget):
],
)
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
-def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, backend):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, backend, pipelined):
if device == "cuda" and dtype == torch.float32 and backend == "flash_attn_4":
pytest.skip("FA4 requires 16-bit operands")
m = model(dtype=dtype).to(device)
@@ -96,7 +100,9 @@ def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, back
m = OuroForCausalLM(tiny_ouro_config(hidden_size=256, head_dim=64)).to(
device=device, dtype=dtype
)
- e = engine(m, k=3, cache_config=CacheConfig(128, 2), attention_backend=backend)
+ e = engine(
+ m, k=3, pipelined=pipelined, cache_config=CacheConfig(128, 2), attention_backend=backend
+ )
cache = e.cache_manager
prefix = [3, 4, 5]
current = 7
@@ -113,23 +119,25 @@ def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, back
work = Counter()
def record(hidden, ids, positions, depth, **kwargs):
- work[depth] += len(ids)
out = core(hidden, ids, positions, depth, **kwargs)
for row, pos in enumerate(positions):
- seen[depth, pos] = out[row].clone()
+ row_depth = kwargs["batch"].rows[row][1]
+ work[row_depth] += 1
+ seen[row_depth, pos] = out[row].clone()
return out
runner._core = record
- head_outputs = []
- hook = m.lm_head.register_forward_hook(
- lambda module, inputs, output: head_outputs.append(output.detach().clone())
- )
- try:
- runner.execute(SchedulerOutput(Stage.SPECULATIVE, [ScheduledItem(r, 3, 4)]))
- finally:
- hook.remove()
- # Include the actual batched LM-head GEMM in the numeric comparison.
- verified_logits = head_outputs[-1]
+ # Capture the actual target logits even when heads run during the pipeline.
+ verify = runner._verify
+ verified = []
+
+ def record_verify(items, drafted, proposals, logits):
+ verified.append(logits.clone())
+ return verify(items, drafted, proposals, logits)
+
+ runner._verify = record_verify
+ runner.execute(SchedulerOutput(Stage.SPECULATIVE, [ScheduledItem(r, 3, 4)]))
+ verified_logits = verified[0]
# The shallow head outputs determine each subsequent input in this round.
tokens = prefix + [current]
for pos in range(3, 6):
@@ -221,9 +229,10 @@ def test_greedy_accept_commits_target_prefix_through_first_mismatch(k):
assert tokens == candidates[:rejected] + [7]
-def test_eos_in_accepted_prefix_never_delivers_following_tokens(monkeypatch):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_eos_in_accepted_prefix_never_delivers_following_tokens(monkeypatch, pipelined):
m = model()
- e = engine(m)
+ e = engine(m, pipelined=pipelined)
e.add_request("r", [2, 3], SamplingParams(max_tokens=12))
while not e.step():
pass
@@ -255,14 +264,15 @@ def test_truncation_removes_sparse_validity_without_destroying_prefix():
@pytest.mark.parametrize("prefix_cache,incremental", [(False, True), (True, False), (True, True)])
-def test_prefix_replay_incremental_growth_cancel_and_reuse(prefix_cache, incremental):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_prefix_replay_incremental_growth_cancel_and_reuse(prefix_cache, incremental, pipelined):
m = model()
kwargs = dict(
cache_config=CacheConfig(
128, 2, enable_prefix_caching=prefix_cache, incremental_allocation=incremental
)
)
- llm = LLM(m, speculative_config=SpeculativeConfig(4), **kwargs)
+ llm = LLM(m, speculative_config=SpeculativeConfig(4, loop_level_pipelining=pipelined), **kwargs)
params = SamplingParams(max_tokens=9, ignore_eos=True)
prompts = [[2, 3, 4, 5, 6], [2, 3, 4, 5, 7]]
expected = LLM(m).generate(prompts, params)
@@ -282,7 +292,8 @@ def test_prefix_replay_incremental_growth_cancel_and_reuse(prefix_cache, increme
@pytest.mark.parametrize("top_k,top_p", [(-1, 1.0), (5, 1.0), (7, 0.8)])
-def test_sampling_seeded_replay_and_request_isolation(top_k, top_p):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_sampling_seeded_replay_and_request_isolation(top_k, top_p, pipelined):
m = model()
prompts = [[2, 3], [7, 8, 9]]
params = SamplingParams(
@@ -290,7 +301,9 @@ def test_sampling_seeded_replay_and_request_isolation(top_k, top_p):
)
def run(prompts):
- return LLM(m, speculative_config=SpeculativeConfig(3)).generate(prompts, params)
+ return LLM(
+ m, speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined)
+ ).generate(prompts, params)
batched = run(prompts)
assert [o.token_ids for o in run(prompts)] == [o.token_ids for o in batched]
@@ -444,8 +457,11 @@ def test_rejects_adaptive_or_wrong_target_requests():
assert not e.has_unfinished_requests()
-def test_dynamic_arrival_budget_and_failure_reclamation(monkeypatch):
- e = engine(model(), scheduler_config=SchedulerConfig(max_num_batched_tokens=5))
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_dynamic_arrival_budget_and_failure_reclamation(monkeypatch, pipelined):
+ e = engine(
+ model(), pipelined=pipelined, scheduler_config=SchedulerConfig(max_num_batched_tokens=5)
+ )
params = SamplingParams(max_tokens=9, ignore_eos=True)
e.add_request("a", [2, 3], params)
while not e.step():
@@ -475,7 +491,8 @@ def test_invalid_k(k):
@pytest.mark.gpu
@pytest.mark.parametrize("backend", ["triton", "flash_attn_4"])
-def test_gpu_sampling_and_greedy_match_replay_with_ragged_requests(backend):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_gpu_sampling_and_greedy_match_replay_with_ragged_requests(backend, pipelined):
torch.manual_seed(123)
m = OuroForCausalLM(tiny_ouro_config(hidden_size=256, head_dim=64)).to(
device="cuda", dtype=torch.bfloat16
@@ -489,7 +506,9 @@ def test_gpu_sampling_and_greedy_match_replay_with_ragged_requests(backend):
def run(speculative, sampling):
return LLM(
m,
- speculative_config=SpeculativeConfig(3) if speculative else None,
+ speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined)
+ if speculative
+ else None,
cache_config=CacheConfig(128, 16),
attention_backend=backend,
).generate(prompts, sampling)
@@ -502,14 +521,17 @@ def run(speculative, sampling):
@pytest.mark.gpu
-def test_greedy_round_reads_back_once_with_ragged_requests():
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_greedy_round_reads_back_once_with_ragged_requests(pipelined):
# 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(tiny_ouro_config(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")
+ e = engine(
+ m, k=3, pipelined=pipelined, 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
@@ -529,7 +551,8 @@ def test_greedy_round_reads_back_once_with_ragged_requests():
@pytest.mark.gpu
@pytest.mark.parametrize("max_graph_rows", [2, 8])
-def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows, pipelined):
m = model(dtype=torch.bfloat16).to("cuda")
prompts = [[2, 3, 4], [7, 8]]
params = SamplingParams(max_tokens=9, ignore_eos=True)
@@ -537,7 +560,7 @@ def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows)
def run(graphs):
llm = LLM(
m,
- speculative_config=SpeculativeConfig(3),
+ speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined),
cache_config=CacheConfig(128, 16),
attention_backend="triton",
execution_config=ExecutionConfig(
diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py
index 2491dcf0..2d26620f 100644
--- a/vllm_rlt/config.py
+++ b/vllm_rlt/config.py
@@ -137,9 +137,12 @@ class SpeculativeConfig:
num_speculative_tokens: int
draft_loops: int = 2
target_loops: int = 4
+ loop_level_pipelining: bool = False
def __post_init__(self):
for name in ("num_speculative_tokens", "draft_loops", "target_loops"):
_positive(name, getattr(self, name))
if self.draft_loops >= self.target_loops:
raise ValueError("draft_loops must be smaller than target_loops")
+ if type(self.loop_level_pipelining) is not bool:
+ raise ValueError("loop_level_pipelining must be a boolean")
diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py
index 9a7df07e..75102930 100644
--- a/vllm_rlt/entrypoints/runtime_args.py
+++ b/vllm_rlt/entrypoints/runtime_args.py
@@ -18,6 +18,11 @@ def add_runtime_args(parser):
)
parser.add_argument("--draft-loops", type=int, default=2)
parser.add_argument("--target-loops", type=int, default=4)
+ parser.add_argument(
+ "--loop-level-pipelining",
+ action="store_true",
+ help="Batch active speculative branches at different recurrent depths",
+ )
parser.add_argument("--enable-prefix-caching", action="store_true")
parser.add_argument("--incremental-kv", action="store_true")
parser.add_argument("--kv-watermark", type=float, default=0.0)
@@ -60,11 +65,14 @@ def runtime_configs(args):
parser = argparse.ArgumentParser(add_help=False)
add_runtime_args(parser)
args = SimpleNamespace(**(vars(parser.parse_args([])) | vars(args)))
+ if args.loop_level_pipelining and args.speculative_tokens is None:
+ raise ValueError("--loop-level-pipelining requires --speculative-tokens")
return dict(
speculative_config=SpeculativeConfig(
num_speculative_tokens=args.speculative_tokens,
draft_loops=args.draft_loops,
target_loops=args.target_loops,
+ loop_level_pipelining=args.loop_level_pipelining,
)
if args.speculative_tokens is not None
else None,
diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py
index 089e708c..fb7bca77 100644
--- a/vllm_rlt/worker/speculative.py
+++ b/vllm_rlt/worker/speculative.py
@@ -28,6 +28,16 @@ class SpeculativeStats:
verified_rows: int = 0
+@dataclass
+class _PipelineBranch:
+ """One speculative position and its current recurrent state."""
+
+ item_index: int
+ offset: int
+ hidden: torch.Tensor
+ depth: int = 0
+
+
def greedy_accept(candidates, targets):
"""Return (tokens, accepted) for one request from host candidate and target IDs.
@@ -109,6 +119,11 @@ def _acquire_generator(self, request, logits):
"""
return self.rng.acquire(request.request_id, request.sampling_params, logits.device)
+ def _core_mixed(self, hidden, ids, positions, depths):
+ """Advance active pipeline branches with one depth per row."""
+ batch = self.cache._prepare_batch(ids, depths, positions)
+ return self._core(hidden, ids, positions, None, batch=batch)
+
def _coda(self, hidden):
return (
self.coda_graphs.run(hidden)
@@ -116,8 +131,204 @@ def _coda(self, hidden):
else self.model.coda(hidden)
)
+ def _greedy_pruning_enabled(self, items):
+ return bool(getattr(self.config, "early_rejection_pruning", False)) and all(
+ item.request.sampling_params.temperature == 0 for item in items
+ )
+
+ @torch.inference_mode()
+ def _execute_loop_level_greedy_pruned(self, batch):
+ """Run loop-level speculation while pruning rejected greedy suffixes."""
+ items = batch.items
+ accepted = [0] * len(items)
+ draft_counts = [0] * len(items)
+ results = [[] for _ in items]
+ candidates = [{} for _ in items]
+ stopped = [False] * len(items)
+
+ initial = self.cache._stage([item.request.input_token_id for item in items], torch.long)
+ hidden = self.model.prelude(initial)
+ active = [_PipelineBranch(i, 0, hidden[i]) for i in range(len(items))]
+ verified_rows = 0
+
+ while active:
+ ids = [items[branch.item_index].request.request_id for branch in active]
+ positions = [items[branch.item_index].token_start + branch.offset for branch in active]
+ depths = [branch.depth for branch in active]
+ hidden = torch.stack([branch.hidden for branch in active])
+ hidden = self._core_mixed(hidden, ids, positions, depths)
+ for row, branch in enumerate(active):
+ branch.hidden = hidden[row]
+ branch.depth += 1
+
+ draft_ready = [
+ branch
+ for branch in active
+ if branch.depth == self.config.draft_loops
+ and branch.offset + 1 < items[branch.item_index].token_count
+ and not stopped[branch.item_index]
+ ]
+ target_ready = [branch for branch in active if branch.depth == self.config.target_loops]
+ ready = draft_ready + target_ready
+ ready_logits = None
+ ready_ids = None
+ ready_ids_host = None
+ if ready:
+ ready_logits = self._coda(torch.stack([branch.hidden for branch in ready]))
+ # Read all draft and target IDs back once for this wave. A
+ # per-row .item() would synchronize the GPU for every branch.
+ ready_ids = ready_logits.argmax(-1)
+ ready_ids_host = ready_ids.tolist()
+
+ # Validate target rows before creating a child from a same-wave
+ # draft row for the same request.
+ target_start = len(draft_ready)
+ for row, branch in enumerate(target_ready):
+ item_index = branch.item_index
+ item = items[item_index]
+ target = ready_ids_host[target_start + row]
+ verified_rows += 1
+ if branch.offset == item.token_count - 1:
+ results[item_index].append(target)
+ stopped[item_index] = True
+ continue
+ candidate = candidates[item_index].get(branch.offset)
+ if candidate is None:
+ raise RuntimeError(
+ f"missing greedy candidate for item {item_index}, offset {branch.offset}"
+ )
+ results[item_index].append(target)
+ accepted[item_index] += candidate == target
+ if candidate != target:
+ stopped[item_index] = True
+
+ # Create children only for requests that survived target validation
+ # in this wave.
+ draft_rows = [
+ (row, branch)
+ for row, branch in enumerate(draft_ready)
+ if not stopped[branch.item_index]
+ ]
+ children = []
+ if draft_rows:
+ row_ids = torch.tensor([row for row, _ in draft_rows], device=self.device)
+ child_tokens = ready_ids[: len(draft_ready)].index_select(0, row_ids)
+ child_hidden = self.model.prelude(child_tokens)
+ for child_row, (draft_row, branch) in enumerate(draft_rows):
+ item_index = branch.item_index
+ candidate = ready_ids_host[draft_row]
+ candidates[item_index][branch.offset] = candidate
+ draft_counts[item_index] += 1
+ children.append(
+ _PipelineBranch(item_index, branch.offset + 1, child_hidden[child_row])
+ )
+
+ active = [
+ branch
+ for branch in active
+ if branch.depth < self.config.target_loops
+ and not stopped[branch.item_index]
+ ]
+ active.extend(children)
+
+ self.stats.rounds += len(items)
+ self.stats.drafted_tokens += sum(draft_counts)
+ self.stats.verified_rows += verified_rows
+ return [
+ SpeculativeResult(results[i], accepted[i], draft_counts[i]) for i in range(len(items))
+ ]
+
+ @torch.inference_mode()
+ def _execute_loop_level(self, batch):
+ """Run a synchronous single-proposal loop-level pipeline.
+
+ Every active branch advances one recurrent depth per step. Once a
+ branch reaches the draft depth it starts its child position while the
+ parent continues toward the full target depth. The finite window is
+ fully verified before returning, so the engine keeps its existing
+ atomic commit and KV rollback boundary.
+ """
+ items = batch.items
+ drafted = []
+ proposals = [[] for _ in items]
+ target_logits = [{} for _ in items]
+
+ initial = self.cache._stage([item.request.input_token_id for item in items], torch.long)
+ hidden = self.model.prelude(initial)
+ active = [_PipelineBranch(i, 0, hidden[i]) for i in range(len(items))]
+
+ while active:
+ # Children start draft_loops steps behind their parents. Thus the
+ # earlier position's KV at this row's depth is already complete.
+ # Use independent query rows: adjacent positions here have different
+ # depth planes and must not form one packed prefill sequence.
+ ids = [items[branch.item_index].request.request_id for branch in active]
+ positions = [items[branch.item_index].token_start + branch.offset for branch in active]
+ depths = [branch.depth for branch in active]
+ hidden = torch.stack([branch.hidden for branch in active])
+ hidden = self._core_mixed(hidden, ids, positions, depths)
+ for row, branch in enumerate(active):
+ branch.hidden = hidden[row]
+ branch.depth += 1
+
+ draft_ready = [
+ branch
+ for branch in active
+ if branch.depth == self.config.draft_loops
+ and branch.offset + 1 < items[branch.item_index].token_count
+ ]
+ target_ready = [branch for branch in active if branch.depth == self.config.target_loops]
+ ready = draft_ready + target_ready
+ ready_logits = None
+ if ready:
+ ready_logits = self._coda(torch.stack([branch.hidden for branch in ready]))
+
+ children = []
+ if draft_ready:
+ draft_logits = ready_logits[: len(draft_ready)]
+ child_tokens = draft_logits.argmax(-1)
+ for row, branch in enumerate(draft_ready):
+ request = items[branch.item_index].request
+ if request.sampling_params.temperature != 0:
+ q = probabilities(draft_logits[row], request.sampling_params)
+ proposals[branch.item_index].append(q)
+ child_tokens[row] = draw(q, generator_for(request, self.device))
+ drafted.append(([branch.item_index for branch in draft_ready], child_tokens))
+ child_hidden = self.model.prelude(child_tokens)
+ children = [
+ _PipelineBranch(
+ branch.item_index,
+ branch.offset + 1,
+ child_hidden[row],
+ )
+ for row, branch in enumerate(draft_ready)
+ ]
+
+ if target_ready:
+ target_start = len(draft_ready)
+ for row, branch in enumerate(target_ready):
+ target_logits[branch.item_index][branch.offset] = ready_logits[
+ target_start + row
+ ].clone()
+
+ active = [branch for branch in active if branch.depth < self.config.target_loops]
+ active.extend(children)
+
+ logits = torch.stack(
+ [
+ target_logits[i][offset]
+ for i, item in enumerate(items)
+ for offset in range(item.token_count)
+ ]
+ )
+ return self._verify(items, drafted, proposals, logits)
+
@torch.inference_mode()
def execute(self, batch):
+ if self.config.loop_level_pipelining:
+ if self._greedy_pruning_enabled(batch.items):
+ return self._execute_loop_level_greedy_pruned(batch)
+ return self._execute_loop_level(batch)
items = batch.items
proposals = [[] for _ in items]
states = [[] for _ in items]
@@ -166,6 +377,10 @@ def execute(self, batch):
packed=True,
)
logits = self._coda(hidden)
+ return self._verify(items, drafted, proposals, logits)
+
+ def _verify(self, items, drafted, proposals, logits):
+ """Collect the finite window and apply the common acceptance rule."""
# One device-to-host read returns every draft ID and, for greedy requests,
# every verification target argmax.
greedy = any(item.request.sampling_params.temperature == 0 for item in items)
@@ -213,5 +428,5 @@ def execute(self, batch):
results.append(SpeculativeResult(tokens, accepted, len(candidates[i])))
self.stats.rounds += len(items)
self.stats.drafted_tokens += sum(len(c) for c in candidates)
- self.stats.verified_rows += len(ids)
+ self.stats.verified_rows += len(logits)
return results
From 9586b4a5466354025c43536c70b7ac9d60de8770 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 07:24:22 +0000
Subject: [PATCH 02/20] feat: add greedy early rejection pruning
---
benchmarks/loop_speculative.py | 11 +++++-
docs/loop_speculative.md | 18 ++++++---
tests/test_loop_speculative.py | 59 +++++++++++++++++++++++++++-
vllm_rlt/config.py | 9 ++++-
vllm_rlt/entrypoints/runtime_args.py | 6 +++
vllm_rlt/worker/speculative.py | 31 ++++++++++-----
6 files changed, 114 insertions(+), 20 deletions(-)
diff --git a/benchmarks/loop_speculative.py b/benchmarks/loop_speculative.py
index 993e5423..c5531e3e 100644
--- a/benchmarks/loop_speculative.py
+++ b/benchmarks/loop_speculative.py
@@ -70,6 +70,11 @@ def main():
parser.add_argument("--cuda-graphs", action="store_true")
parser.add_argument("--k", type=int, default=4)
parser.add_argument("--draft-loops", type=int, default=2)
+ parser.add_argument(
+ "--early-rejection-pruning",
+ action="store_true",
+ help="Enable greedy early rejection pruning for the pipeline engine",
+ )
parser.add_argument("--concurrency", type=int, default=1)
parser.add_argument("--new-tokens", type=int, default=32)
parser.add_argument("--repeats", type=int, default=5)
@@ -109,7 +114,11 @@ def main():
execution_config=execution,
attention_backend=args.backend,
speculative_config=SpeculativeConfig(
- args.k, args.draft_loops, model.config.total_ut_steps, enabled
+ num_speculative_tokens=args.k,
+ draft_loops=args.draft_loops,
+ target_loops=model.config.total_ut_steps,
+ loop_level_pipelining=enabled,
+ early_rejection_pruning=enabled and args.early_rejection_pruning,
),
)
for name, enabled in [("baseline", False), ("pipeline", True)]
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 9faaa0cf..c34ce630 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -18,9 +18,11 @@ SpeculativeConfig(
```
The offline and serving CLI accept `--speculative-tokens 4 --loop-level-pipelining`.
+Greedy-only early suffix pruning can be enabled separately with
+`--early-rejection-pruning`; it is off by default. Sampling continues to use the
+original batched verification path so its RNG ordering is unchanged.
Existing restrictions apply: `last_exited`, full fixed target depth, synchronous
-refill scheduling, and no preemption or PD. Greedy and temperature/top-k/top-p
-sampling use the same acceptance implementation as the original runner.
+refill scheduling, and no preemption or PD.
## Schedule and state
@@ -59,10 +61,14 @@ exception cleanup remain responsible for the request lifecycle.
## Scope and benefit assessment
-This version does **not** implement streaming acceptance, early rejection pruning,
-cross-round hidden-state retention, or the complete continuous LoopSpec algorithm.
-It runs all K+1 positions to target depth even when an earlier proposal will be
-rejected. No new CUDA streams or asynchronous scheduler are introduced.
+The default finite-window path does **not** implement cross-round hidden-state
+retention or the complete continuous LoopSpec algorithm. The optional greedy
+pruning path checks a target row when its coda output is ready and skips later
+suffix branches after rejection. Work that already entered the mixed-depth wave
+cannot be cancelled. Sampling still runs the complete window before
+verification. Greedy target and draft IDs are read back once per mixed-depth
+wave; the host-side rejection decision remains the pruning synchronization
+boundary. No new CUDA streams or asynchronous scheduler are introduced.
For a uniform window, both the original block verifier and this pipeline require
`d*K + D` recurrent calls and `(K+1)*D` recurrent rows. Pipelining changes batch
diff --git a/tests/test_loop_speculative.py b/tests/test_loop_speculative.py
index c31287a7..a2da7e0a 100644
--- a/tests/test_loop_speculative.py
+++ b/tests/test_loop_speculative.py
@@ -11,6 +11,7 @@
from vllm_rlt.core.scheduler import ScheduledItem, SchedulerOutput
from vllm_rlt.entrypoints.runtime_args import add_runtime_args, runtime_configs
from vllm_rlt.request import Request, Stage
+from vllm_rlt.worker.speculative import SpeculativeResult
@pytest.mark.parametrize("draft_depth", [1, 2, 3])
@@ -146,6 +147,57 @@ def controlled(hidden):
assert e.cache_manager.num_used_blocks == 0
+def test_greedy_early_rejection_prunes_unstarted_suffix(monkeypatch):
+ m = model()
+ e = engine(m, pipelined=True, cache_config=CacheConfig(128, 2))
+ runner = e.speculative_runner
+ runner.config = SpeculativeConfig(4, 2, 4, True, True)
+ cache = e.cache_manager
+ request = Request(
+ "r", [2, 3], SamplingParams(max_tokens=12, ignore_eos=True), generated_token_ids=[7]
+ )
+ cache.allocate(request.request_id, 16)
+ for pos, token in enumerate(request.prompt_token_ids):
+ hidden = m.prelude(torch.tensor([token]))
+ for depth in range(4):
+ hidden, _ = m.recurrent(hidden, [request.request_id], [depth], [pos], cache)
+ item = ScheduledItem(request, len(request.prompt_token_ids), 4)
+
+ calls = []
+ core, coda = runner._core_mixed, runner._coda
+ head_rows = []
+
+ def record(hidden, ids, positions, depths):
+ nonlocal head_rows
+ calls.append([(p - item.token_start, d) for p, d in zip(positions, depths)])
+ output = core(hidden, ids, positions, depths)
+ head_rows = [(p - item.token_start, False) for p, d in zip(positions, depths) if d == 1]
+ head_rows += [(p - item.token_start, True) for p, d in zip(positions, depths) if d == 3]
+ return output
+
+ def controlled(hidden):
+ assert len(hidden) == len(head_rows)
+ logits = torch.full((len(hidden), m.config.vocab_size), -torch.inf)
+ for row, (offset, target) in enumerate(head_rows):
+ token = 20 if target and offset == 0 else 10 + offset
+ logits[row, token] = 0
+ return logits
+
+ monkeypatch.setattr(runner, "_core_mixed", record)
+ monkeypatch.setattr(runner, "_coda", controlled)
+ result = runner.execute(SchedulerOutput(Stage.SPECULATIVE, [item]))[0]
+
+ assert result == SpeculativeResult([20], 0, 1)
+ assert runner.stats.verified_rows == 1
+ assert runner.stats.drafted_tokens == 1
+ assert all(offset < 2 for call in calls for offset, _ in call)
+ assert any(offset == 1 and depth == 1 for call in calls for offset, depth in call)
+ assert not any(offset == 1 and depth >= 2 for call in calls for offset, depth in call)
+ assert not any(offset >= 2 for call in calls for offset, _ in call)
+ monkeypatch.setattr(runner, "_core_mixed", core)
+ monkeypatch.setattr(runner, "_coda", coda)
+
+
def test_pipeline_config_and_cli():
assert not SpeculativeConfig(3).loop_level_pipelining
for invalid in (None, 0, 1, "true"):
@@ -154,8 +206,11 @@ def test_pipeline_config_and_cli():
parser = argparse.ArgumentParser()
add_runtime_args(parser)
parser.set_defaults(num_blocks=128, block_size=16)
- args = parser.parse_args(["--speculative-tokens", "3", "--loop-level-pipelining"])
- assert runtime_configs(args)["speculative_config"].loop_level_pipelining
+ args = parser.parse_args(
+ ["--speculative-tokens", "3", "--loop-level-pipelining", "--early-rejection-pruning"]
+ )
+ config = runtime_configs(args)["speculative_config"]
+ assert config.loop_level_pipelining and config.early_rejection_pruning
args = parser.parse_args(["--loop-level-pipelining"])
with pytest.raises(ValueError, match="requires --speculative-tokens"):
runtime_configs(args)
diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py
index 2d26620f..a98cd922 100644
--- a/vllm_rlt/config.py
+++ b/vllm_rlt/config.py
@@ -138,11 +138,16 @@ class SpeculativeConfig:
draft_loops: int = 2
target_loops: int = 4
loop_level_pipelining: bool = False
+ early_rejection_pruning: bool = False
def __post_init__(self):
for name in ("num_speculative_tokens", "draft_loops", "target_loops"):
_positive(name, getattr(self, name))
if self.draft_loops >= self.target_loops:
raise ValueError("draft_loops must be smaller than target_loops")
- if type(self.loop_level_pipelining) is not bool:
- raise ValueError("loop_level_pipelining must be a boolean")
+ for name in (
+ "loop_level_pipelining",
+ "early_rejection_pruning",
+ ):
+ if type(getattr(self, name)) is not bool:
+ raise ValueError(f"{name} must be a boolean")
diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py
index 75102930..824d3385 100644
--- a/vllm_rlt/entrypoints/runtime_args.py
+++ b/vllm_rlt/entrypoints/runtime_args.py
@@ -23,6 +23,11 @@ def add_runtime_args(parser):
action="store_true",
help="Batch active speculative branches at different recurrent depths",
)
+ parser.add_argument(
+ "--early-rejection-pruning",
+ action="store_true",
+ help="Stop greedy loop-level branches after the first rejected candidate",
+ )
parser.add_argument("--enable-prefix-caching", action="store_true")
parser.add_argument("--incremental-kv", action="store_true")
parser.add_argument("--kv-watermark", type=float, default=0.0)
@@ -73,6 +78,7 @@ def runtime_configs(args):
draft_loops=args.draft_loops,
target_loops=args.target_loops,
loop_level_pipelining=args.loop_level_pipelining,
+ early_rejection_pruning=args.early_rejection_pruning,
)
if args.speculative_tokens is not None
else None,
diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py
index fb7bca77..baddbbe3 100644
--- a/vllm_rlt/worker/speculative.py
+++ b/vllm_rlt/worker/speculative.py
@@ -138,7 +138,14 @@ def _greedy_pruning_enabled(self, items):
@torch.inference_mode()
def _execute_loop_level_greedy_pruned(self, batch):
- """Run loop-level speculation while pruning rejected greedy suffixes."""
+ """Run loop-level speculation while pruning rejected greedy suffixes.
+
+ A target row is compared as soon as its coda output is available. A
+ rejected request is removed before the next mixed-depth core wave, so
+ shallow work already in flight is kept but no later suffix work is
+ started. This path is greedy-only; sampling retains the common
+ batched verification path so its RNG ordering is unchanged.
+ """
items = batch.items
accepted = [0] * len(items)
draft_counts = [0] * len(items)
@@ -175,13 +182,15 @@ def _execute_loop_level_greedy_pruned(self, batch):
ready_ids_host = None
if ready:
ready_logits = self._coda(torch.stack([branch.hidden for branch in ready]))
- # Read all draft and target IDs back once for this wave. A
- # per-row .item() would synchronize the GPU for every branch.
+ # One device-to-host transfer per wave supplies both draft
+ # and target IDs. Per-row `.item()` here would synchronize
+ # the GPU once for every target and draft branch.
ready_ids = ready_logits.argmax(-1)
ready_ids_host = ready_ids.tolist()
- # Validate target rows before creating a child from a same-wave
- # draft row for the same request.
+ # Target rows are handled first. If a parent is rejected, a draft
+ # row for the same request in this coda batch must not create its
+ # child position.
target_start = len(draft_ready)
for row, branch in enumerate(target_ready):
item_index = branch.item_index
@@ -189,6 +198,7 @@ def _execute_loop_level_greedy_pruned(self, batch):
target = ready_ids_host[target_start + row]
verified_rows += 1
if branch.offset == item.token_count - 1:
+ # The final row has no candidate: it is the bonus token.
results[item_index].append(target)
stopped[item_index] = True
continue
@@ -202,8 +212,9 @@ def _execute_loop_level_greedy_pruned(self, batch):
if candidate != target:
stopped[item_index] = True
- # Create children only for requests that survived target validation
- # in this wave.
+ # Generate children only for requests that survived target
+ # validation in this wave. Candidate IDs are selected on-device;
+ # their host values come from the single wave-level readback above.
draft_rows = [
(row, branch)
for row, branch in enumerate(draft_ready)
@@ -290,9 +301,11 @@ def _execute_loop_level(self, batch):
for row, branch in enumerate(draft_ready):
request = items[branch.item_index].request
if request.sampling_params.temperature != 0:
- q = probabilities(draft_logits[row], request.sampling_params)
+ q = Sampler.probabilities(draft_logits[row], request.sampling_params)
proposals[branch.item_index].append(q)
- child_tokens[row] = draw(q, generator_for(request, self.device))
+ child_tokens[row] = Sampler.draw(
+ q, self._acquire_generator(request, draft_logits[row])
+ )
drafted.append(([branch.item_index for branch in draft_ready], child_tokens))
child_hidden = self.model.prelude(child_tokens)
children = [
From 9c0dd53327890c1e4dbd5852ff2b0440eb4f9343 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 08:06:43 +0000
Subject: [PATCH 03/20] docs: record rebased loop speculative benchmark
---
benchmarks/loop_speculative.py | 52 ++++++++++++++----
docs/loop_speculative.md | 98 +++++++++++++++++++---------------
2 files changed, 97 insertions(+), 53 deletions(-)
diff --git a/benchmarks/loop_speculative.py b/benchmarks/loop_speculative.py
index c5531e3e..3be2f9e4 100644
--- a/benchmarks/loop_speculative.py
+++ b/benchmarks/loop_speculative.py
@@ -24,19 +24,23 @@
)
from vllm_rlt.engine.llm_engine import LLMEngine
from vllm_rlt.models import OuroForCausalLM, load_tokenizer
-from vllm_rlt.request import Stage
from vllm_rlt.worker.speculative import SpeculativeStats
@torch.inference_mode()
-def trial(engine, prompts, new_tokens):
- params = SamplingParams(max_tokens=new_tokens + 1, ignore_eos=True)
+def trial(engine, prompts, new_tokens, temperature, top_p, top_k):
+ params = SamplingParams(
+ max_tokens=new_tokens + 1,
+ ignore_eos=True,
+ temperature=temperature,
+ top_p=top_p,
+ top_k=top_k,
+ )
for i, prompt in enumerate(prompts):
engine.add_request(str(i), prompt, params)
- # All prompts fit one prefill batch, so no decode can precede this barrier.
+ # Hold decode until all requests have crossed the first-output barrier.
while any(not r.generated_token_ids for r in engine.scheduler.requests.values()):
engine.step()
- assert engine.last_schedule.stage != Stage.SPECULATIVE
engine.speculative_runner.stats = SpeculativeStats()
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
@@ -50,7 +54,11 @@ def trial(engine, prompts, new_tokens):
elapsed = time.perf_counter() - start
committed = len(prompts) * new_tokens
stats = asdict(engine.speculative_runner.stats)
- assert stats["committed_tokens"] == committed
+ if stats["committed_tokens"] != committed:
+ raise RuntimeError(
+ f"benchmark committed {stats['committed_tokens']} tokens, expected {committed}; "
+ f"output lengths={[len(outputs[str(i)]) for i in range(len(prompts))]}"
+ )
return dict(
seconds=elapsed,
committed_tokens=committed,
@@ -78,11 +86,18 @@ def main():
parser.add_argument("--concurrency", type=int, default=1)
parser.add_argument("--new-tokens", type=int, default=32)
parser.add_argument("--repeats", type=int, default=5)
+ parser.add_argument("--temperature", type=float, default=0.0)
+ parser.add_argument("--top-p", type=float, default=1.0)
+ parser.add_argument("--top-k", type=int, default=-1)
args = parser.parse_args()
if not args.model.is_dir():
parser.error("--model must be a complete local checkpoint directory")
if min(args.concurrency, args.new_tokens, args.repeats, args.k) < 1:
parser.error("concurrency, new-tokens, repeats and k must be positive")
+ if args.temperature < 0 or not 0 < args.top_p <= 1:
+ parser.error("temperature must be nonnegative and top-p must be in (0, 1]")
+ if args.top_k == 0 or args.top_k < -1:
+ parser.error("top-k must be -1 or positive")
torch.set_num_threads(1)
texts = (
json.loads(args.prompts.read_text())
@@ -102,8 +117,13 @@ def main():
)
scheduler = SchedulerConfig(
max_num_seqs=args.concurrency,
- max_num_batched_tokens=max(sum(map(len, prompts)), args.concurrency * (args.k + 1)),
+ max_num_batched_tokens=max(
+ sum(map(len, prompts)),
+ args.concurrency * max(map(len, prompts)),
+ args.concurrency * (args.k + 1),
+ ),
prefill_chunk_size=max(map(len, prompts)),
+ max_prefill_batches_before_decode=1024,
)
execution = ExecutionConfig(cuda_graphs=args.cuda_graphs)
engines = {
@@ -126,11 +146,25 @@ def main():
raw = {name: [] for name in engines}
try:
for engine in engines.values():
- trial(engine, prompts, args.new_tokens)
+ trial(
+ engine,
+ prompts,
+ args.new_tokens,
+ args.temperature,
+ args.top_p,
+ args.top_k,
+ )
for repeat in range(args.repeats):
order = list(engines) if repeat % 2 == 0 else list(reversed(engines))
for name in order:
- result = trial(engines[name], prompts, args.new_tokens)
+ result = trial(
+ engines[name],
+ prompts,
+ args.new_tokens,
+ args.temperature,
+ args.top_p,
+ args.top_k,
+ )
raw[name].append(result)
print(
f"{name} repeat={repeat}: {result['tokens_per_second']:.2f} tokens/s",
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index c34ce630..1ca4edf6 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -19,11 +19,27 @@ SpeculativeConfig(
The offline and serving CLI accept `--speculative-tokens 4 --loop-level-pipelining`.
Greedy-only early suffix pruning can be enabled separately with
-`--early-rejection-pruning`; it is off by default. Sampling continues to use the
-original batched verification path so its RNG ordering is unchanged.
+`--early-rejection-pruning`; it is off by default. Top-p and top-k sampling are
+supported by the same loop pipeline, but sampling still verifies the complete
+finite window before applying rejection sampling. Early pruning is intentionally
+greedy-only so it does not change the sampling RNG schedule.
Existing restrictions apply: `last_exited`, full fixed target depth, synchronous
refill scheduling, and no preemption or PD.
+The implementation has three execution choices:
+
+| Mode | Recurrent schedule | Verification | Intended use |
+| --- | --- | --- | --- |
+| Baseline | Complete all draft positions, then complete target depth | One batched verifier window | Reference path and default |
+| Loop | Advance active positions one depth per mixed-depth wave | Reorder completed target rows into the common verifier | Measure mixed-depth batching |
+| Loop + greedy prune | Loop schedule with target-ready rejection checks | Stop future suffix branches after the first mismatch | Optional greedy experiment |
+
+The loop mode is a synchronous pipeline. It does not add CUDA streams, reduce the
+number of recurrent rows, or retain hidden state across speculative rounds. Its
+possible benefit comes from the shape and timing of mixed-depth batches; metadata,
+Python scheduling, Coda fragmentation and dynamic CUDA Graph shapes are additional
+costs.
+
## Schedule and state
For K=2, d=2 and D=4, A is the last committed input, B and C are proposals.
@@ -48,10 +64,12 @@ length for every row. The pipeline uses single-query paged attention, not packed
multi-query prefill. No shallow KV is copied into uncomputed deep planes.
Each position/depth is computed exactly once. Proposals remain on the device and
-are read back with greedy target IDs once per round. Sampling retains the actual
-filtered q distribution and draws all drafts before acceptance, preserving the
-existing per-request RNG ordering. Different GEMM shapes can still change BF16
-results; same-seed agreement is not a general sampling correctness guarantee.
+greedy target IDs are read back once per mixed-depth wave. For sampling, the
+filtered proposal distribution `q` is retained for each draft and the common
+verifier applies rejection sampling with the request's generator from the shared
+`RngRegistry`. Different GEMM shapes and the changed draft/verify interleaving can
+change BF16 results or the fixed-seed token sequence; same-seed agreement is not a
+general sampling correctness guarantee.
Completed target logits are reordered into item/position order for the common
verifier. `SpeculativeResult` and engine commit stay unchanged: keep KV through
@@ -104,43 +122,35 @@ python -m benchmarks.loop_speculative \
--model /path/to/local/Ouro-1.4B \
--cuda-graphs --k 4 --concurrency 1 --new-tokens 32 --repeats 5 \
--output /tmp/loop-speculative.json
+
+# Sampling variant (top-p + top-k):
+python -m benchmarks.loop_speculative \
+ --model /path/to/local/Ouro-1.4B \
+ --cuda-graphs --k 4 --concurrency 8 --new-tokens 32 --repeats 5 \
+ --temperature 0.8 --top-p 0.8 --top-k 7 \
+ --output /tmp/loop-speculative-top-p-top-k.json
```
-The harness warms both engines, completes matched prefill and first-token
-generation before timing, ignores EOS to commit equal work, synchronizes timing
-boundaries, alternates trial order, and saves raw timings, outputs, first
-divergences and runtime metadata. `--prompts` accepts a JSON list of strings;
-`--backend flash_attn_4` requires compatible hardware. Peak memory includes both
-resident engines; the recorded checkpoint config hash is not a weight hash.
-This small run does not replace RFC #43's fixed-input real-model diagnostics,
-sampling distribution checks, long-context/serving matrix or BF16 acceptance
-criteria.
-
-### Local feasibility results (2026-10-07)
-
-Base commit: `ecb1f8b505b7e831815b40aec3b4598619cca23a`, plus this working-tree
-implementation. Local Ouro-1.4B, A100-SXM4-40GB, PyTorch 2.14.0+cu130,
-BF16/Triton, CUDA Graphs, d=2/D=4, 32 new tokens after prefill/first-token,
-five alternating measured repetitions after warmup. GPU 0 was idle before the
-experiment. These are the harness's two short default prompts, cycled to fill
-the batch; C=1 uses only the first. No workload generalization is implied.
-
-| Concurrency | K | Baseline tokens/s | Pipeline tokens/s | Throughput change |
-| --- | --- | --- | --- | --- |
-| 1 | 1 | 53.74 | 53.13 | -1.1% |
-| 1 | 2 | 57.18 | 55.90 | -2.2% |
-| 1 | 4 | 58.59 | 56.92 | -2.8% |
-| 1 | 8 | 44.92 | 43.17 | -3.9% |
-| 8 | 1 | 383.68 | 380.84 | -0.7% |
-| 8 | 2 | 409.14 | 402.28 | -1.7% |
-| 8 | 4 | 420.43 | 409.88 | -2.5% |
-| 8 | 8 | 324.42 | 316.53 | -2.4% |
-
-Throughput is fixed committed tokens divided by median elapsed decode time.
-Every measured sequence agreed with the matching original speculative run.
-This does not prove general BF16 equivalence or agreement with native decoding.
-FA4, long contexts, native real-model baselines and serving latency were not
-evaluated. On this bounded workload the prototype is correct in the exercised
-cases but **does not improve throughput**; it remains opt-in for experimentation.
-Raw JSON from this local run is retained under `artifacts/loop-speculative/`
-(ignored by git), including per-trial timings, outputs and first divergences.
+The harness warms both engines, completes a common prefill barrier before timing,
+ignores EOS to commit equal work, synchronizes timing boundaries, alternates trial
+order, and saves raw timings, outputs, first divergences and runtime metadata.
+`--prompts` accepts a JSON list of strings; `--backend flash_attn_4` requires
+compatible hardware. Peak memory includes both resident engines; the recorded
+checkpoint config hash is not a weight hash. This small run does not replace
+RFC #43's fixed-input real-model diagnostics, sampling distribution checks,
+long-context/serving matrix or BF16 acceptance criteria.
+
+### Rebase verification (2026-10-18)
+
+The current matrix was measured at `9586b4a` on an A100-SXM4-40GB with BF16,
+Triton, CUDA Graphs, `K=4`, `d=2/D=4`, 32 new tokens and five alternating
+repetitions. It covers C=1/4/8/16/32/64/128. Greedy results include baseline,
+loop and opt-in loop + early pruning; sampling results cover top-p (`p=0.8`) and
+top-p + top-k (`p=0.8,k=7`) for baseline and loop. The full tables and raw JSON
+paths are recorded in `notes/vllm-rlt/loop-speculative-rebase-2026-10-18.md`.
+
+The result is workload-dependent: greedy loop improves throughput at C=32/64,
+but loses at C=128; greedy pruning helps C=1 and follows the same trend at larger
+batches. Sampling has mixed results and can produce fixed-seed divergence because
+the loop changes RNG consumption order. The loop remains opt-in pending a broader
+profile and serving workload evaluation.
From f3a6a1dad3974d56cd7cdb8406cd612f6814c543 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 08:31:46 +0000
Subject: [PATCH 04/20] bench: expose CUDA graph batch limit
---
benchmarks/loop_speculative.py | 8 +++++++-
docs/loop_speculative.md | 2 ++
2 files changed, 9 insertions(+), 1 deletion(-)
diff --git a/benchmarks/loop_speculative.py b/benchmarks/loop_speculative.py
index 3be2f9e4..ae0ac223 100644
--- a/benchmarks/loop_speculative.py
+++ b/benchmarks/loop_speculative.py
@@ -89,6 +89,7 @@ def main():
parser.add_argument("--temperature", type=float, default=0.0)
parser.add_argument("--top-p", type=float, default=1.0)
parser.add_argument("--top-k", type=int, default=-1)
+ parser.add_argument("--cuda-graph-max-batch-size", type=int, default=128)
args = parser.parse_args()
if not args.model.is_dir():
parser.error("--model must be a complete local checkpoint directory")
@@ -98,6 +99,8 @@ def main():
parser.error("temperature must be nonnegative and top-p must be in (0, 1]")
if args.top_k == 0 or args.top_k < -1:
parser.error("top-k must be -1 or positive")
+ if args.cuda_graph_max_batch_size < 1:
+ parser.error("cuda-graph-max-batch-size must be positive")
torch.set_num_threads(1)
texts = (
json.loads(args.prompts.read_text())
@@ -125,7 +128,10 @@ def main():
prefill_chunk_size=max(map(len, prompts)),
max_prefill_batches_before_decode=1024,
)
- execution = ExecutionConfig(cuda_graphs=args.cuda_graphs)
+ execution = ExecutionConfig(
+ cuda_graphs=args.cuda_graphs,
+ cuda_graph_max_batch_size=args.cuda_graph_max_batch_size,
+ )
engines = {
name: LLMEngine(
model,
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 1ca4edf6..4de0f12e 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -134,6 +134,8 @@ python -m benchmarks.loop_speculative \
The harness warms both engines, completes a common prefill barrier before timing,
ignores EOS to commit equal work, synchronizes timing boundaries, alternates trial
order, and saves raw timings, outputs, first divergences and runtime metadata.
+For large loop batches, `--cuda-graph-max-batch-size 256` allows Graph capture of
+the 2C mixed-depth wave when C=128; the default is 128.
`--prompts` accepts a JSON list of strings; `--backend flash_attn_4` requires
compatible hardware. Peak memory includes both resident engines; the recorded
checkpoint config hash is not a weight hash. This small run does not replace
From dbaea320013129a084969cd3d7b661d0d53756ea Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 08:39:35 +0000
Subject: [PATCH 05/20] docs: use graph256 result for c128 benchmark
---
docs/loop_speculative.md | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 4de0f12e..d6a19a3d 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -151,8 +151,10 @@ loop and opt-in loop + early pruning; sampling results cover top-p (`p=0.8`) and
top-p + top-k (`p=0.8,k=7`) for baseline and loop. The full tables and raw JSON
paths are recorded in `notes/vllm-rlt/loop-speculative-rebase-2026-10-18.md`.
-The result is workload-dependent: greedy loop improves throughput at C=32/64,
-but loses at C=128; greedy pruning helps C=1 and follows the same trend at larger
-batches. Sampling has mixed results and can produce fixed-seed divergence because
-the loop changes RNG consumption order. The loop remains opt-in pending a broader
-profile and serving workload evaluation.
+The result is workload-dependent: with the default graph limit of 128, the loop
+loses at C=128 because its mixed-depth rows fall back to eager execution. Raising
+`cuda_graph_max_batch_size` to 256 recovers C=128 and gives about +3.4% for the
+loop and +2.8% for loop + greedy pruning. Sampling results below the graph256
+ablation use the original graph limit of 128 and are mixed; they can also produce
+fixed-seed divergence because the loop changes RNG consumption order. The loop
+remains opt-in pending a broader profile and serving workload evaluation.
From 7200a75c2209f0f4ab5668185cee32048b49acf9 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 08:57:00 +0000
Subject: [PATCH 06/20] docs: record graph256 sampling benchmark
---
docs/loop_speculative.md | 11 ++++++-----
1 file changed, 6 insertions(+), 5 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index d6a19a3d..320538b6 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -153,8 +153,9 @@ paths are recorded in `notes/vllm-rlt/loop-speculative-rebase-2026-10-18.md`.
The result is workload-dependent: with the default graph limit of 128, the loop
loses at C=128 because its mixed-depth rows fall back to eager execution. Raising
-`cuda_graph_max_batch_size` to 256 recovers C=128 and gives about +3.4% for the
-loop and +2.8% for loop + greedy pruning. Sampling results below the graph256
-ablation use the original graph limit of 128 and are mixed; they can also produce
-fixed-seed divergence because the loop changes RNG consumption order. The loop
-remains opt-in pending a broader profile and serving workload evaluation.
+`cuda_graph_max_batch_size` to 256 removes that C=128 fallback; the greedy loop is
+about +3.4% faster and loop + greedy pruning about +2.8% faster in the recorded
+run. The C=128 sampling rows use the same graph256 setting, while the other
+sampling points retain graph max 128. Sampling can still produce fixed-seed
+divergence because the loop changes RNG consumption order. The loop remains
+opt-in pending a broader profile and serving workload evaluation.
From 2ce6d21081136609ac8022c9bfec45fbf64eb9b1 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 09:04:11 +0000
Subject: [PATCH 07/20] docs: clarify loop speculative usage and results
---
docs/loop_speculative.md | 77 ++++++++++++++++++++++++++++------------
1 file changed, 54 insertions(+), 23 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 320538b6..948249e6 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -1,6 +1,6 @@
# Loop-level speculative execution
-This opt-in synchronous prototype implements the mixed-depth batching portion of
+This opt-in synchronous path implements the mixed-depth batching portion of
[RFC #43](https://github.com/ThinkFlowLab/vllm-rlt/issues/43): advance shallow work
for later token positions in the same recurrent call as deep work for earlier
positions. It retains the existing finite speculative window and atomic round
@@ -40,7 +40,7 @@ possible benefit comes from the shape and timing of mixed-depth batches; metadat
Python scheduling, Coda fragmentation and dynamic CUDA Graph shapes are additional
costs.
-## Schedule and state
+## Execution and state
For K=2, d=2 and D=4, A is the last committed input, B and C are proposals.
Numbers below are zero-based depths; columns are successive recurrent calls.
@@ -77,7 +77,7 @@ the initial input and accepted candidates; invalidate every depth's rejected
suffix; leave correction/bonus unforwarded. Engine EOS/length handling and
exception cleanup remain responsible for the request lifecycle.
-## Scope and benefit assessment
+## Scope and benefit
The default finite-window path does **not** implement cross-round hidden-state
retention or the complete continuous LoopSpec algorithm. The optional greedy
@@ -96,6 +96,54 @@ verifier's packed FA4 attention. Speedup must be measured; mixing depths alone
does not guarantee an improvement. Reducing round bubbles or wasted suffix work
would require a follow-up design for acceptance and state lifetime.
+## Parameters and results
+
+The reference measurements use an Ouro-1.4B checkpoint on an NVIDIA A100-SXM4-40GB,
+BF16, Triton attention, CUDA Graphs, `K=4`, `draft_loops=2`, `target_loops=4`,
+32 generated tokens and five timed repetitions per point. The benchmark compares
+the original block verifier with the loop pipeline using the same prompts and
+prefill barrier. Throughput is committed decode tokens per second.
+
+For C=1 through C=64, the CUDA Graph limit is 128. The C=128 rows use limit 256
+because the loop's largest mixed-depth wave is approximately `2C` rows. The C=64
+wave is at most 128 rows and therefore remains within the default limit.
+
+### Greedy
+
+| Concurrency | Baseline | Loop | Change | Loop + greedy prune | Change |
+| ---: | ---: | ---: | ---: | ---: | ---: |
+| 1 | 58.6 | 56.9 | -2.8% | 66.7 | +13.8% |
+| 4 | 217.2 | 214.7 | -1.1% | 215.5 | -0.7% |
+| 8 | 420.4 | 409.9 | -2.5% | 412.4 | -1.9% |
+| 16 | 790.8 | 780.1 | -1.4% | 782.3 | -1.0% |
+| 32 | 1321.2 | 1474.1 | +11.6% | 1468.6 | +11.1% |
+| 64 | 2434.1 | 2649.6 | +8.9% | 2625.5 | +7.7% |
+| 128 | 4092.9 | 4231.2 | +3.4% | 4207.7 | +2.8% |
+
+### Sampling
+
+Sampling uses `temperature=0.8`, `top_p=0.8`, with `top_k=-1` for top-p and
+`top_k=7` for top-p + top-k. Early pruning is disabled. The C=128 rows use
+CUDA Graph max 256; all other rows use max 128.
+
+| Concurrency | Top-p baseline | Top-p loop | Change | Top-p + top-k baseline | Top-p + top-k loop | Change |
+| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| 1 | 61.8 | 58.4 | -5.5% | 51.6 | 46.8 | -9.4% |
+| 4 | 205.2 | 159.9 | -22.0% | 140.3 | 138.0 | -1.6% |
+| 8 | 295.1 | 330.2 | +11.9% | 236.8 | 254.9 | +7.7% |
+| 16 | 524.3 | 493.9 | -5.8% | 417.3 | 393.2 | -5.8% |
+| 32 | 669.1 | 728.7 | +8.9% | 520.3 | 533.6 | +2.6% |
+| 64 | 870.7 | 908.3 | +4.3% | 628.2 | 649.4 | +3.4% |
+| 128 | 1057.5 | 953.4 | -9.8% | 730.6 | 805.7 | +10.3% |
+
+The results are workload-dependent. Loop-level scheduling does not reduce the
+number of recurrent rows; its potential benefit comes from mixed-depth batch
+shapes and device utilization. Metadata preparation, Python scheduling, Coda
+fragmentation and graph capture shape can outweigh that benefit. Sampling may
+also diverge from baseline at a fixed seed because draft sampling and rejection
+sampling are interleaved differently; this does not by itself indicate an
+incorrect sampling distribution.
+
## Validation
```bash
@@ -111,9 +159,7 @@ Tests exercise mixed depths including non-divisible d/D, actual hidden/logit/KV
agreement against a serial reference, every rejection boundary and bonus,
greedy/sampling batches, K=1/2/4/8, shortened windows, EOS, prefix reuse,
incremental allocation, cancellation and partial execution failure. CUDA coverage
-checks eager/graph parity and one greedy readback per round. The inherited GPU
-readback test also replaces its obsolete `OuroConfig.tiny()` call with the current
-test helper.
+checks eager/graph parity and one greedy readback per round.
For a bounded real-checkpoint feasibility run:
@@ -142,20 +188,5 @@ checkpoint config hash is not a weight hash. This small run does not replace
RFC #43's fixed-input real-model diagnostics, sampling distribution checks,
long-context/serving matrix or BF16 acceptance criteria.
-### Rebase verification (2026-10-18)
-
-The current matrix was measured at `9586b4a` on an A100-SXM4-40GB with BF16,
-Triton, CUDA Graphs, `K=4`, `d=2/D=4`, 32 new tokens and five alternating
-repetitions. It covers C=1/4/8/16/32/64/128. Greedy results include baseline,
-loop and opt-in loop + early pruning; sampling results cover top-p (`p=0.8`) and
-top-p + top-k (`p=0.8,k=7`) for baseline and loop. The full tables and raw JSON
-paths are recorded in `notes/vllm-rlt/loop-speculative-rebase-2026-10-18.md`.
-
-The result is workload-dependent: with the default graph limit of 128, the loop
-loses at C=128 because its mixed-depth rows fall back to eager execution. Raising
-`cuda_graph_max_batch_size` to 256 removes that C=128 fallback; the greedy loop is
-about +3.4% faster and loop + greedy pruning about +2.8% faster in the recorded
-run. The C=128 sampling rows use the same graph256 setting, while the other
-sampling points retain graph max 128. Sampling can still produce fixed-seed
-divergence because the loop changes RNG consumption order. The loop remains
-opt-in pending a broader profile and serving workload evaluation.
+Raw benchmark JSON is stored under `artifacts/loop-speculative/`; the extended
+run notes include the exact runtime metadata and per-repetition timings.
From 853b653f076fcc567fc8fb9cadb1a58a039af3ab Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 09:57:07 +0000
Subject: [PATCH 08/20] docs: explain issue scope and pruning limits
---
docs/loop_speculative.md | 100 +++++++++++++++++++++++++++++++++++++--
1 file changed, 96 insertions(+), 4 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 948249e6..8ce76510 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -1,7 +1,7 @@
# Loop-level speculative execution
-This opt-in synchronous path implements the mixed-depth batching portion of
-[RFC #43](https://github.com/ThinkFlowLab/vllm-rlt/issues/43): advance shallow work
+This opt-in synchronous path implements the mixed-depth batching portion proposed
+in [Issue #43](https://github.com/ThinkFlowLab/vllm-rlt/issues/43): advance shallow work
for later token positions in the same recurrent call as deep work for earlier
positions. It retains the existing finite speculative window and atomic round
commit. The default remains the original block verifier.
@@ -22,7 +22,9 @@ Greedy-only early suffix pruning can be enabled separately with
`--early-rejection-pruning`; it is off by default. Top-p and top-k sampling are
supported by the same loop pipeline, but sampling still verifies the complete
finite window before applying rejection sampling. Early pruning is intentionally
-greedy-only so it does not change the sampling RNG schedule.
+greedy-only because greedy acceptance is a deterministic candidate/argmax
+comparison, while sampling acceptance requires the target and proposal
+distributions, random draws and ordered correction sampling.
Existing restrictions apply: `last_exited`, full fixed target depth, synchronous
refill scheduling, and no preemption or PD.
@@ -42,6 +44,50 @@ costs.
## Execution and state
+### Loop-level flow
+
+The loop path keeps one branch for each active token position. A branch carries
+its request index, token offset, recurrent hidden state and current loop depth.
+Each wave advances every active branch by one depth. When a branch reaches the
+draft depth, its coda output starts a child branch while the parent continues
+toward the target depth. When a branch reaches the target depth, its target
+logits are saved for the existing verifier.
+
+```mermaid
+flowchart TD
+ A[Start speculative round] --> B[Create one branch per input position]
+ B --> C[Initialize hidden state and depth]
+ C --> D[Advance active branches by one recurrent depth]
+ D --> E{Draft depth reached?}
+ E -- Yes --> F[Run coda and generate draft token]
+ F --> G[Create child branch for next position]
+ G --> H{Target depth reached?}
+ E -- No --> H
+ H -- Yes --> I[Run coda and save target logits]
+ H -- No --> J[Keep branch active]
+ I --> K{Greedy pruning enabled?}
+ K -- Target mismatch --> L[Stop future suffix branches]
+ K -- Match or disabled --> M[Keep request active]
+ J --> N{Active branches remain?}
+ L --> N
+ M --> N
+ N -- Yes --> D
+ N -- No --> O[Reorder target rows by request and position]
+ O --> P{Sampling mode?}
+ P -- Greedy --> Q[Compare candidates with target argmax]
+ P -- Top-p or top-k --> R[Apply rejection sampling using p and q]
+ Q --> S[Commit accepted prefix and correction or bonus]
+ R --> S
+ S --> T[Invalidate rejected KV suffix]
+ T --> U[End speculative round]
+```
+
+This PR currently implements only the synchronous finite-window path. Mixed-
+depth waves are launched sequentially from the host; there are no CUDA streams,
+asynchronous scheduler overlap, cross-round hidden-state retention, preemption,
+or PD execution. The existing `SpeculativeResult` and atomic KV commit/rollback
+boundary are intentionally preserved.
+
For K=2, d=2 and D=4, A is the last committed input, B and C are proposals.
Numbers below are zero-based depths; columns are successive recurrent calls.
@@ -88,6 +134,27 @@ verification. Greedy target and draft IDs are read back once per mixed-depth
wave; the host-side rejection decision remains the pruning synchronization
boundary. No new CUDA streams or asynchronous scheduler are introduced.
+### Why pruning is greedy-only
+
+For greedy decoding, a target row produces one deterministic token ID. Once that
+ID is available, the runner can compare it with the candidate and stop creating
+later branches for that request after the first mismatch. The already submitted
+rows still finish, but suffix rows that have not entered a wave are avoided.
+
+Top-p and top-k use rejection sampling. Acceptance depends on the filtered target
+distribution `p`, the filtered draft distribution `q`, and a random uniform draw;
+on rejection, the correction token is drawn from `normalize(max(p - q, 0))`.
+A target argmax mismatch is therefore not equivalent to rejection: a candidate
+can differ from the target argmax and still be accepted, while a candidate that
+matches the argmax can still be rejected by the random draw. Applying the full
+decision as soon as each target row is ready would also require retaining every
+proposal distribution, preserving the request's RNG order, sampling a residual
+correction token, and invalidating all descendants generated from the rejected
+candidate. The current sampling path instead computes the finite window and
+applies the existing verifier in order, preserving its distribution and RNG
+behavior. Device-side batched acceptance can support sampling pruning in a later
+change.
+
For a uniform window, both the original block verifier and this pipeline require
`d*K + D` recurrent calls and `(K+1)*D` recurrent rows. Pipelining changes batch
shapes and head batching, not total recurrent work or that critical path. Its
@@ -144,6 +211,28 @@ also diverge from baseline at a fixed seed because draft sampling and rejection
sampling are interleaved differently; this does not by itself indicate an
incorrect sampling distribution.
+In summary, the synchronous loop path is a clear win only at selected higher
+concurrency points: greedy throughput improves by 3-12% from C=32 in this
+benchmark, while low-concurrency runs regress by about 1-3%. Greedy pruning is
+strongest for a single request with an early mismatch (+13.8% at C=1), but adds
+little at larger batches because work already submitted in the current wave
+cannot be cancelled. Sampling remains workload-sensitive, ranging from a 22%
+regression to a 12% improvement for top-p and from a 9% regression to a 10%
+improvement for top-p plus top-k. These results show that mixed-depth scheduling
+alone does not guarantee a speedup; asynchronous overlap, state retention,
+device-side acceptance and more efficient dynamic compaction remain follow-up
+work.
+
+At low concurrency, metadata preparation, Python scheduling and fragmented Coda
+work outweigh the mixed-depth benefit. The measured greedy loop becomes faster
+from C=32 onward, reaching its best gain at C=32; at C=128 it remains faster only
+after raising the Graph limit so the `2C` wave can be captured. Greedy pruning is
+helpful for a single request with early rejection, but adds little at larger
+concurrency because much of the suffix work has already entered a wave. Sampling
+is more workload-sensitive: top-p is slower at C=128, while top-p + top-k is
+faster there, showing that the loop benefit is not a general consequence of
+sampling or concurrency alone.
+
## Validation
```bash
@@ -189,4 +278,7 @@ RFC #43's fixed-input real-model diagnostics, sampling distribution checks,
long-context/serving matrix or BF16 acceptance criteria.
Raw benchmark JSON is stored under `artifacts/loop-speculative/`; the extended
-run notes include the exact runtime metadata and per-repetition timings.
+run notes include the runtime metadata and per-repetition timings. This work is
+an opt-in implementation of the loop-level direction in Issue #43; it does not
+claim that mixed-depth scheduling reduces total recurrent work or improves every
+hardware and workload combination.
From 41bf7a60df9be9fa4c7ce40b82b895d03eebf29f Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 10:12:52 +0000
Subject: [PATCH 09/20] docs: add loop pipeline schedule diagram
---
docs/assets/loop-schedule-comparison.svg | 101 +++++++++++++++++++++++
docs/loop_speculative.md | 37 ++-------
2 files changed, 110 insertions(+), 28 deletions(-)
create mode 100644 docs/assets/loop-schedule-comparison.svg
diff --git a/docs/assets/loop-schedule-comparison.svg b/docs/assets/loop-schedule-comparison.svg
new file mode 100644
index 00000000..5ed66215
--- /dev/null
+++ b/docs/assets/loop-schedule-comparison.svg
@@ -0,0 +1,101 @@
+
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 8ce76510..da1740cb 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -53,34 +53,15 @@ draft depth, its coda output starts a child branch while the parent continues
toward the target depth. When a branch reaches the target depth, its target
logits are saved for the existing verifier.
-```mermaid
-flowchart TD
- A[Start speculative round] --> B[Create one branch per input position]
- B --> C[Initialize hidden state and depth]
- C --> D[Advance active branches by one recurrent depth]
- D --> E{Draft depth reached?}
- E -- Yes --> F[Run coda and generate draft token]
- F --> G[Create child branch for next position]
- G --> H{Target depth reached?}
- E -- No --> H
- H -- Yes --> I[Run coda and save target logits]
- H -- No --> J[Keep branch active]
- I --> K{Greedy pruning enabled?}
- K -- Target mismatch --> L[Stop future suffix branches]
- K -- Match or disabled --> M[Keep request active]
- J --> N{Active branches remain?}
- L --> N
- M --> N
- N -- Yes --> D
- N -- No --> O[Reorder target rows by request and position]
- O --> P{Sampling mode?}
- P -- Greedy --> Q[Compare candidates with target argmax]
- P -- Top-p or top-k --> R[Apply rejection sampling using p and q]
- Q --> S[Commit accepted prefix and correction or bonus]
- R --> S
- S --> T[Invalidate rejected KV suffix]
- T --> U[End speculative round]
-```
+
+
+The upper lane shows the baseline's block draft followed by block target
+verification. The lower lane shows the loop pipeline: after A reaches the draft
+depth, B starts at depth 0 while A continues at a deeper level; C starts two
+waves later. The target rows are reordered into the common verifier only after
+their corresponding branches reach the target depth. This is the loop-level
+pipeline added in this PR; the waves are still launched sequentially on the
+current host path.
This PR currently implements only the synchronous finite-window path. Mixed-
depth waves are launched sequentially from the host; there are no CUDA streams,
From e9e5ef1863378ef8b9fdedebdc2f288392e638b4 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 11:47:08 +0000
Subject: [PATCH 10/20] test: document issue 43 loop validation
---
docs/loop_speculative.md | 108 ++++++++++++++++++++++++++------------
tests/test_speculative.py | 92 +++++++++++++++++++++++++++++++-
2 files changed, 165 insertions(+), 35 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index da1740cb..4f2cda74 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -144,11 +144,33 @@ verifier's packed FA4 attention. Speedup must be measured; mixing depths alone
does not guarantee an improvement. Reducing round bubbles or wasted suffix work
would require a follow-up design for acceptance and state lifetime.
+## Correctness checks for Issue #43
+
+The fixed-input oracle test runs the same token chain through native serial
+full-depth execution and the reuse path. It compares every loop-2 hidden row,
+every depth/layer KV row and the final target logits, and checks max absolute
+error, RMS error and target argmax mismatch rate in the assertion contract. The
+FP32 case is used to diagnose indexing or dependency mistakes; the BF16 case uses
+relaxed numerical tolerances to assess real model behavior.
+
+Free greedy generation is checked at `K=1/2/4/8`, with the existing ragged
+single-request cases and an eight-request batch. The test requires exact token
+and exit-depth agreement with native generation and reports the first divergent
+position in the failure message. Sampling is checked independently with small
+analytical distributions after applying the actual top-p and top-k filters; the
+empirical samples must match the filtered target distribution within the
+sampling tolerance. Fixed-seed text equality is not required for sampling.
+
+```bash
+python -m pytest -q tests/test_speculative.py \
+ -k 'reused_hidden_logits or issue43_greedy or filtered_top_p'
+```
+
## Parameters and results
The reference measurements use an Ouro-1.4B checkpoint on an NVIDIA A100-SXM4-40GB,
BF16, Triton attention, CUDA Graphs, `K=4`, `draft_loops=2`, `target_loops=4`,
-32 generated tokens and five timed repetitions per point. The benchmark compares
+32 generated tokens and three timed repetitions per point. The benchmark compares
the original block verifier with the loop pipeline using the same prompts and
prefill barrier. Throughput is committed decode tokens per second.
@@ -160,13 +182,25 @@ wave is at most 128 rows and therefore remains within the default limit.
| Concurrency | Baseline | Loop | Change | Loop + greedy prune | Change |
| ---: | ---: | ---: | ---: | ---: | ---: |
-| 1 | 58.6 | 56.9 | -2.8% | 66.7 | +13.8% |
-| 4 | 217.2 | 214.7 | -1.1% | 215.5 | -0.7% |
-| 8 | 420.4 | 409.9 | -2.5% | 412.4 | -1.9% |
-| 16 | 790.8 | 780.1 | -1.4% | 782.3 | -1.0% |
-| 32 | 1321.2 | 1474.1 | +11.6% | 1468.6 | +11.1% |
-| 64 | 2434.1 | 2649.6 | +8.9% | 2625.5 | +7.7% |
-| 128 | 4092.9 | 4231.2 | +3.4% | 4207.7 | +2.8% |
+| 1 | 58.6 | 56.9 | -2.8% | 66.7 | +13.9% |
+| 4 | 217.1 | 214.7 | -1.1% | 215.5 | -0.7% |
+| 8 | 420.3 | 409.8 | -2.5% | 412.0 | -2.0% |
+| 16 | 790.7 | 780.3 | -1.3% | 782.8 | -1.1% |
+| 32 | 1319.5 | 1473.7 | +11.7% | 1465.0 | +10.4% |
+| 64 | 2428.5 | 2648.4 | +9.1% | 2628.8 | +7.9% |
+| 128 | 4085.1 | 4229.6 | +3.5% | 4203.6 | +3.1% |
+
+The additional greedy counters show where pruning saves work:
+
+| Concurrency | Loop verified rows / committed tokens | Loop + prune verified rows / committed tokens | Loop peak allocated |
+| ---: | ---: | ---: | ---: |
+| 1 | 40 / 32 | 32 / 32 | 2.99 GiB |
+| 4 | 144 / 128 | 128 / 128 | 3.23 GiB |
+| 8 | 288 / 256 | 256 / 256 | 3.54 GiB |
+| 16 | 576 / 512 | 512 / 512 | 4.18 GiB |
+| 32 | 1152 / 1024 | 1024 / 1024 | 5.42 GiB |
+| 64 | 2304 / 2048 | 2048 / 2048 | 7.90 GiB |
+| 128 | 4608 / 4096 | 4096 / 4096 | 12.88 GiB |
### Sampling
@@ -176,21 +210,30 @@ CUDA Graph max 256; all other rows use max 128.
| Concurrency | Top-p baseline | Top-p loop | Change | Top-p + top-k baseline | Top-p + top-k loop | Change |
| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
-| 1 | 61.8 | 58.4 | -5.5% | 51.6 | 46.8 | -9.4% |
-| 4 | 205.2 | 159.9 | -22.0% | 140.3 | 138.0 | -1.6% |
-| 8 | 295.1 | 330.2 | +11.9% | 236.8 | 254.9 | +7.7% |
-| 16 | 524.3 | 493.9 | -5.8% | 417.3 | 393.2 | -5.8% |
-| 32 | 669.1 | 728.7 | +8.9% | 520.3 | 533.6 | +2.6% |
-| 64 | 870.7 | 908.3 | +4.3% | 628.2 | 649.4 | +3.4% |
-| 128 | 1057.5 | 953.4 | -9.8% | 730.6 | 805.7 | +10.3% |
+| 1 | 61.9 | 58.5 | -5.5% | 51.6 | 46.8 | -9.4% |
+| 4 | 205.1 | 159.9 | -22.0% | 139.7 | 137.8 | -1.3% |
+| 8 | 295.4 | 330.2 | +11.8% | 236.3 | 254.2 | +7.5% |
+| 16 | 523.8 | 492.9 | -5.9% | 417.2 | 393.2 | -5.8% |
+| 32 | 662.2 | 723.9 | +9.3% | 519.3 | 532.2 | +2.5% |
+| 64 | 876.7 | 914.1 | +4.3% | 625.2 | 646.8 | +3.5% |
+| 128 | 1046.6 | 940.0 | -10.2% | 727.1 | 801.2 | +10.2% |
+
+The benchmark JSON also records median/min/max decode time, peak allocated
+memory, recurrent rounds, drafted/accepted tokens, verified rows and the first
+divergent output position. Greedy runs have identical outputs at every point;
+sampling divergence is expected when the same request generator is consumed in
+a different draft/verify order. It is not used as a distribution-correctness
+criterion.
The results are workload-dependent. Loop-level scheduling does not reduce the
number of recurrent rows; its potential benefit comes from mixed-depth batch
shapes and device utilization. Metadata preparation, Python scheduling, Coda
-fragmentation and graph capture shape can outweigh that benefit. Sampling may
-also diverge from baseline at a fixed seed because draft sampling and rejection
-sampling are interleaved differently; this does not by itself indicate an
-incorrect sampling distribution.
+fragmentation and graph capture shape can outweigh that benefit. For example,
+the greedy loop reduces verified rows from 4608 to 4096 at C=128 only when
+early pruning is enabled, while the non-pruned loop keeps the same recurrent
+row count as baseline. The sampling rows also show that acceptance rate and
+round count change with the sampled path, so throughput should be read together
+with the recorded counters.
In summary, the synchronous loop path is a clear win only at selected higher
concurrency points: greedy throughput improves by 3-12% from C=32 in this
@@ -204,16 +247,6 @@ alone does not guarantee a speedup; asynchronous overlap, state retention,
device-side acceptance and more efficient dynamic compaction remain follow-up
work.
-At low concurrency, metadata preparation, Python scheduling and fragmented Coda
-work outweigh the mixed-depth benefit. The measured greedy loop becomes faster
-from C=32 onward, reaching its best gain at C=32; at C=128 it remains faster only
-after raising the Graph limit so the `2C` wave can be captured. Greedy pruning is
-helpful for a single request with early rejection, but adds little at larger
-concurrency because much of the suffix work has already entered a wave. Sampling
-is more workload-sensitive: top-p is slower at C=128, while top-p + top-k is
-faster there, showing that the loop benefit is not a general consequence of
-sampling or concurrency alone.
-
## Validation
```bash
@@ -231,18 +264,25 @@ greedy/sampling batches, K=1/2/4/8, shortened windows, EOS, prefix reuse,
incremental allocation, cancellation and partial execution failure. CUDA coverage
checks eager/graph parity and one greedy readback per round.
+The focused Issue #43 checks can be run without the full suite:
+
+```bash
+python -m pytest -q tests/test_speculative.py \
+ -k 'reused_hidden_logits or issue43_greedy or filtered_top_p'
+```
+
For a bounded real-checkpoint feasibility run:
```bash
python -m benchmarks.loop_speculative \
--model /path/to/local/Ouro-1.4B \
- --cuda-graphs --k 4 --concurrency 1 --new-tokens 32 --repeats 5 \
+ --cuda-graphs --k 4 --concurrency 1 --new-tokens 32 --repeats 3 \
--output /tmp/loop-speculative.json
# Sampling variant (top-p + top-k):
python -m benchmarks.loop_speculative \
--model /path/to/local/Ouro-1.4B \
- --cuda-graphs --k 4 --concurrency 8 --new-tokens 32 --repeats 5 \
+ --cuda-graphs --k 4 --concurrency 8 --new-tokens 32 --repeats 3 \
--temperature 0.8 --top-p 0.8 --top-k 7 \
--output /tmp/loop-speculative-top-p-top-k.json
```
@@ -254,9 +294,9 @@ For large loop batches, `--cuda-graph-max-batch-size 256` allows Graph capture o
the 2C mixed-depth wave when C=128; the default is 128.
`--prompts` accepts a JSON list of strings; `--backend flash_attn_4` requires
compatible hardware. Peak memory includes both resident engines; the recorded
-checkpoint config hash is not a weight hash. This small run does not replace
-RFC #43's fixed-input real-model diagnostics, sampling distribution checks,
-long-context/serving matrix or BF16 acceptance criteria.
+checkpoint config hash is not a weight hash. These checks cover fixed-input
+reuse and sampling distribution contracts, but do not replace Issue #43's
+long-context, serving-latency or full BF16 acceptance matrix.
Raw benchmark JSON is stored under `artifacts/loop-speculative/`; the extended
run notes include the runtime metadata and per-repetition timings. This work is
diff --git a/tests/test_speculative.py b/tests/test_speculative.py
index bdcd70ee..01b1601f 100644
--- a/tests/test_speculative.py
+++ b/tests/test_speculative.py
@@ -81,6 +81,31 @@ def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget, pipeline
assert llm.engine.cache_manager.num_used_blocks == 0
+@pytest.mark.parametrize("k", [1, 2, 4, 8])
+def test_issue43_greedy_generation_agrees_at_eight_concurrent_requests(k):
+ """Free generation stays byte-for-byte equal for the issue's K sweep."""
+ m = model()
+ prompts = [[2 + i, 7 - (i % 4), 11 + (i % 3)] for i in range(8)]
+ params = [SamplingParams(max_tokens=7, ignore_eos=True) for _ in prompts]
+ native = LLM(m).generate(prompts, params)
+ speculative = LLM(
+ m,
+ speculative_config=SpeculativeConfig(k, loop_level_pipelining=True),
+ cache_config=CacheConfig(256, 2),
+ ).generate(prompts, params)
+ for expected, actual in zip(native, speculative):
+ first = next(
+ (
+ i
+ for i, (left, right) in enumerate(zip(expected.token_ids, actual.token_ids))
+ if left != right
+ ),
+ None,
+ )
+ assert actual.token_ids == expected.token_ids, f"first divergence at {first}"
+ assert actual.exit_depths == expected.exit_depths
+
+
@pytest.mark.parametrize(
"device,backend",
[
@@ -91,7 +116,9 @@ def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget, pipeline
)
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("pipelined", [False, True])
-def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, backend, pipelined):
+def test_reused_hidden_logits_and_all_kv_match_serial_oracle(
+ dtype, device, backend, pipelined, record_property
+):
if device == "cuda" and dtype == torch.float32 and backend == "flash_attn_4":
pytest.skip("FA4 requires 16-bit operands")
m = model(dtype=dtype).to(device)
@@ -146,14 +173,23 @@ def record_verify(items, drafted, proposals, logits):
oracle.allocate("r", 16)
atol, rtol = (4e-5, 4e-5) if dtype == torch.float32 else (0.08, 0.04)
max_logit_error = 0.0
+ hidden_errors = []
+ kv_errors = []
+ logit_errors = []
+ logit_argmax_mismatches = 0
+ logit_rows = 0
for pos, token in enumerate(tokens):
h = m.prelude(torch.tensor([token], device=device))
for d in range(4):
h, _ = m.recurrent(h, ["r"], [d], [pos], oracle)
if pos >= 3:
+ hidden_errors.append((seen[d, pos] - h[0]).detach().float().reshape(-1))
torch.testing.assert_close(seen[d, pos], h[0], atol=atol, rtol=rtol)
if pos >= 3:
actual, expected = verified_logits[pos - 3], m.coda(h)[0]
+ logit_errors.append((actual - expected).detach().float().reshape(-1))
+ logit_argmax_mismatches += int(actual.argmax() != expected.argmax())
+ logit_rows += 1
max_logit_error = max(max_logit_error, float((actual - expected).abs().max()))
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
assert actual.argmax() == expected.argmax()
@@ -162,6 +198,7 @@ def record_verify(items, drafted, proposals, logits):
for actual, expected in zip(
cache.read(layer, "r", d, 7), oracle.read(layer, "r", d, 7)
):
+ kv_errors.append((actual - expected).detach().float().reshape(-1))
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
# Every input/depth computed exactly once: no shallow recomputation in verify.
assert work == Counter({0: 4, 1: 4, 2: 4, 3: 4})
@@ -176,6 +213,36 @@ def record_verify(items, drafted, proposals, logits):
)
assert max_logit_error < (1e-4 if dtype == torch.float32 else 0.025)
+ # Keep the RFC #43 fixed-input metrics visible in the assertion contract:
+ # max absolute and RMS error are tracked independently for hidden, KV and
+ # final-logit reuse, and target argmaxes must not diverge.
+ def metrics(errors):
+ values = torch.cat(errors)
+ return float(values.abs().max()), float(values.square().mean().sqrt())
+
+ hidden_max, hidden_rms = metrics(hidden_errors)
+ kv_max, kv_rms = metrics(kv_errors)
+ logits_max, logits_rms = metrics(logit_errors)
+ record_property("hidden_max_abs", hidden_max)
+ record_property("hidden_rms", hidden_rms)
+ record_property("kv_max_abs", kv_max)
+ record_property("kv_rms", kv_rms)
+ record_property("logits_max_abs", logits_max)
+ record_property("logits_rms", logits_rms)
+ record_property("target_argmax_mismatch_rate", logit_argmax_mismatches / logit_rows)
+ assert logit_argmax_mismatches == 0, (
+ f"fixed-input target argmax mismatch rate={logit_argmax_mismatches / logit_rows:.6f}"
+ )
+ assert hidden_max <= (4e-5 if dtype == torch.float32 else 0.08), (
+ f"hidden max_abs={hidden_max:.6g}, rms={hidden_rms:.6g}"
+ )
+ assert kv_max <= (4e-5 if dtype == torch.float32 else 0.08), (
+ f"KV max_abs={kv_max:.6g}, rms={kv_rms:.6g}"
+ )
+ assert logits_max <= (1e-4 if dtype == torch.float32 else 0.025), (
+ f"logits max_abs={logits_max:.6g}, rms={logits_rms:.6g}"
+ )
+
@pytest.mark.parametrize("rejected", [0, 1, 2, 3])
def test_each_rejection_position_and_bonus_commit_correct_frontier(rejected, monkeypatch):
@@ -435,6 +502,29 @@ def test_probability_filters_match_expected_and_apply_strict_top_k():
torch.testing.assert_close(p[p > 0], torch.ones(1))
+@pytest.mark.parametrize(
+ "params",
+ [
+ SamplingParams(temperature=0.8, top_p=0.8),
+ SamplingParams(temperature=0.8, top_k=3, top_p=0.9),
+ ],
+)
+def test_rejection_sampling_respects_filtered_top_p_and_top_k(params):
+ """The acceptance rule targets the filtered p distribution, not raw logits."""
+ target_logits = torch.tensor([3.0, 2.1, 1.5, 0.4, -0.2, -1.0])
+ proposal_logits = torch.tensor([1.0, 2.8, 2.0, 1.2, 0.5, -0.5])
+ p = Sampler.probabilities(target_logits, params)
+ q = Sampler.probabilities(proposal_logits, params)
+ generator = torch.Generator().manual_seed(9173)
+ counts = torch.zeros_like(p)
+ for _ in range(16000):
+ candidate = int(torch.multinomial(q, 1, generator=generator))
+ token, _ = rejection_sample(candidate, p, q, generator)
+ counts[token] += 1
+ empirical = counts / counts.sum()
+ torch.testing.assert_close(empirical, p, atol=0.02, rtol=0)
+
+
@pytest.mark.parametrize(
"kwargs",
[
From ada1e4efe3e70288eaef4c5e504d47daa58935df Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 11:58:36 +0000
Subject: [PATCH 11/20] docs: clarify sampling results and throughput units
---
docs/loop_speculative.md | 14 +++++++++-----
1 file changed, 9 insertions(+), 5 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 4f2cda74..a92d5dd8 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -53,8 +53,6 @@ draft depth, its coda output starts a child branch while the parent continues
toward the target depth. When a branch reaches the target depth, its target
logits are saved for the existing verifier.
-
-
The upper lane shows the baseline's block draft followed by block target
verification. The lower lane shows the loop pipeline: after A reaches the draft
depth, B starts at depth 0 while A continues at a deeper level; C starts two
@@ -161,6 +159,11 @@ analytical distributions after applying the actual top-p and top-k filters; the
empirical samples must match the filtered target distribution within the
sampling tolerance. Fixed-seed text equality is not required for sampling.
+In the focused 16,000-draw check, the maximum absolute deviation from the
+filtered target distribution was 0.003085 for top-p and 0.002593 for top-p plus
+top-k; the corresponding RMS deviations were 0.001781 and 0.001368. Both are
+below the test tolerance of 0.02.
+
```bash
python -m pytest -q tests/test_speculative.py \
-k 'reused_hidden_logits or issue43_greedy or filtered_top_p'
@@ -172,7 +175,8 @@ The reference measurements use an Ouro-1.4B checkpoint on an NVIDIA A100-SXM4-40
BF16, Triton attention, CUDA Graphs, `K=4`, `draft_loops=2`, `target_loops=4`,
32 generated tokens and three timed repetitions per point. The benchmark compares
the original block verifier with the loop pipeline using the same prompts and
-prefill barrier. Throughput is committed decode tokens per second.
+prefill barrier. Throughput is aggregate committed decode throughput across all
+concurrent requests, measured in output tokens per second (`output tok/s`).
For C=1 through C=64, the CUDA Graph limit is 128. The C=128 rows use limit 256
because the loop's largest mixed-depth wave is approximately `2C` rows. The C=64
@@ -180,7 +184,7 @@ wave is at most 128 rows and therefore remains within the default limit.
### Greedy
-| Concurrency | Baseline | Loop | Change | Loop + greedy prune | Change |
+| Concurrency | Baseline (output tok/s) | Loop (output tok/s) | Change | Loop + greedy prune (output tok/s) | Change |
| ---: | ---: | ---: | ---: | ---: | ---: |
| 1 | 58.6 | 56.9 | -2.8% | 66.7 | +13.9% |
| 4 | 217.1 | 214.7 | -1.1% | 215.5 | -0.7% |
@@ -208,7 +212,7 @@ Sampling uses `temperature=0.8`, `top_p=0.8`, with `top_k=-1` for top-p and
`top_k=7` for top-p + top-k. Early pruning is disabled. The C=128 rows use
CUDA Graph max 256; all other rows use max 128.
-| Concurrency | Top-p baseline | Top-p loop | Change | Top-p + top-k baseline | Top-p + top-k loop | Change |
+| Concurrency | Top-p baseline (output tok/s) | Top-p loop (output tok/s) | Change | Top-p + top-k baseline (output tok/s) | Top-p + top-k loop (output tok/s) | Change |
| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| 1 | 61.9 | 58.5 | -5.5% | 51.6 | 46.8 | -9.4% |
| 4 | 205.1 | 159.9 | -22.0% | 139.7 | 137.8 | -1.3% |
From e27f9f501ed1ef415f53a34f6b18066350076254 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 11:59:21 +0000
Subject: [PATCH 12/20] docs: clarify aggregate decode throughput
---
docs/loop_speculative.md | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index a92d5dd8..a9dd557b 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -53,13 +53,11 @@ draft depth, its coda output starts a child branch while the parent continues
toward the target depth. When a branch reaches the target depth, its target
logits are saved for the existing verifier.
-The upper lane shows the baseline's block draft followed by block target
-verification. The lower lane shows the loop pipeline: after A reaches the draft
-depth, B starts at depth 0 while A continues at a deeper level; C starts two
-waves later. The target rows are reordered into the common verifier only after
-their corresponding branches reach the target depth. This is the loop-level
-pipeline added in this PR; the waves are still launched sequentially on the
-current host path.
+The baseline executes all draft positions before block target verification. In
+the loop schedule, after A reaches the draft depth, B starts at depth 0 while A
+continues at a deeper level; C starts two waves later. Target rows are reordered
+into the common verifier only after their corresponding branches reach target
+depth. The waves are still launched sequentially on the current host path.
This PR currently implements only the synchronous finite-window path. Mixed-
depth waves are launched sequentially from the host; there are no CUDA streams,
@@ -177,6 +175,8 @@ BF16, Triton attention, CUDA Graphs, `K=4`, `draft_loops=2`, `target_loops=4`,
the original block verifier with the loop pipeline using the same prompts and
prefill barrier. Throughput is aggregate committed decode throughput across all
concurrent requests, measured in output tokens per second (`output tok/s`).
+It counts generated output tokens only, excludes prompt/prefill tokens and is
+not a per-request rate.
For C=1 through C=64, the CUDA Graph limit is 128. The C=128 rows use limit 256
because the loop's largest mixed-depth wave is approximately `2C` rows. The C=64
From 3415844b100e3be722c688d00d0b33849e9f43fc Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 12:00:12 +0000
Subject: [PATCH 13/20] test: record filtered sampling errors
---
tests/test_speculative.py | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/tests/test_speculative.py b/tests/test_speculative.py
index 01b1601f..c17aa40d 100644
--- a/tests/test_speculative.py
+++ b/tests/test_speculative.py
@@ -509,7 +509,7 @@ def test_probability_filters_match_expected_and_apply_strict_top_k():
SamplingParams(temperature=0.8, top_k=3, top_p=0.9),
],
)
-def test_rejection_sampling_respects_filtered_top_p_and_top_k(params):
+def test_rejection_sampling_respects_filtered_top_p_and_top_k(params, record_property):
"""The acceptance rule targets the filtered p distribution, not raw logits."""
target_logits = torch.tensor([3.0, 2.1, 1.5, 0.4, -0.2, -1.0])
proposal_logits = torch.tensor([1.0, 2.8, 2.0, 1.2, 0.5, -0.5])
@@ -522,6 +522,9 @@ def test_rejection_sampling_respects_filtered_top_p_and_top_k(params):
token, _ = rejection_sample(candidate, p, q, generator)
counts[token] += 1
empirical = counts / counts.sum()
+ error = empirical - p
+ record_property("distribution_max_abs_error", float(error.abs().max()))
+ record_property("distribution_rms_error", float(error.square().mean().sqrt()))
torch.testing.assert_close(empirical, p, atol=0.02, rtol=0)
From 68c724bd063ca0afcc2604590744bb298fab234a Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Tue, 20 Oct 2026 07:32:41 +0000
Subject: [PATCH 14/20] docs: correct memory units and summary
---
docs/loop_speculative.md | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index a9dd557b..2abead84 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -196,7 +196,7 @@ wave is at most 128 rows and therefore remains within the default limit.
The additional greedy counters show where pruning saves work:
-| Concurrency | Loop verified rows / committed tokens | Loop + prune verified rows / committed tokens | Loop peak allocated |
+| Concurrency | Loop verified rows / committed tokens | Loop + prune verified rows / committed tokens | Loop peak allocated (GB, decimal) |
| ---: | ---: | ---: | ---: |
| 1 | 40 / 32 | 32 / 32 | 2.99 GiB |
| 4 | 144 / 128 | 128 / 128 | 3.23 GiB |
@@ -242,7 +242,7 @@ with the recorded counters.
In summary, the synchronous loop path is a clear win only at selected higher
concurrency points: greedy throughput improves by 3-12% from C=32 in this
benchmark, while low-concurrency runs regress by about 1-3%. Greedy pruning is
-strongest for a single request with an early mismatch (+13.8% at C=1), but adds
+strongest for a single request with an early mismatch (+13.9% at C=1), but adds
little at larger batches because work already submitted in the current wave
cannot be cancelled. Sampling remains workload-sensitive, ranging from a 22%
regression to a 12% improvement for top-p and from a 9% regression to a 10%
From 0388ea0802c2f75ee1b96c4897baf034a9b8be35 Mon Sep 17 00:00:00 2001
From: 0z5a
Date: Sat, 24 Oct 2026 17:59:30 +0200
Subject: [PATCH 15/20] Run loop-level speculation as a planned submit/collect
round
Port the mixed-depth round onto main's asynchronous speculative runner:
- Plan every wave's metadata and gather/scatter indices on the host and
upload them once per round, instead of building metadata per wave.
- Keep draft IDs on the device; the round reads back once at collection,
so it composes with deferred collection and multiple in-flight submits.
- Run the head only on draft-ready rows inside the loop and once over all
verification rows at the end, replacing per-wave target head calls,
per-branch stacks and logit clones.
- Reuse the block round's collection, so both modes share greedy and
rejection-sampling acceptance and consume each request's generator in
the same order; same-seed sampled outputs now match the block round.
Drop greedy early-rejection pruning: it needs a host readback every wave,
which the asynchronous submission path cannot provide.
Co-Authored-By: Claude Opus 5.5
---
tests/test_loop_speculative.py | 369 ++++++++++++++-------------
tests/test_speculative.py | 101 +++++---
vllm_rlt/config.py | 9 +-
vllm_rlt/entrypoints/runtime_args.py | 6 -
vllm_rlt/worker/speculative.py | 126 ++++++++-
5 files changed, 382 insertions(+), 229 deletions(-)
diff --git a/tests/test_loop_speculative.py b/tests/test_loop_speculative.py
index a2da7e0a..09d7a199 100644
--- a/tests/test_loop_speculative.py
+++ b/tests/test_loop_speculative.py
@@ -1,4 +1,4 @@
-"""Mixed-depth execution, round equivalence and pipeline rollback contracts."""
+"""Mixed-depth schedule, round equivalence and pipeline rollback contracts."""
import argparse
from collections import Counter
@@ -6,196 +6,167 @@
import pytest
import torch
+from tests.helpers import LocalLLM as LLM
from tests.test_speculative import drain, engine, model
-from vllm_rlt import LLM, CacheConfig, SamplingParams, SpeculativeConfig
-from vllm_rlt.core.scheduler import ScheduledItem, SchedulerOutput
+from vllm_rlt import CacheConfig, SamplingParams, SpeculativeConfig
+from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput
from vllm_rlt.entrypoints.runtime_args import add_runtime_args, runtime_configs
-from vllm_rlt.request import Request, Stage
-from vllm_rlt.worker.speculative import SpeculativeResult
+from vllm_rlt.request import Request, RequestPhase
+from vllm_rlt.worker.speculative import loop_waves
+
+
+@pytest.mark.parametrize("draft_loops,target_loops", [(1, 4), (2, 4), (3, 4), (1, 2)])
+@pytest.mark.parametrize("counts", [(1,), (5,), (5, 3, 1), (2, 9)])
+def test_loop_waves_cover_each_depth_once_after_its_history(counts, draft_loops, target_loops):
+ waves = loop_waves(counts, draft_loops, target_loops)
+ assert len(waves) == max((c - 1) * draft_loops + target_loops for c in counts)
+ seen = Counter()
+ for wave, (rows, drafts) in enumerate(waves):
+ for i, offset, depth in rows:
+ assert depth == wave - offset * draft_loops
+ # The previous position at this depth ran in an earlier wave.
+ assert offset == 0 or seen[i, offset - 1, depth]
+ # This position's previous depth ran in the previous wave.
+ assert depth == 0 or seen[i, offset, depth - 1]
+ seen.update(rows)
+ assert list(drafts) == [
+ row
+ for row, (i, offset, depth) in enumerate(rows)
+ if depth == draft_loops - 1 and offset + 1 < counts[i]
+ ]
+ # One draft per request and wave keeps proposals in offset order.
+ assert len({rows[row][0] for row in drafts}) == len(drafts)
+ assert seen == Counter(
+ (i, offset, depth)
+ for i, count in enumerate(counts)
+ for offset in range(count)
+ for depth in range(target_loops)
+ )
+ if max(counts) > 1 and target_loops > draft_loops:
+ assert any(len({depth for _, _, depth in rows}) > 1 for rows, _ in waves)
-@pytest.mark.parametrize("draft_depth", [1, 2, 3])
+def _round(m, pipelined, draft_loops, sampling):
+ e = engine(m, k=4, pipelined=pipelined, cache_config=CacheConfig(128, 2))
+ runner = e.speculative_runner
+ runner.config = SpeculativeConfig(4, draft_loops, 4, pipelined)
+ cache = e.cache_manager
+ spans = []
+ for i, count in enumerate((5, 3, 1)):
+ rid = str(i)
+ prefix = [2, 3 + i]
+ cache.allocate(rid, 16)
+ for pos, token in enumerate(prefix):
+ h = m.prelude(torch.tensor([token]))
+ for depth in range(4):
+ h, _ = m.recurrent(h, [rid], [depth], [pos], cache)
+ params = SamplingParams(
+ temperature=0.8 if sampling and i != 1 else 0,
+ top_k=7,
+ top_p=0.8,
+ seed=42,
+ ignore_eos=True,
+ )
+ request = Request(rid, prefix, params, generated_token_ids=[7 + i])
+ spans.append(ScheduledRequest(request, len(prefix), count))
+ calls = []
+ core = m.recurrent_prepared
+
+ def record(hidden, batch, cache, **kwargs):
+ calls.append([(id(a), depth, pos) for a, depth, pos in batch.rows])
+ return core(hidden, batch, cache, **kwargs)
+
+ heads = []
+ hook = m.lm_head.register_forward_hook(
+ lambda module, inputs, output: heads.append(output.detach().clone())
+ )
+ m.recurrent_prepared = record
+ try:
+ results = runner.submit(SchedulerOutput(RequestPhase.SPECULATIVE, spans)).collect()
+ finally:
+ m.recurrent_prepared = core
+ hook.remove()
+ kv = [
+ tensor.clone()
+ for span in spans
+ for depth in range(4)
+ for layer in range(m.config.num_hidden_layers)
+ for tensor in cache.read(
+ layer, span.request.request_id, depth, span.token_start + span.token_count
+ )
+ ]
+ return results.speculative, heads[-1], kv, runner.stats, calls
+
+
+@pytest.mark.parametrize("draft_loops", [1, 2, 3])
@pytest.mark.parametrize("sampling", [False, True])
-def test_pipeline_round_matches_baseline_and_batches_mixed_depths(draft_depth, sampling):
+def test_loop_round_matches_block_round_and_batches_mixed_depths(draft_loops, sampling):
m = model()
- snapshots = []
- for pipelined in (False, True):
- e = engine(m)
- runner = e.speculative_runner
- runner.config = SpeculativeConfig(4, draft_depth, 4, pipelined)
- cache = e.cache_manager
- items = []
- for i, count in enumerate((5, 3, 1)):
- rid = str(i)
- prefix = [2, 3 + i]
- cache.allocate(rid, 16)
- for pos, token in enumerate(prefix):
- h = m.prelude(torch.tensor([token]))
- for depth in range(4):
- h, _ = m.recurrent(h, [rid], [depth], [pos], cache)
- params = SamplingParams(
- temperature=0.8 if sampling and i != 1 else 0,
- top_k=7,
- top_p=0.8,
- seed=42,
- ignore_eos=True,
- )
- request = Request(rid, prefix, params, generated_token_ids=[7 + i])
- items.append(ScheduledItem(request, len(prefix), count))
- calls, work, heads = [], Counter(), []
- core, verify = runner._core, runner._verify
-
- def record(hidden, ids, positions, depth, **kwargs):
- depths = [row[1] for row in kwargs["batch"].rows]
- calls.append(list(zip(ids, positions, depths)))
- work.update(zip(ids, positions, depths))
- return core(hidden, ids, positions, depth, **kwargs)
-
- def record_verify(items, drafted, proposals, logits):
- heads.append(logits.clone())
- return verify(items, drafted, proposals, logits)
-
- runner._core, runner._verify = record, record_verify
- results = runner.execute(SchedulerOutput(Stage.SPECULATIVE, items))
- expected_work = Counter(
- (item.request.request_id, pos, depth)
- for item in items
- for pos in range(item.token_start, item.token_start + item.token_count)
- for depth in range(4)
- )
- assert work == expected_work # Each position/depth is computed once.
- if pipelined:
- assert any(len({d for _, _, d in call}) > 1 for call in calls)
- # History at a given depth must finish before its later position.
- completed = set()
- for call in calls:
- for rid, pos, depth in call:
- assert pos == 2 or (rid, pos - 1, depth) in completed
- completed.update(call)
- assert len(calls) == draft_depth * 4 + 4
- kv = [
- tensor.clone()
- for item in items
- for depth in range(4)
- for layer in range(m.config.num_hidden_layers)
- for tensor in cache.read(
- layer, item.request.request_id, depth, item.token_start + item.token_count
- )
- ]
- snapshots.append((results, heads[0], kv, runner.stats))
- baseline, pipeline = snapshots
- assert baseline[0] == pipeline[0]
- torch.testing.assert_close(baseline[1], pipeline[1], atol=4e-5, rtol=4e-5)
- for a, b in zip(baseline[2], pipeline[2]):
+ block = _round(m, False, draft_loops, sampling)
+ loop = _round(m, True, draft_loops, sampling)
+ # Same proposals, verification rows, KV and same-seed sampling decisions.
+ assert loop[0] == block[0]
+ torch.testing.assert_close(loop[1], block[1], atol=4e-5, rtol=4e-5)
+ for a, b in zip(loop[2], block[2]):
torch.testing.assert_close(a, b, atol=4e-5, rtol=4e-5)
- assert baseline[3] == pipeline[3]
+ assert loop[3] == block[3]
+ calls = loop[4]
+ work = Counter(row for call in calls for row in call)
+ assert len(work) == sum(work.values()) == (5 + 3 + 1) * 4
+ assert len(calls) == 4 * draft_loops + 4
+ assert any(len({depth for _, depth, _ in call}) > 1 for call in calls)
-@pytest.mark.parametrize("rejected", [0, 1, 2, 3])
-@pytest.mark.parametrize("temperature", [0, 1])
-def test_pipeline_correction_bonus_and_next_round(rejected, temperature, monkeypatch):
+@pytest.mark.parametrize("temperature", [0.0, 0.8])
+def test_loop_round_runs_only_draft_heads_inside_the_loop(temperature, monkeypatch):
m = model()
- e = engine(m, pipelined=True, cache_config=CacheConfig(128, 2))
- e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=12, ignore_eos=True))
+ e = engine(m, k=3, pipelined=True)
+ e.add_request(
+ "r", [2, 3], SamplingParams(max_tokens=9, ignore_eos=True, temperature=temperature)
+ )
while not e.step():
pass
- request = e.scheduler.requests["r"]
- request.sampling_params = SamplingParams(
- max_tokens=12, ignore_eos=True, temperature=temperature, seed=42
- )
- start = request.position
runner = e.speculative_runner
- core, coda = runner._core_mixed, runner._coda
- head_rows = []
-
- def record(hidden, ids, positions, depths):
- nonlocal head_rows
- output = core(hidden, ids, positions, depths)
- # Force proposal/target distributions by semantic position and depth,
- # independent of which other rows share this batched LM-head call.
- head_rows = [
- (p - start, False) for p, d in zip(positions, depths) if d == 1 and p - start < 3
- ]
- head_rows += [(p - start, True) for p, d in zip(positions, depths) if d == 3]
- return output
-
- def controlled(hidden):
- assert len(hidden) == len(head_rows)
- logits = torch.full((len(hidden), m.config.vocab_size), -torch.inf)
- for row, (offset, target) in enumerate(head_rows):
- token = 20 + offset if target and offset >= rejected else 10 + offset
- logits[row, token] = 0
- return logits
-
- monkeypatch.setattr(runner, "_core_mixed", record)
- monkeypatch.setattr(runner, "_coda", controlled)
- output = e.step()[0]
- count = rejected + 1
- assert output.token_ids[-count:] == list(range(10, 10 + rejected)) + [20 + rejected]
- assert runner.stats.accepted_tokens == rejected
- for plane in e.cache_manager._get_allocation("r").written:
- assert all(w.prefix == start + count and not w.pending for w in plane)
- monkeypatch.setattr(runner, "_core_mixed", core)
- monkeypatch.setattr(runner, "_coda", coda)
- # Replay a common committed prefix with native greedy to expose stale KV.
- request.sampling_params = SamplingParams(max_tokens=12, ignore_eos=True)
- expected = LLM(m).generate(
- [output.prompt_token_ids + output.token_ids],
- SamplingParams(max_tokens=12 - len(output.token_ids), ignore_eos=True),
- )[0]
- assert drain(e)["r"].token_ids[len(output.token_ids) :] == expected.token_ids
- assert e.cache_manager.num_used_blocks == 0
+ rounds, heads = [], []
+ submit, coda = runner._submit_loop_round, runner._coda
+ monkeypatch.setattr(runner, "_submit_round", None)
+ monkeypatch.setattr(
+ runner, "_submit_loop_round", lambda *args: rounds.append(1) or submit(*args)
+ )
+ monkeypatch.setattr(runner, "_coda", lambda hidden: heads.append(len(hidden)) or coda(hidden))
+ e.step()
+ # K single-row draft heads, then one head over all K + 1 verification rows.
+ assert rounds == [1] and heads == [1, 1, 1, 4]
-def test_greedy_early_rejection_prunes_unstarted_suffix(monkeypatch):
+def test_loop_generation_agrees_at_eight_concurrent_requests():
m = model()
- e = engine(m, pipelined=True, cache_config=CacheConfig(128, 2))
- runner = e.speculative_runner
- runner.config = SpeculativeConfig(4, 2, 4, True, True)
- cache = e.cache_manager
- request = Request(
- "r", [2, 3], SamplingParams(max_tokens=12, ignore_eos=True), generated_token_ids=[7]
- )
- cache.allocate(request.request_id, 16)
- for pos, token in enumerate(request.prompt_token_ids):
- hidden = m.prelude(torch.tensor([token]))
- for depth in range(4):
- hidden, _ = m.recurrent(hidden, [request.request_id], [depth], [pos], cache)
- item = ScheduledItem(request, len(request.prompt_token_ids), 4)
-
- calls = []
- core, coda = runner._core_mixed, runner._coda
- head_rows = []
-
- def record(hidden, ids, positions, depths):
- nonlocal head_rows
- calls.append([(p - item.token_start, d) for p, d in zip(positions, depths)])
- output = core(hidden, ids, positions, depths)
- head_rows = [(p - item.token_start, False) for p, d in zip(positions, depths) if d == 1]
- head_rows += [(p - item.token_start, True) for p, d in zip(positions, depths) if d == 3]
- return output
-
- def controlled(hidden):
- assert len(hidden) == len(head_rows)
- logits = torch.full((len(hidden), m.config.vocab_size), -torch.inf)
- for row, (offset, target) in enumerate(head_rows):
- token = 20 if target and offset == 0 else 10 + offset
- logits[row, token] = 0
- return logits
+ prompts = [[2 + i, 7 - (i % 4), 11 + (i % 3)] for i in range(8)]
+ params = [SamplingParams(max_tokens=7, ignore_eos=True) for _ in prompts]
+ native = LLM(m).generate(prompts, params)
+ for k in (1, 2, 4, 8):
+ actual = LLM(
+ m,
+ speculative_config=SpeculativeConfig(k, loop_level_pipelining=True),
+ cache_config=CacheConfig(256, 2),
+ ).generate(prompts, params)
+ assert [o.token_ids for o in actual] == [o.token_ids for o in native]
+ assert [o.exit_depths for o in actual] == [o.exit_depths for o in native]
+
+
+def test_loop_sampling_matches_block_round_for_the_same_seed():
+ # Both rounds consume each request's generator in the same order: drafts by
+ # offset, then verification. Same-seed outputs therefore agree.
+ m = model()
+ prompts = [[2, 3], [7, 8, 9], [4]]
+ params = SamplingParams(max_tokens=11, temperature=0.8, top_k=7, top_p=0.8, seed=5)
- monkeypatch.setattr(runner, "_core_mixed", record)
- monkeypatch.setattr(runner, "_coda", controlled)
- result = runner.execute(SchedulerOutput(Stage.SPECULATIVE, [item]))[0]
+ def run(pipelined):
+ return LLM(
+ m, speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined)
+ ).generate(prompts, params)
- assert result == SpeculativeResult([20], 0, 1)
- assert runner.stats.verified_rows == 1
- assert runner.stats.drafted_tokens == 1
- assert all(offset < 2 for call in calls for offset, _ in call)
- assert any(offset == 1 and depth == 1 for call in calls for offset, depth in call)
- assert not any(offset == 1 and depth >= 2 for call in calls for offset, depth in call)
- assert not any(offset >= 2 for call in calls for offset, _ in call)
- monkeypatch.setattr(runner, "_core_mixed", core)
- monkeypatch.setattr(runner, "_coda", coda)
+ assert [o.token_ids for o in run(True)] == [o.token_ids for o in run(False)]
def test_pipeline_config_and_cli():
@@ -206,35 +177,65 @@ def test_pipeline_config_and_cli():
parser = argparse.ArgumentParser()
add_runtime_args(parser)
parser.set_defaults(num_blocks=128, block_size=16)
- args = parser.parse_args(
- ["--speculative-tokens", "3", "--loop-level-pipelining", "--early-rejection-pruning"]
- )
- config = runtime_configs(args)["speculative_config"]
- assert config.loop_level_pipelining and config.early_rejection_pruning
+ args = parser.parse_args(["--speculative-tokens", "3", "--loop-level-pipelining"])
+ assert runtime_configs(args)["speculative_config"].loop_level_pipelining
args = parser.parse_args(["--loop-level-pipelining"])
with pytest.raises(ValueError, match="requires --speculative-tokens"):
runtime_configs(args)
def test_pipeline_partial_execution_failure_releases_kv(monkeypatch):
- e = engine(model(), pipelined=True)
+ m = model()
+ e = engine(m, pipelined=True)
e.add_request("r", [2, 3], SamplingParams(max_tokens=12, ignore_eos=True))
while not e.step():
pass
- runner = e.speculative_runner
- core = runner._core_mixed
+ core = m.recurrent_prepared
calls = 0
- def fail_after_write(*args):
+ def fail_after_write(*args, **kwargs):
nonlocal calls
- result = core(*args)
+ result = core(*args, **kwargs)
calls += 1
if calls == 3:
raise RuntimeError("partial pipeline failure")
return result
- monkeypatch.setattr(runner, "_core_mixed", fail_after_write)
+ monkeypatch.setattr(m, "recurrent_prepared", fail_after_write)
with pytest.raises(RuntimeError, match="partial pipeline failure"):
e.step()
assert not e.has_unfinished_requests()
assert e.cache_manager.num_used_blocks == 0
+
+
+@pytest.mark.parametrize("rejected", [0, 1, 2, 3])
+def test_loop_rejection_rolls_back_and_next_round_matches_native(rejected, monkeypatch):
+ m = model()
+ e = engine(m, k=3, pipelined=True, cache_config=CacheConfig(128, 2))
+ e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=12, ignore_eos=True))
+ while not e.step():
+ pass
+ coda = m.coda
+ calls = 0
+
+ def controlled(h):
+ nonlocal calls
+ logits = torch.full((len(h), m.config.vocab_size), -100.0)
+ if calls < 3:
+ logits[:, 10 + calls] = 100
+ else:
+ for i in range(4):
+ logits[i, 10 + i if i < rejected else 20 + i] = 100
+ calls += 1
+ return logits
+
+ monkeypatch.setattr(m, "coda", controlled)
+ out = e.step()[0]
+ assert out.token_ids[-rejected - 1 :] == list(range(10, 10 + rejected)) + [20 + rejected]
+ monkeypatch.setattr(m, "coda", coda)
+ expected = LLM(m).generate(
+ [out.prompt_token_ids + out.token_ids],
+ SamplingParams(max_tokens=12 - len(out.token_ids), ignore_eos=True),
+ )[0]
+ assert drain(e)["r"].token_ids[len(out.token_ids) :] == expected.token_ids
+ assert e.cache_manager.num_used_blocks == 0
diff --git a/tests/test_speculative.py b/tests/test_speculative.py
index bf7dbfa9..683aa65e 100644
--- a/tests/test_speculative.py
+++ b/tests/test_speculative.py
@@ -29,8 +29,10 @@ def model(seed=123, dtype=torch.float32):
return OuroForCausalLM(tiny_ouro_config()).to(dtype=dtype)
-def engine(m, k=3, **kwargs):
- return LLMEngine(m, speculative_config=SpeculativeConfig(k), **kwargs)
+def engine(m, k=3, *, pipelined=False, **kwargs):
+ return LLMEngine(
+ m, speculative_config=SpeculativeConfig(k, loop_level_pipelining=pipelined), **kwargs
+ )
def drain(e):
@@ -70,14 +72,15 @@ def test_speculative_prefill_routes_by_declared_capabilities(
@pytest.mark.parametrize("k", [1, 2, 4, 8])
@pytest.mark.parametrize("budget", [1, 3, 32])
-def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget, pipelined):
m = model()
prompts = [[2], [3, 4, 5, 6, 7], [9, 3]]
params = [SamplingParams(max_tokens=n, ignore_eos=True) for n in [1, 13, 8]]
expected = LLM(m).generate(prompts, params)
llm = LLM(
m,
- speculative_config=SpeculativeConfig(k),
+ speculative_config=SpeculativeConfig(k, loop_level_pipelining=pipelined),
cache_config=CacheConfig(256, 2),
scheduler_config=SchedulerConfig(max_num_batched_tokens=budget, prefill_chunk_size=2),
)
@@ -96,7 +99,8 @@ def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget):
],
)
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
-def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, backend):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, backend, pipelined):
if device == "cuda" and dtype == torch.float32 and backend == "flash_attn_4":
pytest.skip("FA4 requires 16-bit operands")
m = model(dtype=dtype).to(device)
@@ -105,7 +109,9 @@ def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, back
m = OuroForCausalLM(tiny_ouro_config(hidden_size=256, head_dim=64)).to(
device=device, dtype=dtype
)
- e = engine(m, k=3, cache_config=CacheConfig(128, 2), attention_backend=backend)
+ e = engine(
+ m, pipelined=pipelined, k=3, cache_config=CacheConfig(128, 2), attention_backend=backend
+ )
cache = e.cache_manager
prefix = [3, 4, 5]
current = 7
@@ -182,9 +188,12 @@ def record(hidden, batch, cache, **kwargs):
@pytest.mark.parametrize("rejected", [0, 1, 2, 3])
-def test_each_rejection_position_and_bonus_commit_correct_frontier(rejected, monkeypatch):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_each_rejection_position_and_bonus_commit_correct_frontier(
+ rejected, monkeypatch, pipelined
+):
m = model()
- e = engine(m, k=3)
+ e = engine(m, pipelined=pipelined, k=3)
e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=12, ignore_eos=True))
while not e.step():
pass
@@ -233,9 +242,10 @@ def test_greedy_accept_commits_target_prefix_through_first_mismatch(k):
assert tokens == candidates[:rejected] + [7]
-def test_eos_in_accepted_prefix_never_delivers_following_tokens(monkeypatch):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_eos_in_accepted_prefix_never_delivers_following_tokens(monkeypatch, pipelined):
m = model()
- e = engine(m)
+ e = engine(m, pipelined=pipelined)
e.add_request("r", [2, 3], SamplingParams(max_tokens=12))
while not e.step():
pass
@@ -267,14 +277,15 @@ def test_truncation_removes_sparse_validity_without_destroying_prefix():
@pytest.mark.parametrize("prefix_cache,incremental", [(False, True), (True, False), (True, True)])
-def test_prefix_replay_incremental_growth_cancel_and_reuse(prefix_cache, incremental):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_prefix_replay_incremental_growth_cancel_and_reuse(prefix_cache, incremental, pipelined):
m = model()
kwargs = dict(
cache_config=CacheConfig(
128, 2, enable_prefix_caching=prefix_cache, incremental_allocation=incremental
)
)
- llm = LLM(m, speculative_config=SpeculativeConfig(4), **kwargs)
+ llm = LLM(m, speculative_config=SpeculativeConfig(4, loop_level_pipelining=pipelined), **kwargs)
params = SamplingParams(max_tokens=9, ignore_eos=True)
prompts = [[2, 3, 4, 5, 6], [2, 3, 4, 5, 7]]
expected = LLM(m).generate(prompts, params)
@@ -294,7 +305,8 @@ def test_prefix_replay_incremental_growth_cancel_and_reuse(prefix_cache, increme
@pytest.mark.parametrize("top_k,top_p", [(-1, 1.0), (5, 1.0), (7, 0.8)])
-def test_sampling_seeded_replay_and_request_isolation(top_k, top_p):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_sampling_seeded_replay_and_request_isolation(top_k, top_p, pipelined):
m = model()
prompts = [[2, 3], [7, 8, 9]]
params = SamplingParams(
@@ -302,7 +314,9 @@ def test_sampling_seeded_replay_and_request_isolation(top_k, top_p):
)
def run(prompts):
- return LLM(m, speculative_config=SpeculativeConfig(3)).generate(prompts, params)
+ return LLM(
+ m, speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined)
+ ).generate(prompts, params)
batched = run(prompts)
assert [o.token_ids for o in run(prompts)] == [o.token_ids for o in batched]
@@ -455,8 +469,11 @@ def test_rejects_adaptive_or_wrong_target_requests():
assert not e.has_unfinished_requests()
-def test_dynamic_arrival_budget_and_failure_reclamation(monkeypatch):
- e = engine(model(), scheduler_config=SchedulerConfig(max_num_batched_tokens=5))
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_dynamic_arrival_budget_and_failure_reclamation(monkeypatch, pipelined):
+ e = engine(
+ model(), pipelined=pipelined, scheduler_config=SchedulerConfig(max_num_batched_tokens=5)
+ )
params = SamplingParams(max_tokens=9, ignore_eos=True)
e.add_request("a", [2, 3], params)
while not e.step():
@@ -486,7 +503,8 @@ def test_invalid_k(k):
@pytest.mark.gpu
@pytest.mark.parametrize("backend", ["triton", "flash_attn_4"])
-def test_gpu_sampling_and_greedy_match_replay_with_ragged_requests(backend):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_gpu_sampling_and_greedy_match_replay_with_ragged_requests(backend, pipelined):
torch.manual_seed(123)
m = OuroForCausalLM(tiny_ouro_config(hidden_size=256, head_dim=64)).to(
device="cuda", dtype=torch.bfloat16
@@ -500,7 +518,9 @@ def test_gpu_sampling_and_greedy_match_replay_with_ragged_requests(backend):
def run(speculative, sampling):
return LLM(
m,
- speculative_config=SpeculativeConfig(3) if speculative else None,
+ speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined)
+ if speculative
+ else None,
cache_config=CacheConfig(128, 16),
attention_backend=backend,
).generate(prompts, sampling)
@@ -513,14 +533,17 @@ def run(speculative, sampling):
@pytest.mark.gpu
-def test_greedy_round_reads_back_once_with_ragged_requests():
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_greedy_round_reads_back_once_with_ragged_requests(pipelined):
# 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(tiny_ouro_config(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")
+ e = engine(
+ m, pipelined=pipelined, 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
@@ -540,7 +563,8 @@ def test_greedy_round_reads_back_once_with_ragged_requests():
@pytest.mark.gpu
@pytest.mark.parametrize("max_graph_rows", [2, 8])
-def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows, pipelined):
m = model(dtype=torch.bfloat16).to("cuda")
prompts = [[2, 3, 4], [7, 8]]
params = SamplingParams(max_tokens=9, ignore_eos=True)
@@ -548,7 +572,7 @@ def test_speculative_cuda_graph_matches_eager_with_ragged_replay(max_graph_rows)
def run(graphs):
llm = LLM(
m,
- speculative_config=SpeculativeConfig(3),
+ speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined),
cache_config=CacheConfig(128, 16),
attention_backend="triton",
execution_config=ExecutionConfig(
@@ -581,7 +605,10 @@ def run(graphs):
pytest.param("cuda", "triton", True, marks=pytest.mark.gpu),
],
)
-def test_priority_preempts_only_between_speculative_rounds_and_resumes(device, backend, graphs):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_priority_preempts_only_between_speculative_rounds_and_resumes(
+ device, backend, graphs, pipelined
+):
m = (
model()
if device == "cpu"
@@ -601,6 +628,7 @@ def test_priority_preempts_only_between_speculative_rounds_and_resumes(device, b
e = engine(
m,
k=3,
+ pipelined=pipelined,
cache_config=CacheConfig(64, 16 if graphs else 2, incremental_allocation=True),
scheduler_config=SchedulerConfig(
max_num_seqs=1,
@@ -629,7 +657,8 @@ def test_priority_preempts_only_between_speculative_rounds_and_resumes(device, b
@pytest.mark.parametrize("temperature", [0, 0.8])
-def test_speculative_preemption_preserves_request_rng(temperature):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_speculative_preemption_preserves_request_rng(temperature, pipelined):
m = model()
params = {
"low": SamplingParams(
@@ -641,7 +670,9 @@ def test_speculative_preemption_preserves_request_rng(temperature):
}
prompts = {"low": [2, 3, 4], "high": [7, 8, 9]}
expected = {
- rid: LLM(m, speculative_config=SpeculativeConfig(3)).generate([prompt], params[rid])[0]
+ rid: LLM(
+ m, speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined)
+ ).generate([prompt], params[rid])[0]
for rid, prompt in prompts.items()
}
e = engine(
@@ -709,12 +740,16 @@ def test_abort_suspended_speculative_request_allows_id_reuse():
@pytest.mark.parametrize("chunk_size", [1, 3, 8])
@pytest.mark.parametrize("budget", [2, 7])
-def test_split_prefill_and_speculation_share_forward_and_match_atomic(chunk_size, budget):
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_split_prefill_and_speculation_share_forward_and_match_atomic(
+ chunk_size, budget, pipelined
+):
from vllm_rlt.request import Operation
def run(atomic):
e = engine(
model(),
+ pipelined=pipelined,
cache_config=CacheConfig(256, 2),
scheduler_config=SchedulerConfig(
max_num_batched_tokens=budget,
@@ -772,8 +807,11 @@ def test_common_submit_rejects_speculation_without_controller():
@pytest.mark.parametrize("temperature", [0.0, 0.7])
-def test_speculative_submit_defers_collection_and_collects_once(monkeypatch, temperature):
- e = engine(model())
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_speculative_submit_defers_collection_and_collects_once(
+ monkeypatch, temperature, pipelined
+):
+ e = engine(model(), pipelined=pipelined)
original = e.speculative_runner._collect_round
collections = []
@@ -841,8 +879,11 @@ def forbidden(*args, **kwargs):
@pytest.mark.gpu
@pytest.mark.parametrize("use_uva", [False, True])
-def test_speculative_input_slots_survive_three_submits_before_collection(use_uva):
- e = engine(model().cuda(), execution_config=ExecutionConfig(use_uva=use_uva))
+@pytest.mark.parametrize("pipelined", [False, True])
+def test_speculative_input_slots_survive_three_submits_before_collection(use_uva, pipelined):
+ e = engine(
+ model().cuda(), pipelined=pipelined, execution_config=ExecutionConfig(use_uva=use_uva)
+ )
try:
for i in range(3):
e.add_request(str(i), [i + 2, i + 3], SamplingParams(max_tokens=12, ignore_eos=True))
diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py
index 88d54f91..7dd0278a 100644
--- a/vllm_rlt/config.py
+++ b/vllm_rlt/config.py
@@ -152,16 +152,11 @@ class SpeculativeConfig:
draft_loops: int = 2
target_loops: int = 4
loop_level_pipelining: bool = False
- early_rejection_pruning: bool = False
def __post_init__(self):
for name in ("num_speculative_tokens", "draft_loops", "target_loops"):
_positive(name, getattr(self, name))
if self.draft_loops >= self.target_loops:
raise ValueError("draft_loops must be smaller than target_loops")
- for name in (
- "loop_level_pipelining",
- "early_rejection_pruning",
- ):
- if type(getattr(self, name)) is not bool:
- raise ValueError(f"{name} must be a boolean")
+ if type(self.loop_level_pipelining) is not bool:
+ raise ValueError("loop_level_pipelining must be a boolean")
diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py
index 16c9aaea..c519f98c 100644
--- a/vllm_rlt/entrypoints/runtime_args.py
+++ b/vllm_rlt/entrypoints/runtime_args.py
@@ -23,11 +23,6 @@ def add_runtime_args(parser):
action="store_true",
help="Batch active speculative branches at different recurrent depths",
)
- parser.add_argument(
- "--early-rejection-pruning",
- action="store_true",
- help="Stop greedy loop-level branches after the first rejected candidate",
- )
parser.add_argument("--enable-prefix-caching", action="store_true")
parser.add_argument("--incremental-kv", action="store_true")
parser.add_argument("--kv-watermark", type=float, default=0.0)
@@ -98,7 +93,6 @@ def runtime_configs(args):
draft_loops=args.draft_loops,
target_loops=args.target_loops,
loop_level_pipelining=args.loop_level_pipelining,
- early_rejection_pruning=args.early_rejection_pruning,
)
if args.speculative_tokens is not None
else None,
diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py
index c799e883..06cc2fc0 100644
--- a/vllm_rlt/worker/speculative.py
+++ b/vllm_rlt/worker/speculative.py
@@ -6,6 +6,7 @@
from contextlib import nullcontext
from dataclasses import dataclass
+from functools import lru_cache
import torch
@@ -63,6 +64,30 @@ def rejection_sample(candidate, target, proposal, generator):
return int(Sampler.draw(residual / total, generator).item()), False
+@lru_cache(maxsize=64)
+def loop_waves(token_counts, draft_loops, target_loops):
+ """Return the mixed-depth rows of one loop-level round, one tuple per wave.
+
+ Offset o of a request enters at wave o * draft_loops, after its parent's draft
+ head chose its input, and runs one depth per wave up to target_loops. Rows are
+ (request, offset, depth), ordered by request and offset. ``drafts`` lists the
+ rows whose state after the wave feeds the draft head for the next offset.
+ The schedule depends only on token counts, never on token values.
+ """
+ waves = []
+ for wave in range(max((count - 1) * draft_loops + target_loops for count in token_counts)):
+ rows, drafts = [], []
+ for i, count in enumerate(token_counts):
+ first = max(0, -(-(wave - target_loops + 1) // draft_loops))
+ for offset in range(first, min(count, wave // draft_loops + 1)):
+ depth = wave - offset * draft_loops
+ if depth == draft_loops - 1 and offset + 1 < count:
+ drafts.append(len(rows))
+ rows.append((i, offset, depth))
+ waves.append((tuple(rows), tuple(drafts)))
+ return tuple(waves)
+
+
@dataclass
class SpeculativeSubmission:
runner: "SpeculativeRunner"
@@ -211,8 +236,11 @@ def submit(self, batch):
if self.compute_stream is not None
else nullcontext()
):
- inputs, steps, verification = self._prepare_inputs(batch)
- submission = self._submit_round(batch, inputs, steps, verification)
+ if self.config.loop_level_pipelining:
+ submission = self._submit_loop_round(batch, *self._prepare_loop_inputs(batch))
+ else:
+ inputs, steps, verification = self._prepare_inputs(batch)
+ submission = self._submit_round(batch, inputs, steps, verification)
self.submissions.complete(submission.event)
return submission
except Exception:
@@ -263,6 +291,100 @@ def _submit_round(self, batch, inputs, steps, verification):
self.stats.verified_rows += sum(span.token_count for span in scheduled_requests)
return SpeculativeSubmission(self, batch, logits, drafted, proposals, event)
+ def _prepare_loop_inputs(self, batch):
+ """Plan every wave of a loop-level round and publish one snapshot per dtype.
+
+ Branch states live in one buffer whose rows follow verification order:
+ contiguous by request, then by offset. Each wave gathers its rows, runs
+ one mixed-depth recurrent call and scatters the states back.
+ """
+ for buffer in self.input_buffers.values():
+ buffer.reset()
+ requests = batch.requests
+ counts = tuple(span.token_count for span in requests)
+ starts = [0]
+ for count in counts:
+ starts.append(starts[-1] + count)
+ int64 = self.input_buffers[torch.int64]
+ tokens = int64.append([span.request.input_token_id for span in requests])
+ firsts = int64.append(starts[:-1])
+ planned = []
+ for rows, drafts in loop_waves(counts, self.config.draft_loops, self.config.target_loops):
+ layouts = self.metadata.build(
+ [requests[i].request.request_id for i, _, _ in rows],
+ [[depth for _, _, depth in rows]],
+ [requests[i].token_start + offset for i, offset, _ in rows],
+ )
+ metadata = stage_metadata(
+ layouts, lambda values, dtype: self.input_buffers[dtype].append(values)
+ )
+ flat = [starts[i] + offset for i, offset, _ in rows]
+ draft = None
+ if drafts:
+ draft = (
+ [rows[row][0] for row in drafts],
+ int64.append(drafts),
+ int64.append([flat[row] + 1 for row in drafts]),
+ )
+ planned.append((int64.append(flat), metadata, draft))
+ uploaded = {dtype: buffer.upload() for dtype, buffer in self.input_buffers.items()}
+
+ def bind(tensor):
+ return uploaded[tensor.dtype].as_strided(
+ tensor.shape, tensor.stride(), tensor.storage_offset()
+ )
+
+ waves = [
+ (
+ bind(flat),
+ tuple(self.metadata.bind(kv, uploaded) for kv in metadata),
+ None if draft is None else (draft[0], bind(draft[1]), bind(draft[2])),
+ )
+ for flat, metadata, draft in planned
+ ]
+ return bind(tokens), bind(firsts), starts[-1], waves
+
+ def _submit_loop_round(self, batch, tokens, firsts, total_rows, waves):
+ """Advance all positions one depth per wave; children start at draft depth.
+
+ A child starts draft_loops waves after its parent, so the parent's KV at
+ each depth is written before the child attends to it. Only draft-ready rows
+ run the head inside the loop; target logits come from one head call after
+ the last wave. Draft IDs stay on the device, as in the block round.
+ """
+ scheduled_requests = batch.requests
+ proposals = [[] for _ in scheduled_requests]
+ drafted = []
+ hidden = self.runtime.forward(ForwardBatch(Operation.PRELUDE, tokens=tokens)).hidden
+ states = hidden.new_empty((total_rows, hidden.shape[-1]))
+ states.index_copy_(0, firsts, hidden)
+ for flat, metadata, draft in waves:
+ hidden = self._recurrent(states.index_select(0, flat), metadata)
+ states.index_copy_(0, flat, hidden)
+ if draft is None:
+ continue
+ owners, rows, children = draft
+ logits = self._coda(hidden.index_select(0, rows))
+ inputs = logits.argmax(-1)
+ for row, i in enumerate(owners):
+ request = scheduled_requests[i].request
+ if request.sampling_params.temperature != 0:
+ q = Sampler.probabilities(logits[row], request.sampling_params)
+ proposals[i].append(q)
+ inputs[row] = Sampler.draw(q, self._acquire_generator(request, logits[row]))
+ drafted.append((owners, inputs))
+ child = self.runtime.forward(ForwardBatch(Operation.PRELUDE, tokens=inputs)).hidden
+ states.index_copy_(0, children, child)
+ logits = self._coda(states)
+ event = None
+ if self.device.type == "cuda":
+ event = torch.cuda.Event()
+ event.record(torch.cuda.current_stream(self.device))
+ self.stats.rounds += len(scheduled_requests)
+ self.stats.drafted_tokens += sum(len(rows) for rows, _ in drafted)
+ self.stats.verified_rows += total_rows
+ return SpeculativeSubmission(self, batch, logits, drafted, proposals, event)
+
@torch.inference_mode()
def _collect_round(self, submission):
scheduled_requests = submission.batch.requests
From 1d7c4010b2d9bba1a3f841654ad694f9270c516b Mon Sep 17 00:00:00 2001
From: 0z5a
Date: Sat, 24 Oct 2026 17:59:30 +0200
Subject: [PATCH 16/20] Pair benchmark trials and report interval estimates
Give every request a fixed seed so sampled arms draw the same random
numbers, report the median of per-repeat paired speedups with a bootstrap
interval, quartiles, step time, acceptance and CUDA graph counters, and
default to 128 new tokens and seven repeats.
Co-Authored-By: Claude Opus 5.5
---
benchmarks/loop_speculative.py | 98 ++++++++++++++++++++++++----------
1 file changed, 69 insertions(+), 29 deletions(-)
diff --git a/benchmarks/loop_speculative.py b/benchmarks/loop_speculative.py
index ae0ac223..93d8b794 100644
--- a/benchmarks/loop_speculative.py
+++ b/benchmarks/loop_speculative.py
@@ -1,12 +1,16 @@
-"""Matched synchronous baseline/pipeline feasibility check on a local Ouro model.
+"""Paired block/loop-level speculative decode comparison on a local Ouro model.
-Exclude prefill, interleave trials, and retain raw timings and generated tokens.
-This small experiment is not the full RFC #43 accuracy/performance acceptance matrix.
+Both engines share one resident model. Prefill is excluded, trial order alternates,
+and every request has a fixed seed, so both arms see the same RNG draws: sampled
+runs compare equal acceptance instead of two different random acceptance paths.
+Raw timings, generated tokens and CUDA graph counters are retained. This is not
+the full RFC #43 accuracy/performance acceptance matrix.
"""
import argparse
import hashlib
import json
+import random
import statistics
import subprocess
import time
@@ -27,16 +31,26 @@
from vllm_rlt.worker.speculative import SpeculativeStats
+def graph_counters(runner):
+ return {
+ name: None
+ if graphs is None
+ else dict(captures=graphs.captures, replays=graphs.replays, fallbacks=graphs.fallbacks)
+ for name, graphs in (("recurrent", runner.graphs), ("coda", runner.coda_graphs))
+ }
+
+
@torch.inference_mode()
-def trial(engine, prompts, new_tokens, temperature, top_p, top_k):
- params = SamplingParams(
- max_tokens=new_tokens + 1,
- ignore_eos=True,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- )
+def trial(engine, prompts, new_tokens, temperature, top_p, top_k, seed):
for i, prompt in enumerate(prompts):
+ params = SamplingParams(
+ max_tokens=new_tokens + 1,
+ ignore_eos=True,
+ temperature=temperature,
+ top_p=top_p,
+ top_k=top_k,
+ seed=seed + i,
+ )
engine.add_request(str(i), prompt, params)
# Hold decode until all requests have crossed the first-output barrier.
while any(not r.generated_token_ids for r in engine.scheduler.requests.values()):
@@ -46,7 +60,9 @@ def trial(engine, prompts, new_tokens, temperature, top_p, top_k):
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
outputs = {}
+ steps = 0
while engine.has_unfinished_requests():
+ steps += 1
for output in engine.step():
if output.finished:
outputs[output.request_id] = output.token_ids
@@ -61,6 +77,7 @@ def trial(engine, prompts, new_tokens, temperature, top_p, top_k):
)
return dict(
seconds=elapsed,
+ steps=steps,
committed_tokens=committed,
tokens_per_second=committed / elapsed,
peak_allocated_bytes=torch.cuda.max_memory_allocated(),
@@ -78,14 +95,10 @@ def main():
parser.add_argument("--cuda-graphs", action="store_true")
parser.add_argument("--k", type=int, default=4)
parser.add_argument("--draft-loops", type=int, default=2)
- parser.add_argument(
- "--early-rejection-pruning",
- action="store_true",
- help="Enable greedy early rejection pruning for the pipeline engine",
- )
parser.add_argument("--concurrency", type=int, default=1)
- parser.add_argument("--new-tokens", type=int, default=32)
- parser.add_argument("--repeats", type=int, default=5)
+ parser.add_argument("--new-tokens", type=int, default=128)
+ parser.add_argument("--repeats", type=int, default=7)
+ parser.add_argument("--seed", type=int, default=0, help="Seed of request 0; request i adds i")
parser.add_argument("--temperature", type=float, default=0.0)
parser.add_argument("--top-p", type=float, default=1.0)
parser.add_argument("--top-k", type=int, default=-1)
@@ -144,7 +157,6 @@ def main():
draft_loops=args.draft_loops,
target_loops=model.config.total_ut_steps,
loop_level_pipelining=enabled,
- early_rejection_pruning=enabled and args.early_rejection_pruning,
),
)
for name, enabled in [("baseline", False), ("pipeline", True)]
@@ -159,6 +171,7 @@ def main():
args.temperature,
args.top_p,
args.top_k,
+ args.seed,
)
for repeat in range(args.repeats):
order = list(engines) if repeat % 2 == 0 else list(reversed(engines))
@@ -170,23 +183,45 @@ def main():
args.temperature,
args.top_p,
args.top_k,
+ args.seed,
)
raw[name].append(result)
print(
f"{name} repeat={repeat}: {result['tokens_per_second']:.2f} tokens/s",
flush=True,
)
+ graphs = {name: graph_counters(e.speculative_runner) for name, e in engines.items()}
finally:
for engine in engines.values():
engine.close()
- summary = {
- name: dict(
- median_seconds=statistics.median(r["seconds"] for r in trials),
- min_seconds=min(r["seconds"] for r in trials),
- max_seconds=max(r["seconds"] for r in trials),
+
+ def acceptance(stats):
+ return stats["accepted_tokens"] / max(stats["drafted_tokens"], 1)
+
+ summary = {}
+ for name, trials in raw.items():
+ seconds = sorted(r["seconds"] for r in trials)
+ quartiles = statistics.quantiles(seconds, n=4) if len(seconds) > 1 else seconds * 3
+ summary[name] = dict(
+ median_seconds=statistics.median(seconds),
+ q1_seconds=quartiles[0],
+ q3_seconds=quartiles[2],
+ min_seconds=seconds[0],
+ max_seconds=seconds[-1],
+ median_tokens_per_second=statistics.median(r["tokens_per_second"] for r in trials),
+ median_step_ms=statistics.median(1e3 * r["seconds"] / r["steps"] for r in trials),
+ acceptance=[acceptance(r["stats"]) for r in trials],
)
- for name, trials in raw.items()
- }
+ # Repeats pair adjacent baseline/pipeline trials; report the paired ratio and a
+ # bootstrap interval of its median rather than a ratio of independent medians.
+ ratios = [b["seconds"] / p["seconds"] for b, p in zip(raw["baseline"], raw["pipeline"])]
+ rng = random.Random(0)
+ boot = sorted(statistics.median(rng.choices(ratios, k=len(ratios))) for _ in range(2000))
+ paired = dict(
+ ratios=ratios,
+ median=statistics.median(ratios),
+ ci95=[boot[49], boot[1949]],
+ )
agreement = []
for baseline, pipeline in zip(raw["baseline"], raw["pipeline"]):
agreement.append(
@@ -212,16 +247,21 @@ def main():
scheduler=asdict(scheduler),
execution=asdict(execution),
summary=summary,
- speedup=summary["baseline"]["median_seconds"] / summary["pipeline"]["median_seconds"],
+ speedup=paired["median"],
+ paired_speedup=paired,
+ cuda_graphs=graphs,
first_divergent_output_positions=agreement,
trials=raw,
- limitations="Small greedy feasibility check; not BF16 acceptance validation. "
+ limitations="Paired decode-throughput check; not BF16 acceptance validation. "
"Peak memory includes both resident engines. Checkpoint config hash does not hash weights.",
)
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(report, indent=2) + "\n")
print(
- json.dumps(dict(summary=summary, speedup=report["speedup"], divergence=agreement), indent=2)
+ json.dumps(
+ dict(summary=summary, paired_speedup=paired, graphs=graphs, divergence=agreement),
+ indent=2,
+ )
)
From c29bad4e6018de09ee0bd8c651e56871d7605761 Mon Sep 17 00:00:00 2001
From: 0z5a
Date: Sat, 24 Oct 2026 18:21:31 +0200
Subject: [PATCH 17/20] Expose the CUDA graph cache limit in the benchmark
Ragged rounds can exhaust the default 16 graphs, after which both rounds
fall back to eager execution and the comparison no longer measures the
schedules. Report runs with zero fallbacks for both arms.
Co-Authored-By: Claude Opus 5.5
---
benchmarks/loop_speculative.py | 11 +++++++++--
1 file changed, 9 insertions(+), 2 deletions(-)
diff --git a/benchmarks/loop_speculative.py b/benchmarks/loop_speculative.py
index 93d8b794..3b422135 100644
--- a/benchmarks/loop_speculative.py
+++ b/benchmarks/loop_speculative.py
@@ -103,6 +103,12 @@ def main():
parser.add_argument("--top-p", type=float, default=1.0)
parser.add_argument("--top-k", type=int, default=-1)
parser.add_argument("--cuda-graph-max-batch-size", type=int, default=128)
+ parser.add_argument(
+ "--cuda-graph-max-graphs",
+ type=int,
+ default=16,
+ help="Raise it when ragged rounds exhaust the cache; check the reported fallbacks",
+ )
args = parser.parse_args()
if not args.model.is_dir():
parser.error("--model must be a complete local checkpoint directory")
@@ -112,8 +118,8 @@ def main():
parser.error("temperature must be nonnegative and top-p must be in (0, 1]")
if args.top_k == 0 or args.top_k < -1:
parser.error("top-k must be -1 or positive")
- if args.cuda_graph_max_batch_size < 1:
- parser.error("cuda-graph-max-batch-size must be positive")
+ if min(args.cuda_graph_max_batch_size, args.cuda_graph_max_graphs) < 1:
+ parser.error("cuda-graph-max-batch-size and cuda-graph-max-graphs must be positive")
torch.set_num_threads(1)
texts = (
json.loads(args.prompts.read_text())
@@ -144,6 +150,7 @@ def main():
execution = ExecutionConfig(
cuda_graphs=args.cuda_graphs,
cuda_graph_max_batch_size=args.cuda_graph_max_batch_size,
+ cuda_graph_max_graphs=args.cuda_graph_max_graphs,
)
engines = {
name: LLMEngine(
From 2a38a997a05120af94c4d3e3922af29d8f45eb65 Mon Sep 17 00:00:00 2001
From: 0z5a
Date: Sat, 24 Oct 2026 18:51:30 +0200
Subject: [PATCH 18/20] Restore the filtered top-p/top-k rejection-sampling
test
The merge took main's speculative tests; keep this branch's check that
rejection sampling recovers the filtered target distribution.
Co-Authored-By: Claude Opus 5.5
---
tests/test_speculative.py | 26 ++++++++++++++++++++++++++
1 file changed, 26 insertions(+)
diff --git a/tests/test_speculative.py b/tests/test_speculative.py
index 683aa65e..1c47af62 100644
--- a/tests/test_speculative.py
+++ b/tests/test_speculative.py
@@ -448,6 +448,32 @@ def test_probability_filters_match_expected_and_apply_strict_top_k():
torch.testing.assert_close(p[p > 0], torch.ones(1))
+@pytest.mark.parametrize(
+ "params",
+ [
+ SamplingParams(temperature=0.8, top_p=0.8),
+ SamplingParams(temperature=0.8, top_k=3, top_p=0.9),
+ ],
+)
+def test_rejection_sampling_respects_filtered_top_p_and_top_k(params, record_property):
+ """The acceptance rule targets the filtered p distribution, not raw logits."""
+ target_logits = torch.tensor([3.0, 2.1, 1.5, 0.4, -0.2, -1.0])
+ proposal_logits = torch.tensor([1.0, 2.8, 2.0, 1.2, 0.5, -0.5])
+ p = Sampler.probabilities(target_logits, params)
+ q = Sampler.probabilities(proposal_logits, params)
+ generator = torch.Generator().manual_seed(9173)
+ counts = torch.zeros_like(p)
+ for _ in range(16000):
+ candidate = int(torch.multinomial(q, 1, generator=generator))
+ token, _ = rejection_sample(candidate, p, q, generator)
+ counts[token] += 1
+ empirical = counts / counts.sum()
+ error = empirical - p
+ record_property("distribution_max_abs_error", float(error.abs().max()))
+ record_property("distribution_rms_error", float(error.square().mean().sqrt()))
+ torch.testing.assert_close(empirical, p, atol=0.02, rtol=0)
+
+
@pytest.mark.parametrize(
"kwargs",
[
From 80a737d047d3ad451615bc6287a60585a61b8559 Mon Sep 17 00:00:00 2001
From: 0z5a
Date: Sat, 24 Oct 2026 18:51:30 +0200
Subject: [PATCH 19/20] Document the loop round, its cost model and RTX 5090
measurements
Describe the planned submit/collect round and why same-seed outputs only
match exactly in FP32. Replace the pruning sections and A100 tables with a
cost model: both rounds launch the same recurrent calls over the same rows,
so any difference comes from GEMM kernel selection at C, 2C and 5C rows,
plus two graph-cache effects that can masquerade as gains. Add paired
per-round measurements, the linear-layer prediction and the graph-limit
control on an RTX 5090.
Co-Authored-By: Claude Opus 5.5
---
docs/assets/loop-schedule-comparison.svg | 6 +-
docs/loop_speculative.md | 425 +++++++++--------------
2 files changed, 157 insertions(+), 274 deletions(-)
diff --git a/docs/assets/loop-schedule-comparison.svg b/docs/assets/loop-schedule-comparison.svg
index 5ed66215..b496729a 100644
--- a/docs/assets/loop-schedule-comparison.svg
+++ b/docs/assets/loop-schedule-comparison.svg
@@ -13,7 +13,6 @@
.draft { fill: #dbeafe; stroke: #2563eb; stroke-width: 1.5; }
.deep { fill: #fef3c7; stroke: #d97706; stroke-width: 1.5; }
.verify { fill: #dcfce7; stroke: #16a34a; stroke-width: 1.5; }
- .prune { fill: #fee2e2; stroke: #dc2626; stroke-width: 1.5; }
.arrow { fill: none; stroke: #64748b; stroke-width: 2; marker-end: url(#arrow); }
.thin-arrow { fill: none; stroke: #94a3b8; stroke-width: 1.5; marker-end: url(#arrow); }
@@ -87,9 +86,9 @@
- reorder + verify
+ one head + verify
- target rows arrive
+ after the last waveshallow / draft depth
@@ -97,5 +96,4 @@
deep / target depthcommon verifier and commit
- Early prune: only unscheduled suffix branches can be removed.
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 2abead84..1fb37c44 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -1,10 +1,10 @@
# Loop-level speculative execution
-This opt-in synchronous path implements the mixed-depth batching portion proposed
-in [Issue #43](https://github.com/ThinkFlowLab/vllm-rlt/issues/43): advance shallow work
-for later token positions in the same recurrent call as deep work for earlier
-positions. It retains the existing finite speculative window and atomic round
-commit. The default remains the original block verifier.
+This opt-in round implements the mixed-depth batching part of
+[Issue #43](https://github.com/ThinkFlowLab/vllm-rlt/issues/43): shallow work for
+later token positions runs in the same recurrent call as deeper work for earlier
+positions. It keeps the finite speculative window, the shared verifier and the
+atomic round commit. The default remains the block round.
## Enable
@@ -18,55 +18,16 @@ SpeculativeConfig(
```
The offline and serving CLI accept `--speculative-tokens 4 --loop-level-pipelining`.
-Greedy-only early suffix pruning can be enabled separately with
-`--early-rejection-pruning`; it is off by default. Top-p and top-k sampling are
-supported by the same loop pipeline, but sampling still verifies the complete
-finite window before applying rejection sampling. Early pruning is intentionally
-greedy-only because greedy acceptance is a deterministic candidate/argmax
-comparison, while sampling acceptance requires the target and proposal
-distributions, random draws and ordered correction sampling.
-Existing restrictions apply: `last_exited`, full fixed target depth, synchronous
-refill scheduling, and no preemption or PD.
-
-The implementation has three execution choices:
-
-| Mode | Recurrent schedule | Verification | Intended use |
-| --- | --- | --- | --- |
-| Baseline | Complete all draft positions, then complete target depth | One batched verifier window | Reference path and default |
-| Loop | Advance active positions one depth per mixed-depth wave | Reorder completed target rows into the common verifier | Measure mixed-depth batching |
-| Loop + greedy prune | Loop schedule with target-ready rejection checks | Stop future suffix branches after the first mismatch | Optional greedy experiment |
-
-The loop mode is a synchronous pipeline. It does not add CUDA streams, reduce the
-number of recurrent rows, or retain hidden state across speculative rounds. Its
-possible benefit comes from the shape and timing of mixed-depth batches; metadata,
-Python scheduling, Coda fragmentation and dynamic CUDA Graph shapes are additional
-costs.
-
-## Execution and state
-
-### Loop-level flow
-
-The loop path keeps one branch for each active token position. A branch carries
-its request index, token offset, recurrent hidden state and current loop depth.
-Each wave advances every active branch by one depth. When a branch reaches the
-draft depth, its coda output starts a child branch while the parent continues
-toward the target depth. When a branch reaches the target depth, its target
-logits are saved for the existing verifier.
-
-The baseline executes all draft positions before block target verification. In
-the loop schedule, after A reaches the draft depth, B starts at depth 0 while A
-continues at a deeper level; C starts two waves later. Target rows are reordered
-into the common verifier only after their corresponding branches reach target
-depth. The waves are still launched sequentially on the current host path.
-
-This PR currently implements only the synchronous finite-window path. Mixed-
-depth waves are launched sequentially from the host; there are no CUDA streams,
-asynchronous scheduler overlap, cross-round hidden-state retention, preemption,
-or PD execution. The existing `SpeculativeResult` and atomic KV commit/rollback
-boundary are intentionally preserved.
-
-For K=2, d=2 and D=4, A is the last committed input, B and C are proposals.
-Numbers below are zero-based depths; columns are successive recurrent calls.
+The restrictions of fixed-depth self-speculation apply unchanged. The loop round
+is a drop-in replacement for the block round inside `SpeculativeRunner.submit`, so
+it composes with deferred collection, several in-flight submissions, preemption
+between rounds and CUDA graphs.
+
+## Schedule
+
+For K=2, d=2 (`draft_loops`) and D=4 (`target_loops`), A is the last committed
+input and B and C are proposals. Numbers are zero-based depths; columns are
+successive recurrent calls ("waves").
| Input position | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
@@ -74,182 +35,120 @@ Numbers below are zero-based depths; columns are successive recurrent calls.
| B | | | 0 | 1 | 2 | 3 | | |
| C | | | | | 0 | 1 | 2 | 3 |
-After A completes depth 1, its shallow head proposes B. Call 3 batches A at
-depth 2 and B at depth 0. After B completes depth 1, it proposes C while A's
-full-depth head produces the distribution that verifies B. These head rows are
-batched too. C's full-depth distribution supplies the bonus token.
-
-Each branch stores its item index, position offset, hidden state and completed
-loop count. A child starts d calls behind its parent. Therefore, at every layer,
-the same-depth KV for earlier positions is already available before the child's
-attention. `_prepare_batch` supplies a separate depth plane and causal context
-length for every row. The pipeline uses single-query paged attention, not packed
-multi-query prefill. No shallow KV is copied into uncomputed deep planes.
-
-Each position/depth is computed exactly once. Proposals remain on the device and
-greedy target IDs are read back once per mixed-depth wave. For sampling, the
-filtered proposal distribution `q` is retained for each draft and the common
-verifier applies rejection sampling with the request's generator from the shared
-`RngRegistry`. Different GEMM shapes and the changed draft/verify interleaving can
-change BF16 results or the fixed-seed token sequence; same-seed agreement is not a
-general sampling correctness guarantee.
-
-Completed target logits are reordered into item/position order for the common
-verifier. `SpeculativeResult` and engine commit stay unchanged: keep KV through
-the initial input and accepted candidates; invalidate every depth's rejected
-suffix; leave correction/bonus unforwarded. Engine EOS/length handling and
-exception cleanup remain responsible for the request lifecycle.
-
-## Scope and benefit
-
-The default finite-window path does **not** implement cross-round hidden-state
-retention or the complete continuous LoopSpec algorithm. The optional greedy
-pruning path checks a target row when its coda output is ready and skips later
-suffix branches after rejection. Work that already entered the mixed-depth wave
-cannot be cancelled. Sampling still runs the complete window before
-verification. Greedy target and draft IDs are read back once per mixed-depth
-wave; the host-side rejection decision remains the pruning synchronization
-boundary. No new CUDA streams or asynchronous scheduler are introduced.
-
-### Why pruning is greedy-only
-
-For greedy decoding, a target row produces one deterministic token ID. Once that
-ID is available, the runner can compare it with the candidate and stop creating
-later branches for that request after the first mismatch. The already submitted
-rows still finish, but suffix rows that have not entered a wave are avoided.
-
-Top-p and top-k use rejection sampling. Acceptance depends on the filtered target
-distribution `p`, the filtered draft distribution `q`, and a random uniform draw;
-on rejection, the correction token is drawn from `normalize(max(p - q, 0))`.
-A target argmax mismatch is therefore not equivalent to rejection: a candidate
-can differ from the target argmax and still be accepted, while a candidate that
-matches the argmax can still be rejected by the random draw. Applying the full
-decision as soon as each target row is ready would also require retaining every
-proposal distribution, preserving the request's RNG order, sampling a residual
-correction token, and invalidating all descendants generated from the rejected
-candidate. The current sampling path instead computes the finite window and
-applies the existing verifier in order, preserving its distribution and RNG
-behavior. Device-side batched acceptance can support sampling pruning in a later
-change.
-
-For a uniform window, both the original block verifier and this pipeline require
-`d*K + D` recurrent calls and `(K+1)*D` recurrent rows. Pipelining changes batch
-shapes and head batching, not total recurrent work or that critical path. Its
-cost also includes extra per-step metadata preparation and loss of the original
-verifier's packed FA4 attention. Speedup must be measured; mixing depths alone
-does not guarantee an improvement. Reducing round bubbles or wasted suffix work
-would require a follow-up design for acceptance and state lifetime.
-
-## Correctness checks for Issue #43
-
-The fixed-input oracle test runs the same token chain through native serial
-full-depth execution and the reuse path. It compares every loop-2 hidden row,
-every depth/layer KV row and the final target logits, and checks max absolute
-error, RMS error and target argmax mismatch rate in the assertion contract. The
-FP32 case is used to diagnose indexing or dependency mistakes; the BF16 case uses
-relaxed numerical tolerances to assess real model behavior.
-
-Free greedy generation is checked at `K=1/2/4/8`, with the existing ragged
-single-request cases and an eight-request batch. The test requires exact token
-and exit-depth agreement with native generation and reports the first divergent
-position in the failure message. Sampling is checked independently with small
-analytical distributions after applying the actual top-p and top-k filters; the
-empirical samples must match the filtered target distribution within the
-sampling tolerance. Fixed-seed text equality is not required for sampling.
-
-In the focused 16,000-draw check, the maximum absolute deviation from the
-filtered target distribution was 0.003085 for top-p and 0.002593 for top-p plus
-top-k; the corresponding RMS deviations were 0.001781 and 0.001368. Both are
-below the test tolerance of 0.02.
-
-```bash
-python -m pytest -q tests/test_speculative.py \
- -k 'reused_hidden_logits or issue43_greedy or filtered_top_p'
-```
-
-## Parameters and results
-
-The reference measurements use an Ouro-1.4B checkpoint on an NVIDIA A100-SXM4-40GB,
-BF16, Triton attention, CUDA Graphs, `K=4`, `draft_loops=2`, `target_loops=4`,
-32 generated tokens and three timed repetitions per point. The benchmark compares
-the original block verifier with the loop pipeline using the same prompts and
-prefill barrier. Throughput is aggregate committed decode throughput across all
-concurrent requests, measured in output tokens per second (`output tok/s`).
-It counts generated output tokens only, excludes prompt/prefill tokens and is
-not a per-request rate.
-
-For C=1 through C=64, the CUDA Graph limit is 128. The C=128 rows use limit 256
-because the loop's largest mixed-depth wave is approximately `2C` rows. The C=64
-wave is at most 128 rows and therefore remains within the default limit.
-
-### Greedy
-
-| Concurrency | Baseline (output tok/s) | Loop (output tok/s) | Change | Loop + greedy prune (output tok/s) | Change |
-| ---: | ---: | ---: | ---: | ---: | ---: |
-| 1 | 58.6 | 56.9 | -2.8% | 66.7 | +13.9% |
-| 4 | 217.1 | 214.7 | -1.1% | 215.5 | -0.7% |
-| 8 | 420.3 | 409.8 | -2.5% | 412.0 | -2.0% |
-| 16 | 790.7 | 780.3 | -1.3% | 782.8 | -1.1% |
-| 32 | 1319.5 | 1473.7 | +11.7% | 1465.0 | +10.4% |
-| 64 | 2428.5 | 2648.4 | +9.1% | 2628.8 | +7.9% |
-| 128 | 4085.1 | 4229.6 | +3.5% | 4203.6 | +3.1% |
-
-The additional greedy counters show where pruning saves work:
-
-| Concurrency | Loop verified rows / committed tokens | Loop + prune verified rows / committed tokens | Loop peak allocated (GB, decimal) |
-| ---: | ---: | ---: | ---: |
-| 1 | 40 / 32 | 32 / 32 | 2.99 GiB |
-| 4 | 144 / 128 | 128 / 128 | 3.23 GiB |
-| 8 | 288 / 256 | 256 / 256 | 3.54 GiB |
-| 16 | 576 / 512 | 512 / 512 | 4.18 GiB |
-| 32 | 1152 / 1024 | 1024 / 1024 | 5.42 GiB |
-| 64 | 2304 / 2048 | 2048 / 2048 | 7.90 GiB |
-| 128 | 4608 / 4096 | 4096 / 4096 | 12.88 GiB |
-
-### Sampling
-
-Sampling uses `temperature=0.8`, `top_p=0.8`, with `top_k=-1` for top-p and
-`top_k=7` for top-p + top-k. Early pruning is disabled. The C=128 rows use
-CUDA Graph max 256; all other rows use max 128.
-
-| Concurrency | Top-p baseline (output tok/s) | Top-p loop (output tok/s) | Change | Top-p + top-k baseline (output tok/s) | Top-p + top-k loop (output tok/s) | Change |
-| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
-| 1 | 61.9 | 58.5 | -5.5% | 51.6 | 46.8 | -9.4% |
-| 4 | 205.1 | 159.9 | -22.0% | 139.7 | 137.8 | -1.3% |
-| 8 | 295.4 | 330.2 | +11.8% | 236.3 | 254.2 | +7.5% |
-| 16 | 523.8 | 492.9 | -5.9% | 417.2 | 393.2 | -5.8% |
-| 32 | 662.2 | 723.9 | +9.3% | 519.3 | 532.2 | +2.5% |
-| 64 | 876.7 | 914.1 | +4.3% | 625.2 | 646.8 | +3.5% |
-| 128 | 1046.6 | 940.0 | -10.2% | 727.1 | 801.2 | +10.2% |
-
-The benchmark JSON also records median/min/max decode time, peak allocated
-memory, recurrent rounds, drafted/accepted tokens, verified rows and the first
-divergent output position. Greedy runs have identical outputs at every point;
-sampling divergence is expected when the same request generator is consumed in
-a different draft/verify order. It is not used as a distribution-correctness
-criterion.
-
-The results are workload-dependent. Loop-level scheduling does not reduce the
-number of recurrent rows; its potential benefit comes from mixed-depth batch
-shapes and device utilization. Metadata preparation, Python scheduling, Coda
-fragmentation and graph capture shape can outweigh that benefit. For example,
-the greedy loop reduces verified rows from 4608 to 4096 at C=128 only when
-early pruning is enabled, while the non-pruned loop keeps the same recurrent
-row count as baseline. The sampling rows also show that acceptance rate and
-round count change with the sampled path, so throughput should be read together
-with the recorded counters.
-
-In summary, the synchronous loop path is a clear win only at selected higher
-concurrency points: greedy throughput improves by 3-12% from C=32 in this
-benchmark, while low-concurrency runs regress by about 1-3%. Greedy pruning is
-strongest for a single request with an early mismatch (+13.9% at C=1), but adds
-little at larger batches because work already submitted in the current wave
-cannot be cancelled. Sampling remains workload-sensitive, ranging from a 22%
-regression to a 12% improvement for top-p and from a 9% regression to a 10%
-improvement for top-p plus top-k. These results show that mixed-depth scheduling
-alone does not guarantee a speedup; asynchronous overlap, state retention,
-device-side acceptance and more efficient dynamic compaction remain follow-up
-work.
+
+
+Offset `o` enters at wave `o*d`, after its parent's draft head chose its input, and
+advances one depth per wave. A child therefore starts `d` waves after its parent,
+so at every layer the same-depth KV of earlier positions is written before the
+child attends to it. Rows use single-query paged attention with their own depth
+plane and context length; adjacent positions are never packed into one prefill
+sequence. Each position and depth is computed exactly once.
+
+`loop_waves(token_counts, d, D)` returns the schedule. It depends only on the
+token counts of the round, never on token values, so the whole round is planned
+before any kernel runs.
+
+## Execution
+
+- `_prepare_loop_inputs` builds every wave's KV metadata and gather/scatter
+ indices on the host and uploads them in one snapshot per dtype, like the block
+ round's `_prepare_inputs`.
+- Branch states live in one buffer whose rows follow verification order
+ (contiguous by request, then offset). Each wave gathers its rows, runs one
+ mixed-depth recurrent call and scatters the states back.
+- After a wave, only rows that reached draft depth run the head. Their IDs (and,
+ for sampling, the filtered proposal `q` and the draw) stay on the device and
+ feed the children's prelude.
+- After the last wave, one head call produces all `K+1` verification rows per
+ request. Collection is the block round's `_collect_round`, so greedy acceptance,
+ rejection sampling and the single device-to-host read are shared.
+
+Both rounds consume each request's generator in the same order: drafts by
+offset, then verification. On CPU in FP32 the two rounds produce identical tokens
+for the same seeds. On GPU in BF16 the two rounds run GEMMs with different row
+counts, which can select different kernels and round differently, so near-tied
+argmax or sampling decisions can diverge (see the measurements). This is the
+same class of difference as between batch sizes in ordinary decoding.
+
+Greedy early-rejection pruning is not supported: deciding per wave whether to
+start a child needs a host readback every wave, which the asynchronous
+submission path does not allow.
+
+## Cost model
+
+For a uniform round (every request has `K+1` rows), both rounds launch the same
+number of recurrent calls, `d*K + D`, over the same `(K+1)*D*C` rows for `C`
+requests. Only the row counts per call differ. With K=4, d=2, D=4:
+
+| Round | Recurrent calls by row count |
+| --- | --- |
+| Block | 10 calls of `C` rows (drafts), 2 calls of `5C` rows (verification) |
+| Loop | 4 calls of `C` rows, 8 calls of `2C` rows |
+
+With `c(M)` the cost of one recurrent call over `M` rows, the loop round is faster
+only if `c(5C) - c(C) > 4 * (c(2C) - c(C))`. If `c` is affine in `M` — flat when
+weight-bandwidth bound, linear when compute bound — both sides are equal, so the
+schedule has no intrinsic advantage. Gains and losses come from where GEMM kernel
+selection changes `c(M)` between `C`, `2C` and `5C`, which depends on the GPU,
+the BLAS library and the model shape. Treat the mode as hardware-specific and
+measure it on the target deployment.
+
+Two configuration effects can masquerade as schedule gains:
+
+- `cuda_graph_max_batch_size`: the block round's verification calls have `5C`
+ rows and fall back to eager execution once `5C` exceeds the limit, while the
+ loop round's `2C`-row waves still replay graphs.
+- `cuda_graph_max_graphs`: ragged rounds near the end of generation produce more
+ distinct wave sizes in the loop round and can exhaust the graph cache.
+
+The benchmark reports graph fallbacks for both rounds; compare runs where both
+are zero.
+
+## Measurements (RTX 5090)
+
+Ouro-1.4B (`574fa66c`), BF16, Triton attention, CUDA graphs with
+`--cuda-graph-max-batch-size 1024 --cuda-graph-max-graphs 128` (no fallbacks in
+either round), K=4, d=2, D=4, 16 prompts cycled over C requests, 128 new tokens,
+7 alternating paired repeats; torch 2.13.0+cu130, Triton 3.7.1, driver 580.
+
+Per-round time compares the schedules directly. Committed throughput also
+depends on acceptance, which differs once BF16 outputs diverge.
+
+| C | Block ms/round | Loop ms/round | Round-time ratio | `c(M)` prediction | Throughput ratio | Acceptance block / loop | Diverged request trials |
+| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| 1 | 44.80 | 52.97 | 0.846 | 0.846 | 0.846 | 0.884 / 0.884 | 0 / 7 |
+| 4 | 56.24 | 58.98 | 0.953 | 0.955 | 0.926 | 0.839 / 0.800 | 21 / 28 |
+| 8 | 59.26 | 61.28 | 0.967 | 0.980 | 1.024 | 0.804 / 0.828 | 28 / 56 |
+| 16 | 63.39 | 59.97 | 1.057 | 1.227 | 1.028 | 0.824 / 0.822 | 63 / 112 |
+| 32 | 63.13 | 63.14 | 1.000 | 1.067 | 1.088 | 0.800 / 0.836 | 126 / 224 |
+| 64 | 75.45 | 80.56 | 0.937 | 0.922 | 0.910 | 0.823 / 0.829 | 280 / 448 |
+
+Ratios are block time over loop time (above 1 means the loop round is faster).
+Bootstrap 95% intervals of the paired throughput ratio are narrower than ±0.01
+at every point. Every repeat replays the same prompts and seeds, so these
+intervals reflect timing jitter only, not prompt or seed variance. C=128 did not fit with two resident engines on the 32 GB card.
+
+The prediction uses a linear-layer-only proxy: graph-replayed time of all
+linear layers of one recurrent pass for M rows (`c(1)=1.90`, `c(2..16)≈2.66`,
+`c(20..32)≈2.0`, `c(64)=2.14`, `c(80..192)≈3.0-3.3`, `c(256)=4.46`,
+`c(320)=4.77`, `c(640)=7.94` ms). It excludes attention, norms and the head,
+so it matches best at small C. Nsight Compute shows the steps come from kernel
+selection: M=1 uses a GEMV kernel, M=2-16 a 16x16 WMMA GEMM that is about 48%
+slower, M=20-32 a split-K 32x32 WMMA GEMM that is faster than the GEMV, and M=80
+switches to a larger tile again.
+
+Top-p (temperature 0.8, top-p 0.8, fixed per-request seeds):
+
+| C | Block ms/round | Loop ms/round | Round-time ratio | Throughput ratio |
+| ---: | ---: | ---: | ---: | ---: |
+| 1 | 47.72 | 56.43 | 0.846 | 0.846 |
+| 16 | 73.52 | 71.97 | 1.021 | 0.985 |
+| 64 | 123.66 | 128.16 | 0.965 | 0.965 |
+
+With the graph batch limit at 128 and C=32, the block round's 160-row
+verification calls fall back to eager (448 recurrent fallbacks) and its round
+time rises from 63.13 to 66.0 ms, while the loop round stays at 63.01 ms: a 4.7%
+"gain" that is entirely the graph limit.
## Validation
@@ -259,51 +158,37 @@ python -m ruff check .
python -m ruff format --check vllm_rlt tests benchmarks
# On an allocated CUDA device with Triton:
-python -m pytest -q tests/test_speculative.py --run-gpu -m gpu -k 'not flash_attn_4'
+python -m pytest -q --run-gpu -m gpu -k 'not flash_attn_4' \
+ tests/test_speculative.py tests/test_loop_speculative.py
```
-Tests exercise mixed depths including non-divisible d/D, actual hidden/logit/KV
-agreement against a serial reference, every rejection boundary and bonus,
-greedy/sampling batches, K=1/2/4/8, shortened windows, EOS, prefix reuse,
-incremental allocation, cancellation and partial execution failure. CUDA coverage
-checks eager/graph parity and one greedy readback per round.
+`tests/test_speculative.py` runs the existing speculative contracts in both modes:
+native greedy agreement with ragged rounds and budgets, the serial hidden/logit/KV
+oracle, every rejection boundary and bonus, EOS, prefix reuse and incremental
+allocation, seeded sampling, preemption between rounds, split prefill, deferred
+collection, several in-flight submissions, and CUDA graph/eager parity.
+`tests/test_loop_speculative.py` checks the schedule itself, per-round equivalence
+with the block round (results, verification logits, KV and stats) for d=1/2/3 and
+sampling, the head calls inside the loop, CLI/config handling, failure cleanup and
+rollback.
-The focused Issue #43 checks can be run without the full suite:
-
-```bash
-python -m pytest -q tests/test_speculative.py \
- -k 'reused_hidden_logits or issue43_greedy or filtered_top_p'
-```
-
-For a bounded real-checkpoint feasibility run:
+Paired benchmark:
```bash
python -m benchmarks.loop_speculative \
- --model /path/to/local/Ouro-1.4B \
- --cuda-graphs --k 4 --concurrency 1 --new-tokens 32 --repeats 3 \
- --output /tmp/loop-speculative.json
-
-# Sampling variant (top-p + top-k):
-python -m benchmarks.loop_speculative \
- --model /path/to/local/Ouro-1.4B \
- --cuda-graphs --k 4 --concurrency 8 --new-tokens 32 --repeats 3 \
- --temperature 0.8 --top-p 0.8 --top-k 7 \
- --output /tmp/loop-speculative-top-p-top-k.json
+ --model /path/to/local/Ouro-1.4B --prompts prompts.json \
+ --cuda-graphs --cuda-graph-max-batch-size 1024 --cuda-graph-max-graphs 128 \
+ --k 4 --concurrency 16 --new-tokens 128 --repeats 7 \
+ --output /tmp/loop-speculative-c16.json
```
-The harness warms both engines, completes a common prefill barrier before timing,
-ignores EOS to commit equal work, synchronizes timing boundaries, alternates trial
-order, and saves raw timings, outputs, first divergences and runtime metadata.
-For large loop batches, `--cuda-graph-max-batch-size 256` allows Graph capture of
-the 2C mixed-depth wave when C=128; the default is 128.
-`--prompts` accepts a JSON list of strings; `--backend flash_attn_4` requires
-compatible hardware. Peak memory includes both resident engines; the recorded
-checkpoint config hash is not a weight hash. These checks cover fixed-input
-reuse and sampling distribution contracts, but do not replace Issue #43's
-long-context, serving-latency or full BF16 acceptance matrix.
-
-Raw benchmark JSON is stored under `artifacts/loop-speculative/`; the extended
-run notes include the runtime metadata and per-repetition timings. This work is
-an opt-in implementation of the loop-level direction in Issue #43; it does not
-claim that mixed-depth scheduling reduces total recurrent work or improves every
-hardware and workload combination.
+Both engines share one resident model. Prefill is excluded, trial order
+alternates and every request has a fixed seed, so sampled arms draw the same
+random numbers. The report contains the median of per-repeat paired ratios with a
+bootstrap interval, quartiles, per-round time, acceptance, CUDA graph counters,
+first divergent output positions and raw trials. Peak memory includes both
+engines; the checkpoint config hash is not a weight hash.
+
+This mode does not reduce recurrent work or the round's critical path, and does
+not implement cross-round hidden-state retention or the continuous LoopSpec
+algorithm from Issue #43.
From 3f10ece367e7e6640746ad54566ec12efa937901 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 25 Oct 2026 03:56:12 +0000
Subject: [PATCH 20/20] docs: keep loop diagram local
---
docs/assets/loop-schedule-comparison.svg | 99 ------------------------
docs/loop_speculative.md | 2 -
2 files changed, 101 deletions(-)
delete mode 100644 docs/assets/loop-schedule-comparison.svg
diff --git a/docs/assets/loop-schedule-comparison.svg b/docs/assets/loop-schedule-comparison.svg
deleted file mode 100644
index b496729a..00000000
--- a/docs/assets/loop-schedule-comparison.svg
+++ /dev/null
@@ -1,99 +0,0 @@
-
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 1fb37c44..8ab49dd1 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -35,8 +35,6 @@ successive recurrent calls ("waves").
| B | | | 0 | 1 | 2 | 3 | | |
| C | | | | | 0 | 1 | 2 | 3 |
-
-
Offset `o` enters at wave `o*d`, after its parent's draft head chose its input, and
advances one depth per wave. A child therefore starts `d` waves after its parent,
so at every layer the same-depth KV of earlier positions is written before the
+
+The upper lane shows the baseline's block draft followed by block target
+verification. The lower lane shows the loop pipeline: after A reaches the draft
+depth, B starts at depth 0 while A continues at a deeper level; C starts two
+waves later. The target rows are reordered into the common verifier only after
+their corresponding branches reach the target depth. This is the loop-level
+pipeline added in this PR; the waves are still launched sequentially on the
+current host path.
This PR currently implements only the synchronous finite-window path. Mixed-
depth waves are launched sequentially from the host; there are no CUDA streams,
From e9e5ef1863378ef8b9fdedebdc2f288392e638b4 Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 11:47:08 +0000
Subject: [PATCH 10/20] test: document issue 43 loop validation
---
docs/loop_speculative.md | 108 ++++++++++++++++++++++++++------------
tests/test_speculative.py | 92 +++++++++++++++++++++++++++++++-
2 files changed, 165 insertions(+), 35 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index da1740cb..4f2cda74 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -144,11 +144,33 @@ verifier's packed FA4 attention. Speedup must be measured; mixing depths alone
does not guarantee an improvement. Reducing round bubbles or wasted suffix work
would require a follow-up design for acceptance and state lifetime.
+## Correctness checks for Issue #43
+
+The fixed-input oracle test runs the same token chain through native serial
+full-depth execution and the reuse path. It compares every loop-2 hidden row,
+every depth/layer KV row and the final target logits, and checks max absolute
+error, RMS error and target argmax mismatch rate in the assertion contract. The
+FP32 case is used to diagnose indexing or dependency mistakes; the BF16 case uses
+relaxed numerical tolerances to assess real model behavior.
+
+Free greedy generation is checked at `K=1/2/4/8`, with the existing ragged
+single-request cases and an eight-request batch. The test requires exact token
+and exit-depth agreement with native generation and reports the first divergent
+position in the failure message. Sampling is checked independently with small
+analytical distributions after applying the actual top-p and top-k filters; the
+empirical samples must match the filtered target distribution within the
+sampling tolerance. Fixed-seed text equality is not required for sampling.
+
+```bash
+python -m pytest -q tests/test_speculative.py \
+ -k 'reused_hidden_logits or issue43_greedy or filtered_top_p'
+```
+
## Parameters and results
The reference measurements use an Ouro-1.4B checkpoint on an NVIDIA A100-SXM4-40GB,
BF16, Triton attention, CUDA Graphs, `K=4`, `draft_loops=2`, `target_loops=4`,
-32 generated tokens and five timed repetitions per point. The benchmark compares
+32 generated tokens and three timed repetitions per point. The benchmark compares
the original block verifier with the loop pipeline using the same prompts and
prefill barrier. Throughput is committed decode tokens per second.
@@ -160,13 +182,25 @@ wave is at most 128 rows and therefore remains within the default limit.
| Concurrency | Baseline | Loop | Change | Loop + greedy prune | Change |
| ---: | ---: | ---: | ---: | ---: | ---: |
-| 1 | 58.6 | 56.9 | -2.8% | 66.7 | +13.8% |
-| 4 | 217.2 | 214.7 | -1.1% | 215.5 | -0.7% |
-| 8 | 420.4 | 409.9 | -2.5% | 412.4 | -1.9% |
-| 16 | 790.8 | 780.1 | -1.4% | 782.3 | -1.0% |
-| 32 | 1321.2 | 1474.1 | +11.6% | 1468.6 | +11.1% |
-| 64 | 2434.1 | 2649.6 | +8.9% | 2625.5 | +7.7% |
-| 128 | 4092.9 | 4231.2 | +3.4% | 4207.7 | +2.8% |
+| 1 | 58.6 | 56.9 | -2.8% | 66.7 | +13.9% |
+| 4 | 217.1 | 214.7 | -1.1% | 215.5 | -0.7% |
+| 8 | 420.3 | 409.8 | -2.5% | 412.0 | -2.0% |
+| 16 | 790.7 | 780.3 | -1.3% | 782.8 | -1.1% |
+| 32 | 1319.5 | 1473.7 | +11.7% | 1465.0 | +10.4% |
+| 64 | 2428.5 | 2648.4 | +9.1% | 2628.8 | +7.9% |
+| 128 | 4085.1 | 4229.6 | +3.5% | 4203.6 | +3.1% |
+
+The additional greedy counters show where pruning saves work:
+
+| Concurrency | Loop verified rows / committed tokens | Loop + prune verified rows / committed tokens | Loop peak allocated |
+| ---: | ---: | ---: | ---: |
+| 1 | 40 / 32 | 32 / 32 | 2.99 GiB |
+| 4 | 144 / 128 | 128 / 128 | 3.23 GiB |
+| 8 | 288 / 256 | 256 / 256 | 3.54 GiB |
+| 16 | 576 / 512 | 512 / 512 | 4.18 GiB |
+| 32 | 1152 / 1024 | 1024 / 1024 | 5.42 GiB |
+| 64 | 2304 / 2048 | 2048 / 2048 | 7.90 GiB |
+| 128 | 4608 / 4096 | 4096 / 4096 | 12.88 GiB |
### Sampling
@@ -176,21 +210,30 @@ CUDA Graph max 256; all other rows use max 128.
| Concurrency | Top-p baseline | Top-p loop | Change | Top-p + top-k baseline | Top-p + top-k loop | Change |
| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
-| 1 | 61.8 | 58.4 | -5.5% | 51.6 | 46.8 | -9.4% |
-| 4 | 205.2 | 159.9 | -22.0% | 140.3 | 138.0 | -1.6% |
-| 8 | 295.1 | 330.2 | +11.9% | 236.8 | 254.9 | +7.7% |
-| 16 | 524.3 | 493.9 | -5.8% | 417.3 | 393.2 | -5.8% |
-| 32 | 669.1 | 728.7 | +8.9% | 520.3 | 533.6 | +2.6% |
-| 64 | 870.7 | 908.3 | +4.3% | 628.2 | 649.4 | +3.4% |
-| 128 | 1057.5 | 953.4 | -9.8% | 730.6 | 805.7 | +10.3% |
+| 1 | 61.9 | 58.5 | -5.5% | 51.6 | 46.8 | -9.4% |
+| 4 | 205.1 | 159.9 | -22.0% | 139.7 | 137.8 | -1.3% |
+| 8 | 295.4 | 330.2 | +11.8% | 236.3 | 254.2 | +7.5% |
+| 16 | 523.8 | 492.9 | -5.9% | 417.2 | 393.2 | -5.8% |
+| 32 | 662.2 | 723.9 | +9.3% | 519.3 | 532.2 | +2.5% |
+| 64 | 876.7 | 914.1 | +4.3% | 625.2 | 646.8 | +3.5% |
+| 128 | 1046.6 | 940.0 | -10.2% | 727.1 | 801.2 | +10.2% |
+
+The benchmark JSON also records median/min/max decode time, peak allocated
+memory, recurrent rounds, drafted/accepted tokens, verified rows and the first
+divergent output position. Greedy runs have identical outputs at every point;
+sampling divergence is expected when the same request generator is consumed in
+a different draft/verify order. It is not used as a distribution-correctness
+criterion.
The results are workload-dependent. Loop-level scheduling does not reduce the
number of recurrent rows; its potential benefit comes from mixed-depth batch
shapes and device utilization. Metadata preparation, Python scheduling, Coda
-fragmentation and graph capture shape can outweigh that benefit. Sampling may
-also diverge from baseline at a fixed seed because draft sampling and rejection
-sampling are interleaved differently; this does not by itself indicate an
-incorrect sampling distribution.
+fragmentation and graph capture shape can outweigh that benefit. For example,
+the greedy loop reduces verified rows from 4608 to 4096 at C=128 only when
+early pruning is enabled, while the non-pruned loop keeps the same recurrent
+row count as baseline. The sampling rows also show that acceptance rate and
+round count change with the sampled path, so throughput should be read together
+with the recorded counters.
In summary, the synchronous loop path is a clear win only at selected higher
concurrency points: greedy throughput improves by 3-12% from C=32 in this
@@ -204,16 +247,6 @@ alone does not guarantee a speedup; asynchronous overlap, state retention,
device-side acceptance and more efficient dynamic compaction remain follow-up
work.
-At low concurrency, metadata preparation, Python scheduling and fragmented Coda
-work outweigh the mixed-depth benefit. The measured greedy loop becomes faster
-from C=32 onward, reaching its best gain at C=32; at C=128 it remains faster only
-after raising the Graph limit so the `2C` wave can be captured. Greedy pruning is
-helpful for a single request with early rejection, but adds little at larger
-concurrency because much of the suffix work has already entered a wave. Sampling
-is more workload-sensitive: top-p is slower at C=128, while top-p + top-k is
-faster there, showing that the loop benefit is not a general consequence of
-sampling or concurrency alone.
-
## Validation
```bash
@@ -231,18 +264,25 @@ greedy/sampling batches, K=1/2/4/8, shortened windows, EOS, prefix reuse,
incremental allocation, cancellation and partial execution failure. CUDA coverage
checks eager/graph parity and one greedy readback per round.
+The focused Issue #43 checks can be run without the full suite:
+
+```bash
+python -m pytest -q tests/test_speculative.py \
+ -k 'reused_hidden_logits or issue43_greedy or filtered_top_p'
+```
+
For a bounded real-checkpoint feasibility run:
```bash
python -m benchmarks.loop_speculative \
--model /path/to/local/Ouro-1.4B \
- --cuda-graphs --k 4 --concurrency 1 --new-tokens 32 --repeats 5 \
+ --cuda-graphs --k 4 --concurrency 1 --new-tokens 32 --repeats 3 \
--output /tmp/loop-speculative.json
# Sampling variant (top-p + top-k):
python -m benchmarks.loop_speculative \
--model /path/to/local/Ouro-1.4B \
- --cuda-graphs --k 4 --concurrency 8 --new-tokens 32 --repeats 5 \
+ --cuda-graphs --k 4 --concurrency 8 --new-tokens 32 --repeats 3 \
--temperature 0.8 --top-p 0.8 --top-k 7 \
--output /tmp/loop-speculative-top-p-top-k.json
```
@@ -254,9 +294,9 @@ For large loop batches, `--cuda-graph-max-batch-size 256` allows Graph capture o
the 2C mixed-depth wave when C=128; the default is 128.
`--prompts` accepts a JSON list of strings; `--backend flash_attn_4` requires
compatible hardware. Peak memory includes both resident engines; the recorded
-checkpoint config hash is not a weight hash. This small run does not replace
-RFC #43's fixed-input real-model diagnostics, sampling distribution checks,
-long-context/serving matrix or BF16 acceptance criteria.
+checkpoint config hash is not a weight hash. These checks cover fixed-input
+reuse and sampling distribution contracts, but do not replace Issue #43's
+long-context, serving-latency or full BF16 acceptance matrix.
Raw benchmark JSON is stored under `artifacts/loop-speculative/`; the extended
run notes include the runtime metadata and per-repetition timings. This work is
diff --git a/tests/test_speculative.py b/tests/test_speculative.py
index bdcd70ee..01b1601f 100644
--- a/tests/test_speculative.py
+++ b/tests/test_speculative.py
@@ -81,6 +81,31 @@ def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget, pipeline
assert llm.engine.cache_manager.num_used_blocks == 0
+@pytest.mark.parametrize("k", [1, 2, 4, 8])
+def test_issue43_greedy_generation_agrees_at_eight_concurrent_requests(k):
+ """Free generation stays byte-for-byte equal for the issue's K sweep."""
+ m = model()
+ prompts = [[2 + i, 7 - (i % 4), 11 + (i % 3)] for i in range(8)]
+ params = [SamplingParams(max_tokens=7, ignore_eos=True) for _ in prompts]
+ native = LLM(m).generate(prompts, params)
+ speculative = LLM(
+ m,
+ speculative_config=SpeculativeConfig(k, loop_level_pipelining=True),
+ cache_config=CacheConfig(256, 2),
+ ).generate(prompts, params)
+ for expected, actual in zip(native, speculative):
+ first = next(
+ (
+ i
+ for i, (left, right) in enumerate(zip(expected.token_ids, actual.token_ids))
+ if left != right
+ ),
+ None,
+ )
+ assert actual.token_ids == expected.token_ids, f"first divergence at {first}"
+ assert actual.exit_depths == expected.exit_depths
+
+
@pytest.mark.parametrize(
"device,backend",
[
@@ -91,7 +116,9 @@ def test_greedy_matches_native_with_ragged_rounds_and_limits(k, budget, pipeline
)
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("pipelined", [False, True])
-def test_reused_hidden_logits_and_all_kv_match_serial_oracle(dtype, device, backend, pipelined):
+def test_reused_hidden_logits_and_all_kv_match_serial_oracle(
+ dtype, device, backend, pipelined, record_property
+):
if device == "cuda" and dtype == torch.float32 and backend == "flash_attn_4":
pytest.skip("FA4 requires 16-bit operands")
m = model(dtype=dtype).to(device)
@@ -146,14 +173,23 @@ def record_verify(items, drafted, proposals, logits):
oracle.allocate("r", 16)
atol, rtol = (4e-5, 4e-5) if dtype == torch.float32 else (0.08, 0.04)
max_logit_error = 0.0
+ hidden_errors = []
+ kv_errors = []
+ logit_errors = []
+ logit_argmax_mismatches = 0
+ logit_rows = 0
for pos, token in enumerate(tokens):
h = m.prelude(torch.tensor([token], device=device))
for d in range(4):
h, _ = m.recurrent(h, ["r"], [d], [pos], oracle)
if pos >= 3:
+ hidden_errors.append((seen[d, pos] - h[0]).detach().float().reshape(-1))
torch.testing.assert_close(seen[d, pos], h[0], atol=atol, rtol=rtol)
if pos >= 3:
actual, expected = verified_logits[pos - 3], m.coda(h)[0]
+ logit_errors.append((actual - expected).detach().float().reshape(-1))
+ logit_argmax_mismatches += int(actual.argmax() != expected.argmax())
+ logit_rows += 1
max_logit_error = max(max_logit_error, float((actual - expected).abs().max()))
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
assert actual.argmax() == expected.argmax()
@@ -162,6 +198,7 @@ def record_verify(items, drafted, proposals, logits):
for actual, expected in zip(
cache.read(layer, "r", d, 7), oracle.read(layer, "r", d, 7)
):
+ kv_errors.append((actual - expected).detach().float().reshape(-1))
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
# Every input/depth computed exactly once: no shallow recomputation in verify.
assert work == Counter({0: 4, 1: 4, 2: 4, 3: 4})
@@ -176,6 +213,36 @@ def record_verify(items, drafted, proposals, logits):
)
assert max_logit_error < (1e-4 if dtype == torch.float32 else 0.025)
+ # Keep the RFC #43 fixed-input metrics visible in the assertion contract:
+ # max absolute and RMS error are tracked independently for hidden, KV and
+ # final-logit reuse, and target argmaxes must not diverge.
+ def metrics(errors):
+ values = torch.cat(errors)
+ return float(values.abs().max()), float(values.square().mean().sqrt())
+
+ hidden_max, hidden_rms = metrics(hidden_errors)
+ kv_max, kv_rms = metrics(kv_errors)
+ logits_max, logits_rms = metrics(logit_errors)
+ record_property("hidden_max_abs", hidden_max)
+ record_property("hidden_rms", hidden_rms)
+ record_property("kv_max_abs", kv_max)
+ record_property("kv_rms", kv_rms)
+ record_property("logits_max_abs", logits_max)
+ record_property("logits_rms", logits_rms)
+ record_property("target_argmax_mismatch_rate", logit_argmax_mismatches / logit_rows)
+ assert logit_argmax_mismatches == 0, (
+ f"fixed-input target argmax mismatch rate={logit_argmax_mismatches / logit_rows:.6f}"
+ )
+ assert hidden_max <= (4e-5 if dtype == torch.float32 else 0.08), (
+ f"hidden max_abs={hidden_max:.6g}, rms={hidden_rms:.6g}"
+ )
+ assert kv_max <= (4e-5 if dtype == torch.float32 else 0.08), (
+ f"KV max_abs={kv_max:.6g}, rms={kv_rms:.6g}"
+ )
+ assert logits_max <= (1e-4 if dtype == torch.float32 else 0.025), (
+ f"logits max_abs={logits_max:.6g}, rms={logits_rms:.6g}"
+ )
+
@pytest.mark.parametrize("rejected", [0, 1, 2, 3])
def test_each_rejection_position_and_bonus_commit_correct_frontier(rejected, monkeypatch):
@@ -435,6 +502,29 @@ def test_probability_filters_match_expected_and_apply_strict_top_k():
torch.testing.assert_close(p[p > 0], torch.ones(1))
+@pytest.mark.parametrize(
+ "params",
+ [
+ SamplingParams(temperature=0.8, top_p=0.8),
+ SamplingParams(temperature=0.8, top_k=3, top_p=0.9),
+ ],
+)
+def test_rejection_sampling_respects_filtered_top_p_and_top_k(params):
+ """The acceptance rule targets the filtered p distribution, not raw logits."""
+ target_logits = torch.tensor([3.0, 2.1, 1.5, 0.4, -0.2, -1.0])
+ proposal_logits = torch.tensor([1.0, 2.8, 2.0, 1.2, 0.5, -0.5])
+ p = Sampler.probabilities(target_logits, params)
+ q = Sampler.probabilities(proposal_logits, params)
+ generator = torch.Generator().manual_seed(9173)
+ counts = torch.zeros_like(p)
+ for _ in range(16000):
+ candidate = int(torch.multinomial(q, 1, generator=generator))
+ token, _ = rejection_sample(candidate, p, q, generator)
+ counts[token] += 1
+ empirical = counts / counts.sum()
+ torch.testing.assert_close(empirical, p, atol=0.02, rtol=0)
+
+
@pytest.mark.parametrize(
"kwargs",
[
From ada1e4efe3e70288eaef4c5e504d47daa58935df Mon Sep 17 00:00:00 2001
From: yudeng <2272776615@qq.com>
Date: Sun, 18 Oct 2026 11:58:36 +0000
Subject: [PATCH 11/20] docs: clarify sampling results and throughput units
---
docs/loop_speculative.md | 14 +++++++++-----
1 file changed, 9 insertions(+), 5 deletions(-)
diff --git a/docs/loop_speculative.md b/docs/loop_speculative.md
index 4f2cda74..a92d5dd8 100644
--- a/docs/loop_speculative.md
+++ b/docs/loop_speculative.md
@@ -53,8 +53,6 @@ draft depth, its coda output starts a child branch while the parent continues
toward the target depth. When a branch reaches the target depth, its target
logits are saved for the existing verifier.
-