[Bugfix][Rust Frontend][Multimodal] Align DeepSeek V4.1 and Kimi K3 media with rendered placeholders (#56366)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Codex <noreply@openai.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-09-14 10:37:59 +00:00
parent 0ca7aef4e3
commit 7b1ea3f524
18 changed files with 415 additions and 85 deletions
@@ -34,8 +34,8 @@ pub use parser::{ParserSelection, validate_parser_overrides};
 pub use renderer::hf::ChatTemplateContentFormatOption;
 pub use renderer::{
    ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer,
-    DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt,
-    RendererSelection,
+    DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, MediaPartSource,
+    RenderedPrompt, RendererSelection,
 };
 pub use request::{
    ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
@@ -112,6 +112,10 @@ impl ChatRequestProcessor {
        request: &ChatRequest,
        rendered: RenderedPrompt,
    ) -> Result<(Prompt, Option<MmFeatures>)> {
+        let media_is_empty = match &rendered.media_order {
+            Some(media_order) => media_order.is_empty(),
+            None => !request.has_multimodal(),
+        };
        match self.model_dtype {
            Some(model_dtype) => {
                multimodal::finalize_rendered_prompt(
@@ -122,7 +126,7 @@ impl ChatRequestProcessor {
                )
                .await
            }
-            None if !request.has_multimodal() => Ok((rendered.prompt, None)),
+            None if media_is_empty => Ok((rendered.prompt, None)),
            None => Err(Error::UnsupportedMultimodalRenderer),
        }
    }
@@ -16,7 +16,7 @@ use std::fs;
 use std::path::Path;
 use std::sync::{Arc, LazyLock};

-use itertools::izip;
+use itertools::{Either, izip};
 use llm_multimodal::{
    AsyncMultiModalTracker, AudioClip, AudioPreProcessor, EncoderFieldLayouts, FieldLayout,
    ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata,
@@ -33,7 +33,7 @@ use vllm_text::Prompt;
 use vllm_text::tokenizer::{DynTokenizer, Tokenizer};

 use crate::error::{Error, Result, bail_multimodal, multimodal};
-use crate::renderer::RenderedPrompt;
+use crate::renderer::{MediaPartSource, RenderedPrompt};
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};

 mod audio;
@@ -549,7 +549,8 @@ pub(crate) async fn finalize_rendered_prompt(
    info: Option<&MultimodalModelInfo>,
    model_dtype: ModelDtype,
 ) -> Result<(Prompt, Option<MmFeatures>)> {
-    if !request.has_multimodal() {
+    let media_parts = extract_media_parts(request, rendered.media_order.as_deref())?;
+    if media_parts.is_empty() {
        return Ok((rendered.prompt, None));
    }
    let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?;
@@ -561,63 +562,92 @@ pub(crate) async fn finalize_rendered_prompt(
            .map_err(|error| multimodal!("{error}"))?,
        Prompt::TokenIds(token_ids) => token_ids,
    };
-    let media_parts = extract_media_parts(request)?;
    let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?;

    Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared)))
 }

-/// Extract media parts from chat messages in message/content order.
-///
-/// Assistant history is skipped because generated assistant blocks are already
-/// represented as text for prompt rendering in this crate.
-fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
-    let mut all_parts = Vec::new();
-    for message in &request.messages {
-        let content = match message {
-            ChatMessage::System { content }
-            | ChatMessage::Developer { content, .. }
-            | ChatMessage::User { content }
-            | ChatMessage::ToolResponse { content, .. } => content,
-            ChatMessage::Assistant { .. } => continue,
-        };
-        let ChatContent::Parts(parts) = content else {
-            continue;
-        };
-        for part in parts {
-            match part {
-                ChatContentPart::Text { .. } => {}
-                ChatContentPart::ImageUrl {
-                    image_url,
-                    detail,
-                    uuid,
-                } => all_parts.push(MediaContentPart::ImageUrl {
-                    url: image_url.clone(),
-                    detail: *detail,
-                    uuid: uuid.clone(),
-                }),
-                ChatContentPart::VideoUrl { video_url, uuid } => {
-                    all_parts.push(MediaContentPart::VideoUrl {
-                        url: video_url.clone(),
-                        uuid: uuid.clone(),
-                    })
-                }
-                ChatContentPart::InputAudio { data, format, uuid } => {
-                    all_parts.push(MediaContentPart::AudioUrl {
-                        url: input_audio_data_url(data, format.as_deref())?,
-                        uuid: uuid.clone(),
-                    })
-                }
-                ChatContentPart::AudioUrl { audio_url, uuid } => {
-                    all_parts.push(MediaContentPart::AudioUrl {
-                        url: audio_url.clone(),
-                        uuid: uuid.clone(),
-                    })
+/// Resolve media parts in the placeholder order reported by the renderer.
+fn extract_media_parts(
+    request: &ChatRequest,
+    media_order: Option<&[MediaPartSource]>,
+) -> Result<Vec<MediaContentPart>> {
+    let parts = match media_order {
+        // Resolve exactly the sources selected by the renderer, in placeholder order.
+        Some(media_order) => Either::Left(media_order.iter().map(|source| -> Result<_> {
+            let message = request.messages.get(source.message_index).ok_or_else(|| {
+                multimodal!(
+                    "renderer reported missing chat message {} for multimodal content",
+                    source.message_index
+                )
+            })?;
+            let content = match message {
+                ChatMessage::System { content }
+                | ChatMessage::Developer { content, .. }
+                | ChatMessage::User { content }
+                | ChatMessage::ToolResponse { content, .. } => content,
+                ChatMessage::Assistant { .. } => {
+                    bail_multimodal!("renderer reported multimodal assistant content")
                }
+            };
+            let ChatContent::Parts(parts) = content else {
+                bail_multimodal!("renderer reported multimodal content in a text-only chat message")
+            };
+            let part = parts.get(source.content_part_index).ok_or_else(|| {
+                multimodal!(
+                    "renderer reported missing content part {} in chat message {}",
+                    source.content_part_index,
+                    source.message_index
+                )
+            })?;
+            Ok(part)
+        })),
+        // Fall back to all media in request message and content-part order (e.g. Jinja).
+        None => Either::Right(request.messages.iter().flat_map(|message| {
+            let content = match message {
+                ChatMessage::System { content }
+                | ChatMessage::Developer { content, .. }
+                | ChatMessage::User { content }
+                | ChatMessage::ToolResponse { content, .. } => Some(content),
+                ChatMessage::Assistant { .. } => None,
+            };
+            match content {
+                Some(ChatContent::Parts(parts)) => parts.as_slice(),
+                _ => &[],
            }
-        }
-    }
-    Ok(all_parts)
+            .iter()
+            .filter(|part| part.is_multimodal())
+            .map(Ok)
+        })),
+    };
+    parts
+        .map(|part| match part? {
+            ChatContentPart::Text { .. } => {
+                bail_multimodal!("renderer reported a text part as multimodal content")
+            }
+            ChatContentPart::ImageUrl {
+                image_url,
+                detail,
+                uuid,
+            } => Ok(MediaContentPart::ImageUrl {
+                url: image_url.clone(),
+                detail: *detail,
+                uuid: uuid.clone(),
+            }),
+            ChatContentPart::VideoUrl { video_url, uuid } => Ok(MediaContentPart::VideoUrl {
+                url: video_url.clone(),
+                uuid: uuid.clone(),
+            }),
+            ChatContentPart::InputAudio { data, format, uuid } => Ok(MediaContentPart::AudioUrl {
+                url: input_audio_data_url(data, format.as_deref())?,
+                uuid: uuid.clone(),
+            }),
+            ChatContentPart::AudioUrl { audio_url, uuid } => Ok(MediaContentPart::AudioUrl {
+                url: audio_url.clone(),
+                uuid: uuid.clone(),
+            }),
+        })
+        .collect()
 }

 /// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`.
@@ -1043,6 +1073,75 @@ mod tests {
        assert!(input_audio_data_url("AAAA", Some("flac")).is_err());
    }

+    #[test]
+    fn extract_media_parts_uses_request_order_or_explicit_selection() {
+        let request = ChatRequest {
+            messages: vec![
+                ChatMessage::user(vec![
+                    ChatContentPart::text("before image"),
+                    ChatContentPart::ImageUrl {
+                        image_url: "data:image/png;base64,image-b".to_string(),
+                        detail: None,
+                        uuid: Some("image-b".to_string()),
+                    },
+                ]),
+                ChatMessage::user(vec![ChatContentPart::ImageUrl {
+                    image_url: "data:image/png;base64,image-a".to_string(),
+                    detail: None,
+                    uuid: Some("image-a".to_string()),
+                }]),
+                ChatMessage::assistant_text("answer"),
+                ChatMessage::user("follow-up"),
+            ],
+            ..ChatRequest::for_test()
+        };
+        let media_order = [
+            MediaPartSource {
+                message_index: 1,
+                content_part_index: 0,
+            },
+            MediaPartSource {
+                message_index: 0,
+                content_part_index: 1,
+            },
+        ];
+
+        let selections = [None, Some(&media_order[..]), Some(&[][..])];
+        let uuids = selections.map(|order| {
+            extract_media_parts(&request, order)
+                .unwrap()
+                .into_iter()
+                .map(|part| match part {
+                    MediaContentPart::ImageUrl { uuid, .. } => uuid,
+                    _ => panic!("expected image URL"),
+                })
+                .collect::<Vec<_>>()
+        });
+
+        expect_test::expect![[r#"
+            [
+                [
+                    Some(
+                        "image-b",
+                    ),
+                    Some(
+                        "image-a",
+                    ),
+                ],
+                [
+                    Some(
+                        "image-a",
+                    ),
+                    Some(
+                        "image-b",
+                    ),
+                ],
+                [],
+            ]
+        "#]]
+        .assert_debug_eq(&uuids);
+    }
+
    fn image_url_part() -> MediaContentPart {
        MediaContentPart::ImageUrl {
            url: "https://example.com/image.png".to_string(),
@@ -16,6 +16,7 @@ use serde_json_fmt::JsonFormat;

 use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER;

+use super::MediaPartSource;
 use crate::error::{Error, Result};
 use crate::request::{
    ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ReasoningEffort,
@@ -87,6 +88,14 @@ struct RenderedToolSchema<'a> {

 /// Render one chat request into the final prompt string.
 pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Result<String> {
+    render_request_with_media_order(request, dialect).map(|(prompt, _)| prompt)
+}
+
+/// Render one chat request and record media in placeholder order.
+pub(super) fn render_request_with_media_order(
+    request: &ChatRequest,
+    dialect: DsDialect,
+) -> Result<(String, Vec<MediaPartSource>)> {
    let (thinking_mode, reasoning_effort_prompt) = match dialect {
        DsDialect::V4 => {
            resolve_thinking_options(request).map(|(mode, prompt)| (mode, Cow::Borrowed(prompt)))?
@@ -112,6 +121,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
        dialect,
    );
    let mut out = String::from(BOS_TOKEN);
+    let mut media_order = Vec::new();
    if dialect == DsDialect::V41
        && (thinking_mode == ThinkingMode::Thinking
            || synthetic_tool_system
@@ -126,7 +136,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
    let mut request_tools_attached = false;
    let mut render_index = 0isize;
    if synthetic_tool_system {
-        render_system_message(&mut out, None, request_tools, dialect)?;
+        render_system_message(&mut out, None, request_tools, dialect, &mut media_order)?;
        request_tools_attached = true;
        render_index += 1;
    }
@@ -157,14 +167,22 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
                } else {
                    &[]
                };
-                render_system_message(&mut out, Some(content), tools, dialect)?;
+                render_system_message(
+                    &mut out,
+                    Some((message_index, content)),
+                    tools,
+                    dialect,
+                    &mut media_order,
+                )?;
            }
            ChatMessage::Developer { content, tools } => {
                render_developer_message(
                    &mut out,
+                    message_index,
                    content,
                    tools.as_deref().unwrap_or(&[]),
                    dialect,
+                    &mut media_order,
                )?;
            }
            ChatMessage::User { .. } | ChatMessage::ToolResponse { .. } => {
@@ -174,6 +192,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
                    message_index,
                    &last_tool_call_order,
                    dialect,
+                    &mut media_order,
                )?;
            }
            ChatMessage::Assistant { content } => {
@@ -223,7 +242,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
        }
    }

-    Ok(out)
+    Ok((out, media_order))
 }

 /// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer
@@ -488,12 +507,13 @@ fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> {
 /// Render a system turn, optionally followed by the V4 tool preamble.
 fn render_system_message(
    out: &mut String,
-    content: Option<&ChatContent>,
+    content: Option<(usize, &ChatContent)>,
    tools: &[ChatTool],
    dialect: DsDialect,
+    media_order: &mut Vec<MediaPartSource>,
 ) -> Result<()> {
-    if let Some(content) = content {
-        write_chat_content(out, content, dialect)?;
+    if let Some((message_index, content)) = content {
+        write_chat_content(out, content, dialect, message_index, media_order)?;
    }
    if !tools.is_empty() {
        out.push_str("\n\n");
@@ -505,9 +525,11 @@ fn render_system_message(
 /// Developer messages are rendered as user-like turns with optional tools.
 fn render_developer_message(
    out: &mut String,
+    message_index: usize,
    content: &ChatContent,
    tools: &[ChatTool],
    dialect: DsDialect,
+    media_order: &mut Vec<MediaPartSource>,
 ) -> Result<()> {
    if content.is_empty() {
        return Err(Error::ChatTemplate(
@@ -516,7 +538,7 @@ fn render_developer_message(
    }

    out.push_str(USER_SP_TOKEN);
-    write_chat_content(out, content, dialect)?;
+    write_chat_content(out, content, dialect, message_index, media_order)?;
    if !tools.is_empty() {
        out.push_str("\n\n");
        render_tools(out, tools, dialect)?;
@@ -531,6 +553,7 @@ fn render_user_content_block(
    message_index: usize,
    tool_call_order: &HashMap<String, usize>,
    dialect: DsDialect,
+    media_order: &mut Vec<MediaPartSource>,
 ) -> Result<()> {
    let (block_start, block_end) = user_content_block_bounds(messages, message_index);
    let mut sorted_tool_indices =
@@ -542,7 +565,9 @@ fn render_user_content_block(
            out.push_str("\n\n");
        }
        match &messages[message_index] {
-            ChatMessage::User { content } => write_chat_content(out, content, dialect)?,
+            ChatMessage::User { content } => {
+                write_chat_content(out, content, dialect, message_index, media_order)?;
+            }
            ChatMessage::ToolResponse { .. } => {
                let sorted_index = sorted_tool_indices
                    .next()
@@ -550,7 +575,7 @@ fn render_user_content_block(
                let ChatMessage::ToolResponse { content, .. } = &messages[sorted_index] else {
                    unreachable!("sorted tool response index should reference a tool message");
                };
-                write_tool_result(out, content, dialect)?;
+                write_tool_result(out, content, dialect, sorted_index, media_order)?;
            }
            _ => unreachable!("user content block should only contain user content messages"),
        }
@@ -597,9 +622,15 @@ fn sorted_tool_response_indices(
 }

 /// Render one tool response payload inside a V4 `<tool_result>` block.
-fn write_tool_result(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> {
+fn write_tool_result(
+    out: &mut String,
+    content: &ChatContent,
+    dialect: DsDialect,
+    message_index: usize,
+    media_order: &mut Vec<MediaPartSource>,
+) -> Result<()> {
    out.push_str("<tool_result>");
-    write_chat_content(out, content, dialect)?;
+    write_chat_content(out, content, dialect, message_index, media_order)?;
    out.push_str("</tool_result>");
    Ok(())
 }
@@ -727,7 +758,13 @@ fn encode_arguments_to_dsml(
 /// The V4.1 dialect inlines the image placeholder at each image part's
 /// position unconditionally (matching the Python encoding's
 /// `IMAGE_PLACEHOLDER`); other dialects reject multimodal parts.
-fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> {
+fn write_chat_content(
+    out: &mut String,
+    content: &ChatContent,
+    dialect: DsDialect,
+    message_index: usize,
+    media_order: &mut Vec<MediaPartSource>,
+) -> Result<()> {
    match content {
        ChatContent::Text(text) => out.push_str(text),
        ChatContent::Parts(parts) => {
@@ -738,6 +775,10 @@ fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialec
                match part {
                    ChatContentPart::ImageUrl { .. } if dialect == DsDialect::V41 => {
                        out.push_str(DEEPSEEK_V41_IMAGE_PLACEHOLDER);
+                        media_order.push(MediaPartSource {
+                            message_index,
+                            content_part_index: index,
+                        });
                    }
                    _ => out.push_str(part.as_text()?),
                }
@@ -26,6 +26,7 @@ impl ChatRenderer for DeepSeekV32ChatRenderer {

        Ok(RenderedPrompt {
            prompt: Prompt::Text(encoding::render_request(request)?),
+            media_order: None,
            effective_template_kwargs: request_template_kwargs(request),
        })
    }
@@ -24,6 +24,7 @@ impl ChatRenderer for DeepSeekV4ChatRenderer {

        Ok(RenderedPrompt {
            prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V4)?),
+            media_order: None,
            effective_template_kwargs: request_template_kwargs(request),
        })
    }
@@ -21,9 +21,12 @@ impl DeepSeekV41ChatRenderer {
 impl ChatRenderer for DeepSeekV41ChatRenderer {
    fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
        request.validate()?;
+        let (prompt, media_order) =
+            deepseek::render_request_with_media_order(request, DsDialect::V41)?;

        Ok(RenderedPrompt {
-            prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V41)?),
+            prompt: Prompt::Text(prompt),
+            media_order: Some(media_order),
            effective_template_kwargs: request_template_kwargs(request),
        })
    }
@@ -8,7 +8,7 @@ use serde_json::{Value, json};

 use super::DeepSeekV41ChatRenderer;
 use crate::ChatRenderer;
-use crate::event::AssistantContentBlock;
+use crate::event::{AssistantContentBlock, AssistantToolCall};
 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort};

@@ -83,6 +83,82 @@ fn inlines_image_placeholder_at_image_part_positions() {
    );
 }

+#[test]
+fn media_order_follows_reordered_tool_results() {
+    let renderer = DeepSeekV41ChatRenderer::new();
+    let request = ChatRequest {
+        messages: vec![
+            ChatMessage::assistant_blocks(vec![
+                AssistantContentBlock::ToolCall(AssistantToolCall {
+                    id: "call-a".to_string(),
+                    name: "tool-a".to_string(),
+                    arguments: "{}".to_string(),
+                }),
+                AssistantContentBlock::ToolCall(AssistantToolCall {
+                    id: "call-b".to_string(),
+                    name: "tool-b".to_string(),
+                    arguments: "{}".to_string(),
+                }),
+            ]),
+            ChatMessage::tool_response(
+                vec![ChatContentPart::text("result-b"), image_part("image-b")],
+                "call-b",
+            ),
+            ChatMessage::tool_response(
+                vec![ChatContentPart::text("result-a"), image_part("image-a")],
+                "call-a",
+            ),
+        ],
+        ..ChatRequest::for_test()
+    };
+
+    let rendered = renderer.render(&request).unwrap();
+    let media_order = rendered.media_order.unwrap();
+    let prompt = rendered.prompt.into_text().unwrap();
+
+    assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap());
+    expect![[r#"
+        [
+            MediaPartSource {
+                message_index: 2,
+                content_part_index: 1,
+            },
+            MediaPartSource {
+                message_index: 1,
+                content_part_index: 1,
+            },
+        ]
+    "#]]
+    .assert_debug_eq(&media_order);
+}
+
+#[test]
+fn media_order_excludes_dropped_historical_developer_content() {
+    let renderer = DeepSeekV41ChatRenderer::new();
+    let request = ChatRequest {
+        messages: vec![
+            ChatMessage::developer(vec![image_part("historical-image")], None),
+            ChatMessage::user("question"),
+        ],
+        ..ChatRequest::for_test()
+    };
+
+    let rendered = renderer.render(&request).unwrap();
+    let media_order = rendered.media_order.unwrap();
+    let prompt = rendered.prompt.into_text().unwrap();
+
+    assert!(!prompt.contains("<｜deepseek_image｜>"));
+    assert!(media_order.is_empty());
+}
+
+fn image_part(uuid: &str) -> ChatContentPart {
+    ChatContentPart::ImageUrl {
+        image_url: format!("data:image/png;base64,{uuid}"),
+        detail: None,
+        uuid: Some(uuid.to_string()),
+    }
+}
+
 #[test]
 fn maps_reasoning_effort_to_reference_numeric_budget() {
    for (effort, budget) in [
@@ -116,6 +116,7 @@ impl ChatRenderer for HarmonyChatRenderer {
    fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
        Ok(RenderedPrompt {
            prompt: Prompt::TokenIds(self.render_token_ids(request)?),
+            media_order: None,
            effective_template_kwargs: request_template_kwargs(request),
        })
    }
@@ -217,6 +217,7 @@ impl HfChatRenderer {

        Ok(RenderedPrompt {
            prompt: Prompt::Text(prompt),
+            media_order: None,
            effective_template_kwargs,
        })
    }
@@ -346,6 +346,7 @@ impl ChatRenderer for InklingChatRenderer {

        Ok(RenderedPrompt {
            prompt: Prompt::TokenIds(out),
+            media_order: None,
            effective_template_kwargs,
        })
    }
@@ -10,6 +10,7 @@ use std::collections::HashMap;
 use serde_json::{Map, Value, json};
 use vllm_tokenizer::Tokenizer;

+use super::super::MediaPartSource;
 use crate::error::{Error, Result};
 use crate::request::{
    ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice,
@@ -29,6 +30,7 @@ const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"];
 pub(super) struct K3TokenWriter<'a> {
    tokenizer: &'a dyn Tokenizer,
    token_ids: Vec<u32>,
+    media_order: Vec<MediaPartSource>,
 }

 impl<'a> K3TokenWriter<'a> {
@@ -36,6 +38,7 @@ impl<'a> K3TokenWriter<'a> {
        Self {
            tokenizer,
            token_ids: Vec::new(),
+            media_order: Vec::new(),
        }
    }

@@ -55,13 +58,25 @@ impl<'a> K3TokenWriter<'a> {
        Ok(())
    }

-    pub(super) fn finish(self) -> Vec<u32> {
-        self.token_ids
+    fn image(&mut self, message_index: usize, content_part_index: usize) -> Result<()> {
+        self.control(IMAGE_PLACEHOLDER)?;
+        self.media_order.push(MediaPartSource {
+            message_index,
+            content_part_index,
+        });
+        Ok(())
+    }
+
+    pub(super) fn finish(self) -> (Vec<u32>, Vec<MediaPartSource>) {
+        (self.token_ids, self.media_order)
    }
 }

-/// Render and tokenize one chat request using K3's segment-aware contract.
-pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result<Vec<u32>> {
+/// Render and tokenize one chat request while recording media in placeholder order.
+pub(super) fn render_request_with_media_order(
+    request: &ChatRequest,
+    tokenizer: &dyn Tokenizer,
+) -> Result<(Vec<u32>, Vec<MediaPartSource>)> {
    let thinking = thinking_enabled(request)?;
    let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?;
    let tools = request_tools(request);
@@ -126,7 +141,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
            resolved.sort_by_key(|item| (item.0, item.1));
        }

-        for (position, _, content, name) in resolved {
+        for (position, message_index, content, name) in resolved {
            let tool_name = name.as_deref().ok_or_else(|| {
                Error::ChatTemplate(
                    "Kimi K3 tool messages need a resolvable tool name: \
@@ -135,7 +150,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
                        .to_string(),
                )
            })?;
-            write_tool_message(out, tool_name, position, &content)?;
+            write_tool_message(out, tool_name, position, message_index, &content)?;
        }
        Ok(())
    };
@@ -161,14 +176,14 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
            } if !local_tools.is_empty() => {
                write_tool_declare(&mut out, local_tools, true)?;
                if !content_is_empty(content) {
-                    write_role_message(&mut out, "system", None, content)?;
+                    write_role_message(&mut out, "system", None, message_index, content)?;
                }
            }
            ChatMessage::System { content } | ChatMessage::Developer { content, .. } => {
-                write_role_message(&mut out, "system", None, content)?;
+                write_role_message(&mut out, "system", None, message_index, content)?;
            }
            ChatMessage::User { content } => {
-                write_role_message(&mut out, "user", None, content)?;
+                write_role_message(&mut out, "user", None, message_index, content)?;
            }
            ChatMessage::Assistant { content } => {
                tool_call_id_index.clear();
@@ -364,6 +379,7 @@ fn write_role_message(
    out: &mut K3TokenWriter<'_>,
    role: &str,
    name: Option<&str>,
+    message_index: usize,
    content: &ChatContent,
 ) -> Result<()> {
    let mut attrs = vec![("role", role.to_string())];
@@ -372,7 +388,7 @@ fn write_role_message(
    }
    let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect();
    write_open_tag(out, "message", &attr_refs)?;
-    write_content(out, content)?;
+    write_content(out, message_index, content)?;
    write_close_tag(out, "message")?;
    out.control(END_OF_MSG)
 }
@@ -381,6 +397,7 @@ fn write_tool_message(
    out: &mut K3TokenWriter<'_>,
    tool_name: &str,
    index: usize,
+    message_index: usize,
    content: &ChatContent,
 ) -> Result<()> {
    let index_str = index.to_string();
@@ -389,7 +406,7 @@ fn write_tool_message(
        "message",
        &[("role", "tool"), ("tool", tool_name), ("index", &index_str)],
    )?;
-    write_content(out, content)?;
+    write_content(out, message_index, content)?;
    write_close_tag(out, "message")?;
    out.control(END_OF_MSG)
 }
@@ -471,14 +488,20 @@ fn write_assistant_tool_call(
    write_close_tag(out, "call")
 }

-fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> {
+fn write_content(
+    out: &mut K3TokenWriter<'_>,
+    message_index: usize,
+    content: &ChatContent,
+) -> Result<()> {
    match content {
        ChatContent::Text(text) => write_text_with_images(out, text),
        ChatContent::Parts(parts) => {
-            for part in parts {
+            for (content_part_index, part) in parts.iter().enumerate() {
                match part {
                    ChatContentPart::Text { text } => write_text_with_images(out, text)?,
-                    ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?,
+                    ChatContentPart::ImageUrl { .. } => {
+                        out.image(message_index, content_part_index)?;
+                    }
                    ChatContentPart::VideoUrl { .. } => {
                        return Err(Error::UnsupportedMultimodalContent("video_url"));
                    }
@@ -30,9 +30,12 @@ impl KimiK3ChatRenderer {
 impl ChatRenderer for KimiK3ChatRenderer {
    fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
        request.validate()?;
+        let (token_ids, media_order) =
+            encoding::render_request_with_media_order(request, self.tokenizer.as_ref())?;

        Ok(RenderedPrompt {
-            prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?),
+            prompt: Prompt::TokenIds(token_ids),
+            media_order: Some(media_order),
            effective_template_kwargs: request_template_kwargs(request),
        })
    }
@@ -284,6 +284,65 @@ fn partial_tool_results_keep_assistant_call_position() {
    );
 }

+#[test]
+fn media_order_follows_reordered_tool_results() {
+    let mut request = crate::request::ChatRequest::for_test();
+    request.messages = vec![
+        ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::ToolCall(AssistantToolCall {
+                id: "call-a".to_string(),
+                name: "tool-a".to_string(),
+                arguments: "{}".to_string(),
+            }),
+            AssistantContentBlock::ToolCall(AssistantToolCall {
+                id: "call-b".to_string(),
+                name: "tool-b".to_string(),
+                arguments: "{}".to_string(),
+            }),
+        ]),
+        ChatMessage::tool_response(
+            vec![ChatContentPart::text("result-b"), image_part("image-b")],
+            "call-b",
+        ),
+        ChatMessage::tool_response(
+            vec![ChatContentPart::text("result-a"), image_part("image-a")],
+            "call-a",
+        ),
+    ];
+    request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+    let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()));
+
+    let rendered = renderer.render(&request).unwrap();
+    let media_order = rendered.media_order.unwrap();
+    let Prompt::TokenIds(token_ids) = rendered.prompt else {
+        panic!("kimi k3 renderer should return token IDs")
+    };
+    let prompt = test_tokenizer().decode(&token_ids, false).unwrap();
+
+    assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap());
+    expect![[r#"
+        [
+            MediaPartSource {
+                message_index: 2,
+                content_part_index: 1,
+            },
+            MediaPartSource {
+                message_index: 1,
+                content_part_index: 1,
+            },
+        ]
+    "#]]
+    .assert_debug_eq(&media_order);
+}
+
+fn image_part(uuid: &str) -> ChatContentPart {
+    ChatContentPart::ImageUrl {
+        image_url: format!("data:image/png;base64,{uuid}"),
+        detail: None,
+        uuid: Some(uuid.to_string()),
+    }
+}
+
 #[test]
 fn defaults_thinking_effort_to_max() {
    let rendered = render_request(&crate::request::ChatRequest::for_test());
@@ -30,11 +30,24 @@ pub use inkling::InklingChatRenderer;
 pub use kimi_k3::KimiK3ChatRenderer;
 pub use selection::RendererSelection;

+/// Location of one multimodal content part in the source chat request.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct MediaPartSource {
+    /// Zero-based message index in the source request.
+    pub message_index: usize,
+    /// Zero-based content-part index within the source message.
+    pub content_part_index: usize,
+}
+
 /// Rendered chat prompt submitted to the text backend.
 #[derive(Debug, Clone, PartialEq)]
 pub struct RenderedPrompt {
    /// The rendered prompt, either as text or already tokenized.
    pub prompt: Prompt,
+    /// Media sources in rendered placeholder order.
+    /// `None` uses the message and content-part order in the request (e.g. Jinja).
+    /// `Some` uses exactly the listed sources; an empty list omits all media.
+    pub media_order: Option<Vec<MediaPartSource>>,
    /// Effective chat-template kwargs visible to the renderer after applying
    /// server defaults, request overrides, and typed reasoning controls.
    pub effective_template_kwargs: HashMap<String, Value>,
@@ -250,6 +250,7 @@ impl ChatRenderer for FakeChatBackend {

        Ok(RenderedPrompt {
            prompt: Prompt::Text(prompt),
+            media_order: None,
            effective_template_kwargs: request.chat_options.template_kwargs.clone(),
        })
    }
@@ -252,6 +252,7 @@ impl ChatRenderer for FakeTextBackend {
    fn render(&self, _request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
        Ok(RenderedPrompt {
            prompt: Prompt::Text(String::new()),
+            media_order: None,
            effective_template_kwargs: Default::default(),
        })
    }
@@ -186,6 +186,7 @@ impl ChatRenderer for FakeChatBackend {
        }
        Ok(RenderedPrompt {
            prompt: Prompt::Text(prompt),
+            media_order: None,
            effective_template_kwargs: Default::default(),
        })
    }
@@ -540,6 +540,7 @@ impl ChatRenderer for FakeChatBackend {
        }
        Ok(vllm_chat::RenderedPrompt {
            prompt: Prompt::Text(prompt),
+            media_order: None,
            effective_template_kwargs: Default::default(),
        })
    }