From 89186614332f34941bc66a520d356c082ad72c40 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 03:41:11 +0800 Subject: [PATCH 01/29] Use dense Numba arrays for exact ROSA --- src/rosa/_numba_backend.py | 89 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 src/rosa/_numba_backend.py diff --git a/src/rosa/_numba_backend.py b/src/rosa/_numba_backend.py new file mode 100644 index 0000000..c977cdd --- /dev/null +++ b/src/rosa/_numba_backend.py @@ -0,0 +1,89 @@ +"""Optional Numba fast path for exact CPU ROSA inference.""" + +from __future__ import annotations + +import numpy as np +import torch +from numba import njit, prange +from torch import Tensor + + +@njit(cache=True, nogil=True) +def _predict_row(tokens: np.ndarray) -> np.ndarray: + n = tokens.shape[0] + vocabulary_size = int(tokens.max()) + 1 + max_states = 2 * n + 1 + transitions = np.full((max_states, vocabulary_size), -1, dtype=np.int32) + suffix_link = np.full(max_states, -1, dtype=np.int32) + length = np.zeros(max_states, dtype=np.int32) + latest_end = np.full(max_states, -1, dtype=np.int32) + predicted = np.full(n, -1, dtype=np.int64) + last = 0 + size = 1 + + for i in range(n): + token = int(tokens[i]) + current = size + size += 1 + length[current] = length[last] + 1 + state = last + + while state != -1 and transitions[state, token] == -1: + transitions[state, token] = current + state = suffix_link[state] + + if state == -1: + suffix_link[current] = 0 + else: + target = transitions[state, token] + if length[state] + 1 == length[target]: + suffix_link[current] = target + else: + clone = size + size += 1 + transitions[clone, :] = transitions[target, :] + length[clone] = length[state] + 1 + suffix_link[clone] = suffix_link[target] + latest_end[clone] = latest_end[target] + while state != -1 and transitions[state, token] == target: + transitions[state, token] = clone + state = suffix_link[state] + suffix_link[target] = clone + suffix_link[current] = clone + + last = current + state = last + while state != -1: + if length[state] > 0 and latest_end[state] >= 0: + predicted[i] = tokens[latest_end[state] + 1] + break + state = suffix_link[state] + + state = last + while state != -1: + latest_end[state] = i + state = suffix_link[state] + + return predicted + + +@njit(cache=True, nogil=True, parallel=True) +def _predict_batch(tokens: np.ndarray) -> np.ndarray: + output = np.empty(tokens.shape, dtype=np.int64) + for batch_index in prange(tokens.shape[0]): + output[batch_index] = _predict_row(tokens[batch_index]) + return output + + +def predict_exact(tokens: Tensor) -> Tensor: + """Return exact ROSA predictions using an optional CPU Numba backend.""" + + squeeze = tokens.ndim == 1 + if squeeze: + tokens = tokens.unsqueeze(0) + if tokens.ndim != 2: + raise ValueError("tokens must have shape [N] or [B, N]") + device = tokens.device + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + output = torch.from_numpy(_predict_batch(cpu_tokens.numpy())).to(device) + return output[0] if squeeze else output From 1748a0f9e731dc1666f55d8a3a5694492869f923 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 03:44:36 +0800 Subject: [PATCH 02/29] Use sparse Numba transitions --- src/rosa/_numba_backend.py | 116 ++++++++++++++++++++++++++++++++++--- 1 file changed, 109 insertions(+), 7 deletions(-) diff --git a/src/rosa/_numba_backend.py b/src/rosa/_numba_backend.py index c977cdd..206507a 100644 --- a/src/rosa/_numba_backend.py +++ b/src/rosa/_numba_backend.py @@ -8,18 +8,79 @@ from torch import Tensor +@njit(cache=True, nogil=True, inline="always") +def _find_transition( + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + state: int, + token: int, +) -> int: + edge = head[state] + while edge != -1: + if edge_token[edge] == token: + return edge_target[edge] + edge = edge_next[edge] + return -1 + + +@njit(cache=True, nogil=True, inline="always") +def _add_transition( + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + edge_count: int, + state: int, + token: int, + target: int, +) -> int: + if edge_count >= edge_token.shape[0]: + raise RuntimeError("suffix automaton transition capacity exceeded") + edge_token[edge_count] = token + edge_target[edge_count] = target + edge_next[edge_count] = head[state] + head[state] = edge_count + return edge_count + 1 + + +@njit(cache=True, nogil=True, inline="always") +def _replace_transition( + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + state: int, + token: int, + target: int, +) -> None: + edge = head[state] + while edge != -1: + if edge_token[edge] == token: + edge_target[edge] = target + return + edge = edge_next[edge] + raise RuntimeError("suffix automaton transition not found") + + @njit(cache=True, nogil=True) def _predict_row(tokens: np.ndarray) -> np.ndarray: n = tokens.shape[0] vocabulary_size = int(tokens.max()) + 1 max_states = 2 * n + 1 - transitions = np.full((max_states, vocabulary_size), -1, dtype=np.int32) + max_edges = 4 * n + vocabulary_size + 1 + head = np.full(max_states, -1, dtype=np.int32) + edge_token = np.empty(max_edges, dtype=np.int32) + edge_target = np.empty(max_edges, dtype=np.int32) + edge_next = np.empty(max_edges, dtype=np.int32) suffix_link = np.full(max_states, -1, dtype=np.int32) length = np.zeros(max_states, dtype=np.int32) latest_end = np.full(max_states, -1, dtype=np.int32) predicted = np.full(n, -1, dtype=np.int64) last = 0 size = 1 + edge_count = 0 for i in range(n): token = int(tokens[i]) @@ -28,25 +89,66 @@ def _predict_row(tokens: np.ndarray) -> np.ndarray: length[current] = length[last] + 1 state = last - while state != -1 and transitions[state, token] == -1: - transitions[state, token] = current + while ( + state != -1 + and _find_transition(head, edge_token, edge_target, edge_next, state, token) + == -1 + ): + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + edge_count, + state, + token, + current, + ) state = suffix_link[state] if state == -1: suffix_link[current] = 0 else: - target = transitions[state, token] + target = _find_transition( + head, edge_token, edge_target, edge_next, state, token + ) if length[state] + 1 == length[target]: suffix_link[current] = target else: clone = size size += 1 - transitions[clone, :] = transitions[target, :] length[clone] = length[state] + 1 suffix_link[clone] = suffix_link[target] latest_end[clone] = latest_end[target] - while state != -1 and transitions[state, token] == target: - transitions[state, token] = clone + edge = head[target] + while edge != -1: + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + edge_count, + clone, + edge_token[edge], + edge_target[edge], + ) + edge = edge_next[edge] + while ( + state != -1 + and _find_transition( + head, edge_token, edge_target, edge_next, state, token + ) + == target + ): + _replace_transition( + head, + edge_token, + edge_target, + edge_next, + state, + token, + clone, + ) state = suffix_link[state] suffix_link[target] = clone suffix_link[current] = clone From ce812fedaaf2761da56618b16df9fe573d0d77a9 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 04:27:52 +0800 Subject: [PATCH 03/29] Avoid Numba parallel overhead on small inputs --- src/rosa/_numba_backend.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/rosa/_numba_backend.py b/src/rosa/_numba_backend.py index 206507a..778a3ad 100644 --- a/src/rosa/_numba_backend.py +++ b/src/rosa/_numba_backend.py @@ -177,6 +177,14 @@ def _predict_batch(tokens: np.ndarray) -> np.ndarray: return output +@njit(cache=True, nogil=True) +def _predict_serial_batch(tokens: np.ndarray) -> np.ndarray: + output = np.empty(tokens.shape, dtype=np.int64) + for batch_index in range(tokens.shape[0]): + output[batch_index] = _predict_row(tokens[batch_index]) + return output + + def predict_exact(tokens: Tensor) -> Tensor: """Return exact ROSA predictions using an optional CPU Numba backend.""" @@ -187,5 +195,12 @@ def predict_exact(tokens: Tensor) -> Tensor: raise ValueError("tokens must have shape [N] or [B, N]") device = tokens.device cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() - output = torch.from_numpy(_predict_batch(cpu_tokens.numpy())).to(device) + cpu_array = cpu_tokens.numpy() + if cpu_array.shape[0] == 1: + output_array = _predict_row(cpu_array[0]).reshape(1, -1) + elif cpu_array.size <= 4096: + output_array = _predict_serial_batch(cpu_array) + else: + output_array = _predict_batch(cpu_array) + output = torch.from_numpy(output_array).to(device) return output[0] if squeeze else output From 0b57a07b74d7847da65a4e7c9395789d2e5ba76a Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 04:44:04 +0800 Subject: [PATCH 04/29] Integrate optional Numba inference backend --- .github/workflows/ci.yml | 2 +- pyproject.toml | 3 + src/rosa/_numba_backend.py | 28 +++-- tests/test_numba_backend.py | 59 ++++++++++ uv.lock | 227 +++++++++++++++++++++++++++++++++++- 5 files changed, 307 insertions(+), 12 deletions(-) create mode 100644 tests/test_numba_backend.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8e0029c..a9bcac6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,7 +34,7 @@ jobs: enable-cache: true - name: Install locked dependencies - run: uv sync --locked --all-groups + run: uv sync --locked --all-groups --extra numba - name: Lint run: uv run ruff check . diff --git a/pyproject.toml b/pyproject.toml index 78c16b8..32cc8fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,9 @@ maintainers = [ ] dependencies = ["torch"] +[project.optional-dependencies] +numba = ["numba>=0.66"] + [project.urls] Repository = "https://github.com/aabbdev/rosa" Issues = "https://github.com/aabbdev/rosa/issues" diff --git a/src/rosa/_numba_backend.py b/src/rosa/_numba_backend.py index 778a3ad..a8475ad 100644 --- a/src/rosa/_numba_backend.py +++ b/src/rosa/_numba_backend.py @@ -9,7 +9,7 @@ @njit(cache=True, nogil=True, inline="always") -def _find_transition( +def _find_transition( # pragma: no cover - executed as compiled Numba code head: np.ndarray, edge_token: np.ndarray, edge_target: np.ndarray, @@ -20,13 +20,13 @@ def _find_transition( edge = head[state] while edge != -1: if edge_token[edge] == token: - return edge_target[edge] + return int(edge_target[edge]) edge = edge_next[edge] return -1 @njit(cache=True, nogil=True, inline="always") -def _add_transition( +def _add_transition( # pragma: no cover - executed as compiled Numba code head: np.ndarray, edge_token: np.ndarray, edge_target: np.ndarray, @@ -46,7 +46,7 @@ def _add_transition( @njit(cache=True, nogil=True, inline="always") -def _replace_transition( +def _replace_transition( # pragma: no cover - executed as compiled Numba code head: np.ndarray, edge_token: np.ndarray, edge_target: np.ndarray, @@ -65,13 +65,14 @@ def _replace_transition( @njit(cache=True, nogil=True) -def _predict_row(tokens: np.ndarray) -> np.ndarray: +def _predict_row( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, +) -> np.ndarray: n = tokens.shape[0] - vocabulary_size = int(tokens.max()) + 1 max_states = 2 * n + 1 - max_edges = 4 * n + vocabulary_size + 1 + max_edges = 4 * n + 1 head = np.full(max_states, -1, dtype=np.int32) - edge_token = np.empty(max_edges, dtype=np.int32) + edge_token = np.empty(max_edges, dtype=np.int64) edge_target = np.empty(max_edges, dtype=np.int32) edge_next = np.empty(max_edges, dtype=np.int32) suffix_link = np.full(max_states, -1, dtype=np.int32) @@ -170,7 +171,9 @@ def _predict_row(tokens: np.ndarray) -> np.ndarray: @njit(cache=True, nogil=True, parallel=True) -def _predict_batch(tokens: np.ndarray) -> np.ndarray: +def _predict_batch( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, +) -> np.ndarray: output = np.empty(tokens.shape, dtype=np.int64) for batch_index in prange(tokens.shape[0]): output[batch_index] = _predict_row(tokens[batch_index]) @@ -178,7 +181,9 @@ def _predict_batch(tokens: np.ndarray) -> np.ndarray: @njit(cache=True, nogil=True) -def _predict_serial_batch(tokens: np.ndarray) -> np.ndarray: +def _predict_serial_batch( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, +) -> np.ndarray: output = np.empty(tokens.shape, dtype=np.int64) for batch_index in range(tokens.shape[0]): output[batch_index] = _predict_row(tokens[batch_index]) @@ -194,6 +199,9 @@ def predict_exact(tokens: Tensor) -> Tensor: if tokens.ndim != 2: raise ValueError("tokens must have shape [N] or [B, N]") device = tokens.device + if tokens.shape[1] == 0: + output = torch.empty(tokens.shape, dtype=torch.long, device=device) + return output[0] if squeeze else output cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() cpu_array = cpu_tokens.numpy() if cpu_array.shape[0] == 1: diff --git a/tests/test_numba_backend.py b/tests/test_numba_backend.py new file mode 100644 index 0000000..04eafe6 --- /dev/null +++ b/tests/test_numba_backend.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import random +import unittest + +import torch + +from rosa import reference_rosa + +try: + from rosa._numba_backend import predict_exact +except ModuleNotFoundError as error: + if error.name not in {"numba", "numpy"}: + raise + predict_exact = None + + +@unittest.skipIf(predict_exact is None, "rosa-torch[numba] is not installed") +class TestNumbaBackend(unittest.TestCase): + def assert_matches_reference(self, tokens: torch.Tensor) -> None: + expected, _, _ = reference_rosa(tokens) + assert predict_exact is not None + self.assertTrue(torch.equal(predict_exact(tokens), expected)) + + def test_empty_squeezed_and_batched_inputs(self) -> None: + self.assert_matches_reference(torch.empty(0, dtype=torch.long)) + self.assert_matches_reference(torch.empty((3, 0), dtype=torch.long)) + + squeezed = torch.tensor([0, 1, 0, 2, 0], dtype=torch.long) + self.assert_matches_reference(squeezed) + self.assertEqual(tuple(predict_exact(squeezed).shape), (5,)) + + def test_serial_and_parallel_dispatch_match_reference(self) -> None: + generator = torch.Generator().manual_seed(20260811) + for shape in ((1, 128), (8, 128), (8, 512), (8, 513)): + tokens = torch.randint(256, shape, generator=generator, dtype=torch.long) + self.assert_matches_reference(tokens) + + def test_random_negative_and_large_token_ids(self) -> None: + rng = random.Random(20260811) + alphabet = (-10_000_000, -1, 0, 2**31, 10**12) + for length in range(1, 40): + with self.subTest(length=length): + tokens = torch.tensor( + [[rng.choice(alphabet) for _ in range(length)]], dtype=torch.long + ) + self.assert_matches_reference(tokens) + + def test_shape_validation(self) -> None: + assert predict_exact is not None + with self.assertRaisesRegex(ValueError, "shape"): + predict_exact(torch.zeros(1, 2, 3, dtype=torch.long)) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_cuda_round_trip_matches_reference(self) -> None: + tokens = torch.tensor( + [[0, 1, 0, 2, 0, 1, 0, 3]], dtype=torch.long, device="cuda" + ) + self.assert_matches_reference(tokens) diff --git a/uv.lock b/uv.lock index a326bdf..b1cb8fa 100644 --- a/uv.lock +++ b/uv.lock @@ -243,6 +243,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "llvmlite" +version = "0.48.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/dc/a0/acc8ffcd5bdc63df0097e22c719bfcd61b604358343089313a8aebbb24ab/llvmlite-0.48.0.tar.gz", hash = "sha256:543b19f9ef8f3c7c60d1468191e4ee1b1537bf9f8a3d56f64c0ddd98de92edd2", size = 184016, upload-time = "2026-07-02T20:20:05.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/4e/32543c42568fb321b3bdfcf9106e4116ab8f5a7bbcfd9ecf5569b0c07d83/llvmlite-0.48.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:614aad57df707e3172efd5165f2aa7da6a0c6897e40dce590bf756396815ba76", size = 40480650, upload-time = "2026-07-01T18:41:01.945Z" }, + { url = "https://files.pythonhosted.org/packages/a9/0d/6aa48abd423067139a129d1434b77bbcc56080db51d12a88510bb491ca3d/llvmlite-0.48.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:13532f248960ba888ad5ab8150494e2f3a3d20e5f59f264e63741ea5b0ba844c", size = 59890118, upload-time = "2026-07-01T18:41:10.608Z" }, + { url = "https://files.pythonhosted.org/packages/5a/c7/aa917444d871a79608af49149de1b28764e87d2ab41f933c5cd02431d03d/llvmlite-0.48.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ee0c77685a18f5fca994ae21d0763007fca5c5c64b41de37accc78b69079176", size = 58343459, upload-time = "2026-07-01T18:41:06.21Z" }, + { url = "https://files.pythonhosted.org/packages/c5/2b/ceee1cdc263617109d514ac4d1b31f10a282662740ff7d5777baae25b3b5/llvmlite-0.48.0-cp310-cp310-win_amd64.whl", hash = "sha256:02853fe4214acb3780fc920c3fee10564b61d58a35e1b78afcc8a546c2deaba3", size = 41864734, upload-time = "2026-07-01T18:41:14.746Z" }, + { url = "https://files.pythonhosted.org/packages/9a/55/595981f14fbae9ba966feb12af552b1fe69889e44e64ac883a731ed335e0/llvmlite-0.48.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:56a7e24607d3f02d7b1bae8d29c7e1e423d53143d68b072999777f19678fe77b", size = 40480651, upload-time = "2026-07-01T18:41:18.438Z" }, + { url = "https://files.pythonhosted.org/packages/26/08/0109d1b9cb3f4603f3890e30bc66c65332b79185f12a045343b2ae431f67/llvmlite-0.48.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6fa532d6bb3fd3f0803567c736401c54aecfe1a396d3ad25d2440d220e09f0e7", size = 59890118, upload-time = "2026-07-01T18:41:28.184Z" }, + { url = "https://files.pythonhosted.org/packages/02/eb/c5281be180c789cdffbf45b671884c57d7e61345ef3b0f643a4965e108e8/llvmlite-0.48.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:979a66a3f28a02565383ff463527dce78e9b856298872a361283132488e83591", size = 58343458, upload-time = "2026-07-01T18:41:23.397Z" }, + { url = "https://files.pythonhosted.org/packages/aa/f7/b3222b13f2d424dae3c9e63fde476af25ebccf1f3faf0b52d1b79fc15c70/llvmlite-0.48.0-cp311-cp311-win_amd64.whl", hash = "sha256:efaee0276e5e17c2b99b92e0c974bd484ef5977cf5dbc9168e82b71578edb47f", size = 41864734, upload-time = "2026-07-01T18:41:31.932Z" }, + { url = "https://files.pythonhosted.org/packages/92/a2/28696a9e61e245d1a79816d29d106692a90a2b6e7d78c98b326db70827af/llvmlite-0.48.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:d66c3beb4209087ddd4cf4ed2a0856b6887e6a913bdcf1aacfec9851cf2cba4e", size = 40480651, upload-time = "2026-07-01T18:41:35.694Z" }, + { url = "https://files.pythonhosted.org/packages/80/f2/72409351db66d0a317ec5087e076f31fb7b773a640db8a90ce6b5cac9edd/llvmlite-0.48.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:416fa4c2c66c2c6dc6d0a402648c19206e548efa0aa1eff01ad5cdad0af8217d", size = 59890118, upload-time = "2026-07-01T18:41:44.886Z" }, + { url = "https://files.pythonhosted.org/packages/3a/27/5ae2f3722606360480707adb47f001ad89df8251d06b14ee80336e660b66/llvmlite-0.48.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f5e5a5131045b72345c71062ea1a91910dde913792b6c9b28ebb2c1c0a712e98", size = 58343459, upload-time = "2026-07-01T18:41:40.306Z" }, + { url = "https://files.pythonhosted.org/packages/16/78/d824ffff7521cd140dc2006e44ce2bc82e64b48d1b32e90e956308c85a74/llvmlite-0.48.0-cp312-cp312-win_amd64.whl", hash = "sha256:d45c7541a80934ec6d8ab0defe67439494ecd2193cbf852a44ba827808976ac1", size = 41865022, upload-time = "2026-07-01T18:41:48.663Z" }, + { url = "https://files.pythonhosted.org/packages/9c/23/fe9316d14626b42c73ef0b502e724705a6ee9450afe53759c0a99c37c2d7/llvmlite-0.48.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:a83a99ef0c05b4ccddf9b6218ed9fe84b653a0caf7c1d9dbe148d6d16c67f518", size = 40480652, upload-time = "2026-07-01T18:41:52.216Z" }, + { url = "https://files.pythonhosted.org/packages/1b/4a/90715fa12006d681270b08d881195b6fab3ec39572e048764a1f7f59fed7/llvmlite-0.48.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8761b9e522f55207e24424fcd98370289eec2710bf8e915c82d1053f642450dc", size = 59890120, upload-time = "2026-07-01T18:42:00.748Z" }, + { url = "https://files.pythonhosted.org/packages/70/5e/7b3e20d64650ca3c80af0cdb664ec4b575ec83d9d4dd05bea8bd31f9bbb6/llvmlite-0.48.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2fe5cb59b2063bfa039dcb8ca6481c0181bf552f340d10dcf61d7996a665556e", size = 58343457, upload-time = "2026-07-01T18:41:56.41Z" }, + { url = "https://files.pythonhosted.org/packages/17/97/5a430055d1838cf1fb7a01cfa943300f5e4c026fc6333a522c5e4a03b0c1/llvmlite-0.48.0-cp313-cp313-win_amd64.whl", hash = "sha256:91c7e24e74cde3f02b88aa5acca678373f9e069f3b98531b3dbb3a142d9d10bb", size = 41865022, upload-time = "2026-07-01T18:42:04.57Z" }, + { url = "https://files.pythonhosted.org/packages/8d/8e/8170f2e0c217f88069c333d85bb976e536b332aecfcce606ddbdb249385f/llvmlite-0.48.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:321f1ac39b462603f0b589751aecf2d237d056f6d005749c1752b6f23ec3f074", size = 40480650, upload-time = "2026-07-01T18:42:07.935Z" }, + { url = "https://files.pythonhosted.org/packages/d6/e1/05b50692b647cac3c18200ac485b04f342f00ed173c9cc46767274469a15/llvmlite-0.48.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:05f0103c8f2f96a37441337e3643863c01b8e83e530aff38960dcb383c54a065", size = 59890115, upload-time = "2026-07-01T18:42:17.805Z" }, + { url = "https://files.pythonhosted.org/packages/f7/c3/470b8c4ff9ae2db2f9cf5c3e73de76ed908a32788ae9eb5602d43e6a476b/llvmlite-0.48.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:37d66fae72802175b0bfe1ea06e624b51e2d7aee6c3c34bbd09739b8f88e8e0b", size = 58343457, upload-time = "2026-07-01T18:42:13.217Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2d/6a5171fb7236ac0895e1a02ccba3735bf291e8597239aa6421894d3c0ba8/llvmlite-0.48.0-cp314-cp314-win_amd64.whl", hash = "sha256:966dcab0a598e2bd8fb5f2cc082cf7b07bae564fc485a3a8692393caf986facf", size = 42986372, upload-time = "2026-07-01T18:42:21.483Z" }, + { url = "https://files.pythonhosted.org/packages/94/e3/7a93e09c9f94e637ca90209ceef0334a9a1d45b0bdb7c92ff922d25d6187/llvmlite-0.48.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:7a5c413317050a1d67c34708bde97707f9b2257ef1017f7532d21fe7d9a9ff30", size = 40480654, upload-time = "2026-07-01T18:42:25.076Z" }, + { url = "https://files.pythonhosted.org/packages/27/98/a29133b4728671a175f7d616fab8b1c6e1d8c269d1523581d3160697bfb1/llvmlite-0.48.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d9f952dff6c350c529997423d4fa43abae9722a884ac7ffafc37a3af676e7db", size = 59890119, upload-time = "2026-07-01T18:42:33.88Z" }, + { url = "https://files.pythonhosted.org/packages/1a/cf/7aac11a1f1c7ec54b60c7f6814e87561fb6b55b2f290455d7941eb113420/llvmlite-0.48.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:054aa7d46595935565f276cf0c1659b4f10929c996dd4a606875fae26fba2a23", size = 58343460, upload-time = "2026-07-01T18:42:29.545Z" }, + { url = "https://files.pythonhosted.org/packages/db/41/b96f440c7df5ebba07872cad4e30fbc3560387755b1ea0b629adb76d5ca8/llvmlite-0.48.0-cp314-cp314t-win_amd64.whl", hash = "sha256:d0b3c61aac83b42fb48cc96bffbf57c81b82b2aa92276b7ed6420c814629a99a", size = 42986383, upload-time = "2026-07-01T18:42:37.544Z" }, +] + [[package]] name = "markupsafe" version = "3.0.3" @@ -386,6 +418,190 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/17/66/1ed71f1f529b8ca727d42c7ceb9db0bef145ce4a13dfc86fb50aa44f3be6/nodejs_wheel_binaries-24.16.0-py2.py3-none-win_arm64.whl", hash = "sha256:8308940b5edd0a50dc5267ea36ba21c9f668e83fe0d9f293937174d3a7e31c36", size = 39714528, upload-time = "2026-05-30T16:52:06.421Z" }, ] +[[package]] +name = "numba" +version = "0.66.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "llvmlite" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ae/a0/570e3dc53e5602b49108f62a13e529f1eec8bfc7ef37d49c825924dcf546/numba-0.66.0.tar.gz", hash = "sha256:b900e63a0e26c05ea9a6d5a3a5a0a177cb64c5011887bf43edb8c3ed2c38d363", size = 2806181, upload-time = "2026-07-01T23:12:46.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/48/d139bde40f2359351bfe26ee1b261937f458ac177ab810d4f045ae1c9d92/numba-0.66.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:27951c47e0def9bf8afe580eb961102902e2fd23cb77924b7d9d7cc0f8b444cb", size = 2727368, upload-time = "2026-07-01T23:12:04.282Z" }, + { url = "https://files.pythonhosted.org/packages/36/e4/b780bfa9191410da50ba249cb3248a75014e17f611e72709cbddcb21f42d/numba-0.66.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc408c54b450f41582f4be1608f8981c1dcc44c7f40355cc150dd93015753407", size = 3803554, upload-time = "2026-07-01T23:12:06.379Z" }, + { url = "https://files.pythonhosted.org/packages/1c/b2/a051b96626bdf5c4d8fa6b8d450605c09638d85dc872ab63ef9a67096dca/numba-0.66.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7c14c044c06b453ec3fa7715dfe75425e2ba72c73377a7ffde6d9ec511dfd94c", size = 3510065, upload-time = "2026-07-01T23:12:08.051Z" }, + { url = "https://files.pythonhosted.org/packages/34/01/24dcdc3e919522e2efbd92969c281ff40deb1d5f8a994bcd0057081c158c/numba-0.66.0-cp310-cp310-win_amd64.whl", hash = "sha256:2338cc0d43609fe448930848fd35a5bc688761b986f81b597a6f45cc0f8c9577", size = 2780379, upload-time = "2026-07-01T23:12:09.772Z" }, + { url = "https://files.pythonhosted.org/packages/9e/02/970796b4daa709604cde22e87a7cda9bde473c278ea4a75f59fe38cee47f/numba-0.66.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:bbd531c327557a9004507fa6bff06c53ab51a7a5776b75261bb9cef1efe2b2ea", size = 2727049, upload-time = "2026-07-01T23:12:11.296Z" }, + { url = "https://files.pythonhosted.org/packages/8c/99/33a6ed9c1a0b5e42efa98eb0edf617d61dca576c82625947377b1d4540c9/numba-0.66.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc6629becb21a867d85401ec89f426dd24c484a4193ade8a38309debfd1529ca", size = 3808870, upload-time = "2026-07-01T23:12:12.944Z" }, + { url = "https://files.pythonhosted.org/packages/04/20/8c51126025211659235b8de2866dfa226984ae0c8273461a3cf374716741/numba-0.66.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aac69f3ccb8af100f5913c1241edc9692bad1cdd2508721713f426eb06c9a659", size = 3514498, upload-time = "2026-07-01T23:12:15.307Z" }, + { url = "https://files.pythonhosted.org/packages/5e/c9/9476940bc6d5caf5c0cf2e4c5feecbf01244bbe6f914614082dd7a3e520e/numba-0.66.0-cp311-cp311-win_amd64.whl", hash = "sha256:fb601841d9e02e6237bb6522e36d0741614be3cfe2b482a6f00a41b5ba209443", size = 2780225, upload-time = "2026-07-01T23:12:16.924Z" }, + { url = "https://files.pythonhosted.org/packages/62/a3/70deb7f88461c1cd5d16aa990c2380604102661a427667b8950dcdccc27f/numba-0.66.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:53ca5900b7cab15109796030113a6b28576bae5ad7bb507ad6dd1360ddd81ba4", size = 2727264, upload-time = "2026-07-01T23:12:18.669Z" }, + { url = "https://files.pythonhosted.org/packages/2d/55/25c319845e9a4e08f16611ddbda56a192eb7b6ed13e1a2bff2da272ffb97/numba-0.66.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0999e3ee1b18c48e1fb51d11af35ef59852c7f4f50569c9550c25faef0616ad1", size = 3866252, upload-time = "2026-07-01T23:12:20.429Z" }, + { url = "https://files.pythonhosted.org/packages/71/ef/a82d6fd6bf1b0fe461651e924d3647eeec9ac17f8eee4896264bf7480930/numba-0.66.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:efe0d2d5099790df945e0cb6e1b3104bd965d7bbfac50d62f1d5d1d6ade0825d", size = 3566974, upload-time = "2026-07-01T23:12:22.116Z" }, + { url = "https://files.pythonhosted.org/packages/fc/eb/9e6171e378822ab191c7abcfd3d8cfc8644516f6c7834c22e210e4acc070/numba-0.66.0-cp312-cp312-win_amd64.whl", hash = "sha256:b075a4e7ebc43dc6294f223e2821659656209fd5e0ce53245877c23d66d6e1a9", size = 2797403, upload-time = "2026-07-01T23:12:23.724Z" }, + { url = "https://files.pythonhosted.org/packages/03/52/176c02d005c5c5143cde10a85bbcdcb6236d9e34c3aac089380e0506cd1d/numba-0.66.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:380b2556a2019ccd1e956ae77dd257eaa39403f7520768b626d44b755112785e", size = 2727084, upload-time = "2026-07-01T23:12:25.434Z" }, + { url = "https://files.pythonhosted.org/packages/44/b5/e930010965568fe7f2c6c962fd2849d458cb9f62c3ab7584af8a19a2b40a/numba-0.66.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:939316d5d8619751207b8972a67852b5a7646665298cb4de693cd6bf135152f4", size = 3873663, upload-time = "2026-07-01T23:12:27.308Z" }, + { url = "https://files.pythonhosted.org/packages/d0/ec/5b51457cbe96e4831141d83e892e65191b23a1b78728456c62909d231ace/numba-0.66.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdf506775d9f02eb92a87bf5c5b1e0d25506fd18cafd769f4ed914a8feac73e7", size = 3573529, upload-time = "2026-07-01T23:12:28.944Z" }, + { url = "https://files.pythonhosted.org/packages/83/7e/cea7710e96913d3c7f2999f16db1b28e6c5be5171cbf40f77f98333a7243/numba-0.66.0-cp313-cp313-win_amd64.whl", hash = "sha256:c5bfe5350284509ab0474390321454c3a8627a188af5b68c910e83df3e2db4a7", size = 2797247, upload-time = "2026-07-01T23:12:30.774Z" }, + { url = "https://files.pythonhosted.org/packages/96/7a/7e0e73550eb4e41ede6e72fb5371f4539537a4d770a3b73fa9b61aea0622/numba-0.66.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:46ae5f2b19e2af3c33c2df100306a90ea2f981c8158b0390f8bf6c20eee7357e", size = 2727296, upload-time = "2026-07-01T23:12:32.39Z" }, + { url = "https://files.pythonhosted.org/packages/0f/26/885774c006de6620ed3d10f45d8e20fe0b8e6aad6d573211a2cbc8b3e528/numba-0.66.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e2b101f23b8b63978d334574d2039f27f0dccfe1d891756f33a2e2f3e4c88cf4", size = 3842720, upload-time = "2026-07-01T23:12:33.938Z" }, + { url = "https://files.pythonhosted.org/packages/93/99/edebf7de890b73973d839dd971cf73734adfb81ffa1b4504f84b9059c3e5/numba-0.66.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:63b943eb2c9ba371908ce2cd6dfc643db51fc40f7966993376a1701bc922f537", size = 3543537, upload-time = "2026-07-01T23:12:35.566Z" }, + { url = "https://files.pythonhosted.org/packages/66/c5/b46ad28ac3681d035ea21365c5e052149062e1a0a9affd0563d2760ea6ff/numba-0.66.0-cp314-cp314-win_amd64.whl", hash = "sha256:bd57790acd20f6a468e0ad333ef6b82355e309a92310fb7dff80e919f01a21a9", size = 2799250, upload-time = "2026-07-01T23:12:37.154Z" }, + { url = "https://files.pythonhosted.org/packages/10/6f/5e77a7397a37dd16f57a7b72e7e470db5227b68e3639df0d13a8e674883d/numba-0.66.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:db7735d15ea17a283d6485b9fa3504769f78fd86e5146638ad5e8da57c031b9e", size = 2730342, upload-time = "2026-07-01T23:12:38.758Z" }, + { url = "https://files.pythonhosted.org/packages/39/fd/e9c9680a3813f3d781c20e5d53c1074801b787d4feecca0472fdd7c05ce1/numba-0.66.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:651a2b53298340956db26ecbe7ab043106b50a40c2807e66d617a4917245a4ab", size = 3878695, upload-time = "2026-07-01T23:12:40.302Z" }, + { url = "https://files.pythonhosted.org/packages/61/3a/9b363287b85fcd4537ea3878793822878b2ac1008a78159d2096fea628de/numba-0.66.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8c1144ba1720ea59ad79f4f488ed54d149b2613b357f7e445678b7d0739c70e9", size = 3596323, upload-time = "2026-07-01T23:12:42.805Z" }, + { url = "https://files.pythonhosted.org/packages/4c/f2/dca53d50b8f2289dd01954ace9da261e0487d5b74b188b4304e4ecc3492c/numba-0.66.0-cp314-cp314t-win_amd64.whl", hash = "sha256:d426178fb991a85714c43112a8ea7b9d9579ea856ad8dcdb9c1c3941903ba5be", size = 2804772, upload-time = "2026-07-01T23:12:44.399Z" }, +] + +[[package]] +name = "numpy" +version = "2.2.6" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +sdist = { url = "https://files.pythonhosted.org/packages/76/21/7d2a95e4bba9dc13d043ee156a356c0a8f0c6309dff6b21b4d71a073b8a8/numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd", size = 20276440, upload-time = "2025-05-17T22:38:04.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/3e/ed6db5be21ce87955c0cbd3009f2803f59fa08df21b5df06862e2d8e2bdd/numpy-2.2.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b412caa66f72040e6d268491a59f2c43bf03eb6c96dd8f0307829feb7fa2b6fb", size = 21165245, upload-time = "2025-05-17T21:27:58.555Z" }, + { url = "https://files.pythonhosted.org/packages/22/c2/4b9221495b2a132cc9d2eb862e21d42a009f5a60e45fc44b00118c174bff/numpy-2.2.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e41fd67c52b86603a91c1a505ebaef50b3314de0213461c7a6e99c9a3beff90", size = 14360048, upload-time = "2025-05-17T21:28:21.406Z" }, + { url = "https://files.pythonhosted.org/packages/fd/77/dc2fcfc66943c6410e2bf598062f5959372735ffda175b39906d54f02349/numpy-2.2.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:37e990a01ae6ec7fe7fa1c26c55ecb672dd98b19c3d0e1d1f326fa13cb38d163", size = 5340542, upload-time = "2025-05-17T21:28:30.931Z" }, + { url = "https://files.pythonhosted.org/packages/7a/4f/1cb5fdc353a5f5cc7feb692db9b8ec2c3d6405453f982435efc52561df58/numpy-2.2.6-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:5a6429d4be8ca66d889b7cf70f536a397dc45ba6faeb5f8c5427935d9592e9cf", size = 6878301, upload-time = "2025-05-17T21:28:41.613Z" }, + { url = "https://files.pythonhosted.org/packages/eb/17/96a3acd228cec142fcb8723bd3cc39c2a474f7dcf0a5d16731980bcafa95/numpy-2.2.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:efd28d4e9cd7d7a8d39074a4d44c63eda73401580c5c76acda2ce969e0a38e83", size = 14297320, upload-time = "2025-05-17T21:29:02.78Z" }, + { url = "https://files.pythonhosted.org/packages/b4/63/3de6a34ad7ad6646ac7d2f55ebc6ad439dbbf9c4370017c50cf403fb19b5/numpy-2.2.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc7b73d02efb0e18c000e9ad8b83480dfcd5dfd11065997ed4c6747470ae8915", size = 16801050, upload-time = "2025-05-17T21:29:27.675Z" }, + { url = "https://files.pythonhosted.org/packages/07/b6/89d837eddef52b3d0cec5c6ba0456c1bf1b9ef6a6672fc2b7873c3ec4e2e/numpy-2.2.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:74d4531beb257d2c3f4b261bfb0fc09e0f9ebb8842d82a7b4209415896adc680", size = 15807034, upload-time = "2025-05-17T21:29:51.102Z" }, + { url = "https://files.pythonhosted.org/packages/01/c8/dc6ae86e3c61cfec1f178e5c9f7858584049b6093f843bca541f94120920/numpy-2.2.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8fc377d995680230e83241d8a96def29f204b5782f371c532579b4f20607a289", size = 18614185, upload-time = "2025-05-17T21:30:18.703Z" }, + { url = "https://files.pythonhosted.org/packages/5b/c5/0064b1b7e7c89137b471ccec1fd2282fceaae0ab3a9550f2568782d80357/numpy-2.2.6-cp310-cp310-win32.whl", hash = "sha256:b093dd74e50a8cba3e873868d9e93a85b78e0daf2e98c6797566ad8044e8363d", size = 6527149, upload-time = "2025-05-17T21:30:29.788Z" }, + { url = "https://files.pythonhosted.org/packages/a3/dd/4b822569d6b96c39d1215dbae0582fd99954dcbcf0c1a13c61783feaca3f/numpy-2.2.6-cp310-cp310-win_amd64.whl", hash = "sha256:f0fd6321b839904e15c46e0d257fdd101dd7f530fe03fd6359c1ea63738703f3", size = 12904620, upload-time = "2025-05-17T21:30:48.994Z" }, + { url = "https://files.pythonhosted.org/packages/da/a8/4f83e2aa666a9fbf56d6118faaaf5f1974d456b1823fda0a176eff722839/numpy-2.2.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f9f1adb22318e121c5c69a09142811a201ef17ab257a1e66ca3025065b7f53ae", size = 21176963, upload-time = "2025-05-17T21:31:19.36Z" }, + { url = "https://files.pythonhosted.org/packages/b3/2b/64e1affc7972decb74c9e29e5649fac940514910960ba25cd9af4488b66c/numpy-2.2.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c820a93b0255bc360f53eca31a0e676fd1101f673dda8da93454a12e23fc5f7a", size = 14406743, upload-time = "2025-05-17T21:31:41.087Z" }, + { url = "https://files.pythonhosted.org/packages/4a/9f/0121e375000b5e50ffdd8b25bf78d8e1a5aa4cca3f185d41265198c7b834/numpy-2.2.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3d70692235e759f260c3d837193090014aebdf026dfd167834bcba43e30c2a42", size = 5352616, upload-time = "2025-05-17T21:31:50.072Z" }, + { url = "https://files.pythonhosted.org/packages/31/0d/b48c405c91693635fbe2dcd7bc84a33a602add5f63286e024d3b6741411c/numpy-2.2.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:481b49095335f8eed42e39e8041327c05b0f6f4780488f61286ed3c01368d491", size = 6889579, upload-time = "2025-05-17T21:32:01.712Z" }, + { url = "https://files.pythonhosted.org/packages/52/b8/7f0554d49b565d0171eab6e99001846882000883998e7b7d9f0d98b1f934/numpy-2.2.6-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b64d8d4d17135e00c8e346e0a738deb17e754230d7e0810ac5012750bbd85a5a", size = 14312005, upload-time = "2025-05-17T21:32:23.332Z" }, + { url = "https://files.pythonhosted.org/packages/b3/dd/2238b898e51bd6d389b7389ffb20d7f4c10066d80351187ec8e303a5a475/numpy-2.2.6-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba10f8411898fc418a521833e014a77d3ca01c15b0c6cdcce6a0d2897e6dbbdf", size = 16821570, upload-time = "2025-05-17T21:32:47.991Z" }, + { url = "https://files.pythonhosted.org/packages/83/6c/44d0325722cf644f191042bf47eedad61c1e6df2432ed65cbe28509d404e/numpy-2.2.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bd48227a919f1bafbdda0583705e547892342c26fb127219d60a5c36882609d1", size = 15818548, upload-time = "2025-05-17T21:33:11.728Z" }, + { url = "https://files.pythonhosted.org/packages/ae/9d/81e8216030ce66be25279098789b665d49ff19eef08bfa8cb96d4957f422/numpy-2.2.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9551a499bf125c1d4f9e250377c1ee2eddd02e01eac6644c080162c0c51778ab", size = 18620521, upload-time = "2025-05-17T21:33:39.139Z" }, + { url = "https://files.pythonhosted.org/packages/6a/fd/e19617b9530b031db51b0926eed5345ce8ddc669bb3bc0044b23e275ebe8/numpy-2.2.6-cp311-cp311-win32.whl", hash = "sha256:0678000bb9ac1475cd454c6b8c799206af8107e310843532b04d49649c717a47", size = 6525866, upload-time = "2025-05-17T21:33:50.273Z" }, + { url = "https://files.pythonhosted.org/packages/31/0a/f354fb7176b81747d870f7991dc763e157a934c717b67b58456bc63da3df/numpy-2.2.6-cp311-cp311-win_amd64.whl", hash = "sha256:e8213002e427c69c45a52bbd94163084025f533a55a59d6f9c5b820774ef3303", size = 12907455, upload-time = "2025-05-17T21:34:09.135Z" }, + { url = "https://files.pythonhosted.org/packages/82/5d/c00588b6cf18e1da539b45d3598d3557084990dcc4331960c15ee776ee41/numpy-2.2.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41c5a21f4a04fa86436124d388f6ed60a9343a6f767fced1a8a71c3fbca038ff", size = 20875348, upload-time = "2025-05-17T21:34:39.648Z" }, + { url = "https://files.pythonhosted.org/packages/66/ee/560deadcdde6c2f90200450d5938f63a34b37e27ebff162810f716f6a230/numpy-2.2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:de749064336d37e340f640b05f24e9e3dd678c57318c7289d222a8a2f543e90c", size = 14119362, upload-time = "2025-05-17T21:35:01.241Z" }, + { url = "https://files.pythonhosted.org/packages/3c/65/4baa99f1c53b30adf0acd9a5519078871ddde8d2339dc5a7fde80d9d87da/numpy-2.2.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:894b3a42502226a1cac872f840030665f33326fc3dac8e57c607905773cdcde3", size = 5084103, upload-time = "2025-05-17T21:35:10.622Z" }, + { url = "https://files.pythonhosted.org/packages/cc/89/e5a34c071a0570cc40c9a54eb472d113eea6d002e9ae12bb3a8407fb912e/numpy-2.2.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:71594f7c51a18e728451bb50cc60a3ce4e6538822731b2933209a1f3614e9282", size = 6625382, upload-time = "2025-05-17T21:35:21.414Z" }, + { url = "https://files.pythonhosted.org/packages/f8/35/8c80729f1ff76b3921d5c9487c7ac3de9b2a103b1cd05e905b3090513510/numpy-2.2.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2618db89be1b4e05f7a1a847a9c1c0abd63e63a1607d892dd54668dd92faf87", size = 14018462, upload-time = "2025-05-17T21:35:42.174Z" }, + { url = "https://files.pythonhosted.org/packages/8c/3d/1e1db36cfd41f895d266b103df00ca5b3cbe965184df824dec5c08c6b803/numpy-2.2.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd83c01228a688733f1ded5201c678f0c53ecc1006ffbc404db9f7a899ac6249", size = 16527618, upload-time = "2025-05-17T21:36:06.711Z" }, + { url = "https://files.pythonhosted.org/packages/61/c6/03ed30992602c85aa3cd95b9070a514f8b3c33e31124694438d88809ae36/numpy-2.2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:37c0ca431f82cd5fa716eca9506aefcabc247fb27ba69c5062a6d3ade8cf8f49", size = 15505511, upload-time = "2025-05-17T21:36:29.965Z" }, + { url = "https://files.pythonhosted.org/packages/b7/25/5761d832a81df431e260719ec45de696414266613c9ee268394dd5ad8236/numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de", size = 18313783, upload-time = "2025-05-17T21:36:56.883Z" }, + { url = "https://files.pythonhosted.org/packages/57/0a/72d5a3527c5ebffcd47bde9162c39fae1f90138c961e5296491ce778e682/numpy-2.2.6-cp312-cp312-win32.whl", hash = "sha256:4eeaae00d789f66c7a25ac5f34b71a7035bb474e679f410e5e1a94deb24cf2d4", size = 6246506, upload-time = "2025-05-17T21:37:07.368Z" }, + { url = "https://files.pythonhosted.org/packages/36/fa/8c9210162ca1b88529ab76b41ba02d433fd54fecaf6feb70ef9f124683f1/numpy-2.2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c1f9540be57940698ed329904db803cf7a402f3fc200bfe599334c9bd84a40b2", size = 12614190, upload-time = "2025-05-17T21:37:26.213Z" }, + { url = "https://files.pythonhosted.org/packages/f9/5c/6657823f4f594f72b5471f1db1ab12e26e890bb2e41897522d134d2a3e81/numpy-2.2.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0811bb762109d9708cca4d0b13c4f67146e3c3b7cf8d34018c722adb2d957c84", size = 20867828, upload-time = "2025-05-17T21:37:56.699Z" }, + { url = "https://files.pythonhosted.org/packages/dc/9e/14520dc3dadf3c803473bd07e9b2bd1b69bc583cb2497b47000fed2fa92f/numpy-2.2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:287cc3162b6f01463ccd86be154f284d0893d2b3ed7292439ea97eafa8170e0b", size = 14143006, upload-time = "2025-05-17T21:38:18.291Z" }, + { url = "https://files.pythonhosted.org/packages/4f/06/7e96c57d90bebdce9918412087fc22ca9851cceaf5567a45c1f404480e9e/numpy-2.2.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f1372f041402e37e5e633e586f62aa53de2eac8d98cbfb822806ce4bbefcb74d", size = 5076765, upload-time = "2025-05-17T21:38:27.319Z" }, + { url = "https://files.pythonhosted.org/packages/73/ed/63d920c23b4289fdac96ddbdd6132e9427790977d5457cd132f18e76eae0/numpy-2.2.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:55a4d33fa519660d69614a9fad433be87e5252f4b03850642f88993f7b2ca566", size = 6617736, upload-time = "2025-05-17T21:38:38.141Z" }, + { url = "https://files.pythonhosted.org/packages/85/c5/e19c8f99d83fd377ec8c7e0cf627a8049746da54afc24ef0a0cb73d5dfb5/numpy-2.2.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92729c95468a2f4f15e9bb94c432a9229d0d50de67304399627a943201baa2f", size = 14010719, upload-time = "2025-05-17T21:38:58.433Z" }, + { url = "https://files.pythonhosted.org/packages/19/49/4df9123aafa7b539317bf6d342cb6d227e49f7a35b99c287a6109b13dd93/numpy-2.2.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bc23a79bfabc5d056d106f9befb8d50c31ced2fbc70eedb8155aec74a45798f", size = 16526072, upload-time = "2025-05-17T21:39:22.638Z" }, + { url = "https://files.pythonhosted.org/packages/b2/6c/04b5f47f4f32f7c2b0e7260442a8cbcf8168b0e1a41ff1495da42f42a14f/numpy-2.2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3143e4451880bed956e706a3220b4e5cf6172ef05fcc397f6f36a550b1dd868", size = 15503213, upload-time = "2025-05-17T21:39:45.865Z" }, + { url = "https://files.pythonhosted.org/packages/17/0a/5cd92e352c1307640d5b6fec1b2ffb06cd0dabe7d7b8227f97933d378422/numpy-2.2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4f13750ce79751586ae2eb824ba7e1e8dba64784086c98cdbbcc6a42112ce0d", size = 18316632, upload-time = "2025-05-17T21:40:13.331Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3b/5cba2b1d88760ef86596ad0f3d484b1cbff7c115ae2429678465057c5155/numpy-2.2.6-cp313-cp313-win32.whl", hash = "sha256:5beb72339d9d4fa36522fc63802f469b13cdbe4fdab4a288f0c441b74272ebfd", size = 6244532, upload-time = "2025-05-17T21:43:46.099Z" }, + { url = "https://files.pythonhosted.org/packages/cb/3b/d58c12eafcb298d4e6d0d40216866ab15f59e55d148a5658bb3132311fcf/numpy-2.2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b0544343a702fa80c95ad5d3d608ea3599dd54d4632df855e4c8d24eb6ecfa1c", size = 12610885, upload-time = "2025-05-17T21:44:05.145Z" }, + { url = "https://files.pythonhosted.org/packages/6b/9e/4bf918b818e516322db999ac25d00c75788ddfd2d2ade4fa66f1f38097e1/numpy-2.2.6-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0bca768cd85ae743b2affdc762d617eddf3bcf8724435498a1e80132d04879e6", size = 20963467, upload-time = "2025-05-17T21:40:44Z" }, + { url = "https://files.pythonhosted.org/packages/61/66/d2de6b291507517ff2e438e13ff7b1e2cdbdb7cb40b3ed475377aece69f9/numpy-2.2.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fc0c5673685c508a142ca65209b4e79ed6740a4ed6b2267dbba90f34b0b3cfda", size = 14225144, upload-time = "2025-05-17T21:41:05.695Z" }, + { url = "https://files.pythonhosted.org/packages/e4/25/480387655407ead912e28ba3a820bc69af9adf13bcbe40b299d454ec011f/numpy-2.2.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:5bd4fc3ac8926b3819797a7c0e2631eb889b4118a9898c84f585a54d475b7e40", size = 5200217, upload-time = "2025-05-17T21:41:15.903Z" }, + { url = "https://files.pythonhosted.org/packages/aa/4a/6e313b5108f53dcbf3aca0c0f3e9c92f4c10ce57a0a721851f9785872895/numpy-2.2.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:fee4236c876c4e8369388054d02d0e9bb84821feb1a64dd59e137e6511a551f8", size = 6712014, upload-time = "2025-05-17T21:41:27.321Z" }, + { url = "https://files.pythonhosted.org/packages/b7/30/172c2d5c4be71fdf476e9de553443cf8e25feddbe185e0bd88b096915bcc/numpy-2.2.6-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1dda9c7e08dc141e0247a5b8f49cf05984955246a327d4c48bda16821947b2f", size = 14077935, upload-time = "2025-05-17T21:41:49.738Z" }, + { url = "https://files.pythonhosted.org/packages/12/fb/9e743f8d4e4d3c710902cf87af3512082ae3d43b945d5d16563f26ec251d/numpy-2.2.6-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f447e6acb680fd307f40d3da4852208af94afdfab89cf850986c3ca00562f4fa", size = 16600122, upload-time = "2025-05-17T21:42:14.046Z" }, + { url = "https://files.pythonhosted.org/packages/12/75/ee20da0e58d3a66f204f38916757e01e33a9737d0b22373b3eb5a27358f9/numpy-2.2.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:389d771b1623ec92636b0786bc4ae56abafad4a4c513d36a55dce14bd9ce8571", size = 15586143, upload-time = "2025-05-17T21:42:37.464Z" }, + { url = "https://files.pythonhosted.org/packages/76/95/bef5b37f29fc5e739947e9ce5179ad402875633308504a52d188302319c8/numpy-2.2.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8e9ace4a37db23421249ed236fdcdd457d671e25146786dfc96835cd951aa7c1", size = 18385260, upload-time = "2025-05-17T21:43:05.189Z" }, + { url = "https://files.pythonhosted.org/packages/09/04/f2f83279d287407cf36a7a8053a5abe7be3622a4363337338f2585e4afda/numpy-2.2.6-cp313-cp313t-win32.whl", hash = "sha256:038613e9fb8c72b0a41f025a7e4c3f0b7a1b5d768ece4796b674c8f3fe13efff", size = 6377225, upload-time = "2025-05-17T21:43:16.254Z" }, + { url = "https://files.pythonhosted.org/packages/67/0e/35082d13c09c02c011cf21570543d202ad929d961c02a147493cb0c2bdf5/numpy-2.2.6-cp313-cp313t-win_amd64.whl", hash = "sha256:6031dd6dfecc0cf9f668681a37648373bddd6421fff6c66ec1624eed0180ee06", size = 12771374, upload-time = "2025-05-17T21:43:35.479Z" }, + { url = "https://files.pythonhosted.org/packages/9e/3b/d94a75f4dbf1ef5d321523ecac21ef23a3cd2ac8b78ae2aac40873590229/numpy-2.2.6-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0b605b275d7bd0c640cad4e5d30fa701a8d59302e127e5f79138ad62762c3e3d", size = 21040391, upload-time = "2025-05-17T21:44:35.948Z" }, + { url = "https://files.pythonhosted.org/packages/17/f4/09b2fa1b58f0fb4f7c7963a1649c64c4d315752240377ed74d9cd878f7b5/numpy-2.2.6-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:7befc596a7dc9da8a337f79802ee8adb30a552a94f792b9c9d18c840055907db", size = 6786754, upload-time = "2025-05-17T21:44:47.446Z" }, + { url = "https://files.pythonhosted.org/packages/af/30/feba75f143bdc868a1cc3f44ccfa6c4b9ec522b36458e738cd00f67b573f/numpy-2.2.6-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce47521a4754c8f4593837384bd3424880629f718d87c5d44f8ed763edd63543", size = 16643476, upload-time = "2025-05-17T21:45:11.871Z" }, + { url = "https://files.pythonhosted.org/packages/37/48/ac2a9584402fb6c0cd5b5d1a91dcf176b15760130dd386bbafdbfe3640bf/numpy-2.2.6-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d042d24c90c41b54fd506da306759e06e568864df8ec17ccc17e9e884634fd00", size = 12812666, upload-time = "2025-05-17T21:45:31.426Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.6" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.11'", +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/49/ec46835a70be8fa6446c495126ac84fdb28cb2558e1620ffb87a10c8b64c/numpy-2.4.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0280e0356c0829a18d9de1cb7eee50ec22ca639878d7240307ca0943d73cd2c4", size = 16969194, upload-time = "2026-05-18T23:33:13.503Z" }, + { url = "https://files.pythonhosted.org/packages/0e/0d/f5957185c0ee2f3e12f78715aa9e3b353fd83633316c8532b38faa37e3f6/numpy-2.4.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:110f8b71aacb688ec69062bb7f6938a0f8acb01b7c1c4beb453c65b6d234584d", size = 14964111, upload-time = "2026-05-18T23:33:17.795Z" }, + { url = "https://files.pythonhosted.org/packages/ad/40/40a40ee0ddf7ceb782c49af278894b686e586d65d8c1889c8b5da01a3d7d/numpy-2.4.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:4cfe66903cc32a9921a6733d96b19bb6abf310397581bbad89c228f5abaf0ee8", size = 5469159, upload-time = "2026-05-18T23:33:20.654Z" }, + { url = "https://files.pythonhosted.org/packages/63/13/f9a8046535cb21deae82f8d03de9617e08882d274fad2539630761888228/numpy-2.4.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:8155154c7c691289fe18f510b5d4657c68c67989f293f0535a91360392ff6538", size = 6798936, upload-time = "2026-05-18T23:33:22.987Z" }, + { url = "https://files.pythonhosted.org/packages/33/a8/6fa8c1a345a8c85dbb21932c447bee07c30a2c2a3f31e369c0a84b300147/numpy-2.4.6-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ab0a9c4ffb1a6d95ef519fe4247dba8eb6b18ad93999f76b7f657039acabd47", size = 15966692, upload-time = "2026-05-18T23:33:26.62Z" }, + { url = "https://files.pythonhosted.org/packages/02/03/74fe2a4cb3817d94d86402f2506554130a2f01414e299b5a843e5a8a957f/numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:89cd468399cfd2504718f0ba50e410dca55a170b61a02ad92bb18c8a65186e93", size = 16918164, upload-time = "2026-05-18T23:33:29.955Z" }, + { url = "https://files.pythonhosted.org/packages/c5/80/3615be3313f7e7696609bc194b9f0101da809df79e859bdb84e0cd043f46/numpy-2.4.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c2d37ab77531417474168eb79d6d80b14f821a966818505d03013d0833edb7a8", size = 17322877, upload-time = "2026-05-18T23:33:34.724Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ac/a691e0fe2675e370d0e08ff905adc49a1c8830e8cae03efe4477e92cd55d/numpy-2.4.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f407cb6b8e9d6d8c626bc73c945db1706035af8fd632295547bf1c9e46d092d6", size = 18651487, upload-time = "2026-05-18T23:33:38.217Z" }, + { url = "https://files.pythonhosted.org/packages/15/a7/9bc1cd626d7bf6869bfedf27b91b6ab5dd607758bf8e959d6fa80c6a59cb/numpy-2.4.6-cp311-cp311-win32.whl", hash = "sha256:ddea102b48f9e339f3948bf22040944184627a30fdf7f858667673b9c5f033c8", size = 6233945, upload-time = "2026-05-18T23:33:41.331Z" }, + { url = "https://files.pythonhosted.org/packages/c5/31/7fc6239c12bce7e931463251cca4426c465e1876ba3cc785402ef4dd8f4e/numpy-2.4.6-cp311-cp311-win_amd64.whl", hash = "sha256:1e254a00cdf42b1e4d5b3d68d33af63268d41340d8885df2ab6470f2e1500147", size = 12608406, upload-time = "2026-05-18T23:33:44.131Z" }, + { url = "https://files.pythonhosted.org/packages/27/83/140f85a466595a16382996a1bf06b2b54bcd597488921b0c9daaeeda72af/numpy-2.4.6-cp311-cp311-win_arm64.whl", hash = "sha256:ed9749eef4cbd126da3dc1d6bcb3a57f5eb7ac6a6484146bdbf743f552dfc577", size = 10479528, upload-time = "2026-05-18T23:33:50.725Z" }, + { url = "https://files.pythonhosted.org/packages/95/2a/3d7b5ac8aac24feaf9ad7ed58f45b0bbc06d37e4338ae84c9f2298b570f9/numpy-2.4.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:001fbb8e08d942dd57599e781f2472269ee7f2755fae407b4f67b2f0b17da3f1", size = 16689119, upload-time = "2026-05-18T23:33:54.065Z" }, + { url = "https://files.pythonhosted.org/packages/ea/12/92c4c131527599e8288d6918e888d88726f84d805d784b771f32408aeaef/numpy-2.4.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ebfb099f8dcf083deef3ac1ca4c1503f387cf76296fcb3816b66f5ecb5f54fdb", size = 14699246, upload-time = "2026-05-18T23:33:57.621Z" }, + { url = "https://files.pythonhosted.org/packages/ad/fe/c0a6b7b2ca128a8fb228575147073b660656734b8ebe4d76c8fd748dcc79/numpy-2.4.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:3213d622a0283a39a93d188f3cf72b26862df52fbb4ca3697f51705016523d41", size = 5204410, upload-time = "2026-05-18T23:34:00.302Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d4/9770d14ba719432bb90a421bfd443872ed0f70f7264b64bec12ea363d5fd/numpy-2.4.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:357cc07a6d7b0b182ff02249616a03742827ebb1277546b5c7cd7f7620a45698", size = 6551240, upload-time = "2026-05-18T23:34:02.852Z" }, + { url = "https://files.pythonhosted.org/packages/c9/c6/50a46a6205feba2343f1d6d17438107c5dc491ed1c736e6ea68689fd906b/numpy-2.4.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f9fb9157b4ce2971008323afe46053787b526ef624fea915b261468a8421a0f", size = 15671012, upload-time = "2026-05-18T23:34:05.485Z" }, + { url = "https://files.pythonhosted.org/packages/99/60/14115e6364fa676c5397c2ad3004e527e9aa487abf5d0706ec81bbd08529/numpy-2.4.6-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90f9849678c75fe7afa2d348ac842c168b0a4d3d61919687216dfc547976d853", size = 16645538, upload-time = "2026-05-18T23:34:09.265Z" }, + { url = "https://files.pythonhosted.org/packages/ae/c5/693cbe59e57db94d2231fa519ca3978dc9e19da5a8f088588f5c6e947ff2/numpy-2.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c1a2af6c6ef86344a6b0db6b97834208bf598db514f2b155042439b62605601a", size = 17020706, upload-time = "2026-05-18T23:34:13.053Z" }, + { url = "https://files.pythonhosted.org/packages/ef/fc/85b7c4eff9b4966ade25c2273cf7e7012e92366c032058653934b37de044/numpy-2.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e5805d5a22fd19c8ccff10a9561f9df94436b0545619ea579db2d3c35294bce2", size = 18368541, upload-time = "2026-05-18T23:34:17.024Z" }, + { url = "https://files.pythonhosted.org/packages/f6/81/e1b27545deedce7f4a0b348618c6b62d74e36a4dc9ccd42f3eb2f85eee32/numpy-2.4.6-cp312-cp312-win32.whl", hash = "sha256:e3eeb0aabd6bd5ce64faae67e9935203a6991b4bc2a485a767fbafb2c5125f45", size = 5962825, upload-time = "2026-05-18T23:34:20.3Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ca/feab00bd44aa5fe1ad2c18f08b4d3bb92e26484b0b1d1443897809ed528c/numpy-2.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:d8e8286dd7cea7895157318d1b91cdacac64c479f3cbc8dce548331728484751", size = 12321687, upload-time = "2026-05-18T23:34:23.095Z" }, + { url = "https://files.pythonhosted.org/packages/63/cf/5a6d34850a39d1093558564f77ee8e8e0bee5061151b8f05a55711001ec7/numpy-2.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:4081eb135ac24158bd51cdfbef16f1c64df7063b1143f24731387137c092bec8", size = 10221482, upload-time = "2026-05-18T23:34:25.876Z" }, + { url = "https://files.pythonhosted.org/packages/fb/82/bdab26d7438c6791ca31b7c024ca37c1eab8b726ba236129005cd4a06e45/numpy-2.4.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:511dbaf848decaaaf4b4ca48032619fb3138710c4bf7da7617765edad1ef96b0", size = 16684648, upload-time = "2026-05-18T23:34:29.41Z" }, + { url = "https://files.pythonhosted.org/packages/1b/30/a80189bcc7f5e4258b3fbc3968d909d1756f54d023299ecc39ad6fdb9ef8/numpy-2.4.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bf162abab1c1a736333192707cef898e735a5ca00f38f27eeedf44b39d9e85eb", size = 14693902, upload-time = "2026-05-18T23:34:33.013Z" }, + { url = "https://files.pythonhosted.org/packages/97/12/70b5d0d7c15e1ebb8a6a84a8caa1d19e181d84fb58bb6d70aca29099dec1/numpy-2.4.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:043191bfa8eab18c776647b62723ac9dddece59743b13f49b2016094129c2b3f", size = 5198992, upload-time = "2026-05-18T23:34:36.132Z" }, + { url = "https://files.pythonhosted.org/packages/ba/8c/ebd2a8f8a83541f8d38cc5667e8c2b69cecfd30da6e45693e8158857d44b/numpy-2.4.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:6180d8b35af935aed8ece3a85e0a43f87393ae0ac87c8d2c8bd2c993f7270ef3", size = 6546944, upload-time = "2026-05-18T23:34:38.484Z" }, + { url = "https://files.pythonhosted.org/packages/bb/c5/7b863a97a91671a0338f4253bd3b5a3d3852f0692dae91711c9f4a10e787/numpy-2.4.6-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72fbe16c6fac95aedf5937fa873445cec2110be35d8a4e9433d7501fd98dae6b", size = 15669392, upload-time = "2026-05-18T23:34:41.257Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9d/3584b9984ca4c047aea75214ce1a4c4c73d849bd71b604264b7f5653f8a8/numpy-2.4.6-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a7830bab239b79cda9c08c2da014761cafb48da6150e1da17ac06283f43b6089", size = 16633220, upload-time = "2026-05-18T23:34:45.075Z" }, + { url = "https://files.pythonhosted.org/packages/05/ae/7c67fba23bd98caec7c99261f3a16072ade14813486b0282cb29846de832/numpy-2.4.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ef4aea96ce4d3b074422cb4f2f64e216bf9e213004bb58ecfdf50ea02ea8eb9a", size = 17020800, upload-time = "2026-05-18T23:34:49.065Z" }, + { url = "https://files.pythonhosted.org/packages/d9/5d/3b6725cb31d983c5e66916f5d36f6d7e5521129e4c4404d64f918292a5b6/numpy-2.4.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dfa20cc6ca228e6b155b11da03825975ce66aea520985dbbddf0f2a5a495c605", size = 18357600, upload-time = "2026-05-18T23:34:52.709Z" }, + { url = "https://files.pythonhosted.org/packages/f7/da/2ccc6c2fe8898dee01d90c75c5f5f914a23daf99e3e0f59516a08760c8b5/numpy-2.4.6-cp313-cp313-win32.whl", hash = "sha256:56b39e5e0622a09a25bf5baf62f4bcf0cb8a41ae6e2819cf49bbc5a74c083f91", size = 5961134, upload-time = "2026-05-18T23:34:55.618Z" }, + { url = "https://files.pythonhosted.org/packages/b5/cd/9cc4dc876fb065d5c220aae4d5e14826b2715331bb7618ce1fb07a679d99/numpy-2.4.6-cp313-cp313-win_amd64.whl", hash = "sha256:c4fc99836233ea196540b17ab0983aff60ed07941751930f5f4d05bc3b3b7359", size = 12318598, upload-time = "2026-05-18T23:34:58.928Z" }, + { url = "https://files.pythonhosted.org/packages/39/1e/c0bcba1f8694116485fe28fd1be698c278fcda4141c5b0e53a2aed8b12a8/numpy-2.4.6-cp313-cp313-win_arm64.whl", hash = "sha256:a7c711e21628b52034bb5ab8d1bce291f752fcc5e92accc615778acee1ff4778", size = 10222272, upload-time = "2026-05-18T23:35:02.167Z" }, + { url = "https://files.pythonhosted.org/packages/63/6d/cc5619247c8f4204e507f5883528372e4ac4bb189e579fb859a12e480b1f/numpy-2.4.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:112b06a867b235ef466ed3508ddf0238050df9c727cafb5301ac385b899189a1", size = 14821197, upload-time = "2026-05-18T23:35:05.468Z" }, + { url = "https://files.pythonhosted.org/packages/00/58/f1c39161c87d9e9bed660f1ed4bafc0e403d5ec9650b6dd77aead07d489b/numpy-2.4.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:eaf7fa2de5c0be8ae6ff8e9bea2ccd725e980541244521d8d4b5f3354a27babe", size = 5326287, upload-time = "2026-05-18T23:35:08.693Z" }, + { url = "https://files.pythonhosted.org/packages/af/57/3917ab0fd97f271a8694513581b8a36c655f111c446852c302f04ccdb6fc/numpy-2.4.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:7265a2f3d436e54ef9f2b52b5c937e6be778781bd97a590319d7348f1c1ca997", size = 6646763, upload-time = "2026-05-18T23:35:11.459Z" }, + { url = "https://files.pythonhosted.org/packages/eb/0f/037e64c494b67581ae18193d770adef354c41f3f2c8ebf865602d949bf8f/numpy-2.4.6-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20", size = 15728070, upload-time = "2026-05-18T23:35:14.79Z" }, + { url = "https://files.pythonhosted.org/packages/21/a6/5d2bae9c9542eb4df16dc9c46dc79c186e9bad53805dfa5399a6023c6db0/numpy-2.4.6-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ede83e07a75dd06bc501566c1eca2afc0d61677c1472ac9ad93fdee6e638a48d", size = 16681752, upload-time = "2026-05-18T23:35:18.836Z" }, + { url = "https://files.pythonhosted.org/packages/92/14/23d1dfb410ae362cd59ce53e936b1513d545eb40db3949ced632e19a459e/numpy-2.4.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:68bb27509ac1b9a3443094260f6326150663b06abe40b73a2f81160623da5b67", size = 17086024, upload-time = "2026-05-18T23:35:22.52Z" }, + { url = "https://files.pythonhosted.org/packages/4b/6e/23595a2c642cdf3bc567877064bdd7f91c8b0038a4453cf2daf7248eafe9/numpy-2.4.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:a0df0043bdb289bde1f62da130d20df23d58b45429f752bc7a8fc5325a225ecd", size = 18403398, upload-time = "2026-05-18T23:35:26.398Z" }, + { url = "https://files.pythonhosted.org/packages/8a/90/0ac3bc947217e66dec77e7cbc6a1979d1af70b6461b82f620d3bccd5e4c8/numpy-2.4.6-cp313-cp313t-win32.whl", hash = "sha256:29a287e0cf63ff528da061de6b9f64a4618da591ca1046aafc54062e40ca7eab", size = 6084971, upload-time = "2026-05-18T23:35:29.387Z" }, + { url = "https://files.pythonhosted.org/packages/77/71/5673e351671a1d2bd6063b91b44f70c0affea7d1516fa7a6572941ba4aa1/numpy-2.4.6-cp313-cp313t-win_amd64.whl", hash = "sha256:25c692919ac5a01f170a3bfcd62d745b24fd095c353d50812637d6fcab442e75", size = 12458532, upload-time = "2026-05-18T23:35:32.175Z" }, + { url = "https://files.pythonhosted.org/packages/3f/88/19d3503c5046e688f049274b27a3ef3d771152fa80d3ba3d01a3dff61abe/numpy-2.4.6-cp313-cp313t-win_arm64.whl", hash = "sha256:1e978ec1e8bd0e0e4de6bb75de9d30cbb74db6b6a2bb727618613703ca0167dd", size = 10291881, upload-time = "2026-05-18T23:35:35.465Z" }, + { url = "https://files.pythonhosted.org/packages/f8/91/3ab2044d05fd16d343c5ac2e69b127f1b2854040dd20b193257c78028bd3/numpy-2.4.6-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:06ca2f61ec4385a07a6977c55ba998a4466c123642b4a32694d3128fce18c079", size = 16683458, upload-time = "2026-05-18T23:35:38.353Z" }, + { url = "https://files.pythonhosted.org/packages/8e/62/764ce66fa4147ae6d73071a3abf804ffe606f174618697c571acdf26a7c9/numpy-2.4.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:38efbc8de75c7a0fc1ac190162d892787f3f47b57cc291231aafee36b80982b7", size = 14704559, upload-time = "2026-05-18T23:35:42.14Z" }, + { url = "https://files.pythonhosted.org/packages/60/61/23f27c172f022e04025b7dc2367f4d63c1a398120607ec896228649a6f48/numpy-2.4.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d581b735e177fdcdce6fed8e7e8880a3fb6ee4e3653a3ac6af01c6f4c03effc5", size = 5209716, upload-time = "2026-05-18T23:35:45.377Z" }, + { url = "https://files.pythonhosted.org/packages/03/71/21cf70dc6ea3e3acb95fc53a265b2fc248b981f0194ceb5b475271b8809d/numpy-2.4.6-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:0a041d3d761dc3c35cc56ce0351506a02bcbc25f7b169f652435141a17db9096", size = 6543947, upload-time = "2026-05-18T23:35:47.926Z" }, + { url = "https://files.pythonhosted.org/packages/d5/91/64288395ee1799bd2e0b04a305dce9666da90c961e1f3fe982a05ee1c036/numpy-2.4.6-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40fdc1ae7125e518ea98e53e69a4ebc27e1fd50510c47b7ea130cf21e5e1d42b", size = 15685197, upload-time = "2026-05-18T23:35:50.863Z" }, + { url = "https://files.pythonhosted.org/packages/f3/eb/ebffaa97dc55502df69584a8f0dcf07f69a3e0b3e2323670a2722db9aa39/numpy-2.4.6-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a2c306dea656c12c68f51f4cea133cbe78ca7435eb28c735eac1d3ebe73be6e8", size = 16638245, upload-time = "2026-05-18T23:35:54.752Z" }, + { url = "https://files.pythonhosted.org/packages/b8/0b/54f9da33128d7e350fab89c7455902eeae70349ee52bddb448dc4a576f45/numpy-2.4.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:33111801a01c12a8a1e3721f0a9232f8cfc8ae2c6b7098167e6f623c6073f402", size = 17036587, upload-time = "2026-05-18T23:35:58.355Z" }, + { url = "https://files.pythonhosted.org/packages/b6/f0/fdebc1052db1cc37c64beb22072d67cd6d1c71adca1299f53dec2b5e20d3/numpy-2.4.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ae506e6902902557576a26ff33eda8695e7ecb3cb36c3b573a0765dee114ebdb", size = 18363226, upload-time = "2026-05-18T23:36:02.845Z" }, + { url = "https://files.pythonhosted.org/packages/aa/b4/298628d98c72b57e57f7165ae6a481a1deaf6f3c28262a6e4c739c275930/numpy-2.4.6-cp314-cp314-win32.whl", hash = "sha256:aaf159caa35993cb1f56fb9b8e4610d35758e7ca005412eb1daa856a78c9c4b1", size = 6010196, upload-time = "2026-05-18T23:36:05.92Z" }, + { url = "https://files.pythonhosted.org/packages/df/ac/46de6dda46478f7942f839e094970be2d4a861e005c4b3bf07c92e291a09/numpy-2.4.6-cp314-cp314-win_amd64.whl", hash = "sha256:b507f5c4c1d508876d1819b6bf9a49d365b96320b5d4993426b33a23ca4b8261", size = 12450334, upload-time = "2026-05-18T23:36:09.107Z" }, + { url = "https://files.pythonhosted.org/packages/78/92/b8b798ac784102c0da830d2257d59358e3d3d90d1e2b3f2575dad976c5cf/numpy-2.4.6-cp314-cp314-win_arm64.whl", hash = "sha256:6f41ae150c4e32db4f3310cdaf64b1593a03dbabe29eec77fc9b50fe64061df6", size = 10495678, upload-time = "2026-05-18T23:36:12.766Z" }, + { url = "https://files.pythonhosted.org/packages/30/34/ec28d1aa8115971537c01469ab2011ee96827930f0a124de1000cc2a7ed7/numpy-2.4.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ece3d2cfe132e7d51f44a832b303895e6f2d499c5e74dfbdb06ee246147a304a", size = 14823672, upload-time = "2026-05-18T23:36:16.473Z" }, + { url = "https://files.pythonhosted.org/packages/16/bd/f6d1fede4e54e8042a7ff97bb495510f3c220f94bcd9e8b228e87c92cc0d/numpy-2.4.6-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:e3e5193ef5a3dc73bceee50f7fdc2c90dbb76c42df8d8fae3d1067a583df579e", size = 5328731, upload-time = "2026-05-18T23:36:19.767Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f0/e105b9e2fd728a9910103884decd6951d9dd73896b914a98d9a231de02ee/numpy-2.4.6-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:17f9ade344e7d9b464a084d69bcf18fc691cb1db67c62ed80820bf4926d78f0e", size = 6649805, upload-time = "2026-05-18T23:36:22.266Z" }, + { url = "https://files.pythonhosted.org/packages/82/dd/1206a7ca6ab15e3f02069707ca96222e202af681bb73756da7527f3cb837/numpy-2.4.6-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9cd5ffd25db4e7ba6a375693b3fc0fc1791ec636c17db3720da19bde7180ec43", size = 15730496, upload-time = "2026-05-18T23:36:25.713Z" }, + { url = "https://files.pythonhosted.org/packages/51/e7/38d3ea825dcab85a591734decb2f6c67caa7c8367d374df1a1c3842f9b07/numpy-2.4.6-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7d92c3819208a60205a12a245c91ad70cb0a85336659b19b834205573ac8456e", size = 16679616, upload-time = "2026-05-18T23:36:29.652Z" }, + { url = "https://files.pythonhosted.org/packages/93/b7/caabfdf53edf663e0b4eb74d7d405d83baef09eb5e83bcd32d601d72b93e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e85b752a1e912b70eaad4fafbd4d1238007ab221de2009b9a2f5ae7461239895", size = 17085145, upload-time = "2026-05-18T23:36:33.449Z" }, + { url = "https://files.pythonhosted.org/packages/f9/45/68d7c33a6bcf3e5aa3bdbd57a367e6f615286dfd6482f97e8ffeb734306e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:29cb7f67d10b479ff07c17d33e39f78c07f71c40ef30d63c153d340e96cd3fb4", size = 18403813, upload-time = "2026-05-18T23:36:37.369Z" }, + { url = "https://files.pythonhosted.org/packages/9c/50/0753655aa844c99cd9e018aacf76f130f1bd81d881bb74bc0aef5d73a8ba/numpy-2.4.6-cp314-cp314t-win32.whl", hash = "sha256:260a5d70215b61ab4fadf5c7baacd64821842975eea312125ed3c39a6391b063", size = 6156982, upload-time = "2026-05-18T23:36:40.817Z" }, + { url = "https://files.pythonhosted.org/packages/b2/d4/7c67becf668f973cb490cec3e98dfd799d866f9c989a54d355672cfa0db6/numpy-2.4.6-cp314-cp314t-win_amd64.whl", hash = "sha256:81a1cca95ed5bb92aa8b10dd2cdc9a0d3853a50fad926c28b5d7e8ea54389627", size = 12638908, upload-time = "2026-05-18T23:36:43.996Z" }, + { url = "https://files.pythonhosted.org/packages/43/bb/e1c71a4295b1b1d1393d50dbb4f2a36283c6859d9d3892e84f00ec5a91d5/numpy-2.4.6-cp314-cp314t-win_arm64.whl", hash = "sha256:0c9136e14ed34a9e343a31c533d78a9813a69a3148332bce5e9821cb2f996e66", size = 10565867, upload-time = "2026-05-18T23:36:47.114Z" }, + { url = "https://files.pythonhosted.org/packages/de/12/b422cc84439adc0d00de605bf4a308890ae5c26f2c71fbd73e5d08fbb0dd/numpy-2.4.6-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:55cced7c52e981362f708ad635198e97a752dfba412cc03c23bbf3bd8d5cd662", size = 16847511, upload-time = "2026-05-18T23:36:50.673Z" }, + { url = "https://files.pythonhosted.org/packages/44/53/f481bef68011740f8849418d82db07230e825013f31f4eef5ba5b805316a/numpy-2.4.6-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:d6da64deb6b8ed903e7560180a92f2d804ee1ba5eeb849ac2748b8c1aba1f6d7", size = 14889064, upload-time = "2026-05-18T23:36:53.879Z" }, + { url = "https://files.pythonhosted.org/packages/7f/57/42ed575c10ced8af951d426bc4e1f8aff16fd851db33f067036215a7f860/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_arm64.whl", hash = "sha256:68a5124b13fa6cc2086764a20005d30bc0548146f7f5322f02fce212ca14317f", size = 5394157, upload-time = "2026-05-18T23:36:57.194Z" }, + { url = "https://files.pythonhosted.org/packages/6a/ef/f66cc724fcc36c1e364c67f51ae9146090b8b584f27d58b97fdae3edd737/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_x86_64.whl", hash = "sha256:948424b06129ce883307e8cff868c31396d8dc7630a59c61d70d98dbe70f222c", size = 6708728, upload-time = "2026-05-18T23:36:59.575Z" }, + { url = "https://files.pythonhosted.org/packages/1a/9c/c531f2293b91265d8b48e9b329f54fdd7ffae73cb4134ea10cca4237e9cc/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5dbbdb29840ca3d91ee0fece42fc29278886d908280bfec0a5846c6f901a3eb0", size = 15798374, upload-time = "2026-05-18T23:37:02.674Z" }, + { url = "https://files.pythonhosted.org/packages/1a/b0/413077f6b1153ed3cba361401c6783bbad6114804a000cc22eb71c13e190/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ad03c0965fb3c692200e74d458ca28c1dbb4ce96f9a479a8aa041ad5fabca02", size = 16747286, upload-time = "2026-05-18T23:37:06.327Z" }, + { url = "https://files.pythonhosted.org/packages/15/ce/e5ec180bc41812edcd8daeb8639d205622c0e8c02259d8ab25a0201b3c2a/numpy-2.4.6-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:2803abfebfc990042cd494d8ce2d5f82e9d847af6d35ec486923aa19dbad5e73", size = 12504263, upload-time = "2026-05-18T23:37:09.715Z" }, +] + [[package]] name = "nvidia-cublas" version = "13.1.1.3" @@ -564,6 +780,11 @@ dependencies = [ { name = "torch" }, ] +[package.optional-dependencies] +numba = [ + { name = "numba" }, +] + [package.dev-dependencies] dev = [ { name = "coverage" }, @@ -572,7 +793,11 @@ dev = [ ] [package.metadata] -requires-dist = [{ name = "torch" }] +requires-dist = [ + { name = "numba", marker = "extra == 'numba'", specifier = ">=0.66" }, + { name = "torch" }, +] +provides-extras = ["numba"] [package.metadata.requires-dev] dev = [ From bcd61aeb9637d3164902e1e8cf4db01ac7938b45 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:03:24 +0800 Subject: [PATCH 05/29] Use a Link-Cut Tree for online ROSA --- src/rosa/_stateful_numba.py | 634 ++++++++++++++++++++++++++++++++++++ 1 file changed, 634 insertions(+) create mode 100644 src/rosa/_stateful_numba.py diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py new file mode 100644 index 0000000..14dbddd --- /dev/null +++ b/src/rosa/_stateful_numba.py @@ -0,0 +1,634 @@ +"""Stateful exact ROSA inference using a suffix automaton and Link-Cut Tree.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch +from numba import njit +from torch import Tensor + + +@njit(cache=True, nogil=True, inline="always") +def _find_transition( # pragma: no cover - executed as compiled Numba code + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + state: int, + token: int, +) -> int: + edge = head[state] + while edge != -1: + if edge_token[edge] == token: + return int(edge_target[edge]) + edge = edge_next[edge] + return -1 + + +@njit(cache=True, nogil=True, inline="always") +def _add_transition( # pragma: no cover - executed as compiled Numba code + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + edge_count: int, + state: int, + token: int, + target: int, +) -> int: + if edge_count >= edge_token.shape[0]: + raise RuntimeError("suffix automaton transition capacity exceeded") + edge_token[edge_count] = token + edge_target[edge_count] = target + edge_next[edge_count] = head[state] + head[state] = edge_count + return edge_count + 1 + + +@njit(cache=True, nogil=True, inline="always") +def _replace_transition( # pragma: no cover - executed as compiled Numba code + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + state: int, + token: int, + target: int, +) -> None: + edge = head[state] + while edge != -1: + if edge_token[edge] == token: + edge_target[edge] = target + return + edge = edge_next[edge] + raise RuntimeError("suffix automaton transition not found") + + +@njit(cache=True, nogil=True, inline="always") +def _lct_is_aux_root( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + node: int, +) -> bool: + p = parent[node] + return p == -1 or (left[p] != node and right[p] != node) + + +@njit(cache=True, nogil=True, inline="always") +def _lct_apply( # pragma: no cover - executed as compiled Numba code + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, + assigned: int, +) -> None: + if node != -1: + value[node] = assigned + lazy[node] = assigned + lazy_valid[node] = 1 + + +@njit(cache=True, nogil=True, inline="always") +def _lct_push( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, +) -> None: + if lazy_valid[node] != 0: + assigned = lazy[node] + _lct_apply(value, lazy, lazy_valid, left[node], assigned) + _lct_apply(value, lazy, lazy_valid, right[node], assigned) + lazy_valid[node] = 0 + + +@njit(cache=True, nogil=True, inline="always") +def _lct_rotate( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + node: int, +) -> None: + p = parent[node] + g = parent[p] + if left[p] == node: + middle = right[node] + right[node] = p + left[p] = middle + else: + middle = left[node] + left[node] = p + right[p] = middle + if middle != -1: + parent[middle] = p + parent[p] = node + parent[node] = g + if g != -1: + if left[g] == p: + left[g] = node + elif right[g] == p: + right[g] = node + + +@njit(cache=True, nogil=True) +def _lct_splay( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, + stack: np.ndarray, +) -> None: + depth = 0 + ancestor = node + stack[depth] = ancestor + depth += 1 + while not _lct_is_aux_root(left, right, parent, ancestor): + ancestor = parent[ancestor] + stack[depth] = ancestor + depth += 1 + while depth > 0: + depth -= 1 + _lct_push(left, right, value, lazy, lazy_valid, stack[depth]) + + while not _lct_is_aux_root(left, right, parent, node): + p = parent[node] + if not _lct_is_aux_root(left, right, parent, p): + g = parent[p] + if (left[p] == node) == (left[g] == p): + _lct_rotate(left, right, parent, p) + else: + _lct_rotate(left, right, parent, node) + _lct_rotate(left, right, parent, node) + + +@njit(cache=True, nogil=True) +def _lct_access( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, + stack: np.ndarray, +) -> None: + last = -1 + current = node + while current != -1: + _lct_splay(left, right, parent, value, lazy, lazy_valid, current, stack) + right[current] = last + if last != -1: + parent[last] = current + last = current + current = parent[current] + _lct_splay(left, right, parent, value, lazy, lazy_valid, node, stack) + + +@njit(cache=True, nogil=True) +def _lct_point_query( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, + stack: np.ndarray, +) -> int: + _lct_access(left, right, parent, value, lazy, lazy_valid, node, stack) + return int(value[node]) + + +@njit(cache=True, nogil=True) +def _lct_path_assign( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, + assigned: int, + stack: np.ndarray, +) -> None: + _lct_access(left, right, parent, value, lazy, lazy_valid, node, stack) + _lct_apply(value, lazy, lazy_valid, node, assigned) + + +@njit(cache=True, nogil=True) +def _lct_cut_parent( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, + stack: np.ndarray, +) -> None: + _lct_access(left, right, parent, value, lazy, lazy_valid, node, stack) + ancestors = left[node] + left[node] = -1 + if ancestors != -1: + parent[ancestors] = -1 + + +@njit(cache=True, nogil=True) +def _lct_link_parent( # pragma: no cover - executed as compiled Numba code + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + value: np.ndarray, + lazy: np.ndarray, + lazy_valid: np.ndarray, + node: int, + represented_parent: int, + stack: np.ndarray, +) -> None: + _lct_access(left, right, parent, value, lazy, lazy_valid, node, stack) + parent[node] = represented_parent + + +@njit(cache=True, nogil=True) +def _step_row( # pragma: no cover - executed as compiled Numba code + token: int, + position: int, + history: np.ndarray, + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + suffix_link: np.ndarray, + length: np.ndarray, + lct_left: np.ndarray, + lct_right: np.ndarray, + lct_parent: np.ndarray, + lct_value: np.ndarray, + lct_lazy: np.ndarray, + lct_lazy_valid: np.ndarray, + lct_stack: np.ndarray, + last: int, + size: int, + edge_count: int, +) -> tuple[int, int, int, int]: + history[position] = token + if size >= head.shape[0]: + raise RuntimeError("suffix automaton state capacity exceeded") + current = size + size += 1 + length[current] = length[last] + 1 + state = last + + while ( + state != -1 + and _find_transition(head, edge_token, edge_target, edge_next, state, token) + == -1 + ): + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + edge_count, + state, + token, + current, + ) + state = suffix_link[state] + + if state == -1: + suffix_link[current] = 0 + _lct_link_parent( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + current, + 0, + lct_stack, + ) + else: + target = _find_transition( + head, edge_token, edge_target, edge_next, state, token + ) + if length[state] + 1 == length[target]: + suffix_link[current] = target + _lct_link_parent( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + current, + target, + lct_stack, + ) + else: + if size >= head.shape[0]: + raise RuntimeError("suffix automaton state capacity exceeded") + clone = size + size += 1 + length[clone] = length[state] + 1 + old_parent = suffix_link[target] + suffix_link[clone] = old_parent + clone_value = _lct_point_query( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + target, + lct_stack, + ) + lct_value[clone] = clone_value + edge = head[target] + while edge != -1: + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + edge_count, + clone, + edge_token[edge], + edge_target[edge], + ) + edge = edge_next[edge] + while ( + state != -1 + and _find_transition( + head, edge_token, edge_target, edge_next, state, token + ) + == target + ): + _replace_transition( + head, + edge_token, + edge_target, + edge_next, + state, + token, + clone, + ) + state = suffix_link[state] + + _lct_link_parent( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + clone, + old_parent, + lct_stack, + ) + _lct_cut_parent( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + target, + lct_stack, + ) + suffix_link[target] = clone + _lct_link_parent( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + target, + clone, + lct_stack, + ) + suffix_link[current] = clone + _lct_link_parent( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + current, + clone, + lct_stack, + ) + + last = current + matched = suffix_link[current] + source = -1 + if matched != 0: + source = _lct_point_query( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + matched, + lct_stack, + ) + prediction = -1 + if source >= 0: + prediction = int(history[source + 1]) + _lct_path_assign( + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + current, + position, + lct_stack, + ) + return prediction, last, size, edge_count + + +@njit(cache=True, nogil=True) +def _step_batch_kernel( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, + position: int, + history: np.ndarray, + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + suffix_link: np.ndarray, + length: np.ndarray, + lct_left: np.ndarray, + lct_right: np.ndarray, + lct_parent: np.ndarray, + lct_value: np.ndarray, + lct_lazy: np.ndarray, + lct_lazy_valid: np.ndarray, + lct_stack: np.ndarray, + last: np.ndarray, + size: np.ndarray, + edge_count: np.ndarray, +) -> np.ndarray: + output = np.empty(tokens.shape[0], dtype=np.int64) + for batch_index in range(tokens.shape[0]): + prediction, new_last, new_size, new_edge_count = _step_row( + int(tokens[batch_index]), + position, + history[batch_index], + head[batch_index], + edge_token[batch_index], + edge_target[batch_index], + edge_next[batch_index], + suffix_link[batch_index], + length[batch_index], + lct_left[batch_index], + lct_right[batch_index], + lct_parent[batch_index], + lct_value[batch_index], + lct_lazy[batch_index], + lct_lazy_valid[batch_index], + lct_stack[batch_index], + int(last[batch_index]), + int(size[batch_index]), + int(edge_count[batch_index]), + ) + output[batch_index] = prediction + last[batch_index] = new_last + size[batch_index] = new_size + edge_count[batch_index] = new_edge_count + return output + + +@dataclass +class _StatefulInferenceState: + """Fixed-capacity, independently batched exact ROSA inference state.""" + + batch_size: int + max_length: int + position: int + history: np.ndarray + head: np.ndarray + edge_token: np.ndarray + edge_target: np.ndarray + edge_next: np.ndarray + suffix_link: np.ndarray + length: np.ndarray + lct_left: np.ndarray + lct_right: np.ndarray + lct_parent: np.ndarray + lct_value: np.ndarray + lct_lazy: np.ndarray + lct_lazy_valid: np.ndarray + lct_stack: np.ndarray + last: np.ndarray + size: np.ndarray + edge_count: np.ndarray + + +def _init_inference_state( + batch_size: int, + max_length: int, +) -> _StatefulInferenceState: + """Allocate a fixed-capacity CPU state suitable for repeated steps.""" + + if batch_size <= 0: + raise ValueError("batch_size must be > 0") + if max_length <= 0: + raise ValueError("max_length must be > 0") + max_states = 2 * max_length + 1 + max_edges = 4 * max_length + 1 + state_shape = (batch_size, max_states) + edge_shape = (batch_size, max_edges) + suffix_link = np.full(state_shape, -1, dtype=np.int32) + return _StatefulInferenceState( + batch_size=batch_size, + max_length=max_length, + position=0, + history=np.empty((batch_size, max_length), dtype=np.int64), + head=np.full(state_shape, -1, dtype=np.int32), + edge_token=np.empty(edge_shape, dtype=np.int64), + edge_target=np.empty(edge_shape, dtype=np.int32), + edge_next=np.empty(edge_shape, dtype=np.int32), + suffix_link=suffix_link, + length=np.zeros(state_shape, dtype=np.int32), + lct_left=np.full(state_shape, -1, dtype=np.int32), + lct_right=np.full(state_shape, -1, dtype=np.int32), + lct_parent=np.full(state_shape, -1, dtype=np.int32), + lct_value=np.full(state_shape, -1, dtype=np.int64), + lct_lazy=np.empty(state_shape, dtype=np.int64), + lct_lazy_valid=np.zeros(state_shape, dtype=np.uint8), + lct_stack=np.empty(state_shape, dtype=np.int32), + last=np.zeros(batch_size, dtype=np.int32), + size=np.ones(batch_size, dtype=np.int32), + edge_count=np.zeros(batch_size, dtype=np.int32), + ) + + +def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: + """Consume one token per batch row and return exact top-1 predictions.""" + + if tokens.ndim == 0 and state.batch_size == 1: + tokens = tokens.unsqueeze(0) + if tokens.ndim != 1 or tokens.shape[0] != state.batch_size: + raise ValueError("tokens must have shape [batch_size]") + if state.position >= state.max_length: + raise RuntimeError("inference state capacity exceeded") + device = tokens.device + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + output = _step_batch_kernel( + cpu_tokens.numpy(), + state.position, + state.history, + state.head, + state.edge_token, + state.edge_target, + state.edge_next, + state.suffix_link, + state.length, + state.lct_left, + state.lct_right, + state.lct_parent, + state.lct_value, + state.lct_lazy, + state.lct_lazy_valid, + state.lct_stack, + state.last, + state.size, + state.edge_count, + ) + state.position += 1 + return torch.from_numpy(output).to(device) + + +def predict_exact_stateful(tokens: Tensor) -> Tensor: + """Return exact top-1 ROSA predictions through the stateful Numba backend.""" + + squeeze = tokens.ndim == 1 + if squeeze: + tokens = tokens.unsqueeze(0) + if tokens.ndim != 2: + raise ValueError("tokens must have shape [N] or [B, N]") + device = tokens.device + batch_size, length = tokens.shape + if length == 0: + output = torch.empty(tokens.shape, dtype=torch.long, device=device) + return output[0] if squeeze else output + state = _init_inference_state(batch_size, length) + output = torch.empty(tokens.shape, dtype=torch.long, device=device) + for position in range(length): + output[:, position] = _forward_step(state, tokens[:, position]) + return output[0] if squeeze else output From f5c2628be939dedf71a4a7e2a81aff67f27a892a Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:04:46 +0800 Subject: [PATCH 06/29] Fuse stateful ROSA prefill in Numba --- src/rosa/_stateful_numba.py | 72 +++++++++++++++++++++++++++++++++++-- 1 file changed, 69 insertions(+), 3 deletions(-) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 14dbddd..d979052 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -513,6 +513,53 @@ def _step_batch_kernel( # pragma: no cover - executed as compiled Numba code return output +@njit(cache=True, nogil=True) +def _replay_kernel( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, + history: np.ndarray, + head: np.ndarray, + edge_token: np.ndarray, + edge_target: np.ndarray, + edge_next: np.ndarray, + suffix_link: np.ndarray, + length: np.ndarray, + lct_left: np.ndarray, + lct_right: np.ndarray, + lct_parent: np.ndarray, + lct_value: np.ndarray, + lct_lazy: np.ndarray, + lct_lazy_valid: np.ndarray, + lct_stack: np.ndarray, + last: np.ndarray, + size: np.ndarray, + edge_count: np.ndarray, +) -> np.ndarray: + output = np.empty(tokens.shape, dtype=np.int64) + for position in range(tokens.shape[1]): + output[:, position] = _step_batch_kernel( + tokens[:, position], + position, + history, + head, + edge_token, + edge_target, + edge_next, + suffix_link, + length, + lct_left, + lct_right, + lct_parent, + lct_value, + lct_lazy, + lct_lazy_valid, + lct_stack, + last, + size, + edge_count, + ) + return output + + @dataclass class _StatefulInferenceState: """Fixed-capacity, independently batched exact ROSA inference state.""" @@ -628,7 +675,26 @@ def predict_exact_stateful(tokens: Tensor) -> Tensor: output = torch.empty(tokens.shape, dtype=torch.long, device=device) return output[0] if squeeze else output state = _init_inference_state(batch_size, length) - output = torch.empty(tokens.shape, dtype=torch.long, device=device) - for position in range(length): - output[:, position] = _forward_step(state, tokens[:, position]) + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + output_array = _replay_kernel( + cpu_tokens.numpy(), + state.history, + state.head, + state.edge_token, + state.edge_target, + state.edge_next, + state.suffix_link, + state.length, + state.lct_left, + state.lct_right, + state.lct_parent, + state.lct_value, + state.lct_lazy, + state.lct_lazy_valid, + state.lct_stack, + state.last, + state.size, + state.edge_count, + ) + output = torch.from_numpy(output_array).to(device) return output[0] if squeeze else output From d6ddc692e1a12574b73636f3b66e093f16700698 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:20:46 +0800 Subject: [PATCH 07/29] Expose production stateful ROSA inference --- .github/workflows/ci.yml | 30 +++++++ README.md | 59 +++++++++++-- src/rosa/__init__.py | 167 +++++++++++++++++++++++++++++++++++- src/rosa/_stateful_numba.py | 4 +- tests/test_inference.py | 134 +++++++++++++++++++++++++++++ tests/test_numba_backend.py | 27 ++++++ 6 files changed, 408 insertions(+), 13 deletions(-) create mode 100644 tests/test_inference.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a9bcac6..fe3ec43 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,3 +54,33 @@ jobs: - name: Build distributions if: matrix.python-version == '3.10' run: uv build + + base-install: + name: Base install without Numba + runs-on: ubuntu-latest + timeout-minutes: 10 + + steps: + - name: Check out repository + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Install uv and Python + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + version: "0.8.15" + python-version: "3.10" + enable-cache: true + + - name: Install base dependencies only + run: uv sync --locked --all-groups + + - name: Verify automatic Python fallback + run: | + uv run python - <<'PY' + import torch + from rosa import forward_step, init_inference_state + + state = init_inference_state(1, max_length=2) + assert state.backend == "python" + assert forward_step(state, torch.tensor(0)).item() == -1 + PY diff --git a/README.md b/README.md index 9dcbb31..e8a87e8 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,7 @@ The virtual-candidate branch and neural value residual have independent curricul - Python 3.10+ - PyTorch +- Optional Numba backend for production exact inference - `coverage`, Ruff, and Pyright for development Install the published package from PyPI: @@ -53,6 +54,12 @@ Install the published package from PyPI: uv add rosa-torch ``` +Install the stateful Link-Cut Tree backend with: + +```bash +uv add 'rosa-torch[numba]' +``` + Install the package and its locked development dependencies with [uv](https://docs.astral.sh/uv/): ```bash @@ -73,16 +80,49 @@ PEP 517-compatible Python package manager. ├── README.md ├── src │ └── rosa -│ └── __init__.py +│ ├── __init__.py +│ ├── _numba_backend.py +│ └── _stateful_numba.py └── tests ├── __init__.py ├── run_coverage.py + ├── test_inference.py + ├── test_numba_backend.py └── test_rosa.py ``` -The implementation is distributed as an installable `rosa` package while -remaining in one source module to keep the exact suffix-automaton and neural -retrieval paths easy to inspect together. +The implementation is distributed as an installable `rosa` package. The core +neural path remains in `__init__.py`; optional compiled inference kernels are +isolated in private backend modules and loaded lazily. + +## Stateful exact inference + +Use one explicit state per independent decoding stream. The automaton remains +on CPU, while CUDA token inputs receive CUDA predictions through a single +batch transfer per step. + +```python +import torch + +from rosa import forward_step, init_inference_state + +state = init_inference_state( + batch_size=2, + max_length=32_768, + backend="auto", # "numba" when installed, otherwise exact Python +) + +for token in generated_token_ids: # each tensor has shape [2] + predicted_token = forward_step(state, token) + +state.reset() +``` + +Capacity is fixed at initialization for predictable memory use. Exceeding it +raises `RuntimeError` before mutation. States are mutable, isolated, and must +not be shared concurrently between decoding requests. `forward_step` implements +exact top-1 ROSA; rich multi-candidate training remains on the full-sequence +`ROSA` path. ## Quick start @@ -263,11 +303,12 @@ tensors are returned to the original PyTorch device. Accelerator backends such as TileLang or Triton should optimize only the differentiable tensor path around the automaton. -The current implementation performs this CPU work synchronously and rebuilds -the automaton for each full-sequence call. Production autoregressive inference -can improve throughput with a stateful CPU worker whose automaton updates are -pipelined alongside GPU layers, while preserving the same exact candidate -semantics and PyTorch fallback. +The stateful inference API retains the suffix automaton across decoding steps. +Its Numba backend uses a rooted Link-Cut Tree for lazy suffix-path timestamp +updates, replacing the previous quadratic eager propagation with amortized +`O(log N)` updates. Full stateful prefill is fused into one compiled replay +kernel; the explicit Python fallback preserves exact semantics without making +Numba a base dependency. ## Design guarantees diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py index 4635c90..57eb60e 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -10,8 +10,8 @@ import math from collections.abc import Sequence -from dataclasses import dataclass -from typing import cast +from dataclasses import dataclass, field +from typing import Any, Literal, cast import torch import torch.nn.functional as F @@ -25,11 +25,14 @@ "EXACT_KIND", "NULL_KIND", "ROSA", + "ROSAInferenceState", "VIRTUAL_KIND", "HardCandidates", "ROSAOutput", "build_hard_candidates", "build_virtual_pool_indices", + "forward_step", + "init_inference_state", "reference_rosa", ] @@ -246,6 +249,46 @@ def write_current_end(self, end_position: int) -> None: v = self.link[v] +@dataclass +class _PythonInferenceState: + batch_size: int + max_length: int + position: int + automata: list[_OnlineSuffixAutomaton] + history: list[list[int]] + + +def _init_python_inference_state( + batch_size: int, + max_length: int, +) -> _PythonInferenceState: + return _PythonInferenceState( + batch_size=batch_size, + max_length=max_length, + position=0, + automata=[_OnlineSuffixAutomaton(max_length, 1) for _ in range(batch_size)], + history=[[] for _ in range(batch_size)], + ) + + +def _python_forward_step(state: _PythonInferenceState, tokens: Tensor) -> Tensor: + if state.position >= state.max_length: + raise RuntimeError("inference state capacity exceeded") + device = tokens.device + predictions: list[int] = [] + for batch_index, token in enumerate(tokens.detach().cpu().tolist()): + value = int(token) + history = state.history[batch_index] + automaton = state.automata[batch_index] + history.append(value) + automaton.extend(value) + candidates = automaton.read_candidates(1, 1) + predictions.append(history[candidates[0][0] + 1] if candidates else -1) + automaton.write_current_end(state.position) + state.position += 1 + return torch.tensor(predictions, dtype=torch.long, device=device) + + def _build_single_hard_candidates( tokens: list[int], suffix_k: int, @@ -884,3 +927,123 @@ def combine_losses( + balance_weight * aux_losses["code_balance"] + virtual_weight * aux_losses["virtual_usage"] ) + + +InferenceBackend = Literal["auto", "python", "numba"] + + +@dataclass(slots=True) +class ROSAInferenceState: + """Mutable, fixed-capacity state for exact autoregressive ROSA inference. + + Create states with :func:`init_inference_state`. Each state is independent, + CPU-resident, and intended to be mutated by one decoding stream at a time. + """ + + batch_size: int + max_length: int + backend: Literal["python", "numba"] + _impl: object = field(repr=False) + + @property + def position(self) -> int: + """Number of tokens consumed by every batch row.""" + + return int(cast(Any, self._impl).position) + + def reset(self) -> None: + """Reset all batch rows while retaining the configured capacity.""" + + self._impl = _make_inference_impl( + self.batch_size, + self.max_length, + self.backend, + ) + + +def _make_inference_impl( + batch_size: int, + max_length: int, + backend: Literal["python", "numba"], +) -> object: + if backend == "python": + return _init_python_inference_state(batch_size, max_length) + try: + from ._stateful_numba import _init_inference_state + except ModuleNotFoundError as error: + if error.name not in {"numba", "numpy"}: + raise + raise ImportError( + "Numba inference requires `pip install rosa-torch[numba]`" + ) from error + return _init_inference_state(batch_size, max_length) + + +def init_inference_state( + batch_size: int, + max_length: int = 8192, + *, + backend: InferenceBackend = "auto", +) -> ROSAInferenceState: + """Create an explicit state for exact top-1 autoregressive ROSA inference. + + ``backend="auto"`` selects the Link-Cut Tree Numba backend when the + optional dependency is installed and otherwise uses the exact Python + fallback. Capacity is fixed so memory use and failure behavior remain + predictable in production. + """ + + if batch_size <= 0: + raise ValueError("batch_size must be > 0") + if max_length <= 0: + raise ValueError("max_length must be > 0") + if backend not in {"auto", "python", "numba"}: + raise ValueError("backend must be 'auto', 'python', or 'numba'") + + selected: Literal["python", "numba"] + if backend == "auto": + try: + impl = _make_inference_impl(batch_size, max_length, "numba") + selected = "numba" + except ImportError: + impl = _make_inference_impl(batch_size, max_length, "python") + selected = "python" + else: + selected = backend + impl = _make_inference_impl(batch_size, max_length, selected) + return ROSAInferenceState(batch_size, max_length, selected, impl) + + +def forward_step(state: ROSAInferenceState, token: Tensor) -> Tensor: + """Consume one token per batch row and return exact ROSA predictions. + + ``token`` must be an integer tensor shaped ``[batch_size]``. A scalar is + also accepted when ``batch_size == 1`` and produces a scalar prediction. + The state remains on CPU; the prediction is returned on the token device. + """ + + if not isinstance(state, ROSAInferenceState): + raise TypeError("state must be a ROSAInferenceState") + if not isinstance(token, Tensor): + raise TypeError("token must be a torch.Tensor") + squeeze = token.ndim == 0 + if squeeze and state.batch_size == 1: + token = token.unsqueeze(0) + if token.ndim != 1 or token.shape[0] != state.batch_size: + raise ValueError("token must have shape [batch_size]") + if token.dtype not in { + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + }: + raise TypeError("token must use an integer dtype") + + if state.backend == "python": + output = _python_forward_step(cast(_PythonInferenceState, state._impl), token) + else: + from ._stateful_numba import _forward_step + + output = _forward_step(cast(Any, state._impl), token) + return output[0] if squeeze else output diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index d979052..6b2de21 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -362,8 +362,8 @@ def _step_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_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..5567f27 --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import unittest +from itertools import product +from typing import Any, Literal, cast +from unittest.mock import patch + +import torch + +from rosa import ( + ROSAInferenceState, + forward_step, + init_inference_state, + reference_rosa, +) + + +class TestStatefulInference(unittest.TestCase): + def collect( + self, + tokens: torch.Tensor, + *, + backend: str, + ) -> tuple[ROSAInferenceState, torch.Tensor]: + selected = cast(Literal["auto", "python", "numba"], backend) + state = init_inference_state(tokens.shape[0], tokens.shape[1], backend=selected) + output = torch.stack( + [ + forward_step(state, tokens[:, position]) + for position in range(tokens.shape[1]) + ], + dim=1, + ) + return state, output + + def test_python_and_numba_match_reference_step_by_step(self) -> None: + generator = torch.Generator().manual_seed(20260811) + cases = [ + torch.zeros((2, 64), dtype=torch.long), + torch.arange(128).remainder(2).reshape(2, 64), + torch.randint(17, (3, 63), generator=generator), + torch.tensor([[-10_000_000, -1, 2**31, 10**12, -1, 2**31]]), + ] + for backend in ("python", "numba"): + for case_index, tokens in enumerate(cases): + with self.subTest(backend=backend, case=case_index): + state, output = self.collect(tokens, backend=backend) + expected, _, _ = reference_rosa(tokens) + self.assertTrue(torch.equal(output, expected)) + self.assertEqual(state.position, tokens.shape[1]) + + def test_scalar_reset_isolation_and_capacity(self) -> None: + first = init_inference_state(1, 3, backend="numba") + second = init_inference_state(1, 3, backend="numba") + sequence = [0, 1, 0] + first_pass = [ + forward_step(first, torch.tensor(token)).item() for token in sequence + ] + self.assertEqual(second.position, 0) + first.reset() + second_pass = [ + forward_step(first, torch.tensor(token)).item() for token in sequence + ] + self.assertEqual(first_pass, second_pass) + self.assertEqual(first.position, 3) + with self.assertRaisesRegex(RuntimeError, "capacity"): + forward_step(first, torch.tensor(2)) + + python_state = init_inference_state(1, 1, backend="python") + forward_step(python_state, torch.tensor(0)) + with self.assertRaisesRegex(RuntimeError, "capacity"): + forward_step(python_state, torch.tensor(1)) + + def test_exhaustive_binary_clone_sequences(self) -> None: + rows = list(product(range(2), repeat=10)) + tokens = torch.tensor(rows, dtype=torch.long) + _, output = self.collect(tokens, backend="numba") + expected, _, _ = reference_rosa(tokens) + self.assertTrue(torch.equal(output, expected)) + + def test_auto_backend_and_validation(self) -> None: + state = init_inference_state(1, backend="auto") + self.assertIn(state.backend, {"python", "numba"}) + self.assertIn("backend=", repr(state)) + + for batch_size, max_length, message in ( + (0, 1, "batch_size"), + (1, 0, "max_length"), + ): + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + init_inference_state(batch_size, max_length) + with self.assertRaisesRegex(ValueError, "backend"): + init_inference_state(1, backend=cast(Any, "invalid")) + + with self.assertRaisesRegex(TypeError, "state"): + forward_step(cast(Any, object()), torch.tensor([1])) + with self.assertRaisesRegex(TypeError, "Tensor"): + forward_step(state, cast(Any, 1)) + with self.assertRaisesRegex(ValueError, "shape"): + forward_step(state, torch.tensor([1, 2])) + with self.assertRaisesRegex(TypeError, "integer"): + forward_step(state, torch.tensor([1.0])) + + def test_numba_missing_fallback_and_explicit_error(self) -> None: + real_import = __import__ + + def missing_numba(name: str, *args: Any, **kwargs: Any) -> Any: + if name.endswith("_stateful_numba"): + raise ModuleNotFoundError("No module named 'numba'", name="numba") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=missing_numba): + fallback = init_inference_state(1, backend="auto") + self.assertEqual(fallback.backend, "python") + with self.assertRaisesRegex(ImportError, "rosa-torch\\[numba\\]"): + init_inference_state(1, backend="numba") + + def unexpected_missing(name: str, *args: Any, **kwargs: Any) -> Any: + if name.endswith("_stateful_numba"): + raise ModuleNotFoundError("unexpected", name="unexpected") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=unexpected_missing): + with self.assertRaises(ModuleNotFoundError): + init_inference_state(1, backend="numba") + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_cuda_token_returns_cuda_prediction(self) -> None: + tokens = torch.tensor([[0, 1, 0, 2]], device="cuda") + _, output = self.collect(tokens, backend="numba") + expected, _, _ = reference_rosa(tokens) + self.assertEqual(output.device.type, "cuda") + self.assertTrue(torch.equal(output, expected)) diff --git a/tests/test_numba_backend.py b/tests/test_numba_backend.py index 04eafe6..2e036f2 100644 --- a/tests/test_numba_backend.py +++ b/tests/test_numba_backend.py @@ -9,6 +9,11 @@ try: from rosa._numba_backend import predict_exact + from rosa._stateful_numba import ( + _forward_step, + _init_inference_state, + predict_exact_stateful, + ) except ModuleNotFoundError as error: if error.name not in {"numba", "numpy"}: raise @@ -51,6 +56,28 @@ def test_shape_validation(self) -> None: with self.assertRaisesRegex(ValueError, "shape"): predict_exact(torch.zeros(1, 2, 3, dtype=torch.long)) + def test_stateful_private_validation_and_full_replay(self) -> None: + with self.assertRaisesRegex(ValueError, "batch_size"): + _init_inference_state(0, 1) + with self.assertRaisesRegex(ValueError, "max_length"): + _init_inference_state(1, 0) + + state = _init_inference_state(1, 1) + self.assertEqual(_forward_step(state, torch.tensor(0)).shape, (1,)) + with self.assertRaisesRegex(ValueError, "shape"): + _forward_step(state, torch.tensor([1, 2])) + + empty = torch.empty(0, dtype=torch.long) + self.assertTrue(torch.equal(predict_exact_stateful(empty), empty)) + empty_batch = torch.empty((2, 0), dtype=torch.long) + self.assertEqual(tuple(predict_exact_stateful(empty_batch).shape), (2, 0)) + with self.assertRaisesRegex(ValueError, "shape"): + predict_exact_stateful(torch.zeros(1, 2, 3, dtype=torch.long)) + + tokens = torch.tensor([0, 1, 0, 2, 0], dtype=torch.long) + expected, _, _ = reference_rosa(tokens) + self.assertTrue(torch.equal(predict_exact_stateful(tokens), expected)) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") def test_cuda_round_trip_matches_reference(self) -> None: tokens = torch.tensor( From 792bc0f16f2bb594f869ceec5b45c68cac2bae40 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:34:14 +0800 Subject: [PATCH 08/29] Use open-addressed SAM transitions --- src/rosa/_stateful_numba.py | 127 +++++++++++++++++++++++++++--------- 1 file changed, 95 insertions(+), 32 deletions(-) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 6b2de21..dddcd34 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -10,20 +10,36 @@ from torch import Tensor +@njit(cache=True, nogil=True, inline="always") +def _transition_hash( # pragma: no cover - executed as compiled Numba code + state: int, + token: int, +) -> np.uint64: + """Mix the complete signed state/token key into a 64-bit hash.""" + + value = np.uint64(token) + value ^= np.uint64(state) + np.uint64(0x9E3779B97F4A7C15) + value = (value ^ (value >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) + value = (value ^ (value >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) + return value ^ (value >> np.uint64(31)) + + @njit(cache=True, nogil=True, inline="always") def _find_transition( # pragma: no cover - executed as compiled Numba code - 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, state: int, token: int, ) -> int: - edge = head[state] - while edge != -1: - if edge_token[edge] == token: - return int(edge_target[edge]) - edge = edge_next[edge] + """Return the edge index for ``(state, token)``, or -1 when absent.""" + + mask = hash_state.shape[0] - 1 + slot = np.int64(_transition_hash(state, token) & np.uint64(mask)) + while hash_state[slot] != -1: + if hash_state[slot] == state and hash_token[slot] == token: + return int(hash_edge[slot]) + slot = (slot + 1) & mask return -1 @@ -33,6 +49,9 @@ def _add_transition( # pragma: no cover - executed as compiled Numba code edge_token: np.ndarray, edge_target: np.ndarray, edge_next: np.ndarray, + hash_state: np.ndarray, + hash_token: np.ndarray, + hash_edge: np.ndarray, edge_count: int, state: int, token: int, @@ -44,26 +63,32 @@ def _add_transition( # pragma: no cover - executed as compiled Numba code edge_target[edge_count] = target edge_next[edge_count] = head[state] head[state] = edge_count + mask = hash_state.shape[0] - 1 + slot = np.int64(_transition_hash(state, token) & np.uint64(mask)) + while hash_state[slot] != -1: + if hash_state[slot] == state and hash_token[slot] == token: + raise RuntimeError("duplicate suffix automaton transition") + slot = (slot + 1) & mask + hash_state[slot] = state + hash_token[slot] = token + hash_edge[slot] = edge_count return edge_count + 1 @njit(cache=True, nogil=True, inline="always") def _replace_transition( # pragma: no cover - executed as compiled Numba code - 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, state: int, token: int, target: int, ) -> None: - edge = head[state] - while edge != -1: - if edge_token[edge] == token: - edge_target[edge] = target - return - edge = edge_next[edge] - raise RuntimeError("suffix automaton transition not found") + edge = _find_transition(hash_state, hash_token, hash_edge, state, token) + if edge == -1: + raise RuntimeError("suffix automaton transition not found") + edge_target[edge] = target @njit(cache=True, nogil=True, inline="always") @@ -266,6 +291,9 @@ def _step_row( # pragma: no cover - executed as compiled Numba code 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, @@ -289,14 +317,16 @@ def _step_row( # pragma: no cover - executed as compiled Numba code while ( state != -1 - and _find_transition(head, edge_token, edge_target, edge_next, state, token) - == -1 + and _find_transition(hash_state, hash_token, hash_edge, state, token) == -1 ): edge_count = _add_transition( head, edge_token, edge_target, edge_next, + hash_state, + hash_token, + hash_edge, edge_count, state, token, @@ -318,9 +348,8 @@ def _step_row( # pragma: no cover - executed as compiled Numba code lct_stack, ) else: - target = _find_transition( - head, edge_token, edge_target, edge_next, state, token - ) + transition = _find_transition(hash_state, hash_token, hash_edge, state, token) + target = int(edge_target[transition]) if length[state] + 1 == length[target]: suffix_link[current] = target _lct_link_parent( @@ -360,29 +389,35 @@ def _step_row( # pragma: no cover - executed as compiled Numba code edge_token, edge_target, edge_next, + hash_state, + hash_token, + hash_edge, edge_count, clone, int(edge_token[edge]), int(edge_target[edge]), ) edge = edge_next[edge] + transition = _find_transition( + hash_state, hash_token, hash_edge, state, token + ) while ( - state != -1 - and _find_transition( - head, edge_token, edge_target, edge_next, state, token - ) - == target + state != -1 and transition != -1 and edge_target[transition] == target ): _replace_transition( - head, - edge_token, edge_target, - edge_next, + hash_state, + hash_token, + hash_edge, state, token, clone, ) state = suffix_link[state] + if state != -1: + transition = _find_transition( + hash_state, hash_token, hash_edge, state, token + ) _lct_link_parent( lct_left, @@ -470,6 +505,9 @@ def _step_batch_kernel( # pragma: no cover - executed as compiled Numba code 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, @@ -493,6 +531,9 @@ def _step_batch_kernel( # pragma: no cover - executed as compiled Numba code 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], @@ -521,6 +562,9 @@ def _replay_kernel( # pragma: no cover - executed as compiled Numba code 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, @@ -544,6 +588,9 @@ def _replay_kernel( # pragma: no cover - executed as compiled Numba code edge_token, edge_target, edge_next, + hash_state, + hash_token, + hash_edge, suffix_link, length, lct_left, @@ -572,6 +619,9 @@ class _StatefulInferenceState: 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 @@ -598,8 +648,12 @@ def _init_inference_state( raise ValueError("max_length must be > 0") max_states = 2 * max_length + 1 max_edges = 4 * max_length + 1 + # Open addressing needs a power-of-two capacity; at most half the slots + # can be occupied even if the conservative edge bound is reached. + hash_capacity = 1 << (2 * max_edges - 1).bit_length() state_shape = (batch_size, max_states) edge_shape = (batch_size, max_edges) + hash_shape = (batch_size, hash_capacity) suffix_link = np.full(state_shape, -1, dtype=np.int32) return _StatefulInferenceState( batch_size=batch_size, @@ -610,6 +664,9 @@ def _init_inference_state( edge_token=np.empty(edge_shape, dtype=np.int64), edge_target=np.empty(edge_shape, dtype=np.int32), edge_next=np.empty(edge_shape, dtype=np.int32), + hash_state=np.full(hash_shape, -1, dtype=np.int32), + hash_token=np.empty(hash_shape, dtype=np.int64), + hash_edge=np.empty(hash_shape, dtype=np.int32), suffix_link=suffix_link, length=np.zeros(state_shape, dtype=np.int32), lct_left=np.full(state_shape, -1, dtype=np.int32), @@ -644,6 +701,9 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: 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, @@ -683,6 +743,9 @@ def predict_exact_stateful(tokens: Tensor) -> Tensor: 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, From 420363fc9a25b8969c9e04c8fde832dc88ecfd8b Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:37:40 +0800 Subject: [PATCH 09/29] Fuse stateful prefill into one kernel --- src/rosa/__init__.py | 50 +++++++++++++++++++++++++++++++++++++ src/rosa/_stateful_numba.py | 40 +++++++++++++++++++++++++++++ tests/test_inference.py | 44 ++++++++++++++++++++++++++++++++ tests/test_numba_backend.py | 10 ++++++++ 4 files changed, 144 insertions(+) diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py index 57eb60e..157161c 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -33,6 +33,7 @@ "build_virtual_pool_indices", "forward_step", "init_inference_state", + "prefill", "reference_rosa", ] @@ -1047,3 +1048,52 @@ def forward_step(state: ROSAInferenceState, token: Tensor) -> Tensor: output = _forward_step(cast(Any, state._impl), token) return output[0] if squeeze else output + + +def prefill(state: ROSAInferenceState, tokens: Tensor) -> Tensor: + """Consume an initial context and return exact ROSA predictions. + + The state must be empty. ``tokens`` uses shape ``[batch_size, N]``; a + one-dimensional context is accepted when ``batch_size == 1``. The Numba + backend fuses the complete replay into one compiled call. + """ + + if not isinstance(state, ROSAInferenceState): + raise TypeError("state must be a ROSAInferenceState") + if not isinstance(tokens, Tensor): + raise TypeError("tokens must be a torch.Tensor") + squeeze = tokens.ndim == 1 + if squeeze and state.batch_size == 1: + tokens = tokens.unsqueeze(0) + if tokens.ndim != 2 or tokens.shape[0] != state.batch_size: + raise ValueError("tokens must have shape [batch_size, sequence_length]") + if tokens.dtype not in { + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + }: + raise TypeError("tokens must use an integer dtype") + if state.position != 0: + raise RuntimeError("prefill requires an empty inference state") + if tokens.shape[1] > state.max_length: + raise RuntimeError("inference state capacity exceeded") + + if state.backend == "python": + backend_state = cast(_PythonInferenceState, state._impl) + if tokens.shape[1] == 0: + output = torch.empty(tokens.shape, dtype=torch.long, device=tokens.device) + else: + output = torch.stack( + [ + _python_forward_step(backend_state, tokens[:, position]) + for position in range(tokens.shape[1]) + ], + dim=1, + ) + else: + from ._stateful_numba import _prefill + + output = _prefill(cast(Any, state._impl), tokens) + return output[0] if squeeze else output diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index dddcd34..1d7b0d8 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -721,6 +721,46 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: return torch.from_numpy(output).to(device) +def _prefill(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: + """Consume a full initial context through one fused compiled replay.""" + + if state.position != 0: + raise RuntimeError("prefill requires an empty inference state") + if tokens.ndim != 2 or tokens.shape[0] != state.batch_size: + raise ValueError("tokens must have shape [batch_size, sequence_length]") + if tokens.shape[1] > state.max_length: + raise RuntimeError("inference state capacity exceeded") + device = tokens.device + if tokens.shape[1] == 0: + return torch.empty(tokens.shape, dtype=torch.long, device=device) + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + output_array = _replay_kernel( + cpu_tokens.numpy(), + 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.lct_value, + state.lct_lazy, + state.lct_lazy_valid, + state.lct_stack, + state.last, + state.size, + state.edge_count, + ) + state.position = tokens.shape[1] + return torch.from_numpy(output_array).to(device) + + def predict_exact_stateful(tokens: Tensor) -> Tensor: """Return exact top-1 ROSA predictions through the stateful Numba backend.""" diff --git a/tests/test_inference.py b/tests/test_inference.py index 5567f27..ed3de02 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -11,6 +11,7 @@ ROSAInferenceState, forward_step, init_inference_state, + prefill, reference_rosa, ) @@ -78,6 +79,49 @@ def test_exhaustive_binary_clone_sequences(self) -> None: expected, _, _ = reference_rosa(tokens) self.assertTrue(torch.equal(output, expected)) + def test_prefill_then_continuation_matches_reference(self) -> None: + generator = torch.Generator().manual_seed(20260811) + tokens = torch.randint(257, (3, 96), generator=generator) + for backend in ("python", "numba"): + with self.subTest(backend=backend): + state = init_inference_state(3, 96, backend=backend) + initial = prefill(state, tokens[:, :64]) + continuation = torch.stack( + [forward_step(state, tokens[:, index]) for index in range(64, 96)], + dim=1, + ) + output = torch.cat((initial, continuation), dim=1) + expected, _, _ = reference_rosa(tokens) + self.assertTrue(torch.equal(output, expected)) + self.assertEqual(state.position, 96) + + scalar_state = init_inference_state(1, 4, backend="numba") + scalar_tokens = torch.tensor([0, 1, 0, 2]) + self.assertEqual(tuple(prefill(scalar_state, scalar_tokens).shape), (4,)) + + def test_prefill_validation_empty_and_capacity(self) -> None: + state = init_inference_state(2, 3, backend="numba") + empty = prefill(state, torch.empty((2, 0), dtype=torch.long)) + self.assertEqual(tuple(empty.shape), (2, 0)) + self.assertEqual(state.position, 0) + python_state = init_inference_state(2, 3, backend="python") + python_empty = prefill(python_state, torch.empty((2, 0), dtype=torch.long)) + self.assertEqual(tuple(python_empty.shape), (2, 0)) + with self.assertRaisesRegex(ValueError, "shape"): + prefill(state, torch.zeros((1, 2), dtype=torch.long)) + with self.assertRaisesRegex(TypeError, "integer"): + prefill(state, torch.zeros((2, 2))) + with self.assertRaisesRegex(RuntimeError, "capacity"): + prefill(state, torch.zeros((2, 4), dtype=torch.long)) + + prefill(state, torch.zeros((2, 1), dtype=torch.long)) + with self.assertRaisesRegex(RuntimeError, "empty"): + prefill(state, torch.zeros((2, 1), dtype=torch.long)) + with self.assertRaisesRegex(TypeError, "state"): + prefill(cast(Any, object()), torch.zeros((2, 1), dtype=torch.long)) + with self.assertRaisesRegex(TypeError, "Tensor"): + prefill(state, cast(Any, [[0], [0]])) + def test_auto_backend_and_validation(self) -> None: state = init_inference_state(1, backend="auto") self.assertIn(state.backend, {"python", "numba"}) diff --git a/tests/test_numba_backend.py b/tests/test_numba_backend.py index 2e036f2..4588621 100644 --- a/tests/test_numba_backend.py +++ b/tests/test_numba_backend.py @@ -12,6 +12,7 @@ from rosa._stateful_numba import ( _forward_step, _init_inference_state, + _prefill, predict_exact_stateful, ) except ModuleNotFoundError as error: @@ -67,6 +68,15 @@ def test_stateful_private_validation_and_full_replay(self) -> None: with self.assertRaisesRegex(ValueError, "shape"): _forward_step(state, torch.tensor([1, 2])) + prefill_state = _init_inference_state(1, 1) + with self.assertRaisesRegex(ValueError, "shape"): + _prefill(prefill_state, torch.zeros((2, 1), dtype=torch.long)) + with self.assertRaisesRegex(RuntimeError, "capacity"): + _prefill(prefill_state, torch.zeros((1, 2), dtype=torch.long)) + _prefill(prefill_state, torch.zeros((1, 1), dtype=torch.long)) + with self.assertRaisesRegex(RuntimeError, "empty"): + _prefill(prefill_state, torch.zeros((1, 0), dtype=torch.long)) + empty = torch.empty(0, dtype=torch.long) self.assertTrue(torch.equal(predict_exact_stateful(empty), empty)) empty_batch = torch.empty((2, 0), dtype=torch.long) From ae48b0dedb49b80f641f716b2309a99ff0690d51 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:41:30 +0800 Subject: [PATCH 10/29] Use pinned buffers for CUDA steps --- src/rosa/_stateful_numba.py | 83 +++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 1d7b0d8..4b17ae2 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -634,6 +634,13 @@ class _StatefulInferenceState: last: np.ndarray size: np.ndarray edge_count: np.ndarray + cuda_device: torch.device | None + cuda_input: Tensor | None + cuda_input_event: torch.cuda.Event | None + cuda_outputs: list[Tensor] | None + cuda_output_events: list[torch.cuda.Event] | None + cuda_output_used: list[bool] | None + cuda_output_slot: int def _init_inference_state( @@ -679,9 +686,83 @@ def _init_inference_state( last=np.zeros(batch_size, dtype=np.int32), size=np.ones(batch_size, dtype=np.int32), edge_count=np.zeros(batch_size, dtype=np.int32), + cuda_device=None, + cuda_input=None, + cuda_input_event=None, + cuda_outputs=None, + cuda_output_events=None, + cuda_output_used=None, + cuda_output_slot=0, ) +def _cuda_forward_step( # pragma: no cover - exercised on CUDA CI/benchmarks + state: _StatefulInferenceState, + tokens: Tensor, +) -> Tensor: + device = tokens.device + if state.cuda_device != device: + state.cuda_device = device + state.cuda_input = torch.empty( + state.batch_size, dtype=torch.long, device="cpu", pin_memory=True + ) + state.cuda_input_event = torch.cuda.Event() + state.cuda_outputs = [ + torch.empty( + state.batch_size, dtype=torch.long, device="cpu", pin_memory=True + ) + for _ in range(4) + ] + state.cuda_output_events = [torch.cuda.Event() for _ in range(4)] + state.cuda_output_used = [False] * 4 + state.cuda_output_slot = 0 + + assert state.cuda_input is not None + assert state.cuda_input_event is not None + assert state.cuda_outputs is not None + assert state.cuda_output_events is not None + assert state.cuda_output_used is not None + stream = torch.cuda.current_stream(device) + state.cuda_input.copy_(tokens, non_blocking=True) + state.cuda_input_event.record(stream) + state.cuda_input_event.synchronize() + output = _step_batch_kernel( + state.cuda_input.numpy(), + state.position, + 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.lct_value, + state.lct_lazy, + state.lct_lazy_valid, + state.lct_stack, + state.last, + state.size, + state.edge_count, + ) + slot = state.cuda_output_slot + if state.cuda_output_used[slot]: + state.cuda_output_events[slot].synchronize() + pinned_output = state.cuda_outputs[slot] + pinned_output.copy_(torch.from_numpy(output)) + result = pinned_output.to(device, non_blocking=True) + state.cuda_output_events[slot].record(stream) + state.cuda_output_used[slot] = True + state.cuda_output_slot = (slot + 1) % len(state.cuda_outputs) + state.position += 1 + return result + + def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: """Consume one token per batch row and return exact top-1 predictions.""" @@ -692,6 +773,8 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: if state.position >= state.max_length: raise RuntimeError("inference state capacity exceeded") device = tokens.device + if device.type == "cuda": # pragma: no cover - exercised on CUDA benchmarks + return _cuda_forward_step(state, tokens) cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() output = _step_batch_kernel( cpu_tokens.numpy(), From 6faf6e319b6d46fb206ed9c48f03b278a47a023f Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:42:16 +0800 Subject: [PATCH 11/29] Revert "Use pinned buffers for CUDA steps" This reverts commit ae48b0dedb49b80f641f716b2309a99ff0690d51. --- src/rosa/_stateful_numba.py | 83 ------------------------------------- 1 file changed, 83 deletions(-) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 4b17ae2..1d7b0d8 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -634,13 +634,6 @@ class _StatefulInferenceState: last: np.ndarray size: np.ndarray edge_count: np.ndarray - cuda_device: torch.device | None - cuda_input: Tensor | None - cuda_input_event: torch.cuda.Event | None - cuda_outputs: list[Tensor] | None - cuda_output_events: list[torch.cuda.Event] | None - cuda_output_used: list[bool] | None - cuda_output_slot: int def _init_inference_state( @@ -686,83 +679,9 @@ def _init_inference_state( last=np.zeros(batch_size, dtype=np.int32), size=np.ones(batch_size, dtype=np.int32), edge_count=np.zeros(batch_size, dtype=np.int32), - cuda_device=None, - cuda_input=None, - cuda_input_event=None, - cuda_outputs=None, - cuda_output_events=None, - cuda_output_used=None, - cuda_output_slot=0, ) -def _cuda_forward_step( # pragma: no cover - exercised on CUDA CI/benchmarks - state: _StatefulInferenceState, - tokens: Tensor, -) -> Tensor: - device = tokens.device - if state.cuda_device != device: - state.cuda_device = device - state.cuda_input = torch.empty( - state.batch_size, dtype=torch.long, device="cpu", pin_memory=True - ) - state.cuda_input_event = torch.cuda.Event() - state.cuda_outputs = [ - torch.empty( - state.batch_size, dtype=torch.long, device="cpu", pin_memory=True - ) - for _ in range(4) - ] - state.cuda_output_events = [torch.cuda.Event() for _ in range(4)] - state.cuda_output_used = [False] * 4 - state.cuda_output_slot = 0 - - assert state.cuda_input is not None - assert state.cuda_input_event is not None - assert state.cuda_outputs is not None - assert state.cuda_output_events is not None - assert state.cuda_output_used is not None - stream = torch.cuda.current_stream(device) - state.cuda_input.copy_(tokens, non_blocking=True) - state.cuda_input_event.record(stream) - state.cuda_input_event.synchronize() - output = _step_batch_kernel( - state.cuda_input.numpy(), - state.position, - 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.lct_value, - state.lct_lazy, - state.lct_lazy_valid, - state.lct_stack, - state.last, - state.size, - state.edge_count, - ) - slot = state.cuda_output_slot - if state.cuda_output_used[slot]: - state.cuda_output_events[slot].synchronize() - pinned_output = state.cuda_outputs[slot] - pinned_output.copy_(torch.from_numpy(output)) - result = pinned_output.to(device, non_blocking=True) - state.cuda_output_events[slot].record(stream) - state.cuda_output_used[slot] = True - state.cuda_output_slot = (slot + 1) % len(state.cuda_outputs) - state.position += 1 - return result - - def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: """Consume one token per batch row and return exact top-1 predictions.""" @@ -773,8 +692,6 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: if state.position >= state.max_length: raise RuntimeError("inference state capacity exceeded") device = tokens.device - if device.type == "cuda": # pragma: no cover - exercised on CUDA benchmarks - return _cuda_forward_step(state, tokens) cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() output = _step_batch_kernel( cpu_tokens.numpy(), From a795dcf1f4d9dad2f581f4bf9697daa3d18af2d0 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:45:18 +0800 Subject: [PATCH 12/29] Parallelize large stateful batches --- src/rosa/_stateful_numba.py | 65 +++++++++++++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 2 deletions(-) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 1d7b0d8..a1b89a6 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -6,7 +6,7 @@ import numpy as np import torch -from numba import njit +from numba import njit, prange from torch import Tensor @@ -554,6 +554,64 @@ def _step_batch_kernel( # pragma: no cover - executed as compiled Numba code return output +@njit(cache=True, nogil=True, parallel=True) +def _step_batch_parallel_kernel( # pragma: no cover - compiled Numba code + tokens: np.ndarray, + position: 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, + lct_value: np.ndarray, + lct_lazy: np.ndarray, + lct_lazy_valid: np.ndarray, + lct_stack: np.ndarray, + last: np.ndarray, + size: np.ndarray, + edge_count: np.ndarray, +) -> np.ndarray: + output = np.empty(tokens.shape[0], dtype=np.int64) + for batch_index in prange(tokens.shape[0]): + prediction, new_last, new_size, new_edge_count = _step_row( + int(tokens[batch_index]), + position, + 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], + lct_value[batch_index], + lct_lazy[batch_index], + lct_lazy_valid[batch_index], + lct_stack[batch_index], + int(last[batch_index]), + int(size[batch_index]), + int(edge_count[batch_index]), + ) + output[batch_index] = prediction + last[batch_index] = new_last + size[batch_index] = new_size + edge_count[batch_index] = new_edge_count + return output + + @njit(cache=True, nogil=True) def _replay_kernel( # pragma: no cover - executed as compiled Numba code tokens: np.ndarray, @@ -693,7 +751,10 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: raise RuntimeError("inference state capacity exceeded") device = tokens.device cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() - output = _step_batch_kernel( + step_kernel = ( + _step_batch_parallel_kernel if state.batch_size >= 8 else _step_batch_kernel + ) + output = step_kernel( cpu_tokens.numpy(), state.position, state.history, From f23eea09b93b05553ad4cd5f8adfa493ee963ed6 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 05:56:19 +0800 Subject: [PATCH 13/29] Revert "Parallelize large stateful batches" This reverts commit a795dcf1f4d9dad2f581f4bf9697daa3d18af2d0. --- src/rosa/_stateful_numba.py | 65 ++----------------------------------- 1 file changed, 2 insertions(+), 63 deletions(-) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index a1b89a6..1d7b0d8 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -6,7 +6,7 @@ import numpy as np import torch -from numba import njit, prange +from numba import njit from torch import Tensor @@ -554,64 +554,6 @@ def _step_batch_kernel( # pragma: no cover - executed as compiled Numba code return output -@njit(cache=True, nogil=True, parallel=True) -def _step_batch_parallel_kernel( # pragma: no cover - compiled Numba code - tokens: np.ndarray, - position: 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, - lct_value: np.ndarray, - lct_lazy: np.ndarray, - lct_lazy_valid: np.ndarray, - lct_stack: np.ndarray, - last: np.ndarray, - size: np.ndarray, - edge_count: np.ndarray, -) -> np.ndarray: - output = np.empty(tokens.shape[0], dtype=np.int64) - for batch_index in prange(tokens.shape[0]): - prediction, new_last, new_size, new_edge_count = _step_row( - int(tokens[batch_index]), - position, - 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], - lct_value[batch_index], - lct_lazy[batch_index], - lct_lazy_valid[batch_index], - lct_stack[batch_index], - int(last[batch_index]), - int(size[batch_index]), - int(edge_count[batch_index]), - ) - output[batch_index] = prediction - last[batch_index] = new_last - size[batch_index] = new_size - edge_count[batch_index] = new_edge_count - return output - - @njit(cache=True, nogil=True) def _replay_kernel( # pragma: no cover - executed as compiled Numba code tokens: np.ndarray, @@ -751,10 +693,7 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: raise RuntimeError("inference state capacity exceeded") device = tokens.device cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() - step_kernel = ( - _step_batch_parallel_kernel if state.batch_size >= 8 else _step_batch_kernel - ) - output = step_kernel( + output = _step_batch_kernel( cpu_tokens.numpy(), state.position, state.history, From 133a6f8967a05fa7451b78af499cc099f58b4e96 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:01:40 +0800 Subject: [PATCH 14/29] Use optional native stateful stepping --- src/rosa/_stateful_numba.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 1d7b0d8..f7feb88 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Any import numpy as np import torch @@ -634,6 +635,7 @@ class _StatefulInferenceState: last: np.ndarray size: np.ndarray edge_count: np.ndarray + native_state: Any def _init_inference_state( @@ -679,9 +681,28 @@ def _init_inference_state( last=np.zeros(batch_size, dtype=np.int32), size=np.ones(batch_size, dtype=np.int32), edge_count=np.zeros(batch_size, dtype=np.int32), + native_state=None, ) +def _native_step( # pragma: no cover - optional native companion + state: _StatefulInferenceState, + cpu_tokens: Tensor, +) -> np.ndarray | None: + if state.native_state is False: + return None + if state.native_state is None: + try: + from rosa_native_step import ( # type: ignore[reportMissingImports] + NativeState, + ) + except ModuleNotFoundError: + state.native_state = False + return None + state.native_state = NativeState(state) + return state.native_state.step(cpu_tokens.numpy()) + + def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: """Consume one token per batch row and return exact top-1 predictions.""" @@ -693,6 +714,9 @@ def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: raise RuntimeError("inference state capacity exceeded") device = tokens.device cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + native_output = _native_step(state, cpu_tokens) + if native_output is not None: # pragma: no cover - optional native companion + return torch.from_numpy(native_output).to(device) output = _step_batch_kernel( cpu_tokens.numpy(), state.position, From cf8f8e3f4d0c58fc746308198c7b4312e258bb8e Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:13:50 +0800 Subject: [PATCH 15/29] Build stateful ROSA prefill in bulk --- src/rosa/_stateful_numba.py | 388 +++++++++++++++++++++++++++++++++++- 1 file changed, 385 insertions(+), 3 deletions(-) diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index f7feb88..213d18e 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -608,6 +608,390 @@ def _replay_kernel( # pragma: no cover - executed as compiled Numba code return output +@njit(cache=True, nogil=True, inline="always") +def _range_max( # pragma: no cover - executed as compiled Numba code + tree: np.ndarray, + base: int, + start: int, + stop: int, +) -> int: + result = -1 + left = base + start + right = base + stop + while left < right: + if left & 1: + if tree[left] > result: + result = int(tree[left]) + left += 1 + if right & 1: + right -= 1 + if tree[right] > result: + result = int(tree[right]) + left >>= 1 + right >>= 1 + return result + + +@njit(cache=True, nogil=True, inline="always") +def _point_max( # pragma: no cover - executed as compiled Numba code + tree: np.ndarray, + base: int, + index: int, + value: int, +) -> None: + node = base + index + if tree[node] >= value: + return + tree[node] = value + node >>= 1 + while node != 0: + updated = tree[node << 1] + if tree[(node << 1) | 1] > updated: + updated = tree[(node << 1) | 1] + if tree[node] == updated: + break + tree[node] = updated + node >>= 1 + + +@njit(cache=True, nogil=True, inline="always") +def _fenwick_prefix( # pragma: no cover - executed as compiled Numba code + tree: np.ndarray, + stop: int, +) -> int: + result = 0 + node = stop + while node > 0: + result += int(tree[node]) + node -= node & -node + return result + + +@njit(cache=True, nogil=True, inline="always") +def _fenwick_add( # pragma: no cover - executed as compiled Numba code + tree: np.ndarray, + index: int, +) -> None: + node = index + 1 + while node < tree.shape[0]: + tree[node] += 1 + node += node & -node + + +@njit(cache=True, nogil=True, inline="always") +def _fenwick_select( # pragma: no cover - executed as compiled Numba code + tree: np.ndarray, + rank: int, +) -> int: + """Return the zero-based position of a one-based order statistic.""" + + node = 0 + step = 1 + size = tree.shape[0] - 1 + while (step << 1) <= size: + step <<= 1 + while step != 0: + candidate = node + step + if candidate <= size and tree[candidate] < rank: + node = candidate + rank -= int(tree[candidate]) + step >>= 1 + return node + + +@njit(cache=True, nogil=True, inline="always") +def _tree_lca( # pragma: no cover - executed as compiled Numba code + up: np.ndarray, + tin: np.ndarray, + tout: np.ndarray, + first: int, + second: int, +) -> int: + if tin[first] <= tin[second] and tout[second] <= tout[first]: + return first + if tin[second] <= tin[first] and tout[first] <= tout[second]: + return second + current = first + for level in range(up.shape[0] - 1, -1, -1): + ancestor = up[level, current] + if ancestor != -1 and not ( + tin[ancestor] <= tin[second] and tout[second] <= tout[ancestor] + ): + current = ancestor + return int(up[0, current]) + + +@njit(cache=True, nogil=True) +def _bulk_prefill_row( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, + 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, + lct_value: np.ndarray, + lct_lazy_valid: np.ndarray, +) -> tuple[np.ndarray, int, int, int]: + """Build the final SAM, answer prefixes offline, and seed the LCT.""" + + token_count = tokens.shape[0] + output = np.full(token_count, -1, dtype=np.int64) + prefix_state = np.empty(token_count, dtype=np.int32) + last = 0 + size = 1 + edge_count = 0 + + for position in range(token_count): + token = int(tokens[position]) + history[position] = token + if size >= head.shape[0]: + raise RuntimeError("suffix automaton state capacity exceeded") + current = size + size += 1 + length[current] = length[last] + 1 + state = last + while ( + state != -1 + and _find_transition(hash_state, hash_token, hash_edge, state, token) == -1 + ): + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + hash_state, + hash_token, + hash_edge, + edge_count, + state, + token, + current, + ) + state = suffix_link[state] + + if state == -1: + suffix_link[current] = 0 + else: + transition = _find_transition( + hash_state, hash_token, hash_edge, state, token + ) + target = int(edge_target[transition]) + if length[state] + 1 == length[target]: + suffix_link[current] = target + else: + if size >= head.shape[0]: + raise RuntimeError("suffix automaton state capacity exceeded") + clone = size + size += 1 + length[clone] = length[state] + 1 + suffix_link[clone] = suffix_link[target] + edge = head[target] + while edge != -1: + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + hash_state, + hash_token, + hash_edge, + edge_count, + clone, + int(edge_token[edge]), + int(edge_target[edge]), + ) + edge = edge_next[edge] + transition = _find_transition( + hash_state, hash_token, hash_edge, state, token + ) + while ( + state != -1 + and transition != -1 + and edge_target[transition] == target + ): + _replace_transition( + edge_target, + hash_state, + hash_token, + hash_edge, + state, + token, + clone, + ) + state = suffix_link[state] + if state != -1: + transition = _find_transition( + hash_state, hash_token, hash_edge, state, token + ) + suffix_link[target] = clone + suffix_link[current] = clone + + last = current + prefix_state[position] = current + + first_child = np.full(size, -1, dtype=np.int32) + next_sibling = np.full(size, -1, dtype=np.int32) + for node in range(1, size): + parent = suffix_link[node] + next_sibling[node] = first_child[parent] + first_child[parent] = node + tin = np.empty(size, dtype=np.int32) + tout = np.empty(size, dtype=np.int32) + euler_node = np.empty(size, dtype=np.int32) + dfs_nodes = np.empty(size, dtype=np.int32) + dfs_next = np.empty(size, dtype=np.int32) + depth = 0 + timer = 0 + dfs_nodes[0] = 0 + dfs_next[0] = first_child[0] + tin[0] = timer + euler_node[timer] = 0 + timer += 1 + while depth >= 0: + child = dfs_next[depth] + if child == -1: + tout[dfs_nodes[depth]] = timer + depth -= 1 + else: + dfs_next[depth] = next_sibling[child] + depth += 1 + dfs_nodes[depth] = child + dfs_next[depth] = first_child[child] + tin[child] = timer + euler_node[timer] = child + timer += 1 + + levels = 1 + span = 1 + while span < size: + levels += 1 + span <<= 1 + up = np.full((levels, size), -1, dtype=np.int32) + for node in range(size): + up[0, node] = suffix_link[node] + for level in range(1, levels): + for node in range(size): + ancestor = up[level - 1, node] + if ancestor != -1: + up[level, node] = up[level - 1, ancestor] + + base = 1 + while base < size: + base <<= 1 + active = np.full(base << 1, -1, dtype=np.int64) + active_order = np.zeros(size + 1, dtype=np.int32) + for position in range(token_count): + node = int(suffix_link[prefix_state[position]]) + if node != -1: + # Among a DFS-ordered set, one of the predecessor/successor has + # the deepest LCA with the query node. Thus two order-statistic + # lookups replace a path of segment-tree range probes. + node_tin = int(tin[node]) + preceding_count = _fenwick_prefix(active_order, node_tin + 1) + best = 0 + if preceding_count > 0: + preceding_tin = _fenwick_select(active_order, preceding_count) + best = _tree_lca(up, tin, tout, node, int(euler_node[preceding_tin])) + before_count = _fenwick_prefix(active_order, node_tin) + if before_count < position: + following_tin = _fenwick_select(active_order, before_count + 1) + candidate = _tree_lca( + up, tin, tout, node, int(euler_node[following_tin]) + ) + if length[candidate] > length[best]: + best = candidate + node = best + source = _range_max(active, base, int(tin[node]), int(tout[node])) + if node != 0 and source >= 0: + output[position] = history[source + 1] + _point_max(active, base, int(tin[prefix_state[position]]), position) + _fenwick_add(active_order, int(tin[prefix_state[position]])) + + latest_end = np.full(size, -1, dtype=np.int64) + for position in range(token_count): + latest_end[prefix_state[position]] = position + counts = np.zeros(token_count + 1, dtype=np.int32) + for node in range(size): + counts[length[node]] += 1 + for index in range(1, counts.shape[0]): + counts[index] += counts[index - 1] + order = np.empty(size, dtype=np.int32) + for node in range(size - 1, -1, -1): + node_length = length[node] + counts[node_length] -= 1 + order[counts[node_length]] = node + for index in range(size - 1, 0, -1): + node = order[index] + parent = suffix_link[node] + if latest_end[node] > latest_end[parent]: + latest_end[parent] = latest_end[node] + for node in range(size): + lct_left[node] = -1 + lct_right[node] = -1 + lct_parent[node] = suffix_link[node] + lct_value[node] = latest_end[node] + lct_lazy_valid[node] = 0 + + return output, last, size, edge_count + + +@njit(cache=True, nogil=True) +def _bulk_prefill_kernel( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, + 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, + lct_value: np.ndarray, + lct_lazy_valid: np.ndarray, + last: np.ndarray, + size: np.ndarray, + edge_count: np.ndarray, +) -> np.ndarray: + output = np.empty(tokens.shape, dtype=np.int64) + for batch_index in range(tokens.shape[0]): + row_output, row_last, row_size, row_edge_count = _bulk_prefill_row( + tokens[batch_index], + 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], + lct_value[batch_index], + lct_lazy_valid[batch_index], + ) + output[batch_index] = row_output + last[batch_index] = row_last + size[batch_index] = row_size + edge_count[batch_index] = row_edge_count + return output + + @dataclass class _StatefulInferenceState: """Fixed-capacity, independently batched exact ROSA inference state.""" @@ -758,7 +1142,7 @@ def _prefill(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: if tokens.shape[1] == 0: return torch.empty(tokens.shape, dtype=torch.long, device=device) cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() - output_array = _replay_kernel( + output_array = _bulk_prefill_kernel( cpu_tokens.numpy(), state.history, state.head, @@ -774,9 +1158,7 @@ def _prefill(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: state.lct_right, state.lct_parent, state.lct_value, - state.lct_lazy, state.lct_lazy_valid, - state.lct_stack, state.last, state.size, state.edge_count, From 3c2028105871316be1c48e5fa6036e8798493933 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:31:50 +0800 Subject: [PATCH 16/29] Package ROSA native acceleration for 0.2 --- .github/workflows/ci.yml | 31 +++ README.md | 9 + native/.gitignore | 5 + native/README.md | 63 ++++++ native/pyproject.toml | 36 +++ native/setup.py | 16 ++ native/src/rosa_native_step.cpp | 376 ++++++++++++++++++++++++++++++++ native/tests/smoke.py | 45 ++++ pyproject.toml | 2 +- src/rosa/_stateful_numba.py | 2 + tests/test_numba_backend.py | 1 + uv.lock | 2 +- 12 files changed, 586 insertions(+), 2 deletions(-) create mode 100644 native/.gitignore create mode 100644 native/README.md create mode 100644 native/pyproject.toml create mode 100644 native/setup.py create mode 100644 native/src/rosa_native_step.cpp create mode 100644 native/tests/smoke.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fe3ec43..9a8c709 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -84,3 +84,34 @@ jobs: assert state.backend == "python" assert forward_step(state, torch.tensor(0)).item() == -1 PY + + native-smoke: + name: Native companion Python ${{ matrix.python-version }} + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.14"] + + steps: + - name: Check out repository + uses: actions/checkout@3d3c42e5aac5ba805825a76410c181273ba90b1 # v7.0.1 + + - name: Install uv and Python + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + version: "0.8.15" + python-version: ${{ matrix.python-version }} + enable-cache: true + + - name: Install ROSA with Numba + run: uv sync --locked --all-groups --extra numba + + - name: Build and install native companion + run: | + uv build --python .venv/bin/python --wheel native --out-dir native/dist + uv pip install --python .venv/bin/python --no-deps --force-reinstall native/dist/*.whl + + - name: Run native smoke test + run: uv run python native/tests/smoke.py diff --git a/README.md b/README.md index e8a87e8..256bed8 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,15 @@ Install the stateful Link-Cut Tree backend with: uv add 'rosa-torch[numba]' ``` +For the lowest CPU step latency, install a locally built native companion wheel +(or a published wheel once multi-ABI releases are enabled): + +```bash +uv pip install native/dist/rosa_torch_native-0.2.0-*.whl +``` + +The stateful backend detects it lazily and otherwise falls back to Numba. + Install the package and its locked development dependencies with [uv](https://docs.astral.sh/uv/): ```bash diff --git a/native/.gitignore b/native/.gitignore new file mode 100644 index 0000000..d458baa --- /dev/null +++ b/native/.gitignore @@ -0,0 +1,5 @@ +build/ +dist/ +*.egg-info/ +*.so +*.pyd diff --git a/native/README.md b/native/README.md new file mode 100644 index 0000000..2226303 --- /dev/null +++ b/native/README.md @@ -0,0 +1,63 @@ +# `rosa-torch-native` + +Compagnon natif optionnel de [`rosa-torch`](https://github.com/aabbdev/rosa) +pour l'étape d'inférence exacte SAM + Link-Cut Tree sur CPU. La distribution +installe le module d'extension importable `rosa_native_step`; elle ne remplace +pas le package Python principal. + +Le cœur C++ est la copie exacte du prototype validé +`benchmark/production-opt-20260811/native/native_step.cpp`. Il lie une fois les +tableaux NumPy d'un `_StatefulInferenceState`, les modifie sur place et libère +le GIL pendant le calcul. Il n'inclut ni n'appelle libtorch. La dépendance +runtime `rosa-torch[numba]>=0.2,<0.3` fournit le contrat d'état compatible, +PyTorch, NumPy et Numba. + +Le constructeur valide intégralement formes, types, compteurs et version ABI +avant de conserver les pointeurs. L'ABI d'état native actuelle vaut `1`. + +## Installation et utilisation + +Installez un wheel correspondant à la version de Python et à la plateforme : + +```bash +python -m pip install rosa-torch-native +``` + +`rosa-torch` détecte automatiquement le module dans son backend d'inférence +Numba. L'API bas niveau reste disponible pour diagnostic : + +```python +from rosa_native_step import NativeState +``` + +`NativeState(state).step(tokens_numpy)` attend un vecteur NumPy contigu +convertible en `int64`, de forme `[batch_size]`. L'objet conserve une référence +à l'état Python et expose sa `position` en lecture seule. + +## Construction locale isolée + +Le backend PEP 517 est setuptools, avec pybind11 uniquement comme dépendance de +construction. `Pybind11Extension` sélectionne C++17 et setuptools fournit les +options d'extension propres à macOS/Linux; `-O3` et `NDEBUG` sont ajoutés sur +les deux plateformes. + +Depuis la racine du dépôt : + +```bash +uv build --wheel native --out-dir /tmp/rosa-native-dist +uv run --isolated \ + --with '.[numba]' \ + --with /tmp/rosa-native-dist/rosa_torch_native-0.2.0-*.whl \ + native/tests/smoke.py +``` + +Le smoke force Numba comme oracle, laisse le chemin ROSA courant charger le +compagnon, puis compare les prédictions étape par étape. + +## Publication multi-plateforme + +Étape suivante : ajouter un workflow `cibuildwheel` dédié après avoir fixé la +matrice Python/architectures, les cibles Linux (manylinux) et la politique de +publication. Aucun workflow n'est ajouté ici afin de ne pas publier une +matrice non validée; chaque wheel contient du code natif et doit être produit +séparément par ABI et plateforme. diff --git a/native/pyproject.toml b/native/pyproject.toml new file mode 100644 index 0000000..3f2f74e --- /dev/null +++ b/native/pyproject.toml @@ -0,0 +1,36 @@ +[build-system] +requires = [ + "setuptools>=68", + "wheel", + "pybind11>=2.11,<4", +] +build-backend = "setuptools.build_meta" + +[project] +name = "rosa-torch-native" +version = "0.2.0" +description = "Optional native CPU step companion for rosa-torch" +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +classifiers = [ + "Development Status :: 3 - Alpha", + "Programming Language :: C++", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "rosa-torch[numba]>=0.2,<0.3", +] + +[project.urls] +Repository = "https://github.com/aabbdev/rosa" +Issues = "https://github.com/aabbdev/rosa/issues" + +[tool.setuptools] +include-package-data = false diff --git a/native/setup.py b/native/setup.py new file mode 100644 index 0000000..cd04b57 --- /dev/null +++ b/native/setup.py @@ -0,0 +1,16 @@ +from pybind11.setup_helpers import Pybind11Extension, build_ext +from setuptools import setup + +extension = Pybind11Extension( + "rosa_native_step", + ["src/rosa_native_step.cpp"], + cxx_std=17, + define_macros=[("NDEBUG", "1")], + extra_compile_args=["-O3"], +) + +setup( + ext_modules=[extension], + cmdclass={"build_ext": build_ext}, + zip_safe=False, +) diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp new file mode 100644 index 0000000..f4e6b19 --- /dev/null +++ b/native/src/rosa_native_step.cpp @@ -0,0 +1,376 @@ +#include +#include + +#include +#include +#include + +namespace py = pybind11; + +class NativeState { + public: + explicit NativeState(py::object state) : state_(std::move(state)) { + const int64_t abi = py::cast(state_.attr("native_abi_version")); + if (abi != 1) throw py::value_error("unsupported native state ABI"); + history_ = bind("history"); + head_ = bind("head"); + edge_token_ = bind("edge_token"); + edge_target_ = bind("edge_target"); + edge_next_ = bind("edge_next"); + hash_state_ = bind("hash_state"); + hash_token_ = bind("hash_token"); + hash_edge_ = bind("hash_edge"); + suffix_link_ = bind("suffix_link"); + length_ = bind("length"); + left_ = bind("lct_left"); + right_ = bind("lct_right"); + parent_ = bind("lct_parent"); + value_ = bind("lct_value"); + lazy_ = bind("lct_lazy"); + lazy_valid_ = bind("lct_lazy_valid"); + stack_ = bind("lct_stack"); + last_ = bind("last"); + size_ = bind("size"); + edge_count_ = bind("edge_count"); + batch_ = py::cast(state_.attr("batch_size")); + max_length_ = py::cast(state_.attr("max_length")); + position_ = py::cast(state_.attr("position")); + state_capacity_ = head_.shape(1); + edge_capacity_ = edge_token_.shape(1); + hash_capacity_ = hash_state_.shape(1); + validate_shapes(); + } + + py::array_t step( + py::array_t tokens) { + if (tokens.ndim() != 1 || tokens.shape(0) != batch_) { + throw py::value_error("tokens must be contiguous int64 [batch_size]"); + } + if (position_ >= max_length_) { + throw std::runtime_error("inference state capacity exceeded"); + } + py::array_t output(batch_); + const int64_t* in = tokens.data(); + int64_t* out = output.mutable_data(); + { + py::gil_scoped_release release; + for (int64_t b = 0; b < batch_; ++b) out[b] = step_row(b, in[b]); + } + ++position_; + state_.attr("position") = py::int_(position_); + return output; + } + + int64_t position() const { return position_; } + + private: + template + py::array_t bind(const char* name) { + py::object object = state_.attr(name); + if (!py::isinstance>(object)) { + throw py::type_error(std::string(name) + " has an unexpected dtype"); + } + auto array = py::cast>(object); + if (!array.writeable()) throw py::value_error(std::string(name) + " is readonly"); + return array; + } + + void validate_shapes() { + if (batch_ <= 0 || max_length_ <= 0 || position_ < 0 || + position_ > max_length_ || state_capacity_ <= 0 || edge_capacity_ <= 0 || + hash_capacity_ <= 0) { + throw py::value_error("incompatible _StatefulInferenceState layout"); + } + if (!matrix_shape(history_, batch_, max_length_) || + !matrix_shape(head_, batch_, state_capacity_) || + !matrix_shape(edge_token_, batch_, edge_capacity_) || + !matrix_shape(edge_target_, batch_, edge_capacity_) || + !matrix_shape(edge_next_, batch_, edge_capacity_) || + !matrix_shape(hash_state_, batch_, hash_capacity_) || + !matrix_shape(hash_token_, batch_, hash_capacity_) || + !matrix_shape(hash_edge_, batch_, hash_capacity_) || + !matrix_shape(suffix_link_, batch_, state_capacity_) || + !matrix_shape(length_, batch_, state_capacity_) || + !matrix_shape(left_, batch_, state_capacity_) || + !matrix_shape(right_, batch_, state_capacity_) || + !matrix_shape(parent_, batch_, state_capacity_) || + !matrix_shape(value_, batch_, state_capacity_) || + !matrix_shape(lazy_, batch_, state_capacity_) || + !matrix_shape(lazy_valid_, batch_, state_capacity_) || + !matrix_shape(stack_, batch_, state_capacity_) || + !vector_shape(last_, batch_) || !vector_shape(size_, batch_) || + !vector_shape(edge_count_, batch_)) { + throw py::value_error("incompatible _StatefulInferenceState layout"); + } + if ((hash_capacity_ & (hash_capacity_ - 1)) != 0) { + throw py::value_error("hash capacity must be a power of two"); + } + for (int64_t b = 0; b < batch_; ++b) { + if (last_.data()[b] < 0 || last_.data()[b] >= state_capacity_ || + size_.data()[b] < 1 || size_.data()[b] > state_capacity_ || + edge_count_.data()[b] < 0 || edge_count_.data()[b] > edge_capacity_) { + throw py::value_error("incompatible _StatefulInferenceState counters"); + } + } + } + + template + bool matrix_shape(const py::array_t& array, + int64_t rows, int64_t columns) const { + return array.ndim() == 2 && array.shape(0) == rows && + array.shape(1) == columns; + } + + template + bool vector_shape(const py::array_t& array, + int64_t length) const { + return array.ndim() == 1 && array.shape(0) == length; + } + + inline uint64_t transition_hash(int32_t state, int64_t token) const { + uint64_t v = static_cast(token); + v ^= static_cast(state) + UINT64_C(0x9E3779B97F4A7C15); + v = (v ^ (v >> 30)) * UINT64_C(0xBF58476D1CE4E5B9); + v = (v ^ (v >> 27)) * UINT64_C(0x94D049BB133111EB); + return v ^ (v >> 31); + } + + inline int64_t idx(int64_t b, int64_t width, int64_t i) const { + return b * width + i; + } + + int32_t find_transition(int64_t b, int32_t state, int64_t token) const { + const int32_t* hs = hash_state_.data(); + const int64_t* ht = hash_token_.data(); + const int32_t* he = hash_edge_.data(); + int64_t slot = static_cast(transition_hash(state, token) & + static_cast(hash_capacity_ - 1)); + while (hs[idx(b, hash_capacity_, slot)] != -1) { + const int64_t at = idx(b, hash_capacity_, slot); + if (hs[at] == state && ht[at] == token) return he[at]; + slot = (slot + 1) & (hash_capacity_ - 1); + } + return -1; + } + + int32_t add_transition(int64_t b, int32_t count, int32_t state, + int64_t token, int32_t target) { + if (count >= edge_capacity_) throw std::runtime_error("transition capacity exceeded"); + int32_t* head = head_.mutable_data(); + int64_t* et = edge_token_.mutable_data(); + int32_t* eg = edge_target_.mutable_data(); + int32_t* en = edge_next_.mutable_data(); + const int64_t edge_at = idx(b, edge_capacity_, count); + const int64_t state_at = idx(b, state_capacity_, state); + et[edge_at] = token; + eg[edge_at] = target; + en[edge_at] = head[state_at]; + head[state_at] = count; + int64_t slot = static_cast(transition_hash(state, token) & + static_cast(hash_capacity_ - 1)); + int32_t* hs = hash_state_.mutable_data(); + int64_t* ht = hash_token_.mutable_data(); + int32_t* he = hash_edge_.mutable_data(); + while (hs[idx(b, hash_capacity_, slot)] != -1) { + const int64_t at = idx(b, hash_capacity_, slot); + if (hs[at] == state && ht[at] == token) + throw std::runtime_error("duplicate suffix automaton transition"); + slot = (slot + 1) & (hash_capacity_ - 1); + } + const int64_t at = idx(b, hash_capacity_, slot); + hs[at] = state; + ht[at] = token; + he[at] = count; + return count + 1; + } + + void replace_transition(int64_t b, int32_t state, int64_t token, + int32_t target) { + const int32_t edge = find_transition(b, state, token); + if (edge == -1) throw std::runtime_error("transition not found"); + edge_target_.mutable_data()[idx(b, edge_capacity_, edge)] = target; + } + + inline bool is_aux_root(int64_t b, int32_t node) const { + const int32_t* left = left_.data(); + const int32_t* right = right_.data(); + const int32_t* parent = parent_.data(); + const int64_t at = idx(b, state_capacity_, node); + const int32_t p = parent[at]; + return p == -1 || (left[idx(b, state_capacity_, p)] != node && + right[idx(b, state_capacity_, p)] != node); + } + + inline void apply(int64_t b, int32_t node, int64_t assigned) { + if (node == -1) return; + const int64_t at = idx(b, state_capacity_, node); + value_.mutable_data()[at] = assigned; + lazy_.mutable_data()[at] = assigned; + lazy_valid_.mutable_data()[at] = 1; + } + + inline void push(int64_t b, int32_t node) { + const int64_t at = idx(b, state_capacity_, node); + if (lazy_valid_.data()[at]) { + const int64_t assigned = lazy_.data()[at]; + apply(b, left_.data()[at], assigned); + apply(b, right_.data()[at], assigned); + lazy_valid_.mutable_data()[at] = 0; + } + } + + inline void rotate(int64_t b, int32_t node) { + int32_t* left = left_.mutable_data(); + int32_t* right = right_.mutable_data(); + int32_t* parent = parent_.mutable_data(); + const auto at = [&](int32_t n) { return idx(b, state_capacity_, n); }; + const int32_t p = parent[at(node)], g = parent[at(p)]; + int32_t middle; + if (left[at(p)] == node) { + middle = right[at(node)]; right[at(node)] = p; left[at(p)] = middle; + } else { + middle = left[at(node)]; left[at(node)] = p; right[at(p)] = middle; + } + if (middle != -1) parent[at(middle)] = p; + parent[at(p)] = node; parent[at(node)] = g; + if (g != -1) { + if (left[at(g)] == p) left[at(g)] = node; + else if (right[at(g)] == p) right[at(g)] = node; + } + } + + void splay(int64_t b, int32_t node) { + int32_t* stack = stack_.mutable_data() + b * state_capacity_; + int32_t depth = 0, ancestor = node; + stack[depth++] = ancestor; + while (!is_aux_root(b, ancestor)) { + ancestor = parent_.data()[idx(b, state_capacity_, ancestor)]; + stack[depth++] = ancestor; + } + while (depth > 0) push(b, stack[--depth]); + while (!is_aux_root(b, node)) { + const int32_t p = parent_.data()[idx(b, state_capacity_, node)]; + if (!is_aux_root(b, p)) { + const int32_t g = parent_.data()[idx(b, state_capacity_, p)]; + if ((left_.data()[idx(b, state_capacity_, p)] == node) == + (left_.data()[idx(b, state_capacity_, g)] == p)) rotate(b, p); + else rotate(b, node); + } + rotate(b, node); + } + } + + void access(int64_t b, int32_t node) { + int32_t last = -1, current = node; + while (current != -1) { + splay(b, current); + right_.mutable_data()[idx(b, state_capacity_, current)] = last; + if (last != -1) parent_.mutable_data()[idx(b, state_capacity_, last)] = current; + last = current; + current = parent_.data()[idx(b, state_capacity_, current)]; + } + splay(b, node); + } + + int64_t point_query(int64_t b, int32_t node) { + access(b, node); + return value_.data()[idx(b, state_capacity_, node)]; + } + void path_assign(int64_t b, int32_t node, int64_t assigned) { + access(b, node); apply(b, node, assigned); + } + void cut_parent(int64_t b, int32_t node) { + access(b, node); + const int64_t at = idx(b, state_capacity_, node); + const int32_t ancestors = left_.data()[at]; + left_.mutable_data()[at] = -1; + if (ancestors != -1) parent_.mutable_data()[idx(b, state_capacity_, ancestors)] = -1; + } + void link_parent(int64_t b, int32_t node, int32_t represented_parent) { + access(b, node); + parent_.mutable_data()[idx(b, state_capacity_, node)] = represented_parent; + } + + int64_t step_row(int64_t b, int64_t token) { + int32_t* last_a = last_.mutable_data(); + int32_t* size_a = size_.mutable_data(); + int32_t* edge_count_a = edge_count_.mutable_data(); + int32_t last = last_a[b], size = size_a[b], edge_count = edge_count_a[b]; + history_.mutable_data()[idx(b, max_length_, position_)] = token; + if (size >= state_capacity_) throw std::runtime_error("state capacity exceeded"); + const int32_t current = size++; + length_.mutable_data()[idx(b, state_capacity_, current)] = + length_.data()[idx(b, state_capacity_, last)] + 1; + int32_t state = last; + while (state != -1 && find_transition(b, state, token) == -1) { + edge_count = add_transition(b, edge_count, state, token, current); + state = suffix_link_.data()[idx(b, state_capacity_, state)]; + } + if (state == -1) { + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = 0; + link_parent(b, current, 0); + } else { + int32_t transition = find_transition(b, state, token); + const int32_t target = edge_target_.data()[idx(b, edge_capacity_, transition)]; + if (length_.data()[idx(b, state_capacity_, state)] + 1 == + length_.data()[idx(b, state_capacity_, target)]) { + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = target; + link_parent(b, current, target); + } else { + if (size >= state_capacity_) throw std::runtime_error("state capacity exceeded"); + const int32_t clone = size++; + length_.mutable_data()[idx(b, state_capacity_, clone)] = + length_.data()[idx(b, state_capacity_, state)] + 1; + const int32_t old_parent = suffix_link_.data()[idx(b, state_capacity_, target)]; + suffix_link_.mutable_data()[idx(b, state_capacity_, clone)] = old_parent; + value_.mutable_data()[idx(b, state_capacity_, clone)] = point_query(b, target); + int32_t edge = head_.data()[idx(b, state_capacity_, target)]; + while (edge != -1) { + edge_count = add_transition( + b, edge_count, clone, edge_token_.data()[idx(b, edge_capacity_, edge)], + edge_target_.data()[idx(b, edge_capacity_, edge)]); + edge = edge_next_.data()[idx(b, edge_capacity_, edge)]; + } + transition = find_transition(b, state, token); + while (state != -1 && transition != -1 && + edge_target_.data()[idx(b, edge_capacity_, transition)] == target) { + replace_transition(b, state, token, clone); + state = suffix_link_.data()[idx(b, state_capacity_, state)]; + if (state != -1) transition = find_transition(b, state, token); + } + link_parent(b, clone, old_parent); + cut_parent(b, target); + suffix_link_.mutable_data()[idx(b, state_capacity_, target)] = clone; + link_parent(b, target, clone); + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = clone; + link_parent(b, current, clone); + } + } + last = current; + const int32_t matched = suffix_link_.data()[idx(b, state_capacity_, current)]; + int64_t source = -1; + if (matched != 0) source = point_query(b, matched); + int64_t prediction = -1; + if (source >= 0) prediction = history_.data()[idx(b, max_length_, source + 1)]; + path_assign(b, current, position_); + last_a[b] = last; size_a[b] = size; edge_count_a[b] = edge_count; + return prediction; + } + + py::object state_; + py::array_t history_, edge_token_, hash_token_, + value_, lazy_; + py::array_t head_, edge_target_, edge_next_, + hash_state_, hash_edge_, suffix_link_, length_, left_, right_, parent_, stack_, + last_, size_, edge_count_; + py::array_t lazy_valid_; + int64_t batch_, max_length_, position_, state_capacity_, edge_capacity_, hash_capacity_; +}; + +PYBIND11_MODULE(rosa_native_step, m) { + m.doc() = "Exact CPU SAM+LCT step prototype (no libtorch calls in core)"; + py::class_(m, "NativeState") + .def(py::init(), py::keep_alive<1, 2>()) + .def("step", &NativeState::step) + .def_property_readonly("position", &NativeState::position); +} diff --git a/native/tests/smoke.py b/native/tests/smoke.py new file mode 100644 index 0000000..22d492e --- /dev/null +++ b/native/tests/smoke.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import numpy as np +import rosa_native_step +import torch + +from rosa._stateful_numba import _forward_step, _init_inference_state, _prefill + + +def main() -> None: + tokens = torch.tensor( + [[0, 1, 0, 1, 2, 0, 1, 0, 1, 3, -1, 2**31, -1, 7, -1, 7]] * 2, + dtype=torch.long, + ) + oracle = _init_inference_state(tokens.shape[0], tokens.shape[1]) + candidate = _init_inference_state(tokens.shape[0], tokens.shape[1]) + oracle.native_state = False + + split = 6 + assert torch.equal( + _prefill(oracle, tokens[:, :split]), + _prefill(candidate, tokens[:, :split]), + ) + for position in range(split, tokens.shape[1]): + expected = _forward_step(oracle, tokens[:, position]) + actual = _forward_step(candidate, tokens[:, position]) + assert torch.equal(actual, expected), (position, actual, expected) + + assert isinstance(candidate.native_state, rosa_native_step.NativeState) + assert candidate.native_state.position == tokens.shape[1] + assert candidate.position == tokens.shape[1] + + malformed = _init_inference_state(2, 4) + malformed.edge_target = np.empty((2, 0), dtype=np.int32) + try: + rosa_native_step.NativeState(malformed) + except ValueError as error: + assert "layout" in str(error) + else: + raise AssertionError("malformed native state was accepted") + print("rosa_native_step smoke: ok") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 32cc8fc..9c256df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "uv_build" [project] name = "rosa-torch" -version = "0.1.1" +version = "0.2.0" description = "Independent PyTorch implementation of RWKV-8 ROSA with exact suffix-automaton retrieval" readme = "README.md" requires-python = ">=3.10" diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 213d18e..10a5cd3 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -998,6 +998,7 @@ class _StatefulInferenceState: batch_size: int max_length: int + native_abi_version: int position: int history: np.ndarray head: np.ndarray @@ -1044,6 +1045,7 @@ def _init_inference_state( return _StatefulInferenceState( batch_size=batch_size, max_length=max_length, + native_abi_version=1, position=0, history=np.empty((batch_size, max_length), dtype=np.int64), head=np.full(state_shape, -1, dtype=np.int32), diff --git a/tests/test_numba_backend.py b/tests/test_numba_backend.py index 4588621..7c9d574 100644 --- a/tests/test_numba_backend.py +++ b/tests/test_numba_backend.py @@ -64,6 +64,7 @@ def test_stateful_private_validation_and_full_replay(self) -> None: _init_inference_state(1, 0) state = _init_inference_state(1, 1) + state.native_state = False self.assertEqual(_forward_step(state, torch.tensor(0)).shape, (1,)) with self.assertRaisesRegex(ValueError, "shape"): _forward_step(state, torch.tensor([1, 2])) diff --git a/uv.lock b/uv.lock index b1cb8fa..ba71646 100644 --- a/uv.lock +++ b/uv.lock @@ -774,7 +774,7 @@ nodejs = [ [[package]] name = "rosa-torch" -version = "0.1.1" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "torch" }, From a8bad508bb158d09ca5bcdadb74217458d5784fe Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:47:11 +0800 Subject: [PATCH 17/29] Build stateful prefill in native code --- native/src/rosa_native_step.cpp | 451 ++++++++++++++++++++++++++------ native/tests/smoke.py | 98 ++++++- src/rosa/_stateful_numba.py | 25 ++ 3 files changed, 499 insertions(+), 75 deletions(-) diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp index f4e6b19..728185f 100644 --- a/native/src/rosa_native_step.cpp +++ b/native/src/rosa_native_step.cpp @@ -1,17 +1,20 @@ #include #include +#include #include #include #include +#include namespace py = pybind11; class NativeState { - public: +public: explicit NativeState(py::object state) : state_(std::move(state)) { const int64_t abi = py::cast(state_.attr("native_abi_version")); - if (abi != 1) throw py::value_error("unsupported native state ABI"); + if (abi != 1) + throw py::value_error("unsupported native state ABI"); history_ = bind("history"); head_ = bind("head"); edge_token_ = bind("edge_token"); @@ -41,8 +44,8 @@ class NativeState { validate_shapes(); } - py::array_t step( - py::array_t tokens) { + py::array_t + step(py::array_t tokens) { if (tokens.ndim() != 1 || tokens.shape(0) != batch_) { throw py::value_error("tokens must be contiguous int64 [batch_size]"); } @@ -50,35 +53,74 @@ class NativeState { throw std::runtime_error("inference state capacity exceeded"); } py::array_t output(batch_); - const int64_t* in = tokens.data(); - int64_t* out = output.mutable_data(); + const int64_t *in = tokens.data(); + int64_t *out = output.mutable_data(); { py::gil_scoped_release release; - for (int64_t b = 0; b < batch_; ++b) out[b] = step_row(b, in[b]); + for (int64_t b = 0; b < batch_; ++b) + out[b] = step_row(b, in[b]); } ++position_; state_.attr("position") = py::int_(position_); return output; } + py::array_t prefill(py::array tokens_object) { + if (!py::isinstance>(tokens_object)) { + throw py::type_error("tokens must have dtype int64"); + } + if ((tokens_object.flags() & py::array::c_style) == 0) { + throw py::value_error("tokens must be C-contiguous"); + } + auto tokens = + py::cast>(tokens_object); + if (tokens.ndim() != 2 || tokens.shape(0) != batch_) { + throw py::value_error( + "tokens must be contiguous int64 [batch_size, sequence_length]"); + } + if (position_ != 0) { + throw std::runtime_error("prefill requires an empty inference state"); + } + const int64_t token_count = tokens.shape(1); + if (token_count > max_length_) { + throw std::runtime_error("inference state capacity exceeded"); + } + py::array_t output({batch_, token_count}); + if (token_count == 0) + return output; + const int64_t *in = tokens.data(); + int64_t *out = output.mutable_data(); + { + py::gil_scoped_release release; + for (int64_t b = 0; b < batch_; ++b) { + prefill_row(b, in + b * token_count, token_count, + out + b * token_count); + } + } + position_ = token_count; + state_.attr("position") = py::int_(position_); + return output; + } + int64_t position() const { return position_; } - private: +private: template - py::array_t bind(const char* name) { + py::array_t bind(const char *name) { py::object object = state_.attr(name); if (!py::isinstance>(object)) { throw py::type_error(std::string(name) + " has an unexpected dtype"); } auto array = py::cast>(object); - if (!array.writeable()) throw py::value_error(std::string(name) + " is readonly"); + if (!array.writeable()) + throw py::value_error(std::string(name) + " is readonly"); return array; } void validate_shapes() { if (batch_ <= 0 || max_length_ <= 0 || position_ < 0 || - position_ > max_length_ || state_capacity_ <= 0 || edge_capacity_ <= 0 || - hash_capacity_ <= 0) { + position_ > max_length_ || state_capacity_ <= 0 || + edge_capacity_ <= 0 || hash_capacity_ <= 0) { throw py::value_error("incompatible _StatefulInferenceState layout"); } if (!matrix_shape(history_, batch_, max_length_) || @@ -115,14 +157,14 @@ class NativeState { } template - bool matrix_shape(const py::array_t& array, + bool matrix_shape(const py::array_t &array, int64_t rows, int64_t columns) const { return array.ndim() == 2 && array.shape(0) == rows && array.shape(1) == columns; } template - bool vector_shape(const py::array_t& array, + bool vector_shape(const py::array_t &array, int64_t length) const { return array.ndim() == 1 && array.shape(0) == length; } @@ -140,37 +182,41 @@ class NativeState { } int32_t find_transition(int64_t b, int32_t state, int64_t token) const { - const int32_t* hs = hash_state_.data(); - const int64_t* ht = hash_token_.data(); - const int32_t* he = hash_edge_.data(); - int64_t slot = static_cast(transition_hash(state, token) & - static_cast(hash_capacity_ - 1)); + const int32_t *hs = hash_state_.data(); + const int64_t *ht = hash_token_.data(); + const int32_t *he = hash_edge_.data(); + int64_t slot = + static_cast(transition_hash(state, token) & + static_cast(hash_capacity_ - 1)); while (hs[idx(b, hash_capacity_, slot)] != -1) { const int64_t at = idx(b, hash_capacity_, slot); - if (hs[at] == state && ht[at] == token) return he[at]; + if (hs[at] == state && ht[at] == token) + return he[at]; slot = (slot + 1) & (hash_capacity_ - 1); } return -1; } - int32_t add_transition(int64_t b, int32_t count, int32_t state, - int64_t token, int32_t target) { - if (count >= edge_capacity_) throw std::runtime_error("transition capacity exceeded"); - int32_t* head = head_.mutable_data(); - int64_t* et = edge_token_.mutable_data(); - int32_t* eg = edge_target_.mutable_data(); - int32_t* en = edge_next_.mutable_data(); + int32_t add_transition(int64_t b, int32_t count, int32_t state, int64_t token, + int32_t target) { + if (count >= edge_capacity_) + throw std::runtime_error("transition capacity exceeded"); + int32_t *head = head_.mutable_data(); + int64_t *et = edge_token_.mutable_data(); + int32_t *eg = edge_target_.mutable_data(); + int32_t *en = edge_next_.mutable_data(); const int64_t edge_at = idx(b, edge_capacity_, count); const int64_t state_at = idx(b, state_capacity_, state); et[edge_at] = token; eg[edge_at] = target; en[edge_at] = head[state_at]; head[state_at] = count; - int64_t slot = static_cast(transition_hash(state, token) & - static_cast(hash_capacity_ - 1)); - int32_t* hs = hash_state_.mutable_data(); - int64_t* ht = hash_token_.mutable_data(); - int32_t* he = hash_edge_.mutable_data(); + int64_t slot = + static_cast(transition_hash(state, token) & + static_cast(hash_capacity_ - 1)); + int32_t *hs = hash_state_.mutable_data(); + int64_t *ht = hash_token_.mutable_data(); + int32_t *he = hash_edge_.mutable_data(); while (hs[idx(b, hash_capacity_, slot)] != -1) { const int64_t at = idx(b, hash_capacity_, slot); if (hs[at] == state && ht[at] == token) @@ -187,14 +233,15 @@ class NativeState { void replace_transition(int64_t b, int32_t state, int64_t token, int32_t target) { const int32_t edge = find_transition(b, state, token); - if (edge == -1) throw std::runtime_error("transition not found"); + if (edge == -1) + throw std::runtime_error("transition not found"); edge_target_.mutable_data()[idx(b, edge_capacity_, edge)] = target; } inline bool is_aux_root(int64_t b, int32_t node) const { - const int32_t* left = left_.data(); - const int32_t* right = right_.data(); - const int32_t* parent = parent_.data(); + const int32_t *left = left_.data(); + const int32_t *right = right_.data(); + const int32_t *parent = parent_.data(); const int64_t at = idx(b, state_capacity_, node); const int32_t p = parent[at]; return p == -1 || (left[idx(b, state_capacity_, p)] != node && @@ -202,7 +249,8 @@ class NativeState { } inline void apply(int64_t b, int32_t node, int64_t assigned) { - if (node == -1) return; + if (node == -1) + return; const int64_t at = idx(b, state_capacity_, node); value_.mutable_data()[at] = assigned; lazy_.mutable_data()[at] = assigned; @@ -220,41 +268,52 @@ class NativeState { } inline void rotate(int64_t b, int32_t node) { - int32_t* left = left_.mutable_data(); - int32_t* right = right_.mutable_data(); - int32_t* parent = parent_.mutable_data(); + int32_t *left = left_.mutable_data(); + int32_t *right = right_.mutable_data(); + int32_t *parent = parent_.mutable_data(); const auto at = [&](int32_t n) { return idx(b, state_capacity_, n); }; const int32_t p = parent[at(node)], g = parent[at(p)]; int32_t middle; if (left[at(p)] == node) { - middle = right[at(node)]; right[at(node)] = p; left[at(p)] = middle; + middle = right[at(node)]; + right[at(node)] = p; + left[at(p)] = middle; } else { - middle = left[at(node)]; left[at(node)] = p; right[at(p)] = middle; + middle = left[at(node)]; + left[at(node)] = p; + right[at(p)] = middle; } - if (middle != -1) parent[at(middle)] = p; - parent[at(p)] = node; parent[at(node)] = g; + if (middle != -1) + parent[at(middle)] = p; + parent[at(p)] = node; + parent[at(node)] = g; if (g != -1) { - if (left[at(g)] == p) left[at(g)] = node; - else if (right[at(g)] == p) right[at(g)] = node; + if (left[at(g)] == p) + left[at(g)] = node; + else if (right[at(g)] == p) + right[at(g)] = node; } } void splay(int64_t b, int32_t node) { - int32_t* stack = stack_.mutable_data() + b * state_capacity_; + int32_t *stack = stack_.mutable_data() + b * state_capacity_; int32_t depth = 0, ancestor = node; stack[depth++] = ancestor; while (!is_aux_root(b, ancestor)) { ancestor = parent_.data()[idx(b, state_capacity_, ancestor)]; stack[depth++] = ancestor; } - while (depth > 0) push(b, stack[--depth]); + while (depth > 0) + push(b, stack[--depth]); while (!is_aux_root(b, node)) { const int32_t p = parent_.data()[idx(b, state_capacity_, node)]; if (!is_aux_root(b, p)) { const int32_t g = parent_.data()[idx(b, state_capacity_, p)]; if ((left_.data()[idx(b, state_capacity_, p)] == node) == - (left_.data()[idx(b, state_capacity_, g)] == p)) rotate(b, p); - else rotate(b, node); + (left_.data()[idx(b, state_capacity_, g)] == p)) + rotate(b, p); + else + rotate(b, node); } rotate(b, node); } @@ -265,7 +324,8 @@ class NativeState { while (current != -1) { splay(b, current); right_.mutable_data()[idx(b, state_capacity_, current)] = last; - if (last != -1) parent_.mutable_data()[idx(b, state_capacity_, last)] = current; + if (last != -1) + parent_.mutable_data()[idx(b, state_capacity_, last)] = current; last = current; current = parent_.data()[idx(b, state_capacity_, current)]; } @@ -277,14 +337,16 @@ class NativeState { return value_.data()[idx(b, state_capacity_, node)]; } void path_assign(int64_t b, int32_t node, int64_t assigned) { - access(b, node); apply(b, node, assigned); + access(b, node); + apply(b, node, assigned); } void cut_parent(int64_t b, int32_t node) { access(b, node); const int64_t at = idx(b, state_capacity_, node); const int32_t ancestors = left_.data()[at]; left_.mutable_data()[at] = -1; - if (ancestors != -1) parent_.mutable_data()[idx(b, state_capacity_, ancestors)] = -1; + if (ancestors != -1) + parent_.mutable_data()[idx(b, state_capacity_, ancestors)] = -1; } void link_parent(int64_t b, int32_t node, int32_t represented_parent) { access(b, node); @@ -292,12 +354,13 @@ class NativeState { } int64_t step_row(int64_t b, int64_t token) { - int32_t* last_a = last_.mutable_data(); - int32_t* size_a = size_.mutable_data(); - int32_t* edge_count_a = edge_count_.mutable_data(); + int32_t *last_a = last_.mutable_data(); + int32_t *size_a = size_.mutable_data(); + int32_t *edge_count_a = edge_count_.mutable_data(); int32_t last = last_a[b], size = size_a[b], edge_count = edge_count_a[b]; history_.mutable_data()[idx(b, max_length_, position_)] = token; - if (size >= state_capacity_) throw std::runtime_error("state capacity exceeded"); + if (size >= state_capacity_) + throw std::runtime_error("state capacity exceeded"); const int32_t current = size++; length_.mutable_data()[idx(b, state_capacity_, current)] = length_.data()[idx(b, state_capacity_, last)] + 1; @@ -311,32 +374,40 @@ class NativeState { link_parent(b, current, 0); } else { int32_t transition = find_transition(b, state, token); - const int32_t target = edge_target_.data()[idx(b, edge_capacity_, transition)]; + const int32_t target = + edge_target_.data()[idx(b, edge_capacity_, transition)]; if (length_.data()[idx(b, state_capacity_, state)] + 1 == length_.data()[idx(b, state_capacity_, target)]) { suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = target; link_parent(b, current, target); } else { - if (size >= state_capacity_) throw std::runtime_error("state capacity exceeded"); + if (size >= state_capacity_) + throw std::runtime_error("state capacity exceeded"); const int32_t clone = size++; length_.mutable_data()[idx(b, state_capacity_, clone)] = length_.data()[idx(b, state_capacity_, state)] + 1; - const int32_t old_parent = suffix_link_.data()[idx(b, state_capacity_, target)]; - suffix_link_.mutable_data()[idx(b, state_capacity_, clone)] = old_parent; - value_.mutable_data()[idx(b, state_capacity_, clone)] = point_query(b, target); + const int32_t old_parent = + suffix_link_.data()[idx(b, state_capacity_, target)]; + suffix_link_.mutable_data()[idx(b, state_capacity_, clone)] = + old_parent; + value_.mutable_data()[idx(b, state_capacity_, clone)] = + point_query(b, target); int32_t edge = head_.data()[idx(b, state_capacity_, target)]; while (edge != -1) { - edge_count = add_transition( - b, edge_count, clone, edge_token_.data()[idx(b, edge_capacity_, edge)], - edge_target_.data()[idx(b, edge_capacity_, edge)]); + edge_count = + add_transition(b, edge_count, clone, + edge_token_.data()[idx(b, edge_capacity_, edge)], + edge_target_.data()[idx(b, edge_capacity_, edge)]); edge = edge_next_.data()[idx(b, edge_capacity_, edge)]; } transition = find_transition(b, state, token); while (state != -1 && transition != -1 && - edge_target_.data()[idx(b, edge_capacity_, transition)] == target) { + edge_target_.data()[idx(b, edge_capacity_, transition)] == + target) { replace_transition(b, state, token, clone); state = suffix_link_.data()[idx(b, state_capacity_, state)]; - if (state != -1) transition = find_transition(b, state, token); + if (state != -1) + transition = find_transition(b, state, token); } link_parent(b, clone, old_parent); cut_parent(b, target); @@ -347,24 +418,255 @@ class NativeState { } } last = current; - const int32_t matched = suffix_link_.data()[idx(b, state_capacity_, current)]; + const int32_t matched = + suffix_link_.data()[idx(b, state_capacity_, current)]; int64_t source = -1; - if (matched != 0) source = point_query(b, matched); + if (matched != 0) + source = point_query(b, matched); int64_t prediction = -1; - if (source >= 0) prediction = history_.data()[idx(b, max_length_, source + 1)]; + if (source >= 0) + prediction = history_.data()[idx(b, max_length_, source + 1)]; path_assign(b, current, position_); - last_a[b] = last; size_a[b] = size; edge_count_a[b] = edge_count; + last_a[b] = last; + size_a[b] = size; + edge_count_a[b] = edge_count; return prediction; } + void prefill_row(int64_t b, const int64_t *tokens, int64_t token_count, + int64_t *output) { + std::fill(output, output + token_count, int64_t{-1}); + std::vector prefix_state(token_count); + int32_t last = 0, size = 1, edge_count = 0; + + // Build exactly the same final suffix automaton as the bulk Numba path, + // deliberately postponing all link-cut-tree work until the final tree is + // known. + for (int64_t position = 0; position < token_count; ++position) { + const int64_t token = tokens[position]; + history_.mutable_data()[idx(b, max_length_, position)] = token; + if (size >= state_capacity_) + throw std::runtime_error("suffix automaton state capacity exceeded"); + const int32_t current = size++; + length_.mutable_data()[idx(b, state_capacity_, current)] = + length_.data()[idx(b, state_capacity_, last)] + 1; + int32_t state = last; + while (state != -1 && find_transition(b, state, token) == -1) { + edge_count = add_transition(b, edge_count, state, token, current); + state = suffix_link_.data()[idx(b, state_capacity_, state)]; + } + if (state == -1) { + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = 0; + } else { + int32_t transition = find_transition(b, state, token); + const int32_t target = + edge_target_.data()[idx(b, edge_capacity_, transition)]; + if (length_.data()[idx(b, state_capacity_, state)] + 1 == + length_.data()[idx(b, state_capacity_, target)]) { + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = + target; + } else { + if (size >= state_capacity_) + throw std::runtime_error( + "suffix automaton state capacity exceeded"); + const int32_t clone = size++; + length_.mutable_data()[idx(b, state_capacity_, clone)] = + length_.data()[idx(b, state_capacity_, state)] + 1; + suffix_link_.mutable_data()[idx(b, state_capacity_, clone)] = + suffix_link_.data()[idx(b, state_capacity_, target)]; + int32_t edge = head_.data()[idx(b, state_capacity_, target)]; + while (edge != -1) { + edge_count = add_transition( + b, edge_count, clone, + edge_token_.data()[idx(b, edge_capacity_, edge)], + edge_target_.data()[idx(b, edge_capacity_, edge)]); + edge = edge_next_.data()[idx(b, edge_capacity_, edge)]; + } + transition = find_transition(b, state, token); + while (state != -1 && transition != -1 && + edge_target_.data()[idx(b, edge_capacity_, transition)] == + target) { + replace_transition(b, state, token, clone); + state = suffix_link_.data()[idx(b, state_capacity_, state)]; + if (state != -1) + transition = find_transition(b, state, token); + } + suffix_link_.mutable_data()[idx(b, state_capacity_, target)] = clone; + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = clone; + } + } + last = current; + prefix_state[position] = current; + } + + std::vector first_child(size, -1), next_sibling(size, -1); + for (int32_t node = 1; node < size; ++node) { + const int32_t p = suffix_link_.data()[idx(b, state_capacity_, node)]; + next_sibling[node] = first_child[p]; + first_child[p] = node; + } + std::vector tin(size), tout(size), euler_node(size), + dfs_nodes(size), dfs_next(size); + int32_t depth = 0, timer = 0; + dfs_nodes[0] = 0; + dfs_next[0] = first_child[0]; + tin[0] = timer; + euler_node[timer++] = 0; + while (depth >= 0) { + const int32_t child = dfs_next[depth]; + if (child == -1) { + tout[dfs_nodes[depth]] = timer; + --depth; + } else { + dfs_next[depth] = next_sibling[child]; + ++depth; + dfs_nodes[depth] = child; + dfs_next[depth] = first_child[child]; + tin[child] = timer; + euler_node[timer++] = child; + } + } + + int32_t levels = 1; + for (int32_t span = 1; span < size; span <<= 1) + ++levels; + std::vector up(static_cast(levels) * size, -1); + const auto up_at = [size](int32_t level, int32_t node) { + return static_cast(level) * size + node; + }; + for (int32_t node = 0; node < size; ++node) + up[up_at(0, node)] = suffix_link_.data()[idx(b, state_capacity_, node)]; + for (int32_t level = 1; level < levels; ++level) { + for (int32_t node = 0; node < size; ++node) { + const int32_t ancestor = up[up_at(level - 1, node)]; + if (ancestor != -1) + up[up_at(level, node)] = up[up_at(level - 1, ancestor)]; + } + } + const auto lca = [&](int32_t first, int32_t second) { + if (tin[first] <= tin[second] && tout[second] <= tout[first]) + return first; + if (tin[second] <= tin[first] && tout[first] <= tout[second]) + return second; + int32_t current = first; + for (int32_t level = levels - 1; level >= 0; --level) { + const int32_t ancestor = up[up_at(level, current)]; + if (ancestor != -1 && + !(tin[ancestor] <= tin[second] && tout[second] <= tout[ancestor])) + current = ancestor; + } + return up[up_at(0, current)]; + }; + + int32_t base = 1; + while (base < size) + base <<= 1; + std::vector active(static_cast(base) * 2, -1); + std::vector fenwick(size + 1, 0); + const auto fenwick_prefix = [&](int32_t index) { + int32_t total = 0; + while (index > 0) { + total += fenwick[index]; + index -= index & -index; + } + return total; + }; + const auto fenwick_select = [&](int32_t rank) { + int32_t node = 0, step = 1; + while ((step << 1) <= size) + step <<= 1; + while (step != 0) { + const int32_t candidate = node + step; + if (candidate <= size && fenwick[candidate] < rank) { + node = candidate; + rank -= fenwick[candidate]; + } + step >>= 1; + } + return node; + }; + const auto range_max = [&](int32_t left, int32_t right) { + int64_t result = -1; + for (left += base, right += base; left < right; left >>= 1, right >>= 1) { + if (left & 1) + result = std::max(result, active[left++]); + if (right & 1) + result = std::max(result, active[--right]); + } + return result; + }; + + for (int64_t position = 0; position < token_count; ++position) { + int32_t node = + suffix_link_.data()[idx(b, state_capacity_, prefix_state[position])]; + if (node != -1) { + const int32_t node_tin = tin[node]; + const int32_t preceding_count = fenwick_prefix(node_tin + 1); + int32_t best = 0; + if (preceding_count > 0) + best = lca(node, euler_node[fenwick_select(preceding_count)]); + const int32_t before_count = fenwick_prefix(node_tin); + if (before_count < position) { + const int32_t candidate = + lca(node, euler_node[fenwick_select(before_count + 1)]); + if (length_.data()[idx(b, state_capacity_, candidate)] > + length_.data()[idx(b, state_capacity_, best)]) + best = candidate; + } + node = best; + const int64_t source = range_max(tin[node], tout[node]); + if (node != 0 && source >= 0) + output[position] = history_.data()[idx(b, max_length_, source + 1)]; + } + int32_t tree_node = base + tin[prefix_state[position]]; + active[tree_node] = std::max(active[tree_node], position); + while ((tree_node >>= 1) != 0) + active[tree_node] = + std::max(active[tree_node << 1], active[(tree_node << 1) | 1]); + for (int32_t index = tin[prefix_state[position]] + 1; index <= size; + index += index & -index) + ++fenwick[index]; + } + + std::vector latest_end(size, -1); + for (int64_t position = 0; position < token_count; ++position) + latest_end[prefix_state[position]] = position; + std::vector counts(token_count + 1, 0), order(size); + for (int32_t node = 0; node < size; ++node) + ++counts[length_.data()[idx(b, state_capacity_, node)]]; + for (size_t i = 1; i < counts.size(); ++i) + counts[i] += counts[i - 1]; + for (int32_t node = size - 1; node >= 0; --node) { + const int32_t node_length = length_.data()[idx(b, state_capacity_, node)]; + order[--counts[node_length]] = node; + } + for (int32_t index = size - 1; index > 0; --index) { + const int32_t node = order[index]; + const int32_t p = suffix_link_.data()[idx(b, state_capacity_, node)]; + latest_end[p] = std::max(latest_end[p], latest_end[node]); + } + for (int32_t node = 0; node < size; ++node) { + const int64_t at = idx(b, state_capacity_, node); + left_.mutable_data()[at] = -1; + right_.mutable_data()[at] = -1; + parent_.mutable_data()[at] = suffix_link_.data()[at]; + value_.mutable_data()[at] = latest_end[node]; + lazy_valid_.mutable_data()[at] = 0; + } + last_.mutable_data()[b] = last; + size_.mutable_data()[b] = size; + edge_count_.mutable_data()[b] = edge_count; + } + py::object state_; py::array_t history_, edge_token_, hash_token_, value_, lazy_; py::array_t head_, edge_target_, edge_next_, - hash_state_, hash_edge_, suffix_link_, length_, left_, right_, parent_, stack_, - last_, size_, edge_count_; + hash_state_, hash_edge_, suffix_link_, length_, left_, right_, parent_, + stack_, last_, size_, edge_count_; py::array_t lazy_valid_; - int64_t batch_, max_length_, position_, state_capacity_, edge_capacity_, hash_capacity_; + int64_t batch_, max_length_, position_, state_capacity_, edge_capacity_, + hash_capacity_; }; PYBIND11_MODULE(rosa_native_step, m) { @@ -372,5 +674,6 @@ PYBIND11_MODULE(rosa_native_step, m) { py::class_(m, "NativeState") .def(py::init(), py::keep_alive<1, 2>()) .def("step", &NativeState::step) + .def("prefill", &NativeState::prefill) .def_property_readonly("position", &NativeState::position); } diff --git a/native/tests/smoke.py b/native/tests/smoke.py index 22d492e..2e19c9a 100644 --- a/native/tests/smoke.py +++ b/native/tests/smoke.py @@ -7,6 +7,49 @@ from rosa._stateful_numba import _forward_step, _init_inference_state, _prefill +def assert_same_initialized_state(oracle: object, candidate: object) -> None: + assert oracle.position == candidate.position + for batch in range(oracle.batch_size): + size = int(oracle.size[batch]) + edges = int(oracle.edge_count[batch]) + assert int(candidate.size[batch]) == size + assert int(candidate.edge_count[batch]) == edges + assert int(candidate.last[batch]) == int(oracle.last[batch]) + for name in ("history",): + assert np.array_equal( + getattr(candidate, name)[batch, : oracle.position], + getattr(oracle, name)[batch, : oracle.position], + ) + for name in ("head", "hash_state"): + assert np.array_equal( + getattr(candidate, name)[batch], getattr(oracle, name)[batch] + ) + for name in ("edge_token", "edge_target", "edge_next"): + assert np.array_equal( + getattr(candidate, name)[batch, :edges], + getattr(oracle, name)[batch, :edges], + ) + occupied = oracle.hash_state[batch] != -1 + for name in ("hash_token", "hash_edge"): + assert np.array_equal( + getattr(candidate, name)[batch, occupied], + getattr(oracle, name)[batch, occupied], + ) + for name in ( + "suffix_link", + "length", + "lct_left", + "lct_right", + "lct_parent", + "lct_value", + "lct_lazy_valid", + ): + assert np.array_equal( + getattr(candidate, name)[batch, :size], + getattr(oracle, name)[batch, :size], + ), name + + def main() -> None: tokens = torch.tensor( [[0, 1, 0, 1, 2, 0, 1, 0, 1, 3, -1, 2**31, -1, 7, -1, 7]] * 2, @@ -17,9 +60,12 @@ def main() -> None: oracle.native_state = False split = 6 + candidate.native_state = rosa_native_step.NativeState(candidate) assert torch.equal( _prefill(oracle, tokens[:, :split]), - _prefill(candidate, tokens[:, :split]), + torch.from_numpy( + candidate.native_state.prefill(tokens[:, :split].contiguous().numpy()) + ), ) for position in range(split, tokens.shape[1]): expected = _forward_step(oracle, tokens[:, position]) @@ -30,6 +76,56 @@ def main() -> None: assert candidate.native_state.position == tokens.shape[1] assert candidate.position == tokens.shape[1] + generator = torch.Generator().manual_seed(20260811) + cases = [ + torch.randint(-3, 9, (3, 257), generator=generator, dtype=torch.long), + torch.tensor([[1, 2, 1, 2] * 64, [7] * 256], dtype=torch.long), + ] + for case in cases: + oracle = _init_inference_state(case.shape[0], case.shape[1] + 8) + candidate = _init_inference_state(case.shape[0], case.shape[1] + 8) + oracle.native_state = False + expected = _prefill(oracle, case) + actual = _prefill(candidate, case) + assert torch.equal(actual, expected) + assert isinstance(candidate.native_state, rosa_native_step.NativeState) + assert_same_initialized_state(oracle, candidate) + # The offline-built LCT must be immediately usable by streaming. + continuation = torch.randint( + -3, 9, (case.shape[0], 8), generator=generator, dtype=torch.long + ) + for position in range(continuation.shape[1]): + expected_step = _forward_step(oracle, continuation[:, position]) + actual_step = _forward_step(candidate, continuation[:, position]) + assert torch.equal(actual_step, expected_step), position + + validation = _init_inference_state(2, 4) + native_validation = rosa_native_step.NativeState(validation) + for invalid in ( + np.zeros((2, 2), dtype=np.int32), + np.zeros((2, 2), dtype=np.int64)[:, ::2], + np.zeros((3, 2), dtype=np.int64), + ): + try: + native_validation.prefill(invalid) + except (TypeError, ValueError): + pass + else: + raise AssertionError("invalid prefill input was accepted") + try: + native_validation.prefill(np.zeros((2, 5), dtype=np.int64)) + except RuntimeError as error: + assert "capacity" in str(error) + else: + raise AssertionError("over-capacity prefill was accepted") + native_validation.prefill(np.zeros((2, 1), dtype=np.int64)) + try: + native_validation.prefill(np.zeros((2, 1), dtype=np.int64)) + except RuntimeError as error: + assert "empty" in str(error) + else: + raise AssertionError("prefill accepted a non-empty state") + malformed = _init_inference_state(2, 4) malformed.edge_target = np.empty((2, 0), dtype=np.int32) try: diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 10a5cd3..510d63d 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -1089,6 +1089,28 @@ def _native_step( # pragma: no cover - optional native companion return state.native_state.step(cpu_tokens.numpy()) +def _native_prefill( # pragma: no cover - optional native companion + state: _StatefulInferenceState, + cpu_tokens: Tensor, +) -> np.ndarray | None: + if state.native_state is False: + return None + if state.native_state is None: + try: + from rosa_native_step import ( # type: ignore[reportMissingImports] + NativeState, + ) + except ModuleNotFoundError: + state.native_state = False + return None + state.native_state = NativeState(state) + native_prefill = getattr(state.native_state, "prefill", None) + if native_prefill is None: + state.native_state = False + return None + return native_prefill(cpu_tokens.numpy()) + + def _forward_step(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: """Consume one token per batch row and return exact top-1 predictions.""" @@ -1144,6 +1166,9 @@ def _prefill(state: _StatefulInferenceState, tokens: Tensor) -> Tensor: if tokens.shape[1] == 0: return torch.empty(tokens.shape, dtype=torch.long, device=device) cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + native_output = _native_prefill(state, cpu_tokens) + if native_output is not None: # pragma: no cover - optional native companion + return torch.from_numpy(native_output).to(device) output_array = _bulk_prefill_kernel( cpu_tokens.numpy(), state.history, From 2555552ae778facb478562bf3cd8007cfab38b4b Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:50:41 +0800 Subject: [PATCH 18/29] Track rich stateful ROSA candidates --- src/rosa/_stateful_candidates_numba.py | 933 +++++++++++++++++++++++++ tests/test_stateful_candidates.py | 133 ++++ 2 files changed, 1066 insertions(+) create mode 100644 src/rosa/_stateful_candidates_numba.py create mode 100644 tests/test_stateful_candidates.py diff --git a/src/rosa/_stateful_candidates_numba.py b/src/rosa/_stateful_candidates_numba.py new file mode 100644 index 0000000..f5084aa --- /dev/null +++ b/src/rosa/_stateful_candidates_numba.py @@ -0,0 +1,933 @@ +"""Stateful exact top-R ROSA candidates without eager suffix propagation. + +The suffix automaton is augmented with a Link-Cut Tree. A write to the +current suffix chain is represented by a lazy tag containing the newest +bounded occurrence prefix and an unbounded frequency delta. Tag composition +is exact: a newer prefix is prepended to the older pending prefix, while the +frequency deltas are added. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch +from numba import njit +from torch import Tensor + +from ._stateful_numba import ( + _add_transition, + _find_transition, + _replace_transition, +) + + +@njit(cache=True, nogil=True, inline="always") +def _is_aux_root( + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + node: int, +) -> bool: # pragma: no cover - executed as compiled Numba code + ancestor = parent[node] + return ancestor == -1 or (left[ancestor] != node and right[ancestor] != node) + + +@njit(cache=True, nogil=True) +def _apply_tag( + left: np.ndarray, + right: np.ndarray, + occurrences: np.ndarray, + occurrence_size: np.ndarray, + frequency: np.ndarray, + lazy_prefix: np.ndarray, + lazy_size: np.ndarray, + lazy_delta: np.ndarray, + node: int, + prefix: np.ndarray, + prefix_size: int, + delta: int, +) -> None: # pragma: no cover - executed as compiled Numba code + """Apply and compose ``(newest_prefix, frequency_delta)`` at ``node``.""" + + if node == -1: + return + capacity = occurrences.shape[1] + take = min(prefix_size, capacity) + + old_size = int(occurrence_size[node]) + updated_size = min(capacity, take + old_size) + for index in range(updated_size - 1, take - 1, -1): + occurrences[node, index] = occurrences[node, index - take] + for index in range(take): + occurrences[node, index] = prefix[index] + occurrence_size[node] = updated_size + frequency[node] += delta + + old_lazy_size = int(lazy_size[node]) + updated_lazy_size = min(capacity, take + old_lazy_size) + for index in range(updated_lazy_size - 1, take - 1, -1): + lazy_prefix[node, index] = lazy_prefix[node, index - take] + for index in range(take): + lazy_prefix[node, index] = prefix[index] + lazy_size[node] = updated_lazy_size + lazy_delta[node] += delta + + +@njit(cache=True, nogil=True) +def _push( + left: np.ndarray, + right: np.ndarray, + occurrences: np.ndarray, + occurrence_size: np.ndarray, + frequency: np.ndarray, + lazy_prefix: np.ndarray, + lazy_size: np.ndarray, + lazy_delta: np.ndarray, + node: int, +) -> None: # pragma: no cover - executed as compiled Numba code + size = int(lazy_size[node]) + delta = int(lazy_delta[node]) + if size != 0 or delta != 0: + _apply_tag( + left, + right, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + int(left[node]), + lazy_prefix[node], + size, + delta, + ) + _apply_tag( + left, + right, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + int(right[node]), + lazy_prefix[node], + size, + delta, + ) + lazy_size[node] = 0 + lazy_delta[node] = 0 + + +@njit(cache=True, nogil=True, inline="always") +def _rotate( + left: np.ndarray, + right: np.ndarray, + parent: np.ndarray, + node: int, +) -> None: # pragma: no cover - executed as compiled Numba code + p = int(parent[node]) + g = int(parent[p]) + if left[p] == node: + middle = int(right[node]) + right[node] = p + left[p] = middle + else: + middle = int(left[node]) + left[node] = p + right[p] = middle + if middle != -1: + parent[middle] = p + parent[p] = node + parent[node] = g + if g != -1: + if left[g] == p: + left[g] = node + elif right[g] == p: + right[g] = node + + +@njit(cache=True, nogil=True) +def _splay( + left: np.ndarray, + right: np.ndarray, + 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, + node: int, + stack: np.ndarray, +) -> None: # pragma: no cover - executed as compiled Numba code + depth = 0 + ancestor = node + stack[depth] = ancestor + depth += 1 + while not _is_aux_root(left, right, parent, ancestor): + ancestor = int(parent[ancestor]) + stack[depth] = ancestor + depth += 1 + while depth > 0: + depth -= 1 + _push( + left, + right, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + int(stack[depth]), + ) + + while not _is_aux_root(left, right, parent, node): + p = int(parent[node]) + if not _is_aux_root(left, right, parent, p): + g = int(parent[p]) + if (left[p] == node) == (left[g] == p): + _rotate(left, right, parent, p) + else: + _rotate(left, right, parent, node) + _rotate(left, right, parent, node) + + +@njit(cache=True, nogil=True) +def _access( + left: np.ndarray, + right: np.ndarray, + 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, + node: int, + stack: np.ndarray, +) -> None: # pragma: no cover - executed as compiled Numba code + last = -1 + current = node + while current != -1: + _splay( + left, + right, + parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + current, + stack, + ) + right[current] = last + if last != -1: + parent[last] = current + last = current + current = int(parent[current]) + _splay( + left, + right, + parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + node, + stack, + ) + + +@njit(cache=True, nogil=True) +def _materialize( + left: np.ndarray, + right: np.ndarray, + 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, + node: int, + stack: np.ndarray, +) -> None: # pragma: no cover - executed as compiled Numba code + _access( + left, + right, + parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + node, + stack, + ) + + +@njit(cache=True, nogil=True) +def _cut_parent( + left: np.ndarray, + right: np.ndarray, + 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, + node: int, + stack: np.ndarray, +) -> None: # pragma: no cover - executed as compiled Numba code + _materialize( + left, + right, + parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + node, + stack, + ) + ancestors = int(left[node]) + left[node] = -1 + if ancestors != -1: + parent[ancestors] = -1 + + +@njit(cache=True, nogil=True) +def _link_parent( + left: np.ndarray, + right: np.ndarray, + 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, + node: int, + represented_parent: int, + stack: np.ndarray, +) -> None: # pragma: no cover - executed as compiled Numba code + _materialize( + left, + right, + parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + node, + stack, + ) + parent[node] = represented_parent + + +@njit(cache=True, nogil=True) +def _path_write( + left: np.ndarray, + right: np.ndarray, + 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, + node: int, + position: int, + stack: np.ndarray, +) -> None: # pragma: no cover - executed as compiled Numba code + _materialize( + left, + right, + parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + node, + stack, + ) + prefix = np.empty(1, dtype=np.int64) + prefix[0] = position + _apply_tag( + left, + right, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + node, + prefix, + 1, + 1, + ) + + +@njit(cache=True, nogil=True) +def _step_row( + token: int, + position: int, + 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: int, + size: int, + edge_count: int, + output_source: np.ndarray, + output_length: np.ndarray, + output_state: np.ndarray, + output_frequency: np.ndarray, +) -> tuple[int, int, int, int]: # pragma: no cover - compiled Numba code + history[position] = token + if size >= head.shape[0]: + raise RuntimeError("suffix automaton state capacity exceeded") + current = size + size += 1 + length[current] = length[last] + 1 + state = last + + while ( + state != -1 + and _find_transition(hash_state, hash_token, hash_edge, state, token) == -1 + ): + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + hash_state, + hash_token, + hash_edge, + edge_count, + state, + token, + current, + ) + state = int(suffix_link[state]) + + if state == -1: + suffix_link[current] = 0 + _link_parent( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + current, + 0, + lct_stack, + ) + else: + transition = _find_transition(hash_state, hash_token, hash_edge, state, token) + target = int(edge_target[transition]) + if length[state] + 1 == length[target]: + suffix_link[current] = target + _link_parent( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + current, + target, + lct_stack, + ) + else: + if size >= head.shape[0]: + raise RuntimeError("suffix automaton state capacity exceeded") + clone = size + size += 1 + length[clone] = length[state] + 1 + old_parent = int(suffix_link[target]) + suffix_link[clone] = old_parent + + # Materialize q before changing represented-tree edges. The clone + # receives q's exact bounded newest prefix and full count, but no + # pending tag because it has no represented children yet. + _materialize( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + target, + lct_stack, + ) + clone_occurrence_size = int(occurrence_size[target]) + occurrence_size[clone] = clone_occurrence_size + for index in range(clone_occurrence_size): + occurrences[clone, index] = occurrences[target, index] + frequency[clone] = frequency[target] + lazy_size[clone] = 0 + lazy_delta[clone] = 0 + + edge = int(head[target]) + while edge != -1: + edge_count = _add_transition( + head, + edge_token, + edge_target, + edge_next, + hash_state, + hash_token, + hash_edge, + edge_count, + clone, + int(edge_token[edge]), + int(edge_target[edge]), + ) + edge = int(edge_next[edge]) + + transition = _find_transition( + hash_state, hash_token, hash_edge, state, token + ) + while ( + state != -1 and transition != -1 and edge_target[transition] == target + ): + _replace_transition( + edge_target, + hash_state, + hash_token, + hash_edge, + state, + token, + clone, + ) + state = int(suffix_link[state]) + if state != -1: + transition = _find_transition( + hash_state, hash_token, hash_edge, state, token + ) + + _link_parent( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + clone, + old_parent, + lct_stack, + ) + _cut_parent( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + target, + lct_stack, + ) + suffix_link[target] = clone + _link_parent( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + target, + clone, + lct_stack, + ) + suffix_link[current] = clone + _link_parent( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + current, + clone, + lct_stack, + ) + + 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]) + + _path_write( + lct_left, + lct_right, + lct_parent, + occurrences, + occurrence_size, + frequency, + lazy_prefix, + lazy_size, + lazy_delta, + current, + position, + lct_stack, + ) + return candidate_count, last, size, edge_count + + +@njit(cache=True, nogil=True) +def _step_batch_kernel( + tokens: np.ndarray, + position: int, + 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]: + slots = suffix_k * occurrences_r + batch_size = tokens.shape[0] + source = np.full((batch_size, slots), -1, dtype=np.int64) + match_length = np.zeros((batch_size, slots), dtype=np.int64) + state_id = np.full((batch_size, slots), -1, dtype=np.int64) + candidate_frequency = np.zeros((batch_size, slots), dtype=np.int64) + count = np.zeros(batch_size, dtype=np.int32) + for batch_index in range(batch_size): + row_count, row_last, row_size, row_edge_count = _step_row( + int(tokens[batch_index]), + 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]), + source[batch_index], + match_length[batch_index], + state_id[batch_index], + candidate_frequency[batch_index], + ) + count[batch_index] = row_count + 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.""" + + batch_size: int + max_length: int + suffix_k: int + occurrences_r: int + position: 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 + + +@dataclass(frozen=True) +class CandidateStep: + """Exact hard candidates emitted while consuming one token per row.""" + + source_index: Tensor + match_length: Tensor + state_id: Tensor + frequency: Tensor + mask: Tensor + rosa_slot: Tensor + rosa_source_index: Tensor + rosa_match_length: Tensor + rosa_predicted_tokens: Tensor + + +def init_candidate_state( + batch_size: int, + max_length: int, + *, + suffix_k: int = 16, + occurrences_r: int = 4, +) -> CandidateState: + """Allocate an exact bounded-candidate state backed by CPU tensors.""" + + if batch_size <= 0: + raise ValueError("batch_size must be > 0") + if max_length <= 0: + raise ValueError("max_length must be > 0") + if suffix_k <= 0: + raise ValueError("suffix_k must be > 0") + if occurrences_r <= 0: + raise ValueError("occurrences_r must be > 0") + max_states = 2 * max_length + 1 + max_edges = 4 * max_length + 1 + hash_capacity = 1 << (2 * max_edges - 1).bit_length() + state_shape = (batch_size, max_states) + edge_shape = (batch_size, max_edges) + hash_shape = (batch_size, hash_capacity) + occurrence_shape = (batch_size, max_states, occurrences_r) + return CandidateState( + batch_size=batch_size, + max_length=max_length, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + position=0, + history=np.empty((batch_size, max_length), dtype=np.int64), + head=np.full(state_shape, -1, dtype=np.int32), + edge_token=np.empty(edge_shape, dtype=np.int64), + edge_target=np.empty(edge_shape, dtype=np.int32), + edge_next=np.empty(edge_shape, dtype=np.int32), + hash_state=np.full(hash_shape, -1, dtype=np.int32), + hash_token=np.empty(hash_shape, dtype=np.int64), + hash_edge=np.empty(hash_shape, dtype=np.int32), + suffix_link=np.full(state_shape, -1, dtype=np.int32), + length=np.zeros(state_shape, dtype=np.int32), + lct_left=np.full(state_shape, -1, dtype=np.int32), + lct_right=np.full(state_shape, -1, dtype=np.int32), + lct_parent=np.full(state_shape, -1, dtype=np.int32), + occurrences=np.full(occurrence_shape, -1, dtype=np.int64), + occurrence_size=np.zeros(state_shape, dtype=np.int32), + frequency=np.zeros(state_shape, dtype=np.int64), + lazy_prefix=np.full(occurrence_shape, -1, dtype=np.int64), + lazy_size=np.zeros(state_shape, dtype=np.int32), + lazy_delta=np.zeros(state_shape, dtype=np.int64), + lct_stack=np.empty(state_shape, dtype=np.int32), + last=np.zeros(batch_size, dtype=np.int32), + size=np.ones(batch_size, dtype=np.int32), + edge_count=np.zeros(batch_size, dtype=np.int32), + ) + + +_INTEGER_DTYPES = { + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, +} + + +def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateStep: + """Consume one token per row and return exact top-R candidates for K suffixes.""" + + if not isinstance(state, CandidateState): + raise TypeError("state must be a CandidateState") + if not isinstance(tokens, Tensor): + raise TypeError("tokens must be a Tensor") + if tokens.ndim == 0 and state.batch_size == 1: + tokens = tokens.unsqueeze(0) + if tokens.ndim != 1 or tokens.shape[0] != state.batch_size: + raise ValueError("tokens must have shape [batch_size]") + if tokens.dtype not in _INTEGER_DTYPES: + raise TypeError("tokens must use an integer dtype") + if state.position >= state.max_length: + raise RuntimeError("candidate state capacity exceeded") + + device = tokens.device + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + source, match_length, state_id, frequency, count = _step_batch_kernel( + cpu_tokens.numpy(), + state.position, + 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 += 1 + + slots = state.suffix_k * state.occurrences_r + slot_index = np.arange(slots, dtype=np.int32)[None, :] + mask = slot_index < count[:, None] + rosa_source = source[:, 0].copy() + rosa_length = match_length[:, 0].copy() + rosa_slot = np.where(count > 0, 0, -1).astype(np.int64) + rosa_predicted = np.full(state.batch_size, -1, dtype=np.int64) + for batch_index in range(state.batch_size): + if count[batch_index] > 0: + source_position = int(rosa_source[batch_index]) + rosa_predicted[batch_index] = state.history[ + batch_index, source_position + 1 + ] + + return CandidateStep( + source_index=torch.from_numpy(source).to(device), + match_length=torch.from_numpy(match_length).to(device), + state_id=torch.from_numpy(state_id).to(device), + frequency=torch.from_numpy(frequency).to(device), + mask=torch.from_numpy(mask).to(device), + rosa_slot=torch.from_numpy(rosa_slot).to(device), + rosa_source_index=torch.from_numpy(rosa_source).to(device), + rosa_match_length=torch.from_numpy(rosa_length).to(device), + rosa_predicted_tokens=torch.from_numpy(rosa_predicted).to(device), + ) diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py new file mode 100644 index 0000000..d4dbbb8 --- /dev/null +++ b/tests/test_stateful_candidates.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import unittest +from itertools import product + +import torch + +from rosa import build_hard_candidates +from rosa._stateful_candidates_numba import ( + CandidateState, + CandidateStep, + forward_candidates_step, + init_candidate_state, +) + +_FIELDS = ( + "source_index", + "match_length", + "state_id", + "frequency", + "mask", + "rosa_slot", + "rosa_source_index", + "rosa_match_length", + "rosa_predicted_tokens", +) + + +class TestStatefulCandidates(unittest.TestCase): + def assert_matches_oracle( + self, + tokens: torch.Tensor, + *, + suffix_k: int, + occurrences_r: int, + split_at: int | None = None, + ) -> CandidateState: + state = init_candidate_state( + tokens.shape[0], + tokens.shape[1], + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + steps: list[CandidateStep] = [] + boundary = tokens.shape[1] if split_at is None else split_at + for position in range(boundary): + steps.append(forward_candidates_step(state, tokens[:, position])) + # Deliberately retain and continue the same mutable state across calls. + for position in range(boundary, tokens.shape[1]): + steps.append(forward_candidates_step(state, tokens[:, position])) + + expected = build_hard_candidates( + tokens, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + for field in _FIELDS: + with self.subTest(field=field, suffix_k=suffix_k, r=occurrences_r): + actual = torch.stack([getattr(step, field) for step in steps], dim=1) + self.assertTrue(torch.equal(actual, getattr(expected, field))) + self.assertEqual(state.position, tokens.shape[1]) + return state + + def test_exhaustive_binary_clones_for_multiple_k_and_r(self) -> None: + rows = list(product(range(2), repeat=9)) + tokens = torch.tensor(rows, dtype=torch.long) + for suffix_k, occurrences_r in ((1, 1), (2, 3), (4, 2), (5, 4)): + state = self.assert_matches_oracle( + tokens, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + # Binary strings exercise SAM clone creation in aggregate. + self.assertTrue(torch.from_numpy(state.size > tokens.shape[1] + 1).any()) + + def test_newest_first_frequency_beyond_r_and_deduplication(self) -> None: + tokens = torch.tensor([[0, 1, 0, 2, 0, 3, 0, 4, 0]], dtype=torch.long) + state = init_candidate_state(1, 9, suffix_k=6, occurrences_r=3) + result: CandidateStep | None = None + for position in range(tokens.shape[1]): + result = forward_candidates_step(state, tokens[:, position]) + assert result is not None + + self.assertEqual(result.source_index[0, :3].tolist(), [6, 4, 2]) + self.assertEqual(result.frequency[0, :3].tolist(), [4, 4, 4]) + valid_sources = result.source_index[0][result.mask[0]].tolist() + self.assertEqual(len(valid_sources), len(set(valid_sources))) + self.assert_matches_oracle(tokens, suffix_k=6, occurrences_r=3) + + def test_batched_continuation_matches_offline_oracle(self) -> None: + generator = torch.Generator().manual_seed(20260811) + random_tokens = torch.randint(7, (4, 73), generator=generator) + repetitive = torch.arange(73).remainder(3).unsqueeze(0) + tokens = torch.cat((random_tokens, repetitive), dim=0) + self.assert_matches_oracle( + tokens, + suffix_k=7, + occurrences_r=5, + split_at=41, + ) + + def test_scalar_batch_and_validation(self) -> None: + state = init_candidate_state(1, 2, suffix_k=2, occurrences_r=2) + first = forward_candidates_step(state, torch.tensor(7)) + self.assertEqual(tuple(first.source_index.shape), (1, 4)) + forward_candidates_step(state, torch.tensor([7])) + with self.assertRaisesRegex(RuntimeError, "capacity"): + forward_candidates_step(state, torch.tensor([7])) + + for kwargs, message in ( + ({"batch_size": 0, "max_length": 1}, "batch_size"), + ({"batch_size": 1, "max_length": 0}, "max_length"), + ({"batch_size": 1, "max_length": 1, "suffix_k": 0}, "suffix_k"), + ( + {"batch_size": 1, "max_length": 1, "occurrences_r": 0}, + "occurrences_r", + ), + ): + with self.subTest(kwargs=kwargs): + with self.assertRaisesRegex(ValueError, message): + init_candidate_state(**kwargs) + + shape_state = init_candidate_state(2, 1) + with self.assertRaisesRegex(ValueError, "shape"): + forward_candidates_step(shape_state, torch.tensor([1])) + with self.assertRaisesRegex(TypeError, "integer"): + forward_candidates_step(shape_state, torch.tensor([1.0, 2.0])) + with self.assertRaisesRegex(TypeError, "CandidateState"): + forward_candidates_step(object(), torch.tensor([1])) # type: ignore[arg-type] + + +if __name__ == "__main__": + unittest.main() From a2c75ad21b1b064a123e8b820b42fc7dc09fc2cc Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 07:10:48 +0800 Subject: [PATCH 19/29] Support recyclable ragged inference slots --- native/src/rosa_native_step.cpp | 127 +++++++++++++++++++- native/tests/smoke.py | 78 +++++++++++++ src/rosa/_stateful_numba.py | 100 ++++++++++++++++ src/rosa/ragged.py | 199 ++++++++++++++++++++++++++++++++ tests/test_numba_backend.py | 4 + tests/test_ragged.py | 124 ++++++++++++++++++++ 6 files changed, 627 insertions(+), 5 deletions(-) create mode 100644 src/rosa/ragged.py create mode 100644 tests/test_ragged.py diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp index 728185f..957b613 100644 --- a/native/src/rosa_native_step.cpp +++ b/native/src/rosa_native_step.cpp @@ -42,10 +42,31 @@ class NativeState { edge_capacity_ = edge_token_.shape(1); hash_capacity_ = hash_state_.shape(1); validate_shapes(); + positions_ = py::array_t(batch_); + std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, + position_); + if (py::hasattr(state_, "positions")) { + ragged_mode_ = true; + py::object object = state_.attr("positions"); + if (!py::isinstance>(object)) + throw py::type_error("positions has an unexpected dtype"); + positions_ = py::cast>(object); + if (!positions_.writeable() || !vector_shape(positions_, batch_)) + throw py::value_error("positions must be writable contiguous int64 [batch_size]"); + } + occupied_slots_.resize(batch_); + for (int64_t b = 0; b < batch_; ++b) { + for (int64_t slot = 0; slot < hash_capacity_; ++slot) { + if (hash_state_.data()[idx(b, hash_capacity_, slot)] != -1) + occupied_slots_[b].push_back(static_cast(slot)); + } + } } py::array_t step(py::array_t tokens) { + if (ragged_mode_) + throw std::runtime_error("uniform step is unavailable on a ragged state"); if (tokens.ndim() != 1 || tokens.shape(0) != batch_) { throw py::value_error("tokens must be contiguous int64 [batch_size]"); } @@ -58,14 +79,66 @@ class NativeState { { py::gil_scoped_release release; for (int64_t b = 0; b < batch_; ++b) - out[b] = step_row(b, in[b]); + out[b] = step_row(b, in[b], position_); } ++position_; + std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, + position_); state_.attr("position") = py::int_(position_); return output; } + py::array_t step_masked(py::array tokens_object, + py::array active_object, + py::array reset_object) { + if (!ragged_mode_) + throw std::runtime_error("step_masked requires a ragged state"); + auto tokens = checked_vector(tokens_object, "tokens", "int64"); + const bool active_bool = py::isinstance>(active_object); + const bool active_u8 = py::isinstance>(active_object); + const bool reset_bool = py::isinstance>(reset_object); + const bool reset_u8 = py::isinstance>(reset_object); + if ((!active_bool && !active_u8) || (!reset_bool && !reset_u8)) + throw py::type_error("active and reset must have dtype bool or uint8"); + if ((active_object.flags() & py::array::c_style) == 0 || + (reset_object.flags() & py::array::c_style) == 0 || + active_object.ndim() != 1 || reset_object.ndim() != 1 || + active_object.shape(0) != batch_ || reset_object.shape(0) != batch_) + throw py::value_error( + "active and reset must be contiguous [batch_size]"); + std::vector active(batch_), reset(batch_); + for (int64_t b = 0; b < batch_; ++b) { + active[b] = active_bool + ? static_cast(active_object.data())[b] + : static_cast(active_object.data())[b] != 0; + reset[b] = reset_bool + ? static_cast(reset_object.data())[b] + : static_cast(reset_object.data())[b] != 0; + } + for (int64_t b = 0; b < batch_; ++b) { + if (active[b] && !reset[b] && positions_.data()[b] >= max_length_) + throw std::runtime_error("inference state capacity exceeded"); + } + py::array_t output(batch_); + std::fill(output.mutable_data(), output.mutable_data() + batch_, int64_t{-1}); + { + py::gil_scoped_release release; + for (int64_t b = 0; b < batch_; ++b) { + if (!active[b]) + continue; + if (reset[b]) + reset_row(b); + const int64_t position = positions_.data()[b]; + output.mutable_data()[b] = step_row(b, tokens.data()[b], position); + positions_.mutable_data()[b] = position + 1; + } + } + return output; + } + py::array_t prefill(py::array tokens_object) { + if (ragged_mode_) + throw std::runtime_error("prefill is unavailable on a ragged state"); if (!py::isinstance>(tokens_object)) { throw py::type_error("tokens must have dtype int64"); } @@ -98,13 +171,28 @@ class NativeState { } } position_ = token_count; + std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, + position_); state_.attr("position") = py::int_(position_); return output; } int64_t position() const { return position_; } + py::array_t positions() const { return positions_; } private: + template + py::array_t + checked_vector(py::array object, const char *name, const char *dtype) const { + if (!py::isinstance>(object)) + throw py::type_error(std::string(name) + " must have dtype " + dtype); + if ((object.flags() & py::array::c_style) == 0 || object.ndim() != 1 || + object.shape(0) != batch_) + throw py::value_error(std::string(name) + + " must be contiguous [batch_size]"); + return py::cast>(object); + } + template py::array_t bind(const char *name) { py::object object = state_.attr(name); @@ -227,9 +315,33 @@ class NativeState { hs[at] = state; ht[at] = token; he[at] = count; + occupied_slots_[b].push_back(static_cast(slot)); return count + 1; } + void reset_row(int64_t b) { + const int32_t used_states = size_.data()[b]; + for (int32_t state = 0; state < used_states; ++state) { + const int64_t at = idx(b, state_capacity_, state); + head_.mutable_data()[at] = -1; + suffix_link_.mutable_data()[at] = -1; + length_.mutable_data()[at] = 0; + left_.mutable_data()[at] = -1; + right_.mutable_data()[at] = -1; + parent_.mutable_data()[at] = -1; + value_.mutable_data()[at] = -1; + lazy_valid_.mutable_data()[at] = 0; + } + for (const int32_t slot : occupied_slots_[b]) { + hash_state_.mutable_data()[idx(b, hash_capacity_, slot)] = -1; + } + occupied_slots_[b].clear(); + last_.mutable_data()[b] = 0; + size_.mutable_data()[b] = 1; + edge_count_.mutable_data()[b] = 0; + positions_.mutable_data()[b] = 0; + } + void replace_transition(int64_t b, int32_t state, int64_t token, int32_t target) { const int32_t edge = find_transition(b, state, token); @@ -353,12 +465,12 @@ class NativeState { parent_.mutable_data()[idx(b, state_capacity_, node)] = represented_parent; } - int64_t step_row(int64_t b, int64_t token) { + int64_t step_row(int64_t b, int64_t token, int64_t position) { int32_t *last_a = last_.mutable_data(); int32_t *size_a = size_.mutable_data(); int32_t *edge_count_a = edge_count_.mutable_data(); int32_t last = last_a[b], size = size_a[b], edge_count = edge_count_a[b]; - history_.mutable_data()[idx(b, max_length_, position_)] = token; + history_.mutable_data()[idx(b, max_length_, position)] = token; if (size >= state_capacity_) throw std::runtime_error("state capacity exceeded"); const int32_t current = size++; @@ -426,7 +538,7 @@ class NativeState { int64_t prediction = -1; if (source >= 0) prediction = history_.data()[idx(b, max_length_, source + 1)]; - path_assign(b, current, position_); + path_assign(b, current, position); last_a[b] = last; size_a[b] = size; edge_count_a[b] = edge_count; @@ -665,6 +777,9 @@ class NativeState { hash_state_, hash_edge_, suffix_link_, length_, left_, right_, parent_, stack_, last_, size_, edge_count_; py::array_t lazy_valid_; + py::array_t positions_; + std::vector> occupied_slots_; + bool ragged_mode_ = false; int64_t batch_, max_length_, position_, state_capacity_, edge_capacity_, hash_capacity_; }; @@ -674,6 +789,8 @@ PYBIND11_MODULE(rosa_native_step, m) { py::class_(m, "NativeState") .def(py::init(), py::keep_alive<1, 2>()) .def("step", &NativeState::step) + .def("step_masked", &NativeState::step_masked) .def("prefill", &NativeState::prefill) - .def_property_readonly("position", &NativeState::position); + .def_property_readonly("position", &NativeState::position) + .def_property_readonly("positions", &NativeState::positions); } diff --git a/native/tests/smoke.py b/native/tests/smoke.py index 2e19c9a..70a1728 100644 --- a/native/tests/smoke.py +++ b/native/tests/smoke.py @@ -5,6 +5,7 @@ import torch from rosa._stateful_numba import _forward_step, _init_inference_state, _prefill +from rosa.ragged import RaggedInferenceState def assert_same_initialized_state(oracle: object, candidate: object) -> None: @@ -134,6 +135,83 @@ def main() -> None: assert "layout" in str(error) else: raise AssertionError("malformed native state was accepted") + + native_ragged = RaggedInferenceState(5, 40, use_native=True) + fallback_ragged = RaggedInferenceState(5, 40, use_native=False) + for tick in range(40): + step_tokens = torch.tensor( + [(tick * 3 + row * 5) % 11 - 2 for row in range(5)], dtype=torch.long + ) + active = torch.tensor( + [(tick + row) % (row + 2) != 0 for row in range(5)], dtype=torch.uint8 + ) + reset = torch.tensor( + [tick in (8 + row, 20 + row, 32 + row) for row in range(5)], + dtype=torch.uint8, + ) + expected = fallback_ragged.step_masked(step_tokens, active, reset) + actual = native_ragged.step_masked(step_tokens, active, reset) + assert torch.equal(actual, expected), (tick, actual, expected) + assert torch.equal(native_ragged.positions, fallback_ragged.positions) + assert native_ragged.using_native + + direct_state = _init_inference_state(2, 4) + direct_state.positions = np.zeros(2, dtype=np.int64) + direct = rosa_native_step.NativeState(direct_state) + direct_output = direct.step_masked( + np.array([7, 9], dtype=np.int64), + np.array([True, False], dtype=np.bool_), + np.array([False, True], dtype=np.bool_), + ) + assert direct_output.tolist() == [-1, -1] + assert direct.positions.tolist() == [1, 0] + try: + direct.step(np.array([1, 2], dtype=np.int64)) + except RuntimeError as error: + assert "uniform" in str(error) + else: + raise AssertionError("ragged native state accepted a uniform step") + try: + direct.prefill(np.zeros((2, 1), dtype=np.int64)) + except RuntimeError as error: + assert "ragged" in str(error) + else: + raise AssertionError("ragged native state accepted prefill") + + uniform = rosa_native_step.NativeState(_init_inference_state(2, 4)) + try: + uniform.step_masked( + np.array([1, 2], dtype=np.int64), + np.ones(2, dtype=np.uint8), + np.zeros(2, dtype=np.uint8), + ) + except RuntimeError as error: + assert "ragged" in str(error) + else: + raise AssertionError("uniform native state accepted a masked step") + for invalid_tokens, invalid_active, invalid_reset in ( + ( + np.zeros(2, dtype=np.int32), + np.ones(2, dtype=np.uint8), + np.zeros(2, dtype=np.uint8), + ), + ( + np.zeros(2, dtype=np.int64), + np.ones(2, dtype=np.int64), + np.zeros(2, dtype=np.uint8), + ), + ( + np.zeros(2, dtype=np.int64), + np.ones(3, dtype=np.uint8), + np.zeros(2, dtype=np.uint8), + ), + ): + try: + direct.step_masked(invalid_tokens, invalid_active, invalid_reset) + except (TypeError, ValueError): + pass + else: + raise AssertionError("invalid masked-step input was accepted") print("rosa_native_step smoke: ok") diff --git a/src/rosa/_stateful_numba.py b/src/rosa/_stateful_numba.py index 510d63d..9ae3ba3 100644 --- a/src/rosa/_stateful_numba.py +++ b/src/rosa/_stateful_numba.py @@ -555,6 +555,106 @@ def _step_batch_kernel( # pragma: no cover - executed as compiled Numba code return output +@njit(cache=True, nogil=True) +def _step_masked_kernel( # pragma: no cover - executed as compiled Numba code + tokens: np.ndarray, + active: np.ndarray, + reset: np.ndarray, + positions: np.ndarray, + 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, + lct_value: np.ndarray, + lct_lazy: np.ndarray, + lct_lazy_valid: np.ndarray, + lct_stack: np.ndarray, + last: np.ndarray, + size: np.ndarray, + edge_count: np.ndarray, +) -> np.ndarray: + output = np.full(tokens.shape[0], -1, dtype=np.int64) + for row in range(tokens.shape[0]): + if active[row] == 0: + continue + if reset[row] != 0: + used_states = int(size[row]) + used_edges = int(edge_count[row]) + mask = hash_state.shape[1] - 1 + occupied_slots = np.empty(used_edges, dtype=np.int64) + occupied_count = 0 + for state in range(used_states): + edge = int(head[row, state]) + while edge != -1: + token = int(edge_token[row, edge]) + slot = np.int64(_transition_hash(state, token) & np.uint64(mask)) + while hash_state[row, slot] != -1: + if ( + hash_state[row, slot] == state + and hash_token[row, slot] == token + ): + occupied_slots[occupied_count] = slot + occupied_count += 1 + break + slot = (slot + 1) & mask + edge = int(edge_next[row, edge]) + for index in range(occupied_count): + hash_state[row, occupied_slots[index]] = -1 + for state in range(used_states): + head[row, state] = -1 + suffix_link[row, state] = -1 + length[row, state] = 0 + lct_left[row, state] = -1 + lct_right[row, state] = -1 + lct_parent[row, state] = -1 + lct_value[row, state] = -1 + lct_lazy_valid[row, state] = 0 + last[row] = 0 + size[row] = 1 + edge_count[row] = 0 + positions[row] = 0 + position = int(positions[row]) + prediction, new_last, new_size, new_edge_count = _step_row( + int(tokens[row]), + position, + history[row], + head[row], + edge_token[row], + edge_target[row], + edge_next[row], + hash_state[row], + hash_token[row], + hash_edge[row], + suffix_link[row], + length[row], + lct_left[row], + lct_right[row], + lct_parent[row], + lct_value[row], + lct_lazy[row], + lct_lazy_valid[row], + lct_stack[row], + int(last[row]), + int(size[row]), + int(edge_count[row]), + ) + output[row] = prediction + last[row] = new_last + size[row] = new_size + edge_count[row] = new_edge_count + positions[row] = position + 1 + return output + + @njit(cache=True, nogil=True) def _replay_kernel( # pragma: no cover - executed as compiled Numba code tokens: np.ndarray, diff --git a/src/rosa/ragged.py b/src/rosa/ragged.py new file mode 100644 index 0000000..b258ac3 --- /dev/null +++ b/src/rosa/ragged.py @@ -0,0 +1,199 @@ +"""Dynamic ragged batching for exact stateful ROSA inference. + +This module is lazy: importing it does not import Numba or the optional native +companion. Constructing a state requires the ``rosa-torch[numba]`` extra; when +the native companion is installed, one masked native call handles every row. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from torch import Tensor + + +class RaggedInferenceState: + """Fixed-slot state with independent positions and per-step row masks. + + ``reset`` is applied only to active rows and happens before their token is + consumed. Inactive rows always return ``-1`` and are not mutated. + """ + + def __init__( + self, + batch_size: int, + max_length: int, + *, + use_native: bool = True, + ) -> None: + try: + from ._stateful_numba import _init_inference_state + except ModuleNotFoundError as error: # pragma: no cover - optional extra + if error.name == "numba": + raise ModuleNotFoundError( + "ragged inference requires the 'rosa-torch[numba]' extra" + ) from error + raise + + self._state = _init_inference_state(batch_size, max_length) + self._positions = np.zeros(batch_size, dtype=np.int64) + # NativeState discovers this optional ABI extension without changing + # the established scalar-position state layout. + setattr(self._state, "positions", self._positions) + self._use_native = use_native + self._native: Any = None + + @property + def batch_size(self) -> int: + """Number of reusable row slots.""" + + return int(self._state.batch_size) + + @property + def max_length(self) -> int: + """Independent token capacity of each row slot.""" + + return int(self._state.max_length) + + @property + def positions(self) -> Tensor: + """Current consumed-token count for each row, returned as a copy.""" + + return torch.from_numpy(self._positions.copy()) + + @property + def using_native(self) -> bool: + """Whether the masked native companion has been selected.""" + + return self._native not in (None, False) + + def _native_state(self) -> Any: + if not self._use_native or self._native is False: + return None + if self._native is None: + try: + from rosa_native_step import ( # type: ignore[reportMissingImports] + NativeState, + ) + except ModuleNotFoundError: # pragma: no cover - optional companion + self._native = False + return None + candidate = NativeState( # pragma: no cover - optional native companion + self._state + ) + if getattr(candidate, "step_masked", None) is None: + self._native = False + return None + self._native = candidate + return self._native + + def step_masked( + self, + tokens: Tensor, + active: Tensor, + reset: Tensor | None = None, + ) -> Tensor: + """Consume tokens for active rows, optionally recycling selected slots.""" + + if not isinstance(tokens, Tensor): + raise TypeError("tokens must be a Tensor") + if tokens.ndim == 0 and self.batch_size == 1: + tokens = tokens.unsqueeze(0) + if tokens.ndim != 1 or tokens.shape[0] != self.batch_size: + raise ValueError("tokens must have shape [batch_size]") + if tokens.dtype not in { + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + }: + raise TypeError("tokens must use an integer dtype") + active_cpu = self._mask(active, "active") + if reset is None: + reset_cpu = torch.zeros(self.batch_size, dtype=torch.uint8) + else: + reset_cpu = self._mask(reset, "reset") + active_array = active_cpu.numpy() + reset_array = reset_cpu.numpy() + exhausted = ( + (active_array != 0) + & (reset_array == 0) + & (self._positions >= self.max_length) + ) + if bool(np.any(exhausted)): + raise RuntimeError("inference state capacity exceeded") + + device = tokens.device + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + native = self._native_state() + if native is not None: # pragma: no cover - optional native companion + output = native.step_masked(cpu_tokens.numpy(), active_array, reset_array) + else: + from ._stateful_numba import _step_masked_kernel + + state = self._state + output = _step_masked_kernel( + cpu_tokens.numpy(), + active_array, + reset_array, + self._positions, + 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.lct_value, + state.lct_lazy, + state.lct_lazy_valid, + state.lct_stack, + state.last, + state.size, + state.edge_count, + ) + return torch.from_numpy(output).to(device) + + def step( + self, + tokens: Tensor, + active: Tensor | None = None, + reset: Tensor | None = None, + ) -> Tensor: + """Convenience alias; all rows are active when ``active`` is omitted.""" + + if active is None: + active = torch.ones(self.batch_size, dtype=torch.uint8) + return self.step_masked(tokens, active, reset) + + def _mask(self, mask: Tensor, name: str) -> Tensor: + if mask.ndim == 0 and self.batch_size == 1: + mask = mask.unsqueeze(0) + if mask.ndim != 1 or mask.shape[0] != self.batch_size: + raise ValueError(f"{name} must have shape [batch_size]") + if mask.dtype not in (torch.bool, torch.uint8): + raise TypeError(f"{name} must have dtype bool or uint8") + return mask.detach().to(device="cpu", dtype=torch.uint8).contiguous() + + +def init_ragged_state( + batch_size: int, + max_length: int, + *, + use_native: bool = True, +) -> RaggedInferenceState: + """Create a dynamic ragged inference state.""" + + return RaggedInferenceState(batch_size, max_length, use_native=use_native) + + +__all__ = ["RaggedInferenceState", "init_ragged_state"] diff --git a/tests/test_numba_backend.py b/tests/test_numba_backend.py index 7c9d574..743bbf0 100644 --- a/tests/test_numba_backend.py +++ b/tests/test_numba_backend.py @@ -66,6 +66,10 @@ def test_stateful_private_validation_and_full_replay(self) -> None: state = _init_inference_state(1, 1) state.native_state = False self.assertEqual(_forward_step(state, torch.tensor(0)).shape, (1,)) + + prefill_state = _init_inference_state(1, 2) + prefill_state.native_state = False + self.assertEqual(_prefill(prefill_state, torch.tensor([[0, 1]])).shape, (1, 2)) with self.assertRaisesRegex(ValueError, "shape"): _forward_step(state, torch.tensor([1, 2])) diff --git a/tests/test_ragged.py b/tests/test_ragged.py new file mode 100644 index 0000000..336cac5 --- /dev/null +++ b/tests/test_ragged.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import sys +import types +import unittest +from unittest.mock import patch + +import torch + +from rosa.ragged import RaggedInferenceState, init_ragged_state + + +class TestRaggedInference(unittest.TestCase): + def test_mask_reset_and_recycling_match_single_row_oracles(self) -> None: + from rosa._stateful_numba import _forward_step, _init_inference_state + + ragged = RaggedInferenceState(4, 40, use_native=False) + oracles = [_init_inference_state(1, 40) for _ in range(4)] + for oracle in oracles: + oracle.native_state = False + generator = torch.Generator().manual_seed(20260811) + for tick in range(30): + tokens = torch.randint(-2, 7, (4,), generator=generator) + active = torch.tensor([(tick + row) % (row + 2) != 0 for row in range(4)]) + reset = torch.tensor( + [tick in (9 + row, 20 + row) for row in range(4)], + dtype=torch.bool, + ) + before = ragged.positions + actual = ragged.step_masked(tokens, active, reset) + expected = torch.full((4,), -1, dtype=torch.long) + for row in range(4): + if not active[row]: + self.assertEqual(ragged.positions[row], before[row]) + continue + if reset[row]: + oracles[row] = _init_inference_state(1, 40) + oracles[row].native_state = False + expected[row] = _forward_step(oracles[row], tokens[row]).item() + self.assertTrue(torch.equal(actual, expected), tick) + self.assertEqual(ragged.positions.tolist(), [o.position for o in oracles]) + + def test_capacity_is_per_row_and_reset_precedes_consumption(self) -> None: + state = init_ragged_state(2, 2, use_native=False) + state.step(torch.tensor([1, 4]), torch.tensor([1, 0], dtype=torch.uint8)) + state.step(torch.tensor([2, 5]), torch.tensor([1, 0], dtype=torch.uint8)) + self.assertEqual(state.positions.tolist(), [2, 0]) + with self.assertRaisesRegex(RuntimeError, "capacity"): + state.step(torch.tensor([3, 6]), torch.tensor([1, 0], dtype=torch.uint8)) + output = state.step( + torch.tensor([3, 6]), + torch.tensor([1, 1], dtype=torch.uint8), + torch.tensor([1, 0], dtype=torch.uint8), + ) + self.assertEqual(output.tolist(), [-1, -1]) + self.assertEqual(state.positions.tolist(), [1, 1]) + + def test_inactive_reset_is_ignored_and_inputs_are_validated(self) -> None: + state = RaggedInferenceState(2, 3, use_native=False) + state.step(torch.tensor([1, 2])) + before = state.positions + output = state.step_masked( + torch.tensor([3, 4]), + torch.tensor([False, True]), + torch.tensor([True, False]), + ) + self.assertEqual(output[0], -1) + self.assertEqual(state.positions[0], before[0]) + with self.assertRaisesRegex(ValueError, "tokens"): + state.step(torch.zeros((2, 1), dtype=torch.long)) + with self.assertRaisesRegex(TypeError, "Tensor"): + state.step_masked(object(), torch.ones(2, dtype=torch.bool)) # type: ignore[arg-type] + with self.assertRaisesRegex(TypeError, "integer"): + state.step(torch.tensor([1.9, 2.1])) + with self.assertRaisesRegex(ValueError, "active"): + state.step_masked(torch.zeros(2, dtype=torch.long), torch.ones(3)) + with self.assertRaisesRegex(TypeError, "active"): + state.step_masked(torch.zeros(2, dtype=torch.long), torch.ones(2)) + with self.assertRaisesRegex(TypeError, "reset"): + state.step_masked( + torch.zeros(2, dtype=torch.long), + torch.ones(2, dtype=torch.bool), + torch.ones(2, dtype=torch.int64), + ) + + scalar = RaggedInferenceState(1, 1, use_native=False) + self.assertEqual( + scalar.step_masked(torch.tensor(1), torch.tensor(True)).ndim, 1 + ) + + def test_missing_companion_falls_back_exactly(self) -> None: + with patch.dict(sys.modules, {"rosa_native_step": None}): + state = RaggedInferenceState(1, 4) + actual = torch.stack( + [state.step(torch.tensor(token)) for token in (0, 1, 0, 2)] + ).flatten() + self.assertEqual(actual.tolist(), [-1, -1, 1, -1]) + self.assertFalse(state.using_native) + + class OldNativeState: + def __init__(self, state: object) -> None: + self.state = state + + old_module = types.SimpleNamespace(NativeState=OldNativeState) + with patch.dict(sys.modules, {"rosa_native_step": old_module}): + old_state = RaggedInferenceState(1, 1) + self.assertIsNone(old_state._native_state()) + self.assertFalse(old_state.using_native) + + class CurrentNativeState(OldNativeState): + def step_masked(self) -> None: + return None + + current_module = types.SimpleNamespace(NativeState=CurrentNativeState) + with patch.dict(sys.modules, {"rosa_native_step": current_module}): + current_state = RaggedInferenceState(1, 1) + selected = current_state._native_state() + self.assertIsInstance(selected, CurrentNativeState) + self.assertIs(current_state._native_state(), selected) + self.assertTrue(current_state.using_native) + + +if __name__ == "__main__": + unittest.main() From 9d6b62a7b43068982e0cf35a05f61dc9ff02e811 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 07:11:25 +0800 Subject: [PATCH 20/29] Add structured differentiable candidates --- src/rosa/__init__.py | 252 +++++++++++++++++++++++-- src/rosa/_stateful_candidates_numba.py | 4 +- tests/test_rosa.py | 167 ++++++++++++++++ tests/test_stateful_candidates.py | 21 ++- 4 files changed, 428 insertions(+), 16 deletions(-) diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py index 157161c..9cdbca2 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -31,7 +31,9 @@ "ROSAOutput", "build_hard_candidates", "build_virtual_pool_indices", + "forward_candidates_step", "forward_step", + "init_candidate_state", "init_inference_state", "prefill", "reference_rosa", @@ -495,6 +497,10 @@ def __init__( soft_verify_window: int = 32, virtual_candidates: int = 4, virtual_pool_size: int = 64, + dense_recent_candidates: int = 0, + sparse_old_candidates: int = 0, + sparse_old_pool_size: int = 64, + soft_candidates_forward: bool = False, selector_dim: int = 128, token_temperature: float = 1.0, retrieval_temperature: float = 1.0, @@ -518,6 +524,14 @@ def __init__( ) 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: + raise ValueError( + "sparse_old_pool_size must be >= sparse_old_candidates >= 0" + ) + if not isinstance(soft_candidates_forward, bool): + raise TypeError("soft_candidates_forward must be a bool") if selector_dim <= 0: raise ValueError("selector_dim must be > 0") if token_temperature <= 0 or retrieval_temperature <= 0: @@ -540,6 +554,10 @@ def __init__( self.soft_verify_window = soft_verify_window self.virtual_candidates = virtual_candidates self.virtual_pool_size = virtual_pool_size + self.dense_recent_candidates = dense_recent_candidates + self.sparse_old_candidates = sparse_old_candidates + self.sparse_old_pool_size = sparse_old_pool_size + self.soft_candidates_forward = soft_candidates_forward self.selector_dim = selector_dim self.token_temperature = token_temperature self.retrieval_temperature = retrieval_temperature @@ -692,6 +710,81 @@ def _candidate_symbolic_values( e2 = g2 @ self.symbol_embedding_2.weight return (e1 + e2) * mask.unsqueeze(-1).to(e1.dtype) + def _hybrid_soft_candidates( + self, + st1: Tensor, + st2: Tensor, + exact_source: Tensor, + exact_mask: Tensor, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Build separately budgeted recent and old soft-only candidates.""" + + bsz, n, _ = exact_source.shape + device = exact_source.device + dense_count = self.dense_recent_candidates + sparse_count = self.sparse_old_candidates + + positions = torch.arange(n, device=device).view(1, n, 1) + if dense_count: + offsets = torch.arange(1, dense_count + 1, device=device).view(1, 1, -1) + dense_source = (positions - offsets).expand(bsz, -1, -1) + 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, n, 0), dtype=torch.long, device=device) + dense_mask = torch.empty((bsz, n, 0), dtype=torch.bool, device=device) + + if not sparse_count: + sparse_source = torch.empty((bsz, n, 0), dtype=torch.long, device=device) + sparse_mask = torch.empty((bsz, n, 0), dtype=torch.bool, device=device) + return dense_source, dense_mask, sparse_source, sparse_mask + + pool_size = self.sparse_old_pool_size + # H is constant: evenly spaced anchors cover the admissible old range. + # Build every position in one tensor operation; a Python loop here + # would launch O(N) tiny CUDA kernels. The strict upper bound keeps old + # and recent quotas disjoint. + old_count = (torch.arange(n, device=device) - dense_count).clamp_min(0) + anchor_rank = torch.arange(pool_size, device=device).view(1, -1) + valid_anchors = anchor_rank < old_count.clamp_max(pool_size).view(-1, 1) + if pool_size == 1: + position_anchors = torch.zeros((n, 1), dtype=torch.long, device=device) + else: + spread = torch.round( + anchor_rank * (old_count - 1).clamp_min(0).view(-1, 1) / (pool_size - 1) + ).to(torch.long) + position_anchors = torch.where( + old_count.view(-1, 1) <= pool_size, + anchor_rank.expand(n, -1), + spread, + ) + pool = position_anchors.unsqueeze(0).expand(bsz, -1, -1) + pool_mask = valid_anchors.unsqueeze(0).expand(bsz, -1, -1) + hard_duplicate = ( + pool.unsqueeze(-1) == exact_source.unsqueeze(-2) + ) & exact_mask.unsqueeze(-2) + dense_duplicate = ( + pool.unsqueeze(-1) == dense_source.unsqueeze(-2) + ) & dense_mask.unsqueeze(-2) + pool_mask = ( + pool_mask & ~hard_duplicate.any(dim=-1) & ~dense_duplicate.any(dim=-1) + ) + pool_score = self._soft_match(st1, st2, pool, pool_mask).masked_fill( + ~pool_mask, -1e9 + ) + + # Stable lexicographic order: score descending, then source descending. + recency_order = torch.argsort(pool, dim=-1, descending=True, stable=True) + ordered_score = pool_score.gather(-1, recency_order) + score_order = torch.argsort(ordered_score, dim=-1, descending=True, stable=True) + selected = recency_order.gather(-1, score_order)[..., :sparse_count] + sparse_source = pool.gather(-1, selected) + sparse_mask = pool_mask.gather(-1, selected) + return dense_source, dense_mask, sparse_source, sparse_mask + def forward( self, z_a: Tensor, @@ -726,24 +819,59 @@ def forward( virtual_mask = virtual_mask & (self.virtual_scale > 0) virtual_router = virtual_router * self.virtual_scale + dense_source, dense_mask, sparse_source, sparse_mask = ( + self._hybrid_soft_candidates(st1, st2, exact_source, exact_mask) + ) + bsz, n, _ = z_a.shape null_source = torch.full((bsz, n, 1), -1, dtype=torch.long, device=z_a.device) null_mask = torch.ones((bsz, n, 1), dtype=torch.bool, device=z_a.device) - source = torch.cat([exact_source, virtual_source, null_source], dim=-1) - mask = torch.cat([exact_mask, virtual_mask, null_mask], dim=-1) + 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 + ) exact_kind = torch.full_like(exact_source, EXACT_KIND) virtual_kind = torch.full_like(virtual_source, VIRTUAL_KIND) + dense_kind = torch.full_like(dense_source, VIRTUAL_KIND) + sparse_kind = torch.full_like(sparse_source, VIRTUAL_KIND) null_kind = torch.full_like(null_source, NULL_KIND) - kind = torch.cat([exact_kind, virtual_kind, null_kind], dim=-1) + kind = torch.cat( + [exact_kind, virtual_kind, dense_kind, sparse_kind, null_kind], dim=-1 + ) zeros_virtual = torch.zeros_like(virtual_source) + zeros_dense = torch.zeros_like(dense_source) + zeros_sparse = torch.zeros_like(sparse_source) zeros_null = torch.zeros_like(null_source) hard_match_length = torch.cat( - [hard.match_length, zeros_virtual, zeros_null], dim=-1 + [ + hard.match_length, + zeros_virtual, + zeros_dense, + zeros_sparse, + zeros_null, + ], + dim=-1, ) frequency = torch.cat( - [hard.frequency, torch.ones_like(virtual_source), zeros_null], dim=-1 + [ + hard.frequency, + torch.ones_like(virtual_source), + torch.ones_like(dense_source), + torch.ones_like(sparse_source), + zeros_null, + ], + dim=-1, ) non_null_mask = mask & (kind != NULL_KIND) @@ -809,15 +937,32 @@ def forward( # Curriculum gate for virtual candidates. At zero, the virtual branch # is effectively disabled while exact suffix and NULL candidates remain. virtual_log_gate = torch.log(self.virtual_scale.clamp_min(1e-6)) - rosa_prior = rosa_prior + is_virtual * virtual_log_gate + 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) - chosen = scores.argmax(dim=-1) + hard_scores = scores + if not self.soft_candidates_forward: + soft_start = exact_slots + self.virtual_candidates + soft_end = ( + soft_start + self.dense_recent_candidates + self.sparse_old_candidates + ) + 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) - st_weights = hard_weights + soft_weights - soft_weights.detach() + 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() next_position = torch.where(non_null_mask, source + 1, torch.zeros_like(source)) symbolic_value = self._candidate_symbolic_values( @@ -832,7 +977,48 @@ def forward( value_gate = value_gate * non_null_mask.to(z_a.dtype) neural_value = self.value_proj(next_z) * value_gate.unsqueeze(-1) candidate_value = symbolic_value + self.neural_value_scale * neural_value - retrieved = torch.sum(st_weights.unsqueeze(-1) * candidate_value, dim=-2) + hybrid_enabled = bool( + self.dense_recent_candidates or self.sparse_old_candidates + ) + if hybrid_enabled and not self.soft_candidates_forward: + # Recreate the historical [hard, legacy virtual, NULL] reduction + # exactly, then replace only its soft backward correction with the + # full-union correction. The parenthesized delta is exactly zero + # in forward arithmetic. + 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_a)) updated = z_b + read_gate * self.out_proj(retrieved) @@ -866,9 +1052,7 @@ def forward( ) 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 : exact_slots + self.virtual_candidates - ] + 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(), @@ -980,6 +1164,50 @@ def _make_inference_impl( return _init_inference_state(batch_size, max_length) +def init_candidate_state( + batch_size: int, + max_length: int, + *, + suffix_k: int = 16, + occurrences_r: int = 4, +) -> Any: + """Allocate the exact bounded rich-candidate inference state. + + This optional path requires the ``numba`` extra. It tracks K suffix states, + top-R occurrences and unbounded frequencies without eager suffix-chain + propagation. + """ + + try: + from ._stateful_candidates_numba import init_candidate_state as initialize + except ModuleNotFoundError as error: + if error.name == "numba": + raise RuntimeError( + "rich stateful candidates require the 'numba' extra" + ) from error + raise # pragma: no cover - unrelated optional import failure + return initialize( + batch_size, + max_length, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + + +def forward_candidates_step(state: object, tokens: Tensor) -> Any: + """Consume one token and return exact bounded hard candidates.""" + + try: + from ._stateful_candidates_numba import forward_candidates_step as step + except ModuleNotFoundError as error: + if error.name == "numba": + raise RuntimeError( + "rich stateful candidates require the 'numba' extra" + ) from error + raise # pragma: no cover - unrelated optional import failure + return step(cast(Any, state), tokens) + + 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 f5084aa..f3cd800 100644 --- a/src/rosa/_stateful_candidates_numba.py +++ b/src/rosa/_stateful_candidates_numba.py @@ -694,7 +694,9 @@ def _step_batch_kernel( last: np.ndarray, size: np.ndarray, edge_count: np.ndarray, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: +) -> tuple[ # pragma: no cover - executed as compiled Numba code + np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray +]: slots = suffix_k * occurrences_r batch_size = tokens.shape[0] source = np.full((batch_size, slots), -1, dtype=np.int64) diff --git a/tests/test_rosa.py b/tests/test_rosa.py index ebe680a..1dc6c42 100644 --- a/tests/test_rosa.py +++ b/tests/test_rosa.py @@ -176,6 +176,16 @@ def test_constructor_validations(self) -> None: dict(d_model=4, virtual_candidates=4, virtual_pool_size=3), "virtual_pool_size", ), + (dict(d_model=4, dense_recent_candidates=-1), "dense_recent"), + (dict(d_model=4, sparse_old_candidates=-1), "sparse_old"), + ( + dict( + d_model=4, + sparse_old_candidates=2, + sparse_old_pool_size=1, + ), + "sparse_old_pool_size", + ), (dict(d_model=4, selector_dim=0), "selector_dim"), (dict(d_model=4, token_temperature=0), "temperatures"), (dict(d_model=4, retrieval_temperature=0), "temperatures"), @@ -192,6 +202,8 @@ def test_constructor_validations(self) -> None: self.assertRaisesRegex(ValueError, pattern), ): ROSA(**kwargs) + with self.assertRaisesRegex(TypeError, "soft_candidates_forward"): + ROSA(d_model=4, soft_candidates_forward=1) # type: ignore[arg-type] def test_setters_property_and_encode_validation(self) -> None: model = ROSA( @@ -431,6 +443,161 @@ def test_external_code_logits_receive_gradient(self) -> None: float(logits[0].grad.abs().sum() + logits[1].grad.abs().sum()), 0.0 ) + def test_hybrid_union_quotas_causality_dedup_and_constant_budget(self) -> None: + model = ROSA( + d_model=8, + codebook_sizes=(4, 4), + suffix_k=2, + occurrences_r=2, + soft_verify_window=3, + virtual_candidates=1, + virtual_pool_size=2, + dense_recent_candidates=3, + sparse_old_candidates=2, + sparse_old_pool_size=4, + selector_dim=8, + virtual_scale=0.0, + ) + for n in (7, 12): + tokens = torch.arange(n, dtype=torch.long).unsqueeze(0) + logits = factor_logits_from_tokens(tokens, (4, 4)) + out = model(torch.randn(1, n, 8), code_logits=logits) + # K*R hard, one legacy virtual, D dense, S sparse, NULL. + self.assertEqual(out.candidate_source_index.shape[-1], 4 + 1 + 3 + 2 + 1) + dense = slice(5, 8) + sparse = slice(8, 10) + for i in range(n): + dense_valid = out.candidate_source_index[0, i, dense][ + out.candidate_mask[0, i, dense] + ] + sparse_valid = out.candidate_source_index[0, i, sparse][ + out.candidate_mask[0, i, sparse] + ] + self.assertEqual(dense_valid.numel(), min(3, i)) + self.assertEqual(sparse_valid.numel(), min(2, max(0, i - 3))) + if dense_valid.numel(): + self.assertTrue(torch.all(dense_valid < i)) + self.assertEqual( + dense_valid.tolist(), list(range(i - 1, max(-1, i - 4), -1)) + ) + if sparse_valid.numel(): + self.assertTrue(torch.all(sparse_valid < i - 3)) + valid_source = out.candidate_source_index[0, i][ + out.candidate_mask[0, i] & (out.candidate_source_index[0, i] >= 0) + ] + self.assertEqual(len(valid_source), len(set(valid_source.tolist()))) + + # All sparse scores tie for unique symbols: stable secondary ordering + # chooses the newest anchors, [6, 4], from [0, 2, 4, 6]. + self.assertEqual(out.candidate_source_index[0, 10, 8:10].tolist(), [6, 4]) + + one_anchor = self.make_model( + dense_recent_candidates=1, + sparse_old_candidates=1, + sparse_old_pool_size=1, + ) + one_anchor( + torch.randn(1, 4, 8), + code_logits=factor_logits_from_tokens( + torch.arange(4, dtype=torch.long).unsqueeze(0), (2, 3) + ), + ) + + def test_soft_only_union_preserves_hard_forward_and_opt_in_can_win(self) -> None: + torch.manual_seed(2026) + tokens = torch.tensor([[0, 1, 0, 2, 3, 1, 4, 5]], dtype=torch.long) + logits = factor_logits_from_tokens(tokens, (2, 3), hi=4.0, lo=-4.0) + z_a = torch.randn(1, tokens.shape[1], 8) + z_b = torch.randn_like(z_a) + baseline = self.make_model(virtual_scale=0.0) + union = self.make_model( + virtual_scale=0.0, + dense_recent_candidates=2, + sparse_old_candidates=2, + sparse_old_pool_size=4, + soft_candidates_forward=False, + ) + union.load_state_dict(baseline.state_dict()) + expected = baseline(z_a, z_b=z_b, code_logits=logits) + actual = union(z_a, z_b=z_b, code_logits=logits) + for name in ( + "updated", + "retrieved", + "chosen_source_index", + "chosen_token", + "chosen_match_length", + "hard_rosa_source_index", + "hard_rosa_predicted_tokens", + ): + self.assertTrue( + torch.equal(getattr(actual, name), getattr(expected, name)), name + ) + + opt_in = self.make_model( + learned_residual_scale=1.0, + virtual_scale=0.0, + dense_recent_candidates=1, + sparse_old_candidates=0, + soft_candidates_forward=True, + ) + zero_learned_scorer(opt_in) + with torch.no_grad(): + opt_in.kind_bias[VIRTUAL_KIND] = 30.0 + opt_in.kind_bias[NULL_KIND] = -30.0 + unique = torch.tensor([[0, 1, 2, 3, 4, 5]], dtype=torch.long) + unique_logits = factor_logits_from_tokens(unique, (2, 3)) + opted = opt_in(torch.zeros(1, 6, 8), code_logits=unique_logits) + self.assertTrue(opted.chosen_is_virtual[0, 1:].all()) + self.assertEqual(opted.chosen_source_index[0, 1:].tolist(), [0, 1, 2, 3, 4]) + + def test_recent_and_old_almost_matches_receive_targeted_gradient(self) -> None: + torch.manual_seed(2026) + tokens = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7]], dtype=torch.long) + # A small margin keeps hard argmaxes distinct while exposing useful + # overlap to the soft backward path. + logits = factor_logits_from_tokens( + tokens, (2, 4), hi=0.1, lo=0.0, requires_grad=True + ) + hard = build_hard_candidates(tokens, suffix_k=1, occurrences_r=1) + self.assertFalse(hard.mask[0, 7, 0]) + model = ROSA( + d_model=8, + codebook_sizes=(2, 4), + suffix_k=1, + occurrences_r=1, + soft_verify_window=3, + virtual_candidates=1, + virtual_pool_size=2, + dense_recent_candidates=2, + sparse_old_candidates=1, + sparse_old_pool_size=4, + selector_dim=8, + learned_residual_scale=1.0, + virtual_scale=0.0, + soft_candidates_forward=False, + ) + out = model(torch.randn(1, 8, 8), code_logits=logits) + # Layout: hard[1], legacy[1], dense[2], sparse[1], NULL. + recent_position, recent_slot = 7, 2 + old_position, old_slot = 7, 4 + self.assertEqual( + int(out.candidate_source_index[0, recent_position, recent_slot]), 6 + ) + self.assertEqual(int(out.candidate_source_index[0, old_position, old_slot]), 4) + self.assertLess(4, old_position - model.dense_recent_candidates) + recent_loss = -torch.log(out.soft_weights[0, recent_position, recent_slot]) + old_loss = -torch.log(out.soft_weights[0, old_position, old_slot]) + recent_gradient = torch.autograd.grad(recent_loss, logits, retain_graph=True) + old_gradient = torch.autograd.grad(old_loss, logits) + self.assertGreater( + float(sum(gradient[0, 4:8].abs().sum() for gradient in recent_gradient)), + 1e-8, + ) + self.assertGreater( + float(sum(gradient[0, 1:8].abs().sum() for gradient in old_gradient)), + 1e-8, + ) + def test_combine_losses_validation(self) -> None: base = torch.tensor(2.0) aux = { diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py index d4dbbb8..d54608f 100644 --- a/tests/test_stateful_candidates.py +++ b/tests/test_stateful_candidates.py @@ -1,16 +1,20 @@ from __future__ import annotations +import sys import unittest from itertools import product +from unittest.mock import patch import torch -from rosa import build_hard_candidates +from rosa import ( + build_hard_candidates, + forward_candidates_step, + init_candidate_state, +) from rosa._stateful_candidates_numba import ( CandidateState, CandidateStep, - forward_candidates_step, - init_candidate_state, ) _FIELDS = ( @@ -127,6 +131,17 @@ def test_scalar_batch_and_validation(self) -> None: forward_candidates_step(shape_state, torch.tensor([1.0, 2.0])) with self.assertRaisesRegex(TypeError, "CandidateState"): forward_candidates_step(object(), torch.tensor([1])) # type: ignore[arg-type] + with self.assertRaisesRegex(TypeError, "Tensor"): + forward_candidates_step(state, object()) # type: ignore[arg-type] + + def test_public_wrapper_reports_missing_numba(self) -> None: + with patch.dict(sys.modules): + sys.modules.pop("rosa._stateful_candidates_numba", None) + sys.modules["numba"] = None + with self.assertRaisesRegex(RuntimeError, "numba"): + init_candidate_state(1, 1) + with self.assertRaisesRegex(RuntimeError, "numba"): + forward_candidates_step(object(), torch.tensor([1])) if __name__ == "__main__": From 528c6d9e6b8fbc7862dc4d18cd5625c04f9c1b9a Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 07:33:25 +0800 Subject: [PATCH 21/29] Use native rich candidates in ROSA forward --- native/README.md | 22 ++ native/benchmark_candidates.py | 62 ++++ native/src/rosa_native_step.cpp | 452 +++++++++++++++++++++++++ native/tests/candidate_smoke.py | 128 +++++++ src/rosa/__init__.py | 92 ++++- src/rosa/_stateful_candidates_numba.py | 89 +++-- tests/test_rosa.py | 142 ++++++++ tests/test_stateful_candidates.py | 1 + 8 files changed, 955 insertions(+), 33 deletions(-) create mode 100644 native/benchmark_candidates.py create mode 100644 native/tests/candidate_smoke.py diff --git a/native/README.md b/native/README.md index 2226303..56b7794 100644 --- a/native/README.md +++ b/native/README.md @@ -15,6 +15,19 @@ PyTorch, NumPy et Numba. Le constructeur valide intégralement formes, types, compteurs et version ABI avant de conserver les pointeurs. L'ABI d'état native actuelle vaut `1`. +Le module expose aussi `NativeCandidateState` pour l'état riche exact de +`rosa._stateful_candidates_numba`. Son `step` batch maintient les mêmes K +suffixes, R occurrences les plus récentes, fréquences non bornées et tags LCT +`newest-prefix + delta`. La capacité R reste possédée par les tableaux NumPy du +`CandidateState` Python, dont l'objet natif conserve la durée de vie. La +capacité peut être détectée via la présence de `NativeCandidateState` et +`candidate_abi_version == 1`. + +Cette première ABI riche expose `step`, `reset` global et `position`. Elle +n'expose pas encore de préremplissage riche ni de reset/continuation masqué par +ligne; ces opérations nécessitent un contrat de positions par ligne distinct +de `CandidateState.position`. + ## Installation et utilisation Installez un wheel correspondant à la version de Python et à la plateforme : @@ -34,6 +47,11 @@ from rosa_native_step import NativeState convertible en `int64`, de forme `[batch_size]`. L'objet conserve une référence à l'état Python et expose sa `position` en lecture seule. +`NativeCandidateState(candidate_state).step(tokens_numpy)` attend strictement +un vecteur NumPy C-contigu `int64` et renvoie le tuple bas niveau +`(source, match_length, state_id, frequency, count)`. `reset()` recycle tout le +batch en temps proportionnel aux états et slots de hachage réellement occupés. + ## Construction locale isolée Le backend PEP 517 est setuptools, avec pybind11 uniquement comme dépendance de @@ -54,6 +72,10 @@ uv run --isolated \ Le smoke force Numba comme oracle, laisse le chemin ROSA courant charger le compagnon, puis compare les prédictions étape par étape. +Le smoke riche et le benchmark direct contre Numba se lancent respectivement +avec `native/tests/candidate_smoke.py` et `native/benchmark_candidates.py` dans +le même environnement isolé. + ## Publication multi-plateforme Étape suivante : ajouter un workflow `cibuildwheel` dédié après avoir fixé la diff --git a/native/benchmark_candidates.py b/native/benchmark_candidates.py new file mode 100644 index 0000000..a4f691b --- /dev/null +++ b/native/benchmark_candidates.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import statistics +import time + +import numpy as np +import rosa_native_step +import torch + +from rosa._stateful_candidates_numba import ( + forward_candidates_step, + init_candidate_state, +) + + +def measure_native(tokens: torch.Tensor, suffix_k: int, occurrences_r: int) -> float: + started = time.perf_counter_ns() + state = init_candidate_state( + tokens.shape[0], + tokens.shape[1], + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + native = rosa_native_step.NativeCandidateState(state) + for position in range(tokens.shape[1]): + native.step(np.ascontiguousarray(tokens[:, position].numpy())) + return (time.perf_counter_ns() - started) / 1e6 + + +def measure_numba(tokens: torch.Tensor, suffix_k: int, occurrences_r: int) -> float: + started = time.perf_counter_ns() + state = init_candidate_state( + tokens.shape[0], + tokens.shape[1], + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + for position in range(tokens.shape[1]): + forward_candidates_step(state, tokens[:, position]) + return (time.perf_counter_ns() - started) / 1e6 + + +def main() -> None: + batch_size, length, suffix_k, occurrences_r = 8, 4096, 16, 4 + tokens = torch.randint( + 128, + (batch_size, length), + generator=torch.Generator().manual_seed(20260811), + ) + warm = init_candidate_state(1, 2, suffix_k=1, occurrences_r=1) + forward_candidates_step(warm, torch.tensor([0])) + native_ms = [measure_native(tokens, suffix_k, occurrences_r) for _ in range(5)] + numba_ms = [measure_numba(tokens, suffix_k, occurrences_r) for _ in range(5)] + native = statistics.median(native_ms) + numba = statistics.median(numba_ms) + print(f"native_ms={native:.6f}") + print(f"numba_ms={numba:.6f}") + print(f"native_vs_numba={numba / native:.6f}x") + + +if __name__ == "__main__": + main() diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp index 957b613..eb90e6b 100644 --- a/native/src/rosa_native_step.cpp +++ b/native/src/rosa_native_step.cpp @@ -784,6 +784,452 @@ class NativeState { hash_capacity_; }; +class NativeCandidateState { +public: + explicit NativeCandidateState(py::object state) : state_(std::move(state)) { + if (!py::hasattr(state_, "native_candidate_abi_version") || + py::cast(state_.attr("native_candidate_abi_version")) != 1) + throw py::value_error("unsupported native candidate state ABI"); + history_ = bind("history"); + head_ = bind("head"); + edge_token_ = bind("edge_token"); + edge_target_ = bind("edge_target"); + edge_next_ = bind("edge_next"); + hash_state_ = bind("hash_state"); + hash_token_ = bind("hash_token"); + hash_edge_ = bind("hash_edge"); + suffix_link_ = bind("suffix_link"); + length_ = bind("length"); + left_ = bind("lct_left"); + right_ = bind("lct_right"); + parent_ = bind("lct_parent"); + occurrences_ = bind("occurrences"); + occurrence_size_ = bind("occurrence_size"); + frequency_ = bind("frequency"); + lazy_prefix_ = bind("lazy_prefix"); + lazy_size_ = bind("lazy_size"); + lazy_delta_ = bind("lazy_delta"); + stack_ = bind("lct_stack"); + last_ = bind("last"); + size_ = bind("size"); + edge_count_ = bind("edge_count"); + batch_ = py::cast(state_.attr("batch_size")); + max_length_ = py::cast(state_.attr("max_length")); + suffix_k_ = py::cast(state_.attr("suffix_k")); + occurrences_r_ = py::cast(state_.attr("occurrences_r")); + position_ = py::cast(state_.attr("position")); + state_capacity_ = head_.shape(1); + edge_capacity_ = edge_token_.shape(1); + hash_capacity_ = hash_state_.shape(1); + validate_shapes(); + occupied_slots_.resize(batch_); + for (int64_t b = 0; b < batch_; ++b) + for (int64_t slot = 0; slot < hash_capacity_; ++slot) + if (hash_state_.data()[idx(b, hash_capacity_, slot)] != -1) + occupied_slots_[b].push_back(static_cast(slot)); + } + + py::tuple step(py::array tokens_object) { + if (!py::isinstance>(tokens_object)) + throw py::type_error("tokens must have dtype int64"); + if ((tokens_object.flags() & py::array::c_style) == 0 || + tokens_object.ndim() != 1 || tokens_object.shape(0) != batch_) + throw py::value_error("tokens must be contiguous int64 [batch_size]"); + if (position_ >= max_length_) + throw std::runtime_error("candidate state capacity exceeded"); + auto tokens = py::cast>(tokens_object); + const int64_t slots = suffix_k_ * occurrences_r_; + py::array_t source({batch_, slots}), match_length({batch_, slots}), + state_id({batch_, slots}), candidate_frequency({batch_, slots}); + py::array_t count(batch_); + std::fill(source.mutable_data(), source.mutable_data() + batch_ * slots, + int64_t{-1}); + std::fill(match_length.mutable_data(), + match_length.mutable_data() + batch_ * slots, int64_t{0}); + std::fill(state_id.mutable_data(), state_id.mutable_data() + batch_ * slots, + int64_t{-1}); + std::fill(candidate_frequency.mutable_data(), + candidate_frequency.mutable_data() + batch_ * slots, int64_t{0}); + { + py::gil_scoped_release release; + for (int64_t b = 0; b < batch_; ++b) + count.mutable_data()[b] = + step_row(b, tokens.data()[b], position_, + source.mutable_data() + b * slots, + match_length.mutable_data() + b * slots, + state_id.mutable_data() + b * slots, + candidate_frequency.mutable_data() + b * slots); + } + ++position_; + state_.attr("position") = py::int_(position_); + return py::make_tuple(source, match_length, state_id, candidate_frequency, + count); + } + + void reset() { + { + py::gil_scoped_release release; + for (int64_t b = 0; b < batch_; ++b) + reset_row(b); + } + position_ = 0; + state_.attr("position") = py::int_(0); + } + + int64_t position() const { return position_; } + +private: + template py::array_t bind(const char *name) { + py::object object = state_.attr(name); + if (!py::isinstance>(object)) + throw py::type_error(std::string(name) + " has an unexpected dtype"); + auto array = py::cast>(object); + if (!array.writeable()) + throw py::value_error(std::string(name) + " is readonly"); + return array; + } + template + bool matrix_shape(const py::array_t &a, int64_t rows, + int64_t columns) const { + return a.ndim() == 2 && a.shape(0) == rows && a.shape(1) == columns; + } + template + bool vector_shape(const py::array_t &a, + int64_t length) const { + return a.ndim() == 1 && a.shape(0) == length; + } + void validate_shapes() { + if (batch_ <= 0 || max_length_ <= 0 || suffix_k_ <= 0 || + occurrences_r_ <= 0 || position_ < 0 || position_ > max_length_ || + state_capacity_ <= 0 || edge_capacity_ <= 0 || hash_capacity_ <= 0 || + (hash_capacity_ & (hash_capacity_ - 1)) != 0) + throw py::value_error("incompatible CandidateState layout"); + if (!matrix_shape(history_, batch_, max_length_) || + !matrix_shape(head_, batch_, state_capacity_) || + !matrix_shape(edge_token_, batch_, edge_capacity_) || + !matrix_shape(edge_target_, batch_, edge_capacity_) || + !matrix_shape(edge_next_, batch_, edge_capacity_) || + !matrix_shape(hash_state_, batch_, hash_capacity_) || + !matrix_shape(hash_token_, batch_, hash_capacity_) || + !matrix_shape(hash_edge_, batch_, hash_capacity_) || + !matrix_shape(suffix_link_, batch_, state_capacity_) || + !matrix_shape(length_, batch_, state_capacity_) || + !matrix_shape(left_, batch_, state_capacity_) || + !matrix_shape(right_, batch_, state_capacity_) || + !matrix_shape(parent_, batch_, state_capacity_) || + occurrences_.ndim() != 3 || occurrences_.shape(0) != batch_ || + occurrences_.shape(1) != state_capacity_ || + occurrences_.shape(2) != occurrences_r_ || + lazy_prefix_.ndim() != 3 || lazy_prefix_.shape(0) != batch_ || + lazy_prefix_.shape(1) != state_capacity_ || + lazy_prefix_.shape(2) != occurrences_r_ || + !matrix_shape(occurrence_size_, batch_, state_capacity_) || + !matrix_shape(frequency_, batch_, state_capacity_) || + !matrix_shape(lazy_size_, batch_, state_capacity_) || + !matrix_shape(lazy_delta_, batch_, state_capacity_) || + !matrix_shape(stack_, batch_, state_capacity_) || + !vector_shape(last_, batch_) || !vector_shape(size_, batch_) || + !vector_shape(edge_count_, batch_)) + throw py::value_error("incompatible CandidateState layout"); + for (int64_t b = 0; b < batch_; ++b) + if (last_.data()[b] < 0 || last_.data()[b] >= state_capacity_ || + size_.data()[b] < 1 || size_.data()[b] > state_capacity_ || + edge_count_.data()[b] < 0 || edge_count_.data()[b] > edge_capacity_) + throw py::value_error("incompatible CandidateState counters"); + } + inline int64_t idx(int64_t b, int64_t width, int64_t i) const { + return b * width + i; + } + inline int64_t occ_idx(int64_t b, int32_t node, int64_t i) const { + return (b * state_capacity_ + node) * occurrences_r_ + i; + } + inline uint64_t transition_hash(int32_t state, int64_t token) const { + uint64_t v = static_cast(token); + v ^= static_cast(state) + UINT64_C(0x9E3779B97F4A7C15); + v = (v ^ (v >> 30)) * UINT64_C(0xBF58476D1CE4E5B9); + v = (v ^ (v >> 27)) * UINT64_C(0x94D049BB133111EB); + return v ^ (v >> 31); + } + int32_t find_transition(int64_t b, int32_t state, int64_t token) const { + int64_t slot = static_cast(transition_hash(state, token) & + static_cast(hash_capacity_ - 1)); + while (hash_state_.data()[idx(b, hash_capacity_, slot)] != -1) { + const int64_t at = idx(b, hash_capacity_, slot); + if (hash_state_.data()[at] == state && hash_token_.data()[at] == token) + return hash_edge_.data()[at]; + slot = (slot + 1) & (hash_capacity_ - 1); + } + return -1; + } + int32_t add_transition(int64_t b, int32_t count, int32_t state, + int64_t token, int32_t target) { + if (count >= edge_capacity_) + throw std::runtime_error("transition capacity exceeded"); + const int64_t edge_at = idx(b, edge_capacity_, count); + edge_token_.mutable_data()[edge_at] = token; + edge_target_.mutable_data()[edge_at] = target; + edge_next_.mutable_data()[edge_at] = head_.data()[idx(b, state_capacity_, state)]; + head_.mutable_data()[idx(b, state_capacity_, state)] = count; + int64_t slot = static_cast(transition_hash(state, token) & + static_cast(hash_capacity_ - 1)); + while (hash_state_.data()[idx(b, hash_capacity_, slot)] != -1) { + const int64_t at = idx(b, hash_capacity_, slot); + if (hash_state_.data()[at] == state && hash_token_.data()[at] == token) + throw std::runtime_error("duplicate suffix automaton transition"); + slot = (slot + 1) & (hash_capacity_ - 1); + } + const int64_t at = idx(b, hash_capacity_, slot); + hash_state_.mutable_data()[at] = state; + hash_token_.mutable_data()[at] = token; + hash_edge_.mutable_data()[at] = count; + occupied_slots_[b].push_back(static_cast(slot)); + return count + 1; + } + void replace_transition(int64_t b, int32_t state, int64_t token, + int32_t target) { + const int32_t edge = find_transition(b, state, token); + if (edge == -1) + throw std::runtime_error("transition not found"); + edge_target_.mutable_data()[idx(b, edge_capacity_, edge)] = target; + } + inline bool is_aux_root(int64_t b, int32_t node) const { + const int64_t at = idx(b, state_capacity_, node); + const int32_t p = parent_.data()[at]; + return p == -1 || (left_.data()[idx(b, state_capacity_, p)] != node && + right_.data()[idx(b, state_capacity_, p)] != node); + } + void apply_tag(int64_t b, int32_t node, const int64_t *prefix, + int32_t prefix_size, int64_t delta) { + if (node == -1) + return; + const int32_t take = std::min(prefix_size, occurrences_r_); + const int64_t at = idx(b, state_capacity_, node); + const int32_t old_size = occurrence_size_.data()[at]; + const int32_t updated = std::min(occurrences_r_, take + old_size); + for (int32_t i = updated - 1; i >= take; --i) + occurrences_.mutable_data()[occ_idx(b, node, i)] = + occurrences_.data()[occ_idx(b, node, i - take)]; + for (int32_t i = 0; i < take; ++i) + occurrences_.mutable_data()[occ_idx(b, node, i)] = prefix[i]; + occurrence_size_.mutable_data()[at] = updated; + frequency_.mutable_data()[at] += delta; + const int32_t old_lazy = lazy_size_.data()[at]; + const int32_t updated_lazy = + std::min(occurrences_r_, take + old_lazy); + for (int32_t i = updated_lazy - 1; i >= take; --i) + lazy_prefix_.mutable_data()[occ_idx(b, node, i)] = + lazy_prefix_.data()[occ_idx(b, node, i - take)]; + for (int32_t i = 0; i < take; ++i) + lazy_prefix_.mutable_data()[occ_idx(b, node, i)] = prefix[i]; + lazy_size_.mutable_data()[at] = updated_lazy; + lazy_delta_.mutable_data()[at] += delta; + } + void push(int64_t b, int32_t node) { + const int64_t at = idx(b, state_capacity_, node); + const int32_t size = lazy_size_.data()[at]; + const int64_t delta = lazy_delta_.data()[at]; + if (size != 0 || delta != 0) { + const int64_t *prefix = lazy_prefix_.data() + occ_idx(b, node, 0); + apply_tag(b, left_.data()[at], prefix, size, delta); + apply_tag(b, right_.data()[at], prefix, size, delta); + lazy_size_.mutable_data()[at] = 0; + lazy_delta_.mutable_data()[at] = 0; + } + } + void rotate(int64_t b, int32_t node) { + const auto at = [&](int32_t n) { return idx(b, state_capacity_, n); }; + int32_t *left = left_.mutable_data(), *right = right_.mutable_data(), + *parent = parent_.mutable_data(); + const int32_t p = parent[at(node)], g = parent[at(p)]; + int32_t middle; + if (left[at(p)] == node) { + middle = right[at(node)]; right[at(node)] = p; left[at(p)] = middle; + } else { + middle = left[at(node)]; left[at(node)] = p; right[at(p)] = middle; + } + if (middle != -1) parent[at(middle)] = p; + parent[at(p)] = node; parent[at(node)] = g; + if (g != -1) { + if (left[at(g)] == p) left[at(g)] = node; + else if (right[at(g)] == p) right[at(g)] = node; + } + } + void splay(int64_t b, int32_t node) { + int32_t *stack = stack_.mutable_data() + b * state_capacity_; + int32_t depth = 0, ancestor = node; + stack[depth++] = ancestor; + while (!is_aux_root(b, ancestor)) { + ancestor = parent_.data()[idx(b, state_capacity_, ancestor)]; + stack[depth++] = ancestor; + } + while (depth > 0) push(b, stack[--depth]); + while (!is_aux_root(b, node)) { + const int32_t p = parent_.data()[idx(b, state_capacity_, node)]; + if (!is_aux_root(b, p)) { + const int32_t g = parent_.data()[idx(b, state_capacity_, p)]; + if ((left_.data()[idx(b, state_capacity_, p)] == node) == + (left_.data()[idx(b, state_capacity_, g)] == p)) rotate(b, p); + else rotate(b, node); + } + rotate(b, node); + } + } + void access(int64_t b, int32_t node) { + int32_t last = -1, current = node; + while (current != -1) { + splay(b, current); + right_.mutable_data()[idx(b, state_capacity_, current)] = last; + if (last != -1) parent_.mutable_data()[idx(b, state_capacity_, last)] = current; + last = current; + current = parent_.data()[idx(b, state_capacity_, current)]; + } + splay(b, node); + } + void materialize(int64_t b, int32_t node) { access(b, node); } + void cut_parent(int64_t b, int32_t node) { + materialize(b, node); + const int64_t at = idx(b, state_capacity_, node); + const int32_t ancestors = left_.data()[at]; + left_.mutable_data()[at] = -1; + if (ancestors != -1) parent_.mutable_data()[idx(b, state_capacity_, ancestors)] = -1; + } + void link_parent(int64_t b, int32_t node, int32_t represented_parent) { + materialize(b, node); + parent_.mutable_data()[idx(b, state_capacity_, node)] = represented_parent; + } + void path_write(int64_t b, int32_t node, int64_t position) { + materialize(b, node); + apply_tag(b, node, &position, 1, 1); + } + int32_t step_row(int64_t b, int64_t token, int64_t position, + 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], + edge_count = edge_count_.data()[b]; + history_.mutable_data()[idx(b, max_length_, position)] = token; + if (size >= state_capacity_) throw std::runtime_error("state capacity exceeded"); + const int32_t current = size++; + length_.mutable_data()[idx(b, state_capacity_, current)] = + length_.data()[idx(b, state_capacity_, last)] + 1; + int32_t state = last; + while (state != -1 && find_transition(b, state, token) == -1) { + edge_count = add_transition(b, edge_count, state, token, current); + state = suffix_link_.data()[idx(b, state_capacity_, state)]; + } + if (state == -1) { + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = 0; + link_parent(b, current, 0); + } else { + int32_t transition = find_transition(b, state, token); + const int32_t target = edge_target_.data()[idx(b, edge_capacity_, transition)]; + if (length_.data()[idx(b, state_capacity_, state)] + 1 == + length_.data()[idx(b, state_capacity_, target)]) { + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = target; + link_parent(b, current, target); + } else { + if (size >= state_capacity_) throw std::runtime_error("state capacity exceeded"); + const int32_t clone = size++; + length_.mutable_data()[idx(b, state_capacity_, clone)] = + length_.data()[idx(b, state_capacity_, state)] + 1; + const int32_t old_parent = suffix_link_.data()[idx(b, state_capacity_, target)]; + suffix_link_.mutable_data()[idx(b, state_capacity_, clone)] = old_parent; + materialize(b, target); + const int32_t copied = occurrence_size_.data()[idx(b, state_capacity_, target)]; + occurrence_size_.mutable_data()[idx(b, state_capacity_, clone)] = copied; + for (int32_t i = 0; i < copied; ++i) + occurrences_.mutable_data()[occ_idx(b, clone, i)] = occurrences_.data()[occ_idx(b, target, i)]; + frequency_.mutable_data()[idx(b, state_capacity_, clone)] = frequency_.data()[idx(b, state_capacity_, target)]; + lazy_size_.mutable_data()[idx(b, state_capacity_, clone)] = 0; + lazy_delta_.mutable_data()[idx(b, state_capacity_, clone)] = 0; + int32_t edge = head_.data()[idx(b, state_capacity_, target)]; + while (edge != -1) { + edge_count = add_transition(b, edge_count, clone, + edge_token_.data()[idx(b, edge_capacity_, edge)], + edge_target_.data()[idx(b, edge_capacity_, edge)]); + edge = edge_next_.data()[idx(b, edge_capacity_, edge)]; + } + transition = find_transition(b, state, token); + while (state != -1 && transition != -1 && + edge_target_.data()[idx(b, edge_capacity_, transition)] == target) { + replace_transition(b, state, token, clone); + state = suffix_link_.data()[idx(b, state_capacity_, state)]; + if (state != -1) transition = find_transition(b, state, token); + } + link_parent(b, clone, old_parent); + cut_parent(b, target); + suffix_link_.mutable_data()[idx(b, state_capacity_, target)] = clone; + link_parent(b, target, clone); + suffix_link_.mutable_data()[idx(b, state_capacity_, current)] = clone; + link_parent(b, current, clone); + } + } + last = current; + int32_t candidate_count = 0, states_with_history = 0, node = last; + while (node != -1 && states_with_history < suffix_k_) { + const int64_t at = idx(b, state_capacity_, node); + if (length_.data()[at] > 0) { + materialize(b, node); + const int32_t node_occurrences = occurrence_size_.data()[at]; + if (node_occurrences > 0) { + ++states_with_history; + const int32_t take = std::min(occurrences_r_, node_occurrences); + for (int32_t occurrence_index = 0; occurrence_index < take; ++occurrence_index) { + const int64_t source = occurrences_.data()[occ_idx(b, node, occurrence_index)]; + bool duplicate = false; + for (int32_t seen = 0; seen < candidate_count; ++seen) + if (source_out[seen] == source) { duplicate = true; break; } + if (!duplicate) { + source_out[candidate_count] = source; + length_out[candidate_count] = length_.data()[at]; + state_out[candidate_count] = node; + frequency_out[candidate_count] = frequency_.data()[at]; + ++candidate_count; + } + } + } + } + node = suffix_link_.data()[idx(b, state_capacity_, node)]; + } + path_write(b, current, position); + last_.mutable_data()[b] = last; + size_.mutable_data()[b] = size; + edge_count_.mutable_data()[b] = edge_count; + return candidate_count; + } + void reset_row(int64_t b) { + const int32_t used_states = size_.data()[b]; + for (int32_t node = 0; node < used_states; ++node) { + const int64_t at = idx(b, state_capacity_, node); + head_.mutable_data()[at] = -1; + suffix_link_.mutable_data()[at] = -1; + length_.mutable_data()[at] = 0; + left_.mutable_data()[at] = -1; + right_.mutable_data()[at] = -1; + parent_.mutable_data()[at] = -1; + occurrence_size_.mutable_data()[at] = 0; + frequency_.mutable_data()[at] = 0; + lazy_size_.mutable_data()[at] = 0; + lazy_delta_.mutable_data()[at] = 0; + } + for (int32_t slot : occupied_slots_[b]) + hash_state_.mutable_data()[idx(b, hash_capacity_, slot)] = -1; + occupied_slots_[b].clear(); + last_.mutable_data()[b] = 0; + size_.mutable_data()[b] = 1; + edge_count_.mutable_data()[b] = 0; + } + + py::object state_; + py::array_t history_, edge_token_, hash_token_, + occurrences_, frequency_, lazy_prefix_, lazy_delta_; + py::array_t head_, edge_target_, edge_next_, + hash_state_, hash_edge_, suffix_link_, length_, left_, right_, parent_, + occurrence_size_, lazy_size_, stack_, last_, size_, edge_count_; + std::vector> occupied_slots_; + int64_t batch_, max_length_, suffix_k_, occurrences_r_, position_, + state_capacity_, edge_capacity_, hash_capacity_; +}; + PYBIND11_MODULE(rosa_native_step, m) { m.doc() = "Exact CPU SAM+LCT step prototype (no libtorch calls in core)"; py::class_(m, "NativeState") @@ -793,4 +1239,10 @@ PYBIND11_MODULE(rosa_native_step, m) { .def("prefill", &NativeState::prefill) .def_property_readonly("position", &NativeState::position) .def_property_readonly("positions", &NativeState::positions); + py::class_(m, "NativeCandidateState") + .def(py::init(), py::keep_alive<1, 2>()) + .def("step", &NativeCandidateState::step) + .def("reset", &NativeCandidateState::reset) + .def_property_readonly("position", &NativeCandidateState::position); + m.attr("candidate_abi_version") = py::int_(1); } diff --git a/native/tests/candidate_smoke.py b/native/tests/candidate_smoke.py new file mode 100644 index 0000000..fbc0c17 --- /dev/null +++ b/native/tests/candidate_smoke.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import gc +import weakref +from itertools import product + +import numpy as np +import rosa_native_step +import torch + +from rosa._stateful_candidates_numba import ( + CandidateStep, + forward_candidates_step, + init_candidate_state, +) + + +def assert_step_equal(actual: tuple[np.ndarray, ...], expected: CandidateStep) -> None: + expected_arrays = ( + expected.source_index.numpy(), + expected.match_length.numpy(), + expected.state_id.numpy(), + expected.frequency.numpy(), + ) + for candidate, oracle in zip(actual[:4], expected_arrays, strict=True): + assert np.array_equal(candidate, oracle) + assert np.array_equal(actual[4], expected.mask.sum(dim=1).numpy().astype(np.int32)) + + +def compare(tokens: torch.Tensor, suffix_k: int, occurrences_r: int) -> None: + oracle = init_candidate_state( + tokens.shape[0], + tokens.shape[1], + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + candidate = init_candidate_state( + tokens.shape[0], + tokens.shape[1], + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + native = rosa_native_step.NativeCandidateState(candidate) + for position in range(tokens.shape[1]): + expected = forward_candidates_step(oracle, tokens[:, position]) + column = np.ascontiguousarray(tokens[:, position].numpy()) + assert_step_equal(native.step(column), expected) + assert native.position == candidate.position == tokens.shape[1] + + native.reset() + assert native.position == candidate.position == 0 + replay_oracle = init_candidate_state( + tokens.shape[0], + tokens.shape[1], + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + for position in range(tokens.shape[1]): + expected = forward_candidates_step(replay_oracle, tokens[:, position]) + column = np.ascontiguousarray(tokens[:, position].numpy()) + assert_step_equal(native.step(column), expected) + + +def main() -> None: + assert rosa_native_step.candidate_abi_version == 1 + binary = torch.tensor(list(product(range(2), repeat=9)), dtype=torch.long) + for suffix_k, occurrences_r in ((1, 1), (2, 3), (4, 2), (5, 4)): + compare(binary, suffix_k, occurrences_r) + + generator = torch.Generator().manual_seed(20260811) + random_tokens = torch.randint(-3, 9, (7, 193), generator=generator) + compare(random_tokens, 7, 5) + + # A native wrapper may take over an already-mutated Numba state. + oracle = init_candidate_state(7, 193, suffix_k=7, occurrences_r=5) + candidate = init_candidate_state(7, 193, suffix_k=7, occurrences_r=5) + for position in range(83): + forward_candidates_step(oracle, random_tokens[:, position]) + forward_candidates_step(candidate, random_tokens[:, position]) + continuation = rosa_native_step.NativeCandidateState(candidate) + for position in range(83, random_tokens.shape[1]): + expected = forward_candidates_step(oracle, random_tokens[:, position]) + column = np.ascontiguousarray(random_tokens[:, position].numpy()) + assert_step_equal(continuation.step(column), expected) + + state = init_candidate_state(2, 4, suffix_k=3, occurrences_r=2) + state_ref = weakref.ref(state) + native = rosa_native_step.NativeCandidateState(state) + del state + gc.collect() + assert state_ref() is not None + native.step(np.array([1, 2], dtype=np.int64)) + + for invalid in ( + np.zeros(2, dtype=np.int32), + np.zeros((2, 1), dtype=np.int64), + np.zeros(4, dtype=np.int64)[::2], + ): + try: + native.step(invalid) + except (TypeError, ValueError): + pass + else: + raise AssertionError("invalid candidate-step input was accepted") + + malformed = init_candidate_state(2, 4) + malformed.native_candidate_abi_version = 2 + try: + rosa_native_step.NativeCandidateState(malformed) + except ValueError as error: + assert "ABI" in str(error) + else: + raise AssertionError("unsupported candidate ABI was accepted") + + malformed = init_candidate_state(2, 4) + malformed.occurrences = np.empty((2, 1, 4), dtype=np.int64) + try: + rosa_native_step.NativeCandidateState(malformed) + except ValueError as error: + assert "layout" in str(error) + else: + raise AssertionError("malformed candidate layout was accepted") + + print("rosa_native_step candidate smoke: ok") + + +if __name__ == "__main__": + main() diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py index 9cdbca2..22945c8 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -390,6 +390,85 @@ def build_hard_candidates( ) +CandidateBackend = Literal["auto", "python", "stateful"] + + +def _build_stateful_hard_candidates( + tokens: Tensor, + suffix_k: int, + occurrences_r: int, +) -> HardCandidates: + """Replay a full sequence through the exact bounded stateful backend.""" + + from ._stateful_candidates_numba import forward_candidates_step as step + from ._stateful_candidates_numba import init_candidate_state as initialize + + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + batch_size, sequence_length = cpu_tokens.shape + state = initialize( + batch_size, + sequence_length, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + steps = [ + step(state, cpu_tokens[:, position]) for position in range(sequence_length) + ] + stacked = { + name: torch.stack([getattr(item, name) for item in steps], dim=1) + for name in HardCandidates.__dataclass_fields__ + } + candidate_fields = torch.stack( + [ + stacked["source_index"], + stacked["match_length"], + stacked["state_id"], + stacked["frequency"], + ] + ).to(tokens.device) + rosa_fields = torch.stack( + [ + stacked["rosa_slot"], + stacked["rosa_source_index"], + stacked["rosa_match_length"], + stacked["rosa_predicted_tokens"], + ] + ).to(tokens.device) + return HardCandidates( + source_index=candidate_fields[0], + match_length=candidate_fields[1], + state_id=candidate_fields[2], + frequency=candidate_fields[3], + mask=stacked["mask"].to(tokens.device), + rosa_slot=rosa_fields[0], + rosa_source_index=rosa_fields[1], + rosa_match_length=rosa_fields[2], + rosa_predicted_tokens=rosa_fields[3], + ) + + +def _build_forward_hard_candidates( + tokens: Tensor, + suffix_k: int, + occurrences_r: int, + backend: CandidateBackend, +) -> HardCandidates: + """Dispatch ROSA forward candidates without weakening backend failures.""" + + if backend == "python": + return build_hard_candidates(tokens, suffix_k, occurrences_r) + try: + return _build_stateful_hard_candidates(tokens, 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 build_hard_candidates(tokens, suffix_k, occurrences_r) + + def _virtual_pool_single(i: int, pool_size: int) -> list[int]: """Causal bounded pool: half recent positions, half history anchors.""" @@ -512,6 +591,7 @@ def __init__( learned_residual_scale: float = 0.0, virtual_scale: float = 0.0, neural_value_scale: float = 0.0, + candidate_backend: CandidateBackend = "auto", ) -> None: super().__init__() if d_model <= 0: @@ -546,6 +626,10 @@ def __init__( raise ValueError("virtual_scale must be in [0, 1]") if not 0.0 <= neural_value_scale <= 1.0: raise ValueError("neural_value_scale must be in [0, 1]") + if candidate_backend not in {"auto", "python", "stateful"}: + raise ValueError( + "candidate_backend must be 'auto', 'python', or 'stateful'" + ) self.d_model = d_model self.codebook_sizes = (int(codebook_sizes[0]), int(codebook_sizes[1])) @@ -558,6 +642,7 @@ def __init__( self.sparse_old_candidates = sparse_old_candidates self.sparse_old_pool_size = sparse_old_pool_size self.soft_candidates_forward = soft_candidates_forward + self.candidate_backend: CandidateBackend = candidate_backend self.selector_dim = selector_dim self.token_temperature = token_temperature self.retrieval_temperature = retrieval_temperature @@ -804,10 +889,11 @@ def forward( # 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_hard_candidates( + hard = _build_forward_hard_candidates( hard_tokens, - suffix_k=self.suffix_k, - occurrences_r=self.occurrences_r, + self.suffix_k, + self.occurrences_r, + self.candidate_backend, ) exact_source = hard.source_index exact_mask = hard.mask diff --git a/src/rosa/_stateful_candidates_numba.py b/src/rosa/_stateful_candidates_numba.py index f3cd800..c4500c9 100644 --- a/src/rosa/_stateful_candidates_numba.py +++ b/src/rosa/_stateful_candidates_numba.py @@ -10,6 +10,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Any import numpy as np import torch @@ -749,6 +750,7 @@ def _step_batch_kernel( class CandidateState: """Fixed-capacity tensor state for exact online hard candidates.""" + native_candidate_abi_version: int batch_size: int max_length: int suffix_k: int @@ -777,6 +779,7 @@ class CandidateState: last: np.ndarray size: np.ndarray edge_count: np.ndarray + native_state: Any @dataclass(frozen=True) @@ -819,6 +822,7 @@ def init_candidate_state( hash_shape = (batch_size, hash_capacity) occurrence_shape = (batch_size, max_states, occurrences_r) return CandidateState( + native_candidate_abi_version=1, batch_size=batch_size, max_length=max_length, suffix_k=suffix_k, @@ -847,6 +851,7 @@ def init_candidate_state( last=np.zeros(batch_size, dtype=np.int32), size=np.ones(batch_size, dtype=np.int32), edge_count=np.zeros(batch_size, dtype=np.int32), + native_state=None, ) @@ -859,6 +864,26 @@ def init_candidate_state( } +def _native_candidate_step( # pragma: no cover - optional native companion + state: CandidateState, + cpu_tokens: Tensor, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None: + if state.native_state is False: + return None + if state.native_state is None: + try: + import rosa_native_step # type: ignore[reportMissingImports] + except ModuleNotFoundError: + state.native_state = False + return None + native_type = getattr(rosa_native_step, "NativeCandidateState", None) + if native_type is None: + state.native_state = False + return None + state.native_state = native_type(state) + return state.native_state.step(cpu_tokens.numpy()) + + def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateStep: """Consume one token per row and return exact top-R candidates for K suffixes.""" @@ -877,36 +902,40 @@ def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateS device = tokens.device cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() - source, match_length, state_id, frequency, count = _step_batch_kernel( - cpu_tokens.numpy(), - state.position, - 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 += 1 + native_output = _native_candidate_step(state, cpu_tokens) + if native_output is None: + source, match_length, state_id, frequency, count = _step_batch_kernel( + cpu_tokens.numpy(), + state.position, + 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 += 1 + else: + source, match_length, state_id, frequency, count = native_output slots = state.suffix_k * state.occurrences_r slot_index = np.arange(slots, dtype=np.int32)[None, :] diff --git a/tests/test_rosa.py b/tests/test_rosa.py index 1dc6c42..00ce94a 100644 --- a/tests/test_rosa.py +++ b/tests/test_rosa.py @@ -1,16 +1,20 @@ from __future__ import annotations +import copy import random import unittest +from unittest.mock import patch import torch import torch.nn.functional as F +import rosa from rosa import ( NULL_KIND, ROSA, VIRTUAL_KIND, _balance_kl, + _build_forward_hard_candidates, _gather_sequence, _st_categorical, _virtual_pool_single, @@ -195,6 +199,7 @@ def test_constructor_validations(self) -> None: (dict(d_model=4, learned_residual_scale=-0.1), "learned_residual_scale"), (dict(d_model=4, virtual_scale=1.1), "virtual_scale"), (dict(d_model=4, neural_value_scale=2.0), "neural_value_scale"), + (dict(d_model=4, candidate_backend="invalid"), "candidate_backend"), ] for kwargs, pattern in invalid_calls: with ( @@ -266,6 +271,143 @@ def make_model(self, **overrides) -> ROSA: kwargs.update(overrides) return ROSA(**kwargs) + def assert_nested_equal(self, actual, expected, name: str) -> None: + if isinstance(actual, torch.Tensor): + self.assertTrue(torch.equal(actual, expected), name) + elif isinstance(actual, tuple): + self.assertEqual(len(actual), len(expected), name) + for index, (actual_item, expected_item) in enumerate( + zip(actual, expected, strict=True) + ): + self.assert_nested_equal(actual_item, expected_item, f"{name}[{index}]") + elif isinstance(actual, dict): + self.assertEqual(actual.keys(), expected.keys(), name) + for key in actual: + self.assert_nested_equal(actual[key], expected[key], f"{name}.{key}") + else: + self.assertEqual(actual, expected, name) + + def test_python_and_stateful_backends_match_all_fields_outputs_and_gradients( + self, + ) -> None: + torch.manual_seed(20260811) + tokens = torch.randint(6, (2, 13)) + eager_hard = _build_forward_hard_candidates(tokens, 4, 3, "python") + stateful_hard = _build_forward_hard_candidates(tokens, 4, 3, "stateful") + for name in eager_hard.__dataclass_fields__: + self.assertTrue( + torch.equal(getattr(eager_hard, name), getattr(stateful_hard, name)), + name, + ) + + python_model = self.make_model( + candidate_backend="python", + learned_residual_scale=1.0, + neural_value_scale=1.0, + ) + stateful_model = copy.deepcopy(python_model) + stateful_model.candidate_backend = "stateful" + z_python = torch.randn(2, 13, 8, requires_grad=True) + z_stateful = z_python.detach().clone().requires_grad_() + logits_python = factor_logits_from_tokens( + tokens, (2, 3), hi=0.2, lo=-0.1, requires_grad=True + ) + logits_stateful = tuple( + item.detach().clone().requires_grad_() for item in logits_python + ) + python_output = python_model(z_python, code_logits=logits_python) + stateful_output = stateful_model(z_stateful, code_logits=logits_stateful) + for name in python_output.__dataclass_fields__: + self.assert_nested_equal( + getattr(stateful_output, name), getattr(python_output, name), name + ) + + python_loss = python_output.updated.square().mean() + sum( + python_output.aux_losses.values() + ) + stateful_loss = stateful_output.updated.square().mean() + sum( + stateful_output.aux_losses.values() + ) + python_loss.backward() + stateful_loss.backward() + assert z_stateful.grad is not None + assert z_python.grad is not None + self.assertTrue(torch.equal(z_stateful.grad, z_python.grad)) + for actual, expected in zip(logits_stateful, logits_python, strict=True): + assert actual.grad is not None + assert expected.grad is not None + self.assertTrue(torch.equal(actual.grad, expected.grad)) + for (actual_name, actual), (expected_name, expected) in zip( + stateful_model.named_parameters(), + python_model.named_parameters(), + strict=True, + ): + self.assertEqual(actual_name, expected_name) + self.assertEqual(actual.grad is None, expected.grad is None, actual_name) + if actual.grad is not None: + assert expected.grad is not None + self.assertTrue(torch.equal(actual.grad, expected.grad), actual_name) + + def test_stateful_forward_does_not_call_eager_or_suffix_write(self) -> None: + tokens = torch.tensor([[0, 1, 0, 2, 0, 1]], dtype=torch.long) + logits = factor_logits_from_tokens(tokens, (2, 3)) + expected = _build_forward_hard_candidates(tokens, 3, 2, "python") + model = self.make_model( + suffix_k=3, + occurrences_r=2, + candidate_backend="stateful", + ) + with ( + patch("rosa.build_hard_candidates", side_effect=AssertionError("eager")), + patch.object( + rosa._OnlineSuffixAutomaton, + "write_current_end", + side_effect=AssertionError("suffix write"), + ), + ): + hard = _build_forward_hard_candidates(tokens, 3, 2, "stateful") + output = model(torch.randn(1, 6, 8), code_logits=logits) + for name in expected.__dataclass_fields__: + self.assertTrue( + torch.equal(getattr(hard, name), getattr(expected, name)), name + ) + self.assertTrue( + torch.equal(output.hard_rosa_source_index, expected.rosa_source_index) + ) + + def test_auto_fallback_is_limited_to_missing_optional_dependencies(self) -> None: + tokens = torch.tensor([[0, 1, 0]], dtype=torch.long) + expected = build_hard_candidates(tokens, 2, 2) + for dependency in ("numba", "numpy"): + missing = ModuleNotFoundError( + f"No module named '{dependency}'", name=dependency + ) + with patch("rosa._build_stateful_hard_candidates", side_effect=missing): + actual = _build_forward_hard_candidates(tokens, 2, 2, "auto") + self.assertTrue( + torch.equal(actual.source_index, expected.source_index), dependency + ) + with self.assertRaisesRegex(RuntimeError, "numba.*extra"): + _build_forward_hard_candidates(tokens, 2, 2, "stateful") + + unrelated = ModuleNotFoundError("No module named 'other'", name="other") + with ( + patch("rosa._build_stateful_hard_candidates", side_effect=unrelated), + self.assertRaises(ModuleNotFoundError), + ): + _build_forward_hard_candidates(tokens, 2, 2, "auto") + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_stateful_full_sequence_matches_eager_on_cuda(self) -> None: + tokens = torch.randint(5, (3, 31), device="cuda") + expected = _build_forward_hard_candidates(tokens, 5, 3, "python") + actual = _build_forward_hard_candidates(tokens, 5, 3, "stateful") + for name in expected.__dataclass_fields__: + with self.subTest(name=name): + self.assertTrue( + torch.equal(getattr(actual, name), getattr(expected, name)) + ) + def test_zero_residual_is_exact_rosa_even_with_virtuals(self) -> None: tokens = torch.tensor( [[0, 1, 0, 2, 0, 1, 0, 3, 0], [4, 4, 2, 4, 4, 2, 1, 4, 4]], diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py index d54608f..32569c9 100644 --- a/tests/test_stateful_candidates.py +++ b/tests/test_stateful_candidates.py @@ -45,6 +45,7 @@ def assert_matches_oracle( suffix_k=suffix_k, occurrences_r=occurrences_r, ) + state.native_state = False steps: list[CandidateStep] = [] boundary = tokens.shape[1] if split_at is None else split_at for position in range(boundary): From cd09e7e4fe766f2000ae55ea7dd133bfbc8bddbf Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 07:59:41 +0800 Subject: [PATCH 22/29] Add native rich prefill and ragged stepping --- native/README.md | 12 +- native/src/rosa_native_step.cpp | 169 ++++++++- native/tests/candidate_smoke.py | 113 ++++++ src/rosa/_stateful_candidates_numba.py | 506 +++++++++++++++++++++++-- tests/test_stateful_candidates.py | 191 ++++++++++ 5 files changed, 957 insertions(+), 34 deletions(-) diff --git a/native/README.md b/native/README.md index 56b7794..d1d5b56 100644 --- a/native/README.md +++ b/native/README.md @@ -23,10 +23,14 @@ suffixes, R occurrences les plus récentes, fréquences non bornées et tags LCT capacité peut être détectée via la présence de `NativeCandidateState` et `candidate_abi_version == 1`. -Cette première ABI riche expose `step`, `reset` global et `position`. Elle -n'expose pas encore de préremplissage riche ni de reset/continuation masqué par -ligne; ces opérations nécessitent un contrat de positions par ligne distinct -de `CandidateState.position`. +L'ABI riche conserve `step`, `reset` global et `position`, et détecte par +capacité les extensions `prefill`, `step_masked`, `reset_masked` et +`positions`. `prefill` émet les cinq tableaux natifs à chaque position dans un +seul appel C++ et laisse l'état continuable. Le mode ragged possède une position +par ligne; les chemins uniformes et ragged sont volontairement incompatibles +afin qu'une seule autorité de position existe à tout instant. Le wrapper Python +retombe exactement sur Numba lorsqu'un ancien wheel ABI 1 ne fournit pas ces +méthodes optionnelles. ## Installation et utilisation diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp index eb90e6b..133ac0e 100644 --- a/native/src/rosa_native_step.cpp +++ b/native/src/rosa_native_step.cpp @@ -822,6 +822,22 @@ class NativeCandidateState { edge_capacity_ = edge_token_.shape(1); hash_capacity_ = hash_state_.shape(1); validate_shapes(); + ragged_mode_ = py::hasattr(state_, "ragged_mode") && + py::cast(state_.attr("ragged_mode")); + positions_ = py::array_t(batch_); + std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, + position_); + if (py::hasattr(state_, "positions")) { + positions_ = bind("positions"); + if (!vector_shape(positions_, batch_)) + throw py::value_error( + "positions must be contiguous int64 [batch_size]"); + } else if (ragged_mode_) { + throw py::value_error("ragged candidate state requires positions"); + } + for (int64_t b = 0; b < batch_; ++b) + if (positions_.data()[b] < 0 || positions_.data()[b] > max_length_) + throw py::value_error("candidate positions are outside capacity"); occupied_slots_.resize(batch_); for (int64_t b = 0; b < batch_; ++b) for (int64_t slot = 0; slot < hash_capacity_; ++slot) @@ -830,6 +846,8 @@ class NativeCandidateState { } py::tuple step(py::array tokens_object) { + if (ragged_mode_) + throw std::runtime_error("uniform step is unavailable on a ragged candidate state"); if (!py::isinstance>(tokens_object)) throw py::type_error("tokens must have dtype int64"); if ((tokens_object.flags() & py::array::c_style) == 0 || @@ -861,22 +879,164 @@ class NativeCandidateState { candidate_frequency.mutable_data() + b * slots); } ++position_; + std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, + position_); state_.attr("position") = py::int_(position_); return py::make_tuple(source, match_length, state_id, candidate_frequency, count); } void reset() { + if (ragged_mode_) + throw std::runtime_error("uniform reset is unavailable on a ragged candidate state"); { py::gil_scoped_release release; for (int64_t b = 0; b < batch_; ++b) reset_row(b); } position_ = 0; + std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, + int64_t{0}); state_.attr("position") = py::int_(0); } + py::tuple step_masked(py::array tokens_object, py::array active_object, + py::array reset_object) { + if (!ragged_mode_) + throw std::runtime_error("step_masked requires a ragged candidate state"); + if (!py::isinstance>(tokens_object) || + (tokens_object.flags() & py::array::c_style) == 0 || + tokens_object.ndim() != 1 || tokens_object.shape(0) != batch_) + throw py::value_error("tokens must be contiguous int64 [batch_size]"); + const bool active_bool = py::isinstance>(active_object); + const bool active_u8 = py::isinstance>(active_object); + const bool reset_bool = py::isinstance>(reset_object); + const bool reset_u8 = py::isinstance>(reset_object); + if ((!active_bool && !active_u8) || (!reset_bool && !reset_u8)) + throw py::type_error("active and reset must have dtype bool or uint8"); + if ((active_object.flags() & py::array::c_style) == 0 || + (reset_object.flags() & py::array::c_style) == 0 || + active_object.ndim() != 1 || reset_object.ndim() != 1 || + active_object.shape(0) != batch_ || reset_object.shape(0) != batch_) + throw py::value_error("active and reset must be contiguous [batch_size]"); + auto tokens = py::cast>(tokens_object); + std::vector active(batch_), reset(batch_); + for (int64_t b = 0; b < batch_; ++b) { + active[b] = active_bool ? static_cast(active_object.data())[b] + : static_cast(active_object.data())[b] != 0; + reset[b] = reset_bool ? static_cast(reset_object.data())[b] + : static_cast(reset_object.data())[b] != 0; + const int64_t next_position = reset[b] ? 0 : positions_.data()[b]; + if (active[b] && next_position < 0) + throw std::runtime_error("candidate position must be non-negative"); + if (active[b] && next_position >= max_length_) + throw std::runtime_error("candidate state capacity exceeded"); + } + const int64_t slots = suffix_k_ * occurrences_r_; + py::array_t source({batch_, slots}), match_length({batch_, slots}), + state_id({batch_, slots}), candidate_frequency({batch_, slots}); + py::array_t count(batch_); + std::fill(source.mutable_data(), source.mutable_data() + batch_ * slots, int64_t{-1}); + std::fill(match_length.mutable_data(), match_length.mutable_data() + batch_ * slots, int64_t{0}); + std::fill(state_id.mutable_data(), state_id.mutable_data() + batch_ * slots, int64_t{-1}); + std::fill(candidate_frequency.mutable_data(), candidate_frequency.mutable_data() + batch_ * slots, int64_t{0}); + std::fill(count.mutable_data(), count.mutable_data() + batch_, int32_t{0}); + { + py::gil_scoped_release release; + for (int64_t b = 0; b < batch_; ++b) { + if (!active[b]) + continue; + if (reset[b]) { + reset_row(b); + positions_.mutable_data()[b] = 0; + } + const int64_t position = positions_.data()[b]; + count.mutable_data()[b] = step_row( + b, tokens.data()[b], position, source.mutable_data() + b * slots, + match_length.mutable_data() + b * slots, + state_id.mutable_data() + b * slots, + candidate_frequency.mutable_data() + b * slots); + positions_.mutable_data()[b] = position + 1; + } + } + return py::make_tuple(source, match_length, state_id, candidate_frequency, + count); + } + + void reset_masked(py::array reset_object) { + if (!ragged_mode_) + throw std::runtime_error("reset_masked requires a ragged candidate state"); + const bool reset_bool = py::isinstance>(reset_object); + const bool reset_u8 = py::isinstance>(reset_object); + if ((!reset_bool && !reset_u8)) + throw py::type_error("reset must have dtype bool or uint8"); + if ((reset_object.flags() & py::array::c_style) == 0 || + reset_object.ndim() != 1 || reset_object.shape(0) != batch_) + throw py::value_error("reset must be contiguous [batch_size]"); + std::vector reset(batch_); + for (int64_t b = 0; b < batch_; ++b) + reset[b] = reset_bool ? static_cast(reset_object.data())[b] + : static_cast(reset_object.data())[b] != 0; + { + py::gil_scoped_release release; + for (int64_t b = 0; b < batch_; ++b) + if (reset[b]) { + reset_row(b); + positions_.mutable_data()[b] = 0; + } + } + } + + py::tuple prefill(py::array tokens_object) { + if (ragged_mode_) + throw std::runtime_error("prefill is unavailable on a ragged candidate state"); + if (!py::isinstance>(tokens_object)) + throw py::type_error("tokens must have dtype int64"); + if ((tokens_object.flags() & py::array::c_style) == 0 || + tokens_object.ndim() != 2 || tokens_object.shape(0) != batch_) + throw py::value_error( + "tokens must be contiguous int64 [batch_size, sequence_length]"); + if (position_ != 0) + throw std::runtime_error("prefill requires an empty candidate state"); + auto tokens = py::cast>(tokens_object); + const int64_t sequence_length = tokens.shape(1); + if (sequence_length > max_length_) + throw std::runtime_error("candidate state capacity exceeded"); + const int64_t slots = suffix_k_ * occurrences_r_; + py::array_t source({batch_, sequence_length, slots}), + match_length({batch_, sequence_length, slots}), + state_id({batch_, sequence_length, slots}), + candidate_frequency({batch_, sequence_length, slots}); + py::array_t count({batch_, sequence_length}); + const int64_t output_size = batch_ * sequence_length * slots; + 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_ * sequence_length, int32_t{0}); + { + py::gil_scoped_release release; + for (int64_t position = 0; position < sequence_length; ++position) + for (int64_t b = 0; b < batch_; ++b) { + 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, + source.mutable_data() + output_at, + match_length.mutable_data() + output_at, + state_id.mutable_data() + output_at, + candidate_frequency.mutable_data() + output_at); + } + } + position_ = sequence_length; + std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, + position_); + state_.attr("position") = py::int_(position_); + return py::make_tuple(source, match_length, state_id, candidate_frequency, + count); + } + int64_t position() const { return position_; } + py::array_t positions() const { return positions_; } private: template py::array_t bind(const char *name) { @@ -1221,13 +1381,14 @@ class NativeCandidateState { py::object state_; py::array_t history_, edge_token_, hash_token_, - occurrences_, frequency_, lazy_prefix_, lazy_delta_; + occurrences_, frequency_, lazy_prefix_, lazy_delta_, positions_; py::array_t head_, edge_target_, edge_next_, hash_state_, hash_edge_, suffix_link_, length_, left_, right_, parent_, occurrence_size_, lazy_size_, stack_, last_, size_, edge_count_; std::vector> occupied_slots_; int64_t batch_, max_length_, suffix_k_, occurrences_r_, position_, state_capacity_, edge_capacity_, hash_capacity_; + bool ragged_mode_ = false; }; PYBIND11_MODULE(rosa_native_step, m) { @@ -1243,6 +1404,10 @@ PYBIND11_MODULE(rosa_native_step, m) { .def(py::init(), py::keep_alive<1, 2>()) .def("step", &NativeCandidateState::step) .def("reset", &NativeCandidateState::reset) - .def_property_readonly("position", &NativeCandidateState::position); + .def("step_masked", &NativeCandidateState::step_masked) + .def("reset_masked", &NativeCandidateState::reset_masked) + .def("prefill", &NativeCandidateState::prefill) + .def_property_readonly("position", &NativeCandidateState::position) + .def_property_readonly("positions", &NativeCandidateState::positions); m.attr("candidate_abi_version") = py::int_(1); } diff --git a/native/tests/candidate_smoke.py b/native/tests/candidate_smoke.py index fbc0c17..d2b97df 100644 --- a/native/tests/candidate_smoke.py +++ b/native/tests/candidate_smoke.py @@ -11,7 +11,10 @@ from rosa._stateful_candidates_numba import ( CandidateStep, forward_candidates_step, + forward_candidates_step_masked, init_candidate_state, + prefill_candidates, + reset_candidates_masked, ) @@ -83,6 +86,87 @@ def main() -> None: column = np.ascontiguousarray(random_tokens[:, position].numpy()) assert_step_equal(continuation.step(column), expected) + # One native prefill call emits every historical field and leaves a + # continuation-compatible state, including clone-heavy/random inputs. + prefix_length = 137 + prefill_oracle = init_candidate_state(7, 193, suffix_k=7, occurrences_r=5) + prefill_oracle.native_state = False + expected_prefill = prefill_candidates( + prefill_oracle, random_tokens[:, :prefix_length] + ) + prefill_state = init_candidate_state(7, 193, suffix_k=7, occurrences_r=5) + native_prefill = rosa_native_step.NativeCandidateState(prefill_state) + actual_prefill = native_prefill.prefill( + np.ascontiguousarray(random_tokens[:, :prefix_length].numpy()) + ) + expected_arrays = ( + expected_prefill.source_index.numpy(), + expected_prefill.match_length.numpy(), + expected_prefill.state_id.numpy(), + expected_prefill.frequency.numpy(), + expected_prefill.mask.sum(dim=2).numpy().astype(np.int32), + ) + for actual, expected in zip(actual_prefill, expected_arrays, strict=True): + assert np.array_equal(actual, expected) + assert native_prefill.position == prefill_state.position == prefix_length + assert np.array_equal(native_prefill.positions, np.full(7, prefix_length)) + for position in range(prefix_length, random_tokens.shape[1]): + expected = forward_candidates_step(prefill_oracle, random_tokens[:, position]) + assert_step_equal( + native_prefill.step( + np.ascontiguousarray(random_tokens[:, position].numpy()) + ), + expected, + ) + + # Ragged active/reset/recycle follows independent per-row positions. The + # Numba fallback is the exact oracle and inactive rows emit empty outputs. + ragged_oracle = init_candidate_state( + 5, 17, suffix_k=5, occurrences_r=3, ragged=True + ) + ragged_oracle.native_state = False + ragged_state = init_candidate_state(5, 17, suffix_k=5, occurrences_r=3, ragged=True) + native_ragged = rosa_native_step.NativeCandidateState(ragged_state) + ragged_generator = np.random.default_rng(1804) + for iteration in range(73): + token_values = ragged_generator.integers(-2, 7, size=5, dtype=np.int64) + active = ragged_generator.random(5) < 0.72 + reset = np.logical_and(active, ragged_generator.random(5) < 0.13) + full = np.logical_and(active, ragged_oracle.positions >= 17) + reset = np.logical_or(reset, full) + expected = forward_candidates_step_masked( + ragged_oracle, + torch.from_numpy(token_values), + torch.from_numpy(active), + torch.from_numpy(reset), + ) + actual = native_ragged.step_masked(token_values, active, reset) + assert_step_equal(actual, expected) + assert np.array_equal(native_ragged.positions, ragged_oracle.positions) + reset_rows = np.array([True, False, True, False, True]) + reset_candidates_masked(ragged_oracle, torch.from_numpy(reset_rows)) + native_ragged.reset_masked(reset_rows) + assert np.array_equal(native_ragged.positions, ragged_oracle.positions) + assert np.array_equal(ragged_state.size, ragged_oracle.size) + + for action in ( + lambda: native_ragged.step(np.zeros(5, dtype=np.int64)), + lambda: native_ragged.prefill(np.zeros((5, 1), dtype=np.int64)), + lambda: rosa_native_step.NativeCandidateState( + init_candidate_state(1, 2) + ).step_masked( + np.zeros(1, dtype=np.int64), + np.ones(1, dtype=np.bool_), + np.zeros(1, dtype=np.bool_), + ), + ): + try: + action() + except RuntimeError: + pass + else: + raise AssertionError("uniform/ragged candidate modes were mixed") + state = init_candidate_state(2, 4, suffix_k=3, occurrences_r=2) state_ref = weakref.ref(state) native = rosa_native_step.NativeCandidateState(state) @@ -91,6 +175,35 @@ def main() -> None: assert state_ref() is not None native.step(np.array([1, 2], dtype=np.int64)) + # ABI 1 compatibility: pre-positions uniform states remain accepted. + legacy = init_candidate_state(1, 2) + del legacy.positions + legacy_native = rosa_native_step.NativeCandidateState(legacy) + legacy_native.step(np.array([1], dtype=np.int64)) + + invalid_position = init_candidate_state(1, 2, ragged=True) + invalid_position.positions[0] = -1 + try: + rosa_native_step.NativeCandidateState(invalid_position) + except ValueError as error: + assert "positions" in str(error) + else: + raise AssertionError("negative candidate position was accepted") + + runtime_position = init_candidate_state(1, 2, ragged=True) + runtime_native = rosa_native_step.NativeCandidateState(runtime_position) + runtime_position.positions[0] = -1 + try: + runtime_native.step_masked( + np.array([1], dtype=np.int64), + np.array([True]), + np.array([False]), + ) + except RuntimeError as error: + assert "non-negative" in str(error) + else: + raise AssertionError("mutated negative candidate position was consumed") + for invalid in ( np.zeros(2, dtype=np.int32), np.zeros((2, 1), dtype=np.int64), diff --git a/src/rosa/_stateful_candidates_numba.py b/src/rosa/_stateful_candidates_numba.py index c4500c9..2b27fe0 100644 --- a/src/rosa/_stateful_candidates_numba.py +++ b/src/rosa/_stateful_candidates_numba.py @@ -746,6 +746,212 @@ def _step_batch_kernel( return source, match_length, state_id, candidate_frequency, count +@njit(cache=True, nogil=True) +def _reset_candidate_rows_kernel( + reset: np.ndarray, + head: np.ndarray, + hash_state: np.ndarray, + suffix_link: np.ndarray, + length: np.ndarray, + lct_left: np.ndarray, + lct_right: np.ndarray, + lct_parent: np.ndarray, + occurrence_size: np.ndarray, + frequency: np.ndarray, + lazy_size: np.ndarray, + lazy_delta: np.ndarray, + last: np.ndarray, + size: np.ndarray, + edge_count: np.ndarray, + positions: np.ndarray, +) -> None: # pragma: no cover - executed as compiled Numba code + for batch_index in range(reset.shape[0]): + if not reset[batch_index]: + continue + used_states = int(size[batch_index]) + for node in range(used_states): + head[batch_index, node] = -1 + suffix_link[batch_index, node] = -1 + length[batch_index, node] = 0 + lct_left[batch_index, node] = -1 + lct_right[batch_index, node] = -1 + lct_parent[batch_index, node] = -1 + occurrence_size[batch_index, node] = 0 + frequency[batch_index, node] = 0 + lazy_size[batch_index, node] = 0 + lazy_delta[batch_index, node] = 0 + for slot in range(hash_state.shape[1]): + hash_state[batch_index, slot] = -1 + last[batch_index] = 0 + size[batch_index] = 1 + edge_count[batch_index] = 0 + positions[batch_index] = 0 + + +@njit(cache=True, nogil=True) +def _step_masked_batch_kernel( + tokens: np.ndarray, + active: np.ndarray, + positions: 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 + slots = suffix_k * occurrences_r + batch_size = tokens.shape[0] + source = np.full((batch_size, slots), -1, dtype=np.int64) + match_length = np.zeros((batch_size, slots), dtype=np.int64) + state_id = np.full((batch_size, slots), -1, dtype=np.int64) + candidate_frequency = np.zeros((batch_size, slots), dtype=np.int64) + count = np.zeros(batch_size, dtype=np.int32) + for batch_index in range(batch_size): + if not active[batch_index]: + continue + row_count, row_last, row_size, row_edge_count = _step_row( + int(tokens[batch_index]), + int(positions[batch_index]), + 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]), + source[batch_index], + match_length[batch_index], + state_id[batch_index], + candidate_frequency[batch_index], + ) + count[batch_index] = row_count + last[batch_index] = row_last + size[batch_index] = row_size + edge_count[batch_index] = row_edge_count + positions[batch_index] += 1 + return source, match_length, state_id, candidate_frequency, count + + +@njit(cache=True, nogil=True) +def _prefill_candidate_kernel( + tokens: 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 + slots = suffix_k * occurrences_r + source = np.full((batch_size, sequence_length, slots), -1, dtype=np.int64) + match_length = np.zeros((batch_size, sequence_length, slots), dtype=np.int64) + state_id = np.full((batch_size, sequence_length, slots), -1, dtype=np.int64) + candidate_frequency = np.zeros((batch_size, sequence_length, slots), dtype=np.int64) + count = np.zeros((batch_size, sequence_length), dtype=np.int32) + for position in range(sequence_length): + for batch_index in range(batch_size): + 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]), + source[batch_index, position], + match_length[batch_index, position], + state_id[batch_index, position], + candidate_frequency[batch_index, position], + ) + count[batch_index, position] = row_count + 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.""" @@ -756,6 +962,8 @@ class CandidateState: suffix_k: int occurrences_r: int position: int + ragged_mode: bool + positions: np.ndarray history: np.ndarray head: np.ndarray edge_token: np.ndarray @@ -803,6 +1011,7 @@ def init_candidate_state( *, suffix_k: int = 16, occurrences_r: int = 4, + ragged: bool = False, ) -> CandidateState: """Allocate an exact bounded-candidate state backed by CPU tensors.""" @@ -828,6 +1037,8 @@ def init_candidate_state( suffix_k=suffix_k, occurrences_r=occurrences_r, position=0, + ragged_mode=ragged, + positions=np.zeros(batch_size, dtype=np.int64), history=np.empty((batch_size, max_length), dtype=np.int64), head=np.full(state_shape, -1, dtype=np.int32), edge_token=np.empty(edge_shape, dtype=np.int64), @@ -884,19 +1095,90 @@ def _native_candidate_step( # pragma: no cover - optional native companion return state.native_state.step(cpu_tokens.numpy()) -def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateStep: - """Consume one token per row and return exact top-R candidates for K suffixes.""" +def _native_candidate_call( # pragma: no cover - optional native companion + state: CandidateState, + method: str, + *args: np.ndarray, +) -> Any | None: + """Call a post-ABI-1 capability, falling back for an older installed wheel.""" + if state.native_state is False: + return None + if state.native_state is None: + try: + import rosa_native_step # type: ignore[reportMissingImports] + except ModuleNotFoundError: + state.native_state = False + return None + native_type = getattr(rosa_native_step, "NativeCandidateState", None) + if native_type is None: + state.native_state = False + return None + state.native_state = native_type(state) + native_method = getattr(state.native_state, method, None) + if native_method is None: + state.native_state = False + return None + return native_method(*args) + + +def _validate_candidate_tokens( + state: CandidateState, tokens: Tensor, *, sequence: bool = False +) -> tuple[Tensor, bool]: if not isinstance(state, CandidateState): raise TypeError("state must be a CandidateState") if not isinstance(tokens, Tensor): raise TypeError("tokens must be a Tensor") - if tokens.ndim == 0 and state.batch_size == 1: + scalar = tokens.ndim == 0 and state.batch_size == 1 and not sequence + if scalar: tokens = tokens.unsqueeze(0) - if tokens.ndim != 1 or tokens.shape[0] != state.batch_size: - raise ValueError("tokens must have shape [batch_size]") + expected_ndim = 2 if sequence else 1 + if tokens.ndim != expected_ndim or tokens.shape[0] != state.batch_size: + shape = "[batch_size, sequence_length]" if sequence else "[batch_size]" + raise ValueError(f"tokens must have shape {shape}") if tokens.dtype not in _INTEGER_DTYPES: raise TypeError("tokens must use an integer dtype") + return tokens, scalar + + +def _candidate_step_from_arrays( + state: CandidateState, + arrays: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], + device: torch.device, +) -> CandidateStep: + source, match_length, state_id, frequency, count = arrays + slots = state.suffix_k * state.occurrences_r + slot_shape = (1,) * count.ndim + (slots,) + slot_index = np.arange(slots, dtype=np.int32).reshape(slot_shape) + mask = slot_index < count[..., None] + rosa_source = source[..., 0].copy() + rosa_length = match_length[..., 0].copy() + rosa_slot = np.where(count > 0, 0, -1).astype(np.int64) + rosa_predicted = np.full(count.shape, -1, dtype=np.int64) + for index in np.ndindex(count.shape): + if count[index] > 0: + batch_index = index[0] + source_position = int(rosa_source[index]) + rosa_predicted[index] = state.history[batch_index, source_position + 1] + return CandidateStep( + source_index=torch.from_numpy(source).to(device), + match_length=torch.from_numpy(match_length).to(device), + state_id=torch.from_numpy(state_id).to(device), + frequency=torch.from_numpy(frequency).to(device), + mask=torch.from_numpy(mask).to(device), + rosa_slot=torch.from_numpy(rosa_slot).to(device), + rosa_source_index=torch.from_numpy(rosa_source).to(device), + rosa_match_length=torch.from_numpy(rosa_length).to(device), + rosa_predicted_tokens=torch.from_numpy(rosa_predicted).to(device), + ) + + +def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateStep: + """Consume one token per row and return exact top-R candidates for K suffixes.""" + + tokens, _ = _validate_candidate_tokens(state, tokens) + if state.ragged_mode: + raise RuntimeError("uniform step is unavailable on a ragged candidate state") if state.position >= state.max_length: raise RuntimeError("candidate state capacity exceeded") @@ -934,31 +1216,199 @@ def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateS state.edge_count, ) state.position += 1 + state.positions.fill(state.position) else: source, match_length, state_id, frequency, count = native_output + state.positions.fill(state.position) + return _candidate_step_from_arrays( + state, (source, match_length, state_id, frequency, count), device + ) - slots = state.suffix_k * state.occurrences_r - slot_index = np.arange(slots, dtype=np.int32)[None, :] - mask = slot_index < count[:, None] - rosa_source = source[:, 0].copy() - rosa_length = match_length[:, 0].copy() - rosa_slot = np.where(count > 0, 0, -1).astype(np.int64) - rosa_predicted = np.full(state.batch_size, -1, dtype=np.int64) - for batch_index in range(state.batch_size): - if count[batch_index] > 0: - source_position = int(rosa_source[batch_index]) - rosa_predicted[batch_index] = state.history[ - batch_index, source_position + 1 - ] - return CandidateStep( - source_index=torch.from_numpy(source).to(device), - match_length=torch.from_numpy(match_length).to(device), - state_id=torch.from_numpy(state_id).to(device), - frequency=torch.from_numpy(frequency).to(device), - mask=torch.from_numpy(mask).to(device), - rosa_slot=torch.from_numpy(rosa_slot).to(device), - rosa_source_index=torch.from_numpy(rosa_source).to(device), - rosa_match_length=torch.from_numpy(rosa_length).to(device), - rosa_predicted_tokens=torch.from_numpy(rosa_predicted).to(device), +def reset_candidates_masked(state: CandidateState, reset: Tensor) -> None: + """Reset selected ragged rows without reallocating their fixed-capacity storage.""" + + if not isinstance(state, CandidateState): + raise TypeError("state must be a CandidateState") + if not state.ragged_mode: + raise RuntimeError("reset_masked requires a ragged candidate state") + if not isinstance(reset, Tensor): + raise TypeError("reset must be a Tensor") + if reset.ndim == 0 and state.batch_size == 1: + reset = reset.unsqueeze(0) + if reset.ndim != 1 or reset.shape[0] != state.batch_size: + raise ValueError("reset must have shape [batch_size]") + if reset.dtype not in (torch.bool, torch.uint8): + raise TypeError("reset must have dtype bool or uint8") + cpu_reset = reset.detach().to(device="cpu", dtype=torch.bool).contiguous() + native = _native_candidate_call(state, "reset_masked", cpu_reset.numpy()) + if native is None: # pragma: no branch - native capability is optional + _reset_candidate_rows_kernel( + cpu_reset.numpy(), + state.head, + state.hash_state, + state.suffix_link, + state.length, + state.lct_left, + state.lct_right, + state.lct_parent, + state.occurrence_size, + state.frequency, + state.lazy_size, + state.lazy_delta, + state.last, + state.size, + state.edge_count, + state.positions, + ) + + +def forward_candidates_step_masked( + state: CandidateState, + tokens: Tensor, + active: Tensor, + reset: Tensor | None = None, +) -> CandidateStep: + """Consume tokens only on active rows, optionally recycling active rows first.""" + + tokens, _ = _validate_candidate_tokens(state, tokens) + if not state.ragged_mode: + raise RuntimeError("step_masked requires a ragged candidate state") + if not isinstance(active, Tensor): + raise TypeError("active must be a Tensor") + if active.ndim == 0 and state.batch_size == 1: + active = active.unsqueeze(0) + if active.ndim != 1 or active.shape[0] != state.batch_size: + raise ValueError("active must have shape [batch_size]") + if active.dtype not in (torch.bool, torch.uint8): + raise TypeError("active must have dtype bool or uint8") + if reset is None: + reset = torch.zeros_like(active, dtype=torch.bool) + if not isinstance(reset, Tensor): + raise TypeError("reset must be a Tensor") + if reset.ndim == 0 and state.batch_size == 1: + reset = reset.unsqueeze(0) + if reset.ndim != 1 or reset.shape[0] != state.batch_size: + raise ValueError("reset must have shape [batch_size]") + if reset.dtype not in (torch.bool, torch.uint8): + raise TypeError("reset must have dtype bool or uint8") + device = tokens.device + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + cpu_active = active.detach().to(device="cpu", dtype=torch.bool).contiguous() + cpu_reset = reset.detach().to(device="cpu", dtype=torch.bool).contiguous() + effective_reset = np.logical_and(cpu_active.numpy(), cpu_reset.numpy()) + future_positions = np.where(effective_reset, 0, state.positions) + if np.any(np.logical_and(cpu_active.numpy(), future_positions < 0)): + raise RuntimeError("candidate position must be non-negative") + if np.any(np.logical_and(cpu_active.numpy(), future_positions >= state.max_length)): + raise RuntimeError("candidate state capacity exceeded") + native_output = _native_candidate_call( + state, + "step_masked", + cpu_tokens.numpy(), + cpu_active.numpy(), + cpu_reset.numpy(), ) + if native_output is None: + _reset_candidate_rows_kernel( + effective_reset, + state.head, + state.hash_state, + state.suffix_link, + state.length, + state.lct_left, + state.lct_right, + state.lct_parent, + state.occurrence_size, + state.frequency, + state.lazy_size, + state.lazy_delta, + state.last, + state.size, + state.edge_count, + state.positions, + ) + native_output = _step_masked_batch_kernel( + cpu_tokens.numpy(), + cpu_active.numpy(), + state.positions, + 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, + ) + return _candidate_step_from_arrays(state, native_output, device) + + +def prefill_candidates(state: CandidateState, tokens: Tensor) -> CandidateStep: + """Consume a complete uniform sequence and emit candidates at every position.""" + + tokens, _ = _validate_candidate_tokens(state, tokens, sequence=True) + 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() + native_output = _native_candidate_call(state, "prefill", cpu_tokens.numpy()) + if native_output is None: + native_output = _prefill_candidate_kernel( + cpu_tokens.numpy(), + 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_stateful_candidates.py b/tests/test_stateful_candidates.py index 32569c9..d85015a 100644 --- a/tests/test_stateful_candidates.py +++ b/tests/test_stateful_candidates.py @@ -15,6 +15,12 @@ from rosa._stateful_candidates_numba import ( CandidateState, CandidateStep, + forward_candidates_step_masked, + prefill_candidates, + reset_candidates_masked, +) +from rosa._stateful_candidates_numba import ( + init_candidate_state as init_candidate_state_internal, ) _FIELDS = ( @@ -104,6 +110,112 @@ def test_batched_continuation_matches_offline_oracle(self) -> None: split_at=41, ) + def test_prefill_emits_every_position_and_continues(self) -> None: + generator = torch.Generator().manual_seed(1804) + prefix = torch.randint(5, (5, 61), generator=generator) + continuation = torch.randint(5, (5, 17), generator=generator) + state = init_candidate_state(5, 78, suffix_k=6, occurrences_r=4) + state.native_state = False + actual = prefill_candidates(state, prefix) + expected = build_hard_candidates(prefix, suffix_k=6, occurrences_r=4) + for field in _FIELDS: + self.assertTrue( + torch.equal(getattr(actual, field), getattr(expected, field)) + ) + self.assertEqual(state.position, prefix.shape[1]) + self.assertEqual(state.positions.tolist(), [prefix.shape[1]] * 5) + + steps = [ + forward_candidates_step(state, continuation[:, position]) + for position in range(continuation.shape[1]) + ] + full = torch.cat((prefix, continuation), dim=1) + expected_full = build_hard_candidates(full, suffix_k=6, occurrences_r=4) + for field in _FIELDS: + actual_tail = torch.stack([getattr(step, field) for step in steps], dim=1) + self.assertTrue( + torch.equal( + actual_tail, getattr(expected_full, field)[:, prefix.shape[1] :] + ) + ) + + def test_ragged_inactive_reset_capacity_and_recycle(self) -> None: + state = init_candidate_state_internal( + 3, 5, suffix_k=4, occurrences_r=3, ragged=True + ) + state.native_state = False + histories: list[list[int]] = [[], [], []] + schedule = ( + ([1, 8, 3], [1, 1, 0], [0, 0, 0]), + ([2, 9, 4], [1, 0, 1], [0, 0, 0]), + ([1, 7, 3], [1, 1, 1], [0, 0, 0]), + ([2, 6, 5], [1, 1, 0], [0, 1, 0]), + ([1, 6, 3], [1, 1, 1], [0, 0, 1]), + ) + for token_values, active_values, reset_values in schedule: + tokens = torch.tensor(token_values) + active = torch.tensor(active_values, dtype=torch.bool) + reset = torch.tensor(reset_values, dtype=torch.bool) + actual = forward_candidates_step_masked(state, tokens, active, reset) + for batch_index in range(3): + if not active_values[batch_index]: + self.assertFalse(bool(actual.mask[batch_index].any())) + continue + if reset_values[batch_index]: + histories[batch_index].clear() + histories[batch_index].append(token_values[batch_index]) + oracle = build_hard_candidates( + torch.tensor([histories[batch_index]]), + suffix_k=4, + occurrences_r=3, + ) + for field in _FIELDS: + self.assertTrue( + torch.equal( + getattr(actual, field)[batch_index], + getattr(oracle, field)[0, -1], + ), + field, + ) + self.assertEqual(state.positions.tolist(), [5, 2, 1]) + reset_candidates_masked(state, torch.tensor([True, False, True])) + self.assertEqual(state.positions.tolist(), [0, 2, 0]) + recycled = forward_candidates_step_masked( + state, + torch.tensor([4, 0, 4]), + torch.tensor([True, False, True]), + ) + self.assertFalse(bool(recycled.mask.any())) + + full = init_candidate_state_internal(1, 1, ragged=True) + full.native_state = False + forward_candidates_step_masked(full, torch.tensor([1]), torch.tensor([True])) + with self.assertRaisesRegex(RuntimeError, "capacity"): + forward_candidates_step_masked( + full, torch.tensor([2]), torch.tensor([True]) + ) + # Reset and consume in one operation must be allowed at full capacity. + forward_candidates_step_masked( + full, + torch.tensor([2]), + torch.tensor([True]), + torch.tensor([True]), + ) + + def test_uniform_and_ragged_modes_cannot_mix(self) -> None: + uniform = init_candidate_state(1, 2) + ragged = init_candidate_state_internal(1, 2, ragged=True) + with self.assertRaisesRegex(RuntimeError, "ragged"): + forward_candidates_step_masked( + uniform, torch.tensor([1]), torch.tensor([True]) + ) + with self.assertRaisesRegex(RuntimeError, "ragged"): + forward_candidates_step(ragged, torch.tensor([1])) + with self.assertRaisesRegex(RuntimeError, "ragged"): + prefill_candidates(ragged, torch.tensor([[1]])) + with self.assertRaisesRegex(RuntimeError, "ragged"): + reset_candidates_masked(uniform, torch.tensor([True])) + def test_scalar_batch_and_validation(self) -> None: state = init_candidate_state(1, 2, suffix_k=2, occurrences_r=2) first = forward_candidates_step(state, torch.tensor(7)) @@ -135,6 +247,85 @@ def test_scalar_batch_and_validation(self) -> None: with self.assertRaisesRegex(TypeError, "Tensor"): forward_candidates_step(state, object()) # type: ignore[arg-type] + ragged = init_candidate_state_internal(1, 3, ragged=True) + ragged.native_state = False + with self.assertRaisesRegex(TypeError, "CandidateState"): + reset_candidates_masked(object(), torch.tensor([True])) # type: ignore[arg-type] + with self.assertRaisesRegex(TypeError, "Tensor"): + reset_candidates_masked(ragged, object()) # type: ignore[arg-type] + reset_candidates_masked(ragged, torch.tensor(True)) + with self.assertRaisesRegex(ValueError, "shape"): + reset_candidates_masked(ragged, torch.tensor([True, False])) + with self.assertRaisesRegex(TypeError, "bool or uint8"): + reset_candidates_masked(ragged, torch.tensor([1])) + + with self.assertRaisesRegex(TypeError, "active"): + forward_candidates_step_masked( + ragged, + torch.tensor([1]), + object(), # type: ignore[arg-type] + ) + forward_candidates_step_masked(ragged, torch.tensor([1]), torch.tensor(True)) + with self.assertRaisesRegex(ValueError, "active"): + forward_candidates_step_masked( + ragged, torch.tensor([1]), torch.tensor([True, False]) + ) + with self.assertRaisesRegex(TypeError, "active"): + forward_candidates_step_masked(ragged, torch.tensor([1]), torch.tensor([1])) + with self.assertRaisesRegex(TypeError, "reset"): + forward_candidates_step_masked( + ragged, + torch.tensor([1]), + torch.tensor([True]), + object(), # type: ignore[arg-type] + ) + forward_candidates_step_masked( + ragged, + torch.tensor([1]), + torch.tensor([True]), + torch.tensor(True), + ) + with self.assertRaisesRegex(ValueError, "reset"): + forward_candidates_step_masked( + ragged, + torch.tensor([1]), + torch.tensor([True]), + torch.tensor([True, False]), + ) + with self.assertRaisesRegex(TypeError, "reset"): + forward_candidates_step_masked( + ragged, + torch.tensor([1]), + torch.tensor([True]), + torch.tensor([1]), + ) + + prefill_state = init_candidate_state_internal(1, 1) + prefill_state.native_state = False + with self.assertRaisesRegex(ValueError, "shape"): + prefill_candidates(prefill_state, torch.tensor([1])) + with self.assertRaisesRegex(TypeError, "integer"): + prefill_candidates(prefill_state, torch.tensor([[1.0]])) + prefill_candidates(prefill_state, torch.tensor([[1]])) + with self.assertRaisesRegex(RuntimeError, "empty"): + prefill_candidates(prefill_state, torch.tensor([[1]])) + too_long = init_candidate_state_internal(1, 1) + with self.assertRaisesRegex(RuntimeError, "capacity"): + prefill_candidates(too_long, torch.tensor([[1, 2]])) + + negative = init_candidate_state_internal(1, 2, ragged=True) + negative.native_state = False + negative.positions[0] = -1 + with self.assertRaisesRegex(RuntimeError, "non-negative"): + forward_candidates_step_masked( + negative, torch.tensor([1]), torch.tensor([True]) + ) + + old_capability = init_candidate_state_internal(1, 2) + old_capability.native_state = object() + prefill_candidates(old_capability, torch.tensor([[0, 1]])) + self.assertIs(old_capability.native_state, False) + def test_public_wrapper_reports_missing_numba(self) -> None: with patch.dict(sys.modules): sys.modules.pop("rosa._stateful_candidates_numba", None) From a71a3bda5b65d2f7664f1098b41d8e05cb5fac56 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:00:01 +0800 Subject: [PATCH 23/29] Unify top-one rich and ragged inference state --- src/rosa/__init__.py | 330 +++++++++++++++++++++++++++++--- tests/test_unified_inference.py | 222 +++++++++++++++++++++ 2 files changed, 526 insertions(+), 26 deletions(-) create mode 100644 tests/test_unified_inference.py diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py index 22945c8..12b9113 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -25,6 +25,7 @@ "EXACT_KIND", "NULL_KIND", "ROSA", + "InferenceOutput", "ROSAInferenceState", "VIRTUAL_KIND", "HardCandidates", @@ -1201,6 +1202,19 @@ def combine_losses( InferenceBackend = Literal["auto", "python", "numba"] +InferenceMode = Literal["top1", "rich"] + + +@dataclass(frozen=True, slots=True) +class InferenceOutput: + """Result returned by the unified inference-state methods. + + ``candidates`` is ``None`` in top-1 mode, a ``CandidateStep`` for one + decoding step, and ``HardCandidates`` for a prefill in rich mode. + """ + + predicted_tokens: Tensor + candidates: object | None @dataclass(slots=True) @@ -1214,14 +1228,47 @@ class ROSAInferenceState: batch_size: int max_length: int backend: Literal["python", "numba"] + mode: InferenceMode + ragged: bool + suffix_k: int + occurrences_r: int _impl: object = field(repr=False) @property def position(self) -> int: """Number of tokens consumed by every batch row.""" + if self.ragged: + raise AttributeError( + "position is undefined for ragged states; use positions" + ) return int(cast(Any, self._impl).position) + @property + def positions(self) -> Tensor: + """Consumed-token counts for every row, always returned as a copy.""" + + 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 torch.full((self.batch_size,), self.position, dtype=torch.long) + + def step( + self, + tokens: Tensor, + active: Tensor | None = None, + reset: Tensor | None = None, + ) -> InferenceOutput: + """Consume one token per selected row using the configured mode.""" + + return _inference_step(self, tokens, active=active, reset=reset) + + def prefill(self, tokens: Tensor) -> InferenceOutput: + """Consume an initial dense context and return every step result.""" + + return _inference_prefill(self, tokens) + def reset(self) -> None: """Reset all batch rows while retaining the configured capacity.""" @@ -1229,6 +1276,10 @@ def reset(self) -> None: self.batch_size, self.max_length, self.backend, + self.mode, + self.ragged, + self.suffix_k, + self.occurrences_r, ) @@ -1236,7 +1287,38 @@ def _make_inference_impl( batch_size: int, max_length: int, backend: Literal["python", "numba"], + mode: InferenceMode = "top1", + ragged: bool = False, + suffix_k: int = 16, + occurrences_r: int = 4, ) -> object: + if mode == "rich": + if backend != "numba": + raise ValueError( # pragma: no cover - validated by public initializer + "rich inference requires backend='auto' or 'numba'" + ) + if ragged: + return init_candidate_state( + batch_size, + max_length, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ragged=True, + ) + return init_candidate_state( + batch_size, + max_length, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ) + if ragged: + if backend != "numba": + raise ValueError( # pragma: no cover - validated by public initializer + "ragged inference requires backend='auto' or 'numba'" + ) + from .ragged import init_ragged_state + + return init_ragged_state(batch_size, max_length) if backend == "python": return _init_python_inference_state(batch_size, max_length) try: @@ -1256,6 +1338,7 @@ def init_candidate_state( *, suffix_k: int = 16, occurrences_r: int = 4, + ragged: bool = False, ) -> Any: """Allocate the exact bounded rich-candidate inference state. @@ -1272,6 +1355,14 @@ def init_candidate_state( "rich stateful candidates require the 'numba' extra" ) from error raise # pragma: no cover - unrelated optional import failure + if ragged: + return initialize( + batch_size, + max_length, + suffix_k=suffix_k, + occurrences_r=occurrences_r, + ragged=True, + ) return initialize( batch_size, max_length, @@ -1283,6 +1374,14 @@ def init_candidate_state( def forward_candidates_step(state: object, tokens: Tensor) -> Any: """Consume one token and return exact bounded hard candidates.""" + if isinstance(state, ROSAInferenceState): + if state.mode != "rich": + raise ValueError("forward_candidates_step requires a rich state") + candidates = state.step(tokens).candidates + if candidates is None: # pragma: no cover - guarded by mode + raise RuntimeError("rich inference did not return candidates") + return candidates + try: from ._stateful_candidates_numba import forward_candidates_step as step except ModuleNotFoundError as error: @@ -1299,13 +1398,17 @@ def init_inference_state( max_length: int = 8192, *, backend: InferenceBackend = "auto", + mode: InferenceMode = "top1", + ragged: bool = False, + suffix_k: int = 16, + occurrences_r: int = 4, ) -> ROSAInferenceState: - """Create an explicit state for exact top-1 autoregressive ROSA inference. + """Create an exact top-1 or rich autoregressive ROSA inference state. ``backend="auto"`` selects the Link-Cut Tree Numba backend when the optional dependency is installed and otherwise uses the exact Python - fallback. Capacity is fixed so memory use and failure behavior remain - predictable in production. + fallback for uniform top-1 inference. Rich and ragged modes require the + Numba extra. Capacity is fixed so memory use remains predictable. """ if batch_size <= 0: @@ -1314,9 +1417,39 @@ def init_inference_state( raise ValueError("max_length must be > 0") if backend not in {"auto", "python", "numba"}: raise ValueError("backend must be 'auto', 'python', or 'numba'") + if mode not in {"top1", "rich"}: + raise ValueError("mode must be 'top1' or 'rich'") + if suffix_k <= 0: + raise ValueError("suffix_k must be > 0") + if occurrences_r <= 0: + raise ValueError("occurrences_r must be > 0") + if (mode == "rich" or ragged) and backend == "python": + feature = "rich" if mode == "rich" else "ragged" + raise ValueError(f"{feature} inference does not support backend='python'") selected: Literal["python", "numba"] if backend == "auto": + if mode == "rich" or ragged: + selected = "numba" + impl = _make_inference_impl( + batch_size, + max_length, + selected, + mode, + ragged, + suffix_k, + occurrences_r, + ) + return ROSAInferenceState( + batch_size, + max_length, + selected, + mode, + ragged, + suffix_k, + occurrences_r, + impl, + ) try: impl = _make_inference_impl(batch_size, max_length, "numba") selected = "numba" @@ -1325,20 +1458,30 @@ def init_inference_state( selected = "python" else: selected = backend - impl = _make_inference_impl(batch_size, max_length, selected) - return ROSAInferenceState(batch_size, max_length, selected, impl) - - -def forward_step(state: ROSAInferenceState, token: Tensor) -> Tensor: - """Consume one token per batch row and return exact ROSA predictions. + impl = _make_inference_impl( + batch_size, + max_length, + selected, + mode, + ragged, + suffix_k, + occurrences_r, + ) + return ROSAInferenceState( + batch_size, + max_length, + selected, + mode, + ragged, + suffix_k, + occurrences_r, + impl, + ) - ``token`` must be an integer tensor shaped ``[batch_size]``. A scalar is - also accepted when ``batch_size == 1`` and produces a scalar prediction. - The state remains on CPU; the prediction is returned on the token device. - """ - if not isinstance(state, ROSAInferenceState): - raise TypeError("state must be a ROSAInferenceState") +def _normalize_step_tokens( + state: ROSAInferenceState, token: Tensor +) -> tuple[Tensor, bool]: if not isinstance(token, Tensor): raise TypeError("token must be a torch.Tensor") squeeze = token.ndim == 0 @@ -1354,26 +1497,114 @@ def forward_step(state: ROSAInferenceState, token: Tensor) -> Tensor: torch.int64, }: raise TypeError("token must use an integer dtype") + return token, squeeze - if state.backend == "python": + +def _rich_ragged_step( + state: ROSAInferenceState, + tokens: Tensor, + active: Tensor | None, + reset: Tensor | None, +) -> Any: + from ._stateful_candidates_numba import forward_candidates_step_masked + + active_mask = ( + torch.ones(state.batch_size, dtype=torch.bool) if active is None else active + ) + return forward_candidates_step_masked( + cast(Any, state._impl), tokens, active_mask, reset + ) + + +def _inference_step( + state: ROSAInferenceState, + token: Tensor, + *, + active: Tensor | None = None, + reset: Tensor | None = None, +) -> InferenceOutput: + if not isinstance(state, ROSAInferenceState): # pragma: no cover - method only + raise TypeError("state must be a ROSAInferenceState") + token, squeeze = _normalize_step_tokens(state, token) + if not state.ragged and (active is not None or reset is not None): + raise ValueError("active and reset are only valid for ragged states") + + candidates: object | None = None + if state.mode == "rich": + if state.ragged: + candidates = _rich_ragged_step(state, token, active, reset) + else: + candidates = forward_candidates_step(state._impl, token) + output = cast(Any, candidates).rosa_predicted_tokens + elif state.ragged: + output = cast(Any, state._impl).step(token, active=active, reset=reset) + elif state.backend == "python": output = _python_forward_step(cast(_PythonInferenceState, state._impl), token) else: from ._stateful_numba import _forward_step output = _forward_step(cast(Any, state._impl), token) - return output[0] if squeeze else output + if squeeze: + output = output[0] + if candidates is not None: + from ._stateful_candidates_numba import CandidateStep + + candidate_data = cast(Any, candidates) + candidates = CandidateStep( + *( + getattr(candidate_data, name)[0] + for name in candidate_data.__dataclass_fields__ + ) + ) + return InferenceOutput(output, candidates) -def prefill(state: ROSAInferenceState, tokens: Tensor) -> Tensor: - """Consume an initial context and return exact ROSA predictions. +def forward_step(state: ROSAInferenceState, token: Tensor) -> Tensor: + """Consume one token per batch row and return exact ROSA predictions. - The state must be empty. ``tokens`` uses shape ``[batch_size, N]``; a - one-dimensional context is accepted when ``batch_size == 1``. The Numba - backend fuses the complete replay into one compiled call. + ``token`` must be an integer tensor shaped ``[batch_size]``. A scalar is + also accepted when ``batch_size == 1`` and produces a scalar prediction. + The state remains on CPU; the prediction is returned on the token device. """ if not isinstance(state, ROSAInferenceState): raise TypeError("state must be a ROSAInferenceState") + return state.step(token).predicted_tokens + + +def _empty_rich_candidates( + state: ROSAInferenceState, device: torch.device +) -> HardCandidates: + slots = state.suffix_k * state.occurrences_r + candidate_shape = (state.batch_size, 0, slots) + rosa_shape = (state.batch_size, 0) + source = torch.empty(candidate_shape, dtype=torch.long, device=device) + zeros = torch.empty(candidate_shape, dtype=torch.long, device=device) + return HardCandidates( + source, + zeros.clone(), + zeros.clone(), + zeros.clone(), + torch.empty(candidate_shape, dtype=torch.bool, device=device), + torch.empty(rosa_shape, dtype=torch.long, device=device), + torch.empty(rosa_shape, dtype=torch.long, device=device), + torch.empty(rosa_shape, dtype=torch.long, device=device), + torch.empty(rosa_shape, dtype=torch.long, device=device), + ) + + +def _stack_rich_steps(steps: list[Any], dim: int) -> HardCandidates: + return HardCandidates( + *( + torch.stack([getattr(step, name) for step in steps], dim=dim) + for name in HardCandidates.__dataclass_fields__ + ) + ) + + +def _inference_prefill(state: ROSAInferenceState, tokens: Tensor) -> InferenceOutput: + if not isinstance(state, ROSAInferenceState): # pragma: no cover - method only + raise TypeError("state must be a ROSAInferenceState") if not isinstance(tokens, Tensor): raise TypeError("tokens must be a torch.Tensor") squeeze = tokens.ndim == 1 @@ -1389,12 +1620,13 @@ def prefill(state: ROSAInferenceState, tokens: Tensor) -> Tensor: torch.int64, }: raise TypeError("tokens must use an integer dtype") - if state.position != 0: + if bool(torch.any(state.positions != 0)): raise RuntimeError("prefill requires an empty inference state") if tokens.shape[1] > state.max_length: raise RuntimeError("inference state capacity exceeded") - if state.backend == "python": + candidates: HardCandidates | None = None + if state.mode == "top1" and not state.ragged and state.backend == "python": backend_state = cast(_PythonInferenceState, state._impl) if tokens.shape[1] == 0: output = torch.empty(tokens.shape, dtype=torch.long, device=tokens.device) @@ -1406,8 +1638,54 @@ def prefill(state: ROSAInferenceState, tokens: Tensor) -> Tensor: ], dim=1, ) - else: + elif state.mode == "top1" and not state.ragged: from ._stateful_numba import _prefill output = _prefill(cast(Any, state._impl), tokens) - return output[0] if squeeze else output + elif state.mode == "rich" and not state.ragged: + from ._stateful_candidates_numba import prefill_candidates + + rich_output = prefill_candidates(cast(Any, state._impl), tokens) + output = rich_output.rosa_predicted_tokens + candidates = HardCandidates( + *( + getattr(rich_output, name) + for name in HardCandidates.__dataclass_fields__ + ) + ) + elif tokens.shape[1] == 0: + output = torch.empty(tokens.shape, dtype=torch.long, device=tokens.device) + if state.mode == "rich": + candidates = _empty_rich_candidates(state, tokens.device) + else: + step_outputs = [ + state.step(tokens[:, position]) for position in range(tokens.shape[1]) + ] + output = torch.stack([item.predicted_tokens for item in step_outputs], dim=1) + if state.mode == "rich": + candidates = _stack_rich_steps( + [cast(Any, item.candidates) for item in step_outputs], 1 + ) + if squeeze: + output = output[0] + if candidates is not None: + candidates = HardCandidates( + *( + getattr(candidates, name)[0] + for name in candidates.__dataclass_fields__ + ) + ) + return InferenceOutput(output, candidates) + + +def prefill(state: ROSAInferenceState, tokens: Tensor) -> Tensor: + """Consume an initial context and return exact ROSA predictions. + + The state must be empty. ``tokens`` uses shape ``[batch_size, N]``; a + one-dimensional context is accepted when ``batch_size == 1``. The Numba + backend fuses the complete replay into one compiled call. + """ + + if not isinstance(state, ROSAInferenceState): + raise TypeError("state must be a ROSAInferenceState") + return state.prefill(tokens).predicted_tokens diff --git a/tests/test_unified_inference.py b/tests/test_unified_inference.py new file mode 100644 index 0000000..2e112a6 --- /dev/null +++ b/tests/test_unified_inference.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +import unittest +from typing import Any, cast + +import torch + +from rosa import ( + HardCandidates, + InferenceOutput, + build_hard_candidates, + forward_candidates_step, + forward_step, + init_candidate_state, + init_inference_state, + prefill, + reference_rosa, +) + + +class TestUnifiedInferenceState(unittest.TestCase): + 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 + ) + expected, _, _ = reference_rosa(tokens) + for mode in ("top1", "rich"): + with self.subTest(mode=mode): + state = init_inference_state( + 2, + 6, + mode=mode, # type: ignore[arg-type] + suffix_k=3, + occurrences_r=2, + ) + initial = state.prefill(tokens[:, :4]) + self.assertIsInstance(initial, InferenceOutput) + continuation = torch.stack( + [state.step(tokens[:, index]).predicted_tokens for index in (4, 5)], + dim=1, + ) + self.assertTrue( + torch.equal( + torch.cat((initial.predicted_tokens, continuation), dim=1), + expected, + ) + ) + self.assertEqual(state.positions.tolist(), [6, 6]) + if mode == "top1": + self.assertIsNone(initial.candidates) + else: + self.assertIsInstance(initial.candidates, HardCandidates) + oracle = build_hard_candidates( + tokens[:, :4], suffix_k=3, occurrences_r=2 + ) + assert isinstance(initial.candidates, HardCandidates) + for name in HardCandidates.__dataclass_fields__: + self.assertTrue( + torch.equal( + getattr(initial.candidates, name), getattr(oracle, name) + ), + name, + ) + state.reset() + self.assertEqual(state.positions.tolist(), [0, 0]) + self.assertTrue( + torch.equal(state.prefill(tokens).predicted_tokens, expected) + ) + + def test_legacy_top1_and_candidate_wrappers_remain_compatible(self) -> None: + tokens = torch.tensor([[0, 1, 0, 2]], dtype=torch.long) + top1 = init_inference_state(1, 4, backend="numba") + self.assertTrue(torch.equal(prefill(top1, tokens), reference_rosa(tokens)[0])) + + rich = init_inference_state(1, 4, mode="rich", suffix_k=2, occurrences_r=2) + candidate = forward_candidates_step(rich, tokens[:, 0]) + self.assertEqual(tuple(candidate.source_index.shape), (1, 4)) + self.assertEqual(forward_step(rich, tokens[:, 1]).shape, (1,)) + + old_state = init_candidate_state(1, 1, suffix_k=2, occurrences_r=2) + old_output = forward_candidates_step(old_state, torch.tensor([7])) + self.assertEqual(tuple(old_output.source_index.shape), (1, 4)) + + def test_positions_are_copies_and_uniform_position_is_preserved(self) -> None: + uniform = init_inference_state(2, 2) + forward_step(uniform, torch.tensor([1, 2])) + positions = uniform.positions + positions[0] = 99 + self.assertEqual(uniform.positions.tolist(), [1, 1]) + self.assertEqual(uniform.position, 1) + + ragged = init_inference_state(2, 2, ragged=True) + ragged.step(torch.tensor([1, 2]), active=torch.tensor([True, False])) + ragged_positions = ragged.positions + ragged_positions[0] = 99 + self.assertEqual(ragged.positions.tolist(), [1, 0]) + with self.assertRaisesRegex(AttributeError, "positions"): + _ = ragged.position + + def test_top1_does_not_allocate_rich_occurrence_storage(self) -> None: + top1 = init_inference_state(2, 32, backend="numba") + rich = init_inference_state(2, 32, mode="rich") + self.assertFalse(hasattr(top1._impl, "occurrences")) + self.assertTrue(hasattr(rich._impl, "occurrences")) + + def test_ragged_top1_reset_and_inactive_rows(self) -> None: + state = init_inference_state(2, 3, ragged=True) + first = state.step(torch.tensor([5, 7]), active=torch.tensor([True, False])) + self.assertEqual(first.predicted_tokens.tolist(), [-1, -1]) + second = state.step( + torch.tensor([5, 7]), + active=torch.tensor([True, True]), + reset=torch.tensor([True, False]), + ) + self.assertEqual(second.predicted_tokens.tolist(), [-1, -1]) + self.assertEqual(state.positions.tolist(), [1, 1]) + state.reset() + self.assertEqual(state.positions.tolist(), [0, 0]) + + def test_rich_ragged_matches_independent_row_states(self) -> None: + state = init_inference_state( + 2, + 4, + mode="rich", + ragged=True, + suffix_k=2, + occurrences_r=2, + ) + row_states = [ + init_candidate_state(1, 4, suffix_k=2, occurrences_r=2) for _ in range(2) + ] + schedule = [ + (torch.tensor([0, 4]), torch.tensor([1, 0]), torch.tensor([0, 0])), + (torch.tensor([1, 4]), torch.tensor([1, 1]), torch.tensor([0, 0])), + (torch.tensor([0, 5]), torch.tensor([1, 1]), torch.tensor([0, 1])), + ] + for tokens, active, reset in schedule: + actual = state.step(tokens, active=active.bool(), reset=reset.bool()) + assert actual.candidates is not None + for row in range(2): + if not bool(active[row]): + self.assertEqual(actual.predicted_tokens[row].item(), -1) + continue + if bool(reset[row]): + row_states[row] = init_candidate_state( + 1, 4, suffix_k=2, occurrences_r=2 + ) + expected = forward_candidates_step(row_states[row], tokens[row]) + for name in expected.__dataclass_fields__: + self.assertTrue( + torch.equal( + getattr(actual.candidates, name)[row], + getattr(expected, name)[0], + ), + name, + ) + self.assertEqual(state.positions.tolist(), [3, 1]) + + def test_rich_scalar_prefill(self) -> None: + rich = init_inference_state(1, 3, mode="rich", suffix_k=2, occurrences_r=2) + scalar_step = rich.step(torch.tensor(0)) + self.assertEqual(scalar_step.predicted_tokens.ndim, 0) + self.assertEqual(cast(Any, scalar_step.candidates).source_index.ndim, 1) + rich.reset() + output = rich.prefill(torch.tensor([0, 1, 0])) + self.assertEqual(tuple(output.predicted_tokens.shape), (3,)) + assert isinstance(output.candidates, HardCandidates) + self.assertEqual(tuple(output.candidates.source_index.shape), (3, 4)) + rich.reset() + empty = rich.prefill(torch.empty(0, dtype=torch.long)) + assert isinstance(empty.candidates, HardCandidates) + self.assertEqual(tuple(empty.candidates.source_index.shape), (0, 4)) + + def test_mode_validation_and_uniform_masks(self) -> None: + with self.assertRaisesRegex(ValueError, "mode"): + init_inference_state(1, mode="other") # type: ignore[arg-type] + with self.assertRaisesRegex(ValueError, "suffix_k"): + init_inference_state(1, mode="rich", suffix_k=0) + with self.assertRaisesRegex(ValueError, "occurrences_r"): + init_inference_state(1, mode="rich", occurrences_r=0) + with self.assertRaisesRegex(ValueError, "does not support"): + init_inference_state(1, mode="rich", backend="python") + with self.assertRaisesRegex(ValueError, "ragged"): + init_inference_state(1, ragged=True, backend="python") + state = init_inference_state(1) + with self.assertRaisesRegex(ValueError, "ragged"): + state.step(torch.tensor([1]), active=torch.tensor([True])) + with self.assertRaisesRegex(ValueError, "rich"): + forward_candidates_step(state, torch.tensor([1])) + + rich_ragged = init_inference_state(1, 2, mode="rich", ragged=True) + with self.assertRaisesRegex(TypeError, "active"): + rich_ragged.step(torch.tensor([1]), active=object()) # type: ignore[arg-type] + with self.assertRaisesRegex(ValueError, "active"): + rich_ragged.step(torch.tensor([1]), active=torch.tensor([True, False])) + with self.assertRaisesRegex(TypeError, "active"): + rich_ragged.step(torch.tensor([1]), active=torch.tensor([1])) + + empty_ragged = init_inference_state(1, 2, mode="rich", ragged=True) + empty = empty_ragged.prefill(torch.empty((1, 0), dtype=torch.long)) + self.assertEqual(tuple(empty.predicted_tokens.shape), (1, 0)) + self.assertIsInstance(empty.candidates, HardCandidates) + + rich_prefill = init_inference_state( + 1, 3, mode="rich", ragged=True, suffix_k=2, occurrences_r=2 + ).prefill(torch.tensor([[0, 1, 0]])) + self.assertEqual(rich_prefill.predicted_tokens.tolist(), [[-1, -1, 1]]) + self.assertIsInstance(rich_prefill.candidates, HardCandidates) + + top1_ragged = init_inference_state(1, 2, ragged=True) + top1_empty = top1_ragged.prefill(torch.empty((1, 0), dtype=torch.long)) + self.assertEqual(tuple(top1_empty.predicted_tokens.shape), (1, 0)) + self.assertIsNone(top1_empty.candidates) + top1_full = init_inference_state(1, 2, ragged=True).prefill( + torch.tensor([[0, 0]]) + ) + self.assertEqual(top1_full.predicted_tokens.tolist(), [[-1, 0]]) + self.assertIsNone(top1_full.candidates) + + +if __name__ == "__main__": + unittest.main() From 9a7b9b7dfa144bda521a69f891981261ebc1e299 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:05:49 +0800 Subject: [PATCH 24/29] Document unified rich inference --- README.md | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/README.md b/README.md index 256bed8..706b85d 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,8 @@ PEP 517-compatible Python package manager. │ └── rosa │ ├── __init__.py │ ├── _numba_backend.py +│ ├── _stateful_candidates_numba.py +│ ├── ragged.py │ └── _stateful_numba.py └── tests ├── __init__.py @@ -133,6 +135,35 @@ not be shared concurrently between decoding requests. `forward_step` implements exact top-1 ROSA; rich multi-candidate training remains on the full-sequence `ROSA` path. +The same facade also exposes exact rich candidates and independently advancing +rows without allocating rich storage for top-1 states: + +```python +rich = init_inference_state( + batch_size=8, + max_length=32_768, + mode="rich", + ragged=True, + suffix_k=16, + occurrences_r=4, +) + +result = rich.step( + token_ids, + active=active_rows, + reset=recycled_rows, +) +predicted = result.predicted_tokens +candidates = result.candidates +positions = rich.positions +``` + +`mode="top1"` remains the default. Uniform states expose scalar `position`; +all states expose a copied `positions` tensor. Rich and ragged modes require +the `numba` extra and automatically use compatible native companion methods +when installed. Legacy `forward_step`, `prefill`, `init_candidate_state`, and +`forward_candidates_step` remain supported. + ## Quick start ```python @@ -156,6 +187,7 @@ model = ROSA( learned_residual_scale=0.0, virtual_scale=0.0, neural_value_scale=0.0, + candidate_backend="auto", # stateful rich backend, Python oracle fallback ) z_a = torch.randn(batch_size, sequence_length, d_model, requires_grad=True) From 2cc75ccf5066731c2d6bd8efdd3ab6c6b9a3bb12 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:43:45 +0800 Subject: [PATCH 25/29] Use fused rich prefill in ROSA forward --- src/rosa/__init__.py | 52 ++++++++++++++------------------------------ tests/test_rosa.py | 50 +++++++++++++++++++++++++++++++++++++----- 2 files changed, 61 insertions(+), 41 deletions(-) diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py index 12b9113..c8c693e 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -399,52 +399,32 @@ def _build_stateful_hard_candidates( suffix_k: int, occurrences_r: int, ) -> HardCandidates: - """Replay a full sequence through the exact bounded stateful backend.""" + """Prefill a full sequence through the exact bounded stateful backend.""" - from ._stateful_candidates_numba import forward_candidates_step as step from ._stateful_candidates_numba import init_candidate_state as initialize + from ._stateful_candidates_numba import prefill_candidates - cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() - batch_size, sequence_length = cpu_tokens.shape + squeeze = tokens.ndim == 1 + if squeeze: + tokens = tokens.unsqueeze(0) + if tokens.ndim != 2: + raise ValueError("tokens must have shape [N] or [B, N]") + + batch_size, sequence_length = tokens.shape state = initialize( batch_size, sequence_length, suffix_k=suffix_k, occurrences_r=occurrences_r, ) - steps = [ - step(state, cpu_tokens[:, position]) for position in range(sequence_length) - ] - stacked = { - name: torch.stack([getattr(item, name) for item in steps], dim=1) - for name in HardCandidates.__dataclass_fields__ - } - candidate_fields = torch.stack( - [ - stacked["source_index"], - stacked["match_length"], - stacked["state_id"], - stacked["frequency"], - ] - ).to(tokens.device) - rosa_fields = torch.stack( - [ - stacked["rosa_slot"], - stacked["rosa_source_index"], - stacked["rosa_match_length"], - stacked["rosa_predicted_tokens"], - ] - ).to(tokens.device) + candidates = prefill_candidates(state, tokens) + result = HardCandidates( + *(getattr(candidates, name) for name in HardCandidates.__dataclass_fields__) + ) + if not squeeze: + return result return HardCandidates( - source_index=candidate_fields[0], - match_length=candidate_fields[1], - state_id=candidate_fields[2], - frequency=candidate_fields[3], - mask=stacked["mask"].to(tokens.device), - rosa_slot=rosa_fields[0], - rosa_source_index=rosa_fields[1], - rosa_match_length=rosa_fields[2], - rosa_predicted_tokens=rosa_fields[3], + *(getattr(result, name)[0] for name in result.__dataclass_fields__) ) diff --git a/tests/test_rosa.py b/tests/test_rosa.py index 00ce94a..9b20c8c 100644 --- a/tests/test_rosa.py +++ b/tests/test_rosa.py @@ -349,6 +349,8 @@ def test_python_and_stateful_backends_match_all_fields_outputs_and_gradients( self.assertTrue(torch.equal(actual.grad, expected.grad), actual_name) def test_stateful_forward_does_not_call_eager_or_suffix_write(self) -> None: + from rosa._stateful_candidates_numba import prefill_candidates + tokens = torch.tensor([[0, 1, 0, 2, 0, 1]], dtype=torch.long) logits = factor_logits_from_tokens(tokens, (2, 3)) expected = _build_forward_hard_candidates(tokens, 3, 2, "python") @@ -359,22 +361,60 @@ def test_stateful_forward_does_not_call_eager_or_suffix_write(self) -> None: ) with ( patch("rosa.build_hard_candidates", side_effect=AssertionError("eager")), + patch( + "rosa._stateful_candidates_numba.forward_candidates_step", + side_effect=AssertionError("step"), + ), + patch( + "rosa._stateful_candidates_numba.prefill_candidates", + wraps=prefill_candidates, + ) as prefill_mock, patch.object( rosa._OnlineSuffixAutomaton, "write_current_end", side_effect=AssertionError("suffix write"), ), ): - hard = _build_forward_hard_candidates(tokens, 3, 2, "stateful") output = model(torch.randn(1, 6, 8), code_logits=logits) - for name in expected.__dataclass_fields__: - self.assertTrue( - torch.equal(getattr(hard, name), getattr(expected, name)), name - ) + prefill_mock.assert_called_once() self.assertTrue( torch.equal(output.hard_rosa_source_index, expected.rosa_source_index) ) + def test_stateful_prefill_reuses_full_sequence_candidate_tensors(self) -> None: + from rosa._stateful_candidates_numba import prefill_candidates + + tokens = torch.tensor([[0, 1, 0, 2, 0, 1]], dtype=torch.long) + captured = None + + def capture_prefill(state, full_tokens): + nonlocal captured + captured = prefill_candidates(state, full_tokens) + return captured + + with patch( + "rosa._stateful_candidates_numba.prefill_candidates", + side_effect=capture_prefill, + ): + hard = _build_forward_hard_candidates(tokens, 3, 2, "stateful") + + assert captured is not None + for name in hard.__dataclass_fields__: + self.assertIs(getattr(hard, name), getattr(captured, name), name) + + 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") + actual = _build_forward_hard_candidates(tokens, 3, 2, "stateful") + for name in expected.__dataclass_fields__: + self.assertTrue( + torch.equal(getattr(actual, name), getattr(expected, name)), name + ) + with self.assertRaisesRegex(ValueError, r"\[N\].*\[B, N\]"): + _build_forward_hard_candidates( + torch.zeros((1, 1, 1), dtype=torch.long), 3, 2, "stateful" + ) + def test_auto_fallback_is_limited_to_missing_optional_dependencies(self) -> None: tokens = torch.tensor([[0, 1, 0]], dtype=torch.long) expected = build_hard_candidates(tokens, 2, 2) From d0e93b6f7bf71ef5a38f1329ed99e9cfcb49c0ae Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:04:59 +0800 Subject: [PATCH 26/29] Parallelize native batch prefill --- native/README.md | 24 +++ native/src/rosa_native_step.cpp | 364 +++++++++++++++++++++++++++----- native/tests/candidate_smoke.py | 66 ++++++ 3 files changed, 406 insertions(+), 48 deletions(-) diff --git a/native/README.md b/native/README.md index d1d5b56..2f4d1cb 100644 --- a/native/README.md +++ b/native/README.md @@ -56,6 +56,30 @@ un vecteur NumPy C-contigu `int64` et renvoie le tuple bas niveau `(source, match_length, state_id, frequency, count)`. `reset()` recycle tout le batch en temps proportionnel aux états et slots de hachage réellement occupés. +Chaque état natif crée à la demande un petit pool C++17 persistant, sans +dépendance, qui parallélise les lignes indépendantes des prefill top-1 et +riches. Les prefills +conservent aussi leur chemin série pour les batchs inférieurs à 4, où le réveil d'un +worker coûte plus cher que le calcul mesuré. Les steps uniformes ne l'emploient +qu'à partir d'un batch de 64 afin de conserver le +chemin série peu coûteux des petits batchs. Le pool est limité par le batch, le +nombre de CPU annoncé et 16 threads; `ROSA_NATIVE_THREADS=1` force le chemin +série, et une valeur entière positive fixe une limite plus basse. Une valeur +absente sélectionne automatiquement la limite disponible, tandis qu'une valeur +invalide retombe prudemment à un thread. Seules les chaînes non vides composées +exclusivement de chiffres et représentant une valeur strictement positive sont +acceptées. + +Le pool par état évite le cycle de vie fragile d'un singleton d'extension et +n'appelle jamais Python depuis ses workers. Aucun worker n'est donc créé pour +un état ragged ni pour un batch inférieur à 4 qui n'utilise jamais de prefill. +Son coût éventuel est la création d'au plus 15 threads au premier appel assez +grand; ils sont réutilisés jusqu'à sa destruction. Les appels mutateurs +concurrents sur une même instance sont +sérialisés par un mutex acquis GIL libéré, sans modifier les allocations ni les +tableaux de sortie publics. Les exceptions C++ des workers sont capturées puis +relancées sur le thread appelant. + ## Construction locale isolée Le backend PEP 517 est setuptools, avec pybind11 uniquement comme dépendance de diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp index 133ac0e..ada5383 100644 --- a/native/src/rosa_native_step.cpp +++ b/native/src/rosa_native_step.cpp @@ -2,13 +2,193 @@ #include #include +#include +#include +#include #include +#include +#include +#include +#include #include #include +#include +#include #include namespace py = pybind11; +namespace { + +constexpr size_t kMaximumNativeThreads = 16; + +size_t native_thread_count(int64_t rows) { + if (rows <= 1) + return 1; + size_t available = std::thread::hardware_concurrency(); + if (available == 0) + available = 1; + size_t requested = available; + if (const char *value = std::getenv("ROSA_NATIVE_THREADS")) { + try { + const std::string text(value); + if (text.empty() || + !std::all_of(text.begin(), text.end(), [](unsigned char character) { + return character >= '0' && character <= '9'; + })) + throw std::invalid_argument("thread count must contain only digits"); + size_t consumed = 0; + const unsigned long parsed = std::stoul(text, &consumed); + if (consumed == text.size() && parsed > 0) + requested = static_cast(parsed); + else + requested = 1; + } catch (const std::exception &) { + requested = 1; + } + } + return std::max( + 1, std::min({requested, available, kMaximumNativeThreads, + static_cast(rows)})); +} + +// A pool belongs to one native state. Its workers never touch Python and stay +// alive across step/prefill calls; this avoids singleton shutdown ordering and +// lets ROSA_NATIVE_THREADS be selected independently when each state is made. +class RowThreadPool { +public: + explicit RowThreadPool(int64_t rows) : thread_count_(native_thread_count(rows)) { + try { + for (size_t worker = 1; worker < thread_count_; ++worker) + workers_.emplace_back([this] { worker_loop(); }); + } catch (...) { + // Without this cleanup, unwinding a partially constructed vector of + // joinable std::threads calls std::terminate. Retain a serial pool. + { + std::lock_guard lock(mutex_); + stopping_ = true; + ++generation_; + } + work_ready_.notify_all(); + for (std::thread &worker : workers_) + worker.join(); + workers_.clear(); + thread_count_ = 1; + return; + } + std::unique_lock lock(mutex_); + workers_ready_.wait(lock, + [this] { return ready_workers_ == workers_.size(); }); + } + + RowThreadPool(const RowThreadPool &) = delete; + RowThreadPool &operator=(const RowThreadPool &) = delete; + + size_t worker_count() const { return workers_.size(); } + + ~RowThreadPool() { + { + std::lock_guard lock(mutex_); + stopping_ = true; + ++generation_; + } + work_ready_.notify_all(); + for (std::thread &worker : workers_) + worker.join(); + } + + template + void parallel_for_rows(int64_t rows, int64_t minimum_parallel_rows, + Function &&function) { + if (rows < minimum_parallel_rows || workers_.empty()) { + for (int64_t row = 0; row < rows; ++row) + function(row); + return; + } + + { + std::lock_guard lock(mutex_); + next_row_.store(0, std::memory_order_relaxed); + end_row_ = rows; + cancelled_.store(false, std::memory_order_relaxed); + exception_ = nullptr; + function_ = std::forward(function); + active_workers_ = workers_.size(); + ++generation_; + } + work_ready_.notify_all(); + run_rows(); + { + std::unique_lock lock(mutex_); + work_done_.wait(lock, [this] { return active_workers_ == 0; }); + function_ = nullptr; + if (exception_) + std::rethrow_exception(exception_); + } + } + +private: + void capture_exception() { + cancelled_.store(true, std::memory_order_relaxed); + std::lock_guard lock(mutex_); + if (!exception_) + exception_ = std::current_exception(); + } + + void run_rows() { + try { + while (!cancelled_.load(std::memory_order_relaxed)) { + const int64_t row = next_row_.fetch_add(1, std::memory_order_relaxed); + if (row >= end_row_) + break; + function_(row); + } + } catch (...) { + capture_exception(); + } + } + + void worker_loop() { + size_t observed_generation = 0; + { + std::lock_guard lock(mutex_); + ++ready_workers_; + workers_ready_.notify_one(); + } + for (;;) { + { + std::unique_lock lock(mutex_); + work_ready_.wait(lock, [this, observed_generation] { + return stopping_ || generation_ != observed_generation; + }); + if (stopping_) + return; + observed_generation = generation_; + } + run_rows(); + { + std::lock_guard lock(mutex_); + if (--active_workers_ == 0) + work_done_.notify_one(); + } + } + } + + size_t thread_count_; + std::vector workers_; + std::mutex mutex_; + std::condition_variable work_ready_, work_done_, workers_ready_; + std::function function_; + std::atomic next_row_{0}; + std::atomic cancelled_{false}; + int64_t end_row_ = 0; + size_t generation_ = 0, active_workers_ = 0, ready_workers_ = 0; + bool stopping_ = false; + std::exception_ptr exception_; +}; + +} // namespace + class NativeState { public: explicit NativeState(py::object state) : state_(std::move(state)) { @@ -70,21 +250,25 @@ class NativeState { if (tokens.ndim() != 1 || tokens.shape(0) != batch_) { throw py::value_error("tokens must be contiguous int64 [batch_size]"); } - if (position_ >= max_length_) { - throw std::runtime_error("inference state capacity exceeded"); - } py::array_t output(batch_); const int64_t *in = tokens.data(); int64_t *out = output.mutable_data(); + ensure_pool(64); + std::unique_lock call_lock; { py::gil_scoped_release release; - for (int64_t b = 0; b < batch_; ++b) + call_lock = std::unique_lock(call_mutex_); + if (position_ >= max_length_) + throw std::runtime_error("inference state capacity exceeded"); + parallel_for_rows(64, [&](int64_t b) { out[b] = step_row(b, in[b], position_); + }); } ++position_; std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, position_); state_.attr("position") = py::int_(position_); + call_lock.unlock(); return output; } @@ -115,24 +299,33 @@ class NativeState { ? static_cast(reset_object.data())[b] : static_cast(reset_object.data())[b] != 0; } - for (int64_t b = 0; b < batch_; ++b) { - if (active[b] && !reset[b] && positions_.data()[b] >= max_length_) - throw std::runtime_error("inference state capacity exceeded"); - } py::array_t output(batch_); std::fill(output.mutable_data(), output.mutable_data() + batch_, int64_t{-1}); + std::vector next_positions(batch_); + std::unique_lock call_lock; { py::gil_scoped_release release; + call_lock = std::unique_lock(call_mutex_); + for (int64_t b = 0; b < batch_; ++b) { + if (active[b] && !reset[b] && positions_.data()[b] >= max_length_) + throw std::runtime_error("inference state capacity exceeded"); + } for (int64_t b = 0; b < batch_; ++b) { + const int64_t current_position = + active[b] && reset[b] ? 0 : positions_.data()[b]; + next_positions[b] = current_position; if (!active[b]) continue; if (reset[b]) reset_row(b); - const int64_t position = positions_.data()[b]; - output.mutable_data()[b] = step_row(b, tokens.data()[b], position); - positions_.mutable_data()[b] = position + 1; + output.mutable_data()[b] = + step_row(b, tokens.data()[b], current_position); + next_positions[b] = current_position + 1; } } + std::copy(next_positions.begin(), next_positions.end(), + positions_.mutable_data()); + call_lock.unlock(); return output; } @@ -151,36 +344,56 @@ class NativeState { throw py::value_error( "tokens must be contiguous int64 [batch_size, sequence_length]"); } - if (position_ != 0) { - throw std::runtime_error("prefill requires an empty inference state"); - } const int64_t token_count = tokens.shape(1); - if (token_count > max_length_) { - throw std::runtime_error("inference state capacity exceeded"); - } py::array_t output({batch_, token_count}); - if (token_count == 0) - return output; const int64_t *in = tokens.data(); int64_t *out = output.mutable_data(); + ensure_pool(4); + std::unique_lock call_lock; { py::gil_scoped_release release; - for (int64_t b = 0; b < batch_; ++b) { - prefill_row(b, in + b * token_count, token_count, - out + b * token_count); - } + call_lock = std::unique_lock(call_mutex_); + if (position_ != 0) + throw std::runtime_error("prefill requires an empty inference state"); + if (token_count > max_length_) + throw std::runtime_error("inference state capacity exceeded"); + if (token_count > 0) + parallel_for_rows(4, [&](int64_t b) { + prefill_row(b, in + b * token_count, token_count, + out + b * token_count); + }); } position_ = token_count; std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, position_); state_.attr("position") = py::int_(position_); + call_lock.unlock(); return output; } int64_t position() const { return position_; } py::array_t positions() const { return positions_; } + size_t worker_count() const { + return row_pool_ ? row_pool_->worker_count() : 0; + } private: + void ensure_pool(int64_t minimum_parallel_rows) { + if (batch_ >= minimum_parallel_rows && !row_pool_) + row_pool_ = std::make_unique(batch_); + } + + template + void parallel_for_rows(int64_t minimum_parallel_rows, Function &&function) { + if (batch_ < minimum_parallel_rows) { + for (int64_t b = 0; b < batch_; ++b) + function(b); + return; + } + row_pool_->parallel_for_rows(batch_, minimum_parallel_rows, + std::forward(function)); + } + template py::array_t checked_vector(py::array object, const char *name, const char *dtype) const { @@ -339,7 +552,6 @@ class NativeState { last_.mutable_data()[b] = 0; size_.mutable_data()[b] = 1; edge_count_.mutable_data()[b] = 0; - positions_.mutable_data()[b] = 0; } void replace_transition(int64_t b, int32_t state, int64_t token, @@ -779,6 +991,8 @@ class NativeState { py::array_t lazy_valid_; py::array_t positions_; std::vector> occupied_slots_; + std::unique_ptr row_pool_; + std::mutex call_mutex_; bool ragged_mode_ = false; int64_t batch_, max_length_, position_, state_capacity_, edge_capacity_, hash_capacity_; @@ -853,8 +1067,6 @@ class NativeCandidateState { if ((tokens_object.flags() & py::array::c_style) == 0 || tokens_object.ndim() != 1 || tokens_object.shape(0) != batch_) throw py::value_error("tokens must be contiguous int64 [batch_size]"); - if (position_ >= max_length_) - throw std::runtime_error("candidate state capacity exceeded"); auto tokens = py::cast>(tokens_object); const int64_t slots = suffix_k_ * occurrences_r_; py::array_t source({batch_, slots}), match_length({batch_, slots}), @@ -868,20 +1080,27 @@ class NativeCandidateState { int64_t{-1}); std::fill(candidate_frequency.mutable_data(), candidate_frequency.mutable_data() + batch_ * slots, int64_t{0}); + ensure_pool(64); + std::unique_lock call_lock; { py::gil_scoped_release release; - for (int64_t b = 0; b < batch_; ++b) + call_lock = std::unique_lock(call_mutex_); + if (position_ >= max_length_) + throw std::runtime_error("candidate state capacity exceeded"); + parallel_for_rows(64, [&](int64_t b) { count.mutable_data()[b] = step_row(b, tokens.data()[b], position_, source.mutable_data() + b * slots, match_length.mutable_data() + b * slots, state_id.mutable_data() + b * slots, candidate_frequency.mutable_data() + b * slots); + }); } ++position_; std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, position_); state_.attr("position") = py::int_(position_); + call_lock.unlock(); return py::make_tuple(source, match_length, state_id, candidate_frequency, count); } @@ -889,8 +1108,10 @@ class NativeCandidateState { void reset() { if (ragged_mode_) throw std::runtime_error("uniform reset is unavailable on a ragged candidate state"); + std::unique_lock call_lock; { py::gil_scoped_release release; + call_lock = std::unique_lock(call_mutex_); for (int64_t b = 0; b < batch_; ++b) reset_row(b); } @@ -898,6 +1119,7 @@ class NativeCandidateState { std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, int64_t{0}); state_.attr("position") = py::int_(0); + call_lock.unlock(); } py::tuple step_masked(py::array tokens_object, py::array active_object, @@ -926,11 +1148,6 @@ class NativeCandidateState { : static_cast(active_object.data())[b] != 0; reset[b] = reset_bool ? static_cast(reset_object.data())[b] : static_cast(reset_object.data())[b] != 0; - const int64_t next_position = reset[b] ? 0 : positions_.data()[b]; - if (active[b] && next_position < 0) - throw std::runtime_error("candidate position must be non-negative"); - if (active[b] && next_position >= max_length_) - throw std::runtime_error("candidate state capacity exceeded"); } const int64_t slots = suffix_k_ * occurrences_r_; py::array_t source({batch_, slots}), match_length({batch_, slots}), @@ -941,24 +1158,38 @@ class NativeCandidateState { std::fill(state_id.mutable_data(), state_id.mutable_data() + batch_ * slots, int64_t{-1}); std::fill(candidate_frequency.mutable_data(), candidate_frequency.mutable_data() + batch_ * slots, int64_t{0}); std::fill(count.mutable_data(), count.mutable_data() + batch_, int32_t{0}); + std::vector next_positions(batch_); + std::unique_lock call_lock; { py::gil_scoped_release release; + call_lock = std::unique_lock(call_mutex_); + for (int64_t b = 0; b < batch_; ++b) { + const int64_t next_position = reset[b] ? 0 : positions_.data()[b]; + if (active[b] && next_position < 0) + throw std::runtime_error("candidate position must be non-negative"); + if (active[b] && next_position >= max_length_) + throw std::runtime_error("candidate state capacity exceeded"); + } for (int64_t b = 0; b < batch_; ++b) { + const int64_t current_position = + active[b] && reset[b] ? 0 : positions_.data()[b]; + next_positions[b] = current_position; if (!active[b]) continue; - if (reset[b]) { + if (reset[b]) reset_row(b); - positions_.mutable_data()[b] = 0; - } - const int64_t position = positions_.data()[b]; count.mutable_data()[b] = step_row( - b, tokens.data()[b], position, source.mutable_data() + b * slots, + b, tokens.data()[b], current_position, + source.mutable_data() + b * slots, match_length.mutable_data() + b * slots, state_id.mutable_data() + b * slots, candidate_frequency.mutable_data() + b * slots); - positions_.mutable_data()[b] = position + 1; + next_positions[b] = current_position + 1; } } + std::copy(next_positions.begin(), next_positions.end(), + positions_.mutable_data()); + call_lock.unlock(); return py::make_tuple(source, match_length, state_id, candidate_frequency, count); } @@ -974,17 +1205,25 @@ class NativeCandidateState { reset_object.ndim() != 1 || reset_object.shape(0) != batch_) throw py::value_error("reset must be contiguous [batch_size]"); std::vector reset(batch_); + std::vector next_positions(batch_); for (int64_t b = 0; b < batch_; ++b) reset[b] = reset_bool ? static_cast(reset_object.data())[b] : static_cast(reset_object.data())[b] != 0; + std::unique_lock call_lock; { py::gil_scoped_release release; - for (int64_t b = 0; b < batch_; ++b) + call_lock = std::unique_lock(call_mutex_); + for (int64_t b = 0; b < batch_; ++b) { + next_positions[b] = positions_.data()[b]; if (reset[b]) { reset_row(b); - positions_.mutable_data()[b] = 0; + next_positions[b] = 0; } + } } + std::copy(next_positions.begin(), next_positions.end(), + positions_.mutable_data()); + call_lock.unlock(); } py::tuple prefill(py::array tokens_object) { @@ -996,12 +1235,8 @@ class NativeCandidateState { tokens_object.ndim() != 2 || tokens_object.shape(0) != batch_) throw py::value_error( "tokens must be contiguous int64 [batch_size, sequence_length]"); - if (position_ != 0) - throw std::runtime_error("prefill requires an empty candidate state"); auto tokens = py::cast>(tokens_object); const int64_t sequence_length = tokens.shape(1); - if (sequence_length > max_length_) - throw std::runtime_error("candidate state capacity exceeded"); const int64_t slots = suffix_k_ * occurrences_r_; py::array_t source({batch_, sequence_length, slots}), match_length({batch_, sequence_length, slots}), @@ -1014,10 +1249,17 @@ class NativeCandidateState { 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_ * sequence_length, int32_t{0}); + ensure_pool(4); + std::unique_lock call_lock; { py::gil_scoped_release release; - for (int64_t position = 0; position < sequence_length; ++position) - for (int64_t b = 0; b < batch_; ++b) { + call_lock = std::unique_lock(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"); + parallel_for_rows(4, [&](int64_t b) { + for (int64_t position = 0; position < sequence_length; ++position) { 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, @@ -1026,19 +1268,40 @@ class NativeCandidateState { state_id.mutable_data() + output_at, candidate_frequency.mutable_data() + output_at); } + }); } position_ = sequence_length; std::fill(positions_.mutable_data(), positions_.mutable_data() + batch_, position_); state_.attr("position") = py::int_(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 positions() const { return positions_; } + size_t worker_count() const { + return row_pool_ ? row_pool_->worker_count() : 0; + } private: + void ensure_pool(int64_t minimum_parallel_rows) { + if (batch_ >= minimum_parallel_rows && !row_pool_) + row_pool_ = std::make_unique(batch_); + } + + template + void parallel_for_rows(int64_t minimum_parallel_rows, Function &&function) { + if (batch_ < minimum_parallel_rows) { + for (int64_t b = 0; b < batch_; ++b) + function(b); + return; + } + row_pool_->parallel_for_rows(batch_, minimum_parallel_rows, + std::forward(function)); + } + template py::array_t bind(const char *name) { py::object object = state_.attr(name); if (!py::isinstance>(object)) @@ -1386,6 +1649,8 @@ class NativeCandidateState { hash_state_, hash_edge_, suffix_link_, length_, left_, right_, parent_, occurrence_size_, lazy_size_, stack_, last_, size_, edge_count_; std::vector> occupied_slots_; + std::unique_ptr row_pool_; + std::mutex call_mutex_; int64_t batch_, max_length_, suffix_k_, occurrences_r_, position_, state_capacity_, edge_capacity_, hash_capacity_; bool ragged_mode_ = false; @@ -1399,7 +1664,8 @@ PYBIND11_MODULE(rosa_native_step, m) { .def("step_masked", &NativeState::step_masked) .def("prefill", &NativeState::prefill) .def_property_readonly("position", &NativeState::position) - .def_property_readonly("positions", &NativeState::positions); + .def_property_readonly("positions", &NativeState::positions) + .def_property_readonly("worker_count", &NativeState::worker_count); py::class_(m, "NativeCandidateState") .def(py::init(), py::keep_alive<1, 2>()) .def("step", &NativeCandidateState::step) @@ -1408,6 +1674,8 @@ PYBIND11_MODULE(rosa_native_step, m) { .def("reset_masked", &NativeCandidateState::reset_masked) .def("prefill", &NativeCandidateState::prefill) .def_property_readonly("position", &NativeCandidateState::position) - .def_property_readonly("positions", &NativeCandidateState::positions); + .def_property_readonly("positions", &NativeCandidateState::positions) + .def_property_readonly("worker_count", + &NativeCandidateState::worker_count); m.attr("candidate_abi_version") = py::int_(1); } diff --git a/native/tests/candidate_smoke.py b/native/tests/candidate_smoke.py index d2b97df..1017cbd 100644 --- a/native/tests/candidate_smoke.py +++ b/native/tests/candidate_smoke.py @@ -1,7 +1,10 @@ from __future__ import annotations import gc +import os +import threading import weakref +from concurrent.futures import ThreadPoolExecutor from itertools import product import numpy as np @@ -66,6 +69,69 @@ def compare(tokens: torch.Tensor, suffix_k: int, occurrences_r: int) -> None: def main() -> None: assert rosa_native_step.candidate_abi_version == 1 + + # 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)) + assert small_pool.worker_count == 0 + small_pool.prefill(np.zeros((3, 4), dtype=np.int64)) + assert small_pool.worker_count == 0 + + previous_threads = os.environ.get("ROSA_NATIVE_THREADS") + os.environ["ROSA_NATIVE_THREADS"] = "-2" + try: + invalid_limit_pool = rosa_native_step.NativeCandidateState( + init_candidate_state(16, 4) + ) + assert invalid_limit_pool.worker_count == 0 + invalid_limit_pool.prefill(np.zeros((16, 4), dtype=np.int64)) + assert invalid_limit_pool.worker_count == 0 + finally: + if previous_threads is None: + os.environ.pop("ROSA_NATIVE_THREADS", None) + else: + os.environ["ROSA_NATIVE_THREADS"] = previous_threads + + ragged_pool = rosa_native_step.NativeCandidateState( + init_candidate_state(16, 4, ragged=True) + ) + ragged_pool.step_masked( + np.zeros(16, dtype=np.int64), + np.ones(16, dtype=np.bool_), + np.zeros(16, dtype=np.bool_), + ) + assert ragged_pool.worker_count == 0 + + # Concurrent mutators serialize while the GIL is released; publication of + # position remains atomic from Python's perspective. + concurrent_state = init_candidate_state(4, 8) + concurrent_native = rosa_native_step.NativeCandidateState(concurrent_state) + repeated = np.arange(4, dtype=np.int64) + sequential_state = init_candidate_state(4, 8) + sequential_native = rosa_native_step.NativeCandidateState(sequential_state) + expected_steps = [sequential_native.step(repeated) for _ in range(2)] + started = threading.Barrier(3) + + def concurrent_step() -> tuple[np.ndarray, ...]: + started.wait() + return concurrent_native.step(repeated) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(concurrent_step) for _ in range(2)] + started.wait() + actual_steps = [future.result() for future in futures] + assert concurrent_native.position == concurrent_state.position == 2 + assert all( + any( + all( + np.array_equal(actual, expected) + for actual, expected in zip(step, candidate, strict=True) + ) + for candidate in expected_steps + ) + for step in actual_steps + ) + binary = torch.tensor(list(product(range(2), repeat=9)), dtype=torch.long) for suffix_k, occurrences_r in ((1, 1), (2, 3), (4, 2), (5, 4)): compare(binary, suffix_k, occurrences_r) From a7082431364b4d2b79b0102add380c749ca0c35d Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:19:17 +0800 Subject: [PATCH 27/29] Add reusable native candidate buffers --- README.md | 15 ++ native/src/rosa_native_step.cpp | 177 ++++++++++++++---- native/tests/candidate_smoke.py | 55 ++++++ src/rosa/__init__.py | 55 ++++++ src/rosa/_stateful_candidates_numba.py | 237 +++++++++++++++++++++++++ tests/test_stateful_candidates.py | 55 ++++++ tests/test_unified_inference.py | 25 +++ 7 files changed, 584 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 706b85d..7a39a2a 100644 --- a/README.md +++ b/README.md @@ -164,6 +164,21 @@ the `numba` extra and automatically use compatible native companion methods when installed. Legacy `forward_step`, `prefill`, `init_candidate_state`, and `forward_candidates_step` remain supported. +Latency-sensitive uniform rich inference can opt into caller-owned output +storage and avoid the five native NumPy allocations on every token: + +```python +from rosa import init_candidate_buffers + +rich = init_inference_state(8, 32_768, mode="rich") +buffers = init_candidate_buffers(rich) +result = rich.step_into(token_ids, buffers) +``` + +The returned candidate tensors alias `buffers` and are valid until those +buffers are reused. The regular `step` API continues to return independently +owned snapshots suitable for retention. + ## Quick start ```python diff --git a/native/src/rosa_native_step.cpp b/native/src/rosa_native_step.cpp index ada5383..a1b506b 100644 --- a/native/src/rosa_native_step.cpp +++ b/native/src/rosa_native_step.cpp @@ -1060,26 +1060,35 @@ class NativeCandidateState { } py::tuple step(py::array tokens_object) { - if (ragged_mode_) - throw std::runtime_error("uniform step is unavailable on a ragged candidate state"); - if (!py::isinstance>(tokens_object)) - throw py::type_error("tokens must have dtype int64"); - if ((tokens_object.flags() & py::array::c_style) == 0 || - tokens_object.ndim() != 1 || tokens_object.shape(0) != batch_) - throw py::value_error("tokens must be contiguous int64 [batch_size]"); - auto tokens = py::cast>(tokens_object); const int64_t slots = suffix_k_ * occurrences_r_; py::array_t source({batch_, slots}), match_length({batch_, slots}), state_id({batch_, slots}), candidate_frequency({batch_, slots}); py::array_t count(batch_); - std::fill(source.mutable_data(), source.mutable_data() + batch_ * slots, - int64_t{-1}); - std::fill(match_length.mutable_data(), - match_length.mutable_data() + batch_ * slots, int64_t{0}); - std::fill(state_id.mutable_data(), state_id.mutable_data() + batch_ * slots, - int64_t{-1}); - std::fill(candidate_frequency.mutable_data(), - candidate_frequency.mutable_data() + batch_ * slots, int64_t{0}); + step_into(tokens_object, source, match_length, state_id, + candidate_frequency, count); + return py::make_tuple(source, match_length, state_id, candidate_frequency, + count); + } + + void step_into(py::array tokens_object, py::array source_object, + py::array match_length_object, py::array state_id_object, + py::array candidate_frequency_object, + py::array count_object) { + if (ragged_mode_) + throw std::runtime_error("uniform step is unavailable on a ragged candidate state"); + const int64_t slots = suffix_k_ * occurrences_r_; + auto tokens = checked_input(tokens_object, "tokens", {batch_}); + auto source = checked_output(source_object, "source", {batch_, slots}); + auto match_length = checked_output( + match_length_object, "length", {batch_, slots}); + auto state_id = checked_output(state_id_object, "state", {batch_, slots}); + auto candidate_frequency = checked_output( + candidate_frequency_object, "frequency", {batch_, slots}); + auto count = checked_output(count_object, "count", {batch_}); + validate_disjoint({tokens_object, source_object, match_length_object, + state_id_object, candidate_frequency_object, + count_object}); + validate_runtime_positions(); ensure_pool(64); std::unique_lock call_lock; { @@ -1087,6 +1096,14 @@ class NativeCandidateState { call_lock = std::unique_lock(call_mutex_); if (position_ >= max_length_) throw std::runtime_error("candidate state capacity exceeded"); + std::fill(source.mutable_data(), source.mutable_data() + batch_ * slots, + int64_t{-1}); + std::fill(match_length.mutable_data(), + match_length.mutable_data() + batch_ * slots, int64_t{0}); + std::fill(state_id.mutable_data(), state_id.mutable_data() + batch_ * slots, + int64_t{-1}); + std::fill(candidate_frequency.mutable_data(), + candidate_frequency.mutable_data() + batch_ * slots, int64_t{0}); parallel_for_rows(64, [&](int64_t b) { count.mutable_data()[b] = step_row(b, tokens.data()[b], position_, @@ -1101,8 +1118,6 @@ class NativeCandidateState { position_); state_.attr("position") = py::int_(position_); call_lock.unlock(); - return py::make_tuple(source, match_length, state_id, candidate_frequency, - count); } void reset() { @@ -1227,28 +1242,51 @@ class NativeCandidateState { } py::tuple prefill(py::array tokens_object) { - if (ragged_mode_) - throw std::runtime_error("prefill is unavailable on a ragged candidate state"); - if (!py::isinstance>(tokens_object)) - throw py::type_error("tokens must have dtype int64"); - if ((tokens_object.flags() & py::array::c_style) == 0 || - tokens_object.ndim() != 2 || tokens_object.shape(0) != batch_) - throw py::value_error( - "tokens must be contiguous int64 [batch_size, sequence_length]"); - auto tokens = py::cast>(tokens_object); - const int64_t sequence_length = tokens.shape(1); + if (!py::isinstance>(tokens_object) || + tokens_object.ndim() != 2) + throw py::type_error("tokens must be a two-dimensional int64 array"); + const int64_t sequence_length = tokens_object.shape(1); const int64_t slots = suffix_k_ * occurrences_r_; py::array_t source({batch_, sequence_length, slots}), match_length({batch_, sequence_length, slots}), state_id({batch_, sequence_length, slots}), candidate_frequency({batch_, sequence_length, slots}); py::array_t count({batch_, sequence_length}); + prefill_into(tokens_object, source, match_length, state_id, + candidate_frequency, count); + return py::make_tuple(source, match_length, state_id, candidate_frequency, + count); + } + + void prefill_into(py::array tokens_object, py::array source_object, + py::array match_length_object, py::array state_id_object, + py::array candidate_frequency_object, + py::array count_object) { + if (ragged_mode_) + throw std::runtime_error("prefill is unavailable on a ragged candidate state"); + if (!py::isinstance>(tokens_object) || + tokens_object.ndim() != 2 || tokens_object.shape(0) != batch_) + throw py::value_error("tokens must be contiguous int64 [batch_size, sequence_length]"); + const int64_t sequence_length = tokens_object.shape(1); + const int64_t slots = suffix_k_ * occurrences_r_; + auto tokens = checked_input(tokens_object, "tokens", + {batch_, sequence_length}); + auto source = checked_output( + source_object, "source", {batch_, sequence_length, slots}); + auto match_length = checked_output( + match_length_object, "length", {batch_, sequence_length, slots}); + auto state_id = checked_output( + state_id_object, "state", {batch_, sequence_length, slots}); + auto candidate_frequency = checked_output( + candidate_frequency_object, "frequency", + {batch_, sequence_length, slots}); + auto count = checked_output(count_object, "count", + {batch_, sequence_length}); + validate_disjoint({tokens_object, source_object, match_length_object, + state_id_object, candidate_frequency_object, + count_object}); + validate_runtime_positions(); const int64_t output_size = batch_ * sequence_length * slots; - 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_ * sequence_length, int32_t{0}); ensure_pool(4); std::unique_lock call_lock; { @@ -1258,6 +1296,11 @@ class NativeCandidateState { 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_ * sequence_length, int32_t{0}); parallel_for_rows(4, [&](int64_t b) { for (int64_t position = 0; position < sequence_length; ++position) { const int64_t output_at = (b * sequence_length + position) * slots; @@ -1275,8 +1318,6 @@ class NativeCandidateState { position_); state_.attr("position") = py::int_(position_); call_lock.unlock(); - return py::make_tuple(source, match_length, state_id, candidate_frequency, - count); } int64_t position() const { return position_; } @@ -1286,6 +1327,69 @@ class NativeCandidateState { } private: + template + py::array_t + checked_input(py::array object, const char *name, + std::initializer_list shape) const { + if (!py::isinstance>(object)) + throw py::type_error(std::string(name) + " has an unexpected dtype"); + if ((object.flags() & py::array::c_style) == 0 || + object.ndim() != static_cast(shape.size())) + throw py::value_error(std::string(name) + " must be C-contiguous with the expected shape"); + int64_t dimension = 0; + for (const int64_t extent : shape) + if (object.shape(dimension++) != extent) + throw py::value_error(std::string(name) + " must be C-contiguous with the expected shape"); + return py::cast>(object); + } + + template + py::array_t + checked_output(py::array object, const char *name, + std::initializer_list shape) const { + auto output = checked_input(object, name, shape); + if (!output.writeable()) + throw py::value_error(std::string(name) + " must be writable"); + return output; + } + + static bool overlaps(const py::array &first, const py::array &second) { + if (first.nbytes() == 0 || second.nbytes() == 0) + return false; + const auto first_begin = reinterpret_cast(first.data()); + const auto second_begin = reinterpret_cast(second.data()); + const auto first_end = first_begin + static_cast(first.nbytes()); + const auto second_end = second_begin + static_cast(second.nbytes()); + return first_begin < second_end && second_begin < first_end; + } + + void validate_disjoint(std::initializer_list call_arrays) const { + const std::vector arrays(call_arrays); + for (size_t first = 0; first < arrays.size(); ++first) + for (size_t second = first + 1; second < arrays.size(); ++second) + if (overlaps(arrays[first], arrays[second])) + throw py::value_error("tokens and output buffers must not overlap"); + const py::array state_arrays[] = { + history_, head_, edge_token_, edge_target_, edge_next_, hash_state_, + hash_token_, hash_edge_, suffix_link_, length_, left_, right_, parent_, + occurrences_, occurrence_size_, frequency_, lazy_prefix_, lazy_size_, + lazy_delta_, stack_, last_, size_, edge_count_, positions_}; + for (const py::array &array : arrays) + for (const py::array &state_array : state_arrays) + if (overlaps(array, state_array)) + throw py::value_error("tokens and output buffers must not overlap candidate state storage"); + } + + void validate_runtime_positions() const { + for (int64_t b = 0; b < batch_; ++b) { + const int64_t position = positions_.data()[b]; + if (position < 0 || position > max_length_) + throw py::value_error("candidate positions are outside capacity"); + if (!ragged_mode_ && position != position_) + throw py::value_error("uniform candidate positions are inconsistent"); + } + } + void ensure_pool(int64_t minimum_parallel_rows) { if (batch_ >= minimum_parallel_rows && !row_pool_) row_pool_ = std::make_unique(batch_); @@ -1656,6 +1760,7 @@ class NativeCandidateState { bool ragged_mode_ = false; }; +// Allocating candidate entry points delegate to their caller-owned variants. PYBIND11_MODULE(rosa_native_step, m) { m.doc() = "Exact CPU SAM+LCT step prototype (no libtorch calls in core)"; py::class_(m, "NativeState") @@ -1669,10 +1774,12 @@ PYBIND11_MODULE(rosa_native_step, m) { py::class_(m, "NativeCandidateState") .def(py::init(), py::keep_alive<1, 2>()) .def("step", &NativeCandidateState::step) + .def("step_into", &NativeCandidateState::step_into) .def("reset", &NativeCandidateState::reset) .def("step_masked", &NativeCandidateState::step_masked) .def("reset_masked", &NativeCandidateState::reset_masked) .def("prefill", &NativeCandidateState::prefill) + .def("prefill_into", &NativeCandidateState::prefill_into) .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 1017cbd..c65a21a 100644 --- a/native/tests/candidate_smoke.py +++ b/native/tests/candidate_smoke.py @@ -70,6 +70,61 @@ def compare(tokens: torch.Tensor, suffix_k: int, occurrences_r: int) -> None: def main() -> None: assert rosa_native_step.candidate_abi_version == 1 + into_state = init_candidate_state(2, 3, suffix_k=2, occurrences_r=2) + into_native = rosa_native_step.NativeCandidateState(into_state) + into_arrays = ( + np.empty((2, 4), dtype=np.int64), + np.empty((2, 4), dtype=np.int64), + np.empty((2, 4), dtype=np.int64), + np.empty((2, 4), dtype=np.int64), + np.empty(2, dtype=np.int32), + ) + first_tokens = np.array([0, 3], dtype=np.int64) + into_native.step_into(first_tokens, *into_arrays) + allocating_state = init_candidate_state(2, 3, suffix_k=2, occurrences_r=2) + allocating_native = rosa_native_step.NativeCandidateState(allocating_state) + allocated = allocating_native.step(first_tokens) + assert all( + np.array_equal(actual, expected) + for actual, expected in zip(into_arrays, allocated, strict=True) + ) + try: + overlap_state = rosa_native_step.NativeCandidateState( + init_candidate_state(2, 1, suffix_k=2, occurrences_r=2) + ) + overlap_state.step_into( + first_tokens, + into_arrays[0], + into_arrays[0], + into_arrays[2], + into_arrays[3], + into_arrays[4], + ) + except ValueError as error: + assert "overlap" in str(error) + else: + raise AssertionError("overlapping outputs were accepted") + + prefix = np.array([[0, 1, 0], [3, 3, 4]], dtype=np.int64) + prefill_state = rosa_native_step.NativeCandidateState( + init_candidate_state(2, 3, suffix_k=2, occurrences_r=2) + ) + prefill_arrays = ( + np.empty((2, 3, 4), dtype=np.int64), + np.empty((2, 3, 4), dtype=np.int64), + np.empty((2, 3, 4), dtype=np.int64), + np.empty((2, 3, 4), dtype=np.int64), + np.empty((2, 3), dtype=np.int32), + ) + prefill_state.prefill_into(prefix, *prefill_arrays) + prefill_allocating = rosa_native_step.NativeCandidateState( + init_candidate_state(2, 3, suffix_k=2, occurrences_r=2) + ).prefill(prefix) + assert all( + np.array_equal(actual, expected) + for actual, expected in zip(prefill_arrays, prefill_allocating, 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 c8c693e..f746db2 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -33,8 +33,10 @@ "build_hard_candidates", "build_virtual_pool_indices", "forward_candidates_step", + "forward_candidates_step_into", "forward_step", "init_candidate_state", + "init_candidate_buffers", "init_inference_state", "prefill", "reference_rosa", @@ -1244,6 +1246,14 @@ def step( return _inference_step(self, tokens, active=active, reset=reset) + def step_into(self, tokens: Tensor, buffers: object) -> InferenceOutput: + """Consume one rich uniform step into opt-in reusable buffers. + + Candidate tensors in the result alias ``buffers`` until its next use. + """ + + return _inference_step_into(self, tokens, buffers) + def prefill(self, tokens: Tensor) -> InferenceOutput: """Consume an initial dense context and return every step result.""" @@ -1351,6 +1361,18 @@ def init_candidate_state( ) +def init_candidate_buffers(state: object, *, sequence_length: int | None = None) -> Any: + """Allocate reusable packed CPU outputs for rich ``*_into`` calls.""" + + if isinstance(state, ROSAInferenceState): + if state.mode != "rich": + raise ValueError("candidate buffers require a rich state") + state = state._impl + from ._stateful_candidates_numba import init_candidate_buffers as initialize + + return initialize(cast(Any, state), sequence_length=sequence_length) + + def forward_candidates_step(state: object, tokens: Tensor) -> Any: """Consume one token and return exact bounded hard candidates.""" @@ -1373,6 +1395,19 @@ def forward_candidates_step(state: object, tokens: Tensor) -> Any: return step(cast(Any, state), tokens) +def forward_candidates_step_into(state: object, tokens: Tensor, buffers: object) -> Any: + """Consume one rich step into explicitly reusable caller-owned buffers.""" + + if isinstance(state, ROSAInferenceState): + candidates = state.step_into(tokens, buffers).candidates + if candidates is None: # pragma: no cover - guarded by step_into + raise RuntimeError("rich inference did not return candidates") + return candidates + from ._stateful_candidates_numba import forward_candidates_step_into as step + + return step(cast(Any, state), tokens, cast(Any, buffers)) + + def init_inference_state( batch_size: int, max_length: int = 8192, @@ -1539,6 +1574,26 @@ def _inference_step( return InferenceOutput(output, candidates) +def _inference_step_into( + state: ROSAInferenceState, token: Tensor, buffers: object +) -> InferenceOutput: + if state.mode != "rich" or state.ragged: + raise ValueError("step_into requires a uniform rich inference state") + token, squeeze = _normalize_step_tokens(state, token) + from ._stateful_candidates_numba import forward_candidates_step_into as step + + candidates = step(cast(Any, state._impl), token, cast(Any, buffers)) + output = candidates.rosa_predicted_tokens + if squeeze: + output = output[0] + from ._stateful_candidates_numba import CandidateStep + + candidates = CandidateStep( + *(getattr(candidates, name)[0] for name in candidates.__dataclass_fields__) + ) + return InferenceOutput(output, candidates) + + def forward_step(state: ROSAInferenceState, token: Tensor) -> Tensor: """Consume one token per batch row and return exact ROSA predictions. diff --git a/src/rosa/_stateful_candidates_numba.py b/src/rosa/_stateful_candidates_numba.py index 2b27fe0..3bd5278 100644 --- a/src/rosa/_stateful_candidates_numba.py +++ b/src/rosa/_stateful_candidates_numba.py @@ -1005,6 +1005,47 @@ class CandidateStep: rosa_predicted_tokens: Tensor +@dataclass(slots=True) +class CandidateBuffers: + """Reusable packed CPU storage for opt-in immediate candidate consumption. + + Results returned by ``*_into`` alias these arrays and remain valid only + until the buffers are reused. The historical allocating APIs never alias + a ``CandidateBuffers`` instance. + """ + + source_index: np.ndarray + match_length: np.ndarray + state_id: np.ndarray + frequency: np.ndarray + count: np.ndarray + + +def init_candidate_buffers( + state: CandidateState, *, sequence_length: int | None = None +) -> CandidateBuffers: + """Allocate packed caller-owned buffers for ``step_into`` or ``prefill_into``.""" + + if not isinstance(state, CandidateState): + raise TypeError("state must be a CandidateState") + if sequence_length is not None and sequence_length < 0: + raise ValueError("sequence_length must be >= 0") + slots = state.suffix_k * state.occurrences_r + prefix = ( + (state.batch_size,) + if sequence_length is None + else (state.batch_size, sequence_length) + ) + candidate_shape = (*prefix, slots) + return CandidateBuffers( + np.empty(candidate_shape, dtype=np.int64), + np.empty(candidate_shape, dtype=np.int64), + np.empty(candidate_shape, dtype=np.int64), + np.empty(candidate_shape, dtype=np.int64), + np.empty(prefix, dtype=np.int32), + ) + + def init_candidate_state( batch_size: int, max_length: int, @@ -1173,6 +1214,202 @@ def _candidate_step_from_arrays( ) +def _validate_candidate_buffers( + state: CandidateState, + buffers: CandidateBuffers, + tokens: np.ndarray, + *, + sequence_length: int | None = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + if not isinstance(buffers, CandidateBuffers): + raise TypeError("buffers must be CandidateBuffers") + slots = state.suffix_k * state.occurrences_r + prefix = ( + (state.batch_size,) + if sequence_length is None + else (state.batch_size, sequence_length) + ) + expected = ( + ("source_index", buffers.source_index, np.dtype(np.int64), (*prefix, slots)), + ("match_length", buffers.match_length, np.dtype(np.int64), (*prefix, slots)), + ("state_id", buffers.state_id, np.dtype(np.int64), (*prefix, slots)), + ("frequency", buffers.frequency, np.dtype(np.int64), (*prefix, slots)), + ("count", buffers.count, np.dtype(np.int32), prefix), + ) + arrays: list[np.ndarray] = [] + for name, array, dtype, shape in expected: + if not isinstance(array, np.ndarray): + raise TypeError(f"buffers.{name} must be a numpy.ndarray") + if array.dtype != dtype: + raise TypeError(f"buffers.{name} has an unexpected dtype") + if array.shape != shape or not array.flags.c_contiguous: + raise ValueError(f"buffers.{name} must be C-contiguous with shape {shape}") + if not array.flags.writeable: + raise ValueError(f"buffers.{name} must be writable") + arrays.append(array) + all_call_arrays = [tokens, *arrays] + for first_index, first in enumerate(all_call_arrays): + for second in all_call_arrays[first_index + 1 :]: + if np.shares_memory(first, second): + raise ValueError("tokens and candidate buffers must not overlap") + state_arrays = [ + value for value in state.__dict__.values() if isinstance(value, np.ndarray) + ] + for array in all_call_arrays: + if any(np.shares_memory(array, state_array) for state_array in state_arrays): + raise ValueError( + "candidate buffers must not overlap candidate state storage" + ) + return arrays[0], arrays[1], arrays[2], arrays[3], arrays[4] + + +def _native_candidate_into( + state: CandidateState, + method: str, + tokens: np.ndarray, + arrays: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], +) -> bool: # pragma: no cover - optional native companion + if state.native_state is False: + return False + if state.native_state is None: + try: + import rosa_native_step # type: ignore[reportMissingImports] + except ModuleNotFoundError: + state.native_state = False + return False + native_type = getattr(rosa_native_step, "NativeCandidateState", None) + if native_type is None: + state.native_state = False + return False + state.native_state = native_type(state) + native_method = getattr(state.native_state, method, None) + if native_method is not None: + native_method(tokens, *arrays) + return True + + # An ABI-1 wheel predating caller-owned buffers remains a valid exact + # backend: allocate through its historical method, then copy into storage. + allocating_method = getattr(state.native_state, method.removesuffix("_into"), None) + if allocating_method is None: + return False + allocated = allocating_method(tokens) + for destination, source in zip(arrays, allocated, strict=True): + np.copyto(destination, source) + return True + + +def forward_candidates_step_into( + state: CandidateState, tokens: Tensor, buffers: CandidateBuffers +) -> CandidateStep: + """Consume one token into reusable CPU buffers. + + The returned tensors alias ``buffers`` and are intended for immediate + consumption. Use :func:`forward_candidates_step` for retained snapshots. + """ + + tokens, _ = _validate_candidate_tokens(state, tokens) + if state.ragged_mode: + raise RuntimeError("uniform step is unavailable on a ragged candidate state") + if state.position >= state.max_length: + raise RuntimeError("candidate state capacity exceeded") + device = tokens.device + cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() + token_array = cpu_tokens.numpy() + arrays = _validate_candidate_buffers(state, buffers, token_array) + if not _native_candidate_into(state, "step_into", token_array, arrays): + allocated = _step_batch_kernel( + token_array, + state.position, + 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, + ) + for destination, source in zip(arrays, allocated, strict=True): + np.copyto(destination, source) + state.position += 1 + state.positions.fill(state.position) + else: + state.positions.fill(state.position) + return _candidate_step_from_arrays(state, arrays, device) + + +def prefill_candidates_into( + state: CandidateState, tokens: Tensor, buffers: CandidateBuffers +) -> CandidateStep: + """Consume a dense context into reusable packed CPU buffers.""" + + tokens, _ = _validate_candidate_tokens(state, tokens, sequence=True) + 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() + token_array = cpu_tokens.numpy() + arrays = _validate_candidate_buffers( + state, buffers, token_array, sequence_length=sequence_length + ) + if not _native_candidate_into(state, "prefill_into", token_array, arrays): + allocated = _prefill_candidate_kernel( + token_array, + 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, + ) + for destination, source in zip(arrays, allocated, strict=True): + np.copyto(destination, source) + state.position = sequence_length + state.positions.fill(state.position) + return _candidate_step_from_arrays(state, arrays, device) + + def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateStep: """Consume one token per row and return exact top-R candidates for K suffixes.""" diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py index d85015a..816481e 100644 --- a/tests/test_stateful_candidates.py +++ b/tests/test_stateful_candidates.py @@ -5,11 +5,14 @@ from itertools import product from unittest.mock import patch +import numpy as np import torch from rosa import ( build_hard_candidates, forward_candidates_step, + forward_candidates_step_into, + init_candidate_buffers, init_candidate_state, ) from rosa._stateful_candidates_numba import ( @@ -17,6 +20,7 @@ CandidateStep, forward_candidates_step_masked, prefill_candidates, + prefill_candidates_into, reset_candidates_masked, ) from rosa._stateful_candidates_numba import ( @@ -139,6 +143,57 @@ def test_prefill_emits_every_position_and_continues(self) -> None: ) ) + 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) + state.native_state = False + buffers = init_candidate_buffers(state) + retained = forward_candidates_step_into(state, tokens[:, 0], buffers) + retained_source = retained.source_index.clone() + current = forward_candidates_step_into(state, tokens[:, 1], buffers) + expected = build_hard_candidates(tokens[:, :2], suffix_k=3, occurrences_r=2) + for field in _FIELDS: + self.assertTrue( + torch.equal(getattr(current, field), getattr(expected, field)[:, 1]), + field, + ) + self.assertFalse(torch.equal(retained.source_index, retained_source)) + + prefill_state = init_candidate_state(2, 4, suffix_k=3, occurrences_r=2) + prefill_state.native_state = False + prefill_buffers = init_candidate_buffers(prefill_state, sequence_length=4) + actual = prefill_candidates_into(prefill_state, tokens, prefill_buffers) + expected = build_hard_candidates(tokens, suffix_k=3, occurrences_r=2) + for field in _FIELDS: + self.assertTrue( + torch.equal(getattr(actual, field), getattr(expected, field)) + ) + + def test_caller_owned_buffer_validation_and_allocating_snapshots(self) -> None: + state = init_candidate_state(1, 2, suffix_k=2, occurrences_r=2) + state.native_state = False + first = forward_candidates_step(state, torch.tensor([0])) + snapshot = first.source_index.clone() + forward_candidates_step(state, torch.tensor([0])) + self.assertTrue(torch.equal(first.source_index, snapshot)) + + wrong_state = init_candidate_state(1, 2, suffix_k=2, occurrences_r=2) + wrong_state.native_state = False + buffers = init_candidate_buffers(wrong_state) + buffers.source_index = np.empty((1, 3), dtype=np.int64) + with self.assertRaisesRegex(ValueError, "shape"): + forward_candidates_step_into(wrong_state, torch.tensor([0]), buffers) + + buffers = init_candidate_buffers(wrong_state) + buffers.count.flags.writeable = False + with self.assertRaisesRegex(ValueError, "writable"): + forward_candidates_step_into(wrong_state, torch.tensor([0]), buffers) + + with self.assertRaisesRegex(TypeError, "CandidateState"): + init_candidate_buffers(object()) + with self.assertRaisesRegex(ValueError, "sequence_length"): + init_candidate_buffers(wrong_state, sequence_length=-1) + def test_ragged_inactive_reset_capacity_and_recycle(self) -> None: state = init_candidate_state_internal( 3, 5, suffix_k=4, occurrences_r=3, ragged=True diff --git a/tests/test_unified_inference.py b/tests/test_unified_inference.py index 2e112a6..5586034 100644 --- a/tests/test_unified_inference.py +++ b/tests/test_unified_inference.py @@ -11,6 +11,7 @@ build_hard_candidates, forward_candidates_step, forward_step, + init_candidate_buffers, init_candidate_state, init_inference_state, prefill, @@ -81,6 +82,30 @@ def test_legacy_top1_and_candidate_wrappers_remain_compatible(self) -> None: old_output = forward_candidates_step(old_state, torch.tensor([7])) self.assertEqual(tuple(old_output.source_index.shape), (1, 4)) + def test_rich_step_into_uses_explicit_ephemeral_storage(self) -> None: + rich = init_inference_state(2, 2, mode="rich", suffix_k=2, occurrences_r=2) + buffers = init_candidate_buffers(rich) + first = rich.step_into(torch.tensor([0, 3]), buffers) + self.assertIsNotNone(first.candidates) + second = rich.step_into(torch.tensor([0, 3]), buffers) + expected = build_hard_candidates( + torch.tensor([[0, 0], [3, 3]]), suffix_k=2, occurrences_r=2 + ) + assert second.candidates is not None + for field in HardCandidates.__dataclass_fields__: + self.assertTrue( + torch.equal( + getattr(second.candidates, field), getattr(expected, field)[:, 1] + ), + field, + ) + + top1 = init_inference_state(1, 1) + with self.assertRaisesRegex(ValueError, "rich"): + init_candidate_buffers(top1) + with self.assertRaisesRegex(ValueError, "uniform rich"): + top1.step_into(torch.tensor([0]), buffers) + def test_positions_are_copies_and_uniform_position_is_preserved(self) -> None: uniform = init_inference_state(2, 2) forward_step(uniform, torch.tensor([1, 2])) From 3b43225fe73d484209e5f928fc5f814c23c30ee2 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:54:08 +0800 Subject: [PATCH 28/29] Reuse candidate tensor views --- src/rosa/_stateful_candidates_numba.py | 187 +++++++++++++++++++++++-- tests/test_stateful_candidates.py | 174 +++++++++++++++++++++++ tests/test_unified_inference.py | 11 ++ 3 files changed, 357 insertions(+), 15 deletions(-) diff --git a/src/rosa/_stateful_candidates_numba.py b/src/rosa/_stateful_candidates_numba.py index 3bd5278..8cf04b4 100644 --- a/src/rosa/_stateful_candidates_numba.py +++ b/src/rosa/_stateful_candidates_numba.py @@ -9,7 +9,8 @@ from __future__ import annotations -from dataclasses import dataclass +from collections.abc import Callable +from dataclasses import dataclass, field from typing import Any import numpy as np @@ -1019,6 +1020,25 @@ class CandidateBuffers: state_id: np.ndarray frequency: np.ndarray count: np.ndarray + _validation_signature: tuple[int, ...] | None = field( + default=None, init=False, repr=False + ) + _output_ranges: tuple[tuple[int, int], ...] = field( + default=(), init=False, repr=False + ) + _state_ranges: tuple[tuple[int, int], ...] = field( + default=(), init=False, repr=False + ) + _tensor_signature: tuple[int, ...] | None = field( + default=None, init=False, repr=False + ) + _numpy_outputs: tuple[np.ndarray, ...] = field(default=(), init=False, repr=False) + _cpu_outputs: tuple[Tensor, ...] = field(default=(), init=False, repr=False) + _device_outputs: dict[tuple[str, int | None], tuple[Tensor, ...]] = field( + default_factory=dict, init=False, repr=False + ) + _slot_index: np.ndarray | None = field(default=None, init=False, repr=False) + _valid: np.ndarray | None = field(default=None, init=False, repr=False) def init_candidate_buffers( @@ -1214,6 +1234,76 @@ def _candidate_step_from_arrays( ) +def _candidate_step_from_buffers( + state: CandidateState, + buffers: CandidateBuffers, + arrays: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], + device: torch.device, +) -> CandidateStep: + """Build an ephemeral result while reusing every NumPy/Torch output.""" + + source, match_length, state_id, frequency, count = arrays + signature = tuple(id(array) for array in arrays) + if buffers._tensor_signature != signature: + slots = state.suffix_k * state.occurrences_r + slot_shape = (1,) * count.ndim + (slots,) + buffers._slot_index = np.arange(slots, dtype=np.int32).reshape(slot_shape) + mask = np.empty(source.shape, dtype=np.bool_) + rosa_slot = np.empty(count.shape, dtype=np.int64) + rosa_predicted = np.empty(count.shape, dtype=np.int64) + buffers._valid = np.empty(count.shape, dtype=np.bool_) + buffers._numpy_outputs = ( + source, + match_length, + state_id, + frequency, + mask, + rosa_slot, + source[..., 0], + match_length[..., 0], + rosa_predicted, + ) + buffers._cpu_outputs = tuple( + torch.from_numpy(array) for array in buffers._numpy_outputs + ) + buffers._device_outputs.clear() + buffers._tensor_signature = signature + + mask = buffers._numpy_outputs[4] + rosa_slot = buffers._numpy_outputs[5] + rosa_source = buffers._numpy_outputs[6] + rosa_predicted = buffers._numpy_outputs[8] + assert buffers._slot_index is not None + assert buffers._valid is not None + np.less(buffers._slot_index, count[..., None], out=mask) + np.greater(count, 0, out=buffers._valid) + rosa_slot.fill(-1) + np.copyto(rosa_slot, 0, where=buffers._valid) + rosa_predicted.fill(-1) + for index in np.ndindex(count.shape): + if buffers._valid[index]: + batch_index = index[0] + source_position = int(rosa_source[index]) + rosa_predicted[index] = state.history[batch_index, source_position + 1] + + if device.type == "cpu": + outputs = buffers._cpu_outputs + else: # pragma: no cover - exercised by CUDA validation + device_key = (device.type, device.index) + outputs = buffers._device_outputs.get(device_key, ()) + if not outputs: + outputs = tuple( + torch.empty_like(output, device=device) + for output in buffers._cpu_outputs + ) + buffers._device_outputs[device_key] = outputs + for destination, source_tensor in zip( + outputs, buffers._cpu_outputs, strict=True + ): + destination.copy_(source_tensor, non_blocking=False) + return CandidateStep(*outputs) + + def _validate_candidate_buffers( state: CandidateState, buffers: CandidateBuffers, @@ -1247,20 +1337,63 @@ def _validate_candidate_buffers( if not array.flags.writeable: raise ValueError(f"buffers.{name} must be writable") arrays.append(array) - all_call_arrays = [tokens, *arrays] - for first_index, first in enumerate(all_call_arrays): - for second in all_call_arrays[first_index + 1 :]: - if np.shares_memory(first, second): - raise ValueError("tokens and candidate buffers must not overlap") + return arrays[0], arrays[1], arrays[2], arrays[3], arrays[4] + + +def _validate_candidate_buffer_overlap( + state: CandidateState, + buffers: CandidateBuffers, + tokens: np.ndarray, + arrays: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], +) -> None: + """Validate fallback buffers; current native companions do this in C++.""" + state_arrays = [ value for value in state.__dict__.values() if isinstance(value, np.ndarray) ] - for array in all_call_arrays: - if any(np.shares_memory(array, state_array) for state_array in state_arrays): - raise ValueError( - "candidate buffers must not overlap candidate state storage" + + def memory_range(array: np.ndarray) -> tuple[int, int]: + begin = int(array.__array_interface__["data"][0]) + return begin, begin + array.nbytes + + signature_arrays = [*arrays, *state_arrays] + signature = tuple( + item for array in signature_arrays for item in (id(array), *memory_range(array)) + ) + if buffers._validation_signature != signature: + output_ranges = tuple(memory_range(array) for array in arrays) + state_ranges = tuple(memory_range(array) for array in state_arrays) + + def overlaps(first: tuple[int, int], second: tuple[int, int]) -> bool: + return ( + first[0] < second[1] + and second[0] < first[1] + and first[0] != first[1] + and second[0] != second[1] ) - return arrays[0], arrays[1], arrays[2], arrays[3], arrays[4] + + for first_index, first in enumerate(output_ranges): + if any( + overlaps(first, second) for second in output_ranges[first_index + 1 :] + ): + raise ValueError("candidate output buffers must not overlap") + if any(overlaps(first, state_range) for state_range in state_ranges): + raise ValueError( + "candidate buffers must not overlap candidate state storage" + ) + buffers._validation_signature = signature + buffers._output_ranges = output_ranges + buffers._state_ranges = state_ranges + + token_range = memory_range(tokens) + if any( + token_range[0] < end + and begin < token_range[1] + and token_range[0] != token_range[1] + and begin != end + for begin, end in (*buffers._output_ranges, *buffers._state_ranges) + ): + raise ValueError("tokens and candidate storage must not overlap") def _native_candidate_into( @@ -1268,6 +1401,7 @@ def _native_candidate_into( method: str, tokens: np.ndarray, arrays: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], + validate_fallback: Callable[[], None], ) -> bool: # pragma: no cover - optional native companion if state.native_state is False: return False @@ -1292,6 +1426,7 @@ def _native_candidate_into( allocating_method = getattr(state.native_state, method.removesuffix("_into"), None) if allocating_method is None: return False + validate_fallback() allocated = allocating_method(tokens) for destination, source in zip(arrays, allocated, strict=True): np.copyto(destination, source) @@ -1316,7 +1451,18 @@ def forward_candidates_step_into( cpu_tokens = tokens.detach().to(device="cpu", dtype=torch.long).contiguous() token_array = cpu_tokens.numpy() arrays = _validate_candidate_buffers(state, buffers, token_array) - if not _native_candidate_into(state, "step_into", token_array, arrays): + overlap_validated = False + + def validate_overlap() -> None: + nonlocal overlap_validated + if not overlap_validated: + _validate_candidate_buffer_overlap(state, buffers, token_array, arrays) + overlap_validated = True + + if not _native_candidate_into( + state, "step_into", token_array, arrays, validate_overlap + ): + validate_overlap() allocated = _step_batch_kernel( token_array, state.position, @@ -1352,7 +1498,7 @@ def forward_candidates_step_into( state.positions.fill(state.position) else: state.positions.fill(state.position) - return _candidate_step_from_arrays(state, arrays, device) + return _candidate_step_from_buffers(state, buffers, arrays, device) def prefill_candidates_into( @@ -1374,7 +1520,18 @@ def prefill_candidates_into( arrays = _validate_candidate_buffers( state, buffers, token_array, sequence_length=sequence_length ) - if not _native_candidate_into(state, "prefill_into", token_array, arrays): + overlap_validated = False + + def validate_overlap() -> None: + nonlocal overlap_validated + if not overlap_validated: + _validate_candidate_buffer_overlap(state, buffers, token_array, arrays) + overlap_validated = True + + if not _native_candidate_into( + state, "prefill_into", token_array, arrays, validate_overlap + ): + validate_overlap() allocated = _prefill_candidate_kernel( token_array, state.suffix_k, @@ -1407,7 +1564,7 @@ def prefill_candidates_into( np.copyto(destination, source) state.position = sequence_length state.positions.fill(state.position) - return _candidate_step_from_arrays(state, arrays, device) + return _candidate_step_from_buffers(state, buffers, arrays, device) def forward_candidates_step(state: CandidateState, tokens: Tensor) -> CandidateStep: diff --git a/tests/test_stateful_candidates.py b/tests/test_stateful_candidates.py index 816481e..e609da7 100644 --- a/tests/test_stateful_candidates.py +++ b/tests/test_stateful_candidates.py @@ -151,6 +151,8 @@ def test_caller_owned_step_and_prefill_buffers_are_exact(self) -> None: retained = forward_candidates_step_into(state, tokens[:, 0], buffers) retained_source = retained.source_index.clone() current = forward_candidates_step_into(state, tokens[:, 1], buffers) + self.assertIs(retained.source_index, current.source_index) + self.assertIs(retained.mask, current.mask) expected = build_hard_candidates(tokens[:, :2], suffix_k=3, occurrences_r=2) for field in _FIELDS: self.assertTrue( @@ -169,6 +171,25 @@ def test_caller_owned_step_and_prefill_buffers_are_exact(self) -> None: torch.equal(getattr(actual, field), getattr(expected, field)) ) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is unavailable") + def test_caller_owned_buffers_reuse_cuda_storage(self) -> None: + tokens = torch.tensor([[0, 0], [3, 3]], device="cuda") + state = init_candidate_state(2, 2, suffix_k=2, occurrences_r=2) + buffers = init_candidate_buffers(state) + first = forward_candidates_step_into(state, tokens[:, 0], buffers) + source_pointer = first.source_index.data_ptr() + second = forward_candidates_step_into(state, tokens[:, 1], buffers) + self.assertEqual(second.source_index.device.type, "cuda") + self.assertEqual(second.source_index.data_ptr(), source_pointer) + expected = build_hard_candidates(tokens.cpu(), suffix_k=2, occurrences_r=2) + for field in _FIELDS: + self.assertTrue( + torch.equal( + getattr(second, field).cpu(), getattr(expected, field)[:, 1] + ), + field, + ) + def test_caller_owned_buffer_validation_and_allocating_snapshots(self) -> None: state = init_candidate_state(1, 2, suffix_k=2, occurrences_r=2) state.native_state = False @@ -189,11 +210,164 @@ def test_caller_owned_buffer_validation_and_allocating_snapshots(self) -> None: with self.assertRaisesRegex(ValueError, "writable"): forward_candidates_step_into(wrong_state, torch.tensor([0]), buffers) + for replacement, error in ( + (object(), TypeError), + (np.empty((1, 4), dtype=np.int32), TypeError), + (np.empty((1, 8), dtype=np.int64)[:, ::2], ValueError), + ): + buffers = init_candidate_buffers(wrong_state) + buffers.source_index = replacement # type: ignore[assignment] + with self.assertRaises(error): + forward_candidates_step_into(wrong_state, torch.tensor([0]), buffers) + + buffers = init_candidate_buffers(wrong_state) + buffers.match_length = buffers.source_index + with self.assertRaisesRegex(ValueError, "overlap"): + forward_candidates_step_into(wrong_state, torch.tensor([0]), buffers) + + token_overlap_state = init_candidate_state(1, 1, suffix_k=1, occurrences_r=1) + token_overlap_state.native_state = False + buffers = init_candidate_buffers(token_overlap_state) + with self.assertRaisesRegex(ValueError, "tokens"): + forward_candidates_step_into( + token_overlap_state, + torch.from_numpy(buffers.source_index.reshape(-1)), + buffers, + ) + + buffers = init_candidate_buffers(wrong_state) + buffers.source_index = wrong_state.lazy_prefix.reshape(1, -1)[:, :4] + with self.assertRaisesRegex(ValueError, "state storage"): + forward_candidates_step_into(wrong_state, torch.tensor([0]), buffers) + + with self.assertRaisesRegex(TypeError, "CandidateBuffers"): + forward_candidates_step_into( + wrong_state, + torch.tensor([0]), + object(), # type: ignore[arg-type] + ) + with self.assertRaisesRegex(TypeError, "CandidateState"): init_candidate_buffers(object()) with self.assertRaisesRegex(ValueError, "sequence_length"): init_candidate_buffers(wrong_state, sequence_length=-1) + ragged = init_candidate_state_internal(1, 2, ragged=True) + with self.assertRaisesRegex(RuntimeError, "ragged"): + forward_candidates_step_into( + ragged, torch.tensor([0]), init_candidate_buffers(ragged) + ) + with self.assertRaisesRegex(RuntimeError, "ragged"): + prefill_candidates_into( + ragged, + torch.tensor([[0]]), + init_candidate_buffers(ragged, sequence_length=1), + ) + + full = init_candidate_state(1, 1) + full.native_state = False + full_buffers = init_candidate_buffers(full) + forward_candidates_step_into(full, torch.tensor([0]), full_buffers) + with self.assertRaisesRegex(RuntimeError, "capacity"): + forward_candidates_step_into(full, torch.tensor([0]), full_buffers) + + nonempty = init_candidate_state(1, 2) + nonempty.native_state = False + forward_candidates_step_into( + nonempty, torch.tensor([0]), init_candidate_buffers(nonempty) + ) + with self.assertRaisesRegex(RuntimeError, "empty"): + prefill_candidates_into( + nonempty, + torch.tensor([[0]]), + init_candidate_buffers(nonempty, sequence_length=1), + ) + + short = init_candidate_state(1, 1) + with self.assertRaisesRegex(RuntimeError, "capacity"): + prefill_candidates_into( + short, + torch.tensor([[0, 1]]), + init_candidate_buffers(short, sequence_length=2), + ) + + native_like = init_candidate_state(1, 1) + native_buffers = init_candidate_buffers(native_like, sequence_length=1) + + def fill_like_native( + candidate_state: CandidateState, + method: str, + tokens: np.ndarray, + arrays: tuple[np.ndarray, ...], + validate_fallback: object, + ) -> bool: + self.assertEqual(method, "prefill_into") + self.assertTrue(callable(validate_fallback)) + validate_fallback() # type: ignore[operator] + validate_fallback() # type: ignore[operator] + for array in arrays: + array.fill(0) + candidate_state.position = tokens.shape[1] + return True + + with patch( + "rosa._stateful_candidates_numba._native_candidate_into", + side_effect=fill_like_native, + ): + prefill_candidates_into(native_like, torch.tensor([[0]]), native_buffers) + self.assertEqual(native_like.positions.tolist(), [1]) + + old_state = init_candidate_state(1, 1, suffix_k=2, occurrences_r=2) + + class AllocatingOnlyNative: + def step(self, tokens: np.ndarray) -> tuple[np.ndarray, ...]: + old_state.position += 1 + old_state.positions.fill(old_state.position) + return ( + np.full((1, 4), -1, dtype=np.int64), + np.zeros((1, 4), dtype=np.int64), + np.full((1, 4), -1, dtype=np.int64), + np.zeros((1, 4), dtype=np.int64), + np.zeros(1, dtype=np.int32), + ) + + old_state.native_state = AllocatingOnlyNative() + old_result = forward_candidates_step_into( + old_state, torch.tensor([0]), init_candidate_buffers(old_state) + ) + self.assertFalse(bool(old_result.mask.any())) + + missing_state = init_candidate_state(1, 1) + missing_state.native_state = object() + missing_result = forward_candidates_step_into( + missing_state, torch.tensor([0]), init_candidate_buffers(missing_state) + ) + self.assertEqual(missing_result.rosa_slot.tolist(), [-1]) + + repeated_validation = init_candidate_state(1, 1) + repeated_validation.native_state = False + + def miss_after_validating( + candidate_state: CandidateState, + method: str, + tokens: np.ndarray, + arrays: tuple[np.ndarray, ...], + validate_fallback: object, + ) -> bool: + validate_fallback() # type: ignore[operator] + validate_fallback() # type: ignore[operator] + return False + + with patch( + "rosa._stateful_candidates_numba._native_candidate_into", + side_effect=miss_after_validating, + ): + forward_candidates_step_into( + repeated_validation, + torch.tensor([0]), + init_candidate_buffers(repeated_validation), + ) + def test_ragged_inactive_reset_capacity_and_recycle(self) -> None: state = init_candidate_state_internal( 3, 5, suffix_k=4, occurrences_r=3, ragged=True diff --git a/tests/test_unified_inference.py b/tests/test_unified_inference.py index 5586034..afcaa93 100644 --- a/tests/test_unified_inference.py +++ b/tests/test_unified_inference.py @@ -10,6 +10,7 @@ InferenceOutput, build_hard_candidates, forward_candidates_step, + forward_candidates_step_into, forward_step, init_candidate_buffers, init_candidate_state, @@ -100,12 +101,22 @@ def test_rich_step_into_uses_explicit_ephemeral_storage(self) -> None: field, ) + wrapper = init_inference_state(1, 1, mode="rich", suffix_k=2, occurrences_r=2) + wrapped = forward_candidates_step_into( + wrapper, torch.tensor(0), init_candidate_buffers(wrapper) + ) + self.assertEqual(wrapped.source_index.ndim, 1) + top1 = init_inference_state(1, 1) with self.assertRaisesRegex(ValueError, "rich"): init_candidate_buffers(top1) with self.assertRaisesRegex(ValueError, "uniform rich"): top1.step_into(torch.tensor([0]), buffers) + rich_ragged = init_inference_state(1, 1, mode="rich", ragged=True) + with self.assertRaisesRegex(ValueError, "uniform rich"): + rich_ragged.step_into(torch.tensor([0]), buffers) + def test_positions_are_copies_and_uniform_position_is_preserved(self) -> None: uniform = init_inference_state(2, 2) forward_step(uniform, torch.tensor([1, 2])) From 7864fb0adeab12492a734379fb92c14457cb0c33 Mon Sep 17 00:00:00 2001 From: Lucas <30107107+aabbdev@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:24:31 +0800 Subject: [PATCH 29/29] Release rosa-torch 0.2.0 --- CHANGELOG.md | 54 +++++ README.md | 59 +++++- pyproject.toml | 2 + src/rosa/__init__.py | 275 +++++++++++++++++++++---- tests/test_rosa.py | 480 +++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 828 insertions(+), 42 deletions(-) create mode 100644 CHANGELOG.md diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..3ec6404 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,54 @@ +# Changelog + +All notable changes to `rosa-torch` are documented here. The project follows +semantic versioning while it remains in the 0.x development series. + +## 0.2.0 — 2026-08-11 + +### Added + +- Unified `ROSAInferenceState` facade for exact top-1 and rich candidate modes, + with uniform and ragged batches, prefill, continuation, reset, and row + recycling. +- Exact bounded rich candidate state with top-K suffix states, top-R newest + occurrences, frequencies, masks, and compatibility wrappers. +- Optional native C++ companion capabilities for top-1/rich step and prefill, + ragged operations, lazy persistent batch workers, and ABI-1-safe fallback. +- Caller-owned `step_into` and `prefill_into` candidate buffers with persistent + NumPy/Torch views for latency-sensitive inference. +- Opt-in `compile_soft_match=True` Inductor island with per-signature cache + isolation and eager forward fallback. + +### Changed + +- Replaced eager suffix-chain occurrence propagation with rooted Link-Cut Tree + lazy path updates, reducing exact online updates to amortized `O(log N)`. +- `ROSA.forward` now uses fused stateful rich prefill when available while + retaining `build_hard_candidates` as an independent exact oracle. +- Moved selector, value, virtual-key, and symbolic projections before candidate + gather, avoiding repeated candidate-wise matrix multiplications. +- Gather chosen symbolic IDs directly instead of constructing temporary + one-hot sequences. +- Parallelized native full-context batch prefill while retaining serial paths + for small batches where dispatch overhead dominates. + +### Compatibility + +- The distribution remains `rosa-torch`; imports remain `from rosa import ...`. +- Python 3.10+ remains supported. +- Existing `forward_step`, `prefill`, `init_candidate_state`, and + `forward_candidates_step` entry points remain available. +- Numba and the native companion remain optional; exact Python/Numba fallbacks + are preserved when newer native capabilities are unavailable. +- `torch.compile` acceleration is opt-in because first-call and new-shape + compilation can take several seconds. + +## 0.1.1 — 2026-08-10 + +- Corrected attribution and clarified that this is an independent + implementation of RWKV-8 ROSA. + +## 0.1.0 — 2026-08-10 + +- Initial PyPI release of the differentiable exact suffix-automaton retrieval + module. diff --git a/README.md b/README.md index 7a39a2a..d8c1e53 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,9 @@ # ROSA as Differentiable Sparse Retrieval with an Exact Suffix Automaton +[![PyPI](https://img.shields.io/pypi/v/rosa-torch.svg)](https://pypi.org/project/rosa-torch/) +[![Python](https://img.shields.io/pypi/pyversions/rosa-torch.svg)](https://pypi.org/project/rosa-torch/) +[![CI](https://github.com/aabbdev/rosa/actions/workflows/ci.yml/badge.svg)](https://github.com/aabbdev/rosa/actions/workflows/ci.yml) + This repository is an independent PyTorch implementation and differentiable extension of **RWKV-8 ROSA (Rapid Online Suffix Automaton)**, described by [Bo Peng (BlinkDL)](https://github.com/BlinkDL) in @@ -27,8 +31,32 @@ The design avoids a trainable dense automaton transition tensor and avoids dense - Learned read gate before the retrieved value is added to the target stream. - Exact ROSA prior plus a learned residual candidate score. - Auxiliary losses for ROSA distillation, hard/soft consistency, codebook balance, and virtual-candidate usage. +- Unified exact `top1`/`rich`, uniform/ragged stateful inference facade. +- Rooted Link-Cut Tree updates with fused full-context prefill. +- Optional C++ companion with parallel batch prefill and reusable output buffers. +- Candidate-wise projections eliminated from the differentiable tensor path. +- 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) +for compatibility notes and the complete release summary. + ## Core scoring rule Candidate ranking is deliberately residual around standard ROSA behavior: @@ -60,13 +88,20 @@ Install the stateful Link-Cut Tree backend with: uv add 'rosa-torch[numba]' ``` -For the lowest CPU step latency, install a locally built native companion wheel -(or a published wheel once multi-ABI releases are enabled): +For the lowest CPU step latency, build and install the optional native companion +locally. `rosa-torch-native` is not currently published on PyPI because it +requires per-platform and per-Python ABI wheels: ```bash +git clone https://github.com/aabbdev/rosa.git +cd rosa +uv build --wheel native --out-dir native/dist uv pip install native/dist/rosa_torch_native-0.2.0-*.whl ``` +The native sources are available from the Git repository and are not included +in the pure-Python `rosa-torch` source distribution on PyPI. + The stateful backend detects it lazily and otherwise falls back to Numba. Install the package and its locked development dependencies with [uv](https://docs.astral.sh/uv/): @@ -87,6 +122,11 @@ PEP 517-compatible Python package manager. . ├── pyproject.toml ├── README.md +├── CHANGELOG.md +├── native +│ ├── pyproject.toml +│ ├── src/rosa_native_step.cpp +│ └── tests ├── src │ └── rosa │ ├── __init__.py @@ -99,6 +139,7 @@ PEP 517-compatible Python package manager. ├── run_coverage.py ├── test_inference.py ├── test_numba_backend.py + ├── test_ragged.py └── test_rosa.py ``` @@ -217,6 +258,20 @@ print(out.chosen_source_index.shape) # [B, N] print(out.hard_rosa_match_length.shape) # [B, N] ``` +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: + +```python +compiled_rosa = ROSA(d_model=64, compile_soft_match=True) +# Run representative forward and backward calls during application warm-up. +``` + +The compiled path reuses one callable per verification window. A compilation +or execution failure during the forward falls back to eager only for that input +signature; other devices and shapes remain eligible for compilation. Deferred +AOTAutograd errors raised during backward are propagated rather than retried. + `z_a` is used to derive the internal symbolic stream and retrieval decisions. `z_b` is the stream receiving the gated retrieval residual. If `z_b` is omitted, `z_a` is used as the target stream as well. ## External code logits diff --git a/pyproject.toml b/pyproject.toml index 9c256df..5ade85d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +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" "Original ROSA description" = "https://www.rwkv.com/#rwkv-8-explained" [dependency-groups] @@ -53,3 +54,4 @@ pythonVersion = "3.10" [tool.uv.build-backend] module-name = "rosa" +source-include = ["CHANGELOG.md"] diff --git a/src/rosa/__init__.py b/src/rosa/__init__.py index f746db2..803c33c 100644 --- a/src/rosa/__init__.py +++ b/src/rosa/__init__.py @@ -9,6 +9,7 @@ from __future__ import annotations import math +import threading from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any, Literal, cast @@ -512,6 +513,205 @@ def _gather_sequence(x: Tensor, index: Tensor) -> Tensor: return x[batch, safe] +def _soft_match_torch( + st1: Tensor, + st2: Tensor, + source_index: Tensor, + candidate_mask: Tensor, + window: int, +) -> Tensor: + """Pure-Torch soft suffix verification with a statically bounded window.""" + + bsz, n, candidates = source_index.shape + positions = torch.arange(n, device=source_index.device).view(1, n, 1) + positions = positions.expand(bsz, n, candidates) + survival = torch.ones((bsz, n, 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 +# signature lock. +_SOFT_MATCH_COMPILE_ENABLED = True +_SOFT_MATCH_COMPILED: dict[int, Any] = {} +_SoftMatchTensorSignature = tuple[ + str, + int | None, + torch.dtype, + torch.layout, + tuple[int, ...], + tuple[int, ...], + int, + bool, +] +_SoftMatchSignature = tuple[ + int, + bool, + bool, + tuple[ + _SoftMatchTensorSignature, + _SoftMatchTensorSignature, + _SoftMatchTensorSignature, + _SoftMatchTensorSignature, + ], +] +_SOFT_MATCH_COMPILE_FAILURES: set[_SoftMatchSignature] = set() +_SOFT_MATCH_COMPILE_READY: set[_SoftMatchSignature] = set() +_SOFT_MATCH_SIGNATURE_LOCKS: dict[_SoftMatchSignature, threading.Lock] = {} +_SOFT_MATCH_WINDOW_LOCKS: dict[int, threading.Lock] = {} +_SOFT_MATCH_COMPILE_MASTER_LOCK = threading.Lock() + + +def _soft_match_signature( + st1: Tensor, + st2: Tensor, + source_index: Tensor, + candidate_mask: Tensor, + window: int, +) -> _SoftMatchSignature: + def tensor_signature(tensor: Tensor) -> _SoftMatchTensorSignature: + return ( + tensor.device.type, + tensor.device.index, + tensor.dtype, + tensor.layout, + tuple(tensor.shape), + tuple(tensor.stride()), + int(tensor.storage_offset()), + tensor.requires_grad, + ) + + tensors = (st1, st2, source_index, candidate_mask) + return ( + window, + torch.is_grad_enabled(), + torch.is_inference_mode_enabled(), + tuple(tensor_signature(tensor) for tensor in tensors), # type: ignore[return-value] + ) + + +def _clear_soft_match_compile_cache() -> None: + """Clear all compiled soft-match process state, primarily for tests.""" + + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + _SOFT_MATCH_COMPILED.clear() + _SOFT_MATCH_COMPILE_FAILURES.clear() + _SOFT_MATCH_COMPILE_READY.clear() + _SOFT_MATCH_SIGNATURE_LOCKS.clear() + _SOFT_MATCH_WINDOW_LOCKS.clear() + + +def _soft_match( + st1: Tensor, + st2: Tensor, + source_index: Tensor, + candidate_mask: Tensor, + window: int, +) -> Tensor: + """Use a static compiled island with eager fallback for forward failures.""" + + if not _SOFT_MATCH_COMPILE_ENABLED: + return _soft_match_torch(st1, st2, source_index, candidate_mask, window) + + signature = _soft_match_signature(st1, st2, source_index, candidate_mask, window) + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + signature_lock = _SOFT_MATCH_SIGNATURE_LOCKS.setdefault( + signature, threading.Lock() + ) + window_lock = _SOFT_MATCH_WINDOW_LOCKS.setdefault(window, threading.Lock()) + if signature in _SOFT_MATCH_COMPILE_FAILURES: + failed = True + ready = False + compiled = None + else: + failed = False + ready = signature in _SOFT_MATCH_COMPILE_READY + compiled = _SOFT_MATCH_COMPILED.get(window) + if failed: + return _soft_match_torch(st1, st2, source_index, candidate_mask, window) + if ready: + assert compiled is not None + try: + return compiled(st1, st2, source_index, candidate_mask) + except Exception: + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + _SOFT_MATCH_COMPILE_FAILURES.add(signature) + return _soft_match_torch(st1, st2, source_index, candidate_mask, window) + + with signature_lock: + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + if signature in _SOFT_MATCH_COMPILE_FAILURES: # pragma: no cover - race + failed = True + ready = False + compiled = None + else: + failed = False + ready = signature in _SOFT_MATCH_COMPILE_READY + compiled = _SOFT_MATCH_COMPILED.get(window) + if not failed and not ready and compiled is None: + with window_lock: + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + compiled = _SOFT_MATCH_COMPILED.get(window) + if compiled is None: # pragma: no branch - concurrent recheck + try: + compiled = torch.compile( + lambda a, b, source, mask: _soft_match_torch( + a, b, source, mask, window + ), + fullgraph=True, + dynamic=False, + ) + except Exception: + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + _SOFT_MATCH_COMPILE_FAILURES.add(signature) + return _soft_match_torch( + st1, st2, source_index, candidate_mask, window + ) + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + _SOFT_MATCH_COMPILED[window] = compiled + + if not failed and not ready: + assert compiled is not None + try: + result = compiled(st1, st2, source_index, candidate_mask) + except Exception: + # Keep the window callable: an unsupported shape, dtype, or + # device must not disable its other specializations. + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + _SOFT_MATCH_COMPILE_FAILURES.add(signature) + return _soft_match_torch(st1, st2, source_index, candidate_mask, window) + + # torch.compile is lazy, so the signature becomes ready only after + # its first forward execution succeeds. Deferred AOT backward + # errors are intentionally propagated to the caller. + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + _SOFT_MATCH_COMPILE_READY.add(signature) + return result + + if failed: # pragma: no cover - concurrent recheck + return _soft_match_torch(st1, st2, source_index, candidate_mask, window) + assert ready and compiled is not None + try: + return compiled(st1, st2, source_index, candidate_mask) + except Exception: # pragma: no cover - concurrent runtime failure + with _SOFT_MATCH_COMPILE_MASTER_LOCK: + _SOFT_MATCH_COMPILE_FAILURES.add(signature) + return _soft_match_torch(st1, st2, source_index, candidate_mask, window) + + def _st_categorical( logits: Tensor, temperature: float ) -> tuple[Tensor, Tensor, Tensor]: @@ -575,6 +775,7 @@ def __init__( virtual_scale: float = 0.0, neural_value_scale: float = 0.0, candidate_backend: CandidateBackend = "auto", + compile_soft_match: bool = False, ) -> None: super().__init__() if d_model <= 0: @@ -595,6 +796,8 @@ def __init__( ) if not isinstance(soft_candidates_forward, bool): raise TypeError("soft_candidates_forward must be a bool") + if not isinstance(compile_soft_match, bool): + raise TypeError("compile_soft_match must be a bool") if selector_dim <= 0: raise ValueError("selector_dim must be > 0") if token_temperature <= 0 or retrieval_temperature <= 0: @@ -625,6 +828,7 @@ def __init__( self.sparse_old_candidates = sparse_old_candidates self.sparse_old_pool_size = sparse_old_pool_size self.soft_candidates_forward = soft_candidates_forward + self.compile_soft_match = compile_soft_match self.candidate_backend: CandidateBackend = candidate_backend self.selector_dim = selector_dim self.token_temperature = token_temperature @@ -722,23 +926,14 @@ def _soft_match( source_index: Tensor, candidate_mask: Tensor, ) -> Tensor: - bsz, n, candidates = source_index.shape - positions = torch.arange(n, device=source_index.device).view(1, n, 1) - positions = positions.expand(bsz, n, candidates) - survival = torch.ones((bsz, n, candidates), dtype=st1.dtype, device=st1.device) - score = torch.zeros_like(survival) - for r in range(self.soft_verify_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 + implementation = _soft_match if self.compile_soft_match else _soft_match_torch + return implementation( + st1, + st2, + source_index, + candidate_mask, + self.soft_verify_window, + ) def _virtual_candidates( self, @@ -754,9 +949,8 @@ def _virtual_candidates( ) & exact_mask.unsqueeze(-2) pool_mask = pool_mask & ~duplicate.any(dim=-1) - pool_z = _gather_sequence(z_a, pool) q = self.virtual_query(z_a).unsqueeze(-2) - k = self.virtual_key(pool_z) + k = self._virtual_pool_keys(z_a, pool) router_score = (q * k).sum(-1) / math.sqrt(self.selector_dim) masked = router_score.masked_fill(~pool_mask, -1e9) _, top_idx = torch.topk(masked, k=self.virtual_candidates, dim=-1) @@ -765,6 +959,12 @@ def _virtual_candidates( selected_router = router_score.gather(-1, top_idx) return selected_source, selected_mask, selected_router + def _virtual_pool_keys(self, z_a: Tensor, pool: Tensor) -> Tensor: + return _gather_sequence(self.virtual_key(z_a), pool) + + def _candidate_selector_keys(self, z_a: Tensor, source: Tensor) -> Tensor: + return _gather_sequence(self.selector_key(z_a), source) + def _candidate_symbolic_values( self, st1: Tensor, @@ -772,11 +972,14 @@ def _candidate_symbolic_values( next_position: Tensor, mask: Tensor, ) -> Tensor: - g1 = _gather_sequence(st1, next_position) - g2 = _gather_sequence(st2, next_position) - e1 = g1 @ self.symbol_embedding_1.weight - e2 = g2 @ self.symbol_embedding_2.weight - return (e1 + e2) * mask.unsqueeze(-1).to(e1.dtype) + symbol_sequence = ( + st1 @ self.symbol_embedding_1.weight + st2 @ self.symbol_embedding_2.weight + ) + symbolic_value = _gather_sequence(symbol_sequence, next_position) + return symbolic_value * mask.unsqueeze(-1).to(symbolic_value.dtype) + + def _candidate_neural_values(self, z_a: Tensor, next_position: Tensor) -> Tensor: + return _gather_sequence(self.value_proj(z_a), next_position) def _hybrid_soft_candidates( self, @@ -982,10 +1185,9 @@ def forward( dim=-1, ) - source_z = _gather_sequence(z_a, source) query = self.selector_query(z_a).unsqueeze(-2) - key = self.selector_key(source_z) - semantic = (query * key).sum(-1) / math.sqrt(self.selector_dim) + cand_key = self._candidate_selector_keys(z_a, source) + 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_a).squeeze(-1) @@ -1037,14 +1239,13 @@ def forward( symbolic_value = self._candidate_symbolic_values( st1, st2, next_position, non_null_mask ) - next_z = _gather_sequence(z_a, next_position) - cand_key = self.selector_key(source_z) 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(next_z) * value_gate.unsqueeze(-1) + 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 @@ -1098,16 +1299,10 @@ def forward( 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] - # Gather both factors to make the chosen token explicitly causal and - # avoid depending on any flattened-token table. - c1 = _gather_sequence( - F.one_hot(chosen_id1, self.codebook_sizes[0]).to(z_a.dtype), - chosen_next, - ).argmax(-1) - c2 = _gather_sequence( - F.one_hot(chosen_id2, self.codebook_sizes[1]).to(z_a.dtype), - chosen_next, - ).argmax(-1) + # The IDs are already discrete. Gather them directly instead of + # allocating two one-hot sequences only to immediately argmax them. + 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 diff --git a/tests/test_rosa.py b/tests/test_rosa.py index 9b20c8c..b74c94f 100644 --- a/tests/test_rosa.py +++ b/tests/test_rosa.py @@ -2,7 +2,9 @@ import copy import random +import threading import unittest +from concurrent.futures import ThreadPoolExecutor from unittest.mock import patch import torch @@ -15,7 +17,11 @@ VIRTUAL_KIND, _balance_kl, _build_forward_hard_candidates, + _clear_soft_match_compile_cache, _gather_sequence, + _soft_match, + _soft_match_signature, + _soft_match_torch, _st_categorical, _virtual_pool_single, build_hard_candidates, @@ -56,6 +62,36 @@ def zero_learned_scorer(model: ROSA) -> None: model.null_head.bias.zero_() +class _GatherFirstROSA(ROSA): + """Local pre-fusion oracle retaining candidate-first projection order.""" + + def _virtual_pool_keys(self, z_a: torch.Tensor, pool: torch.Tensor) -> torch.Tensor: + return self.virtual_key(_gather_sequence(z_a, pool)) + + def _candidate_selector_keys( + self, z_a: torch.Tensor, source: torch.Tensor + ) -> torch.Tensor: + return self.selector_key(_gather_sequence(z_a, source)) + + def _candidate_symbolic_values( + self, + st1: torch.Tensor, + st2: torch.Tensor, + next_position: torch.Tensor, + mask: torch.Tensor, + ) -> torch.Tensor: + g1 = _gather_sequence(st1, next_position) + g2 = _gather_sequence(st2, next_position) + e1 = g1 @ self.symbol_embedding_1.weight + e2 = g2 @ self.symbol_embedding_2.weight + return (e1 + e2) * mask.unsqueeze(-1).to(e1.dtype) + + def _candidate_neural_values( + self, z_a: torch.Tensor, next_position: torch.Tensor + ) -> torch.Tensor: + return self.value_proj(_gather_sequence(z_a, next_position)) + + class TestReferenceROSA(unittest.TestCase): def test_reference_squeeze_batch_and_validation(self) -> None: one_d = torch.tensor([0, 1, 0, 2], dtype=torch.long) @@ -165,6 +201,261 @@ def test_gather_st_and_balance_helpers(self) -> None: self.assertAlmostEqual(float(_balance_kl(uniform)), 0.0, places=6) self.assertGreater(float(_balance_kl(peaked)), 1.0) + def test_compiled_soft_match_matches_eager_across_specializations(self) -> None: + devices = [torch.device("cpu")] + if torch.cuda.is_available(): + devices.append(torch.device("cuda")) + + for device in devices: + with self.subTest(device=device.type): + _clear_soft_match_compile_cache() + cached_callable = None + for shape in ((2, 7, 5), (1, 11, 3)): + batch, length, candidates = shape + torch.manual_seed(20260811 + length) + actual_st1 = torch.softmax( + torch.randn(batch, length, 4, device=device), dim=-1 + ).requires_grad_() + actual_st2 = torch.softmax( + torch.randn(batch, length, 3, device=device), dim=-1 + ).requires_grad_() + expected_st1 = actual_st1.detach().clone().requires_grad_() + expected_st2 = actual_st2.detach().clone().requires_grad_() + source = torch.randint( + -1, length, (batch, length, candidates), device=device + ) + mask = torch.rand(batch, length, candidates, device=device) > 0.2 + + actual = _soft_match(actual_st1, actual_st2, source, mask, window=4) + expected = _soft_match_torch( + expected_st1, expected_st2, source, mask, window=4 + ) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + actual.square().sum().backward() + expected.square().sum().backward() + assert actual_st1.grad is not None + assert actual_st2.grad is not None + assert expected_st1.grad is not None + assert expected_st2.grad is not None + torch.testing.assert_close( + actual_st1.grad, expected_st1.grad, rtol=2e-5, atol=2e-6 + ) + torch.testing.assert_close( + actual_st2.grad, expected_st2.grad, rtol=2e-5, atol=2e-6 + ) + + signature = _soft_match_signature( + actual_st1, actual_st2, source, mask, 4 + ) + if signature not in rosa._SOFT_MATCH_COMPILE_FAILURES: + self.assertEqual(len(rosa._SOFT_MATCH_COMPILED), 1) + self.assertIn(signature, rosa._SOFT_MATCH_COMPILE_READY) + if cached_callable is None: + cached_callable = rosa._SOFT_MATCH_COMPILED[4] + else: + self.assertIs(cached_callable, rosa._SOFT_MATCH_COMPILED[4]) + + def test_soft_match_compile_failure_falls_back_before_caching(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 + expected = _soft_match_torch(st1, st2, source, mask, window=3) + with patch("rosa.torch.compile", side_effect=RuntimeError("unavailable")): + actual = _soft_match(st1, st2, source, mask, window=3) + signature = _soft_match_signature(st1, st2, source, mask, 3) + self.assertTrue(torch.equal(actual, expected)) + self.assertNotIn(3, rosa._SOFT_MATCH_COMPILED) + self.assertIn(signature, rosa._SOFT_MATCH_COMPILE_FAILURES) + _clear_soft_match_compile_cache() + + def test_soft_match_disabled_and_ready_failure_fall_back(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 + expected = _soft_match_torch(st1, st2, source, mask, window=3) + with patch("rosa._SOFT_MATCH_COMPILE_ENABLED", False): + disabled = _soft_match(st1, st2, source, mask, window=3) + self.assertTrue(torch.equal(disabled, expected)) + + signature = _soft_match_signature(st1, st2, source, mask, 3) + + def fail(*args: torch.Tensor) -> torch.Tensor: + raise RuntimeError("cached specialization failed") + + rosa._SOFT_MATCH_COMPILED[3] = fail + rosa._SOFT_MATCH_COMPILE_READY.add(signature) + actual = _soft_match(st1, st2, source, mask, window=3) + self.assertTrue(torch.equal(actual, expected)) + self.assertIn(signature, rosa._SOFT_MATCH_COMPILE_FAILURES) + _clear_soft_match_compile_cache() + + def test_soft_match_compile_initialization_is_thread_safe(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 + start = threading.Barrier(2) + compile_calls = 0 + forward_calls = 0 + calls_lock = threading.Lock() + + def compiled(*args: torch.Tensor) -> torch.Tensor: + nonlocal forward_calls + with calls_lock: + forward_calls += 1 + return _soft_match_torch(*args, window=3) + + def compile_once(*args: object, **kwargs: object) -> object: + nonlocal compile_calls + with calls_lock: + compile_calls += 1 + return compiled + + def invoke() -> torch.Tensor: + start.wait() + return _soft_match(st1, st2, source, mask, window=3) + + with ( + patch("rosa.torch.compile", side_effect=compile_once), + ThreadPoolExecutor(max_workers=2) as executor, + ): + results = [ + future.result() + for future in (executor.submit(invoke), executor.submit(invoke)) + ] + + self.assertEqual(compile_calls, 1) + self.assertEqual(forward_calls, 2) + self.assertTrue(torch.equal(results[0], results[1])) + signature = _soft_match_signature(st1, st2, source, mask, 3) + self.assertIn(signature, rosa._SOFT_MATCH_COMPILE_READY) + + def test_soft_match_failure_does_not_poison_other_signature(self) -> None: + _clear_soft_match_compile_cache() + calls: list[tuple[int, ...]] = [] + + def compiled( + st1: torch.Tensor, + st2: torch.Tensor, + source: torch.Tensor, + mask: torch.Tensor, + ) -> torch.Tensor: + calls.append(tuple(source.shape)) + if source.shape[1] == 5: + raise RuntimeError("unsupported shape") + return _soft_match_torch(st1, st2, source, mask, window=3) + + def inputs(length: int) -> tuple[torch.Tensor, ...]: + st1 = torch.softmax(torch.randn(1, length, 3), dim=-1) + st2 = torch.softmax(torch.randn(1, length, 2), dim=-1) + source = torch.randint(-1, length, (1, length, 4)) + return st1, st2, source, source >= 0 + + failing = inputs(5) + working = inputs(7) + with patch("rosa.torch.compile", return_value=compiled) as compile_mock: + failed_result = _soft_match(*failing, window=3) + working_result = _soft_match(*working, window=3) + retried_result = _soft_match(*failing, window=3) + + self.assertEqual(compile_mock.call_count, 1) + self.assertEqual(calls, [(1, 5, 4), (1, 7, 4)]) + self.assertTrue( + torch.equal(failed_result, _soft_match_torch(*failing, window=3)) + ) + self.assertTrue( + torch.equal(retried_result, _soft_match_torch(*failing, window=3)) + ) + self.assertTrue( + torch.equal(working_result, _soft_match_torch(*working, window=3)) + ) + self.assertIn(3, rosa._SOFT_MATCH_COMPILED) + self.assertIn( + _soft_match_signature(*failing, window=3), + rosa._SOFT_MATCH_COMPILE_FAILURES, + ) + self.assertIn( + _soft_match_signature(*working, window=3), + rosa._SOFT_MATCH_COMPILE_READY, + ) + + def test_soft_match_signature_distinguishes_device(self) -> None: + cpu = torch.empty(1, 2, 3) + cpu_source = torch.empty(1, 2, 4, dtype=torch.long) + meta = torch.empty(1, 2, 3, device="meta") + meta_source = torch.empty(1, 2, 4, dtype=torch.long, device="meta") + cpu_signature = _soft_match_signature(cpu, cpu, cpu_source, cpu_source >= 0, 3) + meta_signature = _soft_match_signature( + meta, meta, meta_source, meta_source >= 0, 3 + ) + self.assertNotEqual(cpu_signature, meta_signature) + + def test_soft_match_signature_distinguishes_layout_and_grad_mode(self) -> None: + st1 = torch.randn(1, 5, 3, requires_grad=True) + st2 = torch.randn(1, 5, 2, requires_grad=True) + source = torch.zeros(1, 5, 4, dtype=torch.long) + mask = torch.ones_like(source, dtype=torch.bool) + baseline = _soft_match_signature(st1, st2, source, mask, 3) + noncontiguous_st2 = torch.randn(1, 2, 5).transpose(1, 2).requires_grad_() + self.assertNotEqual( + baseline, + _soft_match_signature(st1, noncontiguous_st2, source, mask, 3), + ) + self.assertNotEqual( + baseline, + _soft_match_signature(st1.detach(), st2, source, mask, 3), + ) + with torch.no_grad(): + no_grad = _soft_match_signature(st1, st2, source, mask, 3) + self.assertNotEqual(baseline, no_grad) + with torch.inference_mode(): + inference = _soft_match_signature(st1, st2, source, mask, 3) + self.assertNotEqual(no_grad, inference) + + def test_soft_match_backward_compile_error_is_propagated(self) -> None: + _clear_soft_match_compile_cache() + st1 = torch.randn(1, 3, 2, requires_grad=True) + st2 = torch.randn(1, 3, 2, requires_grad=True) + source = torch.zeros(1, 3, 1, dtype=torch.long) + mask = torch.ones_like(source, dtype=torch.bool) + + class BackwardFailure(torch.autograd.Function): + @staticmethod + def forward(ctx: object, value: torch.Tensor) -> torch.Tensor: + return value.sum() + + @staticmethod + def backward(ctx: object, grad: torch.Tensor) -> torch.Tensor: + raise RuntimeError("deferred AOT backward failure") + + def compiled(*args: torch.Tensor) -> torch.Tensor: + return BackwardFailure.apply(args[0]) + + with patch("rosa.torch.compile", return_value=compiled): + result = _soft_match(st1, st2, source, mask, window=2) + + with self.assertRaisesRegex(RuntimeError, "deferred AOT backward failure"): + result.backward() + + def test_rosa_soft_match_is_eager_by_default(self) -> None: + model = ROSA(d_model=4, soft_verify_window=2) + st1 = torch.softmax(torch.randn(1, 4, 3), dim=-1) + st2 = torch.softmax(torch.randn(1, 4, 2), dim=-1) + source = torch.randint(-1, 4, (1, 4, 2)) + mask = source >= 0 + expected = _soft_match_torch(st1, st2, source, mask, window=2) + + with patch("rosa.torch.compile") as compile_mock: + actual = model._soft_match(st1, st2, source, mask) + + compile_mock.assert_not_called() + self.assertTrue(torch.equal(actual, expected)) + class TestROSAConfiguration(unittest.TestCase): def test_constructor_validations(self) -> None: @@ -209,6 +500,8 @@ def test_constructor_validations(self) -> None: ROSA(**kwargs) with self.assertRaisesRegex(TypeError, "soft_candidates_forward"): ROSA(d_model=4, soft_candidates_forward=1) # type: ignore[arg-type] + with self.assertRaisesRegex(TypeError, "compile_soft_match"): + ROSA(d_model=4, compile_soft_match=1) # type: ignore[arg-type] def test_setters_property_and_encode_validation(self) -> None: model = ROSA( @@ -222,6 +515,24 @@ def test_setters_property_and_encode_validation(self) -> None: selector_dim=5, ) self.assertEqual(model.vocab_size, 6) + self.assertFalse(model.compile_soft_match) + self.assertTrue(ROSA(d_model=4, compile_soft_match=True).compile_soft_match) + positional = ROSA( + 4, + (2, 2), + 2, + 2, + 3, + 2, + 4, + 0, + 0, + 4, + False, + 5, + ) + self.assertEqual(positional.selector_dim, 5) + self.assertFalse(positional.compile_soft_match) model.set_learned_residual_scale(0.4) model.set_virtual_scale(0.5) model.set_neural_value_scale(0.6) @@ -348,6 +659,175 @@ def test_python_and_stateful_backends_match_all_fields_outputs_and_gradients( assert expected.grad is not None self.assertTrue(torch.equal(actual.grad, expected.grad), actual_name) + def test_project_before_gather_matches_gather_first_oracle(self) -> None: + # Moving a linear projection across a gather changes GEMM row batching, + # so FP32 accumulation may differ slightly. These tolerances cover that + # expected roundoff while remaining tight enough to catch path changes. + devices = [torch.device("cpu")] + if torch.cuda.is_available(): + devices.append(torch.device("cuda")) + + for device in devices: + with self.subTest(device=device.type): + torch.manual_seed(20260811) + optimized = self.make_model( + dense_recent_candidates=3, + sparse_old_candidates=2, + sparse_old_pool_size=5, + learned_residual_scale=1.0, + virtual_scale=1.0, + neural_value_scale=1.0, + read_gate_bias=0.0, + value_gate_bias=0.0, + candidate_backend="python", + compile_soft_match=False, + ).to(device) + oracle = _GatherFirstROSA( + d_model=8, + codebook_sizes=(2, 3), + suffix_k=5, + occurrences_r=3, + soft_verify_window=6, + virtual_candidates=2, + virtual_pool_size=6, + dense_recent_candidates=3, + sparse_old_candidates=2, + sparse_old_pool_size=5, + selector_dim=8, + token_temperature=0.2, + retrieval_temperature=0.7, + learned_residual_scale=1.0, + virtual_scale=1.0, + neural_value_scale=1.0, + read_gate_bias=0.0, + value_gate_bias=0.0, + candidate_backend="python", + compile_soft_match=False, + ).to(device) + oracle.load_state_dict(optimized.state_dict()) + + z_optimized = torch.randn(2, 13, 8, device=device, requires_grad=True) + z_oracle = z_optimized.detach().clone().requires_grad_() + target = torch.randn_like(z_optimized) + logits_optimized = tuple( + torch.randn(2, 13, size, device=device, requires_grad=True) + for size in (2, 3) + ) + logits_oracle = tuple( + item.detach().clone().requires_grad_() for item in logits_optimized + ) + + actual = optimized(z_optimized, code_logits=logits_optimized) + expected = oracle(z_oracle, code_logits=logits_oracle) + exact_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", + ) + for name in exact_fields: + self.assertTrue( + torch.equal(getattr(actual, name), getattr(expected, name)), + name, + ) + + rtol, atol = (1e-4, 2e-5) if device.type == "cuda" else (3e-5, 3e-6) + + def assert_close_nested( + actual_value, expected_value, name: str + ) -> None: + if isinstance(actual_value, torch.Tensor): + if actual_value.is_floating_point(): + torch.testing.assert_close( + actual_value, + expected_value, + rtol=rtol, + atol=atol, + msg=name, + ) + else: + self.assertTrue( + torch.equal(actual_value, expected_value), name + ) + elif isinstance(actual_value, tuple): + for index, (actual_item, expected_item) in enumerate( + zip(actual_value, expected_value, strict=True) + ): + assert_close_nested( + actual_item, expected_item, f"{name}[{index}]" + ) + elif isinstance(actual_value, dict): + self.assertEqual(actual_value.keys(), expected_value.keys()) + for key in actual_value: + assert_close_nested( + actual_value[key], expected_value[key], f"{name}.{key}" + ) + else: + self.assertEqual(actual_value, expected_value, name) + + for name in actual.__dataclass_fields__: + assert_close_nested( + getattr(actual, name), getattr(expected, name), name + ) + + actual_loss = F.mse_loss(actual.updated, target) + sum( + actual.aux_losses.values() + ) + expected_loss = F.mse_loss(expected.updated, target) + sum( + expected.aux_losses.values() + ) + torch.testing.assert_close( + actual_loss, expected_loss, rtol=rtol, atol=atol + ) + actual_loss.backward() + expected_loss.backward() + + assert z_optimized.grad is not None + assert z_oracle.grad is not None + torch.testing.assert_close( + z_optimized.grad, z_oracle.grad, rtol=rtol, atol=atol + ) + for actual_logits, expected_logits in zip( + logits_optimized, logits_oracle, strict=True + ): + assert actual_logits.grad is not None + assert expected_logits.grad is not None + torch.testing.assert_close( + actual_logits.grad, + expected_logits.grad, + rtol=rtol, + atol=atol, + ) + for (actual_name, actual_parameter), ( + expected_name, + expected_parameter, + ) in zip( + optimized.named_parameters(), oracle.named_parameters(), strict=True + ): + self.assertEqual(actual_name, expected_name) + self.assertEqual( + actual_parameter.grad is None, + expected_parameter.grad is None, + actual_name, + ) + if actual_parameter.grad is not None: + assert expected_parameter.grad is not None + torch.testing.assert_close( + actual_parameter.grad, + expected_parameter.grad, + rtol=rtol, + atol=atol, + msg=actual_name, + ) + def test_stateful_forward_does_not_call_eager_or_suffix_write(self) -> None: from rosa._stateful_candidates_numba import prefill_candidates