Research/rosa torch fusion 20260811

.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 .
@@ -54,3 +54,64 @@ 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
+
+  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

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.

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:
@@ -45,6 +73,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 +82,28 @@ Install the published package from PyPI:
 uv add rosa-torch
 ```
 
+Install the stateful Link-Cut Tree backend with:
+
+```bash
+uv add 'rosa-torch[numba]'
+```
+
+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/):
 
 ```bash
@@ -71,18 +122,103 @@ 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
+│       ├── __init__.py
+│       ├── _numba_backend.py
+│       ├── _stateful_candidates_numba.py
+│       ├── ragged.py
+│       └── _stateful_numba.py
 └── tests
     ├── __init__.py
     ├── run_coverage.py
+    ├── test_inference.py
+    ├── test_numba_backend.py
+    ├── test_ragged.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.
+
+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.
+
+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
 
@@ -107,6 +243,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)
@@ -121,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
@@ -263,11 +414,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
 

native/.gitignore

@@ -0,0 +1,5 @@
+build/
+dist/
+*.egg-info/
+*.so
+*.pyd