# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # # SPDX-License-Identifier: Apache-2.0 from math import prod from typing import Any, ClassVar, Optional from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator def _is_subconfig(annotation: Any) -> bool: """True if the annotation is a `BaseModel` subclass. `issubclass` raises TypeError on generic aliases like `tuple[int, int]` (pydantic's BaseModel metaclass rejects them), so we guard it. """ try: return issubclass(annotation, BaseModel) except TypeError: return False import ttnn from models.common.modules.moe.tt_moe_decode_config_schemas import ( ActivationFunction, CoreCoord, CoreRangeSet, DispatchAlgorithm, MemoryConfig, Topology, WorkerMode, ) def _default_fast_reduce_output_memory_config() -> ttnn.MemoryConfig: """`fast_reduce_output_memory_config` from test_optimized_moe_decode_block.py.""" return ttnn.MemoryConfig( ttnn.BufferType.L1, ttnn.NdShardSpec( ttnn.Shape([1, 32, 128]), ttnn.CoreRangeSet( [ ttnn.CoreRange(ttnn.CoreCoord(2, 0), ttnn.CoreCoord(2, 0)), ttnn.CoreRange(ttnn.CoreCoord(2, 5), ttnn.CoreCoord(2, 5)), ttnn.CoreRange(ttnn.CoreCoord(3, 0), ttnn.CoreCoord(3, 0)), ttnn.CoreRange(ttnn.CoreCoord(3, 5), ttnn.CoreCoord(3, 5)), ttnn.CoreRange(ttnn.CoreCoord(6, 0), ttnn.CoreCoord(6, 0)), ttnn.CoreRange(ttnn.CoreCoord(6, 5), ttnn.CoreCoord(6, 5)), ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(0, 0)), ] ), ttnn.ShardOrientation.ROW_MAJOR, ttnn.ShardDistributionStrategy.ROUND_ROBIN_1D, ), ) def _default_dispatch_input_expert_scores_memory_config( tokens_per_device: int, select_experts_k: int ) -> ttnn.MemoryConfig: """`dispatch_input_expert_scores_memory_config` from test_optimized_moe_decode_block.py. Height-sharded L1, shard shape `[1, select_experts_k]` (seq=1 for decode), spread over a core grid sized from `tokens_per_device`. """ num_cores_y = min(8, tokens_per_device) num_cores_x = (tokens_per_device + num_cores_y - 1) // num_cores_y shard_spec = ttnn.ShardSpec( ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(num_cores_x - 1, num_cores_y - 1))}), [1, select_experts_k], ttnn.ShardOrientation.ROW_MAJOR, ) return ttnn.MemoryConfig( ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, shard_spec, ) _POST_COMBINE_TILIZE_MAX_CORES_X = 7 # original (deepseek) used 7 cores @ 1024 wide for hidden=7168 _POST_COMBINE_TILIZE_SHARD_WIDTH_MULTIPLE = ( 128 # kernel reads 4 tiles wide; shard_width must be a multiple of 4*TILE_SIZE ) # RS-axis size that selects the fused deepseek_moe_reduce_scatter path (consumes a list of # N pre-split outputs from fast_reduce_nc_fused). Any other size falls through to a single # input ttnn.reduce_scatter, so fast_reduce should produce one output instead of N. _DEEPSEEK_RS_DP_DIM = 8 _WH_MAX_CORE_GRID_Y = 9 def _default_post_combine_tilize_memory_config( effective_experts_k: int, hidden_size: int ) -> Optional[ttnn.MemoryConfig]: """`post_combine_tilize_output_memory_config` from test_optimized_moe_decode_block.py. Each expert is one shard-row, `num_cores_x` cores wide (the hidden split). Rows stack down a column up to the device height (`_WH_MAX_CORE_GRID_Y + 1`); when `effective_experts_k` exceeds that, the overflow wraps into additional column-bands of width `num_cores_x`. `num_cores_x` is reduced as needed so the bands (`num_cores_x * num_bands`) fit the device width. `num_cores_x` and the inner shard width are chosen so the shards evenly tile `hidden_size` and the width is a multiple of 128 (the kernel's 4-tile read granularity). Prefer the widest grid that satisfies all of these. Returns None when no `num_cores_x` yields a width that is a multiple of 128 — caller is expected to fall back to the `tilize_with_val_padding` path. """ usable_rows = _WH_MAX_CORE_GRID_Y # 0-indexed inclusive max → row count usable_cols = _POST_COMBINE_TILIZE_MAX_CORES_X num_bands = (effective_experts_k + usable_rows - 1) // usable_rows # Each band is num_cores_x wide; all bands must fit the compute width side by side. max_cores_x = usable_cols // num_bands num_cores_x = None for n in range(max_cores_x, 0, -1): if hidden_size % n != 0: continue width = hidden_size // n if width % _POST_COMBINE_TILIZE_SHARD_WIDTH_MULTIPLE == 0: num_cores_x = n break if num_cores_x is None: return None shard_width = hidden_size // num_cores_x # Lay experts into column-bands: band b holds experts [b*usable_dim, ...), stacked # down y, at x-offset b*num_cores_x. A single band (effective_experts_k ≤ usable_rows) # reproduces the original single-rectangle grid exactly. core_ranges = [] for band in range(num_bands): first_expert = band * usable_rows band_height = min(usable_rows, effective_experts_k - first_expert) x0 = band * num_cores_x core_ranges.append( ttnn.CoreRange( ttnn.CoreCoord(x0, 0), ttnn.CoreCoord(x0 + num_cores_x - 1, band_height - 1), ) ) return ttnn.MemoryConfig( buffer_type=ttnn.BufferType.L1, nd_shard_spec=ttnn.NdShardSpec( shard_shape=[ttnn.TILE_SIZE, shard_width], grid=ttnn.CoreRangeSet(core_ranges), orientation=ttnn.ShardOrientation.ROW_MAJOR, ), ) def _fill_default_if_missing(data: dict, sub_field_name: str, mem_field_name: str, default_value: Any) -> None: """Set `data[sub_field_name][mem_field_name] = default_value` only if not provided. Handles dict/BaseModel/None inputs for the sub-config. User-provided values are preserved; only None or missing keys get the default. """ cur = data.get(sub_field_name) if cur is None: data[sub_field_name] = {mem_field_name: default_value} elif isinstance(cur, dict): if cur.get(mem_field_name) is None: cur[mem_field_name] = default_value elif isinstance(cur, BaseModel): if getattr(cur, mem_field_name, None) is None: data[sub_field_name] = cur.model_copy(update={mem_field_name: default_value}) class _TTOpKwargs(BaseModel): """Base for sub-configs whose `model_dump()` is splatted into a ttnn op call. `arbitrary_types_allowed` lets us hold ttnn objects (MemoryConfig, CoreRangeSet, enum values, etc.) as field values. `frozen` prevents accidental mutation after construction. `populate_by_name` lets users construct via either the field name or its `validation_alias` (used to receive renamed values from the parent). Shared fields adopted from `TTMoEDecodeConfig` are declared via `adopt_fields()` — the parent's pre-validator injects each name in that set into this sub-config's input. For fields whose name in the sub-config differs from the parent (e.g. reduce-scatter `cluster_axis` ← parent's `rs_cluster_axis`), declare the field with `Field(validation_alias='')` and list the parent name in `adopt_fields()`. """ model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True, populate_by_name=True) @classmethod def adopt_fields(cls) -> set[str]: """Parent field/derived-value names this sub-config wants at validation.""" return set() class DispatchConfig(_TTOpKwargs): """Kwargs spread into `ttnn.experimental.all_to_all_dispatch_metadata`.""" shared_expert_ids: Optional[list[int]] = None cluster_axis: int # all_to_all_dispatch_metadata auto-detects num_links when None. num_links: Optional[int] = None drain_sync_tilizer_core: Optional[CoreCoord] = None worker_mode: WorkerMode = ttnn.WorkerMode.DIRECT dispatch_algorithm: DispatchAlgorithm = ttnn.DispatchAlgorithm.SPARSE_MCAST_SHORTEST_PATH @classmethod def adopt_fields(cls) -> set[str]: return {"cluster_axis", "num_links"} class ComputeConfig(_TTOpKwargs): """Kwargs spread into `ttnn.experimental.moe_compute`.""" output_height_shard_dim: int = 4 cluster_axis: int mux_core_range_set: CoreRangeSet = Field( default_factory=lambda: ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(1, 1), ttnn.CoreCoord(3, 3))}) ) has_bias: bool activation_type: ActivationFunction intermediate_size: int # Splatted into moe_compute as `num_shared_experts_per_device`. This is the *physical* # per-device shared-expert count, derived by the parent from `shared_expert_ids_to_devices` # (see TTMoEDecodeConfig._shared_experts_per_device) — NOT the logical `num_shared_experts`. num_shared_experts_per_device: int @classmethod def adopt_fields(cls) -> set[str]: return {"cluster_axis", "has_bias", "num_shared_experts_per_device"} class PostCombineTilizeConfig(_TTOpKwargs): """Kwargs spread into `ttnn.experimental.deepseek_moe_post_combine_tilize` (taken when `batch_per_device == ttnn.TILE_SIZE`). `output_memory_config` defaults to the data-dependent NdShard config from test_optimized_moe_decode_block.py (filled in by the parent validator using `effective_experts_k`). """ output_memory_config: Optional[MemoryConfig] = None class TilizeWithValPaddingConfig(_TTOpKwargs): """Kwargs spread into `ttnn.tilize_with_val_padding` (taken when `batch_per_device != ttnn.TILE_SIZE`). `memory_config` shares the same default as `PostCombineTilizeConfig.output_memory_config`. """ memory_config: Optional[MemoryConfig] = None class ReduceConfig(_TTOpKwargs): """Kwargs spread into `ttnn.experimental.deepseek_moe_fast_reduce_nc_fused`.""" reduce_dim: int = 0 cluster_axis: int split_size: int output_memory_config: MemoryConfig = Field(default_factory=_default_fast_reduce_output_memory_config) num_shared_experts: int shared_expert_scale: float @classmethod def adopt_fields(cls) -> set[str]: return {"cluster_axis", "split_size", "num_shared_experts"} class DeepseekMoEReduceScatterConfig(_TTOpKwargs): """Kwargs spread into `ttnn.experimental.deepseek_moe_reduce_scatter`. `cluster_axis` is adopted from the parent's `rs_cluster_axis` (= `1 - top.cluster_axis`) via a `validation_alias` — reduce-scatter runs on the replicated axis, not the dispatch axis. """ output_memory_config: MemoryConfig = Field(default_factory=lambda: ttnn.DRAM_MEMORY_CONFIG) dim: int = -1 # num_links is auto-detected by the op when None; topology must be Ring (op requirement). num_links: Optional[int] = None topology: Topology = ttnn.Topology.Ring cluster_axis: int = Field(validation_alias="rs_cluster_axis") @classmethod def adopt_fields(cls) -> set[str]: return {"num_links", "topology", "rs_cluster_axis"} class ReduceScatterConfig(_TTOpKwargs): """Kwargs spread into `ttnn.reduce_scatter` (generic fallback path). `cluster_axis` is adopted from the parent's `rs_cluster_axis` via a `validation_alias`, same reasoning as the deepseek path. """ dim: int = -1 # ttnn.reduce_scatter accepts None for both (auto-determines num_links / fabric topology), # so leave these unset rather than forcing a value the caller may not know. num_links: Optional[int] = None cluster_axis: int = Field(validation_alias="rs_cluster_axis") topology: Optional[Topology] = None memory_config: Optional[MemoryConfig] = None @classmethod def adopt_fields(cls) -> set[str]: return {"num_links", "topology", "rs_cluster_axis"} class ExpertStateConfig(_TTOpKwargs): """Routing / weight-prep params for `_TTMoEDecodeState`.""" mesh_shape: tuple[int, int] cluster_axis: int has_bias: bool expert_mapping: list[int] | str num_routed_experts: int num_shared_experts: int shared_expert_ids_to_devices: Optional[dict[int, list[int]] | str] = None @model_validator(mode="before") @classmethod def _expand_convenience_keywords(cls, data: Any) -> Any: """Expand the string-form convenience keywords for `expert_mapping` and `shared_expert_ids_to_devices` into their concrete dict/list forms. Runs before pydantic validates the field types, so we can replace strings in the input dict without bumping into the frozen-model assignment ban. """ if not isinstance(data, dict): return data mesh_shape = data.get("mesh_shape") cluster_axis = data.get("cluster_axis") num_routed = data.get("num_routed_experts") num_shared = data.get("num_shared_experts", 0) if mesh_shape is None or cluster_axis is None or num_routed is None: return data # let pydantic raise about missing fields num_devices = prod(mesh_shape) if num_routed % num_devices != 0: raise ValueError( f"num_routed_experts ({num_routed}) must be evenly divisible " f"by num_devices ({num_devices})" ) # expert_mapping: "sequential" → linearized mesh-coord per expert, following # the convention in `test_optimized_moe_decode_block.get_linearized_mesh_coord`. # cluster_axis=1 (row-major dispatch): linear_id = expert_id // experts_per_device. # cluster_axis=0 (column-major dispatch): adjacent experts straddle clusters — # linear_id = device_within_cluster * num_replicated + cluster_id. mapping = data.get("expert_mapping") if isinstance(mapping, str): if mapping == "sequential": experts_per_device = num_routed // num_devices if cluster_axis == 1: data["expert_mapping"] = [e // experts_per_device for e in range(num_routed)] elif cluster_axis == 0: num_dispatch = mesh_shape[0] num_replicated = num_devices // num_dispatch experts_per_cluster = num_routed // num_replicated new_mapping = [] for e in range(num_routed): cluster_id = e // experts_per_cluster expert_within_cluster = e % experts_per_cluster device_within_cluster = expert_within_cluster // experts_per_device new_mapping.append(device_within_cluster * num_replicated + cluster_id) data["expert_mapping"] = new_mapping else: raise ValueError(f"Unsupported cluster_axis={cluster_axis} for sequential expert_mapping") else: raise ValueError(f"Unknown expert_mapping keyword: {mapping!r}") # shared_expert_ids_to_devices: "fully_replicated" → {num_routed+i: list(range(num_devices)) for i ...} shared = data.get("shared_expert_ids_to_devices") if isinstance(shared, str): if num_shared <= 0: raise ValueError(f"shared_expert_ids_to_devices={shared!r} given but num_shared_experts={num_shared}") if shared == "fully_replicated": data["shared_expert_ids_to_devices"] = { num_routed + i: list(range(num_devices)) for i in range(num_shared) } else: raise ValueError(f"Unknown shared_expert_ids_to_devices keyword: {shared!r}") elif isinstance(shared, dict): if num_shared != len(shared): raise ValueError( f"shared_expert_ids_to_devices has {len(shared)} entries but " f"num_shared_experts={num_shared}" ) return data @classmethod def adopt_fields(cls) -> set[str]: return {"mesh_shape", "cluster_axis", "has_bias", "num_routed_experts", "num_shared_experts"} class BuffersConfig(_TTOpKwargs): """Sizing and placement params for `_TTMoEDecodeBuffers` allocation.""" mesh_shape: tuple[int, int] cluster_axis: int batch_per_device: int hidden_size: int effective_experts_k: int shard_dim: int = 0 # Tilize drain core for the MoE compute op. Must be None (default); TTMoEDecode # resolves it dynamically via ttnn.experimental.get_moe_tilize_drain_core(), which # accounts for the actual device grid, mux exclusion zones, and hidden-size-dependent # core layout. Pinning an explicit core is not supported. compute_tilize_drain_core: Optional[CoreCoord] = None @classmethod def adopt_fields(cls) -> set[str]: return {"mesh_shape", "cluster_axis", "batch_per_device", "hidden_size", "effective_experts_k"} class TTMoEDecodeConfig(BaseModel): """Top-level config for `TTMoEDecode`. Source of truth for shared and derived fields. The pre-validator builds a dict of all parent-visible values (real top-level fields plus derived ones like `rs_cluster_axis`, `split_size`, `effective_experts_k`) and asks each sub-config — via its `adopt_fields()` classmethod — which of those it wants injected before pydantic validates it. """ model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) # Mesh / shape mesh_shape: tuple[int, int] cluster_axis: int batch_per_device: int hidden_size: int select_experts_k: int num_routed_experts: int num_shared_experts: int has_bias: bool # Post-combine path selection. True → fused `deepseek_moe_post_combine_tilize` # (requires `batch_per_device == TILE_SIZE` . # False → generic `tilize_with_val_padding` fallback. Auto-computed in the pre-validator from preconditions. use_post_combine_tilize: Optional[bool] = None # Number of outputs fast_reduce_nc_fused will produce. Equals the RS-axis size # only when the deepseek RS path is taken; otherwise 1 (downstream ttnn.reduce_scatter # consumes a single tensor and splits internally). Auto-computed in pre-validator. num_fast_reduce_outputs: Optional[int] = None # Shared ccl kwargs adopted by sub-configs num_links: Optional[int] = None # Runtime fabric topology (caller-detected). Required: it drives both the fast-reduce # output layout and the reduce-scatter op selection, so we never silently guess it. topology: Topology # Top-level memory configs dispatch_input_memory_config: MemoryConfig = Field(default_factory=lambda: ttnn.L1_MEMORY_CONFIG) dispatch_input_expert_scores_memory_config: Optional[MemoryConfig] = None # Per-op sub-configs dispatch: DispatchConfig compute: ComputeConfig post_combine_tilize: PostCombineTilizeConfig tilize_with_val_padding: TilizeWithValPaddingConfig reduce: ReduceConfig deepseek_moe_reduce_scatter: DeepseekMoEReduceScatterConfig reduce_scatter: ReduceScatterConfig experts: ExpertStateConfig buffers: BuffersConfig # Top-level fields exposed to sub-configs as-is (passthrough); their values # are copied straight from input. Derived/special-case values are computed # from these in `_adoptable()` below. _ADOPTABLE_PASSTHROUGH: ClassVar[tuple[str, ...]] = ( "mesh_shape", "cluster_axis", "batch_per_device", "hidden_size", "select_experts_k", "num_routed_experts", "num_shared_experts", "has_bias", "num_links", "topology", ) @staticmethod def _shared_experts_per_device(shared_expert_ids_to_devices: Any, num_devices: int, num_shared_experts: int) -> int: """Number of *physical* shared experts resident on each device. This is what `moe_compute` wants as `num_shared_experts_per_device`, and it is NOT the logical `num_shared_experts`. It is derived from `shared_expert_ids_to_devices` (global shared-expert id → list of hosting device linear ids): fully-replicated placement puts every shared expert on every device (per-device count == num_shared_experts), while a distributed placement (each shared expert on a subset of devices) yields a smaller, uniform per-device count. e.g. 2 shared experts each residing on half the devices → 1 per device. """ if not num_shared_experts: return 0 # None or the "fully_replicated" convenience keyword (not yet expanded at this # point): every shared expert resides on every device. if shared_expert_ids_to_devices is None or isinstance(shared_expert_ids_to_devices, str): return num_shared_experts counts = [0] * num_devices for devices in shared_expert_ids_to_devices.values(): for device_id in devices: counts[device_id] += 1 if len(set(counts)) != 1: raise ValueError( "shared_expert_ids_to_devices must place an equal number of shared experts on every " f"device to derive num_shared_experts_per_device; got per-device counts {counts}" ) return counts[0] @classmethod def _adoptable(cls, data: dict[str, Any]) -> dict[str, Any]: """Full lookup of values sub-configs may adopt by name. For passthrough fields, fall back to the field's declared default when missing — pydantic only fills defaults *after* this mode='before' validator, so we apply them ourselves so sub-config adoption sees them. """ resolved: dict[str, Any] = {} for name in cls._ADOPTABLE_PASSTHROUGH: finfo = cls.model_fields[name] if name in data and data[name] is not None: resolved[name] = data[name] elif finfo.default_factory is not None: resolved[name] = finfo.default_factory() elif not finfo.is_required() and finfo.default is not None: # A None default means "unset" — don't adopt it, or it would clobber the # sub-config's own default (e.g. DeepseekMoEReduceScatterConfig.topology=Ring) # with None and fail that sub-config's type check. resolved[name] = finfo.default # else: required & missing — bail check upstream returned already # Pick num_fast_reduce_outputs and split_size based on which post-compute RS # path forward() will take: # - rs_axis == _DEEPSEEK_RS_DP_DIM AND topology is Ring → deepseek_moe_reduce_scatter # consumes the full list, so fast_reduce pre-splits into N=rs_axis_size outputs. # Each per-output split must be TILE_SIZE-aligned (fast_reduce's constraint). # - otherwise → ttnn.reduce_scatter (or SKIP) takes a single tensor and does # the per-device split internally. fast_reduce produces N=1 wide output # covering the whole hidden dim; padding it to a multiple of # num_replicated * TILE_SIZE keeps each post-RS device chunk tile-aligned. # deepseek_moe_reduce_scatter only supports Ring, so on any other topology (or when # the caller left topology unset → None) we must use the generic single-output layout. num_replicated = resolved["mesh_shape"][1 - resolved["cluster_axis"]] if num_replicated == _DEEPSEEK_RS_DP_DIM and resolved.get("topology") == ttnn.Topology.Ring: num_fast_reduce_outputs = num_replicated align_unit = ttnn.TILE_SIZE else: num_fast_reduce_outputs = 1 align_unit = ttnn.TILE_SIZE * num_replicated pre_split_chunk = resolved["hidden_size"] // num_fast_reduce_outputs split_size = ((pre_split_chunk + align_unit - 1) // align_unit) * align_unit # Derive the *physical* per-device shared-expert count moe_compute wants from the # experts sub-config's device assignment. Use only `.values()` (device-id lists), # so the dict/str/None and JSON-stringified-key forms all work uniformly. experts_data = data.get("experts") if isinstance(experts_data, BaseModel): shared_expert_ids_to_devices = getattr(experts_data, "shared_expert_ids_to_devices", None) elif isinstance(experts_data, dict): shared_expert_ids_to_devices = experts_data.get("shared_expert_ids_to_devices") else: shared_expert_ids_to_devices = None num_shared_experts_per_device = cls._shared_experts_per_device( shared_expert_ids_to_devices, prod(resolved["mesh_shape"]), resolved["num_shared_experts"] ) return { **resolved, # derived "rs_cluster_axis": 1 - resolved["cluster_axis"], "split_size": split_size, "num_fast_reduce_outputs": num_fast_reduce_outputs, "effective_experts_k": resolved["select_experts_k"] + resolved["num_shared_experts"], "num_shared_experts_per_device": num_shared_experts_per_device, } @model_validator(mode="before") @classmethod def _adopt_shared_into_subconfigs(cls, data: Any) -> Any: if not isinstance(data, dict): return data # Sub-configs whose required fields are entirely either (a) adopted from # the parent or (b) defaulted at the sub-config level can be omitted from # input — slot in an empty dict so the adoption + default-fill logic below # has something to populate. for field_name in ( "dispatch", "compute", "post_combine_tilize", "tilize_with_val_padding", "deepseek_moe_reduce_scatter", "reduce_scatter", "buffers", ): data.setdefault(field_name, {}) # bail if any *truly required* passthrough is missing — let pydantic raise if any(data.get(name) is None for name in cls._ADOPTABLE_PASSTHROUGH if cls.model_fields[name].is_required()): return data adoptable = cls._adoptable(data) for field_name, field_info in cls.model_fields.items(): annotation = field_info.annotation if not _is_subconfig(annotation): continue if not hasattr(annotation, "adopt_fields"): continue additions = {name: adoptable[name] for name in annotation.adopt_fields() if name in adoptable} if not additions: continue cur = data.get(field_name) if isinstance(cur, BaseModel): data[field_name] = {**cur.model_dump(), **additions} elif isinstance(cur, dict): data[field_name] = {**cur, **additions} # else: leave it; pydantic will produce a typed error # Fill in data-dependent memory config defaults where the user didn't provide one. if data.get("dispatch_input_expert_scores_memory_config") is None: data["dispatch_input_expert_scores_memory_config"] = _default_dispatch_input_expert_scores_memory_config( tokens_per_device=adoptable["batch_per_device"], # seq=1 for decode select_experts_k=adoptable["select_experts_k"], ) post_combine_default = _default_post_combine_tilize_memory_config( adoptable["effective_experts_k"], adoptable["hidden_size"] ) # Only fill the post_combine_tilize default if we actually have one — otherwise # leave the field None and let the tilize_with_val_padding fallback own it. if post_combine_default is not None: _fill_default_if_missing(data, "post_combine_tilize", "output_memory_config", post_combine_default) _fill_default_if_missing(data, "tilize_with_val_padding", "memory_config", post_combine_default) # `deepseek_moe_post_combine_tilize` requires batch_per_device == TILE_SIZE AND # a sharded memory config (either default-computable or user-supplied). if data.get("use_post_combine_tilize") is None: user_pct = data.get("post_combine_tilize") user_mc_present = (isinstance(user_pct, dict) and user_pct.get("output_memory_config") is not None) or ( isinstance(user_pct, BaseModel) and getattr(user_pct, "output_memory_config", None) is not None ) data["use_post_combine_tilize"] = adoptable["batch_per_device"] == ttnn.TILE_SIZE and ( post_combine_default is not None or user_mc_present ) else: raise ValidationError("post combine tilize is not user configurable") if data.get("num_fast_reduce_outputs") is None: data["num_fast_reduce_outputs"] = adoptable["num_fast_reduce_outputs"] # Auto-fill `dispatch.shared_expert_ids` when shared experts exist. The dispatch # kernel needs this list to identify which expert ids are shared (broadcast to # the local replica) vs. routed (cross-cluster send). Without it, dispatch # treats shared-expert ids as routed, mis-routes them, and stomps memory. # By contract (enforced in `map_shared_experts`), shared expert ids are always # `[num_routed, num_routed + num_shared)`, regardless of device assignment, so # we derive them directly rather than depending on the experts sub-config # convenience-keyword expansion (which runs after this validator). num_shared = adoptable.get("num_shared_experts", 0) if num_shared > 0: num_routed = adoptable["num_routed_experts"] shared_ids = list(range(num_routed, num_routed + num_shared)) dispatch_data = data.get("dispatch") if isinstance(dispatch_data, dict) and dispatch_data.get("shared_expert_ids") is None: dispatch_data["shared_expert_ids"] = shared_ids elif isinstance(dispatch_data, BaseModel) and getattr(dispatch_data, "shared_expert_ids", None) is None: data["dispatch"] = dispatch_data.model_copy(update={"shared_expert_ids": shared_ids}) return data # ---- Test-only mesh slicing ---- def with_mesh_shape(self, target_mesh_shape: tuple[int, int]) -> "TTMoEDecodeConfig": """Return a new config sized for a column of the original mesh. Lets a config that targets a full e.g. 16x4 mesh run on a single column (16x1) so correctness tests can run on smaller hardware. The `cluster_axis` dim must match the original; the replicated dim must be a divisor of the original. Total expert count scales down proportionally so each device keeps the same expert load (`num_routed_experts / num_devices` is invariant). Derived fields (split_size, num_fast_reduce_outputs, the sub-configs' adopted values, data-dependent memory-config defaults, etc.) are recomputed by round-tripping through `model_validate` with the new mesh and expert count. """ if tuple(target_mesh_shape) == self.mesh_shape: return self if target_mesh_shape[self.cluster_axis] != self.mesh_shape[self.cluster_axis]: raise ValueError( f"cluster_axis dim must be preserved when slicing: " f"target[{self.cluster_axis}]={target_mesh_shape[self.cluster_axis]} != " f"self[{self.cluster_axis}]={self.mesh_shape[self.cluster_axis]}" ) other = 1 - self.cluster_axis if self.mesh_shape[other] % target_mesh_shape[other] != 0: raise ValueError( f"replicated dim must be a divisor of the original: " f"self[{other}]={self.mesh_shape[other]} not divisible by " f"target[{other}]={target_mesh_shape[other]}" ) new_num_routed = self.num_routed_experts * target_mesh_shape[other] // self.mesh_shape[other] if new_num_routed % prod(target_mesh_shape) != 0: raise ValueError( f"scaled num_routed_experts ({new_num_routed}) is not divisible by " f"target num_devices ({prod(target_mesh_shape)})" ) data = self.model_dump(mode="json", exclude_defaults=True, exclude_none=True) data["mesh_shape"] = list(target_mesh_shape) data["num_routed_experts"] = new_num_routed # Force re-derivation of mesh-derived top-level fields. self._drop_derived_fields(data) # Re-derive expert routing — the dumped expert_mapping is the expanded # "sequential" list for the old mesh and won't address the sliced experts. # Explicit (non-sequential) mappings aren't supported by the slicer. experts_data = data.setdefault("experts", {}) if isinstance(experts_data, dict): experts_data["expert_mapping"] = "sequential" if self.num_shared_experts > 0: experts_data["shared_expert_ids_to_devices"] = "fully_replicated" # Clear the auto-filled `dispatch.shared_expert_ids` so it re-derives against # the new num_routed_experts. if isinstance(data.get("dispatch"), dict): data["dispatch"].pop("shared_expert_ids", None) return type(self).model_validate(data) @staticmethod def _drop_derived_fields(data: dict) -> None: """In-place: remove top-level fields the pre-validator derives (and rejects as input) so they re-derive when dumped data is fed back through `model_validate`. """ for derived in ("num_fast_reduce_outputs", "use_post_combine_tilize"): data.pop(derived, None) # ---- YAML round-trip ---- def to_yaml(self, *, exclude_defaults: bool = True, exclude_none: bool = True, **dump_kwargs) -> str: """Serialize to a minimal YAML string. Defaults to skipping fields equal to their default and `None` fields, so the output only carries values the user actually configured. ttnn objects (MemoryConfig, CoreCoord, etc.) are serialized via their `to_json` representation; enums by their name. """ import yaml data = self.model_dump(mode="json", exclude_defaults=exclude_defaults, exclude_none=exclude_none, **dump_kwargs) self._drop_derived_fields(data) return yaml.safe_dump(data, sort_keys=False) @classmethod def from_yaml(cls, yaml_str: str, topology: ttnn.Topology) -> "TTMoEDecodeConfig": """Load a config from YAML. `topology` is the runtime fabric topology (the caller detects it from the device). It overrides any `topology` authored in the YAML, since the layout and reduce-scatter path depend on the actual hardware, not the saved value. The pre-validators on each field re-hydrate ttnn objects from their dict / string representations via `from_json` (and enums by their name). """ import yaml data = yaml.safe_load(yaml_str) data["topology"] = topology return cls.model_validate(data)