From ab4476de98d92f3f4133efc83c5faccbf196af54 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:57:25 +0800 Subject: [PATCH 1/3] Prepare rosa-torch 0.3.0 release --- .github/workflows/ci.yml | 6 +-- .github/workflows/publish.yml | 7 +-- CHANGELOG.md | 28 +++++++++++ README.md | 34 ++++++------- pyproject.toml | 4 +- tests/test_rlbwt_backend.py | 42 +++++++++++++++- tests/test_stateful_candidates.py | 79 +++++++++++++++++++++++++++++++ uv.lock | 2 +- 8 files changed, 174 insertions(+), 28 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9a8c709..6e51525 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index a5b73ff..bf52fbb 100644 --- a/.github/workflows/publish.yml +++ b/.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/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ec6404..48c1763 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index 9b6bfc4..6ab4106 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 5ade85d..7df6873 100644 --- a/pyproject.toml +++ b/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] diff --git a/tests/test_rlbwt_backend.py b/tests/test_rlbwt_backend.py index 17c9149..6425129 100644 --- a/tests/test_rlbwt_backend.py +++ b/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 diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py index e609da7..9b85f3c 100644 --- a/tests/test_stateful_candidates.py +++ b/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) diff --git a/uv.lock b/uv.lock index ba71646..2cd12eb 100644 --- a/uv.lock +++ b/uv.lock @@ -774,7 +774,7 @@ nodejs = [ [[package]] name = "rosa-torch" -version = "0.2.0" +version = "0.3.0" source = { editable = "." } dependencies = [ { name = "torch" }, From ef8baac511fa6248b5834b99e8ac4b75643197b2 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:02:44 +0800 Subject: [PATCH 2/3] Make compile cache coverage deterministic --- tests/test_rosa.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_rosa.py b/tests/test_rosa.py index b74c94f..86ebdfa 100644 --- a/tests/test_rosa.py +++ b/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") From 416ad14067ca096dd133377486a173335db24872 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:07:38 +0800 Subject: [PATCH 3/3] Cover concurrent compile cache readiness --- tests/test_rosa.py | 52 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tests/test_rosa.py b/tests/test_rosa.py index 86ebdfa..b3cbbbe 100644 --- a/tests/test_rosa.py +++ b/tests/test_rosa.py @@ -343,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, ...]] = []