trnsparse 0.4.2__tar.gz → 0.6.0__tar.gz

{trnsparse-0.4.2 → trnsparse-0.6.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.6.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.6.0}/CHANGELOG.md

@@ -5,6 +5,159 @@ 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).
 
+## [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.6.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.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.4.2
+Version: 0.6.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.6.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.6.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.6.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.6.0}/docs/sparse_attention.md

@@ -142,21 +142,88 @@ The win is at long sequences. At `seq_len=1024`, a window-2 mask still stores
 the per-block compute becomes the bottleneck, which is where the Tensor Engine
 thrives.
 
-## What's next
+## Tiled two-pass path (v0.4.3)
 
-The current path materializes the full `(seq_len, seq_len)` score matrix to
-compute attention weights. A production path would:
+`block_sparse_attention_tiled` eliminates the `O(seq_len²)` score intermediate
+without any new NKI primitives. The algorithm uses two passes over the nonzero
+blocks:
 
-1. Iterate over nonzero blocks in `mask_bsr`.
-2. For each block `(i, j)`, load `Q[i*b:(i+1)*b]` and `K[j*b:(j+1)*b]` into
-   SBUF, compute the score tile, and apply softmax over the block-row.
-3. Multiply the score tile by `V[j*b:(j+1)*b]` and accumulate into the output.
+**Pass 1 — per-block statistics**: For each stored block `(m, ki)`, compute
+the score tile `Q[m] @ K[ki].T` and extract its row-wise max (`tile_max`,
+shape `(block_size,)`) and stable row-wise exp-sum (`tile_sumexp`). Every
+`(m, ki)` pair is independent — no iteration-carried state.
 
-This fused tile-level kernel avoids the `O(seq_len²)` intermediate entirely.
-It's the same architectural opportunity as [on-chip iterative solvers](iterative_solvers.md):
-load A once, iterate on-chip. The NKI building block is available (the BSR
-kernel in `nki/kernels.py`); the row-wise softmax over variable numbers of
-tiles is the authoring challenge.
+**Host reduction**: For each block-row m, take the max of `tile_max` across all
+its nonzero k-blocks (`row_max`) and compute the corrected denominator
+`row_denom`. This is O(n_stored_blocks × block_size) work — at seq_len=4096
+with window=2, about 160 block pairs × 128 = 20K operations, negligible.
 
-A runnable reference for the current (non-fused) path is in
+**Pass 2 — stable softmax + accumulation**: Recompute each score tile, apply
+`exp(score - row_max) / row_denom`, and accumulate `weights @ V_block` into
+the output via a PSUM-style loop.
+
+```python
+out = trnsparse.block_sparse_attention_tiled(Q, K, V, mask_bsr)
+```
+
+Memory: `n_stored_blocks × block_size × 4` bytes for stats (5 KB at seq_len=512,
+window=1) vs `seq_len² × 4` bytes for the naive dense score matrix (1 MB at
+seq_len=512, 64 MB at seq_len=4096).
+
+**On CPU**, the Python-level loops make this path slower than the vectorized
+naive implementation — that's expected. The value of `block_sparse_attention_tiled`
+is algorithmic: it validates the two-pass decomposition that the NKI kernel pair
+(`_attn_stats_kernel` + `_attn_out_kernel`) will implement on hardware, where
+per-block `nc_matmul` calls amortize the dispatch cost and the O(seq_len²)
+allocation is the binding constraint.
+
+The NKI implementation path is documented in
+[#25](https://github.com/trnsci/trnsparse/issues/25).
+
+## NKI kernel pair (v0.4.4)
+
+`block_sparse_attention_tiled` now routes through the NKI kernel pair on the
+`nki` backend. The `pytorch` backend retains the Python-loop reference path.
+
+`_attn_stats_kernel` (pass 1) and `_attn_out_kernel` (pass 2) implement the
+two-pass decomposition in `trnsparse/nki/kernels.py`:
+
+- **Pass 1** (`_attn_stats_kernel`): each `(m, ki)` block pair is independent;
+  no carry. Uses `nl.max` and `nl.sum` reductions within the 128×128 SBUF tile.
+  Q block is loaded once as a transposed tile (stationary) and reused across all
+  ki blocks in the row.
+- **Pass 2** (`_attn_out_kernel`): recomputes scores from Q and K to avoid
+  storing the full score tensor; uses `row_max` / `row_denom` from the host
+  reduction (loaded by static `affine_range` offset — legal). Accumulates
+  `weights @ V` into a PSUM tile that spans all ki blocks.
+
+**Constraint:** `head_dim ≤ 128` in v0.4.4 (`nc_matmul` partition limit).
+`head_dim=256` requires K-tiling and is a follow-up.
+
+A runnable reference for both the naive and tiled paths is in
 [`examples/block_sparse_attention.py`](https://github.com/trnsci/trnsparse/blob/main/examples/block_sparse_attention.py).
+
+## Autograd (v0.5.0)
+
+`block_sparse_attention_tiled` is now differentiable. When any of Q, K, or V
+has `requires_grad=True`, the call is routed through `_AttnTiledFunction`,
+which provides a tiled backward using the Flash Attention delta identity.
+
+**Backward algorithm** (single pass over nonzero blocks, no O(seq_len²) intermediate):
+
+```
+D_i = dO_i · O_i                   # row-wise dot — Flash "delta"
+For each stored block (m, ki):
+    P   = stable_softmax(Q_m @ K_ki.T * scale)
+    dP  = dO_m @ V_ki.T
+    dS  = P * (dP - D_m)           # softmax backward
+    dQ_m  += dS @ K_ki * scale
+    dK_ki += dS.T @ Q_m * scale
+    dV_ki += P.T @ dO_m
+```
+
+`torch.autograd.gradcheck` passes at `atol=1e-3` for local-window and dilated
+patterns (float64, seq_len=256, head_dim=32).
+
+**NKI backward kernel** is a follow-up (v0.5.x). The current backward runs on
+CPU/PyTorch even when the forward used the NKI kernel pair.