Skip to content

trnsparse 0.3.0 — BSR is the Trainium-native sparse format

Choose a tag to compare

@scttfrdmn scttfrdmn released this 13 Apr 21:43
· 35 commits to main since this release

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 at block_size=128. Conversions to/from CSRMatrix and dense.
  • bsr_spmm(A_bsr, B) with NKI + PyTorch dispatch; NKI path wraps _BSRSpMMFunction (suite's second torch.autograd.Function-backed kernel after v0.2.0 CSR).
  • Hardware-validated on trn1.2xlarge — 7/7 @pytest.mark.neuron tests including torch.autograd.gradcheck.
  • Benchmarks populated with real trn1 numbers (docs/benchmarks.md).
  • sparse_add no longer materializes N×N dense intermediate (closes #8).
  • density_screen test 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

#8, #9, #10, #12, #18.

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.