Add relative position attention kernel and bindings

cactus-engine/src/model.cpp

@@ -4968,6 +4968,7 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& a
 
         if (!advanced) time_index += 1;
     }
+    audio_enc->graph->release_runtime_buffers();
 
     if (stream) {
         stream->initialized = true;

cactus-graph/cactus_graph.h

@@ -147,7 +147,8 @@ enum class OpType {
     QKV_TQ_FUSED,
     PROJECTION_PAIR_TQ_FUSED,
     CONV1D_CAUSAL_CHANNEL_FIRST,
-    LOGITS_TQ_SOFTCAP
+    LOGITS_TQ_SOFTCAP,
+    REL_POS_ATTENTION
 };
 
 struct PrecisionTraits {
@@ -652,6 +653,9 @@ class CactusGraph {
         bool additive_mask = false, size_t position_offset = 0, size_t window_size = 0,
         float logit_cap = 0.0f);
     size_t rel_pos_bias(size_t query, size_t relative_key, float scale, ComputeBackend backend = cactus_default_backend());
+    size_t rel_pos_attention(size_t query, size_t key, size_t value, size_t rel_query, size_t relative_key,
+                             size_t key_mask, float scale, size_t window_size,
+                             ComputeBackend backend = cactus_default_backend());
     size_t attention_int8_hybrid(
         size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset,
         const int8_t* cached_keys, const int8_t* cached_values,
@@ -1174,6 +1178,10 @@ CACTUS_FFI_EXPORT int cactus_graph_attention(
     cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, float scale, bool is_causal, size_t position_offset, size_t window_size, bool use_mask, cactus_node_t mask, bool additive_mask, cactus_node_t* out);
 CACTUS_FFI_EXPORT int cactus_graph_rel_pos_bias(
     cactus_graph_t graph, cactus_node_t query, cactus_node_t relative_key, float scale, cactus_node_t* out);
+CACTUS_FFI_EXPORT int cactus_graph_rel_pos_attention(
+    cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value,
+    cactus_node_t rel_query, cactus_node_t relative_key, bool use_key_mask, cactus_node_t key_mask,
+    float scale, size_t window_size, cactus_node_t* out);
 CACTUS_FFI_EXPORT int cactus_graph_attention_int8_hybrid(
     cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, float scale, size_t position_offset,
     const int8_t* cached_keys, const int8_t* cached_values, const float* k_scales, const float* v_scales,

cactus-graph/src/builder.cpp

@@ -612,6 +612,24 @@ size_t CactusGraph::rel_pos_bias(size_t query, size_t relative_key, float scale,
     return tag_backend(add_node(OpType::REL_POS_BIAS, {query, relative_key}, {B, H, T, T}, params), backend);
 }
 
+size_t CactusGraph::rel_pos_attention(size_t query, size_t key, size_t value, size_t rel_query, size_t relative_key,
+                                      size_t key_mask, float scale, size_t window_size, ComputeBackend backend) {
+    const auto& q = get_output_buffer(query);
+    const auto& r = get_output_buffer(relative_key);
+    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]) {
+        throw std::runtime_error("rel_pos_attention expects query [B, T, H, D % 16 == 0] and relative_key [1, R, H, D]");
+    }
+    const size_t reach = (window_size == 0 || window_size >= q.shape[1]) ? q.shape[1] - 1 : window_size;
+    if (r.shape[1] % 2 == 0 || (r.shape[1] - 1) / 2 < reach) {
+        throw std::runtime_error("rel_pos_attention relative_key must hold an odd number of offsets covering the window");
+    }
+    std::vector<size_t> inputs = {query, key, value, rel_query, relative_key};
+    if (key_mask != static_cast<size_t>(-1)) inputs.push_back(key_mask);
+    OpParams params{.scale = scale, .window_size = window_size, .output_precision = q.precision};
+    return tag_backend(add_node(OpType::REL_POS_ATTENTION, inputs, q.shape, params), backend);
+}
+
 size_t CactusGraph::attention_int8_hybrid(size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset,
                                           const int8_t* cached_keys, const int8_t* cached_values,
                                           const float* k_scales, const float* v_scales,

cactus-graph/src/execute.cpp

@@ -52,6 +52,7 @@ DECLARE_COMPUTE(compute_softmax_node);
 DECLARE_COMPUTE(compute_attention_node);
 DECLARE_COMPUTE(compute_attention_int8_hybrid_node);
 DECLARE_COMPUTE(compute_rel_pos_bias_node);
+DECLARE_COMPUTE(compute_rel_pos_attention_node);
 DECLARE_COMPUTE(compute_layernorm_node);
 DECLARE_COMPUTE(compute_conv1d_causal_node);
 DECLARE_COMPUTE(compute_conv1d_k3_node);
@@ -117,7 +118,7 @@ DECLARE_COMPUTE(compute_conv1d_causal_channel_first_node);
 extern void shrink_thread_local_buffers();
 #undef DECLARE_COMPUTE
 
-static constexpr int OP_TYPE_COUNT = static_cast<int>(OpType::LOGITS_TQ_SOFTCAP) + 1;
+static constexpr int OP_TYPE_COUNT = static_cast<int>(OpType::REL_POS_ATTENTION) + 1;
 static_assert(OP_TYPE_COUNT <= 256, "OpType dispatch table overflow");
 static ComputeFn dispatch_flat[OP_TYPE_COUNT] = {};
 
@@ -186,6 +187,7 @@ static bool init_dispatch() {
     dispatch_flat[static_cast<int>(OpType::ATTENTION)] = compute_attention_node;
     dispatch_flat[static_cast<int>(OpType::ATTENTION_INT8_HYBRID)] = compute_attention_int8_hybrid_node;
     dispatch_flat[static_cast<int>(OpType::REL_POS_BIAS)] = compute_rel_pos_bias_node;
+    dispatch_flat[static_cast<int>(OpType::REL_POS_ATTENTION)] = compute_rel_pos_attention_node;
     dispatch_flat[static_cast<int>(OpType::CONV1D_CAUSAL)] = compute_conv1d_causal_node;
     dispatch_flat[static_cast<int>(OpType::CONV1D_CAUSAL_CHANNEL_FIRST)] = compute_conv1d_causal_channel_first_node;
     dispatch_flat[static_cast<int>(OpType::CONV1D_K3)] = compute_conv1d_k3_node;
@@ -311,7 +313,8 @@ static const char* op_type_names[] = {
     "EXPAND",
     "MASKED_SELECT_PREFIX",
     "QKV_TQ_FUSED", "PROJECTION_PAIR_TQ_FUSED", "CONV1D_CAUSAL_CHANNEL_FIRST",
-    "LOGITS_TQ_SOFTCAP"
+    "LOGITS_TQ_SOFTCAP",
+    "REL_POS_ATTENTION"
 };
 
 static const char* get_op_name(OpType op) {
@@ -487,6 +490,8 @@ std::vector<size_t> infer_output_shape(const GraphNode& node, const nodes_vector
             }
             return out;
         }
+        case OpType::REL_POS_ATTENTION:
+            return in(0);
         case OpType::ATTENTION: case OpType::ATTENTION_CACHED: case OpType::ATTENTION_INT8_HYBRID: {
             std::vector<size_t> out = in(0);
             if (node.params.v_head_dim > 0) out.back() = node.params.v_head_dim;
@@ -969,6 +974,22 @@ static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const n
                 (uint32_t)B, (uint32_t)T, (uint32_t)H, (uint32_t)D, (uint32_t)R,
                 Rb == 1 ? 0 : 1, node.params.scale);
         }
+        case OpType::REL_POS_ATTENTION: {
+            if (node.input_ids.size() < 5) return false;
+            const auto& q = get_input(node, 0, nodes, map);
+            const auto& k = get_input(node, 1, nodes, map);
+            const auto& v = get_input(node, 2, nodes, map);
+            const auto& qv = get_input(node, 3, nodes, map);
+            const auto& r = get_input(node, 4, nodes, map);
+            const BufferDesc* mk = node.input_ids.size() > 5 ? &get_input(node, 5, nodes, map) : nullptr;
+            if (!fp16(q) || !fp16(k) || !fp16(v) || !fp16(qv) || !fp16(r) || !fp16(out)) return false;
+            if (mk && !fp16(*mk)) return false;
+            size_t B = q.shape[0], T = q.shape[1], H = q.shape[2], D = q.shape[3];
+            return cactus_metal_encode_rel_pos_attention_f16(out.get_data(), q.get_data(), k.get_data(),
+                v.get_data(), qv.get_data(), r.get_data(), mk ? mk->get_data() : nullptr,
+                (uint32_t)B, (uint32_t)T, (uint32_t)H, (uint32_t)D, (uint32_t)r.shape[1],
+                (uint32_t)node.params.window_size, node.params.scale);
+        }
         case OpType::BATCHNORM: {
             if (node.input_ids.size() != 5) return false;
             const auto& x = get_input(node, 0, nodes, map);

cactus-graph/src/graph_ffi.cpp

@@ -1203,6 +1203,22 @@ int cactus_graph_attention(cactus_graph_t graph, cactus_node_t query, cactus_nod
     }
 }
 
+int cactus_graph_rel_pos_attention(cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value,
+                                   cactus_node_t rel_query, cactus_node_t relative_key, bool use_key_mask,
+                                   cactus_node_t key_mask, float scale, size_t window_size, cactus_node_t* out) {
+    if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rel_pos_attention");
+    try {
+        *out = static_cast<cactus_node_t>(as_graph(graph)->graph.rel_pos_attention(
+            static_cast<size_t>(query), static_cast<size_t>(key), static_cast<size_t>(value),
+            static_cast<size_t>(rel_query), static_cast<size_t>(relative_key),
+            use_key_mask ? static_cast<size_t>(key_mask) : static_cast<size_t>(-1), scale, window_size));
+        return 0;
+    } catch (const std::exception& e) {
+        last_error_message = e.what();
+        return -1;
+    }
+}
+
 int cactus_graph_rel_pos_bias(cactus_graph_t graph, cactus_node_t query, cactus_node_t relative_key, float scale, cactus_node_t* out) {
     if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rel_pos_bias");
     try {

cactus-graph/src/io.cpp

@@ -278,10 +278,10 @@ namespace {
         GraphFile::NodeEntry node;
         node.index = read_u32(in);
         uint32_t op_type_val = read_u32(in);
-        if (op_type_val > static_cast<uint32_t>(OpType::LOGITS_TQ_SOFTCAP)) {
+        if (op_type_val > static_cast<uint32_t>(OpType::REL_POS_ATTENTION)) {
             throw std::runtime_error(
                 "Graph file corrupted: invalid op type " + std::to_string(op_type_val) +
-                " > " + std::to_string(static_cast<uint32_t>(OpType::LOGITS_TQ_SOFTCAP)));
+                " > " + std::to_string(static_cast<uint32_t>(OpType::REL_POS_ATTENTION)));
         }
         node.op_type = static_cast<OpType>(op_type_val);
         node.inputs = read_u32_vector(in);

cactus-graph/src/ops_nn.cpp

@@ -836,6 +836,19 @@ void compute_rel_pos_bias_node(GraphNode& node, const std::vector<std::unique_pt
         });
 }
 
+void compute_rel_pos_attention_node(GraphNode& node, const std::vector<std::unique_ptr<GraphNode>>& nodes,
+                                    const std::unordered_map<size_t, size_t>& node_index_map) {
+    const auto& q = get_input(node, 0, nodes, node_index_map);
+    const auto& r = get_input(node, 4, nodes, node_index_map);
+    const BufferDesc* mask = node.input_ids.size() == 6 ? &get_input(node, 5, nodes, node_index_map) : nullptr;
+    cactus_rel_pos_attention_f16(
+        q.data_as<__fp16>(), get_input(node, 1, nodes, node_index_map).data_as<__fp16>(),
+        get_input(node, 2, nodes, node_index_map).data_as<__fp16>(),
+        get_input(node, 3, nodes, node_index_map).data_as<__fp16>(), r.data_as<__fp16>(),
+        mask ? mask->data_as<__fp16>() : nullptr, node.output_buffer.data_as<__fp16>(),
+        q.shape[0], q.shape[1], q.shape[2], q.shape[3], r.shape[1], node.params.scale, node.params.window_size);
+}
+
 void compute_attention_node(GraphNode& node, const std::vector<std::unique_ptr<GraphNode>>& nodes, const std::unordered_map<size_t, size_t>& node_index_map) {
     if (node.input_ids.size() < 3 || node.input_ids.size() > 4) {
         throw std::runtime_error("Attention operation requires 3 or 4 inputs (query, key, value[, mask]), got " +

cactus-graph/src/param_io.cpp

@@ -204,6 +204,7 @@ const Schema& op_schema(OpType op_type) {
         {OpType::TOPK, {{ParamField::TopK, FieldPersistence::Persistent}}},
         {OpType::ATTENTION, {{ParamField::Scale, FieldPersistence::Persistent}, {ParamField::PositionOffset, FieldPersistence::Persistent}, {ParamField::WindowSize, FieldPersistence::Persistent}, {ParamField::IsCausal, FieldPersistence::Persistent}, {ParamField::AttentionMaskIsAdditive, FieldPersistence::Persistent}, {ParamField::LogitCap, FieldPersistence::Persistent}, {ParamField::Backend, FieldPersistence::Persistent}}},
         {OpType::REL_POS_BIAS, {{ParamField::Scale, FieldPersistence::Persistent}}},
+        {OpType::REL_POS_ATTENTION, {{ParamField::Scale, FieldPersistence::Persistent}, {ParamField::WindowSize, FieldPersistence::Persistent}}},
         {OpType::ATTENTION_INT8_HYBRID, {
             {ParamField::Scale, FieldPersistence::Persistent},
             {ParamField::PositionOffset, FieldPersistence::Persistent},

cactus-graph/tests/test_metal_parity.cpp

@@ -381,6 +381,25 @@ bool parity_flash_attention() {
     return true;
 }
 
+bool parity_rel_pos_attention() {
+    for (auto [T, window] : {std::pair<size_t, size_t>{150, 20}, {70, 0}, {37, 50}}) {
+        ParityCase c;
+        const size_t H = 2, D = 128;
+        const size_t R = window > 0 ? 2 * window + 1 : 2 * T - 1;
+        for (int i = 0; i < 4; ++i) c.add_input({1, T, H, D});
+        c.add_input({1, R, H, D});
+        std::vector<__fp16> mask(T, static_cast<__fp16>(0.0f));
+        for (size_t j = T - 4; j < T; ++j) mask[j] = static_cast<__fp16>(-10000.0f);
+        c.add_input_data({1, T}, mask);
+        c.tolerance = 2e-2f;
+        c.build = [D, window](CactusGraph& g, const std::vector<size_t>& in) {
+            return g.rel_pos_attention(in[0], in[1], in[2], in[3], in[4], in[5], 1.0f / std::sqrt((float)D), window);
+        };
+        if (!c.check()) return false;
+    }
+    return true;
+}
+
 bool parity_attention_causal() {
     ParityCase c;
     const size_t T = 37, H = 4, HKV = 2, D = 32;
@@ -562,6 +581,7 @@ int main() {
     runner.run_test("elemwise_chain (fused)", parity_elemwise_chain());
     runner.run_test("flash_attention (T%64!=0)", parity_flash_attention());
     runner.run_test("attention_causal", parity_attention_causal());
+    runner.run_test("rel_pos_attention", parity_rel_pos_attention());
     runner.run_test("equality_exact", parity_equality_exact());
     runner.run_test("sliding_window_cache (ring wrap)", parity_sliding_window_cache());
     runner.run_test("cache_growth", parity_cache_growth());

cactus-graph/tests/test_nn.cpp

@@ -320,6 +320,76 @@ bool test_attention() {
     return true;
 }
 
+static std::vector<float> run_rel_pos_graph(bool fused, size_t T, size_t window) {
+    const size_t B = 1, H = 2, D = 32;
+    const size_t full_rel = 2 * T - 1;
+    const size_t rel_rows = (fused && window > 0) ? 2 * window + 1 : full_rel;
+    std::mt19937 gen(7);
+    std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+    auto random = [&](size_t n) {
+        std::vector<__fp16> v(n);
+        for (auto& x : v) x = static_cast<__fp16>(dis(gen));
+        return v;
+    };
+    std::vector<__fp16> q = random(B * T * H * D), k = random(B * T * H * D), v = random(B * T * H * D);
+    std::vector<__fp16> qv = random(B * T * H * D), rel_full = random(full_rel * H * D);
+    std::vector<__fp16> key_mask(B * T, static_cast<__fp16>(0.0f));
+    for (size_t j = T - 5; j < T; ++j) key_mask[j] = static_cast<__fp16>(-10000.0f);
+    const float scale = 1.0f / std::sqrt(static_cast<float>(D));
+
+    CactusGraph g;
+    size_t iq = g.input({B, T, H, D}, Precision::FP16);
+    size_t ik = g.input({B, T, H, D}, Precision::FP16);
+    size_t iv = g.input({B, T, H, D}, Precision::FP16);
+    size_t iqv = g.input({B, T, H, D}, Precision::FP16);
+    size_t irel = g.input({1, rel_rows, H, D}, Precision::FP16);
+    size_t out;
+    std::vector<__fp16> band;
+    size_t iband = 0, imask;
+    if (fused) {
+        imask = g.input({B, T}, Precision::FP16);
+        out = g.rel_pos_attention(iq, ik, iv, iqv, irel, imask, scale, window);
+    } else {
+        imask = g.input({B, 1, 1, T}, Precision::FP16);
+        size_t mask = g.add(g.rel_pos_bias(iqv, irel, scale), imask);
+        if (window > 0) {
+            band.assign(T * T, static_cast<__fp16>(0.0f));
+            for (size_t t = 0; t < T; ++t)
+                for (size_t j = 0; j < T; ++j)
+                    if ((t > j ? t - j : j - t) > window) band[t * T + j] = static_cast<__fp16>(-INFINITY);
+            iband = g.input({1, 1, T, T}, Precision::FP16);
+            mask = g.add(mask, iband);
+        }
+        out = g.attention_masked(iq, ik, iv, mask, scale, false, cactus_default_backend(), true);
+    }
+    std::vector<__fp16> rel(rel_rows * H * D);
+    const size_t offset = (full_rel - rel_rows) / 2;
+    std::copy(rel_full.begin() + offset * H * D, rel_full.begin() + (offset + rel_rows) * H * D, rel.begin());
+    g.set_input(iq, q.data(), Precision::FP16);
+    g.set_input(ik, k.data(), Precision::FP16);
+    g.set_input(iv, v.data(), Precision::FP16);
+    g.set_input(iqv, qv.data(), Precision::FP16);
+    g.set_input(irel, rel.data(), Precision::FP16);
+    g.set_input(imask, key_mask.data(), Precision::FP16);
+    if (!band.empty()) g.set_input(iband, band.data(), Precision::FP16);
+    g.execute();
+    const __fp16* p = static_cast<const __fp16*>(g.get_output(out));
+    std::vector<float> result(B * T * H * D);
+    for (size_t i = 0; i < result.size(); ++i) result[i] = static_cast<float>(p[i]);
+    return result;
+}
+
+bool test_rel_pos_attention_matches_unfused() {
+    for (auto [T, window] : {std::pair<size_t, size_t>{70, 0}, {90, 25}, {40, 50}}) {
+        std::vector<float> expected = run_rel_pos_graph(false, T, window);
+        std::vector<float> actual = run_rel_pos_graph(true, T, window);
+        for (size_t i = 0; i < expected.size(); ++i) {
+            if (std::abs(expected[i] - actual[i]) > 5e-3f) return false;
+        }
+    }
+    return true;
+}
+
 bool test_reduction_operations() {
     TestUtils::FP16TestFixture fixture("Reduction Operations");
 
@@ -719,6 +789,7 @@ int main() {
     runner.run_test("Softmax", test_softmax());
     runner.run_test("Attention", test_attention());
     runner.run_test("Attention INT8 Hybrid", test_attention_int8_hybrid());
+    runner.run_test("Rel Pos Attention", test_rel_pos_attention_matches_unfused());
     runner.run_test("Reduction Operations", test_reduction_operations());
     runner.run_test("Mean Operations", test_mean_operations());
     runner.run_test("Variance Operations", test_variance_operations());

cactus-kernels/cactus_kernels.h

@@ -490,6 +490,22 @@ void cactus_attention_f16(
     size_t v_head_dim = 0,
     float logit_cap = 0.0f);
 
+void cactus_rel_pos_attention_f16(
+    const __fp16* query,
+    const __fp16* key,
+    const __fp16* value,
+    const __fp16* rel_query,
+    const __fp16* rel_key,
+    const __fp16* key_mask,
+    __fp16* output,
+    size_t batch_size,
+    size_t seq_len,
+    size_t num_heads,
+    size_t head_dim,
+    size_t rel_len,
+    float scale,
+    size_t window_size);
+
 void cactus_attention_hybrid_int8_fp16(
     const __fp16* queries,
     const int8_t* keys_cached,

cactus-kernels/metal_backend.h

@@ -221,6 +221,10 @@ bool cactus_metal_encode_conv_cache_append(void* out, const void* src, void* rin
 bool cactus_metal_encode_rel_pos_bias(void* y, const void* q, const void* r,
     uint32_t B, uint32_t T, uint32_t H, uint32_t D, uint32_t R, int r_batched, float scale);
 
+bool cactus_metal_encode_rel_pos_attention_f16(void* out, const void* q, const void* k, const void* v,
+    const void* qv, const void* r, const void* mask, uint32_t B, uint32_t T, uint32_t H, uint32_t D,
+    uint32_t R, uint32_t window, float scale);
+
 bool cactus_metal_encode_gemv_bias(void* out, const void* x, const void* w, const void* bias,
     uint32_t K, uint32_t N, int pretransposed);