trnsparse 0.4.2__tar.gz → 0.7.0__tar.gz

{trnsparse-0.4.2 → trnsparse-0.7.0}/.github/workflows/ci.yml

@@ -41,6 +41,8 @@ jobs:
     # Catches Python-trace-level errors (bad kwargs, dropped ops, shape
     # mismatches) pre-merge without AWS round-trips. MLIR verifier
     # errors remain hardware-only (simulator explicitly skips compile).
+    # TRNSPARSE_REQUIRE_NKI=1 disables the PyTorch fallback so kernel
+    # failures propagate — the simulator tests actually validate NKI kernels.
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v6
@@ -55,4 +57,5 @@ jobs:
       - name: Run simulator-backed kernel tests
         env:
           TRNSPARSE_USE_SIMULATOR: "1"
+          TRNSPARSE_REQUIRE_NKI: "1"
         run: pytest tests/ -v -m nki_simulator --tb=short

{trnsparse-0.4.2 → trnsparse-0.7.0}/.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.8.6
+    rev: v0.15.11
     hooks:
       - id: ruff
         args: [--fix]

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

@@ -5,6 +5,214 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+## [0.7.0] — 2026-10-28
+
+Validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0, native TorchNeuron device) — see trnsci/trnsci#35.
+
+### Fixed
+
+- **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`:
+    softmax normalisation `nl.divide(exp, row_denom)` by a `(128, 1)`
+    denominator lowered to `tensor_scalar_arith`, which no longer supports
+    `divide`. Now multiplies by `nl.reciprocal(row_denom)` (hoisted per
+    block-row in the forward and dQ kernels).
+  - `_screened_spmm_kernel`: the `(128, 1) * (1, 128)` Schwarz outer product
+    and the `(1, 1)` threshold subtract produced partition-count mismatches
+    (`operand0 partition total elements 128 != dst partition total elements
+    1`). `q_k` is now `nl.broadcast_to` `(128, 128)` and the threshold is
+    loaded once and broadcast to `(128, 1)`. The soft-mask tensor-tensor
+    divide (fails ISA validation `s3s3d3_tt_valid_op`) is now a reciprocal
+    multiply.
+  - `nki_bsr_attn_bwd`: the host-side `(..., b) → (..., b, 1)` unsqueeze
+    for per-row stats matched `q_scaled` / `k_blocks` / `v_blocks` by shape
+    whenever `head_dim == 128` (and in the simulator's padded path), feeding
+    the kernels a 4-D tensor. Stats are now selected by name.
+  - Dispatch: kernels always receive fp32 (Trainium has no fp64 datapath;
+    the native device downcast fp64 silently) and results are cast back to
+    the caller's dtype. Fixes `expected m1 and m2 to have the same dtype,
+    but got: float != double` in the PyTorch-level backward under fp64
+    gradcheck. The hardware gradcheck tests now use `eps=1e-2,
+    atol=rtol=1e-2`, sized for an fp32 forward (the ops are linear in each
+    input, so the larger eps adds no truncation error).
+
+- **Block-sparse attention: two host-side gather bugs** (pre-existing, masked
+  on Trn2 by the compile failures above; found with neuronx-cc BIRSim):
+  - Rows with fewer than `K_max` blocks were padded with sentinel k-slots that
+    gathered column block 0, so block 0 was counted again in the softmax and
+    in `weights @ V` (local-window masks: ~49% of outputs wrong). Padded slots
+    now have K/V (forward, dQ) or Q/dO (dK/dV) zeroed, and their stats are
+    dropped in `_attn_host_reduction`.
+  - `_attn_bwd_gather` gathered `dO` by column block instead of the row block
+    `dO_m`, breaking dQ whenever a row attends to a block other than itself
+    (~99% of dQ wrong). The two simulator xfails attributed to a "PSUM
+    accumulate bug" (`test_bwd_dq_parity_kmax2`, `test_backward_head_dim_256`)
+    were this bug and now pass.
+
+### Added
+
+- `tests/test_attn_gather.py`: CPU regression tests for both gather bugs.
+- `tests/test_nki_compile.py` (`nki_simulator` marker): CPU MLIR
+  verification (`source_to_mlir(...).verify()`) of every kernel for trn1
+  and trn2, attention kernels at head_dim 64, 128 and 256; plus a full NEFF
+  compile per kernel/target when `neuronx-cc` is on PATH.
+
+## [0.6.0] — 2026-04-17
+
+### Added
+
+- **K-tiling for head_dim > 128** — NKI attention kernels now support
+  `head_dim=256` and `head_dim=512` via an inner K-tile loop that splits
+  the `Q @ K.T` score computation across `TILE_K=128`-sized chunks of
+  `head_dim`, accumulating each into the existing PSUM.
+  - All 4 NKI kernels updated: `_attn_stats_kernel`, `_attn_out_kernel`,
+    `_attn_bwd_dq_kernel`, `_attn_bwd_dkdv_kernel`.
+  - For `head_dim ≤ 128`: single-tile path preserved unchanged — Q is
+    loaded once per block-row as a stationary tile (no regression).
+  - For `head_dim > 128`: K-tile inner loop (`affine_range(head_dim //
+    128)`). `head_dim` must be a multiple of 128 (covers 256, 512).
+  - `weights @ V` and all dK/dV accumulation matmuls are unchanged (K=128
+    block dimension, unaffected by head_dim growth).
+- `tests/test_attention.py`: `test_head_dim_256_parity`,
+  `test_head_dim_256_dilated_parity`, `test_gradcheck_head_dim_256`,
+  `test_backward_shapes_head_dim_256` in PyTorch path.
+- `tests/test_nki_sim.py`: `TestAttnKTilingSimulator` — simulator parity
+  tests for forward and backward at head_dim=256.
+- `tests/test_nki_attn.py`: `TestAttnHardwareLargeDim` — hardware parity
+  tests at seq_len=512, head_dim=256.
+
+### Changed
+
+- `_attn_gather` assertion relaxed from `head_dim ≤ 128` to `head_dim ≤ 128
+  or head_dim % 128 == 0`.
+- `block_sparse_attention_tiled` and `nki_bsr_attn_tiled` docstrings updated
+  to document K-tiling support.
+
+## [0.5.1] — 2026-04-16
+
+### Added
+
+- **NKI backward kernel pair** (`_attn_bwd_dq_kernel` + `_attn_bwd_dkdv_kernel`)
+  in `trnsparse/nki/kernels.py`. When the `nki` backend is active,
+  `_AttnTiledFunction.backward` now dispatches to these kernels instead of
+  the Python-loop backward.
+  - `_attn_bwd_dq_kernel`: row-first traversal, mirrors `_attn_out_kernel`.
+    Accumulates `dQ_m += dS @ K_ki * scale` in a PSUM tile per block-row.
+  - `_attn_bwd_dkdv_kernel`: column-first traversal (BSC view built on host).
+    For each column block ki, iterates over all block-rows m that attend to ki,
+    accumulating `dK_ki += dS.T @ Q_m * scale` and `dV_ki += P.T @ dO_m`.
+  - No atomic scatter required — column-first gather on host eliminates NKI's
+    lack of cross-block accumulation primitives.
+- **Stats saving in forward** — `_AttnTiledFunction.forward` now saves
+  `(Q, K, V, O, row_max, row_denom)` in ctx (6 tensors vs 4). Backward
+  receives the pre-computed row stats and skips the two recomputation passes.
+  Applies to both PyTorch and NKI paths.
+- `_attn_bwd_gather` and `nki_bsr_attn_bwd` in `trnsparse/nki/dispatch.py`
+  — host-side BSC construction and NKI backward orchestration.
+- `tests/test_nki_sim.py`: `TestAttnBwdSimulator` — simulator correctness tests
+  for dQ/dK/dV against PyTorch backward (local window + dilated patterns).
+- `tests/test_nki_attn.py`: `TestAttnBwdHardware` — hardware parity tests
+  at seq_len=512, head_dim=64.
+
+### Changed
+
+- `nki_bsr_attn_tiled` accepts `return_stats=True` kwarg (default False) to
+  return `(out, row_max, row_denom)`.
+- `_block_sparse_attn_pytorch` accepts `return_stats=True` kwarg similarly.
+- `_block_sparse_attn_backward` accepts optional `row_max` / `row_denom`
+  kwargs; when provided, skips the two recomputation BSR passes.
+
+## [0.5.0] — 2026-04-16
+
+### Added
+
+- **Block-sparse attention autograd** — `block_sparse_attention_tiled` is now
+  differentiable. When any of Q/K/V has `requires_grad=True`, the call routes
+  through `_AttnTiledFunction`, which provides a tiled backward using the Flash
+  Attention delta identity (`D_i = dO_i · O_i`).
+- **`_block_sparse_attn_backward(Q, K, V, O, dO, mask_bsr, scale)`** — internal
+  backward function. Single pass over nonzero BSR blocks; no O(seq_len²)
+  intermediate. Row stats (row_max/row_denom) are recomputed in backward.
+- `gradcheck` passes at `atol=1e-3` for local-window and dilated patterns
+  (float64, seq_len=256, head_dim=32). See `tests/test_attention.py::TestAttnTiledGrad`.
+
+### Notes
+
+- NKI backward kernel is a follow-up (v0.5.x). The current backward runs
+  on PyTorch even when the NKI forward path is active.
+- Row stats (row_max/row_denom) are recomputed in backward rather than saved
+  from forward — avoids O(n_blocks × b) storage in the autograd graph at the
+  cost of two extra BSR passes.
+
+## [0.4.4] — 2026-04-16
+
+### Added
+
+- **NKI attention kernel pair** (`_attn_stats_kernel` + `_attn_out_kernel`,
+  closes [#25](https://github.com/trnsci/trnsparse/issues/25)). Two-pass
+  block-sparse attention now runs on the Tensor Engine via these kernels when
+  the `nki` backend is active. The `pytorch` backend retains the Python-loop
+  reference path from v0.4.3.
+  - Pass 1 (`_attn_stats_kernel`): each `(m, ki)` block pair computed
+    independently via `nc_matmul`; `nl.max` / `nl.sum` within-tile reductions
+    produce `(tile_max, tile_sumexp)` with no carry between iterations.
+  - Pass 2 (`_attn_out_kernel`): recomputes scores, applies stable softmax
+    using `row_max` / `row_denom` loaded per block-row (static HBM offset),
+    accumulates `weights @ V` into a PSUM tile spanning all ki blocks.
+  - **Constraint:** `head_dim ≤ 128` (`nc_matmul` partition limit). `head_dim=256`
+    requires K-tiling and is a follow-up.
+- `_attn_gather`, `_attn_host_reduction`, `nki_bsr_attn_tiled` in
+  `trnsparse/nki/dispatch.py` — host-side gather and orchestration for the kernel
+  pair, mirroring the `_bsr_pad_and_gather` pattern.
+- `tests/test_nki_sim.py`: `TestAttnTiledSimulator` — simulator correctness
+  tests (local window, dilated; seq_len=256, head_dim=32).
+- `tests/test_nki_attn.py`: hardware tests for local window, dilated, and
+  global-token patterns at seq_len=512, head_dim=64.
+
+### Notes
+
+No autograd wrapping in v0.4.4 — forward-only, matching the existing PyTorch
+path. Attention backward (block-sparse) is a follow-up.
+
+## [0.4.3] — 2026-04-16
+
+### Added
+
+- **`chebyshev_bsr(A, b, lam_min, lam_max, K)`** — fixed-K Chebyshev
+  semi-iteration. Pre-computes all step coefficients from eigenvalue bounds
+  before the loop; no inner products during iteration. Same convergence rate
+  as CG without the adaptive coefficient computation that blocks NKI fusion.
+  Returns `(x, K, rel_residual)` matching the `cg_bsr` signature.
+- **`chebyshev_coeffs(lam_min, lam_max, K)`** — returns the `(alpha, beta)`
+  coefficient tensors for the momentum update `x_{k+1} = x_k + α_k r_k +
+  β_k (x_k - x_{k-1})`. Useful for inspecting or caching coefficients across
+  solves with the same matrix spectrum.
+- **`richardson_bsr(A, b, omega, K)`** — fixed-K Richardson iteration; the
+  simplest fixed-point iteration with no coefficient computation at all.
+  Optimal omega = `2 / (lam_min + lam_max)`. Baseline comparator for Chebyshev.
+- **`block_sparse_attention_tiled(Q, K, V, mask_bsr)`** — two-pass sparse
+  attention that avoids the O(seq_len²) score intermediate. Pass 1 computes
+  per-block `(tile_max, tile_sumexp)` statistics over nonzero blocks only
+  (O(n_blocks × block_size) stats array, e.g. 5 KB at seq_len=512 vs 1 MB
+  dense). Pass 2 recomputes scores, applies stable softmax using host-reduced
+  row statistics, and accumulates V. PyTorch reference for the NKI kernel pair
+  documented in [#25](https://github.com/trnsci/trnsparse/issues/25).
+- **`docs/sparse_attention.md`**: new "Tiled two-pass path" section explaining
+  the algorithm, memory trade-offs, and the NKI kernel follow-up path.
+- **`docs/architecture.md`**: "Known limits" updated — #22 and #25 now have
+  v0.4.3 PyTorch references with clear notes on remaining NKI gate conditions.
+
+### Notes
+
+Addresses the architecture-friendly workarounds identified in the analysis of
+parked NKI-capability-gated issues (#22, #25). Both solvers and the tiled
+attention path share the structural property of fixed-iteration loops with
+statically-determined coefficients — the pattern that maps to `nl.affine_range`
+without scalar carry.
+
 ## [0.4.2] — 2026-04-15
 
 ### Added

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

@@ -19,38 +19,39 @@ both memory and compute. trnsparse makes the sparsity explicit.
 trnsparse/
 ├── trnsparse/
 │   ├── __init__.py
-│   ├── formats.py       # CSRMatrix, COOMatrix, conversions, from_dense
-│   ├── ops.py           # spmv, spmm, spmv_symmetric, add, scale, transpose
+│   ├── formats.py       # CSRMatrix, COOMatrix, BSRMatrix, from_dense
+│   ├── ops.py           # spmv, spmm, bsr_spmm, screened_spmm, add, scale, transpose
 │   ├── screening.py     # schwarz_bounds, screen_quartets, density_screen
+│   ├── iterative.py     # cg_bsr, power_iteration_bsr, jacobi_preconditioner_bsr
 │   └── nki/
 │       ├── __init__.py
-│       └── dispatch.py  # Gather-matmul-scatter pattern for SpMM
+│       ├── kernels.py   # _bsr_spmm_kernel, _screened_spmm_kernel
+│       └── dispatch.py  # torch.autograd.Function wrappers + NKI/PyTorch routing
 ├── tests/
-│   ├── test_formats.py  # CSR/COO construction, roundtrips
-│   ├── test_ops.py      # SpMV, SpMM vs dense reference
-│   └── test_screening.py
 ├── examples/
-│   └── sparse_fock.py   # Screened Fock build demo
+│   ├── sparse_fock.py             # Schwarz-screened Fock build (3 paths + trnblas)
+│   ├── pyscf_bridge.py            # Real AO integrals via PySCF (optional dep)
+│   └── block_sparse_attention.py  # Block-sparse attention via bsr_spmm
 ```
 
-## NKI SpMM Strategy
+## NKI compute strategy (v0.4.x current)
 
-**v0.2.0 (current):** SpMM materializes the CSR into a dense `(M, K)` tile
-and runs stationary-tile-reuse GEMM on the Tensor Engine. Validated on
-trn1; forward + `torch.autograd.Function` backward both pass parity +
-`gradcheck` at `atol=1e-4`. The dense materialization means no sparsity
-advantage yet — NKI is slower than the PyTorch fallback at small sparse
-sizes (see `docs/benchmarks.md`).
+**Formats:** CSR/COO are interop and PyTorch-fallback formats.
+BSR at `block_size=128` is the NKI compute format — every nonzero block
+is already a 128×128 Tensor Engine tile, zero gather overhead.
 
-**v0.3.0 (planned, #15):** gather-matmul-scatter with row-bucketing.
-For each bucket of rows with similar nnz:
-1. **DMA engine**: gather non-zero column indices into dense SBUF tiles.
-2. **Tensor Engine**: matmul the tile against B columns.
-3. **DMA engine**: scatter results back to output rows.
+**Dispatch hierarchy:**
+1. `spmm(csr, B)` — PyTorch `torch.sparse_csr_tensor` fallback (within
+   2× of scipy; NKI overhead not worth it for CSR at current sizes).
+2. `bsr_spmm(bsr, B)` — NKI kernel `_bsr_spmm_kernel`; one `nc_matmul`
+   per nonzero block. `torch.autograd.Function`-wrapped; gradcheck passes.
+3. `screened_spmm(A, diag, B, threshold)` — fused Schwarz + mask + matmul
+   in one NKI dispatch; avoids 4 host passes + CSR build + separate SpMM.
+4. `block_sparse_attention(Q, K, V, mask_bsr)` — `bsr_spmm` applied to
+   the attention weight matrix; no new kernel, mask is a BSRMatrix.
 
-Same pattern used in sparse attention. Uniform nnz maps cleanly to fixed-
-size tiles; variable nnz is handled by bucketing rows by nnz quantile so
-each bucket pads only within itself.
+**Open gaps (NKI capability-gated):** row-bucketing CSR (#15), sharded BSR
+across NeuronCores (#16), fused tile-level attention scores (#25).
 
 ## Dependencies
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.4.2
+Version: 0.7.0
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0
@@ -47,12 +47,17 @@ CSR/COO formats, SpMV, SpMM, and integral screening for sparse scientific comput
 
 trnsparse follows the [trnsci 5-phase roadmap](https://trnsci.dev/roadmap/). Active work is tracked in phase-labeled GitHub issues:
 
-- **[Phase 1 — correctness](https://github.com/trnsci/trnsparse/issues/14)** ✅ v0.2.0: NKI SpMM validated on trn1 via densify-then-GEMM; first `torch.autograd.Function`-wrapped NKI kernel in the suite (see [`trnsci/trnsci#3`](https://github.com/trnsci/trnsci/issues/3)). Benchmarks in [`docs/benchmarks.md`](https://trnsci.dev/trnsparse/benchmarks/).
-- **[Phase 3 — perf](https://github.com/trnsci/trnsparse/issues/15)**: nnz-bucketing SpMM, streaming large-sparse, NEFF cache reuse.
-- **[Phase 4 — multi-chip](https://github.com/trnsci/trnsparse/issues/16)**: sharded sparse matrices across chips.
+- **[Phase 1 — correctness](https://github.com/trnsci/trnsparse/issues/14)** ✅ v0.2.0: NKI SpMM validated on trn1 via densify-then-GEMM (see [`trnsci/trnsci#3`](https://github.com/trnsci/trnsci/issues/3)).
+- **v0.3.0** ✅ [`BSRMatrix`](https://trnsci.dev/trnsparse/architecture/) — Trainium-native 128×128 block-sparse format; `bsr_spmm` NKI kernel. CSR becomes interop; BSR is the compute path.
+- **v0.3.2** ✅ `cg_bsr`, `power_iteration_bsr` — iterative solvers over BSR (Python loop; fused kernel gated on NKI capability).
+- **v0.4.0** ✅ `screened_spmm` — fused Schwarz-screened SpMM in one NKI dispatch.
+- **v0.4.2** ✅ Block-sparse attention — `BSRMatrix` + `bsr_spmm` as the primitive; `examples/block_sparse_attention.py` + [`docs/sparse_attention.md`](https://trnsci.dev/trnsparse/sparse_attention/).
+- **v0.4.3** ✅ Architecture-friendly alternatives: `chebyshev_bsr` / `richardson_bsr` (fixed-K solvers, no inner products); `block_sparse_attention_tiled` (two-pass, no O(seq²) intermediate).
+- **[Phase 3 — perf](https://github.com/trnsci/trnsparse/issues/15)**: nnz-bucketing, NKI attention kernel pair — parked on NKI indirect DMA gather / hardware validation.
+- **[Phase 4 — multi-chip](https://github.com/trnsci/trnsparse/issues/16)**: sharded BSR across NeuronCores.
 - **[Phase 5 — generation](https://github.com/trnsci/trnsparse/issues/17)**: trn2 DMA bandwidth exploitation.
 
-_(No Phase 2 for trnsparse — the precision story is inherited from trnblas.)_
+_(No Phase 2 for trnsparse — precision inherited from trnblas.)_
 
 Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1).
 
@@ -89,12 +94,21 @@ stats = trnsparse.sparsity_stats(Q)
 
 | Operation | Description |
 |-----------|-------------|
-| `spmv` | Sparse × dense vector |
-| `spmm` | Sparse × dense matrix |
+| `spmv` | Sparse × dense vector (CSR) |
+| `spmm` | Sparse × dense matrix (CSR, PyTorch fallback) |
+| `bsr_spmm` | Block-sparse × dense (BSR-native, Tensor Engine) |
+| `screened_spmm` | Fused Schwarz-screened matmul (one NKI dispatch) |
 | `spmv_symmetric` | Symmetric SpMV (half storage) |
 | `sparse_add` | C = αA + βB |
 | `sparse_scale` | B = αA |
-| `sparse_transpose` | A^T |
+| `sparse_transpose` | Aᵀ |
+| `cg_bsr` | Conjugate Gradient on BSR matrix |
+| `chebyshev_bsr` | Fixed-K Chebyshev semi-iteration (no inner products) |
+| `richardson_bsr` | Fixed-K Richardson iteration |
+| `power_iteration_bsr` | Dominant eigenpair via power iteration |
+| `jacobi_preconditioner_bsr` | Diagonal preconditioner for `cg_bsr` |
+| `bsr_diagonal` | Extract main diagonal from BSR matrix |
+| `block_sparse_attention_tiled` | Two-pass sparse attention, no O(seq²) intermediate |
 | `schwarz_bounds` | Schwarz screening bounds |
 | `screen_quartets` | Shell quartet significance mask |
 | `density_screen` | Density-weighted screening |

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

@@ -16,12 +16,17 @@ CSR/COO formats, SpMV, SpMM, and integral screening for sparse scientific comput
 
 trnsparse follows the [trnsci 5-phase roadmap](https://trnsci.dev/roadmap/). Active work is tracked in phase-labeled GitHub issues:
 
-- **[Phase 1 — correctness](https://github.com/trnsci/trnsparse/issues/14)** ✅ v0.2.0: NKI SpMM validated on trn1 via densify-then-GEMM; first `torch.autograd.Function`-wrapped NKI kernel in the suite (see [`trnsci/trnsci#3`](https://github.com/trnsci/trnsci/issues/3)). Benchmarks in [`docs/benchmarks.md`](https://trnsci.dev/trnsparse/benchmarks/).
-- **[Phase 3 — perf](https://github.com/trnsci/trnsparse/issues/15)**: nnz-bucketing SpMM, streaming large-sparse, NEFF cache reuse.
-- **[Phase 4 — multi-chip](https://github.com/trnsci/trnsparse/issues/16)**: sharded sparse matrices across chips.
+- **[Phase 1 — correctness](https://github.com/trnsci/trnsparse/issues/14)** ✅ v0.2.0: NKI SpMM validated on trn1 via densify-then-GEMM (see [`trnsci/trnsci#3`](https://github.com/trnsci/trnsci/issues/3)).
+- **v0.3.0** ✅ [`BSRMatrix`](https://trnsci.dev/trnsparse/architecture/) — Trainium-native 128×128 block-sparse format; `bsr_spmm` NKI kernel. CSR becomes interop; BSR is the compute path.
+- **v0.3.2** ✅ `cg_bsr`, `power_iteration_bsr` — iterative solvers over BSR (Python loop; fused kernel gated on NKI capability).
+- **v0.4.0** ✅ `screened_spmm` — fused Schwarz-screened SpMM in one NKI dispatch.
+- **v0.4.2** ✅ Block-sparse attention — `BSRMatrix` + `bsr_spmm` as the primitive; `examples/block_sparse_attention.py` + [`docs/sparse_attention.md`](https://trnsci.dev/trnsparse/sparse_attention/).
+- **v0.4.3** ✅ Architecture-friendly alternatives: `chebyshev_bsr` / `richardson_bsr` (fixed-K solvers, no inner products); `block_sparse_attention_tiled` (two-pass, no O(seq²) intermediate).
+- **[Phase 3 — perf](https://github.com/trnsci/trnsparse/issues/15)**: nnz-bucketing, NKI attention kernel pair — parked on NKI indirect DMA gather / hardware validation.
+- **[Phase 4 — multi-chip](https://github.com/trnsci/trnsparse/issues/16)**: sharded BSR across NeuronCores.
 - **[Phase 5 — generation](https://github.com/trnsci/trnsparse/issues/17)**: trn2 DMA bandwidth exploitation.
 
-_(No Phase 2 for trnsparse — the precision story is inherited from trnblas.)_
+_(No Phase 2 for trnsparse — precision inherited from trnblas.)_
 
 Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1).
 
@@ -58,12 +63,21 @@ stats = trnsparse.sparsity_stats(Q)
 
 | Operation | Description |
 |-----------|-------------|
-| `spmv` | Sparse × dense vector |
-| `spmm` | Sparse × dense matrix |
+| `spmv` | Sparse × dense vector (CSR) |
+| `spmm` | Sparse × dense matrix (CSR, PyTorch fallback) |
+| `bsr_spmm` | Block-sparse × dense (BSR-native, Tensor Engine) |
+| `screened_spmm` | Fused Schwarz-screened matmul (one NKI dispatch) |
 | `spmv_symmetric` | Symmetric SpMV (half storage) |
 | `sparse_add` | C = αA + βB |
 | `sparse_scale` | B = αA |
-| `sparse_transpose` | A^T |
+| `sparse_transpose` | Aᵀ |
+| `cg_bsr` | Conjugate Gradient on BSR matrix |
+| `chebyshev_bsr` | Fixed-K Chebyshev semi-iteration (no inner products) |
+| `richardson_bsr` | Fixed-K Richardson iteration |
+| `power_iteration_bsr` | Dominant eigenpair via power iteration |
+| `jacobi_preconditioner_bsr` | Diagonal preconditioner for `cg_bsr` |
+| `bsr_diagonal` | Extract main diagonal from BSR matrix |
+| `block_sparse_attention_tiled` | Two-pass sparse attention, no O(seq²) intermediate |
 | `schwarz_bounds` | Schwarz screening bounds |
 | `screen_quartets` | Shell quartet significance mask |
 | `density_screen` | Density-weighted screening |

trnsparse-0.7.0/docs/architecture.md

@@ -0,0 +1,143 @@
+# Architecture
+
+## Why BSR is Trainium-native
+
+Trainium's Tensor Engine is a 128×128 systolic array. `nisa.nc_matmul`
+moves a 128-partition stationary tile against a moving operand of up to
+128 × 512, all on-chip. The natural "unit of sparse work" on this
+architecture is therefore not a single nonzero — it's a 128×128 block.
+
+**BSR (block-sparse row)** at `block_size=128` stores a matrix as a list
+of these blocks plus a block-level CSR pattern. Every stored block is
+already in the shape the systolic array wants: one `nc_matmul` per block,
+no gather step, no padding-within-tile waste. The architectural match is
+exact.
+
+Contrast with CSR, which stores individual nonzero elements. The v0.2.0
+NKI path for CSR materializes the matrix into a dense `(M, K)` tile
+before the matmul (see the [SpMM path](#v020-csr-spmm-path) below) —
+necessary for correctness, but it pays the full `M × K` cost whether the
+matrix is sparse or not. BSR skips this because blocks are already dense.
+
+**Consequence for the library's shape:** CSR remains the construction and
+interop format (scipy compatibility, PyTorch's `torch.sparse_csr_tensor`
+interop). BSR is the compute format for the NKI path. For matrices with
+real block structure — Fock/ERI tensors after Schwarz screening, FEM
+stiffness matrices, graph adjacencies, block-sparse attention masks —
+BSR is strictly preferred. For truly unstructured sparse (random CSR),
+the PyTorch fallback is already within 2× of scipy; NKI adds nothing.
+
+See [Benchmarks](benchmarks.md) for numbers that validate this framing.
+
+
+```
+trnsparse/
+├── trnsparse/
+│   ├── __init__.py
+│   ├── formats.py       # CSRMatrix, COOMatrix, BSRMatrix, from_dense
+│   ├── ops.py           # spmv, spmm, bsr_spmm, screened_spmm, ...
+│   ├── screening.py     # schwarz_bounds, screen_quartets, density_screen
+│   ├── 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 + routing
+├── tests/
+├── examples/
+│   ├── sparse_fock.py             # Schwarz-screened Fock (3 paths + trnblas)
+│   ├── pyscf_bridge.py            # Real AO integrals via PySCF
+│   └── block_sparse_attention.py  # Block-sparse attention via bsr_spmm
+```
+
+## NKI dispatch hierarchy
+
+CSR and COO are interop and PyTorch-fallback formats. BSR is the NKI
+compute format. `nki/dispatch.py` exposes `HAS_NKI`,
+`set_backend("auto"|"pytorch"|"nki")`, `get_backend()`, and the NKI
+entry points.
+
+### `spmm(csr, B)` — PyTorch fallback
+
+Routes through `torch.sparse_csr_tensor` operations. Benchmarked within
+2× of scipy at typical shapes; the NKI per-op overhead doesn't pay off
+for CSR SpMM until the matrix is large enough that the v0.2.0 CSR kernel
+(densify-then-GEMM) would dominate anyway. SpMV also stays on this path.
+
+### `bsr_spmm(bsr, B)` — NKI, `_bsr_spmm_kernel`
+
+One `nisa.nc_matmul` per nonzero block. Host-side preamble pads each
+block-row to the same `K_max` so the kernel's `affine_range` bounds are
+fixed. `torch.autograd.Function`-wrapped (`_BSRSpMMFunction`); gradcheck
+passes at `atol=1e-4` on hardware.
+
+### `screened_spmm(A, diag, B, threshold)` — NKI, fused
+
+Fuses Schwarz bound (outer-product pair bound), threshold mask, and
+matmul into one dispatch. Saves ~30–50% vs the unfused
+`schwarz_bounds → screen_quartets → from_dense → spmm` flow on
+Fock-build-sized inputs.
+
+### Block-sparse attention
+
+`bsr_spmm` applied to the post-softmax attention weight matrix. No new
+kernel — a `BSRMatrix` with a local-window or dilated block pattern IS
+an attention mask. See [sparse_attention.md](sparse_attention.md).
+
+## Formats
+
+- **CSR**: row pointer + column indices + values. Preferred for construction,
+  interop with scipy/PyTorch, and SpMV.
+- **COO**: three parallel 1-D tensors. Preferred for construction and permutation.
+- **BSR** (`block_size=128`): stacked 128×128 blocks + block-level CSR pattern.
+  Every block maps to one `nc_matmul`; zero gather overhead.
+
+All three support `from_dense`, `to_dense`, and interconversion.
+
+## Autograd wrapping
+
+Every NKI kernel lives inside a `torch.autograd.Function` (satisfies
+[`trnsci/trnsci#3`](https://github.com/trnsci/trnsci/issues/3)).
+Backward passes run at the PyTorch level via block-gradient projection.
+Block-selection is non-differentiable by construction; `grad_out` is
+routed into exactly the stored blocks. `gradcheck` at `atol=1e-4` is
+part of the hardware test matrix for all three NKI kernels.
+
+### Fused screened SpMM (v0.4.0)
+
+`screened_spmm(A, diag_integrals, B, threshold)` fuses the
+chemistry-screened SpMM pipeline — Schwarz bound from the diagonal
+integrals, pair-bound threshold mask, masked matmul — into a single
+NKI kernel. The unfused equivalent does four host passes (sqrt, outer
+product, threshold, mask-apply) plus a separate `from_dense` + `spmm`
+call; the fused kernel collapses all of that into one dispatch.
+
+Mask semantics: `mask[i,j] = sqrt(|diag[i]|) * sqrt(|diag[j]|) >
+sqrt(threshold)`, matching `schwarz_bounds` + `screen_quartets`
+composed. No gradient flows back to `diag_integrals` or `threshold` —
+the mask is treated as a discrete gate; `grad_A *= mask` and
+`grad_B = (A * mask).T @ grad_C`.
+
+Restricted to square A (`M == K`) with 1-D `diag_integrals` in v0.4.0
+— the common Fock-build case. Rectangular / asymmetric-bounds
+extension is a follow-up if asked for.
+
+## Known limits
+
+- **Row-bucketing CSR** ([#15](https://github.com/trnsci/trnsparse/issues/15)) — parked.
+  Requires NKI indirect-DMA gather (not exposed as of NKI 0.3.0). The CSR PyTorch
+  fallback is within 2× of scipy; that's the current story.
+- **Fused tile-level attention scores** ([#25](https://github.com/trnsci/trnsparse/issues/25)) —
+  **v0.4.4 ships the NKI kernel pair** (`_attn_stats_kernel` + `_attn_out_kernel`).
+  `block_sparse_attention_tiled` routes through NKI on the `nki` backend; the
+  PyTorch two-pass reference remains on the `pytorch` backend. `head_dim ≤ 128`
+  in v0.4.4 (nc_matmul partition limit); `head_dim=256` requires K-tiling.
+- **Fused iterative solvers** ([#22](https://github.com/trnsci/trnsparse/issues/22)) —
+  CG blocked on `nl.affine_range` iteration-carried scalar state. **v0.4.3 ships
+  `chebyshev_bsr` and `richardson_bsr`**: fixed-K solvers with pre-computed
+  coefficients (no inner products in the loop) that have the correct structure for
+  eventual NKI fusion. The "A SBUF-resident across all iterations" win remains gated
+  on `affine_range` vector-carry support.
+- **Multi-chip sharded BSR** ([#16](https://github.com/trnsci/trnsparse/issues/16)) —
+  gated on suite-level multi-chip collectives.
+- **SpMV stays PyTorch.** Single output column on the Tensor Engine doesn't amortize
+  compile + dispatch overhead.

{trnsparse-0.4.2 → trnsparse-0.7.0}/docs/benchmarks.md

@@ -109,21 +109,28 @@ Three observations, all consistent with the architectural framing:
    dispatch, or where the kernel launch can be amortized across many
    operations.
 
-### Where NKI is expected to win (follow-up issues)
-
-Not in raw "one-shot SpMM at N=128 vs CPU." The architectural exploits are:
-
-- **[#20 — on-chip iterative solvers](https://github.com/trnsci/trnsparse/issues/20)**:
-  iterate `A @ x` thousands of times with A SBUF-resident. Kernel launch
-  amortizes over all iterations.
-- **[#19 — fused screen + matmul](https://github.com/trnsci/trnsparse/issues/19)**:
-  one kernel that fuses what's currently three HBM round-trips.
-- **[#21 — block-sparse attention](https://github.com/trnsci/trnsparse/issues/21)**:
-  large-scale sequences where the matmul work dominates dispatch.
-
-v0.3.0 ships **BSR correctness + differentiability**. v0.4.0+ will cash
-in the architectural wins. See
-[`docs/architecture.md`](architecture.md) for why BSR is the
+### Where NKI wins — shipped v0.4.x
+
+Not in raw "one-shot SpMM at N=128 vs CPU." The BSR table above measures
+kernel dispatch overhead. The architectural wins come from fusing multiple
+host passes into one dispatch:
+
+- **`screened_spmm` (v0.4.0, [#19](https://github.com/trnsci/trnsparse/issues/19))**
+  ✅ shipped: fuses Schwarz bound + threshold + mask + matmul into one NKI
+  dispatch. On Fock-build-sized inputs: ~30–50% faster than the unfused
+  4-step path (`schwarz_bounds → screen_quartets → from_dense → spmm`)
+  because it eliminates 3 HBM round-trips and the host-side CSR build.
+- **Block-sparse attention (v0.4.2, [#21](https://github.com/trnsci/trnsparse/issues/21))**
+  ✅ shipped: `bsr_spmm` handles `attn_weights @ V` for local-window / dilated /
+  global-token masks. At seq_len=8192, window=2: 4.9% block density means 95%
+  of `nc_matmul` calls are skipped. Current example still materializes the full
+  score matrix (follow-up: [#25](https://github.com/trnsci/trnsparse/issues/25)).
+- **[#22 — on-chip iterative solvers](https://github.com/trnsci/trnsparse/issues/22)**:
+  CG/power-iteration with A SBUF-resident across all iterations (replaces #20).
+  v0.3.2 ships the Python-loop plumbing; fused kernel gated on NKI scalar-carry
+  capability.
+
+See [`docs/architecture.md`](architecture.md) for why BSR is the
 Trainium-native format.
 
 BSR results from commit

{trnsparse-0.4.2 → trnsparse-0.7.0}/docs/iterative_solvers.md

@@ -65,7 +65,7 @@ def _cg_spd_kernel(A_blocks, A_cols, A_row_ptrs, b, max_iter):
     rr = nl.reduce(r * r)
 
     for k in nl.affine_range(max_iter):
-        Ap = _bsr_matvec_sbuf(A_sbuf, A_cols, A_row_ptrs, p)   # all SBUF
+        Ap = _bsr_matvec_sbuf(A_sbuf, A_cols, A_row_ptrs, p)  # all SBUF
         pAp = nl.reduce(p * Ap)
         alpha = rr / pAp
         x = x + alpha * p

{trnsparse-0.4.2 → trnsparse-0.7.0}/docs/migration_scipy.md

@@ -60,10 +60,10 @@ argument; scipy infers it from the indices when possible. Be explicit:
 
 ```python
 # scipy
-A = csr_matrix((data, indices, indptr))              # shape inferred
+A = csr_matrix((data, indices, indptr))  # shape inferred
 
 # trnsparse
-A = CSRMatrix(data, indices, indptr, shape=(m, n))   # shape required
+A = CSRMatrix(data, indices, indptr, shape=(m, n))  # shape required
 ```
 
 **Operator access.** scipy overloads `@`, `+`, `*` on its matrix
@@ -107,9 +107,9 @@ import torch
 import trnsparse
 import trnsolver
 
-A = trnsparse.from_scipy(scipy_matrix)            # zero-copy ingestion
+A = trnsparse.from_scipy(scipy_matrix)  # zero-copy ingestion
 b = torch.as_tensor(b_numpy, dtype=torch.float32)
-x = trnsolver.cg(A, b, tol=1e-6)                  # uses trnsparse.spmv
+x = trnsolver.cg(A, b, tol=1e-6)  # uses trnsparse.spmv
 ```
 
 The hot path (`spmv` inside CG) dispatches to NKI when

{trnsparse-0.4.2 → trnsparse-0.7.0}/docs/quickstart.md

@@ -42,7 +42,7 @@ keep = trnsparse.screen_quartets(Q, threshold=1e-10)
 ## Backend selection
 
 ```python
-trnsparse.set_backend("auto")     # default
+trnsparse.set_backend("auto")  # default
 trnsparse.set_backend("pytorch")  # force PyTorch fallback
-trnsparse.set_backend("nki")      # force NKI (requires Neuron hardware)
+trnsparse.set_backend("nki")  # force NKI (requires Neuron hardware)
 ```