One graph: base logits for all n_positions hidden columns (a single lm_head matmul), then for rows [first_corrected, n_positions) the low-rank Markov correction chained along the main path - bias_i = markov_w2 . markov_w1[prev] corrected_i = base_i + bias_i tok_i = argmax(corrected_i) (feeds the next step's get_rows) The chain seed is an I32 graph input; markov_w1 doubles as the previ
| 215 | // corrected_i = base_i + bias_i |
| 216 | // tok_i = argmax(corrected_i) (feeds the next step's get_rows) |
| 217 | // The chain seed is an I32 graph input; markov_w1 doubles as the previous- |
| 218 | // token embedding table. Rows below first_corrected keep the uncorrected base. |
| 219 | bool build_markov_chain_graph(const DraftWeights & dw, |
| 220 | ggml_tensor * lm_head, |
| 221 | int n_positions, int first_corrected, |
| 222 | bool corrected_are_outputs, |
| 223 | bool confidence_are_outputs, |
| 224 | std::vector<uint8_t> & arena, |
| 225 | MarkovChainGraph & out) { |
| 226 | const int hdim = dw.n_embd; |
| 227 | const int vocab = (int)lm_head->ne[1]; |
| 228 | const int n_corr = n_positions - first_corrected; |
| 229 | if (n_positions <= 0 || n_corr <= 0) return false; |
| 230 | const bool have_confidence = confidence_are_outputs && |
| 231 | dw.dspark.confidence_w != nullptr && |
| 232 | dw.dspark.confidence_b != nullptr && |
| 233 | (dw.dspark.confidence_dim == hdim || |
| 234 | dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); |
| 235 | |
| 236 | const size_t arena_size = ggml_tensor_overhead() * (size_t)(64 + 16 * n_corr) + |
| 237 | ggml_graph_overhead_custom(512, false) + 2 * 1024 * 1024; |
| 238 | if (arena.size() < arena_size) arena.resize(arena_size); |
| 239 | |
| 240 | ggml_init_params ip{}; |
| 241 | ip.mem_size = arena.size(); |
| 242 | ip.mem_buffer = arena.data(); |
| 243 | ip.no_alloc = true; |
| 244 | out.ctx = ggml_init(ip); |
| 245 | if (!out.ctx) return false; |
| 246 | out.gf = ggml_new_graph_custom(out.ctx, 512, false); |
| 247 | |
| 248 | out.inp_hidden = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hdim, n_positions); |
| 249 | if (have_confidence) { |
| 250 | out.inp_confidence_hidden = |
| 251 | ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hdim, n_positions); |
| 252 | ggml_set_input(out.inp_confidence_hidden); |
| 253 | } |
| 254 | out.inp_seed = ggml_new_tensor_1d(out.ctx, GGML_TYPE_I32, 1); |
| 255 | ggml_set_input(out.inp_hidden); |
| 256 | ggml_set_input(out.inp_seed); |
| 257 | |
| 258 | out.base = ggml_mul_mat(out.ctx, lm_head, out.inp_hidden); |
| 259 | if (first_corrected > 0) { |
| 260 | // The uncorrected rows are read back by the caller. |
| 261 | ggml_set_output(out.base); |
| 262 | ggml_build_forward_expand(out.gf, out.base); |
| 263 | } |
| 264 | |
| 265 | ggml_tensor * prev_ids = out.inp_seed; |
| 266 | out.toks.assign((size_t)n_corr, nullptr); |
| 267 | out.corrected.assign((size_t)n_corr, nullptr); |
| 268 | out.confidence.assign((size_t)n_corr, nullptr); |
| 269 | for (int i = 0; i < n_corr; ++i) { |
| 270 | const int row = first_corrected + i; |
| 271 | ggml_tensor * prev_emb = ggml_get_rows(out.ctx, dw.dspark.markov_w1, prev_ids); |
| 272 | ggml_tensor * bias = ggml_mul_mat(out.ctx, dw.dspark.markov_w2, prev_emb); |
| 273 | ggml_tensor * base_i = ggml_view_2d(out.ctx, out.base, vocab, 1, |
| 274 | out.base->nb[1], (size_t)row * out.base->nb[1]); |
no test coverage detected