trnsparse 0.8.0__tar.gz → 0.9.0__tar.gz

{trnsparse-0.8.0 → trnsparse-0.9.0}/CHANGELOG.md

@@ -7,6 +7,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.9.0] — 2026-10-30
+
+Validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0, native TorchNeuron; LNC=1 and LNC=2) — see trnsci/trnsci#35.
+
+Headline performance: static BSR SpMM and block-sparse attention split across both LNC=2 cores, with bf16 operands on fp32 PSUM; the fp32 BSR crossover against dense aten at n=8192, width 512 rises from ~27% to ~49% block density, and at seq 8192 attention forward plus backward beats unmasked dense SDPA for all three masks (#36).
+
+### Changed
+
+#### Native path: static kernels on both LNC=2 cores, bf16 operands (#30)
+
+- `bsr_spmm_static`, `attn_fwd_static` and `attn_bwd_static` split their
+  output tiles between the two physical cores of an LNC=2 logical core
+  (`program_id`), balanced by stored blocks; each output tile has one
+  writer. `TRNSPARSE_CORE_SPLIT=0` turns the split off.
+- bf16 inputs (all operands bf16) run bf16 matmuls with fp32 PSUM
+  accumulation; scores and softmax statistics stay fp32, and the attention
+  statistics returned by `trnsparse::bsr_attn` are fp32 for bf16 / fp16 Q.
+  Mixed fp32/bf16 operands run in fp32.
+- `benchmarks/bench_native_trn2.py`: `--dtype bf16` (matched-precision
+  dense aten baselines) and `--fast` (device columns only).
+
 ## [0.8.0] — 2026-10-30
 
 Native TorchNeuron dispatch (Trn2/Trn3) alongside the PyTorch/XLA path (Trn1, maintenance); validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0) — see trnsci/trnsci#35.

{trnsparse-0.8.0 → trnsparse-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.8.0
+Version: 0.9.0
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0

{trnsparse-0.8.0 → trnsparse-0.9.0}/benchmarks/bench_native_trn2.py

@@ -15,6 +15,17 @@ Sections:
     attn    Block-sparse attention (local / dilated / global masks), forward
             and forward+backward, vs dense SDPA on the device and on CPU.
 
+Options for the bsr and attn sections:
+    --dtype bf16  run the operands in bf16 (fp32 PSUM inside the kernels)
+                  and the dense aten baselines in bf16 too, so each side has
+                  matched precision. Errors are against fp64 on the
+                  bf16-rounded inputs.
+    --fast        only the device columns (trnsparse e2e, static kernel,
+                  dense aten); skips the gather path and the CPU baselines.
+    TRNSPARSE_CORE_SPLIT=0 in the environment keeps each static kernel on
+    one physical core's partition (the whole kernel on both cores, as
+    before the split), for A/B runs.
+
 Timing: each callable is run ``--warmup`` times (the first call compiles),
 then ``--reps`` times with ``torch.neuron.synchronize()`` after every call;
 the median is reported in milliseconds. Structure stays on the host, dense
@@ -106,21 +117,26 @@ def random_bsr(n, density, seed=0):
     )
 
 
+_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16}
+
+
 def section_bsr(args):
     densities = [float(x) for x in args.densities.split(",")]
     widths = [int(x) for x in args.widths.split(",")]
     n = args.size
+    dt = _DTYPES[args.dtype]
     for width in widths:
         g = torch.Generator().manual_seed(1)
-        Bc = torch.randn(n, width, generator=g)
-        Bd = Bc.to(DEV)
+        Bc = torch.randn(n, width, generator=g).to(dt).float()  # values exact in dt
+        Bd = Bc.to(DEV, dt)
         dense_dev_ms = None
         cpu_dense_ms = None
         for dens in densities:
             A = random_bsr(n, dens)
+            A.blocks = A.blocks.to(dt).float()
             Ac = A.to_dense()
             Ad_dev = trnsparse.BSRMatrix(
-                blocks=A.blocks.to(DEV),
+                blocks=A.blocks.to(DEV, dt),
                 block_col_indices=A.block_col_indices,
                 block_row_ptrs=A.block_row_ptrs,
                 shape=A.shape,
@@ -130,6 +146,8 @@ def section_bsr(args):
             K_max = max(counts)
             row = {
                 "section": "bsr",
+                "dtype": args.dtype,
+                "core_split": _core_split(),
                 "n": n,
                 "width": width,
                 "block_density": dens,
@@ -141,13 +159,28 @@ def section_bsr(args):
             out = trnsparse.bsr_spmm(Ad_dev, Bd)
             row["rel_err"] = rel_err(out, ref)
             row["trnsparse_ms"] = safe_bench(lambda: trnsparse.bsr_spmm(Ad_dev, Bd), reps=args.reps)
+            ops = _d._get_native_ops()
+            if args.fast:
+                rp, cl = _d._bsr_structure_lists(A)
+                bl = Ad_dev.blocks
+                row["static_kernel_ms"] = safe_bench(
+                    lambda: ops["bsr_spmm_static"](bl, Bd, rp, cl), reps=args.reps
+                )
+                Ad = Ac.to(DEV, dt)
+                if dense_dev_ms is None:
+                    dense_dev_ms = safe_bench(lambda: Ad @ Bd, reps=args.reps)
+                row["dense_dev_ms"] = dense_dev_ms
+                Ad = None
+                if row["trnsparse_ms"] and dense_dev_ms:
+                    row["speedup_vs_dense_dev"] = dense_dev_ms / row["trnsparse_ms"]
+                emit(row, args.out)
+                continue
 
             # Split: device-side prep (gather) vs the kernel launch.
             def prep():
                 return _d._bsr_pad_and_gather(_d._BlockView(Ad_dev, Ad_dev.blocks), Bd)
 
             bp, bg, *_ = prep()
-            ops = _d._get_native_ops()
             # Gather path (main before trnsparse#30 static kernels): device-side
             # pad + gather, then the K_max-padded kernel.
             row["gather_prep_ms"] = safe_bench(prep, reps=args.reps)
@@ -160,7 +193,7 @@ def section_bsr(args):
                     lambda: ops["bsr_spmm_static"](bl, Bd, rp, cl), reps=args.reps
                 )
 
-            Ad = Ac.to(DEV)
+            Ad = Ac.to(DEV, dt)
             if dense_dev_ms is None:
                 dense_dev_ms = safe_bench(lambda: Ad @ Bd, reps=args.reps)
             row["dense_dev_ms"] = dense_dev_ms
@@ -302,11 +335,12 @@ def section_attn(args):
     seq = args.size
     nb = seq // P
     F = torch.nn.functional
+    dt = _DTYPES[args.dtype]
     for hd in [int(x) for x in args.head_dims.split(",")]:
         g = torch.Generator().manual_seed(0)
-        Qc, Kc, Vc = (torch.randn(seq, hd, generator=g) for _ in range(3))
-        dOc = torch.randn(seq, hd, generator=g)
-        Qd, Kd, Vd, dOd = (t.to(DEV) for t in (Qc, Kc, Vc, dOc))
+        Qc, Kc, Vc = (torch.randn(seq, hd, generator=g).to(dt).float() for _ in range(3))
+        dOc = torch.randn(seq, hd, generator=g).to(dt).float()
+        Qd, Kd, Vd, dOd = (t.to(DEV, dt) for t in (Qc, Kc, Vc, dOc))
 
         def sdpa_fwd(q, k, v, m=None):
             return F.scaled_dot_product_attention(q[None], k[None], v[None], attn_mask=m)[0]
@@ -320,7 +354,9 @@ def section_attn(args):
             return q.grad
 
         dense_fb = safe_bench(dense_fwd_bwd, reps=args.reps)
-        cpu_fwd = safe_bench(lambda: sdpa_fwd(Qc, Kc, Vc), warmup=1, reps=args.cpu_reps)
+        cpu_fwd = None
+        if not args.fast:
+            cpu_fwd = safe_bench(lambda: sdpa_fwd(Qc, Kc, Vc), warmup=1, reps=args.cpu_reps)
         for kind in args.masks.split(","):
             bm = block_mask(kind, nb)
             mb = mask_bsr_from_block_mask(bm)
@@ -328,6 +364,8 @@ def section_attn(args):
             counts = bm.sum(1).tolist()
             row = {
                 "section": "attn",
+                "dtype": args.dtype,
+                "core_split": _core_split(),
                 "seq": seq,
                 "head_dim": hd,
                 "mask": kind,
@@ -365,14 +403,21 @@ def section_attn(args):
             row["sdpa_masked_dev_fwd_ms"] = safe_bench(
                 lambda: sdpa_fwd(Qd, Kd, Vd, em_d), reps=args.reps
             )
-            row["sdpa_cpu_fwd_ms"] = cpu_fwd
-            row["sdpa_masked_cpu_fwd_ms"] = safe_bench(
-                lambda: sdpa_fwd(Qc, Kc, Vc, elem_mask), warmup=1, reps=args.cpu_reps
-            )
+            if not args.fast:
+                row["sdpa_cpu_fwd_ms"] = cpu_fwd
+                row["sdpa_masked_cpu_fwd_ms"] = safe_bench(
+                    lambda: sdpa_fwd(Qc, Kc, Vc, elem_mask), warmup=1, reps=args.cpu_reps
+                )
             emit(row, args.out)
             em_d = None
 
 
+def _core_split():
+    from trnsparse.nki import kernels
+
+    return kernels._core_split_enabled()
+
+
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--section", required=True, choices=["bsr", "csr", "screen", "attn"])
@@ -383,6 +428,8 @@ def main():
     ap.add_argument("--head-dims", default="64,128")
     ap.add_argument("--masks", default="local,dilated,global")
     ap.add_argument("--no-bwd", action="store_true")
+    ap.add_argument("--dtype", default="fp32", choices=sorted(_DTYPES))
+    ap.add_argument("--fast", action="store_true")
     ap.add_argument("--reps", type=int, default=10)
     ap.add_argument("--cpu-reps", type=int, default=3)
     ap.add_argument("--out", default="")

{trnsparse-0.8.0 → trnsparse-0.9.0}/docs/architecture.md

@@ -97,6 +97,24 @@ cost is one compile per distinct structure (seconds), which a reused pattern
 amortizes. `TRNSPARSE_STATIC_STRUCTURE=0` restores the gather path. Measured
 crossovers against dense aten: `docs/benchmarks.md`.
 
+Under LNC=2 (the Trn2 default) NKI traces each kernel once per physical
+core, with `program_id` 0 and 1 as Python ints. The static kernels split
+their output tiles between the two traces, balanced by stored blocks:
+
+- SpMM: block rows and 512-column tiles.
+- Attention forward: block rows.
+- Attention backward: one pool of dQ block rows and dK/dV block columns.
+
+Each output tile has exactly one writer. `TRNSPARSE_CORE_SPLIT=0` turns the
+split off, for A/B measurements.
+
+The static kernels take fp32 or bf16 operands. When every operand is bf16,
+the matmuls run in bf16 and accumulate in fp32 PSUM. Scores, softmax
+statistics and score-space arithmetic stay fp32; only P and dS are rounded
+to bf16 for their matmuls. Outputs come back in bf16, and the attention
+statistics (`row_max`, `row_denom`) are always fp32. Mixed fp32/bf16
+operands run in fp32.
+
 ## Formats
 
 - **CSR**: row pointer + column indices + values. Preferred for construction,

{trnsparse-0.8.0 → trnsparse-0.9.0}/pyproject.toml

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

trnsparse-0.9.0/tests/test_core_split.py

@@ -0,0 +1,53 @@
+"""CPU tests for the static kernels' work split across physical cores (#30)."""
+
+from __future__ import annotations
+
+import pytest
+
+from trnsparse.nki import kernels as K
+
+
+@pytest.mark.parametrize("n_prog", [1, 2, 3])
+@pytest.mark.parametrize(
+    "costs",
+    [
+        [],
+        [1.0],
+        [3, 3, 3, 3],
+        [41, 3, 3, 3, 3, 3, 3, 3, 3, 3],  # global-token row first
+        [0.25, 5, 0.25, 7, 2, 2, 9],
+    ],
+)
+def test_assign_units_partitions(costs, n_prog):
+    owned = [K._assign_units(costs, n_prog, pid) for pid in range(n_prog)]
+    flat = sorted(u for o in owned for u in o)
+    assert flat == list(range(len(costs)))  # every unit exactly once: no shared writes
+    for o in owned:
+        assert o == sorted(o)
+
+
+def test_assign_units_balances_skew():
+    costs = [41, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]
+    loads = [sum(costs[u] for u in K._assign_units(costs, 2, pid)) for pid in range(2)]
+    assert max(loads) == 41  # the heavy row alone; the rest go to the other core
+    even = [4] * 10
+    loads = [sum(even[u] for u in K._assign_units(even, 2, pid)) for pid in range(2)]
+    assert loads == [20, 20]
+
+
+def test_assign_units_deterministic():
+    costs = [5, 1, 5, 1, 2, 2]
+    assert K._assign_units(costs, 2, 0) == K._assign_units(list(costs), 2, 0)
+
+
+def test_core_split_env(monkeypatch):
+    monkeypatch.setenv("TRNSPARSE_CORE_SPLIT", "0")
+    assert not K._core_split_enabled()
+    monkeypatch.setenv("TRNSPARSE_CORE_SPLIT", "1")
+    assert K._core_split_enabled()
+    monkeypatch.delenv("TRNSPARSE_CORE_SPLIT")
+    assert K._core_split_enabled()
+
+
+def test_row_block_counts():
+    assert K._row_block_counts((0, 2, 2, 5)) == [2, 0, 3]

{trnsparse-0.8.0 → trnsparse-0.9.0}/tests/test_native_dispatch.py

@@ -845,6 +845,59 @@ class TestNativePathOnCpu:
         for g_, w_ in zip(got, want, strict=True):
             torch.testing.assert_close(g_, w_, atol=1e-4, rtol=1e-4)
 
+    def test_static_kernels_bf16_operand_contract(self, cpu_native, monkeypatch):
+        """bf16 inputs reach the static kernels as bf16 operands; statistics stay fp32."""
+        seen = {}
+
+        def recorder(name, ref, out_dtype_of):
+            def run(*args):
+                seen[name] = [a.dtype for a in args if isinstance(a, torch.Tensor)]
+                f32 = [a.float() if isinstance(a, torch.Tensor) else a for a in args]
+                res = ref(*f32)
+                dt = out_dtype_of(args)
+                if isinstance(res, tuple):
+                    return (res[0].to(dt),) + tuple(
+                        r.to(dt if name == "bwd" else torch.float32) for r in res[1:]
+                    )
+                return res.to(dt)
+
+            return run
+
+        for attr, name, ref, odt in (
+            ("_bsr_spmm_static_kernel", "spmm", _ref_bsr_spmm_static, lambda a: a[1].dtype),
+            ("_attn_fwd_static_kernel", "fwd", _ref_attn_fwd_static, lambda a: a[2].dtype),
+            ("_attn_bwd_static_kernel", "bwd", _ref_attn_bwd_static, lambda a: a[4].dtype),
+        ):
+            fn = recorder(name, ref, odt)
+            monkeypatch.setattr(kernels, attr, fn, raising=False)
+            monkeypatch.setattr(dispatch, attr, fn, raising=False)
+
+        bf = torch.bfloat16
+        g = torch.Generator().manual_seed(21)
+        dense = torch.randn(256, 256, generator=g)
+        dense[:128, 128:] = 0
+        A = trnsparse.BSRMatrix.from_dense(dense, block_size=P)
+        A.blocks = A.blocks.to(bf)
+        C = trnsparse.bsr_spmm(A, torch.randn(256, 64, generator=g).to(bf))
+        assert C.dtype == bf and seen["spmm"] == [bf, bf]
+        # Mixed precision runs in fp32 and returns the promoted dtype.
+        C = trnsparse.bsr_spmm(A, torch.randn(256, 64, generator=g))
+        assert C.dtype == torch.float32 and seen["spmm"] == [torch.float32] * 2
+
+        seq_len, hd = 2 * P, 64
+        mask = trnsparse.BSRMatrix.from_dense(torch.ones(seq_len, seq_len), block_size=P)
+        Q, K, V, dO = (torch.randn(seq_len, hd, generator=g).to(bf) for _ in range(4))
+        out, rm, rd = dispatch.nki_bsr_attn_tiled(Q, K, V, mask, return_stats=True)
+        assert out.dtype == bf and seen["fwd"] == [bf] * 3
+        assert rm.dtype == torch.float32 and rd.dtype == torch.float32
+        o2, rm2, rd2 = torch.ops.trnsparse.bsr_attn(
+            Q, K, V, mask.block_col_indices, mask.block_row_ptrs, P
+        )
+        assert o2.dtype == bf and rm2.dtype == torch.float32 and rd2.dtype == torch.float32
+        dQ, dK, dV = dispatch.nki_bsr_attn_bwd(Q, K, V, dO, out, mask, rm, rd)
+        assert seen["bwd"] == [bf] * 7 + [torch.float32] * 6
+        assert {t.dtype for t in (dQ, dK, dV)} == {bf}
+
     def test_column_structure(self):
         assert dispatch._column_structure([0, 2, 2, 5], [0, 1, 0, 1, 2], 3) == (
             [0, 2, 4, 5],

{trnsparse-0.8.0 → trnsparse-0.9.0}/tests/test_nki_compile.py

@@ -183,3 +183,119 @@ def test_kernel_neff_compiles(case, target, tmp_path):
     )
     driver._compile_bir_to_neff(bir, opts, inputs)
     assert (tmp_path / "file.neff").stat().st_size > 0
+
+
+# ── Static kernels on both physical cores, fp32 and bf16 (#30) ───────────────
+#
+# Under LNC=2 the static kernels are traced once per physical core and split
+# their output tiles between the two traces. Verify both traces for every
+# kernel and input dtype, including structures with fewer units than cores
+# (one block row: the BSR kernel splits columns instead) and skewed rows.
+
+
+def _bf16(*shape):
+    import ml_dtypes
+
+    return _f32(*shape).astype(ml_dtypes.bfloat16)
+
+
+def _static_lnc2_cases():
+    cases = {}
+    for dtype, mk in (("f32", _f32), ("bf16", _bf16)):
+        cases[("_bsr_spmm_static_kernel", f"skew-n512-{dtype}")] = lambda mk=mk: (
+            mk(6, _P, _P),
+            mk(4 * _P, 512),
+            (0, 4, 4, 5, 6),
+            (0, 1, 2, 3, 0, 3),
+        )
+        cases[("_bsr_spmm_static_kernel", f"n1024-{dtype}")] = lambda mk=mk: (
+            mk(3, _P, _P),
+            mk(3 * _P, 1024),
+            (0, 2, 2, 3),
+            (0, 2, 2),
+        )
+        cases[("_bsr_spmm_static_kernel", f"onerow-n128-{dtype}")] = lambda mk=mk: (
+            mk(2, _P, _P),
+            mk(2 * _P, 128),
+            (0, 2),
+            (0, 1),
+        )
+        for hd in (64, 128):
+            cases[("_attn_fwd_static_kernel", f"hd{hd}-{dtype}")] = lambda hd=hd, mk=mk: (
+                mk(hd, 3 * _P),
+                mk(hd, 3 * _P),
+                mk(3 * _P, hd),
+                (0, 2, 2, 5),
+                (0, 1, 0, 1, 2),
+            )
+            cases[("_attn_bwd_static_kernel", f"hd{hd}-{dtype}")] = lambda hd=hd, mk=mk: (
+                mk(hd, 3 * _P),
+                mk(hd, 3 * _P),
+                mk(hd, 3 * _P),
+                mk(hd, 3 * _P),
+                mk(3 * _P, hd),
+                mk(3 * _P, hd),
+                mk(3 * _P, hd),
+                _f32(3 * _P, 1),
+                np.abs(_f32(3 * _P, 1)),
+                _f32(3 * _P, 1),
+                _f32(1, 3 * _P),
+                np.abs(_f32(1, 3 * _P)),
+                _f32(1, 3 * _P),
+                (0, 2, 2, 5),
+                (0, 1, 0, 1, 2),
+                (0, 2, 4, 5),
+                (0, 2, 0, 2, 2),
+                0.125,
+            )
+    return cases
+
+
+_STATIC_LNC2 = _static_lnc2_cases()
+
+
+@pytest.mark.parametrize("lnc", [1, 2])
+@pytest.mark.parametrize("case", sorted(_STATIC_LNC2), ids=lambda c: f"{c[0]}-{c[1]}")
+def test_static_kernel_lnc_mlir_verifies(case, lnc):
+    from trnsparse.nki import kernels
+
+    if case[1].endswith("bf16"):
+        pytest.importorskip("ml_dtypes")
+    kernel = getattr(kernels, case[0])
+    result = source_to_mlir(
+        kernel, *_STATIC_LNC2[case](), target="trn2", lnc=lnc, frontend=FrontendKind.TRACER
+    )
+    result.verify()
+
+
+@pytest.mark.skipif(shutil.which("neuronx-cc") is None, reason="neuronx-cc not on PATH")
+@pytest.mark.parametrize("case", sorted(_STATIC_LNC2), ids=lambda c: f"{c[0]}-{c[1]}")
+def test_static_kernel_lnc2_neff_compiles(case, tmp_path):
+    """Full CPU NEFF compile of both per-core traces (LNC=2, trn2)."""
+    try:
+        from nki.compiler import driver
+        from nki.compiler.frontend import TracerFrontend
+        from nki.compiler.ncc_driver import CompileOptions
+    except ImportError as e:
+        pytest.skip(f"nki compiler driver API not available: {e}")
+    if not hasattr(driver, "_compile_bir_to_neff"):
+        pytest.skip("nki.compiler.driver._compile_bir_to_neff not available in this nki")
+    if case[1].endswith("bf16"):
+        pytest.importorskip("ml_dtypes")
+
+    from trnsparse.nki import kernels
+
+    kernel = getattr(kernels, case[0])
+    fn = getattr(kernel, "func", kernel)
+    inputs = dict(zip(inspect.signature(fn).parameters, _STATIC_LNC2[case](), strict=True))
+    opts = CompileOptions(
+        target="trn2",
+        lnc=2,
+        artifacts_dir=str(tmp_path),
+        output_path=str(tmp_path / "file.neff"),
+    )
+    bir = driver.compile_to_bir(
+        fn, frontend=TracerFrontend(), inputs=inputs, compile_opts=opts, enable_cache=False
+    )
+    driver._compile_bir_to_neff(bir, opts, inputs)
+    assert (tmp_path / "file.neff").stat().st_size > 0

{trnsparse-0.8.0 → trnsparse-0.9.0}/tests/test_nki_native.py

@@ -234,3 +234,107 @@ class TestTorchCompile:
 
         out = f(A.values, A.col_indices, A.row_ptrs, B.to(DEV))
         torch.testing.assert_close(out.cpu(), dense @ B, atol=ATOL, rtol=RTOL)
+
+
+def _rel(x, ref):
+    x, ref = x.detach().cpu().double(), ref.detach().cpu().double()
+    return float((x - ref).norm() / ref.norm())
+
+
+def _skewed_bsr(n_rows, n_cols, seed):
+    """Block row 0 dense (a global-token row), the rest one or two blocks; a middle row empty."""
+    g = torch.Generator().manual_seed(seed)
+    bm = torch.zeros(n_rows, n_cols, dtype=torch.bool)
+    bm[0, :] = True
+    for i in range(1, n_rows):
+        bm[i, i % n_cols] = True
+        bm[i, (3 * i + 1) % n_cols] = True
+    if n_rows > 2:
+        bm[n_rows // 2, :] = False
+    dense = torch.randn(n_rows * P, n_cols * P, generator=g)
+    dense *= bm.repeat_interleave(P, 0).repeat_interleave(P, 1)
+    return dense, trnsparse.BSRMatrix.from_dense(dense, block_size=P)
+
+
+class TestTwoCoreAndBf16:
+    """Static kernels split across both physical cores (LNC=2), fp32 and bf16 (#30).
+
+    A core that skipped its share would leave whole output tiles unwritten,
+    so the fp32 cases use tight tolerances over every row and column tile.
+    """
+
+    @pytest.mark.parametrize(
+        "n_rows,n_cols,N",
+        [(6, 6, 128), (6, 4, 1024), (1, 3, 512), (1, 2, 128), (9, 5, 200)],
+    )
+    def test_bsr_static_fp32_every_tile(self, n_rows, n_cols, N):
+        dense, A = _skewed_bsr(n_rows, n_cols, seed=n_rows * 10 + N)
+        B = torch.randn(n_cols * P, N, generator=torch.Generator().manual_seed(1))
+        A.blocks = A.blocks.to(DEV)
+        C = trnsparse.bsr_spmm(A, B.to(DEV))
+        assert C.dtype == torch.float32
+        torch.testing.assert_close(C.cpu(), dense @ B, atol=ATOL, rtol=RTOL)
+
+    @pytest.mark.parametrize("n_rows,n_cols,N", [(6, 6, 128), (6, 4, 1024), (1, 3, 512)])
+    def test_bsr_static_bf16(self, n_rows, n_cols, N):
+        dense, A = _skewed_bsr(n_rows, n_cols, seed=n_rows * 10 + N)
+        B = torch.randn(n_cols * P, N, generator=torch.Generator().manual_seed(1))
+        A.blocks = A.blocks.to(DEV, torch.bfloat16)
+        Bh = B.to(DEV, torch.bfloat16)
+        C = trnsparse.bsr_spmm(A, Bh)
+        assert C.dtype == torch.bfloat16 and C.device.type == DEV
+        ref = dense.bfloat16().double() @ B.bfloat16().double()  # bf16 inputs, exact sum
+        assert _rel(C, ref) < 1e-2
+
+    def test_bsr_static_bf16_grad(self):
+        dense, A = _skewed_bsr(4, 4, seed=2)
+        B = torch.randn(4 * P, 64, generator=torch.Generator().manual_seed(2))
+        A.blocks = A.blocks.to(DEV, torch.bfloat16).requires_grad_(True)
+        Bd = B.to(DEV, torch.bfloat16).requires_grad_(True)
+        trnsparse.bsr_spmm(A, Bd).float().sum().backward()
+        ref = dense.bfloat16().double().T @ torch.ones(4 * P, 64, dtype=torch.float64)
+        assert Bd.grad.dtype == torch.bfloat16
+        assert _rel(Bd.grad, ref) < 1e-2
+
+    @pytest.mark.parametrize("dtype,tol", [(torch.float32, 1e-5), (torch.bfloat16, 3e-2)])
+    @pytest.mark.parametrize("hd", [64, 128])
+    def test_attention_static_two_core(self, dtype, tol, hd):
+        from trnsparse.ops import _block_sparse_attn_backward, _block_sparse_attn_pytorch
+
+        torch.manual_seed(7)
+        seq_len = 5 * P
+        # Global-token block row/column 0 plus a local band: skewed per-row cost.
+        nb = seq_len // P
+        bm = torch.zeros(nb, nb)
+        bm[0, :] = 1
+        bm[:, 0] = 1
+        for i in range(nb):
+            bm[i, max(0, i - 1) : i + 2] = 1
+        mask = trnsparse.BSRMatrix.from_dense(
+            bm.repeat_interleave(P, 0).repeat_interleave(P, 1), block_size=P
+        )
+        Q, K, V, dO = (torch.randn(seq_len, hd).to(dtype) for _ in range(4))
+        Qd, Kd, Vd = (t.to(DEV).requires_grad_(True) for t in (Q, K, V))
+        out = trnsparse.block_sparse_attention_tiled(Qd, Kd, Vd, mask)
+        assert out.dtype == dtype
+        q64, k64, v64, do64 = (t.double() for t in (Q, K, V, dO))
+        scale = hd**-0.5
+        ref, rm, rd = _block_sparse_attn_pytorch(q64, k64, v64, mask, scale, return_stats=True)
+        assert _rel(out, ref) < tol
+        out.backward(dO.to(DEV))
+        dQ, dK, dV = _block_sparse_attn_backward(
+            q64, k64, v64, ref, do64, mask, scale, row_max=rm, row_denom=rd
+        )
+        for got, exp in ((Qd.grad, dQ), (Kd.grad, dK), (Vd.grad, dV)):
+            assert got.dtype == dtype
+            assert _rel(got, exp) < tol
+
+    def test_attention_bf16_stats_are_fp32(self):
+        seq_len, hd = 2 * P, 64
+        mask = _block_causal_mask(seq_len)
+        Q, K, V = (torch.randn(seq_len, hd, dtype=torch.bfloat16).to(DEV) for _ in range(3))
+        out, rm, rd = torch.ops.trnsparse.bsr_attn(
+            Q, K, V, mask.block_col_indices, mask.block_row_ptrs, P
+        )
+        assert out.dtype == torch.bfloat16
+        assert rm.dtype == torch.float32 and rd.dtype == torch.float32

{trnsparse-0.8.0 → trnsparse-0.9.0}/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.8.0"
+__version__ = "0.9.0"
 
 from .formats import BSRMatrix, COOMatrix, CSRMatrix, eye_sparse, from_dense, from_scipy
 from .iterative import (

{trnsparse-0.8.0 → trnsparse-0.9.0}/trnsparse/nki/_native_attn.py

@@ -32,6 +32,15 @@ import torch
 from . import dispatch as _d
 
 
+def _stats_dtype(dtype: torch.dtype) -> torch.dtype:
+    """Softmax statistics dtype: fp32 for low-precision Q (bf16 / fp16), else Q's.
+
+    The backward recomputes P = exp(S - row_max) / row_denom; a bf16
+    row_max would scale every P of the row by up to exp(2^-9 · |row_max|).
+    """
+    return torch.float32 if dtype in (torch.bfloat16, torch.float16) else dtype
+
+
 def _pattern(block_col_indices, block_row_ptrs, seq_len, block_size):
     return _d._MaskPattern(block_col_indices, block_row_ptrs, (seq_len, seq_len), block_size)
 
@@ -47,17 +56,24 @@ def bsr_attn(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     mask = _pattern(block_col_indices, block_row_ptrs, Q.shape[0], block_size)
     out, row_max, row_denom = _d.nki_bsr_attn_tiled(Q, K, V, mask, return_stats=True)
-    # Fixed output contract (Q's dtype and device) to match the fake.
-    return tuple(t.to(device=Q.device, dtype=Q.dtype) for t in (out, row_max, row_denom))
+    # Fixed output contract (Q's device; out in Q's dtype, the statistics in
+    # `_stats_dtype`) to match the fake.
+    sdt = _stats_dtype(Q.dtype)
+    return (
+        out.to(device=Q.device, dtype=Q.dtype),
+        row_max.to(device=Q.device, dtype=sdt),
+        row_denom.to(device=Q.device, dtype=sdt),
+    )
 
 
 @bsr_attn.register_fake
 def _bsr_attn_fake(Q, K, V, block_col_indices, block_row_ptrs, block_size):
     M_tiles = Q.shape[0] // block_size
+    sdt = _stats_dtype(Q.dtype)
     return (
         torch.empty_like(Q),
-        Q.new_empty((M_tiles, block_size)),
-        Q.new_empty((M_tiles, block_size)),
+        Q.new_empty((M_tiles, block_size), dtype=sdt),
+        Q.new_empty((M_tiles, block_size), dtype=sdt),
     )