# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """ RMSNorm2D: RMSNorm for 2D mesh topologies (TG, Galaxy). Execution paths: - decode_forward - Sharded, Linear topology, cluster_axis=1 - prefill_forward - Interleaved, Linear topology, cluster_axis=1 Key design: - Both decode and prefill are ALWAYS distributed for 2D mesh - Uses Linear topology with cluster_axis=1 (gather across columns) - Decode uses sharded configs, Prefill uses interleaved - Separate weight for distributed path (sharded across columns, replicated across rows) """ from dataclasses import dataclass, replace from typing import Optional import ttnn from models.common.lightweightmodule import LightweightModule from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight from models.common.modules.tt_ccl import get_tt_ccl from models.common.tensor_utils import TILE_SIZE # ============================================================================= # Constants # ============================================================================= SHARD_HEIGHT = TILE_SIZE # ============================================================================= # RMSNorm2DConfig # ============================================================================= @dataclass class RMSNorm2DConfig: """ Configuration for RMSNorm2D - only fields relevant to 2D mesh topologies. Both paths are always distributed with Linear topology and cluster_axis=1. Paths: - Path 4: 2D distributed decode (sharded) - requires decode_* configs - Path 5: 2D distributed prefill (interleaved) - no program_config needed Simple usage: config = RMSNorm2DConfig(weight, cluster_shape=(4, 8)) Override any field: config = RMSNorm2DConfig(weight, cluster_shape=(4, 8), eps=1e-6) """ # Required: weight and cluster_shape weight: LazyWeight cluster_shape: tuple[int, int] | None = None # (rows, cols), e.g. (4, 8) for TG # Normalization settings eps: float = 1e-5 add_unit_offset: bool = False # Device and collectives mesh_device: ttnn.MeshDevice | None = None tt_ccl: "TT_CCL | None" = None # type: ignore # Batch size for decode max_batch_size: int = 32 # Input memory configs (for _load_input_device_tensor_2d) decode_input_memcfg: ttnn.MemoryConfig | None = None prefill_input_memcfg: ttnn.MemoryConfig | None = None # 2D distributed decode configs - sharded decode_progcfg: ttnn.LayerNormShardedMultiCoreProgramConfig | None = None decode_stats_memcfg: ttnn.MemoryConfig | None = None # Compute kernel config (only for prefill - decode uses program_config) compute_kernel_config_prefill: ttnn.WormholeComputeKernelConfig | None = None def is_resolved(self) -> bool: """Check if all required fields are resolved.""" required = [ "mesh_device", "cluster_shape", "tt_ccl", # Decode sharded configs "decode_input_memcfg", "decode_progcfg", "decode_stats_memcfg", # Compute config "compute_kernel_config_prefill", # Weight (2D always uses distributed/sharded) "weight", ] return all(getattr(self, f) is not None for f in required) # ============================================================================= # RMSNorm2D # ============================================================================= class RMSNorm2D(LightweightModule): """ RMSNorm for 2D mesh topologies (TG, Galaxy). Both decode and prefill are always distributed with Linear topology. Execution paths (bound at construction): - decode_forward -> _decode_2d_distributed (Path 4) - sharded - prefill_forward -> _prefill_2d_distributed (Path 5) - interleaved Simple API: norm = RMSNorm2D(weight) Power API: config = RMSNorm2DConfig(weight, cluster_shape=(4, 8), max_batch_size=64) norm = RMSNorm2D.from_config(config) """ def __init__(self, weight: LazyWeight): """ Simple API - derives all config from weight. Args: weight: RMSNorm weight tensor of shape (dim,) or (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) Note: Use from_config() to customize eps, add_unit_offset, cluster_shape, or other settings. """ super().__init__() self.config = _resolve_2d_config(RMSNorm2DConfig(weight=weight)) self._device_weights_loaded = False @classmethod def from_config(cls, config: RMSNorm2DConfig): """Power API - any level of customization via config.""" instance = object.__new__(cls) super(RMSNorm2D, instance).__init__() instance.config = _resolve_2d_config(config) instance._device_weights_loaded = False return instance def load_device_weights(self): """Load weights to device lazily.""" if self._device_weights_loaded: return assert self.config.is_resolved(), "config must be resolved before loading device weights!" # Load weight (sharded across columns, replicated across rows) self.weight = self.config.weight.get_device_weight() self._device_weights_loaded = True def decode_forward(self, x: "ttnn.Tensor | LazyWeight") -> ttnn.Tensor: """ 2D distributed decode - sharded for TG/Galaxy. Uses Linear topology, cluster_axis=1, with program_config. Execution: to_sharded -> rms_norm_pre_all_gather(sharded) -> all_gather(cluster_axis=1) -> rms_norm_post_all_gather(sharded) """ self.load_device_weights() x = _load_input_device_tensor_2d(x, self.config, mode="decode") cfg = self.config # Convert to sharded memory config x = ttnn.to_memory_config(x, memory_config=cfg.decode_input_memcfg) # Run distributed rmsnorm part 1 (sharded) tt_stats = ttnn.rms_norm_pre_all_gather(x, program_config=cfg.decode_progcfg) # All gather stats along cluster axis 1 (columns) cluster_axis = 1 tt_stats = ttnn.experimental.all_gather_async( tt_stats, persistent_output_buffer=None, dim=3, multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis), num_links=1, cluster_axis=cluster_axis, topology=ttnn.Topology.Linear, memory_config=cfg.decode_stats_memcfg, barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), chunks_per_sync=10, num_workers_per_link=2, num_buffers_per_channel=2, ) # Run distributed rmsnorm part 2 tt_out = ttnn.rms_norm_post_all_gather( x, epsilon=cfg.eps, weight=self.weight, program_config=cfg.decode_progcfg, stats=tt_stats, ) tt_stats.deallocate(True) return tt_out def prefill_forward(self, x: "ttnn.Tensor | LazyWeight") -> ttnn.Tensor: """ 2D distributed prefill - interleaved for TG/Galaxy. Uses Linear topology, cluster_axis=1. Execution: rms_norm_pre_all_gather -> reshape -> all_gather(cluster_axis=1) -> rms_norm_post_all_gather """ self.load_device_weights() x = _load_input_device_tensor_2d(x, self.config, mode="prefill") cfg = self.config # Run distributed rmsnorm part 1 tt_stats = ttnn.rms_norm_pre_all_gather( x, compute_kernel_config=cfg.compute_kernel_config_prefill, dtype=ttnn.bfloat16 ) # Reshape stats for all_gather (preserve batch dimension for multi-batch prefill) batch = x.shape[0] padded_shape = (batch, 1, x.shape[-2], 32) tt_stats = ttnn.reshape(tt_stats, ttnn.Shape(padded_shape)) # All gather stats along cluster axis 1 (columns) cluster_axis = 1 tt_stats_gathered = ttnn.experimental.all_gather_async( tt_stats, persistent_output_buffer=None, dim=3, multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis), num_links=1, cluster_axis=cluster_axis, topology=ttnn.Topology.Linear, memory_config=ttnn.DRAM_MEMORY_CONFIG, barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), chunks_per_sync=10, num_workers_per_link=2, num_buffers_per_channel=2, ) tt_stats.deallocate(True) # Run distributed rmsnorm part 2 tt_out = ttnn.rms_norm_post_all_gather( x, tt_stats_gathered, epsilon=cfg.eps, weight=self.weight, compute_kernel_config=cfg.compute_kernel_config_prefill, ) tt_stats_gathered.deallocate(True) return tt_out # ========================================================================= # Forward dispatcher # ========================================================================= def forward(self, x: "ttnn.Tensor | LazyWeight", mode: str) -> ttnn.Tensor: """Dispatch to decode_forward or prefill_forward based on mode.""" if mode == "decode": return self.decode_forward(x) else: return self.prefill_forward(x) # ========================================================================= # Factory method for TTTv1 compatibility # ========================================================================= @classmethod def from_model_args( cls, mesh_device, tt_ccl, args, state_dict, weight_cache_path, layer_num: int, weight_key: str, state_dict_prefix: Optional[str] = None, ): """Factory method for backward compatibility with ModelArgs.""" # Validate TG topology cluster_shape = tuple(mesh_device.shape) valid_shapes = [(4, 8), (8, 4)] if cluster_shape not in valid_shapes: raise ValueError( f"RMSNorm2D requires Galaxy topology (4x8 or 8x4), got {cluster_shape}. " "Use RMSNorm1D for non-Galaxy devices." ) # Build weight name from state_dict_prefix and weight_key if state_dict_prefix: weight_name = f"{state_dict_prefix}{weight_key}.weight" else: if layer_num is None: weight_name = f"{weight_key}.weight" else: weight_name = f"layers.{layer_num}.{weight_key}.weight" # Transform weight to expected shape: (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) dim = args.dim torch_weight = ( state_dict[weight_name].unsqueeze(0).view(1, 1, dim).reshape([1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT]) ) # Add offset before creating LazyWeight if args.rms_norm_add_unit_offset: torch_weight = torch_weight + 1.0 # Create LazyWeight (mesh_mapper_config is handled by _resolve_2d_config) lazy_weight = LazyWeight( source=torch_weight, dtype=ttnn.bfloat16, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG, cache_dir_weight_name=(weight_cache_path, weight_name + "_2d") if weight_cache_path else None, ) # Build config (dim derived from weight in _resolve_2d_config) config = RMSNorm2DConfig( weight=lazy_weight, cluster_shape=cluster_shape, eps=args.norm_eps, add_unit_offset=False, # Already applied above mesh_device=mesh_device, tt_ccl=tt_ccl, max_batch_size=args.max_batch_size, ) return cls.from_config(config) # ============================================================================= # Config resolution # ============================================================================= def _resolve_2d_config(config: RMSNorm2DConfig) -> RMSNorm2DConfig: """Resolve config defaults for RMSNorm2D.""" to_set = {} # --- Phase 1: Foundational fields --- # Derive dim from weight (works for any shape: [dim], [1, dim], [1, 1, dim//32, 32], etc.) dim = config.weight.source.numel() # Derive mesh_device from weight mesh_device = config.mesh_device if mesh_device is None: mesh_device = config.weight.device if mesh_device is None: mesh_device = ttnn.GetDefaultDevice() if config.mesh_device is None: to_set["mesh_device"] = mesh_device assert mesh_device is not None, "mesh_device must be available!" # Derive cluster_shape from mesh_device if not provided cluster_shape = config.cluster_shape if cluster_shape is None: cluster_shape = tuple(mesh_device.shape) to_set["cluster_shape"] = cluster_shape assert len(cluster_shape) == 2 and all( d > 1 for d in cluster_shape ), f"cluster_shape must be 2D with both dims > 1, got {cluster_shape}" num_rows, num_cols = cluster_shape # Derive tt_ccl (always needed for 2D) if config.tt_ccl is None: to_set["tt_ccl"] = get_tt_ccl(mesh_device) # --- Phase 2: Compute kernel config (prefill only - decode uses program_config) --- if config.compute_kernel_config_prefill is None: # 2D uses fp32=False, packer_l1=False (from DistributedNorm.ln_cfg) to_set["compute_kernel_config_prefill"] = ttnn.WormholeComputeKernelConfig( math_fidelity=ttnn.MathFidelity.HiFi2, math_approx_mode=False, fp32_dest_acc_en=False, packer_l1_acc=False, ) # --- Phase 3: 2D distributed decode configs (Path 4) --- hidden_size_per_device = dim // num_cols core_grid_ln = ( min(4, hidden_size_per_device // TILE_SIZE // 8), 8, ) num_cores_ln = core_grid_ln[0] * core_grid_ln[1] if config.decode_input_memcfg is None: to_set["decode_input_memcfg"] = ttnn.create_sharded_memory_config( shape=(1, 1, 32, hidden_size_per_device), core_grid=ttnn.CoreGrid(y=core_grid_ln[0], x=core_grid_ln[1]), strategy=ttnn.ShardStrategy.WIDTH, ) if config.decode_progcfg is None: to_set["decode_progcfg"] = ttnn.LayerNormShardedMultiCoreProgramConfig( compute_with_storage_grid_size=(core_grid_ln[1], core_grid_ln[0]), subblock_w=(hidden_size_per_device // num_cores_ln) // TILE_SIZE, block_h=1, block_w=(hidden_size_per_device // num_cores_ln) // TILE_SIZE, inplace=False, ) if config.decode_stats_memcfg is None: # Stats memory: 32 x (32 * num_cols) where num_cols is the cluster width to_set["decode_stats_memcfg"] = ttnn.create_sharded_memory_config( shape=[1, 1, 32, 32 * num_cols], core_grid=ttnn.CoreGrid(y=1, x=1), strategy=ttnn.ShardStrategy.WIDTH, ) # --- Phase 4: Resolve weight (sharded across columns, replicated across rows) --- # 2D always uses distributed paths # Weight shape must be (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT), shard on dim 2 # Reshape weight to expected shape if needed (handles [dim], [1, dim], etc.) assert dim % SHARD_HEIGHT == 0, f"dim must be divisible by SHARD_HEIGHT={SHARD_HEIGHT}, got {dim}" expected_shape = (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) if config.weight.source.shape != expected_shape: transformed_source = config.weight.source.reshape(*expected_shape) else: transformed_source = config.weight.source # Apply unit offset if requested (e.g., for Gemma models) if config.add_unit_offset: transformed_source = transformed_source + 1.0 # Create a new LazyWeight with the transformed source transformed_weight = replace(config.weight, source=transformed_source) # Shard on tensor dim 2 across columns (cluster axis 1), replicate across rows (cluster axis 0) mesh_mapper_config = ttnn.MeshMapperConfig( placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(2)], mesh_shape_override=ttnn.MeshShape(list(cluster_shape)), ) resolved_weight = resolve_lazy_weight( transformed_weight, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG, mesh_mapper_config=mesh_mapper_config, ) to_set["weight"] = resolved_weight return replace(config, **to_set) # ============================================================================= # Input Tensor Utilities # ============================================================================= def _load_input_device_tensor_2d(x: "ttnn.Tensor | LazyWeight", config: RMSNorm2DConfig, mode: str) -> ttnn.Tensor: """ Resolve the input tensor to ttnn tensor if x is a LazyWeight, otherwise return as is. For 2D distributed, shards input on last dim across columns (cluster_axis=1). Args: x: Input tensor (ttnn.Tensor or LazyWeight wrapping torch tensor) config: Resolved RMSNorm2DConfig mode: "decode" or "prefill" Returns: ttnn.Tensor ready for the actual forward computation """ assert mode in ("decode", "prefill"), f"mode must be 'decode' or 'prefill', got {mode}" if isinstance(x, LazyWeight): num_rows, num_cols = config.cluster_shape # Determine memory config based on mode if mode == "decode": mem_cfg = config.decode_input_memcfg or ttnn.DRAM_MEMORY_CONFIG else: mem_cfg = config.prefill_input_memcfg or ttnn.DRAM_MEMORY_CONFIG # Shard input on last dim across columns, replicate across rows mesh_mapper_config = ttnn.MeshMapperConfig( placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(-1)], mesh_shape_override=ttnn.MeshShape(num_rows, num_cols), ) resolved_x = resolve_lazy_weight( x, device=config.mesh_device, memory_config=mem_cfg, mesh_mapper_config=mesh_mapper_config, layout=ttnn.TILE_LAYOUT, ) return resolved_x.get_device_weight() # Already a ttnn.Tensor - return as is assert isinstance(x, ttnn.Tensor), f"x must be ttnn.Tensor or LazyWeight, got {type(x)}" return x