diff --git a/benchmarks/loop_speculative.py b/benchmarks/loop_speculative.py new file mode 100644 index 00000000..3b422135 --- /dev/null +++ b/benchmarks/loop_speculative.py @@ -0,0 +1,276 @@ +"""Paired block/loop-level speculative decode comparison on a local Ouro model. + +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 +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.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, 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()): + engine.step() + engine.speculative_runner.stats = SpeculativeStats() + torch.cuda.synchronize() + 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 + torch.cuda.synchronize() + elapsed = time.perf_counter() - start + committed = len(prompts) * new_tokens + stats = asdict(engine.speculative_runner.stats) + 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, + steps=steps, + 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=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) + 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") + 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") + 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()) + 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 * 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, + cuda_graph_max_batch_size=args.cuda_graph_max_batch_size, + cuda_graph_max_graphs=args.cuda_graph_max_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( + num_speculative_tokens=args.k, + draft_loops=args.draft_loops, + target_loops=model.config.total_ut_steps, + loop_level_pipelining=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, + 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)) + for name in order: + result = trial( + engines[name], + prompts, + args.new_tokens, + 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() + + 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], + ) + # 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( + [ + 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=paired["median"], + paired_speedup=paired, + cuda_graphs=graphs, + first_divergent_output_positions=agreement, + trials=raw, + 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, paired_speedup=paired, graphs=graphs, divergence=agreement), + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/docs/README.md b/docs/README.md index d83666cb..423b0018 100644 --- a/docs/README.md +++ b/docs/README.md @@ -16,6 +16,7 @@ an environment to receiving the first generated response. The | [Runtime execution and configuration](runtime.md) | Batch contracts, preparation, transfers, execution options and compatibility | | [Split prefill and WCPB](split-prefill-wavefront.md) | Chunk/depth scheduling, hidden-state ring, KV dependencies and validation limits | | [Self-speculation](async-speculative.md) | Cross-round asynchronous block execution, token pipelines, and acceptance/rejection timelines | +| [Loop-level speculation](loop_speculative.md) | Opt-in mixed-depth execution, KV dependencies, and feasibility measurements | | [Asynchronous scheduling](https://github.com/InferMatrixLab/vllm-rlt/pull/30) | CPU/GPU pipelining and single-stream or multi-stream execution | | [FlashAttention](https://github.com/InferMatrixLab/vllm-rlt/pull/30) | FA2/FA3/FA4 installation, hardware selection, and constraints | | [Cache and scheduling features](https://github.com/InferMatrixLab/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..8ab49dd1 --- /dev/null +++ b/docs/loop_speculative.md @@ -0,0 +1,192 @@ +# Loop-level speculative execution + +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 + +```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`. +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 | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | +| A | 0 | 1 | 2 | 3 | | | | | +| 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 +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 + +```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 --run-gpu -m gpu -k 'not flash_attn_4' \ + tests/test_speculative.py tests/test_loop_speculative.py +``` + +`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. + +Paired benchmark: + +```bash +python -m benchmarks.loop_speculative \ + --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 +``` + +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. diff --git a/tests/test_loop_speculative.py b/tests/test_loop_speculative.py new file mode 100644 index 00000000..e0d4523d --- /dev/null +++ b/tests/test_loop_speculative.py @@ -0,0 +1,245 @@ +"""Mixed-depth schedule, round equivalence and pipeline rollback contracts.""" + +import argparse +from collections import Counter + +import pytest +import torch + +from tests.helpers import LocalLLM as LLM +from tests.test_speculative import drain, engine, model +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, 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) + + +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=draft_loops, target_loops=4, loop_level_pipelining=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_loop_round_matches_block_round_and_batches_mixed_depths(draft_loops, sampling): + m = model() + 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 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("temperature", [0.0, 0.8]) +def test_loop_round_runs_only_draft_heads_inside_the_loop(temperature, monkeypatch): + m = model() + 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 + runner = e.speculative_runner + 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_loop_generation_agrees_at_eight_concurrent_requests(): + 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) + 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) + + def run(pipelined): + return LLM( + m, speculative_config=SpeculativeConfig(3, loop_level_pipelining=pipelined) + ).generate(prompts, params) + + assert [o.token_ids for o in run(True)] == [o.token_ids for o in run(False)] + + +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) + with pytest.raises(ValueError, match="mutually exclusive"): + SpeculativeConfig(3, mode="pipeline", loop_level_pipelining=True) + + +def test_pipeline_partial_execution_failure_releases_kv(monkeypatch): + 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 + core = m.recurrent_prepared + calls = 0 + + def fail_after_write(*args, **kwargs): + nonlocal calls + result = core(*args, **kwargs) + calls += 1 + if calls == 3: + raise RuntimeError("partial pipeline failure") + return result + + 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 b2593c83..306cb42f 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] @@ -434,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", [ @@ -455,8 +495,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 +529,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 +544,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 +559,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 +589,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 +598,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 +631,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 +654,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 +683,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 +696,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,15 +766,17 @@ def test_abort_suspended_speculative_request_allows_id_reuse(): @pytest.mark.parametrize("chunk_size", [1, 3, 8]) @pytest.mark.parametrize("budget", [2, 7]) +@pytest.mark.parametrize("pipelined", [False, True]) @pytest.mark.parametrize("wavefront", [False, True]) def test_split_prefill_and_speculation_share_forward_and_match_atomic( - chunk_size, budget, wavefront + chunk_size, budget, pipelined, wavefront ): from vllm_rlt.request import Operation def run(wavefront): e = engine( model(), + pipelined=pipelined, cache_config=CacheConfig(256, 2), scheduler_config=SchedulerConfig( max_num_batched_tokens=budget, @@ -780,8 +839,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 = [] @@ -849,8 +911,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 648781d7..d52dfaf2 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -180,13 +180,18 @@ class SpeculativeConfig: draft_loops: int = 2 target_loops: int = 4 mode: str = "block" + 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.mode not in ("block", "pipeline"): raise ValueError("speculative mode must be block or pipeline") + if self.mode == "pipeline" and self.loop_level_pipelining: + raise ValueError("mode=pipeline and loop_level_pipelining are mutually exclusive") if self.mode == "pipeline" and self.target_loops % self.draft_loops: raise ValueError("pipeline draft_loops must divide target_loops") 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 883d884b..2a0b8983 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -19,6 +19,11 @@ def add_runtime_args(parser): parser.add_argument("--speculative-mode", choices=["block", "pipeline"], default="block") 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) @@ -90,12 +95,15 @@ 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, mode=args.speculative_mode, + 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 85f0e235..857823a8 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, replace +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" @@ -259,6 +284,8 @@ def submit(self, batch): ): if self.config.mode == "pipeline": submission = self._submit_pipeline(batch) + elif 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) @@ -456,6 +483,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