diff --git a/src/commands/chat.rs b/src/commands/chat.rs index 63f85fa30..408e7269b 100644 --- a/src/commands/chat.rs +++ b/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}"))?; diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 6433a05d6..319caefc9 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -846,8 +846,16 @@ fn run_generation_mode( sampling_config: &SamplingConfig, vlm_embeddings: Option<&InputEmbeddings>, kv_cache_mode: KVCacheMode, - token_bias: TokenBiasMap, + mut token_bias: TokenBiasMap, ) -> Result<(Vec, 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 diff --git a/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs b/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs index 6f9f4ace0..05c9d6d55 100644 --- a/src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs +++ b/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 { let shape = ffi::array_shape(logits); debug_assert!(shape.len() == 3, "expected [1, seq_len, vocab] logits"); diff --git a/src/lib/mlxcel-core/src/generate.rs b/src/lib/mlxcel-core/src/generate.rs index 03860b3b4..f70500a94 100644 --- a/src/lib/mlxcel-core/src/generate.rs +++ b/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; + /// 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 { + Vec::new() + } + /// Forward with pre-computed embeddings (for VLM prefill) /// Used by: VisionLanguageModel (Gemma3 VLM) fn forward_with_embeddings( diff --git a/src/lib/mlxcel-core/src/sampling.rs b/src/lib/mlxcel-core/src/sampling.rs index eca94f4c1..1cc8830f4 100644 --- a/src/lib/mlxcel-core/src/sampling.rs +++ b/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 { +/// Used by: standard generation, speculative decoding, batch scheduler, MTP +/// verify (Gemma4 target adapter) +pub fn apply_token_bias(logits: &MlxArray, bias: &TokenBiasMap) -> UniquePtr { 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. diff --git a/src/loaded_model.rs b/src/loaded_model.rs index 88fa81339..29a9bd4ad 100644 --- a/src/loaded_model.rs +++ b/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 { + delegate_language_model!(self, output_suppressed_token_ids()) + } + fn make_caches(&self) -> Vec { // Use fully qualified syntax to ensure we call the trait method, not inherent methods delegate_language_model!(self, make_caches()) diff --git a/src/models/gemma4_mtp_target.rs b/src/models/gemma4_mtp_target.rs index c0098cc5c..e5918941c 100644 --- a/src/models/gemma4_mtp_target.rs +++ b/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 { + /// + /// `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 { 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], - _sampler: &SamplingConfig, + sampler: &SamplingConfig, ) -> Result { // 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 diff --git a/src/models/gemma4_mtp_target_tests.rs b/src/models/gemma4_mtp_target_tests.rs index c15f79eaa..5a44a67ca 100644 --- a/src/models/gemma4_mtp_target_tests.rs +++ b/src/models/gemma4_mtp_target_tests.rs @@ -114,6 +114,63 @@ fn argmax_per_position_returns_one_id_per_position() { assert_eq!(argmax, vec![2, 0, 3]); } +#[test] +fn verify_bias_suppresses_argmax_at_leaking_position_b1() { + // issue #350: the B = 1 verify step applies `sampler.token_bias` to the + // logits before `argmax_per_position`. Model the leak with a `[1, 3, 4]` + // verify block where the placeholder id 2 wins position 0 by raw argmax, + // then confirm the suppression bias re-routes ONLY that position to the + // next-best non-suppressed id while the legitimate positions are unchanged. + let _runtime = crate::initialize_runtime(); + + let data: Vec = vec![ + // pos 0: raw argmax 2 (placeholder), next-best 3 + 0.1, 0.2, 0.9, 0.7, // pos 1: legitimate argmax 0 + 0.9, 0.1, 0.2, 0.3, // pos 2: legitimate argmax 3 + 0.1, 0.2, 0.3, 0.9, + ]; + let logits = mlxcel_core::from_slice_f32(&data, &[1, 3, 4]); + + // Raw argmax leaks the placeholder id 2 at position 0. + let raw = Gemma4MtpTargetAdapter::argmax_per_position(logits.as_ref().unwrap()); + assert_eq!(raw, vec![2, 0, 3]); + + // Suppress id 2; the biased argmax must pick the next-best id 3 at + // position 0 while positions 1 and 2 stay put. + let mut bias = mlxcel_core::sampling::TokenBiasMap::new(); + bias.suppress_tokens(&[2]); + let biased_logits = mlxcel_core::sampling::apply_token_bias(logits.as_ref().unwrap(), &bias); + let biased = Gemma4MtpTargetAdapter::argmax_per_position(biased_logits.as_ref().unwrap()); + assert_eq!( + biased, + vec![3, 0, 3], + "suppressing the placeholder id must re-route only the leaking position" + ); +} + +#[test] +fn verify_empty_bias_is_identity_b1() { + // An empty `TokenBiasMap` (every non-multimodal model) must leave the + // per-position argmax bit-identical to the raw path. + let _runtime = crate::initialize_runtime(); + + let data: Vec = vec![ + 0.1, 0.2, 0.9, 0.3, // pos 0 -> 2 + 0.9, 0.1, 0.2, 0.3, // pos 1 -> 0 + 0.1, 0.2, 0.3, 0.9, // pos 2 -> 3 + ]; + let logits = mlxcel_core::from_slice_f32(&data, &[1, 3, 4]); + + let raw = Gemma4MtpTargetAdapter::argmax_per_position(logits.as_ref().unwrap()); + let empty = mlxcel_core::sampling::TokenBiasMap::new(); + let passthrough = mlxcel_core::sampling::apply_token_bias(logits.as_ref().unwrap(), &empty); + let biased = Gemma4MtpTargetAdapter::argmax_per_position(passthrough.as_ref().unwrap()); + assert_eq!( + raw, biased, + "empty bias must be a no-op over the verify argmax" + ); +} + // =========================================================================== // Batched MTP target adapter helper tests // @@ -192,6 +249,43 @@ fn batched_argmax_per_row_returns_b_by_width_ids() { assert_eq!(argmax, vec![vec![2, 0, 3], vec![1, 3, 0]]); } +#[test] +fn batched_verify_bias_suppresses_argmax_per_row() { + // issue #350: the batched verify step applies `sampler.token_bias` to the + // [B, width, vocab] logits before `argmax_per_row`. Confirm a [1, vocab] + // suppression mask broadcasts across the [B, width] grid and re-routes only + // the positions where the suppressed id would otherwise win, in any row. + let _runtime = crate::initialize_runtime(); + + // `[2, 3, 4]`. The placeholder id 2 is the raw winner at row 0 pos 0 and + // row 1 pos 2; every other position wins on a legitimate id. + let data: Vec = vec![ + // row 0 + 0.1, 0.2, 0.9, 0.7, // pos 0 -> raw 2, next-best 3 + 0.9, 0.1, 0.2, 0.3, // pos 1 -> 0 + 0.1, 0.2, 0.3, 0.9, // pos 2 -> 3 + // row 1 + 0.1, 0.9, 0.2, 0.3, // pos 0 -> 1 + 0.1, 0.2, 0.3, 0.9, // pos 1 -> 3 + 0.6, 0.1, 0.9, 0.3, // pos 2 -> raw 2, next-best 0 + ]; + let logits = mlxcel_core::from_slice_f32(&data, &[2, 3, 4]); + + let raw = Gemma4MtpBatchedTargetAdapter::argmax_per_row(logits.as_ref().unwrap(), 2, 3); + assert_eq!(raw, vec![vec![2, 0, 3], vec![1, 3, 2]]); + + let mut bias = mlxcel_core::sampling::TokenBiasMap::new(); + bias.suppress_tokens(&[2]); + let biased_logits = mlxcel_core::sampling::apply_token_bias(logits.as_ref().unwrap(), &bias); + let biased = + Gemma4MtpBatchedTargetAdapter::argmax_per_row(biased_logits.as_ref().unwrap(), 2, 3); + assert_eq!( + biased, + vec![vec![3, 0, 3], vec![1, 3, 0]], + "suppressing id 2 must re-route both rows' leaking positions and leave the rest" + ); +} + #[test] fn batched_slice_shared_kv_zero_rejected_is_identity() { let _runtime = crate::initialize_runtime(); diff --git a/src/models/mod.rs b/src/models/mod.rs index a24558700..1cd5d641b 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -20,6 +20,7 @@ mod detection; mod gemma3n_helpers; mod llama4_helpers; mod model_owned; +pub mod multimodal_placeholders; pub(crate) mod qwen_mrope_state; mod recurrent_snapshot; mod sanitize; @@ -167,6 +168,7 @@ pub use mixtral::MixtralModel; pub use molmo::MolmoModel; pub use molmo2::Molmo2Model; pub use moondream3::Moondream3Model; +pub use multimodal_placeholders::MultimodalPlaceholderTokens; pub use nemotron::NemotronModel; pub use nemotron_h::NemotronHModel; pub use nemotron_nas::NemotronNASModel; diff --git a/src/models/multimodal_placeholders.rs b/src/models/multimodal_placeholders.rs new file mode 100644 index 000000000..8d6a5b158 --- /dev/null +++ b/src/models/multimodal_placeholders.rs @@ -0,0 +1,87 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Multimodal placeholder token ids that must never appear in text output. +//! +//! Multimodal model configs carry `*_token_id` / `*_token_index` fields that +//! mark audio / image / video spans in the INPUT id stream so the runtime can +//! scatter encoded features into those positions. They are not real +//! vocabulary tokens for generation: if one becomes the argmax at a near-tie +//! greedy decode step it leaks into the text output (issue #350). +//! +//! [`MultimodalPlaceholderTokens`] is a family-agnostic collector: a model (or +//! its config) fills in whichever placeholder fields it has, and +//! [`MultimodalPlaceholderTokens::suppressed_ids`] returns the deduplicated, +//! sorted set of valid ids to mask to `-inf` in the output logits. New +//! multimodal families opt in by building this struct from their own config +//! and returning `suppressed_ids()` from +//! [`mlxcel_core::generate::LanguageModel::output_suppressed_token_ids`]. + +/// Reserved multimodal placeholder token ids for a single model. +/// +/// Every field is optional so a family only sets the markers it actually +/// defines (a vision-only model leaves the audio markers `None`, etc.). The +/// real EOS / BOS / normal text ids are deliberately NOT part of this struct: +/// only input-alignment placeholders belong here, so the derived suppression +/// set can never silence end-of-sequence detection. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct MultimodalPlaceholderTokens { + /// Audio soft-token placeholder (`audio_token_id`). + pub audio_token_id: Option, + /// Image soft-token placeholder (`image_token_id`). + pub image_token_id: Option, + /// Video soft-token placeholder (`video_token_id`). + pub video_token_id: Option, + /// Begin-of-audio marker (`boa_token_id`). + pub boa_token_id: Option, + /// Begin-of-image marker (`boi_token_id`). + pub boi_token_id: Option, + /// End-of-audio marker (`eoa_token_id` / `eoa_token_index`). + pub eoa_token_id: Option, + /// End-of-image marker (`eoi_token_id`). + pub eoi_token_id: Option, +} + +impl MultimodalPlaceholderTokens { + /// The deduplicated, ascending set of valid placeholder ids to suppress + /// in the output logits. + /// + /// `None` fields are skipped; negative ids (a sentinel for "unset" in some + /// configs) are dropped. The result is sorted and deduplicated so callers + /// get a stable, allocation-minimal set. Returns an empty vec when no + /// placeholder is defined, which keeps the suppression path zero-cost for + /// non-multimodal models. + pub fn suppressed_ids(&self) -> Vec { + let mut ids: Vec = [ + self.audio_token_id, + self.image_token_id, + self.video_token_id, + self.boa_token_id, + self.boi_token_id, + self.eoa_token_id, + self.eoi_token_id, + ] + .into_iter() + .flatten() + .filter(|&id| id >= 0) + .collect(); + ids.sort_unstable(); + ids.dedup(); + ids + } +} + +#[cfg(test)] +#[path = "multimodal_placeholders_tests.rs"] +mod tests; diff --git a/src/models/multimodal_placeholders_tests.rs b/src/models/multimodal_placeholders_tests.rs new file mode 100644 index 000000000..1e1e9fd45 --- /dev/null +++ b/src/models/multimodal_placeholders_tests.rs @@ -0,0 +1,101 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::MultimodalPlaceholderTokens; + +#[test] +fn empty_struct_yields_no_suppressed_ids() { + assert!( + MultimodalPlaceholderTokens::default() + .suppressed_ids() + .is_empty() + ); +} + +#[test] +fn gemma4_unified_12b_config_values_produce_all_seven_placeholders() { + // The exact ids from models/gemma-4-12b-it-4bit/config.json (issue #350). + let tokens = MultimodalPlaceholderTokens { + audio_token_id: Some(258_881), + image_token_id: Some(258_880), + video_token_id: Some(258_884), + boa_token_id: Some(256_000), + boi_token_id: Some(255_999), + eoa_token_id: Some(258_883), + eoi_token_id: Some(258_882), + }; + // Sorted, deduplicated. + assert_eq!( + tokens.suppressed_ids(), + vec![ + 255_999, 256_000, 258_880, 258_881, 258_882, 258_883, 258_884 + ] + ); +} + +#[test] +fn real_eos_ids_are_never_in_the_suppressed_set() { + // The gemma-4-12b config also carries eos_token_id = [1, 106, 50]; those + // must stay generatable. The placeholder struct intentionally has no EOS + // field, so the derived set can never contain them. + let tokens = MultimodalPlaceholderTokens { + audio_token_id: Some(258_881), + image_token_id: Some(258_880), + video_token_id: Some(258_884), + boa_token_id: Some(256_000), + boi_token_id: Some(255_999), + eoa_token_id: Some(258_883), + eoi_token_id: Some(258_882), + }; + let suppressed = tokens.suppressed_ids(); + for eos in [1, 106, 50] { + assert!( + !suppressed.contains(&eos), + "real EOS id {eos} must not be suppressed" + ); + } +} + +#[test] +fn unset_fields_are_skipped() { + // A vision-only model leaves the audio markers unset. + let tokens = MultimodalPlaceholderTokens { + image_token_id: Some(258_880), + boi_token_id: Some(255_999), + eoi_token_id: Some(258_882), + ..Default::default() + }; + assert_eq!(tokens.suppressed_ids(), vec![255_999, 258_880, 258_882]); +} + +#[test] +fn negative_sentinel_ids_are_dropped() { + let tokens = MultimodalPlaceholderTokens { + image_token_id: Some(258_880), + audio_token_id: Some(-1), + ..Default::default() + }; + assert_eq!(tokens.suppressed_ids(), vec![258_880]); +} + +#[test] +fn duplicate_ids_are_deduplicated() { + // Some configs alias two markers to the same id; the set stays unique. + let tokens = MultimodalPlaceholderTokens { + image_token_id: Some(258_880), + video_token_id: Some(258_880), + ..Default::default() + }; + assert_eq!(tokens.suppressed_ids(), vec![258_880]); +} diff --git a/src/server/batch/scheduler.rs b/src/server/batch/scheduler.rs index 17748c718..4fc45b725 100644 --- a/src/server/batch/scheduler.rs +++ b/src/server/batch/scheduler.rs @@ -301,6 +301,17 @@ pub struct BatchScheduler { /// Empty map = bit-exact baseline path (no sampling change, no alloc). token_bias: TokenBiasMap, + /// Reserved output-illegal token ids for the loaded model (issue #350). + /// + /// Multimodal models (e.g. `gemma4_unified`) carry audio / image / video + /// span placeholder ids that must never be emitted as text. Resolved once + /// at construction from [`LanguageModel::output_suppressed_token_ids`] and + /// force-masked to `-inf` on every sequence's sampling config at + /// [`Self::enqueue_request`] time, unconditionally (suppression always + /// wins over any per-request bias). Empty for non-multimodal models, which + /// keeps the enqueue path zero-cost. + model_output_suppressed: Vec, + // -- — thinking-token budget -- /// Server-wide default thinking-token budget. `None` means unrestricted. /// Per-request `thinking_budget_tokens` overrides this at enqueue time. @@ -825,6 +836,10 @@ impl BatchScheduler { // batch and the prefill queue so requests can be queued while another // sequence is generating. let pool_capacity = max_batch_size + max_queue_depth; + // issue #350: resolve the model's reserved output-illegal placeholder + // ids once, before `model` is moved into the scheduler. Empty for + // non-multimodal models (zero cost on the enqueue path). + let model_output_suppressed = model.output_suppressed_token_ids(); Self { cache_pool: CachePool::new(pool_capacity), prefill_queue: PrefillQueue::with_capacity(max_queue_depth), @@ -847,6 +862,7 @@ impl BatchScheduler { crate::vision::feature_cache::DEFAULT_VISION_CACHE_SIZE, )), token_bias: TokenBiasMap::default(), + model_output_suppressed, reasoning_budget: None, thinking_token_ids: None, prompt_cache: None, @@ -2299,6 +2315,19 @@ impl BatchScheduler { sampling.token_bias = self.token_bias.clone(); } + // issue #350: force-suppress the model's reserved multimodal + // placeholder tokens (audio / image / video span markers) on every + // sequence's output logits. Applied after the lang-bias merge and + // unconditionally, so suppression always wins over any per-request + // bias and a placeholder id can never become the sampled argmax. A + // no-op (and zero alloc) for non-multimodal models whose suppressed + // set is empty. + if !self.model_output_suppressed.is_empty() { + sampling + .token_bias + .suppress_tokens(&self.model_output_suppressed); + } + let is_multimodal = !images.is_empty() || !audio.is_empty() || !videos.is_empty(); // Experimental VLM prompt-prefix cache sharing (#124 step c). Off by diff --git a/src/vision/gemma4_unified.rs b/src/vision/gemma4_unified.rs index 91272fd43..1e2f8f809 100644 --- a/src/vision/gemma4_unified.rs +++ b/src/vision/gemma4_unified.rs @@ -597,6 +597,25 @@ impl LanguageModel for Gemma4UnifiedModel { self.text_model.eos_token_ids_value() } + /// Suppress the reserved audio / image / video span markers in the output + /// logits (issue #350). These ids are parsed from the checkpoint config + /// (`audio_token_id`, `image_token_id`, ...) and stored on the model; they + /// are input-alignment placeholders that must never be sampled as text. + /// The real EOS ids are not included here, so end-of-sequence detection is + /// unaffected. + fn output_suppressed_token_ids(&self) -> Vec { + crate::models::MultimodalPlaceholderTokens { + audio_token_id: Some(self.audio_token_id), + image_token_id: Some(self.image_token_id), + video_token_id: Some(self.video_token_id), + boa_token_id: Some(self.boa_token_id), + boi_token_id: Some(self.boi_token_id), + eoa_token_id: Some(self.eoa_token_id), + eoi_token_id: Some(self.eoi_token_id), + } + .suppressed_ids() + } + fn supports_padded_prefill(&self) -> bool { false } diff --git a/src/vision/gemma4_unified_config.rs b/src/vision/gemma4_unified_config.rs index 8a17b7307..d878d1a74 100644 --- a/src/vision/gemma4_unified_config.rs +++ b/src/vision/gemma4_unified_config.rs @@ -25,6 +25,8 @@ use serde::Deserialize; use serde_json::Value; +use crate::models::multimodal_placeholders::MultimodalPlaceholderTokens; + fn default_rms_norm_eps() -> f32 { 1e-6 } @@ -171,6 +173,28 @@ impl Gemma4UnifiedConfig { .or(self.eoa_token_index) .unwrap_or_else(default_eoa_token_index) } + + /// The reserved multimodal placeholder token ids (audio / image / video + /// span markers) that must never appear in generated text output + /// (issue #350). + /// + /// These are input-alignment placeholders; the runtime scatters encoded + /// features into them during prefill, but they are illegal as generation + /// output. The returned [`MultimodalPlaceholderTokens`] is fed through + /// [`MultimodalPlaceholderTokens::suppressed_ids`] and masked to `-inf` at + /// every decode step. Real EOS ids (`eos_token_id`) are intentionally + /// excluded so end-of-sequence detection is unaffected. + pub fn placeholder_tokens(&self) -> MultimodalPlaceholderTokens { + MultimodalPlaceholderTokens { + audio_token_id: Some(self.audio_token_id), + image_token_id: Some(self.image_token_id), + video_token_id: Some(self.video_token_id), + boa_token_id: Some(self.boa_token_id), + boi_token_id: Some(self.boi_token_id), + eoa_token_id: Some(self.resolve_eoa_token_id()), + eoi_token_id: Some(self.eoi_token_id), + } + } } #[cfg(test)] diff --git a/src/vision/gemma4_unified_config_tests.rs b/src/vision/gemma4_unified_config_tests.rs index a0be33487..8c8df7ef3 100644 --- a/src/vision/gemma4_unified_config_tests.rs +++ b/src/vision/gemma4_unified_config_tests.rs @@ -173,6 +173,56 @@ fn audio_config_optional() { assert!(cfg.audio_config.is_none()); } +#[test] +fn placeholder_tokens_cover_all_seven_markers() { + // issue #350: the config's placeholder ids feed the output-suppression set. + let cfg: Gemma4UnifiedConfig = serde_json::from_value(reference_config()).unwrap(); + let suppressed = cfg.placeholder_tokens().suppressed_ids(); + // audio 258881, image 258880, video 258884, boa 256000, boi 255999, + // eoa 258883 (via eoa_token_index fallback), eoi 258882, returned sorted. + assert_eq!( + suppressed, + vec![ + 255_999, 256_000, 258_880, 258_881, 258_882, 258_883, 258_884 + ] + ); +} + +#[test] +fn output_suppression_masks_placeholders_but_not_eos() { + // The end-to-end shape used at generation time: derive the suppressed set + // from the gemma4_unified config and force it into a TokenBiasMap. Every + // placeholder id must become -inf; the real EOS ids ([1, 106, 50] in the + // gemma-4-12b checkpoint) must stay untouched so end-of-sequence detection + // and normal text generation are unaffected (issue #350 acceptance). + let mut value = reference_config(); + value["eos_token_id"] = serde_json::json!([1, 106, 50]); + let cfg: Gemma4UnifiedConfig = serde_json::from_value(value).unwrap(); + + let mut bias = mlxcel_core::TokenBiasMap::default(); + assert!(bias.is_empty(), "baseline starts empty (zero-cost path)"); + bias.suppress_tokens(&cfg.placeholder_tokens().suppressed_ids()); + + for placeholder in [ + 255_999, 256_000, 258_880, 258_881, 258_882, 258_883, 258_884, + ] { + let b = bias + .get(&placeholder) + .copied() + .unwrap_or_else(|| panic!("placeholder {placeholder} must be in the bias map")); + assert!( + b.is_infinite() && b.is_sign_negative(), + "placeholder {placeholder} must be masked to -inf, got {b}" + ); + } + for eos in [1, 106, 50] { + assert!( + bias.get(&eos).is_none(), + "real EOS id {eos} must never be suppressed" + ); + } +} + #[test] fn vision_config_defaults_apply() { // Minimal vision_config falls back to documented defaults.