feat(multimodal): add Kimi-K2.5 vision model spec and image processor - #1097
ConnorLi96 wants to merge 1 commit into
Conversation
📝 WalkthroughWalkthroughThe PR adds support for the Kimi K2.5 model family by introducing a model processor spec that defines image handling (10-image max, Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the Kimi-K2.5 model by implementing the KimiK25Spec for model registration and the KimiK25Processor for vision processing, which reuses the Qwen-VL base architecture. Feedback identifies a critical field name mismatch for grid information between the processor and the registry, suggests keeping the patches_per_image metadata on the CPU for host-side logic, and recommends overriding get_processed_size to ensure consistent behavior with the underlying dynamic resolution processor.
| fn preprocess( | ||
| &self, | ||
| images: &[DynamicImage], | ||
| config: &PreProcessorConfig, | ||
| ) -> Result<PreprocessedImages, TransformError> { | ||
| self.inner.preprocess(images, config) | ||
| } |
There was a problem hiding this comment.
The QwenVLProcessorBase produces a model-specific field named "image_grid_thw", but the KimiK25Spec defines and expects "grid_thws". This mismatch will prevent the model from receiving the necessary grid information. The key should be renamed in the processor's preprocess method to match the spec.
| fn preprocess( | |
| &self, | |
| images: &[DynamicImage], | |
| config: &PreProcessorConfig, | |
| ) -> Result<PreprocessedImages, TransformError> { | |
| self.inner.preprocess(images, config) | |
| } | |
| fn preprocess( | |
| &self, | |
| images: &[DynamicImage], | |
| config: &PreProcessorConfig, | |
| ) -> Result<PreprocessedImages, TransformError> { | |
| let mut preprocessed = self.inner.preprocess(images, config)?; | |
| if let Some(grid_thw) = preprocessed.model_specific.remove("image_grid_thw") { | |
| preprocessed.model_specific.insert("grid_thws".to_string(), grid_thw); | |
| } | |
| Ok(preprocessed) | |
| } |
References
- For protocol data structures that mirror an external API (e.g., OpenAI), prioritize alignment with the external specification over internal consistency.
| fn keep_on_cpu_keys(&self) -> Vec<String> { | ||
| vec!["grid_thws".to_string()] | ||
| } |
There was a problem hiding this comment.
The patches_per_image field is a metadata tensor used for batching and slicing. Similar to grid_thws, it should be kept on the CPU to ensure it is accessible to the host-side scheduling logic and to avoid unnecessary GPU transfers if the model expects it on the host.
| fn keep_on_cpu_keys(&self) -> Vec<String> { | |
| vec!["grid_thws".to_string()] | |
| } | |
| fn keep_on_cpu_keys(&self) -> Vec<String> { | |
| vec!["grid_thws".to_string(), "patches_per_image".to_string()] | |
| } |
| fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { | ||
| self.inner.calculate_num_tokens(width, height, config) | ||
| } |
There was a problem hiding this comment.
The KimiK25Processor should override get_processed_size to delegate to the inner processor. Since Kimi-K2.5 uses dynamic resolution (NaViT-style), it doesn't have a single fixed processed size, and the base implementation correctly returns None.
| fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { | |
| self.inner.calculate_num_tokens(width, height, config) | |
| } | |
| fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { | |
| self.inner.calculate_num_tokens(width, height, config) | |
| } | |
| fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { | |
| self.inner.get_processed_size(config) | |
| } |
Add KimiK25Spec to the multimodal model registry and KimiK25Processor to the image processor registry, enabling SMG to handle Kimi-K2.5 VLM requests (image placeholder expansion, token computation, raw image byte passthrough for TRT-LLM). Kimi-K2.5 uses MoonViT3d with NaViT-style architecture (patch_size=14, merge_kernel_size=2x2, dynamic resolution), which is structurally similar to Qwen VL. The processor reuses QwenVLProcessorBase with Kimi-specific defaults (mean/std=[0.5,0.5,0.5], max 16384 patches). Signed-off-by: ConnorLi96 <ConnorLi96@users.noreply.github.com> Made-with: Cursor Signed-off-by: ConnorLi96 <ConnorLi96@users.noreply.github.com> Made-with: Cursor
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| ("grid_thws".to_string(), FieldLayout::Batched), | ||
| ("patches_per_image".to_string(), FieldLayout::Batched), | ||
| ]) | ||
| } | ||
|
|
||
| fn keep_on_cpu_keys(&self) -> Vec<String> { | ||
| vec!["grid_thws".to_string()] |
There was a problem hiding this comment.
🔴 Important: Key mismatch — "grid_thws" here doesn't match the actual tensor key emitted by the preprocessor.
KimiK25Processor delegates to QwenVLProcessorBase::preprocess(), which stores the grid tensor under the key "image_grid_thw" (see qwen_vl_base.rs:400). The Qwen VL and Qwen3 VL specs both use "image_grid_thw" in their field_layouts() and keep_on_cpu_keys() for this reason.
Because of this mismatch:
- The
FieldLayout::Batcheddeclaration for"grid_thws"won't match any tensor, so the actual"image_grid_thw"tensor gets the default layout instead ofBatched. - The
keep_on_cpu_keysentry"grid_thws"won't match either, so the grid tensor will be sent to GPU instead of staying on CPU.
| ("grid_thws".to_string(), FieldLayout::Batched), | |
| ("patches_per_image".to_string(), FieldLayout::Batched), | |
| ]) | |
| } | |
| fn keep_on_cpu_keys(&self) -> Vec<String> { | |
| vec!["grid_thws".to_string()] | |
| ("image_grid_thw".to_string(), FieldLayout::Batched), | |
| ("patches_per_image".to_string(), FieldLayout::Batched), | |
| ]) | |
| } | |
| fn keep_on_cpu_keys(&self) -> Vec<String> { | |
| vec!["image_grid_thw".to_string()] |
![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: 2585607c64
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| "pixel_values".to_string(), | ||
| FieldLayout::flat("patches_per_image"), | ||
| ), | ||
| ("grid_thws".to_string(), FieldLayout::Batched), |
There was a problem hiding this comment.
KimiK25Spec::field_layouts declares a batched tensor named grid_thws, but KimiK25Processor delegates to QwenVLProcessorBase::preprocess, which emits image_grid_thw instead. Because the declared key never appears in model_specific_tensors, vLLM falls back to treating image_grid_thw as a shared tensor rather than per-image batched data; with multi-image prompts this misroutes grid metadata and can produce incorrect multimodal packing. Please use the actual emitted key name (and matching CPU hint key) in the spec.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/multimodal/src/registry/kimi_k25.rs`:
- Around line 79-94: The field-layout and keep-on-CPU key names in KimiK25Spec
are wrong: update the key used in field_layouts (function field_layouts) and
keep_on_cpu_keys (method keep_on_cpu_keys) from "grid_thws" to the actual
emitted tensor name "image_grid_thw" so the FieldLayout::Batched hint and CPU
pinning apply to the tensor produced by QwenVLBase (the mismatch between
"grid_thws" and "image_grid_thw" is causing Kimi tensor handling to fail).
🪄 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: 9e906e91-8134-4d62-8808-9eb35d7efcc6
📒 Files selected for processing (5)
crates/multimodal/src/registry/kimi_k25.rscrates/multimodal/src/registry/mod.rscrates/multimodal/src/vision/image_processor.rscrates/multimodal/src/vision/processors/kimi_k25.rscrates/multimodal/src/vision/processors/mod.rs
| fn field_layouts(&self) -> HashMap<String, FieldLayout> { | ||
| // pixel_values is patchified: [total_patches, C, patch_h, patch_w]. | ||
| // grid_thws is [num_images, 3] with (t, h, w) per image. | ||
| 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<String> { | ||
| vec!["grid_thws".to_string()] | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify actual model_specific key naming in QwenVL base and spec contracts.
rg -n -C3 'image_grid_thw|grid_thws|patches_per_image|model_specific' crates/multimodal/src/vision/processors/qwen_vl_base.rs
rg -n -C3 'field_layouts|keep_on_cpu_keys|image_grid_thw|grid_thws|patches_per_image' crates/multimodal/src/registry/qwen_vl.rs crates/multimodal/src/registry/kimi_k25.rsRepository: lightseekorg/smg
Length of output: 5511
Field-layout key name mismatch: grid_thws vs image_grid_thw will break Kimi tensor handling.
The KimiK25Spec declares field layouts and keep-on-CPU rules for "grid_thws" (lines 87, 93), but the underlying QwenVLBase processor emits "image_grid_thw" (qwen_vl_base.rs:400). This key mismatch prevents layout hints from applying to the actual tensor, breaking Kimi request assembly.
🔧 Proposed fix
fn field_layouts(&self) -> HashMap<String, FieldLayout> {
// pixel_values is patchified: [total_patches, C, patch_h, patch_w].
- // grid_thws is [num_images, 3] with (t, h, w) per image.
+ // image_grid_thw is [num_images, 3] with (t, h, w) per image.
HashMap::from([
(
"pixel_values".to_string(),
FieldLayout::flat("patches_per_image"),
),
- ("grid_thws".to_string(), FieldLayout::Batched),
+ ("image_grid_thw".to_string(), FieldLayout::Batched),
("patches_per_image".to_string(), FieldLayout::Batched),
])
}
fn keep_on_cpu_keys(&self) -> Vec<String> {
- vec!["grid_thws".to_string()]
+ vec!["image_grid_thw".to_string()]
}📝 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 field_layouts(&self) -> HashMap<String, FieldLayout> { | |
| // pixel_values is patchified: [total_patches, C, patch_h, patch_w]. | |
| // grid_thws is [num_images, 3] with (t, h, w) per image. | |
| 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<String> { | |
| vec!["grid_thws".to_string()] | |
| } | |
| fn field_layouts(&self) -> HashMap<String, FieldLayout> { | |
| // pixel_values is patchified: [total_patches, C, patch_h, patch_w]. | |
| // image_grid_thw is [num_images, 3] with (t, h, w) per image. | |
| HashMap::from([ | |
| ( | |
| "pixel_values".to_string(), | |
| FieldLayout::flat("patches_per_image"), | |
| ), | |
| ("image_grid_thw".to_string(), FieldLayout::Batched), | |
| ("patches_per_image".to_string(), FieldLayout::Batched), | |
| ]) | |
| } | |
| fn keep_on_cpu_keys(&self) -> Vec<String> { | |
| vec!["image_grid_thw".to_string()] | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@crates/multimodal/src/registry/kimi_k25.rs` around lines 79 - 94, The
field-layout and keep-on-CPU key names in KimiK25Spec are wrong: update the key
used in field_layouts (function field_layouts) and keep_on_cpu_keys (method
keep_on_cpu_keys) from "grid_thws" to the actual emitted tensor name
"image_grid_thw" so the FieldLayout::Batched hint and CPU pinning apply to the
tensor produced by QwenVLBase (the mismatch between "grid_thws" and
"image_grid_thw" is causing Kimi tensor handling to fail).
|
Hi @ConnorLi96, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
2 participants
Problem
SMG has no multimodal support for Kimi-K2.5, which uses MoonViT3d with a NaViT-style architecture.
Solution
Add
KimiK25Specto the model registry andKimiK25Processorto the image processor registry. The processor reusesQwenVLProcessorBasewith Kimi-specific defaults (patch_size=14, merge_size=2, mean/std=[0.5,0.5,0.5], max 16384 patches).Prompt replacements keep 1 placeholder per image — TRT-LLM's
KimiK25InputProcessorhandles expansion to N vision tokens server-side.Changes
crates/multimodal/src/registry/kimi_k25.rs(new): model spec with config-driven placeholder token, field layouts, testscrates/multimodal/src/registry/mod.rs: register speccrates/multimodal/src/vision/processors/kimi_k25.rs(new): image processor delegating to QwenVLProcessorBasecrates/multimodal/src/vision/processors/mod.rs: exportcrates/multimodal/src/vision/image_processor.rs: register processorTest Plan
cargo clippy+cargo +nightly fmtcleanChecklist:
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesMade with Cursor
Summary by CodeRabbit
Release Notes