Skip to content

trnsparse 0.4.2 — block-sparse attention primitive

Choose a tag to compare

View all tags
@scttfrdmn scttfrdmn released this 16 Apr 20:29
· 18 commits to main since this release
v0.4.2
e101f77

What's new

  • 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.
  • 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 (#25).
  • 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.
  • mkdocs.yml: add Iterative Solvers (was missing from nav) and Sparse Attention.

Notes

No API changes, no kernel changes. The claim in #21: bsr_spmm is the block-sparse attention primitive; BSRMatrix captures the mask.

The 128-token granularity at which sparse attention is natural is the Tensor Engine tile — Trainium was built for attention; that design transfers to any block-128-structured sparse workload.

Install

pip install trnsparse==0.4.2

Full changelog: CHANGELOG.md

Assets 2
Loading