@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0377802c20

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Multimodal] Remove legacy multimodal fields in favor of MultiModalFeatureSpec (#24548)

Signed-off-by: sfeng33 <4florafeng@gmail.com>

Flora Feng committed Sep 12, 2025 at 06:42 UTC 0377802c20df2528d895add1f9918b2c7b868eb0

13 files changed +102 -116

tests/v1/core/test_encoder_cache_manager.py

+11 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

4 +from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange

5 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager

@@ -9,8 +10,17 @@ class MockRequest:

11 def __init__(self, request_id, mm_hashes, token_counts):

12 self.request_id = request_id

12 - self.mm_hashes = mm_hashes

13 self._token_counts = token_counts

14 + self.mm_features = []

15 + for i, mm_hash in enumerate(mm_hashes):

16 + feature = MultiModalFeatureSpec(

17 + data=None,

18 + modality="image",

19 + identifier=mm_hash,

20 + mm_position=PlaceholderRange(offset=0,

21 + length=self._token_counts[i]),

22 + )

23 + self.mm_features.append(feature)

25 def get_num_encoder_tokens(self, input_id: int) -> int:

26 return self._token_counts[input_id]

tests/v1/tpu/worker/test_tpu_model_runner.py

+1 -3

@@ -64,9 +64,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

64 NewRequestData(

65 req_id=req_id,

66 prompt_token_ids=[1, 2, 3],

67 - mm_kwargs=[],

68 - mm_hashes=[],

69 - mm_positions=[],

67 + mm_features=[],

68 sampling_params=SamplingParams(),

69 pooling_params=PoolingParams(),

70 block_ids=([0], ), # block_ids should be tuple[list[int]]

tests/v1/worker/test_gpu_input_batch.py

+1 -3

@@ -203,9 +203,7 @@ def _construct_cached_request_state(req_id_suffix: int):

203 prompt_token_ids=prompt_token_ids,

204 sampling_params=_create_sampling_params(),

205 pooling_params=None,

206 - mm_kwargs=[],

207 - mm_positions=[],

208 - mm_hashes=[],

206 + mm_features=[],

207 block_ids=([], ),

208 generator=None,

209 num_computed_tokens=len(output_token_ids),

tests/v1/worker/test_gpu_model_runner.py

+1 -3

@@ -118,9 +118,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

118 NewRequestData(

119 req_id=req_id,

120 prompt_token_ids=[1, 2, 3],

121 - mm_kwargs=[],

122 - mm_hashes=[],

123 - mm_positions=[],

121 + mm_features=[],

122 sampling_params=SamplingParams(),

123 pooling_params=None,

124 block_ids=([0], ),

vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py

+22 -19

@@ -300,11 +300,12 @@ class SharedStorageConnector(KVConnectorBase_V1):

300 total_need_load = 0

301 for new_req in scheduler_output.scheduled_new_reqs:

302 if new_req.req_id in self._requests_need_load:

303 - meta.add_request(token_ids=new_req.prompt_token_ids,

304 - block_ids=new_req.block_ids[0],

305 - block_size=self._block_size,

306 - is_store=False,

307 - mm_hashes=new_req.mm_hashes)

303 + meta.add_request(

304 + token_ids=new_req.prompt_token_ids,

305 + block_ids=new_req.block_ids[0],

306 + block_size=self._block_size,

307 + is_store=False,

308 + mm_hashes=[f.identifier for f in new_req.mm_features])

309 total_need_load += 1

310 else:

311 # NOTE: here, we set the store and load being exclusive,

@@ -312,11 +313,12 @@ class SharedStorageConnector(KVConnectorBase_V1):

313 # NOTE(rob): for this debug implementation, we only cache

314 # the original prompt tokens.

315 if not self._found_match_for_request(new_req):

315 - meta.add_request(token_ids=new_req.prompt_token_ids,

316 - block_ids=new_req.block_ids[0],

317 - block_size=self._block_size,

318 - is_store=True,

319 - mm_hashes=new_req.mm_hashes)

316 + meta.add_request(

317 + token_ids=new_req.prompt_token_ids,

318 + block_ids=new_req.block_ids[0],

319 + block_size=self._block_size,

320 + is_store=True,

321 + mm_hashes=[f.identifier for f in new_req.mm_features])

322

323 cached_reqs = scheduler_output.scheduled_cached_reqs

324 for i, req_id in enumerate(cached_reqs.req_ids):

@@ -341,11 +343,12 @@ class SharedStorageConnector(KVConnectorBase_V1):

343 # of the block_ids for the request.

344 block_ids = new_block_ids[0]

345

344 - meta.add_request(token_ids=token_ids,

345 - block_ids=block_ids,

346 - block_size=self._block_size,

347 - is_store=False,

348 - mm_hashes=request.mm_hashes)

346 + meta.add_request(

347 + token_ids=token_ids,

348 + block_ids=block_ids,

349 + block_size=self._block_size,

350 + is_store=False,

351 + mm_hashes=[f.identifier for f in request.mm_features])

352 total_need_load += 1

353

354 assert total_need_load == len(self._requests_need_load)

@@ -364,10 +367,10 @@ class SharedStorageConnector(KVConnectorBase_V1):

367 """

368 num_tokens_to_check = align_to_block_size(

369 len(request.prompt_token_ids) - 1, self._block_size)

367 - foldername = self._generate_foldername_debug(torch.tensor(

368 - request.prompt_token_ids)[:num_tokens_to_check],

369 - request.mm_hashes,

370 - create_folder=False)

370 + foldername = self._generate_foldername_debug(

371 + torch.tensor(request.prompt_token_ids)[:num_tokens_to_check],

372 + [f.identifier for f in request.mm_features],

373 + create_folder=False)

374 return os.path.exists(foldername)

375

376 def _generate_foldername_debug(

vllm/v1/core/encoder_cache_manager.py

+5 -5

@@ -86,7 +86,7 @@ class EncoderCacheManager:

86 Returns:

87 True if the encoder output for this input is already cached

88 """

89 - mm_hash = request.mm_hashes[input_id]

89 + mm_hash = request.mm_features[input_id].identifier

90 # Not cached at all

91 if mm_hash not in self.cached:

92 return False

@@ -167,7 +167,7 @@ class EncoderCacheManager:

167 This method assumes can_allocate() returned True for the same input.

168 """

169

170 - mm_hash = request.mm_hashes[input_id]

170 + mm_hash = request.mm_features[input_id].identifier

171 request_id = request.request_id

172 if mm_hash not in self.cached:

173 self.cached[mm_hash] = set()

@@ -193,8 +193,8 @@ class EncoderCacheManager:

193 """

194 return {

195 input_id

196 - for input_id in range(len(request.mm_hashes))

197 - if request.mm_hashes[input_id] in self.cached

196 + for input_id in range(len(request.mm_features))

197 + if request.mm_features[input_id].identifier in self.cached

198 }

199

200 def free_encoder_input(self, request: Request, input_id: int) -> None:

@@ -208,7 +208,7 @@ class EncoderCacheManager:

208 `can_allocate`).

209 """

210 req_id = request.request_id

211 - mm_hash = request.mm_hashes[input_id]

211 + mm_hash = request.mm_features[input_id].identifier

212 # The mm_hash not in cache or the req_id set is empty

213 if not self.cached.get(mm_hash, None):

214 return

vllm/v1/core/kv_cache_utils.py

+16 -20

@@ -418,9 +418,9 @@ def need_extra_keys(request: Request) -> bool:

418 # Multimodal requests need to include the MM hash.

419 # LoRA requests need to include the LoRA ID.

420 # Request with provided cache salt need to include the salt.

421 - return bool(request.mm_hashes) or (request.lora_request

422 - is not None) or (request.cache_salt

423 - is not None)

421 + return bool(request.mm_features) or (request.lora_request

422 + is not None) or (request.cache_salt

423 + is not None)

424

425

426 def _gen_mm_extra_hash_keys(request: Request, start_token_idx: int,

@@ -442,32 +442,28 @@ def _gen_mm_extra_hash_keys(request: Request, start_token_idx: int,

442 """

443 extra_keys: list[Any] = []

444

445 - mm_positions, mm_hashes = request.mm_positions, request.mm_hashes

446 - if not mm_positions:

445 + mm_features = request.mm_features

446 + if not mm_features:

447 return extra_keys, start_mm_idx

448

449 - if mm_positions and len(mm_positions) != len(mm_hashes):

450 - raise ValueError(

451 - "The number of multi-modal positions and hashes must match. This "

452 - "is likely because you did not enable MM hashing. "

453 - "Please set `mm_processor_cache_gb > 0`.")

454 -

455 - # Note that we assume mm_positions is sorted by offset.

449 + # Note that we assume mm_features are sorted by mm_position.offset.

450 # We do not need to check all mm inputs if the start token index is out of

451 # range. This usually happens in the late prefill phase and decoding phase.

458 - if mm_positions[-1].offset + mm_positions[-1].length < start_token_idx:

452 + last_pos = mm_features[-1].mm_position

453 + if last_pos.offset + last_pos.length < start_token_idx:

454 return extra_keys, start_mm_idx

455

456 # Support start_mm_idx == -1 to indicate the last mm input.

457 if start_mm_idx < 0:

463 - assert -start_mm_idx <= len(mm_positions)

464 - start_mm_idx = len(mm_positions) + start_mm_idx

458 + assert -start_mm_idx <= len(mm_features)

459 + start_mm_idx = len(mm_features) + start_mm_idx

460

461 curr_mm_idx = start_mm_idx

467 - while mm_positions and curr_mm_idx < len(mm_positions):

468 - assert mm_hashes[curr_mm_idx] is not None

469 - offset = mm_positions[curr_mm_idx].offset

470 - length = mm_positions[curr_mm_idx].length

462 + while mm_features and curr_mm_idx < len(mm_features):

463 + mm_feature = mm_features[curr_mm_idx]

464 + assert mm_feature.identifier is not None

465 + offset = mm_feature.mm_position.offset

466 + length = mm_feature.mm_position.length

467 if end_token_idx > offset:

468 if start_token_idx > offset + length:

469 # This block has passed the current mm input.

@@ -475,7 +471,7 @@ def _gen_mm_extra_hash_keys(request: Request, start_token_idx: int,

471 continue

472

473 # The block contains the current mm input.

478 - extra_keys.append(mm_hashes[curr_mm_idx])

474 + extra_keys.append(mm_feature.identifier)

475

476 if end_token_idx >= offset + length:

477 # If this block contains the end of the current mm input,

vllm/v1/core/sched/output.py

+5 -13

@@ -15,7 +15,7 @@ if TYPE_CHECKING:

15 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

16 KVConnectorMetadata)

17 from vllm.lora.request import LoRARequest

18 - from vllm.multimodal.inputs import MultiModalKwargsItem, PlaceholderRange

18 + from vllm.multimodal.inputs import MultiModalFeatureSpec

19 from vllm.pooling_params import PoolingParams

20 from vllm.sampling_params import SamplingParams

21 from vllm.v1.request import Request

@@ -27,9 +27,7 @@ class NewRequestData:

28 req_id: str

29 prompt_token_ids: list[int]

30 - mm_kwargs: list[MultiModalKwargsItem]

31 - mm_hashes: list[str]

32 - mm_positions: list[PlaceholderRange]

30 + mm_features: list[MultiModalFeatureSpec]

31 sampling_params: Optional[SamplingParams]

32 pooling_params: Optional[PoolingParams]

33 block_ids: tuple[list[int], ...]

@@ -45,9 +43,7 @@ class NewRequestData:

43 return cls(

44 req_id=request.request_id,

45 prompt_token_ids=request.prompt_token_ids,

48 - mm_kwargs=request.mm_kwargs,

49 - mm_hashes=request.mm_hashes,

50 - mm_positions=request.mm_positions,

46 + mm_features=request.mm_features,

47 sampling_params=request.sampling_params,

48 pooling_params=request.pooling_params,

49 block_ids=block_ids,

@@ -59,9 +55,7 @@ class NewRequestData:

55 return (f"NewRequestData("

56 f"req_id={self.req_id},"

57 f"prompt_token_ids={self.prompt_token_ids},"

62 - f"mm_kwargs={self.mm_kwargs},"

63 - f"mm_hashes={self.mm_hashes},"

64 - f"mm_positions={self.mm_positions},"

58 + f"mm_features={self.mm_features},"

59 f"sampling_params={self.sampling_params},"

60 f"block_ids={self.block_ids},"

61 f"num_computed_tokens={self.num_computed_tokens},"

@@ -73,9 +67,7 @@ class NewRequestData:

67 return (f"NewRequestData("

68 f"req_id={self.req_id},"

69 f"prompt_token_ids_len={len(self.prompt_token_ids)},"

76 - f"mm_kwargs={self.mm_kwargs},"

77 - f"mm_hashes={self.mm_hashes},"

78 - f"mm_positions={self.mm_positions},"

70 + f"mm_features={self.mm_features},"

71 f"sampling_params={self.sampling_params},"

72 f"block_ids={self.block_ids},"

73 f"num_computed_tokens={self.num_computed_tokens},"

vllm/v1/core/sched/scheduler.py

+11 -11

@@ -736,18 +736,18 @@ class Scheduler(SchedulerInterface):

736 if num_new_tokens == 0 or not request.has_encoder_inputs:

737 return [], num_new_tokens, encoder_compute_budget

738 encoder_inputs_to_schedule: list[int] = []

739 - mm_positions = request.mm_positions

740 - assert mm_positions is not None

741 - assert len(mm_positions) > 0

739 + mm_features = request.mm_features

740 + assert mm_features is not None

741 + assert len(mm_features) > 0

742

743 # NOTE: since scheduler operates on the request level (possibly with

744 # multiple encoder inputs per request), we need to create temporary

745 # trackers for accounting at the encoder input level.

746 mm_hashes_to_schedule = set()

747 num_tokens_to_schedule = 0

748 - for i, pos_info in enumerate(mm_positions):

749 - start_pos = pos_info.offset

750 - num_encoder_tokens = pos_info.length

748 + for i, mm_feature in enumerate(mm_features):

749 + start_pos = mm_feature.mm_position.offset

750 + num_encoder_tokens = mm_feature.mm_position.length

751

752 # The encoder output is needed if the two ranges overlap:

753 # [num_computed_tokens, num_computed_tokens + num_new_tokens) and

@@ -778,7 +778,7 @@ class Scheduler(SchedulerInterface):

778 if not self.is_encoder_decoder:

779 # We are not using the encoder cache for encoder-decoder models,

780 # yet.

781 - if request.mm_hashes[i] in mm_hashes_to_schedule:

781 + if request.mm_features[i].identifier in mm_hashes_to_schedule:

782 # The same encoder input has already been scheduled in the

783 # current step.

784 continue

@@ -820,7 +820,7 @@ class Scheduler(SchedulerInterface):

820

821 num_tokens_to_schedule += num_encoder_tokens

822 encoder_compute_budget -= num_encoder_tokens

823 - mm_hashes_to_schedule.add(request.mm_hashes[i])

823 + mm_hashes_to_schedule.add(request.mm_features[i].identifier)

824 encoder_inputs_to_schedule.append(i)

825

826 return (

@@ -1048,9 +1048,9 @@ class Scheduler(SchedulerInterface):

1048 # Here, we use list(set) to avoid modifying the set while iterating

1049 # over it.

1050 for input_id in list(cached_encoder_input_ids):

1051 - mm_positions = request.mm_positions[input_id]

1052 - start_pos = mm_positions.offset

1053 - num_tokens = mm_positions.length

1051 + mm_feature = request.mm_features[input_id]

1052 + start_pos = mm_feature.mm_position.offset

1053 + num_tokens = mm_feature.mm_position.length

1054 if self.is_encoder_decoder and request.num_computed_tokens > 0:

1055 # With Whisper, as soon as we've generated a single token,

1056 # we know we're done with the encoder input. Cross Attention

vllm/v1/request.py

+2 -7

@@ -91,11 +91,6 @@ class Request:

91 self.mm_features = mm_features or []

92 self.num_encoder_inputs = len(self.mm_features)

93 self.has_encoder_inputs = self.num_encoder_inputs > 0

94 - # TODO(sfeng33): Remove these legacy fields after clearing out all

95 - # references in scheduler and model runner

96 - self.mm_positions = [f.mm_position for f in self.mm_features]

97 - self.mm_kwargs = [f.data for f in self.mm_features]

98 - self.mm_hashes = [f.identifier for f in self.mm_features]

95 # Read-only views

96 # Prevent directly appending to these lists since

@@ -180,8 +175,8 @@ class Request:

175 return RequestStatus.get_finished_reason(self.status)

176

177 def get_num_encoder_tokens(self, input_id: int) -> int:

183 - assert input_id < len(self.mm_positions)

184 - num_tokens = self.mm_positions[input_id].length

178 + assert input_id < len(self.mm_features)

179 + num_tokens = self.mm_features[input_id].mm_position.length

180 return num_tokens

181

182 def record_event(

vllm/v1/worker/gpu_input_batch.py

+4 -6

@@ -10,8 +10,7 @@ import torch

10 from typing_extensions import deprecated

12 from vllm.lora.request import LoRARequest

13 -from vllm.multimodal.inputs import (MultiModalKwargsItem,

14 - MultiModalKwargsItems, PlaceholderRange)

13 +from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems

14 from vllm.pooling_params import PoolingParams

15 from vllm.sampling_params import SamplingParams, SamplingType

16 from vllm.utils import swap_dict_values

@@ -31,9 +30,7 @@ class CachedRequestState:

31 req_id: str

32 prompt_token_ids: list[int]

34 - mm_kwargs: list[MultiModalKwargsItem]

35 - mm_positions: list[PlaceholderRange]

36 - mm_hashes: list[str]

33 + mm_features: list[MultiModalFeatureSpec]

34 sampling_params: Optional[SamplingParams]

35 pooling_params: Optional[PoolingParams]

36 generator: Optional[torch.Generator]

@@ -60,7 +57,8 @@ class CachedRequestState:

57 "removed in v0.13. Please use `mm_kwargs` instead.")

58 def mm_inputs(self) -> list[MultiModalKwargsItems]:

59 return [

63 - MultiModalKwargsItems.from_seq([item]) for item in self.mm_kwargs

60 + MultiModalKwargsItems.from_seq([f.data]) for f in self.mm_features

61 + if f.data is not None

62 ]

64 def get_token_id(self, idx: int) -> int:

vllm/v1/worker/gpu_model_runner.py

+15 -13

@@ -555,9 +555,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

555 req_state = CachedRequestState(

556 req_id=req_id,

557 prompt_token_ids=new_req_data.prompt_token_ids,

558 - mm_kwargs=new_req_data.mm_kwargs,

559 - mm_positions=new_req_data.mm_positions,

560 - mm_hashes=new_req_data.mm_hashes,

558 + mm_features=new_req_data.mm_features,

559 sampling_params=sampling_params,

560 pooling_params=pooling_params,

561 generator=generator,

@@ -698,7 +696,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

696 second_per_grid_ts = []

697 audio_feature_lengths = []

698 use_audio_in_video = False

701 - for mm_item in req_state.mm_kwargs:

699 + for mm_feature in req_state.mm_features:

700 + mm_item = mm_feature.data

701 + if mm_item is None:

702 + continue

703 mm_input = mm_item.get_data()

704 if (t := mm_input.get("image_grid_thw")) is not None:

705 image_grid_thw.append(t.tolist())

@@ -731,7 +732,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

732

733 mm_kwargs = list[MultiModalKwargsItem]()

734 for req in scheduler_output.scheduled_new_reqs:

734 - mm_kwargs.extend(req.mm_kwargs)

735 + for feature in req.mm_features:

736 + if feature.data is not None:

737 + mm_kwargs.append(feature.data)

738

739 # Input all modalities at once

740 mm_kwargs_combined: BatchedTensorInputs = {}

@@ -1361,10 +1364,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1364 req_state = self.requests[req_id]

1365

1366 for mm_input_id in encoder_input_ids:

1364 - mm_hash = req_state.mm_hashes[mm_input_id]

1365 - mm_kwargs.append(req_state.mm_kwargs[mm_input_id])

1366 - mm_hashes_pos.append(

1367 - (mm_hash, req_state.mm_positions[mm_input_id]))

1367 + mm_feature = req_state.mm_features[mm_input_id]

1368 + mm_hash = mm_feature.identifier

1369 + mm_kwargs.append(mm_feature.data)

1370 + mm_hashes_pos.append((mm_hash, mm_feature.mm_position))

1371

1372 return mm_kwargs, mm_hashes_pos

1373

@@ -1426,9 +1429,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1429 req_state = self.requests[req_id]

1430 num_computed_tokens = \

1431 req_state.num_computed_tokens + shift_computed_tokens

1429 - mm_positions = req_state.mm_positions

1430 - mm_hashes = req_state.mm_hashes

1431 - for i, pos_info in enumerate(mm_positions):

1432 + for mm_feature in req_state.mm_features:

1433 + pos_info = mm_feature.mm_position

1434 start_pos = pos_info.offset

1435 num_encoder_tokens = pos_info.length

1436

@@ -1451,7 +1453,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1453 )

1454 assert start_idx < end_idx

1455

1454 - mm_hash = mm_hashes[i]

1456 + mm_hash = mm_feature.identifier

1457 encoder_output = self.encoder_cache.get(mm_hash, None)

1458 assert encoder_output is not None,\

1459 f"Encoder cache miss for {mm_hash}."

vllm/v1/worker/tpu_model_runner.py

+8 -12

@@ -387,9 +387,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

387 self.requests[req_id] = CachedRequestState(

388 req_id=req_id,

389 prompt_token_ids=new_req_data.prompt_token_ids,

390 - mm_kwargs=new_req_data.mm_kwargs,

391 - mm_positions=new_req_data.mm_positions,

392 - mm_hashes=new_req_data.mm_hashes,

390 + mm_features=new_req_data.mm_features,

391 sampling_params=sampling_params,

392 pooling_params=None,

393 generator=None,

@@ -822,10 +820,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

820 req_state = self.requests[req_id]

821

822 for mm_input_id in encoder_input_ids:

825 - mm_hash = req_state.mm_hashes[mm_input_id]

826 - mm_kwargs.append(req_state.mm_kwargs[mm_input_id])

827 - mm_hashes_pos.append(

828 - (mm_hash, req_state.mm_positions[mm_input_id]))

823 + mm_feature = req_state.mm_features[mm_input_id]

824 + mm_hash = mm_feature.identifier

825 + mm_kwargs.append(mm_feature.data)

826 + mm_hashes_pos.append((mm_hash, mm_feature.mm_position))

827

828 # Batch mm inputs as much as we can: if a request in the batch has

829 # multiple modalities or a different modality than the previous one,

@@ -883,13 +881,12 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

881 req_id]

882 req_state = self.requests[req_id]

883 num_computed_tokens = req_state.num_computed_tokens

886 - mm_positions = req_state.mm_positions

887 - mm_hashes = req_state.mm_hashes

884 # TODO unroll loop and assume/enforce --disable_chunked_mm_input

885 # NOTE (NickLucche) here we diverge from logic in other runners, as

886 # we assume to only have whole mm items to process. Hence we avoid

887 # the intrinsic dynamism that `gather_mm_placeholders` introduces.

892 - for i, pos_info in enumerate(mm_positions):

888 + for mm_feature in req_state.mm_features:

889 + pos_info = mm_feature.mm_position

890 start_pos = pos_info.offset

891 num_encoder_tokens = pos_info.length

892

@@ -904,8 +901,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

901 # The encoder output is already processed and stored

902 # in the decoder's KV cache.

903 continue

907 -

908 - mm_hash = mm_hashes[i]

904 + mm_hash = mm_feature.identifier

905 encoder_output = self.encoder_cache.get(mm_hash, None)

906 assert encoder_output is not None,\

907 f"Encoder cache miss for {mm_hash}."