deepseek_v3.py

# 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 math
import os
from abc import abstractmethod
from dataclasses import InitVar, dataclass
from itertools import islice
from typing import Iterable, List, Optional, Tuple, Union

import jax
import jax.numpy as jnp
from flax import nnx
from flax.typing import Sharding
from jax import lax
from jax.sharding import Mesh
from jax.sharding import PartitionSpec as P
from jaxtyping import Float
from vllm.config import VllmConfig

from tpu_inference import utils
from tpu_inference.distributed.jax_parallel_state import get_pp_group
from tpu_inference.kernels.quantized_matmul.util import quantize_tensor
from tpu_inference.kernels.ragged_paged_attention.v3.kernel import \
    ragged_paged_attention
from tpu_inference.layers.common.attention_interface import mla_attention
from tpu_inference.layers.common.moe import MoEBackend
from tpu_inference.layers.common.quantization import (dequantize_tensor,
                                                      quantize_kv)
from tpu_inference.layers.common.sharding import \
    ShardingAxisNameBase as ShardingAxisName
from tpu_inference.layers.common.utils import cpu_mesh_context
from tpu_inference.layers.jax import JaxModule
from tpu_inference.layers.jax.attention.attention import AttentionMetadata
from tpu_inference.layers.jax.base import _init_fn as init_fn
from tpu_inference.layers.jax.base import create_param, sharded_initializer
from tpu_inference.layers.jax.embed import JaxEmbed
from tpu_inference.layers.jax.layers import FlaxUtils
from tpu_inference.layers.jax.linear import JaxEinsum
from tpu_inference.layers.jax.moe.moe import JaxMoE
from tpu_inference.layers.jax.moe.utils import (get_expert_parallelism,
                                                select_moe_backend)
from tpu_inference.layers.jax.norm import JaxRmsNorm
from tpu_inference.layers.jax.pp_utils import PPMissingLayer, make_layers
from tpu_inference.layers.jax.quantization.configs import QuantizationConfig
from tpu_inference.layers.jax.rope import DeepseekScalingRotaryEmbedding
from tpu_inference.logger import init_logger
from tpu_inference.models.jax.jax_intermediate_tensor import \
    JaxIntermediateTensors
from tpu_inference.models.jax.utils.weight_utils import (JaxAutoWeightsLoader,
                                                         LoadableWithIterator,
                                                         shard_put)

KVCache = Tuple[jax.Array, jax.Array]

logger = init_logger(__name__)


def _weight_init(random_init: bool):
    return sharded_initializer if random_init else nnx.initializers.uniform()


modeling_flax_utils = FlaxUtils()

# TODO: read these configs from HF config.
num_local_experts: int = 256
vocab_size: int = 129280
hidden_size: int = 7168
num_attention_heads: int = 128
num_key_value_heads: int = 128
ffw_intermediate_size: int = 18432
moe_intermediate_size: int = 2048
num_experts_per_token: int = 8
n_group: int = 8
interleave_moe_layer_step: int = 1  # Deepseek V3 has moe_layer_freq=1 in hf config.
hidden_act: str = "silu"
rms_norm_eps: float = 1e-06
routed_scaling_factor: float = 2.5
first_k_dense_replace: int = 3  # replace the first few MOE layers to dense layer.

num_shared_experts = 1
rope_theta = 10000
rope_scaling = {
    "beta_fast": 32,
    "beta_slow": 1,
    "factor": 40,
    "mscale": 1.0,
    "mscale_all_dim": 1.0,
    "original_max_position_embeddings": 4096,
    "type": "yarn"
}
q_lora_rank = 1536
kv_lora_rank = 512
qk_nope_head_dim = 128
qk_rope_head_dim = 64
v_head_dim = 128
expert_axis_name = ShardingAxisName.ATTN_DATA_EXPERT


@dataclass(kw_only=True)
class DeepseekV3BaseAttention(JaxModule):
    """
    Base class containing shared logic for DeepSeek Attention mechanisms.
    Handles initialization of common layers and defines skeleton forward pass.
    """
    # Core configuration
    hidden_size: int
    num_attention_heads: int
    num_key_value_heads: int
    head_dim: int
    rope: DeepseekScalingRotaryEmbedding
    dtype: jnp.dtype
    kv_cache_dtype: str
    mesh: Mesh

    # Attention-specific configuration
    q_lora_rank: int
    kv_lora_rank: int
    qk_nope_head_dim: int
    qk_rope_head_dim: int
    v_head_dim: int
    rms_norm_eps: float

    # Sharding
    rd_sharding: P = P()
    q_da_sharding: P = P()
    ap_sharding: P = P()
    kv_da_sharding: P = P()
    activation_attention_td: P = P()
    activation_q_td: P = P()
    query_tnh: P = P()
    keyvalue_skh: P = P()
    attn_o_tnh: P = P()
    activation_attention_out_td: P = P()
    # Weight initialization
    random_init: bool = False
    rope_mscale_all_dim: float = 1.0

    # RNG for weight initialization
    rngs: InitVar[nnx.Rngs]

    quant_config: Optional[QuantizationConfig] = None

    # Scales for Q/KV quantization (per-tensor)
    _q_scale: float = 1
    _k_scale: float = 1
    _v_scale: float = 1

    prefix: str = ""

    def __post_init__(self, rngs: nnx.Rngs):
        self.N = self.num_attention_heads
        self.K = self.num_key_value_heads
        self.D = self.hidden_size
        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

        if self.rope.scaling_factor <= 1.0:
            yarn_mscale = 1.0
        else:
            yarn_mscale = 0.1 * self.rope_mscale_all_dim * math.log(
                self.rope.scaling_factor) + 1.0

        self.scale = self.qk_head_dim**-0.5 * yarn_mscale**2

        weight_init = _weight_init(self.random_init)

        self.q_a_proj = JaxEinsum(
            einsum_str="TD,DA->TA",
            kernel_shape=(self.D, self.q_lora_rank),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.q_da_sharding),
            prefix=self.prefix + ".q_a_proj",
        )

        self.q_b_proj = JaxEinsum(
            einsum_str="TA,AP->TP",
            kernel_shape=(self.q_lora_rank, self.N * self.qk_head_dim),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.ap_sharding),
            prefix=self.prefix + ".q_b_proj")

        self.kv_a_proj_with_mqa = JaxEinsum(
            einsum_str="SD,DA->SA",
            kernel_shape=(self.D, self.kv_lora_rank + self.qk_rope_head_dim),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init,
                                              self.kv_da_sharding),
            prefix=self.prefix + ".kv_a_proj_with_mqa")

        self.o_proj = JaxEinsum(
            einsum_str="TR,RD->TD",
            kernel_shape=(self.N * self.v_head_dim, self.D),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.rd_sharding),
            prefix=self.prefix + ".o_proj")

        self.q_a_layernorm = JaxRmsNorm(self.q_lora_rank,
                                        epsilon=self.rms_norm_eps,
                                        scale_init=nnx.with_partitioning(
                                            init_fn, (None, )),
                                        param_dtype=self.dtype,
                                        dtype=self.dtype,
                                        quant_config=self.quant_config,
                                        prefix=self.prefix + ".q_a_layernorm",
                                        rngs=rngs)

        self.kv_a_layernorm = JaxRmsNorm(
            self.kv_lora_rank,
            epsilon=self.rms_norm_eps,
            scale_init=nnx.with_partitioning(init_fn, (None, )),
            param_dtype=self.dtype,
            dtype=self.dtype,
            quant_config=self.quant_config,
            prefix=self.prefix + ".kv_a_layernorm",
            rngs=rngs)

        self.kv_cache_quantized_dtype = None
        if self.kv_cache_dtype != "auto":
            self.kv_cache_quantized_dtype = utils.get_jax_dtype_from_str_dtype(
                self.kv_cache_dtype)

        self.kv_b_proj = JaxEinsum(
            einsum_str="SA,AL->SL",
            kernel_shape=(self.kv_lora_rank,
                          self.N * (self.qk_nope_head_dim + self.v_head_dim)),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(init_fn, self.ap_sharding),
            prefix=self.prefix + ".kv_b_proj",
        )

    @abstractmethod
    def compute_q_projection(self, *args) -> jax.Array:
        raise NotImplementedError

    @abstractmethod
    def compute_kv_projection(self, *args) -> Tuple[jax.Array, jax.Array]:
        raise NotImplementedError

    @abstractmethod
    def compute_attention(self, *args) -> Tuple[KVCache, jax.Array]:
        raise NotImplementedError

    def process_output(self, outputs_TNH) -> jax.Array:
        return outputs_TNH

    def __call__(
            self, x: jax.Array, kv_cache: KVCache,
            attention_metadata: AttentionMetadata
    ) -> Tuple[KVCache, jax.Array]:
        """Performs the forward pass of the attention module.  Expects that the
        child class has implemented the `compute_q_projection`, `compute_kv_projection`,
        and `compute_attention` methods.

        Args:
            x: The input tensor of shape `(batch_size, seq_len, d_model)`.
            kv_cache: The key-value cache for storing past attention states.
            attention_metadata: Metadata for attention, such as input positions.

        Returns:
            A tuple containing:
                - The updated KV cache.
                - The attention output tensor of shape
                  `(batch_size, seq_len, d_model)`.
        """

        md = attention_metadata
        x = jnp.asarray(x, self.dtype)
        x_SD = lax.with_sharding_constraint(x, self.activation_attention_td)
        x_q_TD = lax.with_sharding_constraint(x, self.activation_q_td)

        with jax.named_scope("q_proj"):
            q_data = self.compute_q_projection(x_q_TD, md.input_positions)

        with jax.named_scope("kv_proj"):
            kv_data = self.compute_kv_projection(x_SD, md.input_positions)

        with jax.named_scope("attn_op"):
            new_kv_cache, outputs_TNH = self.compute_attention(
                q_data, kv_data, kv_cache, md)

            outputs_TNH = self.process_output(outputs_TNH)

            if outputs_TNH.shape[-1] != self.v_head_dim:
                outputs_TNH = outputs_TNH[..., :self.v_head_dim]

            with jax.named_scope("o_proj"):
                outputs_TR = outputs_TNH.reshape(outputs_TNH.shape[0],
                                                 self.N * self.v_head_dim)
                o_TD = self.o_proj(outputs_TR)

            return new_kv_cache, o_TD


@dataclass(kw_only=True)
class DeepseekV3Attention(DeepseekV3BaseAttention):
    """Standard Multi-Head Attention (MHA) for DeepSeek models."""

    def __post_init__(self, rngs: nnx.Rngs):
        super().__post_init__(rngs)

        weight_init = _weight_init(self.random_init)
        self.kv_b_proj = JaxEinsum(
            einsum_str="SA,AL->SL",
            kernel_shape=(self.kv_lora_rank,
                          self.N * (self.qk_nope_head_dim + self.v_head_dim)),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.ap_sharding),
            prefix=self.prefix + ".kv_b_proj",
        )

    def compute_q_projection(self, x_q_TD: jax.Array,
                             input_positions: jax.Array) -> jax.Array:
        """
        Computes the query projection for MHA.

        Args:
            x_q_TD: The input tensor of shape `(tokens_query, d_model)`.
            input_positions: The input positions tensor of shape `(padded_total_num_scheduled_tokens,)`.

        Returns:
            The query tensor of shape `(tokens_query, num_query_heads, head_dim)`.
        """
        q_TA = self.q_a_proj(x_q_TD)
        q_TA = self.q_a_layernorm(q_TA)
        q_TP = self.q_b_proj(q_TA)
        q_TNH = q_TP.reshape(q_TA.shape[0], self.N, self.qk_head_dim)

        q_nope_TNH = q_TNH[..., :self.qk_nope_head_dim]
        q_rope_TNH = q_TNH[..., self.qk_nope_head_dim:]
        q_rope_TNH = self.rope.apply_rope(input_positions, q_rope_TNH)
        q_TNH = jnp.concatenate([q_nope_TNH, q_rope_TNH], axis=-1)

        return lax.with_sharding_constraint(q_TNH, self.query_tnh)

    def compute_kv_projection(
            self, x_SD: jax.Array,
            input_positions: jax.Array) -> Tuple[jax.Array, jax.Array]:
        """
        Computes the key-value projection for MHA.

        Args:
            x_SD: The input tensor of shape `(tokens_kv, d_model)`.
            input_positions: The input positions tensor of shape `(padded_total_num_scheduled_tokens,)`.

        Returns:
            Tuple of key-value tensors of shape `(tokens_kv, num_query_heads, d_model)`.
        """

        kv_SA = self.kv_a_proj_with_mqa(x_SD)

        k_rope_SH = kv_SA[..., self.kv_lora_rank:]
        k_rope_SNH = k_rope_SH[..., None, :]
        k_rope_SNH = self.rope.apply_rope(input_positions, k_rope_SNH)
        assert k_rope_SNH.shape[1] == 1

        k_rope_SNH = jnp.broadcast_to(
            k_rope_SNH, (k_rope_SNH.shape[0], self.N, self.qk_rope_head_dim))

        kv_SA = kv_SA[..., :self.kv_lora_rank]
        kv_SA = self.kv_a_layernorm(kv_SA)
        kv_SA = lax.with_sharding_constraint(kv_SA, self.keyvalue_skh)

        kv_SL = self.kv_b_proj(kv_SA)
        kv_nope_SNH = kv_SL.reshape(kv_SA.shape[0], self.N,
                                    self.qk_nope_head_dim + self.v_head_dim)

        k_nope_SNH = kv_nope_SNH[..., :self.qk_nope_head_dim]
        v_SNH = kv_nope_SNH[..., self.qk_nope_head_dim:]

        k_SNH = jnp.concatenate([k_nope_SNH, k_rope_SNH], axis=-1)

        # Shard
        k_SNH = lax.with_sharding_constraint(k_SNH, self.keyvalue_skh)
        v_SNH = lax.with_sharding_constraint(v_SNH, self.keyvalue_skh)

        return (k_SNH, v_SNH)

    def compute_attention(self, q_data: jax.Array, kv_data: Tuple[jax.Array,
                                                                  jax.Array],
                          kv_cache: KVCache,
                          md: AttentionMetadata) -> Tuple[jax.Array, KVCache]:
        """
        Computes self-attention for MHA.

        Args:
            q_data: The query tensor of shape `(tokens_query, num_query_heads, head_dim)`.
            kv_data: Tuple of key-value tensors of shape `(tokens_kv, num_query_heads, d_model)`.
            kv_cache: KVCache object.
            md: AttentionMetadata object.

        Returns:
            Tuple of output tensors of shape `(tokens_query, num_query_heads, head_dim)` and KVCache object.
        """

        q_TNH = q_data
        k_SNH, v_SNH = kv_data

        multiple_of_128 = ((self.qk_head_dim - 1) // 128 + 1) * 128
        q_TNH = jnp.pad(q_TNH, ((0, 0), (0, 0),
                                (0, multiple_of_128 - self.qk_head_dim)))
        k_SNH = jnp.pad(k_SNH, ((0, 0), (0, 0),
                                (0, multiple_of_128 - self.qk_head_dim)))
        v_SNH = jnp.pad(v_SNH, ((0, 0), (0, 0),
                                (0, multiple_of_128 - self.v_head_dim)))

        q_scale = k_scale = v_scale = None
        if self.kv_cache_quantized_dtype:
            k_scale = self._k_scale
            v_scale = self._v_scale
            k_SNH, v_SNH = quantize_kv(self.kv_cache_quantized_dtype, k_SNH,
                                       v_SNH, k_scale, v_scale)

        def _ragged_paged_attention(q, k, v, cache, seq_lens, block_tables,
                                    starts, dist):
            return ragged_paged_attention(q,
                                          k,
                                          v,
                                          cache,
                                          seq_lens,
                                          block_tables,
                                          starts,
                                          dist,
                                          sm_scale=self.scale,
                                          q_scale=q_scale,
                                          k_scale=k_scale,
                                          v_scale=v_scale)

        in_specs = (
            self.query_tnh,  # q
            self.keyvalue_skh,  # k
            self.keyvalue_skh,  # v
            P(None, None, ShardingAxisName.ATTN_HEAD),  # kv_cache
            P(),  # md.seq_lens: Replicated
            P(),  # page_indices_flat: Replicated
            P(),  # query_start_loc: Replicated
            P(),  # distribution: Replicated
        )

        out_specs = (self.attn_o_tnh, P(None, None,
                                        ShardingAxisName.ATTN_HEAD))

        output_TNH, kv_cache = jax.jit(
            jax.shard_map(_ragged_paged_attention,
                          mesh=self.mesh,
                          in_specs=in_specs,
                          out_specs=out_specs,
                          check_vma=False))(q_TNH, k_SNH, v_SNH, kv_cache,
                                            md.seq_lens, md.block_tables,
                                            md.query_start_loc,
                                            md.request_distribution)

        return kv_cache, output_TNH


class MLAEinsum(JaxEinsum):
    """Extending JaxEinsum to handle MLA.
    
    This class is used for MLA, where:
    1) the weight is split into k/v parts after loading, and
    2) modify the MLA layer to set k/v weights
    """

    def __init__(self,
                 mla_layer,
                 einsum_str: str,
                 kernel_shape: tuple[int, ...],
                 rngs,
                 bias_shape: Optional[tuple[int, ...]] = None,
                 quant_config: Optional[QuantizationConfig] = None,
                 prefix: str = "",
                 **kwargs):
        super().__init__(einsum_str,
                         kernel_shape,
                         rngs,
                         bias_shape=bias_shape,
                         quant_config=quant_config,
                         prefix=prefix,
                         **kwargs)
        self.loaded = set()
        self.mla_layer = mla_layer
        self.quant_config = quant_config

    def named_children(self):
        # Override, otherwise "mla_layer" will be visited, causing infinite recursion.
        yield from []

    def load_weights(self, weights):
        named_params = dict(self.named_parameters())
        if len(self.loaded) >= 2:
            raise ValueError(
                f"Expect at most 2 params to load for kv_b_proj, already got {self.loaded}, still have {[name for name, _ in weights]} coming."
            )
        for name, weight in weights:
            param = named_params[name]
            weight_loader = getattr(param, "weight_loader")
            weight_loader(param, weight)
            self.loaded.add(name)
        if len(self.loaded) != len(named_params):
            return
        assert self.quant_config is not None
        # After loading, split the weights into k/v
        with cpu_mesh_context():
            dequantized_weight = dequantize_tensor(
                self.weight,
                self.weight_scale_inv,
                (0, 1),
                block_size=None,
            ).T
            A, N, qk_nope_head_dim, v_head_dim = self.mla_layer.kv_lora_rank, self.mla_layer.N, self.mla_layer.qk_nope_head_dim, self.mla_layer.v_head_dim
            if dequantized_weight.shape != (A, N *
                                            (qk_nope_head_dim + v_head_dim)):
                raise ValueError(
                    f"Unexpected weight shape after dequantization: {dequantized_weight.shape}, expected {(A, N * (qk_nope_head_dim + v_head_dim))=}"
                )
            dequantized_weight = dequantized_weight.reshape(
                A, N, qk_nope_head_dim + v_head_dim)
            k_ANH, v_ANH = jnp.split(dequantized_weight, [qk_nope_head_dim],
                                     axis=-1)
            k_ANH_weight, k_ANH_scale = quantize_tensor(k_ANH,
                                                        self.weight.dtype,
                                                        dim=-1)
            v_ANH_weight, v_1NH_scale = quantize_tensor(v_ANH,
                                                        self.weight.dtype,
                                                        dim=0)
            # As of writing, sharded_quantized_batched_matmul expects scale to be
            # a different shape order than weight
            k_N1A_scale = k_ANH_scale.transpose(1, 2, 0)
            v_N1H_scale = v_1NH_scale.transpose(1, 0, 2)
        mla_layer = self.mla_layer
        setattr(
            mla_layer, "k_up_proj",
            JaxEinsum(
                einsum_str="TNH,ANH->TNA",
                kernel_shape=(A, N, qk_nope_head_dim),
                rngs=nnx.Rngs(0),
                prefix=mla_layer.prefix + ".k_up_proj",
                quant_config=self.quant_config,
            ))
        setattr(
            mla_layer, "v_up_proj",
            JaxEinsum(
                einsum_str="TNA,ANH->TNH",
                kernel_shape=(A, N, v_head_dim),
                rngs=nnx.Rngs(0),
                prefix=mla_layer.prefix + ".v_up_proj",
                quant_config=self.quant_config,
            ))
        # Cannot apply anh_sharding to scales, otherwise it complains about shape mismatch.
        mla_layer.k_up_proj.weight.value = shard_put(
            k_ANH_weight, self.mla_layer.anh_sharding)
        mla_layer.k_up_proj.weight_scale_inv.value = shard_put(k_N1A_scale, ())
        mla_layer.v_up_proj.weight.value = shard_put(
            v_ANH_weight, self.mla_layer.anh_sharding)
        mla_layer.v_up_proj.weight_scale_inv.value = shard_put(v_N1H_scale, ())

        delattr(self, 'weight')
        delattr(self, 'weight_scale_inv')
        delattr(self, 'quant_method')


@dataclass(kw_only=True)
class DeepseekV3MLA(DeepseekV3BaseAttention):
    """Multi-Head Latent Attention (MLA) for DeepSeek V3."""
    anh_sharding: Sharding = ()

    def __post_init__(self, rngs: nnx.Rngs):
        super().__post_init__(rngs)

        weight_init = _weight_init(self.random_init)
        self.kv_b_proj = MLAEinsum(
            mla_layer=self,
            einsum_str="SA,AL->SL",
            kernel_shape=(self.kv_lora_rank,
                          self.N * (self.qk_nope_head_dim + self.v_head_dim)),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.ap_sharding),
            prefix=self.prefix + ".kv_b_proj",
        )

    def compute_q_projection(
            self, x_q_TD: jax.Array,
            input_positions: jax.Array) -> Tuple[jax.Array, jax.Array]:
        """
        Computes the query projection for MLA.

        Args:
            x_q_TD: The input tensor of shape `(tokens_query, d_model)`.
            input_positions: The input positions tensor of shape `(padded_total_num_scheduled_tokens,)`.

        Returns:
            A tuple of query tensor of shape `(tokens_query, num_query_heads, q_lora_rank)` and
            rope tensor of shape `(tokens_query, num_query_heads, head_dim)`.
        """
        q_TA = self.q_a_proj(x_q_TD)
        q_TA = self.q_a_layernorm(q_TA)
        q_TP = self.q_b_proj(q_TA)
        q_TNH = q_TP.reshape(q_TA.shape[0], self.N, self.qk_head_dim)

        q_nope_TNH = q_TNH[..., :self.qk_nope_head_dim]
        q_rope_TNH = q_TNH[..., self.qk_nope_head_dim:]
        q_rope_TNH = self.rope.apply_rope(input_positions, q_rope_TNH)

        q_TNA = self.k_up_proj(q_nope_TNH)

        q_TNA = lax.with_sharding_constraint(q_TNA, self.query_tnh)
        return (q_TNA, q_rope_TNH)

    def compute_kv_projection(
            self, x_SD: jax.Array,
            input_positions: jax.Array) -> Tuple[jax.Array, jax.Array]:
        """
        Computes the key-value projection for MLA.

        Args:
            x_SD: The input tensor of shape `(tokens_kv, d_model)`.
            input_positions: The input positions tensor of shape `(padded_total_num_scheduled_tokens,)`.

        Returns:
            A tuple of key-value tensor of shape `(tokens_kv, q_lora_rank)` and
            rope tensor of shape `(tokens_kv, head_dim)`.
        """
        kv_SA = self.kv_a_proj_with_mqa(x_SD)

        k_rope_SH = kv_SA[..., self.kv_lora_rank:]
        k_rope_SNH = k_rope_SH[..., None, :]
        k_rope_SNH = self.rope.apply_rope(input_positions, k_rope_SNH)
        assert k_rope_SNH.shape[1] == 1
        k_rope_SH = k_rope_SNH[:, 0, :]

        kv_SA = kv_SA[..., :self.kv_lora_rank]
        kv_SA = self.kv_a_layernorm(kv_SA)
        kv_SA = lax.with_sharding_constraint(kv_SA, self.keyvalue_skh)

        return (kv_SA, k_rope_SH)

    def compute_attention(self, q_data: Tuple[jax.Array, jax.Array],
                          kv_data: Tuple[jax.Array,
                                         jax.Array], kv_cache: KVCache,
                          md: AttentionMetadata) -> Tuple[KVCache, jax.Array]:
        """
        Computes the attention for MLA.

        Args:
            q_data: A tuple of query tensor of shape `(tokens_query, num_query_heads, q_lora_rank)` and
                rope tensor of shape `(tokens_query, num_query_heads, head_dim)`.
            kv_data: A tuple of key-value tensor of shape `(tokens_kv, q_lora_rank)` and
                rope tensor of shape `(tokens_kv, head_dim)`.
            kv_cache: The key-value cache.
            md: The attention metadata.

        Returns:
            A tuple of key-value cache and output tensor of shape `(tokens_query, num_query_heads, q_lora_rank)`.
        """

        q_TNA, q_rope_TNH = q_data
        k_SA, k_rope_SH = kv_data

        q_scale = k_scale = None
        if self.kv_cache_quantized_dtype:
            k_scale = self._k_scale
            # TODO: May need to apply quantization separately for k_c & k_pe
            k_SA, _ = quantize_kv(self.kv_cache_quantized_dtype,
                                  k_SA,
                                  value=None,
                                  k_scale=k_scale)
            k_rope_SH, _ = quantize_kv(self.kv_cache_quantized_dtype,
                                       k_rope_SH,
                                       value=None,
                                       k_scale=k_scale)

        return mla_attention(q_TNA,
                             q_rope_TNH,
                             k_SA,
                             k_rope_SH,
                             kv_cache,
                             md,
                             self.mesh,
                             self.num_attention_heads,
                             self.qk_nope_head_dim,
                             query_tnh_sharding=self.query_tnh,
                             keyvalue_skh_sharding=self.keyvalue_skh,
                             attn_o_tnh_sharding=self.attn_o_tnh,
                             q_scale=q_scale,
                             k_scale=k_scale,
                             v_scale=k_scale,
                             sm_scale=self.scale)

    def process_output(self, outputs_TNA: jax.Array) -> jax.Array:
        """
        Processes output for MLA specifically.

        Args:
            outputs_TNH: The output tensor of shape `(tokens_query, num_query_heads, q_lora_rank)`.

        Returns:
            The processed output tensor of shape `(tokens_query, num_query_heads, head_dim)`.
        """

        # MLA Specific: Apply V-Up Projection after attention
        # Outputs from MLA kernel are in latent space (TNA), project to TNH
        outputs_TNH = self.v_up_proj(outputs_TNA)
        return outputs_TNH


@dataclass(kw_only=True)
class DeepseekV3MLP(JaxModule):
    """A Gated Feed-Forward Network (FFN) layer.

    This module consists of two linear projections (gating and up-projection),
    an element-wise multiplication of the activated gating projection and the
    up-projection, followed by a final downward projection.

    Attributes:
        sharding_cfg: The configuration for tensor sharding.
    """
    dtype: jnp.dtype
    hidden_act: str
    hidden_size: int
    intermediate_size: int
    df_sharding: P = P()
    fd_sharding: P = P()
    activation_ffw_td: P = P()
    random_init: bool = False
    quant_config: Optional[QuantizationConfig] = None

    rngs: InitVar[nnx.Rngs]

    def __call__(self, x_TD):
        """Performs the forward pass of the FFW layer.

        Args:
            x_TD: The input tensor of shape either `(sequence, d_model)`

        Returns:
            The output tensor of shape `(batch, sequence, d_model)`.
        """
        x_TD = jnp.asarray(x_TD, self.dtype)
        x_TD = lax.with_sharding_constraint(x_TD, self.activation_ffw_td)
        with jax.named_scope("wi_0"):
            gating_TF = self.gate_proj(x_TD)
            activated_gating_TF = modeling_flax_utils.ACT2FN[self.hidden_act](
                gating_TF)
        with jax.named_scope("wi_1"):
            up_proj_TF = self.up_proj(x_TD)
        fuse_TF = activated_gating_TF * up_proj_TF
        with jax.named_scope("wo"):
            output_TD = self.down_proj(fuse_TF)

        return output_TD

    def __post_init__(self, rngs: nnx.Rngs):
        D = self.hidden_size
        F = self.intermediate_size
        weight_init = _weight_init(self.random_init)

        self.gate_proj = JaxEinsum(
            einsum_str="TD,DF->TF",
            kernel_shape=(D, F),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.df_sharding),
        )
        self.up_proj = JaxEinsum(
            einsum_str="TD,DF->TF",
            kernel_shape=(D, F),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.df_sharding),
        )
        self.down_proj = JaxEinsum(
            einsum_str="TF,FD->TD",
            kernel_shape=(F, D),
            rngs=rngs,
            quant_config=self.quant_config,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.fd_sharding),
        )


@dataclass(kw_only=True)
class SharedFusedMoe(JaxMoE):
    """
    Corresponds to vLLM's SharedFusedMoe.
    Handles the routed and shared experts + the relevant forward pass.

    Reference here: https://github.com/vllm-project/vllm/blob/168ee03e1cbba2b962adbc704b16762b266be184/vllm/model_executor/layers/fused_moe/shared_fused_moe.py#L14
    """
    shared_experts: Optional[DeepseekV3MLP] = None

    routed_scaling_factor: float = 1.0

    def __call__(self, x_TD: jax.Array) -> jax.Array:
        # Compute Routed Experts
        final_hidden_states = super().__call__(x_TD)

        # (Maybe) Compute Shared Experts
        if self.shared_experts is not None:
            shared_output = self.shared_experts(x_TD)
            final_hidden_states += shared_output

        return final_hidden_states


class DeepseekV2Moe(JaxModule):
    """Jax implementation of Deepseek MoE layer
    
    vllm ref. https://github.com/vllm-project/vllm/blob/168ee03e1cbba2b962adbc704b16762b266be184/vllm/model_executor/models/deepseek_v2.py#L225
    """

    def __init__(self,
                 *,
                 mesh,
                 dtype,
                 num_expert_parallelism,
                 moe_backend,
                 quant_config,
                 scoring_func,
                 rng,
                 prefix: str = ""):

        self.gate = DeepSeekV3Router(
            hidden_size=hidden_size,
            num_experts=num_local_experts,
            num_experts_per_tok=num_experts_per_token,
            n_groups=n_group,
            topk_groups=4,
            norm_topk_prob=True,
            rngs=rng,
            routed_scaling_factor=routed_scaling_factor,
            dtype=dtype,
            moe_backend=moe_backend,
            activation_ffw_td=(ShardingAxisName.MLP_DATA, None),
            ed_sharding=(None, None),
            e_sharding=(None, ),
            scoring_func=scoring_func,
            quant_config=quant_config)

        # shared experts
        self.shared_experts = DeepseekV3MLP(
            dtype=dtype,
            hidden_act=hidden_act,
            hidden_size=hidden_size,
            intermediate_size=num_shared_experts * moe_intermediate_size,
            rngs=rng,
            activation_ffw_td=P(ShardingAxisName.MLP_DATA, None),
            df_sharding=P(None, ShardingAxisName.ATTN_HEAD),
            fd_sharding=P(ShardingAxisName.ATTN_HEAD, None),
            quant_config=quant_config)

        # routed experts
        if moe_backend == MoEBackend.GMM_TP:
            moe_activation_ffw_td = P(ShardingAxisName.MLP_DATA, None)
            moe_activation_ffw_ted = P(ShardingAxisName.MLP_DATA, None,
                                       ShardingAxisName.MOE_TENSOR)
            moe_edf_sharding = P(None, ShardingAxisName.ATTN_DATA_EXPERT,
                                 ShardingAxisName.MOE_TENSOR)
            moe_efd_sharding = P(None, ShardingAxisName.MOE_TENSOR,
                                 ShardingAxisName.ATTN_DATA_EXPERT)
        else:
            moe_activation_ffw_td = P(ShardingAxisName.MLP_DATA,
                                      ShardingAxisName.MOE_TENSOR)
            moe_activation_ffw_ted = P(ShardingAxisName.MLP_DATA, None,
                                       ShardingAxisName.MOE_TENSOR)
            moe_edf_sharding = P(ShardingAxisName.ATTN_DATA_EXPERT, None, None)
            moe_efd_sharding = P(ShardingAxisName.ATTN_DATA_EXPERT, None, None)

        self.experts = SharedFusedMoe(
            dtype=dtype,
            num_local_experts=num_local_experts,
            apply_expert_weight_before_computation=False,
            expert_axis_name=expert_axis_name,
            num_expert_parallelism=num_expert_parallelism,
            hidden_size=hidden_size,
            intermediate_size_moe=moe_intermediate_size,
            num_experts_per_tok=num_experts_per_token,
            mesh=mesh,
            hidden_act=hidden_act,
            rngs=rng,
            quant_config=quant_config,
            activation_ffw_td=moe_activation_ffw_td,
            activation_ffw_ted=moe_activation_ffw_ted,
            edf_sharding=moe_edf_sharding,
            efd_sharding=moe_efd_sharding,
            moe_backend=moe_backend,
            qwix_quantized_weight_dtype=None,
            # It's abnormal prefix here because we are using dataclass for SharedFusedMoe and JaxMoe.
            # The proper way is to change both to normal class, set prefix=prefix+".mlp" here,
            # then in __init__, pass prefix+".experts" to super().__init__.
            prefix=f"{prefix}.experts",
            router=self.gate,
            shared_experts=self.shared_experts,
            scoring_func=scoring_func,
            routed_scaling_factor=routed_scaling_factor)

    def __call__(self, x_TD: jax.Array):
        return self.experts(x_TD)


class DeepseekV3DecoderLayer(JaxModule):
    """
    Implementats the DecoderLayer for DeepseekV3.
    """

    def __init__(
            self,
            input_layernorm: JaxRmsNorm,
            post_attention_layernorm: JaxRmsNorm,
            self_attn: Union[DeepseekV3Attention, DeepseekV3MLA],

            # MLP can be either the Dense MLP (for first k layers) or SharedFusedMoe
            mlp: nnx.Module | SharedFusedMoe | DeepseekV3MLP,
            prefix: str = ""):
        self.input_layernorm = input_layernorm
        self.post_attention_layernorm = post_attention_layernorm
        self.self_attn = self_attn
        self.mlp = mlp

    def __call__(
        self, x_TD: jax.Array, *, kv_cache: List[jax.Array],
        attention_metadata: AttentionMetadata
    ) -> Tuple[List[jax.Array], jax.Array]:

        # Run Self-Attention
        residual = x_TD
        hidden_states = self.input_layernorm(x_TD)
        new_cache, attn_output = self.self_attn(hidden_states, kv_cache,
                                                attention_metadata)
        hidden_states = residual + attn_output

        # Run MLP/MoE
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        mlp_output = self.mlp(hidden_states)

        # Residual
        hidden_states = residual + mlp_output

        return new_cache, hidden_states


class DeepSeekV3Router(JaxEinsum):
    """Router module for Mixture-of-Experts (MoE) layers.

    This module determines which experts each token should be routed to based on the input.
    """

    def __init__(
            self,
            hidden_size: int,
            num_experts: int,
            num_experts_per_tok: int,
            n_groups: int,
            topk_groups: int,
            norm_topk_prob: bool,
            routed_scaling_factor,
            dtype: jnp.dtype,
            rngs: nnx.Rngs,
            # Sharding Attributes
            activation_ffw_td: P = P(),
            ed_sharding: P = P(),
            e_sharding: P = P(),
            random_init: bool = False,
            quant_config: Optional[QuantizationConfig] = None,
            router_bias_dtype: jnp.dtype = jnp.float32,
            scoring_func: str = "sigmoid",
            moe_backend: MoEBackend = MoEBackend.DENSE_MAT):
        self.hidden_size = hidden_size
        self.num_experts = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.n_groups = n_groups
        self.topk_groups = topk_groups
        self.norm_topk_prob = norm_topk_prob
        self.routed_scaling_factor = routed_scaling_factor
        self.dtype = dtype
        self.activation_ffw_td = activation_ffw_td
        self.ed_sharding = ed_sharding
        self.e_sharding = e_sharding
        self.random_init = random_init
        self.quant_config = quant_config
        self.router_bias_dtype = router_bias_dtype
        self.scoring_func = scoring_func
        self.moe_backend = moe_backend
        """Generates the router kernel (weights and bias) for routing."""
        D = self.hidden_size
        E = self.num_experts
        weight_init = _weight_init(self.random_init)
        JaxEinsum.__init__(
            self,
            einsum_str="TD,DE->TE",
            kernel_shape=(D, E),
            rngs=rngs,
            # DS model has gate weights unquantized, but not mentioned in the config.
            quant_config=None,
            param_dtype=self.dtype,
            kernel_init=nnx.with_partitioning(weight_init, self.ed_sharding),
        )
        self.e_score_correction_bias = create_param(
            rngs,
            shape=(E, ),
            dtype=self.router_bias_dtype,
            sharding=self.e_sharding,
            random_init=self.random_init)

    def get_topk_indices(self, scores_TE: Float) -> Float:
        """Get the topk indices of the scores.

        Args:
            scores_TE: The scores to get the topk indices of. Shape (sequence, num_experts).

        Returns:
            The topk indices of the scores. Shape (sequence, num_experts_per_tok).
        """

        scores_TE = scores_TE + self.e_score_correction_bias
        if self.n_groups > 1:
            experts_per_group = self.num_experts // self.n_groups
            group_scores_TGM = jnp.reshape(
                scores_TE, (-1, self.n_groups, experts_per_group))
            group_scores_TG2 = jax.lax.top_k(group_scores_TGM, k=2)[0]
            group_scores_TG = jnp.sum(group_scores_TG2, axis=-1)
            group_indices = jax.lax.top_k(group_scores_TG,
                                          k=self.topk_groups)[1]

            # Apply mask at the group level before flattening
            mask_TG1 = jax.nn.one_hot(
                group_indices,
                self.n_groups).sum(axis=1)[..., None].astype(jnp.bool_)

            # Apply mask to each group of experts
            group_scores_TGM = jnp.where(mask_TG1, group_scores_TGM, -jnp.inf)

            scores_TE = jnp.reshape(group_scores_TGM, (-1, self.num_experts))

        indices_TX = jax.lax.top_k(scores_TE, k=self.num_experts_per_tok)[1]

        return indices_TX

    def __call__(self, x_TD: Float) -> Tuple[Float, Float]:
        """Routes tokens to top k experts.

        Args:
            x_TD: Input array of shape (sequence, d_model).

        Returns:
            A tuple containing:
                - weights: Normalized weights for selected experts, shape (sequence, num_experts_per_tok).
                - indices: Indices of selected experts, shape (sequence, num_experts_per_tok).
        """
        x_TD = jnp.asarray(x_TD, self.dtype)
        x_TD = jax.lax.with_sharding_constraint(x_TD,
                                                P(*self.activation_ffw_td))

        logits_TE = super().__call__(x_TD).astype(jnp.float32)

        # Apply scoring function (Sigmoid/Softmax) to get probabilities
        if self.scoring_func == "sigmoid":
            probs_TE = jax.nn.sigmoid(logits_TE)
        elif self.scoring_func == "softmax":
            probs_TE = jax.nn.softmax(logits_TE, axis=-1)
        else:
            probs_TE = logits_TE

        # Will add Aux-Loss-Free bias the activation outputs during topk selection.
        topk_indices_TX = self.get_topk_indices(probs_TE)

        # The actual weights do not include the bias terms.
        weights_TX = jnp.take_along_axis(probs_TE, topk_indices_TX, axis=-1)

        if self.norm_topk_prob:
            weights_TX /= jnp.sum(weights_TX, axis=-1)[..., None] + 1e-20

        # Scale expert weights before taking linear combination of experts.
        weights_TX *= self.routed_scaling_factor

        return weights_TX.astype(self.dtype), topk_indices_TX


@dataclass
class DeepSeekV3(JaxModule):

    def __init__(self,
                 vllm_config: VllmConfig,
                 rng: nnx.Rngs,
                 mesh: Mesh,
                 quant_config,
                 prefix: str = ""):
        self.vllm_config = vllm_config

        self.use_mla_kernel: bool = self.vllm_config.model_config.use_mla

        logger.info(f"Is using MLA kernel in DeepSeek: {self.use_mla_kernel}")

        self.mesh = mesh

        self.num_expert_parallelism = get_expert_parallelism(
            expert_axis_name, self.mesh)
        total_tensor_parallelsim = self.vllm_config.sharding_config.tp_size * \
                                        self.vllm_config.sharding_config.attn_dp_size
        self.use_ep = self.num_expert_parallelism > 1 and total_tensor_parallelsim == 1
        self.moe_backend = select_moe_backend(self.use_ep)

        # TODO (jacobplatin): we will resolve this issue in a forthcoming PR that will refactor weight loading
        if vllm_config.load_config.load_format == "dummy" and self.moe_backend in MoEBackend.fused_moe_backends(
        ):
            raise ValueError(
                f"Random / dummy weights are not supported for {MoEBackend.fused_moe_backends()} backends right now."
            )

        self.is_first_rank = get_pp_group().is_first_rank
        self.is_last_rank = get_pp_group().is_last_rank
        hf_config = vllm_config.model_config.hf_config
        dtype = vllm_config.model_config.dtype
        scoring_func = getattr(hf_config, "scoring_func", "sigmoid")

        if self.is_first_rank:
            self.embed_tokens = JaxEmbed(
                num_embeddings=vocab_size,
                features=hf_config.hidden_size,
                param_dtype=dtype,
                dtype=dtype,
                embedding_init=nnx.with_partitioning(
                    init_fn, (ShardingAxisName.MLP_TENSOR, )),
                rngs=rng,
                quant_config=quant_config,
                prefix=prefix + ".embed_tokens",
            )
        else:
            self.embed_tokens = PPMissingLayer()

        self.rope_emb = DeepseekScalingRotaryEmbedding(
            rotary_dim=qk_rope_head_dim,
            rope_theta=rope_theta,
            original_max_position_embeddings=rope_scaling[
                "original_max_position_embeddings"],
            scaling_factor=rope_scaling["factor"],
            dtype=dtype,
            beta_fast=rope_scaling["beta_fast"],
            beta_slow=rope_scaling["beta_slow"],
            mscale_value=rope_scaling["mscale"],
            mscale_all_dim=rope_scaling["mscale_all_dim"],
        )

        def _create_deepseek_attention(
                i: int) -> Union[DeepseekV3MLA, DeepseekV3Attention]:
            if self.use_mla_kernel:
                query_tnh_spec = P(ShardingAxisName.MLP_TENSOR, None, None)
                keyvalue_skh_spec = P(ShardingAxisName.MLP_TENSOR, None)
                attn_o_tnh_spec = P(ShardingAxisName.MLP_TENSOR, None, None)
                anh_sharding = (None, ShardingAxisName.MLP_TENSOR, None)
            else:
                query_tnh_spec = P(None, ShardingAxisName.MLP_TENSOR)
                keyvalue_skh_spec = P(None, ShardingAxisName.MLP_TENSOR)
                attn_o_tnh_spec = P(None, ShardingAxisName.MLP_TENSOR)
            rd_sharding = P(ShardingAxisName.ATTN_HEAD, None)
            ap_sharding = P(None, ShardingAxisName.ATTN_HEAD)
            q_da_sharding = P(None, ShardingAxisName.ATTN_HEAD)
            kv_da_sharding = P(None, ShardingAxisName.ATTN_HEAD)

            if self.vllm_config.additional_config.get("replicate_attn_weights",
                                                      False):
                rd_sharding = ()
                ap_sharding = ()
                q_da_sharding = ()
                kv_da_sharding = ()
                if self.use_mla_kernel:
                    anh_sharding = ()

            attn_cls = None
            if self.use_mla_kernel:
                attn_cls = DeepseekV3MLA
            else:
                attn_cls = DeepseekV3Attention
                assert num_attention_heads == num_key_value_heads, "Expected same number of of attention heads and key value heads for MHA."

            kwargs = dict(
                q_lora_rank=q_lora_rank,
                kv_lora_rank=kv_lora_rank,
                qk_nope_head_dim=qk_nope_head_dim,
                qk_rope_head_dim=qk_rope_head_dim,
                rms_norm_eps=rms_norm_eps,
                v_head_dim=v_head_dim,
                mesh=self.mesh,
                hidden_size=hidden_size,
                num_attention_heads=num_attention_heads,
                num_key_value_heads=1
                if self.use_mla_kernel else num_key_value_heads,
                head_dim=v_head_dim,  # MLA uses v_head_dim as head_dim
                rope=self.rope_emb,
                rope_mscale_all_dim=rope_scaling["mscale_all_dim"],
                dtype=dtype,
                # TODO (jacobplatin): we should refactor this to pass a dtype (or config) directly
                kv_cache_dtype=vllm_config.cache_config.cache_dtype,
                rngs=rng,
                quant_config=quant_config,
                activation_attention_td=P(None, None),
                activation_q_td=P(None, None),
                query_tnh=query_tnh_spec,
                keyvalue_skh=keyvalue_skh_spec,
                activation_attention_out_td=P(None, None),
                attn_o_tnh=attn_o_tnh_spec,
                q_da_sharding=q_da_sharding,
                ap_sharding=ap_sharding,
                kv_da_sharding=kv_da_sharding,
                rd_sharding=rd_sharding,
                prefix=f"{prefix}.layers.{i}.self_attn",
            )
            if self.use_mla_kernel:
                kwargs.update(anh_sharding=anh_sharding)

            return attn_cls(**kwargs)

        def get_decoder_layer(layer_index: int):
            input_layernorm = JaxRmsNorm(
                hidden_size,
                epsilon=rms_norm_eps,
                scale_init=nnx.with_partitioning(init_fn, (None, )),
                dtype=dtype,
                param_dtype=dtype,
                rngs=rng,
                quant_config=quant_config,
            )

            post_attention_layernorm = JaxRmsNorm(
                hidden_size,
                epsilon=rms_norm_eps,
                scale_init=nnx.with_partitioning(init_fn, (None, )),
                dtype=dtype,
                param_dtype=dtype,
                rngs=rng,
                quant_config=quant_config,
            )

            # Logic to determine if this layer is Dense or MoE
            # * The first k layers are always dense.
            # * Subsequent layers are MoE if interleave_moe_layer_step conditions are met
            if layer_index < first_k_dense_replace:
                is_moe_layer = False
            else:
                is_moe_layer = ((layer_index + 1) %
                                interleave_moe_layer_step == 0)

            if not is_moe_layer:
                # Dense Layer (used for first k layers or interleaved dense layers)
                mlp_layer = DeepseekV3MLP(
                    dtype=dtype,
                    hidden_act=hidden_act,
                    hidden_size=hidden_size,
                    intermediate_size=ffw_intermediate_size,
                    rngs=rng,
                    activation_ffw_td=P(ShardingAxisName.MLP_DATA, None),
                    df_sharding=P(None, ShardingAxisName.ATTN_HEAD),
                    fd_sharding=P(ShardingAxisName.ATTN_HEAD, None),
                    quant_config=quant_config)
            else:
                # MoE Layer
                mlp_layer = DeepseekV2Moe(
                    mesh=self.mesh,
                    dtype=dtype,
                    num_expert_parallelism=self.num_expert_parallelism,
                    moe_backend=self.moe_backend,
                    quant_config=quant_config,
                    scoring_func=scoring_func,
                    rng=rng,
                    prefix=f"{prefix}.layers.{layer_index}.mlp")

            return DeepseekV3DecoderLayer(
                input_layernorm=input_layernorm,
                post_attention_layernorm=post_attention_layernorm,
                self_attn=_create_deepseek_attention(layer_index),
                mlp=mlp_layer,
                prefix=f"{prefix}.layers.{layer_index}")

        # hf_config.num_hidden_layers is 61, which ignores the last MTP layer.
        self.start_layer, self.end_layer, self.layers = make_layers(
            hf_config.num_hidden_layers, get_decoder_layer)

        if self.is_last_rank:
            self.norm = JaxRmsNorm(
                hidden_size,
                epsilon=rms_norm_eps,
                dtype=dtype,
                param_dtype=dtype,
                scale_init=nnx.with_partitioning(nnx.initializers.uniform(),
                                                 (None, )),
                rngs=rng,
                quant_config=quant_config,
                prefix=prefix + ".norm",
            )
        else:
            self.norm = PPMissingLayer()

    # For compatibility with flax.
    def apply(self, variables, *args, **kwargs):
        return self.__call__(*args, **kwargs)

    def initialize_cache(self):
        # Initialize RoPE cache once after weights are loaded.
        self.rope_emb.initialize_cache()

    def __call__(
        self,
        kv_caches: List[jax.Array],
        input_ids: Optional[jax.Array],
        attention_metadata: AttentionMetadata,
        inputs_embeds: Optional[jax.Array] = None,
    ) -> Tuple[List[jax.Array], jax.Array]:
        if inputs_embeds is not None:
            x = inputs_embeds
        else:
            x = self.embed_tokens(input_ids)

        for i, layer in enumerate(
                islice(self.layers, self.start_layer, self.end_layer)):
            kv_cache = kv_caches[i]
            kv_cache, x = layer(
                x,
                kv_cache=kv_cache,
                attention_metadata=attention_metadata,
            )
            kv_caches[i] = kv_cache
        x = self.norm(x)
        return kv_caches, x


class DeepseekV3ForCausalLM(JaxModule, LoadableWithIterator):

    def __init__(self, vllm_config: VllmConfig, rng_key: jax.Array,
                 mesh: Mesh) -> None:
        if getattr(vllm_config.model_config, "quantization", None) == "fp8":
            # `get_tpu_quantization_config` returns None for "fp8" because
            # the work in #1623 is not fully merged. So this block overrides
            # the logic to return Fp8Config when model_config indicates fp8.
            # TODO(#1623): Remove this block when `get_tpu_quantization_config`
            # is updated.
            from tpu_inference.layers.jax.quantization.fp8 import Fp8Config
            hg_quant_config = getattr(vllm_config.model_config.hf_config,
                                      "quantization_config", {})
            vllm_config.quant_config = Fp8Config(hg_quant_config)

        self.vllm_config = vllm_config
        rng = nnx.Rngs(rng_key)
        self.mesh = mesh

        self.model = DeepSeekV3(
            vllm_config=vllm_config,
            rng=rng,
            mesh=mesh,
            quant_config=vllm_config.quant_config,
            prefix="model",
        )

        model_config = vllm_config.model_config
        if self.model.is_last_rank:
            vocab_size = model_config.get_vocab_size()
            hidden_size = model_config.hf_config.hidden_size
            self.lm_head = JaxEinsum(
                einsum_str="TD,DV->TV",
                kernel_shape=(hidden_size, vocab_size),
                param_dtype=model_config.dtype,
                dtype=model_config.dtype,
                rngs=rng,
                kernel_init=nnx.with_partitioning(
                    init_fn, (None, ShardingAxisName.MLP_TENSOR)),
                # Same as https://github.com/vllm-project/tpu-inference/issues/1684
                # DS-V3 doesn't quantize lm_head.
                quant_config=None,
                prefix="lm_head",
            )
        else:
            self.lm_head = PPMissingLayer()

    def __call__(
        self,
        kv_caches: List[jax.Array],
        input_ids: jax.Array,
        attention_metadata: AttentionMetadata,
        inputs_embeds: Optional[jax.Array] = None,
        _input_positions=None,
        _layer_name_to_kv_cache=None,
        _lora_metadata=None,
        intermediate_tensors: JaxIntermediateTensors | None = None,
        is_first_rank: bool = True,
        is_last_rank: bool = True,
        *args,
    ) -> Tuple[List[jax.Array], jax.Array | JaxIntermediateTensors,
               List[jax.Array]]:
        if not is_first_rank:
            assert intermediate_tensors is not None
            inputs_embeds = intermediate_tensors["hidden_states"]

        kv_caches, x = self.model(
            kv_caches,
            input_ids,
            attention_metadata,
            inputs_embeds,
        )

        if not is_last_rank:
            x = JaxIntermediateTensors(tensors={"hidden_states": x}, )

        return kv_caches, x, []

    def compute_logits(self, hidden_states: jax.Array) -> jax.Array:
        return self.lm_head(hidden_states)

    def load_weights(self, weights: Iterable) -> set[str]:
        if not isinstance(weights, Iterable):
            # Use next parent class in MRO.
            return super().load_weights(weights)

        start_ignore_layer_num = len(self.model.layers)
        end_ignore_layer_num = 62  # last layer is MTP, we ignore it for now
        loader = JaxAutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head"]
                           if not hasattr(self, 'lm_head') else []),
            skip_substrs=[
                f"layers.{i}"
                for i in range(start_ignore_layer_num, end_ignore_layer_num)
            ],
        )
        loaded = loader.load_weights(weights)

        self.model.initialize_cache()

        # Display model arch
        if os.environ.get("VLLM_LOGGING_LEVEL", "").upper() == "DEBUG":
            logger.debug("Model architecture and parameter dtypes:")
            num_layers_to_display = 5
            should_skip_layer_display = False
            for name, param in self.named_parameters():
                if f"layers.{num_layers_to_display}." in name:
                    should_skip_layer_display = True
                if should_skip_layer_display and "layers." in name:
                    continue
                v: jax.Array = param.value
                logger.debug(f"{name} : {v.dtype}{v.shape} on {v.device}")

        return loaded