Repository navigation
fix(models): suppress gemma4_unified multimodal placeholder tokens - #351
Conversation
gemma4_unified (e.g. gemma-4-12b-it-4bit) reserves audio/image/video span placeholder ids (audio_token_id 258881, image_token_id 258880, boi_token_id 255999, boa_token_id 256000, eoi_token_id 258882, eoa 258883, video_token_id 258884) purely for INPUT alignment. mlxcel did not mask them in the output logits, so at a near-tie greedy step one could become the argmax and leak into generated text (a stray <audio|> mid-sentence). Through /v1/chat/completions the server strips special tokens, so a leak-heavy generation returned EMPTY content with finish_reason=length. Add a generic, opt-in suppression mechanism and wire it for gemma4_unified only: - LanguageModel gains output_suppressed_token_ids() (default empty), the per-model "never generate these" set, mirroring eos_token_ids(). Non-multimodal models return empty and pay zero cost. - TokenBiasMap::suppress_tokens(ids) force-masks each id to f32::NEG_INFINITY (suppression always wins over any prior finite bias). Empty slice is a no-op, preserving the bit-exact baseline. - models::MultimodalPlaceholderTokens is a family-agnostic collector: given the config *_token_id / *_token_index fields it returns the deduplicated, non-negative id set. Gemma4UnifiedConfig::placeholder_tokens() builds it from the parsed config, and the Gemma4UnifiedModel LanguageModel impl returns suppressed_ids() from its stored ids. The set is derived from config, never hardcoded, and excludes the real EOS ids (eos_token_id [1,106,50]), so end-of-sequence detection and normal text are unaffected. Applied at generation time on both decode paths: the CLI generate path merges it into the token-bias map at the top of run_generation_mode (covering the standard, VLM-embedding, and speculative sub-paths), and the server scheduler resolves it once in with_config and force-applies it on every sequence's sampling config in enqueue_request (covering classic decode and the MTP burst, which both sample through sample_token_optimized). Both apply after any lang-bias merge so suppression cannot be displaced by a per-request bias. Tests: TokenBiasMap::suppress_tokens (-inf + zero softmax probability, empty no-op); MultimodalPlaceholderTokens derivation (dedup, negative-id drop, unset-field skip, the seven gemma-4-12b ids); and a gemma4_unified config acceptance test asserting the placeholder ids are masked to -inf while EOS ids [1,106,50] stay untouched.
inureyes
added
type:bug
Implementation Review SummaryIntent
Findings Addressed
Remaining Items (non-blocking, out of #350 scope)
Verification
Real-model validation (hash-table prompt non-empty/coherent via CLI and |
PR #351 masks gemma4_unified multimodal placeholder ids to -inf via the per-step token bias, but two generation paths bypassed that bias and could still leak a placeholder into generated text (issue #350). The server MTP speculative-burst verify step derived target tokens with a raw argmax (Gemma4MtpTargetAdapter::verify_forward, the deferred argmax_from_hidden_positions path, and the batched verify_forward_batched), so only the first bonus token (sampled through sample_token_optimized) was suppressed and every later token could leak. The suppression set is already present in the sampler forwarded from enqueue_request through MtpGenerator, so verify now applies apply_token_bias to the logits before the per-position argmax (and before logprob extraction, so the emitted token and its logprob stay consistent). An empty bias short-circuits, keeping non-multimodal models bit-exact. The interactive chat path (mlxcel generate with no -p, and mlxcel run) built its CxxGenerator without merging the model's suppressed set, so chat replies could leak. run_chat now attaches the suppression bias via with_token_bias, mirroring run_generation_mode. apply_token_bias is promoted from pub(crate) to pub so the MTP target adapter in the binary crate can reuse the existing masking helper. Tests: B=1 verify argmax re-routes a leaking position and is a no-op on an empty bias; batched verify argmax re-routes per row. Validated with cargo check (lib + bins, release), targeted unit tests, and clippy on both crates. Real-model MTP/CLI/server validation is pending.
Implementation Review Update (supersedes the Verification section of my earlier comment)A deeper pass plus the Codex reviewer found two CRITICAL leak paths that my first comment incorrectly marked as covered. Both are now fixed on this branch. CRITICAL findings (fixed)
Verification
Non-blocking (unchanged)
|
Add a comment to `argmax_logits_to_array` in the DFlash verify round-loop noting that the raw argmax carries no token bias. Currently safe because only gemma4_unified has a non-empty suppressed-id set and it runs on the MTP path, not DFlash. If a future model adopts DFlash verify with a non-empty suppressed-id set, the argmax must be preceded by `apply_token_bias` to prevent reintroducing the #350 leak. Also reformats a long `assert_eq!` in gemma4_mtp_target_tests.rs to satisfy `cargo fmt`.
PR FinalizationTestsNo gaps found. The existing 16 suppression tests cover the suppressed-set derivation, EOS exclusion, NEG_INFINITY masking, and the 3 MTP verify-path cases. No new tests were needed. DocumentationNo doc updates warranted. This is a correctness fix with no user-facing surface; there is no existing known-issues or model-support section that applies. Lint / Format
Latent-gap commentAdded a comment to Commit: |
Real-model validation (orchestrator)Placeholder suppression validated on M1 Ultra:
Root-cause note on the earlier "empty content" symptom (it is NOT a placeholder issue): gemma-4-12b-it is a thinking model by design. The server detects its pr-reviewer fixed two CRITICAL leak paths it found (MTP speculative verify and interactive chat), pr-security-checker confirmed MTP exactness consistency (verify applies the same suppression as classic), and 19 tests pass (16 suppression + 3 MTP verify). The #350 placeholder leak is fixed and validated. Minor observations (out of #350 scope, not regressions): gemma-4-12b-it thinking is verbose (a three-sentence request did not finish within 800 tokens), and |
…ablup#351) * fix(models): suppress gemma4_unified multimodal placeholder tokens gemma4_unified (e.g. gemma-4-12b-it-4bit) reserves audio/image/video span placeholder ids (audio_token_id 258881, image_token_id 258880, boi_token_id 255999, boa_token_id 256000, eoi_token_id 258882, eoa 258883, video_token_id 258884) purely for INPUT alignment. mlxcel did not mask them in the output logits, so at a near-tie greedy step one could become the argmax and leak into generated text (a stray <audio|> mid-sentence). Through /v1/chat/completions the server strips special tokens, so a leak-heavy generation returned EMPTY content with finish_reason=length. Add a generic, opt-in suppression mechanism and wire it for gemma4_unified only: - LanguageModel gains output_suppressed_token_ids() (default empty), the per-model "never generate these" set, mirroring eos_token_ids(). Non-multimodal models return empty and pay zero cost. - TokenBiasMap::suppress_tokens(ids) force-masks each id to f32::NEG_INFINITY (suppression always wins over any prior finite bias). Empty slice is a no-op, preserving the bit-exact baseline. - models::MultimodalPlaceholderTokens is a family-agnostic collector: given the config *_token_id / *_token_index fields it returns the deduplicated, non-negative id set. Gemma4UnifiedConfig::placeholder_tokens() builds it from the parsed config, and the Gemma4UnifiedModel LanguageModel impl returns suppressed_ids() from its stored ids. The set is derived from config, never hardcoded, and excludes the real EOS ids (eos_token_id [1,106,50]), so end-of-sequence detection and normal text are unaffected. Applied at generation time on both decode paths: the CLI generate path merges it into the token-bias map at the top of run_generation_mode (covering the standard, VLM-embedding, and speculative sub-paths), and the server scheduler resolves it once in with_config and force-applies it on every sequence's sampling config in enqueue_request (covering classic decode and the MTP burst, which both sample through sample_token_optimized). Both apply after any lang-bias merge so suppression cannot be displaced by a per-request bias. Tests: TokenBiasMap::suppress_tokens (-inf + zero softmax probability, empty no-op); MultimodalPlaceholderTokens derivation (dedup, negative-id drop, unset-field skip, the seven gemma-4-12b ids); and a gemma4_unified config acceptance test asserting the placeholder ids are masked to -inf while EOS ids [1,106,50] stay untouched. * fix(inference): suppress placeholders on MTP verify and chat paths PR lablup#351 masks gemma4_unified multimodal placeholder ids to -inf via the per-step token bias, but two generation paths bypassed that bias and could still leak a placeholder into generated text (issue lablup#350). The server MTP speculative-burst verify step derived target tokens with a raw argmax (Gemma4MtpTargetAdapter::verify_forward, the deferred argmax_from_hidden_positions path, and the batched verify_forward_batched), so only the first bonus token (sampled through sample_token_optimized) was suppressed and every later token could leak. The suppression set is already present in the sampler forwarded from enqueue_request through MtpGenerator, so verify now applies apply_token_bias to the logits before the per-position argmax (and before logprob extraction, so the emitted token and its logprob stay consistent). An empty bias short-circuits, keeping non-multimodal models bit-exact. The interactive chat path (mlxcel generate with no -p, and mlxcel run) built its CxxGenerator without merging the model's suppressed set, so chat replies could leak. run_chat now attaches the suppression bias via with_token_bias, mirroring run_generation_mode. apply_token_bias is promoted from pub(crate) to pub so the MTP target adapter in the binary crate can reuse the existing masking helper. Tests: B=1 verify argmax re-routes a leaking position and is a no-op on an empty bias; batched verify argmax re-routes per row. Validated with cargo check (lib + bins, release), targeted unit tests, and clippy on both crates. Real-model MTP/CLI/server validation is pending. * chore: add lablup#350 latent-gap note to DFlash argmax; fix fmt Add a comment to `argmax_logits_to_array` in the DFlash verify round-loop noting that the raw argmax carries no token bias. Currently safe because only gemma4_unified has a non-empty suppressed-id set and it runs on the MTP path, not DFlash. If a future model adopts DFlash verify with a non-empty suppressed-id set, the argmax must be preceded by `apply_token_bias` to prevent reintroducing the lablup#350 leak. Also reformats a long `assert_eq!` in gemma4_mtp_target_tests.rs to satisfy `cargo fmt`.
Summary
gemma4_unified(e.g.gemma-4-12b-it-4bit) reserves multimodal placeholder ids (audio/image/video span markers) for INPUT alignment only, but mlxcel did not mask them in the output logits. At a near-tie greedy decode step one could become the argmax and leak into generated text (a stray<audio|>mid-sentence on the "hash table" prompt). Through/v1/chat/completionsthe server strips special tokens, so a leak-heavy generation returned EMPTYcontentwithfinish_reason=length. This masks those ids to-infat every decode step via the existing per-step token-bias mechanism, with a generic helper so other multimodal families can opt in later.What changed
LanguageModel::output_suppressed_token_ids()(new trait method, default empty): the per-model "never generate these" set, mirroringeos_token_ids(). Non-multimodal models inherit the empty default and pay zero cost.TokenBiasMap::suppress_tokens(ids): force-masks each id tof32::NEG_INFINITY(suppression always wins over any prior finite bias); empty slice is a no-op that preserves the bit-exact baseline.models::MultimodalPlaceholderTokens: a family-agnostic collector that, given the config*_token_id/*_token_indexfields, returns the deduplicated, non-negative id set.Gemma4UnifiedConfig::placeholder_tokens()builds it from the parsed config, and theGemma4UnifiedModelLanguageModelimpl returnssuppressed_ids()from its stored ids.LoadedModeldelegates the new trait method.run_generation_modemerges the set into the token-bias map before any generator is built; serverBatchSchedulerresolves it once inwith_configand force-applies it on every sequence's sampling config inenqueue_request.Where suppression is applied (both decode paths)
generate:run_generation_modemergesmodel.output_suppressed_token_ids()into the token-bias map at the top, so it reaches the standard, VLM-embedding, and speculative sub-paths (all consume that map viawith_token_bias).enqueue_requestapplies the resolved set to each sequence'sSamplingConfig.token_biasafter the lang-bias merge and unconditionally, so it covers classic decode and the MTP burst (both sample throughsample_token_optimized, which appliesconfig.token_bias). Suppression is added last, so a per-request bias can never displace it.Derivation and safety
The suppressed set is derived from config (
audio_token_id,image_token_id,video_token_id,boa_token_id,boi_token_id,eoa_token_id/eoa_token_index,eoi_token_id), never hardcoded as raw integers in the wiring. The real EOS ids (eos_token_id[1, 106, 50]) are not part of the placeholder struct, so end-of-sequence detection and normal text generation are unaffected. Non-multimodal models return an empty set, soapply_token_biaskeeps its zero-overhead short-circuit.Test plan
cargo test --release -p mlxcel-core sampling::tests::suppress(2 tests:-inf+ zero softmax probability, empty no-op)cargo test --release -p mlxcel --lib -- multimodal_placeholders gemma4_unified_config(14 tests incl.output_suppression_masks_placeholders_but_not_eos: placeholder ids masked to-inf, EOS[1,106,50]untouched)cargo clippy -p mlxcel-core --lib --tests -- -D warningsandcargo clippy -p mlxcel --lib --tests -- -D warnings(clean)cargo fmt --check(clean)mlxcel generate -t 0AND/v1/chat/completions; no regression on clean prompts ("list three primary colors", "write one sentence about the moon").Closes #350