trnsparse 0.3.0 — BSR is the Trainium-native sparse format
v0.3.0 reframes the library around what Trainium uniquely enables.
Trainium's Tensor Engine is a 128×128 systolic array. The natural unit of sparse work on it is not an individual nonzero — it's a 128×128 block. v0.3.0 introduces BSRMatrix and bsr_spmm, where every stored block is already a Tensor-Engine tile and maps to one nc_matmul call with zero gather overhead.
Added
BSRMatrix— block-sparse row format atblock_size=128. Conversions to/fromCSRMatrixand dense.bsr_spmm(A_bsr, B)with NKI + PyTorch dispatch; NKI path wraps_BSRSpMMFunction(suite's secondtorch.autograd.Function-backed kernel after v0.2.0 CSR).- Hardware-validated on
trn1.2xlarge— 7/7@pytest.mark.neurontests includingtorch.autograd.gradcheck. - Benchmarks populated with real trn1 numbers (
docs/benchmarks.md). sparse_addno longer materializesN×Ndense intermediate (closes #8).density_screentest coverage (closes #10).
Architectural thesis
Documented in docs/architecture.md lede: CSR is the construction and interop format; BSR is the NKI compute format. For matrices with real block structure — Fock/ERI tensors after Schwarz screening, FEM stiffness, graph adjacencies, block-sparse attention masks — BSR is strictly preferred. For truly unstructured sparse, the torch.sparse_csr_tensor PyTorch fallback (v0.1.3) is already within 2× of scipy and NKI adds nothing.
Honest reading of the benchmarks
At v0.3.0 scales, NKI dispatch + compilation + HBM round-trips dominate the matmul work. BSR-NKI is ~15-25× slower than BSR-PyTorch at small sizes. The architectural wins live in follow-up issues:
- #19 — fused screen + matmul kernel (eliminates two HBM round-trips)
- #20 — on-chip iterative solvers (CG / power iteration with A SBUF-resident across iterations)
- #21 — block-sparse attention primitive (BSR is the building block for Longformer/BigBird-style sparse transformers)
Closed
Reframed
#15 (CSR row-bucketing) demoted to backlog. Under the architectural frame, the CSR path is served by the PyTorch fallback and BSR is the NKI-side story. Row-bucketing would only help if NKI 2.24 exposed an indirect-DMA primitive, which it doesn't.