1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
// arch/funasr_nano/decoder.h - Qwen3 LM prefill + step graphs.
//
// Pruned fork of arch/qwen3_asr/decoder.h. Same Qwen3 block: pre-LN
// RMSNorm, GQA, per-head Q/K-RMSNorm, NeoX RoPE @ freq_base=1e6, SwiGLU
// MLP, tied lm_head. Audio injection via three-way concat
// [prefix | adaptor_out[:fake_token_len] | suffix].
#pragma once
#include "causal_lm/causal_lm.h"
#include "funasr_nano.h"
#include "ggml.h"
#include "weights.h"
struct ggml_context;
struct ggml_cgraph;
struct ggml_tensor;
namespace transcribe::funasr_nano {
struct DecoderDumps {
ggml_tensor * token_emb = nullptr;
ggml_tensor * audio_injected = nullptr;
ggml_tensor * block_0_out = nullptr;
ggml_tensor * block_last_out = nullptr;
ggml_tensor * out_before_head = nullptr;
ggml_tensor * logits_raw = nullptr;
};
struct PrefillBuild {
ggml_tensor * input_ids_in = nullptr; // [T_prompt] i32
ggml_tensor * audio_in = nullptr; // [hidden, T_audio] f32 (slice of adaptor_out)
ggml_tensor * positions_in = nullptr; // [T_prompt] i32
ggml_tensor * mask_in = nullptr; // [T_prompt, T_prompt] f16
ggml_tensor * out = nullptr; // [vocab]
DecoderDumps dumps{};
ggml_cgraph * graph = nullptr;
int T_prompt = 0;
int T_audio = 0;
int prefix_len = 0;
int suffix_len = 0;
};
PrefillBuild build_prefill_graph(ggml_context * ctx,
const FunAsrNanoWeights & weights,
const FunAsrNanoHParams & hp,
transcribe::causal_lm::KvCache & kv_cache,
int T_prompt,
int T_audio,
int prefix_len,
int suffix_len,
bool use_flash,
bool slice_last);
struct StepBuild {
ggml_tensor * input_id_in = nullptr; // [1] i32
ggml_tensor * position_in = nullptr; // [1] i32
ggml_tensor * kv_idx_in = nullptr; // [1] i64
ggml_tensor * mask_in = nullptr; // [max_n_kv, 1] f16
ggml_tensor * out = nullptr; // [1] i32 — argmax
ggml_tensor * logits = nullptr; // [vocab] f32 — pre-argmax, for dumps
ggml_cgraph * graph = nullptr;
int max_n_kv = 0;
};
StepBuild build_step_graph(ggml_context * ctx,
const FunAsrNanoWeights & weights,
const FunAsrNanoHParams & hp,
transcribe::causal_lm::KvCache & kv_cache,
int max_n_kv,
bool use_flash);
// ---------- Batched prefill / step (offline transcribe_run_batch) ----------
// Mirror arch/qwen3_asr's batched builders; the only family difference is the
// audio block is the adaptor output (already in llm_dim space). See
// causal_lm::block_prefill_batched / block_step_batched.
struct PrefillBuildBatched {
ggml_tensor * input_ids_in = nullptr; // [T_prompt_max, B] i32
// Audio injection via elementwise blend (no set_rows): audio_dense holds
// the audio embeds scattered (host-side) into their prompt positions, zero
// elsewhere; keep_mask is 0 at audio positions and 1 elsewhere. The block
// input is x*keep_mask + audio_dense. Elementwise ops cross the CPU/CUDA
// split (forced by k-quant token_embd get_rows) cleanly, unlike a set_rows.
ggml_tensor * audio_dense_in = nullptr; // [hidden, T_prompt_max*B] f32
ggml_tensor * keep_mask_in = nullptr; // [1, T_prompt_max*B] f32 (0=audio,1=keep)
ggml_tensor * positions_in = nullptr; // [T_prompt_max] i32
ggml_tensor * mask_in = nullptr; // [T_prompt_max, T_prompt_max] f16
ggml_tensor * kv_idx_in = nullptr; // [T_prompt_max, B] i64
ggml_tensor * last_idx_in = nullptr; // [1, B] i32
ggml_tensor * logits = nullptr; // [vocab, B]
ggml_tensor * out = nullptr; // [B] i32 argmax
ggml_cgraph * graph = nullptr;
int T_prompt_max = 0;
int T_audio_max = 0;
int n_batch = 0;
};
PrefillBuildBatched build_prefill_graph_batched(ggml_context * ctx,
const FunAsrNanoWeights & weights,
const FunAsrNanoHParams & hp,
transcribe::causal_lm::KvCache & kv_cache,
int T_prompt_max,
int T_audio_max,
int n_batch,
bool use_flash);
struct StepBuildBatched {
ggml_tensor * input_ids_in = nullptr; // [B] i32
ggml_tensor * position_in = nullptr; // [B] i32
ggml_tensor * kv_idx_in = nullptr; // [1, B] i64
ggml_tensor * mask_in = nullptr; // [max_n_kv, 1, 1, B] f16
ggml_tensor * logits = nullptr; // [vocab, B]
ggml_tensor * out = nullptr; // [B] i32 argmax
ggml_cgraph * graph = nullptr;
int max_n_kv = 0;
int n_batch = 0;
};
StepBuildBatched build_step_graph_batched(ggml_context * ctx,
const FunAsrNanoWeights & weights,
const FunAsrNanoHParams & hp,
transcribe::causal_lm::KvCache & kv_cache,
int max_n_kv,
int n_batch,
bool use_flash);
} // namespace transcribe::funasr_nano