diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 2fc60d1e4..d4c5ef79b 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -4968,6 +4968,7 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a if (!advanced) time_index += 1; } + audio_enc->graph->release_runtime_buffers(); if (stream) { stream->initialized = true; diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 28c7c978a..c5908b1c1 100644 --- a/cactus-graph/cactus_graph.h +++ b/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, diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index 0dd3796e4..9fc1241d6 100644 --- a/cactus-graph/src/builder.cpp +++ b/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 inputs = {query, key, value, rel_query, relative_key}; + if (key_mask != static_cast(-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, diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp index 56830324e..fcbdb628f 100644 --- a/cactus-graph/src/execute.cpp +++ b/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(OpType::LOGITS_TQ_SOFTCAP) + 1; +static constexpr int OP_TYPE_COUNT = static_cast(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(OpType::ATTENTION)] = compute_attention_node; dispatch_flat[static_cast(OpType::ATTENTION_INT8_HYBRID)] = compute_attention_int8_hybrid_node; dispatch_flat[static_cast(OpType::REL_POS_BIAS)] = compute_rel_pos_bias_node; + dispatch_flat[static_cast(OpType::REL_POS_ATTENTION)] = compute_rel_pos_attention_node; dispatch_flat[static_cast(OpType::CONV1D_CAUSAL)] = compute_conv1d_causal_node; dispatch_flat[static_cast(OpType::CONV1D_CAUSAL_CHANNEL_FIRST)] = compute_conv1d_causal_channel_first_node; dispatch_flat[static_cast(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 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 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); diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 223ce8c84..9ad8cea41 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/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(as_graph(graph)->graph.rel_pos_attention( + static_cast(query), static_cast(key), static_cast(value), + static_cast(rel_query), static_cast(relative_key), + use_key_mask ? static_cast(key_mask) : static_cast(-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 { diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index ad4fbf5f0..184a2bd0f 100644 --- a/cactus-graph/src/io.cpp +++ b/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(OpType::LOGITS_TQ_SOFTCAP)) { + if (op_type_val > static_cast(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(OpType::LOGITS_TQ_SOFTCAP))); + " > " + std::to_string(static_cast(OpType::REL_POS_ATTENTION))); } node.op_type = static_cast(op_type_val); node.inputs = read_u32_vector(in); diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index e5eed8d89..a7de35d64 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -836,6 +836,19 @@ void compute_rel_pos_bias_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& 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>& nodes, const std::unordered_map& 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 " + diff --git a/cactus-graph/src/param_io.cpp b/cactus-graph/src/param_io.cpp index 2c58b6020..02fc97220 100644 --- a/cactus-graph/src/param_io.cpp +++ b/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}, diff --git a/cactus-graph/tests/test_metal_parity.cpp b/cactus-graph/tests/test_metal_parity.cpp index ae8f2d548..8391d7135 100644 --- a/cactus-graph/tests/test_metal_parity.cpp +++ b/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{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& 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()); diff --git a/cactus-graph/tests/test_nn.cpp b/cactus-graph/tests/test_nn.cpp index 890631931..37d3bfacf 100644 --- a/cactus-graph/tests/test_nn.cpp +++ b/cactus-graph/tests/test_nn.cpp @@ -320,6 +320,76 @@ bool test_attention() { return true; } +static std::vector 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 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(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(g.get_output(out)); + std::vector result(B * T * H * D); + for (size_t i = 0; i < result.size(); ++i) result[i] = static_cast(p[i]); + return result; +} + +bool test_rel_pos_attention_matches_unfused() { + for (auto [T, window] : {std::pair{70, 0}, {90, 25}, {40, 50}}) { + std::vector expected = run_rel_pos_graph(false, T, window); + std::vector 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()); diff --git a/cactus-kernels/cactus_kernels.h b/cactus-kernels/cactus_kernels.h index eddcc5bd9..f626aab1c 100644 --- a/cactus-kernels/cactus_kernels.h +++ b/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, diff --git a/cactus-kernels/metal_backend.h b/cactus-kernels/metal_backend.h index 9505aad48..16e2e66c3 100644 --- a/cactus-kernels/metal_backend.h +++ b/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); diff --git a/cactus-kernels/src/attention.cpp b/cactus-kernels/src/attention.cpp index 131111a83..38ac10a21 100644 --- a/cactus-kernels/src/attention.cpp +++ b/cactus-kernels/src/attention.cpp @@ -728,3 +728,235 @@ void cactus_attention_f16( } }); } + +namespace { + +constexpr size_t REL_POS_BLOCK_Q = 32; +constexpr size_t REL_POS_CHUNK_BLOCKS = 8; +constexpr CactusThreading::ParallelConfig REL_POS_ATTENTION_PARALLEL{2, 1}; +#ifdef __APPLE__ +constexpr bool REL_POS_PANELS = false; +#else +constexpr bool REL_POS_PANELS = true; +#endif + +inline float32x4_t rel_pos_exp_f32(float32x4_t x) { + x = vmaxq_f32(x, vdupq_n_f32(-87.0f)); + const float32x4_t n = vrndnq_f32(vmulq_n_f32(x, 1.4426950408889634f)); + float32x4_t r = vfmsq_f32(x, n, vdupq_n_f32(0.693145751953125f)); + r = vfmsq_f32(r, n, vdupq_n_f32(1.428606765330187e-06f)); + float32x4_t p = vdupq_n_f32(1.0f / 720.0f); + p = vfmaq_f32(vdupq_n_f32(1.0f / 120.0f), p, r); + p = vfmaq_f32(vdupq_n_f32(1.0f / 24.0f), p, r); + p = vfmaq_f32(vdupq_n_f32(1.0f / 6.0f), p, r); + p = vfmaq_f32(vdupq_n_f32(0.5f), p, r); + p = vfmaq_f32(vdupq_n_f32(1.0f), p, r); + p = vfmaq_f32(vdupq_n_f32(1.0f), p, r); + const int32x4_t e = vshlq_n_s32(vaddq_s32(vcvtq_s32_f32(n), vdupq_n_s32(127)), 23); + return vmulq_f32(p, vreinterpretq_f32_s32(e)); +} + +void rel_pos_rows_to_f32(const __fp16* src, size_t row_stride, size_t rows, size_t dim, float scale, float* dst) { + const float32x4_t s = vdupq_n_f32(scale); + for (size_t r = 0; r < rows; ++r) { + for (size_t i = 0; i < dim; i += 8) { + const float16x8_t v = vld1q_f16(src + r * row_stride + i); + vst1q_f32(dst + r * dim + i, vmulq_f32(vcvt_f32_f16(vget_low_f16(v)), s)); + vst1q_f32(dst + r * dim + i + 4, vmulq_f32(vcvt_f32_f16(vget_high_f16(v)), s)); + } + } +} + +void rel_pos_pack_16_rows(const __fp16* src, size_t row_stride, ptrdiff_t row0, size_t num_rows, size_t dim, + bool panels, float* dst) { + auto load = [&](size_t r, size_t i) { + const ptrdiff_t g = row0 + static_cast(r); + return (g < 0 || g >= static_cast(num_rows)) ? vdupq_n_f32(0.0f) : vcvt_f32_f16(vld1_f16(src + g * row_stride + i)); + }; + for (size_t r = 0; r < 16; r += 4) { + for (size_t i = 0; i < dim; i += 4) { + const float32x4_t v[4] = {load(r, i), load(r + 1, i), load(r + 2, i), load(r + 3, i)}; + if (!panels) { + for (size_t j = 0; j < 4; ++j) vst1q_f32(dst + (r + j) * dim + i, v[j]); + continue; + } + const float32x4x2_t lo = vtrnq_f32(v[0], v[1]); + const float32x4x2_t hi = vtrnq_f32(v[2], v[3]); + vst1q_f32(dst + i * 16 + r, vcombine_f32(vget_low_f32(lo.val[0]), vget_low_f32(hi.val[0]))); + vst1q_f32(dst + (i + 1) * 16 + r, vcombine_f32(vget_low_f32(lo.val[1]), vget_low_f32(hi.val[1]))); + vst1q_f32(dst + (i + 2) * 16 + r, vcombine_f32(vget_high_f32(lo.val[0]), vget_high_f32(hi.val[0]))); + vst1q_f32(dst + (i + 3) * 16 + r, vcombine_f32(vget_high_f32(lo.val[1]), vget_high_f32(hi.val[1]))); + } + } +} + +template +inline void rel_pos_lane_fma(float32x4_t (&acc)[4][4], const float32x4_t (&a)[4], const float* b) { + const float32x4_t bv[4] = {vld1q_f32(b), vld1q_f32(b + 4), vld1q_f32(b + 8), vld1q_f32(b + 12)}; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) acc[i][j] = vfmaq_laneq_f32(acc[i][j], bv[j], a[i], L); +} + +void rel_pos_gemm(const float* a, size_t lda, const float* b, bool b_is_keys, size_t depth, + float* c, size_t ldc, size_t rows, size_t cols) { +#ifdef __APPLE__ + cblas_sgemm(CblasRowMajor, CblasNoTrans, b_is_keys ? CblasTrans : CblasNoTrans, (int)rows, (int)cols, (int)depth, + 1.0f, a, (int)lda, b, (int)(b_is_keys ? depth : cols), 0.0f, c, (int)ldc); +#else + const size_t ldb = b_is_keys ? 16 : cols; + const size_t panel = b_is_keys ? 16 * depth : 16; + for (size_t n = 0; n < cols; n += 16) { + const float* bp = b + (n / 16) * panel; + for (size_t m = 0; m < rows; m += 4) { + float32x4_t acc[4][4] = {}; + for (size_t k = 0; k < depth; k += 4) { + const float32x4_t av[4] = {vld1q_f32(a + m * lda + k), vld1q_f32(a + (m + 1) * lda + k), + vld1q_f32(a + (m + 2) * lda + k), vld1q_f32(a + (m + 3) * lda + k)}; + rel_pos_lane_fma<0>(acc, av, bp + k * ldb); + rel_pos_lane_fma<1>(acc, av, bp + (k + 1) * ldb); + rel_pos_lane_fma<2>(acc, av, bp + (k + 2) * ldb); + rel_pos_lane_fma<3>(acc, av, bp + (k + 3) * ldb); + } + for (size_t i = 0; i < 4; ++i) + for (size_t j = 0; j < 4; ++j) vst1q_f32(c + (m + i) * ldc + n + 4 * j, acc[i][j]); + } + } +#endif +} + +void rel_pos_softmax_row(float* s, const float* rel, const __fp16* key_mask, size_t c_lo, size_t c_hi, size_t len) { + std::fill(s, s + c_lo, 0.0f); + std::fill(s + c_hi, s + len, 0.0f); + size_t c = c_lo; + float32x4_t vmax = vdupq_n_f32(-INFINITY); + for (; c + 4 <= c_hi; c += 4) { + float32x4_t x = vaddq_f32(vld1q_f32(s + c), vld1q_f32(rel + c)); + if (key_mask) x = vaddq_f32(x, vcvt_f32_f16(vld1_f16(key_mask + c))); + vst1q_f32(s + c, x); + vmax = vmaxq_f32(vmax, x); + } + float row_max = vmaxvq_f32(vmax); + for (; c < c_hi; ++c) { + s[c] += rel[c] + (key_mask ? static_cast(key_mask[c]) : 0.0f); + row_max = std::max(row_max, s[c]); + } + const float32x4_t vrow_max = vdupq_n_f32(row_max); + float32x4_t vsum = vdupq_n_f32(0.0f); + for (c = c_lo; c + 4 <= c_hi; c += 4) { + const float32x4_t e = rel_pos_exp_f32(vsubq_f32(vld1q_f32(s + c), vrow_max)); + vst1q_f32(s + c, e); + vsum = vaddq_f32(vsum, e); + } + float sum = vaddvq_f32(vsum); + for (; c < c_hi; ++c) { + s[c] = std::exp(s[c] - row_max); + sum += s[c]; + } + const float inv = 1.0f / sum; + for (c = c_lo; c < c_hi; ++c) s[c] *= inv; +} + +} + +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 +) { + constexpr size_t BQ = REL_POS_BLOCK_Q; + const size_t T = seq_len; + const size_t D = head_dim; + const size_t row_stride = num_heads * D; + const size_t window = (window_size == 0 || window_size >= T) ? T : window_size; + const ptrdiff_t center = static_cast((rel_len - 1) / 2); + const size_t padded_len = (T + 15) / 16 * 16; + const size_t num_blocks = (T + BQ - 1) / BQ; + const size_t workers = CactusThreading::get_thread_pool().num_workers(); + const size_t chunk = std::clamp(batch_size * num_heads * num_blocks / (16 * workers), 1, REL_POS_CHUNK_BLOCKS); + const size_t num_chunks = (num_blocks + chunk - 1) / chunk; + + auto key_start = [&](size_t t0) { return (t0 > window ? t0 - window : 0) / 16 * 16; }; + auto key_end = [&](size_t t0) { return std::min(padded_len, (t0 + BQ + window + 15) / 16 * 16); }; + auto rel_start = [&](size_t t0) { return center + 1 + static_cast(key_start(t0)) - static_cast(t0 + BQ); }; + + CactusThreading::parallel_for(batch_size * num_heads * num_chunks, REL_POS_ATTENTION_PARALLEL, + [&](size_t start, size_t end) { + thread_local std::vector keys, values, rels, q_f, qv_f, out_f, scores, rel; + q_f.resize(BQ * D); + qv_f.resize(BQ * D); + out_f.resize(BQ * D); + + for (size_t item = start; item < end;) { + const size_t bh = item / num_chunks; + const size_t run_end = std::min({end, (bh + 1) * num_chunks, item + REL_POS_CHUNK_BLOCKS / chunk}); + const size_t b = bh / num_heads; + const size_t h = bh % num_heads; + const size_t head_offset = b * T * row_stride + h * D; + const size_t blk_begin = (item % num_chunks) * chunk; + const size_t blk_end = std::min(num_blocks, ((run_end - 1) % num_chunks + 1) * chunk); + item = run_end; + + const size_t k0 = key_start(blk_begin * BQ); + const size_t k1 = key_end((blk_end - 1) * BQ); + const ptrdiff_t g0 = rel_start((blk_end - 1) * BQ); + const ptrdiff_t g1 = rel_start(blk_begin * BQ) + static_cast(key_end(blk_begin * BQ) - k0 + BQ); + keys.resize((k1 - k0) * D); + values.resize((k1 - k0) * D); + rels.resize(static_cast(g1 - g0 + 15) / 16 * 16 * D); + for (size_t k = k0; k < k1; k += 16) { + rel_pos_pack_16_rows(key + head_offset, row_stride, k, T, D, REL_POS_PANELS, keys.data() + (k - k0) * D); + rel_pos_pack_16_rows(value + head_offset, row_stride, k, T, D, false, values.data() + (k - k0) * D); + } + for (ptrdiff_t g = g0; g < g1; g += 16) { + rel_pos_pack_16_rows(rel_key + h * D, row_stride, g, rel_len, D, REL_POS_PANELS, rels.data() + (g - g0) * D); + } + + for (size_t blk = blk_begin; blk < blk_end; ++blk) { + const size_t t0 = blk * BQ; + const size_t nq = std::min(BQ, T - t0); + const size_t ks = key_start(t0); + const size_t len = key_end(t0) - ks; + scores.resize(BQ * len); + rel.resize(BQ * (len + BQ)); + + rel_pos_rows_to_f32(query + head_offset + t0 * row_stride, row_stride, nq, D, scale, q_f.data()); + rel_pos_rows_to_f32(rel_query + head_offset + t0 * row_stride, row_stride, nq, D, scale, qv_f.data()); + rel_pos_gemm(q_f.data(), D, keys.data() + (ks - k0) * D, true, D, scores.data(), len, BQ, len); + rel_pos_gemm(qv_f.data(), D, rels.data() + (rel_start(t0) - g0) * D, true, D, rel.data(), len + BQ, BQ, len + BQ); + + const __fp16* mask_row = key_mask ? key_mask + b * T + ks : nullptr; + for (size_t r = 0; r < BQ; ++r) { + float* row = scores.data() + r * len; + if (r < nq) { + 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); + } else { + std::fill(row, row + len, 0.0f); + } + } + rel_pos_gemm(scores.data(), len, values.data() + (ks - k0) * D, false, len, out_f.data(), D, BQ, D); + + for (size_t r = 0; r < nq; ++r) { + __fp16* o = output + head_offset + (t0 + r) * row_stride; + const float* src = out_f.data() + r * D; + for (size_t d = 0; d < D; d += 8) { + vst1q_f16(o + d, vcombine_f16(vcvt_f16_f32(vld1q_f32(src + d)), vcvt_f16_f32(vld1q_f32(src + d + 4)))); + } + } + } + } + }); +} diff --git a/cactus-kernels/src/cactus_kernels.metal b/cactus-kernels/src/cactus_kernels.metal index 1ef064202..1e12d34a8 100644 --- a/cactus-kernels/src/cactus_kernels.metal +++ b/cactus-kernels/src/cactus_kernels.metal @@ -3229,6 +3229,140 @@ kernel void rel_pos_bias_f16(device const half* q [[buffer(0)]], device const ha for (uint d = 0; d < D; ++d) acc = fma((float)qv[d], (float)rv[d], acc); y[((size_t)b * H + h) * T * T + (size_t)t * T + j] = (half)(acc * scale); } +struct RelPosAttnU { uint T, H, D, R, window, has_mask; float scale; }; + +kernel void rel_pos_attn_f16(device const half* q [[buffer(0)]], device const half* k [[buffer(1)]], + device const half* v [[buffer(2)]], device const half* qv [[buffer(3)]], + device const half* r [[buffer(4)]], device const half* mask [[buffer(5)]], + device half* o [[buffer(6)]], constant RelPosAttnU& U [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], uint tl [[thread_index_in_threadgroup]]) { + const uint D = 128u, D4 = 32u, QB = 8u, NSG = 4u, QT = QB * NSG, BK = 32u, RB = 64u, LDS = 72u, NT = NSG * 32u; + threadgroup half4 kv_s[BK * D4]; + threadgroup half4 r_s[RB * D4]; + threadgroup half4 qv_s[QT * D4]; + const uint sg = tl >> 5, lane = tl & 31u; + const uint T = U.T; + const size_t stride4 = (size_t)U.H * D4; + const uint t_tg = tg.x * QT, t0 = t_tg + sg * QB; + const size_t head4 = (size_t)tg.z * T * stride4 + (size_t)tg.y * D4; + device const half4* q4 = (device const half4*)q + head4; + device const half4* qv4 = (device const half4*)qv + head4; + device const half4* k4 = (device const half4*)k + head4; + device const half4* v4 = (device const half4*)v + head4; + device const half4* r4 = (device const half4*)r + (size_t)tg.y * D4; + const uint W = (U.window == 0u || U.window >= T) ? T : U.window; + const int center = (int)(U.R - 1u) / 2; + threadgroup float* s_sg = (threadgroup float*)r_s + sg * QB * LDS; + threadgroup half* p_sg = (threadgroup half*)((threadgroup float*)r_s + NSG * QB * LDS) + sg * QB * BK; + threadgroup const half* kv_h = (threadgroup const half*)kv_s; + threadgroup const half* r_h = (threadgroup const half*)r_s; + threadgroup const half* qv_h = (threadgroup const half*)qv_s + sg * QB * D; + + for (uint i = tl; i < QT * D4; i += NT) { + const uint t = t_tg + i / D4; + r_s[i] = t < T ? q4[(size_t)t * stride4 + i % D4] : half4(0.0h); + qv_s[i] = t < T ? qv4[(size_t)t * stride4 + i % D4] : half4(0.0h); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_matrix qm[16]; + simdgroup_matrix acc[16]; + for (uint c = 0; c < 16u; ++c) { + simdgroup_load(qm[c], r_h + sg * QB * D + c * 8u, D); + acc[c] = make_filled_simdgroup_matrix(0.0f); + } + float m_row = -INFINITY, l_row = 0.0f, corr_row = 1.0f; + + const uint band = QB * (NSG - 1u - sg); + const uint kstart = t_tg > W ? t_tg - W : 0u; + const uint kend = min(T, t_tg + QT + W); + for (uint j0 = kstart; j0 < kend; j0 += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g0 = center + (int)j0 - (int)(t_tg + QT - 1u); + for (uint i = tl; i < BK * D4; i += NT) { + const uint j = j0 + i / D4; + kv_s[i] = j < T ? k4[(size_t)j * stride4 + i % D4] : half4(0.0h); + } + for (uint i = tl; i < RB * D4; i += NT) { + const int g = g0 + (int)(i / D4); + r_s[i] = (g >= 0 && g < (int)U.R) ? r4[(size_t)g * stride4 + i % D4] : half4(0.0h); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_matrix sm[4], gm[5]; + for (uint n = 0; n < 4u; ++n) sm[n] = make_filled_simdgroup_matrix(0.0f); + for (uint n = 0; n < 5u; ++n) gm[n] = make_filled_simdgroup_matrix(0.0f); + for (uint c = 0; c < 16u; ++c) { + simdgroup_matrix bm, qvm; + simdgroup_load(qvm, qv_h + c * 8u, D); + for (uint n = 0; n < 4u; ++n) { + simdgroup_load(bm, kv_h + n * 8u * D + c * 8u, D, ulong2(0, 0), true); + simdgroup_multiply_accumulate(sm[n], qm[c], bm, sm[n]); + } + for (uint n = 0; n < 5u; ++n) { + simdgroup_load(bm, r_h + (band + n * 8u) * D + c * 8u, D, ulong2(0, 0), true); + simdgroup_multiply_accumulate(gm[n], qvm, bm, gm[n]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint n = 0; n < 4u; ++n) simdgroup_store(sm[n], s_sg + n * 8u, LDS); + for (uint n = 0; n < 5u; ++n) simdgroup_store(gm[n], s_sg + BK + n * 8u, LDS); + for (uint i = tl; i < BK * D4; i += NT) { + 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); + + const uint j = j0 + lane; + const float mv = (U.has_mask && j < T) ? (float)mask[(size_t)tg.z * T + j] : 0.0f; + bool rescale = false; + for (uint i = 0; i < QB; ++i) { + const uint t = t0 + i; + const bool inside = j < T && t < T && j + W >= t && j <= t + W; + const float sv = inside ? (s_sg[i * LDS + lane] + s_sg[i * LDS + BK + lane + 7u - i]) * U.scale + mv : -INFINITY; + const float mi = simd_shuffle(m_row, (ushort)i); + const float mnew = max(mi, simd_max(sv)); + const float corr = mnew == -INFINITY ? 1.0f : exp(mi - mnew); + const float pv = sv == -INFINITY ? 0.0f : exp(sv - mnew); + const float psum = simd_sum(pv); + if ((lane & 7u) == i) { m_row = mnew; l_row = l_row * corr + psum; corr_row = corr; } + p_sg[i * BK + lane] = (half)pv; + rescale = rescale || corr != 1.0f; + } + if (rescale) { + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = lane; e < 64u; e += 32u) { + const float c = simd_shuffle(corr_row, (ushort)(e / 8u)); + s_sg[e] = (e / 8u == e % 8u) ? c : 0.0f; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + simdgroup_matrix dm; + simdgroup_load(dm, s_sg, 8u); + for (uint c = 0; c < 16u; ++c) simdgroup_multiply(acc[c], dm, acc[c]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint kk = 0; kk < BK; kk += 8u) { + simdgroup_matrix pm, vm; + simdgroup_load(pm, p_sg + kk, BK); + for (uint c = 0; c < 16u; ++c) { + simdgroup_load(vm, kv_h + kk * D + c * 8u, D); + simdgroup_multiply_accumulate(acc[c], pm, vm, acc[c]); + } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint c = 0; c < 16u; ++c) { + simdgroup_store(acc[c], s_sg, 8u); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = lane; e < 64u; e += 32u) { + const uint i = e / 8u, t = t0 + i; + const float li = simd_shuffle(l_row, (ushort)i); + if (t < T) o[(head4 + (size_t)t * stride4) * 4u + c * 8u + e % 8u] = (half)(li > 0.0f ? s_sg[e] / li : 0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } +} + kernel void conv_cache_append_f16(device const uchar* src [[buffer(0)]], device half* ring [[buffer(1)]], device half* out [[buffer(2)]], constant uint& hd [[buffer(3)]], constant uint& ws [[buffer(4)]], diff --git a/cactus-kernels/src/metal_backend.mm b/cactus-kernels/src/metal_backend.mm index cb55e36e1..7dd7725ce 100644 --- a/cactus-kernels/src/metal_backend.mm +++ b/cactus-kernels/src/metal_backend.mm @@ -64,7 +64,7 @@ id psoBatchnorm=nil, psoGroupnorm=nil, psoBiasAddRows=nil, psoEwChain=nil; id psoAttnFlash=nil, psoRmsSimd=nil, psoScatterRows=nil; id psoTranspose2d=nil, psoBcastRows=nil, psoRmsAddSimd=nil; - id psoConvCacheAppend=nil, psoRelPosBias=nil, psoGemvBias=nil; + id psoConvCacheAppend=nil, psoRelPosBias=nil, psoRelPosAttn=nil, psoGemvBias=nil; id psoTopkRows=nil, psoMoeT=nil, psoMoeUp=nil; id psoMoeT2=nil, psoMoeDownAcc=nil, psoRopePair=nil; id psoRmsScale=nil, psoSoftmaxTopk=nil, psoRopePairRms=nil; @@ -129,7 +129,7 @@ psoAttnFlash=pso("attn_flash_f16"); psoRmsSimd=pso("rms_norm_simd_f16"); psoScatterRows=pso("strided_scatter_rows_f16"); psoTranspose2d=pso("transpose2d_f16"); psoBcastRows=pso("bcast_binary_rows_f16"); psoRmsAddSimd=pso("rms_norm_add_simd_f16"); - psoConvCacheAppend=pso("conv_cache_append_f16"); psoRelPosBias=pso("rel_pos_bias_f16"); + psoConvCacheAppend=pso("conv_cache_append_f16"); psoRelPosBias=pso("rel_pos_bias_f16"); psoRelPosAttn=pso("rel_pos_attn_f16"); psoGemvBias=pso("gemv_bias_f16"); psoTopkRows=pso("topk_rows_f16"); psoMoeT=pso("cq4_moe_transform"); psoMoeUp=pso("cq4_moe_gemv_up"); psoMoeT2=pso("cq4_moe_transform2"); @@ -2380,6 +2380,24 @@ bool cactus_metal_encode_rel_pos_bias(void* y, const void* q, const void* r, return true; } +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) { + if (!ctx().ok || !ctx().psoRelPosAttn || D != 128) return false; + ensureEncoder(); + struct { uint32_t T, H, D, R, window, has_mask; float scale; } U = { T, H, D, R, window, mask ? 1u : 0u, scale }; + const size_t bytes = (size_t)B*T*H*D*2; + [g_enc setComputePipelineState:ctx().psoRelPosAttn]; + setBufAt(q, bytes, 0); setBufAt(k, bytes, 1); setBufAt(v, bytes, 2); setBufAt(qv, bytes, 3); + setBufAt(r, (size_t)R*H*D*2, 4); + if (mask) setBufAt(mask, (size_t)B*T*2, 5); + else [g_enc setBuffer:ctx().dummy offset:0 atIndex:5]; + setBufAt(out, bytes, 6); + [g_enc setBytes:&U length:sizeof(U) atIndex:7]; + [g_enc dispatchThreadgroups:MTLSizeMake((T+31)/32, H, B) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; +} + bool cactus_metal_encode_conv_cache_append(void* out, const void* src, void* ring, uint32_t hd, uint32_t ws, uint32_t nnew, uint32_t head0, uint32_t count_new, uint32_t num_rows, int src_f32) { diff --git a/cactus-kernels/src/metal_backend_stub.cpp b/cactus-kernels/src/metal_backend_stub.cpp index 9ecafc05b..07013390f 100644 --- a/cactus-kernels/src/metal_backend_stub.cpp +++ b/cactus-kernels/src/metal_backend_stub.cpp @@ -111,6 +111,8 @@ bool cactus_metal_encode_conv_cache_append(void*, const void*, void*, uint32_t, uint32_t, uint32_t, uint32_t, int) { return false; } bool cactus_metal_encode_rel_pos_bias(void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, int, float) { return false; } +bool cactus_metal_encode_rel_pos_attention_f16(void*, const void*, const void*, const void*, const void*, + const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float) { return false; } bool cactus_metal_encode_gemv_bias(void*, const void*, const void*, const void*, uint32_t, uint32_t, int) { return false; } bool cactus_metal_encode_rope_pair(void*, const void*, const void*, const void*, diff --git a/cactus-kernels/tests/test_attention.cpp b/cactus-kernels/tests/test_attention.cpp index 699159815..b82dbc9ac 100644 --- a/cactus-kernels/tests/test_attention.cpp +++ b/cactus-kernels/tests/test_attention.cpp @@ -120,6 +120,48 @@ bool test_attention_f16_long_decode() { return true; } +static bool check_rel_pos_attention(size_t seq, size_t window, size_t rel_len, size_t padded_keys) { + const size_t heads = 2, dim = 16, stride = heads * dim, n = seq * stride; + std::vector<__fp16> q(n), k(n), v(n), qv(n), out(n), rel(rel_len * stride); + std::vector<__fp16> mask(seq, static_cast<__fp16>(0.0f)); + for (auto* t : {&q, &k, &v, &qv, &rel}) fill_random_fp16(*t, -1.0f, 1.0f); + for (size_t j = seq - padded_keys; j < seq; ++j) mask[j] = static_cast<__fp16>(-10000.0f); + const float scale = 0.25f; + cactus_rel_pos_attention_f16(q.data(), k.data(), v.data(), qv.data(), rel.data(), mask.data(), out.data(), + 1, seq, heads, dim, rel_len, scale, window); + + const size_t center = (rel_len - 1) / 2; + for (size_t h = 0; h < heads; ++h) { + for (size_t t = 0; t < seq; ++t) { + const size_t j_lo = (window > 0 && t > window) ? t - window : 0; + const size_t j_hi = window > 0 ? std::min(seq, t + window + 1) : seq; + std::vector p(j_hi - j_lo); + double max_score = -INFINITY, sum = 0.0; + for (size_t j = j_lo; j < j_hi; ++j) { + double score = 0.0; + for (size_t d = 0; d < dim; ++d) { + score += static_cast(q[t * stride + h * dim + d]) * static_cast(k[j * stride + h * dim + d]) + + static_cast(qv[t * stride + h * dim + d]) + * static_cast(rel[(center + j - t) * stride + h * dim + d]); + } + p[j - j_lo] = score * scale + static_cast(mask[j]); + max_score = std::max(max_score, p[j - j_lo]); + } + for (double& x : p) { x = std::exp(x - max_score); sum += x; } + for (size_t d = 0; d < dim; ++d) { + double expected = 0.0; + for (size_t j = j_lo; j < j_hi; ++j) expected += p[j - j_lo] / sum * static_cast(v[j * stride + h * dim + d]); + if (std::abs(static_cast(out[t * stride + h * dim + d]) - expected) > 2e-3) return false; + } + } + } + return true; +} + +bool test_rel_pos_attention() { + return check_rel_pos_attention(150, 20, 41, 7) && check_rel_pos_attention(70, 0, 139, 3); +} + bool test_attention_f16_streaming_additive_mask() { const size_t batch = 1, seq = 64, heads = 2, dim = 16; std::vector<__fp16> q(batch * seq * heads * dim), k(batch * seq * heads * dim); @@ -235,6 +277,7 @@ int main() { runner.run_test("attention_f16", test_attention_f16()); runner.run_test("attention_f16_long_decode", test_attention_f16_long_decode()); runner.run_test("attention_f16_streaming_additive_mask", test_attention_f16_streaming_additive_mask()); + runner.run_test("rel_pos_attention", test_rel_pos_attention()); runner.print_benchmarks_header(); runner.run_bench("benchmarks", run_benchmarks()); runner.print_summary(); diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index fc9115713..4b519d27e 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -568,6 +568,11 @@ class cactus_tensor_info_t(ctypes.Structure): cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_rel_pos_bias.restype = ctypes.c_int +_lib.cactus_graph_rel_pos_attention.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, + ctypes.c_bool, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_rel_pos_attention.restype = ctypes.c_int _lib.cactus_graph_attention_int8_hybrid.argtypes = [ cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.POINTER(ctypes.c_int8), ctypes.POINTER(ctypes.c_int8), @@ -2570,6 +2575,31 @@ def rel_pos_bias(self, query, relative_key, scale, backend=None): raise RuntimeError(_err("graph_rel_pos_bias failed")) return self._apply_backend(self._tensor_from_node(out.value), backend) + def rel_pos_attention(self, query, key, value, rel_query, relative_key, key_mask, scale, window_size=0, backend=None): + query = self._ensure_tensor(query) + key = self._ensure_tensor(key) + value = self._ensure_tensor(value) + rel_query = self._ensure_tensor(rel_query) + relative_key = self._ensure_tensor(relative_key) + key_mask = None if key_mask is None else self._ensure_tensor(key_mask) + out = cactus_node_t() + rc = _lib.cactus_graph_rel_pos_attention( + self.h, + cactus_node_t(query.id), + cactus_node_t(key.id), + cactus_node_t(value.id), + cactus_node_t(rel_query.id), + cactus_node_t(relative_key.id), + ctypes.c_bool(key_mask is not None), + cactus_node_t(0 if key_mask is None else key_mask.id), + ctypes.c_float(float(scale)), + ctypes.c_size_t(int(window_size)), + ctypes.byref(out), + ) + if rc != 0: + raise RuntimeError(_err("graph_rel_pos_attention failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + def attention_int8_hybrid(self, query, key_new, value_new, scale, position_offset, cached_keys, cached_values, k_scales, v_scales, cache_len, num_kv_heads, head_dim, window_size=0, backend=None): diff --git a/python/cactus/convert/tests/test_parakeet_local_attention.py b/python/cactus/convert/tests/test_parakeet_local_attention.py new file mode 100644 index 000000000..5ee2907a5 --- /dev/null +++ b/python/cactus/convert/tests/test_parakeet_local_attention.py @@ -0,0 +1,125 @@ +from types import SimpleNamespace + +import numpy as np +import torch + +from cactus.transpile.tdt_runtime import ParakeetTDTSelfAttention +from cactus.transpile.tdt_runtime import rel_pos_attention + + +def _random_inputs(seq_len, rel_len, *, heads=2, dim=16, padded=3, seed=0): + generator = torch.Generator().manual_seed(seed) + shape = (1, seq_len, heads, dim) + tensors = [torch.randn(shape, generator=generator) for _ in range(4)] + relative_key = torch.randn((1, rel_len, heads, dim), generator=generator) + key_mask = torch.zeros((1, seq_len)) + key_mask[:, seq_len - padded:] = -10000.0 + return (*tensors, relative_key, key_mask) + + +def _dense_reference(query, key, value, rel_query, relative_key, key_mask, scale, window): + seq_len = query.shape[1] + center = (relative_key.shape[1] - 1) // 2 + out = torch.empty_like(query) + for t in range(seq_len): + for h in range(query.shape[2]): + scores = [] + for j in range(seq_len): + if window > 0 and abs(t - j) > window: + scores.append(float("-inf")) + continue + content = torch.dot(query[0, t, h], key[0, j, h]) + position = torch.dot(rel_query[0, t, h], relative_key[0, center - (t - j), h]) + scores.append(float((content + position) * scale + key_mask[0, j])) + weights = torch.softmax(torch.tensor(scores), dim=0) + out[0, t, h] = weights @ value[0, :, h] + return out + + +def test_rel_pos_attention_matches_dense_reference(): + for seq_len, window, rel_len in ((40, 6, 13), (40, 0, 79), (300, 20, 41)): + inputs = _random_inputs(seq_len, rel_len) + expected = _dense_reference(*inputs, 0.25, window) + actual = rel_pos_attention(*inputs, 0.25, window) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + +def _attention_layer(hidden=32, heads=4): + torch.manual_seed(0) + state_dict = {f"layer.{name}.weight": torch.randn(hidden, hidden) * 0.2 + for name in ("linear_q", "linear_k", "linear_v", "linear_out", "linear_pos")} + state_dict["layer.pos_bias_u"] = torch.randn(heads, hidden // heads) + state_dict["layer.pos_bias_v"] = torch.randn(heads, hidden // heads) + config = SimpleNamespace( + hidden_dim=hidden, + attention_heads=heads, + attention_head_dim=hidden // heads, + attention_scale=(hidden // heads) ** -0.5, + ) + return ParakeetTDTSelfAttention(config, "layer", state_dict).eval() + + +def test_local_attention_exports_single_semantic_op(): + from cactus.transpile.capture_pytorch import capture_model + + class _Windowed(torch.nn.Module): + def __init__(self): + super().__init__() + self.attention = _attention_layer() + + def forward(self, x, mask): + return self.attention(x, mask, 4) + + captured = capture_model(_Windowed().eval(), (torch.randn(1, 12, 32), torch.zeros(1, 1, 1, 12)), strict=False) + nodes = [ + captured.ir_graph.nodes[node_id] + for node_id in captured.ir_graph.order + if captured.ir_graph.nodes[node_id].op == "rel_pos_attention" + ] + assert len(nodes) == 1 + assert nodes[0].attrs == {"scale": 8 ** -0.5, "window_size": 4} + assert len(nodes[0].inputs) == 6 + + +def test_engine_op_matches_torch_reference(): + from cactus.bindings.cactus import Graph + + for seq_len, window, rel_len in ((150, 20, 41), (70, 0, 139)): + inputs = _random_inputs(seq_len, rel_len, dim=32, padded=5, seed=1) + inputs = tuple(tensor.half().float() for tensor in inputs) + expected = rel_pos_attention(*inputs, 32 ** -0.5, window) + + graph = Graph() + nodes = [graph.input(tuple(tensor.shape)) for tensor in inputs] + out = graph.rel_pos_attention(*nodes, 32 ** -0.5, window) + for node, tensor in zip(nodes, inputs): + graph.set_input(node, tensor.numpy().astype(np.float16)) + graph.execute() + np.testing.assert_allclose(out.numpy().astype(np.float32), expected.numpy(), atol=3e-3) + + +def test_chunked_subsampling_matches_single_pass(): + from cactus.transpile.tdt_runtime import ParakeetTDTPreEncode + + torch.manual_seed(0) + channels, mels, hidden = 4, 16, 8 + state_dict = { + "encoder.pre_encode.conv.0.weight": torch.randn(channels, 1, 3, 3), + "encoder.pre_encode.conv.0.bias": torch.randn(channels), + "encoder.pre_encode.out.weight": torch.randn(hidden, channels * 2), + "encoder.pre_encode.out.bias": torch.randn(hidden), + } + for index in (2, 5): + state_dict[f"encoder.pre_encode.conv.{index}.weight"] = torch.randn(channels, 1, 3, 3) + state_dict[f"encoder.pre_encode.conv.{index}.bias"] = torch.randn(channels) + state_dict[f"encoder.pre_encode.conv.{index + 1}.weight"] = torch.randn(channels, channels, 1, 1) + state_dict[f"encoder.pre_encode.conv.{index + 1}.bias"] = torch.randn(channels) + state_dict = {name: tensor * 0.3 for name, tensor in state_dict.items()} + config = SimpleNamespace(subsampling_conv_channels=channels, num_mel_bins=mels, hidden_dim=hidden) + pre_encode = ParakeetTDTPreEncode(config, state_dict).eval() + features = torch.randn(1, 7003, mels) + with torch.no_grad(): + chunked = pre_encode(features) + single = pre_encode._subsample(features) + assert chunked.shape == single.shape + torch.testing.assert_close(chunked, single, atol=1e-5, rtol=1e-5) diff --git a/python/cactus/transpile/aten_ops.py b/python/cactus/transpile/aten_ops.py index 9096b8515..29df500af 100644 --- a/python/cactus/transpile/aten_ops.py +++ b/python/cactus/transpile/aten_ops.py @@ -54,6 +54,7 @@ def canonical_torch_op(target: Any) -> str: "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", # Keep the legacy importer key until the importer itself is renamed. "aten.diff.default": "aten.diff.default", } diff --git a/python/cactus/transpile/canonicalize/cleanup.py b/python/cactus/transpile/canonicalize/cleanup.py index c68e00447..c3b9f0f38 100644 --- a/python/cactus/transpile/canonicalize/cleanup.py +++ b/python/cactus/transpile/canonicalize/cleanup.py @@ -132,6 +132,7 @@ "lfm2_moe_layer_gated", "qwen2_moe_layer_gated", "gemma4_moe_layer_gated", + "rel_pos_attention", "rope", } diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py index 9666f7c03..f70469d59 100644 --- a/python/cactus/transpile/importers.py +++ b/python/cactus/transpile/importers.py @@ -1860,6 +1860,23 @@ def import_gemma4_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op): op="gemma4_moe_layer_gated", has_expert_bias=False, activation="gelu") +def import_rel_pos_attention(ir, node, ctx, *, shape, dtype, torch_op): + if len(node.args) != 8: + ctx.fail(f"rel_pos_attention expected 8 args, got {len(node.args)}") + ir_node = IRNode( + id=node_id(node), + op="rel_pos_attention", + inputs=[value_id(node.args[i], ctx) for i in range(6)], + outputs=[value_id(node, ctx)], + attrs={ + "scale": float(extract_literals(node.args[6])), + "window_size": int(extract_literals(node.args[7])), + }, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + OP_IMPORTERS = { "arange": import_arange, "identity": import_identity, @@ -1957,4 +1974,5 @@ def import_gemma4_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op): "lfm2_moe_layer_gated": import_lfm2_moe_layer_gated, "qwen2_moe_layer_gated": import_qwen2_moe_layer_gated, "gemma4_moe_layer_gated": import_gemma4_moe_layer_gated, + "rel_pos_attention": import_rel_pos_attention, } diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index 8896291d2..7b321a936 100644 --- a/python/cactus/transpile/lower.py +++ b/python/cactus/transpile/lower.py @@ -2218,6 +2218,26 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> ) return [g.rel_pos_bias(query, relative_key, float(node.attrs.get("scale", 1.0)))] + if op == "rel_pos_attention": + query, key, value, rel_query, relative_key, key_mask = ( + _ensure_fp16_tensor(g, _tensor(env, input_id)) for input_id in node.inputs + ) + mask_shape = (int(query.shape[0]), int(query.shape[1])) + if tuple(int(dim) for dim in key_mask.shape) != mask_shape: + key_mask = g.reshape(key_mask, mask_shape) + return [ + g.rel_pos_attention( + query, + key, + value, + rel_query, + relative_key, + key_mask, + float(node.attrs["scale"]), + int(node.attrs.get("window_size", 0)), + ) + ] + if op in {"gated_deltanet_prefill", "gated_deltanet_decode"}: x = _tensor(env, node.inputs[0]) qkv_weight = _tensor(env, node.inputs[1]) diff --git a/python/cactus/transpile/ops.py b/python/cactus/transpile/ops.py index 4dcdd243a..a6102ecf1 100644 --- a/python/cactus/transpile/ops.py +++ b/python/cactus/transpile/ops.py @@ -147,6 +147,7 @@ class OpSchema: OpSchema("rms_norm", num_inputs=2, attrs=("eps",), backend_op="rms_norm"), OpSchema("rope", num_inputs=1, attrs=("theta", "position_offset"), backend_op="rope"), OpSchema("rel_pos_bias", num_inputs=2, attrs=("scale",), backend_op="rel_pos_bias"), + OpSchema("rel_pos_attention", num_inputs=6, attrs=("scale", "window_size"), backend_op="rel_pos_attention"), OpSchema( "attention", num_inputs=3, diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 70b02d485..baa74e848 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -20,6 +20,57 @@ from cactus.transpile.model_profiles import PARAKEET_TDT_PROFILE +_FULL_ATTENTION_MAX_ENCODER_FRAMES = 688 +_LOCAL_ATTENTION_WINDOW = 256 +_REL_POS_ATTENTION_QUERY_BLOCK = 256 +_PRE_ENCODE_CHUNK_FRAMES = 3000 +_PRE_ENCODE_HALO_FRAMES = 16 + + +@torch.library.custom_op("cactus_transpile::rel_pos_attention", mutates_args=()) +def rel_pos_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + rel_query: torch.Tensor, + relative_key: torch.Tensor, + key_mask: torch.Tensor, + scale: float, + window: int, +) -> torch.Tensor: + seq_len = int(query.shape[1]) + rel_len = int(relative_key.shape[1]) + center = (rel_len - 1) // 2 + reach = seq_len - 1 if window <= 0 or window >= seq_len else window + q = query.permute(0, 2, 1, 3).float() + k = key.permute(0, 2, 1, 3).float() + v = value.permute(0, 2, 1, 3).float() + q_rel = rel_query.permute(0, 2, 1, 3).float() + r = relative_key.permute(0, 2, 1, 3).float() + mask = key_mask.float() + out = torch.empty_like(q) + for t0 in range(0, seq_len, _REL_POS_ATTENTION_QUERY_BLOCK): + t1 = min(seq_len, t0 + _REL_POS_ATTENTION_QUERY_BLOCK) + k0 = max(0, t0 - reach) + k1 = min(seq_len, t1 + reach) + offset = ( + 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) + content = q[:, :, t0:t1] @ k[:, :, k0:k1].transpose(-1, -2) + position = (q_rel[:, :, t0:t1] @ r.transpose(-1, -2)).gather(-1, rel_index) + scores = (content + position) * scale + mask[:, None, None, k0:k1] + scores = scores.masked_fill(offset.abs() > reach, float("-inf")) + out[:, :, t0:t1] = scores.softmax(dim=-1) @ v[:, :, k0:k1] + return out.permute(0, 2, 1, 3).to(query.dtype).contiguous() + + +@rel_pos_attention.register_fake +def _rel_pos_attention_fake(query, key, value, rel_query, relative_key, key_mask, scale, window): + return torch.empty_like(query) + + def _cfg_get(config: dict[str, Any], key: str, default: Any = None) -> Any: value = config.get(key, default) return default if value is None else value @@ -404,19 +455,6 @@ def _relative_position_embeddings(*, seq_len: int, hidden_dim: int, device: torc return embeddings.to(dtype=dtype) -def _relative_position_bias(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: - batch, heads, seq_len, _ = query.shape - scores = torch.matmul(query, relative_key.transpose(-1, -2)) - rel_index = ( - torch.arange(seq_len, device=query.device).view(seq_len, 1) - - torch.arange(seq_len, device=query.device).view(1, seq_len) - + (seq_len - 1) - ) - rel_index = rel_index.view(1, 1, seq_len, seq_len).expand(batch, heads, seq_len, seq_len) - gathered = scores.gather(-1, rel_index) - return gathered * float(scale) - - class ParakeetTDTFeedForward(nn.Module): def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): super().__init__() @@ -458,7 +496,7 @@ def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, self.pos_bias_v = nn.Parameter(state_dict[f"{prefix}.pos_bias_v"].clone()) self.config = config - def forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, attn_mask: torch.Tensor, window: int = 0) -> torch.Tensor: batch, seq_len, _ = x.shape num_heads = self.config.attention_heads head_dim = self.config.attention_head_dim @@ -467,39 +505,29 @@ def forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: k = self.linear_k(x).view(batch, seq_len, num_heads, head_dim) v = self.linear_v(x).view(batch, seq_len, num_heads, head_dim) + reach = window if window > 0 else seq_len - 1 rel_pos = _relative_position_embeddings( - seq_len=seq_len, + seq_len=reach + 1, hidden_dim=self.config.hidden_dim, device=x.device, dtype=x.dtype, ) - rel_k = self.linear_pos(rel_pos).view(1, 2 * seq_len - 1, num_heads, head_dim) + rel_k = self.linear_pos(rel_pos).view(1, 2 * reach + 1, num_heads, head_dim) q_u = q + self.pos_bias_u.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) q_v = q + self.pos_bias_v.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) - q_u_heads = q_u.permute(0, 2, 1, 3) - q_v_heads = q_v.permute(0, 2, 1, 3) - k_heads = k.permute(0, 2, 1, 3) - v_heads = v.permute(0, 2, 1, 3) - rel_k_heads = rel_k.permute(0, 2, 1, 3) - - rel_bias = _relative_position_bias( - q_v_heads, - rel_k_heads, - scale=self.config.attention_scale, - ) - rel_bias = rel_bias + attn_mask - attn = F.scaled_dot_product_attention( - q_u_heads, - k_heads, - v_heads, - attn_mask=rel_bias, - dropout_p=0.0, - is_causal=False, + attn = rel_pos_attention( + q_u, + k, + v, + q_v, + rel_k, + attn_mask.reshape(batch, seq_len), + self.config.attention_scale, + window, ) - attn = attn.permute(0, 2, 1, 3).reshape(batch, seq_len, self.config.hidden_dim) - return self.linear_out(attn) + return self.linear_out(attn.reshape(batch, seq_len, self.config.hidden_dim)) class ParakeetTDTConformerConv(nn.Module): @@ -584,9 +612,9 @@ def __init__(self, config: ParakeetTDTConfig, layer_index: int, state_dict: dict module.bias.data.copy_(state_dict[f"{prefix}.{name}.bias"].to(dtype=module.bias.dtype)) self.config = config - def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, attn_mask: torch.Tensor, window: int = 0) -> torch.Tensor: x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act) - x = x + self.self_attn(self.norm_self_att(x), attn_mask) + x = x + self.self_attn(self.norm_self_att(x), attn_mask, window) x = x + self.conv(self.norm_conv(x), conv_mask) x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act) return self.norm_out(x) @@ -625,6 +653,19 @@ def __init__(self, config: ParakeetTDTConfig, state_dict: dict[str, torch.Tensor def forward(self, input_features: torch.Tensor) -> torch.Tensor: if input_features.ndim != 3: raise ValueError(f"expected input_features [batch, frames, mels], got {tuple(input_features.shape)}") + frames = int(input_features.shape[1]) + if frames <= _PRE_ENCODE_CHUNK_FRAMES: + return self._subsample(input_features) + chunks = [] + for start in range(0, frames, _PRE_ENCODE_CHUNK_FRAMES): + end = min(frames, start + _PRE_ENCODE_CHUNK_FRAMES) + low = max(0, start - _PRE_ENCODE_HALO_FRAMES) + out = self._subsample(input_features[:, low : min(frames, end + _PRE_ENCODE_HALO_FRAMES)]) + first = (start - low) // 8 + chunks.append(out[:, first : first + (end - start + 7) // 8]) + return torch.cat(chunks, dim=1) + + def _subsample(self, input_features: torch.Tensor) -> torch.Tensor: x = input_features.unsqueeze(1) x = F.relu(self.conv[0](x)) x = F.relu(self.conv[3](self.conv[2](x))) @@ -647,8 +688,9 @@ def forward(self, input_features: torch.Tensor, input_features_mask: torch.Tenso pad_mask = input_features_mask[:, :: self.config.subsampling_factor] conv_mask = pad_mask.unsqueeze(1) attn_mask = ((pad_mask - 1.0) * 10000.0).unsqueeze(1).unsqueeze(1) + window = 0 if int(x.shape[1]) <= _FULL_ATTENTION_MAX_ENCODER_FRAMES else _LOCAL_ATTENTION_WINDOW for layer in self.layers: - x = layer(x, conv_mask, attn_mask) + x = layer(x, conv_mask, attn_mask, window) return x