diff --git a/python/pyproject.toml b/python/pyproject.toml index 3c753fa2d8..9d2c6fc5e0 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -66,6 +66,9 @@ multimodal = [ "torchvision", ] +[tool.setuptools.package-data] +sgl_jax = ["**/*.yaml"] + [tool.setuptools.packages.find] exclude = [ "assets*", diff --git a/python/sgl_jax/srt/configs/model_config.py b/python/sgl_jax/srt/configs/model_config.py index 1aa9edeb39..88240992fe 100644 --- a/python/sgl_jax/srt/configs/model_config.py +++ b/python/sgl_jax/srt/configs/model_config.py @@ -59,7 +59,6 @@ def __init__( moe_backend: str | MoEBackend = MoEBackend.AUTO, model_sub_dir: str | None = None, ) -> None: - self.model_path = model_path self.model_sub_dir = model_sub_dir self.revision = revision @@ -696,6 +695,7 @@ def is_generation_model(model_architectures: list[str], is_embedding: bool = Fal "Qwen2AudioForConditionalGeneration", "Qwen2VLForConditionalGeneration", "Qwen2_5_VLForConditionalGeneration", + "Qwen3VLForConditionalGeneration", "KimiVLForConditionalGeneration", "InternVLChatModel", "Phi4MMForCausalLM", diff --git a/python/sgl_jax/srt/models/registry.py b/python/sgl_jax/srt/models/registry.py index 2690d7f1c4..57780a7ce5 100644 --- a/python/sgl_jax/srt/models/registry.py +++ b/python/sgl_jax/srt/models/registry.py @@ -84,14 +84,14 @@ def import_model_classes(): entry = module.EntryClass if isinstance(entry, list): # To support multiple model classes in one module for tmp in entry: - assert ( - tmp.__name__ not in model_arch_name_to_cls - ), f"Duplicated model implementation for {tmp.__name__}" + assert tmp.__name__ not in model_arch_name_to_cls, ( + f"Duplicated model implementation for {tmp.__name__}" + ) model_arch_name_to_cls[tmp.__name__] = tmp else: - assert ( - entry.__name__ not in model_arch_name_to_cls - ), f"Duplicated model implementation for {entry.__name__}" + assert entry.__name__ not in model_arch_name_to_cls, ( + f"Duplicated model implementation for {entry.__name__}" + ) model_arch_name_to_cls[entry.__name__] = entry return model_arch_name_to_cls diff --git a/python/sgl_jax/srt/multimodal/configs/qwen_vl/qwen3_vl_config.py b/python/sgl_jax/srt/multimodal/configs/qwen_vl/qwen3_vl_config.py new file mode 100644 index 0000000000..53b435f951 --- /dev/null +++ b/python/sgl_jax/srt/multimodal/configs/qwen_vl/qwen3_vl_config.py @@ -0,0 +1,176 @@ +from dataclasses import dataclass, field + +from sgl_jax.srt.multimodal.configs.multimodal_base_config import MultiModalModelConfigs + +@dataclass +class Qwen3VLVisionConfig: + """Vision encoder configuration for Qwen3-VL.""" + + depth: int = 24 + hidden_size: int = 1024 + intermediate_size: int = 4096 + num_heads: int = 16 + in_channels: int = 3 + patch_size: int = 16 + temporal_patch_size: int = 2 + spatial_merge_size: int = 2 + out_hidden_size: int = 2048 + num_position_embeddings: int = 2304 + deepstack_visual_indexes: tuple = (5, 11, 17) + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-6 + rope_theta: float = 10000.0 + + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_heads + + @classmethod + def qwen3vl_2b(cls): + return cls( + depth=24, + hidden_size=1024, + intermediate_size=4096, + num_heads=16, + out_hidden_size=2048, + deepstack_visual_indexes=(5, 11, 17), + ) + + @classmethod + def qwen3vl_4b(cls): + return cls( + depth=24, + hidden_size=1024, + intermediate_size=4096, + num_heads=16, + out_hidden_size=2560, + deepstack_visual_indexes=(5, 11, 17), + ) + + @classmethod + def qwen3vl_8b(cls): + return cls( + depth=27, + hidden_size=1152, + intermediate_size=4304, + num_heads=16, + out_hidden_size=4096, + deepstack_visual_indexes=(8, 16, 24), + ) + + @classmethod + def qwen3vl_32b(cls): + return cls( + depth=27, + hidden_size=1152, + intermediate_size=4304, + num_heads=16, + out_hidden_size=5120, + deepstack_visual_indexes=(8, 16, 24), + ) + + +@dataclass +class Qwen3VLTextConfig: + """Text decoder configuration for Qwen3-VL.""" + + vocab_size: int = 151936 + hidden_size: int = 2048 + intermediate_size: int = 6144 + num_hidden_layers: int = 28 + num_attention_heads: int = 16 + num_key_value_heads: int = 8 + head_dim: int = 128 + hidden_act: str = "silu" + rms_norm_eps: float = 1e-6 + rope_theta: float = 5_000_000 + mrope_section: tuple = (24, 20, 20) # T, H, W partitions of head_dim + attention_bias: bool = False + tie_word_embeddings: bool = True + + @classmethod + def qwen3vl_2b(cls): + return cls( + hidden_size=2048, + intermediate_size=6144, + num_hidden_layers=28, + num_attention_heads=16, + num_key_value_heads=8, + tie_word_embeddings=True, + ) + + @classmethod + def qwen3vl_4b(cls): + return cls( + hidden_size=2560, + intermediate_size=9728, + num_hidden_layers=36, + num_attention_heads=32, + num_key_value_heads=8, + tie_word_embeddings=True, + ) + + @classmethod + def qwen3vl_8b(cls): + return cls( + hidden_size=4096, + intermediate_size=12288, + num_hidden_layers=36, + num_attention_heads=32, + num_key_value_heads=8, + tie_word_embeddings=False, + ) + + @classmethod + def qwen3vl_32b(cls): + return cls( + hidden_size=5120, + intermediate_size=25600, + num_hidden_layers=64, + num_attention_heads=64, + num_key_value_heads=8, + tie_word_embeddings=False, + ) + + +@dataclass +class Qwen3VLConfig(MultiModalModelConfigs): + """Combined configuration for Qwen3-VL model.""" + + vision_config: Qwen3VLVisionConfig = field(default_factory=Qwen3VLVisionConfig) + text_config: Qwen3VLTextConfig = field(default_factory=Qwen3VLTextConfig) + image_token_id: int = 151655 + video_token_id: int = 151656 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + + @classmethod + def qwen3vl_2b(cls): + """Qwen3-VL 2B configuration.""" + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_2b(), + text_config=Qwen3VLTextConfig.qwen3vl_2b(), + ) + + @classmethod + def qwen3vl_4b(cls): + """Qwen3-VL 4B configuration.""" + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_4b(), + text_config=Qwen3VLTextConfig.qwen3vl_4b(), + ) + + @classmethod + def qwen3vl_8b(cls): + """Qwen3-VL 8B configuration.""" + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_8b(), + text_config=Qwen3VLTextConfig.qwen3vl_8b(), + ) + + @classmethod + def qwen3vl_32b(cls): + return cls( + vision_config=Qwen3VLVisionConfig.qwen3vl_32b(), + text_config=Qwen3VLTextConfig.qwen3vl_32b(), + ) diff --git a/python/sgl_jax/srt/multimodal/entrypoint/http_server.py b/python/sgl_jax/srt/multimodal/entrypoint/http_server.py index 81d8edbaa3..3095508a51 100644 --- a/python/sgl_jax/srt/multimodal/entrypoint/http_server.py +++ b/python/sgl_jax/srt/multimodal/entrypoint/http_server.py @@ -421,10 +421,11 @@ def _execute_multimodal_server_warmup( # Send a warmup request # For Wan models, send an image generation request - if "Qwen2.5-VL" in server_args.model_path: + if "Qwen3-VL" in server_args.model_path or "Qwen2.5-VL" in server_args.model_path: + model_name = "Qwen/Qwen3-VL" if "Qwen3-VL" in server_args.model_path else "Qwen/Qwen2.5-VL" request_endpoint = "/v1/chat/completions" json_data = { - "model": "Qwen/Qwen2.5-VL", + "model": model_name, "messages": [ { "role": "user", diff --git a/python/sgl_jax/srt/multimodal/manager/stage.py b/python/sgl_jax/srt/multimodal/manager/stage.py index 3531208c97..6175c55a2f 100644 --- a/python/sgl_jax/srt/multimodal/manager/stage.py +++ b/python/sgl_jax/srt/multimodal/manager/stage.py @@ -20,6 +20,10 @@ from sgl_jax.srt.multimodal.models.qwen2_5VL.qwen2_5_vl_generation import ( Qwen2_5_VL_Generation, ) +from sgl_jax.srt.multimodal.models.qwen3_VL.qwen3_vl_vit import Qwen3_VL_VisionModel +from sgl_jax.srt.multimodal.models.qwen3_VL.qwen3_vl_generation import ( + Qwen3_VL_Generation, +) from sgl_jax.srt.multimodal.models.wan.diffusion.wan_dit import ( WanDualTransformer3DModel, WanTransformer3DModel, @@ -189,6 +193,10 @@ def get_model_class(name: str): return Qwen2_5_VL_Generation elif name == "Qwen2_5_VL_VisionModel": return Qwen2_5_VL_VisionModel + elif name == "Qwen3_VL_Generation": + return Qwen3_VL_Generation + elif name == "Qwen3_VL_VisionModel": + return Qwen3_VL_VisionModel elif name == "Qwen2ForCausalLM": return Qwen2ForCausalLM else: diff --git a/python/sgl_jax/srt/multimodal/models/qwen3_VL/__init__.py b/python/sgl_jax/srt/multimodal/models/qwen3_VL/__init__.py new file mode 100644 index 0000000000..a1a563ad51 --- /dev/null +++ b/python/sgl_jax/srt/multimodal/models/qwen3_VL/__init__.py @@ -0,0 +1,43 @@ +"""Qwen3-VL model implementation for SGLang-JAX. + +This module provides the Qwen3-VL multimodal model for high-performance +distributed inference on TPUs. + +Components: +- Qwen3_VL_VisionModel: Vision encoder with DeepStack feature extraction +- Qwen3_VL_Generation: Text decoder with M-RoPE for conditional generation +""" + +from sgl_jax.srt.multimodal.models.qwen3_VL.qwen3_vl_generation import ( + MRotaryEmbedding, + Qwen3_VL_Generation, + Qwen3_VL_Model, +) +from sgl_jax.srt.multimodal.models.qwen3_VL.qwen3_vl_vit import ( + Qwen3_VL_VisionModel, + Qwen3_VL_VisionTransformer, + Qwen3_VLImageInputs, + Qwen3_VLVisionAttention, + Qwen3_VLVisionBlock, + Qwen3_VLVisionMLP, + Qwen3_VLVisionPatchEmbed, + Qwen3_VLVisionPatchMerger, + Qwen3_VLVisionRotaryEmbedding, +) + +__all__ = [ + # Vision components + "Qwen3_VL_VisionModel", + "Qwen3_VL_VisionTransformer", + "Qwen3_VLVisionPatchEmbed", + "Qwen3_VLVisionRotaryEmbedding", + "Qwen3_VLVisionMLP", + "Qwen3_VLVisionAttention", + "Qwen3_VLVisionBlock", + "Qwen3_VLVisionPatchMerger", + "Qwen3_VLImageInputs", + # Generation components + "MRotaryEmbedding", + "Qwen3_VL_Model", + "Qwen3_VL_Generation", +] diff --git a/python/sgl_jax/srt/multimodal/models/qwen3_VL/qwen3_vl_generation.py b/python/sgl_jax/srt/multimodal/models/qwen3_VL/qwen3_vl_generation.py new file mode 100644 index 0000000000..affcdb30ff --- /dev/null +++ b/python/sgl_jax/srt/multimodal/models/qwen3_VL/qwen3_vl_generation.py @@ -0,0 +1,727 @@ +"""Qwen3-VL Generation Model for SGLang-JAX. + +Self-contained text decoder — no dependency on qwen2.py. Uses sglang-jax +sharded layers (LinearBase, Embed, RMSNorm, RadixAttention) with Qwen3-VL +specific features: + - Q/K normalization (RMSNorm on query/key before RoPE) + - M-RoPE (Multimodal Rotary Position Embeddings) + +Reference: qwen3-vl/q3vljax/qwen3_vl/modeling.py +""" + +import logging + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx + +from sgl_jax.srt.configs.model_config import ModelConfig +from sgl_jax.srt.layers.embeddings import Embed, ParallelLMHead, apply_rotary_emb +from sgl_jax.srt.layers.layernorm import RMSNorm +from sgl_jax.srt.layers.linear import LinearBase +from sgl_jax.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor +from sgl_jax.srt.layers.radix_attention import RadixAttention +from sgl_jax.srt.mem_cache.memory_pool import KVCache +from sgl_jax.srt.model_executor.forward_batch_info import ForwardBatch +from sgl_jax.srt.utils.weight_utils import WeightLoader, WeightMapping + +logger = logging.getLogger(__name__) + + +# ============================================================================= +# M-RoPE utilities +# ============================================================================= + + +def _apply_interleaved_rope(x: jax.Array, mrope_section: list[int]) -> jax.Array: + """Apply interleaved MRoPE layout. + + Reorganizes frequency layout from chunked [TTT...HHH...WWW] to + interleaved [THTHWHTHW...TT] pattern. + + Args: + x: Frequencies of shape (3, seq_len, head_dim // 2) for T, H, W + mrope_section: Section sizes (e.g., [24, 20, 20]) + + Returns: + Interleaved frequencies of shape (seq_len, head_dim // 2) + """ + x_t = x[0] + x_t = x_t.at[..., 1 : mrope_section[1] * 3 : 3].set(x[1, ..., 1 : mrope_section[1] * 3 : 3]) + x_t = x_t.at[..., 2 : mrope_section[2] * 3 : 3].set(x[2, ..., 2 : mrope_section[2] * 3 : 3]) + return x_t + + +class MRotaryEmbedding: + """Rotary Embedding with Multimodal Sections for Qwen3-VL. + + Implements the M-RoPE mechanism that partitions the head dimension + into sections for temporal (T), height (H), and width (W) positions. + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: jnp.dtype, + mrope_section: list[int], + mrope_interleaved: bool = True, + ) -> None: + del max_position_embeddings + self.head_size = head_size + self.rotary_dim = rotary_dim + self.base = base + self.is_neox_style = is_neox_style + self.dtype = dtype + self.mrope_section = list(mrope_section) + self.mrope_interleaved = mrope_interleaved + + inv_freq_np = 1.0 / (base ** (np.arange(0, rotary_dim, 2, dtype=np.float32) / rotary_dim)) + self._inv_freq_np = inv_freq_np + + # Validate and adjust section sizes + expected_sum = rotary_dim // 2 + actual_sum = sum(self.mrope_section) + if actual_sum != expected_sum: + logger.warning( + "MRoPE section sum mismatch: expected %s, got %s. Adjusting.", + expected_sum, + actual_sum, + ) + if actual_sum > 0: + scale_factor = expected_sum / actual_sum + self.mrope_section = [ + max(1, int(section * scale_factor)) for section in self.mrope_section + ] + current_sum = sum(self.mrope_section) + if current_sum != expected_sum: + self.mrope_section[-1] += expected_sum - current_sum + else: + self.mrope_section = [expected_sum // len(self.mrope_section)] * len( + self.mrope_section + ) + remainder = expected_sum % len(self.mrope_section) + for i in range(remainder): + self.mrope_section[i] += 1 + + def __call__( + self, + positions: jax.Array, + query: jax.Array, + key: jax.Array, + ) -> tuple[jax.Array, jax.Array]: + """Apply M-RoPE to query and key. + + Args: + positions: Position IDs. Shape (seq,) for text-only or (3, seq) for multimodal + query: Query tensor of shape (seq, hidden_dim) + key: Key tensor of shape (seq, hidden_dim) + + Returns: + Rotated query and key tensors + """ + inv_freq = jnp.asarray(self._inv_freq_np, dtype=self.dtype) + + if positions.ndim == 1: + # Text-only: simple 1D positions + freqs = jnp.einsum("n,d->nd", positions.astype(jnp.float32), inv_freq) + cos = jnp.cos(freqs).astype(self.dtype) + sin = jnp.sin(freqs).astype(self.dtype) + else: + # Multimodal: 3D positions (T, H, W) + freqs = jnp.einsum("tn,d->tnd", positions.astype(jnp.float32), inv_freq) + cos = jnp.cos(freqs).astype(self.dtype) + sin = jnp.sin(freqs).astype(self.dtype) + + if self.mrope_interleaved: + cos = _apply_interleaved_rope(cos, self.mrope_section) + sin = _apply_interleaved_rope(sin, self.mrope_section) + else: + cos_slices = [] + sin_slices = [] + offset = 0 + for i, section in enumerate(self.mrope_section): + cos_slices.append(cos[i, :, offset : offset + section]) + sin_slices.append(sin[i, :, offset : offset + section]) + offset += section + cos = jnp.concatenate(cos_slices, axis=-1) + sin = jnp.concatenate(sin_slices, axis=-1) + + num_tokens = positions.shape[-1] + query_shape = query.shape + query = query.reshape(num_tokens, -1, self.head_size) + query_rot = query[..., : self.rotary_dim] + query_pass = query[..., self.rotary_dim :] + query_rot = apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) + query = jnp.concatenate((query_rot, query_pass), axis=-1).reshape(query_shape) + + key_shape = key.shape + key = key.reshape(num_tokens, -1, self.head_size) + key_rot = key[..., : self.rotary_dim] + key_pass = key[..., self.rotary_dim :] + key_rot = apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) + key = jnp.concatenate((key_rot, key_pass), axis=-1).reshape(key_shape) + + return query, key + + +# ============================================================================= +# Config helper +# ============================================================================= + + +def _get_mrope_section(config) -> list[int]: + """Extract mrope_section from either our config or HuggingFace config. + + Our Qwen3VLTextConfig: config.mrope_section = (24, 20, 20) + HuggingFace config: config.rope_scaling = {'mrope_section': [24, 20, 20], ...} + """ + # Our config: direct attribute + if hasattr(config, "mrope_section") and config.mrope_section: + return list(config.mrope_section) + # HF config: nested in rope_scaling dict + rope_scaling = getattr(config, "rope_scaling", None) or {} + return rope_scaling.get("mrope_section", [24, 20, 20]) + + +# ============================================================================= +# Model components (self-contained, using sglang-jax sharded layers) +# ============================================================================= + + +class Qwen3VL_MLP(nnx.Module): + """SiLU-gated MLP for Qwen3-VL text decoder.""" + + def __init__(self, config, mesh, layer_id: int = 0, dtype=jnp.bfloat16): + self.gate_proj = LinearBase( + input_size=config.hidden_size, + output_size=config.intermediate_size, + kernel_axes=(None, "tensor"), + use_bias=False, + params_dtype=dtype, + mesh=mesh, + ) + self.up_proj = LinearBase( + input_size=config.hidden_size, + output_size=config.intermediate_size, + kernel_axes=(None, "tensor"), + use_bias=False, + params_dtype=dtype, + mesh=mesh, + ) + self.down_proj = LinearBase( + input_size=config.intermediate_size, + output_size=config.hidden_size, + kernel_axes=("tensor", None), + use_bias=False, + params_dtype=dtype, + mesh=mesh, + ) + + def __call__(self, hidden_states: jax.Array) -> jax.Array: + gate, _ = self.gate_proj(hidden_states) + up, _ = self.up_proj(hidden_states) + output, _ = self.down_proj(jax.nn.silu(gate) * up) + return output + + +class Qwen3VL_Attention(nnx.Module): + """Qwen3-VL text decoder attention with Q/K norms, GQA, and M-RoPE.""" + + def __init__(self, config, mesh, layer_id: int = 0, dtype=jnp.bfloat16): + self.layer_id = layer_id + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.q_head_num = config.num_attention_heads + self.kv_head_num = config.num_key_value_heads + self.q_size = self.q_head_num * self.head_dim + self.kv_size = self.kv_head_num * self.head_dim + self.scaling = self.head_dim**-0.5 + + use_bias = getattr(config, "attention_bias", False) + + self.q_proj = LinearBase( + input_size=config.hidden_size, + output_size=self.q_size, + use_bias=use_bias, + kernel_axes=(None, "tensor"), + params_dtype=dtype, + mesh=mesh, + ) + self.k_proj = LinearBase( + input_size=config.hidden_size, + output_size=self.kv_size, + use_bias=use_bias, + kernel_axes=(None, "tensor"), + params_dtype=dtype, + mesh=mesh, + ) + self.v_proj = LinearBase( + input_size=config.hidden_size, + output_size=self.kv_size, + use_bias=use_bias, + kernel_axes=(None, "tensor"), + params_dtype=dtype, + mesh=mesh, + ) + self.o_proj = LinearBase( + input_size=self.q_size, + output_size=config.hidden_size, + use_bias=False, + kernel_axes=("tensor", None), + params_dtype=dtype, + mesh=mesh, + ) + + # Qwen3-VL specific: Q/K normalization before RoPE + self.q_norm = RMSNorm( + self.head_dim, + epsilon=config.rms_norm_eps, + param_dtype=dtype, + ) + self.k_norm = RMSNorm( + self.head_dim, + epsilon=config.rms_norm_eps, + param_dtype=dtype, + ) + + # M-RoPE + mrope_section = _get_mrope_section(config) + rope_theta = getattr(config, "rope_theta", 5_000_000) + + self.rotary_emb = MRotaryEmbedding( + head_size=self.head_dim, + rotary_dim=self.head_dim, + max_position_embeddings=32768, + base=rope_theta, + is_neox_style=True, + dtype=dtype, + mrope_section=mrope_section, + ) + + self.attn = RadixAttention( + num_heads=self.q_head_num, + head_dim=self.head_dim, + scaling=self.scaling, + num_kv_heads=self.kv_head_num, + layer_id=layer_id, + ) + + def __call__( + self, + positions: jax.Array, + hidden_states: jax.Array, + forward_batch: ForwardBatch, + token_to_kv_pool: KVCache, + ) -> tuple[jax.Array, jax.Array]: + q, _ = self.q_proj(hidden_states) + k, _ = self.k_proj(hidden_states) + v, _ = self.v_proj(hidden_states) + + q = q.reshape(-1, self.q_head_num, self.head_dim) + k = k.reshape(-1, self.kv_head_num, self.head_dim) + v = v.reshape(-1, self.kv_head_num, self.head_dim) + + # Qwen3-VL specific: normalize Q/K before RoPE + q = self.q_norm(q) + k = self.k_norm(k) + + q, k = self.rotary_emb(positions, q, k) + attn_output, kv_fused = self.attn(q, k, v, forward_batch, token_to_kv_pool) + + output, _ = self.o_proj(attn_output) + return output, kv_fused + + +class Qwen3VL_DecoderLayer(nnx.Module): + """Single decoder layer for Qwen3-VL.""" + + def __init__(self, config, mesh, layer_id: int = 0, dtype=jnp.bfloat16): + self.layer_id = layer_id + self.hidden_size = config.hidden_size + + self.self_attn = Qwen3VL_Attention( + config=config, + mesh=mesh, + layer_id=layer_id, + dtype=dtype, + ) + self.mlp = Qwen3VL_MLP( + config=config, + mesh=mesh, + layer_id=layer_id, + dtype=dtype, + ) + self.input_layernorm = RMSNorm( + config.hidden_size, + epsilon=config.rms_norm_eps, + param_dtype=dtype, + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, + epsilon=config.rms_norm_eps, + param_dtype=dtype, + ) + + def __call__( + self, + positions: jax.Array, + hidden_states: jax.Array, + forward_batch: ForwardBatch, + token_to_kv_pool: KVCache, + residual: jax.Array | None = None, + ): + layer_callback_flag = [] + + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states += residual + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, kv_fused = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + token_to_kv_pool=token_to_kv_pool, + ) + + hidden_states += residual + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + + return hidden_states, residual, kv_fused, layer_callback_flag + + +class Qwen3_VL_Model(nnx.Module): + """Qwen3-VL text decoder with M-RoPE (self-contained, no Qwen2 dependency).""" + + def __init__(self, config, mesh, dtype=jnp.bfloat16): + self.config = config + + self.embed_tokens = Embed( + num_embeddings=config.vocab_size, + features=config.hidden_size, + dtype=dtype, + kernel_axes=("tensor", None), + param_dtype=dtype, + mesh=mesh, + ) + + self.layers = nnx.data( + [ + Qwen3VL_DecoderLayer( + config=config, + layer_id=i, + dtype=dtype, + mesh=mesh, + ) + for i in range(config.num_hidden_layers) + ] + ) + + self.norm = RMSNorm( + config.hidden_size, + epsilon=config.rms_norm_eps, + param_dtype=dtype, + ) + + # Store mrope_section for position routing + self._mrope_section = _get_mrope_section(config) + + def __call__( + self, + forward_batch: ForwardBatch, + token_to_kv_pool: KVCache, + ): + residual = None + + # Use input embeddings if provided (for multimodal prefill) + input_embeds = ( + forward_batch.input_embedding + if forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed() + else None + ) + hidden_states = ( + self.embed_tokens(forward_batch.input_ids) if input_embeds is None else input_embeds + ) + + # Use M-RoPE positions if available + rope_positions = ( + forward_batch.mrope_positions + if self._mrope_section and forward_batch.mrope_positions is not None + else forward_batch.positions + ) + + layers_kv_fused = [] + layers_callback_flag = [] + + for layer in self.layers: + hidden_states, residual, kv_fused, callback_flag = layer( + rope_positions, + hidden_states, + forward_batch, + token_to_kv_pool, + residual, + ) + layers_kv_fused.append(kv_fused) + layers_callback_flag.extend(callback_flag) + + if residual is not None: + hidden_states += residual + hidden_states = self.norm(hidden_states) + + return hidden_states, layers_kv_fused, layers_callback_flag + + +# ============================================================================= +# Top-level generation model +# ============================================================================= + + +class Qwen3_VL_Generation(nnx.Module): + """Qwen3-VL model for conditional generation. + + Self-contained implementation (no Qwen2 dependency). + Accepts either our Qwen3VLConfig or a HuggingFace config with .text_config. + + Architecture: + - Vision encoder (separate module): Processes images/videos to embeddings + - Language model (self.model): Generates text with M-RoPE + + Usage Pattern: + 1. PREFILL (once per image): + - Process vision with Qwen3_VL_VisionModel + - Merge embeddings with get_input_embeddings() + - Call __call__() with merged embeddings + 2. DECODE (many times for text generation): + - Call __call__() without embeddings (uses text tokens only) + """ + + def __init__(self, config=None, dtype=None, mesh=None): + super().__init__() + self.mesh = mesh + self.config = config + self.dtype = dtype or jnp.bfloat16 + + # Extract text config (works with both our Qwen3VLConfig and HF config) + self.text_config = getattr(config, "text_config", None) or config + + self.model = Qwen3_VL_Model(self.text_config, mesh=mesh, dtype=self.dtype) + + if not getattr(self.text_config, "tie_word_embeddings", False): + self.lm_head = ParallelLMHead( + self.text_config.vocab_size, + self.text_config.hidden_size, + dtype=self.dtype, + param_dtype=self.dtype, + kernel_axes=("tensor", None), + ) + + self.logits_processor = LogitsProcessor(self.text_config.vocab_size, mesh=self.mesh) + + # Multimodal token IDs + self.image_token_id = getattr(self.config, "image_token_id", 151655) + self.video_token_id = getattr(self.config, "video_token_id", 151656) + + def load_weights(self, model_config: ModelConfig): + """Load model weights from safetensors.""" + loader = WeightLoader( + model=self, + model_config=model_config, + mesh=self.mesh, + dtype=self.dtype, + ) + weight_mappings = self._create_qwen3_weight_mappings() + loader.load_weights_from_safetensors(weight_mappings) + logger.info("Qwen3-VL (LLM) weights loaded successfully!") + + def _create_qwen3_weight_mappings(self) -> dict: + """Create weight mappings for text decoder. + + HF safetensors keys use ``model.language_model.`` prefix; + our JAX model paths use ``model.`` (no language_model). + """ + mappings = { + "model.language_model.embed_tokens.weight": WeightMapping( + target_path="model.embed_tokens.embedding", + sharding=("tensor", None), + transpose=False, + ), + "model.language_model.norm.weight": WeightMapping( + target_path="model.norm.scale", + sharding=(None,), + transpose=False, + ), + } + + if not getattr(self.text_config, "tie_word_embeddings", False): + mappings["lm_head.weight"] = WeightMapping( + target_path="lm_head.embedding", + sharding=("tensor", None), + transpose=False, + ) + + num_layers = self.text_config.num_hidden_layers + for layer_idx in range(num_layers): + mappings.update(self._create_layer_mappings(layer_idx)) + + return mappings + + def _create_layer_mappings(self, layer_idx: int) -> dict: + """Create weight mappings for a single decoder layer.""" + prefix = f"model.language_model.layers.{layer_idx}" + target_prefix = f"model.layers.{layer_idx}" + + mappings = { + # Layer norms + f"{prefix}.input_layernorm.weight": WeightMapping( + target_path=f"{target_prefix}.input_layernorm.scale", + sharding=(None,), + transpose=False, + ), + f"{prefix}.post_attention_layernorm.weight": WeightMapping( + target_path=f"{target_prefix}.post_attention_layernorm.scale", + sharding=(None,), + transpose=False, + ), + # Attention projections + f"{prefix}.self_attn.q_proj.weight": WeightMapping( + target_path=f"{target_prefix}.self_attn.q_proj.weight", + sharding=(None, "tensor"), + transpose=True, + head_dim_padding=True, + kv_head_padding=False, + ), + f"{prefix}.self_attn.k_proj.weight": WeightMapping( + target_path=f"{target_prefix}.self_attn.k_proj.weight", + sharding=(None, "tensor"), + transpose=True, + head_dim_padding=True, + kv_head_padding=True, + ), + f"{prefix}.self_attn.v_proj.weight": WeightMapping( + target_path=f"{target_prefix}.self_attn.v_proj.weight", + sharding=(None, "tensor"), + transpose=True, + head_dim_padding=True, + kv_head_padding=True, + ), + f"{prefix}.self_attn.o_proj.weight": WeightMapping( + target_path=f"{target_prefix}.self_attn.o_proj.weight", + sharding=("tensor", None), + transpose=True, + head_dim_padding=True, + kv_head_padding=False, + ), + # Q/K norms (Qwen3-VL specific) + f"{prefix}.self_attn.q_norm.weight": WeightMapping( + target_path=f"{target_prefix}.self_attn.q_norm.scale", + sharding=(None,), + transpose=False, + ), + f"{prefix}.self_attn.k_norm.weight": WeightMapping( + target_path=f"{target_prefix}.self_attn.k_norm.scale", + sharding=(None,), + transpose=False, + ), + # MLP + f"{prefix}.mlp.gate_proj.weight": WeightMapping( + target_path=f"{target_prefix}.mlp.gate_proj.weight", + sharding=(None, "tensor"), + transpose=True, + ), + f"{prefix}.mlp.up_proj.weight": WeightMapping( + target_path=f"{target_prefix}.mlp.up_proj.weight", + sharding=(None, "tensor"), + transpose=True, + ), + f"{prefix}.mlp.down_proj.weight": WeightMapping( + target_path=f"{target_prefix}.mlp.down_proj.weight", + sharding=("tensor", None), + transpose=True, + ), + } + + # Add bias mappings if attention uses bias + if getattr(self.text_config, "attention_bias", False): + mappings.update( + { + f"{prefix}.self_attn.q_proj.bias": WeightMapping( + target_path=f"{target_prefix}.self_attn.q_proj.bias", + sharding=(None,), + transpose=False, + head_dim_padding=True, + kv_head_padding=False, + ), + f"{prefix}.self_attn.k_proj.bias": WeightMapping( + target_path=f"{target_prefix}.self_attn.k_proj.bias", + sharding=(None,), + transpose=False, + head_dim_padding=True, + kv_head_padding=True, + ), + f"{prefix}.self_attn.v_proj.bias": WeightMapping( + target_path=f"{target_prefix}.self_attn.v_proj.bias", + sharding=(None,), + transpose=False, + head_dim_padding=True, + kv_head_padding=True, + ), + } + ) + + return mappings + + def get_embed_and_head(self): + """Get embedding and lm_head weights for tied embeddings handling.""" + if getattr(self.text_config, "tie_word_embeddings", False): + weight = self.model.embed_tokens.embedding.value + return (weight, weight) + return (self.model.embed_tokens.embedding.value, self.lm_head.embedding.value) + + def set_embed_and_head( + self, + embed_weight: jax.Array | None = None, + head_weight: jax.Array | None = None, + ) -> None: + """Set embedding and lm_head weights.""" + if embed_weight is not None: + self.model.embed_tokens.embedding.value = embed_weight + if head_weight is not None: + self.lm_head.embedding.value = head_weight + + def __call__( + self, + forward_batch: ForwardBatch, + token_to_kv_pool: KVCache, + logits_metadata: LogitsMetadata, + ): + """Forward pass for text generation. + + Args: + forward_batch: Batch information including input_ids and positions + token_to_kv_pool: KV cache for inference + logits_metadata: Metadata for logits processing + + Returns: + Tuple of (logits, layers_kv_fused, layers_callback_flag, None) + """ + hidden_states, layers_kv_fused, layers_callback_flag = self.model( + forward_batch, token_to_kv_pool + ) + + if not getattr(self.text_config, "tie_word_embeddings", False): + output = self.logits_processor(hidden_states, self.lm_head, logits_metadata) + else: + output = self.logits_processor(hidden_states, self.model.embed_tokens, logits_metadata) + + return output, layers_kv_fused, layers_callback_flag, None diff --git a/python/sgl_jax/srt/multimodal/models/qwen3_VL/qwen3_vl_vit.py b/python/sgl_jax/srt/multimodal/models/qwen3_VL/qwen3_vl_vit.py new file mode 100644 index 0000000000..e37366aa98 --- /dev/null +++ b/python/sgl_jax/srt/multimodal/models/qwen3_VL/qwen3_vl_vit.py @@ -0,0 +1,958 @@ +"""Qwen3-VL Vision Transformer for SGLang-JAX. + +This module implements the vision encoder for Qwen3-VL with DeepStack feature extraction. +Key differences from Qwen2.5-VL: +- DeepStack: Extracts features at intermediate layers for early LLM fusion +- Full attention: No windowed attention in vision encoder +- Simpler position embeddings: Bilinear interpolation + 2D RoPE +""" + +import logging +import math +from functools import partial +from typing import Literal, TypedDict + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +from jax.sharding import Mesh +from transformers import modeling_flax_utils + +from sgl_jax.srt.layers.embeddings import Embed +from sgl_jax.srt.multimodal.configs.qwen_vl.qwen3_vl_config import ( + Qwen3VLConfig, + Qwen3VLVisionConfig, +) +from sgl_jax.srt.utils.jax_utils import is_tpu_runtime +from sgl_jax.srt.utils.weight_utils import WeightLoader, WeightMapping + +_FLASH_MHA = None + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +init_fn = nnx.initializers.uniform() + + +def _get_flash_mha(): + global _FLASH_MHA + if _FLASH_MHA is None: + from flash_attn_jax import flash_mha as _FLASH_MHA + return _FLASH_MHA + + +class Qwen3_VLImagePixelInputs(TypedDict): + type: Literal["pixel_values"] + pixel_values: jax.Array + image_grid_thw: tuple[tuple[int, int, int], ...] + + +class Qwen3_VLImageEmbeddingInputs(TypedDict): + type: Literal["image_embeds"] + image_embeds: jax.Array + image_grid_thw: jax.Array + + +Qwen3_VLImageInputs = Qwen3_VLImagePixelInputs | Qwen3_VLImageEmbeddingInputs + + +def apply_rotary_pos_emb_vision( + q: jax.Array, k: jax.Array, cos: jax.Array, sin: jax.Array +) -> tuple[jax.Array, jax.Array]: + """Apply rotary position embeddings to query and key for vision. + + Uses the rotate_half formulation: x * cos + rotate_half(x) * sin + """ + # q, k: (seq, heads, head_dim) + # cos, sin: (seq, head_dim) + half_dim = q.shape[-1] // 2 + + # Split into real and imaginary parts + q1, q2 = q[..., :half_dim], q[..., half_dim:] + k1, k2 = k[..., :half_dim], k[..., half_dim:] + + # cos, sin need to be (seq, 1, half_dim) for broadcasting + cos = cos[:, None, :] # (seq, 1, head_dim) + sin = sin[:, None, :] + cos1, cos2 = cos[..., :half_dim], cos[..., half_dim:] + sin1, sin2 = sin[..., :half_dim], sin[..., half_dim:] + + # Apply rotation: [x1, x2] -> [x1*cos - x2*sin, x2*cos + x1*sin] + q_rot = jnp.concatenate([q1 * cos1 - q2 * sin1, q2 * cos2 + q1 * sin2], axis=-1) + k_rot = jnp.concatenate([k1 * cos1 - k2 * sin1, k2 * cos2 + k1 * sin2], axis=-1) + + return q_rot, k_rot + + +def vision_attention( + q: jax.Array, + k: jax.Array, + v: jax.Array, + scale: float, +) -> jax.Array: + """Compute vision attention. + + Full attention (no windowing) for Qwen3-VL. + + Args: + q, k, v: Input tensors of shape [seq_len, num_heads, head_dim] + scale: Attention scale factor (1/sqrt(head_dim)) + + Returns: + Output tensor of shape [seq_len, num_heads, head_dim] + """ + if not is_tpu_runtime(): + # GPU: use flash_mha with batch dim + flash_mha = _get_flash_mha() + original_dtype = q.dtype + if q.dtype not in [jnp.bfloat16, jnp.float16]: + q = q.astype(jnp.bfloat16) + k = k.astype(jnp.bfloat16) + v = v.astype(jnp.bfloat16) + + # Add batch dimension for flash_mha: (1, seq, heads, dim) + q = q[None, :, :, :] + k = k[None, :, :, :] + v = v[None, :, :, :] + + output = flash_mha(q, k, v, softmax_scale=scale, is_causal=False) + output = output[0] # Remove batch dim + + if output.dtype != original_dtype: + output = output.astype(original_dtype) + return output + else: + # TPU: native attention + # Transpose to (heads, seq, dim) for matmul + q = jnp.transpose(q, (1, 0, 2)) + k = jnp.transpose(k, (1, 0, 2)) + v = jnp.transpose(v, (1, 0, 2)) + + attn_weights = jnp.matmul(q, k.transpose(0, 2, 1)) * scale + attn_weights = jax.nn.softmax(attn_weights.astype(jnp.float32), axis=-1).astype(q.dtype) + output = jnp.matmul(attn_weights, v) + + # Transpose back to (seq, heads, dim) + return jnp.transpose(output, (1, 0, 2)) + + +class Qwen3_VLVisionPatchEmbed(nnx.Module): + """3D Convolutional patch embedding for vision input.""" + + def __init__( + self, + config: Qwen3VLVisionConfig, + dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs = None, + ) -> None: + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.hidden_size = config.hidden_size + kernel_size = (config.temporal_patch_size, config.patch_size, config.patch_size) + + self.proj = nnx.Conv( + in_features=config.in_channels, + out_features=config.hidden_size, + kernel_size=kernel_size, + strides=kernel_size, + use_bias=True, + param_dtype=dtype, + rngs=rngs or nnx.Rngs(0), + ) + + def __call__(self, x: jax.Array) -> jax.Array: + # x: (num_patches, in_channels * temporal_patch_size * patch_size * patch_size) + L, dim = x.shape + C = dim // (self.temporal_patch_size * self.patch_size * self.patch_size) + + # Reshape to (L, C, T, H, W) then transpose to (L, T, H, W, C) for Conv + x = x.reshape(L, C, self.temporal_patch_size, self.patch_size, self.patch_size) + x = jnp.transpose(x, (0, 2, 3, 4, 1)) + + # Apply conv: (L, T, H, W, C) -> (L, 1, 1, 1, hidden_size) + x = self.proj(x) + return x.reshape(L, self.hidden_size) + + +class Qwen3_VLVisionRotaryEmbedding(nnx.Module): + """Rotary position embedding for vision encoder.""" + + def __init__(self, dim: int, theta: float = 10000.0): + self.dim = dim + self.theta = theta + + def __call__(self, seq_len: int) -> jax.Array: + inv_freq = 1.0 / (self.theta ** (jnp.arange(0, self.dim, 2, dtype=jnp.float32) / self.dim)) + seq = jnp.arange(seq_len, dtype=jnp.float32) + freqs = jnp.outer(seq, inv_freq) + return freqs + + +class Qwen3_VLVisionMLP(nnx.Module): + """Vision encoder MLP with GELU activation.""" + + def __init__( + self, + config: Qwen3VLVisionConfig, + dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs = None, + ): + _rngs = rngs or nnx.Rngs(0) + self.linear_fc1 = nnx.Linear( + config.hidden_size, + config.intermediate_size, + use_bias=True, + param_dtype=dtype, + rngs=_rngs, + ) + self.linear_fc2 = nnx.Linear( + config.intermediate_size, + config.hidden_size, + use_bias=True, + param_dtype=dtype, + rngs=_rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.linear_fc1(x) + x = nnx.gelu(x, approximate=True) + x = self.linear_fc2(x) + return x + + +class Qwen3_VLVisionAttention(nnx.Module): + """Vision encoder multi-head attention with RoPE.""" + + def __init__( + self, + config: Qwen3VLVisionConfig, + dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs = None, + ): + self.hidden_size = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = config.head_dim + self.scale = self.head_dim**-0.5 + + _rngs = rngs or nnx.Rngs(0) + self.qkv_proj = nnx.Linear( + self.hidden_size, + 3 * self.hidden_size, + use_bias=True, + param_dtype=dtype, + rngs=_rngs, + ) + self.proj = nnx.Linear( + self.hidden_size, + self.hidden_size, + use_bias=True, + param_dtype=dtype, + rngs=_rngs, + ) + + def __call__( + self, + hidden_states: jax.Array, + position_embeddings: tuple[jax.Array, jax.Array], + ) -> jax.Array: + seq_len = hidden_states.shape[0] + cos, sin = position_embeddings + + # QKV projection: (seq, hidden) -> (seq, 3, heads, head_dim) + qkv = self.qkv_proj(hidden_states).reshape(seq_len, 3, self.num_heads, self.head_dim) + q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] # Each: (seq, heads, head_dim) + + # Apply RoPE + q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + + # Compute attention + out = vision_attention(q, k, v, self.scale) + + # Reshape and project + out = out.reshape(seq_len, -1) + return self.proj(out) + + +class Qwen3_VLVisionBlock(nnx.Module): + """Single transformer block for vision encoder.""" + + def __init__( + self, + config: Qwen3VLVisionConfig, + dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs = None, + ): + _rngs = rngs or nnx.Rngs(0) + self.norm1 = nnx.LayerNorm( + config.hidden_size, + epsilon=config.layer_norm_eps, + dtype=dtype, + rngs=_rngs, + ) + self.norm2 = nnx.LayerNorm( + config.hidden_size, + epsilon=config.layer_norm_eps, + dtype=dtype, + rngs=_rngs, + ) + self.attn = Qwen3_VLVisionAttention(config, dtype=dtype, rngs=rngs) + self.mlp = Qwen3_VLVisionMLP(config, dtype=dtype, rngs=rngs) + + def __call__( + self, + hidden_states: jax.Array, + position_embeddings: tuple[jax.Array, jax.Array], + ) -> jax.Array: + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.attn(hidden_states, position_embeddings) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class Qwen3_VLVisionPatchMerger(nnx.Module): + """Merge spatial patches after vision encoding. + + Two modes: + - use_postshuffle_norm=False: For main merger (norm before shuffle) + - use_postshuffle_norm=True: For DeepStack mergers (norm after shuffle) + """ + + def __init__( + self, + config: Qwen3VLVisionConfig, + use_postshuffle_norm: bool = False, + dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs = None, + ): + self.spatial_merge_size = config.spatial_merge_size + merge_factor = config.spatial_merge_size**2 + self.hidden_merged = config.hidden_size * merge_factor + self.use_postshuffle_norm = use_postshuffle_norm + + norm_dim = self.hidden_merged if use_postshuffle_norm else config.hidden_size + + _rngs = rngs or nnx.Rngs(0) + self.norm = nnx.LayerNorm( + norm_dim, + epsilon=config.layer_norm_eps, + dtype=dtype, + rngs=_rngs, + ) + self.linear_fc1 = nnx.Linear( + self.hidden_merged, + self.hidden_merged, + use_bias=True, + param_dtype=dtype, + rngs=_rngs, + ) + self.linear_fc2 = nnx.Linear( + self.hidden_merged, + config.out_hidden_size, + use_bias=True, + param_dtype=dtype, + rngs=_rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + if not self.use_postshuffle_norm: + x = self.norm(x) + + # Reshape to merge spatial patches + merge_factor = self.spatial_merge_size**2 + n_patches = x.shape[0] // merge_factor + x = x.reshape(n_patches, -1) + + if self.use_postshuffle_norm: + x = self.norm(x) + + x = self.linear_fc1(x) + x = nnx.gelu(x) + x = self.linear_fc2(x) + return x + + +class Qwen3_VL_VisionTransformer(nnx.Module): + """Complete vision transformer with DeepStack feature extraction.""" + + def __init__( + self, + config: Qwen3VLVisionConfig, + dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs = None, + mesh: Mesh = None, + ): + self.config = config + self.dtype = dtype + self.spatial_merge_size = config.spatial_merge_size + self.spatial_merge_unit = config.spatial_merge_size**2 + + self.patch_embed = Qwen3_VLVisionPatchEmbed(config, dtype=dtype, rngs=rngs) + + # Position embedding (learnable, for interpolation) + self.num_grid_per_side = int(config.num_position_embeddings**0.5) + _rngs = rngs or nnx.Rngs(0) + self.pos_embed = nnx.Embed( + num_embeddings=config.num_position_embeddings, + features=config.hidden_size, + param_dtype=dtype, + rngs=_rngs, + ) + + # Rotary embeddings + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen3_VLVisionRotaryEmbedding(head_dim // 2, theta=config.rope_theta) + + # Transformer blocks + self.blocks = nnx.List( + [Qwen3_VLVisionBlock(config, dtype=dtype, rngs=rngs) for _ in range(config.depth)] + ) + + # Main merger + self.merger = Qwen3_VLVisionPatchMerger( + config, + use_postshuffle_norm=False, + dtype=dtype, + rngs=rngs, + ) + + # DeepStack mergers (extract at intermediate layers) + self.deepstack_visual_indexes = config.deepstack_visual_indexes + self.deepstack_merger_list = nnx.List( + [ + Qwen3_VLVisionPatchMerger( + config, + use_postshuffle_norm=True, + dtype=dtype, + rngs=rngs, + ) + for _ in range(len(config.deepstack_visual_indexes)) + ] + ) + + def _fast_pos_embed_interpolate(self, grid_thw: tuple[tuple[int, int, int], ...]) -> jax.Array: + """Bilinear interpolation for position embeddings.""" + all_pos_embeds = [] + + for t, h, w in grid_thw: + # Create interpolation indices + h_idxs = jnp.linspace(0, self.num_grid_per_side - 1, h) + w_idxs = jnp.linspace(0, self.num_grid_per_side - 1, w) + + h_floor = jnp.floor(h_idxs).astype(jnp.int32) + w_floor = jnp.floor(w_idxs).astype(jnp.int32) + h_ceil = jnp.clip(h_floor + 1, 0, self.num_grid_per_side - 1) + w_ceil = jnp.clip(w_floor + 1, 0, self.num_grid_per_side - 1) + + dh = h_idxs - h_floor + dw = w_idxs - w_floor + + # 2D grid indices for 4 corners + base_h = h_floor * self.num_grid_per_side + base_h_ceil = h_ceil * self.num_grid_per_side + + idx00 = (base_h[:, None] + w_floor[None, :]).flatten() + idx01 = (base_h[:, None] + w_ceil[None, :]).flatten() + idx10 = (base_h_ceil[:, None] + w_floor[None, :]).flatten() + idx11 = (base_h_ceil[:, None] + w_ceil[None, :]).flatten() + + # Weights for bilinear interpolation + w00 = ((1 - dh)[:, None] * (1 - dw)[None, :]).flatten() + w01 = ((1 - dh)[:, None] * dw[None, :]).flatten() + w10 = (dh[:, None] * (1 - dw)[None, :]).flatten() + w11 = (dh[:, None] * dw[None, :]).flatten() + + # Lookup and interpolate + pos_embeds = ( + self.pos_embed(idx00) * w00[:, None] + + self.pos_embed(idx01) * w01[:, None] + + self.pos_embed(idx10) * w10[:, None] + + self.pos_embed(idx11) * w11[:, None] + ) + + # Repeat for temporal dimension and apply spatial merge permutation + pos_embeds = pos_embeds.reshape(h, w, -1) + if t > 1: + pos_embeds = jnp.tile(pos_embeds[None], (t, 1, 1, 1)) + else: + pos_embeds = pos_embeds[None] + + # Permute for spatial merge + merge_size = self.spatial_merge_size + merged_h, merged_w = h // merge_size, w // merge_size + pos_embeds = pos_embeds.reshape(t, merged_h, merge_size, merged_w, merge_size, -1) + pos_embeds = pos_embeds.transpose(0, 1, 3, 2, 4, 5) + pos_embeds = pos_embeds.reshape(-1, pos_embeds.shape[-1]) + + all_pos_embeds.append(pos_embeds) + + return jnp.concatenate(all_pos_embeds, axis=0) + + def _rot_pos_emb( + self, grid_thw: tuple[tuple[int, int, int], ...] + ) -> tuple[jax.Array, jax.Array]: + """Compute rotary position embeddings.""" + merge_size = self.spatial_merge_size + all_embeddings = [] + + for grid_t, grid_h, grid_w in grid_thw: + merged_h, merged_w = grid_h // merge_size, grid_w // merge_size + + # Compute position indices for each patch + block_rows = jnp.arange(merged_h) + block_cols = jnp.arange(merged_w) + intra_row = jnp.arange(merge_size) + intra_col = jnp.arange(merge_size) + + # Full resolution positions + row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None] + col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :] + + row_idx = jnp.broadcast_to( + row_idx, (merged_h, merged_w, merge_size, merge_size) + ).reshape(-1) + col_idx = jnp.broadcast_to( + col_idx, (merged_h, merged_w, merge_size, merge_size) + ).reshape(-1) + + # Repeat for temporal dimension + if grid_t > 1: + row_idx = jnp.tile(row_idx, grid_t) + col_idx = jnp.tile(col_idx, grid_t) + + # Create frequency table + max_hw = max(grid_h, grid_w) + head_dim = self.config.hidden_size // self.config.num_heads + freq_table = self.rotary_pos_emb(max_hw) # (max_hw, rotary_dim//2) + + # Lookup embeddings + row_emb = freq_table[row_idx] + col_emb = freq_table[col_idx] + + # Concatenate and double + emb = jnp.concatenate([row_emb, col_emb], axis=-1) + emb = jnp.concatenate([emb, emb], axis=-1) + + all_embeddings.append(emb) + + all_emb = jnp.concatenate(all_embeddings, axis=0) + cos = jnp.cos(all_emb) + sin = jnp.sin(all_emb) + return cos, sin + + def __call__( + self, + hidden_states: jax.Array, + grid_thw: tuple[tuple[int, int, int], ...], + ) -> tuple[jax.Array, list[jax.Array]]: + """Forward pass through vision transformer. + + Args: + hidden_states: Flattened pixel values (num_patches, patch_dim) + grid_thw: Grid dimensions for each image/video + + Returns: + Tuple of (merged_features, deepstack_features) + """ + hidden_states = self.patch_embed(hidden_states) + seq_len = hidden_states.shape[0] + + # Position embeddings with bilinear interpolation + pos_embeds = self._fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds[:seq_len] + + # RoPE embeddings + cos, sin = self._rot_pos_emb(grid_thw) + position_embeddings = (cos[:seq_len], sin[:seq_len]) + + # Process through transformer blocks + deepstack_features = [] + for layer_idx, block in enumerate(self.blocks): + hidden_states = block(hidden_states, position_embeddings) + + # Extract DeepStack features at specified layers + if layer_idx in self.deepstack_visual_indexes: + ds_idx = list(self.deepstack_visual_indexes).index(layer_idx) + deepstack_features.append(self.deepstack_merger_list[ds_idx](hidden_states)) + + # Final merger + hidden_states = self.merger(hidden_states) + + return hidden_states, deepstack_features + + +class Qwen3_VL_VisionModel(nnx.Module): + """Qwen3-VL Vision Model with weight loading support.""" + + def __init__( + self, + config: Qwen3VLVisionConfig, + dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs = None, + mesh: Mesh = None, + ) -> None: + self.config = config + self.dtype = dtype + self.mesh = mesh + self.visual = Qwen3_VL_VisionTransformer( + config=config, + dtype=dtype, + rngs=rngs, + mesh=mesh, + ) + logger.info("Qwen3_VL_VisionModel initialized with dtype %s", dtype) + + def load_weights(self, model_config) -> None: + """Load model weights from safetensors.""" + if not hasattr(self, "text_embed"): + self.text_embed = Embed( + num_embeddings=model_config.vocab_size, + features=model_config.text_hidden_size, + dtype=self.dtype, + param_dtype=self.dtype, + kernel_axes=(None, None), + mesh=self.mesh, + ) + + loader = WeightLoader( + model=self, + model_config=model_config, + mesh=self.mesh, + dtype=self.dtype, + ) + weight_mappings = self._create_vision_weight_mappings() + + if self.mesh is not None: + with self.mesh: + loader.load_weights_from_safetensors(weight_mappings) + else: + loader.load_weights_from_safetensors(weight_mappings) + + logger.info("Qwen3-VL Vision weights loaded successfully!") + + def _create_vision_weight_mappings(self) -> dict: + """Create weight mappings for vision encoder.""" + mappings = {} + + # Text embedding (for multimodal fusion later) + mappings["model.embed_tokens.weight"] = WeightMapping( + target_path="text_embed.embedding", + sharding=(None, None), + transpose=False, + ) + + # Patch embedding (Conv3D) + mappings["model.visual.patch_embed.proj.weight"] = WeightMapping( + target_path="visual.patch_embed.proj.kernel", + sharding=(None, None, None, None, None), + transpose_axes=(2, 3, 4, 1, 0), + ) + mappings["model.visual.patch_embed.proj.bias"] = WeightMapping( + target_path="visual.patch_embed.proj.bias", + sharding=(None,), + transpose=False, + ) + + # Position embedding + mappings["model.visual.pos_embed.weight"] = WeightMapping( + target_path="visual.pos_embed.embedding", + sharding=(None, None), + transpose=False, + ) + + # Main merger + mappings["model.visual.merger.norm.weight"] = WeightMapping( + target_path="visual.merger.norm.scale", + sharding=(None,), + transpose=False, + ) + mappings["model.visual.merger.norm.bias"] = WeightMapping( + target_path="visual.merger.norm.bias", + sharding=(None,), + transpose=False, + ) + mappings["model.visual.merger.linear_fc1.weight"] = WeightMapping( + target_path="visual.merger.linear_fc1.kernel", + sharding=(None, None), + transpose=True, + ) + mappings["model.visual.merger.linear_fc1.bias"] = WeightMapping( + target_path="visual.merger.linear_fc1.bias", + sharding=(None,), + transpose=False, + ) + mappings["model.visual.merger.linear_fc2.weight"] = WeightMapping( + target_path="visual.merger.linear_fc2.kernel", + sharding=(None, None), + transpose=True, + ) + mappings["model.visual.merger.linear_fc2.bias"] = WeightMapping( + target_path="visual.merger.linear_fc2.bias", + sharding=(None,), + transpose=False, + ) + + # DeepStack mergers + for i in range(len(self.visual.deepstack_visual_indexes)): + prefix = f"model.visual.deepstack_merger_list.{i}" + target_prefix = f"visual.deepstack_merger_list.{i}" + + mappings[f"{prefix}.norm.weight"] = WeightMapping( + target_path=f"{target_prefix}.norm.scale", + sharding=(None,), + transpose=False, + ) + mappings[f"{prefix}.norm.bias"] = WeightMapping( + target_path=f"{target_prefix}.norm.bias", + sharding=(None,), + transpose=False, + ) + mappings[f"{prefix}.linear_fc1.weight"] = WeightMapping( + target_path=f"{target_prefix}.linear_fc1.kernel", + sharding=(None, None), + transpose=True, + ) + mappings[f"{prefix}.linear_fc1.bias"] = WeightMapping( + target_path=f"{target_prefix}.linear_fc1.bias", + sharding=(None,), + transpose=False, + ) + mappings[f"{prefix}.linear_fc2.weight"] = WeightMapping( + target_path=f"{target_prefix}.linear_fc2.kernel", + sharding=(None, None), + transpose=True, + ) + mappings[f"{prefix}.linear_fc2.bias"] = WeightMapping( + target_path=f"{target_prefix}.linear_fc2.bias", + sharding=(None,), + transpose=False, + ) + + # Vision blocks + num_layers = self.config.depth + for layer_idx in range(num_layers): + mappings.update(self._create_vision_block_mappings(layer_idx)) + + return mappings + + def _create_vision_block_mappings(self, layer_idx: int) -> dict: + """Create weight mappings for a single vision block.""" + prefix = f"model.visual.blocks.{layer_idx}" + target_prefix = f"visual.blocks.{layer_idx}" + + return { + # Layer norms + f"{prefix}.norm1.weight": WeightMapping( + target_path=f"{target_prefix}.norm1.scale", + sharding=(None,), + transpose=False, + ), + f"{prefix}.norm1.bias": WeightMapping( + target_path=f"{target_prefix}.norm1.bias", + sharding=(None,), + transpose=False, + ), + f"{prefix}.norm2.weight": WeightMapping( + target_path=f"{target_prefix}.norm2.scale", + sharding=(None,), + transpose=False, + ), + f"{prefix}.norm2.bias": WeightMapping( + target_path=f"{target_prefix}.norm2.bias", + sharding=(None,), + transpose=False, + ), + # QKV projection + f"{prefix}.attn.qkv.weight": WeightMapping( + target_path=f"{target_prefix}.attn.qkv_proj.kernel", + sharding=(None, None), + transpose=True, + ), + f"{prefix}.attn.qkv.bias": WeightMapping( + target_path=f"{target_prefix}.attn.qkv_proj.bias", + sharding=(None,), + transpose=False, + ), + # Output projection + f"{prefix}.attn.proj.weight": WeightMapping( + target_path=f"{target_prefix}.attn.proj.kernel", + sharding=(None, None), + transpose=True, + ), + f"{prefix}.attn.proj.bias": WeightMapping( + target_path=f"{target_prefix}.attn.proj.bias", + sharding=(None,), + transpose=False, + ), + # MLP + f"{prefix}.mlp.linear_fc1.weight": WeightMapping( + target_path=f"{target_prefix}.mlp.linear_fc1.kernel", + sharding=(None, None), + transpose=True, + ), + f"{prefix}.mlp.linear_fc1.bias": WeightMapping( + target_path=f"{target_prefix}.mlp.linear_fc1.bias", + sharding=(None,), + transpose=False, + ), + f"{prefix}.mlp.linear_fc2.weight": WeightMapping( + target_path=f"{target_prefix}.mlp.linear_fc2.kernel", + sharding=(None, None), + transpose=True, + ), + f"{prefix}.mlp.linear_fc2.bias": WeightMapping( + target_path=f"{target_prefix}.mlp.linear_fc2.bias", + sharding=(None,), + transpose=False, + ), + } + + def _validate_and_reshape_mm_tensor(self, mm_input: object, name: str) -> jax.Array: + if isinstance(mm_input, list): + arrays_to_concat = [jnp.asarray(item) for item in mm_input] + return jnp.concatenate(arrays_to_concat, axis=0) + + if hasattr(mm_input, "ndim"): + array_input = jnp.asarray(mm_input) + if array_input.ndim == 2: + return array_input + if array_input.ndim == 3: + return array_input.reshape(-1, array_input.shape[-1]) + + raise ValueError(f"Incorrect type of {name}. Got type: {type(mm_input)}") + + def _parse_and_validate_image_input( + self, + image_grid_thw: tuple[tuple[int, int, int], ...], + **kwargs: object, + ) -> Qwen3_VLImageInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + image_embeds = kwargs.pop("image_embeds", None) + + if pixel_values is None and image_embeds is None: + return None + + if pixel_values is not None: + pixel_values = self._validate_and_reshape_mm_tensor(pixel_values, "image pixel values") + return Qwen3_VLImagePixelInputs( + type="pixel_values", + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + ) + + return None + + def _parse_and_validate_multimodal_inputs( + self, + image_grid_thw: tuple[tuple[int, int, int], ...], + **kwargs: object, + ) -> dict: + mm_input_by_modality = {} + + for input_key in kwargs: + if ( + input_key in ("pixel_values", "image_embeds") + and "image" not in mm_input_by_modality + ): + mm_input_by_modality["image"] = self._parse_and_validate_image_input( + image_grid_thw, **kwargs + ) + return mm_input_by_modality + + def get_single_image_embedding( + self, + image_pixel_values: jax.Array, + image_grid_thw: tuple[int, int, int], + ) -> jax.Array: + hidden_states, _ = self.visual(image_pixel_values, (image_grid_thw,)) + return hidden_states + + def _process_image_input( + self, + image_input: Qwen3_VLImageInputs, + ) -> tuple[jax.Array, ...]: + grid_thw = image_input["image_grid_thw"] + + if image_input["type"] == "image_embeds": + image_embeds = image_input["image_embeds"].astype(self.dtype) + else: + pixel_values = image_input["pixel_values"] + image_embeds = [] + current_idx = 0 + for image_thw in grid_thw: + t, h, w = image_thw + image_size = t * h * w + end_idx = current_idx + image_size + image_pixel_values = pixel_values[current_idx:end_idx, :] + image_embeds.append(self.get_single_image_embedding(image_pixel_values, image_thw)) + current_idx = end_idx + image_embeds = jnp.concatenate(image_embeds, axis=0) + + merge_size = self.visual.config.spatial_merge_size + sizes = np.prod(np.array(grid_thw, dtype=np.int64), axis=-1) // merge_size // merge_size + + if sizes.size == 0: + return () + if sizes.size == 1: + return (image_embeds,) + + split_indices = np.cumsum(sizes)[:-1] + return tuple(jnp.split(image_embeds, split_indices)) + + def get_multimodal_embeddings( + self, + image_grid_thw: tuple[tuple[int, int, int], ...], + **kwargs: object, + ) -> list[jax.Array]: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(image_grid_thw, **kwargs) + if not mm_input_by_modality: + return [] + + multimodal_embeddings: tuple[jax.Array, ...] = () + + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "image": + vision_embeddings = self._process_image_input(multimodal_input) + multimodal_embeddings += vision_embeddings + + return list(multimodal_embeddings) + + def __call__( + self, + pixel_values: jax.Array, + image_grid_thw: tuple[tuple[int, int, int], ...] = None, + video_grid_thw: tuple[tuple[int, int, int], ...] = None, + ) -> jax.Array: + """Encode vision inputs to embeddings. + + Args: + pixel_values: Pixel values [num_patches, channels * patch_size^2] + image_grid_thw: Grid dimensions for each image + video_grid_thw: Grid dimensions for each video + + Returns: + Vision embeddings [total_patches, hidden_dim] + """ + combined_grid_thw = [] + if image_grid_thw: + combined_grid_thw.extend(image_grid_thw) + if video_grid_thw: + combined_grid_thw.extend(video_grid_thw) + + if not combined_grid_thw: + return jnp.zeros((0, self.config.hidden_size), dtype=pixel_values.dtype) + + combined_grid_thw = tuple(combined_grid_thw) + vision_embeds_list = self.get_multimodal_embeddings( + image_grid_thw=combined_grid_thw, + pixel_values=pixel_values, + ) + return jnp.concatenate(vision_embeds_list, axis=0) diff --git a/python/sgl_jax/srt/multimodal/models/qwen3_VL/run_inference.py b/python/sgl_jax/srt/multimodal/models/qwen3_VL/run_inference.py new file mode 100644 index 0000000000..2a8027dbbc --- /dev/null +++ b/python/sgl_jax/srt/multimodal/models/qwen3_VL/run_inference.py @@ -0,0 +1,431 @@ +#!/usr/bin/env python3 +""" +Qwen3-VL Standalone Inference Script (Option B) + +Demonstrates loading PyTorch safetensors weights into the sglang-jax Qwen3-VL +model and running inference using inputs processed by the HuggingFace processor. + +This script validates: +1. Weight loading for both vision encoder and text decoder +2. Vision encoding (pixel_values → vision embeddings) +3. Input preparation using AutoProcessor (PyTorch CPU tensors → JAX arrays) + +For full end-to-end generation (with KV cache, auto-regressive decoding), +use the sglang-jax server (Option A): + python -m sgl_jax --model-path Qwen/Qwen3-VL-2B-Instruct --multimodal + +Usage: + python -m sgl_jax.srt.multimodal.models.qwen3_VL.run_inference \ + --model-path Qwen/Qwen3-VL-2B-Instruct + + # With local weights: + python -m sgl_jax.srt.multimodal.models.qwen3_VL.run_inference \ + --model-path /path/to/qwen3-vl-2b + + # Vision encoding test with image: + python -m sgl_jax.srt.multimodal.models.qwen3_VL.run_inference \ + --model-path Qwen/Qwen3-VL-2B-Instruct \ + --image-url "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.jpeg" +""" + +import argparse +import json +import logging +import os +import time + +import jax +import jax.numpy as jnp +from flax import nnx + +logger = logging.getLogger(__name__) + + +def download_model(model_id: str) -> str: + """Download model from HuggingFace Hub. Returns local path.""" + from huggingface_hub import snapshot_download + + print(f" Downloading model: {model_id}") + local_path = snapshot_download(model_id) + print(f" Model downloaded to: {local_path}") + return local_path + + +def load_processor(model_path: str): + """Load the HuggingFace processor for tokenization and image processing.""" + from transformers import AutoProcessor + + return AutoProcessor.from_pretrained(model_path) + + +def detect_config_size(model_path: str) -> str: + """Detect model size from config.json.""" + config_path = os.path.join(model_path, "config.json") + if os.path.exists(config_path): + with open(config_path) as f: + config = json.load(f) + num_hidden_layers = config.get("num_hidden_layers", 28) + # 2B has 28 layers, 4B has 36, 8B has 36, 32B has 64 + hidden_size = config.get("hidden_size", 2048) + if num_hidden_layers <= 28: + return "2b" + elif hidden_size <= 2560: + return "4b" + elif hidden_size <= 4096: + return "8b" + else: + return "32b" + return "2b" + + +def create_model_config(model_path: str, dtype): + """Create Qwen3VLConfig for model initialization. + + Returns: + qwen3_config: Full Qwen3VLConfig (contains .vision_config and .text_config) + """ + from sgl_jax.srt.multimodal.configs.qwen_vl.qwen3_vl_config import Qwen3VLConfig + + size = detect_config_size(model_path) + print(f" Detected model size: {size}") + + size_to_config = { + "2b": Qwen3VLConfig.qwen3vl_2b, + "4b": Qwen3VLConfig.qwen3vl_4b, + "8b": Qwen3VLConfig.qwen3vl_8b, + "32b": Qwen3VLConfig.qwen3vl_32b, + } + return size_to_config[size]() + + +def load_vision_model(model_path: str, qwen3_config, mesh, dtype): + """Load the Qwen3-VL vision encoder. + + Args: + model_path: Path to model weights directory + qwen3_config: Full Qwen3VLConfig (we extract .vision_config from it) + mesh: JAX device mesh + dtype: Model dtype + """ + from sgl_jax.srt.multimodal.models.qwen3_VL.qwen3_vl_vit import ( + Qwen3_VL_VisionModel, + ) + + # Extract the vision sub-config — VisionModel expects Qwen3VLVisionConfig, + # NOT the full Qwen3VLConfig + vis_cfg = qwen3_config.vision_config + text_cfg = qwen3_config.text_config + + print( + f" Vision config: depth={vis_cfg.depth}, hidden={vis_cfg.hidden_size}, " + f"out_hidden={vis_cfg.out_hidden_size}" + ) + + print(" Initializing vision model...") + t0 = time.time() + + # Must use jax.set_mesh() (not `with mesh:`) — this is the pattern used by + # JAXModelLoader._get_model(). jax.set_mesh() propagates into eval_shape; + # plain `with mesh:` does not. + with jax.set_mesh(mesh): + vision_model = nnx.eval_shape( + lambda: Qwen3_VL_VisionModel( + config=vis_cfg, + dtype=dtype, + rngs=nnx.Rngs(0), + mesh=mesh, + ) + ) + + # load_weights expects a config object with .model_path, .vocab_size, + # .text_hidden_size (used to create the text_embed layer) + class VisionWeightConfig: + def __init__(self, path, text_config, vision_config, dt): + self.model_path = path + self.dtype = dt + self.vocab_size = text_config.vocab_size + self.text_hidden_size = vision_config.out_hidden_size + self.quantization_config = None + + vm_config = VisionWeightConfig(model_path, text_cfg, vis_cfg, dtype) + vision_model.load_weights(vm_config) + + print(f" Vision model loaded in {time.time() - t0:.2f}s") + return vision_model + + +def load_generation_model(model_path: str, qwen3_config, mesh, dtype): + """Load the Qwen3-VL text decoder. + + Args: + model_path: Path to model weights directory + qwen3_config: Full Qwen3VLConfig (contains .text_config) + mesh: JAX device mesh + dtype: Model dtype + """ + from sgl_jax.srt.multimodal.models.qwen3_VL.qwen3_vl_generation import ( + Qwen3_VL_Generation, + ) + + text_cfg = qwen3_config.text_config + print( + f" Text config: layers={text_cfg.num_hidden_layers}, " + f"hidden={text_cfg.hidden_size}, vocab={text_cfg.vocab_size}" + ) + + print(" Initializing generation model...") + t0 = time.time() + + with jax.set_mesh(mesh): + gen_model = nnx.eval_shape( + lambda: Qwen3_VL_Generation( + config=qwen3_config, + dtype=dtype, + mesh=mesh, + ) + ) + + # WeightLoader needs .model_path, .num_hidden_layers, .quantization_config, + # and several ModelConfig methods. Proxy simple attributes to text_config; + # add required methods explicitly (they don't exist on the dataclass). + class GenModelConfig: + def __init__(self, path, text_config, dt): + self.model_path = path + self.dtype = dt + self.quantization_config = None + self._text_config = text_config + + def get_total_num_kv_heads(self): + return self._text_config.num_key_value_heads + + def needs_kv_head_replication(self, tensor_parallel_size): + return tensor_parallel_size > self.get_total_num_kv_heads() + + def get_num_kv_head_replicas(self, tensor_parallel_size): + total = self.get_total_num_kv_heads() + if tensor_parallel_size > total: + return (tensor_parallel_size + total - 1) // total + return 1 + + def get_kv_padding_strategy(self): + # GQA (num_kv_heads < num_attention_heads) → replicate + if self._text_config.num_key_value_heads < self._text_config.num_attention_heads: + return "replicate" + return "zero" + + def __getattr__(self, name): + return getattr(self._text_config, name) + + gm_config = GenModelConfig(model_path, text_cfg, dtype) + gen_model.load_weights(gm_config) + + print(f" Generation model loaded in {time.time() - t0:.2f}s") + return gen_model + + +def run_text_test(processor, model_path: str, gen_model, mesh, dtype): + """Test text tokenization and embedding lookup.""" + print("\n--- Text Input Test ---") + messages = [ + {"role": "user", "content": [{"type": "text", "text": "What is the capital of France?"}]}, + ] + + inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt", + ) + + input_ids = jnp.array(inputs["input_ids"].numpy()) + print(f" Input IDs shape: {input_ids.shape}") + print(f" Input IDs (first 20): {input_ids[0, :20].tolist()}") + print(f" Total tokens: {input_ids.shape[1]}") + + # Test embedding lookup + embed_tokens = gen_model.model.embed_tokens + embeddings = embed_tokens(input_ids) + print(f" Embedding shape: {embeddings.shape}") + print(f" Embedding dtype: {embeddings.dtype}") + print(f" Embedding norm (first token): {float(jnp.linalg.norm(embeddings[0, 0])):.6f}") + print(" ✓ Text embedding lookup successful!") + + +def run_vision_test(processor, vision_model, image_url: str): + """Test vision encoding with an image.""" + print("\n--- Vision Encoding Test ---") + messages = [ + { + "role": "user", + "content": [ + {"type": "image", "image": image_url}, + {"type": "text", "text": "Describe this image."}, + ], + } + ] + + print(f" Processing image: {image_url}") + inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + return_tensors="pt", + ) + + input_ids = jnp.array(inputs["input_ids"].numpy()) + print(f" Input IDs shape: {input_ids.shape}") + + if "pixel_values" not in inputs: + print(" ✗ No pixel_values found in inputs. Vision processing may have failed.") + return + + pixel_values = jnp.array(inputs["pixel_values"].numpy()) + image_grid_thw = inputs["image_grid_thw"].numpy() + + print(f" Pixel values shape: {pixel_values.shape}") + print(f" Pixel values dtype: {pixel_values.dtype}") + print(f" Image grid THW: {image_grid_thw}") + + # Convert image_grid_thw to tuple format + grid_thw_tuple = tuple(tuple(x) for x in image_grid_thw.tolist()) + + # Run vision encoder + print(" Running vision encoder...") + t0 = time.time() + vision_embeddings = vision_model( + pixel_values=pixel_values, + image_grid_thw=grid_thw_tuple, + ) + jax.block_until_ready(vision_embeddings) + vision_time = time.time() - t0 + + print(f" Vision embeddings shape: {vision_embeddings.shape}") + print(f" Vision embeddings dtype: {vision_embeddings.dtype}") + print( + f" Vision embeddings norm (mean): {float(jnp.mean(jnp.linalg.norm(vision_embeddings, axis=-1))):.6f}" + ) + print(f" Vision encoding time: {vision_time:.2f}s") + print(" ✓ Vision encoding successful!") + + +def main(): + parser = argparse.ArgumentParser( + description="Qwen3-VL Standalone Inference - Weight Loading & Vision Encoding Test" + ) + parser.add_argument( + "--model-path", + type=str, + default="Qwen/Qwen3-VL-2B-Instruct", + help="HuggingFace model ID or local path", + ) + parser.add_argument( + "--dtype", + type=str, + default="bfloat16", + choices=["bfloat16", "float32"], + help="Model dtype", + ) + parser.add_argument( + "--image-url", + type=str, + default=None, + help="Image URL for vision encoding test", + ) + parser.add_argument( + "--skip-generation", + action="store_true", + help="Skip loading the generation model (only test vision encoder)", + ) + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + + dtype_map = {"bfloat16": jnp.bfloat16, "float32": jnp.float32} + dtype = dtype_map[args.dtype] + + print("=" * 60) + print("Qwen3-VL Standalone Inference (sglang-jax)") + print("=" * 60) + print(f"JAX devices: {jax.devices()}") + print(f"Device count: {jax.device_count()}") + print(f"Backend: {jax.default_backend()}") + + # Create mesh with axis names matching sglang-jax convention ("data", "tensor"). + # Use single device for standalone testing; multi-device is handled by Option A. + # axis_types must be Explicit (matching create_device_mesh in mesh_utils.py). + import numpy as np + + single_device = np.array(jax.devices()[:1]).reshape(1, 1) + mesh = jax.sharding.Mesh( + single_device, + axis_names=("data", "tensor"), + axis_types=(jax.sharding.AxisType.Explicit, jax.sharding.AxisType.Explicit), + ) + print(f"Mesh: {mesh}") + + # Resolve model path (download if needed) + model_path = args.model_path + if not os.path.exists(model_path): + print("\n1. Downloading model...") + model_path = download_model(model_path) + else: + print(f"\n1. Using local model: {model_path}") + + # Load configs + print("\n2. Loading configs...") + qwen3_config = create_model_config(model_path, dtype) + + # Load processor + print("\n3. Loading processor...") + processor = load_processor(model_path) + + # Load vision model + print("\n4. Loading vision model...") + vision_model = load_vision_model(model_path, qwen3_config, mesh, dtype) + + # Load generation model (optional) + gen_model = None + if not args.skip_generation: + print("\n5. Loading generation model...") + gen_model = load_generation_model(model_path, qwen3_config, mesh, dtype) + + # ========================================================================= + # Test 1: Text Input Processing & Embedding + # ========================================================================= + if gen_model is not None: + run_text_test(processor, model_path, gen_model, mesh, dtype) + + # ========================================================================= + # Test 2: Vision Encoding + # ========================================================================= + if args.image_url: + run_vision_test(processor, vision_model, args.image_url) + else: + print("\n--- Vision Test Skipped (no --image-url provided) ---") + print(" Pass --image-url to test vision encoding") + + # ========================================================================= + # Summary + # ========================================================================= + print("\n" + "=" * 60) + print("Summary") + print("=" * 60) + print(" ✓ Model weights loaded from safetensors") + print(" ✓ HuggingFace processor can prepare inputs") + if gen_model: + print(" ✓ Text decoder embeddings verified") + if args.image_url: + print(" ✓ Vision encoder produces embeddings") + print() + print("For full auto-regressive generation, use the sglang-jax server:") + print(f" python -m sgl_jax --model-path {args.model_path} --multimodal") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/python/sgl_jax/srt/multimodal/models/static_configs/qwen3_vl_stage_config.yaml b/python/sgl_jax/srt/multimodal/models/static_configs/qwen3_vl_stage_config.yaml new file mode 100644 index 0000000000..b583e82a09 --- /dev/null +++ b/python/sgl_jax/srt/multimodal/models/static_configs/qwen3_vl_stage_config.yaml @@ -0,0 +1,27 @@ +model_arch: Qwen3-VL +stage_args: + - stage_id: 0 + stage_sub_dir: "" + runtime: + num_tpus: 1 + max_batch_size: 1 + scheduler: vit + final_output: false + scheduler_params: {} + model_class: Qwen3_VL_VisionModel + + - stage_id: 1 + stage_sub_dir: "" + runtime: + num_tpus: 1 + max_batch_size: 1 + scheduler: auto_regressive + scheduler_params: {} + model_class: Qwen3_VL_Generation + final_output: true + +# stage0: ViT + # - Call encode_vision() to get vision embeddings + # - Call get_input_embeddings() to merge vision + text embeddings + +# Stage 1: prefill and decode diff --git a/python/sgl_jax/srt/multimodal/models/static_configs/qwen3_vl_stage_config_tp4.yaml b/python/sgl_jax/srt/multimodal/models/static_configs/qwen3_vl_stage_config_tp4.yaml new file mode 100644 index 0000000000..8e1b9836a1 --- /dev/null +++ b/python/sgl_jax/srt/multimodal/models/static_configs/qwen3_vl_stage_config_tp4.yaml @@ -0,0 +1,24 @@ +model_arch: Qwen3-VL +stage_args: + - stage_id: 0 + stage_sub_dir: "" + runtime: + num_tpus: 1 + max_batch_size: 1 + scheduler: vit + final_output: false + scheduler_params: {} + model_class: Qwen3_VL_VisionModel + + - stage_id: 1 + stage_sub_dir: "" + runtime: + num_tpus: 4 + max_batch_size: 1 + scheduler: auto_regressive + scheduler_params: {} + model_class: Qwen3_VL_Generation + final_output: true + +# stage0: ViT (single device) +# Stage 1: prefill and decode (TP=4 for larger models like 8B/32B) diff --git a/python/sgl_jax/srt/multimodal/models/static_configs/yaml_registry.py b/python/sgl_jax/srt/multimodal/models/static_configs/yaml_registry.py index dc2045daa3..5865125b44 100644 --- a/python/sgl_jax/srt/multimodal/models/static_configs/yaml_registry.py +++ b/python/sgl_jax/srt/multimodal/models/static_configs/yaml_registry.py @@ -39,12 +39,22 @@ class StageConfigRegistry: "Qwen2.5-VL-32B-Instruct": "qwen2_5_vl_stage_config_tp4.yaml", "Qwen/Qwen2.5-VL-72B-Instruct": "qwen2_5_vl_stage_config_tp4.yaml", "Qwen2.5-VL-72B-Instruct": "qwen2_5_vl_stage_config_tp4.yaml", + # Qwen3-VL series + "Qwen/Qwen3-VL-2B-Instruct": "qwen3_vl_stage_config.yaml", + "Qwen3-VL-2B-Instruct": "qwen3_vl_stage_config.yaml", + "Qwen/Qwen3-VL-4B-Instruct": "qwen3_vl_stage_config.yaml", + "Qwen3-VL-4B-Instruct": "qwen3_vl_stage_config.yaml", + "Qwen/Qwen3-VL-8B-Instruct": "qwen3_vl_stage_config_tp4.yaml", + "Qwen3-VL-8B-Instruct": "qwen3_vl_stage_config_tp4.yaml", + "Qwen/Qwen3-VL-32B-Instruct": "qwen3_vl_stage_config_tp4.yaml", + "Qwen3-VL-32B-Instruct": "qwen3_vl_stage_config_tp4.yaml", } # Keyword patterns for fallback matching (order matters - more specific first) _KEYWORD_PATTERNS: list[tuple[str, str]] = [ ("Wan2.2", "wan2_2_stage_config.yaml"), ("Wan2.1", "wan2_1_stage_config.yaml"), + ("Qwen3-VL", "qwen3_vl_stage_config.yaml"), ("Qwen2.5-VL", "qwen2_5_vl_stage_config.yaml"), ]