Skip to content

refactor(preprocessor)!: let workers declare their per-image token #12281

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
furionw wants to merge 7 commits into from
Closed

refactor(preprocessor)!: let workers declare their per-image token #12281

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b0b81d6
refactor(preprocessor)!: let workers declare their per-image token
furionw Jul 28, 2026
9e17b77
fix(preprocessor): warn when a declared image token is not atomic
furionw Jul 28, 2026
2cbeddc
revert(preprocessor): drop the declared-image-token atomicity warning
furionw Jul 28, 2026
0a68adb
fix(sglang): keep the image-token probe quiet for models without images
furionw Jul 28, 2026
cd92ad2
fix(sglang): only read the image token for models that need it
furionw Jul 28, 2026
94fac66
refactor(vllm)!: convert the K3 pad worker-side instead of declaring it
furionw Jul 29, 2026
d369e47
perf(vllm): splice around K3 pads instead of walking every token
furionw Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions components/src/dynamo/sglang/register.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
102 changes: 101 additions & 1 deletion components/src/dynamo/vllm/multimodal_utils/request_processor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
furionw marked this conversation as resolved.

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
Comment thread
furionw marked this conversation as resolved.
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(
Expand Down Expand Up @@ -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:
Expand Down
79 changes: 79 additions & 0 deletions components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_request_processor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
10 changes: 0 additions & 10 deletions lib/bindings/python/rust/llm/local_model.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>) {
self.inner.kv_state_endpoint = kv_state_endpoint.as_deref().map(EndpointId::from);
Expand Down Expand Up @@ -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<String> {
self.inner
Expand Down
24 changes: 0 additions & 24 deletions lib/llm/src/local_model/runtime_config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(),
Expand Down
Loading
Loading