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 9, 03:45 UTC
[Models]: lfm2_siglip2 return intermediate encoder layers (#33370)
Signed-off-by: Eduardo Salinas <edus@microsoft.com>
Eduardo Salinas committed
Feb 1, 2026 at 01:17 UTC
302ecf64ff22e6e2f59d108e95f0c66d96c9ffc0
1 file changed
+76
-7
vllm/model_executor/models/lfm2_siglip2.py
+76
-7
@@ -22,7 +22,11 @@ from vllm.model_executor.layers.linear import (
22
from vllm.model_executor.layers.quantization import QuantizationConfig
23
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24
25
-from .vision import is_vit_use_data_parallel, should_torch_compile_mm_vit
25
+from .vision import (
26
+ is_vit_use_data_parallel,
27
+ resolve_visual_encoder_outputs,
28
+ should_torch_compile_mm_vit,
29
+)
30
31
32
class Siglip2VisionEmbeddings(nn.Module):
@@ -331,10 +335,17 @@ class Siglip2Encoder(nn.Module):
335
self,
336
config: Siglip2VisionConfig,
337
quant_config: QuantizationConfig | None = None,
338
+ num_hidden_layers_override: int | None = None,
339
prefix: str = "",
340
):
341
super().__init__()
342
self.config = config
343
+
344
+ if num_hidden_layers_override is None:
345
+ num_hidden_layers = config.num_hidden_layers
346
+ else:
347
+ num_hidden_layers = num_hidden_layers_override
348
+
349
self.layers = nn.ModuleList(
350
[
351
Siglip2EncoderLayer(
@@ -342,7 +353,7 @@ class Siglip2Encoder(nn.Module):
353
quant_config=quant_config,
354
prefix=f"{prefix}.layers.{idx}",
355
)
345
- for idx in range(config.num_hidden_layers)
356
+ for idx in range(num_hidden_layers)
357
]
358
)
359
@@ -351,15 +362,21 @@ class Siglip2Encoder(nn.Module):
362
inputs_embeds: torch.Tensor,
363
cu_seqlens: torch.Tensor,
364
max_seqlen: int | torch.Tensor,
354
- ) -> torch.Tensor:
365
+ return_all_hidden_states: bool = False,
366
+ ) -> torch.Tensor | list[torch.Tensor]:
367
+ hidden_states_pool = [inputs_embeds]
368
hidden_states = inputs_embeds
369
+
370
for encoder_layer in self.layers:
357
- layer_outputs = encoder_layer(
371
+ hidden_states = encoder_layer(
372
hidden_states,
373
cu_seqlens=cu_seqlens,
374
max_seqlen=max_seqlen,
375
)
362
- hidden_states = layer_outputs
376
+ if return_all_hidden_states:
377
+ hidden_states_pool.append(hidden_states)
378
+ if return_all_hidden_states:
379
+ return hidden_states_pool
380
return hidden_states
381
382
@@ -368,6 +385,8 @@ class Siglip2VisionTransformer(nn.Module):
385
self,
386
config: Siglip2VisionConfig,
387
quant_config: QuantizationConfig | None = None,
388
+ num_hidden_layers_override: int | None = None,
389
+ require_post_norm: bool | None = None,
390
prefix: str = "",
391
):
392
super().__init__()
@@ -381,6 +400,7 @@ class Siglip2VisionTransformer(nn.Module):
400
self.encoder = Siglip2Encoder(
401
config,
402
quant_config=quant_config,
403
+ num_hidden_layers_override=num_hidden_layers_override,
404
prefix=f"{prefix}.encoder",
405
)
406
num_hidden_layers = config.num_hidden_layers
@@ -390,7 +410,13 @@ class Siglip2VisionTransformer(nn.Module):
410
f"layers, but you requested {len(self.encoder.layers)} layers."
411
)
412
393
- self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
413
+ if require_post_norm is None:
414
+ require_post_norm = len(self.encoder.layers) == num_hidden_layers
415
+
416
+ if require_post_norm:
417
+ self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
418
+ else:
419
+ self.post_layernorm = None
420
421
def get_input_embeddings(self):
422
return self.embeddings
@@ -401,19 +427,34 @@ class Siglip2VisionTransformer(nn.Module):
427
spatial_shapes: torch.LongTensor,
428
cu_seqlens: torch.Tensor,
429
max_seqlen: torch.Tensor,
430
+ select_layers: list[int] | None = None,
431
) -> torch.Tensor:
432
r"""
433
spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
434
Tensor containing the spatial dimensions (height, width)
435
of the input images.
436
+ select_layers (`list[int]` or `None`, defaults to `None`):
437
+ Layer indices to select hidden states from. Supports negative
438
+ indices (e.g., -1 for last layer, -2 for second-to-last).
439
+ If None, returns the last layer output.
440
"""
441
hidden_states = self.embeddings(pixel_values_packed, spatial_shapes)
442
+
443
encoder_outputs = self.encoder(
444
inputs_embeds=hidden_states,
445
cu_seqlens=cu_seqlens,
446
max_seqlen=max_seqlen,
447
+ return_all_hidden_states=select_layers is not None,
448
)
416
- return self.post_layernorm(encoder_outputs)
449
+
450
+ encoder_outputs = resolve_visual_encoder_outputs(
451
+ encoder_outputs,
452
+ self.post_layernorm,
453
+ select_layers=select_layers,
454
+ max_possible_layers=self.config.num_hidden_layers,
455
+ )
456
+
457
+ return encoder_outputs
458
459
460
class Siglip2Model(torch.nn.Module):
@@ -421,6 +462,8 @@ class Siglip2Model(torch.nn.Module):
462
self,
463
config: Siglip2VisionConfig,
464
quant_config: QuantizationConfig | None = None,
465
+ num_hidden_layers_override: int | None = None,
466
+ require_post_norm: bool | None = None,
467
prefix: str = "",
468
):
469
super().__init__()
@@ -428,6 +471,8 @@ class Siglip2Model(torch.nn.Module):
471
self.vision_model = Siglip2VisionTransformer(
472
config,
473
quant_config=quant_config,
474
+ num_hidden_layers_override=num_hidden_layers_override,
475
+ require_post_norm=require_post_norm,
476
prefix=f"{prefix}.vision_model",
477
)
478
@@ -437,12 +482,22 @@ class Siglip2Model(torch.nn.Module):
482
spatial_shapes: torch.LongTensor,
483
cu_seqlens: torch.Tensor,
484
max_seqlen: torch.Tensor,
485
+ select_layers: list[int] | None = None,
486
) -> torch.Tensor:
487
+ """Forward pass through the vision model.
488
+
489
+ Args:
490
+ select_layers: Layer indices to select hidden states from.
491
+ Supports negative indices (e.g., [-2] for second-to-last).
492
+ If None, returns the last layer output with post_layernorm.
493
+ Multiple layers can be selected and will be concatenated.
494
+ """
495
return self.vision_model(
496
pixel_values_packed=pixel_values_packed,
497
spatial_shapes=spatial_shapes,
498
cu_seqlens=cu_seqlens,
499
max_seqlen=max_seqlen,
500
+ select_layers=select_layers,
501
)
502
503
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -454,8 +509,22 @@ class Siglip2Model(torch.nn.Module):
509
]
510
params_dict = dict(self.named_parameters())
511
loaded_params: set[str] = set()
512
+ layer_count = len(self.vision_model.encoder.layers)
513
514
for name, loaded_weight in weights:
515
+ # post_layernorm is optional in Siglip2Model
516
+ if (
517
+ name.startswith("vision_model.post_layernorm")
518
+ and self.vision_model.post_layernorm is None
519
+ ):
520
+ continue
521
+
522
+ # omit layers when num_hidden_layers_override is set
523
+ if name.startswith("vision_model.encoder.layers"):
524
+ layer_idx = int(name.split(".")[3])
525
+ if layer_idx >= layer_count:
526
+ continue
527
+
528
for param_name, weight_name, shard_id in stacked_params_mapping:
529
if weight_name not in name:
530
continue