fix(models): suppress gemma4_unified multimodal placeholder tokens

src/commands/chat.rs

@@ -63,6 +63,7 @@ use mlxcel::server::model_provider::model_worker::StreamingDecodeState;
 use mlxcel::tokenizer::{MlxcelTokenizer, load_tokenizer};
 use mlxcel::{CxxGenerator, LanguageModel, SamplingConfig, initialize_runtime, load_model};
 use mlxcel_core::cache::KVCacheMode;
+use mlxcel_core::sampling::TokenBiasMap;
 
 /// Triple-quote fence that opens / closes an ollama-style multiline input
 /// block.
@@ -228,10 +229,22 @@ pub fn run_chat(opts: ChatOptions) -> Result<()> {
     }
     let sampling_config = build_sampling_config(sampling);
 
+    // issue #350: suppress this model's reserved multimodal placeholder ids
+    // (audio / image / video span markers) in the interactive chat generator,
+    // mirroring `run_generation_mode` for the one-shot path. `CxxGenerator`'s
+    // `compose_sampling` injects this cached bias whenever the per-call sampling
+    // config carries no token bias of its own (the chat config's stays empty),
+    // so a placeholder id can never leak into a chat reply. Zero-cost for
+    // non-multimodal models: the suppressed set is empty, the bias map stays
+    // empty, and `apply_token_bias` short-circuits.
+    let mut output_suppression = TokenBiasMap::new();
+    output_suppression.suppress_tokens(&model.output_suppressed_token_ids());
+
     // One generator for the whole session. `generate_streaming` resets its KV
     // cache per call, so re-rendering the full transcript each turn preserves
     // context without forking the generation loop.
-    let mut generator = CxxGenerator::new_with_kv_mode(model.num_layers(), opts.kv_cache_mode);
+    let mut generator = CxxGenerator::new_with_kv_mode(model.num_layers(), opts.kv_cache_mode)
+        .with_token_bias(output_suppression);
 
     let mut editor = DefaultEditor::new()
         .map_err(|e| anyhow!("failed to initialize the interactive line editor: {e}"))?;

src/commands/generate.rs

@@ -846,8 +846,16 @@ fn run_generation_mode<M: LanguageModel>(
     sampling_config: &SamplingConfig,
     vlm_embeddings: Option<&InputEmbeddings>,
     kv_cache_mode: KVCacheMode,
-    token_bias: TokenBiasMap,
+    mut token_bias: TokenBiasMap,
 ) -> Result<(Vec<i32>, GenerationStats)> {
+    // issue #350: mask this model's reserved multimodal placeholder tokens
+    // (audio / image / video span markers) to -inf in the output logits so
+    // they can never leak into generated text. Merged into the token-bias map
+    // here, before any generator is built, so it reaches every sub-path below
+    // (speculative, VLM-embedding, and standard). No-op for non-multimodal
+    // models whose suppressed set is empty.
+    token_bias.suppress_tokens(&model.output_suppressed_token_ids());
+
     let output = if let Some(ref draft_model_path) = args.model.draft_model {
         // resolve the effective DrafterKind from
         // (a) the explicit `--draft-kind` CLI flag, OR

src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs

@@ -910,6 +910,16 @@ impl DFlashGenerator {
 /// Equivalent to upstream `sampler(verify_out.logits)` with the greedy
 /// `sampler = argmax(axis=-1)`. Stochastic samplers are out of scope
 /// for this sub-issue (sub-9 covers stochastic DFlash parity).
+///
+/// LATENT GAP (issue #350): this argmax is raw — it applies no token bias.
+/// Today that is safe because only `gemma4_unified` carries a non-empty
+/// `output_suppressed_token_ids` set, and gemma4_unified uses the MTP verify
+/// path (not DFlash). If a model with a non-empty suppressed-id set ever
+/// adopts the DFlash verify path, this call must be preceded by
+/// `mlxcel_core::sampling::apply_token_bias(logits, token_bias)` — exactly as
+/// `Gemma4MtpTargetAdapter::argmax_from_hidden_positions` does — so that
+/// suppressed placeholder token ids cannot win a verify position and
+/// reintroduce the #350 leak.
 fn argmax_logits_to_array(logits: &MlxArray, seq_len: i32) -> UniquePtr<MlxArray> {
     let shape = ffi::array_shape(logits);
     debug_assert!(shape.len() == 3, "expected [1, seq_len, vocab] logits");

src/lib/mlxcel-core/src/generate.rs

@@ -265,6 +265,30 @@ pub trait LanguageModel {
     /// Get the EOS token IDs for this model
     fn eos_token_ids(&self) -> Vec<i32>;
 
+    /// Token ids this model must never emit as generated text output.
+    ///
+    /// Multimodal models reserve placeholder ids (audio / image / video span
+    /// markers, e.g. Gemma 4 Unified's `audio_token_id`, `image_token_id`,
+    /// `boi_token_id`, ...) purely for INPUT alignment. If one becomes the
+    /// argmax at a near-tie decode step it corrupts the text stream
+    /// (issue #350). Generation paths mask these to `f32::NEG_INFINITY` in
+    /// the per-step [`crate::sampling::TokenBiasMap`] so they can never be
+    /// sampled.
+    ///
+    /// The default is empty: non-multimodal models have nothing to suppress
+    /// and pay zero cost (the bias map stays empty and `apply_token_bias`
+    /// short-circuits). Only families with reserved output-illegal ids
+    /// override this, and they must return ONLY those placeholder ids, never
+    /// real EOS or normal text ids.
+    ///
+    /// Used by: CLI `generate` (`run_generation_mode`) and the server batch
+    /// scheduler, which merge the returned ids into the effective
+    /// `TokenBiasMap` (via [`crate::sampling::TokenBiasMap::suppress_tokens`])
+    /// at generator / scheduler construction.
+    fn output_suppressed_token_ids(&self) -> Vec<i32> {
+        Vec::new()
+    }
+
     /// Forward with pre-computed embeddings (for VLM prefill)
     /// Used by: VisionLanguageModel (Gemma3 VLM)
     fn forward_with_embeddings(

src/lib/mlxcel-core/src/sampling.rs

@@ -134,6 +134,28 @@ impl TokenBiasMap {
         self.byte_fragment_ids.insert(token_id);
     }
 
+    /// Permanently suppress every token id in `ids` by forcing its bias to
+    /// `f32::NEG_INFINITY` (sampled probability becomes 0).
+    ///
+    /// Suppression always wins: an existing finite bias for the same id is
+    /// overwritten. This is the mechanism the generation paths use to mask a
+    /// model's reserved output-illegal tokens (multimodal placeholder ids,
+    /// issue #350) so they can never become the argmax at a near-tie decode
+    /// step. Negative or out-of-range ids are stored but ignored when the
+    /// bias is applied (see [`apply_token_bias`]).
+    ///
+    /// An empty slice is a no-op, so a non-multimodal model (whose
+    /// suppressed set is empty) keeps the bit-exact zero-overhead baseline:
+    /// the map stays empty and `apply_token_bias` short-circuits.
+    ///
+    /// Used by: CLI `generate` (`run_generation_mode`) and the server batch
+    /// scheduler (`enqueue_request`).
+    pub fn suppress_tokens(&mut self, ids: &[i32]) {
+        for &id in ids {
+            self.entries.insert(id, f32::NEG_INFINITY);
+        }
+    }
+
     /// Returns `true` when no bias entries are stored.
     pub fn is_empty(&self) -> bool {
         self.entries.is_empty()
@@ -187,8 +209,9 @@ impl TokenBiasMap {
 /// Invalid token ids (negative, or `>= vocab_size`) are silently ignored —
 /// no panic, no error.
 ///
-/// Used by: standard generation, speculative decoding, batch scheduler
-pub(crate) fn apply_token_bias(logits: &MlxArray, bias: &TokenBiasMap) -> UniquePtr<MlxArray> {
+/// Used by: standard generation, speculative decoding, batch scheduler, MTP
+/// verify (Gemma4 target adapter)
+pub fn apply_token_bias(logits: &MlxArray, bias: &TokenBiasMap) -> UniquePtr<MlxArray> {
     if bias.is_empty() {
         return ffi::copy(logits);
     }
@@ -1634,6 +1657,49 @@ mod tests {
         );
     }
 
+    #[test]
+    fn suppress_tokens_forces_neg_inf_and_drives_probability_to_zero() {
+        // issue #350: suppress_tokens masks each id to -inf so it can never be
+        // sampled, while leaving other tokens (e.g. real EOS) untouched.
+        let mut bias = TokenBiasMap::new();
+        // Simulate a model's reserved multimodal placeholder ids alongside an
+        // existing finite bias that suppression must override.
+        bias.insert(2, 5.0);
+        bias.suppress_tokens(&[1, 2, 4]);
+
+        for id in [1, 2, 4] {
+            let b = *bias.get(&id).expect("suppressed id present");
+            assert!(
+                b.is_infinite() && b.is_sign_negative(),
+                "token {id} must be -inf, got {b}"
+            );
+        }
+        // An id that was never suppressed stays absent (not silenced).
+        assert!(bias.get(&3).is_none(), "untouched token must stay absent");
+
+        // After softmax the suppressed indices carry zero probability.
+        let logits = ffi::from_slice_f32(&[1.0, 1.0, 1.0, 1.0, 1.0], &[1, 5]);
+        let biased = apply_token_bias(&logits, &bias);
+        let probs = ffi::softmax(&biased, -1);
+        ffi::eval(&probs);
+        for id in [1, 2, 4] {
+            assert_eq!(
+                logit_at(&probs, id),
+                0.0,
+                "probability at suppressed token {id} must be 0.0"
+            );
+        }
+        // The non-suppressed token 3 keeps positive probability.
+        assert!(logit_at(&probs, 3) > 0.0, "token 3 must remain reachable");
+    }
+
+    #[test]
+    fn suppress_tokens_empty_slice_is_noop() {
+        let mut bias = TokenBiasMap::new();
+        bias.suppress_tokens(&[]);
+        assert!(bias.is_empty(), "empty suppression keeps the baseline path");
+    }
+
     #[test]
     fn apply_token_bias_multiple_entries() {
         // Multiple entries are applied independently and correctly.

src/loaded_model.rs

@@ -287,6 +287,10 @@ impl LanguageModel for LoadedModel {
         delegate_language_model!(self, eos_token_ids())
     }
 
+    fn output_suppressed_token_ids(&self) -> Vec<i32> {
+        delegate_language_model!(self, output_suppressed_token_ids())
+    }
+
     fn make_caches(&self) -> Vec<mlxcel_core::layers::KVCache> {
         // Use fully qualified syntax to ensure we call the trait method, not inherent methods
         delegate_language_model!(self, make_caches())

src/models/gemma4_mtp_target.rs

@@ -264,12 +264,28 @@ impl<'a> Gemma4MtpTargetAdapter<'a> {
     /// early-stop on the first mismatch: for the small Gemma 4 MTP block sizes
     /// we use today, avoiding `K` separate cxx/MLX calls is more important
     /// than skipping the tail projection on low-accept rounds.
-    fn argmax_from_hidden_positions(&self, hidden_full: &MlxArray) -> Vec<i32> {
+    ///
+    /// `token_bias` is the model's output-suppression set (issue #350). When it
+    /// is non-empty the projected logits are biased before the argmax so a
+    /// suppressed placeholder id can never win a deferred-greedy verify
+    /// position. An empty map (every non-multimodal model) takes the raw argmax,
+    /// preserving the bit-exact baseline.
+    fn argmax_from_hidden_positions(
+        &self,
+        hidden_full: &MlxArray,
+        token_bias: &mlxcel_core::sampling::TokenBiasMap,
+    ) -> Vec<i32> {
         let shape = mlxcel_core::array_shape(hidden_full);
         debug_assert_eq!(shape.len(), 3, "hidden must be 3-D [B, T, H]");
         let expected_len = shape[1].max(0) as usize;
         let logits = self.wrapper.speculative_logits_from_hidden(hidden_full);
-        let argmax = mlxcel_core::argmax_last_axis(logits.as_ref().unwrap());
+        let argmax = if token_bias.is_empty() {
+            mlxcel_core::argmax_last_axis(logits.as_ref().unwrap())
+        } else {
+            let biased =
+                mlxcel_core::sampling::apply_token_bias(logits.as_ref().unwrap(), token_bias);
+            mlxcel_core::argmax_last_axis(&biased)
+        };
         mlxcel_core::eval(&argmax);
         materialize_argmax_i32_vec(&argmax, expected_len)
     }
@@ -560,7 +576,7 @@ impl<'a> MtpTarget for Gemma4MtpTargetAdapter<'a> {
             );
             None
         } else {
-            Some(self.wrapper.forward_with_speculative_sinks(
+            let raw_logits = self.wrapper.forward_with_speculative_sinks(
                 &verify_arr,
                 None,
                 None,
@@ -569,7 +585,23 @@ impl<'a> MtpTarget for Gemma4MtpTargetAdapter<'a> {
                 None,
                 Some(&mut sinks),
                 None,
-            ))
+            );
+            // issue #350: mask the model's output-illegal placeholder ids on the
+            // verify logits BEFORE the per-position argmax AND the logprob
+            // extraction, so a suppressed id can never win a near-tie verify
+            // position (the leak) and the emitted token plus its logprob stay
+            // consistent. The bias is the model's suppressed-id set, forwarded
+            // from the server `enqueue_request` through `MtpGenerator` into
+            // `sampler`. An empty map (every non-multimodal model) short-circuits
+            // here to the raw logits, so the baseline stays bit-exact.
+            if sampler.token_bias.is_empty() {
+                Some(raw_logits)
+            } else {
+                Some(mlxcel_core::sampling::apply_token_bias(
+                    &raw_logits,
+                    &sampler.token_bias,
+                ))
+            }
         };
 
         // Greedy-parity gate: pull the per-position argmax tokens from
@@ -605,7 +637,7 @@ impl<'a> MtpTarget for Gemma4MtpTargetAdapter<'a> {
             .next_back()
             .expect("hidden sink must carry at least one entry");
         let target_tokens = if target_tokens.is_empty() && use_deferred_greedy {
-            self.argmax_from_hidden_positions(hidden_full.as_ref().unwrap())
+            self.argmax_from_hidden_positions(hidden_full.as_ref().unwrap(), &sampler.token_bias)
         } else {
             target_tokens
         };
@@ -1733,7 +1765,7 @@ impl<'a> MtpTarget for Gemma4MtpBatchedTargetAdapter<'a> {
     fn verify_forward_batched(
         &self,
         verify_input_per_row: &[Vec<i32>],
-        _sampler: &SamplingConfig,
+        sampler: &SamplingConfig,
     ) -> Result<MtpBatchedVerifyForwardOutput, DrafterError> {
         // Build the rectangular [B, block_size] verify batch.
         let (verify_arr, width) = Self::rectangular_input(verify_input_per_row, self.batch_size)?;
@@ -1772,6 +1804,20 @@ impl<'a> MtpTarget for Gemma4MtpBatchedTargetAdapter<'a> {
         let (logits, hidden_full, shared_kv) =
             self.batched_sink_forward_with_mask(&verify_arr, None, lp_ref, Some(&valid_ends));
 
+        // issue #350: mask the model's output-illegal placeholder ids on the
+        // [B, width, vocab] verify logits BEFORE the per-row argmax, so a
+        // suppressed id can never win a near-tie verify position (the leak).
+        // `apply_token_bias` broadcasts the [1, vocab] mask across the
+        // [B, width] grid; the bias is forwarded from the server `enqueue_request`
+        // (the batched burst shares one sampler across rows). An empty map (every
+        // non-multimodal model) short-circuits to the raw logits, so the baseline
+        // stays bit-exact.
+        let logits = if sampler.token_bias.is_empty() {
+            logits
+        } else {
+            mlxcel_core::sampling::apply_token_bias(&logits, &sampler.token_bias)
+        };
+
         // Greedy-parity: per-row argmax over the [B, width, vocab]
         // logits. At temperature 0 this matches the drafter-less
         // target's own argmax extension, identically to the B = 1