diff --git a/benchmarks/loopcd_scale.py b/benchmarks/loopcd_scale.py new file mode 100644 index 00000000..f3c38a84 --- /dev/null +++ b/benchmarks/loopcd_scale.py @@ -0,0 +1,314 @@ +"""Finite closed-loop native E2E for fixed-depth Ouro LoopCD.""" + +import argparse +import hashlib +import json +import math +import statistics +import subprocess +import time +from collections import Counter +from dataclasses import asdict +from pathlib import Path + +import torch + +from vllm_rlt import ( + LLM, + CacheConfig, + ExecutionConfig, + LoopCDParams, + SamplingParams, + SchedulerConfig, +) +from vllm_rlt.request import Stage + + +def distribution(values): + ordered = sorted(values) + return { + "n": len(values), + "mean": statistics.mean(values), + "median": statistics.median(values), + "p95": ordered[math.ceil(0.95 * len(values)) - 1], + "p99": ordered[math.ceil(0.99 * len(values)) - 1], + "min": ordered[0], + "max": ordered[-1], + } + + +def graph_counts(engine): + bank = engine.model_runner.graphs + return ( + dict(captures=bank.captures, replays=bank.replays, fallbacks=bank.fallbacks) + if bank + else dict(captures=0, replays=0, fallbacks=0) + ) + + +def synchronize(engine): + if engine.model_runner.device.type == "cuda": + torch.cuda.synchronize(engine.model_runner.device) + + +@torch.inference_mode() +def trial(engine, prompts, sampling, concurrency): + """Refill reused IDs only after consuming the complete output batch.""" + if engine.has_unfinished_requests() or engine.cache_manager.num_used_blocks: + raise ValueError("trial requires a drained engine") + if not 1 <= concurrency <= len(prompts) or sampling.max_tokens < 2: + raise ValueError("require 1 <= C <= requests and at least two output tokens") + if not sampling.ignore_eos or sampling.min_loops != sampling.max_loops: + raise ValueError("fixed work requires ignore_eos and fixed decode depth") + runner, cache = engine.model_runner, engine.cache_manager + before_graph = graph_counts(engine) + before_work = dict(runner.loopcd_stats) + synchronize(engine) + if runner.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(runner.device) + start = time.perf_counter() + active, traces, free = {}, {}, list(range(concurrency)) + next_request = kv_peak = 0 + previous_decode_rows = 0 + resident, outstanding, core_rows = Counter(), Counter(), Counter() + phases = Counter() + + def fill(): + nonlocal next_request + while free and next_request < len(prompts): + slot = free.pop() + rid = f"slot-{slot}" + prompt = prompts[next_request] + engine.add_request(rid, prompt, sampling) + active[rid] = next_request + traces[next_request] = dict( + ordinal=next_request, + request_id=rid, + prompt=prompt, + submitted_s=time.perf_counter() - start, + admitted_s=None, + token_ids=[], + exit_depths=[], + delivery_s=[], + ) + next_request += 1 + + fill() + bound = len(prompts) * ( + max(map(len, prompts)) * engine.prefill_depth + + sampling.max_tokens * (sampling.max_loops + 4) + + 10 + ) + for _ in range(bound): + outstanding[len(active)] += 1 + outputs = engine.step() + now = time.perf_counter() - start + requests = engine.scheduler.requests + resident[ + sum(r.stage not in (Stage.WAITING, Stage.RECEIVING) for r in requests.values()) + ] += 1 + for rid, request in requests.items(): + row = traces[active[rid]] + if request.stage not in (Stage.WAITING, Stage.RECEIVING) and row["admitted_s"] is None: + row["admitted_s"] = now + schedule = engine.last_schedule + if schedule is not None: + phases[f"{schedule.stage.value}:{len(schedule.items)}"] += 1 + work = runner.loopcd_stats + submitted = work["decode_core_rows"] - before_work["decode_core_rows"] + # Counts are taken from actual model runner row counters, not B/C settings. + if submitted != previous_decode_rows: + core_rows[submitted - previous_decode_rows] += 1 + previous_decode_rows = submitted + kv_peak = max(kv_peak, cache.num_used_blocks) + released = [] + for output in outputs: + row = traces[active[output.request_id]] + count = len(row["token_ids"]) + if ( + output.prompt_token_ids != row["prompt"] + or output.token_ids[:count] != row["token_ids"] + ): + raise RuntimeError("request identity or output prefix changed") + row["delivery_s"].extend([now] * (len(output.token_ids) - count)) + row["token_ids"], row["exit_depths"] = output.token_ids, output.exit_depths + if output.finished: + expected = [engine.prefill_depth] + [sampling.max_loops] * (sampling.max_tokens - 1) + if ( + len(output.token_ids) != sampling.max_tokens + or output.exit_depths != expected + or output.finish_reason != "length" + ): + raise RuntimeError("fixed work/depth/finish contract failed") + row.update(finished_s=now, finish_reason=output.finish_reason) + released.append(int(output.request_id.split("-")[1])) + del active[output.request_id] + free.extend(sorted(released)) + fill() + if not active: + break + else: + raise RuntimeError("finite work bound exhausted") + synchronize(engine) + elapsed = time.perf_counter() - start + if engine.has_unfinished_requests() or cache.num_used_blocks or runner.loopcd_references: + raise RuntimeError("KV/reference/request leak after drain") + rows = [traces[i] for i in range(len(prompts))] + after_graph = graph_counts(engine) + return dict( + elapsed_s=elapsed, + tps=len(prompts) * sampling.max_tokens / elapsed, + rps=len(prompts) / elapsed, + requests=rows, + ttft_s=distribution([r["delivery_s"][0] - r["submitted_s"] for r in rows]), + latency_s=distribution([r["finished_s"] - r["submitted_s"] for r in rows]), + tpot_s=distribution( + [(r["delivery_s"][-1] - r["delivery_s"][0]) / (sampling.max_tokens - 1) for r in rows] + ), + itl_s=distribution( + [b - a for r in rows for a, b in zip(r["delivery_s"], r["delivery_s"][1:])] + ), + admission_wait_s=distribution([r["admitted_s"] - r["submitted_s"] for r in rows]), + resident_histogram=dict(resident), + outstanding_histogram=dict(outstanding), + decode_effective_rows_histogram=dict(core_rows), + schedule_histogram=dict(phases), + kv_peak_blocks=kv_peak, + kv_pool_bytes=cache.num_blocks * cache.bytes_per_block, + loopcd_work={ + k: v - before_work[k] + for k, v in runner.loopcd_stats.items() + if k != "reference_peak_bytes" + }, + reference_peak_bytes=runner.loopcd_stats["reference_peak_bytes"], + graph={k: after_graph[k] - before_graph[k] for k in after_graph}, + peak_allocated_bytes=torch.cuda.max_memory_allocated(runner.device) + if runner.device.type == "cuda" + else None, + peak_reserved_bytes=torch.cuda.max_memory_reserved(runner.device) + if runner.device.type == "cuda" + else None, + drained=True, + ) + + +def signature(result): + return [ + (r["ordinal"], r["request_id"], r["token_ids"], r["exit_depths"], r["finish_reason"]) + for r in result["requests"] + ] + + +def run(engine, prompts, sampling, concurrency, output, warmups=5, repetitions=5): + output.mkdir(parents=True, exist_ok=False) + expected = None + summary = [] + for index in range(warmups + repetitions): + result = trial(engine, prompts, sampling, concurrency) + result.update(index=index, phase="warmup" if index < warmups else "measured") + current = signature(result) + if expected is not None and current != expected: + result["repeat_parity"] = False + else: + result["repeat_parity"] = True + expected = current + (output / f"trial-{index:02d}.json").write_text(json.dumps(result, indent=2) + "\n") + if not result["repeat_parity"]: + raise RuntimeError("repeat token/depth/ID/finish mismatch; raw result preserved") + if index >= warmups and result["graph"]["captures"]: + raise RuntimeError("cold graph capture during measured work; raw result preserved") + summary.append({k: v for k, v in result.items() if k != "requests"}) + print( + json.dumps( + {"index": index, "phase": result["phase"], "tps": result["tps"], "drained": True} + ), + flush=True, + ) + (output / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--revision", default="574fa66cb8bf5abdc979642d01cf2b79b16bfab1") + parser.add_argument("--batch", type=int, required=True) + parser.add_argument("--concurrency", type=int, required=True) + parser.add_argument("--input-lengths", type=int, nargs="+", required=True) + parser.add_argument("--output-tokens", type=int, default=128) + parser.add_argument("--depth", type=int, choices=[2, 3, 4], default=4) + parser.add_argument("--strength", type=float, choices=[0.0, 0.3], default=0.3) + parser.add_argument("--kv-gib", type=float, required=True) + parser.add_argument("--graph", action="store_true") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if not 1 <= args.batch <= args.concurrency <= 128 or min(args.input_lengths) < 1: + parser.error("require 1 <= B <= C <= 128 and positive input lengths") + if not Path(args.model).is_dir(): + parser.error("--model must be a verified local checkpoint directory") + import vllm_rlt + + source_root = Path(__file__).resolve().parents[1] + if Path(vllm_rlt.__file__).resolve().parents[1] != source_root: + raise RuntimeError("imported vllm_rlt belongs to another checkout") + subprocess.run(["git", "diff", "--quiet", "HEAD"], cwd=source_root, check=True) + source_sha = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=source_root, text=True + ).strip() + torch.set_num_threads(1) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + llm = LLM( + args.model, + revision=args.revision, + device="cuda", + dtype=torch.bfloat16, + attention_backend="triton", + cache_config=CacheConfig(kv_cache_memory_bytes=int(args.kv_gib * 1024**3)), + scheduler_config=SchedulerConfig( + max_num_seqs=args.batch, + max_num_batched_tokens=max(128, args.batch), + prefill_chunk_size=128, + ), + execution_config=ExecutionConfig( + loopcd=True, prefill_depth=4, cuda_graphs=args.graph, cuda_graph_max_graphs=64 + ), + ) + from vllm_rlt.models import load_tokenizer + + tokenizer = load_tokenizer(args.model, revision=args.revision) + seed_ids = tokenizer.encode( + "The scientist studied the stars and recorded the results. ", add_special_tokens=False + ) + prompts = [ + (seed_ids * math.ceil(length / len(seed_ids)))[:length] + for length in ( + args.input_lengths * math.ceil(4 * args.concurrency / len(args.input_lengths)) + )[: 4 * args.concurrency] + ] + sampling = SamplingParams( + max_tokens=args.output_tokens, + min_loops=args.depth, + max_loops=args.depth, + ignore_eos=True, + temperature=0.0, + loopcd=LoopCDParams(strength=args.strength) if args.strength else None, + ) + manifest = dict( + source_sha=source_sha, + args={**vars(args), "output": str(args.output)}, + sampling=asdict(sampling), + torch=torch.__version__, + cuda=torch.version.cuda, + device=torch.cuda.get_device_name(), + prompt_sha256=hashlib.sha256(json.dumps(prompts).encode()).hexdigest(), + memory_plan=llm.engine.memory_plan, + scope="native fixed work; no task quality or HTTP claim", + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.with_suffix(".manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + run(llm.engine, prompts, sampling, args.concurrency, args.output) + llm.close() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/results/loopcd-c32/summary.json b/benchmarks/results/loopcd-c32/summary.json new file mode 100644 index 00000000..23a65ade --- /dev/null +++ b/benchmarks/results/loopcd-c32/summary.json @@ -0,0 +1,305 @@ +{ + "schema_version": 1, + "tested_source_revision": "e8aa1ee57734bc5c1118c143ebdb68d0b53210a1", + "production_delta_since_task0": false, + "h100_fixed_work": { + "hardware": "one H100", + "cuda_version": "13.0", + "dependencies": { + "python": "3.12.1", + "torch": "2.12.1+cu130", + "triton": "3.7.1", + "transformers": "4.54.1", + "tokenizers": "0.21.4", + "safetensors": "0.8.0", + "huggingface-hub": "0.35.3", + "numpy": "2.4.6", + "host_compiler": "GCC 11.5.0" + }, + "model": "ByteDance/Ouro-1.4B", + "model_revision": "574fa66cb8bf5abdc979642d01cf2b79b16bfab1", + "source_sha": "e8aa1ee57734bc5c1118c143ebdb68d0b53210a1", + "dtype": "float32", + "attention_backend": "triton", + "prefill_depth": 4, + "B": 16, + "initial_total_C": 32, + "actual_A": 16, + "actual_S": 16, + "P": 512, + "F": 128, + "fresh_processes": 6, + "warmups_per_process": 5, + "measurements_per_process": 5, + "trials": 60, + "requests": 1920, + "output_tokens": 245760, + "measured_requests": 960, + "measured_output_tokens": 122880, + "order": "ABC/CBA with A=D4 off, B=D3 off, C=D3 guided", + "rows": { + "D4 true off": { + "mean_TPS": 179.2798970731368, + "mean_TTFT_ms": 9133.832768897173, + "mean_TPOT_ms": 63.18140338326404, + "mean_native_latency_ms": 17157.870998571707, + "mean_native_ITL_p99_ms": 303.6823431029916, + "matched_fresh_TPS_ratios_vs_D4": [ + 1.0, + 1.0 + ], + "mean_TPS_gain_percent_vs_D4": 0.0, + "peak_reference_bytes": 0, + "peak_allocated_delta_vs_D4_bytes": 0, + "peak_reserved_delta_vs_D4_bytes": 0 + }, + "D3 true off": { + "mean_TPS": 212.989612611624, + "mean_TTFT_ms": 8191.6508564029755, + "mean_TPOT_ms": 49.251064304719584, + "mean_native_latency_ms": 14446.536023102362, + "mean_native_ITL_p99_ms": 241.91237840568647, + "matched_fresh_TPS_ratios_vs_D4": [ + 1.1826452043318265, + 1.1934382505108443 + ], + "mean_TPS_gain_percent_vs_D4": 18.802841862819353, + "peak_reference_bytes": 0, + "peak_allocated_delta_vs_D4_bytes": 0, + "peak_reserved_delta_vs_D4_bytes": 0 + }, + "D3 guided ref1 w0.3": { + "mean_TPS": 207.10076799059408, + "mean_TTFT_ms": 8371.483820582307, + "mean_TPOT_ms": 51.022023927212494, + "mean_native_latency_ms": 14851.280859338294, + "mean_native_ITL_p99_ms": 246.3158268976258, + "matched_fresh_TPS_ratios_vs_D4": [ + 1.156436304849139, + 1.1539199043585666 + ], + "mean_TPS_gain_percent_vs_D4": 15.518120755116138, + "peak_reference_bytes": 131072, + "peak_allocated_delta_vs_D4_bytes": 131072, + "peak_reserved_delta_vs_D4_bytes": 2097152 + } + }, + "mean_TPS_scope": "Arithmetic mean of the two fresh process means; each contains five measured equal-work trials, five warmup trials retained separately", + "speed_gain_formula": "(arm_mean_TPS / D4_mean_TPS - 1) * 100", + "fresh_ratio_scope": "Each candidate fresh process divided by matching fresh D4 mean; range of two fresh ratios, not a confidence interval" + }, + "h100_fp32_qualification": { + "cases": 12, + "readouts": 108, + "forced_continuation_inputs": 96, + "ownership_checks": 21, + "key_checks": 82944, + "value_checks": 82944, + "max_abs": { + "state": 6.29425048828125e-05, + "key": 3.0994415283203125e-05, + "value": 1.0848045349121094e-05, + "logits": 5.698204040527344e-05, + "selected_lp": 2.86102294921875e-05, + "readout_state": 0.0 + }, + "atol": 0.0001, + "rtol": 0.0001 + }, + "h100_quality_confirmation": { + "required_outputs": 1959, + "complete_outputs": 1827, + "not_run_budget_outputs": 132, + "generation_counts": [ + { + "arm_id": "A_P4D4_off", + "task": "gsm8k", + "required_n": 505, + "complete_n": 505, + "not_run_budget_n": 0 + }, + { + "arm_id": "A_P4D4_off", + "task": "humanevalplus", + "required_n": 148, + "complete_n": 148, + "not_run_budget_n": 0 + }, + { + "arm_id": "D_P4D3_on", + "task": "gsm8k", + "required_n": 505, + "complete_n": 505, + "not_run_budget_n": 0 + }, + { + "arm_id": "D_P4D3_on", + "task": "humanevalplus", + "required_n": 148, + "complete_n": 148, + "not_run_budget_n": 0 + }, + { + "arm_id": "C_P4D3_off", + "task": "gsm8k", + "required_n": 505, + "complete_n": 505, + "not_run_budget_n": 0 + }, + { + "arm_id": "C_P4D3_off", + "task": "humanevalplus", + "required_n": 148, + "complete_n": 16, + "not_run_budget_n": 132 + } + ], + "mathematics": { + "denominator": 505, + "arms": [ + { + "arm_id": "A_P4D4_off", + "n": 505, + "completed": 505, + "missing_or_unrun": 0, + "correct": 316, + "accuracy": 0.6257425742574257, + "hit_max_tokens": 0, + "status": "COMPLETE" + }, + { + "arm_id": "D_P4D3_on", + "n": 505, + "completed": 505, + "missing_or_unrun": 0, + "correct": 301, + "accuracy": 0.596039603960396, + "hit_max_tokens": 4, + "status": "COMPLETE" + }, + { + "arm_id": "C_P4D3_off", + "n": 505, + "completed": 505, + "missing_or_unrun": 0, + "correct": 287, + "accuracy": 0.5683168316831683, + "hit_max_tokens": 2, + "status": "COMPLETE" + } + ], + "comparisons": [ + { + "candidate": "D_P4D3_on", + "baseline": "A_P4D4_off", + "n": 505, + "wins": 31, + "losses": 46, + "family_comparisons": 2, + "difference": -0.0297029702970297, + "lower": -0.09008988535465567, + "upper": 0.031290344982920076, + "status": "UNCERTAIN", + "scope": "primary mathematics confirmation; frozen family M2" + }, + { + "candidate": "D_P4D3_on", + "baseline": "C_P4D3_off", + "n": 505, + "wins": 44, + "losses": 30, + "family_comparisons": 2, + "difference": 0.027722772277227723, + "lower": -0.03222355049747422, + "upper": 0.08709923492766208, + "status": "UNCERTAIN", + "scope": "exploratory same-depth attribution; not an additional confirmatory claim" + } + ] + }, + "code": { + "denominator": 148, + "arms": [ + { + "arm_id": "A_P4D4_off", + "n": 148, + "generated": 148, + "scored": 148, + "correct": 1, + "status": "COMPLETE" + }, + { + "arm_id": "D_P4D3_on", + "n": 148, + "generated": 148, + "scored": 148, + "correct": 0, + "status": "COMPLETE" + }, + { + "arm_id": "C_P4D3_off", + "n": 148, + "generated": 16, + "scored": 0, + "correct": null, + "status": "NOT_RUN" + } + ], + "comparisons": [ + { + "candidate": "D_P4D3_on", + "baseline": "A_P4D4_off", + "n": 148, + "scope": "primary code confirmation; frozen family M2", + "status": "UNCERTAIN", + "interval": { + "n": 148, + "wins": 0, + "losses": 1, + "family_comparisons": 2, + "difference": -0.006756756756756757, + "lower": -0.0474900317031247, + "upper": 0.03366805626254722 + } + }, + { + "candidate": "D_P4D3_on", + "baseline": "C_P4D3_off", + "n": 148, + "scope": "exploratory same-depth attribution; no additional confirmation claim", + "status": "INCOMPLETE", + "interval": null + } + ], + "execution": "Original cooperative20M opcodes/5s per call; complete canonical reference; no repairs; not official EvalPlus timeout execution" + }, + "noninferiority_margin": 0.01 + }, + "h20_fixed_work": { + "policies": [ + "D4_off", + "D3_off", + "D3_guided" + ], + "mean_tokens_per_second": [ + 144.177, + 167.121, + 164.701 + ], + "mean_paired_speed_ratio_vs_D4": [ + 1.0, + 1.159342, + 1.142596 + ], + "guided_vs_same_depth_paired_throughput_change_percent": -1.436, + "comparison_estimator": "mean of ten matched elapsed-time ratios; separate device profile" + }, + "scope": { + "full_model_C64_C128": false, + "HTTP": false, + "engine_async": false, + "quality_noninferiority_established": false, + "VRAM_saving_established": false, + "whole_task_complete": false + } +} diff --git a/docs/loopcd-ouro.md b/docs/loopcd-ouro.md new file mode 100644 index 00000000..389b45bd --- /dev/null +++ b/docs/loopcd-ouro.md @@ -0,0 +1,150 @@ +# Initial LoopCD support for Ouro + +Task 0 of [RFC #85](https://github.com/ThinkFlowLab/vllm-rlt/issues/85) +adds fixed-depth synchronous Ouro readout guidance from +[LoopCD, arXiv:2610.02185v1](https://arxiv.org/abs/2610.02185). +Guidance is off by default. The recurrent stack, normalized feedback, gate, +KV writes and sampling transforms retain their original behavior. + +```text +Same token prefix -> original recurrent stack, updates 1 ... R + | b | R + v v + owned normalized h_b normalized h_R + | | + original head original head + | | + z_b z_R + `--------> z_R + w*(z_R-z_b) -> sampler + +Reference owner = Request object + token position + output index + depth +Next token clears the old reference; completion/abort releases it before reuse. +``` + +The reference depth is one-based and must precede the actual prefill/decode +depth. `two_head` projects both normalized states and combines logits in FP32. +It is the enabled default. Optional `linear_fused` projects +`h_R + w*(h_R-h_b)` through the existing linear head, without another norm. +Its algebraic identity does not imply BF16 equality or identical low-margin +argmax decisions; it remains an explicit experimental choice. + +```python +from vllm_rlt import ExecutionConfig, LoopCDParams, SamplingParams + +execution = ExecutionConfig(loopcd=True, prefill_depth=4) +sampling = SamplingParams( + min_loops=2, + max_loops=2, + exit_threshold=1.0, + loopcd=LoopCDParams(reference_loop=1, strength=0.3), +) +``` + +This is fixed prefill P4/decode D2, not a P2/D2 run. The first emitted token +records depth 4 and subsequent tokens depth 2. `prefill_reference_loop` can +freeze a separate prompt reference. A zero strength follows the original +readout without reference allocation. + +Nonzero guidance requires `ExecutionConfig(loopcd=True)` and synchronous +`last_exited` KV with fixed decode depth and the original `ouro` exit mode. +Prefix caching, prefill UVA, asynchronous scheduling, speculation and P/D +transfer are rejected before request allocation. Other model families are +rejected. Single-worker preemption is supported with eager execution only; +combining guidance, preemption and CUDA Graphs remains rejected until separate +qualification. Official-checkpoint eager/Graph readout qualification is tracked +in Task 1. + +Automatic memory planning reserves owned references and conservative two-head +FP32 scratch before planning KV. Explicit KV budgets still require physical +device admission. Existing release-event ordering protects live references; +request-ID reuse cannot inherit another Request object's reference. + +## Correctness scope + +`tests/test_loopcd.py` checks the first readout against the repository's +independent functional dense Ouro oracle across prompt chunks and P/D depths. +Fixed P=D teacher-forced continuation checks all vocabulary logits, selected +log-probabilities, raw normalized state and top-1 decisions through abort and +request-ID reuse. It also covers zero-strength equivalence, refill/drain, +configuration rejection, stale owners and memory reservation. These use real +in-memory tiny FP32 CPU models. The affine FP32/FP64 test does not qualify BF16. + +Official-checkpoint continuation and fixed-work results are tracked in +[Task 0's audited results](loopcd-scale-results.md), including its retained +BF16 independent-oracle failure and uncertain quality confirmation. +Preemption qualification is described below. + + +## Single-worker preemption contract + +Enable the existing scheduler path with +`SchedulerConfig(enable_preemption=True)`. Suspension occurs only at the +existing safe scheduler boundaries; pending outputs and selected requests remain +ineligible. The Request object and its RNG stay owned by the scheduler. + +At first-token CODA, or after the decode reference depth has been reached, the +owned normalized reference is copied to a private CPU tensor alongside the +existing KV, hidden-state and input-token snapshot. Its Request identity, token +position, output index and one-based reference depth are retained. Resume checks +these fields before KV allocation, then restores the reference on the cache +device before re-enqueuing the saved stage. It does not replay recurrent work or +change the guidance equation. + +A partial prompt chunk, next-token PRELUDE, or decode before the reference depth +does not need the old reference: the resumed computation captures a fresh one. +A capacity-blocked resume keeps only the CPU snapshot. Abort discards both the +snapshot and runner-owned reference; a reused request ID gets a new Request +object and cannot inherit the old guidance state. + +`tests/test_loopcd_preemption.py` covers suspend/resume logits and seeded +sampling, all relevant stage boundaries, asymmetric P4/D2, repeated preemption, +capacity-blocked restore, cancellation, ID reuse and rejection of corrupted +reference identity/metadata. Scheduler-driven priority and KV-pressure tests +also check output parity and resource cleanup on CPU and Triton FP32 CUDA. +The CUDA cases cover static buffers off and on. The CUDA tests use tiny random checkpoints and +require the `gpu` marker. + +`tests/test_loopcd_preemption_continuation.py` uses the stateful +`OuroContinuation` oracle for P4D2/P4D3, with resumed CODA, PRELUDE and +RECURRENT stages, delayed KV allocation, cancellation and ID reuse. At each +readout it compares final hidden state, every historical K/V plane, guided +vocabulary logits and selected-token log probability. The original +`atol=rtol=1e-4` gates are unchanged. Reference-hidden error is recorded as a +separate diagnostic; saved references are checked bit-exactly. + +Official Ouro-1.4B FP32/Triton passes all 12 combinations of decode depth 2/3, +prompt length 3/4/5 and guidance strength 0/0.3 in two fresh processes: +264 readouts including the reused requests. Native BF16 B1 priority resumption +also has bit-exact vocabulary logits with static buffers off/on in both +processes. This BF16 parity result does not qualify the retained independent +BF16 numerical failure in [Task 0's results](loopcd-scale-results.md). + +Earlier official FP32 KV-pressure runs match tokens, exit depths and recurrent +work; score differences up to `2.4e-5` also occur in an uninterrupted batching +control. BF16 multi-request pressure remains unqualified. + +### Fixed-work preemption cost + +The RTX 5090 comparison uses FP32/eager Triton, P4D3, two_head/ref1/strength0.3, +B16, initial-total C32, P512/F128 and a common 18 GiB KV pool. The interrupted +arm suspends one request at completed-step CODA before output index 4, then +uses the native FCFS scheduler to restore it. Two reversed configuration orders +each run five warmups and five measurements per arm. The four-arm run retains +80 trials; the two guided arms have ten measured trials each. + +| Guided arm | Mean tokens/s | Mean native request latency | +| --- | ---: | ---: | +| Uninterrupted | 224.60 | 13.67 s | +| One controlled preemption | 214.55 | 14.55 s | + +All 20 trials per guided arm match request IDs, output tokens, exit depths and +finish state; model-row work is unchanged. Batch regrouping changes head calls +from 1670 to 1694. The snapshot contains 960 MiB of KV tensors and an 8 KiB +guidance reference, plus hidden state and written metadata. Actual successful +restore callbacks average 115.1 ms; suspension averages 546.3 ms. + +The paused request's next-token delivery gap averages 7.63 s, reaching 7.74 s; +it waits at the back of the FCFS queue. Copy time alone therefore understates +its interruption. These timings include native scheduling, KV transfer and +batch regrouping; they do not isolate the reference-copy cost or estimate HTTP +SLOs. See [raw trials, source hashes and retained diagnostics](https://github.com/prettygirlisnotme/vllm-rlt/releases/tag/loopcd-preemption-evidence-20261010). diff --git a/docs/loopcd-scale-results.md b/docs/loopcd-scale-results.md new file mode 100644 index 00000000..9e5f9c69 --- /dev/null +++ b/docs/loopcd-scale-results.md @@ -0,0 +1,93 @@ +# Audited Ouro LoopCD results + +The tested source is `e8aa1ee57734bc5c1118c143ebdb68d0b53210a1`, a direct +child of the original Task 0 source `d700ecca`. That increment adds the +mixed-depth continuation oracle, scale-validation entry point and documentation; +it does not change the production readout, recurrent update or feature guards. +The checkpoint is `ByteDance/Ouro-1.4B` at +`574fa66cb8bf5abdc979642d01cf2b79b16bfab1`. + +## Public loading and full-model correctness + +The public model entry passed the frozen tokenizer and parameter mapping checks +for FP32 and BF16: 52 tokenizer records, 269 parameter mappings and four facade +cases per dtype. Loading success is separate from mixed-depth numerical +qualification. + +Full-model FP32 mixed-depth qualification passed with both separately qualified +Torch and Triton attention paths. The original Triton protocol includes 12 +cases, 108 readouts, 96 forced continuation inputs, all four historical KV +planes and 21 ownership/lifecycle checks at unchanged `atol=rtol=1e-4`. +The same original FP32/Triton protocol subsequently passed on one H100, with +independent source/raw evidence review and selected-logprob scalar recomputation. + +The full-model BF16/Triton mixed-depth gate failed at its first P4 prefill +readout: 998/2048 hidden-state elements were outside the frozen tolerance, +with maximum absolute discrepancy 0.23046875. Subsequent checks were not run. +This failure is retained; neither loading checks nor tiny-model tests qualify +the failed full-model BF16 path. The earlier BF16 `linear_fused` rejection also +remains in force. `two_head` stays the reference, default-off implementation. + +## Completed C32 native fixed-work comparisons + +All arms use FP32/eager Triton, prefill depth 4, B16, initial-total C32, +512 prompt tokens and 128 emitted tokens per request. Each profile uses its +own unchanged common KV pool. Observed resident A and active decode S are 16. +Two fresh starts reverse A/B/C and C/B/A order; each arm/process includes +five warmups and five measurements. Each complete profile contains six +processes, 60 trials, 1920 requests and 245760 output tokens. The measured +subset contains 30 trials, 960 requests and 122880 tokens. + +| Profile | Policy | Mean tokens/s | TTFT mean ms | TPOT mean ms | Native request latency mean ms | +| --- | --- | ---: | ---: | ---: | ---: | +| H20 | D4, true off | 144.177 | 11945.174 | 74.588 | 21417.862 | +| H20 | D3, true off | 167.121 | 10934.280 | 59.380 | 18475.506 | +| H20 | D3, two_head/ref1/strength0.3 | 164.701 | 11076.655 | 60.521 | 18762.769 | +| One H100 | D4, true off | 179.28 | 9133.83 | 63.18 | 17157.87 | +| One H100 | D3, true off | 212.99 | 8191.65 | 49.25 | 14446.54 | +| One H100 | D3, two_head/ref1/strength0.3 | 207.10 | 8371.48 | 51.02 | 14851.28 | + +The H20 comparison uses the mean of ten matched elapsed-time ratios across +the two fresh starts: D3-off/D4 is 1.159342x and D3-guided/D4 is 1.142596x. +Guidance reduces paired throughput by 1.436% relative to the same-depth +D3-off control. The H100 comparison uses ratios of mean TPS: D3-off/D4 +is +18.80%, D3-guided/D4 is +15.52%; guided fresh-process gains are +15.39%-15.64%. These statistics are different estimators and the two-start +range is not a confidence interval. Device profiles are not pooled. + +Request IDs, emitted tokens/depths, finish state, retained references, drain, +native timestamps and KV accounting passed the original immutable raw auditor. +The complete results were independently verified off device. Guidance adds +128 KiB peak allocated reference state and 2 MiB peak reserved memory in these +common-pool comparisons; there is no measured VRAM saving. + +## Quality and remaining scope + +The separate FP32 development screen generated and scored all 128 questions +under six policies. D4-off and selected D3-guided both scored 80/128, but paired +uncertainty did not establish the frozen one-percentage-point non-inferiority +criterion. Natural-stop D3-guided timing retained two output-limit hits and +was slower than D4 in that screen; fixed-length cost cannot remove those tails. + +The independent H100 confirmation retains all 505 GSM8K and 148 HumanEval+ +prompt IDs and three arms. It completed 1827/1959 outputs, with 132 D3-off +code outputs not run within the original budget. D4-off mathematics scored +316/505 and D3-guided 301/505: -2.970 percentage points, frozen paired M2 +interval [-9.009, +3.129]. Code scored 1/148 versus 0/148 under the unchanged +cooperative scorer. Both primary quality comparisons remain uncertain; +the incomplete matched-off code subset is not scored as a full arm. + +Completed C32 results qualify this instrumented native fixed-work scope. +They do not establish full-model C64/C128, steady 4C refill, HTTP serving, +engine-async, new CUDA Graph performance, training convergence or quality +non-inferiority. Tiny BF16 eager/Graph C32/C64/C128 checks are retained as +separate tiny-model qualification. The broader Task 1-4 matrix remains open. + +Downstream preemption, async and P/D work should retain the ownership and +numerical gates before enabling each feature. Complete the frozen confirmation +and qualify feature-specific full-model parity before extending the final +concurrency/context campaign. Existing successful and failed results remain +separate from any new candidate or runtime. + +[Machine-readable scientific summary](../benchmarks/results/loopcd-c32/summary.json) +and [scale-validation protocol](loopcd-scale.md). diff --git a/docs/loopcd-scale.md b/docs/loopcd-scale.md new file mode 100644 index 00000000..6d86fe58 --- /dev/null +++ b/docs/loopcd-scale.md @@ -0,0 +1,81 @@ +# Ouro LoopCD scale qualification + +This continuation starts at PR #86 head +`d700ecca664da2e8f3122f66ad9d45de4295e45a`. Production model, engine and +feature guards are unchanged. Nonzero guidance still requires synchronous +LAST_EXITED execution. Running the experiment as a background job does not +enable asynchronous engine scheduling. + +## Independent continuation + +`tests/reference/continuation.py` retains independent K/V tensors for each +depth and layer. The prompt runs at P4; eight subsequent teacher-forced inputs +run at D2 or D3. Skipped depths receive only the newly computed positions from +the last executed depth. The P4 prompt history is retained. The final sampled +output is not forwarded and therefore does not create a KV position. + +The 12 FP32 CPU cases cover prompt lengths 3/4/5 around the block/chunk size4, +both decode depths, and guidance strengths0/.3. They compare recurrent state, +every historical K/V plane, vocabulary logits and selected log probability +with fixed atol=rtol=1e-4. They use a tiny initialized model, not official +checkpoint weights. There is no new BF16 tolerance or quality claim. + +## Closed-loop driver + +`python -m benchmarks.loopcd_scale --help` lists the implemented interface. +Each trial submits C outstanding requests, immediately refills completed slots +until 4C requests have completed, then drains. Request IDs are reused after +consuming the full output batch. Fixed output count, actual P/D depths, finish +reason, repeat token identity, empty references and full KV reclamation are +checked. Every request and trial is retained, including a failing trial. + +The driver reports scheduler B, configured/outstanding C, observed resident A, +and actual decode core rows S separately. A CPU capacity regression tests +B16/C32, B32/C64 and B64/C128 with both full B residency and an eight-request +physical KV limit. It asserts actual A and S reach the expected residency, +alongside queued/refill behavior and exact repeat outputs. + +Each process uses one model/engine, five warmups and five measured trials. +TTFT, TPOT, ITL, admission delay, completion latency, TPS/RPS, native work, +physical KV pool/peak, reference high water, CUDA allocated/reserved peaks and +Graph capture/replay/fallback deltas are recorded. Percentiles use nearest +rank. A cold capture during measured work invalidates the case; a fallback is +charged to elapsed time. Model loading, tokenization and diagnostics are +outside the native timing window. Prompts are deterministic tokenized repeated +text: these are fixed-work load tests, not task-quality or HTTP SLO estimates. + +Example, after resource admission and official input verification: + +```bash +python -m benchmarks.loopcd_scale \ + --model /absolute/path/to/pinned-ouro \ + --batch 32 --concurrency 64 --input-lengths 512 2048 \ + --output-tokens 128 --depth 4 --strength 0.3 --kv-gib 60 \ + --output /absolute/path/to/new-result +``` + +The prepared H20 queue covers B/C=1/1,32/32,32/64,64/64,64/128 and a32/32 +input2048 case. It compares P4D4 off/two_head in eager/Graph mode with two +fresh process starts and reversed configuration order. The fixed60GiB KV pool +targets the observed96GB H20 and is not suitable for the currently occupied +Thor. It is a planned pool, not an observed capacity or OOM result. + +No reduced-depth quality candidate has been selected. P4D2/P4D3 CLI options +support later separate qualification; the prepared timing queue does not +claim quality-preserving depth reduction. The earlier BF16 linear_fused +failure is unchanged; two_head remains the guidance baseline. + +## Actual validation status, 2026-10-06 + +- Local: independent continuation12 passed, capacity/refill6 passed, existing + LoopCD regression62 passed. Earlier three refill cases also passed before + tightening the capacity coverage and are not added to the final case count. +- Thor: original12 continuation plus3 refill cases passed in12.27s; CUDA was + uninitialized before/after. The stronger6-case capacity run also passed in18.20s with CUDA hidden. + Both controllers and test children exited naturally0. +- Official-weight high-concurrency GPU execution: **NOT_RUN**. The H20 lacks + the public loader/tokenizer dependencies; isolated-environment authorization + is pending. The inherited Thor Qwen job retains its original GPU lock. +- These additions are preparation/correctness evidence. They do not close + public loader qualification, official mixed-depth correctness, quality, + high-concurrency GPU performance or Task4. diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..234e77dc --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Repository-owned correctness fixtures and reference oracles.""" diff --git a/tests/reference/continuation.py b/tests/reference/continuation.py new file mode 100644 index 00000000..64ff1f13 --- /dev/null +++ b/tests/reference/continuation.py @@ -0,0 +1,104 @@ +"""Independent Ouro continuation with retained per-depth historical K/V.""" + +import torch +from torch.nn import functional as F + + +class OuroContinuation: + """Functional oracle; no production layers, cache, runner or guidance helper.""" + + def __init__(self, model): + self.config = model.config + self.weights = model.state_dict() + self.kv = {} + self.position = 0 + + @torch.inference_mode() + def forward(self, token_ids, depth): + c, w = self.config, self.weights + hidden = w["model.embed_tokens.weight"][token_ids] + length = len(token_ids) + positions = torch.arange(self.position, self.position + length, device=hidden.device) + frequency = c.rope_theta ** ( + -torch.arange(0, c.head_dim, 2, device=hidden.device).float() / c.head_dim + ) + angles = positions.float()[:, None] * frequency[None, :] + cos = torch.cat([angles.cos(), angles.cos()], -1)[:, None].to(hidden.dtype) + sin = torch.cat([angles.sin(), angles.sin()], -1)[:, None].to(hidden.dtype) + future = ( + torch.arange(self.position + length, device=hidden.device)[None, :] + > (positions[:, None]) + ) + + def norm(x, name): + value = x.float() / (x.float().square().mean(-1, keepdim=True) + c.rms_norm_eps).sqrt() + return value.to(x.dtype) * w[name + ".weight"] + + def linear(x, name): + return x @ w[name + ".weight"].T + + def rotary(x): + split = c.head_dim // 2 + return x * cos + torch.cat([-x[..., split:], x[..., :split]], -1) * sin + + outputs = [] + for loop in range(depth): + for layer in range(c.num_hidden_layers): + prefix = f"model.layers.{layer}" + x = norm(hidden, prefix + ".input_layernorm") + q = rotary( + linear(x, prefix + ".self_attn.q_proj").reshape( + length, c.num_attention_heads, c.head_dim + ) + ) + k = rotary( + linear(x, prefix + ".self_attn.k_proj").reshape( + length, c.num_key_value_heads, c.head_dim + ) + ) + v = linear(x, prefix + ".self_attn.v_proj").reshape( + length, c.num_key_value_heads, c.head_dim + ) + previous = self.kv.get((loop, layer)) + keys, values = ( + (k, v) + if previous is None + else (torch.cat([previous[0], k]), torch.cat([previous[1], v])) + ) + self.kv[loop, layer] = (keys, values) + groups = c.num_attention_heads // c.num_key_value_heads + scores = torch.einsum("thd,shd->hts", q, keys.repeat_interleave(groups, 1)) + scores = scores / c.head_dim**0.5 + probabilities = scores.masked_fill(future, -torch.inf).float().softmax(-1) + attended = torch.einsum( + "hts,shd->thd", + probabilities.to(hidden.dtype), + values.repeat_interleave(groups, 1), + ).reshape(length, -1) + hidden = hidden + norm( + linear(attended, prefix + ".self_attn.o_proj"), prefix + ".input_layernorm_2" + ) + x = norm(hidden, prefix + ".post_attention_layernorm") + mlp = F.silu(linear(x, prefix + ".mlp.gate_proj")) * linear( + x, prefix + ".mlp.up_proj" + ) + hidden = hidden + norm( + linear(mlp, prefix + ".mlp.down_proj"), prefix + ".post_attention_layernorm_2" + ) + hidden = norm(hidden, "model.norm") + outputs.append((hidden.clone(), linear(hidden, "lm_head"))) + # Only newly computed positions are filled. Prompt P4 history survives D2/D3. + for loop in range(depth, c.total_ut_steps): + for layer in range(c.num_hidden_layers): + k, v = self.kv[depth - 1, layer] + previous = self.kv.get((loop, layer)) + self.kv[loop, layer] = ( + (k[-length:].clone(), v[-length:].clone()) + if previous is None + else ( + torch.cat([previous[0], k[-length:]]), + torch.cat([previous[1], v[-length:]]), + ) + ) + self.position += length + return outputs diff --git a/tests/test_loopcd.py b/tests/test_loopcd.py new file mode 100644 index 00000000..48b1dd06 --- /dev/null +++ b/tests/test_loopcd.py @@ -0,0 +1,288 @@ +"""Independent readout oracles and request-lifecycle checks for Ouro LoopCD.""" + +from dataclasses import replace + +import pytest +import torch +from torch.nn import functional as F + +from tests.helpers import tiny_ouro_config +from tests.reference import dense_reference +from vllm_rlt import CacheConfig, ExecutionConfig, LoopCDParams, SamplingParams, SchedulerConfig +from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.models import OuroForCausalLM +from vllm_rlt.request import Request, Stage +from vllm_rlt.worker.loopcd import Reference, extrapolate, reserved_bytes + + +def model_for(): + torch.manual_seed(42) + return OuroForCausalLM(tiny_ouro_config()).eval() + + +def params_for(**kwargs): + return SamplingParams( + min_loops=4, max_loops=4, max_tokens=3, ignore_eos=True, loopcd=LoopCDParams(), **kwargs + ) + + +def engine_for(model, chunk=4, **kwargs): + return LLMEngine( + model, + cache_config=CacheConfig(num_blocks=256, block_size=2), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=8, prefill_chunk_size=chunk + ), + execution_config=ExecutionConfig(loopcd=True), + **kwargs, + ) + + +@pytest.mark.parametrize("dtype,tolerance", [(torch.float64, 3e-14), (torch.float32, 8e-6)]) +@pytest.mark.parametrize("strength", [0, 0.3, 0.5, 2.0]) +def test_affine_logits_folding_against_two_projections(dtype, tolerance, strength): + generator = torch.Generator().manual_seed(17) + final, reference = [torch.randn(8, 64, generator=generator, dtype=dtype) for _ in range(2)] + weight = torch.randn(257, 64, generator=generator, dtype=dtype) + bias = torch.randn(257, generator=generator, dtype=dtype) + expected = (1 + strength) * F.linear(final, weight, bias) - strength * F.linear( + reference, weight, bias + ) + actual = F.linear(extrapolate(final, reference, strength), weight, bias) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize( + "values", + [ + {"reference_loop": 0}, + {"reference_loop": True}, + {"reference_loop": None}, + {"prefill_reference_loop": 0}, + {"strength": -1}, + {"strength": float("nan")}, + {"strength": float("inf")}, + {"strength_mode": "margin"}, + {"mode": "hidden"}, + {"implementation": "unknown"}, + ], +) +def test_invalid_config(values): + with pytest.raises(ValueError): + LoopCDParams(**values) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_zero_strength_matches_unguided_readout_without_references(enabled): + model = model_for() + results = [] + for cd in (None, LoopCDParams(strength=0)): + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=256, block_size=2), + execution_config=ExecutionConfig(loopcd=enabled), + ) + engine.add_request("a", [4, 7, 3], replace(params_for(), loopcd=cd)) + while engine.has_unfinished_requests(): + output = engine.step() + results.append(output[-1].token_ids) + assert engine.model_runner.loopcd_stats["captures"] == 0 + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + assert results[0] == results[1] + + +@pytest.mark.parametrize("prefill,decode", [(4, 4), (4, 2), (3, 2)]) +@pytest.mark.parametrize("implementation", ["two_head", "linear_fused"]) +@pytest.mark.parametrize("chunk", [1, 3, 8]) +@torch.inference_mode() +def test_first_readout_matches_independent_dense_oracle(prefill, decode, implementation, chunk): + model = model_for() + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=256, block_size=2), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=8, prefill_chunk_size=chunk + ), + execution_config=ExecutionConfig(loopcd=True, prefill_depth=prefill), + ) + prompt = torch.tensor([4, 7, 9, 3, 12]) + params = replace( + params_for(), + min_loops=decode, + max_loops=decode, + loopcd=LoopCDParams(prefill_reference_loop=2, implementation=implementation), + ) + depths = dense_reference(model, prompt, prefill) + expected = extrapolate(depths[-1][2][-1], depths[1][2][-1], 0.3) + engine.add_request("a", prompt.tolist(), params) + request = engine.scheduler.requests["a"] + while request.stage != Stage.CODA: + engine.step() + torch.testing.assert_close(request.hidden_state, depths[-1][0][-1], atol=4e-6, rtol=4e-5) + ref = engine.model_runner.loopcd_references["a"] + ref.check(request, 4, 0, 2) + raw = request.hidden_state.clone() + seen = [] + + def sample(logits, request): + seen.append(logits.clone()) + return logits.argmax() + + engine.model_runner._sample_tensor = sample + engine.step() + torch.testing.assert_close(seen[0], expected, atol=5e-6, rtol=5e-5) + torch.testing.assert_close(request.hidden_state, raw, atol=0, rtol=0) + while engine.has_unfinished_requests(): + output = engine.step() + assert output[-1].exit_depths == [prefill, decode, decode] + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + + +def test_refill_decode_abort_and_request_id_reuse(): + engine = engine_for(model_for(), 2) + params = params_for() + engine.add_request("a", [4, 7, 3], params) + engine.add_request("b", [9, 2], params) + while not engine.model_runner.loopcd_references: + engine.step() + old = engine.scheduler.requests["a"] + engine.abort_request("a") + engine.add_request("a", [12, 3, 7, 9], params) + assert engine.scheduler.requests["a"] is not old + completed = {} + while engine.has_unfinished_requests(): + for output in engine.step(): + for ref in engine.model_runner.loopcd_references.values(): + assert ref.owner is engine.scheduler.requests[ref.owner.request_id] + if output.finished: + completed[output.request_id] = output + assert set(completed) == {"a", "b"} + assert all(x.exit_depths == [4, 4, 4] for x in completed.values()) + assert engine.model_runner.loopcd_stats["guided_rows"] > 0 + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + + +@pytest.mark.parametrize("field", ["owner", "position", "output_index", "loop"]) +def test_stale_reference_rejected(field): + request = Request("a", [1, 2], SamplingParams()) + ref = Reference(request, 1, 0, 1, torch.ones(3)) + ref.check(request, 1, 0, 1) + value = Request("a", [1, 2], SamplingParams()) if field == "owner" else 5 + with pytest.raises(RuntimeError, match="mismatch"): + replace(ref, **{field: value}).check(request, 1, 0, 1) + + +@pytest.mark.parametrize("case", ["unreserved", "reference", "early_exit", "prefix"]) +def test_unsupported_admission_rejected_before_request_allocation(case): + engine = LLMEngine( + model_for(), + cache_config=CacheConfig( + num_blocks=256, block_size=2, enable_prefix_caching=case == "prefix" + ), + execution_config=ExecutionConfig(loopcd=case != "unreserved"), + ) + params = params_for() + if case == "reference": + params = replace(params, loopcd=replace(params.loopcd, reference_loop=4)) + elif case == "early_exit": + params = replace(params, min_loops=2) + with pytest.raises(ValueError, match="LoopCD"): + engine.add_request("a", [4, 7], params) + assert engine.cache_manager.num_used_blocks == 0 + assert not engine.scheduler.requests + + +@pytest.mark.parametrize("feature", ["async", "prefix", "shared"]) +def test_shallow_prefill_capability_checked_before_allocation(feature): + with pytest.raises(ValueError, match="prefill_depth override"): + LLMEngine( + model_for(), + execution_config=ExecutionConfig(prefill_depth=2, async_scheduling=feature == "async"), + cache_config=CacheConfig( + enable_prefix_caching=feature == "prefix", + layout="shared" if feature == "shared" else "last_exited", + ), + ) + + +def test_reference_and_two_head_scratch_budgeted_before_auto_kv(): + model = model_for() + scheduler = SchedulerConfig(max_num_seqs=4, max_num_batched_tokens=8) + assert reserved_bytes(model.config, scheduler, ExecutionConfig(), 2) == 0 + assert reserved_bytes(model.config, scheduler, ExecutionConfig(loopcd=True), 2) == ( + (4 + 3 * 8) * model.config.hidden_size * 2 + 3 * 8 * model.config.vocab_size * 4 + ) + + +@pytest.mark.parametrize("depth", [2, 4]) +@pytest.mark.parametrize( + "strength,implementation", + [ + (0.0, "two_head"), + (0.3, "two_head"), + (0.3, "linear_fused"), + (0.5, "two_head"), + (0.5, "linear_fused"), + ], +) +@torch.inference_mode() +def test_teacher_forced_continuation_against_dense_oracle(depth, strength, implementation): + model = model_for() + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=64, block_size=2), + scheduler_config=SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=16, prefill_chunk_size=3 + ), + execution_config=ExecutionConfig(loopcd=True, prefill_depth=depth), + ) + sampling = replace( + params_for(), + min_loops=depth, + max_loops=depth, + loopcd=LoopCDParams(strength=strength, implementation=implementation), + ) + prompt, candidate = [4, 7, 9, 3, 12], [17, 8, 2] + seen, owners = [], [] + + def observe(logits, request): + index = len(request.generated_token_ids) + if index == 0: + owners.append(request) + prefix = torch.tensor(prompt + request.generated_token_ids) + dense = dense_reference(model, prefix, depth) + expected = extrapolate(dense[-1][2][-1], dense[0][2][-1], strength) + torch.testing.assert_close(request.hidden_state, dense[-1][0][-1], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(logits, expected, atol=1e-5, rtol=1e-5) + selected = candidate[index] + torch.testing.assert_close( + logits.float().log_softmax(-1)[selected], + expected.float().log_softmax(-1)[selected], + atol=1e-5, + rtol=1e-5, + ) + assert logits.argmax() == expected.argmax() + seen.append(index) + return torch.tensor(selected) + + engine.model_runner._sample_tensor = observe + for run in range(2): + if run: + engine.add_request("a", [12, 3, 7], sampling) + interrupted = engine.scheduler.requests["a"] + while interrupted.stage != Stage.CODA: + engine.step() + engine.abort_request("a") + engine.add_request("a", prompt, sampling) + while engine.has_unfinished_requests(): + output = engine.step() + assert output[-1].token_ids == candidate + assert output[-1].exit_depths == [depth] * 3 + assert output[-1].finish_reason == "length" + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + assert seen == [0, 1, 2] * 2 and owners[0] is not owners[1] + assert not torch.cuda.is_initialized() diff --git a/tests/test_loopcd_continuation.py b/tests/test_loopcd_continuation.py new file mode 100644 index 00000000..0410c240 --- /dev/null +++ b/tests/test_loopcd_continuation.py @@ -0,0 +1,83 @@ +"""Mixed-depth teacher forcing checks state, guided readout and every KV plane.""" + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from tests.reference.continuation import OuroContinuation +from vllm_rlt import CacheConfig, ExecutionConfig, LoopCDParams, SamplingParams, SchedulerConfig +from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.models import OuroForCausalLM +from vllm_rlt.request import Stage + + +@pytest.mark.parametrize("depth", [2, 3]) +@pytest.mark.parametrize("length", [3, 4, 5]) +@pytest.mark.parametrize("strength", [0.0, 0.3]) +@torch.inference_mode() +def test_retained_p4_history(depth, length, strength): + torch.manual_seed(1729) + model = OuroForCausalLM(tiny_ouro_config()).eval() + reference = OuroContinuation(model) + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=128, block_size=4), + scheduler_config=SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=8, prefill_chunk_size=4 + ), + execution_config=ExecutionConfig(loopcd=True, prefill_depth=4), + ) + prompt, continuation = [4, 7, 9, 3, 12][:length], [17, 8, 2, 11, 5, 7, 13, 6, 9] + sampling = SamplingParams( + max_tokens=len(continuation), + min_loops=depth, + max_loops=depth, + ignore_eos=True, + loopcd=LoopCDParams(strength=strength), + ) + seen = [] + + def sample(logits, request): + seen.append(logits.detach().clone()) + return torch.tensor(continuation[len(seen) - 1], device=logits.device) + + engine.model_runner._sample_tensor = sample + engine.add_request("reuse", prompt, sampling) + for index, token in enumerate(continuation): + expected = reference.forward( + torch.tensor(prompt if index == 0 else [continuation[index - 1]]), + 4 if index == 0 else depth, + ) + request = engine.scheduler.requests["reuse"] + while request.stage != Stage.CODA: + engine.step() + torch.testing.assert_close(request.hidden_state, expected[-1][0][-1], atol=1e-4, rtol=1e-4) + cache = engine.cache_manager + allocation = cache._allocations["reuse"] + for (loop, layer), (keys, values) in reference.kv.items(): + for position in range(reference.position): + block = allocation.block_tables[loop][position // cache.block_size] + offset = position % cache.block_size + torch.testing.assert_close( + cache.key_cache[block, layer, offset], keys[position], atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + cache.value_cache[block, layer, offset], values[position], atol=1e-4, rtol=1e-4 + ) + expected_logits = expected[-1][1][-1].float() + expected_logits = expected_logits + strength * ( + expected_logits - expected[0][1][-1].float() + ) + engine.step() + torch.testing.assert_close(seen[-1], expected_logits, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + seen[-1].float().log_softmax(-1)[token], + expected_logits.log_softmax(-1)[token], + atol=1e-4, + rtol=1e-4, + ) + assert not engine.has_unfinished_requests() + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + assert reference.position == length + len(continuation) - 1 + engine.close() diff --git a/tests/test_loopcd_preemption.py b/tests/test_loopcd_preemption.py new file mode 100644 index 00000000..a15b4e6a --- /dev/null +++ b/tests/test_loopcd_preemption.py @@ -0,0 +1,342 @@ +"""LoopCD state preservation through synchronous Ouro preemption.""" + +from dataclasses import replace + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + LoopCDParams, + SamplingParams, + SchedulerConfig, +) +from vllm_rlt.config import ExitConfig +from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.models import OuroForCausalLM +from vllm_rlt.request import Request, Stage +from vllm_rlt.worker.loopcd import validate + +PROMPT = [4, 7, 9, 3, 12] +BOUNDARIES = [ + ("partial_prefill", False), + ("first_coda", True), + ("next_prelude", False), + ("before_reference", False), + ("at_reference", True), + ("after_reference", True), + ("decode_coda", True), +] + + +def model_for(): + torch.manual_seed(42) + return OuroForCausalLM(tiny_ouro_config()).eval() + + +def gpu_model_for(): + torch.manual_seed(42) + return OuroForCausalLM(tiny_ouro_config(head_dim=64)).to(device="cuda").eval() + + +def params_for(strength=0.3, decode=4, implementation="two_head"): + return SamplingParams( + min_loops=decode, + max_loops=decode, + max_tokens=4, + ignore_eos=True, + temperature=0.8, + seed=17, + loopcd=LoopCDParams( + reference_loop=1 if decode == 2 else 2, + prefill_reference_loop=2, + strength=strength, + implementation=implementation, + ), + ) + + +def engine_for(model, mode="refill", attention_backend="torch", static_buffers=False): + return LLMEngine( + model, + attention_backend=attention_backend, + cache_config=CacheConfig(num_blocks=128, block_size=2), + scheduler_config=SchedulerConfig( + max_num_seqs=4, + max_num_batched_tokens=8, + prefill_chunk_size=2, + enable_preemption=True, + mode=mode, + ), + execution_config=ExecutionConfig( + loopcd=True, prefill_depth=4, static_buffers=static_buffers + ), + ) + + +def at_boundary(request, boundary): + if boundary == "partial_prefill": + return request.stage == Stage.PREFILL and request.num_prefilled_tokens == 2 + if boundary == "first_coda": + return request.stage == Stage.CODA and request.num_scheduled_outputs == 0 + if request.num_scheduled_outputs != 1: + return False + if boundary == "next_prelude": + return request.stage == Stage.PRELUDE + if boundary == "decode_coda": + return request.stage == Stage.CODA + reference_loop = request.sampling_params.loopcd.reference_loop + depth = { + "before_reference": reference_loop - 1, + "at_reference": reference_loop, + "after_reference": reference_loop + 1, + }[boundary] + return request.stage == Stage.RECURRENT and request.loops_done == depth + + +def advance_to(engine, boundary): + request = engine.scheduler.requests["a"] + for _ in range(100): + if at_boundary(request, boundary): + return request + engine.step() + pytest.fail(f"boundary not reached: {boundary}") + + +def suspend(engine): + # Exercise the existing victim-selection and WAITING/resume path at a + # completed scheduler boundary, without changing subsequent batch shape. + requester = Request("pressure", [2], SamplingParams()) + assert engine.preemption.preempt(requester) + request = engine.scheduler.requests["a"] + assert request.stage == Stage.WAITING + assert request.hidden_state is None + assert "a" not in engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + return engine.preemption.snapshots["a"] + + +def drain(engine): + result = None + for _ in range(100): + if not engine.has_unfinished_requests(): + assert result is not None + assert not engine.preemption.snapshots + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + return result + for output in engine.step(): + if output.finished: + result = output + pytest.fail("engine did not drain") + + +@pytest.mark.parametrize("boundary,has_reference", BOUNDARIES) +@pytest.mark.parametrize("mode", ["refill", "no_refill"]) +@pytest.mark.parametrize("implementation", ["two_head", "linear_fused"]) +@torch.inference_mode() +def test_suspend_resume_preserves_all_logits_rng_and_recurrent_work( + boundary, has_reference, mode, implementation +): + model = model_for() + runs = [] + for preempt in (False, True): + engine = engine_for(model, mode) + seen = [] + sample = engine.model_runner._sample_tensor + + def observe(logits, request): + seen.append(logits.clone()) + return sample(logits, request) + + engine.model_runner._sample_tensor = observe + engine.add_request("a", PROMPT, params_for(implementation=implementation)) + request = advance_to(engine, boundary) + if preempt: + old = engine.model_runner.loopcd_references.get("a") + state = suspend(engine) + reference = state["loopcd_reference"] + assert (reference is not None) == has_reference + if reference is not None: + assert reference.owner is request + assert reference.hidden.device.type == "cpu" + assert reference.hidden.data_ptr() != old.hidden.data_ptr() + torch.testing.assert_close(reference.hidden, old.hidden, atol=0, rtol=0) + engine.step() # Native scheduler callback resumes the saved stage. + assert engine.preemption.resumptions == 1 + assert not engine.preemption.snapshots + output = drain(engine) + runs.append((output, torch.stack(seen), engine.model_runner.loopcd_stats)) + assert runs[0][0].token_ids == runs[1][0].token_ids + assert runs[0][0].exit_depths == runs[1][0].exit_depths == [4, 4, 4, 4] + torch.testing.assert_close(runs[0][1], runs[1][1], atol=0, rtol=0) + for key in ("captures", "prefill_core_rows", "decode_core_rows", "guided_rows"): + assert runs[0][2][key] == runs[1][2][key] + + +@pytest.mark.parametrize("strength", [0.0, 0.3]) +@torch.inference_mode() +def test_partial_decode_depth_and_repeated_preemption(strength): + model = model_for() + outputs = [] + for preempt in (False, True): + engine = engine_for(model) + engine.add_request("a", PROMPT, params_for(strength=strength, decode=2)) + if preempt: + for boundary in ("partial_prefill", "first_coda", "at_reference", "decode_coda"): + advance_to(engine, boundary) + state = suspend(engine) + if strength == 0: + assert state["loopcd_reference"] is None + engine.step() + assert engine.preemption.preemptions == engine.preemption.resumptions == 4 + output = drain(engine) + outputs.append(output) + assert outputs[0].token_ids == outputs[1].token_ids + assert outputs[0].exit_depths == outputs[1].exit_depths == [4, 2, 2, 2] + + +@torch.inference_mode() +def test_blocked_restore_keeps_cpu_reference_until_capacity_returns(monkeypatch): + engine = engine_for(model_for()) + engine.add_request("a", PROMPT, params_for()) + request = advance_to(engine, "after_reference") + state = suspend(engine) + allocate = engine.cache_manager.allocate + with monkeypatch.context() as patch: + patch.setattr(engine.cache_manager, "allocate", lambda *args, **kwargs: False) + assert engine.preemption.resume(request) is False + assert engine.preemption.snapshots["a"] is state + assert state["loopcd_reference"].owner is request + assert "a" not in engine.model_runner.loopcd_references + assert request.hidden_state is None + assert engine.cache_manager.num_used_blocks == 0 + assert engine.cache_manager.allocate == allocate + drain(engine) + assert engine.preemption.resumptions == 1 + + +@pytest.mark.parametrize("field", ["owner", "position", "output_index", "loop", "missing"]) +@torch.inference_mode() +def test_corrupt_suspended_reference_rejected_before_kv_allocation(field): + engine = engine_for(model_for()) + engine.add_request("a", PROMPT, params_for()) + request = advance_to(engine, "after_reference") + state = suspend(engine) + reference = state["loopcd_reference"] + value = Request("a", PROMPT, params_for()) if field == "owner" else 999 + state["loopcd_reference"] = None if field == "missing" else replace(reference, **{field: value}) + with pytest.raises(RuntimeError, match="LoopCD reference"): + engine.preemption.resume(request) + assert engine.cache_manager.num_used_blocks == 0 + assert "a" not in engine.model_runner.loopcd_references + engine.abort_request("a") + assert not engine.preemption.snapshots + assert not engine.has_unfinished_requests() + + +@torch.inference_mode() +def test_abort_suspended_request_and_id_reuse_cannot_inherit_reference(): + model = model_for() + engine = engine_for(model) + engine.add_request("a", PROMPT, params_for()) + old = advance_to(engine, "after_reference") + suspend(engine) + engine.abort_request("a") + assert not engine.preemption.snapshots + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + engine.add_request("a", [12, 3, 7], params_for()) + assert engine.scheduler.requests["a"] is not old + reused = drain(engine) + baseline = engine_for(model) + baseline.add_request("a", [12, 3, 7], params_for()) + expected = drain(baseline) + assert reused.token_ids == expected.token_ids + assert reused.exit_depths == expected.exit_depths + + +def test_graph_preemption_stays_rejected_until_separate_qualification(): + with pytest.raises(ValueError, match="preemption currently requires eager"): + validate( + model_for(), + params_for(), + ExecutionConfig(loopcd=True, cuda_graphs=True), + CacheConfig(), + SchedulerConfig(enable_preemption=True), + ExitConfig(), + None, + ) + + +DEVICE_CASES = [ + pytest.param("cpu", False, id="cpu"), + pytest.param("cuda", False, marks=pytest.mark.gpu, id="cuda-static-off"), + pytest.param("cuda", True, marks=pytest.mark.gpu, id="cuda-static-on"), +] + + +@pytest.mark.parametrize("device,static_buffers", DEVICE_CASES) +@pytest.mark.parametrize("policy", ["priority", "pressure"]) +@torch.inference_mode() +def test_guided_requests_survive_scheduler_driven_preemption(policy, device, static_buffers): + if device == "cpu": + model = model_for() + attention_backend = "torch" + else: + model = gpu_model_for() + attention_backend = "triton" + engine = LLMEngine( + model, + attention_backend=attention_backend, + cache_config=CacheConfig( + num_blocks=64 if policy == "priority" else 24, + block_size=2, + incremental_allocation=policy == "pressure", + ), + scheduler_config=SchedulerConfig( + max_num_seqs=1 if policy == "priority" else 3, + max_num_batched_tokens=8, + prefill_chunk_size=2, + policy="priority" if policy == "priority" else "fcfs", + enable_preemption=True, + ), + execution_config=ExecutionConfig( + loopcd=True, prefill_depth=4, static_buffers=static_buffers + ), + ) + params = params_for() + if policy == "priority": + engine.add_request("a", PROMPT, replace(params, priority=10)) + advance_to(engine, "first_coda") + engine.add_request("b", [9, 2], replace(params, priority=-10)) + expected_ids = {"a", "b"} + else: + for name in ("a", "b", "c"): + engine.add_request(name, PROMPT, params) + expected_ids = {"a", "b", "c"} + completed = {} + for _ in range(1000): + if not engine.has_unfinished_requests(): + break + for output in engine.step(): + for reference in engine.model_runner.loopcd_references.values(): + assert reference.owner is engine.scheduler.requests[reference.owner.request_id] + if output.finished: + completed[output.request_id] = output + assert set(completed) == expected_ids + assert engine.preemption.preemptions > 0 + assert engine.preemption.preemptions == engine.preemption.resumptions + assert not engine.preemption.snapshots + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + if policy == "priority": + assert list(completed) == ["b", "a"] + baseline = engine_for(model, attention_backend=attention_backend, static_buffers=static_buffers) + baseline.add_request("a", PROMPT, params) + expected = drain(baseline) + for name in ("a",) if policy == "priority" else ("a", "b", "c"): + assert completed[name].token_ids == expected.token_ids + assert completed[name].exit_depths == expected.exit_depths diff --git a/tests/test_loopcd_preemption_continuation.py b/tests/test_loopcd_preemption_continuation.py new file mode 100644 index 00000000..acb39093 --- /dev/null +++ b/tests/test_loopcd_preemption_continuation.py @@ -0,0 +1,299 @@ +"""Mixed-depth oracle checks across preemption and request-ID reuse.""" + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from tests.reference.continuation import OuroContinuation +from tests.test_loopcd_preemption import suspend +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + LoopCDParams, + SamplingParams, + SchedulerConfig, +) +from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.models import OuroForCausalLM +from vllm_rlt.request import Stage + +PROMPT = [4, 7, 9, 3, 12] +CONTINUATION = [17, 8, 2, 11, 5, 7, 13, 6, 9] +REUSE_PROMPT = [2, 5, 8, 1] +REUSE_TOKENS = [11, 13] +PREFILL_DEPTH = 4 +ATOL = RTOL = 1e-4 + + +def _engine(model, attention_backend): + return LLMEngine( + model, + attention_backend=attention_backend, + cache_config=CacheConfig(num_blocks=128, block_size=2), + scheduler_config=SchedulerConfig( + max_num_seqs=4, + max_num_batched_tokens=8, + prefill_chunk_size=2, + enable_preemption=True, + ), + execution_config=ExecutionConfig(loopcd=True, prefill_depth=PREFILL_DEPTH), + ) + + +def _params(depth, strength, max_tokens): + return SamplingParams( + max_tokens=max_tokens, + min_loops=depth, + max_loops=depth, + ignore_eos=True, + loopcd=LoopCDParams(strength=strength), + ) + + +def _plane_error(cache, oracle, request_id): + """One gather per plane, not one sync per position.""" + device = cache.device + positions = torch.arange(oracle.position, device=device) + worst = 0.0 + for loop in range(cache.max_loops): + table = torch.tensor( + cache.get_block_table(request_id, loop), device=device, dtype=torch.long + ) + blocks = table[positions // cache.block_size] + offsets = positions % cache.block_size + keys = cache.key_cache[blocks, :, offsets].permute(1, 0, 2, 3).float().cpu() + values = cache.value_cache[blocks, :, offsets].permute(1, 0, 2, 3).float().cpu() + expected_keys = ( + torch.stack([oracle.kv[(loop, layer)][0] for layer in range(cache.num_layers)]) + .float() + .cpu() + ) + expected_values = ( + torch.stack([oracle.kv[(loop, layer)][1] for layer in range(cache.num_layers)]) + .float() + .cpu() + ) + worst = max( + worst, + (keys - expected_keys).abs().max().item(), + (values - expected_values).abs().max().item(), + ) + torch.testing.assert_close(keys, expected_keys, atol=ATOL, rtol=RTOL) + torch.testing.assert_close(values, expected_values, atol=ATOL, rtol=RTOL) + return worst + + +def _check(engine, oracle, logits, request, expected, forced, strength, worst): + """Assert independent state/KV/logits/logprob, recording maxima first.""" + state = request.hidden_state.float().cpu() + final_hidden = expected[-1][0][-1].float().cpu() + worst["state"] = max(worst["state"], (state - final_hidden).abs().max().item()) + torch.testing.assert_close(state, final_hidden, atol=ATOL, rtol=RTOL) + worst["kv"] = max(worst["kv"], _plane_error(engine.cache_manager, oracle, request.request_id)) + if strength != 0: + reference = engine.model_runner.loopcd_references.get(request.request_id) + assert reference is not None + assert reference.owner is request + # The oracle gates the resulting readout; reference copy/restore is lossless. + worst["reference"] = max( + worst["reference"], + (reference.hidden.float().cpu() - expected[0][0][-1].float().cpu()).abs().max().item(), + ) + final_logits = expected[-1][1][-1].float().cpu() + guided = final_logits + strength * (final_logits - expected[0][1][-1].float().cpu()) + row = logits.float().cpu() + worst["logits"] = max(worst["logits"], (row - guided).abs().max().item()) + torch.testing.assert_close(row, guided, atol=ATOL, rtol=RTOL) + chosen = row.log_softmax(-1)[forced] + oracle_logp = guided.log_softmax(-1)[forced] + worst["logprob"] = max(worst["logprob"], (chosen - oracle_logp).abs().item()) + torch.testing.assert_close(chosen, oracle_logp, atol=ATOL, rtol=RTOL) + return torch.tensor(forced, device=logits.device) + + +def _expected(oracle, prompt, forced, index, depth, device): + tokens = prompt if index == 0 else [forced[index - 1]] + feed_depth = PREFILL_DEPTH if index == 0 else depth + return oracle.forward(torch.tensor(tokens, device=device), feed_depth) + + +@torch.inference_mode() +def check_preempted_continuation(model, depth, length, strength=0.3, attention_backend="torch"): + """Check independent hidden/KV/readout gates across native resumption. + + Reference-hidden error is diagnostic; snapshot copies are checked bit-exactly. + Checks run at sampling because resumed CODA may finish in the same step. + """ + prompt = PROMPT[:length] + device = next(model.parameters()).device + engine = _engine(model, attention_backend) + worst = {"state": 0.0, "kv": 0.0, "logits": 0.0, "logprob": 0.0, "reference": 0.0} + result = None + try: + oracle = OuroContinuation(model) + expected, seen = [], [] + stops = {"first_coda": False, "next_prelude": False, "recurrent": False} + + def main_sample(logits, request): + index = len(seen) + sampled = _check( + engine, + oracle, + logits, + request, + expected[index], + CONTINUATION[index], + strength, + worst, + ) + seen.append(logits.detach().clone()) + return sampled + + engine.model_runner._sample_tensor = main_sample + engine.add_request("a", prompt, _params(depth, strength, len(CONTINUATION))) + completed = None + for _ in range(400): + if not engine.has_unfinished_requests(): + break + while len(expected) <= len(seen): + index = len(expected) + expected.append(_expected(oracle, prompt, CONTINUATION, index, depth, device)) + request = engine.scheduler.requests.get("a") + if ( + request is not None + and not stops["first_coda"] + and (request.stage == Stage.CODA and request.num_scheduled_outputs == 0) + ): + live = engine.model_runner.loopcd_references.get("a") + if strength != 0: + assert live is not None and live.owner is request + + state = suspend(engine) + if strength != 0: + torch.testing.assert_close( + state["loopcd_reference"].hidden, live.hidden.cpu(), atol=0, rtol=0 + ) + blocker = ( + engine.cache_manager.num_blocks // engine.cache_manager.storage_depths + ) * engine.cache_manager.block_size + assert engine.cache_manager.allocate("blocker", blocker) + assert engine.preemption.resume(request) is False + assert engine.preemption.snapshots["a"] is state + reference = state["loopcd_reference"] + assert reference is None if strength == 0 else reference.owner is request + assert not engine.cache_manager.has_allocation("a") + assert request.hidden_state is None + assert "a" not in engine.model_runner.loopcd_references + engine.cache_manager.free("blocker") + stops["first_coda"] = True + elif ( + request is not None + and not stops["next_prelude"] + and (request.stage == Stage.PRELUDE and request.num_scheduled_outputs == 1) + ): + suspend(engine) + stops["next_prelude"] = True + elif ( + request is not None + and not stops["recurrent"] + and (request.stage == Stage.RECURRENT and request.loops_done == 1) + ): + suspend(engine) + stops["recurrent"] = True + for output in engine.step(): + if output.finished: + completed = output + + assert completed is not None + assert completed.token_ids == CONTINUATION + assert completed.exit_depths == [PREFILL_DEPTH] + [depth] * (len(CONTINUATION) - 1) + assert len(seen) == len(CONTINUATION) + assert all(stops.values()) + assert not engine.preemption.snapshots + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + + engine.model_runner._sample_tensor = lambda logits, request: torch.tensor( + 0, device=logits.device + ) + engine.add_request("a", prompt, _params(depth, strength, len(REUSE_TOKENS))) + old_request = engine.scheduler.requests["a"] + for _ in range(200): + current = engine.scheduler.requests["a"] + if current.stage == Stage.CODA and current.num_scheduled_outputs == 0: + break + engine.step() + assert current.stage == Stage.CODA and current.num_scheduled_outputs == 0 + suspend(engine) + engine.abort_request("a") + assert not engine.preemption.snapshots + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + + reuse_oracle = OuroContinuation(model) + reuse_expected, reuse_seen = [], [] + + def reuse_sample(logits, request): + index = len(reuse_seen) + sampled = _check( + engine, + reuse_oracle, + logits, + request, + reuse_expected[index], + REUSE_TOKENS[index], + strength, + worst, + ) + reuse_seen.append(logits.detach().clone()) + return sampled + + engine.model_runner._sample_tensor = reuse_sample + engine.add_request("a", REUSE_PROMPT, _params(depth, strength, len(REUSE_TOKENS))) + assert engine.scheduler.requests["a"] is not old_request + reused = None + for _ in range(200): + if not engine.has_unfinished_requests(): + break + while len(reuse_expected) <= len(reuse_seen): + index = len(reuse_expected) + reuse_expected.append( + _expected(reuse_oracle, REUSE_PROMPT, REUSE_TOKENS, index, depth, device) + ) + for output in engine.step(): + if output.finished: + reused = output + + assert reused is not None + assert reused.token_ids == REUSE_TOKENS + assert reused.exit_depths == [PREFILL_DEPTH, depth] + assert len(reuse_seen) == len(REUSE_TOKENS) + assert not engine.preemption.snapshots + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + result = { + "outputs": len(seen), + "reuse_outputs": len(reuse_seen), + "exit_depths": completed.exit_depths, + "preemptions": engine.preemption.preemptions, + "resumptions": engine.preemption.resumptions, + **{f"{key}_max_error": value for key, value in worst.items()}, + } + finally: + engine.close() + return result + + +@pytest.mark.parametrize("depth", [2, 3]) +@pytest.mark.parametrize("length", [3, 5]) +@torch.inference_mode() +def test_preempted_continuation_cpu(depth, length): + torch.manual_seed(1729) + model = OuroForCausalLM(tiny_ouro_config()).eval() + result = check_preempted_continuation( + model, depth, length, strength=0.3, attention_backend="torch" + ) + assert result["outputs"] == len(CONTINUATION) + assert result["reuse_outputs"] == len(REUSE_TOKENS) + assert result["preemptions"] == 4 + assert result["resumptions"] == 3 diff --git a/tests/test_loopcd_scale.py b/tests/test_loopcd_scale.py new file mode 100644 index 00000000..0046dd94 --- /dev/null +++ b/tests/test_loopcd_scale.py @@ -0,0 +1,36 @@ +import pytest +import torch + +from benchmarks.loopcd_scale import signature, trial +from tests.helpers import tiny_ouro_config +from vllm_rlt import CacheConfig, ExecutionConfig, LoopCDParams, SamplingParams, SchedulerConfig +from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.models import OuroForCausalLM + + +@pytest.mark.parametrize("batch,concurrency", [(16, 32), (32, 64), (64, 128)]) +@pytest.mark.parametrize("capacity_rows", [8, None]) +@torch.inference_mode() +def test_closed_loop_refill_drains_and_reuses_ids(batch, concurrency, capacity_rows): + torch.manual_seed(1729) + resident_limit = capacity_rows or batch + engine = LLMEngine( + OuroForCausalLM(tiny_ouro_config()).eval(), + cache_config=CacheConfig(num_blocks=8 * resident_limit, block_size=4), + scheduler_config=SchedulerConfig(max_num_seqs=batch, max_num_batched_tokens=batch), + execution_config=ExecutionConfig(loopcd=True, prefill_depth=4), + ) + prompts = [[4, 7, 3]] * (4 * concurrency) + sampling = SamplingParams( + max_tokens=3, min_loops=2, max_loops=2, ignore_eos=True, loopcd=LoopCDParams() + ) + first, second = [trial(engine, prompts, sampling, concurrency) for _ in range(2)] + assert signature(first) == signature(second) + assert len(first["requests"]) == 4 * concurrency + assert max(first["outstanding_histogram"]) == concurrency + assert max(first["resident_histogram"]) == resident_limit + assert max(first["decode_effective_rows_histogram"]) == resident_limit + assert first["drained"] and second["drained"] + assert first["kv_peak_blocks"] == 8 * resident_limit + assert first["graph"] == dict(captures=0, replays=0, fallbacks=0) + engine.close() diff --git a/tests/test_scheduler_selection.py b/tests/test_scheduler_selection.py index fccfe0f4..713eadf9 100644 --- a/tests/test_scheduler_selection.py +++ b/tests/test_scheduler_selection.py @@ -9,15 +9,16 @@ import pytest import torch +from tests.helpers import tiny_ouro_config from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig from vllm_rlt.engine.llm_engine import LLMEngine -from vllm_rlt.models import OuroConfig, OuroForCausalLM +from vllm_rlt.models import OuroForCausalLM from vllm_rlt.request import Stage def model(): torch.manual_seed(123) - return OuroForCausalLM(OuroConfig.tiny()) + return OuroForCausalLM(tiny_ouro_config()) def engine(max_num_seqs=2): diff --git a/vllm_rlt/__init__.py b/vllm_rlt/__init__.py index d0b346bc..1ca604d7 100644 --- a/vllm_rlt/__init__.py +++ b/vllm_rlt/__init__.py @@ -10,7 +10,7 @@ from vllm_rlt.entrypoints.llm import LLM from vllm_rlt.profiling import ProfileConfig from vllm_rlt.request import RequestOutput -from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.sampling_params import LoopCDParams, SamplingParams __all__ = [ "LLM", @@ -21,5 +21,6 @@ "SchedulerConfig", "SpeculativeConfig", "SamplingParams", + "LoopCDParams", "RequestOutput", ] diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 2491dcf0..4a78f6f6 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -108,9 +108,14 @@ class ExecutionConfig: cuda_graph_max_graphs: int = 16 cuda_graph_memory_reserve_bytes: int = 1024**3 prefill_uva: bool = False + loopcd: bool = False + prefill_depth: int | None = None def __post_init__(self): + if self.prefill_depth is not None: + _positive("prefill_depth", self.prefill_depth) for name in ( + "loopcd", "prefill_uva", "async_scheduling", "multi_stream", diff --git a/vllm_rlt/core/scheduling_policy.py b/vllm_rlt/core/scheduling_policy.py index 411a9625..d4117206 100644 --- a/vllm_rlt/core/scheduling_policy.py +++ b/vllm_rlt/core/scheduling_policy.py @@ -52,6 +52,13 @@ def schedule( scheduler._admit() if q[Stage.PREFILL]: return scheduler._take(Stage.PREFILL) + # Admission can restore a snapshot directly to a boundary stage. + if q[Stage.PRELUDE]: + return scheduler._take(Stage.PRELUDE) + if q[Stage.CODA] and ( + len(q[Stage.CODA]) >= self.config.min_coda_batch_size or not q[Stage.RECURRENT] + ): + return scheduler._take(Stage.CODA) if q[Stage.RECURRENT]: return scheduler._take(Stage.RECURRENT) return None diff --git a/vllm_rlt/engine/llm_engine.py b/vllm_rlt/engine/llm_engine.py index a7708960..aed15b83 100644 --- a/vllm_rlt/engine/llm_engine.py +++ b/vllm_rlt/engine/llm_engine.py @@ -9,6 +9,7 @@ from vllm_rlt.profiling import Profiler from vllm_rlt.request import FinishReason, Request, RequestOutput, Stage from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.worker.loopcd import reserved_bytes, validate from vllm_rlt.worker.model_runner import ModelRunner from vllm_rlt.worker.speculative import SpeculativeRunner @@ -35,6 +36,18 @@ def __init__( config = model.config self.exit_config = exit_config or ExitConfig() self.execution_config = execution_config or ExecutionConfig() + self.prefill_depth = self.execution_config.prefill_depth or config.total_ut_steps + if self.prefill_depth > config.total_ut_steps: + raise ValueError("prefill_depth exceeds the model's supported depth") + if self.execution_config.prefill_depth is not None and ( + cache_config.layout != "last_exited" + or cache_config.enable_prefix_caching + or self.execution_config.async_scheduling + or speculative_config is not None + ): + raise ValueError( + "prefill_depth override requires synchronous last_exited without prefix/speculation" + ) self.speculative_config = speculative_config if speculative_config is not None: if cache_config.layout != "last_exited": @@ -90,6 +103,9 @@ def __init__( self.execution_config, attention.info["selected_backend"], ) + self.memory_plan["loopcd_reserve_bytes"] = reserved_bytes( + config, scheduler_config, self.execution_config, parameter.element_size() + ) self.cache_manager = KVCacheManager( num_layers=config.num_hidden_layers, num_kv_heads=config.num_key_value_heads, @@ -153,6 +169,15 @@ def add_request( max_loops = params.max_loops or config.total_ut_steps if max_loops > config.total_ut_steps or params.min_loops > max_loops: raise ValueError("requested loop bounds exceed the model's supported depth") + validate( + self.model, + params, + self.execution_config, + self.cache_manager, + self.scheduler.config, + self.exit_config, + self.speculative_config, + ) if self.speculative_config is not None and ( max_loops != self.speculative_config.target_loops or params.exit_threshold != 1.0 ): @@ -170,7 +195,7 @@ def add_request( trace = self._exit_traces[key] if ( len(trace) < params.max_tokens - or trace[0] != config.total_ut_steps + or trace[0] != self.prefill_depth or any( type(d) is not int or not params.min_loops <= d <= max_loops for d in trace[1 : params.max_tokens] @@ -313,7 +338,7 @@ def _update(self, batch, result) -> list[RequestOutput]: event, ) if request.num_prefilled_tokens == len(request.prompt_token_ids): - request.loops_done = self.model.config.total_ut_steps + request.loops_done = self.prefill_depth self.scheduler.enqueue(request, Stage.CODA) else: self.scheduler.enqueue(request, Stage.PREFILL) diff --git a/vllm_rlt/engine/preemption.py b/vllm_rlt/engine/preemption.py index 3ece7f02..42ce0823 100644 --- a/vllm_rlt/engine/preemption.py +++ b/vllm_rlt/engine/preemption.py @@ -5,9 +5,12 @@ semantics. Snapshots are bounded by the admitted request population. """ +from dataclasses import replace + import torch from vllm_rlt.request import Stage +from vllm_rlt.worker.loopcd import preemption_reference class PreemptionManager: @@ -84,8 +87,14 @@ def preempt(self, requester, *, priority_only=False, excluded=frozenset()): e = self.engine e.model_runner.synchronize() cache = e.cache_manager + reference = preemption_reference( + victim, victim.stage, e.model_runner.loopcd_references.get(victim.request_id) + ) snapshot = dict( stage=victim.stage, + loopcd_reference=None + if reference is None + else replace(reference, hidden=reference.hidden.cpu().clone()), kv=cache.snapshot(victim.request_id), hidden=None if victim.hidden_state is None else victim.hidden_state.cpu().clone(), token=None @@ -115,6 +124,7 @@ def resume(self, request): if state is None: return None e, cache = self.engine, self.engine.cache_manager + reference = preemption_reference(request, state["stage"], state["loopcd_reference"]) kv = state["kv"] frontier = min(kv.max_tokens, kv.pages * cache.block_size) if not cache.allocate(request.request_id, kv.max_tokens, initial_tokens=frontier): @@ -124,6 +134,10 @@ def resume(self, request): request.input_token_tensor = ( None if state["token"] is None else state["token"].to(cache.device) ) + if reference is not None: + e.model_runner.loopcd_references[request.request_id] = replace( + reference, hidden=reference.hidden.to(cache.device) + ) # Restored buffers become visible to every runner stream before resumption. if cache.device.type == "cuda": torch.cuda.current_stream(cache.device).synchronize() diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index 59b5da94..991bae1a 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -301,6 +301,8 @@ def add_request(self, request_id, prompt_token_ids, sampling_params=None, *, tra if len(self.transfers) >= self.config.max_pending_requests: raise ValueError("PD pending-request limit reached") params = sampling_params or SamplingParams() + if params.loopcd is not None and params.loopcd.strength != 0: + raise ValueError("LoopCD P/D reference transfer is not implemented") cfg = self.model.config if not prompt_token_ids or any( type(t) is not int or not 0 <= t < cfg.vocab_size for t in prompt_token_ids diff --git a/vllm_rlt/sampling_params.py b/vllm_rlt/sampling_params.py index 30009bde..5a8e6e21 100644 --- a/vllm_rlt/sampling_params.py +++ b/vllm_rlt/sampling_params.py @@ -2,6 +2,38 @@ from dataclasses import dataclass +@dataclass(frozen=True) +class LoopCDParams: + """Fixed guidance after complete one-based recurrent updates.""" + + mode: str = "logits" + reference_loop: int = 1 + strength: float = 0.3 + implementation: str = "two_head" + strength_mode: str = "fixed" + prefill_reference_loop: int | None = None + + def __post_init__(self): + if self.mode != "logits": + raise ValueError("Ouro LoopCD supports logits mode only") + if self.strength_mode != "fixed": + raise ValueError("LoopCD supports fixed strength only") + if not math.isfinite(self.strength) or self.strength < 0: + raise ValueError("LoopCD strength must be finite and nonnegative") + for value in (self.reference_loop, self.prefill_reference_loop): + if value is not None and (type(value) is not int or value < 1): + raise ValueError("LoopCD reference_loop must be a positive integer") + if self.reference_loop is None: + raise ValueError("LoopCD reference_loop is required") + allowed = ("two_head", "linear_fused") + if self.implementation not in allowed: + raise ValueError("LoopCD implementation does not match guidance mode") + + @property + def prefill_loop(self): + return self.prefill_reference_loop or self.reference_loop + + @dataclass(frozen=True) class SamplingParams: max_tokens: int = 16 @@ -14,8 +46,11 @@ class SamplingParams: exit_threshold: float = 1.0 ignore_eos: bool = False priority: int = 0 + loopcd: LoopCDParams | None = None def __post_init__(self): + if self.loopcd is not None and not isinstance(self.loopcd, LoopCDParams): + raise ValueError("loopcd must be LoopCDParams or None") if type(self.priority) is not int: raise ValueError("priority must be an integer") for name in ("max_tokens", "min_loops", "max_loops"): diff --git a/vllm_rlt/worker/buffers.py b/vllm_rlt/worker/buffers.py index 167ca0a5..00b12f05 100644 --- a/vllm_rlt/worker/buffers.py +++ b/vllm_rlt/worker/buffers.py @@ -5,6 +5,7 @@ import torch from vllm_rlt.core.kv_cache_manager import _PreparedKVBatch +from vllm_rlt.worker.loopcd import reserved_bytes def execution_buffer_bytes(config, scheduler, cache, execution, element_size): @@ -13,7 +14,7 @@ def execution_buffer_bytes(config, scheduler, cache, execution, element_size): rows = 1 << (rows - 1).bit_length() width = math.ceil(config.max_position_embeddings / cache.block_size) hidden = config.hidden_size * element_size - total = 0 + total = reserved_bytes(config, scheduler, execution, element_size) if execution.static_buffers: # Existing prefill/eager workspaces remain separate from async banks. total += 4 * rows * (hidden + 4 * width + 36) + scheduler.max_num_seqs * hidden diff --git a/vllm_rlt/worker/loopcd.py b/vllm_rlt/worker/loopcd.py new file mode 100644 index 00000000..60cf5901 --- /dev/null +++ b/vllm_rlt/worker/loopcd.py @@ -0,0 +1,90 @@ +"""Readout-only fixed LoopCD for Ouro; original recurrent feedback is preserved.""" + +from dataclasses import dataclass + +import torch + +from vllm_rlt.models import OuroForCausalLM +from vllm_rlt.request import Request, Stage + + +@dataclass(frozen=True) +class Reference: + owner: Request + position: int + output_index: int + loop: int + hidden: torch.Tensor + + def check(self, request, position, output_index, loop): + if self.owner is not request or (self.position, self.output_index, self.loop) != ( + position, + output_index, + loop, + ): + raise RuntimeError("LoopCD reference owner/token/loop mismatch") + + +def active(params): + return params.loopcd is not None and params.loopcd.strength != 0 + + +def validate(model, params, execution, cache, scheduler, exit_config, speculative): + if not active(params): + return + if type(model) is not OuroForCausalLM: + raise ValueError("LoopCD currently supports Ouro only") + if not execution.loopcd: + raise ValueError("LoopCD requires ExecutionConfig(loopcd=True) for memory reservation") + if ( + execution.async_scheduling + or execution.prefill_uva + or cache.layout != "last_exited" + or cache.enable_prefix_caching + or speculative is not None + ): + raise ValueError("LoopCD requires synchronous last_exited without prefix/P-D/speculation") + if scheduler.enable_preemption and execution.cuda_graphs: + raise ValueError("LoopCD preemption currently requires eager execution") + maximum = params.max_loops or model.config.total_ut_steps + prefill = execution.prefill_depth or model.config.total_ut_steps + if params.min_loops != maximum or params.exit_threshold != 1 or exit_config.mode != "ouro": + raise ValueError("LoopCD requires fixed decode depth and ouro exit mode") + cd = params.loopcd + if not (1 <= cd.reference_loop < maximum and 1 <= cd.prefill_loop < prefill): + raise ValueError("LoopCD requires 1 <= reference_loop < actual prefill/decode depth") + + +def preemption_reference(request, stage, reference): + """Validate only the reference that the suspended token still needs. + + A partial prompt chunk and a next-token prelude produce a fresh reference. + After the decode reference depth, or at readout, exact restoration is needed. + """ + if not active(request.sampling_params): + return None + cd = request.sampling_params.loopcd + needed = stage == Stage.CODA or ( + stage == Stage.RECURRENT and request.loops_done >= cd.reference_loop + ) + if not needed: + return None + if reference is None: + raise RuntimeError("LoopCD reference missing at preemption") + index = request.num_scheduled_outputs + loop = cd.prefill_loop if index == 0 else cd.reference_loop + reference.check(request, request.position, index, loop) + return reference + + +def extrapolate(final, reference, strength): + return final + strength * (final - reference) + + +def reserved_bytes(config, scheduler, execution, element_size): + if not execution.loopcd: + return 0 + rows = scheduler.max_num_batched_tokens + states = (scheduler.max_num_seqs + 3 * rows) * config.hidden_size * element_size + # Two-head reference logits and FP32 combination coexist with normal logits. + return states + 3 * rows * config.vocab_size * 4 diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index 055d7f35..355ee458 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -12,6 +12,7 @@ from vllm_rlt.request import Request, Stage from vllm_rlt.worker.buffers import Workspace from vllm_rlt.worker.cuda_graph import RecurrentGraphs +from vllm_rlt.worker.loopcd import Reference, active, extrapolate from vllm_rlt.worker.sampler import Sampler @@ -82,6 +83,21 @@ def __init__( self.sampler = Sampler(self.device) self.exit_config = exit_config or ExitConfig() self.execution_config = execution_config or ExecutionConfig() + self.loopcd_references = {} + self.loopcd_stats = dict( + captures=0, + guided_rows=0, + head_calls=0, + head_rows=0, + skipped_prompt_head_calls=0, + skipped_prompt_head_rows=0, + reference_peak_bytes=0, + prefill_core_rows=0, + prefill_submitted_rows=0, + decode_core_rows=0, + decode_submitted_rows=0, + ) + self.prefill_depth = self.execution_config.prefill_depth or model.config.total_ut_steps scheduler = scheduler_config or SchedulerConfig() self.graphs = ( RecurrentGraphs( @@ -255,23 +271,98 @@ def _prefill(self, batch): ids = [i.request.request_id for i in active] positions = [i.token_start + offset for i in active] tokens = [i.request.prompt_token_ids[p] for i, p in zip(active, positions)] - hidden = self._prefill_tokens(ids, positions, tokens) + hidden = self._prefill_tokens(ids, positions, tokens, [i.request for i in active]) for row, item in enumerate(active): self._save(item.request, hidden[row]) return - ids, positions, tokens = [], [], [] + ids, positions, tokens, requests = [], [], [], [] for item in batch.items: start, count, request = item.token_start, item.token_count, item.request ids.extend([request.request_id] * count) positions.extend(range(start, start + count)) tokens.extend(request.prompt_token_ids[start : start + count]) - hidden = self._prefill_tokens(ids, positions, tokens) + requests.extend([request] * count) + hidden = self._prefill_tokens(ids, positions, tokens, requests) offset = 0 for item in batch.items: offset += item.token_count self._save(item.request, hidden[offset - 1]) - def _prefill_tokens(self, ids, positions, tokens): + def _capture_loopcd(self, requests, hidden, depths, positions, indices, *, prefill=False): + if not self.execution_config.loopcd: + return + for row, (request, depth, position, index) in enumerate( + zip(requests, depths, positions, indices) + ): + if not active(request.sampling_params): + continue + if prefill and row + 1 < len(requests) and requests[row + 1] is request: + continue + cd = request.sampling_params.loopcd + loop = cd.prefill_loop if prefill else cd.reference_loop + if depth + 1 == loop: + self.loopcd_references[request.request_id] = Reference( + request, position, index, loop, hidden[row].clone() + ) + self.loopcd_stats["captures"] += 1 + size = sum( + ref.hidden.numel() * ref.hidden.element_size() + for ref in self.loopcd_references.values() + ) + self.loopcd_stats["reference_peak_bytes"] = max( + self.loopcd_stats["reference_peak_bytes"], size + ) + + def _count_core(self, effective, submitted, *, prefill=False): + if self.execution_config.loopcd: + phase = "prefill" if prefill else "decode" + self.loopcd_stats[phase + "_core_rows"] += effective + self.loopcd_stats[phase + "_submitted_rows"] += submitted + + def _loopcd_readout(self, hidden, requests, positions, indices): + if not self.execution_config.loopcd: + return self.model.coda(hidden) + references = [] + for request, position, index in zip(requests, positions, indices): + ref = None + if active(request.sampling_params): + cd = request.sampling_params.loopcd + loop = cd.prefill_loop if index == 0 else cd.reference_loop + ref = self.loopcd_references.get(request.request_id) + if ref is None: + raise RuntimeError("LoopCD reference missing at readout") + ref.check(request, position, index, loop) + references.append(ref) + guided = hidden + for row, ref in enumerate(references): + if ref is None or requests[row].sampling_params.loopcd.implementation == "two_head": + continue + if guided is hidden: + guided = hidden.clone() + guided[row] = extrapolate( + hidden[row], ref.hidden, requests[row].sampling_params.loopcd.strength + ) + logits = self.model.coda(guided) + self.loopcd_stats["head_calls"] += 1 + self.loopcd_stats["head_rows"] += logits.shape[0] + weak_rows = [ + row + for row, ref in enumerate(references) + if ref is not None and requests[row].sampling_params.loopcd.implementation == "two_head" + ] + if weak_rows: + weak = self.model.coda(torch.stack([references[row].hidden for row in weak_rows])) + logits = logits.float() + for row, value in zip(weak_rows, weak): + logits[row] = extrapolate( + logits[row], value.float(), requests[row].sampling_params.loopcd.strength + ) + self.loopcd_stats["head_calls"] += 1 + self.loopcd_stats["head_rows"] += len(weak_rows) + self.loopcd_stats["guided_rows"] += sum(ref is not None for ref in references) + return logits + + def _prefill_tokens(self, ids, positions, tokens, requests=()): cache = self.cache_manager if cache.layout == "last_exited" and cache.attention_capabilities.packed_prefill: # Prefill has genuinely ragged query sequences. Do not pad token rows @@ -290,7 +381,7 @@ def _prefill_tokens(self, ids, positions, tokens): else torch.tensor(tokens, device=self.device, dtype=torch.long) ) hidden = self.model.prelude(tensor) - for depth in range(self.model.config.total_ut_steps): + for depth in range(self.prefill_depth): metadata = ( bank.metadata(depth) if bank @@ -301,6 +392,15 @@ def _prefill_tokens(self, ids, positions, tokens): hidden, _ = self.model.recurrent_prepared( hidden, metadata, cache, compute_gate=False ) + self._count_core(len(ids), len(ids), prefill=True) + self._capture_loopcd( + requests, + hidden, + [depth] * len(requests), + positions, + [0] * len(requests), + prefill=True, + ) if bank: bank.release() return hidden @@ -312,11 +412,20 @@ def _prefill_tokens(self, ids, positions, tokens): else torch.tensor(tokens, device=self.device, dtype=torch.long) ) hidden = self.model.prelude(tensor) - for depth in range(self.model.config.total_ut_steps): + for depth in range(self.prefill_depth): # Same workspace metadata can be refilled only once previous DMA is done. if workspace: workspace.acquire() hidden, _ = self._core(hidden, ids, [depth] * len(ids), positions, workspace, size) + self._count_core(len(ids), size, prefill=True) + self._capture_loopcd( + requests, + hidden, + [depth] * len(requests), + positions, + [0] * len(requests), + prefill=True, + ) if workspace: workspace.release() return hidden @@ -359,6 +468,8 @@ def _execute(self, batch, prepared=None): routing = prepared.routing if prepared is not None else None workspace = None if routing is not None else self._workspace(group) if batch.stage == Stage.PRELUDE: + for request in requests: + self.loopcd_references.pop(request.request_id, None) if prepared is not None: # Feed the sampled GPU IDs directly to embedding. No .item(), # .tolist(), or CPU token roundtrip on the dependency path. @@ -415,12 +526,29 @@ def _execute(self, batch, prepared=None): workspace, size, ) + self._count_core(len(requests), len(hidden)) + self._capture_loopcd( + requests, + hidden, + prepared.depths if prepared else [r.loops_done for r in requests], + prepared.positions if prepared else [r.position for r in requests], + prepared.output_indices + if prepared + else [r.num_scheduled_outputs for r in requests], + ) self._save_batch(requests, hidden, routing) if self.lookahead_head is not None: logits = self.lookahead_head(hidden).squeeze(-1) result = logits[: len(requests)].float().sigmoid() if logits is not None else None elif batch.stage == Stage.CODA: - logits = self.model.coda(hidden) + logits = self._loopcd_readout( + hidden, + requests, + prepared.positions if prepared else [r.position for r in requests], + prepared.output_indices + if prepared + else [r.num_scheduled_outputs for r in requests], + ) result = torch.stack( [self._sample_tensor(row, r) for row, r in zip(logits, requests)] ) @@ -581,6 +709,7 @@ def release(self, request_id): event = self.events.pop(request_id, None) if event is not None: event.synchronize() + self.loopcd_references.pop(request_id, None) if self.async_state is not None: self.async_state.release(request_id) return