model: support JetVLM (#13289)

python/sglang/srt/configs/__init__.py

@@ -7,6 +7,7 @@ from sglang.srt.configs.exaone import ExaoneConfig
 from sglang.srt.configs.falcon_h1 import FalconH1Config
 from sglang.srt.configs.janus_pro import MultiModalityConfig
 from sglang.srt.configs.jet_nemotron import JetNemotronConfig
+from sglang.srt.configs.jet_vlm import JetVLMConfig
 from sglang.srt.configs.kimi_linear import KimiLinearConfig
 from sglang.srt.configs.kimi_vl import KimiVLConfig
 from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
@@ -40,4 +41,5 @@ __all__ = [
     "FalconH1Config",
     "NemotronHConfig",
     "JetNemotronConfig",
+    "JetVLMConfig",
 ]

python/sglang/srt/configs/jet_vlm.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.siglip import SiglipVisionConfig
+
+from sglang.srt.configs.jet_nemotron import JetNemotronConfig
+from sglang.srt.configs.mamba_utils import Mamba2CacheParams
+
+
+class JetVLMConfig(PretrainedConfig):
+    model_type = "jet_vlm"
+    sub_configs = {
+        "text_config": JetNemotronConfig,
+        "vision_config": SiglipVisionConfig,
+    }
+    _auto_class = "AutoConfig"
+
+    def __init__(
+        self,
+        *,
+        text_config: dict[str, Any] | None = None,
+        vision_config: dict[str, Any] | None = None,
+        image_token_id: int | None = None,
+        video_token_id: int | None = None,
+        **kwargs,
+    ):
+        self.text_config = (
+            JetNemotronConfig(**text_config)
+            if text_config is not None
+            else JetNemotronConfig()
+        )
+        self.vision_config = (
+            SiglipVisionConfig(**vision_config)
+            if vision_config is not None
+            else SiglipVisionConfig()
+        )
+
+        self.image_token_id = image_token_id if image_token_id is not None else -1
+        self.video_token_id = video_token_id if video_token_id is not None else -1
+
+        super().__init__(**kwargs)
+
+    @property
+    def full_attention_layer_ids(self) -> list[int]:
+        return self.text_config.full_attention_layer_ids
+
+    @property
+    def linear_layer_ids(self) -> list[int]:
+        return self.text_config.linear_layer_ids
+
+    @property
+    def mamba2_cache_params(self) -> Mamba2CacheParams:
+        return self.text_config.mamba2_cache_params

python/sglang/srt/configs/model_config.py

@@ -955,6 +955,7 @@ multimodal_model_archs = [
     "NVILAForConditionalGeneration",
     "NVILALiteForConditionalGeneration",
     "DeepseekOCRForCausalLM",
+    "JetVLMForConditionalGeneration",
 ]
 
 

python/sglang/srt/model_executor/model_runner.py

@@ -32,6 +32,7 @@ import torch.distributed as dist
 from sglang.srt.configs import (
     FalconH1Config,
     JetNemotronConfig,
+    JetVLMConfig,
     KimiLinearConfig,
     NemotronHConfig,
     Qwen3NextConfig,
@@ -1416,7 +1417,7 @@ class ModelRunner:
     @property
     def hybrid_gdn_config(self):
         config = self.model_config.hf_config
-        if isinstance(config, Qwen3NextConfig | JetNemotronConfig):
+        if isinstance(config, Qwen3NextConfig | JetNemotronConfig | JetVLMConfig):
             return config
         return None
 

python/sglang/srt/models/jet_nemotron.py

@@ -545,6 +545,9 @@ class JetNemotronForCausalLM(nn.Module):
         else:
             return self.pooler(hidden_states, forward_batch)
 
+    def get_input_embeddings(self) -> nn.Module:
+        return self.model.embed_tokens
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         stacked_params_mapping: list[tuple[str, str, str | int]] = [
             # (param_name, shard_weight_name, shard_id)

python/sglang/srt/models/jet_vlm.py

@@ -0,0 +1,143 @@
+import math
+from collections.abc import Iterable
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.siglip import SiglipVisionModel
+
+import sglang.srt.managers.mm_utils as mm_utils
+import sglang.srt.model_loader.weight_utils as weight_utils
+import sglang.srt.utils as utils
+from sglang.srt.configs.jet_vlm import JetVLMConfig
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
+from sglang.srt.managers.schedule_batch import (
+    Modality,
+    MultimodalDataItem,
+    MultimodalInputs,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.models.jet_nemotron import JetNemotronForCausalLM
+
+MM_HIDDEN_SIZE = 1152
+
+
+class JetVLMDownSample2x2BlockFix(nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        _, seq_len, _ = x.shape
+
+        feat_size = math.isqrt(seq_len)
+
+        features = einops.rearrange(x, "b (h w) d -> b h w d", h=feat_size, w=feat_size)
+
+        if feat_size % 2 == 1:
+            features = F.pad(features, (0, 0, 0, 1, 0, 1))
+
+        features = einops.rearrange(
+            features, "b (h p1) (w p2) d -> b (h w) (p1 p2 d)", p1=2, p2=2
+        )
+
+        return features
+
+
+class JetVLMMultiModalProjector(nn.Module):
+    def __init__(self, config: JetVLMConfig) -> None:
+        super().__init__()
+
+        self.layers = nn.Sequential(
+            JetVLMDownSample2x2BlockFix(),
+            nn.LayerNorm(MM_HIDDEN_SIZE * 4),
+            nn.Linear(MM_HIDDEN_SIZE * 4, config.text_config.hidden_size),
+            nn.GELU(),
+            nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.layers(x)
+
+
+class JetVLMForConditionalGeneration(nn.Module):
+    def __init__(
+        self,
+        config: JetVLMConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+
+        self.vision_tower = SiglipVisionModel(config.vision_config)
+        self.mm_projector = JetVLMMultiModalProjector(config)
+        self.llm = JetNemotronForCausalLM(
+            config=config.text_config,
+            quant_config=quant_config,
+            prefix=utils.add_prefix("llm", prefix),
+        )
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        positions: Tensor,
+        forward_batch: ForwardBatch,
+        get_embedding: bool = False,
+    ) -> LogitsProcessorOutput:
+        output = mm_utils.general_mm_embed_routine(
+            input_ids=input_ids,
+            forward_batch=forward_batch,
+            language_model=self.llm,
+            data_embedding_funcs={
+                Modality.IMAGE: self.get_image_feature,
+                Modality.VIDEO: self.get_image_feature,
+            },
+            get_embedding=get_embedding,
+            positions=positions,
+        )
+
+        assert isinstance(output, LogitsProcessorOutput)
+
+        return output
+
+    def get_image_feature(self, mm_input: list[MultimodalDataItem]) -> Tensor:
+        pixel_values = torch.cat([torch.tensor(x.feature) for x in mm_input], dim=0)
+
+        vision_tower_output: BaseModelOutputWithPooling = self.vision_tower(
+            pixel_values,
+            output_hidden_states=True,
+        )
+        assert vision_tower_output.hidden_states is not None
+
+        vision_features = vision_tower_output.hidden_states[-2]
+
+        vision_features = self.mm_projector(vision_features)
+
+        vision_features = einops.rearrange(vision_features, "n p d -> (n p) d")
+
+        return vision_features
+
+    def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> None:
+        params_dict = dict(self.named_parameters())
+
+        for name, loaded_weight in weights:
+            if name.startswith("llm."):
+                self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(
+                    param, "weight_loader", weight_utils.default_weight_loader
+                )
+                weight_loader(param, loaded_weight)
+
+    def pad_input_ids(
+        self, input_ids: list[int], mm_inputs: MultimodalInputs
+    ) -> list[int]:
+        pattern = MultiModalityDataPaddingPatternMultimodalTokens()
+        return pattern.pad_input_tokens(input_ids, mm_inputs)
+
+
+EntryClass = [JetVLMForConditionalGeneration]

python/sglang/srt/multimodal/processors/nvila.py

@@ -6,6 +6,7 @@ from transformers.processing_utils import ProcessorMixin
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from sglang.srt.managers.io_struct import GenerateReqInput
+from sglang.srt.models.jet_vlm import JetVLMForConditionalGeneration
 from sglang.srt.models.nvila import NVILAForConditionalGeneration
 from sglang.srt.models.nvila_lite import NVILALiteForConditionalGeneration
 from sglang.srt.multimodal.processors.base_processor import (
@@ -21,6 +22,7 @@ class NVILAMultimodalProcessor(BaseMultimodalProcessor):
     models: list[type[nn.Module]] = [
         NVILAForConditionalGeneration,
         NVILALiteForConditionalGeneration,
+        JetVLMForConditionalGeneration,
     ]
 
     def __init__(

python/sglang/srt/utils/hf_transformers_utils.py

@@ -45,6 +45,7 @@ from sglang.srt.configs import (
     ExaoneConfig,
     FalconH1Config,
     JetNemotronConfig,
+    JetVLMConfig,
     KimiLinearConfig,
     KimiVLConfig,
     LongcatFlashConfig,
@@ -79,6 +80,7 @@ _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
     NemotronHConfig,
     DeepseekVLV2Config,
     JetNemotronConfig,
+    JetVLMConfig,
 ]
 
 _CONFIG_REGISTRY = {