File size: 5,205 Bytes
2415c4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
#
# SPDX-License-Identifier: Apache-2.0
"""Config for ``TTMoEGate`` (the routing front-end).
Parsed from the same per-model YAML that ``TTMoEDecodeConfig`` reads — the gate-specific keys
(``n_group``/``score_func``/``routed_scaling_factor``/``score_correction_bias``/``gate_proj_bias``/
``gate_matmul_compute``) are exactly the extra keys ``TTMoEDecodeConfig`` ignores. Mirrors the
``TTMoEDecode(mesh_device, config, torch_weights...)`` entry-point convention: ``TTMoEGate`` takes this
config + the torch gate weight/bias(es), and builds everything on device internally.
"""
from __future__ import annotations
import yaml
from pydantic import BaseModel, ConfigDict
# Gate-relevant keys read from a model YAML; every other key (decode/compute/reduce/...) is ignored.
_GATE_YAML_KEYS = (
"num_routed_experts",
"select_experts_k",
"hidden_size",
"batch_per_device",
"n_group",
"score_func",
"softmax_position",
"routed_scaling_factor",
"score_correction_bias",
"gate_proj_bias",
"gate_matmul_high_fidelity",
"gate_matmul_compute",
"gate_matmul_auto_program_config",
"gate_matmul_program_config",
)
class TTMoEGateConfig(BaseModel):
"""Gate shape + routing semantics for ``TTMoEGate``.
``batch_per_device`` is OPTIONAL: when set, ``TTMoEGate`` sizes its per-token op buffers to it, so the
per-forward ``ttnn.slice`` of those buffers is a free full-tensor slice (zero perf cost — decode path).
When ``None``, the buffers are sized to the device core count and the slice trims ``cores →
batch_per_iter`` (the variable-/large-batch fallback). Either way the result is identical.
"""
model_config = ConfigDict(frozen=True)
# --- shape ---
num_routed_experts: int
select_experts_k: int
hidden_size: int
batch_per_device: int | None = None # set → buffers sized to it (free slice); None → sized to cores
# --- gate semantics (see TTMoEGate) ---
n_group: int = 1 # 1 → generalized (ungrouped) op; 8 → deepseek (grouped) op
score_func: str = "softmax" # "softmax" | "sigmoid" | "sqrtsoftplus"
# softmax ONLY (ignored otherwise) — where the softmax is applied relative to the top-k. With no score-
# correction bias both are mathematically equal (the global Z cancels under renorm), differing only in
# how the exp is computed (a selection bias would make them diverge, but softmax models are bias-free):
# "post" (default) — rank by the raw logit, then softmax OVER THE SELECTED top-k (the kernel op's
# exp-over-selected / ttnn.softmax(sel) in the fallback). The original behavior.
# "pre" — softmax over ALL experts UP FRONT (same slot as sqrtsoftplus), then LINEAR renorm
# over the selected top-k. The exp is the numerically-stable full softmax
# (max-subtracted over all experts) instead of the kernel's exp-over-raw-logit.
softmax_position: str = "post" # "post" | "pre"
routed_scaling_factor: float = 1.0
score_correction_bias: bool = False # deepseek/noaux_tc: e_score_correction_bias (selection-only)
gate_proj_bias: bool = (
False # gpt-oss: model has a router LINEAR/projection bias b (logits = Wx + b → selection AND weights)
)
# router-matmul COMPUTE kernel fidelity. Default HIGH fidelity (HiFi2 + fp32 accumulate): the deep gate
# reduction drifts under the ttnn default (LoFi + bf16 accumulate) and flips near-tied experts at the
# top-k boundary, so every gate gets high fidelity (the matmul is tiny → the cost is negligible). Set
# `gate_matmul_high_fidelity: false` to opt out (ttnn's default fidelity).
gate_matmul_high_fidelity: bool = True
# OPTIONAL full override of the compute kernel: a dict (keys math_fidelity / math_approx_mode /
# fp32_dest_acc_en / packer_l1_acc override the HiFi2 defaults). Wins over gate_matmul_high_fidelity
# when set; None → use the flag above. See TTMoEGate._matmul_compute_config.
gate_matmul_compute: dict | None = None
# router-matmul PROGRAM config (kernel grid / 2D-mcast blocking). Default AUTO: derived from the gate
# matmul shape + device grid (TTMoEGate._derive_program_config) — a tuned config materially cuts the gate
# matmul latency vs ttnn's auto pick. Set `gate_matmul_auto_program_config: false` for ttnn's auto-select.
gate_matmul_auto_program_config: bool = True
# OPTIONAL full override of the program config: a dict whose `type` names the ttnn program-config class
# (default MatmulMultiCoreReuseMultiCastProgramConfig) and remaining keys are its kwargs
# (compute_with_storage_grid_size auto-filled). Wins over the auto flag when set; None → auto / ttnn.
gate_matmul_program_config: dict | None = None
eps: float = 1e-20
@classmethod
def from_yaml(cls, yaml_text: str) -> "TTMoEGateConfig":
"""Build from a model YAML, picking only the gate-relevant keys (defaults fill the rest)."""
raw = yaml.safe_load(yaml_text)
return cls(**{k: raw[k] for k in _GATE_YAML_KEYS if k in raw})
|