Add initial LoopCD support for Ouro and correctness checks

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

tests/test_engine.py

@@ -1,10 +1,19 @@
 import math
+from dataclasses import replace
 
 import pytest
 import torch
 
 from tests.helpers import tiny_ouro_config
-from vllm_rlt import LLM, CacheConfig, SamplingParams, SchedulerConfig
+from vllm_rlt import (
+    LLM,
+    CacheConfig,
+    ExecutionConfig,
+    ExitConfig,
+    LoopCDParams,
+    SamplingParams,
+    SchedulerConfig,
+)
 from vllm_rlt.engine.llm_engine import LLMEngine
 from vllm_rlt.models import OuroForCausalLM
 from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase
@@ -415,3 +424,140 @@ def test_finish_reason_preserves_public_string_values(reason):
     output = RequestOutput.from_request(request)
     assert output.finished and type(output.finish_reason) is str
     assert output.finish_reason == reason.value
+
+
+@pytest.mark.parametrize("strength", [None, 0.0, 0.3])
+@pytest.mark.parametrize("max_loops", [3, None])
+@torch.inference_mode()
+def test_decode_deeper_than_prefill_rejected_before_admission(strength, max_loops):
+    guidance = None if strength is None else LoopCDParams(strength=strength)
+    engine = LLMEngine(
+        tiny_model(),
+        cache_config=CacheConfig(num_blocks=128, block_size=4),
+        execution_config=ExecutionConfig(loopcd=bool(strength), prefill_depth=2),
+    )
+    params = SamplingParams(
+        max_tokens=3,
+        min_loops=max_loops or 4,
+        max_loops=max_loops,
+        ignore_eos=True,
+        loopcd=guidance,
+    )
+    with pytest.raises(ValueError, match="decode depth .* exceeds prefill_depth 2"):
+        engine.add_request("r", [4, 7, 9], params)
+    assert not engine.scheduler.requests
+    assert not engine.has_unfinished_requests()
+    assert engine.cache_manager.num_used_blocks == 0
+    assert not engine.model_runner.loopcd_references
+
+    # Rejected admission must leave the ID available for a valid continuation.
+    engine.add_request("r", [4, 7, 9], replace(params, min_loops=2, max_loops=2))
+    outputs, _ = drain(engine)
+    assert outputs["r"].exit_depths == [2, 2, 2]
+    assert len(outputs["r"].token_ids) == 3
+    assert engine.cache_manager.num_used_blocks == 0
+    assert not engine.model_runner.loopcd_references
+    engine.close()
+
+
+@pytest.mark.parametrize("strength", [None, 0.0, 0.3])
+@torch.inference_mode()
+def test_prefill_history_allows_shallower_decode(strength):
+    engine = LLMEngine(
+        tiny_model(),
+        cache_config=CacheConfig(num_blocks=128, block_size=4),
+        execution_config=ExecutionConfig(loopcd=bool(strength), prefill_depth=4),
+    )
+    engine.add_request(
+        "r",
+        [4, 7, 9],
+        SamplingParams(
+            max_tokens=3,
+            min_loops=3,
+            max_loops=3,
+            ignore_eos=True,
+            loopcd=None if strength is None else LoopCDParams(strength=strength),
+        ),
+    )
+    outputs, _ = drain(engine)
+    assert outputs["r"].exit_depths == [4, 3, 3]
+    assert len(outputs["r"].token_ids) == 3
+    assert engine.cache_manager.num_used_blocks == 0
+    assert not engine.model_runner.loopcd_references
+    engine.close()
+
+
+@pytest.mark.parametrize("mode", ["ouro", "ouro_delayed", "random_lookahead"])
+def test_prefill_history_checks_adaptive_decode_bound(mode):
+    engine = LLMEngine(
+        tiny_model(),
+        execution_config=ExecutionConfig(prefill_depth=2),
+        exit_config=ExitConfig(mode=mode),
+    )
+    params = SamplingParams(max_tokens=3, min_loops=1, max_loops=4, exit_threshold=0.7)
+    with pytest.raises(ValueError, match="decode depth 4 exceeds prefill_depth 2"):
+        engine.add_request("r", [4, 7, 9], params)
+    assert not engine.scheduler.requests
+    assert engine.cache_manager.num_used_blocks == 0
+    engine.close()
+
+
+@pytest.mark.parametrize("trace", [[2, 3, 2], [2, 2, 3]])
+def test_prefill_history_rejects_deeper_trace_before_admission(trace):
+    engine = LLMEngine(
+        tiny_model(),
+        execution_config=ExecutionConfig(prefill_depth=2),
+        exit_config=ExitConfig(mode="trace", depths_by_request={"profile": trace}),
+    )
+    params = SamplingParams(max_tokens=3, min_loops=1, max_loops=4)
+    with pytest.raises(ValueError, match="decode depth 3 exceeds prefill_depth 2"):
+        engine.add_request("r", [4, 7, 9], params, trace_id="profile")
+    assert not engine.scheduler.requests
+    assert engine.cache_manager.num_used_blocks == 0
+    engine.close()
+
+
+@torch.inference_mode()
+def test_prefill_history_uses_only_requested_trace_depths():
+    engine = LLMEngine(
+        tiny_model(),
+        execution_config=ExecutionConfig(prefill_depth=2),
+        exit_config=ExitConfig(mode="trace", depths_by_request={"profile": [2, 1, 2, 4]}),
+    )
+    engine.add_request(
+        "r",
+        [4, 7, 9],
+        SamplingParams(max_tokens=3, min_loops=1, ignore_eos=True),
+        trace_id="profile",
+    )
+    outputs, _ = drain(engine)
+    assert outputs["r"].exit_depths == [2, 1, 2]
+    assert len(outputs["r"].token_ids) == 3
+    assert engine.cache_manager.num_used_blocks == 0
+    engine.close()
+
+
+@pytest.mark.parametrize("strength", [None, 0.0, 0.3])
+@torch.inference_mode()
+def test_prefill_only_output_does_not_require_decode_history(strength):
+    engine = LLMEngine(
+        tiny_model(),
+        execution_config=ExecutionConfig(loopcd=bool(strength), prefill_depth=2),
+    )
+    engine.add_request(
+        "r",
+        [4, 7, 9],
+        SamplingParams(
+            max_tokens=1,
+            min_loops=4,
+            max_loops=4,
+            ignore_eos=True,
+            loopcd=None if strength is None else LoopCDParams(strength=strength),
+        ),
+    )
+    outputs, _ = drain(engine)
+    assert outputs["r"].exit_depths == [2]
+    assert len(outputs["r"].token_ids) == 1
+    assert engine.cache_manager.num_used_blocks == 0
+    assert not engine.model_runner.loopcd_references
+    engine.close()