fix: place Gemma 4 server <|audio|> block inside the user turn (#440)

* fix: place Gemma 4 server <|audio|> block inside the user turn

Gemma 4 audio understanding over the server (POST /v1/chat/completions with an input_audio part) emitted 0 tokens. The server renders chat messages text-only (MessageContent::text() drops input_audio content parts), so the chat template never produced an <|audio|> marker; expand_gemma4_audio_tokens_for_server then found no placeholder and used its "insert before the final token" fallback, which lands the BOA + AUDIO*N + EOA block inside the trailing model turn of the Gemma generation prompt and forces an immediate EOS. This is the server half of the placement defect that #436 fixed for the CLI.

expand_gemma4_audio_tokens_for_server now places the audio block at the token-stream level, the same layer the image path uses via expand_gemma4_image_tokens / apply_image_token_blocks, so the prefix-cache-sensitive render path in chat_request.rs is untouched. When a rendered <|audio|> placeholder is present it is wrapped in place; otherwise the block is spliced in immediately before the last <end_of_turn> token, which closes the latest user turn in a Gemma generation prompt, so the audio stays inside the user turn. The model worker resolves the <end_of_turn> id from the tokenizer and threads it through both Gemma 4 audio paths: prepare_gemma4_audio_embeddings (Conformer VLM) and prepare_gemma4_unified_audio_embeddings (encoder-free Unified). The legacy before-the-final-token insertion stays only as a last resort for prompts with no <end_of_turn> marker.

Non-audio requests are byte-identical: only the two audio embedding paths change and chat_request.rs is not touched, so text-only, image-only, and non-Gemma server requests keep their exact prompt and prompt-cache prefix behavior. New unit tests cover in-place wrap, user-turn insertion, multi-turn last-user-turn targeting, and both last-resort fallbacks.

* fix(server): resolve Gemma 4 end-of-turn token so audio lands in the user turn

The server places the Gemma 4 audio block before the last end-of-turn token (issue #437), but `resolve_end_of_turn_token_id` only looked up "<end_of_turn>". Gemma 4 renamed the marker to "<turn|>" (id 106; start-of-turn is "<|turn>"), so the lookup returned None and the placement silently degraded to the model-turn fallback, still yielding 0 tokens. Try both "<end_of_turn>" (Gemma 2/3) and "<turn|>" (Gemma 4) so the audio block lands inside the last user turn. Verified end to end: a chat-completions request with an input_audio part to gemma-4-e2b now transcribes (0 -> 120 tokens) instead of emitting EOS at prefill.

* test: add regression tests for resolve_end_of_turn_token_id

Add three unit tests to src/server/model_worker_tests.rs covering the function that had the bug (resolve_end_of_turn_token_id had no tests before this PR):

- resolve_end_of_turn_id_handles_gemma4_turn_marker: a tokenizer with "<turn|>" = id 106 but no "<end_of_turn>" must return Some(106). This is the exact regression guard for the #437 bug, where the old implementation only looked up "<end_of_turn>" and returned None for Gemma 4 tokenizers.
- resolve_end_of_turn_id_handles_gemma23_end_of_turn_marker: a tokenizer with "<end_of_turn>" = id 107 must return Some(107).
- resolve_end_of_turn_id_returns_none_when_no_marker_present: a tokenizer with neither marker must return None so the caller keeps its own fallback.

Each test builds a minimal HuggingFace BPE stub (via two helper functions, stub_eot_tokenizer and stub_tokenizer_no_eot) that mirrors the pattern of MlxcelTokenizer::stub_with_byte_fallback.

Also remove the stale "not yet wired through the chat prompt" note from docs/supported-models.md and replace it with a brief description of what the fix delivers.

Author: inureyes

docs/supported-models.md

@@ -71,7 +71,7 @@ Implemented VLM variants include:
 
 - Gemma 3 VL, Gemma 3n VL, Gemma 4 VL
 - Gemma 4 Unified (`gemma4_unified`): encoder-free text + image + audio + video. Patch-projection vision embedder and waveform-chunk audio path feed the shared Gemma 4 backbone, with blockwise bidirectional attention over image/video token spans during prefill. Video is handled as images-per-frame: frames are extracted with `ffmpeg` (uniform sampling, default 2.0 fps), patchified through the same vision embedder with a per-frame `vision_soft_tokens_per_video_frame` budget (70), and scattered into `video_token_id` placeholder spans. Video is available on both the CLI (`--video`) and the server (`video_url` content blocks). The prompt grows by ~70 soft tokens per sampled frame, so a multi-second clip at the default fps expands past the model's 1024-token sliding window. That over-window single-pass prefill is handled correctly: the windowed prefill mask spans the full prompt (the rotating cache keeps every prefill key and only trims to the window for the decode step), so long clips decode coherently without lowering `--fps`.
-- Gemma 4 VL audio: the Gemma 4 VL checkpoints that ship a Conformer audio tower (for example the `e2b` / `e4b` instruct models) take spoken audio from the CLI with `--audio <path>` and transcribe or answer questions about it. Input audio is resampled to 16 kHz before the Conformer encoder, so a clip at any source rate produces the encoder frame count the duration-based audio-token budget expects. Server-side audio input for this path is not yet wired through the chat prompt.
+- Gemma 4 VL audio: the Gemma 4 VL checkpoints that ship a Conformer audio tower (for example the `e2b` / `e4b` instruct models) take spoken audio from the CLI with `--audio <path>` and transcribe or answer questions about it. Input audio is resampled to 16 kHz before the Conformer encoder, so a clip at any source rate produces the encoder frame count the duration-based audio-token budget expects. Server-side `input_audio` in `POST /v1/chat/completions` is also supported: the audio block is spliced inside the last user turn, before the `<turn|>` end-of-turn marker that Gemma 4 uses (id 106).
 - Llama 4 VLM
 - LLaVA and LLaVA-Bunny
 - Aya Vision and PaliGemma

src/multimodal/vlm_runtime.rs

@@ -1089,40 +1089,77 @@ pub fn expand_gemma4_unified_video_tokens(
     Ok(())
 }
 
-/// Expand audio token placeholder in prompt tokens for server requests.
+/// Expand the Gemma 4 audio placeholder in server prompt tokens into a full
+/// `boa + audio_token*N + eoa` block.
 ///
-/// Replaces the first `audio_token_id` with `boa + audio_token*N + eoa`.
-/// If no audio placeholder is found, inserts before the last token.
+/// Placement order, most specific first:
+///
+/// 1. **In place.** If a rendered `<|audio|>` (`audio_token_id`) placeholder is
+///    already present (e.g. a future template path emits one into the user
+///    turn), wrap the first occurrence as `boa + audio_token*N + eoa`. This is
+///    the same in-place expansion the CLI relies on after its chat template
+///    renders the marker.
+/// 2. **End of the last user turn.** The server renders chat messages
+///    text-only (`MessageContent::text()` drops `input_audio` content parts),
+///    so no `<|audio|>` marker reaches the prompt. When `end_of_turn_token_id`
+///    is supplied, insert the audio block immediately before the *last*
+///    `<end_of_turn>` token. In a Gemma generation prompt
+///    (`<start_of_turn>user\n…<end_of_turn>\n<start_of_turn>model\n`) the final
+///    `<end_of_turn>` closes the latest user turn (the pending model turn has
+///    no closing marker yet), so the block lands inside the user turn. This is
+///    the fix for issue #437: the previous "before the last token" fallback put
+///    the audio inside the model turn, which forced an immediate EOS (0-token
+///    output).
+/// 3. **Before the final token.** Last-resort fallback for prompts with no
+///    `<end_of_turn>` marker (e.g. `--no-chat-template`). Preserves the
+///    historical behavior so a non-Gemma-shaped prompt still gets an audio
+///    block rather than none.
 pub fn expand_gemma4_audio_tokens_for_server(
     prompt_tokens: &mut Vec<i32>,
     audio_token_id: i32,
     boa_token_id: i32,
     eoa_token_id: i32,
     num_audio_tokens: usize,
+    end_of_turn_token_id: Option<i32>,
 ) {
-    let mut expanded = Vec::with_capacity(prompt_tokens.len() + num_audio_tokens + 2);
-    let mut found = false;
-    for &token in prompt_tokens.iter() {
-        if token == audio_token_id && !found {
-            found = true;
-            expanded.push(boa_token_id);
-            expanded.extend(std::iter::repeat_n(audio_token_id, num_audio_tokens));
-            expanded.push(eoa_token_id);
-        } else {
-            expanded.push(token);
+    let build_block = || {
+        let mut block = Vec::with_capacity(num_audio_tokens + 2);
+        block.push(boa_token_id);
+        block.extend(std::iter::repeat_n(audio_token_id, num_audio_tokens));
+        block.push(eoa_token_id);
+        block
+    };
+
+    // 1. Wrap a rendered `<|audio|>` placeholder in place.
+    if prompt_tokens.contains(&audio_token_id) {
+        let mut expanded = Vec::with_capacity(prompt_tokens.len() + num_audio_tokens + 2);
+        let mut found = false;
+        for &token in prompt_tokens.iter() {
+            if token == audio_token_id && !found {
+                found = true;
+                expanded.extend(build_block());
+            } else {
+                expanded.push(token);
+            }
         }
+        *prompt_tokens = expanded;
+        return;
     }
-    // If no placeholder found, insert before last token
-    if !found {
-        let last = expanded.pop();
-        expanded.push(boa_token_id);
-        expanded.extend(std::iter::repeat_n(audio_token_id, num_audio_tokens));
-        expanded.push(eoa_token_id);
-        if let Some(tok) = last {
-            expanded.push(tok);
-        }
+
+    // 2. Insert inside the last user turn (before its closing `<end_of_turn>`).
+    if let Some(eot) = end_of_turn_token_id
+        && let Some(pos) = prompt_tokens.iter().rposition(|&token| token == eot)
+    {
+        prompt_tokens.splice(pos..pos, build_block());
+        return;
+    }
+
+    // 3. Last resort: insert before the final token.
+    let last = prompt_tokens.pop();
+    prompt_tokens.extend(build_block());
+    if let Some(tok) = last {
+        prompt_tokens.push(tok);
     }
-    *prompt_tokens = expanded;
 }
 
 #[cfg(test)]

src/multimodal/vlm_runtime_tests.rs

@@ -13,7 +13,8 @@
 // limitations under the License.
 
 use super::{
-    VlmPreparationSummary, expand_gemma4_unified_video_tokens, expand_gemma4_video_tokens,
+    VlmPreparationSummary, expand_gemma4_audio_tokens_for_server,
+    expand_gemma4_unified_video_tokens, expand_gemma4_video_tokens,
     format_molmo_v1_prompt_for_processor, prepared_embedding_refs,
     shift_molmo_v1_image_input_idx_for_bos, should_prepare_vlm_embeddings,
 };
@@ -221,3 +222,83 @@ fn expand_gemma4_unified_video_tokens_no_op_when_videos_empty() {
     expand_gemma4_unified_video_tokens(&mut prompt, 100, 201, 202, &[]).unwrap();
     assert_eq!(prompt, original);
 }
+
+// Audio token ids used across the audio-placement tests below.
+const AUDIO: i32 = 50; // audio_token_id
+const BOA: i32 = 51; // boa_token_id
+const EOA: i32 = 52; // eoa_token_id
+const EOT: i32 = 106; // <end_of_turn>
+const SOT: i32 = 105; // <start_of_turn>
+
+#[test]
+fn server_audio_wraps_rendered_placeholder_in_place() {
+    // A rendered `<|audio|>` (AUDIO) already sits in the user turn: wrap the
+    // first occurrence as BOA + AUDIO*N + EOA in place, leaving the turn
+    // structure intact.
+    // [BOS, <sot>, user_text, AUDIO, <eot>, <sot>, model]
+    let mut prompt = vec![2, SOT, 7, AUDIO, EOT, SOT, 8];
+    expand_gemma4_audio_tokens_for_server(&mut prompt, AUDIO, BOA, EOA, 3, Some(EOT));
+    assert_eq!(
+        prompt,
+        vec![2, SOT, 7, BOA, AUDIO, AUDIO, AUDIO, EOA, EOT, SOT, 8]
+    );
+}
+
+#[test]
+fn server_audio_inserts_inside_last_user_turn() {
+    // Text-only server render: no AUDIO placeholder. The block must land
+    // before the user turn's closing `<end_of_turn>` so it stays in the user
+    // turn, not the model turn (issue #437).
+    // [BOS, <sot>, user, text, <eot>, <sot>, model]
+    let mut prompt = vec![2, SOT, 11, 7, EOT, SOT, 8];
+    expand_gemma4_audio_tokens_for_server(&mut prompt, AUDIO, BOA, EOA, 3, Some(EOT));
+    // Block spliced before the EOT at index 4, i.e. after the user text.
+    assert_eq!(
+        prompt,
+        vec![2, SOT, 11, 7, BOA, AUDIO, AUDIO, AUDIO, EOA, EOT, SOT, 8]
+    );
+    // The audio block precedes the (only) `<end_of_turn>`, never the trailing
+    // `<start_of_turn>model` generation prompt.
+    let eot_pos = prompt.iter().position(|&t| t == EOT).unwrap();
+    let eoa_pos = prompt.iter().position(|&t| t == EOA).unwrap();
+    assert!(eoa_pos < eot_pos);
+}
+
+#[test]
+fn server_audio_targets_the_last_user_turn_in_multi_turn() {
+    // Multi-turn prompt: user / model / user, then the generation prompt. The
+    // block must go before the LAST `<end_of_turn>` (closing the latest user
+    // turn), not an earlier one.
+    // idx: 0   1    2   3    4    5   6    7    8   9    10   11
+    //     [BOS,<sot>,u1, EOT,<sot>,m1, EOT,<sot>,u2, EOT,<sot>,model]
+    let mut prompt = vec![2, SOT, 21, EOT, SOT, 31, EOT, SOT, 22, EOT, SOT, 8];
+    expand_gemma4_audio_tokens_for_server(&mut prompt, AUDIO, BOA, EOA, 2, Some(EOT));
+    assert_eq!(
+        prompt,
+        vec![
+            2, SOT, 21, EOT, SOT, 31, EOT, SOT, 22, BOA, AUDIO, AUDIO, EOA, EOT, SOT, 8
+        ]
+    );
+    // The first two `<end_of_turn>` markers stay adjacent to their turn text;
+    // only the final user turn gains the audio block.
+    assert_eq!(prompt[3], EOT);
+    assert_eq!(prompt[6], EOT);
+}
+
+#[test]
+fn server_audio_falls_back_before_last_token_without_end_of_turn() {
+    // No placeholder and no `<end_of_turn>` id (e.g. --no-chat-template): keep
+    // the historical "before the final token" insertion as a last resort.
+    let mut prompt = vec![2, 11, 8];
+    expand_gemma4_audio_tokens_for_server(&mut prompt, AUDIO, BOA, EOA, 2, None);
+    assert_eq!(prompt, vec![2, 11, BOA, AUDIO, AUDIO, EOA, 8]);
+}
+
+#[test]
+fn server_audio_falls_back_when_end_of_turn_absent_from_prompt() {
+    // `<end_of_turn>` id is known but the prompt does not contain it: the
+    // last-resort fallback still fires rather than dropping the audio block.
+    let mut prompt = vec![2, 11, 8];
+    expand_gemma4_audio_tokens_for_server(&mut prompt, AUDIO, BOA, EOA, 2, Some(EOT));
+    assert_eq!(prompt, vec![2, 11, BOA, AUDIO, AUDIO, EOA, 8]);
+}

src/server/model_worker.rs

@@ -957,12 +957,28 @@ pub(crate) fn prepare_request_vlm_embeddings(
 
     // Audio-only or audio+images for Gemma4 / Gemma4 Unified
     if !audio.is_empty() {
-        if let Some(embeddings) =
-            prepare_gemma4_unified_audio_embeddings(model, prompt_tokens, images, audio)?
-        {
+        // The server renders chat messages text-only, so the prompt carries no
+        // `<|audio|>` marker (issue #437). Resolve the Gemma `<end_of_turn>`
+        // id so the per-family audio expansion can place the audio block inside
+        // the last user turn instead of the model turn (which forces an
+        // immediate EOS / 0-token output).
+        let end_of_turn_token_id = resolve_end_of_turn_token_id(tokenizer);
+        if let Some(embeddings) = prepare_gemma4_unified_audio_embeddings(
+            model,
+            prompt_tokens,
+            images,
+            audio,
+            end_of_turn_token_id,
+        )? {
             return Ok(Some(embeddings));
         }
-        match prepare_gemma4_audio_embeddings(model, prompt_tokens, images, audio)? {
+        match prepare_gemma4_audio_embeddings(
+            model,
+            prompt_tokens,
+            images,
+            audio,
+            end_of_turn_token_id,
+        )? {
             Some(embeddings) => return Ok(Some(embeddings)),
             None => {
                 // Model does not support audio (not Gemma4 or no audio tower).
@@ -1045,6 +1061,42 @@ pub(crate) fn prepare_request_vlm_embeddings(
     Ok(None)
 }
 
+/// Resolve the Gemma `<end_of_turn>` token id from the tokenizer.
+///
+/// The server flattens chat messages to text-only, so an `input_audio` request
+/// produces a prompt with no `<|audio|>` marker. Knowing the `<end_of_turn>`
+/// id lets [`crate::vlm_runtime::expand_gemma4_audio_tokens_for_server`] place
+/// the audio block inside the last user turn instead of the model turn (issue
+/// #437). Returns `None` when the marker is not a single token in this
+/// tokenizer (non-Gemma tokenizers), in which case the caller keeps the legacy
+/// "before the final token" insertion.
+fn resolve_end_of_turn_token_id(tokenizer: &MlxcelTokenizer) -> Option<i32> {
+    // The end-of-turn marker differs across Gemma generations: Gemma 2/3 use
+    // "<end_of_turn>", while Gemma 4 renamed it to "<turn|>" (id 106, and
+    // "<|turn>" for start-of-turn). Try both so the audio block lands inside
+    // the last user turn on every Gemma checkpoint; "<end_of_turn>" is tried
+    // first because that is the value carried by the non-Gemma-4 templates.
+    const EOT_CANDIDATES: &[&str] = &["<end_of_turn>", "<turn|>"];
+    if let Some(hf) = tokenizer.hf_tokenizer() {
+        for candidate in EOT_CANDIDATES {
+            if let Some(id) = hf.token_to_id(candidate) {
+                return Some(id as i32);
+            }
+        }
+    }
+    // SentencePiece / Tiktoken fallback: accept only when the literal marker
+    // encodes to exactly one token, so a tokenizer that splits it into pieces
+    // does not yield a bogus mid-vocabulary id.
+    for candidate in EOT_CANDIDATES {
+        if let Ok(ids) = tokenizer.encode(candidate, false)
+            && ids.len() == 1
+        {
+            return Some(ids[0] as i32);
+        }
+    }
+    None
+}
+
 /// Process audio (and optionally images) for Gemma4 VLM models.
 ///
 /// Returns `Ok(None)` if the model is not a Gemma4 VLM with audio support.
@@ -1053,6 +1105,7 @@ fn prepare_gemma4_audio_embeddings(
     prompt_tokens: &mut Vec<i32>,
     images: &[Vec<u8>],
     audio_data: &[Vec<u8>],
+    end_of_turn_token_id: Option<i32>,
 ) -> Result<Option<InputEmbeddings>> {
     use crate::audio;
 
@@ -1087,13 +1140,15 @@ fn prepare_gemma4_audio_embeddings(
 
     let num_audio_tokens = audio::compute_audio_num_tokens(samples.len(), sample_rate, 40, 750);
 
-    // Expand audio tokens: BOA + AUDIO*N + EOA
+    // Expand audio tokens: BOA + AUDIO*N + EOA, placed inside the last user
+    // turn (issue #437).
     crate::vlm_runtime::expand_gemma4_audio_tokens_for_server(
         prompt_tokens,
         gemma4_vl.audio_token_id,
         gemma4_vl.boa_token_id,
         gemma4_vl.eoa_token_id,
         num_audio_tokens,
+        end_of_turn_token_id,
     );
 
     // `AudioFeatureExtractor::extract` assumes a 16 kHz waveform (160-sample
@@ -1163,6 +1218,7 @@ fn prepare_gemma4_unified_audio_embeddings(
     prompt_tokens: &mut Vec<i32>,
     images: &[Vec<u8>],
     audio_data: &[Vec<u8>],
+    end_of_turn_token_id: Option<i32>,
 ) -> Result<Option<InputEmbeddings>> {
     use crate::audio;
 
@@ -1201,6 +1257,7 @@ fn prepare_gemma4_unified_audio_embeddings(
         unified.boa_token_id,
         unified.eoa_token_id,
         num_audio_tokens,
+        end_of_turn_token_id,
     );
 
     // Process images alongside audio (encoder-free patch projector).