feat(sparse): genuinely sparse kernels — spmm_csr/coo + sddmm + bsmm on both devices - #180
Merged
Merged
Conversation
…on both devices
Closes the S-series `sparse` backend_kernel axis: spmm_csr / spmm_coo / sddmm /
bsmm now run on both x86 AVX-512 and ROCm gfx1151 with kernels that ITERATE the
nonzero structure (not densify-then-GEMM), each on-device-validated vs numpy.
Kernels:
- x86 `avx512_sparse_f32.cpp` — spmm_csr (per-row AXPY, length-N SIMD inner +
scalar tail over the CSR nonzeros), sddmm (sampled dense-dense: only mask≠0
entries do the length-K AVX-512 FMA dot; B passed transposed so both dot
operands are contiguous). Added to the elementwise .so.
- ROCm `generate-rocm-spmm-kernel` / `generate-rocm-sddmm-kernel` — one thread
per output element; spmm iterates indptr/indices (i32 memrefs) over the row's
nonzeros; sddmm guards the K-dot with scf.if on mask≠0 (genuinely skips masked
zeros). New ODS ops ROCM_Spmm/SddmmKernelOp; both lower gpu→ROCDL→hsaco.
Runtime:
- `_sparse_compute` dispatches spmm_csr/coo (COO→CSR fold on host) / sddmm / bsmm
over device kernels. bsmm = block-sparse → dense GEMM (x86 f32 microkernel;
ROCm bf16 WMMA matmul). Lanes `x86_sparse_compiled` / `rocm_sparse_compiled`.
Validation: test_{x86,rocm}_sparse_compiled.py — x86 8 / gfx1151 6 (incl. codegen
gate), all pass on real hardware. ruff/mypy clean.
Coverage: backend_kernel reference→fused for the 4 sparse ops on both devices;
dashboards regenerated, drift clean. The residual `partial` is the distributed
sharding axis (mesh-gated), unchanged.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 28, 2026
…ion + native fused-kernel inventory (#187) * docs(roadmap): S-series gap-closure plan + native fused-kernel inventory Plan for driving the remaining S-series primitives to real on-device execution on the two devices we can prove on today — x86 AVX-512 + ROCm gfx1151 (RDNA 3.5 WMMA). Corrects the "144 reference" framing to the per-device truth (the universal backend_kernel gate reads 0/474 by design): ~127 ops execute on both devices, ~171 on neither, ~10 ROCm-only (WMMA attention lanes needing an x86 partner). Triages the ~171 both-device gap into: - Tier 0 host/structural (layout_transform/indexing) → classify not_applicable, no kernel; - Tier 1 mesh-gated transport (collective/moe_transport) → stays gated; - Tier 2 easy elementwise/predicate → extend existing lanes; - Tier 3 composable (optimizer steps incl. fused AdamW, norm, complex/GA, losses); - Tier 4 new kernel families (attention, conv, sort/top_k, MLA loop_nest, RNG). Records the NATIVE (AVX-512 + RDNA WMMA) planned fused-kernel inventory — flash_attn / MLA / NSA / lightning / kimi / swiglu / fused AdamW / conv / sort / Philox RNG — the executable-now companion to the hardware-gated WGMMA/MFMA inventories (nvidia_cuda13 / rocm_mfma), aligned to the canonical family names guarded by test_kernel_inventory_and_lit_fixtures.py. Sequences phased PRs A→G using the per-family cadence proven across #180–#186. Linked from ROADMAP_AUDIT. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(roadmap): add §6 Tier-0 compiler foundation plan Reframe Tier 0 from "just mark not_applicable" to the compiler-foundation workstream it actually is. Verified state: 28 of 33 structural ops have NO Graph IR op (only transpose/reshape/cast/select/masked_fill do) — they live only in the Python frontend + numpy runtime, so they can't enter a compiled pipeline, fuse, or lower to device movement. Autodiff VJPs are largely complete (incl. gather/ scatter/masked_fill — corrected an earlier mis-read); JVP deferred. §6 specifies what the compiler / runtime / stdlib each need: - 6.A Graph IR: ODS ops (reuse tensor/linalg dialect where it models the op; native tessera.layout/index ops only for roll/flip/repeat/tile/pack + ops that must carry a Tessera layout/distribution attr) + shape verify + canonicalizers + frontend emission. - 6.B Passes: view-vs-copy analysis, layout propagation/canonicalize, bufferization to strided memref, movement lowering — makes 0-view ops compile to stride rewrites (the real meaning of not_applicable). - 6.C Runtime: a backend-agnostic view descriptor ABI (generalize the Apple-only reshape_view), a strided-copy movement lane (0-move), and gather/scatter(+atomic) kernels (0-reduce) on x86 + ROCm via the proven cadence. - 6.D Contracts: formalize the UNSPECIFIED vmap/batching rule + sharding rule + flip transpose_rule (VJPs already exist). - 6.E Stdlib: pack/unpack, gather lowering-class fix, JVPs. Splits the kernel axis into 0-view (not_applicable), 0-move (movement lane), 0-reduce (real kernel) — so the disposition is honest per sub-class, not a blanket not_applicable. Sequenced F1→F5 (Graph IR → view passes → runtime ABI → gather/ scatter kernels → contracts). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(roadmap): consolidate into one S-series closure plan — deep per-tier dives + Tier 5 Single closure plan: per-tier deep dives (§7) mirroring the §6 Tier-0 foundation lens (compiler / runtime / stdlib + kernel + deps + validation + disposition), plus a consolidated P0–P15 phase order (§8) that folds the compute phases together with the Tier-0 foundation steps F1–F5. - Tier 1 (transport): NOT greenfield — collective/moe ops already have ODS + mock-collective execution + sharding contract; the ONLY gap is real multi-GPU HW. Stays partial (single-rank ref + mock-mesh), ungates at Phase H. No kernel. - Tier 2 (elementwise/predicate): per-lane breakdown (unary+new digamma/lgamma poly cores, binary atan2/floor_div/mod, predicate isnan/isinf/isfinite float→i8, popcount VPOPCNTDQ, reduce/segment_reduce, ntk_rope on the rope lane). - Tier 3: optimizer steps are STATEFUL (need an m/v + step-counter runtime ABI, in-place multi-output) — fused AdamW called out; norm composes on reduce; complex reuses the FFT interleaved-f32 substrate; EBM/diffusion losses split by RNG dependency. - Tier 4: attention is mostly an x86-PARTNER problem (the 10 ROCm-only ops are attention WMMA lanes); state-recurrence variants reuse the selective_ssm scan substrate; NSA depends on top_k; RNG (Philox) unblocks the sampling losses. - Tier 5 (NEW): geometric_algebra / Clifford (18 ops) — distinct kernel pattern (table-driven bilinear contraction over the multivector Cayley table, not flat elementwise nor dense GEMM). CPU reference exists; ROCm is the headline gap. §8 P0–P15 front-loads honesty/structural-IR/elementwise/AdamW before the marquee attention+MLA work (which sits behind its substrate); transport (P15) is the only phase blocked on hardware this repo lacks. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(roadmap): add required audit-doc frontmatter to the closure plan last_updated / audit_role: plan / plan_state: open — satisfies the test_audit_docs + test_docs_freshness frontmatter contract. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: gstoner <angstroms01@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes the S-series
sparsebackend_kernel axis. spmm_csr / spmm_coo / sddmm / bsmm now run on both x86 AVX-512 and ROCm gfx1151 with kernels that iterate the nonzero structure (not densify-then-GEMM) — each on-device-validated vs numpy. First device lane for this family (was reference-only).Kernels
x86 (
avx512_sparse_f32.cpp, added to the elementwise .so):spmm_csr— per-row AXPY over the CSR nonzeros (length-N SIMD inner + scalar tail).sddmm— sampled dense-dense: onlymask≠0entries do the length-K AVX-512 FMA dot; B passed transposed so both dot operands are contiguous.ROCm (
generate-rocm-spmm-kernel/generate-rocm-sddmm-kernel, new ODS opsROCM_Spmm/SddmmKernelOp, both lower gpu→ROCDL→hsaco):scf.ifonmask≠0(genuinely skips masked-zero entries — a real GPU sampling win).Runtime
_sparse_computedispatches spmm_csr/coo (COO→CSR fold on host) / sddmm / bsmm over the device kernels. bsmm = block-sparse → dense GEMM (x86 f32 microkernel; ROCm bf16 WMMA matmul). Lanesx86_sparse_compiled/rocm_sparse_compiled.Validation
test_x86_sparse_compiled.py(8) +test_rocm_sparse_compiled.py(6, incl. codegen gate) — all vs numpy, all pass on real hardware (gfx1151 + AVX-512). ruff/mypy clean.Coverage (honest)
backend_kernel flips reference→fused for the 4 sparse ops on both devices. Dashboards regenerated, drift clean. Residual
partialis the distributed sharding axis (mesh-gated), unchanged.First of the remaining reference-only categories — next: state_space, state_update, linalg, moe-compute.
🤖 Generated with Claude Code