diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index 1268bdaae158..60ac61a11e21 100644 --- a/components/src/dynamo/sglang/register.py +++ b/components/src/dynamo/sglang/register.py @@ -376,12 +376,6 @@ async def _get_runtime_config( runtime_config.enable_local_indexer = ( dynamo_args.enable_local_indexer and not is_decode_worker ) - # SGLang's multimodal processors expand a single image pad token to the - # image's feature count and never rebuild the model's native media - # sequence, so the frontend must send the pad rather than the placeholder - # marker. Engines that re-derive the sequence from the marker leave this - # false. - runtime_config.expands_image_pad_token = True start_dp_rank, end_dp_rank = model_card_dp_rank_bounds(server_args) registered_dp_size = end_dp_rank - start_dp_rank diff --git a/components/src/dynamo/vllm/multimodal_utils/request_processor.py b/components/src/dynamo/vllm/multimodal_utils/request_processor.py index 8bf8bacd7420..233097f357e8 100644 --- a/components/src/dynamo/vllm/multimodal_utils/request_processor.py +++ b/components/src/dynamo/vllm/multimodal_utils/request_processor.py @@ -255,6 +255,104 @@ def __init__( ) self.use_unified_vision_chunk = use_unified_vision_chunk + def _kimi_k3_pad_expansion(self) -> Optional[tuple[int, list[int]]]: + """``(pad_id, native_ids)`` for Kimi-K3, or None for every other model. + + The frontend emits one ``<|media_pad|>`` per image for every engine. + vLLM's K3 processor instead matches the checkpoint's + ``<|kimi_image_placeholder|>`` and expands *that* into + ``<|media_begin|>image WxH<|media_content|>...<|media_end|>``, so the + two forms are converted here rather than diverging in the renderer. + + Converting in this direction is what makes the contract reliable: + ``<|media_pad|>`` is a single vocabulary id, so locating it is exact. + ``<|kimi_image_placeholder|>`` is a plain string that is not in the + vocabulary -- its token boundaries shift with surrounding text -- so a + frontend emitting it would leave the worker nothing dependable to find. + + Both values come from checkpoint metadata, so no registration field or + engine introspection is required. + """ + if getattr(self, "_k3_expansion_resolved", False): + return self._k3_expansion + self._k3_expansion_resolved = True + self._k3_expansion = None + try: + model_config = getattr( + getattr(self.engine_client, "vllm_config", None), "model_config", None + ) + hf_config = getattr(model_config, "hf_config", None) + if getattr(hf_config, "model_type", None) != "kimi_k3": + return None + pad_id = hf_config.media_placeholder_token_id + tokenizer = self.engine_client.get_tokenizer() + native_ids = list( + tokenizer.encode(hf_config.image_placeholder, add_special_tokens=False) + ) + if not isinstance(pad_id, int) or not native_ids: + logger.warning( + "Kimi-K3 placeholder metadata unusable (pad_id=%r, native_ids=%r); " + "leaving prompt token ids untouched", + pad_id, + native_ids, + ) + return None + self._k3_expansion = (pad_id, native_ids) + except Exception as e: + logger.warning("Could not resolve the Kimi-K3 placeholder mapping: %s", e) + return self._k3_expansion + + def _expand_kimi_k3_pads( + self, token_ids: list[int], multi_modal_data: Optional[dict[str, Any]] + ) -> list[int]: + """Replace each structural pad id with the checkpoint-native sequence. + + Raw-media path only -- inputs that already carry processed multimodal + state are left alone. During rollout a prompt may already be in the + native form, in which case there are no pads and this is a no-op. + """ + expansion = self._kimi_k3_pad_expansion() + if expansion is None or not multi_modal_data: + return token_ids + pad_id, native_ids = expansion + + # Prompts here reach 100k+ tokens while pads number in the single + # digits, so locate the rare token and splice around it rather than + # walking every id in Python. `list.index` short-circuits at the first + # pad, so the no-pad case (already-native prompt) costs one C-level + # scan and zero copies. + try: + first = token_ids.index(pad_id) + except ValueError: + return token_ids # already in native form -- nothing to replace + + pad_positions = [first] + while True: + try: + pad_positions.append(token_ids.index(pad_id, pad_positions[-1] + 1)) + except ValueError: + break + + images = multi_modal_data.get("image") + expected = len(images) if isinstance(images, (list, tuple)) else 1 + if len(pad_positions) != expected: + # Guessing here would silently misalign images against their + # embedding slots, so refuse instead. + raise ValueError( + f"Kimi-K3 prompt carries {len(pad_positions)} <|media_pad|> " + f"token(s) but {expected} image(s) were supplied; refusing to " + f"expand." + ) + + expanded: list[int] = [] + prev = 0 + for position in pad_positions: + expanded.extend(token_ids[prev:position]) + expanded.extend(native_ids) + prev = position + 1 + expanded.extend(token_ids[prev:]) + return expanded + @staticmethod def _multimodal_disabled_error() -> ValueError: return ValueError( @@ -562,7 +660,9 @@ def build_tokens_prompt( ) prompt_kwargs: dict[str, Any] = { - "prompt_token_ids": request["token_ids"], + "prompt_token_ids": self._expand_kimi_k3_pads( + request["token_ids"], multi_modal_data + ), "multi_modal_data": multi_modal_data, } if mm_uuids is not None: diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_request_processor.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_request_processor.py index 635ef83f635b..1fe2ce62d232 100644 --- a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_request_processor.py +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_request_processor.py @@ -819,3 +819,82 @@ def test_qwen_handoff_accepts_encoder_embeddings(): "image_grid_thw": [[1, 16, 16]], "embeddings_shape": [1, 256, 1024], } + + +# --- Kimi-K3 structural-pad -> checkpoint-native expansion ------------------- +# +# The frontend emits one <|media_pad|> per image for every engine. vLLM's K3 +# processor matches the checkpoint's <|kimi_image_placeholder|> instead, so the +# worker converts. See VllmMultimodalRequestProcessor._expand_kimi_k3_pads. + +_K3_PAD_ID = 163605 +_K3_NATIVE_IDS = [27, 91, 74, 30223, 11947, 114136, 91, 29] + + +def _k3_processor(model_type: str = "kimi_k3") -> mod.VllmMultimodalRequestProcessor: + engine_client = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + model_type=model_type, + media_placeholder_token_id=_K3_PAD_ID, + image_placeholder="<|kimi_image_placeholder|>", + use_unified_vision_chunk=False, + ) + ) + ), + get_tokenizer=lambda: SimpleNamespace( + encode=lambda text, add_special_tokens=False: list(_K3_NATIVE_IDS) + ), + ) + return mod.VllmMultimodalRequestProcessor( + model="moonshot-ai/Kimi-K3", + engine_client=engine_client, + enable_multimodal=True, + ) + + +def test_k3_pad_expands_to_the_checkpoint_native_sequence(): + proc = _k3_processor() + out = proc._expand_kimi_k3_pads([1, _K3_PAD_ID, 2], {"image": ["img"]}) + + assert out == [1, *_K3_NATIVE_IDS, 2] + + +def test_k3_expansion_is_one_per_image(): + proc = _k3_processor() + out = proc._expand_kimi_k3_pads([_K3_PAD_ID, 7, _K3_PAD_ID], {"image": ["a", "b"]}) + + assert out.count(_K3_PAD_ID) == 0 + assert out == [*_K3_NATIVE_IDS, 7, *_K3_NATIVE_IDS] + + +def test_k3_mismatched_pad_count_is_rejected_not_guessed(): + """Misaligning images against embedding slots is worse than failing.""" + proc = _k3_processor() + with pytest.raises(ValueError, match="refusing to expand"): + proc._expand_kimi_k3_pads([_K3_PAD_ID], {"image": ["a", "b"]}) + + +def test_k3_already_native_prompt_is_untouched(): + """Rollout compatibility: a frontend still emitting the native form has no + pads to replace, so it passes through.""" + proc = _k3_processor() + native = [1, *_K3_NATIVE_IDS, 2] + + assert proc._expand_kimi_k3_pads(native, {"image": ["img"]}) == native + + +def test_non_k3_models_are_never_rewritten(): + proc = _k3_processor(model_type="qwen3_vl") + tokens = [1, _K3_PAD_ID, 2] + + assert proc._expand_kimi_k3_pads(tokens, {"image": ["img"]}) == tokens + + +def test_k3_without_raw_media_is_untouched(): + """Processed-input paths must not be rewritten.""" + proc = _k3_processor() + tokens = [1, _K3_PAD_ID, 2] + + assert proc._expand_kimi_k3_pads(tokens, None) == tokens diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index cee1679b5b63..751d1d2c8e91 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -157,11 +157,6 @@ impl ModelRuntimeConfig { self.inner.enable_local_indexer = enable_local_indexer; } - #[setter] - fn set_expands_image_pad_token(&mut self, expands_image_pad_token: bool) { - self.inner.expands_image_pad_token = expands_image_pad_token; - } - #[setter] fn set_kv_state_endpoint(&mut self, kv_state_endpoint: Option) { self.inner.kv_state_endpoint = kv_state_endpoint.as_deref().map(EndpointId::from); @@ -245,11 +240,6 @@ impl ModelRuntimeConfig { self.inner.enable_local_indexer } - #[getter] - fn expands_image_pad_token(&self) -> bool { - self.inner.expands_image_pad_token - } - #[getter] fn kv_state_endpoint(&self) -> Option { self.inner diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index d0d2011ec9f6..57ed1b7afed7 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -171,29 +171,6 @@ pub struct ModelRuntimeConfig { #[serde(default = "default_local_indexer")] pub enable_local_indexer: bool, - /// Whether this engine expands a single image pad token to the image's - /// feature count instead of rebuilding the model's native media sequence. - /// - /// Multimodal families whose prompt carries one placeholder marker per - /// image (currently Kimi-K3) need that marker turned into the model's real - /// media sequence before the vision embeddings can bind to it, and engines - /// split on who does it: - /// - /// - `false` (default): the engine re-derives the media sequence from the - /// frontend's placeholder marker, reading image dimensions from the - /// multimodal payload. The preprocessor passes the marker through - /// untouched; pre-substituting here would collide with the engine's own - /// expansion. - /// - `true`: the engine only repeats a pad token up to the feature count - /// and never constructs the media sequence, so it needs exactly one pad - /// token per image in the prompt. The preprocessor substitutes the - /// formatter's `image_pad_token()` for the marker. - /// - /// Ignored for formatters that declare no pad token, which is every family - /// except Kimi-K3. - #[serde(default)] - pub expands_image_pad_token: bool, - /// Endpoint whose event sources describe this worker's KV state. /// /// When unset, consumers use the worker's serving endpoint. This keeps existing @@ -298,7 +275,6 @@ impl Default for ModelRuntimeConfig { exclude_tools_when_tool_choice_none: default_exclude_tools_when_tool_choice_none(), data_parallel_start_rank: default_data_parallel_start_rank(), data_parallel_size: default_data_parallel_size(), - expands_image_pad_token: false, enable_local_indexer: true, kv_state_endpoint: None, runtime_data: HashMap::new(), diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index cb8c90278321..88d139e9557a 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -897,24 +897,13 @@ impl OpenAIPreprocessor { let mut builder = self.builder(request)?; let template_start = Instant::now(); - let mut formatted_prompt = { + let formatted_prompt = { let _nvtx = dynamo_nvtx_range!("preprocess.template"); self.apply_template(request) .with_context(|| "Failed to apply prompt template")? }; TEMPLATE_SECONDS.observe(template_start.elapsed().as_secs_f64()); - // Engines that only repeat a pad token to the image's feature count - // never build the model's native media sequence, so hand them the pad - // instead of the placeholder marker. Engines that rebuild the sequence - // themselves keep the marker (the default) — substituting here would - // leave them nothing to match on. - if self.runtime_config.expands_image_pad_token - && let Some(prompt) = formatted_prompt.as_mut() - { - self.substitute_image_pad_token(prompt); - } - // Generic reasoning parsers start from ``; MiniMax M3 starts // from ``. If the chat template injected that opener at the // end of the prompt, the model completion starts mid-reasoning. @@ -1289,48 +1278,6 @@ impl OpenAIPreprocessor { } } - /// Swap each image-placeholder marker segment for the formatter's pad - /// token, for engines that expand a pad rather than rebuilding the model's - /// media sequence from the marker. - /// - /// A strict no-op for formatters that declare no pad token (every family - /// except Kimi K3) and for requests rendered as raw text, which have no - /// segment boundaries to rewrite. Cardinality is unchanged — one marker - /// becomes one pad — so the engine's own placeholder count still holds. - fn substitute_image_pad_token(&self, prompt: &mut RenderedPrompt) { - let (Some(marker), Some(pad)) = ( - self.formatter.image_placeholder_template(), - self.formatter.image_pad_token(), - ) else { - return; - }; - Self::swap_marker_segments(marker, pad, prompt); - } - - /// Segment-rewriting core of [`Self::substitute_image_pad_token`]. - fn swap_marker_segments(marker: &str, pad: &str, prompt: &mut RenderedPrompt) { - let Some(segments) = prompt.segments() else { - return; // raw_prompt path → no segments to rewrite - }; - if !segments.iter().any(|seg| seg.text == marker) { - return; - } - let rewritten = segments - .iter() - .map(|seg| { - if seg.text == marker { - crate::tokenizers::EncodeSegment { - text: pad.to_string(), - allow_special: true, - } - } else { - seg.clone() - } - }) - .collect(); - *prompt = RenderedPrompt::segmented(rewritten); - } - /// Replace inline `data:` URLs with empty strings in message content parts. /// Preserves HTTP(S) URLs, text content, and overall message structure. fn strip_inline_data_urls(messages: &mut serde_json::Value) { @@ -3945,64 +3892,6 @@ mod tests { use crate::protocols::common::preprocessor::MultimodalData; use crate::protocols::common::{OutputOptions, SamplingOptions, StopConditions}; - const K3_MARKER: &str = "<|kimi_image_placeholder|>"; - const K3_PAD: &str = "<|media_pad|>"; - - fn seg(text: &str, allow_special: bool) -> crate::tokenizers::EncodeSegment { - crate::tokenizers::EncodeSegment { - text: text.to_string(), - allow_special, - } - } - - #[test] - fn swap_marker_segments_replaces_each_marker_with_one_pad() { - let mut prompt = RenderedPrompt::segmented(vec![ - seg("<|open|>message role=\"user\"<|sep|>", true), - seg(K3_MARKER, true), - seg("and", false), - seg(K3_MARKER, true), - seg("compare them", false), - ]); - - OpenAIPreprocessor::swap_marker_segments(K3_MARKER, K3_PAD, &mut prompt); - - let segments = prompt.segments().expect("still segmented"); - let pads = segments.iter().filter(|s| s.text == K3_PAD).count(); - // Cardinality is preserved 1:1 -- the engine expands each pad itself. - assert_eq!(pads, 2); - assert!(!segments.iter().any(|s| s.text == K3_MARKER)); - // The pad must stay a special segment or it tokenizes as literal text. - assert!(segments.iter().all(|s| s.text != K3_PAD || s.allow_special)); - // Surrounding structure and ordinary text are untouched. - assert_eq!(segments[2].text, "and"); - assert!(!segments[2].allow_special); - } - - #[test] - fn swap_marker_segments_is_a_noop_without_markers() { - let original = vec![seg("<|open|>message<|sep|>", true), seg("hello", false)]; - let mut prompt = RenderedPrompt::segmented(original.clone()); - - OpenAIPreprocessor::swap_marker_segments(K3_MARKER, K3_PAD, &mut prompt); - - let segments = prompt.segments().expect("still segmented"); - assert_eq!(segments.len(), original.len()); - assert_eq!(segments[0].text, original[0].text); - assert_eq!(segments[1].text, original[1].text); - } - - #[test] - fn swap_marker_segments_leaves_raw_text_prompts_alone() { - // The raw_prompt path has no segment boundaries, so a marker inside user - // text must never be rewritten into prompt structure. - let mut prompt = RenderedPrompt::text(format!("please describe {K3_MARKER}")); - - OpenAIPreprocessor::swap_marker_segments(K3_MARKER, K3_PAD, &mut prompt); - - assert_eq!(prompt.as_str(), format!("please describe {K3_MARKER}")); - } - #[test] fn prompt_invalid_request_maps_to_invalid_argument() { let error = PromptRenderError::invalid_request("unsupported model parameter").into();