Source code
Revision control
Copy as Markdown
Other Tools
#include "relpos_attention.hpp"
#include "ggml_graph.hpp"
#include "backend.hpp"
#include "ggml.h"
#include <cassert>
#include <cmath>
#include <cstring>
#include <string>
#include <vector>
namespace pk {
// Weights from the loader are referenced DIRECTLY as graph leaves via the shared
// pk::clone_weight (backend.cpp; zero-copy via the loader's CPU backend buffer).
// Allowlisted attention linears (linear_q/k/v/out/pos.weight) may be f16/q8_0
// and are fed into ggml_mul_mat, which dequantizes src0 on the fly. pos_bias_u/v
// and every other weight stay F32. A std::string-name overload keeps the call
// sites (pre + suffix) unchanged.
static ggml_tensor* clone_weight(ggml_context* ctx, const ModelLoader& ml,
const std::string& name) {
return pk::clone_weight(ctx, ml, name.c_str());
}
RelPosAttention::RelPosAttention(const ModelLoader& ml, int layer_idx)
: ml_(ml), layer_idx_(layer_idx) {
d_model_ = (int)ml.config().d_model;
n_heads_ = (int)ml.config().n_heads;
assert(n_heads_ > 0 && d_model_ % n_heads_ == 0);
d_head_ = d_model_ / n_heads_;
// Chunked-limited attention (NeMo att_context_style=="chunked_limited",
// e.g. parakeet_realtime_eou_120m-v1 with att_context_size=[70,1]). The
// offline forward applies the SAME additive -inf mask NeMo builds in
// ConformerEncoder._create_masks. Offline models use "regular" (full context).
chunked_limited_ = (ml.config().att_context_style == "chunked_limited" &&
ml.config().att_context_right >= 0);
att_left_ = ml.config().att_context_left;
att_right_ = ml.config().att_context_right;
}
ggml_tensor* RelPosAttention::build_graph(ggml_context* ctx, ggml_tensor* xt,
int T, ggml_tensor* pe, int pos_len,
int valid_len,
GraphInputPool& pool,
int att_left, int att_right) const {
// Scalar (B=1) builder: the verbatim v1 2-D/3D relative-position attention
// graph. The single-clip conformer layer routes here so B=1 runs the lean
// graph and is bit-exact with v1. build_graph_batched below serves B>1.
const int D = d_model_;
const int H = n_heads_;
const int dk = d_head_;
const float scale = 1.0f / std::sqrt((float)dk);
assert(pos_len == 2 * T - 1);
const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn.";
const ModelLoader& ml = ml_;
// ---- linear projections (nn.Linear: ggml W ne=[in,out]) ----
// The bias is added only when requested AND present: NeMo configures the
// attention linears with bias=False in some checkpoints
// (parakeet-tdt-0.6b-v2/-v3) and bias=True in others (110m).
auto linear = [&](const char* w, const char* b, ggml_tensor* in) {
ggml_tensor* W = clone_weight(ctx, ml, pre + w);
ggml_tensor* y = ggml_mul_mat(ctx, W, in); // [out, *]
if (b && ml.tensor(pre + b)) {
ggml_tensor* B = clone_weight(ctx, ml, pre + b);
y = ggml_add(ctx, y, B); // broadcast [out] over cols
}
return y;
};
ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); // [D, T]
ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt); // [D, T]
ggml_tensor* v = linear("linear_v.weight", "linear_v.bias", xt); // [D, T]
ggml_tensor* p = linear("linear_pos.weight", nullptr, pe); // [D, P]
// ---- split into heads: [D, *] -> [dk, H, *] -> [dk, *, H] ----
auto to_heads = [&](ggml_tensor* t, int n) {
t = ggml_reshape_3d(ctx, t, dk, H, n); // [dk, H, n]
t = ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H]
return t;
};
ggml_tensor* qh = to_heads(q, T); // [dk, T, H]
ggml_tensor* kh = to_heads(k, T); // [dk, T, H]
ggml_tensor* vh = to_heads(v, T); // [dk, T, H]
ggml_tensor* ph = to_heads(p, pos_len); // [dk, P, H]
// ---- pos_bias_u/v: ne [dk, H] -> [dk, 1, H] to broadcast over T ----
ggml_tensor* bu = clone_weight(ctx, ml, pre + "pos_bias_u"); // [dk, H]
ggml_tensor* bv = clone_weight(ctx, ml, pre + "pos_bias_v"); // [dk, H]
bu = ggml_reshape_3d(ctx, bu, dk, 1, H);
bv = ggml_reshape_3d(ctx, bv, dk, 1, H);
ggml_tensor* qu = ggml_add(ctx, qh, bu); // [dk, T, H]
ggml_tensor* qv = ggml_add(ctx, qh, bv); // [dk, T, H]
// ---- ac = q_u @ k^T : ggml_mul_mat([dk,T,H],[dk,T,H]) -> [T_k, T_q, H] ----
ggml_tensor* ac = ggml_mul_mat(ctx, kh, qu); // [T(key), T(query), H]
// ---- bd = q_v @ p^T -> [P(pos), T(query), H], then rel_shift -> [T,T,H] ----
ggml_tensor* bd = ggml_mul_mat(ctx, ph, qv); // [P, T, H]
bd = ggml_pad_ext(ctx, bd, /*lp0*/1, /*rp0*/0, 0,0, 0,0, 0,0); // [P+1=2T, T, H]
bd = ggml_reshape_3d(ctx, bd, T, 2 * T, H); // [T, 2T, H]
bd = ggml_view_3d(ctx, bd, T, 2 * T - 1, H,
bd->nb[1], bd->nb[2], bd->nb[1]); // [T, 2T-1, H]
bd = ggml_cont(ctx, bd);
bd = ggml_reshape_3d(ctx, bd, 2 * T - 1, T, H); // [2T-1, T, H]
bd = ggml_view_3d(ctx, bd, T, T, H, bd->nb[1], bd->nb[2], 0);
bd = ggml_cont(ctx, bd);
// ---- scores = ac + bd ; softmax(scores*scale + mask) ----
ggml_tensor* scores = ggml_add(ctx, ac, bd); // [T_k, T_q, H]
// Additive mask [T_k, T_q]: 0 where query qi may attend to key kj, -inf
// otherwise. (1) pad mask: key kj valid iff kj < valid_len. (2) chunked-
// limited window for streaming models. See header / NeMo _create_masks.
const int chunk_size = chunked_limited_ ? (att_right_ + 1) : 0;
const int left_chunks = (chunked_limited_ && chunk_size > 0)
? (att_left_ / chunk_size) : 0;
std::vector<float>& mask_host = pool.alloc_f32((size_t)T * T);
{
float* md = mask_host.data();
const float ninf = -INFINITY;
for (int qi = 0; qi < T; ++qi) {
const int cq = chunked_limited_ ? (qi / chunk_size) : 0;
for (int kj = 0; kj < T; ++kj) {
bool ok = (kj < valid_len);
if (ok && chunked_limited_) {
const int ck = kj / chunk_size;
const int diff = cq - ck;
ok = (diff >= 0 && diff <= left_chunks);
}
// Symmetric sliding window (NeMo rel_pos_local_attn): keep only
// keys within [qi-att_left, qi+att_right].
if (ok && att_left >= 0) {
const int rel = qi - kj;
ok = (rel <= att_left) && (rel >= -att_right);
}
md[(size_t)qi * T + kj] = ok ? 0.0f : ninf;
}
}
}
int64_t mask_ne[2] = {T, T};
ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, mask_ne,
mask_host.data(), mask_host.size() * sizeof(float));
ggml_tensor* attn = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); // [T_k, T_q, H]
// ---- context = attn @ v -> [dk, T_q, H] ----
ggml_tensor* vtk = ggml_cont(ctx, ggml_permute(ctx, vh, 1, 0, 2, 3)); // [T_k, dk, H]
ggml_tensor* ctxh = ggml_mul_mat(ctx, vtk, attn); // [dk, T_q, H]
// ---- concat heads: [dk, T, H] -> [dk, H, T] -> [D, T] ----
ggml_tensor* merged = ggml_cont(ctx, ggml_permute(ctx, ctxh, 0, 2, 1, 3)); // [dk, H, T]
merged = ggml_reshape_2d(ctx, merged, D, T); // [D, T]
// Zero the context for PADDED query rows (NeMo masks padded query rows fully
// -> output reduces to linear_out.bias). Apply a query-row mask [1, T].
if (valid_len < T) {
std::vector<float>& qmask_host = pool.alloc_f32(T);
for (int qi = 0; qi < T; ++qi)
qmask_host[qi] = (qi < valid_len) ? 1.0f : 0.0f;
int64_t qm_ne[2] = {1, T};
ggml_tensor* qmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, qm_ne,
qmask_host.data(), qmask_host.size() * sizeof(float));
merged = ggml_mul(ctx, merged, qmask); // broadcast over D
}
// ---- output projection ----
return linear("linear_out.weight", "linear_out.bias", merged); // [D, T]
}
ggml_tensor* RelPosAttention::build_graph_batched(
ggml_context* ctx, ggml_tensor* xt, int T, int B, ggml_tensor* pe,
int pos_len, const std::vector<int>& valid_len,
GraphInputPool& pool) const {
const int D = d_model_;
const int H = n_heads_;
const int dk = d_head_;
const float scale = 1.0f / std::sqrt((float)dk);
assert(pos_len == 2 * T - 1);
assert((int)valid_len.size() == B);
const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn.";
const ModelLoader& ml = ml_;
// ---- linear projections (nn.Linear: ggml W ne=[in,out]) ----
// The bias is added only when requested AND present: NeMo configures the
// attention linears with bias=False in some checkpoints
// (parakeet-tdt-0.6b-v2/-v3) and bias=True in others (110m).
auto linear = [&](const char* w, const char* b, ggml_tensor* in) {
ggml_tensor* W = clone_weight(ctx, ml, pre + w);
ggml_tensor* y = ggml_mul_mat(ctx, W, in); // [out, *]
if (b && ml.tensor(pre + b)) {
ggml_tensor* B = clone_weight(ctx, ml, pre + b);
y = ggml_add(ctx, y, B); // broadcast [out] over cols
}
return y;
};
// xt is [D, T, B]; mul_mat batches over ne2 -> q/k/v are [D, T, B]. pe is
// shared [D, P] (NO batch) -> p is [D, P].
ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); // [D, T, B]
ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt); // [D, T, B]
ggml_tensor* v = linear("linear_v.weight", "linear_v.bias", xt); // [D, T, B]
ggml_tensor* p = linear("linear_pos.weight", nullptr, pe); // [D, P]
// ---- split into heads (batched): [D, n, B] -> [dk, H, n, B] -> [dk, n, H, B] ----
auto to_heads_b = [&](ggml_tensor* t, int n) {
t = ggml_reshape_4d(ctx, t, dk, H, n, B); // [dk, H, n, B]
t = ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H, B]
return t;
};
// p is shared (no batch) -> keep the 3D head-split: [dk, P, H].
auto to_heads = [&](ggml_tensor* t, int n) {
t = ggml_reshape_3d(ctx, t, dk, H, n); // [dk, H, n]
t = ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H]
return t;
};
ggml_tensor* qh = to_heads_b(q, T); // [dk, T, H, B]
ggml_tensor* kh = to_heads_b(k, T); // [dk, T, H, B]
ggml_tensor* vh = to_heads_b(v, T); // [dk, T, H, B]
ggml_tensor* ph = to_heads(p, pos_len); // [dk, P, H] (ne3=1, broadcast over B)
// ---- pos_bias_u/v: ne [dk, H] -> [dk, 1, H, 1] to broadcast over T and B ----
ggml_tensor* bu = clone_weight(ctx, ml, pre + "pos_bias_u"); // [dk, H]
ggml_tensor* bv = clone_weight(ctx, ml, pre + "pos_bias_v"); // [dk, H]
bu = ggml_reshape_4d(ctx, bu, dk, 1, H, 1);
bv = ggml_reshape_4d(ctx, bv, dk, 1, H, 1);
ggml_tensor* qu = ggml_add(ctx, qh, bu); // [dk, T, H, B]
ggml_tensor* qv = ggml_add(ctx, qh, bv); // [dk, T, H, B]
// ---- ac = q_u @ k^T : mul_mat([dk,T,H,B],[dk,T,H,B]) -> [T_k, T_q, H, B] ----
ggml_tensor* ac = ggml_mul_mat(ctx, kh, qu); // [T(key), T(query), H, B]
// ---- bd = q_v @ p^T -> [P, T_q, H, B], then rel_shift -> [T, T, H, B] ----
// ph is [dk, P, H] (ne3=1) and broadcasts over the batch (ne3=B) of qv.
ggml_tensor* bd = ggml_mul_mat(ctx, ph, qv); // [P, T, H, B]
// 4D rel-shift: identical ne0/ne1 stride+offset arithmetic as the 3D path,
// with the batch axis threaded through every reshape/view via ne3=B and
// nb[3]. ne2/ne3 are passive (no offset on them).
bd = ggml_pad_ext(ctx, bd, /*lp0*/1, /*rp0*/0, 0,0, 0,0, 0,0); // [2T, T, H, B]
bd = ggml_reshape_4d(ctx, bd, T, 2 * T, H, B); // [T, 2T, H, B]
bd = ggml_view_4d(ctx, bd, T, 2 * T - 1, H, B,
bd->nb[1], bd->nb[2], bd->nb[3], bd->nb[1]); // [T, 2T-1, H, B]
bd = ggml_cont(ctx, bd);
bd = ggml_reshape_4d(ctx, bd, 2 * T - 1, T, H, B); // [2T-1, T, H, B]
bd = ggml_view_4d(ctx, bd, T, T, H, B,
bd->nb[1], bd->nb[2], bd->nb[3], 0); // [T, T, H, B]
bd = ggml_cont(ctx, bd);
// ---- scores = ac + bd ; softmax(scores*scale + mask) ----
ggml_tensor* scores = ggml_add(ctx, ac, bd); // [T_k, T_q, H, B]
// Per-item additive mask [T_k, T_q, 1, B]: 0 where query qi may attend to key
// kj, -inf otherwise. (1) pad mask: key kj valid iff kj < valid_len[b].
// (2) chunked-limited window for streaming models. See NeMo _create_masks.
// Mask shape is [T, T, 1, B] (ne2=1) so soft_max_ext broadcasts it over the
// head axis (ne2=H) while indexing per item on ne3=B. Verified against the
// ggml CPU kernel ggml_compute_forward_soft_max_f32: it reads the mask at
// i12 = i02 % ne12 (head, here ne12=1 -> always 0) and i13 = i03 % ne13
// (batch, here ne13=B -> exact per-item), and ggml_soft_max_impl asserts
// a->ne[2] % mask->ne[2] == 0 and a->ne[3] % mask->ne[3] == 0.
const int chunk_size = chunked_limited_ ? (att_right_ + 1) : 0;
const int left_chunks = (chunked_limited_ && chunk_size > 0)
? (att_left_ / chunk_size) : 0;
std::vector<float>& mask_host = pool.alloc_f32((size_t)B * T * T);
{
float* md = mask_host.data();
const float ninf = -INFINITY;
for (int b = 0; b < B; ++b) {
const int vl = valid_len[b];
for (int qi = 0; qi < T; ++qi) {
const int cq = chunked_limited_ ? (qi / chunk_size) : 0;
for (int kj = 0; kj < T; ++kj) {
bool ok = (kj < vl);
if (ok && chunked_limited_) {
const int ck = kj / chunk_size;
const int diff = cq - ck;
ok = (diff >= 0 && diff <= left_chunks);
}
md[(size_t)b * T * T + (size_t)qi * T + kj] = ok ? 0.0f : ninf;
}
}
}
}
int64_t mask_ne[4] = {T, T, 1, B};
ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 4, mask_ne,
mask_host.data(), mask_host.size() * sizeof(float));
ggml_tensor* attn = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); // [T_k, T_q, H, B]
// ---- context = attn @ v -> [dk, T_q, H, B] ----
ggml_tensor* vtk = ggml_cont(ctx, ggml_permute(ctx, vh, 1, 0, 2, 3)); // [T_k, dk, H, B]
ggml_tensor* ctxh = ggml_mul_mat(ctx, vtk, attn); // [dk, T_q, H, B]
// ---- concat heads: [dk, T, H, B] -> [dk, H, T, B] -> [D, T, B] ----
ggml_tensor* merged = ggml_cont(ctx, ggml_permute(ctx, ctxh, 0, 2, 1, 3)); // [dk, H, T, B]
merged = ggml_reshape_3d(ctx, merged, D, T, B); // [D, T, B]
// Zero the context for PADDED query rows (NeMo masks padded query rows fully
// -> output reduces to linear_out.bias). Apply a per-item query-row mask
// [1, T, B] (broadcast over D). Emit only when some item has valid_len < T.
bool any_pad = false;
for (int b = 0; b < B; ++b) any_pad = any_pad || (valid_len[b] < T);
if (any_pad) {
std::vector<float>& qmask_host = pool.alloc_f32((size_t)B * T);
for (int b = 0; b < B; ++b)
for (int qi = 0; qi < T; ++qi)
qmask_host[(size_t)b * T + qi] = (qi < valid_len[b]) ? 1.0f : 0.0f;
int64_t qm_ne[3] = {1, T, B};
ggml_tensor* qmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, qm_ne,
qmask_host.data(), qmask_host.size() * sizeof(float));
merged = ggml_mul(ctx, merged, qmask); // broadcast over D
}
// ---- output projection ----
return linear("linear_out.weight", "linear_out.bias", merged); // [D, T, B]
}
ggml_tensor* RelPosAttention::build_graph_batched_local(
ggml_context* ctx, ggml_tensor* xt, int T, int B, ggml_tensor* pe,
int pos_len, const std::vector<int>& valid_len,
int att_left, int att_right, GraphInputPool& pool) const {
const int D = d_model_, H = n_heads_, dk = d_head_;
const int P = pos_len;
const float scale = 1.0f / std::sqrt((float)dk);
assert(att_left >= 0 && att_right >= 0);
assert(P == att_left + att_right + 1);
assert((int)valid_len.size() == B);
const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn.";
const ModelLoader& ml = ml_;
auto linear = [&](const char* wn, const char* bn, ggml_tensor* in) {
ggml_tensor* W = clone_weight(ctx, ml, pre + wn);
ggml_tensor* y = ggml_mul_mat(ctx, W, in);
if (bn && ml.tensor(pre + bn)) y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + bn));
return y;
};
ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt); // [D, T, B]
ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt);
ggml_tensor* v = linear("linear_v.weight", "linear_v.bias", xt);
ggml_tensor* p = linear("linear_pos.weight", nullptr, pe); // [D, P]
auto to_heads_b = [&](ggml_tensor* t, int n) {
t = ggml_reshape_4d(ctx, t, dk, H, n, B);
return ggml_cont(ctx, ggml_permute(ctx, t, 0, 2, 1, 3)); // [dk, n, H, B]
};
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_b(q, T), *kh = to_heads_b(k, T);
ggml_tensor* vh = to_heads_b(v, T), *php = to_heads(p, P); // ph shared (ne3=1)
ggml_tensor* bu = ggml_reshape_4d(ctx, clone_weight(ctx, ml, pre + "pos_bias_u"), dk, 1, H, 1);
ggml_tensor* bv = ggml_reshape_4d(ctx, clone_weight(ctx, ml, pre + "pos_bias_v"), dk, 1, H, 1);
ggml_tensor* qu = ggml_add(ctx, qh, bu); // [dk, T, H, B]
ggml_tensor* qv = ggml_add(ctx, qh, bv); // [dk, T, H, B]
// Pad K/V along time (ne1); view offset c -> key (t - att_left + c).
ggml_tensor* kpad = ggml_pad_ext(ctx, kh, 0,0, att_left,att_right, 0,0, 0,0); // [dk, T+P-1, H, B]
ggml_tensor* vpad = ggml_pad_ext(ctx, vh, 0,0, att_left,att_right, 0,0, 0,0);
// Banded content scores ac[c, t, H, B]; stack on ne0=c.
ggml_tensor* ac = nullptr;
for (int c = 0; c < P; ++c) {
ggml_tensor* kc = ggml_view_4d(ctx, kpad, dk, T, H, B,
kpad->nb[1], kpad->nb[2], kpad->nb[3], (size_t)c * kpad->nb[1]);
ggml_tensor* acc = ggml_sum_rows(ctx, ggml_mul(ctx, qu, kc)); // [1, T, H, B]
ac = ac ? ggml_concat(ctx, ac, acc, 0) : acc;
}
ggml_tensor* bd = ggml_mul_mat(ctx, php, qv); // [P, T, H, B] (php broadcasts over B)
ggml_tensor* scores = ggml_add(ctx, ac, bd); // [P, T, H, B]
// Per-item band mask [P, T, 1, B].
std::vector<float>& mh = pool.alloc_f32((size_t)B * T * P);
for (int b = 0; b < B; ++b) {
const int vl = valid_len[b];
for (int t = 0; t < T; ++t)
for (int c = 0; c < P; ++c) {
const int key = t - att_left + c;
mh[(size_t)b * T * P + (size_t)t * P + c] = (key >= 0 && key < vl) ? 0.0f : -INFINITY;
}
}
int64_t mne[4] = {P, T, 1, B};
ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 4, mne,
mh.data(), mh.size() * sizeof(float));
ggml_tensor* prob = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); // softmax over c
// context[dk, t, H, B] = sum_c prob[c, t] * v[t-att_left+c].
ggml_tensor* context = nullptr;
for (int c = 0; c < P; ++c) {
ggml_tensor* vc = ggml_view_4d(ctx, vpad, dk, T, H, B,
vpad->nb[1], vpad->nb[2], vpad->nb[3], (size_t)c * vpad->nb[1]);
ggml_tensor* pc = ggml_view_4d(ctx, prob, 1, T, H, B,
prob->nb[1], prob->nb[2], prob->nb[3], (size_t)c * prob->nb[0]); // [1,T,H,B]
ggml_tensor* term = ggml_mul(ctx, vc, pc);
context = context ? ggml_add(ctx, context, term) : term;
}
// Merge heads [dk, T, H, B] -> [dk, H, T, B] -> [D, T, B].
ggml_tensor* merged = ggml_cont(ctx, ggml_permute(ctx, context, 0, 2, 1, 3));
merged = ggml_reshape_3d(ctx, merged, D, T, B);
bool any_pad = false;
for (int b = 0; b < B; ++b) any_pad = any_pad || (valid_len[b] < T);
if (any_pad) {
std::vector<float>& qm = pool.alloc_f32((size_t)B * T);
for (int b = 0; b < B; ++b)
for (int t = 0; t < T; ++t) qm[(size_t)b * T + t] = (t < valid_len[b]) ? 1.0f : 0.0f;
int64_t qne[3] = {1, T, B};
ggml_tensor* qmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, qne,
qm.data(), qm.size() * sizeof(float));
merged = ggml_mul(ctx, merged, qmask);
}
return linear("linear_out.weight", "linear_out.bias", merged); // [D, T, B]
}
void RelPosAttention::forward(const std::vector<float>& x, int T,
const std::vector<float>& pos_emb, int pos_len,
int valid_len,
std::vector<float>& out) const {
const int D = d_model_;
assert((int)x.size() == T * D);
assert((int)pos_emb.size() == pos_len * D);
assert(pos_len == 2 * T - 1);
// Thin wrapper over the graph-builder: build JUST the attention sub-graph
// with x/pos_emb fed as inputs and compute it on the persistent Backend.
// Used by the unit test (the fused conformer layer uses build_graph).
GraphInputPool pool;
bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/4,
[&](ggml_context* ctx) -> ggml_tensor* {
int64_t xt_ne[2] = {D, T};
ggml_tensor* xt = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, xt_ne,
x.data(), (size_t)T * D * sizeof(float));
int64_t pe_ne[2] = {D, pos_len};
ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
pos_emb.data(), (size_t)pos_len * D * sizeof(float));
return build_graph(ctx, xt, T, pe, pos_len, valid_len, pool);
}, out);
assert(ok && "relpos attention graph failed");
(void)ok;
}
ggml_tensor* RelPosAttention::build_graph_local(ggml_context* ctx, ggml_tensor* xt,
int T, ggml_tensor* pe, int pos_len,
int valid_len, int att_left, int att_right,
GraphInputPool& pool) const {
const int D = d_model_, H = n_heads_, dk = d_head_;
const int P = pos_len; // window width = att_left+att_right+1
const float scale = 1.0f / std::sqrt((float)dk);
assert(att_left >= 0 && att_right >= 0);
assert(P == att_left + att_right + 1);
// Exact NeMo rel_pos_local_attn (RelPositionMultiHeadAttentionLongformer),
// computed in O(T*window) via pad-and-shift instead of NeMo's skew/chunk
// tricks. For query t and window column c in [0, P), the key is
// (t - att_left + c). NeMo's local pos is ordered index0 = +att_left .. last
// = -att_right, so column c uses pos row c directly (matrix_bd = q_v . p^T,
// added 1:1 to the banded content scores).
{
const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn.";
const ModelLoader& ml = ml_;
auto linear = [&](const char* wn, const char* bn, ggml_tensor* in) {
ggml_tensor* W = clone_weight(ctx, ml, pre + wn);
ggml_tensor* y = ggml_mul_mat(ctx, W, in);
if (bn && ml.tensor(pre + bn)) y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + bn));
return y;
};
ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt);
ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt);
ggml_tensor* v = linear("linear_v.weight", "linear_v.bias", xt);
ggml_tensor* p = linear("linear_pos.weight", nullptr, pe);
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, T), *kh = to_heads(k, T);
ggml_tensor* vh = to_heads(v, T), *php = to_heads(p, P);
ggml_tensor* bu = ggml_reshape_3d(ctx, clone_weight(ctx, ml, pre + "pos_bias_u"), dk, 1, H);
ggml_tensor* bv = ggml_reshape_3d(ctx, clone_weight(ctx, ml, pre + "pos_bias_v"), dk, 1, H);
ggml_tensor* qu = ggml_add(ctx, qh, bu); // [dk, T, H]
ggml_tensor* qv = ggml_add(ctx, qh, bv); // [dk, T, H]
// Pad K/V along time (ne1): att_left on the left, att_right on the right,
// so view offset c yields key (t - att_left + c).
ggml_tensor* kpad = ggml_pad_ext(ctx, kh, 0,0, att_left,att_right, 0,0, 0,0); // [dk, T+P-1, H]
ggml_tensor* vpad = ggml_pad_ext(ctx, vh, 0,0, att_left,att_right, 0,0, 0,0);
// Banded content scores ac[c, t, H] = q_u[t] . k[t-att_left+c]; stack on ne0=c.
ggml_tensor* ac = nullptr;
for (int c = 0; c < P; ++c) {
ggml_tensor* kc = ggml_view_3d(ctx, kpad, dk, T, H, kpad->nb[1], kpad->nb[2],
(size_t)c * kpad->nb[1]);
ggml_tensor* acc = ggml_sum_rows(ctx, ggml_mul(ctx, qu, kc)); // [1, T, H]
ac = ac ? ggml_concat(ctx, ac, acc, 0) : acc;
}
// Positional scores bd[c, t, H] = q_v[t] . p[c] (direct, no rel-shift).
ggml_tensor* bd = ggml_mul_mat(ctx, php, qv); // [P, T, H]
ggml_tensor* scores = ggml_add(ctx, ac, bd); // [P, T, H]
// Band mask [P, T]: 0 if key in [0, valid_len), else -inf (covers the
// out-of-sequence window corners and pad frames).
std::vector<float>& mh = pool.alloc_f32((size_t)P * T);
for (int t = 0; t < T; ++t)
for (int c = 0; c < P; ++c) {
const int key = t - att_left + c;
mh[(size_t)t * P + c] = (key >= 0 && key < valid_len) ? 0.0f : -INFINITY;
}
int64_t mne[2] = {P, T};
ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, mne,
mh.data(), mh.size() * sizeof(float));
ggml_tensor* prob = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); // softmax over c
// context[dk, t, H] = sum_c prob[c, t] * v[t-att_left+c].
ggml_tensor* context = nullptr;
for (int c = 0; c < P; ++c) {
ggml_tensor* vc = ggml_view_3d(ctx, vpad, dk, T, H, vpad->nb[1], vpad->nb[2],
(size_t)c * vpad->nb[1]);
ggml_tensor* pc = ggml_view_3d(ctx, prob, 1, T, H, prob->nb[1], prob->nb[2],
(size_t)c * prob->nb[0]); // [1, T, H]
ggml_tensor* term = ggml_mul(ctx, vc, pc); // broadcast pc over dk
context = context ? ggml_add(ctx, context, term) : term;
}
// Merge heads [dk, T, H] -> [dk, H, T] -> [D, T].
ggml_tensor* merged = ggml_cont(ctx, ggml_permute(ctx, context, 0, 2, 1, 3));
merged = ggml_reshape_2d(ctx, merged, D, T);
if (valid_len < T) { // zero padded query rows -> output = linear_out.bias
std::vector<float>& qm = pool.alloc_f32(T);
for (int t = 0; t < T; ++t) qm[t] = (t < valid_len) ? 1.0f : 0.0f;
int64_t qne[2] = {1, T};
ggml_tensor* qmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, qne,
qm.data(), qm.size() * sizeof(float));
merged = ggml_mul(ctx, merged, qmask);
}
ggml_tensor* Wo = clone_weight(ctx, ml, pre + "linear_out.weight");
ggml_tensor* y = ggml_mul_mat(ctx, Wo, merged);
if (ml.tensor(pre + "linear_out.bias"))
y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + "linear_out.bias"));
return y; // [D, T]
}
}
void RelPosAttention::forward_local(const std::vector<float>& x, int T,
const std::vector<float>& pos_emb, int pos_len,
int valid_len, int att_left, int att_right,
std::vector<float>& out) const {
const int D = d_model_;
assert((int)x.size() == T * D);
assert((int)pos_emb.size() == pos_len * D);
// Thin wrapper over build_graph_local: feed x/pos_emb as graph inputs and
// compute the banded attention sub-graph on the persistent Backend. The
// fused conformer encoder calls build_graph_local directly.
GraphInputPool pool;
bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/4,
[&](ggml_context* ctx) -> ggml_tensor* {
int64_t xt_ne[2] = {D, T};
ggml_tensor* xt = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, xt_ne,
x.data(), (size_t)T * D * sizeof(float));
int64_t pe_ne[2] = {D, pos_len};
ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
pos_emb.data(), (size_t)pos_len * D * sizeof(float));
return build_graph_local(ctx, xt, T, pe, pos_len, valid_len,
att_left, att_right, pool);
}, out);
assert(ok && "relpos local attention graph failed");
(void)ok;
}
ggml_tensor* RelPosAttention::build_graph_local_chunked(
ggml_context* ctx, ggml_tensor* xt, int T, ggml_tensor* pe, int pos_len,
int valid_len, int att_left, int att_right, GraphInputPool& pool,
int chunk) const {
const int D = d_model_, H = n_heads_, dk = d_head_;
const int P = pos_len; // window width = att_left+att_right+1
const float scale = 1.0f / std::sqrt((float)dk);
assert(att_left >= 0 && att_right >= 0);
assert(P == att_left + att_right + 1);
// Tile time into chunks of C frames (G chunks, Tp = G*C padded length). Each
// chunk carries its own C+P-1 keys/values (the P-1 halo overlaps the next
// chunk), so a query in chunk g only attends within g. Default C spans the
// window so the halo is one chunk wide.
int C = chunk > 0 ? chunk : (att_left + att_right);
if (C < 1) C = 1;
const int G = (T + C - 1) / C;
const int Tp = G * C;
const int Lk = (C + P - 1) * G; // dense length the chunk VIEW needs
const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".self_attn.";
const ModelLoader& ml = ml_;
auto linear = [&](const char* wn, const char* bn, ggml_tensor* in) {
ggml_tensor* W = clone_weight(ctx, ml, pre + wn);
ggml_tensor* y = ggml_mul_mat(ctx, W, in);
if (bn && ml.tensor(pre + bn)) y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + bn));
return y;
};
ggml_tensor* q = linear("linear_q.weight", "linear_q.bias", xt);
ggml_tensor* k = linear("linear_k.weight", "linear_k.bias", xt);
ggml_tensor* v = linear("linear_v.weight", "linear_v.bias", xt);
ggml_tensor* p = linear("linear_pos.weight", nullptr, pe);
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, T), *kh = to_heads(k, T);
ggml_tensor* vh = to_heads(v, T), *php = to_heads(p, P);
ggml_tensor* bu = ggml_reshape_3d(ctx, clone_weight(ctx, ml, pre + "pos_bias_u"), dk, 1, H);
ggml_tensor* bv = ggml_reshape_3d(ctx, clone_weight(ctx, ml, pre + "pos_bias_v"), dk, 1, H);
ggml_tensor* qu = ggml_add(ctx, qh, bu); // [dk, T, H]
ggml_tensor* qv = ggml_add(ctx, qh, bv); // [dk, T, H]
// ---- Content scores ac[c,t,H] via chunked matmul + diagonal skew-view ----
// Pad queries to Tp and reshape into non-overlapping chunks [dk, C, G, H].
ggml_tensor* qu_p = (Tp > T) ? ggml_pad_ext(ctx, qu, 0,0, 0,Tp-T, 0,0, 0,0) : qu;
ggml_tensor* qu_c = ggml_reshape_4d(ctx, qu_p, dk, C, G, H);
// Pad keys (left att_left, right att_right) then OVER-pad to Lk so the
// overlapping chunk view's dense ne-product fits ggml's bounds check.
ggml_tensor* kpad = ggml_pad_ext(ctx, kh, 0,0, att_left,att_right, 0,0, 0,0); // [dk,T+P-1,H]
if (Lk > (int)kpad->ne[1]) kpad = ggml_pad_ext(ctx, kpad, 0,0, 0,Lk-(int)kpad->ne[1], 0,0, 0,0);
// Overlapping key chunks [dk, C+P-1, G, H]: chunk g advances C along time.
ggml_tensor* kchunk = ggml_view_4d(ctx, kpad, dk, C+P-1, G, H,
kpad->nb[1], (size_t)C*kpad->nb[1], kpad->nb[2], 0);
kchunk = ggml_cont(ctx, kchunk);
// Per-chunk q.k block [C+P-1, C, G, H]: sc[j,i,g] = k[gC+j] . qu[gC+i].
ggml_tensor* sc = ggml_mul_mat(ctx, kchunk, qu_c);
// Diagonal skew: ac_band[c,i,g] = sc[i+c, i, g] -> [P, C, G, H], nb1 walks (C+P).
ggml_tensor* acb = ggml_view_4d(ctx, sc, P, C, G, H,
(size_t)(C+P)*sc->nb[0], sc->nb[2], sc->nb[3], 0);
acb = ggml_cont(ctx, acb);
acb = ggml_reshape_3d(ctx, acb, P, Tp, H);
ggml_tensor* ac = (Tp > T) ? ggml_view_3d(ctx, acb, P, T, H, acb->nb[1], acb->nb[2], 0) : acb;
// ---- Positional scores bd[c,t,H] = qv[t].p[c] (same as build_graph_local) ----
ggml_tensor* bd = ggml_mul_mat(ctx, php, qv); // [P, T, H]
ggml_tensor* scores = ggml_add(ctx, ac, bd); // [P, T, H]
// Band mask [P, T]: 0 if key in [0, valid_len), else -inf.
std::vector<float>& mh = pool.alloc_f32((size_t)P * T);
for (int t = 0; t < T; ++t)
for (int c = 0; c < P; ++c) {
const int key = t - att_left + c;
mh[(size_t)t * P + c] = (key >= 0 && key < valid_len) ? 0.0f : -INFINITY;
}
int64_t mne[2] = {P, T};
ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, mne,
mh.data(), mh.size() * sizeof(float));
ggml_tensor* prob = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); // softmax over c
// ---- Context[dk,t,H] = sum_c prob[c,t] v[t-att_left+c] via inverse-skew + matmul ----
// Pad prob to Tp, chunk [P, C, G, H], inverse-skew to a banded [C+P-1, C, G, H].
ggml_tensor* prob_p = (Tp > T) ? ggml_pad_ext(ctx, prob, 0,0, 0,Tp-T, 0,0, 0,0) : prob;
ggml_tensor* prob_c = ggml_reshape_4d(ctx, prob_p, P, C, G, H);
ggml_tensor* probpad = ggml_pad_ext(ctx, prob_c, 0,C, 0,0, 0,0, 0,0); // ne0 P->C+P
// Pfull[j,i,g] = prob_c[j-i, i, g] (skew view; upper off-band already zero
// from the pad, lower off-band masked below).
ggml_tensor* pfull = ggml_view_4d(ctx, probpad, C+P-1, C, G, H,
(size_t)(C+P-1)*probpad->nb[0], probpad->nb[2], probpad->nb[3], 0);
pfull = ggml_cont(ctx, pfull);
std::vector<float>& b01 = pool.alloc_f32((size_t)(C+P-1) * C);
for (int i = 0; i < C; ++i)
for (int j = 0; j < C+P-1; ++j) {
const int rel = j - i;
b01[(size_t)i * (C+P-1) + j] = (rel >= 0 && rel < P) ? 1.0f : 0.0f;
}
int64_t bne[2] = {C+P-1, C};
ggml_tensor* band01 = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, bne,
b01.data(), b01.size() * sizeof(float));
pfull = ggml_mul(ctx, pfull, band01); // zero the lower off-band (broadcast over G,H)
// Over-padded transposed V chunks [C+P-1, dk, G, H]: Vchunk[j,d,g]=v[gC+j].
ggml_tensor* vpad = ggml_pad_ext(ctx, vh, 0,0, att_left,att_right, 0,0, 0,0); // [dk,T+P-1,H]
if (Lk > (int)vpad->ne[1]) vpad = ggml_pad_ext(ctx, vpad, 0,0, 0,Lk-(int)vpad->ne[1], 0,0, 0,0);
ggml_tensor* vpt = ggml_cont(ctx, ggml_permute(ctx, vpad, 1, 0, 2, 3)); // [Lk, dk, H]
ggml_tensor* vchunk = ggml_view_4d(ctx, vpt, C+P-1, dk, G, H,
vpt->nb[1], (size_t)C*vpt->nb[0], vpt->nb[2], 0);
vchunk = ggml_cont(ctx, vchunk);
// context_g[d,i] = sum_j Vchunk[j,d] Pfull[j,i] -> [dk, C, G, H].
ggml_tensor* cc = ggml_mul_mat(ctx, vchunk, pfull);
cc = ggml_reshape_3d(ctx, cc, dk, Tp, H);
ggml_tensor* context = (Tp > T) ? ggml_view_3d(ctx, cc, dk, T, H, cc->nb[1], cc->nb[2], 0) : cc;
// Merge heads [dk,T,H] -> [dk,H,T] -> [D,T]; mask padded query rows; linear_out.
ggml_tensor* merged = ggml_cont(ctx, ggml_permute(ctx, context, 0, 2, 1, 3));
merged = ggml_reshape_2d(ctx, merged, D, T);
if (valid_len < T) {
std::vector<float>& qm = pool.alloc_f32(T);
for (int t = 0; t < T; ++t) qm[t] = (t < valid_len) ? 1.0f : 0.0f;
int64_t qne[2] = {1, T};
ggml_tensor* qmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, qne,
qm.data(), qm.size() * sizeof(float));
merged = ggml_mul(ctx, merged, qmask);
}
ggml_tensor* Wo = clone_weight(ctx, ml, pre + "linear_out.weight");
ggml_tensor* y = ggml_mul_mat(ctx, Wo, merged);
if (ml.tensor(pre + "linear_out.bias"))
y = ggml_add(ctx, y, clone_weight(ctx, ml, pre + "linear_out.bias"));
return y; // [D, T]
}
ggml_tensor* RelPosAttention::build_graph_batched_local_chunked(
ggml_context* ctx, ggml_tensor* xt, int T, int B, ggml_tensor* pe,
int pos_len, const std::vector<int>& valid_len, int att_left, int att_right,
GraphInputPool& pool, int chunk) const {
const int D = d_model_;
assert((int)valid_len.size() == B);
// Run the O(1) chunk kernel per item (the 4D chunk graph can't also carry a
// batch dim), then stack the per-item [D,T] outputs back into [D,T,B].
ggml_tensor* out = nullptr;
for (int b = 0; b < B; ++b) {
ggml_tensor* xb = ggml_view_2d(ctx, xt, D, T, xt->nb[1], (size_t)b * xt->nb[2]);
xb = ggml_cont(ctx, xb); // linear() mul_mat wants a dense [D,T] item
ggml_tensor* yb = build_graph_local_chunked(ctx, xb, T, pe, pos_len,
valid_len[b], att_left, att_right, pool, chunk); // [D,T]
yb = ggml_reshape_3d(ctx, yb, D, T, 1);
out = out ? ggml_concat(ctx, out, yb, 2) : yb;
}
return out; // [D, T, B]
}
void RelPosAttention::forward_local_chunked(const std::vector<float>& x, int T,
const std::vector<float>& pos_emb, int pos_len,
int valid_len, int att_left, int att_right,
std::vector<float>& out, int chunk) const {
const int D = d_model_;
assert((int)x.size() == T * D);
assert((int)pos_emb.size() == pos_len * D);
GraphInputPool pool;
bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/4,
[&](ggml_context* ctx) -> ggml_tensor* {
int64_t xt_ne[2] = {D, T};
ggml_tensor* xt = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, xt_ne,
x.data(), (size_t)T * D * sizeof(float));
int64_t pe_ne[2] = {D, pos_len};
ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
pos_emb.data(), (size_t)pos_len * D * sizeof(float));
return build_graph_local_chunked(ctx, xt, T, pe, pos_len, valid_len,
att_left, att_right, pool, chunk);
}, out);
assert(ok && "relpos local chunked attention graph failed");
(void)ok;
}
} // namespace pk