Skip to content

Commit 88ad4b9

Browse files
crispasr integrationclaude
crispasr integration
and
claude
committed
perf(firered): batch beam decoder — 4.5× faster beam search
Restructure the AED beam-search decoder to process all active beams through each weight matrix in a single batched cpu_matmul_bt call (M=n_active) instead of beam_size separate M=1 calls. Each weight matrix is now read once per step regardless of beam count. Self-attention and cross-attention scoring remain per-beam (different KV histories), but the dominant cost — 8 weight matmuls per layer × 16 layers — is fully batched. JFK 11s Q4_K beam=3 on VPS CPU (4 threads): Before: 137.8s decode, 4922 ms/token After: 30.8s decode, 1099 ms/token (4.48×) Transcript: byte-identical Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 0d44b9c commit 88ad4b9

1 file changed

Lines changed: 132 additions & 96 deletions

File tree

‎src/firered_asr.cpp‎

100755100644
Lines changed: 132 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -2174,96 +2174,120 @@ static char* firered_asr_transcribe_impl(struct firered_asr_context* ctx, const
21742174
}
21752175
}
21762176

2177-
for (int bi = 0; bi < beam_size; bi++) {
2178-
if (beams[bi].finished || beams[bi].score <= -1e8f)
2179-
continue;
2180-
int cur_token = beams[bi].tokens.back();
2177+
// Collect active (non-finished) beam indices
2178+
std::vector<int> active_beams;
2179+
active_beams.reserve(beam_size);
2180+
for (int bi = 0; bi < beam_size; bi++)
2181+
if (!beams[bi].finished && beams[bi].score > -1e8f)
2182+
active_beams.push_back(bi);
2183+
const int n_active = (int)active_beams.size();
2184+
if (n_active == 0)
2185+
break;
21812186

2182-
// Embed
2183-
std::vector<float> x(d);
2187+
// Batched hidden state: X[n_active, d] — one row per active beam
2188+
std::vector<float> X(n_active * d);
2189+
2190+
// Embed all active beams
2191+
for (int a = 0; a < n_active; a++) {
2192+
int cur_token = beams[active_beams[a]].tokens.back();
2193+
float* xi = X.data() + a * d;
21842194
for (int i = 0; i < d; i++) {
2185-
x[i] = emb_w[cur_token * d + i] * scale;
2195+
xi[i] = emb_w[cur_token * d + i] * scale;
21862196
if (step < (int)(m.dec.pe->ne[1]))
2187-
x[i] += pe_dec[step * d + i];
2197+
xi[i] += pe_dec[step * d + i];
21882198
}
2199+
}
21892200

2190-
for (int li = 0; li < hp.n_layers_dec; li++) {
2191-
auto& c = dec_cache[li];
2192-
2193-
// === Self-attention (causal, attend to history) ===
2194-
{
2195-
std::vector<float> xn(d);
2196-
cpu_layernorm(x.data(), c.sattn_norm_w.data(), c.sattn_norm_b.data(), xn.data(), 1, d);
2197-
2198-
// Q/K/V projections via parallelized matmul
2199-
std::vector<float> Q_sa(d), K_cur(d), V_cur(d);
2200-
cpu_matmul_bt(xn.data(), c.sattn_w_qs.data(), Q_sa.data(), 1, d, d);
2201-
if (!c.sattn_w_qs_b.empty())
2202-
for (int i = 0; i < d; i++)
2203-
Q_sa[i] += c.sattn_w_qs_b[i];
2204-
cpu_matmul_bt(xn.data(), c.sattn_w_ks.data(), K_cur.data(), 1, d, d);
2205-
cpu_matmul_bt(xn.data(), c.sattn_w_vs.data(), V_cur.data(), 1, d, d);
2206-
if (!c.sattn_w_vs_b.empty())
2207-
for (int i = 0; i < d; i++)
2208-
V_cur[i] += c.sattn_w_vs_b[i];
2209-
if (!beams[bi].sa_k[li].unique())
2210-
beams[bi].sa_k[li] = std::make_shared<std::vector<float>>(*beams[bi].sa_k[li]);
2211-
if (!beams[bi].sa_v[li].unique())
2212-
beams[bi].sa_v[li] = std::make_shared<std::vector<float>>(*beams[bi].sa_v[li]);
2213-
auto& sa_k_hist = *beams[bi].sa_k[li];
2214-
auto& sa_v_hist = *beams[bi].sa_v[li];
2215-
sa_k_hist.insert(sa_k_hist.end(), K_cur.begin(), K_cur.end());
2216-
sa_v_hist.insert(sa_v_hist.end(), V_cur.begin(), V_cur.end());
2217-
2218-
int n_hist = (int)(sa_k_hist.size() / d); // step + 1
2201+
// Batched scratch buffers (allocated once, reused per layer)
2202+
std::vector<float> XN(n_active * d);
2203+
std::vector<float> Q_batch(n_active * d), K_batch(n_active * d), V_batch(n_active * d);
2204+
std::vector<float> sa_out_batch(n_active * d);
2205+
std::vector<float> fc_batch(n_active * d);
2206+
std::vector<float> Qx_batch(n_active * d);
2207+
std::vector<float> xattn_out_batch(n_active * d);
22192208

2220-
// Multi-head attention: Q[d] @ K_history[n_hist, d]
2221-
std::vector<float> sa_out(d, 0);
2222-
for (int h = 0; h < nh_dec; h++) {
2223-
std::vector<float> scores(n_hist);
2224-
for (int t = 0; t < n_hist; t++) {
2225-
double s = 0;
2226-
for (int dd = 0; dd < hd_dec; dd++)
2227-
s += Q_sa[h * hd_dec + dd] * sa_k_hist[t * d + h * hd_dec + dd];
2228-
scores[t] = (float)s * inv_sqrt_hd_dec;
2229-
}
2230-
cpu_softmax_rows(scores.data(), 1, n_hist);
2231-
for (int dd = 0; dd < hd_dec; dd++) {
2232-
double s = 0;
2233-
for (int t = 0; t < n_hist; t++)
2234-
s += scores[t] * sa_v_hist[t * d + h * hd_dec + dd];
2235-
sa_out[h * hd_dec + dd] = (float)s;
2236-
}
2237-
}
2238-
2239-
// FC + residual
2240-
std::vector<float> sa_fc(d);
2241-
cpu_matmul_bt(sa_out.data(), c.sattn_fc_w.data(), sa_fc.data(), 1, d, d);
2209+
for (int li = 0; li < hp.n_layers_dec; li++) {
2210+
auto& c = dec_cache[li];
2211+
2212+
// === Self-attention (causal, attend to history) ===
2213+
// Batched LayerNorm
2214+
for (int a = 0; a < n_active; a++)
2215+
cpu_layernorm(X.data() + a * d, c.sattn_norm_w.data(), c.sattn_norm_b.data(), XN.data() + a * d, 1,
2216+
d);
2217+
2218+
// Batched QKV projections: read each weight matrix once for all beams
2219+
cpu_matmul_bt(XN.data(), c.sattn_w_qs.data(), Q_batch.data(), n_active, d, d);
2220+
if (!c.sattn_w_qs_b.empty())
2221+
for (int a = 0; a < n_active; a++)
2222+
for (int i = 0; i < d; i++)
2223+
Q_batch[a * d + i] += c.sattn_w_qs_b[i];
2224+
cpu_matmul_bt(XN.data(), c.sattn_w_ks.data(), K_batch.data(), n_active, d, d);
2225+
cpu_matmul_bt(XN.data(), c.sattn_w_vs.data(), V_batch.data(), n_active, d, d);
2226+
if (!c.sattn_w_vs_b.empty())
2227+
for (int a = 0; a < n_active; a++)
22422228
for (int i = 0; i < d; i++)
2243-
x[i] += sa_fc[i] + (c.sattn_fc_b.empty() ? 0 : c.sattn_fc_b[i]);
2229+
V_batch[a * d + i] += c.sattn_w_vs_b[i];
2230+
2231+
// Per-beam self-attention scoring (different KV history per beam)
2232+
for (int a = 0; a < n_active; a++) {
2233+
int bi = active_beams[a];
2234+
if (!beams[bi].sa_k[li].unique())
2235+
beams[bi].sa_k[li] = std::make_shared<std::vector<float>>(*beams[bi].sa_k[li]);
2236+
if (!beams[bi].sa_v[li].unique())
2237+
beams[bi].sa_v[li] = std::make_shared<std::vector<float>>(*beams[bi].sa_v[li]);
2238+
auto& sa_k_hist = *beams[bi].sa_k[li];
2239+
auto& sa_v_hist = *beams[bi].sa_v[li];
2240+
sa_k_hist.insert(sa_k_hist.end(), K_batch.data() + a * d, K_batch.data() + (a + 1) * d);
2241+
sa_v_hist.insert(sa_v_hist.end(), V_batch.data() + a * d, V_batch.data() + (a + 1) * d);
2242+
2243+
int n_hist = (int)(sa_k_hist.size() / d);
2244+
float* Q_sa = Q_batch.data() + a * d;
2245+
float* sa_out = sa_out_batch.data() + a * d;
2246+
memset(sa_out, 0, d * sizeof(float));
2247+
for (int h = 0; h < nh_dec; h++) {
2248+
std::vector<float> scores(n_hist);
2249+
for (int t = 0; t < n_hist; t++)
2250+
scores[t] = cpu_dot(Q_sa + h * hd_dec, sa_k_hist.data() + t * d + h * hd_dec, hd_dec) *
2251+
inv_sqrt_hd_dec;
2252+
cpu_softmax_rows(scores.data(), 1, n_hist);
2253+
for (int dd = 0; dd < hd_dec; dd++) {
2254+
double s = 0;
2255+
for (int t = 0; t < n_hist; t++)
2256+
s += scores[t] * sa_v_hist[t * d + h * hd_dec + dd];
2257+
sa_out[h * hd_dec + dd] = (float)s;
2258+
}
22442259
}
2260+
}
22452261

2246-
// === Cross-attention: attend to encoder output (pre-computed K/V) ===
2247-
std::vector<float> xn(d);
2248-
cpu_layernorm(x.data(), c.xattn_norm_w.data(), c.xattn_norm_b.data(), xn.data(), 1, d);
2249-
2250-
std::vector<float> Qx(d);
2251-
cpu_matmul_bt(xn.data(), c.xattn_w_qs.data(), Qx.data(), 1, d, d);
2252-
if (!c.xattn_w_qs_b.empty())
2262+
// Batched SA FC projection + residual
2263+
cpu_matmul_bt(sa_out_batch.data(), c.sattn_fc_w.data(), fc_batch.data(), n_active, d, d);
2264+
for (int a = 0; a < n_active; a++)
2265+
for (int i = 0; i < d; i++)
2266+
X[a * d + i] += fc_batch[a * d + i] + (c.sattn_fc_b.empty() ? 0 : c.sattn_fc_b[i]);
2267+
2268+
// === Cross-attention: attend to encoder output (pre-computed K/V) ===
2269+
// Batched LayerNorm
2270+
for (int a = 0; a < n_active; a++)
2271+
cpu_layernorm(X.data() + a * d, c.xattn_norm_w.data(), c.xattn_norm_b.data(), XN.data() + a * d, 1,
2272+
d);
2273+
2274+
// Batched cross-attention Q projection
2275+
cpu_matmul_bt(XN.data(), c.xattn_w_qs.data(), Qx_batch.data(), n_active, d, d);
2276+
if (!c.xattn_w_qs_b.empty())
2277+
for (int a = 0; a < n_active; a++)
22532278
for (int i = 0; i < d; i++)
2254-
Qx[i] += c.xattn_w_qs_b[i];
2279+
Qx_batch[a * d + i] += c.xattn_w_qs_b[i];
22552280

2256-
int nh_dec = hp.n_head_dec;
2257-
int hd_dec = d / nh_dec;
2258-
std::vector<float> attn_out(d, 0);
2281+
// Per-beam cross-attention scoring (shared K_enc/V_enc)
2282+
for (int a = 0; a < n_active; a++) {
2283+
float* Qx = Qx_batch.data() + a * d;
2284+
float* attn_out = xattn_out_batch.data() + a * d;
2285+
memset(attn_out, 0, d * sizeof(float));
22592286
for (int h = 0; h < nh_dec; h++) {
22602287
std::vector<float> scores(T_sub);
2261-
for (int t = 0; t < T_sub; t++) {
2262-
double s = 0;
2263-
for (int dd = 0; dd < hd_dec; dd++)
2264-
s += Qx[h * hd_dec + dd] * K_enc[li][t * d + h * hd_dec + dd];
2265-
scores[t] = (float)s * inv_sqrt_hd_dec;
2266-
}
2288+
for (int t = 0; t < T_sub; t++)
2289+
scores[t] = cpu_dot(Qx + h * hd_dec, K_enc[li].data() + t * d + h * hd_dec, hd_dec) *
2290+
inv_sqrt_hd_dec;
22672291
cpu_softmax_rows(scores.data(), 1, T_sub);
22682292
for (int dd = 0; dd < hd_dec; dd++) {
22692293
double s = 0;
@@ -2272,39 +2296,51 @@ static char* firered_asr_transcribe_impl(struct firered_asr_context* ctx, const
22722296
attn_out[h * hd_dec + dd] = (float)s;
22732297
}
22742298
}
2299+
}
22752300

2276-
// FC output projection + residual
2277-
std::vector<float> fc_out(d);
2278-
cpu_matmul_bt(attn_out.data(), c.xattn_fc_w.data(), fc_out.data(), 1, d, d);
2301+
// Batched cross-attention FC + residual
2302+
cpu_matmul_bt(xattn_out_batch.data(), c.xattn_fc_w.data(), fc_batch.data(), n_active, d, d);
2303+
for (int a = 0; a < n_active; a++)
22792304
for (int i = 0; i < d; i++)
2280-
x[i] += fc_out[i] + (c.xattn_fc_b.empty() ? 0 : c.xattn_fc_b[i]);
2281-
2282-
// MLP: LN → Linear(d→4d) → GELU → Linear(4d→d) + residual
2283-
cpu_layernorm(x.data(), c.mlp_norm_w.data(), c.mlp_norm_b.data(), xn.data(), 1, d);
2284-
std::vector<float> h_up(c.di);
2285-
cpu_matmul_bt(xn.data(), c.mlp_w1.data(), h_up.data(), 1, d, c.di);
2305+
X[a * d + i] += fc_batch[a * d + i] + (c.xattn_fc_b.empty() ? 0 : c.xattn_fc_b[i]);
2306+
2307+
// === MLP: LN → Linear(d→4d) → GELU → Linear(4d→d) + residual ===
2308+
for (int a = 0; a < n_active; a++)
2309+
cpu_layernorm(X.data() + a * d, c.mlp_norm_w.data(), c.mlp_norm_b.data(), XN.data() + a * d, 1, d);
2310+
std::vector<float> h_up_batch(n_active * c.di);
2311+
cpu_matmul_bt(XN.data(), c.mlp_w1.data(), h_up_batch.data(), n_active, d, c.di);
2312+
for (int a = 0; a < n_active; a++) {
2313+
float* h_up = h_up_batch.data() + a * c.di;
22862314
if (!c.mlp_b1.empty())
22872315
for (int i = 0; i < c.di; i++)
22882316
h_up[i] += c.mlp_b1[i];
22892317
for (int i = 0; i < c.di; i++) {
22902318
float v = h_up[i];
22912319
h_up[i] = 0.5f * v * (1.0f + tanhf(0.7978845608f * (v + 0.044715f * v * v * v)));
22922320
}
2293-
std::vector<float> mlp_out(d);
2294-
cpu_matmul_bt(h_up.data(), c.mlp_w2.data(), mlp_out.data(), 1, c.di, d);
2321+
}
2322+
std::vector<float> mlp_out_batch(n_active * d);
2323+
cpu_matmul_bt(h_up_batch.data(), c.mlp_w2.data(), mlp_out_batch.data(), n_active, c.di, d);
2324+
for (int a = 0; a < n_active; a++) {
2325+
float* mlp_out = mlp_out_batch.data() + a * d;
22952326
if (!c.mlp_b2.empty())
22962327
for (int i = 0; i < d; i++)
22972328
mlp_out[i] += c.mlp_b2[i];
22982329
for (int i = 0; i < d; i++)
2299-
x[i] += mlp_out[i];
2330+
X[a * d + i] += mlp_out[i];
23002331
}
2332+
}
23012333

2302-
// Final LN + projection
2303-
std::vector<float> xn(d);
2304-
cpu_layernorm(x.data(), norm_w.data(), norm_b.data(), xn.data(), 1, d);
2305-
std::vector<float> logits(odim);
2306-
if (!project_decoder_logits(xn.data(), logits))
2307-
cpu_matmul_bt(xn.data(), prj_w.data(), logits.data(), 1, d, odim);
2334+
// Batched final LN + logit projection
2335+
for (int a = 0; a < n_active; a++)
2336+
cpu_layernorm(X.data() + a * d, norm_w.data(), norm_b.data(), XN.data() + a * d, 1, d);
2337+
std::vector<float> logits_batch(n_active * odim);
2338+
cpu_matmul_bt(XN.data(), prj_w.data(), logits_batch.data(), n_active, d, odim);
2339+
2340+
// Per-beam top-k + candidate generation
2341+
for (int a = 0; a < n_active; a++) {
2342+
int bi = active_beams[a];
2343+
float* logits = logits_batch.data() + a * odim;
23082344

23092345
float mx = logits[0];
23102346
for (int i = 1; i < odim; i++)
@@ -2338,7 +2374,7 @@ static char* firered_asr_transcribe_impl(struct firered_asr_context* ctx, const
23382374
cands.push_back({bi, top_idx[k], beams[bi].score + top_score[k], top_score[k]});
23392375
}
23402376
}
2341-
} // end per-beam computation
2377+
} // end batched beam computation
23422378

23432379
for (int bi = 0; bi < beam_size; bi++)
23442380
if (beams[bi].finished)

0 commit comments

Comments
 (0)