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

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 10, 22:45 UTC

Feature/isaac 0.1 (#28367)

Signed-off-by: oscardev256 <42308241+oscardev256@users.noreply.github.com> Signed-off-by: Oscar Gonzalez <ogonzal6@alumni.jh.edu> Signed-off-by: Yang <lymailforjob@gmail.com> Co-authored-by: Yang <lymailforjob@gmail.com>

oscardev256 committed Dec 25, 2025 at 21:49 UTC b7165d53c689ed11f3f9a984a4b7f45accbc211f

11 files changed +1800 -2

docs/models/supported_models.md

@@ -688,6 +688,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

688 | `H2OVLChatModel` | H2OVL | T + IE+ | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ |

689 | `HunYuanVLForConditionalGeneration` | HunyuanOCR | T + IE+ | `tencent/HunyuanOCR`, etc. | ✅︎ | ✅︎ |

690 | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | |

691 +| `IsaacForConditionalGeneration` | Isaac | T + I+ | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ |

692 | `InternS1ForConditionalGeneration` | Intern-S1 | T + IE+ + VE+ | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |

693 | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + IE+ + (VE+) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |

694 | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + IE+ + VE+ | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |

requirements/test.in

@@ -56,3 +56,5 @@ pydantic>=2.12 # 2.11 leads to error on python 3.13

56 decord==0.6.0

57 terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test

58 gpt-oss >= 0.0.7; python_version > '3.11'

59 +

60 +perceptron # required for isaac test

requirements/test.txt

+19 -2

@@ -135,6 +135,7 @@ cloudpickle==3.1.1

135 # via mlflow-skinny

136 colorama==0.4.6

137 # via

138 + # perceptron

139 # sacrebleu

140 # schemathesis

141 # tqdm-multiprocess

@@ -302,6 +303,8 @@ h11==0.14.0

303 # via

304 # httpcore

305 # uvicorn

306 +h2==4.3.0

307 + # via httpx

308 h5py==3.13.0

309 # via terratorch

310 harfile==0.3.0

@@ -310,6 +313,8 @@ hf-xet==1.1.7

313 # via huggingface-hub

314 hiredis==3.0.0

315 # via tensorizer

316 +hpack==4.1.0

317 + # via h2

318 html2text==2025.4.15

319 # via gpt-oss

320 httpcore==1.0.6

@@ -317,6 +322,7 @@ httpcore==1.0.6

322 httpx==0.27.2

323 # via

324 # -r requirements/test.in

325 + # perceptron

326 # schemathesis

327 huggingface-hub==0.34.3

328 # via

@@ -338,6 +344,8 @@ hydra-core==1.3.2

344 # via

345 # lightly

346 # lightning

347 +hyperframe==6.1.0

348 + # via h2

349 hypothesis==6.131.0

350 # via

351 # hypothesis-graphql

@@ -549,6 +557,7 @@ numpy==1.26.4

557 # pandas

558 # patsy

559 # peft

560 + # perceptron

561 # pycocotools

562 # pyogrio

563 # rasterio

@@ -702,6 +711,8 @@ peft==0.16.0

711 # via

712 # -r requirements/test.in

713 # lm-eval

714 +perceptron==0.1.4

715 + # via -r requirements/test.in

716 pillow==10.4.0

717 # via

718 # genai-perf

@@ -709,6 +720,7 @@ pillow==10.4.0

720 # lightly-utils

721 # matplotlib

722 # mistral-common

723 + # perceptron

724 # scikit-image

725 # segmentation-models-pytorch

726 # sentence-transformers

@@ -952,6 +964,7 @@ rich==13.9.4

964 # genai-perf

965 # lightning

966 # mteb

967 + # perceptron

968 # typer

969 rioxarray==0.19.0

970 # via terratorch

@@ -1024,7 +1037,9 @@ shapely==2.1.1

1037 # geopandas

1038 # torchgeo

1039 shellingham==1.5.4

1027 - # via typer

1040 + # via

1041 + # perceptron

1042 + # typer

1043 six==1.16.0

1044 # via

1045 # junit-xml

@@ -1218,7 +1233,9 @@ typepy==1.3.2

1233 # pytablewriter

1234 # tabledata

1235 typer==0.15.2

1221 - # via fastsafetensors

1236 + # via

1237 + # fastsafetensors

1238 + # perceptron

1239 types-python-dateutil==2.9.0.20241206

1240 # via arrow

1241 typeshed-client==2.8.2

tests/models/multimodal/generation/test_common.py

+25

@@ -529,6 +529,31 @@ VLM_TEST_SETTINGS = {

529 use_tokenizer_eos=True,

530 auto_cls=AutoModelForImageTextToText,

531 ),

532 + "isaac": VLMTestInfo(

533 + models=["PerceptronAI/Isaac-0.1"],

534 + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),

535 + prompt_formatter=lambda img_prompt: (

536 + f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n"

537 + ),

538 + img_idx_to_prompt=lambda idx: "<image>",

539 + single_image_prompts=IMAGE_ASSETS.prompts(

540 + {

541 + "stop_sign": "<vlm_image>Please describe the image shortly.",

542 + "cherry_blossom": "<vlm_image>Please infer the season with reason.",

543 + }

544 + ),

545 + multi_image_prompt=(

546 + "Picture 1: <vlm_image>\n"

547 + "Picture 2: <vlm_image>\n"

548 + "Describe these two images with one paragraph respectively."

549 + ),

550 + enforce_eager=False,

551 + max_model_len=4096,

552 + max_num_seqs=2,

553 + hf_model_kwargs={"device_map": "auto"},

554 + patch_hf_runner=model_utils.isaac_patch_hf_runner,

555 + image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],

556 + ),

557 "kimi_vl": VLMTestInfo(

558 models=["moonshotai/Kimi-VL-A3B-Instruct"],

559 test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),

tests/models/multimodal/generation/vlm_utils/model_utils.py

+177

@@ -522,6 +522,183 @@ def h2ovl_patch_hf_runner(hf_model: HfRunner) -> HfRunner:

522 return hf_model

523

524

525 +def isaac_patch_hf_runner(hf_model: HfRunner) -> HfRunner:

526 + """Patch HF runner for Isaac:

527 + 1) Move processor outputs to model device

528 + 2) Ensure IsaacModel.forward returns hidden_states

529 + for compatibility with hidden_states_to_seq_logprobs()

530 + """

531 +

532 + from perceptron.tensorstream import TextType

533 + from perceptron.tensorstream.ops import compute_mrope_pos_tensor, modality_mask

534 + from transformers.modeling_outputs import BaseModelOutputWithPast

535 +

536 + def compute_position_ids_input_ids(input_ids: torch.Tensor) -> torch.Tensor:

537 + """

538 + Create 3D positional indices for token input.

539 + """

540 + batch_size, seq_length = input_ids.shape

541 + position_ids = torch.arange(seq_length, device=input_ids.device)

542 + position_ids = position_ids.view(1, -1).expand(batch_size, -1)

543 + position_ids = position_ids.unsqueeze(2).expand(-1, -1, 3) # Add 3D for MRoPE

544 + return position_ids

545 +

546 + model_device = next(hf_model.model.parameters()).device

547 +

548 + # ----------------------------

549 + # 1) Patch processor: move BatchFeature input_ids and TensorStream to model device

550 + # ----------------------------

551 + original_processor = hf_model.processor

552 +

553 + def patched_processor(*args, **kwargs):

554 + result = original_processor(*args, **kwargs)

555 + for k, v in result.data.items():

556 + result[k] = v.to(model_device)

557 + return result

558 +

559 + hf_model.processor = patched_processor

560 +

561 + tokenizer = AutoTokenizer.from_pretrained(

562 + hf_model.model_name, trust_remote_code=True

563 + )

564 +

565 + original_generate = hf_model.model.generate

566 +

567 + def patched_generate(*args, **kwargs):

568 + kwargs["pad_token_id"] = tokenizer.eos_token_id

569 + kwargs["eos_token_id"] = tokenizer.eos_token_id

570 + return original_generate(*args, **kwargs)

571 +

572 + hf_model.model.generate = patched_generate

573 +

574 + # ----------------------------

575 + # 2) Patch IsaacModel.forward: add hidden_states to the output

576 + # ----------------------------

577 + isaac_model = hf_model.model.model

578 +

579 + def patched_forward(

580 + self,

581 + input_ids=None,

582 + tensor_stream=None,

583 + attention_mask=None,

584 + position_ids=None,

585 + modality_tensor=None,

586 + past_key_values=None,

587 + inputs_embeds=None,

588 + use_cache=None,

589 + output_hidden_states=None,

590 + return_dict=None,

591 + cache_position=None,

592 + **kwargs,

593 + ):

594 + """

595 + Forward pass with MRoPE position embeddings.

596 + Computes position embeddings once and passes them through all layers.

597 + """

598 + output_hidden_states = (

599 + output_hidden_states

600 + if output_hidden_states is not None

601 + else self.config.output_hidden_states

602 + )

603 + use_cache = use_cache if use_cache is not None else self.config.use_cache

604 + return_dict = (

605 + return_dict if return_dict is not None else self.config.use_return_dict

606 + )

607 +

608 + # Get inputs

609 + if tensor_stream is not None and inputs_embeds is not None:

610 + raise ValueError("You cannot specify both tensor_stream and inputs_embeds")

611 + elif tensor_stream is not None:

612 + # Embed TensorStream directly

613 + inputs_embeds = self.embed_stream(tensor_stream)

614 + # Create modality tensor if not provided

615 + if modality_tensor is None:

616 + modality_tensor = modality_mask(tensor_stream)

617 + elif input_ids is not None and inputs_embeds is not None:

618 + raise ValueError(

619 + "You cannot specify both input_ids and inputs_embeds at the same time"

620 + )

621 + elif input_ids is not None:

622 + inputs_embeds = self.embed_tokens(input_ids)

623 + # Create text modality tensor if not provided

624 + if modality_tensor is None:

625 + batch_size, seq_length = input_ids.shape

626 + modality_tensor = torch.full(

627 + (batch_size, seq_length),

628 + TextType.text.value,

629 + device=input_ids.device,

630 + dtype=torch.long,

631 + )

632 + elif inputs_embeds is None:

633 + raise ValueError(

634 + "You have to specify either tensor_stream, input_ids or inputs_embeds"

635 + )

636 +

637 + # Create default position_ids if not provided

638 + if position_ids is None:

639 + if tensor_stream is not None:

640 + position_ids = compute_mrope_pos_tensor(tensor_stream) # (B,L,3)

641 + else:

642 + position_ids = compute_position_ids_input_ids(input_ids)

643 +

644 + # Compute MRoPE position embeddings if we have custom rotary_emb

645 + cos, sin = self.rotary_emb(position_ids, modality_tensor)

646 + cos = cos.to(inputs_embeds.dtype)

647 + sin = sin.to(inputs_embeds.dtype)

648 +

649 + # Prepare attention mask

650 + if attention_mask is not None:

651 + attention_mask = self._update_causal_mask(

652 + attention_mask, inputs_embeds, cache_position, past_key_values, False

653 + )

654 +

655 + # Initialize and collect hidden states

656 + hidden_states = inputs_embeds

657 + hidden_states_list: list[torch.Tensor] = []

658 +

659 + if output_hidden_states:

660 + hidden_states_list.append(hidden_states)

661 +

662 + for decoder_layer in self.layers:

663 + layer_outputs = decoder_layer(

664 + hidden_states,

665 + attention_mask=attention_mask,

666 + position_ids=position_ids,

667 + past_key_value=past_key_values,

668 + use_cache=use_cache,

669 + cache_position=cache_position,

670 + position_embeddings=(cos, sin),

671 + **kwargs,

672 + )

673 +

674 + hidden_states = (

675 + layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs

676 + )

677 +

678 + if output_hidden_states:

679 + hidden_states_list.append(hidden_states)

680 +

681 + # Final layer norm

682 + hidden_states = self.norm(hidden_states)

683 +

684 + if output_hidden_states:

685 + hidden_states_list.append(hidden_states)

686 +

687 + # Convert to tuple or None

688 + all_hidden_states = tuple(hidden_states_list) if output_hidden_states else None

689 +

690 + # Include hiden_states for compatibility with hidden_states_to_seq_logprobs()

691 + return BaseModelOutputWithPast(

692 + last_hidden_state=hidden_states,

693 + past_key_values=past_key_values,

694 + hidden_states=all_hidden_states,

695 + )

696 +

697 + isaac_model.forward = types.MethodType(patched_forward, isaac_model)

698 +

699 + return hf_model

700 +

701 +

702 def skyworkr1v_patch_hf_runner(hf_model: HfRunner) -> HfRunner:

703 """Patches and returns an instance of the HfRunner to use for SkyworkR1V."""

704

tests/models/registry.py

@@ -662,6 +662,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {

662 "HuggingFaceM4/Idefics3-8B-Llama3",

663 extras={"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"},

664 ),

665 + "IsaacForConditionalGeneration": _HfExamplesInfo(

666 + "PerceptronAI/Isaac-0.1",

667 + trust_remote_code=True,

668 + ),

669 "InternS1ForConditionalGeneration": _HfExamplesInfo(

670 "internlm/Intern-S1", trust_remote_code=True

671 ),

vllm/model_executor/models/isaac.py new

+1480

@@ -0,0 +1,1480 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +from __future__ import annotations

4 +

5 +import math

6 +from collections.abc import Iterable, Iterator, Mapping, Sequence

7 +from typing import Annotated, Any

8 +

9 +import numpy as np

10 +import PIL.Image

11 +import torch

12 +import torch.nn as nn

13 +import torch.nn.functional as F

14 +from einops import rearrange

15 +from transformers.image_processing_utils import BatchFeature

16 +from transformers.tokenization_utils import TensorType

17 +from typing_extensions import TypedDict, Unpack

18 +

19 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention

20 +from vllm.config import MultiModalConfig, VllmConfig

21 +from vllm.config.model import ModelConfig

22 +from vllm.distributed import parallel_state

23 +from vllm.distributed import utils as dist_utils

24 +from vllm.model_executor.layers.linear import (

25 + ColumnParallelLinear,

26 + QKVParallelLinear,

27 + ReplicatedLinear,

28 + RowParallelLinear,

29 +)

30 +from vllm.model_executor.layers.quantization import QuantizationConfig

31 +from vllm.model_executor.model_loader.weight_utils import (

32 + default_weight_loader,

33 +)

34 +from vllm.model_executor.models.interfaces import (

35 + MultiModalEmbeddings,

36 + SupportsLoRA,

37 + SupportsMRoPE,

38 + SupportsMultiModal,

39 + SupportsPP,

40 +)

41 +from vllm.model_executor.models.module_mapping import MultiModelKeys

42 +from vllm.model_executor.models.siglip import SiglipMLP

43 +from vllm.model_executor.models.utils import (

44 + AutoWeightsLoader,

45 + WeightsMapper,

46 + init_vllm_registered_model,

47 + maybe_prefix,

48 +)

49 +from vllm.multimodal import MULTIMODAL_REGISTRY

50 +from vllm.multimodal.inputs import (

51 + MultiModalDataDict,

52 + MultiModalFeatureSpec,

53 + MultiModalFieldConfig,

54 + MultiModalKwargsItems,

55 +)

56 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems

57 +from vllm.multimodal.processing import (

58 + BaseMultiModalProcessor,

59 + BaseProcessingInfo,

60 + PromptReplacement,

61 + PromptUpdate,

62 + PromptUpdateDetails,

63 +)

64 +from vllm.multimodal.profiling import BaseDummyInputsBuilder

65 +from vllm.sequence import IntermediateTensors

66 +from vllm.tokenizers import get_tokenizer

67 +from vllm.tokenizers.hf import get_cached_tokenizer

68 +from vllm.transformers_utils.configs import (

69 + IsaacConfig,

70 + PixelShuffleSiglip2VisionConfig,

71 +)

72 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

73 +

74 +

75 +def create_cumulative_seq_lengths(

76 + seq_sizes: torch.Tensor, device: torch.device

77 +) -> tuple[torch.Tensor, torch.Tensor]:

78 + """Create cumulative sequence lengths for variable-length attention."""

79 + cu_seqlens = torch.zeros(len(seq_sizes) + 1, dtype=torch.int32, device=device)

80 + cu_seqlens[1:] = seq_sizes.cumsum(0)

81 + max_seqlen = (

82 + seq_sizes.max()

83 + if len(seq_sizes) > 0

84 + else torch.tensor(0, dtype=torch.int32, device=device)

85 + )

86 + return cu_seqlens, max_seqlen

87 +

88 +

89 +class Siglip2VariableSequenceEmbeddings(nn.Module):

90 + def __init__(self, config: PixelShuffleSiglip2VisionConfig):

91 + super().__init__()

92 + self.config = config

93 + self.embed_dim = config.hidden_size

94 + self.patch_size = config.patch_size

95 +

96 + self.patch_embedding = ReplicatedLinear(

97 + input_size=config.num_channels * self.patch_size * self.patch_size,

98 + output_size=self.embed_dim,

99 + return_bias=False,

100 + )

101 +

102 + self.num_patches = config.num_patches

103 + self.position_embedding_size = int(self.num_patches**0.5)

104 + self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)

105 +

106 + def positional_embeddings(

107 + self, packed_seq_patches: tuple[torch.Tensor, torch.Tensor, torch.Tensor]

108 + ) -> torch.Tensor:

109 + # Prepare positional embeddings grid: (1, embed_dim, h, w)

110 + positional_embeddings = (

111 + self.position_embedding.weight.reshape(

112 + self.position_embedding_size, self.position_embedding_size, -1

113 + )

114 + .permute(2, 0, 1)

115 + .unsqueeze(0)

116 + )

117 +

118 + _seq_patches, _seq_sizes, spatial_shapes = packed_seq_patches

119 + pos_embeds_list = []

120 + mode = "bilinear"

121 + align_corners = False

122 + antialias = True

123 + for spatial_shape in spatial_shapes:

124 + height, width = int(spatial_shape[0]), int(spatial_shape[1])

125 + # Guard to ensure height and width are positive for torch.compile

126 + if height > 0 and width > 0:

127 + resized_pos_embed = F.interpolate(

128 + positional_embeddings,

129 + size=(height, width),

130 + mode=mode,

131 + align_corners=align_corners,

132 + antialias=antialias,

133 + )

134 + # Reshape from (1, embed_dim, height, width) to

135 + # (height*width, embed_dim)

136 + resized_pos_embed = resized_pos_embed.reshape(

137 + self.embed_dim, height * width

138 + ).transpose(0, 1)

139 + else:

140 + # Fallback - should never happen in practice

141 + resized_pos_embed = positional_embeddings.reshape(

142 + self.embed_dim,

143 + self.position_embedding_size * self.position_embedding_size,

144 + ).transpose(0, 1)[: height * width]

145 + pos_embeds_list.append(resized_pos_embed)

146 +

147 + # Concatenate all positional embeddings along the sequence dimension

148 + pos_embeds = torch.cat(pos_embeds_list, dim=0)

149 + return pos_embeds

150 +

151 + def forward(

152 + self, packed_seq_patches: tuple[torch.Tensor, torch.Tensor, torch.Tensor]

153 + ):

154 + seq_patches, _seq_sizes, _spatial_shapes = packed_seq_patches

155 +

156 + target_weight = self.patch_embedding.weight

157 + seq_patches = seq_patches.to(

158 + device=target_weight.device, dtype=target_weight.dtype

159 + )

160 + patch_embeds = self.patch_embedding(seq_patches)

161 + pos_embeds = self.positional_embeddings(packed_seq_patches)

162 +

163 + # Flatten patch embeddings to match positional embeddings format

164 + if patch_embeds.dim() == 3:

165 + patch_embeds = patch_embeds.view(-1, patch_embeds.size(-1))

166 +

167 + # Add positional embeddings to patch embeddings

168 + embeddings = patch_embeds + pos_embeds

169 + return embeddings

170 +

171 +

172 +def create_pixel_shuffle_index_map(

173 + seq_sizes: torch.Tensor,

174 + token_grids: torch.Tensor,

175 + scale_factor: int = 1,

176 + device: torch.device | None = None,

177 +) -> torch.Tensor:

178 + """

179 + Build a gather-index map that tells us, for every *output* token after

180 + pixel-shuffle, which `scale_factor**2` *input* tokens are being merged.

181 +

182 + Args

183 + ----

184 + seq_sizes : (num_images,) - #patches in each image (row-major order)

185 + token_grids : (num_images,2) - (height, width) for every image

186 + scale_factor : spatial down-scale factor (≥2)

187 + device : (optional) overrides `seq_sizes.device`

188 +

189 + Returns

190 + -------

191 + gather_idx : (new_total_seq_len, scale_factor**2) int64 tensor.

192 + gather_idx[i, j] is the *flat* index into the *original*

193 + packed sequence for the j-th sub-patch that forms the

194 + i-th output token.

195 + """

196 + if device is None:

197 + device = seq_sizes.device

198 +

199 + r = int(scale_factor)

200 + if r < 2:

201 + raise ValueError("`scale_factor` must be ≥ 2")

202 +

203 + # Safety: all spatial dims must be divisible by r

204 + # Cannot run under torch compile fullgraph mode hence

205 + if not torch.compiler.is_compiling() and not (

206 + (token_grids[:, 0] % r == 0).all() and (token_grids[:, 1] % r == 0).all()

207 + ):

208 + raise AssertionError(

209 + "Every (H,W) in `token_grids` must be divisible by "

210 + f"scale_factor={r}, got {token_grids.tolist()}"

211 + )

212 +

213 + gather_chunks: list[torch.Tensor] = []

214 + tok_offset = 0

215 +

216 + for seq_len, (h, w) in zip(seq_sizes.tolist(), token_grids.tolist(), strict=False):

217 + # Build the (H, W) grid of flat indices for this image

218 + grid = torch.arange(seq_len, device=device, dtype=torch.int64) + tok_offset

219 + grid = grid.view(h, w) # (H, W)

220 +

221 + # -------- identical ordering to your fixed-res routine --------

222 + # Step 1: split width into blocks of r

223 + grid = grid.view(h, w // r, r) # (H, W/r, r)

224 + # Step 2: now split height into blocks of r

225 + grid = grid.view(h // r, r, w // r, r) # (H/r, r, W/r, r)

226 + # Step 3: final permutation to (H/r, W/r, r, r)

227 + grid = grid.permute(0, 2, 1, 3).contiguous() # (H/r, W/r, r, r)

228 + # Step 4: each (r, r) block forms one output token

229 + gather_chunks.append(grid.reshape(-1, r * r)) # (H*W / r², r²)

230 +

231 + tok_offset += seq_len

232 +

233 + # Concatenate over all images in the packed batch

234 + gather_idx = torch.cat(gather_chunks, dim=0) # (Σ_i HᵢWᵢ/r², r²)

235 + return gather_idx

236 +

237 +

238 +def pixel_shuffle_varlen(

239 + x: torch.Tensor,

240 + token_grids: torch.Tensor,

241 + scale_factor: int = 1,

242 +) -> torch.Tensor:

243 + r"""Apply pixel shuffle to a packed vision sequence without unpacking per image.

244 +

245 + Args:

246 + x (`torch.Tensor`):

247 + Concatenated vision embeddings. Accepts `(seq_len, hidden_size)` or

248 + `(1, seq_len, hidden_size)` shapes produced by stacking image

249 + patches.

250 + token_grids (`torch.Tensor`):

251 + Integer tensor of shape `(num_images, 2)` whose rows give the

252 + `(height, width)` patch grid sizes corresponding to each image

253 + segment inside `x`.

254 + scale_factor (`int`, *optional*, defaults to 1):

255 + Spatial down-sampling factor specific to pixel shuffle. Values

256 + greater than one merge `scale_factor**2` neighboring patches into a

257 + single embedding channel-group.

258 +

259 + Returns:

260 + `torch.Tensor`: Pixel-shuffled embeddings with shape matching the input

261 + convention: `(seq_len, hidden_size * scale_factor**2)` when the input

262 + was 2D, or `(1, seq_len, hidden_size * scale_factor**2)` if the

263 + singleton batch dimension was present.

264 +

265 + Raises:

266 + ValueError: If more than one batch item is provided.

267 + """

268 + keep_batch_dim = x.dim() == 3

269 + if keep_batch_dim:

270 + if x.size(0) != 1:

271 + raise AssertionError("Packed sequence is expected to have batch_size == 1")

272 + x_ = x.squeeze(0) # (seq, embed)

273 + else:

274 + x_ = x # (seq, embed)

275 +

276 + embed_dim = x_.size(-1)

277 + r = int(scale_factor)

278 +

279 + # Calculate seq_sizes from token_grids

280 + seq_sizes = torch.prod(token_grids, dim=-1)

281 +

282 + # Build index map and gather in one go

283 + gather_idx = create_pixel_shuffle_index_map(

284 + seq_sizes=seq_sizes,

285 + token_grids=token_grids,

286 + scale_factor=r,

287 + device=x_.device,

288 + ) # (new_seq, r²)

289 +

290 + # Gather → (new_seq, r², embed_dim)

291 + gathered = x_[gather_idx] # fancy indexing keeps gradient

292 +

293 + # Merge the r² group dimension into channels to finish the shuffle

294 + out = gathered.reshape(gathered.size(0), embed_dim * r * r)

295 +

296 + # Restore batch dimension if needed

297 + if keep_batch_dim:

298 + out = out.unsqueeze(0)

299 + return out

300 +

301 +

302 +# ============================================================================

303 +# Configuration

304 +# ============================================================================

305 +

306 +MAX_PIXELS = 60_000_000 # 60-megapixel ceiling ≈ 8200 × 7300 px

307 +

308 +# Vision preprocessing constants

309 +VISION_MEAN = (0.5, 0.5, 0.5)

310 +VISION_STD = (0.5, 0.5, 0.5)

311 +VISION_SCALE = 1 / 255

312 +

313 +

314 +def _make_writeable(arr: np.ndarray) -> np.ndarray:

315 + """Return *arr* itself if it is already writeable, otherwise try to flip the

316 + write flag in-place and finally fall back to `arr.copy()`.

317 + This guarantees the buffer handed to `torch.from_numpy()` is always

318 + writeable, silencing the PyTorch warning about undefined behaviour.

319 + """

320 + if arr.flags.writeable:

321 + return arr

322 +

323 + # First, try the cheap path — in-place flag toggle (works for mmap'd arrays

324 + # and some shared memory buffers):

325 + try:

326 + arr.setflags(write=True)

327 + return arr # success: no data copy

328 + except ValueError:

329 + # Buffer is inherently read-only (e.g. backed by PyAV / PIL): make copy

330 + return arr.copy()

331 +

332 +

333 +def extract_image_pil(image: PIL.Image.Image) -> torch.Tensor | None:

334 + if image.width * image.height > MAX_PIXELS:

335 + raise ValueError(

336 + f"Image (w={image.width}, h={image.height}) > MAX=`{MAX_PIXELS}`"

337 + )

338 + img = image if image.mode == "RGB" else image.convert("RGB")

339 + arr = np.asarray(img)

340 + arr = _make_writeable(arr)

341 + return torch.from_numpy(arr)

342 +

343 +

344 +def get_image_size_for_max_num_patches(

345 + image_height: int,

346 + image_width: int,

347 + patch_size: int,

348 + max_num_patches: int,

349 + min_num_patches: int | None = None,

350 + eps: float = 1e-5,

351 + pixel_shuffle_scale: int = 1,

352 +) -> tuple[int, int]:

353 + r"""Compute a target resolution whose patch grid satisfies patching parametrization.

354 +

355 + Args:

356 + image_height (`int`):

357 + Height in pixels of the source image prior to any resizing.

358 + image_width (`int`):

359 + Width in pixels of the source image prior to any resizing.

360 + patch_size (`int`):

361 + Size of the square patch used by the vision encoder.

362 + max_num_patches (`int`):

363 + Upper bound on `(height / patch_size) * (width / patch_size)` after

364 + resizing.

365 + min_num_patches (`int`, *optional*):

366 + Lower bound on the number of patches. When provided the image will

367 + be scaled up if necessary.

368 + eps (`float`, *optional*, defaults to 1e-5):

369 + Convergence tolerance for the internal binary search to determine

370 + the target dimensions.

371 + pixel_shuffle_scale (`int`, *optional*, defaults to 1):

372 + Additional stride multiplier applied when pixel shuffle later

373 + reduces spatial resolution.

374 +

375 + Returns:

376 + `tuple[int, int]`: Height and width (in pixels) that are multiples of

377 + `patch_size * pixel_shuffle_scale` and respect both the maximum and

378 + optional minimum patch-count constraints.

379 + """

380 +

381 + def get_scaled_image_size(scale, original_size, patch_size, pixel_shuffle_scale):

382 + scaled_size = scale * original_size

383 + divisor = patch_size * pixel_shuffle_scale

384 + scaled_size = math.ceil(scaled_size / divisor) * divisor

385 + scaled_size = max(divisor, scaled_size)

386 + return int(scaled_size)

387 +

388 + # Ensure divisibility

389 + divisor = patch_size * pixel_shuffle_scale

390 + adjusted_height = math.ceil(image_height / divisor) * divisor

391 + adjusted_height = max(divisor, adjusted_height)

392 + adjusted_width = math.ceil(image_width / divisor) * divisor

393 + adjusted_width = max(divisor, adjusted_width)

394 +

395 + num_patches = (adjusted_height / patch_size) * (adjusted_width / patch_size)

396 +

397 + if min_num_patches is not None and num_patches < min_num_patches:

398 + # Scale up

399 + scale_min, scale_max = 1.0, 100.0

400 + while (scale_max - scale_min) >= eps:

401 + scale = (scale_min + scale_max) / 2

402 + target_height = get_scaled_image_size(

403 + scale, image_height, patch_size, pixel_shuffle_scale

404 + )

405 + target_width = get_scaled_image_size(

406 + scale, image_width, patch_size, pixel_shuffle_scale

407 + )

408 + num_patches = (target_height / patch_size) * (target_width / patch_size)

409 + if num_patches >= min_num_patches:

410 + scale_max = scale

411 + else:

412 + scale_min = scale

413 + scale = scale_max

414 + target_height = get_scaled_image_size(

415 + scale, image_height, patch_size, pixel_shuffle_scale

416 + )

417 + target_width = get_scaled_image_size(

418 + scale, image_width, patch_size, pixel_shuffle_scale

419 + )

420 + return target_height, target_width

421 + elif num_patches <= max_num_patches:

422 + return adjusted_height, adjusted_width

423 + else:

424 + # Scale down

425 + scale_min, scale_max = eps / 10, 1.0

426 + while (scale_max - scale_min) >= eps:

427 + scale = (scale_min + scale_max) / 2

428 + target_height = get_scaled_image_size(

429 + scale, image_height, patch_size, pixel_shuffle_scale

430 + )

431 + target_width = get_scaled_image_size(

432 + scale, image_width, patch_size, pixel_shuffle_scale

433 + )

434 + num_patches = (target_height / patch_size) * (target_width / patch_size)

435 + if num_patches <= max_num_patches:

436 + scale_min = scale

437 + else:

438 + scale_max = scale

439 + scale = scale_min

440 + target_height = get_scaled_image_size(

441 + scale, image_height, patch_size, pixel_shuffle_scale

442 + )

443 + target_width = get_scaled_image_size(

444 + scale, image_width, patch_size, pixel_shuffle_scale

445 + )

446 + return target_height, target_width

447 +

448 +

449 +_MEAN_TENSOR = torch.tensor(VISION_MEAN, dtype=torch.float32).view(1, 1, 1, -1)

450 +_STD_TENSOR = torch.tensor(VISION_STD, dtype=torch.float32).view(1, 1, 1, -1)

451 +

452 +

453 +def _resolve_vision_token_id(model_config: ModelConfig, vision_token: str) -> int:

454 + tokenizer_name = model_config.tokenizer or model_config.model

455 + tokenizer = get_cached_tokenizer(

456 + get_tokenizer(

457 + tokenizer_name,

458 + tokenizer_mode=model_config.tokenizer_mode,

459 + trust_remote_code=model_config.trust_remote_code,

460 + revision=model_config.tokenizer_revision or model_config.revision,

461 + )

462 + )

463 + return tokenizer.encode(vision_token, add_special_tokens=False)[0]

464 +

465 +

466 +def prepare_image_tensor(

467 + image: torch.Tensor,

468 + scale: float = VISION_SCALE,

469 +) -> torch.Tensor:

470 + r"""Standardize RGB images prior to patch extraction via rescaling and whitening.

471 +

472 + Args:

473 + image (`torch.Tensor`):

474 + Tensor with shape `(..., height, width, 3)` containing RGB values.

475 + The tensor is converted to floating point if needed.

476 + scale (`float`, *optional*, defaults to `VISION_SCALE`):

477 + Scalar multiplier applied before normalization.

478 + Returns:

479 + `torch.Tensor`: Normalized tensor with the same shape as the input and

480 + dtype `torch.float32`.

481 + """

482 + if not torch.is_floating_point(image):

483 + image = image.float()

484 + rescaled = image * scale

485 +

486 + # Use precomputed tensors and move to the correct device if needed

487 + mean_tensor = _MEAN_TENSOR.to(image.device)

488 + std_tensor = _STD_TENSOR.to(image.device)

489 +

490 + normalized = (rescaled - mean_tensor) / std_tensor

491 + return normalized

492 +

493 +

494 +def patchify_vision(image: torch.Tensor, patch_size: int) -> torch.Tensor:

495 + r"""Convert normalized images into flattened ViT-style patches.

496 +

497 + Args:

498 + image (`torch.Tensor`):

499 + Tensor of shape `(num_images, height, width, channels)`.

500 + patch_size (`int`):

501 + Edge length of the square patches

502 +

503 + Returns:

504 + `torch.Tensor`:

505 + Patch tensor where each position stores the flattened pixels

506 + belonging to that patch.

507 +

508 + Raises:

509 + ValueError: If `height` or `width` is not divisible by `patch_size`.

510 + """

511 + num_images, height, width, channels = image.shape

512 + if height % patch_size or width % patch_size:

513 + raise ValueError(

514 + "Dimensions of images "

515 + f"{image.shape} are not divisible by patch_size={patch_size}."

516 + )

517 + patches = image.reshape(

518 + num_images,

519 + height // patch_size,

520 + patch_size,

521 + width // patch_size,

522 + patch_size,

523 + channels,

524 + )

525 + patches = patches.permute(0, 1, 3, 2, 4, 5)

526 + patches = patches.reshape(

527 + num_images,

528 + height // patch_size,

529 + width // patch_size,

530 + channels * patch_size * patch_size,

531 + )

532 + return patches

533 +

534 +

535 +def process_vision_for_patches(

536 + images: torch.Tensor,

537 + patch_size: int,

538 + max_num_patches: int,

539 + min_num_patches: int | None = None,

540 + pixel_shuffle_scale: int = 1,

541 +) -> tuple[torch.Tensor, list[int]]:

542 + r"""Resize, normalize, and patchify RGB images for the vision encoder.

543 +

544 + Args:

545 + images (`torch.Tensor`):

546 + Either `(height, width, channels)` for a single image or

547 + `(num_images, height, width, channels)` for a batch. Channels are

548 + expected to be RGB.

549 + patch_size (`int`):

550 + Edge length of square patches; implictly controls resize grid granularity.

551 + max_num_patches (`int`):

552 + Maximum number of patches allowed after resizing.

553 + min_num_patches (`int`, *optional*):

554 + Minimum number of patches. If provided, the routine upsamples images

555 + as needed to satisfy the lower bound.

556 + pixel_shuffle_scale (`int`, *optional*, defaults to 1):

557 + Pixel shuffle scale factor; influences the target grid that the

558 + function produces.

559 +

560 + Returns:

561 + `tuple[torch.Tensor, list[int]]`: A pair `(patches, dims_virtual)`

562 + where `patches` has shape `(num_images, target_h / patch_size, target_w

563 + / patch_size, channels * patch_size**2)` and `dims_virtual` encodes

564 + effective `(images, height, width)` dimensions after optional pixel

565 + shuffling.

566 + """

567 + # Add batch dim if single image

568 + if images.dim() == 3:

569 + images = images.unsqueeze(0)

570 +

571 + # Permute to channel first for resize

572 + images = images.permute(0, 3, 1, 2)

573 +

574 + # Get target dimensions

575 + _, _, orig_height, orig_width = images.shape

576 + target_height, target_width = get_image_size_for_max_num_patches(

577 + orig_height,

578 + orig_width,

579 + patch_size,

580 + max_num_patches,

581 + min_num_patches=min_num_patches,

582 + pixel_shuffle_scale=pixel_shuffle_scale,

583 + )

584 +

585 + # Resize

586 + images = F.interpolate(

587 + images,

588 + size=(target_height, target_width),

589 + mode="bilinear",

590 + align_corners=False,

591 + )

592 +

593 + # Back to channel last

594 + images = images.permute(0, 2, 3, 1)

595 +

596 + # Normalize

597 + images = prepare_image_tensor(images)

598 +

599 + # Patchify

600 + patches = patchify_vision(images, patch_size=patch_size)

601 +

602 + # Calculate dimensions for the patches

603 + n_images, h_patches, w_patches, _ = patches.shape

604 + dims_virtual = (

605 + [1, h_patches, w_patches]

606 + if pixel_shuffle_scale == 1

607 + else [1, h_patches // pixel_shuffle_scale, w_patches // pixel_shuffle_scale]

608 + )

609 +

610 + return patches, dims_virtual

611 +

612 +

613 +class IsaacImageProcessorKwargs(TypedDict, total=False):

614 + patch_size: int

615 + max_num_patches: int

616 + min_num_patches: int

617 + pixel_shuffle_scale: int

618 +

619 +

620 +class IsaacImageProcessor:

621 + patch_size = 16

622 + max_num_patches = 6144

623 + min_num_patches = 256

624 + pixel_shuffle_scale = 2

625 +

626 + valid_kwargs = IsaacImageProcessorKwargs

627 + model_input_names = ["pixel_values", "image_grid_thw"]

628 +

629 + def __init__(self, kwargs):

630 + self.patch_size = kwargs.pop("patch_size", self.patch_size)

631 + self.vision_max_num_patches = kwargs.pop(

632 + "vision_max_num_patches", self.max_num_patches

633 + )

634 + self.vision_min_num_patches = kwargs.pop(

635 + "vision_min_num_patches", self.min_num_patches

636 + )

637 + self.pixel_shuffle_scale = kwargs.pop("pixel_shuffle_scale", 2)

638 +

639 + def preprocess(

640 + self,

641 + images: list[torch.Tensor],

642 + return_tensors: str | TensorType | None,

643 + **kwargs: Unpack[IsaacImageProcessorKwargs],

644 + ) -> BatchFeature:

645 + """Preprocess images into format compatibile with vLLM input processing."""

646 +

647 + all_pixel_values: list[torch.Tensor] = []

648 + all_image_grids: list[torch.Tensor] = []

649 +

650 + for image in images:

651 + image_tensor = extract_image_pil(image)

652 +

653 + patches, dims_virtual = process_vision_for_patches(

654 + image_tensor,

655 + patch_size=self.patch_size,

656 + max_num_patches=self.vision_max_num_patches,

657 + min_num_patches=self.vision_min_num_patches,

658 + pixel_shuffle_scale=self.pixel_shuffle_scale,

659 + )

660 +

661 + # Isaac packs a dummy temporal dim for images

662 + patches = patches.unsqueeze(1) # [N, T=1, Hp, Wp, D]

663 +

664 + hp, wp, dim = patches.shape[-3], patches.shape[-2], patches.shape[-1]

665 + current_num_patches = hp * wp

666 + pixel_values = patches.reshape(current_num_patches, dim) # [N_tokens, D]

667 +

668 + # Use real patch dimensions for image_grid_thw, not virtual dimensions

669 + # This ensures the vision model receives correct grid info for pixel shuffle

670 + dims_real = [1, hp, wp] # Real patch dimensions

671 + image_grid_thw = torch.tensor(dims_real).unsqueeze(0)

672 +

673 + all_pixel_values.append(pixel_values)

674 + all_image_grids.append(image_grid_thw)

675 +

676 + if all_pixel_values:

677 + final_pixel_values = torch.cat(all_pixel_values, dim=0)

678 + final_image_grids = torch.cat(all_image_grids, dim=0)

679 + else:

680 + final_pixel_values = torch.empty(0, 0)

681 + final_image_grids = torch.empty(0, 3)

682 +

683 + return BatchFeature(

684 + data={

685 + "pixel_values": final_pixel_values,

686 + "image_grid_thw": final_image_grids,

687 + },

688 + tensor_type=return_tensors,

689 + )

690 +

691 +

692 +class IsaacProcessor:

693 + """Processor wrapper (tokenizer + IsaacImageProcessor)."""

694 +

695 + def __init__(self, image_processor=None, tokenizer=None, **kwargs):

696 + self.image_token = kwargs.pop("image_token", "<image>")

697 + self.image_processor = image_processor or IsaacImageProcessor(kwargs)

698 + self.tokenizer = tokenizer

699 +

700 + def __call__(self, text=None, images=None, **kwargs) -> BatchFeature:

701 + result = {}

702 +

703 + if images is not None:

704 + image_inputs = self.image_processor.preprocess(images, **kwargs)

705 + image_grid_thw = image_inputs["image_grid_thw"]

706 + result.update(image_inputs)

707 +

708 + if text is not None:

709 + if not isinstance(text, list):

710 + text = [text]

711 +

712 + text = text.copy() # below lines change text in-place

713 + merge_length = self.image_processor.pixel_shuffle_scale**2

714 + index = 0

715 + for i in range(len(text)):

716 + while self.image_token in text[i]:

717 + num_image_tokens = image_grid_thw[index].prod() // merge_length

718 + text[i] = text[i].replace(

719 + self.image_token, "<|placeholder|>" * num_image_tokens, 1

720 + )

721 + index += 1

722 + text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>")

723 +

724 + if text is not None:

725 + result.update(self.tokenizer(text, **kwargs))

726 +

727 + return BatchFeature(result)

728 +

729 + def apply_chat_template(

730 + self,

731 + messages: list[dict[str, Any]],

732 + tokenize: bool = False,

733 + add_generation_prompt: bool = False,

734 + **kwargs,

735 + ) -> Any:

736 + # Convert mixed content messages to simple text format

737 + processed_messages = []

738 +

739 + for message in messages:

740 + if "content" in message and isinstance(message["content"], list):

741 + # Handle mixed content (text + image)

742 + text_parts = []

743 + for content_item in message["content"]:

744 + if content_item.get("type") == "text":

745 + text_parts.append(content_item.get("text", ""))

746 + elif content_item.get("type") == "image":

747 + # Replace image with vision token

748 + text_parts.append(self.image_token)

749 +

750 + processed_message = {

751 + "role": message.get("role", "user"),

752 + "content": "".join(text_parts),

753 + }

754 + processed_messages.append(processed_message)

755 + else:

756 + # Regular text message

757 + processed_messages.append(message)

758 +

759 + return self.tokenizer.apply_chat_template(

760 + processed_messages,

761 + tokenize=tokenize,

762 + add_generation_prompt=add_generation_prompt,

763 + **kwargs,

764 + )

765 +

766 +

767 +class IsaacProcessingInfo(BaseProcessingInfo):

768 + def get_hf_config(self) -> IsaacConfig:

769 + if hasattr(self.ctx, "get_hf_config"):

770 + original_config = self.ctx.get_hf_config()

771 + # Map HF config parameters to our vLLM config parameters

772 + return IsaacConfig(

773 + # Vision parameters - map from HF names

774 + vision_config=getattr(original_config, "vision_config", None),

775 + vision_patch_size=getattr(original_config, "video_patch_size", 16),

776 + vision_max_num_patches=getattr(

777 + original_config, "vision_max_num_patches", 256

778 + ),

779 + vision_min_num_patches=getattr(

780 + original_config, "vision_min_num_patches", None

781 + ),

782 + pixel_shuffle_scale=getattr(original_config, "pixel_shuffle_scale", 1),

783 + max_sequence_length=getattr(

784 + original_config, "max_sequence_length", 16384

785 + ),

786 + vision_token=getattr(original_config, "vision_token", "<image>"),

787 + vision_attn_implementation=getattr(

788 + original_config, "vision_attn_implementation", None

789 + ),

790 + )

791 + return IsaacConfig()

792 +

793 + def get_hf_processor(self, **kwargs) -> IsaacProcessor:

794 + hf_config = self.get_hf_config()

795 + processor_kwargs = {

796 + "image_token": hf_config.vision_token,

797 + }

798 + processor_kwargs.update(kwargs)

799 + return self.ctx.get_hf_processor(IsaacProcessor, **processor_kwargs)

800 +

801 + def get_tokenizer(self):

802 + return self.ctx.tokenizer

803 +

804 + def get_image_size_with_most_features(self) -> ImageSize:

805 + hf_config = self.get_hf_config()

806 + # Get target dimensions

807 + target_height, target_width = get_image_size_for_max_num_patches(

808 + 9999999,

809 + 9999999,

810 + hf_config.video_patch_size,

811 + hf_config.vision_max_num_patches,

812 + min_num_patches=hf_config.vision_min_num_patches,

813 + pixel_shuffle_scale=hf_config.pixel_shuffle_scale,

814 + )

815 + return ImageSize(width=target_width, height=target_height)

816 +

817 + def get_image_processor(self, **kwargs) -> IsaacImageProcessor:

818 + return self.get_hf_processor(**kwargs).image_processor

819 +

820 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

821 + return {"image": None}

822 +

823 + def get_mm_max_tokens_per_item(

824 + self,

825 + seq_len: int,

826 + mm_counts: Mapping[str, int],

827 + ) -> Mapping[str, int]:

828 + hf_config = self.get_hf_config()

829 + num_vision_tokens = hf_config.vision_max_num_patches // (

830 + hf_config.pixel_shuffle_scale**2

831 + )

832 + return {"image": num_vision_tokens}

833 +

834 +

835 +class IsaacDummyInputsBuilder(BaseDummyInputsBuilder[IsaacProcessingInfo]):

836 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

837 + num_images = mm_counts.get("image", 0)

838 +

839 + hf_processor = self.info.get_hf_processor()

840 + image_token: str = hf_processor.image_token

841 +

842 + return image_token * num_images

843 +

844 + def get_dummy_mm_data(

845 + self,

846 + seq_len: int,

847 + mm_counts: Mapping[str, int],

848 + mm_options: Mapping[str] | None = None,

849 + ) -> MultiModalDataDict:

850 + num_images = mm_counts.get("image", 0)

851 +

852 + target_width, target_height = self.info.get_image_size_with_most_features()

853 + image_overrides = mm_options.get("image") if mm_options else None

854 +

855 + return {

856 + "image": self._get_dummy_images(

857 + width=target_width,

858 + height=target_height,

859 + num_images=num_images,

860 + overrides=image_overrides,

861 + ),

862 + }

863 +

864 +

865 +class IsaacImagePixelInputs(TensorSchema):

866 + """

867 + Schema for validating Isaac image inputs.

868 +

869 + Dimensions:

870 + - np: Number of patches

871 + - d: Patch dimension

872 + - ni: Number of images

873 +

874 + The schema enforces:

875 + - pixel_values must be 2D: (num_patches, patch_dim)

876 + - image_grid_thw must be 2D: (num_images, 3)

877 + where 3 represents [T, H, W]

878 + """

879 +

880 + pixel_values: Annotated[

881 + torch.Tensor,

882 + TensorShape("np", "d"),

883 + ]

884 +

885 + image_grid_thw: Annotated[

886 + torch.Tensor,

887 + TensorShape("ni", 3),

888 + ]

889 +

890 +

891 +class IsaacMultiModalProcessor(BaseMultiModalProcessor):

892 + def _get_mm_fields_config(

893 + self,

894 + hf_inputs: BatchFeature,

895 + hf_processor_mm_kwargs: Mapping[str, object],

896 + ) -> Mapping[str, MultiModalFieldConfig]:

897 + # Configure multimodal fields for Isaac model

898 + image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))

899 + image_grid_sizes = image_grid_thw.prod(-1)

900 +

901 + return {

902 + "pixel_values": MultiModalFieldConfig.flat_from_sizes(

903 + "image", image_grid_sizes

904 + ),

905 + "image_grid_thw": MultiModalFieldConfig.batched("image"),

906 + }

907 +

908 + def _get_prompt_updates(

909 + self,

910 + mm_items: MultiModalDataItems,

911 + hf_processor_mm_kwargs: Mapping[str, Any],

912 + out_mm_kwargs: MultiModalKwargsItems,

913 + ) -> Sequence[PromptUpdate]:

914 + image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)

915 +

916 + pixel_shuffle_scale = getattr(image_processor, "pixel_shuffle_scale", 2)

917 + merge_length = pixel_shuffle_scale**2

918 +

919 + def get_replacement_isaac(item_idx: int):

920 + out_item = out_mm_kwargs["image"][item_idx]

921 + grid_thw = out_item["image_grid_thw"].data

922 + assert isinstance(grid_thw, torch.Tensor)

923 +

924 + feature_size = int(grid_thw.prod()) // merge_length

925 + repl_full = "<|image_pad|>" * feature_size

926 + return PromptUpdateDetails.select_text(repl_full, "<|image_pad|>")

927 +

928 + return [

929 + PromptReplacement(

930 + modality="image",

931 + target="<image>",

932 + replacement=get_replacement_isaac,

933 + )

934 + ]

935 +

936 +

937 +class Siglip2VisionAttention(nn.Module):

938 + def __init__(

939 + self,

940 + config: PixelShuffleSiglip2VisionConfig,

941 + quant_config: QuantizationConfig | None = None,

942 + *,

943 + prefix: str = "",

944 + multimodal_config: MultiModalConfig | None = None,

945 + ) -> None:

946 + super().__init__()

947 +

948 + use_data_parallel = (

949 + multimodal_config.mm_encoder_tp_mode == "data"

950 + if multimodal_config

951 + else False

952 + )

953 + self.tp_size = (

954 + 1

955 + if use_data_parallel

956 + else parallel_state.get_tensor_model_parallel_world_size()

957 + )

958 + self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

959 + self.hidden_size_per_attention_head = dist_utils.divide(

960 + config.hidden_size, config.num_attention_heads

961 + )

962 + self.num_attention_heads_per_partition = dist_utils.divide(

963 + config.num_attention_heads, self.tp_size

964 + )

965 +

966 + self.qkv_proj = QKVParallelLinear(

967 + hidden_size=config.hidden_size,

968 + head_size=self.hidden_size_per_attention_head,

969 + total_num_heads=config.num_attention_heads,

970 + total_num_kv_heads=config.num_attention_heads,

971 + bias=True,

972 + quant_config=quant_config,

973 + prefix=f"{prefix}.qkv_proj",

974 + disable_tp=use_data_parallel,

975 + )

976 + self.out_proj = RowParallelLinear(

977 + input_size=config.hidden_size,

978 + output_size=config.hidden_size,

979 + quant_config=quant_config,

980 + prefix=f"{prefix}.out_proj",

981 + disable_tp=use_data_parallel,

982 + )

983 +

984 + self.attn = MMEncoderAttention(

985 + num_heads=self.num_attention_heads_per_partition,

986 + head_size=self.hidden_size_per_attention_head,

987 + prefix=f"{prefix}.attn",

988 + multimodal_config=multimodal_config,

989 + )

990 +

991 + def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:

992 + seq_len, bs, _ = qkv.shape

993 + q, k, v = qkv.chunk(3, dim=2)

994 + new_shape = (

995 + seq_len,

996 + bs,

997 + self.num_attention_heads_per_partition,

998 + self.hidden_size_per_attention_head,

999 + )

1000 + q, k, v = (x.view(*new_shape) for x in (q, k, v))

1001 + return q, k, v

1002 +

1003 + def forward(

1004 + self,

1005 + hidden_states: torch.Tensor,

1006 + *,

1007 + cu_seqlens: torch.Tensor,

1008 + max_seqlen: torch.Tensor | None,

1009 + ) -> torch.Tensor:

1010 + batch_size, _, _ = hidden_states.shape

1011 + if batch_size != 1:

1012 + raise ValueError("packed variable-length attention expects batch_size=1")

1013 + x = rearrange(hidden_states, "b s d -> s b d")

1014 + x, _ = self.qkv_proj(x)

1015 + q, k, v = self.split_qkv(x)

1016 + q, k, v = (rearrange(t, "s b h d -> b s h d") for t in (q, k, v))

1017 +

1018 + context_layer = self.attn(

1019 + query=q,

1020 + key=k,

1021 + value=v,

1022 + cu_seqlens=cu_seqlens,

1023 + max_seqlen=max_seqlen,

1024 + )

1025 + context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous()

1026 +

1027 + output, _ = self.out_proj(context_layer)

1028 + output = rearrange(output, "s b d -> b s d")

1029 + return output

1030 +

1031 +

1032 +class Siglip2EncoderLayer(nn.Module):

1033 + def __init__(

1034 + self,

1035 + config: PixelShuffleSiglip2VisionConfig,

1036 + quant_config: QuantizationConfig | None = None,

1037 + *,

1038 + prefix: str = "",

1039 + multimodal_config: MultiModalConfig | None = None,

1040 + ) -> None:

1041 + super().__init__()

1042 + self.embed_dim = config.hidden_size

1043 + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

1044 + self.self_attn = Siglip2VisionAttention(

1045 + config,

1046 + quant_config=quant_config,

1047 + prefix=f"{prefix}.self_attn",

1048 + multimodal_config=multimodal_config,

1049 + )

1050 + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

1051 + self.mlp = SiglipMLP(

1052 + config,

1053 + quant_config=quant_config,

1054 + prefix=f"{prefix}.mlp",

1055 + )

1056 +

1057 + def forward(

1058 + self,

1059 + hidden_states: torch.Tensor,

1060 + *,

1061 + cu_seqlens: torch.Tensor,

1062 + max_seqlen: torch.Tensor | None,

1063 + ) -> torch.Tensor:

1064 + residual = hidden_states

1065 +

1066 + hidden_states = self.layer_norm1(hidden_states)

1067 + hidden_states = self.self_attn(

1068 + hidden_states=hidden_states,

1069 + cu_seqlens=cu_seqlens,

1070 + max_seqlen=max_seqlen,

1071 + )

1072 + hidden_states = residual + hidden_states

1073 +

1074 + residual = hidden_states

1075 + hidden_states = self.layer_norm2(hidden_states)

1076 + hidden_states = self.mlp(hidden_states)

1077 + hidden_states = residual + hidden_states

1078 +

1079 + return hidden_states

1080 +

1081 +

1082 +class Siglip2Encoder(nn.Module):

1083 + def __init__(

1084 + self,

1085 + config: PixelShuffleSiglip2VisionConfig,

1086 + quant_config: QuantizationConfig | None = None,

1087 + *,

1088 + prefix: str = "",

1089 + multimodal_config: MultiModalConfig | None = None,

1090 + ) -> None:

1091 + super().__init__()

1092 + self.config = config

1093 + self.layers = nn.ModuleList(

1094 + [

1095 + Siglip2EncoderLayer(

1096 + config,

1097 + quant_config=quant_config,

1098 + prefix=f"{prefix}.layers.{layer_idx}",

1099 + multimodal_config=multimodal_config,

1100 + )

1101 + for layer_idx in range(config.num_hidden_layers)

1102 + ]

1103 + )

1104 +

1105 + def forward(

1106 + self,

1107 + inputs_embeds: torch.Tensor,

1108 + *,

1109 + cu_seqlens: torch.Tensor | None = None,

1110 + max_seqlen: torch.Tensor | None = None,

1111 + ) -> torch.Tensor:

1112 + hidden_states = inputs_embeds

1113 + for encoder_layer in self.layers:

1114 + hidden_states = encoder_layer(

1115 + hidden_states,

1116 + cu_seqlens=cu_seqlens,

1117 + max_seqlen=max_seqlen,

1118 + )

1119 + return hidden_states

1120 +

1121 +

1122 +class Siglip2VisionTransformer(nn.Module):

1123 + def __init__(

1124 + self,

1125 + config: PixelShuffleSiglip2VisionConfig,

1126 + quant_config: QuantizationConfig | None = None,

1127 + prefix: str = "",

1128 + multimodal_config: MultiModalConfig | None = None,

1129 + ):

1130 + super().__init__()

1131 + self.config = config

1132 + self.quant_config = quant_config

1133 + embed_dim = config.hidden_size

1134 +

1135 + self.embeddings = Siglip2VariableSequenceEmbeddings(config)

1136 + self.pixel_shuffle_scale_factor = config.pixel_shuffle_scale_factor

1137 + self.encoder = Siglip2Encoder(

1138 + config,

1139 + quant_config=quant_config,

1140 + prefix=f"{prefix}.encoder",

1141 + multimodal_config=multimodal_config,

1142 + )

1143 + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

1144 +

1145 + def forward(

1146 + self,

1147 + packed_seq_patches: tuple[torch.Tensor, torch.Tensor],

1148 + ) -> torch.Tensor:

1149 + r"""

1150 + spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):

1151 + Tensor containing the spatial dimensions (height, width)

1152 + of the input images.

1153 + """

1154 +

1155 + seq_patches, token_grids = packed_seq_patches

1156 + seq_sizes = torch.prod(token_grids, dim=-1)

1157 +

1158 + # Get embeddings from packed sequence

1159 + hidden_states = self.embeddings((seq_patches, seq_sizes, token_grids))

1160 +

1161 + # Add a pseudo batch dimension for the encoder

1162 + hidden_states = hidden_states.unsqueeze(0)

1163 +

1164 + cu_seqlens, max_seqlen = create_cumulative_seq_lengths(

1165 + seq_sizes, hidden_states.device

1166 + )

1167 +

1168 + hidden_states = self.encoder(

1169 + inputs_embeds=hidden_states,

1170 + cu_seqlens=cu_seqlens,

1171 + max_seqlen=max_seqlen,

1172 + )

1173 + hidden_states = self.post_layernorm(hidden_states)

1174 +

1175 + if self.pixel_shuffle_scale_factor > 1:

1176 + hidden_states = pixel_shuffle_varlen(

1177 + x=hidden_states,

1178 + token_grids=token_grids,

1179 + scale_factor=self.pixel_shuffle_scale_factor,

1180 + )

1181 + # Remove the pseudo batch dimension we added earlier

1182 + hidden_states = hidden_states.squeeze(0)

1183 +

1184 + # return last_hidden_state

1185 + return hidden_states

1186 +

1187 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1188 + stacked_params_mapping = [

1189 + # (param_name, shard_name, shard_id)

1190 + ("qkv_proj", "q_proj", "q"),

1191 + ("qkv_proj", "k_proj", "k"),

1192 + ("qkv_proj", "v_proj", "v"),

1193 + ]

1194 + params_dict = dict(self.named_parameters())

1195 + loaded_params: set[str] = set()

1196 +

1197 + for name, loaded_weight in weights:

1198 + for param_name, weight_name, shard_id in stacked_params_mapping:

1199 + if weight_name not in name:

1200 + continue

1201 + name = name.replace(weight_name, param_name)

1202 +

1203 + param = params_dict[name]

1204 + weight_loader = param.weight_loader

1205 + weight_loader(param, loaded_weight, shard_id)

1206 + break

1207 + else:

1208 + param = params_dict[name]

1209 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

1210 + weight_loader(param, loaded_weight)

1211 + loaded_params.add(name)

1212 + return loaded_params

1213 +

1214 +

1215 +class IsaacVisionEmbedding(nn.Module):

1216 + def __init__(

1217 + self,

1218 + vision_cfg: PixelShuffleSiglip2VisionConfig,

1219 + hidden_dim: int,

1220 + output_dim: int,

1221 + quant_config: QuantizationConfig | None = None,

1222 + multimodal_config: MultiModalConfig | None = None,

1223 + prefix: str = "",

1224 + ):

1225 + super().__init__()

1226 + self.transformer = Siglip2VisionTransformer(

1227 + vision_cfg,

1228 + quant_config=quant_config,

1229 + prefix=maybe_prefix(prefix, "0"),

1230 + multimodal_config=multimodal_config,

1231 + )

1232 + self.linear_fc1 = ColumnParallelLinear(

1233 + hidden_dim,

1234 + 4 * hidden_dim,

1235 + bias=False,

1236 + quant_config=quant_config,

1237 + prefix=maybe_prefix(prefix, "1"),

1238 + return_bias=False,

1239 + )

1240 + self.act = nn.SiLU()

1241 + self.linear_fc2 = RowParallelLinear(

1242 + 4 * hidden_dim,

1243 + output_dim,

1244 + bias=False,

1245 + quant_config=quant_config,

1246 + prefix=maybe_prefix(prefix, "3"),

1247 + return_bias=False,

1248 + )

1249 +

1250 + def forward(

1251 + self, packed_seq_patches: tuple[torch.Tensor, torch.Tensor]

1252 + ) -> torch.Tensor:

1253 + hidden_states = self.transformer(packed_seq_patches)

1254 + hidden_states = self.linear_fc1(hidden_states)

1255 + hidden_states = self.act(hidden_states)

1256 + hidden_states = self.linear_fc2(hidden_states)

1257 + return hidden_states

1258 +

1259 +

1260 +@MULTIMODAL_REGISTRY.register_processor(

1261 + IsaacMultiModalProcessor,

1262 + info=IsaacProcessingInfo,

1263 + dummy_inputs=IsaacDummyInputsBuilder,

1264 +)

1265 +class IsaacForConditionalGeneration(

1266 + nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE

1267 +):

1268 + packed_modules_mapping = {

1269 + "qkv_proj": [

1270 + "q_proj",

1271 + "k_proj",

1272 + "v_proj",

1273 + ],

1274 + "gate_up_proj": [

1275 + "gate_proj",

1276 + "up_proj",

1277 + ],

1278 + }

1279 +

1280 + supports_encoder_tp_data = True

1281 +

1282 + # To ensure correct weight loading and mapping.

1283 + hf_to_vllm_mapper = WeightsMapper(

1284 + orig_to_new_prefix={

1285 + "lm_head.": "language_model.lm_head.",

1286 + "model.vision_embedding.0": "vision_embedding.transformer",

1287 + "model.vision_embedding.1": "vision_embedding.linear_fc1",

1288 + "model.vision_embedding.2": "vision_embedding.act",

1289 + "model.vision_embedding.3": "vision_embedding.linear_fc2",

1290 + "model.vision_embedding.": "vision_embedding.",

1291 + "model.": "language_model.model.",

1292 + }

1293 + )

1294 +

1295 + @classmethod

1296 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1297 + if modality.startswith("image"):

1298 + return "<image>"

1299 +

1300 + raise ValueError("Only image modality is supported")

1301 +

1302 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):

1303 + super().__init__()

1304 + config: IsaacConfig = vllm_config.model_config.hf_config

1305 + quant_config = vllm_config.quant_config

1306 + self.config = config

1307 + self.multimodal_config = vllm_config.model_config.multimodal_config

1308 +

1309 + head_dim = config.head_dim

1310 + calculated_mrope_section = [

1311 + head_dim // 4, # 2x more for temporal dim

1312 + head_dim // 8,

1313 + head_dim // 8,

1314 + ]

1315 +

1316 + self.vision_token_id = _resolve_vision_token_id(

1317 + vllm_config.model_config, config.vision_token

1318 + )

1319 + config.image_token_id = self.vision_token_id

1320 +

1321 + config.rope_scaling["mrope_section"] = calculated_mrope_section

1322 + self.language_model = init_vllm_registered_model(

1323 + vllm_config=vllm_config,

1324 + architectures=["Qwen3ForCausalLM"],

1325 + prefix=maybe_prefix(prefix, "language_model"),

1326 + )

1327 + self.make_empty_intermediate_tensors = (

1328 + self.language_model.make_empty_intermediate_tensors

1329 + )

1330 +

1331 + vision_cfg = config.vision_config

1332 + if vision_cfg is None:

1333 + raise ValueError("IsaacConfig should always have vision_config")

1334 + attn_impl = (

1335 + config.vision_attn_implementation

1336 + if config.vision_attn_implementation is not None

1337 + else getattr(config, "_attn_implementation", None)

1338 + )

1339 + if attn_impl is not None:

1340 + vision_cfg._attn_implementation = attn_impl

1341 +

1342 + hidden_dim = vision_cfg.hidden_size * (vision_cfg.pixel_shuffle_scale_factor**2)

1343 + self.vision_embedding = IsaacVisionEmbedding(

1344 + vision_cfg=vision_cfg,

1345 + hidden_dim=hidden_dim,

1346 + output_dim=config.hidden_size,

1347 + quant_config=quant_config,

1348 + multimodal_config=self.multimodal_config,

1349 + prefix=maybe_prefix(prefix, "vision_embedding"),

1350 + )

1351 +

1352 + def iter_mm_grid_hw(

1353 + self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec]

1354 + ) -> Iterator[tuple[int, int, int]]:

1355 + spatial_merge_size = self.config.vision_config.pixel_shuffle_scale_factor

1356 + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):

1357 + offset = mm_feature.mm_position.offset

1358 + if mm_feature.modality == "image":

1359 + t, h, w = mm_feature.data["image_grid_thw"].data.tolist()

1360 + assert t == 1, f"Image must have 1 frame, got {t}"

1361 + yield offset, h // spatial_merge_size, w // spatial_merge_size

1362 + else:

1363 + raise ValueError(f"Unsupported modality: {mm_feature.modality}")

1364 +

1365 + def get_mrope_input_positions(

1366 + self,

1367 + input_tokens: list[int],

1368 + mm_features: list[MultiModalFeatureSpec],

1369 + ) -> tuple[torch.Tensor, int]:

1370 + llm_pos_ids_list = []

1371 + st = 0

1372 + for offset, llm_grid_h, llm_grid_w in self.iter_mm_grid_hw(

1373 + input_tokens, mm_features

1374 + ):

1375 + text_len = offset - st

1376 + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0

1377 + llm_pos_ids_list.append(

1378 + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx

1379 + )

1380 +

1381 + grid_indices = np.indices((1, llm_grid_h, llm_grid_w)).reshape(3, -1)

1382 + grid_indices[0, :] = grid_indices[0, :] + text_len + st_idx

1383 + llm_pos_ids_list.append(grid_indices)

1384 + st = offset + llm_grid_h * llm_grid_w

1385 +

1386 + if st < len(input_tokens):

1387 + st_idx = llm_pos_ids_list[-1][0, -1] + 1 if len(llm_pos_ids_list) > 0 else 0

1388 + text_len = len(input_tokens) - st

1389 + llm_pos_ids_list.append(

1390 + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx

1391 + )

1392 +

1393 + llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)

1394 + mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()

1395 +

1396 + return torch.from_numpy(llm_positions), mrope_position_delta

1397 +

1398 + def _parse_and_validate_image_input(

1399 + self, **kwargs: object

1400 + ) -> IsaacImagePixelInputs | None:

1401 + pixel_values = kwargs.get("pixel_values")

1402 + image_grid_thw = kwargs.get("image_grid_thw")

1403 + if pixel_values is None or image_grid_thw is None:

1404 + return None

1405 +

1406 + # TensorSchema will automatically validate shapes on initialization

1407 + return IsaacImagePixelInputs(

1408 + pixel_values=pixel_values,

1409 + image_grid_thw=image_grid_thw,

1410 + )

1411 +

1412 + def _process_image_input(

1413 + self,

1414 + image_input: IsaacImagePixelInputs,

1415 + ) -> tuple[torch.Tensor, ...]:

1416 + pixel_values = image_input["pixel_values"]

1417 + image_grid_thw = image_input["image_grid_thw"]

1418 + if pixel_values.numel() == 0:

1419 + return ()

1420 +

1421 + device = next(self.language_model.parameters()).device

1422 + dtype = self.vision_embedding.linear_fc1.weight.dtype

1423 + pixel_values = pixel_values.to(device=device, dtype=dtype)

1424 + spatial_grids = image_grid_thw[:, 1:3].to(device, dtype=torch.int32)

1425 +

1426 + vision_embeddings = self.vision_embedding((pixel_values, spatial_grids))

1427 + merge_size = self.config.vision_config.pixel_shuffle_scale_factor

1428 + sizes = spatial_grids.prod(-1) // (merge_size * merge_size)

1429 + return tuple(vision_embeddings.split(sizes.tolist()))

1430 +

1431 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:

1432 + image_input = self._parse_and_validate_image_input(**kwargs)

1433 + if image_input is None:

1434 + return ()

1435 + return self._process_image_input(image_input)

1436 +

1437 + def get_multimodal_embeddings(

1438 + self, **kwargs: object

1439 + ) -> MultiModalEmbeddings | None:

1440 + # Backward compatibility for older runners.

1441 + embeddings = self.embed_multimodal(**kwargs)

1442 + if not embeddings:

1443 + return []

1444 + return embeddings

1445 +

1446 + def get_language_model(self) -> torch.nn.Module:

1447 + return self.language_model

1448 +

1449 + def forward(

1450 + self,

1451 + input_ids: torch.Tensor,

1452 + positions: torch.Tensor,

1453 + intermediate_tensors: IntermediateTensors | None = None,

1454 + inputs_embeds: torch.Tensor | None = None,

1455 + **kwargs: object,

1456 + ) -> torch.Tensor | IntermediateTensors:

1457 + return self.language_model(

1458 + input_ids=input_ids,

1459 + positions=positions,

1460 + intermediate_tensors=intermediate_tensors,

1461 + inputs_embeds=inputs_embeds,

1462 + **kwargs,

1463 + )

1464 +

1465 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

1466 + return self.language_model.compute_logits(hidden_states)

1467 +

1468 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1469 + loader = AutoWeightsLoader(self)

1470 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1471 +

1472 + def get_mm_mapping(self) -> MultiModelKeys:

1473 + """

1474 + Get the module prefix in multimodal models

1475 + """

1476 + return MultiModelKeys.from_string_field(

1477 + language_model="language_model",

1478 + connector="vision_embedding.linear_fc2", # The final linear layer

1479 + tower_model="vision_embedding",

1480 + )

vllm/model_executor/models/registry.py

@@ -333,6 +333,7 @@ _MULTIMODAL_MODELS = {

333 "idefics3",

334 "Idefics3ForConditionalGeneration",

335 ),

336 + "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"),

337 "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), # noqa: E501

338 "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),

339 "KeyeVL1_5ForConditionalGeneration": (

vllm/transformers_utils/config.py

@@ -81,6 +81,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

81 deepseek_v32="DeepseekV3Config",

82 flex_olmo="FlexOlmoConfig",

83 hunyuan_vl="HunYuanVLConfig",

84 + isaac="IsaacConfig",

85 kimi_linear="KimiLinearConfig",

86 kimi_vl="KimiVLConfig",

87 RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)

vllm/transformers_utils/configs/__init__.py

@@ -25,6 +25,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

25 "HunYuanVLConfig": "vllm.transformers_utils.configs.hunyuan_vl",

26 "HunYuanVLTextConfig": "vllm.transformers_utils.configs.hunyuan_vl",

27 "HunYuanVLVisionConfig": "vllm.transformers_utils.configs.hunyuan_vl",

28 + "IsaacConfig": "vllm.transformers_utils.configs.isaac",

29 # RWConfig is for the original tiiuae/falcon-40b(-instruct) and

30 # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the

31 # `FalconConfig` class from the official HuggingFace transformers library.

@@ -41,6 +42,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

42 "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h",

43 "Olmo3Config": "vllm.transformers_utils.configs.olmo3",

44 "OvisConfig": "vllm.transformers_utils.configs.ovis",

45 + "PixelShuffleSiglip2VisionConfig": "vllm.transformers_utils.configs.isaac",

46 "RadioConfig": "vllm.transformers_utils.configs.radio",

47 "SpeculatorsConfig": "vllm.transformers_utils.configs.speculators.base",

48 "UltravoxConfig": "vllm.transformers_utils.configs.ultravox",

@@ -65,6 +67,7 @@ __all__ = [

67 "HunYuanVLConfig",

68 "HunYuanVLTextConfig",

69 "HunYuanVLVisionConfig",

70 + "IsaacConfig",

71 "RWConfig",

72 "JAISConfig",

73 "Lfm2MoeConfig",

@@ -78,6 +81,7 @@ __all__ = [

81 "NemotronHConfig",

82 "Olmo3Config",

83 "OvisConfig",

84 + "PixelShuffleSiglip2VisionConfig",

85 "RadioConfig",

86 "SpeculatorsConfig",

87 "UltravoxConfig",

vllm/transformers_utils/configs/isaac.py new

+86

@@ -0,0 +1,86 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +from __future__ import annotations

4 +

5 +from transformers import Qwen3Config

6 +from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig

7 +

8 +

9 +class PixelShuffleSiglip2VisionConfig(Siglip2VisionConfig):

10 + """Vision configuration for Isaac with Pixel Shuffle support.

11 +

12 + Extends Siglip2VisionConfig with additional fields for pixel shuffle.

13 + """

14 +

15 + model_type = "pixel_shuffle_siglip2"

16 + base_config_key = "vision_config"

17 +

18 + def __init__(

19 + self,

20 + pixel_shuffle_scale_factor: int = 1,

21 + num_patches: int = 256,

22 + **kwargs,

23 + ):

24 + super().__init__(**kwargs)

25 +

26 + # Add our custom fields

27 + self.pixel_shuffle_scale_factor = pixel_shuffle_scale_factor

28 + self.num_patches = num_patches

29 +

30 +

31 +class IsaacConfig(Qwen3Config):

32 + """Configuration class for Isaac multimodal model."""

33 +

34 + model_type = "isaac"

35 + sub_configs = {"vision_config": PixelShuffleSiglip2VisionConfig}

36 +

37 + def __init__(

38 + self,

39 + vision_config=None,

40 + vision_patch_size: int = 16,

41 + vision_max_num_patches: int = 256,

42 + vision_min_num_patches: int | None = None,

43 + pixel_shuffle_scale: int = 1,

44 + max_sequence_length: int = 16384,

45 + vision_token: str = "<image>",

46 + vision_attn_implementation: str | None = None,

47 + **kwargs,

48 + ):

49 + super().__init__(**kwargs)

50 +

51 + # EventStreamProcessor parameters (for backward compatibility)

52 + self.video_patch_size = vision_patch_size

53 + self.vision_max_num_patches = vision_max_num_patches

54 + self.vision_min_num_patches = vision_min_num_patches

55 + self.pixel_shuffle_scale = pixel_shuffle_scale

56 +

57 + # Processing parameters

58 + self.max_sequence_length = max_sequence_length

59 + self.vision_token = vision_token

60 +

61 + # Handle vision config - PixelShuffleSiglip2VisionConfig instance

62 + if isinstance(vision_config, dict):

63 + self.vision_config = PixelShuffleSiglip2VisionConfig(**vision_config)

64 + elif vision_config is None:

65 + self.vision_config = PixelShuffleSiglip2VisionConfig()

66 + else:

67 + self.vision_config = vision_config

68 +

69 + # Ensure compatibility with pretrained checkpoints

70 + self.vision_config.pixel_shuffle_scale_factor = getattr(

71 + self.vision_config,

72 + "pixel_shuffle_scale_factor",

73 + pixel_shuffle_scale,

74 + )

75 + self.vision_config.num_patches = getattr(

76 + self.vision_config,

77 + "num_patches",

78 + vision_max_num_patches,

79 + )

80 + self.vision_attn_implementation = vision_attn_implementation

81 +

82 +

83 +__all__ = [

84 + "IsaacConfig",

85 + "PixelShuffleSiglip2VisionConfig",

86 +]