diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs index 994e5cf06..0caa8cbbe 100644 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ b/crates/multimodal/src/registry/kimi_k25.rs @@ -28,16 +28,14 @@ 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 — 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") || 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 +124,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..b328d4ee6 --- /dev/null +++ b/crates/multimodal/src/registry/kimi_k3.rs @@ -0,0 +1,333 @@ +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. +/// +/// 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. +/// +/// 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 { + /// 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 into token ids. + /// + /// The dimension text sits between two special tokens, which are hard + /// 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 { + 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)?; + + // `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() + .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}, + }; + + /// 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; + /// Byte-encoder offset, chosen so text ids cannot collide with media 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. + // Pads start after `<|media_begin|>`, the dims, 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..c37a21f6e 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 text as one id per byte, offset by `base`. + /// + /// 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 } } 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, diff --git a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs index 4e0dc0a28..13222aaf4 100644 --- a/crates/tokenizer/src/encoders/kimi_k3_xtml.rs +++ b/crates/tokenizer/src/encoders/kimi_k3_xtml.rs @@ -28,13 +28,13 @@ //! 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 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}; @@ -47,7 +47,17 @@ 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. +/// +/// `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 — 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 @@ -57,6 +67,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 +123,10 @@ 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 the bare 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` is not read by the reference renderer and is ignored. if thinking { if let Some(effort_val) = params .template_kwargs @@ -122,6 +149,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 +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 emit the bare `<|kimi_image_placeholder|>` marker; real -/// `image_prompts` substitution is deferred (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), @@ -989,8 +1018,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 +1030,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..8f34be181 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,9 @@ 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 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 080e39cb2..f4c0cfda0 100644 --- a/crates/tokenizer/tests/kimi_k3_renderer.rs +++ b/crates/tokenizer/tests/kimi_k3_renderer.rs @@ -180,6 +180,11 @@ 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, 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(); @@ -205,5 +210,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}" + ); }