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
131
132
133
134
135
136
137
138
139
140
141
142
// arch/canary/decoder.h - Canary autoregressive Transformer decoder graph builder.
//
// Untied LM head; per-sublayer dumps at layers {0, n_layers/2, n_layers-1}.
#pragma once
#include "ggml.h"
#include <cstdint>
struct ggml_context;
struct ggml_tensor;
struct ggml_cgraph;
namespace transcribe::canary {
struct CanaryHParams;
struct CanaryWeights;
struct CanaryKvCache;
struct DecoderDumps {
ggml_tensor * token_emb = nullptr;
ggml_tensor * pos_emb = nullptr;
ggml_tensor * embed_norm = nullptr;
// For canary's per-sublayer dumps: indexed by layer, three slots each.
// Slot 0 = post self_attn residual, slot 1 = post cross_attn residual,
// slot 2 = post ffn residual (i.e. block out).
ggml_tensor * sub_self[64] = {};
ggml_tensor * sub_cross[64] = {};
ggml_tensor * sub_ffn[64] = {};
ggml_tensor * out_before_head = nullptr;
ggml_tensor * logits_raw = nullptr; // pre-softmax
ggml_tensor * logits = nullptr; // post-softmax (when log_softmax fires)
};
struct DecoderBuild {
ggml_tensor * token_ids_in = nullptr;
ggml_tensor * pos_ids_in = nullptr;
ggml_tensor * encoder_out_in = nullptr;
ggml_tensor * out = nullptr; // logits [vocab, seq]
ggml_tensor * argmax_out = nullptr; // [seq] i32 when skip_log_softmax
DecoderDumps dumps{};
ggml_cgraph * graph = nullptr;
};
// Build a graph that computes cross-attention K/V for all decoder
// layers from the encoder output, writing them into the cross-attn
// KV cache.
DecoderBuild build_cross_kv_graph(ggml_context * ctx,
const CanaryWeights & w,
const CanaryHParams & hp,
CanaryKvCache & kv_cache,
int T_enc);
// Build a decoder graph that uses the KV cache. Works for both prompt
// pass (n_tokens > 1, n_past == 0) and step pass (n_tokens == 1).
//
// skip_log_softmax: when true, skip the log-softmax tail and add a GPU
// argmax over the last new-token position (saves a vocab-wide softmax
// per step). The pre-softmax logits remain available via dumps.logits_raw.
DecoderBuild build_decoder_graph_kv(ggml_context * ctx,
const CanaryWeights & w,
const CanaryHParams & hp,
CanaryKvCache & kv_cache,
int n_tokens,
int n_past,
int T_enc,
bool skip_log_softmax = false,
bool use_flash = true);
// Static-topology single-token decoder graph. Built once per utterance
// after the prompt pass and reused for every step in the autoregressive
// loop. Compared to build_decoder_graph_kv with n_tokens=1, this graph
// is independent of n_past:
// - KV cache writes go through ggml_set_rows with kv_idx_in as the
// runtime row index (instead of ggml_cpy at a build-time-baked
// view offset).
// - Self-attn K/V reads span the full [0, max_n_kv) window; mask_in
// gates which positions are valid each step.
// Caller pre-allocates the graph + sched once and only updates the
// runtime inputs (token_id_in, pos_id_in, kv_idx_in, mask_in) per step.
struct StepBuild {
ggml_tensor * token_id_in = nullptr; // i32 [1]
ggml_tensor * pos_id_in = nullptr; // i32 [1]
ggml_tensor * kv_idx_in = nullptr; // i64 [1]
ggml_tensor * mask_in = nullptr; // f16 [max_n_kv, 1]
ggml_tensor * argmax_out = nullptr; // i32 [1]
int max_n_kv = 0;
ggml_cgraph * graph = nullptr;
};
StepBuild build_step_graph(ggml_context * ctx,
const CanaryWeights & w,
const CanaryHParams & hp,
CanaryKvCache & kv_cache,
int max_n_kv,
int T_enc,
bool use_flash = true);
// ---------------------------------------------------------------------------
// Offline batched decode (B utterances). Mirrors src/arch/cohere.
// ---------------------------------------------------------------------------
// Batched cross-attention K/V. encoder_out_in is [hidden, T_enc_max, B]
// (each utterance right-padded); per-layer K/V written to the batched cross
// cache slab (layer, b). Padded frames produce bias-only K/V discarded by
// the cross-pad mask at attention time.
DecoderBuild build_cross_kv_graph_batched(ggml_context * ctx,
const CanaryWeights & w,
const CanaryHParams & hp,
CanaryKvCache & kv_cache,
int T_enc_max,
int n_batch);
// Static-topology batched single-step graph for B utterances, reused for
// the prompt feed and generation. Self-attn writes self-KV at kv_idx[b] and
// reads [0,max_n_kv) (self_mask gates valid slots); cross-attn reads
// [0,T_enc_max) (cross_mask gates each utterance's real frames).
struct StepBuildBatched {
ggml_tensor * token_ids_in = nullptr; // i32 [B]
ggml_tensor * pos_ids_in = nullptr; // i32 [B]
ggml_tensor * kv_idx_in = nullptr; // i64 [1, B]
ggml_tensor * self_mask_in = nullptr; // f16 [max_n_kv, 1, 1, B]
ggml_tensor * cross_mask_in = nullptr; // f16 [T_enc_max, 1, 1, B]
ggml_tensor * argmax_out = nullptr; // i32 [B]
int max_n_kv = 0;
int n_batch = 0;
ggml_cgraph * graph = nullptr;
};
StepBuildBatched build_step_graph_batched(ggml_context * ctx,
const CanaryWeights & w,
const CanaryHParams & hp,
CanaryKvCache & kv_cache,
int max_n_kv,
int T_enc_max,
int n_batch,
bool use_flash = true);
} // namespace transcribe::canary