perf: compute the streaming linear_pos projection once per window length

src/streaming_encoder.cpp

@@ -91,7 +91,8 @@ struct StreamLayerCapture {
 //   * The NEXT caches are emitted as captured outputs (cap->next_channel/_time):
 //       next_channel = cat([cache_channel[Tc:], attn_in[:Tc]])  (keep last cache_len)
 //       next_time    = last LP cols of [cache_time(LP) ; glu(Tc)]
-// `xt` is the chunk input [D, Tc]; `pe` is the pos-emb [D, pos_len] for
+// `xt` is the chunk input [D, Tc]; `pos_heads` is this layer's linear_pos(pos_emb)
+// split into heads, host [dk, pos_len, H] from StreamingEncoder::pos_heads, for
 // pos_len = 2*(Tc+cache_len)-1. `cache_ch_t` is the K/V cache input [D, cache_len];
 // `cache_t_t` is the conv cache input [D, LP]. Returns the layer output [D, Tc].
 //
@@ -100,7 +101,7 @@ struct StreamLayerCapture {
 // with the cache reads/updates moved in-graph, so the numerics are preserved.
 static ggml_tensor* build_stream_layer(
         ggml_context* ctx, const ModelLoader& ml, int layer_idx,
-        ggml_tensor* xt, int Tc, ggml_tensor* pe, int pos_len,
+        ggml_tensor* xt, int Tc, const float* pos_heads, int pos_len,
         int cache_len, int clc_len, int n_heads, int d_model, int conv_kernel,
         int att_left, int att_right, const std::string& conv_norm_type,
         ggml_tensor* cache_ch_t, ggml_tensor* cache_t_t,
@@ -162,15 +163,17 @@ static ggml_tensor* build_stream_layer(
         ggml_tensor* q = lin("linear_q.weight", "linear_q.bias", attn_in); // [D, Tc]
         ggml_tensor* k = lin("linear_k.weight", "linear_k.bias", kvt);     // [D, Tk]
         ggml_tensor* v = lin("linear_v.weight", "linear_v.bias", kvt);     // [D, Tk]
-        ggml_tensor* p = lin("linear_pos.weight", nullptr, pe);            // [D, P]
 
         auto to_heads = [&](ggml_tensor* t, int n) {
             t = ggml_reshape_3d(ctx, t, dk, H, n);
             return ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H]
         };
         ggml_tensor* qh = to_heads(q, Tc);
         ggml_tensor* kh = to_heads(k, Tk);
-        ggml_tensor* ph = to_heads(p, pos_len);
+        // linear_pos(pos_emb) split into heads, precomputed per window length.
+        int64_t ph_ne[3] = {dk, pos_len, H};
+        ggml_tensor* ph = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, ph_ne, pos_heads,
+                                                 (size_t)dk * pos_len * H * sizeof(float)); // [dk, P, H]
 
         ggml_tensor* bu = clone_weight(ctx, ml, ap + "pos_bias_u");
         ggml_tensor* bv = clone_weight(ctx, ml, ap + "pos_bias_v");
@@ -340,6 +343,37 @@ static ggml_tensor* build_stream_layer(
     return r; // [D, Tc]
 }
 
+// The relative-position projection linear_pos(pos_emb) depends only on the attention
+// window length (cache_len + Tc), never on the audio, yet the streaming forward used to
+// recompute it in every layer for every chunk: on a 24-layer, D=1024 model with ~145
+// positions that is ~7 GFLOP per chunk. It is computed once per window length for the
+// session, already split into heads, and fed to the chunk graph as an input. Same ggml
+// ops on the same inputs, so the outputs are unchanged.
+const std::vector<std::vector<float>>& StreamingEncoder::pos_heads(int window) {
+    auto it = pos_heads_.find(window);
+    if (it != pos_heads_.end()) return it->second;
+
+    const int D = d_model_, H = n_heads_, dk = D / H;
+    const int pos_len = 2 * window - 1;
+    std::vector<float> pe_host;
+    rel_pos_encoding(window, D, pe_host); // row-major [pos_len, D]
+    std::vector<std::vector<float>> heads(n_layers_);
+    for (int i = 0; i < n_layers_; ++i) {
+        const std::string w = "encoder.layers." + std::to_string(i) + ".self_attn.linear_pos.weight";
+        bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/0, [&](ggml_context* ctx) -> ggml_tensor* {
+            int64_t pe_ne[2] = {D, pos_len};
+            ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne, pe_host.data(),
+                                                     pe_host.size() * sizeof(float));
+            ggml_tensor* p = ggml_mul_mat(ctx, clone_weight(ctx, ml_, w), pe); // [D, P]
+            p = ggml_reshape_3d(ctx, p, dk, H, pos_len);
+            return ggml_cont(ctx, ggml_permute(ctx, p, 0, 2, 1, 3));         // [dk, P, H]
+        }, heads[i]);
+        if (!ok || (int)heads[i].size() != dk * pos_len * H)
+            throw std::runtime_error("streaming encoder: linear_pos precompute failed");
+    }
+    return pos_heads_.emplace(window, std::move(heads)).first->second;
+}
+
 std::vector<float> StreamingEncoder::step(const std::vector<float>& mel_chunk_frames,
                                           int n_mel_frames, bool is_last,
                                           int& n_valid_out) {
@@ -378,6 +412,7 @@ std::vector<float> StreamingEncoder::step(const std::vector<float>& mel_chunk_fr
     const int clc_len_now = clc_len_;   // snapshot for the in-graph mask
 
     const std::string cnt = ml_.config().conv_norm_type;
+    const std::vector<std::vector<float>>& pos_heads = this->pos_heads(Pn);
     StreamLayerCapture caps[/*max layers*/64];
     assert(n_layers_ <= 64);
 
@@ -414,13 +449,6 @@ std::vector<float> StreamingEncoder::step(const std::vector<float>& mel_chunk_fr
                 x = ggml_scale(ctx, x, std::sqrt((float)D));
             }
 
-            // pos_emb [D, pos_len] fed as a graph input (host-computed once).
-            std::vector<float>& pe_host = pool.alloc_f32();
-            rel_pos_encoding(Pn, D, pe_host); // row-major [pos_len, D]
-            int64_t pe_ne[2] = {D, pos_len};
-            ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
-                                  pe_host.data(), pe_host.size() * sizeof(float));
-
             // Conformer layer stack (conv + attention caches threaded in-graph).
             const int LP = left_pad_;
             for (int i = 0; i < n_layers_; ++i) {
@@ -435,7 +463,7 @@ std::vector<float> StreamingEncoder::step(const std::vector<float>& mel_chunk_fr
                         GGML_TYPE_F32, 2, t_ne, cache_time_[i].data(),
                         cache_time_[i].size() * sizeof(float));
 
-                x = build_stream_layer(ctx, ml_, i, x, Tc, pe, pos_len,
+                x = build_stream_layer(ctx, ml_, i, x, Tc, pos_heads[i].data(), pos_len,
                                        cache_len, clc_len_now, n_heads_, D,
                                        conv_kernel_, att_left_, att_right_, cnt,
                                        cache_ch_t, cache_t_t, pool, &caps[i]);

src/streaming_encoder.hpp

@@ -1,5 +1,6 @@
 #pragma once
 #include "model_loader.hpp"
+#include <map>
 #include <vector>
 namespace pk {
 
@@ -91,6 +92,10 @@ class StreamingEncoder {
     std::vector<std::vector<float>> cache_time_;
     // cache_last_channel[layer]: row-major [last_channel_cache, d_model]
     std::vector<std::vector<float>> cache_channel_;
+    // linear_pos(pos_emb) split into heads, [layer] -> [dk, pos_len, H], per window
+    // length Tc + cache_len (see pos_heads()).
+    std::map<int, std::vector<std::vector<float>>> pos_heads_;
+    const std::vector<std::vector<float>>& pos_heads(int window);
 };
 
 } // namespace pk