code/models/common/modules/lm_head/lm_head_1d.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /common /modules /lm_head /lm_head_1d.py

Add files using upload-large-folder tool

656477e verified 20 days ago

24.4 kB

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
TTTv2-style LM Head module for 1D-topology devices: N150 (1x1), N300 (1x2), T3K (1x8).
Computes logits over the vocabulary by splitting the output projection into
weight chunks that fit in L1, running linear ops per chunk, and concatenating.
Each device holds the correct logits for its vocab shard (column-parallel linear,
no partial sums). The caller handles all_gather for multi-device argmax.
Execution path:
  for each (weight, pc): linear(x, weight) → sharded_to_interleaved → append
  → concat
"""

import math
from dataclasses import dataclass, replace
from pathlib import Path
from typing import List

import ttnn
from models.common.lightweightmodule import LightweightModule
from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight
from models.common.tensor_utils import TILE_SIZE
from models.common.tensor_utils import nearest_32 as _nearest_32

# =============================================================================
# Config dataclass
# =============================================================================


@dataclass
class LMHead1DConfig:
    """
    Configuration for LMHead1D.
    Simple usage (pre-split weights):
        config = LMHead1DConfig(output_weights=[w1, w2, w3])
    The program_configs and other fields are auto-computed from weights if None.
    """

    # Required: output projection weights (already split for L1 fit)
    output_weights: List[LazyWeight]

    # Optional: device
    mesh_device: ttnn.MeshDevice | None = None

    # Optional: derived from weights if None
    dim: int | None = None

    # Optional: batch/tile config
    max_batch_size: int = 32

    # Optional: power-user overrides
    program_configs: List | None = None
    compute_kernel_config: ttnn.DeviceComputeKernelConfig | None = None
    # Logical output width for each split on one device. A physical weight shard
    # may be tile-padded beyond this width; forward trims it before concat.
    output_split_sizes: List[int] | None = None
    lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b
    output_memcfg: ttnn.MemoryConfig | None = None

    # Input memory config (None = DRAM interleaved for simple API, width-sharded for from_model_args)
    input_memcfg: ttnn.MemoryConfig | None = None

    # Weight memory configs (None = auto-compute)
    weights_memcfgs: List[ttnn.MemoryConfig] | None = None

    def is_resolved(self) -> bool:
        return all(getattr(self, f) is not None for f in self.__dataclass_fields__)


# =============================================================================
# LMHead1D
# =============================================================================


class LMHead1D(LightweightModule):
    """
    LM Head for non-TG (1D) devices.
    Splits vocabulary projection into L1-sized chunks, runs linear per chunk,
    and concatenates. Each device holds correct logits for its vocab shard;
    caller handles all_gather for multi-device argmax.
    Simple API:
        lm_head = LMHead1D(output_weights=[w1, w2])
    Power API:
        config = LMHead1DConfig(output_weights=[w1, w2], lm_head_dtype=ttnn.bfloat16)
        lm_head = LMHead1D.from_config(config)
    Execution path:
      for each (w, pc): linear(x, w) → sharded_to_interleaved → concat
    """

    def __init__(self, output_weights: List[LazyWeight]):
        super().__init__()
        self.config = resolve_lm_head_1d_arch_config(LMHead1DConfig(output_weights=output_weights))
        self._device_weights_loaded = False

    @classmethod
    def from_config(cls, config: LMHead1DConfig):
        instance = object.__new__(cls)
        super(LMHead1D, instance).__init__()
        instance.config = resolve_lm_head_1d_arch_config(config)
        instance._device_weights_loaded = False
        return instance

    def load_device_weights(self):
        if self._device_weights_loaded:
            return
        self.output_weights = [w.get_device_weight() for w in self.config.output_weights]
        self._device_weights_loaded = True

    def forward(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor:
        """
        Compute logits over vocabulary.
        Args:
            x: Input hidden states, shape [1, 1, batch_rows, dim].
        Returns:
            Logits tensor, shape [1, 1, batch_rows, padded_vocab / num_devices] per device.
        """
        self.load_device_weights()
        x = _load_input_device_tensor(x, self.config)
        cfg = self.config

        outputs = []
        for weight, pc, logical_width in zip(self.output_weights, cfg.program_configs, cfg.output_split_sizes):
            if pc is not None:
                # DRAM-sharded path (from_model_args): width-sharded output → interleaved
                output = ttnn.linear(
                    x,
                    weight,
                    compute_kernel_config=cfg.compute_kernel_config,
                    program_config=pc,
                    memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG,
                    dtype=cfg.lm_head_dtype,
                )
                output = ttnn.sharded_to_interleaved(output, memory_config=cfg.output_memcfg)
            else:
                # Auto path (simple API): interleaved output directly
                output = ttnn.linear(
                    x,
                    weight,
                    compute_kernel_config=cfg.compute_kernel_config,
                    memory_config=cfg.output_memcfg,
                    dtype=cfg.lm_head_dtype,
                )
            if output.shape[-1] != logical_width:
                output = ttnn.slice(
                    output,
                    (0, 0, 0, 0),
                    (output.shape[0], output.shape[1], output.shape[2], logical_width),
                )
            outputs.append(output)

        # Concatenate splits
        output = ttnn.concat(outputs, dim=-1, memory_config=cfg.output_memcfg)

        return output

    # [INFO] this is the entry point for TTTv1 model_config.py and will retire with TTTv1
    @classmethod
    def from_model_args(
        cls,
        mesh_device,
        args,
        state_dict,
        state_dict_prefix,
        weight_cache_path,
        max_columns_per_device,
        dtype=None,
        model_config=None,
        tt_ccl=None,
    ):
        """Factory method for backward compatibility with ModelArgs.
        Note: tt_ccl is accepted for signature compatibility with TTTv1 LMHead
        but is not used -- 1D LMHead does not need CCL (column-parallel, no partial sums).
        """
        if args.is_galaxy:
            raise ValueError("LMHead1D cannot be used for Galaxy devices.")

        import torch

        vocab_size = args.vocab_size
        num_devices = mesh_device.get_num_devices()
        dim = args.dim
        padded_vocab_size = math.ceil(vocab_size / (TILE_SIZE * num_devices)) * (TILE_SIZE * num_devices)
        size_per_device = padded_vocab_size // num_devices
        num_splits = math.ceil(size_per_device / max_columns_per_device)
        split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1)
        split_sizes.append(size_per_device - sum(split_sizes))

        # Build output weights
        source_weight = state_dict[f"{state_dict_prefix}output.weight"]
        if tuple(source_weight.shape) != (vocab_size, dim):
            raise ValueError(
                f"LMHead1D output weight must have shape {(vocab_size, dim)}, got {tuple(source_weight.shape)}"
            )
        torch_output_weights = source_weight.permute(1, 0)
        if vocab_size < padded_vocab_size:
            padding_size = padded_vocab_size - vocab_size
            torch_output_weights = torch.cat(
                [
                    torch_output_weights,
                    torch.zeros(torch_output_weights.shape[0], padding_size, dtype=torch_output_weights.dtype),
                ],
                dim=-1,
            )

        # DRAM grid for weight memory configs
        dram_size = mesh_device.dram_grid_size()
        dram_grid = ttnn.CoreRangeSet(
            {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))}
        )

        cache_dir = None if args.dummy_weights else Path(weight_cache_path) / "lm_head"

        output_weights = []
        weights_memcfgs = []
        for i, split_size in enumerate(split_sizes):
            device_splits = []
            physical_split_size = math.ceil(split_size / TILE_SIZE) * TILE_SIZE
            for device_idx in range(num_devices):
                start = device_idx * size_per_device + sum(split_sizes[:i])
                end = start + split_size
                device_split = torch_output_weights[:, start:end]
                if split_size < physical_split_size:
                    device_split = torch.cat(
                        [
                            device_split,
                            torch.zeros(dim, physical_split_size - split_size, dtype=device_split.dtype),
                        ],
                        dim=-1,
                    )
                device_splits.append(device_split)
            combined_split = torch.cat(device_splits, dim=-1)

            mem_cfg = _create_dram_sharded_mem_config(
                k=dim,
                n=math.ceil(combined_split.shape[-1] / num_devices),
                dram_grid=dram_grid,
                tile_size=TILE_SIZE,
                dram_cores=dram_size.x,
            )
            weights_memcfgs.append(mem_cfg)

            w_dtype = dtype if dtype is not None else ttnn.bfloat8_b
            output_weights.append(
                LazyWeight(
                    source=combined_split,
                    dtype=w_dtype,
                    device=mesh_device,
                    mesh_mapper_config=ttnn.MeshMapperConfig(
                        placements=[ttnn.PlacementShard(-1)],
                        mesh_shape_override=ttnn.MeshShape([num_devices]),
                    ),
                    layout=ttnn.TILE_LAYOUT,
                    memory_config=mem_cfg,
                    cache_dir_weight_name=(
                        (cache_dir, f"output_split_{i}_logical_{split_size}_physical_{combined_split.shape[-1]}")
                        if cache_dir
                        else None
                    ),
                )
            )

        # Program configs - use the args.lm_head_core_grid which TTTv1 carefully computes
        tile_padded_batch_rows = TILE_SIZE * math.ceil(args.max_batch_size / TILE_SIZE)
        lm_head_core_grid = args.lm_head_core_grid
        program_configs = [
            args.dram_matmul_config(
                tile_padded_batch_rows,
                dim,
                ss,
                lm_head_core_grid.num_cores,
            )
            for ss in split_sizes
        ]

        # Get input memory config: width-sharded across lm_head_core_grid
        # (DRAM-sharded matmul requires width-sharded input in L1)
        input_memcfg = ttnn.create_sharded_memory_config(
            (
                tile_padded_batch_rows,
                _nearest_32(dim // lm_head_core_grid.num_cores),
            ),
            lm_head_core_grid,
            ttnn.ShardStrategy.WIDTH,
            ttnn.ShardOrientation.ROW_MAJOR,
            use_height_and_width_as_shard_shape=True,
        )

        config = LMHead1DConfig(
            output_weights=output_weights,
            mesh_device=mesh_device,
            dim=dim,
            max_batch_size=args.max_batch_size,
            program_configs=program_configs,
            output_split_sizes=split_sizes,
            lm_head_dtype=getattr(args, "lm_head_dtype", ttnn.bfloat8_b),
            output_memcfg=ttnn.L1_MEMORY_CONFIG,
            input_memcfg=input_memcfg,
            weights_memcfgs=weights_memcfgs,
        )
        return cls.from_config(config)


# =============================================================================
# Config resolution
# =============================================================================


def _compute_kernel_config_hifi2(arch) -> ttnn.DeviceComputeKernelConfig:
    """Construct the TTTv1 LM-head recipe for the selected architecture."""
    return ttnn.init_device_compute_kernel_config(
        arch,
        math_fidelity=ttnn.MathFidelity.HiFi2,
        math_approx_mode=False,
        fp32_dest_acc_en=False,
        packer_l1_acc=True,
    )


def _validate_lm_head_compute_kernel_config(compute_kernel_config) -> None:
    if compute_kernel_config is None:
        raise ValueError("LMHead1D architecture config requires compute_kernel_config")
    if not isinstance(compute_kernel_config, ttnn.WormholeComputeKernelConfig):
        raise TypeError(
            "LMHead1D compute_kernel_config must be a concrete TTNN compute kernel config, "
            f"got {type(compute_kernel_config).__name__}"
        )


def _derive_lm_head_mesh_device(config: LMHead1DConfig) -> ttnn.MeshDevice:
    if not config.output_weights:
        raise ValueError("LMHead1D requires at least one output weight")
    mesh_device = config.mesh_device
    if mesh_device is None:
        mesh_device = config.output_weights[0].device
    if mesh_device is None:
        mesh_device = ttnn.GetDefaultDevice()
    if mesh_device is None:
        raise ValueError("LMHead1D mesh_device must be available")
    for split_index, weight in enumerate(config.output_weights):
        if weight.device is not None and weight.device is not mesh_device:
            raise ValueError(f"LMHead1D mesh_device must match output weight {split_index} device")
    return mesh_device


def _validate_lm_head_program_configs(config: LMHead1DConfig) -> None:
    num_splits = len(config.output_weights)
    if num_splits == 0:
        raise ValueError("LMHead1D requires at least one output weight")
    if config.program_configs is None or len(config.program_configs) != num_splits:
        raise ValueError("LMHead1D requires exactly one program config per output-weight split")
    if config.weights_memcfgs is None or len(config.weights_memcfgs) != num_splits:
        raise ValueError("LMHead1D requires exactly one weight memory config per output-weight split")
    if config.output_split_sizes is None or len(config.output_split_sizes) != num_splits:
        raise ValueError("LMHead1D requires exactly one logical output width per output-weight split")

    num_devices = config.mesh_device.get_num_devices()
    physical_widths = []
    for split_index, (weight, logical_width) in enumerate(zip(config.output_weights, config.output_split_sizes)):
        source_shape = weight.source.shape
        if len(source_shape) < 2 or source_shape[-2] != config.dim:
            raise ValueError(
                f"LMHead1D split {split_index} weight K dimension must equal dim={config.dim}, got {source_shape}"
            )
        combined_width = source_shape[-1]
        if combined_width % num_devices:
            raise ValueError(
                f"LMHead1D split {split_index} combined physical width {combined_width} "
                f"must be divisible by {num_devices} devices"
            )
        physical_width = combined_width // num_devices
        if physical_width % TILE_SIZE:
            raise ValueError(
                f"LMHead1D split {split_index} physical per-device width {physical_width} must be tile aligned"
            )
        if not isinstance(logical_width, int) or logical_width <= 0 or logical_width > physical_width:
            raise ValueError(
                f"LMHead1D split {split_index} logical width must satisfy 0 < logical <= {physical_width}, "
                f"got {logical_width}"
            )
        physical_widths.append(physical_width)

    has_explicit_program = any(program_config is not None for program_config in config.program_configs)
    if has_explicit_program:
        if config.input_memcfg is None or not config.input_memcfg.is_sharded():
            raise ValueError("LMHead1D DRAM-sharded programs require a sharded input_memcfg")
        if config.input_memcfg.buffer_type != ttnn.BufferType.L1:
            raise ValueError("LMHead1D DRAM-sharded programs require input_memcfg in L1")
        if config.input_memcfg.memory_layout != ttnn.TensorMemoryLayout.WIDTH_SHARDED:
            raise ValueError("LMHead1D DRAM-sharded programs require width-sharded input_memcfg")

    for split_index, program_config in enumerate(config.program_configs):
        if program_config is None:
            continue
        if not isinstance(program_config, ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig):
            raise TypeError(
                "LMHead1D explicit program configs must use the DRAM-sharded matmul path; "
                f"split {split_index} has {type(program_config).__name__}"
            )
        if program_config.in0_block_w <= 0 or program_config.per_core_M <= 0 or program_config.per_core_N <= 0:
            raise ValueError(f"LMHead1D split {split_index} program dimensions must be positive")
        dim_tiles = config.dim // TILE_SIZE
        if dim_tiles % program_config.in0_block_w != 0:
            raise ValueError(f"LMHead1D split {split_index} dim tiles must be divisible by in0_block_w")
        expected_per_core_m = math.ceil(config.max_batch_size / TILE_SIZE)
        if program_config.per_core_M != expected_per_core_m:
            raise ValueError(
                f"LMHead1D split {split_index} per_core_M={program_config.per_core_M} "
                f"does not match padded batch tiles {expected_per_core_m}"
            )

        weight_memcfg = config.weights_memcfgs[split_index]
        if not weight_memcfg.is_sharded():
            raise ValueError(f"LMHead1D split {split_index} DRAM-sharded program requires sharded weight memory")
        if weight_memcfg.buffer_type != ttnn.BufferType.DRAM:
            raise ValueError(f"LMHead1D split {split_index} weight memory must reside in DRAM")
        if weight_memcfg.memory_layout != ttnn.TensorMemoryLayout.WIDTH_SHARDED:
            raise ValueError(f"LMHead1D split {split_index} weight memory must be width sharded")
        # per_core_N sizes output storage shards, independently of the input
        # shards and DRAM readers. Matmul places these on the full compute grid;
        # per_core_M covers the padded batch, so only one row of blocks is needed.
        output_cores = math.ceil(physical_widths[split_index] / (program_config.per_core_N * TILE_SIZE))
        grid = config.mesh_device.compute_with_storage_grid_size()
        available_cores = grid.x * grid.y
        if output_cores > available_cores:
            raise ValueError(
                f"LMHead1D split {split_index} program requires {output_cores} output storage cores, "
                f"but the device compute grid has only {available_cores}"
            )


def resolve_lm_head_1d_arch_config(config: LMHead1DConfig) -> LMHead1DConfig:
    """Return an independent, fully resolved config after one architecture query."""
    mesh_device = _derive_lm_head_mesh_device(config)
    arch = mesh_device.arch()

    if arch not in (ttnn.device.Arch.WORMHOLE_B0, ttnn.device.Arch.BLACKHOLE):
        raise ValueError(f"Unsupported LMHead1D architecture: {arch}")
    requested_compute = config.compute_kernel_config or _compute_kernel_config_hifi2(arch)
    try:
        compute_kernel_config = ttnn.init_device_compute_kernel_config(arch, requested_compute)
    except (TypeError, ValueError) as error:
        raise ValueError(f"Invalid LMHead1D compute recipe for {arch}: {error}") from error

    _validate_lm_head_compute_kernel_config(compute_kernel_config)
    resolved_common = _resolve_lm_head_1d_config(config)
    resolved_common = replace(resolved_common, compute_kernel_config=compute_kernel_config)
    _validate_lm_head_program_configs(resolved_common)
    return resolved_common


def _resolve_lm_head_1d_config(config: LMHead1DConfig) -> LMHead1DConfig:
    """Resolve defaults for LMHead1DConfig."""
    if not config.output_weights:
        raise ValueError("LMHead1D requires at least one output weight")
    to_set = {}

    # Mesh device
    mesh_device = _derive_lm_head_mesh_device(config)
    if config.mesh_device is None:
        to_set["mesh_device"] = mesh_device

    # Dim
    dim = config.dim
    if dim is None:
        dim = config.output_weights[0].source.shape[-2]
        to_set["dim"] = dim

    # Output memcfg
    if config.output_memcfg is None:
        to_set["output_memcfg"] = ttnn.L1_MEMORY_CONFIG

    # Input memcfg
    if config.input_memcfg is None:
        to_set["input_memcfg"] = ttnn.DRAM_MEMORY_CONFIG

    # Program configs
    num_devices = mesh_device.get_num_devices()

    if config.program_configs is None:
        # Use None program configs for auto-resolve (let ttnn.linear auto-select).
        # DRAM-sharded program configs require matching DRAM-sharded weight memory configs,
        # which are only set up correctly via from_model_args.
        pcs = [None for _ in config.output_weights]
        to_set["program_configs"] = pcs

    if config.output_split_sizes is None:
        split_sizes = []
        for split_index, weight in enumerate(config.output_weights):
            combined_width = weight.source.shape[-1]
            if combined_width % num_devices:
                raise ValueError(
                    f"LMHead1D split {split_index} combined physical width {combined_width} "
                    f"must be divisible by {num_devices} devices"
                )
            split_sizes.append(combined_width // num_devices)
        to_set["output_split_sizes"] = split_sizes
    elif len(config.output_split_sizes) != len(config.output_weights):
        raise ValueError("LMHead1D requires exactly one logical output split size per output weight")

    # Weight memory configs + resolve LazyWeights
    if config.weights_memcfgs is None:
        # Use regular DRAM for auto-resolve (DRAM sharded requires DRAM-core-aligned padding
        # which is handled by from_model_args when it explicitly provides weights_memcfgs)
        memcfgs = [ttnn.DRAM_MEMORY_CONFIG for _ in config.output_weights]
        to_set["weights_memcfgs"] = memcfgs

    weights_memcfgs = (
        config.weights_memcfgs if config.weights_memcfgs is not None else to_set.get("weights_memcfgs", [])
    )

    resolved_weights = []
    for i, w in enumerate(config.output_weights):
        mem_cfg = weights_memcfgs[i] if i < len(weights_memcfgs) else ttnn.DRAM_MEMORY_CONFIG
        resolved_weights.append(
            resolve_lazy_weight(
                w,
                device=mesh_device,
                memory_config=mem_cfg,
                mesh_mapper_config=ttnn.MeshMapperConfig(
                    placements=[ttnn.PlacementShard(-1)],
                    mesh_shape_override=ttnn.MeshShape([num_devices]),
                ),
                layout=ttnn.TILE_LAYOUT,
                dtype=config.lm_head_dtype,
            )
        )
    to_set["output_weights"] = resolved_weights

    resolved = replace(config, **to_set)
    return resolved


def _load_input_device_tensor(x: ttnn.Tensor | LazyWeight, config: LMHead1DConfig) -> ttnn.Tensor:
    """Resolve input tensor using config's input_memcfg."""
    if isinstance(x, LazyWeight):
        resolved_x = resolve_lazy_weight(
            x,
            device=config.mesh_device,
            memory_config=config.input_memcfg,
            mesh_mapper_config=None,
            layout=ttnn.TILE_LAYOUT,
        )
        return resolved_x.get_device_weight()
    assert isinstance(x, ttnn.Tensor)
    return x


# =============================================================================
# Config helper functions (adapted from TTTv1 model_config.py)
# =============================================================================


def _create_dram_sharded_mem_config(
    k: int, n: int, dram_grid: ttnn.CoreRangeSet, tile_size: int = TILE_SIZE, dram_cores: int = 12
) -> ttnn.MemoryConfig:
    padded_size = math.ceil(n / (tile_size * dram_cores)) * (tile_size * dram_cores)
    shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_size // dram_cores), ttnn.ShardOrientation.ROW_MAJOR)
    return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec)

	# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
	# SPDX-License-Identifier: Apache-2.0

	"""
	TTTv2-style LM Head module for 1D-topology devices: N150 (1x1), N300 (1x2), T3K (1x8).

	Computes logits over the vocabulary by splitting the output projection into
	weight chunks that fit in L1, running linear ops per chunk, and concatenating.

	Each device holds the correct logits for its vocab shard (column-parallel linear,
	no partial sums). The caller handles all_gather for multi-device argmax.

	Execution path:
	for each (weight, pc): linear(x, weight) → sharded_to_interleaved → append
	→ concat
	"""

	import math
	from dataclasses import dataclass, replace
	from pathlib import Path
	from typing import List

	import ttnn
	from models.common.lightweightmodule import LightweightModule
	from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight
	from models.common.tensor_utils import TILE_SIZE
	from models.common.tensor_utils import nearest_32 as _nearest_32

	# =============================================================================
	# Config dataclass
	# =============================================================================


	@dataclass
	class LMHead1DConfig:
	"""
	Configuration for LMHead1D.

	Simple usage (pre-split weights):
	config = LMHead1DConfig(output_weights=[w1, w2, w3])

	The program_configs and other fields are auto-computed from weights if None.
	"""

	# Required: output projection weights (already split for L1 fit)
	output_weights: List[LazyWeight]

	# Optional: device
	mesh_device: ttnn.MeshDevice \| None = None

	# Optional: derived from weights if None
	dim: int \| None = None

	# Optional: batch/tile config
	max_batch_size: int = 32

	# Optional: power-user overrides
	program_configs: List \| None = None
	compute_kernel_config: ttnn.DeviceComputeKernelConfig \| None = None
	# Logical output width for each split on one device. A physical weight shard
	# may be tile-padded beyond this width; forward trims it before concat.
	output_split_sizes: List[int] \| None = None
	lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b
	output_memcfg: ttnn.MemoryConfig \| None = None

	# Input memory config (None = DRAM interleaved for simple API, width-sharded for from_model_args)
	input_memcfg: ttnn.MemoryConfig \| None = None

	# Weight memory configs (None = auto-compute)
	weights_memcfgs: List[ttnn.MemoryConfig] \| None = None

	def is_resolved(self) -> bool:
	return all(getattr(self, f) is not None for f in self.__dataclass_fields__)


	# =============================================================================
	# LMHead1D
	# =============================================================================


	class LMHead1D(LightweightModule):
	"""
	LM Head for non-TG (1D) devices.

	Splits vocabulary projection into L1-sized chunks, runs linear per chunk,
	and concatenates. Each device holds correct logits for its vocab shard;
	caller handles all_gather for multi-device argmax.

	Simple API:
	lm_head = LMHead1D(output_weights=[w1, w2])

	Power API:
	config = LMHead1DConfig(output_weights=[w1, w2], lm_head_dtype=ttnn.bfloat16)
	lm_head = LMHead1D.from_config(config)

	Execution path:
	for each (w, pc): linear(x, w) → sharded_to_interleaved → concat
	"""

	def __init__(self, output_weights: List[LazyWeight]):
	super().__init__()
	self.config = resolve_lm_head_1d_arch_config(LMHead1DConfig(output_weights=output_weights))
	self._device_weights_loaded = False

	@classmethod
	def from_config(cls, config: LMHead1DConfig):
	instance = object.__new__(cls)
	super(LMHead1D, instance).__init__()
	instance.config = resolve_lm_head_1d_arch_config(config)
	instance._device_weights_loaded = False
	return instance

	def load_device_weights(self):
	if self._device_weights_loaded:
	return
	self.output_weights = [w.get_device_weight() for w in self.config.output_weights]
	self._device_weights_loaded = True

	def forward(self, x: ttnn.Tensor \| LazyWeight) -> ttnn.Tensor:
	"""
	Compute logits over vocabulary.

	Args:
	x: Input hidden states, shape [1, 1, batch_rows, dim].

	Returns:
	Logits tensor, shape [1, 1, batch_rows, padded_vocab / num_devices] per device.
	"""
	self.load_device_weights()
	x = _load_input_device_tensor(x, self.config)
	cfg = self.config

	outputs = []
	for weight, pc, logical_width in zip(self.output_weights, cfg.program_configs, cfg.output_split_sizes):
	if pc is not None:
	# DRAM-sharded path (from_model_args): width-sharded output → interleaved
	output = ttnn.linear(
	x,
	weight,
	compute_kernel_config=cfg.compute_kernel_config,
	program_config=pc,
	memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG,
	dtype=cfg.lm_head_dtype,
	)
	output = ttnn.sharded_to_interleaved(output, memory_config=cfg.output_memcfg)
	else:
	# Auto path (simple API): interleaved output directly
	output = ttnn.linear(
	x,
	weight,
	compute_kernel_config=cfg.compute_kernel_config,
	memory_config=cfg.output_memcfg,
	dtype=cfg.lm_head_dtype,
	)
	if output.shape[-1] != logical_width:
	output = ttnn.slice(
	output,
	(0, 0, 0, 0),
	(output.shape[0], output.shape[1], output.shape[2], logical_width),
	)
	outputs.append(output)

	# Concatenate splits
	output = ttnn.concat(outputs, dim=-1, memory_config=cfg.output_memcfg)

	return output

	# [INFO] this is the entry point for TTTv1 model_config.py and will retire with TTTv1
	@classmethod
	def from_model_args(
	cls,
	mesh_device,
	args,
	state_dict,
	state_dict_prefix,
	weight_cache_path,
	max_columns_per_device,
	dtype=None,
	model_config=None,
	tt_ccl=None,
	):
	"""Factory method for backward compatibility with ModelArgs.

	Note: tt_ccl is accepted for signature compatibility with TTTv1 LMHead
	but is not used -- 1D LMHead does not need CCL (column-parallel, no partial sums).
	"""
	if args.is_galaxy:
	raise ValueError("LMHead1D cannot be used for Galaxy devices.")

	import torch

	vocab_size = args.vocab_size
	num_devices = mesh_device.get_num_devices()
	dim = args.dim
	padded_vocab_size = math.ceil(vocab_size / (TILE_SIZE * num_devices)) * (TILE_SIZE * num_devices)
	size_per_device = padded_vocab_size // num_devices
	num_splits = math.ceil(size_per_device / max_columns_per_device)
	split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1)
	split_sizes.append(size_per_device - sum(split_sizes))

	# Build output weights
	source_weight = state_dict[f"{state_dict_prefix}output.weight"]
	if tuple(source_weight.shape) != (vocab_size, dim):
	raise ValueError(
	f"LMHead1D output weight must have shape {(vocab_size, dim)}, got {tuple(source_weight.shape)}"
	)
	torch_output_weights = source_weight.permute(1, 0)
	if vocab_size < padded_vocab_size:
	padding_size = padded_vocab_size - vocab_size
	torch_output_weights = torch.cat(
	[
	torch_output_weights,
	torch.zeros(torch_output_weights.shape[0], padding_size, dtype=torch_output_weights.dtype),
	],
	dim=-1,
	)

	# DRAM grid for weight memory configs
	dram_size = mesh_device.dram_grid_size()
	dram_grid = ttnn.CoreRangeSet(
	{ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))}
	)

	cache_dir = None if args.dummy_weights else Path(weight_cache_path) / "lm_head"

	output_weights = []
	weights_memcfgs = []
	for i, split_size in enumerate(split_sizes):
	device_splits = []
	physical_split_size = math.ceil(split_size / TILE_SIZE) * TILE_SIZE
	for device_idx in range(num_devices):
	start = device_idx * size_per_device + sum(split_sizes[:i])
	end = start + split_size
	device_split = torch_output_weights[:, start:end]
	if split_size < physical_split_size:
	device_split = torch.cat(
	[
	device_split,
	torch.zeros(dim, physical_split_size - split_size, dtype=device_split.dtype),
	],
	dim=-1,
	)
	device_splits.append(device_split)
	combined_split = torch.cat(device_splits, dim=-1)

	mem_cfg = _create_dram_sharded_mem_config(
	k=dim,
	n=math.ceil(combined_split.shape[-1] / num_devices),
	dram_grid=dram_grid,
	tile_size=TILE_SIZE,
	dram_cores=dram_size.x,
	)
	weights_memcfgs.append(mem_cfg)

	w_dtype = dtype if dtype is not None else ttnn.bfloat8_b
	output_weights.append(
	LazyWeight(
	source=combined_split,
	dtype=w_dtype,
	device=mesh_device,
	mesh_mapper_config=ttnn.MeshMapperConfig(
	placements=[ttnn.PlacementShard(-1)],
	mesh_shape_override=ttnn.MeshShape([num_devices]),
	),
	layout=ttnn.TILE_LAYOUT,
	memory_config=mem_cfg,
	cache_dir_weight_name=(
	(cache_dir, f"output_split_{i}_logical_{split_size}_physical_{combined_split.shape[-1]}")
	if cache_dir
	else None
	),
	)
	)

	# Program configs - use the args.lm_head_core_grid which TTTv1 carefully computes
	tile_padded_batch_rows = TILE_SIZE * math.ceil(args.max_batch_size / TILE_SIZE)
	lm_head_core_grid = args.lm_head_core_grid
	program_configs = [
	args.dram_matmul_config(
	tile_padded_batch_rows,
	dim,
	ss,
	lm_head_core_grid.num_cores,
	)
	for ss in split_sizes
	]

	# Get input memory config: width-sharded across lm_head_core_grid
	# (DRAM-sharded matmul requires width-sharded input in L1)
	input_memcfg = ttnn.create_sharded_memory_config(
	(
	tile_padded_batch_rows,
	_nearest_32(dim // lm_head_core_grid.num_cores),
	),
	lm_head_core_grid,
	ttnn.ShardStrategy.WIDTH,
	ttnn.ShardOrientation.ROW_MAJOR,
	use_height_and_width_as_shard_shape=True,
	)

	config = LMHead1DConfig(
	output_weights=output_weights,
	mesh_device=mesh_device,
	dim=dim,
	max_batch_size=args.max_batch_size,
	program_configs=program_configs,
	output_split_sizes=split_sizes,
	lm_head_dtype=getattr(args, "lm_head_dtype", ttnn.bfloat8_b),
	output_memcfg=ttnn.L1_MEMORY_CONFIG,
	input_memcfg=input_memcfg,
	weights_memcfgs=weights_memcfgs,
	)
	return cls.from_config(config)


	# =============================================================================
	# Config resolution
	# =============================================================================


	def _compute_kernel_config_hifi2(arch) -> ttnn.DeviceComputeKernelConfig:
	"""Construct the TTTv1 LM-head recipe for the selected architecture."""
	return ttnn.init_device_compute_kernel_config(
	arch,
	math_fidelity=ttnn.MathFidelity.HiFi2,
	math_approx_mode=False,
	fp32_dest_acc_en=False,
	packer_l1_acc=True,
	)


	def _validate_lm_head_compute_kernel_config(compute_kernel_config) -> None:
	if compute_kernel_config is None:
	raise ValueError("LMHead1D architecture config requires compute_kernel_config")
	if not isinstance(compute_kernel_config, ttnn.WormholeComputeKernelConfig):
	raise TypeError(
	"LMHead1D compute_kernel_config must be a concrete TTNN compute kernel config, "
	f"got {type(compute_kernel_config).__name__}"
	)


	def _derive_lm_head_mesh_device(config: LMHead1DConfig) -> ttnn.MeshDevice:
	if not config.output_weights:
	raise ValueError("LMHead1D requires at least one output weight")
	mesh_device = config.mesh_device
	if mesh_device is None:
	mesh_device = config.output_weights[0].device
	if mesh_device is None:
	mesh_device = ttnn.GetDefaultDevice()
	if mesh_device is None:
	raise ValueError("LMHead1D mesh_device must be available")
	for split_index, weight in enumerate(config.output_weights):
	if weight.device is not None and weight.device is not mesh_device:
	raise ValueError(f"LMHead1D mesh_device must match output weight {split_index} device")
	return mesh_device


	def _validate_lm_head_program_configs(config: LMHead1DConfig) -> None:
	num_splits = len(config.output_weights)
	if num_splits == 0:
	raise ValueError("LMHead1D requires at least one output weight")
	if config.program_configs is None or len(config.program_configs) != num_splits:
	raise ValueError("LMHead1D requires exactly one program config per output-weight split")
	if config.weights_memcfgs is None or len(config.weights_memcfgs) != num_splits:
	raise ValueError("LMHead1D requires exactly one weight memory config per output-weight split")
	if config.output_split_sizes is None or len(config.output_split_sizes) != num_splits:
	raise ValueError("LMHead1D requires exactly one logical output width per output-weight split")

	num_devices = config.mesh_device.get_num_devices()
	physical_widths = []
	for split_index, (weight, logical_width) in enumerate(zip(config.output_weights, config.output_split_sizes)):
	source_shape = weight.source.shape
	if len(source_shape) < 2 or source_shape[-2] != config.dim:
	raise ValueError(
	f"LMHead1D split {split_index} weight K dimension must equal dim={config.dim}, got {source_shape}"
	)
	combined_width = source_shape[-1]
	if combined_width % num_devices:
	raise ValueError(
	f"LMHead1D split {split_index} combined physical width {combined_width} "
	f"must be divisible by {num_devices} devices"
	)
	physical_width = combined_width // num_devices
	if physical_width % TILE_SIZE:
	raise ValueError(
	f"LMHead1D split {split_index} physical per-device width {physical_width} must be tile aligned"
	)
	if not isinstance(logical_width, int) or logical_width <= 0 or logical_width > physical_width:
	raise ValueError(
	f"LMHead1D split {split_index} logical width must satisfy 0 < logical <= {physical_width}, "
	f"got {logical_width}"
	)
	physical_widths.append(physical_width)

	has_explicit_program = any(program_config is not None for program_config in config.program_configs)
	if has_explicit_program:
	if config.input_memcfg is None or not config.input_memcfg.is_sharded():
	raise ValueError("LMHead1D DRAM-sharded programs require a sharded input_memcfg")
	if config.input_memcfg.buffer_type != ttnn.BufferType.L1:
	raise ValueError("LMHead1D DRAM-sharded programs require input_memcfg in L1")
	if config.input_memcfg.memory_layout != ttnn.TensorMemoryLayout.WIDTH_SHARDED:
	raise ValueError("LMHead1D DRAM-sharded programs require width-sharded input_memcfg")

	for split_index, program_config in enumerate(config.program_configs):
	if program_config is None:
	continue
	if not isinstance(program_config, ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig):
	raise TypeError(
	"LMHead1D explicit program configs must use the DRAM-sharded matmul path; "
	f"split {split_index} has {type(program_config).__name__}"
	)
	if program_config.in0_block_w <= 0 or program_config.per_core_M <= 0 or program_config.per_core_N <= 0:
	raise ValueError(f"LMHead1D split {split_index} program dimensions must be positive")
	dim_tiles = config.dim // TILE_SIZE
	if dim_tiles % program_config.in0_block_w != 0:
	raise ValueError(f"LMHead1D split {split_index} dim tiles must be divisible by in0_block_w")
	expected_per_core_m = math.ceil(config.max_batch_size / TILE_SIZE)
	if program_config.per_core_M != expected_per_core_m:
	raise ValueError(
	f"LMHead1D split {split_index} per_core_M={program_config.per_core_M} "
	f"does not match padded batch tiles {expected_per_core_m}"
	)

	weight_memcfg = config.weights_memcfgs[split_index]
	if not weight_memcfg.is_sharded():
	raise ValueError(f"LMHead1D split {split_index} DRAM-sharded program requires sharded weight memory")
	if weight_memcfg.buffer_type != ttnn.BufferType.DRAM:
	raise ValueError(f"LMHead1D split {split_index} weight memory must reside in DRAM")
	if weight_memcfg.memory_layout != ttnn.TensorMemoryLayout.WIDTH_SHARDED:
	raise ValueError(f"LMHead1D split {split_index} weight memory must be width sharded")
	# per_core_N sizes output storage shards, independently of the input
	# shards and DRAM readers. Matmul places these on the full compute grid;
	# per_core_M covers the padded batch, so only one row of blocks is needed.
	output_cores = math.ceil(physical_widths[split_index] / (program_config.per_core_N * TILE_SIZE))
	grid = config.mesh_device.compute_with_storage_grid_size()
	available_cores = grid.x * grid.y
	if output_cores > available_cores:
	raise ValueError(
	f"LMHead1D split {split_index} program requires {output_cores} output storage cores, "
	f"but the device compute grid has only {available_cores}"
	)


	def resolve_lm_head_1d_arch_config(config: LMHead1DConfig) -> LMHead1DConfig:
	"""Return an independent, fully resolved config after one architecture query."""
	mesh_device = _derive_lm_head_mesh_device(config)
	arch = mesh_device.arch()

	if arch not in (ttnn.device.Arch.WORMHOLE_B0, ttnn.device.Arch.BLACKHOLE):
	raise ValueError(f"Unsupported LMHead1D architecture: {arch}")
	requested_compute = config.compute_kernel_config or _compute_kernel_config_hifi2(arch)
	try:
	compute_kernel_config = ttnn.init_device_compute_kernel_config(arch, requested_compute)
	except (TypeError, ValueError) as error:
	raise ValueError(f"Invalid LMHead1D compute recipe for {arch}: {error}") from error

	_validate_lm_head_compute_kernel_config(compute_kernel_config)
	resolved_common = _resolve_lm_head_1d_config(config)
	resolved_common = replace(resolved_common, compute_kernel_config=compute_kernel_config)
	_validate_lm_head_program_configs(resolved_common)
	return resolved_common


	def _resolve_lm_head_1d_config(config: LMHead1DConfig) -> LMHead1DConfig:
	"""Resolve defaults for LMHead1DConfig."""
	if not config.output_weights:
	raise ValueError("LMHead1D requires at least one output weight")
	to_set = {}

	# Mesh device
	mesh_device = _derive_lm_head_mesh_device(config)
	if config.mesh_device is None:
	to_set["mesh_device"] = mesh_device

	# Dim
	dim = config.dim
	if dim is None:
	dim = config.output_weights[0].source.shape[-2]
	to_set["dim"] = dim

	# Output memcfg
	if config.output_memcfg is None:
	to_set["output_memcfg"] = ttnn.L1_MEMORY_CONFIG

	# Input memcfg
	if config.input_memcfg is None:
	to_set["input_memcfg"] = ttnn.DRAM_MEMORY_CONFIG

	# Program configs
	num_devices = mesh_device.get_num_devices()

	if config.program_configs is None:
	# Use None program configs for auto-resolve (let ttnn.linear auto-select).
	# DRAM-sharded program configs require matching DRAM-sharded weight memory configs,
	# which are only set up correctly via from_model_args.
	pcs = [None for _ in config.output_weights]
	to_set["program_configs"] = pcs

	if config.output_split_sizes is None:
	split_sizes = []
	for split_index, weight in enumerate(config.output_weights):
	combined_width = weight.source.shape[-1]
	if combined_width % num_devices:
	raise ValueError(
	f"LMHead1D split {split_index} combined physical width {combined_width} "
	f"must be divisible by {num_devices} devices"
	)
	split_sizes.append(combined_width // num_devices)
	to_set["output_split_sizes"] = split_sizes
	elif len(config.output_split_sizes) != len(config.output_weights):
	raise ValueError("LMHead1D requires exactly one logical output split size per output weight")

	# Weight memory configs + resolve LazyWeights
	if config.weights_memcfgs is None:
	# Use regular DRAM for auto-resolve (DRAM sharded requires DRAM-core-aligned padding
	# which is handled by from_model_args when it explicitly provides weights_memcfgs)
	memcfgs = [ttnn.DRAM_MEMORY_CONFIG for _ in config.output_weights]
	to_set["weights_memcfgs"] = memcfgs

	weights_memcfgs = (
	config.weights_memcfgs if config.weights_memcfgs is not None else to_set.get("weights_memcfgs", [])
	)

	resolved_weights = []
	for i, w in enumerate(config.output_weights):
	mem_cfg = weights_memcfgs[i] if i < len(weights_memcfgs) else ttnn.DRAM_MEMORY_CONFIG
	resolved_weights.append(
	resolve_lazy_weight(
	w,
	device=mesh_device,
	memory_config=mem_cfg,
	mesh_mapper_config=ttnn.MeshMapperConfig(
	placements=[ttnn.PlacementShard(-1)],
	mesh_shape_override=ttnn.MeshShape([num_devices]),
	),
	layout=ttnn.TILE_LAYOUT,
	dtype=config.lm_head_dtype,
	)
	)
	to_set["output_weights"] = resolved_weights

	resolved = replace(config, **to_set)
	return resolved


	def _load_input_device_tensor(x: ttnn.Tensor \| LazyWeight, config: LMHead1DConfig) -> ttnn.Tensor:
	"""Resolve input tensor using config's input_memcfg."""
	if isinstance(x, LazyWeight):
	resolved_x = resolve_lazy_weight(
	x,
	device=config.mesh_device,
	memory_config=config.input_memcfg,
	mesh_mapper_config=None,
	layout=ttnn.TILE_LAYOUT,
	)
	return resolved_x.get_device_weight()
	assert isinstance(x, ttnn.Tensor)
	return x


	# =============================================================================
	# Config helper functions (adapted from TTTv1 model_config.py)
	# =============================================================================


	def _create_dram_sharded_mem_config(
	k: int, n: int, dram_grid: ttnn.CoreRangeSet, tile_size: int = TILE_SIZE, dram_cores: int = 12
	) -> ttnn.MemoryConfig:
	padded_size = math.ceil(n / (tile_size * dram_cores)) * (tile_size * dram_cores)
	shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_size // dram_cores), ttnn.ShardOrientation.ROW_MAJOR)
	return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec)