# 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 os
import re
from dataclasses import dataclass
from typing import List, Optional, Tuple

import jax
import jax.numpy as jnp
import torch
from flax import nnx
from flax.typing import PRNGKey
from jax.sharding import Mesh, NamedSharding
from jax.sharding import PartitionSpec as P
from torchax.ops.mappings import j2t_dtype
from vllm.config import VllmConfig

from tpu_inference import utils
from tpu_inference.layers.common.quantization import u8_unpack_e2m1
from tpu_inference.layers.common.sharding import ShardingAxisName
from tpu_inference.layers.jax.attention.attention import AttentionMetadata
from tpu_inference.layers.jax.attention.deepseek_v3_attention import MLA
from tpu_inference.layers.jax.constants import KVCacheType
from tpu_inference.layers.jax.layers import DenseFFW, Embedder, LMhead, RMSNorm
from tpu_inference.layers.jax.moe.deepseek_v3_moe import DeepSeekV3Router
from tpu_inference.layers.jax.moe.moe import MoE
from tpu_inference.layers.jax.moe.utils import MoEBackend, select_moe_backend
from tpu_inference.layers.jax.transformer_block import (
    SharedExpertsTransformerBlock, TransformerBlock)
from tpu_inference.logger import init_logger
from tpu_inference.models.jax.utils.weight_utils import (BaseWeightLoader,
                                                         get_param,
                                                         print_param_info)

logger = init_logger(__name__)

# A map from JAX dtype to the corresponding PyTorch integer dtype for raw memory viewing.
DTYPE_VIEW_MAP = {
    jnp.dtype(jnp.float8_e4m3fn): torch.uint8,
    jnp.dtype(jnp.bfloat16): torch.uint16,
    jnp.dtype(jnp.float32): torch.uint32,
}


@dataclass
class DeepSeekV3WeightLoader(BaseWeightLoader):

    def __init__(self,
                 vllm_config: VllmConfig,
                 num_layers,
                 hidden_size,
                 q_lora_rank,
                 kv_lora_rank,
                 attn_heads,
                 qk_nope_head_dim,
                 qk_rope_head_dim,
                 v_head_dim,
                 num_local_experts,
                 model_dtype,
                 use_mla_kernel=False):
        super().__init__(vllm_config, framework="pt")
        self.num_layers = num_layers
        self.is_verbose = vllm_config.additional_config.get(
            "is_verbose", None) is not None
        self.num_routed_experts = num_local_experts
        self.attn_heads = attn_heads
        self.qk_nope_head_dim = qk_nope_head_dim
        self.v_head_dim = v_head_dim
        self.kv_lora_rank = kv_lora_rank
        self.model_dtype = model_dtype
        self.use_mla_kernel = use_mla_kernel

        self._transpose_map = {
            # dense mlp
            r"mlp\.down_proj": (1, 0),
            r"mlp\.gate_proj": (1, 0),
            r"mlp\.up_proj": (1, 0),
            # mla
            r"q_a_proj": (1, 0),
            r"q_b_proj": (1, 0),
            r"kv_a_proj_with_mqa": (1, 0),
            r"kv_b_proj": (1, 0),
            r"k_b_proj": (2, 0, 1),  # used for MLA kernel
            r"v_b_proj": (2, 0, 1),  # used for MLA kernel
            r"o_proj": (1, 0),
            # moe
            r"mlp\.gate\.weight": (1, 0),
            r"mlp\.experts\.\d+\.gate_proj": (0, 2, 1),
            r"mlp\.experts\.\d+\.down_proj": (0, 2, 1),
            r"mlp\.experts\.\d+\.up_proj": (0, 2, 1),
            r"mlp\.shared_experts\.down_proj": (1, 0),
            r"mlp\.shared_experts\.gate_proj": (1, 0),
            r"mlp\.shared_experts\.up_proj": (1, 0),
            # lm_head
            r"lm_head\.weight": (1, 0)
        }

        # Set the mappings from loaded parameter keys to standardized names.
        self._loaded_to_standardized_keys = {
            # encode & decode
            "model.embed_tokens.weight":
            "embedder.input_embedding_table_VD",
            "lm_head.weight":
            "lm_head.input_embedding_table_DV",
            # final norm
            "model.norm.weight":
            "final_norm.scale",
            # norm in transformer blocks
            "model.layers.*.input_layernorm.weight":
            "layers.*.pre_attention_norm.scale",
            "model.layers.*.post_attention_layernorm.weight":
            "layers.*.pre_mlp_norm.scale",
            # attention (MLA)
            "model.layers.*.self_attn.q_a_layernorm.weight":
            "layers.*.attn.q_rms_norm.scale",
            "model.layers.*.self_attn.kv_a_layernorm.weight":
            "layers.*.attn.kv_rms_norm.scale",
            "model.layers.*.self_attn.q_a_proj.weight":
            "layers.*.attn.kernel_q_down_proj_DA",
            "model.layers.*.self_attn.q_b_proj.weight":
            "layers.*.attn.kernel_q_up_proj_AP",
            "model.layers.*.self_attn.kv_a_proj_with_mqa.weight":
            "layers.*.attn.kernel_kv_down_proj_DA",
            "model.layers.*.self_attn.kv_b_proj.weight":
            "layers.*.attn.kernel_kv_up_proj_AL",
            "model.layers.*.self_attn.o_proj.weight":
            "layers.*.attn.kernel_o_proj_RD",
            # Dense ffw
            "model.layers.*.mlp.gate_proj.weight":
            "layers.*.custom_module.kernel_gating_DF",
            "model.layers.*.mlp.up_proj.weight":
            "layers.*.custom_module.kernel_up_proj_DF",
            "model.layers.*.mlp.down_proj.weight":
            "layers.*.custom_module.kernel_down_proj_FD",
            # MOE(routed experts)
            "model.layers.*.mlp.gate.weight":
            "layers.*.custom_module.router.kernel_DE",
            "model.layers.*.mlp.gate.e_score_correction_bias":
            "layers.*.custom_module.router.bias_E",
            "model.layers.*.mlp.experts.*.gate_proj.weight":
            "layers.*.custom_module.kernel_gating_EDF",
            "model.layers.*.mlp.experts.*.down_proj.weight":
            "layers.*.custom_module.kernel_down_proj_EFD",
            "model.layers.*.mlp.experts.*.up_proj.weight":
            "layers.*.custom_module.kernel_up_proj_EDF",
            "model.layers.*.mlp.experts.*.gate_upproj_fused.weight":
            "layers.*.custom_module.kernel_gating_upproj_E2DF",
            # MOE(shared experts)
            "model.layers.*.mlp.shared_experts.down_proj.weight":
            "layers.*.shared_experts.kernel_down_proj_FD",
            "model.layers.*.mlp.shared_experts.gate_proj.weight":
            "layers.*.shared_experts.kernel_gating_DF",
            "model.layers.*.mlp.shared_experts.up_proj.weight":
            "layers.*.shared_experts.kernel_up_proj_DF",
        }
        if self.use_mla_kernel:
            self._loaded_to_standardized_keys.update({
                "model.layers.*.self_attn.k_b_proj.weight":
                "layers.*.attn.kernel_k_up_proj_ANH",
                "model.layers.*.self_attn.v_b_proj.weight":
                "layers.*.attn.kernel_v_up_proj_ANH",
            })
        # TODO (jacobplatin): we should not be hard-coding these
        self.scale_dtype, self.quant_dtype = jnp.bfloat16, jnp.float8_e4m3fn

        self.is_model_quantized = not vllm_config.additional_config.get(
            "skip_quantization", False)

        # TODO (jacobplatin): remove this once the JAX path refactor is done
        self.is_native_fp8_model = False

        if self.is_model_quantized:
            # NOTE: this is only needed for pre-quantized models when doing random weight loading
            # because the scales that Qwix configures by default don't necessarily match the
            # scales in practice
            # TODO (jacobplatin): remove or clean this up
            self.scale_shape_map_for_random_weight_loading = {
                # MoE experts (3D)
                "custom_module.kernel_down_proj_EFD": (256, 8, 7168),
                "custom_module.kernel_gating_EDF": (256, 28, 2048),
                "custom_module.kernel_up_proj_EDF": (256, 28, 2048),
                # Shared experts (2D)
                "shared_experts.kernel_down_proj_FD": (8, 7168),
                "shared_experts.kernel_gating_DF": (28, 2048),
                "shared_experts.kernel_up_proj_DF": (28, 2048),
                # Dense FFW (2D)
                "custom_module.kernel_gating_DF": (28, 18432),
                "custom_module.kernel_up_proj_DF": (28, 18432),
                "custom_module.kernel_down_proj_FD": (72, 7168),
                # Attention (3D for MLA, 2D for the rest)
                "attn.kernel_q_down_proj_DA": (28, 1536),
                "attn.kernel_q_up_proj_AP": (6, 24576),
                "attn.kernel_kv_down_proj_DA": (28, 576),
                "attn.kernel_kv_up_proj_AL": (2, 32768),
                "attn.kernel_o_proj_RD": (64, 7168),
                "attn.kernel_k_up_proj_ANH": (2, 128, 128),  # MLA
                "attn.kernel_v_up_proj_ANH": (2, 128, 128),  # MLA
            }

            # TODO (jacobplatin): remove this check eventually!
            assert self.quant_dtype == jnp.float8_e4m3fn, f"Expected quant_dtype to be float8_e4m3fn for DeepSeek but got {self.quant_dtype}"

            quantization_block_sizes = vllm_config.model_config.hf_config.quantization_config[
                "weight_block_size"]
            # TODO (jacobplatin): remove this once the JAX path refactor is done
            if quantization_block_sizes == [128, 128]:
                assert len(
                    quantization_block_sizes
                ) == 2, f"Expected only 2 quantization block sizes but got {quantization_block_sizes}"
                self.quantization_block_size_n = quantization_block_sizes[0]
                self.quantization_block_size_k = quantization_block_sizes[1]
                self.is_native_fp8_model = True

    def map_loaded_to_standardized_name(self, loaded_key: str) -> str:
        # Find the corresponding model key using the HF key
        if "layer" in loaded_key:
            # extract layer number and replace it with *
            layer_num = re.search(r"layers\.(\d+)", loaded_key).group(1)
            layer_key = re.sub(r"layers\.\d+", "layers.*", loaded_key)
            # extract expert number if exists and replace it with *
            if "experts" in loaded_key and "shared_experts" not in loaded_key:
                layer_key = re.sub(r"experts\.\d+", "experts.*", layer_key)
            # get standardized key and replace * with layer number.
            mapped_key = self._loaded_to_standardized_keys.get(
                layer_key, loaded_key)
            mapped_key = re.sub(r"layers\.\*", f"layers.{layer_num}",
                                mapped_key)
        else:
            mapped_key = self._loaded_to_standardized_keys.get(
                loaded_key, loaded_key)
        return mapped_key

    def _transpose_params(self, param_key: str, param_tensor: jax.Array):
        for key, value in self._transpose_map.items():
            if re.search(key, param_key):
                return jnp.transpose(param_tensor, value)
        return param_tensor  # Base case / no-op

    def _process_moe_weights(self, loaded_name, loaded_weight, weights_dict):
        layer_num = re.search(r"layers\.(\d+)", loaded_name).group(1)
        expert_num_str = re.search(r"experts\.(\d+)", loaded_name).group(1)
        expert_idx = int(expert_num_str)

        if layer_num not in weights_dict:
            weights_dict[layer_num] = ([None] * self.num_routed_experts, 0)

        expert_list, count = weights_dict[layer_num]

        expert_list[expert_idx] = loaded_weight
        count += 1
        weights_dict[layer_num] = (expert_list, count)

        if count == self.num_routed_experts:
            stacked_weights = torch.stack(expert_list, axis=0)
            del weights_dict[layer_num]
            return stacked_weights
        return None

    def _load_individual_weight(self,
                                name,
                                weight,
                                model_params,
                                model_mesh,
                                scale=None) -> Tuple[int, int]:
        """
        Loads a single weight into the model.

        NOTE: if using the base quantized model, it is assumed that the Qwix abstract
        pass has been run and that the model weights are thus QArrays, which we
        will then load the weights/scales into.

        Args:
            name: The name of the weight.
            weight: The weight to load.
            model_params: The model parameters.
            model_mesh: The model mesh.
            scale: The scale of the weight (if using the pre-quantized model).

        Returns:
            Tuple[int, int]: The size (in bytes) for the given layer overall and per shard.
                NOTE: if using the pre-quantized model (with Qwix), we'll include the scale size as well.
        """
        mapped_name = self.map_loaded_to_standardized_name(name)
        base_model_weight = get_param(model_params, mapped_name)
        model_weight = base_model_weight.array.qvalue if hasattr(
            base_model_weight, "array") else base_model_weight
        sharding = base_model_weight.array.qvalue.sharding if hasattr(
            base_model_weight, "array") else base_model_weight.sharding

        # Convert weights from torch into numpy
        if weight.dtype == torch.uint8 and scale is not None:
            # Assume packed FP4 format when uint8 weights with scale provided
            weight_jax_u8 = jnp.array(weight.cpu().numpy())
            weight_np = u8_unpack_e2m1(weight_jax_u8)
            scale = scale.to(torch.float32).numpy().astype(self.scale_dtype)
        else:
            cast_type = model_weight.value.dtype
            # Special-case: FP4 values stored as FP8 for compatibility.
            # If the model expects float4_e2m1fn but the checkpoint provides FP8,
            # convert by numeric value (float32) then cast to float4.
            if cast_type == jnp.float4_e2m1fn and weight.dtype == torch.float8_e4m3fn:
                weight_np = jnp.array(weight.float().numpy()).astype(cast_type)
            else:
                torch_view_type = DTYPE_VIEW_MAP.get(jnp.dtype(cast_type))

                if torch_view_type:
                    # Avoid unnecessary upcasting and mem copy by viewing the tensor's
                    # raw data as integers before converting to a JAX array.
                    weight_np = jnp.array(
                        weight.view(torch_view_type).numpy()).view(cast_type)
                else:
                    raise ValueError(
                        f"Unsupported dtype for tensor conversion: {cast_type}"
                    )

            if scale is not None:
                scale = scale.to(torch.float32).numpy().astype(
                    self.scale_dtype)
        weight_np = self._transpose_params(name, weight_np)
        if scale is not None:
            scale = self._transpose_params(name, scale)
            # Ensure scale is broadcastable to weight_np by repeating per-axis.
            weight_shape = weight_np.shape
            scale_shape = scale.shape
            if len(weight_shape) == len(scale_shape):
                new_scale = scale
                # TODO (jacobplatin): remove once the refactor is complete
                if self.is_native_fp8_model:
                    for idx, (weight_dim, scale_dim) in enumerate(
                            zip(weight_shape, scale_shape)):
                        if weight_dim // self.quantization_block_size_n != scale_dim and weight_dim // scale_dim != 1:
                            old_scale_shape = scale.shape
                            scale = scale.repeat(
                                self.quantization_block_size_n,
                                axis=idx)[:, :weight_dim]
                            logger.warning(
                                f"Got a weight with shape {weight_shape} and scale with shape {old_scale_shape} "
                                f"where the scale_dim {scale_dim} does not match the weight_dim {weight_dim} "
                                f"multiplied by the quantization block size {self.quantization_block_size_n}. "
                                f"Repeating the scale to new shape {scale.shape} along axis {idx} with repeat size {self.quantization_block_size_n}."
                            )
                else:
                    for wdim, sdim in zip(weight_shape, scale_shape):
                        if (wdim % sdim != 0):
                            raise ValueError(
                                f"Weight dim {wdim} is not divisible by scale dim {sdim} for weight {name} with shape {weight_shape} and scale {scale_shape}!"
                            )
                    if scale_shape != new_scale.shape:
                        logger.warning(
                            f"Adjusted scale shape {scale_shape} to {new_scale.shape} to match weight {weight_shape}"
                        )
                    scale = new_scale
            else:
                raise ValueError(
                    f"Scale rank {scale_shape} does not match weight rank {weight_shape}"
                )

        if model_weight.value.shape != weight_np.shape:
            raise ValueError(
                f"Loaded shape for {name}: {weight_np.shape} "
                f"does not match model shape for {mapped_name}: {model_weight.value.shape}!"
            )

        def get_slice(index):
            return weight_np[index]

        def get_slice_scale(index):
            # ruff: noqa: F821
            return scale[index]

        sharded_array = jax.make_array_from_callback(
            weight_np.shape, NamedSharding(model_mesh, P(*sharding)),
            get_slice)

        if scale is not None:
            maybe_sharded_scale = scale
            # Since, by default, we'll use the same sharding as the weights, we might
            # encounter an error where the smaller/different sharding dim isn't divisible
            # by the parallel size
            # NOTE: Qwix expert dangyi@ mentioned that we don't need to worry about accuracy
            # impacts when sharing scales
            try:
                maybe_sharded_scale = jax.make_array_from_callback(
                    scale.shape, NamedSharding(model_mesh, P(*sharding)),
                    get_slice_scale)
            except ValueError:
                logger.warning(
                    f"Could not create sharded scale for {name} with shape {scale.shape} and sharding {sharding}, skipping sharding..."
                )
            assert base_model_weight.array.scale.value.dtype == maybe_sharded_scale.dtype, f"Expected dtype for model weight scale with name {mapped_name} and dtype ({base_model_weight.array.scale.value.dtype}) to match that of the incoming weight scale ({maybe_sharded_scale.dtype})"
            assert base_model_weight.array.qvalue.value.dtype == sharded_array.dtype, f"Expected dtype for model weight with name {mapped_name} and dtype ({base_model_weight.array.qvalue.value.dtype}) to match that of the incoming weight ({sharded_array.dtype})"
            base_model_weight.array.scale.value = maybe_sharded_scale
            base_model_weight.array.qvalue.value = sharded_array
        else:
            assert model_weight.value.dtype == sharded_array.dtype, f"Expected dtype for model weight with name {mapped_name} and dtype ({model_weight.value.dtype}) to match that of the incoming weight ({sharded_array.dtype})"
            model_weight.value = sharded_array

        model_weight_size_bytes = model_weight.nbytes / 1e9
        model_weight_local_size_bytes = model_weight.addressable_shards[
            0].data.nbytes / 1e9

        if scale is not None:
            model_weight_size_bytes += base_model_weight.array.scale.nbytes / 1e9
            model_weight_local_size_bytes += base_model_weight.array.scale.addressable_shards[
                0].data.nbytes / 1e9

        if self.is_verbose:
            logger.info(f"Memory usage after loading in {name}: "
                        f"hbm={utils.hbm_usage_gb(jax.local_devices())}Gb")
            print_param_info(model_weight, name)
            if scale is not None:
                print_param_info(base_model_weight.array.scale,
                                 "scale for " + name)

        del weight, scale
        return model_weight_size_bytes, model_weight_local_size_bytes

    def load_weights(self, model_for_loading: nnx.Module):
        model_params = nnx.state(model_for_loading)
        logger.warning(
            f"loaded_to_standardized_keys: {self._loaded_to_standardized_keys}"
        )
        cumulative_global_memory = 0
        cumulative_local_memory = 0
        mlp_experts_gate_proj_weights = {}
        mlp_experts_gate_proj_scales = {}
        mlp_experts_up_proj_weights = {}
        mlp_experts_up_proj_scales = {}
        mlp_experts_down_proj_weights = {}
        mlp_experts_down_proj_scales = {}
        stacked_tensors = {}
        quantized_weights = {}
        quantized_scales = {}
        with jax.default_device(jax.devices("cpu")[0]):
            for loaded_name, loaded_weight in self.get_weights_iterator():
                # Skip if the model has fewer layers than original.
                if re.search(r"layers\.(\d+)", loaded_name):
                    layer_num = re.search(r"layers\.(\d+)",
                                          loaded_name).group(1)
                    if int(layer_num) >= self.num_layers:
                        del loaded_weight
                        continue
                if 'layers.61' in loaded_name:
                    # skip loading MTP module.
                    del loaded_weight
                    continue
                if re.search(r"experts\.(\d+)", loaded_name):
                    expert_num = re.search(r"experts\.(\d+)",
                                           loaded_name).group(1)
                    if int(expert_num) >= self.num_routed_experts:
                        del loaded_weight
                        continue
                # NOTE: loaded_weight will be a Torch tensor, so we need to convert it to the
                # equivalent jnp dtype
                # TODO (jacobplatin): refactor this so that we instead change / update `model_weights_generator`
                # instead of checking "weight_scale_inv" and assuming quantization method is fp8
                scale = None
                # Mixed quantization: accept both fp8 and packed fp4 (uint8) tensors
                allowed_quant_dtypes = {
                    j2t_dtype(self.quant_dtype.dtype), torch.uint8
                }
                if loaded_weight.dtype in allowed_quant_dtypes:
                    if self.is_model_quantized:
                        scale_name = loaded_name.replace(
                            ".weight", ".weight_scale_inv")
                        if scale_name in quantized_scales:
                            scale = quantized_scales[scale_name]
                            del quantized_scales[scale_name]
                        else:
                            quantized_weights[loaded_name] = loaded_weight
                            continue
                    else:
                        quantized_weights[loaded_name] = loaded_weight
                        continue

                if loaded_name.endswith(".weight_scale_inv"):
                    if self.is_model_quantized:
                        weight_name = loaded_name.replace(
                            ".weight_scale_inv", ".weight")
                        if weight_name in quantized_weights:
                            scale = loaded_weight
                            loaded_weight = quantized_weights[weight_name]
                            loaded_name = weight_name
                            del quantized_weights[weight_name]
                        else:
                            quantized_scales[loaded_name] = loaded_weight
                            continue
                    # In the case that we don't want to use the quantized weights,
                    # we'll dequantize the weights using the loaded scale on-the-fly
                    else:
                        # assuming weights are loaded before scales.
                        weight_name = loaded_name.replace(
                            ".weight_scale_inv", ".weight")
                        loaded_weight = weights_dequant_cpu(
                            quantized_weights[weight_name], loaded_weight,
                            self.model_dtype)
                        loaded_name = weight_name
                        del quantized_weights[weight_name]
                # concat mlp.experts weights
                stacked_scales = None
                stacked_weights = None
                if "mlp.experts" in loaded_name:
                    if "down_proj" in loaded_name:
                        proj_type = "down_proj"
                        stacked_weights = self._process_moe_weights(
                            loaded_name, loaded_weight,
                            mlp_experts_down_proj_weights)
                        if scale is not None:
                            stacked_scales = self._process_moe_weights(
                                loaded_name, scale,
                                mlp_experts_down_proj_scales)
                    if "gate_proj" in loaded_name:
                        proj_type = "gate_proj"
                        stacked_weights = self._process_moe_weights(
                            loaded_name, loaded_weight,
                            mlp_experts_gate_proj_weights)
                        if scale is not None:
                            stacked_scales = self._process_moe_weights(
                                loaded_name, scale,
                                mlp_experts_gate_proj_scales)
                    if "up_proj" in loaded_name:
                        proj_type = "up_proj"
                        stacked_weights = self._process_moe_weights(
                            loaded_name, loaded_weight,
                            mlp_experts_up_proj_weights)
                        if scale is not None:
                            stacked_scales = self._process_moe_weights(
                                loaded_name, scale, mlp_experts_up_proj_scales)
                    if stacked_weights is not None:
                        stacked_tensors[layer_num +
                                        proj_type] = (stacked_weights,
                                                      stacked_scales)
                        if all(f"{layer_num}{p}" in stacked_tensors
                               for p in ["gate_proj", "up_proj", "down_proj"]):

                            gate_w, gate_s = stacked_tensors.pop(layer_num +
                                                                 "gate_proj")
                            up_w, up_s = stacked_tensors.pop(layer_num +
                                                             "up_proj")
                            down_w, down_s = stacked_tensors.pop(layer_num +
                                                                 "down_proj")

                            is_moe_kernel = model_for_loading.moe_backend in [
                                MoEBackend.FUSED_MOE, MoEBackend.VLLM_MOE
                            ]
                            gate_name = loaded_name.replace(
                                proj_type, "gate_proj")
                            up_name = loaded_name.replace(proj_type, "up_proj")
                            down_name = loaded_name.replace(
                                proj_type, "down_proj")
                            if is_moe_kernel:
                                # (E, D, F) -> (E, 2, D, F)
                                fused_w = torch.stack([gate_w, up_w], dim=1)
                                fused_s = torch.stack(
                                    [gate_s, up_s], dim=1
                                ) if gate_s is not None and up_s is not None else None

                                fused_name = loaded_name.replace(
                                    proj_type, "gate_upproj_fused")
                                weight_bytes, weight_shards = self._load_individual_weight(
                                    fused_name,
                                    fused_w,
                                    model_params,
                                    model_for_loading.mesh,
                                    scale=fused_s)

                                weight_bytes_down, weight_shards_down = self._load_individual_weight(
                                    loaded_name,
                                    down_w,
                                    model_params,
                                    model_for_loading.mesh,
                                    scale=down_s)

                                # Update cumulative memory
                                weight_bytes += weight_bytes_down
                                weight_shards += weight_shards_down

                            else:

                                weight_bytes, weight_shards = self._load_individual_weight(
                                    gate_name,
                                    gate_w,
                                    model_params,
                                    model_for_loading.mesh,
                                    scale=gate_s)

                                weight_bytes_up, weight_shards_up = self._load_individual_weight(
                                    up_name,
                                    up_w,
                                    model_params,
                                    model_for_loading.mesh,
                                    scale=up_s)
                                weight_bytes += weight_bytes_up
                                weight_shards += weight_shards_up

                                weight_bytes_down, weight_shards_down = self._load_individual_weight(
                                    down_name,
                                    down_w,
                                    model_params,
                                    model_for_loading.mesh,
                                    scale=down_s)
                                weight_bytes += weight_bytes_down
                                weight_shards += weight_shards_down
                        else:
                            continue
                        if self.is_verbose:
                            cumulative_global_memory += weight_bytes
                            cumulative_local_memory += weight_shards
                            logger.info(
                                f"Cumulative global memory: {cumulative_global_memory} GB"
                            )
                            logger.info(
                                f"Cumulative local memory: {cumulative_local_memory} GB"
                            )
                    else:
                        continue
                else:
                    if self.use_mla_kernel and "kv_b_proj" in loaded_name:
                        # loaded_weight shape: (num_heads * (d_k + d_v), kv_lora_rank)
                        # scale shape: (num_heads * (d_k + d_v) / block_n, kv_lora_rank / block_k)
                        # Reshape to (num_heads, (d_k + d_v), kv_lora_rank) and split
                        weight_reshaped = loaded_weight.view(
                            self.attn_heads,
                            self.qk_nope_head_dim + self.v_head_dim,
                            self.kv_lora_rank)

                        k_weight = weight_reshaped[:, :self.
                                                   qk_nope_head_dim, :]
                        v_weight = weight_reshaped[:,
                                                   self.qk_nope_head_dim:, :]

                        loaded_weights_list = [k_weight, v_weight]
                        loaded_names = [
                            loaded_name.replace("kv_b_proj", "k_b_proj"),
                            loaded_name.replace("kv_b_proj", "v_b_proj")
                        ]

                        scales_list = [None, None]
                        if scale is not None:
                            # TODO (jacobplatin): remove once refactor happens
                            if self.is_native_fp8_model:
                                bn = self.quantization_block_size_n
                                bk = self.quantization_block_size_k
                                scale_reshaped = scale.view(
                                    self.attn_heads,
                                    (self.qk_nope_head_dim + self.v_head_dim)
                                    // bn, self.kv_lora_rank // bk)

                                k_scale = scale_reshaped[:, :self.
                                                         qk_nope_head_dim //
                                                         bn, :]
                                v_scale = scale_reshaped[:, self.
                                                         qk_nope_head_dim //
                                                         bn:, :]
                            else:
                                assert loaded_weight.shape[0] == scale.shape[0]
                                block_size_k = loaded_weight.shape[
                                    1] // scale.shape[1]
                                assert block_size_k > 0, f"Expected non-zero block size but got {block_size_k}!"
                                scale_reshaped = scale.view(
                                    self.attn_heads,
                                    (self.qk_nope_head_dim + self.v_head_dim),
                                    self.kv_lora_rank // block_size_k)

                                k_scale = scale_reshaped[:, :self.
                                                         qk_nope_head_dim, :]
                                v_scale = scale_reshaped[:, self.
                                                         qk_nope_head_dim:, :]
                            scales_list = [k_scale, v_scale]

                    else:
                        loaded_weights_list = [loaded_weight]
                        loaded_names = [loaded_name]
                        scales_list = [scale]

                    for loaded_name, loaded_weight, scale in zip(
                            loaded_names, loaded_weights_list, scales_list):

                        weight_bytes, weight_shards = self._load_individual_weight(
                            loaded_name,
                            loaded_weight,
                            model_params,
                            model_for_loading.mesh,
                            scale=scale)
                        if self.is_verbose:
                            cumulative_global_memory += weight_bytes
                            cumulative_local_memory += weight_shards
                            logger.info(
                                f"Cumulative global memory: {cumulative_global_memory} GB"
                            )
                            logger.info(
                                f"Cumulative local memory: {cumulative_local_memory} GB"
                            )

        del mlp_experts_gate_proj_weights
        del mlp_experts_up_proj_weights
        del mlp_experts_down_proj_weights
        del quantized_weights
        del quantized_scales
        # TODO: validate that all of the model_params were accounted for as well.
        nnx.update(model_for_loading, model_params)


@dataclass
class DeepSeekV3(nnx.Module):
    WeightLoader = DeepSeekV3WeightLoader

    def __init__(self,
                 vllm_config: VllmConfig,
                 rng: jax.Array,
                 mesh: Mesh,
                 force_random_weights: bool = False):
        assert mesh is not None

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

        # NOTE: the default is 61
        num_layers: int = vllm_config.model_config.hf_config.num_hidden_layers
        num_local_experts: int = 256

        vocab_size: int = 129280
        hidden_size: int = 7168
        # NOTE: this dtype may be implicitly overriden if using to Qwix to load in the quantized weights
        dtype: jnp.dtype = jnp.bfloat16
        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
        first_k_dense_replace: int = 3  # replace the first few MOE layers to dense layer.
        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}")

        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

        self.mesh = mesh

        self.moe_backend = select_moe_backend()

        self.weight_loader = self.WeightLoader(
            vllm_config=vllm_config,
            num_layers=num_layers,
            hidden_size=hidden_size,
            q_lora_rank=q_lora_rank,
            kv_lora_rank=kv_lora_rank,
            attn_heads=num_attention_heads,
            qk_nope_head_dim=qk_nope_head_dim,
            qk_rope_head_dim=qk_rope_head_dim,
            v_head_dim=v_head_dim,
            num_local_experts=num_local_experts,
            model_dtype=dtype,
            use_mla_kernel=self.use_mla_kernel)

        self.embedder = Embedder(vocab_size=vocab_size,
                                 hidden_size=hidden_size,
                                 dtype=dtype,
                                 rngs=self.rng,
                                 vd_sharding=(ShardingAxisName.MLP_TENSOR,
                                              None))

        self.layers = []

        def _create_mla() -> MLA:
            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)

            else:
                query_tnh_spec = P(None, ShardingAxisName.MLP_TENSOR, None)
                keyvalue_skh_spec = P(None, ShardingAxisName.MLP_TENSOR, None)
                attn_o_tnh_spec = P(None, ShardingAxisName.MLP_TENSOR, None)

            return MLA(
                rope_theta=rope_theta,
                rope_scaling=rope_scaling,
                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,
                use_mla_kernel=self.use_mla_kernel,
                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
                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=self.rng,
                activation_attention_td=(None, None),
                activation_q_td=(None, None),
                query_tnh=query_tnh_spec,
                keyvalue_skh=keyvalue_skh_spec,
                activation_attention_out_td=(None, None),
                attn_o_tnh=attn_o_tnh_spec,
                q_da_sharding=(None, ShardingAxisName.VOCAB),
                ap_sharding=(None, ShardingAxisName.MLP_TENSOR),
                anh_sharding=(None, ShardingAxisName.MLP_TENSOR, None),
                kv_da_sharding=(None, ShardingAxisName.VOCAB),
                rd_sharding=(ShardingAxisName.MLP_TENSOR, None))

        for i in range(first_k_dense_replace):
            block = TransformerBlock(
                pre_attention_norm=RMSNorm(
                    dims=hidden_size,
                    epsilon=rms_norm_eps,
                    with_scale=True,
                    dtype=dtype,
                    rngs=self.rng,
                ),
                pre_mlp_norm=RMSNorm(
                    dims=hidden_size,
                    epsilon=rms_norm_eps,
                    with_scale=True,
                    dtype=dtype,
                    rngs=self.rng,
                ),
                attn=_create_mla(),
                custom_module=DenseFFW(
                    dtype=dtype,
                    hidden_act=hidden_act,
                    hidden_size=hidden_size,
                    intermediate_size=ffw_intermediate_size,
                    rngs=self.rng,
                    activation_ffw_td=(ShardingAxisName.MLP_DATA, None),
                    df_sharding=(None, ShardingAxisName.MLP_TENSOR),
                    fd_sharding=(ShardingAxisName.MLP_TENSOR, None)))

            self.layers.append(block)

        for i in range(first_k_dense_replace, num_layers):
            is_moe_layer = ((i + 1) % interleave_moe_layer_step == 0)
            router = 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=self.rng,
                routed_scaling_factor=2.5,
                dtype=dtype,
                moe_backend=self.moe_backend,
                activation_ffw_td=(ShardingAxisName.MLP_DATA, None),
                ed_sharding=(ShardingAxisName.MLP_TENSOR, None),
                e_sharding=(ShardingAxisName.MLP_TENSOR, ))
            custom_module = MoE(
                dtype=dtype,
                num_local_experts=num_local_experts,
                apply_expert_weight_before_computation=False,
                hidden_size=hidden_size,
                intermediate_size_moe=moe_intermediate_size,
                num_experts_per_tok=num_experts_per_token,
                mesh=self.mesh,
                hidden_act=hidden_act,
                rngs=self.rng,
                activation_ffw_td=(ShardingAxisName.MLP_DATA, None),
                activation_ffw_ted=(ShardingAxisName.MLP_DATA, None, None),
                edf_sharding=(ShardingAxisName.MLP_TENSOR, None, None),
                efd_sharding=(ShardingAxisName.MLP_TENSOR, None, None),
                moe_backend=self.moe_backend,
                quantized_dtype=self.weight_loader.quant_dtype
                if self.weight_loader.is_model_quantized else None,
                router=router) if is_moe_layer else DenseFFW(
                    dtype=dtype,
                    hidden_act=hidden_act,
                    hidden_size=hidden_size,
                    intermediate_size=ffw_intermediate_size,
                    rngs=self.rng,
                    activation_ffw_td=(ShardingAxisName.MLP_DATA, None),
                    df_sharding=(None, ShardingAxisName.MLP_TENSOR),
                    fd_sharding=(ShardingAxisName.MLP_TENSOR, None))

            shared_experts = DenseFFW(
                dtype=dtype,
                hidden_act=hidden_act,
                hidden_size=hidden_size,
                intermediate_size=num_shared_experts * moe_intermediate_size,
                rngs=self.rng,
                activation_ffw_td=(ShardingAxisName.MLP_DATA, None),
                df_sharding=(None, ShardingAxisName.MLP_TENSOR),
                fd_sharding=(ShardingAxisName.MLP_TENSOR, None))

            pre_attention_norm = RMSNorm(
                dims=hidden_size,
                rngs=self.rng,
                epsilon=rms_norm_eps,
                with_scale=True,
                dtype=dtype,
            )

            pre_mlp_norm = RMSNorm(
                dims=hidden_size,
                rngs=self.rng,
                epsilon=rms_norm_eps,
                with_scale=True,
                dtype=dtype,
            )

            block = SharedExpertsTransformerBlock(
                custom_module=custom_module,
                attn=_create_mla(),
                pre_attention_norm=pre_attention_norm,
                pre_mlp_norm=pre_mlp_norm,
                shared_experts=shared_experts)
            self.layers.append(block)

        self.final_norm = RMSNorm(
            dims=hidden_size,
            rngs=self.rng,
            epsilon=rms_norm_eps,
            with_scale=True,
            dtype=dtype,
        )

        self.lm_head = LMhead(vocab_size=vocab_size,
                              hidden_size=hidden_size,
                              dtype=dtype,
                              rngs=self.rng,
                              vd_sharding=(ShardingAxisName.MLP_TENSOR, None),
                              dv_sharding=(None, ShardingAxisName.MLP_TENSOR))

        if os.environ.get("VLLM_LOGGING_LEVEL", "").upper() == "DEBUG":
            self._print_model_architecture()

    def _print_model_architecture(self):
        num_display_layers = 5

        logger.debug("### Embedding ###")
        nnx.display(self.embedder)

        logger.debug(f"\n### First {num_display_layers} Layers ###")
        # Loop through the slice and display each layer
        for i, layer in enumerate(self.layers[:num_display_layers]):
            logger.debug(f"\n--- Layer {i} ---")
            nnx.display(layer)

        logger.debug("\n### LM Head ###")
        nnx.display(self.lm_head)

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

    def load_weights(self, rng: PRNGKey, cache_dir: Optional[str] = None):
        # NOTE: Since we are using nnx.eval_shape to init the model,
        # we have to pass dynamic arrays here for __call__'s usage.
        self.rng = nnx.Rngs(rng)
        self.weight_loader.load_weights(self)
        self.initialize_cache()

    def initialize_cache(self):
        # Initialize RoPE caches after weights are loaded and before JIT compilation.
        for layer in self.layers:
            if hasattr(layer, 'attn') and hasattr(layer.attn, 'rope'):
                if hasattr(layer.attn.rope, 'initialize_cache'):
                    layer.attn.rope.initialize_cache(self.mesh)

    def __call__(
        self,
        kv_caches: List[jax.Array],
        input_ids: jax.Array,
        attention_metadata: AttentionMetadata,
        *args,
    ) -> Tuple[List[KVCacheType], jax.Array, List[jax.Array]]:
        is_prefill = False
        x = self.embedder.encode(input_ids)
        for (i, block) in enumerate(self.layers):
            kv_cache = kv_caches[i]
            new_kv_cache, x = block(x, is_prefill, kv_cache,
                                    attention_metadata)
            kv_caches[i] = new_kv_cache

        final_activation = self.final_norm(x)

        return kv_caches, final_activation, []

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


def weights_dequant_cpu(x: torch.Tensor,
                        s: torch.Tensor,
                        output_dtype: jnp.dtype,
                        block_size: int = 128) -> torch.Tensor:
    assert x.dim() == 2 and s.dim() == 2, "Both x and s must be 2D tensors"
    M, N = x.shape

    x = x.to(torch.float32)
    s = s.to(torch.float32)
    y = torch.empty_like(x)

    M_main = (M // block_size) * block_size
    N_main = (N // block_size) * block_size

    if M_main > 0 and N_main > 0:
        x_main = x[:M_main, :N_main]
        s_main = s[:(M // block_size), :(N // block_size)]

        x_reshaped = x_main.view(M // block_size, block_size, N // block_size,
                                 block_size).permute(0, 2, 1, 3)
        s_reshaped = s_main.view(M // block_size, N // block_size, 1, 1)
        y_main = (x_reshaped * s_reshaped).permute(0, 2, 1,
                                                   3).reshape(M_main, N_main)

        y[:M_main, :N_main] = y_main

    if N_main < N:
        for i in range(0, M_main, block_size):
            block = x[i:i + block_size, N_main:N]
            scale = s[i // block_size, N // block_size]
            y[i:i + block_size, N_main:N] = block * scale

    if M_main < M:
        for j in range(0, N, block_size):
            block = x[M_main:M, j:j + block_size]
            scale = s[M // block_size, j // block_size]
            y[M_main:M, j:j + block_size] = block * scale

    return y.to(j2t_dtype(jnp.dtype(output_dtype)))
