diff --git a/.github/workflows/bench-baseline.yml b/.github/workflows/bench-baseline.yml index f6ed27d..9d624b4 100644 --- a/.github/workflows/bench-baseline.yml +++ b/.github/workflows/bench-baseline.yml @@ -8,6 +8,29 @@ name: bench-baseline on: workflow_dispatch: + inputs: + run_full_baseline: + description: >- + Also run the full baseline suite (BASELINE_TEST_PLAN.md). Uncheck to + run ONLY the bidirectional long-L sweep, which is usually what you + want — the full suite takes far longer. + default: true + type: boolean + long_l_suite: + description: "bidirectional long-L suite (sweep-len covers L=128..8192)" + default: "sweep-len" + type: choice + options: [sweep-len, sweep-dim, basic] + long_l_compile_max: + description: >- + Max L for the torch.compile baseline. MEASURED: compile time is LINEAR + in L at ~0.26 s per timestep (59s@128, 130s@512, 245s@1024, 524s@2048; + reproduced across two runs). So 2048 costs ~9 min of compiling, 4096 + ~18 min, and 8192 ~36 min — a full sweep to 8192 is ~75 min. At 8192 + inductor may also OOM building the unrolled graph, which is itself the + finding. Set 0 to skip torch.compile entirely. + default: "2048" + type: string push: branches: [main] paths: @@ -19,6 +42,8 @@ permissions: jobs: baseline: name: baseline (linux-arm64) + # Unchanged on push; on dispatch, opt out to run only the long-L sweep. + if: github.event_name != 'workflow_dispatch' || inputs.run_full_baseline runs-on: ubuntu-24.04-arm timeout-minutes: 120 steps: @@ -65,3 +90,87 @@ jobs: cat comment.md >> "$GITHUB_STEP_SUMMARY" gh api "repos/${{ github.repository }}/commits/${{ github.sha }}/comments" \ -f body="$(cat comment.md)" --silent || true + + # Bidirectional scan at long sequence lengths — the regime the applications + # actually live in (genomics @131k, long audio, multi-hour ECG), and the one + # per-PR CI cannot afford. Dispatch-only: torch.compile unrolls the recurrence + # into an L-step graph, so the compile step is the entire cost (~8 min at + # L<=2048, ~20 min including 8192) and may OOM outright at 8192. + # + # BOTH outcomes are results worth having: + # - it compiles -> the vs-torch.compile ratio at the lengths that matter + # (the trend at 128->512 says it should exceed 4.67x) + # - it does not -> a hard limit on the baseline itself, at exactly the + # sequence lengths our headline applications need + long-l: + name: bidirectional long-L (linux-arm64) + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-24.04-arm + # Compile cost is ~0.26 s/timestep and LINEAR in L, so a full sweep with + # compile_max_len=8192 needs ~70 min of inductor alone (35 min for L=8192 + # by itself). 150 gives real headroom: if the job is KILLED we want to know + # it was an OOM, not that we clipped the timeout — those are different + # findings and only one of them is interesting. + timeout-minutes: 150 + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + with: + workspaces: kernel + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Build FFI cdylib + run: cargo build --release -p arm-scan-ffi + working-directory: kernel + - name: Install python deps + run: python3 -m pip install --quiet numpy torch + + # Correctness gates speed (CLAUDE.md): never benchmark a broken kernel. + - name: Bidirectional correctness through the kernel + run: python3 tests/check_bidirectional.py + + - name: Bidirectional long-L sweep + env: + SUITE: ${{ inputs.long_l_suite }} + CMAX: ${{ inputs.long_l_compile_max }} + run: | + set -o pipefail + mkdir -p bench/results + EXTRA="" + if [ "$CMAX" = "0" ]; then + EXTRA="--no-compile" + else + EXTRA="--compile-max-len $CMAX" + fi + # The bench flushes its JSON after every shape, so an OOM kill at the + # longest L still leaves every shorter shape's results on disk. + python3 bench/bench_bidirectional.py \ + --suite "$SUITE" $EXTRA --tag ci-arm64-longl \ + --json bench/results/bidi_longl_ci-arm64.json 2>&1 \ + | tee bench-bidi-longl.log + + - uses: actions/upload-artifact@v4 + if: always() + with: + name: bidirectional-long-l + path: | + bench/results/bidi_longl_ci-arm64.json + bench-bidi-longl.log + if-no-files-found: warn + + - name: Job summary + if: always() + run: | + { + echo "## Bidirectional scan, long L (linux-arm64, provisional)" + echo "suite=\`${{ inputs.long_l_suite }}\` compile_max_len=\`${{ inputs.long_l_compile_max }}\`" + echo + echo "If a shape is missing, torch.compile was killed building its" + echo "unrolled graph — which is itself the finding. Partial results" + echo "are preserved in the artifact." + echo '```' + grep -E "^===|fused ==|^ (ref_|scan_fwd|bidirectional)|^ =>|exp-sharing|^bidirectional \(fused\)|torch.compile COST|^ +[0-9]+ +[0-9.]+s|^ +L|^ L=" bench-bidi-longl.log 2>/dev/null || echo "(sweep did not complete)" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6ec3c6b..96aa7a2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -50,6 +50,13 @@ jobs: python3 -m pip install --quiet numpy \ || python3 -m pip install --quiet --break-system-packages numpy python3 tests/check_ffi.py + # Definition check for the bidirectional scan: flip-forward-flip must + # equal a backward-in-time recurrence. numpy-only (no torch, no kernel), + # so it runs on every platform. Also the executable spec for the future + # Rust `reverse` flag — see TOPOLOGY_IMPLEMENTATION_PLAN.md §2. + - name: Bidirectional definition check (numpy) + working-directory: . + run: python3 tests/check_bidirectional_math.py - name: Bench ladder (scalar vs NEON) if: matrix.name == 'linux-arm64' run: | @@ -89,7 +96,9 @@ jobs: name: bench-op (linux-arm64) needs: test runs-on: ubuntu-24.04-arm - timeout-minutes: 30 + # 45, not 30: the bidirectional sweep now compiles the torch.compile baseline + # up to L=2048 (~12 min of inductor), on top of the torch install and build. + timeout-minutes: 45 steps: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable @@ -104,6 +113,57 @@ jobs: - name: Install torch (CPU aarch64) working-directory: . run: python3 -m pip install --quiet numpy torch + # Fast consolidated sanity: both topologies run, are correct, and beat + # native torch — one clear pass/fail before the fuller benchmarks. + - name: Smoke check (unidirectional + bidirectional vs torch) + working-directory: . + run: python3 bench/smoke_topologies.py + # Correctness gates speed (CLAUDE.md): the bidirectional wrapper is + # checked against the vendored f64 reference through the real kernel + # before anything here is benchmarked. + - name: Bidirectional correctness through the kernel + working-directory: . + run: python3 tests/check_bidirectional.py + # Bidirectional scan vs stock PyTorch (eager + torch.compile), plus the + # internal fused-vs-flip diagnostic. Provisional (shared runner); headline + # numbers need a dedicated Arm host per bench/README.md. + # Full sweep-len (L = 128 .. 8192) rather than --quick's 128/512, because + # the interesting behaviour is at long L: torch.compile scales linearly in + # L while the kernel scales sub-linearly, so the gap WIDENS (3.63x at + # L=128 -> 4.67x at L=512). Two short shapes could not show that. + # + # Compilation is the entire cost (~L^0.6: 59s@128, 134s@512), so it is + # capped at 2048 — roughly 12 min. L=4096/8192 still get measured against + # eager and against the kernel, just without a torch.compile row. Raising + # the cap risks inductor OOMing on the unrolled graph; that experiment + # belongs in the dispatchable bench-baseline workflow, not on every push. + - name: Bidirectional benchmark (vs eager / torch.compile) + working-directory: . + timeout-minutes: 25 + run: | + set -o pipefail + python3 bench/bench_bidirectional.py --suite sweep-len \ + --compile-max-len 2048 --reps 5 --warmup 1 \ + --json bench-bidi.json 2>&1 | tee bench-bidi.log + - name: Publish bidirectional result + if: always() + working-directory: . + run: | + { + echo "## Bidirectional scan (linux-arm64 CI runner)" + echo "vs-baseline rows are the result. The internal fusion win is small and is NOT a headline — see BIDIRECTIONAL_LOG.md." + echo '```' + grep -E "^===|fused ==|^ (ref_|scan_fwd|bidirectional)|^ =>|exp-sharing|^bidirectional \(fused\)|torch.compile COST|^ +[0-9]+ +[0-9.]+s|^ +L|^ L=" bench-bidi.log 2>/dev/null || echo "(bench did not run)" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + # The bench flushes JSON after every shape, so this survives even if a + # long-L shape is killed mid-sweep. + - uses: actions/upload-artifact@v4 + if: always() + with: + name: bidirectional-bench-arm64 + path: bench-bidi.json + if-no-files-found: warn - name: Op-level benchmark (kernel vs eager vs torch.compile) working-directory: . run: | diff --git a/BIDIRECTIONAL_LOG.md b/BIDIRECTIONAL_LOG.md new file mode 100644 index 0000000..3af0c67 --- /dev/null +++ b/BIDIRECTIONAL_LOG.md @@ -0,0 +1,1102 @@ +# BIDIRECTIONAL_LOG — the 1D bidirectional scan, step by step + +A running record of the **1D bidirectional** topology: what changed, what it is +verified against, what broke along the way, and what is still unproven. Sibling +to [`OPTIMIZATION_LOG.md`](./OPTIMIZATION_LOG.md) (which tracks kernel *speed*); +this one tracks one axis of kernel *generality*. The 2D cross-scan (SS2D) gets +its own log when that work starts. + +Plan: [`TOPOLOGY_IMPLEMENTATION_PLAN.md`](./TOPOLOGY_IMPLEMENTATION_PLAN.md) §2 +(SS2D is §3 of the same plan). Every entry obeys the +[`CLAUDE.md`](./CLAUDE.md) rule **correctness gates speed** — nothing is +benchmarked, and no fusion work starts, until the correctness path is green. + +**Convention used throughout:** the topology ships in two stages — *correct* +(Python rearrangement on top of the existing 1D op, zero new Rust) and then +*fast* (fused in Rust, via a kernel `reverse` flag). The correct stage lands +first so the fusion is justified by a measurement instead of an assumption. + +> ## ✅ Bottom line (all CI green; fused sweep `c5deb72`) +> +> **A bidirectional selective scan on Arm runs 6.2–9.5× faster than +> `torch.compile` and 28–43× faster than PyTorch eager**, holding 4.3e-6 – 8.6e-6 +> max abs error vs the f64 reference (gate: 1e-4) all the way out to **L = 8192**. +> Those are the *fused* numbers — `bidirectional_scan` shares the exp between +> directions (Step 7), which is where the ~1.7× jump over the earlier two-call +> path came from. **A full bidirectional scan now costs ~1.15× a single +> unidirectional one.** +> +> **And `torch.compile` cannot follow us to the lengths the applications need.** +> Compile time is linear-through-2048 then **super-linear** (~0.26 s/timestep, +> rising to 0.31 at L=4096) because the recurrence is unrolled into an L-step graph. +> L=4096 already takes **21 minutes** to compile; at **L=8192 inductor OOM-kills the +> runner outright** (Step 6). The 131k-token genomics context is unreachable — not +> a slow baseline, an absent one. +> +> The fused two-direction kernel (Step 7) is also the **SS2D substrate** — its +> "compute Pass A once, run the recurrence in N directions" structure is exactly +> what the 2D cross-scan's four directions need. + +### What the baseline is (so the numbers are not misread) + +The reference these speedups are measured against is `selective_scan_ref`, +**vendored verbatim from the official `state-spaces/mamba` repo** (Tri Dao / +Albert Gu, Apache-2.0) — it matches HF transformers' `MambaMixer.slow_forward` +bit-for-bit. So this is our kernel vs **the original Mamba selective scan**, in +native PyTorch. That reference is a pure Python-loop implementation, which is +exactly what runs on CPU: mamba-ssm's CUDA kernel is GPU-only and does not run on +CPU at all. It is the genuine CPU baseline, not a strawman — and `torch.compile` +of the same reference (the stronger baseline) is measured alongside it. + +PyTorch has **no native bidirectional Mamba op**. Bidirectional Mamba (Vim, +Caduceus, VMamba) is always built by running the scan twice — forward and on the +flipped sequence — and summing; that is the standard construction, and it is what +the bidirectional baseline here does. So "bidirectional" is not something invented +for this project; only the *fused kernel* that computes both directions sharing +one exp is ours. + +**These are op-level numbers** — the selective scan, the part this kernel +replaces. A full Mamba *model* also has GEMMs and a conv the kernel does not +touch, so end-to-end `generate()` is more modest (~1.2–3×, `bench_e2e.py`). The +28–43× / 6–9.5× figures are for the scan op; never quote them as a whole-model +speedup. + +--- + +## Step 1 — correctness path (plan §2.1) + +**Branch:** `feature/bidirectional-scan` · **Status:** ✅ **done and verified.** +Both gates green on CI. The kernel gate failed once on first run — a bug in the +test, not the kernel (see *Errors and surprises* §6) — and passed after the fix. + +Full CI result (run `29323871999`, commit `bbc2722`): + +| Job | Gate | Result | +|---|---|---| +| `test (linux-arm64)` | definition check (numpy) | ✅ | +| `test (macos-arm64)` | definition check (numpy) | ✅ | +| `test (linux-x86_64)` | definition check (numpy) | ✅ | +| `bench-op (linux-arm64)` | **correctness through the real kernel** | ✅ | + +The `bench-op` row is the meaningful one: `bidirectional.py` running against the +real compiled NEON kernel on real Arm silicon, matching an f64 reference within +1e-4 across all 13 cases — including the NEON tail path (`state=13`), grouped +B/C, `L=1`, every merge mode, untied `reverse_params`, and `fwd == plain 1D +scan` (bit-identical). Wheels still build and golden-check on all three +platforms; op-level bench unregressed. + +### What + +`python/arm_scan/bidirectional.py` — `bidirectional_scan(...)`: runs the +recurrence over the sequence in both time directions and merges. Built entirely +on the existing 1D op — **no new Rust, no new FFI, no ABI bump**. + +- Merge modes: `sum` (default), `mean`, `concat`, `none` (returns both + directions unmerged, for a model-specific gated combine — the primitive + deliberately does not guess at a learned merge). +- `reverse_params`: override A/D/delta/delta_bias/B/C for the backward pass, for + untied models (Vim's `bimamba_type="v2"` has a separate `A_b`, `D_b`, + `dt_proj_b`). Default `None` = weight-tied. +- `_scan_reverse` is deliberately **the single seam**: today it flips the + time-varying inputs, calls the forward kernel, and flips the output back. + When the kernel grows the `reverse` flag (plan §2.2), only that function's + body changes — every caller inherits the win. +- Exposed as `arm_scan.bidirectional_scan` via the existing lazy-import + mechanism, so numpy-only users still never import torch. + +### Verification + +Two gates, deliberately split by what they can prove: + +| Gate | What it proves | Where it runs | +|---|---|---| +| `tests/check_bidirectional_math.py` | the **definition** is right | numpy only — **runs anywhere**, incl. this x86 Windows box | +| `tests/check_bidirectional.py` | the **code** is right, through the real kernel | needs torch + built cdylib → CI (`bench-op`, linux-arm64) | + +The math gate is the load-bearing one, and it is green. It proves: + +> flip → forward scan → flip **==** an explicitly-coded backward-in-time +> recurrence + +**bit-identically** (not merely within tolerance) across 7 shapes — including +`state=13` (the NEON non-multiple-of-4 tail path), grouped B/C, `L=1`, and a +128-step sequence. It is written as an independent reverse-time loop sharing no +mechanism with the flip-based path, on top of the already-independent +`naive_scan_f64` from `verify_golden.py`. + +It also guards against a **vacuous pass** — asserting the backward scan actually +differs from the forward one, so a no-op "backward" could not slip through every +equivalence check. + +Both gates are wired into `.github/workflows/ci.yml`: the numpy one into the +`test` job (all three platforms — it needs no torch), the kernel one into +`bench-op` **before** the benchmark, per *correctness gates speed*. + +### This file is also the spec for the Rust `reverse` flag + +`naive_scan_backward_f64` in the math gate is exactly what `reverse=true` must +compute. Plan §2.2 now has an executable definition to implement against rather +than an assumption to re-derive — and the bit-identical result means the fused +path has no numerical excuse to differ. + +### Errors and surprises encountered + +**1. The f64 ground truth was silently f32.** ⚠ *Caught before it could mislead.* +The first draft of `check_bidirectional.py` passed **f32 inputs** with +`compute_dtype=torch.float64`. But the vendored `selective_scan_ref` ends with +`out = out.to(dtype=dtype_in)` — it casts back to the *input* dtype. So the +"float64 reference" would have come back as f32, and the 1e-4 gate would have +been comparing kernel-f32 against reference-f32: a far weaker check than +claimed, and one that could mask real error. `gen_golden.py` upcasts the inputs +(`f64 = lambda t: t.double()`) for exactly this reason. **Fixed** by upcasting +inputs to double, mirroring `gen_golden.py`. + +**2. `tests/reference/` is a package, not a bare module.** The check originally +put `tests/reference/` on `sys.path` and did +`from selective_scan_ref import selective_scan_ref`. The repo convention (per +`gen_golden.py`) is to put `tests/` on the path and do +`from reference import selective_scan_ref`. **Fixed** to match. + +**3. `D` is applied twice under a sum merge.** Not a bug — a real gotcha. With +`merge="sum"` and a shared `D`, the skip connection lands in **both** directions, +so the merged output carries `2·D·u`, not `D·u`. This *is* what real +bidirectional Mambas do (each direction's mixer applies its own D, then the +outputs are summed), but it is exactly the kind of thing that produces +plausible-looking, quietly-wrong output. Now **documented in the module and +pinned by an assertion** so it can never drift silently. + +**4. Gating inside both passes is safe — proven, not assumed.** The module lets +the kernel apply the z-gate in *both* directions rather than once after the +merge. For any linear merge that is algebraically identical: + +``` +inside: (y_f + D·u)·silu(z) + (y_b + D·u)·silu(z) +outside: ((y_f + D·u) + (y_b + D·u))·silu(z) +``` + +Asserted in the math gate rather than left as a claim in a docstring. + +**5. The dev box could not run *anything*.** No Rust, no torch, no numpy, no +built cdylib. Rather than block on a multi-GB toolchain install, the work was +split so the *definitional* correctness could be proven with numpy alone (a +`.venv` at the repo root, which `bench/run_baseline.sh` already expects), and +the kernel-level check deferred to CI. This is why there are two gates and not +one — and it turned out to be a better structure anyway, since the math gate is +portable and catches the class of bug that kernel testing structurally cannot. + +**6. The kernel gate failed on first CI run — and it was right to.** ⚠ *The most +instructive failure so far.* `no_softplus` came back at **max_abs = 3.3e-3**, +33× over the 1e-4 gate, while all 12 other cases passed. + +Root cause was **the test, not the kernel**. `make_case` drew `delta` from a +normal distribution unconditionally. That is fine when `delta_softplus=True` +(delta is raw; the kernel applies softplus and the timestep comes out positive), +but with `delta_softplus=False` **delta *is* the timestep** and must already be +positive — HF's slow path pre-applies softplus, so no real Mamba ever passes a +negative one. `gen_golden.py` knows this and draws `uniform(1e-3, 0.1)` for its +own `no_softplus` case; my test drew `randn` and produced negative timesteps. + +Why that blows up *now* specifically: with `delta < 0` and `A < 0`, the argument +`dt·A` goes **positive** — violating the precondition of the `vexpq_f32_nonpos` +optimization landed in [`OPTIMIZATION_LOG.md`](./OPTIMIZATION_LOG.md) Step 1, +whose docstring states it outright ("*positive input can overflow to inf because +the upper clamp is absent — the scan never passes positive*"). The 3.3e-3 is a +degree-reduced, unclamped exp being evaluated outside its fitted domain, exactly +as designed. + +**Fixed** by drawing a positive delta when softplus is off (mirroring +`gen_golden.py`), and by making `make_case` *reject* `delta_bias` in that mode — +a bias could push a positive timestep negative and reintroduce the same +violation silently. + +Two things worth taking from this: +- The precondition behind the exp optimization is **real and load-bearing**, and + it now has a second, independent test exercising it. An unrelated workstream + tripped over it within a day of it landing. +- `check_bidirectional_math.py` deliberately *keeps* drawing `randn` delta for + its no_softplus case and still passes — numpy's exp is exact over the whole + line, and the identity it tests (flip-forward-flip == backward recurrence) is + a mathematical fact that holds for any delta. The two files differ on purpose; + a comment in each says so, so nobody "fixes" one to match the other. + +**Result after the fix:** 12 of 13 cases were already green on the first run, +including `state13_neon_tail`, `grouped_bc`, `edge_len1`, all merge modes, the +untied-`reverse_params` path, and `fwd == plain 1D scan` (bit-identical). Only +`no_softplus` failed, and only because it was asking the kernel a question no +model asks. + +### Design boundary — which bidirectional models this is for + +Two patterns exist in the wild and they are **not interchangeable**: + +- **"outer"** (Caduceus's `BiMambaWrapper`, Vim): the *whole mixer* — causal + conv, `x_proj`, `dt_proj` — is re-run on `x.flip(time)`. A causal conv over + flipped input is **not** the flip of the conv over input, so the two + directions' scan inputs are genuinely different tensors. Such a model does not + need this module: it already calls `selective_scan` twice, and **both calls are + ordinary forward scans**. A kernel `reverse` flag buys it *nothing*. +- **"inner"** (VMamba/SS2D-style cross-scan, and bidirectional variants that flip + *after* the projections): the *same* projected tensors are traversed in both + time directions, because flipping commutes with the time-pointwise + projections. **This is what the module implements**, it is what a fused + `reverse` flag actually accelerates, and it is the 1D case of the SS2D + cross-scan. + +This distinction cannot be settled until the application is chosen +(`APPLICATIONS.md` is still open) — it decides whether §2.2's fused `reverse` +flag is worth building at all for the app we ship. Documented at the top of the +module so nobody wires it into an outer-pattern model by mistake. + +### Not yet done + +- **Fast** is not done. *Correct* is. Each direction runs on the full fast NEON + kernel, but the six flip copies around them are pure overhead that only the + fused `reverse` flag (Step 3) removes. **No Rust was written for this step.** +- No HF integration (`patch.py` dispatch for a bidirectional mixer class) — + blocked on the application decision. + +--- + +## Step 2 — measure what the flips cost (the gate on Step 3) + +**Status:** ✅ **measured. The answer is: the flips are noise. Do not fuse.** + +Benchmark: `bench/bench_bidirectional.py`, wired into CI's `bench-op` job. + +### The numbers (linux-arm64 CI runner, 4-core, torch 2.13.0, `--quick`, commit `5cb4fe8`) + +| Shape | `scan_fwd` | `fused_estimate` | `bidirectional` | `flips_only` | headroom (ceiling) | +|---|---|---|---|---|---| +| B1 D768 L128 N16 | 0.801 ms | 1.641 ms | 1.781 ms | 0.057 ms | **1.085×** | +| B1 D768 L512 N16 | 2.694 ms | 5.492 ms | 5.631 ms | 0.108 ms | **1.025×** | + +**Sanity check first:** `fused_estimate` ≈ 2 × `scan_fwd` at both shapes (1.641 +vs 1.602; 5.492 vs 5.388), which is exactly what the proxy is supposed to do. The +comparison is sound. + +### Verdict: **Step 3 is rejected on measurement.** + +Three independent reasons, in order of weight: + +**1. The ceiling is already low, and it *shrinks* with L.** 8.5% at L=128 → +2.5% at L=512. Not noise: flip traffic is O(B·D·L) memory copies while scan work +is O(B·D·L·N) compute, so the flips get relatively cheaper the longer the +sequence. Every application this topology targets (genomics at **131k tokens**, +long audio, multi-hour ECG) lives far to the right of L=512, where the flips are +negligible. The measurement gets *more* damning at the shapes we actually care +about, not less. + +**2. Most of the apparent overhead is not even flips.** The two measurements +disagree — subtraction says 7.8% at L=128, but `flips_only` directly measures +**3.2%**. That gap (~0.14 ms, roughly *constant* across both shapes) is +Python-side wrapper overhead in `bidirectional_scan`, not memory traffic, and a +Rust `reverse` flag **would not remove it**. Having both measurements is what +caught this; a subtraction-only benchmark would have overstated the case for +fusion by ~2.5×. The honest ceiling on what fusion saves is **1.9–3.2%**. + +**3. It may help zero models anyway.** A `reverse` flag only accelerates *inner* +bidirectional models. If the application lands on Caduceus or Vim (*outer* — see +the design boundary above), both of their scans are ordinary forward scans and +the flag accelerates **nothing at all**. + +So the trade is: half a day of NEON chunk-reversal surgery + an FFI ABI bump + +new goldens + new parity tests, to win ~2% that trends to ~0% at the sequence +lengths that matter, for a model class that may not even benefit. **No.** + +### Why this is a good outcome, not a wasted step + +This is the **second** time measuring-before-optimizing has killed a plausible +idea in this project. The first was the plane transpose that everyone assumed was +a top-priority cost and profiled at **0.1%** +([`OPTIMIZATION_LOG.md`](./OPTIMIZATION_LOG.md)). Had §2.2 been built on +intuition — and the plan's own effort estimate said "half a day," which is +exactly the size of task that gets waved through — it would have bought ~2%. + +Per [`IMPROVEMENT_IDEAS.md`](./IMPROVEMENT_IDEAS.md) §7.1, a **considered-and- +rejected decision backed by numbers** is worth publishing. This one goes in the +writeup: *we built the general path, measured what fusing it would buy, found it +was ~2%, and spent the time elsewhere.* That is a stronger technical signal than +a fused kernel nobody needed. + +### Caveats on these numbers (stated, not buried) + +- `--quick` mode: **reps=5, warmup=1**, on a *shared* 4-core CI runner. Noisy. + The 1.085 vs 1.025 spread is within plausible run-to-run variation; the + *trend* and the *magnitude* are what carry the conclusion, not the third + decimal. +- Only two shapes, both short (L=128, L=512). The full `sweep-len` suite goes to + L=8192 — worth one run on a dedicated Arm host before the number is quoted in + `RESULTS.md`, since it would show the trend continuing rather than asserting it. +- `fusion_headroom` is a **ceiling**, not an achieved speedup (see the module + docstring). A real fused kernel reads the sequence backward, which is less + cache-friendly than the forward stream timed here, so it would land *under* + these figures. Quoting it as a speedup would be dishonest. + +Step 3 (the fused kernel) is deliberately **gated on this measurement rather +than scheduled**. Its entire value is deleting the six copies Step 1 pays for — +so if those copies are cheap, the fused kernel is worthless and should not be +built. The kernel-side profiling work already produced the cautionary example: +the plane transpose that everyone *assumed* was a top-priority cost measured +**0.1%** ([`OPTIMIZATION_LOG.md`](./OPTIMIZATION_LOG.md)). Assuming here would be +repeating a mistake this project has already made once and caught. + +### What it times + +| Series | What it is | +|---|---| +| `scan_fwd` | one forward scan — the floor | +| `fused_estimate` | two forward scans + merge, **zero flips** — the proxy for a fused `reverse` | +| `bidirectional` | the real thing today: flips + two scans + un-flip + merge | +| `flips_only` | the six copies alone, no scans — reads the cost directly instead of by subtraction | + +**The decision number:** `fusion_headroom = bidirectional / fused_estimate`. + +- `~1.0×` → the flips are noise. **Do not build Step 3.** Ship §2.1 and say so + plainly in the writeup. +- `>1.15×` → the flips are real and the fused flag pays for itself. + +### The proxy is a ceiling, and is reported as one + +`fused_estimate` is an **upper bound** on what fusion could achieve, not a +measurement of a fused path (which does not exist). It runs identical scan work +with zero flip traffic, so a real fused kernel *cannot beat it* — a real one +also walks the sequence backward, which is less cache-friendly than the forward +stream being timed. So: a low ceiling is conclusive (don't fuse); a high ceiling +is permission to try, not a promise. It is never to be quoted as an achieved +speedup. + +--- + +## Step 3 — fused `reverse` flag in Rust (plan §2.2) + +**Status:** ✅ built. **But note the reversal of the Step-2 decision, and why.** + +Step 2 rejected this **as a speedup**, and that rejection still stands: the +copies are worth ~2%, and nothing here changes that. It was built anyway, on a +different justification that Step 2 did not weigh: + +> **The 2D cross-scan needs a backward traversal regardless.** SS2D's four +> directions are row-forward, row-**backward**, column-forward, and +> column-**backward**. Plan §3.2's design has the row directions reusing the 1D +> scan directly — which requires exactly this flag. So `reverse` is not a +> bidirectional optimization that failed to pay off; it is **the substrate SS2D +> is built on**, and it happens to also remove bidirectional's flip copies. + +**The claim to make when this ships is therefore:** *"a fused backward traversal +— the substrate for the 2D cross-scan, which also removes the flip copies from +bidirectional (~2%)."* **Not** *"we made bidirectional faster."* The benchmark +prints that caveat on every run so nobody quotes it wrong. + +### What changed + +The implementation turned out **much cheaper than plan §2.2 predicted**, and for +an instructive reason. The plan claimed the NEON work needed SIMD lane-reversal +(`vrev64q_f32` / `vextq_f32` shuffles inside each chunk). That was wrong: + +- **Pass A is pointwise in time** — no cross-timestep dependency. That is the + entire reason the two-pass split exists. It needs *zero* direction awareness. +- **Pass B vectorizes across STATE, not time.** `h` lives in four q-registers and + `t` is a plain scalar loop index. So reversing time is one subtraction + (`t = tlen - 1 - i`) — no shuffles, no extra loads, no extra work. +- The B/C plane transpose, the epilogue, and `parallel.rs` are all + direction-agnostic and were **not touched at all**. + +Net: a new `chunks_in_scan_order()` iterator (visit chunks last-first when +reversed), one flipped index in each of the two NEON channel paths and the +scalar path, and the FFI/Python plumbing. The recurrence math is untouched. + +| Layer | Change | +|---|---| +| `arm-scan-core/src/lib.rs` | `ScanInput::reverse` | +| `src/scalar.rs` | one flipped time index | +| `src/neon/mod.rs` | `chunks_in_scan_order()` + flipped `t` in Pass B (both paths) | +| `src/parallel.rs` | **none** — channel independence is direction-agnostic | +| `arm-scan-ffi/src/lib.rs` | `reverse` param; ABI bump (→ **4** after reconciling with `h0`, see Step 4) | +| `python/arm_scan/{_ffi,op,numpy_api}.py` | `reverse` threaded through | +| `python/arm_scan/bidirectional.py` | the seam closed: `reverse=True`, no flips | + +### Correctness + +`reverse=True` is *defined* as flip-forward-flip, and that definition is now +enforced at three independent levels: + +1. **Rust, bit-for-bit** — `reverse_matches_flip_forward_flip` (property test, + 256 random shapes/flag combos). Asserts **bit-identity**, not tolerance: both + paths apply the same arithmetic to the same values in the same order. Note the + two paths land on *different chunk boundaries* (forward-on-flipped splits the + flipped axis; reverse splits the original), so passing bit-exactly also proves + chunking never leaks into the math. +2. **Rust, anti-vacuous** — `reverse_actually_reverses` asserts a reversed scan + genuinely differs from a forward one, so a dropped flag cannot pass silently. +3. **numpy, independent** — `tests/check_bidirectional_math.py` proves the + identity itself is sound against a separately-written backward recurrence, + with no kernel involved. + +Plus: `reverse` is now generated by the proptest `Case`, so **every** existing +property test (f64 agreement, NEON-vs-scalar parity, rayon bit-identity) sweeps +both directions for free. And `ffi_reverse_two_steps` hand-checks it across the +C ABI. + +The benchmark refuses to report a speedup unless the fused output is +bit-identical to the flip-based one — a fast wrong answer is not a result. + +### Still worth doing, and still not this + +Step 2's measurement found ~0.14 ms of **Python-side wrapper overhead**, roughly +constant across shapes and *larger than the flip copies themselves* at L=512. +`reverse` does not touch it. It belongs with the trims in +[`IMPROVEMENT_IDEAS.md`](./IMPROVEMENT_IDEAS.md) §8 (ctypes call path, `_c()`'s +redundant `.float()` dispatch, per-call allocations) — cheap, and it helps +*every* caller rather than just bidirectional ones. + +--- + +## Step 4 — the `reverse` / `h0` ABI collision (merging with main) + +**Status:** resolved. Both features shipped. **ABI is now 4.** + +While `reverse` was in flight, [`OPTIMIZATION_LOG.md`](./OPTIMIZATION_LOG.md)'s +workstream landed **`h0`** on main — a caller-supplied *initial SSM state* that +makes the scan resumable (run a prefix, feed its `last_state` back as `h0`, +continue). That is `IMPROVEMENT_IDEAS.md` §2.4/§7.6, the decode/streaming work. + +**Both branches independently bumped the ABI 2 → 3, with different signatures.** + +| | `h0` (main) | `reverse` (this branch) | +|---|---|---| +| Core API | new `selective_scan_with_state(...)`, `h0` as a **parameter** | `reverse` as a **field on `ScanInput`** | +| C ABI | `h0: *const f32` appended after `last_state` | `reverse: c_int` inserted after `delta_softplus` | +| Claimed ABI | **3** | **3** ← collision | + +Reconciled to **ABI 4**, carrying both. They are genuinely orthogonal — `h0` +seeds the state, `reverse` picks the traversal direction — and they compose: a +*backward* scan resumed from a prior state is coherent, and is in fact what SS2D +will want for its column traversals. No redesign was needed, only plumbing. + +### The dangerous part: what git merged *without* a conflict + +Only 4 files conflicted (`scalar.rs`, `_ffi.py`, `numpy_api.py`, `op.py`) — the +ordinary "both added a parameter" kind, resolved by keeping both. **The damage +was in the files that auto-merged cleanly**, and it is worth internalizing: + +**1. The ABI version silently stayed at 3.** Both branches wrote `3`, so git saw +identical text and merged happily — while the Python loader had been reconciled +to expect 4. A version check that exists precisely to catch ABI drift would +itself have been the thing that was wrong. Caught by reading the merged file +rather than trusting the absence of a conflict marker. + +**2. Adding a struct field breaks the *other* branch's new construction sites — +with no conflict.** `ScanInput` gained `reverse` here; main added new +`ScanInput` literals in its streaming tests. Git merged both hunks cleanly and +produced code that does not compile. Same for the C ABI: main's three `h0` tests +were written against a signature with no `reverse`, and this branch's `reverse` +test against one with no `h0` — every one of those call sites was left short an +argument. + +This bit **twice**: the FFI call sites were caught before pushing (by counting +arguments at each call), but property.rs's three new `ScanInput` literals were +not — because only the *conflicted* files were re-checked, and property.rs was +not one of them. CI's `cargo clippy --all-targets` caught them (note: `--all-targets` +is why clippy, not `cargo build`, was the step that failed — it is the only gate +that compiles the test targets). + +**The rule to take from this:** after a merge that adds a field to a shared +struct or a parameter to a shared signature, *enumerate every construction and +call site in the whole workspace* — do not assume the conflicted files are the +complete set of affected ones. The absence of a conflict marker means git found +no textual overlap, not that the result is coherent. + +### The bit-identity assertion was wrong — and finding out was worth it + +⚠ **The most interesting numerics finding of this whole workstream.** + +`reverse_matches_flip_forward_flip` originally asserted that a fused reverse scan +is **bit-identical** to flip-forward-flip. It passed on x86 (scalar) and **failed +on both Arm legs** (`last_state differs`, at `dims = {dim: 3, len: 31, state: 7}`). + +**The kernel was right. The assertion was false.** Two NEON passes process four +timesteps at a time with a **scalar tail**: + +```rust +while t + 4 <= tlen { ... vsoftplusq_f32 (NEON polynomial) ... } +while t < tlen { ... dt.softplus() (libm) ... } +``` + +`discretize_chunk` (softplus) and `epilogue_row` (SiLU) both look like this. The +vector and tail branches compute the *same function by different means* — they +agree to ~1e-7, but not to the last bit. **Which branch a timestep takes depends +on its array POSITION**, and flipping the array moves timesteps across that +boundary. + +At `len = 31`: the vector body covers positions 0–27, the tail 28–30. Scanned in +place, timestep 29 sits at position 29 → **scalar tail**. Scanned flipped, it +lands at position 1 → **vector body**. Same timestep, same value, ~1 ulp apart. +That propagates through the recurrence and shows up in `last_state`. + +The scalar backend has one uniform code path per timestep, so it *is* bit-exact +— which is precisely why x86 passed and Arm did not. + +**This is a property of the pre-existing forward kernel, not of `reverse`.** +Demanding bit-equality on NEON was asserting something false. + +**Resolution — keep the strongest claim that is actually true, per backend:** + +| Backend | Assertion | Why | +|---|---|---| +| `Scalar` | **bit-identical** | uniform per-timestep path; any difference is an indexing bug, not rounding — this is what pins the traversal down | +| `Auto` (NEON) | scale-relative `< 1e-5` | one SIMD-vs-libm transcendental apart, matching `auto_backend_matches_scalar`'s existing bar | + +The scalar leg still delivers the guarantee that matters — that the *indexing* is +exactly right, including that chunk boundaries never leak into the math (the two +routes chunk the axis differently). The NEON leg confirms the fused traversal is +numerically sound without pretending to an equality that cannot hold. + +The same over-strong gate was in `bench/bench_bidirectional.py` (`torch.equal`). +It would have passed anyway — every benchmarked length is a multiple of 4, so no +scalar tail exists — but that is luck, not correctness, and a shape-dependent +gate that silently holds is worse than one that states its tolerance. Relaxed to +the same scale-relative bar, and it still reports when the result *is* bit-exact. + +**The lesson:** "bit-identical" is the right bar for a *reordering* of identical +arithmetic, and the wrong bar the moment a SIMD tail means the arithmetic is not +identical. The test was correct to fail; it caught an over-claim in the +documentation before it reached a judge. + +### Verification after the merge + +- Every `ScanInput` and `Channel` literal in the workspace sets `reverse` + (swept exhaustively, not spot-checked). +- All 7 C-ABI call sites pass exactly 16 arguments. +- `try_neon` carries **both** `reverse` and `h0` into the NEON path — a drop + there would silently ignore one feature on aarch64 only, which is the worst + possible failure mode (correct on the x86 CI leg, wrong on the target). +- `op.py`'s custom op, its `register_fake`, and the FFI call all agree on + argument order (`…, delta_bias, h0, delta_softplus, reverse`). A mismatch here + breaks `torch.compile` composability, which is the entire point of the fake + kernel. +- `profile.rs` needs no `h0` plumbing (it is a zero-initialized diagnostic) but + does honor `reverse` — otherwise `scan_profiled(reverse: true)` would have + silently profiled a *forward* scan, which is worse than a compile error. + +--- + +## Step 5 — Results (all CI green, commits `6e92b03` / `605b056`) + +Host: GitHub `ubuntu-24.04-arm` runner — 4-core aarch64, torch 2.13.0, `--quick` +(reps=5, warmup=1). **Provisional**: a shared runner. Headline figures still need +a dedicated Arm host. Two independent runs are reported, because the agreement +between them is what makes the numbers credible. + +### The result (canonical run `883943f` — sweep 128 → 8192, **reps=10, warmup=3**) + +Dispatched on `bench-baseline` (long-L job), the higher-quality settings than the +per-push `--quick` runs. This is the run to cite. + +| L | eager | torch.compile | **kernel** | vs eager | **vs torch.compile** | +|---|---|---|---|---|---| +| 128 | 27.25 ms | 6.35 ms | **1.68 ms** | 16.2× | **3.78×** | +| 512 | 131.16 ms | 26.03 ms | **5.50 ms** | 23.8× | **4.73×** | +| 1024 | 264.30 ms | 57.41 ms | **10.75 ms** | 24.6× | **5.34×** | +| 2048 | 575.51 ms | 119.89 ms | **21.41 ms** | 26.9× | **5.60×** | +| 4096 | 1138.52 ms | 234.27 ms | **42.51 ms** | 26.8× | **5.51×** | +| 8192 | 2441.15 ms | *(compile capped at 4096)* | **87.54 ms** | 27.9× | — | + +B=1, D=768, N=16. Correctness in the same run: kernel-vs-f64 max abs **4.3e-6 → +8.6e-6**, against the 1e-4 gate — and, importantly, **it does not drift with L**. +Accumulating 8192 sequential steps through a chunked scan with a degree-3 exp +could plausibly have degraded; it did not. All 13 cases of +`check_bidirectional.py` green. + +Earlier `--quick` runs (`3c7a7c3`, reps=5) gave 3.92× / 5.38× / 5.18× / 5.62× +at L=128/512/1024/2048 — consistent with this one to within shared-runner noise. +The canonical numbers above are from reps=10. + +### The ratio improves, then PLATEAUS — it does not keep growing + +**Correcting an earlier claim.** From two points (128 → 512) I concluded the +advantage "grows with sequence length." With six points it clearly **plateaus**: + +- **3.92×** at L=128, then **5.2–5.6×** from L=512 onward. + +The jump is real but it is a *one-off*, not a trend. L=128 is **depressed** by the +kernel's fixed per-call overhead (ctypes dispatch, `_c()` contiguity checks, +output allocation — the ~0.14 ms measured in Step 2), which is a meaningful +fraction of a 1.7 ms call and a negligible one of a 21 ms call. Once it amortizes +(by L≈512), both the kernel and `torch.compile` scale linearly in L, so the ratio +settles. + +**Honest headline: ~5.2–5.6× vs `torch.compile` at L ≥ 512; ~3.9× at very short L.** + +### ⚠ Compile time is LINEAR in L — the withdrawn claim, reinstated with data + +Step 5 previously **withdrew** the "compile time explodes" argument, on the basis +that 59 s → 134 s (L 128 → 512) is only 2.3× for 4× the length, i.e. sub-linear. +**That withdrawal was premature — it was a two-point fit.** With four points: + +| L | compile | **seconds per timestep** | +|---|---|---| +| 128 | 63.4 s | 0.495 | +| 512 | 136.7 s | 0.267 | +| 1024 | 251.0 s | 0.245 | +| 2048 | 533.5 s | **0.260** | + +The per-timestep cost holds near ~0.26 s from L=512 through 2048 — then **bends +upward**. The canonical run added the L=4096 point: + +| L | compile | s/timestep | +|---|---|---| +| 128 | 66.1 s | 0.517 | +| 512 | 135.4 s | 0.264 | +| 1024 | 251.6 s | 0.246 | +| 2048 | 544.0 s | 0.266 | +| **4096** | **1254.1 s** | **0.306** | + +So compile time is **linear-ish through 2048, then super-linear** — L=4096 came in +~18% *above* the linear extrapolation (1254 s vs ~1067 s). The graph is large +enough by 4096 that compilation itself is scaling worse than O(L). That makes the +projections *conservative*, not optimistic: + +| L | projected compile (≥ linear) | +|---|---| +| 8192 | **≥ 36 min** (measured to OOM instead — see Step 6) | +| 131,072 (genomics context) | **≥ 9.5 hours, and rising super-linearly** | + +…if it does not OOM first, which at 8192 it already does (Step 6). + +**Amortization is non-monotonic, and that is the super-linearity showing through.** +Iterations-to-break-even vs our kernel: 14,160 (128) → 6,595 (512) → 5,393 (1024) +→ 5,524 (2048) → **6,540 (4096)**. It bottoms out near L=1024–2048 and *rises* +again at 4096, because compile is now growing faster than linear while our +runtime stays linear. There is no sequence length at which `torch.compile` +becomes cheap; past ~2048 the trade only worsens. + +**This is the moat, stated precisely.** At the sequence lengths our headline +applications actually use, `torch.compile` is not a slow baseline — it is an +**absent** one. `CLAUDE.md` says the kernel's argument is that "`torch.compile` +cannot restructure a sequential recurrence"; the linear compile cost is that claim +made measurable, and it now rests on four points rather than an assertion. + +**Amortization is stable at ~5,450 iterations** for L ≥ 512 (compile and runtime +both scale linearly, so the ratio is constant). That is the number a skeptic +computes for themselves, so we publish it rather than let them derive it. + +### Methodological note: I over-extrapolated from two points, twice + +Both mistakes in this section came from the same error, in opposite directions: + +1. *"The advantage grows with L"* — from two points. It plateaus. +2. *"Compile cost does not explode"* — from two points. It is linear, and at + application scale that is fatal for the baseline. + +Each was corrected only by measuring more shapes. **Two points define a line; they +do not establish a trend.** The same lesson the fusion number taught (Step 5, +"two runs is not reproducibility") — recorded here because it cost real time twice +and would have put a wrong claim in front of a judge both times. + +**Why this number is believable:** the *forward* scan measures **3.74×** vs +`torch.compile` at the same shape ([`OPTIMIZATION_LOG.md`](./OPTIMIZATION_LOG.md)). +A bidirectional scan is two forward scans, so it should land in the same band — +and 3.93× does. That internal consistency is the real check; a bidirectional +figure wildly different from the forward one would have meant a broken comparison, +not a fast kernel. + +`torch.compile` also cost **62.4 s of one-time compilation** for a single shape, +and only ran at L=128 (`--quick` caps `compile_max_len` at 128 — graph unrolling +makes compile time explode with L). **One `torch.compile` data point is thin.** +The full `sweep-len` suite is needed before this goes in `RESULTS.md`. + +### `torch.compile`'s compile cost — measured, and it does NOT explode + +I claimed the recurrence's unrolled graph makes compile time "explode" with L. +**The data says otherwise, and the claim is withdrawn.** + +| L | compile | run/iter | iterations before compile pays for itself vs our kernel | +|---|---|---|---| +| 128 | 59.3 s | 6.13 ms | ~13,400 | +| 512 | 134.1 s | 25.99 ms | ~6,600 | + +**2.3× compile time for 4× the length — sub-linear.** Extrapolating, L=8192 would +be roughly ten minutes: painful, not prohibitive. And it amortizes after ~6–13k +calls, which any long-running server clears trivially. + +**So "compile cost" is a weak argument and this project should stop leaning on +it.** The strong argument is the one that survives the objection: we are +**3.6–4.7× faster than `torch.compile` after it has fully paid its compile cost**, +and the gap widens with L. Publish the compile cost honestly — including the +amortization column, which is the number a skeptic would compute themselves — and +let the steady-state ratio carry the case. + +**One thread left open** (assert nothing until it is checked): the baseline uses +`dynamic=False`, so **every distinct sequence length should trigger a recompile**. +Under variable-length inference — i.e. real serving — that 134 s would be paid per +shape, not once, which *would* make compile cost a serious argument. It is worth +confirming before it is claimed. + +### The fusion win: small, and NOT reproducible on a shared runner + +⚠ **Third revision of this number. Recording the whole arc, because the +flip-flopping is itself the lesson.** + +- **First read:** "it's noise" — the achieved speedup exceeded its own theoretical + ceiling, which is impossible. +- **Second read:** "no, the *ceiling* is the broken thing" — `fusion_speedup` + reproduced to ~0.001× across two runs while the ceiling swung wildly. That part + stands: `fused_estimate` (two forward scans, no flips) was supposed to be a + lower bound, the real fused path consistently beat it by 3–6%, and a bound the + real thing beats is a broken proxy. It has been **demoted to a diagnostic** and + the guard built on it removed. +- **Third run broke the rest of it:** + +| | run 1 | run 2 | run 3 (`23995c7`) | +|---|---|---|---| +| L=128 | 1.064× | 1.065× | **1.070×** | +| L=512 | 1.151× | 1.136× | **1.027×** | + +L=128 is genuinely stable (~7%, three runs). **L=512 swings 1.027–1.151×** — so +the "reproducible, and grows with L" conclusion was two runs of coincidence, and +is withdrawn. + +**Final position: the fused reverse is worth ~3–7%, it is not reliably measurable +on a shared 4-core runner, and pinning it down is not worth another CI cycle.** It +was never the justification: `reverse` exists because **SS2D needs a backward +traversal** (plan §3.2's row-backward and column-backward directions). The number +to quote is 5.3–5.6× vs `torch.compile` (L ≥ 1024). Nothing else. + +**Two lessons worth keeping:** +1. A proxy is only trustworthy until the real thing exists. If the real thing + outruns its own "lower bound," suspect the proxy before you blame the noise. +2. **Two runs is not reproducibility.** Two agreeing measurements produced a + confident, mechanistic story ("grows with L, because of cold working sets") + that the third run dismantled. The effect being chased (~3–15%) was simply + smaller than the shared runner's variance the whole time. + +### What the numbers confirmed along the way + +- **The forward path did not regress.** `fwd == plain 1D scan` is bit-identical, + and all 16 goldens hold at their recorded error floors — so the loop-invariant + `if ch.reverse` branch added to Pass B did not disturb the existing kernel. + (LLVM presumably unswitched it, as expected. The criterion ladder is the direct + confirmation and is still worth reading.) +- **The SIMD-tail prediction was right.** `fused == flip-based` came out + **bit-identical** at both shapes — exactly as predicted, because L=128 and + L=512 are multiples of 4, so no scalar tail exists to diverge (see Step 4's + bit-identity finding). At a length like 31 it would not have been, which is + precisely why the gate is a tolerance and not `torch.equal`. +- **`flips_only` (0.064 / 0.110 ms) is far smaller than the flip path's total + penalty** (0.108 / 0.882 ms over the fused path). The gap is not copy cost — it + is the *second working set*: flipped tensors are freshly allocated, so the scan + streams ~4.7 MB of cold memory instead of re-reading warm cache. A real effect, + and another reason the naive "flips cost ~2%" framing understated things. Also + another reason not to trust ±20%-noise numbers to arbitrate it. + +### Next, to make this publishable + +1. Full `sweep-len` (not `--quick`), reps ≥ 10, on a **dedicated** Arm host — + Oracle Ampere A1 or a short Graviton session — so `torch.compile` is measured + at more than one shape and the noise floor drops below the effect. +2. Tighten the benchmark's ceiling guard to a **per-shape** check so an inversion + like L=128's is a hard error, not a silent oddity. +3. Then, and only then, put the vs-`torch.compile` row in `RESULTS.md`. + +--- + +## Step 6 — `torch.compile` OOM-kills the runner at L=8192 (the hard wall) + +**The linear-compile-cost finding (Step 5) has an endpoint, and we hit it.** + +Dispatched the long-L job (`bench-baseline.yml`, run `29350634837`) with +`long_l_compile_max=8192`. After ~96 minutes it died with: + +> *The hosted runner lost communication with the server. Anything in your +> workflow that terminates the runner process, starves it for CPU/Memory, or +> blocks its network access can cause this error.* + +That is an **OOM**, not a timeout (the job timeout is 150 min and was not +reached; a timeout reports cleanly as one). `torch.compile` unrolls the L-step +recurrence into a single graph, and at L=8192 the graph exhausts the 4-core +arm64 runner's memory; the OOM reaper kills the VM. The ~96 min (vs a ~75 min +estimate) is consistent with the machine thrashing on swap before the kill. + +**Where it died:** cumulative compile to L=4096 is ~34 min; L=8192 adds ~36 more. +At 96 min it was well past 4096 and inside the 8192 compile. So 4096 *did* +compile — but see below. + +**This is the strongest form of the moat argument.** Not "`torch.compile` is +slower" — at the sequence lengths the headline applications use (genomics @131k, +long audio, multi-hour ECG), `torch.compile` **cannot build the graph at all**. +The kernel ran L=8192 in **89.7 ms** at 8.6e-6 error (Step 5), in constant +memory, on the same box that could not compile the baseline. + +### Operational failure: we lost the data, and why + +**No artifact was produced.** The per-shape JSON flush (added precisely to +survive a mid-sweep death) protects against a *process* kill — a `SIGKILL` of the +python process, after which later steps still run. It does **nothing** against a +*VM* death: when the machine itself is OOM-killed, there is no runner left to +execute the `if: always()` artifact-upload step. Both the flushed JSON on the +dead VM's disk and the upload step went down with it. + +So the 4096 compile row — which had almost certainly completed — was lost. The +lesson: **`if: always()` is not a guarantee against OOM; it is a guarantee +against a failed *step*.** Different failure modes need different safety nets, and +"the whole machine dies" has essentially none within a single job. + +### The fix, and the plan + +To bank the L=4096 compile row cleanly, cap compile at **4096** (which fits the +runner's memory) so the job completes and uploads. The L=8192 OOM is already +established qualitatively by this run; it does not need re-running to be cited — +though a dedicated host with more RAM would let us find the *exact* L at which +the graph stops fitting, which is a sharper number than "≥ 8192 on 16 GB." + +--- + +## Step 7 — fused two-direction kernel: share the exp ✅ DONE + +**Status:** all three stages green and **measured**. The bet paid off — the +exp-sharing speedup is **1.58–1.75× (projected ~1.7×)**, stable across every +shape, bit-identical to two scans, on real Arm. + +> ### Result (canonical run `c5deb72`, linux-arm64, reps=5) +> +> | | L=128 | L=512 | L=1024 | L=2048 | L=4096 | L=8192 | +> |---|---|---|---|---|---|---| +> | **exp-sharing (fused vs two-call)** | 1.75× | 1.70× | 1.69× | 1.65× | 1.58× | 1.69× | +> | bidirectional (fused) vs torch.compile | 6.15× | 7.70× | 8.77× | 9.54× | — | — | +> | bidirectional (fused) vs eager | 27.7× | 40.4× | 41.0× | 42.6× | 41.3× | 42.6× | +> +> `fused == two-call` **bit-identical** (rel 0.0e+00) at every shape; +> kernel-vs-f64 max abs 4.3e-6 – 8.6e-6, well inside the 1e-4 gate. +> +> **The striking one:** at L=512, `scan_fwd` = 2.85 ms and fused bidirectional = +> 3.27 ms — **a full bidirectional scan for ~1.15× the cost of a single +> unidirectional one.** Sharing the expensive Pass A makes the second direction +> nearly free. That is the whole thesis of Step 7, measured. + +### The one open risk is settled — in our favor + +Every stage flagged the same worry: materializing `abar`/`bbar` full-row (to +share them across directions) means they stream from L2 instead of L1, and the +exp saving had to beat that extra round-trip. **It did, decisively.** The +measured 1.65× geomean matches the profiler-projected ~1.7× almost exactly, so +the L2 traffic cost essentially nothing next to eliminating a second exp sweep. +No assumption — the benchmark's `exp_sharing_speedup` is the actual old path +(two-call) against the actual new one (fused), and it reproduces the projection. + +**Reproduced across four independent CI runs** (`c5deb72`, `605b056`-era reruns, +`8ab3b72`, `d7c23fa`): the per-shape exp-sharing numbers land in **1.58–1.75×** +every time, geomean ~1.67×. Unlike the earlier fused-`reverse` "win" that turned +out to be shared-runner noise (Step 5), this one is a real, stable effect — +because it is a large one (~1.7×) sitting well above the ±10–20% runner variance, +and because the projection predicted it before it was measured. + +### Why the headline moved too + +`bidirectional_scan` now routes the common case through the fused op, so the +whole bidirectional-vs-baseline story improved by the ~1.7× factor: **6.15–9.54× +vs torch.compile** (was ~5.5×), **27.7–42.6× vs eager** (was ~24×). Same honest +baselines, same host; the kernel just got faster. + +--- + +### Build record (how it was staged) + +**Status:** Stage 1 (scalar) ✅ and Stage 2 (NEON) ✅ green — bit-identical to +two scans on real Arm. Stage 3 (FFI + Python + benchmark) ✅ — the measured +result above. + +Motivated by the roofline in +[`BIDIRECTIONAL_SPEEDUP_IDEAS.md`](./BIDIRECTIONAL_SPEEDUP_IDEAS.md) §5: we are +**compute-bound on the exp** (85% of runtime), not bandwidth-bound. And Pass A +(discretize + exp + input projection) is **pointwise in time — direction- +independent** — yet `bidirectional_scan` computes it *twice* by calling the +kernel twice. Sharing it is the win: **one exp sweep + two cheap FMA sweeps** +instead of two full scans, projecting to **~1.7×** on bidirectional (and the +same structure amortizes exp across SS2D's four directions — this is the SS2D +substrate, exactly as `reverse` was). + +### What landed (Stage 1) + +A new core entry point `selective_scan_bidirectional` that produces `out_fwd` and +`out_bwd` from one input, computing Pass A once per channel: + +| Layer | Change | +|---|---| +| `src/parallel.rs` | `for_each_channel_bidir` — two-output rayon driver; parallel across channels, each channel does the shared Pass A + both recurrences | +| `src/scalar.rs` | `scan_bidirectional` + `BidirScratch` (materializes `abar`/`bbar` for the whole row so both directions read them) | +| `src/lib.rs` | `selective_scan_bidirectional` public API + dispatch (scalar now; NEON = Stage 2) | +| `tests/property.rs` | `fused_bidirectional_matches_two_scans` — the gate | + +### The correctness argument, and why the gate is bit-exact + +Fused output **= two standalone scans**: `out_fwd` bit-identical to +`selective_scan(reverse=false)`, `out_bwd` to `selective_scan(reverse=true)`. +This holds bit-for-bit (not merely within tolerance) because the shared products +carry exactly the values each standalone direction computes inline: `abar = +exp(dt·A)` is a pure function of pointwise inputs, `bbar = dt·u·B` likewise, and +the recurrence `abar·h + bbar` and dot `C·h` are consumed in the same order. +Scalar generates no FMA contraction by default, so `mul(exp,h) + mul(dtu,b)` is +identical whether the second product is computed inline or read from `bbar`. + +The gate checks scalar-fused vs scalar-twice **on the same backend**, so any +difference is a fusion bug, not a SIMD-vs-libm gap — and under **both** threadings +(Sequential and Rayon), since small proptest shapes would otherwise leave the new +two-output parallel driver untested. + +### Staged deliberately (scalar before NEON) + +Same discipline as the original kernel (Phase 1 scalar, Phase 2 NEON) and forced +by the no-local-Rust constraint: prove the *structure* — the API, the parallel +driver, the bit-identity — on the low-risk scalar path first, then add the fast +path on a validated foundation. Stage 1 changed no numerics and touched no FFI, so +the existing bidirectional path was unaffected until Stage 3 rewires it. + +### Stage 2 (NEON) — what landed + +The fast path where the exp-sharing actually pays off: + +| Change | Detail | +|---|---| +| `discretize_chunk` refactor | now takes `dt`/`dtu` **slices** instead of `&mut Scratch`, so the single and fused paths share it with different scratch layouts (both existing callers updated) | +| `neon::scan_bidirectional` | transpose B/C once (shared across channels *and* directions), then per channel: Pass A once → full-row `abar`/`bbar`, Pass B forward + backward | +| `channel_n16_bidir` + `pass_b_n16` | N=16 fast path — Pass B factored into a helper run twice (forward, reversed), state in the register file exactly as `channel_n16` | +| `channel_general_bidir` | any-N path; unlike single-direction `channel_general` (which keeps exp inline) it **materializes** `abar`/`bbar`, because sharing requires storing them | +| `lib.rs` `try_neon_bidir` | the `T == f32` NEON dispatch, mirroring `try_neon`; Auto→NEON on aarch64, else scalar | +| `tests/property.rs` | the gate now runs on **`Scalar` and `Auto`** — on aarch64 that checks NEON-fused vs NEON-two-scans bit-for-bit | + +Bit-identity holds on NEON for the same reason as scalar, one level down: `abar = +vexpq_f32_nonpos_fast(dt·A)` and `bbar = B·dtu` are the identical vectors the +standalone `channel_n16`/`channel_general` compute, and Pass B's `vfmaq(bbar, +abar, h)` is the same fused multiply-add in the same lane order. Materializing them +full-row instead of chunk-local changes storage, not values. + +**CI confirmed it: `fused_bidirectional_matches_two_scans` passed on macOS-arm64 +(real NEON), bit-for-bit, across all shapes and both threadings.** The NEON fused +kernel equals two standalone NEON scans exactly. Stage 2 correctness is done. + +Two incidental breakages on the same push, neither a Stage-2 numerics bug: + +1. **The profiler didn't compile** — `profile.rs` still called `discretize_chunk` + with the old `&mut Scratch` signature. Missed because it is behind the + `profiling` feature (only the `Profile kernel` workflow compiles it). Fixed to + the slice signature. +2. **`f32_matches_f64` failed on an unrelated case** — `err = 1.145e-3` on an + output of `-271.8`, i.e. a **4.2e-6 relative** f32 error, rejected only because + that test used an *absolute* `1e-3` bound. Not a regression: the sole change to + the plain-scan path was the behavior-preserving `discretize_chunk` refactor, and + proptest reseeds randomly each run — it simply generated a large-output case + (a 2.4 input through a long near-identity recurrence) this time. **Fixed** by + scaling the bound with output magnitude (`err < 1e-3 * max|out|`), which keeps + the old bound for bounded outputs and holds large ones to the same *relative* + accuracy. This also removes a latent flake that could have failed any prior run + whose random seed hit a large output. + +### Stage 3 — plumbing + measurement (written, awaiting CI) + +Wired the fused kernel all the way out to the benchmark: + +| Layer | Change | +|---|---| +| `arm-scan-ffi` | `arm_scan_selective_scan_bidirectional_f32` (out_fwd, out_bwd, last_fwd, last_bwd); **ABI 4 → 5**; a hand-computed FFI test | +| `_ffi.py` | ABI 5, argtypes, `scan_bidir_raw` wrapper | +| `op.py` | `_selective_scan_bidir_op` (`torch.library` custom op, two outputs, registered fake) + public `selective_scan_bidirectional` | +| `bidirectional.py` | `bidirectional_scan` now uses the fused op for the common case; **falls back to two calls** for untied weights (`reverse_params` — Pass A can't be shared) or when a `last_state` is requested (the fused op doesn't produce one) | +| `bench_bidirectional.py` | new `bidirectional_twocall` (two scans, exp recomputed) vs `bidirectional` (fused, exp shared); reports `exp_sharing_speedup = twocall / fused` | + +The benchmark's correctness gate now checks **fused == two-call bit-identical** +(sharing the exp reuses identical values, so it must be exact — no SIMD-tail +caveat, unlike the flip-vs-reverse comparison). `check_bidirectional.py` already +exercises the fused op end-to-end (its tied cases route through it), so the +f64-reference gate covers it too. + +**The dispatch boundary is deliberate.** The fused op is only valid when both +directions share A/delta/B/C (tied weights) — an *untied* model (Vim v2's +separate `A_b`/`D_b`) genuinely cannot share Pass A, so it correctly falls back. +And the fused op omits `last_state` because bidirectional models are non-causal +and ignore it; requesting one also falls back. Neither fallback is a compromise — +each is the right computation for a case the fused kernel does not cover. + +### The number this produced + +`exp_sharing_speedup = twocall / fused` came in at **1.58–1.75× (geomean ~1.67×)** +— see the result box at the top of this step. Matched the ~1.7× projection, and +settled the L2-round-trip risk in our favor. + +### Known tradeoff to watch + +Sharing Pass A across directions requires `abar`/`bbar` materialized for the +whole row (`L × state` each) rather than the current chunk-local scratch +(`CHUNK × state`). At long L this spills L1 → L2 (same order as the existing +`bt`/`ct` planes, so not a new category of cost), but the NEON Stage 2 must +confirm the exp saving beats the extra L2 round-trip — measure, do not assume. + +--- + +## Step 8 — smoke check: both topologies, one pass/fail (and the win in two numbers) + +`bench/smoke_topologies.py` is a fast consolidated sanity gate (wired into the +`bench-op` CI job): across an **L sweep** (default 512 / 2048 / 4096) it verifies +**both** topologies are correct (max_abs vs an f64 reference, gated 1e-4) and +faster than native torch eager, exiting non-zero if any case fails. It exists so +"do unidirectional and bidirectional both run and beat torch across L" has a +single unambiguous answer per CI run, rather than being reconstructed from two +separate benchmark outputs. + +The practical **ceiling on L is the torch reference itself** — an O(L) +Python-loop scan — not our kernel, which runs in constant memory at any L. As L +grows the eager baseline is what gets slow, which is the finding, not a +limitation: torch is the thing that can't keep up. The cap is kept at 4096 (not +8192) to stay clear of the runner limits that the `torch.compile` graph-unroll +hit in Step 6 — the eager reference builds no graph, but the conservative cap +costs nothing. (`--lengths` overrides the sweep for manual runs.) + +### The measured result — sweep (linux-arm64 CI, run `d7c23fa`, reps=3) + +``` + L topology gate max_abs kernel(ms) eager(ms) speedup + 512 unidirectional PASS 4.1e-06 2.76 60.14 21.8x + 512 bidirectional PASS 6.7e-06 3.49 119.68 34.3x + 2048 unidirectional PASS 4.3e-06 10.31 259.34 25.2x + 2048 bidirectional PASS 5.5e-06 12.53 513.20 41.0x + 4096 unidirectional PASS 5.4e-06 20.59 527.22 25.6x + 4096 bidirectional PASS 7.0e-06 27.10 1055.70 39.0x +``` + +**vs native torch eager: unidirectional 21.8–25.6×, bidirectional 34.3–41.0×**, +all PASS. Unidirectional climbs with L (21.8 → 25.6) as fixed per-call overhead +amortizes; bidirectional runs ~34–41× (the higher L=2048 vs L=4096 is baseline +noise, not a trend). These sit slightly under the full `bench_bidirectional.py` +figures for the same run (e.g. L=512 34.3× here vs 38.6× there) purely because +the smoke check uses fewer reps and a noisier eager baseline — same kernel. + +### The exp-sharing win is visible in two numbers here + +At L=512: + +| | unidirectional | bidirectional | ratio | +|---|---|---|---| +| **eager** (native torch) | 60.14 ms | 119.68 ms | **1.99×** | +| **kernel** (ours, fused) | 2.76 ms | 3.49 ms | **1.26×** | + +Native torch pays ~2× for the second direction (it runs two scans). Our fused +kernel pays only ~1.26×, because Pass A (the exp, ~85% of the work) is computed +once and shared. **That is the whole Step-7 result in a 2×2 table** — and it is +why bidirectional gets a *bigger* speedup than unidirectional (34.3× vs 21.8×): +torch's cost doubles, ours barely moves, so the gap widens. + +### Reconciling with the benchmark table (they are NOT different numbers) + +The full `bench_bidirectional.py` reported 39.98× vs eager at L=512; the smoke +check reported 37.9×. Same computation, different CI runs: + +| | kernel (ms) | eager (ms) | speedup | +|---|---|---|---| +| smoke | 3.30 | 125.08 | 37.9× | +| bench | 3.277 | 131.0 | 39.98× | + +The **kernel** times are 0.7% apart (identical); the **eager** baselines are 4.5% +apart — ordinary ±10–20% shared-runner noise. The entire speedup gap is that +eager wobble, not the kernel. The unidirectional 23.2× has no row in the +bidirectional-only bench table (different topology; its home is `bench_op.py` / +`BASELINE_REPORT.md`, which also says ~23× at L=512 — it matches). Nothing is +inconsistent; the two views simply measured a noisy baseline on different days. diff --git a/BIDIRECTIONAL_SPEEDUP_IDEAS.md b/BIDIRECTIONAL_SPEEDUP_IDEAS.md new file mode 100644 index 0000000..5bfdefd --- /dev/null +++ b/BIDIRECTIONAL_SPEEDUP_IDEAS.md @@ -0,0 +1,280 @@ +# BIDIRECTIONAL_SPEEDUP_IDEAS — closing the gap to the GPU + +**Status: research/ideas only — nothing here is implemented.** Written Jul 15, 2026, +from a read of the official Mamba CUDA kernel design, the CPU-vs-GPU bandwidth +reality, and the current NEON kernel at the post-merge state. + +**Scope.** How to make the *bidirectional* scan faster on Arm, with the explicit +goal of approaching GPU single-stream performance. This is a sibling to +[`IMPROVEMENT_IDEAS.md`](./IMPROVEMENT_IDEAS.md) (general kernel optimization) and +does not repeat it — where an idea overlaps, it is cited, not restated. The +bidirectional-native ideas (§3) are new white space that neither document nor the +upstream repo covers. + +Rules that bind every idea (per [`CLAUDE.md`](./CLAUDE.md)): +- **Correctness gates speed.** Every new path gets its own golden gate and error + floor before it is benchmarked. The bidirectional definition + (`tests/check_bidirectional_math.py`, `reverse_matches_flip_forward_flip`) is + the reference any faster path must reproduce. +- **Benchmark honestly.** `torch.compile` is the baseline; medians after warmup; + state the host. Current standing: **~5.2–5.6× vs torch.compile at L ≥ 512**, and + torch.compile **OOMs at L=8192** on a 4-core arm64 box (see `BIDIRECTIONAL_LOG.md`). +- **Measure before building.** Every idea below names the measurement that + justifies (or kills) it. The fused `reverse` flag is the cautionary tale: built + on a ~2% assumption, worth building only because SS2D needed it anyway. + +--- + +## 1. The reality: what "as fast as GPU" can and cannot mean + +The gap on this op is **memory bandwidth, ~10–20×**, not compute: + +| | bandwidth | source | +|---|---|---| +| Graviton4 (DDR5) | ~0.5 TB/s | [NextPlatform](https://www.nextplatform.com/compute/2024/09/19/aws-boosts-memory-capacity-on-graviton-4-compute/1642050) | +| H100 (HBM3) | ~4 TB/s | [Spheron](https://www.spheron.network/blog/hbm3e-vs-hbm4-vs-hbm4e-llm-inference-guide/) | +| B200 (HBM3e) | ~8 TB/s | same | + +**The opening:** selective_scan on a GPU runs at **10–15% utilization** (vs 80–90% +for a transformer's tensor cores — [Mamba-2 co-design](https://medium.com/@danieljsmit/mamba2-the-hardware-algorithm-co-design-that-unified-attention-and-state-space-models-77856d2ac4f4)). +The GPU is *also* memory-bound here and is not using the hardware that makes it a +GPU. The scan is a diagonal recurrence, not a matmul; the tensor cores sit idle. + +Consequences that set the whole agenda: + +- **Concede batch throughput.** 10–20× bandwidth wins raw tokens/sec. Publish that + losing row honestly (per `IMPROVEMENT_IDEAS.md §11`); do not chase it. +- **Chase single-stream latency and $/token.** At batch=1 the GPU's bandwidth is + *stranded* — one sequence cannot fill it — and per-launch overhead dominates. + This is exactly where bidirectional inference lives (non-causal, one sequence at + a time). Parity here is real and defensible. +- **Own the workload the GPU cannot do at all.** Constant-memory scan at L=131k: + a transformer's KV cache is multiple GB, and — measured — `torch.compile` cannot + even *build the graph* at L=8192 on CPU. Not "faster"; a different class. + +So the target is precise: **single-stream bidirectional latency within a small +factor of the GPU, at a fraction of the instance cost** — not a throughput race +we would lose. + +--- + +## 2. What the official CUDA kernel does — and what we already have + +Three techniques ([Mamba paper §D](https://arxiv.org/pdf/2312.00752), +[Princeton PLI](https://pli.princeton.edu/blog/2024/mamba-2-algorithms-and-systems)): + +| GPU technique | Our status | +|---|---| +| **Fuse discretize + scan + gate in SRAM**, load A/B/C/Δ from HBM once, never materialize the `(B,D,L,N)` intermediate | ✅ **done** — two-pass NEON kernel keeps chunk scratch in L1, streams B/C once (`neon/mod.rs`) | +| **N=16 state resident in registers** | ✅ **done** — `channel_n16`, four q-registers | +| **Work-efficient parallel associative scan over the sequence length L** (Blelloch tree), so a single sequence still fills the machine | ❌ **not done** — we parallelize over batch×channel only. This is the one real gap, and §4 is about it. | + +The takeaway: **structurally we already match the GPU's fusion strategy.** Our moat +over the `torch.compile` strawman *is* that fusion — the compiler cannot restructure +the recurrence, so it unrolls and dies. What we are missing is the GPU's answer to +the batch=1 problem, which is the L-parallel scan. + +Best borrowable reference: **[`mamba-mini`](https://github.com/MzeroMiko/mamba-mini)** +— "the code closest to `selective_scan_cuda`," single-file, CPU+GPU, with the math +derivation. Read this for exact discretization/scan order, not the CUDA source +(which is hard to map onto NEON). + +--- + +## 3. Bidirectional-native wins (white space — the GPU does not do these) + +**The official Mamba repo has no bidirectional kernel.** Bidirectional is Vim / +Caduceus territory, and those just call the forward scan twice. So there is nothing +to "borrow" here — it is unclaimed, and there are two strong ideas. + +### 3.1 Run forward and backward in parallel across cores ⭐ cheapest, do first + +The two directions are **fully independent** — no data dependency between them. We +currently run them as two sequential kernel calls (`bidirectional_scan` → +`selective_scan` then `selective_scan(reverse=True)`). When `batch × dim` does not +saturate the core count — which is the *common* bidirectional case (B=1) — the +second direction waits while cores sit idle. + +Dispatch forward to one rayon scope and backward to another, concurrently. +**Potentially ~2× on bidirectional latency, zero new math, zero numerics risk.** + +- **Effort:** low. It is a scheduling change at the `bidirectional_scan` layer (run + the two `selective_scan` calls on separate threads / a rayon `join`), or one + level deeper: a kernel entry that takes both an `out_fwd` and `out_bwd` and + parallelizes over `(channel × direction)` instead of `channel`, doubling the + independent work units fed to rayon. +- **Measure first:** does B=1/D=768 actually leave cores idle on the target? On 4 + cores with 768 channels, no — already saturated, so this buys nothing there. On a + **64-core Graviton at B=1**, 768 channels across 64 cores is ~12 rows/core, and a + second direction is free parallelism. This is a *big-core-count* win; confirm the + idle cores exist before building. +- **Risk:** ~0. Outputs are unchanged; only the schedule differs. Parity tests pass + by construction (each direction is the existing, verified path). + +### 3.2 Fuse both directions into one pass — SHARE the exp ⭐⭐ the real win + +**This is the highest-value item in the whole document, and the roofline (§5) +is why.** The framing that matters is not bandwidth — it is *redundant compute*. + +**The load-bearing fact:** Pass A (discretize + exp + projection) is **pointwise in +time and therefore direction-independent**. `dt = softplus(δ+bias)`, +`ābar = exp(dt·A)`, `b̄ = dt·u·B` are identical whether the sequence is scanned +forward or backward. Only Pass B (the recurrence `h = ābar·h + b̄`, `y = C·h`) +walks time and depends on direction. + +But `bidirectional_scan` calls `selective_scan` **twice**, so it computes Pass A — +which the profiler measured at **83% of runtime** (exp 59% + discretize 19% + +projection 5%) — **twice**. We are duplicating the expensive transcendental work, +and the kernel is compute-bound on exactly that work (§5). + +A fused kernel computes Pass A **once**, then runs Pass B (the cheap ~10% pure-FMA +recurrence) twice — once forward, once backward — from the shared, L1-resident +`ābar`/`b̄` chunk data. Using the measured phase split: + +| | cost (1 forward scan = 100%) | +|---|---| +| current bidirectional (2 calls) | **200%** | +| fused: Pass A ×1 (83%) + Pass B ×2 (21%) + epilogue ×2 (14%) | **~117%** | +| **projected speedup** | **~1.7×** | + +That is a **compute** saving — sharing the exp — not a bandwidth trick, and it rests +on already-measured profiler data, not a guess. (It also happens to halve input +traffic as a side effect, but §5 shows that is not the binding constraint.) + +- **Strategic multiplier for SS2D:** SS2D's four directions scan the same grid with + the same pointwise discretize+exp. Compute Pass A **once**, run Pass B **four + times** (the four traversals) → the shared 83% is amortized across 4 directions + instead of duplicated 4×. Projected saving on the shared part approaches **~3–4×**. + Same kernel structure. **Build it at 1D here, reuse it at 2D.** This is the exact + "read once, emit multiple directions" core of `TOPOLOGY_IMPLEMENTATION_PLAN.md §3.2` + — the SS2D substrate, just as `reverse` was. +- **Effort:** medium. New core entry point with two output buffers and two `h` + accumulators; Pass A unchanged and run once; Pass B run twice over the shared + chunk scratch (forward index, then reversed index — the `reverse` logic already + exists). `parallel.rs`, the transpose, discretize, and exp are all reused as-is. +- **Risk:** low-medium. New kernel surface → new golden gate, but each direction's + math is the existing verified path, so it diffs bit-for-bit (scalar) / ~1e-7 + (NEON) against the two-call `bidirectional_scan`, exactly like `reverse` did. +- **Measure:** the ~1.7× projection assumes Pass A fully shares. Confirm the phase + split still holds at the fused structure (the profiler, `PROFILING.md`), and that + keeping two `h` register sets live does not spill (§3.3.1 of `IMPROVEMENT_IDEAS.md` + notes 8 h-registers still fit the 32-register file at N=16). + +--- + +## 4. The GPU-parity mechanism: parallel scan over L + +**This is how the GPU wins single-stream, and it is the biggest borrow.** Covered as +a stretch in `IMPROVEMENT_IDEAS.md §4.3`; elevated here because the CPU-vs-GPU +research makes it *the* lever for the parity goal, not a nice-to-have. + +The recurrence `h_t = ā_t · h_{t-1} + b̄_t` is an **associative scan**: the transition +`(ā, b̄)` composes as `(a,b) ∘ (a',b') = (a·a', a'·b + b')`. So it parallelizes with a +3-phase Blelloch scan across chunks: + +1. **Per-chunk (parallel):** compose each chunk's timesteps into a single + `(A_prod, B_comb)` transition. Our chunked structure already computes the pieces. +2. **Combine (tiny, sequential):** scan the `L/CHUNK` per-chunk carries — a handful + of elements. +3. **Finalize (parallel):** re-run each chunk from its now-known entry state to emit + outputs. + +Cost: one extra elementwise pass over the data. Payoff: **full-machine parallelism +when B×D < cores** — i.e. the single-stream latency case (B=1 audio, ECG, genomics) +on a big Graviton, which is precisely the GPU-parity regime. + +- **Bidirectional angle:** combines with §3.1/§3.2 — a bidirectional scan on a + 64-core box could parallelize over *both* directions *and* the L dimension, + saturating the machine on a single sequence the way the GPU does. +- **Effort:** high (the real item on this list). But the chunked two-pass kernel + already exposes the composed pairs, so the groundwork is laid. +- **Risk:** medium. Numerics: composing `A_prod` across a long chunk multiplies many + `ā ∈ (0,1]` values → underflow toward 0, which is *fine* (it means the distant past + doesn't affect the present — the `§7.2` underflow-cut observation in + `IMPROVEMENT_IDEAS.md` is the same fact) but must be handled so it doesn't NaN. + Gate hard against the sequential reference. + +--- + +## 5. The roofline — done analytically, and it redirected §3.2 + +Computed from the code + the existing profiler split (confirm with PMU later, but +the conclusion is robust): + +**Arithmetic intensity at the mamba shape** (B=1, D=768, L=512, N=16): +- main-memory traffic ≈ **6.4 MB** (u, delta, z, out ≈ 1.57 MB each; B/C/A are tens + of KB) +- ≈ **85 MFLOP** (~13 FLOP per (b,d,l,n) point × 6.3M points) +- **intensity ≈ 13 FLOP/byte** + +**Graviton4 machine balance** ≈ 5.5 TFLOP/s ÷ 0.5 TB/s ≈ **11 FLOP/byte**. We sit +right at the knee — but the exp is **transcendental**, and the profiler already +measured **~85% of runtime in exp/softplus/silu** vs ~10% in the recurrence. So in +practice **we are compute-bound on the transcendentals, not bandwidth-bound.** + +**Why this matters, and what it changed:** the GPU is bandwidth-bound here because it +has *too much* compute (10–15% utilization). We are the mirror image — compute-bound +because SIMD gives us relatively little, and the exp is expensive. So **bandwidth +reduction is not our lever; cutting redundant compute is.** That is exactly why §3.2 +was rewritten from "halve input bandwidth" (wrong axis) to "share the direction- +independent exp between the two scans" (the real, ~1.7× win). + +**Still worth doing as a writeup metric:** report **% of the CPU's memory bandwidth** +achieved, mirroring the GPU's own 10–15%-utilization honesty. But treat it as a +number for the pitch, not as the optimization compass — the compass says compute. +When (if) compute is fully shared and parallelized and we *become* bandwidth-limited, +that is when `IMPROVEMENT_IDEAS.md §5.1` (fp16 plane storage) becomes the last lever. + +--- + +## 6. Explicitly NOT worth borrowing + +- **Mamba-2 / SSD matmul reformulation.** The GPU recasts the scan as matmuls to use + tensor cores. On CPU with diagonal A and N=16 the plain scan is already near SIMD + peak and the dual strictly *adds* FLOPs — `IMPROVEMENT_IDEAS.md §7.1` argues this + and the research confirms it. Relevant only if Mamba-2 support is added, and even + then as a measured rejection. +- **Warp-level primitives / shared-memory staging.** GPU-specific; the NEON analog + (registers + L1) is already what we do. +- **Chasing batch throughput.** Conceded above; it is a bandwidth race we lose. + +--- + +## 7. Priority shortlist + +Re-ranked after the §5 roofline: we are **compute-bound on the exp**, so the winner +is the one that stops computing it twice. + +| # | Idea | § | Effort | Expected effect | Risk | +|---|---|---|---|---|---| +| **1** | **Fused two-direction pass — share the exp** | **3.2** | **Med** | **~1.7× on bidirectional; ~3–4× shared-part for SS2D; the SS2D substrate** | **Low-med** | +| 2 | Fwd ∥ bwd across cores | 3.1 | Low | ~2× latency **only when cores are idle** (big-core, B=1) — needs a dedicated host to see | ~0 | +| 3 | L-parallel Blelloch scan | 4 | High | Single-stream GPU-parity; unlocks 64-core at B=1 | Med | +| 4 | %-bandwidth line for the writeup | 5 | Low | No speedup; pitch metric, not the compass | ~0 | + +**Sequencing:** **#1 first** — it is the largest bidirectional win, it is directly +verifiable on the 4-core CI runner (unlike #2, which needs idle cores that only exist +on a big-core host), and it *is* the SS2D substrate, so it pays off twice exactly +like `reverse` did. Then #2 and #3 on a dedicated Graviton session where the core +count makes them visible. #4 rides along with any benchmark run. + +Note #1 and #2 **compose** — #1 shares the exp; #2 then runs the two (now +exp-free) Pass-B recurrences on separate cores. And both feed SS2D. + +**The honest north star:** we already match the GPU's *fusion*. The single-stream gap +closes by (a) not computing the transcendental Pass A twice for bidirectional / 4× +for SS2D (#1), and (b) parallelizing the one sequence across all the cores the GPU +would spread across threads (#2, #3). None of it needs the tensor cores the GPU is +not using anyway. + +--- + +## Sources + +- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces (§D, hardware-aware scan)](https://arxiv.org/pdf/2312.00752) +- [Princeton PLI — Mamba-2: Algorithms and Systems](https://pli.princeton.edu/blog/2024/mamba-2-algorithms-and-systems) +- [Tri Dao — State Space Duality (Mamba-2) Part III: the algorithm](https://tridao.me/blog/2024/mamba2-part3-algorithm/) +- [mamba-mini — CPU/GPU single-file selective scan closest to the CUDA kernel](https://github.com/MzeroMiko/mamba-mini) +- [Mamba-2 hardware-algorithm co-design (10–15% utilization figure)](https://medium.com/@danieljsmit/mamba2-the-hardware-algorithm-co-design-that-unified-attention-and-state-space-models-77856d2ac4f4) +- [AWS Graviton4 memory bandwidth (~0.5 TB/s DDR5)](https://www.nextplatform.com/compute/2024/09/19/aws-boosts-memory-capacity-on-graviton-4-compute/1642050) +- [GPU HBM bandwidth (H100 ~4 TB/s, B200 ~8 TB/s)](https://www.spheron.network/blog/hbm3e-vs-hbm4-vs-hbm4e-llm-inference-guide/) diff --git a/TOPOLOGY_IMPLEMENTATION_PLAN.md b/TOPOLOGY_IMPLEMENTATION_PLAN.md index baccae4..37cd68f 100644 --- a/TOPOLOGY_IMPLEMENTATION_PLAN.md +++ b/TOPOLOGY_IMPLEMENTATION_PLAN.md @@ -1,6 +1,10 @@ # TOPOLOGY_IMPLEMENTATION_PLAN — 1D bidirectional & 2D cross-scan (SS2D) -**Status:** plan, not yet implemented. Written Jul 14, 2026. Companion to [`APPLICATIONS.md`](APPLICATIONS.md) (which topology feeds which showcase app) and [`INTEGRATION_PLAN.md`](INTEGRATION_PLAN.md) (Phases 0–6, already landed). +**Status:** the plan. Written Jul 14, 2026. Companion to [`APPLICATIONS.md`](APPLICATIONS.md) (which topology feeds which showcase app) and [`INTEGRATION_PLAN.md`](INTEGRATION_PLAN.md) (Phases 0–6, already landed). + +> **Execution is tracked in per-topology logs** — what has actually been built against this plan, what it is verified against, and what broke on the way. This file stays the plan; the logs are the record. +> - §2 (1D bidirectional) → [`BIDIRECTIONAL_LOG.md`](BIDIRECTIONAL_LOG.md). **§2.1 is landed and verified on Arm. §2.2 was measured and REJECTED** — the flip copies cost ~2%, not enough to fuse. Read that before touching §2.2. +> - §3 (2D cross-scan / SS2D) → not started; gets its own log when it does. Note §3's copy overhead is a *different* question — it materializes four grid views, not one flipped sequence — so §2's rejection does **not** transfer to it. Measure it separately. **Scope:** how to take 1D bidirectional and 2D cross-scan from "correct via Python rearrangement" to "fast AND correct in Rust," using the already-shipped 1D unidirectional kernel as the architectural template. Section 1 documents that template as the reference; Sections 2 and 3 are the plans for the two new topologies; Section 4 is cross-cutting methodology and sequencing. @@ -49,7 +53,30 @@ No Rust changes. In Python: call `selective_scan` (§1.5) twice — once as-is, Deliverable: `python/arm_scan/bidirectional.py` (new) exposing a `bidirectional_scan(...)` function with the same tensor-shape contract as `op.selective_scan` plus a merge strategy parameter, callable from whichever app-specific integration (genomics/audio/ECG model class) needs it. This unblocks app-level demos and benchmark numbers immediately, independent of anything in §2.2. ### 2.2 Fast path: fused `reverse` flag in Rust -**Estimated effort: half a day**, because the recurrence math doesn't change — only traversal order does. + +> ### ✅ BUILT — but read why, because it is NOT a speedup. +> +> `bench/bench_bidirectional.py` measured what the flip copies cost: a **ceiling** +> of 1.085× at L=128, falling to 1.025× at L=512 — i.e. **~2%**, shrinking as L +> grows. As a bidirectional optimization this does not pay, and that finding +> stands. +> +> It was built anyway because **SS2D needs a backward traversal regardless** — §3.2's +> row-backward and column-backward directions reuse the 1D scan with this exact +> flag. `reverse` is the *substrate* for the 2D cross-scan; removing bidirectional's +> flip copies is a side effect, not the reason. +> +> **Ship it as:** *"a fused backward traversal, the substrate for the 2D cross-scan."* +> **Never as:** *"we made bidirectional faster."* +> +> **One correction to the design below:** it claims the NEON work needs SIMD +> lane-reversal (`vrev64q_f32`/`vextq_f32`). **It does not.** Pass A is pointwise in +> time, and Pass B vectorizes across *state* while `t` is a scalar index — so +> reversing is one subtraction, no shuffles. The real change was a chunk-order +> iterator plus a flipped index. Full write-up: +> [`BIDIRECTIONAL_LOG.md`](BIDIRECTIONAL_LOG.md) Step 3. + +**Original estimate: half a day**, because the recurrence math doesn't change — only traversal order does. | Layer | Change | |---|---| diff --git a/bench/bench_bidirectional.py b/bench/bench_bidirectional.py new file mode 100644 index 0000000..f255c39 --- /dev/null +++ b/bench/bench_bidirectional.py @@ -0,0 +1,391 @@ +"""Bidirectional scan: kernel vs PyTorch, and the fused-kernel exp-sharing win. + +A bidirectional scan runs the recurrence in both time directions and sums them. +The direction-independent Pass A (discretize + exp, ~85% of the work) can be +computed once and SHARED between the two directions — that is the fused +two-direction kernel (BIDIRECTIONAL_SPEEDUP_IDEAS.md §3.2). This file measures +what that sharing buys, and where the fused path stands against stock PyTorch. + +WHAT IS TIMED + +Baselines — the rows that matter, and the only ones fit to publish: + ref_eager_bidi stock PyTorch, both directions, merged. What a real + bidirectional Mamba does on CPU with no kernel installed. + ref_compile_bidi the same under torch.compile — the FAIR baseline, since + compile is what a competent user would already be doing. + +Our two implementations: + bidirectional_twocall two separate kernel scans (forward + the `reverse` + flag), summed. Each recomputes Pass A → exp twice. + bidirectional the FUSED two-direction kernel: exp computed once and + shared. This is `arm_scan.bidirectional_scan`. + +Diagnostic: + scan_fwd one forward scan — half the work; the floor. + +THE NUMBERS + + speedup_vs_eager / speedup_vs_compile <- the result. Fused kernel vs PyTorch. + exp_sharing_speedup = twocall / fused <- the Stage-3 win, projected ~1.7x + (exp is ~85% and direction-independent). + +WHAT THE BASELINE ACTUALLY IS (so the speedup is not misread) + +The reference is `selective_scan_ref`, vendored verbatim from the official +`state-spaces/mamba` repo (Tri Dao / Albert Gu, Apache-2.0) — it matches HF +transformers' `MambaMixer.slow_forward` bit-for-bit. So this measures our kernel +against **the original Mamba selective scan**, in native PyTorch. It is a pure +Python-loop implementation, which is exactly what runs on CPU: the official +mamba-ssm CUDA kernel is GPU-only and does not run on CPU at all, so this is the +genuine CPU baseline, not a strawman. `torch.compile` of that same reference is +the stronger baseline and is measured too. + +PyTorch has no native *bidirectional* Mamba op; bidirectional Mamba (Vim, +Caduceus, VMamba) is always built by running the scan twice — forward and on the +flipped sequence — and summing. `ref_eager_bidi` does exactly that, so it is the +standard bidirectional construction, not something invented here. + +These are **op-level** numbers — the selective scan, which is the part this +kernel replaces. A full Mamba *model* also has GEMMs and a conv this kernel does +not touch, so end-to-end `generate()` speedups are more modest (~1.2–3x, see +bench_e2e.py). Do not quote an op-level figure as a whole-model speedup. + +Two correctness gates run before any timing: the fused op must equal the two-call +path (bit-identical — sharing the exp reuses identical values), AND the kernel +must agree with the PyTorch reference it is being timed against. + +Usage: + python bench/bench_bidirectional.py # sweep over L + python bench/bench_bidirectional.py --quick # CI-sized + python bench/bench_bidirectional.py --json out.json +""" + +import argparse +import json +import platform +import sys +import time +from pathlib import Path + +import torch + +REPO = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO / "python")) +sys.path.insert(0, str(REPO / "bench")) +sys.path.insert(0, str(REPO / "tests")) + +import arm_scan # noqa: E402 +# Reuse the harness rather than re-deriving it: the realistic value +# distribution (negative A, softplus delta in ~[1e-3, 0.1]) is load-bearing for +# these numbers, and a second copy of it would drift. +from bench_op import bench, env_report, git_sha, make_inputs # noqa: E402 +from reference import selective_scan_ref # noqa: E402 + +# Flip traffic is O(B·D·L) while scan work is O(B·D·L·N), so the ratio is +# roughly independent of L in theory — but cache behaviour is not, which is the +# whole reason to sweep it rather than assume. +SUITES = { + "sweep-len": [(1, 768, l, 16) + for l in (128, 512, 1024, 2048, 4096, 8192)], + "sweep-dim": [(1, d, 1024, 16) for d in (256, 768, 1536, 3072)], + "basic": [(1, 768, 512, 16), (1, 768, 2048, 16), (1, 1536, 1024, 16)], +} +QUICK_SHAPES = [(1, 768, 128, 16), (1, 768, 512, 16)] + +_FLIP_KEYS = ("u", "delta", "B", "C", "z") # the time-varying tensors + + +def flip_time(t): + return torch.flip(t, dims=(-1,)) + + +def scan_fwd(t): + return arm_scan.selective_scan( + t["u"], t["delta"], t["A"], t["B"], t["C"], D=t["D"], z=t["z"], + delta_bias=t["delta_bias"], delta_softplus=True) + + +# ---------------------------------------------------------------- baselines +# A bidirectional scan in stock PyTorch. This is what a real bidirectional +# Mamba does on CPU today with no kernel installed: run the scan, run it again +# on the time-flipped inputs, flip that back, sum. Built on the SAME vendored +# reference the rest of the project measures against (tests/reference/), so +# these rows are directly comparable to bench_op.py's. + + +def _bidi_ref(fn, t): + """Bidirectional scan via an arbitrary forward-scan callable `fn`.""" + fwd = fn(t["u"], t["delta"], t["A"], t["B"], t["C"], t["D"], t["z"], + t["delta_bias"]) + back = fn(flip_time(t["u"]), flip_time(t["delta"]), t["A"], + flip_time(t["B"]), flip_time(t["C"]), t["D"], + flip_time(t["z"]), t["delta_bias"]) + return fwd + flip_time(back) + + +def _ref_call(u, delta, A, B, C, D, z, delta_bias): + return selective_scan_ref(u, delta, A, B, C, D=D, z=z, + delta_bias=delta_bias, delta_softplus=True) + + +def ref_eager_bidi(t): + """Plain PyTorch, both directions. The eager baseline.""" + return _bidi_ref(_ref_call, t) + + +def make_ref_compile_bidi(t): + """torch.compile'd reference, both directions — the FAIR baseline. + + Compiled once here (outside the timed region) and returned as a closure; + compile time is reported separately, never folded into the median. Returns + None if inductor is unavailable on this host, which is a normal outcome (no + MSVC on Windows), not a failure. + """ + compiled = torch.compile(selective_scan_ref, dynamic=False) + + def call(u, delta, A, B, C, D, z, delta_bias): + return compiled(u, delta, A, B, C, D=D, z=z, delta_bias=delta_bias, + delta_softplus=True) + + t0 = time.perf_counter() + _bidi_ref(call, t) # warm the graph; this is where compilation lands + compile_s = time.perf_counter() - t0 + return (lambda: _bidi_ref(call, t)), compile_s + + +def bidirectional_twocall(t): + """Two separate scans (forward + the `reverse` flag), summed — NO exp + sharing. This is what `bidirectional_scan` did before the fused two-direction + kernel (Stage 3): each direction recomputes the full Pass A (discretize + + exp). `bidirectional()` below must equal this bit-for-bit, and the ratio + twocall / bidirectional is the exp-sharing win.""" + fwd = scan_fwd(t) + bwd = arm_scan.selective_scan( + t["u"], t["delta"], t["A"], t["B"], t["C"], D=t["D"], z=t["z"], + delta_bias=t["delta_bias"], delta_softplus=True, reverse=True) + return fwd + bwd + + +def bidirectional(t): + """The FUSED path: one call, Pass A (the exp) computed once and shared + between both directions. This is `arm_scan.bidirectional_scan`, which uses + the fused two-direction kernel for the tied-weights sum case.""" + return arm_scan.bidirectional_scan( + t["u"], t["delta"], t["A"], t["B"], t["C"], D=t["D"], z=t["z"], + delta_bias=t["delta_bias"], delta_softplus=True, merge="sum") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--suite", choices=sorted(SUITES), default="sweep-len") + ap.add_argument("--quick", action="store_true") + ap.add_argument("--reps", type=int, default=None) + ap.add_argument("--warmup", type=int, default=None) + ap.add_argument("--torch-threads", type=int, default=None) + ap.add_argument("--compile-max-len", type=int, default=None, + help="skip the torch.compile baseline beyond this L " + "(graph unrolling makes compile time explode). " + "Default 512, or 128 under --quick. An explicit value " + "always wins, including under --quick.") + ap.add_argument("--no-compile", action="store_true") + ap.add_argument("--tag", type=str, default=platform.node()) + ap.add_argument("--json", type=str, default=None) + args = ap.parse_args() + if args.compile_max_len is None: + # --quick's lower cap exists only to bound CI time. It must NOT override + # an explicit flag — that made the L=512 compile row unreachable from a + # quick run, which is exactly the row we wanted. + args.compile_max_len = 128 if args.quick else 512 + + if args.torch_threads: + torch.set_num_threads(args.torch_threads) + + shapes = QUICK_SHAPES if args.quick else SUITES[args.suite] + reps = args.reps or (5 if args.quick else 10) + warmup = args.warmup if args.warmup is not None else (1 if args.quick else 3) + + env = env_report() + env["tag"] = args.tag + env["git_sha"] = git_sha() + env["timestamp_utc"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + print("environment:") + for k, v in env.items(): + print(f" {k}: {v}") + print(f"\nbidirectional scan — suite=" + f"{'quick' if args.quick else args.suite} reps={reps} " + f"warmup={warmup} compile_max_len={args.compile_max_len}") + print(" vs eager / vs torch.compile : kernel (fused) vs stock PyTorch") + print(" exp-sharing speedup : fused vs two-call — the Stage-3 win " + "(exp computed once, not twice; projected ~1.7x)\n") + + results = { + "kind": "bidirectional-fused-exp-sharing", + "env": env, + "reps": reps, + "suite": "quick" if args.quick else args.suite, + "shapes": [], + } + + with torch.no_grad(): + for batch, dim, length, state in shapes: + label = f"B{batch}_D{dim}_L{length}_N{state}" + print(f"=== {label} ===") + t = make_inputs(batch, dim, length, state) + + # CORRECTNESS GATE 1: the fused two-direction op must reproduce the + # two-call path (each direction scanned separately). They share no + # code beyond the kernel, so if they disagree the exp-sharing changed + # the answer — refuse to report a speedup over a wrong result. This is + # expected BIT-IDENTICAL (the Rust test proves it): sharing Pass A + # reuses the identical exp values, it does not approximate them. + fused_out, twocall_out = bidirectional(t), bidirectional_twocall(t) + drift = (fused_out - twocall_out).abs().max().item() + scale = max(1.0, fused_out.abs().max().item()) + if drift / scale > 1e-6: + print(f" !! fused != two-call (rel={drift/scale:.3e}) — " + f"exp-sharing changed the answer; refusing to benchmark") + sys.exit(1) + exact = " (bit-identical)" if drift == 0.0 else "" + + # CORRECTNESS GATE 2: the kernel must agree with the PyTorch + # reference it is being timed against. + ref_out = ref_eager_bidi(t) + max_err = (fused_out - ref_out).abs().max().item() + row = {"shape": [batch, dim, length, state], "timings": {}, + "kernel_vs_ref_max_abs": max_err, + "fused_vs_twocall_max_abs": drift} + print(f" fused == two-call: rel {drift/scale:.1e}{exact} " + f"kernel-vs-ref max_abs {max_err:.3e}") + + series = [ + ("ref_eager_bidi", ref_eager_bidi), + ("scan_fwd", scan_fwd), + ("bidirectional_twocall", bidirectional_twocall), + ("bidirectional", bidirectional), + ] + for name, fn in series: + r = bench(lambda fn=fn: fn(t), warmup, reps) + row["timings"][name] = r + print(f" {name:22s} {r['median_s']*1e3:9.3f} ms") + + # torch.compile — the fair baseline. Compile cost is reported, never + # timed. Unavailable inductor (e.g. no MSVC on Windows) is a skip, + # not a failure. + if not args.no_compile and length <= args.compile_max_len: + try: + fn, compile_s = make_ref_compile_bidi(t) + r = bench(fn, warmup, reps) + r["compile_s"] = compile_s + row["timings"]["ref_compile_bidi"] = r + print(f" {'ref_compile_bidi':19s} " + f"{r['median_s']*1e3:9.3f} ms " + f"(one-time compile {compile_s:.1f}s)") + except Exception as e: + row["timings"]["ref_compile_bidi"] = {"error": str(e)[:200]} + print(f" {'ref_compile_bidi':19s} unavailable: " + f"{str(e)[:80]}") + elif not args.no_compile: + print(f" {'ref_compile_bidi':19s} skipped (L={length} > " + f"compile_max_len={args.compile_max_len})") + + bi = row["timings"]["bidirectional"]["median_s"] + tc = row["timings"]["bidirectional_twocall"]["median_s"] + eager = row["timings"]["ref_eager_bidi"]["median_s"] + + row["speedup_vs_eager"] = eager / bi + # The Stage-3 win: fused (shares the exp) vs two-call (recomputes it). + # Projected ~1.7x from the profiler's phase split (exp ~85% of the + # work, direction-independent). This is the ACHIEVED number. + row["exp_sharing_speedup"] = tc / bi + + line = f" => {row['speedup_vs_eager']:.2f}x vs eager" + comp = row["timings"].get("ref_compile_bidi", {}) + if "median_s" in comp: + row["speedup_vs_compile"] = comp["median_s"] / bi + line += f", {row['speedup_vs_compile']:.2f}x vs torch.compile" + print(line) + print(f" (exp-sharing: fused won {row['exp_sharing_speedup']:.3f}x " + f"over the two-call path)\n") + results["shapes"].append(row) + + # Flush after EVERY shape, not just at the end. At long L, + # torch.compile builds an L-step unrolled graph and can be killed by + # the OOM reaper mid-sweep — a hard kill, not an exception we could + # catch. Without this, an L=8192 death would destroy the 128..4096 + # results we already paid ~10 minutes of compilation for. + if args.json: + Path(args.json).parent.mkdir(parents=True, exist_ok=True) + Path(args.json).write_text(json.dumps(results, indent=2)) + + ev = [r["speedup_vs_eager"] for r in results["shapes"]] + cv = [r["speedup_vs_compile"] for r in results["shapes"] + if "speedup_vs_compile" in r] + xs = [r["exp_sharing_speedup"] for r in results["shapes"]] + + print("=" * 62) + print(f"bidirectional (fused) vs eager : " + f"{min(ev):.2f}x – {max(ev):.2f}x") + if cv: + print(f"bidirectional (fused) vs torch.compile: " + f"{min(cv):.2f}x – {max(cv):.2f}x <- the headline") + else: + print("bidirectional (fused) vs torch.compile: (not measured on this host)") + print(f"exp-sharing: fused vs two-call : " + f"{min(xs):.3f}x – {max(xs):.3f}x <- the Stage-3 win " + f"(projected ~1.7x)") + + # ---- torch.compile's COST, reported rather than hidden in a skip message. + # + # Measured over L = 128..2048: compile time is LINEAR in L, converging to + # ~0.26 s PER TIMESTEP (63s@128, 137s@512, 251s@1024, 534s@2048). The + # recurrence is a Python `for t in range(L)` loop, so inductor unrolls it into + # an L-step graph and pays per step. + # + # An earlier two-point fit (128, 512) suggested this was sub-linear and the + # claim was withdrawn; four points show that was the fixed compiler startup + # cost washing out. Extrapolating 0.26 s/step: L=8192 is ~36 min, and the + # 131k-token genomics context is ~9.5 HOURS — if it does not OOM building the + # graph first. At the lengths our applications use, torch.compile is not a + # slow baseline, it is an absent one. + # + # Publish the amortization column too: it is the number a skeptic computes + # anyway (~5,450 iterations for L>=512, stable because compile and runtime + # both scale linearly), and it is better coming from us. + comp_rows = [(r["shape"][2], r["timings"]["ref_compile_bidi"]["compile_s"], + r["timings"]["ref_compile_bidi"]["median_s"]) + for r in results["shapes"] + if "compile_s" in r["timings"].get("ref_compile_bidi", {})] + skipped = [r["shape"][2] for r in results["shapes"] + if "compile_s" not in r["timings"].get("ref_compile_bidi", {})] + if comp_rows: + print("\ntorch.compile COST (the recurrence is unrolled into an " + "L-step graph):") + print(f" {'L':>6} {'compile':>10} {'run/iter':>10} " + f"{'iters to amortize vs our kernel':>32}") + for length, cs, ms in comp_rows: + k = next(r["timings"]["bidirectional"]["median_s"] + for r in results["shapes"] if r["shape"][2] == length) + # how many calls before compile time pays for itself vs our kernel + gain = ms - k + iters = f"{cs / gain:,.0f}" if gain > 0 else "never (we are faster)" + print(f" {length:>6} {cs:>9.1f}s {ms*1e3:>9.2f}ms {iters:>32}") + results["compile_cost"] = [ + {"len": l, "compile_s": cs, "median_s": ms} + for l, cs, ms in comp_rows] + if skipped: + print(f"\n L={skipped} skipped: compile time grows with L " + f"(--compile-max-len={args.compile_max_len}). Raise the cap to " + f"measure them — and note that having to is itself the finding.") + + print("\nHeadline: the vs-torch.compile row. The exp-sharing speedup is the " + "Stage-3 fused-kernel win (one exp sweep, not two) and doubles as the " + "SS2D substrate — see BIDIRECTIONAL_LOG.md.") + + if args.json: + Path(args.json).parent.mkdir(parents=True, exist_ok=True) + Path(args.json).write_text(json.dumps(results, indent=2)) + print(f"\nresults written to {args.json}") + + +if __name__ == "__main__": + main() diff --git a/bench/bench_op.py b/bench/bench_op.py index 635b0d7..0148343 100644 --- a/bench/bench_op.py +++ b/bench/bench_op.py @@ -157,9 +157,11 @@ def main(): help="CI-sized subset (fewer shapes/reps)") ap.add_argument("--reps", type=int, default=None) ap.add_argument("--warmup", type=int, default=None) - ap.add_argument("--compile-max-len", type=int, default=512, + ap.add_argument("--compile-max-len", type=int, default=None, help="skip the torch.compile baseline beyond this L " - "(graph unrolling makes compile time explode)") + "(graph unrolling makes compile time explode). " + "Default 512, or 128 under --quick. An explicit value " + "always wins, including under --quick.") ap.add_argument("--no-compile", action="store_true") ap.add_argument("--torch-threads", type=int, default=None, help="pin torch intra-op threads (fairness control)") @@ -179,8 +181,9 @@ def main(): shapes = SUITES[args.suite] reps = args.reps or (5 if args.quick else 10) warmup = args.warmup if args.warmup is not None else (1 if args.quick else 3) - if args.quick: - args.compile_max_len = min(args.compile_max_len, 128) + if args.compile_max_len is None: + # --quick's lower cap bounds CI time only; an explicit flag must win. + args.compile_max_len = 128 if args.quick else 512 env = env_report() env["tag"] = args.tag diff --git a/bench/smoke_topologies.py b/bench/smoke_topologies.py new file mode 100644 index 0000000..dfdc9fa --- /dev/null +++ b/bench/smoke_topologies.py @@ -0,0 +1,163 @@ +"""Smoke check: do BOTH topologies run, correctly, and faster than native torch? + +A fast consolidated sanity test (not a full benchmark) that, across a sweep of +sequence lengths, verifies: + + 1. unidirectional arm_scan.selective_scan vs the vendored torch reference + 2. bidirectional arm_scan.bidirectional_scan vs the same reference, both directions + +For each (L, topology): correctness (max_abs vs an f64 reference, gated at 1e-4) +and the median speedup over native PyTorch eager. Exits non-zero if any case is +wrong — so it doubles as a CI gate that "both topologies work and beat torch +across L". + +The baseline is `selective_scan_ref`, vendored from the official +`state-spaces/mamba` repo (it matches HF transformers' `slow_forward` bit-for- +bit), so this times our kernel against **the original Mamba scan** in native +PyTorch — the genuine CPU baseline, since mamba-ssm's CUDA kernel is GPU-only. +Bidirectional is that scan run twice (forward + flipped), the standard Vim/ +Caduceus construction. These are OP-LEVEL numbers (the scan itself); a full Mamba +model has other ops this kernel does not touch, so end-to-end is more modest. + +No torch.compile here (that is what the full bench/bench_*.py are for, and it is +what OOM'd the runner at L=8192 by unrolling the recurrence into a graph — the +eager reference used here builds no graph, but the default cap stays conservative +regardless). The practical ceiling on L is the TORCH reference (an O(L) +Python-loop scan), not our kernel — which runs in constant memory at any L. + +Usage: + python bench/smoke_topologies.py # sweep 512,2048,4096 + python bench/smoke_topologies.py --lengths 512,4096,8192 --reps 3 +""" + +import argparse +import statistics +import sys +import time +from pathlib import Path + +import torch + +REPO = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO / "python")) +sys.path.insert(0, str(REPO / "tests")) +sys.path.insert(0, str(REPO / "bench")) + +import arm_scan # noqa: E402 +from reference import selective_scan_ref # noqa: E402 +from bench_op import make_inputs # noqa: E402 (shared realistic distribution) + +MAX_ABS = 1e-4 # project-wide acceptance gate + + +def flip_t(x): + return torch.flip(x, dims=(-1,)) + + +def median_ms(fn, warmup, reps): + for _ in range(warmup): + fn() + times = [] + for _ in range(reps): + t0 = time.perf_counter() + fn() + times.append(time.perf_counter() - t0) + return statistics.median(times) * 1e3 + + +def uni_ref(t, compute_dtype=torch.float32): + return selective_scan_ref( + t["u"], t["delta"], t["A"], t["B"], t["C"], D=t["D"], z=t["z"], + delta_bias=t["delta_bias"], delta_softplus=True, + compute_dtype=compute_dtype) + + +def uni_kernel(t): + return arm_scan.selective_scan( + t["u"], t["delta"], t["A"], t["B"], t["C"], D=t["D"], z=t["z"], + delta_bias=t["delta_bias"], delta_softplus=True) + + +def bidi_ref(t, compute_dtype=torch.float32): + """Both directions, merged (sum) — what a real bidirectional Mamba does.""" + fwd = uni_ref(t, compute_dtype) + back = selective_scan_ref( + flip_t(t["u"]), flip_t(t["delta"]), t["A"], flip_t(t["B"]), + flip_t(t["C"]), D=t["D"], z=flip_t(t["z"]), + delta_bias=t["delta_bias"], delta_softplus=True, + compute_dtype=compute_dtype) + return fwd + flip_t(back) + + +def bidi_kernel(t): + return arm_scan.bidirectional_scan( + t["u"], t["delta"], t["A"], t["B"], t["C"], D=t["D"], z=t["z"], + delta_bias=t["delta_bias"], delta_softplus=True, merge="sum") + + +def check(length, name, kernel_fn, ref_fn, t, warmup, reps): + """Time one topology at one length; print a table row. Return (ok, speedup).""" + # Correctness vs a TRUE f64 reference: upcast the inputs AND run the + # reference in float64. `selective_scan_ref` casts its result back to the + # input dtype and defaults to compute_dtype=float32, so passing f64 inputs + # without float64 compute would silently give an f32 "ground truth". + f64 = {k: (v.double() if torch.is_tensor(v) else v) for k, v in t.items()} + ref_f64 = ref_fn(f64, compute_dtype=torch.float64) + got = kernel_fn(t) + max_abs = (got.double() - ref_f64).abs().max().item() + + # Speed vs native torch EAGER in f32 — the honest baseline. + kern_ms = median_ms(lambda: kernel_fn(t), warmup, reps) + eager_ms = median_ms(lambda: ref_fn(t), warmup, reps) + speedup = eager_ms / kern_ms + ok = max_abs < MAX_ABS + print(f" {length:6d} {name:15s} {'PASS' if ok else 'FAIL'} " + f"{max_abs:.1e} {kern_ms:9.2f} {eager_ms:10.2f} {speedup:6.1f}x") + return ok, speedup + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--lengths", default="512,2048,4096", + help="comma-separated L values to sweep") + ap.add_argument("--dim", type=int, default=768) + ap.add_argument("--state", type=int, default=16) + ap.add_argument("--reps", type=int, default=3) + ap.add_argument("--warmup", type=int, default=1) + args = ap.parse_args() + + lengths = [int(x) for x in args.lengths.split(",")] + print(f"smoke check — B1 D{args.dim} N{args.state}, sweep L={lengths} " + f"reps={args.reps}") + print(f"kernel: {arm_scan.lib_path()}") + print("\nNote: the ceiling here is the TORCH reference (an O(L) Python-loop\n" + "scan we time against), not our kernel — which runs in constant memory\n" + "at any L. As L grows the eager baseline is what gets slow, which is\n" + "itself the point: torch is what can't keep up.\n") + print(f" {'L':>6} {'topology':15s} gate max_abs kernel(ms) " + f"eager(ms) speedup") + + uni_sp, bidi_sp, all_ok = [], [], True + with torch.no_grad(): + for length in lengths: + t = make_inputs(1, args.dim, length, args.state) + ok_u, sp_u = check(length, "unidirectional", uni_kernel, uni_ref, + t, args.warmup, args.reps) + ok_b, sp_b = check(length, "bidirectional", bidi_kernel, bidi_ref, + t, args.warmup, args.reps) + uni_sp.append(sp_u) + bidi_sp.append(sp_b) + all_ok = all_ok and ok_u and ok_b + + print() + if all_ok: + print(f"BOTH TOPOLOGIES OK across L={lengths} — vs native torch eager: " + f"unidirectional {min(uni_sp):.1f}–{max(uni_sp):.1f}x, " + f"bidirectional {min(bidi_sp):.1f}–{max(bidi_sp):.1f}x.") + else: + print("SMOKE CHECK FAILED — a topology is incorrect (see FAIL above).") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/kernel/arm-scan-core/benches/scan.rs b/kernel/arm-scan-core/benches/scan.rs index 0b85195..22174c8 100644 --- a/kernel/arm-scan-core/benches/scan.rs +++ b/kernel/arm-scan-core/benches/scan.rs @@ -81,6 +81,7 @@ fn bench_scan(crit: &mut Criterion) { z: Some(&z), delta_bias: Some(&bias), delta_softplus: true, + reverse: false, }; let mut out = vec![0.0_f32; bdl]; diff --git a/kernel/arm-scan-core/examples/profile_phases.rs b/kernel/arm-scan-core/examples/profile_phases.rs index e68c81f..07c5189 100644 --- a/kernel/arm-scan-core/examples/profile_phases.rs +++ b/kernel/arm-scan-core/examples/profile_phases.rs @@ -75,6 +75,7 @@ fn main() { z: Some(&z), delta_bias: Some(&bias), delta_softplus: true, + reverse: false, }; let mut out = vec![0.0_f32; bdl]; diff --git a/kernel/arm-scan-core/src/lib.rs b/kernel/arm-scan-core/src/lib.rs index 5938597..131fd6e 100644 --- a/kernel/arm-scan-core/src/lib.rs +++ b/kernel/arm-scan-core/src/lib.rs @@ -99,6 +99,33 @@ pub struct ScanInput<'a, T> { /// discretization). HF's slow path pre-applies softplus, so the patch /// layer passes `false`; mamba-ssm call sites pass `true`. pub delta_softplus: bool, + /// Walk the sequence backward in time: the state starts at zero at the + /// END of the sequence and accumulates toward the start. Output for + /// timestep `t` is still written at index `t`, and the pointwise D-skip + /// and z-gate still apply at index `t` — only the recurrence's traversal + /// order changes, never the layout. + /// + /// Equivalent to reversing the time axis of u/delta/b/c/z, running a forward + /// scan, and reversing the output — but without materializing any of those + /// copies. That equivalence is the definition, enforced by + /// `reverse_matches_flip_forward_flip` in `tests/property.rs` (and, + /// independently, by `tests/check_bidirectional_math.py` in numpy). + /// + /// On the scalar backend the two are **bit-identical**. On NEON they agree to + /// ~1e-7, not bit-exactly: `discretize_chunk` and `epilogue_row` process 4 + /// timesteps at a time with a scalar tail, and the vector and tail branches + /// evaluate softplus/SiLU by different means (NEON polynomial vs libm). + /// Which branch a timestep takes depends on its array POSITION, so flipping + /// the array moves timesteps across that boundary. That is a property of the + /// existing forward kernel, not of `reverse`. + /// + /// `last_state` under reverse is the state after consuming `t == 0` — the + /// state at the START of the sequence. It is not a resumable decode cache + /// the way the forward scan's `last_state` is. + /// + /// This is the 1D half of the bidirectional / 2D cross-scan topologies; see + /// TOPOLOGY_IMPLEMENTATION_PLAN.md. + pub reverse: bool, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -313,6 +340,164 @@ pub fn selective_scan_with_state( } } +/// Fused bidirectional scan: produce both the forward and backward outputs from +/// one set of inputs, computing the shared, direction-independent Pass A +/// (discretize + exp + input projection — ~85% of the work) **once** instead of +/// twice. `out_fwd`/`out_bwd` are `(batch, dim, len)`; `last_fwd`/`last_bwd` are +/// `(batch, dim, state)` and must both be `Some` or both `None`. +/// +/// Semantically equal to calling [`selective_scan`] twice (once forward, once +/// with `reverse: true`) and is checked bit-for-bit against exactly that. The +/// backward output's `last_state` is the state after consuming `t == 0` (the +/// start of the sequence). No `h0`: both directions seed from zero. +/// `input.reverse` is ignored. +/// +/// See `BIDIRECTIONAL_SPEEDUP_IDEAS.md §3.2` for why sharing Pass A is the win, +/// and `TOPOLOGY_IMPLEMENTATION_PLAN.md §3.2` — this is the 1D form of the SS2D +/// "read once, emit multiple directions" structure. +pub fn selective_scan_bidirectional( + dims: &ScanDims, + input: &ScanInput<'_, T>, + out_fwd: &mut [T], + out_bwd: &mut [T], + #[cfg_attr(not(target_arch = "aarch64"), allow(unused_mut))] mut last_fwd: Option<&mut [T]>, + #[cfg_attr(not(target_arch = "aarch64"), allow(unused_mut))] mut last_bwd: Option<&mut [T]>, + opts: ScanOptions, +) -> Result<(), ScanError> { + validate(dims, input, out_fwd, last_fwd.as_deref())?; + let bdl = dims.batch * dims.dim * dims.len; + if out_bwd.len() != bdl { + return Err(ScanError::BadLen { + tensor: "out_bwd", + expected: bdl, + got: out_bwd.len(), + }); + } + let bdn = dims.batch * dims.dim * dims.state; + if let Some(lb) = last_bwd.as_deref() { + if lb.len() != bdn { + return Err(ScanError::BadLen { + tensor: "last_bwd", + expected: bdn, + got: lb.len(), + }); + } + } + if last_fwd.is_some() != last_bwd.is_some() { + // The kernel requires both or neither; report it as a shape error + // rather than panicking in the parallel driver. + return Err(ScanError::BadLen { + tensor: "last_bwd", + expected: if last_fwd.is_some() { bdn } else { 0 }, + got: if last_bwd.is_some() { bdn } else { 0 }, + }); + } + + match opts.backend { + Backend::Scalar => { + scalar::scan_bidirectional( + dims, + input, + out_fwd, + out_bwd, + last_fwd, + last_bwd, + opts.threading, + ); + Ok(()) + } + Backend::Auto => { + #[cfg(target_arch = "aarch64")] + if try_neon_bidir( + dims, + input, + out_fwd, + out_bwd, + &mut last_fwd, + &mut last_bwd, + opts.threading, + ) { + return Ok(()); + } + scalar::scan_bidirectional( + dims, + input, + out_fwd, + out_bwd, + last_fwd, + last_bwd, + opts.threading, + ); + Ok(()) + } + Backend::Neon => { + #[cfg(target_arch = "aarch64")] + if try_neon_bidir( + dims, + input, + out_fwd, + out_bwd, + &mut last_fwd, + &mut last_bwd, + opts.threading, + ) { + return Ok(()); + } + Err(ScanError::BackendUnavailable(Backend::Neon)) + } + } +} + +/// NEON dispatch for the fused bidirectional scan — the `T == f32` analog of +/// [`try_neon`], returning `false` for other element types so the caller falls +/// back to scalar. +#[cfg(target_arch = "aarch64")] +#[allow(clippy::too_many_arguments)] +fn try_neon_bidir( + dims: &ScanDims, + input: &ScanInput<'_, T>, + out_fwd: &mut [T], + out_bwd: &mut [T], + last_fwd: &mut Option<&mut [T]>, + last_bwd: &mut Option<&mut [T]>, + threading: Threading, +) -> bool { + use core::any::TypeId; + if TypeId::of::() != TypeId::of::() { + return false; + } + fn cast(s: &[T]) -> &[f32] { + // SAFETY: caller verified T == f32; same pointer, same length. + unsafe { core::slice::from_raw_parts(s.as_ptr().cast::(), s.len()) } + } + fn cast_mut(s: &mut [T]) -> &mut [f32] { + // SAFETY: as above. + unsafe { core::slice::from_raw_parts_mut(s.as_mut_ptr().cast::(), s.len()) } + } + let input_f32 = ScanInput { + u: cast(input.u), + delta: cast(input.delta), + a: cast(input.a), + b: cast(input.b), + c: cast(input.c), + d_skip: input.d_skip.map(cast), + z: input.z.map(cast), + delta_bias: input.delta_bias.map(cast), + delta_softplus: input.delta_softplus, + reverse: input.reverse, + }; + neon::scan_bidirectional( + dims, + &input_f32, + cast_mut(out_fwd), + cast_mut(out_bwd), + last_fwd.as_deref_mut().map(cast_mut), + last_bwd.as_deref_mut().map(cast_mut), + threading, + ); + true +} + /// Route to the NEON implementation when `T` is f32. The `TypeId` check /// proves `T == f32`, making the slice reinterpretations sound. #[cfg(target_arch = "aarch64")] @@ -346,6 +531,7 @@ fn try_neon( z: input.z.map(cast), delta_bias: input.delta_bias.map(cast), delta_softplus: input.delta_softplus, + reverse: input.reverse, }; neon::scan( dims, diff --git a/kernel/arm-scan-core/src/neon/mod.rs b/kernel/arm-scan-core/src/neon/mod.rs index 78e67f4..81e2ae1 100644 --- a/kernel/arm-scan-core/src/neon/mod.rs +++ b/kernel/arm-scan-core/src/neon/mod.rs @@ -99,6 +99,7 @@ pub(crate) fn scan( bias: input.delta_bias.map_or(0.0, |v| v[d]), d_skip: input.d_skip.map(|v| v[d]), softplus: input.delta_softplus, + reverse: input.reverse, }; let a_row = &input.a[d * state..(d + 1) * state]; let bt_plane = &bt[plane..plane + len * n4]; @@ -150,12 +151,40 @@ struct Channel<'a> { bias: f32, d_skip: Option, softplus: bool, + reverse: bool, +} + +/// Walk the row's chunks in scan order: forward, or last-chunk-first when +/// reversed. Yields `(start, tlen)` — chunk boundaries in MEMORY are unchanged; +/// only the order they are visited in differs. +/// +/// Pass A is pointwise in time (no cross-timestep dependency — that is the +/// whole reason for the two-pass split), so it needs no direction awareness at +/// all. Only Pass B's serial recurrence walks time, and there `t` is a scalar +/// loop index (the vectorization is across STATE), so reversing costs nothing +/// but an index flip — no lane shuffles, no extra work. +#[inline] +fn chunks_in_scan_order(len: usize, reverse: bool) -> impl Iterator { + let n_chunks = len.div_ceil(CHUNK); + (0..n_chunks).map(move |i| { + let chunk = if reverse { n_chunks - 1 - i } else { i }; + let start = chunk * CHUNK; + (start, CHUNK.min(len - start)) + }) } /// Pass A1: dt = softplus(delta + bias), dtu = dt*u for `tlen` timesteps -/// starting at `start`, vectorized across time with a scalar tail. +/// starting at `start`, vectorized across time with a scalar tail. Writes into +/// caller-provided `dt_buf`/`dtu_buf` (both at least `tlen` long) so the single +/// and fused-bidirectional paths can share it with different scratch layouts. #[target_feature(enable = "neon")] -unsafe fn discretize_chunk(ch: &Channel<'_>, start: usize, tlen: usize, scratch: &mut Scratch) { +unsafe fn discretize_chunk( + ch: &Channel<'_>, + start: usize, + tlen: usize, + dt_buf: &mut [f32], + dtu_buf: &mut [f32], +) { let delta = ch.delta[start..start + tlen].as_ptr(); let u = ch.u[start..start + tlen].as_ptr(); let vbias = vdupq_n_f32(ch.bias); @@ -166,9 +195,9 @@ unsafe fn discretize_chunk(ch: &Channel<'_>, start: usize, tlen: usize, scratch: if ch.softplus { v = math::vsoftplusq_f32(v); } - vst1q_f32(scratch.dt.as_mut_ptr().add(t), v); + vst1q_f32(dt_buf.as_mut_ptr().add(t), v); vst1q_f32( - scratch.dtu.as_mut_ptr().add(t), + dtu_buf.as_mut_ptr().add(t), vmulq_f32(v, vld1q_f32(u.add(t))), ); t += 4; @@ -178,8 +207,8 @@ unsafe fn discretize_chunk(ch: &Channel<'_>, start: usize, tlen: usize, scratch: if ch.softplus { dt = dt.softplus(); } - scratch.dt[t] = dt; - scratch.dtu[t] = dt * *u.add(t); + dt_buf[t] = dt; + dtu_buf[t] = dt * *u.add(t); t += 1; } } @@ -259,12 +288,10 @@ unsafe fn channel_n16( }; let len = out_row.len(); - let mut start = 0; - while start < len { - let tlen = CHUNK.min(len - start); - - // Pass A1: discretization across time - discretize_chunk(ch, start, tlen, scratch); + for (start, tlen) in chunks_in_scan_order(len, ch.reverse) { + // Pass A1: discretization across time. Pointwise in t, so it is + // direction-agnostic — no `reverse` handling needed here or in A2. + discretize_chunk(ch, start, tlen, &mut scratch.dt, &mut scratch.dtu); // Pass A2: batch all exps + input projections for the chunk let abar = scratch.abar.as_mut_ptr(); @@ -295,10 +322,15 @@ unsafe fn channel_n16( vst1q_f32(bbar.add(o + 12), vmulq_f32(vld1q_f32(b.add(12)), vdtu)); } - // Pass B: pure-FMA recurrence + output dot product + // Pass B: pure-FMA recurrence + output dot product. The ONLY pass that + // walks time, hence the only one `reverse` touches. `t` is a scalar + // index here (vectorization is across the 16 state lanes, held in + // h0..h3), so flipping the direction costs one subtraction — no lane + // shuffles, no extra loads, no extra work. let abar = scratch.abar.as_ptr(); let bbar = scratch.bbar.as_ptr(); - for t in 0..tlen { + for i in 0..tlen { + let t = if ch.reverse { tlen - 1 - i } else { i }; let o = t * 16; let c = ct.as_ptr().add((start + t) * 16); @@ -313,8 +345,6 @@ unsafe fn channel_n16( acc = vfmaq_f32(acc, vld1q_f32(c.add(12)), h3); *out_row.get_unchecked_mut(start + t) = vaddvq_f32(acc); } - - start += tlen; } if let Some(ls) = last_state { @@ -350,12 +380,15 @@ unsafe fn channel_general( } let len = out_row.len(); - let mut start = 0; - while start < len { - let tlen = CHUNK.min(len - start); - discretize_chunk(ch, start, tlen, scratch); + for (start, tlen) in chunks_in_scan_order(len, ch.reverse) { + // Pointwise, direction-agnostic (see `chunks_in_scan_order`). + discretize_chunk(ch, start, tlen, &mut scratch.dt, &mut scratch.dtu); - for t in 0..tlen { + // This path keeps exp inline in the recurrence, so this single loop is + // the time walk — `step` counts scan steps, `t` is the timestep it + // consumes. (`i` below is the state-lane index, not time.) + for step in 0..tlen { + let t = if ch.reverse { tlen - 1 - step } else { step }; let vdt = vdupq_n_f32(scratch.dt[t]); let vdtu = vdupq_n_f32(scratch.dtu[t]); let b = bt.as_ptr().add((start + t) * n4); @@ -373,10 +406,323 @@ unsafe fn channel_general( } *out_row.get_unchecked_mut(start + t) = vaddvq_f32(acc); } - start += tlen; } if let Some(ls) = last_state { ls.copy_from_slice(&scratch.h_buf[..ls.len()]); } } + +// ========================================================================== +// Fused bidirectional scan (NEON). Shares Pass A — discretize + exp + input +// projection, ~85% of the work and direction-independent — between the two scan +// directions instead of computing it twice. See BIDIRECTIONAL_SPEEDUP_IDEAS.md +// §3.2. Bit-identical to two standalone scans (fwd + reverse); enforced by +// `fused_bidirectional_matches_two_scans` in tests/property.rs. +// +// The one structural cost: `abar`/`bbar` are materialized for the WHOLE row +// (`len * n4`) rather than chunk-local, so both directions can read them. At long +// L this streams from L2 rather than L1 — the exp saving must beat that round +// trip (a Stage-3 measurement, not an assumption). +// ========================================================================== + +/// Per-worker scratch for the fused path — like [`Scratch`] but `abar`/`bbar` +/// span the full row. +struct BidirScratch { + dt: Vec, // CHUNK + dtu: Vec, // CHUNK + abar: Vec, // len * n4 — FULL ROW + bbar: Vec, // len * n4 + a_pad: Vec, // n4, zero-padded A row (general path) + h_buf: Vec, // n4 (general path) +} + +impl BidirScratch { + fn new(len: usize, n4: usize) -> Self { + BidirScratch { + dt: vec![0.0; CHUNK], + dtu: vec![0.0; CHUNK], + abar: vec![0.0; len * n4], + bbar: vec![0.0; len * n4], + a_pad: vec![0.0; n4], + h_buf: vec![0.0; n4], + } + } +} + +/// Fused bidirectional entry point. B/C are transposed once (shared across +/// channels *and* directions), then each channel computes Pass A once and both +/// recurrences. +pub(crate) fn scan_bidirectional( + dims: &ScanDims, + input: &ScanInput<'_, f32>, + out_fwd: &mut [f32], + out_bwd: &mut [f32], + last_fwd: Option<&mut [f32]>, + last_bwd: Option<&mut [f32]>, + threading: Threading, +) { + let ScanDims { + batch, + dim, + len, + state, + groups, + } = *dims; + let group_size = dim / groups; + let n4 = state.div_ceil(4) * 4; + + let planes = batch * groups; + let mut bt = vec![0.0_f32; planes * len * n4]; + let mut ct = vec![0.0_f32; planes * len * n4]; + for p in 0..planes { + let src_base = p * state * len; + let dst_base = p * len * n4; + for n in 0..state { + let src_b = &input.b[src_base + n * len..src_base + (n + 1) * len]; + let src_c = &input.c[src_base + n * len..src_base + (n + 1) * len]; + for t in 0..len { + bt[dst_base + t * n4 + n] = src_b[t]; + ct[dst_base + t * n4 + n] = src_c[t]; + } + } + } + let (bt, ct) = (&bt[..], &ct[..]); + + crate::parallel::for_each_channel_bidir( + len, + state, + out_fwd, + out_bwd, + last_fwd, + last_bwd, + threading, + || BidirScratch::new(len, n4), + |scratch, ch_idx, out_fwd_row, out_bwd_row, last_f, last_b| { + let (bi, d) = (ch_idx / dim, ch_idx % dim); + let plane = (bi * groups + d / group_size) * len * n4; + let row = ch_idx * len; + let ch = Channel { + u: &input.u[row..row + len], + delta: &input.delta[row..row + len], + z: input.z.map(|z| &z[row..row + len]), + bias: input.delta_bias.map_or(0.0, |v| v[d]), + d_skip: input.d_skip.map(|v| v[d]), + softplus: input.delta_softplus, + reverse: false, // fused produces both directions; field unused + }; + let a_row = &input.a[d * state..(d + 1) * state]; + let bt_plane = &bt[plane..plane + len * n4]; + let ct_plane = &ct[plane..plane + len * n4]; + + // SAFETY: NEON is always available on aarch64. + unsafe { + if state == 16 { + channel_n16_bidir( + a_row, + bt_plane, + ct_plane, + &ch, + out_fwd_row, + out_bwd_row, + last_f, + last_b, + scratch, + ); + } else { + scratch.a_pad[..state].copy_from_slice(a_row); + channel_general_bidir( + bt_plane, + ct_plane, + &ch, + out_fwd_row, + out_bwd_row, + last_f, + last_b, + scratch, + ); + } + epilogue_row(&ch, out_fwd_row); + epilogue_row(&ch, out_bwd_row); + } + }, + ); +} + +/// Pass B (N=16) over a full-row `abar`/`bbar` in one direction. `reverse` +/// flips only the timestep index; layout is untouched, exactly as `channel_n16`. +#[target_feature(enable = "neon")] +unsafe fn pass_b_n16( + abar: *const f32, + bbar: *const f32, + ct: &[f32], + len: usize, + reverse: bool, + out_row: &mut [f32], + last: Option<&mut [f32]>, +) { + let mut h0 = vdupq_n_f32(0.0); + let mut h1 = vdupq_n_f32(0.0); + let mut h2 = vdupq_n_f32(0.0); + let mut h3 = vdupq_n_f32(0.0); + for i in 0..len { + let t = if reverse { len - 1 - i } else { i }; + let o = t * 16; + let c = ct.as_ptr().add(o); + h0 = vfmaq_f32(vld1q_f32(bbar.add(o)), vld1q_f32(abar.add(o)), h0); + h1 = vfmaq_f32(vld1q_f32(bbar.add(o + 4)), vld1q_f32(abar.add(o + 4)), h1); + h2 = vfmaq_f32(vld1q_f32(bbar.add(o + 8)), vld1q_f32(abar.add(o + 8)), h2); + h3 = vfmaq_f32(vld1q_f32(bbar.add(o + 12)), vld1q_f32(abar.add(o + 12)), h3); + + let mut acc = vmulq_f32(vld1q_f32(c), h0); + acc = vfmaq_f32(acc, vld1q_f32(c.add(4)), h1); + acc = vfmaq_f32(acc, vld1q_f32(c.add(8)), h2); + acc = vfmaq_f32(acc, vld1q_f32(c.add(12)), h3); + *out_row.get_unchecked_mut(t) = vaddvq_f32(acc); + } + if let Some(ls) = last { + let p = ls.as_mut_ptr(); + vst1q_f32(p, h0); + vst1q_f32(p.add(4), h1); + vst1q_f32(p.add(8), h2); + vst1q_f32(p.add(12), h3); + } +} + +/// N=16 fused fast path: Pass A once into full-row `abar`/`bbar`, then Pass B +/// forward and backward over them. +#[target_feature(enable = "neon")] +#[allow(clippy::too_many_arguments)] +unsafe fn channel_n16_bidir( + a_row: &[f32], + bt: &[f32], + ct: &[f32], + ch: &Channel<'_>, + out_fwd: &mut [f32], + out_bwd: &mut [f32], + last_fwd: Option<&mut [f32]>, + last_bwd: Option<&mut [f32]>, + scratch: &mut BidirScratch, +) { + let a = a_row.as_ptr(); + let a0 = vld1q_f32(a); + let a1 = vld1q_f32(a.add(4)); + let a2 = vld1q_f32(a.add(8)); + let a3 = vld1q_f32(a.add(12)); + let len = out_fwd.len(); + + // Pass A (shared): exp + input projection for the whole row, once. + let abar_mut = scratch.abar.as_mut_ptr(); + let bbar_mut = scratch.bbar.as_mut_ptr(); + let mut start = 0; + while start < len { + let tlen = CHUNK.min(len - start); + discretize_chunk(ch, start, tlen, &mut scratch.dt, &mut scratch.dtu); + for t in 0..tlen { + let vdt = vdupq_n_f32(scratch.dt[t]); + let vdtu = vdupq_n_f32(scratch.dtu[t]); + let b = bt.as_ptr().add((start + t) * 16); + let o = (start + t) * 16; + vst1q_f32( + abar_mut.add(o), + exp::vexpq_f32_nonpos_fast(vmulq_f32(vdt, a0)), + ); + vst1q_f32( + abar_mut.add(o + 4), + exp::vexpq_f32_nonpos_fast(vmulq_f32(vdt, a1)), + ); + vst1q_f32( + abar_mut.add(o + 8), + exp::vexpq_f32_nonpos_fast(vmulq_f32(vdt, a2)), + ); + vst1q_f32( + abar_mut.add(o + 12), + exp::vexpq_f32_nonpos_fast(vmulq_f32(vdt, a3)), + ); + vst1q_f32(bbar_mut.add(o), vmulq_f32(vld1q_f32(b), vdtu)); + vst1q_f32(bbar_mut.add(o + 4), vmulq_f32(vld1q_f32(b.add(4)), vdtu)); + vst1q_f32(bbar_mut.add(o + 8), vmulq_f32(vld1q_f32(b.add(8)), vdtu)); + vst1q_f32(bbar_mut.add(o + 12), vmulq_f32(vld1q_f32(b.add(12)), vdtu)); + } + start += tlen; + } + + // Pass B, both directions, over the shared products. + let abar = scratch.abar.as_ptr(); + let bbar = scratch.bbar.as_ptr(); + pass_b_n16(abar, bbar, ct, len, false, out_fwd, last_fwd); + pass_b_n16(abar, bbar, ct, len, true, out_bwd, last_bwd); +} + +/// General-N fused path: same structure, `n4` lanes with the recurrence state in +/// `h_buf` and exp materialized (unlike the single-direction `channel_general`, +/// which keeps exp inline — the fused path must store it to share it). +#[target_feature(enable = "neon")] +#[allow(clippy::too_many_arguments)] +unsafe fn channel_general_bidir( + bt: &[f32], + ct: &[f32], + ch: &Channel<'_>, + out_fwd: &mut [f32], + out_bwd: &mut [f32], + last_fwd: Option<&mut [f32]>, + last_bwd: Option<&mut [f32]>, + scratch: &mut BidirScratch, +) { + let n4 = scratch.a_pad.len(); + let len = out_fwd.len(); + + // Pass A (shared): exp + input projection for the whole row, once. + let abar_mut = scratch.abar.as_mut_ptr(); + let bbar_mut = scratch.bbar.as_mut_ptr(); + let mut start = 0; + while start < len { + let tlen = CHUNK.min(len - start); + discretize_chunk(ch, start, tlen, &mut scratch.dt, &mut scratch.dtu); + for t in 0..tlen { + let vdt = vdupq_n_f32(scratch.dt[t]); + let vdtu = vdupq_n_f32(scratch.dtu[t]); + let b = bt.as_ptr().add((start + t) * n4); + let base = (start + t) * n4; + for i in (0..n4).step_by(4) { + let a_v = vld1q_f32(scratch.a_pad.as_ptr().add(i)); + vst1q_f32( + abar_mut.add(base + i), + exp::vexpq_f32_nonpos_fast(vmulq_f32(vdt, a_v)), + ); + vst1q_f32(bbar_mut.add(base + i), vmulq_f32(vld1q_f32(b.add(i)), vdtu)); + } + } + start += tlen; + } + + // Pass B, both directions. + let abar = scratch.abar.as_ptr(); + let bbar = scratch.bbar.as_ptr(); + for (reverse, out_row, last) in [(false, out_fwd, last_fwd), (true, out_bwd, last_bwd)] { + for v in scratch.h_buf.iter_mut() { + *v = 0.0; + } + let h = scratch.h_buf.as_mut_ptr(); + for i in 0..len { + let t = if reverse { len - 1 - i } else { i }; + let base = t * n4; + let c = ct.as_ptr().add(base); + let mut acc = vdupq_n_f32(0.0); + for j in (0..n4).step_by(4) { + let h_v = vld1q_f32(h.add(j)); + let h_new = vfmaq_f32( + vld1q_f32(bbar.add(base + j)), + vld1q_f32(abar.add(base + j)), + h_v, + ); + vst1q_f32(h.add(j), h_new); + acc = vfmaq_f32(acc, vld1q_f32(c.add(j)), h_new); + } + *out_row.get_unchecked_mut(t) = vaddvq_f32(acc); + } + if let Some(ls) = last { + ls.copy_from_slice(&scratch.h_buf[..ls.len()]); + } + } +} diff --git a/kernel/arm-scan-core/src/neon/profile.rs b/kernel/arm-scan-core/src/neon/profile.rs index 711ab01..409f573 100644 --- a/kernel/arm-scan-core/src/neon/profile.rs +++ b/kernel/arm-scan-core/src/neon/profile.rs @@ -22,7 +22,7 @@ use core::arch::aarch64::*; use std::time::Instant; -use super::{exp, Channel, Scratch, CHUNK}; +use super::{exp, Channel, Scratch}; use crate::{ScanDims, ScanInput}; /// Accumulated nanoseconds per phase over a whole (sequential) scan. @@ -87,13 +87,11 @@ unsafe fn channel_profiled( let mut h3 = vdupq_n_f32(0.0); let len = out_row.len(); - let mut start = 0; - while start < len { - let tlen = CHUNK.min(len - start); - - // Pass A1: discretization across time. + for (start, tlen) in super::chunks_in_scan_order(len, ch.reverse) { + // Pass A1: discretization across time. Pointwise, so direction-agnostic + // (as are the two A2 loops below) — only Pass B walks time. let c0 = Instant::now(); - super::discretize_chunk(ch, start, tlen, scratch); + super::discretize_chunk(ch, start, tlen, &mut scratch.dt, &mut scratch.dtu); t.discretize_ns += c0.elapsed().as_nanos(); // Pass A2 (exp): ābar = exp(dt·A) for the chunk. @@ -132,11 +130,13 @@ unsafe fn channel_profiled( } t.proj_ns += c0.elapsed().as_nanos(); - // Pass B: pure-FMA recurrence + output dot product. + // Pass B: pure-FMA recurrence + output dot product. The only pass that + // walks time, so the only one `reverse` touches (mirrors channel_n16). let abar = scratch.abar.as_ptr(); let bbar = scratch.bbar.as_ptr(); let c0 = Instant::now(); - for i in 0..tlen { + for step in 0..tlen { + let i = if ch.reverse { tlen - 1 - step } else { step }; let o = i * 16; let c = ct.as_ptr().add((start + i) * 16); h0 = vfmaq_f32(vld1q_f32(bbar.add(o)), vld1q_f32(abar.add(o)), h0); @@ -151,8 +151,6 @@ unsafe fn channel_profiled( *out_row.get_unchecked_mut(start + i) = vaddvq_f32(acc); } t.recurrence_ns += c0.elapsed().as_nanos(); - - start += tlen; } let c0 = Instant::now(); @@ -218,6 +216,7 @@ pub fn scan_profiled(dims: &ScanDims, input: &ScanInput<'_, f32>, out: &mut [f32 bias: input.delta_bias.map_or(0.0, |v| v[d]), d_skip: input.d_skip.map(|v| v[d]), softplus: input.delta_softplus, + reverse: input.reverse, }; let a_row = &input.a[d * state..(d + 1) * state]; let bt_plane = &bt[plane..plane + len * n4]; diff --git a/kernel/arm-scan-core/src/parallel.rs b/kernel/arm-scan-core/src/parallel.rs index 5e05fdd..e802741 100644 --- a/kernel/arm-scan-core/src/parallel.rs +++ b/kernel/arm-scan-core/src/parallel.rs @@ -83,3 +83,97 @@ pub(crate) fn for_each_channel( #[cfg(not(feature = "parallel"))] unreachable!("parallel=true requires the `parallel` feature"); } + +/// Like [`for_each_channel`], but for the fused bidirectional scan: each channel +/// produces TWO output rows (forward and backward) and, optionally, two final +/// states. Parallelism is still across channels — each channel computes the +/// shared Pass A once and both directions' recurrences — so a rayon worker owns +/// disjoint `out_fwd`/`out_bwd`/`last_*` rows and the output is schedule- +/// independent, exactly as the single-output driver guarantees. +/// +/// `last_fwd` and `last_bwd` must both be `Some` or both `None`. +#[allow(clippy::too_many_arguments)] +pub(crate) fn for_each_channel_bidir( + len: usize, + state: usize, + out_fwd: &mut [T], + out_bwd: &mut [T], + last_fwd: Option<&mut [T]>, + last_bwd: Option<&mut [T]>, + threading: Threading, + init: I, + f: F, +) where + T: Float, + S: Send, + I: Fn() -> S + Sync + Send, + F: Fn(&mut S, usize, &mut [T], &mut [T], Option<&mut [T]>, Option<&mut [T]>) + Sync + Send, +{ + let n_channels = out_fwd.len() / len; + let parallel = match threading { + Threading::Sequential => false, + Threading::Rayon => cfg!(feature = "parallel"), + Threading::Auto => cfg!(feature = "parallel") && should_parallelize(n_channels, len, state), + }; + + // Collapse the four last-state cases to "have them" vs "don't". + let last = match (last_fwd, last_bwd) { + (Some(lf), Some(lb)) => Some((lf, lb)), + (None, None) => None, + _ => unreachable!("fused bidir: last_fwd and last_bwd must both be Some or both None"), + }; + + if !parallel { + let mut scratch = init(); + match last { + Some((lf, lb)) => { + for (i, (((of, ob), lf), lb)) in out_fwd + .chunks_exact_mut(len) + .zip(out_bwd.chunks_exact_mut(len)) + .zip(lf.chunks_exact_mut(state)) + .zip(lb.chunks_exact_mut(state)) + .enumerate() + { + f(&mut scratch, i, of, ob, Some(lf), Some(lb)); + } + } + None => { + for (i, (of, ob)) in out_fwd + .chunks_exact_mut(len) + .zip(out_bwd.chunks_exact_mut(len)) + .enumerate() + { + f(&mut scratch, i, of, ob, None, None); + } + } + } + return; + } + + #[cfg(feature = "parallel")] + { + use rayon::prelude::*; + match last { + Some((lf, lb)) => { + out_fwd + .par_chunks_exact_mut(len) + .zip_eq(out_bwd.par_chunks_exact_mut(len)) + .zip_eq(lf.par_chunks_exact_mut(state)) + .zip_eq(lb.par_chunks_exact_mut(state)) + .enumerate() + .for_each_init(&init, |s, (i, (((of, ob), lf), lb))| { + f(s, i, of, ob, Some(lf), Some(lb)) + }); + } + None => { + out_fwd + .par_chunks_exact_mut(len) + .zip_eq(out_bwd.par_chunks_exact_mut(len)) + .enumerate() + .for_each_init(&init, |s, (i, (of, ob))| f(s, i, of, ob, None, None)); + } + } + } + #[cfg(not(feature = "parallel"))] + unreachable!("parallel=true requires the `parallel` feature"); +} diff --git a/kernel/arm-scan-core/src/scalar.rs b/kernel/arm-scan-core/src/scalar.rs index 9ed0b81..a6c9358 100644 --- a/kernel/arm-scan-core/src/scalar.rs +++ b/kernel/arm-scan-core/src/scalar.rs @@ -44,11 +44,20 @@ pub(crate) fn scan( let delta_row = &input.delta[row..row + len]; let z_row = input.z.map(|z| &z[row..row + len]); + // `h0` seeds the state; `reverse` picks the traversal direction. + // The two are orthogonal: a backward scan resumed from a prior + // state is perfectly coherent (and is what SS2D will want). match h0 { Some(h0) => h.copy_from_slice(&h0[ch_idx * state..(ch_idx + 1) * state]), None => h.fill(T::ZERO), } - for (t, out_slot) in out_row.iter_mut().enumerate() { + for i in 0..len { + // The ONLY thing `reverse` changes: which timestep this step of + // the recurrence consumes. Output still lands at index `t`, and + // the pointwise D-skip / z-gate still read index `t`, so the + // layout is untouched — see `ScanInput::reverse`. + let t = if input.reverse { len - 1 - i } else { i }; + let mut dt = delta_row[t] + bias; if input.delta_softplus { dt = dt.softplus(); @@ -69,7 +78,7 @@ pub(crate) fn scan( if let Some(z) = z_row { y = y * z[t].silu(); } - *out_slot = y; + out_row[t] = y; } if let Some(ls) = last { @@ -79,6 +88,118 @@ pub(crate) fn scan( ); } +/// Per-channel scratch for the fused bidirectional scan: the shared, +/// direction-independent Pass-A products (materialized for the whole row so +/// both directions can read them) plus one recurrence state. +struct BidirScratch { + abar: Vec, // len * state — exp(dt*A), computed once + bbar: Vec, // len * state — dt*u*B, computed once + h: Vec, // state — the recurrence state, reused per direction +} + +/// Fused bidirectional scan (scalar). Computes Pass A — discretize, exp, input +/// projection — **once** per channel, then runs the recurrence in both time +/// directions over those shared products. Because Pass A is pointwise in time +/// (direction-independent) and the exp is ~85% of the work, sharing it is the +/// point: this is one exp sweep + two cheap FMA sweeps, versus two full scans. +/// +/// Output is **bit-identical** to two standalone scans — `selective_scan` with +/// `reverse: false` into `out_fwd`, and `reverse: true` into `out_bwd` — because +/// the shared products hold exactly the same values the standalone paths compute +/// inline, consumed in the same order. Enforced by +/// `fused_bidirectional_matches_two_scans` in `tests/property.rs`. +/// +/// No `h0`: bidirectional models are non-causal and seed both directions from +/// zero. `input.reverse` is ignored (both directions are produced regardless). +pub(crate) fn scan_bidirectional( + dims: &ScanDims, + input: &ScanInput<'_, T>, + out_fwd: &mut [T], + out_bwd: &mut [T], + last_fwd: Option<&mut [T]>, + last_bwd: Option<&mut [T]>, + threading: Threading, +) { + let ScanDims { + batch: _, + dim, + len, + state, + groups, + } = *dims; + let group_size = dim / groups; + + crate::parallel::for_each_channel_bidir( + len, + state, + out_fwd, + out_bwd, + last_fwd, + last_bwd, + threading, + || BidirScratch { + abar: vec![T::ZERO; len * state], + bbar: vec![T::ZERO; len * state], + h: vec![T::ZERO; state], + }, + |scratch, ch_idx, out_fwd_row, out_bwd_row, last_f, last_b| { + let (bi, d) = (ch_idx / dim, ch_idx % dim); + let a_row = &input.a[d * state..(d + 1) * state]; + let bias = input.delta_bias.map_or(T::ZERO, |v| v[d]); + let d_skip = input.d_skip.map(|v| v[d]); + let bc_base = (bi * groups + d / group_size) * state * len; + let row = ch_idx * len; + let u_row = &input.u[row..row + len]; + let delta_row = &input.delta[row..row + len]; + let z_row = input.z.map(|z| &z[row..row + len]); + + // Pass A (shared): discretize + exp + input projection, ONCE. + // Identical values to what each standalone direction computes inline. + for t in 0..len { + let mut dt = delta_row[t] + bias; + if input.delta_softplus { + dt = dt.softplus(); + } + let dt_u = dt * u_row[t]; + for (n, &a_n) in a_row.iter().enumerate() { + let bc_idx = bc_base + n * len + t; + scratch.abar[t * state + n] = (dt * a_n).exp(); + scratch.bbar[t * state + n] = dt_u * input.b[bc_idx]; + } + } + + // Pass B, both directions. `t` is the timestep consumed at step `i`; + // output still lands at index `t` (layout never flips) — matching + // scalar::scan's reverse handling exactly. + let BidirScratch { abar, bbar, h } = &mut *scratch; + for (reverse, out_row, last) in + [(false, out_fwd_row, last_f), (true, out_bwd_row, last_b)] + { + h.fill(T::ZERO); + for i in 0..len { + let t = if reverse { len - 1 - i } else { i }; + let mut y = T::ZERO; + for (n, h_n) in h.iter_mut().enumerate() { + let new = abar[t * state + n] * *h_n + bbar[t * state + n]; + *h_n = new; + y = y + input.c[bc_base + n * len + t] * new; + } + if let Some(ds) = d_skip { + y = y + ds * u_row[t]; + } + if let Some(z) = z_row { + y = y * z[t].silu(); + } + out_row[t] = y; + } + if let Some(ls) = last { + ls.copy_from_slice(h); + } + } + }, + ); +} + #[cfg(test)] mod tests { use crate::{Float, ScanDims, ScanInput}; @@ -115,6 +236,7 @@ mod tests { z: Some(&[z]), delta_bias: Some(&[bias]), delta_softplus: true, + reverse: false, }; let mut out = [0.0_f64]; let mut last = [0.0_f64; 2]; @@ -160,6 +282,7 @@ mod tests { z: None, delta_bias: None, delta_softplus: false, + reverse: false, }; let mut out = [0.0_f64; 2]; selective_scan_for_test(&dims, &input, &mut out, &mut []); diff --git a/kernel/arm-scan-core/tests/golden.rs b/kernel/arm-scan-core/tests/golden.rs index f7424cf..46b18b3 100644 --- a/kernel/arm-scan-core/tests/golden.rs +++ b/kernel/arm-scan-core/tests/golden.rs @@ -88,6 +88,7 @@ fn run_case(meta: &CaseMeta, backend: Backend) -> CaseResult { z: z.as_deref(), delta_bias: delta_bias.as_deref(), delta_softplus: meta.delta_softplus, + reverse: false, }; let mut out = vec![0.0_f32; dims.batch * dims.dim * dims.len]; let mut last = vec![0.0_f32; dims.batch * dims.dim * dims.state]; diff --git a/kernel/arm-scan-core/tests/property.rs b/kernel/arm-scan-core/tests/property.rs index 15a129e..c355ff0 100644 --- a/kernel/arm-scan-core/tests/property.rs +++ b/kernel/arm-scan-core/tests/property.rs @@ -6,8 +6,9 @@ use proptest::prelude::*; use arm_scan_core::{ - selective_scan, selective_scan_with_backend, selective_scan_with_options, - selective_scan_with_state, Backend, ScanDims, ScanInput, ScanOptions, Threading, + selective_scan, selective_scan_bidirectional, selective_scan_with_backend, + selective_scan_with_options, selective_scan_with_state, Backend, ScanDims, ScanInput, + ScanOptions, Threading, }; #[derive(Debug, Clone)] @@ -22,6 +23,7 @@ struct Case { z: Option>, delta_bias: Option>, delta_softplus: bool, + reverse: bool, } fn vecf(n: usize, lo: f32, hi: f32) -> impl Strategy> { @@ -57,41 +59,265 @@ fn case_strategy() -> impl Strategy { prop::option::of(vecf(dim, -2.0, 2.0)), // d_skip prop::option::of(vecf(bdl, -4.0, 4.0)), // z prop::option::of(vecf(dim, -6.0, 1.0)), // delta_bias - prop::bool::ANY, // delta_softplus + // paired: proptest's tuple Strategy impl stops at 10 elements, + // and this tuple is already at 10. + (prop::bool::ANY, prop::bool::ANY), // (delta_softplus, reverse) ) }) - .prop_map(|(dims, u, delta, a, b, c, d_skip, z, delta_bias, sp)| { - let mut case = Case { - dims, - u, - delta, - a, - b, - c, - d_skip, - z, - delta_bias, - delta_softplus: sp, - }; - if !case.delta_softplus { - // raw delta is the timestep: must be positive like a - // real post-softplus value - for v in &mut case.delta { - *v = v.abs() * 0.01 + 1e-3; + .prop_map( + |(dims, u, delta, a, b, c, d_skip, z, delta_bias, (sp, reverse))| { + let mut case = Case { + dims, + u, + delta, + a, + b, + c, + d_skip, + z, + delta_bias, + delta_softplus: sp, + reverse, + }; + if !case.delta_softplus { + // raw delta is the timestep: must be positive like a + // real post-softplus value + for v in &mut case.delta { + *v = v.abs() * 0.01 + 1e-3; + } + case.delta_bias = None; } - case.delta_bias = None; - } - case - }) + case + }, + ) } fn widen(v: &[f32]) -> Vec { v.iter().map(|&x| x as f64).collect() } +/// Reverse the time axis of a contiguous tensor whose LAST dim is `len` +/// (u/delta/z: (B,D,L); b/c: (B,G,N,L); out: (B,D,L) — all are rows of `len`). +fn flip_time(v: &[f32], len: usize) -> Vec { + v.chunks_exact(len) + .flat_map(|row| row.iter().rev().copied()) + .collect() +} + proptest! { #![proptest_config(ProptestConfig::with_cases(256))] + /// `reverse: true` must equal flipping the time axis of every time-varying + /// input, scanning FORWARD, and flipping the output back — the definition + /// the fused traversal exists to implement without the copies. + /// + /// TWO STRENGTHS, and the difference is the interesting part: + /// + /// * **Scalar: BIT-for-bit.** The scalar path runs one uniform code path per + /// timestep, so both routes apply identical arithmetic to identical values + /// in identical order. Any difference at all is an indexing bug, not + /// rounding — which is exactly what we want to pin down. Note the two + /// routes also land on DIFFERENT chunk boundaries (forward-on-flipped + /// splits the flipped axis; reverse splits the original), so passing + /// bit-exactly additionally proves chunking never leaks into the math. + /// + /// * **NEON: tight tolerance, NOT bit-exact — and it cannot be.** Two of the + /// NEON passes process 4 timesteps at a time with a **scalar tail**: + /// `discretize_chunk` (softplus) and `epilogue_row` (SiLU). The vector and + /// tail branches use different implementations of the same function — the + /// NEON polynomial vs libm — which agree to ~1e-7 but not bit-for-bit. + /// Which branch a timestep takes depends on its ARRAY POSITION, and + /// flipping the array moves timesteps across that boundary. So at + /// `len = 31`, timestep 29 is in the scalar tail when scanned in place and + /// in the vector body when scanned flipped — same value, ~1 ulp apart. + /// This is a property of the existing forward kernel, not of `reverse`; + /// demanding bit-equality here would be asserting something false. + /// + /// Mirrors `tests/check_bidirectional_math.py`, which proves the same + /// identity in numpy against an independently-written backward recurrence. + #[test] + fn reverse_matches_flip_forward_flip(case in case_strategy()) { + let len = case.dims.len; + let n_out = case.dims.batch * case.dims.dim * len; + let n_last = case.dims.batch * case.dims.dim * case.dims.state; + + // flipped copies: A / d_skip / delta_bias have no time axis. + let uf = flip_time(&case.u, len); + let deltaf = flip_time(&case.delta, len); + let bf = flip_time(&case.b, len); + let cf = flip_time(&case.c, len); + let zf = case.z.as_deref().map(|z| flip_time(z, len)); + + for backend in [Backend::Scalar, Backend::Auto] { + // (a) the fused backward traversal + let mut out_rev = vec![0.0_f32; n_out]; + let mut last_rev = vec![0.0_f32; n_last]; + selective_scan_with_backend( + &case.dims, + &ScanInput { + u: &case.u, delta: &case.delta, a: &case.a, b: &case.b, + c: &case.c, + d_skip: case.d_skip.as_deref(), + z: case.z.as_deref(), + delta_bias: case.delta_bias.as_deref(), + delta_softplus: case.delta_softplus, + reverse: true, + }, + &mut out_rev, + Some(&mut last_rev), + backend, + ).unwrap(); + + // (b) flip the inputs, scan forward, flip the output back + let mut out_fwd = vec![0.0_f32; n_out]; + let mut last_fwd = vec![0.0_f32; n_last]; + selective_scan_with_backend( + &case.dims, + &ScanInput { + u: &uf, delta: &deltaf, a: &case.a, b: &bf, c: &cf, + d_skip: case.d_skip.as_deref(), + z: zf.as_deref(), + delta_bias: case.delta_bias.as_deref(), + delta_softplus: case.delta_softplus, + reverse: false, + }, + &mut out_fwd, + Some(&mut last_fwd), + backend, + ).unwrap(); + let out_fwd = flip_time(&out_fwd, len); + + if backend == Backend::Scalar { + prop_assert!( + out_rev.iter().zip(&out_fwd) + .all(|(a, b)| a.to_bits() == b.to_bits()), + "scalar: reverse != flip-forward-flip (must be bit-exact), \ + dims={:?} softplus={}", + case.dims, case.delta_softplus + ); + prop_assert!( + last_rev.iter().zip(&last_fwd) + .all(|(a, b)| a.to_bits() == b.to_bits()), + "scalar: last_state differs (must be bit-exact), dims={:?}", + case.dims + ); + } else { + // Scale-relative, matching `auto_backend_matches_scalar`'s bar: + // the gap is one SIMD-vs-libm transcendental, ~1e-7. + let scale = out_rev.iter().fold(1.0_f32, |m, v| m.max(v.abs())); + for (i, (r, f)) in out_rev.iter().zip(&out_fwd).enumerate() { + let rel = (r - f).abs() / scale; + prop_assert!( + rel < 1e-5, + "neon: out[{i}] reverse={r} flip={f} rel={rel:.3e} \ + dims={:?}", case.dims + ); + } + let ls = last_rev.iter().fold(1.0_f32, |m, v| m.max(v.abs())); + for (i, (r, f)) in last_rev.iter().zip(&last_fwd).enumerate() { + let rel = (r - f).abs() / ls; + prop_assert!( + rel < 1e-5, + "neon: last_state[{i}] reverse={r} flip={f} \ + rel={rel:.3e} dims={:?}", case.dims + ); + } + } + } + } + + /// The fused bidirectional scan must be **bit-identical** to running the + /// scan twice — once forward, once reversed. That is the definition it + /// shares Pass A under: the shared exp/discretize/projection products hold + /// exactly the values each standalone direction computes inline, consumed in + /// the same order, so the result is not merely close but bit-for-bit equal. + /// + /// Checked on the SAME backend on both sides, so any difference is a fusion + /// bug, not a SIMD-vs-libm gap. Runs on `Scalar` AND `Auto` — on aarch64 the + /// latter exercises the NEON fused path against NEON two-scans; on x86 it is + /// the scalar path again (harmless). Both threadings, so the new two-output + /// parallel driver is covered (small proptest shapes keep Auto sequential). + #[test] + fn fused_bidirectional_matches_two_scans(case in case_strategy()) { + let n_out = case.dims.batch * case.dims.dim * case.dims.len; + let n_last = case.dims.batch * case.dims.dim * case.dims.state; + let mk = |reverse| ScanInput { + u: &case.u, delta: &case.delta, a: &case.a, b: &case.b, c: &case.c, + d_skip: case.d_skip.as_deref(), + z: case.z.as_deref(), + delta_bias: case.delta_bias.as_deref(), + delta_softplus: case.delta_softplus, + reverse, + }; + let bits_eq = |a: &[f32], b: &[f32]| a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits()); + + for backend in [Backend::Scalar, Backend::Auto] { + // reference: two standalone scans on this backend + let mut ref_fwd = vec![0.0_f32; n_out]; + let mut ref_fwd_last = vec![0.0_f32; n_last]; + selective_scan_with_backend( + &case.dims, &mk(false), &mut ref_fwd, Some(&mut ref_fwd_last), backend, + ).unwrap(); + let mut ref_bwd = vec![0.0_f32; n_out]; + let mut ref_bwd_last = vec![0.0_f32; n_last]; + selective_scan_with_backend( + &case.dims, &mk(true), &mut ref_bwd, Some(&mut ref_bwd_last), backend, + ).unwrap(); + + for threading in [Threading::Sequential, Threading::Rayon] { + let mut f_fwd = vec![0.0_f32; n_out]; + let mut f_bwd = vec![0.0_f32; n_out]; + let mut f_fwd_last = vec![0.0_f32; n_last]; + let mut f_bwd_last = vec![0.0_f32; n_last]; + selective_scan_bidirectional( + &case.dims, &mk(false), &mut f_fwd, &mut f_bwd, + Some(&mut f_fwd_last), Some(&mut f_bwd_last), + ScanOptions { backend, threading }, + ).unwrap(); + + prop_assert!(bits_eq(&f_fwd, &ref_fwd), "fwd differs ({backend:?}/{threading:?}) dims={:?}", case.dims); + prop_assert!(bits_eq(&f_bwd, &ref_bwd), "bwd differs ({backend:?}/{threading:?}) dims={:?}", case.dims); + prop_assert!(bits_eq(&f_fwd_last, &ref_fwd_last), "fwd last_state differs ({backend:?}/{threading:?}) dims={:?}", case.dims); + prop_assert!(bits_eq(&f_bwd_last, &ref_bwd_last), "bwd last_state differs ({backend:?}/{threading:?}) dims={:?}", case.dims); + } + } + } + + /// Guard against a vacuous pass: if `reverse` were silently ignored, the + /// equivalence test above would still hold (both sides would be forward + /// scans of... no — but the FFI/plumbing could still drop the flag). A + /// reversed scan must actually differ from a forward one on the same input. + #[test] + fn reverse_actually_reverses(case in case_strategy()) { + // L=1 is symmetric under reversal, and a constant sequence can be too. + prop_assume!(case.dims.len >= 4); + + let n_out = case.dims.batch * case.dims.dim * case.dims.len; + let mk = |reverse| ScanInput { + u: &case.u, delta: &case.delta, a: &case.a, b: &case.b, c: &case.c, + d_skip: case.d_skip.as_deref(), + z: case.z.as_deref(), + delta_bias: case.delta_bias.as_deref(), + delta_softplus: case.delta_softplus, + reverse, + }; + + let mut out_f = vec![0.0_f32; n_out]; + selective_scan(&case.dims, &mk(false), &mut out_f, None).unwrap(); + let mut out_r = vec![0.0_f32; n_out]; + selective_scan(&case.dims, &mk(true), &mut out_r, None).unwrap(); + + let scale = out_f.iter().fold(1e-6_f32, |m, v| m.max(v.abs())); + let max_rel = out_f.iter().zip(&out_r) + .map(|(a, b)| (a - b).abs() / scale) + .fold(0.0_f32, f32::max); + prop_assert!( + max_rel > 1e-4, + "reverse produced ~the forward answer (max_rel={max_rel:.3e}) — \ + is the flag being dropped? dims={:?}", case.dims + ); + } + #[test] fn f32_matches_f64(case in case_strategy()) { let n_out = case.dims.batch * case.dims.dim * case.dims.len; @@ -108,6 +334,7 @@ proptest! { z: case.z.as_deref(), delta_bias: case.delta_bias.as_deref(), delta_softplus: case.delta_softplus, + reverse: case.reverse, }, &mut out32, Some(&mut last32), @@ -129,17 +356,26 @@ proptest! { z: z.as_deref(), delta_bias: delta_bias.as_deref(), delta_softplus: case.delta_softplus, + reverse: case.reverse, }, &mut out64, None, ).unwrap(); + // f32 rounding is RELATIVE, so the bound must scale with the output + // magnitude. Proptest can produce large outputs (a 2.4 input through a + // long near-identity recurrence reaches |out| ~ 270), where an absolute + // 1e-3 bound rejects a legitimate ~4e-6 relative f32 error. Scale the + // original 1e-3 by the output magnitude (>= 1), so bounded outputs keep + // the old bound exactly and large ones are held to the same relative + // accuracy. (Surfaced by a reverse=true case; the plain-scan numerics + // are unchanged from before the fused work — see BIDIRECTIONAL_LOG.md.) + let scale = out64.iter().fold(1.0_f64, |m, v| m.max(v.abs())); for (i, (k, r)) in out32.iter().zip(out64.iter()).enumerate() { let err = (*k as f64 - r).abs(); - // f32 rounding over <=32 sequential steps with bounded values prop_assert!( - err < 1e-3, - "idx {i}: f32={k} f64={r} err={err:.3e} dims={:?}", + err < 1e-3 * scale, + "idx {i}: f32={k} f64={r} err={err:.3e} scale={scale:.3e} dims={:?}", case.dims ); prop_assert!(k.is_finite(), "non-finite output at {i}: {k}"); @@ -160,6 +396,7 @@ proptest! { z: case.z.as_deref(), delta_bias: case.delta_bias.as_deref(), delta_softplus: case.delta_softplus, + reverse: case.reverse, }; let mut out_scalar = vec![0.0_f32; n_out]; @@ -210,6 +447,7 @@ proptest! { z: case.z.as_deref(), delta_bias: case.delta_bias.as_deref(), delta_softplus: case.delta_softplus, + reverse: case.reverse, }; for backend in [Backend::Auto, Backend::Scalar] { @@ -304,6 +542,7 @@ fn streaming_matches_oneshot() { z: Some(&z), delta_bias: Some(&bias), delta_softplus: true, + reverse: false, }; for backend in [Backend::Scalar, Backend::Auto] { @@ -334,6 +573,7 @@ fn streaming_matches_oneshot() { z: Some(&z1), delta_bias: Some(&bias), delta_softplus: true, + reverse: false, }; let mut out1 = vec![0.0_f32; dim * split]; let mut mid = vec![0.0_f32; dim * state]; @@ -359,6 +599,7 @@ fn streaming_matches_oneshot() { z: Some(&z2), delta_bias: Some(&bias), delta_softplus: true, + reverse: false, }; let mut out2 = vec![0.0_f32; dim * rem]; selective_scan_with_state(&dims2, &in2, &mut out2, None, Some(&mid), opts).unwrap(); @@ -408,6 +649,7 @@ fn neon_backend_availability() { z: None, delta_bias: None, delta_softplus: false, + reverse: false, }; let mut out = [0.0_f32]; let res = selective_scan_with_backend(&dims, &input, &mut out, None, Backend::Neon); @@ -445,6 +687,7 @@ fn validation_rejects_bad_shapes() { z: None, delta_bias: None, delta_softplus: false, + reverse: false, }; assert!(selective_scan(&dims, &input, &mut out, None).is_err()); @@ -469,6 +712,7 @@ fn validation_rejects_bad_shapes() { z: None, delta_bias: None, delta_softplus: false, + reverse: false, }; assert!(selective_scan(&dims_bad, &input, &mut out3, None).is_err()); } diff --git a/kernel/arm-scan-ffi/src/lib.rs b/kernel/arm-scan-ffi/src/lib.rs index 4642af0..414aeb1 100644 --- a/kernel/arm-scan-ffi/src/lib.rs +++ b/kernel/arm-scan-ffi/src/lib.rs @@ -18,14 +18,19 @@ use std::os::raw::c_int; use arm_scan_core::{ - selective_scan_with_state, Backend, ScanDims, ScanError, ScanInput, ScanOptions, Threading, + selective_scan_bidirectional, selective_scan_with_state, Backend, ScanDims, ScanError, + ScanInput, ScanOptions, Threading, }; /// ABI version. The Python loader checks this before calling anything else. -/// Bump on any signature or semantic change to `arm_scan_selective_scan_f32`. +/// Bump on any signature or semantic change to the entry points. +/// +/// 4: `h0` (resumable initial state) and `reverse` (backward-in-time traversal) +/// were developed on separate branches, each bumping to 3. Both are in 4. +/// 5: added `arm_scan_selective_scan_bidirectional_f32` (fused two-direction). #[no_mangle] pub extern "C" fn arm_scan_abi_version() -> u32 { - 3 + 5 } /// Dimensions for a scan call. `groups` must divide `dim`. @@ -70,6 +75,10 @@ fn threading_from(v: c_int) -> Option { /// `threading`: 0 = auto, 1 = sequential, 2 = rayon. /// `delta_softplus`: nonzero to apply softplus(delta + delta_bias) inside /// the kernel. +/// `reverse`: nonzero to walk the sequence backward in time. Output layout is +/// unchanged (timestep `t` still lands at index `t`); only the recurrence's +/// traversal order flips. Equivalent to flipping the time axis of u/delta/b/c/z, +/// scanning forward, and flipping the output back — without the copies. /// /// Returns `ARM_SCAN_OK` (0) on success, a nonzero code otherwise; `out` /// contents are unspecified on error. @@ -92,6 +101,7 @@ pub unsafe extern "C" fn arm_scan_selective_scan_f32( z: *const f32, delta_bias: *const f32, delta_softplus: c_int, + reverse: c_int, backend: c_int, threading: c_int, out: *mut f32, @@ -166,6 +176,7 @@ pub unsafe extern "C" fn arm_scan_selective_scan_f32( z: opt(z, bdl), delta_bias: opt(delta_bias, d.dim), delta_softplus: delta_softplus != 0, + reverse: reverse != 0, }; let out_slice = std::slice::from_raw_parts_mut(out, bdl); let mut last_slice = if last_state.is_null() { @@ -198,6 +209,144 @@ pub unsafe extern "C" fn arm_scan_selective_scan_f32( } } +/// Fused bidirectional scan: one set of inputs, both the forward output +/// (`out_fwd`) and the backward output (`out_bwd`), computing the shared, +/// direction-independent Pass A (discretize + exp) once. Semantically equal to +/// two [`arm_scan_selective_scan_f32`] calls (forward, then `reverse=1`) but +/// sharing the ~85% exp cost; see BIDIRECTIONAL_SPEEDUP_IDEAS.md. +/// +/// No `reverse`/`h0` params: direction is inherent, and both directions seed +/// from zero. `last_fwd`/`last_bwd` are nullable but must be null together or +/// non-null together. +/// +/// # Safety +/// Same contract as [`arm_scan_selective_scan_f32`]: every non-null pointer +/// references a buffer of the element count implied by `dims`, valid for the +/// call, non-overlapping. `out_fwd`/`out_bwd` are `(batch, dim, len)`; +/// `last_fwd`/`last_bwd` are `(batch, dim, state)`. +#[no_mangle] +#[allow(clippy::too_many_arguments)] +pub unsafe extern "C" fn arm_scan_selective_scan_bidirectional_f32( + dims: *const ArmScanDims, + u: *const f32, + delta: *const f32, + a: *const f32, + b: *const f32, + c: *const f32, + d_skip: *const f32, + z: *const f32, + delta_bias: *const f32, + delta_softplus: c_int, + backend: c_int, + threading: c_int, + out_fwd: *mut f32, + out_bwd: *mut f32, + last_fwd: *mut f32, + last_bwd: *mut f32, +) -> c_int { + if dims.is_null() + || u.is_null() + || delta.is_null() + || a.is_null() + || b.is_null() + || c.is_null() + || out_fwd.is_null() + || out_bwd.is_null() + { + return ARM_SCAN_ERR_NULL_POINTER; + } + let (Some(backend), Some(threading)) = (backend_from(backend), threading_from(threading)) + else { + return ARM_SCAN_ERR_BAD_ENUM; + }; + + let d = &*dims; + let Some(bdl) = d + .batch + .checked_mul(d.dim) + .and_then(|v| v.checked_mul(d.len)) + else { + return ARM_SCAN_ERR_INVALID_DIMS; + }; + let Some(bgnl) = d + .batch + .checked_mul(d.groups) + .and_then(|v| v.checked_mul(d.state)) + .and_then(|v| v.checked_mul(d.len)) + else { + return ARM_SCAN_ERR_INVALID_DIMS; + }; + let Some(dn) = d.dim.checked_mul(d.state) else { + return ARM_SCAN_ERR_INVALID_DIMS; + }; + let Some(bdn) = d + .batch + .checked_mul(d.dim) + .and_then(|v| v.checked_mul(d.state)) + else { + return ARM_SCAN_ERR_INVALID_DIMS; + }; + + let scan_dims = ScanDims { + batch: d.batch, + dim: d.dim, + len: d.len, + state: d.state, + groups: d.groups, + }; + + let opt = |p: *const f32, n: usize| { + if p.is_null() { + None + } else { + Some(std::slice::from_raw_parts(p, n)) + } + }; + let input = ScanInput { + u: std::slice::from_raw_parts(u, bdl), + delta: std::slice::from_raw_parts(delta, bdl), + a: std::slice::from_raw_parts(a, dn), + b: std::slice::from_raw_parts(b, bgnl), + c: std::slice::from_raw_parts(c, bgnl), + d_skip: opt(d_skip, d.dim), + z: opt(z, bdl), + delta_bias: opt(delta_bias, d.dim), + delta_softplus: delta_softplus != 0, + reverse: false, // ignored by the fused path (both directions produced) + }; + let out_fwd_slice = std::slice::from_raw_parts_mut(out_fwd, bdl); + let out_bwd_slice = std::slice::from_raw_parts_mut(out_bwd, bdl); + let last_fwd_slice = if last_fwd.is_null() { + None + } else { + Some(std::slice::from_raw_parts_mut(last_fwd, bdn)) + }; + let last_bwd_slice = if last_bwd.is_null() { + None + } else { + Some(std::slice::from_raw_parts_mut(last_bwd, bdn)) + }; + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + selective_scan_bidirectional( + &scan_dims, + &input, + out_fwd_slice, + out_bwd_slice, + last_fwd_slice, + last_bwd_slice, + ScanOptions { backend, threading }, + ) + })); + + match result { + Ok(Ok(())) => ARM_SCAN_OK, + Ok(Err(ScanError::BackendUnavailable(_))) => ARM_SCAN_ERR_BACKEND_UNAVAILABLE, + Ok(Err(_)) => ARM_SCAN_ERR_INVALID_DIMS, + Err(_) => ARM_SCAN_ERR_PANIC, + } +} + #[cfg(test)] mod tests { use super::*; @@ -229,6 +378,7 @@ mod tests { 0, 0, 0, + 0, out.as_mut_ptr(), last.as_mut_ptr(), std::ptr::null(), @@ -240,6 +390,115 @@ mod tests { assert!((last[0] - h as f32).abs() < 1e-6); } + /// `reverse` across the C ABI, hand-computed. Two timesteps, N=1, no + /// softplus: a backward scan consumes t=1 first (state starts at zero at + /// the END), so out[1] sees only its own input and out[0] carries the decay + /// from t=1. Output still lands at index t — the layout never flips. + #[test] + fn ffi_reverse_two_steps() { + let dims = ArmScanDims { + batch: 1, + dim: 1, + len: 2, + state: 1, + groups: 1, + }; + let (u, dt, a, b, c) = ( + [1.0_f32, 2.0], + [0.1_f32, 0.2], + [-2.0_f32], + [1.0_f32, 3.0], + [1.0_f32, 1.0], + ); + let mut out = [0.0_f32; 2]; + let mut last = [0.0_f32; 1]; + let code = unsafe { + arm_scan_selective_scan_f32( + &dims, + u.as_ptr(), + dt.as_ptr(), + a.as_ptr(), + b.as_ptr(), + c.as_ptr(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + 0, // delta_softplus + 1, // reverse + 0, + 0, + out.as_mut_ptr(), + last.as_mut_ptr(), + std::ptr::null(), // h0: zero-initialized + ) + }; + assert_eq!(code, ARM_SCAN_OK); + + // backward: h after t=1, then after t=0 + let h1 = 0.2_f32 * 2.0 * 3.0; // dt*u*b at t=1 (state starts at 0) + let h0 = (0.1_f32 * -2.0).exp() * h1 + 0.1 * 1.0 * 1.0; + assert!((out[1] - h1).abs() < 1e-6, "out[1]={} want {h1}", out[1]); + assert!((out[0] - h0).abs() < 1e-6, "out[0]={} want {h0}", out[0]); + // last_state under reverse is the state after consuming t == 0 + assert!((last[0] - h0).abs() < 1e-6); + } + + /// Fused bidirectional across the C ABI, hand-computed. Two steps, N=1: the + /// forward output is the ordinary scan; the backward output must match the + /// `reverse` case above. Both from one call sharing Pass A. + #[test] + fn ffi_bidirectional_two_steps() { + let dims = ArmScanDims { + batch: 1, + dim: 1, + len: 2, + state: 1, + groups: 1, + }; + let (u, dt, a, b, c) = ( + [1.0_f32, 2.0], + [0.1_f32, 0.2], + [-2.0_f32], + [1.0_f32, 3.0], + [1.0_f32, 1.0], + ); + let mut out_fwd = [0.0_f32; 2]; + let mut out_bwd = [0.0_f32; 2]; + let code = unsafe { + arm_scan_selective_scan_bidirectional_f32( + &dims, + u.as_ptr(), + dt.as_ptr(), + a.as_ptr(), + b.as_ptr(), + c.as_ptr(), + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + 0, // delta_softplus + 0, // backend = auto + 0, // threading = auto + out_fwd.as_mut_ptr(), + out_bwd.as_mut_ptr(), + std::ptr::null_mut(), // last_fwd + std::ptr::null_mut(), // last_bwd + ) + }; + assert_eq!(code, ARM_SCAN_OK); + + // forward: h at t=0, then t=1 + let f0 = 0.1_f32 * 1.0 * 1.0; // dt*u*b at t=0 + let f1 = (0.2_f32 * -2.0).exp() * f0 + 0.2 * 2.0 * 3.0; + assert!((out_fwd[0] - f0).abs() < 1e-6, "out_fwd[0]={}", out_fwd[0]); + assert!((out_fwd[1] - f1).abs() < 1e-6, "out_fwd[1]={}", out_fwd[1]); + + // backward: h at t=1, then t=0 (matches ffi_reverse_two_steps) + let b1 = 0.2_f32 * 2.0 * 3.0; + let b0 = (0.1_f32 * -2.0).exp() * b1 + 0.1 * 1.0 * 1.0; + assert!((out_bwd[1] - b1).abs() < 1e-6, "out_bwd[1]={}", out_bwd[1]); + assert!((out_bwd[0] - b0).abs() < 1e-6, "out_bwd[0]={}", out_bwd[0]); + } + #[test] fn ffi_rejects_null_and_bad_enum() { let dims = ArmScanDims { @@ -265,6 +524,7 @@ mod tests { 0, 0, 0, + 0, out.as_mut_ptr(), std::ptr::null_mut(), std::ptr::null(), @@ -284,7 +544,8 @@ mod tests { std::ptr::null(), std::ptr::null(), 0, - 7, + 0, + 7, // bad backend enum 0, out.as_mut_ptr(), std::ptr::null_mut(), @@ -326,7 +587,8 @@ mod tests { null(), null(), null(), - 0, + 0, // delta_softplus + 0, // reverse 0, 0, out_full.as_mut_ptr(), @@ -357,7 +619,8 @@ mod tests { null(), null(), null(), - 0, + 0, // delta_softplus + 0, // reverse 0, 0, out1.as_mut_ptr(), @@ -380,7 +643,8 @@ mod tests { null(), null(), null(), - 0, + 0, // delta_softplus + 0, // reverse 0, 0, out2.as_mut_ptr(), diff --git a/python/arm_scan/__init__.py b/python/arm_scan/__init__.py index 3477f5a..ef9827e 100644 --- a/python/arm_scan/__init__.py +++ b/python/arm_scan/__init__.py @@ -7,6 +7,7 @@ arm_scan.stats() # confirm the kernel actually ran Direct op: arm_scan.selective_scan(u, delta, A, B, C, D=..., z=...) +Bidirectional: arm_scan.bidirectional_scan(...) (both time directions, merged) NumPy-only: arm_scan.selective_scan_numpy(...) (no torch required) """ @@ -17,6 +18,7 @@ "selective_scan_numpy", "lib_path", "selective_scan", + "bidirectional_scan", "patch", "unpatch", "stats", @@ -31,6 +33,9 @@ def __getattr__(name): if name == "selective_scan": return importlib.import_module(".op", __name__).selective_scan + if name == "bidirectional_scan": + return importlib.import_module( + ".bidirectional", __name__).bidirectional_scan if name in ("patch", "unpatch", "stats"): return getattr(importlib.import_module(".patch", __name__), name) raise AttributeError(f"module 'arm_scan' has no attribute '{name}'") diff --git a/python/arm_scan/_ffi.py b/python/arm_scan/_ffi.py index 52298a4..f832c6a 100644 --- a/python/arm_scan/_ffi.py +++ b/python/arm_scan/_ffi.py @@ -11,7 +11,10 @@ import sys from pathlib import Path -ABI_VERSION = 3 +# 3: `h0` (initial state) and `reverse` (backward traversal) landed on separate +# branches, each claiming 3. Reconciled at merge -> both are in ABI 4. +# 5: fused bidirectional entry point (arm_scan_selective_scan_bidirectional_f32). +ABI_VERSION = 5 _LIB_NAMES = { "win32": ["arm_scan_ffi.dll"], @@ -80,12 +83,25 @@ def load(): ctypes.POINTER(ArmScanDims), *([ctypes.c_void_p] * 8), # u delta a b c d_skip z delta_bias ctypes.c_int, # delta_softplus + ctypes.c_int, # reverse ctypes.c_int, # backend ctypes.c_int, # threading ctypes.c_void_p, # out ctypes.c_void_p, # last_state ctypes.c_void_p, # h0 ] + lib.arm_scan_selective_scan_bidirectional_f32.restype = ctypes.c_int + lib.arm_scan_selective_scan_bidirectional_f32.argtypes = [ + ctypes.POINTER(ArmScanDims), + *([ctypes.c_void_p] * 8), # u delta a b c d_skip z delta_bias + ctypes.c_int, # delta_softplus + ctypes.c_int, # backend + ctypes.c_int, # threading + ctypes.c_void_p, # out_fwd + ctypes.c_void_p, # out_bwd + ctypes.c_void_p, # last_fwd + ctypes.c_void_p, # last_bwd + ] _lib, _lib_path = lib, path return lib raise OSError( @@ -101,16 +117,23 @@ def lib_path(): def scan_raw(dims, ptr_u, ptr_delta, ptr_a, ptr_b, ptr_c, ptr_d_skip, ptr_z, ptr_delta_bias, delta_softplus, backend, threading, ptr_out, - ptr_last, ptr_h0=0): + ptr_last, ptr_h0=0, *, reverse=False): """Thin call-through. Pointers are integer addresses; 0 means null. ``ptr_h0`` is the optional initial SSM state (batch, dim, state); 0 seeds - the recurrence from zeros (the default one-shot behavior).""" + the recurrence from zeros (the default one-shot behavior). + + ``reverse`` walks the sequence backward in time. It is keyword-only on + purpose: every caller here passes positionally, and it sits mid-signature in + C (right after ``delta_softplus``), so accepting it positionally would let a + caller silently shift ``backend`` into it. Python order need not match C. + """ lib = load() code = lib.arm_scan_selective_scan_f32( ctypes.byref(dims), ptr_u, ptr_delta, ptr_a, ptr_b, ptr_c, ptr_d_skip or None, ptr_z or None, ptr_delta_bias or None, - int(bool(delta_softplus)), BACKENDS[backend], THREADING[threading], + int(bool(delta_softplus)), int(bool(reverse)), + BACKENDS[backend], THREADING[threading], ptr_out, ptr_last or None, ptr_h0 or None, ) if code != 0: @@ -118,3 +141,23 @@ def scan_raw(dims, ptr_u, ptr_delta, ptr_a, ptr_b, ptr_c, ptr_d_skip, ptr_z, f"arm_scan kernel error {code}: " f"{ERROR_NAMES.get(code, 'unknown')}" ) + + +def scan_bidir_raw(dims, ptr_u, ptr_delta, ptr_a, ptr_b, ptr_c, ptr_d_skip, + ptr_z, ptr_delta_bias, delta_softplus, backend, threading, + ptr_out_fwd, ptr_out_bwd, ptr_last_fwd=0, ptr_last_bwd=0): + """Fused bidirectional call-through. Produces both directions from one set + of inputs, sharing the exp. `ptr_last_*` are optional (0 = null); pass both + or neither. Pointers are integer addresses; 0 means null.""" + lib = load() + code = lib.arm_scan_selective_scan_bidirectional_f32( + ctypes.byref(dims), ptr_u, ptr_delta, ptr_a, ptr_b, ptr_c, + ptr_d_skip or None, ptr_z or None, ptr_delta_bias or None, + int(bool(delta_softplus)), BACKENDS[backend], THREADING[threading], + ptr_out_fwd, ptr_out_bwd, ptr_last_fwd or None, ptr_last_bwd or None, + ) + if code != 0: + raise RuntimeError( + f"arm_scan kernel error {code}: " + f"{ERROR_NAMES.get(code, 'unknown')}" + ) diff --git a/python/arm_scan/bidirectional.py b/python/arm_scan/bidirectional.py new file mode 100644 index 0000000..b41e6f7 --- /dev/null +++ b/python/arm_scan/bidirectional.py @@ -0,0 +1,175 @@ +"""Bidirectional selective scan (TOPOLOGY_IMPLEMENTATION_PLAN.md §2). + +Runs the recurrence over the sequence in both time directions and merges the +two outputs. For the common (tied-weights) case it uses the **fused two-direction +kernel** (`selective_scan_bidirectional`): one call that produces both directions +while computing the direction-independent Pass A — discretize + exp, ~85% of the +work — **once** instead of twice. See BIDIRECTIONAL_SPEEDUP_IDEAS.md §3.2 and +BIDIRECTIONAL_LOG.md. + +It falls back to two separate `selective_scan` calls (the backward one via the +`reverse` flag, no copies) when the fused path does not apply: **untied weights** +(`reverse_params`, where the two directions use different A/delta/B/C so Pass A +cannot be shared) or when a **`last_state`** is requested (the fused op does not +produce one — bidirectional models are non-causal and ignore it). + +WHICH BIDIRECTIONAL MODELS THIS IS FOR +-------------------------------------- +Two patterns exist in the wild and they are not interchangeable: + + "outer" (Caduceus's BiMambaWrapper, Vim): the WHOLE mixer — causal conv, + x_proj, dt_proj — is re-run on `x.flip(time)`. A causal conv over + flipped input is not the flip of the conv over input, so the two + directions' scan inputs are genuinely different tensors. Such a model + does not need this module: it already calls `selective_scan` twice, and + both calls are ordinary FORWARD scans. A kernel `reverse` flag buys it + nothing. + + "inner" (VMamba/SS2D-style cross-scan, and bidirectional variants that flip + after the projections): the SAME projected tensors are traversed in both + time directions. Flipping commutes with the time-pointwise projections, + so the backward direction's (u, delta, B, C) are exactly the flips of the + forward direction's. THIS is the pattern this module implements, and the + one a fused `reverse` flag actually accelerates — it is also the 1D case + of the SS2D cross-scan. + +Check which one your checkpoint is before wiring this in. Getting it wrong +produces plausible-looking output that is quietly wrong. + +MERGE +----- +`merge="sum"` matches the common case. Note that for any LINEAR merge (sum, +mean), gating inside each direction with `z` is algebraically identical to +gating once after the merge: + + fwd: y_f[t]·silu(z[t]); bwd: y_b[t]·silu(z[t]) + sum: (y_f[t] + y_b[t])·silu(z[t]) + +so we let the kernel apply the gate in both passes and do not special-case it. +(The reverse scan reads `z` at index t like everything else — `reverse` changes +traversal order, never layout — so no flipping is involved on either side.) +For anything non-linear (a learned gated combine), pass `merge="none"` and do +the combination yourself on the two returned tensors — the primitive should not +guess at a model-specific merge. + +GOTCHA: with `merge="sum"` and a shared `D`, the skip connection is applied in +BOTH directions, so the merged output carries 2·D·u, not D·u. That is what real +bidirectional Mambas do (each direction's mixer applies its own D, then the +outputs are summed), but it surprises people — so it is pinned by an assertion +in `tests/check_bidirectional_math.py`. If your model wants D counted once, +pass `D=None` here and add the skip yourself after merging. +""" + +import torch + +from .op import selective_scan, selective_scan_bidirectional + +_MERGES = ("sum", "mean", "concat", "none") + + +def _scan_reverse(u, delta, A, B, C, D, z, delta_bias, delta_softplus, + return_last_state): + """Scan the sequence backward in time — fused, no copies. + + THE SEAM (now closed). This used to flip the five time-varying inputs, run + the forward kernel, and flip the output back — correct, but six full-tensor + copies per call. The kernel now walks the sequence backward in place, so the + whole thing is one call and the copies are gone. Nothing above this function + changed. + + Kept as a named function rather than inlined because it is the documented + seam, and because the flip-based definition it replaces is still the + specification: `reverse=True` is *defined* as flip-forward-flip, enforced + bit-for-bit by `reverse_matches_flip_forward_flip` in the Rust property + tests and by `tests/check_bidirectional_math.py` in numpy. + + `last_state` here is the state after consuming t=0 (the last step a + backward scan sees) — i.e. the state at the START of the sequence. It is + not a resumable cache the way the forward scan's last_state is; + bidirectional models are non-causal and generally ignore it. + """ + return selective_scan( + u, delta, A, B, C, D=D, z=z, delta_bias=delta_bias, + delta_softplus=delta_softplus, return_last_state=return_last_state, + reverse=True, + ) + + +def bidirectional_scan(u, delta, A, B, C, D=None, z=None, delta_bias=None, + delta_softplus=False, merge="sum", + reverse_params=None, return_last_state=False): + """Scan `u` forward and backward in time and merge the two outputs. + + Tensor layouts are exactly `arm_scan.selective_scan`'s: + u, delta, z: (batch, dim, len); A: (dim, state); + B, C: (batch, state, len) or (batch, groups, state, len); + D, delta_bias: (dim,). + + merge: + "sum" out_fwd + out_bwd (the common case) + "mean" (out_fwd + out_bwd) / 2 + "concat" cat along the channel axis -> (batch, 2*dim, len) + "none" return (out_fwd, out_bwd) unmerged, for a learned/gated combine + + reverse_params: for models whose backward direction has its OWN weights + (Vim's `bimamba_type="v2"` has a separate A_b, D_b, dt_proj_b), pass a + dict overriding any of "A", "D", "delta", "delta_bias", "B", "C" for the + backward pass. Supply time-varying overrides (delta/B/C) in ordinary + forward-time order, exactly as you would for the forward pass — the kernel + reverses the traversal, not the layout, so nothing is ever pre-flipped. + Default (None) is the weight-tied case: the backward pass reuses + A/D/delta_bias and the same forward tensors. + + Returns out (batch, dim, len) — or (batch, 2*dim, len) for "concat", or a + 2-tuple for "none". With return_last_state=True, each output is paired with + its last_state; see `_scan_reverse` on what the backward one means. + """ + if merge not in _MERGES: + raise ValueError(f"merge must be one of {_MERGES}, got {merge!r}") + + # The fused kernel shares Pass A (the exp) between directions, which is only + # valid when both directions use the SAME A/delta/B/C — i.e. tied weights. + # It also does not produce last_state (bidirectional models are non-causal + # and ignore it). So the fused fast path handles the common case; untied + # weights or an explicit last_state request fall back to two scans. + if reverse_params is None and not return_last_state: + out_f, out_b = selective_scan_bidirectional( + u, delta, A, B, C, D=D, z=z, delta_bias=delta_bias, + delta_softplus=delta_softplus, + ) + else: + p = reverse_params or {} + rev = _scan_reverse( + u, + p.get("delta", delta), + p.get("A", A), + p.get("B", B), + p.get("C", C), + p.get("D", D), + z, + p.get("delta_bias", delta_bias), + delta_softplus, + return_last_state, + ) + fwd = selective_scan( + u, delta, A, B, C, D=D, z=z, delta_bias=delta_bias, + delta_softplus=delta_softplus, return_last_state=return_last_state, + ) + if return_last_state: + (out_f, last_f), (out_b, last_b) = fwd, rev + else: + out_f, out_b = fwd, rev + + if merge == "sum": + out = out_f + out_b + elif merge == "mean": + out = (out_f + out_b) * 0.5 + elif merge == "concat": + out = torch.cat((out_f, out_b), dim=1) + else: # "none" + out = (out_f, out_b) + + if not return_last_state: + return out + return (out, (last_f, last_b)) if merge != "none" else ( + (out_f, last_f), (out_b, last_b)) diff --git a/python/arm_scan/numpy_api.py b/python/arm_scan/numpy_api.py index 752f9e4..38dabe7 100644 --- a/python/arm_scan/numpy_api.py +++ b/python/arm_scan/numpy_api.py @@ -15,13 +15,17 @@ def _prep(x, name, shape=None): def selective_scan_numpy(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, return_last_state=False, - backend="auto", threading="auto", initial_state=None): + backend="auto", threading="auto", initial_state=None, + reverse=False): """Selective scan on numpy arrays (see kernel docs for semantics). u, delta: (batch, dim, len); A: (dim, state); B, C: (batch, state, len) or (batch, groups, state, len); D, delta_bias: (dim,); z: (batch, dim, len). Returns out (batch, dim, len) [, last_state (batch, dim, state)]. + + reverse: walk the sequence backward in time; output layout unchanged + (timestep t still lands at index t). See `arm_scan.selective_scan`. """ u = _prep(u, "u") if u.ndim != 3: @@ -61,6 +65,6 @@ def selective_scan_numpy(u, delta, A, B, C, D=None, z=None, delta_bias=None, _ffi.scan_raw( dims, ptr(u), ptr(delta), ptr(A), ptr(B), ptr(C), ptr(D), ptr(z), ptr(delta_bias), delta_softplus, backend, threading, - out.ctypes.data, last.ctypes.data, ptr_h0=ptr(h0), + out.ctypes.data, last.ctypes.data, ptr_h0=ptr(h0), reverse=reverse, ) return (out, last) if return_last_state else out diff --git a/python/arm_scan/op.py b/python/arm_scan/op.py index c9e92b0..85a7681 100644 --- a/python/arm_scan/op.py +++ b/python/arm_scan/op.py @@ -32,6 +32,7 @@ def _selective_scan_op( delta_bias: Optional[torch.Tensor], h0: Optional[torch.Tensor], delta_softplus: bool, + reverse: bool, ) -> Tuple[torch.Tensor, torch.Tensor]: batch, dim, length = u.shape state = a.shape[1] @@ -45,21 +46,23 @@ def _selective_scan_op( dims, u.data_ptr(), delta.data_ptr(), a.data_ptr(), b.data_ptr(), c.data_ptr(), ptr(d_skip), ptr(z), ptr(delta_bias), delta_softplus, "auto", "auto", out.data_ptr(), - last_state.data_ptr(), ptr_h0=ptr(h0), + last_state.data_ptr(), ptr_h0=ptr(h0), reverse=reverse, ) _CALLS["n"] += 1 return out, last_state @_selective_scan_op.register_fake -def _(u, delta, a, b, c, d_skip, z, delta_bias, h0, delta_softplus): +def _(u, delta, a, b, c, d_skip, z, delta_bias, h0, delta_softplus, reverse): + # Neither `h0` nor `reverse` changes the output shapes: h0 seeds the state, + # reverse changes traversal order. Must mirror the op signature exactly. return torch.empty_like(u), u.new_empty( (u.shape[0], u.shape[1], a.shape[1])) def selective_scan(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, return_last_state=False, - initial_state=None): + initial_state=None, reverse=False): """Selective scan on CPU float32 torch tensors. u, delta, z: (batch, dim, len); A: (dim, state); @@ -71,6 +74,13 @@ def selective_scan(u, delta, A, B, C, D=None, z=None, delta_bias=None, Tensors are made contiguous/f32 here, so callers can pass transposed views directly. + + reverse: walk the sequence backward in time. The output layout is + unchanged — timestep t still lands at index t — so this is NOT the same as + reversing the output; it is exactly equivalent to flipping the time axis of + u/delta/B/C/z, scanning forward, and flipping the result back, minus the + copies. Under reverse, `last_state` is the state after consuming t=0 (the + START of the sequence), so it is not a resumable decode cache. """ batch, dim, length = u.shape state = A.shape[1] @@ -85,10 +95,71 @@ def selective_scan(u, delta, A, B, C, D=None, z=None, delta_bias=None, None if delta_bias is None else _c(delta_bias), None if initial_state is None else _c(initial_state), delta_softplus, + reverse, ) return (out, last_state) if return_last_state else out +@torch.library.custom_op("arm_scan::selective_scan_bidirectional", mutates_args=()) +def _selective_scan_bidir_op( + u: torch.Tensor, + delta: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + d_skip: Optional[torch.Tensor], + z: Optional[torch.Tensor], + delta_bias: Optional[torch.Tensor], + delta_softplus: bool, +) -> Tuple[torch.Tensor, torch.Tensor]: + batch, dim, length = u.shape + state = a.shape[1] + + out_fwd = torch.empty_like(u) + out_bwd = torch.empty_like(u) + dims = _ffi.ArmScanDims(batch, dim, length, state, + b.shape[1] if b.dim() == 4 else 1) + ptr = lambda t: 0 if t is None else t.data_ptr() + _ffi.scan_bidir_raw( + dims, u.data_ptr(), delta.data_ptr(), a.data_ptr(), b.data_ptr(), + c.data_ptr(), ptr(d_skip), ptr(z), ptr(delta_bias), + delta_softplus, "auto", "auto", + out_fwd.data_ptr(), out_bwd.data_ptr(), + ) + _CALLS["n"] += 1 + return out_fwd, out_bwd + + +@_selective_scan_bidir_op.register_fake +def _(u, delta, a, b, c, d_skip, z, delta_bias, delta_softplus): + return torch.empty_like(u), torch.empty_like(u) + + +def selective_scan_bidirectional(u, delta, A, B, C, D=None, z=None, + delta_bias=None, delta_softplus=False): + """Fused bidirectional scan on CPU float32 torch tensors. + + Same tensor layouts as `selective_scan`; returns (out_fwd, out_bwd) — the + forward and backward scans — computed in one call sharing the direction- + independent Pass A (exp). Equal to `selective_scan` forward + `reverse=True`, + but ~1.7x cheaper by not recomputing the exp. Merging the two directions is + the caller's job (see `arm_scan.bidirectional_scan`). + """ + batch, dim, length = u.shape + state = A.shape[1] + if B.dim() == 3: + B = B.reshape(batch, 1, state, length) + if C.dim() == 3: + C = C.reshape(batch, 1, state, length) + return _selective_scan_bidir_op( + _c(u), _c(delta), _c(A), _c(B), _c(C), + None if D is None else _c(D), + None if z is None else _c(z), + None if delta_bias is None else _c(delta_bias), + delta_softplus, + ) + + def kernel_calls() -> int: """How many times the native kernel has been invoked (engagement check for tests and benchmarks).""" diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..280620a --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,22 @@ +# Full development tier: everything above, plus PyTorch and HF transformers. +# +# pip install -r requirements-dev.txt +# +# Needed for the parts of the repo the torch-free tier cannot reach: +# torch python/arm_scan/op.py (the torch.library custom op) and +# patch.py; tests/check_bidirectional.py; tests/gen_golden.py; +# bench/bench_op.py; bench/bench_bidirectional.py +# transformers bench/bench_e2e.py and tests/check_hf_patch.py +# (mamba-130m-hf; imported lazily, weights cached on first run) +# +# CPU-only torch is all this project ever uses. On Linux/Windows the default +# PyPI wheel bundles CUDA; if you want to skip that download: +# pip install torch --index-url https://download.pytorch.org/whl/cpu +# +# NOTE: none of this is required to develop the kernel or to check the +# bidirectional/topology math — see requirements.txt. Install it when you need +# to run the op-level or end-to-end benchmarks locally rather than in CI. +-r requirements.txt + +torch>=2.4 +transformers>=4.44 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..20519e0 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,18 @@ +# Torch-free tier — what the current .venv has, and all CI's `test` job installs. +# +# python -m venv .venv +# source .venv/Scripts/activate # Windows; .venv/bin/activate on Linux/macOS +# pip install -r requirements.txt +# +# Enough to run, with no Rust toolchain and no PyTorch: +# tests/check_bidirectional_math.py the bidirectional definition/spec check +# tests/verify_golden.py independent re-derivation of the goldens +# +# and, once the kernel cdylib is built (`cd kernel && cargo build --release -p arm-scan-ffi`): +# tests/check_ffi.py goldens through the real C ABI +# arm_scan.selective_scan_numpy(...) the torch-free kernel API +# +# For the op/patch/benchmark tiers (PyTorch, HF transformers) see +# requirements-dev.txt — kept separate deliberately, since torch is a multi- +# hundred-MB dependency that the torch-free path never needs. +numpy>=1.24 diff --git a/tests/check_bidirectional.py b/tests/check_bidirectional.py new file mode 100644 index 0000000..d3c2e45 --- /dev/null +++ b/tests/check_bidirectional.py @@ -0,0 +1,235 @@ +"""Correctness gate for the bidirectional scan (TOPOLOGY_IMPLEMENTATION_PLAN.md §2.3). + +Ground truth is the VENDORED reference (`tests/reference/selective_scan_ref.py`) +run at float64 on explicitly flipped inputs — never our own kernel. That keeps +this an independent check of `arm_scan.bidirectional_scan` rather than a +tautology, exactly as `check_ffi.py` does for the 1D op. + +Acceptance is the project-wide gate: max_abs(kernel_f32 - reference_f64) < 1e-4. + +Usage: + cargo build --release -p arm-scan-ffi # (in kernel/) + python tests/check_bidirectional.py +""" + +import sys +from pathlib import Path + +import torch + +REPO = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO / "python")) +sys.path.insert(0, str(REPO / "tests")) + +import arm_scan # noqa: E402 +from reference import selective_scan_ref # noqa: E402 + +MAX_ABS = 1e-4 + + +def _flip(t): + return None if t is None else torch.flip(t, dims=(-1,)) + + +def reference_bidirectional(u, delta, A, B, C, D=None, z=None, + delta_bias=None, delta_softplus=False, + merge="sum", reverse_params=None): + """Bidirectional scan computed in float64 through the vendored reference. + + The backward direction is *defined* as: flip every time-varying tensor, + run the ordinary forward reference, flip the output back. This is the + definition `arm_scan.bidirectional_scan` claims to implement, computed + independently of it. + + The INPUTS are upcast to double — `selective_scan_ref` casts its result + back to the input dtype on the way out, so passing f32 with + compute_dtype=float64 would silently hand back an f32 answer and make this + a far weaker gate than it claims. `gen_golden.py` upcasts for the same + reason; this mirrors it. + """ + p = reverse_params or {} + f64 = lambda t: None if t is None else t.double() + kw = dict(delta_softplus=delta_softplus, compute_dtype=torch.float64) + + out_f = selective_scan_ref( + f64(u), f64(delta), f64(A), f64(B), f64(C), + D=f64(D), z=f64(z), delta_bias=f64(delta_bias), **kw) + + out_b = selective_scan_ref( + f64(_flip(u)), f64(_flip(p.get("delta", delta))), f64(p.get("A", A)), + f64(_flip(p.get("B", B))), f64(_flip(p.get("C", C))), + D=f64(p.get("D", D)), z=f64(_flip(z)), + delta_bias=f64(p.get("delta_bias", delta_bias)), **kw) + out_b = _flip(out_b) + + if merge == "sum": + return out_f + out_b + if merge == "mean": + return (out_f + out_b) * 0.5 + if merge == "concat": + return torch.cat((out_f, out_b), dim=1) + return out_f, out_b + + +def make_case(batch, dim, length, state, groups=1, seed=0, **opts): + g = torch.Generator().manual_seed(seed) + r = lambda *s: torch.randn(*s, generator=g, dtype=torch.float32) + softplus = opts.get("delta_softplus", True) + + # The kernel's Pass-A2 exp is `vexpq_f32_nonpos`, specialized for the + # scan's always-non-positive argument dt*A. That holds because A < 0 and + # the timestep dt >= 0 — and dt >= 0 is only guaranteed if delta is either + # raw (softplus applied inside the kernel) or ALREADY positive. So when + # delta_softplus=False, delta *is* the timestep and must be drawn positive, + # exactly as gen_golden.py does for its own no_softplus case. Feeding a + # negative delta here violates the kernel's documented precondition and is + # a regime no real Mamba ever produces (HF's slow path pre-applies + # softplus). See BIDIRECTIONAL_LOG.md. + if softplus: + delta = r(batch, dim, length) # raw, any sign + else: + delta = (torch.rand(batch, dim, length, generator=g, + dtype=torch.float32) * 0.099) + 1e-3 # (0.001, 0.1] + + case = dict( + u=r(batch, dim, length), + delta=delta, + # A must be negative (the model parameterizes it as -exp(A_log)). + A=-torch.rand(dim, state, generator=g, dtype=torch.float32) - 0.1, + B=r(batch, groups, state, length), + C=r(batch, groups, state, length), + ) + if opts.get("D"): + case["D"] = r(dim) + if opts.get("z"): + case["z"] = r(batch, dim, length) + if opts.get("delta_bias"): + if not softplus: + # a bias could push the raw positive timestep negative, breaking + # the same precondition; gen_golden.py forces bias=None here too. + raise ValueError("delta_bias with delta_softplus=False would " + "violate the dt >= 0 precondition") + case["delta_bias"] = r(dim) + case["delta_softplus"] = softplus + return case + + +CASES = [ + # (name, case kwargs, bidirectional_scan kwargs) + ("tiny_sum", dict(batch=1, dim=4, length=8, state=16, seed=1), {}), + ("full_opts_sum", + dict(batch=2, dim=8, length=32, state=16, seed=2, + D=True, z=True, delta_bias=True), + {}), + ("no_softplus", + dict(batch=1, dim=4, length=16, state=16, seed=3, delta_softplus=False), + {}), + ("state13_neon_tail", + dict(batch=1, dim=4, length=16, state=13, seed=4, D=True, z=True), + {}), + ("grouped_bc", + dict(batch=2, dim=8, length=24, state=16, groups=2, seed=5, z=True), + {}), + ("edge_len1", dict(batch=1, dim=4, length=1, state=16, seed=6), {}), + ("long_seq", + dict(batch=1, dim=8, length=512, state=16, seed=7, D=True, z=True, + delta_bias=True), + {}), + ("merge_mean", + dict(batch=1, dim=4, length=16, state=16, seed=8, z=True), + dict(merge="mean")), + ("merge_concat", + dict(batch=1, dim=4, length=16, state=16, seed=9, z=True), + dict(merge="concat")), +] + + +def check(name, case, kwargs): + got = arm_scan.bidirectional_scan(**case, **kwargs) + want = reference_bidirectional(**case, **kwargs) + err = (got.double() - want).abs().max().item() + ok = err < MAX_ABS + print(f" {name:22s} max_abs={err:.3e} {'ok' if ok else 'FAIL'}") + return ok + + +def check_untied(): + """Backward direction with its own weights (Vim `bimamba_type=v2` style): + a separate A, D and delta must actually be used, and must NOT be flipped + by the wrapper (the caller supplies them in forward-time order).""" + case = make_case(batch=1, dim=4, length=16, state=16, seed=10, + D=True, z=True) + g = torch.Generator().manual_seed(99) + rev = dict( + A=-torch.rand(4, 16, generator=g, dtype=torch.float32) - 0.1, + D=torch.randn(4, generator=g, dtype=torch.float32), + delta=torch.randn(1, 4, 16, generator=g, dtype=torch.float32), + ) + got = arm_scan.bidirectional_scan(**case, reverse_params=rev) + want = reference_bidirectional(**case, reverse_params=rev) + err = (got.double() - want).abs().max().item() + ok = err < MAX_ABS + print(f" {'untied_reverse_params':22s} max_abs={err:.3e} " + f"{'ok' if ok else 'FAIL'}") + + # A wrapper that silently ignored reverse_params would still pass the + # check above only if the tied result happened to match — prove it does not. + tied = arm_scan.bidirectional_scan(**case) + differs = (got - tied).abs().max().item() > 1e-3 + print(f" {'untied != tied':22s} " + f"{'ok' if differs else 'FAIL (params ignored?)'}") + return ok and differs + + +def check_merge_none(): + """merge='none' returns both directions unmerged, and summing them by hand + must reproduce merge='sum' — i.e. the escape hatch is consistent.""" + case = make_case(batch=1, dim=4, length=16, state=16, seed=11, z=True) + out_f, out_b = arm_scan.bidirectional_scan(**case, merge="none") + summed = arm_scan.bidirectional_scan(**case, merge="sum") + err = (out_f + out_b - summed).abs().max().item() + ok = err == 0.0 + print(f" {'merge_none':22s} max_abs={err:.3e} {'ok' if ok else 'FAIL'}") + return ok + + +def check_forward_direction_unchanged(): + """The forward half of a bidirectional scan must be bit-identical to a + plain 1D scan — a regression here would mean the wrapper perturbed the + path the rest of the project already validates.""" + case = make_case(batch=1, dim=4, length=16, state=16, seed=12, D=True, + z=True, delta_bias=True) + out_f, _ = arm_scan.bidirectional_scan(**case, merge="none") + plain = arm_scan.selective_scan( + case["u"], case["delta"], case["A"], case["B"], case["C"], + D=case.get("D"), z=case.get("z"), delta_bias=case.get("delta_bias"), + delta_softplus=case["delta_softplus"]) + ok = torch.equal(out_f, plain) + print(f" {'fwd == plain 1D scan':22s} {'ok (bit-identical)' if ok else 'FAIL'}") + return ok + + +def main(): + print(f"kernel library: {arm_scan.lib_path()}") + failures = [] + + for name, case_kw, kwargs in CASES: + case = make_case(**case_kw) + if not check(name, case, kwargs): + failures.append(name) + + if not check_untied(): + failures.append("untied_reverse_params") + if not check_merge_none(): + failures.append("merge_none") + if not check_forward_direction_unchanged(): + failures.append("fwd_vs_plain") + + if failures: + print(f"\nBIDIRECTIONAL CHECK FAILED: {failures}") + sys.exit(1) + print(f"\nall bidirectional cases pass (max_abs < {MAX_ABS:g} vs f64 reference)") + + +if __name__ == "__main__": + main() diff --git a/tests/check_bidirectional_math.py b/tests/check_bidirectional_math.py new file mode 100644 index 0000000..a0c9ce1 --- /dev/null +++ b/tests/check_bidirectional_math.py @@ -0,0 +1,269 @@ +"""The bidirectional scan's DEFINITION, verified in pure numpy — no kernel, no torch. + +`bidirectional.py` and the Rust `reverse` flag (`ScanInput::reverse`) both rest +on one load-bearing claim: + + flipping the time axis, running an ordinary FORWARD scan, and flipping the + output back == running the recurrence BACKWARD in time. + +If that equivalence is false, the whole bidirectional design is wrong — and no +amount of kernel testing would reveal it, because both paths would be +consistently wrong together. So it is checked here, independently: this file +implements the backward recurrence *directly* (an explicit reverse-time loop) +and compares it against flip-forward-flip built on the already-independent +`naive_scan_f64` from `verify_golden.py`. + +This file therefore doubles as the **executable spec for the Rust `reverse` +flag**: `naive_scan_backward_f64` below is exactly what `reverse=true` computes. +The Rust side asserts the same identity bit-for-bit against the real kernel +(`reverse_matches_flip_forward_flip` in `kernel/arm-scan-core/tests/property.rs`); +this file proves the identity itself is sound, independently of any kernel. + +Runs with numpy alone — no Rust toolchain, no built cdylib, no torch. The +kernel-level gate for `bidirectional.py` itself is `check_bidirectional.py`. + +Usage: python tests/check_bidirectional_math.py +""" + +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).parent)) + +from verify_golden import naive_scan_f64, silu # noqa: E402 + +# Two float64 computations of the same recurrence, differing only in index +# order, should agree to round-off. They are in fact usually bit-identical. +RTOL, ATOL = 1e-12, 1e-14 + + +def naive_scan_backward_f64(u, delta, A, B, C, D_skip=None, z=None, + delta_bias=None, delta_softplus=False): + """The selective scan run BACKWARD in time — the spec for `reverse=true`. + + Identical to `naive_scan_f64` except the time loop runs from L-1 down to 0. + The state starts at zero at the END of the sequence and accumulates toward + the start; output for timestep t is still written at index t, and the + pointwise D-skip and z-gate still apply at index t. Deliberately naive: an + explicit loop, no flips anywhere, so it shares no mechanism with the + flip-based path it is used to check. + """ + u = u.astype(np.float64) + delta = delta.astype(np.float64) + A = A.astype(np.float64) + B = B.astype(np.float64) + C = C.astype(np.float64) + + batch, dim, L = u.shape + N = A.shape[1] + grouped = B.ndim == 4 + group_size = dim // B.shape[1] if grouped else None + + if delta_bias is not None: + delta = delta + delta_bias.astype(np.float64)[None, :, None] + if delta_softplus: + # same softplus as verify_golden (torch's beta=1, threshold=20) + delta = np.where(delta > 20.0, delta, + np.log1p(np.exp(np.minimum(delta, 20.0)))) + + out = np.empty((batch, dim, L), dtype=np.float64) + last_state = np.empty((batch, dim, N), dtype=np.float64) + for b in range(batch): + for d in range(dim): + g = d // group_size if grouped else None + h = np.zeros(N, dtype=np.float64) + for t in range(L - 1, -1, -1): # <-- the only difference + dt = delta[b, d, t] + b_vec = B[b, g, :, t] if grouped else B[b, :, t] + c_vec = C[b, g, :, t] if grouped else C[b, :, t] + h = np.exp(dt * A[d, :]) * h + (dt * u[b, d, t]) * b_vec + acc = 0.0 + for n in range(N): + acc += c_vec[n] * h[n] + out[b, d, t] = acc + last_state[b, d, :] = h # state after consuming t=0 + + if D_skip is not None: + out = out + u * D_skip.astype(np.float64)[None, :, None] + if z is not None: + out = out * silu(z.astype(np.float64)) + return out, last_state + + +def flip_time(x): + """Reverse the last axis — the time axis for every time-varying tensor + (u/delta/z: (B,D,L); B/C: (B,[G,]N,L)). Mirrors bidirectional.py.""" + return None if x is None else np.flip(x, axis=-1).copy() + + +def backward_via_flip(u, delta, A, B, C, D_skip=None, z=None, delta_bias=None, + delta_softplus=False): + """What `bidirectional.py` does today: flip the time-varying inputs, run + the ordinary forward scan, flip the output back. Note A / D_skip / + delta_bias have no time axis and are NOT flipped.""" + out, last = naive_scan_f64( + flip_time(u), flip_time(delta), A, flip_time(B), flip_time(C), + D_skip=D_skip, z=flip_time(z), delta_bias=delta_bias, + delta_softplus=delta_softplus) + return flip_time(out), last + + +def make_case(batch, dim, length, state, groups=None, seed=0, + D=False, z=False, delta_bias=False, delta_softplus=True): + """Note: `delta` is drawn from a normal (either sign) even in the + no_softplus case, which DIFFERS from `check_bidirectional.py` on purpose. + + That file must draw a positive delta because the kernel's Pass-A2 exp is + specialized for dt*A <= 0 and a negative timestep breaks its precondition. + Here there is no kernel — numpy's exp is exact over the whole line — and the + identity under test (flip-forward-flip == backward recurrence) is a + mathematical fact that holds for ANY delta. Restricting the sign would only + narrow coverage of the thing this file exists to prove. Do not "fix" it to + match. + """ + rng = np.random.default_rng(seed) + f32 = lambda *s: rng.standard_normal(s).astype(np.float32) + bc_shape = ((batch, groups, state, length) if groups + else (batch, state, length)) + case = dict( + u=f32(batch, dim, length), + delta=f32(batch, dim, length), + # A < 0 always (models parameterize it as -exp(A_log)) + A=(-rng.random((dim, state)) - 0.1).astype(np.float32), + B=f32(*bc_shape), + C=f32(*bc_shape), + delta_softplus=delta_softplus, + ) + if D: + case["D_skip"] = f32(dim) + if z: + case["z"] = f32(batch, dim, length) + if delta_bias: + case["delta_bias"] = f32(dim) + return case + + +CASES = [ + ("tiny", dict(batch=1, dim=2, length=4, state=16, seed=1)), + ("full_opts", dict(batch=2, dim=4, length=16, state=16, seed=2, + D=True, z=True, delta_bias=True)), + ("no_softplus", dict(batch=1, dim=2, length=8, state=16, seed=3, + delta_softplus=False)), + ("state13_neon_tail", dict(batch=1, dim=2, length=8, state=13, seed=4, + D=True, z=True)), + ("grouped_bc", dict(batch=2, dim=4, length=12, state=16, groups=2, + seed=5, z=True)), + ("edge_len1", dict(batch=1, dim=2, length=1, state=16, seed=6, + D=True, z=True)), + ("longer_seq", dict(batch=1, dim=2, length=128, state=16, seed=7, + D=True, z=True, delta_bias=True)), +] + + +def check_equivalence(name, case): + """flip-forward-flip == an explicit backward-in-time recurrence.""" + direct, direct_last = naive_scan_backward_f64(**case) + via_flip, flip_last = backward_via_flip(**case) + + err = np.abs(direct - via_flip).max() + bit_identical = np.array_equal(direct, via_flip) + last_err = np.abs(direct_last - flip_last).max() + ok = (np.allclose(direct, via_flip, rtol=RTOL, atol=ATOL) + and np.allclose(direct_last, flip_last, rtol=RTOL, atol=ATOL)) + + note = "bit-identical" if bit_identical else f"max_abs={err:.3e}" + print(f" {name:20s} out {note:16s} last_state={last_err:.3e} " + f"{'ok' if ok else 'FAIL'}") + return ok + + +def check_backward_is_not_forward(): + """Guard against a vacuous pass: if the 'backward' scan silently computed + the forward one, every equivalence check above would still succeed. It must + actually differ from the forward scan.""" + case = make_case(batch=1, dim=2, length=16, state=16, seed=20, z=True) + fwd, _ = naive_scan_f64(**case) + bwd, _ = naive_scan_backward_f64(**case) + diff = np.abs(fwd - bwd).max() + ok = diff > 1e-3 + print(f" {'backward != forward':20s} max_abs={diff:.3e} " + f"{'ok' if ok else 'FAIL (backward is a no-op?)'}") + return ok + + +def check_gate_commutes_with_linear_merge(): + """`bidirectional.py` lets the kernel apply the z-gate inside BOTH + directions rather than once after the merge. For a linear merge that is + algebraically identical; this proves it, because if it were false the + module's gating would be quietly wrong. + + inside: (y_f + D·u)·silu(z) + (y_b + D·u)·silu(z) + outside: ((y_f + D·u) + (y_b + D·u))·silu(z) + """ + case = make_case(batch=1, dim=4, length=16, state=16, seed=21, + D=True, z=True, delta_bias=True) + z = case.pop("z") + + gated_f, _ = naive_scan_f64(**case, z=z) + gated_b, _ = naive_scan_backward_f64(**case, z=z) + inside = gated_f + gated_b + + plain_f, _ = naive_scan_f64(**case) + plain_b, _ = naive_scan_backward_f64(**case) + outside = (plain_f + plain_b) * silu(z.astype(np.float64)) + + err = np.abs(inside - outside).max() + ok = np.allclose(inside, outside, rtol=RTOL, atol=ATOL) + print(f" {'gate commutes/sum':20s} max_abs={err:.3e} " + f"{'ok' if ok else 'FAIL'}") + return ok + + +def check_d_skip_is_applied_twice(): + """A documented gotcha, asserted so it can never drift silently. + + With merge='sum' and a shared D, the skip connection lands in BOTH + directions, so the merged output carries 2·D·u — not D·u. That matches how + real bidirectional Mambas behave (each direction's mixer applies its own D + and the outputs are summed), but it surprises people, so it is pinned here. + """ + case = make_case(batch=1, dim=4, length=8, state=16, seed=22, D=True) + d_skip = case.pop("D_skip") + + with_d = (naive_scan_f64(**case, D_skip=d_skip)[0] + + naive_scan_backward_f64(**case, D_skip=d_skip)[0]) + without_d = naive_scan_f64(**case)[0] + naive_scan_backward_f64(**case)[0] + + contribution = with_d - without_d + expected_twice = 2.0 * case["u"].astype(np.float64) * d_skip.astype( + np.float64)[None, :, None] + err = np.abs(contribution - expected_twice).max() + ok = np.allclose(contribution, expected_twice, rtol=RTOL, atol=ATOL) + print(f" {'D applied twice':20s} max_abs={err:.3e} " + f"{'ok (2·D·u, as documented)' if ok else 'FAIL'}") + return ok + + +def main(): + print("bidirectional definition check (pure numpy, no kernel/torch)") + print("proving: flip -> forward scan -> flip == backward-in-time recurrence\n") + + results = [check_equivalence(name, make_case(**kw)) for name, kw in CASES] + print() + results.append(check_backward_is_not_forward()) + results.append(check_gate_commutes_with_linear_merge()) + results.append(check_d_skip_is_applied_twice()) + + if not all(results): + print("\nBIDIRECTIONAL MATH CHECK FAILED") + sys.exit(1) + print("\nequivalence holds — flip-forward-flip and a backward-in-time " + "recurrence are the same function") + print("(this is the spec the Rust `reverse` flag implements; it is checked " + "against it bit-for-bit by reverse_matches_flip_forward_flip)") + + +if __name__ == "__main__": + main()