Repository navigation
Add relative position attention kernel and bindings - #822
Merged
Merged
Conversation
- Implemented the `rel_pos_attn_f16` kernel in Metal for efficient relative position attention computation. - Updated the Metal backend to include the new kernel and its associated pipeline state. - Added a new function `cactus_metal_encode_rel_pos_attention_f16` to encode the relative position attention operation. - Created a corresponding stub function for the new operation in the Metal backend stub. - Developed unit tests to validate the functionality of the relative position attention against a dense reference implementation. - Enhanced Python bindings to expose the new relative position attention operation to the user. - Updated the transpiler to handle the new operation, including importers and lowering logic. - Ensured compatibility with existing attention mechanisms by integrating the new operation into the attention layer. Signed-off-by: jakmro <kubamroz124@gmail.com>
jakmro
force-pushed
the
parakeet-local-attention
branch
from
September 26, 2026 02:57
2198343 to
6738254
Compare
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Critical correctness, synchronization, importer, and input-validation issues remain unresolved.
Review effort: Lite
Findings: 5
Open (5)
What changed in this PR
Adds fused relative-position attention across Python transpilation, graph execution, CPU/Metal kernels, bindings, and Parakeet integration.
Changes:
- Adds schemas, importers, lowering, serialization, and Python bindings.
- Implements CPU and Metal attention kernels with graph dispatch.
- Adds runtime, graph, parity, and kernel tests.
| File | Summary |
|---|---|
python/cactus/transpile/tdt_runtime.py |
Adds runtime attention integration. |
python/cactus/transpile/ops.py |
Registers the operation schema. |
python/cactus/transpile/lower.py |
Lowers the new IR operation. |
python/cactus/transpile/importers.py |
Imports relative attention nodes. |
python/cactus/transpile/canonicalize/cleanup.py |
Preserves the operation during cleanup. |
python/cactus/transpile/aten_ops.py |
Maps the custom operation. |
python/cactus/convert/tests/test_parakeet_local_attention.py |
Tests runtime, capture, engine, and chunking behavior. |
python/cactus/bindings/cactus.py |
Exposes the graph binding. |
cactus-kernels/tests/test_attention.cpp |
Tests the CPU kernel. |
cactus-kernels/src/metal_backend.mm |
Adds Metal pipeline encoding. |
cactus-kernels/src/metal_backend_stub.cpp |
Adds non-Metal stub support. |
cactus-kernels/src/cactus_kernels.metal |
Implements the Metal kernel. |
cactus-kernels/src/attention.cpp |
Implements the CPU kernel. |
cactus-kernels/metal_backend.h |
Declares the Metal encoder. |
cactus-kernels/cactus_kernels.h |
Declares the CPU API. |
cactus-graph/tests/test_nn.cpp |
Tests fused and unfused graph equivalence. |
cactus-graph/tests/test_metal_parity.cpp |
Tests Metal parity. |
cactus-graph/src/param_io.cpp |
Adds serialized operation parameters. |
cactus-graph/src/ops_nn.cpp |
Adds graph node execution. |
cactus-graph/src/io.cpp |
Extends valid operation types. |
cactus-graph/src/graph_ffi.cpp |
Adds the FFI entry point. |
cactus-graph/src/execute.cpp |
Adds dispatch and backend encoding. |
cactus-graph/src/builder.cpp |
Adds graph construction and validation. |
cactus-graph/cactus_graph.h |
Defines the operation and public API. |
cactus-engine/src/model.cpp |
Releases encoder runtime buffers. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+619
to
+620
| if (q.shape.size() != 4 || q.shape[3] % 16 != 0 || r.shape.size() != 4 || r.shape[0] != 1 | ||
| || r.shape[2] != q.shape[2] || r.shape[3] != q.shape[3]) { |
| const size_t t = t0 + r; | ||
| const size_t j_lo = t > window ? t - window : 0; | ||
| const size_t j_hi = std::min(T, t + window + 1); | ||
| rel_pos_softmax_row(row, rel.data() + r * (len + BQ) + (BQ - 1 - r), mask_row, j_lo - ks, j_hi - ks, len); |
| const uint j = j0 + i / D4; | ||
| kv_s[i] = j < T ? v4[(size_t)j * stride4 + i % D4] : half4(0.0h); | ||
| } | ||
| simdgroup_barrier(mem_flags::mem_threadgroup); |
| "qwen2_moe_layer_gated.default": "qwen2_moe_layer_gated", | ||
| "cactus_transpile.gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated", | ||
| "gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated", | ||
| "rel_pos_attention.default": "rel_pos_attention", |
| torch.arange(k0, k1, device=query.device).view(1, -1) | ||
| - torch.arange(t0, t1, device=query.device).view(-1, 1) | ||
| ) | ||
| rel_index = (center + offset).clamp(0, rel_len - 1).expand(q.shape[0], q.shape[1], -1, -1) |
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
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.
rel_pos_attn_f16kernel in Metal for efficient relative position attention computation.cactus_metal_encode_rel_pos_attention_f16to encode the relative position attention operation.