Release rosa-torch 0.3.0

.github/workflows/ci.yml

@@ -24,7 +24,7 @@ jobs:
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4
 
       - name: Install uv and Python
         uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
@@ -62,7 +62,7 @@ jobs:
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4
 
       - name: Install uv and Python
         uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
@@ -96,7 +96,7 @@ jobs:
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@3d3c42e5aac5ba805825a76410c181273ba90b1 # v7.0.1
+        uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4
 
       - name: Install uv and Python
         uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0

.github/workflows/publish.yml

@@ -10,12 +10,13 @@ permissions:
 jobs:
   build:
     name: Build distributions
+    if: ${{ startsWith(github.event.release.tag_name, 'v') && !startsWith(github.event.release.tag_name, 'native-v') }}
     runs-on: ubuntu-latest
     timeout-minutes: 15
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4
 
       - name: Install uv and Python
         uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
@@ -28,7 +29,7 @@ jobs:
         run: uv build
 
       - name: Upload distributions
-        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
         with:
           name: python-distributions
           path: dist/
@@ -47,7 +48,7 @@ jobs:
 
     steps:
       - name: Download distributions
-        uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
+        uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4
         with:
           name: python-distributions
           path: dist/

CHANGELOG.md

@@ -3,6 +3,34 @@
 All notable changes to `rosa-torch` are documented here. The project follows
 semantic versioning while it remains in the 0.x development series.
 
+## 0.3.0 — 2026-08-18
+
+### Added
+
+- Exact online RLBWT inference backends for top-1 dense workloads, including a
+  Python semantic oracle and optional fused native implementations.
+- Compact exact `rlbwt_compact256` storage for long contexts with adaptive
+  leaves, packed position arrays, and bounded owned memory.
+- Explicit opt-in Monte-Carlo RLBWT variants with independently seeded suffix
+  fingerprints and clearly separated backend names.
+- Native RLBWT smoke tests covering prefill, continuation, reset, compact token
+  validation, and agreement with the exact Python oracle.
+
+### Changed
+
+- Reworked native RLBWT storage around unified cache-sized leaves and adaptive
+  representations for BWT, position, and longest-common-suffix data.
+- Allocated long-context history and tree arenas lazily from live length rather
+  than configured capacity.
+- Kept `backend="auto"` on the production suffix-automaton path; all RLBWT
+  backends remain explicit opt-in choices.
+
+### Compatibility
+
+- The public `rosa` imports and existing stateful inference APIs are unchanged.
+- The distribution remains `rosa-torch` and supports Python 3.10+.
+- The optional native companion remains separately versioned and optional.
+
 ## 0.2.0 — 2026-08-11
 
 ### Added

README.md

@@ -38,23 +38,23 @@ 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.2.0
-
-Version 0.2.0 turns the original differentiable prototype into a unified
-training and inference package:
-
-- exact stateful inference now scales with amortized `O(log N)` suffix-path
-  updates instead of eager linear propagation;
-- one facade covers top-1 and rich candidates, dense and ragged batches,
-  prefill, continuation, reset, and row recycling;
-- the optional native companion accelerates rich/top-1 prefill, parallel batch
-  work, and caller-owned `step_into` buffers while retaining exact fallbacks;
-- `ROSA.forward` uses fused rich candidate prefill and preserves the independent
-  Python oracle;
-- projections are performed before candidate gather, and an opt-in compiled
-  soft-match island accelerates warmed fixed-shape training workloads.
-
-See the [changelog](https://github.com/aabbdev/rosa/blob/v0.2.0/CHANGELOG.md)
+## 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)
 for compatibility notes and the complete release summary.
 
 ## Core scoring rule

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "rosa-torch"
-version = "0.2.0"
+version = "0.3.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.2.0/CHANGELOG.md"
+Changelog = "https://github.com/aabbdev/rosa/blob/v0.3.0/CHANGELOG.md"
 "Original ROSA description" = "https://www.rwkv.com/#rwkv-8-explained"
 
 [dependency-groups]

tests/test_rlbwt_backend.py

@@ -198,10 +198,20 @@ def unexpected_missing(name: str, *args: Any, **kwargs: Any) -> Any:
 
     def test_native_dispatch_with_capability_stub(self) -> None:
         class StubNativeRLBWTState:
-            def __init__(self, batch_size: int, max_length: int) -> None:
+            def __init__(
+                self,
+                batch_size: int,
+                max_length: int,
+                vocabulary_size: int | None = None,
+                lanes: int | None = None,
+                seed: int | None = None,
+            ) -> None:
                 self.batch_size = batch_size
                 self.max_length = max_length
                 self.position = 0
+                self.vocabulary_size = vocabulary_size
+                self.lanes = lanes
+                self.seed = seed
 
             def step(self, tokens: np.ndarray) -> np.ndarray:
                 self.position += 1
@@ -211,8 +221,19 @@ def prefill(self, tokens: np.ndarray) -> np.ndarray:
                 self.position = tokens.shape[1]
                 return tokens.copy()
 
+        class StubNativeRLBWTStateMC(StubNativeRLBWTState):
+            def __init__(
+                self, batch_size: int, max_length: int, lanes: int, seed: int
+            ) -> None:
+                super().__init__(batch_size, max_length, lanes=lanes, seed=seed)
+
         capability = SimpleNamespace(
-            rlbwt_abi_version=1, NativeRLBWTState=StubNativeRLBWTState
+            rlbwt_abi_version=1,
+            NativeRLBWTState=StubNativeRLBWTState,
+            rlbwt_compact_abi_version=1,
+            NativeRLBWTCompactState=StubNativeRLBWTState,
+            rlbwt_mc_abi_version=1,
+            NativeRLBWTStateMC=StubNativeRLBWTStateMC,
         )
         with patch.dict("sys.modules", {"rosa_native_step": capability}):
             state = init_inference_state(2, 3, backend="rlbwt_native")
@@ -230,6 +251,23 @@ def prefill(self, tokens: np.ndarray) -> np.ndarray:
                 tuple(prefill(empty, torch.empty(0, dtype=torch.long)).shape), (0,)
             )
 
+            compact = init_inference_state(1, 3, backend="rlbwt_compact256")
+            self.assertEqual(compact._impl.vocabulary_size, 256)
+            compact_tokens = torch.tensor([[0, 255]])
+            self.assertTrue(
+                torch.equal(prefill(compact, compact_tokens), compact_tokens)
+            )
+            compact.reset()
+            self.assertEqual(forward_step(compact, torch.tensor([7])).item(), 7)
+            with self.assertRaisesRegex(ValueError, r"\[0, 255\]"):
+                forward_step(compact, torch.tensor([256]))
+
+            for backend, lanes in (("rlbwt_mc128", 2), ("rlbwt_mc192", 3)):
+                mc = init_inference_state(1, 2, backend=backend)
+                self.assertEqual(mc._impl.lanes, lanes)
+                self.assertEqual(mc._impl.seed, 20260811)
+                self.assertEqual(prefill(mc, torch.tensor([[5]])).item(), 5)
+
     def test_native_backend_matches_python_oracle_when_available(self) -> None:
         try:
             import rosa_native_step

tests/test_rosa.py

@@ -283,6 +283,14 @@ def test_soft_match_disabled_and_ready_failure_fall_back(self) -> None:
 
         signature = _soft_match_signature(st1, st2, source, mask, 3)
 
+        def succeed(*args: torch.Tensor) -> torch.Tensor:
+            return _soft_match_torch(*args, window=3)
+
+        rosa._SOFT_MATCH_COMPILED[3] = succeed
+        rosa._SOFT_MATCH_COMPILE_READY.add(signature)
+        cached = _soft_match(st1, st2, source, mask, window=3)
+        self.assertTrue(torch.equal(cached, expected))
+
         def fail(*args: torch.Tensor) -> torch.Tensor:
             raise RuntimeError("cached specialization failed")
 
@@ -335,6 +343,58 @@ def invoke() -> torch.Tensor:
         signature = _soft_match_signature(st1, st2, source, mask, 3)
         self.assertIn(signature, rosa._SOFT_MATCH_COMPILE_READY)
 
+    def test_soft_match_waiter_reuses_newly_ready_specialization(self) -> None:
+        _clear_soft_match_compile_cache()
+        st1 = torch.softmax(torch.randn(1, 5, 3), dim=-1)
+        st2 = torch.softmax(torch.randn(1, 5, 2), dim=-1)
+        source = torch.randint(-1, 5, (1, 5, 4))
+        mask = source >= 0
+        signature = _soft_match_signature(st1, st2, source, mask, 3)
+        compiled_started = threading.Event()
+        second_waiting = threading.Event()
+        release_compiled = threading.Event()
+
+        class SignalingLock:
+            def __init__(self) -> None:
+                self._lock = threading.Lock()
+                self._count_lock = threading.Lock()
+                self._entries = 0
+
+            def __enter__(self) -> None:
+                with self._count_lock:
+                    self._entries += 1
+                    if self._entries == 2:
+                        second_waiting.set()
+                self._lock.acquire()
+
+            def __exit__(self, *args: object) -> None:
+                self._lock.release()
+
+        rosa._SOFT_MATCH_SIGNATURE_LOCKS[signature] = SignalingLock()
+
+        def compiled(*args: torch.Tensor) -> torch.Tensor:
+            compiled_started.set()
+            if not release_compiled.wait(timeout=5):
+                raise RuntimeError("timed out waiting for concurrent caller")
+            return _soft_match_torch(*args, window=3)
+
+        with (
+            patch("rosa.torch.compile", return_value=compiled) as compile_mock,
+            ThreadPoolExecutor(max_workers=2) as executor,
+        ):
+            first = executor.submit(_soft_match, st1, st2, source, mask, 3)
+            self.assertTrue(compiled_started.wait(timeout=5))
+            second = executor.submit(_soft_match, st1, st2, source, mask, 3)
+            try:
+                self.assertTrue(second_waiting.wait(timeout=5))
+            finally:
+                release_compiled.set()
+            first_result = first.result()
+            second_result = second.result()
+
+        self.assertEqual(compile_mock.call_count, 1)
+        self.assertTrue(torch.equal(first_result, second_result))
+
     def test_soft_match_failure_does_not_poison_other_signature(self) -> None:
         _clear_soft_match_compile_cache()
         calls: list[tuple[int, ...]] = []

tests/test_stateful_candidates.py

@@ -143,6 +143,85 @@ def test_prefill_emits_every_position_and_continues(self) -> None:
                 )
             )
 
+    def test_allocating_native_dispatch_paths(self) -> None:
+        def outputs(
+            state: CandidateState, sequence_length: int | None = None
+        ) -> tuple[np.ndarray, ...]:
+            slots = state.suffix_k * state.occurrences_r
+            prefix = (
+                (state.batch_size,)
+                if sequence_length is None
+                else (state.batch_size, sequence_length)
+            )
+            shape = (*prefix, slots)
+            return (
+                np.full(shape, -1, dtype=np.int64),
+                np.zeros(shape, dtype=np.int64),
+                np.full(shape, -1, dtype=np.int64),
+                np.zeros(shape, dtype=np.int64),
+                np.zeros(prefix, dtype=np.int32),
+            )
+
+        uniform = init_candidate_state(2, 2, suffix_k=2, occurrences_r=2)
+
+        def native_step(
+            state: CandidateState, tokens: torch.Tensor
+        ) -> tuple[np.ndarray, ...]:
+            state.position += 1
+            return outputs(state)
+
+        with patch(
+            "rosa._stateful_candidates_numba._native_candidate_step",
+            side_effect=native_step,
+        ):
+            result = forward_candidates_step(uniform, torch.tensor([1, 2]))
+        self.assertFalse(bool(result.mask.any()))
+        self.assertEqual(uniform.positions.tolist(), [1, 1])
+
+        ragged = init_candidate_state_internal(
+            2, 2, suffix_k=2, occurrences_r=2, ragged=True
+        )
+
+        def native_masked(
+            state: CandidateState, method: str, *arrays: np.ndarray
+        ) -> tuple[np.ndarray, ...] | None:
+            self.assertEqual(method, "step_masked")
+            active = arrays[1].astype(bool)
+            reset = arrays[2].astype(bool)
+            state.positions[np.logical_and(active, reset)] = 0
+            state.positions[active] += 1
+            return outputs(state)
+
+        with patch(
+            "rosa._stateful_candidates_numba._native_candidate_call",
+            side_effect=native_masked,
+        ):
+            result = forward_candidates_step_masked(
+                ragged,
+                torch.tensor([1, 2]),
+                torch.tensor([True, False]),
+                torch.tensor([True, False]),
+            )
+        self.assertFalse(bool(result.mask.any()))
+        self.assertEqual(ragged.positions.tolist(), [1, 0])
+
+        prefilled = init_candidate_state(2, 2, suffix_k=2, occurrences_r=2)
+
+        def native_prefill(
+            state: CandidateState, method: str, tokens: np.ndarray
+        ) -> tuple[np.ndarray, ...]:
+            self.assertEqual(method, "prefill")
+            state.position = tokens.shape[1]
+            return outputs(state, tokens.shape[1])
+
+        with patch(
+            "rosa._stateful_candidates_numba._native_candidate_call",
+            side_effect=native_prefill,
+        ):
+            result = prefill_candidates(prefilled, torch.tensor([[1, 2], [3, 4]]))
+        self.assertFalse(bool(result.mask.any()))
+        self.assertEqual(prefilled.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)