tpu-inference 0.11.1__py3-none-any.whl

tpu_inference/layers/jax/sharding.py

@@ -0,0 +1,406 @@
+import json
+from dataclasses import dataclass
+
+import numpy as np
+from jax.sharding import Mesh
+from vllm.config import VllmConfig
+
+BATCH_AXIS_NAME = 'data'
+SEQUENCE_AXIS_NAME = 'data'
+DATA_AXIS_NAME = 'data'
+ATTN_HEAD_AXIS_NAME = 'model'
+ATTN_TENSOR_AXIS_NAME = None
+MLP_TENSOR_AXIS_NAME = ('model', 'expert')
+MOE_TENSOR_AXIS_NAME = 'model'
+EXPERT_AXIS_NAME = 'expert'
+VOCAB_AXIS_NAME = ('data', 'expert', 'model')
+
+
+@dataclass
+class ShardingStrategy:
+    """Defines the high-level parallelism strategy.
+
+    This class specifies how many ways each type of parallelism (tensor, expert,
+    sequence, data) should be distributed across the available devices.
+
+    Attributes:
+        tensor_parallelism: The degree of tensor parallelism (e.g., splitting
+            weights of a single layer).
+        expert_parallelism: The degree of expert parallelism for MoE models.
+        sequence_parallelism: The degree of sequence parallelism (splitting
+            activations along the sequence length dimension).
+        data_parallelism: The degree of data parallelism (splitting the batch
+            across devices).
+    """
+    tensor_parallelism: int = 1
+    expert_parallelism: int = 1
+    sequence_parallelism: int = 1
+    data_parallelism: int = 1
+
+
+#TODO split this into block unique sharding config, i.e. attentionShardingConfig, MoEShardingConfig
+@dataclass
+class ShardingRulesConfig:
+    """Holds detailed sharding configurations for individual tensors, namely logical rules.
+
+    Each attribute in this class corresponds to a specific weight or activation
+    tensor within a transformer model. The value of each attribute is a
+    tuple of logical mesh axis names (e.g., 'dp', 'sp', 'tp'), which defines
+    how the corresponding tensor's dimensions are partitioned across the device mesh.
+    The dimension order in the attribute name (e.g., `btd` for batch, sequence,
+    d_model) maps directly to the sharding tuple.
+
+    TODO: update the mesh axis names to be clear and reduce confusion between prefill & generate
+    """
+
+    # Activation for attn input: (Batch * Sequence, Dim)
+    activation_attention_td: tuple = (None, None)
+    # Activation for attn out: (Batch * Sequence, Dim)
+    activation_attention_out_td: tuple = (None, None)
+    # Activation for q projection input: (Batch * Sequence, Dim)
+    activation_q_td: tuple = (None, None)
+    # Attention Out activation after projection: (Batch * Sequence, NumHeads, HeadDim)
+    attn_o_tnh: tuple = (None, None, None)
+    # Q vector: (Batch * Sequence, NumHeads, HeadDim)
+    query_tnh: tuple = (None, None, None)
+    # K/V vector: (Batch * Sequence, NumKVHeads, HeadDim)
+    keyvalue_skh: tuple = (None, None, None)
+
+    # Attention Q weight: (Dim, NumHeads, HeadDim)
+    attn_q_weight_dnh: tuple = (None, None, None)
+    # Attention K weight: (Dim, NumKVHeads, HeadDim)
+    attn_k_weight_dkh: tuple = (None, None, None)
+    # Attention V weight: (Dim, NumKVHeads, HeadDim)
+    attn_v_weight_dkh: tuple = (None, None, None)
+    # Attention Out weight: (NumHeads, HeadDim, Dim)
+    attn_o_weight_nhd: tuple = (None, None, None)
+
+    # Activation for ffw input: (Batch * Sequence, Dim)
+    activation_ffw_td: tuple = (None, None)
+
+    # Activation for ffw input: (Batch * Sequence, Expert, Dim)
+    activation_ffw_ted: tuple = (None, None, None)
+
+    # FFW hidden activation: (Batch * Sequence, FfwDim)
+    ffw_hidden_tf: tuple = (None, None)
+
+    # FFW up/gate weight: (Dim, FfwDim)
+    ffw_weight_df: tuple = (None, None)
+    # FFW down weight: (FfwDim, Dim)
+    ffw_weight_fd: tuple = (None, None)
+    # MoE gate/up weights: (NumExperts, Dim, FfwDim)
+    moe_weights_edf: tuple = (None, None, None)
+    # MoE down weights: (NumExperts, FfwDim, Dim)
+    moe_weights_efd: tuple = (None, None, None)
+    # MoE router weights: (Dim, NumExperts)
+    moe_router_de: tuple = (None, None)
+    # MoE router bias weights: (NumExperts,)
+    moe_router_bias_e: tuple = (None, )
+
+    # Embedding weight: (VocabSize, Dim)
+    emb_weight_vd: tuple = (None, None)
+    # Activation between layers: (Batch * Sequence, Dim)
+    activation_td: tuple = (None, None)
+    # Final activation before logits: (Batch * Sequence, Dim)
+    prelogit_td: tuple = (None, None)
+    # Logit activation: (Batch * Sequence, VocabSize)
+    logits_tv: tuple = (None, None)
+    # RMS norm scale weight: (Dim,)
+    norm_scale: tuple = (None)
+    # Vocab projection weight (tied embeddings): (Dim, VocabSize)
+    vocab_vd: tuple = (None, None)
+    vocab_dv: tuple = (None, None)
+
+
+class ShardingConfig:
+    """Container for operation-specific sharding configurations.
+
+    This class holds two separate `ShardingRulesConfig` objects, one for the
+    'prefill' phase and one for the 'generate' (or decode) phase of model
+    execution. This allows tailoring sharding strategies to the different
+    computational patterns of each phase.
+
+    Example Sharding Strategy and Configuration:
+
+    Sharding Strategy defines the high-level parallelism dimensions.
+    For a device mesh like `Mesh((2, 4, 4, 4), ('data', 'seq', 'expert', 'tensor'))` on 128 devices:
+    - data: Data Parallelism (2-way)
+    - seq: Sequence Parallelism (4-way)
+    - expert: Expert Parallelism (4-way)
+    - tensor: Tensor Parallelism (4-way)
+
+    ShardingConfig then maps tensor dimensions to these logical mesh axes.
+    For example, a tensor with shape (Batch, Sequence, Dimension) could be sharded
+    differently for prefill and decode/generate operations:
+
+    - Prefill (long sequences, small batch):
+    Sharding sequence dim on the 'sp' axis is often efficient.
+    `prefill_rules.activation_attention_btd = (None, 'seq', 'tensor')`
+
+    - Generate (short sequences, large batch):
+    Sharding batch dim on the 'dp' axis is often efficient.
+    `generate_rules.activation_attention_btd = ('data', None, 'tensor')`
+    """
+
+    def __init__(self,
+                 prefill_rules=None,
+                 generate_rules=None,
+                 default_rules_cls=ShardingRulesConfig):
+        """Initializes the ShardingConfig.
+
+        Args:
+            prefill_rules: An `ShardingRulesConfig` for the prefill phase.
+                If None, a default config is created.
+            generate_rules: An `ShardingRulesConfig` for the generate phase.
+                If None, a default config is created.
+            default_rules_cls: The default sharding rules (class) to use.
+        """
+        # Use a factory pattern to avoid mutable default arguments
+        self.default_rules_cls = default_rules_cls
+        self.prefill_rules = prefill_rules if prefill_rules is not None else default_rules_cls(
+        )
+        self.generate_rules = generate_rules if generate_rules is not None else default_rules_cls(
+        )
+
+
+def build_mesh(devices, strategy: dict[str, int]) -> Mesh:
+    """Constructs a JAX device mesh from a sharding strategy.
+
+    This method creates a logical grid of devices based on the parallelism
+    degrees defined in the strategy. The logical axis names ('dp', 'ep',
+    'sp', 'tp') are used to map tensor dimensions to the physical device grid.
+
+    Args:
+        strategy: A dictionary from upper level config.
+
+    Returns:
+        A JAX `Mesh` object.
+    """
+
+    axis_order = {
+        "data": strategy.get("data_parallelism", 1),
+        "expert": strategy.get("expert_parallelism", 1),
+        "seq": strategy.get("sequence_parallelism", 1),
+        "model": strategy.get("tensor_parallelism", 1),
+    }
+    # TODO: add logic to infer axis when the degree is -1
+    mesh_axis_names = []
+    mesh_shape = []
+    for axis, dim in axis_order.items():
+        mesh_axis_names.append(axis)
+        mesh_shape.append(dim)
+
+    if not mesh_shape:
+        mesh_shape = [1]
+        mesh_axis_names = [
+            'data'
+        ]  # default to data parallelism if no other strategy is specified
+
+    devices = np.asarray(devices).reshape(mesh_shape)
+    return Mesh(devices, axis_names=tuple(mesh_axis_names))
+
+
+class Sharding:
+    """Generates and manages sharding configurations based on a high-level strategy.
+
+    This class populates a `ShardingConfig` with detailed tensor sharding
+    rules for both prefill and generation phases. It also allows for runtime
+    overrides of these rules.
+
+    Attributes:
+        sharding_cfg: The generated `ShardingConfig` with detailed rules.
+    """
+
+    def __init__(self,
+                 prefill_rules: dict | None = None,
+                 generate_rules: dict | None = None,
+                 default_rules_cls=ShardingRulesConfig,
+                 vllm_config: VllmConfig = None):
+        """Initializes the Sharding manager.
+
+        Args:
+            prefill_rules: A dictionary of overrides for the prefill
+                sharding config. Keys are attribute names in `ShardingRulesConfig`,
+                and values are the new sharding tuples.
+            generate_rules: A dictionary of overrides for the generate
+                sharding config.
+        """
+        self.vllm_config = vllm_config
+        self.default_rules_cls = default_rules_cls
+        self.sharding_cfg = self.make_sharding_config(
+            default_rules_cls=default_rules_cls,
+            prefill_overrides=prefill_rules,
+            generate_overrides=generate_rules)
+
+    def _get_overrides(self, sharding_phase: str):
+        """Return the overrides from the vLLM config for the given sharding phase."""
+        overrides = {}
+        try:
+            overrides = self.vllm_config.additional_config["sharding"][
+                "logical_rules"]["all"]
+        except KeyError:
+            pass
+
+        try:
+            additional_overrides = self.vllm_config.additional_config[
+                "sharding"]["logical_rules"][f"{sharding_phase}"]
+            overrides.update(additional_overrides)
+        except KeyError:
+            pass
+        return overrides
+
+    def __str__(self):
+        """Succinct representation of relevant Sharding settings and overrides."""
+        output_str = f"  Using {self.default_rules_cls.__name__} logical rules.\n"
+        output_str += f"  {self.__class__.__name__:} overrides:\n"
+        output_str += f"    prefill logical_rule overrides:\n    {json.dumps(self._get_overrides('prefill'), indent=4, default=str)}\n\n"
+        output_str += f"    generate logical_rule overrides:\n    {json.dumps(self._get_overrides('generate'), indent=4, default=str)}\n\n"
+        return output_str
+
+    def validate_sharding_strategy(self, ):
+        """Validates if the sharding strategy is compatible with the environment.
+
+        This method is a placeholder now, and will check if the product of parallelism degrees
+        matches the number of available devices.
+        """
+        #TODO: check num_devices % parallelism == 0
+        #TODO: check num_devices == multiply(parallelism(with inferred))
+        return
+
+    def get_sharding_cfg(self) -> ShardingConfig:
+        """Returns the generated sharding configuration."""
+        return self.sharding_cfg
+
+    def _apply_overrides(self, config_obj: ShardingRulesConfig,
+                         overrides: dict | None):
+        """Applies runtime overrides to a sharding configuration object.
+
+        Args:
+            config_obj: The sharding configuration object (e.g., prefill_rules)
+                to be updated.
+            overrides: A dictionary where keys are attribute names of the config
+                object and values are the new sharding tuples.
+
+        Raises:
+            AttributeError: If a key in the overrides dictionary is not a valid
+                attribute of the configuration object.
+        """
+        for key, value in overrides.items():
+            if hasattr(config_obj, key):
+                setattr(config_obj, key, value)
+            else:
+                # Raise an error for invalid keys to prevent silent failures
+                raise AttributeError(
+                    f"'{key}' is not a valid attribute of {type(config_obj).__name__}"
+                )
+
+    def _make_default_sharding_config(self, prefill_rules, generate_rules):
+
+        # Populate Prefill Config
+        # During prefill, sequence length is long, so we shard along the sequence axis.
+        prefill_rules.activation_attention_td = (DATA_AXIS_NAME,
+                                                 ATTN_TENSOR_AXIS_NAME)
+        prefill_rules.activation_attention_out_td = (DATA_AXIS_NAME,
+                                                     ATTN_TENSOR_AXIS_NAME)
+        prefill_rules.activation_q_td = (DATA_AXIS_NAME, ATTN_TENSOR_AXIS_NAME)
+        #TODO: the default qkv and kvcache is sharded on head dim
+        # We may change it after we finalize the KVCache design
+        prefill_rules.attn_o_tnh = (DATA_AXIS_NAME, ATTN_HEAD_AXIS_NAME, None)
+        prefill_rules.query_tnh = (DATA_AXIS_NAME, ATTN_HEAD_AXIS_NAME, None)
+        prefill_rules.keyvalue_skh = (DATA_AXIS_NAME, ATTN_HEAD_AXIS_NAME,
+                                      None)
+
+        # Populate Generate (Decode) Config
+        # During decode, batch size is the large dimension, so we shard along the batch axis.
+        generate_rules.activation_attention_td = (DATA_AXIS_NAME,
+                                                  ATTN_TENSOR_AXIS_NAME)
+        generate_rules.activation_attention_out_td = (DATA_AXIS_NAME,
+                                                      ATTN_TENSOR_AXIS_NAME)
+        generate_rules.activation_q_td = (DATA_AXIS_NAME,
+                                          ATTN_TENSOR_AXIS_NAME)
+        #TODO: the default qkv and kvcache is sharded on head dim
+        # We may change it after we finalize the KVCache design
+        generate_rules.attn_o_tnh = (DATA_AXIS_NAME, ATTN_HEAD_AXIS_NAME, None)
+        generate_rules.query_tnh = (DATA_AXIS_NAME, ATTN_HEAD_AXIS_NAME, None)
+        generate_rules.keyvalue_skh = (DATA_AXIS_NAME, ATTN_HEAD_AXIS_NAME,
+                                       None)
+        generate_rules.attn_q_weight_dnh = (None, ATTN_HEAD_AXIS_NAME,
+                                            ATTN_TENSOR_AXIS_NAME)
+        generate_rules.attn_k_weight_dkh = (None, ATTN_HEAD_AXIS_NAME,
+                                            ATTN_TENSOR_AXIS_NAME)
+        generate_rules.attn_v_weight_dkh = (None, ATTN_HEAD_AXIS_NAME,
+                                            ATTN_TENSOR_AXIS_NAME)
+        generate_rules.attn_o_weight_nhd = (ATTN_HEAD_AXIS_NAME, None,
+                                            ATTN_TENSOR_AXIS_NAME)
+        generate_rules.activation_ffw_td = (DATA_AXIS_NAME, None)
+        generate_rules.activation_ffw_ted = (DATA_AXIS_NAME, EXPERT_AXIS_NAME,
+                                             None)
+        generate_rules.ffw_hidden_tf = (DATA_AXIS_NAME, MLP_TENSOR_AXIS_NAME)
+        # FFW weights are typically sharded along the hidden dimension (F).
+        generate_rules.ffw_weight_df = (None, MLP_TENSOR_AXIS_NAME)
+        generate_rules.ffw_weight_fd = (MLP_TENSOR_AXIS_NAME, None)
+        # MoE weights are sharded along the expert axis and the hidden dimension.
+        generate_rules.moe_weights_edf = (EXPERT_AXIS_NAME, None,
+                                          MOE_TENSOR_AXIS_NAME)
+        generate_rules.moe_weights_efd = (EXPERT_AXIS_NAME,
+                                          MOE_TENSOR_AXIS_NAME, None)
+        generate_rules.moe_router_de = (None, EXPERT_AXIS_NAME)
+
+        # Embedding weight: (VocabSize, Dim)
+        generate_rules.emb_weight_vd = (MLP_TENSOR_AXIS_NAME, None)
+        generate_rules.activation_td = (DATA_AXIS_NAME, ATTN_TENSOR_AXIS_NAME)
+        generate_rules.prelogit_td = (DATA_AXIS_NAME, ATTN_TENSOR_AXIS_NAME)
+        generate_rules.logits_tv = (DATA_AXIS_NAME, MLP_TENSOR_AXIS_NAME)
+        generate_rules.vocab_vd = (VOCAB_AXIS_NAME, None)
+        generate_rules.vocab_dv = (None, VOCAB_AXIS_NAME)
+
+    def make_sharding_config(
+            self,
+            default_rules_cls: ShardingRulesConfig,
+            prefill_overrides: dict | None = None,
+            generate_overrides: dict | None = None) -> ShardingConfig:
+        """Creates the detailed `ShardingConfig` with specific partitioning rules
+        and applies any runtime overrides.
+
+        This method populates the `prefill_rules` and
+        `generate_rules` with hardcoded sharding rules that are generally
+        effective for transformer models, and then updates them with any provided
+        overrides.
+
+        Args:
+            prefill_overrides: A dictionary with attribute names and their new values
+                for the prefill sharding configuration.
+            generate_overrides: A dictionary with attribute names and their new values
+                for the generate sharding configuration.
+
+        Returns:
+            The populated and overridden `ShardingConfig` object.
+        """
+        #TODO: organize into update_prefill() and update_decode for each axis
+        #TODO: verify the sharding axes
+        sharding_cfg = ShardingConfig(default_rules_cls=default_rules_cls)
+        prefill_rules = sharding_cfg.prefill_rules
+        generate_rules = sharding_cfg.generate_rules
+
+        # Extract the overrides from the vllm_config if they are not provided programatically.
+        if prefill_overrides is None:
+            prefill_overrides = self._get_overrides("prefill")
+        if generate_overrides is None:
+            generate_overrides = self._get_overrides("generate")
+
+        # Apply default sharding configs
+        self._make_default_sharding_config(prefill_rules, generate_rules)
+
+        # Apply overriding the runtime sharding rules
+        self._apply_overrides(prefill_rules, prefill_overrides)
+        self._apply_overrides(generate_rules, generate_overrides)
+
+        return sharding_cfg
+
+    #TODO: Add __repr__
+
+
+class ShardingInfo:
+    #TODO a sharding info class for visualizing & debugging the sharding performance
+    # Will implement it for the next version
+    pass

tpu_inference/layers/jax/transformer_block.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+# Flax and JAX sharding imports
+import jax
+from flax import nnx
+
+from tpu_inference.layers.jax.attention.attention import (AttentionMetadata,
+                                                          KVCache)
+from tpu_inference.layers.jax.layers import DenseFFW
+from tpu_inference.layers.jax.moe.moe import MoE
+
+
+@dataclass(kw_only=True)
+class TransformerBlock(nnx.Module):
+    """
+    A heavy weight module which serves as the stateful live blocks in serving
+
+    custom_module can be either a dense module (i.e., DenseFFW) or MoE.
+    """
+    pre_attention_norm: nnx.Module
+    pre_mlp_norm: nnx.Module
+    custom_module: nnx.Module
+    attn: nnx.Module
+    use_attention_rope: bool = True
+    quant: Any | None = None
+
+    def __call__(
+            self, x_TD: jax.Array, is_prefill: bool, kv_cache: KVCache,
+            attention_metadata: AttentionMetadata
+    ) -> Tuple[KVCache, jax.Array]:
+        # Attn Block
+        attn_residual_TD = x_TD
+        x_TD = self.pre_attention_norm(x_TD)
+        new_cache, attn_output_TD = self.attn(x_TD, is_prefill, kv_cache,
+                                              attention_metadata,
+                                              self.use_attention_rope)
+        attn_output_TD += attn_residual_TD
+
+        # FFW Block
+        ffw_residual_TD = attn_output_TD
+        normed_ffw_input_TD = self.pre_mlp_norm(attn_output_TD)
+        logits_TD = self.custom_module(normed_ffw_input_TD)
+        logits_TD += ffw_residual_TD
+        return new_cache, logits_TD
+
+
+@dataclass(kw_only=True)
+class SharedExpertsTransformerBlock(TransformerBlock):
+    """Create a modified TransformerBlock that sums MoE layer output with shared expert output."""
+    shared_experts: nnx.Module
+
+    def __call__(self, x_TD, is_prefill, kv_cache, attention_metadata):
+        # Attn Block
+        attn_residual_TD = x_TD
+        x_TD = self.pre_attention_norm(x_TD)
+        new_cache, attn_output_TD = self.attn(x_TD, is_prefill, kv_cache,
+                                              attention_metadata,
+                                              self.use_attention_rope)
+        attn_output_TD += attn_residual_TD
+
+        # FFW Block
+        ffw_residual_TD = attn_output_TD
+        normed_ffw_input_TD = self.pre_mlp_norm(attn_output_TD)
+        if isinstance(self.custom_module, MoE):
+            logits_TD = self.custom_module(normed_ffw_input_TD)
+            # Add the shared expert outputs to the MoE outputs.
+            shared_expert_output_TD = self.shared_experts(normed_ffw_input_TD)
+            logits_TD += shared_expert_output_TD
+        elif isinstance(self.custom_module, DenseFFW):
+            logits_TD = self.custom_module(normed_ffw_input_TD)
+        else:
+            raise ValueError(
+                f"Invalid custom moduel type: {type(self.custom_module)}")
+        logits_TD += ffw_residual_TD
+        return new_cache, logits_TD

tpu_inference/layers/vllm/attention.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools
+from typing import Optional, Tuple
+
+import jax
+import torch
+from jax.sharding import Mesh
+from torchax.interop import jax_view, torch_view
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionLayer, AttentionType)
+
+from tpu_inference import utils
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.jax.attention_interface import attention
+from tpu_inference.logger import init_logger
+from tpu_inference.models.vllm.vllm_model_wrapper_context import \
+    get_vllm_model_wrapper_context
+
+logger = init_logger(__name__)
+
+
+class PallasAttentionBackend(AttentionBackend):
+
+    @staticmethod
+    def get_name() -> str:
+        return "PALLAS"
+
+    @staticmethod
+    def get_impl_cls() -> type["PallasAttentionBackendImpl"]:
+        return PallasAttentionBackendImpl
+
+
+class PallasAttentionBackendImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[list[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        logits_soft_cap: Optional[float] = None,
+        attn_type: str = AttentionType.DECODER,
+        kv_sharing_target_layer_name: Optional[int] = None,
+        use_irope: bool = False,
+    ) -> None:
+        if use_irope:
+            logger.warning_once(
+                "Using irope in Pallas is not supported yet, it will fall back "
+                "to global attention for long context.")
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        self.sliding_window = sliding_window
+        self.logits_soft_cap = logits_soft_cap
+        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
+
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+        if alibi_slopes is not None:
+            raise NotImplementedError("Alibi slopes is not supported.")
+        self.kv_cache_quantized_dtype = None
+        if kv_cache_dtype != "auto":
+            self.kv_cache_quantized_dtype = utils.get_jax_dtype_from_str_dtype(
+                kv_cache_dtype)
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError("Encoder self-attention and "
+                                      "encoder/decoder cross-attention "
+                                      "are not implemented for "
+                                      "PallasAttentionBackendImpl")
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        output: Optional[torch.Tensor] = None,
+        output_scale: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if output_scale is not None:
+            raise NotImplementedError(
+                "fused output quantization is not yet supported for "
+                "PallasAttentionBackendImpl")
+
+        if kv_cache.numel():
+            raise RuntimeError(
+                "KV cache from vLLM Attention layer should be empty but has "
+                "the size of %s.", kv_cache.numel())
+
+        del kv_cache  # Use kv_cache from vllm wrapper context values instead.
+
+        vllm_model_wrapper_context = get_vllm_model_wrapper_context()
+        kv_cache_index = vllm_model_wrapper_context.layer_name_to_kvcache_index[
+            layer.layer_name]
+        kv_cache = vllm_model_wrapper_context.kv_caches[kv_cache_index]
+
+        mesh = vllm_model_wrapper_context.mesh
+
+        query, key, value = jax_view(query), jax_view(key), jax_view(value)
+        q_scale = k_scale = v_scale = None
+        if self.kv_cache_quantized_dtype:
+            key, value = utils.quantize_kv(key, value,
+                                           self.kv_cache_quantized_dtype,
+                                           layer._k_scale_float,
+                                           layer._v_scale_float)
+            # TODO(kyuyeunk): Enable w8a8 when VREG spill issue is resolved.
+            # q_scale = layer._q_scale_float
+            k_scale = layer._k_scale_float
+            v_scale = layer._v_scale_float
+
+        new_kv_cache, outputs = _jax_attn_func(kv_cache, query, key, value,
+                                               attn_metadata, mesh, self.scale,
+                                               self.head_size, self.num_heads,
+                                               self.num_kv_heads, q_scale,
+                                               k_scale, v_scale)
+        vllm_model_wrapper_context.kv_caches[kv_cache_index] = new_kv_cache
+
+        return torch_view(outputs)
+
+
+@functools.partial(
+    jax.jit,
+    static_argnums=(
+        5, 6, 7, 8, 9, 10, 11, 12
+    ),  # mesh, scale, head_size, num_heads, num_kv_heads, q_scale, k_scale, v_scale
+    donate_argnums=(0, ),  # donate kv_cache
+)
+def _jax_attn_func(
+    kv_cache: jax.Array,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    attention_metadata: AttentionMetadata,
+    mesh: Mesh,
+    scale: float,
+    head_size: int,
+    num_heads: int,
+    num_kv_heads: int,
+    q_scale: Optional[float] = None,
+    k_scale: Optional[float] = None,
+    v_scale: Optional[float] = None,
+) -> Tuple[jax.Array, jax.Array]:
+    del scale  # Unused for now, as the attention function applies a default scale.
+
+    # Get shapes from vllm
+    q_len, q_compute_dim = q.shape
+    k_len, k_compute_dim = k.shape
+    assert k.shape == v.shape
+    assert q_compute_dim == head_size * num_heads
+    assert k_compute_dim == head_size * num_kv_heads
+
+    # Convert the shapes from vLLM's convetion to what the attention function expects
+    # bs, num_heads, q_len, head_size
+    q = q.reshape(q_len, num_heads, head_size)
+    # bs, num_kv_heads, k_len, head_size
+    k = k.reshape(k_len, num_kv_heads, head_size)
+    v = v.reshape(k_len, num_kv_heads, head_size)
+
+    new_kv_cache, outputs = attention(
+        kv_cache,
+        q,
+        k,
+        v,
+        attention_metadata,
+        mesh,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+
+    # Convert the shape back to vLLM's convention
+    assert outputs.shape[0] == q_len
+    assert outputs.shape[1] == num_heads
+    assert outputs.shape[2] == head_size
+    outputs = outputs.reshape(q_len, q_compute_dim)
+
+    return new_kv_cache, outputs