tpu-inference 0.11.1rc1__py3-none-any.whl → 0.11.1rc3__py3-none-any.whl

tpu_inference/layers/jax/moe/deepseek_v3_moe.py

@@ -0,0 +1,608 @@
+import enum
+from dataclasses import InitVar, dataclass
+from functools import partial
+from typing import Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from flax.typing import Sharding
+from jax.sharding import PartitionSpec
+from jaxtyping import Float
+from qwix._src.core.ragged_dot import ragged_dot as qwix_ragged_dot
+from qwix._src.providers import ptq
+
+from tpu_inference.layers.jax.base import create_param
+from tpu_inference.layers.jax.layers import FlaxUtils
+from tpu_inference.layers.jax.moe.moe import MoE
+from tpu_inference.models.jax.utils.quantization.quantization_utils import (
+    manually_quantize_qwix_activation, manually_quantize_qwix_weight)
+
+modeling_flax_utils = FlaxUtils()
+
+
+@dataclass
+class DeepSeekV3Router(nnx.Module):
+    """Router module for Mixture-of-Experts (MoE) layers.
+
+    This module determines which experts each token should be routed to based on the input.
+
+    """
+
+    hidden_size: int
+    num_experts: int
+    num_experts_per_tok: int
+    n_groups: int
+    topk_groups: int
+    norm_topk_prob: bool
+    routed_scaling_factor: float
+    dtype: jnp.dtype
+    rngs: InitVar[nnx.Rngs]
+
+    # Sharding Attributes
+    activation_ffw_td: Sharding = ()
+    ed_sharding: Sharding = ()
+    e_sharding: Sharding = ()
+
+    random_init: bool = False
+
+    router_bias_dtype: jnp.dtype = jnp.float32
+
+    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.bias_E
+        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)
+            indices = jax.lax.top_k(group_scores_TG, k=self.topk_groups)[1]
+
+            mask_TG = jnp.any(jnp.arange(
+                self.n_groups)[:, None] == indices[..., None, :],
+                              axis=-1)
+            mask_TE = jnp.repeat(mask_TG,
+                                 scores_TE.shape[-1] // mask_TG.shape[-1], -1)
+            scores_TE = jnp.where(mask_TE, scores_TE, 0.0)
+
+        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 = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+
+        scores_TE = jnp.einsum("TD,DE -> TE", x_TD, self.kernel_DE.value)
+        scores_TE = nnx.sigmoid(scores_TE)
+
+        original_scores_TE = scores_TE
+        topk_indices_TX = self.get_topk_indices(scores_TE)
+        weights_TX = jnp.take_along_axis(original_scores_TE,
+                                         topk_indices_TX,
+                                         axis=-1)
+
+        if self.norm_topk_prob:
+            weights_TX /= jnp.sum(weights_TX, axis=-1)[..., None] + 1e-20
+
+        weights_TX *= self.routed_scaling_factor
+
+        return weights_TX, topk_indices_TX
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        """Generates the router kernel (weights and bias) for routing."""
+        D = self.hidden_size
+        E = self.num_experts
+        self.kernel_DE = create_param(rngs,
+                                      shape=(D, E),
+                                      dtype=self.dtype,
+                                      sharding=self.ed_sharding,
+                                      random_init=self.random_init)
+        self.bias_E = create_param(rngs,
+                                   shape=(E, ),
+                                   dtype=self.router_bias_dtype,
+                                   sharding=self.e_sharding,
+                                   random_init=self.random_init)
+
+
+@dataclass(kw_only=True)
+class SparseMoE(MoE):
+    """Mixture-of-Experts (MoE) Routed MLP Layer.
+
+    This module implements a Sparse MoE layer with a router and multiple expert MLPs.
+
+    Attributes:
+        num_experts_per_tok: The number of experts each token is routed to.
+        tile_size: A tuple (batch, activation_dim, weight_dim) for GMM tiling.
+        use_megablox: If True, uses the MegaBlox GMM kernel.
+        mesh: The device mesh.
+        # TODO: need to redesign this I/O for parallelism
+        num_expert_parallelism: The size of the 'expert' mesh dimension.
+        # TODO: determine if we get it from external or extrat it in MoE class
+        is_batch_sharded_by_expert: True if batch is sharded over 'expert' dim.
+    """
+    num_experts_per_tok: int
+    #TODO: tile size is (tile_batch_seq, tile_activation_dim, tile_weight_dim,) from MaxText
+    tile_size: tuple[int, int, int] = (128, 64, 128)
+    use_megablox: bool = False
+    mesh: jax.sharding.Mesh
+    # This should be set if and only if you have quantized your model (via Qwix)
+    quantized_dtype: Optional[jnp.dtype] = None
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        super().__post_init__(rngs)
+
+        # Derive the expert sharding
+        self.expert_axis_name = self.edf_sharding[0]
+        if self.expert_axis_name is None:
+            self.num_expert_parallelism = 1
+        else:
+            self.num_expert_parallelism = self.mesh.shape[
+                self.expert_axis_name]
+
+        # Derive if data is sharded by expert
+        self.data_axis_name = self.activation_ffw_td[0]
+        self.is_batch_sharded_by_expert = (
+            self.expert_axis_name is not None) and (self.expert_axis_name
+                                                    == self.data_axis_name)
+
+    def _sort_activations(self, inputs: jax.Array,
+                          sort_indices: jax.Array) -> jax.Array:
+        """Sorts activations(inputs) by `sort_indices` for the forward pass."""
+        return inputs[sort_indices, ...]
+
+    @staticmethod
+    def get_all_to_all_params(
+        all_shards_group_sizes,
+        shard_id,
+        num_expert_parallelism,
+        is_batch_sharded=True,
+    ):
+        """Generates params for ragged_all_to_all communication."""
+
+        class TransformStrategy(enum.Enum):
+            INPUT_OFFSET = enum.auto()
+            SEND_SIZE = enum.auto()
+            OUTPUT_OFFSET = enum.auto()
+            RECV_SIZE = enum.auto()
+
+        def transform_array(input_array, shard_id, strategy, is_batch_sharded):
+            if is_batch_sharded:
+                if strategy == TransformStrategy.INPUT_OFFSET:
+                    local_array = input_array[shard_id]
+                    return jnp.concatenate(
+                        (jnp.array([0]), jnp.cumsum(local_array)[:-1]))
+                elif strategy == TransformStrategy.SEND_SIZE:
+                    return input_array[shard_id]
+                elif strategy == TransformStrategy.OUTPUT_OFFSET:
+                    zero_row = jnp.zeros((1, ) + input_array.shape[1:],
+                                         dtype=input_array.dtype)
+                    array_with_zeros = jnp.concatenate((zero_row, input_array),
+                                                       axis=0)
+                    cumulated_array = jnp.cumsum(array_with_zeros,
+                                                 axis=0,
+                                                 dtype=input_array.dtype)
+                    return cumulated_array[shard_id]
+                elif strategy == TransformStrategy.RECV_SIZE:
+                    return input_array[:, shard_id]
+                else:
+                    raise ValueError(
+                        f"Unknown transform array strategy: {strategy}")
+            else:
+                if strategy == TransformStrategy.INPUT_OFFSET:
+                    return jnp.zeros(num_expert_parallelism,
+                                     dtype=input_array.dtype)
+                elif strategy == TransformStrategy.SEND_SIZE:
+                    return jnp.repeat(input_array[shard_id],
+                                      num_expert_parallelism)
+                elif strategy == TransformStrategy.OUTPUT_OFFSET:
+                    output_offset = jnp.concatenate(
+                        (jnp.array([0]),
+                         jnp.cumsum(input_array[:-1])))[shard_id]
+                    return jnp.repeat(output_offset, num_expert_parallelism)
+                elif strategy == TransformStrategy.RECV_SIZE:
+                    return input_array
+                else:
+                    raise ValueError(
+                        f"Unknown transform array strategy: {strategy}")
+
+        input_offsets = transform_array(all_shards_group_sizes, shard_id,
+                                        TransformStrategy.INPUT_OFFSET,
+                                        is_batch_sharded)
+        send_sizes = transform_array(all_shards_group_sizes, shard_id,
+                                     TransformStrategy.SEND_SIZE,
+                                     is_batch_sharded)
+        output_offsets = transform_array(all_shards_group_sizes, shard_id,
+                                         TransformStrategy.OUTPUT_OFFSET,
+                                         is_batch_sharded)
+        recv_sizes = transform_array(all_shards_group_sizes, shard_id,
+                                     TransformStrategy.RECV_SIZE,
+                                     is_batch_sharded)
+        return input_offsets, send_sizes, output_offsets, recv_sizes
+
+    def _local_permute(
+        self,
+        inputs,
+        global_group_sizes,
+        local_expert_size,
+        shard_index,
+        is_offset=False,
+        global_sorted_experts=None,
+    ):
+        """Permutes tokens locally within an expert shard."""
+        # global_group_sizes: (tokens parallelism, num_total_experts)
+        # all_shard_local_sizes: (tokens parallelism, num local experts in the shard)
+        all_shard_local_sizes = jax.lax.dynamic_slice_in_dim(
+            global_group_sizes,
+            shard_index * local_expert_size,
+            local_expert_size,
+            axis=1,
+        )
+        local_sizes = all_shard_local_sizes.reshape(-1)
+
+        # local_group_size: (tokens parallelism, )
+        local_group_size = jnp.sum(all_shard_local_sizes, axis=0)
+
+        # When token replicated in devices
+        if is_offset:
+            global_sorted_shard_assignments = jnp.floor_divide(
+                global_sorted_experts, local_expert_size)
+            expert_indices = jnp.where(
+                global_sorted_shard_assignments == shard_index,
+                jnp.mod(global_sorted_experts, local_expert_size),
+                local_expert_size,
+            )
+
+        # When token sharded in devices
+        else:
+            base_indices = jnp.mod(jnp.arange(local_sizes.shape[0]),
+                                   local_expert_size)
+            expert_indices = jnp.repeat(base_indices,
+                                        local_sizes,
+                                        total_repeat_length=inputs.shape[0])
+
+        sorted_indices = jnp.argsort(expert_indices)
+        # sort the inputs based on the local expert_indices
+        sorted_inputs = self._sort_activations(inputs, sorted_indices)
+        # sortted local expert id from 0 to local expert size
+        sorted_experts_ids = expert_indices[sorted_indices]
+        return (
+            sorted_inputs,
+            sorted_indices,
+            local_group_size,
+            sorted_experts_ids,
+        )
+
+    def _permute(self, inputs_TD: Float, selected_experts_TX: jax.Array):
+        """Global permute: Sorts tokens by assigned expert."""
+        # suffix t = T * X = total_assignments for the local tokens(T) on this device.
+        total_tokens = inputs_TD.shape[0]
+        flat_expert_indices = selected_experts_TX.flatten()
+        sort_indices_t = jnp.argsort(flat_expert_indices)
+
+        replicated_inputs_tD = jnp.repeat(inputs_TD,
+                                          self.num_experts_per_tok,
+                                          axis=0)
+        sorted_inputs_tD = self._sort_activations(replicated_inputs_tD,
+                                                  sort_indices_t)
+
+        # number of tokens assigned to each expert
+        group_sizes_E = jnp.bincount(flat_expert_indices,
+                                     length=self.num_local_experts)
+
+        expert_ids = jnp.arange(self.num_local_experts)
+        total_assignments = total_tokens * self.num_experts_per_tok
+        sorted_expert_assignments_t = jnp.repeat(
+            expert_ids,
+            repeats=group_sizes_E,
+            total_repeat_length=total_assignments)
+
+        return (
+            sorted_inputs_tD,
+            sort_indices_t,
+            group_sizes_E,
+            sorted_expert_assignments_t,
+        )
+
+    def _unpermute(self, processed_tokens: jax.Array, sort_indices: jax.Array,
+                   router_weights_TX: jax.Array):
+        """Unsorts tokens to their original order and combines expert outputs with router's weight."""
+        with jax.named_scope("unpermute"):
+            unsorted_tokens_tD = self._sort_activations(
+                processed_tokens, jnp.argsort(sort_indices))
+            reshaped_tokens_TXD = unsorted_tokens_tD.reshape(
+                -1, self.num_experts_per_tok, self.hidden_size)
+        with jax.named_scope("combine_weights"):
+            output_TD = jnp.einsum(
+                "TXD,TX -> TD",
+                reshaped_tokens_TXD.astype(jnp.float32),
+                router_weights_TX.astype(jnp.float32),
+                precision='float32',
+            )
+
+        return output_TD.astype(self.dtype)
+
+    def _gmm(self, inputs, kernel, group_sizes):
+        """Performs Grouped Matrix Multiply."""
+        num_rows = inputs.shape[0]
+        pad_amount = (self.tile_size[0] -
+                      num_rows % self.tile_size[0]) % self.tile_size[0]
+        if pad_amount > 0:
+            inputs = jnp.pad(inputs, ((0, pad_amount), (0, 0)))
+
+        if self.use_megablox:
+            #TODO: megablox is used in MaxText, keep a placeholder here for future implement
+            raise NotImplementedError(
+                "MegaBlox kernel call is not implemented.")
+        else:
+            inputs = manually_quantize_qwix_activation(
+                inputs, "ragged_dot", jnp.float8_e4m3fn, [0], {},
+                "absmax") if self.quantized_dtype else inputs
+            ragged_dot_func = qwix_ragged_dot if self.quantized_dtype else jax.lax.ragged_dot
+            output = ragged_dot_func(
+                lhs=inputs,
+                rhs=kernel,
+                group_sizes=group_sizes,
+                preferred_element_type=self.dtype,
+            )
+
+        if pad_amount > 0:
+            output = output[:num_rows, :]
+        return output
+
+    @staticmethod
+    def _distributed_sparse_moe_fwd(
+        self,
+        x_TD: jax.Array,
+        router_weights_TX: jax.Array,
+        selected_experts_TX: jax.Array,
+        kernel_gating: jax.Array,
+        kernel_up_proj: jax.Array,
+        kernel_down_proj: jax.Array,
+    ):
+        """
+        The sparse MoE forward pass with fully distributed logic.
+        This assumes it is running within a distributed TPU.
+        """
+
+        # 1. Global Permute, perpute all tokens across shards
+        (
+            sorted_inputs,
+            global_sort_indices,
+            global_group_sizes,
+            global_sorted_experts,
+        ) = self._permute(x_TD, selected_experts_TX)
+
+        # TODO: update to 'expert' after we enable expert parallelism, currently experts are sharded along model axis
+        # or we sould derive it from the model init
+        expert_shard_id = jax.lax.axis_index(self.expert_axis_name)
+        local_expert_size = self.num_local_experts // self.num_expert_parallelism
+
+        if self.num_expert_parallelism > 1:
+            if self.is_batch_sharded_by_expert:
+                # When token sharded in devices
+                # In this path, we assume the data(tokens) are fully sharded on expert, namely data_axis_name == expert_axis_name
+
+                # 2a. Send Tokens To Experts (All-to-All)
+                # Gather group sizes from all data shards
+                # all_shards_group_sizes: (data parallelism = expert parallelism, number of total experts )
+                all_shards_group_sizes = jax.lax.all_gather(
+                    global_group_sizes, axis_name=self.data_axis_name)
+
+                # all_shards_group_sizes_per_expert_shard[i][j] = # tokens on shard[i] to be sent to expert shard[j]
+                all_shards_group_sizes_per_expert_shard = jnp.sum(
+                    all_shards_group_sizes.reshape(
+                        self.num_expert_parallelism,  # data parallelism
+                        self.num_expert_parallelism,  # expert parallelism
+                        local_expert_size  # Experts per shard
+                    ),
+                    axis=2)
+                input_offsets, send_sizes, output_offsets, recv_sizes = self.get_all_to_all_params(
+                    all_shards_group_sizes_per_expert_shard, expert_shard_id,
+                    self.num_expert_parallelism)
+                # Estimate buffer size
+                local_total_assignments = x_TD.shape[
+                    0] * self.num_experts_per_tok
+                global_total_assignments = local_total_assignments * self.num_expert_parallelism
+                output_shape_est = jnp.zeros(
+                    (global_total_assignments, self.hidden_size),
+                    dtype=sorted_inputs.dtype)
+
+                inputs_after_all2all = jax.lax.ragged_all_to_all(
+                    sorted_inputs,
+                    output_shape_est,
+                    input_offsets,
+                    send_sizes,
+                    output_offsets,
+                    recv_sizes,
+                    axis_name=self.expert_axis_name)
+
+                # 3a. Local Permute
+                # Get full group sizes from all shards
+                full_global_group_sizes = jax.lax.all_gather(
+                    global_group_sizes, axis_name=self.expert_axis_name)
+                (
+                    compute_inputs,
+                    local_sorted_indices,
+                    compute_group_sizes,
+                    compute_expert_ids,
+                ) = self._local_permute(
+                    inputs_after_all2all,
+                    full_global_group_sizes,
+                    local_expert_size,
+                    shard_index=expert_shard_id,
+                    is_offset=False,
+                )
+
+            else:
+                # When token replicated in devices
+
+                # 2. No send all-to-all needed, as the tokens are sorted and replicated on all devices
+                # 3b. Local "Permute"
+                (
+                    compute_inputs,
+                    local_sorted_indices,
+                    compute_group_sizes,
+                    compute_expert_ids,
+                ) = self._local_permute(
+                    sorted_inputs,
+                    global_group_sizes[None, :],
+                    local_expert_size,
+                    shard_index=expert_shard_id,
+                    is_offset=True,
+                    global_sorted_experts=global_sorted_experts,
+                )
+
+                # Calculate group sizes for return all-to-all
+                reshaped_group_sizes = jnp.sum(global_group_sizes.reshape(
+                    -1, local_expert_size),
+                                               axis=1)
+                mask = compute_expert_ids < local_expert_size
+                compute_inputs = compute_inputs * mask[..., None]
+
+        else:
+            # --- NO EXPERT PARALLELISM ---
+            compute_inputs = sorted_inputs
+            compute_group_sizes = global_group_sizes
+            compute_expert_ids = global_sorted_experts
+            local_sorted_indices = jnp.arange(sorted_inputs.shape[0])
+
+        # 4. Compute: Apply experts using Grouped Matrix Multiply
+        with jax.named_scope("gating"):
+            # compute_inputs: (local total assignments, D)
+            gating_TEF = self._gmm(compute_inputs, kernel_gating,
+                                   compute_group_sizes)
+            activated_gating_TEF = modeling_flax_utils.ACT2FN[self.hidden_act](
+                gating_TEF)
+
+        with jax.named_scope("up_projection"):
+            up_proj_TEF = self._gmm(compute_inputs, kernel_up_proj,
+                                    compute_group_sizes)
+
+        fuse_TEF = activated_gating_TEF * up_proj_TEF
+
+        with jax.named_scope("down_projection"):
+            # intermediate_output: (local total assignments, D)
+            intermediate_output = self._gmm(fuse_TEF, kernel_down_proj,
+                                            compute_group_sizes)
+
+        # 5. Return Results (All-to-All)
+        if self.num_expert_parallelism > 1:
+            local_total_assignments = x_TD.shape[0] * self.num_experts_per_tok
+            output_shape = jnp.zeros(
+                (local_total_assignments, self.hidden_size),
+                dtype=intermediate_output.dtype)
+
+            if self.is_batch_sharded_by_expert:
+                # When token sharded in devices
+                # Unsort locally before sending back
+                local_output = self._sort_activations(
+                    intermediate_output, jnp.argsort(local_sorted_indices))
+
+                input_offsets, send_sizes, output_offsets, recv_sizes = self.get_all_to_all_params(
+                    jnp.transpose(all_shards_group_sizes),
+                    expert_shard_id,
+                    self.num_expert_parallelism,
+                )
+                final_intermediate_output = jax.lax.ragged_all_to_all(
+                    local_output,
+                    output_shape,
+                    input_offsets,
+                    send_sizes,
+                    output_offsets,
+                    recv_sizes,
+                    axis_name=self.expert_axis_name)
+            else:
+                # When token replicated in devices
+                input_offsets, send_sizes, output_offsets, recv_sizes = self.get_all_to_all_params(
+                    reshaped_group_sizes,
+                    expert_shard_id,
+                    self.num_expert_parallelism,
+                    is_batch_sharded=False,
+                )
+                final_intermediate_output = jax.lax.ragged_all_to_all(
+                    intermediate_output,
+                    output_shape,
+                    input_offsets,
+                    send_sizes,
+                    output_offsets,
+                    recv_sizes,
+                    axis_name=self.expert_axis_name)
+        else:
+            final_intermediate_output = intermediate_output
+
+        # 6. Global Unpermute (on the data shard)
+        with jax.named_scope("unpermute"):
+            output_TD = self._unpermute(final_intermediate_output,
+                                        global_sort_indices, router_weights_TX)
+
+        return output_TD
+
+    def __call__(self, x_TD: Float):
+        """Performs the forward pass of the Sparse MoE layer."""
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+        router_weights_TX, selected_experts_TX = self.router(x_TD)
+
+        in_specs = (
+            PartitionSpec(),  # Replicated `self`
+            PartitionSpec(*self.activation_ffw_td),  # Sharded x_TD
+            PartitionSpec(),  # Replicated router_weights_TX
+            PartitionSpec(),  # Replicated selected_experts_TX
+            PartitionSpec(*self.edf_sharding),  # Sharded gating kernel
+            PartitionSpec(*self.edf_sharding),  # Sharded up-projection kernel
+            PartitionSpec(
+                *self.efd_sharding),  # Sharded down-projection kernel
+        )
+        out_specs = PartitionSpec(*self.activation_ffw_td)
+
+        mapped_moe_fwd = partial(jax.experimental.shard_map.shard_map,
+                                 mesh=self.mesh,
+                                 in_specs=in_specs,
+                                 out_specs=out_specs,
+                                 check_rep=False)(
+                                     SparseMoE._distributed_sparse_moe_fwd)
+
+        kernel_gating_EDF = self.kernel_gating_EDF.value
+        kernel_up_proj_EDF = self.kernel_up_proj_EDF.value
+        kernel_down_proj_EFD = self.kernel_down_proj_EFD.value
+
+        if self.quantized_dtype:
+            if not isinstance(kernel_gating_EDF, ptq.WithAux):
+                kernel_gating_EDF = manually_quantize_qwix_weight(
+                    kernel_gating_EDF, self.quantized_dtype, [0, 2], {},
+                    "absmax")
+            if not isinstance(kernel_up_proj_EDF, ptq.WithAux):
+                kernel_up_proj_EDF = manually_quantize_qwix_weight(
+                    kernel_up_proj_EDF, self.quantized_dtype, [0, 2], {},
+                    "absmax")
+            if not isinstance(kernel_down_proj_EFD, ptq.WithAux):
+                kernel_down_proj_EFD = manually_quantize_qwix_weight(
+                    kernel_down_proj_EFD, self.quantized_dtype, [0, 1], {},
+                    "absmax")
+            kernel_gating_EDF = kernel_gating_EDF.array
+            kernel_up_proj_EDF = kernel_up_proj_EDF.array
+            kernel_down_proj_EFD = kernel_down_proj_EFD.array
+
+        return mapped_moe_fwd(self, x_TD, router_weights_TX,
+                              selected_experts_TX, kernel_gating_EDF,
+                              kernel_up_proj_EDF, kernel_down_proj_EFD)