From bbcf135f84ff95e11d65d7a626a44100b2343362 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 15:52:07 +0800
Subject: [PATCH 01/10] Release one-shot candidate state storage

---
 src/rosa/__init__.py | 24 +++++++-----
 tests/test_rosa.py   | 87 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 9 deletions(-)

diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py
index 8d37823..cfa6c9b 100644
--- a/src/rosa/__init__.py
+++ b/src/rosa/__init__.py
@@ -420,15 +420,21 @@ def _build_stateful_hard_candidates(
         suffix_k=suffix_k,
         occurrences_r=occurrences_r,
     )
-    candidates = prefill_candidates(state, tokens)
-    result = HardCandidates(
-        *(getattr(candidates, name) for name in HardCandidates.__dataclass_fields__)
-    )
-    if not squeeze:
-        return result
-    return HardCandidates(
-        *(getattr(result, name)[0] for name in result.__dataclass_fields__)
-    )
+    try:
+        candidates = prefill_candidates(state, tokens)
+        result = HardCandidates(
+            *(
+                getattr(candidates, name)
+                for name in HardCandidates.__dataclass_fields__
+            )
+        )
+        if not squeeze:
+            return result
+        return HardCandidates(
+            *(getattr(result, name)[0] for name in result.__dataclass_fields__)
+        )
+    finally:
+        state.native_state = None
 
 
 def _build_forward_hard_candidates(
diff --git a/tests/test_rosa.py b/tests/test_rosa.py
index b3cbbbe..abc7fc5 100644
--- a/tests/test_rosa.py
+++ b/tests/test_rosa.py
@@ -1,12 +1,15 @@
 from __future__ import annotations
 
 import copy
+import gc
 import random
 import threading
 import unittest
+import weakref
 from concurrent.futures import ThreadPoolExecutor
 from unittest.mock import patch
 
+import numpy as np
 import torch
 import torch.nn.functional as F
 
@@ -942,6 +945,90 @@ def capture_prefill(state, full_tokens):
         for name in hard.__dataclass_fields__:
             self.assertIs(getattr(hard, name), getattr(captured, name), name)
 
+    def test_stateful_one_shot_releases_native_state_and_storage(self) -> None:
+        try:
+            import rosa_native_step
+        except ModuleNotFoundError:
+            self.skipTest("native companion is unavailable")
+        if getattr(rosa_native_step, "NativeCandidateState", None) is None:
+            self.skipTest("native candidate backend is unavailable")
+
+        from rosa._stateful_candidates_numba import (
+            init_candidate_state,
+            prefill_candidates,
+        )
+
+        tokens = torch.tensor([[0, 1, 0, 2, 0, 1]], dtype=torch.long)
+        expected = build_hard_candidates(tokens, suffix_k=3, occurrences_r=2)
+        state_ref = None
+        array_refs = []
+        native_used = False
+
+        def tracked_initialize(*args, **kwargs):
+            nonlocal state_ref, array_refs
+            state = init_candidate_state(*args, **kwargs)
+            state_ref = weakref.ref(state)
+            array_refs = [
+                weakref.ref(value)
+                for value in vars(state).values()
+                if isinstance(value, np.ndarray)
+            ]
+            return state
+
+        def tracked_prefill(state, full_tokens):
+            nonlocal native_used
+            candidates = prefill_candidates(state, full_tokens)
+            native_used = state.native_state not in (None, False)
+            return candidates
+
+        with (
+            patch(
+                "rosa._stateful_candidates_numba.init_candidate_state",
+                side_effect=tracked_initialize,
+            ),
+            patch(
+                "rosa._stateful_candidates_numba.prefill_candidates",
+                side_effect=tracked_prefill,
+            ),
+        ):
+            actual = rosa._build_stateful_hard_candidates(tokens, 3, 2)
+
+        self.assertTrue(native_used)
+        for name in expected.__dataclass_fields__:
+            self.assertTrue(
+                torch.equal(getattr(actual, name), getattr(expected, name)), name
+            )
+        gc.collect()
+        assert state_ref is not None
+        self.assertIsNone(state_ref())
+        self.assertTrue(array_refs)
+        self.assertTrue(all(array_ref() is None for array_ref in array_refs))
+
+    def test_stateful_one_shot_detaches_native_state_on_exception(self) -> None:
+        from rosa._stateful_candidates_numba import init_candidate_state
+
+        state = init_candidate_state(1, 3, suffix_k=2, occurrences_r=2)
+
+        class NativeOwner:
+            def __init__(self, candidate_state):
+                self.candidate_state = candidate_state
+
+        native = NativeOwner(state)
+        state.native_state = native
+        with (
+            patch(
+                "rosa._stateful_candidates_numba.init_candidate_state",
+                return_value=state,
+            ),
+            patch(
+                "rosa._stateful_candidates_numba.prefill_candidates",
+                side_effect=RuntimeError("prefill failed"),
+            ),
+            self.assertRaisesRegex(RuntimeError, "prefill failed"),
+        ):
+            rosa._build_stateful_hard_candidates(torch.tensor([[0, 1, 0]]), 2, 2)
+        self.assertIsNone(state.native_state)
+
     def test_stateful_full_sequence_preserves_scalar_batch_squeeze(self) -> None:
         tokens = torch.tensor([0, 1, 0, 2, 0, 1], dtype=torch.long)
         expected = _build_forward_hard_candidates(tokens, 3, 2, "python")

From 06111b39d0b7d8774e4201052343be32f708b994 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 16:03:18 +0800
Subject: [PATCH 02/10] Close persistent inference states deterministically

---
 src/rosa/__init__.py                   | 42 +++++++++++++++++---
 src/rosa/_stateful_candidates_numba.py | 25 ++++++++++++
 src/rosa/_stateful_numba.py            | 17 ++++++++-
 src/rosa/ragged.py                     | 18 +++++++++
 tests/test_ragged.py                   |  9 +++++
 tests/test_stateful_candidates.py      | 31 +++++++++++++++
 tests/test_unified_inference.py        | 53 ++++++++++++++++++++++++++
 7 files changed, 188 insertions(+), 7 deletions(-)

diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py
index cfa6c9b..9085555 100644
--- a/src/rosa/__init__.py
+++ b/src/rosa/__init__.py
@@ -434,7 +434,7 @@ def _build_stateful_hard_candidates(
             *(getattr(result, name)[0] for name in result.__dataclass_fields__)
         )
     finally:
-        state.native_state = None
+        state.close()
 
 
 def _build_forward_hard_candidates(
@@ -1432,26 +1432,34 @@ class ROSAInferenceState:
     ragged: bool
     suffix_k: int
     occurrences_r: int
-    _impl: object = field(repr=False)
+    _impl: object | None = field(repr=False)
+
+    def _require_impl(self) -> object:
+        impl = self._impl
+        if impl is None:
+            raise RuntimeError("state is closed")
+        return impl
 
     @property
     def position(self) -> int:
         """Number of tokens consumed by every batch row."""
 
+        impl = self._require_impl()
         if self.ragged:
             raise AttributeError(
                 "position is undefined for ragged states; use positions"
             )
-        return int(cast(Any, self._impl).position)
+        return int(cast(Any, impl).position)
 
     @property
     def positions(self) -> Tensor:
         """Consumed-token counts for every row, always returned as a copy."""
 
+        impl = self._require_impl()
         if self.ragged:
             if self.mode == "top1":
-                return cast(Any, self._impl).positions.clone()
-            return torch.from_numpy(cast(Any, self._impl).positions.copy())
+                return cast(Any, impl).positions.clone()
+            return torch.from_numpy(cast(Any, impl).positions.copy())
         return torch.full((self.batch_size,), self.position, dtype=torch.long)
 
     def step(
@@ -1462,6 +1470,7 @@ def step(
     ) -> InferenceOutput:
         """Consume one token per selected row using the configured mode."""
 
+        self._require_impl()
         return _inference_step(self, tokens, active=active, reset=reset)
 
     def step_into(self, tokens: Tensor, buffers: object) -> InferenceOutput:
@@ -1470,16 +1479,37 @@ def step_into(self, tokens: Tensor, buffers: object) -> InferenceOutput:
         Candidate tensors in the result alias ``buffers`` until its next use.
         """
 
+        self._require_impl()
         return _inference_step_into(self, tokens, buffers)
 
     def prefill(self, tokens: Tensor) -> InferenceOutput:
         """Consume an initial dense context and return every step result."""
 
+        self._require_impl()
         return _inference_prefill(self, tokens)
 
+    def close(self) -> None:
+        """Release persistent backend storage; repeated calls are safe."""
+
+        impl = self._impl
+        if impl is None:
+            return
+        self._impl = None
+        close = getattr(impl, "close", None)
+        if close is not None:
+            close()
+
+    def __enter__(self) -> ROSAInferenceState:
+        self._require_impl()
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.close()
+
     def reset(self) -> None:
-        """Reset all batch rows while retaining the configured capacity."""
+        """Replace storage with a fresh open state, reopening a closed state."""
 
+        self.close()
         self._impl = _make_inference_impl(
             self.batch_size,
             self.max_length,
diff --git a/src/rosa/_stateful_candidates_numba.py b/src/rosa/_stateful_candidates_numba.py
index 8cf04b4..061bd07 100644
--- a/src/rosa/_stateful_candidates_numba.py
+++ b/src/rosa/_stateful_candidates_numba.py
@@ -989,6 +989,28 @@ class CandidateState:
     size: np.ndarray
     edge_count: np.ndarray
     native_state: Any
+    _closed: bool = field(default=False, init=False, repr=False)
+
+    def __getattribute__(self, name: str) -> Any:
+        if name in {"position", "positions"} and object.__getattribute__(
+            self, "_closed"
+        ):
+            raise RuntimeError("state is closed")
+        return object.__getattribute__(self, name)
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("state is closed")
+
+    def close(self) -> None:
+        """Release native ownership deterministically; repeated calls are safe."""
+
+        if self._closed:
+            return
+        self._closed = True
+        # NativeCandidateState retains this Python state. Detaching this end of
+        # the relationship breaks the cycle without requiring a native API.
+        self.native_state = None
 
 
 @dataclass(frozen=True)
@@ -1048,6 +1070,7 @@ def init_candidate_buffers(
 
     if not isinstance(state, CandidateState):
         raise TypeError("state must be a CandidateState")
+    state._ensure_open()
     if sequence_length is not None and sequence_length < 0:
         raise ValueError("sequence_length must be >= 0")
     slots = state.suffix_k * state.occurrences_r
@@ -1188,6 +1211,7 @@ def _validate_candidate_tokens(
 ) -> tuple[Tensor, bool]:
     if not isinstance(state, CandidateState):
         raise TypeError("state must be a CandidateState")
+    state._ensure_open()
     if not isinstance(tokens, Tensor):
         raise TypeError("tokens must be a Tensor")
     scalar = tokens.ndim == 0 and state.batch_size == 1 and not sequence
@@ -1624,6 +1648,7 @@ def reset_candidates_masked(state: CandidateState, reset: Tensor) -> None:
 
     if not isinstance(state, CandidateState):
         raise TypeError("state must be a CandidateState")
+    state._ensure_open()
     if not state.ragged_mode:
         raise RuntimeError("reset_masked requires a ragged candidate state")
     if not isinstance(reset, Tensor):
diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py
index 9ae3ba3..a079c42 100644
--- a/src/rosa/_stateful_numba.py
+++ b/src/rosa/_stateful_numba.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import numpy as np
@@ -1121,6 +1121,19 @@ class _StatefulInferenceState:
     size: np.ndarray
     edge_count: np.ndarray
     native_state: Any
+    _closed: bool = field(default=False, init=False, repr=False)
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("state is closed")
+
+    def close(self) -> None:
+        """Break the optional native ownership cycle idempotently."""
+
+        if self._closed:
+            return
+        self._closed = True
+        self.native_state = None
 
 
 def _init_inference_state(
@@ -1214,6 +1227,7 @@ def _native_prefill(  # pragma: no cover - optional native companion
 def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor:
     """Consume one token per batch row and return exact top-1 predictions."""
 
+    state._ensure_open()
     if tokens.ndim == 0 and state.batch_size == 1:
         tokens = tokens.unsqueeze(0)
     if tokens.ndim != 1 or tokens.shape[0] != state.batch_size:
@@ -1256,6 +1270,7 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor:
 def _prefill(state: _StatefulInferenceState, tokens: Tensor) -> Tensor:
     """Consume a full initial context through one fused compiled replay."""
 
+    state._ensure_open()
     if state.position != 0:
         raise RuntimeError("prefill requires an empty inference state")
     if tokens.ndim != 2 or tokens.shape[0] != state.batch_size:
diff --git a/src/rosa/ragged.py b/src/rosa/ragged.py
index b258ac3..f7f70b8 100644
--- a/src/rosa/ragged.py
+++ b/src/rosa/ragged.py
@@ -44,6 +44,21 @@ def __init__(
         setattr(self._state, "positions", self._positions)
         self._use_native = use_native
         self._native: Any = None
+        self._closed = False
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("state is closed")
+
+    def close(self) -> None:
+        """Release optional native ownership cycles idempotently."""
+
+        if self._closed:
+            return
+        self._closed = True
+        self._native = None
+        self._state.native_state = None
+        self._state.close()
 
     @property
     def batch_size(self) -> int:
@@ -61,6 +76,7 @@ def max_length(self) -> int:
     def positions(self) -> Tensor:
         """Current consumed-token count for each row, returned as a copy."""
 
+        self._ensure_open()
         return torch.from_numpy(self._positions.copy())
 
     @property
@@ -70,6 +86,7 @@ def using_native(self) -> bool:
         return self._native not in (None, False)
 
     def _native_state(self) -> Any:
+        self._ensure_open()
         if not self._use_native or self._native is False:
             return None
         if self._native is None:
@@ -97,6 +114,7 @@ def step_masked(
     ) -> Tensor:
         """Consume tokens for active rows, optionally recycling selected slots."""
 
+        self._ensure_open()
         if not isinstance(tokens, Tensor):
             raise TypeError("tokens must be a Tensor")
         if tokens.ndim == 0 and self.batch_size == 1:
diff --git a/tests/test_ragged.py b/tests/test_ragged.py
index 336cac5..bb36d8f 100644
--- a/tests/test_ragged.py
+++ b/tests/test_ragged.py
@@ -11,6 +11,15 @@
 
 
 class TestRaggedInference(unittest.TestCase):
+    def test_close_is_idempotent_and_rejects_use(self) -> None:
+        state = init_ragged_state(1, 2, use_native=False)
+        state.close()
+        state.close()
+        with self.assertRaisesRegex(RuntimeError, "^state is closed$"):
+            _ = state.positions
+        with self.assertRaisesRegex(RuntimeError, "^state is closed$"):
+            state.step(torch.tensor([0]))
+
     def test_mask_reset_and_recycling_match_single_row_oracles(self) -> None:
         from rosa._stateful_numba import _forward_step, _init_inference_state
 
diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py
index 9b85f3c..6947120 100644
--- a/tests/test_stateful_candidates.py
+++ b/tests/test_stateful_candidates.py
@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import gc
 import sys
 import unittest
+import weakref
 from itertools import product
 from unittest.mock import patch
 
@@ -41,6 +43,35 @@
 
 
 class TestStatefulCandidates(unittest.TestCase):
+    def test_close_is_idempotent_breaks_cycles_and_rejects_use(self) -> None:
+        state = init_candidate_state_internal(1, 2, suffix_k=2, occurrences_r=2)
+        buffers = init_candidate_buffers(state)
+
+        class NativeOwner:
+            def __init__(self, candidate_state: CandidateState) -> None:
+                self.candidate_state = candidate_state
+
+        state.native_state = NativeOwner(state)
+        native_ref = weakref.ref(state.native_state)
+        state.close()
+        state.close()
+        self.assertIsNone(native_ref())
+        with self.assertRaisesRegex(RuntimeError, "^state is closed$"):
+            _ = state.position
+        with self.assertRaisesRegex(RuntimeError, "^state is closed$"):
+            _ = state.positions
+        with self.assertRaisesRegex(RuntimeError, "^state is closed$"):
+            forward_candidates_step(state, torch.tensor([0]))
+        with self.assertRaisesRegex(RuntimeError, "^state is closed$"):
+            forward_candidates_step_into(state, torch.tensor([0]), buffers)
+        with self.assertRaisesRegex(RuntimeError, "^state is closed$"):
+            prefill_candidates(state, torch.tensor([[0]]))
+
+        state_ref = weakref.ref(state)
+        del state
+        gc.collect()
+        self.assertIsNone(state_ref())
+
     def assert_matches_oracle(
         self,
         tokens: torch.Tensor,
diff --git a/tests/test_unified_inference.py b/tests/test_unified_inference.py
index afcaa93..fef5ec2 100644
--- a/tests/test_unified_inference.py
+++ b/tests/test_unified_inference.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import gc
 import unittest
+import weakref
 from typing import Any, cast
 
 import torch
@@ -21,6 +23,57 @@
 
 
 class TestUnifiedInferenceState(unittest.TestCase):
+    def test_close_context_manager_and_reset_lifetime(self) -> None:
+        tokens = torch.tensor([[0, 1, 0]], dtype=torch.long)
+        expected = reference_rosa(tokens)[0]
+
+        with init_inference_state(1, 3, mode="rich", suffix_k=2) as managed:
+            buffers = init_candidate_buffers(managed)
+            self.assertTrue(
+                torch.equal(managed.prefill(tokens).predicted_tokens, expected)
+            )
+        for access in (
+            lambda: managed.position,
+            lambda: managed.positions,
+            lambda: managed.step(torch.tensor([0])),
+            lambda: managed.step_into(torch.tensor([0]), buffers),
+            lambda: managed.prefill(tokens),
+        ):
+            with self.subTest(access=access), self.assertRaisesRegex(
+                RuntimeError, "^state is closed$"
+            ):
+                access()
+        managed.close()
+
+        for mode in ("top1", "rich"):
+            with self.subTest(mode=mode):
+                state = init_inference_state(1, 3, mode=mode)  # type: ignore[arg-type]
+                old_impl = state._impl
+                assert old_impl is not None
+
+                class NativeOwner:
+                    def __init__(self, impl: object) -> None:
+                        self.impl = impl
+
+                old_impl.native_state = NativeOwner(old_impl)  # type: ignore[attr-defined]
+                impl_ref = weakref.ref(old_impl)
+                array_ref = weakref.ref(old_impl.history)  # type: ignore[attr-defined]
+                del old_impl
+                state.reset()
+                gc.collect()
+                self.assertIsNone(impl_ref())
+                self.assertIsNone(array_ref())
+                self.assertEqual(state.positions.tolist(), [0])
+                self.assertTrue(
+                    torch.equal(state.prefill(tokens).predicted_tokens, expected)
+                )
+                state.close()
+                state.reset()
+                self.assertEqual(state.position, 0)
+                self.assertTrue(
+                    torch.equal(state.prefill(tokens).predicted_tokens, expected)
+                )
+
     def test_uniform_mode_matrix_prefill_continuation_and_reset(self) -> None:
         tokens = torch.tensor(
             [[0, 1, 0, 2, 0, 3], [4, 4, 5, 4, 4, 6]], dtype=torch.long

From 822bfef3ce3ce01de2470c6962b0d63e9eb96fa2 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 16:44:33 +0800
Subject: [PATCH 03/10] Remove native state ownership cycles

---
 native/src/rosa_native_step.cpp | 178 ++++++++++++++++++--------------
 native/tests/candidate_smoke.py |  27 ++++-
 native/tests/smoke.py           |  37 +++++++
 3 files changed, 165 insertions(+), 77 deletions(-)

diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp
index 9c32cc1..848ce6b 100644
--- a/native/src/rosa_native_step.cpp
+++ b/native/src/rosa_native_step.cpp
@@ -267,37 +267,46 @@ class RowThreadPool {
   std::exception_ptr exception_;
 };
 
+py::object make_owner_weakref(py::handle owner) {
+  if (!PyType_SUPPORTS_WEAKREFS(Py_TYPE(owner.ptr())))
+    return py::none();
+  PyObject *reference = PyWeakref_NewRef(owner.ptr(), nullptr);
+  if (reference != nullptr)
+    return py::reinterpret_steal<py::object>(reference);
+  throw py::error_already_set();
+}
+
 } // namespace
 
 class NativeState {
 public:
-  explicit NativeState(py::object state) : state_(std::move(state)) {
-    const int64_t abi = py::cast<int64_t>(state_.attr("native_abi_version"));
+  explicit NativeState(py::object state) : owner_ref_(make_owner_weakref(state)) {
+    const int64_t abi = py::cast<int64_t>(state.attr("native_abi_version"));
     if (abi != 1)
       throw py::value_error("unsupported native state ABI");
-    history_ = bind<int64_t>("history");
-    head_ = bind<int32_t>("head");
-    edge_token_ = bind<int64_t>("edge_token");
-    edge_target_ = bind<int32_t>("edge_target");
-    edge_next_ = bind<int32_t>("edge_next");
-    hash_state_ = bind<int32_t>("hash_state");
-    hash_token_ = bind<int64_t>("hash_token");
-    hash_edge_ = bind<int32_t>("hash_edge");
-    suffix_link_ = bind<int32_t>("suffix_link");
-    length_ = bind<int32_t>("length");
-    left_ = bind<int32_t>("lct_left");
-    right_ = bind<int32_t>("lct_right");
-    parent_ = bind<int32_t>("lct_parent");
-    value_ = bind<int64_t>("lct_value");
-    lazy_ = bind<int64_t>("lct_lazy");
-    lazy_valid_ = bind<uint8_t>("lct_lazy_valid");
-    stack_ = bind<int32_t>("lct_stack");
-    last_ = bind<int32_t>("last");
-    size_ = bind<int32_t>("size");
-    edge_count_ = bind<int32_t>("edge_count");
-    batch_ = py::cast<int64_t>(state_.attr("batch_size"));
-    max_length_ = py::cast<int64_t>(state_.attr("max_length"));
-    position_ = py::cast<int64_t>(state_.attr("position"));
+    history_ = bind<int64_t>(state, "history");
+    head_ = bind<int32_t>(state, "head");
+    edge_token_ = bind<int64_t>(state, "edge_token");
+    edge_target_ = bind<int32_t>(state, "edge_target");
+    edge_next_ = bind<int32_t>(state, "edge_next");
+    hash_state_ = bind<int32_t>(state, "hash_state");
+    hash_token_ = bind<int64_t>(state, "hash_token");
+    hash_edge_ = bind<int32_t>(state, "hash_edge");
+    suffix_link_ = bind<int32_t>(state, "suffix_link");
+    length_ = bind<int32_t>(state, "length");
+    left_ = bind<int32_t>(state, "lct_left");
+    right_ = bind<int32_t>(state, "lct_right");
+    parent_ = bind<int32_t>(state, "lct_parent");
+    value_ = bind<int64_t>(state, "lct_value");
+    lazy_ = bind<int64_t>(state, "lct_lazy");
+    lazy_valid_ = bind<uint8_t>(state, "lct_lazy_valid");
+    stack_ = bind<int32_t>(state, "lct_stack");
+    last_ = bind<int32_t>(state, "last");
+    size_ = bind<int32_t>(state, "size");
+    edge_count_ = bind<int32_t>(state, "edge_count");
+    batch_ = py::cast<int64_t>(state.attr("batch_size"));
+    max_length_ = py::cast<int64_t>(state.attr("max_length"));
+    position_ = py::cast<int64_t>(state.attr("position"));
     state_capacity_ = head_.shape(1);
     edge_capacity_ = edge_token_.shape(1);
     hash_capacity_ = hash_state_.shape(1);
@@ -305,9 +314,9 @@ class NativeState {
     positions_ = py::array_t<int64_t>(batch_);
     std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
               position_);
-    if (py::hasattr(state_, "positions")) {
+    if (py::hasattr(state, "positions")) {
       ragged_mode_ = true;
-      py::object object = state_.attr("positions");
+      py::object object = state.attr("positions");
       if (!py::isinstance<py::array_t<int64_t>>(object))
         throw py::type_error("positions has an unexpected dtype");
       positions_ = py::cast<py::array_t<int64_t, py::array::c_style>>(object);
@@ -347,7 +356,7 @@ class NativeState {
     ++position_;
     std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
               position_);
-    state_.attr("position") = py::int_(position_);
+    sync_owner_position();
     call_lock.unlock();
     return output;
   }
@@ -446,7 +455,7 @@ class NativeState {
     position_ = token_count;
     std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
               position_);
-    state_.attr("position") = py::int_(position_);
+    sync_owner_position();
     call_lock.unlock();
     return output;
   }
@@ -487,8 +496,9 @@ class NativeState {
   }
 
   template <typename T>
-  py::array_t<T, py::array::c_style> bind(const char *name) {
-    py::object object = state_.attr(name);
+  py::array_t<T, py::array::c_style> bind(const py::object &state,
+                                          const char *name) {
+    py::object object = state.attr(name);
     if (!py::isinstance<py::array_t<T>>(object)) {
       throw py::type_error(std::string(name) + " has an unexpected dtype");
     }
@@ -498,6 +508,14 @@ class NativeState {
     return array;
   }
 
+  void sync_owner_position() {
+    if (owner_ref_.is_none())
+      return;
+    py::object owner = owner_ref_();
+    if (!owner.is_none())
+      owner.attr("position") = py::int_(position_);
+  }
+
   void validate_shapes() {
     if (batch_ <= 0 || max_length_ <= 0 || position_ < 0 ||
         position_ > max_length_ || state_capacity_ <= 0 ||
@@ -1062,7 +1080,7 @@ class NativeState {
     edge_count_.mutable_data()[b] = edge_count;
   }
 
-  py::object state_;
+  py::object owner_ref_;
   py::array_t<int64_t, py::array::c_style> history_, edge_token_, hash_token_,
       value_, lazy_;
   py::array_t<int32_t, py::array::c_style> head_, edge_target_, edge_next_,
@@ -1080,49 +1098,50 @@ class NativeState {
 
 class NativeCandidateState {
 public:
-  explicit NativeCandidateState(py::object state) : state_(std::move(state)) {
-    if (!py::hasattr(state_, "native_candidate_abi_version") ||
-        py::cast<int64_t>(state_.attr("native_candidate_abi_version")) != 1)
+  explicit NativeCandidateState(py::object state)
+      : owner_ref_(make_owner_weakref(state)) {
+    if (!py::hasattr(state, "native_candidate_abi_version") ||
+        py::cast<int64_t>(state.attr("native_candidate_abi_version")) != 1)
       throw py::value_error("unsupported native candidate state ABI");
-    history_ = bind<int64_t>("history");
-    head_ = bind<int32_t>("head");
-    edge_token_ = bind<int64_t>("edge_token");
-    edge_target_ = bind<int32_t>("edge_target");
-    edge_next_ = bind<int32_t>("edge_next");
-    hash_state_ = bind<int32_t>("hash_state");
-    hash_token_ = bind<int64_t>("hash_token");
-    hash_edge_ = bind<int32_t>("hash_edge");
-    suffix_link_ = bind<int32_t>("suffix_link");
-    length_ = bind<int32_t>("length");
-    left_ = bind<int32_t>("lct_left");
-    right_ = bind<int32_t>("lct_right");
-    parent_ = bind<int32_t>("lct_parent");
-    occurrences_ = bind<int64_t>("occurrences");
-    occurrence_size_ = bind<int32_t>("occurrence_size");
-    frequency_ = bind<int64_t>("frequency");
-    lazy_prefix_ = bind<int64_t>("lazy_prefix");
-    lazy_size_ = bind<int32_t>("lazy_size");
-    lazy_delta_ = bind<int64_t>("lazy_delta");
-    stack_ = bind<int32_t>("lct_stack");
-    last_ = bind<int32_t>("last");
-    size_ = bind<int32_t>("size");
-    edge_count_ = bind<int32_t>("edge_count");
-    batch_ = py::cast<int64_t>(state_.attr("batch_size"));
-    max_length_ = py::cast<int64_t>(state_.attr("max_length"));
-    suffix_k_ = py::cast<int64_t>(state_.attr("suffix_k"));
-    occurrences_r_ = py::cast<int64_t>(state_.attr("occurrences_r"));
-    position_ = py::cast<int64_t>(state_.attr("position"));
+    history_ = bind<int64_t>(state, "history");
+    head_ = bind<int32_t>(state, "head");
+    edge_token_ = bind<int64_t>(state, "edge_token");
+    edge_target_ = bind<int32_t>(state, "edge_target");
+    edge_next_ = bind<int32_t>(state, "edge_next");
+    hash_state_ = bind<int32_t>(state, "hash_state");
+    hash_token_ = bind<int64_t>(state, "hash_token");
+    hash_edge_ = bind<int32_t>(state, "hash_edge");
+    suffix_link_ = bind<int32_t>(state, "suffix_link");
+    length_ = bind<int32_t>(state, "length");
+    left_ = bind<int32_t>(state, "lct_left");
+    right_ = bind<int32_t>(state, "lct_right");
+    parent_ = bind<int32_t>(state, "lct_parent");
+    occurrences_ = bind<int64_t>(state, "occurrences");
+    occurrence_size_ = bind<int32_t>(state, "occurrence_size");
+    frequency_ = bind<int64_t>(state, "frequency");
+    lazy_prefix_ = bind<int64_t>(state, "lazy_prefix");
+    lazy_size_ = bind<int32_t>(state, "lazy_size");
+    lazy_delta_ = bind<int64_t>(state, "lazy_delta");
+    stack_ = bind<int32_t>(state, "lct_stack");
+    last_ = bind<int32_t>(state, "last");
+    size_ = bind<int32_t>(state, "size");
+    edge_count_ = bind<int32_t>(state, "edge_count");
+    batch_ = py::cast<int64_t>(state.attr("batch_size"));
+    max_length_ = py::cast<int64_t>(state.attr("max_length"));
+    suffix_k_ = py::cast<int64_t>(state.attr("suffix_k"));
+    occurrences_r_ = py::cast<int64_t>(state.attr("occurrences_r"));
+    position_ = py::cast<int64_t>(state.attr("position"));
     state_capacity_ = head_.shape(1);
     edge_capacity_ = edge_token_.shape(1);
     hash_capacity_ = hash_state_.shape(1);
     validate_shapes();
-    ragged_mode_ = py::hasattr(state_, "ragged_mode") &&
-                   py::cast<bool>(state_.attr("ragged_mode"));
+    ragged_mode_ = py::hasattr(state, "ragged_mode") &&
+                   py::cast<bool>(state.attr("ragged_mode"));
     positions_ = py::array_t<int64_t>(batch_);
     std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
               position_);
-    if (py::hasattr(state_, "positions")) {
-      positions_ = bind<int64_t>("positions");
+    if (py::hasattr(state, "positions")) {
+      positions_ = bind<int64_t>(state, "positions");
       if (!vector_shape(positions_, batch_))
         throw py::value_error(
             "positions must be contiguous int64 [batch_size]");
@@ -1196,7 +1215,7 @@ class NativeCandidateState {
     ++position_;
     std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
               position_);
-    state_.attr("position") = py::int_(position_);
+    sync_owner_position();
     call_lock.unlock();
   }
 
@@ -1213,7 +1232,7 @@ class NativeCandidateState {
     position_ = 0;
     std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
               int64_t{0});
-    state_.attr("position") = py::int_(0);
+    sync_owner_position();
     call_lock.unlock();
   }
 
@@ -1396,7 +1415,7 @@ class NativeCandidateState {
     position_ = sequence_length;
     std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
               position_);
-    state_.attr("position") = py::int_(position_);
+    sync_owner_position();
     call_lock.unlock();
   }
 
@@ -1486,8 +1505,10 @@ class NativeCandidateState {
                                  std::forward<Function>(function));
   }
 
-  template <typename T> py::array_t<T, py::array::c_style> bind(const char *name) {
-    py::object object = state_.attr(name);
+  template <typename T>
+  py::array_t<T, py::array::c_style> bind(const py::object &state,
+                                          const char *name) {
+    py::object object = state.attr(name);
     if (!py::isinstance<py::array_t<T>>(object))
       throw py::type_error(std::string(name) + " has an unexpected dtype");
     auto array = py::cast<py::array_t<T, py::array::c_style>>(object);
@@ -1495,6 +1516,13 @@ class NativeCandidateState {
       throw py::value_error(std::string(name) + " is readonly");
     return array;
   }
+  void sync_owner_position() {
+    if (owner_ref_.is_none())
+      return;
+    py::object owner = owner_ref_();
+    if (!owner.is_none())
+      owner.attr("position") = py::int_(position_);
+  }
   template <typename T>
   bool matrix_shape(const py::array_t<T, py::array::c_style> &a, int64_t rows,
                     int64_t columns) const {
@@ -1826,7 +1854,7 @@ class NativeCandidateState {
     edge_count_.mutable_data()[b] = 0;
   }
 
-  py::object state_;
+  py::object owner_ref_;
   py::array_t<int64_t, py::array::c_style> history_, edge_token_, hash_token_,
       occurrences_, frequency_, lazy_prefix_, lazy_delta_, positions_;
   py::array_t<int32_t, py::array::c_style> head_, edge_target_, edge_next_,
@@ -5026,7 +5054,7 @@ class NativeRLBWTStateMC : public NativeRLBWTState {
 PYBIND11_MODULE(rosa_native_step, m) {
   m.doc() = "Exact CPU SAM+LCT and RLBWT backends (no libtorch calls)";
   py::class_<NativeState>(m, "NativeState")
-      .def(py::init<py::object>(), py::keep_alive<1, 2>())
+      .def(py::init<py::object>())
       .def("step", &NativeState::step)
       .def("step_masked", &NativeState::step_masked)
       .def("prefill", &NativeState::prefill)
@@ -5034,7 +5062,7 @@ PYBIND11_MODULE(rosa_native_step, m) {
       .def_property_readonly("positions", &NativeState::positions)
       .def_property_readonly("worker_count", &NativeState::worker_count);
   py::class_<NativeCandidateState>(m, "NativeCandidateState")
-      .def(py::init<py::object>(), py::keep_alive<1, 2>())
+      .def(py::init<py::object>())
       .def("step", &NativeCandidateState::step)
       .def("step_into", &NativeCandidateState::step_into)
       .def("reset", &NativeCandidateState::reset)
diff --git a/native/tests/candidate_smoke.py b/native/tests/candidate_smoke.py
index c65a21a..829ce34 100644
--- a/native/tests/candidate_smoke.py
+++ b/native/tests/candidate_smoke.py
@@ -3,6 +3,7 @@
 import gc
 import os
 import threading
+import types
 import weakref
 from concurrent.futures import ThreadPoolExecutor
 from itertools import product
@@ -290,11 +291,33 @@ def concurrent_step() -> tuple[np.ndarray, ...]:
 
     state = init_candidate_state(2, 4, suffix_k=3, occurrences_r=2)
     state_ref = weakref.ref(state)
+    history_ref = weakref.ref(state.history)
     native = rosa_native_step.NativeCandidateState(state)
+    state.position = 3
+    native.step(np.array([1, 2], dtype=np.int64))
+    assert state.position == native.position == 1
     del state
     gc.collect()
-    assert state_ref() is not None
-    native.step(np.array([1, 2], dtype=np.int64))
+    assert state_ref() is None
+    assert history_ref() is not None
+    native.step(np.array([3, 4], dtype=np.int64))
+    assert native.position == 2
+    assert history_ref()[:, :2].tolist() == [[1, 3], [2, 4]]
+
+    cyclic_owner = init_candidate_state(1, 2)
+    cyclic_wrapper = rosa_native_step.NativeCandidateState(cyclic_owner)
+    cyclic_owner.native_state = cyclic_wrapper
+    cyclic_owner_ref = weakref.ref(cyclic_owner)
+    cyclic_wrapper_ref = weakref.ref(cyclic_wrapper)
+    del cyclic_owner, cyclic_wrapper
+    gc.collect()
+    assert cyclic_owner_ref() is None
+    assert cyclic_wrapper_ref() is None
+
+    nonweak_owner = types.SimpleNamespace(**vars(init_candidate_state(1, 2)))
+    nonweak = rosa_native_step.NativeCandidateState(nonweak_owner)
+    nonweak.step(np.array([11], dtype=np.int64))
+    assert nonweak.position == 1
 
     # ABI 1 compatibility: pre-positions uniform states remain accepted.
     legacy = init_candidate_state(1, 2)
diff --git a/native/tests/smoke.py b/native/tests/smoke.py
index 70a1728..4d18228 100644
--- a/native/tests/smoke.py
+++ b/native/tests/smoke.py
@@ -1,5 +1,9 @@
 from __future__ import annotations
 
+import gc
+import types
+import weakref
+
 import numpy as np
 import rosa_native_step
 import torch
@@ -77,6 +81,39 @@ def main() -> None:
     assert candidate.native_state.position == tokens.shape[1]
     assert candidate.position == tokens.shape[1]
 
+    # The wrapper owns the NumPy buffers, but only weakly observes the Python
+    # state for position publication.  It therefore survives its owner, and a
+    # normal owner -> wrapper link does not form an uncollectable cycle.
+    detached_owner = _init_inference_state(2, 4)
+    detached_owner_ref = weakref.ref(detached_owner)
+    detached_history_ref = weakref.ref(detached_owner.history)
+    detached = rosa_native_step.NativeState(detached_owner)
+    detached_owner.position = 3
+    detached.step(np.array([3, 5], dtype=np.int64))
+    assert detached_owner.position == detached.position == 1
+    del detached_owner
+    gc.collect()
+    assert detached_owner_ref() is None
+    assert detached_history_ref() is not None
+    detached.step(np.array([7, 9], dtype=np.int64))
+    assert detached.position == 2
+    assert detached_history_ref()[:, :2].tolist() == [[3, 7], [5, 9]]
+
+    cyclic_owner = _init_inference_state(1, 2)
+    cyclic_wrapper = rosa_native_step.NativeState(cyclic_owner)
+    cyclic_owner.native_state = cyclic_wrapper
+    cyclic_owner_ref = weakref.ref(cyclic_owner)
+    cyclic_wrapper_ref = weakref.ref(cyclic_wrapper)
+    del cyclic_owner, cyclic_wrapper
+    gc.collect()
+    assert cyclic_owner_ref() is None
+    assert cyclic_wrapper_ref() is None
+
+    nonweak_owner = types.SimpleNamespace(**vars(_init_inference_state(1, 2)))
+    nonweak = rosa_native_step.NativeState(nonweak_owner)
+    nonweak.step(np.array([11], dtype=np.int64))
+    assert nonweak.position == 1
+
     generator = torch.Generator().manual_seed(20260811)
     cases = [
         torch.randint(-3, 9, (3, 257), generator=generator, dtype=torch.long),

From efa5a8eaeffdb1368483fae848b36b0e6447f730 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 17:11:57 +0800
Subject: [PATCH 04/10] Handle non-weakref native state owners

---
 native/src/rosa_native_step.cpp |  4 ++++
 native/tests/candidate_smoke.py | 21 ++++++++++++++++++---
 native/tests/smoke.py           | 21 ++++++++++++++++++---
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp
index 848ce6b..c63b6bb 100644
--- a/native/src/rosa_native_step.cpp
+++ b/native/src/rosa_native_step.cpp
@@ -273,6 +273,10 @@ py::object make_owner_weakref(py::handle owner) {
   PyObject *reference = PyWeakref_NewRef(owner.ptr(), nullptr);
   if (reference != nullptr)
     return py::reinterpret_steal<py::object>(reference);
+  if (PyErr_ExceptionMatches(PyExc_TypeError)) {
+    PyErr_Clear();
+    return py::none();
+  }
   throw py::error_already_set();
 }
 
diff --git a/native/tests/candidate_smoke.py b/native/tests/candidate_smoke.py
index 829ce34..cecd77c 100644
--- a/native/tests/candidate_smoke.py
+++ b/native/tests/candidate_smoke.py
@@ -3,7 +3,6 @@
 import gc
 import os
 import threading
-import types
 import weakref
 from concurrent.futures import ThreadPoolExecutor
 from itertools import product
@@ -22,6 +21,16 @@
 )
 
 
+def nonweak_owner(source: object) -> object:
+    class NonWeakOwner:
+        __slots__ = tuple(vars(source))
+
+    owner = NonWeakOwner()
+    for name, value in vars(source).items():
+        setattr(owner, name, value)
+    return owner
+
+
 def assert_step_equal(actual: tuple[np.ndarray, ...], expected: CandidateStep) -> None:
     expected_arrays = (
         expected.source_index.numpy(),
@@ -314,8 +323,14 @@ def concurrent_step() -> tuple[np.ndarray, ...]:
     assert cyclic_owner_ref() is None
     assert cyclic_wrapper_ref() is None
 
-    nonweak_owner = types.SimpleNamespace(**vars(init_candidate_state(1, 2)))
-    nonweak = rosa_native_step.NativeCandidateState(nonweak_owner)
+    nonweak_state = nonweak_owner(init_candidate_state(1, 2))
+    try:
+        weakref.ref(nonweak_state)
+    except TypeError:
+        pass
+    else:
+        raise AssertionError("slotted owner unexpectedly supports weak references")
+    nonweak = rosa_native_step.NativeCandidateState(nonweak_state)
     nonweak.step(np.array([11], dtype=np.int64))
     assert nonweak.position == 1
 
diff --git a/native/tests/smoke.py b/native/tests/smoke.py
index 4d18228..ab2589b 100644
--- a/native/tests/smoke.py
+++ b/native/tests/smoke.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import gc
-import types
 import weakref
 
 import numpy as np
@@ -12,6 +11,16 @@
 from rosa.ragged import RaggedInferenceState
 
 
+def nonweak_owner(source: object) -> object:
+    class NonWeakOwner:
+        __slots__ = tuple(vars(source))
+
+    owner = NonWeakOwner()
+    for name, value in vars(source).items():
+        setattr(owner, name, value)
+    return owner
+
+
 def assert_same_initialized_state(oracle: object, candidate: object) -> None:
     assert oracle.position == candidate.position
     for batch in range(oracle.batch_size):
@@ -109,8 +118,14 @@ def main() -> None:
     assert cyclic_owner_ref() is None
     assert cyclic_wrapper_ref() is None
 
-    nonweak_owner = types.SimpleNamespace(**vars(_init_inference_state(1, 2)))
-    nonweak = rosa_native_step.NativeState(nonweak_owner)
+    nonweak_state = nonweak_owner(_init_inference_state(1, 2))
+    try:
+        weakref.ref(nonweak_state)
+    except TypeError:
+        pass
+    else:
+        raise AssertionError("slotted owner unexpectedly supports weak references")
+    nonweak = rosa_native_step.NativeState(nonweak_state)
     nonweak.step(np.array([11], dtype=np.int64))
     assert nonweak.position == 1
 

From bf50aa19e4e659ff7cff4920317ad40208362f87 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 17:40:13 +0800
Subject: [PATCH 05/10] Add opt-in query-only retrieval scoring

---
 src/rosa/__init__.py | 456 +++++++++++++++++++++++++++++++++++++++++++
 tests/test_rosa.py   | 162 +++++++++++++++
 2 files changed, 618 insertions(+)

diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py
index 9085555..a8d25ff 100644
--- a/src/rosa/__init__.py
+++ b/src/rosa/__init__.py
@@ -547,6 +547,36 @@ def _soft_match_torch(
     return score
 
 
+def _soft_match_at_positions_torch(
+    st1: Tensor,
+    st2: Tensor,
+    query_positions: Tensor,
+    source_index: Tensor,
+    candidate_mask: Tensor,
+    window: int,
+) -> Tensor:
+    """Soft suffix verification for selected real sequence positions."""
+
+    bsz, queries, candidates = source_index.shape
+    positions = query_positions.unsqueeze(-1).expand(bsz, queries, candidates)
+    survival = torch.ones(
+        (bsz, queries, candidates), dtype=st1.dtype, device=st1.device
+    )
+    score = torch.zeros_like(survival)
+    for r in range(window):
+        left_idx = positions - r
+        right_idx = source_index - r
+        valid = candidate_mask & (left_idx >= 0) & (right_idx >= 0)
+        left1 = _gather_sequence(st1, left_idx)
+        right1 = _gather_sequence(st1, right_idx)
+        left2 = _gather_sequence(st2, left_idx)
+        right2 = _gather_sequence(st2, right_idx)
+        eq = (left1 * right1).sum(-1) * (left2 * right2).sum(-1)
+        survival = survival * eq * valid.to(eq.dtype)
+        score = score + survival
+    return score
+
+
 # A callable owns shape specializations while the outer cache keeps exactly one
 # full-graph compiler island per static verification window. Specializations
 # are initialized once per input signature; already-ready calls do not take the
@@ -1062,11 +1092,424 @@ def _hybrid_soft_candidates(
         sparse_mask = pool_mask.gather(-1, selected)
         return dense_source, dense_mask, sparse_source, sparse_mask
 
+    @staticmethod
+    def _validate_query_positions(query_positions: Tensor, z_a: Tensor) -> None:
+        if not isinstance(query_positions, Tensor):
+            raise TypeError("query_positions must be a Tensor")
+        if query_positions.dtype != torch.long:
+            raise TypeError("query_positions must have dtype torch.long")
+        if query_positions.ndim != 2 or query_positions.shape[0] != z_a.shape[0]:
+            raise ValueError("query_positions must have shape [B, Q]")
+        if query_positions.shape[1] <= 0:
+            raise ValueError("query_positions must contain at least one query")
+        if query_positions.device != z_a.device:
+            raise ValueError("query_positions must be on the same device as z_a")
+        if bool(((query_positions < 0) | (query_positions >= z_a.shape[1])).any()):
+            raise ValueError("query_positions values must be in [0, N)")
+        ordered = query_positions.sort(dim=1).values
+        if ordered.shape[1] > 1 and bool((ordered[:, 1:] == ordered[:, :-1]).any()):
+            raise ValueError("query_positions must be unique within each batch row")
+
+    @staticmethod
+    def _scatter_queries(
+        values: Tensor,
+        query_positions: Tensor,
+        sequence_length: int,
+        fill_value: float | int | bool,
+    ) -> Tensor:
+        shape = (values.shape[0], sequence_length, *values.shape[2:])
+        result = torch.full(shape, fill_value, dtype=values.dtype, device=values.device)
+        index = query_positions.reshape(
+            values.shape[0], values.shape[1], *([1] * (values.ndim - 2))
+        ).expand_as(values)
+        return result.scatter(1, index, values)
+
+    def _forward_query_only(
+        self,
+        z_a: Tensor,
+        z_b: Tensor,
+        code_logits: tuple[Tensor, Tensor] | None,
+        query_positions: Tensor,
+    ) -> ROSAOutput:
+        """Run candidate activations only at selected sequence positions."""
+
+        (soft1, soft2), (st1, st2), hard_tokens = self.encode(z_a, code_logits)
+        hard = _build_forward_hard_candidates(
+            hard_tokens,
+            self.suffix_k,
+            self.occurrences_r,
+            self.candidate_backend,
+        )
+        bsz, n, _ = z_a.shape
+        queries = query_positions.shape[1]
+        z_query = _gather_sequence(z_a, query_positions)
+        exact_source = _gather_sequence(hard.source_index, query_positions)
+        exact_mask = _gather_sequence(hard.mask, query_positions)
+        exact_match_length = _gather_sequence(hard.match_length, query_positions)
+        exact_frequency = _gather_sequence(hard.frequency, query_positions)
+        exact_slots = exact_source.shape[-1]
+
+        # The pool is discrete and cheap to build for N; floating candidate
+        # projections below are restricted to [B, Q, C, *].
+        virtual_pool = _gather_sequence(
+            build_virtual_pool_indices(bsz, n, self.virtual_pool_size, z_a.device),
+            query_positions,
+        )
+        virtual_pool_mask = virtual_pool >= 0
+        virtual_duplicate = (
+            virtual_pool.unsqueeze(-1) == exact_source.unsqueeze(-2)
+        ) & exact_mask.unsqueeze(-2)
+        virtual_pool_mask = virtual_pool_mask & ~virtual_duplicate.any(dim=-1)
+        virtual_query = self.virtual_query(z_query).unsqueeze(-2)
+        virtual_keys = self.virtual_key(_gather_sequence(z_a, virtual_pool))
+        virtual_router_all = (virtual_query * virtual_keys).sum(-1) / math.sqrt(
+            self.selector_dim
+        )
+        virtual_masked = virtual_router_all.masked_fill(~virtual_pool_mask, -1e9)
+        _, virtual_top = torch.topk(virtual_masked, k=self.virtual_candidates, dim=-1)
+        virtual_source = virtual_pool.gather(-1, virtual_top)
+        virtual_mask = virtual_pool_mask.gather(-1, virtual_top)
+        virtual_router = virtual_router_all.gather(-1, virtual_top)
+        virtual_mask = virtual_mask & (self.virtual_scale > 0)
+        virtual_router = virtual_router * self.virtual_scale
+
+        dense_count = self.dense_recent_candidates
+        sparse_count = self.sparse_old_candidates
+        if dense_count:
+            offsets = torch.arange(1, dense_count + 1, device=z_a.device).view(1, 1, -1)
+            dense_source = query_positions.unsqueeze(-1) - offsets
+            dense_mask = dense_source >= 0
+            dense_duplicate = (
+                dense_source.unsqueeze(-1) == exact_source.unsqueeze(-2)
+            ) & exact_mask.unsqueeze(-2)
+            dense_mask = dense_mask & ~dense_duplicate.any(dim=-1)
+        else:
+            dense_source = torch.empty(
+                (bsz, queries, 0), dtype=torch.long, device=z_a.device
+            )
+            dense_mask = torch.empty(
+                (bsz, queries, 0), dtype=torch.bool, device=z_a.device
+            )
+
+        if sparse_count:
+            pool_size = self.sparse_old_pool_size
+            old_count = (query_positions - dense_count).clamp_min(0)
+            anchor_rank = torch.arange(pool_size, device=z_a.device).view(1, 1, -1)
+            sparse_pool_mask = anchor_rank < old_count.clamp_max(pool_size).unsqueeze(
+                -1
+            )
+            if pool_size == 1:
+                sparse_pool = torch.zeros(
+                    (bsz, queries, 1), dtype=torch.long, device=z_a.device
+                )
+            else:
+                spread = torch.round(
+                    anchor_rank
+                    * (old_count - 1).clamp_min(0).unsqueeze(-1)
+                    / (pool_size - 1)
+                ).to(torch.long)
+                sparse_pool = torch.where(
+                    old_count.unsqueeze(-1) <= pool_size,
+                    anchor_rank.expand(bsz, queries, -1),
+                    spread,
+                )
+            hard_duplicate = (
+                sparse_pool.unsqueeze(-1) == exact_source.unsqueeze(-2)
+            ) & exact_mask.unsqueeze(-2)
+            dense_duplicate = (
+                sparse_pool.unsqueeze(-1) == dense_source.unsqueeze(-2)
+            ) & dense_mask.unsqueeze(-2)
+            sparse_pool_mask = (
+                sparse_pool_mask
+                & ~hard_duplicate.any(dim=-1)
+                & ~dense_duplicate.any(dim=-1)
+            )
+            sparse_pool_score = _soft_match_at_positions_torch(
+                st1,
+                st2,
+                query_positions,
+                sparse_pool,
+                sparse_pool_mask,
+                self.soft_verify_window,
+            ).masked_fill(~sparse_pool_mask, -1e9)
+            recency_order = torch.argsort(
+                sparse_pool, dim=-1, descending=True, stable=True
+            )
+            ordered_score = sparse_pool_score.gather(-1, recency_order)
+            score_order = torch.argsort(
+                ordered_score, dim=-1, descending=True, stable=True
+            )
+            sparse_selected = recency_order.gather(-1, score_order)[..., :sparse_count]
+            sparse_source = sparse_pool.gather(-1, sparse_selected)
+            sparse_mask = sparse_pool_mask.gather(-1, sparse_selected)
+        else:
+            sparse_source = torch.empty(
+                (bsz, queries, 0), dtype=torch.long, device=z_a.device
+            )
+            sparse_mask = torch.empty(
+                (bsz, queries, 0), dtype=torch.bool, device=z_a.device
+            )
+
+        null_source = torch.full(
+            (bsz, queries, 1), -1, dtype=torch.long, device=z_a.device
+        )
+        null_mask = torch.ones((bsz, queries, 1), dtype=torch.bool, device=z_a.device)
+        source = torch.cat(
+            [exact_source, virtual_source, dense_source, sparse_source, null_source],
+            dim=-1,
+        )
+        mask = torch.cat(
+            [exact_mask, virtual_mask, dense_mask, sparse_mask, null_mask], dim=-1
+        )
+        kind = torch.cat(
+            [
+                torch.full_like(exact_source, EXACT_KIND),
+                torch.full_like(virtual_source, VIRTUAL_KIND),
+                torch.full_like(dense_source, VIRTUAL_KIND),
+                torch.full_like(sparse_source, VIRTUAL_KIND),
+                torch.full_like(null_source, NULL_KIND),
+            ],
+            dim=-1,
+        )
+        hard_match_length = torch.cat(
+            [
+                exact_match_length,
+                torch.zeros_like(virtual_source),
+                torch.zeros_like(dense_source),
+                torch.zeros_like(sparse_source),
+                torch.zeros_like(null_source),
+            ],
+            dim=-1,
+        )
+        frequency = torch.cat(
+            [
+                exact_frequency,
+                torch.ones_like(virtual_source),
+                torch.ones_like(dense_source),
+                torch.ones_like(sparse_source),
+                torch.zeros_like(null_source),
+            ],
+            dim=-1,
+        )
+        non_null_mask = mask & (kind != NULL_KIND)
+        soft_match = _soft_match_at_positions_torch(
+            st1,
+            st2,
+            query_positions,
+            source,
+            non_null_mask,
+            self.soft_verify_window,
+        )
+
+        safe_source = source.clamp_min(0)
+        age = (query_positions.unsqueeze(-1) - safe_source).clamp_min(0)
+        log_len = torch.log1p(hard_match_length.to(z_a.dtype))
+        log_age = torch.log1p(age.to(z_a.dtype)) / max(math.log1p(n), 1.0)
+        log_freq = torch.log1p(frequency.to(z_a.dtype))
+        is_exact = (kind == EXACT_KIND).to(z_a.dtype)
+        is_virtual = (kind == VIRTUAL_KIND).to(z_a.dtype)
+        is_null = (kind == NULL_KIND).to(z_a.dtype)
+        candidate_number = torch.arange(source.shape[-1], device=z_a.device).view(
+            1, 1, -1
+        )
+        exact_rosa_slot = hard.rosa_slot.gather(1, query_positions).unsqueeze(-1)
+        is_rosa = ((candidate_number == exact_rosa_slot) & (exact_rosa_slot >= 0)).to(
+            z_a.dtype
+        )
+        router_feature = torch.zeros_like(source, dtype=z_a.dtype)
+        router_feature[..., exact_slots : exact_slots + self.virtual_candidates] = (
+            virtual_router
+        )
+        features = torch.stack(
+            [
+                log_len,
+                log_age,
+                log_freq,
+                soft_match / float(self.soft_verify_window),
+                is_exact,
+                is_virtual,
+                is_null,
+                is_rosa,
+                router_feature,
+            ],
+            dim=-1,
+        )
+
+        query = self.selector_query(z_query).unsqueeze(-2)
+        candidate_z = _gather_sequence(z_a, source)
+        cand_key = self.selector_key(candidate_z)
+        semantic = (query * cand_key).sum(-1) / math.sqrt(self.selector_dim)
+        learned = semantic + self.feature_mlp(features).squeeze(-1)
+        learned = learned + self.kind_bias[kind]
+        learned[..., -1] = learned[..., -1] + self.null_head(z_query).squeeze(-1)
+        tie = self.tie_break_scale * (safe_source.to(z_a.dtype) + 1.0) / (n + 1.0)
+        rosa_prior = torch.where(
+            kind == EXACT_KIND,
+            hard_match_length.to(z_a.dtype) + tie,
+            torch.full_like(learned, self.virtual_prior_bias),
+        )
+        rosa_prior = torch.where(
+            kind == NULL_KIND,
+            torch.full_like(rosa_prior, self.null_prior_bias),
+            rosa_prior,
+        )
+        virtual_log_gate = torch.log(self.virtual_scale.clamp_min(1e-6))
+        legacy_virtual = torch.zeros_like(is_virtual)
+        legacy_virtual[..., exact_slots : exact_slots + self.virtual_candidates] = 1.0
+        rosa_prior = rosa_prior + legacy_virtual * virtual_log_gate
+        scores = rosa_prior + self.learned_residual_scale * learned
+        scores = scores.masked_fill(~mask, -1e9)
+
+        soft_weights = F.softmax(scores / self.retrieval_temperature, dim=-1)
+        hard_scores = scores
+        if not self.soft_candidates_forward:
+            soft_start = exact_slots + self.virtual_candidates
+            soft_end = soft_start + dense_count + sparse_count
+            soft_only = torch.zeros_like(mask)
+            soft_only[..., soft_start:soft_end] = True
+            hard_scores = scores.masked_fill(soft_only, -1e9)
+        chosen = hard_scores.argmax(dim=-1)
+        hard_weights = F.one_hot(chosen, num_classes=scores.shape[-1]).to(z_a.dtype)
+        if dense_count or sparse_count:
+            st_weights = hard_weights + (soft_weights - soft_weights.detach())
+        else:
+            st_weights = hard_weights + soft_weights - soft_weights.detach()
+
+        next_position = torch.where(non_null_mask, source + 1, torch.zeros_like(source))
+        next_st1 = _gather_sequence(st1, next_position)
+        next_st2 = _gather_sequence(st2, next_position)
+        symbolic_value = (
+            next_st1 @ self.symbol_embedding_1.weight
+            + next_st2 @ self.symbol_embedding_2.weight
+        ) * non_null_mask.unsqueeze(-1).to(z_a.dtype)
+        query_expanded = query.expand_as(cand_key)
+        value_gate = torch.sigmoid(
+            self.value_gate_head(torch.cat([query_expanded, cand_key], dim=-1))
+        ).squeeze(-1)
+        value_gate = value_gate * non_null_mask.to(z_a.dtype)
+        neural_value = self.value_proj(_gather_sequence(z_a, next_position))
+        neural_value = neural_value * value_gate.unsqueeze(-1)
+        candidate_value = symbolic_value + self.neural_value_scale * neural_value
+        if (dense_count or sparse_count) and not self.soft_candidates_forward:
+            historical_end = exact_slots + self.virtual_candidates
+            historical_scores = torch.cat(
+                [scores[..., :historical_end], scores[..., -1:]], dim=-1
+            )
+            historical_values = torch.cat(
+                [
+                    candidate_value[..., :historical_end, :],
+                    candidate_value[..., -1:, :],
+                ],
+                dim=-2,
+            )
+            historical_soft = F.softmax(
+                historical_scores / self.retrieval_temperature, dim=-1
+            )
+            historical_chosen = historical_scores.argmax(dim=-1)
+            historical_hard = F.one_hot(
+                historical_chosen, num_classes=historical_scores.shape[-1]
+            ).to(z_a.dtype)
+            historical_st = historical_hard + historical_soft - historical_soft.detach()
+            historical_retrieved = torch.sum(
+                historical_st.unsqueeze(-1) * historical_values, dim=-2
+            )
+            union_soft_retrieved = torch.sum(
+                soft_weights.unsqueeze(-1) * candidate_value, dim=-2
+            )
+            historical_soft_retrieved = torch.sum(
+                historical_soft.unsqueeze(-1) * historical_values, dim=-2
+            )
+            backward_delta = union_soft_retrieved - historical_soft_retrieved
+            retrieved = historical_retrieved + backward_delta - backward_delta.detach()
+        else:
+            retrieved = torch.sum(st_weights.unsqueeze(-1) * candidate_value, dim=-2)
+
+        read_gate = torch.sigmoid(self.read_gate_head(z_query))
+        updated_query = _gather_sequence(
+            z_b, query_positions
+        ) + read_gate * self.out_proj(retrieved)
+        update_index = query_positions.unsqueeze(-1).expand_as(updated_query)
+        updated = z_b.scatter(1, update_index, updated_query)
+        chosen_source = source.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
+        chosen_kind = kind.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
+        chosen_match = hard_match_length.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
+        chosen_next = (chosen_source + 1).clamp(min=0, max=n - 1)
+        chosen_id1 = hard_tokens // self.codebook_sizes[1]
+        chosen_id2 = hard_tokens % self.codebook_sizes[1]
+        c1 = _gather_sequence(chosen_id1.unsqueeze(-1), chosen_next).squeeze(-1)
+        c2 = _gather_sequence(chosen_id2.unsqueeze(-1), chosen_next).squeeze(-1)
+        chosen_token = c1 * self.codebook_sizes[1] + c2
+        chosen_token = torch.where(
+            chosen_kind == NULL_KIND, -torch.ones_like(chosen_token), chosen_token
+        )
+
+        eps = 1e-9
+        rosa_target = torch.where(
+            exact_rosa_slot.squeeze(-1) >= 0,
+            exact_rosa_slot.squeeze(-1),
+            torch.full_like(chosen, scores.shape[-1] - 1),
+        )
+        rosa_prob = soft_weights.gather(-1, rosa_target.unsqueeze(-1)).squeeze(-1)
+        hard_prob = soft_weights.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
+        virtual_slice = soft_weights[..., exact_slots:-1]
+        aux_losses = {
+            "rosa_distillation": -torch.log(rosa_prob.clamp_min(eps)).mean(),
+            "hard_soft_consistency": -torch.log(hard_prob.clamp_min(eps)).mean(),
+            "code_balance": 0.5 * (_balance_kl(soft1) + _balance_kl(soft2)),
+            "virtual_usage": virtual_slice.sum(-1).mean(),
+        }
+
+        scatter = self._scatter_queries
+        return ROSAOutput(
+            updated=updated,
+            retrieved=scatter(retrieved, query_positions, n, 0),
+            hard_tokens=hard_tokens,
+            code_soft=(soft1, soft2),
+            code_st=(st1, st2),
+            candidate_source_index=scatter(source, query_positions, n, -1),
+            candidate_kind=scatter(kind, query_positions, n, -1),
+            candidate_mask=scatter(mask, query_positions, n, False),
+            candidate_scores=scatter(scores, query_positions, n, -torch.inf),
+            soft_weights=scatter(soft_weights, query_positions, n, 0),
+            hard_weights=scatter(hard_weights, query_positions, n, 0),
+            chosen_candidate=scatter(chosen, query_positions, n, -1),
+            chosen_source_index=scatter(chosen_source, query_positions, n, -1),
+            chosen_token=scatter(chosen_token, query_positions, n, -1),
+            chosen_match_length=scatter(chosen_match, query_positions, n, 0),
+            chosen_is_virtual=scatter(
+                chosen_kind == VIRTUAL_KIND, query_positions, n, False
+            ),
+            hard_rosa_source_index=scatter(
+                hard.rosa_source_index.gather(1, query_positions),
+                query_positions,
+                n,
+                -1,
+            ),
+            hard_rosa_predicted_tokens=scatter(
+                hard.rosa_predicted_tokens.gather(1, query_positions),
+                query_positions,
+                n,
+                -1,
+            ),
+            hard_rosa_match_length=scatter(
+                hard.rosa_match_length.gather(1, query_positions),
+                query_positions,
+                n,
+                0,
+            ),
+            soft_match_score=scatter(soft_match, query_positions, n, 0),
+            read_gate=scatter(read_gate, query_positions, n, 0),
+            value_gate=scatter(value_gate, query_positions, n, 0),
+            aux_losses=aux_losses,
+        )
+
     def forward(
         self,
         z_a: Tensor,
         z_b: Tensor | None = None,
         code_logits: tuple[Tensor, Tensor] | None = None,
+        *,
+        query_positions: Tensor | None = None,
     ) -> ROSAOutput:
         if z_a.ndim != 3 or z_a.shape[-1] != self.d_model:
             raise ValueError("z_a must have shape [B, N, d_model]")
@@ -1077,6 +1520,19 @@ def forward(
         elif z_b.shape != z_a.shape:
             raise ValueError("z_b must have the same shape as z_a")
 
+        if query_positions is not None:
+            self._validate_query_positions(query_positions, z_a)
+            n = z_a.shape[1]
+            if query_positions.shape[1] == n:
+                full_positions = torch.arange(n, device=z_a.device).expand(
+                    z_a.shape[0], -1
+                )
+                if torch.equal(query_positions, full_positions):
+                    # Route the canonical full query set through the untouched
+                    # historical implementation, including operation order.
+                    return self.forward(z_a, z_b, code_logits)
+            return self._forward_query_only(z_a, z_b, code_logits, query_positions)
+
         (soft1, soft2), (st1, st2), hard_tokens = self.encode(z_a, code_logits)
         # Keep the exact, non-differentiable automaton on CPU as proposed for
         # RWKV-8 ROSA. Accelerator backends may optimize the tensor path around
diff --git a/tests/test_rosa.py b/tests/test_rosa.py
index abc7fc5..93dfbf3 100644
--- a/tests/test_rosa.py
+++ b/tests/test_rosa.py
@@ -661,6 +661,168 @@ def assert_nested_equal(self, actual, expected, name: str) -> None:
         else:
             self.assertEqual(actual, expected, name)
 
+    def test_query_positions_validation_and_keyword_only_signature(self) -> None:
+        model = self.make_model(candidate_backend="python")
+        z = torch.randn(2, 7, 8)
+        with self.assertRaisesRegex(TypeError, "must be a Tensor"):
+            model(z, query_positions=[[1], [2]])  # type: ignore[arg-type]
+        with self.assertRaisesRegex(TypeError, "dtype torch.long"):
+            model(z, query_positions=torch.zeros(2, 1))
+        for positions, message in (
+            (torch.zeros(2, dtype=torch.long), r"\[B, Q\]"),
+            (torch.zeros(1, 1, dtype=torch.long), r"\[B, Q\]"),
+            (torch.empty(2, 0, dtype=torch.long), "at least one"),
+            (torch.tensor([[0, 7], [1, 2]]), r"\[0, N\)"),
+            (torch.tensor([[1, 1], [2, 3]]), "unique"),
+        ):
+            with (
+                self.subTest(message=message),
+                self.assertRaisesRegex(ValueError, message),
+            ):
+                model(z, query_positions=positions)
+        with self.assertRaises(TypeError):
+            model(z, None, None, torch.tensor([[1], [2]]))
+
+    def test_query_positions_matches_full_gradients_and_sentinels(self) -> None:
+        devices = [torch.device("cpu")]
+        if torch.cuda.is_available():
+            devices.append(torch.device("cuda"))
+        discrete_fields = (
+            "hard_tokens",
+            "candidate_source_index",
+            "candidate_kind",
+            "candidate_mask",
+            "chosen_candidate",
+            "chosen_source_index",
+            "chosen_token",
+            "chosen_match_length",
+            "chosen_is_virtual",
+            "hard_rosa_source_index",
+            "hard_rosa_predicted_tokens",
+            "hard_rosa_match_length",
+        )
+        float_fields = (
+            "updated",
+            "retrieved",
+            "candidate_scores",
+            "soft_weights",
+            "hard_weights",
+            "soft_match_score",
+            "read_gate",
+            "value_gate",
+        )
+        for device in devices:
+            for backend in ("python", "stateful"):
+                with self.subTest(device=device, backend=backend):
+                    torch.manual_seed(20260819)
+                    full_model = self.make_model(
+                        candidate_backend=backend,
+                        learned_residual_scale=1.0,
+                        neural_value_scale=1.0,
+                        dense_recent_candidates=2,
+                        sparse_old_candidates=1,
+                        sparse_old_pool_size=4,
+                    ).to(device)
+                    query_model = copy.deepcopy(full_model)
+                    z_full = torch.randn(2, 17, 8, device=device, requires_grad=True)
+                    z_query = z_full.detach().clone().requires_grad_()
+                    positions = torch.tensor([[1, 6, 15], [2, 9, 16]], device=device)
+                    full = full_model(z_full)
+                    query = query_model(z_query, query_positions=positions)
+                    batch = torch.arange(2, device=device).unsqueeze(1)
+                    for name in discrete_fields:
+                        expected = getattr(full, name)
+                        actual = getattr(query, name)
+                        if name != "hard_tokens":
+                            expected = expected[batch, positions]
+                            actual = actual[batch, positions]
+                        self.assertTrue(torch.equal(actual, expected), name)
+                    for name in float_fields:
+                        expected = getattr(full, name)[batch, positions]
+                        actual = getattr(query, name)[batch, positions]
+                        torch.testing.assert_close(
+                            actual, expected, rtol=2e-6, atol=2e-7
+                        )
+
+                    query_mask = torch.zeros((2, 17), dtype=torch.bool, device=device)
+                    query_mask.scatter_(1, positions, True)
+                    non_query = ~query_mask
+                    self.assertTrue(
+                        torch.equal(query.updated[non_query], z_query[non_query])
+                    )
+                    for name in ("retrieved", "soft_weights", "hard_weights"):
+                        self.assertTrue(
+                            torch.count_nonzero(getattr(query, name)[non_query]) == 0
+                        )
+                    self.assertTrue(
+                        torch.isneginf(query.candidate_scores[non_query]).all()
+                    )
+                    self.assertTrue(
+                        (query.candidate_source_index[non_query] == -1).all()
+                    )
+                    self.assertTrue((query.candidate_kind[non_query] == -1).all())
+                    self.assertFalse(query.candidate_mask[non_query].any())
+                    for name in (
+                        "chosen_candidate",
+                        "chosen_source_index",
+                        "chosen_token",
+                        "hard_rosa_source_index",
+                        "hard_rosa_predicted_tokens",
+                    ):
+                        self.assertTrue(
+                            (getattr(query, name)[non_query] == -1).all(), name
+                        )
+                    for name in (
+                        "chosen_match_length",
+                        "hard_rosa_match_length",
+                        "soft_match_score",
+                        "read_gate",
+                        "value_gate",
+                    ):
+                        self.assertTrue(
+                            torch.count_nonzero(getattr(query, name)[non_query]) == 0,
+                            name,
+                        )
+
+                    def selected_loss(output):
+                        return (
+                            output.updated[batch, positions].square().mean()
+                            + output.retrieved[batch, positions].square().mean()
+                            + output.soft_weights[batch, positions].square().mean()
+                            + sum(item.square().mean() for item in output.code_soft)
+                        )
+
+                    selected_loss(full).backward()
+                    selected_loss(query).backward()
+                    torch.testing.assert_close(
+                        z_query.grad, z_full.grad, rtol=2e-6, atol=2e-7
+                    )
+                    for (_, full_parameter), (_, query_parameter) in zip(
+                        full_model.named_parameters(),
+                        query_model.named_parameters(),
+                        strict=True,
+                    ):
+                        self.assertEqual(
+                            full_parameter.grad is None, query_parameter.grad is None
+                        )
+                        if full_parameter.grad is not None:
+                            torch.testing.assert_close(
+                                query_parameter.grad,
+                                full_parameter.grad,
+                                rtol=2e-6,
+                                atol=2e-7,
+                            )
+
+    def test_query_positions_full_arange_is_bit_exact_legacy_route(self) -> None:
+        model = self.make_model(candidate_backend="python")
+        z_a = torch.randn(2, 9, 8)
+        z_b = torch.randn_like(z_a)
+        legacy = model(z_a, z_b, None)
+        positions = torch.arange(9).expand(2, -1)
+        query = model(z_a, z_b, None, query_positions=positions)
+        for name in legacy.__dataclass_fields__:
+            self.assert_nested_equal(getattr(query, name), getattr(legacy, name), name)
+
     def test_python_and_stateful_backends_match_all_fields_outputs_and_gradients(
         self,
     ) -> None:

From 7d728db12974d7ae1a81d0452629a83eab60d51d Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 18:14:31 +0800
Subject: [PATCH 06/10] Reuse exact prepared hard candidates

---
 src/rosa/__init__.py | 237 ++++++++++++++++++++++++++++++++++---
 tests/test_rosa.py   | 272 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 493 insertions(+), 16 deletions(-)

diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py
index a8d25ff..0073882 100644
--- a/src/rosa/__init__.py
+++ b/src/rosa/__init__.py
@@ -30,6 +30,7 @@
     "ROSAInferenceState",
     "VIRTUAL_KIND",
     "HardCandidates",
+    "PreparedHardCandidates",
     "ROSAOutput",
     "build_hard_candidates",
     "build_virtual_pool_indices",
@@ -137,6 +138,62 @@ class HardCandidates:
     rosa_predicted_tokens: Tensor
 
 
+@dataclass(frozen=True, slots=True)
+class PreparedHardCandidates:
+    """Immutable, reusable snapshot of exact hard candidate construction.
+
+    Tensor version counters and identities are captured so in-place mutation or
+    replacement of either the token snapshot or any candidate field is rejected
+    before the object can influence a forward pass.
+    """
+
+    hard_tokens: Tensor
+    candidates: HardCandidates
+    suffix_k: int
+    occurrences_r: int
+    candidate_backend: CandidateBackend
+    shape: tuple[int, int]
+    device: torch.device
+    tensor_versions: tuple[int | None, ...]
+    _tensor_ids: tuple[int, ...] = field(repr=False)
+    _tensor_snapshots: tuple[Tensor, ...] = field(repr=False)
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> PreparedHardCandidates:
+        snapshot = self.hard_tokens.detach().clone()
+        candidates = HardCandidates(
+            *(
+                getattr(self.candidates, name).detach().clone()
+                for name in HardCandidates.__dataclass_fields__
+            )
+        )
+        tensors = (snapshot,) + tuple(
+            getattr(candidates, name) for name in HardCandidates.__dataclass_fields__
+        )
+        copied = PreparedHardCandidates(
+            hard_tokens=snapshot,
+            candidates=candidates,
+            suffix_k=self.suffix_k,
+            occurrences_r=self.occurrences_r,
+            candidate_backend=self.candidate_backend,
+            shape=self.shape,
+            device=self.device,
+            tensor_versions=tuple(_tensor_version(tensor) for tensor in tensors),
+            _tensor_ids=tuple(id(tensor) for tensor in tensors),
+            _tensor_snapshots=tuple(tensor.detach().clone() for tensor in tensors),
+        )
+        memo[id(self)] = copied
+        return copied
+
+
+def _tensor_version(tensor: Tensor) -> int | None:
+    """Return a mutation counter, or ``None`` for inference tensors."""
+
+    try:
+        return tensor._version
+    except RuntimeError:
+        return None
+
+
 @dataclass
 class ROSAOutput:
     """Complete ROSA forward result."""
@@ -423,10 +480,7 @@ def _build_stateful_hard_candidates(
     try:
         candidates = prefill_candidates(state, tokens)
         result = HardCandidates(
-            *(
-                getattr(candidates, name)
-                for name in HardCandidates.__dataclass_fields__
-            )
+            *(getattr(candidates, name) for name in HardCandidates.__dataclass_fields__)
         )
         if not squeeze:
             return result
@@ -955,6 +1009,143 @@ def encode(
         hard_tokens = id1 * self.codebook_sizes[1] + id2
         return (p1, p2), (st1, st2), hard_tokens
 
+    @staticmethod
+    def _prepared_tensors(prepared: PreparedHardCandidates) -> tuple[Tensor, ...]:
+        hard = prepared.candidates
+        return (
+            prepared.hard_tokens,
+            hard.source_index,
+            hard.match_length,
+            hard.state_id,
+            hard.frequency,
+            hard.mask,
+            hard.rosa_slot,
+            hard.rosa_source_index,
+            hard.rosa_match_length,
+            hard.rosa_predicted_tokens,
+        )
+
+    def prepare_hard_candidates(self, hard_tokens: Tensor) -> PreparedHardCandidates:
+        """Build exact candidates once for reuse by compatible ROSA modules."""
+
+        if not isinstance(hard_tokens, Tensor):
+            raise TypeError("hard_tokens must be a Tensor")
+        if hard_tokens.dtype != torch.long:
+            raise TypeError("hard_tokens must have dtype torch.long")
+        if hard_tokens.ndim != 2:
+            raise ValueError("hard_tokens must have shape [B, N]")
+        if hard_tokens.shape[0] <= 0 or hard_tokens.shape[1] <= 0:
+            raise ValueError("hard_tokens dimensions must be > 0")
+        if hard_tokens.device != self.learned_residual_scale.device:
+            raise ValueError("hard_tokens must be on the same device as ROSA")
+
+        snapshot = hard_tokens.detach().clone()
+        candidates = _build_forward_hard_candidates(
+            snapshot,
+            self.suffix_k,
+            self.occurrences_r,
+            self.candidate_backend,
+        )
+        tensors = (snapshot,) + tuple(
+            getattr(candidates, name) for name in HardCandidates.__dataclass_fields__
+        )
+        return PreparedHardCandidates(
+            hard_tokens=snapshot,
+            candidates=candidates,
+            suffix_k=self.suffix_k,
+            occurrences_r=self.occurrences_r,
+            candidate_backend=self.candidate_backend,
+            shape=(int(snapshot.shape[0]), int(snapshot.shape[1])),
+            device=snapshot.device,
+            tensor_versions=tuple(_tensor_version(tensor) for tensor in tensors),
+            _tensor_ids=tuple(id(tensor) for tensor in tensors),
+            _tensor_snapshots=tuple(tensor.detach().clone() for tensor in tensors),
+        )
+
+    def _validate_prepared_hard_candidates(
+        self,
+        prepared: PreparedHardCandidates,
+        hard_tokens: Tensor,
+    ) -> HardCandidates:
+        if not isinstance(prepared, PreparedHardCandidates):
+            raise TypeError("hard_candidates must be a PreparedHardCandidates")
+        if prepared.suffix_k != self.suffix_k:
+            raise ValueError("hard_candidates suffix_k does not match ROSA")
+        if prepared.occurrences_r != self.occurrences_r:
+            raise ValueError("hard_candidates occurrences_r does not match ROSA")
+        if prepared.candidate_backend != self.candidate_backend:
+            raise ValueError("hard_candidates candidate_backend does not match ROSA")
+        expected_shape = (int(hard_tokens.shape[0]), int(hard_tokens.shape[1]))
+        if prepared.shape != expected_shape:
+            raise ValueError("hard_candidates B/N shape does not match forward")
+        if prepared.device != hard_tokens.device:
+            raise ValueError("hard_candidates device does not match forward")
+
+        tensors = self._prepared_tensors(prepared)
+        if not all(isinstance(tensor, Tensor) for tensor in tensors):
+            raise TypeError("hard_candidates fields must all be Tensors")
+        if (
+            len(prepared.tensor_versions) != len(tensors)
+            or len(prepared._tensor_ids) != len(tensors)
+            or len(prepared._tensor_snapshots) != len(tensors)
+        ):
+            raise ValueError("hard_candidates tensor metadata is invalid")
+        if (
+            prepared.hard_tokens.dtype != torch.long
+            or tuple(prepared.hard_tokens.shape) != expected_shape
+            or prepared.hard_tokens.device != prepared.device
+        ):
+            raise ValueError("hard_candidates hard_tokens metadata is invalid")
+        slots = self.suffix_k * self.occurrences_r
+        slot_shape = (*expected_shape, slots)
+        row_shape = expected_shape
+        expected = {
+            "source_index": (slot_shape, torch.long),
+            "match_length": (slot_shape, torch.long),
+            "state_id": (slot_shape, torch.long),
+            "frequency": (slot_shape, torch.long),
+            "mask": (slot_shape, torch.bool),
+            "rosa_slot": (row_shape, torch.long),
+            "rosa_source_index": (row_shape, torch.long),
+            "rosa_match_length": (row_shape, torch.long),
+            "rosa_predicted_tokens": (row_shape, torch.long),
+        }
+        for name, (shape, dtype) in expected.items():
+            tensor = getattr(prepared.candidates, name)
+            if not isinstance(tensor, Tensor):
+                raise TypeError(f"hard_candidates.{name} must be a Tensor")
+            if tuple(tensor.shape) != shape:
+                raise ValueError(f"hard_candidates.{name} has invalid shape")
+            if tensor.dtype != dtype:
+                raise TypeError(f"hard_candidates.{name} has invalid dtype")
+            if tensor.device != prepared.device:
+                raise ValueError(f"hard_candidates.{name} has invalid device")
+
+        for tensor, identity, version, snapshot in zip(
+            tensors,
+            prepared._tensor_ids,
+            prepared.tensor_versions,
+            prepared._tensor_snapshots,
+            strict=True,
+        ):
+            if not isinstance(snapshot, Tensor):
+                raise TypeError("hard_candidates integrity snapshots must be Tensors")
+            if (
+                snapshot.shape != tensor.shape
+                or snapshot.dtype != tensor.dtype
+                or snapshot.device != tensor.device
+            ):
+                raise ValueError("hard_candidates tensor metadata is invalid")
+            if (
+                id(tensor) != identity
+                or _tensor_version(tensor) != version
+                or not torch.equal(tensor, snapshot)
+            ):
+                raise ValueError("hard_candidates was mutated")
+        if not torch.equal(prepared.hard_tokens, hard_tokens):
+            raise ValueError("hard_candidates is stale for encoded hard_tokens")
+        return prepared.candidates
+
     def _soft_match(
         self,
         st1: Tensor,
@@ -1130,15 +1321,20 @@ def _forward_query_only(
         z_b: Tensor,
         code_logits: tuple[Tensor, Tensor] | None,
         query_positions: Tensor,
+        hard_candidates: PreparedHardCandidates | None,
     ) -> ROSAOutput:
         """Run candidate activations only at selected sequence positions."""
 
         (soft1, soft2), (st1, st2), hard_tokens = self.encode(z_a, code_logits)
-        hard = _build_forward_hard_candidates(
-            hard_tokens,
-            self.suffix_k,
-            self.occurrences_r,
-            self.candidate_backend,
+        hard = (
+            _build_forward_hard_candidates(
+                hard_tokens,
+                self.suffix_k,
+                self.occurrences_r,
+                self.candidate_backend,
+            )
+            if hard_candidates is None
+            else self._validate_prepared_hard_candidates(hard_candidates, hard_tokens)
         )
         bsz, n, _ = z_a.shape
         queries = query_positions.shape[1]
@@ -1510,6 +1706,7 @@ def forward(
         code_logits: tuple[Tensor, Tensor] | None = None,
         *,
         query_positions: Tensor | None = None,
+        hard_candidates: PreparedHardCandidates | None = None,
     ) -> ROSAOutput:
         if z_a.ndim != 3 or z_a.shape[-1] != self.d_model:
             raise ValueError("z_a must have shape [B, N, d_model]")
@@ -1530,18 +1727,26 @@ def forward(
                 if torch.equal(query_positions, full_positions):
                     # Route the canonical full query set through the untouched
                     # historical implementation, including operation order.
-                    return self.forward(z_a, z_b, code_logits)
-            return self._forward_query_only(z_a, z_b, code_logits, query_positions)
+                    return self.forward(
+                        z_a, z_b, code_logits, hard_candidates=hard_candidates
+                    )
+            return self._forward_query_only(
+                z_a, z_b, code_logits, query_positions, hard_candidates
+            )
 
         (soft1, soft2), (st1, st2), hard_tokens = self.encode(z_a, code_logits)
         # Keep the exact, non-differentiable automaton on CPU as proposed for
         # RWKV-8 ROSA. Accelerator backends may optimize the tensor path around
         # it, but must not silently replace this exact discrete control path.
-        hard = _build_forward_hard_candidates(
-            hard_tokens,
-            self.suffix_k,
-            self.occurrences_r,
-            self.candidate_backend,
+        hard = (
+            _build_forward_hard_candidates(
+                hard_tokens,
+                self.suffix_k,
+                self.occurrences_r,
+                self.candidate_backend,
+            )
+            if hard_candidates is None
+            else self._validate_prepared_hard_candidates(hard_candidates, hard_tokens)
         )
         exact_source = hard.source_index
         exact_mask = hard.mask
diff --git a/tests/test_rosa.py b/tests/test_rosa.py
index 93dfbf3..414487e 100644
--- a/tests/test_rosa.py
+++ b/tests/test_rosa.py
@@ -7,6 +7,7 @@
 import unittest
 import weakref
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import FrozenInstanceError
 from unittest.mock import patch
 
 import numpy as np
@@ -18,6 +19,7 @@
     NULL_KIND,
     ROSA,
     VIRTUAL_KIND,
+    PreparedHardCandidates,
     _balance_kl,
     _build_forward_hard_candidates,
     _clear_soft_match_compile_cache,
@@ -661,6 +663,276 @@ def assert_nested_equal(self, actual, expected, name: str) -> None:
         else:
             self.assertEqual(actual, expected, name)
 
+    def test_prepared_candidates_are_bit_exact_and_shared_without_rebuild(self) -> None:
+        torch.manual_seed(20260819)
+        baseline_model = self.make_model(
+            candidate_backend="python",
+            learned_residual_scale=1.0,
+            neural_value_scale=1.0,
+        )
+        prepared_model = copy.deepcopy(baseline_model)
+        second_consumer = copy.deepcopy(baseline_model)
+        tokens = torch.randint(6, (2, 13))
+        baseline_z = torch.randn(2, 13, 8, requires_grad=True)
+        prepared_z = baseline_z.detach().clone().requires_grad_()
+        logits_base = factor_logits_from_tokens(
+            tokens, (2, 3), hi=0.4, lo=-0.2, requires_grad=True
+        )
+        logits_prepared = (
+            logits_base[0].detach().clone().requires_grad_(),
+            logits_base[1].detach().clone().requires_grad_(),
+        )
+
+        baseline = baseline_model(baseline_z, code_logits=logits_base)
+        baseline_loss = baseline.updated.square().sum() + sum(
+            value.square().sum() for value in baseline.aux_losses.values()
+        )
+        baseline_loss.backward()
+
+        calls = 0
+        original = rosa._build_forward_hard_candidates
+
+        def counted(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            return original(*args, **kwargs)
+
+        with patch("rosa._build_forward_hard_candidates", side_effect=counted):
+            hard_tokens = prepared_model.encode(prepared_z, logits_prepared)[2]
+            prepared = prepared_model.prepare_hard_candidates(hard_tokens)
+            self.assertIsInstance(prepared, PreparedHardCandidates)
+            self.assertFalse(prepared.hard_tokens.requires_grad)
+            self.assertEqual(calls, 1)
+            actual = prepared_model(
+                prepared_z, code_logits=logits_prepared, hard_candidates=prepared
+            )
+            # Different logits with the same argmax remain valid for another model.
+            shifted_logits = tuple(value + 0.125 for value in logits_prepared)
+            shared = second_consumer(
+                prepared_z.detach(),
+                code_logits=shifted_logits,
+                hard_candidates=prepared,
+            )
+            self.assertEqual(calls, 1)
+
+        for name in baseline.__dataclass_fields__:
+            self.assert_nested_equal(
+                getattr(actual, name), getattr(baseline, name), name
+            )
+        self.assertTrue(torch.equal(shared.hard_tokens, actual.hard_tokens))
+        self.assertTrue(
+            torch.equal(shared.candidate_source_index, actual.candidate_source_index)
+        )
+        actual_loss = actual.updated.square().sum() + sum(
+            value.square().sum() for value in actual.aux_losses.values()
+        )
+        self.assertTrue(torch.equal(actual_loss, baseline_loss))
+        actual_loss.backward()
+        prepared_z_grad = prepared_z.grad
+        baseline_z_grad = baseline_z.grad
+        assert prepared_z_grad is not None
+        assert baseline_z_grad is not None
+        self.assertTrue(torch.equal(prepared_z_grad, baseline_z_grad))
+        for actual_logit, baseline_logit in zip(
+            logits_prepared, logits_base, strict=True
+        ):
+            actual_logit_grad = actual_logit.grad
+            baseline_logit_grad = baseline_logit.grad
+            assert actual_logit_grad is not None
+            assert baseline_logit_grad is not None
+            self.assertTrue(torch.equal(actual_logit_grad, baseline_logit_grad))
+        for (_, actual_parameter), (_, baseline_parameter) in zip(
+            prepared_model.named_parameters(),
+            baseline_model.named_parameters(),
+            strict=True,
+        ):
+            self.assertEqual(
+                actual_parameter.grad is None, baseline_parameter.grad is None
+            )
+            actual_parameter_grad = actual_parameter.grad
+            baseline_parameter_grad = baseline_parameter.grad
+            if actual_parameter_grad is not None:
+                assert baseline_parameter_grad is not None
+                self.assertTrue(
+                    torch.equal(actual_parameter_grad, baseline_parameter_grad)
+                )
+
+    def test_prepared_candidates_combine_with_query_positions(self) -> None:
+        model = self.make_model(candidate_backend="python", learned_residual_scale=1.0)
+        z = torch.randn(2, 11, 8)
+        tokens = torch.randint(6, (2, 11))
+        logits = factor_logits_from_tokens(tokens, (2, 3))
+        prepared = model.prepare_hard_candidates(model.encode(z, logits)[2])
+        positions = torch.tensor([[1, 5, 10], [0, 4, 8]])
+        expected = model(z, code_logits=logits, query_positions=positions)
+        with patch(
+            "rosa._build_forward_hard_candidates",
+            side_effect=AssertionError("unexpected rebuild"),
+        ):
+            actual = model(
+                z,
+                code_logits=logits,
+                query_positions=positions,
+                hard_candidates=prepared,
+            )
+        for name in expected.__dataclass_fields__:
+            self.assert_nested_equal(
+                getattr(actual, name), getattr(expected, name), name
+            )
+
+    def test_prepared_candidates_strict_rejections(self) -> None:
+        model = self.make_model(candidate_backend="python")
+        z = torch.randn(2, 9, 8)
+        tokens = torch.randint(6, (2, 9))
+        logits = factor_logits_from_tokens(tokens, (2, 3))
+
+        def prepare():
+            return model.prepare_hard_candidates(model.encode(z, logits)[2])
+
+        def forged(field, replacement):
+            base = prepare()
+            candidates = copy.deepcopy(base.candidates)
+            setattr(candidates, field, replacement)
+            snapshot = base.hard_tokens.clone()
+            tensors = (snapshot,) + tuple(
+                getattr(candidates, name) for name in candidates.__dataclass_fields__
+            )
+            return PreparedHardCandidates(
+                snapshot,
+                candidates,
+                base.suffix_k,
+                base.occurrences_r,
+                base.candidate_backend,
+                base.shape,
+                base.device,
+                tuple(tensor._version for tensor in tensors),
+                tuple(id(tensor) for tensor in tensors),
+                tuple(tensor.clone() for tensor in tensors),
+            )
+
+        with self.assertRaisesRegex(TypeError, "PreparedHardCandidates"):
+            model(z, code_logits=logits, hard_candidates=prepare().candidates)  # type: ignore[arg-type]
+        with self.assertRaises(FrozenInstanceError):
+            prepare().suffix_k = 99  # type: ignore[misc]
+
+        stale_logits = factor_logits_from_tokens(tokens.clone(), (2, 3))
+        # Force one encoded token to differ without changing shapes.
+        stale_logits[0].data[0, 0].fill_(-20.0)
+        stale_id = (int(tokens[0, 0] // 3) + 1) % 2
+        stale_logits[0].data[0, 0, stale_id] = 20.0
+        with self.assertRaisesRegex(ValueError, "stale"):
+            model(z, code_logits=stale_logits, hard_candidates=prepare())
+
+        for attribute, value, message in (
+            ("suffix_k", model.suffix_k + 1, "suffix_k"),
+            ("occurrences_r", model.occurrences_r + 1, "occurrences_r"),
+            ("candidate_backend", "stateful", "candidate_backend"),
+        ):
+            consumer = copy.deepcopy(model)
+            setattr(consumer, attribute, value)
+            with self.assertRaisesRegex(ValueError, message):
+                consumer(z, code_logits=logits, hard_candidates=prepare())
+
+        short_z = z[:, :-1]
+        short_logits = tuple(value[:, :-1] for value in logits)
+        with self.assertRaisesRegex(ValueError, "B/N"):
+            model(short_z, code_logits=short_logits, hard_candidates=prepare())
+
+        with self.assertRaisesRegex(ValueError, "source_index.*shape"):
+            model(
+                z,
+                code_logits=logits,
+                hard_candidates=forged(
+                    "source_index", torch.zeros(2, 9, 14, dtype=torch.long)
+                ),
+            )
+        with self.assertRaisesRegex(TypeError, "match_length.*dtype"):
+            model(
+                z,
+                code_logits=logits,
+                hard_candidates=forged(
+                    "match_length", torch.zeros(2, 9, 15, dtype=torch.float32)
+                ),
+            )
+        with self.assertRaisesRegex(ValueError, "source_index.*device"):
+            model(
+                z,
+                code_logits=logits,
+                hard_candidates=forged(
+                    "source_index",
+                    torch.empty(2, 9, 15, dtype=torch.long, device="meta"),
+                ),
+            )
+
+        for field, replacement in (
+            ("source_index", torch.zeros(2, 9, 1, dtype=torch.long)),
+            ("match_length", torch.zeros(2, 9, 15, dtype=torch.float32)),
+        ):
+            prepared = prepare()
+            setattr(prepared.candidates, field, replacement)
+            with self.assertRaises((TypeError, ValueError)):
+                model(z, code_logits=logits, hard_candidates=prepared)
+
+        prepared = prepare()
+        prepared.candidates.mask.logical_not_()
+        with self.assertRaisesRegex(ValueError, "mutated"):
+            model(z, code_logits=logits, hard_candidates=prepared)
+        prepared = prepare()
+        prepared.candidates.mask.data.logical_not_()
+        with self.assertRaisesRegex(ValueError, "mutated"):
+            model(z, code_logits=logits, hard_candidates=prepared)
+        prepared = prepare()
+        prepared.candidates.source_index.numpy()[0, 0, 0] = 42
+        with self.assertRaisesRegex(ValueError, "mutated"):
+            model(z, code_logits=logits, hard_candidates=prepared)
+        prepared = prepare()
+        prepared.hard_tokens.add_(1)
+        with self.assertRaisesRegex(ValueError, "mutated"):
+            model(z, code_logits=logits, hard_candidates=prepared)
+
+        if torch.cuda.is_available():
+            cuda_model = copy.deepcopy(model).cuda()
+            with self.assertRaisesRegex(ValueError, "device"):
+                cuda_model(
+                    z.cuda(),
+                    code_logits=tuple(value.cuda() for value in logits),
+                    hard_candidates=prepare(),
+                )
+
+    def test_prepared_candidates_have_explicit_external_ownership(self) -> None:
+        model = self.make_model(candidate_backend="python")
+        z = torch.randn(1, 8, 8)
+        tokens = torch.randint(6, (1, 8))
+        logits = factor_logits_from_tokens(tokens, (2, 3))
+        prepared = model.prepare_hard_candidates(model.encode(z, logits)[2])
+        state_keys = set(model.state_dict())
+        self.assertFalse(
+            any("prepared" in key or "candidate" in key for key in state_keys)
+        )
+        copied_model = copy.deepcopy(model)
+        copied_prepared = copy.deepcopy(prepared)
+        self.assertIsNot(copied_prepared.hard_tokens, prepared.hard_tokens)
+        copied_model(z, code_logits=logits, hard_candidates=copied_prepared)
+        model.to("cpu")
+        self.assertEqual(prepared.device.type, "cpu")
+        model(z, code_logits=logits, hard_candidates=prepared)
+
+    def test_prepared_candidates_support_inference_tensors(self) -> None:
+        model = self.make_model(candidate_backend="python")
+        z = torch.randn(1, 8, 8)
+        tokens = torch.randint(6, (1, 8))
+        logits = factor_logits_from_tokens(tokens, (2, 3))
+        with torch.inference_mode():
+            hard_tokens = model.encode(z, logits)[2]
+            prepared = model.prepare_hard_candidates(hard_tokens)
+            copied = copy.deepcopy(prepared)
+            expected = model(z, code_logits=logits)
+            actual = model(z, code_logits=logits, hard_candidates=copied)
+        for name in expected.__dataclass_fields__:
+            self.assert_nested_equal(
+                getattr(actual, name), getattr(expected, name), name
+            )
+
     def test_query_positions_validation_and_keyword_only_signature(self) -> None:
         model = self.make_model(candidate_backend="python")
         z = torch.randn(2, 7, 8)

From 6d48fb3a6b1ec7db1d1413d5d65041581b1e63ec Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 18:58:59 +0800
Subject: [PATCH 07/10] Support minimal exact candidate budgets

---
 src/rosa/__init__.py | 123 +++++++++++++++++++++++--------------
 tests/test_rosa.py   | 143 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 220 insertions(+), 46 deletions(-)

diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py
index 0073882..3ff1695 100644
--- a/src/rosa/__init__.py
+++ b/src/rosa/__init__.py
@@ -810,7 +810,7 @@ def _st_categorical(
     soft = F.softmax(logits / temperature, dim=-1)
     ids = logits.argmax(dim=-1)
     hard = F.one_hot(ids, num_classes=logits.shape[-1]).to(logits.dtype)
-    st = hard + soft - soft.detach()
+    st = hard + (soft - soft.detach())
     return soft, st, ids
 
 
@@ -876,8 +876,8 @@ def __init__(
             raise ValueError(
                 "suffix_k, occurrences_r and soft_verify_window must be > 0"
             )
-        if virtual_candidates <= 0 or virtual_pool_size < virtual_candidates:
-            raise ValueError("virtual_pool_size must be >= virtual_candidates > 0")
+        if virtual_candidates < 0 or virtual_pool_size < virtual_candidates:
+            raise ValueError("virtual_pool_size must be >= virtual_candidates >= 0")
         if dense_recent_candidates < 0:
             raise ValueError("dense_recent_candidates must be >= 0")
         if sparse_old_candidates < 0 or sparse_old_pool_size < sparse_old_candidates:
@@ -1169,6 +1169,13 @@ def _virtual_candidates(
         exact_mask: Tensor,
     ) -> tuple[Tensor, Tensor, Tensor]:
         bsz, n, _ = z_a.shape
+        if self.virtual_candidates == 0:
+            shape = (bsz, n, 0)
+            return (
+                torch.empty(shape, dtype=torch.long, device=z_a.device),
+                torch.empty(shape, dtype=torch.bool, device=z_a.device),
+                torch.empty(shape, dtype=z_a.dtype, device=z_a.device),
+            )
         pool = build_virtual_pool_indices(bsz, n, self.virtual_pool_size, z_a.device)
         pool_mask = pool >= 0
         duplicate = (
@@ -1315,6 +1322,17 @@ def _scatter_queries(
         ).expand_as(values)
         return result.scatter(1, index, values)
 
+    def _attach_skipped_virtual_gradients(
+        self, retrieved: Tensor, updated: Tensor
+    ) -> tuple[Tensor, Tensor]:
+        if self.virtual_candidates or not torch.is_grad_enabled():
+            return retrieved, updated
+        virtual_zero = retrieved.new_zeros(())
+        for module in (self.virtual_query, self.virtual_key):
+            for parameter in module.parameters():
+                virtual_zero = virtual_zero + parameter.sum() * 0
+        return retrieved + virtual_zero, updated + virtual_zero
+
     def _forward_query_only(
         self,
         z_a: Tensor,
@@ -1345,29 +1363,37 @@ def _forward_query_only(
         exact_frequency = _gather_sequence(hard.frequency, query_positions)
         exact_slots = exact_source.shape[-1]
 
-        # The pool is discrete and cheap to build for N; floating candidate
-        # projections below are restricted to [B, Q, C, *].
-        virtual_pool = _gather_sequence(
-            build_virtual_pool_indices(bsz, n, self.virtual_pool_size, z_a.device),
-            query_positions,
-        )
-        virtual_pool_mask = virtual_pool >= 0
-        virtual_duplicate = (
-            virtual_pool.unsqueeze(-1) == exact_source.unsqueeze(-2)
-        ) & exact_mask.unsqueeze(-2)
-        virtual_pool_mask = virtual_pool_mask & ~virtual_duplicate.any(dim=-1)
-        virtual_query = self.virtual_query(z_query).unsqueeze(-2)
-        virtual_keys = self.virtual_key(_gather_sequence(z_a, virtual_pool))
-        virtual_router_all = (virtual_query * virtual_keys).sum(-1) / math.sqrt(
-            self.selector_dim
-        )
-        virtual_masked = virtual_router_all.masked_fill(~virtual_pool_mask, -1e9)
-        _, virtual_top = torch.topk(virtual_masked, k=self.virtual_candidates, dim=-1)
-        virtual_source = virtual_pool.gather(-1, virtual_top)
-        virtual_mask = virtual_pool_mask.gather(-1, virtual_top)
-        virtual_router = virtual_router_all.gather(-1, virtual_top)
-        virtual_mask = virtual_mask & (self.virtual_scale > 0)
-        virtual_router = virtual_router * self.virtual_scale
+        if self.virtual_candidates:
+            # The pool is discrete and cheap to build for N; floating candidate
+            # projections below are restricted to [B, Q, C, *].
+            virtual_pool = _gather_sequence(
+                build_virtual_pool_indices(bsz, n, self.virtual_pool_size, z_a.device),
+                query_positions,
+            )
+            virtual_pool_mask = virtual_pool >= 0
+            virtual_duplicate = (
+                virtual_pool.unsqueeze(-1) == exact_source.unsqueeze(-2)
+            ) & exact_mask.unsqueeze(-2)
+            virtual_pool_mask = virtual_pool_mask & ~virtual_duplicate.any(dim=-1)
+            virtual_query = self.virtual_query(z_query).unsqueeze(-2)
+            virtual_keys = self.virtual_key(_gather_sequence(z_a, virtual_pool))
+            virtual_router_all = (virtual_query * virtual_keys).sum(-1) / math.sqrt(
+                self.selector_dim
+            )
+            virtual_masked = virtual_router_all.masked_fill(~virtual_pool_mask, -1e9)
+            _, virtual_top = torch.topk(
+                virtual_masked, k=self.virtual_candidates, dim=-1
+            )
+            virtual_source = virtual_pool.gather(-1, virtual_top)
+            virtual_mask = virtual_pool_mask.gather(-1, virtual_top)
+            virtual_router = virtual_router_all.gather(-1, virtual_top)
+            virtual_mask = virtual_mask & (self.virtual_scale > 0)
+            virtual_router = virtual_router * self.virtual_scale
+        else:
+            shape = (bsz, queries, 0)
+            virtual_source = torch.empty(shape, dtype=torch.long, device=z_a.device)
+            virtual_mask = torch.empty(shape, dtype=torch.bool, device=z_a.device)
+            virtual_router = torch.empty(shape, dtype=z_a.dtype, device=z_a.device)
 
         dense_count = self.dense_recent_candidates
         sparse_count = self.sparse_old_candidates
@@ -1566,10 +1592,7 @@ def _forward_query_only(
             hard_scores = scores.masked_fill(soft_only, -1e9)
         chosen = hard_scores.argmax(dim=-1)
         hard_weights = F.one_hot(chosen, num_classes=scores.shape[-1]).to(z_a.dtype)
-        if dense_count or sparse_count:
-            st_weights = hard_weights + (soft_weights - soft_weights.detach())
-        else:
-            st_weights = hard_weights + soft_weights - soft_weights.detach()
+        st_weights = hard_weights + (soft_weights - soft_weights.detach())
 
         next_position = torch.where(non_null_mask, source + 1, torch.zeros_like(source))
         next_st1 = _gather_sequence(st1, next_position)
@@ -1605,7 +1628,9 @@ def _forward_query_only(
             historical_hard = F.one_hot(
                 historical_chosen, num_classes=historical_scores.shape[-1]
             ).to(z_a.dtype)
-            historical_st = historical_hard + historical_soft - historical_soft.detach()
+            historical_st = historical_hard + (
+                historical_soft - historical_soft.detach()
+            )
             historical_retrieved = torch.sum(
                 historical_st.unsqueeze(-1) * historical_values, dim=-2
             )
@@ -1616,7 +1641,9 @@ def _forward_query_only(
                 historical_soft.unsqueeze(-1) * historical_values, dim=-2
             )
             backward_delta = union_soft_retrieved - historical_soft_retrieved
-            retrieved = historical_retrieved + backward_delta - backward_delta.detach()
+            retrieved = historical_retrieved + (
+                backward_delta - backward_delta.detach()
+            )
         else:
             retrieved = torch.sum(st_weights.unsqueeze(-1) * candidate_value, dim=-2)
 
@@ -1626,6 +1653,7 @@ def _forward_query_only(
         ) + read_gate * self.out_proj(retrieved)
         update_index = query_positions.unsqueeze(-1).expand_as(updated_query)
         updated = z_b.scatter(1, update_index, updated_query)
+        retrieved, updated = self._attach_skipped_virtual_gradients(retrieved, updated)
         chosen_source = source.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
         chosen_kind = kind.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
         chosen_match = hard_match_length.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
@@ -1752,11 +1780,17 @@ def forward(
         exact_mask = hard.mask
         exact_slots = exact_source.shape[-1]
 
-        virtual_source, virtual_mask, virtual_router = self._virtual_candidates(
-            z_a, exact_source, exact_mask
-        )
-        virtual_mask = virtual_mask & (self.virtual_scale > 0)
-        virtual_router = virtual_router * self.virtual_scale
+        if self.virtual_candidates:
+            virtual_source, virtual_mask, virtual_router = self._virtual_candidates(
+                z_a, exact_source, exact_mask
+            )
+            virtual_mask = virtual_mask & (self.virtual_scale > 0)
+            virtual_router = virtual_router * self.virtual_scale
+        else:
+            shape = (*exact_source.shape[:2], 0)
+            virtual_source = torch.empty(shape, dtype=torch.long, device=z_a.device)
+            virtual_mask = torch.empty(shape, dtype=torch.bool, device=z_a.device)
+            virtual_router = torch.empty(shape, dtype=z_a.dtype, device=z_a.device)
 
         dense_source, dense_mask, sparse_source, sparse_mask = (
             self._hybrid_soft_candidates(st1, st2, exact_source, exact_mask)
@@ -1894,13 +1928,9 @@ def forward(
             hard_scores = scores.masked_fill(soft_only, -1e9)
         chosen = hard_scores.argmax(dim=-1)
         hard_weights = F.one_hot(chosen, num_classes=scores.shape[-1]).to(z_a.dtype)
-        if self.dense_recent_candidates or self.sparse_old_candidates:
-            # Parenthesizing the zero-valued correction makes the forward
-            # exactly one-hot while retaining the full-union softmax backward.
-            st_weights = hard_weights + (soft_weights - soft_weights.detach())
-        else:
-            # Preserve the historical arithmetic when both new budgets are 0.
-            st_weights = hard_weights + soft_weights - soft_weights.detach()
+        # Parenthesizing the zero-valued correction makes the forward exactly
+        # one-hot while retaining the full-union softmax backward.
+        st_weights = hard_weights + (soft_weights - soft_weights.detach())
 
         next_position = torch.where(non_null_mask, source + 1, torch.zeros_like(source))
         symbolic_value = self._candidate_symbolic_values(
@@ -1940,7 +1970,9 @@ def forward(
             historical_hard = F.one_hot(
                 historical_chosen, num_classes=historical_scores.shape[-1]
             ).to(z_a.dtype)
-            historical_st = historical_hard + historical_soft - historical_soft.detach()
+            historical_st = historical_hard + (
+                historical_soft - historical_soft.detach()
+            )
             historical_retrieved = torch.sum(
                 historical_st.unsqueeze(-1) * historical_values, dim=-2
             )
@@ -1959,6 +1991,7 @@ def forward(
 
         read_gate = torch.sigmoid(self.read_gate_head(z_a))
         updated = z_b + read_gate * self.out_proj(retrieved)
+        retrieved, updated = self._attach_skipped_virtual_gradients(retrieved, updated)
 
         chosen_source = source.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
         chosen_kind = kind.gather(-1, chosen.unsqueeze(-1)).squeeze(-1)
diff --git a/tests/test_rosa.py b/tests/test_rosa.py
index 414487e..259624e 100644
--- a/tests/test_rosa.py
+++ b/tests/test_rosa.py
@@ -531,7 +531,7 @@ def test_constructor_validations(self) -> None:
             (dict(d_model=4, suffix_k=0), "suffix_k"),
             (dict(d_model=4, occurrences_r=0), "suffix_k"),
             (dict(d_model=4, soft_verify_window=0), "suffix_k"),
-            (dict(d_model=4, virtual_candidates=0), "virtual_pool_size"),
+            (dict(d_model=4, virtual_candidates=-1), "virtual_pool_size"),
             (
                 dict(d_model=4, virtual_candidates=4, virtual_pool_size=3),
                 "virtual_pool_size",
@@ -663,6 +663,147 @@ def assert_nested_equal(self, actual, expected, name: str) -> None:
         else:
             self.assertEqual(actual, expected, name)
 
+    def test_zero_virtual_candidates_skip_virtual_work_and_have_exact_width(
+        self,
+    ) -> None:
+        model = self.make_model(
+            suffix_k=2,
+            occurrences_r=3,
+            virtual_candidates=0,
+            dense_recent_candidates=2,
+            sparse_old_candidates=1,
+            sparse_old_pool_size=3,
+        )
+        z = torch.randn(2, 9, 8)
+        positions = torch.tensor([[1, 5], [2, 8]])
+        virtual_calls = 0
+
+        def virtual_hook(_module, _args, _output) -> None:
+            nonlocal virtual_calls
+            virtual_calls += 1
+
+        hooks = [
+            model.virtual_query.register_forward_hook(virtual_hook),
+            model.virtual_key.register_forward_hook(virtual_hook),
+        ]
+        try:
+            with (
+                patch.object(
+                    model,
+                    "_virtual_candidates",
+                    side_effect=AssertionError("unexpected virtual candidates"),
+                ),
+                patch(
+                    "rosa.build_virtual_pool_indices",
+                    side_effect=AssertionError("unexpected virtual pool"),
+                ),
+                patch(
+                    "rosa.torch.topk",
+                    side_effect=AssertionError("unexpected virtual topk"),
+                ),
+            ):
+                full = model(z)
+                query_only = model(z, query_positions=positions)
+
+                for output_name, query in (
+                    ("retrieved", None),
+                    ("updated", None),
+                    ("retrieved", positions),
+                    ("updated", positions),
+                ):
+                    model.zero_grad(set_to_none=True)
+                    output = model(z, query_positions=query)
+                    getattr(output, output_name).sum().backward()
+                    for parameter in (
+                        model.virtual_query.weight,
+                        model.virtual_key.weight,
+                    ):
+                        self.assertIsNotNone(parameter.grad)
+                        assert parameter.grad is not None
+                        self.assertEqual(torch.count_nonzero(parameter.grad).item(), 0)
+
+                with torch.no_grad():
+                    no_grad_full = model(z)
+                    no_grad_query = model(z, query_positions=positions)
+                with torch.inference_mode():
+                    inference_full = model(z)
+                    inference_query = model(z, query_positions=positions)
+        finally:
+            for hook in hooks:
+                hook.remove()
+
+        expected_width = 2 * 3 + 2 + 1 + 1
+        self.assertEqual(full.candidate_source_index.shape[-1], expected_width)
+        self.assertEqual(query_only.candidate_source_index.shape[-1], expected_width)
+        self.assertEqual(full.value_gate.shape[-1], expected_width)
+        self.assertEqual(query_only.value_gate.shape[-1], expected_width)
+        self.assertEqual(virtual_calls, 0)
+        for reference, no_grad, inference in (
+            (full, no_grad_full, inference_full),
+            (query_only, no_grad_query, inference_query),
+        ):
+            for name in ("retrieved", "updated"):
+                self.assertTrue(
+                    torch.equal(getattr(reference, name), getattr(no_grad, name))
+                )
+                self.assertTrue(
+                    torch.equal(getattr(reference, name), getattr(inference, name))
+                )
+
+    def test_zero_virtual_checkpoint_compatibility_and_runtime_reactivation(
+        self,
+    ) -> None:
+        source = self.make_model(
+            suffix_k=4,
+            occurrences_r=4,
+            virtual_candidates=1,
+        )
+        target = self.make_model(
+            suffix_k=1,
+            occurrences_r=1,
+            virtual_candidates=0,
+        )
+        incompatible = target.load_state_dict(source.state_dict(), strict=True)
+        self.assertEqual(incompatible.missing_keys, [])
+        self.assertEqual(incompatible.unexpected_keys, [])
+        self.assertIn("virtual_query.weight", target.state_dict())
+        self.assertIn("virtual_key.weight", target.state_dict())
+
+        z = torch.randn(1, 8, 8)
+        hard_tokens = target.encode(z)[2]
+        prepared_k1r1 = target.prepare_hard_candidates(hard_tokens)
+        prepared_k4r4 = source.prepare_hard_candidates(hard_tokens)
+        with self.assertRaisesRegex(ValueError, "suffix_k"):
+            target(z, hard_candidates=prepared_k4r4)
+
+        target.suffix_k = 2
+        with self.assertRaisesRegex(ValueError, "suffix_k"):
+            target(z, hard_candidates=prepared_k1r1)
+        target.suffix_k = 1
+        target.occurrences_r = 2
+        with self.assertRaisesRegex(ValueError, "occurrences_r"):
+            target(z, hard_candidates=prepared_k1r1)
+        target.occurrences_r = 1
+
+        calls = 0
+
+        def virtual_hook(_module, _args, _output) -> None:
+            nonlocal calls
+            calls += 1
+
+        hooks = [
+            target.virtual_query.register_forward_hook(virtual_hook),
+            target.virtual_key.register_forward_hook(virtual_hook),
+        ]
+        try:
+            target.virtual_candidates = 1
+            output = target(z)
+        finally:
+            for hook in hooks:
+                hook.remove()
+        self.assertEqual(output.candidate_source_index.shape[-1], 1 * 1 + 1 + 1)
+        self.assertEqual(calls, 2)
+
     def test_prepared_candidates_are_bit_exact_and_shared_without_rebuild(self) -> None:
         torch.manual_seed(20260819)
         baseline_model = self.make_model(

From c62a095ad503e46cb21d69ae9b9840d589fc8063 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 19:21:35 +0800
Subject: [PATCH 08/10] Emit exact candidates only at query positions

---
 native/src/rosa_native_step.cpp        | 110 ++++++++++-
 native/tests/candidate_smoke.py        |  22 +++
 src/rosa/__init__.py                   | 130 ++++++++++++-
 src/rosa/_stateful_candidates_numba.py | 260 +++++++++++++++++++++----
 tests/test_rosa.py                     |  25 +++
 tests/test_stateful_candidates.py      | 121 ++++++++++++
 6 files changed, 624 insertions(+), 44 deletions(-)

diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp
index c63b6bb..b752ad5 100644
--- a/native/src/rosa_native_step.cpp
+++ b/native/src/rosa_native_step.cpp
@@ -1210,6 +1210,7 @@ class NativeCandidateState {
       parallel_for_rows(64, [&](int64_t b) {
         count.mutable_data()[b] =
             step_row(b, tokens.data()[b], position_,
+                     true,
                      source.mutable_data() + b * slots,
                      match_length.mutable_data() + b * slots,
                      state_id.mutable_data() + b * slots,
@@ -1298,6 +1299,7 @@ class NativeCandidateState {
           reset_row(b);
         count.mutable_data()[b] = step_row(
             b, tokens.data()[b], current_position,
+            true,
             source.mutable_data() + b * slots,
             match_length.mutable_data() + b * slots,
             state_id.mutable_data() + b * slots,
@@ -1409,6 +1411,7 @@ class NativeCandidateState {
           const int64_t output_at = (b * sequence_length + position) * slots;
           count.mutable_data()[b * sequence_length + position] = step_row(
               b, tokens.data()[b * sequence_length + position], position,
+              true,
               source.mutable_data() + output_at,
               match_length.mutable_data() + output_at,
               state_id.mutable_data() + output_at,
@@ -1423,6 +1426,109 @@ class NativeCandidateState {
     call_lock.unlock();
   }
 
+  py::tuple prefill_selected(py::array tokens_object,
+                             py::array query_positions_object) {
+    if (ragged_mode_)
+      throw std::runtime_error(
+          "prefill is unavailable on a ragged candidate state");
+    if (!py::isinstance<py::array_t<int64_t>>(tokens_object) ||
+        tokens_object.ndim() != 2 || tokens_object.shape(0) != batch_)
+      throw py::value_error(
+          "tokens must be contiguous int64 [batch_size, sequence_length]");
+    if (!py::isinstance<py::array_t<int64_t>>(query_positions_object) ||
+        query_positions_object.ndim() != 2 ||
+        query_positions_object.shape(0) != batch_)
+      throw py::value_error(
+          "query_positions must be contiguous int64 [batch_size, query_count]");
+    const int64_t sequence_length = tokens_object.shape(1);
+    const int64_t query_count = query_positions_object.shape(1);
+    if (query_count <= 0)
+      throw py::value_error(
+          "query_positions must contain at least one query");
+    auto tokens = checked_input<int64_t>(tokens_object, "tokens",
+                                         {batch_, sequence_length});
+    auto query_positions = checked_input<int64_t>(
+        query_positions_object, "query_positions", {batch_, query_count});
+
+    std::vector<std::vector<std::pair<int64_t, int64_t>>> ordered_queries(
+        static_cast<size_t>(batch_));
+    for (int64_t b = 0; b < batch_; ++b) {
+      auto &row = ordered_queries[static_cast<size_t>(b)];
+      row.reserve(static_cast<size_t>(query_count));
+      for (int64_t q = 0; q < query_count; ++q) {
+        const int64_t position =
+            query_positions.data()[b * query_count + q];
+        if (position < 0 || position >= sequence_length)
+          throw py::value_error("query_positions values must be in [0, N)");
+        row.emplace_back(position, q);
+      }
+      std::sort(row.begin(), row.end());
+      for (int64_t q = 1; q < query_count; ++q)
+        if (row[static_cast<size_t>(q - 1)].first ==
+            row[static_cast<size_t>(q)].first)
+          throw py::value_error(
+              "query_positions must be unique within each batch row");
+    }
+
+    const int64_t slots = suffix_k_ * occurrences_r_;
+    py::array_t<int64_t> source({batch_, query_count, slots}),
+        match_length({batch_, query_count, slots}),
+        state_id({batch_, query_count, slots}),
+        candidate_frequency({batch_, query_count, slots});
+    py::array_t<int32_t> count({batch_, query_count});
+    const int64_t output_size = batch_ * query_count * slots;
+    validate_runtime_positions();
+    ensure_pool(4);
+    std::unique_lock<std::mutex> call_lock;
+    {
+      py::gil_scoped_release release;
+      call_lock = std::unique_lock<std::mutex>(call_mutex_);
+      if (position_ != 0)
+        throw std::runtime_error("prefill requires an empty candidate state");
+      if (sequence_length > max_length_)
+        throw std::runtime_error("candidate state capacity exceeded");
+      std::fill(source.mutable_data(), source.mutable_data() + output_size,
+                int64_t{-1});
+      std::fill(match_length.mutable_data(),
+                match_length.mutable_data() + output_size, int64_t{0});
+      std::fill(state_id.mutable_data(), state_id.mutable_data() + output_size,
+                int64_t{-1});
+      std::fill(candidate_frequency.mutable_data(),
+                candidate_frequency.mutable_data() + output_size, int64_t{0});
+      std::fill(count.mutable_data(),
+                count.mutable_data() + batch_ * query_count, int32_t{0});
+      parallel_for_rows(4, [&](int64_t b) {
+        const auto &row = ordered_queries[static_cast<size_t>(b)];
+        int64_t next_query = 0;
+        for (int64_t position = 0; position < sequence_length; ++position) {
+          const bool emit = next_query < query_count &&
+                            row[static_cast<size_t>(next_query)].first == position;
+          const int64_t output_query =
+              emit ? row[static_cast<size_t>(next_query)].second : 0;
+          const int64_t output_at =
+              (b * query_count + output_query) * slots;
+          const int32_t row_count = step_row(
+              b, tokens.data()[b * sequence_length + position], position, emit,
+              source.mutable_data() + output_at,
+              match_length.mutable_data() + output_at,
+              state_id.mutable_data() + output_at,
+              candidate_frequency.mutable_data() + output_at);
+          if (emit) {
+            count.mutable_data()[b * query_count + output_query] = row_count;
+            ++next_query;
+          }
+        }
+      });
+    }
+    position_ = sequence_length;
+    std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_,
+              position_);
+    sync_owner_position();
+    call_lock.unlock();
+    return py::make_tuple(source, match_length, state_id, candidate_frequency,
+                          count);
+  }
+
   int64_t position() const { return position_; }
   py::array_t<int64_t> positions() const { return positions_; }
   size_t worker_count() const {
@@ -1741,6 +1847,7 @@ class NativeCandidateState {
     apply_tag(b, node, &position, 1, 1);
   }
   int32_t step_row(int64_t b, int64_t token, int64_t position,
+                   bool emit_output,
                    int64_t *source_out, int64_t *length_out,
                    int64_t *state_out, int64_t *frequency_out) {
     int32_t last = last_.data()[b], size = size_.data()[b],
@@ -1804,7 +1911,7 @@ class NativeCandidateState {
     }
     last = current;
     int32_t candidate_count = 0, states_with_history = 0, node = last;
-    while (node != -1 && states_with_history < suffix_k_) {
+    while (emit_output && node != -1 && states_with_history < suffix_k_) {
       const int64_t at = idx(b, state_capacity_, node);
       if (length_.data()[at] > 0) {
         materialize(b, node);
@@ -5074,6 +5181,7 @@ PYBIND11_MODULE(rosa_native_step, m) {
       .def("reset_masked", &NativeCandidateState::reset_masked)
       .def("prefill", &NativeCandidateState::prefill)
       .def("prefill_into", &NativeCandidateState::prefill_into)
+      .def("prefill_selected", &NativeCandidateState::prefill_selected)
       .def_property_readonly("position", &NativeCandidateState::position)
       .def_property_readonly("positions", &NativeCandidateState::positions)
       .def_property_readonly("worker_count",
diff --git a/native/tests/candidate_smoke.py b/native/tests/candidate_smoke.py
index cecd77c..dbce7d7 100644
--- a/native/tests/candidate_smoke.py
+++ b/native/tests/candidate_smoke.py
@@ -135,6 +135,28 @@ def main() -> None:
         for actual, expected in zip(prefill_arrays, prefill_allocating, strict=True)
     )
 
+    selected_queries = np.array([[2, 0], [1, 2]], dtype=np.int64)
+    selected_state = rosa_native_step.NativeCandidateState(
+        init_candidate_state(2, 4, suffix_k=2, occurrences_r=2)
+    )
+    selected = selected_state.prefill_selected(prefix, selected_queries)
+    batch = np.arange(2)[:, None]
+    assert all(
+        np.array_equal(actual, expected[batch, selected_queries])
+        for actual, expected in zip(selected, prefill_allocating, strict=True)
+    )
+    continuation_tokens = np.array([2, 5], dtype=np.int64)
+    selected_tail = selected_state.step(continuation_tokens)
+    full_state = rosa_native_step.NativeCandidateState(
+        init_candidate_state(2, 4, suffix_k=2, occurrences_r=2)
+    )
+    full_state.prefill(prefix)
+    full_tail = full_state.step(continuation_tokens)
+    assert all(
+        np.array_equal(actual, expected)
+        for actual, expected in zip(selected_tail, full_tail, strict=True)
+    )
+
     # Pools are lazy, never useful below the prefill threshold, and invalid
     # thread limits (including signed strings) select the serial fallback.
     small_pool = rosa_native_step.NativeCandidateState(init_candidate_state(3, 4))
diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py
index 3ff1695..9dec797 100644
--- a/src/rosa/__init__.py
+++ b/src/rosa/__init__.py
@@ -41,6 +41,7 @@
     "init_candidate_buffers",
     "init_inference_state",
     "prefill",
+    "prefill_candidates_selected",
     "reference_rosa",
 ]
 
@@ -513,6 +514,81 @@ def _build_forward_hard_candidates(
         return build_hard_candidates(tokens, suffix_k, occurrences_r)
 
 
+def _select_hard_candidates(
+    candidates: HardCandidates, query_positions: Tensor
+) -> HardCandidates:
+    """Gather a full candidate object into caller-specified query order."""
+
+    batch = torch.arange(
+        query_positions.shape[0], device=query_positions.device
+    ).unsqueeze(1)
+    return HardCandidates(
+        *(
+            value[batch, query_positions]
+            for value in (
+                getattr(candidates, name)
+                for name in HardCandidates.__dataclass_fields__
+            )
+        )
+    )
+
+
+def _build_stateful_hard_candidates_selected(
+    tokens: Tensor,
+    query_positions: Tensor,
+    suffix_k: int,
+    occurrences_r: int,
+) -> HardCandidates:
+    """Ingest a full context but materialize stateful candidates only for Q."""
+
+    from ._stateful_candidates_numba import init_candidate_state as initialize
+    from ._stateful_candidates_numba import prefill_candidates_selected
+
+    batch_size, sequence_length = tokens.shape
+    state = initialize(
+        batch_size,
+        sequence_length,
+        suffix_k=suffix_k,
+        occurrences_r=occurrences_r,
+    )
+    try:
+        candidates = prefill_candidates_selected(state, tokens, query_positions)
+        return HardCandidates(
+            *(getattr(candidates, name) for name in HardCandidates.__dataclass_fields__)
+        )
+    finally:
+        state.close()
+
+
+def _build_forward_hard_candidates_selected(
+    tokens: Tensor,
+    query_positions: Tensor,
+    suffix_k: int,
+    occurrences_r: int,
+    backend: CandidateBackend,
+) -> HardCandidates:
+    """Build Q-only candidates, allowing Python to use full-build plus gather."""
+
+    if backend == "python":
+        return _select_hard_candidates(
+            build_hard_candidates(tokens, suffix_k, occurrences_r), query_positions
+        )
+    try:
+        return _build_stateful_hard_candidates_selected(
+            tokens, query_positions, suffix_k, occurrences_r
+        )
+    except ModuleNotFoundError as error:
+        if error.name not in {"numba", "numpy"}:
+            raise
+        if backend == "stateful":
+            raise RuntimeError(
+                "stateful candidate backend requires the 'numba' extra"
+            ) from error
+        return _select_hard_candidates(
+            build_hard_candidates(tokens, suffix_k, occurrences_r), query_positions
+        )
+
+
 def _virtual_pool_single(i: int, pool_size: int) -> list[int]:
     """Causal bounded pool: half recent positions, half history anchors."""
 
@@ -1344,9 +1420,11 @@ def _forward_query_only(
         """Run candidate activations only at selected sequence positions."""
 
         (soft1, soft2), (st1, st2), hard_tokens = self.encode(z_a, code_logits)
+        selected_hard = hard_candidates is None
         hard = (
-            _build_forward_hard_candidates(
+            _build_forward_hard_candidates_selected(
                 hard_tokens,
+                query_positions,
                 self.suffix_k,
                 self.occurrences_r,
                 self.candidate_backend,
@@ -1357,10 +1435,24 @@ def _forward_query_only(
         bsz, n, _ = z_a.shape
         queries = query_positions.shape[1]
         z_query = _gather_sequence(z_a, query_positions)
-        exact_source = _gather_sequence(hard.source_index, query_positions)
-        exact_mask = _gather_sequence(hard.mask, query_positions)
-        exact_match_length = _gather_sequence(hard.match_length, query_positions)
-        exact_frequency = _gather_sequence(hard.frequency, query_positions)
+        exact_source = (
+            hard.source_index
+            if selected_hard
+            else _gather_sequence(hard.source_index, query_positions)
+        )
+        exact_mask = (
+            hard.mask if selected_hard else _gather_sequence(hard.mask, query_positions)
+        )
+        exact_match_length = (
+            hard.match_length
+            if selected_hard
+            else _gather_sequence(hard.match_length, query_positions)
+        )
+        exact_frequency = (
+            hard.frequency
+            if selected_hard
+            else _gather_sequence(hard.frequency, query_positions)
+        )
         exact_slots = exact_source.shape[-1]
 
         if self.virtual_candidates:
@@ -1534,7 +1626,11 @@ def _forward_query_only(
         candidate_number = torch.arange(source.shape[-1], device=z_a.device).view(
             1, 1, -1
         )
-        exact_rosa_slot = hard.rosa_slot.gather(1, query_positions).unsqueeze(-1)
+        exact_rosa_slot = (
+            hard.rosa_slot
+            if selected_hard
+            else hard.rosa_slot.gather(1, query_positions)
+        ).unsqueeze(-1)
         is_rosa = ((candidate_number == exact_rosa_slot) & (exact_rosa_slot >= 0)).to(
             z_a.dtype
         )
@@ -1704,19 +1800,25 @@ def _forward_query_only(
                 chosen_kind == VIRTUAL_KIND, query_positions, n, False
             ),
             hard_rosa_source_index=scatter(
-                hard.rosa_source_index.gather(1, query_positions),
+                hard.rosa_source_index
+                if selected_hard
+                else hard.rosa_source_index.gather(1, query_positions),
                 query_positions,
                 n,
                 -1,
             ),
             hard_rosa_predicted_tokens=scatter(
-                hard.rosa_predicted_tokens.gather(1, query_positions),
+                hard.rosa_predicted_tokens
+                if selected_hard
+                else hard.rosa_predicted_tokens.gather(1, query_positions),
                 query_positions,
                 n,
                 -1,
             ),
             hard_rosa_match_length=scatter(
-                hard.rosa_match_length.gather(1, query_positions),
+                hard.rosa_match_length
+                if selected_hard
+                else hard.rosa_match_length.gather(1, query_positions),
                 query_positions,
                 n,
                 0,
@@ -2378,6 +2480,16 @@ def forward_candidates_step_into(state: object, tokens: Tensor, buffers: object)
     return step(cast(Any, state), tokens, cast(Any, buffers))
 
 
+def prefill_candidates_selected(
+    state: object, tokens: Tensor, query_positions: Tensor
+) -> Any:
+    """Ingest N rich-candidate tokens and return outputs only for ``[B, Q]``."""
+
+    from ._stateful_candidates_numba import prefill_candidates_selected as prefill
+
+    return prefill(cast(Any, state), tokens, query_positions)
+
+
 def init_inference_state(
     batch_size: int,
     max_length: int = 8192,
diff --git a/src/rosa/_stateful_candidates_numba.py b/src/rosa/_stateful_candidates_numba.py
index 061bd07..d3ffc7c 100644
--- a/src/rosa/_stateful_candidates_numba.py
+++ b/src/rosa/_stateful_candidates_numba.py
@@ -417,6 +417,7 @@ def _step_row(
     last: int,
     size: int,
     edge_count: int,
+    emit_output: bool,
     output_source: np.ndarray,
     output_length: np.ndarray,
     output_state: np.ndarray,
@@ -615,40 +616,41 @@ def _step_row(
 
     last = current
     candidate_count = 0
-    states_with_history = 0
-    node = last
-    while node != -1 and states_with_history < suffix_k:
-        if length[node] > 0:
-            _materialize(
-                lct_left,
-                lct_right,
-                lct_parent,
-                occurrences,
-                occurrence_size,
-                frequency,
-                lazy_prefix,
-                lazy_size,
-                lazy_delta,
-                node,
-                lct_stack,
-            )
-            node_occurrences = int(occurrence_size[node])
-            if node_occurrences > 0:
-                states_with_history += 1
-                for occurrence_index in range(min(occurrences_r, node_occurrences)):
-                    source = int(occurrences[node, occurrence_index])
-                    duplicate = False
-                    for seen_index in range(candidate_count):
-                        if output_source[seen_index] == source:
-                            duplicate = True
-                            break
-                    if not duplicate:
-                        output_source[candidate_count] = source
-                        output_length[candidate_count] = length[node]
-                        output_state[candidate_count] = node
-                        output_frequency[candidate_count] = frequency[node]
-                        candidate_count += 1
-        node = int(suffix_link[node])
+    if emit_output:
+        states_with_history = 0
+        node = last
+        while node != -1 and states_with_history < suffix_k:
+            if length[node] > 0:
+                _materialize(
+                    lct_left,
+                    lct_right,
+                    lct_parent,
+                    occurrences,
+                    occurrence_size,
+                    frequency,
+                    lazy_prefix,
+                    lazy_size,
+                    lazy_delta,
+                    node,
+                    lct_stack,
+                )
+                node_occurrences = int(occurrence_size[node])
+                if node_occurrences > 0:
+                    states_with_history += 1
+                    for occurrence_index in range(min(occurrences_r, node_occurrences)):
+                        source = int(occurrences[node, occurrence_index])
+                        duplicate = False
+                        for seen_index in range(candidate_count):
+                            if output_source[seen_index] == source:
+                                duplicate = True
+                                break
+                        if not duplicate:
+                            output_source[candidate_count] = source
+                            output_length[candidate_count] = length[node]
+                            output_state[candidate_count] = node
+                            output_frequency[candidate_count] = frequency[node]
+                            candidate_count += 1
+            node = int(suffix_link[node])
 
     _path_write(
         lct_left,
@@ -735,6 +737,7 @@ def _step_batch_kernel(
             int(last[batch_index]),
             int(size[batch_index]),
             int(edge_count[batch_index]),
+            True,
             source[batch_index],
             match_length[batch_index],
             state_id[batch_index],
@@ -860,6 +863,7 @@ def _step_masked_batch_kernel(
             int(last[batch_index]),
             int(size[batch_index]),
             int(edge_count[batch_index]),
+            True,
             source[batch_index],
             match_length[batch_index],
             state_id[batch_index],
@@ -941,6 +945,7 @@ def _prefill_candidate_kernel(
                 int(last[batch_index]),
                 int(size[batch_index]),
                 int(edge_count[batch_index]),
+                True,
                 source[batch_index, position],
                 match_length[batch_index, position],
                 state_id[batch_index, position],
@@ -953,6 +958,98 @@ def _prefill_candidate_kernel(
     return source, match_length, state_id, candidate_frequency, count
 
 
+@njit(cache=True, nogil=True)
+def _prefill_candidate_selected_kernel(
+    tokens: np.ndarray,
+    query_positions: np.ndarray,
+    query_order: np.ndarray,
+    suffix_k: int,
+    occurrences_r: int,
+    history: np.ndarray,
+    head: np.ndarray,
+    edge_token: np.ndarray,
+    edge_target: np.ndarray,
+    edge_next: np.ndarray,
+    hash_state: np.ndarray,
+    hash_token: np.ndarray,
+    hash_edge: np.ndarray,
+    suffix_link: np.ndarray,
+    length: np.ndarray,
+    lct_left: np.ndarray,
+    lct_right: np.ndarray,
+    lct_parent: np.ndarray,
+    occurrences: np.ndarray,
+    occurrence_size: np.ndarray,
+    frequency: np.ndarray,
+    lazy_prefix: np.ndarray,
+    lazy_size: np.ndarray,
+    lazy_delta: np.ndarray,
+    lct_stack: np.ndarray,
+    last: np.ndarray,
+    size: np.ndarray,
+    edge_count: np.ndarray,
+) -> tuple[
+    np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray
+]:  # pragma: no cover
+    batch_size, sequence_length = tokens.shape
+    query_count = query_positions.shape[1]
+    slots = suffix_k * occurrences_r
+    source = np.full((batch_size, query_count, slots), -1, dtype=np.int64)
+    match_length = np.zeros((batch_size, query_count, slots), dtype=np.int64)
+    state_id = np.full((batch_size, query_count, slots), -1, dtype=np.int64)
+    candidate_frequency = np.zeros((batch_size, query_count, slots), dtype=np.int64)
+    count = np.zeros((batch_size, query_count), dtype=np.int32)
+    for batch_index in range(batch_size):
+        next_query = 0
+        for position in range(sequence_length):
+            emit = (
+                next_query < query_count
+                and query_positions[batch_index, next_query] == position
+            )
+            output_index = int(query_order[batch_index, next_query]) if emit else 0
+            row_count, row_last, row_size, row_edge_count = _step_row(
+                int(tokens[batch_index, position]),
+                position,
+                suffix_k,
+                occurrences_r,
+                history[batch_index],
+                head[batch_index],
+                edge_token[batch_index],
+                edge_target[batch_index],
+                edge_next[batch_index],
+                hash_state[batch_index],
+                hash_token[batch_index],
+                hash_edge[batch_index],
+                suffix_link[batch_index],
+                length[batch_index],
+                lct_left[batch_index],
+                lct_right[batch_index],
+                lct_parent[batch_index],
+                occurrences[batch_index],
+                occurrence_size[batch_index],
+                frequency[batch_index],
+                lazy_prefix[batch_index],
+                lazy_size[batch_index],
+                lazy_delta[batch_index],
+                lct_stack[batch_index],
+                int(last[batch_index]),
+                int(size[batch_index]),
+                int(edge_count[batch_index]),
+                emit,
+                source[batch_index, output_index],
+                match_length[batch_index, output_index],
+                state_id[batch_index, output_index],
+                candidate_frequency[batch_index, output_index],
+            )
+            if emit:
+                count[batch_index, output_index] = row_count
+                next_query += 1
+            last[batch_index] = row_last
+            size[batch_index] = row_size
+            edge_count[batch_index] = row_edge_count
+    return source, match_length, state_id, candidate_frequency, count
+
+
 @dataclass
 class CandidateState:
     """Fixed-capacity tensor state for exact online hard candidates."""
@@ -1226,6 +1323,29 @@ def _validate_candidate_tokens(
     return tokens, scalar
 
 
+def _validate_query_positions(
+    state: CandidateState,
+    tokens: Tensor,
+    query_positions: Tensor,
+) -> None:
+    if not isinstance(query_positions, Tensor):
+        raise TypeError("query_positions must be a Tensor")
+    if query_positions.dtype != torch.long:
+        raise TypeError("query_positions must have dtype torch.long")
+    if query_positions.ndim != 2 or query_positions.shape[0] != state.batch_size:
+        raise ValueError("query_positions must have shape [B, Q]")
+    if query_positions.shape[1] == 0:
+        raise ValueError("query_positions must contain at least one query")
+    if query_positions.device != tokens.device:
+        raise ValueError("query_positions must be on the same device as tokens")
+    sequence_length = tokens.shape[1]
+    if bool(((query_positions < 0) | (query_positions >= sequence_length)).any()):
+        raise ValueError("query_positions values must be in [0, N)")
+    ordered = query_positions.sort(dim=1).values
+    if ordered.shape[1] > 1 and bool((ordered[:, 1:] == ordered[:, :-1]).any()):
+        raise ValueError("query_positions must be unique within each batch row")
+
+
 def _candidate_step_from_arrays(
     state: CandidateState,
     arrays: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray],
@@ -1828,6 +1948,78 @@ def prefill_candidates(state: CandidateState, tokens: Tensor) -> CandidateStep:
     return _candidate_step_from_arrays(state, native_output, device)
 
 
+def prefill_candidates_selected(
+    state: CandidateState,
+    tokens: Tensor,
+    query_positions: Tensor,
+) -> CandidateStep:
+    """Consume N tokens while emitting exact candidates only at ``[B, Q]``.
+
+    Query order is preserved independently for every batch row. The state is
+    left at N exactly as after :func:`prefill_candidates`, so decoding may
+    continue with :func:`forward_candidates_step`.
+    """
+
+    tokens, _ = _validate_candidate_tokens(state, tokens, sequence=True)
+    _validate_query_positions(state, tokens, query_positions)
+    if state.ragged_mode:
+        raise RuntimeError("prefill is unavailable on a ragged candidate state")
+    if state.position != 0:
+        raise RuntimeError("prefill requires an empty candidate state")
+    sequence_length = tokens.shape[1]
+    if sequence_length > state.max_length:
+        raise RuntimeError("candidate state capacity exceeded")
+
+    device = tokens.device
+    cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous()
+    cpu_queries = query_positions.detach().to(device="cpu").contiguous()
+    token_array = cpu_tokens.numpy()
+    query_array = cpu_queries.numpy()
+
+    # Capability detection happens before either backend mutates the shared
+    # ABI-1 arrays. An older wheel can therefore fall back to selected Numba
+    # ingestion rather than replaying or gathering a full native prefill.
+    native_output = _native_candidate_call(
+        state, "prefill_selected", token_array, query_array
+    )
+    if native_output is None:
+        query_order = np.argsort(query_array, axis=1, kind="stable")
+        sorted_queries = np.take_along_axis(query_array, query_order, axis=1)
+        native_output = _prefill_candidate_selected_kernel(
+            token_array,
+            sorted_queries,
+            query_order,
+            state.suffix_k,
+            state.occurrences_r,
+            state.history,
+            state.head,
+            state.edge_token,
+            state.edge_target,
+            state.edge_next,
+            state.hash_state,
+            state.hash_token,
+            state.hash_edge,
+            state.suffix_link,
+            state.length,
+            state.lct_left,
+            state.lct_right,
+            state.lct_parent,
+            state.occurrences,
+            state.occurrence_size,
+            state.frequency,
+            state.lazy_prefix,
+            state.lazy_size,
+            state.lazy_delta,
+            state.lct_stack,
+            state.last,
+            state.size,
+            state.edge_count,
+        )
+        state.position = sequence_length
+    state.positions.fill(state.position)
+    return _candidate_step_from_arrays(state, native_output, device)
+
+
 # Explicit aliases keep naming discoverable while preserving the original API.
 prefill_candidate_state = prefill_candidates
 reset_candidate_rows = reset_candidates_masked
diff --git a/tests/test_rosa.py b/tests/test_rosa.py
index 259624e..4ecf3ce 100644
--- a/tests/test_rosa.py
+++ b/tests/test_rosa.py
@@ -1096,6 +1096,31 @@ def test_query_positions_validation_and_keyword_only_signature(self) -> None:
         with self.assertRaises(TypeError):
             model(z, None, None, torch.tensor([[1], [2]]))
 
+    def test_stateful_query_only_uses_selected_prefill_builder(self) -> None:
+        from rosa._stateful_candidates_numba import prefill_candidates_selected
+
+        model = self.make_model(
+            candidate_backend="stateful", learned_residual_scale=1.0
+        )
+        z = torch.randn(2, 9, 8, requires_grad=True)
+        positions = torch.tensor([[8, 0, 4], [1, 8, 0]])
+        with (
+            patch(
+                "rosa._stateful_candidates_numba.prefill_candidates",
+                side_effect=AssertionError("full stateful prefill"),
+            ),
+            patch(
+                "rosa._stateful_candidates_numba.prefill_candidates_selected",
+                wraps=prefill_candidates_selected,
+            ) as selected,
+        ):
+            output = model(z, query_positions=positions)
+            loss = output.updated.square().mean() + sum(output.aux_losses.values())
+            loss.backward()
+        selected.assert_called_once()
+        self.assertEqual(output.candidate_source_index.shape[:2], (2, 9))
+        self.assertIsNotNone(z.grad)
+
     def test_query_positions_matches_full_gradients_and_sentinels(self) -> None:
         devices = [torch.device("cpu")]
         if torch.cuda.is_available():
diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py
index 6947120..a984c06 100644
--- a/tests/test_stateful_candidates.py
+++ b/tests/test_stateful_candidates.py
@@ -23,6 +23,7 @@
     forward_candidates_step_masked,
     prefill_candidates,
     prefill_candidates_into,
+    prefill_candidates_selected,
     reset_candidates_masked,
 )
 from rosa._stateful_candidates_numba import (
@@ -43,6 +44,12 @@
 
 
 class TestStatefulCandidates(unittest.TestCase):
+    def assert_step_equal(self, actual: CandidateStep, expected: CandidateStep) -> None:
+        for name in _FIELDS:
+            self.assertTrue(
+                torch.equal(getattr(actual, name), getattr(expected, name)), name
+            )
+
     def test_close_is_idempotent_breaks_cycles_and_rejects_use(self) -> None:
         state = init_candidate_state_internal(1, 2, suffix_k=2, occurrences_r=2)
         buffers = init_candidate_buffers(state)
@@ -174,6 +181,120 @@ def test_prefill_emits_every_position_and_continues(self) -> None:
                 )
             )
 
+    def test_selected_prefill_is_q_only_ordered_and_continuable(self) -> None:
+        tokens_cpu = torch.tensor(
+            [[0, 1, 0, 2, 0, 1, 0, 3, 0], [3, 3, 4, 3, 5, 3, 4, 3, 3]]
+        )
+        query_cases = (
+            torch.tensor([[8], [0]]),
+            torch.tensor([[8, 0, 4], [1, 8, 0]]),
+            torch.tensor([list(reversed(range(9))), list(range(9))]),
+        )
+        devices = [torch.device("cpu")]
+        if torch.cuda.is_available():
+            devices.append(torch.device("cuda"))
+        backends = ["numba"]
+        try:
+            import rosa_native_step
+        except ModuleNotFoundError:
+            rosa_native_step = None
+        if rosa_native_step is not None and hasattr(
+            rosa_native_step.NativeCandidateState, "prefill_selected"
+        ):
+            backends.append("native")
+
+        for device, backend, queries_cpu in product(devices, backends, query_cases):
+            with self.subTest(device=device, backend=backend, queries=queries_cpu):
+                tokens = tokens_cpu.to(device)
+                queries = queries_cpu.to(device)
+                full_state = init_candidate_state_internal(
+                    2, 10, suffix_k=3, occurrences_r=2
+                )
+                selected_state = init_candidate_state_internal(
+                    2, 10, suffix_k=3, occurrences_r=2
+                )
+                if backend == "numba":
+                    full_state.native_state = False
+                    selected_state.native_state = False
+                else:
+                    assert rosa_native_step is not None
+                    full_state.native_state = rosa_native_step.NativeCandidateState(
+                        full_state
+                    )
+                    selected_state.native_state = rosa_native_step.NativeCandidateState(
+                        selected_state
+                    )
+                full = prefill_candidates(full_state, tokens)
+                selected = prefill_candidates_selected(selected_state, tokens, queries)
+                batch = torch.arange(2, device=device).unsqueeze(1)
+                expected = CandidateStep(
+                    *(getattr(full, field)[batch, queries] for field in _FIELDS)
+                )
+                self.assert_step_equal(selected, expected)
+                self.assertEqual(selected.source_index.shape[:2], queries.shape)
+                self.assertEqual(selected_state.position, 9)
+                self.assertEqual(selected_state.positions.tolist(), [9, 9])
+
+                continuation = torch.tensor([2, 4], device=device)
+                self.assert_step_equal(
+                    forward_candidates_step(selected_state, continuation),
+                    forward_candidates_step(full_state, continuation),
+                )
+
+    def test_selected_prefill_validation_and_old_wheel_fallback(self) -> None:
+        tokens = torch.tensor([[0, 1, 0], [2, 2, 3]])
+
+        class OldWheel:
+            def prefill(self, _tokens: np.ndarray) -> tuple[np.ndarray, ...]:
+                raise AssertionError("full prefill fallback must not be called")
+
+        old_state = init_candidate_state_internal(2, 3, suffix_k=2, occurrences_r=2)
+        old_state.native_state = OldWheel()
+        expected_state = init_candidate_state_internal(
+            2, 3, suffix_k=2, occurrences_r=2
+        )
+        expected_state.native_state = False
+        queries = torch.tensor([[2, 0], [1, 2]])
+        actual = prefill_candidates_selected(old_state, tokens, queries)
+        full = prefill_candidates(expected_state, tokens)
+        batch = torch.arange(2).unsqueeze(1)
+        self.assert_step_equal(
+            actual,
+            CandidateStep(*(getattr(full, field)[batch, queries] for field in _FIELDS)),
+        )
+        self.assertIs(old_state.native_state, False)
+
+        invalid = (
+            ([[0], [1]], TypeError, "must be a Tensor"),
+            (torch.zeros(2, 1), TypeError, "dtype torch.long"),
+            (torch.zeros(2, dtype=torch.long), ValueError, r"\[B, Q\]"),
+            (torch.zeros(1, 1, dtype=torch.long), ValueError, r"\[B, Q\]"),
+            (torch.empty(2, 0, dtype=torch.long), ValueError, "at least one"),
+            (torch.tensor([[0, 3], [1, 2]]), ValueError, r"\[0, N\)"),
+            (torch.tensor([[1, 1], [0, 2]]), ValueError, "unique"),
+        )
+        for positions, error_type, message in invalid:
+            state = init_candidate_state_internal(2, 3)
+            state.native_state = False
+            with (
+                self.subTest(message=message),
+                self.assertRaisesRegex(error_type, message),
+            ):
+                prefill_candidates_selected(state, tokens, positions)  # type: ignore[arg-type]
+            self.assertEqual(state.position, 0)
+            self.assertEqual(state.size.tolist(), [1, 1])
+
+        nonempty = init_candidate_state_internal(2, 4)
+        nonempty.native_state = False
+        forward_candidates_step(nonempty, tokens[:, 0])
+        with self.assertRaisesRegex(RuntimeError, "empty candidate state"):
+            prefill_candidates_selected(nonempty, tokens, queries)
+
+        if torch.cuda.is_available():
+            device_state = init_candidate_state_internal(2, 3)
+            with self.assertRaisesRegex(ValueError, "same device"):
+                prefill_candidates_selected(device_state, tokens.cuda(), queries)
+
     def test_allocating_native_dispatch_paths(self) -> None:
         def outputs(
             state: CandidateState, sequence_length: int | None = None

From 60f9f3f4bd9b0b0f9f4bda4c89926307b417b241 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 19:39:22 +0800
Subject: [PATCH 09/10] Skip inactive neural value projection

---
 src/rosa/__init__.py |  32 +++++--
 tests/test_rosa.py   | 200 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 225 insertions(+), 7 deletions(-)

diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py
index 9dec797..9d64371 100644
--- a/src/rosa/__init__.py
+++ b/src/rosa/__init__.py
@@ -1291,6 +1291,14 @@ def _candidate_symbolic_values(
     def _candidate_neural_values(self, z_a: Tensor, next_position: Tensor) -> Tensor:
         return _gather_sequence(self.value_proj(z_a), next_position)
 
+    def _attach_skipped_value_projection_gradient(self, values: Tensor) -> Tensor:
+        if not torch.is_grad_enabled():
+            return values
+        value_zero = values.new_zeros(())
+        for parameter in self.value_proj.parameters():
+            value_zero = value_zero + parameter.reshape(-1)[:0].sum()
+        return values + value_zero
+
     def _hybrid_soft_candidates(
         self,
         st1: Tensor,
@@ -1406,7 +1414,7 @@ def _attach_skipped_virtual_gradients(
         virtual_zero = retrieved.new_zeros(())
         for module in (self.virtual_query, self.virtual_key):
             for parameter in module.parameters():
-                virtual_zero = virtual_zero + parameter.sum() * 0
+                virtual_zero = virtual_zero + parameter.reshape(-1)[:0].sum()
         return retrieved + virtual_zero, updated + virtual_zero
 
     def _forward_query_only(
@@ -1702,9 +1710,14 @@ def _forward_query_only(
             self.value_gate_head(torch.cat([query_expanded, cand_key], dim=-1))
         ).squeeze(-1)
         value_gate = value_gate * non_null_mask.to(z_a.dtype)
-        neural_value = self.value_proj(_gather_sequence(z_a, next_position))
-        neural_value = neural_value * value_gate.unsqueeze(-1)
-        candidate_value = symbolic_value + self.neural_value_scale * neural_value
+        if self.neural_value_scale.item() == 0.0:
+            candidate_value = self._attach_skipped_value_projection_gradient(
+                symbolic_value
+            )
+        else:
+            neural_value = self.value_proj(_gather_sequence(z_a, next_position))
+            neural_value = neural_value * value_gate.unsqueeze(-1)
+            candidate_value = symbolic_value + self.neural_value_scale * neural_value
         if (dense_count or sparse_count) and not self.soft_candidates_forward:
             historical_end = exact_slots + self.virtual_candidates
             historical_scores = torch.cat(
@@ -2043,9 +2056,14 @@ def forward(
             self.value_gate_head(torch.cat([query_expanded, cand_key], dim=-1))
         ).squeeze(-1)
         value_gate = value_gate * non_null_mask.to(z_a.dtype)
-        neural_value = self._candidate_neural_values(z_a, next_position)
-        neural_value = neural_value * value_gate.unsqueeze(-1)
-        candidate_value = symbolic_value + self.neural_value_scale * neural_value
+        if self.neural_value_scale.item() == 0.0:
+            candidate_value = self._attach_skipped_value_projection_gradient(
+                symbolic_value
+            )
+        else:
+            neural_value = self._candidate_neural_values(z_a, next_position)
+            neural_value = neural_value * value_gate.unsqueeze(-1)
+            candidate_value = symbolic_value + self.neural_value_scale * neural_value
         hybrid_enabled = bool(
             self.dense_recent_candidates or self.sparse_old_candidates
         )
diff --git a/tests/test_rosa.py b/tests/test_rosa.py
index 4ecf3ce..1972eee 100644
--- a/tests/test_rosa.py
+++ b/tests/test_rosa.py
@@ -804,6 +804,206 @@ def virtual_hook(_module, _args, _output) -> None:
         self.assertEqual(output.candidate_source_index.shape[-1], 1 * 1 + 1 + 1)
         self.assertEqual(calls, 2)
 
+    def test_zero_neural_value_scale_skips_projection_and_reactivates(self) -> None:
+        torch.manual_seed(20260811)
+        model = self.make_model(neural_value_scale=0.0)
+        z = torch.randn(2, 9, 8)
+        positions = torch.tensor([[1, 6], [2, 8]])
+        projection_calls = 0
+
+        def projection_hook(_module, _args, _output) -> None:
+            nonlocal projection_calls
+            projection_calls += 1
+
+        hook = model.value_proj.register_forward_hook(projection_hook)
+        try:
+            for query_positions in (None, positions):
+                with self.subTest(scale=0.0, query_positions=query_positions):
+                    model.zero_grad(set_to_none=True)
+                    with patch.object(
+                        model,
+                        "_candidate_neural_values",
+                        side_effect=AssertionError("unexpected neural values"),
+                    ):
+                        output = model(z, query_positions=query_positions)
+                        output.updated.square().mean().backward()
+                        with torch.no_grad():
+                            model(z, query_positions=query_positions)
+                    for parameter in model.value_proj.parameters():
+                        self.assertIsNotNone(parameter.grad)
+                        assert parameter.grad is not None
+                        self.assertEqual(torch.count_nonzero(parameter.grad).item(), 0)
+            self.assertEqual(projection_calls, 0)
+
+            model.set_neural_value_scale(1.0)
+            for query_positions in (None, positions):
+                with self.subTest(scale=1.0, query_positions=query_positions):
+                    model.zero_grad(set_to_none=True)
+                    output = model(z, query_positions=query_positions)
+                    output.updated.square().mean().backward()
+                    for parameter in model.value_proj.parameters():
+                        self.assertIsNotNone(parameter.grad)
+                        assert parameter.grad is not None
+                        self.assertGreater(float(parameter.grad.abs().sum()), 0.0)
+            self.assertEqual(projection_calls, 2)
+        finally:
+            hook.remove()
+
+    def test_skipped_value_projection_zero_is_numerically_dormant(self) -> None:
+        model = self.make_model(neural_value_scale=0.0).to(dtype=torch.bfloat16)
+        symbolic = torch.randn(2, 3, 4, 8, dtype=torch.bfloat16)
+
+        for corruption in ("maximum", "non_finite"):
+            with self.subTest(corruption=corruption), torch.no_grad():
+                model.value_proj.weight.fill_(torch.finfo(torch.bfloat16).max)
+                if corruption == "non_finite":
+                    flat_weight = model.value_proj.weight.reshape(-1)
+                    flat_weight[0] = torch.nan
+                    flat_weight[1] = torch.inf
+                    flat_weight[2] = -torch.inf
+            model.zero_grad(set_to_none=True)
+            candidate_value = model._attach_skipped_value_projection_gradient(symbolic)
+            self.assertTrue(torch.equal(candidate_value, symbolic))
+            self.assertTrue(torch.isfinite(candidate_value).all())
+            candidate_value.float().sum().backward()
+            self.assertIsNotNone(model.value_proj.weight.grad)
+            assert model.value_proj.weight.grad is not None
+            self.assertEqual(
+                torch.count_nonzero(model.value_proj.weight.grad).item(), 0
+            )
+
+    def test_skipped_virtual_zero_is_numerically_dormant(self) -> None:
+        model = self.make_model(virtual_candidates=0).to(dtype=torch.bfloat16)
+        retrieved = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+        updated = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+
+        for corruption in ("maximum", "non_finite"):
+            with self.subTest(corruption=corruption), torch.no_grad():
+                for parameter in (
+                    model.virtual_query.weight,
+                    model.virtual_key.weight,
+                ):
+                    parameter.fill_(torch.finfo(torch.bfloat16).max)
+                    if corruption == "non_finite":
+                        flat_parameter = parameter.reshape(-1)
+                        flat_parameter[0] = torch.nan
+                        flat_parameter[1] = torch.inf
+                        flat_parameter[2] = -torch.inf
+            model.zero_grad(set_to_none=True)
+            actual_retrieved, actual_updated = model._attach_skipped_virtual_gradients(
+                retrieved, updated
+            )
+            self.assertTrue(torch.equal(actual_retrieved, retrieved))
+            self.assertTrue(torch.equal(actual_updated, updated))
+            self.assertTrue(torch.isfinite(actual_retrieved).all())
+            self.assertTrue(torch.isfinite(actual_updated).all())
+            (actual_retrieved.float().sum() + actual_updated.float().sum()).backward()
+            for parameter in (
+                model.virtual_query.weight,
+                model.virtual_key.weight,
+            ):
+                self.assertIsNotNone(parameter.grad)
+                assert parameter.grad is not None
+                self.assertEqual(torch.count_nonzero(parameter.grad).item(), 0)
+
+    def test_zero_neural_value_scale_matches_explicit_legacy_oracle(self) -> None:
+        torch.manual_seed(20260811)
+        base = self.make_model(
+            learned_residual_scale=1.0,
+            neural_value_scale=0.0,
+            value_gate_bias=0.0,
+        )
+        positions = torch.tensor([[1, 6], [2, 8]])
+
+        def legacy_oracle(
+            model: ROSA,
+            z: torch.Tensor,
+            query_positions: torch.Tensor | None,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]:
+            output = model(z, query_positions=query_positions)
+            non_null = output.candidate_mask & (output.candidate_kind != NULL_KIND)
+            next_position = torch.where(
+                non_null,
+                output.candidate_source_index + 1,
+                torch.zeros_like(output.candidate_source_index),
+            )
+            if query_positions is None:
+                symbol_sequence = (
+                    output.code_st[0] @ model.symbol_embedding_1.weight
+                    + output.code_st[1] @ model.symbol_embedding_2.weight
+                )
+                symbolic = _gather_sequence(symbol_sequence, next_position)
+                neural = model.value_proj(_gather_sequence(z, next_position))
+            else:
+                next_st1 = _gather_sequence(output.code_st[0], next_position)
+                next_st2 = _gather_sequence(output.code_st[1], next_position)
+                symbolic = (
+                    next_st1 @ model.symbol_embedding_1.weight
+                    + next_st2 @ model.symbol_embedding_2.weight
+                )
+                neural = _gather_sequence(model.value_proj(z), next_position)
+            symbolic = symbolic * non_null.unsqueeze(-1).to(z.dtype)
+            neural = neural * output.value_gate.unsqueeze(-1)
+            candidate = symbolic + model.neural_value_scale * neural
+            st_weights = output.hard_weights + (
+                output.soft_weights - output.soft_weights.detach()
+            )
+            retrieved = (st_weights.unsqueeze(-1) * candidate).sum(dim=-2)
+            updated = z + output.read_gate * model.out_proj(retrieved)
+            loss = (
+                updated.square().mean()
+                + retrieved.square().mean()
+                + output.value_gate.square().mean()
+            )
+            return updated, retrieved, loss, output
+
+        for query_positions in (None, positions):
+            with self.subTest(query_positions=query_positions):
+                optimized_model = copy.deepcopy(base)
+                oracle_model = copy.deepcopy(base)
+                z_optimized = torch.randn(2, 9, 8, requires_grad=True)
+                z_oracle = z_optimized.detach().clone().requires_grad_()
+
+                optimized = optimized_model(
+                    z_optimized, query_positions=query_positions
+                )
+                optimized_loss = (
+                    optimized.updated.square().mean()
+                    + optimized.retrieved.square().mean()
+                    + optimized.value_gate.square().mean()
+                )
+                oracle_updated, oracle_retrieved, oracle_loss, oracle = legacy_oracle(
+                    oracle_model, z_oracle, query_positions
+                )
+
+                self.assertTrue(torch.equal(optimized.updated, oracle_updated))
+                self.assertTrue(torch.equal(optimized.retrieved, oracle_retrieved))
+                self.assertTrue(torch.equal(optimized.value_gate, oracle.value_gate))
+                self.assertTrue(torch.equal(optimized_loss, oracle_loss))
+
+                optimized_loss.backward()
+                oracle_loss.backward()
+                torch.testing.assert_close(
+                    z_optimized.grad, z_oracle.grad, rtol=0, atol=0
+                )
+                for (optimized_name, optimized_parameter), (
+                    oracle_name,
+                    oracle_parameter,
+                ) in zip(
+                    optimized_model.named_parameters(),
+                    oracle_model.named_parameters(),
+                    strict=True,
+                ):
+                    self.assertEqual(optimized_name, oracle_name)
+                    self.assertIsNotNone(optimized_parameter.grad, optimized_name)
+                    self.assertIsNotNone(oracle_parameter.grad, oracle_name)
+                    torch.testing.assert_close(
+                        optimized_parameter.grad,
+                        oracle_parameter.grad,
+                        rtol=0,
+                        atol=0,
+                    )
+
     def test_prepared_candidates_are_bit_exact_and_shared_without_rebuild(self) -> None:
         torch.manual_seed(20260819)
         baseline_model = self.make_model(

From 63ab093e0e59723607ecde05234b5dfcf494e7f6 Mon Sep 17 00:00:00 2001
From: Lucas <30107107+aabbdev@users.noreply.github.com>
Date: Wed, 19 Aug 2026 21:27:45 +0800
Subject: [PATCH 10/10] Prepare 0.4.0 release

---
 CHANGELOG.md                      |  42 ++++++++++
 README.md                         |  80 ++++++++++++++----
 native/pyproject.toml             |   7 +-
 native/uv.lock                    |  24 ++++--
 pyproject.toml                    |   4 +-
 src/rosa/_numba_backend.py        |   4 +-
 tests/test_numba_backend.py       |   7 ++
 tests/test_rosa.py                | 134 +++++++++++++++++++++++++++++-
 tests/test_stateful_candidates.py |  49 +++++++++++
 tests/test_unified_inference.py   |   5 +-
 uv.lock                           |   2 +-
 11 files changed, 324 insertions(+), 34 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6a23b7a..41cf58f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,48 @@
 All notable changes to `rosa-torch` are documented here. The project follows
 semantic versioning while it remains in the 0.x development series.
 
+## 0.4.0 — 2026-08-19
+
+### Added
+
+- Query-position execution through `ROSA.forward(..., query_positions=...)`,
+  with full-shape public outputs and candidate computation restricted to Q.
+- Reusable `PreparedHardCandidates` snapshots with strict token, geometry,
+  backend, device, and mutation validation.
+- Exact native and Numba selected-prefill APIs that ingest N tokens, emit only
+  Q candidate rows, and preserve exact continuation.
+- Deterministic `close()` and context-manager support for persistent inference
+  states.
+
+### Changed
+
+- Removed native Python-owner reference cycles and made persistent state cleanup
+  deterministic.
+- Allowed `virtual_candidates=0` while preserving checkpoint compatibility and
+  expected zero gradients for inactive parameters.
+- Skipped inactive neural value projections without changing forward values or
+  training gradient coverage.
+- Preserved an exact one-hot straight-through forward while retaining the soft
+  backward path.
+
+### Performance
+
+- Reduced selected native candidate prefill from 15.65 ms to 5.72 ms on the
+  B16/N512/Q8 reference workload, while reducing candidate output storage from
+  about 4.66 MiB to 70 KiB.
+- Reduced the validated K4/R4/V1 query workload from 99.86 ms to 45.69 ms with
+  a task-specific K1/R1/V0 configuration.
+- Validated exact N32768/B1 query evaluation at 86.21 ms and 112.3 MiB peak CUDA
+  allocation on the reference system.
+
+### Compatibility
+
+- Existing calls without `query_positions` or `hard_candidates` keep the
+  historical execution path and output shapes.
+- Default candidate budgets and `backend="auto"` behavior are unchanged.
+- `rosa-torch-native 0.4.0` remains optional and requires
+  `rosa-torch>=0.4,<0.5` plus NumPy.
+
 ## 0.3.0 — 2026-08-18
 
 ### Added
diff --git a/README.md b/README.md
index 1fb6475..5018190 100644
--- a/README.md
+++ b/README.md
@@ -38,23 +38,24 @@ The design avoids a trainable dense automaton transition tensor and avoids dense
 - Optional shape-specialized `torch.compile` soft-match acceleration.
 - 100% statement and branch coverage for the `rosa` package.
 
-## What's new in 0.3.0
-
-Version 0.3.0 adds exact long-context RLBWT inference while preserving the
-unified training and inference API introduced in 0.2.0:
-
-- `backend="rlbwt"` provides a Python semantic oracle for exact online top-1
-  retrieval;
-- `backend="rlbwt_native"` fuses the same state machine in the optional C++
-  companion;
-- `backend="rlbwt_compact256"` adds compact exact storage for vocabularies up
-  to 256 IDs and very long configured contexts;
-- explicit `rlbwt_mc128` and `rlbwt_mc192` variants offer opt-in probabilistic
-  acceleration without changing exact `auto` dispatch;
-- lazy arenas and adaptive packed storage keep allocation tied to live context
-  length rather than maximum capacity.
-
-See the [changelog](https://github.com/aabbdev/rosa/blob/v0.3.0/CHANGELOG.md)
+## What's new in 0.4.0
+
+Version 0.4.0 improves exact differentiable retrieval and state lifecycle:
+
+- `ROSA.forward(..., query_positions=...)` restricts candidate scoring and
+  value retrieval to selected positions while preserving full-shape outputs;
+- exact stateful prefill emits candidates only at those positions, with native
+  and Numba implementations and exact continuation after the context;
+- `PreparedHardCandidates` allows compatible consumers to share one exact hard
+  candidate construction;
+- `virtual_candidates=0` removes the virtual branch without changing model
+  parameters or checkpoint compatibility;
+- persistent inference states provide deterministic `close()` and context
+  manager cleanup;
+- inactive neural value projections are skipped while retaining the expected
+  zero gradients during training.
+
+See the [changelog](https://github.com/aabbdev/rosa/blob/v0.4.0/CHANGELOG.md)
 for compatibility notes and the complete release summary.
 
 ## Core scoring rule
@@ -166,6 +167,14 @@ for token in generated_token_ids:  # each tensor has shape [2]
     predicted_token = forward_step(state, token)
 
 state.reset()
+state.close()
+```
+
+States can also be closed automatically:
+
+```python
+with init_inference_state(2, 32_768) as state:
+    predicted_token = forward_step(state, token_ids)
 ```
 
 Experimental top-1 RLBWT backends are also available: `rlbwt` is the Python
@@ -262,6 +271,43 @@ print(out.chosen_source_index.shape)  # [B, N]
 print(out.hard_rosa_match_length.shape)  # [B, N]
 ```
 
+### Query-only retrieval
+
+When losses or outputs are needed at a small set of positions, pass one unique
+position per row in a `[B, Q]` `torch.long` tensor:
+
+```python
+query_positions = torch.tensor([[31, 63], [15, 63]], device=z_a.device)
+out = model(z_a, z_b=z_b, query_positions=query_positions)
+```
+
+The encoder and exact automaton still consume all `N` positions. Candidate
+scoring, value retrieval, and their intermediate tensors use `Q`; public
+outputs retain `[B, N, ...]` shapes with sentinel values outside the selected
+positions. Mask invalid positions before reducing fields such as
+`candidate_scores`, whose sentinel is `-inf`.
+
+### Reusing exact hard candidates
+
+Compatible forwards can share one detached exact candidate snapshot:
+
+```python
+_, _, hard_tokens = model.encode(z_a)
+prepared = model.prepare_hard_candidates(hard_tokens)
+
+out_a = model(z_a, hard_candidates=prepared)
+out_b = model(z_a, hard_candidates=prepared)
+```
+
+Consumers must use identical hard tokens, `suffix_k`, `occurrences_r`, backend,
+shape, and device. Stale or mutated snapshots are rejected before use.
+
+For workloads whose retrieval quality has been validated with one suffix state
+and one occurrence, `suffix_k=1`, `occurrences_r=1`, and
+`virtual_candidates=0` provide the smallest exact top-1 candidate geometry.
+The default budgets remain unchanged because larger candidate sets provide
+additional training and ranking alternatives.
+
 ROSA uses the eager bounded differentiable `_soft_match` implementation by
 default. Set `compile_soft_match=True` to opt into a static `torch.compile`
 island, then warm every expected device, dtype, and shape bucket before serving:
diff --git a/native/pyproject.toml b/native/pyproject.toml
index 9a72540..b042948 100644
--- a/native/pyproject.toml
+++ b/native/pyproject.toml
@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rosa-torch-native"
-version = "0.3.0"
+version = "0.4.0"
 description = "Optional native CPU inference companion for rosa-torch"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -27,13 +27,16 @@ classifiers = [
 ]
 dependencies = [
     "numpy>=1.24",
-    "rosa-torch>=0.3,<0.4",
+    "rosa-torch>=0.4,<0.5",
 ]
 
 [project.urls]
 Repository = "https://github.com/aabbdev/rosa"
 Issues = "https://github.com/aabbdev/rosa/issues"
 
+[tool.uv.sources]
+rosa-torch = { path = ".." }
+
 [tool.setuptools]
 include-package-data = false
 
diff --git a/native/uv.lock b/native/uv.lock
index f497b62..f4a93c8 100644
--- a/native/uv.lock
+++ b/native/uv.lock
@@ -612,19 +612,29 @@ wheels = [
 
 [[package]]
 name = "rosa-torch"
-version = "0.3.0"
-source = { registry = "https://pypi.org/simple" }
+version = "0.4.0"
+source = { directory = "../" }
 dependencies = [
     { name = "torch" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/46/20/88b88b9b6307f72a2c3892405897efd370c7614abf49489569069ba7340d/rosa_torch-0.3.0.tar.gz", hash = "sha256:19009a36049c994bbc2950d1c9c9f113819f047dccc235a88ccd5d70eb35a677", size = 46505, upload-time = "2026-08-18T13:15:59.421Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/53/bc/5477145f95f309938cd271e48892ca3721a1f425551ed0600f9c005aef3d/rosa_torch-0.3.0-py3-none-any.whl", hash = "sha256:8c143a8cd681ec30ca7783b80254057f750cc371542d35f2504ecf123a3b14a5", size = 48752, upload-time = "2026-08-18T13:15:58.144Z" },
+
+[package.metadata]
+requires-dist = [
+    { name = "numba", marker = "extra == 'numba'", specifier = ">=0.66" },
+    { name = "torch" },
+]
+provides-extras = ["numba"]
+
+[package.metadata.requires-dev]
+dev = [
+    { name = "coverage", specifier = ">=7.10" },
+    { name = "pyright", extras = ["nodejs"], specifier = ">=1.1.400" },
+    { name = "ruff", specifier = ">=0.11" },
 ]
 
 [[package]]
 name = "rosa-torch-native"
-version = "0.3.0"
+version = "0.4.0"
 source = { editable = "." }
 dependencies = [
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
@@ -636,7 +646,7 @@ dependencies = [
 [package.metadata]
 requires-dist = [
     { name = "numpy", specifier = ">=1.24" },
-    { name = "rosa-torch", specifier = ">=0.3,<0.4" },
+    { name = "rosa-torch", directory = "../" },
 ]
 
 [[package]]
diff --git a/pyproject.toml b/pyproject.toml
index 7df6873..e3cc6f2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "rosa-torch"
-version = "0.3.0"
+version = "0.4.0"
 description = "Independent PyTorch implementation of RWKV-8 ROSA with exact suffix-automaton retrieval"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -32,7 +32,7 @@ numba = ["numba>=0.66"]
 [project.urls]
 Repository = "https://github.com/aabbdev/rosa"
 Issues = "https://github.com/aabbdev/rosa/issues"
-Changelog = "https://github.com/aabbdev/rosa/blob/v0.3.0/CHANGELOG.md"
+Changelog = "https://github.com/aabbdev/rosa/blob/v0.4.0/CHANGELOG.md"
 "Original ROSA description" = "https://www.rwkv.com/#rwkv-8-explained"
 
 [dependency-groups]
diff --git a/src/rosa/_numba_backend.py b/src/rosa/_numba_backend.py
index a8475ad..e29055e 100644
--- a/src/rosa/_numba_backend.py
+++ b/src/rosa/_numba_backend.py
@@ -130,8 +130,8 @@ def _predict_row(  # pragma: no cover - executed as compiled Numba code
                         edge_next,
                         edge_count,
                         clone,
-                        edge_token[edge],
-                        edge_target[edge],
+                        int(edge_token[edge]),
+                        int(edge_target[edge]),
                     )
                     edge = edge_next[edge]
                 while (
diff --git a/tests/test_numba_backend.py b/tests/test_numba_backend.py
index 743bbf0..270dcc1 100644
--- a/tests/test_numba_backend.py
+++ b/tests/test_numba_backend.py
@@ -93,6 +93,13 @@ def test_stateful_private_validation_and_full_replay(self) -> None:
         expected, _, _ = reference_rosa(tokens)
         self.assertTrue(torch.equal(predict_exact_stateful(tokens), expected))
 
+    def test_private_state_close_is_idempotent_and_rejects_use(self) -> None:
+        state = _init_inference_state(1, 1)
+        state.close()
+        state.close()
+        with self.assertRaisesRegex(RuntimeError, "state is closed"):
+            _forward_step(state, torch.tensor([0]))
+
     @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable")
     def test_cuda_round_trip_matches_reference(self) -> None:
         tokens = torch.tensor(
diff --git a/tests/test_rosa.py b/tests/test_rosa.py
index 1972eee..ee437d0 100644
--- a/tests/test_rosa.py
+++ b/tests/test_rosa.py
@@ -7,7 +7,7 @@
 import unittest
 import weakref
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import FrozenInstanceError
+from dataclasses import FrozenInstanceError, replace
 from unittest.mock import patch
 
 import numpy as np
@@ -22,6 +22,7 @@
     PreparedHardCandidates,
     _balance_kl,
     _build_forward_hard_candidates,
+    _build_forward_hard_candidates_selected,
     _clear_soft_match_compile_cache,
     _gather_sequence,
     _soft_match,
@@ -1240,6 +1241,109 @@ def forged(field, replacement):
                     hard_candidates=prepare(),
                 )
 
+    def test_prepared_candidates_cover_all_metadata_guards(self) -> None:
+        model = self.make_model(candidate_backend="python")
+        tokens = torch.zeros((1, 3), dtype=torch.long)
+
+        for invalid, error_type, message in (
+            (object(), TypeError, "must be a Tensor"),
+            (tokens.float(), TypeError, "dtype torch.long"),
+            (tokens[0], ValueError, r"\[B, N\]"),
+            (torch.empty((1, 0), dtype=torch.long), ValueError, "dimensions"),
+            (
+                torch.empty((1, 3), dtype=torch.long, device="meta"),
+                ValueError,
+                "same device",
+            ),
+        ):
+            with (
+                self.subTest(message=message),
+                self.assertRaisesRegex(error_type, message),
+            ):
+                model.prepare_hard_candidates(invalid)  # type: ignore[arg-type]
+
+        def prepare() -> PreparedHardCandidates:
+            return model.prepare_hard_candidates(tokens)
+
+        with self.assertRaisesRegex(ValueError, "device does not match"):
+            model._validate_prepared_hard_candidates(
+                replace(prepare(), device=torch.device("meta")), tokens
+            )
+
+        non_tensor_field = prepare()
+        non_tensor_field.candidates.source_index = object()  # type: ignore[assignment]
+        with self.assertRaisesRegex(TypeError, "fields must all be Tensors"):
+            model._validate_prepared_hard_candidates(non_tensor_field, tokens)
+
+        with self.assertRaisesRegex(ValueError, "metadata is invalid"):
+            model._validate_prepared_hard_candidates(
+                replace(prepare(), tensor_versions=()), tokens
+            )
+        with self.assertRaisesRegex(ValueError, "hard_tokens metadata"):
+            model._validate_prepared_hard_candidates(
+                replace(prepare(), hard_tokens=tokens.float()), tokens
+            )
+
+        candidate_type = prepare()
+        original_tensors = model._prepared_tensors(candidate_type)
+        candidate_type.candidates.source_index = object()  # type: ignore[assignment]
+        with (
+            patch.object(model, "_prepared_tensors", return_value=original_tensors),
+            self.assertRaisesRegex(TypeError, "source_index.*Tensor"),
+        ):
+            model._validate_prepared_hard_candidates(candidate_type, tokens)
+
+        snapshots = prepare()
+        with self.assertRaisesRegex(TypeError, "snapshots must be Tensors"):
+            model._validate_prepared_hard_candidates(
+                replace(
+                    snapshots,
+                    _tensor_snapshots=(object(), *snapshots._tensor_snapshots[1:]),
+                ),
+                tokens,
+            )
+        with self.assertRaisesRegex(ValueError, "tensor metadata is invalid"):
+            model._validate_prepared_hard_candidates(
+                replace(
+                    snapshots,
+                    _tensor_snapshots=(
+                        torch.empty(0, dtype=torch.long),
+                        *snapshots._tensor_snapshots[1:],
+                    ),
+                ),
+                tokens,
+            )
+
+    def test_selected_builder_optional_dependency_fallbacks(self) -> None:
+        tokens = torch.tensor([[0, 1, 0]])
+        positions = torch.tensor([[2, 0]])
+
+        with (
+            patch(
+                "rosa._build_stateful_hard_candidates_selected",
+                side_effect=ModuleNotFoundError("unexpected", name="unexpected"),
+            ),
+            self.assertRaises(ModuleNotFoundError),
+        ):
+            _build_forward_hard_candidates_selected(tokens, positions, 2, 1, "auto")
+
+        missing_numba = ModuleNotFoundError("missing numba", name="numba")
+        with patch(
+            "rosa._build_stateful_hard_candidates_selected", side_effect=missing_numba
+        ):
+            with self.assertRaisesRegex(RuntimeError, "numba.*extra"):
+                _build_forward_hard_candidates_selected(
+                    tokens, positions, 2, 1, "stateful"
+                )
+            actual = _build_forward_hard_candidates_selected(
+                tokens, positions, 2, 1, "auto"
+            )
+        expected = _build_forward_hard_candidates_selected(
+            tokens, positions, 2, 1, "python"
+        )
+        for name in expected.__dataclass_fields__:
+            self.assertTrue(torch.equal(getattr(actual, name), getattr(expected, name)))
+
     def test_prepared_candidates_have_explicit_external_ownership(self) -> None:
         model = self.make_model(candidate_backend="python")
         z = torch.randn(1, 8, 8)
@@ -1295,6 +1399,34 @@ def test_query_positions_validation_and_keyword_only_signature(self) -> None:
                 model(z, query_positions=positions)
         with self.assertRaises(TypeError):
             model(z, None, None, torch.tensor([[1], [2]]))
+        with self.assertRaisesRegex(ValueError, "same device"):
+            model(
+                z,
+                query_positions=torch.empty((2, 1), dtype=torch.long, device="meta"),
+            )
+
+    def test_private_zero_virtual_and_query_only_branch_edges(self) -> None:
+        no_virtual = self.make_model(virtual_candidates=0)
+        z = torch.randn(1, 4, 8)
+        source, mask, score = no_virtual._virtual_candidates(
+            z,
+            torch.empty((1, 4, 0), dtype=torch.long),
+            torch.empty((1, 4, 0), dtype=torch.bool),
+        )
+        self.assertEqual(source.shape, (1, 4, 0))
+        self.assertEqual(mask.shape, source.shape)
+        self.assertEqual(score.shape, source.shape)
+
+        one_anchor = self.make_model(
+            virtual_candidates=0,
+            dense_recent_candidates=1,
+            sparse_old_candidates=1,
+            sparse_old_pool_size=1,
+            soft_candidates_forward=True,
+        )
+        positions = torch.tensor([[3, 1, 0, 2]])
+        output = one_anchor(z, query_positions=positions)
+        self.assertEqual(output.updated.shape, z.shape)
 
     def test_stateful_query_only_uses_selected_prefill_builder(self) -> None:
         from rosa._stateful_candidates_numba import prefill_candidates_selected
diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py
index a984c06..7361ba2 100644
--- a/tests/test_stateful_candidates.py
+++ b/tests/test_stateful_candidates.py
@@ -284,6 +284,22 @@ def prefill(self, _tokens: np.ndarray) -> tuple[np.ndarray, ...]:
             self.assertEqual(state.position, 0)
             self.assertEqual(state.size.tolist(), [1, 1])
 
+        device_state = init_candidate_state_internal(2, 3)
+        with self.assertRaisesRegex(ValueError, "same device"):
+            prefill_candidates_selected(
+                device_state,
+                tokens,
+                torch.empty((2, 1), dtype=torch.long, device="meta"),
+            )
+
+        ragged = init_candidate_state_internal(2, 3, ragged=True)
+        with self.assertRaisesRegex(RuntimeError, "ragged"):
+            prefill_candidates_selected(ragged, tokens, queries)
+
+        too_short = init_candidate_state_internal(2, 2)
+        with self.assertRaisesRegex(RuntimeError, "capacity"):
+            prefill_candidates_selected(too_short, tokens, queries)
+
         nonempty = init_candidate_state_internal(2, 4)
         nonempty.native_state = False
         forward_candidates_step(nonempty, tokens[:, 0])
@@ -374,6 +390,27 @@ def native_prefill(
         self.assertFalse(bool(result.mask.any()))
         self.assertEqual(prefilled.positions.tolist(), [2, 2])
 
+        selected = init_candidate_state_internal(2, 2, suffix_k=2, occurrences_r=2)
+
+        class SelectedNative:
+            def __init__(self, state: CandidateState) -> None:
+                self.state = state
+
+            def prefill_selected(
+                self, tokens: np.ndarray, queries: np.ndarray
+            ) -> tuple[np.ndarray, ...]:
+                self.state.position = tokens.shape[1]
+                return outputs(self.state, queries.shape[1])
+
+        selected.native_state = SelectedNative(selected)
+        selected_result = prefill_candidates_selected(
+            selected,
+            torch.tensor([[1, 2], [3, 4]]),
+            torch.tensor([[1], [0]]),
+        )
+        self.assertFalse(bool(selected_result.mask.any()))
+        self.assertEqual(selected.positions.tolist(), [2, 2])
+
     def test_caller_owned_step_and_prefill_buffers_are_exact(self) -> None:
         tokens = torch.tensor([[0, 1, 0, 2], [3, 3, 4, 3]])
         state = init_candidate_state(2, 4, suffix_k=3, occurrences_r=2)
@@ -795,6 +832,18 @@ def test_public_wrapper_reports_missing_numba(self) -> None:
             with self.assertRaisesRegex(RuntimeError, "numba"):
                 forward_candidates_step(object(), torch.tensor([1]))
 
+    def test_public_selected_prefill_wrapper_delegates(self) -> None:
+        sentinel = object()
+        with patch(
+            "rosa._stateful_candidates_numba.prefill_candidates_selected",
+            return_value=sentinel,
+        ) as delegated:
+            actual = __import__("rosa").prefill_candidates_selected(
+                object(), torch.tensor([[1]]), torch.tensor([[0]])
+            )
+        self.assertIs(actual, sentinel)
+        delegated.assert_called_once()
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_unified_inference.py b/tests/test_unified_inference.py
index fef5ec2..22daafb 100644
--- a/tests/test_unified_inference.py
+++ b/tests/test_unified_inference.py
@@ -39,8 +39,9 @@ def test_close_context_manager_and_reset_lifetime(self) -> None:
             lambda: managed.step_into(torch.tensor([0]), buffers),
             lambda: managed.prefill(tokens),
         ):
-            with self.subTest(access=access), self.assertRaisesRegex(
-                RuntimeError, "^state is closed$"
+            with (
+                self.subTest(access=access),
+                self.assertRaisesRegex(RuntimeError, "^state is closed$"),
             ):
                 access()
         managed.close()
diff --git a/uv.lock b/uv.lock
index 2cd12eb..fc26004 100644
--- a/uv.lock
+++ b/uv.lock
@@ -774,7 +774,7 @@ nodejs = [
 
 [[package]]
 name = "rosa-torch"
-version = "0.3.0"
+version = "0.4.0"
 source = { editable = "." }
 dependencies = [
     { name = "torch" },