Source code
Revision control
Copy as Markdown
Other Tools
#include "joint.hpp"
#include "backend.hpp"
#include "ggml_graph.hpp"
#include "ggml.h"
#include <cassert>
#include <cstring>
#include <vector>
namespace pk {
Joint::Joint(const ModelLoader& ml) : ml_(ml) {
// Read joint_hidden from the enc weight shape: ne[1] (ggml) = joint_hidden.
ggml_tensor* ew = ml.tensor("joint.enc.weight");
assert(ew && "missing joint.enc.weight");
joint_hidden_ = (int)ew->ne[1];
enc_hidden_ = (int)ew->ne[0];
ggml_tensor* pw = ml.tensor("joint.pred.weight");
assert(pw && "missing joint.pred.weight");
pred_hidden_ = (int)pw->ne[0];
assert((int)pw->ne[1] == joint_hidden_ && "pred/enc joint_hidden mismatch");
vocab_size_ = (int)ml.config().vocab_size;
num_durations_ = (int)ml.config().tdt_durations.size();
V_plus_ = vocab_size_ + 1 + num_durations_;
// Sanity-check the output projection shape (joint_net.2 is kept f32 by the
// converter; enc/pred projections are quantization-allowlisted and may be
// f16/q8_0 — ggml_mul_mat dequantizes those on the fly in the graphs below).
ggml_tensor* wout = ml.tensor("joint.joint_net.2.weight");
assert(wout && "missing joint.joint_net.2.weight");
assert((int)wout->ne[0] == joint_hidden_ && (int)wout->ne[1] == V_plus_ &&
"joint_net.2 weight shape mismatch");
(void)wout;
}
void Joint::precompute_enc_proj(const std::vector<float>& enc, int T, int enc_hidden,
std::vector<float>& enc_proj) const {
assert((int)enc.size() == T * enc_hidden);
assert(enc_hidden == enc_hidden_ && "enc_hidden mismatch");
const int H = joint_hidden_;
// enc_proj[t][h] = enc.weight·enc[t] + enc.bias, over ALL T frames in ONE
// multithreaded ggml matmul on the persistent backend (vs the old per-step
// 1xH matmul that dispatched a graph for every emitted token). Done once per
// utterance, reused by step_logits for every (t,u). ggml_mul_mat dequantizes
// a quantized enc.weight on the fly, so quantized models work here too.
bool ok = pk::run_graph(/*mem ignored*/0, /*n_threads ignored*/0,
[&](ggml_context* ctx) -> ggml_tensor* {
// Input: row-major [T, E] -> ggml ne[0]=E (fastest), ne[1]=T.
int64_t x_ne[2] = { enc_hidden, T };
ggml_tensor* x = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, x_ne,
enc.data(), (size_t)T * enc_hidden * sizeof(float));
// Weight (zero-copy loader leaf): ne[0]=E, ne[1]=H -> mul_mat -> [H, T].
ggml_tensor* W = pk::clone_weight(ctx, ml_, "joint.enc.weight");
ggml_tensor* y = ggml_mul_mat(ctx, W, x);
ggml_tensor* b = pk::clone_weight(ctx, ml_, "joint.enc.bias");
y = ggml_add(ctx, y, b);
return y; // [H, T] ggml memory -> enc_proj[t*H + h]
}, enc_proj);
assert(ok && "enc_proj graph failed");
(void)H;
}
void Joint::step_logits(const float* enc_proj_t,
const float* g, int pred_hidden,
std::vector<float>& logits) const {
assert(pred_hidden == pred_hidden_ && "pred_hidden mismatch");
const int H = joint_hidden_;
// Per-step joint on the PERSISTENT backend (one reused graph; no per-call
// ggml_init/gallocr churn — Backend::compute reuses the backend + gallocr).
// The two matmuls (pred.weight: P->H and joint_net.2.weight: H->V) are the
// hot cost and are memory-bandwidth bound (the H->V weight is ~2.6 MB read
// every step); ggml's matmul parallelizes the output dimension across the
// worker threads, hitting much higher aggregate memory bandwidth than a
// single-threaded C++ matvec — measured ~15x faster per step (26us vs 389us)
// on the 110m. The enc_proj input is the precomputed projection row for t.
bool ok = pk::run_graph(0, 0,
[&](ggml_context* ctx) -> ggml_tensor* {
// enc_proj row for frame t: [H].
int64_t ep_ne[1] = { H };
ggml_tensor* ep = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, ep_ne,
enc_proj_t, (size_t)H * sizeof(float));
// pred-net output g: [P].
int64_t g_ne[1] = { pred_hidden_ };
ggml_tensor* gv = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, g_ne,
g, (size_t)pred_hidden_ * sizeof(float));
// pred_proj = pred.weight·g + pred.bias (P->H). Weight ne=[P,H].
ggml_tensor* Wp = pk::clone_weight(ctx, ml_, "joint.pred.weight");
ggml_tensor* pp = ggml_mul_mat(ctx, Wp, gv); // [H]
ggml_tensor* bp = pk::clone_weight(ctx, ml_, "joint.pred.bias");
pp = ggml_add(ctx, pp, bp);
// f = ReLU(enc_proj + pred_proj)
ggml_tensor* f = ggml_relu(ctx, ggml_add(ctx, ep, pp)); // [H]
// logits = joint_net.2.weight·f + joint_net.2.bias (H->V). Weight ne=[H,V].
ggml_tensor* Wo = pk::clone_weight(ctx, ml_, "joint.joint_net.2.weight");
ggml_tensor* y = ggml_mul_mat(ctx, Wo, f); // [V]
ggml_tensor* bo = pk::clone_weight(ctx, ml_, "joint.joint_net.2.bias");
y = ggml_add(ctx, y, bo);
return y; // [V_plus]
}, logits);
assert(ok && "step_logits graph failed");
}
void Joint::step_logits_batch(const float* enc_proj_gathered,
const float* g, int pred_hidden, int n,
std::vector<float>& logits) const {
assert(pred_hidden == pred_hidden_ && "pred_hidden mismatch");
const int H = joint_hidden_;
// Batched per-step joint over N items on the PERSISTENT backend. Mirrors
// step_logits with a batch axis (ggml ne1 = N): each of the two matmuls is
// applied across all N columns at once, and the biases broadcast over N.
// N=1 reduces exactly to step_logits. The gathered enc_proj input holds one
// joint_hidden row per item (item k at offset k*H), and g holds one
// pred_hidden vector per item (item k at offset k*pred_hidden).
bool ok = pk::run_graph(0, 0,
[&](ggml_context* ctx) -> ggml_tensor* {
// Gathered enc_proj rows: [H, N].
int64_t ep_ne[2] = { H, n };
ggml_tensor* ep = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ep_ne,
enc_proj_gathered, (size_t)H * n * sizeof(float));
// Batched pred-net output g: [P, N].
int64_t g_ne[2] = { pred_hidden_, n };
ggml_tensor* gv = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, g_ne,
g, (size_t)pred_hidden_ * n * sizeof(float));
// pred_proj = pred.weight·g + pred.bias (P->H). Weight ne=[P,H].
ggml_tensor* Wp = pk::clone_weight(ctx, ml_, "joint.pred.weight");
ggml_tensor* pp = ggml_mul_mat(ctx, Wp, gv); // [H, N]
ggml_tensor* bp = pk::clone_weight(ctx, ml_, "joint.pred.bias");
pp = ggml_add(ctx, pp, bp); // bp [H] broadcasts over N
// f = ReLU(enc_proj + pred_proj)
ggml_tensor* f = ggml_relu(ctx, ggml_add(ctx, ep, pp)); // [H, N]
// logits = joint_net.2.weight·f + joint_net.2.bias (H->V). Weight ne=[H,V].
ggml_tensor* Wo = pk::clone_weight(ctx, ml_, "joint.joint_net.2.weight");
ggml_tensor* y = ggml_mul_mat(ctx, Wo, f); // [V, N]
ggml_tensor* bo = pk::clone_weight(ctx, ml_, "joint.joint_net.2.bias");
y = ggml_add(ctx, y, bo); // bo [V] broadcasts over N
return y; // [V_plus, N]
}, logits);
assert(ok && "step_logits_batch graph failed");
}
void Joint::forward(const std::vector<float>& enc, int T, int enc_hidden,
const std::vector<float>& pred, int U, int pred_hidden,
std::vector<float>& logits, int& V_plus_out) const {
assert((int)enc.size() == T * enc_hidden);
assert((int)pred.size() == U * pred_hidden);
const int V = V_plus_;
V_plus_out = V;
// Compute the full [T, U, V_plus] grid by reusing the SAME kernels the greedy
// fast path uses: precompute enc_proj over all T frames once, then emit the
// per-(t,u) logits via step_logits. This keeps the batch path (test_joint /
// test_transducer_core parity vs NeMo) numerically identical to the per-step
// decode path. (Used only by the parity tests; the real decode loops call
// precompute_enc_proj + step_logits directly.)
std::vector<float> enc_proj;
precompute_enc_proj(enc, T, enc_hidden, enc_proj);
const int H = joint_hidden_;
logits.resize((size_t)T * U * V);
std::vector<float> step;
for (int t = 0; t < T; ++t) {
const float* ep_t = enc_proj.data() + (size_t)t * H;
for (int u = 0; u < U; ++u) {
step_logits(ep_t, pred.data() + (size_t)u * pred_hidden, pred_hidden, step);
std::memcpy(logits.data() + ((size_t)t * U + u) * V,
step.data(), (size_t)V * sizeof(float));
}
}
}
} // namespace pk