# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # # SPDX-License-Identifier: Apache-2.0 """TTMoEGate — the routing front-end that feeds ``TTMoEDecode``. ``TTMoEDecode.forward`` consumes a pre-computed routing decision per token (``tt_scores``/``tt_indices`` of shape ``[batch, 1, 1, select_experts_k]`` — the ``all_to_all_dispatch_metadata`` op requires the token, index, and score tensors to share their first three dims, and the token tensor is ``[batch, 1, 1, hidden]``); this module PRODUCES it from the router logits: hidden ──[router matmul: ttnn.matmul(h, W_gate)]──▶ logits[1,1,batch,num_experts] ──[reshape→16×16 height-shard, one token/core]──▶ gate-op input ──[generalized_moe_gate: (sigmoid|softmax)+top-k+normalize]──▶ (scores,indices)[batch,32,32] ──[slice :k + view to the decode layout]──▶ (weights, indices)[batch,1,1,k] The device flow mirrors ``models/demos/deepseek_v3/tt/moe_gate.py`` (decode path): the gate op is single-token-per-core, so a variable ``batch`` is padded up to a multiple of the core count and processed in ``batch_per_iter`` chunks, reusing the bias / input-indices / output buffers. External tensors are plain L1-interleaved; the height-sharded layout the op needs is produced inside ``forward``. Gate-op selection (via ``n_group``): - ``n_group=1`` → ungrouped global top-k. With k ∈ {4,6,8} and ≤512 experts it uses the ``generalized_moe_gate`` kernel (configurable ``topk`` + sigmoid/softmax; softmax-no-bias models default to ``enable_sigmoid=False, output_softmax=True`` since softmax-over-selected ≡ softmax-front + renormalize. ``config.softmax_position="pre"`` instead applies a full softmax over all experts up front — mathematically equal absent a score-correction bias (the global Z cancels under renorm), a numerically-stable exp — and takes the linear-renorm tail like sqrtsoftplus). Any other k, OR >512 experts, takes a pure-ttnn fallback (``matmul → topk → softmax``) — so there is NO expert-count ceiling for n_group=1. - ``n_group=8`` → ``deepseek_moe_gate`` (grouped; sigmoid + bias). The kernel is HARDWIRED to 256 experts as 8 groups × 32, emitting top-8, so this path is EXACTLY 256-experts-select-8 (see ``__init__``). Expert layout (kernel-op path only): ≤256 runs on one 256-face (fewer experts padded up with phantom ``_PAD_NEG`` experts that rank last); 257-512 uses the op's 2-block combine path (n_group=1 only). """ from __future__ import annotations import torch import ttnn from models.common.modules.moe.tt_moe_gate_config import TTMoEGateConfig _TOKEN_SHAPE = (16, 16) # 256 experts laid out as a 16×16 face per token _SHARD_SHAPE = (32, 32) # one 32×32 tile per core (the op's per-token shard) _FACE_EXPERTS = _TOKEN_SHAPE[0] * _TOKEN_SHAPE[1] # 256: the op's fixed face size; <256 models pad up to it # Phantom (padding) experts get this large-negative value in BOTH the logit and the bias. Selection ranks # by `score + bias` (score = logit for softmax, sigmoid(logit) for sigmoid); driving the bias to -1e9 makes # phantom experts rank last in EITHER mode (a -1e9 logit alone isn't enough for sigmoid, since sigmoid(-1e9)=0 # could still sit above a real expert with a negative bias). They never enter the top-k, so they never reach # the output weights or the softmax exp. _PAD_NEG = -1e9 # top-k values the generalized_moe_gate KERNEL op supports. Its finalize rank-mask keeps the right ranks # ONLY for these (the op's C++ validation + unit tests cover exactly {4,6,8}): topk 1-3 leave ranks 0-3 # unmasked (wrong normalize) and 5/7 are untested. Any other k (n_group=1) takes the pure-ttnn fallback, # which handles any k. Keep in sync with the op's TT_FATAL (generalized_moe_gate_device_operation.cpp). _KERNEL_TOPK = (4, 6, 8) class TTMoEGate: """Produce ``(weights, indices)`` for ``TTMoEDecode`` from hidden states.""" def __init__( self, mesh_device, config: TTMoEGateConfig, torch_gate_weight: torch.Tensor, # [hidden_size, num_experts] router projection # TWO distinct biases (both [num_experts], both optional): # torch_gate_bias = score-CORRECTION bias (deepseek/sigmoid): added INSIDE the op to the # scores for SELECTION ranking only; the output weights stay unbiased. # torch_gate_proj_bias = router LINEAR bias (gpt-oss): added to the LOGITS (Wx + b) BEFORE the op, # so it flows into BOTH selection AND the softmax/normalized weights. torch_gate_bias: torch.Tensor | None = None, torch_gate_proj_bias: torch.Tensor | None = None, ): # Entry point mirrors TTMoEDecode: a config + the torch gate weight/bias(es) in; everything # device-side is built here. Unpack the config into the locals the rest of __init__ uses. self.config = config # The config DECLARES which bias tensors a model has (`score_correction_bias` = the selection-only # correction bias `torch_gate_bias`; `gate_proj_bias` = the router-LINEAR/projection bias `torch_gate_proj_bias`). # Nothing downstream re-checks: a missing bias is SILENTLY swapped for zeros (correction, line ~211) # or a bias-free matmul (router, forward), and a stray bias is used regardless of the flag — both # change routing while looking valid. Enforce the contract here so a miswired gate fails at # construction, not as a quietly-wrong route. (Validate BOTH directions: required AND forbidden.) if config.score_correction_bias and torch_gate_bias is None: raise ValueError("config.score_correction_bias=True requires torch_gate_bias, but none was given") if not config.score_correction_bias and torch_gate_bias is not None: raise ValueError("config.score_correction_bias=False forbids torch_gate_bias, but one was given") if config.gate_proj_bias and torch_gate_proj_bias is None: raise ValueError("config.gate_proj_bias=True requires torch_gate_proj_bias, but none was given") if not config.gate_proj_bias and torch_gate_proj_bias is not None: raise ValueError("config.gate_proj_bias=False forbids torch_gate_proj_bias, but one was given") # ...and the right SHAPE: both biases are 1-D [num_routed_experts] (the device upload reshapes to # [1, N] / F.pads the last dim assuming exactly that). A [1, N] / [N, 1] / wrong-length tensor would # mis-pad or mis-broadcast into the wrong experts silently, so reject anything but [num_experts]. for _name, _tensor in (("torch_gate_bias", torch_gate_bias), ("torch_gate_proj_bias", torch_gate_proj_bias)): if _tensor is not None and tuple(_tensor.shape) != (config.num_routed_experts,): raise ValueError( f"{_name} must be shape [{config.num_routed_experts}] (num_routed_experts); " f"got {tuple(_tensor.shape)}" ) num_experts = config.num_routed_experts select_experts_k = config.select_experts_k hidden_size = config.hidden_size n_group = config.n_group score_func = config.score_func scaling_factor = config.routed_scaling_factor eps = config.eps # `score_func` maps to (a) the score nonlinearity and (b) two op flags -- enable_sigmoid (apply # sigmoid inside the op) and output_softmax (softmax over the selected top-k vs linear renorm): # "softmax" → score=logit, enable_sigmoid=False, output_softmax=True (no bias) # "sigmoid" → score=sigmoid(logit), enable_sigmoid=True, output_softmax=False (+bias) # "sqrtsoftplus" → score=sqrt(softplus(l)), enable_sigmoid=False, output_softmax=False (+bias) # sqrtsoftplus (deepseek-v4) has no in-kernel op, so it is applied EXTERNALLY in forward() via # ttnn.sqrt(ttnn.softplus(.)) and fed to the op with enable_sigmoid=False (the op then just adds # bias, ranks, and linearly renormalizes -- same tail as the sigmoid path). # `softmax_position` (softmax only) moves the softmax exp relative to the top-k. With NO score- # correction bias the two placements are mathematically equal (the global Z cancels under renorm), # differing only in how the exp is computed; they would DIVERGE with a selection bias, since ranking # by softmax(logit)+bias ≠ logit+bias (softmax is nonlinear) picks different experts — but softmax # models are bias-free, so this is moot in practice. # "post" (default) → the mapping above (output_softmax=True: exp-over-selected). # "pre" → softmax over ALL experts is applied EXTERNALLY in forward() (same slot as # sqrtsoftplus) and fed to the op with enable_sigmoid=False, output_softmax=False # (LINEAR renorm over the selected) — so the exp is the stable full softmax, # not the kernel's exp-over-raw-logit. Same op flags/tail as sqrtsoftplus. if score_func not in ("softmax", "sigmoid", "sqrtsoftplus"): raise ValueError(f"score_func must be 'softmax' | 'sigmoid' | 'sqrtsoftplus'; got {score_func!r}") if score_func == "softmax" and config.softmax_position not in ("pre", "post"): raise ValueError(f"softmax_position must be 'pre' | 'post'; got {config.softmax_position!r}") # n_group picks the backing op: 1 = generalized/ungrouped (kernel op or ttnn fallback), # 8 = deepseek_moe_gate (grouped). Only these two are wired. if n_group not in (1, 8): raise NotImplementedError( f"n_group must be 1 (generalized/ungrouped) or 8 (deepseek grouped); got {n_group}" ) if num_experts < 1: raise ValueError(f"num_experts must be ≥ 1; got {num_experts}") # n_group=8 → the deepseek grouped op. It is HARDWIRED to 256 experts laid out as 8 groups × 32, # emitting top-8 (top-2-sum per group → top-4 groups → top-8 of 128). It takes NO group/expert-count/k # argument — grouped_golden bakes in the top-2/top-4/top-8 constants, and the op unit test # (models/demos/deepseek_v3_b1/tests/unit_tests/test_deepseek_moe_gate.py) only ever feeds (batch,8,32). # So n_group=8 is EXACTLY 256-experts-select-8: a <256 grouped model has ≠32 experts/group (padding to # 256 would re-group it → wrong top-8) and k≠8 isn't emitted, so neither is supported. (No real grouped # model is <256 anyway — deepseek-v3 / ling-1t are both 256.) if n_group == 8 and not (num_experts == 256 and select_experts_k == 8): raise NotImplementedError( f"n_group=8 (deepseek grouped op) is hardwired to 256 experts select-8; got N={num_experts}, k={select_experts_k}" ) # Ungrouped (n_group=1) dispatch — TWO backends, and NO expert-count ceiling (a large N is NOT rejected, # it just takes the fallback): # • kernel op (forward): k ∈ {4,6,8} AND N ≤ 512. The op works in 256-expert blocks (a 16×16 # face): ≤256 pads up to ONE block, 257-512 uses the op's 2-block COMBINE path. Phantom experts # (num_experts..num_blocks*256) get a large-negative logit + bias (_PAD_NEG) so they rank last. # • ttnn fallback (_forward_fallback): EVERYTHING else — k ∉ {4,6,8} (the kernel's rank-mask only # covers 4/6/8, drop threshold = 16*(k-4)) OR N > 512 (beyond the 2-block combine). It is pure # ttnn.matmul → topk → softmax, so it handles any k and any N (e.g. qwen3.5: 512-experts top-10). use_fallback = (n_group == 1) and (select_experts_k not in _KERNEL_TOPK or num_experts > 2 * _FACE_EXPERTS) # Defense-in-depth: the kernel op path (n_group=1 AND not fallback) must only ever see a supported # topk. use_fallback already routes every other k to the pure-ttnn path, so this can only fire if # that routing is later changed inconsistently — fail loudly here instead of feeding the kernel a # topk it silently mis-normalizes (e.g. topk 1-3 → first-4-ranks). The deepseek grouped op # (n_group=8) is a separate kernel hardwired to k=8 (guarded above), so it's exempt. if n_group == 1 and not use_fallback and select_experts_k not in _KERNEL_TOPK: raise ValueError( f"generalized_moe_gate kernel op supports topk {_KERNEL_TOPK}; got {select_experts_k} " "(this k should have taken the ttnn fallback)" ) self.mesh_device = mesh_device self.num_experts = num_experts # op-path block count: 1 (≤256) or 2 (257-512). For the n_group=1 fallback (N>512) these are computed # but UNUSED — the fallback returns early below and runs ttnn.topk over the full unpadded logits. self.num_blocks = (num_experts + _FACE_EXPERTS - 1) // _FACE_EXPERTS self._padded_experts = self.num_blocks * _FACE_EXPERTS self.k = select_experts_k self.hidden_size = hidden_size self.n_group = n_group self.score_func = score_func self.scaling_factor = scaling_factor self.eps = eps self.softmax_position = config.softmax_position # "post" (default) | "pre" — softmax only self.enable_sigmoid = score_func == "sigmoid" # output_softmax = the exp-over-selected tail (in-op / ttnn.softmax(sel)). ONLY softmax-post takes it; # softmax-pre, sigmoid, and sqrtsoftplus all use the LINEAR renorm tail. self.output_softmax = score_func == "softmax" and self.softmax_position == "post" # external score transform applied to the logits before the op (None → the op's own scoring): # sqrtsoftplus → sqrt∘softplus; softmax-pre → full softmax over ALL experts (then linear renorm). if score_func == "sqrtsoftplus": self.score_transform = "sqrtsoftplus" elif score_func == "softmax" and self.softmax_position == "pre": self.score_transform = "softmax" else: self.score_transform = None self.use_fallback = use_fallback # ungrouped + k∉{4,6,8} → pure-ttnn path (no kernel op) grid = mesh_device.compute_with_storage_grid_size() self.num_device_cores = grid.x * grid.y self._grid = grid self.compute_kernel_config = self._matmul_compute_config( config.gate_matmul_compute, config.gate_matmul_high_fidelity ) self.matmul_program_config = self._matmul_program_config( config.gate_matmul_program_config, config.gate_matmul_auto_program_config ) # Op path: pad the router weight (+ linear bias) to the block size UP FRONT so the matmul directly # emits padded_experts-wide logits — this avoids a per-forward ttnn.pad (its dispatch + memory copy # dominated the gate latency). Phantom columns are zero (their logit = 0); the phantom score-correction # bias (_PAD_NEG in tt_bias) is what ranks them last, so the zero logit is harmless. (The fallback runs # ttnn.topk over the REAL experts, so it keeps the unpadded weight — padding would add phantom experts # to the topk.) if not self.use_fallback and num_experts < self._padded_experts: pad_cols = self._padded_experts - num_experts torch_gate_weight = torch.nn.functional.pad(torch_gate_weight, (0, pad_cols)) # [hidden, padded_experts] if torch_gate_proj_bias is not None: torch_gate_proj_bias = torch.nn.functional.pad(torch_gate_proj_bias, (0, pad_cols)) # [padded_experts] # --- router weight (replicated, DRAM) --- self.tt_gate_weight = ttnn.from_torch( torch_gate_weight, # [hidden, num_experts] (op path: padded to [hidden, padded_experts]) dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) # --- router LINEAR bias (gpt-oss): logits = Wx + b, added by ttnn.linear in forward (None → matmul) --- self.tt_gate_proj_bias = None if torch_gate_proj_bias is not None: self.tt_gate_proj_bias = ttnn.from_torch( torch_gate_proj_bias.reshape( 1, -1 ), # [1, num_experts] (op path: [1, padded_experts]) — ttnn.linear bias dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) # --- fallback path: ttnn.topk works on the full logits directly, so none of the op's face/combine/ # output buffers are needed. Build only the (optional) score-correction bias for selection # ([1, num_experts] broadcast-add), then return. --- if self.use_fallback: self.tt_fallback_bias = None if torch_gate_bias is not None: self.tt_fallback_bias = ttnn.from_torch( torch_gate_bias.reshape(1, num_experts).to(torch.float32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) return # --- per-token buffers, sliced to batch_per_iter in forward. Sized to batch_per_device when the # config provides it (decode: the per-forward slice is then a free FULL-tensor slice, zero cost), # else to the device core count (the slice trims cores→batch_per_iter — the variable-batch # fallback). The slice in forward is kept either way; it's just free in the sized-to-batch case. --- buffer_rows = config.batch_per_device if config.batch_per_device is not None else self.num_device_cores self._buffer_rows = buffer_rows mem = self._sharded_mem_config(buffer_rows) # bias + input-indices buffers. Real experts [0:num_experts] get torch_gate_bias (or 0 for no-bias — # a constant shift changes neither selection nor output, so zeros is the correct no-bias case); # phantom experts [num_experts:num_blocks*256] get _PAD_NEG so they always rank last (see _PAD_NEG). real_bias = (torch_gate_bias if torch_gate_bias is not None else torch.zeros(num_experts)).to(torch.float32) bias = torch.full((self._padded_experts,), _PAD_NEG, dtype=torch.float32) bias[:num_experts] = real_bias if self.num_blocks == 1: # single 256-face: (1,16,16) → pad to (1,32,32) → transpose → per core. indices = arange(256). bias_face = torch.nn.functional.pad(bias.reshape(1, 16, 16), (0, 16, 0, 16, 0, 0), "constant", 0).transpose( 1, 2 ) # (1, 32, 32) self.tt_bias = ttnn.from_torch( bias_face.repeat(buffer_rows, 1, 1), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=mem, ) idx = torch.arange(_FACE_EXPERTS, dtype=torch.int32).reshape(1, 16, 16).transpose(1, 2) self.tt_input_indices = ttnn.from_torch( idx.repeat(buffer_rows, 1, 1), dtype=ttnn.uint16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=mem, ) else: # combine path (num_blocks×256): num_blocks 16×16 faces stacked, each block transposed; indices # carry GLOBAL expert ids (block b = arange(256)+b*256). Shard (num_blocks*32, 32), tile (32,32). # Mirrors test_generalized_moe_gate.py::test_generalized_moe_gate_512_global. tile = ttnn.Tile((32, 32)) mem_blocks = self._sharded_mem_config(buffer_rows, (self.num_blocks * 32, 32)) bias_blocks = torch.transpose(bias.reshape(self.num_blocks, 16, 16), -2, -1).contiguous() # (nb,16,16) self.tt_bias = ttnn.from_torch( bias_blocks.unsqueeze(0).repeat(buffer_rows, 1, 1, 1), # (buffer_rows, nb, 16, 16) dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=mem_blocks, tile=tile, ) ar = torch.arange(_FACE_EXPERTS, dtype=torch.int32).reshape(1, 16, 16) offs = (torch.arange(self.num_blocks, dtype=torch.int32) * _FACE_EXPERTS).reshape(self.num_blocks, 1, 1) idx_blocks = torch.transpose(ar + offs, -2, -1).contiguous().to(torch.int32) # (nb, 16, 16) global ids self.tt_input_indices = ttnn.from_torch( idx_blocks.unsqueeze(0).repeat(buffer_rows, 1, 1, 1), dtype=ttnn.uint16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=mem_blocks, tile=tile, ) # preallocated output buffers (filled in place by the op), per core. ROW_MAJOR (like moe_gate.py) # so the final ttnn.view is metadata-only — a TILE buffer routes view through the device # reshape_view op, which rejects uint16. self.tt_output = ttnn.from_torch( torch.zeros((buffer_rows, *_SHARD_SHAPE), dtype=torch.float32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=mesh_device, memory_config=mem, ) self.tt_output_indices = ttnn.from_torch( torch.zeros((buffer_rows, *_SHARD_SHAPE), dtype=torch.int32), dtype=ttnn.uint16, layout=ttnn.ROW_MAJOR_LAYOUT, device=mesh_device, memory_config=mem, ) # ------------------------------------------------------------------ helpers def _matmul_compute_config(self, cfg, high_fidelity): """Build the router-matmul ``compute_kernel_config``. The gate matmul is tiny but its reduction is ``hidden_size``-deep; with the ttnn DEFAULT matmul (LoFi + bf16 accumulate) the accumulated error flips near-tied experts at the top-k boundary (output weights stay ~right, but the selected expert *ids* drift). Whether a given gate actually has a near-tie there is data-dependent (not a clean function of ``hidden_size`` / ``k``), so rather than gate it per-model we DEFAULT every gate to HiFi2 + fp32 accumulation — the matmul is small, so the fidelity cost is negligible and it removes the drift uniformly. Built device-agnostically (``init_device_compute_kernel_config`` picks the right config class per arch), unlike a hardcoded ``WormholeComputeKernelConfig`` (wrong on Blackhole). - ``cfg`` dict → custom: keys ``math_fidelity`` ("HiFi2"|…), ``math_approx_mode``, ``fp32_dest_acc_en``, ``packer_l1_acc`` override the high-fidelity defaults. A dict (or prebuilt config) WINS over ``high_fidelity``. - else ``high_fidelity`` True (default) → built-in HiFi2 + fp32 accumulate + packer L1 acc. - else (``high_fidelity`` False) → ttnn's default (low) matmul fidelity. """ if cfg is not None and not isinstance(cfg, dict): return cfg # already a ttnn compute_kernel_config → use as-is if cfg is None: if not high_fidelity: return None # opt out → ttnn's default matmul fidelity cfg = {} # built-in high-fidelity default (keys below) cfg = dict(cfg) fidelity = cfg.pop("math_fidelity", "HiFi2") if isinstance(fidelity, str): fidelity = getattr(ttnn.MathFidelity, fidelity) return ttnn.init_device_compute_kernel_config( self.mesh_device.arch(), math_fidelity=fidelity, math_approx_mode=cfg.pop("math_approx_mode", True), fp32_dest_acc_en=cfg.pop("fp32_dest_acc_en", True), packer_l1_acc=cfg.pop("packer_l1_acc", True), ) def _matmul_program_config(self, cfg, auto): """Build the router-matmul ``program_config``. The gate matmul is tiny (``batch×hidden @ hidden×experts``) but ttnn's auto-selected program config is generic; a hand-tuned 2D-multicast config (kernel grid + block/subblock sizing) materially cuts its latency. The tuning is pure shape/grid math (no data dependence), so by DEFAULT we DERIVE it from the gate matmul shape + the device grid (see ``_derive_program_config``) rather than hand-listing it per model. - ``cfg`` dict → full override: key ``type`` names the ttnn program-config class (default "MatmulMultiCoreReuseMultiCastProgramConfig"); the remaining keys are its kwargs (e.g. ``in0_block_w``, ``per_core_M``, ``per_core_N``, ``out_subblock_*``). ``compute_with_storage_grid_size`` is auto-filled when omitted (yaml stays arch-agnostic). A dict (or prebuilt config) WINS over ``auto``. - else ``auto`` True (default) → the derived 2D-mcast config. - else (``auto`` False) → ttnn auto-selects the program config. """ if cfg is not None and not isinstance(cfg, dict): return cfg # already a ttnn program config → use as-is if cfg is None: if not auto: return None # let ttnn auto-select cfg = self._derive_program_config() cfg = dict(cfg) pc_class = getattr(ttnn, cfg.pop("type", "MatmulMultiCoreReuseMultiCastProgramConfig")) cfg.setdefault("compute_with_storage_grid_size", self.mesh_device.compute_with_storage_grid_size()) return pc_class(**cfg) def _derive_program_config(self) -> dict: """Derive the 2D-mcast blocking for the gate matmul (``batch×hidden @ hidden×padded_experts``) from its shape + the device grid. Matches the previously hand-tuned per-model configs, except the 256-expert gate now uses ``per_core_N=1`` instead of ``2`` (see the ``per_core_N`` bullet). ``M_tiles = ceil(batch/32)`` (=1 for the batch-32 decode gate); ``K_tiles = hidden/32``; ``N_tiles = padded_experts/32`` (256 experts → 8, the 512-combine path → 16). • ``in0_block_w`` = largest divisor of ``K_tiles`` that is ≤ 32 (the widest K-block within the 32-tile cap; a non-1024-multiple hidden gives e.g. 20 / 22 / 30 instead of 32). • ``per_core_N`` = smallest divisor ≥ 1 of ``N_tiles`` that fits the grid (``N_tiles / per_core_N ≤ grid.x``) — i.e. spread N across as many cores as the grid holds. 256 experts (``N_tiles=8``) on an 8-wide grid → ``1`` (8 N-cores); the 512-combine path (``N_tiles=16``) → ``2`` (16 tiles can't map to 16 cores, so 8). NOTE: ``per_core_N=1`` measured faster than ``2`` for the tiny (``M_tiles=1``) 256 gate. • ``out_block_{h,w}`` = ``per_core_{M,N}``; ``out_subblock_{h,w}`` = largest factors with ``h * w ≤ dst_subblock_cap`` — the DEST subblock budget = HALF of the DEST tile capacity (MATH/PACK double-buffer): 8 tiles for bf16 accumulation, but **4 for fp32 accumulation** (an fp32 tile takes 2× the DEST space, so half as many fit). This matches ttnn's own ``get_subblock_sizes`` rule (``h*w ≤ 4`` iff ``fp32_dest_acc_en``). The gate DEFAULTS to fp32 accumulation (see ``_matmul_compute_config``), so the cap is normally 4. """ # DEST subblock tile budget: fp32 accumulation halves the usable DEST (8->4 tiles). Read the flag off # the already-built compute config (None = ttnn default = bf16 accumulate). Must match the cap ttnn's # own auto-derivation enforces, else a hand-built subblock could exceed the fp32 DEST budget. fp32_dest_acc_en = bool(getattr(self.compute_kernel_config, "fp32_dest_acc_en", False)) dst_subblock_cap = 4 if fp32_dest_acc_en else 8 batch = self.config.batch_per_device or 32 m_tiles = (batch + 31) // 32 k_tiles = self.hidden_size // 32 n_tiles = self._padded_experts // 32 grid_x = self._grid.x in0_block_w = max(d for d in range(1, 33) if k_tiles % d == 0) per_core_n = next((d for d in range(1, n_tiles + 1) if n_tiles % d == 0 and n_tiles // d <= grid_x), n_tiles) # out_subblock_w capped at dst_subblock_cap too: out_subblock_h can be 1, so w alone can equal the # h*w budget — without the cap a wide per_core_n (e.g. 8) would pick w=8,h=1 = 8 tiles, busting fp32's 4. out_subblock_w = next((w for w in range(min(per_core_n, dst_subblock_cap), 0, -1) if per_core_n % w == 0), 1) out_subblock_h = next( (h for h in range(m_tiles, 0, -1) if m_tiles % h == 0 and h * out_subblock_w <= dst_subblock_cap), 1 ) return { "in0_block_w": in0_block_w, "out_subblock_h": out_subblock_h, "out_subblock_w": out_subblock_w, "out_block_h": m_tiles, "out_block_w": per_core_n, "per_core_M": m_tiles, "per_core_N": per_core_n, "transpose_mcast": False, "fused_activation": None, } def _sharded_mem_config(self, num_cores: int, shard_shape: tuple = _SHARD_SHAPE) -> ttnn.MemoryConfig: # one token per core; shard_shape is (32,32) for a single 256-block, (num_blocks*32, 32) for the # combine path (num_blocks tiles stacked per core). core_grid = ttnn.num_cores_to_corerangeset(num_cores, ttnn.CoreCoord(self._grid.x, self._grid.y), row_wise=True) return ttnn.MemoryConfig( ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(core_grid, shard_shape, ttnn.ShardOrientation.ROW_MAJOR), ) # ------------------------------------------------------------------ device def forward(self, tt_x: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: """``tt_x``: hidden states ``[1, 1, batch, hidden_size]`` (L1/DRAM interleaved) → ``(weights, indices)`` each ``[batch, 1, 1, select_experts_k]`` — the layout ``TTMoEDecode``/``all_to_all_dispatch_metadata`` consumes (batch in dim 0, sharing the token tensor's first three dims). The batch-in-dim-2 → batch-in-dim-0 move is a metadata-only view (ROW_MAJOR, identical memory order).""" if self.use_fallback: return self._forward_fallback(tt_x) # 1) router projection → logits [1, 1, batch, num_experts]. With a router LINEAR bias (gpt-oss) it's # Wx + b via ttnn.linear; otherwise the proven bias-free matmul (unchanged for every other model). # compute_kernel_config + program_config are per-model (see __init__): a deep, tie-sensitive gate # wants HiFi2 + fp32, and a tuned 2D-mcast program_config can cut the matmul latency (else None). if self.tt_gate_proj_bias is None: logits = ttnn.matmul( tt_x, self.tt_gate_weight, compute_kernel_config=self.compute_kernel_config, program_config=self.matmul_program_config, ) else: logits = ttnn.linear( tt_x, self.tt_gate_weight, bias=self.tt_gate_proj_bias, compute_kernel_config=self.compute_kernel_config, program_config=self.matmul_program_config, ) # logits is already padded_experts-wide (the router weight was padded up front, see __init__) — no # per-forward ttnn.pad. Phantom columns are zero; the phantom bias (_PAD_NEG in tt_bias) ranks them # last regardless of their (bounded) score, so they never enter the top-k. # 1a) external score transform fed to the op with enable_sigmoid=False + the LINEAR renorm tail. Both # transforms turn `logits` into the per-expert score; the op then adds bias, ranks, gathers, and # linearly renormalizes over the selected. (Phantom cols are 0 → bounded score → ranked out by # the _PAD_NEG bias.) # • sqrtsoftplus (deepseek-v4): score = sqrt(softplus(logit)). # • softmax-pre (softmax_position="pre"): score = softmax(logit) over ALL (padded) experts. The # linear renorm over the selected cancels softmax's global Z, so this is mathematically equal # to softmax-over-selected (softmax-post) absent a selection bias — but the exp is the stable # full softmax, not the kernel's exp-over-raw-logit. (Phantom cols get a bounded softmax prob; # Z cancels regardless.) if self.score_transform == "sqrtsoftplus": # standard softplus log(1+e^x): beta=1, threshold=20 (matches torch's defaults / the golden). logits = ttnn.sqrt(ttnn.softplus(logits, beta=1.0, threshold=20.0)) elif self.score_transform == "softmax": logits = ttnn.softmax(logits, dim=-1) # softmax over the full padded width; renorm tail cancels Z total_batch = logits.shape[2] # 2) the gate op is one-token-per-core: process in equal chunks of ≤ chunk tokens, reusing the # bias/indices/output buffers. chunk = min(cores, buffer_rows): bounded by cores (one token/core) # AND by the buffers' row count, so the per-iter slice always fits. When buffer_rows == # batch_per_device == this iter's batch (decode), the slice is a free full-tensor slice. chunk = min(self.num_device_cores, self._buffer_rows) num_iters = (total_batch + chunk - 1) // chunk padding = (num_iters - (total_batch % num_iters)) % num_iters batch_per_iter = (total_batch + padding) // num_iters if padding != 0: logits = ttnn.pad(logits, [(0, 0), (0, 0), (0, padding), (0, 0)], 0) # input/bias/indices live on (num_blocks*32, 32) shards (1 token/core, num_blocks tiles); the output # is always a single (32,32) tile per token (the combine writes the final global top-k into row 0). mem_in = self._sharded_mem_config(batch_per_iter, (self.num_blocks * 32, 32)) mem_out = self._sharded_mem_config(batch_per_iter, _SHARD_SHAPE) if self.num_blocks == 1: reshaped_shape = (batch_per_iter, *_TOKEN_SHAPE) # (bpi, 16, 16) bias = ttnn.slice(self.tt_bias, [0, 0, 0], [batch_per_iter, 16, 16], memory_config=mem_in) in_idx = ttnn.slice(self.tt_input_indices, [0, 0, 0], [batch_per_iter, 16, 16], memory_config=mem_in) else: reshaped_shape = (batch_per_iter, self.num_blocks, *_TOKEN_SHAPE) # (bpi, nb, 16, 16) bias = ttnn.slice( self.tt_bias, [0, 0, 0, 0], [batch_per_iter, self.num_blocks, 16, 16], memory_config=mem_in ) in_idx = ttnn.slice( self.tt_input_indices, [0, 0, 0, 0], [batch_per_iter, self.num_blocks, 16, 16], memory_config=mem_in ) out = ttnn.slice(self.tt_output, [0, 0, 0], [batch_per_iter, 32, 32], memory_config=mem_out) out_idx = ttnn.slice(self.tt_output_indices, [0, 0, 0], [batch_per_iter, 32, 32], memory_config=mem_out) weights_chunks, indices_chunks = [], [] for start in range(0, total_batch + padding, batch_per_iter): cur = logits[:, :, start : start + batch_per_iter, :] cur = ttnn.reshape(cur, reshaped_shape) # (bpi,16,16) or (bpi,num_blocks,16,16) for combine cur = ttnn.to_memory_config(cur, memory_config=mem_in) # height-shard: one token/core if self.n_group == 8: # deepseek grouped top-8 (sigmoid + bias); no top-k / output-softmax knobs. Call the C++ op # directly (ttnn namespace, not a demo package), symmetric with the generalized op below. w, idx = ttnn.experimental.deepseek.moe.deepseek_moe_gate( cur, bias_tensor=bias, input_indices_tensor=in_idx, output_tensor=out, output_indices_tensor=out_idx, eps=self.eps, scaling_factor=self.scaling_factor, enable_sigmoid=self.enable_sigmoid, ) else: w, idx = ttnn.experimental.deepseek.moe.generalized_moe_gate( cur, bias_tensor=bias, input_indices_tensor=in_idx, output_tensor=out, output_indices_tensor=out_idx, eps=self.eps, scaling_factor=self.scaling_factor, enable_sigmoid=self.enable_sigmoid, topk=self.k, output_softmax=self.output_softmax, ) weights_chunks.append(ttnn.to_memory_config(w, memory_config=ttnn.L1_MEMORY_CONFIG)) indices_chunks.append(ttnn.to_memory_config(idx, memory_config=ttnn.L1_MEMORY_CONFIG)) ttnn.deallocate(cur) ttnn.deallocate(logits) weights = weights_chunks[0] if num_iters == 1 else ttnn.concat(weights_chunks, dim=0) indices = indices_chunks[0] if num_iters == 1 else ttnn.concat(indices_chunks, dim=0) # 3) take the top-k (ranks 0..k-1 are the first k cols of row 0) → [batch,1,1,k], the layout # TTMoEDecode/all_to_all_dispatch_metadata wants (batch in dim 0). ROW_MAJOR output buffers make # these views metadata-only (no device reshape_view) — the slice is [total_batch, k] and dim 0 is # already the batch, so the view to [total_batch,1,1,k] is the same memory in the same order (the # uint16 indices reshape without a dtype round-trip — same as moe_gate.py). weights = ttnn.view(weights[:total_batch, 0, : self.k], (total_batch, 1, 1, self.k)) indices = ttnn.view(indices[:total_batch, 0, : self.k], (total_batch, 1, 1, self.k)) return weights, indices def _forward_fallback(self, tt_x: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: """Pure-ttnn ungrouped gate for cases the kernel op can't do (k∉{4,6,8}, e.g. qwen3.5 top-10). Mirrors the op's logic with stock ttnn ops, no 256-face / combine / padding: matmul → score transform → (+ score-correction bias for selection) → topk → gather unbiased scores → normalize (softmax-over-selected if output_softmax, else linear renorm) at the END → × scale. Output ``(weights, indices)`` each ``[batch, 1, 1, k]`` (the decode/dispatch layout, like the op path), indices uint16.""" # 1) logits = Wx (+ router LINEAR bias). compute_kernel_config per-model (HiFi2 for deep gates). if self.tt_gate_proj_bias is None: logits = ttnn.matmul( tt_x, self.tt_gate_weight, compute_kernel_config=self.compute_kernel_config, program_config=self.matmul_program_config, ) else: logits = ttnn.linear( tt_x, self.tt_gate_weight, bias=self.tt_gate_proj_bias, compute_kernel_config=self.compute_kernel_config, program_config=self.matmul_program_config, ) # 2) score = transform(logits): sqrt(softplus) | full softmax (softmax-pre) | sigmoid | # identity (softmax-post / sigmoid-less, ranks by the raw logit). softmax-pre does the exp here via # the stable full softmax over ALL experts; step 6 then LINEAR-renorms over the selected, which # cancels Z → mathematically equal to softmax-post (absent a selection bias). if self.score_transform == "sqrtsoftplus": scores = ttnn.sqrt(ttnn.softplus(logits, beta=1.0, threshold=20.0)) elif self.score_transform == "softmax": scores = ttnn.softmax(logits, dim=-1) elif self.enable_sigmoid: scores = ttnn.sigmoid(logits) else: scores = logits # 3) add the score-CORRECTION bias for SELECTION ranking (None → skip; qwen3.5 is bias-free). rank_key = scores if self.tt_fallback_bias is None else ttnn.add(scores, self.tt_fallback_bias) # 4) top-k over ALL experts → [1, 1, batch, k]. ttnn.topk handles any k / expert count, ungrouped. topk_vals, topk_idx = ttnn.topk(rank_key, self.k, dim=-1) # 5) weights = the UNBIASED scores at the selected experts. With a correction bias the topk values are # the BIASED ranking key, so gather the unbiased scores; without one they're already the scores. sel = topk_vals if self.tt_fallback_bias is None else ttnn.gather(scores, dim=-1, index=topk_idx) # 6) normalize at the END (unified): softmax-over-selected (output_softmax) or linear renorm; × scale. if self.output_softmax: weights = ttnn.softmax(sel, dim=-1) else: weights = ttnn.div(sel, ttnn.sum(sel, dim=-1, keepdim=True)) if self.scaling_factor != 1.0: weights = ttnn.multiply(weights, self.scaling_factor) indices = ttnn.typecast(topk_idx, ttnn.uint16) # reshape [1,1,batch,k] → [batch,1,1,k] to match the op path / the decode dispatch layout. These are # TILE-layout topk/softmax outputs, so a view can't move `batch` out of the tiled dims directly; go to # ROW_MAJOR first (which all_to_all_dispatch_metadata needs anyway), and then — as on the op path — the # view is metadata-only (identical byte order). batch lives in dim 2 here. batch = weights.shape[2] weights = ttnn.view(ttnn.to_layout(weights, ttnn.ROW_MAJOR_LAYOUT), (batch, 1, 1, self.k)) indices = ttnn.view(ttnn.to_layout(indices, ttnn.ROW_MAJOR_LAYOUT), (batch, 1, 1, self.k)) return weights, indices # ------------------------------------------------------------------ golden @staticmethod def golden( hidden: torch.Tensor, # [batch, hidden] gate_weight: torch.Tensor, # [hidden, num_experts] gate_bias: torch.Tensor | None, # [num_experts] score-correction bias (selection-only) or None *, select_experts_k: int, score_func: str = "softmax", softmax_position: str = "post", # softmax only: "post" = exp-over-selected; "pre" = softmax over all + renorm scaling_factor: float = 1.0, eps: float = 1e-20, n_group: int = 1, gate_proj_bias: torch.Tensor | None = None, # [num_experts] router LINEAR bias (gpt-oss): logits = Wx + b ) -> tuple[torch.Tensor, torch.Tensor]: """PyTorch reference: hidden → logits → gate → (scores[batch,k], indices[batch,k]). n_group==8 (grouped) delegates to ``TTMoEGate.grouped_golden`` (below); n_group==1 (ungrouped) is inlined below. Both are self-contained — no dependency on any demo package.""" batch = hidden.shape[0] logits = hidden.float() @ gate_weight.float() # [batch, num_experts] # router LINEAR bias (gpt-oss) is part of the projection: logits = Wx + b. It flows into BOTH # selection and the softmax weights (unlike the score-correction `gate_bias`, which is selection-only). if gate_proj_bias is not None: logits = logits + gate_proj_bias.float() bias = gate_bias.float() if gate_bias is not None else torch.zeros(logits.shape[-1]) enable_sigmoid = score_func == "sigmoid" # external score transform (sqrtsoftplus, deepseek-v4): mirror forward() — apply sqrt∘softplus up # front and feed the op golden with enable_sigmoid=False, so it just adds bias / ranks / renorms. if score_func == "sqrtsoftplus": logits = torch.sqrt(torch.nn.functional.softplus(logits)) if n_group == 8: # deepseek grouped top-8: 8 groups of 32 experts (top-2-sum per group → top-4 groups → top-8). # The golden's grouping is driven by shape[-2], so reshape to the LOGICAL (8, 32) — NOT the # (16, 16) tile face (that's only the device's physical layout; the kernel still groups the 256 # experts as 8×32). Bias likewise reshapes to (8, 32) (the device uploads it transposed in the # 16×16 tile, but the golden works in logical expert order). See the op unit test: # models/demos/deepseek_v3_b1/tests/unit_tests/test_deepseek_moe_gate.py (input_shape=(b,8,32)). return TTMoEGate.grouped_golden( logits.reshape(batch, 8, 32), bias.reshape(8, 32), eps=eps, scaling_factor=scaling_factor, enable_sigmoid=enable_sigmoid, ) # ungrouped global top-k (inlined — the generalized op is no longer a deepseek-v3 dependency): # rank by (score + bias), gather the UNBIASED score at the selected experts, normalize, scale. # score over ALL experts: sigmoid → sigmoid(logit); softmax-pre → softmax(logit) over all experts # (mirrors forward's external transform); softmax-post / sqrtsoftplus → the (already-transformed) logit. if enable_sigmoid: scores = torch.sigmoid(logits) # [batch, num_experts] elif score_func == "softmax" and softmax_position == "pre": scores = torch.softmax(logits, dim=-1) else: scores = logits _, indices = torch.topk(scores + bias, select_experts_k, dim=-1, sorted=True) # bias broadcasts [num_experts] sel = torch.gather(scores, -1, indices) # tail: softmax-POST is the only exp-over-selected; softmax-pre / sigmoid / sqrtsoftplus → linear renorm # (softmax-pre already did the exp up front, so its renorm cancels the global Z → mathematically equal # to softmax-post absent a selection bias). weights = torch.exp(sel) if (score_func == "softmax" and softmax_position == "post") else sel return weights / (weights.sum(-1, keepdim=True) + eps) * scaling_factor, indices @staticmethod def grouped_golden( input_tensor: torch.Tensor, # [batch, n_group, group_size] per-expert scores/logits (e.g. [batch, 8, 32]) bias_tensor: torch.Tensor, # [n_group, group_size] or [batch, n_group, group_size] selection bias eps: float = 1e-20, scaling_factor: float = 2.5, enable_sigmoid: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """PyTorch reference for the DeepSeek GROUPED gate (n_group=8): 256 experts as 8 groups × 32 — per group take the top-2 bias-corrected scores and sum them, pick the top-4 groups by that sum, then the top-8 (by bias-corrected score) over those 4 groups (128 experts), gather the UNBIASED score at the chosen experts, linearly renormalize and scale. Returns (scores[batch, 8], global_indices[batch, 8]). Mirrors the device op ``ttnn.experimental.deepseek.moe.deepseek_moe_gate``; SINGLE source of truth, shared by ``TTMoEGate.golden`` (n_group=8) and the op unit test (test_generalized_moe_gate.py).""" row_offsets = torch.arange(input_tensor.shape[-2]) * input_tensor.shape[-1] # group g -> base id g*32 batch_idx = torch.arange(input_tensor.shape[0]).unsqueeze(-1) scores = torch.sigmoid(input_tensor) if enable_sigmoid else input_tensor bias_scores = scores + bias_tensor sorted_bias, sorted_indices = torch.sort(bias_scores, dim=-1, descending=True) sorted_scores = torch.gather(scores, dim=-1, index=sorted_indices) sorted_indices = sorted_indices + row_offsets.view(1, -1, 1) # local -> global expert id top2_sum = sorted_bias[:, :, 0] + sorted_bias[:, :, 1] # per-group top-2-sum _, sorted_top2_indices = torch.sort(top2_sum, dim=-1, descending=True) # rank groups by it top4_values = sorted_bias[batch_idx, sorted_top2_indices[:, :4]].flatten(1) # top-4 groups' bias scores top4_scores = sorted_scores[batch_idx, sorted_top2_indices[:, :4]].flatten(1) top4_indices = sorted_indices[batch_idx, sorted_top2_indices[:, :4]].flatten(1) _, top8_pos = torch.topk(top4_values, 8, dim=-1, sorted=True) # top-8 of the 128 by bias score top8_scores = torch.gather(top4_scores, dim=-1, index=top8_pos) # UNBIASED scores at the chosen experts top8_indices = torch.gather(top4_indices, dim=-1, index=top8_pos) return top8_scores / (torch.sum(top8_scores, dim=-1, keepdim=True) + eps) * scaling_factor, top8_indices