trnsparse 0.4.1__tar.gz → 0.4.2__tar.gz

{trnsparse-0.4.1 → trnsparse-0.4.2}/CHANGELOG.md

@@ -5,6 +5,30 @@ 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.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.egg-info → trnsparse-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.4.1
+Version: 0.4.2
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0
@@ -102,3 +102,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.4.1 → trnsparse-0.4.2}/README.md

@@ -71,3 +71,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.4.2/docs/sparse_attention.md

@@ -0,0 +1,162 @@
+# Block-sparse attention
+
+## The connection
+
+Transformer attention computes `O = softmax(Q @ K.T / √d) @ V`. The expensive
+part at long sequence lengths is the `(seq_len, seq_len)` score matrix and the
+final `attn_weights @ V` matmul. Sparse attention restricts which positions can
+attend to each other, zeroing out large regions of the score matrix. What
+remains is a sparse × dense matmul: exactly `bsr_spmm`.
+
+The block size that matters is 128. Trainium's Tensor Engine is a 128-partition
+systolic array; `nc_matmul` consumes a 128×K×N tile in one call. A local-window
+attention mask with window granularity of 128 tokens has exactly `2w+1` nonzero
+128×128 blocks per block-row — each nonzero block maps one-to-one to one
+`nc_matmul` call, with no gather overhead. The tile size of the mask and the
+tile size of the hardware coincide.
+
+cuSPARSE's BSR is a specialization added on top of CSR. Here it's the other
+way: BSR is what the attention mask asks for, and the hardware is built around
+that shape.
+
+## Building patterns
+
+### Local window (Longformer-style)
+
+Each token block attends to its `window` nearest neighbors in block space:
+
+```python
+import trnsparse
+
+def local_window_mask(seq_len: int, block_size: int, window: int) -> trnsparse.BSRMatrix:
+    """BSRMatrix for sliding-window attention."""
+    import torch
+    n_blocks = seq_len // block_size
+    block_mask = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
+    for i in range(n_blocks):
+        lo, hi = max(0, i - window), min(n_blocks, i + window + 1)
+        block_mask[i, lo:hi] = True
+    mask_dense = block_mask.repeat_interleave(block_size, 0).repeat_interleave(block_size, 1)
+    return trnsparse.BSRMatrix.from_dense(mask_dense.float(), block_size=block_size)
+```
+
+At `seq_len=4096` and `window=2`: `5` nonzero blocks per row, `5/32 ≈ 15.6%`
+block density vs 100% for dense attention. Memory for the stored blocks:
+`32 rows × 5 blocks × 128 × 128 × 4 bytes ≈ 10 MB` vs 64 MB for the full
+float32 attention matrix.
+
+### Dilated sparse
+
+Every `stride`-th block column is attended to. Block `(i, j)` is nonzero
+iff `(i - j) % stride == 0`:
+
+```python
+def dilated_mask(seq_len: int, block_size: int, stride: int) -> trnsparse.BSRMatrix:
+    import torch
+    n_blocks = seq_len // block_size
+    block_mask = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
+    for i in range(n_blocks):
+        for j in range(n_blocks):
+            if (i - j) % stride == 0:
+                block_mask[i, j] = True
+    mask_dense = block_mask.repeat_interleave(block_size, 0).repeat_interleave(block_size, 1)
+    return trnsparse.BSRMatrix.from_dense(mask_dense.float(), block_size=block_size)
+```
+
+### Global tokens (BigBird-style)
+
+The first `n_global` token-blocks attend to and are attended to by every
+block. Remaining block-rows follow a local window:
+
+```python
+def global_token_mask(
+    seq_len: int, block_size: int, window: int, n_global: int
+) -> trnsparse.BSRMatrix:
+    import torch
+    n_blocks = seq_len // block_size
+    block_mask = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
+    block_mask[:n_global, :] = True   # global rows attend everywhere
+    block_mask[:, :n_global] = True   # global columns attended by all
+    for i in range(n_global, n_blocks):
+        lo, hi = max(0, i - window), min(n_blocks, i + window + 1)
+        block_mask[i, lo:hi] = True
+    mask_dense = block_mask.repeat_interleave(block_size, 0).repeat_interleave(block_size, 1)
+    return trnsparse.BSRMatrix.from_dense(mask_dense.float(), block_size=block_size)
+```
+
+## The computation
+
+Given a `mask_bsr` constructed above:
+
+```python
+import torch
+
+def block_sparse_attention(Q, K, V, mask_bsr):
+    """Q, K, V: (seq_len, head_dim) → (seq_len, head_dim)."""
+    scale = Q.shape[-1] ** -0.5
+    scores = (Q @ K.T) * scale
+
+    # Mask out unattended positions before softmax.
+    mask_dense = mask_bsr.to_dense().bool()
+    masked_scores = scores.masked_fill(~mask_dense, float("-inf"))
+    attn_weights = torch.softmax(masked_scores, dim=-1)
+
+    # Positions outside the mask go to zero after softmax; from_dense
+    # with threshold=0 will still store them unless they're exactly zero.
+    # Use threshold slightly above zero to keep only the attended blocks.
+    attn_bsr = trnsparse.BSRMatrix.from_dense(
+        attn_weights, block_size=mask_bsr.block_size, threshold=1e-9
+    )
+    return trnsparse.bsr_spmm(attn_bsr, V)
+```
+
+**What this materializes**: the full `(seq_len, seq_len)` score matrix. At
+`seq_len=4096` that's 64 MB of float32 even before the matmul. See
+[What's next](#whats-next) for the fused-tile path that avoids this.
+
+**Gradients**: `bsr_spmm` is differentiable (see `architecture.md`). The
+backward through `block_sparse_attention` works out-of-the-box with
+`torch.autograd`; the block-selection step (`from_dense`) is non-differentiable
+by construction and is treated as a constant by the autograd graph.
+
+## Block density arithmetic
+
+For `seq_len = S`, `block_size = b`, `window = w` local-window mask:
+
+```
+blocks per row   = 2w + 1  (clamped at edges)
+block density    = (2w+1) / (S/b)
+```
+
+| seq_len | window | n_blocks_per_row | block density |
+|--------:|-------:|-----------------:|:--------------|
+| 1024    | 2      | 5                | 39.1%         |
+| 2048    | 2      | 5                | 19.5%         |
+| 4096    | 2      | 5                | 9.8%          |
+| 4096    | 4      | 9                | 17.6%         |
+| 8192    | 2      | 5                | 4.9%          |
+
+The win is at long sequences. At `seq_len=1024`, a window-2 mask still stores
+40% of blocks — dispatch overhead dominates on Trainium at that size. At
+`seq_len=8192`, 5% density means 95% of the `nc_matmul` calls are skipped and
+the per-block compute becomes the bottleneck, which is where the Tensor Engine
+thrives.
+
+## What's next
+
+The current path materializes the full `(seq_len, seq_len)` score matrix to
+compute attention weights. A production path would:
+
+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.
+
+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.
+
+A runnable reference for the current (non-fused) path is in
+[`examples/block_sparse_attention.py`](https://github.com/trnsci/trnsparse/blob/main/examples/block_sparse_attention.py).

trnsparse-0.4.2/examples/block_sparse_attention.py

@@ -0,0 +1,284 @@
+"""Block-sparse attention via BSRMatrix and bsr_spmm.
+
+The final matmul in sparse attention — `attn_weights @ V` — is a sparse ×
+dense matrix multiplication. When the attention mask has 128×128 block
+structure (local window, dilated, global tokens), that matmul is exactly
+`bsr_spmm(attn_weights_bsr, V)`. No new kernel; BSRMatrix captures the
+sparsity pattern, and the Tensor Engine sees a sequence of dense 128×128
+tiles with zero gather overhead.
+
+Three mask patterns are demonstrated:
+
+  local  — sliding-window (Longformer-style): each token attends to the
+            nearest `window` blocks. Block (i, j) is stored iff |i-j| ≤ window.
+  dilated — every `stride`-th block column is stored. Covers dilated sparse
+            attention variants.
+  global — first `n_global` token blocks attend to and are attended to by
+            everyone; remaining rows follow a local window.
+
+Usage:
+    python examples/block_sparse_attention.py --demo
+    python examples/block_sparse_attention.py --seq-len 512 --pattern dilated
+    python examples/block_sparse_attention.py --seq-len 1024 --head-dim 64 --pattern global
+
+Note on materialization: this example computes the full score matrix
+`Q @ K.T` before masking. A production path would compute only the nonzero
+tiles of S = Q @ K.T, avoiding the O(seq_len²) intermediate. That fused
+tile-level score computation is the architectural follow-up; here the claim
+being demonstrated is that `bsr_spmm` handles the attn_weights × V step.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+# Allow running from repo root without install.
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+import trnsparse
+
+# ---------------------------------------------------------------------------
+# Mask pattern helpers
+# ---------------------------------------------------------------------------
+
+
+def _local_window_mask(seq_len: int, block_size: int, window: int) -> torch.Tensor:
+    """Boolean mask for local-window (Longformer-style) attention.
+
+    Block (i, j) in block-row i, block-column j is nonzero iff |i - j| ≤ window.
+
+    Returns:
+        bool tensor of shape (seq_len, seq_len).
+    """
+    assert seq_len % block_size == 0, "seq_len must be a multiple of block_size"
+    n_blocks = seq_len // block_size
+    block_mask = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
+    for i in range(n_blocks):
+        lo = max(0, i - window)
+        hi = min(n_blocks, i + window + 1)
+        block_mask[i, lo:hi] = True
+    # Expand each block indicator to (block_size, block_size)
+    return block_mask.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)
+
+
+def _dilated_mask(seq_len: int, block_size: int, stride: int) -> torch.Tensor:
+    """Boolean mask for dilated sparse attention.
+
+    Block (i, j) is nonzero iff (i - j) % stride == 0 (i.e., block column j
+    is a multiple of stride away from block row i, including j == i).
+
+    Returns:
+        bool tensor of shape (seq_len, seq_len).
+    """
+    assert seq_len % block_size == 0
+    n_blocks = seq_len // block_size
+    block_mask = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
+    for i in range(n_blocks):
+        for j in range(n_blocks):
+            if (i - j) % stride == 0:
+                block_mask[i, j] = True
+    return block_mask.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)
+
+
+def _global_token_mask(seq_len: int, block_size: int, window: int, n_global: int) -> torch.Tensor:
+    """Boolean mask with global tokens + local window (BigBird-style).
+
+    First `n_global` token-blocks attend to all blocks and are attended to
+    by all blocks. Remaining rows follow a local window of `window` blocks.
+
+    Returns:
+        bool tensor of shape (seq_len, seq_len).
+    """
+    assert seq_len % block_size == 0
+    n_blocks = seq_len // block_size
+    assert 0 <= n_global <= n_blocks
+    block_mask = torch.zeros(n_blocks, n_blocks, dtype=torch.bool)
+    # Global rows: attend everywhere
+    block_mask[:n_global, :] = True
+    # Global columns: attended to by everyone
+    block_mask[:, :n_global] = True
+    # Local window for non-global rows
+    for i in range(n_global, n_blocks):
+        lo = max(0, i - window)
+        hi = min(n_blocks, i + window + 1)
+        block_mask[i, lo:hi] = True
+    return block_mask.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)
+
+
+# ---------------------------------------------------------------------------
+# Attention computation
+# ---------------------------------------------------------------------------
+
+
+def block_sparse_attention(
+    Q: torch.Tensor,
+    K: torch.Tensor,
+    V: torch.Tensor,
+    mask_bsr: trnsparse.BSRMatrix,
+) -> torch.Tensor:
+    """Compute sparse attention output using bsr_spmm for the final matmul.
+
+    Args:
+        Q: (seq_len, head_dim) query tensor.
+        K: (seq_len, head_dim) key tensor.
+        V: (seq_len, head_dim) value tensor.
+        mask_bsr: BSRMatrix encoding which attention blocks are nonzero.
+            The blocks themselves are ignored — only the sparsity pattern
+            (block_row_ptrs, block_col_indices) is used. The actual weights
+            come from softmax(Q @ K.T).
+
+    Returns:
+        (seq_len, head_dim) attention output.
+
+    Note: this materializes the full (seq_len, seq_len) score matrix Q @ K.T
+    before masking. See module docstring on the production alternative.
+    """
+    seq_len, head_dim = Q.shape
+    scale = head_dim**-0.5
+
+    # Full score matrix — O(seq_len²). A production path computes only the
+    # nonzero tiles of S to avoid this intermediate.
+    scores = (Q @ K.T) * scale  # (seq_len, seq_len)
+
+    # Apply additive mask: positions outside the pattern → -inf before softmax.
+    mask_dense = mask_bsr.to_dense().bool()
+    masked_scores = scores.masked_fill(~mask_dense, float("-inf"))
+
+    # Row-wise softmax over all attended positions (positions outside mask → 0).
+    attn_weights = torch.softmax(masked_scores, dim=-1)  # (seq_len, seq_len)
+
+    # Convert the post-softmax weights to BSR (same pattern, real values now).
+    attn_bsr = trnsparse.BSRMatrix.from_dense(attn_weights, block_size=mask_bsr.block_size)
+
+    # Final matmul via bsr_spmm: (seq_len, seq_len) sparse × (seq_len, head_dim) dense.
+    return trnsparse.bsr_spmm(attn_bsr, V)  # (seq_len, head_dim)
+
+
+def _dense_reference(
+    Q: torch.Tensor,
+    K: torch.Tensor,
+    V: torch.Tensor,
+    mask: torch.Tensor,
+) -> torch.Tensor:
+    """Dense reference: full softmax(Q @ K.T) @ V with additive mask."""
+    scale = Q.shape[-1] ** -0.5
+    scores = (Q @ K.T) * scale
+    masked = scores.masked_fill(~mask, float("-inf"))
+    attn = torch.softmax(masked, dim=-1)
+    return attn @ V
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--demo", action="store_true", help="run a small demo (seq_len=512)")
+    parser.add_argument("--seq-len", type=int, default=512)
+    parser.add_argument("--head-dim", type=int, default=64)
+    parser.add_argument(
+        "--pattern",
+        choices=["local", "dilated", "global"],
+        default="local",
+        help="attention mask pattern",
+    )
+    parser.add_argument(
+        "--window", type=int, default=2, help="block window for local/global patterns"
+    )
+    parser.add_argument("--stride", type=int, default=2, help="block stride for dilated pattern")
+    parser.add_argument(
+        "--n-global", type=int, default=2, help="global token blocks for global pattern"
+    )
+    parser.add_argument("--block-size", type=int, default=128, help="BSR block size")
+    args = parser.parse_args()
+
+    if args.demo:
+        args.seq_len = 512
+        args.head_dim = 64
+        args.pattern = "local"
+        args.window = 1
+
+    seq_len = args.seq_len
+    head_dim = args.head_dim
+    block_size = args.block_size
+
+    if seq_len % block_size != 0:
+        print(f"Error: seq_len ({seq_len}) must be a multiple of block_size ({block_size}).")
+        sys.exit(1)
+
+    torch.manual_seed(42)
+    Q = torch.randn(seq_len, head_dim)
+    K = torch.randn(seq_len, head_dim)
+    V = torch.randn(seq_len, head_dim)
+
+    # Build the mask.
+    if args.pattern == "local":
+        mask = _local_window_mask(seq_len, block_size, args.window)
+        pattern_desc = f"local window (w={args.window})"
+    elif args.pattern == "dilated":
+        mask = _dilated_mask(seq_len, block_size, args.stride)
+        pattern_desc = f"dilated (stride={args.stride})"
+    else:
+        mask = _global_token_mask(seq_len, block_size, args.window, args.n_global)
+        pattern_desc = f"global tokens (n_global={args.n_global}, w={args.window})"
+
+    mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)
+    n_blocks_total = (seq_len // block_size) ** 2
+    block_density = mask_bsr.n_blocks / n_blocks_total
+
+    print("Block-sparse attention via bsr_spmm:")
+    print(f"  seq_len:      {seq_len}")
+    print(f"  head_dim:     {head_dim}")
+    print(f"  block_size:   {block_size}")
+    print(f"  pattern:      {pattern_desc}")
+    print(
+        f"  n_blocks:     {mask_bsr.n_blocks} / {n_blocks_total}  ({block_density:.1%} block density)"
+    )
+    print()
+
+    # Dense reference.
+    out_dense = _dense_reference(Q, K, V, mask)
+
+    # BSR path.
+    out_bsr = block_sparse_attention(Q, K, V, mask_bsr)
+
+    max_err = (out_bsr - out_dense).abs().max().item()
+    print(f"  Numerical agreement: max |Δ| = {max_err:.2e}")
+    assert max_err < 1e-3, f"BSR and dense outputs disagree: max_err={max_err}"
+
+    # Timing: attn_weights @ V, BSR vs dense.
+    # Pre-compute attn_weights for a fair comparison of just the matmul step.
+    scale = head_dim**-0.5
+    masked_scores = (Q @ K.T * scale).masked_fill(~mask, float("-inf"))
+    attn_weights = torch.softmax(masked_scores, dim=-1)
+    attn_bsr_timed = trnsparse.BSRMatrix.from_dense(attn_weights, block_size=block_size)
+
+    N_REPS = 20
+
+    t0 = time.perf_counter()
+    for _ in range(N_REPS):
+        _ = trnsparse.bsr_spmm(attn_bsr_timed, V)
+    t_bsr = (time.perf_counter() - t0) / N_REPS
+
+    t0 = time.perf_counter()
+    for _ in range(N_REPS):
+        _ = attn_weights @ V
+    t_dense = (time.perf_counter() - t0) / N_REPS
+
+    print(f"  Timing (attn_weights @ V, {N_REPS}-rep mean):")
+    print(f"    bsr_spmm:   {t_bsr * 1e6:8.1f} μs")
+    print(f"    dense @:    {t_dense * 1e6:8.1f} μs")
+    print()
+    print("  Note: on CPU the PyTorch fallback is used. bsr_spmm wins at")
+    print("  large seq_len on Trainium, where dispatch overhead amortizes")
+    print("  against the tile-level nc_matmul throughput advantage.")
+
+
+if __name__ == "__main__":
+    main()

{trnsparse-0.4.1 → trnsparse-0.4.2}/mkdocs.yml

@@ -24,6 +24,8 @@ nav:
   - Architecture: architecture.md
   - Migration from scipy.sparse: migration_scipy.md
   - Benchmarks: benchmarks.md
+  - Iterative Solvers: iterative_solvers.md
+  - Sparse Attention: sparse_attention.md
   - Deployment:
       - AWS Setup for Neuron CI: aws_setup.md
 

{trnsparse-0.4.1 → trnsparse-0.4.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trnsparse"
-version = "0.4.1"
+version = "0.4.2"
 description = "Sparse matrix operations for AWS Trainium via NKI"
 readme = "README.md"
 license = "Apache-2.0"

trnsparse-0.4.2/tests/test_attention.py

@@ -0,0 +1,93 @@
+"""Tests for block-sparse attention example and pattern helpers."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+import trnsparse
+
+EXAMPLES = Path(__file__).resolve().parent.parent / "examples"
+sys.path.insert(0, str(EXAMPLES))
+
+import block_sparse_attention as demo  # noqa: E402
+
+
+class TestMaskShapes:
+    """Pattern helper sanity checks — block count and coverage."""
+
+    def test_local_window_mask_shape(self):
+        seq_len, block_size, window = 512, 128, 2
+        mask = demo._local_window_mask(seq_len, block_size, window)
+        assert mask.shape == (seq_len, seq_len)
+        assert mask.dtype == torch.bool
+        # Each block-row should have at most 2*window+1 nonzero blocks.
+        n_blocks = seq_len // block_size
+        bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)
+        for i in range(n_blocks):
+            row_count = (bsr.block_row_ptrs[i + 1] - bsr.block_row_ptrs[i]).item()
+            assert row_count <= 2 * window + 1
+
+    def test_local_window_mask_symmetric(self):
+        mask = demo._local_window_mask(256, 128, 1)
+        assert (mask == mask.T).all(), "local window mask should be symmetric"
+
+    def test_dilated_mask_shape(self):
+        seq_len, block_size, stride = 512, 128, 2
+        mask = demo._dilated_mask(seq_len, block_size, stride)
+        assert mask.shape == (seq_len, seq_len)
+        n_blocks = seq_len // block_size
+        # For stride=2 and n_blocks=4: blocks 0,2 from row 0; blocks 1,3 from row 1; etc.
+        bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)
+        expected_per_row = (n_blocks + stride - 1) // stride
+        for i in range(n_blocks):
+            row_count = (bsr.block_row_ptrs[i + 1] - bsr.block_row_ptrs[i]).item()
+            assert row_count == expected_per_row
+
+    def test_global_token_mask_global_rows_full(self):
+        seq_len, block_size, n_global = 384, 128, 1
+        mask = demo._global_token_mask(seq_len, block_size, window=1, n_global=n_global)
+        # First n_global*block_size rows must be fully attended.
+        assert mask[: n_global * block_size, :].all()
+        # First n_global*block_size columns must be fully attended by all rows.
+        assert mask[:, : n_global * block_size].all()
+
+
+class TestAttentionParity:
+    """bsr_spmm attention matches dense reference."""
+
+    def _run_parity(self, seq_len, head_dim, block_size, mask):
+        torch.manual_seed(0)
+        Q = torch.randn(seq_len, head_dim)
+        K = torch.randn(seq_len, head_dim)
+        V = torch.randn(seq_len, head_dim)
+
+        mask_bsr = trnsparse.BSRMatrix.from_dense(mask.float(), block_size=block_size)
+        out_bsr = demo.block_sparse_attention(Q, K, V, mask_bsr)
+        out_dense = demo._dense_reference(Q, K, V, mask)
+
+        torch.testing.assert_close(out_bsr, out_dense, atol=1e-4, rtol=1e-4)
+        assert out_bsr.shape == (seq_len, head_dim)
+
+    def test_local_window_parity(self):
+        seq_len, head_dim, block_size = 256, 32, 128
+        mask = demo._local_window_mask(seq_len, block_size, window=1)
+        self._run_parity(seq_len, head_dim, block_size, mask)
+
+    def test_dilated_parity(self):
+        seq_len, head_dim, block_size = 256, 32, 128
+        mask = demo._dilated_mask(seq_len, block_size, stride=2)
+        self._run_parity(seq_len, head_dim, block_size, mask)
+
+    def test_global_token_parity(self):
+        seq_len, head_dim, block_size = 384, 32, 128
+        mask = demo._global_token_mask(seq_len, block_size, window=1, n_global=1)
+        self._run_parity(seq_len, head_dim, block_size, mask)
+
+    def test_full_attention_is_dense(self):
+        """A fully-nonzero BSR mask should match plain dense attention."""
+        seq_len, head_dim, block_size = 256, 32, 128
+        mask = torch.ones(seq_len, seq_len, dtype=torch.bool)
+        self._run_parity(seq_len, head_dim, block_size, mask)

{trnsparse-0.4.1 → trnsparse-0.4.2}/trnsparse/__init__.py

@@ -5,7 +5,7 @@ CSR/COO formats, SpMV, SpMM, and integral screening for
 sparse scientific computing. Part of the trnsci scientific computing suite.
 """
 
-__version__ = "0.4.1"
+__version__ = "0.4.2"
 
 from .formats import BSRMatrix, COOMatrix, CSRMatrix, eye_sparse, from_dense, from_scipy
 from .iterative import bsr_diagonal, cg_bsr, jacobi_preconditioner_bsr, power_iteration_bsr

{trnsparse-0.4.1 → trnsparse-0.4.2/trnsparse.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.4.1
+Version: 0.4.2
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0
@@ -102,3 +102,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.4.1 → trnsparse-0.4.2}/trnsparse.egg-info/SOURCES.txt

@@ -26,6 +26,8 @@ docs/installation.md
 docs/iterative_solvers.md
 docs/migration_scipy.md
 docs/quickstart.md
+docs/sparse_attention.md
+examples/block_sparse_attention.py
 examples/sparse_fock.py
 infra/terraform/.terraform.lock.hcl
 infra/terraform/README.md
@@ -35,6 +37,7 @@ scripts/run_benchmarks.sh
 scripts/run_neuron_tests.sh
 scripts/run_simulator_tests.sh
 tests/conftest.py
+tests/test_attention.py
 tests/test_bsr.py
 tests/test_examples.py
 tests/test_formats.py