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

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, 01:30 UTC

[Model] Molmo2: Enable quantized weight mapping for vision backbone (#32385)

Signed-off-by: kimheesu <wlskaka4@gmail.com>

Kim Hee Su committed Jan 17, 2026 at 18:33 UTC 1646fea672c939fa1204ae7ad0b7e87c072b3615

1 file changed +23 -8

vllm/model_executor/models/molmo2.py

+23 -8

@@ -67,13 +67,13 @@ from vllm.multimodal.parse import (

67 MultiModalDataParser,

68 )

69 from vllm.multimodal.processing import (

70 - BaseDummyInputsBuilder,

70 BaseMultiModalProcessor,

71 BaseProcessingInfo,

72 PromptReplacement,

73 PromptUpdate,

74 PromptUpdateDetails,

75 )

76 +from vllm.multimodal.processing.dummy_inputs import BaseDummyInputsBuilder

77 from vllm.sequence import IntermediateTensors

78 from vllm.utils.math_utils import round_down

79 from vllm.utils.tensor_schema import TensorSchema, TensorShape

@@ -295,6 +295,7 @@ class ViTMLP(nn.Module):

295 hidden_dim: int,

296 hidden_act: str,

297 quant_config: QuantizationConfig | None = None,

298 + prefix: str = "",

299 ) -> None:

300 super().__init__()

301 self.w1 = ColumnParallelLinear(

@@ -302,6 +303,7 @@ class ViTMLP(nn.Module):

303 hidden_dim,

304 bias=True,

305 quant_config=quant_config,

306 + prefix=f"{prefix}.w1",

307 )

308 # Activation function.

309 self.act = get_act_fn(hidden_act)

@@ -310,6 +312,7 @@ class ViTMLP(nn.Module):

312 dim,

313 bias=True,

314 quant_config=quant_config,

315 + prefix=f"{prefix}.w2",

316 )

317

318 def forward(self, x: torch.Tensor) -> torch.Tensor:

@@ -364,12 +367,14 @@ class ViTMultiHeadDotProductAttention(nn.Module):

367 self.total_num_kv_heads,

368 bias=use_bias,

369 quant_config=quant_config,

370 + prefix=f"{prefix}.merged_qkv",

371 )

372 self.wo = RowParallelLinear(

373 self.total_num_heads * self.head_dim,

374 self.hidden_size,

375 bias=use_bias,

376 quant_config=quant_config,

377 + prefix=f"{prefix}.wo",

378 )

379 self.scale = self.head_dim**-0.5

380 self.attn = MMEncoderAttention(

@@ -414,6 +419,7 @@ class Molmo2VisionBlock(nn.Module):

419 hidden_dim=config.intermediate_size,

420 hidden_act=config.hidden_act,

421 quant_config=quant_config,

422 + prefix=f"{prefix}.feed_forward",

423 )

424 self.attention_norm = nn.LayerNorm(

425 config.hidden_size,

@@ -548,6 +554,7 @@ class ImagePoolingAttention(nn.Module):

554 use_bias: bool = True,

555 use_pytorch_sdpa: bool = False,

556 quant_config: QuantizationConfig | None = None,

557 + prefix: str = "",

558 ) -> None:

559 super().__init__()

560

@@ -579,18 +586,21 @@ class ImagePoolingAttention(nn.Module):

586 self.total_num_heads * self.head_dim,

587 bias=use_bias,

588 quant_config=quant_config,

589 + prefix=f"{prefix}.q_proj",

590 )

591 self.merged_kv = MergedColumnParallelLinear(

592 self.input_dim,

593 [self.total_num_kv_heads * self.head_dim] * 2,

594 bias=use_bias,

595 quant_config=quant_config,

596 + prefix=f"{prefix}.merged_kv",

597 )

598 self.o_proj = RowParallelLinear(

599 self.total_num_heads * self.head_dim,

600 self.hidden_size,

601 bias=use_bias,

602 quant_config=quant_config,

603 + prefix=f"{prefix}.o_proj",

604 )

605 self.scale = self.head_dim**-0.5

606 self.use_pytorch_sdpa = use_pytorch_sdpa

@@ -672,6 +682,7 @@ class ImageProjectorMLP(nn.Module):

682 output_dim: int,

683 hidden_act: str,

684 quant_config: QuantizationConfig | None = None,

685 + prefix: str = "",

686 ) -> None:

687 super().__init__()

688

@@ -680,6 +691,7 @@ class ImageProjectorMLP(nn.Module):

691 [hidden_dim] * 2,

692 bias=False,

693 quant_config=quant_config,

694 + prefix=f"{prefix}.merged_linear",

695 )

696 # Activation function.

697 assert hidden_act == "silu"

@@ -691,6 +703,7 @@ class ImageProjectorMLP(nn.Module):

703 output_dim,

704 bias=False,

705 quant_config=quant_config,

706 + prefix=f"{prefix}.down_proj",

707 )

708

709 def forward(self, x: torch.Tensor) -> torch.Tensor:

@@ -745,6 +758,7 @@ class Molmo2VisionBackbone(nn.Module, SupportsQuant):

758 head_dim=adapter_config.head_dim,

759 use_pytorch_sdpa=adapter_config.pooling_attention_mask,

760 quant_config=quant_config,

761 + prefix=f"{prefix}.image_pooling_2d",

762 )

763 self.image_projector = ImageProjectorMLP(

764 input_dim=adapter_config.hidden_size,

@@ -752,6 +766,7 @@ class Molmo2VisionBackbone(nn.Module, SupportsQuant):

766 output_dim=adapter_config.text_hidden_size,

767 hidden_act=adapter_config.hidden_act,

768 quant_config=quant_config,

769 + prefix=f"{prefix}.image_projector",

770 )

771

772 @property

@@ -2438,13 +2453,13 @@ class Molmo2ForConditionalGeneration(

2453 hf_to_vllm_mapper = WeightsMapper(

2454 orig_to_new_substr={

2455 # vision backbone mapping

2441 - "image_pooling_2d.wq.": "image_pooling_2d.q_proj.",

2442 - "image_pooling_2d.wk.": "image_pooling_2d.k_proj.",

2443 - "image_pooling_2d.wv.": "image_pooling_2d.v_proj.",

2444 - "image_pooling_2d.wo.": "image_pooling_2d.o_proj.",

2445 - "image_projector.w1.": "image_projector.gate_proj.",

2446 - "image_projector.w3.": "image_projector.up_proj.",

2447 - "image_projector.w2.": "image_projector.down_proj.",

2456 + "image_pooling_2d.wq": "image_pooling_2d.q_proj",

2457 + "image_pooling_2d.wk": "image_pooling_2d.k_proj",

2458 + "image_pooling_2d.wv": "image_pooling_2d.v_proj",

2459 + "image_pooling_2d.wo": "image_pooling_2d.o_proj",

2460 + "image_projector.w1": "image_projector.gate_proj",

2461 + "image_projector.w3": "image_projector.up_proj",

2462 + "image_projector.w2": "image_projector.down_proj",

2463 # language backbone mapping

2464 "att_proj": "qkv_proj",

2465 "attn_out": "o_proj",