tpu-inference 0.11.1.dev202511220812__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tpu_inference/runner/compilation_manager.py

@@ -1,13 +1,27 @@
-import os
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import time
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
 import numpy as np
-import vllm.envs as envs
+import vllm.envs as vllm_envs
 from jax.sharding import NamedSharding, PartitionSpec
 
+import tpu_inference.envs as envs
 from tpu_inference.core.disagg_utils import is_disagg_enabled
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.layers.common.sharding import ShardingAxisName
@@ -15,6 +29,8 @@ from tpu_inference.layers.jax.sample.sampling import sample
 from tpu_inference.layers.jax.sample.sampling_metadata import \
     TPUSupportedSamplingMetadata
 from tpu_inference.logger import init_logger
+from tpu_inference.models.jax.jax_intermediate_tensor import \
+    JaxIntermediateTensors
 from tpu_inference.utils import device_array
 
 if TYPE_CHECKING:
@@ -30,10 +46,12 @@ class CompilationManager:
 
     def __init__(self, runner: "TPUModelRunner"):
         self.runner = runner
-        if not envs.VLLM_DISABLE_COMPILE_CACHE:
+        self._sampling_precompiled = False
+        self._gather_logprobs_precompiled = False
+        if not vllm_envs.VLLM_DISABLE_COMPILE_CACHE:
             logger.info("Enabling JAX compile cache.")
             jax.config.update("jax_compilation_cache_dir",
-                              envs.VLLM_XLA_CACHE_PATH)
+                              vllm_envs.VLLM_XLA_CACHE_PATH)
 
     def _create_dummy_tensor(self,
                              shape: Tuple[int, ...],
@@ -67,8 +85,7 @@ class CompilationManager:
         logger.info("Compilation finished in %.2f [secs].", end - start)
 
     def capture_model(self) -> None:
-        if os.getenv("SKIP_JAX_PRECOMPILE",
-                     False) or self.runner.model_config.enforce_eager:
+        if envs.SKIP_JAX_PRECOMPILE or self.runner.model_config.enforce_eager:
             return
         logger.info("Precompile all the subgraphs with possible input shapes.")
 
@@ -81,11 +98,17 @@ class CompilationManager:
                 self._precompile_backbone_with_inputs_embeds()
             if self.runner.scheduler_config.async_scheduling:
                 self._precompile_substitute_placeholder_token()
+            if not self.runner.is_last_rank:
+                return
             self._precompile_select_from_array()
             self._precompile_compute_logits()
+            # Skip sampling if already precompiled before KV cache allocation
+            if not self._sampling_precompiled:
+                self._precompile_sampling()
             self._precompile_disagg_utils()
-            self._precompile_sampling()
-            self._precompile_gather_logprobs()
+            # Skip gather_logprobs if already precompiled before KV cache allocation
+            if not self._gather_logprobs_precompiled:
+                self._precompile_gather_logprobs()
             self._precompile_structured_decoding()
             if self.runner.speculative_config:
                 self._precompile_speculative_decoding()
@@ -104,7 +127,7 @@ class CompilationManager:
 
             self._run_compilation(
                 "input_embeddings_merger",
-                self.runner.get_input_embeddings_fn,
+                self.runner.embed_input_ids_fn,
                 self.runner.state,
                 dummy_input_ids,
                 dummy_multimodal_embeddings,
@@ -113,15 +136,22 @@ class CompilationManager:
 
             self._run_compilation(
                 "input_embeddings_merger_text_only",
-                self.runner.get_input_embeddings_fn,
+                self.runner.embed_input_ids_fn,
                 self.runner.state,
                 dummy_input_ids,
                 None,
                 num_tokens=num_tokens,
             )
 
-    def _precompile_backbone_helper(self, name, *, input_ids, positions,
-                                    inputs_embeds) -> None:
+    def _precompile_backbone_helper(self,
+                                    name,
+                                    *,
+                                    input_ids,
+                                    positions,
+                                    inputs_embeds,
+                                    intermediate_tensors=None,
+                                    is_first_rank=True,
+                                    is_last_rank=True) -> None:
         num_tokens = None
         if input_ids is not None:
             num_tokens = input_ids.shape[0]
@@ -181,10 +211,14 @@ class CompilationManager:
             inputs_embeds,
             layer_name_to_kvcache_index,
             lora_metadata,
+            intermediate_tensors,
+            is_first_rank,
+            is_last_rank,
         ):
             kv_caches, hidden_states, _ = self.runner.model_fn(
                 state, kv_caches, input_ids, attention_metadata, inputs_embeds,
-                positions, layer_name_to_kvcache_index, lora_metadata)
+                positions, layer_name_to_kvcache_index, lora_metadata,
+                intermediate_tensors, is_first_rank, is_last_rank)
             self.runner.kv_caches = kv_caches
             return hidden_states
 
@@ -207,6 +241,9 @@ class CompilationManager:
                 inputs_embeds,
                 tuple(self.runner.layer_name_to_kvcache_index.items()),
                 lora_metadata,
+                intermediate_tensors,
+                is_first_rank,
+                is_last_rank,
                 num_tokens=num_tokens,
             )
 
@@ -257,6 +294,7 @@ class CompilationManager:
                 )
 
     def _precompile_backbone_text_only(self) -> None:
+        hidden_size = self.runner.model_config.get_hidden_size()
         for num_tokens in self.runner.num_tokens_paddings:
             dp_sharding = NamedSharding(
                 self.runner.mesh, PartitionSpec(ShardingAxisName.ATTN_DATA, )
@@ -266,10 +304,28 @@ class CompilationManager:
                                                   dp_sharding)
             positions = self._create_dummy_tensor((num_tokens, ), jnp.int32,
                                                   dp_sharding)
-            self._precompile_backbone_helper("backbone",
-                                             input_ids=input_ids,
-                                             positions=positions,
-                                             inputs_embeds=None)
+            is_first_rank = self.runner.is_first_rank
+            is_last_rank = self.runner.is_last_rank
+            if is_first_rank:
+                intermediate_tensors = None
+            else:
+                hidden_states = self._create_dummy_tensor(
+                    (num_tokens, hidden_size), jnp.bfloat16)
+                residual = self._create_dummy_tensor((num_tokens, hidden_size),
+                                                     jnp.bfloat16)
+                intermediate_tensors = JaxIntermediateTensors(
+                    tensors={
+                        "hidden_states": hidden_states,
+                        "residual": residual
+                    })
+            self._precompile_backbone_helper(
+                f"worker{self.runner.rank} backbone",
+                input_ids=input_ids,
+                positions=positions,
+                inputs_embeds=None,
+                intermediate_tensors=intermediate_tensors,
+                is_first_rank=is_first_rank,
+                is_last_rank=is_last_rank)
 
     def _precompile_backbone_with_inputs_embeds(self) -> None:
         hidden_size = self.runner.model_config.get_hidden_size()
@@ -283,10 +339,28 @@ class CompilationManager:
             else:
                 positions = self._create_dummy_tensor((num_tokens, ),
                                                       jnp.int32)
-            self._precompile_backbone_helper("backbone with embeds",
-                                             input_ids=None,
-                                             positions=positions,
-                                             inputs_embeds=inputs_embeds)
+            is_first_rank = self.runner.is_first_rank
+            is_last_rank = self.runner.is_last_rank
+            if not is_first_rank:
+                hidden_states = self._create_dummy_tensor(
+                    (num_tokens, hidden_size), jnp.bfloat16)
+                residual = self._create_dummy_tensor((num_tokens, hidden_size),
+                                                     jnp.bfloat16)
+                intermediate_tensors = JaxIntermediateTensors(
+                    tensors={
+                        "hidden_states": hidden_states,
+                        "residual": residual
+                    })
+            else:
+                intermediate_tensors = None
+            self._precompile_backbone_helper(
+                f"worker{self.runner.rank} backbone with embeds",
+                input_ids=None,
+                positions=positions,
+                inputs_embeds=inputs_embeds,
+                intermediate_tensors=intermediate_tensors,
+                is_first_rank=is_first_rank,
+                is_last_rank=is_last_rank)
 
     def _precompile_select_from_array_helper(
         self,
@@ -354,7 +428,7 @@ class CompilationManager:
             self.runner.mesh, PartitionSpec(ShardingAxisName.ATTN_DATA, None))
         dp_size = self.runner.vllm_config.sharding_config.total_dp_size
         self._precompile_select_from_array_helper(
-            name="select all logits",
+            name=f"worker{self.runner.rank} select all logits",
             source_paddings=self.runner.num_tokens_paddings,
             indices_paddings=index_paddings,
             hidden_dim=hsize,
@@ -365,7 +439,8 @@ class CompilationManager:
         if self.runner.speculative_config:
             vocab_size = self.runner.model_config.get_vocab_size()
             self._precompile_select_from_array_helper(
-                name="select bonus tokens for spec decoding",
+                name=
+                f"worker{self.runner.rank} select bonus tokens for spec decoding",
                 source_paddings=self.runner.num_logits_paddings,
                 indices_paddings=self.runner.num_reqs_paddings,
                 hidden_dim=vocab_size,
@@ -373,7 +448,8 @@ class CompilationManager:
                                              PartitionSpec(None, "model")),
             )
             self._precompile_select_from_array_helper(
-                name="select target tokens for spec decoding",
+                name=
+                f"worker{self.runner.rank} select target tokens for spec decoding",
                 source_paddings=self.runner.num_logits_paddings,
                 indices_paddings=self.runner.num_logits_paddings,
                 hidden_dim=vocab_size,
@@ -396,7 +472,7 @@ class CompilationManager:
                     np.array([num_reqs], dtype=np.int32)):
                 lora_metadata = self.runner.lora_utils.extract_lora_metadata()
                 self._run_compilation(
-                    "compute_logits",
+                    f"worker{self.runner.rank} compute_logits",
                     self.runner.compute_logits_fn,
                     self.runner.state,
                     hidden_states,
@@ -410,43 +486,48 @@ class CompilationManager:
         for num_reqs in self.runner.num_reqs_paddings:
             logits_sharding = NamedSharding(
                 self.runner.mesh,
-                PartitionSpec(ShardingAxisName.ATTN_DATA, "model"))
+                PartitionSpec(ShardingAxisName.MLP_DATA,
+                              ShardingAxisName.MLP_TENSOR))
             dp_size = self.runner.vllm_config.sharding_config.total_dp_size
             sampling_metadata_sharding = NamedSharding(
                 self.runner.mesh, PartitionSpec(
-                    ShardingAxisName.ATTN_DATA)) if dp_size > 1 else None
+                    ShardingAxisName.MLP_DATA)) if dp_size > 1 else None
             logits = self._create_dummy_tensor((num_reqs, hsize), jnp.bfloat16,
                                                logits_sharding)
             for do_sampling in (True, False):
-                if do_sampling:
-                    temperature = np.full((num_reqs, ), 0.7, dtype=np.float32)
-                    top_k = np.full((num_reqs, ), 20, dtype=np.int32)
-                    top_p = np.full((num_reqs, ), 0.8, dtype=np.float32)
-                    (temperature, top_k,
-                     top_p) = device_array(self.runner.mesh,
-                                           (temperature, top_k, top_p),
-                                           sharding=sampling_metadata_sharding)
-                else:
-                    temperature = None
-                    top_k = None
-                    top_p = None
-
-                sampling_metadata = TPUSupportedSamplingMetadata(
-                    temperature=temperature,
-                    top_k=top_k,
-                    top_p=top_p,
-                    do_sampling=do_sampling,
-                )
-                self._run_compilation(
-                    "sample",
-                    sample,
-                    self.runner.rng_params_for_sampling,
-                    self.runner.mesh,
-                    logits,
-                    sampling_metadata,
-                    num_reqs=num_reqs,
-                    do_sampling=do_sampling,
-                )
+                for logprobs in (True, False):
+                    if do_sampling:
+                        temperature = np.full((num_reqs, ),
+                                              0.7,
+                                              dtype=np.float32)
+                        top_k = np.full((num_reqs, ), 20, dtype=np.int32)
+                        top_p = np.full((num_reqs, ), 0.8, dtype=np.float32)
+                        (temperature, top_k, top_p) = device_array(
+                            self.runner.mesh, (temperature, top_k, top_p),
+                            sharding=sampling_metadata_sharding)
+                    else:
+                        temperature = None
+                        top_k = None
+                        top_p = None
+
+                    sampling_metadata = TPUSupportedSamplingMetadata(
+                        temperature=temperature,
+                        top_k=top_k,
+                        top_p=top_p,
+                        do_sampling=do_sampling,
+                        logprobs=logprobs)
+                    self._run_compilation(
+                        f"worker{self.runner.rank} sample",
+                        sample,
+                        self.runner.rng_params_for_sampling,
+                        self.runner.mesh,
+                        logits,
+                        sampling_metadata,
+                        num_reqs=num_reqs,
+                        do_sampling=do_sampling,
+                    )
+
+        self._sampling_precompiled = True
 
     def _precompile_disagg_utils(self) -> None:
         if not is_disagg_enabled():
@@ -476,10 +557,18 @@ class CompilationManager:
         logger.info("Compiling gather_logprobs with different input shapes.")
         hsize = self.runner.model_config.get_vocab_size()
         for num_reqs in self.runner.num_reqs_paddings:
-            logits = self._create_dummy_tensor((num_reqs, hsize), jnp.bfloat16)
-            token_ids = self._create_dummy_tensor((num_reqs, ), jnp.int32)
+            logits_sharding = NamedSharding(
+                self.runner.mesh,
+                PartitionSpec(ShardingAxisName.MLP_DATA,
+                              ShardingAxisName.MLP_TENSOR))
+            token_ids_sharding = NamedSharding(
+                self.runner.mesh, PartitionSpec(ShardingAxisName.MLP_DATA, ))
+            logits = self._create_dummy_tensor((num_reqs, hsize), jnp.bfloat16,
+                                               logits_sharding)
+            token_ids = self._create_dummy_tensor((num_reqs, ), jnp.int32,
+                                                  token_ids_sharding)
             self._run_compilation(
-                "gather_logprobs",
+                f"worker{self.runner.rank} gather_logprobs",
                 self.runner._compute_and_gather_logprobs,
                 logits,
                 token_ids,
@@ -487,6 +576,8 @@ class CompilationManager:
                 num_reqs=num_reqs,
             )
 
+        self._gather_logprobs_precompiled = True
+
     def _precompile_speculative_decoding(self) -> None:
         logger.info(
             "Compiling speculative_decoding with different input shapes.")
@@ -531,7 +622,7 @@ class CompilationManager:
                             do_sampling=do_sampling)
 
                     self._run_compilation(
-                        compilation_name,
+                        f"worker{self.runner.rank} {compilation_name}",
                         self.runner.rejection_sampler,
                         draft_token_ids,
                         num_draft_tokens,
@@ -548,7 +639,9 @@ class CompilationManager:
     def _precompile_eagle3_helpers(self) -> None:
         logger.info(
             "Compiling eagle3 jitted helpers with different input shapes.")
-        hidden_size = self.runner.model_config.get_hidden_size()
+        target_hidden_size = self.runner.model_config.get_hidden_size()
+        draft_hidden_size = self.runner.speculative_config.draft_model_config.get_hidden_size(
+        )
         dtype = self.runner.model_config.dtype
 
         num_kv_cache_groups = len(self.runner.kv_cache_config.kv_cache_groups)
@@ -595,10 +688,11 @@ class CompilationManager:
 
         for num_logits in self.runner.num_logits_paddings:
             hidden_states = self._create_dummy_tensor(
-                (num_logits, hidden_size), jnp.bfloat16)
+                (num_logits, draft_hidden_size), jnp.bfloat16)
             self._run_compilation(
                 "eagle3_get_draft_token_ids",
                 self.runner.drafter._get_draft_token_ids,
+                self.runner.drafter.state,
                 hidden_states,
                 num_logits=num_logits,
             )
@@ -606,8 +700,8 @@ class CompilationManager:
         input_ids_loop = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32,
             NamedSharding(self.runner.mesh, PartitionSpec()))
-        target_hidden_state_loop = self._create_dummy_tensor(
-            (self.runner.max_num_reqs, hidden_size), dtype,
+        draft_hidden_state_loop = self._create_dummy_tensor(
+            (self.runner.max_num_reqs, draft_hidden_size), dtype,
             NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
         next_token_ids = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32)
@@ -615,9 +709,12 @@ class CompilationManager:
             (self.runner.max_num_reqs, ), jnp.int32)
         for num_tokens in self.runner.num_tokens_paddings:
             aux_hidden_states = [
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
             ]
 
             positions = self._create_dummy_tensor((num_tokens, ), jnp.int32)
@@ -640,23 +737,23 @@ class CompilationManager:
                 num_reqs,
             ):
                 target_hidden_states, input_ids, last_token_indices, _ = self.runner.drafter._filter_token_and_prepare_initial_inputs(
-                    token_indices, query_start_loc, seq_lens, input_ids,
-                    aux_hidden_states, attention_metadata, next_token_ids,
-                    num_reqs)
+                    self.runner.drafter.state, token_indices, query_start_loc,
+                    seq_lens, input_ids, aux_hidden_states, attention_metadata,
+                    next_token_ids, num_reqs)
                 return target_hidden_states, input_ids, last_token_indices
 
             input_ids = self._create_dummy_tensor((num_tokens, ), jnp.int32)
             aux_hidden_states = [
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
             ]
@@ -688,17 +785,17 @@ class CompilationManager:
                 state,
                 kv_caches,
                 input_ids,
-                target_hidden_states,
+                draft_hidden_states,
                 attention_metadata,
             ):
                 kv_caches, hidden_states, _ = self.runner.drafter.model_fn(
-                    state, kv_caches, input_ids, target_hidden_states,
+                    state, kv_caches, input_ids, draft_hidden_states,
                     attention_metadata)
                 self.runner.kv_caches = kv_caches
                 return hidden_states
 
-            target_hidden_states = self._create_dummy_tensor(
-                (num_tokens, hidden_size), dtype,
+            draft_hidden_states = self._create_dummy_tensor(
+                (num_tokens, draft_hidden_size), dtype,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, "model")))
             input_ids = self._create_dummy_tensor(
                 (num_tokens, ), jnp.int32,
@@ -709,7 +806,7 @@ class CompilationManager:
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids,
-                target_hidden_states,
+                draft_hidden_states,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
@@ -719,6 +816,7 @@ class CompilationManager:
             self._run_compilation(
                 "eagle3_prepare_hidden_states_and_input_ids",
                 self.runner.drafter._prepare_hidden_states_and_input_ids,
+                self.runner.drafter.state,
                 aux_hidden_states,
                 query_start_loc,
                 target_token_ids,
@@ -741,18 +839,19 @@ class CompilationManager:
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids_loop,
-                target_hidden_state_loop,
+                draft_hidden_state_loop,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
 
             hidden_states = self._create_dummy_tensor(
-                (num_tokens, hidden_size), jnp.bfloat16,
+                (num_tokens, draft_hidden_size), jnp.bfloat16,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
 
             self._run_compilation(
                 "eagle3_select_inputs_for_loop_speculation",
                 self.runner.drafter._select_inputs_for_loop_speculation,
+                self.runner.drafter.state,
                 positions,
                 hidden_states,
                 hidden_states,
@@ -763,6 +862,7 @@ class CompilationManager:
             self._run_compilation(
                 "eagle3_select_draft_token_ids",
                 self.runner.drafter._select_draft_token_ids,
+                self.runner.drafter.state,
                 hidden_states,
                 last_token_indices,
                 num_tokens=num_tokens,

tpu_inference/runner/kv_cache.py

@@ -1,3 +1,17 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any, List
 
 import jax
@@ -7,6 +21,7 @@ from jax._src import dtypes
 from jax.sharding import Mesh, NamedSharding, PartitionSpec
 from torchax.ops.mappings import t2j_dtype
 
+import tpu_inference.kernels.mla.v1.kernel as mla
 import tpu_inference.kernels.ragged_paged_attention.v3.kernel as rpa
 import tpu_inference.kernels.ragged_paged_attention.v3.kernel_hd64 as rpa_hd64
 from tpu_inference.layers.common.sharding import ShardingAxisName
@@ -17,9 +32,13 @@ logger = init_logger(__name__)
 DEFAULT_KV_CACHE_DTYPE = jnp.bfloat16
 
 
-def get_kv_cache_shape_with_mesh(mesh: Mesh, total_num_pages: int,
-                                 page_size: int, actual_num_kv_heads: int,
-                                 actual_head_dim: int, kv_dtype: any):
+def get_kv_cache_shape_with_mesh(mesh: Mesh,
+                                 total_num_pages: int,
+                                 page_size: int,
+                                 actual_num_kv_heads: int,
+                                 actual_head_dim: int,
+                                 kv_dtype: any,
+                                 use_mla: bool = False):
     """Gets the KV cache shape based on the mesh configuration."""
 
     model_cnt = mesh.shape["model"]
@@ -28,15 +47,21 @@ def get_kv_cache_shape_with_mesh(mesh: Mesh, total_num_pages: int,
     # specific model, rather than being determined by the head_dim. If new
     # models are introduced with a head_dim of 64, this will require additional
     # model-specific adjustments.
-    get_kv_cache_shape_fn = (
-        rpa_hd64.get_kv_cache_shape if actual_head_dim == 64 \
-            else rpa.get_kv_cache_shape
-    )
-    shape = list(
-        get_kv_cache_shape_fn(total_num_pages, page_size,
-                              actual_num_kv_heads // model_cnt,
-                              actual_head_dim, kv_dtype))
-    shape[2] *= model_cnt
+    if use_mla:
+        get_kv_cache_shape_fn = mla.get_kv_cache_shape
+        shape = list(
+            get_kv_cache_shape_fn(total_num_pages, page_size, actual_head_dim,
+                                  kv_dtype))
+    else:
+        get_kv_cache_shape_fn = (
+            rpa_hd64.get_kv_cache_shape if actual_head_dim == 64 \
+                else rpa.get_kv_cache_shape
+        )
+        shape = list(
+            get_kv_cache_shape_fn(total_num_pages, page_size,
+                                  actual_num_kv_heads // model_cnt,
+                                  actual_head_dim, kv_dtype))
+        shape[2] *= model_cnt
     return tuple(shape)
 
 
@@ -48,6 +73,7 @@ def create_kv_caches(
     mesh: Mesh,
     layer_names: List[str],
     cache_dtype: jnp.dtype = DEFAULT_KV_CACHE_DTYPE,
+    use_mla: bool = False,
 ) -> List[jax.Array]:
     """
     Creates a list of KV cache where each array mapps to single attention layer.
@@ -74,12 +100,16 @@ def create_kv_caches(
 
     cache_shape = get_kv_cache_shape_with_mesh(mesh, num_blocks, block_size,
                                                num_kv_heads, head_size,
-                                               cache_dtype)
+                                               cache_dtype, use_mla)
 
-    sharding = NamedSharding(
-        mesh,
-        PartitionSpec(ShardingAxisName.ATTN_DATA, None,
-                      ShardingAxisName.ATTN_HEAD))
+    if use_mla:
+        sharding = NamedSharding(mesh,
+                                 PartitionSpec(ShardingAxisName.MLP_TENSOR))
+    else:
+        sharding = NamedSharding(
+            mesh,
+            PartitionSpec(ShardingAxisName.ATTN_DATA, None,
+                          ShardingAxisName.ATTN_HEAD))
 
     def _allocate() -> jax.Array:
         return jnp.empty(
@@ -94,7 +124,8 @@ def create_kv_caches(
     return kv_caches
 
 
-def get_rpa_page_size_bytes(mesh: Mesh, kv_cache_specs: dict[str, Any]) -> int:
+def get_attention_page_size_bytes(mesh: Mesh,
+                                  kv_cache_specs: dict[str, Any]) -> int:
     """
     Calculate KV cache page size of RPA kernel.
 
@@ -107,14 +138,16 @@ def get_rpa_page_size_bytes(mesh: Mesh, kv_cache_specs: dict[str, Any]) -> int:
     """
 
     # Import it here to avoid circular import.
-    from vllm.v1.kv_cache_interface import AttentionSpec
+    from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
 
     page_size_bytes_set = set()
     for kv_cache_spec in kv_cache_specs.values():
         assert isinstance(kv_cache_spec, AttentionSpec)
 
         dtype = t2j_dtype(kv_cache_spec.dtype)
-        bits = dtypes.bit_width(dtype)
+        bits = (dtypes.bit_width(dtype) if hasattr(dtypes, "bit_width") else
+                dtypes.itemsize_bits(dtype))
+        use_mla = isinstance(kv_cache_spec, MLAAttentionSpec)
 
         kv_cache_shape = get_kv_cache_shape_with_mesh(
             mesh=mesh,
@@ -123,6 +156,7 @@ def get_rpa_page_size_bytes(mesh: Mesh, kv_cache_specs: dict[str, Any]) -> int:
             actual_num_kv_heads=kv_cache_spec.num_kv_heads,
             actual_head_dim=kv_cache_spec.head_size,
             kv_dtype=dtype,
+            use_mla=use_mla,
         )
         page_size_bytes = (bits * np.prod(kv_cache_shape)) // 8
         page_size_bytes_set.add(page_size_bytes)