trnsparse 0.7.0__tar.gz → 0.8.0__tar.gz

{trnsparse-0.7.0 → trnsparse-0.8.0}/CHANGELOG.md

@@ -7,12 +7,112 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.8.0] — 2026-10-30
+
+Native TorchNeuron dispatch (Trn2/Trn3) alongside the PyTorch/XLA path (Trn1, maintenance); validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0) — see trnsci/trnsci#35.
+
+Headline performance: static-structure BSR SpMM, 6.8× faster than dense aten `matmul` at 1% block density, n=8192 (#34).
+
+### Changed
+
+#### Native path: static-structure kernels (#30)
+
+On TorchNeuron the sparse structure is host-resident, so the native path now
+hands it to the kernels as compile-time constants (`int[]` op arguments,
+tuples inside the kernel) instead of padded, gathered device tensors.
+Measured on Trn2 (`benchmarks/bench_native_trn2.py`; tables on #30):
+
+- `bsr_spmm` → `trnsparse::bsr_spmm_static`: walks exactly the stored
+  blocks and addresses B by block column; no `(M_tiles, K_max)` padding and
+  no device-side gather. 8192×8192 @ 8192×512: 3.1–12.5 ms → 0.7–3.9 ms;
+  beats dense aten `matmul` on the device up to ~13% (width 128) to ~34%
+  (width 512) block density at n = 8192.
+- Block-sparse attention forward → `trnsparse::attn_fwd_static`: one launch,
+  two passes with the softmax statistics on chip; no K/V gather, no host
+  reduction, no per-block score round trip through HBM. Backward →
+  `trnsparse::attn_bwd_static`: one launch for dQ (row-first) and dK/dV
+  (column-first). seq 8192 global mask, fwd+bwd: 127–195 ms → ~4 ms.
+- Screened SpMM masks each tile in the transposed frame, so the masked tile
+  is directly the stationary operand (no HBM round trip per tile).
+- `spmm` with host CSR scatters the nnz triplets on the device when that is
+  cheaper than copying a host-built dense matrix.
+
+One compile per distinct structure (seconds; cached in-process by the
+structure). `TRNSPARSE_STATIC_STRUCTURE=0` restores the gather path;
+structures above `TRNSPARSE_STATIC_MAX_UNROLL` (default 16384 unrolled
+block matmuls) use it automatically.
+
+### Added
+
+- `benchmarks/bench_native_trn2.py`: warm-latency benchmark for the native
+  path against dense aten on the device, CPU, and `torch.sparse` on CPU.
+
+#### Native PyTorch dispatch on TorchNeuron (#29, #30)
+
+On TorchNeuron (Neuron SDK 2.33+, `torch-neuronx` 2.14, PyTorch 2.14,
+Trn2/Trn3), every NKI kernel is now registered as a PyTorch custom op with
+`torch_neuronx.nki_op` and dispatched on `device="neuron"`, with no
+`torch_xla` import. Ops: `trnsparse::spmm_dense`, `bsr_spmm`,
+`screened_spmm`, `attn_stats`, `attn_out`, `attn_bwd_dq`, `attn_bwd_dkdv`,
+plus the composite `trnsparse::bsr_attn` / `bsr_attn_bwd`. Kernel bodies
+are unchanged.
+
+- Autograd on the native path uses `register_autograd` on the ops instead
+  of `torch.autograd.Function`: `spmm_dense`, `bsr_spmm` and
+  `screened_spmm` carry their analytic adjoints, and `bsr_attn` carries the
+  NKI dQ / dK-dV backward (via `bsr_attn_bwd`). Gradients reach
+  `CSRMatrix.values`, `BSRMatrix.blocks` and the dense operands through the
+  host-side densify / pad / gather. `_SpMMFunction`, `_BSRSpMMFunction`,
+  `_ScreenedSpMMFunction` and `_AttnTiledFunction` remain for the XLA path,
+  the simulator and the PyTorch backend. The native BSR backward has no
+  per-block Python loop.
+- Placement (#30): the `neuron` device has no sparse CSR tensors
+  (aws-neuron-sdk#1442), so sparse structure stays on the host and only
+  dense data moves. `spmm`, `spmv`, `spmv_symmetric`, `bsr_spmm`,
+  `screened_spmm`, `block_sparse_attention_tiled`, `cg_bsr` and
+  `jacobi_preconditioner_bsr` accept host CSR/BSR structure with dense
+  operands (and BSR `blocks`) on `neuron`, and return results on the dense
+  operand's device. A host CSR against a device operand is densified on the
+  host without a torch sparse tensor (`formats._csr_to_dense_host`, static
+  shapes), so `spmm` with host-CSR structure traces under
+  `torch.compile(..., fullgraph=True)`. Index tensors found on a device are
+  read through host copies. Documented in `docs/installation.md` and the
+  format and op docstrings.
+- The dispatch path is detected on the first kernel launch: native when
+  `torch_neuronx` exports `nki_op` and a `neuron` device is available, XLA
+  when `torch_xla` is installed, otherwise none. The XLA path (Trn1) is
+  unchanged, as are the PyTorch fallback and the `TRNSPARSE_USE_SIMULATOR`
+  path. All kernel launches share one helper (`dispatch._launch`, or its
+  halves `_to_device` / `_run` where several attention kernels share
+  operands).
+- Custom ops are registered lazily (`trnsparse/nki/_native.py` and
+  `_native_attn.py` are imported only when the native path is selected).
+- `TRNSCI_NKI_DISPATCH=auto|native|xla|cpu` forces the dispatch path (a
+  suite-wide variable). Forcing an unavailable path raises instead of
+  falling back. `TRNSCI_XLA_DEPRECATION=1` makes the XLA path emit a
+  one-time `DeprecationWarning` (off by default).
+- New `[torchneuron]` extra (`torch>=2.14`, `torch-neuronx>=2.14`,
+  `nki>=0.6`); `[neuron]` stays the PyTorch/XLA lineage.
+- Tests: `conftest.py` sets `NKI_ENABLE_TRACE_CACHE=0` by default.
+  `tests/test_native_dispatch.py` (CPU) covers path selection, launch
+  routing, lazy registration, op schemas and fake shapes, the registered
+  backward formulas, the native wrappers end to end on CPU with real custom
+  ops, and a CPU `fullgraph=True` trace of host-CSR `spmm`.
+  `tests/test_nki_native.py` (`@pytest.mark.neuron`, native only) covers
+  host-structure / device-operand calls, gradients, and
+  `torch.compile(backend="neuron", fullgraph=True)` of `spmm` with
+  host-CSR structure. Not yet run on hardware.
+
 ## [0.7.0] — 2026-10-28
 
 Validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0, native TorchNeuron device) — see trnsci/trnsci#35.
 
 ### Fixed
 
+- PyTorch reference attention (`_block_sparse_attn_pytorch` and the
+  recompute branch of `_block_sparse_attn_backward`) allocates on Q's
+  device instead of the default device.
+
 - **NKI 0.7.0 / neuronx-cc 2.28 kernel compatibility (#29).** Kernels that
   passed `nki.simulate` failed to compile on Trn2 (Neuron SDK 2.33):
   - `_attn_out_kernel`, `_attn_bwd_dq_kernel`, `_attn_bwd_dkdv_kernel`:

{trnsparse-0.7.0 → trnsparse-0.8.0}/CLAUDE.md

@@ -25,8 +25,10 @@ trnsparse/
 │   ├── iterative.py     # cg_bsr, power_iteration_bsr, jacobi_preconditioner_bsr
 │   └── nki/
 │       ├── __init__.py
-│       ├── kernels.py   # _bsr_spmm_kernel, _screened_spmm_kernel
-│       └── dispatch.py  # torch.autograd.Function wrappers + NKI/PyTorch routing
+│       ├── kernels.py       # _bsr_spmm_kernel, _screened_spmm_kernel, attention kernels
+│       ├── dispatch.py      # dispatch-path selection, _launch, autograd.Function wrappers (XLA/CPU)
+│       ├── _native.py       # TorchNeuron nki_op custom ops + register_autograd, loaded lazily
+│       └── _native_attn.py  # trnsparse::bsr_attn / bsr_attn_bwd composite ops (native path)
 ├── tests/
 ├── examples/
 │   ├── sparse_fock.py             # Schwarz-screened Fock build (3 paths + trnblas)

{trnsparse-0.7.0 → trnsparse-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.7.0
+Version: 0.8.0
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0
@@ -22,6 +22,10 @@ Provides-Extra: neuron
 Requires-Dist: nki>=0.3.0; extra == "neuron"
 Requires-Dist: neuronxcc>=2.24; extra == "neuron"
 Requires-Dist: torch-neuronx>=2.9; extra == "neuron"
+Provides-Extra: torchneuron
+Requires-Dist: torch>=2.14; extra == "torchneuron"
+Requires-Dist: torch-neuronx>=2.14; extra == "torchneuron"
+Requires-Dist: nki>=0.6; extra == "torchneuron"
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-benchmark>=4.0; extra == "dev"
@@ -65,8 +69,34 @@ Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1)
 
 ```bash
 pip install trnsparse
+
+# Trn2/Trn3 with TorchNeuron (native PyTorch, Neuron SDK 2.33+)
+pip install "trnsparse[torchneuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
+
+# Trn1 / PyTorch-XLA lineage
+pip install "trnsparse[neuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
 ```
 
+Kernels dispatch on whichever Neuron PyTorch lineage is installed:
+
+- **TorchNeuron (native)** — Trn2/Trn3, Neuron SDK 2.33+, Python ≥ 3.12,
+  `torch` 2.14, `torch-neuronx` 2.14.*. Each kernel is a PyTorch custom op
+  (`torch.ops.trnsparse.*`) on the `neuron` device, with its backward
+  attached via `register_autograd`; no `torch_xla`.
+- **PyTorch/XLA** — Trn1 (and any host on the `torch-neuronx` 2.9 lineage).
+  Kernels dispatch through `torch_xla`, unchanged.
+
+**Placement on the native path:** the `neuron` device has no sparse CSR
+tensors, so keep sparse *structure* (CSR `row_ptrs` / `col_indices`, BSR
+`block_row_ptrs` / `block_col_indices`) on the host and put dense operands
+(and BSR `blocks`) on `device="neuron"`. Results come back on the dense
+operand's device. See [installation](docs/installation.md#placement-sparse-structure-on-the-host).
+
+Set `TRNSCI_NKI_DISPATCH=native|xla|cpu` (shared by all trnsci libraries)
+to force a path. With an editable install (`pip install -e`), set
+`NKI_ENABLE_TRACE_CACHE=0`: the NKI kernel cache only versions installed
+wheels and can serve stale kernels.
+
 ## Usage
 
 ```python

{trnsparse-0.7.0 → trnsparse-0.8.0}/README.md

@@ -34,8 +34,34 @@ Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1)
 
 ```bash
 pip install trnsparse
+
+# Trn2/Trn3 with TorchNeuron (native PyTorch, Neuron SDK 2.33+)
+pip install "trnsparse[torchneuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
+
+# Trn1 / PyTorch-XLA lineage
+pip install "trnsparse[neuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
 ```
 
+Kernels dispatch on whichever Neuron PyTorch lineage is installed:
+
+- **TorchNeuron (native)** — Trn2/Trn3, Neuron SDK 2.33+, Python ≥ 3.12,
+  `torch` 2.14, `torch-neuronx` 2.14.*. Each kernel is a PyTorch custom op
+  (`torch.ops.trnsparse.*`) on the `neuron` device, with its backward
+  attached via `register_autograd`; no `torch_xla`.
+- **PyTorch/XLA** — Trn1 (and any host on the `torch-neuronx` 2.9 lineage).
+  Kernels dispatch through `torch_xla`, unchanged.
+
+**Placement on the native path:** the `neuron` device has no sparse CSR
+tensors, so keep sparse *structure* (CSR `row_ptrs` / `col_indices`, BSR
+`block_row_ptrs` / `block_col_indices`) on the host and put dense operands
+(and BSR `blocks`) on `device="neuron"`. Results come back on the dense
+operand's device. See [installation](docs/installation.md#placement-sparse-structure-on-the-host).
+
+Set `TRNSCI_NKI_DISPATCH=native|xla|cpu` (shared by all trnsci libraries)
+to force a path. With an editable install (`pip install -e`), set
+`NKI_ENABLE_TRACE_CACHE=0`: the NKI kernel cache only versions installed
+wheels and can serve stale kernels.
+
 ## Usage
 
 ```python

trnsparse-0.8.0/benchmarks/bench_native_trn2.py

@@ -0,0 +1,397 @@
+"""Warm-latency benchmark for the native TorchNeuron path (Trn2/Trn3).
+
+Not a pytest-benchmark file: a standalone script, one section per process,
+so each run stays short and a compile failure does not take the rest down.
+Every row is printed as one JSON line (and appended to ``--out``).
+
+    NKI_ENABLE_TRACE_CACHE=0 python benchmarks/bench_native_trn2.py --section bsr --size 4096
+
+Sections:
+    bsr     BSR (128x128) SpMM vs dense aten matmul on the device, CPU dense,
+            and torch.sparse BSR/CSR on CPU. Block densities x RHS widths.
+    csr     CSR SpMM (host structure, densify, NKI dense kernel) at element
+            densities, vs the same baselines.
+    screen  Screened SpMM vs dense masked matmul.
+    attn    Block-sparse attention (local / dilated / global masks), forward
+            and forward+backward, vs dense SDPA on the device and on CPU.
+
+Timing: each callable is run ``--warmup`` times (the first call compiles),
+then ``--reps`` times with ``torch.neuron.synchronize()`` after every call;
+the median is reported in milliseconds. Structure stays on the host, dense
+data (blocks, operands) lives on the device, as `docs/installation.md`
+recommends for the native path.
+"""
+
+# ruff: noqa: B023  (every closure is benchmarked before its loop variables change)
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+import trnsparse  # noqa: E402
+from trnsparse.nki import dispatch as _d  # noqa: E402
+
+DEV = "neuron"
+P = 128
+
+
+def _sync():
+    if hasattr(torch, "neuron"):
+        torch.neuron.synchronize()
+
+
+def bench(fn, warmup=2, reps=10):
+    for _ in range(warmup):
+        fn()
+    _sync()
+    ts = []
+    for _ in range(reps):
+        t0 = time.perf_counter()
+        fn()
+        _sync()
+        ts.append((time.perf_counter() - t0) * 1e3)
+    return statistics.median(ts)
+
+
+def safe_bench(fn, **kw):
+    try:
+        return bench(fn, **kw)
+    except Exception as e:  # report, don't abort the section
+        print(f"# bench failed: {type(e).__name__}: {str(e)[:200]}", file=sys.stderr)
+        return None
+
+
+def rel_err(x, ref):
+    x = x.detach().cpu().double()
+    ref = ref.detach().cpu().double()
+    return float((x - ref).norm() / ref.norm().clamp_min(1e-30))
+
+
+def emit(row, out):
+    line = json.dumps(row)
+    print(line, flush=True)
+    if out:
+        with open(out, "a") as f:
+            f.write(line + "\n")
+
+
+# ── BSR ──────────────────────────────────────────────────────────────────────
+
+
+def random_bsr(n, density, seed=0):
+    g = torch.Generator().manual_seed(seed)
+    nb = n // P
+    total = nb * nb
+    k = max(1, round(density * total))
+    flat = torch.randperm(total, generator=g)[:k].sort().values
+    rows, cols = flat // nb, flat % nb
+    row_ptrs = torch.zeros(nb + 1, dtype=torch.long)
+    row_ptrs[1:] = torch.bincount(rows, minlength=nb).cumsum(0)
+    blocks = torch.randn(k, P, P, generator=g) / math.sqrt(P)
+    return trnsparse.BSRMatrix(
+        blocks=blocks,
+        block_col_indices=cols.to(torch.long),
+        block_row_ptrs=row_ptrs,
+        shape=(n, n),
+        block_size=P,
+    )
+
+
+def section_bsr(args):
+    densities = [float(x) for x in args.densities.split(",")]
+    widths = [int(x) for x in args.widths.split(",")]
+    n = args.size
+    for width in widths:
+        g = torch.Generator().manual_seed(1)
+        Bc = torch.randn(n, width, generator=g)
+        Bd = Bc.to(DEV)
+        dense_dev_ms = None
+        cpu_dense_ms = None
+        for dens in densities:
+            A = random_bsr(n, dens)
+            Ac = A.to_dense()
+            Ad_dev = trnsparse.BSRMatrix(
+                blocks=A.blocks.to(DEV),
+                block_col_indices=A.block_col_indices,
+                block_row_ptrs=A.block_row_ptrs,
+                shape=A.shape,
+                block_size=P,
+            )
+            counts = (A.block_row_ptrs[1:] - A.block_row_ptrs[:-1]).tolist()
+            K_max = max(counts)
+            row = {
+                "section": "bsr",
+                "n": n,
+                "width": width,
+                "block_density": dens,
+                "n_blocks": A.n_blocks,
+                "K_max": K_max,
+                "mean_blocks_per_row": sum(counts) / len(counts),
+            }
+            ref = Ac.double() @ Bc.double()
+            out = trnsparse.bsr_spmm(Ad_dev, Bd)
+            row["rel_err"] = rel_err(out, ref)
+            row["trnsparse_ms"] = safe_bench(lambda: trnsparse.bsr_spmm(Ad_dev, Bd), reps=args.reps)
+
+            # Split: device-side prep (gather) vs the kernel launch.
+            def prep():
+                return _d._bsr_pad_and_gather(_d._BlockView(Ad_dev, Ad_dev.blocks), Bd)
+
+            bp, bg, *_ = prep()
+            ops = _d._get_native_ops()
+            # Gather path (main before trnsparse#30 static kernels): device-side
+            # pad + gather, then the K_max-padded kernel.
+            row["gather_prep_ms"] = safe_bench(prep, reps=args.reps)
+            row["gather_kernel_ms"] = safe_bench(lambda: ops["bsr_spmm"](bp, bg), reps=args.reps)
+            bp = bg = None  # free the gathered tiles
+            if "bsr_spmm_static" in ops:
+                rp, cl = _d._bsr_structure_lists(A)
+                bl = Ad_dev.blocks
+                row["static_kernel_ms"] = safe_bench(
+                    lambda: ops["bsr_spmm_static"](bl, Bd, rp, cl), reps=args.reps
+                )
+
+            Ad = Ac.to(DEV)
+            if dense_dev_ms is None:
+                dense_dev_ms = safe_bench(lambda: Ad @ Bd, reps=args.reps)
+            row["dense_dev_ms"] = dense_dev_ms
+            Ad = None
+            if cpu_dense_ms is None:
+                cpu_dense_ms = safe_bench(lambda: Ac @ Bc, warmup=1, reps=args.cpu_reps)
+            row["cpu_dense_ms"] = cpu_dense_ms
+            try:
+                Absr = Ac.to_sparse_bsr((P, P))
+                row["cpu_torch_bsr_ms"] = safe_bench(
+                    lambda: Absr @ Bc, warmup=1, reps=args.cpu_reps
+                )
+            except Exception as e:
+                row["cpu_torch_bsr_ms"] = f"unsupported: {type(e).__name__}"
+            Acsr = Ac.to_sparse_csr()
+            row["cpu_torch_csr_ms"] = safe_bench(lambda: Acsr @ Bc, warmup=1, reps=args.cpu_reps)
+            # trnsparse on CPU (pytorch backend): the loop-over-blocks fallback.
+            prev = trnsparse.get_backend()
+            trnsparse.set_backend("pytorch")
+            row["cpu_trnsparse_ms"] = safe_bench(
+                lambda: trnsparse.bsr_spmm(A, Bc), warmup=1, reps=args.cpu_reps
+            )
+            trnsparse.set_backend(prev)
+            if row["trnsparse_ms"] and row["dense_dev_ms"]:
+                row["speedup_vs_dense_dev"] = row["dense_dev_ms"] / row["trnsparse_ms"]
+            emit(row, args.out)
+
+
+# ── CSR ──────────────────────────────────────────────────────────────────────
+
+
+def section_csr(args):
+    densities = [float(x) for x in args.densities.split(",")]
+    widths = [int(x) for x in args.widths.split(",")]
+    n = args.size
+    for width in widths:
+        Bc = torch.randn(n, width, generator=torch.Generator().manual_seed(1))
+        Bd = Bc.to(DEV)
+        for dens in densities:
+            g = torch.Generator().manual_seed(0)
+            Ac = torch.randn(n, n, generator=g) * (torch.rand(n, n, generator=g) < dens)
+            A = trnsparse.from_dense(Ac)
+            row = {"section": "csr", "n": n, "width": width, "density": dens, "nnz": A.nnz}
+            out = trnsparse.spmm(A, Bd)
+            row["rel_err"] = rel_err(out, Ac.double() @ Bc.double())
+            row["trnsparse_ms"] = safe_bench(lambda: trnsparse.spmm(A, Bd), reps=args.reps)
+            Ad = Ac.to(DEV)
+            row["dense_dev_ms"] = safe_bench(lambda: Ad @ Bd, reps=args.reps)
+            row["densify_host_h2d_ms"] = safe_bench(
+                lambda: trnsparse.formats._csr_to_dense_host(A).to(DEV), reps=args.reps
+            )
+            if hasattr(_d, "_csr_to_dense_on"):
+                row["densify_on_device_ms"] = safe_bench(
+                    lambda: _d._csr_to_dense_on(A, DEV), reps=args.reps
+                )
+            row["cpu_dense_ms"] = safe_bench(lambda: Ac @ Bc, warmup=1, reps=args.cpu_reps)
+            Acsr = Ac.to_sparse_csr()
+            row["cpu_torch_csr_ms"] = safe_bench(lambda: Acsr @ Bc, warmup=1, reps=args.cpu_reps)
+            emit(row, args.out)
+            Ad = None
+
+
+# ── Screened SpMM ────────────────────────────────────────────────────────────
+
+
+def section_screen(args):
+    widths = [int(x) for x in args.widths.split(",")]
+    n = args.size
+    g = torch.Generator().manual_seed(0)
+    # Exponentially decaying Schwarz bounds, so a threshold keeps a band.
+    pos = torch.arange(n, dtype=torch.float32)
+    diag = torch.exp(-pos / (n / 8)) + 1e-12
+    Ac = torch.randn(n, n, generator=g)
+    for width in widths:
+        Bc = torch.randn(n, width, generator=g)
+        Bd, Ad, diag_d = Bc.to(DEV), Ac.to(DEV), diag.to(DEV)
+        for thr in [float(x) for x in args.thresholds.split(",")]:
+            Q = diag.sqrt()
+            mask = (Q[:, None] * Q[None, :]) > math.sqrt(thr)
+            ref = (Ac.double() * mask) @ Bc.double()
+            row = {
+                "section": "screen",
+                "n": n,
+                "width": width,
+                "threshold": thr,
+                "kept_fraction": float(mask.float().mean()),
+            }
+            out = trnsparse.screened_spmm(Ad, diag_d, Bd, thr)
+            row["rel_err"] = rel_err(out, ref)
+            row["trnsparse_ms"] = safe_bench(
+                lambda: trnsparse.screened_spmm(Ad, diag_d, Bd, thr), reps=args.reps
+            )
+            mask_d = mask.to(DEV, torch.float32)
+            row["dense_masked_dev_ms"] = safe_bench(lambda: (Ad * mask_d) @ Bd, reps=args.reps)
+            row["dense_dev_ms"] = safe_bench(lambda: Ad @ Bd, reps=args.reps)
+            mf = mask.float()
+            row["cpu_masked_ms"] = safe_bench(lambda: (Ac * mf) @ Bc, warmup=1, reps=args.cpu_reps)
+            emit(row, args.out)
+
+
+# ── Attention ────────────────────────────────────────────────────────────────
+
+
+def block_mask(kind, nb):
+    m = torch.zeros(nb, nb, dtype=torch.bool)
+    w = 1  # local window radius in blocks (3 blocks per row)
+    if kind == "local":
+        for i in range(nb):
+            m[i, max(0, i - w) : i + w + 1] = True
+    elif kind == "dilated":
+        stride = 4
+        for i in range(nb):
+            m[i, (torch.arange(nb) - i) % stride == 0] = True
+    elif kind == "global":
+        for i in range(nb):
+            m[i, max(0, i - w) : i + w + 1] = True
+        m[0, :] = True
+        m[:, 0] = True
+    else:
+        raise ValueError(kind)
+    return m
+
+
+def mask_bsr_from_block_mask(bm):
+    nb = bm.shape[0]
+    row_ptrs = torch.zeros(nb + 1, dtype=torch.long)
+    row_ptrs[1:] = bm.sum(1).cumsum(0)
+    cols = bm.nonzero()[:, 1].to(torch.long)
+    return trnsparse.BSRMatrix(
+        blocks=torch.ones(len(cols), 1, 1),  # values unused by attention
+        block_col_indices=cols,
+        block_row_ptrs=row_ptrs,
+        shape=(nb * P, nb * P),
+        block_size=P,
+    )
+
+
+def section_attn(args):
+    seq = args.size
+    nb = seq // P
+    F = torch.nn.functional
+    for hd in [int(x) for x in args.head_dims.split(",")]:
+        g = torch.Generator().manual_seed(0)
+        Qc, Kc, Vc = (torch.randn(seq, hd, generator=g) for _ in range(3))
+        dOc = torch.randn(seq, hd, generator=g)
+        Qd, Kd, Vd, dOd = (t.to(DEV) for t in (Qc, Kc, Vc, dOc))
+
+        def sdpa_fwd(q, k, v, m=None):
+            return F.scaled_dot_product_attention(q[None], k[None], v[None], attn_mask=m)[0]
+
+        dense_fwd = safe_bench(lambda: sdpa_fwd(Qd, Kd, Vd), reps=args.reps)
+
+        def dense_fwd_bwd(m=None):
+            q, k, v = (t.detach().requires_grad_() for t in (Qd, Kd, Vd))
+            o = sdpa_fwd(q, k, v, m)
+            o.backward(dOd)
+            return q.grad
+
+        dense_fb = safe_bench(dense_fwd_bwd, reps=args.reps)
+        cpu_fwd = safe_bench(lambda: sdpa_fwd(Qc, Kc, Vc), warmup=1, reps=args.cpu_reps)
+        for kind in args.masks.split(","):
+            bm = block_mask(kind, nb)
+            mb = mask_bsr_from_block_mask(bm)
+            elem_mask = bm.repeat_interleave(P, 0).repeat_interleave(P, 1)
+            counts = bm.sum(1).tolist()
+            row = {
+                "section": "attn",
+                "seq": seq,
+                "head_dim": hd,
+                "mask": kind,
+                "block_density": float(bm.float().mean()),
+                "K_max": max(counts),
+                "mean_blocks_per_row": sum(counts) / len(counts),
+            }
+            ref = sdpa_fwd(Qc.double(), Kc.double(), Vc.double(), elem_mask)
+            out = trnsparse.block_sparse_attention_tiled(Qd, Kd, Vd, mb)
+            row["fwd_rel_err"] = rel_err(out, ref)
+            row["trnsparse_fwd_ms"] = safe_bench(
+                lambda: trnsparse.block_sparse_attention_tiled(Qd, Kd, Vd, mb), reps=args.reps
+            )
+
+            def ts_fwd_bwd():
+                q, k, v = (t.detach().requires_grad_() for t in (Qd, Kd, Vd))
+                o = trnsparse.block_sparse_attention_tiled(q, k, v, mb)
+                o.backward(dOd)
+                return q.grad, k.grad, v.grad
+
+            if not args.no_bwd:
+                try:
+                    gq, gk, gv = ts_fwd_bwd()
+                    qr, kr, vr = (t.double().requires_grad_() for t in (Qc, Kc, Vc))
+                    sdpa_fwd(qr, kr, vr, elem_mask).backward(dOc.double())
+                    row["bwd_rel_err"] = max(
+                        rel_err(gq, qr.grad), rel_err(gk, kr.grad), rel_err(gv, vr.grad)
+                    )
+                    row["trnsparse_fwd_bwd_ms"] = safe_bench(ts_fwd_bwd, reps=args.reps)
+                except Exception as e:
+                    row["trnsparse_fwd_bwd_ms"] = f"failed: {type(e).__name__}: {str(e)[:120]}"
+            row["sdpa_dense_dev_fwd_ms"] = dense_fwd
+            row["sdpa_dense_dev_fwd_bwd_ms"] = dense_fb
+            em_d = elem_mask.to(DEV)
+            row["sdpa_masked_dev_fwd_ms"] = safe_bench(
+                lambda: sdpa_fwd(Qd, Kd, Vd, em_d), reps=args.reps
+            )
+            row["sdpa_cpu_fwd_ms"] = cpu_fwd
+            row["sdpa_masked_cpu_fwd_ms"] = safe_bench(
+                lambda: sdpa_fwd(Qc, Kc, Vc, elem_mask), warmup=1, reps=args.cpu_reps
+            )
+            emit(row, args.out)
+            em_d = None
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--section", required=True, choices=["bsr", "csr", "screen", "attn"])
+    ap.add_argument("--size", type=int, default=1024)
+    ap.add_argument("--densities", default="0.01,0.05,0.1,0.25,0.5")
+    ap.add_argument("--widths", default="128,512")
+    ap.add_argument("--thresholds", default="1e-12,1e-4,1e-2")
+    ap.add_argument("--head-dims", default="64,128")
+    ap.add_argument("--masks", default="local,dilated,global")
+    ap.add_argument("--no-bwd", action="store_true")
+    ap.add_argument("--reps", type=int, default=10)
+    ap.add_argument("--cpu-reps", type=int, default=3)
+    ap.add_argument("--out", default="")
+    args = ap.parse_args()
+    trnsparse.set_backend("nki")
+    {"bsr": section_bsr, "csr": section_csr, "screen": section_screen, "attn": section_attn}[
+        args.section
+    ](args)
+
+
+if __name__ == "__main__":
+    main()

{trnsparse-0.7.0 → trnsparse-0.8.0}/docs/api.md

@@ -31,3 +31,22 @@
 - `set_backend("auto" | "pytorch" | "nki")` — select compute backend
 - `get_backend()` — current backend
 - `HAS_NKI` — module-level flag set at import time
+
+On Neuron hardware the NKI backend runs on one of two `torch-neuronx`
+lineages, detected on the first kernel launch (force with
+`TRNSCI_NKI_DISPATCH=native|xla|cpu`):
+
+- **Native (TorchNeuron, Trn2/Trn3)** — each kernel is a custom op on
+  `device="neuron"`: `torch.ops.trnsparse.spmm_dense`, `bsr_spmm`,
+  `screened_spmm`, `attn_stats`, `attn_out`, `attn_bwd_dq`,
+  `attn_bwd_dkdv`, and the composite `bsr_attn` / `bsr_attn_bwd`. Backwards
+  are attached with `torch.library.register_autograd`; `spmm` also traces
+  under `torch.compile(backend="neuron", fullgraph=True)`. The ops are
+  registered only once this path is selected.
+- **XLA (Trn1)** — kernels run on the `torch_xla` device, as before.
+
+**Placement on the native path:** sparse structure (CSR `row_ptrs` /
+`col_indices` / `values`, BSR `block_row_ptrs` / `block_col_indices`) stays
+on the host; dense operands and BSR `blocks` may be on `neuron`; results are
+on the dense operand's device. See
+[Installation](installation.md#placement-sparse-structure-on-the-host).

{trnsparse-0.7.0 → trnsparse-0.8.0}/docs/architecture.md

@@ -83,6 +83,20 @@ Fock-build-sized inputs.
 kernel — a `BSRMatrix` with a local-window or dilated block pattern IS
 an attention mask. See [sparse_attention.md](sparse_attention.md).
 
+### Static structure on the native path (Trn2/Trn3)
+
+TorchNeuron has no device sparse CSR, so the structure stays on the host
+(#30). The native path turns that into an advantage: `bsr_spmm` and
+block-sparse attention pass `block_row_ptrs` / `block_col_indices` to the
+kernel as compile-time constants (`trnsparse::bsr_spmm_static`,
+`attn_fwd_static`, `attn_bwd_static`). The trace unrolls exactly the stored
+blocks with static DMA addresses: no `(M_tiles, K_max)` padding (one dense
+global-token row no longer makes every row cost `K_max` blocks), no gathered
+copies of B / K / V, and the attention softmax statistics stay on chip. The
+cost is one compile per distinct structure (seconds), which a reused pattern
+amortizes. `TRNSPARSE_STATIC_STRUCTURE=0` restores the gather path. Measured
+crossovers against dense aten: `docs/benchmarks.md`.
+
 ## Formats
 
 - **CSR**: row pointer + column indices + values. Preferred for construction,