fix(kimi-k3): move the media wrapper and thinking-effort default to the prompt-encoding layer - #1995
Conversation
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 <rukeyang@gmail.com>
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 <rukeyang@gmail.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
github-actions
Bot
added
tokenizer
📝 WalkthroughWalkthroughAdds a dedicated Kimi K3 multimodal processor with image wrapper generation and registry routing. Kimi K3 XTML rendering now applies a served ChangesKimi K3 multimodal registry and prompt processing
Kimi K3 served rendering
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Request
participant TiktokenTokenizer
participant KimiK3XTML
Request->>TiktokenTokenizer: apply KimiK3Xtml chat template
TiktokenTokenizer->>KimiK3XTML: call served renderer
KimiK3XTML->>KimiK3XTML: resolve explicit effort or default to max
KimiK3XTML-->>TiktokenTokenizer: return rendered XTML
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean PR. The two-layer split (bare renderer vs. served entry point with effort default) is well-motivated, the precedence chain is correct, the K3 vision spec builds the reference wrapper faithfully, and test coverage hits all the key scenarios — matcher routing, wrapper layout, per-image dimensions, missing tokens, and effort precedence. No issues found.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/multimodal/src/registry/kimi_k25.rs`:
- Around line 130-148: The test kimi_k3_does_not_use_the_k25_spec must directly
verify KimiK25VisionSpec::matches rejects each K3 metadata case before checking
registry routing. Construct the metadata for each model_id, assert the K2.5 spec
does not match it, then retain the existing registry lookup assertion that
routes to kimi_k3.
In `@crates/multimodal/src/registry/kimi_k3.rs`:
- Around line 112-115: Update the comment near the upstream media-count
validation to explicitly prefix the cardinality assumption with “INVARIANT:”.
Keep the existing explanation unchanged and do not alter the surrounding logic.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: f2b3cb48-0a98-4c17-80e3-b729576b3544
📒 Files selected for processing (7)
crates/multimodal/src/registry/kimi_k25.rscrates/multimodal/src/registry/kimi_k3.rscrates/multimodal/src/registry/mod.rscrates/multimodal/src/registry/traits.rscrates/tokenizer/src/encoders/kimi_k3_xtml.rscrates/tokenizer/src/tiktoken.rscrates/tokenizer/tests/kimi_k3_renderer.rs
| 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}"); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Directly test that K2.5 rejects K3 metadata.
This registry lookup still passes if KimiK25VisionSpec::matches() regresses, since K3 is registered first. Assert the K2.5 spec rejects each case before checking registry routing.
Proposed test strengthening
+ use super::KimiK25VisionSpec;
use crate::{
- registry::{test_helpers::*, ModelMetadata, ModelRegistry},
+ registry::{test_helpers::*, ModelMetadata, ModelProcessorSpec, ModelRegistry},
types::ImageSize,
};
...
let metadata = ModelMetadata {
model_id,
tokenizer: &tokenizer,
config: &config,
};
+ let k25 = KimiK25VisionSpec;
+ assert!(
+ !k25.matches(&metadata),
+ "K3 metadata must not match the K2.5 spec: {model_id}"
+ );
let spec = registry.lookup(&metadata).expect("kimi_k3 spec");📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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}"); | |
| } | |
| 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)]); | |
| let config = json!({ | |
| "model_type": "kimi_k3", | |
| "media_placeholder_token_id": 163605 | |
| }); | |
| let registry = ModelRegistry::new(); | |
| for model_id in ["moonshotai/Kimi-K3", "internal/checkpoint-final"] { | |
| let metadata = ModelMetadata { | |
| model_id, | |
| tokenizer: &tokenizer, | |
| config: &config, | |
| }; | |
| let k25 = KimiK25VisionSpec; | |
| assert!( | |
| !k25.matches(&metadata), | |
| "K3 metadata must not match the K2.5 spec: {model_id}" | |
| ); | |
| let spec = registry.lookup(&metadata).expect("kimi_k3 spec"); | |
| assert_eq!(spec.name(), "kimi_k3", "model_id {model_id}"); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/multimodal/src/registry/kimi_k25.rs` around lines 130 - 148, The test
kimi_k3_does_not_use_the_k25_spec must directly verify
KimiK25VisionSpec::matches rejects each K3 metadata case before checking
registry routing. Construct the metadata for each model_id, assert the K2.5 spec
does not match it, then retain the existing registry lookup assertion that
routes to kimi_k3.
| // 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. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Mark the upstream cardinality assumption as an invariant.
Prefix the upstream-validation assumption with INVARIANT: (or remove the assertion from this comment) to follow the repository’s safe-code invariant convention.
Based on learnings, use the marker INVARIANT: to document assumptions in safe code.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/multimodal/src/registry/kimi_k3.rs` around lines 112 - 115, Update the
comment near the upstream media-count validation to explicitly prefix the
cardinality assumption with “INVARIANT:”. Keep the existing explanation
unchanged and do not alter the surrounding logic.
Source: Learnings
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 <rukeyang@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 879016805a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| .filter(|e| matches!(*e, "low" | "high" | "max")) | ||
| .or(default_effort) |
There was a problem hiding this comment.
When the served TiktokenTokenizer path receives the valid OpenAI value reasoning_effort="medium", this filter removes it and .or(default_effort) immediately substitutes max, so K3 gets a max-effort directive despite the documented bridge semantics that unsupported K3 levels emit no directive. The existing medium test only exercises the lower-level renderer without a default and therefore misses this regression; distinguish an absent reasoning_effort key from a present but unmappable value before applying the default.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/multimodal/src/registry/kimi_k25.rs (1)
31-38: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRestrict the matcher to Kimi-K2.5 models.
id.contains("k2")also accepts other Kimi K2-prefixed model IDs, contradicting the K2.5-only routing contract and potentially sending them through the wrong vision processor. Match canonical K2.5 identifiers explicitly and add a negative test forKimi-K2.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/multimodal/src/registry/kimi_k25.rs` around lines 31 - 38, Update the model matcher near metadata.model_id and config_model_type to accept only canonical Kimi-K2.5 identifiers, removing the broad id.contains("k2") check. Preserve the existing kimi_k25 config-model-type match, and add a negative test confirming that “Kimi-K2” is rejected.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@crates/multimodal/src/registry/kimi_k25.rs`:
- Around line 31-38: Update the model matcher near metadata.model_id and
config_model_type to accept only canonical Kimi-K2.5 identifiers, removing the
broad id.contains("k2") check. Preserve the existing kimi_k25 config-model-type
match, and add a negative test confirming that “Kimi-K2” is rejected.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b6eeea3b-4d7f-4fbc-9c88-7ddc90eedc4c
📒 Files selected for processing (6)
crates/multimodal/src/registry/kimi_k25.rscrates/multimodal/src/registry/kimi_k3.rscrates/multimodal/src/registry/mod.rscrates/tokenizer/src/encoders/kimi_k3_xtml.rscrates/tokenizer/src/tiktoken.rscrates/tokenizer/tests/kimi_k3_renderer.rs
Node validation: prompt shape and token parityRan this branch on the B300 node against the real Setup. One vLLM gRPC engine (
Reference. Before handing the GPUs to the gRPC engine, I captured the same requests against vLLM's HTTP server ( for a 1024×768 image — 1036 pads, and 3 marker tokens + 6 tokens of Token parity
new matches the reference on 8/8. old matches on 2/8 — exactly the two cases that pass an explicit effort and contain no image, i.e. the two things the old build already got right. Every baseline gap decomposes exactly, with no remainder:
The per-image term tracks the actual dimensions rather than being a constant, which is the part that matters. Gateway debug logging agrees from the inside — for the 4000×3000 case the new build reports Width×height ordering, and pre-resize dimensionsToken counts cannot catch a Strip images make this unambiguous: the same two numbers in both orders, and few enough pads that retrieval is reliable.
The mirror pair both read back correctly, so this is not the model guessing an orientation. This also settles pre- vs post-resize: a 4000×64 strip is heavily resized for patching (~380 pads), yet the wrapper still states the original 4000×64 — One caveatSMG does not forward |
Vendor-verifier accuracy checkRan the two image benchmarks from MoonshotAI/Kimi-Vendor-Verifier against this branch, to confirm the new media wrapper doesn't cost anything end-to-end.
The three OCRBench numbers sit inside one standard error of each other, so this reads as no regression rather than an improvement. The middle row is the useful one for this PR specifically: it sends no Setup:
Two things to note about the numbers rather than the code:
|
1 participant
Motivation
Two divergences from the K3 reference implementation, both in prompt encoding rather than in the vision processor.
1. The media wrapper carried no dimensions. K3's image prompt is one wrapper per image:
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 — and K2.5's jinja chat template emits its own, dimensionless wrapper. Every K3 image prompt was 9 tokens short of the reference and told the model nothing about the image's original size.2. Every request was missing the thinking-effort directive. The K3 checkpoint splits rendering across two layers.
encoding_k3.build_chat_segmentsinjects no directive; the entry point above it,tokenization_kimi.apply_chat_template, runskwargs.setdefault("thinking_effort", "max")first. vLLM calls the latter, so every served K3 request carries the directive. SMG's renderer is a faithful port of the lower layer and was wired directly intoTiktokenTokenizer::apply_chat_template, so it omitted the directive — a 67-token divergence on every request, image or not.Approach
Wrapper. It cannot be built while rendering: the chat template runs before any media is fetched, so the dimensions do not exist yet. A new
KimiK3VisionSpecbuilds it inprompt_replacements, from the sizes the preprocessor reports — the same place vLLM builds it (kimi_k3.py::_get_prompt_updates). The renderer emits a bare<|media_pad|>anchor per image and prompt expansion replaces it with the full block.with_feature_spankeeps the encoder-feature positions on the pad run alone; the surrounding wrapper is text.The K2.5 matcher is narrowed to K2.5 and K3 is registered ahead of it. The two families share the MoonViT transport layout but not a prompt shape — implicit sharing is exactly what produced the pixel-pipeline divergence fixed in #1984.
Thinking effort.
apply_kimi_k3_xtmlstays a faithfulbuild_chat_segmentsport; a newapply_kimi_k3_xtml_with_effort_defaultmodels the served layer and is whatTiktokenTokenizer::apply_chat_templatenow calls. The default is threaded through the fallback branch rather than pre-seedingtemplate_kwargs, which preserves the precedence explicitthinking_effort> OpenAIreasoning_effort> default — pre-seeding would have permanently starved thereasoning_effortbridge. All seven golden fixtures stay byte-valid.Verification against the checkpoint
Run inside the K3 container on a B300 node, against
moonshotai/Kimi-K3:[<|media_begin|>] + encode("image {w}x{h}") + [<|media_content|>, <|media_pad|>, <|media_end|>]is byte-identical to the reference's one-shot encoding ofmake_image_prompt(w, h). This confirms the one real assumption in the spec: the media tokens are hard segment boundaries for the tiktoken encoder, so the dimension text can be encoded on its own. Wrapper cost beyond the bare pad: 9 tokens.tokenization_kimi.apply_chat_templateoutput is exactlybuild_chat_segmentsoutput prefixed by thethinking_effort=maxdirective, 67 tokens.Together these close the token accounting on the reference request: 1092 (SMG before) + 9 + 67 = 1168, matching what vLLM sends.
Behaviour change
Text-only K3 requests get 67 tokens longer, and image requests get 9 more per image. Both are the point — that is what the reference serves.
Tests
registry::kimi_k3: matcher by model_id and model_type; the exact wrapper layout and feature span; per-image dimensions across a two-image batch; a checkpoint missing the structural tokens fails loudly rather than silently emitting a bare pad run.registry::kimi_k25: K3 no longer resolves to the K2.5 spec, by model_id and by model_type.encoders::kimi_k3_xtml: the served path defaults tomax; the default yields to both an explicitthinking_effortand an OpenAIreasoning_effort; the default is suppressed when thinking is off.tests/kimi_k3_renderer.rs: the end-to-end no-chat-template case now asserts the directive is present and is the only addition.Full workspace suite: 99 suites, 4151 passed, 0 failed.
Note
cargo clippy --all-featurescould not be run locally — it pulls the optionalopencvdependency and the dev machine has neitherpkg-confignoropencvinstalled.cargo clippy --workspace --all-targets -- -D warningsis clean, and no changed code sits behind acfg(feature)gate. CI covers the all-features build.Summary by CodeRabbit
thinking-efforttomaxwhen no effort is specified.