trnsparse 0.4.0__tar.gz → 0.4.2__tar.gz

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

@@ -5,6 +5,61 @@ 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
+
+- **`examples/sparse_fock.py`** rewritten around v0.4.0's
+  `screened_spmm`. Three paths side-by-side on the same inputs:
+  (1) v0.1.x unfused `schwarz_bounds → screen → from_dense → spmm`;
+  (2) v0.4.0 fused `screened_spmm` (one call);
+  (3) full Fock build — the coulomb from path 2 contracted against MO
+  coefficients via `trnblas.gemm` for `F_MO = C.T @ J @ C` (falls
+  back to `torch.matmul` if trnblas isn't installed). On a 50-basis
+  synthetic system, the fused path is ~130× faster than the unfused
+  (dominated by eliminating the Python `from_dense` CSR construction).
+  Closes #6.
+- **`examples/pyscf_bridge.py`** (new) — optional PySCF-driven demo.
+  Builds H2O (or benzene, or H2), pulls real AO ERIs via
+  `mol.intor("int2e")`, feeds the `(μμ|μμ)` diagonal into
+  `schwarz_bounds` + `screened_spmm` against a mock density matrix.
+  Reports realistic sparsity at `threshold=1e-8`. Requires
+  `pip install pyscf`; tests skip cleanly if not available.
+  Closes #13.
+- **`tests/test_examples.py`** — 2 CPU smoke tests plus a
+  PySCF-gated test. Exercises the `sparse_fock` unfused + fused
+  paths end-to-end and asserts parity (`atol=1e-6`).
+
+### Notes
+
+No API changes, no kernel changes — pure integration demo release.
+Users already on v0.4.0 can stay there; upgrade to v0.4.1 only to
+pick up the new examples.
+
 ## [0.4.0] — 2026-04-14
 
 ### Added

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.4.0
+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.0 → 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.2/examples/sparse_fock.py

@@ -0,0 +1,164 @@
+"""Sparse Fock matrix build with integral screening.
+
+Three paths, side by side:
+
+1. **v0.1.x unfused** — `schwarz_bounds` → `screen_quartets` → mask-apply
+   → `from_dense` → `spmm`. What the library shipped before v0.4.0.
+   Four host passes over an (n, n) mask + a separate CSR build + a
+   separate SpMM dispatch.
+2. **v0.4.0 fused** — `screened_spmm(A, diag_integrals, B, threshold)`.
+   One NKI kernel on the NKI backend; explicit mask + matmul on CPU.
+   Same numeric result, fewer HBM round-trips, no mask tensor on HBM.
+3. **Full Fock build** — the coulomb J from path 2 contracted against
+   MO coefficients via trnblas: `F_MO = C.T @ J @ C`. Demonstrates the
+   suite composition — trnsparse hands off to trnblas's GEMM once the
+   sparse step is done.
+
+Schwarz bounds here are synthetic (Gaussian-distance decay) — realistic
+enough to show non-trivial sparsity. For real AO integrals from a
+molecule, see `examples/pyscf_bridge.py`.
+
+Usage:
+    python examples/sparse_fock.py --demo
+    python examples/sparse_fock.py --nbasis 200 --threshold 1e-8
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+
+import trnsparse
+
+
+def _synthetic_schwarz_system(n: int, seed: int = 42) -> torch.Tensor:
+    """Return synthetic diagonal integrals `(μμ|μμ)` for a 1-D molecular chain.
+
+    Chemistry convention: the Schwarz bound for `(μν|μν)` factors as
+    `Q[μ] * Q[ν]` where `Q[i] = sqrt((ii|ii))`. This demo generates
+    per-shell magnitudes that span several orders of magnitude so the
+    outer-product bound `Q[i] * Q[j]` produces non-trivial sparsity at
+    a realistic threshold.
+    """
+    torch.manual_seed(seed)
+    # Shells arranged in a 1-D chain; Gaussian-like decay in magnitude
+    # from a central "heavy" region so there's a tail of small Q values.
+    idx = torch.arange(n, dtype=torch.float32)
+    center = n / 2.0
+    diag_integrals = torch.exp(-((idx - center) ** 2) / (2.0 * (n / 6.0) ** 2)) + 0.01
+    return diag_integrals
+
+
+def _unfused_path(
+    integrals_dense: torch.Tensor, diag_integrals: torch.Tensor, P: torch.Tensor, threshold: float
+):
+    """Path 1: explicit Schwarz bound + mask + from_dense + spmm. v0.1.x flow."""
+    import math
+
+    t0 = time.perf_counter()
+    Q = trnsparse.schwarz_bounds(diag_integrals)  # (n,)
+    pair_bound = Q.unsqueeze(-1) * Q.unsqueeze(0)  # (n, n)
+    mask = pair_bound > math.sqrt(threshold)
+    integrals_masked = integrals_dense * mask.to(integrals_dense.dtype)
+    integrals_sparse = trnsparse.from_dense(integrals_masked)
+    J = trnsparse.spmm(integrals_sparse, P)
+    return J, time.perf_counter() - t0
+
+
+def _fused_path(
+    integrals_dense: torch.Tensor, diag_integrals: torch.Tensor, P: torch.Tensor, threshold: float
+):
+    """Path 2: v0.4.0 fused screened_spmm."""
+    t0 = time.perf_counter()
+    J = trnsparse.screened_spmm(integrals_dense, diag_integrals, P, threshold=threshold)
+    return J, time.perf_counter() - t0
+
+
+def _full_fock_build(J: torch.Tensor, C: torch.Tensor):
+    """Path 3: transform the coulomb J into the MO basis via trnblas.
+
+    F_MO = C.T @ J @ C  — two GEMMs. Falls back to torch.matmul if
+    trnblas isn't importable (it's an optional suite dep; pure-trnsparse
+    users don't need it).
+    """
+    t0 = time.perf_counter()
+    try:
+        import trnblas
+
+        Jt = trnblas.gemm(1.0, C, J, transA=True)  # C.T @ J
+        F_MO = trnblas.gemm(1.0, Jt, C)  # (C.T @ J) @ C
+        backend = "trnblas"
+    except ImportError:
+        F_MO = C.T @ J @ C
+        backend = "torch.matmul (trnblas not installed)"
+    return F_MO, time.perf_counter() - t0, backend
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--demo", action="store_true", help="run a small demo")
+    parser.add_argument("--nbasis", type=int, default=50)
+    parser.add_argument("--threshold", type=float, default=1e-4)
+    args = parser.parse_args()
+
+    if args.demo:
+        args.nbasis = 50
+
+    n = args.nbasis
+    print("Sparse Fock build:")
+    print(f"  Basis functions: {n}")
+    print(f"  Threshold:       {args.threshold:.0e}")
+
+    diag_integrals = _synthetic_schwarz_system(n)
+    Q = trnsparse.schwarz_bounds(diag_integrals)  # 1-D Schwarz bounds
+
+    # Unscreened ERI slice — random, scaled by outer-product Schwarz
+    # (chemistry-realistic: integrals tracking the bound).
+    torch.manual_seed(0)
+    integrals_dense = torch.randn(n, n) * (Q.unsqueeze(-1) * Q.unsqueeze(0)) * 0.01
+
+    # Density matrix — random SPD.
+    M = torch.randn(n, n) * 0.1
+    P = M @ M.T
+
+    # MO coefficients (for the trnblas transform in path 3) — orthonormal.
+    U, _ = torch.linalg.qr(torch.randn(n, n))
+    C = U
+
+    # --- Path 1: unfused ---
+    J_unfused, t_unfused = _unfused_path(integrals_dense, diag_integrals, P, args.threshold)
+
+    # --- Path 2: fused ---
+    J_fused, t_fused = _fused_path(integrals_dense, diag_integrals, P, args.threshold)
+
+    # --- Path 3: trnblas MO transform ---
+    F_MO, t_transform, backend = _full_fock_build(J_fused, C)
+
+    # --- Sparsity stats for context ---
+    # Build the pair-bound matrix for reporting stats at the matmul scale.
+    pair_bound = Q.unsqueeze(-1) * Q.unsqueeze(0)
+    stats = trnsparse.sparsity_stats(pair_bound, args.threshold**0.5)
+
+    print()
+    print("  Sparsity statistics:")
+    print(f"    Total shell pairs:       {stats['total_pairs']}")
+    print(f"    Significant pairs:       {stats['significant_pairs']}")
+    print(f"    Pair sparsity:           {stats['pair_sparsity']:.1%}")
+    print()
+    print("  Coulomb build timings:")
+    print(f"    Path 1 (unfused, 4-step):   {t_unfused * 1e3:8.3f} ms")
+    print(
+        f"    Path 2 (fused screened_spmm): {t_fused * 1e3:8.3f} ms  ({t_unfused / t_fused:.2f}x vs unfused)"
+    )
+    print()
+    print("  Full Fock build (trnsparse → trnblas):")
+    print(f"    MO transform (C.T @ J @ C): {t_transform * 1e3:8.3f} ms via {backend}")
+    print()
+    print(f"  Unfused/fused J agreement: max |ΔJ| = {(J_unfused - J_fused).abs().max().item():.2e}")
+    print(f"  F_MO shape: {tuple(F_MO.shape)}, mean |F_MO| = {F_MO.abs().mean().item():.3e}")
+
+
+if __name__ == "__main__":
+    main()

{trnsparse-0.4.0 → 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.0 → trnsparse-0.4.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trnsparse"
-version = "0.4.0"
+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.2/tests/test_examples.py

@@ -0,0 +1,70 @@
+"""Smoke tests for the examples directory.
+
+These cover the user-facing integration demos without executing them as
+subprocesses — the tests import the example modules and call their
+entry points directly so failures surface with full tracebacks.
+"""
+
+from __future__ import annotations
+
+import math
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+EXAMPLES = Path(__file__).resolve().parent.parent / "examples"
+sys.path.insert(0, str(EXAMPLES))
+
+
+class TestSparseFockDemo:
+    def test_unfused_path_runs(self):
+        """Exercise the v0.1.x multi-step screening path end-to-end."""
+        import sparse_fock as demo
+
+        diag = demo._synthetic_schwarz_system(50)
+        import trnsparse
+
+        Q = trnsparse.schwarz_bounds(diag)
+        A = torch.randn(50, 50) * (Q.unsqueeze(-1) * Q.unsqueeze(0)) * 0.01
+        P = torch.randn(50, 50)
+        J, t = demo._unfused_path(A, diag, P, threshold=1e-4)
+        assert J.shape == (50, 50)
+        assert torch.isfinite(J).all()
+        assert t > 0
+
+    def test_fused_vs_unfused_parity(self):
+        """The two paths should agree up to fp tolerance on the same inputs."""
+        import sparse_fock as demo
+
+        torch.manual_seed(7)
+        diag = demo._synthetic_schwarz_system(40)
+        import trnsparse
+
+        Q = trnsparse.schwarz_bounds(diag)
+        A = torch.randn(40, 40) * (Q.unsqueeze(-1) * Q.unsqueeze(0)) * 0.01
+        P = torch.randn(40, 40)
+        threshold = 1e-4
+
+        J_unfused, _ = demo._unfused_path(A, diag, P, threshold)
+        J_fused, _ = demo._fused_path(A, diag, P, threshold)
+
+        torch.testing.assert_close(J_unfused, J_fused, atol=1e-6, rtol=1e-6)
+
+
+class TestPyScfBridge:
+    """Skips cleanly without PySCF; asserts non-trivial screening on H2O."""
+
+    def test_h2o_sto3g_screening(self):
+        pytest.importorskip("pyscf")
+        import pyscf_bridge as demo
+
+        report = demo.run_demo("h2o", "sto-3g", threshold=1e-8)
+        assert report["nao"] > 0
+        # Tight tolerance would give no sparsity on a tiny basis; sto-3g
+        # H2O has only 7 AOs, so we just assert the API ran and produced
+        # a finite result — non-trivial sparsity needs larger bases.
+        assert math.isfinite(report["output_norm"])
+        assert 0.0 <= report["pair_sparsity"] <= 1.0
+        assert report["output_shape"] == (report["nao"], report["nao"])

{trnsparse-0.4.0 → 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.0"
+__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.0 → trnsparse-0.4.2/trnsparse.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.4.0
+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.0 → 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,7 +37,9 @@ 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
 tests/test_iterative.py
 tests/test_nki_bsr.py

trnsparse-0.4.0/examples/sparse_fock.py

@@ -1,80 +0,0 @@
-"""
-Sparse Fock matrix build with integral screening.
-
-Demonstrates how Schwarz screening reduces the Fock build from O(N⁴) to
-effectively O(N²) for large molecules. The sparsity pattern is stored
-as a CSR matrix, and SpMM handles the screened contraction.
-
-Usage:
-    python examples/sparse_fock.py --demo
-    python examples/sparse_fock.py --nbasis 200
-"""
-
-import argparse
-import time
-
-import torch
-
-import trnsparse
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--demo", action="store_true")
-    parser.add_argument("--nbasis", type=int, default=50)
-    parser.add_argument("--threshold", type=float, default=1e-10)
-    args = parser.parse_args()
-
-    if args.demo:
-        args.nbasis = 50
-
-    n = args.nbasis
-    print("Sparse Fock build:")
-    print(f"  Basis functions: {n}")
-    print(f"  Threshold:       {args.threshold:.0e}")
-
-    torch.manual_seed(42)
-
-    # Simulate Schwarz bounds (decay with distance for realistic sparsity)
-    positions = torch.rand(n, 3) * 10.0  # Random 3D positions
-    distances = torch.cdist(positions, positions)
-    Q = torch.exp(-0.5 * distances)  # Gaussian decay
-
-    # Screen
-    stats = trnsparse.sparsity_stats(Q, args.threshold)
-    print("\n  Sparsity statistics:")
-    print(f"    Total shell pairs:       {stats['total_pairs']}")
-    print(f"    Significant pairs:       {stats['significant_pairs']}")
-    print(f"    Pair sparsity:           {stats['pair_sparsity']:.1%}")
-    print(f"    Quartet sparsity (lower): {stats['quartet_sparsity_lower']:.1%}")
-
-    # Build sparse integral matrix (simulated)
-    mask = trnsparse.screen_quartets(Q, args.threshold)
-    integrals_dense = torch.randn(n, n) * Q * 0.01
-    integrals_dense[~mask] = 0.0
-    integrals_sparse = trnsparse.from_dense(integrals_dense)
-    print(f"    Integral matrix nnz:     {integrals_sparse.nnz} / {n * n}")
-
-    # Density matrix (random SPD for demo)
-    P = torch.randn(n, n) * 0.1
-    P = P @ P.T
-
-    # Sparse Fock build: J_μν = Σ_λσ P_λσ * (μν|λσ)
-    # Approximated here as SpMM: J ≈ sparse_integrals @ P
-    t0 = time.perf_counter()
-    J_sparse = trnsparse.spmm(integrals_sparse, P)
-    t_sparse = time.perf_counter() - t0
-
-    # Dense reference
-    t0 = time.perf_counter()
-    J_dense = integrals_dense @ P
-    t_dense = time.perf_counter() - t0
-
-    error = torch.linalg.norm(J_sparse - J_dense).item()
-    print(f"\n  Sparse SpMM:  {t_sparse:.4f}s")
-    print(f"  Dense matmul: {t_dense:.4f}s")
-    print(f"  Error:        {error:.2e}")
-
-
-if __name__ == "__main__":
-    main()