tpu-inference 0.11.1.dev202511130813__py3-none-any.whl → 0.11.1.dev202511220812__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

@@ -8,8 +8,8 @@ from transformers import Qwen2Config, modeling_flax_utils
 from vllm.config import VllmConfig
 
 from tpu_inference import utils
+from tpu_inference.layers.common.attention_interface import attention
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
-from tpu_inference.layers.jax.attention_interface import attention
 from tpu_inference.layers.jax.rope_interface import apply_rope
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.utils.weight_utils import (get_default_maps,
@@ -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

@@ -14,9 +14,9 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
 from vllm.config import VllmConfig
 
 from tpu_inference import utils as utils
-from tpu_inference.layers.common.attention_metadata import AttentionMetadata
-from tpu_inference.layers.jax.attention_interface import \
+from tpu_inference.layers.common.attention_interface import \
     sharded_flash_attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.qwen2 import Qwen2ForCausalLM
 # from vllm.model_executor.models.interfaces import MultiModalEmbeddings
@@ -1061,7 +1061,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,

tpu_inference/models/jax/qwen3.py

@@ -8,8 +8,8 @@ from transformers import Qwen3Config
 from vllm.config import VllmConfig
 
 from tpu_inference import utils
+from tpu_inference.layers.common.attention_interface import attention
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
-from tpu_inference.layers.jax.attention_interface import attention
 from tpu_inference.layers.jax.rope_interface import apply_rope
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.qwen2 import Qwen2DecoderLayer
@@ -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/weight_utils.py

@@ -18,7 +18,7 @@ from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 from safetensors import safe_open
 
-from tpu_inference import utils
+from tpu_inference import envs, utils
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.utils import file_utils
 
@@ -197,12 +197,11 @@ def shard_put(x: jax.Array, shardings, mesh: jax.sharding.Mesh) -> jax.Array:
         return jax.device_put(x, shardings)
 
 
-def get_default_maps(vllm_config, mesh: Mesh,
+def get_default_maps(model_config, mesh: Mesh,
                      name_map: dict[str, str]) -> MetadataMap:
     """Load weights from one model weights file to the model, run on single thread."""
     sharding_size = mesh.shape["model"]
 
-    model_config = vllm_config.model_config
     hf_config = model_config.hf_config
 
     num_heads = hf_config.num_attention_heads
@@ -273,7 +272,8 @@ def _load_hf_weights_on_thread(vllm_config,
                                weights_file: str,
                                filter_regex: str | None = None,
                                keep_original_dtype_keys_regex: list[str]
-                               | None = None):
+                               | None = None,
+                               exclude_regex: list[str] | None = None):
     name_map = metadata_map.name_map
     reshape_keys = metadata_map.reshape_map
     bias_reshape_keys = metadata_map.bias_reshape_map
@@ -298,6 +298,18 @@ def _load_hf_weights_on_thread(vllm_config,
     for hf_key, hf_weight in model_weights_single_file_generator(
             weights_file, framework="flax", filter_regex=filter_regex):
 
+        # Check if the key should be excluded
+        if exclude_regex:
+            should_exclude = False
+            for pattern in exclude_regex:
+                if re.search(pattern, hf_key):
+                    logger.info(
+                        f"Excluding {hf_key} based on pattern {pattern}")
+                    should_exclude = True
+                    break
+            if should_exclude:
+                continue
+
         # Check if the key should retain its original dtype
         keep_original_dtype = False
         if keep_original_dtype_keys_regex:
@@ -408,7 +420,8 @@ def load_hf_weights(vllm_config,
                     mesh: Mesh,
                     filter_regex: str | None = None,
                     is_draft_model: bool = False,
-                    keep_original_dtype_keys_regex: list[str] | None = None):
+                    keep_original_dtype_keys_regex: list[str] | None = None,
+                    exclude_regex: list[str] | None = None):
     """Load weights from all model weights files to the model, run in multi threads."""
     if is_draft_model:
         model_path = vllm_config.speculative_config.draft_model_config.model
@@ -421,7 +434,7 @@ def load_hf_weights(vllm_config,
     # NOTE(xiang): Disable multi-threading mode if running on multi-host.
     # Because multi-threading would cause different JAX processes to load
     # different weights at the same time.
-    if os.environ.get("TPU_MULTIHOST_BACKEND", "").lower() == "ray":
+    if envs.TPU_MULTIHOST_BACKEND == "ray":
         max_workers = 1
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
         futures = [
@@ -433,8 +446,8 @@ def load_hf_weights(vllm_config,
                 mesh,
                 weights_file,
                 filter_regex=filter_regex,
-                keep_original_dtype_keys_regex=keep_original_dtype_keys_regex)
-            for weights_file in weights_files
+                keep_original_dtype_keys_regex=keep_original_dtype_keys_regex,
+                exclude_regex=exclude_regex) for weights_file in weights_files
         ]
         for future in futures:
             future.result()

tpu_inference/models/vllm/vllm_model_wrapper.py

@@ -25,6 +25,8 @@ from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.layers.vllm.quantization import get_tpu_quantization_config
 from tpu_inference.layers.vllm.sharding import shard_model_to_tpu
 from tpu_inference.logger import init_logger
+from tpu_inference.models.jax.jax_intermediate_tensor import \
+    JaxIntermediateTensors
 from tpu_inference.models.vllm.vllm_model_wrapper_context import (
     get_vllm_model_wrapper_context, set_vllm_model_wrapper_context)
 from tpu_inference.runner.lora_utils import replace_lora_metadata
@@ -89,13 +91,14 @@ class VllmModelWrapper:
             slice_config = self.vllm_config.device_config.slice
             modified_slice_config = True
             self.vllm_config.device_config.slice = None
+        self.vllm_config.compilation_config.static_forward_context.clear()
+
         vllm_config_for_load = copy.deepcopy(self.vllm_config)
         if modified_slice_config:
             self.vllm_config.device_config.slice = slice_config
         assert self.vllm_config.model_config.dtype in TORCH_DTYPE_TO_JAX, "The model_config.dtype must be a PyTorch dtype."
         vllm_config_for_load.device_config.device = "cpu"
         # Clearing the cached compilation config, otherwise vllm model init will fail
-        vllm_config_for_load.compilation_config.static_forward_context.clear()
 
         # When expert parallelism is enabled, vLLM loads weight in sharding
         # aware manner. Since tpu-inference has its own sharding logic, this
@@ -117,7 +120,7 @@ class VllmModelWrapper:
 
         # Load the vLLM model and wrap it into a new model whose forward
         # function can calculate the hidden_state and logits.
-        with load_context:
+        with load_context, jax.default_device(jax.devices("cpu")[0]):
             vllm_model = vllm_get_model(vllm_config=vllm_config_for_load)
         lora_manager = None
         if vllm_config_for_load.lora_config is not None:
@@ -149,7 +152,8 @@ class VllmModelWrapper:
                 "xla_tpu_reduce_scatter_collective_matmul_mode":
                 "post_spmd_conservative"
             },
-            static_argnames=("layer_name_to_kvcache_index", ),
+            static_argnames=("layer_name_to_kvcache_index", "is_first_rank",
+                             "is_last_rank"),
         )
         def step_fun(
             params_and_buffers,  # This has been wrapped into torchax TorchValue
@@ -157,8 +161,12 @@ class VllmModelWrapper:
             input_ids: jax.Array,
             attn_metadata: AttentionMetadata,
             input_embeds: jax.Array,
+            input_positions: jax.Array,
             layer_name_to_kvcache_index: Sequence[Tuple[str, int]],
             lora_metadata,
+            intermediate_tensors: JaxIntermediateTensors = None,
+            is_first_rank: bool = True,
+            is_last_rank: bool = True,
             *args,
         ) -> Tuple[List[jax.Array], jax.Array]:
             layer_name_to_kvcache_index = dict(layer_name_to_kvcache_index)
@@ -173,12 +181,14 @@ class VllmModelWrapper:
                 # torch_view in order to call the Torch function.
                 original_lora_metadata = replace_lora_metadata(
                     self.model, lora_metadata, self.vllm_config.lora_config)
-                hidden_states = torch.func.functional_call(
+                if not is_first_rank:
+                    intermediate_tensors = intermediate_tensors.to_torch()
+                output_from_torch = torch.func.functional_call(
                     self.model,
                     torch_view(params_and_buffers),
                     kwargs={
                         "input_ids": torch_view(input_ids),
-                        "positions": torch_view(attn_metadata.input_positions),
+                        "positions": torch_view(input_positions),
                         "intermediate_tensors": None,
                         "inputs_embeds": None,
                     },
@@ -188,11 +198,13 @@ class VllmModelWrapper:
                                       self.vllm_config.lora_config)
                 vllm_model_wrapper_context = get_vllm_model_wrapper_context()
                 new_kv_caches = vllm_model_wrapper_context.kv_caches
-            # Wrap the hidden_states from torch land into a JaxValue for the jax
-            # code to consume.
-            hidden_states = jax_view(hidden_states)
-
-            return new_kv_caches, hidden_states, []
+            # Wrap the output(hidden states or intermediate tensor)
+            # from torch land into a JaxValue for the jax code to consume.
+            if not is_last_rank:
+                output = JaxIntermediateTensors.from_torch(output_from_torch)
+            else:
+                output = jax_view(output_from_torch)
+            return new_kv_caches, output, []
 
         return step_fun
 

tpu_inference/platforms/tpu_platform.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import os
-from typing import TYPE_CHECKING, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Union, cast
 
 import jax.numpy as jnp
 import vllm.envs as vllm_envs
@@ -12,7 +11,7 @@ from vllm.platforms.interface import Platform, PlatformEnum
 from vllm.sampling_params import SamplingParams, SamplingType
 
 from tpu_inference import envs
-from tpu_inference.layers.jax.sharding import ShardingConfigManager
+from tpu_inference.layers.common.sharding import ShardingConfigManager
 from tpu_inference.logger import init_logger
 
 if TYPE_CHECKING:
@@ -57,7 +56,8 @@ class TpuPlatform(Platform):
     def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
                              dtype: jnp.dtype, kv_cache_dtype: Optional[str],
                              block_size: int, use_v1: bool, use_mla: bool,
-                             has_sink: bool, use_sparse: bool) -> str:
+                             has_sink: bool, use_sparse: bool,
+                             attn_type: Any) -> str:
         from vllm.attention.backends.registry import _Backend
         if selected_backend != _Backend.PALLAS:
             logger.info("Cannot use %s backend on TPU.", selected_backend)
@@ -182,10 +182,16 @@ class TpuPlatform(Platform):
         parallel_config.worker_cls = \
                         "tpu_inference.worker.tpu_worker.TPUWorker"
 
-        multihost_backend = os.environ.get("TPU_MULTIHOST_BACKEND", "").lower()
+        multihost_backend = envs.TPU_MULTIHOST_BACKEND
         if not multihost_backend:  # Single host
-            logger.info("Force using UniProcExecutor for JAX on single host.")
-            parallel_config.distributed_executor_backend = "uni"
+            if parallel_config.pipeline_parallel_size == 1:
+                logger.info("Force using UniProcExecutor for JAX on \
+                        single host without pipeline parallelism.")
+                parallel_config.distributed_executor_backend = "uni"
+            else:
+                logger.info("Force using MultiprocExecutor for JAX on \
+                        single host with pipeline parallelism.")
+                parallel_config.distributed_executor_backend = "mp"
         elif multihost_backend == "ray":
             from tpu_inference.executors.ray_distributed_executor import \
                 RayDistributedExecutor
@@ -260,3 +266,7 @@ class TpuPlatform(Platform):
         Returns if the current platform needs to sync weight loader.
         """
         return True
+
+    @classmethod
+    def support_hybrid_kv_cache(cls) -> bool:
+        return True