code/models/tt_transformers/tt/model.py

# SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0


import math

import torch
from tqdm import tqdm

import ttnn
from models.common.lightweightmodule import LightweightModule
from models.common.rmsnorm import RMSNorm
from models.common.sampling.generator import SamplingGenerator
from models.common.sampling.tt_sampling import TOPK_MAX_WIDTH, TTSampling
from models.tt_transformers.tt.ccl import TT_CCL
from models.tt_transformers.tt.common import Mode, copy_host_to_device
from models.tt_transformers.tt.decoder import TransformerBlock
from models.tt_transformers.tt.distributed_norm import DistributedNorm
from models.tt_transformers.tt.embedding import Embedding, ScaledEmbedding
from models.tt_transformers.tt.lm_head import LMHead
from models.tt_transformers.tt.model_config import TensorGroup
from models.tt_transformers.tt.rope import HfRotarySetup, RotarySetup


def _get_trace_rope_table_len(max_seq_len, trace_prefill_seq_lens):
    if not trace_prefill_seq_lens:
        return max_seq_len

    slice_alignment = math.lcm(*trace_prefill_seq_lens)
    min_table_len = max_seq_len + max(trace_prefill_seq_lens)
    return ((min_table_len + slice_alignment - 1) // slice_alignment) * slice_alignment


def _prefill_rope_setups_to_pad(rope_setup, rope_local_setup, rope_setup_class):
    """Select the RoPE setups whose shared prefill tables Transformer slices.

    A caller-supplied rope_setup_class builds its prefill cosine and sine mats
    per request on the host, inside its own prepare_inputs_prefill, and its
    forward never calls Transformer._slice_prefill_rot_mats. Such a setup owns
    no cos_matrix_prefill to pad. rope_local_setup always comes from the
    built-in classes, so it always owns one.
    """
    rope_setups = [] if rope_setup_class is not None else [rope_setup]
    if rope_local_setup is not None:
        rope_setups.append(rope_local_setup)
    return rope_setups


def _pad_prefill_rope_tables(rope_setups, max_seq_len, trace_prefill_seq_lens):
    table_len = _get_trace_rope_table_len(max_seq_len, trace_prefill_seq_lens)
    pad_len = table_len - max_seq_len
    if pad_len == 0:
        return

    padding = [(0, 0), (0, 0), (0, pad_len), (0, 0)]
    for rope_setup in rope_setups:
        rope_setup.cos_matrix_prefill = ttnn.pad(
            rope_setup.cos_matrix_prefill,
            padding=padding,
            value=0.0,
        )
        rope_setup.sin_matrix_prefill = ttnn.pad(
            rope_setup.sin_matrix_prefill,
            padding=padding,
            value=0.0,
        )


class Transformer(LightweightModule):
    def __init__(
        self,
        args,
        dtype,
        mesh_device,
        state_dict,
        weight_cache_path,
        paged_attention_config=None,
        use_paged_kv_cache=False,
        attention_class=None,
        rope_setup_class=None,
        block_class=None,
        lm_head_cls=None,
        final_norm_builder=None,
        prefetcher=None,
    ):
        super().__init__()
        self.args = args
        self.vocab_size = args.vocab_size
        assert self.vocab_size > 0
        self.n_layers = args.n_layers
        self.mesh_device = mesh_device
        self.dtype = dtype
        self.model_config = args.get_model_config()
        self.grid_size = self.args.max_grid_size
        state_dict_prefix = args.get_state_dict_prefix("", None)
        self.decoders_optimizations = args.decoders_optimizations
        self.prefetcher = prefetcher
        self.tt_ccl = TT_CCL(self.mesh_device)
        # Runtime bounds for the post-prefill tail's slice. Allocated here, before any trace exists,
        # and rewritten in place per call - see process_logits_after_prefill_trace.
        # These buffers belong to this model/DP lane. Calls on a lane enqueue the
        # copy and slice on the same command queue, in order. Concurrent host
        # calls on the same model instance are not supported.
        self._tail_slice_start = ttnn.from_torch(
            torch.zeros(4, dtype=torch.int32),
            device=self.mesh_device,
            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
        )
        self._tail_slice_end = ttnn.from_torch(
            torch.zeros(4, dtype=torch.int32),
            device=self.mesh_device,
            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
        )

        embd_kwargs = {
            "mesh_device": mesh_device,
            "args": args,
            "weight_cache_path": args.weight_cache_path(dtype),
            "state_dict": state_dict,
            "dtype": ttnn.bfloat16,  # Row major layout requires bfloat16
        }
        if self.args.embed_scale is not None:
            embd_cls = ScaledEmbedding
            embd_kwargs["embed_scale"] = self.args.embed_scale
        else:
            embd_cls = Embedding
        self.embd = embd_cls(**embd_kwargs)

        DefaultRopeSetup = HfRotarySetup if self.args.use_hf_rope else RotarySetup
        ActualRopeSetupClass = rope_setup_class if rope_setup_class is not None else DefaultRopeSetup
        # NoPE global layers (EXAONE-4.x): full-attention layers apply no rotary at
        # all, so the global setup's cos/sin are neutralized to the identity. Only
        # the Meta-style RotarySetup implements this.
        use_global_nope = getattr(args, "use_global_nope", False)
        if use_global_nope and self.args.use_hf_rope:
            raise NotImplementedError("use_global_nope (NoPE global layers) requires the Meta-style RotarySetup")
        global_rope_kwargs = {"nope": True} if use_global_nope else {}
        self.rope_setup = ActualRopeSetupClass(
            device=mesh_device,
            batch_size=args.max_batch_size,
            head_dim=args.head_dim,
            max_seq_len=args.max_seq_len,
            rope_theta=args.rope_theta,
            rope_scaling=args.rope_scaling,
            use_qk_fused=args.use_qk_fused,
            prefetcher=prefetcher,
            **global_rope_kwargs,
        )

        if args.rope_theta_local:
            self.rope_local_setup = DefaultRopeSetup(
                mesh_device,
                args.max_batch_size,
                args.head_dim,
                args.max_seq_len,
                args.rope_theta_local,
                # Most hybrid models (Gemma-3) use unscaled rope on sliding layers;
                # EXAONE-4.x instead applies its llama3-scaled rope there and sets
                # rope_scaling_local (the global layers being NoPE).
                rope_scaling=getattr(args, "rope_scaling_local", None),
                use_qk_fused=args.use_qk_fused,
                prefetcher=None,
            )

        # Dynamic starts share one table across fixed-width trace buckets. The
        # tail prevents out-of-range reads and the common multiple preserves
        # the tensor-bound slice partition geometry for every traced length.
        _pad_prefill_rope_tables(
            _prefill_rope_setups_to_pad(
                self.rope_setup,
                getattr(self, "rope_local_setup", None),
                rope_setup_class,
            ),
            args.max_seq_len,
            args.trace_prefill_supported_seq_lens,
        )

        self.trans_mats_dict = self.rope_setup.get_both_trans_mats()

        # Device tensors used to build dynamic slice params for prefill RoPE slicing.
        # Keeps chunk_start_idx-driven slicing inside the traced graph.
        self._tt_seq_len_buffer = ttnn.from_torch(
            torch.tensor([1, 1, self.args.max_seq_len, self.args.head_dim], dtype=torch.int32),
            device=self.mesh_device,
            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
        )
        self._tt_slice_start_zeros_4 = ttnn.from_torch(
            torch.tensor([0, 0, 0, 0], dtype=torch.int32),
            device=self.mesh_device,
            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
        )

        # Model-family dispatch (Command-R / cohere): swap the decoder block, final
        # norm and LM head. Lazy imports keep the experimental cohere module out of the default path.
        ActualBlockClass = block_class
        ActualLMHeadCls = lm_head_cls
        final_norm_builder_resolved = final_norm_builder
        if str(getattr(self.args, "model_type", None) or "").lower() == "cohere":
            from models.experimental.cohere.tt.cohere_decoder import CohereDecoderLayer
            from models.experimental.cohere.tt.cohere_lm_head import CohereLMHead
            from models.experimental.cohere.tt.cohere_norm import build_cohere_final_norm

            ActualBlockClass = ActualBlockClass or CohereDecoderLayer
            ActualLMHeadCls = ActualLMHeadCls or CohereLMHead
            final_norm_builder_resolved = final_norm_builder_resolved or build_cohere_final_norm
        ActualBlockClass = ActualBlockClass or TransformerBlock
        ActualLMHeadCls = ActualLMHeadCls or LMHead

        self.layers = [
            ActualBlockClass(
                args=args,
                mesh_device=mesh_device,
                tt_ccl=self.tt_ccl,
                dtype=dtype,
                state_dict=state_dict,
                weight_cache_path=weight_cache_path,
                layer_num=i,
                transformation_mats=self.trans_mats_dict,
                paged_attention_config=paged_attention_config,
                use_paged_kv_cache=use_paged_kv_cache,
                attention_class=attention_class,
                prefetcher=prefetcher,
            )
            for i in tqdm(range(self.n_layers))
        ]
        self.norm = (
            final_norm_builder_resolved(
                args=args,
                mesh_device=mesh_device,
                state_dict=state_dict,
                weight_cache_path=weight_cache_path,
                dtype=dtype,
                tt_ccl=self.tt_ccl,
            )
            if final_norm_builder_resolved is not None
            else DistributedNorm(
                RMSNorm(
                    device=mesh_device,
                    dim=args.dim,
                    eps=args.norm_eps,
                    state_dict=state_dict,
                    state_dict_prefix=args.get_state_dict_prefix("", None),
                    weight_cache_path=None if args.dummy_weights else weight_cache_path,
                    weight_dtype=ttnn.bfloat16,
                    weight_key="norm",
                    add_unit_offset=self.args.rms_norm_add_unit_offset,
                    is_distributed=self.args.is_distributed_norm,
                    ccl_topology=self.args.ccl_topology(),
                    tt_ccl=self.tt_ccl,
                ),
                args,
                tt_ccl=self.tt_ccl,
                prefetcher=prefetcher,
                TG=args.is_galaxy,
            )
        )  # close the final_norm_builder_resolved conditional-expression paren

        self.lm_head = ActualLMHeadCls(
            args=args,
            mesh_device=mesh_device,
            tt_ccl=self.tt_ccl,
            dtype=dtype,
            state_dict=state_dict,
            state_dict_prefix=state_dict_prefix,
            weight_cache_path=weight_cache_path,
            max_columns_per_device=self.args.max_columns_per_device_lm_head,
            prefetcher=prefetcher,
        )

        # Initialize on-device sampling if supported
        # Sampling on device is supported only if each device holds at most TOPK_MAX_WIDTH logits.
        # On a single device TTSampling cuts the padded vocab into as many same-device chunks as
        # needed (power-of-two, each <= TOPK_MAX_WIDTH), so any vocab it can cut tile-aligned is
        # supported (#53064); anything it cannot falls back to host sampling.
        padded_vocab_size = getattr(self.args, "padded_vocab_size", None) or self.args.vocab_size
        if list(self.mesh_device.shape) != [1, 1]:
            vocab_fits_on_device = padded_vocab_size // self.args.num_devices <= TOPK_MAX_WIDTH
        else:
            vocab_fits_on_device = TTSampling.num_single_device_vocab_splits(padded_vocab_size) is not None
        self._supports_on_device_sampling = prefetcher is None and vocab_fits_on_device
        if self._supports_on_device_sampling:
            self.sampling = SamplingGenerator(
                args=args,
                mesh_device=mesh_device,
                tt_ccl=self.tt_ccl,
            )
        else:
            self.sampling = None

    def update_weights(
        self,
        hf_state_dict: dict[str, ttnn.Tensor],
        *,
        hf_rope: bool = False,
    ) -> None:
        """In-place replace every weight from an HF-keyed dict of on-device 4D
        ttnn tensors (replicated, DRAM-interleaved, TILE, bf16). Keys follow HF
        safetensors naming; shapes are HF Linear/gamma/embedding wrapped in two
        leading unit dims.

        Strict by construction: every required key must be present (missing ->
        ``KeyError``) and every provided key consumed by exactly one leaf
        ``.update()`` (extras -> ``ValueError``). No "loose" mode -- silent
        partial updates are an expensive class of bug.

        ``hf_rope=False`` (default): caller has already permuted Q/K rows into
        this model's convention (right for the ttml -> TTT transfer, both store
        Meta-permuted rows). ``hf_rope=True`` defers HF -> Meta permutation to
        ``Attention.update`` (currently raises -- kernel not wired up).

        Tied embeddings: the protocol still requires both
        ``model.embed_tokens.weight`` and ``lm_head.weight`` (typically the same
        source tensor), keeping dispatch one-to-one with device buffers.

        Every existing buffer keeps its device allocation, so captured traces
        and the prefetcher's recorded addresses stay valid.
        """
        unconsumed = set(hf_state_dict.keys())

        def consume(key: str) -> ttnn.Tensor:
            if key not in hf_state_dict:
                raise KeyError(f"Transformer.update_weights: missing required HF key {key!r}")
            unconsumed.discard(key)
            return hf_state_dict[key]

        # Top-level (always required).
        self.embd.update(embed_tokens=consume("model.embed_tokens.weight"))
        self.norm.update(weight=consume("model.norm.weight"))
        self.lm_head.update(weight=consume("lm_head.weight"))

        # Per-layer: prefix-strip into a layer-local dict, dispatch.
        for i, block in enumerate(self.layers):
            prefix = f"model.layers.{i}."
            layer_dict = {}
            for key in list(hf_state_dict.keys()):
                if key.startswith(prefix):
                    layer_dict[key[len(prefix) :]] = hf_state_dict[key]
                    unconsumed.discard(key)
            block.update_weights(layer_dict, hf_rope=hf_rope)

        if unconsumed:
            sample = sorted(unconsumed)[:10]
            raise ValueError(
                f"Transformer.update_weights: {len(unconsumed)} HF key(s) not "
                f"consumed by any leaf .update(). This usually means a typo, "
                f"a stray weight, or a layer-index off-by-one. "
                f"Showing up to 10: {sample}"
            )

    def process_logits_after_prefill_trace(self, logits, last_token_idx):
        get_last_token = (last_token_idx // 32) * 32
        seq_len = int(logits.shape[-2])
        # Pass the offset as a runtime argument rather than a compile-time attribute. With literal
        # bounds every distinct prompt offset compiles its own slice program, and since this runs
        # after the prefill traces are captured - once per data-parallel group - that was the single
        # largest source of buffers left live across trace replays on a DP run. Warmup cannot cover
        # it either: it only ever sees bucket-length mock prompts, and real prompts are shorter.
        #
        # The tensor-args path needs the slice tile-aligned, which this one already is: it takes 32
        # rows starting at a multiple of 32. num_devices splits the sequence into equal parts, so
        # seq_len // 32 gives exactly the 32-row window, and the program then keys on the padded
        # prefill bucket instead of the offset.
        if seq_len % 32 == 0:
            for device_tensor, values in (
                (self._tail_slice_start, [0, 0, get_last_token, 0]),
                (self._tail_slice_end, [1, 1, get_last_token + 32, int(logits.shape[-1])]),
            ):
                ttnn.copy_host_to_device_tensor(
                    ttnn.from_torch(
                        torch.tensor(values, dtype=torch.int32),
                        mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
                    ),
                    device_tensor,
                )
            logits = ttnn.slice(
                input_tensor=logits,
                starts=self._tail_slice_start,
                ends=self._tail_slice_end,
                slice_dim=2,
                num_devices=seq_len // 32,
            )
        else:
            logits = ttnn.slice(
                logits,
                (0, 0, get_last_token, 0),
                (1, 1, get_last_token + 32, logits.shape[-1]),
            )
        logits = self._apply_norm_and_lm_head(logits)
        return logits

    def extract_last_tokens_batched_prefill(
        self, hidden_states, last_token_idx_list, padded_batch, prefill_seq_len, target_batch=None
    ):
        """Extract each user's last-token hidden state from batched prefill output.

        Reads hidden states to host, extracts the relevant row for each user,
        and sends the combined tensor back to device with the correct column-sharded
        mesh mapping (ShardTensorToMesh dim=-1) so the DistributedNorm all-gather
        produces the correct full hidden dim.

        Args:
            hidden_states: [padded_batch, 1, prefill_seq_len, dim_per_device] on device (column-sharded, TILE_LAYOUT)
            last_token_idx_list: list of length padded_batch with per-user last token positions
            padded_batch: number of slots (typically 32)
            prefill_seq_len: padded sequence length per user

        Returns:
            user_tokens: [1, 1, target_batch or padded_batch, dim_per_device] per device,
            column-sharded, TILE_LAYOUT
        """
        active_indices = [lt for lt in last_token_idx_list if lt > 0]
        all_same = len(set(active_indices)) <= 1

        if all_same and active_indices:
            common_last = active_indices[0]
            get_last = (common_last // 32) * 32
            R = common_last % 32
            block = ttnn.slice(
                hidden_states,
                (0, 0, get_last, 0),
                (padded_batch, 1, get_last + 32, hidden_states.shape[-1]),
            )
        else:
            block = hidden_states
            R = None

        host_tensors = [ttnn.to_torch(dt) for dt in ttnn.get_device_tensors(block)]
        host_full = torch.cat(host_tensors, dim=-1)

        if R is not None:
            combined = host_full[:, :, R : R + 1, :].reshape(1, 1, padded_batch, -1).contiguous()
        else:
            rows = []
            for slot in range(padded_batch):
                lt_idx = last_token_idx_list[slot]
                rows.append(host_full[slot : slot + 1, :, lt_idx : lt_idx + 1, :])
            combined = torch.cat(rows, dim=0).reshape(1, 1, padded_batch, -1).contiguous()

        target_batch = padded_batch if target_batch is None else target_batch
        if target_batch < padded_batch:
            raise ValueError(f"target_batch {target_batch} must be >= padded_batch {padded_batch}")
        if target_batch > padded_batch:
            padded_combined = torch.zeros(
                1,
                1,
                target_batch,
                combined.shape[-1],
                dtype=combined.dtype,
            )
            padded_combined[:, :, :padded_batch, :] = combined
            combined = padded_combined

        user_tokens = ttnn.from_torch(
            combined,
            device=self.mesh_device,
            dtype=ttnn.bfloat16,
            layout=ttnn.TILE_LAYOUT,
            mesh_mapper=ttnn.ShardTensorToMesh(self.mesh_device, dim=-1),
        )
        return user_tokens

    def process_logits_after_batched_prefill(self, hidden_states, last_token_idx_list, padded_batch, prefill_seq_len):
        """Extract last tokens and run norm + lm_head once for all users."""
        user_tokens = self.extract_last_tokens_batched_prefill(
            hidden_states, last_token_idx_list, padded_batch, prefill_seq_len
        )
        return self._apply_norm_and_lm_head(user_tokens)

    def _apply_final_logit_softcapping(self, logits):
        """Gemma-2 final logit soft-capping: logits -> tanh(logits / cap) * cap.

        No-op unless args.final_logit_softcapping is set (only Gemma-2 sets it), so
        this leaves every other model's output path unchanged.
        """
        cap = self.args.final_logit_softcapping
        if cap is None or cap <= 0:
            return logits
        logits = ttnn.multiply(logits, 1.0 / cap)
        logits = ttnn.tanh(logits)
        logits = ttnn.multiply(logits, cap)
        return logits

    def _apply_norm_and_lm_head(self, x):
        """Shared norm + lm_head for prefill logit processing. Input: [1, 1, 32, hidden_dim]."""
        x = self.norm(
            x, mode=Mode.PREFILL, norm_config=self.args.get_norm_config("lm_head", Mode.PREFILL, self.prefetcher)
        )
        lm_head_input_mem_cfg = self.args.get_lm_head_input_mem_config(Mode.PREFILL, None)
        if lm_head_input_mem_cfg.is_sharded():
            x = ttnn.interleaved_to_sharded(x, lm_head_input_mem_cfg)
        logits = self.lm_head(x)
        logits = self._apply_final_logit_softcapping(logits)
        logits = ttnn.to_memory_config(logits, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        return logits

    def process_hidden_states_after_prefill_trace(self, hidden_states, last_token_idx):
        """
        Process hidden states after prefill trace, stopping before LM head.
        Returns hidden states (after norm) instead of logits.
        Used for embedding models that need hidden states rather than logits.
        """
        get_last_token = (last_token_idx // 32) * 32
        hidden_states = ttnn.slice(
            hidden_states,
            (0, 0, get_last_token, 0),
            (1, 1, get_last_token + 32, hidden_states.shape[-1]),
        )
        # Apply norm (this is the final layer norm before LM head)
        hidden_states = self.norm(hidden_states, mode="prefill")
        # Convert to row major layout for output (but don't apply LM head)
        hidden_states = ttnn.to_layout(
            hidden_states, layout=ttnn.ROW_MAJOR_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG
        )
        return hidden_states

    def prepare_prefill_inputs_trace(
        self,
        tokens,
        page_table=None,
        chunk_page_table=None,
        chunk_start_idx=0,
        batch_size=1,
        user_id=0,
        **kwargs,
    ):
        """
        Inputs are torch tensors or python types. This function returns ttnn
        tensors on host.
        """
        host_inputs = self.prepare_inputs_prefill(
            tokens,
            page_table=page_table,
            chunk_page_table=chunk_page_table,
            chunk_start_idx=chunk_start_idx,
            trace_enabled=True,
            batch_size=batch_size,
            user_id=user_id,
        )
        return host_inputs

    def transform_and_embed_prefill_inputs_device(
        self,
        tokens,
        tt_page_table,
        tt_chunk_page_table,
        tt_chunk_start_idx,
    ):
        tt_tokens = self.embd(tokens)
        tt_tokens = ttnn.unsqueeze_to_4D(tt_tokens)
        return tt_tokens, tt_page_table, tt_chunk_page_table, tt_chunk_start_idx

    def prepare_inputs_prefill(
        self,
        tokens,
        start_pos=0,
        page_table=None,
        chunk_page_table=None,
        chunk_start_idx=None,
        trace_enabled=False,
        last_token_idx=None,
        global_user_id=None,
        batch_size=1,
        user_id=0,
        **kwargs,
    ):
        """
        Inputs are torch tensors or python types. This function returns ttnn
        tensors on device if trace is disabled or on host if trace is enabled.
        TODO: Debate whether this function is responsible for padding
        """

        # We set the device to None if trace is enabled so we keep the tensors on host instead of sending it to the device (None - keeps on host, device - sends to specified device)
        # We will send them to device later (copy_host_to_device)
        device = None if trace_enabled else self.mesh_device

        assert tokens.dim() == 2, "tokens must be a 2D tensor"
        # For batched prefill, tokens come in as [padded_batch, S]
        # Each user's tokens are at their slot index in dimension 0
        # Reshape to [1, 1, 1, padded_batch * S] for embedding
        if batch_size > 1:
            # Tokens are in slot-based format [padded_batch, S_per_user]
            S = tokens.shape[-1]  # Per-user sequence length
            tokens = tokens.reshape(1, 1, 1, -1)  # Flatten to [1, 1, 1, padded_batch * S]
        else:
            tokens = tokens.reshape(1, 1, 1, -1)
            S = tokens.shape[-1]
        tokens = ttnn.from_torch(
            tokens,
            device=device,
            dtype=ttnn.uint32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
        )

        # self.embd expects that tokens are on device ; if trace is enabled, the tensors will be later on device, so we will do these 2 steps when we copy the tokens to the device
        if not trace_enabled:
            tokens_embd = self.embd(tokens)
            tokens_embd = ttnn.unsqueeze_to_4D(tokens_embd)

        # Slice the rot mats to the prefill seqlen
        trace_mat_len = self.rope_setup.cos_matrix_prefill.shape[2]
        mat_len = self.args.max_seq_len
        seq_len = last_token_idx + 1 if last_token_idx is not None else S
        assert mat_len >= seq_len, f"Sequence length {seq_len} exceeds max seq len {mat_len}"

        required_end = start_pos + S
        pad_len = max(0, required_end - mat_len)

        # We set the end_pos to max_seq_len so that we don't create a new tensor for the whole cos_matrix and sin_matrix
        # In case of trace, we will use the whole matrix for all seq_lens supported by trace
        prefill_start_pos = 0 if trace_enabled else start_pos
        slice_end = trace_mat_len if trace_enabled else min(mat_len, required_end)

        cos_slice = self.rope_setup.cos_matrix_prefill[:, :, prefill_start_pos:slice_end, :]
        sin_slice = self.rope_setup.sin_matrix_prefill[:, :, prefill_start_pos:slice_end, :]

        if pad_len > 0:
            # Padding: [(before, after), ...] for each dim; pad at end of 3rd dim (dim=2) by pad_len
            padding = [(0, 0)] * 4
            padding[2] = (0, pad_len)
            cos_slice = ttnn.pad(cos_slice, padding=padding, value=0.0)
            sin_slice = ttnn.pad(sin_slice, padding=padding, value=0.0)

        tt_rot_mats_prefill_global = [cos_slice, sin_slice]

        if hasattr(self, "rope_local_setup"):
            local_trace_mat_len = self.rope_local_setup.cos_matrix_prefill.shape[2]
            local_mat_len = self.args.max_seq_len
            local_required_end = start_pos + S
            local_pad_len = max(0, local_required_end - local_mat_len)
            local_slice_end = local_trace_mat_len if trace_enabled else min(local_mat_len, local_required_end)

            local_cos_slice = self.rope_local_setup.cos_matrix_prefill[:, :, prefill_start_pos:local_slice_end, :]
            local_sin_slice = self.rope_local_setup.sin_matrix_prefill[:, :, prefill_start_pos:local_slice_end, :]

            if local_pad_len > 0:
                # Pad at end of 3rd dim (dim=2) by local_pad_len
                local_padding = [(0, 0)] * 4
                local_padding[2] = (0, local_pad_len)
                local_cos_slice = ttnn.pad(local_cos_slice, padding=local_padding, value=0.0)
                local_sin_slice = ttnn.pad(local_sin_slice, padding=local_padding, value=0.0)

            tt_rot_mats_prefill_local = [local_cos_slice, local_sin_slice]
        else:
            tt_rot_mats_prefill_local = None

        if page_table is not None:
            # For batched prefill, replicate page_table to all devices (same as single-user path)
            # The KV cache fill will loop over users and use batch_idx=user_id for each
            tt_page_table = ttnn.from_torch(
                page_table,
                device=device,
                dtype=ttnn.int32,
                layout=ttnn.ROW_MAJOR_LAYOUT,
                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
            )
        else:
            tt_page_table = None

        if chunk_page_table is not None:
            tt_chunk_page_table = ttnn.from_torch(
                chunk_page_table,
                device=device,
                dtype=ttnn.int32,
                layout=ttnn.ROW_MAJOR_LAYOUT,
                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
            )
        else:
            tt_chunk_page_table = None

        if chunk_start_idx is not None and int(chunk_start_idx) > 0:
            chunk_start_idx_tensor = torch.tensor([chunk_start_idx], dtype=torch.int32)
            tt_chunk_start_idx = ttnn.from_torch(
                chunk_start_idx_tensor,
                device=device,
                dtype=ttnn.int32,
                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
            )
        else:
            tt_chunk_start_idx = None

        return (
            tokens if trace_enabled else tokens_embd,
            tt_rot_mats_prefill_global,
            tt_rot_mats_prefill_local,
            tt_page_table,
            tt_chunk_page_table,
            tt_chunk_start_idx,
        )

    def prepare_inputs_decode(self, *inputs):
        """
        Inputs are torch tensors or python types. This function returns ttnn
        tensors on device.
        Its implementation can take advantage of a few other functions which the
        model must implement.
        """
        host_inputs = self.prepare_decode_inputs_host(*inputs)
        device_inputs = copy_host_to_device(host_inputs, mesh_device=self.mesh_device)  # Helper function
        return device_inputs

    def prepare_decode_inputs_host(self, tokens, current_pos, page_table=None):
        """
        Inputs are torch tensors or python types. Outputs are ttnn tensors on host.
        NOTE: Tokens and current_pos are padded to batch
        """
        B = tokens.shape[0]
        assert current_pos.shape[0] == B, "Batch size mismatch"
        assert (
            B == self.args.max_batch_size
        ), f"Batch size {B} must be equal to max_batch_size {self.args.max_batch_size}"

        # Necessary padding to be full tile sized when on device
        tokens = torch.nn.functional.pad(tokens.view(-1), (0, 32 - len(tokens)), "constant", 0)
        tokens = ttnn.from_torch(
            tokens,
            device=None,
            dtype=ttnn.uint32,
            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
        )
        tokens = ttnn.unsqueeze_to_4D(tokens)

        rot_current_pos = torch.maximum(
            current_pos, torch.tensor(0, dtype=torch.int64)
        )  # Ensure position indices are non-negative
        rope_idxs = self.rope_setup.get_rot_idxs(rot_current_pos, on_host=True)

        current_pos_tt = ttnn.from_torch(
            current_pos,
            device=None,
            dtype=ttnn.int32,
            mesh_mapper=ttnn.ShardTensor2dMesh(
                self.mesh_device,
                dims=(None, 0) if (self.args.is_galaxy and B > 1) else (None, None),
                mesh_shape=self.args.cluster_shape,
            ),
        )

        if page_table is not None:
            page_table = ttnn.from_torch(
                page_table,
                device=None,
                dtype=ttnn.int32,
                mesh_mapper=ttnn.ShardTensor2dMesh(
                    self.mesh_device,
                    dims=(None, -2) if (self.args.is_galaxy and B > 1) else (None, None),
                    mesh_shape=self.args.cluster_shape,
                ),
            )
        return tokens, current_pos_tt, rope_idxs, page_table

    def _transform_decode_inputs_device(
        self,
        tokens,
    ):
        """
        Inputs are ttnn tensors on device. This function applies any on-device
        transformations which should happen before forward decode.
        For example: tilize, reshape, shard.
        Return transformed device tensors

        Embed tokens
        """
        decode_residual_mem_cfg = self.args.get_residual_mem_config(Mode.DECODE, self.prefetcher)
        tt_tokens = self.embd(
            tokens,
            memory_config=ttnn.DRAM_MEMORY_CONFIG if self.prefetcher is None else decode_residual_mem_cfg,
        )
        tt_tokens = ttnn.unsqueeze_to_4D(tt_tokens)
        tt_tokens = ttnn.to_memory_config(tt_tokens, decode_residual_mem_cfg)
        return tt_tokens

    def concat_host_output(self, tt_out, is_log_probs=False):
        """
        Concatenate the output of the devices into a single host tensor.
        """
        torch_out_tensors = [ttnn.to_torch(x) for x in ttnn.get_device_tensors(tt_out)]
        if self.args.is_galaxy:
            row_dim, col_dim = (3, 1)
        else:
            row_dim, col_dim = (1, -1)

        rows, cols = self.args.cluster_shape
        mesh_shape = [torch_out_tensors[i : i + cols] for i in range(0, len(torch_out_tensors), cols)]
        if is_log_probs:
            row_concatenated = []
            for row in mesh_shape:
                row_reshaped = [tensor.reshape(1, 1, -1, 1) for tensor in row]
                row_concatenated.append(torch.cat(row_reshaped, dim=col_dim))
        else:
            row_concatenated = [torch.cat(row, dim=col_dim) for row in mesh_shape]

        return torch.cat(row_concatenated, dim=row_dim)

    def process_output_prefill(self, tt_out, last_token_idx):
        """
        Input is ttnn host tensor of logits. Output is torch logits tensor.
        NOTE: In this model, prefill always uses get_last_token
        """
        assert tt_out.storage_type() == ttnn.StorageType.HOST, "Expected host tensor"
        return self.concat_host_output(tt_out)[0, 0, last_token_idx, : self.vocab_size]

    def process_output_prefill_hidden_states(self, tt_out, last_token_idx):
        """
        Input is ttnn host tensor of hidden states (after norm, before LM head).
        Output is torch hidden states tensor of shape [hidden_size].
        Used for embedding models.
        """
        assert tt_out.storage_type() == ttnn.StorageType.HOST, "Expected host tensor"
        # Extract the last token's hidden state
        # Shape: [batch=1, head=1, seq, hidden_dim] -> [hidden_dim]
        # For hidden states, if they're replicated across devices (not sharded),
        # we should take just the first device's output to avoid incorrect concatenation.
        # If sharded, concat_host_output will properly concatenate them.
        concatenated = self.concat_host_output(tt_out)
        # Check if concatenation resulted in oversized tensor (replicated case)
        # If so, take only the first device's portion (first self.args.dim elements)
        if concatenated.shape[-1] > self.args.dim:
            # Hidden states are replicated, take first device's output
            return concatenated[0, 0, last_token_idx, : self.args.dim]
        else:
            # Hidden states are sharded, concatenation is correct
            return concatenated[0, 0, last_token_idx, :]

    def process_output_decode(self, tt_out, B, S=1, is_tokens=False, is_log_probs=False):
        """
        Input is ttnn host tensor of logits if is_tokens=False, otherwise tokens. Output is the corresponding torch tensor.
        """
        if is_tokens or is_log_probs:
            # Pad to 32 to match the expected batch size for decode operations (tiles are 32x32)
            padded_batch_size = 32
            if not is_log_probs:
                tt_out = ttnn.reshape(tt_out, ttnn.Shape([1, 1, padded_batch_size, 1]))
            return self.concat_host_output(tt_out, is_log_probs)[0, 0, :B, 0]
        if self.args.num_devices > 1:
            tt_out = ttnn.to_torch(ttnn.get_device_tensors(tt_out)[0]).float()
        else:
            tt_out = ttnn.to_torch(tt_out).float()
        tt_out = tt_out[:, :, :B, : self.vocab_size].view(B, S, -1)
        return tt_out

    def ttnn_prefill_forward(
        self,
        x,
        rot_mats_global=None,
        rot_mats_local=None,
        user_id=0,
        page_table=None,
        chunk_page_table=None,
        chunk_start_idx=None,
        get_last_token=-1,
        kv_cache=None,
        batch_size=1,
        page_tables_per_layer=None,
    ):
        """
        This method will take device tensors and any other args to run forward.
        It returns ttnn device tensors.
        """
        if page_tables_per_layer is None:
            # vLLM hybrid bridges (HybridAttentionForCausalLM subclasses) stash
            # the per-layer list on the model handle for the duration of a
            # forward call rather than threading the kwarg through Generator's
            # many ttnn_prefill_forward call sites. Pick it up here when set.
            page_tables_per_layer = getattr(self, "_active_page_tables_per_layer", None)
        page_tables_per_layer = self._page_tables_to_ttnn(page_tables_per_layer)
        return self.forward(
            x,
            current_pos=None,
            rot_mats_global=rot_mats_global,
            rot_mats_local=rot_mats_local,
            user_id=user_id,
            mode=Mode.PREFILL,
            page_table=page_table,
            chunk_page_table=chunk_page_table,
            chunk_start_idx=chunk_start_idx,
            get_last_token=get_last_token,
            kv_cache=kv_cache,
            batch_size=batch_size,
            page_tables_per_layer=page_tables_per_layer,
        )

    def _page_table_mesh_mapper(self, B):
        """Mesh mapper for per-layer page tables, matching the layout that
        :meth:`prepare_decode_inputs_host` uses for the legacy single
        ``page_table`` kwarg: shard the batch dim across mesh axis 1 on
        Galaxy when ``B>1``, replicate otherwise. The hybrid bridge
        chunks the global page table per-DP before calling into a
        submesh, so ``B`` here is the per-DP batch — same value the
        legacy path sees on entry to ``prepare_decode_inputs_host``.
        """
        return ttnn.ShardTensor2dMesh(
            self.mesh_device,
            dims=(None, -2) if (self.args.is_galaxy and B > 1) else (None, None),
            mesh_shape=self.args.cluster_shape,
        )

    def _page_tables_to_ttnn(self, page_tables_per_layer):
        """Resolve a per-layer list of ``torch.Tensor`` page tables to a
        list of *persistent* ttnn device tensors (allocate-only).

        Tracing bakes each input tensor's device address into the captured
        graph; replaying the trace reads from those exact addresses
        regardless of any new ttnn objects created on the Python side.
        Allocating fresh device tensors on every call would therefore
        make traced inference read stale memory at the original
        addresses, so we lazily allocate one persistent device tensor per
        layer on first use and *only* update contents from outside the
        traced ``ttnn_*_forward`` calls (writes are forbidden during trace
        capture). The hybrid bridge calls
        :meth:`update_persistent_per_layer_page_tables` *before* invoking
        ``Generator``'s decode/prefill which executes traces — that's
        where content updates happen.

        First call (warmup compile) populates the persistent buffers from
        the input torch tensors; subsequent calls return the existing
        buffers unchanged. ``None`` entries propagate; already-ttnn
        entries pass through.
        """
        if page_tables_per_layer is None:
            return None
        persistent = getattr(self, "_persistent_per_layer_page_tables", None)
        n = len(page_tables_per_layer)
        if persistent is None or len(persistent) != n:
            persistent = []
            for pt in page_tables_per_layer:
                if pt is None:
                    persistent.append(None)
                    continue
                if isinstance(pt, ttnn.Tensor):
                    persistent.append(pt)
                    continue
                persistent.append(
                    ttnn.from_torch(
                        pt,
                        device=self.mesh_device,
                        dtype=ttnn.int32,
                        layout=ttnn.ROW_MAJOR_LAYOUT,
                        mesh_mapper=self._page_table_mesh_mapper(pt.shape[0]),
                    )
                )
            self._persistent_per_layer_page_tables = persistent
        return persistent

    def update_persistent_per_layer_page_tables(self, page_tables_per_layer):
        """Update content of persistent per-layer page_table device
        tensors in place. Called by the hybrid bridge *before* invoking
        ``Generator``'s decode/prefill so traced replay observes the new
        block IDs at the captured addresses. Must be called outside trace
        capture (writes forbidden inside).

        No-op if persistent tensors haven't been allocated yet (first
        call goes through :meth:`_page_tables_to_ttnn`'s allocation).
        """
        if page_tables_per_layer is None:
            return
        persistent = getattr(self, "_persistent_per_layer_page_tables", None)
        if persistent is None or len(persistent) != len(page_tables_per_layer):
            return
        for i, pt in enumerate(page_tables_per_layer):
            if pt is None or persistent[i] is None or isinstance(pt, ttnn.Tensor):
                continue
            host_pt = ttnn.from_torch(
                pt,
                device=None,
                dtype=ttnn.int32,
                layout=ttnn.ROW_MAJOR_LAYOUT,
                mesh_mapper=self._page_table_mesh_mapper(pt.shape[0]),
            )
            ttnn.copy_host_to_device_tensor(host_pt, persistent[i])

    def _increment_decode_positions_device(self, current_pos, rot_mat_idxs):
        ttnn.plus_one(current_pos, skip_negative_entries=True)
        ttnn.plus_one(rot_mat_idxs)

    def _slice_prefill_rot_mats(self, rot_mats, chunk_start_idx, prefill_seq_len):
        """Slice full prefill RoPE mats to the traced prefill sequence length."""
        if rot_mats is None or chunk_start_idx is None or not isinstance(chunk_start_idx, ttnn.Tensor):
            return rot_mats

        full_rot_cos, full_rot_sin = rot_mats[0], rot_mats[1]
        full_seq_len = full_rot_cos.shape[2]
        if prefill_seq_len <= 0:
            raise ValueError(f"Prefill sequence length must be positive, got {prefill_seq_len}")
        if full_rot_sin.shape[2] != full_seq_len:
            raise ValueError(
                f"Prefill RoPE cosine and sine sequence lengths must match, got "
                f"{full_seq_len} and {full_rot_sin.shape[2]}"
            )
        if full_seq_len == prefill_seq_len:
            return rot_mats
        if full_seq_len % prefill_seq_len != 0:
            raise ValueError(
                f"Full RoPE sequence length {full_seq_len} must be evenly divisible by "
                f"prefill sequence length {prefill_seq_len}"
            )
        # Tensor-bound slice fixes output geometry as input length divided by
        # num_devices; this argument is a partition count, not the mesh width.
        num_partitions = full_seq_len // prefill_seq_len

        z = self._tt_slice_start_zeros_4
        tt_slice_starts = ttnn.concat([z[0:2], chunk_start_idx, z[3:4]], dim=0)

        rot_cos_slice = ttnn.slice(
            input_tensor=full_rot_cos,
            starts=tt_slice_starts,
            ends=self._tt_seq_len_buffer,
            slice_dim=2,
            num_devices=num_partitions,
        )
        rot_sin_slice = ttnn.slice(
            input_tensor=full_rot_sin,
            starts=tt_slice_starts,
            ends=self._tt_seq_len_buffer,
            slice_dim=2,
            num_devices=num_partitions,
        )
        return (rot_cos_slice, rot_sin_slice)

    def ttnn_decode_forward(
        self,
        x,
        current_pos,
        rot_mat_idxs=None,
        page_table=None,
        kv_cache=None,
        on_device_logits=False,
        page_tables_per_layer=None,
    ):
        """
        This method will take device tensors and any other args to run forward.
        It returns ttnn device tensors.
        """
        rot_mats_global = self.rope_setup.get_rot_mats(rot_mat_idxs)
        rot_mats_local = self.rope_local_setup.get_rot_mats(rot_mat_idxs) if hasattr(self, "rope_local_setup") else None

        x_embed = self._transform_decode_inputs_device(x)

        if page_tables_per_layer is None:
            # See ttnn_prefill_forward: hybrid bridges stash the per-layer list
            # on the model when active, since Generator doesn't thread the kwarg.
            page_tables_per_layer = getattr(self, "_active_page_tables_per_layer", None)
        page_tables_per_layer = self._page_tables_to_ttnn(page_tables_per_layer)

        tt_logits = self.forward(
            x_embed,
            current_pos,
            rot_mats_global=rot_mats_global,
            rot_mats_local=rot_mats_local,
            mode=Mode.DECODE,
            page_table=page_table,
            kv_cache=kv_cache,
            page_tables_per_layer=page_tables_per_layer,
        )

        if on_device_logits:
            assert self.sampling is not None, (
                "ttnn_decode_forward got on_device_logits=True but no on-device sampling "
                "module exists (self.sampling is None)."
            )
            self._increment_decode_positions_device(current_pos, rot_mat_idxs)
            return tt_logits

        # Gather the output across all devices and untilize the tensor (for argmax)
        if self.args.num_devices > 1:
            cluster_axis = 0 if self.args.is_galaxy else None
            num_links = 2 if self.args.is_galaxy else 1
            tt_logits = ttnn.experimental.all_gather_async(
                tt_logits,
                persistent_output_buffer=None,
                dim=3,
                multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis),
                num_links=num_links,
                memory_config=tt_logits.memory_config() if self.prefetcher is None else ttnn.DRAM_MEMORY_CONFIG,
                cluster_axis=cluster_axis,
                topology=self.args.ccl_topology(),
                barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis),
                chunks_per_sync=10,
                num_workers_per_link=2,
                num_buffers_per_channel=2,
                subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
            )

        tt_logits = ttnn.untilize(
            tt_logits,
            use_multicore=True,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            sub_core_grids=self.prefetcher.all_worker_cores_range_set if self.prefetcher is not None else None,
        )

        return tt_logits, None

    def switch_mode(self, mode: Mode):
        if self.prefetcher is not None:
            self.prefetcher.init(mode)
            self.prefetcher.prefetch()

    def forward(
        self,
        x: ttnn.Tensor,
        current_pos,
        rot_mats_global=None,
        rot_mats_local=None,
        user_id=0,
        mode: Mode = Mode.DECODE,
        page_table=None,
        chunk_page_table=None,
        chunk_start_idx=None,
        get_last_token=-1,
        kv_cache=None,
        batch_size=1,
        page_tables_per_layer=None,
    ):
        if mode == Mode.DECODE:
            # Run prefetcher if it is enabled
            if self.prefetcher is not None:
                self.prefetcher.run()

        if mode == Mode.PREFILL:
            # For traced prefill, keep RoPE slicing in-graph and driven by the
            # on-device chunk_start_idx input. Batched prefill arrives flattened
            # to [1, 1, batch_size * S_per_user, dim] and each TransformerBlock
            # restores the batch dimension before attention, so the RoPE slice
            # width is the per-user length, not the flattened one.
            prefill_seq_len = x.shape[2] // batch_size
            rot_mats_global = self._slice_prefill_rot_mats(rot_mats_global, chunk_start_idx, prefill_seq_len)
            if rot_mats_local is not None:
                rot_mats_local = self._slice_prefill_rot_mats(rot_mats_local, chunk_start_idx, prefill_seq_len)

        if page_tables_per_layer is not None and len(page_tables_per_layer) != len(self.layers):
            raise ValueError(
                f"page_tables_per_layer has {len(page_tables_per_layer)} entries "
                f"but model has {len(self.layers)} layers"
            )

        for i, layer in enumerate(self.layers):
            # No-op if callers already provide the right memory config
            activation_dtype = self.args.decoders_optimizations.get_tensor_dtype(
                decoder_id=i, tensor=TensorGroup.ACTIVATION
            )

            if mode == Mode.DECODE and not self.args.is_galaxy:
                x = ttnn.to_memory_config(
                    x,
                    self.args.get_residual_mem_config(mode, self.prefetcher),
                    activation_dtype,
                )
            elif activation_dtype is not None and x.dtype != activation_dtype:
                x = ttnn.typecast(x, activation_dtype)

            # vLLM hybrid kv-cache-groups: each attention layer gets its own
            # paged pool (sliding-window vs full-attention have different
            # block counts). When ``page_tables_per_layer`` is None we fall
            # back to broadcasting the single ``page_table`` to every layer
            # — byte-equivalent to the pre-hybrid path used by every legacy
            # caller (demos, unit tests, non-hybrid vLLM bridges).
            layer_page_table = page_tables_per_layer[i] if page_tables_per_layer is not None else page_table

            x = layer(
                x,
                current_pos,
                rot_mats_global=rot_mats_global,
                rot_mats_local=rot_mats_local,
                user_id=user_id,
                mode=mode,
                page_table=layer_page_table,
                chunk_page_table=chunk_page_table,
                chunk_start_idx=chunk_start_idx,
                kv_cache=kv_cache[i] if kv_cache is not None else None,
                batch_size=batch_size,
            )

        if mode == Mode.DECODE:
            if self.prefetcher is not None:
                self.prefetcher.stop()

        if mode == Mode.PREFILL and get_last_token == -1:
            return x

        # Slicing the tensor to the nearest ceiling/floor multiples of 32 for the prefill_len, to get the last token
        if get_last_token != -1:
            seq_len = int(x.shape[2])
            if seq_len % 32 == 0:
                # Runtime bounds, as in process_logits_after_prefill_trace: with literal bounds every
                # distinct prompt offset compiles its own slice program. Untraced-prefill models
                # (Gemma-3) reach this slice on every real request, after warmup has recorded their
                # decode traces, and warmup only ever sees bucket-length mock prompts - measured on
                # Gemma-3-27B DP-4 as the last 8 buffers left live across trace replays.
                for device_tensor, values in (
                    (self._tail_slice_start, [0, 0, get_last_token, 0]),
                    (self._tail_slice_end, [1, 1, get_last_token + 32, int(x.shape[-1])]),
                ):
                    ttnn.copy_host_to_device_tensor(
                        ttnn.from_torch(
                            torch.tensor(values, dtype=torch.int32),
                            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
                        ),
                        device_tensor,
                    )
                x = ttnn.slice(
                    input_tensor=x,
                    starts=self._tail_slice_start,
                    ends=self._tail_slice_end,
                    slice_dim=2,
                    num_devices=seq_len // 32,
                )
            else:
                x = ttnn.slice(x, (0, 0, get_last_token, 0), (1, 1, get_last_token + 32, x.shape[-1]))

        # Output norm
        x = self.norm(x, mode=mode, norm_config=self.args.get_norm_config("lm_head", mode, self.prefetcher))

        lm_head_input_mem_cfg = self.args.get_lm_head_input_mem_config(
            mode, None if mode == Mode.PREFILL else self.prefetcher
        )
        if mode == Mode.PREFILL and lm_head_input_mem_cfg.is_sharded():
            x = ttnn.interleaved_to_sharded(x, lm_head_input_mem_cfg)
        if mode == Mode.DECODE and self.prefetcher is not None:
            x = ttnn.to_memory_config(x, self.args.get_lm_head_input_mem_config(mode, self.prefetcher))

        x = self.lm_head(x)
        x = self._apply_final_logit_softcapping(x)
        if mode == Mode.PREFILL:
            x = ttnn.to_memory_config(x, memory_config=ttnn.DRAM_MEMORY_CONFIG)

        return x