@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3fb7bb46d6

[Perf] Avoid repeated multimodal prompt update scans (#51774)

Signed-off-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>

Tianyu Guo committed Aug 11, 2026 at 17:23 UTC 3fb7bb46d6202e23988c96746891ec788ee6aaf5

2 files changed +226 -121

tests/multimodal/test_processing.py

+33

@@ -1099,6 +1099,39 @@ def test_apply_matches_no_match_exits_quickly():

1099 assert "".join(result) == long_prompt

1100

1101

1102 +def test_apply_matches_many_shared_targets_scales_linearly():

1103 + """Shared replacement targets must not trigger per-item rescanning."""

1104 + replacement = [1] * 50

1105 + update = PromptReplacement("image", [0], replacement)

1106 +

1107 + def measure(item_count: int) -> float:

1108 + mm_prompt_updates = {

1109 + "image": [[update.resolve(item_idx)] for item_idx in range(item_count)]

1110 + }

1111 + prompt = [0] * item_count

1112 +

1113 + start = time.perf_counter()

1114 + result, match_result = apply_token_matches(

1115 + prompt,

1116 + mm_prompt_updates,

1117 + tokenizer=None,

1118 + )

1119 + elapsed = time.perf_counter() - start

1120 +

1121 + assert len(result) == item_count * len(replacement)

1122 + assert all(token_id == 1 for token_id in result)

1123 + assert match_result == {"image": [0] * item_count}

1124 +

1125 + return elapsed

1126 +

1127 + measure(100)

1128 + small_time = measure(1_000)

1129 + large_time = measure(4_000)

1130 +

1131 + time_ratio = large_time / small_time

1132 + assert time_ratio < 8, f"Expected linear scaling, got {time_ratio:.1f}x"

1133 +

1134 +

1135 def test_iter_token_matches_rejects_negative_start_idx():

1136 with pytest.raises(ValueError, match="non-negative"):

1137 list(iter_token_matches([1, 2, 3], [2], start_idx=-1))

vllm/multimodal/processing/processor.py

+193 -121

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from abc import ABC, abstractmethod

4 -from collections import defaultdict

4 +from collections import defaultdict, deque

5 from collections.abc import Callable, Generator, ItemsView, Iterable, Mapping, Sequence

6 from dataclasses import dataclass, field, replace

7 from enum import Enum

@@ -706,91 +706,191 @@ class PlaceholderFeaturesInfo:

706 )

707

708

709 -_MatchToApply = tuple[tuple[str, int], tuple[PromptTargetMatch, int]]

709 +class _MatchedUpdate(NamedTuple):

710 + """A resolved update selected for a match in the original prompt."""

711

712 + priority: int

713 + """The original item order used to preserve match tie-breaking."""

714

712 -def _find_matches(

713 - prompt: _S,

715 + update: ResolvedPromptUpdate

716 + """The selected update for the multimodal item."""

717 +

718 + update_idx: int

719 + """The selected update's index within the item's alternatives."""

720 +

721 + match: PromptTargetMatch

722 + """The target range in the original prompt."""

723 +

724 +

725 +_UpdateQueue: TypeAlias = deque[tuple[int, Sequence[ResolvedPromptUpdate]]]

726 +"""

727 +Items with the same ordered match rules, stored as `(priority, alternatives)`.

728 +"""

729 +

730 +_QueueMatch: TypeAlias = tuple[_UpdateQueue, PromptTargetMatch, int]

731 +"""A queue together with its next match and selected alternative index."""

732 +

733 +

734 +def _target_key(target: UpdateTarget) -> tuple[str, object]:

735 + """Return a hashable key that preserves target matching semantics."""

736 + if isinstance(target, PromptIndex):

737 + return ("index", id(target))

738 + if isinstance(target, str):

739 + return ("text", target)

740 +

741 + return ("tokens", tuple(target))

742 +

743 +

744 +def _compile_prompt_update_queues(

745 mm_prompt_updates: "MultiModalPromptUpdates",

746 +) -> dict[

747 + tuple[tuple[UpdateMode, tuple[str, object]], ...],

748 + _UpdateQueue,

749 +]:

750 + """Group items with identical match rules into ordered queues."""

751 + queues_by_signature = dict[

752 + tuple[tuple[UpdateMode, tuple[str, object]], ...],

753 + _UpdateQueue,

754 + ]()

755 + priority = 0

756 +

757 + for modality_updates in mm_prompt_updates.values():

758 + for updates in modality_updates:

759 + signature = tuple(

760 + (update.mode, _target_key(update.target)) for update in updates

761 + )

762 + queues_by_signature.setdefault(signature, deque()).append(

763 + (priority, updates)

764 + )

765 + priority += 1

766 +

767 + return queues_by_signature

768 +

769 +

770 +def _find_queue_match(

771 + queue: _UpdateQueue,

772 + prompt: _S,

773 tokenizer: TokenizerLike | None,

774 *,

717 - prev_end_idx: int = 0,

718 - current_result: "MultiModalPromptUpdatesApplyResult",

719 -) -> tuple[UpdateMode | None, list[_MatchToApply]]:

720 - mode: UpdateMode | None = None

721 - mm_matches = dict[tuple[str, int], tuple[PromptTargetMatch, int]]()

722 -

723 - # Items commonly share the same target object (a static PromptUpdate

724 - # target resolves to the same object for every item); scan the prompt once

725 - # per distinct target instead of once per item. The cache is keyed by

726 - # object identity: it implies value equality, costs O(1) per lookup

727 - # (hashing a long target value per lookup would defeat the point), and the

728 - # target objects are kept alive by `mm_prompt_updates` for the duration of

729 - # this call. Items with equal-valued but distinct target objects simply

730 - # miss the cache and scan independently.

731 - # NOTE: ResolvedPromptUpdate.iter_matches depends only on the target,

732 - # prompt, tokenizer and start_idx, so updates sharing a target share the

733 - # same first match within this invocation. Revisit this cache if matching

734 - # ever depends on other update attributes.

735 - first_match_by_target_id = dict[int, PromptTargetMatch | None]()

736 -

737 - def get_first_match(update: "ResolvedPromptUpdate") -> PromptTargetMatch | None:

738 - target = update.target

739 - if isinstance(target, PromptIndex): # Matches are not cacheable

740 - return next(

741 - update.iter_matches(prompt, tokenizer, start_idx=prev_end_idx),

742 - None,

743 - )

775 + start_idx: int,

776 + mode: UpdateMode | None = None,

777 +) -> tuple[PromptTargetMatch, int] | None:

778 + """Find the first matching alternative for the next queued item."""

779 + _, updates = queue[0]

780 + for update_idx, update in enumerate(updates):

781 + if mode is not None and update.mode != mode:

782 + continue

783 +

784 + match = next(

785 + update.iter_matches(prompt, tokenizer, start_idx=start_idx),

786 + None,

787 + )

788 + if match is not None:

789 + return match, update_idx

790

745 - key = id(target)

746 - if key not in first_match_by_target_id:

747 - first_match_by_target_id[key] = next(

748 - update.iter_matches(prompt, tokenizer, start_idx=prev_end_idx),

749 - None,

750 - )

791 + return None

792

752 - return first_match_by_target_id[key]

793

754 - for modality, modality_updates in mm_prompt_updates.items():

755 - for item_idx, item_updates in enumerate(modality_updates):

756 - if current_result[modality][item_idx] is not None:

757 - continue # Updates have already been applied for this item

794 +def _next_priority(queue: _UpdateQueue) -> int:

795 + """Return the original priority of the next queued item."""

796 + priority, _ = queue[0]

797 + return priority

798

759 - for update_idx, update in enumerate(item_updates):

760 - if (modality, item_idx) in mm_matches:

761 - break # Already found a match for this item

799

763 - match = get_first_match(update)

764 - if match is None:

765 - continue

800 +def _plan_prompt_updates(

801 + prompt: _S,

802 + mm_prompt_updates: "MultiModalPromptUpdates",

803 + tokenizer: TokenizerLike | None,

804 +) -> tuple[list[_MatchedUpdate], "MultiModalPromptUpdatesApplyResult"]:

805 + """Plan non-overlapping prompt updates before rendering the output."""

806 + queues = list(_compile_prompt_update_queues(mm_prompt_updates).values())

807 + result: MultiModalPromptUpdatesApplyResult = {

808 + modality: [None] * len(items) for modality, items in mm_prompt_updates.items()

809 + }

810 + planned_updates = list[_MatchedUpdate]()

811 + prev_end_idx = 0

812

767 - # All matches should share the same mode

768 - if mode is None:

769 - mode = update.mode

770 - elif mode != update.mode:

771 - continue

813 + while queues:

814 + first_matches = list[_QueueMatch]()

815 + for queue in queues:

816 + queue_match = _find_queue_match(

817 + queue,

818 + prompt,

819 + tokenizer,

820 + start_idx=prev_end_idx,

821 + )

822 + if queue_match is not None:

823 + prompt_match, update_idx = queue_match

824 + first_matches.append((queue, prompt_match, update_idx))

825

773 - mm_matches[(modality, item_idx)] = match, update_idx

826 + if not first_matches:

827 + break

828

775 - # Prioritize earlier matches

776 - matches_to_apply = sorted(mm_matches.items(), key=lambda item: item[1][0])

829 + mode_queue, _, mode_update_idx = min(

830 + first_matches,

831 + key=lambda item: _next_priority(item[0]),

832 + )

833 + _, mode_updates = mode_queue[0]

834 + mode = mode_updates[mode_update_idx].mode

835

778 - # To avoid conflicts, only replace one non-empty item at a time

779 - if mode == UpdateMode.REPLACE:

780 - matches_to_apply_ = list[_MatchToApply]()

781 - has_non_empty_matches = False

836 + mode_matches = list[_QueueMatch]()

837 + for queue, prompt_match, first_update_idx in first_matches:

838 + if queue[0][1][first_update_idx].mode == mode:

839 + mode_matches.append((queue, prompt_match, first_update_idx))

840 + continue

841 +

842 + queue_match = _find_queue_match(

843 + queue,

844 + prompt,

845 + tokenizer,

846 + start_idx=prev_end_idx,

847 + mode=mode,

848 + )

849 + if queue_match is not None:

850 + prompt_match, update_idx = queue_match

851 + mode_matches.append((queue, prompt_match, update_idx))

852 +

853 + updates_to_apply = list[_MatchedUpdate]()

854 + non_empty_replacements = list[_QueueMatch]()

855 + for queue, match, update_idx in mode_matches:

856 + if mode == UpdateMode.REPLACE and match.start_idx != match.end_idx:

857 + non_empty_replacements.append((queue, match, update_idx))

858 + else:

859 + while queue:

860 + priority, updates = queue.popleft()

861 + updates_to_apply.append(

862 + _MatchedUpdate(

863 + priority=priority,

864 + update=updates[update_idx],

865 + update_idx=update_idx,

866 + match=match,

867 + )

868 + )

869

783 - for item in matches_to_apply:

784 - _, (match, _) = item

785 - if match.start_idx == match.end_idx:

786 - matches_to_apply_.append(item)

787 - elif not has_non_empty_matches:

788 - has_non_empty_matches = True

789 - matches_to_apply_.append(item)

870 + if non_empty_replacements:

871 + queue, match, update_idx = min(

872 + non_empty_replacements,

873 + key=lambda item: (item[1], _next_priority(item[0])),

874 + )

875 + priority, updates = queue.popleft()

876 + updates_to_apply.append(

877 + _MatchedUpdate(

878 + priority=priority,

879 + update=updates[update_idx],

880 + update_idx=update_idx,

881 + match=match,

882 + )

883 + )

884

791 - matches_to_apply = matches_to_apply_

885 + updates_to_apply.sort(key=lambda item: (item.match, item.priority))

886 + for matched_update in updates_to_apply:

887 + update = matched_update.update

888 + result[update.modality][update.item_idx] = matched_update.update_idx

889 + prev_end_idx = matched_update.match.end_idx

890 + planned_updates.extend(updates_to_apply)

891 + queues = [queue for queue in queues if queue]

892

793 - return mode, matches_to_apply

893 + return planned_updates, result

894

895

896 def _all_items_found(

@@ -808,65 +908,37 @@ def _apply_matches(

908 mm_prompt_updates: "MultiModalPromptUpdates",

909 tokenizer: TokenizerLike | None,

910 ) -> tuple[list[_S], "MultiModalPromptUpdatesApplyResult"]:

811 - mm_item_counts = {m: len(items) for m, items in mm_prompt_updates.items()}

812 -

911 out_seqs = list[str | list[int]]()

814 - out_result: MultiModalPromptUpdatesApplyResult = {

815 - m: [None] * len(items) for m, items in mm_prompt_updates.items()

816 - }

817 -

818 - # Early exit if no items to find

819 - mm_found_counts = {

820 - m: sum(r is not None for r in res) for m, res in out_result.items()

821 - }

822 - if _all_items_found(mm_item_counts, mm_found_counts):

823 - return [prompt], out_result

912 + matched_updates, result = _plan_prompt_updates(

913 + prompt,

914 + mm_prompt_updates,

915 + tokenizer,

916 + )

917

918 prev_end_idx = 0

826 - while True:

827 - mode, matches_to_apply = _find_matches(

828 - prompt,

829 - mm_prompt_updates,

830 - tokenizer,

831 - prev_end_idx=prev_end_idx,

832 - current_result=out_result,

833 - )

834 -

835 - if mode is None:

836 - break # No more matches to find

837 -

838 - for (modality, item_idx), (match, update_idx) in matches_to_apply:

839 - matched_update = mm_prompt_updates[modality][item_idx][update_idx]

840 - matched_content = matched_update.content.full

841 -

842 - if mode == UpdateMode.INSERT:

843 - end_idx_to_insert = match.end_idx

844 - elif mode == UpdateMode.REPLACE:

845 - end_idx_to_insert = match.start_idx

846 - else:

847 - assert_never(mode)

848 -

849 - out_seqs.append(prompt[prev_end_idx:end_idx_to_insert])

850 - out_seqs.append(

851 - _seq2text(tokenizer, matched_content)

852 - if isinstance(prompt, str)

853 - else _seq2tokens(tokenizer, matched_content)

854 - )

855 - out_result[modality][item_idx] = update_idx

856 -

857 - # Exclude overlapping matches

858 - prev_end_idx = match.end_idx

919 + for matched_update in matched_updates:

920 + update = matched_update.update

921 + match = matched_update.match

922 + matched_content = update.content.full

923 +

924 + if update.mode == UpdateMode.INSERT:

925 + end_idx_to_insert = match.end_idx

926 + elif update.mode == UpdateMode.REPLACE:

927 + end_idx_to_insert = match.start_idx

928 + else:

929 + assert_never(update.mode)

930

860 - # Early exit if all items found

861 - mm_found_counts = {

862 - m: sum(r is not None for r in res) for m, res in out_result.items()

863 - }

864 - if _all_items_found(mm_item_counts, mm_found_counts):

865 - break

931 + out_seqs.append(prompt[prev_end_idx:end_idx_to_insert])

932 + out_seqs.append(

933 + _seq2text(tokenizer, matched_content)

934 + if isinstance(prompt, str)

935 + else _seq2tokens(tokenizer, matched_content)

936 + )

937 + prev_end_idx = match.end_idx

938

939 out_seqs.append(prompt[prev_end_idx:])

940

869 - return cast(list[_S], out_seqs), out_result

941 + return cast(list[_S], out_seqs), result

942

943

944 def apply_token_matches(