From 586d74e7f4d16f9b9686ec1faa4f6c4641141486 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Tue, 28 Jul 2026 15:54:48 +0800 Subject: [PATCH 1/2] avoid long-form OOMs in MOSS and Voxtral --- src/arch/moss/decoder.cpp | 96 +++++++++++ src/arch/moss/decoder.h | 36 ++++ src/arch/moss/model.cpp | 261 ++++++++++++++++++++-------- src/arch/voxtral/decoder.cpp | 121 +++++++++++++ src/arch/voxtral/decoder.h | 40 +++++ src/arch/voxtral/model.cpp | 278 ++++++++++++++++++++++-------- src/causal_lm/causal_lm.cpp | 52 ++++++ src/causal_lm/causal_lm.h | 35 ++++ tests/CMakeLists.txt | 15 ++ tests/prefill_chunk_mask_unit.cpp | 155 +++++++++++++++++ 10 files changed, 946 insertions(+), 143 deletions(-) create mode 100644 tests/prefill_chunk_mask_unit.cpp diff --git a/src/arch/moss/decoder.cpp b/src/arch/moss/decoder.cpp index 3a7b3dec..0623c177 100644 --- a/src/arch/moss/decoder.cpp +++ b/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(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, diff --git a/src/arch/moss/decoder.h b/src/arch/moss/decoder.h index b06b8ed6..c812518b 100644 --- a/src/arch/moss/decoder.h +++ b/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 diff --git a/src/arch/moss/model.cpp b/src/arch/moss/model.cpp index abe4b3d7..f1330ca5 100644 --- a/src/arch/moss/model.cpp +++ b/src/arch/moss/model.cpp @@ -565,6 +565,93 @@ transcribe_status encode_one(MossSession * cc, // Build audio_dense [hidden, T_prompt] + keep_mask [1, T_prompt] host-side by // scattering enc_out columns into the audio-pad prompt positions. +// Chunked prefill: push the prompt through the decoder in blocks against the +// growing KV cache, and return the final position's logits. +// +// A single-shot prefill runs every prompt token as one flash-attention call +// (Metal asserts ne01 < 65536, so ~82 min of audio aborts the process) with a +// dense [T_prompt, T_prompt] f16 mask (17 GB at 2 h, twice over counting the +// host copy). Chunking bounds the query rows to `chunk_size` and the mask to +// [n_past + T_chunk, T_chunk]. Prompts that fit in one chunk never come here +// — see run() — so the numerics recorded by the golden dumps are untouched. +transcribe_status prefill_chunked(MossSession * cc, + MossModel * cm, + const std::vector & prompt_ids, + const std::vector & audio_dense, + const std::vector & keep_mask, + int chunk_size, + std::vector & out_logits) { + const int T_prompt = static_cast(prompt_ids.size()); + const int hidden = cm->hparams.dec_hidden; + const int n_chunks = (T_prompt + chunk_size - 1) / chunk_size; + + // Sized for the widest chunk we will build (the last chunk sees the whole + // KV window), then refilled in place — one allocation, not one per chunk. + std::vector mask(static_cast(T_prompt) * std::min(chunk_size, T_prompt)); + std::vector positions(chunk_size); + std::vector kv_idx(chunk_size); + + log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, "moss prefill: %d tokens in %d chunks of %d", T_prompt, n_chunks, chunk_size); + + for (int c = 0; c < n_chunks; ++c) { + const int n_past = c * chunk_size; + const int T_chunk = std::min(chunk_size, T_prompt - n_past); + const int max_n_kv = n_past + T_chunk; + const bool last = (c == n_chunks - 1); + + if (const transcribe_status st = reset_compute_ctx(cc, 16); st != TRANSCRIBE_OK) { + return st; + } + PrefillChunkBuild pb = build_prefill_chunk_graph(cc->compute_ctx, cm->weights, cm->hparams, cc->kv_cache, + T_chunk, max_n_kv, cc->decoder_use_flash, last); + if (pb.graph == nullptr || (last && pb.out == nullptr)) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(cc->sched); + if (!ggml_backend_sched_alloc_graph(cc->sched, pb.graph)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss prefill: chunk %d/%d graph alloc failed", c + 1, n_chunks); + return TRANSCRIBE_ERR_OOM; + } + + ggml_backend_tensor_set(pb.input_ids_in, prompt_ids.data() + n_past, 0, + static_cast(T_chunk) * sizeof(int32_t)); + ggml_backend_tensor_set(pb.audio_dense_in, audio_dense.data() + static_cast(n_past) * hidden, 0, + static_cast(T_chunk) * hidden * sizeof(float)); + ggml_backend_tensor_set(pb.keep_mask_in, keep_mask.data() + n_past, 0, + static_cast(T_chunk) * sizeof(float)); + + for (int i = 0; i < T_chunk; ++i) { + positions[i] = n_past + i; + kv_idx[i] = n_past + i; + } + ggml_backend_tensor_set(pb.positions_in, positions.data(), 0, static_cast(T_chunk) * sizeof(int32_t)); + ggml_backend_tensor_set(pb.kv_idx_in, kv_idx.data(), 0, static_cast(T_chunk) * sizeof(int64_t)); + + causal_lm::fill_prefill_chunk_mask(mask.data(), max_n_kv, T_chunk, n_past); + ggml_backend_tensor_set(pb.mask_in, mask.data(), 0, + static_cast(max_n_kv) * T_chunk * sizeof(ggml_fp16_t)); + + if (ggml_backend_sched_graph_compute(cc->sched, pb.graph) != GGML_STATUS_SUCCESS) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss prefill: chunk %d/%d compute failed", c + 1, n_chunks); + return TRANSCRIBE_ERR_GGUF; + } + + cc->kv_cache.n = max_n_kv; + cc->kv_cache.head = max_n_kv; + + if (last) { + if (transcribe::debug::enabled()) { + // The one dump point chunked prefill can honour: same tensor, + // same shape, same meaning as the single-shot path, which is + // what the chunked-vs-single parity check compares. + transcribe::debug::dump_tensor("dec.logits_raw", pb.out, "dec.logits_raw"); + } + ggml_backend_tensor_get(pb.out, out_logits.data(), 0, out_logits.size() * sizeof(float)); + } + } + return TRANSCRIBE_OK; +} + void build_injection(int hidden, int T_prompt, const std::vector & enc_out, @@ -685,14 +772,11 @@ transcribe_status run(transcribe_session * session, // for long-form far exceeds the k_max_new floor. Clamp to the context. const int gen_budget = std::min(ceiling - T_prompt, std::max(k_max_new, 2 * T_enc + 128)); - // KV cache (grow-to-fit, pow2, clamped to ceiling). - int want_n_ctx = 1024; - while (want_n_ctx < T_prompt + gen_budget) { - want_n_ctx *= 2; - } - if (want_n_ctx > ceiling) { - want_n_ctx = ceiling; - } + // KV cache (grow-to-fit, clamped to ceiling). Short inputs retain the old + // 1K/2K/4K buckets; longer ones grow in 4K steps so crossing 32K does not + // jump straight to the full 131072-token, ~14 GiB worst case. Still + // grow-only, so repeat runs on one session reuse the largest allocation. + const int want_n_ctx = causal_lm::pick_kv_cache_context(T_prompt + gen_budget, ceiling); if (cc->kv_cache.ctx != nullptr && cc->kv_cache.n_ctx < want_n_ctx) { cc->kv_cache.free(); } @@ -711,77 +795,93 @@ transcribe_status run(transcribe_session * session, cc->kv_cache.head = 0; } - // Prefill. - if (const transcribe_status st = reset_compute_ctx(cc, 16); st != TRANSCRIBE_OK) { - return st; - } - const bool slice_last = !dumps_on; - PrefillBuild pb = build_prefill_graph(cc->compute_ctx, cm->weights, cm->hparams, cc->kv_cache, T_prompt, - cc->decoder_use_flash, slice_last); - if (pb.graph == nullptr || pb.out == nullptr) { - return TRANSCRIBE_ERR_GGUF; - } - ggml_backend_sched_reset(cc->sched); - if (!ggml_backend_sched_alloc_graph(cc->sched, pb.graph)) { - return TRANSCRIBE_ERR_OOM; - } + // Prefill. Prompts that fit in one chunk keep the original single-shot + // graph verbatim — that is the graph every golden dump and tolerance was + // recorded against. Longer prompts go chunk-by-chunk, which is the only + // way they run at all: a single flash-attention call is capped at 65535 + // query rows on Metal (~82 min of audio), and the dense [T, T] mask is + // O(T^2) on both the host and the device. + const int vocab = cm->hparams.dec_vocab_size; + const int chunk_size = causal_lm::prefill_chunk_size(); + std::vector logits(vocab); + std::vector audio_dense, keep_mask; + build_injection(cm->hparams.dec_hidden, T_prompt, cc->enc_host, audio_positions, audio_dense, keep_mask); + + int64_t t_prefill_us = 0; + if (T_prompt > chunk_size) { + if (dumps_on) { + log_msg(TRANSCRIBE_LOG_LEVEL_WARN, + "moss run: T_prompt=%d exceeds the %d-token prefill chunk — only dec.logits_raw is dumped " + "(per-chunk hidden states have no single-shot counterpart)", + T_prompt, chunk_size); + } + const int64_t t_pf0 = perf_debug ? ggml_time_us() : 0; + if (const transcribe_status st = + prefill_chunked(cc, cm, prompt_ids, audio_dense, keep_mask, chunk_size, logits); + st != TRANSCRIBE_OK) { + return st; + } + t_prefill_us = perf_debug ? (ggml_time_us() - t_pf0) : 0; + } else { + if (const transcribe_status st = reset_compute_ctx(cc, 16); st != TRANSCRIBE_OK) { + return st; + } + const bool slice_last = !dumps_on; + PrefillBuild pb = build_prefill_graph(cc->compute_ctx, cm->weights, cm->hparams, cc->kv_cache, T_prompt, + cc->decoder_use_flash, slice_last); + if (pb.graph == nullptr || pb.out == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(cc->sched); + if (!ggml_backend_sched_alloc_graph(cc->sched, pb.graph)) { + return TRANSCRIBE_ERR_OOM; + } - ggml_backend_tensor_set(pb.input_ids_in, prompt_ids.data(), 0, prompt_ids.size() * sizeof(int32_t)); - { - std::vector audio_dense, keep_mask; - build_injection(cm->hparams.dec_hidden, T_prompt, cc->enc_host, audio_positions, audio_dense, keep_mask); + ggml_backend_tensor_set(pb.input_ids_in, prompt_ids.data(), 0, prompt_ids.size() * sizeof(int32_t)); ggml_backend_tensor_set(pb.audio_dense_in, audio_dense.data(), 0, audio_dense.size() * sizeof(float)); ggml_backend_tensor_set(pb.keep_mask_in, keep_mask.data(), 0, keep_mask.size() * sizeof(float)); - } - { - std::vector positions(T_prompt); - for (int i = 0; i < T_prompt; ++i) { - positions[i] = i; - } - ggml_backend_tensor_set(pb.positions_in, positions.data(), 0, positions.size() * sizeof(int32_t)); - } - { - const ggml_fp16_t mz = ggml_fp32_to_fp16(0.0f); - const ggml_fp16_t mn = ggml_fp32_to_fp16(-INFINITY); - std::vector mask(static_cast(T_prompt) * T_prompt, mn); - for (int r = 0; r < T_prompt; ++r) { - for (int c = 0; c <= r; ++c) { - mask[static_cast(r) * T_prompt + c] = mz; + { + std::vector positions(T_prompt); + for (int i = 0; i < T_prompt; ++i) { + positions[i] = i; } + ggml_backend_tensor_set(pb.positions_in, positions.data(), 0, positions.size() * sizeof(int32_t)); + } + { + std::vector mask(static_cast(T_prompt) * T_prompt); + causal_lm::fill_prefill_chunk_mask(mask.data(), T_prompt, T_prompt, /*n_past=*/0); + ggml_backend_tensor_set(pb.mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); } - ggml_backend_tensor_set(pb.mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); - } - const int64_t t_pf0 = perf_debug ? ggml_time_us() : 0; - if (ggml_backend_sched_graph_compute(cc->sched, pb.graph) != GGML_STATUS_SUCCESS) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss run: prefill compute failed"); - return TRANSCRIBE_ERR_GGUF; - } - const int64_t t_prefill_us = perf_debug ? (ggml_time_us() - t_pf0) : 0; - cc->kv_cache.n = T_prompt; - cc->kv_cache.head = T_prompt; + const int64_t t_pf0 = perf_debug ? ggml_time_us() : 0; + if (ggml_backend_sched_graph_compute(cc->sched, pb.graph) != GGML_STATUS_SUCCESS) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss run: prefill compute failed"); + return TRANSCRIBE_ERR_GGUF; + } + t_prefill_us = perf_debug ? (ggml_time_us() - t_pf0) : 0; + cc->kv_cache.n = T_prompt; + cc->kv_cache.head = T_prompt; - if (dumps_on) { - auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { - if (t != nullptr) { - transcribe::debug::dump_tensor(name, t, stage); + if (dumps_on) { + auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { + if (t != nullptr) { + transcribe::debug::dump_tensor(name, t, stage); + } + }; + try_dump("dec.token_emb", pb.dumps.token_emb, "dec.token_emb"); + try_dump("dec.audio_injected", pb.dumps.audio_injected, "dec.audio_injected"); + try_dump("dec.block.0.out", pb.dumps.block_0_out, "dec.block.0"); + { + char nm[64]; + std::snprintf(nm, sizeof(nm), "dec.block.%d.out", cm->hparams.dec_n_layers - 1); + try_dump(nm, pb.dumps.block_last_out, "dec.block.last"); } - }; - try_dump("dec.token_emb", pb.dumps.token_emb, "dec.token_emb"); - try_dump("dec.audio_injected", pb.dumps.audio_injected, "dec.audio_injected"); - try_dump("dec.block.0.out", pb.dumps.block_0_out, "dec.block.0"); - { - char nm[64]; - std::snprintf(nm, sizeof(nm), "dec.block.%d.out", cm->hparams.dec_n_layers - 1); - try_dump(nm, pb.dumps.block_last_out, "dec.block.last"); + try_dump("dec.out_before_head", pb.dumps.out_before_head, "dec.out_before_head"); + try_dump("dec.logits_raw", pb.dumps.logits_raw, "dec.logits_raw"); } - try_dump("dec.out_before_head", pb.dumps.out_before_head, "dec.out_before_head"); - try_dump("dec.logits_raw", pb.dumps.logits_raw, "dec.logits_raw"); - } - const int vocab = cm->hparams.dec_vocab_size; - std::vector logits(vocab); - ggml_backend_tensor_get(pb.out, logits.data(), 0, logits.size() * sizeof(float)); + ggml_backend_tensor_get(pb.out, logits.data(), 0, logits.size() * sizeof(float)); + } std::vector generated_ids; int32_t next_tok = argmax_vec(logits); @@ -1004,6 +1104,29 @@ transcribe_status run_batch(transcribe_session * session, if (!cc->decoder_use_flash || transcribe::debug::enabled() || n == 1) { return run_batch_serial(cc, pcm, n_samples, n, params); } + // The batched prefill is still single-shot ([T_max, T_max] mask, every + // query row in one flash-attention call), so a long utterance would trip + // the same 65535-query-row limit chunked prefill exists to avoid. Prompt + // length is a pure function of the sample count, so predict it here — no + // encoder pass needed — and hand the whole batch to the serial path, + // which goes through run() and therefore chunks. + { + const int chunk_size = causal_lm::prefill_chunk_size(); + for (int b = 0; b < n; ++b) { + if (pcm[b] == nullptr || n_samples[b] <= 0) { + continue; + } + std::vector ids, positions; + build_prompt_tokens(cm->hparams, audio_token_length(n_samples[b], cm->hparams), ids, positions); + if (static_cast(ids.size()) > chunk_size) { + log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, + "moss run_batch: utterance %d needs %zu prompt tokens (> %d) — running the batch serially so " + "prefill can chunk", + b, ids.size(), chunk_size); + return run_batch_serial(cc, pcm, n_samples, n, params); + } + } + } transcribe::debug::init(); const int hidden = cm->hparams.dec_hidden; diff --git a/src/arch/voxtral/decoder.cpp b/src/arch/voxtral/decoder.cpp index fab496bf..c08217a8 100644 --- a/src/arch/voxtral/decoder.cpp +++ b/src/arch/voxtral/decoder.cpp @@ -230,6 +230,127 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, return pb; } +PrefillChunkBuild build_prefill_chunk_graph(ggml_context * ctx, + const VoxtralWeights & weights, + const VoxtralHParams & hp, + transcribe::causal_lm::KvCache & kv_cache, + int T_chunk, + int max_n_kv, + int pre_n, + int aud_n, + int suf_n, + 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, "voxtral decoder: invalid chunk (T_chunk=%d max_n_kv=%d)", T_chunk, + max_n_kv); + return pb; + } + if (pre_n < 0 || aud_n < 0 || suf_n < 0 || pre_n + aud_n + suf_n != T_chunk) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral decoder: chunk segments %d+%d+%d != T_chunk(%d)", pre_n, aud_n, + suf_n, T_chunk); + return pb; + } + if (kv_cache.self_k == nullptr || kv_cache.self_v == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral decoder: kv_cache not initialized"); + return pb; + } + if (max_n_kv > kv_cache.n_ctx) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral decoder: max_n_kv=%d exceeds kv_cache.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 auto bp = to_block_params(hp); + const float rms_eps = hp.dec_rms_norm_eps; + + // A chunk landing wholly inside the audio run needs no token embeddings + // at all. Leave the tensor out rather than creating one the graph never + // reaches — the scheduler only allocates tensors it can see, so an unused + // input would have no buffer to upload into. + if (pre_n + suf_n > 0) { + 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); + } + + if (aud_n > 0) { + pb.enc_out_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, aud_n); + named(pb.enc_out_in, "dec.chunk.enc_out"); + ggml_set_input(pb.enc_out_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, /*size=*/16384, /*grads=*/false); + if (gf == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral decoder: ggml_new_graph_custom failed"); + return pb; + } + pb.graph = gf; + + // Same three-way splice as the single-shot builder, restricted to the part + // of each run that falls inside this chunk. + ggml_tensor * token_emb = + (pb.input_ids_in != nullptr) ? ggml_get_rows(ctx, weights.dec_embed.token_w, pb.input_ids_in) : nullptr; + const size_t emb_elem = (token_emb != nullptr) ? ggml_element_size(token_emb) : 0; + + ggml_tensor * x = nullptr; + if (pre_n > 0) { + x = ggml_cont(ctx, ggml_view_2d(ctx, token_emb, hidden, pre_n, emb_elem * hidden, /*off=*/0)); + } + if (aud_n > 0) { + x = (x == nullptr) ? pb.enc_out_in : ggml_concat(ctx, x, pb.enc_out_in, /*dim=*/1); + } + if (suf_n > 0) { + ggml_tensor * x_suffix = ggml_cont(ctx, ggml_view_2d(ctx, token_emb, hidden, suf_n, emb_elem * hidden, + emb_elem * hidden * static_cast(pre_n + aud_n))); + x = (x == nullptr) ? x_suffix : ggml_concat(ctx, x, x_suffix, /*dim=*/1); + } + + 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); + + ggml_tensor * last_x = ggml_view_2d(ctx, x, hidden, 1, ggml_element_size(x) * hidden, + ggml_element_size(x) * hidden * static_cast(T_chunk - 1)); + last_x = ggml_cont(ctx, last_x); + + ggml_tensor * logits = ggml_mul_mat(ctx, weights.dec_embed.output_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 { + ggml_build_forward_expand(gf, x); + } + return pb; +} + StepBuild build_step_graph(ggml_context * ctx, const VoxtralWeights & weights, const VoxtralHParams & hp, diff --git a/src/arch/voxtral/decoder.h b/src/arch/voxtral/decoder.h index ab23f13c..094608ae 100644 --- a/src/arch/voxtral/decoder.h +++ b/src/arch/voxtral/decoder.h @@ -54,6 +54,46 @@ struct PrefillBuild { int suffix_len = 0; }; +// ---------- Chunked prefill ---------- +// +// One chunk of a long prompt, run against the KV earlier chunks wrote. Used +// instead of build_prefill_graph when T_prompt exceeds +// causal_lm::prefill_chunk_size(); shorter prompts keep the single-shot graph +// the golden dumps were recorded against. +// +// Voxtral's decoder context is 131072 tokens but the Metal flash-attention +// kernel asserts ne01 < 65536 query rows, so a single-shot prefill aborts the +// process somewhere past ~87 min of audio; the dense [T, T] mask is the other +// half of the problem. See src/arch/moss/decoder.h — same defect, same fix. +// +// A chunk is a contiguous slice of [prefix | audio | suffix], so it holds at +// most one run of each, in that order: `pre_n` token-embedding rows, then +// `aud_n` encoder rows, then `suf_n` token-embedding rows, summing to T_chunk. +struct PrefillChunkBuild { + ggml_tensor * input_ids_in = nullptr; // [T_chunk] i32 + ggml_tensor * enc_out_in = nullptr; // [dec_hidden, aud_n] f32 (null when aud_n == 0) + 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; +}; + +PrefillChunkBuild build_prefill_chunk_graph(ggml_context * ctx, + const VoxtralWeights & weights, + const VoxtralHParams & hp, + transcribe::causal_lm::KvCache & kv_cache, + int T_chunk, + int max_n_kv, + int pre_n, + int aud_n, + int suf_n, + bool use_flash, + bool want_logits); + // Prefill graph: token-embed the prompt, splice proj.out over the audio // placeholder run, run the LM blocks (writing KV [0,T_prompt)), final // RMSNorm + UNTIED head, output last-position logits. diff --git a/src/arch/voxtral/model.cpp b/src/arch/voxtral/model.cpp index 93c86bea..0f55fa97 100644 --- a/src/arch/voxtral/model.cpp +++ b/src/arch/voxtral/model.cpp @@ -101,19 +101,106 @@ int pick_decode_budget(int n_audio, int t_prompt, int model_max) { return budget; } -// Pick a KV-cache context length that fits `needed` (prompt + decode budget). -// Rounds up to a power of two (so the step graph's `max_n_kv`, computed the same -// way, never exceeds the allocation) and clamps to the trained -// max_position_embeddings — the real ceiling, since RoPE is only valid there. -int pick_kv_ctx(int needed, int model_max) { - int want = 1024; - while (want < needed) { - want *= 2; - } - if (want > model_max) { - want = model_max; - } - return want; +// Chunked prefill — see decoder.h. Walks the prompt in blocks against the +// growing KV cache and returns the final position's logits. The prompt is +// laid out [prefix | audio | suffix], so each chunk holds at most one run of +// each and the audio rows it needs are a contiguous slice of enc_host. +transcribe_status prefill_chunked(VoxtralSession * cc, + VoxtralModel * cm, + const std::vector & prompt_ids, + int prefix_len, + int T_enc, + int chunk_size, + std::vector & out_logits) { + const int T_prompt = static_cast(prompt_ids.size()); + const int hidden = cm->hparams.dec_hidden; + const int n_chunks = (T_prompt + chunk_size - 1) / chunk_size; + const int aud_lo = prefix_len; // first audio position + const int aud_hi = prefix_len + T_enc; // one past the last + + std::vector mask(static_cast(T_prompt) * std::min(chunk_size, T_prompt)); + std::vector positions(chunk_size); + std::vector kv_idx(chunk_size); + + log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, "voxtral prefill: %d tokens in %d chunks of %d", T_prompt, n_chunks, + chunk_size); + + for (int c = 0; c < n_chunks; ++c) { + const int a = c * chunk_size; + const int T_chunk = std::min(chunk_size, T_prompt - a); + const int b = a + T_chunk; + const int max_n_kv = b; + const bool last = (c == n_chunks - 1); + + // Overlap of [a, b) with each of the three runs. + const int pre_n = std::max(0, std::min(b, aud_lo) - a); + const int aud_n = std::max(0, std::min(b, aud_hi) - std::max(a, aud_lo)); + const int suf_n = T_chunk - pre_n - aud_n; + + if (cc->compute_ctx != nullptr) { + ggml_free(cc->compute_ctx); + cc->compute_ctx = nullptr; + } + ggml_init_params ip{}; + ip.mem_size = 64 * 1024 * 1024; + ip.no_alloc = true; + cc->compute_ctx = ggml_init(ip); + if (cc->compute_ctx == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral prefill: ggml_init failed — out of memory."); + return TRANSCRIBE_ERR_OOM; + } + + PrefillChunkBuild pb = + build_prefill_chunk_graph(cc->compute_ctx, cm->weights, cm->hparams, cc->kv_cache, T_chunk, max_n_kv, pre_n, + aud_n, suf_n, cc->decoder_use_flash, last); + if (pb.graph == nullptr || (last && pb.out == nullptr)) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(cc->sched); + if (!ggml_backend_sched_alloc_graph(cc->sched, pb.graph)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral prefill: chunk %d/%d graph alloc failed — out of memory.", + c + 1, n_chunks); + return TRANSCRIBE_ERR_OOM; + } + + if (pb.input_ids_in != nullptr) { // null when the chunk is pure audio + ggml_backend_tensor_set(pb.input_ids_in, prompt_ids.data() + a, 0, + static_cast(T_chunk) * sizeof(int32_t)); + } + if (aud_n > 0) { + const size_t row0 = static_cast(std::max(a, aud_lo) - aud_lo); + ggml_backend_tensor_set(pb.enc_out_in, cc->enc_host.data() + row0 * hidden, 0, + static_cast(aud_n) * hidden * sizeof(float)); + } + + for (int i = 0; i < T_chunk; ++i) { + positions[i] = a + i; + kv_idx[i] = a + i; + } + ggml_backend_tensor_set(pb.positions_in, positions.data(), 0, static_cast(T_chunk) * sizeof(int32_t)); + ggml_backend_tensor_set(pb.kv_idx_in, kv_idx.data(), 0, static_cast(T_chunk) * sizeof(int64_t)); + + causal_lm::fill_prefill_chunk_mask(mask.data(), max_n_kv, T_chunk, /*n_past=*/a); + ggml_backend_tensor_set(pb.mask_in, mask.data(), 0, + static_cast(max_n_kv) * T_chunk * sizeof(ggml_fp16_t)); + + if (const ggml_status gs = ggml_backend_sched_graph_compute(cc->sched, pb.graph); gs != GGML_STATUS_SUCCESS) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral prefill: chunk %d/%d compute failed (%d)", c + 1, n_chunks, + static_cast(gs)); + return TRANSCRIBE_ERR_GGUF; + } + + cc->kv_cache.n = max_n_kv; + cc->kv_cache.head = max_n_kv; + + if (last) { + if (transcribe::debug::enabled()) { + transcribe::debug::dump_tensor("dec.logits_raw", pb.out, "dec.logits_raw"); + } + ggml_backend_tensor_get(pb.out, out_logits.data(), 0, out_logits.size() * sizeof(float)); + } + } + return TRANSCRIBE_OK; } // Input-length contract (see docs/input-limits.md). Hard-context-cap family: @@ -696,7 +783,7 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_INPUT_TOO_LONG; } const int max_new = pick_decode_budget(n_audio_total, T_prompt, model_max); - const int want_ctx = pick_kv_ctx(T_prompt + max_new, model_max); + const int want_ctx = causal_lm::pick_kv_cache_context(T_prompt + max_new, model_max); if (cc->kv_cache.n_ctx < want_ctx) { const ggml_type kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; cc->kv_cache.free(); @@ -736,76 +823,91 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_OOM; } } - PrefillBuild pb = build_prefill_graph(cc->compute_ctx, cm->weights, cm->hparams, cc->kv_cache, T_prompt, - n_audio_total, prefix_len, suffix_len, cc->decoder_use_flash, slice_last); - if (pb.graph == nullptr || pb.out == nullptr) { - return TRANSCRIBE_ERR_GGUF; - } - - ggml_backend_sched_reset(cc->sched); - if (!ggml_backend_sched_alloc_graph(cc->sched, pb.graph)) { - transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "voxtral run: prefill graph allocation failed — out of memory. " - "Lower transcribe_session_params.n_ctx or shorten the audio."); - return TRANSCRIBE_ERR_OOM; - } + const int vocab = cm->hparams.dec_vocab_size; + const int chunk_size = causal_lm::prefill_chunk_size(); + std::vector logits(vocab); - ggml_backend_tensor_set(pb.input_ids_in, prompt_ids.data(), 0, prompt_ids.size() * sizeof(int32_t)); - ggml_backend_tensor_set(pb.enc_out_in, cc->enc_host.data(), 0, cc->enc_host.size() * sizeof(float)); - { - std::vector positions(T_prompt); - for (int i = 0; i < T_prompt; ++i) { - positions[i] = i; + if (T_prompt > chunk_size) { + // Long prompt: chunk it. A single flash-attention call is capped at + // 65535 query rows on Metal, and the [T, T] mask is O(T^2) on both + // host and device — see decoder.h. + if (dumps_on) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_WARN, + "voxtral run: T_prompt=%d exceeds the %d-token prefill chunk — only " + "dec.logits_raw is dumped (per-chunk hidden states have no single-shot " + "counterpart)", + T_prompt, chunk_size); } - ggml_backend_tensor_set(pb.positions_in, positions.data(), 0, positions.size() * sizeof(int32_t)); - } - { - const ggml_fp16_t mz = ggml_fp32_to_fp16(0.0f); - const ggml_fp16_t mn = ggml_fp32_to_fp16(-INFINITY); - std::vector mask(static_cast(T_prompt) * T_prompt, mn); - for (int r = 0; r < T_prompt; ++r) { - for (int col = 0; col <= r; ++col) { - mask[static_cast(r) * T_prompt + col] = mz; - } + set_sched_threads(cc->sched, cc->n_threads); + if (const transcribe_status st = + prefill_chunked(cc, cm, prompt_ids, prefix_len, n_audio_total, chunk_size, logits); + st != TRANSCRIBE_OK) { + return st; + } + } else { + PrefillBuild pb = build_prefill_graph(cc->compute_ctx, cm->weights, cm->hparams, cc->kv_cache, T_prompt, + n_audio_total, prefix_len, suffix_len, cc->decoder_use_flash, slice_last); + if (pb.graph == nullptr || pb.out == nullptr) { + return TRANSCRIBE_ERR_GGUF; } - ggml_backend_tensor_set(pb.mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); - } - set_sched_threads(cc->sched, cc->n_threads); - if (const ggml_status gs = ggml_backend_sched_graph_compute(cc->sched, pb.graph); gs != GGML_STATUS_SUCCESS) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral run: prefill compute failed (%d)", static_cast(gs)); - return TRANSCRIBE_ERR_GGUF; - } - cc->kv_cache.n = T_prompt; - cc->kv_cache.head = T_prompt; + ggml_backend_sched_reset(cc->sched); + if (!ggml_backend_sched_alloc_graph(cc->sched, pb.graph)) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "voxtral run: prefill graph allocation failed — out of memory. " + "Lower transcribe_session_params.n_ctx or shorten the audio."); + return TRANSCRIBE_ERR_OOM; + } - if (dumps_on) { - auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { - if (t != nullptr) { - transcribe::debug::dump_tensor(name, t, stage); - } - }; - try_dump("dec.token_emb", pb.dumps.token_emb, "dec.token_emb"); - try_dump("dec.audio_injected", pb.dumps.audio_injected, "dec.audio_injected"); - try_dump("dec.block.0.out", pb.dumps.block_0_out, "dec.block.0"); + ggml_backend_tensor_set(pb.input_ids_in, prompt_ids.data(), 0, prompt_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(pb.enc_out_in, cc->enc_host.data(), 0, cc->enc_host.size() * sizeof(float)); { - char nm[64]; - std::snprintf(nm, sizeof(nm), "dec.block.%d.out", cm->hparams.dec_n_layers / 2); - try_dump(nm, pb.dumps.block_mid_out, "dec.block.mid"); + std::vector positions(T_prompt); + for (int i = 0; i < T_prompt; ++i) { + positions[i] = i; + } + ggml_backend_tensor_set(pb.positions_in, positions.data(), 0, positions.size() * sizeof(int32_t)); } { - char nm[64]; - std::snprintf(nm, sizeof(nm), "dec.block.%d.out", cm->hparams.dec_n_layers - 1); - try_dump(nm, pb.dumps.block_last_out, "dec.block.last"); + std::vector mask(static_cast(T_prompt) * T_prompt); + causal_lm::fill_prefill_chunk_mask(mask.data(), T_prompt, T_prompt, /*n_past=*/0); + ggml_backend_tensor_set(pb.mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); } - try_dump("dec.out_before_head", pb.dumps.out_before_head, "dec.out_before_head"); - try_dump("dec.logits_raw", pb.dumps.logits_raw, "dec.logits_raw"); - } + set_sched_threads(cc->sched, cc->n_threads); - // ----- First token (argmax of prefill logits) ----- - const int vocab = cm->hparams.dec_vocab_size; - std::vector logits(vocab); - ggml_backend_tensor_get(pb.out, logits.data(), 0, logits.size() * sizeof(float)); + if (const ggml_status gs = ggml_backend_sched_graph_compute(cc->sched, pb.graph); gs != GGML_STATUS_SUCCESS) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral run: prefill compute failed (%d)", static_cast(gs)); + return TRANSCRIBE_ERR_GGUF; + } + cc->kv_cache.n = T_prompt; + cc->kv_cache.head = T_prompt; + + if (dumps_on) { + auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { + if (t != nullptr) { + transcribe::debug::dump_tensor(name, t, stage); + } + }; + try_dump("dec.token_emb", pb.dumps.token_emb, "dec.token_emb"); + try_dump("dec.audio_injected", pb.dumps.audio_injected, "dec.audio_injected"); + try_dump("dec.block.0.out", pb.dumps.block_0_out, "dec.block.0"); + { + char nm[64]; + std::snprintf(nm, sizeof(nm), "dec.block.%d.out", cm->hparams.dec_n_layers / 2); + try_dump(nm, pb.dumps.block_mid_out, "dec.block.mid"); + } + { + char nm[64]; + std::snprintf(nm, sizeof(nm), "dec.block.%d.out", cm->hparams.dec_n_layers - 1); + try_dump(nm, pb.dumps.block_last_out, "dec.block.last"); + } + try_dump("dec.out_before_head", pb.dumps.out_before_head, "dec.out_before_head"); + try_dump("dec.logits_raw", pb.dumps.logits_raw, "dec.logits_raw"); + } + + // ----- First token (argmax of prefill logits) ----- + ggml_backend_tensor_get(pb.out, logits.data(), 0, logits.size() * sizeof(float)); + } auto argmax = [&](const std::vector & v) -> int32_t { int32_t best = 0; float best_v = v[0]; @@ -996,6 +1098,34 @@ transcribe_status run_batch(transcribe_session * session, transcribe::debug::init(); const auto & hp = cm->hparams; + // The batched prefill is still single-shot ([T_max, T_max] mask, all query + // rows in one flash-attention call), so a long utterance would trip the + // same 65535-query-row limit chunked prefill exists to avoid. Audio-token + // count is a pure function of the sample count, so screen for it here and + // send the batch down the serial path, which goes through run() and + // therefore chunks. + { + const int chunk_size = causal_lm::prefill_chunk_size(); + int spc = hp.fe_n_samples; + if (spc <= 0) { + spc = (hp.fe_chunk_length > 0 ? hp.fe_chunk_length : 30) * hp.fe_sample_rate; + } + const int per_chunk = hp.audio_tokens_per_chunk(); + for (int b = 0; b < n; ++b) { + if (pcm[b] == nullptr || n_samples[b] <= 0 || spc <= 0 || per_chunk <= 0) { + continue; + } + const int64_t audio_tok = static_cast((n_samples[b] + spc - 1) / spc) * per_chunk; + if (audio_tok > chunk_size) { // conservative: prompt >= audio tokens + log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, + "voxtral run_batch: utterance %d needs >= %lld prompt tokens (> %d) — running the batch " + "serially so prefill can chunk", + b, static_cast(audio_tok), chunk_size); + return run_batch_serial(cc, pcm, n_samples, n, params); + } + } + } + // ----- Prompt mode (uniform across the batch) ----- const bool translate = (params != nullptr && params->task == TRANSCRIBE_TASK_TRANSLATE); std::string instruction; diff --git a/src/causal_lm/causal_lm.cpp b/src/causal_lm/causal_lm.cpp index f9510882..f8ecb2da 100644 --- a/src/causal_lm/causal_lm.cpp +++ b/src/causal_lm/causal_lm.cpp @@ -6,6 +6,7 @@ #include "ggml-backend.h" #include "ggml.h" #include "transcribe-backend.h" +#include "transcribe-env.h" #include "transcribe-log.h" #include "transcribe-session.h" @@ -13,6 +14,7 @@ #include #include #include +#include #include namespace transcribe::causal_lm { @@ -161,6 +163,56 @@ bool kv_init_batched(KvCache & cache, return true; } +int pick_kv_cache_context(int needed, int model_max) { + if (model_max <= 0) { + return 0; + } + + constexpr int k_min_context = 1024; + constexpr int k_large_step = 4096; + + int64_t want = k_min_context; + if (needed > k_large_step) { + want = (static_cast(needed) + k_large_step - 1) / k_large_step * k_large_step; + } else { + while (want < needed) { + want *= 2; + } + } + return static_cast(std::min(want, model_max)); +} + +void fill_prefill_chunk_mask(ggml_fp16_t * dst, int max_n_kv, int T_chunk, int n_past) { + if (dst == nullptr || max_n_kv <= 0 || T_chunk <= 0 || n_past < 0 || n_past + T_chunk > max_n_kv) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "causal_lm prefill mask: bad geometry (max_n_kv=%d T_chunk=%d n_past=%d)", + max_n_kv, T_chunk, n_past); + return; + } + const ggml_fp16_t keep = ggml_fp32_to_fp16(0.0f); + const ggml_fp16_t drop = ggml_fp32_to_fp16(-INFINITY); + for (int q = 0; q < T_chunk; ++q) { + ggml_fp16_t * row = dst + static_cast(q) * max_n_kv; + // Absolute position of this query; it sees KV rows [0, last]. + const int last = n_past + q; + std::fill(row, row + last + 1, keep); + std::fill(row + last + 1, row + max_n_kv, drop); + } +} + +int prefill_chunk_size() { + int chunk = k_prefill_chunk_default; + if (const char * s = transcribe::env::str("TRANSCRIBE_PREFILL_CHUNK")) { + const long v = std::strtol(s, nullptr, 10); + if (v > 0 && v <= k_prefill_chunk_max) { + chunk = static_cast(v); + } else { + log_msg(TRANSCRIBE_LOG_LEVEL_WARN, "TRANSCRIBE_PREFILL_CHUNK=%s out of range (1..%d) — using %d", s, + k_prefill_chunk_max, chunk); + } + } + return chunk; +} + // Block forward — prefill. ggml_tensor * block_prefill(ggml_context * ctx, ggml_cgraph * gf, diff --git a/src/causal_lm/causal_lm.h b/src/causal_lm/causal_lm.h index c83f394a..46999d84 100644 --- a/src/causal_lm/causal_lm.h +++ b/src/causal_lm/causal_lm.h @@ -156,6 +156,41 @@ ggml_tensor * block_step(ggml_context * ctx, ggml_tensor * kv_idx, // [1] i64 bool use_flash); +// Pick a reusable KV-cache context for one utterance. Preserve the historical +// 1K/2K/4K buckets for short inputs, then grow in 4K increments to avoid the +// large memory cliffs caused by power-of-two rounding. The result is capped at +// model_max; callers reject inputs that do not fit before calling this helper. +int pick_kv_cache_context(int needed, int model_max); + +// Host-side causal mask for one CHUNKED-PREFILL chunk, to be uploaded into +// the [max_n_kv, T_chunk] f16 mask that block_step_n takes (ne[0] = KV +// extent, ne[1] = queries; element (k, q) lives at q*max_n_kv + k). +// +// Query q of the chunk sits at absolute position n_past + q, so it may attend +// to KV rows [0, n_past + q] and nothing beyond. The result is a trapezoid, +// not the square a single-shot prefill uses: rows below n_past are fully +// visible to every query, and the triangular part only covers this chunk. +// n_past == 0 && max_n_kv == T_chunk reproduces the plain causal triangle. +// +// `dst` must have room for max_n_kv * T_chunk entries. Callers are expected +// to reuse one buffer sized for the largest chunk rather than allocating per +// chunk. Requires 0 <= n_past and n_past + T_chunk <= max_n_kv. +void fill_prefill_chunk_mask(ggml_fp16_t * dst, int max_n_kv, int T_chunk, int n_past); + +// Default chunk width for chunked prefill, and the env override that changes +// it (TRANSCRIBE_PREFILL_CHUNK). The override exists so the chunked path can +// be forced on short audio for parity testing against single-shot prefill; +// it is not a user-facing tuning knob. +// +// The cap matters for two independent reasons: the Metal flash-attention +// kernel asserts ne01 < 65536 (query rows per call), and the prefill mask is +// O(T_prompt * chunk) — at 2048 that is ~380 MB for a 2 h clip instead of the +// ~17 GB a single-shot [T, T] mask would need. +constexpr int k_prefill_chunk_default = 2048; +constexpr int k_prefill_chunk_max = 32768; + +int prefill_chunk_size(); + // Block forward (multi-position step). Like block_step but processes // T_seq positions in one forward: writes T_seq rows at the indices in // `kv_idx` (i64 [T_seq]) and reads the full [0, max_n_kv) window. `mask` diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index bec7501c..b0095476 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -186,6 +186,21 @@ transcribe_apply_warnings(transcribe_batch_mask_unit) add_test(NAME transcribe_batch_mask_unit COMMAND transcribe_batch_mask_unit) +# ----------------------------------------------------------------------------- +# Chunked-prefill causal mask geometry (pure host, no model) + +add_executable(transcribe_prefill_chunk_mask_unit + prefill_chunk_mask_unit.cpp) + +target_link_libraries(transcribe_prefill_chunk_mask_unit PRIVATE transcribe ggml) + +target_include_directories(transcribe_prefill_chunk_mask_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_prefill_chunk_mask_unit) + +add_test(NAME transcribe_prefill_chunk_mask_unit COMMAND transcribe_prefill_chunk_mask_unit) + # ----------------------------------------------------------------------------- # MOSS diarized-transcript parser unit test (pure host, no model) # ----------------------------------------------------------------------------- diff --git a/tests/prefill_chunk_mask_unit.cpp b/tests/prefill_chunk_mask_unit.cpp new file mode 100644 index 00000000..973ee019 --- /dev/null +++ b/tests/prefill_chunk_mask_unit.cpp @@ -0,0 +1,155 @@ +// Geometry of the chunked-prefill causal mask. +// +// A single-shot prefill masks a square: query r sees keys [0, r]. Chunked +// prefill masks a trapezoid instead — query q of a chunk starting at n_past +// sits at absolute position n_past + q and sees keys [0, n_past + q], so the +// rows below n_past are visible to everything and only the tail is +// triangular. Getting that boundary wrong by one leaks exactly one future +// token per row, which does not crash and barely moves a short transcript, so +// it needs a direct test rather than an end-to-end one. + +#include "causal_lm/causal_lm.h" +#include "ggml.h" + +#include +#include +#include + +namespace { + +int g_failures = 0; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); \ + ++g_failures; \ + } \ + } while (0) + +bool is_keep(ggml_fp16_t v) { + return ggml_fp16_to_fp32(v) == 0.0f; +} + +bool is_drop(ggml_fp16_t v) { + const float f = ggml_fp16_to_fp32(v); + return std::isinf(f) && f < 0.0f; +} + +void check_kv_context(int needed, int model_max, int expected) { + const int got = transcribe::causal_lm::pick_kv_cache_context(needed, model_max); + if (got != expected) { + std::fprintf(stderr, "FAIL kv context (needed=%d model_max=%d): got %d, expected %d\n", needed, model_max, got, + expected); + ++g_failures; + } +} + +// Every entry of a [max_n_kv, T_chunk] mask matches the causal rule exactly. +void check_chunk(int max_n_kv, int T_chunk, int n_past) { + std::vector mask(static_cast(max_n_kv) * T_chunk, ggml_fp32_to_fp16(123.0f)); + transcribe::causal_lm::fill_prefill_chunk_mask(mask.data(), max_n_kv, T_chunk, n_past); + + int bad = 0; + for (int q = 0; q < T_chunk; ++q) { + for (int k = 0; k < max_n_kv; ++k) { + const ggml_fp16_t v = mask[static_cast(q) * max_n_kv + k]; + const bool want = (k <= n_past + q); // visible iff at or before this query + if (want ? !is_keep(v) : !is_drop(v)) { + ++bad; + } + } + } + if (bad != 0) { + std::fprintf(stderr, "FAIL mask(max_n_kv=%d T_chunk=%d n_past=%d): %d wrong entries\n", max_n_kv, T_chunk, + n_past, bad); + ++g_failures; + } +} + +} // namespace + +int main() { + using transcribe::causal_lm::fill_prefill_chunk_mask; + + // Keep short-audio allocations at their historical sizes, but switch to + // fixed 4K growth above that so a one-token boundary crossing does not + // double a large cache. + check_kv_context(1, 131072, 1024); + check_kv_context(1024, 131072, 1024); + check_kv_context(1025, 131072, 2048); + check_kv_context(2048, 131072, 2048); + check_kv_context(2049, 131072, 4096); + check_kv_context(4096, 131072, 4096); + check_kv_context(4097, 131072, 8192); + check_kv_context(8192, 131072, 8192); + check_kv_context(8193, 131072, 12288); + check_kv_context(32768, 131072, 32768); + check_kv_context(32769, 131072, 36864); + check_kv_context(65537, 131072, 69632); + check_kv_context(131071, 131072, 131072); + check_kv_context(4097, 6000, 6000); // trained/user ceiling wins + check_kv_context(1, 512, 512); + check_kv_context(1, 0, 0); + + // n_past == 0 with a full-width window is the plain causal triangle, i.e. + // exactly what the single-shot prefill path uploads. + check_chunk(/*max_n_kv=*/8, /*T_chunk=*/8, /*n_past=*/0); + check_chunk(1, 1, 0); + + // Interior chunks: a rectangle of history plus a triangle of self. + check_chunk(/*max_n_kv=*/16, /*T_chunk=*/8, /*n_past=*/8); + check_chunk(/*max_n_kv=*/24, /*T_chunk=*/8, /*n_past=*/16); + + // Ragged final chunk (T_chunk smaller than the others), and the + // single-token final chunk that a prompt of chunk*n+1 tokens produces. + check_chunk(/*max_n_kv=*/20, /*T_chunk=*/4, /*n_past=*/16); + check_chunk(/*max_n_kv=*/17, /*T_chunk=*/1, /*n_past=*/16); + + // Odd sizes, so nothing depends on power-of-two alignment. + check_chunk(/*max_n_kv=*/23, /*T_chunk=*/7, /*n_past=*/16); + check_chunk(/*max_n_kv=*/227, /*T_chunk=*/35, /*n_past=*/192); + + // Walking the whole prompt one chunk at a time must tile the square + // triangle a single-shot prefill would have built — no gaps, no overlap, + // no leaked future. Reassemble it and compare against the reference. + { + const int T = 37, chunk = 8; + // Reference: full [T, T] causal mask. + std::vector full(static_cast(T) * T); + fill_prefill_chunk_mask(full.data(), T, T, 0); + + std::vector scratch(static_cast(T) * chunk); + for (int n_past = 0; n_past < T; n_past += chunk) { + const int T_chunk = std::min(chunk, T - n_past); + const int max_n_kv = n_past + T_chunk; + fill_prefill_chunk_mask(scratch.data(), max_n_kv, T_chunk, n_past); + for (int q = 0; q < T_chunk; ++q) { + const int row = n_past + q; + for (int k = 0; k < T; ++k) { + // Beyond this chunk's window the key is simply not read. + const ggml_fp16_t got = (k < max_n_kv) ? scratch[static_cast(q) * max_n_kv + k] + : ggml_fp32_to_fp16(-INFINITY); + const ggml_fp16_t want = full[static_cast(row) * T + k]; + CHECK(is_keep(got) == is_keep(want)); + } + } + } + } + + // Bad geometry is rejected, not written past the end of the buffer. The + // canary stays untouched because the helper bails before filling. + { + std::vector mask(16, ggml_fp32_to_fp16(7.0f)); + fill_prefill_chunk_mask(mask.data(), /*max_n_kv=*/4, /*T_chunk=*/4, /*n_past=*/4); // n_past+T > max_n_kv + CHECK(ggml_fp16_to_fp32(mask[0]) == 7.0f); + fill_prefill_chunk_mask(mask.data(), /*max_n_kv=*/4, /*T_chunk=*/4, /*n_past=*/-1); + CHECK(ggml_fp16_to_fp32(mask[0]) == 7.0f); + fill_prefill_chunk_mask(nullptr, 4, 4, 0); // must not crash + } + + if (g_failures == 0) { + std::printf("OK\n"); + } + return g_failures == 0 ? 0 : 1; +} From d312564fd6850f6c4b8d2b79b422d99bc3a74a54 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Tue, 28 Jul 2026 18:08:10 +0800 Subject: [PATCH 2/2] format --- tests/prefill_chunk_mask_unit.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/prefill_chunk_mask_unit.cpp b/tests/prefill_chunk_mask_unit.cpp index 973ee019..9f6d67f3 100644 --- a/tests/prefill_chunk_mask_unit.cpp +++ b/tests/prefill_chunk_mask_unit.cpp @@ -114,7 +114,7 @@ int main() { // triangle a single-shot prefill would have built — no gaps, no overlap, // no leaked future. Reassemble it and compare against the reference. { - const int T = 37, chunk = 8; + const int T = 37, chunk = 8; // Reference: full [T, T] causal mask. std::vector full(static_cast(T) * T); fill_prefill_chunk_mask(full.data(), T, T, 0); @@ -128,8 +128,8 @@ int main() { const int row = n_past + q; for (int k = 0; k < T; ++k) { // Beyond this chunk's window the key is simply not read. - const ggml_fp16_t got = (k < max_n_kv) ? scratch[static_cast(q) * max_n_kv + k] - : ggml_fp32_to_fp16(-INFINITY); + const ggml_fp16_t got = + (k < max_n_kv) ? scratch[static_cast(q) * max_n_kv + k] : ggml_fp32_to_fp16(-INFINITY); const ggml_fp16_t want = full[static_cast(row) * T + k]; CHECK(is_keep(got) == is_keep(want)); }