enne2/qwentts.cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

abi, pipeline, cli: conform qwentts on omnivoice convention

Browse Source
This commit is contained in:
Pascal
2026-05-14 21:42:52 +02:00
parent 259e7059f6
commit 38bf6d762a
11 changed files with 602 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/code-predictor-forward.h
+54 -8
View File
@@ -51,10 +51,27 @@ struct CodePredictorOutput {
std::vector<int32_t> codes;
};
// Manual F32 attention chain for the code predictor block. Same shape
// contract as talker_attn_f32: q [hd, T, n_q_heads], k/v [hd, T_full,
// n_kv], output [hd, n_q_heads, T]. Used when use_flash_attn is false.
static struct ggml_tensor * code_predictor_attn_f32(struct ggml_context * ctx,
struct ggml_tensor * q,
struct ggml_tensor * k,
struct ggml_tensor * v,
struct ggml_tensor * mask,
float scale) {
struct ggml_tensor * scores = ggml_mul_mat(ctx, k, q);
scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f);
struct ggml_tensor * vt = ggml_cont(ctx, ggml_transpose(ctx, v));
struct ggml_tensor * out = ggml_mul_mat(ctx, vt, scores);
return ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3));
}
// One Qwen3 decoder block, KV cached. K and V for the T fresh positions
// are written into the cache at [n_past, n_past+T) on dim 1; the
// attention reads the contiguous slice [0, n_past+T). Returns the layer
// output [hidden, T].
// output [hidden, T]. use_flash_attn and clamp_fp16 follow the same
// contract as in talker-forward.h.
static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * ctx,
const CodePredictorWeights * cw,
const TalkerLayer & layer,
@@ -65,6 +82,8 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
struct ggml_tensor * v_cache,
int n_past,
int T,
bool use_flash_attn,
bool clamp_fp16,
struct ggml_cgraph * gf) {
const int n_q_heads = cw->num_attention_heads;
const int n_kv = cw->num_key_value_heads;
@@ -114,16 +133,32 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
// Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads] for flash_attn_ext.
struct ggml_tensor * q_p = ggml_permute(ctx, q, 0, 2, 1, 3);
// Fused flash attention. Matches the working acestep qw3lm_build_attn
// pattern, fixes the Vulkan autoregressive decode bug.
// Clamp V before attention when clamp_fp16 is set, same rationale
// as the talker block: sub Ampere CUDA tensor cores accumulate in
// FP16 and a V projection overflow corrupts everything downstream.
if (clamp_fp16) {
v_full = ggml_clamp(ctx, v_full, -65504.0f, 65504.0f);
}
// Attention: fused flash kernel or manual F32 chain. Matches the
// working acestep qw3lm_build_attn pattern on the fused branch and
// the omnivoice qwen3_attn_f32 helper on the manual one.
float scale = 1.0f / sqrtf((float) hd);
struct ggml_tensor * attn = ggml_flash_attn_ext(ctx, q_p, k_full, v_full, mask, scale, 0.0f, 0.0f);
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
struct ggml_tensor * attn;
if (use_flash_attn) {
attn = ggml_flash_attn_ext(ctx, q_p, k_full, v_full, mask, scale, 0.0f, 0.0f);
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32);
} else {
attn = code_predictor_attn_f32(ctx, q_p, k_full, v_full, mask, scale);
}
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn);
x = ggml_add(ctx, x, o);
if (clamp_fp16) {
x = ggml_clamp(ctx, x, -65504.0f, 65504.0f);
}
struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps);
h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w);
@@ -135,6 +170,9 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu);
x = ggml_add(ctx, x, mlp);
if (clamp_fp16) {
x = ggml_clamp(ctx, x, -65504.0f, 65504.0f);
}
return x;
}
@@ -142,6 +180,7 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
// position `n_past`, run all 5 layers, and pull the logits for the last
// position through lm_head[g_head]. The cache is written as a side
// effect so subsequent decode steps can append a single token.
// use_flash_attn / clamp_fp16 are forwarded as is to every layer.
static bool code_predictor_run(const CodePredictorWeights * cw,
KVCache * kv,
ggml_backend_sched_t sched,
@@ -150,6 +189,8 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
int n_past,
int talker_hidden,
int g_head,
bool use_flash_attn,
bool clamp_fp16,
std::vector<float> * logits_out) {
const int vocab = cw->vocab_size;
const int n_layers = cw->num_hidden_layers;
@@ -188,7 +229,7 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
for (int l = 0; l < n_layers; l++) {
h = code_predictor_layer_forward(gctx, cw, cw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
kv->v[(size_t) l], n_past, T, gf);
kv->v[(size_t) l], n_past, T, use_flash_attn, clamp_fp16, gf);
}
struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps);
@@ -280,6 +321,7 @@ static void embed_row_from_backend(struct ggml_tensor * t, int row_id, int dim,
// is the Philox subsequence of the c0 sample for this step; the 15
// acoustic samples consume subseq_base + 1 .. subseq_base + 15.
// Returns the full vector of 16 codes. dump_dir may be NULL.
// use_flash_attn / clamp_fp16 are forwarded as is to every internal run.
static bool code_predictor_step(const TalkerWeights * tw,
const CodePredictorWeights * cw,
KVCache * kv,
@@ -291,6 +333,8 @@ static bool code_predictor_step(const TalkerWeights * tw,
float top_p,
int64_t seed,
int64_t subseq_base,
bool use_flash_attn,
bool clamp_fp16,
const char * dump_dir,
CodePredictorOutput * out) {
// sub_input slots live at the talker hidden dimension: both the
@@ -316,7 +360,8 @@ static bool code_predictor_step(const TalkerWeights * tw,
embed_row_from_backend(tw->codec_embedding, c0, talker_hidden, prefill_input.data() + (size_t) talker_hidden);
std::vector<float> logits;
if (!code_predictor_run(cw, kv, sched, prefill_input.data(), 2, 0, talker_hidden, 0, &logits)) {
if (!code_predictor_run(cw, kv, sched, prefill_input.data(), 2, 0, talker_hidden, 0, use_flash_attn, clamp_fp16,
&logits)) {
return false;
}
{
@@ -340,7 +385,8 @@ static bool code_predictor_step(const TalkerWeights * tw,
for (int g = 1; g < n_acoustic; g++) {
embed_row_from_backend(cw->codec_embedding[(size_t) (g - 1)], out->codes[(size_t) g], talker_hidden,
step_input.data());
if (!code_predictor_run(cw, kv, sched, step_input.data(), 1, kv->cur_len, talker_hidden, g, &logits)) {
if (!code_predictor_run(cw, kv, sched, step_input.data(), 1, kv->cur_len, talker_hidden, g, use_flash_attn,
clamp_fp16, &logits)) {
return false;
}
float u_g = 0.0f;