From d700ecca664da2e8f3122f66ad9d45de4295e45a Mon Sep 17 00:00:00 2001 From: 0z5a Date: Mon, 5 Oct 2026 02:21:41 +0800 Subject: [PATCH 1/5] Add initial LoopCD readout support and correctness checks for Ouro Signed-off-by: 0z5a --- docs/loopcd-ouro.md | 77 +++++++++ tests/__init__.py | 1 + tests/test_loopcd.py | 289 ++++++++++++++++++++++++++++++++ vllm_rlt/__init__.py | 3 +- vllm_rlt/config.py | 5 + vllm_rlt/engine/llm_engine.py | 29 +++- vllm_rlt/pd/engine.py | 2 + vllm_rlt/sampling_params.py | 35 ++++ vllm_rlt/worker/buffers.py | 3 +- vllm_rlt/worker/loopcd.py | 69 ++++++++ vllm_rlt/worker/model_runner.py | 143 +++++++++++++++- 11 files changed, 645 insertions(+), 11 deletions(-) create mode 100644 docs/loopcd-ouro.md create mode 100644 tests/__init__.py create mode 100644 tests/test_loopcd.py create mode 100644 vllm_rlt/worker/loopcd.py diff --git a/docs/loopcd-ouro.md b/docs/loopcd-ouro.md new file mode 100644 index 00000000..01fc0d36 --- /dev/null +++ b/docs/loopcd-ouro.md @@ -0,0 +1,77 @@ +# 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, preemption, prefill UVA, asynchronous scheduling, speculation +and P/D transfer are rejected before request allocation. Other model families +are rejected. Existing CUDA Graph machinery is unchanged; official-checkpoint +eager/Graph qualification belongs to Task 1. Task 1's single-worker Ouro +preemption slice must establish its save/rebuild, cancellation and ownership +contract before relaxing the admission guard. + +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 Ouro-1.4B checkpoint +`574fa66cb8bf5abdc979642d01cf2b79b16bfab1` qualification, BF16 fused development +measurements, generation/multiple-choice quality and matched GPU cost remain +pending. Existing exploratory Huginn results are separate evidence. This initial +PR does not establish official Ouro quality, learning gains, throughput gains, +long-context behavior or compatibility with the rejected combinations. 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/test_loopcd.py b/tests/test_loopcd.py new file mode 100644 index 00000000..2ad397b5 --- /dev/null +++ b/tests/test_loopcd.py @@ -0,0 +1,289 @@ +"""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", "preemption"]) +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" + ), + scheduler_config=SchedulerConfig(enable_preemption=case == "preemption"), + 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/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/engine/llm_engine.py b/vllm_rlt/engine/llm_engine.py index 8a7bb559..fbfb32dc 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": @@ -75,6 +88,9 @@ def __init__( num_blocks, self.memory_plan = plan_cache( model, cache_config, scheduler_config, self.execution_config, attention_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, @@ -143,6 +159,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 ): @@ -160,7 +185,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] @@ -303,7 +328,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/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..d27582c4 --- /dev/null +++ b/vllm_rlt/worker/loopcd.py @@ -0,0 +1,69 @@ +"""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 + + +@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 scheduler.enable_preemption + or speculative is not None + ): + raise ValueError( + "LoopCD requires synchronous last_exited without prefix/preemption/P-D/speculation" + ) + 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 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 b0025137..ee387f0f 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 getattr(cache.attention, "generation", None) == 4: # 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)] ) @@ -586,6 +714,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 From e8aa1ee57734bc5c1118c143ebdb68d0b53210a1 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Tue, 6 Oct 2026 00:37:07 +0800 Subject: [PATCH 2/5] Add stateful mixed-depth LoopCD oracle and closed-loop scale validation Signed-off-by: 0z5a --- benchmarks/loopcd_scale.py | 314 ++++++++++++++++++++++++++++++ docs/loopcd-scale.md | 81 ++++++++ tests/reference/continuation.py | 104 ++++++++++ tests/test_loopcd_continuation.py | 83 ++++++++ tests/test_loopcd_scale.py | 36 ++++ 5 files changed, 618 insertions(+) create mode 100644 benchmarks/loopcd_scale.py create mode 100644 docs/loopcd-scale.md create mode 100644 tests/reference/continuation.py create mode 100644 tests/test_loopcd_continuation.py create mode 100644 tests/test_loopcd_scale.py 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/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/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_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_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() From bc1a926d7285760f2333977febbdc1a8edec813e Mon Sep 17 00:00:00 2001 From: 0z5a Date: Fri, 9 Oct 2026 16:02:09 +0200 Subject: [PATCH 3/5] Document audited full-model LoopCD C32 results and qualification limits Signed-off-by: 0z5a --- benchmarks/results/loopcd-c32/summary.json | 305 +++++++++++++++++++++ docs/loopcd-scale-results.md | 93 +++++++ 2 files changed, 398 insertions(+) create mode 100644 benchmarks/results/loopcd-c32/summary.json create mode 100644 docs/loopcd-scale-results.md 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-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). From 48932b2360f6ee750c154b268dae71fef8701782 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sat, 10 Oct 2026 17:01:39 +0200 Subject: [PATCH 4/5] Reject decode depths without initialized prefill KV Validate possible decode depth at request admission, including disabled and zero-strength LoopCD. Use only the requested trace prefix, and preserve requests that generate solely the prefill output. Add 18 admission and continuation regressions. CPU suite: 567 passed, 29 skipped, 114 GPU cases deselected; Ruff lint and format checks passed. --- tests/test_engine.py | 148 +++++++++++++++++++++++++++++++++- vllm_rlt/engine/llm_engine.py | 9 +++ 2 files changed, 156 insertions(+), 1 deletion(-) diff --git a/tests/test_engine.py b/tests/test_engine.py index 75199bcf..8dc32790 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -1,10 +1,19 @@ import math +from dataclasses import replace import pytest import torch from tests.helpers import tiny_ouro_config -from vllm_rlt import LLM, CacheConfig, SamplingParams, SchedulerConfig +from vllm_rlt import ( + LLM, + CacheConfig, + ExecutionConfig, + ExitConfig, + LoopCDParams, + SamplingParams, + SchedulerConfig, +) from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM from vllm_rlt.request import FinishReason, RequestOutput, Stage @@ -324,3 +333,140 @@ def test_finish_reason_preserves_public_string_values(reason): output = RequestOutput.from_request(request) assert output.finished and type(output.finish_reason) is str assert output.finish_reason == reason.value + + +@pytest.mark.parametrize("strength", [None, 0.0, 0.3]) +@pytest.mark.parametrize("max_loops", [3, None]) +@torch.inference_mode() +def test_decode_deeper_than_prefill_rejected_before_admission(strength, max_loops): + guidance = None if strength is None else LoopCDParams(strength=strength) + engine = LLMEngine( + tiny_model(), + cache_config=CacheConfig(num_blocks=128, block_size=4), + execution_config=ExecutionConfig(loopcd=bool(strength), prefill_depth=2), + ) + params = SamplingParams( + max_tokens=3, + min_loops=max_loops or 4, + max_loops=max_loops, + ignore_eos=True, + loopcd=guidance, + ) + with pytest.raises(ValueError, match="decode depth .* exceeds prefill_depth 2"): + engine.add_request("r", [4, 7, 9], params) + assert not engine.scheduler.requests + assert not engine.has_unfinished_requests() + assert engine.cache_manager.num_used_blocks == 0 + assert not engine.model_runner.loopcd_references + + # Rejected admission must leave the ID available for a valid continuation. + engine.add_request("r", [4, 7, 9], replace(params, min_loops=2, max_loops=2)) + outputs, _ = drain(engine) + assert outputs["r"].exit_depths == [2, 2, 2] + assert len(outputs["r"].token_ids) == 3 + assert engine.cache_manager.num_used_blocks == 0 + assert not engine.model_runner.loopcd_references + engine.close() + + +@pytest.mark.parametrize("strength", [None, 0.0, 0.3]) +@torch.inference_mode() +def test_prefill_history_allows_shallower_decode(strength): + engine = LLMEngine( + tiny_model(), + cache_config=CacheConfig(num_blocks=128, block_size=4), + execution_config=ExecutionConfig(loopcd=bool(strength), prefill_depth=4), + ) + engine.add_request( + "r", + [4, 7, 9], + SamplingParams( + max_tokens=3, + min_loops=3, + max_loops=3, + ignore_eos=True, + loopcd=None if strength is None else LoopCDParams(strength=strength), + ), + ) + outputs, _ = drain(engine) + assert outputs["r"].exit_depths == [4, 3, 3] + assert len(outputs["r"].token_ids) == 3 + assert engine.cache_manager.num_used_blocks == 0 + assert not engine.model_runner.loopcd_references + engine.close() + + +@pytest.mark.parametrize("mode", ["ouro", "ouro_delayed", "random_lookahead"]) +def test_prefill_history_checks_adaptive_decode_bound(mode): + engine = LLMEngine( + tiny_model(), + execution_config=ExecutionConfig(prefill_depth=2), + exit_config=ExitConfig(mode=mode), + ) + params = SamplingParams(max_tokens=3, min_loops=1, max_loops=4, exit_threshold=0.7) + with pytest.raises(ValueError, match="decode depth 4 exceeds prefill_depth 2"): + engine.add_request("r", [4, 7, 9], params) + assert not engine.scheduler.requests + assert engine.cache_manager.num_used_blocks == 0 + engine.close() + + +@pytest.mark.parametrize("trace", [[2, 3, 2], [2, 2, 3]]) +def test_prefill_history_rejects_deeper_trace_before_admission(trace): + engine = LLMEngine( + tiny_model(), + execution_config=ExecutionConfig(prefill_depth=2), + exit_config=ExitConfig(mode="trace", depths_by_request={"profile": trace}), + ) + params = SamplingParams(max_tokens=3, min_loops=1, max_loops=4) + with pytest.raises(ValueError, match="decode depth 3 exceeds prefill_depth 2"): + engine.add_request("r", [4, 7, 9], params, trace_id="profile") + assert not engine.scheduler.requests + assert engine.cache_manager.num_used_blocks == 0 + engine.close() + + +@torch.inference_mode() +def test_prefill_history_uses_only_requested_trace_depths(): + engine = LLMEngine( + tiny_model(), + execution_config=ExecutionConfig(prefill_depth=2), + exit_config=ExitConfig(mode="trace", depths_by_request={"profile": [2, 1, 2, 4]}), + ) + engine.add_request( + "r", + [4, 7, 9], + SamplingParams(max_tokens=3, min_loops=1, ignore_eos=True), + trace_id="profile", + ) + outputs, _ = drain(engine) + assert outputs["r"].exit_depths == [2, 1, 2] + assert len(outputs["r"].token_ids) == 3 + assert engine.cache_manager.num_used_blocks == 0 + engine.close() + + +@pytest.mark.parametrize("strength", [None, 0.0, 0.3]) +@torch.inference_mode() +def test_prefill_only_output_does_not_require_decode_history(strength): + engine = LLMEngine( + tiny_model(), + execution_config=ExecutionConfig(loopcd=bool(strength), prefill_depth=2), + ) + engine.add_request( + "r", + [4, 7, 9], + SamplingParams( + max_tokens=1, + min_loops=4, + max_loops=4, + ignore_eos=True, + loopcd=None if strength is None else LoopCDParams(strength=strength), + ), + ) + outputs, _ = drain(engine) + assert outputs["r"].exit_depths == [2] + assert len(outputs["r"].token_ids) == 1 + assert engine.cache_manager.num_used_blocks == 0 + assert not engine.model_runner.loopcd_references + engine.close() diff --git a/vllm_rlt/engine/llm_engine.py b/vllm_rlt/engine/llm_engine.py index aed15b83..f238c1b5 100644 --- a/vllm_rlt/engine/llm_engine.py +++ b/vllm_rlt/engine/llm_engine.py @@ -204,6 +204,15 @@ def add_request( raise ValueError( "exit trace must cover every output and obey prefill/decode bounds" ) + # The first output comes from prefill. Later tokens require prompt KV + # history at every recurrent depth they can visit. + if params.max_tokens > 1: + decode_depth = max(trace[1 : params.max_tokens]) if trace else max_loops + if decode_depth > self.prefill_depth: + raise ValueError( + f"decode depth {decode_depth} exceeds prefill_depth {self.prefill_depth}; " + "increase prefill_depth or reduce decode depth" + ) capacity = len(prompt_token_ids) + params.max_tokens - 1 if capacity > config.max_position_embeddings: raise ValueError("prompt plus decode positions exceed the model context length") From aeb14a61add3894d46b3be1ab7066591a53c6093 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 1 Nov 2026 23:32:13 +0100 Subject: [PATCH 5/5] Keep the LoopCD GPU-session tests independent of test order Under --run-gpu, earlier tests in the same process initialize CUDA, so the CPU oracle test now checks that it did not change CUDA initialization instead of asserting CUDA was never initialized. With main's bucketed coda graphs and power-of-two padding, one graph covers every batch of the request-reuse test, so it no longer expects an eager fallback; it still requires replays and compares graph and eager outputs. Signed-off-by: 0z5a --- tests/test_loopcd.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_loopcd.py b/tests/test_loopcd.py index 180dd8e2..5f7abd78 100644 --- a/tests/test_loopcd.py +++ b/tests/test_loopcd.py @@ -165,7 +165,9 @@ def test_loopcd_coda_graphs_match_eager_through_request_reuse(implementation): rounds.append(finished) if graphs: coda = engine.model_runner.coda_graphs - assert coda.captures == 1 and coda.replays > 1 and coda.fallbacks > 0 + # Power-of-two padding lets the one bucketed coda graph cover every + # batch here, so no row count falls back to eager. + assert coda.captures == 1 and coda.replays > 1 results.append(rounds) assert results[0] == results[1] @@ -320,6 +322,9 @@ def test_reference_and_two_head_scratch_budgeted_before_auto_kv(): ) @torch.inference_mode() def test_teacher_forced_continuation_against_dense_oracle(depth, strength, implementation): + # CPU-only: this test must not initialize CUDA. A --run-gpu session may + # already have initialized it in an earlier test of the same process. + cuda_was_initialized = torch.cuda.is_initialized() model = model_for() engine = LLMEngine( model, @@ -376,4 +381,4 @@ def observe(logits, params, generator=None): 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() + assert torch.cuda.is_initialized() == cuda_was_initialized