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

tpu_inference/layers/jax/binary_search.py

@@ -0,0 +1,295 @@
+# Copyright 2024 The T5X Authors.
+#
+# 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.
+"""Binary search over float32 bits.
+
+Includes fast algorithms top-k masking and top-p masking on probability
+distributions.
+"""
+
+from typing import Callable, Sequence
+
+import jax
+from jax import lax
+from jax import numpy as jnp
+
+
+def int32_bsearch(batch_shape: Sequence[int],
+                  predicate: Callable[[jnp.ndarray], jnp.ndarray]):
+    """Batched binary search over int32 values.
+
+    For each element of the batch, search for the largest int32 (closest to
+    positive infinity) for which the predicate is False. If the predicate is
+    always True, returns the minimum int32 value.
+
+    Args:
+      batch_shape: Shape of the search that we're batching over.
+      predicate: the query we're searching for. For every batch element, this is
+        required to be a monotonic function from int32 to bool. In other words,
+        the predicate must return False for all numbers <= some threshold and then
+        return True for all numbers > that threshold. The threshold may be
+        different for different elements of the batch.
+
+    Returns:
+      For each element of the batch, the largest int32 for which the predicate
+      returns False. Shape: batch_shape.
+    """
+    current_bits = jnp.zeros(batch_shape, dtype=jnp.int32)
+
+    # bit 31 is special, because it compares in the opposite order of all other
+    # bits. we use uint32 due to numpy promotion/casting rules.
+    midpoint = current_bits
+    predicate_satisfied = predicate(midpoint)
+    current_bits = current_bits | jnp.where(predicate_satisfied,
+                                            jnp.uint32(1 << 31), jnp.uint32(0))
+    del midpoint, predicate_satisfied
+
+    def loop_body(i, current_bits):
+        bit_index = 30 - i
+        bit = jnp.int32(1 << bit_index)
+        midpoint = current_bits | bit
+        predicate_satisfied = predicate(midpoint)
+        current_bits = current_bits | jnp.where(predicate_satisfied,
+                                                jnp.int32(0), bit)
+        return current_bits
+
+    current_bits = lax.fori_loop(0, 31, loop_body, current_bits)
+    return current_bits
+
+
+def _monotonic_int32_to_float32_bit_pattern(x: int) -> int:
+    """Converts an int32 to a float32 bit pattern with consistent ordering.
+
+    This function is the unique function that is monotonic with respect to the
+    floating point total order, see
+    https://en.wikipedia.org/wiki/IEEE_754#Total-ordering_predicate. Note that
+    this function returns an int32, not a float32. For the function that returns
+    float32, see `monotonic_int32_to_float32`.
+
+    Args:
+      x: int bit pattern.
+
+    Returns:
+      Bit pattern of a float32 number.
+    """
+    non_sign_bits = jnp.int32((1 << 31) - 1)
+    # See
+    # https://stackoverflow.com/questions/20097380/iee-754-total-order-in-standard-c11
+    # for the relationship between int32 order and f32 total order, including
+    # the "xor trick".
+
+    # Flip the sort order for numbers where the sign bit is set. On int32,
+    # the bit pattern with sign bit set and all other bits clear is the most
+    # negative bit pattern (it's int32::MIN), whereas on float32 it's the least
+    # negative bit pattern (it's -0.0). Flipping all the non-sign bits makes the
+    # int32 sort order consistent with the float32 sort order.
+    x = x ^ jnp.where(x < 0, non_sign_bits, jnp.int32(0))
+    return x
+
+
+def _monotonic_int32_to_float32(x: int) -> jax.Array:
+    """Converts an int32 to a float32 with consistent ordering.
+
+    This function is the unique function that is monotonic with respect to the
+    floating point total order, see
+    https://en.wikipedia.org/wiki/IEEE_754#Total-ordering_predicate.
+
+    Args:
+      x: int bit pattern.
+
+    Returns:
+      float32 number with consistent ordering.
+    """
+    x = _monotonic_int32_to_float32_bit_pattern(x)
+    return lax.bitcast_convert_type(x, jnp.float32)
+
+
+def float32_bsearch(batch_shape, predicate):
+    """Binary search on finite float32 numbers.
+
+    For each element of the batch, this function searches for the largest finite
+    non-NaN float32 for which the predicate is False.
+
+    Args:
+      batch_shape: Shape of the search that we're batching over.
+      predicate: the query we're searching for. This is required to be monotonic
+        with respect to the floating point order, i.e. it must be False for all
+        numbers <= a threshold, and then True for all numbers > the threshold. The
+        threshold may be different for different elements of the batch.
+
+    Returns:
+      For each element of the batch, the largest float32 for which the predicate
+      returns False. Shape: f32[batch_shape].
+    """
+    exponent_bits = jnp.int32((1 << 31) - (1 << (31 - 8)))
+
+    def int32_predicate(x):
+        x = _monotonic_int32_to_float32_bit_pattern(x)
+        is_finite = (x & exponent_bits) != exponent_bits
+
+        # Non-finite numbers (infinity and NaN) are at the very extremes of the
+        # int32 range, i.e. they include int32::MAX and int32::MIN, plus the numbers
+        # adjacent to them. For the nonfinite numbers touching int32::MIN, we
+        # arrange for them to return False from the predicate, and for the nonfinite
+        # numbers touching int32::MAX, we arrange for them to return True from the
+        # predicate. x>=0 is an easy way to achieve that.
+        predicate_on_nonfinite = x >= 0
+        x_float32 = lax.bitcast_convert_type(x, jnp.float32)
+        return jnp.where(is_finite, predicate(x_float32),
+                         predicate_on_nonfinite)
+
+    # We search over bit patterns, which requires bit shifting and ordering of bit
+    # patterns. This is natively supported on int32 but not on float32.
+    # Additionally, it's more common to reason about int32 bit arithmetic and
+    # ordering than float32 bit arithmetic and ordering, so we do the core of our
+    # search in int32. Additionally, this allows us to test the underlying binary
+    # search on int32 values.
+    #
+    # The function _monotonic_int32_to_float32 encapsulates all of the knowledge
+    # we need about float32 bit patterns.
+    result = int32_bsearch(batch_shape, int32_predicate)
+    return _monotonic_int32_to_float32(result)
+
+
+def topk_mask(x: jnp.ndarray, k: int, replace_val: jnp.ndarray) -> jnp.ndarray:
+    """Sets everything to replace_val, except the top k values per batch element.
+
+    Sharding considerations: this function does 32 reductions over the vocab_size
+    axis of the input array. To avoid excessive latency from these reductions, you
+    should ensure that the vocab_size axis is unsharded on input to this function.
+    Prefer to shard the batch axes instead.
+
+    Scratchpad memory considerations: this function is most efficient if the
+    entire input array can fit in a fast memory tier. To help ensure this, you may
+    wish to split the batch axes into microbatches and the microbatches in a
+    sequential loop.
+
+    Args:
+      x: Values before masking. [batch..., vocab_size]
+      k: Number of masked values to return. In presence of ties, more than k
+        values might be returned.
+      replace_val: For the masked values of x, what to overwrite them with.
+
+    Returns:
+      masked version of x. [batch..., vocab_size]
+    """
+    batch_shape = tuple(list(x.shape)[:-1])  # [batch...]
+
+    x_for_loop = x
+    reduce_axis = x.ndim - 1
+    if x.ndim > 1:
+        # We're going to be doing 32 reductions over 'reduce_axis'. Generally,
+        # reductions over the last dimension are the most expensive, because they
+        # involve reducing across vector lanes, which is often not efficient. So
+        # we transpose the reduce_axis to be the second-last dimension, to avoid
+        # this inefficiency.
+        #
+        # Normaly the XLA compiler would automatically perform this optimization,
+        # but it doesn't yet see through loops to do so. So we do it ourselves.
+        x_for_loop = jnp.swapaxes(x_for_loop, -1, -2)
+        reduce_axis = x.ndim - 2
+
+    # x: [batch..., vocab_size, batch]
+    def predicate(threshold):
+        # threshold: [batch...]
+
+        # Since we've negated, we now want a predicate that is True for small
+        # numbers and False for large numbers. The result of the bsearch is the
+        # smallest float32 for which the predicate is False.
+        threshold = -threshold
+
+        threshold = lax.expand_dims(threshold, (reduce_axis, ))
+        # threshold: [batch..., 1, last_batch]
+
+        # count_ge: [batch...]
+        count_gt = jnp.sum(x_for_loop > threshold, axis=reduce_axis)
+
+        return count_gt >= k
+
+    # cutoff: [batch...]
+    cutoff = float32_bsearch(batch_shape, predicate)
+    cutoff = -cutoff
+    # cutoff: [batch..., 1]
+    cutoff = lax.expand_dims(cutoff, (cutoff.ndim, ))
+    return jnp.where(x >= cutoff, x, jnp.full_like(x, replace_val))
+
+
+def topp_mask(logits: jnp.ndarray, p: float,
+              replace_val: jnp.ndarray) -> jnp.ndarray:
+    """Applies top-p masking to logits.
+
+    Masks logits down to the smallest set of choices, such that the total
+    probability mass is >= p. Values in this set are left as they are. All other
+    values are set with `replace_val`.
+
+    Sharding considerations: this function does 33 reductions over the vocab_size
+    axis of the input array. To avoid excessive latency from these reductions, you
+    should ensure that the vocab_size axis is unsharded on input to this function.
+    Prefer to shard the batch axes instead.
+
+    Scratchpad memory considerations: this function is most efficient if the
+    entire input array can fit in a fast memory tier. To help ensure this, you may
+    wish to split the batch axes into microbatches and the microbatches in a
+    sequential loop.
+
+    Args:
+      logits: Logits before masking. [batch..., vocab_size]
+      p: Minimum probability mass requested.
+      replace_val: For the masked values of logits, what to overwrite them with.
+
+    Returns:
+      masked version of x. [batch..., vocab_size]
+    """
+    batch_shape = tuple(list(logits.shape)[:-1])  # [batch...]
+
+    probs = jax.nn.softmax(logits, axis=-1)
+
+    probs_for_reduction = probs
+    reduce_axis = probs_for_reduction.ndim - 1
+    if probs_for_reduction.ndim > 1:
+        # We're going to be doing 33 reductions over 'reduce_axis'. Generally,
+        # reductions over the last dimension are the most expensive, because they
+        # involve reducing across vector lanes, which is often not efficient. So
+        # we transpose the reduce_axis to be the second-last dimension, to avoid
+        # this inefficiency.
+        probs_for_reduction = jnp.swapaxes(probs_for_reduction, -1, -2)
+        reduce_axis = probs_for_reduction.ndim - 2
+
+    # As we increase the threshold, the probability mass decreases, and the number
+    # selected decreases.
+    #
+    # We want the largest threshold with the probability mass >= p. Binary search
+    # searches for when the predicate is False, so we negate the output of the
+    # predicate, i.e. probability mass < p.
+
+    # probs_for_reduction: [batch..., vocab_size, batch]
+    def predicate(threshold):
+        # threshold: [batch...]
+        threshold = lax.expand_dims(threshold, (reduce_axis, ))
+        # threshold: [batch..., 1, last_batch]
+
+        # count_ge: [batch...]
+        probability_mass = jnp.sum(
+            jnp.where(probs_for_reduction >= threshold, probs_for_reduction,
+                      0.0),
+            axis=reduce_axis,
+        )
+
+        return probability_mass < p
+
+    # threshold: [batch...]
+    threshold = float32_bsearch(batch_shape, predicate)
+    # threshold: [batch..., 1]
+    threshold = lax.expand_dims(threshold, (threshold.ndim, ))
+    return jnp.where(probs >= threshold, logits,
+                     jnp.full_like(logits, replace_val))

tpu_inference/layers/jax/constants.py

@@ -0,0 +1,88 @@
+"""
+Current Used Abbreviation for Tensor Dimensions:
+B: Batch size
+T: Sequence Length (for Query tensors)
+S: Sequence Length (for Key/Value tensors)
+D: d_model, the embedding dimension of the model
+F: d_ff, the hidden dimension of the feed-forward MLP layers
+V: Vocab Size
+H: Dimension of each attention head
+N: Number of query heads in Attention
+Q: Number of query heads (synonymous with N)
+K: Number of Key/Value heads in Attention
+C: Expert capacity in Mixture-of-Experts models
+X: Number of activated experts per token in MoE
+G: Number of groups in Grouped-Query Attention
+E: Total number of experts in MoE
+"""
+
+import enum
+from typing import Tuple, TypeAlias
+
+import jax
+
+KVCacheType: TypeAlias = Tuple[jax.Array, jax.Array]
+
+
+class RouterType(enum.Enum):
+    """Enum for router types."""
+    TOP_K = 'top_k'
+
+
+class OPERATION_MODE(enum.Enum):
+    PREFILL = 1
+    DECODE = 2
+
+
+class HuggingFaceArgNames(enum.Enum):
+    ## Modeling params
+    HIDDEN_ACT: str = "hidden_act"
+    HIDDEN_SIZE: str = "hidden_size"
+    NUM_HIDDEN_LAYERS: str = "num_hidden_layers"
+    RMS_NORM_EPS: str = "rms_norm_eps"
+    ROPE_SCALING: str = "rope_scaling"
+    ROPE_THETA: str = "rope_theta"
+    VOCAB_SIZE: str = "vocab_size"
+
+    # Block parameters
+    SHARED_EXPERTS: str = "shared_experts"
+
+    # FFW params
+    INTERMEDIATE_SIZE: str = "intermediate_size"
+
+    # Attention params
+    HEAD_DIM: str = "head_dim"
+    NUM_ATTENTION_HEADS: str = "num_attention_heads"
+    NUM_KEY_VALUE_HEADS: str = "num_key_value_heads"
+    ATTENTION_DROPOUT: str = "attention_dropout"
+    ATTENTION_BIAS: str = "attention_bias"
+    ATTENTION_CHUNK_SIZE: str = "attention_chunk_size"
+
+    ## Llama4 Attention Params
+    USE_QK_NORM: str = "use_qk_norm"
+    TEMPERATURE_TUNING: str = "temperature_tuning"
+    TEMPERATURE_TUNING_SCALE: str = "temperature_tuning_scale"
+    TEMPERATURE_TUNING_FLOOR_SCALE: str = "temperature_tuning_floor_scale"
+
+    # MLA params
+    KV_LORA_RANK: str = "kv_lora_rank"
+    Q_LORA_RANK: str = "q_lora_rank"
+    QK_NOPE_HEAD_DIM: str = "qk_nope_head_dim"
+    QK_ROPE_HEAD_DIM: str = "qk_rope_head_dim"
+    V_HEAD_DIM: str = "v_head_dim"
+
+    # MoE
+    INTERMEDIATE_SIZE_MOE: str = "intermediate_size_moe"
+    NUM_LOCAL_EXPERTS: str = "num_local_experts"  # Llama moe
+    NUM_EXPERTS_PER_TOKEN: str = "num_experts_per_token"
+    NUM_ROUTED_EXPERTS: str = "n_routed_experts"  # Deepseek moe
+    NUM_SHARED_ROUTED_EXPERTS: str = "n_shared_experts"
+    NUM_GROUPS: str = "n_group"
+    ROUTED_SCALING_FACTOR: str = "routed_scaling_factor"
+    TOPK_GROUP: str = "topk_group"
+    NORM_TOPK_PROB: str = "norm_topk_prob"
+    SCORING_FUNCTION: str = "scoring_func"
+
+    ## Sampling params
+    BOS_TOKEN_ID: str = "bos_token_id"
+    EOS_TOKEN_ID: str = "eos_token_id"

tpu_inference/layers/jax/layers.py

@@ -0,0 +1,301 @@
+from dataclasses import InitVar, dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from flax.typing import Sharding
+from jaxtyping import Float, Int
+
+from tpu_inference.layers.jax.base import create_param
+
+
+# A dummy for modeling_flax_utils which might contain activation functions
+class FlaxUtils:
+    """A dummy class to namespace activation functions, mimicking external utilities."""
+    ACT2FN = {
+        'silu': nnx.silu,
+        'gelu': nnx.gelu,
+        'relu': nnx.relu,
+        'sigmoid': nnx.sigmoid,
+        'softmax': nnx.softmax
+    }
+
+
+modeling_flax_utils = FlaxUtils()
+
+
+@dataclass
+class RuntimeParams:
+    """A container for runtime parameters needed by neural network blocks.
+
+    This dataclass acts as a flexible container to pass objects that are only
+    available at runtime (like a pre-allocated KV cache or dynamic sharding
+    configurations) into the initialization of stateful modules. This avoids
+    having to update the constructor signature of every module when a new
+    runtime dependency is introduced.
+
+    Attributes:
+        kv_cache: The key-value cache object for attention layers.
+        sharding_cfg: The configuration for tensor sharding.
+        quantization: Configuration for quantization schemes.
+    """
+    kv_cache: Any = None
+    sharding_cfg: Any = None
+    quantization: Any = None
+
+
+@dataclass(kw_only=True)
+class RMSNorm(nnx.Module):
+    """An implementation of Root Mean Square Layer Normalization.
+
+    Attributes:
+        dims: The feature dimension to normalize over.
+        epsilon: A small float added to the variance to avoid division by zero.
+        with_scale: If True, learns a multiplicative scale parameter.
+        dtype: The data type for computations.
+    """
+    dims: int
+    activation_ffw_td: Sharding = ()
+    random_init: bool = False
+    epsilon: float = 1e-6
+    with_scale: bool = True
+    dtype: Any = jnp.float32
+
+    rngs: InitVar[nnx.Rngs]
+
+    def __call__(self, x_TD: Float, op_mode='generate') -> Float:
+        """Applies RMS Normalization to the input tensor.
+
+        Args:
+            x_TD: The input tensor. The normalization is applied over the last dimension.
+
+        Returns:
+            The normalized tensor with the same shape as the input.
+        """
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+
+        with jax.named_scope("rms_norm_variance"):
+            var_T1 = jnp.mean(jnp.square(x_TD), axis=-1, keepdims=True)
+        with jax.named_scope("rms_norm_rsqrt"):
+            normed_x_TD = x_TD * jax.lax.rsqrt(var_T1 + self.epsilon)
+
+        with jax.named_scope("rms_norm_scale_apply"):
+            normed_x_TD *= self.scale.value
+        normed_x_TD = nnx.with_sharding_constraint(normed_x_TD,
+                                                   self.activation_ffw_td)
+        return normed_x_TD.astype(self.dtype)
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        self.scale = create_param(rngs,
+                                  shape=(self.dims, ),
+                                  dtype=self.dtype,
+                                  random_init=self.random_init)
+
+
+@dataclass(kw_only=True)
+class DenseFFW(nnx.Module):
+    """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: Sharding = ()
+    fd_sharding: Sharding = ()
+    activation_ffw_td: Sharding = ()
+    random_init: bool = False
+
+    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)`.
+        """
+        # TODO consider to create factories for einsum(?)
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.activation_ffw_td)
+        with jax.named_scope("wi_0"):
+            gating_TF = jnp.einsum('TD,DF -> TF', x_TD,
+                                   self.kernel_gating_DF.value)
+            activated_gating_TF = modeling_flax_utils.ACT2FN[self.hidden_act](
+                gating_TF)
+        with jax.named_scope("wi_1"):
+            up_proj_TF = jnp.einsum('TD,DF -> TF', x_TD,
+                                    self.kernel_up_proj_DF.value)
+        fuse_TF = activated_gating_TF * up_proj_TF
+        with jax.named_scope("wo"):
+            output_TD = jnp.einsum('TF,FD -> TD', fuse_TF,
+                                   self.kernel_down_proj_FD.value)
+
+        return output_TD
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        D = self.hidden_size
+        F = self.intermediate_size
+
+        self.kernel_gating_DF = create_param(rngs,
+                                             shape=(D, F),
+                                             dtype=self.dtype,
+                                             sharding=self.df_sharding,
+                                             random_init=self.random_init)
+        self.kernel_up_proj_DF = create_param(rngs,
+                                              shape=(D, F),
+                                              dtype=self.dtype,
+                                              sharding=self.df_sharding,
+                                              random_init=self.random_init)
+        self.kernel_down_proj_FD = create_param(rngs,
+                                                shape=(F, D),
+                                                dtype=self.dtype,
+                                                sharding=self.fd_sharding,
+                                                random_init=self.random_init)
+
+
+@dataclass(kw_only=True)
+class Embedder(nnx.Module):
+    """A module for token embedding and, optionally, decoding (tied embeddings).
+
+    This class handles both the "encoding" step of converting token IDs to dense
+    vectors and the "decoding" step of projecting model outputs back to logits
+    over the vocabulary.
+
+    """
+    vocab_size: int
+    hidden_size: int
+    dtype: jnp.dtype
+    prelogit_td: Sharding = ()
+    vd_sharding: Sharding = ()
+    random_init: bool = False
+    normalize_embeddings: bool = False
+
+    rngs: InitVar[nnx.Rngs]
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        self.input_embedding_table_VD = create_param(
+            rngs,
+            shape=(self.vocab_size, self.hidden_size),
+            sharding=self.vd_sharding,
+            dtype=self.dtype,
+            random_init=self.random_init)
+
+    def __call__(self, x, decode=False):
+        """Dispatches to either the encode or decode method.
+
+        Args:
+            x: The input tensor. Either token IDs for encoding or hidden states
+                for decoding.
+            decode: A boolean flag. If False (default), performs encoding. If
+                True, performs decoding.
+
+        Returns:
+            Either embedding vectors or logit scores.
+        """
+        if decode:
+            return self.decode(x)
+        else:
+            return self.encode(x)
+
+    def decode(self, x_TD: Float) -> Float:
+        """Projects hidden states to vocabulary logits.
+
+        Args:
+            x_TD: The input tensor of hidden states from the model backbone, with
+                shape `(sequence, d_model)`.
+
+        Returns:
+            The output logits over the vocabulary, with shape
+            `(sequence, vocab_size)`.
+        """
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.prelogit_td)
+
+        with jax.named_scope("embedder_decode_projection"):
+            logits_TV = jnp.einsum('VD,TD -> TV',
+                                   self.input_embedding_table_VD.value, x_TD)
+        return logits_TV
+
+    def encode(self, x_T: Int) -> Float:
+        """Converts integer token IDs to dense embedding vectors.
+
+        Args:
+            x_T: The input tensor of token IDs, with shape `(sequence, )`.
+
+        Returns:
+            The corresponding embedding vectors, with shape
+            `(batch, sequence, d_model)`.
+        """
+        with jax.named_scope("embedder_encode_lookup"):
+            embedding_TD = jnp.take(self.input_embedding_table_VD.value,
+                                    x_T,
+                                    axis=0)
+
+        if self.normalize_embeddings:
+            with jax.named_scope("embedder_normalize_embeddings"):
+                embedding_TD *= jnp.sqrt(self.hidden_size).astype(self.dtype)
+        return embedding_TD
+
+
+@dataclass(kw_only=True)
+class LMhead(Embedder):
+    """
+    An Embedder that uses a (D, V) shaped embedding table, inheriting from
+    the base Embedder class.
+
+    This implementation overrides the kernel generation, encoding, and decoding
+    methods to work with the transposed embedding matrix layout.
+    """
+    dv_sharding: Sharding
+
+    def __post_init__(self, rngs: nnx.Rngs):
+        self.input_embedding_table_DV = create_param(
+            rngs,
+            shape=(self.hidden_size, self.vocab_size),
+            sharding=self.dv_sharding,
+            dtype=self.dtype,
+            random_init=self.random_init)
+
+    def __call__(self, x):
+        """Dispatches to decode method.
+
+        Args:
+            x: The input tensor. Either token IDs for encoding or hidden states
+                for decoding.
+            decode: A boolean flag. If False (default), performs encoding. If
+                True, performs decoding.
+
+        Returns:
+            Either embedding vectors or logit scores.
+        """
+        return self.decode(x)
+
+    def decode(self, x_TD: Float) -> Float:
+        """Projects hidden states to vocabulary logits.
+
+        Args:
+            x_TD: The input tensor of hidden states from the model backbone, with
+                shape `(sequence, d_model)`.
+
+        Returns:
+            The output logits over the vocabulary, with shape
+            `(sequence, vocab_size)`.
+        """
+        x_TD = jnp.asarray(x_TD, self.dtype)
+        x_TD = nnx.with_sharding_constraint(x_TD, self.prelogit_td)
+
+        with jax.named_scope("lmhead_decode_projection"):
+            logits_TV = jnp.einsum('DV,TD -> TV',
+                                   self.input_embedding_table_DV.value, x_TD)
+        return logits_TV

tpu_inference/layers/jax/misc.py

@@ -0,0 +1,16 @@
+import math
+from typing import Tuple
+
+import jax
+from jax.sharding import NamedSharding
+from jax.sharding import PartitionSpec as P
+
+
+# TODO(xiang): move this to weight_utils.py
+def shard_put(x: jax.Array, sharding_names: Tuple[str, ...] | P,
+              mesh: jax.sharding.Mesh) -> jax.Array:
+    # Single device sharding requires this special handling
+    # to avoid the recursive jit error.
+    if math.prod(mesh.axis_sizes) == 1:
+        return jax.device_put(x, mesh.devices.flatten()[0])
+    return jax.device_put(x, NamedSharding(mesh, P(*sharding_names)))