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_engine.py b/tests/test_engine.py index 1ad00979..bb6e3007 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -1,10 +1,19 @@ import math +from dataclasses import replace import pytest import torch from tests.helpers import tiny_ouro_config -from vllm_rlt import LLM, CacheConfig, SamplingParams, SchedulerConfig +from vllm_rlt import ( + LLM, + CacheConfig, + ExecutionConfig, + ExitConfig, + LoopCDParams, + SamplingParams, + SchedulerConfig, +) from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM from vllm_rlt.request import FinishReason, 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() diff --git a/tests/test_loopcd.py b/tests/test_loopcd.py new file mode 100644 index 00000000..5f7abd78 --- /dev/null +++ b/tests/test_loopcd.py @@ -0,0 +1,384 @@ +"""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, RequestPhase +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("implementation", ["two_head", "linear_fused"]) +def test_padded_coda_projects_only_real_requests(implementation): + model = model_for() + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=256, block_size=2), + scheduler_config=SchedulerConfig(max_num_seqs=3, max_num_batched_tokens=4), + execution_config=ExecutionConfig( + loopcd=True, static_buffers=True, pad_to_power_of_two=True + ), + ) + rows = [] + hook = model.lm_head.register_forward_pre_hook( + lambda _module, inputs: rows.append(len(inputs[0])) + ) + for index in range(3): + engine.add_request( + str(index), + [4 + index], + replace(params_for(), max_tokens=1, loopcd=LoopCDParams(implementation=implementation)), + ) + try: + while engine.has_unfinished_requests(): + engine.step() + finally: + hook.remove() + assert rows == ([3, 3] if implementation == "two_head" else [3]) + assert not engine.model_runner.loopcd_references + assert engine.cache_manager.num_used_blocks == 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize("implementation", ["two_head", "linear_fused"]) +def test_loopcd_coda_graphs_match_eager_through_request_reuse(implementation): + torch.manual_seed(42) + model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.bfloat16).eval() + results = [] + for graphs in (False, True): + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=256, block_size=4), + scheduler_config=SchedulerConfig(max_num_seqs=3, max_num_batched_tokens=4), + execution_config=ExecutionConfig( + loopcd=True, + static_buffers=True, + pad_to_power_of_two=True, + cuda_graphs=graphs, + cuda_graph_max_graphs=1, + ), + attention_backend="triton", + ) + rounds = [] + for _ in range(2): + for index in range(3): + engine.add_request( + str(index), + [4 + index, 7, 3], + replace( + params_for(), + max_tokens=2 + index, + loopcd=LoopCDParams(implementation=implementation), + ), + ) + finished = {} + while engine.has_unfinished_requests(): + for output in engine.step(): + if output.finished: + finished[output.request_id] = (output.token_ids, output.exit_depths) + assert len(finished) == 3 and engine.cache_manager.num_used_blocks == 0 + assert not engine.model_runner.loopcd_references + rounds.append(finished) + if graphs: + coda = engine.model_runner.coda_graphs + # Power-of-two padding lets the one bucketed coda graph cover every + # batch here, so no row count falls back to eager. + assert coda.captures == 1 and coda.replays > 1 + results.append(rounds) + assert results[0] == results[1] + + +@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.phase != RequestPhase.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, params, generator=None): + seen.append(logits.clone()) + return logits.argmax() + + engine.model_runner.sampler.sample = 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", "split", "stage"] +) +def test_unsupported_admission_rejected_before_request_allocation(case): + from vllm_rlt.engine.stage_config import StageConfig + + 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", wavefront_chunked_prefill=case == "split" + ), + execution_config=ExecutionConfig(loopcd=case != "unreserved"), + stage_config=StageConfig(role="prefill" if case == "stage" else "prefill_decode"), + ) + 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", "split"]) +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", + ), + scheduler_config=SchedulerConfig(wavefront_chunked_prefill=feature == "split"), + 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): + # CPU-only: this test must not initialize CUDA. A --run-gpu session may + # already have initialized it in an earlier test of the same process. + cuda_was_initialized = torch.cuda.is_initialized() + model = model_for() + engine = LLMEngine( + model, + 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, params, generator=None): + request = engine.scheduler.requests["a"] + 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.sampler.sample = observe + for run in range(2): + if run: + engine.add_request("a", [12, 3, 7], sampling) + interrupted = engine.scheduler.requests["a"] + while interrupted.phase != RequestPhase.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 torch.cuda.is_initialized() == cuda_was_initialized diff --git a/tests/test_loopcd_continuation.py b/tests/test_loopcd_continuation.py new file mode 100644 index 00000000..c3645bab --- /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 RequestPhase + + +@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, params, generator=None): + seen.append(logits.detach().clone()) + return torch.tensor(continuation[len(seen) - 1], device=logits.device) + + engine.model_runner.sampler.sample = 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.phase != RequestPhase.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/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 648781d7..02e044f9 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -151,9 +151,14 @@ class ExecutionConfig: # Select packed prefill metadata staging (and require packed attention). # The historical name is retained; use_uva controls its transfer mechanism. 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", "use_uva", "async_scheduling", diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 12aa040c..ac9e097c 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -1067,7 +1067,9 @@ def _apply_prefill_progress( request.request_id, request.prompt_token_ids, request.num_prefilled_tokens, event ) if final: - request.loops_done = self.total_ut_steps + # The request's prompt-depth frontier may be shallower than + # the model's maximum decode depth. + request.loops_done = len(request.prefill_depth_ranges) or self.total_ut_steps self.enqueue(request, self._prompt_destination(request) if final else RequestPhase.PREFILL) return PrefillProgress(True, final) diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..6c46bb74 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -21,6 +21,7 @@ from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase from vllm_rlt.sampling_params import SamplingParams from vllm_rlt.worker.executor import Executor +from vllm_rlt.worker.loopcd import active, reserved_bytes, validate from vllm_rlt.worker.model_runner import ModelRunner from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput from vllm_rlt.worker.sampler import RngRegistry @@ -73,6 +74,20 @@ def __init__( staged_prefill = requires_staged_prefill( scheduler_config, produces_tokens=self.stage_config.produces_tokens ) + 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 ( + staged_prefill + or 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 atomic synchronous last_exited " + "without prefix/speculation" + ) if getattr(model, "requires_boundary_kv", False): if staged_prefill: raise ValueError( @@ -178,6 +193,9 @@ def __init__( speculative=speculative_config, staged_prefill=staged_prefill, ) + 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, @@ -325,6 +343,19 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i 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 active(params) and ( + self.connector is not None or self.stage_config.role != "prefill_decode" + ): + raise ValueError("LoopCD P/D reference transfer is not implemented") if self.speculative_config is not None and ( max_loops != self.speculative_config.target_loops or params.exit_threshold != 1.0 ): @@ -342,7 +373,7 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i 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] @@ -351,6 +382,15 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i raise ValueError( "exit trace must cover every output and obey prefill/decode bounds" ) + # The first output comes from prefill. Later tokens require prompt KV + # history at every recurrent depth they can visit. + if params.max_tokens > 1: + decode_depth = max(trace[1 : params.max_tokens]) if trace else max_loops + if decode_depth > self.prefill_depth: + raise ValueError( + f"decode depth {decode_depth} exceeds prefill_depth {self.prefill_depth}; " + "increase prefill_depth or reduce decode depth" + ) capacity = len(prompt_token_ids) + params.max_tokens - 1 if capacity > config.max_position_embeddings: raise ValueError("prompt plus decode positions exceed the model context length") @@ -370,7 +410,7 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i # Prefill always runs the model's full depth, whatever decode # loop bound the request asks for. prefill_depth_ranges=Request.prefill_boundaries( - len(prompt_token_ids), config.total_ut_steps + len(prompt_token_ids), self.prefill_depth ), ) ) diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index 47f58ad1..33e88283 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -275,6 +275,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 self.speculative_config is not None: if ( diff --git a/vllm_rlt/sampling_params.py b/vllm_rlt/sampling_params.py index 30009bde..b534909d 100644 --- a/vllm_rlt/sampling_params.py +++ b/vllm_rlt/sampling_params.py @@ -2,6 +2,42 @@ from dataclasses import dataclass +@dataclass(frozen=True) +class LoopCDParams: + """Fixed guidance after complete one-based recurrent updates. + + Nonzero guidance requires atomic synchronous prefill and local execution; + split prefill and cross-stage reference transport are not supported. + """ + + 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 +50,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 2829b0d3..dc764736 100644 --- a/vllm_rlt/worker/buffers.py +++ b/vllm_rlt/worker/buffers.py @@ -5,6 +5,7 @@ import torch +from vllm_rlt.worker.loopcd import reserved_bytes from vllm_rlt.worker.metadata import KVLayout, KVMetadata from vllm_rlt.worker.submission_queue import MAX_CONCURRENT_SUBMISSIONS, SubmissionQueue from vllm_rlt.worker.tensors import UvaBackedTensor, copy_tensor, use_uva_on @@ -33,7 +34,7 @@ def execution_buffer_bytes( 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 buffer sets. total += 4 * rows * (hidden + 4 * width + 36) + scheduler.max_num_seqs * hidden diff --git a/vllm_rlt/worker/executor.py b/vllm_rlt/worker/executor.py index bad146d4..e2a8183b 100644 --- a/vllm_rlt/worker/executor.py +++ b/vllm_rlt/worker/executor.py @@ -101,6 +101,12 @@ def submit(self, scheduled_batch): raise ValueError("speculative submission requires a configured round controller") return self.speculation.submit(scheduled_batch) batch = build_execution_batch(scheduled_batch, self.runner.model.config.total_ut_steps) + if ( + batch.operation == Operation.PRELUDE + and batch.loop_count + and batch.loop_count != self.runner.prefill_depth + ): + batch = replace(batch, loop_count=self.runner.prefill_depth) return self.runner.submit(scheduled_batch, batch=batch) diff --git a/vllm_rlt/worker/loopcd.py b/vllm_rlt/worker/loopcd.py new file mode 100644 index 00000000..35121c81 --- /dev/null +++ b/vllm_rlt/worker/loopcd.py @@ -0,0 +1,72 @@ +"""Readout-only fixed LoopCD for Ouro; original recurrent feedback is preserved.""" + +from dataclasses import dataclass + +import torch + +from vllm_rlt.core.scheduling_policy import requires_staged_prefill +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 ( + requires_staged_prefill(scheduler) + or 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 atomic 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 ada23d20..2a968b14 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -17,7 +17,8 @@ stage_metadata, ) from vllm_rlt.worker.cuda_graph import CodaGraphs, RecurrentGraphs -from vllm_rlt.worker.executor import ForwardBatch, ModelRuntime +from vllm_rlt.worker.executor import ForwardBatch, ForwardResult, ModelRuntime +from vllm_rlt.worker.loopcd import Reference, active, extrapolate from vllm_rlt.worker.output import ModelRunnerOutput from vllm_rlt.worker.sampler import RngRegistry, Sampler from vllm_rlt.worker.submission_queue import MAX_CONCURRENT_SUBMISSIONS, SubmissionQueue @@ -173,6 +174,21 @@ def __init__( self.rng = rng if rng is not None else RngRegistry() 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( @@ -185,8 +201,16 @@ def __init__( else None ) self.coda_graphs = ( - CodaGraphs(model, self.execution_config, cache_manager) - if self.graphs is not None and getattr(model, "requires_boundary_kv", False) + CodaGraphs( + model, + self.execution_config, + cache_manager if getattr(model, "requires_boundary_kv", False) else None, + ) + if self.graphs is not None + and ( + getattr(model, "requires_boundary_kv", False) + or (self.execution_config.loopcd and not self.execution_config.async_scheduling) + ) else None ) self.input_buffers: dict[torch.dtype, StagedInputTensor] = { @@ -614,6 +638,142 @@ def _gather(self, requests, workspace, size): hidden[row].copy_(request.hidden_state) return hidden + 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 _coda(self, hidden): + forward = ForwardBatch(Operation.CODA, hidden=hidden, graphs=self.coda_graphs) + return self.runtime.forward(self.runtime.prepare(forward)).logits + + def _loopcd_readout(self, hidden, requests, positions, indices): + # Padding is needed by the core, but never by the output projections. + hidden = hidden[: len(requests)] + if not self.execution_config.loopcd: + return self._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._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._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 _forward_loopcd(self, batch, prepared, entry): + """Observe original recurrence boundaries without changing recurrent feedback.""" + forward = entry.forward + if forward.operation == Operation.CODA: + if not forward.compute_logits or forward.prelude_coda_kv_metadata is not None: + return self.runtime.forward(forward) + return ForwardResult( + hidden=forward.hidden, + logits=self._loopcd_readout( + forward.hidden, + [span.request for span in batch.spans], + prepared.positions, + prepared.output_indices, + ), + ) + if forward.operation == Operation.PRELUDE and batch.loop_count: + requests = entry.requests + positions = [ + position + for span in batch.spans + for position in range(span.token_start, span.token_start + span.token_count) + ] + # LoopCD is admitted only for atomic LAST_EXITED prefill, whose rows + # retain span order. Run the same prepared calls one depth at a time + # so the selected reference is captured before a later loop overwrites it. + result = self.runtime.forward( + replace(forward, metadata=(), graphs=None, graph_calls=()) + ) + for depth, kv in enumerate(forward.metadata): + recurrent = ForwardBatch(Operation.RECURRENT, hidden=result.hidden, metadata=(kv,)) + result = self.runtime.forward(recurrent) + self._count_core(len(requests), len(result.hidden), prefill=True) + self._capture_loopcd( + requests, + result.hidden, + [depth] * len(requests), + positions, + [0] * len(requests), + prefill=True, + ) + return result + if forward.operation == Operation.PRELUDE: + for span in batch.spans: + self.loopcd_references.pop(span.request.request_id, None) + result = self.runtime.forward(forward) + if forward.operation == Operation.RECURRENT: + self._count_core(batch.num_rows, len(result.hidden)) + self._capture_loopcd( + [span.request for span in batch.spans], + result.hidden, + prepared.depths, + prepared.positions, + prepared.output_indices, + ) + return result + def _prepare_prefill_coda(self, ids, positions, requests, kv=None): if not getattr(self.model, "requires_boundary_kv", False): return None, None @@ -1103,7 +1263,11 @@ def _execute_batch(self, batch: ExecutionBatch, prepared): forward = entry.forward if entry.hidden_rows is not None: forward = replace(forward, hidden=hidden[entry.hidden_rows]) - result = self.runtime.forward(forward) + result = ( + self._forward_loopcd(batch, prepared, replace(entry, forward=forward)) + if self.execution_config.loopcd + else self.runtime.forward(forward) + ) hidden = result.hidden if forward.operation in (Operation.PRELUDE, Operation.RECURRENT): if entry.outputs is None: @@ -1364,6 +1528,7 @@ def _free_runner_slot(self, request_id): event = self.events.pop(request_id, None) if event is not None: event.synchronize() + self.loopcd_references.pop(request_id, None) self.prompt_token_owners.pop(request_id, None) if self.async_state is not None: self.async_state.release(request_id)