tpu-inference 0.11.1.dev202511180814__py3-none-any.whl → 0.12.0.dev20251213__py3-none-any.whl

tpu_inference/models/jax/llama_guard_4.py

@@ -0,0 +1,361 @@
+import re
+from typing import Any, List, Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+import torch
+from flax import nnx
+from flax.typing import PRNGKey
+from jax.sharding import Mesh
+from jax.sharding import PartitionSpec as P
+from vllm.config import VllmConfig
+
+from tpu_inference.layers.jax.attention.attention import AttentionMetadata
+from tpu_inference.layers.jax.attention.llama4_attention import Llama4Attention
+from tpu_inference.layers.jax.constants import KVCacheType
+from tpu_inference.layers.jax.layers import DenseFFW, Embedder, LMhead, RMSNorm
+from tpu_inference.layers.jax.misc import shard_put
+from tpu_inference.layers.jax.transformer_block import TransformerBlock
+from tpu_inference.logger import init_logger
+from tpu_inference.models.jax.utils.weight_utils import (
+    get_param, model_weights_generator, print_param_info, reshape_params,
+    transpose_params)
+
+logger = init_logger(__name__)
+
+
+class LlamaGuard4ForCausalLM(nnx.Module):
+
+    def __init__(self,
+                 vllm_config: VllmConfig,
+                 rng: PRNGKey,
+                 mesh: Mesh,
+                 force_random_weights: bool = False):
+        logger.warning(
+            "🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨\n"
+            "Llama Guard 4 (JAX) is WIP: Only the text modality is currently implemented.  "
+            "Multimodal inputs will fail.\n"
+            "🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨")
+        assert mesh is not None
+
+        self.vllm_config = vllm_config
+        self.vllm_config.model_config.dtype = torch.bfloat16
+        model_config = vllm_config.model_config
+        text_config = model_config.hf_config.text_config
+
+        self.mesh = mesh
+        self.is_verbose = getattr(self.vllm_config.additional_config,
+                                  "is_verbose", False)
+
+        self.use_qk_norm = getattr(text_config, "use_qk_norm", True)
+
+        vocab_size = model_config.get_vocab_size()
+        self.hidden_size = model_config.get_hidden_size()
+
+        self.dtype: jnp.dtype = jnp.bfloat16
+
+        self.num_layers: int = getattr(text_config, "num_layers", 48)
+        hidden_act: str = getattr(text_config, "hidden_act", "silu")
+
+        rms_norm_eps = getattr(text_config, "rms_norm_eps", 1e-5)
+        self.num_attention_heads = getattr(text_config, "num_attention_heads",
+                                           40)
+        self.num_key_value_heads = getattr(text_config, "num_key_value_heads",
+                                           8)
+        self.head_dim = getattr(text_config, "head_dim", 128)
+
+        intermediate_size = getattr(text_config, "intermediate_size", 8192)
+
+        self.rope_theta_text = getattr(text_config, "rope_theta", 500000.0)
+        self.rope_scaling = getattr(text_config, "rope_scaling")
+
+        self.rng = nnx.Rngs(rng)
+
+        self.embedder = Embedder(
+            vocab_size=vocab_size,
+            hidden_size=self.hidden_size,
+            dtype=self.dtype,
+            vd_sharding=(('data', 'model'), None),
+            rngs=self.rng,
+            random_init=force_random_weights,
+        )
+
+        self.layers = []
+
+        for i in range(self.num_layers):
+            use_attention_rope = True
+
+            custom_module = DenseFFW(dtype=self.dtype,
+                                     hidden_act=hidden_act,
+                                     hidden_size=self.hidden_size,
+                                     intermediate_size=intermediate_size,
+                                     random_init=force_random_weights,
+                                     rngs=self.rng,
+                                     df_sharding=P(None, 'model'),
+                                     fd_sharding=P('model', None),
+                                     activation_ffw_td=P('data', None))
+
+            attn = Llama4Attention(
+                hidden_size=self.hidden_size,
+                dtype=self.dtype,
+                num_attention_heads=self.num_attention_heads,
+                num_key_value_heads=self.num_key_value_heads,
+                head_dim=self.head_dim,
+                rope_theta=self.rope_theta_text,
+                rope_scaling={
+                    "scale_factor":
+                    self.rope_scaling["factor"],
+                    "low_freq_factor":
+                    self.rope_scaling["low_freq_factor"],
+                    "high_freq_factor":
+                    self.rope_scaling["high_freq_factor"],
+                    "original_max_position_embeddings":
+                    self.rope_scaling["original_max_position_embeddings"]
+                },
+                rngs=self.rng,
+                rope_input_ordering="interleaved",
+                # TODO (jacobplatin): we should refactor this to pass a dtype (or config) directly
+                kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+                temperature_tuning=True,
+                temperature_tuning_scale=0.1,
+                temperature_tuning_floor_scale=8192,
+                use_qk_norm=self.use_qk_norm,
+                attention_chunk_size=None if use_attention_rope else 8192,
+                mesh=self.mesh,
+                random_init=force_random_weights,
+                activation_attention_td=('data', 'model'),
+                activation_q_td=('data', 'model'),
+                query_tnh=P('data', 'model', None),
+                keyvalue_skh=P('data', 'model', None),
+                activation_attention_out_td=('data', 'model'),
+                attn_o_tnh=P('data', 'model', None),
+                dnh_sharding=(None, 'model', None),
+                dkh_sharding=(None, 'model', None),
+                nhd_sharding=('model', None, None),
+            )
+
+            pre_attention_norm = RMSNorm(
+                dims=self.hidden_size,
+                random_init=force_random_weights,
+                epsilon=rms_norm_eps,
+                rngs=self.rng,
+                activation_ffw_td=('data', None),
+                with_scale=True,
+                dtype=self.dtype,
+            )
+
+            pre_mlp_norm = RMSNorm(
+                dims=self.hidden_size,
+                activation_ffw_td=('data', None),
+                epsilon=rms_norm_eps,
+                rngs=self.rng,
+                with_scale=True,
+                dtype=self.dtype,
+                random_init=force_random_weights,
+            )
+
+            block = TransformerBlock(custom_module=custom_module,
+                                     attn=attn,
+                                     pre_attention_norm=pre_attention_norm,
+                                     pre_mlp_norm=pre_mlp_norm,
+                                     use_attention_rope=use_attention_rope)
+            self.layers.append(block)
+
+        self.final_norm = RMSNorm(
+            dims=self.hidden_size,
+            activation_ffw_td=P(),
+            epsilon=rms_norm_eps,
+            rngs=self.rng,
+            with_scale=True,
+            dtype=self.dtype,
+            random_init=force_random_weights,
+        )
+
+        self.lm_head = LMhead(vocab_size=vocab_size,
+                              hidden_size=self.hidden_size,
+                              dtype=self.dtype,
+                              rngs=self.rng,
+                              vd_sharding=(('data', 'model'), None),
+                              dv_sharding=(None, ('data', 'model')),
+                              random_init=force_random_weights)
+        if self.is_verbose:
+            self._print_model_architecture()
+
+    def _print_model_architecture(self):
+
+        logger.info("### Embedding ###")
+        nnx.display(self.embedder)
+
+        logger.info("\n### Layers ###")
+        for i, layer in enumerate(self.layers):
+            logger.info(f"\n--- Layer {i} ---")
+            nnx.display(layer)
+
+        logger.info("\n### LM Head ###")
+        nnx.display(self.lm_head)
+
+    def load_weights(self, rng: jax.Array, cache_dir: Optional[str] = None):
+        self.rng = nnx.Rngs(rng)
+
+        weight_loader = LlamaGuard4WeightLoader(
+            vllm_config=self.vllm_config,
+            hidden_size=self.hidden_size,
+            attn_heads=self.num_attention_heads,
+            num_key_value_heads=self.num_key_value_heads,
+            attn_head_dim=self.head_dim)
+        weight_loader.load_weights(self)
+
+    def __call__(
+        self,
+        kv_caches: List[jax.Array],
+        input_ids: jax.Array,
+        attention_metadata: AttentionMetadata,
+        inputs_embeds: Optional[jax.Array] = None,
+        layer_metadata_tuple: Optional[Tuple] = None,
+        lora_metadata: Optional[Any] = None,
+        *args,
+    ) -> Tuple[List[KVCacheType], jax.Array]:
+        is_prefill = False
+
+        if inputs_embeds is not None:
+            x_TD = inputs_embeds
+        elif input_ids is not None:
+            x_TD = self.embedder.encode(input_ids)
+        else:
+            raise ValueError(
+                "Cannot run forward pass: Both input_ids and inputs_embeds are None."
+            )
+
+        for (i, block) in enumerate(self.layers):
+            kv_cache = kv_caches[i]
+            new_kv_cache, x_TD = block(x_TD, is_prefill, kv_cache,
+                                       attention_metadata)
+            jax.block_until_ready(x_TD)
+            kv_caches[i] = new_kv_cache
+
+        final_activation_TD = self.final_norm(x_TD)
+
+        return kv_caches, final_activation_TD, []
+
+    def compute_logits(self, hidden_states: jax.Array) -> jax.Array:
+        logits_TV = jnp.dot(hidden_states,
+                            self.lm_head.input_embedding_table_DV.value)
+        return logits_TV
+
+    def get_input_embeddings(
+            self,
+            input_ids: jax.Array,
+            multimodal_embeddings: Optional[List[jax.Array]] = None
+    ) -> jax.Array:
+        """
+        Computes the embeddings for text input (used for input to fusion).
+        """
+        return self.embedder.encode(input_ids)
+
+
+class LlamaGuard4WeightLoader:
+
+    def __init__(self, vllm_config: VllmConfig, hidden_size, attn_heads,
+                 num_key_value_heads, attn_head_dim):
+        self.names_and_weights_generator = model_weights_generator(
+            model_name_or_path=vllm_config.model_config.model,
+            framework="flax",
+            filter_regex="language_model",
+            download_dir=vllm_config.load_config.download_dir)
+        self.is_verbose = getattr(vllm_config.additional_config, "is_verbose",
+                                  False)
+        self._transpose_map = {
+            "q_proj": (2, 0, 1),
+            "k_proj": (2, 0, 1),
+            "v_proj": (2, 0, 1),
+            "o_proj": (1, 2, 0),
+            "lm_head": (1, 0),
+            "feed_forward.down_proj": (1, 0),
+            "feed_forward.gate_proj": (1, 0),
+            "feed_forward.up_proj": (1, 0),
+            "mlp.down_proj": (1, 0),
+            "mlp.gate_proj": (1, 0),
+            "mlp.up_proj": (1, 0),
+        }
+        self._weight_shape_map = {
+            "q_proj": (attn_heads, attn_head_dim, hidden_size),
+            "k_proj": (num_key_value_heads, attn_head_dim, hidden_size),
+            "v_proj": (num_key_value_heads, attn_head_dim, hidden_size),
+            "o_proj": (hidden_size, attn_heads, attn_head_dim),
+        }
+
+        self._loaded_to_standardized_keys = {
+            "language_model.model.embed_tokens.weight":
+            "embedder.input_embedding_table_VD",
+            "language_model.lm_head.weight":
+            "lm_head.input_embedding_table_DV",
+            "language_model.model.norm.weight":
+            "final_norm.scale",
+            "language_model.model.layers.*.input_layernorm.weight":
+            "layers.*.pre_attention_norm.scale",
+            "language_model.model.layers.*.post_attention_layernorm.weight":
+            "layers.*.pre_mlp_norm.scale",
+            "language_model.model.layers.*.self_attn.q_proj.weight":
+            "layers.*.attn.kernel_q_proj_DNH",
+            "language_model.model.layers.*.self_attn.k_proj.weight":
+            "layers.*.attn.kernel_k_proj_DKH",
+            "language_model.model.layers.*.self_attn.v_proj.weight":
+            "layers.*.attn.kernel_v_proj_DKH",
+            "language_model.model.layers.*.self_attn.o_proj.weight":
+            "layers.*.attn.kernel_o_proj_NHD",
+            "language_model.model.layers.*.feed_forward.gate_proj.weight":
+            "layers.*.custom_module.kernel_gating_DF",
+            "language_model.model.layers.*.feed_forward.up_proj.weight":
+            "layers.*.custom_module.kernel_up_proj_DF",
+            "language_model.model.layers.*.feed_forward.down_proj.weight":
+            "layers.*.custom_module.kernel_down_proj_FD",
+        }
+
+    def map_loaded_to_standardized_name(self, loaded_key: str) -> str:
+        if "layer" in loaded_key:
+            layer_num = re.search(r"layers\.(\d+)", loaded_key).group(1)
+            layer_key = re.sub(r"layers\.\d+", "layers.*", loaded_key)
+            mapped_key = self._loaded_to_standardized_keys.get(
+                layer_key, loaded_key)
+            mapped_key = re.sub(r"layers\.\*", f"layers.{layer_num}",
+                                mapped_key)
+        else:
+            mapped_key = self._loaded_to_standardized_keys.get(
+                loaded_key, loaded_key)
+        return mapped_key
+
+    def load_weights(self, model_for_loading: nnx.Module):
+        model_params = nnx.state(model_for_loading)
+        with jax.default_device(jax.devices("cpu")[0]):
+            for loaded_name, loaded_weight in self.names_and_weights_generator:
+                if loaded_name.endswith(".bias"):
+                    continue
+                if "vision_model" in loaded_name or "multi_modal_projector" in loaded_name:
+                    continue
+
+                mapped_name = self.map_loaded_to_standardized_name(loaded_name)
+                model_weight = get_param(model_params, mapped_name)
+
+                if not loaded_name.endswith(".bias"):
+                    # For other layers, continue to use the transpose_params helper.
+                    loaded_weight = reshape_params(loaded_name, loaded_weight,
+                                                   self._weight_shape_map)
+                    loaded_weight = transpose_params(loaded_name,
+                                                     loaded_weight,
+                                                     self._transpose_map)
+                if model_weight.value.shape != loaded_weight.shape:
+                    raise ValueError(
+                        f"Loaded shape for {loaded_name}: {loaded_weight.shape} "
+                        f"does not match model shape for {mapped_name}: {model_weight.value.shape}!"
+                    )
+                logger.debug(
+                    f"Transformed parameter {loaded_name} to {mapped_name}: {loaded_weight.shape} --> {model_weight.value.shape}"
+                )
+
+                model_weight.value = shard_put(loaded_weight,
+                                               model_weight.sharding,
+                                               mesh=model_for_loading.mesh)
+                if self.is_verbose:
+                    print_param_info(model_weight, loaded_name)
+
+        nnx.update(model_for_loading, model_params)

tpu_inference/models/jax/qwen2.py

@@ -368,7 +368,8 @@ class Qwen2ForCausalLM(nnx.Module):
                 "lm_head": "model.lm_head",
             })
 
-        metadata_map = get_default_maps(self.vllm_config, self.mesh, mappings)
+        metadata_map = get_default_maps(self.vllm_config.model_config,
+                                        self.mesh, mappings)
         load_hf_weights(vllm_config=self.vllm_config,
                         model=self,
                         metadata_map=metadata_map,

tpu_inference/models/jax/qwen2_5_vl.py

@@ -486,6 +486,11 @@ class Qwen2_5_VisionTransformer(nnx.Module):
             dtype=dtype,
             rngs=rngs)
 
+        additional_config = getattr(vllm_config, "additional_config",
+                                    None) or {}
+        self.enable_dynamic_image_sizes = additional_config.get(
+            "enable_dynamic_image_sizes", False)
+
     def rotary_pos_emb_thw(self, t, h, w):
         hpos_ids, wpos_ids = jnp.indices((h, w))
         hpos_ids = hpos_ids.reshape(
@@ -579,21 +584,7 @@ class Qwen2_5_VisionTransformer(nnx.Module):
         seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
         return max_seqlen, seqlens
 
-    def __call__(self, x: jax.Array, grid_thw: tuple[tuple[int, int,
-                                                           int]]) -> jax.Array:
-        # x: pixel_values: jax.Array
-        # """Shape:
-        # `(num_patches, num_channels * patch_size * patch_size)`
-        # """
-
-        # grid_thw: image_grid_thw: jax.Array
-        # """Shape: `(num_images, 3)`
-        # This should be in `(grid_t, grid_h, grid_w)` format.
-        # """
-        hidden_states = self.patch_embed(x)
-
-        # num of patches
-        seq_len = x.shape[0]
+    def compute_aux_arrays(self, grid_thw: tuple[tuple[int, int, int]]):
         # num of images/videoes
         num_grids = len(grid_thw)
 
@@ -638,6 +629,42 @@ class Qwen2_5_VisionTransformer(nnx.Module):
         cu_seqlens = jnp.pad(cu_seqlens, ((1, 0), ),
                              mode='constant',
                              constant_values=0)
+        return window_index, rotary_pos_emb, cu_seqlens, cu_window_seqlens
+
+    def pad_inputs(self, x, window_index, rotary_pos_emb, cu_seqlens,
+                   cu_window_seqlens):
+        # padding
+        num_patches = int(rotary_pos_emb.shape[0])
+        bucket_num_patches = 1 << (num_patches - 1).bit_length()
+        num_tokens = window_index.shape[0]
+        bucket_num_tokens = bucket_num_patches // self.spatial_merge_unit
+        vit_merger_window_size = (self.window_size //
+                                  self.spatial_merge_size // self.patch_size)
+        max_windows = (bucket_num_tokens // vit_merger_window_size) + 2
+
+        rotary_pos_emb = jnp.pad(rotary_pos_emb,
+                                 ((0, bucket_num_patches - num_patches),
+                                  (0, 0)))
+        window_index = jnp.concatenate([
+            window_index,
+            jnp.arange(num_tokens, bucket_num_tokens, dtype=jnp.int32)
+        ])
+        cu_window_seqlens = jnp.append(cu_window_seqlens, bucket_num_patches)
+        pad_w = max(0, max_windows + 1 - cu_window_seqlens.shape[0])
+        cu_window_seqlens = jnp.pad(cu_window_seqlens, (0, pad_w), mode='edge')
+        cu_seqlens = jnp.append(cu_seqlens, bucket_num_patches)
+
+        x_padded = jnp.pad(x, ((0, bucket_num_patches - x.shape[0]), (0, 0)))
+
+        return x_padded, window_index, rotary_pos_emb, cu_seqlens, cu_window_seqlens, num_tokens
+
+    def compute_hidden_states(self, x: jax.Array, window_index: jax.Array,
+                              rotary_pos_emb: jax.Array, cu_seqlens: jax.Array,
+                              cu_window_seqlens: jax.Array) -> jax.Array:
+        hidden_states = self.patch_embed(x)
+
+        # num of patches
+        seq_len = x.shape[0]
 
         hidden_states = hidden_states.reshape(
             seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
@@ -664,6 +691,48 @@ class Qwen2_5_VisionTransformer(nnx.Module):
         hidden_states = hidden_states[reverse_indices, :]
         return hidden_states
 
+    @jax.jit
+    def encode_padded_jit(self, x_padded, window_index, rotary_pos_emb,
+                          cu_seqlens, cu_window_seqlens):
+        return self.compute_hidden_states(x_padded, window_index,
+                                          rotary_pos_emb, cu_seqlens,
+                                          cu_window_seqlens)
+
+    @partial(
+        jax.jit,
+        static_argnames=("grid_thw", ),
+    )
+    def encode_jit(self, x, grid_thw):
+        window_index, rotary_pos_emb, cu_seqlens, cu_window_seqlens = self.compute_aux_arrays(
+            grid_thw)
+        return self.compute_hidden_states(x, window_index, rotary_pos_emb,
+                                          cu_seqlens, cu_window_seqlens)
+
+    def __call__(self, x: jax.Array, grid_thw: tuple[tuple[int, int,
+                                                           int]]) -> jax.Array:
+        # x: pixel_values: jax.Array
+        # """Shape:
+        # `(num_patches, num_channels * patch_size * patch_size)`
+        # """
+
+        # grid_thw: image_grid_thw: jax.Array
+        # """Shape: `(num_images, 3)`
+        # This should be in `(grid_t, grid_h, grid_w)` format.
+        # """
+        if self.enable_dynamic_image_sizes:
+            window_index, rotary_pos_emb, cu_seqlens, cu_window_seqlens = self.compute_aux_arrays(
+                grid_thw)
+            x_padded, window_index, rotary_pos_emb, cu_seqlens, cu_window_seqlens, num_tokens = self.pad_inputs(
+                x, window_index, rotary_pos_emb, cu_seqlens, cu_window_seqlens)
+
+            hidden_states = self.encode_padded_jit(x_padded, window_index,
+                                                   rotary_pos_emb, cu_seqlens,
+                                                   cu_window_seqlens)
+            return hidden_states[:num_tokens]
+
+        else:
+            return self.encode_jit(x, grid_thw)
+
 
 class Qwen2_5_VLForConditionalGeneration(nnx.Module):
 
@@ -888,10 +957,6 @@ class Qwen2_5_VLForConditionalGeneration(nnx.Module):
             #         "video"] = self._parse_and_validate_video_input(**kwargs)
         return mm_input_by_modality
 
-    @partial(
-        jax.jit,
-        static_argnames=("image_grid_thw", ),
-    )
     def get_single_image_embedding(self, image_pixel_values, image_grid_thw):
         return self.visual(image_pixel_values, (image_grid_thw, ))
 
@@ -1061,7 +1126,8 @@ class Qwen2_5_VLForConditionalGeneration(nnx.Module):
                 "lm_head": "language_model.model.lm_head",
             })
 
-        metadata_map = get_default_maps(self.vllm_config, self.mesh, mappings)
+        metadata_map = get_default_maps(self.vllm_config.model_config,
+                                        self.mesh, mappings)
         load_hf_weights(vllm_config=self.vllm_config,
                         model=self,
                         metadata_map=metadata_map,
@@ -1071,33 +1137,82 @@ class Qwen2_5_VLForConditionalGeneration(nnx.Module):
         self,
         run_compilation_fn: Callable,
     ) -> None:
-        image_shapes = []
-        if (warmup_config := self.vllm_config.additional_config.get(
-                "vision_warmup_config")):
-            image_shapes = warmup_config.get("image_shapes")
-
         vc = self.vllm_config.model_config.hf_config.vision_config
-        factor = vc.patch_size * vc.spatial_merge_size
-        for input_hw in image_shapes:
-            if not isinstance(input_hw, list) or len(input_hw) != 2:
-                logger.warning(f"Skipping invalid shape {input_hw}.")
-                continue
-            h_input, w_input = input_hw
-            h_processed = round(h_input / factor) * factor
-            w_processed = round(w_input / factor) * factor
-            t, h, w = 1, h_processed // vc.patch_size, w_processed // vc.patch_size
-            grid_thw = (t, h, w)
-            num_patches = t * h * w
-            patch_input_dim = vc.in_channels * vc.temporal_patch_size * vc.patch_size * vc.patch_size
-
-            dummy_pixel_values = jnp.ones(
-                (num_patches, patch_input_dim),
-                self.vllm_config.model_config.dtype,
-            )
-            dummy_grid_thw = grid_thw
+        patch_input_dim = vc.in_channels * vc.temporal_patch_size * vc.patch_size * vc.patch_size
+        if self.visual.enable_dynamic_image_sizes:
+            spatial_merge_unit = vc.spatial_merge_size**2
+            max_num_batched_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
+            mm_kwargs = self.vllm_config.model_config.multimodal_config.mm_processor_kwargs or {}
+            limit_pixels = float(mm_kwargs.get("max_pixels", float('inf')))
+
+            max_patches = int(
+                min(max_num_batched_tokens * spatial_merge_unit,
+                    limit_pixels / (vc.patch_size**2)))
+
+            num_patches_paddings = [
+                1 << i for i in range(4, (max_patches - 1).bit_length() + 1)
+            ]
+            rotary_dim = vc.hidden_size // vc.num_heads // 2
+            vit_merger_window_size = (vc.window_size //
+                                      vc.spatial_merge_size // vc.patch_size)
+
+            for num_patches in num_patches_paddings:
+                dummy_x_padded = jnp.ones(
+                    (num_patches, patch_input_dim),
+                    dtype=self.vllm_config.model_config.dtype)
+
+                num_tokens = num_patches // spatial_merge_unit
+                dummy_window_index = jnp.arange(num_tokens, dtype=jnp.int32)
+
+                dummy_rotary_pos_emb = jnp.ones(
+                    (num_patches, rotary_dim),
+                    dtype=self.vllm_config.model_config.dtype)
+
+                dummy_cu_seqlens = jnp.array([0, num_patches, num_patches],
+                                             dtype=jnp.int32)
+
+                max_windows = (num_tokens // vit_merger_window_size) + 2
+                patches_per_window = (vit_merger_window_size**
+                                      2) * spatial_merge_unit
+                dummy_cu_window_seqlens = jnp.arange(
+                    max_windows + 1, dtype=jnp.int32) * patches_per_window
+                dummy_cu_window_seqlens = jnp.minimum(dummy_cu_window_seqlens,
+                                                      num_patches)
+
+                run_compilation_fn("vision_encoder_padded",
+                                   self.visual.encode_padded_jit,
+                                   dummy_x_padded,
+                                   dummy_window_index,
+                                   dummy_rotary_pos_emb,
+                                   dummy_cu_seqlens,
+                                   dummy_cu_window_seqlens,
+                                   num_patches=num_patches)
+        else:
+            image_shapes = []
+            if (warmup_config := self.vllm_config.additional_config.get(
+                    "vision_warmup_config")):
+                image_shapes = warmup_config.get("image_shapes")
+
+            factor = vc.patch_size * vc.spatial_merge_size
+            for input_hw in image_shapes:
+                if not isinstance(input_hw, list) or len(input_hw) != 2:
+                    logger.warning(f"Skipping invalid shape {input_hw}.")
+                    continue
+                h_input, w_input = input_hw
+                h_processed = round(h_input / factor) * factor
+                w_processed = round(w_input / factor) * factor
+                t, h, w = 1, h_processed // vc.patch_size, w_processed // vc.patch_size
+                grid_thw = (t, h, w)
+                num_patches = t * h * w
+
+                dummy_pixel_values = jnp.ones(
+                    (num_patches, patch_input_dim),
+                    self.vllm_config.model_config.dtype,
+                )
+                dummy_grid_thw = (grid_thw, )
 
-            run_compilation_fn("single_image_encoder",
-                               self.get_single_image_embedding,
-                               dummy_pixel_values,
-                               dummy_grid_thw,
-                               image_shape=input_hw)
+                run_compilation_fn("vision_encoder",
+                                   self.visual.encode_jit,
+                                   dummy_pixel_values,
+                                   dummy_grid_thw,
+                                   image_shape=input_hw)

tpu_inference/models/jax/qwen3.py

@@ -295,7 +295,8 @@ class Qwen3ForCausalLM(nnx.Module):
                 "lm_head": "model.lm_head",
             })
 
-        metadata_map = get_default_maps(self.vllm_config, self.mesh, mappings)
+        metadata_map = get_default_maps(self.vllm_config.model_config,
+                                        self.mesh, mappings)
         load_hf_weights(vllm_config=self.vllm_config,
                         model=self,
                         metadata_map=metadata_map,

tpu_inference/models/jax/utils/quantization/quantization_utils.py

@@ -154,12 +154,9 @@ def qwix_quantize_nnx_model(model: nnx.Module, qwix_config: List[dict],
     logger.info(f"Memory usage before applying quantization of params: "
                 f"hbm={utils.hbm_usage_gb(jax.local_devices())}Gb")
 
-    # TODO (jacobplatin): we should refactor this to pass a dtype (or config) directly
-    kv_cache_jnp_dtype = utils.get_jax_dtype_from_str_dtype(kv_cache_dtype)
-
-    # Handle the case where kv_cache_dtype is "auto"
-    if kv_cache_jnp_dtype is None:
-        assert kv_cache_dtype == "auto", "kv_cache_dtype must be 'auto' if kv_cache_jnp_dtype is None"
+    if kv_cache_dtype != "auto":
+        kv_cache_jnp_dtype = utils.to_jax_dtype(kv_cache_dtype)
+    else:
         kv_cache_jnp_dtype = DEFAULT_KV_CACHE_DTYPE
 
     kv_caches = create_kv_caches(
@@ -169,9 +166,11 @@ def qwix_quantize_nnx_model(model: nnx.Module, qwix_config: List[dict],
         head_size=kv_cache_head_size,
         mesh=mesh,
         layer_names=[f"layer.{i}" for i in range(num_hidden_layers)],
-        cache_dtype=kv_cache_jnp_dtype)
+        cache_dtype=kv_cache_jnp_dtype,
+        use_mla=model.vllm_config.model_config.use_mla,
+    )
 
-    dp_size = mesh.shape.get("data", 1) * mesh.shape.get("attn", 1)
+    dp_size = model.vllm_config.sharding_config.total_dp_size
 
     # NOTE: the inputs don't need to match the actual ones, as long as the consumed weights are the same
     input_ids = jax.random.randint(rng,