avoid long-form crashes in MOSS and Voxtral (#115)

* avoid long-form OOMs in MOSS and Voxtral

* format

Author: cjpais

src/arch/moss/decoder.cpp

@@ -179,6 +179,102 @@ PrefillBuild build_prefill_graph(ggml_context *                   ctx,
     return pb;
 }
 
+PrefillChunkBuild build_prefill_chunk_graph(ggml_context *                   ctx,
+                                            const MossWeights &              weights,
+                                            const MossHParams &              hp,
+                                            transcribe::causal_lm::KvCache & kv_cache,
+                                            int                              T_chunk,
+                                            int                              max_n_kv,
+                                            bool                             use_flash,
+                                            bool                             want_logits) {
+    PrefillChunkBuild pb{};
+    pb.T_chunk  = T_chunk;
+    pb.max_n_kv = max_n_kv;
+    if (ctx == nullptr || T_chunk <= 0 || max_n_kv < T_chunk) {
+        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss decoder: invalid chunk (T_chunk=%d max_n_kv=%d)", T_chunk, max_n_kv);
+        return pb;
+    }
+    if (kv_cache.self_k == nullptr || kv_cache.self_v == nullptr) {
+        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss decoder: kv_cache not initialized");
+        return pb;
+    }
+    if (max_n_kv > kv_cache.n_ctx) {
+        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss decoder: max_n_kv=%d exceeds n_ctx=%d", max_n_kv, kv_cache.n_ctx);
+        return pb;
+    }
+
+    const int64_t hidden  = hp.dec_hidden;
+    const int64_t vocab   = hp.dec_vocab_size;
+    const int     n_layer = hp.dec_n_layers;
+    const float   rms_eps = hp.dec_rms_norm_eps;
+    const auto    bp      = to_block_params(hp);
+
+    pb.input_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T_chunk);
+    named(pb.input_ids_in, "dec.chunk.input_ids");
+    ggml_set_input(pb.input_ids_in);
+
+    pb.audio_dense_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, T_chunk);
+    named(pb.audio_dense_in, "dec.chunk.audio_dense");
+    ggml_set_input(pb.audio_dense_in);
+
+    pb.keep_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, T_chunk);
+    named(pb.keep_mask_in, "dec.chunk.keep_mask");
+    ggml_set_input(pb.keep_mask_in);
+
+    pb.positions_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T_chunk);
+    named(pb.positions_in, "dec.chunk.positions");
+    ggml_set_input(pb.positions_in);
+
+    pb.kv_idx_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, T_chunk);
+    named(pb.kv_idx_in, "dec.chunk.kv_idx");
+    ggml_set_input(pb.kv_idx_in);
+
+    pb.mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, max_n_kv, T_chunk);
+    named(pb.mask_in, "dec.chunk.attn_mask");
+    ggml_set_input(pb.mask_in);
+
+    ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false);
+    if (gf == nullptr) {
+        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss decoder: ggml_new_graph_custom failed");
+        return pb;
+    }
+    pb.graph = gf;
+
+    ggml_tensor * token_emb = ggml_get_rows(ctx, weights.dec_embed.token_w, pb.input_ids_in);
+
+    // Same blend injection as the single-shot path, on this chunk's slice.
+    ggml_tensor * x = ggml_add(ctx, ggml_mul(ctx, token_emb, pb.keep_mask_in), pb.audio_dense_in);
+
+    for (int il = 0; il < n_layer; ++il) {
+        x = causal_lm::block_step_n(ctx, gf, x, to_block_view(weights.dec_blocks[il]), bp, kv_cache, il, T_chunk,
+                                    max_n_kv, pb.mask_in, pb.positions_in, pb.kv_idx_in, use_flash);
+    }
+
+    if (want_logits) {
+        x = ggml_mul(ctx, ggml_rms_norm(ctx, x, rms_eps), weights.dec_final.norm_w);
+
+        // Only the chunk's last position feeds the head; the rest of the
+        // prompt has already done its job by writing KV.
+        ggml_tensor * last_x = ggml_view_2d(ctx, x, hidden, 1, ggml_element_size(x) * hidden,
+                                            ggml_element_size(x) * hidden * static_cast<size_t>(T_chunk - 1));
+        last_x               = ggml_cont(ctx, last_x);
+
+        ggml_tensor * logits = ggml_mul_mat(ctx, weights.dec_embed.token_w, last_x);
+        logits               = ggml_reshape_1d(ctx, logits, vocab);
+        named(logits, "dec.logits_raw");
+        transcribe::debug::mark_tensor_for_dump(logits);
+
+        pb.out = logits;
+        ggml_set_output(pb.out);
+        ggml_build_forward_expand(gf, pb.out);
+    } else {
+        // No output tensor to pull the graph through: expand on the last
+        // block's hidden state so every KV write in the chain is scheduled.
+        ggml_build_forward_expand(gf, x);
+    }
+    return pb;
+}
+
 StepBuild build_step_graph(ggml_context *                   ctx,
                            const MossWeights &              weights,
                            const MossHParams &              hp,

src/arch/moss/decoder.h

@@ -52,6 +52,42 @@ PrefillBuild build_prefill_graph(ggml_context *                   ctx,
                                  bool                             use_flash,
                                  bool                             slice_last);
 
+// ---------- Chunked prefill ----------
+//
+// One chunk of a long prompt, run against the KV already written by earlier
+// chunks. Used instead of build_prefill_graph when T_prompt exceeds
+// causal_lm::prefill_chunk_size(); a prompt that fits in one chunk still goes
+// through build_prefill_graph unchanged, so every existing golden dump and
+// tolerance keeps pinning exactly the graph it was recorded against.
+//
+// Bounds two things that a single-shot prefill does not: the flash-attention
+// query-row count (Metal asserts ne01 < 65536) and the causal mask, which is
+// [n_past + T_chunk, T_chunk] here instead of [T_prompt, T_prompt].
+struct PrefillChunkBuild {
+    ggml_tensor * input_ids_in   = nullptr;  // [T_chunk] i32
+    ggml_tensor * audio_dense_in = nullptr;  // [hidden, T_chunk] f32
+    ggml_tensor * keep_mask_in   = nullptr;  // [1, T_chunk] f32
+    ggml_tensor * positions_in   = nullptr;  // [T_chunk] i32
+    ggml_tensor * kv_idx_in      = nullptr;  // [T_chunk] i64
+    ggml_tensor * mask_in        = nullptr;  // [max_n_kv, T_chunk] f16
+    ggml_tensor * out            = nullptr;  // [vocab] last-position logits (final chunk only)
+    ggml_cgraph * graph          = nullptr;
+    int           T_chunk        = 0;
+    int           max_n_kv       = 0;
+};
+
+// `want_logits` is true only for the chunk containing the final prompt
+// position: earlier chunks exist purely to populate the KV cache, so building
+// the tied lm_head for them would cost a [vocab, T_chunk] matmul for nothing.
+PrefillChunkBuild build_prefill_chunk_graph(ggml_context *                   ctx,
+                                            const MossWeights &              weights,
+                                            const MossHParams &              hp,
+                                            transcribe::causal_lm::KvCache & kv_cache,
+                                            int                              T_chunk,
+                                            int                              max_n_kv,
+                                            bool                             use_flash,
+                                            bool                             want_logits);
+
 struct StepBuild {
     ggml_tensor * input_id_in = nullptr;  // [1] i32
     ggml_tensor * position_in = nullptr;  // [1] i32