trnsparse 0.4.1__tar.gz → 0.6.0__tar.gz

{trnsparse-0.4.1 → 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.1 → 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.1 → trnsparse-0.6.0}/CHANGELOG.md

@@ -5,6 +5,183 @@ 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
+
+- **`examples/block_sparse_attention.py`** — block-sparse attention
+  reference using `BSRMatrix` + `bsr_spmm`. Three mask patterns
+  (local window, dilated, global tokens); verifies against a dense
+  reference; reports block density and timing for the `bsr_spmm` step.
+  Closes [#21](https://github.com/trnsci/trnsparse/issues/21).
+- **`docs/sparse_attention.md`** — writeup: how BSR-128 maps to
+  Longformer/BigBird-style attention masks, block density arithmetic,
+  pattern construction helpers, and the fused-tile follow-up.
+- **`docs/iterative_solvers.md`** now linked in `mkdocs.yml` nav
+  (was present but missing from navigation).
+- **`tests/test_attention.py`** — 8 CPU tests: mask shape/symmetry
+  checks + parity against dense reference at `atol=1e-4` for all three
+  patterns and the full-attention edge case.
+
+### Notes
+
+No API changes, no kernel changes. The claim in #21 is: `bsr_spmm` is
+the block-sparse attention primitive; `BSRMatrix` captures the mask.
+The example and docs make that explicit.
+
 ## [0.4.1] — 2026-04-14
 
 ### Added

{trnsparse-0.4.1 → 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.1/trnsparse.egg-info → trnsparse-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.4.1
+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 |
@@ -102,3 +116,16 @@ stats = trnsparse.sparsity_stats(Q)
 ## License
 
 Apache 2.0 — Copyright 2026 Scott Friedman
+
+
+## Disclaimer
+
+trnsci is an **independent open-source project**. It is not sponsored by, endorsed by, or affiliated with Amazon.com, Inc., Amazon Web Services, Inc., or Annapurna Labs Ltd.
+
+"AWS", "Amazon", "Trainium", "Inferentia", "NeuronCore", "Neuron SDK", and related identifiers are trademarks of their respective owners and are used here solely for descriptive and interoperability purposes. Use does not imply endorsement, partnership, or any other relationship.
+
+All work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on [trnsci.dev](https://trnsci.dev) are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.
+
+Feedback directed at the Neuron SDK or Trainium hardware is good-faith ecosystem commentary from independent users. It is not privileged information, is not pre-reviewed by AWS, and should not be read as authoritative about product roadmap, behavior, or quality.
+
+For official AWS guidance, see [aws-neuron documentation](https://awsdocs-neuron.readthedocs-hosted.com/) and the [AWS Trainium product page](https://aws.amazon.com/ai/machine-learning/trainium/).

trnsparse-0.6.0/README.md

@@ -0,0 +1,100 @@
+# trnsparse
+
+[![CI](https://github.com/trnsci/trnsparse/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trnsparse/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/trnsci/trnsparse/graph/badge.svg)](https://codecov.io/gh/trnsci/trnsparse)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![PyPI](https://img.shields.io/pypi/v/trnsparse)](https://pypi.org/project/trnsparse/)
+[![Python](https://img.shields.io/pypi/pyversions/trnsparse)](https://pypi.org/project/trnsparse/)
+[![License](https://img.shields.io/github/license/trnsci/trnsparse)](LICENSE)
+[![Docs](https://img.shields.io/badge/docs-trnsci.dev-blue)](https://trnsci.dev/trnsparse/)
+
+Sparse matrix operations for AWS Trainium via NKI.
+
+CSR/COO formats, SpMV, SpMM, and integral screening for sparse scientific computing on Trainium. Part of the trnsci scientific computing suite ([github.com/trnsci](https://github.com/trnsci)).
+
+## Current phase
+
+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 (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 — precision inherited from trnblas.)_
+
+Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1).
+
+## Install
+
+```bash
+pip install trnsparse
+```
+
+## Usage
+
+```python
+import torch
+import trnsparse
+
+# Dense → sparse
+A = torch.randn(100, 100)
+A[torch.abs(A) < 1.0] = 0.0
+csr = trnsparse.from_dense(A)
+
+# SpMV: y = A @ x
+y = trnsparse.spmv(csr, x, alpha=2.0)
+
+# SpMM: C = A @ B
+C = trnsparse.spmm(csr, B)
+
+# Integral screening
+Q = trnsparse.schwarz_bounds(diagonal_integrals)
+mask = trnsparse.screen_quartets(Q, threshold=1e-10)
+stats = trnsparse.sparsity_stats(Q)
+```
+
+## Operations
+
+| Operation | Description |
+|-----------|-------------|
+| `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ᵀ |
+| `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 |
+
+## License
+
+Apache 2.0 — Copyright 2026 Scott Friedman
+
+
+## Disclaimer
+
+trnsci is an **independent open-source project**. It is not sponsored by, endorsed by, or affiliated with Amazon.com, Inc., Amazon Web Services, Inc., or Annapurna Labs Ltd.
+
+"AWS", "Amazon", "Trainium", "Inferentia", "NeuronCore", "Neuron SDK", and related identifiers are trademarks of their respective owners and are used here solely for descriptive and interoperability purposes. Use does not imply endorsement, partnership, or any other relationship.
+
+All work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on [trnsci.dev](https://trnsci.dev) are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.
+
+Feedback directed at the Neuron SDK or Trainium hardware is good-faith ecosystem commentary from independent users. It is not privileged information, is not pre-reviewed by AWS, and should not be read as authoritative about product roadmap, behavior, or quality.
+
+For official AWS guidance, see [aws-neuron documentation](https://awsdocs-neuron.readthedocs-hosted.com/) and the [AWS Trainium product page](https://aws.amazon.com/ai/machine-learning/trainium/).

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.1 → 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