Repository navigation
Add topology-derived sparse attention kernel - #22
ThomasRaoux merged 8 commits into
Conversation
|
Requesting review: this PR moves the earlier triton-lang/triton#10768 work here following the repo-fit feedback and now includes focused correctness/validation tests plus a checked-in benchmark runner. @ThomasRaoux, if there is a better reviewer for kernels examples, could you route this to the right maintainer? |
ThomasRaoux
left a comment
ThomasRaoux
left a comment
There was a problem hiding this comment.
looks good to me, just one comment
Thanks, fixed. I moved the optional pandas dependency handling into benchmark_utils.compare_benchmarks, where pandas is actually used, and restored benchmarking/init.py to the normal exports. Added a regression test for importing benchmark_utils without pandas installed. any other blockers? @ThomasRaoux how can I move forward? |
should we just make panda an explicit dependency? I thought it was already the case |
I see two reasonable designs:
i am specifically working with second design in plan But well design one is also valid anyone operating triton must have pandas' installed . |
|
@ThomasRaoux can we move for merge if no blockers or caveats remain please !! |
A port of the Triton kernel merged as triton-lang/kernels#22. The Python original runs on CUDA; this runs anywhere aether-core does, including no_std. The port keeps the original's decomposition, which is what makes it testable. The CSR block schedule is combinatorial, so it is checked by exact set equality against the lower-triangular CSR the Python builder emits. The kernel is numeric, so it is checked against dense masked attention to 1e-12. The block loop carries a running maximum and denominator per query row and rescales by exp(m_prev - m_new), so the working set is one score tile rather than a seq x seq matrix. Measured block reduction at 16 blocks with local_radius 1, sink 1, topk 2: 56 of 136 scheduled blocks, 58.8%. The Triton PR measured 56.6% at seq 1024 and 80.9% at seq 4096 on an RTX 4060; this asserts the direction at a size a unit test can run and does not restate wall-clock numbers measured on hardware this workspace cannot reach. Salience is the elder rule, so every non-zero score is asserted against this crate's persistence engine rather than a second implementation of H0. Exactly one block scores zero, which follows from an invariant every merge preserves: each component holds exactly one block that has never been written. One caveat the port surfaced. Per-block salience is not permutation-equivariant: when two components tie on size, which is absorbed depends on index order, so the same centroid scores differently depending on where it sits, and the zero-scoring block moves. The multiset is invariant, being the H0 barcode. The Triton original has the same tie-breaking. A fix needs a tie-break on centroid content rather than index.
…rators
Converts any model into a world model without touching its weights: learns a
linear operator on persistence-barcode signatures of the model's own
activations, rolls it forward ~880x cheaper than a forward pass, and carries a
sheaf-consistency gate that detects drift without invoking the model.
Transplanted from the Faraday/Hamilton coupling-operator construction
(barcode -> Hilbert embedding -> least-squares operator -> Banach fixed point).
Measured, at matched state dimension:
- S2-Rips entity corpus: reads the island partition at 0.855 vs 0.469 for a
same-dimension linear channel (0.487 majority baseline), and forecasts it
16 steps out above chance where the linear channel is below it
- Lorenz: ~30% better rollouts
- distilgpt2: no gain at any layer -- reported, not buried
The topological channel earns its dimensions when the state is a set of
interacting entities whose interaction threshold is a fixed physical distance.
Also:
- control.py: CEM-MPC with topological costs and gate-based refusal
- schedule.py: reproduces triton-lang/kernels#22 block salience
bit-identically and 13x faster (MST edges, not all pairs)
- H0 via MST reproduces the mujoco#3396 DSU island partition exactly
21 runnable checks, no test framework required.
…or the merged CSR kernel The family's only structural property a sigmoid forget gate lacks is an exact zero, and every bed so far ran on a dense kernel, where a closed gate cannot change the result, only the cost. The ruling's hypothesis moves the claim to cost: the continuous part of the gate is a forget-gate bias, and the exact-zero part is a learned segment boundary that triton-lang/kernels#22, a merged causal CSR block-schedule kernel, can consume with zero error instead of the 0.0009 its persistence schedule carries. The block counts were recomputed at S 4096 and block 64: 2,080 dense causal blocks, 288 under fixed 512-token segments and 160 under fixed 256-token ones, all as stated. At the gate's measured density of 0.0024 flags per token, random boundaries visit 478 blocks on average over 200 draws, which is more than the 398 the persistence schedule visits - the ruling's text said fewer, and the record corrects it. The hypothesis is recorded as open and unbound, with its kill row: a learned flag under an L0 target of one boundary per 512 tokens must beat fixed 512-token segments at equal visited blocks and stay within seed spread of dense FoX. Limits: no row has run; the kernel has no backward pass, so training would run dense and price the cost forward only. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ics to the docs The README now opens with a program and its output instead of an audit. A periodic signal is delay-embedded and a seal loop widens the radius until the cloud is one piece, giving Betti numbers [1, 1, 0]. A hoop is certified through a square with lk 1. A threshold is refused three times before it is proven at radius 0.25. Short sections follow on the idea, on what the language can answer, and on where it stands; the technical paper continues below them. The per-module derivations, the source repositories and the table of upstream pull requests were in the README. They now live only in the documentation site's Integrated Mathematics pages, and the README links to them. The scheduled-attention kernel is described as the Rust rebuild of the author's own triton-lang/kernels#22. tests/tour.rs runs examples/tour.aegis. It checks that the README quotes the file verbatim and that the quoted results hold. 428 passing and 80 ignored under `cargo test --workspace --exclude aether-kernel`; readme_claims, features_doc and doc_ratchet pass.
Summary
[seq, dim]and[batch, heads, seq, dim]inputs with the same CSR schedule shared across batch/head lanes.Closes #21.
This was originally attempted in triton-lang/triton#10768. Maintainer feedback there said the Triton core repo was probably not the right place and that a repo of Triton kernels would be a better fit, so this PR moves the contribution here instead of reopening the core PR.
Local validation
wsl.exe -d Ubuntu -- bash -lc "cd /mnt/c/Users/seal/Documents/GitHub/kernels && /home/seal/.cache/codex-triton-topology-venv/bin/python -m py_compile benchmarking/__init__.py benchmarking/topology_sparse_attention.py kernels/topology_sparse_attention.py test/test_topology_sparse_attention.py && /home/seal/.cache/codex-triton-topology-venv/bin/python -m pytest -s --tb=short test/test_topology_sparse_attention.py"Result: 17 passed.
git diff --check -- .Result: no whitespace errors. Git emitted local CRLF warnings only.
Benchmark smoke command:
wsl.exe -d Ubuntu -- bash -lc "cd /mnt/c/Users/seal/Documents/GitHub/kernels && /home/seal/.cache/codex-triton-topology-venv/bin/python -m benchmarking.topology_sparse_attention --seq 1024 --rounds 3"Result: command completed and printed the markdown benchmark table.
Local benchmark
To reproduce the full table:
Measured on NVIDIA GeForce RTX 4060 Laptop GPU, PyTorch 2.12.1+cu130, Triton 3.7.1,
dim=64,block_size=64, 50 timing rounds.The main Triton-side comparison is
Triton sparse vs dense CSR: both sides use this PR's scheduled-attention Triton kernel, but the dense CSR path visits every causal block. The SDPA number is a full-causal optimized baseline, not the same sparse mask; the 1024-token case is slower than SDPA, so this should not be read as a universal SDPA speedup claim.