code/models/tt_transformers/tt/attention.py

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

# SPDX-License-Identifier: Apache-2.0

import math

import torch

import ttnn
from models.common.lightweightmodule import LightweightModule
from models.common.rmsnorm import RMSNorm
from models.common.utility_functions import copy_to_buffer, nearest_32
from models.tt_transformers.tt.ccl import tt_all_gather, tt_all_reduce
from models.tt_transformers.tt.common import Mode
from models.tt_transformers.tt.model_config import OpGroup, TensorGroup, num_to_corerange


class Attention(LightweightModule):
    def __init__(
        self,
        mesh_device,
        tt_ccl,
        args,
        state_dict,
        weight_cache_path,
        layer_num,
        dtype,
        transformation_mats,
        configuration,
        paged_attention_config=None,
        use_paged_kv_cache=False,
        prefetcher=None,
    ):
        super().__init__()
        self.args = args
        self.mesh_device = mesh_device
        self.tt_ccl = tt_ccl
        self.num_devices = configuration.num_devices
        self.prefetcher = prefetcher
        self.TG = self.num_devices == 32
        self.hidden_size = configuration.dim
        self.n_heads = configuration.n_heads
        self.head_dim = configuration.head_dim
        self.max_seq_len = configuration.max_seq_len
        self.max_batch_size = configuration.max_batch_size
        self.n_kv_heads = configuration.n_kv_heads
        self.paged_attention_config = paged_attention_config
        self.min_kv_prefill_shard_seqlen = configuration.min_kv_prefill_shard_seqlen
        self.ccl_dtype = configuration.ccl_dtype
        self.MAX_QKV_MM_SEQ_LEN = configuration.MAX_QKV_MM_SEQ_LEN
        self.tile_size = configuration.tile_size
        self.rms_norm_add_unit_offset = configuration.rms_norm_add_unit_offset
        self.num_device_groups = self.num_devices // self.n_kv_heads
        self.num_devices_per_group = self.n_kv_heads if self.TG else self.num_devices
        self.batch_size_per_device_group = (
            max(self.max_batch_size // self.num_device_groups, 1) if self.TG else self.max_batch_size
        )

        self.n_local_heads = self.n_heads // self.num_devices_per_group
        self.n_local_kv_heads = self.n_kv_heads // self.num_devices_per_group

        self.use_qk_fused = configuration.use_qk_fused
        self.use_hf_rope = configuration.use_hf_rope
        self.arch_name = configuration.arch_name
        # TODO: Fix this once all-gather supports < tile_size
        if self.TG:
            weight = torch.zeros(1, 32, 8, 32)
            for i in range(32):
                col = i % 4  # This determines which group of 8 to select
                weight[:, i, :, col * 8 : (col + 1) * 8] = torch.eye(8)

            self.slice_mat = ttnn.from_torch(
                weight,
                dtype=ttnn.bfloat4_b,
                layout=ttnn.TILE_LAYOUT,
                device=self.mesh_device,
                mesh_mapper=ttnn.ShardTensorToMesh(self.mesh_device, dim=1),
            )
            user_selection_matrix = torch.eye(8, 8)
            user_selection_matrix = torch.nn.functional.pad(user_selection_matrix, (0, 24), "constant", 0)  # (8, 32)
            user_selection_matrix = [user_selection_matrix] * 4
            user_selection_matrix = torch.block_diag(*user_selection_matrix)  # (32, 128)
            self.user_selection_matrix = ttnn.from_torch(
                user_selection_matrix,
                dtype=ttnn.bfloat4_b,
                layout=ttnn.TILE_LAYOUT,
                device=self.mesh_device,
                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
            )

        self.dtype = dtype

        self.max_seq_len = configuration.max_seq_len
        self.grid_size = configuration.max_grid_size

        self.compute_kernel_config_hifi2 = configuration.compute_kernel_config_hifi2
        self.compute_kernel_config_hifi2_fp16 = configuration.compute_kernel_config_hifi2_fp16

        self.compute_kernel_config_hifi4 = configuration.compute_kernel_config_hifi4

        self.transformation_mats = transformation_mats
        self.is_sliding = (
            configuration.layer_types[layer_num] == "sliding_attention" if configuration.layer_types else False
        )
        self.sliding_window = configuration.sliding_window if self.is_sliding else None

        self.model_config = configuration.get_model_config()
        self.ccl_topology = configuration.ccl_topology()
        self.is_multichip = configuration.is_multichip

        # When prefetcher is enabled, use consistent dtypes across all layers to avoid
        # race conditions caused by different block sizes
        use_prefetcher = prefetcher is not None

        decoders_optimizations = self.args.decoders_optimizations
        self.activation_dtype = decoders_optimizations.get_tensor_dtype(
            decoder_id=layer_num, tensor=TensorGroup.ACTIVATION, prefetcher=use_prefetcher
        )
        self.wqkv_dtype = decoders_optimizations.get_tensor_dtype(
            decoder_id=layer_num, tensor=TensorGroup.WQKV, prefetcher=use_prefetcher
        )
        self.wo_dtype = decoders_optimizations.get_tensor_dtype(
            decoder_id=layer_num, tensor=TensorGroup.WO, prefetcher=use_prefetcher
        )
        self.kv_cache_dtype = decoders_optimizations.get_tensor_dtype(
            decoder_id=layer_num, tensor=TensorGroup.KV_CACHE, prefetcher=use_prefetcher
        )
        self.li_qkv_decode_compute_kernel_cfg = decoders_optimizations.get_math_fidelity(
            decoder_id=layer_num, op=OpGroup.LI_QKV_DECODE, configuration=configuration
        )
        self.sdpa_decode_compute_kernel_cfg = decoders_optimizations.get_math_fidelity(
            decoder_id=layer_num, op=OpGroup.SDPA_DECODE, configuration=configuration
        )
        # Some architectures (e.g. Gemma-2, head_dim=256) need the flash-decode op's
        # internal default compute config instead of the framework's forced
        # fp32_dest_acc_en/packer_l1_acc, which corrupts the cross-chunk online-softmax
        # reduction for large head_dim. Config-gated flag; False for every other model.
        if getattr(configuration, "sdpa_decode_use_default_compute_config", False):
            self.sdpa_decode_compute_kernel_cfg = None
        self.li_o_decode_compute_kernel_cfg = decoders_optimizations.get_math_fidelity(
            decoder_id=layer_num, op=OpGroup.LI_O_DECODE, configuration=configuration
        )
        self.sdpa_prefill_compute_kernel_cfg = decoders_optimizations.get_math_fidelity(
            decoder_id=layer_num, op=OpGroup.SDPA_PREFILL, configuration=configuration
        )
        self.li_qkv_prefill_compute_kernel_cfg = decoders_optimizations.get_math_fidelity(
            decoder_id=layer_num, op=OpGroup.LI_QKV_PREFILL, configuration=configuration
        )
        self.li_o_prefill_compute_kernel_cfg = decoders_optimizations.get_math_fidelity(
            decoder_id=layer_num, op=OpGroup.LI_O_PREFILL, configuration=configuration
        )

        layer_name = configuration.get_state_dict_prefix(self.__class__.__name__, layer_num)
        if configuration.dummy_weights or (weight_cache_path is None):
            cache_name = lambda _: None
        else:
            cache_name = lambda name: weight_cache_path / (f"{layer_name}.{name}")

        # Select rotary embedding implementation for decode
        if self.use_hf_rope and self.use_qk_fused:
            raise NotImplementedError("Fused QK is not implemented for HF-style rope")
        if self.use_hf_rope:
            self.rotary_embedding_decode = self._hf_rope_decode
        elif self.use_qk_fused:
            self.rotary_embedding_decode = self._mllama_rope_fused_qk_decode
        else:
            self.rotary_embedding_decode = self._mllama_rope_decode

        # Select rotary embedding implementation for prefill
        if self.use_hf_rope:
            self.rotary_embedding_prefill = self._hf_rope_prefill
        else:
            self.rotary_embedding_prefill = self._mllama_rope_prefill

        wq_str = f"{layer_name}.wq"
        wk_str = f"{layer_name}.wk"
        wv_str = f"{layer_name}.wv"
        wo_str = f"{layer_name}.wo"
        q_norm_str = f"{layer_name}.q_norm"
        k_norm_str = f"{layer_name}.k_norm"

        # Initialize bias tensors as None
        self.wqkv_bias_decode = None
        self.wqkv_bias_prefill = None

        # Create combined QKV bias if present in state dict
        if f"{wq_str}.bias" in state_dict:
            bias_num_devices = self.num_devices_per_group if self.TG else configuration.num_devices
            qkv_bias = torch.concat(
                [
                    torch.concat(
                        [
                            torch.chunk(state_dict[f"{wq_str}.bias"], bias_num_devices)[i],
                            torch.chunk(state_dict[f"{wk_str}.bias"], bias_num_devices)[i],
                            torch.chunk(state_dict[f"{wv_str}.bias"], bias_num_devices)[i],
                        ],
                        dim=-1,
                    )
                    for i in range(bias_num_devices)
                ],
                dim=-1,
            )
            # Galaxy QKV matmuls shard the input dimension over mesh columns, so the
            # column all-reduce sums one partial result from every column. The bias is
            # replicated over those columns and added to each partial before the
            # collective; divide it by the number of contributors so the reduced
            # result contains exactly one copy of the model bias.
            if self.TG:
                qkv_bias = qkv_bias / configuration.cluster_shape[1]
            bias_mesh_mapper = (
                ttnn.ShardTensor2dMesh(
                    self.mesh_device,
                    dims=(-1, None),
                    mesh_shape=configuration.cluster_shape,
                )
                if self.TG
                else ttnn.ShardTensorToMesh(self.mesh_device, dim=-1)
            )
            bias_cache_suffix = f"sharded_2d_col_reduce_{configuration.cluster_shape[1]}" if self.TG else "sharded"
            # Prefill can use broadcasting on the bias add so wants a 1d tensor
            self.wqkv_bias_prefill = ttnn.as_tensor(
                qkv_bias,
                device=self.mesh_device,
                mesh_mapper=bias_mesh_mapper,
                dtype=ttnn.bfloat16,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                layout=ttnn.TILE_LAYOUT,
                cache_file_name=cache_name(f"wqkv_bias_prefill_{bias_cache_suffix}"),
            )
            # as_tensor returns (32, dim) which is incorrect, this reshape updates the padded size to the correct size
            self.wqkv_bias_prefill = ttnn.reshape(
                self.wqkv_bias_prefill,
                (1, 1, 1, self.wqkv_bias_prefill.shape[-1]),
                (1, 1, self.wqkv_bias_prefill.shape[-2], self.wqkv_bias_prefill.shape[-1]),
            )

            # Broadcasting does not seem to be supported inside execute_trace so expand to the whole batch size
            # Create a list of bias tensors for each multiple of tile_size up to max_batch_size
            self.wqkv_bias_decode = []
            for batch_size in range(
                configuration.tile_size,
                configuration.tile_padded_batch_rows + configuration.tile_size,
                configuration.tile_size,
            ):
                qkv_bias_decode = qkv_bias.unsqueeze(0).expand(batch_size, -1)
                bias_tensor = ttnn.as_tensor(
                    qkv_bias_decode,
                    device=self.mesh_device,
                    mesh_mapper=bias_mesh_mapper,
                    dtype=ttnn.bfloat16,
                    memory_config=ttnn.DRAM_MEMORY_CONFIG,
                    layout=ttnn.TILE_LAYOUT,
                    cache_file_name=cache_name(f"wqkv_bias_decode_{bias_cache_suffix}_{batch_size}"),
                )
                self.wqkv_bias_decode.append(bias_tensor)

        # when splitting the devices, we need to make sure that the number of heads is divisible by the number of devices
        assert self.n_heads % self.num_devices_per_group == 0
        assert self.n_kv_heads % self.num_devices_per_group == 0
        assert configuration.qkv_size % self.num_devices_per_group == 0
        assert configuration.dim % self.num_devices_per_group == 0

        # wqkv: 4096 x 3072 (2 devices): width-sharded on 12 banks, 3072 over 12 banks.
        wqkv_mem_config = configuration.create_dram_sharded_mem_config(
            configuration.dim, configuration.qkv_size // configuration.num_devices
        )

        qkv_list = []
        for i in range(self.num_devices_per_group):
            # Chunk weights
            wq_selected = torch.chunk(state_dict[f"{wq_str}.weight"], self.num_devices_per_group, dim=0)[i]
            wk_selected = torch.chunk(state_dict[f"{wk_str}.weight"], self.num_devices_per_group, dim=0)[i]
            wv_selected = torch.chunk(state_dict[f"{wv_str}.weight"], self.num_devices_per_group, dim=0)[i]

            # Transpose the selected chunks
            wq = torch.transpose(wq_selected, -2, -1)
            wk = torch.transpose(wk_selected, -2, -1)
            wv = torch.transpose(wv_selected, -2, -1)

            qkv = torch.cat([wq, wk, wv], dim=-1)
            qkv_list.append(qkv)

        qkv_cat = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0)

        self.wqkv = ttnn.as_tensor(
            qkv_cat,
            dtype=self.wqkv_dtype,
            layout=ttnn.TILE_LAYOUT,
            device=self.mesh_device,
            memory_config=ttnn.DRAM_MEMORY_CONFIG if self.TG else wqkv_mem_config,
            mesh_mapper=ttnn.ShardTensor2dMesh(
                self.mesh_device, dims=(3, 2) if self.TG else (2, 3), mesh_shape=configuration.cluster_shape
            ),
            cache_file_name=cache_name("wqkv_sharded_2d"),
        )

        def norm_reshard(x, norm, mode, norm_config):
            """Hack until RMSNorm supports height-sharded output config"""
            if mode == Mode.DECODE:
                mem_cfg = x.memory_config()
                x = ttnn.to_memory_config(x, ttnn.L1_MEMORY_CONFIG, dtype=x.dtype)
            x = norm(x, mode, norm_config=norm_config)
            if mode == Mode.DECODE:
                x = ttnn.to_memory_config(x, mem_cfg, dtype=x.dtype)
            return x

        if f"{q_norm_str}.weight" in state_dict:
            fn_q_norm = RMSNorm(
                device=self.mesh_device,
                dim=self.head_dim,
                eps=configuration.norm_eps,
                state_dict=state_dict,
                state_dict_prefix=None,  # we already prefix q_norm_str
                weight_cache_path=None if configuration.dummy_weights else weight_cache_path,
                weight_dtype=ttnn.bfloat16,
                weight_key=q_norm_str,
                add_unit_offset=self.rms_norm_add_unit_offset,
                is_distributed=False,
                tt_ccl=self.tt_ccl,
            )
            self.q_norm = lambda x, mode, norm_config: norm_reshard(x, fn_q_norm, mode, norm_config)
        else:
            self.q_norm = lambda x, mode, norm_config: x

        if f"{k_norm_str}.weight" in state_dict:
            fn_k_norm = RMSNorm(
                device=self.mesh_device,
                dim=self.head_dim,
                eps=configuration.norm_eps,
                state_dict=state_dict,
                state_dict_prefix=None,  # we already prefix k_norm_str
                weight_cache_path=None if configuration.dummy_weights else weight_cache_path,
                weight_dtype=ttnn.bfloat16,
                weight_key=k_norm_str,
                add_unit_offset=self.rms_norm_add_unit_offset,
                is_distributed=False,
                tt_ccl=self.tt_ccl,
            )
            self.k_norm = lambda x, mode, norm_config: norm_reshard(x, fn_k_norm, mode, norm_config)
        else:
            self.k_norm = lambda x, mode, norm_config: x

        # For ring topology we can use all gather matmul for wo
        self.use_fused_all_gather_matmul = self.args.use_fused_all_gather_matmul
        pt_wo = state_dict[f"{wo_str}.weight"].transpose(-1, -2).unsqueeze(0).unsqueeze(0)

        wo_mem_config = configuration.create_dram_sharded_mem_config(
            (configuration.n_heads * configuration.head_dim) // configuration.num_devices, configuration.dim
        )

        def get_wo_mesh_mapper():
            if self.use_fused_all_gather_matmul or self.TG:
                return ttnn.ShardTensor2dMesh(
                    self.mesh_device,
                    dims=(2, 3),
                    mesh_shape=configuration.cluster_shape,
                )
            return ttnn.ShardTensorToMesh(self.mesh_device, dim=2)

        if self.prefetcher is not None:
            self.wo_sharded_ring = ttnn.as_tensor(
                pt_wo,
                dtype=self.wo_dtype,
                layout=ttnn.TILE_LAYOUT,
                device=self.mesh_device,
                memory_config=self.args.get_sharded_wo_ring_mem_config(),
                mesh_mapper=get_wo_mesh_mapper(),
                cache_file_name=(cache_name("wo_sharded_ring")),
            )

        def get_wo_memory_config():
            if self.use_fused_all_gather_matmul or self.TG:
                return ttnn.DRAM_MEMORY_CONFIG
            else:
                return wo_mem_config

        self.wo = ttnn.as_tensor(
            pt_wo,
            dtype=self.wo_dtype,
            layout=ttnn.TILE_LAYOUT,
            device=self.mesh_device,
            memory_config=get_wo_memory_config(),
            mesh_mapper=get_wo_mesh_mapper(),
            cache_file_name=(
                cache_name("wo_width_sharded_2d") if (self.use_fused_all_gather_matmul or self.TG) else cache_name("wo")
            ),
        )
        if not use_paged_kv_cache:
            # vLLM provides its own kv cache
            self.init_kv_cache(configuration, weight_cache_path)

        if configuration.query_pre_attn_scalar is not None:
            self.scale = configuration.query_pre_attn_scalar**-0.5
        else:
            self.scale = self.head_dim**-0.5

        # Insert the tensors into the prefetcher only in decode mode, we do not use prefetcher in prefill mode
        if self.prefetcher is not None:

            def register_weights():
                self.prefetcher.insert_tensor(self.wqkv)
                self.prefetcher.insert_tensor(self.wo_sharded_ring)

            self.prefetcher.register_callback(register_weights)

    def init_kv_cache(self, configuration, weight_cache_path):
        """
        Generates empty KV cache and pushed to device memory
        """

        if self.paged_attention_config:
            cache_k = torch.zeros(
                (
                    self.paged_attention_config.max_num_blocks,
                    self.n_local_kv_heads,
                    self.paged_attention_config.block_size,
                    self.head_dim,
                )
            )
            cache_v = torch.zeros(
                (
                    self.paged_attention_config.max_num_blocks,
                    self.n_local_kv_heads,
                    self.paged_attention_config.block_size,
                    self.head_dim,
                )
            )
        else:
            cache_k = torch.zeros(
                (
                    self.batch_size_per_device_group,
                    self.n_local_kv_heads,
                    self.max_seq_len,
                    self.head_dim,
                )
            )
            cache_v = torch.zeros(
                (
                    self.batch_size_per_device_group,
                    self.n_local_kv_heads,
                    self.max_seq_len,
                    self.head_dim,
                )
            )

        self.layer_past = [
            ttnn.as_tensor(
                k_or_v,
                dtype=self.kv_cache_dtype,
                layout=self.args.get_attn_weights_layout(),
                device=self.mesh_device,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
                cache_file_name=(
                    f"{weight_cache_path}/kvcache_{k_or_v.shape}"
                    if weight_cache_path and not configuration.dummy_weights
                    else None
                ),
            )
            for k_or_v in [cache_k, cache_v]
        ]

    def _update_wqkv(self, tensor: ttnn.Tensor) -> None:
        """In-place replace ``self.wqkv`` via ``ttnn.copy``. Caller must match
        the constructor's ``self.wqkv``: shape ``(1, 1, H, qkv_size_per_device)``,
        TILE, ``ShardTensor2dMesh(dims=(2, 3))`` (or ``(3, 2)`` on TG).
        """
        copy_to_buffer(tensor, self.wqkv, self.wqkv_dtype)

    def _update_wo(self, tensor: ttnn.Tensor) -> None:
        """In-place replace ``self.wo`` (and ``self.wo_sharded_ring`` when the
        prefetcher is enabled) via ``ttnn.copy``. Both must stay in sync: the
        decode path reads ``wo_sharded_ring`` instead of ``wo`` when
        ``self.prefetcher is not None``.
        """
        copy_to_buffer(tensor, self.wo, self.wo_dtype)
        wo_sharded_ring = getattr(self, "wo_sharded_ring", None)
        if wo_sharded_ring is not None:
            copy_to_buffer(tensor, wo_sharded_ring, self.wo_dtype)

    def _update_wqkv_bias(
        self,
        q_proj_bias: ttnn.Tensor,
        k_proj_bias: ttnn.Tensor,
        v_proj_bias: ttnn.Tensor,
    ) -> None:
        """In-place replace ``self.wqkv_bias_prefill`` and ``self.wqkv_bias_decode``
        from separate HF Q/K/V biases. Not yet implemented (needs the
        constructor's per-device slice+concat and decode batch broadcast).
        """
        raise NotImplementedError("wqkv_bias update is not yet implemented")

    def _update_qk_norm(self, which: str, weight: ttnn.Tensor) -> None:
        """In-place replace the optional Q-/K-RMSNorm gamma. Not yet implemented:
        ``self.q_norm``/``self.k_norm`` are lambdas closing over an internal
        ``RMSNorm`` not exposed on ``self``, so we can't delegate to
        ``RMSNorm.update`` without a small refactor.
        """
        raise NotImplementedError(f"{which} update is not yet implemented")

    def update(
        self,
        *,
        q_proj: ttnn.Tensor,
        k_proj: ttnn.Tensor,
        v_proj: ttnn.Tensor,
        o_proj: ttnn.Tensor,
        q_proj_bias: ttnn.Tensor | None = None,
        k_proj_bias: ttnn.Tensor | None = None,
        v_proj_bias: ttnn.Tensor | None = None,
        q_norm: ttnn.Tensor | None = None,
        k_norm: ttnn.Tensor | None = None,
        hf_rope: bool = False,
    ) -> None:
        """In-place replace the on-device attention weights via ``ttnn.copy``.

        HF-format input contract (see ``LLAMA_WEIGHT_TRANSFER.md``): required
        ``self_attn.{q,k,v,o}_proj.weight`` plus optional biases/QK-norm gammas,
        as HF Linear shapes wrapped in two leading unit dims, bf16, TILE,
        DRAM-interleaved, replicated.

        ``hf_rope`` (default ``False``): whether Q/K rows are in HF split-half
        rotary order or already in this model's convention. ``False`` is right
        for the ttml -> tt-transformers transfer (both store Meta-permuted Q/K
        for Llama-3.2-1B-Instruct). When it disagrees with ``self.use_hf_rope``
        an HF<->Meta row permutation is required; not implemented (raises).

        Conversion (single-device only): transpose each proj (mirrors the
        constructor's ``torch.transpose``), concat Q/K/V along the output dim,
        then ``_update_wqkv``/``_update_wo`` ``ttnn.copy`` into existing buffers,
        preserving buffer addresses (so captured traces and the prefetcher's
        recorded addresses stay valid). The multi-device per-device
        chunk-transpose-concat is not yet implemented.

        Raises ``ValueError`` when optional inputs don't match construction-time
        presence (e.g. biases passed but the module has no QKV bias).
        """
        if hf_rope != self.use_hf_rope:
            raise NotImplementedError(
                "Attention.update: HF<->Meta Q/K row permutation is not yet "
                f"implemented (hf_rope={hf_rope}, self.use_hf_rope={self.use_hf_rope}). "
                "Permute Q/K rows on the caller side, or extend update() with the "
                "on-device row-permutation kernel."
            )
        if self.num_devices_per_group != 1:
            raise NotImplementedError(
                "Attention.update for num_devices_per_group > 1 is not yet "
                "implemented; the multi-device path needs the constructor's "
                "per-device chunk-transpose-concat (see __init__'s qkv_list loop)."
            )

        q_t = ttnn.transpose(q_proj, -2, -1)
        k_t = ttnn.transpose(k_proj, -2, -1)
        v_t = ttnn.transpose(v_proj, -2, -1)
        wqkv_fused = ttnn.concat([q_t, k_t, v_t], dim=-1)
        self._update_wqkv(wqkv_fused)

        wo_internal = ttnn.transpose(o_proj, -2, -1)
        self._update_wo(wo_internal)

        bias_present = (q_proj_bias is not None) or (k_proj_bias is not None) or (v_proj_bias is not None)
        if bias_present:
            if self.wqkv_bias_prefill is None:
                raise ValueError(
                    "Q/K/V projection biases were supplied but Attention was " "constructed without a QKV bias."
                )
            if q_proj_bias is None or k_proj_bias is None or v_proj_bias is None:
                raise ValueError("Q/K/V projection biases must all be supplied together.")
            self._update_wqkv_bias(q_proj_bias, k_proj_bias, v_proj_bias)
        else:
            if self.wqkv_bias_prefill is not None:
                raise ValueError(
                    "Attention was constructed with a QKV bias but Q/K/V " "projection biases were not supplied."
                )

        if q_norm is not None:
            self._update_qk_norm("q_norm", q_norm)
        if k_norm is not None:
            self._update_qk_norm("k_norm", k_norm)

    def to_qk_fused_memory_config(self, q_tensor: ttnn.Tensor, k_tensor: ttnn.Tensor):
        """
        Convert Q and K tensors to height-sharded memory layouts suitable for
        fused QK ops such as rotary_embedding_llama_fused_qk and the subsequent
        QK matmul/attention score computation.

        This function:
        - Infers the number of Q heads and KV heads from the input tensors
        - Shards Q and K along the batch dimension using HEIGHT sharding
        - Places Q and K on disjoint core regions to avoid overlap within sub_core_grids
        - Uses row-major shard orientation with explicit shard shapes

        The resulting memory layouts are compatible with fused attention
        kernels that expect Q and K to be distributed across separate
        core ranges while preserving per-head contiguity.

        Args:
            q_tensor (ttnn.Tensor):
                Query tensor with shape [..., batch, num_q_heads, head_dim].

            k_tensor (ttnn.Tensor):
                Key tensor with shape [..., batch, num_kv_heads, head_dim].

            sub_core_grids (ttnn.CoreRangeSet):
                The available core grids to place Q and K tensors on.

        Returns:
            Tuple[ttnn.Tensor, ttnn.Tensor]:
                (q_tensor, k_tensor) converted to sharded memory configurations.
        """
        n_q_heads = q_tensor.shape[2]
        n_kv_heads = k_tensor.shape[2]
        q_batch = q_tensor.shape[1]
        k_batch = k_tensor.shape[1]
        assert q_batch == k_batch

        row_size = 8  # We assume a row size of 8 cores
        k_start_core = ttnn.CoreCoord(q_batch % row_size, q_batch // row_size)

        q_core_grid = ttnn.CoreRangeSet({num_to_corerange(q_batch)})
        k_core_grid = ttnn.CoreRangeSet({num_to_corerange(k_batch, start_core=k_start_core)})

        q_mem_config = ttnn.create_sharded_memory_config(
            shape=(nearest_32(n_q_heads), self.head_dim),
            core_grid=q_core_grid,
            strategy=ttnn.ShardStrategy.HEIGHT,
            orientation=ttnn.ShardOrientation.ROW_MAJOR,
            use_height_and_width_as_shard_shape=True,
        )
        k_mem_config = ttnn.create_sharded_memory_config(
            shape=(nearest_32(n_kv_heads), self.head_dim),
            core_grid=k_core_grid,
            strategy=ttnn.ShardStrategy.HEIGHT,
            orientation=ttnn.ShardOrientation.ROW_MAJOR,
            use_height_and_width_as_shard_shape=True,
        )
        q_tensor = ttnn.to_memory_config(q_tensor, q_mem_config)
        k_tensor = ttnn.to_memory_config(k_tensor, k_mem_config)
        return q_tensor, k_tensor

    def _mllama_rope_decode(self, q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos):
        # Q Rotary Embeddings
        q_heads_1BQD = ttnn.experimental.rotary_embedding_llama(
            q_heads_pre_rot_1BQD, rot_mats[0], rot_mats[1], self.transformation_mats["decode"], is_decode_mode=True
        )

        # K Rotary Embeddings
        k_heads_1BKD = ttnn.experimental.rotary_embedding_llama(
            k_heads_pre_rot_1BKD, rot_mats[0], rot_mats[1], self.transformation_mats["decode"], is_decode_mode=True
        )
        return q_heads_1BQD, k_heads_1BKD

    def _mllama_rope_fused_qk_decode(self, q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos):
        q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD = self.to_qk_fused_memory_config(
            q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD
        )

        q_heads_1BQD, k_heads_1BKD = ttnn.experimental.rotary_embedding_llama_fused_qk(
            q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats[0], rot_mats[1], self.transformation_mats["decode"]
        )
        return q_heads_1BQD, k_heads_1BKD

    def _hf_rope_decode(self, q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos):
        if q_heads_pre_rot_1BQD.dtype != ttnn.bfloat16:
            q_heads_pre_rot_1BQD = ttnn.typecast(q_heads_pre_rot_1BQD, dtype=ttnn.bfloat16)
        if k_heads_pre_rot_1BKD.dtype != ttnn.bfloat16:
            k_heads_pre_rot_1BKD = ttnn.typecast(k_heads_pre_rot_1BKD, dtype=ttnn.bfloat16)

        q_heads_1BQD = ttnn.experimental.rotary_embedding_hf(
            q_heads_pre_rot_1BQD,
            rot_mats[0],
            rot_mats[1],
            is_decode_mode=True,
        )
        k_heads_1BKD = ttnn.experimental.rotary_embedding_hf(
            k_heads_pre_rot_1BKD,
            rot_mats[0],
            rot_mats[1],
            is_decode_mode=True,
        )
        return q_heads_1BQD, k_heads_1BKD

    def _mllama_rope_prefill(self, q_heads_1QSD_pre_rot, k_heads_1KSD_pre_rot, rot_mats):
        q_heads_1QSD = ttnn.experimental.rotary_embedding_llama(
            q_heads_1QSD_pre_rot,
            rot_mats[0],
            rot_mats[1],
            self.transformation_mats["prefill"],
            is_decode_mode=False,
        )

        k_heads_1KSD = ttnn.experimental.rotary_embedding_llama(
            k_heads_1KSD_pre_rot,
            rot_mats[0],
            rot_mats[1],
            self.transformation_mats["prefill"],
            is_decode_mode=False,
        )

        return q_heads_1QSD, k_heads_1KSD

    def _hf_rope_prefill(self, q_heads_1QSD_pre_rot, k_heads_1KSD_pre_rot, rot_mats):
        if q_heads_1QSD_pre_rot.dtype != ttnn.bfloat16:
            q_heads_1QSD_pre_rot = ttnn.typecast(q_heads_1QSD_pre_rot, dtype=ttnn.bfloat16)

        q_heads_1QSD = ttnn.experimental.rotary_embedding_hf(
            q_heads_1QSD_pre_rot,
            rot_mats[0],
            rot_mats[1],
            is_decode_mode=False,
        )

        if k_heads_1KSD_pre_rot.dtype != ttnn.bfloat16:
            k_heads_1KSD_pre_rot = ttnn.typecast(k_heads_1KSD_pre_rot, dtype=ttnn.bfloat16)

        k_heads_1KSD = ttnn.experimental.rotary_embedding_hf(
            k_heads_1KSD_pre_rot,
            rot_mats[0],
            rot_mats[1],
            is_decode_mode=False,
        )

        return q_heads_1QSD, k_heads_1KSD

    def forward_decode(self, x: ttnn.Tensor, current_pos, rot_mats=None, page_table=None, kv_cache=None) -> ttnn.Tensor:
        """
        x: (seq_len, 1, batch, dim)
        current_pos: (batch_size), current token position in the sequence for each user
        """

        ###
        # QKV matmuls
        # Use HiFi2 for DRAM-sharded matmuls as they are otherwise flop-bound. Loses 1 bit of activation precision.
        ###
        if self.TG and self.prefetcher is None:
            x_interleaved = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG)
            xqkv_fused_sharded = ttnn.linear(
                x_interleaved,
                self.wqkv,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                program_config=None,
                compute_kernel_config=self.li_qkv_decode_compute_kernel_cfg,
                dtype=self.ccl_dtype,
            )
            ttnn.deallocate(x_interleaved)
        else:
            xqkv_fused_sharded = ttnn.linear(
                x,
                self.wqkv,
                memory_config=self.args.get_attn_qkv_mm_mem_config(Mode.DECODE, self.prefetcher),
                program_config=self.args.get_attn_qkv_program_config(Mode.DECODE, 1, self.prefetcher),
                compute_kernel_config=self.li_qkv_decode_compute_kernel_cfg,
                dtype=self.ccl_dtype if self.TG else self.activation_dtype or ttnn.bfloat16,
                global_cb=self.prefetcher.global_cb if self.prefetcher is not None else None,
                sub_device_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
            )
        # FIXME: File bug against dram-sharded matmuls with bias
        if self.wqkv_bias_decode:
            # select the bias tensor based on the number of tiles in the rows
            # WARNING: must not change the batch size between compiling and executing a trace
            num_tiles = int(math.ceil(xqkv_fused_sharded.shape[-2] / self.tile_size))
            xqkv_fused_sharded = xqkv_fused_sharded + self.wqkv_bias_decode[num_tiles - 1]

        ttnn.deallocate(x)
        qkv_all_reduce_mem_cfg = self.args.get_attn_qkv_all_reduce_output_mem_config(
            Mode.DECODE, list(self.mesh_device.shape)[1], self.prefetcher
        )
        xqkv_fused = tt_all_reduce(
            xqkv_fused_sharded,
            self.mesh_device,
            self.tt_ccl,
            cluster_axis=1,
            memory_config=qkv_all_reduce_mem_cfg
            if qkv_all_reduce_mem_cfg is not None
            else xqkv_fused_sharded.memory_config(),
            sharded=not (self.TG and self.prefetcher is None),
            dtype=self.ccl_dtype,
            topology=self.ccl_topology,
            subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
        )
        if self.TG:
            # TODO: Slice the fused_query_key_value tensor get batch=8
            xqkv_fused = ttnn.matmul(
                self.slice_mat,
                xqkv_fused,
                dtype=ttnn.bfloat16,
                memory_config=self.args.get_attn_create_head_input_mem_config(Mode.DECODE),
            )
        else:
            # bfloat16 is required by nlp_create_qkv_heads_decode
            if self.prefetcher is None:
                xqkv_fused = ttnn.sharded_to_interleaved(xqkv_fused_sharded, ttnn.L1_MEMORY_CONFIG, ttnn.bfloat16)
                ttnn.deallocate(xqkv_fused_sharded)
            else:
                xqkv_fused = xqkv_fused_sharded
        # Reshape such that true unpadded batch is tracked in shape
        fqkv_shape = xqkv_fused.shape
        xqkv_fused = ttnn.reshape(
            xqkv_fused, (1, 1, self.batch_size_per_device_group, fqkv_shape[3]), (1, 1, 32, fqkv_shape[3])
        )

        ###
        # Reshape and rotary embeddings
        ###
        (
            q_heads_pre_rot_1BQD,
            k_heads_pre_rot_1BKD,
            v_heads_1BKD,
        ) = ttnn.experimental.nlp_create_qkv_heads_decode(
            xqkv_fused,
            num_heads=self.n_local_heads,
            num_kv_heads=self.n_local_kv_heads,
            memory_config=self.args.get_attn_create_head_output_mem_config(Mode.DECODE, self.prefetcher),
        )
        norm_config = self.args.get_norm_config("attn", Mode.DECODE, None)
        q_heads_pre_rot_1BQD = self.q_norm(q_heads_pre_rot_1BQD, mode=Mode.DECODE, norm_config=norm_config)
        k_heads_pre_rot_1BKD = self.k_norm(k_heads_pre_rot_1BKD, mode=Mode.DECODE, norm_config=norm_config)
        ttnn.deallocate(xqkv_fused)

        # Q, K Rotary Embeddings
        q_heads_1BQD, k_heads_1BKD = self.rotary_embedding_decode(
            q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos
        )

        ttnn.deallocate(q_heads_pre_rot_1BQD)
        ttnn.deallocate(k_heads_pre_rot_1BKD)
        ###
        # KV update
        ###
        if kv_cache:
            keys = kv_cache[0]
            values = kv_cache[1]
        else:
            keys = self.layer_past[0]
            values = self.layer_past[1]

        # k_heads, [seqlen, n_kv_heads, bsz, head_dim]
        # v_heads [seqlen, n_kv_heads, bsz, head_dim]
        # keys, [max_batch_size, n_kv_heads // configuration.num_devices, max_seq_len, head_dim]

        if self.use_qk_fused:
            ttnn.experimental.paged_fused_update_cache(
                keys, k_heads_1BKD, values, v_heads_1BKD, update_idxs_tensor=current_pos, page_table=page_table
            )
        else:
            ttnn.experimental.paged_update_cache(
                keys, k_heads_1BKD, update_idxs_tensor=current_pos, page_table=page_table
            )
            ttnn.experimental.paged_update_cache(
                values, v_heads_1BKD, update_idxs_tensor=current_pos, page_table=page_table
            )
        ttnn.deallocate(k_heads_1BKD)
        ttnn.deallocate(v_heads_1BKD)
        # NOTE: Varying the batch size will result in slightly different outputs.
        # For example, a prompt w/ 1 user vs, the same prompt repeated N times for N users, will produce different outputs
        # This is because the SDPA op in decode mode has different number of reductions depending on batch size
        # Which leads to slightly different outputs from attention (due to accumulated errors)
        sdpa_decode_prog_cfg = self.args.get_attn_sdpa_decode_program_config(self.prefetcher)
        if page_table is not None:
            attn_output_1G4D = ttnn.transformer.paged_scaled_dot_product_attention_decode(
                q_heads_1BQD,
                keys,
                values,
                page_table_tensor=page_table,
                cur_pos_tensor=current_pos,
                scale=self.scale,
                sliding_window_size=self.sliding_window,
                program_config=sdpa_decode_prog_cfg,
                compute_kernel_config=self.sdpa_decode_compute_kernel_cfg,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
            )
        else:
            attn_output_1G4D = ttnn.transformer.scaled_dot_product_attention_decode(
                q_heads_1BQD,
                keys,
                values,
                cur_pos_tensor=current_pos,
                scale=self.scale,
                sliding_window_size=self.sliding_window,
                program_config=sdpa_decode_prog_cfg,
                compute_kernel_config=self.sdpa_decode_compute_kernel_cfg,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,  # FIXME: why not L1 height sharded e.g. SCORES_BATCHED_MM_OUTPUT_MEMCFG?
            )

        ttnn.deallocate(q_heads_1BQD)
        attn_output_11BH = ttnn.to_memory_config(
            attn_output_1G4D,
            memory_config=self.args.get_attn_sdpa_output_mem_config(
                Mode.DECODE, self.batch_size_per_device_group, self.prefetcher
            ),
        )

        attn_output_cat = ttnn.experimental.nlp_concat_heads_decode(
            attn_output_11BH,
            num_heads=self.n_local_heads,
            sub_core_grids=self.prefetcher.all_worker_cores_range_set if self.prefetcher is not None else None,
        )
        ttnn.deallocate(attn_output_11BH)
        ttnn.deallocate(attn_output_1G4D)

        if self.use_fused_all_gather_matmul or self.prefetcher is not None:
            attn_output_cat = ttnn.to_memory_config(
                attn_output_cat,
                self.args.get_attn_concat_heads_output_mem_config(Mode.DECODE, self.prefetcher),
            )

            # Fused AGMM only valid for ring topology
            if self.ccl_topology == ttnn.Topology.Ring and self.prefetcher is None:
                _, dense_out_sharded = ttnn.experimental.all_gather_matmul_async(
                    attn_output_cat,
                    self.wo,
                    persistent_output_buffer=None,
                    dim=3,
                    multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(),
                    all_gather_core_grid_offset=(0, 4),
                    barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(),
                    num_links=self.model_config["ATTN_AGMM_CONFIG"]["num_links"],
                    memory_config_ag=self.args.get_attn_all_gather_output_mem_config(Mode.DECODE, None),
                    memory_config_mm=self.args.get_attn_dense_output_mem_config(Mode.DECODE, None),
                    program_config=self.args.get_attn_all_gather_matmul_program_config(Mode.DECODE, None),
                    compute_kernel_config=self.compute_kernel_config_hifi2,
                    chunks_per_sync=self.model_config["ATTN_AGMM_CONFIG"]["chunks_per_sync"],
                    num_workers_per_link=self.model_config["ATTN_AGMM_CONFIG"]["num_workers_per_link"],
                    num_buffers_per_channel=2,
                    subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
                )
            else:
                all_gather_output = ttnn.experimental.all_gather_async(
                    attn_output_cat,
                    persistent_output_buffer=None,
                    dim=3,
                    multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(),
                    num_links=1,
                    topology=self.ccl_topology,
                    memory_config=self.args.get_attn_all_gather_output_mem_config(Mode.DECODE, self.prefetcher),
                    barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(),
                    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,
                )
                dense_out_sharded = ttnn.linear(
                    all_gather_output,
                    self.wo_sharded_ring if self.prefetcher is not None else self.wo,
                    memory_config=self.args.get_attn_dense_output_mem_config(Mode.DECODE, self.prefetcher),
                    program_config=self.args.get_attn_all_gather_matmul_program_config(Mode.DECODE, self.prefetcher),
                    compute_kernel_config=self.li_o_decode_compute_kernel_cfg,
                    global_cb=self.prefetcher.global_cb if self.prefetcher is not None else None,
                    sub_device_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
                )
                ttnn.deallocate(all_gather_output)
            ttnn.deallocate(attn_output_cat)
            dense_out_sharded = ttnn.to_memory_config(
                dense_out_sharded,
                self.args.get_attn_dense_output_mem_config(Mode.DECODE, self.prefetcher),
            )
            return dense_out_sharded

        else:
            attn_output = tt_all_gather(
                attn_output_cat,
                self.mesh_device,
                self.tt_ccl,
                dim=2,
                cluster_axis=1,
                memory_config=self.args.get_attn_gather_users_mem_config(
                    Mode.DECODE, list(self.mesh_device.shape)[1], self.prefetcher
                ),
                sharded=True,
                subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
                # dtype=self.ccl_dtype,  # Running bf16 until we have SDPA output bfp8 df; otherwise we have two sharded to interleaved/interleaved to sharded conversions
            )
            if self.TG:
                attn_output = ttnn.to_memory_config(attn_output, ttnn.L1_MEMORY_CONFIG)
                # user_selection_matrix = [1, 1, 32, 128]
                # user_selection_matrix @ activation -> [1, 1, 32, 128] * [1, 1, 128, 2048] -> [1, 1, 32, 2048]
                attn_output = ttnn.matmul(
                    self.user_selection_matrix,
                    attn_output,
                    core_grid=ttnn.CoreGrid(y=4, x=8),
                    dtype=ttnn.bfloat16,
                    memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG,
                )

            # TODO: Fix this once self.TG supports dram-sharded matmuls
            dense_out_sharded = ttnn.linear(
                attn_output,
                self.wo,
                core_grid=ttnn.CoreGrid(y=4, x=8) if self.TG else None,
                program_config=self.args.get_attn_wo_program_config(Mode.DECODE, 1, self.prefetcher),
                memory_config=self.args.get_attn_wo_output_mem_config(Mode.DECODE, self.prefetcher),
                dtype=ttnn.bfloat8_b if self.TG else None,
                compute_kernel_config=self.li_o_decode_compute_kernel_cfg,
                global_cb=self.prefetcher.global_cb if self.prefetcher is not None else None,
                sub_device_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
            )

            ttnn.deallocate(attn_output_cat)

            # All reduce
            dense_out_reduced = tt_all_reduce(
                dense_out_sharded,
                self.mesh_device,
                self.tt_ccl,
                cluster_axis=0,
                dim=0 if (self.TG and self.hidden_size < 8192) else 3,
                topology=self.ccl_topology,
                memory_config=self.args.get_attn_all_reduce_output_mem_config(
                    Mode.DECODE, self.hidden_size, list(self.mesh_device.shape)[0], self.prefetcher
                ),
                sharded=True,
                dtype=self.ccl_dtype,
                use_composite=True if self.hidden_size == 8192 else False,
                subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None,
            )

            if not self.TG:
                dense_out_reduced = ttnn.to_memory_config(
                    dense_out_reduced, self.args.get_attn_dense_output_mem_config(Mode.DECODE, None)
                )

            return dense_out_reduced

    def forward_prefill(
        self,
        x_11SH,
        rot_mats,
        user_id: int = 0,
        page_table=None,
        chunk_page_table=None,
        chunk_start_idx=None,
        kv_cache=None,
    ):
        # For batched prefill, x_11SH has shape [B, 1, S, H] where B is batch_size
        # concat before QKV matmul, then reshape back to batch after
        batch_size = x_11SH.shape[0]
        if batch_size > 1:
            # Concatenate batch dimension into sequence for matmul compatibility
            x_11SH = ttnn.reshape(x_11SH, [1, 1, x_11SH.shape[-2] * x_11SH.shape[-3] * x_11SH.shape[-4], -1])

        seq_len = x_11SH.shape[-2]
        original_seq_len = seq_len  # Track original for later unpadding
        assert seq_len % 128 == 0 and seq_len > 0, "Seqlen must be divisible by 128"
        ###
        # QKV matmuls
        ###

        # reshaping long sequence to matmul fit on device
        # Pad seq_len to nearest multiple of MAX_QKV_MM_SEQ_LEN if needed
        if seq_len > self.MAX_QKV_MM_SEQ_LEN and seq_len % self.MAX_QKV_MM_SEQ_LEN != 0:
            padded_seq_len = (
                (seq_len + self.MAX_QKV_MM_SEQ_LEN - 1) // self.MAX_QKV_MM_SEQ_LEN
            ) * self.MAX_QKV_MM_SEQ_LEN
            pad_len = padded_seq_len - seq_len
            x_11SH = ttnn.pad(x_11SH, padding=[(0, 0), (0, 0), (0, pad_len), (0, 0)], value=0.0)
            seq_len = padded_seq_len

        if seq_len > self.MAX_QKV_MM_SEQ_LEN:
            x_11SH = ttnn.reshape(x_11SH, [1, seq_len // self.MAX_QKV_MM_SEQ_LEN, self.MAX_QKV_MM_SEQ_LEN, -1])

        if self.args.use_minimal_qkv_prefill_matmul(seq_len):
            xqkv_fused = ttnn.experimental.minimal_matmul(
                x_11SH,
                self.wqkv,
                compute_kernel_config=self.li_qkv_prefill_compute_kernel_cfg,
                config=self.args.get_attn_qkv_program_config(Mode.PREFILL, seq_len, None),
            )
        else:
            xqkv_fused = ttnn.linear(
                x_11SH,
                self.wqkv,
                dtype=self.ccl_dtype if self.TG else self.activation_dtype or ttnn.bfloat16,
                memory_config=self.args.get_attn_qkv_mm_mem_config(Mode.PREFILL, None),
                compute_kernel_config=self.li_qkv_prefill_compute_kernel_cfg,
                program_config=self.args.get_attn_qkv_program_config(Mode.PREFILL, seq_len, None),
            )

        # FIXME: surely ttnn.linear bias should work?
        if self.wqkv_bias_prefill is not None:
            xqkv_fused = xqkv_fused + self.wqkv_bias_prefill

        xqkv_fused = tt_all_reduce(
            xqkv_fused,
            self.mesh_device,
            self.tt_ccl,
            cluster_axis=1,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            dtype=self.ccl_dtype,
        )

        if seq_len > self.MAX_QKV_MM_SEQ_LEN:
            xqkv_fused = ttnn.reshape(xqkv_fused, [1, 1, seq_len, -1])

        # Slice back to original seq_len if we padded earlier
        if original_seq_len != seq_len:
            xqkv_fused = xqkv_fused[:, :, :original_seq_len, :]
            seq_len = original_seq_len

        if batch_size > 1:
            xqkv_fused = ttnn.reshape(xqkv_fused, [batch_size, 1, seq_len // batch_size, -1])

        ttnn.deallocate(x_11SH)

        # split qkv into heads
        (
            q_heads_1QSD_pre_rot,
            k_heads_1KSD_pre_rot,
            v_heads_1VSD,
        ) = ttnn.experimental.nlp_create_qkv_heads(
            xqkv_fused,
            num_heads=self.n_local_heads,
            num_kv_heads=self.n_local_kv_heads,
            transpose_k_heads=False,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )

        norm_config = self.args.get_norm_config("attn", Mode.PREFILL, None)
        q_heads_1QSD_pre_rot = self.q_norm(q_heads_1QSD_pre_rot, mode=Mode.PREFILL, norm_config=norm_config)
        k_heads_1KSD_pre_rot = self.k_norm(k_heads_1KSD_pre_rot, mode=Mode.PREFILL, norm_config=norm_config)

        ttnn.deallocate(xqkv_fused)

        ###
        # Rotary embeddings
        ###

        # Apply rotary embeddings using the selected implementation
        q_heads_1QSD, k_heads_1KSD = self.rotary_embedding_prefill(q_heads_1QSD_pre_rot, k_heads_1KSD_pre_rot, rot_mats)
        ttnn.deallocate(q_heads_1QSD_pre_rot)
        ttnn.deallocate(k_heads_1KSD_pre_rot)

        # Fill KV-Cache
        if kv_cache:
            keys_BKSD, values_BKSD = kv_cache[0], kv_cache[1]
        else:
            keys_BKSD, values_BKSD = self.layer_past[0], self.layer_past[1]

        k_heads_1KSD_8b = ttnn.typecast(k_heads_1KSD, dtype=keys_BKSD.dtype)
        ttnn.deallocate(k_heads_1KSD)

        # sharding k_fill to deal with update_cache memory limitation
        if seq_len >= self.min_kv_prefill_shard_seqlen and not self.TG and page_table is None:
            k_fill = ttnn.interleaved_to_sharded(k_heads_1KSD_8b, self.args.get_attn_kv_prefill_mem_config(seq_len))
        else:
            k_fill = k_heads_1KSD_8b

        v_heads_1VSD_8b = ttnn.typecast(v_heads_1VSD, dtype=values_BKSD.dtype)

        ttnn.deallocate(v_heads_1VSD)

        # sharding v_fill to deal with update_cache memory limitation
        if seq_len >= self.min_kv_prefill_shard_seqlen and not self.TG and page_table is None:
            v_fill = ttnn.interleaved_to_sharded(v_heads_1VSD_8b, self.args.get_attn_kv_prefill_mem_config(seq_len))
        else:
            v_fill = v_heads_1VSD_8b

        if self.TG:
            k_fill = self.prefill_prepare_tensor_for_kv_cache(k_fill, user_id)
            v_fill = self.prefill_prepare_tensor_for_kv_cache(v_fill, user_id)
        if page_table is not None:
            # In the case that the tokens have been padded along the seq len dimension, we need to fill the cache with the unpadded k/v values.
            # Assume that the page table does not have padding, so we can use it to get the unpadded page len.
            block_size = keys_BKSD.shape[2]
            # If chunked prefill, use chunk_page_table if given, otherwise use page_table.
            fill_page_table = chunk_page_table if chunk_page_table is not None else page_table

        if batch_size > 1:
            # For batched prefill, fill every slot of the padded batch separately.
            # k_fill/v_fill have shape [padded_batch, n_kv_heads, seq_len_per_user, head_dim]
            # The paged_fill_cache kernel reads batch_idx_ptr[0] for all positions,
            # so we must call it once per slot with that slot's K/V slice.
            #
            # IMPORTANT: this loop is captured into the prefill trace, whose identity is the
            # padded batch width, not the set of active slots. Looping over the active slots
            # (user_id) would freeze that slot set into the trace, and a later replay for a
            # different slot set would never write K/V for the new slots. Empty slots have
            # page_table rows of -1, which the kernel treats as "skip" without any write, so
            # filling all padded_batch rows is both safe and slot-agnostic.
            seq_len_per_user = k_fill.shape[2]
            page_len = fill_page_table.shape[1] * block_size

            for slot_idx in range(batch_size):
                # Extract this slot's K/V slice: [1, n_kv_heads, seq_len_per_user, head_dim]
                k_user = k_fill[slot_idx : slot_idx + 1, :, :, :]
                v_user = v_fill[slot_idx : slot_idx + 1, :, :, :]

                # Slice to page length if needed (same as single-user path)
                k_user_sliced = k_user[:, :, :page_len, :] if page_len < seq_len_per_user else k_user
                v_user_sliced = v_user[:, :, :page_len, :] if page_len < seq_len_per_user else v_user

                # Fill cache for this specific slot with scalar batch_idx
                ttnn.experimental.paged_fill_cache(keys_BKSD, k_user_sliced, fill_page_table, batch_idx=slot_idx)
                ttnn.experimental.paged_fill_cache(values_BKSD, v_user_sliced, fill_page_table, batch_idx=slot_idx)
        elif page_table is not None:
            # Single user path with page_table
            page_len = fill_page_table.shape[1] * block_size
            k_fill_sliced = k_fill[:, :, :page_len, :] if page_len < k_fill.shape[2] else k_fill
            v_fill_sliced = v_fill[:, :, :page_len, :] if page_len < v_fill.shape[2] else v_fill
            ttnn.experimental.paged_fill_cache(keys_BKSD, k_fill_sliced, fill_page_table, batch_idx=user_id)
            ttnn.experimental.paged_fill_cache(values_BKSD, v_fill_sliced, fill_page_table, batch_idx=user_id)
        else:
            # Single user path without page_table
            ttnn.fill_cache(
                keys_BKSD,
                k_fill,
                user_id % self.batch_size_per_device_group,
            )
            ttnn.fill_cache(
                values_BKSD,
                v_fill,
                user_id % self.batch_size_per_device_group,
            )
        if seq_len >= self.min_kv_prefill_shard_seqlen and not self.TG and page_table is None:
            ttnn.deallocate(k_fill)
            ttnn.deallocate(v_fill)

        # SDPA
        q_heads_1QSD_8b = ttnn.typecast(q_heads_1QSD, dtype=self.activation_dtype or ttnn.bfloat8_b)
        ttnn.deallocate(q_heads_1QSD)

        if chunk_start_idx is not None:
            if self.sliding_window is not None:
                raise NotImplementedError("Sliding window not supported for chunked prefill SDPA")
            if isinstance(chunk_start_idx, ttnn.Tensor):
                attn_output_84SD = ttnn.transformer.chunked_scaled_dot_product_attention(
                    input_tensor_q=q_heads_1QSD_8b,
                    input_tensor_k=keys_BKSD,
                    input_tensor_v=values_BKSD,
                    page_table_tensor=page_table,
                    chunk_start_idx=None,
                    chunk_start_idx_tensor=chunk_start_idx,
                    compute_kernel_config=self.sdpa_prefill_compute_kernel_cfg,
                    program_config=self.args.get_attn_sdpa_program_config(Mode.PREFILL, seq_len, 0, None),
                )
            else:
                attn_output_84SD = ttnn.transformer.chunked_scaled_dot_product_attention(
                    input_tensor_q=q_heads_1QSD_8b,
                    input_tensor_k=keys_BKSD,
                    input_tensor_v=values_BKSD,
                    page_table_tensor=page_table,
                    chunk_start_idx=chunk_start_idx,
                    compute_kernel_config=self.sdpa_prefill_compute_kernel_cfg,
                    program_config=self.args.get_attn_sdpa_program_config(Mode.PREFILL, seq_len, chunk_start_idx, None),
                )
        else:
            # For batched prefill, the actual per-user seq_len is seq_len // batch_size
            # since the tensors have shape [batch_size, n_heads, seq_len_per_user, head_dim]
            sdpa_seq_len = seq_len // batch_size if batch_size > 1 else seq_len
            attn_output_84SD = ttnn.transformer.scaled_dot_product_attention(
                q_heads_1QSD_8b,
                k_heads_1KSD_8b,
                v_heads_1VSD_8b,
                is_causal=True,
                sliding_window_size=self.sliding_window,
                scale=self.scale,
                compute_kernel_config=self.sdpa_prefill_compute_kernel_cfg,
                program_config=self.args.get_attn_sdpa_program_config(Mode.PREFILL, sdpa_seq_len, None, None),
            )

        # deallocate keys and values
        ttnn.deallocate(q_heads_1QSD_8b)
        ttnn.deallocate(k_heads_1KSD_8b)
        ttnn.deallocate(v_heads_1VSD_8b)

        # For single-user prefill, reshape to expected format for nlp_concat_heads
        # For batched prefill (batch_size > 1), skip this reshape - nlp_concat_heads handles [B, H, S, D]
        # IMPORTANT: Reshaping [B, H, S, D] to [1, H, B*S, D] BEFORE concat_heads would scramble data
        # because batch and sequence dimensions are separated by heads. Must reshape AFTER concat_heads.
        if batch_size == 1:
            attn_output_1QSD = ttnn.reshape(attn_output_84SD, [1, self.n_local_heads, -1, self.head_dim])
        else:
            attn_output_1QSD = attn_output_84SD

        ###
        # Output matmul
        ###
        attn_output_11SH = ttnn.experimental.nlp_concat_heads(
            attn_output_1QSD,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )
        ttnn.deallocate(attn_output_1QSD)

        # For batched prefill, reshape to concatenate batch dimension into sequence
        # This MUST happen AFTER nlp_concat_heads to preserve correct data layout
        # nlp_concat_heads outputs [B, 1, S_per_user, H*D], reshape to [1, 1, B*S, H*D]
        if batch_size > 1:
            attn_output_11SH = ttnn.reshape(attn_output_11SH, [1, 1, seq_len, -1])

        # reshaping long sequence to matmul fit on device
        if seq_len > 1024:
            attn_output_11SH = ttnn.reshape(attn_output_11SH, [1, seq_len // 1024, 1024, -1])

        # Non fused All Gather Matmul
        if self.use_fused_all_gather_matmul:  # is true for Ring topology
            attn_output_11SH = ttnn.experimental.all_gather_async(
                attn_output_11SH,
                persistent_output_buffer=None,
                dim=3,
                multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(),
                num_links=1,
                topology=self.ccl_topology,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(),
                chunks_per_sync=10,
                num_workers_per_link=2,
                num_buffers_per_channel=2,
            )

        output_11SH = ttnn.linear(
            attn_output_11SH,
            self.wo,
            compute_kernel_config=self.li_o_prefill_compute_kernel_cfg,
            dtype=self.activation_dtype or ttnn.bfloat8_b,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
            program_config=self.args.get_attn_wo_program_config(Mode.PREFILL, seq_len, None),
        )

        if seq_len > 1024:
            output_11SH = ttnn.reshape(output_11SH, [1, 1, seq_len, -1])
        ttnn.deallocate(attn_output_11SH)

        # Reduce-scatter
        if not self.use_fused_all_gather_matmul:
            output_11SH = tt_all_reduce(
                output_11SH,
                self.mesh_device,
                self.tt_ccl,
                cluster_axis=0,
                dim=0 if self.TG else 3,
                topology=self.ccl_topology,
                memory_config=ttnn.DRAM_MEMORY_CONFIG,
                dtype=self.ccl_dtype,
            )

        return output_11SH

    def forward(
        self,
        x,
        current_pos,
        rot_mats=None,
        user_id=0,
        mode=Mode.DECODE,
        page_table=None,
        chunk_page_table=None,
        chunk_start_idx=None,
        kv_cache=None,
    ):
        if mode == Mode.PREFILL:
            return self.forward_prefill(
                x,
                rot_mats,
                user_id,
                page_table=page_table,
                chunk_page_table=chunk_page_table,
                chunk_start_idx=chunk_start_idx,
                kv_cache=kv_cache,
            )
        else:
            return self.forward_decode(x, current_pos, rot_mats, page_table=page_table, kv_cache=kv_cache)

    def prefill_prepare_tensor_for_kv_cache(self, key_or_value_layer, user_id):
        tensor_copy = ttnn.clone(key_or_value_layer)
        # key_or_value_layer.deallocate(True)
        # Get all tensors from multi-device tensor
        tensors = ttnn.get_device_tensors(tensor_copy)
        # Get only tensors from specific column chips
        # Get every 4th tensor starting from user_id // 8
        single_column_tensors = tensors[user_id // self.batch_size_per_device_group :: 4]
        # Create multi-device tensor
        multi_device_tensor = ttnn.combine_device_tensors(tensors=single_column_tensors)

        return multi_device_tensor