Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ requires per-platform and per-Python ABI wheels:
```bash
git clone https://github.com/aabbdev/rosa.git
cd rosa
uv sync --extra numba
uv build --wheel native --out-dir native/dist
uv pip install native/dist/rosa_torch_native-0.2.0-*.whl
```
Expand Down
191 changes: 100 additions & 91 deletions native/README.md
Original file line number Diff line number Diff line change
@@ -1,113 +1,122 @@
# `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`.

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`.

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

Installez un wheel correspondant à la version de Python et à la plateforme :
Optional native companion for
[`rosa-torch`](https://github.com/aabbdev/rosa), accelerating the exact CPU
suffix-automaton and Link-Cut Tree inference step. The distribution installs
the importable `rosa_native_step` extension module; it does not replace the
main Python package.

The C++ core implements the validated exact production state machine. It binds
the NumPy arrays of a `_StatefulInferenceState` once, updates them in place,
and releases the GIL during computation. It neither includes nor calls
libtorch. The runtime dependency `rosa-torch[numba]>=0.2,<0.3` provides the
compatible state contract together with PyTorch, NumPy, and Numba.

The constructor validates every shape, dtype, counter, and ABI version before
retaining any pointer. The current native state ABI is `1`.

The module also exposes `NativeCandidateState` for the exact rich state in
`rosa._stateful_candidates_numba`. Its batched `step` maintains the same K
suffixes, R newest occurrences, unbounded frequencies, and
`newest-prefix + delta` LCT tags as the Python/Numba implementation. The R
capacity remains owned by the NumPy arrays of the Python `CandidateState`, whose
lifetime is retained by the native object. Capability can be detected through
the presence of `NativeCandidateState` and `candidate_abi_version == 1`.

The rich ABI retains `step`, global `reset`, and `position`, while detecting the
optional `prefill`, `step_masked`, `reset_masked`, `positions`, `step_into`, and
`prefill_into` capabilities. `prefill` emits all five native arrays at every
position in one C++ call and leaves the state ready for continuation. The
caller-owned `*_into` methods avoid repeated output allocation after validating
dtype, shape, contiguity, writability, and memory overlap.

Ragged mode stores one position per row. Uniform and ragged paths are
intentionally incompatible so that exactly one position authority exists at
any time. The Python wrapper falls back exactly to Numba when an older ABI-1
wheel does not provide a newer optional method.

## Installation and usage

`rosa-torch-native` is not currently published on PyPI. Build a wheel from a
Git checkout using the target Python interpreter, then install the matching
wheel for the current platform and ABI:

```bash
python -m pip install rosa-torch-native
git clone https://github.com/aabbdev/rosa.git
cd rosa
uv sync --extra numba
uv build --python .venv/bin/python --wheel native --out-dir native/dist
uv pip install --python .venv/bin/python \
native/dist/rosa_torch_native-0.2.0-*.whl
```

`rosa-torch` détecte automatiquement le module dans son backend d'inférence
Numba. L'API bas niveau reste disponible pour diagnostic :
`rosa-torch` detects the extension automatically from its Numba inference
backend. The low-level API remains available for diagnostics:

```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.

`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.

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
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 :
`NativeState(state).step(tokens_numpy)` expects a contiguous NumPy vector of
shape `[batch_size]` that is convertible to `int64`. The object retains a
reference to the Python state and exposes its read-only `position`.

`NativeCandidateState(candidate_state).step(tokens_numpy)` requires a
C-contiguous `int64` NumPy vector and returns the low-level tuple
`(source, match_length, state_id, frequency, count)`. `reset()` recycles the
whole batch in time proportional to the state nodes and hash slots that are
actually occupied.

Each native state lazily creates a small dependency-free persistent C++17
thread pool. It parallelizes independent batch rows for top-1 and rich prefill.
Prefill keeps a serial path below batch 4, where waking a worker costs more than
the measured computation. Uniform steps use the pool only from batch 64 onward
to preserve the inexpensive serial path for small batches.

The pool is bounded by the batch size, reported CPU count, and 16 total threads.
`ROSA_NATIVE_THREADS=1` forces the serial path; another positive integer sets a
lower cap. An unset variable selects the available limit automatically,
while an invalid value falls back conservatively to one thread. Only non-empty
digit strings representing a strictly positive integer are accepted.

Using one pool per state avoids fragile extension-singleton shutdown ordering,
and workers never call Python. No worker is created for a ragged state or a
batch below 4 that never uses prefill. A sufficiently large first call may
create at most 15 workers, which are reused until state destruction. Concurrent
mutating calls on the same instance are serialized by a mutex acquired while
the GIL is released. Public output allocation and array ownership remain
unchanged. Worker exceptions are captured and rethrown on the calling thread.

## Isolated local build

The PEP 517 backend is setuptools, with setuptools, wheel, and pybind11 as build
dependencies. `Pybind11Extension` selects C++17, while setuptools provides the
platform extension flags for macOS and Linux. Both platforms additionally use
`-O3` and `NDEBUG`.

From the repository root:

```bash
uv build --wheel native --out-dir /tmp/rosa-native-dist
uv sync --extra numba
uv build --python .venv/bin/python --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.
The smoke test forces Numba as the oracle, lets the regular ROSA path load the
companion, and compares predictions step by step.

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é.
Run the rich smoke test and direct Numba comparison with
`native/tests/candidate_smoke.py` and `native/benchmark_candidates.py`,
respectively, in the same isolated environment.

## Publication multi-plateforme
## Multi-platform publication

É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.
The next publication step is a dedicated `cibuildwheel` workflow after defining
the supported Python and architecture matrix, manylinux targets, and release
policy. No such workflow is included yet because an unvalidated matrix must not
be published. Every wheel contains native code and must be built separately for
each Python ABI and platform.
Loading