code/models/tt_dit/layers/linear.py

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

# SPDX-License-Identifier: Apache-2.0

import math

import torch

import ttnn
from models.common.utility_functions import is_blackhole

from ..utils.matmul import get_fabric_agmm_config, get_fused_mmrs_config, get_matmul_config, get_matmul_core_grid
from ..utils.tensor import prepare_for_fused_swiglu
from .module import Module, Parameter

# Fidelity per weight dtype. Quantized dtypes (e.g. bfloat8_b) are deliberately absent: callers
# look up with .get(dtype, HiFi2). This compute_config is a dead default for any op handed an
# explicit compute_kernel_config (every LTX matmul is), so the fallback only has to construct.
MATH_FIDELITY = {
    ttnn.bfloat16: ttnn.MathFidelity.HiFi2,
    ttnn.float32: ttnn.MathFidelity.HiFi4,
}

# Activation strings accepted by Linear / ColParallelLinear `activation_fn`,
# mapped to the values the matmul fused-activation path expects. Each value is
# either a bare ttnn.UnaryOpType (no parameter) or a (UnaryOpType, param0)
# tuple; nanobind's implicit caster handles both forms.
#
# "gelu":      exact GELU (piecewise CDF / FP32 erf), matches F.gelu().
# "gelu_fast": 6-segment piecewise-linear LUT, ~1% absolute error vs exact GELU.
# "gelu_tanh": FP32 tanh approximation, matches F.gelu(approximate="tanh").
_FUSED_GELU_VARIANTS = {
    "gelu": (ttnn.UnaryOpType.GELU, False),
    "gelu_fast": (ttnn.UnaryOpType.GELU, True),
    "gelu_tanh": ttnn.UnaryOpType.GELU_TANH,
}


def maybe_cast_activation(x: ttnn.Tensor, activation_dtype) -> ttnn.Tensor:
    """Cast an activation that is about to cross the fabric, if a quant config asked for it.

    Must be applied BEFORE the collective, never after: the win is in the bytes the gather moves,
    and the gather's page size is the tile size of the gathered dtype (bfloat8_b tiles are 1088 B
    vs bfloat16's 2048 B). A cast placed after the gather buys nothing and costs a full pass.

    The op-level constraint that makes this legal: the AG-matmul validates the activation and the
    weight dtypes independently, so a bf8 activation composes with the bf16-weight carve-out the
    fused addcmul epilogue requires.
    """
    if activation_dtype is None or x.get_dtype() == activation_dtype:
        return x
    return ttnn.typecast(x, activation_dtype)


def resolve_output_dtype(dtype, x: ttnn.Tensor):
    """Pin a block-float-fed matmul's output back to bf16 unless the caller asked for something else.

    Called only by a linear whose quant config opted in (``pin_output_bf16``), so no other
    model's default output dtype (``output_dtype.value_or(in0.dtype())``) changes. Keyed on the input's
    dtype so it covers an input that arrived block-float from upstream (e.g. the gate projection fed
    the shared bf8 activation), not just one this linear cast itself; without the pin a bf8 activation
    would push downstream into the residual stream and ``DistributedRMSNorm``, which rejects anything
    but bf16.
    """
    if dtype is None and x.get_dtype() in (ttnn.bfloat8_b, ttnn.bfloat4_b):
        return ttnn.bfloat16
    return dtype


class Linear(Module):
    """
    Linear layer with replicated weights
    """

    def __init__(
        self,
        in_features,
        out_features,
        bias=True,
        activation_fn=None,
        dtype=ttnn.bfloat16,
        mesh_device=None,
        # Branch addition kept over main: the H3 / Qwen3-VL layers pass an explicit config for
        # the sites that need more precision than the shared default.
        compute_kernel_config=None,
    ):
        super().__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.activation_fn = activation_fn
        self.fused_activation_fn = None
        self.fuse_swiglu = False
        if self.activation_fn == "swiglu":
            # Double out features for the packed [gate|up] swiglu weight.
            self.out_features = self.out_features * 2
            self.fuse_swiglu = True
            self.activation_fn = None
        elif self.activation_fn in _FUSED_GELU_VARIANTS:
            self.fused_activation_fn = _FUSED_GELU_VARIANTS[self.activation_fn]
            self.activation_fn = None
        self.mesh_device = mesh_device

        """
        NOTE: This is the special config which attains good correctness
        HiFi2 + packer_l1_acc + bf16 acc in a fused linear (matmul + bias) with unfused non-approx activation
        """
        self.compute_config = compute_kernel_config or ttnn.init_device_compute_kernel_config(
            mesh_device.arch(),
            math_fidelity=MATH_FIDELITY.get(dtype, ttnn.MathFidelity.HiFi2),
            math_approx_mode=False,
            fp32_dest_acc_en=True,
            packer_l1_acc=True,
        )

        self.weight = Parameter(total_shape=[self.in_features, self.out_features], device=mesh_device, dtype=dtype)
        self.bias = Parameter(total_shape=[1, self.out_features], device=mesh_device, dtype=dtype) if bias else None

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        if "weight" in state:
            weight = state["weight"].transpose(0, 1)
            if self.fuse_swiglu:
                weight = prepare_for_fused_swiglu(weight, ndev=1)
            state["weight"] = weight
        if "bias" in state:
            bias = state["bias"].reshape(1, -1)
            if self.fuse_swiglu:
                bias = prepare_for_fused_swiglu(bias, ndev=1)
            state["bias"] = bias

    def forward(self, x: ttnn.Tensor, compute_kernel_config=None, dtype=None, default_block_size=None) -> ttnn.Tensor:
        M, K, N = x.padded_shape[-2], x.padded_shape[-1], self.weight.data.padded_shape[-1]
        core_grid = get_matmul_core_grid(self.mesh_device)
        matmul_config = get_matmul_config(M, K, N, core_grid, default_block_size)
        output = ttnn.experimental.minimal_matmul(
            input_tensor=x,
            weight_tensor=self.weight.data,
            bias_tensor=self.bias.data if self.bias is not None else None,
            config=matmul_config,
            fused_activation=self.fused_activation_fn,
            compute_kernel_config=compute_kernel_config or self.compute_config,
            dtype=dtype,
            fuse_swiglu=self.fuse_swiglu,
        )

        return _apply_activation_fn(output, self.activation_fn)


def gelu_decomposed(x: ttnn.Tensor) -> ttnn.Tensor:
    # GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2)))
    # ttnn.gelu is the same, but avoiding for potential issues (see ttnn.layernorm)
    # Use a single scratch buffer that's reused for every intermediate so peak
    # DRAM is x + scratch (2x input) instead of the naive 6x.
    sqrt_2 = math.sqrt(2.0)
    tmp = ttnn.multiply(x, 1.0 / sqrt_2)
    ttnn.erf(tmp, output_tensor=tmp)
    ttnn.add(tmp, 1.0, output_tensor=tmp)
    ttnn.multiply(x, tmp, output_tensor=tmp)
    ttnn.multiply(tmp, 0.5, output_tensor=tmp)
    return tmp


class ColParallelLinear(Module):
    """
    Linear layer with column parallel weights
    """

    def __init__(
        self,
        in_features,
        out_features,
        bias=True,
        activation_fn=None,
        dtype=ttnn.bfloat16,
        mesh_device=None,
        mesh_axis=0,
        fsdp_mesh_axis=None,
        ccl_manager=None,
        chunks=None,
        chunk_sizes=None,
        activation_dtype=None,
        # Branch addition kept over main: the H3 / Qwen3-VL layers pass an explicit config for
        # the sites that need more precision than the shared default.
        compute_kernel_config=None,
        pin_output_bf16=False,
    ):
        super().__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.activation_fn = activation_fn
        self.fused_activation_fn = None
        self.fuse_swiglu = False
        if self.activation_fn == "swiglu":
            # Double out features for the packed [gate|up] swiglu weight.
            self.out_features = self.out_features * 2
            self.fuse_swiglu = True
            self.activation_fn = None
        elif self.activation_fn in _FUSED_GELU_VARIANTS:
            self.fused_activation_fn = _FUSED_GELU_VARIANTS[self.activation_fn]
            self.activation_fn = None
        self.mesh_device = mesh_device

        self.mesh_axis = mesh_axis
        self.fsdp_mesh_axis = fsdp_mesh_axis
        self.ccl_manager = ccl_manager
        self.chunks = chunks
        # Per-chunk output widths in ELEMENTS (global, pre-TP-shard); None => uniform N/chunks.
        self.chunk_sizes = chunk_sizes

        if self.fsdp_mesh_axis is not None:
            assert self.mesh_axis != self.fsdp_mesh_axis
            assert self.ccl_manager is not None

        # Optional cast of the *input* activation, set by a quant config. This is the only Linear
        # variant that honours it, because it is the only one whose input crosses the fabric: at
        # TP>1 the input is the payload of the fused all-gather, and the gather's page size follows
        # the dtype of the gathered tensor. Casting a RowParallel/replicated Linear's input would
        # buy matmul-internal precision only while paying a full typecast pass — and RowParallel's
        # input is the 4x-wide FFN intermediate, so that trade is strictly negative.
        self.activation_dtype = activation_dtype
        # Pin a bf8/bf4-fed output back to bf16 (see resolve_output_dtype). Set by the quant config on
        # the linears on its path; off elsewhere so no other model's default output dtype changes.
        self.pin_output_bf16 = pin_output_bf16

        self.compute_config = compute_kernel_config or ttnn.init_device_compute_kernel_config(
            mesh_device.arch(),
            math_fidelity=MATH_FIDELITY.get(dtype, ttnn.MathFidelity.HiFi2),
            math_approx_mode=False,
            fp32_dest_acc_en=True,
            packer_l1_acc=True,
        )

        self.weight = Parameter(
            total_shape=[self.in_features, self.out_features],
            mesh_axes=[fsdp_mesh_axis, mesh_axis],
            device=mesh_device,
            dtype=dtype,
        )
        self.bias = (
            Parameter(total_shape=[1, self.out_features], mesh_axes=[None, mesh_axis], device=mesh_device, dtype=dtype)
            if bias
            else None
        )

        self._mesh_axis_size = self.mesh_device.shape[self.mesh_axis] if self.mesh_axis is not None else 1

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        weight = state.pop("weight", None)
        bias = state.pop("bias", None)

        def permute_for_swiglu(tensor):
            assert self.activation_fn == "swiglu"
            ndev = self._mesh_axis_size
            tensor = tensor.reshape(-1, 2, ndev, tensor.shape[-1] // 2 // ndev)
            tensor = tensor.permute(0, 2, 1, 3)
            tensor = tensor.reshape(-1, self.out_features)
            assert tensor.shape[0] in [1, self.in_features]
            return tensor

        if weight is not None:
            weight = weight.transpose(0, 1)
            if self.fuse_swiglu:
                weight = prepare_for_fused_swiglu(weight, ndev=self._mesh_axis_size)
            elif self.activation_fn == "swiglu":
                weight = permute_for_swiglu(weight)
            state["weight"] = weight
        if bias is not None:
            bias = bias.reshape(1, -1)
            if self.fuse_swiglu:
                bias = prepare_for_fused_swiglu(bias, ndev=self._mesh_axis_size)
            elif self.activation_fn == "swiglu":
                bias = permute_for_swiglu(bias)
            state["bias"] = bias

    def _forward_fabric_agmm(self, x, weight, fabric_cfg, parallel_config, compute_kernel_config, dtype) -> ttnn.Tensor:
        """Optimized fabric-bound TP all-gather-matmul via strided_all_gather_minimal_matmul_async.

        The matmul runs on ``fabric_cfg.mm_core_grid`` (lower rows); the strided all-gather workers
        run on the rows starting at ``fabric_cfg.ag_core_grid_offset`` (disjoint region). Returns the
        single (chunks==1) matmul output; the op's first output is the gathered-K scratch.

        Under fused SwiGLU the weight is the packed [gate|up] matrix, so ``fabric_cfg`` blocks on the
        doubled width and the returned tensor is half as wide. The op has no dtype override, so the
        output follows the input/weight dtype rather than the caller's requested one.
        """
        mesh_axis = parallel_config.tensor_parallel.mesh_axis
        # The op gathers on dim 3 and fatals unless padded_shape[0] and [1] are both 1, but model
        # activations are rank 3 ([1, seq, K]), whose [1] is the sequence length. Widen here and
        # restore the caller's rank on the way out so this stays a drop-in for the non-fabric path.
        orig_rank = len(x.padded_shape)
        if orig_rank != 4:
            x = ttnn.unsqueeze_to_4D(x)
        if self.fuse_swiglu:
            # The factory partitions gate/up PAIRS across cores, so a pair must never straddle an
            # N block. N_tiles and N_tiles_per_core are even by construction of the packed weight.
            assert (
                fabric_cfg.N_block_size % 2 == 0
            ), f"fuse_swiglu needs an even N_block_size (in tiles), got {fabric_cfg.N_block_size}"
        matmul_config = ttnn.MinimalMatmulConfig(
            M_block_size=fabric_cfg.M_block_size,
            K_block_size=fabric_cfg.K_block_size,
            N_block_size=fabric_cfg.N_block_size,
            subblock_h=fabric_cfg.subblock_h,
            subblock_w=fabric_cfg.subblock_w,
            compute_with_storage_grid_size=fabric_cfg.mm_core_grid,
        )
        ag_persistent_buffer = self.ccl_manager.get_ag_ping_pong_buffer(x.shape, 3, mesh_axis, dtype=x.get_dtype())
        ag_global_semaphores = self.ccl_manager.get_strided_ag_mm_semaphore(mesh_axis, fabric_cfg.num_workers_per_link)
        dram = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM)
        outputs = ttnn.experimental.strided_all_gather_minimal_matmul_async(
            x,
            weight,
            persistent_output_buffer=ag_persistent_buffer,
            dim=3,
            multi_device_global_semaphore=ag_global_semaphores,
            strided_all_gather_core_grid_offset=fabric_cfg.ag_core_grid_offset,
            num_links=self.ccl_manager.num_links,
            memory_config_ag=dram,
            topology=self.ccl_manager.topology,
            cluster_axis=mesh_axis,
            bias=self.bias.data if self.bias is not None else None,
            fused_activation=self.fused_activation_fn,
            config=matmul_config,
            memory_config_mm=dram,
            compute_kernel_config=compute_kernel_config or self.compute_config,
            num_workers_per_link=fabric_cfg.num_workers_per_link,
            num_buffers_per_channel=fabric_cfg.num_buffers_per_channel,
            read_local_slice_from_input=True,
            chunks=1,
            fuse_swiglu=self.fuse_swiglu,
        )
        # Op returns [all_gather_output, matmul_chunk_0]; take the single matmul chunk.
        out = _apply_activation_fn(outputs[1], self.activation_fn)
        if orig_rank != 4:
            out = ttnn.reshape(out, tuple(out.shape)[-orig_rank:])
        return out

    def forward(
        self,
        x: ttnn.Tensor,
        compute_kernel_config=None,
        default_block_size=None,
        parallel_config=None,
        dtype=None,
        addcmul_a=None,
        addcmul_b=None,
        addcmul_scalar: float = 1.0,
        core_grid=None,
        use_heuristic_mmcfg=False,
    ) -> ttnn.Tensor | list[ttnn.Tensor]:
        """
        Expects x to be replicated.
        Return output fractured on columns.
        If chunks is set, returns a list of tensors split along the output dimension.

        `addcmul_a` / `addcmul_b` fuse a gated residual into the matmul epilogue, returning
        `addcmul_a + addcmul_scalar * matmul_result * addcmul_b`. Both must already be at the
        per-TP-device output slice and require `parallel_config`. The Ring path fuses them into
        the all-gather-matmul; other paths route through the addcmul-fused minimal matmul.
        """
        if addcmul_a is not None or addcmul_b is not None:
            if (addcmul_a is None) != (addcmul_b is None):
                msg = "addcmul_a and addcmul_b must be given together"
                raise ValueError(msg)
            if parallel_config is None or parallel_config.tensor_parallel.factor <= 1:
                msg = "fused addcmul needs the all-gather-matmul path; pass parallel_config"
                raise ValueError(msg)
            if self.chunks is not None and self.chunks > 1:
                msg = "fused addcmul is not supported alongside chunked output"
                raise ValueError(msg)

        x = maybe_cast_activation(x, self.activation_dtype)
        if self.pin_output_bf16:
            dtype = resolve_output_dtype(dtype, x)

        if self.fsdp_mesh_axis is not None and self.mesh_device.shape[self.fsdp_mesh_axis] > 1:
            unsqueezed_weight = ttnn.unsqueeze_to_4D(self.weight.data)
            weight = self.ccl_manager.all_gather_persistent_buffer(
                unsqueezed_weight, dim=2, mesh_axis=self.fsdp_mesh_axis
            )

            weight = ttnn.reshape(weight, (weight.shape[-2], weight.shape[-1]))
        else:
            weight = self.weight.data

        parallel_config_tp = parallel_config.tensor_parallel.factor if parallel_config is not None else 1
        needs_gather = x.padded_shape[-1] != weight.padded_shape[-2]  # If gathered, switch to non fused AGMM
        if parallel_config_tp > 1 and self.ccl_manager.topology == ttnn.Topology.Ring and needs_gather:
            M, K, N = x.padded_shape[-2], weight.padded_shape[-2], weight.padded_shape[-1]
            full_grid = self.mesh_device.compute_with_storage_grid_size()

            # Fabric-bound path: known shapes route to the optimized strided all-gather-matmul op.
            # N is the weight width, so a fused-SwiGLU layer keys on its packed [gate|up] width.
            # Restricted to chunks==1; other shapes fall through to the all_gather_minimal_matmul_async
            # path below. The op gathers on dim 3 of a rank-4 input, so every dim above the matmul's
            # (M, K) must be unit; _forward_fabric_agmm widens rank-3 activations to satisfy that.
            fabric_cfg = get_fabric_agmm_config(M, K, N, (self.chunks or 1), full_grid)
            has_unit_batch = len(x.padded_shape) <= 4 and all(d == 1 for d in list(x.padded_shape)[:-2])
            if fabric_cfg is not None and self.chunks in (None, 1) and has_unit_batch and addcmul_a is None:
                return self._forward_fabric_agmm(x, weight, fabric_cfg, parallel_config, compute_kernel_config, dtype)

            core_grid = core_grid or ttnn.CoreCoord(full_grid.x, full_grid.y - 1)
            matmul_config = get_matmul_config(M, K, N, core_grid, default_block_size, use_heuristic=use_heuristic_mmcfg)

            ag_persistent_buffer = self.ccl_manager.get_ag_ping_pong_buffer(
                x.shape, -1, parallel_config.tensor_parallel.mesh_axis, dtype=x.get_dtype()
            )
            ag_global_semaphores = self.ccl_manager.get_ag_ping_pong_semaphore(
                parallel_config.tensor_parallel.mesh_axis
            )
            outputs = ttnn.experimental.all_gather_minimal_matmul_async(
                input_tensor=x,
                weight_tensor=weight,
                bias_tensor=self.bias.data if self.bias is not None else None,
                config=matmul_config,
                fused_activation=self.fused_activation_fn,
                compute_kernel_config=compute_kernel_config or self.compute_config,
                persistent_output_buffer=ag_persistent_buffer,
                multi_device_global_semaphore=ag_global_semaphores,
                num_links=self.ccl_manager.num_links,
                topology=self.ccl_manager.topology,
                cluster_axis=parallel_config.tensor_parallel.mesh_axis,
                barrier_semaphore=None,
                num_workers_per_link=full_grid.x // self.ccl_manager.num_links,
                num_buffers_per_channel=48 if not is_blackhole() else 24,
                chunks=self.chunks if self.chunks is not None else 1,
                # Op's N is per-device, so pass per-device widths (each global width is % TP == 0).
                chunk_sizes=(
                    [w // parallel_config.tensor_parallel.factor for w in self.chunk_sizes] if self.chunk_sizes else []
                ),
                dtype=dtype,
                fuse_swiglu=self.fuse_swiglu,
                scalar=addcmul_scalar if addcmul_a is not None else None,
                addcmul_input_tensor1=addcmul_a,
                addcmul_input_tensor2=addcmul_b,
            )

            if self.chunks is not None and (self.chunks > 1):
                return [_apply_activation_fn(o, self.activation_fn) for o in outputs]
            else:
                output = outputs[0]
        else:
            M, K, N = x.padded_shape[-2], x.padded_shape[-1], weight.padded_shape[-1]
            core_grid = get_matmul_core_grid(self.mesh_device)

            # Gather if needed here. Helps cleanup upstream code
            if needs_gather:
                x = self.ccl_manager.all_gather_persistent_buffer(
                    x, dim=-1, mesh_axis=parallel_config.tensor_parallel.mesh_axis, use_hyperparams=True
                )

            if self.chunks is not None:
                matmul_config = get_matmul_config(M, K, N, core_grid, default_block_size)
                outputs = ttnn.experimental.minimal_matmul_split(
                    x,
                    weight,
                    chunks=self.chunks,
                    dim=-1,
                    bias_tensor=self.bias.data if self.bias is not None else None,
                    fused_activation=self.fused_activation_fn,
                    compute_kernel_config=compute_kernel_config or self.compute_config,
                    config=matmul_config,
                    dtype=dtype,
                    fuse_swiglu=self.fuse_swiglu,
                )
                return [_apply_activation_fn(o, self.activation_fn) for o in outputs]
            if addcmul_a is not None:
                # This branch used to accept addcmul_a/addcmul_b and silently drop them:
                # only the Ring all-gather-matmul above fused them, so on Linear topology
                # the gated residual promised by the docstring was never added (flux2's
                # double blocks lost their attention residual this way). Route through the
                # addcmul-fused minimal matmul instead. That op has no fused-activation
                # support, so reject the combination rather than compute the wrong thing.
                if self.fused_activation_fn is not None or self.fuse_swiglu:
                    msg = "fused addcmul is not supported alongside a fused activation on the minimal_matmul path"
                    raise ValueError(msg)
                # x may have been gathered above, so size the config from its current K.
                matmul_config = get_matmul_config(M, x.padded_shape[-1], N, core_grid, default_block_size)
                output = ttnn.experimental.dit_minimal_matmul_addcmul_fused(
                    x,
                    weight,
                    addcmul_scalar,
                    addcmul_a,
                    addcmul_b,
                    bias_tensor=self.bias.data if self.bias is not None else None,
                    config=matmul_config,
                    compute_kernel_config=compute_kernel_config or self.compute_config,
                    dtype=dtype,
                )
            else:
                matmul_config = get_matmul_config(M, K, N, core_grid, default_block_size)
                output = ttnn.experimental.minimal_matmul(
                    input_tensor=x,
                    weight_tensor=weight,
                    bias_tensor=self.bias.data if self.bias is not None else None,
                    config=matmul_config,
                    fused_activation=self.fused_activation_fn,
                    compute_kernel_config=compute_kernel_config or self.compute_config,
                    dtype=dtype,
                    fuse_swiglu=self.fuse_swiglu,
                )

        return _apply_activation_fn(output, self.activation_fn)


class RowParallelLinear(Module):
    """
    Linear layer with row parallel weights
    """

    def __init__(
        self,
        in_features,
        out_features,
        bias=True,
        dtype=ttnn.bfloat16,
        mesh_device=None,
        mesh_axis=0,
        fsdp_mesh_axis=None,
        ccl_manager=None,
        # Branch addition kept over main: the H3 / Qwen3-VL layers pass an explicit config for
        # the sites that need more precision than the shared default.
        compute_kernel_config=None,
        mm_memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM),
    ):
        super().__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.mesh_device = mesh_device
        self.mesh_axis = mesh_axis
        self.fsdp_mesh_axis = fsdp_mesh_axis
        self.ccl_manager = ccl_manager
        self.mm_memory_config = mm_memory_config

        if self.fsdp_mesh_axis is not None:
            assert self.mesh_axis != self.fsdp_mesh_axis

        self.compute_config = compute_kernel_config or ttnn.init_device_compute_kernel_config(
            mesh_device.arch(),
            math_fidelity=MATH_FIDELITY.get(dtype, ttnn.MathFidelity.HiFi2),
            math_approx_mode=False,
            fp32_dest_acc_en=True,
            packer_l1_acc=True,
        )

        ndev = self.mesh_device.shape[self.mesh_axis] if self.mesh_axis is not None else 1

        self.weight = Parameter(
            total_shape=[self.in_features, self.out_features],
            mesh_axes=[mesh_axis, fsdp_mesh_axis],
            device=mesh_device,
            dtype=dtype,
        )
        self.bias = (
            Parameter(
                total_shape=[1, self.out_features * ndev], mesh_axes=[None, mesh_axis], device=mesh_device, dtype=dtype
            )
            if bias
            else None
        )

        self._mesh_axis_size = ndev

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        if "weight" in state:
            state["weight"] = state["weight"].transpose(0, 1)

        bias = state.pop("bias", None)
        if bias is not None:
            bias = bias.reshape(1, -1)
            if self._mesh_axis_size > 1:
                zero_bias = torch.zeros(1, bias.shape[1] * (self._mesh_axis_size - 1))
                bias = torch.cat([bias, zero_bias], dim=-1)
            state["bias"] = bias

    def forward(
        self,
        x: ttnn.Tensor | list[ttnn.Tensor],
        *,
        compute_kernel_config=None,
        use_persistent_buffer: bool = True,
        default_block_size: tuple = None,
        dtype=None,
        gather_output: bool = False,
    ) -> ttnn.Tensor:
        """
        Expects x to be column fractured.
        x may be a 2-element list [prefix, suffix] for fused concat over K (concat-free).
        Return output fractured on columns.
        """
        if self.fsdp_mesh_axis is not None and self.mesh_device.shape[self.fsdp_mesh_axis] > 1:
            unsqueezed_weight = ttnn.unsqueeze_to_4D(self.weight.data)
            weight = self.ccl_manager.all_gather_persistent_buffer(
                unsqueezed_weight, dim=3, mesh_axis=self.fsdp_mesh_axis
            )

            weight = ttnn.reshape(weight, (weight.shape[-2], weight.shape[-1]))
        else:
            weight = self.weight.data

        if isinstance(x, (list, tuple)):
            assert len(x) == 2, f"RowParallelLinear.forward: list x must be [prefix, suffix], got {len(x)}"
            x, x_second = x
            K = weight.padded_shape[-2]
        else:
            x_second = None
            K = x.padded_shape[-1]

        M, N = x.padded_shape[-2], weight.padded_shape[-1]
        core_grid = get_matmul_core_grid(self.mesh_device)
        matmul_config = get_matmul_config(M, K, N, core_grid, default_block_size)
        output = ttnn.experimental.minimal_matmul(
            input_tensor=[x, x_second] if x_second is not None else x,
            weight_tensor=weight,
            bias_tensor=self.bias.data if self.bias is not None else None,
            config=matmul_config,
            compute_kernel_config=compute_kernel_config or self.compute_config,
            dtype=dtype,
        )

        if self._mesh_axis_size > 1:
            # Reduce over rows when replicating: N may be too narrow to scatter over the mesh axis.
            dim = -2 if gather_output else -1
            output = self.ccl_manager.reduce_scatter(
                output, dim=dim, mesh_axis=self.mesh_axis, use_persistent_buffer=use_persistent_buffer
            )
            if gather_output:
                output = self.ccl_manager.all_gather(
                    output, dim=dim, mesh_axis=self.mesh_axis, use_hyperparams=True, use_persistent_buffer=True
                )

        return output

    def forward_fused_addcmul(
        self,
        x: ttnn.Tensor | list[ttnn.Tensor],
        addcmul_a: ttnn.Tensor,
        addcmul_b: ttnn.Tensor,
        scalar: float = 1.0,
        *,
        compute_kernel_config=None,
        dtype=None,
    ) -> ttnn.Tensor:
        """Fused RowParallel matmul + reduce-scatter + addcmul at the RS final write step.

        Computes: output = addcmul_a + scalar * rs_result * addcmul_b

        ``x`` may be a single tensor or a 2-element list ``[prefix, suffix]`` for fused concat over K.
        The weight must be per-segment tile-padded (see ``prepare_weight_for_concatenated_input``).

        Both addcmul_a and addcmul_b must already be at their per-TP-device slice size
        [D/tp]. The RS kernel fuses the addcmul at the final ring write, eliminating
        extra CCL ops entirely.
        """
        if self.fsdp_mesh_axis is not None and self.mesh_device.shape[self.fsdp_mesh_axis] > 1:
            unsqueezed_weight = ttnn.unsqueeze_to_4D(self.weight.data)
            weight = self.ccl_manager.all_gather_persistent_buffer(
                unsqueezed_weight, dim=3, mesh_axis=self.fsdp_mesh_axis
            )
            weight = ttnn.reshape(weight, (weight.shape[-2], weight.shape[-1]))
        else:
            weight = self.weight.data

        # x: single tensor, or [prefix, suffix] virtually concatenated over K (concat-free).
        if isinstance(x, (list, tuple)):
            assert len(x) == 2, f"forward_fused_addcmul: list x must be exactly [prefix, suffix], got {len(x)}"
            x, x_second = x
        else:
            x_second = None

        # For fused concat the matmul K spans both halves = the weight's K; x is only the prefix half.
        K = weight.padded_shape[-2] if x_second is not None else x.padded_shape[-1]
        M, N = x.padded_shape[-2], weight.padded_shape[-1]
        core_grid = self.mesh_device.compute_with_storage_grid_size()

        needs_reshape = len(x.shape) <= 3
        if needs_reshape:
            x = ttnn.unsqueeze(x, 0)
            if x_second is not None:
                x_second = ttnn.unsqueeze(x_second, 0)
        pre_rs_shape = tuple(list(x.shape)[:-1] + [N])
        _, rs_output_buffer = self.ccl_manager.get_rs_ping_pong_buffer(
            pre_rs_shape, 3, self.mesh_axis, return_intermediate=False
        )
        _, output = ttnn.experimental.minimal_matmul_strided_reduce_scatter_async(
            input_tensor=[x, x_second] if x_second is not None else x,
            weight_tensor=weight,
            dim=3,
            multi_device_global_semaphore=self.ccl_manager.get_rs_ping_pong_semaphore(self.mesh_axis),
            **get_fused_mmrs_config(M, K, N, core_grid, self.ccl_manager.num_links),
            bias=self.bias.data if self.bias is not None else None,
            memory_config_mm=self.mm_memory_config,
            rs_intermediate_mem_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM),
            rs_output_mem_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM),
            topology=self.ccl_manager.topology,
            cluster_axis=self.mesh_axis,
            compute_kernel_config=compute_kernel_config or self.compute_config,
            using_persistent_buffers=True,
            optional_rs_output_tensor=rs_output_buffer,
            fused_ternary_scalar=scalar,
            addcmul_input_tensor1=addcmul_a,
            addcmul_input_tensor2=addcmul_b,
            dtype=dtype,
            mm_progress_counters=self.ccl_manager.get_mm_progress_counters_buffer(),
        )
        if needs_reshape:
            output = ttnn.squeeze(output, 0)
        return output


def _apply_activation_fn(t: ttnn.Tensor, activation_fn: str | None) -> ttnn.Tensor:
    if activation_fn is None:
        return t
    if activation_fn == "silu":
        return ttnn.silu(t)
    if activation_fn == "decomposed_gelu":
        return gelu_decomposed(t)
    if activation_fn == "quick_gelu":
        return t * ttnn.sigmoid(1.702 * t)  # quick approx gelu
    if activation_fn == "swiglu":
        t, gate = ttnn.chunk(t, 2, -1)
        return ttnn.multiply_(t, ttnn.silu(gate, output_tensor=gate))

    msg = f"Activation function {activation_fn} not supported"
    raise ValueError(msg)


def prepare_chunked_linear_output(
    state: dict[str, torch.Tensor], *, prefix: str, device_count: int, chunks: int
) -> None:
    weight_key = f"{prefix}.weight"
    bias_key = f"{prefix}.bias"

    weight = state.get(weight_key)
    bias = state.get(bias_key)

    if weight is not None:
        _, in_dim = weight.shape
        weight = weight.reshape([chunks, device_count, -1, in_dim]).transpose(0, 1).reshape([-1, in_dim])
        state[weight_key] = weight

    if bias is not None:
        bias = state[bias_key].reshape([chunks, device_count, -1]).transpose(0, 1).reshape([-1])
        state[bias_key] = bias


# =====================================================================
# LoRA-aware Linear variants
# =====================================================================
# Each variant subclasses its base Linear + the shared LoRAMixin. The
# mixin offers two execution paths chosen at construction with
# ``lora_mode`` ('fuse' or 'runtime'); see models/tt_dit/layers/lora.py
# for the trade-offs.
from .lora import LoRAMixin  # noqa: E402


class LoRALinear(LoRAMixin, Linear):
    def __init__(self, *args, lora_mode: str = "fuse", **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._init_lora_state(mode=lora_mode)


class LoRAColParallelLinear(LoRAMixin, ColParallelLinear):
    def __init__(self, *args, lora_mode: str = "fuse", **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._init_lora_state(mode=lora_mode)


class LoRARowParallelLinear(LoRAMixin, RowParallelLinear):
    def __init__(self, *args, lora_mode: str = "fuse", **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # Runtime mode lacks the all-reduce the base path performs via
        # reduce_scatter, so the delta and base sit at different mesh layouts.
        if lora_mode == "runtime" and self._mesh_axis_size > 1:
            raise ValueError(
                "LoRARowParallelLinear with lora_mode='runtime' is unsupported "
                f"at TP>1 (mesh_axis_size={self._mesh_axis_size}); use lora_mode='fuse'"
            )
        self._init_lora_state(mode=lora_mode)