From d700ecca664da2e8f3122f66ad9d45de4295e45a Mon Sep 17 00:00:00 2001 From: 0z5a Date: Mon, 5 Oct 2026 02:21:41 +0800 Subject: [PATCH 1/7] 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 aaf7943ccabffcec2702d2855ebd9d3a85bf5f23 Mon Sep 17 00:00:00 2001 From: Dragonborn <1078757180@qq.com> Date: Mon, 5 Oct 2026 02:02:03 +0000 Subject: [PATCH 2/7] Support eager preemption for Ouro LoopCD --- docs/loopcd-ouro.md | 40 +++- tests/test_loopcd.py | 3 +- tests/test_loopcd_preemption.py | 319 +++++++++++++++++++++++++++++ vllm_rlt/core/scheduling_policy.py | 7 + vllm_rlt/engine/preemption.py | 13 ++ vllm_rlt/worker/loopcd.py | 31 ++- 6 files changed, 400 insertions(+), 13 deletions(-) create mode 100644 tests/test_loopcd_preemption.py diff --git a/docs/loopcd-ouro.md b/docs/loopcd-ouro.md index 01fc0d36..788c490e 100644 --- a/docs/loopcd-ouro.md +++ b/docs/loopcd-ouro.md @@ -47,12 +47,12 @@ 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. +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 @@ -75,3 +75,31 @@ 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. + + +## 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. Graph preemption, async execution, P/D handoff, +prefix caching and speculative decoding are outside this contract. diff --git a/tests/test_loopcd.py b/tests/test_loopcd.py index 2ad397b5..48b1dd06 100644 --- a/tests/test_loopcd.py +++ b/tests/test_loopcd.py @@ -175,14 +175,13 @@ def test_stale_reference_rejected(field): replace(ref, **{field: value}).check(request, 1, 0, 1) -@pytest.mark.parametrize("case", ["unreserved", "reference", "early_exit", "prefix", "preemption"]) +@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" ), - scheduler_config=SchedulerConfig(enable_preemption=case == "preemption"), execution_config=ExecutionConfig(loopcd=case != "unreserved"), ) params = params_for() diff --git a/tests/test_loopcd_preemption.py b/tests/test_loopcd_preemption.py new file mode 100644 index 00000000..6a44c8c8 --- /dev/null +++ b/tests/test_loopcd_preemption.py @@ -0,0 +1,319 @@ +"""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 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"): + return LLMEngine( + model, + 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), + ) + + +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. + engine.scheduler.selected_request_ids.clear() + 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, + ) + + +@pytest.mark.parametrize("policy", ["priority", "pressure"]) +@torch.inference_mode() +def test_guided_requests_survive_scheduler_driven_preemption(policy): + model = model_for() + engine = LLMEngine( + model, + 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), + ) + 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) + 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/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/preemption.py b/vllm_rlt/engine/preemption.py index fc61b42d..0f466b36 100644 --- a/vllm_rlt/engine/preemption.py +++ b/vllm_rlt/engine/preemption.py @@ -6,10 +6,12 @@ """ from copy import deepcopy +from dataclasses import replace import torch from vllm_rlt.request import Stage +from vllm_rlt.worker.loopcd import preemption_reference class PreemptionManager: @@ -78,6 +80,9 @@ def preempt(self, requester, *, priority_only=False): 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) + ) allocation = cache._get_allocation(victim.request_id) blocks = [b for table in allocation.block_tables for b in table] # Copy views one page at a time: a pressure recovery must not allocate @@ -89,6 +94,9 @@ def preempt(self, requester, *, priority_only=False): values[row].copy_(cache.value_cache[block]) snapshot = dict( stage=victim.stage, + loopcd_reference=None + if reference is None + else replace(reference, hidden=reference.hidden.cpu().clone()), maximum=allocation.max_tokens, pages=len(allocation.block_tables[0]), written=deepcopy(allocation.written), @@ -122,6 +130,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"]) frontier = min(state["maximum"], state["pages"] * cache.block_size) if not cache.allocate(request.request_id, state["maximum"], initial_tokens=frontier): return False @@ -135,6 +144,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/worker/loopcd.py b/vllm_rlt/worker/loopcd.py index d27582c4..60cf5901 100644 --- a/vllm_rlt/worker/loopcd.py +++ b/vllm_rlt/worker/loopcd.py @@ -5,7 +5,7 @@ import torch from vllm_rlt.models import OuroForCausalLM -from vllm_rlt.request import Request +from vllm_rlt.request import Request, Stage @dataclass(frozen=True) @@ -41,12 +41,11 @@ def validate(model, params, execution, cache, scheduler, exit_config, speculativ 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" - ) + 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": @@ -56,6 +55,28 @@ def validate(model, params, execution, cache, scheduler, exit_config, speculativ 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) From bf4c8871893a64a3a16805506f354bb1903e974b Mon Sep 17 00:00:00 2001 From: Dragonborn <1078757180@qq.com> Date: Mon, 5 Oct 2026 02:02:03 +0000 Subject: [PATCH 3/7] Fix outdated tiny Ouro helper in speculative test --- tests/test_speculative.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_speculative.py b/tests/test_speculative.py index 412a8b13..162f54ec 100644 --- a/tests/test_speculative.py +++ b/tests/test_speculative.py @@ -395,7 +395,7 @@ def test_greedy_round_reads_back_once_with_ragged_requests(): # Draft IDs stay on the device; every speculative round, including ragged and # budget-shortened ones, synchronizes with the host exactly once. torch.manual_seed(123) - m = OuroForCausalLM(OuroConfig.tiny(hidden_size=256, head_dim=64)).to( + m = OuroForCausalLM(tiny_ouro_config(hidden_size=256, head_dim=64)).to( device="cuda", dtype=torch.bfloat16 ) e = engine(m, k=3, cache_config=CacheConfig(128, 16), attention_backend="triton") From e8aa1ee57734bc5c1118c143ebdb68d0b53210a1 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Tue, 6 Oct 2026 00:37:07 +0800 Subject: [PATCH 4/7] 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 c3a9c07cfa754817843e4eec24097df5f179c633 Mon Sep 17 00:00:00 2001 From: Dragonborn <1078757180@qq.com> Date: Wed, 7 Oct 2026 09:39:15 +0000 Subject: [PATCH 5/7] Test scheduler-driven LoopCD preemption on CUDA --- docs/loopcd-ouro.md | 14 +++++++++++-- tests/test_loopcd_preemption.py | 36 +++++++++++++++++++++++++++------ 2 files changed, 42 insertions(+), 8 deletions(-) diff --git a/docs/loopcd-ouro.md b/docs/loopcd-ouro.md index 788c490e..5d3685d3 100644 --- a/docs/loopcd-ouro.md +++ b/docs/loopcd-ouro.md @@ -101,5 +101,15 @@ 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. Graph preemption, async execution, P/D handoff, -prefix caching and speculative decoding are outside this contract. +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. + +Separate official Ouro-1.4B runs cover FP32/BF16 priority preemption and FP32 +multi-request KV pressure. Priority logits match exactly. Pressure tokens, +exit depths and recurrent-work counts match uninterrupted runs; score +differences up to `2.4e-5` also occur in an unpreempted batching control. +Official BF16 multi-request pressure and throughput remain unmeasured. +Graph preemption, async execution, P/D handoff, prefix caching and speculative +decoding are outside this contract. diff --git a/tests/test_loopcd_preemption.py b/tests/test_loopcd_preemption.py index 6a44c8c8..a59cf20f 100644 --- a/tests/test_loopcd_preemption.py +++ b/tests/test_loopcd_preemption.py @@ -36,6 +36,11 @@ def model_for(): 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, @@ -53,9 +58,10 @@ def params_for(strength=0.3, decode=4, implementation="two_head"): ) -def engine_for(model, mode="refill"): +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, @@ -64,7 +70,9 @@ def engine_for(model, mode="refill"): enable_preemption=True, mode=mode, ), - execution_config=ExecutionConfig(loopcd=True, prefill_depth=4), + execution_config=ExecutionConfig( + loopcd=True, prefill_depth=4, static_buffers=static_buffers + ), ) @@ -264,12 +272,26 @@ def test_graph_preemption_stays_rejected_until_separate_qualification(): ) +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): - model = model_for() +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, @@ -282,7 +304,9 @@ def test_guided_requests_survive_scheduler_driven_preemption(policy): policy="priority" if policy == "priority" else "fcfs", enable_preemption=True, ), - execution_config=ExecutionConfig(loopcd=True, prefill_depth=4), + execution_config=ExecutionConfig( + loopcd=True, prefill_depth=4, static_buffers=static_buffers + ), ) params = params_for() if policy == "priority": @@ -311,7 +335,7 @@ def test_guided_requests_survive_scheduler_driven_preemption(policy): assert engine.cache_manager.num_used_blocks == 0 if policy == "priority": assert list(completed) == ["b", "a"] - baseline = engine_for(model) + 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"): From bc1a926d7285760f2333977febbdc1a8edec813e Mon Sep 17 00:00:00 2001 From: 0z5a Date: Fri, 9 Oct 2026 16:02:09 +0200 Subject: [PATCH 6/7] 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 b3f88f8a7333c8aa6ecabc7f5f43bfe017fcfb14 Mon Sep 17 00:00:00 2001 From: Dragonborn <1078757180@qq.com> Date: Sat, 10 Oct 2026 04:21:58 +0000 Subject: [PATCH 7/7] Check LoopCD resume against the continuation oracle --- docs/loopcd-ouro.md | 61 +++- tests/test_loopcd_preemption_continuation.py | 299 +++++++++++++++++++ 2 files changed, 347 insertions(+), 13 deletions(-) create mode 100644 tests/test_loopcd_preemption_continuation.py diff --git a/docs/loopcd-ouro.md b/docs/loopcd-ouro.md index 5d3685d3..389b45bd 100644 --- a/docs/loopcd-ouro.md +++ b/docs/loopcd-ouro.md @@ -69,12 +69,10 @@ 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. +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 @@ -106,10 +104,47 @@ 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. -Separate official Ouro-1.4B runs cover FP32/BF16 priority preemption and FP32 -multi-request KV pressure. Priority logits match exactly. Pressure tokens, -exit depths and recurrent-work counts match uninterrupted runs; score -differences up to `2.4e-5` also occur in an unpreempted batching control. -Official BF16 multi-request pressure and throughput remain unmeasured. -Graph preemption, async execution, P/D handoff, prefix caching and speculative -decoding are outside this contract. +`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/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