diff --git a/README.md b/README.md index ae066e87..78f5db26 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,11 @@ The default runtime uses synchronous execution; Ouro uses its original gate. Advanced execution and cache features are opt-in; see the guides below for supported combinations. +An [experimental deferred Ouro runner](docs/recipes/ouro_deferred.md) groups deep +loops across response positions with true per-loop KV. It provides a standalone +mixed-readout experiment and an eager same-readout reference; real-weight BF16 +quality, memory and speed qualification are still required. + ## Model Support | Checkpoint | Implementation | Documented validation scope | Guide | diff --git a/docs/README.md b/docs/README.md index 02848af0..b65372c5 100644 --- a/docs/README.md +++ b/docs/README.md @@ -20,6 +20,7 @@ an environment to receiving the first generated response. The | [HTTP serving](serving.md) | Completions API, streaming, and service lifecycle | | [Profiling](profiling.md) | Core profiler options, Python/HTTP controls, PD collection, and per-rank archives | | [Ouro-Thinking recipe](recipes/ouro_thinking.md) | Shared 1.4B/2.6B setup, JSON/SSE client, budgets and validation limits | +| [Experimental deferred Ouro](recipes/ouro_deferred.md) | Fixed-interval deep batching, true per-loop KV, same-readout reference and unfinished qualification | | [Accuracy evaluation](accuracy.md) | GSM8K regression setup and comparison methodology | ## Support and Runtime Notes diff --git a/docs/recipes/ouro_deferred.md b/docs/recipes/ouro_deferred.md new file mode 100644 index 00000000..4fb272ae --- /dev/null +++ b/docs/recipes/ouro_deferred.md @@ -0,0 +1,112 @@ +# Experimental deferred Ouro decoding + +This opt-in, synchronous runner groups Ouro's third and fourth loops across +response positions. It is intended for experiments by Ouro rollout developers +working with verl/vime, without importing either training framework. + +## Execution and comparison + +Prompt prefill runs all four loops and produces the first output. At interval +four, the next three response inputs read out after loop two. The fourth input +waits while all four retained loop-two states pass through loop three and then +loop four. The block-end loop-four state produces the next output. Every row +keeps its original token position and its own genuinely computed KV history at +each loop; shallow KV is never propagated into the pending deep planes. + +The main execution variable is the deep batch size. For `N` active requests at +interval `B`, a complete aligned deep block contains `N * B` rows. Eager +execution uses `N` rows on each step. Both execute the same deep rows at a +completed boundary; deferral changes when and how those rows are grouped. +Unfinished deep tails are discarded when a response stops or reaches its limit. + +Use three distinct comparisons: + +- `--interval 1`: four-loop readout at every position, with full-depth histories. +- `--interval 4 --eager`: immediate deep writes under the mixed readout schedule. +- `--interval 4`: deferred deep writes under that same mixed readout schedule. + +The eager/deferred pair isolates execution timing. The interval-one/mixed pair +changes the policy and needs a paired task-quality comparison. BF16 batch shapes +can also change scores and sampled trajectories; CPU causal equality does not +qualify BF16 GPU probability replay. + +## Run an experiment + +Install the normal repository dependencies and provide a local Ouro checkpoint. +The CLI uses BF16 on CUDA and FP32 on CPU. CUDA needs a scheduler-reserved device; +the selected backend must support the device and dtype. + +```bash +python -m vllm_rlt.entrypoints.ouro_deferred /models/Ouro-1.4B \ + --prompt "What is 17 times 23?" --interval 4 --responses 4 \ + --max-tokens 4096 --temperature 1 --seed 42 --backend triton +``` + +Add `--eager` for the same-readout reference, or use `--interval 1` for native +four-loop readout. Repeat the same prompts, limits, seeds and sampling settings +for each comparison. The first token 0 or 2 ends a response. Outputs retain raw +token IDs, actual selected processed logprobs, readout depths, decoded text and +the `stop`/`length` finish reason. Length-limited answers remain in the results. + +The Python entry point accepts token-ID prompts and independent per-answer seeds: + +```python +from vllm_rlt.worker.deferred_ouro import generate_deferred + +outputs = generate_deferred( + model, [[3, 5, 7], [11, 4, 9]], [42, 43], interval=4, + max_tokens=64, temperature=1.0, backend="auto", +) +``` + +Sampling uses the engine's `Sampler` and `RngRegistry` with full vocabulary, +positive temperature and one request-local random stream. Reported logprobs +are the logarithm of the actual processed sampling probability. They describe +the selected mixed readout and its KV history, rather than a separately +recomputed four-loop policy. + +## State ownership and lifecycle + +A runner owns one fixed batch and one four-plane paged cache. Weights must remain +unchanged until it closes. Retained hidden rows have independent storage, so a +finished neighbor cannot keep a larger active-batch tensor alive. `release()` +discards a request's unread deep tail and frees its pages. The context manager +and generation helper release the cohort on completion or error. A recurrent +execution failure closes the runner; create a fresh runner to retry. + +KV reservation covers prompt and executed response inputs. The final sampled +token requires no extra forward position. Full capacity is reserved up front; +this is a batching experiment and does not compress or quantize KV. + +## Validation and unfinished qualification + +Run the causal and lifecycle checks with no GPU access: + +```bash +CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 python -m pytest -q tests/test_deferred_ouro.py +pre-commit run --all-files +``` + +Tests compare selected hidden states and every populated KV plane against the +independent functional dense oracle across intervals 1/2/4/8 and grouped-query +attention configurations. They also check the actual deep row/call counts, +staggered request progress, termination, probability alignment, independent RNG, +bounded hidden storage, failure cleanup and the last executable position. + +This experimental path uses the public native recurrent API and its attention +backend. The normal `LLM`/serving paths keep their existing feature selection. +The experiment does not integrate adaptive gating, speculative verification, +CUDA Graphs, PD, prefix reuse, preemption or weight publication. Those +combinations require separate implementations and validation. + +Real-weight BF16 task-quality, memory and end-to-end throughput comparisons, +plus separate ops-only/full GPU profiles, remain required before qualifying an +acceleration claim. Measure the eager/deferred pair with the same readout policy +and count all GPU work, failure and length-limit costs. Compare the mixed policy +with K4 on the same tasks under an explicit quality budget. The synchronous +Python sampling/output path and per-loop metadata preparation add overhead; +profiling must determine whether larger deep batches cover those costs. + +This supports an experimental execution path related to +[RFC #70](https://github.com/ThinkFlowLab/vllm-rlt/issues/70). It does not complete +training integration or differentiable current-policy probability replay. diff --git a/tests/reference/models.py b/tests/reference/models.py index 0b72ae4f..0002570d 100644 --- a/tests/reference/models.py +++ b/tests/reference/models.py @@ -6,12 +6,13 @@ @torch.inference_mode() -def dense_reference(model, token_ids, loops): +def dense_reference(model, token_ids, loops, *, return_kv=False): """Functional causal oracle: no model modules, paged cache, or recurrent API. Mirrors the equations of the pinned official model, including split-half RoPE, GQA, all four sandwich norms, and norm at each loop boundary. Returns every depth so intermediate hidden states and gate logits are checked. + With return_kv=True, also return the true pre-GQA K/V planes per loop/layer. """ config = model.config weights = model.state_dict() @@ -41,8 +42,9 @@ def rotary(value): rotated = torch.cat([-value[..., split:], value[..., :split]], dim=-1) return value * cos + rotated * sin - outputs = [] + outputs, kv_planes = [], [] for _ in range(loops): + layer_kv = [] for layer in range(config.num_hidden_layers): prefix = f"model.layers.{layer}" attention_input = norm(hidden, prefix + ".input_layernorm") @@ -59,6 +61,8 @@ def rotary(value): v = linear(attention_input, prefix + ".self_attn.v_proj").reshape( length, config.num_key_value_heads, config.head_dim ) + if return_kv: + layer_kv.append((k.clone(), v.clone())) groups = config.num_attention_heads // config.num_key_value_heads k = k.repeat_interleave(groups, dim=1) v = v.repeat_interleave(groups, dim=1) @@ -80,7 +84,9 @@ def rotary(value): linear(hidden, "model.early_exit_gate") + weights["model.early_exit_gate.bias"] ).squeeze(-1) outputs.append((hidden.clone(), gate, linear(hidden, "lm_head"))) - return outputs + if return_kv: + kv_planes.append(layer_kv) + return (outputs, kv_planes) if return_kv else outputs @torch.inference_mode() diff --git a/tests/test_deferred_ouro.py b/tests/test_deferred_ouro.py new file mode 100644 index 00000000..83da9c8b --- /dev/null +++ b/tests/test_deferred_ouro.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of deferred decoding against dense, causal Ouro equations.""" + +import pytest +import torch + +from tests.reference import dense_reference as functional_reference +from vllm_rlt.models.ouro import OuroConfig, OuroForCausalLM +from vllm_rlt.worker.deferred_ouro import DeferredOuroRunner, generate_deferred + + +def make_model(kv_heads=2, loops=4, max_positions=64): + torch.manual_seed(19) + return OuroForCausalLM( + OuroConfig( + vocab_size=31, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=kv_heads, + head_dim=8, + max_position_embeddings=max_positions, + total_ut_steps=loops, + ) + ).eval() + + +@torch.inference_mode() +def dense_reference(model, tokens): + outputs, planes = functional_reference(model, torch.tensor(tokens), 4, return_kv=True) + return [output[0] for output in outputs], planes + + +@pytest.mark.parametrize("interval", [1, 2, 4, 8]) +@pytest.mark.parametrize("kv_heads", [1, 2, 4]) +@pytest.mark.parametrize("defer", [False, True]) +@torch.inference_mode() +def test_readouts_and_all_committed_kv_match_dense(interval, kv_heads, defer): + model = make_model(kv_heads) + prompts = [[3, 5, 7], [11, 4, 9, 6, 8]] + prompt_lengths = [len(prompt) for prompt in prompts] + with DeferredOuroRunner( + model, [32, 32], interval=interval, backend="torch", defer=defer + ) as runner: + hidden = runner.prefill(prompts) + for index, prompt in enumerate(prompts): + expected, _ = dense_reference(model, prompt) + torch.testing.assert_close(hidden[index], expected[3][-1], atol=2e-5, rtol=2e-5) + for step in range(9): + tokens = [3 + step, 17 + step] + hidden, depths = runner.step([0, 1], tokens) + for index, token in enumerate(tokens): + prompts[index].append(token) + expected, planes = dense_reference(model, prompts[index]) + torch.testing.assert_close( + hidden[index], expected[depths[index] - 1][-1], atol=2e-5, rtol=2e-5 + ) + for depth in range(4): + for layer in range(model.config.num_hidden_layers): + actual = runner.cache.read(layer, runner.requests[index].request_id, depth) + committed = ( + prompt_lengths[index] + (step + 1) // interval * interval + if defer and depth >= 2 + else len(prompts[index]) + ) + for cached, reference in zip(actual, planes[depth][layer], strict=True): + assert len(cached) == committed + torch.testing.assert_close( + cached, reference[: len(cached)], atol=2e-5, rtol=2e-5 + ) + assert depths == ([4, 4] if 9 % interval == 0 else [2, 2]) + assert runner.cache.num_free_blocks == runner.cache.num_blocks + assert all(not request.pending_hidden and not request.live for request in runner.requests) + + +@torch.inference_mode() +def test_staggered_blocks_and_request_release(): + model = make_model() + with DeferredOuroRunner(model, [32, 32], interval=4, backend="torch") as runner: + runner.prefill([[3, 5], [7, 9, 11]]) + runner.step([0], [4]) + runner.step([0], [6]) + runner.step([0, 1], [8, 12]) + hidden, depths = runner.step([0, 1], [10, 14]) + assert depths == [4, 2] + expected, _ = dense_reference(model, [3, 5, 4, 6, 8, 10]) + torch.testing.assert_close(hidden[0], expected[3][-1], atol=2e-5, rtol=2e-5) + runner.release([0]) + runner.step([1], [16]) + hidden, depths = runner.step([1], [18]) + assert depths == [4] + expected, _ = dense_reference(model, [7, 9, 11, 12, 14, 16, 18]) + torch.testing.assert_close(hidden[0], expected[3][-1], atol=2e-5, rtol=2e-5) + + +@torch.inference_mode() +def test_stop_tokens_processed_scores_and_batch_shrink(monkeypatch): + model = make_model() + samples = iter([2, 7, 8, 0]) + monkeypatch.setattr(torch, "multinomial", lambda *args, **kwargs: torch.tensor([next(samples)])) + outputs = generate_deferred(model, [[3, 5], [4, 6]], [42, 43], max_tokens=8, backend="torch") + assert [output.token_ids for output in outputs] == [[2], [7, 8, 0]] + assert [output.finish_reason for output in outputs] == ["stop", "stop"] + assert [output.readout_depths for output in outputs] == [[4], [4, 2, 2]] + for prompt, output in zip([[3, 5], [4, 6]], outputs, strict=True): + for index, (token, depth) in enumerate( + zip(output.token_ids, output.readout_depths, strict=True) + ): + levels, _ = dense_reference(model, prompt + output.token_ids[:index]) + expected = model.coda(levels[depth - 1][-1]).float().log_softmax(-1)[token] + assert output.logprobs[index] == pytest.approx(float(expected), abs=2e-5) + + +@torch.inference_mode() +def test_length_limit_and_repeatable_rng(): + model = make_model() + args = dict( + prompts=[[3, 5], [4, 6]], seeds=[42, 43], max_tokens=7, stop_token_ids=(), backend="torch" + ) + first = generate_deferred(model, **args) + second = generate_deferred(model, **args) + assert first == second + assert all(len(output.token_ids) == 7 and output.finish_reason == "length" for output in first) + assert all(output.readout_depths == [4, 2, 2, 2, 4, 2, 2] for output in first) + + +def test_invalid_configuration_and_step_are_rejected(): + with pytest.raises(ValueError, match="four-loop"): + DeferredOuroRunner(make_model(loops=3), [16], backend="torch") + with pytest.raises(ValueError, match="interval"): + DeferredOuroRunner(make_model(), [16], interval=3, backend="torch") + with DeferredOuroRunner(make_model(), [3], backend="torch") as runner: + with pytest.raises(ValueError, match="Invalid request index"): + runner.release([-1]) + assert runner.requests[0].live + with pytest.raises(ValueError, match="after prefill"): + runner.step([0], [1]) + runner.prefill([[3, 4, 5]]) + with pytest.raises(ValueError, match="capacity"): + runner.step([0], [1]) + + +@pytest.mark.parametrize("interval", [1, 2, 4, 8]) +@torch.inference_mode() +def test_deferred_groups_deep_rows_without_reducing_committed_work(interval, monkeypatch): + records = [] + for defer in (False, True): + model = make_model() + calls = [] + original = model.recurrent + + def record(hidden, request_ids, depths, positions, cache, **kwargs): + calls.append((depths[0], len(request_ids))) + return original(hidden, request_ids, depths, positions, cache, **kwargs) + + monkeypatch.setattr(model, "recurrent", record) + with DeferredOuroRunner(model, [32, 32], interval=interval, defer=defer) as runner: + runner.prefill([[3, 5], [7, 9]]) + calls.clear() + for step in range(8): + runner.step([0, 1], [3 + step, 15 + step]) + records.append(calls) + for depth in (2, 3): + eager = [rows for loop, rows in records[0] if loop == depth] + deferred = [rows for loop, rows in records[1] if loop == depth] + assert eager == [2] * 8 + assert deferred == [2 * interval] * (8 // interval) + assert sum(eager) == sum(deferred) == 16 + + +@torch.inference_mode() +def test_pending_states_own_only_their_request_rows(): + model = make_model() + with DeferredOuroRunner(model, [16] * 12) as runner: + runner.prefill([[3, 5]] * 12) + runner.step(list(range(12)), [7] * 12) + runner.release(range(1, 12)) + retained = runner.requests[0].pending_hidden[0] + assert retained.untyped_storage().nbytes() == retained.numel() * retained.element_size() + assert retained.shape == (1, model.config.hidden_size) + + +def test_recurrent_failure_releases_allocations(monkeypatch): + model = make_model() + original = model.recurrent + + def fail(hidden, request_ids, depths, positions, cache, **kwargs): + if depths[0] == 2: + raise RuntimeError("injected recurrent failure") + return original(hidden, request_ids, depths, positions, cache, **kwargs) + + monkeypatch.setattr(model, "recurrent", fail) + with DeferredOuroRunner(model, [16, 16]) as runner: + with pytest.raises(RuntimeError, match="injected recurrent failure"): + runner.prefill([[3, 5], [7, 9]]) + assert runner.cache.num_free_blocks == runner.cache.num_blocks + assert all(not request.live and not request.pending_hidden for request in runner.requests) + + +def test_last_sample_needs_no_extra_kv_position(): + model = make_model(max_positions=16) + output = generate_deferred( + model, [[3, 5]], [42], max_tokens=15, stop_token_ids=(), backend="torch" + )[0] + assert len(output.token_ids) == 15 + assert output.finish_reason == "length" + with pytest.raises(ValueError, match="position limit"): + generate_deferred(model, [[3, 5]], [42], max_tokens=16, backend="torch") + + +def test_request_rng_is_independent_of_batch_members(): + model = make_model() + settings = dict(max_tokens=7, stop_token_ids=(), backend="torch") + grouped = generate_deferred(model, [[3, 5], [4, 6]], [42, 43], **settings) + for prompt, seed, combined in zip([[3, 5], [4, 6]], [42, 43], grouped, strict=True): + alone = generate_deferred(model, [prompt], [seed], **settings)[0] + assert combined.token_ids == alone.token_ids + assert combined.readout_depths == alone.readout_depths + assert combined.finish_reason == alone.finish_reason + torch.testing.assert_close( + torch.tensor(combined.logprobs), torch.tensor(alone.logprobs), atol=2e-5, rtol=2e-5 + ) diff --git a/vllm_rlt/entrypoints/ouro_deferred.py b/vllm_rlt/entrypoints/ouro_deferred.py new file mode 100644 index 00000000..5d2304bd --- /dev/null +++ b/vllm_rlt/entrypoints/ouro_deferred.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run the experimental same-depth-KV, fixed-interval Ouro inference path.""" + +import argparse +import json +from dataclasses import asdict + +import torch +from transformers import AutoTokenizer + +from vllm_rlt.models.ouro import OuroForCausalLM +from vllm_rlt.worker.deferred_ouro import generate_deferred + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model", help="Local Ouro checkpoint directory") + parser.add_argument("--prompt", required=True) + parser.add_argument("--responses", type=int, default=4) + parser.add_argument("--interval", type=int, choices=(1, 2, 4, 8), default=4) + parser.add_argument("--max-tokens", type=int, default=4096) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", default="cuda") + parser.add_argument( + "--eager", + action="store_true", + help="Execute the same readout policy with immediate deep writes", + ) + parser.add_argument("--backend", choices=("auto", "torch", "triton"), default="auto") + args = parser.parse_args() + if args.responses <= 0: + parser.error("--responses must be positive") + device = torch.device(args.device) + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False) + model = OuroForCausalLM.from_pretrained( + args.model, device=device, dtype=torch.bfloat16 if device.type == "cuda" else torch.float32 + ).eval() + prompt = tokenizer.encode( + args.prompt + "\nGive your final answer in \\boxed{}.\n###Response\n", + add_special_tokens=False, + ) + outputs = generate_deferred( + model, + [prompt] * args.responses, + [args.seed + index for index in range(args.responses)], + interval=args.interval, + max_tokens=args.max_tokens, + temperature=args.temperature, + backend=args.backend, + defer=not args.eager, + ) + print( + json.dumps( + [ + { + **asdict(output), + "text": tokenizer.decode(output.token_ids, skip_special_tokens=False), + } + for output in outputs + ], + ensure_ascii=False, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/vllm_rlt/worker/deferred_ouro.py b/vllm_rlt/worker/deferred_ouro.py new file mode 100644 index 00000000..b723fa70 --- /dev/null +++ b/vllm_rlt/worker/deferred_ouro.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Experimental fixed-interval Ouro decoding with exact per-loop KV histories. + +This inference-only runner groups loops three and four across response positions. +It groups genuine deep history without LAST_EXITED propagation copies. +Weights must remain unchanged until the runner is closed. +""" + +import math +from collections.abc import Sequence +from dataclasses import dataclass, field + +import torch + +from vllm_rlt.core.kv_cache_manager import KVCacheManager +from vllm_rlt.models.ouro import OuroForCausalLM +from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.worker.sampler import RngRegistry, Sampler + + +@dataclass +class _Request: + request_id: str + capacity: int + position: int = 0 + response_inputs: int = 0 + pending_hidden: list[torch.Tensor] = field(default_factory=list) + pending_positions: list[int] = field(default_factory=list) + live: bool = False + + +@dataclass +class DeferredOutput: + token_ids: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + readout_depths: list[int] = field(default_factory=list) + finish_reason: str = "length" + + +class DeferredOuroRunner: + """One fixed-weight batch; shallow readout is two and deep readout is four. + + ``defer=False`` is the same-readout eager reference, useful for comparisons. + Positions passed to attention are original request positions, including when + several requests and response blocks are concatenated into one deep batch. + """ + + def __init__( + self, + model: OuroForCausalLM, + capacities: Sequence[int], + *, + interval: int = 4, + backend: str = "auto", + defer: bool = True, + ): + if model.config.total_ut_steps != 4: + raise ValueError("Deferred Ouro decoding requires a four-loop model") + if type(interval) is not int or interval not in (1, 2, 4, 8): + raise ValueError("interval must be one of 1, 2, 4, 8") + if not capacities or any( + type(n) is not int or not 0 < n <= model.config.max_position_embeddings + for n in capacities + ): + raise ValueError("Provide positive capacities within the model's position limit") + self.model, self.interval, self.defer = model, interval, defer + parameter = next(model.parameters()) + self.device = parameter.device + self.prefilled = False + block_size = 16 + self.cache = KVCacheManager( + num_layers=model.config.num_hidden_layers, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + num_blocks=sum(math.ceil(n / block_size) * 4 for n in capacities), + block_size=block_size, + max_loops=4, + device=self.device, + dtype=parameter.dtype, + backend=backend, + layout="last_exited", + ) + self.requests = [_Request(f"deferred-{i}", n) for i, n in enumerate(capacities)] + for request in self.requests: + if not self.cache.allocate(request.request_id, request.capacity): + self.close() + raise RuntimeError("Insufficient KV pages for the request batch") + request.live = True + + def _run_loops(self, hidden, request_ids, positions, depths): + try: + for depth in depths: + hidden, _ = self.model.recurrent( + hidden, + request_ids, + [depth] * len(request_ids), + positions, + self.cache, + compute_gate=False, + ) + return hidden + except Exception: + self.close() + raise + + @torch.inference_mode() + def prefill(self, prompts: Sequence[Sequence[int]]) -> torch.Tensor: + if self.prefilled or not all(request.live for request in self.requests): + raise ValueError("A runner accepts one prefill batch") + if len(prompts) != len(self.requests) or any( + not prompt or len(prompt) > request.capacity + for prompt, request in zip(prompts, self.requests, strict=True) + ): + raise ValueError("Provide one nonempty prompt within each request capacity") + tokens, request_ids, positions, ends = [], [], [], [] + for prompt, request in zip(prompts, self.requests, strict=True): + tokens.extend(prompt) + request_ids.extend([request.request_id] * len(prompt)) + positions.extend(range(len(prompt))) + ends.append(len(tokens) - 1) + request.position = len(prompt) + hidden = self.model.prelude(torch.tensor(tokens, device=self.device, dtype=torch.long)) + hidden = self._run_loops(hidden, request_ids, positions, range(4)) + self.prefilled = True + return hidden[ends] + + @torch.inference_mode() + def step( + self, active: Sequence[int], token_ids: Sequence[int] + ) -> tuple[torch.Tensor, list[int]]: + if ( + not self.prefilled + or not active + or len(active) != len(token_ids) + or len(set(active)) != len(active) + ): + raise ValueError( + "Provide distinct active requests and one token per request after prefill" + ) + if any(type(i) is not int or not 0 <= i < len(self.requests) for i in active): + raise ValueError("Invalid active request index") + requests = [self.requests[i] for i in active] + if any(not request.live or request.position >= request.capacity for request in requests): + raise ValueError("Request is closed or has exhausted its KV capacity") + ids = [request.request_id for request in requests] + positions = [request.position for request in requests] + hidden = self.model.prelude(torch.tensor(token_ids, device=self.device, dtype=torch.long)) + hidden = self._run_loops(hidden, ids, positions, range(2)) + selected = hidden.clone() + readouts, deep_hidden, deep_ids, deep_positions, ends = [], [], [], [], [] + for row, request in enumerate(requests): + request.response_inputs += 1 + depth = 4 if request.response_inputs % self.interval == 0 else 2 + readouts.append(depth) + if self.defer: + request.pending_hidden.append(hidden[row : row + 1].clone()) + request.pending_positions.append(request.position) + if depth == 4: + deep_hidden.extend(request.pending_hidden) + deep_ids.extend([request.request_id] * len(request.pending_positions)) + deep_positions.extend(request.pending_positions) + ends.append((row, len(deep_positions) - 1)) + request.pending_hidden.clear() + request.pending_positions.clear() + else: + deep_hidden.append(hidden[row : row + 1]) + deep_ids.append(request.request_id) + deep_positions.append(request.position) + if depth == 4: + ends.append((row, len(deep_positions) - 1)) + request.position += 1 + if deep_hidden: + deep = self._run_loops(torch.cat(deep_hidden), deep_ids, deep_positions, range(2, 4)) + for row, end in ends: + selected[row] = deep[end] + return selected, readouts + + def release(self, active: Sequence[int]) -> None: + """Discard unread deep tails; never copy shallow KV into deep planes.""" + if any(type(index) is not int or not 0 <= index < len(self.requests) for index in active): + raise ValueError("Invalid request index for release") + for index in active: + request = self.requests[index] + if request.live: + request.pending_hidden.clear() + request.pending_positions.clear() + self.cache.free(request.request_id) + request.live = False + + def close(self) -> None: + self.release(range(len(self.requests))) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + +@torch.inference_mode() +def generate_deferred( + model: OuroForCausalLM, + prompts: Sequence[Sequence[int]], + seeds: Sequence[int], + *, + interval: int = 4, + max_tokens: int = 4096, + temperature: float = 1.0, + stop_token_ids: Sequence[int] = (0, 2), + backend: str = "auto", + defer: bool = True, +) -> list[DeferredOutput]: + """Full-vocabulary sampling with independent per-answer RNG and processed scores. + + The first output uses full four-loop prompt prefill. This is a mixed-readout + inference policy, not the native four-loop policy and not a training adapter. + """ + if ( + not prompts + or len(prompts) != len(seeds) + or any(type(seed) is not int or seed < 0 for seed in seeds) + ): + raise ValueError("Provide one nonnegative seed per nonempty prompt batch") + if ( + type(max_tokens) is not int + or max_tokens <= 0 + or not math.isfinite(temperature) + or temperature <= 0 + ): + raise ValueError("max_tokens and temperature must be positive") + params = [ + SamplingParams(max_tokens=max_tokens, temperature=temperature, seed=seed) for seed in seeds + ] + outputs = [DeferredOutput() for _ in prompts] + registry, sampler = RngRegistry(), Sampler() + with DeferredOuroRunner( + model, + [len(prompt) + max_tokens - 1 for prompt in prompts], + interval=interval, + backend=backend, + defer=defer, + ) as runner: + generators = [ + registry.acquire(request.request_id, parameter, runner.device) + for request, parameter in zip(runner.requests, params, strict=True) + ] + hidden = runner.prefill(prompts) + active, depths = list(range(len(prompts))), [4] * len(prompts) + for step in range(max_tokens): + logits = model.coda(hidden).float() + survivors, tokens, finished = [], [], [] + for row, index in enumerate(active): + probabilities = sampler.probabilities(logits[row], params[index]) + token = int(sampler.draw(probabilities, generators[index])) + output = outputs[index] + output.token_ids.append(token) + output.logprobs.append(float(probabilities[token].log())) + output.readout_depths.append(depths[row]) + if token in stop_token_ids: + output.finish_reason = "stop" + finished.append(index) + elif step + 1 < max_tokens: + survivors.append(index) + tokens.append(token) + runner.release(finished) + if not survivors: + break + hidden, depths = runner.step(survivors, tokens) + active = survivors + return outputs