| |
| |
|
|
| """Qwen3 decoder layer and stack on ttnn -- all three GR00T versions (IMPLEMENTATION_PLAN.md §2.2 row ``tt/qwen3.py``, |
| §3.3 LLM op table, decisions D3 / D4 / D6 / D7 / D8; WP-09). |
| |
| One layer (plan §3.3; mirrors ``reference/qwen3.py`` ``Qwen3DecoderLayer``):: |
| |
| h = rms_norm(x, input_ln.w) HiFi4 fp32acc approx-off (mb2 §B.5) |
| qkv = linear(h, Wqkv [2048, 4096]) table row (m_bucket(S_pad), 2048, 4096, dtype), mb1 Table 2-llm |
| q, k, v = split_heads(qkv, 16, 8) nlp_create_qkv_heads(transpose_k_heads=False), bit-exact (D7) |
| q = rms_norm(q, q_norm.w [1,128]); k = rms_norm(k, k_norm.w) per-head Qwen3 q/k norm BEFORE RoPE |
| q = rope(q, cos, sin[, trans_mat]); k = rope(k, ...) one op per tensor (D8): meta -> rotary_embedding_llama, |
| hf -> rotary_embedding_hf |
| o = sdpa(q, k, v, scale 128**-0.5, is_causal=True) n15 / n17 right padding; (32,32)/(128,128) chunks (§3.4) |
| | sdpa(q, k, v, mask=llm_mask, is_causal=False) n16 text-slot padding: causal AND key-valid mask (WP-02) |
| x = x + linear(concat_heads(o), Wo) GQA 16/8 handled by the SDPA kernel |
| [tap llm_layer_{i}] BEFORE the N1.7 deepstack add (golden README llm_layer0_out) |
| [x = x + deepstack_placed] N1.7 layers 0, 1, 2 (arch-n17 §2.3) |
| h = rms_norm(x, post_attn_ln.w) |
| f = SwiGLU(h): S_pad >= 256 -> minimal_matmul(W_gate_up tile_interleaved, fuse_swiglu=True) -> [S, 6144] |
| S_pad == 128 -> linear(W_gate_up plain) -> slice gate / up -> mul(gate, up, SILU on a) |
| (D3 / mb1 Table 4f: the fused op is 1.3x faster at S = 512 and 12 % slower at S = 128) |
| x = x + linear(f, W_down [6144, 2048]) |
| |
| WP-P (Stage 3, tests/tt/results/p_summary.md): the causal SDPA takes the re-measured allowlist defaults of |
| ``tt.layers.SDPA_CHUNKS`` ((64, 128) at S = 384 / 512: 36.8 / 43.8 µs vs 39.0 / 59.7 µs with (128, 128), same PCC); the |
| L1 residual was measured (-8 / -9 µs per layer) and rejected (L1 clash in the VL-SA ff1, :data:`RESIDUAL_DRAM_MIN_S_PAD`). |
| |
| ``Qwen3StackTT`` runs ``cfg.llm.n_layers_executed`` layers, adds ``deepstack_placed[l]`` after the layers in |
| ``cfg.llm.deepstack_add_after_layers`` (N1.7: (0, 1, 2)) and applies ``rms_norm(llm.final_norm.w)`` iff |
| ``cfg.llm.final_norm_applied`` (N1.5 / N1.6 yes, N1.7 no -- synthesis §0 fact 3). |
| |
| Weights are the ``common.weights`` plan tensors resident in ``tt.weights.TTWeights`` (foundation-f1 §2 names |
| ``llm.layer.{i}.{Wqkv, Wo, W_gate_up, W_down, input_ln.w, post_attn_ln.w, q_norm.w, k_norm.w}``, ``llm.final_norm.w``): |
| under ``TransformOptions(rope_layout="meta")`` the Q / K rows of ``Wqkv`` and the ``q_norm`` / ``k_norm`` vectors are |
| already in the Meta pairwise layout (``reverse_permute`` / ``reverse_permute_1d``), and ``W_gate_up`` is |
| ``tile_interleaved`` iff ``resolve_gate_up_layout(policy, S_pad) == "tile_interleaved"``; the constructors verify both |
| against ``TTWeights.options`` so a mismatched byte layout raises instead of computing garbage. |
| |
| RoPE tables (foundation-f2 §2): ``ModelInputs.rope_cos / rope_sin`` are the per-call ``(1, S_pad, 128)`` fp32 HF |
| half-duplicated tables (N1.6: built from the bf16-rounded ``inv_freq``, f2 §4 F1); :func:`prepare_rope` converts them |
| to the Meta pairwise layout for ``rope_layout="meta"`` (``common.weights.rope_cos_sin_hf_to_meta``) or keeps them for |
| ``"hf"``, uploads them as bf16 TILE ``[1, 1, S_pad, 128]`` (per call into persistent buffers for N1.6, once per shape key |
| otherwise) together with the ``[1, 1, 32, 32]`` transformation matrix of ``rotary_embedding_llama``. |
| |
| ``ttnn`` is imported lazily inside the forward methods / upload helpers; the module, the constructors' shape contract |
| and the host helpers (:func:`rows_to_padded`, :func:`padded_to_rows`, :func:`rope_tables_host`) are CPU-testable. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Dict, List, Optional, Sequence, Tuple |
|
|
| import torch |
|
|
| from models.experimental.gr00t.common.configs import TILE, GR00TConfig |
| from models.experimental.gr00t.common.weights import rope_cos_sin_hf_to_meta |
| from models.experimental.gr00t.tt.device import ( |
| CKC_HIFI2_FP32ACC, |
| CKC_HIFI4_FP32ACC_NOAPPROX, |
| CKC_SDPA_HIFI4, |
| DRAM, |
| EXPECTED_COMPUTE_GRID, |
| L1, |
| core_grid, |
| resolve_ckc, |
| resolve_mem, |
| ) |
| from models.experimental.gr00t.tt.layers import ( |
| WEIGHT_DTYPES, |
| concat_heads, |
| default_sdpa_chunks, |
| dtype_name, |
| linear, |
| matmul_spec, |
| residual_add, |
| rms_norm, |
| rope, |
| sdpa, |
| split_heads, |
| ) |
| from models.experimental.gr00t.tt.policy import TTPolicy, resolve_gate_up_layout |
|
|
| |
| |
| |
| LLM_SDPA_CLASS = "llm_causal" |
| ATTN_MODES: Tuple[str, ...] = ("causal", "mask") |
| SWIGLU_MODES: Tuple[str, ...] = ("fused", "split") |
| SWIGLU_FUSED_MIN_S_PAD = 256 |
| |
| |
| |
| |
| |
| |
| RESIDUAL_DRAM_MIN_S_PAD = 384 |
|
|
| |
| |
| |
| |
| |
| |
| |
| SWIGLU_MINIMAL_BLOCKING: Dict[int, Tuple[int, int, int, int, int]] = { |
| 128: (1, 8, 16, 1, 4), |
| 256: (1, 8, 16, 1, 4), |
| 384: (2, 8, 16, 1, 4), |
| 512: (2, 8, 16, 1, 4), |
| } |
|
|
| LAYER_WEIGHT_STEMS: Tuple[str, ...] = ( |
| "Wqkv", |
| "Wo", |
| "W_gate_up", |
| "W_down", |
| "input_ln.w", |
| "post_attn_ln.w", |
| "q_norm.w", |
| "k_norm.w", |
| ) |
| FINAL_NORM_NAME = "llm.final_norm.w" |
|
|
| |
| |
| LAYER_TAP_NAMES: Tuple[str, ...] = ( |
| "input_layernorm", |
| "qkv", |
| "q_norm", |
| "k_norm", |
| "q_rope", |
| "k_rope", |
| "sdpa", |
| "attn", |
| "residual_attn", |
| "post_attention_layernorm", |
| "ff_hidden", |
| "mlp", |
| ) |
|
|
|
|
| def swiglu_mode(s_pad: int) -> str: |
| """``"fused"`` (``minimal_matmul(fuse_swiglu=True)``, tile-interleaved ``W_gate_up``) for ``S_pad >= 256``, else |
| ``"split"`` (``linear`` -> 2 ``slice`` -> ``mul(SILU)``) -- plan D3, mb1 Table 4f.""" |
| if int(s_pad) <= 0 or int(s_pad) % TILE: |
| raise ValueError(f"S_pad must be a positive multiple of {TILE}, got {s_pad}") |
| return "fused" if int(s_pad) >= SWIGLU_FUSED_MIN_S_PAD else "split" |
|
|
|
|
| def gate_up_layout_for(mode: str) -> str: |
| """The ``TransformOptions.gate_up_layout`` a SwiGLU mode consumes (``"fused"`` -> ``"tile_interleaved"``).""" |
| if mode not in SWIGLU_MODES: |
| raise ValueError(f"swiglu mode must be one of {SWIGLU_MODES}, got {mode!r}") |
| return "tile_interleaved" if mode == "fused" else "plain" |
|
|
|
|
| def residual_mem_for(s_pad: int) -> Any: |
| """Residual-stream memory preset of plan §3.3: DRAM at ``S_pad >= 384``, L1 at ``128`` (and 256); the WP-P L1 |
| measurement and why it was not adopted are recorded at :data:`RESIDUAL_DRAM_MIN_S_PAD`.""" |
| return DRAM if int(s_pad) >= RESIDUAL_DRAM_MIN_S_PAD else L1 |
|
|
|
|
| def llm_sdpa_chunks(s_pad: int, head_dim: int) -> Tuple[int, int]: |
| """Plan §3.4 default ``(q_chunk, k_chunk)`` of the causal LLM SDPA row ``(llm_causal, S_pad, S_pad, head_dim)``; |
| ``KeyError`` for an unlisted S_pad (measure it first, plan D6).""" |
| return default_sdpa_chunks(LLM_SDPA_CLASS, int(s_pad), int(s_pad), int(head_dim)) |
|
|
|
|
| def swiglu_minimal_blocking(s_pad: int) -> Tuple[int, int, int, int, int]: |
| """Frozen ``minimal_matmul`` blocking of the fused SwiGLU at ``S_pad`` (:data:`SWIGLU_MINIMAL_BLOCKING`).""" |
| s = int(s_pad) |
| if s not in SWIGLU_MINIMAL_BLOCKING: |
| raise KeyError( |
| f"no fused-SwiGLU minimal_matmul blocking for S_pad={s}; measure it (mb1 Table 4f method) and add it to " |
| "tt/qwen3.py SWIGLU_MINIMAL_BLOCKING" |
| ) |
| m_blk, k_blk, n_blk, sub_h, sub_w = SWIGLU_MINIMAL_BLOCKING[s] |
| if n_blk % 2: |
| raise ValueError(f"S_pad={s}: fuse_swiglu needs an even N_block, got {n_blk}") |
| return m_blk, k_blk, n_blk, sub_h, sub_w |
|
|
|
|
| def swiglu_minimal_config(s_pad: int, hidden: int, inter: int) -> Any: |
| """``ttnn.MinimalMatmulConfig`` of the fused SwiGLU ``[S_pad, hidden] x [hidden, 2 * inter]`` on the 11x10 grid.""" |
| import ttnn |
|
|
| m_blk, k_blk, n_blk, sub_h, sub_w = swiglu_minimal_blocking(s_pad) |
| n_tiles = (2 * int(inter)) // TILE |
| if (2 * int(inter)) % (2 * TILE) or n_tiles % 2: |
| raise ValueError(f"fuse_swiglu needs 2 * inter = {2 * inter} to be a multiple of {2 * TILE}") |
| if int(hidden) % (k_blk * TILE): |
| raise ValueError(f"K_block {k_blk} tiles does not divide hidden {hidden}") |
| return ttnn.MinimalMatmulConfig( |
| M_block_size=m_blk, |
| K_block_size=k_blk, |
| N_block_size=n_blk, |
| subblock_h=sub_h, |
| subblock_w=sub_w, |
| compute_with_storage_grid_size=core_grid(EXPECTED_COMPUTE_GRID), |
| ) |
|
|
|
|
| def layer_weight_names(layer_idx: int) -> Dict[str, str]: |
| """``{stem: plan name}`` of one layer (foundation-f1 §2: ``llm.layer.{i}.<stem>``).""" |
| if int(layer_idx) < 0: |
| raise ValueError(f"layer_idx must be >= 0, got {layer_idx}") |
| return {stem: f"llm.layer.{int(layer_idx)}.{stem}" for stem in LAYER_WEIGHT_STEMS} |
|
|
|
|
| def stack_weight_names(cfg: GR00TConfig) -> List[str]: |
| """Every plan tensor the stack of ``cfg`` needs (``TTWeights.load(names=...)``), final norm iff applied.""" |
| names: List[str] = [] |
| for i in range(cfg.llm.n_layers_executed): |
| names.extend(layer_weight_names(i).values()) |
| if cfg.llm.final_norm_applied: |
| names.append(FINAL_NORM_NAME) |
| return names |
|
|
|
|
| |
| |
| |
| def rows_to_padded(rows: torch.Tensor, slot_of_token: torch.Tensor, s_pad: int) -> torch.Tensor: |
| """Scatter the ``[S, D]`` (or ``[1, S, D]``) token rows of the reference sequence into the static ``[1, 1, S_pad, |
| D]`` layout: real token ``i`` goes to ``slot_of_token[i]`` (``SlotLayout.slot_of_token``), every other slot is zero |
| (static-shape-plan §4.3; pad rows never reach a real query through the causal / key-valid masks).""" |
| if rows.ndim == 3: |
| if rows.shape[0] != 1: |
| raise ValueError(f"rows must be [S, D] or [1, S, D], got {tuple(rows.shape)}") |
| rows = rows[0] |
| if rows.ndim != 2: |
| raise ValueError(f"rows must be [S, D] or [1, S, D], got {tuple(rows.shape)}") |
| slot = torch.as_tensor(slot_of_token).reshape(-1).long() |
| s, d = rows.shape |
| if slot.numel() != s: |
| raise ValueError(f"slot_of_token has {slot.numel()} entries for {s} rows") |
| if int(s_pad) < s or int(s_pad) % TILE: |
| raise ValueError(f"S_pad={s_pad} must be a multiple of {TILE} and >= S={s}") |
| if s and (int(slot.min()) < 0 or int(slot.max()) >= int(s_pad)): |
| raise ValueError(f"slot_of_token outside [0, {s_pad})") |
| if slot.unique().numel() != s: |
| raise ValueError("slot_of_token has duplicate slots") |
| out = torch.zeros(1, 1, int(s_pad), d, dtype=rows.dtype) |
| out[0, 0, slot] = rows |
| return out |
|
|
|
|
| def padded_to_rows(x: torch.Tensor, slot_of_token: torch.Tensor) -> torch.Tensor: |
| """Inverse of :func:`rows_to_padded`: gather the real token rows of a ``[1, 1, S_pad, D]`` (or ``[1, S_pad, D]``) |
| device layout back into the golden's ``[1, S, D]`` token order.""" |
| if x.ndim == 4: |
| if x.shape[0] != 1 or x.shape[1] != 1: |
| raise ValueError(f"x must be [1, 1, S_pad, D], got {tuple(x.shape)}") |
| x = x[0] |
| if x.ndim != 3 or x.shape[0] != 1: |
| raise ValueError(f"x must be [1, 1, S_pad, D] or [1, S_pad, D], got {tuple(x.shape)}") |
| slot = torch.as_tensor(slot_of_token).reshape(-1).long() |
| if slot.numel() and (int(slot.min()) < 0 or int(slot.max()) >= x.shape[1]): |
| raise ValueError(f"slot_of_token outside [0, {x.shape[1]})") |
| return x[:, slot, :] |
|
|
|
|
| def rope_tables_host(mi: Any, policy: TTPolicy, head_dim: int = 128) -> Tuple[torch.Tensor, torch.Tensor]: |
| """``(cos, sin)`` bf16 ``[1, 1, S_pad, head_dim]`` host tables for ``policy.rope_layout`` from |
| ``ModelInputs.rope_cos / rope_sin`` (``(1, S_pad, head_dim)`` fp32, HF half-duplicated; foundation-f2 §2). |
| |
| ``"meta"``: ``common.weights.rope_cos_sin_hf_to_meta`` (pairwise ``c0, c0, c1, c1, ...``) for |
| ``rotary_embedding_llama``; ``"hf"``: unchanged for ``rotary_embedding_hf`` (plan D8). |
| """ |
| if not isinstance(policy, TTPolicy): |
| raise TypeError(f"policy must be a TTPolicy, got {type(policy).__name__}") |
| cos, sin = mi.rope_cos, mi.rope_sin |
| if not isinstance(cos, torch.Tensor) or not isinstance(sin, torch.Tensor): |
| raise TypeError("ModelInputs.rope_cos / rope_sin must be torch tensors") |
| if cos.shape != sin.shape or cos.ndim != 3 or cos.shape[0] != 1 or cos.shape[-1] != int(head_dim): |
| raise ValueError(f"rope tables must be (1, S_pad, {head_dim}), got {tuple(cos.shape)} / {tuple(sin.shape)}") |
| if cos.shape[1] % TILE: |
| raise ValueError(f"rope table rows {cos.shape[1]} are not a multiple of {TILE}") |
| if not (torch.isfinite(cos).all() and torch.isfinite(sin).all()): |
| raise ValueError("rope tables contain non-finite values") |
| cos, sin = cos.float(), sin.float() |
| if policy.rope_layout == "meta": |
| cos, sin = rope_cos_sin_hf_to_meta(cos, sin) |
| elif policy.rope_layout != "hf": |
| raise ValueError(f"unknown rope_layout {policy.rope_layout!r}") |
| return cos[:, None].to(torch.bfloat16).contiguous(), sin[:, None].to(torch.bfloat16).contiguous() |
|
|
|
|
| def prepare_rope( |
| mi: Any, |
| policy: TTPolicy, |
| device: Any, |
| buffers: Optional[Dict[str, Any]] = None, |
| *, |
| trans_mat: Any = None, |
| memory_config: Any = None, |
| ) -> Dict[str, Any]: |
| """Device RoPE inputs of one call: ``{"cos", "sin", "trans_mat"}`` bf16 TILE tensors (plan §2.2 ``prepare_rope``). |
| |
| * ``buffers is None``: fresh ``ttnn.from_torch`` uploads (DRAM interleaved); ``trans_mat`` is uploaded from |
| ``tt.weights.rope_transformation_matrix()`` when not given (``None`` under ``rope_layout="hf"``, which ignores it). |
| * ``buffers`` given (N1.6 ``per_call_rows``: the tables change with the prompt, so they are written into the |
| persistent buffers whose addresses the trace captured): ``ttnn.copy_host_to_device_tensor`` into ``buffers["cos"]`` |
| / ``buffers["sin"]`` (same shape / dtype / layout, checked) and the dict is returned unchanged otherwise. |
| """ |
| import ttnn |
|
|
| cos_h, sin_h = rope_tables_host(mi, policy) |
| mem = ttnn.DRAM_MEMORY_CONFIG if memory_config is None else memory_config |
| if buffers is None: |
| out: Dict[str, Any] = { |
| "cos": ttnn.from_torch( |
| cos_h, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=mem |
| ), |
| "sin": ttnn.from_torch( |
| sin_h, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=mem |
| ), |
| "trans_mat": None, |
| } |
| if policy.rope_layout == "meta": |
| if trans_mat is None: |
| from models.experimental.gr00t.tt.weights import rope_transformation_matrix |
|
|
| trans_mat = ttnn.from_torch( |
| rope_transformation_matrix(), |
| dtype=ttnn.bfloat16, |
| layout=ttnn.TILE_LAYOUT, |
| device=device, |
| memory_config=mem, |
| ) |
| out["trans_mat"] = trans_mat |
| return out |
| for key in ("cos", "sin", "trans_mat"): |
| if key not in buffers: |
| raise KeyError(f"rope buffers lack {key!r}") |
| for key, host in (("cos", cos_h), ("sin", sin_h)): |
| dev = buffers[key] |
| if tuple(int(d) for d in dev.shape) != tuple(host.shape): |
| raise ValueError(f"rope buffer {key!r} {tuple(dev.shape)} != host table {tuple(host.shape)}") |
| ttnn.copy_host_to_device_tensor(ttnn.from_torch(host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT), dev) |
| if policy.rope_layout == "meta" and buffers["trans_mat"] is None: |
| raise ValueError("rope_layout='meta' needs the persistent trans_mat buffer") |
| return buffers |
|
|
|
|
| |
| |
| |
| def _shape(t: Any) -> Tuple[int, ...]: |
| return tuple(int(d) for d in t.shape) |
|
|
|
|
| def _get_weight(W: Any, name: str) -> Any: |
| getter = getattr(W, "get", None) |
| if getter is None: |
| raise TypeError(f"weights object {type(W).__name__} has no get(name)") |
| t = getter(name) |
| if t is None: |
| raise KeyError(f"plan tensor {name!r} is not loaded") |
| return t |
|
|
|
|
| def _weight_dtype(w: Any) -> str: |
| """``"bf16" | "bfp8_b"`` of a device weight (or of a CPU stand-in carrying a dtype name string).""" |
| dt = getattr(w, "dtype", None) |
| if dt is None: |
| raise TypeError("weight has no dtype") |
| name = dt if isinstance(dt, str) else dtype_name(dt) |
| if name not in WEIGHT_DTYPES: |
| raise ValueError(f"weight dtype {name!r} is not one of {WEIGHT_DTYPES}") |
| return name |
|
|
|
|
| def _check_mat(name: str, t: Any, k: int, n: int) -> None: |
| s = _shape(t) |
| if len(s) == 4 and s[0] == 1 and s[1] == 1: |
| s = s[2:] |
| if s != (k, n): |
| raise ValueError(f"{name} must be [{k}, {n}], got {_shape(t)}") |
|
|
|
|
| def _check_vec(name: str, t: Any, n: int) -> None: |
| s = _shape(t) |
| rows = 1 |
| for d in s[:-1]: |
| rows *= d |
| if s[-1] != n or rows != 1: |
| raise ValueError(f"{name} must be [1, {n}], got {s}") |
|
|
|
|
| def _options_of(W: Any) -> Any: |
| """``TTWeights.options`` (``TransformOptions``) when the weights object carries it, else ``None`` (test dicts).""" |
| return getattr(W, "options", None) |
|
|
|
|
| |
| |
| |
| class Qwen3LayerTT: |
| """One Qwen3 decoder layer of ``cfg.llm`` on the static layout ``dlayout`` (plan §2.2 / §3.3). |
| |
| Args: |
| cfg: ``GR00TConfig`` (``cfg.llm``: hidden 2048, 16 / 8 heads x 128, intermediate 6144, eps 1e-6). |
| layer_idx: executed layer index ``0 .. n_layers_executed - 1``. |
| W: ``tt.weights.TTWeights`` (or any object with ``get(plan_name)`` returning device tensors and, optionally, |
| ``options`` -- the ``TransformOptions`` the bytes were built with). |
| policy: ``TTPolicy`` -- ``rope_layout`` (D8) and ``gate_up_layout`` (D3) must match ``W.options``. |
| dlayout: ``tt.layout.DeviceLayout`` (``s_pad``, ``text_padding``, ``llm_head_dim``, ``mask_presence["llm"]``). |
| mem: memory preset of the intermediates (L1, plan §3 legend). |
| residual_mem: memory preset of the residual outputs (default :func:`residual_mem_for`: DRAM at S >= 384). |
| ckc / norm_ckc / sdpa_ckc: compute-kernel presets of the projections (HiFi2 fp32acc), norms (HiFi4 fp32acc |
| approx-off, mb2 §B.5) and the SDPA (HiFi4, plan §1.2 "LLM SDPA HiFi4"). |
| keep_output: keep (and record as ``debug_taps["llm_layer_{i}"]``) the pre-deepstack-add output even when a |
| deepstack add follows; without it the pre-add tensor is deallocated after the add. |
| collect_taps: keep every intermediate in ``debug_taps`` (:data:`LAYER_TAP_NAMES`) instead of deallocating. |
| """ |
|
|
| def __init__( |
| self, |
| cfg: GR00TConfig, |
| layer_idx: int, |
| W: Any, |
| policy: TTPolicy, |
| dlayout: Any, |
| *, |
| mem: Any = L1, |
| residual_mem: Any = None, |
| ckc: Any = CKC_HIFI2_FP32ACC, |
| norm_ckc: Any = CKC_HIFI4_FP32ACC_NOAPPROX, |
| sdpa_ckc: Any = CKC_SDPA_HIFI4, |
| keep_output: bool = False, |
| collect_taps: bool = False, |
| ) -> None: |
| if not isinstance(cfg, GR00TConfig): |
| raise TypeError(f"cfg must be a GR00TConfig, got {type(cfg).__name__}") |
| if not isinstance(policy, TTPolicy): |
| raise TypeError(f"policy must be a TTPolicy, got {type(policy).__name__}") |
| llm = cfg.llm |
| if not 0 <= int(layer_idx) < llm.n_layers_executed: |
| raise IndexError( |
| f"layer_idx {layer_idx} outside the {llm.n_layers_executed} executed layers of {cfg.version}" |
| ) |
| if llm.n_heads % llm.n_kv_heads: |
| raise ValueError(f"{llm.n_heads} query heads not divisible by {llm.n_kv_heads} kv heads") |
| if not llm.qk_norm: |
| raise ValueError("Qwen3LayerTT implements the q/k-normed Qwen3 layer; cfg.llm.qk_norm is False") |
| if llm.attention_bias: |
| raise ValueError("Qwen3LayerTT has no attention biases (cfg.llm.attention_bias must be False)") |
| if getattr(dlayout, "version", cfg.version) != cfg.version: |
| raise ValueError(f"dlayout is for {dlayout.version}, cfg is {cfg.version}") |
| s_pad = int(dlayout.s_pad) |
| if s_pad <= 0 or s_pad % TILE: |
| raise ValueError(f"dlayout.s_pad={s_pad} is not a positive multiple of {TILE}") |
| if int(dlayout.llm_head_dim) != int(llm.head_dim): |
| raise ValueError(f"dlayout.llm_head_dim {dlayout.llm_head_dim} != cfg.llm.head_dim {llm.head_dim}") |
| text_padding = str(dlayout.text_padding) |
| if text_padding == "right": |
| attn_mode = "causal" |
| elif text_padding == "text_slot": |
| attn_mode = "mask" |
| else: |
| raise ValueError(f"unknown text_padding {text_padding!r}") |
| presence = dict(getattr(dlayout, "mask_presence", {})) |
| if "llm" in presence and bool(presence["llm"]) != (attn_mode == "mask"): |
| raise ValueError( |
| f"dlayout.mask_presence['llm']={presence['llm']} contradicts text_padding={text_padding!r}" |
| ) |
|
|
| self.cfg, self.version, self.layer_idx = cfg, cfg.version, int(layer_idx) |
| self.policy = policy |
| self.s_pad = s_pad |
| self.hidden, self.n_heads, self.n_kv_heads, self.head_dim = ( |
| llm.hidden, |
| llm.n_heads, |
| llm.n_kv_heads, |
| llm.head_dim, |
| ) |
| self.inter, self.eps = llm.intermediate, float(llm.rms_eps) |
| self.qkv_width = (self.n_heads + 2 * self.n_kv_heads) * self.head_dim |
| self.scale = float(self.head_dim) ** -0.5 |
| self.attn_mode = attn_mode |
| self.rope_layout = policy.rope_layout |
| self.swiglu = swiglu_mode(s_pad) |
| self.gate_up_layout = resolve_gate_up_layout(policy, s_pad) |
| if self.gate_up_layout != gate_up_layout_for(self.swiglu): |
| raise ValueError( |
| f"S_pad={s_pad} runs the {self.swiglu!r} SwiGLU, which needs gate_up_layout=" |
| f"{gate_up_layout_for(self.swiglu)!r}; policy resolves to {self.gate_up_layout!r} (plan D3)" |
| ) |
| opts = _options_of(W) |
| if opts is not None: |
| if opts.rope_layout != self.rope_layout: |
| raise ValueError( |
| f"weights were built with rope_layout={opts.rope_layout!r}, policy says {self.rope_layout!r}" |
| ) |
| if opts.gate_up_layout != self.gate_up_layout: |
| raise ValueError( |
| f"weights were built with gate_up_layout={opts.gate_up_layout!r}, S_pad={s_pad} needs " |
| f"{self.gate_up_layout!r}" |
| ) |
| self.sdpa_chunks = llm_sdpa_chunks(s_pad, self.head_dim) |
| self.mem = mem |
| self.residual_mem = residual_mem_for(s_pad) if residual_mem is None else residual_mem |
| self.ckc, self.norm_ckc, self.sdpa_ckc = ckc, norm_ckc, sdpa_ckc |
| self.keep_output, self.collect_taps = bool(keep_output), bool(collect_taps) |
| self.debug_taps: Dict[str, Any] = {} |
|
|
| names = layer_weight_names(self.layer_idx) |
| self.names = names |
| w = {stem: _get_weight(W, name) for stem, name in names.items()} |
| _check_mat(names["Wqkv"], w["Wqkv"], self.hidden, self.qkv_width) |
| _check_mat(names["Wo"], w["Wo"], self.n_heads * self.head_dim, self.hidden) |
| _check_mat(names["W_gate_up"], w["W_gate_up"], self.hidden, 2 * self.inter) |
| _check_mat(names["W_down"], w["W_down"], self.inter, self.hidden) |
| _check_vec(names["input_ln.w"], w["input_ln.w"], self.hidden) |
| _check_vec(names["post_attn_ln.w"], w["post_attn_ln.w"], self.hidden) |
| _check_vec(names["q_norm.w"], w["q_norm.w"], self.head_dim) |
| _check_vec(names["k_norm.w"], w["k_norm.w"], self.head_dim) |
| self.w = w |
| |
| self.dtypes = {stem: _weight_dtype(w[stem]) for stem in ("Wqkv", "Wo", "W_gate_up", "W_down")} |
| matmul_spec(s_pad, self.hidden, self.qkv_width, self.dtypes["Wqkv"]) |
| matmul_spec(s_pad, self.n_heads * self.head_dim, self.hidden, self.dtypes["Wo"]) |
| matmul_spec(s_pad, self.inter, self.hidden, self.dtypes["W_down"]) |
| if self.swiglu == "split": |
| matmul_spec(s_pad, self.hidden, 2 * self.inter, self.dtypes["W_gate_up"]) |
| else: |
| swiglu_minimal_blocking(s_pad) |
| self._swiglu_cfg: Any = None |
|
|
| |
| def _tap(self, name: str, t: Any) -> None: |
| if self.collect_taps: |
| self.debug_taps[name] = t |
|
|
| def _free(self, *ts: Any) -> None: |
| if self.collect_taps: |
| return |
| import ttnn |
|
|
| for t in ts: |
| ttnn.deallocate(t) |
|
|
| def _swiglu_config(self) -> Any: |
| if self._swiglu_cfg is None: |
| self._swiglu_cfg = swiglu_minimal_config(self.s_pad, self.hidden, self.inter) |
| return self._swiglu_cfg |
|
|
| def _check_inputs(self, x: Any, cos: Any, sin: Any, trans_mat: Any, llm_mask: Any, deepstack_add: Any) -> None: |
| xs = _shape(x) |
| if xs != (1, 1, self.s_pad, self.hidden): |
| raise ValueError(f"x must be [1, 1, {self.s_pad}, {self.hidden}], got {xs}") |
| for name, t in (("cos", cos), ("sin", sin)): |
| ts = _shape(t) |
| if ts != (1, 1, self.s_pad, self.head_dim): |
| raise ValueError(f"{name} must be [1, 1, {self.s_pad}, {self.head_dim}] (prepare_rope), got {ts}") |
| if self.rope_layout == "meta" and trans_mat is None: |
| raise ValueError("rope_layout='meta' needs trans_mat (TTWeights.constants(layout)['trans_mat'])") |
| if (llm_mask is not None) != (self.attn_mode == "mask"): |
| raise ValueError( |
| f"attn_mode={self.attn_mode!r} ({self.version} text_padding): llm_mask must be " |
| f"{'given' if self.attn_mode == 'mask' else 'None'} (mask presence is fixed per layout, plan D5)" |
| ) |
| if llm_mask is not None: |
| ms = _shape(llm_mask) |
| if len(ms) != 4 or ms[2] != self.s_pad or ms[3] != self.s_pad: |
| raise ValueError(f"llm_mask must be [1, 1, {self.s_pad}, {self.s_pad}], got {ms}") |
| if deepstack_add is not None and _shape(deepstack_add) != xs: |
| raise ValueError(f"deepstack_add must have the residual shape {xs}, got {_shape(deepstack_add)}") |
|
|
| |
| def attention(self, h: Any, cos: Any, sin: Any, trans_mat: Any, llm_mask: Any) -> Any: |
| """``rms_norm(x)`` input ``h`` -> ``linear(Wo)`` output ``[1, 1, S_pad, 2048]`` (plan §3.3 rows 2-5).""" |
| w = self.w |
| qkv = linear(h, w["Wqkv"], mem=self.mem, ckc=self.ckc) |
| self._tap("qkv", qkv) |
| q, k, v = split_heads(qkv, self.n_heads, self.n_kv_heads, mem=self.mem) |
| self._free(qkv) |
| qn = rms_norm(q, w["q_norm.w"], eps=self.eps, mem=self.mem, ckc=self.norm_ckc) |
| kn = rms_norm(k, w["k_norm.w"], eps=self.eps, mem=self.mem, ckc=self.norm_ckc) |
| self._free(q, k) |
| self._tap("q_norm", qn) |
| self._tap("k_norm", kn) |
| qr = rope(qn, cos, sin, trans_mat, layout=self.rope_layout, mem=self.mem) |
| kr = rope(kn, cos, sin, trans_mat, layout=self.rope_layout, mem=self.mem) |
| self._free(qn, kn) |
| self._tap("q_rope", qr) |
| self._tap("k_rope", kr) |
| o = sdpa( |
| qr, |
| kr, |
| v, |
| scale=self.scale, |
| mask=llm_mask, |
| is_causal=(self.attn_mode == "causal"), |
| chunks=self.sdpa_chunks, |
| ckc=self.sdpa_ckc, |
| mem=self.mem, |
| ) |
| self._free(qr, kr, v) |
| self._tap("sdpa", o) |
| a = concat_heads(o, mem=self.mem) |
| self._free(o) |
| p = linear(a, w["Wo"], mem=self.mem, ckc=self.ckc) |
| self._free(a) |
| self._tap("attn", p) |
| return p |
|
|
| def mlp(self, h: Any) -> Any: |
| """SwiGLU + down projection of the post-attention-norm output ``h`` (plan §3.3 row 6, D3, mb1 Table 4f).""" |
| import ttnn |
|
|
| w = self.w |
| if self.swiglu == "fused": |
| f = ttnn.experimental.minimal_matmul( |
| h, |
| w["W_gate_up"], |
| config=self._swiglu_config(), |
| memory_config=resolve_mem(self.mem), |
| dtype=ttnn.bfloat16, |
| compute_kernel_config=resolve_ckc(self.ckc), |
| fuse_swiglu=True, |
| ) |
| else: |
| gu = linear(h, w["W_gate_up"], mem=self.mem, ckc=self.ckc) |
| gate = ttnn.slice(gu, [0, 0, 0, 0], [1, 1, self.s_pad, self.inter], memory_config=resolve_mem(self.mem)) |
| up = ttnn.slice( |
| gu, [0, 0, 0, self.inter], [1, 1, self.s_pad, 2 * self.inter], memory_config=resolve_mem(self.mem) |
| ) |
| self._free(gu) |
| f = ttnn.mul( |
| gate, |
| up, |
| input_tensor_a_activations=[ttnn.UnaryWithParam(ttnn.UnaryOpType.SILU)], |
| memory_config=resolve_mem(self.mem), |
| ) |
| self._free(gate, up) |
| if _shape(f) != (1, 1, self.s_pad, self.inter): |
| raise RuntimeError(f"SwiGLU output {_shape(f)} != [1, 1, {self.s_pad}, {self.inter}]") |
| self._tap("ff_hidden", f) |
| m = linear(f, w["W_down"], mem=self.mem, ckc=self.ckc) |
| self._free(f) |
| self._tap("mlp", m) |
| return m |
|
|
| |
| def forward( |
| self, |
| x: Any, |
| cos: Any, |
| sin: Any, |
| trans_mat: Any = None, |
| llm_mask: Any = None, |
| deepstack_add: Any = None, |
| ) -> Any: |
| """``x [1, 1, S_pad, 2048]`` bf16 -> the layer output (after the optional deepstack add). |
| |
| ``debug_taps["llm_layer_{i}"]`` is the output **before** the deepstack add (recorded when no add follows or |
| when ``keep_output``); internals are recorded only with ``collect_taps``. |
| """ |
| self._check_inputs(x, cos, sin, trans_mat, llm_mask, deepstack_add) |
| self.debug_taps = {} |
| w = self.w |
| h = rms_norm(x, w["input_ln.w"], eps=self.eps, mem=self.mem, ckc=self.norm_ckc) |
| self._tap("input_layernorm", h) |
| p = self.attention(h, cos, sin, trans_mat, llm_mask) |
| self._free(h) |
| x1 = residual_add(x, p, mem=self.residual_mem) |
| self._free(p) |
| self._tap("residual_attn", x1) |
| h2 = rms_norm(x1, w["post_attn_ln.w"], eps=self.eps, mem=self.mem, ckc=self.norm_ckc) |
| self._tap("post_attention_layernorm", h2) |
| m = self.mlp(h2) |
| self._free(h2) |
| out = residual_add(x1, m, mem=self.residual_mem) |
| self._free(m, x1) |
| tap_name = f"llm_layer_{self.layer_idx}" |
| if deepstack_add is None: |
| self.debug_taps[tap_name] = out |
| return out |
| added = residual_add(out, deepstack_add, mem=self.residual_mem) |
| if self.keep_output or self.collect_taps: |
| self.debug_taps[tap_name] = out |
| else: |
| import ttnn |
|
|
| ttnn.deallocate(out) |
| self.debug_taps["deepstack_out"] = added |
| return added |
|
|
| __call__ = forward |
|
|
| def describe(self) -> Dict[str, Any]: |
| """JSON-able configuration summary (results files / benchmarks).""" |
| return { |
| "version": self.version, |
| "layer_idx": self.layer_idx, |
| "s_pad": self.s_pad, |
| "hidden": self.hidden, |
| "n_heads": self.n_heads, |
| "n_kv_heads": self.n_kv_heads, |
| "head_dim": self.head_dim, |
| "intermediate": self.inter, |
| "eps": self.eps, |
| "scale": self.scale, |
| "attn_mode": self.attn_mode, |
| "sdpa_chunks": list(self.sdpa_chunks), |
| "rope_layout": self.rope_layout, |
| "swiglu": self.swiglu, |
| "gate_up_layout": self.gate_up_layout, |
| "swiglu_blocking": list(swiglu_minimal_blocking(self.s_pad)) if self.swiglu == "fused" else None, |
| "weight_dtypes": dict(self.dtypes), |
| "residual_mem": getattr(self.residual_mem, "kind", str(self.residual_mem)), |
| "mem": getattr(self.mem, "kind", str(self.mem)), |
| } |
|
|
|
|
| |
| |
| |
| class Qwen3StackTT: |
| """``cfg.llm.n_layers_executed`` layers + N1.7 deepstack adds + final RMSNorm iff ``cfg.llm.final_norm_applied`` |
| (plan §2.2 ``Qwen3StackTT``; ``reference/qwen3.py`` ``Qwen3Decoder``). |
| |
| Args: |
| cfg, W, policy, dlayout: as :class:`Qwen3LayerTT`. |
| tap_layers: layers whose pre-deepstack-add output is kept as ``debug_taps["llm_layer_{i}"]`` (default: first |
| and last -- the ``llm_layer_first`` / ``llm_layer_last`` golden taps); other layer outputs are deallocated |
| once consumed. |
| detail_layers: layers built with ``collect_taps=True`` (internals ``llm_layer{i}_<name>``; default none). |
| layer_kw: forwarded to every :class:`Qwen3LayerTT` (``mem``, ``residual_mem``, ``ckc`` ...). |
| """ |
|
|
| def __init__( |
| self, |
| cfg: GR00TConfig, |
| W: Any, |
| policy: TTPolicy, |
| dlayout: Any, |
| *, |
| tap_layers: Optional[Sequence[int]] = None, |
| detail_layers: Sequence[int] = (), |
| **layer_kw: Any, |
| ) -> None: |
| if not isinstance(cfg, GR00TConfig): |
| raise TypeError(f"cfg must be a GR00TConfig, got {type(cfg).__name__}") |
| llm = cfg.llm |
| n = int(llm.n_layers_executed) |
| if n <= 0: |
| raise ValueError(f"{cfg.version}: n_layers_executed must be positive, got {n}") |
| taps = (0, n - 1) if tap_layers is None else tuple(int(i) for i in tap_layers) |
| details = tuple(int(i) for i in detail_layers) |
| for i in taps + details: |
| if not 0 <= i < n: |
| raise IndexError(f"tap / detail layer {i} outside 0..{n - 1}") |
| self.deepstack_add_after_layers: Tuple[int, ...] = tuple(int(i) for i in llm.deepstack_add_after_layers) |
| for i in self.deepstack_add_after_layers: |
| if not 0 <= i < n: |
| raise IndexError(f"deepstack_add_after_layers entry {i} outside 0..{n - 1}") |
| self.cfg, self.version, self.n_layers = cfg, cfg.version, n |
| self.policy, self.dlayout = policy, dlayout |
| self.tap_layers, self.detail_layers = taps, details |
| self.layers: List[Qwen3LayerTT] = [ |
| Qwen3LayerTT(cfg, i, W, policy, dlayout, keep_output=(i in taps), collect_taps=(i in details), **layer_kw) |
| for i in range(n) |
| ] |
| first = self.layers[0] |
| self.s_pad, self.hidden, self.eps = first.s_pad, first.hidden, first.eps |
| self.attn_mode, self.rope_layout, self.swiglu = first.attn_mode, first.rope_layout, first.swiglu |
| self.mem, self.norm_ckc = first.mem, first.norm_ckc |
| self.final_norm_applied = bool(llm.final_norm_applied) |
| self.final_norm_w: Any = None |
| if self.final_norm_applied: |
| self.final_norm_w = _get_weight(W, FINAL_NORM_NAME) |
| _check_vec(FINAL_NORM_NAME, self.final_norm_w, self.hidden) |
| self.debug_taps: Dict[str, Any] = {} |
|
|
| def forward( |
| self, |
| x: Any, |
| cos: Any, |
| sin: Any, |
| trans_mat: Any = None, |
| llm_mask: Any = None, |
| deepstack_placed: Optional[Sequence[Any]] = None, |
| ) -> Any: |
| """``inputs_embeds [1, 1, S_pad, 2048]`` -> ``backbone features [1, 1, S_pad, 2048]``. |
| |
| ``deepstack_placed``: exactly ``len(cfg.llm.deepstack_add_after_layers)`` tensors of the residual shape (the |
| merger outputs placed at the image slots, zeros elsewhere; ``TokenSplicer.place_into_zeros``), ``None`` / empty |
| for N1.5 / N1.6. Taps: ``llm_layer_{i}`` for ``tap_layers`` (pre-add), ``llm_layer{i}_<name>`` internals for |
| ``detail_layers``, ``llm_final_norm`` when applied, ``backbone_features``. |
| """ |
| ds: Tuple[Any, ...] = tuple(deepstack_placed) if deepstack_placed is not None else () |
| if len(ds) != len(self.deepstack_add_after_layers): |
| raise ValueError( |
| f"{self.version}: {len(ds)} deepstack tensors for deepstack_add_after_layers=" |
| f"{self.deepstack_add_after_layers}" |
| ) |
| self.debug_taps = {} |
| import ttnn |
|
|
| cur = x |
| for i, layer in enumerate(self.layers): |
| add = ds[self.deepstack_add_after_layers.index(i)] if i in self.deepstack_add_after_layers else None |
| nxt = layer(cur, cos, sin, trans_mat, llm_mask, add) |
| kept = layer.debug_taps.get(f"llm_layer_{i}") |
| if kept is not None and i in self.tap_layers: |
| self.debug_taps[f"llm_layer_{i}"] = kept |
| if i in self.detail_layers: |
| for name in LAYER_TAP_NAMES: |
| if name in layer.debug_taps: |
| self.debug_taps[f"llm_layer{i}_{name}"] = layer.debug_taps[name] |
| if cur is not x and cur is not self.debug_taps.get(f"llm_layer_{i - 1}"): |
| ttnn.deallocate(cur) |
| cur = nxt |
| if self.final_norm_applied: |
| feats = rms_norm(cur, self.final_norm_w, eps=self.eps, mem=self.layers[-1].residual_mem, ckc=self.norm_ckc) |
| if cur is not x and cur is not self.debug_taps.get(f"llm_layer_{self.n_layers - 1}"): |
| ttnn.deallocate(cur) |
| self.debug_taps["llm_final_norm"] = feats |
| else: |
| feats = cur |
| self.debug_taps["backbone_features"] = feats |
| return feats |
|
|
| __call__ = forward |
|
|
| def describe(self) -> Dict[str, Any]: |
| d = self.layers[0].describe() |
| d.pop("layer_idx") |
| d.update( |
| n_layers=self.n_layers, |
| final_norm_applied=self.final_norm_applied, |
| deepstack_add_after_layers=list(self.deepstack_add_after_layers), |
| tap_layers=list(self.tap_layers), |
| detail_layers=list(self.detail_layers), |
| ) |
| return d |
|
|
|
|
| __all__ = [ |
| "ATTN_MODES", |
| "FINAL_NORM_NAME", |
| "LAYER_TAP_NAMES", |
| "LAYER_WEIGHT_STEMS", |
| "LLM_SDPA_CLASS", |
| "RESIDUAL_DRAM_MIN_S_PAD", |
| "SWIGLU_FUSED_MIN_S_PAD", |
| "SWIGLU_MINIMAL_BLOCKING", |
| "SWIGLU_MODES", |
| "Qwen3LayerTT", |
| "Qwen3StackTT", |
| "gate_up_layout_for", |
| "layer_weight_names", |
| "llm_sdpa_chunks", |
| "padded_to_rows", |
| "prepare_rope", |
| "residual_mem_for", |
| "rope_tables_host", |
| "rows_to_padded", |
| "stack_weight_names", |
| "swiglu_minimal_blocking", |
| "swiglu_minimal_config", |
| "swiglu_mode", |
| ] |
|
|
