From b5f1b4c99fc0dcb4774e53d297074be3f6d0f527 Mon Sep 17 00:00:00 2001 From: key4ng Date: Tue, 28 Jul 2026 23:05:19 -0700 Subject: [PATCH 1/3] fix(kimi-k3): build the reference media wrapper during prompt expansion K3's image prompt is one wrapper per image -- `<|media_begin|>image {w}x{h}<|media_content|><|media_pad|><|media_end|>` -- where the dimensions are the pre-resize decoded size. SMG emitted only the bare `<|media_pad|>` run because K3 was routed to the K2.5 registry spec, whose chat template emits its own dimensionless wrapper. Every K3 image prompt was therefore 9 tokens short of the reference and carried no dimensions at all. The wrapper cannot be built while rendering: the chat template runs before any media is fetched, so the sizes do not exist yet. Give K3 its own spec that builds the block in `prompt_replacements` from the sizes the preprocessor reports -- the same place vLLM builds it (`kimi_k3.py::_get_prompt_updates`). `with_feature_span` keeps the encoder-feature positions on the pad run alone; the wrapper is text. Narrow the K2.5 matcher to K2.5 and register K3 ahead of it. The two families share the MoonViT transport layout but not a prompt shape, and implicit sharing is what produced the pixel-pipeline divergence fixed in PR #1984. Verified against the checkpoint: for 1024x768, 4000x3000, 224x448, 3x4 and 512x512 the token-by-token build is byte-identical to the reference's own encoding of `make_image_prompt`, since the media tokens are hard segment boundaries for the tiktoken encoder. Signed-off-by: key4ng --- crates/multimodal/src/registry/kimi_k25.rs | 45 ++- crates/multimodal/src/registry/kimi_k3.rs | 342 +++++++++++++++++++++ crates/multimodal/src/registry/mod.rs | 32 +- crates/multimodal/src/registry/traits.rs | 2 + 4 files changed, 392 insertions(+), 29 deletions(-) create mode 100644 crates/multimodal/src/registry/kimi_k3.rs diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs index 994e5cf06..4dff24034 100644 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ b/crates/multimodal/src/registry/kimi_k25.rs @@ -28,16 +28,17 @@ impl ModelProcessorSpec for KimiK25VisionSpec { } fn matches(&self, metadata: &ModelMetadata) -> bool { - // Kimi-K3 uses the same `<|media_pad|>` placeholder - // (media_placeholder_token_id 163605), patchification layout and - // prompt-replacement shape, so it shares this spec. Note that the two - // do *not* share a pixel pipeline — see `vision::processors::kimi_k3` - // for the patch budget and alpha-compositing differences. + // K2.5 only. Kimi-K3 shares the `<|media_pad|>` fill token + // (media_placeholder_token_id 163605) and the patchification layout, but + // its prompt carries a per-image `<|media_begin|>image {w}x{h}…` wrapper + // that K2.5's chat template emits itself — see `registry::kimi_k3`. The + // two do not share a pixel pipeline either; `vision::processors::kimi_k3` + // has the patch budget and alpha-compositing differences. let id = metadata.model_id.to_ascii_lowercase(); - id.contains("kimi") && (id.contains("k2") || id.contains("k3")) + id.contains("kimi") && id.contains("k2") || metadata .config_model_type() - .is_some_and(|mt| mt == "kimi_k25" || mt == "kimi_k3") + .is_some_and(|mt| mt == "kimi_k25") } fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { @@ -126,31 +127,25 @@ mod tests { } #[test] - fn kimi_k3_matches_model_id_and_model_type() { + fn kimi_k3_does_not_use_the_k25_spec() { + // K3's prompt carries per-image dimensions that this spec cannot emit, + // so it must route to `kimi_k3` by model_id and by model_type alike. let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); - // Match by model_id containing kimi + k3. let config = json!({ "model_type": "kimi_k3", "media_placeholder_token_id": 163605 }); - let metadata = ModelMetadata { - model_id: "moonshotai/Kimi-K3", - tokenizer: &tokenizer, - config: &config, - }; let registry = ModelRegistry::new(); - let spec = registry - .lookup(&metadata) - .expect("kimi_k3 -> kimi_k25 spec"); - assert_eq!(spec.name(), "kimi_k25"); - // Also match by model_type alone (id without a k3 hint). - let metadata_by_type = ModelMetadata { - model_id: "internal/checkpoint-final", - tokenizer: &tokenizer, - config: &config, - }; - assert!(registry.lookup(&metadata_by_type).is_some()); + for model_id in ["moonshotai/Kimi-K3", "internal/checkpoint-final"] { + let metadata = ModelMetadata { + model_id, + tokenizer: &tokenizer, + config: &config, + }; + let spec = registry.lookup(&metadata).expect("kimi_k3 spec"); + assert_eq!(spec.name(), "kimi_k3", "model_id {model_id}"); + } } #[test] diff --git a/crates/multimodal/src/registry/kimi_k3.rs b/crates/multimodal/src/registry/kimi_k3.rs new file mode 100644 index 000000000..cc53d1798 --- /dev/null +++ b/crates/multimodal/src/registry/kimi_k3.rs @@ -0,0 +1,342 @@ +use std::collections::HashMap; + +use llm_tokenizer::Encoding; +use serde_json::{json, Value}; + +use crate::{ + encoder_inputs::PreprocessedEncoderInputs, + registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, + types::{FieldLayout, Modality, PromptReplacement, TokenId}, +}; + +/// Structural tokens wrapping one Kimi-K3 image, from the checkpoint's +/// `kimi_k3_vision_processing.py::make_image_prompt`: +/// `<|media_begin|>image {width}x{height}<|media_content|><|media_pad|><|media_end|>`. +const MEDIA_BEGIN: &str = "<|media_begin|>"; +const MEDIA_CONTENT: &str = "<|media_content|>"; +const MEDIA_END: &str = "<|media_end|>"; + +/// Kimi-K3. +/// +/// K3 shares K2.5's NaViT transport layout and `<|media_pad|>` fill token, but +/// not its prompt shape: the K3 checkpoint ships no jinja chat template, and its +/// Python renderer splices a per-image `<|media_begin|>image {w}x{h} +/// <|media_content|>…<|media_end|>` block into the raw text at each +/// `<|kimi_image_placeholder|>` marker (`kimi_k3_processor.py::update_raw_text`). +/// K2.5's template emits the wrapper itself and carries no dimensions, so the two +/// families need separate specs. +/// +/// SMG cannot build that block while rendering — the chat template runs before +/// any media is fetched, so the image dimensions do not exist yet. It is built +/// here instead, from the pre-resize sizes the preprocessor reports, which is +/// where vLLM builds it too (`kimi_k3.py::_get_prompt_updates`). The renderer +/// emits a bare `<|media_pad|>` anchor per image and this replacement expands it +/// into the full wrapper. +pub(super) struct KimiK3VisionSpec; + +impl KimiK3VisionSpec { + /// The repeated pad token (`<|media_pad|>`) — `media_placeholder_token_id` in config. + fn pad_token_id(metadata: &ModelMetadata) -> RegistryResult { + metadata + .config_u32(&["media_placeholder_token_id"]) + .map(|v| v as TokenId) + .ok_or_else(|| ModelRegistryError::MissingConfigField { + field: "media_placeholder_token_id".to_string(), + }) + } + + /// Encode ordinary text (no special tokens) into token ids. + /// + /// The dimension text sits between two special tokens, which are hard + /// segment boundaries for the tiktoken encoder, so encoding it on its own + /// yields the same ids the reference gets from encoding the whole block. + fn encode_plain_text(metadata: &ModelMetadata, text: &str) -> RegistryResult> { + let encoding = metadata.tokenizer.encode(text, false).map_err(|_| { + ModelRegistryError::TextEncodingFailed { + spec: "kimi_k3", + text: text.to_string(), + } + })?; + Ok(match encoding { + Encoding::Hf(inner) => inner.get_ids().iter().map(|&id| id as TokenId).collect(), + Encoding::Plain(ids) | Encoding::Tiktoken(ids) => { + ids.into_iter().map(|id| id as TokenId).collect() + } + }) + } +} + +impl ModelProcessorSpec for KimiK3VisionSpec { + fn name(&self) -> &'static str { + "kimi_k3" + } + + fn matches(&self, metadata: &ModelMetadata) -> bool { + let id = metadata.model_id.to_ascii_lowercase(); + id.contains("kimi") && id.contains("k3") + || metadata + .config_model_type() + .is_some_and(|mt| mt == "kimi_k3") + } + + fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok("<|media_pad|>".to_string()) + } + + fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { + Self::pad_token_id(metadata) + } + + fn modality_limits( + &self, + _metadata: &ModelMetadata, + ) -> RegistryResult> { + Ok(HashMap::from([(Modality::Image, 10)])) + } + + fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(json!({})) + } + + fn prompt_replacements( + &self, + metadata: &ModelMetadata, + preprocessed: &PreprocessedEncoderInputs, + ) -> RegistryResult> { + let pad_token_id = Self::pad_token_id(metadata)?; + let placeholder_token = self.placeholder_token(metadata)?; + let media_begin = metadata.token_id(MEDIA_BEGIN)?; + let media_content = metadata.token_id(MEDIA_CONTENT)?; + let media_end = metadata.token_id(MEDIA_END)?; + + // MoonViT reports `item_sizes` as the decoded `(width, height)` before any + // resize, which is exactly the pair the reference prints. The caller + // checks both vectors against the media count, so a short zip here would + // already have been rejected upstream. + preprocessed + .feature_token_counts + .iter() + .zip(&preprocessed.item_sizes) + .map(|(&num_tokens, &(width, height))| { + let dims = Self::encode_plain_text(metadata, &format!("image {width}x{height}"))?; + let mut tokens = Vec::with_capacity(dims.len() + num_tokens + 3); + tokens.push(media_begin); + tokens.extend(dims); + tokens.push(media_content); + // Only the pad run holds encoder features; the wrapper is text. + let feature_offset = tokens.len(); + tokens.extend(std::iter::repeat_n(pad_token_id, num_tokens)); + tokens.push(media_end); + + Ok( + PromptReplacement::sequence(Modality::Image, &placeholder_token, tokens) + .with_feature_span(feature_offset, num_tokens), + ) + }) + .collect() + } + + fn field_layouts(&self) -> HashMap { + // MoonViT patchification, same transport layout as K2.5: + // encoder_input is [total_patches, patch_features], split by patches_per_image. + // grid_thws is [num_images, 3] with (temporal, height, width) grid dimensions. + HashMap::from([ + ( + "pixel_values".to_string(), + FieldLayout::flat("patches_per_image"), + ), + ("grid_thws".to_string(), FieldLayout::Batched), + ("patches_per_image".to_string(), FieldLayout::Batched), + ]) + } + + fn keep_on_cpu_keys(&self) -> Vec { + vec!["grid_thws".to_string()] + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use crate::{ + encoder_inputs::PreprocessedEncoderInputs, + registry::{test_helpers::*, ModelMetadata, ModelRegistry}, + types::{Modality, PlaceholderRange, TokenId}, + }; + + /// Token ids the K3 checkpoint actually assigns to the wrapper. + const MEDIA_BEGIN_ID: u32 = 163602; + const MEDIA_CONTENT_ID: u32 = 163603; + const MEDIA_END_ID: u32 = 163604; + const MEDIA_PAD_ID: u32 = 163605; + /// Offset of the byte encoder used by `TestTokenizer::with_byte_encoder`, + /// chosen so plain-text ids cannot collide with the media token ids. + const TEXT_BASE: u32 = 1000; + + fn k3_tokenizer() -> TestTokenizer { + TestTokenizer::new(&[ + ("<|media_begin|>", MEDIA_BEGIN_ID), + ("<|media_content|>", MEDIA_CONTENT_ID), + ("<|media_end|>", MEDIA_END_ID), + ("<|media_pad|>", MEDIA_PAD_ID), + ]) + .with_byte_encoder(TEXT_BASE) + } + + fn k3_config() -> serde_json::Value { + json!({ + "model_type": "kimi_k3", + "media_placeholder_token_id": MEDIA_PAD_ID, + }) + } + + /// `(width, height)` per item, matching MoonViT's `item_sizes` contract. + fn preprocessed( + sizes: &[(u32, u32)], + feature_token_counts: &[usize], + ) -> PreprocessedEncoderInputs { + PreprocessedEncoderInputs::new( + ndarray::Array4::::zeros((1, 3, 14, 14)), + feature_token_counts.to_vec(), + sizes.to_vec(), + ) + } + + fn text_ids(text: &str) -> Vec { + text.bytes() + .map(|b| (TEXT_BASE + u32::from(b)) as TokenId) + .collect() + } + + #[test] + fn kimi_k3_matches_model_id_and_model_type() { + let tokenizer = k3_tokenizer(); + let config = k3_config(); + let registry = ModelRegistry::new(); + + let metadata = ModelMetadata { + model_id: "moonshotai/Kimi-K3", + tokenizer: &tokenizer, + config: &config, + }; + assert_eq!( + registry.lookup(&metadata).expect("k3 spec").name(), + "kimi_k3" + ); + + // Also match by model_type alone (id without a k3 hint). + let metadata_by_type = ModelMetadata { + model_id: "internal/checkpoint-final", + tokenizer: &tokenizer, + config: &config, + }; + assert_eq!( + registry.lookup(&metadata_by_type).expect("k3 spec").name(), + "kimi_k3" + ); + } + + #[test] + fn kimi_k3_emits_the_reference_media_wrapper() { + let tokenizer = k3_tokenizer(); + let config = k3_config(); + let metadata = ModelMetadata { + model_id: "moonshotai/Kimi-K3", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).expect("k3 spec"); + + let replacements = spec + .prompt_replacements(&metadata, &preprocessed(&[(1024, 768)], &[4])) + .unwrap(); + + assert_eq!(replacements.len(), 1); + let rep = &replacements[0]; + assert_eq!(rep.modality, Modality::Image); + assert_eq!(rep.placeholder_token, "<|media_pad|>"); + + let mut expected = vec![MEDIA_BEGIN_ID as TokenId]; + expected.extend(text_ids("image 1024x768")); + expected.push(MEDIA_CONTENT_ID as TokenId); + expected.extend([MEDIA_PAD_ID as TokenId; 4]); + expected.push(MEDIA_END_ID as TokenId); + assert_eq!(rep.tokens, expected); + + // Only the pad run is an encoder-feature position; the wrapper is text. + // The pads start after `<|media_begin|>`, the dimensions, and + // `<|media_content|>`. + assert_eq!( + rep.feature_ranges, + Some(vec![PlaceholderRange { + offset: 2 + "image 1024x768".len(), + length: 4, + }]) + ); + } + + #[test] + fn kimi_k3_dimensions_are_per_image() { + let tokenizer = k3_tokenizer(); + let config = k3_config(); + let metadata = ModelMetadata { + model_id: "moonshotai/Kimi-K3", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).expect("k3 spec"); + + let replacements = spec + .prompt_replacements( + &metadata, + &preprocessed(&[(4000, 3000), (224, 448)], &[8, 2]), + ) + .unwrap(); + + assert_eq!(replacements.len(), 2); + for (rep, (text, pads)) in replacements + .iter() + .zip([("image 4000x3000", 8usize), ("image 224x448", 2)]) + { + let mut expected = vec![MEDIA_BEGIN_ID as TokenId]; + expected.extend(text_ids(text)); + expected.push(MEDIA_CONTENT_ID as TokenId); + expected.extend(std::iter::repeat_n(MEDIA_PAD_ID as TokenId, pads)); + expected.push(MEDIA_END_ID as TokenId); + assert_eq!(rep.tokens, expected); + assert_eq!( + rep.feature_ranges, + Some(vec![PlaceholderRange { + offset: 2 + text.len(), + length: pads, + }]) + ); + } + } + + #[test] + fn kimi_k3_requires_the_media_tokens_in_the_vocabulary() { + // A checkpoint without the structural tokens must fail loudly rather + // than silently emit a bare pad run. + let tokenizer = TestTokenizer::new(&[("<|media_pad|>", MEDIA_PAD_ID)]); + let config = k3_config(); + let metadata = ModelMetadata { + model_id: "moonshotai/Kimi-K3", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).expect("k3 spec"); + + let err = spec + .prompt_replacements(&metadata, &preprocessed(&[(64, 64)], &[1])) + .unwrap_err(); + assert_eq!( + err.to_string(), + "token '<|media_begin|>' not found in tokenizer vocabulary" + ); + } +} diff --git a/crates/multimodal/src/registry/mod.rs b/crates/multimodal/src/registry/mod.rs index ea167cab0..f4cfe4c7b 100644 --- a/crates/multimodal/src/registry/mod.rs +++ b/crates/multimodal/src/registry/mod.rs @@ -1,5 +1,6 @@ mod inkling; mod kimi_k25; +mod kimi_k3; mod llama4; mod llava; mod phi3_v; @@ -11,6 +12,7 @@ mod traits; use inkling::InklingSpec; use kimi_k25::KimiK25VisionSpec; +use kimi_k3::KimiK3VisionSpec; use llama4::Llama4Spec; use llava::{LlavaNextSpec, LlavaSpec}; use once_cell::sync::Lazy; @@ -33,6 +35,9 @@ impl ModelRegistry { Self { specs: vec![ LazySpec::new(|| Box::new(InklingSpec)), + // Kimi-K3 must be registered before Kimi-K2.5: the two families + // share a transport layout but not a prompt shape. + LazySpec::new(|| Box::new(KimiK3VisionSpec)), LazySpec::new(|| Box::new(KimiK25VisionSpec)), LazySpec::new(|| Box::new(Llama4Spec)), // LlavaNext must be registered before Llava so "llava_next" model_type matches first. @@ -95,6 +100,7 @@ pub(super) mod test_helpers { pub struct TestTokenizer { vocab: HashMap, + text_base: Option, } impl TestTokenizer { @@ -103,13 +109,31 @@ pub(super) mod test_helpers { .iter() .map(|(token, id)| ((*token).to_string(), *id)) .collect(); - Self { vocab } + Self { + vocab, + text_base: None, + } + } + + /// Encode ordinary text as one id per byte, offset by `base`. + /// + /// Off by default (plain text encodes to nothing) so specs that only + /// look up special tokens are unaffected. Specs that splice encoded + /// text into a replacement enable it to assert the exact layout. + pub fn with_byte_encoder(mut self, base: u32) -> Self { + self.text_base = Some(base); + self } } impl Encoder for TestTokenizer { - fn encode(&self, _input: &str, _add_special_tokens: bool) -> anyhow::Result { - Ok(Encoding::Plain(Vec::new())) + fn encode(&self, input: &str, _add_special_tokens: bool) -> anyhow::Result { + let Some(base) = self.text_base else { + return Ok(Encoding::Plain(Vec::new())); + }; + Ok(Encoding::Plain( + input.bytes().map(|b| base + u32::from(b)).collect(), + )) } fn encode_batch( @@ -119,7 +143,7 @@ pub(super) mod test_helpers { ) -> anyhow::Result> { inputs .iter() - .map(|_| self.encode("", add_special_tokens)) + .map(|input| self.encode(input, add_special_tokens)) .collect() } } diff --git a/crates/multimodal/src/registry/traits.rs b/crates/multimodal/src/registry/traits.rs index a4141c7f6..21059f58c 100644 --- a/crates/multimodal/src/registry/traits.rs +++ b/crates/multimodal/src/registry/traits.rs @@ -19,6 +19,8 @@ pub enum ModelRegistryError { TokenNotFound { token: String }, #[error("missing config field '{field}'")] MissingConfigField { field: String }, + #[error("model spec {spec} could not encode '{text}' with the model tokenizer")] + TextEncodingFailed { spec: &'static str, text: String }, #[error("modality {modality} is not supported by model spec {spec}")] UnsupportedModality { spec: &'static str, From d14325c655d01dbb3e420c119624189abc51903f Mon Sep 17 00:00:00 2001 From: key4ng Date: Tue, 28 Jul 2026 23:05:33 -0700 Subject: [PATCH 2/3] fix(kimi-k3): default thinking_effort to max in the served render path The K3 checkpoint splits prompt rendering across two layers: `encoding_k3.build_chat_segments` injects no thinking-effort directive, while the entry point above it, `tokenization_kimi.apply_chat_template`, runs `kwargs.setdefault("thinking_effort", "max")` first. vLLM calls the latter, so every served K3 request carries the directive. SMG called the equivalent of the former and omitted it -- a 67-token divergence on every request, image or not. Keep `apply_kimi_k3_xtml` a faithful `build_chat_segments` port and add `apply_kimi_k3_xtml_with_effort_default` for the served layer, which is what `TiktokenTokenizer::apply_chat_template` now calls. Threading the default through the fallback branch rather than pre-seeding `template_kwargs` preserves the precedence explicit `thinking_effort` > OpenAI `reasoning_effort` > default, and leaves all seven golden fixtures byte-valid. Verified against the checkpoint: the served entry point's output is the bare render prefixed by exactly this directive, 67 tokens. Signed-off-by: key4ng --- crates/tokenizer/src/encoders/kimi_k3_xtml.rs | 130 ++++++++++++++++-- crates/tokenizer/src/tiktoken.rs | 9 +- crates/tokenizer/tests/kimi_k3_renderer.rs | 15 +- 3 files changed, 136 insertions(+), 18 deletions(-) diff --git a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs index 4e0dc0a28..3bdad6808 100644 --- a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs +++ b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs @@ -28,13 +28,18 @@ //! their special-token ids while surrounding text is encoded as ordinary BPE — //! a cross-cutting change across this renderer family, tracked as future work. //! -//! # Deferred (TODO) +//! # Image prompts are not built here //! -//! The following `build_chat_segments` branches are intentionally not ported yet -//! (not exercised by the current gRPC path / golden fixtures): -//! - multimodal `image_prompts` substitution (image parts render as the bare -//! `<|kimi_image_placeholder|>` marker; a separate multimodal task will wire -//! real image prompts through). +//! The reference renders each image as a `<|kimi_image_placeholder|>` marker and +//! a caller-supplied `image_prompts` entry — +//! `<|media_begin|>image {w}x{h}<|media_content|><|media_pad|><|media_end|>` — +//! which `kimi_k3_processor.py::update_raw_text` splices in afterwards. SMG +//! cannot do that here: this renderer runs before any media is fetched, so the +//! dimensions do not exist yet. The gateway instead renders one bare +//! `<|media_pad|>` anchor per image (media parts arrive already flattened into +//! the message string) and expands it into the full block during prompt +//! expansion, from the preprocessor's reported sizes — see +//! `llm_multimodal::registry::kimi_k3`. use anyhow::{anyhow, Result}; use serde_json::{Map, Value}; @@ -47,7 +52,19 @@ const SEP_TOKEN: &str = "<|sep|>"; const END_OF_MSG_TOKEN: &str = "<|end_of_msg|>"; const IMAGE_PLACEHOLDER: &str = "<|kimi_image_placeholder|>"; -/// Render a Kimi-K3 chat prompt to an XTML `String`. +/// The effort the reference applies when a request names none. +/// +/// The two reference layers differ here: `build_chat_segments` injects no +/// directive at all, while the served entry point above it +/// (`tokenization_kimi.apply_chat_template`, which vLLM calls) runs +/// `kwargs.setdefault("thinking_effort", "max")` first. So every served K3 +/// request carries the directive even though the renderer itself never adds one. +pub const DEFAULT_THINKING_EFFORT: &str = "max"; + +/// Render a Kimi-K3 chat prompt to an XTML `String`, exactly as the Python +/// `build_chat_segments` does — including emitting no `thinking-effort` +/// directive when the request names no effort. Callers standing in for the +/// served entry point want [`apply_kimi_k3_xtml_with_effort_default`]. /// /// `params.tools` (when non-empty) produces the leading `tool-declare` system /// message. `params.add_generation_prompt` appends the assistant generation @@ -57,6 +74,24 @@ const IMAGE_PLACEHOLDER: &str = "<|kimi_image_placeholder|>"; /// `response` for both the generation-prompt tail and (per-turn) any prior /// assistant reasoning. pub fn apply_kimi_k3_xtml(messages: &[Value], params: &ChatTemplateParams) -> Result { + render_xtml(messages, params, None) +} + +/// Render as the *served* reference entry point does: like +/// [`apply_kimi_k3_xtml`], except that a request naming no effort falls back to +/// [`DEFAULT_THINKING_EFFORT`] rather than emitting no directive. +pub fn apply_kimi_k3_xtml_with_effort_default( + messages: &[Value], + params: &ChatTemplateParams, +) -> Result { + render_xtml(messages, params, Some(DEFAULT_THINKING_EFFORT)) +} + +fn render_xtml( + messages: &[Value], + params: &ChatTemplateParams, + default_effort: Option<&str>, +) -> Result { // Re-sort tool results by tool_call_id, then normalize each message // (deep-sort tool schemas, coerce tool-call arguments) — both side-effect // free, mirroring the Python entry point. @@ -95,11 +130,12 @@ pub fn apply_kimi_k3_xtml(messages: &[Value], params: &ChatTemplateParams) -> Re // `minimal`/`none` already switch thinking off upstream and `medium` has // no K3 equivalent, so such values emit no directive rather than erroring // on an otherwise-valid OpenAI field. + // 3. Absent both, `default_effort` — `None` for a bare `build_chat_segments` + // port, `Some("max")` for the served entry point. See + // [`DEFAULT_THINKING_EFFORT`]. // - // Absent both, no directive is emitted and the model applies its intrinsic - // `max` default — matching the reference, which injects nothing when - // `thinking_effort` is unset. `preserve_thinking` (which some callers send - // alongside the effort) is not read by the reference renderer and is ignored. + // `preserve_thinking` (which some callers send alongside the effort) is not + // read by the reference renderer and is ignored. if thinking { if let Some(effort_val) = params .template_kwargs @@ -122,6 +158,7 @@ pub fn apply_kimi_k3_xtml(messages: &[Value], params: &ChatTemplateParams) -> Re .and_then(|k| k.get("reasoning_effort")) .and_then(Value::as_str) .filter(|e| matches!(*e, "low" | "high" | "max")) + .or(default_effort) { push_thinking_effort(&mut out, effort); } @@ -315,8 +352,11 @@ fn push_internal_system_message(out: &mut String, message_type: &str, body: &str /// Render `content` (string, or an OpenAI content-part array) into `out`. /// -/// Image parts emit the bare `<|kimi_image_placeholder|>` marker; real -/// `image_prompts` substitution is deferred (see module docs). +/// Image parts in an array emit the reference's bare +/// `<|kimi_image_placeholder|>` marker. The gateway does not reach that branch: +/// K3 reports the `String` content format, so media parts are already flattened +/// into the message text as `<|media_pad|>` anchors before rendering (see module +/// docs). fn push_content(out: &mut String, content: Option<&Value>) { match content { Some(Value::String(s)) => out.push_str(s), @@ -989,8 +1029,9 @@ mod tests { #[test] fn no_effort_directive_when_unspecified() { - // Byte-parity with the reference: absent both keys, nothing is injected - // (the model applies its intrinsic `max` default). + // Byte-parity with `build_chat_segments`: absent both keys, this entry + // point injects nothing. The served wrapper above it does — see + // `served_entry_point_defaults_effort_to_max`. let messages = vec![json!({"role": "user", "content": "Hi"})]; let kwargs = HashMap::new(); let rendered = apply_kimi_k3_xtml(&messages, ¶ms_kw(None, &kwargs, true)).unwrap(); @@ -1000,6 +1041,65 @@ mod tests { ); } + #[test] + fn served_entry_point_defaults_effort_to_max() { + // `tokenization_kimi.apply_chat_template` setdefaults the effort to + // `max`, so a request naming none still gets the directive. + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::new(); + let rendered = + apply_kimi_k3_xtml_with_effort_default(&messages, ¶ms_kw(None, &kwargs, true)) + .unwrap(); + assert!( + rendered.contains("Now the system is invoked with `thinking_effort=max`."), + "served path must default to max: {rendered}" + ); + } + + #[test] + fn served_effort_default_yields_to_a_requested_effort() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + for (kwargs, expected) in [ + ( + HashMap::from([("thinking_effort".to_string(), json!("low"))]), + "low", + ), + ( + HashMap::from([("reasoning_effort".to_string(), json!("high"))]), + "high", + ), + ] { + let rendered = + apply_kimi_k3_xtml_with_effort_default(&messages, ¶ms_kw(None, &kwargs, true)) + .unwrap(); + assert!( + rendered.contains(&format!( + "Now the system is invoked with `thinking_effort={expected}`." + )), + "requested effort must win over the default: {rendered}" + ); + assert!( + !rendered.contains("thinking_effort=max`."), + "default must not also be emitted: {rendered}" + ); + } + } + + #[test] + fn served_effort_default_suppressed_when_thinking_off() { + let messages = vec![json!({"role": "user", "content": "Hi"})]; + let kwargs = HashMap::new(); + let rendered = apply_kimi_k3_xtml_with_effort_default( + &messages, + ¶ms_kw(Some(false), &kwargs, true), + ) + .unwrap(); + assert!( + !rendered.contains("type=\"thinking-effort\""), + "no effort directive when thinking is off: {rendered}" + ); + } + // --- Structural channel ------------------------------------------ #[test] diff --git a/crates/tokenizer/src/tiktoken.rs b/crates/tokenizer/src/tiktoken.rs index 4d78872cf..bf62bfb87 100644 --- a/crates/tokenizer/src/tiktoken.rs +++ b/crates/tokenizer/src/tiktoken.rs @@ -17,7 +17,9 @@ use crate::{ load_chat_template_from_file, ChatTemplateContentFormat, ChatTemplateParams, ChatTemplateState, ThinkingKeyName, ThinkingToggle, }, - encoders::{kimi_k25_tools::apply_kimi_k25_tools, kimi_k3_xtml::apply_kimi_k3_xtml}, + encoders::{ + kimi_k25_tools::apply_kimi_k25_tools, kimi_k3_xtml::apply_kimi_k3_xtml_with_effort_default, + }, factory::discover_chat_template_in_dir, kimi_k2_tokenizer, traits::{Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer as TokenizerTrait}, @@ -554,7 +556,10 @@ impl TokenizerTrait for TiktokenTokenizer { match self.renderer { Renderer::Jinja => self.chat_template.apply(messages, params), Renderer::KimiK25Tools => apply_kimi_k25_tools(&self.chat_template, messages, params), - Renderer::KimiK3Xtml => apply_kimi_k3_xtml(messages, ¶ms), + // This is the layer the K3 checkpoint's own `apply_chat_template` + // sits at, so it applies that wrapper's `thinking_effort` default — + // the bare renderer stays a faithful `build_chat_segments` port. + Renderer::KimiK3Xtml => apply_kimi_k3_xtml_with_effort_default(messages, ¶ms), } } diff --git a/crates/tokenizer/tests/kimi_k3_renderer.rs b/crates/tokenizer/tests/kimi_k3_renderer.rs index 080e39cb2..0faf67ebc 100644 --- a/crates/tokenizer/tests/kimi_k3_renderer.rs +++ b/crates/tokenizer/tests/kimi_k3_renderer.rs @@ -180,6 +180,13 @@ fn thinking_effort_invalid_is_rejected() { /// End-to-end: a tokenizer loaded from a K3 directory (no chat template at all) /// must load successfully, detect the K3 renderer, and render XTML through /// `apply_chat_template`. +/// +/// `apply_chat_template` stands in for the checkpoint's `tokenization_kimi` +/// wrapper, not for `build_chat_segments`, so its output is the plain fixture +/// *plus* the `max` effort directive the wrapper's +/// `setdefault("thinking_effort", "max")` produces. The expected bytes are the +/// `thinking_effort_low` fixture with its effort word swapped — the directive is +/// otherwise identical at every level. #[test] fn tokenizer_loads_and_renders_k3_without_chat_template() { let dir = TempDir::new().unwrap(); @@ -205,5 +212,11 @@ fn tokenizer_loads_and_renders_k3_without_chat_template() { ) .expect("K3 render should succeed"); - assert_eq!(rendered, fixture_text("plain_user_thinking")); + let expected = fixture_text("thinking_effort_low") + .replace("`thinking_effort=low`", "`thinking_effort=max`"); + assert_eq!(rendered, expected); + assert!( + rendered.ends_with(&fixture_text("plain_user_thinking")), + "the directive is the only addition: {rendered}" + ); } From 879016805ac97824d21dad1ba82fae2518f30a13 Mon Sep 17 00:00:00 2001 From: key4ng Date: Tue, 28 Jul 2026 23:18:41 -0700 Subject: [PATCH 3/3] docs(kimi-k3): tighten the comments added for the prompt-encoding fix Trim the doc and inline comments on the new K3 registry spec and the split XTML entry points. The reference details they restate are already in the commit history; keep the part a reader needs at the call site and drop the retelling. Signed-off-by: key4ng --- crates/multimodal/src/registry/kimi_k25.rs | 9 ++-- crates/multimodal/src/registry/kimi_k3.rs | 41 +++++++--------- crates/multimodal/src/registry/mod.rs | 8 ++-- crates/tokenizer/src/encoders/kimi_k3_xtml.rs | 47 +++++++------------ crates/tokenizer/src/tiktoken.rs | 5 +- crates/tokenizer/tests/kimi_k3_renderer.rs | 8 ++-- 6 files changed, 46 insertions(+), 72 deletions(-) diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs index 4dff24034..0caa8cbbe 100644 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ b/crates/multimodal/src/registry/kimi_k25.rs @@ -28,12 +28,9 @@ impl ModelProcessorSpec for KimiK25VisionSpec { } fn matches(&self, metadata: &ModelMetadata) -> bool { - // K2.5 only. Kimi-K3 shares the `<|media_pad|>` fill token - // (media_placeholder_token_id 163605) and the patchification layout, but - // its prompt carries a per-image `<|media_begin|>image {w}x{h}…` wrapper - // that K2.5's chat template emits itself — see `registry::kimi_k3`. The - // two do not share a pixel pipeline either; `vision::processors::kimi_k3` - // has the patch budget and alpha-compositing differences. + // K2.5 only — K3 shares the `<|media_pad|>` fill token and the + // patchification layout, but neither the prompt shape nor the pixel + // pipeline. See `registry::kimi_k3` and `vision::processors::kimi_k3`. let id = metadata.model_id.to_ascii_lowercase(); id.contains("kimi") && id.contains("k2") || metadata diff --git a/crates/multimodal/src/registry/kimi_k3.rs b/crates/multimodal/src/registry/kimi_k3.rs index cc53d1798..b328d4ee6 100644 --- a/crates/multimodal/src/registry/kimi_k3.rs +++ b/crates/multimodal/src/registry/kimi_k3.rs @@ -18,20 +18,14 @@ const MEDIA_END: &str = "<|media_end|>"; /// Kimi-K3. /// -/// K3 shares K2.5's NaViT transport layout and `<|media_pad|>` fill token, but -/// not its prompt shape: the K3 checkpoint ships no jinja chat template, and its -/// Python renderer splices a per-image `<|media_begin|>image {w}x{h} -/// <|media_content|>…<|media_end|>` block into the raw text at each -/// `<|kimi_image_placeholder|>` marker (`kimi_k3_processor.py::update_raw_text`). -/// K2.5's template emits the wrapper itself and carries no dimensions, so the two -/// families need separate specs. +/// Shares K2.5's MoonViT transport layout and `<|media_pad|>` fill token, but +/// not its prompt shape: K3 wraps each image in a block carrying the pre-resize +/// dimensions, while K2.5's chat template emits its own dimensionless wrapper. /// -/// SMG cannot build that block while rendering — the chat template runs before -/// any media is fetched, so the image dimensions do not exist yet. It is built -/// here instead, from the pre-resize sizes the preprocessor reports, which is -/// where vLLM builds it too (`kimi_k3.py::_get_prompt_updates`). The renderer -/// emits a bare `<|media_pad|>` anchor per image and this replacement expands it -/// into the full wrapper. +/// That block cannot be built while rendering — the chat template runs before +/// any media is fetched, so the dimensions do not exist yet. It is built here +/// instead, from the sizes the preprocessor reports, as vLLM does in +/// `kimi_k3.py::_get_prompt_updates`. pub(super) struct KimiK3VisionSpec; impl KimiK3VisionSpec { @@ -45,11 +39,11 @@ impl KimiK3VisionSpec { }) } - /// Encode ordinary text (no special tokens) into token ids. + /// Encode ordinary text into token ids. /// /// The dimension text sits between two special tokens, which are hard - /// segment boundaries for the tiktoken encoder, so encoding it on its own - /// yields the same ids the reference gets from encoding the whole block. + /// segment boundaries for the encoder, so encoding it alone yields the same + /// ids as the reference's one-shot encoding of the whole block. fn encode_plain_text(metadata: &ModelMetadata, text: &str) -> RegistryResult> { let encoding = metadata.tokenizer.encode(text, false).map_err(|_| { ModelRegistryError::TextEncodingFailed { @@ -109,10 +103,9 @@ impl ModelProcessorSpec for KimiK3VisionSpec { let media_content = metadata.token_id(MEDIA_CONTENT)?; let media_end = metadata.token_id(MEDIA_END)?; - // MoonViT reports `item_sizes` as the decoded `(width, height)` before any - // resize, which is exactly the pair the reference prints. The caller - // checks both vectors against the media count, so a short zip here would - // already have been rejected upstream. + // `item_sizes` is the decoded `(width, height)` before any resize — the + // pair the reference prints. The caller already checks both vectors + // against the media count, so a short zip cannot reach here. preprocessed .feature_token_counts .iter() @@ -165,13 +158,12 @@ mod tests { types::{Modality, PlaceholderRange, TokenId}, }; - /// Token ids the K3 checkpoint actually assigns to the wrapper. + /// Wrapper token ids as the K3 checkpoint assigns them. const MEDIA_BEGIN_ID: u32 = 163602; const MEDIA_CONTENT_ID: u32 = 163603; const MEDIA_END_ID: u32 = 163604; const MEDIA_PAD_ID: u32 = 163605; - /// Offset of the byte encoder used by `TestTokenizer::with_byte_encoder`, - /// chosen so plain-text ids cannot collide with the media token ids. + /// Byte-encoder offset, chosen so text ids cannot collide with media ids. const TEXT_BASE: u32 = 1000; fn k3_tokenizer() -> TestTokenizer { @@ -266,8 +258,7 @@ mod tests { assert_eq!(rep.tokens, expected); // Only the pad run is an encoder-feature position; the wrapper is text. - // The pads start after `<|media_begin|>`, the dimensions, and - // `<|media_content|>`. + // Pads start after `<|media_begin|>`, the dims, and `<|media_content|>`. assert_eq!( rep.feature_ranges, Some(vec![PlaceholderRange { diff --git a/crates/multimodal/src/registry/mod.rs b/crates/multimodal/src/registry/mod.rs index f4cfe4c7b..c37a21f6e 100644 --- a/crates/multimodal/src/registry/mod.rs +++ b/crates/multimodal/src/registry/mod.rs @@ -115,11 +115,11 @@ pub(super) mod test_helpers { } } - /// Encode ordinary text as one id per byte, offset by `base`. + /// Encode text as one id per byte, offset by `base`. /// - /// Off by default (plain text encodes to nothing) so specs that only - /// look up special tokens are unaffected. Specs that splice encoded - /// text into a replacement enable it to assert the exact layout. + /// Off by default so specs that only look up special tokens are + /// unaffected; specs that splice encoded text into a replacement + /// enable it to assert the exact layout. pub fn with_byte_encoder(mut self, base: u32) -> Self { self.text_base = Some(base); self diff --git a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs index 3bdad6808..13222aaf4 100644 --- a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs +++ b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs @@ -30,16 +30,11 @@ //! //! # Image prompts are not built here //! -//! The reference renders each image as a `<|kimi_image_placeholder|>` marker and -//! a caller-supplied `image_prompts` entry — -//! `<|media_begin|>image {w}x{h}<|media_content|><|media_pad|><|media_end|>` — -//! which `kimi_k3_processor.py::update_raw_text` splices in afterwards. SMG -//! cannot do that here: this renderer runs before any media is fetched, so the -//! dimensions do not exist yet. The gateway instead renders one bare -//! `<|media_pad|>` anchor per image (media parts arrive already flattened into -//! the message string) and expands it into the full block during prompt -//! expansion, from the preprocessor's reported sizes — see -//! `llm_multimodal::registry::kimi_k3`. +//! The reference splices a per-image `<|media_begin|>image {w}x{h}…` block in +//! afterwards (`kimi_k3_processor.py::update_raw_text`). This renderer runs +//! before any media is fetched, so those dimensions do not exist yet: the +//! gateway emits one bare `<|media_pad|>` anchor per image and expands it +//! during prompt expansion — see `llm_multimodal::registry::kimi_k3`. use anyhow::{anyhow, Result}; use serde_json::{Map, Value}; @@ -54,17 +49,15 @@ const IMAGE_PLACEHOLDER: &str = "<|kimi_image_placeholder|>"; /// The effort the reference applies when a request names none. /// -/// The two reference layers differ here: `build_chat_segments` injects no -/// directive at all, while the served entry point above it -/// (`tokenization_kimi.apply_chat_template`, which vLLM calls) runs -/// `kwargs.setdefault("thinking_effort", "max")` first. So every served K3 -/// request carries the directive even though the renderer itself never adds one. +/// `build_chat_segments` injects no directive; the served entry point above it +/// (`tokenization_kimi.apply_chat_template`, which vLLM calls) first runs +/// `kwargs.setdefault("thinking_effort", "max")`. pub const DEFAULT_THINKING_EFFORT: &str = "max"; /// Render a Kimi-K3 chat prompt to an XTML `String`, exactly as the Python -/// `build_chat_segments` does — including emitting no `thinking-effort` -/// directive when the request names no effort. Callers standing in for the -/// served entry point want [`apply_kimi_k3_xtml_with_effort_default`]. +/// `build_chat_segments` does — emitting no `thinking-effort` directive when +/// the request names none. Callers standing in for the served entry point want +/// [`apply_kimi_k3_xtml_with_effort_default`]. /// /// `params.tools` (when non-empty) produces the leading `tool-declare` system /// message. `params.add_generation_prompt` appends the assistant generation @@ -130,12 +123,10 @@ fn render_xtml( // `minimal`/`none` already switch thinking off upstream and `medium` has // no K3 equivalent, so such values emit no directive rather than erroring // on an otherwise-valid OpenAI field. - // 3. Absent both, `default_effort` — `None` for a bare `build_chat_segments` - // port, `Some("max")` for the served entry point. See - // [`DEFAULT_THINKING_EFFORT`]. + // 3. Absent both, `default_effort` — `None` for the bare port, `Some("max")` + // for the served entry point. See [`DEFAULT_THINKING_EFFORT`]. // - // `preserve_thinking` (which some callers send alongside the effort) is not - // read by the reference renderer and is ignored. + // `preserve_thinking` is not read by the reference renderer and is ignored. if thinking { if let Some(effort_val) = params .template_kwargs @@ -352,11 +343,9 @@ fn push_internal_system_message(out: &mut String, message_type: &str, body: &str /// Render `content` (string, or an OpenAI content-part array) into `out`. /// -/// Image parts in an array emit the reference's bare -/// `<|kimi_image_placeholder|>` marker. The gateway does not reach that branch: -/// K3 reports the `String` content format, so media parts are already flattened -/// into the message text as `<|media_pad|>` anchors before rendering (see module -/// docs). +/// Image parts emit the reference's bare `<|kimi_image_placeholder|>` marker. +/// The gateway does not reach that branch: K3 reports the `String` content +/// format, so media parts arrive already flattened into the message text. fn push_content(out: &mut String, content: Option<&Value>) { match content { Some(Value::String(s)) => out.push_str(s), @@ -1029,7 +1018,7 @@ mod tests { #[test] fn no_effort_directive_when_unspecified() { - // Byte-parity with `build_chat_segments`: absent both keys, this entry + // Byte-parity with `build_chat_segments`: absent both keys this entry // point injects nothing. The served wrapper above it does — see // `served_entry_point_defaults_effort_to_max`. let messages = vec![json!({"role": "user", "content": "Hi"})]; diff --git a/crates/tokenizer/src/tiktoken.rs b/crates/tokenizer/src/tiktoken.rs index bf62bfb87..8f34be181 100644 --- a/crates/tokenizer/src/tiktoken.rs +++ b/crates/tokenizer/src/tiktoken.rs @@ -556,9 +556,8 @@ impl TokenizerTrait for TiktokenTokenizer { match self.renderer { Renderer::Jinja => self.chat_template.apply(messages, params), Renderer::KimiK25Tools => apply_kimi_k25_tools(&self.chat_template, messages, params), - // This is the layer the K3 checkpoint's own `apply_chat_template` - // sits at, so it applies that wrapper's `thinking_effort` default — - // the bare renderer stays a faithful `build_chat_segments` port. + // This is the layer the checkpoint's own `apply_chat_template` sits + // at, so it applies that wrapper's `thinking_effort` default. Renderer::KimiK3Xtml => apply_kimi_k3_xtml_with_effort_default(messages, ¶ms), } } diff --git a/crates/tokenizer/tests/kimi_k3_renderer.rs b/crates/tokenizer/tests/kimi_k3_renderer.rs index 0faf67ebc..f4c0cfda0 100644 --- a/crates/tokenizer/tests/kimi_k3_renderer.rs +++ b/crates/tokenizer/tests/kimi_k3_renderer.rs @@ -182,11 +182,9 @@ fn thinking_effort_invalid_is_rejected() { /// `apply_chat_template`. /// /// `apply_chat_template` stands in for the checkpoint's `tokenization_kimi` -/// wrapper, not for `build_chat_segments`, so its output is the plain fixture -/// *plus* the `max` effort directive the wrapper's -/// `setdefault("thinking_effort", "max")` produces. The expected bytes are the -/// `thinking_effort_low` fixture with its effort word swapped — the directive is -/// otherwise identical at every level. +/// wrapper, so its output is the plain fixture *plus* the `max` effort +/// directive. The expected bytes are the `thinking_effort_low` fixture with its +/// effort word swapped — the directive is identical at every level. #[test] fn tokenizer_loads_and_renders_k3_without_chat_template() { let dir = TempDir::new().unwrap();