diff --git a/code/models/common/models/llama32_1b/model.py b/code/models/common/models/llama32_1b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..d75483410bca1cc1e03f53db716f246cc7222040 --- /dev/null +++ b/code/models/common/models/llama32_1b/model.py @@ -0,0 +1,1138 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2 Llama 3.2-1B-Instruct — native stack (no ``models/tt_transformers`` imports). + +Architecture: standard Llama 1D transformer, same topology family as Llama 3.2-3B / 3.1-8B. + hidden=2048, layers=16, n_heads=32, n_kv_heads=8, head_dim=64, + intermediate=8192, vocab=128256, rope_theta=500000, RoPE llama3-scaled, + tie_word_embeddings=True (the provider explicitly sources the LM head from the + embedding table when the checkpoint declares tied embeddings). + +Mesh compatibility: N150 (1×1), N300 (1×2), T3K (1×8) — 32 heads / 8 KV heads +divide all three. + +TTTv1 source for precision recipes: + ``models/tt_transformers/tt/model_config.py :: DecodersPrecision`` + (Llama-3 group: ``accuracy()`` lines 130-159, ``performance()`` lines 208-218) +""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D +from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + +# ============================================================================= +# Helpers +# ============================================================================= + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +# ============================================================================= +# TransformerBlock1D — single decoder layer +# ============================================================================= + + +@dataclass +class TransformerBlock1DConfig: + attention_norm_config: RMSNorm1DConfig + attention_config: Attention1DConfig + ff_norm_config: RMSNorm1DConfig + mlp_config: MLP1DConfig + + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtype: ttnn.DataType | None = None + + +class TransformerBlock1D(LightweightModule): + def __init__( + self, + attention_norm: RMSNorm1D, + attention: Attention1D, + ff_norm: RMSNorm1D, + feed_forward: MLP1D, + decode_residual_memcfg: ttnn.MemoryConfig | None = None, + prefill_residual_memcfg: ttnn.MemoryConfig | None = None, + activation_dtype: ttnn.DataType | None = None, + ): + super().__init__() + self.attention_norm = attention_norm + self.attention = attention + self.ff_norm = ff_norm + self.feed_forward = feed_forward + self.decode_residual_memcfg = decode_residual_memcfg + self.prefill_residual_memcfg = prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtype = activation_dtype + + @classmethod + def from_config(cls, config: TransformerBlock1DConfig) -> TransformerBlock1D: + return cls( + attention_norm=RMSNorm1D.from_config(config.attention_norm_config), + attention=Attention1D.from_config(config.attention_config), + ff_norm=RMSNorm1D.from_config(config.ff_norm_config), + feed_forward=MLP1D.from_config(config.mlp_config), + decode_residual_memcfg=config.decode_residual_memcfg, + prefill_residual_memcfg=config.prefill_residual_memcfg, + activation_dtype=config.activation_dtype, + ) + + def decode_forward(self, x: ttnn.Tensor, current_pos, rot_mats, page_table) -> ttnn.Tensor: + residual = x + + x = _all_gather_rmsnorm_tensor( + self.attention_norm, x, memory_config=self.attention_norm.config.decode_memory_config + ) + attn_in = self.attention_norm.decode_forward(x) + attn_out = self.attention.decode_forward(attn_in, current_pos, rot_mats, page_table=page_table) + attn_out = ttnn.to_memory_config(attn_out, self.decode_residual_memcfg) + + hidden_states = ttnn.add(residual, attn_out, memory_config=self.decode_residual_memcfg) + residual = hidden_states + + hidden_states = _all_gather_rmsnorm_tensor( + self.ff_norm, hidden_states, memory_config=self.ff_norm.config.decode_memory_config + ) + hidden_states = self.ff_norm.decode_forward(hidden_states) + ttnn.deallocate(attn_out) + hidden_states = self.feed_forward.decode_forward(hidden_states) + + out = ttnn.add( + residual, + hidden_states, + memory_config=self.decode_residual_memcfg, + dtype=self.activation_dtype or ttnn.bfloat16, + ) + return out + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats, + user_id, + page_table, + chunk_page_table, + chunk_start_idx, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + residual = x + + attn_in = self.attention_norm.prefill_forward(x) + attn_in = _all_gather_rmsnorm_tensor(self.attention_norm, attn_in) + attn_out = self.attention.prefill_forward( + attn_in, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + attn_out = ttnn.to_memory_config(attn_out, self.prefill_residual_memcfg) + + hidden_states = ttnn.add(residual, attn_out, memory_config=self.prefill_residual_memcfg) + residual = hidden_states + x.deallocate(True) + + hidden_states = self.ff_norm.prefill_forward(hidden_states) + hidden_states = _all_gather_rmsnorm_tensor(self.ff_norm, hidden_states) + ttnn.deallocate(attn_out) + hidden_states = self.feed_forward.prefill_forward(hidden_states) + + out = ttnn.add( + residual, + hidden_states, + memory_config=self.prefill_residual_memcfg, + dtype=self.activation_dtype or ttnn.bfloat16, + ) + return out + + def forward( + self, + x, + current_pos=None, + rot_mats=None, + user_id=0, + mode="decode", + page_table=None, + chunk_page_table=None, + chunk_start_idx=None, + batch_size: int = 1, + ): + if mode == "prefill": + return self.prefill_forward( + x, rot_mats, user_id, page_table, chunk_page_table, chunk_start_idx, batch_size=batch_size + ) + return self.decode_forward(x, current_pos, rot_mats, page_table) + + +# ============================================================================= +# RMSNorm gather helper +# ============================================================================= + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, x: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + return x + + if memory_config is None: + memory_config = x.memory_config() + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + # CCL pipelining matched to the proven mistral_7b / qwen25_7b N300 recipe + # (num_links=1, chunks_per_sync=24, num_workers_per_link=4) for parity with the + # reference ports. Measured neutral vs the prior (num_links=2, chunks_per_sync=10, + # num_workers_per_link=2) config on N300 batch-1 decode (perf-tuning.md §Axis-3: the + # per-layer gather is invisible in the decode-step budget on N150/N300; it only bites + # at T3K+ scale), but kept to avoid oversubscribing N300's single inter-chip link. + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + + +# ============================================================================= +# PrecisionConfig — TTTv1-matched recipes for Llama 3.2 1B Instruct +# ============================================================================= + +_LOFI_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, +) + + +@dataclass(frozen=True) +class Llama32_1BPrecisionConfig: + """Per-layer precision + math-fidelity recipe for Llama 3.2 1B Instruct. + + Two module-level recipes: :data:`LLAMA32_1B_ACCURACY` and :data:`LLAMA32_1B_PERFORMANCE`. + Pass one to the product adaptor via ``optimizations=``. + + Attention compute-kernel configs are absent: TTTv1 uses HIFI2 QKV/O decode, + HIFI4 SDPA prefill, HIFI2 SDPA decode — matching ``Attention1D``'s TTTv2 defaults. + """ + + wqkv_dtype: ttnn.DataType = ttnn.bfloat8_b + wo_dtype: ttnn.DataType = ttnn.bfloat8_b + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat8_b + mlp_w2_dtype: ttnn.DataType = ttnn.bfloat8_b + # None → MLP1D default HIFI2_FP16 (matches TTTv1 LI_FF1_FF3 / LI_FF2 accuracy) + mlp_ff1_3_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + mlp_ff2_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + + +# TTTv1 DecodersPrecision.accuracy("Llama-3.2-1B-Instruct") (model_config.py Llama-3 group): +# wqkv=BFP8, wo=BFP8, kv_cache=BFP8, mlp_w1_w3=BFP8, mlp_w2=BFP8, +# LI_FF1_FF3=HIFI2_FP16, LI_FF2=HIFI2_FP16, SDPA_prefill=HIFI4 (Attention1D default) +LLAMA32_1B_ACCURACY = Llama32_1BPrecisionConfig() + +# TTTv1 DecodersPrecision.performance("Llama-3.2-1B-Instruct"): +# FF1_FF3 → BFP4, LI_FF1_FF3 → LOFI; all other fields same as accuracy. +LLAMA32_1B_PERFORMANCE = Llama32_1BPrecisionConfig( + mlp_w1_w3_dtype=ttnn.bfloat4_b, + mlp_ff1_3_compute_kernel_cfg=_LOFI_COMPUTE_KERNEL_CFG, +) + + +@dataclass +class Llama32_1BPagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (duck-typed; matches Attention1D's expected interface).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Llama32_1BTransformer1DConfig: + """Complete tensor-graph configuration for Llama-3.2-1B.""" + + n_layers: int + vocab_size: int + max_batch_size: int + max_seq_len: int + dim: int + num_devices: int + mesh_device: ttnn.MeshDevice + embedding_config: Embedding1DConfig + rope_config: Rope1DConfig + block_configs: list[TransformerBlock1DConfig] + norm_config: RMSNorm1DConfig + lm_head_config: LMHead1DConfig + sampling_config: Sampling1DConfig | None = None + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtypes: list[ttnn.DataType | None] = field(default_factory=list) + tt_ccl: Any = None + cache_path: str | None = None + + +@dataclass(frozen=True) +class Llama32_1BModelParameters: + """Resolved checkpoint dimensions used while building module configs.""" + + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + max_batch_size: int + max_seq_len: int + + +@dataclass(frozen=True) +class Llama32_1BLayerWeights: + """Provider-converted tensors for one decoder layer.""" + + wqkv: torch.Tensor + wo: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Llama32_1BWeights: + """Provider-neutral tensor inputs consumed by the model config builder.""" + + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Llama32_1BLayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +# ============================================================================= +# Wormhole tuning +# ============================================================================= + + +@dataclass +class _Llama32_1BWHTuning: + mlp_prefill_len_cutoff: int | None = None + mlp_decode_spill_w1_to_dram: bool = False + # Use ttnn.experimental.minimal_matmul for QKV + W2 prefill matmuls above seq_len > 128 (TTTv1 + # parity, PLAN_01). A/B escape hatch: set DISABLE_MINIMAL_MATMUL=1 to force ttnn.linear. ~Parity on + # tiny 1B (matmuls small vs fixed overhead) but kept on for consistency + long-prompt prefill. + prefill_minimal_matmul: bool = True + + +def _resolve_llama32_1b_wh_tuning(*, num_dev: int, max_batch_size: int) -> _Llama32_1BWHTuning: + """Pick WH tuning knobs. + + TTTv1 (model_config.py): prefill_len_cutoff defaults to 1024 on Wormhole and is only + reduced to 512 for an explicit model/device allowlist (Llama-3.1-8B, Llama-3.2-11B, + Mistral-7B, gemma-3-*, Phi-4 on N150; Qwen2.5-7B/VL-7B/Phi-4 on N300; Mixtral on T3K). + Llama-3.2-1B is in none of them, so it stays at 1024 on N150/N300/T3K — unlike the + llama32_3b port, which uses 512 on N150 as a 3B-specific choice. + """ + t = _Llama32_1BWHTuning() + t.mlp_prefill_len_cutoff = 1024 + t.mlp_decode_spill_w1_to_dram = False + t.prefill_minimal_matmul = not os.environ.get("DISABLE_MINIMAL_MATMUL") + logger.info( + f"MLP tuning for Llama-3.2-1B on {num_dev} device(s): " + f"prefill_len_cutoff={t.mlp_prefill_len_cutoff}, " + f"decode_spill_w1_to_dram={t.mlp_decode_spill_w1_to_dram}" + ) + return t + + +# ============================================================================= +# Layer + head builders +# ============================================================================= + + +def _post_attn_norm_decode_configs( + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding to match MLP1D W1/W3 input.""" + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + memory_config = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return program_config, memory_config + + +def _build_decoder_layer( + *, + idx: int, + weights: Llama32_1BLayerWeights, + mcfg: Llama32_1BModelParameters, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Llama32_1BPrecisionConfig, + paged_block_size: int | None, + paged_max_blocks: int | None, + cache_path: Path | None, + wh: _Llama32_1BWHTuning, + decode_residual_memcfg: ttnn.MemoryConfig, +) -> TransformerBlock1DConfig: + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, + dtype=precision.wqkv_dtype, + cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None, + ) + lazy_wo = _lazy( + weights.wo, + dtype=precision.wo_dtype, + cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None, + ) + + paged_cfg = Llama32_1BPagedAttentionConfig(block_size=paged_block_size, max_num_blocks=paged_max_blocks) + + attention_config = Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=mcfg.n_heads, + n_kv_heads=mcfg.n_kv_heads, + head_dim=mcfg.head_dim, + max_batch_size=mcfg.max_batch_size, + max_seq_len=mcfg.max_seq_len, + use_vllm_paged_kv_cache=True, + paged_attention_config=paged_cfg, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + # TTTv1 parity: Llama-3 family decode SDPA runs HIFI2 with exp_approx_mode=True + # (model_config.py `_default_settings` → SDPA_DECODE=HIFI2, used in BOTH accuracy + # and performance). Attention1D's generic default builds this prog config with + # exp_approx_mode=False, leaving decode SDPA slower than TTTv1. Flip it to match. + decode_sdpa_prg_config=ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=(8, 8), + exp_approx_mode=True, + q_chunk_size=0, + k_chunk_size=0, + ), + prefill_qkv_minimal_matmul=wh.prefill_minimal_matmul, + ) + + mlp_config = MLP1DConfig( + w1=_lazy( + weights.w1, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None, + ), + w2=_lazy( + weights.w2, + dtype=precision.mlp_w2_dtype, + cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None, + ), + w3=_lazy( + weights.w3, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None, + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=mcfg.max_batch_size, + prefill_len_cutoff=wh.mlp_prefill_len_cutoff, + decode_spill_w1_to_dram_before_w3=wh.mlp_decode_spill_w1_to_dram, + w1_w3_dtype=precision.mlp_w1_w3_dtype, + w2_dtype=precision.mlp_w2_dtype, + ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + ff2_compute_kernel_cfg=precision.mlp_ff2_compute_kernel_cfg, + decode_ff2_compute_kernel_cfg=precision.mlp_ff2_compute_kernel_cfg, + prefill_w2_minimal_matmul=wh.prefill_minimal_matmul, + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + dim=mcfg.dim, + hidden_dim=mcfg.hidden_dim, + num_devices=num_dev, + max_batch_size=mcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1DConfig: + lw = _lazy( + weight, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=mcfg.rms_norm_eps, + max_batch_size=mcfg.max_batch_size, + tt_ccl=tt_ccl, + **extra, + ) + + attn_norm = _build_norm(weights.attention_norm, "pre_attn") + ff_norm = _build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ) + + return TransformerBlock1DConfig( + attention_norm_config=attn_norm, + attention_config=attention_config, + ff_norm_config=ff_norm, + mlp_config=mlp_config, + decode_residual_memcfg=decode_residual_memcfg, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtype=None, + ) + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + mcfg: Llama32_1BModelParameters, + lm_head_dtype: ttnn.DataType, + cache_path: Path | None, +) -> LMHead1DConfig: + lm_splits, lm_split_sizes, lm_weights_memcfgs = _build_lm_head_lazy_weights( + mesh_device, + lm_head_weight, + dim=mcfg.dim, + vocab_size=mcfg.vocab_size, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(mcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(mcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, mcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(mcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=mcfg.dim, + max_batch_size=mcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + compute_kernel_config=None, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + ) + + +def _build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """Column-split provider-neutral LM-head weights for ``LMHead1D``.""" + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / TILE_SIZE) * TILE_SIZE + if vocab_size < padded_vocab_size: + torch_w = torch.cat( + [torch_w, torch.zeros(torch_w.shape[0], padded_vocab_size - vocab_size, dtype=torch_w.dtype)], dim=-1 + ) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (TILE_SIZE * dram_size.x)) * (TILE_SIZE * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights = [] + weights_memcfgs = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + device_splits.append(torch_w[:, start : start + split_size]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, f"lm_head_split_{i}_{combined.shape[-1]}") if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs + + +def build_llama32_1b_transformer_1d_config( + *, + mesh_device: ttnn.MeshDevice, + params: Llama32_1BModelParameters, + weights: Llama32_1BWeights, + n_layers: int, + precision: Llama32_1BPrecisionConfig, + cache_path: Path, + paged_attention_config: Llama32_1BPagedAttentionConfig, +) -> Llama32_1BTransformer1DConfig: + """Build the TT model graph from provider-neutral dimensions and tensors.""" + num_devices = mesh_device.get_num_devices() + if num_devices not in (1, 2, 8): + raise ValueError(f"Llama-3.2-1B supports 1, 2, or 8 devices, got {num_devices}") + if params.n_heads % num_devices or params.n_kv_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({params.n_heads}/{params.n_kv_heads}) " + f"must be divisible by device count ({num_devices})" + ) + if len(weights.layers) != n_layers: + raise ValueError(f"Expected {n_layers} decoder layer weight sets, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) if num_devices > 1 else None + topology = default_topology(mesh_device) + embedding_config = Embedding1DConfig( + weights=_lazy( + weights.embedding, + dtype=ttnn.bfloat16, + cache=(cache_path / "embedding", "tok_embeddings"), + ), + mesh_device=mesh_device, + embed_scale=1.0, + ) + rope_config = Rope1DConfig( + cos_matrix=_lazy(weights.rope_cos, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "cos")), + sin_matrix=_lazy(weights.rope_sin, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "sin")), + max_batch_size=params.max_batch_size, + head_dim=params.head_dim, + device=mesh_device, + use_qk_fused=False, + ) + wh = _resolve_llama32_1b_wh_tuning(num_dev=num_devices, max_batch_size=params.max_batch_size) + block_configs = [ + _build_decoder_layer( + idx=index, + weights=weights.layers[index], + mcfg=params, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + paged_block_size=paged_attention_config.block_size, + paged_max_blocks=paged_attention_config.max_num_blocks, + cache_path=cache_path, + wh=wh, + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + ) + for index in range(n_layers) + ] + norm_config = RMSNorm1DConfig( + weight=_lazy(weights.final_norm, dtype=ttnn.bfloat16, cache=(cache_path / "norm", "final")), + mesh_device=mesh_device, + eps=params.rms_norm_eps, + max_batch_size=params.max_batch_size, + tt_ccl=tt_ccl, + ) + lm_head_config = _build_lm_head( + mesh_device=mesh_device, + lm_head_weight=weights.lm_head, + mcfg=params, + lm_head_dtype=precision.lm_head_dtype, + cache_path=cache_path, + ) + sampling_config = Sampling1DConfig( + vocab_size=params.vocab_size, + valid_vocab_size=params.vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=_nearest_32(params.max_batch_size), + allow_force_argmax=False, + pad_to_power_of_2=True, + ) + return Llama32_1BTransformer1DConfig( + n_layers=n_layers, + vocab_size=params.vocab_size, + max_batch_size=params.max_batch_size, + max_seq_len=params.max_seq_len, + dim=params.dim, + num_devices=num_devices, + mesh_device=mesh_device, + embedding_config=embedding_config, + rope_config=rope_config, + block_configs=block_configs, + norm_config=norm_config, + lm_head_config=lm_head_config, + sampling_config=sampling_config, + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtypes=[None] * n_layers, + tt_ccl=tt_ccl, + cache_path=str(cache_path), + ) + + +# ============================================================================= +# Llama32_1BTransformer1D +# ============================================================================= + + +class Llama32_1BTransformer1D(LightweightModule): + """TTTv2 Llama 3.2-1B-Instruct transformer. + + Construct from :class:`Llama32_1BTransformer1DConfig`; the Hugging Face + product adaptor owns checkpoint loading and binds ``model_args``. + """ + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__(self, config: Llama32_1BTransformer1DConfig): + from tqdm import tqdm + + super().__init__() + self.config = config + self.embedding = Embedding1D.from_config(config.embedding_config) + self.rope_setup = RotarySetup1D.from_config(config.rope_config) + self.layers = [ + TransformerBlock1D.from_config(config.block_configs[i]) + for i in tqdm(range(config.n_layers), desc="Building layers") + ] + self.norm = RMSNorm1D.from_config(config.norm_config) + self.lm_head = LMHead1D.from_config(config.lm_head_config) + self.sampling = Sampling1D.from_config(config.sampling_config) if config.sampling_config is not None else None + self.supports_on_device_sampling = self.sampling is not None + self.mesh_device = config.mesh_device + self.tt_ccl = config.tt_ccl or (get_tt_ccl(config.mesh_device) if config.num_devices > 1 else None) + self.vocab_size = config.vocab_size + self.n_layers = config.n_layers + self.num_devices = config.num_devices + self.decode_residual_memcfg = config.decode_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.prefill_residual_memcfg = config.prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtypes = config.activation_dtypes or [None] * config.n_layers + self.model_args = None + + # ========================================================================= + # KV cache binding + # ========================================================================= + + def iter_executor_named_modules(self): + if not hasattr(self, "layers"): + return + for i, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{i}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in self.config.block_configs) + for config in tuple({id(c): c for c in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, block_size=block_size, max_num_blocks=max_num_blocks + ) + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for i, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {i} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {i} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = pair + + # ========================================================================= + # Forward methods — take pre-embedded tensors + # ========================================================================= + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for i, layer in enumerate(self.layers): + x = ttnn.to_memory_config(x, self.decode_residual_memcfg, self.activation_dtypes[i]) + x = layer.decode_forward(x, current_pos, rot_mats, page_table) + + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + x = self.lm_head.forward(x) + return x + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). The batched path always + # returns the full hidden state (get_last_token == -1); the executor does per-slot last-token + # extraction + norm/lm_head so those stages stay bit-identical to the single-user path. + x = x_embed + for i, layer in enumerate(self.layers): + activation_dtype = self.activation_dtypes[i] + if activation_dtype is not None and x.dtype != activation_dtype: + old = x + x = ttnn.typecast(x, activation_dtype) + ttnn.deallocate(old) + x = layer.prefill_forward( + x, + rot_mats, + user_id, + page_table, + chunk_page_table, + chunk_start_idx, + batch_size, + chunk_start_idx_tensor, + ) + + if last_token_index is not None and last_token_slice is None: + raise ValueError("last_token_index is required with a runtime last_token_slice") + if get_last_token == -1 and last_token_slice is None: + return x + + old = x + if last_token_slice is None: + get_last_token_floor = (get_last_token // 32) * 32 + x = ttnn.slice(x, (0, 0, get_last_token_floor, 0), (1, 1, get_last_token_floor + 32, x.shape[-1])) + else: + x = ttnn.slice( + x, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(x.shape[2]) // 32, + ) + ttnn.deallocate(old) + + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + get_last_token_floor = (last_token_idx // 32) * 32 + x = ttnn.slice( + hidden_states, + (0, 0, get_last_token_floor, 0), + (1, 1, get_last_token_floor + 32, hidden_states.shape[-1]), + ) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), + ) + x = ttnn.matmul( + selector, + hidden_states, + dtype=ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + compute_kernel_config=ttnn.init_device_compute_kernel_config( + self.mesh_device.arch(), + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ), + ) + ttnn.deallocate(selector) + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def forward( + self, + x: ttnn.Tensor, + current_pos=None, + rot_mats_global=None, + rot_mats_local=None, + user_id: int = 0, + mode: str = "decode", + page_table=None, + chunk_page_table=None, + chunk_start_idx=None, + get_last_token: int = -1, + batch_size: int = 1, + ) -> ttnn.Tensor: + rot_mats = rot_mats_global + if mode == "prefill": + return self.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + get_last_token=get_last_token, + batch_size=batch_size, + ) + return self.decode_forward(x, current_pos, rot_mats, page_table=page_table) + + # ========================================================================= + # Embedding + output processing helpers (executor contract) + # ========================================================================= + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embedding.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embedding.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) + + +# ============================================================================= +# Public exports +# ============================================================================= + +__all__ = [ + "Llama32_1BPrecisionConfig", + "LLAMA32_1B_ACCURACY", + "LLAMA32_1B_PERFORMANCE", + "Llama32_1BPagedAttentionConfig", + "Llama32_1BLayerWeights", + "Llama32_1BModelParameters", + "Llama32_1BWeights", + "Llama32_1BTransformer1DConfig", + "Llama32_1BTransformer1D", + "build_llama32_1b_transformer_1d_config", + "TransformerBlock1D", + "TransformerBlock1DConfig", + "_all_gather_rmsnorm_tensor", +] diff --git a/code/models/common/models/llama33_70b/README.md b/code/models/common/models/llama33_70b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..2c72085b8529f5dcc219d43590ca39808030b229 --- /dev/null +++ b/code/models/common/models/llama33_70b/README.md @@ -0,0 +1,51 @@ +# Llama 3.3 70B with TTTv2 + +This directory contains the Llama 3.3 70B TTTv2 product path. + +## Construction path + +```text +meta-llama/Llama-3.3-70B-Instruct + -> hf_adaptor.py: provider metadata, tokenizer, and weights + -> model.py: multi-device TTTv2 Llama graph + -> executor.py: thin typed entry point into llama3_executor.py + -> generator.py: vLLM construction, DP lanes, and dispatch +``` + +`model.py` composes reusable embedding, rotary, RMSNorm, attention, MLP, +LM-head, sampling, and TT collective modules. The 70B topology, precision, and +program policy remains model-owned. + +## Executor composition + +The local file retains `Llama33_70BExecutor`, its config, and its builder as +stable imports. `models/common/models/llama3_executor.py` supplies the 70B +family policy and composes `models/common/models/executor.py::ModelExecutor`. + +The family-neutral owner composes paged-KV management, output reading, +prefill/decode runtimes, eager and traced execution, warmup coordination, and +ordered cleanup. The Llama family layer adds only: + +- `SamplingState1D` and complete request history/remap state; +- device-sampling prefill serialization; +- the 70B trace-warmup sequence preference; and +- Q128 top-k priming before traced warmup and after eager warmup. + +## vLLM, DP, and ownership + +The generator builds one model/executor per submesh. DP1 uses one executor; +larger DP uses `LaneGroupExecutor`. `VLLMAdapter` normalizes server calls and +validates the physical KV cache. Each lane owns its TT resources and sampling +state; generator/adapter objects own no TT tensors. + +## Tests + +- `models/common/tests/models/llama33_70b/test_hf_adaptor.py` +- `models/common/tests/models/llama33_70b/test_model_profile.py` +- `models/common/tests/models/llama33_70b/test_demo_contract.py` +- `models/common/tests/models/llama33_70b/test_logits_oracle.py` +- `models/common/tests/models/llama33_70b/test_t3k_batched_prefill_correctness.py` +- `models/common/tests/models/llama33_70b/test_p150x4_smoke.py` +- `models/common/tests/demos/llama33_70b/demo.py` +- `models/common/tests/llm_runtime/test_executor_integration.py` +- `models/common/tests/llm_runtime/test_model_executor.py` diff --git a/code/models/common/models/llama33_70b/__init__.py b/code/models/common/models/llama33_70b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a79f99b792f9a4b7a358f347b970b1dc2f30576c --- /dev/null +++ b/code/models/common/models/llama33_70b/__init__.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +from models.common.models.llama33_70b.hf_adaptor import ( + DEFAULT_HF_MODEL, + Llama33_70BForCausalLM, + Llama33_70BGenerationConfig, + Llama33_70BRuntimeConfig, + from_pretrained, +) +from models.common.models.llama33_70b.executor import ( + Llama33_70BExecutor, + Llama33_70BExecutorConfig, + build_llama33_70b_executor, +) +from models.common.models.llama33_70b.generator import ( + Llama33_70BGenerator, + Llama33_70BGeneratorConfig, + build_llama33_70b_generator, +) +from models.common.models.llama33_70b.model import ( + LLAMA33_70B_ACCURACY, + LLAMA33_70B_PERFORMANCE, + Llama33_70BLayerWeights, + Llama33_70BModelParameters, + Llama33_70BPagedAttentionConfig, + Llama33_70BPrecisionConfig, + Llama33_70BTransformer1D, + Llama33_70BTransformer1DConfig, + Llama33_70BWeights, + build_llama33_70b_transformer_1d_config, +) + +__all__ = [ + "DEFAULT_HF_MODEL", + "LLAMA33_70B_ACCURACY", + "LLAMA33_70B_PERFORMANCE", + "Llama33_70BForCausalLM", + "Llama33_70BGenerationConfig", + "Llama33_70BExecutor", + "Llama33_70BExecutorConfig", + "Llama33_70BGenerator", + "Llama33_70BGeneratorConfig", + "Llama33_70BLayerWeights", + "Llama33_70BModelParameters", + "Llama33_70BPagedAttentionConfig", + "Llama33_70BPrecisionConfig", + "Llama33_70BRuntimeConfig", + "Llama33_70BTransformer1D", + "Llama33_70BTransformer1DConfig", + "Llama33_70BWeights", + "build_llama33_70b_transformer_1d_config", + "build_llama33_70b_executor", + "build_llama33_70b_generator", + "from_pretrained", +] diff --git a/code/models/common/models/llama33_70b/generator.py b/code/models/common/models/llama33_70b/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..f008838d7725e133d14f9e3f16ccd7f61f827be7 --- /dev/null +++ b/code/models/common/models/llama33_70b/generator.py @@ -0,0 +1,506 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Llama 3.3-70B.""" + +from __future__ import annotations + +import atexit +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.llama33_70b.executor import Llama33_70BExecutorConfig, build_llama33_70b_executor +from models.common.models.llama33_70b.hf_adaptor import from_pretrained +from models.common.models.llama33_70b.model import Llama33_70BPagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Llama33_70BGeneratorConfig: + """Validated construction inputs for one vLLM-facing Llama generator.""" + + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Llama33_70BGenerator: + """Adapt vLLM's model interface to the model-owned execution target. + + vLLM constructs this facade with `initialize_vllm_model`, resolves + KV capacity through `allocate_kv_cache`, warms the configured + programs, and then calls `prefill_forward` and + `decode_forward`. Each forward call is normalized by + `VLLMAdapter`, dispatched to the eager or traced executor, and + delegated to ``Llama33_70BExecutor`` or ``LaneGroupExecutor``. + + This class owns dispatch policy but no TT resources. `cleanup` + delegates to the target that owns those resources. + """ + + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "max_device_top_k": 32, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + self._serving_ready_summary_logged = False + self._shutdown_summary_logged = False + self._shutdown_summary_callback = self._log_shutdown_summary + atexit.register(self._shutdown_summary_callback) + + # Public vLLM API + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + """Return the unpadded per-submesh KV token budget for vLLM sizing.""" + + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + device_sampling_enabled: bool = True, + ): + """Build the configured single-lane or data-parallel Llama target.""" + + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_llama33_70b_generator( + Llama33_70BGeneratorConfig( + hf_model=str(hf_model), + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def compile_prefill( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + ) -> None: + """Normalize a vLLM prefill call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.compile_prefill(execution=execution, **normalized) + + def compile_decode( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + ) -> None: + """Normalize a vLLM decode call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.compile_decode(execution=execution, **normalized) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM prefill call.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM decode call.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward( + execution=execution, + read_from_device=read_from_device, + **normalized, + ) + + def read_decode_output( + self, + tt_out: Any, + *, + async_read: bool = False, + ) -> Any: + """Delegate vLLM's raw decode-output read.""" + + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host( + self, + tt_out: Any, + *, + is_tokens: bool = False, + ) -> tuple[Any, Any]: + """Delegate vLLM's asynchronous host-output completion.""" + + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + result = self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + self._log_serving_ready_summary_if_active() + return result + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + result = self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + self._log_serving_ready_summary_if_active() + return result + + def cleanup(self): + """Release every resource owned by the concrete target.""" + + self._log_shutdown_summary() + atexit.unregister(self._shutdown_summary_callback) + return self.target.cleanup() + + # Private implementation + + def _log_serving_ready_summary_if_active(self) -> None: + traced = getattr(self.target, "traced_executor", None) + trace_compiler = getattr(traced, "trace_compiler", None) + if traced is None or self._serving_ready_summary_logged or not getattr(trace_compiler, "trace_active", False): + return + traced.log_runtime_summary(phase="serving_ready") + self._serving_ready_summary_logged = True + + def _log_shutdown_summary(self) -> None: + if self._shutdown_summary_logged: + return + traced = getattr(self.target, "traced_executor", None) + if traced is not None: + traced.log_runtime_summary(phase="shutdown") + self._shutdown_summary_logged = True + + def _select_prefill_execution( + self, + normalized: NormalizedPrefillKwargs, + trace_requested: bool, + ): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_llama33_70b_generator(config: Llama33_70BGeneratorConfig) -> Llama33_70BGenerator: + """Construct lane-local models/executors and compose their shared target surface.""" + + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Llama33_70BPagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + instruct="Instruct" in config.hf_model, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Llama33_70BExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_llama33_70b_executor(llm, executor_config)) + + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Llama33_70BGenerator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + request_state_fields=lane._request_state_fields, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Llama 3.3-70B model must contain at least one attention layer") + + attention_configs = tuple(layer.attention.config for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Llama layer must expose the same KV head shape") + + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/llama33_70b/hf_adaptor.py b/code/models/common/models/llama33_70b/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..a7de941a9c8392290d1c8e90ca8d13963e063d09 --- /dev/null +++ b/code/models/common/models/llama33_70b/hf_adaptor.py @@ -0,0 +1,348 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face product adaptor for the TTTv2 Llama 3.3 70B path.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.models.llama33_70b import weight_utils +from models.common.models.llama33_70b.model import ( + LLAMA33_70B_ACCURACY, + LLAMA33_70B_BH_TP4_CLUSTER_TYPES, + LLAMA33_70B_PERFORMANCE, + Llama33_70BLayerWeights, + Llama33_70BModelParameters, + Llama33_70BPagedAttentionConfig, + Llama33_70BPrecisionConfig, + Llama33_70BTransformer1D, + Llama33_70BWeights, + _llama33_70b_ccl_topology, + build_llama33_70b_transformer_1d_config, +) + +DEFAULT_HF_MODEL = "meta-llama/Llama-3.3-70B-Instruct" + + +@dataclass(frozen=True) +class Llama33_70BGenerationConfig: + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Llama33_70BRuntimeConfig: + model_name: str + model_cache_path: Path | None + max_prefill_chunk_size: int + max_context_len: int + max_seq_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] + trace_prefill_warmup_seq_lens: tuple[int, ...] = () + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + return ( + prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + chat = [] + if isinstance(prompt_text, str): + if system_prompt_text: + chat.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + chat.append({"role": "user", "content": prompt_text}) + else: + chat = prompt_text + try: + return _chat_template_ids(tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True)) + except ValueError: + pass + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Llama33_70BForCausalLM: + model: Llama33_70BTransformer1D + tokenizer: Any + runtime_config: Llama33_70BRuntimeConfig + instruct: bool = True + generation_config: Llama33_70BGenerationConfig = field(default_factory=Llama33_70BGenerationConfig) + + def __post_init__(self): + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + self.generation_config = Llama33_70BGenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=tuple(getattr(self.tokenizer, "stop_tokens", ()) or ()), + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + return encode_prompt( + self.tokenizer, + prompt_text, + system_prompt_text, + instruct=self.instruct if instruct is None else instruct, + ) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_tokenizer(hf_model: str): + tokenizer = AutoTokenizer.from_pretrained(hf_model, local_files_only=os.getenv("CI") == "true") + eos = getattr(tokenizer, "eos_token_id", None) + tokenizer.stop_tokens = [] if eos is None else ([eos] if isinstance(eos, int) else list(eos)) + return tokenizer + + +def _trace_seq_lens(num_devices: int, max_prefill_chunk_size: int, max_seq_len: int) -> tuple[int, ...]: + if num_devices not in (4, 8): + raise ValueError(f"Llama-3.3-70B supports T3K (8 devices) or P150x4 (4 devices), got {num_devices}") + if num_devices == 4: + return (128,) if max_seq_len >= 128 else () + return tuple(length for length in (128, max_prefill_chunk_size) if length <= max_seq_len) + + +def _trace_warmup_seq_lens( + max_prefill_chunk_size: int, + max_seq_len: int, + supported_seq_lens: tuple[int, ...], +) -> tuple[int, ...]: + """Return logical representatives whose first invocation has a trace family.""" + + candidates = (128, max_prefill_chunk_size, 2 * max_prefill_chunk_size) + return tuple( + dict.fromkeys( + length + for length in candidates + if length <= max_seq_len and min(length, max_prefill_chunk_size) in supported_seq_lens + ) + ) + + +def _resolve_supported_sku(*, arch, cluster_type, num_devices: int) -> str: + if arch == ttnn.device.Arch.WORMHOLE_B0 and cluster_type == ttnn.cluster.ClusterType.T3K and num_devices == 8: + return "T3K" + if arch == ttnn.device.Arch.BLACKHOLE and cluster_type in LLAMA33_70B_BH_TP4_CLUSTER_TYPES and num_devices == 4: + return "P150x4" + raise ValueError( + "Llama-3.3-70B supports physical Wormhole T3K (8 devices) or " + "BlackHole P150_X4/P300_X2 as logical P150x4 (4 devices); " + f"got arch={arch}, cluster_type={cluster_type}, num_devices={num_devices}" + ) + + +def _cache_path(hf_model: str, mesh_device, cache_dir: Path | str | None, *, sku: str) -> Path: + if cache_dir is not None: + path = Path(cache_dir) + elif os.getenv("TT_CACHE_PATH"): + path = Path(os.environ["TT_CACHE_PATH"]) + else: + path = Path("model_cache") / hf_model / sku + path.mkdir(parents=True, exist_ok=True) + return path + + +def _stop_token_ids(hf) -> tuple[int, ...]: + eos = getattr(getattr(hf, "generation_config", None), "eos_token_id", None) + if eos is None: + return () + values = (eos,) if isinstance(eos, int) else tuple(eos) + return tuple(dict.fromkeys(int(value) for value in values)) + + +def convert_hf_model_weights( + hf, + hf_config, + *, + n_layers: int, + num_devices: int, + rope_table_len: int, + head_dim: int, +) -> Llama33_70BWeights: + """Extract and convert every HF tensor consumed by the TT graph builder.""" + + base = hf.model + rope_cos, rope_sin = weight_utils.build_rope_cos_sin_torch( + base.rotary_emb, rope_table_len, head_dim, torch.bfloat16 + ) + layers = [] + for layer in base.layers[:n_layers]: + wqkv, wo = weight_utils.attention_wqkv_wo_from_hf_layer(layer.self_attn, num_devices) + w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(layer.mlp) + layers.append( + Llama33_70BLayerWeights( + wqkv=wqkv, + wo=wo, + w1=w1, + w2=w2, + w3=w3, + attention_norm=weight_utils.rms_weight_torch(layer.input_layernorm).to(torch.bfloat16), + ff_norm=weight_utils.rms_weight_torch(layer.post_attention_layernorm).to(torch.bfloat16), + ) + ) + return Llama33_70BWeights( + embedding=weight_utils.embed_tokens_torch(base.embed_tokens), + rope_cos=rope_cos, + rope_sin=rope_sin, + layers=tuple(layers), + final_norm=weight_utils.rms_weight_torch(base.norm).to(torch.bfloat16), + lm_head=hf.lm_head.weight.detach().to(torch.bfloat16).clone(), + ) + + +def from_pretrained( + mesh_device, + *, + hf_model: str = DEFAULT_HF_MODEL, + instruct: bool = True, + max_batch_size: int = 32, + max_seq_len: int = 4096, + optimizations: str | Llama33_70BPrecisionConfig = "accuracy", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config: Llama33_70BPagedAttentionConfig | None = None, + cache_dir: Path | str | None = None, +) -> Llama33_70BForCausalLM: + del dtype + num_devices = mesh_device.get_num_devices() + arch = mesh_device.arch() + sku = _resolve_supported_sku( + arch=arch, + cluster_type=ttnn.cluster.get_cluster_type(), + num_devices=num_devices, + ) + _llama33_70b_ccl_topology(mesh_device) + ttnn.SetDefaultDevice(mesh_device) + load_kwargs = {"local_files_only": os.getenv("CI") == "true"} + hf_config = AutoConfig.from_pretrained(hf_model, **load_kwargs) + if hf_config.num_attention_heads % num_devices or hf_config.num_key_value_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({hf_config.num_attention_heads}/{hf_config.num_key_value_heads}) " + f"must be divisible by device count ({num_devices})" + ) + hf = AutoModelForCausalLM.from_pretrained(hf_model, torch_dtype=torch.bfloat16, **load_kwargs) + hf.eval() + resolved_layers = hf_config.num_hidden_layers if n_layers is None else n_layers + precision = ( + optimizations + if isinstance(optimizations, Llama33_70BPrecisionConfig) + else (LLAMA33_70B_PERFORMANCE if optimizations == "performance" else LLAMA33_70B_ACCURACY) + ) + if not isinstance(precision, Llama33_70BPrecisionConfig): + raise TypeError("optimizations must be 'accuracy', 'performance', or Llama33_70BPrecisionConfig") + cache_path = _cache_path(hf_model, mesh_device, cache_dir, sku=sku) + if paged_attention_config is None: + block_size = 32 + paged_attention_config = Llama33_70BPagedAttentionConfig( + block_size=block_size, + max_num_blocks=((max_seq_len + block_size - 1) // block_size) * max_batch_size, + ) + head_dim = int(getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)) + params = Llama33_70BModelParameters( + dim=hf_config.hidden_size, + n_heads=hf_config.num_attention_heads, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + hidden_dim=hf_config.intermediate_size, + vocab_size=hf_config.vocab_size, + rms_norm_eps=hf_config.rms_norm_eps, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + ) + rope_table_len = ((max(max_seq_len * 2, 8192) + 127) // 128) * 128 + weights = convert_hf_model_weights( + hf, + hf_config, + n_layers=resolved_layers, + num_devices=num_devices, + rope_table_len=rope_table_len, + head_dim=head_dim, + ) + model_config = build_llama33_70b_transformer_1d_config( + mesh_device=mesh_device, + params=params, + weights=weights, + n_layers=resolved_layers, + precision=precision, + cache_path=cache_path, + paged_attention_config=paged_attention_config, + ) + tokenizer = load_tokenizer(hf_model) + stop_token_ids = _stop_token_ids(hf) + if stop_token_ids: + tokenizer.stop_tokens = list(stop_token_ids) + trace_prefill_supported_seq_lens = _trace_seq_lens(num_devices, 2048, max_seq_len) + runtime_config = Llama33_70BRuntimeConfig( + model_name=Path(hf_model).name, + model_cache_path=cache_path, + max_prefill_chunk_size=2048, + max_context_len=int(hf_config.max_position_embeddings), + max_seq_len=max_seq_len, + trace_prefill_supported_seq_lens=trace_prefill_supported_seq_lens, + trace_prefill_warmup_seq_lens=_trace_warmup_seq_lens( + 2048, + max_seq_len, + trace_prefill_supported_seq_lens, + ), + disable_batched_prefill=bool(os.getenv("DISABLE_BATCHED_PREFILL")), + ) + del hf + return Llama33_70BForCausalLM( + model=Llama33_70BTransformer1D(model_config), + tokenizer=tokenizer, + runtime_config=runtime_config, + instruct=instruct, + generation_config=Llama33_70BGenerationConfig(stop_token_ids=stop_token_ids), + ) diff --git a/code/models/common/models/llama33_70b/model.py b/code/models/common/models/llama33_70b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..d3407d05cdc44a2a8f1b43b14ef671be68c65832 --- /dev/null +++ b/code/models/common/models/llama33_70b/model.py @@ -0,0 +1,1312 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2 Llama 3.3-70B-Instruct — native stack (no ``models/tt_transformers`` imports). + +Architecture: standard Llama 1D transformer, same topology as Llama 3.1-8B / 3.2-3B +(no QKV bias, no Q/K norm, GPT-NeoX rotate_half RoPE with llama3 scaling). + hidden=8192, layers=80, n_heads=64, n_kv_heads=8, head_dim=128, + intermediate=28672, vocab=128256, rope_theta=500000, RoPE llama3-scaled (factor=8). + +Mesh compatibility: Wormhole T3K (1×8) and Blackhole P150x4 (1×4), backed by +either physical P150_X4 or P300_X2 hardware. The architecture/SKU profile validates +the exact product, device count, logical mesh shape, Ring topology, and P150 DRAM +width before composing modules. + +TTTv1 source for precision recipes: + ``models/tt_transformers/tt/model_config.py :: DecodersPrecision``. + ``get_base_model_name("…/Llama-3.3-70B-Instruct") == "Llama-3.3-70B"`` is NOT in the + ``Llama-3.1-70B`` special-case list (model_config.py:119), so it resolves to the generic + Llama-3 branch — identical recipe to Llama-3.2-3B: ``accuracy()`` BFP8 attention/KV/MLP + + HIFI2_FP16 FF / HIFI4 SDPA-prefill (lines 130-159); ``performance()`` BFP4 FF1/FF3 + LOFI + (lines 208-218). +""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n, _find_prefill_grid +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D +from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig +from models.common.modules.tt_ccl import get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + +# ============================================================================= +# Helpers +# ============================================================================= + + +LLAMA33_70B_BH_TP4_CLUSTER_TYPES = ( + ttnn.cluster.ClusterType.P150_X4, + ttnn.cluster.ClusterType.P300_X2, +) + + +def _llama33_70b_ccl_topology(mesh_device) -> ttnn.Topology: + """Return Ring only for an exact admitted physical/logical product pairing.""" + + arch = mesh_device.arch() + cluster_type = ttnn.cluster.get_cluster_type() + num_devices = mesh_device.get_num_devices() + mesh_shape = tuple(mesh_device.shape) + if ( + arch == ttnn.device.Arch.WORMHOLE_B0 + and cluster_type == ttnn.cluster.ClusterType.T3K + and num_devices == 8 + and mesh_shape == (1, 8) + ) or ( + arch == ttnn.device.Arch.BLACKHOLE + and cluster_type in LLAMA33_70B_BH_TP4_CLUSTER_TYPES + and num_devices == 4 + and mesh_shape == (1, 4) + ): + return ttnn.Topology.Ring + raise ValueError( + "Llama-3.3-70B CCL requires physical Wormhole T3K/8-device/(1, 8) or " + "BlackHole P150_X4/P300_X2/4-device/(1, 4) Ring geometry; " + f"got arch={arch}, cluster_type={cluster_type}, num_devices={num_devices}, mesh_shape={mesh_shape}" + ) + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +# ============================================================================= +# TransformerBlock1D — single decoder layer +# ============================================================================= + + +@dataclass +class TransformerBlock1DConfig: + attention_norm_config: RMSNorm1DConfig + attention_config: Attention1DConfig + ff_norm_config: RMSNorm1DConfig + mlp_config: MLP1DConfig + + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtype: ttnn.DataType | None = None + + +class TransformerBlock1D(LightweightModule): + def __init__( + self, + attention_norm: RMSNorm1D, + attention: Attention1D, + ff_norm: RMSNorm1D, + feed_forward: MLP1D, + decode_residual_memcfg: ttnn.MemoryConfig | None = None, + prefill_residual_memcfg: ttnn.MemoryConfig | None = None, + activation_dtype: ttnn.DataType | None = None, + ): + super().__init__() + self.attention_norm = attention_norm + self.attention = attention + self.ff_norm = ff_norm + self.feed_forward = feed_forward + self.decode_residual_memcfg = decode_residual_memcfg + self.prefill_residual_memcfg = prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtype = activation_dtype + + @classmethod + def from_config(cls, config: TransformerBlock1DConfig) -> TransformerBlock1D: + return cls( + attention_norm=RMSNorm1D.from_config(config.attention_norm_config), + attention=Attention1D.from_config(config.attention_config), + ff_norm=RMSNorm1D.from_config(config.ff_norm_config), + feed_forward=MLP1D.from_config(config.mlp_config), + decode_residual_memcfg=config.decode_residual_memcfg, + prefill_residual_memcfg=config.prefill_residual_memcfg, + activation_dtype=config.activation_dtype, + ) + + def decode_forward(self, x: ttnn.Tensor, current_pos, rot_mats, page_table) -> ttnn.Tensor: + residual = x + + x = _all_gather_rmsnorm_tensor( + self.attention_norm, x, memory_config=self.attention_norm.config.decode_memory_config + ) + attn_in = self.attention_norm.decode_forward(x) + attn_out = self.attention.decode_forward(attn_in, current_pos, rot_mats, page_table=page_table) + attn_out = ttnn.to_memory_config(attn_out, self.decode_residual_memcfg) + + hidden_states = ttnn.add(residual, attn_out, memory_config=self.decode_residual_memcfg) + residual = hidden_states + + hidden_states = _all_gather_rmsnorm_tensor( + self.ff_norm, hidden_states, memory_config=self.ff_norm.config.decode_memory_config + ) + hidden_states = self.ff_norm.decode_forward(hidden_states) + ttnn.deallocate(attn_out) + hidden_states = self.feed_forward.decode_forward(hidden_states) + + out = ttnn.add( + residual, + hidden_states, + memory_config=self.decode_residual_memcfg, + dtype=self.activation_dtype or ttnn.bfloat16, + ) + return out + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats, + user_id, + page_table, + chunk_page_table, + chunk_start_idx, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + residual = x + + attn_in = self.attention_norm.prefill_forward(x) + attn_in = _all_gather_rmsnorm_tensor(self.attention_norm, attn_in) + attn_out = self.attention.prefill_forward( + attn_in, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + attn_out = ttnn.to_memory_config(attn_out, self.prefill_residual_memcfg) + + hidden_states = ttnn.add(residual, attn_out, memory_config=self.prefill_residual_memcfg) + residual = hidden_states + x.deallocate(True) + + hidden_states = self.ff_norm.prefill_forward(hidden_states) + hidden_states = _all_gather_rmsnorm_tensor(self.ff_norm, hidden_states) + ttnn.deallocate(attn_out) + hidden_states = self.feed_forward.prefill_forward(hidden_states) + + out = ttnn.add( + residual, + hidden_states, + memory_config=self.prefill_residual_memcfg, + dtype=self.activation_dtype or ttnn.bfloat16, + ) + return out + + def forward( + self, + x, + current_pos=None, + rot_mats=None, + user_id=0, + mode="decode", + page_table=None, + chunk_page_table=None, + chunk_start_idx=None, + batch_size: int = 1, + ): + if mode == "prefill": + return self.prefill_forward( + x, rot_mats, user_id, page_table, chunk_page_table, chunk_start_idx, batch_size=batch_size + ) + return self.decode_forward(x, current_pos, rot_mats, page_table) + + +# ============================================================================= +# RMSNorm gather helper +# ============================================================================= + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, x: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + if memory_config is not None: + return ttnn.to_memory_config(x, memory_config) + return x + + if memory_config is None: + memory_config = x.memory_config() + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + # CCL pipelining matched to the proven mistral_7b / qwen25_7b N300 recipe + # (num_links=1, chunks_per_sync=24, num_workers_per_link=4) for parity with the + # reference ports. Measured neutral vs the prior (num_links=2, chunks_per_sync=10, + # num_workers_per_link=2) config on N300 batch-1 decode (perf-tuning.md §Axis-3: the + # per-layer gather is invisible in the decode-step budget on N150/N300; it only bites + # at T3K+ scale), but kept to avoid oversubscribing N300's single inter-chip link. + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=_llama33_70b_ccl_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + + +# ============================================================================= +# PrecisionConfig — TTTv1-matched recipes for Llama 3.3 70B Instruct +# ============================================================================= + +_LOFI_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, +) + +_HIFI2_FP16_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, +) + + +@dataclass(frozen=True) +class Llama33_70BPrecisionConfig: + """Per-layer precision + math-fidelity recipe for Llama 3.3 70B Instruct. + + Two module-level recipes: :data:`LLAMA33_70B_ACCURACY` and :data:`LLAMA33_70B_PERFORMANCE`. + The Hugging Face adaptor selects one while building this provider-neutral graph. + + Attention's six operation slots are materialized by the resolved model profile: + WH preserves the accepted TTTv2 baseline while BH uses the TTTv1 candidate recipe. + """ + + wqkv_dtype: ttnn.DataType = ttnn.bfloat8_b + wo_dtype: ttnn.DataType = ttnn.bfloat8_b + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat8_b + mlp_w2_dtype: ttnn.DataType = ttnn.bfloat8_b + mlp_ff1_3_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig = _HIFI2_FP16_COMPUTE_KERNEL_CFG + mlp_ff2_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig = _HIFI2_FP16_COMPUTE_KERNEL_CFG + + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + + +# TTTv1 DecodersPrecision.accuracy("Llama-3.3-70B-Instruct") (model_config.py Llama-3 group): +# wqkv=BFP8, wo=BFP8, kv_cache=BFP8, mlp_w1_w3=BFP8, mlp_w2=BFP8, +# LI_FF1_FF3=HIFI2_FP16, LI_FF2=HIFI2_FP16, SDPA_prefill=HIFI4 (Attention1D default) +LLAMA33_70B_ACCURACY = Llama33_70BPrecisionConfig() + +# TTTv1 DecodersPrecision.performance("Llama-3.3-70B-Instruct"): +# FF1_FF3 → BFP4, LI_FF1_FF3 → LOFI; all other fields same as accuracy. +LLAMA33_70B_PERFORMANCE = Llama33_70BPrecisionConfig( + mlp_w1_w3_dtype=ttnn.bfloat4_b, + mlp_ff1_3_compute_kernel_cfg=_LOFI_COMPUTE_KERNEL_CFG, + mlp_ff2_compute_kernel_cfg=_HIFI2_FP16_COMPUTE_KERNEL_CFG, +) + + +# ============================================================================= +# Runtime configs +# ============================================================================= + + +@dataclass +class Llama33_70BPagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (duck-typed; matches Attention1D's expected interface).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Llama33_70BTransformer1DConfig: + """Complete provider-neutral tensor-graph configuration for Llama 3.3 70B.""" + + n_layers: int + vocab_size: int + max_batch_size: int + max_seq_len: int + dim: int + num_devices: int + mesh_device: ttnn.MeshDevice + embedding_config: Embedding1DConfig + rope_config: Rope1DConfig + block_configs: list[TransformerBlock1DConfig] + norm_config: RMSNorm1DConfig + lm_head_config: LMHead1DConfig + sampling_config: Sampling1DConfig | None = None + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtypes: list[ttnn.DataType | None] = field(default_factory=list) + tt_ccl: Any = None + cache_path: str | None = None + batched_prefill_selector_compute_kernel_config: ttnn.DeviceComputeKernelConfig | None = None + + +@dataclass(frozen=True) +class Llama33_70BModelParameters: + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + max_batch_size: int + max_seq_len: int + + +@dataclass(frozen=True) +class Llama33_70BLayerWeights: + wqkv: torch.Tensor + wo: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Llama33_70BWeights: + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Llama33_70BLayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +# ============================================================================= +# Model precision profile + architecture/SKU overlay +# ============================================================================= + + +@dataclass(frozen=True, slots=True) +class _Llama33_70BModelProfile: + li_qkv_decode: ttnn.DeviceComputeKernelConfig + sdpa_decode: ttnn.DeviceComputeKernelConfig + li_o_decode: ttnn.DeviceComputeKernelConfig + li_qkv_prefill: ttnn.DeviceComputeKernelConfig + sdpa_prefill: ttnn.DeviceComputeKernelConfig + li_o_prefill: ttnn.DeviceComputeKernelConfig + prefill_ff1_ff3: ttnn.DeviceComputeKernelConfig + prefill_ff2: ttnn.DeviceComputeKernelConfig + decode_ff1_ff3: ttnn.DeviceComputeKernelConfig + decode_ff2: ttnn.DeviceComputeKernelConfig + rmsnorm: ttnn.DeviceComputeKernelConfig + lm_head: ttnn.DeviceComputeKernelConfig + + +@dataclass(frozen=True, slots=True) +class _Llama33_70BSKUOverlay: + mlp_prefill_len_cutoff: int + dram_shard_grid_width: int + prefill_qkv_grid: tuple[int, int] + decode_create_qkv_head_grid: ttnn.CoreGrid | None + decode_transformation_core_grid: ttnn.CoreCoord + lm_head_max_columns_per_device: int + prefill_minimal_matmul: bool + + +@dataclass(frozen=True, slots=True) +class _Llama33_70BComposition: + model: _Llama33_70BModelProfile + sku: _Llama33_70BSKUOverlay + + +def _kernel_config( + arch, + fidelity, + *, + approx: bool, + fp32: bool, + packer: bool, +) -> ttnn.DeviceComputeKernelConfig: + return ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=fidelity, + math_approx_mode=approx, + fp32_dest_acc_en=fp32, + packer_l1_acc=packer, + ) + + +def _copy_profile_kernel(arch, candidate) -> ttnn.DeviceComputeKernelConfig: + if candidate is None: + return _kernel_config(arch, ttnn.MathFidelity.HiFi2, approx=False, fp32=False, packer=True) + return _kernel_config( + arch, + candidate.math_fidelity, + approx=candidate.math_approx_mode, + fp32=candidate.fp32_dest_acc_en, + packer=candidate.packer_l1_acc, + ) + + +def _resolve_llama33_70b_profile( + *, arch, cluster_type, num_devices: int, dram_width: int, precision: Llama33_70BPrecisionConfig +) -> _Llama33_70BComposition: + """Compose the model recipe with the WH baseline or P150x4 SKU overlay.""" + is_wh = arch == ttnn.device.Arch.WORMHOLE_B0 + if not is_wh and arch != ttnn.device.Arch.BLACKHOLE: + raise ValueError(f"Unsupported Llama-3.3-70B architecture: {arch}") + expected_devices = 8 if is_wh else 4 + expected_clusters = (ttnn.cluster.ClusterType.T3K,) if is_wh else LLAMA33_70B_BH_TP4_CLUSTER_TYPES + if cluster_type not in expected_clusters: + raise ValueError( + f"Llama-3.3-70B requires physical cluster in {expected_clusters}, got {cluster_type}; " + "a logical submesh is not SKU-equivalent" + ) + if num_devices != expected_devices: + sku = "T3K" if is_wh else "P150x4" + raise ValueError(f"Llama-3.3-70B {sku} profile requires {expected_devices} devices, got {num_devices}") + if not is_wh and dram_width != 8: + raise ValueError(f"Llama-3.3-70B Blackhole profile requires P150 DRAM width 8, got {dram_width}") + + # WH locks the pre-change TTTv2 baseline. BH adopts the TTTv1 candidate + # recipe: five HiFi2/FP32/approx slots and exact HiFi4 SDPA prefill. + ordinary_attention = _kernel_config(arch, ttnn.MathFidelity.HiFi2, approx=not is_wh, fp32=not is_wh, packer=True) + sdpa_prefill = _kernel_config(arch, ttnn.MathFidelity.HiFi4, approx=False, fp32=True, packer=True) + ff1_ff3 = _copy_profile_kernel(arch, precision.mlp_ff1_3_compute_kernel_cfg) + ff2 = _copy_profile_kernel(arch, precision.mlp_ff2_compute_kernel_cfg) + model = _Llama33_70BModelProfile( + li_qkv_decode=_copy_profile_kernel(arch, ordinary_attention), + sdpa_decode=_copy_profile_kernel(arch, ordinary_attention), + li_o_decode=_copy_profile_kernel(arch, ordinary_attention), + li_qkv_prefill=_copy_profile_kernel(arch, ordinary_attention), + sdpa_prefill=_copy_profile_kernel(arch, sdpa_prefill), + li_o_prefill=_copy_profile_kernel(arch, ordinary_attention), + prefill_ff1_ff3=_copy_profile_kernel(arch, ff1_ff3), + prefill_ff2=_copy_profile_kernel(arch, ff2), + decode_ff1_ff3=_copy_profile_kernel(arch, ff1_ff3), + decode_ff2=_copy_profile_kernel(arch, ff2), + rmsnorm=_kernel_config(arch, ttnn.MathFidelity.HiFi2, approx=False, fp32=True, packer=True), + lm_head=_kernel_config(arch, ttnn.MathFidelity.HiFi2, approx=False, fp32=False, packer=True), + ) + sku = _Llama33_70BSKUOverlay( + mlp_prefill_len_cutoff=1024 if is_wh else 512, + dram_shard_grid_width=8, + prefill_qkv_grid=(8, 8) if is_wh else (8, 10), + decode_create_qkv_head_grid=None if is_wh else ttnn.CoreGrid(y=4, x=8), + decode_transformation_core_grid=ttnn.CoreCoord(8, 8), + lm_head_max_columns_per_device=8192 if is_wh else 128256 // 4 // 8, + prefill_minimal_matmul=not os.environ.get("DISABLE_MINIMAL_MATMUL"), + ) + return _Llama33_70BComposition(model=model, sku=sku) + + +# ============================================================================= +# Layer + head builders +# ============================================================================= + + +def _post_attn_norm_decode_configs( + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding to match MLP1D W1/W3 input.""" + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + memory_config = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return program_config, memory_config + + +def _build_decoder_layer( + *, + idx: int, + weights: Llama33_70BLayerWeights, + mcfg: Llama33_70BModelParameters, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Llama33_70BPrecisionConfig, + paged_attention_config: Llama33_70BPagedAttentionConfig, + cache_path: Path | None, + profile: _Llama33_70BComposition, + decode_residual_memcfg: ttnn.MemoryConfig, +) -> TransformerBlock1DConfig: + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, + dtype=precision.wqkv_dtype, + cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None, + ) + lazy_wo = _lazy( + weights.wo, + dtype=precision.wo_dtype, + cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None, + ) + + attention_config = Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=mcfg.n_heads, + n_kv_heads=mcfg.n_kv_heads, + head_dim=mcfg.head_dim, + max_batch_size=mcfg.max_batch_size, + max_seq_len=mcfg.max_seq_len, + use_vllm_paged_kv_cache=True, + paged_attention_config=paged_attention_config, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + # TTTv1 parity: Llama-3 family decode SDPA runs HIFI2 with exp_approx_mode=True + # (model_config.py `_default_settings` → SDPA_DECODE=HIFI2, used in BOTH accuracy + # and performance). Attention1D's generic default builds this prog config with + # exp_approx_mode=False, leaving decode SDPA slower than TTTv1. Flip it to match. + decode_sdpa_prg_config=ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=(8, 8), + exp_approx_mode=True, + q_chunk_size=0, + k_chunk_size=0, + ), + li_qkv_decode_compute_kernel_cfg=profile.model.li_qkv_decode, + sdpa_decode_compute_kernel_cfg=profile.model.sdpa_decode, + li_o_decode_compute_kernel_cfg=profile.model.li_o_decode, + li_qkv_prefill_compute_kernel_cfg=profile.model.li_qkv_prefill, + sdpa_prefill_compute_kernel_cfg=profile.model.sdpa_prefill, + li_o_prefill_compute_kernel_cfg=profile.model.li_o_prefill, + prefill_qkv_grid=profile.sku.prefill_qkv_grid, + dram_shard_grid_width=profile.sku.dram_shard_grid_width, + decode_create_qkv_head_grid=profile.sku.decode_create_qkv_head_grid, + decode_transformation_core_grid=profile.sku.decode_transformation_core_grid, + prefill_qkv_minimal_matmul=profile.sku.prefill_minimal_matmul, + ) + + padded_hidden_dim = get_padded_hidden_dim(mcfg.hidden_dim, num_dev, TILE_SIZE) + mlp_config = MLP1DConfig( + w1=_lazy( + weights.w1, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None, + ), + w2=_lazy( + weights.w2, + dtype=precision.mlp_w2_dtype, + cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None, + ), + w3=_lazy( + weights.w3, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None, + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=mcfg.max_batch_size, + decode_spill_w1_to_dram_before_w3=False, + w1_w3_dtype=precision.mlp_w1_w3_dtype, + w2_dtype=precision.mlp_w2_dtype, + ff1_3_compute_kernel_cfg=profile.model.prefill_ff1_ff3, + ff2_compute_kernel_cfg=profile.model.prefill_ff2, + decode_ff1_3_compute_kernel_cfg=profile.model.decode_ff1_ff3, + decode_ff2_compute_kernel_cfg=profile.model.decode_ff2, + prefill_len_cutoff=profile.sku.mlp_prefill_len_cutoff, + prefill_dram_shard_grid_width=profile.sku.dram_shard_grid_width, + prefill_ff1_ff3_grid=_find_prefill_grid(8, mcfg.dim // TILE_SIZE), + prefill_ff2_grid=_find_prefill_grid(8, padded_hidden_dim // TILE_SIZE), + prefill_w2_minimal_matmul=profile.sku.prefill_minimal_matmul, + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + dim=mcfg.dim, + hidden_dim=mcfg.hidden_dim, + num_devices=num_dev, + max_batch_size=mcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1DConfig: + lw = _lazy( + weight, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=mcfg.rms_norm_eps, + max_batch_size=mcfg.max_batch_size, + tt_ccl=tt_ccl, + prefill_distributed=num_dev > 1 and mcfg.dim > 4096, + compute_kernel_config=profile.model.rmsnorm, + **extra, + ) + + attn_norm = _build_norm(weights.attention_norm, "pre_attn") + ff_norm = _build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ) + + return TransformerBlock1DConfig( + attention_norm_config=attn_norm, + attention_config=attention_config, + ff_norm_config=ff_norm, + mlp_config=mlp_config, + decode_residual_memcfg=decode_residual_memcfg, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtype=None, + ) + + +def _build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """Build provider-neutral column-split LM-head weights.""" + + num_devices = mesh_device.get_num_devices() + if tuple(lm_head_weight.shape) != (vocab_size, dim): + raise ValueError( + f"Llama 70B LM-head weight must have shape {(vocab_size, dim)}, got {tuple(lm_head_weight.shape)}" + ) + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / (TILE_SIZE * num_devices)) * (TILE_SIZE * num_devices) + if vocab_size < padded_vocab_size: + torch_w = torch.cat( + [torch_w, torch.zeros(torch_w.shape[0], padded_vocab_size - vocab_size, dtype=torch_w.dtype)], dim=-1 + ) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + + output_weights = [] + weights_memcfgs = [] + for split_index, split_size in enumerate(split_sizes): + device_splits = [] + physical_split_size = math.ceil(split_size / TILE_SIZE) * TILE_SIZE + for device_index in range(num_devices): + start = device_index * size_per_device + sum(split_sizes[:split_index]) + device_split = torch_w[:, start : start + split_size] + if split_size < physical_split_size: + device_split = torch.cat( + [device_split, torch.zeros(dim, physical_split_size - split_size, dtype=device_split.dtype)], + dim=-1, + ) + device_splits.append(device_split) + combined = torch.cat(device_splits, dim=-1) + padded_n = math.ceil((combined.shape[-1] // num_devices) / (TILE_SIZE * dram_size.x)) * ( + TILE_SIZE * dram_size.x + ) + shard_spec = ttnn.ShardSpec( + dram_grid, + (dim, padded_n // dram_size.x), + ttnn.ShardOrientation.ROW_MAJOR, + ) + memory_config = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + weights_memcfgs.append(memory_config) + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=memory_config, + cache_dir_weight_name=( + ( + cache_dir, + f"lm_head_split_{split_index}_logical_{split_size}_physical_{combined.shape[-1]}", + ) + if cache_dir + else None + ), + ) + ) + return output_weights, split_sizes, weights_memcfgs + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + mcfg: Llama33_70BModelParameters, + lm_head_dtype: ttnn.DataType, + cache_path: Path | None, + profile: _Llama33_70BComposition, +) -> LMHead1DConfig: + lm_splits, lm_split_sizes, lm_weights_memcfgs = _build_lm_head_lazy_weights( + mesh_device, + lm_head_weight, + dim=mcfg.dim, + vocab_size=mcfg.vocab_size, + max_columns_per_device=profile.sku.lm_head_max_columns_per_device, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(mcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(mcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, mcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(mcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=mcfg.dim, + max_batch_size=mcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + output_split_sizes=lm_split_sizes, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + compute_kernel_config=profile.model.lm_head, + ) + + +# ============================================================================= +# Llama33_70BTransformer1D +# ============================================================================= + + +def build_llama33_70b_transformer_1d_config( + *, + mesh_device: ttnn.MeshDevice, + params: Llama33_70BModelParameters, + weights: Llama33_70BWeights, + n_layers: int, + precision: Llama33_70BPrecisionConfig, + cache_path: Path, + paged_attention_config: Llama33_70BPagedAttentionConfig, +) -> Llama33_70BTransformer1DConfig: + """Build the TT tensor graph from provider-neutral dimensions and tensors.""" + + num_devices = mesh_device.get_num_devices() + arch = mesh_device.arch() + profile = _resolve_llama33_70b_profile( + arch=arch, + cluster_type=ttnn.cluster.get_cluster_type(), + num_devices=num_devices, + dram_width=mesh_device.dram_grid_size().x, + precision=precision, + ) + if params.n_heads % num_devices or params.n_kv_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({params.n_heads}/{params.n_kv_heads}) must be divisible by device count ({num_devices})" + ) + if len(weights.layers) != n_layers: + raise ValueError(f"Expected {n_layers} decoder layer weight sets, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) + topology = _llama33_70b_ccl_topology(mesh_device) + embedding_config = Embedding1DConfig( + weights=_lazy( + weights.embedding, + dtype=ttnn.bfloat16, + cache=(cache_path / "embedding", "tok_embeddings"), + ), + mesh_device=mesh_device, + embed_scale=1.0, + ) + rope_config = Rope1DConfig( + cos_matrix=_lazy(weights.rope_cos, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "cos")), + sin_matrix=_lazy(weights.rope_sin, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "sin")), + max_batch_size=params.max_batch_size, + head_dim=params.head_dim, + device=mesh_device, + use_qk_fused=False, + # Keep decode cos/sin rows on the same 8-wide batch-core mapping as + # create_qkv_heads and Attention's rotary transformation matrix. A + # physical P150 grid is 12-wide; allowing Rope1D to infer that grid + # remaps slot 24 to a different core row than Attention. + core_grid=profile.sku.decode_transformation_core_grid, + ) + block_configs = [ + _build_decoder_layer( + idx=index, + weights=weights.layers[index], + mcfg=params, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + paged_attention_config=paged_attention_config, + cache_path=cache_path, + profile=profile, + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + ) + for index in range(n_layers) + ] + norm_config = RMSNorm1DConfig( + weight=_lazy(weights.final_norm, dtype=ttnn.bfloat16, cache=(cache_path / "norm", "final")), + mesh_device=mesh_device, + eps=params.rms_norm_eps, + max_batch_size=params.max_batch_size, + tt_ccl=tt_ccl, + prefill_distributed=num_devices > 1 and params.dim > 4096, + compute_kernel_config=profile.model.rmsnorm, + ) + lm_head_config = _build_lm_head( + mesh_device=mesh_device, + lm_head_weight=weights.lm_head, + mcfg=params, + lm_head_dtype=precision.lm_head_dtype, + cache_path=cache_path, + profile=profile, + ) + sampling_config = Sampling1DConfig( + vocab_size=params.vocab_size, + valid_vocab_size=params.vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=_nearest_32(params.max_batch_size), + allow_force_argmax=False, + pad_to_power_of_2=True, + ) + return Llama33_70BTransformer1DConfig( + n_layers=n_layers, + vocab_size=params.vocab_size, + max_batch_size=params.max_batch_size, + max_seq_len=params.max_seq_len, + dim=params.dim, + num_devices=num_devices, + mesh_device=mesh_device, + embedding_config=embedding_config, + rope_config=rope_config, + block_configs=block_configs, + norm_config=norm_config, + lm_head_config=lm_head_config, + sampling_config=sampling_config, + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtypes=[None] * n_layers, + tt_ccl=tt_ccl, + cache_path=str(cache_path), + batched_prefill_selector_compute_kernel_config=_kernel_config( + arch, + ttnn.MathFidelity.HiFi4, + approx=False, + fp32=True, + packer=False, + ), + ) + + +class Llama33_70BTransformer1D(LightweightModule): + """Provider-neutral TTTv2 Llama 3.3 70B tensor model.""" + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__(self, config: Llama33_70BTransformer1DConfig): + from tqdm import tqdm + + super().__init__() + self.config = config + self.embedding = Embedding1D.from_config(config.embedding_config) + self.rope_setup = RotarySetup1D.from_config(config.rope_config) + self.layers = [ + TransformerBlock1D.from_config(config.block_configs[index]) + for index in tqdm(range(config.n_layers), desc="Building layers") + ] + self.norm = RMSNorm1D.from_config(config.norm_config) + self.lm_head = LMHead1D.from_config(config.lm_head_config) + self.sampling = Sampling1D.from_config(config.sampling_config) if config.sampling_config is not None else None + self.supports_on_device_sampling = self.sampling is not None + self.mesh_device = config.mesh_device + self.tt_ccl = config.tt_ccl + self.vocab_size = config.vocab_size + self.n_layers = config.n_layers + self.num_devices = config.num_devices + self.decode_residual_memcfg = config.decode_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.prefill_residual_memcfg = config.prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtypes = config.activation_dtypes or [None] * config.n_layers + self.model_args = None + self.batched_prefill_selector_compute_kernel_config = config.batched_prefill_selector_compute_kernel_config + + # ========================================================================= + # KV cache binding + # ========================================================================= + + def iter_executor_named_modules(self): + layers = getattr(self, "layers", ()) + for i, layer in enumerate(layers): + for suffix, submodule in ( + ("attn_norm", layer.attention_norm), + ("attention", layer.attention), + ("ff_norm", layer.ff_norm), + ("mlp", layer.feed_forward), + ): + yield f"layer[{i}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in self.config.block_configs) + for config in tuple({id(c): c for c in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, block_size=block_size, max_num_blocks=max_num_blocks + ) + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for i, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {i} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {i} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = pair + + # ========================================================================= + # Forward methods — take pre-embedded tensors + # ========================================================================= + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for i, layer in enumerate(self.layers): + x = ttnn.to_memory_config(x, self.decode_residual_memcfg, self.activation_dtypes[i]) + x = layer.decode_forward(x, current_pos, rot_mats, page_table) + + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + x = self.lm_head.forward(x) + return x + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). The batched path always + # returns the full hidden state (get_last_token == -1); the executor does per-slot last-token + # extraction + norm/lm_head so those stages stay bit-identical to the single-user path. + x = x_embed + for i, layer in enumerate(self.layers): + activation_dtype = self.activation_dtypes[i] + if activation_dtype is not None and x.dtype != activation_dtype: + old = x + x = ttnn.typecast(x, activation_dtype) + ttnn.deallocate(old) + x = layer.prefill_forward( + x, + rot_mats, + user_id, + page_table, + chunk_page_table, + chunk_start_idx, + batch_size, + chunk_start_idx_tensor, + ) + + if last_token_index is not None and last_token_slice is None: + raise ValueError("last_token_index is required with a runtime last_token_slice") + if get_last_token == -1 and last_token_slice is None: + return x + + old = x + if last_token_slice is None: + get_last_token_floor = (get_last_token // 32) * 32 + x = ttnn.slice(x, (0, 0, get_last_token_floor, 0), (1, 1, get_last_token_floor + 32, x.shape[-1])) + else: + x = ttnn.slice( + x, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(x.shape[2]) // 32, + ) + ttnn.deallocate(old) + + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + get_last_token_floor = (last_token_idx // 32) * 32 + x = ttnn.slice( + hidden_states, + (0, 0, get_last_token_floor, 0), + (1, 1, get_last_token_floor + 32, hidden_states.shape[-1]), + ) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + del last_token_slice, last_token_index + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), + ) + x = ttnn.matmul( + selector, + hidden_states, + dtype=ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + compute_kernel_config=self.batched_prefill_selector_compute_kernel_config, + ) + ttnn.deallocate(selector) + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def forward( + self, + x: ttnn.Tensor, + current_pos=None, + rot_mats_global=None, + rot_mats_local=None, + user_id: int = 0, + mode: str = "decode", + page_table=None, + chunk_page_table=None, + chunk_start_idx=None, + get_last_token: int = -1, + batch_size: int = 1, + ) -> ttnn.Tensor: + rot_mats = rot_mats_global + if mode == "prefill": + return self.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + get_last_token=get_last_token, + batch_size=batch_size, + ) + return self.decode_forward(x, current_pos, rot_mats, page_table=page_table) + + # ========================================================================= + # Embedding + output processing helpers (executor contract) + # ========================================================================= + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embedding.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embedding.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=_llama33_70b_ccl_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) + + +# ============================================================================= +# Public exports +# ============================================================================= + +__all__ = [ + "Llama33_70BPrecisionConfig", + "LLAMA33_70B_ACCURACY", + "LLAMA33_70B_PERFORMANCE", + "Llama33_70BPagedAttentionConfig", + "Llama33_70BLayerWeights", + "Llama33_70BModelParameters", + "Llama33_70BWeights", + "Llama33_70BTransformer1DConfig", + "Llama33_70BTransformer1D", + "build_llama33_70b_transformer_1d_config", + "TransformerBlock1D", + "TransformerBlock1DConfig", + "_all_gather_rmsnorm_tensor", +] diff --git a/code/models/common/models/llama33_70b/weight_utils.py b/code/models/common/models/llama33_70b/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b06738e0d091e603edc969c75dc1673dce03e9dc --- /dev/null +++ b/code/models/common/models/llama33_70b/weight_utils.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side HF weight layout helpers for Llama-3.3-70B-Instruct (TTTv2 port). + +Q/K linear weights use HuggingFace head layout; TTNN RoPE expects Meta layout. +This duplicates the small permute helpers from ``load_checkpoints.reverse_permute`` +without importing ``models/tt_transformers``. + +Llama 3.3 70B has no QKV bias and no Q/K norm. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.modules.lazy_weight import LazyWeight + + +def reverse_permute(tensor: torch.Tensor, n_heads: int, dim1: int, dim2: int) -> torch.Tensor: + """HF Q/K weight rows → Meta layout (RoPE-compatible).""" + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def permute_hf_rope_to_meta_tables(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert HF cos/sin tables to Meta-style tables for ``RotarySetup1D``.""" + if len(cos.shape) == 3: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + return cos, sin + + +def build_rope_cos_sin_torch( + rotary_emb: Any, table_len: int, head_dim: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Build ``[1, 1, table_len, head_dim]`` cos/sin tensors (Meta layout) from HF rotary module.""" + x = torch.zeros(1, 1, table_len, head_dim, dtype=dtype) + position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0) + with torch.no_grad(): + cos_hf, sin_hf = rotary_emb(x, position_ids) + cos_m, sin_m = permute_hf_rope_to_meta_tables(cos_hf.float(), sin_hf.float()) + return cos_m.to(dtype), sin_m.to(dtype) + + +def attention_wqkv_wo_from_hf_layer( + self_attn: Any, + num_devices: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack QKV and WO from an HF Llama attention block for ``Attention1D`` sharding.""" + wq_raw = self_attn.q_proj.weight + wk_raw = self_attn.k_proj.weight + wv_raw = self_attn.v_proj.weight + wo_raw = self_attn.o_proj.weight + + dim = wq_raw.shape[1] + cfg = self_attn.config + n_heads = cfg.num_attention_heads + n_kv_heads = cfg.num_key_value_heads + head_dim = cfg.hidden_size // n_heads + n_heads_times_head_dim = n_heads * head_dim + n_kv_heads_times_head_dim = n_kv_heads * head_dim + + wq_meta = reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim) + wk_meta = reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim) + wv_meta = wv_raw + wo_meta = wo_raw + + wq = wq_meta.T + wk = wk_meta.T + wv = wv_meta.T + wo = wo_meta.T + + qkv_list = [] + for i in range(num_devices): + wq_chunk = torch.chunk(wq, num_devices, dim=1)[i] + wk_chunk = torch.chunk(wk, num_devices, dim=1)[i] + wv_chunk = torch.chunk(wv, num_devices, dim=1)[i] + qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1) + qkv_list.append(qkv) + wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0).clone() + wo = wo.unsqueeze(0).unsqueeze(0).clone() + return wqkv, wo + + +def mlp_weights_from_hf_layer(mlp: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (w1, w2, w3) TTNN layouts: gate^T, down^T, up^T for ``MLP1D``.""" + w1 = mlp.gate_proj.weight.T.contiguous().clone() + w3 = mlp.up_proj.weight.T.contiguous().clone() + w2 = mlp.down_proj.weight.T.contiguous().clone() + return w1, w2, w3 + + +def rms_weight_torch(layernorm: Any) -> torch.Tensor: + return layernorm.weight.detach().float().clone() + + +def embed_tokens_torch(embed: Any) -> torch.Tensor: + w = embed.weight.detach().float().clone() + return w.unsqueeze(0).unsqueeze(0) + + +def build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """Column-split LM head for ``LMHead1D`` (DRAM-sharded matmul chunks per device). + + ``lm_head_weight`` shape: ``[vocab_size, dim]`` (HF). Returns + ``(lazy_weights, split_sizes, weights_memcfgs)`` for ``LMHead1DConfig``. + """ + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / 32) * 32 + if vocab_size < padded_vocab_size: + pad = padded_vocab_size - vocab_size + torch_w = torch.cat([torch_w, torch.zeros(torch_w.shape[0], pad, dtype=torch_w.dtype)], dim=-1) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + tile = ttnn.TILE_SIZE + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (tile * dram_size.x)) * (tile * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights: list[LazyWeight] = [] + weights_memcfgs: list[ttnn.MemoryConfig] = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_splits.append(torch_w[:, start:end]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + name = f"lm_head_split_{i}_{combined.shape[-1]}" + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, name) if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs diff --git a/code/models/common/models/llama3_8b/README.md b/code/models/common/models/llama3_8b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..39cdea628ae88776a846c75f6413dea2bf91523d --- /dev/null +++ b/code/models/common/models/llama3_8b/README.md @@ -0,0 +1,543 @@ +# Llama 3.1 8B with TTTv2 + +This directory contains the model-owned Llama 3.1 8B product path built from +TTTv2 modules and the reusable common LLM runtime. + +The path has four layers: + +```text +model provider / checkpoint + -> hf_adaptor.py: provider metadata, tokenizer, and weight conversion + -> model.py: TTTv2 tensor model assembled from reusable modules + -> executor.py: thin typed entry point into the Llama family executor + -> generator.py: vLLM-facing construction, DP composition, and dispatch +``` + +The most important boundary is between the tensor model and runtime +orchestration: + +- TTTv2 `LightweightModule` objects implement tensor computation. +- [`models/common/llm_runtime`](../../llm_runtime/README.md) implements reusable + execution, tracing, I/O, cache, warmup, and resource mechanics. +- `models/common/models/executor.py::ModelExecutor` composes the common owners. +- `models/common/models/llama3_executor.py::Llama3Executor` supplies the + Llama-8B sampling and prefill policy as a composition facade. +- `Llama3Generator` adapts the resulting target to vLLM. + +## Files + +| File | Responsibility | +| --- | --- | +| `hf_adaptor.py` | Load HF config/tokenizer/weights, convert provider naming/layout, compute Llama 3 RoPE values, and create the product model | +| `model.py` | Build and execute the TTTv2 Llama transformer graph | +| `executor.py` | Preserve the model-local typed builder/import surface over `llama3_executor.py` | +| `generator.py` | Construct lanes, optionally compose DP, normalize vLLM calls, and select eager/traced execution | + +## End-to-end object graph + +For one lane: + +```text +Llama3ForCausalLM +├── tokenizer +├── Llama3RuntimeConfig +└── Llama3Transformer1D + ├── Embedding1D + ├── RotarySetup1D + ├── TransformerBlock1D × N + │ ├── RMSNorm1D + │ ├── Attention1D + │ ├── RMSNorm1D + │ └── MLP1D + ├── RMSNorm1D + ├── LMHead1D + └── optional Sampling1D + +Llama3Executor composition facade +└── ModelExecutor + ├── exact Llama3Transformer1D above + ├── PagedKVCacheManager + ├── OutputReader + ├── PrefillRuntime + ├── DecodeRuntime + ├── ProgramCompiler + ├── EagerExecutor + ├── optional TraceCompiler + ├── optional TracedExecutor over the exact EagerExecutor + └── WarmupCoordinator +``` + +The vLLM-facing graph is: + +```text +Llama3Generator +├── VLLMAdapter +└── target + ├── Llama3Executor when DP = 1 + └── LaneGroupExecutor[Llama3Executor, ...] when DP > 1 +``` + +`Llama3Generator` owns no TT tensors. The lane executors own resources, and a +`LaneGroupExecutor` owns lane/pool lifecycle coordination. + +## Building the tensor model + +### Provider adaptation + +`from_pretrained(...)` in `hf_adaptor.py` is the current Hugging Face provider +entry point. It: + +1. resolves the model ID; +2. loads `AutoConfig` and the tokenizer; +3. derives hidden size, heads, KV heads, layers, vocabulary, norm epsilon, and + context length; +4. computes Llama 3 scaled RoPE cosine/sine tables; +5. loads the HF state dict; +6. splits fused QKV or gate/up weights when necessary; +7. converts Q/K rotary weight layout; +8. maps HF names to the model's Meta-style names; +9. builds `Llama3Transformer1DConfig`; +10. constructs `Llama3Transformer1D`; and +11. returns `Llama3ForCausalLM`, which packages the tensor model, tokenizer, + generation defaults, and `Llama3RuntimeConfig`. + +Provider-facing concerns stop there. Neither `Llama3Executor` nor the common +runtime reads HF config or converts HF weights. + +### TTTv2 module composition + +`build_llama3_transformer_1d_config(...)` translates Llama architecture and +optimization choices into configs for reusable TTTv2 modules: + +- `Embedding1D` +- `RotarySetup1D` +- `RMSNorm1D` +- `Attention1D` +- `MLP1D` +- `LMHead1D` +- optional `Sampling1D` + +`Llama3Transformer1D` constructs these modules. Each +`TransformerBlock1D` performs: + +```text +attention RMSNorm + -> Attention1D + -> residual add + -> feed-forward RMSNorm + -> MLP1D + -> residual add +``` + +The model exposes two graph entry points: + +- `prefill_forward(...)` for one planned regular/batched/chunk invocation; and +- `decode_forward(...)` for one autoregressive step across the fixed lane + capacity. + +It also exposes executor support methods: + +- `iter_executor_named_modules()` yields modules whose input contracts must be + validated during execution; +- `set_kv_cache(cache_or_none)` transactionally binds/unbinds per-layer K/V + tensors; +- embedding and rotary preparation methods stage model inputs; +- prefill post-processing converts a traced hidden body to logits/sampled + output; and +- decode output gathering and position increment helpers support runtime + execution. + +## Constructing the vLLM model + +The public class entry point is: + +```text +Llama3Generator.initialize_vllm_model(...) + -> Llama3GeneratorConfig + -> build_llama3_generator(config) +``` + +`build_llama3_generator(...)` performs the following steps. + +### 1. Resolve lane geometry + +The global vLLM batch is divided evenly by `tt_data_parallel`. For DP1, the +whole mesh is one lane. For DP2/DP4/DP8, the mesh is split into one submesh per +lane. + +Each lane receives: + +- one submesh; +- one fixed per-lane batch capacity; +- the same maximum sequence length; +- the same optimization/precision policy; and +- the same trace and device-sampling policy. + +### 2. Build one product model per lane + +For each submesh: + +```text +from_pretrained(...) + -> Llama3ForCausalLM + -> Llama3Transformer1D on that submesh +``` + +The paged-attention block size is 32. `max_num_blocks` is a safe static +construction ceiling derived from maximum sequence length and per-lane batch +capacity. + +### 3. Build one model-owned executor per lane + +The generator creates `Llama3ExecutorConfig`: + +- `TraceConfig(trace_mode)` +- `WarmupConfig()` +- unresolved `PagedKVCacheConfig` +- device-sampling capability + +It then calls: + +```text +build_llama3_executor(Llama3ForCausalLM, executor_config) + -> llama3_executor.Llama3Executor facade + -> ModelExecutor(model, runtime_config, executor_config, Llama policy) +``` + +The family facade creates native `SamplingState1D` state and resolves the +Llama-8B device-sampling prefill policy. The shared `ModelExecutor` composes +the runtime owners and exposes three execution targets: + +- `eager_execution`: always the one `EagerExecutor`; +- `traced_prefill_execution`: the one `TracedExecutor` when prefill tracing is + configured; and +- `traced_decode_execution`: the same `TracedExecutor` when decode tracing is + configured. + +There is no aggregate executor in `llm_runtime`. The shared composition root +lives in the model layer at `models/common/models/executor.py`. + +### 4. Build the vLLM boundary adapter + +Model metadata is read from the already-built attention configs: + +- layer count; +- KV dtype per layer; +- local KV heads per device; and +- head dimension. + +That metadata resolves `VLLMAdapterConfig`. `VLLMAdapter` then owns only static +vLLM normalization/validation policy; it owns no TT resource. + +### 5. Compose the target + +For DP1, the target is the single `Llama3Executor`. + +For DP greater than one: + +```text +LaneGroupExecutor(lanes) + -> one duck-typed global execution target +``` + +The lane group: + +- assigns prefill rows to lanes from their global slots; +- maps global slots to lane-local slots; +- splits decode into contiguous per-lane batches; +- aggregates outputs in global order; +- replicates cache configuration, warmup, and compilation; and +- coordinates concurrent asynchronous output handling and cleanup. + +Finally: + +```text +Llama3Generator(target, vllm_adapter) +``` + +is returned to vLLM. + +## vLLM lifecycle + +### 1. Model construction uses only a maximum KV ceiling + +At construction, the generator does not know vLLM's final physical block +count. Each lane therefore has: + +```text +PagedKVCacheConfig( + block_size=32, + max_num_blocks=construction_ceiling, + num_blocks=None, +) +``` + +`Llama3Executor` can still construct prefill, decode, and warmup config against +the maximum. This is cheap TTTv2 reconfiguration: no physical KV tensor is +allocated at this point. + +### 2. vLLM resolves physical KV capacity + +vLLM calls: + +```text +Llama3Generator.allocate_kv_cache(kv_cache_shape, dtype, num_layers) +``` + +The call chain is: + +```text +VLLMAdapter.resolve_legacy_kv_cache_config(...) + -> validate physical blocks <= maximum + -> validate local KV heads, block size, head dimension, layer count, dtype + -> return new PagedKVCacheConfig(num_blocks=physical_blocks) + +target.configure_paged_kv_cache(resolved_config) + -> one executor or every DP lane + -> PagedKVCacheManager.configure(...) + -> recompute PageTableLayout for physical capacity + -> replace PrefillRuntimeConfig layout + -> replace DecodeRuntimeConfig layout + -> replace WarmupCoordinatorConfig layout and rebuild coverage plans + +target.allocate_kv_cache() + -> seal runtime geometry + -> allocate per-layer K/V tensors + -> bind tensors to Llama3Transformer1D +``` + +This ordering is important: the physical page-table layout is installed before +allocation, compilation, warmup, or trace capture. + +### 3. Warmup and trace capture + +vLLM calls `warmup_model_prefill(...)` and `warmup_model_decode(...)`. + +Each lane compiles all required program variants. Trace capture waits at the +shared warmup barrier until both configured operation sets are ready. Sampling +buffers are loaded before capture. + +For `trace_mode="all"`, prefill and decode traces are separate artifacts over +the same eager program compiler. This means vLLM may still request eager or +traced execution independently on every forward call. + +### 4. Prefill dispatch + +```text +vLLM + -> Llama3Generator.prefill_forward(...) + -> VLLMAdapter.normalize_prefill(...) + -> bind positional arguments + -> remove known irrelevant compatibility fields + -> require explicit Boolean enable_trace + -> normalize torch dtypes + -> Llama3Generator._select_prefill_execution(...) + -> if trace requested, target.can_trace_prefill(...) + -> cached/chunked/unsupported requests select eager + -> eligible requests select traced + -> target.prefill_forward(execution=selected, ...) + -> Llama3Executor, or LaneGroupExecutor -> each Llama3Executor + -> selected EagerExecutor or TracedExecutor + -> PrefillRuntime + -> Llama3Transformer1D +``` + +The fallback belongs here, at the vLLM/model boundary. `TracedExecutor` never +silently invokes eager execution. + +### 5. Decode dispatch + +```text +vLLM + -> Llama3Generator.decode_forward(...) + -> VLLMAdapter.normalize_decode(...) + -> explicit enable_trace selects: + false -> target.eager_execution + true -> target.traced_decode_execution + -> target.decode_forward(execution=selected, ...) + -> DecodeRuntime + -> Llama3Transformer1D +``` + +Decode trace availability is a static capability. Asking for traced decode +when it was not configured is an error at the vLLM boundary. + +### 6. Asynchronous decode output + +vLLM can request `read_from_device=False`. The executor returns a raw TT output +under an external lease. + +```text +Llama3Generator.read_decode_output(async_read=True) + -> lane target + -> DecodeRuntime.read_decode_output(...) + -> OutputReader.submit(...) + -> host destination + TT completion events + +Llama3Generator.process_decode_output_host(...) + -> DecodeRuntime.process_decode_output_host(...) + -> OutputReader.complete(...) + -> ttnn.event_synchronize(...) + -> normalize output and release the lease +``` + +For DP, the lane group performs the per-lane reads concurrently and aggregates +the completed outputs. + +### 7. Cleanup + +`Llama3Generator.cleanup()` delegates to the target. + +One `Llama3Executor` terminalizes and releases: + +1. externally leased decode outputs; +2. pending output reads; +3. prefill/decode transients; +4. trace resources; +5. program registry state; +6. sampling buffers; and +7. the bound paged KV cache. + +The DP target cleans every lane and then its worker pool. Construction failures +also clean all lanes that were already created. + +## Trace-mode behavior + +Every vLLM forward call carries an explicit `enable_trace` Boolean. + +| Static `trace_mode` | Operation | `enable_trace=False` | `enable_trace=True` | +| --- | --- | --- | --- | +| `none` | prefill or decode | eager | rejected by adapter | +| `decode_only` | prefill | eager | rejected by adapter | +| `decode_only` | decode | eager | traced | +| `all` | decode | eager | traced | +| `all` | eligible regular prefill | eager | traced | +| `all` | cached, chunked, or otherwise trace-ineligible prefill | eager | generator selects eager | + +`trace_mode="all"` is the most flexible serving construction because prefill +and decode artifacts are independent. It supports per-call eager/traced +selection without reconstructing the model. + +## Applying this pattern to another LLM + +The reusable pattern is not “subclass Llama3.” It is: + +```text +provider adapter + -> model-specific TTTv2 graph + -> shared/family model executor or direct runtime composition + -> server-specific facade +``` + +### Model implementation + +A new model should build its tensor graph from reusable TTTv2 modules where +possible. The exact module set may differ: another architecture might use a +different attention implementation, normalization, MLP, MoE, positional +encoding, or output head. + +The tensor model should expose the runtime contract needed by its executor: + +- prefill and decode graph entry points; +- model-owned embedding/input and output-processing helpers; +- module iteration for input-contract validation; +- transactional KV-cache binding; +- per-layer cache metadata; and +- optional device sampling. + +### Model execution composition + +Use the shared `models/common/models/executor.py::ModelExecutor` when the model +fits its established lifecycle. A demonstrated family may add a small policy +facade such as `llama3_executor.py` or `qwen2_executor.py`. + +When a model has genuinely distinct orchestration, its model-local +`executor.py` may instead compose the focused `llm_runtime` modules directly. +Either construction should: + +- translate model metadata into resolved common runtime configs; +- construct one exact eager execution composition; +- optionally construct one trace compiler and one traced executor over it; +- own page-layout sealing and late physical-capacity replacement; +- validate that request cache handles belong to its cache manager; +- expose the duck-typed execution target used by a DP group; and +- be the deterministic cleanup root. + +Do not add a generic aggregate model executor to `llm_runtime`, and do not +force every model through the shared model-layer executor. + +### Server facade + +Create a facade for the target serving system. It should own: + +- external argument normalization; +- external cache-shape adaptation; +- per-call eager/traced selection; +- request-level trace eligibility fallback; +- server-specific async-output conventions; and +- construction of single-lane or DP targets. + +The common prefill/decode/compiler/cache mechanics should not interpret the +server's policy. + +## Extensibility dimensions + +This architecture separates several dimensions that can evolve independently. + +### Other model architectures + +Llama, Mistral, Qwen, Gemma, MoE models, and future architectures can share the +runtime mechanics while owning different TTTv2 module graphs and executors. + +### Other inference servers + +vLLM is one facade. An SGLang integration can build the same model executor and +provide an SGLang-specific adapter for request fields, cache negotiation, +trace selection, and asynchronous output conventions. A direct demo or custom +service can bypass server adapters and call the model-owned executor with an +explicit execution target. + +### Other model providers + +Hugging Face is currently isolated in `hf_adaptor.py`. Another provider can +supply: + +- architecture metadata; +- tokenizer/chat formatting; +- a state-dict reader; +- provider-to-model key and tensor-layout conversion; and +- cache location policy. + +That provider adapter should produce the same model product shape: + +```text +TTTv2 tensor model + tokenizer + model runtime metadata +``` + +The Llama executor and common runtime do not need to know whether weights came +from Hugging Face, a native Meta checkpoint, an internal artifact store, or a +preconverted tensor cache. + +### Other topologies and execution policies + +Mesh topology, tensor parallelism inside modules, data-parallel lane count, +precision/optimization policy, paged-KV capacity, device sampling, warmup +coverage, and trace mode are separate configuration dimensions. A new +combination should normally require new resolved configs and validation, not a +fork of runtime control flow. + +## Practical checklist for a new integration + +1. Build and validate the provider adapter. +2. Construct the TTTv2 tensor model from module configs. +3. Expose model runtime and KV metadata. +4. Select shared model-layer composition, a justified family policy facade, or + direct composition from the common runtime. +5. Test direct eager prefill/decode and cleanup. +6. Add program compilation and warmup coverage. +7. Add trace capture/replay without eager fallback inside `TracedExecutor`. +8. Add late physical KV-capacity resolution. +9. Add a server facade that owns normalization and dispatch. +10. Add DP composition through `LaneGroupExecutor` if required. +11. Validate accuracy, deterministic text quality, sustained TPOT, aggregate + throughput, and cleanup across all supported geometries. diff --git a/code/models/common/models/llama3_8b/executor.py b/code/models/common/models/llama3_8b/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..298b82dcffa3b7c1f9e9ca28922c71b75f1c1f3b --- /dev/null +++ b/code/models/common/models/llama3_8b/executor.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Llama 3.1-8B executor construction entry point.""" + +from models.common.models.llama3_8b.hf_adaptor import Llama3ForCausalLM +from models.common.models.llama3_executor import Llama3Executor, Llama3ExecutorConfig + + +def build_llama3_executor(llm: Llama3ForCausalLM, config: Llama3ExecutorConfig) -> Llama3Executor: + """Build one executor around an already-loaded Llama 3.1-8B adapter.""" + + return Llama3Executor(llm.model, llm.runtime_config, config) diff --git a/code/models/common/models/llama3_8b/generator.py b/code/models/common/models/llama3_8b/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..45260304cee8d5040eba3a7dc9e832ed48b203f5 --- /dev/null +++ b/code/models/common/models/llama3_8b/generator.py @@ -0,0 +1,479 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Llama 3.1-8B.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.llama3_8b.executor import Llama3ExecutorConfig, build_llama3_executor +from models.common.models.llama3_8b.hf_adaptor import from_pretrained +from models.common.models.llama3_8b.model import Llama31_8BPagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Llama3GeneratorConfig: + """Validated construction inputs for one vLLM-facing Llama generator.""" + + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Llama3Generator: + """Adapt vLLM's model interface to the model-owned execution target. + + vLLM constructs this facade with `initialize_vllm_model`, resolves + KV capacity through `allocate_kv_cache`, warms the configured + programs, and then calls `prefill_forward` and + `decode_forward`. Each forward call is normalized by + `VLLMAdapter`, dispatched to the eager or traced executor, and + delegated to ``Llama3Executor`` or ``LaneGroupExecutor``. + + This class owns dispatch policy but no TT resources. `cleanup` + delegates to the target that owns those resources. + """ + + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "max_device_top_k": 32, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + + # Public vLLM API + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + """Return the unpadded per-submesh KV token budget for vLLM sizing.""" + + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + device_sampling_enabled: bool = True, + ): + """Build the configured single-lane or data-parallel Llama target.""" + + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_llama3_generator( + Llama3GeneratorConfig( + hf_model=str(hf_model), + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def compile_prefill( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + ) -> None: + """Normalize a vLLM prefill call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.compile_prefill(execution=execution, **normalized) + + def compile_decode( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + ) -> None: + """Normalize a vLLM decode call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.compile_decode(execution=execution, **normalized) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM prefill call.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM decode call.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward( + execution=execution, + read_from_device=read_from_device, + **normalized, + ) + + def read_decode_output( + self, + tt_out: Any, + *, + async_read: bool = False, + ) -> Any: + """Delegate vLLM's raw decode-output read.""" + + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host( + self, + tt_out: Any, + *, + is_tokens: bool = False, + ) -> tuple[Any, Any]: + """Delegate vLLM's asynchronous host-output completion.""" + + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self): + """Release every resource owned by the concrete target.""" + + return self.target.cleanup() + + # Private implementation + + def _select_prefill_execution( + self, + normalized: NormalizedPrefillKwargs, + trace_requested: bool, + ): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_llama3_generator(config: Llama3GeneratorConfig) -> Llama3Generator: + """Construct lane-local models/executors and compose their shared target surface.""" + + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Llama31_8BPagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + instruct="Instruct" in config.hf_model, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Llama3ExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(include_decode_top_k=config.device_sampling_enabled), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_llama3_executor(llm, executor_config)) + + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Llama3Generator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + request_state_fields=lane._request_state_fields, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Llama model must contain at least one attention layer") + + attention_configs = tuple(layer.attention.config for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Llama layer must expose the same KV head shape") + + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/llama3_8b/hf_adaptor.py b/code/models/common/models/llama3_8b/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..21e07203e8bcec0810a06614cb8e19d98e0a659f --- /dev/null +++ b/code/models/common/models/llama3_8b/hf_adaptor.py @@ -0,0 +1,554 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face adaptor for the TTTv2 Llama-3.1-8B path.""" + +from __future__ import annotations + +import errno +import math +import os +import re +from dataclasses import dataclass, field +from pathlib import Path + +import torch +from loguru import logger +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.device_utils import get_device_name +from models.common.tensor_utils import nearest_multiple + + +@dataclass(frozen=True) +class RopeScaling: + rope_type: str + factor: float + original_max_position_embeddings: int + low_freq_factor: float + high_freq_factor: float + + +def llama3_rope_scaling(rope_parameters: dict) -> RopeScaling: + rope_type = rope_parameters["rope_type"] + if rope_type != "llama3": + raise ValueError(f"Unsupported RoPE scaling type for Llama-3.1-8B TTTv2 path: {rope_type}") + + return RopeScaling( + rope_type=rope_type, + factor=rope_parameters["factor"], + original_max_position_embeddings=rope_parameters["original_max_position_embeddings"], + low_freq_factor=rope_parameters["low_freq_factor"], + high_freq_factor=rope_parameters["high_freq_factor"], + ) + + +def _permute_to_meta_format(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + + return cos.unsqueeze(0).unsqueeze(0), sin.unsqueeze(0).unsqueeze(0) + + +def _gather_cos_sin(position_ids: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + position_id_expanded = position_ids.unsqueeze(1).expand(-1, cos.shape[-1]) + cos = cos.gather(0, position_id_expanded) + sin = sin.gather(0, position_id_expanded) + cos = torch.stack([cos, cos], dim=-1).flatten(-2).unsqueeze(0).unsqueeze(0) + sin = torch.stack([sin, sin], dim=-1).flatten(-2).unsqueeze(0).unsqueeze(0) + return cos, sin + + +def _llama3_scaled_inv_freq(freqs: torch.Tensor, scaling: RopeScaling) -> torch.Tensor: + low_freq_wavelen = scaling.original_max_position_embeddings / scaling.low_freq_factor + high_freq_wavelen = scaling.original_max_position_embeddings / scaling.high_freq_factor + new_freqs = [] + for freq in freqs: + wavelen = 2 * math.pi / freq + if wavelen < high_freq_wavelen: + new_freqs.append(freq) + elif wavelen > low_freq_wavelen: + new_freqs.append(freq / scaling.factor) + else: + smooth = (scaling.original_max_position_embeddings / wavelen - scaling.low_freq_factor) / ( + scaling.high_freq_factor - scaling.low_freq_factor + ) + new_freqs.append((1 - smooth) * freq / scaling.factor + smooth * freq) + return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device) + + +def compute_gather_cos_sin( + dhead: int, end: int, theta: float, rope_scaling: RopeScaling +) -> tuple[torch.Tensor, torch.Tensor]: + seq_len = end // 2 + inv_freq = 1.0 / (theta ** (torch.arange(0, dhead, 2).float() / dhead)) + + if rope_scaling.rope_type != "llama3": + raise ValueError(f"Unsupported RoPE scaling type for Llama-3.1-8B TTTv2 path: {rope_scaling.rope_type}") + inv_freq = _llama3_scaled_inv_freq(inv_freq, rope_scaling) + + t = torch.arange(seq_len * 2.0) + freqs = torch.outer(t, inv_freq).float() + cos, sin = torch.cos(freqs), torch.sin(freqs) + return _gather_cos_sin(torch.arange(seq_len), cos, sin) + + +def should_pad_sampling_logits_to_power_of_2(padded_vocab_size: int, sampling_splits: int) -> bool: + if sampling_splits < 1: + return False + per_device_vocab = padded_vocab_size // sampling_splits + return per_device_vocab > 0 and (per_device_vocab & (per_device_vocab - 1)) != 0 + + +def resolve_hf_model_id(hf_model: str | None = None) -> str: + hf_model = hf_model or os.getenv("HF_MODEL") + if not hf_model: + raise ValueError("Please set HF_MODEL to a HuggingFace name e.g. meta-llama/Llama-3.1-8B-Instruct") + return hf_model + + +def _replace_keys(state_dict, replacements): + output = {} + for key, value in state_dict.items(): + new_key = key + for pattern, repl in replacements: + new_key = re.sub(pattern, repl, new_key) + output[new_key] = value + return output + + +def _standardize_hf_keys(state_dict): + key_meta = "lm_head.weight" + key_hf = "model.embed_tokens.weight" + if key_meta not in state_dict and key_hf in state_dict: + state_dict[key_meta] = state_dict[key_hf] + del state_dict[key_hf] + return state_dict + + +def _split_hf_keys(loaded_weights, n_heads=None, n_kv_heads=None): + converted_weights = {} + for key, tensor in loaded_weights.items(): + if "qkv_proj" in key: + q_key = key.replace("qkv_proj", "q_proj") + k_key = key.replace("qkv_proj", "k_proj") + v_key = key.replace("qkv_proj", "v_proj") + if n_heads is not None and n_kv_heads is not None and n_heads != n_kv_heads: + head_dim = tensor.shape[0] // (n_heads + 2 * n_kv_heads) + q_size = n_heads * head_dim + kv_size = n_kv_heads * head_dim + q_tensor = tensor[:q_size] + k_tensor = tensor[q_size : q_size + kv_size] + v_tensor = tensor[q_size + kv_size : q_size + 2 * kv_size] + else: + q_tensor, k_tensor, v_tensor = torch.split(tensor, tensor.shape[0] // 3, dim=0) + converted_weights[q_key] = q_tensor + converted_weights[k_key] = k_tensor + converted_weights[v_key] = v_tensor + elif "gate_up_proj" in key: + gate_key = key.replace("gate_up_proj", "gate_proj") + up_key = key.replace("gate_up_proj", "up_proj") + gate_tensor, up_tensor = torch.split(tensor, tensor.shape[0] // 2, dim=0) + converted_weights[gate_key] = gate_tensor + converted_weights[up_key] = up_tensor + else: + converted_weights[key] = tensor + return converted_weights + + +def _reverse_permute(tensor, n_heads, dim1, dim2): + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def _reverse_permute_1d(tensor): + dim = tensor.shape[-1] + assert dim % 2 == 0, "Last dimension must be even" + reals = tensor[..., : dim // 2] + imags = tensor[..., dim // 2 :] + return torch.stack((reals, imags), dim=-1).flatten(start_dim=len(tensor.shape) - 1) + + +def _convert_hf_qkv_to_meta_format(loaded_weights, head_dim): + converted_weights = {} + for key, tensor in loaded_weights.items(): + if "q_proj.weight" in key or "k_proj.weight" in key: + n_heads = tensor.shape[0] // head_dim + converted_weights[key] = _reverse_permute(tensor, n_heads, tensor.shape[0], tensor.shape[1]) + elif "q_proj.bias" in key or "k_proj.bias" in key: + n_heads = tensor.shape[0] // head_dim + converted_weights[key] = _reverse_permute(tensor, n_heads, tensor.shape[0], 1).squeeze(-1) + elif "q_norm.weight" in key or "k_norm.weight" in key: + converted_weights[key] = _reverse_permute_1d(tensor) + else: + converted_weights[key] = tensor + return converted_weights + + +def _map_hf_to_meta_keys(loaded_weights): + replacements = [ + ("^emb.weight", "weight"), + ("model.", ""), + ("embed_tokens", "tok_embeddings"), + ("lm_head", "output"), + ("input_layernorm", "attention_norm"), + ("post_attention_layernorm", "ffn_norm"), + ("self_attn", "attention"), + ("mlp", "feed_forward"), + ("gate_proj", "w1"), + ("down_proj", "w2"), + ("up_proj", "w3"), + ("q_proj", "wq"), + ("k_proj", "wk"), + ("v_proj", "wv"), + ("o_proj", "wo"), + ("q_norm", "q_norm"), + ("k_norm", "k_norm"), + ] + return _replace_keys(loaded_weights, replacements) + + +def convert_hf_state_dict_to_meta(state_dict, *, head_dim: int, n_heads: int, n_kv_heads: int): + state_dict = _split_hf_keys(state_dict, n_heads, n_kv_heads) + state_dict = _convert_hf_qkv_to_meta_format(state_dict, head_dim) + return _map_hf_to_meta_keys(state_dict) + + +def load_tokenizer(hf_model: str, *, trust_remote_code: bool = False): + tokenizer = AutoTokenizer.from_pretrained( + hf_model, + local_files_only=os.getenv("CI") == "true", + trust_remote_code=trust_remote_code, + ) + if not hasattr(tokenizer, "stop_tokens") or tokenizer.stop_tokens is None: + tokenizer.stop_tokens = [tokenizer.eos_token_id] + return tokenizer + + +@dataclass(frozen=True) +class Llama3GenerationConfig: + """Text-generation defaults for the Llama 3.1-8B product model.""" + + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Llama3RuntimeConfig: + """Executor/runtime metadata kept outside the tensor graph config.""" + + model_name: str + model_cache_path: Path + max_prefill_chunk_size: int + max_context_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] = (128, 1024) + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len, num_cached_tokens=0): + return ( + num_cached_tokens == 0 + and prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def _encode_prompt_with_chat_template(tokenizer, prompt_text, system_prompt_text=None): + chat = [] + if isinstance(prompt_text, str): + if system_prompt_text: + chat.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + chat.append({"role": "user", "content": prompt_text}) + encoded = tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True) + else: + encoded = tokenizer.apply_chat_template(prompt_text, add_generation_prompt=True, tokenize=True) + return _chat_template_ids(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + try: + return _encode_prompt_with_chat_template(tokenizer, prompt_text, system_prompt_text) + except ValueError as exc: + logger.warning(f"Failed to encode chat prompt, falling back to base encoding: {exc}") + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Llama3ForCausalLM: + """Usable Llama 3.1-8B model product: tokenizer plus TT tensor model. + + The tokenizer interface is intentionally documented rather than enforced + through a Protocol for now. The object must provide encode/decode behavior, + EOS/stop token IDs, and chat-template application for instruct models. + """ + + model: object + tokenizer: object + runtime_config: Llama3RuntimeConfig + instruct: bool + generation_config: Llama3GenerationConfig = field(default_factory=Llama3GenerationConfig) + + def __post_init__(self): + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + stop_tokens = tuple(getattr(self.tokenizer, "stop_tokens", []) or []) + self.generation_config = Llama3GenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=stop_tokens, + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + use_instruct = self.instruct if instruct is None else instruct + return encode_prompt(self.tokenizer, prompt_text, system_prompt_text, instruct=use_instruct) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_converted_state_dict( + hf_model: str, + *, + head_dim: int, + n_heads: int, + n_kv_heads: int, + n_layers: int, + trust_remote_code: bool = False, +): + model = AutoModelForCausalLM.from_pretrained( + hf_model, + torch_dtype="auto", + trust_remote_code=trust_remote_code, + local_files_only=os.getenv("CI") == "true", + ) + state_dict = model.state_dict() + state_dict = _standardize_hf_keys(state_dict) + state_dict = convert_hf_state_dict_to_meta( + state_dict, + head_dim=head_dim, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + ) + for key in list(state_dict.keys()): + if "layers." in key: + layer_num = int(key.split("layers.")[1].split(".")[0]) + if layer_num >= n_layers: + state_dict.pop(key) + return state_dict + + +def _model_cache_path(hf_model: str, mesh_device) -> Path: + cache_path = os.getenv("TT_CACHE_PATH") + device_name = get_device_name(mesh_device) + if not cache_path: + return Path("model_cache") / hf_model / device_name + + configured_path = Path(cache_path) / device_name + try: + configured_path.mkdir(parents=True, exist_ok=True) + return configured_path + except OSError as exc: + if exc.errno not in (errno.EROFS, errno.EACCES, errno.EPERM): + raise + + fallback_root = Path(os.getenv("TT_CACHE_FALLBACK_PATH", "/tmp/tttv2_model_cache")) + fallback_path = fallback_root / Path(hf_model).name / device_name + fallback_path.mkdir(parents=True, exist_ok=True) + logger.warning( + f"Configured TT cache is not writable at {configured_path}; " f"using job-local tensor cache {fallback_path}" + ) + return fallback_path + + +def _max_prefill_chunk_size(mesh_device) -> int: + override = os.getenv("MAX_PREFILL_CHUNK_SIZE") + if override is not None: + return int(override) * 1024 + return { + "N150": 4, + "N300": 64, + "N150x4": 4, + "T3K": 128, + "P150": 4, + "P300": 4, + "P150x4": 128, + }[get_device_name(mesh_device)] * 1024 + + +def _trace_prefill_supported_seq_lens( + device_name: str, max_prefill_chunk_size: int, max_seq_len: int +) -> tuple[int, ...]: + supported_seq_lens_by_device = { + "N150": (128, 1024), + "P150": (128, 1024), + "P300": (128, 1024), + "P150x4": (128, 1024), + "N300": (128, 1024, 2048, 4096, 8192), + "N150x4": (128, 1024, 2048, 4096, 8192), + "T3K": (128, 1024, 2048, 4096, 8192), + } + supported_seq_lens = supported_seq_lens_by_device[device_name] + return tuple(seq_len for seq_len in supported_seq_lens if seq_len <= min(max_prefill_chunk_size, max_seq_len)) + + +def _disable_batched_prefill(mesh_device) -> bool: + """Resolve the model/SKU half of the sequential-prefill policy.""" + + return get_device_name(mesh_device) in {"P150", "P300", "P150x4", "P150x8"} or bool( + os.getenv("DISABLE_BATCHED_PREFILL") + ) + + +def _weight_cache_path(model_cache_path: Path, *, instruct: bool, dtype): + if instruct: + return ( + model_cache_path + / { + ttnn.bfloat16: "tensor_cache_instruct_bf16", + ttnn.bfloat8_b: "tensor_cache_instruct_bfp8", + }[dtype] + ) + return model_cache_path / {ttnn.bfloat16: "tensor_cache_bf16", ttnn.bfloat8_b: "tensor_cache_bfp8"}[dtype] + + +def from_pretrained( + mesh_device, + *, + hf_model: str | None = None, + instruct: bool | None = None, + max_batch_size: int, + max_seq_len: int, + optimizations="performance", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config=None, + converted_state_dict: dict[str, torch.Tensor] | None = None, +): + """Build a product-level TTTv2 Llama-3.1-8B model from an HF checkpoint.""" + from models.common.models.llama3_8b.model import Llama3Transformer1D, build_llama3_transformer_1d_config + + hf_model = resolve_hf_model_id(hf_model) + if instruct is None: + instruct = "Instruct" in Path(hf_model).name + + hf_config = AutoConfig.from_pretrained( + hf_model, + local_files_only=os.getenv("CI") == "true", + ) + text_config = hf_config.to_dict() + model_name = Path(hf_model).name + tokenizer = load_tokenizer(hf_model) + num_hidden_layers = n_layers if n_layers is not None else text_config["num_hidden_layers"] + + rope_cos, rope_sin = compute_gather_cos_sin( + dhead=text_config["hidden_size"] // text_config["num_attention_heads"], + end=2 * max_seq_len, + theta=text_config["rope_parameters"]["rope_theta"], + rope_scaling=llama3_rope_scaling(text_config["rope_parameters"]), + ) + model_cache_path = _model_cache_path(hf_model, mesh_device) + + model_config = build_llama3_transformer_1d_config( + mesh_device=mesh_device, + instruct=instruct, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + model_name=model_name, + dim=text_config["hidden_size"], + n_heads=text_config["num_attention_heads"], + n_kv_heads=text_config["num_key_value_heads"], + n_layers=num_hidden_layers, + head_dim=text_config["hidden_size"] // text_config["num_attention_heads"], + hidden_dim=text_config["intermediate_size"], + vocab_size=text_config["vocab_size"], + norm_eps=text_config["rms_norm_eps"], + padded_vocab_size=nearest_multiple(text_config["vocab_size"], ttnn.TILE_SIZE * mesh_device.get_num_devices()), + rope_cos=rope_cos, + rope_sin=rope_sin, + model_cache_path=model_cache_path, + state_dict=( + converted_state_dict + if converted_state_dict is not None + else load_converted_state_dict( + hf_model, + head_dim=text_config["hidden_size"] // text_config["num_attention_heads"], + n_heads=text_config["num_attention_heads"], + n_kv_heads=text_config["num_key_value_heads"], + n_layers=num_hidden_layers, + ) + ), + optimizations=optimizations, + weight_cache_path=_weight_cache_path(model_cache_path, instruct=instruct, dtype=dtype), + dtype=dtype, + paged_attention_config=paged_attention_config, + pad_logits_to_power_of_2=list(mesh_device.shape) != [1, 1] + and should_pad_sampling_logits_to_power_of_2( + nearest_multiple(text_config["vocab_size"], ttnn.TILE_SIZE * mesh_device.get_num_devices()), + mesh_device.get_num_devices() if list(mesh_device.shape) != [1, 1] else 2, + ), + ) + max_prefill_chunk_size = _max_prefill_chunk_size(mesh_device) + trace_prefill_supported_seq_lens = _trace_prefill_supported_seq_lens( + get_device_name(mesh_device), + max_prefill_chunk_size, + max_seq_len, + ) + runtime_config = Llama3RuntimeConfig( + model_name=model_name, + model_cache_path=model_cache_path, + max_prefill_chunk_size=max_prefill_chunk_size, + max_context_len=text_config["max_position_embeddings"], + trace_prefill_supported_seq_lens=trace_prefill_supported_seq_lens, + # TTTv1 disables batched prefill for Llama-3.1-8B on every supported + # BlackHole SKU because BH prefill reductions are batch-variant. The + # executor independently disables it for device sampling so serving + # also has finite program/trace coverage on every architecture. + disable_batched_prefill=_disable_batched_prefill(mesh_device), + batched_prefill_batched_extract=not os.environ.get("DISABLE_BATCHED_EXTRACT"), + ) + return Llama3ForCausalLM( + model=Llama3Transformer1D(model_config), + tokenizer=tokenizer, + runtime_config=runtime_config, + instruct=instruct, + ) diff --git a/code/models/common/models/llama3_8b/model.py b/code/models/common/models/llama3_8b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..126cd9c95d565ae51c31b5b6d1058cfe131d9e5a --- /dev/null +++ b/code/models/common/models/llama3_8b/model.py @@ -0,0 +1,1902 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2 Llama 3.1-8B Transformer model. + +Model: + Llama3Transformer1D — pure forward methods, no input/output processing + +Executor wrappers live in models/common/models/llama3_8b/executor.py. + +Architecture: + Llama3Transformer1D (1D only — non-TG) + ├── Embedding1D + ├── RotarySetup1D + ├── TransformerBlock1D × n_layers + │ ├── RMSNorm1D (attention_norm) + │ ├── Attention1D + │ ├── RMSNorm1D (ff_norm) + │ └── MLP1D + ├── RMSNorm1D (final norm) + ├── LMHead1D + └── Sampling1D (optional) + +Loop policy functions (run_teacher_forcing, run_perf_benchmark) are in +models/common/models/executor.py. +""" + +import math +import os +from dataclasses import dataclass, field, replace +from pathlib import Path + +import torch + +import ttnn +from models.common.device_utils import get_device_name +from models.common.lightweightmodule import LightweightModule +from models.common.modules.attention.attention_1d import Attention1D, Attention1DConfig +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight as CommonLazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _compute_kernel_config_hifi2 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _create_dram_sharded_mem_config +from models.common.modules.rmsnorm.rmsnorm_1d import SHARD_HEIGHT, RMSNorm1D, RMSNorm1DConfig +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D +from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig +from models.common.modules.tt_ccl import TT_CCL, default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_out_subblock_w, nearest_32, num_to_core_range_set, pad_dim_to_size + + +class LazyWeight(CommonLazyWeight): + """Let equivalent single-device Llama lanes share portable cache files. + The common cache fingerprint includes the concrete mesh-device id. That is + useful for device-bound layouts, but Llama DP lanes serialize host tensors + beneath an already product-qualified ``P150`` cache directory. Reusing an + otherwise identical single-device cache file avoids rebuilding the whole + model once per physical DP lane while retaining the legacy exact path for + writes and every multi-device lookup. + """ + + def _get_cache_fill_path(self, cache_dir, weight_name): + exact_path = super()._get_cache_fill_path(cache_dir, weight_name) + if exact_path is None or exact_path.exists() or self.device is None: + return exact_path + if not hasattr(self.device, "get_num_devices") or self.device.get_num_devices() != 1: + return exact_path + if not hasattr(self.device, "id"): + return exact_path + + device_token = f"device_{self.device.id()}" + if device_token not in exact_path.name: + return exact_path + portable_pattern = exact_path.name.replace(device_token, "device_*", 1) + return next( + (candidate for candidate in sorted(exact_path.parent.glob(portable_pattern)) if candidate.is_file()), + exact_path, + ) + + +# ============================================================================= +# Runtime Config + + +class Llama31DecoderPrecision: + """Per-decoder tensor dtype and math-fidelity selection.""" + + _DTYPES = { + "bfp4": ttnn.bfloat4_b, + "bfp8": ttnn.bfloat8_b, + "bf16": ttnn.bfloat16, + None: None, + } + + @classmethod + def from_string(cls, optimizations: str): + if optimizations == "performance": + return cls.performance + if optimizations == "accuracy": + return cls.accuracy + raise ValueError( + f"Invalid optimization configuration: {optimizations}. Allowed values are 'performance' or 'accuracy'" + ) + + @classmethod + def performance(cls, num_decoders: int, model_name: str): + inst = cls(num_decoders, model_name, cls._performance_settings(model_name)) + if model_name == "Llama-3.1-8B-Instruct" and num_decoders > 31: + inst._tensor_precision[31]["ff1_ff3"] = "bfp8" + inst._op_fidelity[31]["li_ff1_ff3"] = "hifi2fp16" + inst._update_full_name() + inst.__name__ = "performance" + return inst + + @classmethod + def accuracy(cls, num_decoders: int, model_name: str): + inst = cls(num_decoders, model_name, cls._accuracy_settings(model_name)) + inst.__name__ = "accuracy" + return inst + + def __init__(self, num_decoders: int, model_name: str, settings: dict | None = None): + self.model_name = model_name + default_tensor_precision, default_op_fidelity = self._default_settings() + settings = settings or {} + default_tensor_precision.update(settings.get("tensor_precision", {})) + default_op_fidelity.update(settings.get("op_fidelity", {})) + self._tensor_precision = {decoder_id: dict(default_tensor_precision) for decoder_id in range(num_decoders)} + self._op_fidelity = {decoder_id: dict(default_op_fidelity) for decoder_id in range(num_decoders)} + self._update_full_name() + + @staticmethod + def _base_model_name(model_name: str): + for suffix in ("-Instruct", "-instruct"): + if model_name.endswith(suffix): + return model_name[: -len(suffix)] + return model_name + + @classmethod + def _accuracy_settings(cls, model_name: str): + base_model_name = cls._base_model_name(model_name) + if base_model_name.startswith("Llama-3") or base_model_name.startswith("Meta-Llama-3"): + return { + "tensor_precision": { + "wqkv": "bfp8", + "kv_cache": "bfp8", + "wo": "bfp8", + }, + "op_fidelity": { + "li_ff1_ff3": "hifi2fp16", + "li_ff2": "hifi2fp16", + }, + } + return { + "tensor_precision": { + "wqkv": "bf16", + "kv_cache": "bf16", + "wo": "bf16", + }, + "op_fidelity": { + "li_qkv_decode": "hifi4", + "li_qkv_prefill": "hifi4", + "sdpa_decode": "hifi4", + "sdpa_prefill": "hifi4", + "li_o_decode": "hifi4", + "li_o_prefill": "hifi4", + }, + } + + @classmethod + def _performance_settings(cls, model_name: str): + return { + "tensor_precision": {"ff1_ff3": "bfp4"}, + "op_fidelity": {"li_ff1_ff3": "lofi"}, + } + + @staticmethod + def _default_settings(): + return ( + { + "ff1_ff3": "bfp8", + "ff2": "bfp8", + "wqkv": "bfp8", + "wo": "bfp8", + "kv_cache": "bfp8", + "activation": None, + }, + { + "li_ff1_ff3": "hifi2fp16", + "li_ff2": "hifi2fp16", + "li_qkv_decode": "hifi2", + "sdpa_decode": "hifi2", + "li_o_decode": "hifi2", + "li_qkv_prefill": "hifi2", + "sdpa_prefill": "hifi4", + "li_o_prefill": "hifi2", + "accuracy": "hifi4fp32", + }, + ) + + def get_tensor_dtype(self, decoder_id: int, tensor: str, prefetcher: bool = False): + effective_decoder_id = 0 if prefetcher else decoder_id + value = self._tensor_precision.get(effective_decoder_id, {}).get(tensor) + if prefetcher and value is None and tensor != "activation": + return ttnn.bfloat8_b + return self._DTYPES.get(value) + + def get_math_fidelity(self, decoder_id: int, op: str, configuration): + kernel_lookup = { + "lofi": configuration.compute_kernel_config_lofi, + "hifi2": configuration.compute_kernel_config_hifi2, + "hifi2na": configuration.compute_kernel_config_hifi2_na, + "hifi2fp16": configuration.compute_kernel_config_hifi2_fp16, + "hifi2nol1acc": configuration.compute_kernel_config_hifi2_nol1acc, + "hifi4": configuration.compute_kernel_config_hifi4, + "hifi4fp32": configuration.compute_kernel_config_hifi4_fp32, + } + return kernel_lookup[self._op_fidelity[decoder_id][op]] + + def _update_full_name(self): + self._full_name = " | ".join( + f"Decoder {decoder_id}: precision_cfg = {self._tensor_precision[decoder_id]}, fidelity_cfg = {self._op_fidelity[decoder_id]}" + for decoder_id in self._tensor_precision + ) + + +def _base_model_name(model_name: str) -> str: + for suffix in ("-Instruct", "-instruct"): + if model_name.endswith(suffix): + return model_name[: -len(suffix)] + return model_name + + +@dataclass(frozen=True, slots=True) +class _Llama31_8BArchitectureProfile: + """Model/SKU-owned policy layered on top of shared WH/BH legality.""" + + rms_packer_l1_acc: bool + rms_distributed_at_dim_4096: bool + mlp_prefill_len_cutoff: int + mlp_prefill_dram_shard_grid_width: int + mlp_prefill_ff1_ff3_grid: tuple[int, int] + mlp_prefill_ff2_grid: tuple[int, int] + attention_prefill_qkv_grid: tuple[int, int] + attention_decode_create_qkv_head_grid: ttnn.CoreGrid | None + attention_decode_transformation_core_grid: ttnn.CoreCoord | None + enable_minimal_qkv: bool + enable_minimal_ff2: bool + lm_head_max_columns_per_device: int | None + + +def _resolve_llama31_8b_architecture_profile( + *, arch, cluster_type, device_name: str, model_name: str, dram_grid_width: int +) -> _Llama31_8BArchitectureProfile: + """Return the approved model/SKU overlay without querying global architecture state.""" + if arch == ttnn.device.Arch.WORMHOLE_B0: + return _Llama31_8BArchitectureProfile( + rms_packer_l1_acc=False, + rms_distributed_at_dim_4096=True, + mlp_prefill_len_cutoff=( + 512 if device_name == "N150" and _base_model_name(model_name) == "Llama-3.1-8B" else 1024 + ), + mlp_prefill_dram_shard_grid_width=8, + mlp_prefill_ff1_ff3_grid=(8, 8), + mlp_prefill_ff2_grid=(8, 8), + attention_prefill_qkv_grid=(8, 8), + attention_decode_create_qkv_head_grid=None, + attention_decode_transformation_core_grid=None, + enable_minimal_qkv=False, + enable_minimal_ff2=False, + lm_head_max_columns_per_device=None, + ) + if arch == ttnn.device.Arch.BLACKHOLE: + return _Llama31_8BArchitectureProfile( + rms_packer_l1_acc=True, + # The embedding shards the 4096-wide hidden dimension across a + # multi-device mesh, so prefill RMSNorm must all-gather statistics + # for the local slices before the model gathers normalized hidden + # slices. + rms_distributed_at_dim_4096=True, + mlp_prefill_len_cutoff=512, + mlp_prefill_dram_shard_grid_width=dram_grid_width, + mlp_prefill_ff1_ff3_grid=(8, 8), + mlp_prefill_ff2_grid=(8, 8), + attention_prefill_qkv_grid=(8, 10), + attention_decode_create_qkv_head_grid=ttnn.CoreGrid(y=4, x=8), + attention_decode_transformation_core_grid=ttnn.CoreCoord(8, 8), + enable_minimal_qkv=True, + enable_minimal_ff2=True, + lm_head_max_columns_per_device={ + "P100": 16032, + "P150": 16032, + "P300": 16032, + "P150x4": 4008, + "P150x8": 1002, + }.get(device_name), + ) + raise ValueError(f"Unsupported Llama 3.1 8B architecture: {arch}") + + +def _use_distributed_prefill_rmsnorm( + *, num_devices: int, dim: int, architecture_profile: _Llama31_8BArchitectureProfile +) -> bool: + """Resolve the effective model/SKU prefill RMSNorm policy.""" + threshold = 4096 if architecture_profile.rms_distributed_at_dim_4096 else 4097 + return num_devices > 1 and dim >= threshold + + +def _make_llama31_8b_rope_config( + *, + rope_cos, + rope_sin, + max_batch_size: int, + head_dim: int, + mesh_device, + decode_transformation_core_grid, +) -> Rope1DConfig: + """Build RoPE setup on the same decode grid used by attention. + + Fused Q/K decode places the batch-32 Q and K tensors on an 8x8 core + region. Blackhole's physical compute grid is wider, so allowing RoPE to + derive its batch grid from the device would distribute its 64 shards over + a different set of cores. Keep the setup and consuming attention + program on one model-profile-owned grid, matching TTTv1's Blackhole + RotarySetup policy. + """ + return Rope1DConfig( + cos_matrix=LazyWeight(source=rope_cos, device=mesh_device), + sin_matrix=LazyWeight(source=rope_sin, device=mesh_device), + max_batch_size=max_batch_size, + head_dim=head_dim, + device=mesh_device, + use_qk_fused=True, + core_grid=decode_transformation_core_grid, + ) + + +# ============================================================================= +# TransformerBlock1D +# ============================================================================= + + +@dataclass +class TransformerBlock1DConfig: + attention_norm_config: RMSNorm1DConfig + attention_config: Attention1DConfig + ff_norm_config: RMSNorm1DConfig + mlp_config: MLP1DConfig + + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtype: ttnn.DataType | None = None + + +class TransformerBlock1D(LightweightModule): + """Single transformer block for 1D topologies (N150, N300, T3K). + + Happy path (takes pre-built sub-modules): + block = TransformerBlock1D(attn_norm, attention, ff_norm, mlp) + + Power-user path (builds from config): + block = TransformerBlock1D.from_config(config) + """ + + def __init__( + self, + attention_norm: RMSNorm1D, + attention: Attention1D, + ff_norm: RMSNorm1D, + feed_forward: MLP1D, + decode_residual_memcfg: ttnn.MemoryConfig | None = None, + prefill_residual_memcfg: ttnn.MemoryConfig | None = None, + activation_dtype: ttnn.DataType | None = None, + ): + super().__init__() + self.attention_norm = attention_norm + self.attention = attention + self.ff_norm = ff_norm + self.feed_forward = feed_forward + self.decode_residual_memcfg = decode_residual_memcfg + self.prefill_residual_memcfg = prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtype = activation_dtype + + @classmethod + def from_config(cls, config: TransformerBlock1DConfig): + return cls( + attention_norm=RMSNorm1D.from_config(config.attention_norm_config), + attention=Attention1D.from_config(config.attention_config), + ff_norm=RMSNorm1D.from_config(config.ff_norm_config), + feed_forward=MLP1D.from_config(config.mlp_config), + decode_residual_memcfg=config.decode_residual_memcfg, + prefill_residual_memcfg=config.prefill_residual_memcfg, + activation_dtype=config.activation_dtype, + ) + + def decode_forward(self, x: ttnn.Tensor, current_pos, rot_mats, page_table) -> ttnn.Tensor: + residual = x + + x = _all_gather_rmsnorm_tensor( + self.attention_norm, x, memory_config=self.attention_norm.config.decode_memory_config + ) + attn_in = self.attention_norm.decode_forward(x) + attn_out = self.attention.decode_forward(attn_in, current_pos, rot_mats, page_table=page_table) + attn_out = ttnn.to_memory_config(attn_out, self.decode_residual_memcfg) + + hidden_states = ttnn.add(residual, attn_out, memory_config=self.decode_residual_memcfg) + residual = hidden_states + + hidden_states = _all_gather_rmsnorm_tensor( + self.ff_norm, hidden_states, memory_config=self.ff_norm.config.decode_memory_config + ) + hidden_states = self.ff_norm.decode_forward(hidden_states) + ttnn.deallocate(attn_out) + hidden_states = self.feed_forward.decode_forward(hidden_states) + + out = ttnn.add( + residual, + hidden_states, + memory_config=self.decode_residual_memcfg, + dtype=self.activation_dtype or ttnn.bfloat16, + ) + return out + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats, + user_id, + page_table, + chunk_page_table, + chunk_start_idx, + batch_size: int = 1, + chunk_start_idx_tensor=None, + ) -> ttnn.Tensor: + residual = x + + attn_in = self.attention_norm.prefill_forward(x) + attn_in = _all_gather_rmsnorm_tensor(self.attention_norm, attn_in) + if batch_size > 1: + attn_in = ttnn.reshape(attn_in, [batch_size, 1, attn_in.shape[-2] // batch_size, -1]) + attn_out = self.attention.prefill_forward( + attn_in, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + if batch_size > 1: + residual = ttnn.reshape(residual, [1, 1, residual.shape[-2] * residual.shape[-3] * residual.shape[0], -1]) + attn_out = ttnn.to_memory_config(attn_out, self.prefill_residual_memcfg) + + hidden_states = ttnn.add(residual, attn_out, memory_config=self.prefill_residual_memcfg) + residual = hidden_states + x.deallocate(True) + + hidden_states = self.ff_norm.prefill_forward(hidden_states) + hidden_states = _all_gather_rmsnorm_tensor(self.ff_norm, hidden_states) + ttnn.deallocate(attn_out) + hidden_states = self.feed_forward.prefill_forward(hidden_states) + + out = ttnn.add( + residual, + hidden_states, + memory_config=self.prefill_residual_memcfg, + dtype=self.activation_dtype or ttnn.bfloat16, + ) + return out + + def forward( + self, + x, + current_pos=None, + rot_mats=None, + user_id=0, + mode="decode", + page_table=None, + chunk_page_table=None, + chunk_start_idx=None, + batch_size: int = 1, + chunk_start_idx_tensor=None, + ): + if mode == "prefill": + return self.prefill_forward( + x, + rot_mats, + user_id, + page_table, + chunk_page_table, + chunk_start_idx, + batch_size, + chunk_start_idx_tensor, + ) + return self.decode_forward(x, current_pos, rot_mats, page_table) + + +# ============================================================================= +# Llama3Transformer1D +# ============================================================================= + + +@dataclass +class Llama31_8BPagedAttentionConfig: + block_size: int + max_num_blocks: int + + +@dataclass +class Llama3Transformer1DConfig: + """Full TTTv2 model config.""" + + n_layers: int + vocab_size: int + max_batch_size: int + max_seq_len: int + dim: int + num_devices: int + mesh_device: ttnn.MeshDevice + + # Sub-module configs + embedding_config: Embedding1DConfig + rope_config: Rope1DConfig + block_configs: list[TransformerBlock1DConfig] + norm_config: RMSNorm1DConfig + lm_head_config: LMHead1DConfig + sampling_config: Sampling1DConfig | None = None + + # Construction-only architecture compositions paired with the public + # common configs above. + + # Model-level memory configs + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + + # Per-layer activation dtypes (from decoders_optimizations) + activation_dtypes: list[ttnn.DataType | None] = field(default_factory=list) + + # CCL + tt_ccl: TT_CCL | None = None + + # Weight cache path + cache_path: "str | None" = None + + +class Llama3Transformer1D(LightweightModule): + """TTTv2 Llama 3.1-8B Transformer. + + Constructor takes a config and builds everything internally: + model = Llama3Transformer1D(config) + + Public sub-modules (accessible by executor for trace support): + - embedding: Embedding1D + - rope_setup: RotarySetup1D + - layers: list[TransformerBlock1D] + - norm: RMSNorm1D (final) + - lm_head: LMHead1D + - sampling: Sampling1D | None + + Forward methods take pre-embedded tensors. The executor handles + embedding, input preparation, and output processing. + """ + + def __init__(self, config: Llama3Transformer1DConfig): + from tqdm import tqdm + + super().__init__() + self.config = config + + tt_ccl_inst = config.tt_ccl + if tt_ccl_inst is None and config.num_devices > 1: + tt_ccl_inst = get_tt_ccl(config.mesh_device) + + self.embedding = Embedding1D.from_config(config.embedding_config) + self.rope_setup = RotarySetup1D.from_config(config.rope_config) + + self.layers = [ + TransformerBlock1D.from_config(config.block_configs[i]) + for i in tqdm(range(config.n_layers), desc="Building layers") + ] + + self.norm = RMSNorm1D.from_config(config.norm_config) + self.lm_head = LMHead1D.from_config(config.lm_head_config) + + self.sampling = None + if config.sampling_config is not None: + self.sampling = Sampling1D.from_config(config.sampling_config) + self.supports_on_device_sampling = self.sampling is not None + + self.mesh_device = config.mesh_device + self.tt_ccl = tt_ccl_inst + self.vocab_size = config.vocab_size + self.n_layers = config.n_layers + self.num_devices = config.num_devices + self.decode_residual_memcfg = config.decode_residual_memcfg + self.prefill_residual_memcfg = config.prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtypes = config.activation_dtypes or [None] * config.n_layers + + # ========================================================================= + # KV Cache binding + # ========================================================================= + + def iter_executor_named_modules(self): + """Yield named submodules that declare executor input contracts.""" + if not hasattr(self, "layers"): + return + + for i, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{i}].{suffix}", submodule + + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + """Replace provisional external-cache geometry before KV tensors exist.""" + + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer} already has a bound KV cache") + + construction_configs = tuple(block.attention_config for block in self.config.block_configs) + attention_configs = tuple({id(config): config for config in (*construction_configs, *live_configs)}.values()) + for config in attention_configs: + config.paged_attention_config = replace( + config.paged_attention_config, + block_size=block_size, + max_num_blocks=max_num_blocks, + ) + + def set_kv_cache(self, kv_cache: list | None): + """Bind or unbind the static KV-cache pool transactionally.""" + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + + cache_pairs = [] + for i, value in enumerate(kv_cache): + try: + cache_pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {i} must provide an iterable K/V tensor pair") from error + if len(cache_pair) != 2: + raise ValueError(f"kv_cache layer {i} must contain exactly two K/V tensors") + cache_pairs.append(cache_pair) + + for layer, cache_pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = cache_pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = cache_pair + + # ========================================================================= + # Forward methods — take pre-embedded tensors + # ========================================================================= + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + """Decode forward. x_embed is already embedded, unsqueezed, and in decode_residual_memcfg.""" + x = x_embed + + for i, layer in enumerate(self.layers): + x = ttnn.to_memory_config(x, self.decode_residual_memcfg, self.activation_dtypes[i]) + + x = layer.decode_forward(x, current_pos, rot_mats, page_table) + + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + x = self.lm_head.forward(x) + return x + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + """Prefill forward. x_embed is already embedded and unsqueezed to 4D.""" + x = x_embed + + for i, layer in enumerate(self.layers): + activation_dtype = self.activation_dtypes[i] + if activation_dtype is not None and x.dtype != activation_dtype: + old = x + x = ttnn.typecast(x, activation_dtype) + ttnn.deallocate(old) + + x = layer.prefill_forward( + x, + rot_mats, + user_id, + page_table, + chunk_page_table, + chunk_start_idx, + batch_size, + chunk_start_idx_tensor, + ) + + if last_token_index is not None and last_token_slice is None: + raise ValueError("last_token_index is required with a runtime last_token_slice") + if get_last_token == -1 and last_token_slice is None: + return x + + old = x + if last_token_slice is None: + get_last_token_floor = (get_last_token // 32) * 32 + x = ttnn.slice( + x, + (0, 0, get_last_token_floor, 0), + (1, 1, get_last_token_floor + 32, x.shape[-1]), + ) + else: + x = ttnn.slice( + x, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(x.shape[2]) // 32, + ) + ttnn.deallocate(old) + + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + x = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + return x + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + """Convert traced prefill hidden states into logits for the last token block.""" + if last_token_slice is None: + get_last_token_floor = (last_token_idx // 32) * 32 + x = ttnn.slice( + hidden_states, + (0, 0, get_last_token_floor, 0), + (1, 1, get_last_token_floor + 32, hidden_states.shape[-1]), + ) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + x = self.norm.prefill_forward(x) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + x = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + return x + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + """Convert batched prefill hidden states into one logits row per slot.""" + x = self.norm.prefill_forward(hidden_states) + x = _all_gather_rmsnorm_tensor(self.norm, x) + x_split = ttnn.split(x, prefill_seq_len, dim=2) + if last_token_slice is None: + selected = [ + x_user[:, :, last_token_idx : last_token_idx + 1, :] + for x_user, last_token_idx in zip(x_split, last_token_idx_list) + ] + else: + if last_token_index is None: + raise ValueError("last_token_index is required with a runtime last_token_slice") + selected = [] + for x_user in x_split[: len(last_token_idx_list)]: + block = ttnn.slice( + x_user, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=prefill_seq_len // 32, + ) + row = ttnn.embedding(last_token_index, block, layout=ttnn.TILE_LAYOUT) + row = ttnn.unsqueeze_to_4D(row) + ttnn.deallocate(block) + selected.append(row) + x = ttnn.concat(selected, dim=2) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + x = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + return x + + def forward( + self, + x: ttnn.Tensor, + current_pos=None, + rot_mats_global=None, + rot_mats_local=None, + user_id: int = 0, + mode: str = "decode", + page_table=None, + chunk_page_table=None, + chunk_start_idx=None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor=None, + last_token_slice=None, + last_token_index=None, + ) -> ttnn.Tensor: + """Dispatcher for backward compatibility. Llama 3.1-8B has no local rope.""" + rot_mats = rot_mats_global + if mode == "prefill": + return self.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + get_last_token=get_last_token, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + last_token_slice=last_token_slice, + last_token_index=last_token_index, + ) + return self.decode_forward( + x, + current_pos, + rot_mats, + page_table=page_table, + ) + + # ========================================================================= + # Embedding + output processing helpers (called by executor) + # ========================================================================= + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + """Gather prefill RoPE rows from runtime device position indices.""" + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + """Embed tokens and prepare for decode. Returns tensor in decode_residual_memcfg.""" + x = self.embedding.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + x = ttnn.to_memory_config(x, self.decode_residual_memcfg) + return x + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + """Embed tokens for prefill. Returns tensor in DRAM interleaved.""" + x = self.embedding.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return x + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + """All-gather logits across devices and untilize for host argmax.""" + if self.num_devices > 1: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + + logits = ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + return logits + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor): + """Increment decode position counters on device.""" + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) + + +# ============================================================================= +# RMSNorm gather helpers +# ============================================================================= + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, x: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + return x + + if memory_config is None: + memory_config = x.memory_config() + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=tt_ccl.get_num_links(), + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + + +def build_llama3_transformer_1d_config( + *, + mesh_device, + instruct: bool, + max_batch_size: int, + max_seq_len: int, + model_name: str, + dim: int, + n_heads: int, + n_kv_heads: int, + n_layers: int, + head_dim: int, + hidden_dim: int, + vocab_size: int, + norm_eps: float, + padded_vocab_size: int, + rope_cos, + rope_sin, + model_cache_path: str | Path, + state_dict, + optimizations="performance", + weight_cache_path=None, + dtype=None, + paged_attention_config=None, + pad_logits_to_power_of_2=False, +) -> Llama3Transformer1DConfig: + """Build explicit TTTv2 module configs from Llama-3.1-8B construction data.""" + num_devices = mesh_device.get_num_devices() + dram_grid_size = mesh_device.dram_grid_size() + device_name = get_device_name(mesh_device) + cluster_shape = list(mesh_device.shape) + cluster_type = ttnn.cluster.get_cluster_type() + arch = mesh_device.arch() + architecture_profile = _resolve_llama31_8b_architecture_profile( + arch=arch, + cluster_type=cluster_type, + device_name=device_name, + model_name=model_name, + dram_grid_width=dram_grid_size.x, + ) + decode_transformation_core_grid = ( + architecture_profile.attention_decode_transformation_core_grid or mesh_device.compute_with_storage_grid_size() + ) + is_galaxy_cluster = cluster_type in ( + ttnn.cluster.ClusterType.GALAXY, + ttnn.cluster.ClusterType.TG, + ttnn.cluster.ClusterType.BLACKHOLE_GALAXY, + ) + if num_devices == 32: + raise ValueError("Llama3Transformer1D only supports 1D mesh topologies.") + + use_paged_kv_cache = paged_attention_config is not None + + if optimizations is None: + decoder_precision = Llama31DecoderPrecision.performance(n_layers, model_name) + elif isinstance(optimizations, str): + decoder_precision = Llama31DecoderPrecision.from_string(optimizations)(n_layers, model_name) + else: + decoder_precision = optimizations + + assert n_heads % cluster_shape[1] == 0 + assert n_kv_heads % cluster_shape[1] == 0 + + tile_padded_batch_rows = ttnn.TILE_SIZE * int(math.ceil(max_batch_size / ttnn.TILE_SIZE)) + qkv_size = head_dim * (2 * n_kv_heads + n_heads) + min_kv_prefill_shard_seqlen = (ttnn.TILE_SIZE * 8 * 8) / (n_kv_heads // cluster_shape[1]) + compute_kernel_config_lofi = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + compute_kernel_config_hifi2 = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=True, + fp32_dest_acc_en=True, + packer_l1_acc=True, + ) + compute_kernel_config_hifi2_fp16 = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + compute_kernel_config_hifi4 = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=True, + ) + compute_kernel_config_hifi4_fp32 = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi4, + fp32_dest_acc_en=True, + packer_l1_acc=True, + dst_full_sync_en=False, + ) + compute_kernel_config_hifi2_na = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=False, + ) + compute_kernel_config_hifi2_nol1acc = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=True, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ) + + def ccl_topology(): + if cluster_type in ( + ttnn.cluster.ClusterType.P150_X2, + ttnn.cluster.ClusterType.P300_X2, + ttnn.cluster.ClusterType.P150_X4, + ttnn.cluster.ClusterType.P150_X8, + ): + return ttnn.Topology.Ring + if cluster_type == ttnn.cluster.ClusterType.T3K: + return ttnn.Topology.Ring if num_devices >= 8 else ttnn.Topology.Linear + if cluster_type in ( + ttnn.cluster.ClusterType.GALAXY, + ttnn.cluster.ClusterType.TG, + ttnn.cluster.ClusterType.BLACKHOLE_GALAXY, + ): + return ttnn.Topology.Linear + return ttnn.Topology.Linear if num_devices > 1 else None + + use_fused_all_gather_matmul = ( + num_devices == 8 + and not is_galaxy_cluster + and (dim // ttnn.TILE_SIZE // num_devices) % num_devices == 0 + and num_devices > 1 + and ccl_topology() == ttnn.Topology.Ring + ) + + dram_weight_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_grid_size.x - 1, dram_grid_size.y - 1))} + ) + + def find_grid(n): + max_rows = 8 if arch == ttnn.device.Arch.WORMHOLE_B0 else 10 + max_cols = 8 if arch == ttnn.device.Arch.WORMHOLE_B0 else 12 + possible_cores = [k for k in range(1, max_rows * max_cols + 1) if n % k == 0] + possible_cores.sort(key=lambda x: abs(x - 32)) + for cores in possible_cores: + for rows in range(1, max_rows + 1): + if cores % rows == 0: + cols = cores // rows + if cols <= max_cols: + return rows, cols + raise AssertionError(f"Cannot find grid for {n} tiles") + + def find_grid_k_n(k, n): + possible_cores = [c for c in range(1, 65) if k % c == 0 and n % c == 0] + possible_cores.sort(reverse=True) + for cores in possible_cores: + for rows in range(1, 9): + if cores % rows == 0: + cols = cores // rows + if cols <= 8: + return rows, cols + raise AssertionError(f"Cannot find grid for K={k}, N={n}") + + def dram_shard_core_grid_for_k(k): + rows, cols = find_grid(k // ttnn.TILE_SIZE) + return ttnn.CoreGrid(x=cols, y=rows) + + def dram_shard_core_grid_for_k_and_n(k, n): + rows, cols = find_grid_k_n(k // ttnn.TILE_SIZE, n // ttnn.TILE_SIZE) + return ttnn.CoreGrid(x=cols, y=rows) + + def find_largest_divisor(n, max_divisor=8): + for i in range(max_divisor, 0, -1): + if n % i == 0: + return i + return 1 + + def create_dram_sharded_mem_config(k, n, dram_grid=None): + dram_cores = dram_grid_size.x + padded_size = math.ceil(n / (ttnn.TILE_SIZE * dram_cores)) * (ttnn.TILE_SIZE * dram_cores) + grid = dram_grid or dram_weight_grid + shard_spec = ttnn.ShardSpec(grid, (k, padded_size // dram_cores), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + def dram_matmul_config(m, k, n, num_cores=None, fused_activation=None): + if num_cores is None: + num_cores = dram_shard_core_grid_for_k_and_n(k, n).num_cores + return ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig( + in0_block_w=find_largest_divisor(k // (ttnn.TILE_SIZE * num_cores)), + per_core_M=math.ceil(m / ttnn.TILE_SIZE), + per_core_N=math.ceil(n / (ttnn.TILE_SIZE * num_cores)), + fused_activation=fused_activation, + ) + + def create_sharded_norm_config(grid): + block_w = dim // grid.num_cores // ttnn.TILE_SIZE + subblock_w = 4 + while subblock_w > 0: + if block_w % subblock_w == 0: + break + subblock_w -= 1 + return ttnn.LayerNormShardedMultiCoreProgramConfig( + compute_with_storage_grid_size=[grid.x, grid.y], + subblock_w=subblock_w, + block_h=tile_padded_batch_rows // ttnn.TILE_SIZE, + block_w=block_w, + inplace=False, + ) + + def decode_all_gather_matmul_program_config(): + if not use_fused_all_gather_matmul: + return None + do_core_grid_size = (8, 1) + do_per_core_n = dim // num_devices // ttnn.TILE_SIZE // (do_core_grid_size[0] * do_core_grid_size[1]) + return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig( + compute_with_storage_grid_size=do_core_grid_size, + in0_block_w=dim // ttnn.TILE_SIZE // (do_core_grid_size[0] * do_core_grid_size[1]), + out_subblock_h=1, + out_subblock_w=get_out_subblock_w(do_per_core_n, out_subblock_h=1), + per_core_M=tile_padded_batch_rows // ttnn.TILE_SIZE, + per_core_N=do_per_core_n, + fuse_batch=True, + fused_activation=None, + mcast_in0=True, + ) + + def decode_all_gather_matmul_output_mem_config(): + return ttnn.MemoryConfig( + ttnn.TensorMemoryLayout.WIDTH_SHARDED, + ttnn.BufferType.L1, + ttnn.ShardSpec( + num_to_core_range_set(num_devices), + [tile_padded_batch_rows, dim // num_devices], + ttnn.ShardOrientation.ROW_MAJOR, + ), + ) + + def decode_residual_mem_config(): + residual_grid = dram_shard_core_grid_for_k(dim // num_devices) + return ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // residual_grid.num_cores // num_devices), + residual_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + lm_head_num_rows = 8 + lm_head_cores_per_row = 8 + while dim % (ttnn.TILE_SIZE * lm_head_num_rows * lm_head_cores_per_row) != 0: + lm_head_num_rows -= 1 + if lm_head_num_rows == 0: + lm_head_cores_per_row -= 1 + if lm_head_cores_per_row == 0: + raise ValueError("Could not find a valid LM head core grid") + lm_head_num_rows = 8 + lm_head_core_grid = ttnn.CoreGrid(y=lm_head_num_rows, x=lm_head_cores_per_row) + max_columns_per_device_lm_head = ( + architecture_profile.lm_head_max_columns_per_device or 668 * lm_head_core_grid.num_cores + ) + attn_input_grid = dram_shard_core_grid_for_k(dim) + mlp_core_grid = dram_shard_core_grid_for_k_and_n(dim, hidden_dim // num_devices) + mlp2_core_grid = dram_shard_core_grid_for_k_and_n(hidden_dim // num_devices, dim) + + def get_decode_norm_config(norm_type): + if norm_type == "attn": + grid = attn_input_grid + mem = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + elif norm_type == "ff": + grid = mlp_core_grid + mem = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + elif norm_type == "lm_head": + grid = lm_head_core_grid + mem = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, nearest_32(dim // grid.num_cores)), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + else: + raise ValueError(f"Invalid norm_type: {norm_type}") + return { + "sharded_program_config": create_sharded_norm_config(grid), + "sharded_output_config": mem, + "output_mem_config": None, + } + + def get_decode_mlp_ff1_3_prg_config(): + return dram_matmul_config(tile_padded_batch_rows, dim, hidden_dim // cluster_shape[1], mlp_core_grid.num_cores) + + def get_decode_mlp_ff2_prg_config(): + return dram_matmul_config(tile_padded_batch_rows, hidden_dim // cluster_shape[1], dim, mlp2_core_grid.num_cores) + + def get_decode_mlp_binary_mult_mem_config(): + return ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, hidden_dim // cluster_shape[1] // mlp2_core_grid.num_cores), + mlp2_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + def get_tensor_dtype(layer_num, tensor): + return decoder_precision.get_tensor_dtype(layer_num, tensor) + + def get_math_fidelity(layer_num, op): + kernel_lookup = { + "lofi": compute_kernel_config_lofi, + "hifi2": compute_kernel_config_hifi2, + "hifi2na": compute_kernel_config_hifi2_na, + "hifi2fp16": compute_kernel_config_hifi2_fp16, + "hifi2nol1acc": compute_kernel_config_hifi2_nol1acc, + "hifi4": compute_kernel_config_hifi4, + "hifi4fp32": compute_kernel_config_hifi4_fp32, + } + return kernel_lookup[decoder_precision._op_fidelity[layer_num][op]] + + def get_state_dict_prefix(module_name, layer_num): + layer_prefix = f"layers.{layer_num}." if layer_num is not None else "" + module_map = {"MLP": "feed_forward", "Attention": "attention", "TransformerBlock": "", "": ""} + return layer_prefix + module_map[module_name] + + def cache_path(dtype): + cache_path_root = Path(model_cache_path) + if instruct: + return ( + cache_path_root + / { + ttnn.bfloat16: "tensor_cache_instruct_bf16", + ttnn.bfloat8_b: "tensor_cache_instruct_bfp8", + }[dtype] + ) + return cache_path_root / {ttnn.bfloat16: "tensor_cache_bf16", ttnn.bfloat8_b: "tensor_cache_bfp8"}[dtype] + + model_config = { + "SDPA_DECODE_PROGCFG": ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=(8, 8), + exp_approx_mode=False, + q_chunk_size=0, + k_chunk_size=0, + ), + "CREATE_QKV_DECODE_SHARD": ( + ttnn.create_sharded_memory_config( + shape=(ttnn.TILE_SIZE, head_dim), + core_grid=ttnn.CoreGrid(y=4, x=8), + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + if arch == ttnn.device.Arch.BLACKHOLE + else ttnn.L1_HEIGHT_SHARDED_MEMORY_CONFIG + ), + "ATTN_OUTPUT_PROGCFG": dram_matmul_config( + m=tile_padded_batch_rows, + k=(n_heads * head_dim) // num_devices, + n=dim, + num_cores=n_heads // num_devices, + ), + "ATTN_ALL_GATHER_MATMUL_PROGCFG": decode_all_gather_matmul_program_config(), + "ATTN_ALL_GATHER_MATMUL_OUTPUT_MEMCFG": decode_all_gather_matmul_output_mem_config(), + "MLP_RS_CONFIG": { + "chunks_per_sync": 10, + "num_workers_per_link": 2, + "rs_memory_config": ttnn.DRAM_MEMORY_CONFIG, + }, + } + model_config["DECODE_RESIDUAL_MEMCFG"] = decode_residual_mem_config() + + tt_ccl_inst = get_tt_ccl(mesh_device) if num_devices > 1 else None + weight_cache_path = Path(weight_cache_path) if weight_cache_path else None + embedding_cache_path = cache_path(dtype or ttnn.bfloat8_b) + + def mesh_shard(dim: int) -> ttnn.MeshMapperConfig: + return ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(dim)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ) + + def cache_path_for( + base: str | os.PathLike[str] | None, + *parts: str | os.PathLike[str], + ) -> Path | None: + if base is None: + return None + return Path(base).joinpath(*parts) + + def make_embedding_config() -> Embedding1DConfig: + base_name = get_state_dict_prefix("", None) + "tok_embeddings.weight" + torch_weight = state_dict[base_name].unsqueeze(0).unsqueeze(0) + cache_dir = cache_path_for(embedding_cache_path, "embedding") + return Embedding1DConfig( + weights=LazyWeight( + source=torch_weight, + dtype=ttnn.bfloat16, + device=mesh_device, + mesh_mapper_config=mesh_shard(-1), + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(cache_dir, "tok_embeddings") if cache_dir else None, + ), + mesh_device=mesh_device, + weights_dtype=ttnn.bfloat16, + weights_memcfg=ttnn.DRAM_MEMORY_CONFIG, + output_memcfg=ttnn.DRAM_MEMORY_CONFIG, + ) + + def make_rope_config() -> Rope1DConfig: + return _make_llama31_8b_rope_config( + rope_cos=rope_cos, + rope_sin=rope_sin, + max_batch_size=max_batch_size, + head_dim=head_dim, + mesh_device=mesh_device, + decode_transformation_core_grid=decode_transformation_core_grid, + ) + + def norm_weight_name(layer_num: int | None, weight_key: str, state_dict_prefix: str | None = None) -> str: + if state_dict_prefix: + return f"{state_dict_prefix}{weight_key}.weight" + if layer_num is None: + return f"{weight_key}.weight" + return f"layers.{layer_num}.{weight_key}.weight" + + def make_norm_config( + *, + layer_num: int | None, + weight_key: str, + state_dict_prefix: str | None = None, + sharded_program_config=None, + sharded_output_config=None, + ) -> RMSNorm1DConfig: + weight_name = norm_weight_name(layer_num, weight_key, state_dict_prefix) + torch_weight = ( + state_dict[weight_name].unsqueeze(0).view(1, 1, dim).reshape([1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT]) + ) + return RMSNorm1DConfig( + weight=LazyWeight( + source=torch_weight, + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(weight_cache_path, weight_name) if weight_cache_path else None, + mesh_mapper_config=( + ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate()], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ) + if num_devices > 1 + else None + ), + ), + eps=norm_eps, + mesh_device=mesh_device, + tt_ccl=tt_ccl_inst, + max_batch_size=max_batch_size, + prefill_distributed=_use_distributed_prefill_rmsnorm( + num_devices=num_devices, + dim=dim, + architecture_profile=architecture_profile, + ), + decode_program_config=sharded_program_config, + decode_memory_config=sharded_output_config, + compute_kernel_config=ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=architecture_profile.rms_packer_l1_acc, + ), + ) + + def make_attention_config(layer_num: int, transformation_mats: dict[str, ttnn.Tensor]) -> Attention1DConfig: + layer_name = get_state_dict_prefix("Attention", layer_num) + wq_str = f"{layer_name}.wq" + wk_str = f"{layer_name}.wk" + wv_str = f"{layer_name}.wv" + wo_str = f"{layer_name}.wo" + q_norm_str = f"{layer_name}.q_norm" + k_norm_str = f"{layer_name}.k_norm" + + wqkv_dtype = get_tensor_dtype(layer_num, "wqkv") + wo_dtype = get_tensor_dtype(layer_num, "wo") + kv_cache_dtype = get_tensor_dtype(layer_num, "kv_cache") + activation_dtype = get_tensor_dtype(layer_num, "activation") + + qkv_list = [] + for device_idx in range(num_devices): + wq = torch.transpose(torch.chunk(state_dict[f"{wq_str}.weight"], num_devices, dim=0)[device_idx], -2, -1) + wk = torch.transpose(torch.chunk(state_dict[f"{wk_str}.weight"], num_devices, dim=0)[device_idx], -2, -1) + wv = torch.transpose(torch.chunk(state_dict[f"{wv_str}.weight"], num_devices, dim=0)[device_idx], -2, -1) + qkv_list.append(torch.cat([wq, wk, wv], dim=-1)) + qkv_cat = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0) + + wqkv = LazyWeight( + source=qkv_cat, + dtype=wqkv_dtype, + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + memory_config=create_dram_sharded_mem_config(dim, qkv_size // num_devices), + mesh_mapper_config=mesh_shard(-1), + cache_dir_weight_name=(weight_cache_path / layer_name, "wqkv_sharded") if weight_cache_path else None, + ) + wo = LazyWeight( + source=state_dict[f"{wo_str}.weight"].transpose(-1, -2).unsqueeze(0).unsqueeze(0), + dtype=wo_dtype, + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + memory_config=( + ttnn.DRAM_MEMORY_CONFIG + if use_fused_all_gather_matmul + else create_dram_sharded_mem_config((n_heads * head_dim) // num_devices, dim) + ), + mesh_mapper_config=mesh_shard(-1 if use_fused_all_gather_matmul else -2), + cache_dir_weight_name=( + (weight_cache_path / layer_name, "wo_width_sharded" if use_fused_all_gather_matmul else "wo") + if weight_cache_path + else None + ), + ) + + qk_norm_compute_kernel = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ) + + def make_qk_norm_config(name: str) -> RMSNorm1DConfig | None: + weight_name = f"{name}.weight" + if weight_name not in state_dict: + return None + return RMSNorm1DConfig( + weight=LazyWeight( + source=state_dict[weight_name].reshape(1, 1, -1, TILE_SIZE), + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=( + (weight_cache_path / layer_name, name.rsplit(".", 1)[-1]) if weight_cache_path else None + ), + ), + mesh_device=mesh_device, + eps=norm_eps, + decode_in_sharded=False, + decode_out_sharded=False, + prefill_distributed=False, + compute_kernel_config=qk_norm_compute_kernel, + ) + + wqkv_bias = None + if f"{wq_str}.bias" in state_dict: + wqkv_bias = LazyWeight( + source=torch.concat( + [ + torch.concat( + [ + torch.chunk(state_dict[f"{wq_str}.bias"], num_devices)[device_idx], + torch.chunk(state_dict[f"{wk_str}.bias"], num_devices)[device_idx], + torch.chunk(state_dict[f"{wv_str}.bias"], num_devices)[device_idx], + ], + dim=-1, + ) + for device_idx in range(num_devices) + ], + dim=-1, + ) + ) + + scale = head_dim**-0.5 + return Attention1DConfig( + wqkv=wqkv, + wo=wo, + q_norm_config=make_qk_norm_config(q_norm_str), + k_norm_config=make_qk_norm_config(k_norm_str), + wqkv_bias=wqkv_bias, + mesh_device=mesh_device, + tt_ccl=tt_ccl_inst, + topology=ccl_topology(), + dim=dim, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + head_dim=head_dim, + qkv_size=qkv_size, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + scale=scale, + use_qk_fused=True, + use_vllm_paged_kv_cache=use_paged_kv_cache, + paged_attention_config=paged_attention_config, + kv_cache_dtype=kv_cache_dtype, + min_kv_prefill_shard_seqlen=min_kv_prefill_shard_seqlen, + wqkv_dtype=wqkv_dtype, + wo_dtype=wo_dtype, + activation_dtype=activation_dtype, + decode_sdpa_prg_config=model_config.get("SDPA_DECODE_PROGCFG"), + decode_attn_output_prg_config=model_config.get("ATTN_OUTPUT_PROGCFG"), + decode_residual_memcfg=model_config.get("DECODE_RESIDUAL_MEMCFG"), + decode_create_qkv_head_memcfg=model_config.get("CREATE_QKV_DECODE_SHARD"), + use_fused_all_gather_matmul=use_fused_all_gather_matmul, + decode_all_gather_matmul_prg_config=model_config.get("ATTN_ALL_GATHER_MATMUL_PROGCFG"), + decode_all_gather_matmul_memcfg=model_config.get("ATTN_ALL_GATHER_MATMUL_OUTPUT_MEMCFG"), + li_qkv_decode_compute_kernel_cfg=get_math_fidelity(layer_num, "li_qkv_decode"), + sdpa_decode_compute_kernel_cfg=get_math_fidelity(layer_num, "sdpa_decode"), + li_o_decode_compute_kernel_cfg=get_math_fidelity(layer_num, "li_o_decode"), + li_qkv_prefill_compute_kernel_cfg=get_math_fidelity(layer_num, "li_qkv_prefill"), + sdpa_prefill_compute_kernel_cfg=get_math_fidelity(layer_num, "sdpa_prefill"), + li_o_prefill_compute_kernel_cfg=get_math_fidelity(layer_num, "li_o_prefill"), + prefill_qkv_grid=architecture_profile.attention_prefill_qkv_grid, + dram_shard_grid_width=( + 8 if arch == ttnn.device.Arch.WORMHOLE_B0 else architecture_profile.mlp_prefill_dram_shard_grid_width + ), + decode_create_qkv_head_grid=architecture_profile.attention_decode_create_qkv_head_grid, + decode_transformation_core_grid=decode_transformation_core_grid, + prefill_qkv_minimal_matmul=architecture_profile.enable_minimal_qkv, + transformation_mat_decode=transformation_mats.get("decode"), + transformation_mat_prefill=transformation_mats.get("prefill"), + ) + + def make_mlp_config(layer_num: int) -> MLP1DConfig: + state_dict_prefix = get_state_dict_prefix("MLP", layer_num) + ff1_3_dtype = get_tensor_dtype(layer_num, "ff1_ff3") + ff2_dtype = get_tensor_dtype(layer_num, "ff2") + activation_dtype = get_tensor_dtype(layer_num, "activation") + mlp_rs_cfg = model_config.get("MLP_RS_CONFIG", {}) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + w1_w3_mem_config = _create_dram_sharded_mem_config( + k=dim, + n=hidden_dim // num_devices, + dram_grid=dram_grid, + tile_size=TILE_SIZE, + dram_cores=dram_size.x, + ) + w2_mem_config = _create_dram_sharded_mem_config( + k=hidden_dim // num_devices, + n=dim, + dram_grid=dram_grid, + tile_size=TILE_SIZE, + dram_cores=dram_size.x, + ) + cache_dir = cache_path_for(weight_cache_path, state_dict_prefix) + + def make_weight_source(name: str, shard_dim: int): + tensor = torch.transpose(state_dict[f"{state_dict_prefix}.{name}.weight"], -2, -1) + return pad_dim_to_size(tensor, dim=shard_dim, size=hidden_dim) + + return MLP1DConfig( + w1=LazyWeight( + source=make_weight_source("w1", -1), + dtype=ff1_3_dtype, + device=mesh_device, + mesh_mapper_config=mesh_shard(-1), + layout=ttnn.TILE_LAYOUT, + memory_config=w1_w3_mem_config, + cache_dir_weight_name=(cache_dir, "w1_sharded") if cache_dir else None, + ), + w2=LazyWeight( + source=make_weight_source("w2", -2), + dtype=ff2_dtype, + device=mesh_device, + mesh_mapper_config=mesh_shard(-2), + layout=ttnn.TILE_LAYOUT, + memory_config=w2_mem_config, + cache_dir_weight_name=(cache_dir, "w2_sharded") if cache_dir else None, + ), + w3=LazyWeight( + source=make_weight_source("w3", -1), + dtype=ff1_3_dtype, + device=mesh_device, + mesh_mapper_config=mesh_shard(-1), + layout=ttnn.TILE_LAYOUT, + memory_config=w1_w3_mem_config, + cache_dir_weight_name=(cache_dir, "w3_sharded") if cache_dir else None, + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl_inst, + dim=dim, + hidden_dim=hidden_dim, + max_batch_size=max_batch_size, + mlp_activation_type=ttnn.UnaryOpType.SILU, + topology=ccl_topology(), + decode_rs_memory_config=mlp_rs_cfg.get("rs_memory_config", ttnn.L1_MEMORY_CONFIG), + decode_rs_chunks_per_sync=mlp_rs_cfg.get("chunks_per_sync", 1), + decode_rs_num_workers_per_link=mlp_rs_cfg.get("num_workers_per_link", 1), + decode_w1_w3_prg_config=get_decode_mlp_ff1_3_prg_config(), + decode_w2_prg_config=get_decode_mlp_ff2_prg_config(), + decode_mlp2_input_memcfg=get_decode_mlp_binary_mult_mem_config(), + decode_residual_memcfg=decode_residual_mem_config(), + w1_w3_dtype=ff1_3_dtype, + w2_dtype=ff2_dtype, + activation_dtype=activation_dtype, + ff1_3_compute_kernel_cfg=get_math_fidelity(layer_num, "li_ff1_ff3"), + ff2_compute_kernel_cfg=get_math_fidelity(layer_num, "li_ff2"), + decode_ff1_3_compute_kernel_cfg=get_math_fidelity(layer_num, "li_ff1_ff3"), + decode_ff2_compute_kernel_cfg=get_math_fidelity(layer_num, "li_ff2"), + prefill_len_cutoff=architecture_profile.mlp_prefill_len_cutoff, + prefill_dram_shard_grid_width=architecture_profile.mlp_prefill_dram_shard_grid_width, + prefill_ff1_ff3_grid=architecture_profile.mlp_prefill_ff1_ff3_grid, + prefill_ff2_grid=architecture_profile.mlp_prefill_ff2_grid, + prefill_w2_minimal_matmul=architecture_profile.enable_minimal_ff2, + ) + + def make_lm_head_config() -> LMHead1DConfig: + lm_head_padded_vocab_size = math.ceil(vocab_size / (TILE_SIZE * num_devices)) * (TILE_SIZE * num_devices) + size_per_device = lm_head_padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device_lm_head) + split_sizes = [min(size_per_device, max_columns_per_device_lm_head)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + state_dict_prefix = get_state_dict_prefix("", None) + source_weight = state_dict[f"{state_dict_prefix}output.weight"] + if tuple(source_weight.shape) != (vocab_size, dim): + raise ValueError( + f"Llama 8B LM-head weight must have shape {(vocab_size, dim)}, got {tuple(source_weight.shape)}" + ) + torch_output_weights = source_weight.permute(1, 0) + if vocab_size < lm_head_padded_vocab_size: + torch_output_weights = torch.cat( + [ + torch_output_weights, + torch.zeros( + torch_output_weights.shape[0], + lm_head_padded_vocab_size - vocab_size, + dtype=torch_output_weights.dtype, + ), + ], + dim=-1, + ) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + cache_dir = cache_path_for(weight_cache_path, "lm_head") + output_weights = [] + weights_memcfgs = [] + for split_idx, split_size in enumerate(split_sizes): + device_splits = [] + physical_split_size = math.ceil(split_size / TILE_SIZE) * TILE_SIZE + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:split_idx]) + end = start + split_size + device_split = torch_output_weights[:, start:end] + if split_size < physical_split_size: + device_split = torch.cat( + [ + device_split, + torch.zeros(dim, physical_split_size - split_size, dtype=device_split.dtype), + ], + dim=-1, + ) + device_splits.append(device_split) + combined_split = torch.cat(device_splits, dim=-1) + mem_cfg = _create_dram_sharded_mem_config( + k=dim, + n=math.ceil(combined_split.shape[-1] / num_devices), + dram_grid=dram_grid, + tile_size=TILE_SIZE, + dram_cores=dram_size.x, + ) + weights_memcfgs.append(mem_cfg) + output_weights.append( + LazyWeight( + source=combined_split, + dtype=dtype if dtype is not None else ttnn.bfloat8_b, + device=mesh_device, + mesh_mapper_config=mesh_shard(-1), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=( + ( + cache_dir, + f"output_split_{split_idx}_logical_{split_size}_physical_{combined_split.shape[-1]}", + ) + if cache_dir + else None + ), + ) + ) + + lm_head_tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + input_memcfg = ttnn.create_sharded_memory_config( + ( + lm_head_tile_padded_batch_rows, + math.ceil((dim // lm_head_core_grid.num_cores) / TILE_SIZE) * TILE_SIZE, + ), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1DConfig( + output_weights=output_weights, + mesh_device=mesh_device, + dim=dim, + max_batch_size=max_batch_size, + program_configs=[ + dram_matmul_config(lm_head_tile_padded_batch_rows, dim, split_size, lm_head_core_grid.num_cores) + for split_size in split_sizes + ], + output_split_sizes=split_sizes, + output_memcfg=ttnn.L1_MEMORY_CONFIG, + input_memcfg=input_memcfg, + weights_memcfgs=weights_memcfgs, + compute_kernel_config=_compute_kernel_config_hifi2(arch), + ) + + def make_sampling_config() -> Sampling1DConfig | None: + sampling_splits = num_devices if list(mesh_device.shape) != [1, 1] else 2 + if vocab_size // sampling_splits > 64 * 1024: + return None + + return Sampling1DConfig( + vocab_size=padded_vocab_size, + valid_vocab_size=vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl_inst, + max_batch_size=tile_padded_batch_rows, + pad_to_power_of_2=pad_logits_to_power_of_2, + # Decode uses force-argmax for greedy rows; prefill can still force + # the top-k path at the executor call site when a platform needs it. + allow_force_argmax=True, + num_argmax_gather_links=1, + ag_topology=ttnn.Topology.Linear, + argmax_num_workers_per_link=2, + ) + + rope_config = make_rope_config() + trans_mats_dict = RotarySetup1D.from_config(rope_config).get_both_trans_mats() + attn_norm_cfg = get_decode_norm_config("attn") + ff_norm_cfg = get_decode_norm_config("ff") + lm_head_norm_cfg = get_decode_norm_config("lm_head") + activation_dtypes = [get_tensor_dtype(i, "activation") for i in range(n_layers)] + + block_configs = [] + for i in range(n_layers): + attention_norm_config = make_norm_config( + layer_num=i, + weight_key="attention_norm", + sharded_program_config=attn_norm_cfg.get("sharded_program_config"), + sharded_output_config=attn_norm_cfg.get("sharded_output_config"), + ) + attention_config = make_attention_config(i, trans_mats_dict) + ff_norm_config = make_norm_config( + layer_num=i, + weight_key="ffn_norm", + sharded_program_config=ff_norm_cfg.get("sharded_program_config"), + sharded_output_config=ff_norm_cfg.get("sharded_output_config"), + ) + mlp_config = make_mlp_config(i) + block_configs.append( + TransformerBlock1DConfig( + attention_norm_config=attention_norm_config, + attention_config=attention_config, + ff_norm_config=ff_norm_config, + mlp_config=mlp_config, + decode_residual_memcfg=model_config["DECODE_RESIDUAL_MEMCFG"], + activation_dtype=activation_dtypes[i], + ) + ) + + norm_config = make_norm_config( + layer_num=None, + weight_key="norm", + state_dict_prefix=get_state_dict_prefix("", None), + sharded_program_config=lm_head_norm_cfg.get("sharded_program_config"), + sharded_output_config=lm_head_norm_cfg.get("sharded_output_config"), + ) + lm_head_config = make_lm_head_config() + + return Llama3Transformer1DConfig( + n_layers=n_layers, + vocab_size=vocab_size, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + dim=dim, + num_devices=num_devices, + mesh_device=mesh_device, + embedding_config=make_embedding_config(), + rope_config=rope_config, + block_configs=block_configs, + norm_config=norm_config, + lm_head_config=lm_head_config, + sampling_config=make_sampling_config(), + decode_residual_memcfg=model_config["DECODE_RESIDUAL_MEMCFG"], + activation_dtypes=activation_dtypes, + tt_ccl=tt_ccl_inst, + cache_path=str(weight_cache_path) if weight_cache_path else None, + ) diff --git a/code/models/common/models/mistral_7b/README.md b/code/models/common/models/mistral_7b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..5033f04b0a379b225fecaa2f308ef29a8c81b877 --- /dev/null +++ b/code/models/common/models/mistral_7b/README.md @@ -0,0 +1,84 @@ +# Mistral-7B with TTTv2 + +This directory is the model-owned TTTv2 path for the Mistral-7B family. + +It intentionally demonstrates direct executor construction from +`models/common/llm_runtime`. It is not part of the Llama/Qwen executor +consolidation and does not use `models/common/models/executor.py`. + +## Product path + +```text +Hugging Face checkpoint + -> hf_adaptor.py: provider metadata, tokenizer, and weight conversion + -> model.py: Mistral tensor graph composed from TTTv2 modules + -> executor.py: direct composition of common runtime owners for one lane + -> generator.py: vLLM construction, DP composition, and dispatch +``` + +## Files + +| File | Responsibility | +| --- | --- | +| `hf_adaptor.py` | Resolve provider configuration/tokenizer and construct the product model | +| `weight_utils.py` | Convert and map provider weights | +| `model.py` | Build and execute the TTTv2 Mistral transformer graph | +| `executor.py` | Directly compose one execution lane and own its resources | +| `generator.py` | Build lanes, configure the vLLM boundary, and select eager/traced execution | + +## Tensor-module composition + +`model.py` composes: + +- `Embedding1D` +- `RotarySetup1D` +- `RMSNorm1D` +- `Attention1D` +- `MLP1D` +- `LMHead1D` +- optional `Sampling1D` +- common TT collective helpers + +Mistral-specific attention, RoPE, precision, and device-tuning policy remains +model-owned. + +## Direct executor composition + +`Mistral7BExecutor` directly constructs: + +```text +Mistral7B model +├── PagedKVCacheManager +├── OutputReader +├── PrefillRuntime +├── DecodeRuntime +├── ProgramCompiler +├── EagerExecutor +├── optional TraceCompiler +├── optional TracedExecutor +└── WarmupCoordinator +``` + +This is a supported alternative to the shared model-layer `ModelExecutor`. +Models with distinct orchestration may compose the focused `llm_runtime` +modules directly without subclassing or modifying a universal executor. + +The lane executor owns paged KV, compile/trace registries, output leases, +sampling buffers, and deterministic cleanup. The generator owns orchestration +only and does not own TT tensors. + +## vLLM and data parallelism + +`Mistral7BGenerator` builds one model/executor per lane and uses +`LaneGroupExecutor` when `tt_data_parallel > 1`. `VLLMAdapter` normalizes the +server boundary and validates the vLLM-selected KV-cache specification. + +## Tests + +Relevant entry points include: + +- `models/common/tests/models/mistral_7b/test_hf_adaptor.py` +- `models/common/tests/models/mistral_7b/test_demo_contract.py` +- `models/common/tests/models/mistral_7b/test_prefill_last_token_contract.py` +- `models/common/tests/demos/mistral_7b/demo.py` +- `models/common/tests/llm_runtime/test_executor_integration.py` diff --git a/code/models/common/models/mistral_7b/__init__.py b/code/models/common/models/mistral_7b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..46d7229d41c08259862baf965f34f3ca05e3d313 --- /dev/null +++ b/code/models/common/models/mistral_7b/__init__.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +from models.common.models.mistral_7b.executor import ( + Mistral7BExecutor, + Mistral7BExecutorConfig, + build_mistral_7b_executor, +) +from models.common.models.mistral_7b.generator import ( + Mistral7BGenerator, + Mistral7BGeneratorConfig, + build_mistral_7b_generator, +) +from models.common.models.mistral_7b.hf_adaptor import ( + Mistral7BForCausalLM, + Mistral7BRuntimeConfig, +) +from models.common.models.mistral_7b.model import ( + MISTRAL_ACCURACY, + MISTRAL_PERFORMANCE, + Mistral7B, + Mistral7BPagedAttentionConfig, + Mistral7BPrecisionConfig, + Mistral7BTransformerConfig, + build_mistral_7b_transformer_config, +) + +__all__ = [ + "MISTRAL_ACCURACY", + "MISTRAL_PERFORMANCE", + "Mistral7B", + "Mistral7BExecutor", + "Mistral7BExecutorConfig", + "Mistral7BForCausalLM", + "Mistral7BGenerator", + "Mistral7BGeneratorConfig", + "Mistral7BPagedAttentionConfig", + "Mistral7BPrecisionConfig", + "Mistral7BRuntimeConfig", + "Mistral7BTransformerConfig", + "build_mistral_7b_executor", + "build_mistral_7b_generator", + "build_mistral_7b_transformer_config", +] diff --git a/code/models/common/models/mistral_7b/executor.py b/code/models/common/models/mistral_7b/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..57cb9f6021494160a54b0b353a6f8cf3aaca76c1 --- /dev/null +++ b/code/models/common/models/mistral_7b/executor.py @@ -0,0 +1,613 @@ +# SPDX-FileCopyrightText: 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Model-owned Mistral-7B execution composition and cleanup root.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any, Sequence + +import torch + +from models.common.llm_runtime.config import PagedKVCacheConfig, PageTableLayout, TraceConfig, WarmupConfig +from models.common.llm_runtime.decode import DecodeRuntime, DecodeRuntimeConfig +from models.common.llm_runtime.execution import EagerExecutor, TracedExecutor +from models.common.llm_runtime.output_reader import OutputReader +from models.common.llm_runtime.paged_kv_cache import PagedKVCacheManager +from models.common.llm_runtime.prefill.config import PrefillRuntimeConfig +from models.common.llm_runtime.prefill.runtime import PrefillRuntime +from models.common.llm_runtime.program_compiler import ProgramCompiler +from models.common.llm_runtime.tensor_resources import attach_cleanup_failures +from models.common.llm_runtime.trace_compiler import TraceCompiler +from models.common.llm_runtime.warmup import WarmupCoordinator, WarmupCoordinatorConfig +from models.common.models.mistral_7b.hf_adaptor import Mistral7BForCausalLM +from models.common.modules.sampling.sampling_1d import Sampling1D +from models.common.sampling import SamplingParams + + +@dataclass(frozen=True) +class Mistral7BExecutorConfig: + """Immutable aggregate policy paired with one model-owned executor.""" + + trace: TraceConfig + warmup: WarmupConfig + paged_kv_cache: PagedKVCacheConfig + device_sampling_enabled: bool + + def __post_init__(self) -> None: + nested_configs = ( + ("trace", self.trace, TraceConfig), + ("warmup", self.warmup, WarmupConfig), + ("paged_kv_cache", self.paged_kv_cache, PagedKVCacheConfig), + ) + for name, value, expected_type in nested_configs: + if type(value) is not expected_type: + raise TypeError(f"{name} must be exactly {expected_type.__name__}") + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + + +class Mistral7BExecutor: + """Compose every runtime owner for one Mistral-7B execution lane. + + Construction wires one `PrefillRuntime`, `DecodeRuntime`, + `ProgramCompiler`, and `EagerExecutor`. Trace-enabled + configurations add one `TraceCompiler` and one + `TracedExecutor` over that exact eager instance. + + ``Mistral7BGenerator`` or ``LaneGroupExecutor`` first configures and allocates + KV cache, then warms or compiles programs, and finally calls + `prefill_forward` and `decode_forward` with an explicit + execution target. `cleanup` is the deterministic resource-release + root and makes the executor terminal. + """ + + requires_prefill_trace_warmup = True + + def __init__(self, model: Any, runtime_config: Any, config: Mistral7BExecutorConfig) -> None: + if not isinstance(config, Mistral7BExecutorConfig): + raise TypeError("config must be a Mistral7BExecutorConfig") + iter_modules = getattr(model, "iter_executor_named_modules", None) + if not callable(iter_modules): + raise TypeError("model must provide iter_executor_named_modules()") + can_enable_trace = getattr(runtime_config, "can_enable_trace", None) + if not callable(can_enable_trace): + raise TypeError("runtime_config must provide can_enable_trace()") + model_config = getattr(model, "config", None) + mesh_device = getattr(model_config, "mesh_device", None) + if mesh_device is None: + raise ValueError("model.config.mesh_device is required") + + self.model = model + self.runtime_config = runtime_config + self.model_args = runtime_config + self.config = config + self.mesh_device = mesh_device + self.cache_path = getattr(runtime_config, "model_cache_path", None) + self._terminal = False + self._cleaned_up = False + self._sampling_buffers_loaded = False + self._runtime_configuration_sealed = False + self._q128_topk_tile_ends_warmed: set[bool] = set() + + sampling = getattr(model, "sampling", None) + if config.device_sampling_enabled: + if not isinstance(sampling, Sampling1D): + raise TypeError("device sampling requires model.sampling to be a Sampling1D") + is_resolved = getattr(getattr(sampling, "config", None), "is_resolved", None) + if not callable(is_resolved) or not is_resolved(): + raise ValueError("model.sampling must have a resolved Sampling1DConfig") + + self.kv_cache_manager = PagedKVCacheManager(model, config.paged_kv_cache) + self.page_table_layout = self._resolve_page_table_layout() + self.output_reader = OutputReader(mesh_device) + self.prefill_runtime = PrefillRuntime( + PrefillRuntimeConfig.resolve( + model=model, + output_reader=self.output_reader, + page_table_layout=self.page_table_layout, + max_batch_size=int(model.config.max_batch_size), + max_prefill_chunk_size=int(runtime_config.max_prefill_chunk_size), + device_sampling_enabled=config.device_sampling_enabled, + can_enable_trace=runtime_config.can_enable_trace, + supports_batched_prefill=bool(runtime_config.supports_batched_prefill), + disable_batched_prefill=bool(runtime_config.disable_batched_prefill), + max_prefill_batch_size=int(runtime_config.max_prefill_batch_size), + batched_prefill_batched_extract=bool(runtime_config.batched_prefill_batched_extract), + ) + ) + self.decode_runtime = DecodeRuntime( + DecodeRuntimeConfig.resolve( + model=model, + output_reader=self.output_reader, + lane_capacity=int(model.config.max_batch_size), + page_table_layout=self.page_table_layout, + device_sampling_enabled=config.device_sampling_enabled, + force_greedy_top_k=config.warmup.include_decode_top_k, + ) + ) + self.program_compiler = ProgramCompiler(mesh_device, lambda: self.kv_cache_manager.bound_context) + self.eager_executor = EagerExecutor( + prefill=self.prefill_runtime, + decode=self.decode_runtime, + program_compiler=self.program_compiler, + ) + self.trace_compiler: TraceCompiler | None = None + self.traced_executor: TracedExecutor | None = None + if config.trace.mode != "none": + self.trace_compiler = TraceCompiler(self.program_compiler) + self.traced_executor = TracedExecutor( + eager=self.eager_executor, + trace_compiler=self.trace_compiler, + trace_mode=config.trace.mode, + ) + self.eager_execution = self.eager_executor + self.traced_prefill_execution = ( + self.traced_executor if config.trace.prefill_enabled and self.traced_executor is not None else None + ) + self.traced_decode_execution = ( + self.traced_executor if config.trace.decode_enabled and self.traced_executor is not None else None + ) + self._prefill_execution = self.traced_prefill_execution or self.eager_executor + self._decode_execution = self.traced_decode_execution or self.eager_executor + + prefill_sequence_lengths = getattr(runtime_config, "trace_prefill_supported_seq_lens", (128,)) + self.warmup = WarmupCoordinator( + config=WarmupCoordinatorConfig.resolve( + warmup=config.warmup, + trace=config.trace, + prefill=self.prefill_runtime.config, + decode=self.decode_runtime.config, + prefill_sequence_lengths=prefill_sequence_lengths, + ), + execution=self.traced_executor or self.eager_executor, + ensure_sampling_buffers=self._ensure_sampling_buffers, + validate_bound_cache=self._validate_bound_cache, + ) + + # Public model execution API + + @property + def model_config(self): + return self.model.config + + @property + def cluster_shape(self) -> list[int]: + return list(self.mesh_device.shape) + + @property + def paged_kv_cache_config(self) -> PagedKVCacheConfig: + return self.kv_cache_manager.config + + @property + def terminal(self) -> bool: + return self._terminal + + @property + def already_warmed_up_prefill(self) -> bool: + return self.warmup.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value: bool) -> None: + if not isinstance(value, bool): + raise TypeError("already_warmed_up_prefill compatibility value must be bool") + # Compatibility writes are intentionally non-authoritative. Warmup + # coverage is derived from the coordinator's completed case ledger. + return None + + def configure_paged_kv_cache(self, config: PagedKVCacheConfig) -> None: + """Resolve vLLM-owned KV geometry before the first allocation.""" + + self._ensure_active() + if self._runtime_configuration_sealed: + raise RuntimeError("runtime configuration is sealed") + if not isinstance(config, PagedKVCacheConfig): + raise TypeError("config must be a PagedKVCacheConfig") + if self.kv_cache_manager.config.is_resolved(): + raise RuntimeError("paged KV cache configuration is already resolved") + if config.dtype != self.kv_cache_manager.config.dtype: + raise ValueError("resolved paged KV cache cannot change dtype") + if config.memory_config != self.kv_cache_manager.config.memory_config: + raise ValueError("resolved paged KV cache cannot change memory_config") + self.model.configure_paged_attention( + block_size=config.block_size, + max_num_blocks=config.max_num_blocks, + ) + self.kv_cache_manager.configure(config) + self.config = replace(self.config, paged_kv_cache=config) + self._refresh_page_table_layout() + + def allocate_kv_cache( + self, + kv_cache_shape: tuple[int, ...] | None = None, + dtype: torch.dtype | None = None, + num_layers: int | None = None, + ) -> list[list[Any]]: + """Allocate and bind the model-owned paged KV cache.""" + + self._ensure_active() + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if any(supplied): + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + shape = tuple(int(dimension) for dimension in kv_cache_shape) + if len(shape) != 4: + raise ValueError(f"KV cache shape must have rank 4, got {shape}") + expected_layers = len(self.kv_cache_manager.per_layer_dtypes) + if int(num_layers) != expected_layers: + raise ValueError(f"vLLM KV layer count {num_layers} does not match model layer count {expected_layers}") + self.kv_cache_manager.validate_vllm_cache_spec( + block_size=shape[2], + dtype=dtype, + num_blocks=shape[0], + ) + if self.kv_cache_manager.config.num_blocks is None: + self.kv_cache_manager.configure(replace(self.kv_cache_manager.config, num_blocks=shape[0])) + self._refresh_page_table_layout() + elif self.kv_cache_manager.config.num_blocks != shape[0]: + raise ValueError( + f"Paged KV cache is resolved to {self.kv_cache_manager.config.num_blocks} blocks, not {shape[0]}" + ) + if any(tuple(expected) != shape for expected in self.kv_cache_manager.cache_shapes): + raise ValueError( + f"vLLM KV shape {shape} does not match model-derived shapes {self.kv_cache_manager.cache_shapes}" + ) + if not self.kv_cache_manager.config.is_resolved(): + raise RuntimeError("Paged KV cache capacity must be resolved before allocation") + self._seal_runtime_configuration() + return self.kv_cache_manager.allocate() + + def compile_prefill( + self, + *, + tokens: torch.Tensor, # ↓ Core request + page_table: torch.Tensor, + prompt_lens: torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> None: + """Compile prefill on the supplied eager or traced execution target.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._prefill_execution).compile_prefill( + tokens=tokens, + page_table=page_table, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + sampling_params=sampling_params, + ) + + def compile_decode( + self, + *, + tokens: torch.Tensor, # ↓ Core request + start_pos: torch.Tensor, + page_table: torch.Tensor, + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> None: + """Compile decode on the supplied eager or traced execution target.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._decode_execution).compile_decode( + tokens=tokens, + start_pos=start_pos, + page_table=page_table, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + prompt_lens: torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> Any: + """Validate ownership and execute one prefill call.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._prefill_execution).prefill_forward( + tokens=tokens, + page_table=page_table, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + sampling_params=sampling_params, + ) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> Any: + """Validate ownership and execute one decode call.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._decode_execution).decode_forward( + tokens=tokens, + start_pos=start_pos, + page_table=page_table, + sampling_params=sampling_params, + reset_batch=reset_batch, + read_from_device=read_from_device, + ) + + def can_trace_prefill( + self, + *, + tokens: torch.Tensor, # ↓ Core request + prompt_lens: torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + ) -> bool: + """Classify whether the prefill request can use this lane's trace.""" + + if self.traced_executor is None or not self.config.trace.prefill_enabled: + return False + return self.prefill_runtime.can_trace( + tokens=tokens, + prompt_lens=prompt_lens, + start_pos=start_pos, + ) + + def read_decode_output(self, tt_out: Any, *, async_read: bool = False) -> Any: + self._ensure_active() + return self.decode_runtime.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host(self, tt_out: Any, *, is_tokens: bool = False) -> tuple[Any, Any]: + self._ensure_active() + return self.decode_runtime.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + self._ensure_active() + if enable_trace: + self._warmup_q128_topk_tile_ends( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=True, + ) + self.warmup.warmup_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + if not enable_trace: + self._warmup_q128_topk_tile_ends( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=False, + ) + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + self._ensure_active() + return self.warmup.warmup_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self) -> None: + """Release runtime, trace, program, sampling, and KV resources in order.""" + + self._terminal = True + if self._cleaned_up: + return + + failures = [] + actions = [ + self.decode_runtime.drain_external_outputs, + self.output_reader.drain, + self.prefill_runtime.cleanup, + self.decode_runtime.cleanup_transients, + ] + if self.trace_compiler is not None: + actions.append(self.trace_compiler.cleanup) + actions.append(self.program_compiler.cleanup) + if self.config.device_sampling_enabled: + actions.append(self.model.sampling.release) + actions.append(self.kv_cache_manager.release) + + for action in actions: + try: + action() + except BaseException as error: + failures.append(error) + if failures: + _raise_cleanup_failures(failures, "Mistral7BExecutor") + self._cleaned_up = True + + # Private implementation + + def _warmup_q128_topk_tile_ends( + self, + *, + kv_cache: Any, + can_sample_on_device: bool, + enable_trace: bool, + ) -> None: + """Prime Mistral Q128 tile ends in eager and traced execution.""" + + if ( + enable_trace in self._q128_topk_tile_ends_warmed + or not can_sample_on_device + or (enable_trace and self.traced_executor is None) + or 128 not in self.warmup.config.prefill_sequence_lengths + or not self.prefill_runtime.config.static_q128_topk_supported + or self.warmup.config.prime_q128_tile_ends + ): + return + sampling = SamplingParams( + temperature=torch.ones(1), + top_k=torch.full((1,), 32, dtype=torch.int32), + top_p=torch.full((1,), 0.08), + ) + execution = self.traced_executor if enable_trace else self.eager_executor + for sequence_length in (32, 64, 96): + page_table_width = ( + sequence_length + self.page_table_layout.block_size - 1 + ) // self.page_table_layout.block_size + self.compile_prefill( + tokens=torch.zeros((1, sequence_length), dtype=torch.long), + page_table=torch.zeros((1, page_table_width), dtype=torch.int32), + prompt_lens=torch.full((1,), sequence_length, dtype=torch.long), + empty_slots=[0], + kv_cache=kv_cache, + sampling_params=sampling, + execution=execution, + ) + self._q128_topk_tile_ends_warmed.add(enable_trace) + + def _resolve_page_table_layout(self) -> PageTableLayout: + kv_config = self.kv_cache_manager.config + # The direct demo resolves num_blocks=max_num_blocks before constructing + # this executor, so its geometry is final immediately and it physically + # allocates that maximum. vLLM constructs against max_num_blocks only as + # a cheap capacity ceiling: no KV tensor is allocated until vLLM supplies + # num_blocks, _refresh_page_table_layout installs the final geometry, and + # allocate_kv_cache seals all runtime configs. + physical_num_blocks = kv_config.num_blocks or kv_config.max_num_blocks + return PageTableLayout.resolve( + block_size=int(kv_config.block_size), + model_max_sequence_length=int(self.model.config.max_seq_len), + physical_num_blocks=int(physical_num_blocks), + max_prefill_chunk_size=min( + int(self.runtime_config.max_prefill_chunk_size), + int(self.model.config.max_seq_len), + ), + ) + + def _refresh_page_table_layout(self) -> None: + # vLLM reaches this boundary after choosing the physical block count and + # before PagedKVCacheManager.allocate(), compilation, warmup, or tracing. + # Re-resolve complete immutable configs so their normal construction + # checks and capacity ceilings describe vLLM's authoritative geometry. + # Runtime and resource-owner identities remain unchanged. + layout = self._resolve_page_table_layout() + current_prefill = self.prefill_runtime.config + prefill_config = PrefillRuntimeConfig.resolve( + model=self.model, + output_reader=self.output_reader, + page_table_layout=layout, + max_batch_size=current_prefill.max_batch_size, + max_prefill_chunk_size=current_prefill.max_prefill_chunk_size, + device_sampling_enabled=current_prefill.device_sampling_enabled, + can_enable_trace=current_prefill.can_enable_trace, + supports_batched_prefill=current_prefill.supports_batched_prefill, + disable_batched_prefill=current_prefill.disable_batched_prefill, + max_prefill_batch_size=current_prefill.max_prefill_batch_size, + batched_prefill_batched_extract=current_prefill.batched_prefill_batched_extract, + ) + current_decode = self.decode_runtime.config + decode_config = DecodeRuntimeConfig.resolve( + model=self.model, + output_reader=self.output_reader, + lane_capacity=current_decode.lane_capacity, + page_table_layout=layout, + device_sampling_enabled=current_decode.device_sampling_enabled, + force_greedy_top_k=current_decode.force_greedy_top_k, + ) + warmup_config = WarmupCoordinatorConfig.resolve( + warmup=self.warmup.config.warmup, + trace=self.config.trace, + prefill=prefill_config, + decode=decode_config, + prefill_sequence_lengths=self.warmup.config.prefill_sequence_lengths, + ) + + self.prefill_runtime.config = prefill_config + self.decode_runtime.config = decode_config + self.warmup.config = warmup_config + self.page_table_layout = layout + + def _seal_runtime_configuration(self) -> None: + """Seal final geometry immediately before physical KV allocation.""" + + self.warmup.seal_configuration() + self._runtime_configuration_sealed = True + + def _ensure_sampling_for(self, sampling_params: Any) -> None: + if sampling_params is None: + return + if not self.config.device_sampling_enabled: + raise ValueError("sampling parameters were supplied while device sampling is disabled") + self._ensure_sampling_buffers() + + def _ensure_sampling_buffers(self) -> None: + if self._sampling_buffers_loaded: + return + if self.trace_compiler is not None and self.trace_compiler.trace_active: + raise RuntimeError("cannot materialize sampling buffers after trace activation") + self.model.sampling.load_device_buffers() + self._sampling_buffers_loaded = True + + def _validate_bound_cache(self, kv_cache: Any) -> None: + if self.kv_cache_manager.bound_context is None: + raise RuntimeError("Paged KV cache must be allocated and bound before execution") + if kv_cache is not None: + self.kv_cache_manager.validate_borrowed_handle(kv_cache) + + def _ensure_active(self) -> None: + if self._terminal: + raise RuntimeError("Mistral7BExecutor is terminal; construct a new executor") + if self.prefill_runtime.transient_orphan_count or self.decode_runtime.transient_orphan_count: + raise RuntimeError("Mistral7BExecutor has unreleased transient resources; clean up this executor") + + +def build_mistral_7b_executor(llm: Mistral7BForCausalLM, config: Mistral7BExecutorConfig) -> Mistral7BExecutor: + """Build one executor around an already-loaded Mistral model adapter.""" + + return Mistral7BExecutor(llm.model, llm.runtime_config, config) + + +def _raise_cleanup_failures(failures: list[BaseException], owner: str) -> None: + primary, *additional = failures + attach_cleanup_failures( + primary, + additional, + note=f"{owner} cleanup also encountered {{count}} failure(s)", + ) + raise primary diff --git a/code/models/common/models/mistral_7b/generator.py b/code/models/common/models/mistral_7b/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..0baa7bd8601b087fa999eb5e45ffec992aef7464 --- /dev/null +++ b/code/models/common/models/mistral_7b/generator.py @@ -0,0 +1,481 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Mistral-7B.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.mistral_7b.executor import Mistral7BExecutorConfig, build_mistral_7b_executor +from models.common.models.mistral_7b.hf_adaptor import from_pretrained +from models.common.models.mistral_7b.model import Mistral7BPagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Mistral7BGeneratorConfig: + """Validated construction inputs for one vLLM-facing Mistral generator.""" + + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + hf_revision: str | None = None + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Mistral7BGenerator: + """Adapt vLLM's model interface to the model-owned execution target. + + vLLM constructs this facade with `initialize_vllm_model`, resolves + KV capacity through `allocate_kv_cache`, warms the configured + programs, and then calls `prefill_forward` and + `decode_forward`. Each forward call is normalized by + `VLLMAdapter`, dispatched to the eager or traced executor, and + delegated to ``Mistral7BExecutor`` or ``LaneGroupExecutor``. + + This class owns dispatch policy but no TT resources. `cleanup` + delegates to the target that owns those resources. + """ + + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + + # Public vLLM API + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + """Return the unpadded per-submesh KV token budget for vLLM sizing.""" + + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + hf_revision: str | None = None, + device_sampling_enabled: bool = True, + ): + """Build the configured single-lane or data-parallel Mistral target.""" + + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_mistral_7b_generator( + Mistral7BGeneratorConfig( + hf_model=str(hf_model), + hf_revision=hf_revision, + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def compile_prefill( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + ) -> None: + """Normalize a vLLM prefill call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.compile_prefill(execution=execution, **normalized) + + def compile_decode( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + ) -> None: + """Normalize a vLLM decode call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.compile_decode(execution=execution, **normalized) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM prefill call.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM decode call.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward( + execution=execution, + read_from_device=read_from_device, + **normalized, + ) + + def read_decode_output( + self, + tt_out: Any, + *, + async_read: bool = False, + ) -> Any: + """Delegate vLLM's raw decode-output read.""" + + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host( + self, + tt_out: Any, + *, + is_tokens: bool = False, + ) -> tuple[Any, Any]: + """Delegate vLLM's asynchronous host-output completion.""" + + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self): + """Release every resource owned by the concrete target.""" + + return self.target.cleanup() + + # Private implementation + + def _select_prefill_execution( + self, + normalized: NormalizedPrefillKwargs, + trace_requested: bool, + ): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_mistral_7b_generator(config: Mistral7BGeneratorConfig) -> Mistral7BGenerator: + """Construct lane-local models/executors and compose their shared target surface.""" + + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Mistral7BPagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + hf_revision=config.hf_revision, + instruct="Instruct" in config.hf_model, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Mistral7BExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_mistral_7b_executor(llm, executor_config)) + + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Mistral7BGenerator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Mistral-7B model must contain at least one attention layer") + + attention_configs = tuple(layer.attention.config for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Mistral layer must expose the same KV head shape") + + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/mistral_7b/hf_adaptor.py b/code/models/common/models/mistral_7b/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..f60375c6c940367ccb4ed4c27333b056570e1e61 --- /dev/null +++ b/code/models/common/models/mistral_7b/hf_adaptor.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face provider boundary for Mistral-7B-Instruct-v0.3.""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.models.mistral_7b import weight_utils +from models.common.models.mistral_7b.model import ( + MISTRAL_ACCURACY, + MISTRAL_PERFORMANCE, + Mistral7B, + Mistral7BLayerWeights, + Mistral7BModelParameters, + Mistral7BPagedAttentionConfig, + Mistral7BPrecisionConfig, + Mistral7BWeights, + build_mistral_7b_transformer_config, +) + +DEFAULT_HF_MODEL = "mistralai/Mistral-7B-Instruct-v0.3" +DEFAULT_HF_REVISION = None + + +@dataclass(frozen=True) +class Mistral7BGenerationConfig: + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Mistral7BRuntimeConfig: + model_name: str + model_cache_path: Path | None + max_prefill_chunk_size: int + max_context_len: int + max_seq_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + del num_cached_tokens + return ( + prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + chat = [] + if isinstance(prompt_text, str): + if system_prompt_text: + chat.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + chat.append({"role": "user", "content": prompt_text}) + else: + chat = prompt_text + try: + return _chat_template_ids(tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True)) + except ValueError: + pass + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Mistral7BForCausalLM: + model: Mistral7B + tokenizer: Any + runtime_config: Mistral7BRuntimeConfig + instruct: bool = True + generation_config: Mistral7BGenerationConfig = field(default_factory=Mistral7BGenerationConfig) + + def __post_init__(self): + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + stops = tuple(getattr(self.tokenizer, "stop_tokens", ()) or ()) + self.generation_config = Mistral7BGenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=stops, + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + return encode_prompt( + self.tokenizer, + prompt_text, + system_prompt_text, + instruct=self.instruct if instruct is None else instruct, + ) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_tokenizer(hf_model: str, hf_revision: str | None = DEFAULT_HF_REVISION): + tokenizer = AutoTokenizer.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=os.getenv("CI") == "true", + ) + eos = getattr(tokenizer, "eos_token_id", None) + tokenizer.stop_tokens = [] if eos is None else ([eos] if isinstance(eos, int) else list(eos)) + return tokenizer + + +def _trace_seq_lens(num_devices: int, max_prefill_chunk_size: int, max_seq_len: int) -> tuple[int, ...]: + allowed = {1: (128,), 2: (128, 1024), 8: (128, 1024)}.get(num_devices, (128,)) + return tuple(length for length in allowed if length <= min(max_prefill_chunk_size, max_seq_len)) + + +def _cache_path(hf_model: str, mesh_device, cache_dir: Path | str | None) -> Path: + if cache_dir is not None: + path = Path(cache_dir) + elif os.getenv("TT_CACHE_PATH"): + path = Path(os.environ["TT_CACHE_PATH"]) + else: + topology = {1: "N150", 2: "N300", 8: "T3K"}.get( + mesh_device.get_num_devices(), f"TP{mesh_device.get_num_devices()}" + ) + path = Path("model_cache") / hf_model / topology + path.mkdir(parents=True, exist_ok=True) + return path + + +def _validate_checkpoint_config(hf_config) -> None: + if hf_config.hidden_size % hf_config.num_attention_heads: + raise ValueError("Mistral hidden_size must be divisible by num_attention_heads") + rope_parameters = getattr(hf_config, "rope_parameters", None) or {} + rope_theta = getattr(hf_config, "rope_theta", None) + if rope_theta is None: + rope_theta = rope_parameters.get("rope_theta", 1_000_000.0) + rope_type = rope_parameters.get("rope_type", "default") + if float(rope_theta) != 1_000_000.0 or rope_type != "default": + raise ValueError("Mistral-7B-Instruct-v0.3 requires plain RoPE theta=1,000,000") + if getattr(hf_config, "sliding_window", None) is not None: + raise ValueError("Mistral-7B-Instruct-v0.3 requires full attention (sliding_window=None)") + if bool(getattr(hf_config, "attention_bias", False)): + raise ValueError("Mistral-7B-Instruct-v0.3 does not use QKV projection bias") + + +def convert_hf_model_weights( + hf, + *, + n_layers: int, + num_devices: int, + rope_table_len: int, + head_dim: int, +) -> Mistral7BWeights: + """Extract and convert all Hugging Face tensors consumed by the TT builder.""" + + base = hf.model + rope_cos, rope_sin = weight_utils.build_rope_cos_sin_torch( + base.rotary_emb, + rope_table_len, + head_dim, + torch.bfloat16, + ) + layers = [] + for layer in base.layers[:n_layers]: + attention = layer.self_attn + if any(getattr(attention, name, None) is not None for name in ("q_norm", "k_norm")): + raise ValueError("Mistral-7B-Instruct-v0.3 does not use QK norm") + if any( + getattr(projection, "bias", None) is not None + for projection in (attention.q_proj, attention.k_proj, attention.v_proj) + ): + raise ValueError("Mistral-7B-Instruct-v0.3 does not use QKV projection bias") + wqkv, wo = weight_utils.attention_wqkv_wo_from_hf_layer(attention, num_devices) + w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(layer.mlp) + layers.append( + Mistral7BLayerWeights( + wqkv=wqkv, + wo=wo, + w1=w1, + w2=w2, + w3=w3, + attention_norm=weight_utils.rms_weight_torch(layer.input_layernorm).to(torch.bfloat16), + ff_norm=weight_utils.rms_weight_torch(layer.post_attention_layernorm).to(torch.bfloat16), + ) + ) + return Mistral7BWeights( + embedding=weight_utils.embed_tokens_torch(base.embed_tokens), + rope_cos=rope_cos, + rope_sin=rope_sin, + layers=tuple(layers), + final_norm=weight_utils.rms_weight_torch(base.norm).to(torch.bfloat16), + lm_head=hf.lm_head.weight.detach().to(torch.bfloat16).clone(), + ) + + +def from_pretrained( + mesh_device, + *, + hf_model: str = DEFAULT_HF_MODEL, + hf_revision: str | None = DEFAULT_HF_REVISION, + instruct: bool = True, + max_batch_size: int = 32, + max_seq_len: int = 4096, + optimizations: str | Mistral7BPrecisionConfig = "accuracy", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config: Mistral7BPagedAttentionConfig | None = None, + cache_dir: Path | str | None = None, +) -> Mistral7BForCausalLM: + del dtype + ttnn.SetDefaultDevice(mesh_device) + hf_config = AutoConfig.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=os.getenv("CI") == "true", + ) + _validate_checkpoint_config(hf_config) + num_devices = mesh_device.get_num_devices() + if hf_config.num_attention_heads % num_devices or hf_config.num_key_value_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({hf_config.num_attention_heads}/{hf_config.num_key_value_heads}) " + f"must be divisible by device count ({num_devices})" + ) + hf = AutoModelForCausalLM.from_pretrained( + hf_model, + revision=hf_revision, + torch_dtype=torch.bfloat16, + local_files_only=os.getenv("CI") == "true", + ) + hf.eval() + resolved_layers = hf_config.num_hidden_layers if n_layers is None else n_layers + if ( + not isinstance(resolved_layers, int) + or isinstance(resolved_layers, bool) + or not 0 < resolved_layers <= hf_config.num_hidden_layers + ): + raise ValueError(f"n_layers must be in [1, {hf_config.num_hidden_layers}]") + precision = ( + optimizations + if isinstance(optimizations, Mistral7BPrecisionConfig) + else (MISTRAL_PERFORMANCE if optimizations == "performance" else MISTRAL_ACCURACY) + ) + if not isinstance(precision, Mistral7BPrecisionConfig) or ( + isinstance(optimizations, str) and optimizations not in ("accuracy", "performance") + ): + raise TypeError("optimizations must be 'accuracy', 'performance', or Mistral7BPrecisionConfig") + + cache_path = _cache_path(hf_model, mesh_device, cache_dir) + if paged_attention_config is None: + block_size = 32 + paged_attention_config = Mistral7BPagedAttentionConfig( + block_size=block_size, + max_num_blocks=((max_seq_len + block_size - 1) // block_size) * max_batch_size, + ) + head_dim = hf_config.hidden_size // hf_config.num_attention_heads + params = Mistral7BModelParameters( + dim=hf_config.hidden_size, + n_heads=hf_config.num_attention_heads, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + hidden_dim=hf_config.intermediate_size, + vocab_size=hf_config.vocab_size, + rms_norm_eps=hf_config.rms_norm_eps, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + ) + rope_table_len = math.ceil(max(max_seq_len * 2, 8192) / 128) * 128 + weights = convert_hf_model_weights( + hf, + n_layers=resolved_layers, + num_devices=num_devices, + rope_table_len=rope_table_len, + head_dim=head_dim, + ) + model_config = build_mistral_7b_transformer_config( + mesh_device=mesh_device, + params=params, + weights=weights, + n_layers=resolved_layers, + precision=precision, + cache_path=cache_path, + paged_attention_config=paged_attention_config, + ) + tokenizer = load_tokenizer(hf_model, hf_revision) + model = Mistral7B(model_config) + max_prefill_chunk_size = 2048 + runtime_config = Mistral7BRuntimeConfig( + model_name=Path(hf_model).name, + model_cache_path=cache_path, + max_prefill_chunk_size=max_prefill_chunk_size, + max_context_len=int(hf_config.max_position_embeddings), + max_seq_len=max_seq_len, + trace_prefill_supported_seq_lens=_trace_seq_lens(num_devices, max_prefill_chunk_size, max_seq_len), + max_prefill_batch_size=8 if num_devices == 1 else 32, + disable_batched_prefill=bool(os.getenv("DISABLE_BATCHED_PREFILL")), + batched_prefill_batched_extract=not bool(os.getenv("DISABLE_BATCHED_EXTRACT")), + ) + del hf + return Mistral7BForCausalLM( + model=model, + tokenizer=tokenizer, + runtime_config=runtime_config, + instruct=instruct, + ) diff --git a/code/models/common/models/mistral_7b/model.py b/code/models/common/models/mistral_7b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..52d430686868417fa8f82b42a0005bd2eb7b4529 --- /dev/null +++ b/code/models/common/models/mistral_7b/model.py @@ -0,0 +1,944 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Mistral-7B-Instruct-v0.3 — native TTTv2 stack (``Embedding1D``, ``RMSNorm1D``, +``Attention1D``, ``MLP1D``, ``RotarySetup1D``, ``LMHead1D``). + +Tensor layout contracts: + - **Prefill** hidden states: ``[1, 1, S, dim]`` TILE, ``S % 128 == 0``. + - **Decode** hidden states: ``[1, 1, B, dim]`` TILE (``B`` padded to tile in modules). + +The tensor graph is provider-neutral. ``hf_adaptor.py`` owns Hugging Face loading and +builds :class:`Mistral7BTransformerConfig`; the model exposes the strict duck-typed +surface consumed by ``models.common.llm_runtime``. +""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.models.mistral_7b import weight_utils +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D +from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + """Minimal LazyWeight; ``Attention1D`` / ``MLP1D`` / ``Embedding1D`` resolvers fill mesh + memory.""" + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +@dataclass +class Mistral7BPagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (``block_size`` / ``max_num_blocks`` only).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Mistral7BTransformerConfig: + """Complete provider-neutral tensor-graph configuration.""" + + n_layers: int + vocab_size: int + max_batch_size: int + max_seq_len: int + dim: int + num_devices: int + mesh_device: ttnn.MeshDevice + embedding_config: Embedding1DConfig + rope_config: Rope1DConfig + block_configs: list[Mistral7BDecoderLayerConfig] + norm_config: RMSNorm1DConfig + lm_head_config: LMHead1DConfig + sampling_config: Sampling1DConfig | None = None + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtypes: list[ttnn.DataType | None] = field(default_factory=list) + tt_ccl: Any = None + cache_path: str | None = None + + +@dataclass(frozen=True) +class Mistral7BModelParameters: + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + max_batch_size: int + max_seq_len: int + + +@dataclass(frozen=True) +class Mistral7BLayerWeights: + wqkv: torch.Tensor + wo: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Mistral7BWeights: + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Mistral7BLayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +@dataclass(frozen=True) +class Mistral7BPrecisionConfig: + """Per-layer precision + math-fidelity recipe for Mistral-7B-Instruct-v0.3. + + Mirrors the fields TTTv1's ``DecodersPrecision`` actually distinguishes for Mistral-7B + (Llama-family group in ``model_config.py:130-159`` for ``accuracy()``, ``:208-218`` for + ``performance()``). Two module-level recipes are exposed: :data:`MISTRAL_ACCURACY` and + :data:`MISTRAL_PERFORMANCE`. Pass one to :meth:`Mistral7BForCausalLM.from_pretrained` via + ``precision=``; use ``dataclasses.replace(MISTRAL_ACCURACY, lm_head_dtype=...)`` to + customize a single field. + + Attention compute-kernel configs (LI_QKV_*, LI_O_*, SDPA_*) are intentionally absent: + TTTv1 leaves them at engine defaults for Mistral-7B (HIFI2 QKV/O, HIFI4 SDPA prefill, + HIFI2 SDPA decode), which coincide with ``Attention1D``'s TTTv2 defaults. Add fields + here only when a future variant overrides them. + """ + + wqkv_dtype: ttnn.DataType = ttnn.bfloat8_b + wo_dtype: ttnn.DataType = ttnn.bfloat8_b + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat8_b + mlp_w2_dtype: ttnn.DataType = ttnn.bfloat8_b + # ``None`` ⇒ MLP1D resolves to HIFI2_FP16 (matches TTTv1 ``LI_FF1_FF3`` / ``LI_FF2`` for accuracy). + mlp_ff1_3_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + mlp_ff2_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + # Not part of TTTv1's ``DecodersPrecision`` — TTTv2 accuracy mode tightens this to bf16 + # to hit PERF.md top-1 thresholds. + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + + +_LOFI_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, +) + + +# TTTv1 ``DecodersPrecision.accuracy("Mistral-7B-Instruct-v0.3")`` (Llama-family group at +# ``model_config.py:130-159``) keeps BFP8 attention + BFP8 KV cache + HIFI2_FP16 MLP; only +# the LM head is tightened to bf16 (TTTv2 addition; required for top-1 in this stack). +MISTRAL_ACCURACY = Mistral7BPrecisionConfig( + lm_head_dtype=ttnn.bfloat16, +) + +# TTTv1 ``DecodersPrecision.performance("Mistral-7B-Instruct-v0.3")`` (``model_config.py:208-218``): +# FF1_FF3 → BFP4 and LI_FF1_FF3 → LOFI. Everything else matches the accuracy base; LM head +# stays BFP8 in perf mode. This single delta is the bulk of TTTv1's perf-mode throughput uplift. +MISTRAL_PERFORMANCE = Mistral7BPrecisionConfig( + mlp_w1_w3_dtype=ttnn.bfloat4_b, + mlp_ff1_3_compute_kernel_cfg=_LOFI_COMPUTE_KERNEL_CFG, +) + + +def _slice_last_token_tile(x: ttnn.Tensor, last_token_idx: int) -> ttnn.Tensor: + """Slice the 32-row tile containing ``last_token_idx`` from ``[1, 1, S, W]``.""" + floor = (last_token_idx // 32) * 32 + return ttnn.slice(x, (0, 0, floor, 0), (1, 1, floor + 32, x.shape[-1])) + + +def _post_attn_norm_decode_configs( + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding so its output matches MLP1D's W1/W3 input. + + Same gotcha as Qwen: MLP1D decode uses ``_dram_shard_core_grid_k_n(dim, padded_hidden / num_devices)`` + for W1/W3 inputs, but the default RMSNorm program config is derived from ``_compute_norm_core_grid(dim)`` + alone — that DRAM-width-shard mismatch silently corrupts decode activations. + """ + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + memory_config = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return program_config, memory_config + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, x: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + return x + + if memory_config is None: + memory_config = x.memory_config() + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=tt_ccl.get_num_links(), + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + # CCL tuning: match the house default (CCL_CHUNKS_PER_SYNC / CCL_NUM_WORKERS_PER_LINK in + # tt_ccl.py = 10 / 2) used by every shared module (Attention1D/MLP1D/RMSNorm1D) and the + # same-architecture llama3_8b reference (32 heads / 8 KV). The original port shipped a + # divergent 24 / 4 here. These two per-layer rmsnorm all-gathers run 64x/decode step, so on + # the T3K 8-device ring their per-op cost is non-trivial; aligning to the validated house + # default is correctness-neutral (all-gather result is identical) and removes an unexplained + # divergence. (The earlier "~17% T3K decode gap vs TTTv1" this was suspected to drive turned + # out to be a cross-box/cross-harness measurement artifact, not a real regression: on equal + # footing -- same box, on-device sampling -- T3K decode is at parity with TTTv1.) + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + + +@dataclass +class _Mistral7BWHTuning: + """Wormhole-specific MLP tuning resolved at build time. + + Mistral 7B (TTTv1 ``DecodersPrecision.accuracy``) keeps HIFI2_FP16 / BFP8 across + attention + FF, so we do **not** override the ``Attention1D`` / ``MLP1D`` math-fidelity + defaults. Only the MLP prefill_len_cutoff is set here, matching TTTv1's WH cutoff for + 7B-class Llama-family models (512 on N150; 1024 on multi-device WH). + """ + + mlp_prefill_len_cutoff: int | None = None + mlp_decode_spill_w1_to_dram: bool = False + # Use ttnn.experimental.minimal_matmul for QKV + W2 prefill matmuls above seq_len > 128 (TTTv1 + # parity, PLAN_01). A/B escape hatch: set DISABLE_MINIMAL_MATMUL=1 to force ttnn.linear. On a 7B + # the prefill matmuls are large, so minimal_matmul is a real prefill-TTFT win (unlike tiny 1B). + prefill_minimal_matmul: bool = True + + +def _resolve_mistral_wh_tuning(*, num_dev: int, max_batch_size: int) -> _Mistral7BWHTuning: + """Pick WH tuning knobs for Mistral-7B-Instruct-v0.3 on N150 / N300. + + TTTv1 ([model_config.py:516-589](../../../tt_transformers/tt/model_config.py#L516-L589)) sets + ``prefill_len_cutoff=1024`` on WH by default and reduces it to 512 on N150 for Llama 3.1-8B / + Mistral-7B / Phi-4. We mirror that here. + """ + t = _Mistral7BWHTuning() + t.mlp_prefill_len_cutoff = 512 if num_dev == 1 else 1024 + # Decode W1→DRAM spill: leave off by default; promote only if N150 batch decode trips L1. + t.mlp_decode_spill_w1_to_dram = False + t.prefill_minimal_matmul = not os.environ.get("DISABLE_MINIMAL_MATMUL") + logger.info( + f"MLP tuning for Mistral-7B on {num_dev} device(s): " + f"prefill_len_cutoff={t.mlp_prefill_len_cutoff}, " + f"decode_spill_w1_to_dram={t.mlp_decode_spill_w1_to_dram}, " + f"prefill_minimal_matmul={t.prefill_minimal_matmul}" + ) + return t + + +@dataclass +class Mistral7BDecoderLayerConfig: + attention_norm_config: RMSNorm1DConfig + attention_config: Attention1DConfig + ff_norm_config: RMSNorm1DConfig + mlp_config: MLP1DConfig + + +def _build_decoder_layer( + *, + idx: int, + weights: Mistral7BLayerWeights, + mcfg: Mistral7BModelParameters, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Mistral7BPrecisionConfig, + paged_cfg: Mistral7BPagedAttentionConfig, + cache_path: Path | None, + wh: _Mistral7BWHTuning, +) -> Mistral7BDecoderLayerConfig: + """Build one decoder-layer config from provider-converted tensors.""" + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, + dtype=precision.wqkv_dtype, + cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None, + ) + lazy_wo = _lazy( + weights.wo, + dtype=precision.wo_dtype, + cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None, + ) + + attention_config = Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=mcfg.n_heads, + n_kv_heads=mcfg.n_kv_heads, + head_dim=mcfg.head_dim, + max_batch_size=mcfg.max_batch_size, + max_seq_len=mcfg.max_seq_len, + q_norm_config=None, + k_norm_config=None, + wqkv_bias=None, + use_vllm_paged_kv_cache=True, + paged_attention_config=paged_cfg, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + prefill_qkv_minimal_matmul=wh.prefill_minimal_matmul, + ) + + mlp_config = MLP1DConfig( + w1=_lazy( + weights.w1, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None, + ), + w2=_lazy( + weights.w2, + dtype=precision.mlp_w2_dtype, + cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None, + ), + w3=_lazy( + weights.w3, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None, + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=mcfg.max_batch_size, + prefill_len_cutoff=wh.mlp_prefill_len_cutoff, + decode_spill_w1_to_dram_before_w3=wh.mlp_decode_spill_w1_to_dram, + w1_w3_dtype=precision.mlp_w1_w3_dtype, + w2_dtype=precision.mlp_w2_dtype, + ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + ff2_compute_kernel_cfg=precision.mlp_ff2_compute_kernel_cfg, + decode_ff2_compute_kernel_cfg=precision.mlp_ff2_compute_kernel_cfg, + prefill_w2_minimal_matmul=wh.prefill_minimal_matmul, + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + dim=mcfg.dim, + hidden_dim=mcfg.hidden_dim, + num_devices=num_dev, + max_batch_size=mcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1DConfig: + lw = _lazy( + weight, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=mcfg.rms_norm_eps, + max_batch_size=mcfg.max_batch_size, + tt_ccl=tt_ccl, + **extra, + ) + + return Mistral7BDecoderLayerConfig( + attention_norm_config=_build_norm(weights.attention_norm, "pre_attn"), + attention_config=attention_config, + ff_norm_config=_build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ), + mlp_config=mlp_config, + ) + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + mcfg: Mistral7BModelParameters, + lm_head_dtype: ttnn.DataType, + cache_path: Path | None, +) -> LMHead1DConfig: + """Build the vocab-sharded LM head with DRAM-matmul program configs.""" + lm_splits, lm_split_sizes, lm_weights_memcfgs = weight_utils.build_lm_head_lazy_weights( + mesh_device, + lm_head_weight, + dim=mcfg.dim, + vocab_size=mcfg.vocab_size, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(mcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(mcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, mcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(mcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=mcfg.dim, + max_batch_size=mcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + compute_kernel_config=None, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + ) + + +def build_mistral_7b_transformer_config( + *, + mesh_device: ttnn.MeshDevice, + params: Mistral7BModelParameters, + weights: Mistral7BWeights, + n_layers: int, + precision: Mistral7BPrecisionConfig, + cache_path: Path | None, + paged_attention_config: Mistral7BPagedAttentionConfig, +) -> Mistral7BTransformerConfig: + """Build the TT tensor graph from provider-neutral dimensions and converted tensors.""" + + num_devices = mesh_device.get_num_devices() + if params.n_heads % num_devices or params.n_kv_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({params.n_heads}/{params.n_kv_heads}) " + f"must be divisible by device count ({num_devices})" + ) + if len(weights.layers) != n_layers: + raise ValueError(f"Expected {n_layers} decoder layer weight sets, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) if num_devices > 1 else None + topology = default_topology(mesh_device) + wh = _resolve_mistral_wh_tuning(num_dev=num_devices, max_batch_size=params.max_batch_size) + embedding_config = Embedding1DConfig( + weights=_lazy( + weights.embedding, + dtype=ttnn.bfloat16, + cache=(cache_path / "embedding", "tok_embeddings") if cache_path else None, + ), + mesh_device=mesh_device, + embed_scale=1.0, + ) + rope_config = Rope1DConfig( + cos_matrix=_lazy( + weights.rope_cos, + dtype=ttnn.bfloat16, + cache=(cache_path / "rope", "cos") if cache_path else None, + ), + sin_matrix=_lazy( + weights.rope_sin, + dtype=ttnn.bfloat16, + cache=(cache_path / "rope", "sin") if cache_path else None, + ), + max_batch_size=params.max_batch_size, + head_dim=params.head_dim, + device=mesh_device, + use_qk_fused=False, + ) + block_configs = [ + _build_decoder_layer( + idx=index, + weights=weights.layers[index], + mcfg=params, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + paged_cfg=paged_attention_config, + cache_path=cache_path, + wh=wh, + ) + for index in range(n_layers) + ] + norm_config = RMSNorm1DConfig( + weight=_lazy( + weights.final_norm, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", "final") if cache_path else None, + ), + mesh_device=mesh_device, + eps=params.rms_norm_eps, + max_batch_size=params.max_batch_size, + tt_ccl=tt_ccl, + ) + lm_head_config = _build_lm_head( + mesh_device=mesh_device, + lm_head_weight=weights.lm_head, + mcfg=params, + lm_head_dtype=precision.lm_head_dtype, + cache_path=cache_path, + ) + sampling_config = Sampling1DConfig( + vocab_size=params.vocab_size, + valid_vocab_size=params.vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=_nearest_32(params.max_batch_size), + allow_force_argmax=False, + pad_to_power_of_2=True, + ) + return Mistral7BTransformerConfig( + n_layers=n_layers, + vocab_size=params.vocab_size, + max_batch_size=params.max_batch_size, + max_seq_len=params.max_seq_len, + dim=params.dim, + num_devices=num_devices, + mesh_device=mesh_device, + embedding_config=embedding_config, + rope_config=rope_config, + block_configs=block_configs, + norm_config=norm_config, + lm_head_config=lm_head_config, + sampling_config=sampling_config, + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtypes=[None] * n_layers, + tt_ccl=tt_ccl, + cache_path=str(cache_path) if cache_path is not None else None, + ) + + +class Mistral7BDecoderLayer(LightweightModule): + def __init__( + self, + *, + input_layernorm: RMSNorm1D, + self_attn: Attention1D, + post_attention_layernorm: RMSNorm1D, + mlp: MLP1D, + ): + super().__init__() + self.input_layernorm = input_layernorm + self.self_attn = self_attn + self.post_attention_layernorm = post_attention_layernorm + self.mlp = mlp + self.attention_norm = input_layernorm + self.attention = self_attn + self.ff_norm = post_attention_layernorm + self.feed_forward = mlp + + @classmethod + def from_config(cls, config: Mistral7BDecoderLayerConfig) -> Mistral7BDecoderLayer: + return cls( + input_layernorm=RMSNorm1D.from_config(config.attention_norm_config), + self_attn=Attention1D.from_config(config.attention_config), + post_attention_layernorm=RMSNorm1D.from_config(config.ff_norm_config), + mlp=MLP1D.from_config(config.mlp_config), + ) + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + *, + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + # Fractured embed/norm activations must be all-gathered to full ``dim`` before + # Attention1D / MLP1D (QKV matmul expects width ``dim``). + r = self.input_layernorm.prefill_forward(x) + r = _all_gather_rmsnorm_tensor(self.input_layernorm, r) + r = self.self_attn.prefill_forward( + r, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + r2 = self.post_attention_layernorm.prefill_forward(h) + r2 = _all_gather_rmsnorm_tensor(self.post_attention_layernorm, r2) + r2 = self.mlp.prefill_forward(r2) + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def decode_forward( + self, + x: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + xa = _all_gather_rmsnorm_tensor( + self.input_layernorm, x, memory_config=self.input_layernorm.config.decode_memory_config + ) + r = self.input_layernorm.forward(xa, "decode") + r = self.self_attn.forward(r, current_pos, rot_mats, mode="decode", page_table=page_table) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + hf = _all_gather_rmsnorm_tensor( + self.post_attention_layernorm, h, memory_config=self.post_attention_layernorm.config.decode_memory_config + ) + r2 = self.post_attention_layernorm.forward(hf, "decode") + r2 = self.mlp.forward(r2, "decode") + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + +class Mistral7B(LightweightModule): + """ + Full decoder for Mistral-7B-Instruct-v0.3 (TTTv2 modules only). + + Prefill/decode on **embedded** activations are orchestrated by ``Mistral7BExecutor``. Token + embedding is ``embed_prefill`` / ``embed_decode``. Bind KV with ``set_kv_cache`` before first forward. + """ + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__(self, config: Mistral7BTransformerConfig): + from tqdm import tqdm + + super().__init__() + self.config = config + self.embed = Embedding1D.from_config(config.embedding_config) + self.embedding = self.embed + self.rope_setup = RotarySetup1D.from_config(config.rope_config) + self.layers = [ + Mistral7BDecoderLayer.from_config(config.block_configs[index]) + for index in tqdm(range(config.n_layers), desc="Building layers") + ] + self.norm = RMSNorm1D.from_config(config.norm_config) + self.lm_head = LMHead1D.from_config(config.lm_head_config) + self.sampling = Sampling1D.from_config(config.sampling_config) if config.sampling_config is not None else None + self.supports_on_device_sampling = self.sampling is not None + self.mesh_device = config.mesh_device + self.tt_ccl = config.tt_ccl or (get_tt_ccl(config.mesh_device) if config.num_devices > 1 else None) + self.vocab_size = config.vocab_size + self.n_layers = config.n_layers + self.num_devices = config.num_devices + self.decode_residual_memcfg = config.decode_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.prefill_residual_memcfg = config.prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtypes = config.activation_dtypes or [None] * config.n_layers + self.model_args = None + + @property + def n_kv_heads(self) -> int: + return self.config.block_configs[0].attention_config.n_kv_heads + + def iter_executor_named_modules(self): + if not hasattr(self, "layers"): + return + for index, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{index}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in self.config.block_configs) + for config in tuple({id(item): item for item in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, + block_size=block_size, + max_num_blocks=max_num_blocks, + ) + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for index, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {index} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {index} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = pair + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). The batched path always + # returns the full hidden state when no runtime last-token slice is supplied; the executor does + # per-slot last-token extraction + norm/lm_head so those stages stay bit-identical to the single-user path. + x = x_embed + for layer in self.layers: + x = layer.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + + if last_token_index is not None and last_token_slice is None: + raise ValueError("last_token_index is required with a runtime last_token_slice") + if get_last_token == -1 and last_token_slice is None: + return x + + old = x + if last_token_slice is None: + x_tile = _slice_last_token_tile(old, get_last_token) + else: + x_tile = ttnn.slice( + old, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(old.shape[2]) // 32, + ) + ttnn.deallocate(old) + if last_token_index is not None: + if x_tile.dtype != ttnn.bfloat16: + old_tile = x_tile + x_tile = ttnn.typecast(x_tile, ttnn.bfloat16) + ttnn.deallocate(old_tile) + old_tile = x_tile + x_tile = ttnn.embedding(last_token_index, x_tile, layout=ttnn.TILE_LAYOUT) + x_tile = ttnn.unsqueeze_to_4D(x_tile) + ttnn.deallocate(old_tile) + return self._last_tile_logits(x_tile) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + x = _slice_last_token_tile(hidden_states, last_token_idx) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + return self._last_tile_logits(x) + + def _last_tile_logits(self, x_tile: ttnn.Tensor) -> ttnn.Tensor: + """Final-norm + all-gather + LM-head on a 32-row tile. ``x_tile`` shape ``[1, 1, 32, dim]``.""" + x = self.norm.prefill_forward(x_tile) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + del last_token_slice, last_token_index + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), + ) + x = ttnn.matmul( + selector, + hidden_states, + dtype=ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + compute_kernel_config=ttnn.init_device_compute_kernel_config( + self.mesh_device.arch(), + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ), + ) + ttnn.deallocate(selector) + return self._last_tile_logits(x) + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for layer in self.layers: + x = layer.decode_forward(x, current_pos, rot_mats, page_table=page_table) + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + return self.lm_head.forward(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + # House-default CCL tuning (10 / 2), matching llama3_8b. Only used in host + # sampling mode (full-vocab gather); the original 24 / 4 is the suspected reason + # T3K host mode was the worst TTTv2 mode. See _all_gather_rmsnorm_tensor for rationale. + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) diff --git a/code/models/common/models/mistral_7b/weight_utils.py b/code/models/common/models/mistral_7b/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..86cbd7d16079849d7eec5b80ce8afe747d91e0bc --- /dev/null +++ b/code/models/common/models/mistral_7b/weight_utils.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side HF weight layout helpers for Mistral-7B-Instruct-v0.3 (TTTv2 port). + +Q/K linear weights use HuggingFace head layout; TTNN RoPE expects Meta layout. +This duplicates the small permute helpers from ``load_checkpoints.reverse_permute`` +without importing ``models/tt_transformers``. + +Mistral 7B v0.3 has no QKV bias and no Q/K norm — those Qwen-only branches are +intentionally absent here. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.modules.lazy_weight import LazyWeight + + +def reverse_permute(tensor: torch.Tensor, n_heads: int, dim1: int, dim2: int) -> torch.Tensor: + """HF Q/K weight rows → Meta layout (RoPE-compatible).""" + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def permute_hf_rope_to_meta_tables(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert HF cos/sin tables to Meta-style tables for ``RotarySetup1D``.""" + if len(cos.shape) == 3: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + return cos, sin + + +def build_rope_cos_sin_torch( + rotary_emb: Any, table_len: int, head_dim: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Build ``[1, 1, table_len, head_dim]`` cos/sin tensors (Meta layout) from HF rotary module.""" + x = torch.zeros(1, 1, table_len, head_dim, dtype=dtype) + position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0) + with torch.no_grad(): + cos_hf, sin_hf = rotary_emb(x, position_ids) + cos_m, sin_m = permute_hf_rope_to_meta_tables(cos_hf.float(), sin_hf.float()) + return cos_m.to(dtype), sin_m.to(dtype) + + +def attention_wqkv_wo_from_hf_layer( + self_attn: Any, + num_devices: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack QKV and WO from an HF Mistral attention block for ``Attention1D`` sharding.""" + wq_raw = self_attn.q_proj.weight + wk_raw = self_attn.k_proj.weight + wv_raw = self_attn.v_proj.weight + wo_raw = self_attn.o_proj.weight + + dim = wq_raw.shape[1] + cfg = self_attn.config + n_heads = cfg.num_attention_heads + n_kv_heads = cfg.num_key_value_heads + head_dim = cfg.hidden_size // n_heads + n_heads_times_head_dim = n_heads * head_dim + n_kv_heads_times_head_dim = n_kv_heads * head_dim + + wq_meta = reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim) + wk_meta = reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim) + wv_meta = wv_raw + wo_meta = wo_raw + + wq = wq_meta.T + wk = wk_meta.T + wv = wv_meta.T + wo = wo_meta.T + + qkv_list = [] + for i in range(num_devices): + wq_chunk = torch.chunk(wq, num_devices, dim=1)[i] + wk_chunk = torch.chunk(wk, num_devices, dim=1)[i] + wv_chunk = torch.chunk(wv, num_devices, dim=1)[i] + qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1) + qkv_list.append(qkv) + wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0).clone() + wo = wo.unsqueeze(0).unsqueeze(0).clone() + return wqkv, wo + + +def mlp_weights_from_hf_layer(mlp: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (w1, w2, w3) TTNN layouts: gate^T, down^T, up^T for ``MLP1D``.""" + w1 = mlp.gate_proj.weight.T.contiguous().clone() + w3 = mlp.up_proj.weight.T.contiguous().clone() + w2 = mlp.down_proj.weight.T.contiguous().clone() + return w1, w2, w3 + + +def rms_weight_torch(layernorm: Any) -> torch.Tensor: + return layernorm.weight.detach().float().clone() + + +def embed_tokens_torch(embed: Any) -> torch.Tensor: + w = embed.weight.detach().float().clone() + return w.unsqueeze(0).unsqueeze(0) + + +def build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """Column-split LM head for ``LMHead1D`` (DRAM-sharded matmul chunks per device). + + ``lm_head_weight`` shape: ``[vocab_size, dim]`` (HF). Returns + ``(lazy_weights, split_sizes, weights_memcfgs)`` for ``LMHead1DConfig``. + """ + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / 32) * 32 + if vocab_size < padded_vocab_size: + pad = padded_vocab_size - vocab_size + torch_w = torch.cat([torch_w, torch.zeros(torch_w.shape[0], pad, dtype=torch_w.dtype)], dim=-1) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + tile = ttnn.TILE_SIZE + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (tile * dram_size.x)) * (tile * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights: list[LazyWeight] = [] + weights_memcfgs: list[ttnn.MemoryConfig] = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_splits.append(torch_w[:, start:end]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + name = f"lm_head_split_{i}_{combined.shape[-1]}" + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, name) if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs diff --git a/code/models/common/models/phi4/README.md b/code/models/common/models/phi4/README.md new file mode 100644 index 0000000000000000000000000000000000000000..dfd61b3e7d996e7b918c4bc5e58d23c39585b3b1 --- /dev/null +++ b/code/models/common/models/phi4/README.md @@ -0,0 +1,84 @@ +# Phi-4 with TTTv2 + +This directory is the model-owned TTTv2 path for Microsoft Phi-4. + +It intentionally demonstrates direct executor construction from +`models/common/llm_runtime`. It is not part of the Llama/Qwen executor +consolidation and does not use `models/common/models/executor.py`. + +## Product path + +```text +Hugging Face checkpoint + -> hf_adaptor.py: provider metadata, tokenizer, and weight conversion + -> model.py: Phi-4 tensor graph composed from TTTv2 modules + -> executor.py: direct composition of common runtime owners for one lane + -> generator.py: vLLM construction, DP composition, and dispatch +``` + +## Files + +| File | Responsibility | +| --- | --- | +| `hf_adaptor.py` | Resolve the Phi-4 provider configuration/tokenizer and build runtime metadata | +| `weight_utils.py` | Convert and map HF weights into the model-owned layout | +| `model.py` | Build and execute the TTTv2 Phi-4 transformer graph | +| `executor.py` | Directly compose one execution lane and own its runtime resources | +| `generator.py` | Build lanes, configure `VLLMAdapter`, and expose serving dispatch | + +## Tensor-module composition + +`model.py` composes: + +- `Embedding1D` +- `RotarySetup1D` +- `RMSNorm1D` +- `Attention1D` +- `MLP1D` +- `LMHead1D` +- optional `Sampling1D` +- common TT collective helpers + +Phi-specific provider mapping, architecture values, precision, and tuning +remain model-owned. + +## Direct executor composition + +`Phi4Executor` directly constructs: + +```text +Phi4Transformer +├── PagedKVCacheManager +├── OutputReader +├── PrefillRuntime +├── DecodeRuntime +├── ProgramCompiler +├── EagerExecutor +├── optional TraceCompiler +├── optional TracedExecutor +└── WarmupCoordinator +``` + +This direct construction path is intentionally retained as an example for new +models whose orchestration does not belong in the shared Llama or Qwen family +composition. It reuses focused runtime modules without adding a protocol, +profile, or executor subclass hierarchy. + +The executor owns the physical KV cache, compile/trace artifacts, output +leases, sampling resources, and cleanup. Cleanup is terminal, ordered, +retryable, and idempotent. + +## vLLM and data parallelism + +`Phi4Generator` builds one executor per lane. DP1 uses the executor directly; +larger DP configurations use `LaneGroupExecutor`. `VLLMAdapter` performs +server-boundary normalization and KV-cache validation. + +## Tests + +Relevant entry points include: + +- `models/common/tests/models/phi4/test_hf_adaptor.py` +- `models/common/tests/models/phi4/test_demo_contract.py` +- `models/common/tests/demos/phi4/demo.py` +- `models/common/tests/llm_runtime/test_executor_integration.py` diff --git a/code/models/common/models/phi4/__init__.py b/code/models/common/models/phi4/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..332d63f4c12feb7966a02400397a12c269113a12 --- /dev/null +++ b/code/models/common/models/phi4/__init__.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +from models.common.models.phi4.executor import Phi4Executor, Phi4ExecutorConfig, build_phi4_executor +from models.common.models.phi4.generator import Phi4Generator, Phi4GeneratorConfig, build_phi4_generator +from models.common.models.phi4.hf_adaptor import ( + DEFAULT_HF_MODEL, + DEFAULT_HF_REVISION, + Phi4ForCausalLM, + Phi4RuntimeConfig, +) +from models.common.models.phi4.model import ( + PHI4_ACCURACY, + PHI4_PERFORMANCE, + Phi4PagedAttentionConfig, + Phi4PrecisionConfig, + Phi4Transformer, + Phi4TransformerConfig, + build_phi4_transformer_config, +) + +__all__ = [ + "DEFAULT_HF_MODEL", + "DEFAULT_HF_REVISION", + "PHI4_ACCURACY", + "PHI4_PERFORMANCE", + "Phi4Transformer", + "Phi4TransformerConfig", + "Phi4PagedAttentionConfig", + "Phi4PrecisionConfig", + "Phi4ForCausalLM", + "Phi4RuntimeConfig", + "Phi4Executor", + "Phi4ExecutorConfig", + "Phi4Generator", + "Phi4GeneratorConfig", + "build_phi4_executor", + "build_phi4_generator", + "build_phi4_transformer_config", +] diff --git a/code/models/common/models/phi4/executor.py b/code/models/common/models/phi4/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..0a3c44928e7ac69cd3c7418b4acb88b432d8d975 --- /dev/null +++ b/code/models/common/models/phi4/executor.py @@ -0,0 +1,613 @@ +# SPDX-FileCopyrightText: 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Model-owned Phi-4 execution composition and cleanup root.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any, Sequence + +import torch + +from models.common.llm_runtime.config import PagedKVCacheConfig, PageTableLayout, TraceConfig, WarmupConfig +from models.common.llm_runtime.decode import DecodeRuntime, DecodeRuntimeConfig +from models.common.llm_runtime.execution import EagerExecutor, TracedExecutor +from models.common.llm_runtime.output_reader import OutputReader +from models.common.llm_runtime.paged_kv_cache import PagedKVCacheManager +from models.common.llm_runtime.prefill.config import PrefillRuntimeConfig +from models.common.llm_runtime.prefill.runtime import PrefillRuntime +from models.common.llm_runtime.program_compiler import ProgramCompiler +from models.common.llm_runtime.tensor_resources import attach_cleanup_failures +from models.common.llm_runtime.trace_compiler import TraceCompiler +from models.common.llm_runtime.warmup import WarmupCoordinator, WarmupCoordinatorConfig +from models.common.models.phi4.hf_adaptor import Phi4ForCausalLM +from models.common.modules.sampling.sampling_1d import Sampling1D +from models.common.sampling import SamplingParams + + +@dataclass(frozen=True) +class Phi4ExecutorConfig: + """Immutable aggregate policy paired with one model-owned executor.""" + + trace: TraceConfig + warmup: WarmupConfig + paged_kv_cache: PagedKVCacheConfig + device_sampling_enabled: bool + + def __post_init__(self) -> None: + nested_configs = ( + ("trace", self.trace, TraceConfig), + ("warmup", self.warmup, WarmupConfig), + ("paged_kv_cache", self.paged_kv_cache, PagedKVCacheConfig), + ) + for name, value, expected_type in nested_configs: + if type(value) is not expected_type: + raise TypeError(f"{name} must be exactly {expected_type.__name__}") + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + + +class Phi4Executor: + """Compose every runtime owner for one Phi-4 execution lane. + + Construction wires one `PrefillRuntime`, `DecodeRuntime`, + `ProgramCompiler`, and `EagerExecutor`. Trace-enabled + configurations add one `TraceCompiler` and one + `TracedExecutor` over that exact eager instance. + + ``Phi4Generator`` or ``LaneGroupExecutor`` first configures and allocates + KV cache, then warms or compiles programs, and finally calls + `prefill_forward` and `decode_forward` with an explicit + execution target. `cleanup` is the deterministic resource-release + root and makes the executor terminal. + """ + + requires_prefill_trace_warmup = True + + def __init__(self, model: Any, runtime_config: Any, config: Phi4ExecutorConfig) -> None: + if not isinstance(config, Phi4ExecutorConfig): + raise TypeError("config must be a Phi4ExecutorConfig") + iter_modules = getattr(model, "iter_executor_named_modules", None) + if not callable(iter_modules): + raise TypeError("model must provide iter_executor_named_modules()") + can_enable_trace = getattr(runtime_config, "can_enable_trace", None) + if not callable(can_enable_trace): + raise TypeError("runtime_config must provide can_enable_trace()") + model_config = getattr(model, "config", None) + mesh_device = getattr(model_config, "mesh_device", None) + if mesh_device is None: + raise ValueError("model.config.mesh_device is required") + + self.model = model + self.runtime_config = runtime_config + self.model_args = runtime_config + self.config = config + self.mesh_device = mesh_device + self.cache_path = getattr(runtime_config, "model_cache_path", None) + self._terminal = False + self._cleaned_up = False + self._sampling_buffers_loaded = False + self._runtime_configuration_sealed = False + self._q128_topk_tile_ends_warmed: set[bool] = set() + + sampling = getattr(model, "sampling", None) + if config.device_sampling_enabled: + if not isinstance(sampling, Sampling1D): + raise TypeError("device sampling requires model.sampling to be a Sampling1D") + is_resolved = getattr(getattr(sampling, "config", None), "is_resolved", None) + if not callable(is_resolved) or not is_resolved(): + raise ValueError("model.sampling must have a resolved Sampling1DConfig") + + self.kv_cache_manager = PagedKVCacheManager(model, config.paged_kv_cache) + self.page_table_layout = self._resolve_page_table_layout() + self.output_reader = OutputReader(mesh_device) + self.prefill_runtime = PrefillRuntime( + PrefillRuntimeConfig.resolve( + model=model, + output_reader=self.output_reader, + page_table_layout=self.page_table_layout, + max_batch_size=int(model.config.max_batch_size), + max_prefill_chunk_size=int(runtime_config.max_prefill_chunk_size), + device_sampling_enabled=config.device_sampling_enabled, + can_enable_trace=runtime_config.can_enable_trace, + supports_batched_prefill=bool(runtime_config.supports_batched_prefill), + disable_batched_prefill=bool(runtime_config.disable_batched_prefill), + max_prefill_batch_size=int(runtime_config.max_prefill_batch_size), + batched_prefill_batched_extract=bool(runtime_config.batched_prefill_batched_extract), + ) + ) + self.decode_runtime = DecodeRuntime( + DecodeRuntimeConfig.resolve( + model=model, + output_reader=self.output_reader, + lane_capacity=int(model.config.max_batch_size), + page_table_layout=self.page_table_layout, + device_sampling_enabled=config.device_sampling_enabled, + force_greedy_top_k=config.warmup.include_decode_top_k, + ) + ) + self.program_compiler = ProgramCompiler(mesh_device, lambda: self.kv_cache_manager.bound_context) + self.eager_executor = EagerExecutor( + prefill=self.prefill_runtime, + decode=self.decode_runtime, + program_compiler=self.program_compiler, + ) + self.trace_compiler: TraceCompiler | None = None + self.traced_executor: TracedExecutor | None = None + if config.trace.mode != "none": + self.trace_compiler = TraceCompiler(self.program_compiler) + self.traced_executor = TracedExecutor( + eager=self.eager_executor, + trace_compiler=self.trace_compiler, + trace_mode=config.trace.mode, + ) + self.eager_execution = self.eager_executor + self.traced_prefill_execution = ( + self.traced_executor if config.trace.prefill_enabled and self.traced_executor is not None else None + ) + self.traced_decode_execution = ( + self.traced_executor if config.trace.decode_enabled and self.traced_executor is not None else None + ) + self._prefill_execution = self.traced_prefill_execution or self.eager_executor + self._decode_execution = self.traced_decode_execution or self.eager_executor + + prefill_sequence_lengths = getattr(runtime_config, "trace_prefill_supported_seq_lens", (128,)) + self.warmup = WarmupCoordinator( + config=WarmupCoordinatorConfig.resolve( + warmup=config.warmup, + trace=config.trace, + prefill=self.prefill_runtime.config, + decode=self.decode_runtime.config, + prefill_sequence_lengths=prefill_sequence_lengths, + ), + execution=self.traced_executor or self.eager_executor, + ensure_sampling_buffers=self._ensure_sampling_buffers, + validate_bound_cache=self._validate_bound_cache, + ) + + # Public model execution API + + @property + def model_config(self): + return self.model.config + + @property + def cluster_shape(self) -> list[int]: + return list(self.mesh_device.shape) + + @property + def paged_kv_cache_config(self) -> PagedKVCacheConfig: + return self.kv_cache_manager.config + + @property + def terminal(self) -> bool: + return self._terminal + + @property + def already_warmed_up_prefill(self) -> bool: + return self.warmup.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value: bool) -> None: + if not isinstance(value, bool): + raise TypeError("already_warmed_up_prefill compatibility value must be bool") + # Compatibility writes are intentionally non-authoritative. Warmup + # coverage is derived from the coordinator's completed case ledger. + return None + + def configure_paged_kv_cache(self, config: PagedKVCacheConfig) -> None: + """Resolve vLLM-owned KV geometry before the first allocation.""" + + self._ensure_active() + if self._runtime_configuration_sealed: + raise RuntimeError("runtime configuration is sealed") + if not isinstance(config, PagedKVCacheConfig): + raise TypeError("config must be a PagedKVCacheConfig") + if self.kv_cache_manager.config.is_resolved(): + raise RuntimeError("paged KV cache configuration is already resolved") + if config.dtype != self.kv_cache_manager.config.dtype: + raise ValueError("resolved paged KV cache cannot change dtype") + if config.memory_config != self.kv_cache_manager.config.memory_config: + raise ValueError("resolved paged KV cache cannot change memory_config") + self.model.configure_paged_attention( + block_size=config.block_size, + max_num_blocks=config.max_num_blocks, + ) + self.kv_cache_manager.configure(config) + self.config = replace(self.config, paged_kv_cache=config) + self._refresh_page_table_layout() + + def allocate_kv_cache( + self, + kv_cache_shape: tuple[int, ...] | None = None, + dtype: torch.dtype | None = None, + num_layers: int | None = None, + ) -> list[list[Any]]: + """Allocate and bind the model-owned paged KV cache.""" + + self._ensure_active() + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if any(supplied): + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + shape = tuple(int(dimension) for dimension in kv_cache_shape) + if len(shape) != 4: + raise ValueError(f"KV cache shape must have rank 4, got {shape}") + expected_layers = len(self.kv_cache_manager.per_layer_dtypes) + if int(num_layers) != expected_layers: + raise ValueError(f"vLLM KV layer count {num_layers} does not match model layer count {expected_layers}") + self.kv_cache_manager.validate_vllm_cache_spec( + block_size=shape[2], + dtype=dtype, + num_blocks=shape[0], + ) + if self.kv_cache_manager.config.num_blocks is None: + self.kv_cache_manager.configure(replace(self.kv_cache_manager.config, num_blocks=shape[0])) + self._refresh_page_table_layout() + elif self.kv_cache_manager.config.num_blocks != shape[0]: + raise ValueError( + f"Paged KV cache is resolved to {self.kv_cache_manager.config.num_blocks} blocks, not {shape[0]}" + ) + if any(tuple(expected) != shape for expected in self.kv_cache_manager.cache_shapes): + raise ValueError( + f"vLLM KV shape {shape} does not match model-derived shapes {self.kv_cache_manager.cache_shapes}" + ) + if not self.kv_cache_manager.config.is_resolved(): + raise RuntimeError("Paged KV cache capacity must be resolved before allocation") + self._seal_runtime_configuration() + return self.kv_cache_manager.allocate() + + def compile_prefill( + self, + *, + tokens: torch.Tensor, # ↓ Core request + page_table: torch.Tensor, + prompt_lens: torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> None: + """Compile prefill on the supplied eager or traced execution target.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._prefill_execution).compile_prefill( + tokens=tokens, + page_table=page_table, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + sampling_params=sampling_params, + ) + + def compile_decode( + self, + *, + tokens: torch.Tensor, # ↓ Core request + start_pos: torch.Tensor, + page_table: torch.Tensor, + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> None: + """Compile decode on the supplied eager or traced execution target.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._decode_execution).compile_decode( + tokens=tokens, + start_pos=start_pos, + page_table=page_table, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + prompt_lens: torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> Any: + """Validate ownership and execute one prefill call.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._prefill_execution).prefill_forward( + tokens=tokens, + page_table=page_table, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + sampling_params=sampling_params, + ) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + execution: EagerExecutor | TracedExecutor | None = None, # ↓ Internal dispatch + ) -> Any: + """Validate ownership and execute one decode call.""" + + self._ensure_active() + self._validate_bound_cache(kv_cache) + self._ensure_sampling_for(sampling_params) + return (execution or self._decode_execution).decode_forward( + tokens=tokens, + start_pos=start_pos, + page_table=page_table, + sampling_params=sampling_params, + reset_batch=reset_batch, + read_from_device=read_from_device, + ) + + def can_trace_prefill( + self, + *, + tokens: torch.Tensor, # ↓ Core request + prompt_lens: torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + ) -> bool: + """Classify whether the prefill request can use this lane's trace.""" + + if self.traced_executor is None or not self.config.trace.prefill_enabled: + return False + return self.prefill_runtime.can_trace( + tokens=tokens, + prompt_lens=prompt_lens, + start_pos=start_pos, + ) + + def read_decode_output(self, tt_out: Any, *, async_read: bool = False) -> Any: + self._ensure_active() + return self.decode_runtime.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host(self, tt_out: Any, *, is_tokens: bool = False) -> tuple[Any, Any]: + self._ensure_active() + return self.decode_runtime.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + self._ensure_active() + if enable_trace: + self._warmup_q128_topk_tile_ends( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=True, + ) + self.warmup.warmup_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + if not enable_trace: + self._warmup_q128_topk_tile_ends( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=False, + ) + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + self._ensure_active() + return self.warmup.warmup_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self) -> None: + """Release runtime, trace, program, sampling, and KV resources in order.""" + + self._terminal = True + if self._cleaned_up: + return + + failures = [] + actions = [ + self.decode_runtime.drain_external_outputs, + self.output_reader.drain, + self.prefill_runtime.cleanup, + self.decode_runtime.cleanup_transients, + ] + if self.trace_compiler is not None: + actions.append(self.trace_compiler.cleanup) + actions.append(self.program_compiler.cleanup) + if self.config.device_sampling_enabled: + actions.append(self.model.sampling.release) + actions.append(self.kv_cache_manager.release) + + for action in actions: + try: + action() + except BaseException as error: + failures.append(error) + if failures: + _raise_cleanup_failures(failures, "Phi4Executor") + self._cleaned_up = True + + # Private implementation + + def _warmup_q128_topk_tile_ends( + self, + *, + kv_cache: Any, + can_sample_on_device: bool, + enable_trace: bool, + ) -> None: + """Prime Phi-4 Q128 tile ends in eager and traced execution.""" + + if ( + enable_trace in self._q128_topk_tile_ends_warmed + or not can_sample_on_device + or (enable_trace and self.traced_executor is None) + or 128 not in self.warmup.config.prefill_sequence_lengths + or not self.prefill_runtime.config.static_q128_topk_supported + or self.warmup.config.prime_q128_tile_ends + ): + return + sampling = SamplingParams( + temperature=torch.ones(1), + top_k=torch.full((1,), 32, dtype=torch.int32), + top_p=torch.full((1,), 0.08), + ) + execution = self.traced_executor if enable_trace else self.eager_executor + for sequence_length in (32, 64, 96): + page_table_width = ( + sequence_length + self.page_table_layout.block_size - 1 + ) // self.page_table_layout.block_size + self.compile_prefill( + tokens=torch.zeros((1, sequence_length), dtype=torch.long), + page_table=torch.zeros((1, page_table_width), dtype=torch.int32), + prompt_lens=torch.full((1,), sequence_length, dtype=torch.long), + empty_slots=[0], + kv_cache=kv_cache, + sampling_params=sampling, + execution=execution, + ) + self._q128_topk_tile_ends_warmed.add(enable_trace) + + def _resolve_page_table_layout(self) -> PageTableLayout: + kv_config = self.kv_cache_manager.config + # The direct demo resolves num_blocks=max_num_blocks before constructing + # this executor, so its geometry is final immediately and it physically + # allocates that maximum. vLLM constructs against max_num_blocks only as + # a cheap capacity ceiling: no KV tensor is allocated until vLLM supplies + # num_blocks, _refresh_page_table_layout installs the final geometry, and + # allocate_kv_cache seals all runtime configs. + physical_num_blocks = kv_config.num_blocks or kv_config.max_num_blocks + return PageTableLayout.resolve( + block_size=int(kv_config.block_size), + model_max_sequence_length=int(self.model.config.max_seq_len), + physical_num_blocks=int(physical_num_blocks), + max_prefill_chunk_size=min( + int(self.runtime_config.max_prefill_chunk_size), + int(self.model.config.max_seq_len), + ), + ) + + def _refresh_page_table_layout(self) -> None: + # vLLM reaches this boundary after choosing the physical block count and + # before PagedKVCacheManager.allocate(), compilation, warmup, or tracing. + # Re-resolve complete immutable configs so their normal construction + # checks and capacity ceilings describe vLLM's authoritative geometry. + # Runtime and resource-owner identities remain unchanged. + layout = self._resolve_page_table_layout() + current_prefill = self.prefill_runtime.config + prefill_config = PrefillRuntimeConfig.resolve( + model=self.model, + output_reader=self.output_reader, + page_table_layout=layout, + max_batch_size=current_prefill.max_batch_size, + max_prefill_chunk_size=current_prefill.max_prefill_chunk_size, + device_sampling_enabled=current_prefill.device_sampling_enabled, + can_enable_trace=current_prefill.can_enable_trace, + supports_batched_prefill=current_prefill.supports_batched_prefill, + disable_batched_prefill=current_prefill.disable_batched_prefill, + max_prefill_batch_size=current_prefill.max_prefill_batch_size, + batched_prefill_batched_extract=current_prefill.batched_prefill_batched_extract, + ) + current_decode = self.decode_runtime.config + decode_config = DecodeRuntimeConfig.resolve( + model=self.model, + output_reader=self.output_reader, + lane_capacity=current_decode.lane_capacity, + page_table_layout=layout, + device_sampling_enabled=current_decode.device_sampling_enabled, + force_greedy_top_k=current_decode.force_greedy_top_k, + ) + warmup_config = WarmupCoordinatorConfig.resolve( + warmup=self.warmup.config.warmup, + trace=self.config.trace, + prefill=prefill_config, + decode=decode_config, + prefill_sequence_lengths=self.warmup.config.prefill_sequence_lengths, + ) + + self.prefill_runtime.config = prefill_config + self.decode_runtime.config = decode_config + self.warmup.config = warmup_config + self.page_table_layout = layout + + def _seal_runtime_configuration(self) -> None: + """Seal final geometry immediately before physical KV allocation.""" + + self.warmup.seal_configuration() + self._runtime_configuration_sealed = True + + def _ensure_sampling_for(self, sampling_params: Any) -> None: + if sampling_params is None: + return + if not self.config.device_sampling_enabled: + raise ValueError("sampling parameters were supplied while device sampling is disabled") + self._ensure_sampling_buffers() + + def _ensure_sampling_buffers(self) -> None: + if self._sampling_buffers_loaded: + return + if self.trace_compiler is not None and self.trace_compiler.trace_active: + raise RuntimeError("cannot materialize sampling buffers after trace activation") + self.model.sampling.load_device_buffers() + self._sampling_buffers_loaded = True + + def _validate_bound_cache(self, kv_cache: Any) -> None: + if self.kv_cache_manager.bound_context is None: + raise RuntimeError("Paged KV cache must be allocated and bound before execution") + if kv_cache is not None: + self.kv_cache_manager.validate_borrowed_handle(kv_cache) + + def _ensure_active(self) -> None: + if self._terminal: + raise RuntimeError("Phi4Executor is terminal; construct a new executor") + if self.prefill_runtime.transient_orphan_count or self.decode_runtime.transient_orphan_count: + raise RuntimeError("Phi4Executor has unreleased transient resources; clean up this executor") + + +def build_phi4_executor(llm: Phi4ForCausalLM, config: Phi4ExecutorConfig) -> Phi4Executor: + """Build one executor around an already-loaded Phi model adapter.""" + + return Phi4Executor(llm.model, llm.runtime_config, config) + + +def _raise_cleanup_failures(failures: list[BaseException], owner: str) -> None: + primary, *additional = failures + attach_cleanup_failures( + primary, + additional, + note=f"{owner} cleanup also encountered {{count}} failure(s)", + ) + raise primary diff --git a/code/models/common/models/phi4/generator.py b/code/models/common/models/phi4/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..06c0888afdb610d02496d3a2a365b25c3c2def67 --- /dev/null +++ b/code/models/common/models/phi4/generator.py @@ -0,0 +1,481 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Phi-4.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.phi4.executor import Phi4ExecutorConfig, build_phi4_executor +from models.common.models.phi4.hf_adaptor import DEFAULT_HF_REVISION, from_pretrained +from models.common.models.phi4.model import Phi4PagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Phi4GeneratorConfig: + """Validated construction inputs for one vLLM-facing Phi generator.""" + + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + hf_revision: str | None = DEFAULT_HF_REVISION + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Phi4Generator: + """Adapt vLLM's model interface to the model-owned execution target. + + vLLM constructs this facade with `initialize_vllm_model`, resolves + KV capacity through `allocate_kv_cache`, warms the configured + programs, and then calls `prefill_forward` and + `decode_forward`. Each forward call is normalized by + `VLLMAdapter`, dispatched to the eager or traced executor, and + delegated to ``Phi4Executor`` or ``LaneGroupExecutor``. + + This class owns dispatch policy but no TT resources. `cleanup` + delegates to the target that owns those resources. + """ + + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + + # Public vLLM API + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + """Return the unpadded per-submesh KV token budget for vLLM sizing.""" + + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + hf_revision: str | None = DEFAULT_HF_REVISION, + device_sampling_enabled: bool = True, + ): + """Build the configured single-lane or data-parallel Phi target.""" + + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_phi4_generator( + Phi4GeneratorConfig( + hf_model=str(hf_model), + hf_revision=hf_revision, + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def compile_prefill( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + ) -> None: + """Normalize a vLLM prefill call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.compile_prefill(execution=execution, **normalized) + + def compile_decode( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + ) -> None: + """Normalize a vLLM decode call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.compile_decode(execution=execution, **normalized) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM prefill call.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM decode call.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward( + execution=execution, + read_from_device=read_from_device, + **normalized, + ) + + def read_decode_output( + self, + tt_out: Any, + *, + async_read: bool = False, + ) -> Any: + """Delegate vLLM's raw decode-output read.""" + + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host( + self, + tt_out: Any, + *, + is_tokens: bool = False, + ) -> tuple[Any, Any]: + """Delegate vLLM's asynchronous host-output completion.""" + + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self): + """Release every resource owned by the concrete target.""" + + return self.target.cleanup() + + # Private implementation + + def _select_prefill_execution( + self, + normalized: NormalizedPrefillKwargs, + trace_requested: bool, + ): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_phi4_generator(config: Phi4GeneratorConfig) -> Phi4Generator: + """Construct lane-local models/executors and compose their shared target surface.""" + + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Phi4PagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + hf_revision=config.hf_revision, + instruct=True, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Phi4ExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_phi4_executor(llm, executor_config)) + + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Phi4Generator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Phi-4 model must contain at least one attention layer") + + attention_configs = tuple(layer.attention.config for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Phi layer must expose the same KV head shape") + + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/phi4/hf_adaptor.py b/code/models/common/models/phi4/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..6b44e8fedaac339a646ce306567abeb797ff9e90 --- /dev/null +++ b/code/models/common/models/phi4/hf_adaptor.py @@ -0,0 +1,332 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face provider boundary for ``microsoft/phi-4``.""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.models.phi4 import weight_utils +from models.common.models.phi4.model import ( + PHI4_ACCURACY, + PHI4_PERFORMANCE, + Phi4LayerWeights, + Phi4ModelParameters, + Phi4PagedAttentionConfig, + Phi4PrecisionConfig, + Phi4Transformer, + Phi4Weights, + build_phi4_transformer_config, +) + +DEFAULT_HF_MODEL = "microsoft/phi-4" +DEFAULT_HF_REVISION = "187ef0342fff0eb3333be9f00389385e95ef0b61" + + +@dataclass(frozen=True) +class Phi4GenerationConfig: + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Phi4RuntimeConfig: + model_name: str + model_cache_path: Path | None + max_prefill_chunk_size: int + max_context_len: int + max_seq_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 8 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + del num_cached_tokens + return ( + prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + messages = [] + if isinstance(prompt_text, str): + if system_prompt_text: + messages.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + messages.append({"role": "user", "content": prompt_text}) + else: + messages = prompt_text + try: + return _chat_template_ids( + tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True) + ) + except ValueError: + pass + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Phi4ForCausalLM: + model: Phi4Transformer + tokenizer: Any + runtime_config: Phi4RuntimeConfig + instruct: bool = True + generation_config: Phi4GenerationConfig = field(default_factory=Phi4GenerationConfig) + + def __post_init__(self): + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + self.generation_config = Phi4GenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=tuple(getattr(self.tokenizer, "stop_tokens", ()) or ()), + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + return encode_prompt( + self.tokenizer, + prompt_text, + system_prompt_text, + instruct=self.instruct if instruct is None else instruct, + ) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_tokenizer(hf_model: str, hf_revision: str | None = DEFAULT_HF_REVISION): + tokenizer = AutoTokenizer.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=os.getenv("CI") == "true", + ) + eos = getattr(tokenizer, "eos_token_id", None) + stops = [] if eos is None else ([eos] if isinstance(eos, int) else list(eos)) + im_end = tokenizer.convert_tokens_to_ids("<|im_end|>") + if isinstance(im_end, int) and im_end >= 0 and im_end not in stops: + stops.append(im_end) + tokenizer.stop_tokens = stops + return tokenizer + + +def _trace_seq_lens(num_devices: int, max_prefill_chunk_size: int, max_seq_len: int) -> tuple[int, ...]: + if num_devices != 2: + raise ValueError(f"Phi-4 supports logical TP2 lanes only, got {num_devices} devices") + return tuple(length for length in (128, 1024) if length <= min(max_prefill_chunk_size, max_seq_len)) + + +def _cache_path(hf_model: str, mesh_device, cache_dir: Path | str | None) -> Path: + if cache_dir is not None: + path = Path(cache_dir) + elif os.getenv("TT_CACHE_PATH"): + path = Path(os.environ["TT_CACHE_PATH"]) + else: + path = Path("model_cache") / hf_model / "N300" + path.mkdir(parents=True, exist_ok=True) + return path + + +def _validate_checkpoint_config(hf_config) -> None: + if hf_config.hidden_size % hf_config.num_attention_heads: + raise ValueError("Phi-4 hidden_size must be divisible by num_attention_heads") + rope_parameters = getattr(hf_config, "rope_parameters", None) or getattr(hf_config, "rope_scaling", None) or {} + rope_theta = getattr(hf_config, "rope_theta", None) or rope_parameters.get("rope_theta", 250_000.0) + partial_rotary_factor = getattr(hf_config, "partial_rotary_factor", None) + if partial_rotary_factor is None: + partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0) + if float(rope_theta) != 250_000.0: + raise ValueError("Phi-4 requires RoPE theta=250,000") + if float(partial_rotary_factor) != 1.0: + raise ValueError("Phi-4 requires full-head RoPE (partial_rotary_factor=1.0)") + if rope_parameters.get("rope_type", "default") != "default": + raise ValueError("Phi-4 requires plain, unscaled RoPE tables") + if bool(getattr(hf_config, "attention_bias", False)): + raise ValueError("Phi-4 does not use QKV projection bias") + + +def convert_hf_model_weights( + hf, + hf_config, + *, + n_layers: int, + num_devices: int, + rope_table_len: int, + head_dim: int, +) -> Phi4Weights: + """Extract fused Phi-4 projections and convert them to provider-neutral tensors.""" + + base = hf.model + rope_cos, rope_sin = weight_utils.build_rope_cos_sin_torch( + base.rotary_emb, rope_table_len, head_dim, torch.bfloat16 + ) + layers = [] + for layer in base.layers[:n_layers]: + attention = layer.self_attn + if any(getattr(attention, name, None) is not None for name in ("q_norm", "k_norm")): + raise ValueError("Phi-4 does not use QK norm") + if getattr(attention.qkv_proj, "bias", None) is not None: + raise ValueError("Phi-4 does not use QKV projection bias") + wqkv, wo = weight_utils.attention_wqkv_wo_from_hf_layer(attention, num_devices) + w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(layer.mlp) + layers.append( + Phi4LayerWeights( + wqkv=wqkv, + wo=wo, + w1=w1, + w2=w2, + w3=w3, + attention_norm=weight_utils.rms_weight_torch(layer.input_layernorm).to(torch.bfloat16), + ff_norm=weight_utils.rms_weight_torch(layer.post_attention_layernorm).to(torch.bfloat16), + ) + ) + if bool(getattr(hf_config, "tie_word_embeddings", False)): + raise ValueError("Phi-4 requires an untied LM head") + return Phi4Weights( + embedding=weight_utils.embed_tokens_torch(base.embed_tokens), + rope_cos=rope_cos, + rope_sin=rope_sin, + layers=tuple(layers), + final_norm=weight_utils.rms_weight_torch(base.norm).to(torch.bfloat16), + lm_head=hf.lm_head.weight.detach().to(torch.bfloat16).clone(), + ) + + +def from_pretrained( + mesh_device, + *, + hf_model: str = DEFAULT_HF_MODEL, + hf_revision: str | None = DEFAULT_HF_REVISION, + instruct: bool = True, + max_batch_size: int = 32, + max_seq_len: int = 4096, + optimizations: str | Phi4PrecisionConfig = "accuracy", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config: Phi4PagedAttentionConfig | None = None, + cache_dir: Path | str | None = None, +) -> Phi4ForCausalLM: + del dtype + num_devices = mesh_device.get_num_devices() + if num_devices != 2: + raise ValueError(f"Phi-4 supports logical TP2 lanes only, got {num_devices} devices") + ttnn.SetDefaultDevice(mesh_device) + load_kwargs = dict(revision=hf_revision, local_files_only=os.getenv("CI") == "true") + hf_config = AutoConfig.from_pretrained(hf_model, **load_kwargs) + _validate_checkpoint_config(hf_config) + hf = AutoModelForCausalLM.from_pretrained(hf_model, torch_dtype=torch.bfloat16, **load_kwargs) + hf.eval() + resolved_layers = hf_config.num_hidden_layers if n_layers is None else n_layers + if ( + not isinstance(resolved_layers, int) + or isinstance(resolved_layers, bool) + or not 0 < resolved_layers <= hf_config.num_hidden_layers + ): + raise ValueError(f"n_layers must be in [1, {hf_config.num_hidden_layers}]") + precision = ( + optimizations + if isinstance(optimizations, Phi4PrecisionConfig) + else (PHI4_PERFORMANCE if optimizations == "performance" else PHI4_ACCURACY) + ) + if not isinstance(precision, Phi4PrecisionConfig) or ( + isinstance(optimizations, str) and optimizations not in ("accuracy", "performance") + ): + raise TypeError("optimizations must be 'accuracy', 'performance', or Phi4PrecisionConfig") + cache_path = _cache_path(hf_model, mesh_device, cache_dir) + if paged_attention_config is None: + block_size = 32 + paged_attention_config = Phi4PagedAttentionConfig( + block_size=block_size, + max_num_blocks=((max_seq_len + block_size - 1) // block_size) * max_batch_size, + ) + head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads) + params = Phi4ModelParameters( + dim=hf_config.hidden_size, + n_heads=hf_config.num_attention_heads, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + hidden_dim=hf_config.intermediate_size, + vocab_size=hf_config.vocab_size, + rms_norm_eps=hf_config.rms_norm_eps, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + ) + rope_table_len = math.ceil(max(max_seq_len * 2, 8192) / 128) * 128 + weights = convert_hf_model_weights( + hf, + hf_config, + n_layers=resolved_layers, + num_devices=num_devices, + rope_table_len=rope_table_len, + head_dim=head_dim, + ) + config = build_phi4_transformer_config( + mesh_device=mesh_device, + params=params, + weights=weights, + n_layers=resolved_layers, + precision=precision, + cache_path=cache_path, + paged_attention_config=paged_attention_config, + ) + tokenizer = load_tokenizer(hf_model, hf_revision) + model = Phi4Transformer(config) + max_prefill_chunk_size = 2048 + runtime_config = Phi4RuntimeConfig( + model_name=Path(hf_model).name, + model_cache_path=cache_path, + max_prefill_chunk_size=max_prefill_chunk_size, + max_context_len=int(hf_config.max_position_embeddings), + max_seq_len=max_seq_len, + trace_prefill_supported_seq_lens=_trace_seq_lens(num_devices, max_prefill_chunk_size, max_seq_len), + disable_batched_prefill=bool(os.getenv("DISABLE_BATCHED_PREFILL")), + batched_prefill_batched_extract=not bool(os.getenv("DISABLE_BATCHED_EXTRACT")), + ) + del hf + return Phi4ForCausalLM(model=model, tokenizer=tokenizer, runtime_config=runtime_config, instruct=instruct) diff --git a/code/models/common/models/phi4/model.py b/code/models/common/models/phi4/model.py new file mode 100644 index 0000000000000000000000000000000000000000..4f7115b937a05fb0b2aaa0c537afd0a1d6a782b6 --- /dev/null +++ b/code/models/common/models/phi4/model.py @@ -0,0 +1,945 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Phi-4 (``microsoft/phi-4``) — native TTTv2 stack (``Embedding1D``, ``RMSNorm1D``, +``Attention1D``, ``MLP1D``, ``RotarySetup1D``, ``LMHead1D``). Primary target: N300 (mesh ``(1, 2)``). + +Phi-4 is a Phi3-architecture 14B dense decoder (40 layers, hidden 5120, 40 attn / 10 KV heads, +head_dim 128, SwiGLU MLP intermediate 17920, vocab 100352, ``rope_theta=250000``, full RoPE, +no QKV bias, no q/k norm, no sliding window). The HF checkpoint ships **fused** ``qkv_proj`` and +``gate_up_proj`` — :mod:`weight_utils` splits them host-side before this builds TTNN modules. + +Tensor layout contracts: + - **Prefill** hidden states: ``[1, 1, S, dim]`` TILE, ``S % 128 == 0``. + - **Decode** hidden states: ``[1, 1, B, dim]`` TILE (``B`` padded to tile in modules). + +The tensor graph is provider-neutral. ``hf_adaptor.py`` owns Hugging Face loading, +checkpoint validation, and tensor conversion; this module only builds and executes +the TT graph. +""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.models.phi4 import weight_utils +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D +from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + """Minimal LazyWeight; ``Attention1D`` / ``MLP1D`` / ``Embedding1D`` resolvers fill mesh + memory.""" + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +@dataclass +class Phi4PagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (``block_size`` / ``max_num_blocks`` only).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Phi4TransformerConfig: + """Complete provider-neutral tensor-graph configuration.""" + + n_layers: int + vocab_size: int + max_batch_size: int + max_seq_len: int + dim: int + num_devices: int + mesh_device: ttnn.MeshDevice + embedding_config: Embedding1DConfig + rope_config: Rope1DConfig + block_configs: list[Phi4DecoderLayerConfig] + norm_config: RMSNorm1DConfig + lm_head_config: LMHead1DConfig + sampling_config: Sampling1DConfig | None = None + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtypes: list[ttnn.DataType | None] = field(default_factory=list) + tt_ccl: Any = None + cache_path: str | None = None + + +@dataclass(frozen=True) +class Phi4ModelParameters: + """Provider-neutral Phi-4 dimensions consumed by the TT config builder.""" + + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + max_batch_size: int + max_seq_len: int + + +@dataclass(frozen=True) +class Phi4LayerWeights: + wqkv: torch.Tensor + wo: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Phi4Weights: + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Phi4LayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +_LOFI_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, +) +"""LoFi + packer L1 acc for the MLP FF1/FF3 matmuls in performance mode. + +Mirrors TTTv1 ``DecodersPrecision.performance("Phi-4")`` (``model_config.py:183-222``): the +default performance branch sets ``FF1_FF3 → BFP4`` and ``LI_FF1_FF3 → LOFI``. +""" + + +@dataclass(frozen=True) +class Phi4PrecisionConfig: + """Per-layer precision + math-fidelity recipe for Phi-4 on N300. + + Mirrors the fields TTTv1's ``DecodersPrecision`` distinguishes for this model. Phi-4 is + special-cased into the **Llama-3 / Mistral-7B / Phi-3-mini accuracy branch** + (``model_config.py:131-161``), NOT the generic Qwen2 (BF16/HiFi4) branch. That branch keeps + **BFP8** WQKV / WO / KV-cache + HIFI2_FP16 MLP accumulation even in accuracy mode (the code + logs *"running out of DRAM … using BFP8 for WQKV"* specifically for phi-4). So: + + * **Accuracy**: everything BFP8; attention stays at the ``Attention1D`` default (HIFI2). + * **Performance** (``model_config.py:183-222``): only ``FF1_FF3 → BFP4`` and + ``LI_FF1_FF3 → LOFI``; everything else identical to accuracy. + + Two module-level recipes are exposed: :data:`PHI4_ACCURACY` (default) and + :data:`PHI4_PERFORMANCE`. The HF adaptor selects one via ``optimizations=``; + use ``dataclasses.replace(PHI4_ACCURACY, ...)`` to customize a single field. + """ + + # Attention weight / KV-cache dtypes. BFP8 in both modes (TTTv1 Llama/Phi accuracy branch). + wqkv_dtype: ttnn.DataType = ttnn.bfloat8_b + wo_dtype: ttnn.DataType = ttnn.bfloat8_b + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + # MLP FF1/FF3 weight dtype. Accuracy keeps BFP8; performance overrides BFP4. + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat8_b + # MLP FF2 (down) weight dtype — BFP8 in both modes. + mlp_w2_dtype: ttnn.DataType = ttnn.bfloat8_b + + # MLP FF1/FF3 compute kernel. ``None`` → MLP1D default (HIFI2_FP16); performance overrides LOFI. + mlp_ff1_3_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + # LM head dtype. BFP8 in both modes (TTTv1 Phi-4 uses BFP8 LM head; no bf16 upgrade). + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + + +# TTTv1 ``DecodersPrecision.accuracy("Phi-4")`` (Llama/Mistral/Phi branch, ``model_config.py:131-161``): +# BFP8 attention weights + BFP8 KV cache + HIFI2_FP16 MLP; attention kernels at TTTv2 HIFI2 default. +PHI4_ACCURACY = Phi4PrecisionConfig() + +# TTTv1 ``DecodersPrecision.performance("Phi-4")`` (``model_config.py:183-222``): FF1_FF3 → BFP4 and +# LI_FF1_FF3 → LOFI; everything else identical to accuracy. +PHI4_PERFORMANCE = Phi4PrecisionConfig( + mlp_w1_w3_dtype=ttnn.bfloat4_b, + mlp_ff1_3_compute_kernel_cfg=_LOFI_COMPUTE_KERNEL_CFG, +) + + +def _slice_last_token_tile(x: ttnn.Tensor, last_token_idx: int) -> ttnn.Tensor: + """Slice the 32-row tile containing ``last_token_idx`` from ``[1, 1, S, W]``. + + Width-sharded LM matmul M tile rows must equal ``LMHead1D`` program-config tile rows. + """ + floor = (last_token_idx // 32) * 32 + return ttnn.slice(x, (0, 0, floor, 0), (1, 1, floor + 32, x.shape[-1])) + + +def _post_attn_norm_decode_configs( + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding so its output matches MLP1D's W1/W3 input. + + MLP1D decode uses ``_dram_shard_core_grid_k_n(dim, padded_hidden / num_devices)`` for W1/W3 + inputs, but the default RMSNorm program config is derived from ``_compute_norm_core_grid(dim)`` + alone. Mismatched DRAM-width-shard between RMSNorm output and MLP1D W1/W3 input silently + corrupts decode activations (observed on the 7B / 14B ports — same shape pattern on N300 Phi-4). + """ + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + memory_config = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return program_config, memory_config + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, x: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + return x + + if memory_config is None: + memory_config = x.memory_config() + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + + +@dataclass +class _Phi4WHTuning: + """Wormhole-specific L1 / cutoff tuning resolved at build time. + + Precision-vs-fidelity knobs live on :class:`Phi4PrecisionConfig`. This dataclass only carries + non-mode-dependent L1 footprint controls: MLP prefill cutoff and the optional W1→DRAM spill on + decode. + """ + + mlp_prefill_len_cutoff: int | None = None + mlp_decode_spill_w1_to_dram: bool = False + # Use ttnn.experimental.minimal_matmul for the QKV + W2 prefill matmuls above seq_len > 128 (TTTv1 + # parity, PLAN_01). A/B escape hatch: set DISABLE_MINIMAL_MATMUL=1 to force ttnn.linear. On a 14B the + # batch-32-ci prefill is matmul-compute-bound (~80% of FLOPs = the 3 MLP matmuls), so minimal_matmul + # is a real prefill-TTFT win — it closes the batch-32-ci TTFT gap vs TTTv1 (which uses minimal_matmul + # for the same matmuls). The shared plumbing (attention_1d.use_minimal_qkv_matmul / mlp_1d + # use_minimal_w2_matmul, both gated seq_len>128) is already in the base; this flag just engages it. + prefill_minimal_matmul: bool = True + + +def _resolve_phi4_wh_tuning(*, num_dev: int, max_batch_size: int) -> _Phi4WHTuning: + """Pick WH L1 tuning knobs for Phi-4 on N300. + + Phi-4's FF matmul is wide (intermediate 17920); mirror the 14B / 7B port empirical L1 cutoff + (``mlp_prefill_len_cutoff=256``). ``mlp_decode_spill_w1_to_dram`` is off; re-evaluate if decode + batch-32 trips L1 circular-buffer validation on N300. + """ + t = _Phi4WHTuning( + mlp_prefill_len_cutoff=256, + mlp_decode_spill_w1_to_dram=False, + ) + t.prefill_minimal_matmul = not os.environ.get("DISABLE_MINIMAL_MATMUL") + logger.info( + f"L1 tuning for Phi-4 on {num_dev} device(s): " + f"prefill_len_cutoff={t.mlp_prefill_len_cutoff}, " + f"decode_spill_w1_to_dram={t.mlp_decode_spill_w1_to_dram}, " + f"prefill_minimal_matmul={t.prefill_minimal_matmul}" + ) + return t + + +@dataclass +class Phi4DecoderLayerConfig: + attention_norm_config: RMSNorm1DConfig + attention_config: Attention1DConfig + ff_norm_config: RMSNorm1DConfig + mlp_config: MLP1DConfig + + +def _build_decoder_layer( + *, + idx: int, + weights: Phi4LayerWeights, + mcfg: Phi4ModelParameters, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Phi4PrecisionConfig, + paged_cfg: Phi4PagedAttentionConfig, + cache_path: Path | None, + wh: _Phi4WHTuning, +) -> Phi4DecoderLayerConfig: + """Build one decoder-layer config from provider-converted tensors.""" + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, + dtype=precision.wqkv_dtype, + cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None, + ) + lazy_wo = _lazy( + weights.wo, + dtype=precision.wo_dtype, + cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None, + ) + + attention_config = Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=mcfg.n_heads, + n_kv_heads=mcfg.n_kv_heads, + head_dim=mcfg.head_dim, + max_batch_size=mcfg.max_batch_size, + max_seq_len=mcfg.max_seq_len, + q_norm_config=None, + k_norm_config=None, + wqkv_bias=None, + use_vllm_paged_kv_cache=True, + paged_attention_config=paged_cfg, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + prefill_qkv_minimal_matmul=wh.prefill_minimal_matmul, + ) + + mlp_config = MLP1DConfig( + w1=_lazy( + weights.w1, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None, + ), + w2=_lazy( + weights.w2, + dtype=precision.mlp_w2_dtype, + cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None, + ), + w3=_lazy( + weights.w3, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None, + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=mcfg.max_batch_size, + prefill_len_cutoff=wh.mlp_prefill_len_cutoff, + w1_w3_dtype=precision.mlp_w1_w3_dtype, + w2_dtype=precision.mlp_w2_dtype, + ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_spill_w1_to_dram_before_w3=wh.mlp_decode_spill_w1_to_dram, + prefill_w2_minimal_matmul=wh.prefill_minimal_matmul, + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + dim=mcfg.dim, + hidden_dim=mcfg.hidden_dim, + num_devices=num_dev, + max_batch_size=mcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1DConfig: + lw = _lazy( + weight, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=mcfg.rms_norm_eps, + max_batch_size=mcfg.max_batch_size, + tt_ccl=tt_ccl, + **extra, + ) + + return Phi4DecoderLayerConfig( + attention_norm_config=_build_norm(weights.attention_norm, "pre_attn"), + attention_config=attention_config, + ff_norm_config=_build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ), + mlp_config=mlp_config, + ) + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + mcfg: Phi4ModelParameters, + lm_head_dtype: ttnn.DataType, + cache_path: Path | None, +) -> LMHead1DConfig: + """Build the vocab-sharded LM head with DRAM-matmul program configs. + + LM head DRAM matmul is sized for decode batch tiles (``max_batch_size``). Prefill logits + use a single 32-row tile via ``post_process_prefill_output`` / :func:`_slice_last_token_tile`. + """ + lm_splits, lm_split_sizes, lm_weights_memcfgs = weight_utils.build_lm_head_lazy_weights( + mesh_device, + lm_head_weight, + dim=mcfg.dim, + vocab_size=mcfg.vocab_size, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(mcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(mcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, mcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(mcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=mcfg.dim, + max_batch_size=mcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + ) + + +def build_phi4_transformer_config( + *, + mesh_device: ttnn.MeshDevice, + params: Phi4ModelParameters, + weights: Phi4Weights, + n_layers: int, + precision: Phi4PrecisionConfig, + cache_path: Path | None, + paged_attention_config: Phi4PagedAttentionConfig, +) -> Phi4TransformerConfig: + """Build the TT tensor graph from provider-neutral dimensions and converted tensors.""" + + num_devices = mesh_device.get_num_devices() + if params.n_heads % num_devices or params.n_kv_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({params.n_heads}/{params.n_kv_heads}) " + f"must be divisible by device count ({num_devices})" + ) + if len(weights.layers) != n_layers: + raise ValueError(f"Expected {n_layers} decoder layer weight sets, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) if num_devices > 1 else None + topology = default_topology(mesh_device) + wh = _resolve_phi4_wh_tuning(num_dev=num_devices, max_batch_size=params.max_batch_size) + block_configs = [ + _build_decoder_layer( + idx=index, + weights=weights.layers[index], + mcfg=params, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + paged_cfg=paged_attention_config, + cache_path=cache_path, + wh=wh, + ) + for index in range(n_layers) + ] + return Phi4TransformerConfig( + n_layers=n_layers, + vocab_size=params.vocab_size, + max_batch_size=params.max_batch_size, + max_seq_len=params.max_seq_len, + dim=params.dim, + num_devices=num_devices, + mesh_device=mesh_device, + embedding_config=Embedding1DConfig( + weights=_lazy( + weights.embedding, + dtype=ttnn.bfloat16, + cache=(cache_path / "embedding", "tok_embeddings") if cache_path else None, + ), + mesh_device=mesh_device, + embed_scale=1.0, + ), + rope_config=Rope1DConfig( + cos_matrix=_lazy( + weights.rope_cos, + dtype=ttnn.bfloat16, + cache=(cache_path / "rope", "cos") if cache_path else None, + ), + sin_matrix=_lazy( + weights.rope_sin, + dtype=ttnn.bfloat16, + cache=(cache_path / "rope", "sin") if cache_path else None, + ), + max_batch_size=params.max_batch_size, + head_dim=params.head_dim, + device=mesh_device, + use_qk_fused=False, + ), + block_configs=block_configs, + norm_config=RMSNorm1DConfig( + weight=_lazy( + weights.final_norm, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", "final") if cache_path else None, + ), + mesh_device=mesh_device, + eps=params.rms_norm_eps, + max_batch_size=params.max_batch_size, + tt_ccl=tt_ccl, + ), + lm_head_config=_build_lm_head( + mesh_device=mesh_device, + lm_head_weight=weights.lm_head, + mcfg=params, + lm_head_dtype=precision.lm_head_dtype, + cache_path=cache_path, + ), + sampling_config=Sampling1DConfig( + vocab_size=params.vocab_size, + valid_vocab_size=params.vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=_nearest_32(params.max_batch_size), + allow_force_argmax=False, + pad_to_power_of_2=True, + ), + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtypes=[None] * n_layers, + tt_ccl=tt_ccl, + cache_path=str(cache_path) if cache_path is not None else None, + ) + + +class Phi4DecoderLayer(LightweightModule): + def __init__( + self, + *, + input_layernorm: RMSNorm1D, + attention: Attention1D, + post_attention_layernorm: RMSNorm1D, + mlp: MLP1D, + ): + super().__init__() + self.input_layernorm = input_layernorm + self.post_attention_layernorm = post_attention_layernorm + self.mlp = mlp + self.attention_norm = input_layernorm + self.attention = attention + self.ff_norm = post_attention_layernorm + self.feed_forward = mlp + + @classmethod + def from_config(cls, config: Phi4DecoderLayerConfig) -> Phi4DecoderLayer: + return cls( + input_layernorm=RMSNorm1D.from_config(config.attention_norm_config), + attention=Attention1D.from_config(config.attention_config), + post_attention_layernorm=RMSNorm1D.from_config(config.ff_norm_config), + mlp=MLP1D.from_config(config.mlp_config), + ) + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + *, + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + # Fractured embed / norm activations must be all-gathered to full ``dim`` before + # Attention1D / MLP1D (QKV matmul expects width ``dim``). + r = self.input_layernorm.prefill_forward(x) + r = _all_gather_rmsnorm_tensor(self.input_layernorm, r) + r = self.attention.prefill_forward( + r, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + r2 = self.post_attention_layernorm.prefill_forward(h) + r2 = _all_gather_rmsnorm_tensor(self.post_attention_layernorm, r2) + r2 = self.mlp.prefill_forward(r2) + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def decode_forward( + self, + x: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + xa = _all_gather_rmsnorm_tensor( + self.input_layernorm, x, memory_config=self.input_layernorm.config.decode_memory_config + ) + r = self.input_layernorm.forward(xa, "decode") + r = self.attention.forward(r, current_pos, rot_mats, mode="decode", page_table=page_table) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + hf = _all_gather_rmsnorm_tensor( + self.post_attention_layernorm, h, memory_config=self.post_attention_layernorm.config.decode_memory_config + ) + r2 = self.post_attention_layernorm.forward(hf, "decode") + r2 = self.mlp.forward(r2, "decode") + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + +class Phi4Transformer(LightweightModule): + """ + Full decoder for Phi-4 (TTTv2 modules only) on N300. + + Prefill/decode on **embedded** activations match the model-owned runtime. Token embedding + is ``embed_prefill`` / ``embed_decode``. Bind KV with ``set_kv_cache`` before first forward. + """ + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__(self, config: Phi4TransformerConfig): + from tqdm import tqdm + + super().__init__() + self.config = config + self.embed = Embedding1D.from_config(config.embedding_config) + self.embedding = self.embed + self.rope_setup = RotarySetup1D.from_config(config.rope_config) + self.layers = [ + Phi4DecoderLayer.from_config(config.block_configs[index]) + for index in tqdm(range(config.n_layers), desc="Building layers") + ] + self.norm = RMSNorm1D.from_config(config.norm_config) + self.lm_head = LMHead1D.from_config(config.lm_head_config) + self.sampling = Sampling1D.from_config(config.sampling_config) if config.sampling_config is not None else None + self.supports_on_device_sampling = self.sampling is not None + self.mesh_device = config.mesh_device + self.tt_ccl = config.tt_ccl or (get_tt_ccl(config.mesh_device) if config.num_devices > 1 else None) + self.vocab_size = config.vocab_size + self.n_layers = config.n_layers + self.num_devices = config.num_devices + self.decode_residual_memcfg = config.decode_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.prefill_residual_memcfg = config.prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtypes = config.activation_dtypes or [None] * config.n_layers + self.model_args = None + + @property + def n_kv_heads(self) -> int: + return self.config.block_configs[0].attention_config.n_kv_heads + + def iter_executor_named_modules(self): + if not hasattr(self, "layers"): + return + for index, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{index}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in self.config.block_configs) + for config in tuple({id(item): item for item in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, + block_size=block_size, + max_num_blocks=max_num_blocks, + ) + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for index, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {index} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {index} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = pair + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). The batched path always + # returns the full hidden state when no runtime last-token slice is supplied; the executor does + # per-slot last-token extraction + norm/lm_head so those stages stay bit-identical to the single-user path. + x = x_embed + for layer in self.layers: + x = layer.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + + if last_token_index is not None and last_token_slice is None: + raise ValueError("last_token_index is required with a runtime last_token_slice") + if get_last_token == -1 and last_token_slice is None: + return x + + old = x + if last_token_slice is None: + x_tile = _slice_last_token_tile(old, get_last_token) + else: + x_tile = ttnn.slice( + old, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(old.shape[2]) // 32, + ) + ttnn.deallocate(old) + if last_token_index is not None: + if x_tile.dtype != ttnn.bfloat16: + old_tile = x_tile + x_tile = ttnn.typecast(x_tile, ttnn.bfloat16) + ttnn.deallocate(old_tile) + old_tile = x_tile + x_tile = ttnn.embedding(last_token_index, x_tile, layout=ttnn.TILE_LAYOUT) + x_tile = ttnn.unsqueeze_to_4D(x_tile) + ttnn.deallocate(old_tile) + return self._last_tile_logits(x_tile) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + x = _slice_last_token_tile(hidden_states, last_token_idx) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + return self._last_tile_logits(x) + + def _last_tile_logits(self, x_tile: ttnn.Tensor) -> ttnn.Tensor: + """Final-norm + all-gather + LM-head on a 32-row tile. ``x_tile`` shape ``[1, 1, 32, dim]``.""" + x = self.norm.prefill_forward(x_tile) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + del last_token_slice, last_token_index + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), + ) + x = ttnn.matmul( + selector, + hidden_states, + dtype=ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + compute_kernel_config=ttnn.init_device_compute_kernel_config( + self.mesh_device.arch(), + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ), + ) + ttnn.deallocate(selector) + return self._last_tile_logits(x) + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for layer in self.layers: + x = layer.decode_forward(x, current_pos, rot_mats, page_table=page_table) + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + return self.lm_head.forward(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) diff --git a/code/models/common/models/phi4/weight_utils.py b/code/models/common/models/phi4/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..898a832065b3608b48664dbb310ce7f5956553bb --- /dev/null +++ b/code/models/common/models/phi4/weight_utils.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side HF weight layout helpers for Phi-4 (TTTv2 port). + +Phi-4 (``transformers.models.phi3``) ships **fused** projections: + - ``self_attn.qkv_proj`` : ``[n_heads·head_dim + 2·n_kv·head_dim, dim]`` + - ``mlp.gate_up_proj`` : ``[2·intermediate_size, dim]`` (HF order: gate first, up second) + +This module splits those into the per-projection tensors the TTTv2 ``Attention1D`` / +``MLP1D`` modules expect, then applies the Llama/Qwen HF→Meta ``reverse_permute`` to Q/K +(Phi-3/Phi-4 use GPT-NeoX ``rotate_half`` RoPE, identical to the Llama path). The permute +helpers are duplicated from ``load_checkpoints.reverse_permute`` to avoid importing +``models/tt_transformers``. + +Phi-4 has **no** QKV bias (``attention_bias=false``) and **no** q/k RMSNorm. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.modules.lazy_weight import LazyWeight + + +def reverse_permute(tensor: torch.Tensor, n_heads: int, dim1: int, dim2: int) -> torch.Tensor: + """HF Q/K weight rows → Meta layout (RoPE-compatible).""" + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def permute_hf_rope_to_meta_tables(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert HF cos/sin tables to Meta-style tables for ``RotarySetup1D`` (see attention tests).""" + if len(cos.shape) == 3: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + return cos, sin + + +def build_rope_cos_sin_torch( + rotary_emb: Any, table_len: int, head_dim: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Build ``[1, 1, table_len, head_dim]`` cos/sin tensors (Meta layout) from HF rotary module.""" + x = torch.zeros(1, 1, table_len, head_dim, dtype=dtype) + position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0) + with torch.no_grad(): + cos_hf, sin_hf = rotary_emb(x, position_ids) + cos_m, sin_m = permute_hf_rope_to_meta_tables(cos_hf.float(), sin_hf.float()) + return cos_m.to(dtype), sin_m.to(dtype) + + +def split_fused_qkv(self_attn: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]: + """Slice Phi-3/Phi-4 fused ``qkv_proj.weight`` into (q, k, v) raw HF weight rows. + + ``qkv_proj`` is ``nn.Linear(dim, n_heads·head_dim + 2·n_kv·head_dim)`` so its weight has + shape ``[op_size, dim]`` laid out as ``[Q rows | K rows | V rows]`` (HF ``Phi3Attention.forward`` + splits along the projection dim in exactly this order). + """ + cfg = self_attn.config + n_heads = cfg.num_attention_heads + n_kv_heads = cfg.num_key_value_heads + head_dim = getattr(self_attn, "head_dim", cfg.hidden_size // n_heads) + q_size = n_heads * head_dim + kv_size = n_kv_heads * head_dim + + qkv = self_attn.qkv_proj.weight + assert qkv.shape[0] == q_size + 2 * kv_size, ( + f"Unexpected qkv_proj rows {qkv.shape[0]} != {q_size + 2 * kv_size} " + f"(n_heads={n_heads}, n_kv={n_kv_heads}, head_dim={head_dim})" + ) + wq_raw = qkv[:q_size] + wk_raw = qkv[q_size : q_size + kv_size] + wv_raw = qkv[q_size + kv_size : q_size + 2 * kv_size] + return wq_raw, wk_raw, wv_raw, n_heads, n_kv_heads, head_dim + + +def attention_wqkv_wo_from_hf_layer( + self_attn: Any, + num_devices: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack Phi-4 attention weights for ``Attention1D`` (per-device interleaved QKV). + + Returns ``(wqkv, wo)`` where ``wqkv`` is ``[1, 1, dim, sum_per_dev·num_devices]`` with each + device's slice laid out as ``[Q_chunk | K_chunk | V_chunk]``, matching + ``get_attention_weights_from_ref_model`` in ``test_attention_1d.py``. Phi-4 has no QKV bias + and no q/k RMSNorm, so neither is returned. + """ + wq_raw, wk_raw, wv_raw, n_heads, n_kv_heads, head_dim = split_fused_qkv(self_attn) + dim = wq_raw.shape[1] + n_heads_times_head_dim = n_heads * head_dim + n_kv_heads_times_head_dim = n_kv_heads * head_dim + + wq_meta = reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim) + wk_meta = reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim) + wv_meta = wv_raw + wo_meta = self_attn.o_proj.weight + + wq = wq_meta.T + wk = wk_meta.T + wv = wv_meta.T + wo = wo_meta.T + + qkv_list = [] + for i in range(num_devices): + wq_chunk = torch.chunk(wq, num_devices, dim=1)[i] + wk_chunk = torch.chunk(wk, num_devices, dim=1)[i] + wv_chunk = torch.chunk(wv, num_devices, dim=1)[i] + qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1) + qkv_list.append(qkv) + wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0).clone() + wo = wo.unsqueeze(0).unsqueeze(0).clone() + return wqkv, wo + + +def mlp_weights_from_hf_layer(mlp: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split Phi-4 fused ``gate_up_proj`` and return (w1, w2, w3) = (gate^T, down^T, up^T). + + HF ``Phi3MLP.forward``: ``gate, up = gate_up_proj(x).chunk(2, -1)``; ``down_proj(up * silu(gate))``. + So the ``gate_up_proj`` weight rows are ``[gate | up]`` — gate is the first ``intermediate_size`` + rows, up the second. + """ + gate_up = mlp.gate_up_proj.weight + inter = gate_up.shape[0] // 2 + assert gate_up.shape[0] == 2 * inter, f"gate_up_proj rows {gate_up.shape[0]} not even" + gate = gate_up[:inter] + up = gate_up[inter:] + w1 = gate.T.contiguous().clone() + w3 = up.T.contiguous().clone() + w2 = mlp.down_proj.weight.T.contiguous().clone() + return w1, w2, w3 + + +def rms_weight_torch(layernorm: Any) -> torch.Tensor: + return layernorm.weight.detach().float().clone() + + +def embed_tokens_torch(embed: Any) -> torch.Tensor: + w = embed.weight.detach().float().clone() + return w.unsqueeze(0).unsqueeze(0) + + +def build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """ + Column-split LM head for ``LMHead1D`` (adapted from ``LMHead1D.from_model_args`` math, no v1 args). + ``lm_head_weight`` shape: ``[vocab_size, dim]`` (HF). + + Returns ``(lazy_weights, split_sizes, weights_memcfgs)`` for ``LMHead1DConfig`` DRAM-sharded matmuls. + """ + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / 32) * 32 + if vocab_size < padded_vocab_size: + pad = padded_vocab_size - vocab_size + torch_w = torch.cat([torch_w, torch.zeros(torch_w.shape[0], pad, dtype=torch_w.dtype)], dim=-1) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + tile = ttnn.TILE_SIZE + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (tile * dram_size.x)) * (tile * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights: list[LazyWeight] = [] + weights_memcfgs: list[ttnn.MemoryConfig] = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_splits.append(torch_w[:, start:end]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + name = f"lm_head_split_{i}_{combined.shape[-1]}" + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, name) if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs diff --git a/code/models/common/models/qwen25_72b/README.md b/code/models/common/models/qwen25_72b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f9197138b5306949803ce9d2461ba0d981ab0074 --- /dev/null +++ b/code/models/common/models/qwen25_72b/README.md @@ -0,0 +1,47 @@ +# Qwen2.5-72B with TTTv2 + +This directory contains the large Qwen2.5-72B TTTv2 product path. + +## Construction path + +```text +HF checkpoint + -> hf_adaptor.py: provider metadata, tokenizer, and weights + -> model.py: multi-device Qwen2.5 tensor graph + -> executor.py: thin family builder plus concrete direct-run helpers + -> generator.py: vLLM construction, DP composition, and dispatch +``` + +The tensor model composes reusable embedding, rotary, RMSNorm, attention, MLP, +LM-head, optional sampling, and collective modules while retaining 72B +topology, precision, and program policy locally. + +## Executor composition + +The public 72B executor/config/builder remains importable from this directory. +Its lifecycle is supplied by `models/common/models/qwen2_executor.py` over +`models/common/models/executor.py::ModelExecutor`. + +The shared owner composes paged-KV management, prefill/decode runtimes, +eager/trace compilation, output handling, warmup, and cleanup. Qwen2.5-72B uses +plain coordinator prefill warmup rather than the 7B Q128 priming hook. + +Concrete `run_prefill`, `run_decode`, `run_lm_head`, last-token slicing, and +deallocation helpers remain model-local because they invoke the concrete +tensor model. + +## vLLM, DP, and ownership + +The generator builds one executor per lane, configures `VLLMAdapter`, and uses +`LaneGroupExecutor` for DP. Each lane owns its KV tensors, compile/trace +artifacts, output leases, sampling buffers, and terminal cleanup. + +## Tests + +- `models/common/tests/models/qwen25_72b/test_hf_adaptor.py` +- `models/common/tests/models/qwen25_72b/test_model_runtime_surface.py` +- `models/common/tests/models/qwen25_72b/test_demo_contract.py` +- `models/common/tests/demos/qwen25_72b/demo.py` +- `models/common/tests/models/test_qwen2_executor_family.py` +- `models/common/tests/llm_runtime/test_executor_integration.py` +- `models/common/tests/llm_runtime/test_model_executor.py` diff --git a/code/models/common/models/qwen25_72b/__init__.py b/code/models/common/models/qwen25_72b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0ab695204a2a8d6d8f18e8ae6aa9f88bb2094aa2 --- /dev/null +++ b/code/models/common/models/qwen25_72b/__init__.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +from models.common.models.qwen25_72b.executor import ( + Qwen25_72BExecutor, + Qwen25_72BExecutorConfig, + build_qwen25_72b_executor, +) +from models.common.models.qwen25_72b.generator import ( + Qwen25_72BGenerator, + Qwen25_72BGeneratorConfig, + build_qwen25_72b_generator, +) +from models.common.models.qwen25_72b.hf_adaptor import Qwen25_72BForCausalLM, from_pretrained +from models.common.models.qwen25_72b.model import ( + QWEN25_72B_ACCURACY, + QWEN25_72B_PERFORMANCE, + Qwen25_72B, + Qwen25_72BConfig, + Qwen25_72BExecutorRuntimeConfig, + Qwen25_72BPagedAttentionConfig, + Qwen25_72BPrecisionConfig, +) + +__all__ = [ + "Qwen25_72B", + "Qwen25_72BConfig", + "Qwen25_72BExecutorRuntimeConfig", + "Qwen25_72BPagedAttentionConfig", + "Qwen25_72BPrecisionConfig", + "QWEN25_72B_ACCURACY", + "QWEN25_72B_PERFORMANCE", + "Qwen25_72BGenerator", + "Qwen25_72BGeneratorConfig", + "Qwen25_72BExecutor", + "Qwen25_72BExecutorConfig", + "Qwen25_72BForCausalLM", + "build_qwen25_72b_executor", + "build_qwen25_72b_generator", + "from_pretrained", +] diff --git a/code/models/common/models/qwen25_72b/executor.py b/code/models/common/models/qwen25_72b/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..9403bf17c2529dd3dc317ffa0d3b1e30c8cfba96 --- /dev/null +++ b/code/models/common/models/qwen25_72b/executor.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen2.5-72B construction and concrete compatibility entry points.""" + +from models.common.models.qwen2_executor import Qwen25_72BExecutor, Qwen25_72BExecutorConfig +from models.common.models.qwen2_executor import build_qwen25_72b_executor as _build_qwen25_72b_executor +from models.common.models.qwen25_72b.hf_adaptor import Qwen25_72BForCausalLM +from models.common.models.qwen25_72b.model import _slice_last_token_tile + + +def build_qwen25_72b_executor( + llm: Qwen25_72BForCausalLM, + config: Qwen25_72BExecutorConfig, +) -> Qwen25_72BExecutor: + return _build_qwen25_72b_executor(llm, config) + + +def run_prefill(model, token_ids_tt, *, start_pos: int = 0): + return model.prefill_from_token_ids(token_ids_tt, start_pos=start_pos) + + +def run_decode(model, token_id_tt, *, current_pos: int): + return model.decode_from_token_ids(token_id_tt, current_pos=current_pos) + + +def run_lm_head(model, hidden_tt): + if len(hidden_tt.shape) == 4 and hidden_tt.shape[2] > 32: + old = hidden_tt + hidden_tt = _slice_last_token_tile(old, hidden_tt.shape[2] - 1) + old.deallocate(True) + return model.lm_logits(hidden_tt) diff --git a/code/models/common/models/qwen25_72b/generator.py b/code/models/common/models/qwen25_72b/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..050b3e706fc742e5c69a5321e0fb759b8593aa76 --- /dev/null +++ b/code/models/common/models/qwen25_72b/generator.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Qwen2.5-72B.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.qwen25_72b.executor import Qwen25_72BExecutorConfig, build_qwen25_72b_executor +from models.common.models.qwen25_72b.hf_adaptor import DEFAULT_HF_REVISION, from_pretrained +from models.common.models.qwen25_72b.model import Qwen25_72BPagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Qwen25_72BGeneratorConfig: + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + hf_revision: str | None = DEFAULT_HF_REVISION + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Qwen25_72BGenerator: + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + device_sampling_enabled: bool = True, + ): + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_qwen25_72b_generator( + Qwen25_72BGeneratorConfig( + hf_model=str(hf_model), + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def prefill_forward( + self, + tokens, + page_table, + *, + enable_trace: bool, + prompt_lens=None, + start_pos=None, + empty_slots=None, + kv_cache=None, + sampling_params=None, + **compatibility_kwargs, + ): + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens, + start_pos, + page_table, + *, + enable_trace: bool, + kv_cache=None, + sampling_params=None, + reset_batch=False, + read_from_device=True, + **compatibility_kwargs, + ): + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward(execution=execution, read_from_device=read_from_device, **normalized) + + def read_decode_output(self, tt_out: Any, *, async_read: bool = False) -> Any: + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host(self, tt_out: Any, *, is_tokens: bool = False) -> tuple[Any, Any]: + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill(self, *, kv_cache, can_sample_on_device: bool, enable_trace: bool) -> None: + return self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def warmup_model_decode( + self, *, kv_cache, max_batch_size: int, num_blocks: int, can_sample_on_device: bool, enable_trace: bool + ) -> None: + return self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self): + return self.target.cleanup() + + def _select_prefill_execution(self, normalized: NormalizedPrefillKwargs, trace_requested: bool): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_qwen25_72b_generator(config: Qwen25_72BGeneratorConfig) -> Qwen25_72BGenerator: + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Qwen25_72BPagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + instruct="Instruct" in config.hf_model, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Qwen25_72BExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(include_decode_top_k=config.device_sampling_enabled), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_qwen25_72b_executor(llm, executor_config)) + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Qwen25_72BGenerator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Qwen2.5-72B model must contain at least one attention layer") + attention_configs = tuple(_attention_config(layer) for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Qwen layer must expose the same KV head shape") + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _attention_config(layer): + attention = getattr(layer, "attention", None) or getattr(layer, "self_attn", None) + if attention is not None: + return attention.config + if hasattr(layer, "config"): + return layer.config + raise AttributeError("Qwen2.5-72B layer must expose attention/self_attn.config or config") + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/qwen25_72b/hf_adaptor.py b/code/models/common/models/qwen25_72b/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..b1d6a9143511bb0b5685b356083facb3deb0b0e5 --- /dev/null +++ b/code/models/common/models/qwen25_72b/hf_adaptor.py @@ -0,0 +1,365 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face product adaptor for Qwen2.5-72B-Instruct.""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.models.qwen25_72b import weight_utils +from models.common.models.qwen25_72b.model import ( + DEFAULT_HF_REVISION, + QWEN25_72B_ACCURACY, + QWEN25_72B_PERFORMANCE, + Qwen25_72B, + Qwen25_72BConfig, + Qwen25_72BLayerWeights, + Qwen25_72BPagedAttentionConfig, + Qwen25_72BPrecisionConfig, + Qwen25_72BWeights, + build_qwen25_72b_model, +) + +DEFAULT_HF_MODEL = "Qwen/Qwen2.5-72B-Instruct" + + +def _local_files_only() -> bool: + return any( + os.getenv(name, "").lower() in {"1", "true", "yes"} for name in ("CI", "HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ) + + +@dataclass(frozen=True) +class Qwen25_72BGenerationConfig: + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Qwen25_72BRuntimeConfig: + model_name: str + model_cache_path: Path | None + max_prefill_chunk_size: int + max_context_len: int + max_seq_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] + n_layers: int + n_kv_heads: int + head_dim: int + max_batch_size: int + cluster_shape: list[int] + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + return ( + prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + chat = [] + if isinstance(prompt_text, str): + if system_prompt_text: + chat.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + chat.append({"role": "user", "content": prompt_text}) + else: + chat = prompt_text + try: + return _chat_template_ids(tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True)) + except ValueError: + pass + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Qwen25_72BForCausalLM: + model: Qwen25_72B + tokenizer: Any + runtime_config: Qwen25_72BRuntimeConfig + instruct: bool = True + generation_config: Qwen25_72BGenerationConfig = field(default_factory=Qwen25_72BGenerationConfig) + + def __post_init__(self) -> None: + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + self.generation_config = Qwen25_72BGenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=tuple(getattr(self.tokenizer, "stop_tokens", ()) or ()), + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + return encode_prompt( + self.tokenizer, + prompt_text, + system_prompt_text, + instruct=self.instruct if instruct is None else instruct, + ) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_tokenizer(hf_model: str, hf_revision: str | None = DEFAULT_HF_REVISION): + tokenizer = AutoTokenizer.from_pretrained( + hf_model, + revision=hf_revision, + trust_remote_code=True, + local_files_only=_local_files_only(), + ) + stop_ids = _qwen_stop_token_ids(tokenizer) + tokenizer.stop_tokens = list(stop_ids) + return tokenizer + + +def _qwen_stop_token_ids(tokenizer) -> tuple[int, ...]: + ids = [] + eos = getattr(tokenizer, "eos_token_id", None) + if eos is not None: + ids.extend([eos] if isinstance(eos, int) else list(eos)) + for token in ("<|im_end|>", "<|im_start|>"): + token_id = tokenizer.convert_tokens_to_ids(token) + if isinstance(token_id, int) and token_id >= 0: + ids.append(token_id) + return tuple(dict.fromkeys(int(value) for value in ids)) + + +def _trace_seq_lens(num_devices: int, max_prefill_chunk_size: int, max_seq_len: int) -> tuple[int, ...]: + if num_devices != 8: + raise ValueError(f"Qwen2.5-72B supports exactly 8 devices (T3K), got {num_devices}") + return tuple(length for length in (128, 1024) if length <= min(max_prefill_chunk_size, max_seq_len)) + + +def _cache_path(hf_model: str, mesh_device, cache_dir: Path | str | None) -> Path: + if cache_dir is not None: + path = Path(cache_dir) + elif os.getenv("TT_CACHE_PATH"): + path = Path(os.environ["TT_CACHE_PATH"]) + else: + path = Path("model_cache") / hf_model / "T3K" + path.mkdir(parents=True, exist_ok=True) + return path + + +def _validate_checkpoint_config(hf_config, *, num_devices: int) -> None: + expected = { + "num_hidden_layers": 80, + "hidden_size": 8192, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "intermediate_size": 29568, + "vocab_size": 152064, + } + actual = {name: getattr(hf_config, name, None) for name in expected} + mismatches = {name: (actual[name], value) for name, value in expected.items() if actual[name] != value} + if mismatches: + raise ValueError(f"Unexpected Qwen2.5-72B geometry: {mismatches}") + if hf_config.num_attention_heads % num_devices or hf_config.num_key_value_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({hf_config.num_attention_heads}/{hf_config.num_key_value_heads}) " + f"must be divisible by device count ({num_devices})" + ) + if bool(getattr(hf_config, "tie_word_embeddings", False)): + raise ValueError("Qwen2.5-72B requires an untied LM head") + + +def convert_hf_model_weights( + hf, + hf_config, + *, + n_layers: int, + num_devices: int, + rope_table_len: int, + head_dim: int, +) -> Qwen25_72BWeights: + base = hf.model + rope_cos, rope_sin = weight_utils.build_rope_cos_sin_torch( + base.rotary_emb, rope_table_len, head_dim, torch.bfloat16 + ) + layers = [] + for layer in base.layers[:n_layers]: + wqkv, wo, q_norm, k_norm, wqkv_bias = weight_utils.attention_wqkv_wo_from_hf_layer(layer.self_attn, num_devices) + if q_norm is not None or k_norm is not None: + raise ValueError("Qwen2.5-72B does not support QK norm weights") + if wqkv_bias is None: + raise ValueError("Qwen2.5-72B requires QKV projection bias") + w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(layer.mlp) + ff_align = 32 * 32 * num_devices + ff_padded = math.ceil(w1.shape[-1] / ff_align) * ff_align + if ff_padded != w1.shape[-1]: + w1 = torch.nn.functional.pad(w1, (0, ff_padded - w1.shape[-1])) + w3 = torch.nn.functional.pad(w3, (0, ff_padded - w3.shape[-1])) + w2 = torch.nn.functional.pad(w2, (0, 0, 0, ff_padded - w2.shape[-2])) + layers.append( + Qwen25_72BLayerWeights( + wqkv=wqkv, + wo=wo, + wqkv_bias=wqkv_bias.to(torch.bfloat16), + w1=w1, + w2=w2, + w3=w3, + attention_norm=weight_utils.rms_weight_torch(layer.input_layernorm).to(torch.bfloat16), + ff_norm=weight_utils.rms_weight_torch(layer.post_attention_layernorm).to(torch.bfloat16), + ) + ) + return Qwen25_72BWeights( + embedding=weight_utils.embed_tokens_torch(base.embed_tokens), + rope_cos=rope_cos, + rope_sin=rope_sin, + layers=tuple(layers), + final_norm=weight_utils.rms_weight_torch(base.norm).to(torch.bfloat16), + lm_head=hf.lm_head.weight.detach().to(torch.bfloat16).clone(), + ) + + +def from_pretrained( + mesh_device, + *, + hf_model: str = DEFAULT_HF_MODEL, + instruct: bool = True, + max_batch_size: int = 32, + max_seq_len: int = 4096, + optimizations: str | Qwen25_72BPrecisionConfig = "accuracy", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config: Qwen25_72BPagedAttentionConfig | None = None, + cache_dir: Path | str | None = None, +) -> Qwen25_72BForCausalLM: + del dtype + num_devices = mesh_device.get_num_devices() + if num_devices != 8: + raise ValueError(f"Qwen2.5-72B supports exactly 8 devices (T3K), got {num_devices}") + ttnn.SetDefaultDevice(mesh_device) + load_kwargs = { + "revision": DEFAULT_HF_REVISION, + "trust_remote_code": True, + "local_files_only": _local_files_only(), + } + hf_config = AutoConfig.from_pretrained(hf_model, **load_kwargs) + _validate_checkpoint_config(hf_config, num_devices=num_devices) + hf = AutoModelForCausalLM.from_pretrained(hf_model, torch_dtype=torch.bfloat16, **load_kwargs) + hf.eval() + resolved_layers = hf_config.num_hidden_layers if n_layers is None else n_layers + precision = ( + optimizations + if isinstance(optimizations, Qwen25_72BPrecisionConfig) + else (QWEN25_72B_PERFORMANCE if optimizations == "performance" else QWEN25_72B_ACCURACY) + ) + if not isinstance(precision, Qwen25_72BPrecisionConfig): + raise TypeError("optimizations must be 'accuracy', 'performance', or Qwen25_72BPrecisionConfig") + cache_path = _cache_path(hf_model, mesh_device, cache_dir) + block_size = 32 + if paged_attention_config is None: + paged_attention_config = Qwen25_72BPagedAttentionConfig( + block_size=block_size, + max_num_blocks=((max_seq_len + block_size - 1) // block_size) * max_batch_size, + ) + head_dim = int(getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)) + rope_table_len = ((max(max_seq_len * 2, 8192) + 127) // 128) * 128 + config = Qwen25_72BConfig( + hf_model_id=hf_model, + dim=hf_config.hidden_size, + n_heads=hf_config.num_attention_heads, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + hidden_dim=hf_config.intermediate_size, + vocab_size=hf_config.vocab_size, + rms_norm_eps=hf_config.rms_norm_eps, + rope_theta=getattr(hf_config, "rope_theta", 1_000_000.0), + num_hidden_layers=resolved_layers, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + rope_table_len=rope_table_len, + ) + weights = convert_hf_model_weights( + hf, + hf_config, + n_layers=resolved_layers, + num_devices=num_devices, + rope_table_len=rope_table_len, + head_dim=head_dim, + ) + model = build_qwen25_72b_model( + mesh_device=mesh_device, + config=config, + weights=weights, + precision=precision, + cache_path=cache_path, + paged_attention_config=paged_attention_config, + ) + tokenizer = load_tokenizer(hf_model, DEFAULT_HF_REVISION) + stop_token_ids = _qwen_stop_token_ids(tokenizer) + runtime_config = Qwen25_72BRuntimeConfig( + model_name=Path(hf_model).name, + model_cache_path=cache_path, + max_prefill_chunk_size=2048, + max_context_len=int(hf_config.max_position_embeddings), + max_seq_len=max_seq_len, + trace_prefill_supported_seq_lens=_trace_seq_lens(num_devices, 2048, max_seq_len), + n_layers=model.config.n_layers, + n_kv_heads=model.config.n_kv_heads, + head_dim=model.config.head_dim, + max_batch_size=max_batch_size, + cluster_shape=list(mesh_device.shape), + kv_cache_dtype=precision.kv_cache_dtype, + disable_batched_prefill=bool(os.getenv("DISABLE_BATCHED_PREFILL")), + batched_prefill_batched_extract=not bool(os.getenv("DISABLE_BATCHED_EXTRACT")), + ) + del hf + return Qwen25_72BForCausalLM( + model=model, + tokenizer=tokenizer, + runtime_config=runtime_config, + instruct=instruct, + generation_config=Qwen25_72BGenerationConfig(stop_token_ids=stop_token_ids), + ) diff --git a/code/models/common/models/qwen25_72b/model.py b/code/models/common/models/qwen25_72b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..88376931fb98218de91614f0b94cf3c9064ffc0a --- /dev/null +++ b/code/models/common/models/qwen25_72b/model.py @@ -0,0 +1,1119 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Qwen2.5-72B-Instruct — native TTTv2 stack (``Embedding1D``, ``RMSNorm1D``, +``Attention1D``, ``MLP1D``, ``RotarySetup1D``, ``LMHead1D``). Targets T3K (mesh ``(1, 8)``). + +Architecturally identical to the Qwen2.5-7B / Qwen2.5-Coder-32B ports (Qwen2 family: +QKV bias, no q/k RMSNorm, GQA, SwiGLU MLP, GPT-NeoX RoPE); only the scale differs +(80 layers, hidden 8192, 64 attention heads, 8 KV heads, intermediate 29568, vocab 152064). + +Tensor layout contracts: + - **Prefill** hidden states: ``[1, 1, S, dim]`` TILE, ``S % 128 == 0``. + - **Decode** hidden states: ``[1, 1, B, dim]`` TILE (``B`` padded to tile in modules). + +Executor contract: pre-embedded forwards, +``set_kv_cache``, ``rope_setup``, ``page_table`` through attention, ``model_args`` holds a +:class:`Qwen25_72BExecutorRuntimeConfig` (not v1 ``ModelArgs``). +""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field, replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any, List + +import torch +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.models.qwen25_72b import weight_utils +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D, prepare_rot_idxs +from models.common.modules.sampling.sampling_1d import Sampling1D +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + +# Pinned HF revision SHA for Qwen/Qwen2.5-72B-Instruct (resolved 2026-06-05). +DEFAULT_HF_REVISION = "495f39366efef23836d0cfae4fbe635880d2be31" + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + """Minimal LazyWeight; ``Attention1D`` / ``MLP1D`` / ``Embedding1D`` resolvers fill mesh + memory.""" + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +@dataclass +class Qwen25_72BPagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (``block_size`` / ``max_num_blocks`` only).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Qwen25_72BExecutorRuntimeConfig: + """Engine-facing runtime knobs. Exposed as ``model.model_args`` for the model-owned executor.""" + + n_layers: int + n_kv_heads: int + head_dim: int + max_batch_size: int + max_seq_len: int + cluster_shape: list[int] + model_name: str = "Qwen2.5-72B-Instruct" + max_prefill_chunk_size: int = 2048 + max_context_len: int = 32768 + model_cache_path: Path | None = None + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + optimizations: Any = None + trace_prefill_supported_seq_lens: tuple[int, ...] = (128, 1024) + # Batched prefill (parity caveat #12): fuse equal-length users into batched passes to close the + # batch-32 TTFT gap. ``supports_batched_prefill`` is the per-model opt-in (the shared engine only + # batches models whose prefill_forward threads ``batch_size`` — Qwen2.5-72B does, below). + # Qwen2.5-72B is a standard dense Qwen2.5 attention (NO QK-norm), so every prefill op is + # row-independent and the batched fold is bit-safe (same as the qwen25_7b / Coder-32B ports). + # ``max_prefill_batch_size`` is the largest supported padded wave; 32 folds batch-32 prefill in ONE + # 32-user pass (TTTv1 structural parity) so the eager norm+lm_head tail + full-vocab readback run + # once instead of 4×. At the natural 128 bucket the fold is 32*128=4096=2*2048, an exact multiple of + # MAX_QKV_MM_SEQ_LEN (reshape-safe), and 4096 % mlp_prefill_len_cutoff(1024) == 0 for the FF reshape; + # the DRAM guard (padded_batch*seq < 128K) passes with 4096. ``disable_batched_prefill`` is the escape + # hatch back to the sequential loop; ``max_prefill_chunk_size`` (above) drives the #45234 + # chunked-prompt decline. + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + # When True (default), batched prefill runs norm+lm_head ONCE per group over the gathered last-token + # rows (TTTv1 parity); False falls back to the bit-identical per-slot path (one lm_head per user). + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + # Mirror TTTv1's prefill-trace gate (model_config.get_trace_prefill_supported_seq_lens): + # only trace the seq lens TTTv1 lists -- bigger seq lens already have small op2op gaps, so + # tracing buys nothing. Qwen2.5-72B has NO model-specific entry in TTTv1, so it falls to the + # family default: T3K -> [128, 1024] (N150 -> [128]; N300/T3K/TG -> [128, 1024]). NB: the + # same-topology Llama-3.3-70B IS restricted to [128] on T3K via a model-specific override, + # but Qwen2.5-72B has no such override, so the default [128, 1024] applies. The earlier + # hardcoded ``return False`` (claimed TT_FATAL under LazyWeight + distributed norms) is + # superseded -- the 1B/3B ports proved that TT_FATAL no longer reproduces. Decode trace + # remains enabled at the engine layer regardless. + num_devices = int(self.cluster_shape[0]) * int(self.cluster_shape[1]) + allowed = self.trace_prefill_supported_seq_lens or {1: (128,), 2: (128, 1024), 8: (128, 1024)}.get( + num_devices, (128,) + ) + return ( + prefill_seq_len in allowed + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +@dataclass +class Qwen25_72BConfig: + """Resolved hyper-parameters for a loaded HF Qwen2.5-72B checkpoint.""" + + hf_model_id: str + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + rope_theta: float + num_hidden_layers: int + max_batch_size: int + max_seq_len: int + rope_table_len: int + num_devices: int = 8 + mesh_device: ttnn.MeshDevice | None = None + n_layers: int | None = None + block_configs: list[Any] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.n_layers is None: + self.n_layers = self.num_hidden_layers + + +@dataclass(frozen=True) +class Qwen25_72BLayerWeights: + wqkv: torch.Tensor + wo: torch.Tensor + wqkv_bias: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Qwen25_72BWeights: + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Qwen25_72BLayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +_LOFI_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, +) +"""LoFi + packer L1 acc for the MLP FF1/FF3 matmuls. + +For Qwen2.5-72B this is used in **both** accuracy and performance recipes — see +:class:`Qwen25_72BPrecisionConfig`. Mirrors TTTv1 ``DecodersPrecision`` which routes the +72B model through the ``model_config.py:119`` ">70B" special case (``FF1_FF3 → BFP4`` and +``LI_FF1_FF3 → LOFI``) in accuracy mode, identical to the generic performance recipe. +""" + + +@dataclass(frozen=True) +class Qwen25_72BPrecisionConfig: + """Per-layer precision + math-fidelity recipe for Qwen2.5-72B-Instruct on T3K. + + Mirrors the fields TTTv1's ``DecodersPrecision`` distinguishes for Qwen2.5-72B. The base + model name resolves to ``Qwen2.5-72B`` (via ``common.get_base_model_name``), which lands in + the **">70B" special case** at ``model_config.py:119``: large models are precision-insensitive, + so TTTv1 ships BFP4 MLPs + BFP8 attention *even in accuracy mode*. The consequence is that for + this model the accuracy and performance recipes are **identical**: + + * **Accuracy** (``model_config.py:118-128``): ``FF1_FF3 → BFP4`` + ``LI_FF1_FF3 → LOFI``; + everything else stays at TTTv1 defaults (BFP8 ``WQKV`` / ``WO`` / ``KV_CACHE`` + HIFI2 + attention; SDPA prefill HIFI4 per ``_default_settings``). + * **Performance** (``model_config.py:208-218``, generic non-7B branch): the same + ``FF1_FF3 → BFP4`` + ``LI_FF1_FF3 → LOFI``; everything else at defaults. + + Two module-level recipes are exposed for symmetry with the other ports — + :data:`QWEN25_72B_ACCURACY` (default) and :data:`QWEN25_72B_PERFORMANCE` — but they carry the + **same** field values here. Pass one to :meth:`Qwen25_72B.from_pretrained` via ``precision=``; + use ``dataclasses.replace(QWEN25_72B_ACCURACY, ...)`` to customize a single field. + + LM head is BFP8 in both modes. The bf16 "accuracy tightening" used on Mistral-7B / Coder-32B + is **not** portable to large Qwen ports: on Qwen3-32B it regressed accuracy-mode token-accuracy + (the bf16 DRAM-sharded decode LM-head matmul over the zero-padded vocab corrupts the logit + ranking). Match TTTv1 — ship BFP8. + + The ``mlp_w2_dtype`` / ``mlp_ff2_compute_kernel_cfg`` fields are absent because TTTv1 leaves + them at engine defaults (``BFP8`` / ``HIFI2_FP16``) in both recipes, and those defaults + coincide with the TTTv2 ``MLP1D`` defaults. + """ + + # Attention weight / KV-cache dtypes. BFP8 in both modes (the ">70B" special case keeps + # BFP8 attention even in accuracy mode — only the MLP drops to BFP4). + wqkv_dtype: ttnn.DataType = ttnn.bfloat8_b + wo_dtype: ttnn.DataType = ttnn.bfloat8_b + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + # MLP FF1/FF3 weight dtype — BFP4 in both modes. + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat4_b + + # MLP FF1/FF3 compute kernel — LOFI in both modes. + mlp_ff1_3_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = _LOFI_COMPUTE_KERNEL_CFG + + # Attention compute kernels left at the Attention1D default (HIFI2 + fp16 dest acc, SDPA + # prefill HIFI4) in both modes — TTTv1 makes no attention-fidelity override for 72B. + attn_li_qkv_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_sdpa_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_li_o_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + # BFP8 LM head in both modes (matches TTTv1; bf16 tightening regresses large Qwen ports). + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + + +# TTTv1 ``DecodersPrecision.accuracy("Qwen2.5-72B-Instruct")`` (``model_config.py:118-128``): +# the ">70B" special case — BFP4 FF1/FF3 + LOFI, BFP8 attention + BFP8 KV cache. Identical to the +# performance recipe for this model (see :class:`Qwen25_72BPrecisionConfig`). +QWEN25_72B_ACCURACY = Qwen25_72BPrecisionConfig() + +# TTTv1 ``DecodersPrecision.performance("Qwen2.5-72B-Instruct")`` (``model_config.py:208-218``): +# FF1_FF3 → BFP4 + LI_FF1_FF3 → LOFI; everything else at defaults. Same field values as accuracy. +QWEN25_72B_PERFORMANCE = Qwen25_72BPrecisionConfig() + + +def _slice_last_token_tile(x: ttnn.Tensor, last_token_idx: int) -> ttnn.Tensor: + """Slice the 32-row tile containing ``last_token_idx`` from ``[1, 1, S, W]``. + + Width-sharded LM matmul M tile rows must equal ``LMHead1D`` program-config tile rows. + """ + floor = (last_token_idx // 32) * 32 + return ttnn.slice(x, (0, 0, floor, 0), (1, 1, floor + 32, x.shape[-1])) + + +def _post_attn_norm_decode_configs( + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding so its output matches MLP1D's W1/W3 input. + + MLP1D decode uses ``_dram_shard_core_grid_k_n(dim, padded_hidden / num_devices)`` for W1/W3 + inputs, but the default RMSNorm program config is derived from ``_compute_norm_core_grid(dim)`` + alone. Mismatched DRAM-width-shard between RMSNorm output and MLP1D W1/W3 input silently + corrupts decode activations (observed on the 7B / Coder-32B ports — same shape pattern on T3K). + """ + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + memory_config = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return program_config, memory_config + + +# Cast the three per-layer prefill all-gathers (the two DistributedNorm activation reconstructions here +# + the pre-WO gather via Attention1DConfig.prefill_ag_ccl_dtype) from bf16 to bfloat8_b. These are the +# largest CCL cost in the batched b32-ci prefill and a bf16 all-gather moves 2× the bytes of a bf8_b one, +# while every matmul they feed (QKV/WO bf8_b, W1/W3 bf4_b) quantizes the activation further regardless — +# so the bf16 gather precision is largely wasted. PREFILL-ONLY: the decode call sites of +# _all_gather_rmsnorm_tensor pass no dtype (byte-identical decode). None disables (byte-identical). +# Qwen2.5-72B is precision-tolerant (the ">70B" recipe ships bf4 MLP + bf8 attention even in accuracy +# mode); gated by token-accuracy + eval-32, not bit-exactness (see PLAN_03's prefill_reduce_ccl_dtype). +_PREFILL_AG_CCL_DTYPE: ttnn.DataType | None = ttnn.bfloat8_b + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, + x: ttnn.Tensor, + *, + memory_config: ttnn.MemoryConfig | None = None, + dtype: ttnn.DataType | None = None, +) -> ttnn.Tensor: + cfg = norm.config + requested_memory_config = memory_config + if memory_config is None: + memory_config = x.memory_config() + + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + if requested_memory_config is not None: + return ttnn.to_memory_config(x, memory_config) + if x.memory_config() != memory_config: + return ttnn.to_memory_config(x, memory_config) + return x + + # Prefill-only opt-in (``dtype`` set only at the prefill call sites): cast the gather input (bf16) to + # a smaller CCL dtype (bfloat8_b) to halve this collective's cross-device payload. Decode call sites + # leave ``dtype=None`` → byte-identical. An explicit typecast is required (to_memory_config does not + # cast an already-DRAM tensor). + if dtype is not None and x.dtype != dtype: + x_cast = ttnn.typecast(x, dtype) + ttnn.deallocate(x) + x = x_cast + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + + +@dataclass +class _Qwen25_72BWHTuning: + """Wormhole-specific L1 / cutoff tuning resolved at build time. + + Precision-vs-fidelity knobs live on :class:`Qwen25_72BPrecisionConfig`. This dataclass only + carries non-mode-dependent L1 footprint controls: MLP prefill cutoff and the optional + W1→DRAM spill on decode. + """ + + mlp_prefill_len_cutoff: int | None = None + mlp_decode_spill_w1_to_dram: bool = False + # Use ttnn.experimental.minimal_matmul for the QKV + W2 prefill matmuls above seq_len > 128 (TTTv1 + # parity, PLAN_01). A/B escape hatch: set DISABLE_MINIMAL_MATMUL=1 to force ttnn.linear. On a 72B the + # batch-32-ci prefill is matmul-compute-bound (~80% of FLOPs = the 3 MLP matmuls), so minimal_matmul + # is a real prefill-TTFT win — it narrows the batch-32-ci TTFT gap vs TTTv1 (which uses minimal_matmul + # for the same matmuls). The shared plumbing (attention_1d.use_minimal_qkv_matmul / mlp_1d + # use_minimal_w2_matmul, both gated seq_len>128) is already in the base; this flag just engages it. + # Qwen2.5-72B is a Qwen2.5 arch: the QKV bias is added AFTER the matmul, so it is unchanged by the + # minimal path (same as the mistral_7b / Coder-32B / deepseek Qwen2 ports). Independent of the + # FF-hidden DRAM-shard pad (a decode fix); minimal_matmul is prefill-only, so decode is unchanged. + prefill_minimal_matmul: bool = True + + +def _resolve_qwen_72b_wh_tuning(*, num_dev: int, max_batch_size: int) -> _Qwen25_72BWHTuning: + """Pick WH L1 tuning knobs for Qwen2.5-72B-Instruct on T3K. + + ``mlp_prefill_len_cutoff=1024`` = the shared engine's own Wormhole default (``mlp_1d.py``: + ``512 if is_blackhole() else 1024``) and TTTv1's ``prefill_len_cutoff``. For the folded batch-32-ci + FF prefill (``[1,1,B*S,dim]``, B*S=32*128=4096 at the natural 128 bucket) this tiles the wide FF + matmul as 4 chunks of 1024 (``per_core_M=4``) instead of 16 chunks of 256 (``per_core_M=1``) — 4× + fewer / 4× larger sub-matmuls on the ~80%-FLOP FF block (better weight-mcast amortization, fewer + device inter-op boundaries), matching TTTv1's blocking. The earlier 256 was inherited-conservative + from the 7B-on-N300 port (per-device FF shard 9472 ≫ this model's T3K shard 3696), so its tighter-L1 + motive does not apply here; 1024 fits — it is proven on the same-topology Llama-3.3-70B (shard 3584) + and the T3K 32B (3456) on this box, and TTTv1 runs 1024 for this exact model. Output is unchanged: + ``in0_block_w`` and the K-contraction order are independent of the M-tiling, so only ``per_core_M`` + changes. ``mlp_decode_spill_w1_to_dram`` starts off; re-evaluate if decode batch-32 trips L1 + circular-buffer validation (per-device FF shard is 8192×3696 in BFP4) — decode is unaffected by the + cutoff (it uses the separate DRAM-sharded prg configs, which never read ``prefill_len_cutoff``). + """ + t = _Qwen25_72BWHTuning( + mlp_prefill_len_cutoff=1024, + mlp_decode_spill_w1_to_dram=False, + ) + t.prefill_minimal_matmul = not os.environ.get("DISABLE_MINIMAL_MATMUL") + logger.info( + f"L1 tuning for Qwen2.5-72B on {num_dev} device(s): " + f"prefill_len_cutoff={t.mlp_prefill_len_cutoff}, " + f"decode_spill_w1_to_dram={t.mlp_decode_spill_w1_to_dram}, " + f"prefill_minimal_matmul={t.prefill_minimal_matmul}" + ) + return t + + +def _build_decoder_layer( + *, + idx: int, + weights: Qwen25_72BLayerWeights, + qcfg: Qwen25_72BConfig, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Qwen25_72BPrecisionConfig, + executor_mode: bool, + paged_cfg: Qwen25_72BPagedAttentionConfig | None, + cache_path: Path | None, + wh: _Qwen25_72BWHTuning, +) -> Qwen25_72BDecoderLayer: + """Construct one decoder layer from provider-neutral tensors.""" + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, dtype=precision.wqkv_dtype, cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None + ) + lazy_wo = _lazy( + weights.wo, dtype=precision.wo_dtype, cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None + ) + + bias_lw = LazyWeight( + source=weights.wqkv_bias, + dtype=ttnn.bfloat16, + cache_dir_weight_name=(cache_path / "attn", f"{prefix}_bias") if cache_path else None, + ) + + attn = Attention1D.from_config( + Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=qcfg.n_heads, + n_kv_heads=qcfg.n_kv_heads, + head_dim=qcfg.head_dim, + max_batch_size=qcfg.max_batch_size, + max_seq_len=qcfg.max_seq_len, + q_norm_config=None, + k_norm_config=None, + wqkv_bias=bias_lw, + use_vllm_paged_kv_cache=executor_mode, + paged_attention_config=paged_cfg, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + li_qkv_prefill_compute_kernel_cfg=precision.attn_li_qkv_kernel_cfg, + li_qkv_decode_compute_kernel_cfg=precision.attn_li_qkv_kernel_cfg, + sdpa_decode_compute_kernel_cfg=precision.attn_sdpa_kernel_cfg, + li_o_prefill_compute_kernel_cfg=precision.attn_li_o_kernel_cfg, + li_o_decode_compute_kernel_cfg=precision.attn_li_o_kernel_cfg, + prefill_qkv_minimal_matmul=wh.prefill_minimal_matmul, + # Route the folded-batch prefill WO projection through minimal_matmul (completes PLAN_01's + # QKV+FF2 minimal plumbing). WO is the least-efficient prefill matmul on ttnn.linear; the + # minimal op (already used for QKV above) recovers most of that gap and makes TTTv2 beat + # TTTv1's ttnn.linear WO. Gated by the same DISABLE_MINIMAL_MATMUL escape hatch. + prefill_wo_minimal_matmul=wh.prefill_minimal_matmul, + # WO-matmul prefill M-chunk cutoff: regroup the folded batch-32 prefill WO matmul into + # 2 chunks of 2048 (per_core_M=8) instead of 4 chunks of 1024, halving the WO weight + # re-stream passes on the folded prefill. Bit-identical M-reblocking (see field doc). + wo_prefill_len_cutoff=2048, + # Cast the pre-WO fused prefill all-gather from bf16 to bfloat8_b (halves this per-layer + # collective's cross-device payload). Paired with the two norm-reconstruction gathers cast + # via _PREFILL_AG_CCL_DTYPE in prefill_forward. Prefill-only (the fused gather is a + # prefill-path method); decode is unaffected. See _PREFILL_AG_CCL_DTYPE for rationale. + prefill_ag_ccl_dtype=_PREFILL_AG_CCL_DTYPE, + ) + ) + + w1, w2, w3 = weights.w1, weights.w2, weights.w3 + # Pad the FF hidden dim to a grid-friendly per-device size so the DRAM-sharded decode FF + # matmuls (W1/W3/W2) use a full 64-core grid instead of 4. The decode grid must divide both + # the K-tile and N-tile counts (in0 is K-width-sharded, weights are N-sharded); with the raw + # per-device hidden 29568/8 -> 3712 = 116 tiles (2^2*29), gcd(dim_tiles=256, 116)=4 -> only 4 + # DRAM readers stream the FF weights, which dominate memory-bound decode. Padding per device to + # 128 tiles (4096, total 32768) gives gcd(256,128)=64 cores -- matching TTTv1's + # default_padded_cores=32 (nearest_multiple(29568, 32*32*8)=32768). The extra columns are zeros + # (silu(0)=0, mul->0, and W2 contracts the padded rows to 0), so decode output is unchanged. + _ff_align = TILE_SIZE * TILE_SIZE * num_dev # 32*32*8 = 8192 on T3K + _ff_pad = math.ceil(w1.shape[-1] / _ff_align) * _ff_align + if _ff_pad != w1.shape[-1]: + w1 = torch.nn.functional.pad(w1, (0, _ff_pad - w1.shape[-1])) + w3 = torch.nn.functional.pad(w3, (0, _ff_pad - w3.shape[-1])) + w2 = torch.nn.functional.pad(w2, (0, 0, 0, _ff_pad - w2.shape[-2])) + mlp = MLP1D.from_config( + MLP1DConfig( + w1=_lazy( + w1, dtype=precision.mlp_w1_w3_dtype, cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None + ), + w2=_lazy(w2, dtype=ttnn.bfloat8_b, cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None), + w3=_lazy( + w3, dtype=precision.mlp_w1_w3_dtype, cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=qcfg.max_batch_size, + prefill_len_cutoff=wh.mlp_prefill_len_cutoff, + w1_w3_dtype=precision.mlp_w1_w3_dtype, + w2_dtype=ttnn.bfloat8_b, + ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_spill_w1_to_dram_before_w3=wh.mlp_decode_spill_w1_to_dram, + prefill_w2_minimal_matmul=wh.prefill_minimal_matmul, + ) + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + dim=qcfg.dim, + # Use the padded FF hidden so the post-attn RMSNorm decode output is width-sharded on the + # SAME (64-core) grid as MLP1D's W1/W3 decode input; a mismatch silently corrupts decode. + hidden_dim=_ff_pad, + num_devices=num_dev, + max_batch_size=qcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1D: + lw = _lazy( + weight, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1D.from_config( + RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=qcfg.rms_norm_eps, + max_batch_size=qcfg.max_batch_size, + tt_ccl=tt_ccl, + **extra, + ) + ) + + return Qwen25_72BDecoderLayer( + input_layernorm=_build_norm(weights.attention_norm, "pre_attn"), + self_attn=attn, + post_attention_layernorm=_build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ), + mlp=mlp, + ) + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + qcfg: Qwen25_72BConfig, + lm_head_dtype: ttnn.DataType, + cache_path: Path | None, +) -> LMHead1D: + """Build the vocab-sharded LM head with DRAM-matmul program configs. + + LM head DRAM matmul is sized for decode batch tiles (``max_batch_size``). Prefill logits + use a single 32-row tile via ``post_process_prefill_output`` / :func:`_slice_last_token_tile`. + """ + lm_splits, lm_split_sizes, lm_weights_memcfgs = weight_utils.build_lm_head_lazy_weights( + mesh_device, + lm_head_weight, + dim=qcfg.dim, + vocab_size=qcfg.vocab_size, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(qcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(qcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, qcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(qcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1D.from_config( + LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=qcfg.dim, + max_batch_size=qcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + ) + ) + + +def build_qwen25_72b_model( + *, + mesh_device: ttnn.MeshDevice, + config: Qwen25_72BConfig, + weights: Qwen25_72BWeights, + precision: Qwen25_72BPrecisionConfig, + cache_path: Path | None, + paged_attention_config: Qwen25_72BPagedAttentionConfig | None, +) -> Qwen25_72B: + """Build the TT tensor graph from provider-neutral Qwen2.5-72B dimensions and tensors.""" + + ttnn.SetDefaultDevice(mesh_device) + num_devices = mesh_device.get_num_devices() + if num_devices != 8: + raise ValueError( + f"Qwen2.5-72B-Instruct port targets T3K (mesh (1, 8) = 8 devices) only. " + f"Got mesh_device with {num_devices} device(s). Open a T3K mesh with MESH_DEVICE=T3K." + ) + if config.n_heads % num_devices != 0 or config.n_kv_heads % num_devices != 0: + raise ValueError( + f"Checkpoint heads ({config.n_heads}/{config.n_kv_heads}) must be divisible by " + f"device count ({num_devices})" + ) + if len(weights.layers) != config.num_hidden_layers: + raise ValueError(f"Expected {config.num_hidden_layers} decoder layer weights, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) + topology = default_topology(mesh_device) + emb = Embedding1D.from_config( + Embedding1DConfig( + weights=_lazy( + weights.embedding, + dtype=ttnn.bfloat16, + cache=(cache_path / "embedding", "tok_embeddings") if cache_path else None, + ), + mesh_device=mesh_device, + embed_scale=1.0, + ) + ) + rope_setup = RotarySetup1D.from_config( + Rope1DConfig( + cos_matrix=_lazy( + weights.rope_cos, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "cos") if cache_path else None + ), + sin_matrix=_lazy( + weights.rope_sin, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "sin") if cache_path else None + ), + max_batch_size=config.max_batch_size, + head_dim=config.head_dim, + device=mesh_device, + use_qk_fused=False, + ) + ) + wh = _resolve_qwen_72b_wh_tuning(num_dev=num_devices, max_batch_size=config.max_batch_size) + layers = [ + _build_decoder_layer( + idx=idx, + weights=weights.layers[idx], + qcfg=config, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + executor_mode=paged_attention_config is not None, + paged_cfg=paged_attention_config, + cache_path=cache_path, + wh=wh, + ) + for idx in range(config.num_hidden_layers) + ] + final_norm = RMSNorm1D.from_config( + RMSNorm1DConfig( + weight=_lazy( + weights.final_norm, dtype=ttnn.bfloat16, cache=(cache_path / "norm", "final") if cache_path else None + ), + mesh_device=mesh_device, + eps=config.rms_norm_eps, + max_batch_size=config.max_batch_size, + tt_ccl=tt_ccl, + ) + ) + lm_head = _build_lm_head( + mesh_device=mesh_device, + lm_head_weight=weights.lm_head, + qcfg=config, + lm_head_dtype=precision.lm_head_dtype, + cache_path=cache_path, + ) + return Qwen25_72B(config, emb, rope_setup, layers, final_norm, lm_head, mesh_device) + + +class Qwen25_72BDecoderLayer(LightweightModule): + def __init__( + self, + *, + input_layernorm: RMSNorm1D, + self_attn: Attention1D, + post_attention_layernorm: RMSNorm1D, + mlp: MLP1D, + ): + super().__init__() + self.input_layernorm = input_layernorm + self.self_attn = self_attn + self.post_attention_layernorm = post_attention_layernorm + self.mlp = mlp + self.attention_norm = input_layernorm + self.attention = self_attn + self.ff_norm = post_attention_layernorm + self.feed_forward = mlp + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + *, + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + # Match Llama ``TransformerBlock1D``: fractured embed / norm activations must be + # all-gathered to full ``dim`` before Attention1D / MLP1D (QKV matmul expects width ``dim``). + r = self.input_layernorm.prefill_forward(x) + r = _all_gather_rmsnorm_tensor(self.input_layernorm, r, dtype=_PREFILL_AG_CCL_DTYPE) + r = self.self_attn.prefill_forward( + r, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + r2 = self.post_attention_layernorm.prefill_forward(h) + r2 = _all_gather_rmsnorm_tensor(self.post_attention_layernorm, r2, dtype=_PREFILL_AG_CCL_DTYPE) + r2 = self.mlp.prefill_forward(r2) + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def decode_forward( + self, + x: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + xa = _all_gather_rmsnorm_tensor( + self.input_layernorm, x, memory_config=self.input_layernorm.config.decode_memory_config + ) + r = self.input_layernorm.forward(xa, "decode") + r = self.self_attn.forward(r, current_pos, rot_mats, mode="decode", page_table=page_table) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + hf = _all_gather_rmsnorm_tensor( + self.post_attention_layernorm, h, memory_config=self.post_attention_layernorm.config.decode_memory_config + ) + r2 = self.post_attention_layernorm.forward(hf, "decode") + r2 = self.mlp.forward(r2, "decode") + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + +class Qwen25_72B(LightweightModule): + """ + Full decoder for Qwen2.5-72B-Instruct (TTTv2 modules only) on T3K. + + Prefill/decode on **embedded** activations match the model-owned executor. Token embedding + is ``embed_prefill`` / ``embed_decode``. Bind KV with ``set_kv_cache`` before first forward. + """ + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__( + self, + cfg: Qwen25_72BConfig, + embed: Embedding1D, + rope_setup: RotarySetup1D, + layers: List[Qwen25_72BDecoderLayer], + norm: RMSNorm1D, + lm_head: LMHead1D, + mesh_device: ttnn.MeshDevice, + ): + super().__init__() + self.cfg = cfg + self.config = cfg + self.config.mesh_device = mesh_device + self.config.n_layers = cfg.num_hidden_layers + self.config.num_devices = mesh_device.get_num_devices() + self.embed = embed + self.embedding = self.embed + self.rope_setup = rope_setup + self.layers = layers + self.norm = norm + self.lm_head = lm_head + self.mesh_device = mesh_device + self.model_args: Qwen25_72BExecutorRuntimeConfig | None = None + + self.vocab_size = cfg.vocab_size + self.n_layers = cfg.num_hidden_layers + self.num_devices = mesh_device.get_num_devices() + self.tt_ccl = get_tt_ccl(mesh_device) if self.num_devices > 1 else None + self.config.num_devices = self.num_devices + self.config.block_configs = [SimpleNamespace(attention_config=layer.self_attn.config) for layer in self.layers] + + # The model owns its sampler; callers pick behavior per request via sampling_params. + # Buffers are lazy (nothing materializes until the first on-device sampled decode), so + # this is harmless when sampling_params is None (host-argmax path, the demo default). + # Qwen2.5-72B targets T3K (8 devices); on-device top-k wins at >=8-device meshes (vocab + # shards 8 ways so per-device ttnn.topk is cheap), see the rebase handoff crossover table. + self.supports_on_device_sampling = self.num_devices >= 1 + self.sampling = ( + Sampling1D( + vocab_size=self.vocab_size, + mesh_device=mesh_device, + tt_ccl=self.tt_ccl, + max_batch_size=_nearest_32(cfg.max_batch_size), + # Qwen72's LM head exposes 128 padded rows, beyond the top-k sampling + # kernel's 32-user limit. Route greedy requests through force-argmax. + allow_force_argmax=True, + pad_to_power_of_2=True, + ) + if self.supports_on_device_sampling + else None + ) + + @property + def n_kv_heads(self) -> int: + return self.cfg.n_kv_heads + + def iter_executor_named_modules(self): + if not hasattr(self, "layers"): + return + for index, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{index}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in getattr(self.config, "block_configs", ())) + for config in tuple({id(item): item for item in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, + block_size=block_size, + max_num_blocks=max_num_blocks, + ) + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for index, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {index} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {index} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = pair + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). The batched path always + # returns the full hidden state (get_last_token == -1); the executor does per-slot last-token + # extraction + norm/lm_head so those stages stay bit-identical to the single-user path. + x = x_embed + for layer in self.layers: + x = layer.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + + if get_last_token == -1: + return x + + # Slice + deallocate the full-sequence buffer before norm/LM head reduces peak L1. + old = x + if last_token_slice is None: + x_tile = _slice_last_token_tile(old, get_last_token) + else: + x_tile = ttnn.slice( + old, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(old.shape[2]) // 32, + ) + ttnn.deallocate(old) + if last_token_index is not None: + if x_tile.dtype != ttnn.bfloat16: + old_tile = x_tile + x_tile = ttnn.typecast(x_tile, ttnn.bfloat16) + ttnn.deallocate(old_tile) + old_tile = x_tile + x_tile = ttnn.embedding(last_token_index, x_tile, layout=ttnn.TILE_LAYOUT) + x_tile = ttnn.unsqueeze_to_4D(x_tile) + ttnn.deallocate(old_tile) + return self._last_tile_logits(x_tile) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + x = _slice_last_token_tile(hidden_states, last_token_idx) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + return self._last_tile_logits(x) + + def _last_tile_logits(self, x_tile: ttnn.Tensor) -> ttnn.Tensor: + """Final-norm + all-gather + LM-head on a 32-row tile. ``x_tile`` shape ``[1, 1, 32, dim]``.""" + x = self.norm.prefill_forward(x_tile) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + del last_token_slice, last_token_index + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(self.mesh_device), + ) + x = ttnn.matmul(selector, hidden_states, memory_config=ttnn.DRAM_MEMORY_CONFIG) + ttnn.deallocate(selector) + return self._last_tile_logits(x) + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for layer in self.layers: + x = layer.decode_forward(x, current_pos, rot_mats, page_table=page_table) + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + return self.lm_head.forward(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) + + def prefill_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, start_pos: int = 0, user_id: int = 0) -> ttnn.Tensor: + """Legacy path: embed + RoPE + blocks + final norm (no page table). For tests only.""" + x = self.embed_prefill(token_ids_tt) + seq_len = x.shape[2] + assert seq_len % 128 == 0, "prefill seq_len must be divisible by 128" + rot = self.rope_setup.prefill_forward(start_pos, seq_len) + h = x + for layer in self.layers: + h = layer.prefill_forward(h, rot, user_id=user_id, page_table=None) + h = self.norm.prefill_forward(h) + return _all_gather_rmsnorm_tensor(self.norm, h) + + def decode_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, current_pos: int) -> ttnn.Tensor: + """Legacy path: single-token decode without paged ``page_table``.""" + x = self.embed.forward(token_ids_tt) + x = ttnn.unsqueeze_to_4D(x) + pos = torch.tensor([current_pos], dtype=torch.long) + rot_idxs = prepare_rot_idxs(self.rope_setup.config, pos, on_host=False) + rot = self.rope_setup.decode_forward(rot_idxs) + cur = ttnn.from_torch( + pos, + device=self.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(self.mesh_device), + ) + h = x + for layer in self.layers: + h = layer.decode_forward(h, cur, rot, page_table=None) + h = _all_gather_rmsnorm_tensor(self.norm, h, memory_config=self.norm.config.decode_memory_config) + return self.norm.forward(h, "decode") + + def lm_logits(self, hidden: ttnn.Tensor) -> ttnn.Tensor: + """Project last hidden to logits (vocab-sharded on multi-device). + + Skip the explicit interleaved→shard if the caller already produced a sharded + input (``decode_from_token_ids`` returns the decode-mode RMSNorm's width-sharded + output, which already matches ``LMHead1D.config.input_memcfg``). + """ + x = hidden + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded() and not x.memory_config().is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + return self.lm_head.forward(x) diff --git a/code/models/common/models/qwen25_72b/weight_utils.py b/code/models/common/models/qwen25_72b/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9bcc71a776e3245bb73640351600adad4669d140 --- /dev/null +++ b/code/models/common/models/qwen25_72b/weight_utils.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side HF weight layout helpers for Qwen2.5-72B-Instruct (TTTv2 port). + +Q/K linear weights use HuggingFace head layout; TTNN RoPE expects Meta layout. +This duplicates the small permute helpers from ``load_checkpoints.reverse_permute`` +without importing ``models/tt_transformers``. Logic is shape-agnostic and mirrors +the Qwen2.5-7B / Qwen2.5-Coder-32B ports; the Qwen2 family checkpoint exposes QKV +biases but no ``q_norm`` / ``k_norm`` (those are Qwen3). +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.modules.lazy_weight import LazyWeight + + +def reverse_permute(tensor: torch.Tensor, n_heads: int, dim1: int, dim2: int) -> torch.Tensor: + """HF Q/K weight rows → Meta layout (RoPE-compatible).""" + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def reverse_permute_1d(tensor: torch.Tensor) -> torch.Tensor: + """Undo HF split of last dim into interleaved complex pairs for norm/bias vectors.""" + shape = tensor.shape + dim = shape[-1] + assert dim % 2 == 0, "Last dimension must be even" + reals = tensor[..., : dim // 2] + imags = tensor[..., dim // 2 :] + interleaved = torch.stack((reals, imags), dim=-1).flatten(start_dim=len(shape) - 1) + return interleaved + + +def permute_hf_rope_to_meta_tables(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert HF cos/sin tables to Meta-style tables for ``RotarySetup1D`` (see attention tests).""" + if len(cos.shape) == 3: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + return cos, sin + + +def build_rope_cos_sin_torch( + rotary_emb: Any, table_len: int, head_dim: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Build ``[1, 1, table_len, head_dim]`` cos/sin tensors (Meta layout) from HF rotary module.""" + x = torch.zeros(1, 1, table_len, head_dim, dtype=dtype) + position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0) + with torch.no_grad(): + cos_hf, sin_hf = rotary_emb(x, position_ids) + cos_m, sin_m = permute_hf_rope_to_meta_tables(cos_hf.float(), sin_hf.float()) + return cos_m.to(dtype), sin_m.to(dtype) + + +def attention_wqkv_wo_from_hf_layer( + self_attn: Any, + num_devices: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Mirror ``get_attention_weights_from_ref_model`` in ``test_attention_1d.py`` (HF module API).""" + wq_raw = self_attn.q_proj.weight + wk_raw = self_attn.k_proj.weight + wv_raw = self_attn.v_proj.weight + wo_raw = self_attn.o_proj.weight + + dim = wq_raw.shape[1] + cfg = self_attn.config + n_heads = cfg.num_attention_heads + n_kv_heads = cfg.num_key_value_heads + head_dim = cfg.hidden_size // n_heads + n_heads_times_head_dim = n_heads * head_dim + n_kv_heads_times_head_dim = n_kv_heads * head_dim + + wq_meta = reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim) + wk_meta = reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim) + wv_meta = wv_raw + wo_meta = wo_raw + + wq = wq_meta.T + wk = wk_meta.T + wv = wv_meta.T + wo = wo_meta.T + + qkv_list = [] + for i in range(num_devices): + wq_chunk = torch.chunk(wq, num_devices, dim=1)[i] + wk_chunk = torch.chunk(wk, num_devices, dim=1)[i] + wv_chunk = torch.chunk(wv, num_devices, dim=1)[i] + qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1) + qkv_list.append(qkv) + wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0).clone() + wo = wo.unsqueeze(0).unsqueeze(0).clone() + + q_norm = None + k_norm = None + if getattr(self_attn, "q_norm", None) is not None: + q_norm = reverse_permute_1d(self_attn.q_norm.weight).clone() + if getattr(self_attn, "k_norm", None) is not None: + k_norm = reverse_permute_1d(self_attn.k_norm.weight).clone() + + wqkv_bias = None + if getattr(self_attn.q_proj, "bias", None) is not None: + bq_raw = self_attn.q_proj.bias + bk_raw = self_attn.k_proj.bias + bv_raw = self_attn.v_proj.bias + bq_meta = reverse_permute_1d(bq_raw.view(n_heads, head_dim)).view(-1) + bk_meta = reverse_permute_1d(bk_raw.view(n_kv_heads, head_dim)).view(-1) + bv_meta = bv_raw + qkv_bias_list = [] + for i in range(num_devices): + bq_chunk = torch.chunk(bq_meta, num_devices, dim=0)[i] + bk_chunk = torch.chunk(bk_meta, num_devices, dim=0)[i] + bv_chunk = torch.chunk(bv_meta, num_devices, dim=0)[i] + qkv_bias_list.append(torch.cat([bq_chunk, bk_chunk, bv_chunk], dim=-1)) + wqkv_bias = torch.cat(qkv_bias_list, dim=-1).clone() + + return wqkv, wo, q_norm, k_norm, wqkv_bias + + +def mlp_weights_from_hf_layer(mlp: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (w1, w2, w3) TTNN layouts: gate^T, down^T, up^T for ``MLP1D``.""" + w1 = mlp.gate_proj.weight.T.contiguous().clone() + w3 = mlp.up_proj.weight.T.contiguous().clone() + w2 = mlp.down_proj.weight.T.contiguous().clone() + return w1, w2, w3 + + +def rms_weight_torch(layernorm: Any) -> torch.Tensor: + return layernorm.weight.detach().float().clone() + + +def embed_tokens_torch(embed: Any) -> torch.Tensor: + w = embed.weight.detach().float().clone() + return w.unsqueeze(0).unsqueeze(0) + + +def build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """ + Column-split LM head for ``LMHead1D`` (adapted from ``LMHead1D.from_model_args`` math, no v1 args). + ``lm_head_weight`` shape: ``[vocab_size, dim]`` (HF). + + Returns ``(lazy_weights, split_sizes, weights_memcfgs)`` for ``LMHead1DConfig`` DRAM-sharded matmuls. + + Qwen2.5-72B vocab is 152064 → 152064 / 8 = 19008 per device, and 19008 % 32 == 0, so the + per-device split is tile-aligned without extra padding (no LM-head tile-alignment hazard). + """ + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / 32) * 32 + if vocab_size < padded_vocab_size: + pad = padded_vocab_size - vocab_size + torch_w = torch.cat([torch_w, torch.zeros(torch_w.shape[0], pad, dtype=torch_w.dtype)], dim=-1) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + tile = ttnn.TILE_SIZE + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (tile * dram_size.x)) * (tile * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights: list[LazyWeight] = [] + weights_memcfgs: list[ttnn.MemoryConfig] = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_splits.append(torch_w[:, start:end]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + name = f"lm_head_split_{i}_{combined.shape[-1]}" + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, name) if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs diff --git a/code/models/common/models/qwen25_72b/weights.py b/code/models/common/models/qwen25_72b/weights.py new file mode 100644 index 0000000000000000000000000000000000000000..22fe6666c2794395839ed8b863e4b0dd07f90162 --- /dev/null +++ b/code/models/common/models/qwen25_72b/weights.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side weight helpers for Qwen2.5-72B-Instruct TTTv2. + +Construction lives in ``models.common.models.qwen25_72b.hf_adaptor``; +low-level permutes and HF tensor layout live in ``weight_utils``. +""" + +from models.common.models.qwen25_72b import weight_utils + +__all__ = ["weight_utils"] diff --git a/code/models/common/models/qwen25_7b/README.md b/code/models/common/models/qwen25_7b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ffb43813529c601b60139c52f73e531890bc2958 --- /dev/null +++ b/code/models/common/models/qwen25_7b/README.md @@ -0,0 +1,44 @@ +# Qwen2.5-7B with TTTv2 + +This directory contains the Qwen2.5-7B TTTv2 product path. + +## Construction path + +```text +HF checkpoint + -> hf_adaptor.py: provider configuration, tokenizer, and weights + -> model.py: Qwen2.5 tensor graph composed from TTTv2 modules + -> executor.py: thin typed entry point into qwen2_executor.py + -> generator.py: vLLM boundary, DP composition, and dispatch +``` + +`model.py` owns Qwen2.5 architecture/tuning while composing reusable embedding, +rotary, RMSNorm, attention, MLP, LM-head, optional sampling, and collective +modules. + +## Executor composition + +The model-local file preserves the Qwen2.5-7B executor/config/builder imports. +The implementation is shared with Qwen2 through +`models/common/models/qwen2_executor.py`, which configures the family-neutral +`ModelExecutor`. + +The common owner composes paged KV, output reading, prefill/decode runtimes, +eager/trace execution, warmup, and cleanup. The family layer retains the +7B Q128 top-k warmup order and narrow pre-history request contract. It does not +silently adopt Qwen3 stateful sampling behavior. + +## vLLM, DP, and ownership + +`generator.py` builds one executor per lane, uses `VLLMAdapter` for external +normalization/cache validation, and composes multiple lanes with +`LaneGroupExecutor`. TT resources remain lane-owned. + +## Tests + +- `models/common/tests/models/qwen25_7b/test_hf_adaptor.py` +- `models/common/tests/models/qwen25_7b/test_demo_contract.py` +- `models/common/tests/demos/qwen25_7b/demo.py` +- `models/common/tests/models/test_qwen2_executor_family.py` +- `models/common/tests/llm_runtime/test_executor_integration.py` +- `models/common/tests/llm_runtime/test_model_executor.py` diff --git a/code/models/common/models/qwen25_7b/__init__.py b/code/models/common/models/qwen25_7b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cd861b72109c2c1fb9ee60d593bc251e49e50425 --- /dev/null +++ b/code/models/common/models/qwen25_7b/__init__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +from models.common.models.qwen25_7b.executor import Qwen25Executor, Qwen25ExecutorConfig +from models.common.models.qwen25_7b.generator import Qwen25Generator, Qwen25GeneratorConfig +from models.common.models.qwen25_7b.hf_adaptor import Qwen25ForCausalLM, Qwen25RuntimeConfig +from models.common.models.qwen25_7b.model import ( + QWEN25_7B_ACCURACY, + QWEN25_7B_PERFORMANCE, + Qwen25PagedAttentionConfig, + Qwen25_7B, + Qwen25_7BPrecisionConfig, + Qwen25_7BTransformerConfig, +) + +__all__ = [ + "QWEN25_7B_ACCURACY", + "QWEN25_7B_PERFORMANCE", + "Qwen25_7B", + "Qwen25_7BPrecisionConfig", + "Qwen25_7BTransformerConfig", + "Qwen25Executor", + "Qwen25ExecutorConfig", + "Qwen25ForCausalLM", + "Qwen25PagedAttentionConfig", + "Qwen25Generator", + "Qwen25GeneratorConfig", + "Qwen25RuntimeConfig", +] diff --git a/code/models/common/models/qwen25_7b/executor.py b/code/models/common/models/qwen25_7b/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..d408feb37c317eff2b2f1d776b663a67112e2f7e --- /dev/null +++ b/code/models/common/models/qwen25_7b/executor.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen2.5-7B construction entry point over the shared Qwen2 executor.""" + +from models.common.models.qwen2_executor import Qwen25Executor, Qwen25ExecutorConfig +from models.common.models.qwen2_executor import build_qwen25_7b_executor as _build_qwen25_7b_executor +from models.common.models.qwen25_7b.hf_adaptor import Qwen25ForCausalLM + + +def build_qwen25_7b_executor(llm: Qwen25ForCausalLM, config: Qwen25ExecutorConfig) -> Qwen25Executor: + """Build one executor around an already-loaded Qwen2.5 model adapter.""" + + return _build_qwen25_7b_executor(llm, config) diff --git a/code/models/common/models/qwen25_7b/generator.py b/code/models/common/models/qwen25_7b/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..ace298acdb874d4afa87791f8ef97c3ea0cbbd8e --- /dev/null +++ b/code/models/common/models/qwen25_7b/generator.py @@ -0,0 +1,481 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Qwen2.5-7B.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.qwen25_7b.executor import Qwen25ExecutorConfig, build_qwen25_7b_executor +from models.common.models.qwen25_7b.hf_adaptor import from_pretrained +from models.common.models.qwen25_7b.model import Qwen25PagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Qwen25GeneratorConfig: + """Validated construction inputs for one vLLM-facing Qwen generator.""" + + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + hf_revision: str | None = None + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Qwen25Generator: + """Adapt vLLM's model interface to the model-owned execution target. + + vLLM constructs this facade with `initialize_vllm_model`, resolves + KV capacity through `allocate_kv_cache`, warms the configured + programs, and then calls `prefill_forward` and + `decode_forward`. Each forward call is normalized by + `VLLMAdapter`, dispatched to the eager or traced executor, and + delegated to ``Qwen25Executor`` or ``LaneGroupExecutor``. + + This class owns dispatch policy but no TT resources. `cleanup` + delegates to the target that owns those resources. + """ + + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + + # Public vLLM API + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + """Return the unpadded per-submesh KV token budget for vLLM sizing.""" + + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + hf_revision: str | None = None, + device_sampling_enabled: bool = True, + ): + """Build the configured single-lane or data-parallel Qwen target.""" + + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_qwen25_7b_generator( + Qwen25GeneratorConfig( + hf_model=str(hf_model), + hf_revision=hf_revision, + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def compile_prefill( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + ) -> None: + """Normalize a vLLM prefill call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.compile_prefill(execution=execution, **normalized) + + def compile_decode( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + ) -> None: + """Normalize a vLLM decode call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.compile_decode(execution=execution, **normalized) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM prefill call.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM decode call.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward( + execution=execution, + read_from_device=read_from_device, + **normalized, + ) + + def read_decode_output( + self, + tt_out: Any, + *, + async_read: bool = False, + ) -> Any: + """Delegate vLLM's raw decode-output read.""" + + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host( + self, + tt_out: Any, + *, + is_tokens: bool = False, + ) -> tuple[Any, Any]: + """Delegate vLLM's asynchronous host-output completion.""" + + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self): + """Release every resource owned by the concrete target.""" + + return self.target.cleanup() + + # Private implementation + + def _select_prefill_execution( + self, + normalized: NormalizedPrefillKwargs, + trace_requested: bool, + ): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_qwen25_7b_generator(config: Qwen25GeneratorConfig) -> Qwen25Generator: + """Construct lane-local models/executors and compose their shared target surface.""" + + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Qwen25PagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + hf_revision=config.hf_revision, + instruct="Instruct" in config.hf_model, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Qwen25ExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_qwen25_7b_executor(llm, executor_config)) + + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Qwen25Generator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Qwen2.5-7B model must contain at least one attention layer") + + attention_configs = tuple(layer.attention.config for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Qwen layer must expose the same KV head shape") + + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/qwen25_7b/hf_adaptor.py b/code/models/common/models/qwen25_7b/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..f67ceaf876a104ec6beb7d507d1f421df1930979 --- /dev/null +++ b/code/models/common/models/qwen25_7b/hf_adaptor.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face product adaptor for the TTTv2 Qwen2.5-7B path.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.models.qwen25_7b import weight_utils +from models.common.models.qwen25_7b.model import ( + QWEN25_7B_ACCURACY, + QWEN25_7B_PERFORMANCE, + Qwen25_7B, + Qwen25_7BLayerWeights, + Qwen25_7BModelParameters, + Qwen25_7BPrecisionConfig, + Qwen25_7BWeights, + Qwen25PagedAttentionConfig, + build_qwen25_7b_transformer_config, +) + +DEFAULT_HF_MODEL = "Qwen/Qwen2.5-7B-Instruct" + + +@dataclass(frozen=True) +class Qwen25GenerationConfig: + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Qwen25RuntimeConfig: + model_name: str + model_cache_path: Path | None + max_prefill_chunk_size: int + max_context_len: int + max_seq_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 8 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + return ( + prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + chat = [] + if isinstance(prompt_text, str): + if system_prompt_text: + chat.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + chat.append({"role": "user", "content": prompt_text}) + else: + chat = prompt_text + try: + return _chat_template_ids(tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True)) + except ValueError: + pass + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Qwen25ForCausalLM: + model: Qwen25_7B + tokenizer: Any + runtime_config: Qwen25RuntimeConfig + instruct: bool = True + generation_config: Qwen25GenerationConfig = field(default_factory=Qwen25GenerationConfig) + + def __post_init__(self): + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + stops = tuple(getattr(self.tokenizer, "stop_tokens", []) or []) + self.generation_config = Qwen25GenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=stops, + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + return encode_prompt( + self.tokenizer, + prompt_text, + system_prompt_text, + instruct=self.instruct if instruct is None else instruct, + ) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_tokenizer(hf_model: str, hf_revision: str | None = None): + tokenizer = AutoTokenizer.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=os.getenv("CI") == "true", + ) + eos = getattr(tokenizer, "eos_token_id", None) + stops = [] if eos is None else ([eos] if isinstance(eos, int) else list(eos)) + im_start = tokenizer.convert_tokens_to_ids("<|im_start|>") + if isinstance(im_start, int) and im_start >= 0 and im_start not in stops: + stops.append(im_start) + tokenizer.stop_tokens = stops + return tokenizer + + +def _trace_seq_lens(num_devices: int, max_prefill_chunk_size: int, max_seq_len: int) -> tuple[int, ...]: + if num_devices != 2: + raise ValueError(f"Qwen2.5-7B supports logical TP2 lanes only, got {num_devices} devices") + allowed = (128, 1024) + return tuple(length for length in allowed if length <= min(max_prefill_chunk_size, max_seq_len)) + + +def _cache_path(hf_model: str, mesh_device, cache_dir: Path | str | None) -> Path: + if cache_dir is not None: + path = Path(cache_dir) + elif os.getenv("TT_CACHE_PATH"): + path = Path(os.environ["TT_CACHE_PATH"]) + else: + path = Path("model_cache") / hf_model / "N300" + path.mkdir(parents=True, exist_ok=True) + return path + + +def convert_hf_model_weights( + hf, + hf_config, + *, + n_layers: int, + num_devices: int, + rope_table_len: int, + head_dim: int, +) -> Qwen25_7BWeights: + """Extract and convert all Hugging Face tensors consumed by the TT builder.""" + base = hf.model + rope_cos, rope_sin = weight_utils.build_rope_cos_sin_torch( + base.rotary_emb, rope_table_len, head_dim, torch.bfloat16 + ) + layers = [] + for layer in base.layers[:n_layers]: + wqkv, wo, q_norm, k_norm, wqkv_bias = weight_utils.attention_wqkv_wo_from_hf_layer(layer.self_attn, num_devices) + if q_norm is not None or k_norm is not None: + raise ValueError("Qwen2.5-7B does not support QK norm weights") + if wqkv_bias is None: + raise ValueError("Qwen2.5-7B requires QKV projection bias") + w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(layer.mlp) + layers.append( + Qwen25_7BLayerWeights( + wqkv=wqkv, + wo=wo, + wqkv_bias=wqkv_bias, + w1=w1, + w2=w2, + w3=w3, + attention_norm=weight_utils.rms_weight_torch(layer.input_layernorm).to(torch.bfloat16), + ff_norm=weight_utils.rms_weight_torch(layer.post_attention_layernorm).to(torch.bfloat16), + ) + ) + + if hf_config.tie_word_embeddings: + raise ValueError("Qwen2.5-7B requires an untied LM head") + return Qwen25_7BWeights( + embedding=weight_utils.embed_tokens_torch(base.embed_tokens), + rope_cos=rope_cos, + rope_sin=rope_sin, + layers=tuple(layers), + final_norm=weight_utils.rms_weight_torch(base.norm).to(torch.bfloat16), + lm_head=hf.lm_head.weight.detach().to(torch.bfloat16).clone(), + ) + + +def from_pretrained( + mesh_device, + *, + hf_model: str = DEFAULT_HF_MODEL, + hf_revision: str | None = None, + instruct: bool = True, + max_batch_size: int = 32, + max_seq_len: int = 4096, + optimizations: str | Qwen25_7BPrecisionConfig = "accuracy", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config: Qwen25PagedAttentionConfig | None = None, + cache_dir: Path | str | None = None, +) -> Qwen25ForCausalLM: + del dtype + if mesh_device.get_num_devices() != 2: + raise ValueError(f"Qwen2.5-7B supports logical TP2 lanes only, got {mesh_device.get_num_devices()} devices") + ttnn.SetDefaultDevice(mesh_device) + hf_config = AutoConfig.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=os.getenv("CI") == "true", + ) + hf = AutoModelForCausalLM.from_pretrained( + hf_model, + revision=hf_revision, + torch_dtype=torch.bfloat16, + local_files_only=os.getenv("CI") == "true", + ) + hf.eval() + resolved_layers = hf_config.num_hidden_layers if n_layers is None else n_layers + precision = ( + optimizations + if isinstance(optimizations, Qwen25_7BPrecisionConfig) + else (QWEN25_7B_PERFORMANCE if optimizations == "performance" else QWEN25_7B_ACCURACY) + ) + if not isinstance(precision, Qwen25_7BPrecisionConfig): + raise TypeError("optimizations must be 'accuracy', 'performance', or Qwen25_7BPrecisionConfig") + cache_path = _cache_path(hf_model, mesh_device, cache_dir) + if paged_attention_config is None: + block_size = 32 + paged_attention_config = Qwen25PagedAttentionConfig( + block_size=block_size, + max_num_blocks=((max_seq_len + block_size - 1) // block_size) * max_batch_size, + ) + head_dim = hf_config.hidden_size // hf_config.num_attention_heads + params = Qwen25_7BModelParameters( + dim=hf_config.hidden_size, + n_heads=hf_config.num_attention_heads, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + hidden_dim=hf_config.intermediate_size, + vocab_size=hf_config.vocab_size, + rms_norm_eps=hf_config.rms_norm_eps, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + ) + rope_table_len = ((max(max_seq_len * 2, 8192) + 127) // 128) * 128 + weights = convert_hf_model_weights( + hf, + hf_config, + n_layers=resolved_layers, + num_devices=mesh_device.get_num_devices(), + rope_table_len=rope_table_len, + head_dim=head_dim, + ) + model_config = build_qwen25_7b_transformer_config( + mesh_device=mesh_device, + params=params, + weights=weights, + n_layers=resolved_layers, + precision=precision, + cache_path=cache_path, + paged_attention_config=paged_attention_config, + model_name=Path(hf_model).name, + ) + tokenizer = load_tokenizer(hf_model, hf_revision) + model = Qwen25_7B(model_config) + max_prefill_chunk_size = 2048 + runtime_config = Qwen25RuntimeConfig( + model_name=Path(hf_model).name, + model_cache_path=cache_path, + max_prefill_chunk_size=max_prefill_chunk_size, + max_context_len=int(hf_config.max_position_embeddings), + max_seq_len=max_seq_len, + trace_prefill_supported_seq_lens=_trace_seq_lens( + mesh_device.get_num_devices(), max_prefill_chunk_size, max_seq_len + ), + disable_batched_prefill=bool(os.getenv("DISABLE_BATCHED_PREFILL")), + ) + del hf + return Qwen25ForCausalLM( + model=model, + tokenizer=tokenizer, + runtime_config=runtime_config, + instruct=instruct, + ) diff --git a/code/models/common/models/qwen25_7b/model.py b/code/models/common/models/qwen25_7b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..69784b1b708de38da2d57d489d625470faeb2dda --- /dev/null +++ b/code/models/common/models/qwen25_7b/model.py @@ -0,0 +1,1155 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Qwen2.5-7B-Instruct — native TTTv2 stack (``Embedding1D``, ``RMSNorm1D``, +``Attention1D``, ``MLP1D``, ``RotarySetup1D``, ``LMHead1D``). + +Tensor layout contracts: + - **Prefill** hidden states: ``[1, 1, S, dim]`` TILE, ``S % 128 == 0``. + - **Decode** hidden states: ``[1, 1, B, dim]`` TILE (``B`` padded to tile in modules). + +The tensor model is provider-neutral. Hugging Face loading and conversion live +in :mod:`models.common.models.qwen25_7b.hf_adaptor`. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D, prepare_rot_idxs +from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + """Minimal LazyWeight; ``Attention1D`` / ``MLP1D`` / ``Embedding1D`` resolvers fill mesh + memory.""" + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +@dataclass +class Qwen25PagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (``block_size`` / ``max_num_blocks`` only).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Qwen25_7BTransformerConfig: + """Complete provider-neutral tensor-graph configuration.""" + + n_layers: int + vocab_size: int + max_batch_size: int + max_seq_len: int + dim: int + num_devices: int + mesh_device: ttnn.MeshDevice + embedding_config: Embedding1DConfig + rope_config: Rope1DConfig + block_configs: list[Qwen25_7BDecoderLayerConfig] + norm_config: RMSNorm1DConfig + lm_head_config: LMHead1DConfig + sampling_config: Sampling1DConfig | None = None + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtypes: list[ttnn.DataType | None] = field(default_factory=list) + tt_ccl: Any = None + cache_path: str | None = None + + +@dataclass(frozen=True) +class Qwen25_7BModelParameters: + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + max_batch_size: int + max_seq_len: int + + +@dataclass(frozen=True) +class Qwen25_7BLayerWeights: + wqkv: torch.Tensor + wo: torch.Tensor + wqkv_bias: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Qwen25_7BWeights: + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Qwen25_7BLayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +@dataclass(frozen=True) +class Qwen25_7BPrecisionConfig: + """Per-layer precision + math-fidelity recipe for Qwen2.5-7B-Instruct. + + Mirrors the fields TTTv1's ``DecodersPrecision`` actually distinguishes for + Qwen2.5-7B (Llama-family group in ``model_config.py`` for both ``accuracy()`` + and ``performance()``). Two module-level recipes are exposed: + :data:`QWEN25_7B_ACCURACY` and :data:`QWEN25_7B_PERFORMANCE`. Pass one to + the product adaptor via ``optimizations=``; use + ``dataclasses.replace(QWEN25_7B_PERFORMANCE, perf_decode_tuning=False)`` to + customize a single field (e.g. disable aggressive decode math during + teacher-forcing parity runs). + + Attention compute-kernel configs (HIFI4 + fp32 dest acc for LI_QKV / LI_O / + SDPA prefill, plus the decode variants and the LoFi perf-decode kernel) are + resolved inside ``_resolve_qwen_wh_tuning`` from + :attr:`perf_decode_tuning`; they are not exposed here because TTTv1 ships + them as a coupled bundle for this model. + """ + + # Qwen2.5-7B keeps BF16 WQKV/WO in BOTH TTTv1 recipes — TTTv2's + # ``Attention1D`` default of BFP8 silently downgrades QKV/WO precision and + # broadens per-layer divergence on this model (see + # debugging/numerical_divergence_vs_hf_2026-05-14.md). + wqkv_dtype: ttnn.DataType = ttnn.bfloat16 + wo_dtype: ttnn.DataType = ttnn.bfloat16 + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + # MLP weights: BFP8 in both modes for Qwen2.5-7B (TTTv1 does not drop these + # to BFP4 for this checkpoint). + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat8_b + mlp_w2_dtype: ttnn.DataType = ttnn.bfloat8_b + + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + + # Decode-side throughput tuning: LoFi attention decode (LI_QKV / SDPA / + # LI_O), HiFi2 MLP decode FF, SDPA exp-approx, no W1→DRAM spill for small + # batch. Performance recipe only — keep ``False`` in accuracy mode and in + # any teacher-forcing path (aggressive decode math drops top-1). + perf_decode_tuning: bool = False + + +# TTTv1 ``DecodersPrecision.accuracy("Qwen2.5-7B-Instruct")`` (Llama-family +# group in ``model_config.py``): BF16 attention, BFP8 MLP, BF16 LM head, BF16 +# KV cache, HiFi4 throughout. ``perf_decode_tuning`` stays off — decode keeps +# HiFi4 + fp32 dest acc. +QWEN25_7B_ACCURACY = Qwen25_7BPrecisionConfig( + kv_cache_dtype=ttnn.bfloat16, + lm_head_dtype=ttnn.bfloat16, +) + +# TTTv1 ``DecodersPrecision.performance("Qwen2.5-7B-Instruct")``: same dtypes +# as accuracy except ``kv_cache_dtype`` / ``lm_head_dtype`` drop to BFP8, plus +# the decode-side throughput bundle (LoFi attn decode + SDPA exp-approx + HiFi2 +# MLP decode FF) is engaged via ``perf_decode_tuning=True``. +QWEN25_7B_PERFORMANCE = Qwen25_7BPrecisionConfig( + perf_decode_tuning=True, +) + + +def _qwen_wh_mlp_matmul_compute_kernel() -> ttnn.WormholeComputeKernelConfig: + """HiFi4 lowers L1 circular-buffer footprint vs HiFi2 for wide FF matmuls on Wormhole.""" + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + + +def _qwen_wh_mlp_decode_matmul_compute_kernel() -> ttnn.WormholeComputeKernelConfig: + """HiFi2 decode FF matmuls: faster steady-state decode while prefill keeps HiFi4 (``_qwen_wh_mlp_matmul_compute_kernel``).""" + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + + +def _qwen_wh_decode_attn_lofi_kernel() -> ttnn.WormholeComputeKernelConfig: + """LoFi decode attention matmuls + SDPA when ``perf_decode_tuning`` is enabled (performance demo only).""" + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=False, + ) + + +def _qwen_wh_attn_hifi4_kernel() -> ttnn.WormholeComputeKernelConfig: + """HiFi4 + fp32 dest acc for Qwen2.5-7B attention matmuls (LI_QKV, LI_O, SDPA). + + Matches the TTTv1 ``ModelArgs.compute_kernel_config_hifi4`` used by both + ``ModelOptimizations.performance`` and ``ModelOptimizations.accuracy`` for + Qwen2.5-7B (see ``models/tt_transformers/tt/model_config.py``). + The TTTv2 ``Attention1D`` defaults are HiFi2 with fp16 accumulation; that + silently downgrades attention prefill QKV/WO and decode QKV/SDPA/WO matmul + precision for this model, producing a broad per-layer divergence vs HF. + """ + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=True, + ) + + +def _slice_last_token_tile(x: ttnn.Tensor, last_token_idx: int) -> ttnn.Tensor: + """Slice the 32-row tile containing ``last_token_idx`` from ``[1, 1, S, W]``. + + Width-sharded LM matmul M tile rows must equal ``LMHead1D`` program-config tile rows. + """ + floor = (last_token_idx // 32) * 32 + return ttnn.slice(x, (0, 0, floor, 0), (1, 1, floor + 32, x.shape[-1])) + + +def _post_attn_norm_decode_configs( + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding so its output matches MLP1D's W1/W3 input. + + MLP1D decode uses ``_dram_shard_core_grid_k_n(dim, padded_hidden / num_devices)`` for W1/W3 + inputs, but the default RMSNorm program config is derived from ``_compute_norm_core_grid(dim)`` + alone. On N300 Qwen2.5-7B that DRAM-width-shard mismatch silently corrupts decode activations. + """ + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + memory_config = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return program_config, memory_config + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, x: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + return x + + if memory_config is None: + memory_config = x.memory_config() + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + + +@dataclass +class _Qwen25WHTuning: + """Wormhole-specific MLP / LM-head / attention tuning resolved at build time. + + Populated for Qwen2.5-7B on N150 / N300 only; otherwise all fields stay ``None`` + (modules fall back to library defaults). + """ + + mlp_prefill_len_cutoff: int | None = None + mlp_ff_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + mlp_decode_spill_w1_to_dram: bool = False + mlp_decode_ff_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + lm_head_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + perf_decode_sdpa_cfg: ttnn.SDPAProgramConfig | None = None + # Attention prefill compute kernels: always HiFi4 + fp32 dest acc on Qwen + # (Attention1D defaults are HiFi2/fp16 acc — too lossy for this model). + attn_li_qkv_prefill_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_li_o_prefill_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + # Attention decode compute kernels: HiFi4 + fp32 dest acc by default; LoFi + # under ``perf_decode_tuning`` (performance demo only). + attn_li_qkv_decode_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_sdpa_decode_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_li_o_decode_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + +def _resolve_qwen_wh_tuning( + *, hf_model_id: str, num_dev: int, max_batch_size: int, perf_decode_tuning: bool +) -> _Qwen25WHTuning: + """Pick WH tuning knobs for Qwen2.5-7B-Instruct on N150 / N300. + + ``get_padded_prefill_len`` maps 129..1024 tokens to a 1024-wide tile. ``MLP1D`` then + reshapes/chunks using ``prefill_len_cutoff``. + + Qwen2.5-7B preserves its validated 256-token MLP prefill cutoff. This halves the + per-kernel M tile for the model's wide per-device FF shard and keeps its established + circular-buffer footprint. + + Output is unchanged by the cutoff: ``in0_block_w`` and the K-contraction order are independent of the + M-tiling, so only ``per_core_M`` changes. Decode never reads ``prefill_len_cutoff``. + + Decode W1→DRAM before W3 avoids L1 overlap between W1 activations and W3 matmul CBs + (N300 batch 32). + """ + t = _Qwen25WHTuning() + model_slug = hf_model_id.split("/")[-1] + if not (model_slug.startswith("Qwen2.5-7B") and num_dev in (1, 2)): + return t + + t.mlp_prefill_len_cutoff = 256 + t.mlp_ff_compute_kernel_cfg = _qwen_wh_mlp_matmul_compute_kernel() + t.mlp_decode_spill_w1_to_dram = max_batch_size >= 16 + t.lm_head_compute_kernel_cfg = _qwen_wh_mlp_matmul_compute_kernel() + t.mlp_decode_ff_compute_kernel_cfg = ( + _qwen_wh_mlp_decode_matmul_compute_kernel() if perf_decode_tuning else t.mlp_ff_compute_kernel_cfg + ) + + # Attention prefill always uses HiFi4 + fp32 dest acc on Qwen (matches TTTv1). + attn_hifi4 = _qwen_wh_attn_hifi4_kernel() + t.attn_li_qkv_prefill_kernel_cfg = attn_hifi4 + t.attn_li_o_prefill_kernel_cfg = attn_hifi4 + # Attention decode defaults to HiFi4 + fp32 dest acc; LoFi only under perf_decode_tuning. + t.attn_li_qkv_decode_kernel_cfg = attn_hifi4 + t.attn_sdpa_decode_kernel_cfg = attn_hifi4 + t.attn_li_o_decode_kernel_cfg = attn_hifi4 + + if perf_decode_tuning: + lo = _qwen_wh_decode_attn_lofi_kernel() + t.perf_decode_sdpa_cfg = ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=(8, 8), + exp_approx_mode=True, + q_chunk_size=0, + k_chunk_size=0, + ) + t.attn_li_qkv_decode_kernel_cfg = lo + t.attn_sdpa_decode_kernel_cfg = lo + t.attn_li_o_decode_kernel_cfg = lo + logger.info( + f"MLP/LM/attention tuning for {hf_model_id} on {num_dev} device(s): " + f"prefill_len_cutoff={t.mlp_prefill_len_cutoff}, FF prefill HiFi4, attn prefill+decode HiFi4+fp32, " + f"decode spill W1→DRAM={t.mlp_decode_spill_w1_to_dram}, " + f"perf_decode_tuning={perf_decode_tuning}" + ) + return t + + +@dataclass +class Qwen25_7BDecoderLayerConfig: + attention_norm_config: RMSNorm1DConfig + attention_config: Attention1DConfig + ff_norm_config: RMSNorm1DConfig + mlp_config: MLP1DConfig + + +def _build_decoder_layer( + *, + idx: int, + weights: Qwen25_7BLayerWeights, + qcfg: Qwen25_7BModelParameters, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Qwen25_7BPrecisionConfig, + paged_cfg: Qwen25PagedAttentionConfig, + cache_path: Path | None, + wh: _Qwen25WHTuning, +) -> Qwen25_7BDecoderLayerConfig: + """Build one decoder-layer config from provider-converted tensors.""" + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, + dtype=precision.wqkv_dtype, + cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None, + ) + lazy_wo = _lazy( + weights.wo, + dtype=precision.wo_dtype, + cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None, + ) + + bias_lw = LazyWeight( + source=weights.wqkv_bias.to(torch.bfloat16), + dtype=ttnn.bfloat16, + cache_dir_weight_name=(cache_path / "attn", f"{prefix}_bias") if cache_path else None, + ) + + attention_config = Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=qcfg.n_heads, + n_kv_heads=qcfg.n_kv_heads, + head_dim=qcfg.head_dim, + max_batch_size=qcfg.max_batch_size, + max_seq_len=qcfg.max_seq_len, + q_norm_config=None, + k_norm_config=None, + wqkv_bias=bias_lw, + use_vllm_paged_kv_cache=True, + paged_attention_config=paged_cfg, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + decode_sdpa_prg_config=wh.perf_decode_sdpa_cfg, + li_qkv_decode_compute_kernel_cfg=wh.attn_li_qkv_decode_kernel_cfg, + sdpa_decode_compute_kernel_cfg=wh.attn_sdpa_decode_kernel_cfg, + li_o_decode_compute_kernel_cfg=wh.attn_li_o_decode_kernel_cfg, + li_qkv_prefill_compute_kernel_cfg=wh.attn_li_qkv_prefill_kernel_cfg, + li_o_prefill_compute_kernel_cfg=wh.attn_li_o_prefill_kernel_cfg, + ) + + mlp_config = MLP1DConfig( + w1=_lazy( + weights.w1, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None, + ), + w2=_lazy( + weights.w2, + dtype=precision.mlp_w2_dtype, + cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None, + ), + w3=_lazy( + weights.w3, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None, + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=qcfg.max_batch_size, + prefill_len_cutoff=wh.mlp_prefill_len_cutoff, + ff1_3_compute_kernel_cfg=wh.mlp_ff_compute_kernel_cfg, + ff2_compute_kernel_cfg=wh.mlp_ff_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=wh.mlp_decode_ff_compute_kernel_cfg, + decode_ff2_compute_kernel_cfg=wh.mlp_decode_ff_compute_kernel_cfg, + decode_spill_w1_to_dram_before_w3=wh.mlp_decode_spill_w1_to_dram, + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + dim=qcfg.dim, + hidden_dim=qcfg.hidden_dim, + num_devices=num_dev, + max_batch_size=qcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1DConfig: + lw = _lazy( + weight, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=qcfg.rms_norm_eps, + max_batch_size=qcfg.max_batch_size, + tt_ccl=tt_ccl, + **extra, + ) + + return Qwen25_7BDecoderLayerConfig( + attention_norm_config=_build_norm(weights.attention_norm, "pre_attn"), + attention_config=attention_config, + ff_norm_config=_build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ), + mlp_config=mlp_config, + ) + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + qcfg: Qwen25_7BModelParameters, + lm_head_dtype: ttnn.DataType, + lm_head_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None, + cache_path: Path | None, +) -> LMHead1DConfig: + """Build the vocab-sharded LM head with DRAM-matmul program configs. + + LM head DRAM matmul is sized for decode batch tiles (``max_batch_size``). Prefill logits + use a single 32-row tile via ``post_process_prefill_output`` / :func:`_slice_last_token_tile`. + """ + lm_splits, lm_split_sizes, lm_weights_memcfgs = _build_lm_head_lazy_weights( + mesh_device, + lm_head_weight, + dim=qcfg.dim, + vocab_size=qcfg.vocab_size, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(qcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(qcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, qcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(qcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=qcfg.dim, + max_batch_size=qcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + compute_kernel_config=lm_head_compute_kernel_cfg, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + ) + + +def _build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """Column-split a provider-neutral LM-head tensor for ``LMHead1D``.""" + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / TILE_SIZE) * TILE_SIZE + if vocab_size < padded_vocab_size: + torch_w = torch.cat( + [torch_w, torch.zeros(torch_w.shape[0], padded_vocab_size - vocab_size, dtype=torch_w.dtype)], dim=-1 + ) + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (TILE_SIZE * dram_size.x)) * (TILE_SIZE * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights = [] + weights_memcfgs = [] + for index, split_size in enumerate(split_sizes): + device_splits = [] + for device_index in range(num_devices): + start = device_index * size_per_device + sum(split_sizes[:index]) + device_splits.append(torch_w[:, start : start + split_size]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, f"lm_head_split_{index}_{combined.shape[-1]}") if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs + + +def build_qwen25_7b_transformer_config( + *, + mesh_device: ttnn.MeshDevice, + params: Qwen25_7BModelParameters, + weights: Qwen25_7BWeights, + n_layers: int, + precision: Qwen25_7BPrecisionConfig, + cache_path: Path, + paged_attention_config: Qwen25PagedAttentionConfig, + model_name: str = "Qwen2.5-7B-Instruct", +) -> Qwen25_7BTransformerConfig: + """Build the TT graph from provider-neutral dimensions and converted tensors.""" + num_devices = mesh_device.get_num_devices() + if num_devices != 2: + raise ValueError(f"Qwen2.5-7B supports logical TP2 lanes only, got {num_devices} devices") + if params.n_heads % num_devices or params.n_kv_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({params.n_heads}/{params.n_kv_heads}) " + f"must be divisible by device count ({num_devices})" + ) + if len(weights.layers) != n_layers: + raise ValueError(f"Expected {n_layers} decoder layer weight sets, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) + topology = default_topology(mesh_device) + wh = _resolve_qwen_wh_tuning( + hf_model_id=model_name, + num_dev=num_devices, + max_batch_size=params.max_batch_size, + perf_decode_tuning=precision.perf_decode_tuning, + ) + embedding_config = Embedding1DConfig( + weights=_lazy(weights.embedding, dtype=ttnn.bfloat16, cache=(cache_path / "embedding", "tok_embeddings")), + mesh_device=mesh_device, + embed_scale=1.0, + ) + rope_config = Rope1DConfig( + cos_matrix=_lazy(weights.rope_cos, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "cos")), + sin_matrix=_lazy(weights.rope_sin, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "sin")), + max_batch_size=params.max_batch_size, + head_dim=params.head_dim, + device=mesh_device, + use_qk_fused=False, + ) + block_configs = [ + _build_decoder_layer( + idx=index, + weights=weights.layers[index], + qcfg=params, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + paged_cfg=paged_attention_config, + cache_path=cache_path, + wh=wh, + ) + for index in range(n_layers) + ] + norm_config = RMSNorm1DConfig( + weight=_lazy(weights.final_norm, dtype=ttnn.bfloat16, cache=(cache_path / "norm", "final")), + mesh_device=mesh_device, + eps=params.rms_norm_eps, + max_batch_size=params.max_batch_size, + tt_ccl=tt_ccl, + ) + lm_head_config = _build_lm_head( + mesh_device=mesh_device, + lm_head_weight=weights.lm_head, + qcfg=params, + lm_head_dtype=precision.lm_head_dtype, + lm_head_compute_kernel_cfg=wh.lm_head_compute_kernel_cfg, + cache_path=cache_path, + ) + sampling_config = Sampling1DConfig( + vocab_size=params.vocab_size, + valid_vocab_size=params.vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=_nearest_32(params.max_batch_size), + allow_force_argmax=False, + pad_to_power_of_2=True, + ) + return Qwen25_7BTransformerConfig( + n_layers=n_layers, + vocab_size=params.vocab_size, + max_batch_size=params.max_batch_size, + max_seq_len=params.max_seq_len, + dim=params.dim, + num_devices=num_devices, + mesh_device=mesh_device, + embedding_config=embedding_config, + rope_config=rope_config, + block_configs=block_configs, + norm_config=norm_config, + lm_head_config=lm_head_config, + sampling_config=sampling_config, + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtypes=[None] * n_layers, + tt_ccl=tt_ccl, + cache_path=str(cache_path), + ) + + +class Qwen25_7BDecoderLayer(LightweightModule): + def __init__( + self, + *, + input_layernorm: RMSNorm1D, + self_attn: Attention1D, + post_attention_layernorm: RMSNorm1D, + mlp: MLP1D, + ): + super().__init__() + self.input_layernorm = input_layernorm + self.self_attn = self_attn + self.post_attention_layernorm = post_attention_layernorm + self.mlp = mlp + self.attention_norm = input_layernorm + self.attention = self_attn + self.ff_norm = post_attention_layernorm + self.feed_forward = mlp + + @classmethod + def from_config(cls, config: Qwen25_7BDecoderLayerConfig) -> Qwen25_7BDecoderLayer: + return cls( + input_layernorm=RMSNorm1D.from_config(config.attention_norm_config), + self_attn=Attention1D.from_config(config.attention_config), + post_attention_layernorm=RMSNorm1D.from_config(config.ff_norm_config), + mlp=MLP1D.from_config(config.mlp_config), + ) + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + *, + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + # Match Llama ``TransformerBlock1D``: fractured embed / norm activations must be + # all-gathered to full ``dim`` before Attention1D / MLP1D (QKV matmul expects width ``dim``). + r = self.input_layernorm.prefill_forward(x) + r = _all_gather_rmsnorm_tensor(self.input_layernorm, r) + r = self.self_attn.prefill_forward( + r, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + r2 = self.post_attention_layernorm.prefill_forward(h) + r2 = _all_gather_rmsnorm_tensor(self.post_attention_layernorm, r2) + r2 = self.mlp.prefill_forward(r2) + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def decode_forward( + self, + x: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + xa = _all_gather_rmsnorm_tensor( + self.input_layernorm, x, memory_config=self.input_layernorm.config.decode_memory_config + ) + r = self.input_layernorm.forward(xa, "decode") + r = self.self_attn.forward(r, current_pos, rot_mats, mode="decode", page_table=page_table) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + hf = _all_gather_rmsnorm_tensor( + self.post_attention_layernorm, h, memory_config=self.post_attention_layernorm.config.decode_memory_config + ) + r2 = self.post_attention_layernorm.forward(hf, "decode") + r2 = self.mlp.forward(r2, "decode") + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + +class Qwen25_7B(LightweightModule): + """ + Full decoder for Qwen2.5-7B-Instruct (TTTv2 modules only). + + Prefill/decode on **embedded** activations match ``EagerLLMExecutor``. Token embedding + is ``embed_prefill`` / ``embed_decode``. Bind KV with ``set_kv_cache`` before first forward. + """ + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__(self, config: Qwen25_7BTransformerConfig): + from tqdm import tqdm + + super().__init__() + self.config = config + self.embed = Embedding1D.from_config(config.embedding_config) + self.embedding = self.embed + self.rope_setup = RotarySetup1D.from_config(config.rope_config) + self.layers = [ + Qwen25_7BDecoderLayer.from_config(config.block_configs[index]) + for index in tqdm(range(config.n_layers), desc="Building layers") + ] + self.norm = RMSNorm1D.from_config(config.norm_config) + self.lm_head = LMHead1D.from_config(config.lm_head_config) + self.sampling = Sampling1D.from_config(config.sampling_config) if config.sampling_config is not None else None + self.supports_on_device_sampling = self.sampling is not None + self.mesh_device = config.mesh_device + self.tt_ccl = config.tt_ccl or get_tt_ccl(config.mesh_device) + self.vocab_size = config.vocab_size + self.n_layers = config.n_layers + self.num_devices = config.num_devices + self.decode_residual_memcfg = config.decode_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.prefill_residual_memcfg = config.prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtypes = config.activation_dtypes or [None] * config.n_layers + self.model_args = None + + @property + def n_kv_heads(self) -> int: + return self.config.block_configs[0].attention_config.n_kv_heads + + def iter_executor_named_modules(self): + if not hasattr(self, "layers"): + return + for index, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{index}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in self.config.block_configs) + for config in tuple({id(item): item for item in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, + block_size=block_size, + max_num_blocks=max_num_blocks, + ) + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for index, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {index} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {index} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = pair + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). Without a runtime + # last-token slice, the batched path returns the full hidden state (get_last_token == -1); + # the executor does per-slot extraction + norm/lm_head so those stages stay bit-identical. + x = x_embed + for layer in self.layers: + x = layer.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + + if get_last_token == -1 and last_token_slice is None: + return x + + # Slice + deallocate the full-sequence buffer before norm/LM head reduces peak L1. + old = x + if last_token_slice is None: + x_tile = _slice_last_token_tile(old, get_last_token) + else: + x_tile = ttnn.slice( + old, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(old.shape[2]) // 32, + ) + ttnn.deallocate(old) + if last_token_index is not None: + if x_tile.dtype != ttnn.bfloat16: + old_tile = x_tile + x_tile = ttnn.typecast(x_tile, ttnn.bfloat16) + ttnn.deallocate(old_tile) + old_tile = x_tile + x_tile = ttnn.embedding(last_token_index, x_tile, layout=ttnn.TILE_LAYOUT) + x_tile = ttnn.unsqueeze_to_4D(x_tile) + ttnn.deallocate(old_tile) + return self._last_tile_logits(x_tile) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + x = _slice_last_token_tile(hidden_states, last_token_idx) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + return self._last_tile_logits(x) + + def _last_tile_logits(self, x_tile: ttnn.Tensor) -> ttnn.Tensor: + """Final-norm + all-gather + LM-head on a 32-row tile. ``x_tile`` shape ``[1, 1, 32, dim]``.""" + x = self.norm.prefill_forward(x_tile) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + del last_token_slice, last_token_index + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), + ) + x = ttnn.matmul( + selector, + hidden_states, + dtype=ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + compute_kernel_config=ttnn.init_device_compute_kernel_config( + self.mesh_device.arch(), + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ), + ) + ttnn.deallocate(selector) + return self._last_tile_logits(x) + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for layer in self.layers: + x = layer.decode_forward(x, current_pos, rot_mats, page_table=page_table) + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + return self.lm_head.forward(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) + + def prefill_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, start_pos: int = 0, user_id: int = 0) -> ttnn.Tensor: + """Legacy path: embed + RoPE + blocks + final norm (no page table). For tests only.""" + x = self.embed_prefill(token_ids_tt) + seq_len = x.shape[2] + assert seq_len % 128 == 0, "prefill seq_len must be divisible by 128" + rot = self.rope_setup.prefill_forward(start_pos, seq_len) + h = x + for layer in self.layers: + h = layer.prefill_forward(h, rot, user_id=user_id, page_table=None) + h = self.norm.prefill_forward(h) + return _all_gather_rmsnorm_tensor(self.norm, h) + + def decode_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, current_pos: int) -> ttnn.Tensor: + """Legacy path: single-token decode without paged ``page_table``.""" + x = self.embed.forward(token_ids_tt) + x = ttnn.unsqueeze_to_4D(x) + pos = torch.tensor([current_pos], dtype=torch.long) + rot_idxs = prepare_rot_idxs(self.rope_setup.config, pos, on_host=False) + rot = self.rope_setup.decode_forward(rot_idxs) + cur = ttnn.from_torch( + pos, + device=self.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(self.mesh_device), + ) + h = x + for layer in self.layers: + h = layer.decode_forward(h, cur, rot, page_table=None) + h = _all_gather_rmsnorm_tensor(self.norm, h, memory_config=self.norm.config.decode_memory_config) + return self.norm.forward(h, "decode") + + def lm_logits(self, hidden: ttnn.Tensor) -> ttnn.Tensor: + """Project last hidden to logits (vocab-sharded on multi-device). + + Skip the explicit interleaved→shard if the caller already produced a sharded + input (``decode_from_token_ids`` returns the decode-mode RMSNorm's width-sharded + output, which already matches ``LMHead1D.config.input_memcfg``). + """ + x = hidden + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded() and not x.memory_config().is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + return self.lm_head.forward(x) + + +__all__ = [ + "QWEN25_7B_ACCURACY", + "QWEN25_7B_PERFORMANCE", + "Qwen25PagedAttentionConfig", + "Qwen25_7B", + "Qwen25_7BDecoderLayer", + "Qwen25_7BDecoderLayerConfig", + "Qwen25_7BLayerWeights", + "Qwen25_7BModelParameters", + "Qwen25_7BPrecisionConfig", + "Qwen25_7BTransformerConfig", + "Qwen25_7BWeights", + "build_qwen25_7b_transformer_config", +] diff --git a/code/models/common/models/qwen25_7b/weight_utils.py b/code/models/common/models/qwen25_7b/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..62ecf0d9cfd14583e96e3b4e946caa3615693395 --- /dev/null +++ b/code/models/common/models/qwen25_7b/weight_utils.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side HF weight layout helpers for Qwen2.x (TTTv2 port). + +Q/K linear weights use HuggingFace head layout; TTNN RoPE expects Meta layout. +This duplicates the small permute helpers from ``load_checkpoints.reverse_permute`` +without importing ``models/tt_transformers``. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.modules.lazy_weight import LazyWeight + + +def reverse_permute(tensor: torch.Tensor, n_heads: int, dim1: int, dim2: int) -> torch.Tensor: + """HF Q/K weight rows → Meta layout (RoPE-compatible).""" + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def reverse_permute_1d(tensor: torch.Tensor) -> torch.Tensor: + """Undo HF split of last dim into interleaved complex pairs for norm/bias vectors.""" + shape = tensor.shape + dim = shape[-1] + assert dim % 2 == 0, "Last dimension must be even" + reals = tensor[..., : dim // 2] + imags = tensor[..., dim // 2 :] + interleaved = torch.stack((reals, imags), dim=-1).flatten(start_dim=len(shape) - 1) + return interleaved + + +def permute_hf_rope_to_meta_tables(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert HF cos/sin tables to Meta-style tables for ``RotarySetup1D`` (see attention tests).""" + if len(cos.shape) == 3: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + return cos, sin + + +def build_rope_cos_sin_torch( + rotary_emb: Any, table_len: int, head_dim: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Build ``[1, 1, table_len, head_dim]`` cos/sin tensors (Meta layout) from HF rotary module.""" + x = torch.zeros(1, 1, table_len, head_dim, dtype=dtype) + position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0) + with torch.no_grad(): + cos_hf, sin_hf = rotary_emb(x, position_ids) + cos_m, sin_m = permute_hf_rope_to_meta_tables(cos_hf.float(), sin_hf.float()) + return cos_m.to(dtype), sin_m.to(dtype) + + +def attention_wqkv_wo_from_hf_layer( + self_attn: Any, + num_devices: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Mirror ``get_attention_weights_from_ref_model`` in ``test_attention_1d.py`` (HF module API).""" + wq_raw = self_attn.q_proj.weight + wk_raw = self_attn.k_proj.weight + wv_raw = self_attn.v_proj.weight + wo_raw = self_attn.o_proj.weight + + dim = wq_raw.shape[1] + cfg = self_attn.config + n_heads = cfg.num_attention_heads + n_kv_heads = cfg.num_key_value_heads + head_dim = cfg.hidden_size // n_heads + n_heads_times_head_dim = n_heads * head_dim + n_kv_heads_times_head_dim = n_kv_heads * head_dim + + wq_meta = reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim) + wk_meta = reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim) + wv_meta = wv_raw + wo_meta = wo_raw + + wq = wq_meta.T + wk = wk_meta.T + wv = wv_meta.T + wo = wo_meta.T + + qkv_list = [] + for i in range(num_devices): + wq_chunk = torch.chunk(wq, num_devices, dim=1)[i] + wk_chunk = torch.chunk(wk, num_devices, dim=1)[i] + wv_chunk = torch.chunk(wv, num_devices, dim=1)[i] + qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1) + qkv_list.append(qkv) + wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0).clone() + wo = wo.unsqueeze(0).unsqueeze(0).clone() + + q_norm = None + k_norm = None + if getattr(self_attn, "q_norm", None) is not None: + q_norm = reverse_permute_1d(self_attn.q_norm.weight).clone() + if getattr(self_attn, "k_norm", None) is not None: + k_norm = reverse_permute_1d(self_attn.k_norm.weight).clone() + + wqkv_bias = None + if getattr(self_attn.q_proj, "bias", None) is not None: + bq_raw = self_attn.q_proj.bias + bk_raw = self_attn.k_proj.bias + bv_raw = self_attn.v_proj.bias + bq_meta = reverse_permute_1d(bq_raw.view(n_heads, head_dim)).view(-1) + bk_meta = reverse_permute_1d(bk_raw.view(n_kv_heads, head_dim)).view(-1) + bv_meta = bv_raw + qkv_bias_list = [] + for i in range(num_devices): + bq_chunk = torch.chunk(bq_meta, num_devices, dim=0)[i] + bk_chunk = torch.chunk(bk_meta, num_devices, dim=0)[i] + bv_chunk = torch.chunk(bv_meta, num_devices, dim=0)[i] + qkv_bias_list.append(torch.cat([bq_chunk, bk_chunk, bv_chunk], dim=-1)) + wqkv_bias = torch.cat(qkv_bias_list, dim=-1).clone() + + return wqkv, wo, q_norm, k_norm, wqkv_bias + + +def mlp_weights_from_hf_layer(mlp: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (w1, w2, w3) TTNN layouts: gate^T, down^T, up^T for ``MLP1D``.""" + w1 = mlp.gate_proj.weight.T.contiguous().clone() + w3 = mlp.up_proj.weight.T.contiguous().clone() + w2 = mlp.down_proj.weight.T.contiguous().clone() + return w1, w2, w3 + + +def rms_weight_torch(layernorm: Any) -> torch.Tensor: + return layernorm.weight.detach().float().clone() + + +def embed_tokens_torch(embed: Any) -> torch.Tensor: + w = embed.weight.detach().float().clone() + return w.unsqueeze(0).unsqueeze(0) + + +def build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """ + Column-split LM head for ``LMHead1D`` (adapted from ``LMHead1D.from_model_args`` math, no v1 args). + ``lm_head_weight`` shape: ``[vocab_size, dim]`` (HF). + + Returns ``(lazy_weights, split_sizes, weights_memcfgs)`` for ``LMHead1DConfig`` DRAM-sharded matmuls. + """ + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / 32) * 32 + if vocab_size < padded_vocab_size: + pad = padded_vocab_size - vocab_size + torch_w = torch.cat([torch_w, torch.zeros(torch_w.shape[0], pad, dtype=torch_w.dtype)], dim=-1) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + tile = ttnn.TILE_SIZE + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (tile * dram_size.x)) * (tile * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights: list[LazyWeight] = [] + weights_memcfgs: list[ttnn.MemoryConfig] = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_splits.append(torch_w[:, start:end]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + name = f"lm_head_split_{i}_{combined.shape[-1]}" + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, name) if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs diff --git a/code/models/common/models/qwen25_coder_32b/README.md b/code/models/common/models/qwen25_coder_32b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..012ac37288d5f8cec572f377eae33ea1073b4e97 --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/README.md @@ -0,0 +1,46 @@ +# Qwen2.5-Coder-32B with TTTv2 + +This directory contains the Qwen2.5-Coder-32B TTTv2 product path. + +## Construction path + +```text +HF checkpoint + -> hf_adaptor.py: provider metadata, tokenizer, and weights + -> model.py: Qwen2.5-Coder tensor graph + -> executor.py: thin family builder plus compatibility/direct-run helpers + -> generator.py: vLLM construction, DP composition, and dispatch +``` + +The model uses reusable embedding, rotary, RMSNorm, attention, MLP, LM-head, +optional sampling, and collective modules. Coder architecture, dimensions, +precision, and tuning remain model-owned. + +## Executor composition + +The lifecycle is shared through `models/common/models/qwen2_executor.py`, which +configures `models/common/models/executor.py::ModelExecutor`. The common owner +constructs paged KV, output reading, prefill/decode runtimes, eager/traced +execution, warmup, and ordered cleanup. + +The model-local file retains its historical executor/config/builder names, +eager/traced compatibility wrappers, compatibility config construction, +direct-run helpers, concrete last-token slicing, and TT deallocation behavior. +It uses plain coordinator prefill warmup and does not adopt Qwen3 native +sampling state during this structural extraction. + +## vLLM, DP, and ownership + +The generator owns vLLM normalization/dispatch and constructs one executor per +lane. `LaneGroupExecutor` provides DP fanout. Lane executors exclusively own TT +resources and cleanup. + +## Tests + +- `models/common/tests/models/qwen25_coder_32b/test_hf_adaptor.py` +- `models/common/tests/models/qwen25_coder_32b/test_model_runtime_surface.py` +- `models/common/tests/models/qwen25_coder_32b/test_demo_contract.py` +- `models/common/tests/demos/qwen25_coder_32b/demo.py` +- `models/common/tests/models/test_qwen2_executor_family.py` +- `models/common/tests/llm_runtime/test_executor_integration.py` +- `models/common/tests/llm_runtime/test_model_executor.py` diff --git a/code/models/common/models/qwen25_coder_32b/__init__.py b/code/models/common/models/qwen25_coder_32b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c951b892b9493e6655c86c1eca7cce3d6723e375 --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/__init__.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +from models.common.models.qwen25_coder_32b.executor import ( + EagerQwen25Coder32BExecutor, + Qwen25Coder32BExecutor, + Qwen25Coder32BExecutorConfig, + TracedQwen25Coder32BExecutor, + build_qwen25_coder_32b_executor, +) +from models.common.models.qwen25_coder_32b.generator import ( + Qwen25Coder32BGenerator, + Qwen25Coder32BGeneratorConfig, + build_qwen25_coder_32b_generator, +) +from models.common.models.qwen25_coder_32b.hf_adaptor import Qwen25Coder32BForCausalLM, from_pretrained +from models.common.models.qwen25_coder_32b.model import ( + QWEN25_CODER_32B_ACCURACY, + QWEN25_CODER_32B_PERFORMANCE, + Qwen25Coder32B, + Qwen25Coder32BConfig, + Qwen25Coder32BLayerWeights, + Qwen25Coder32BPagedAttentionConfig, + Qwen25Coder32BPrecisionConfig, + Qwen25Coder32BWeights, + build_qwen25_coder_32b_model, +) + +__all__ = [ + "Qwen25Coder32B", + "Qwen25Coder32BConfig", + "Qwen25Coder32BPagedAttentionConfig", + "Qwen25Coder32BPrecisionConfig", + "Qwen25Coder32BLayerWeights", + "Qwen25Coder32BWeights", + "build_qwen25_coder_32b_model", + "QWEN25_CODER_32B_ACCURACY", + "QWEN25_CODER_32B_PERFORMANCE", + "Qwen25Coder32BGenerator", + "Qwen25Coder32BGeneratorConfig", + "build_qwen25_coder_32b_generator", + "Qwen25Coder32BExecutor", + "Qwen25Coder32BExecutorConfig", + "build_qwen25_coder_32b_executor", + "EagerQwen25Coder32BExecutor", + "TracedQwen25Coder32BExecutor", + "Qwen25Coder32BForCausalLM", + "from_pretrained", +] diff --git a/code/models/common/models/qwen25_coder_32b/demo.py b/code/models/common/models/qwen25_coder_32b/demo.py new file mode 100644 index 0000000000000000000000000000000000000000..2bb8240c112a6bff4098449ae343518566032fcd --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/demo.py @@ -0,0 +1,594 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Smoke and accuracy tests for ``Qwen25Coder32B`` on T3K. + +Requires Hugging Face weights locally (or network for fresh download). Uses +``MESH_DEVICE=T3K`` only; the port raises on any other mesh. + +Run (T3K, internal KV smoke, 1-layer fast iteration):: + + MESH_DEVICE=T3K HF_MODEL=Qwen/Qwen2.5-Coder-32B-Instruct \\ + QWEN25_CODER_32B_DEMO_NUM_LAYERS=1 \\ + pytest models/common/models/qwen25_coder_32b/demo.py -v -k prefill_smoke + +Executor + paged KV:: + + MESH_DEVICE=T3K HF_MODEL=Qwen/Qwen2.5-Coder-32B-Instruct \\ + QWEN25_CODER_32B_DEMO_NUM_LAYERS=1 \\ + pytest models/common/models/qwen25_coder_32b/demo.py -v -k executor_prefill + +``Attention1D`` shards heads across the mesh: 40 / 8 = 5 attention heads per device and +8 / 8 = 1 KV head per device. +""" + +from __future__ import annotations + +import os + +import pytest +import torch +from transformers import AutoConfig, AutoModelForCausalLM + +import ttnn +from models.common.models.qwen25_coder_32b.executor import ( + EagerQwen25Coder32BExecutor, + TracedQwen25Coder32BExecutor, + run_lm_head, + run_prefill, +) +from models.common.models.qwen25_coder_32b.generator import greedy_argmax_from_logits, greedy_decode_one_step +from models.common.models.qwen25_coder_32b.model import Qwen25Coder32B +from models.common.tests.demos.cleanup_utils import cleanup_model_case +from models.common.tests.demos.run_helpers import make_contiguous_page_table +from models.common.utility_functions import comp_pcc + + +def _skip_unless_t3k(mesh_device: ttnn.MeshDevice, hf_model_id: str) -> None: + """Port targets T3K only; head divisibility is satisfied (40/8, 8/8).""" + n_dev = mesh_device.get_num_devices() + if n_dev != 8: + pytest.skip( + f"Qwen2.5-Coder-32B-Instruct port targets T3K (8 devices) only; got {n_dev}. " f"Set MESH_DEVICE=T3K." + ) + cfg = AutoConfig.from_pretrained(hf_model_id) + n_h, n_kv = cfg.num_attention_heads, cfg.num_key_value_heads + if n_h % n_dev != 0 or n_kv % n_dev != 0: + pytest.skip( + f"Incompatible mesh for {hf_model_id}: {n_dev} devices need " + f"num_attention_heads ({n_h}) and num_key_value_heads ({n_kv}) each divisible by {n_dev}." + ) + + +@pytest.fixture +def device_params(request, galaxy_type): + """Match ``models/tt_transformers/conftest.py`` so ``fabric_config: True`` maps to a real fabric.""" + params = getattr(request, "param", {}).copy() + + mesh_device = {"N150": (1, 1), "N300": (1, 2), "N150x4": (1, 4), "T3K": (1, 8), "TG": (8, 4)}.get( + os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids()) + ) + is_single_device = (mesh_device == (1, 1)) if isinstance(mesh_device, tuple) else (mesh_device == 1) + + if "fabric_config" in params: + if is_single_device: + params["fabric_config"] = None + elif params["fabric_config"] is True: + params["fabric_config"] = ( + ttnn.FabricConfig.FABRIC_1D_RING if galaxy_type == "6U" else ttnn.FabricConfig.FABRIC_1D + ) + + return params + + +pytestmark = [ + pytest.mark.parametrize( + "mesh_device", + [{"T3K": (1, 8)}.get(os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids()))], + indirect=True, + ), + pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True), +] + + +@pytest.fixture(scope="module") +def hf_model_id(): + return os.environ.get("HF_MODEL", "Qwen/Qwen2.5-Coder-32B-Instruct") + + +_slow = pytest.mark.slow + + +@_slow +@pytest.mark.parametrize("seq_len", [128]) +def test_qwen25_coder_32b_prefill_smoke(mesh_device, hf_model_id, seq_len: int, tmp_path_factory): + """One prefill pass (truncated layers via env) and LM head → greedy token (internal KV).""" + _skip_unless_t3k(mesh_device, hf_model_id) + num_layers = int(os.environ.get("QWEN25_CODER_32B_DEMO_NUM_LAYERS", "1")) + cache = tmp_path_factory.mktemp("qwen25_coder_32b_cache") + model = None + try: + model = Qwen25Coder32B.from_pretrained( + mesh_device, + hf_model_id, + max_batch_size=32, + max_seq_len=max(2048, seq_len), + num_layers=num_layers, + cache_dir=cache, + executor_mode=False, + ) + except Exception as e: + pytest.skip(f"Could not build model (weights / memory): {e}") + + ttnn.SetDefaultDevice(mesh_device) + try: + toks = torch.zeros(1, 1, 1, seq_len, dtype=torch.int32) + toks[..., :4] = torch.tensor([1, 2, 3, 4], dtype=torch.int32).view(1, 1, 1, 4) + x = ttnn.from_torch( + toks, + device=mesh_device, + dtype=ttnn.uint32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(mesh_device), + ) + hidden = run_prefill(model, x, start_pos=0) + logits = run_lm_head(model, hidden) + _ = greedy_argmax_from_logits(logits, mesh_device=mesh_device) + finally: + ttnn.SetDefaultDevice(None) + cleanup_model_case(model, mesh_device) + + +@_slow +def test_qwen25_coder_32b_decode_one_step(mesh_device, hf_model_id, tmp_path_factory): + """Prefill 128 tokens, then one greedy decode step at position 128. + + Single-user (``max_batch_size=1``): the legacy ``decode_from_token_ids`` path only + fills one KV slot per call, so the non-paged ``paged_update_cache`` validation + (``num_indices == batch_size``) requires the model be built at batch=1. + """ + _skip_unless_t3k(mesh_device, hf_model_id) + num_layers = int(os.environ.get("QWEN25_CODER_32B_DEMO_NUM_LAYERS", "1")) + cache = tmp_path_factory.mktemp("qwen25_coder_32b_decode_cache") + seq_len = 128 + model = None + try: + model = Qwen25Coder32B.from_pretrained( + mesh_device, + hf_model_id, + max_batch_size=1, + max_seq_len=max(512, seq_len + 8), + num_layers=num_layers, + cache_dir=cache, + executor_mode=False, + ) + except Exception as e: + pytest.skip(f"Could not build model (weights / memory): {e}") + + ttnn.SetDefaultDevice(mesh_device) + try: + toks = torch.zeros(1, 1, 1, seq_len, dtype=torch.int32) + toks[..., :4] = torch.tensor([1, 2, 3, 4], dtype=torch.int32).view(1, 1, 1, 4) + x = ttnn.from_torch( + toks, + device=mesh_device, + dtype=ttnn.uint32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(mesh_device), + ) + _ = run_prefill(model, x, start_pos=0) + _ = greedy_decode_one_step(model, token_id=64, current_pos=seq_len) + finally: + ttnn.SetDefaultDevice(None) + cleanup_model_case(model, mesh_device) + + +@_slow +@pytest.mark.parametrize("seq_len", [128]) +def test_qwen25_coder_32b_executor_prefill_smoke(mesh_device, hf_model_id, seq_len: int, tmp_path_factory): + """``EagerQwen25Coder32BExecutor`` + paged KV: prefill returns host logits.""" + _skip_unless_t3k(mesh_device, hf_model_id) + num_layers = int(os.environ.get("QWEN25_CODER_32B_DEMO_NUM_LAYERS", "1")) + cache = tmp_path_factory.mktemp("qwen25_coder_32b_exec_cache") + model = None + try: + model = Qwen25Coder32B.from_pretrained( + mesh_device, + hf_model_id, + max_batch_size=1, + max_seq_len=max(2048, seq_len), + num_layers=num_layers, + cache_dir=cache, + executor_mode=True, + ) + except Exception as e: + pytest.skip(f"Could not build model: {e}") + + ma = model.model_args + assert ma is not None + block_size = 32 + max_num_blocks = (ma.max_seq_len // block_size) * ma.max_batch_size + kv_shape = (max_num_blocks, ma.n_kv_heads // mesh_device.get_num_devices(), block_size, ma.head_dim) + + ttnn.SetDefaultDevice(mesh_device) + try: + ex = EagerQwen25Coder32BExecutor(model, mesh_device) + kv = ex.allocate_kv_cache(kv_shape, torch.bfloat16, ma.n_layers) + page_table = make_contiguous_page_table(1, ma.max_seq_len, block_size) + toks = torch.zeros(1, seq_len, dtype=torch.long) + toks[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.long) + logits = ex.prefill_forward(toks, page_table=page_table, kv_cache=kv) + assert logits.shape[0] == 1 and logits.shape[-1] == model.vocab_size + except Exception as e: + pytest.skip(f"Executor prefill not runnable: {e}") + finally: + ttnn.SetDefaultDevice(None) + cleanup_model_case(model, mesh_device) + + +@_slow +def test_qwen25_coder_32b_teacher_forcing_prefill_vs_hf(mesh_device, hf_model_id, tmp_path_factory): + """Last-token logits vs HF after prefill (PCC gate, full depth). + + Must run the same number of decoder layers as ``AutoModelForCausalLM``; + comparing a truncated TT stack (``QWEN25_CODER_32B_DEMO_NUM_LAYERS`` for smoke tests) + to full HF logits always fails PCC. Unset ``QWEN25_CODER_32B_DEMO_NUM_LAYERS`` here, + or set it equal to the checkpoint's ``num_hidden_layers``. + """ + _skip_unless_t3k(mesh_device, hf_model_id) + hf_cfg = AutoConfig.from_pretrained(hf_model_id) + n_hf = int(hf_cfg.num_hidden_layers) + env_layers = os.environ.get("QWEN25_CODER_32B_DEMO_NUM_LAYERS") + if env_layers is not None: + num_layers = int(env_layers) + if num_layers != n_hf: + pytest.skip( + "Teacher-forcing PCC is defined against the full HF forward; " + f"unset QWEN25_CODER_32B_DEMO_NUM_LAYERS to use all {n_hf} layers " + f"(currently QWEN25_CODER_32B_DEMO_NUM_LAYERS={num_layers})." + ) + else: + num_layers = n_hf + seq_len = 128 + cache = tmp_path_factory.mktemp("qwen25_coder_32b_tf") + model = None + try: + model = Qwen25Coder32B.from_pretrained( + mesh_device, + hf_model_id, + max_batch_size=1, + max_seq_len=max(512, seq_len), + num_layers=num_layers, + cache_dir=cache, + executor_mode=True, + ) + except Exception as e: + pytest.skip(f"Could not build model: {e}") + + ma = model.model_args + block_size = 32 + max_num_blocks = (ma.max_seq_len // block_size) * ma.max_batch_size + kv_shape = (max_num_blocks, ma.n_kv_heads // mesh_device.get_num_devices(), block_size, ma.head_dim) + + input_ids = torch.zeros(1, seq_len, dtype=torch.long) + input_ids[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.long) + + ttnn.SetDefaultDevice(mesh_device) + try: + hf = AutoModelForCausalLM.from_pretrained(hf_model_id, torch_dtype=torch.bfloat16) + hf.eval() + with torch.no_grad(): + hf_out = hf(input_ids[:, :seq_len]).logits[0, seq_len - 1, :].float() + + ex = EagerQwen25Coder32BExecutor(model, mesh_device) + kv = ex.allocate_kv_cache(kv_shape, torch.bfloat16, ma.n_layers) + page_table = make_contiguous_page_table(1, ma.max_seq_len, block_size) + tt_logits = ex.prefill_forward(input_ids[:, :seq_len], page_table=page_table, kv_cache=kv) + tt_vec = tt_logits[0, 0, :].float() + + ok, pcc = comp_pcc(hf_out, tt_vec, pcc=0.85) + assert ok, f"Prefill last-token PCC too low: {pcc}" + except Exception as e: + pytest.skip(f"Teacher-forcing PCC check not runnable: {e}") + finally: + ttnn.SetDefaultDevice(None) + cleanup_model_case(model, mesh_device) + + +@_slow +def test_qwen25_coder_32b_eager_traced_prefill_logits_match(mesh_device, hf_model_id, tmp_path_factory): + """``EagerQwen25Coder32BExecutor`` vs ``TracedQwen25Coder32BExecutor``: same prefill → host logits within tolerance. + + Builds two 32B models back-to-back — host RAM (~130 GB at bf16 for both HF copies) and + device cache will be tight. Use ``QWEN25_CODER_32B_DEMO_NUM_LAYERS=1`` for fast iteration. + """ + _skip_unless_t3k(mesh_device, hf_model_id) + num_layers = int(os.environ.get("QWEN25_CODER_32B_DEMO_NUM_LAYERS", "1")) + cache_a = tmp_path_factory.mktemp("qwen25_coder_32b_par_e") + cache_b = tmp_path_factory.mktemp("qwen25_coder_32b_par_t") + seq_len = 128 + m_e = m_t = None + try: + m_e = Qwen25Coder32B.from_pretrained( + mesh_device, + hf_model_id, + max_batch_size=1, + max_seq_len=512, + num_layers=num_layers, + cache_dir=cache_a, + executor_mode=True, + ) + m_t = Qwen25Coder32B.from_pretrained( + mesh_device, + hf_model_id, + max_batch_size=1, + max_seq_len=512, + num_layers=num_layers, + cache_dir=cache_b, + executor_mode=True, + ) + except Exception as e: + pytest.skip(f"Could not build models: {e}") + + ma = m_e.model_args + assert ma is not None + block_size = 32 + max_num_blocks = (ma.max_seq_len // block_size) * ma.max_batch_size + kv_shape = (max_num_blocks, ma.n_kv_heads // mesh_device.get_num_devices(), block_size, ma.head_dim) + + ttnn.SetDefaultDevice(mesh_device) + try: + eager_ex = EagerQwen25Coder32BExecutor(m_e, mesh_device) + traced_ex = TracedQwen25Coder32BExecutor(m_t, mesh_device) + kv_e = eager_ex.allocate_kv_cache(kv_shape, torch.bfloat16, ma.n_layers) + kv_t = traced_ex.allocate_kv_cache(kv_shape, torch.bfloat16, ma.n_layers) + page_table = make_contiguous_page_table(1, ma.max_seq_len, block_size) + toks = torch.zeros(1, seq_len, dtype=torch.long) + toks[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.long) + le = eager_ex.prefill_forward(toks, page_table=page_table, kv_cache=kv_e) + lt = traced_ex.prefill_forward(toks, page_table=page_table, kv_cache=kv_t) + diff = (le.float() - lt.float()).abs().max().item() + assert diff < 0.25, f"eager/traced prefill logits max abs diff too large: {diff}" + finally: + ttnn.SetDefaultDevice(None) + cleanup_model_case(m_e, mesh_device) + cleanup_model_case(m_t, mesh_device) + + +@_slow +def test_qwen25_coder_32b_numerical_divergence_vs_hf(mesh_device, hf_model_id, tmp_path_factory): + """Per-boundary PCC of TT prefill activations vs HF; isolate the first diverging op. + + Captures activations on both sides for the same 128-token prompt: + - embed output + - per layer i (0..N-1): post-input-norm, attention-out, attn-residual, post-attn-norm, + mlp-out, layer-out + - final-norm output + + Use the legacy internal-KV path (``prefill_from_token_ids``) — same residual stream + semantics as the executor path, simpler to instrument. Reports a CSV under + ``QWEN25_CODER_32B_NUMDIV_OUT`` (default ``/tmp/qwen25_coder_32b_numdiv.csv``). + + Override env: + QWEN25_CODER_32B_NUMDIV_LAYERS=N to limit layer count (debug). + QWEN25_CODER_32B_NUMDIV_SEQ=S to change seq_len (must be multiple of 128). + """ + import csv + import gc + + from models.common.models.qwen25_coder_32b.model import Qwen25Coder32BDecoderLayer, _all_gather_rmsnorm_tensor + + _skip_unless_t3k(mesh_device, hf_model_id) + hf_cfg = AutoConfig.from_pretrained(hf_model_id) + n_hf = int(hf_cfg.num_hidden_layers) + num_layers = int(os.environ.get("QWEN25_CODER_32B_NUMDIV_LAYERS", n_hf)) + if num_layers > n_hf: + num_layers = n_hf + seq_len = int(os.environ.get("QWEN25_CODER_32B_NUMDIV_SEQ", "128")) + assert seq_len % 128 == 0, f"seq_len must be multiple of 128, got {seq_len}" + + out_csv = os.environ.get("QWEN25_CODER_32B_NUMDIV_OUT", "/tmp/qwen25_coder_32b_numdiv.csv") + cache = tmp_path_factory.mktemp("qwen25_coder_32b_numdiv") + + input_ids = torch.zeros(1, seq_len, dtype=torch.long) + input_ids[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.long) + + # ---------- HF forward + per-layer hooks ---------- + hf_intermediates: dict[str, torch.Tensor] = {} + hf = AutoModelForCausalLM.from_pretrained(hf_model_id, torch_dtype=torch.bfloat16) + hf.eval() + + def _hook(name): + def fn(_mod, _inp, out): + t = out[0] if isinstance(out, tuple) else out + hf_intermediates[name] = t.detach().float().cpu() + + return fn + + handles = [hf.model.embed_tokens.register_forward_hook(_hook("embed"))] + for i, layer in enumerate(hf.model.layers[:num_layers]): + handles.append(layer.input_layernorm.register_forward_hook(_hook(f"L{i:02d}.input_norm"))) + handles.append(layer.self_attn.register_forward_hook(_hook(f"L{i:02d}.self_attn"))) + handles.append(layer.post_attention_layernorm.register_forward_hook(_hook(f"L{i:02d}.post_attn_norm"))) + handles.append(layer.mlp.register_forward_hook(_hook(f"L{i:02d}.mlp"))) + handles.append(layer.register_forward_hook(_hook(f"L{i:02d}.layer_out"))) + handles.append(hf.model.norm.register_forward_hook(_hook("final_norm"))) + + with torch.no_grad(): + if num_layers == n_hf: + hf(input_ids) + else: + x = hf.model.embed_tokens(input_ids) + position_ids = torch.arange(seq_len).unsqueeze(0) + cos, sin = hf.model.rotary_emb(x, position_ids) + attn_mask = torch.zeros(1, 1, seq_len, seq_len, dtype=x.dtype) + attn_mask.masked_fill_( + torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1), float("-inf") + ) + h = x + for layer in hf.model.layers[:num_layers]: + h = layer( + h, + attention_mask=attn_mask, + position_ids=position_ids, + position_embeddings=(cos, sin), + )[0] + h = hf.model.norm(h) + + for h_ in handles: + h_.remove() + del hf + gc.collect() + + # ---------- Build TT model ---------- + model = None + try: + model = Qwen25Coder32B.from_pretrained( + mesh_device, + hf_model_id, + max_batch_size=1, + max_seq_len=max(512, seq_len), + num_layers=num_layers, + cache_dir=cache, + executor_mode=False, + ) + except Exception as e: + pytest.skip(f"Could not build TT model: {e}") + + # ---------- Instrument TT decoder layer prefill_forward ---------- + tt_intermediates: dict[str, torch.Tensor] = {} + orig_layer_prefill = Qwen25Coder32BDecoderLayer.prefill_forward + layer_idx = [0] + + def _stash(name: str, t: ttnn.Tensor) -> None: + # Topology metadata on intermediate tensors is unreliable here (post-RS / post-AG paths can + # leave Shard(2) tags on tensors whose data is actually sharded on dim 3, etc.). Read each + # device shard directly: if all shards are bit-identical, the tensor is replicated → take one; + # otherwise concatenate along the last dim (the only fracturing axis these modules use). + shards = ttnn.get_device_tensors(t) + torch_shards = [ttnn.to_torch(s).detach().float().cpu() for s in shards] + if len(torch_shards) == 1: + host = torch_shards[0] + elif all(torch.equal(torch_shards[0], s) for s in torch_shards[1:]): + host = torch_shards[0] + else: + host = torch.cat(torch_shards, dim=-1) + tt_intermediates[name] = host + + def patched_layer_prefill( + self, + x, + rot_mats, + *, + user_id: int = 0, + page_table=None, + chunk_page_table=None, + chunk_start_idx=None, + ): + i = layer_idx[0] + r = self.input_layernorm.prefill_forward(x) + r = _all_gather_rmsnorm_tensor(self.input_layernorm, r) + _stash(f"L{i:02d}.input_norm", r) + r = self.self_attn.forward( + r, + None, + rot_mats, + mode="prefill", + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + ) + _stash(f"L{i:02d}.self_attn", r) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + r2 = self.post_attention_layernorm.prefill_forward(h) + r2 = _all_gather_rmsnorm_tensor(self.post_attention_layernorm, r2) + _stash(f"L{i:02d}.post_attn_norm", r2) + r2 = self.mlp.prefill_forward(r2) + _stash(f"L{i:02d}.mlp", r2) + out = ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + _stash(f"L{i:02d}.layer_out", out) + layer_idx[0] += 1 + return out + + Qwen25Coder32BDecoderLayer.prefill_forward = patched_layer_prefill + + ttnn.SetDefaultDevice(mesh_device) + try: + tokens_tt = ttnn.from_torch( + input_ids.view(1, 1, 1, seq_len).to(torch.int32), + device=mesh_device, + dtype=ttnn.uint32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(mesh_device), + ) + x_embed = model.embed_prefill(tokens_tt) + _stash("embed", x_embed) + rot = model.rope_setup.prefill_forward(0, seq_len) + h_tt = x_embed + for layer in model.layers: + h_tt = layer.prefill_forward(h_tt, rot, user_id=0, page_table=None) + h_tt = model.norm.prefill_forward(h_tt) + h_tt = _all_gather_rmsnorm_tensor(model.norm, h_tt) + _stash("final_norm", h_tt) + finally: + Qwen25Coder32BDecoderLayer.prefill_forward = orig_layer_prefill + ttnn.SetDefaultDevice(None) + + # ---------- Compare ---------- + rows = [] + keys = sorted(set(hf_intermediates) & set(tt_intermediates)) + ordering = {"embed": 0, "final_norm": 999} + boundary_order = ["input_norm", "self_attn", "post_attn_norm", "mlp", "layer_out"] + + def _sort_key(k: str) -> tuple: + if k in ordering: + return (ordering[k], 0, "") + parts = k.split(".") + if len(parts) == 2 and parts[0].startswith("L"): + try: + li = int(parts[0][1:]) + bi = boundary_order.index(parts[1]) if parts[1] in boundary_order else 99 + return (10 + li, bi, parts[1]) + except ValueError: + pass + return (1000, 0, k) + + keys.sort(key=_sort_key) + + print("\n=== qwen25_coder_32b numerical divergence vs HF ===") + print(f" seq_len={seq_len} num_layers={num_layers} mesh={mesh_device.get_num_devices()}") + print(f" {'boundary':35s} {'PCC':>10s} {'max_abs_diff':>14s} {'shape':>20s}") + for k in keys: + hf_t = hf_intermediates[k] + tt_t = tt_intermediates[k] + raw_tt_shape = tuple(tt_t.shape) + while tt_t.dim() > hf_t.dim(): + tt_t = tt_t.squeeze(0) + if tt_t.shape[-1] == 2 * hf_t.shape[-1]: + tt_t = tt_t[..., : hf_t.shape[-1]] + if tt_t.shape[-2] > hf_t.shape[-2]: + tt_t = tt_t[..., : hf_t.shape[-2], :] + if hf_t.shape[-2] > tt_t.shape[-2]: + hf_t = hf_t[..., : tt_t.shape[-2], :] + if tt_t.shape != hf_t.shape: + print( + f" {k:35s} {'SHAPE_MISMATCH':>10s} raw_tt={raw_tt_shape} tt={tuple(tt_t.shape)} hf={tuple(hf_t.shape)}" + ) + rows.append((k, float("nan"), float("nan"), tuple(tt_t.shape))) + continue + _, pcc = comp_pcc(hf_t, tt_t, pcc=0.99) + max_abs = (hf_t.float() - tt_t.float()).abs().max().item() + rows.append((k, float(pcc), float(max_abs), tuple(tt_t.shape))) + print(f" {k:35s} {pcc:>10.6f} {max_abs:>14.4e} {str(tuple(tt_t.shape)):>20s}") + + with open(out_csv, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["boundary", "pcc", "max_abs_diff", "shape"]) + for row in rows: + w.writerow([row[0], f"{row[1]:.6f}", f"{row[2]:.6e}", str(row[3])]) + print(f" -> wrote {out_csv}") + + if model is not None: + cleanup_model_case(model, mesh_device) diff --git a/code/models/common/models/qwen25_coder_32b/executor.py b/code/models/common/models/qwen25_coder_32b/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..58204f805de46981714abedc43e14bcca35a1b42 --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/executor.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen2.5-Coder-32B construction and compatibility entry points.""" + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, WarmupConfig +from models.common.models.qwen2_executor import Qwen25Coder32BExecutor, Qwen25Coder32BExecutorConfig +from models.common.models.qwen2_executor import build_qwen25_coder_32b_executor as _build_qwen25_coder_32b_executor +from models.common.models.qwen25_coder_32b.hf_adaptor import Qwen25Coder32BForCausalLM +from models.common.models.qwen25_coder_32b.model import _slice_last_token_tile + + +def build_qwen25_coder_32b_executor( + llm: Qwen25Coder32BForCausalLM, + config: Qwen25Coder32BExecutorConfig, +) -> Qwen25Coder32BExecutor: + return _build_qwen25_coder_32b_executor(llm, config) + + +def _compat_executor_config(model, *, trace_mode: str, device_sampling_enabled: bool) -> Qwen25Coder32BExecutorConfig: + runtime_config = getattr(model, "model_args", None) + if runtime_config is None: + raise ValueError("Qwen25Coder32B compatibility executor requires model.model_args") + block_size = 32 + max_seq_len = int(getattr(runtime_config, "max_seq_len", model.config.max_seq_len)) + max_batch_size = int(getattr(runtime_config, "max_batch_size", model.config.max_batch_size)) + max_num_blocks = ((max_seq_len + block_size - 1) // block_size) * max_batch_size + return Qwen25Coder32BExecutorConfig( + trace=TraceConfig(mode=trace_mode), + warmup=WarmupConfig( + prefill_seq_lens=tuple(getattr(runtime_config, "trace_prefill_supported_seq_lens", (128, 1024))), + prefill_batch_sizes=(1,), + include_decode_top_k=device_sampling_enabled, + ), + paged_kv_cache=PagedKVCacheConfig( + block_size=block_size, + max_num_blocks=max_num_blocks, + dtype=getattr(runtime_config, "kv_cache_dtype", ttnn.bfloat8_b), + ), + device_sampling_enabled=device_sampling_enabled, + ) + + +class EagerQwen25Coder32BExecutor(Qwen25Coder32BExecutor): + """Compatibility wrapper over the shared eager runtime.""" + + def __init__(self, model, mesh_device): + del mesh_device + super().__init__( + model, + model.model_args, + _compat_executor_config(model, trace_mode="none", device_sampling_enabled=False), + ) + + +class TracedQwen25Coder32BExecutor(Qwen25Coder32BExecutor): + """Compatibility wrapper over the shared traced runtime.""" + + def __init__( + self, + model, + mesh_device, + ondevice_decode_loop: bool = False, + fast_prefill_last_token: bool = False, + ): + del mesh_device, fast_prefill_last_token + super().__init__( + model, + model.model_args, + _compat_executor_config( + model, + trace_mode="all", + device_sampling_enabled=bool(ondevice_decode_loop), + ), + ) + + +def run_prefill(model, token_ids_tt, *, start_pos: int = 0): + return model.prefill_from_token_ids(token_ids_tt, start_pos=start_pos) + + +def run_decode(model, token_id_tt, *, current_pos: int): + return model.decode_from_token_ids(token_id_tt, current_pos=current_pos) + + +def run_lm_head(model, hidden_tt): + if len(hidden_tt.shape) == 4 and hidden_tt.shape[2] > 32: + old = hidden_tt + hidden_tt = _slice_last_token_tile(old, hidden_tt.shape[2] - 1) + ttnn.deallocate(old) + return model.lm_logits(hidden_tt) diff --git a/code/models/common/models/qwen25_coder_32b/generator.py b/code/models/common/models/qwen25_coder_32b/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..c8cba975ca91690eaca95ee91a4d3d724ed910a5 --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/generator.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Qwen2.5-Coder-32B.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.qwen25_coder_32b.executor import Qwen25Coder32BExecutorConfig, build_qwen25_coder_32b_executor +from models.common.models.qwen25_coder_32b.hf_adaptor import DEFAULT_HF_REVISION, from_pretrained +from models.common.models.qwen25_coder_32b.model import Qwen25Coder32BPagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Qwen25Coder32BGeneratorConfig: + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + hf_revision: str | None = DEFAULT_HF_REVISION + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Qwen25Coder32BGenerator: + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + device_sampling_enabled: bool = True, + ): + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_qwen25_coder_32b_generator( + Qwen25Coder32BGeneratorConfig( + hf_model=str(hf_model), + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def prefill_forward( + self, + tokens, + page_table, + *, + enable_trace: bool, + prompt_lens=None, + start_pos=None, + empty_slots=None, + kv_cache=None, + sampling_params=None, + **compatibility_kwargs, + ): + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens, + start_pos, + page_table, + *, + enable_trace: bool, + kv_cache=None, + sampling_params=None, + reset_batch=False, + read_from_device=True, + **compatibility_kwargs, + ): + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward(execution=execution, read_from_device=read_from_device, **normalized) + + def read_decode_output(self, tt_out: Any, *, async_read: bool = False) -> Any: + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host(self, tt_out: Any, *, is_tokens: bool = False) -> tuple[Any, Any]: + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill(self, *, kv_cache, can_sample_on_device: bool, enable_trace: bool) -> None: + return self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def warmup_model_decode( + self, *, kv_cache, max_batch_size: int, num_blocks: int, can_sample_on_device: bool, enable_trace: bool + ) -> None: + return self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self): + return self.target.cleanup() + + def _select_prefill_execution(self, normalized: NormalizedPrefillKwargs, trace_requested: bool): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_qwen25_coder_32b_generator(config: Qwen25Coder32BGeneratorConfig) -> Qwen25Coder32BGenerator: + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Qwen25Coder32BPagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + hf_revision=config.hf_revision, + instruct="Instruct" in config.hf_model, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Qwen25Coder32BExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(include_decode_top_k=config.device_sampling_enabled), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_qwen25_coder_32b_executor(llm, executor_config)) + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Qwen25Coder32BGenerator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Qwen2.5-Coder-32B model must contain at least one attention layer") + attention_configs = tuple(_attention_config(layer) for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Qwen layer must expose the same KV head shape") + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _attention_config(layer): + attention = getattr(layer, "attention", None) or getattr(layer, "self_attn", None) + if attention is not None: + return attention.config + if hasattr(layer, "config"): + return layer.config + raise AttributeError("Qwen2.5-Coder-32B layer must expose attention/self_attn.config or config") + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/qwen25_coder_32b/hf_adaptor.py b/code/models/common/models/qwen25_coder_32b/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..341d7c6fabe9b9bfa72eaf4f26a0cc2e5afdf123 --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/hf_adaptor.py @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face product adaptor for Qwen2.5-Coder-32B-Instruct.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.models.qwen25_coder_32b import weight_utils +from models.common.models.qwen25_coder_32b.model import ( + DEFAULT_HF_REVISION, + QWEN25_CODER_32B_ACCURACY, + QWEN25_CODER_32B_PERFORMANCE, + Qwen25Coder32B, + Qwen25Coder32BConfig, + Qwen25Coder32BLayerWeights, + Qwen25Coder32BPagedAttentionConfig, + Qwen25Coder32BPrecisionConfig, + Qwen25Coder32BWeights, + build_qwen25_coder_32b_model, +) + +DEFAULT_HF_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct" + + +def _local_files_only() -> bool: + return any( + os.getenv(name, "").lower() in {"1", "true", "yes"} for name in ("CI", "HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ) + + +@dataclass(frozen=True) +class Qwen25Coder32BGenerationConfig: + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Qwen25Coder32BRuntimeConfig: + model_name: str + model_cache_path: Path | None + max_prefill_chunk_size: int + max_context_len: int + max_seq_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] + n_layers: int + n_kv_heads: int + head_dim: int + max_batch_size: int + cluster_shape: list[int] + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + if num_cached_tokens != 0: + return False + return ( + prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + chat = [] + if isinstance(prompt_text, str): + if system_prompt_text: + chat.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + chat.append({"role": "user", "content": prompt_text}) + else: + chat = prompt_text + try: + return _chat_template_ids(tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True)) + except ValueError: + pass + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Qwen25Coder32BForCausalLM: + model: Qwen25Coder32B + tokenizer: Any + runtime_config: Qwen25Coder32BRuntimeConfig + instruct: bool = True + generation_config: Qwen25Coder32BGenerationConfig = field(default_factory=Qwen25Coder32BGenerationConfig) + + def __post_init__(self) -> None: + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + self.generation_config = Qwen25Coder32BGenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=tuple(getattr(self.tokenizer, "stop_tokens", ()) or ()), + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + return encode_prompt( + self.tokenizer, + prompt_text, + system_prompt_text, + instruct=self.instruct if instruct is None else instruct, + ) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_tokenizer(hf_model: str, hf_revision: str | None = DEFAULT_HF_REVISION): + tokenizer = AutoTokenizer.from_pretrained( + hf_model, + revision=hf_revision, + trust_remote_code=True, + local_files_only=_local_files_only(), + ) + tokenizer.stop_tokens = list(_qwen_stop_token_ids(tokenizer)) + return tokenizer + + +def _qwen_stop_token_ids(tokenizer) -> tuple[int, ...]: + ids = [] + eos = getattr(tokenizer, "eos_token_id", None) + if eos is not None: + ids.extend([eos] if isinstance(eos, int) else list(eos)) + for token in ("<|im_end|>", "<|im_start|>"): + token_id = tokenizer.convert_tokens_to_ids(token) + if isinstance(token_id, int) and token_id >= 0: + ids.append(token_id) + return tuple(dict.fromkeys(int(value) for value in ids)) + + +def _trace_seq_lens(num_devices: int, max_prefill_chunk_size: int, max_seq_len: int) -> tuple[int, ...]: + if num_devices != 8: + raise ValueError(f"Qwen2.5-Coder-32B supports exactly 8 devices (T3K), got {num_devices}") + return tuple(length for length in (128, 1024) if length <= min(max_prefill_chunk_size, max_seq_len)) + + +def _cache_path(hf_model: str, mesh_device, cache_dir: Path | str | None) -> Path: + if cache_dir is not None: + path = Path(cache_dir) + elif os.getenv("TT_CACHE_PATH"): + path = Path(os.environ["TT_CACHE_PATH"]) + else: + path = Path("model_cache") / hf_model / "T3K" + path.mkdir(parents=True, exist_ok=True) + return path + + +def _validate_checkpoint_config(hf_config, *, num_devices: int) -> None: + expected = { + "num_hidden_layers": 64, + "hidden_size": 5120, + "num_attention_heads": 40, + "num_key_value_heads": 8, + "intermediate_size": 27648, + "vocab_size": 152064, + } + actual = {name: getattr(hf_config, name, None) for name in expected} + mismatches = {name: (actual[name], value) for name, value in expected.items() if actual[name] != value} + if mismatches: + raise ValueError(f"Unexpected Qwen2.5-Coder-32B geometry: {mismatches}") + if hf_config.num_attention_heads % num_devices or hf_config.num_key_value_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({hf_config.num_attention_heads}/{hf_config.num_key_value_heads}) " + f"must be divisible by device count ({num_devices})" + ) + if bool(getattr(hf_config, "tie_word_embeddings", False)): + raise ValueError("Qwen2.5-Coder-32B requires an untied LM head") + + +def convert_hf_model_weights( + hf, + hf_config, + *, + n_layers: int, + num_devices: int, + rope_table_len: int, + head_dim: int, +) -> Qwen25Coder32BWeights: + base = hf.model + rope_cos, rope_sin = weight_utils.build_rope_cos_sin_torch( + base.rotary_emb, rope_table_len, head_dim, torch.bfloat16 + ) + layers = [] + for layer in base.layers[:n_layers]: + wqkv, wo, q_norm, k_norm, wqkv_bias = weight_utils.attention_wqkv_wo_from_hf_layer(layer.self_attn, num_devices) + if q_norm is not None or k_norm is not None: + raise ValueError("Qwen2.5-Coder-32B does not support QK norm weights") + if wqkv_bias is None: + raise ValueError("Qwen2.5-Coder-32B requires QKV projection bias") + w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(layer.mlp) + layers.append( + Qwen25Coder32BLayerWeights( + wqkv=wqkv, + wo=wo, + wqkv_bias=wqkv_bias, + w1=w1, + w2=w2, + w3=w3, + attention_norm=weight_utils.rms_weight_torch(layer.input_layernorm).to(torch.bfloat16), + ff_norm=weight_utils.rms_weight_torch(layer.post_attention_layernorm).to(torch.bfloat16), + ) + ) + if hf_config.tie_word_embeddings: + raise ValueError("Qwen2.5-Coder-32B requires an untied LM head") + return Qwen25Coder32BWeights( + embedding=weight_utils.embed_tokens_torch(base.embed_tokens), + rope_cos=rope_cos, + rope_sin=rope_sin, + layers=tuple(layers), + final_norm=weight_utils.rms_weight_torch(base.norm).to(torch.bfloat16), + lm_head=hf.lm_head.weight.detach().to(torch.bfloat16).clone(), + ) + + +def from_pretrained( + mesh_device, + *, + hf_model: str = DEFAULT_HF_MODEL, + hf_revision: str | None = DEFAULT_HF_REVISION, + instruct: bool = True, + max_batch_size: int = 32, + max_seq_len: int = 4096, + optimizations: str | Qwen25Coder32BPrecisionConfig = "accuracy", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config: Qwen25Coder32BPagedAttentionConfig | None = None, + cache_dir: Path | str | None = None, +) -> Qwen25Coder32BForCausalLM: + del dtype + num_devices = mesh_device.get_num_devices() + if num_devices != 8: + raise ValueError(f"Qwen2.5-Coder-32B supports exactly 8 devices (T3K), got {num_devices}") + ttnn.SetDefaultDevice(mesh_device) + hf_config = AutoConfig.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=_local_files_only(), + ) + _validate_checkpoint_config(hf_config, num_devices=num_devices) + hf = AutoModelForCausalLM.from_pretrained( + hf_model, + revision=hf_revision, + torch_dtype=torch.bfloat16, + local_files_only=_local_files_only(), + ) + hf.eval() + resolved_layers = hf_config.num_hidden_layers if n_layers is None else n_layers + precision = ( + optimizations + if isinstance(optimizations, Qwen25Coder32BPrecisionConfig) + else (QWEN25_CODER_32B_PERFORMANCE if optimizations == "performance" else QWEN25_CODER_32B_ACCURACY) + ) + if not isinstance(precision, Qwen25Coder32BPrecisionConfig): + raise TypeError("optimizations must be 'accuracy', 'performance', or Qwen25Coder32BPrecisionConfig") + if not (1 <= resolved_layers <= hf_config.num_hidden_layers): + raise ValueError(f"n_layers must be in [1, {hf_config.num_hidden_layers}], got {resolved_layers}") + cache_path = _cache_path(hf_model, mesh_device, cache_dir) + if paged_attention_config is None: + block_size = 32 + blocks_per_user = (max_seq_len + block_size - 1) // block_size + paged_attention_config = Qwen25Coder32BPagedAttentionConfig( + block_size=block_size, + max_num_blocks=blocks_per_user * max_batch_size, + ) + rope_table_len = max(max_seq_len * 2, 8192) + rope_table_len = (rope_table_len + 127) // 128 * 128 + head_dim = hf_config.hidden_size // hf_config.num_attention_heads + weights = convert_hf_model_weights( + hf, + hf_config, + n_layers=resolved_layers, + num_devices=num_devices, + rope_table_len=rope_table_len, + head_dim=head_dim, + ) + del hf + config = Qwen25Coder32BConfig( + hf_model_id=hf_model, + dim=hf_config.hidden_size, + n_heads=hf_config.num_attention_heads, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + hidden_dim=hf_config.intermediate_size, + vocab_size=hf_config.vocab_size, + rms_norm_eps=hf_config.rms_norm_eps, + rope_theta=getattr(hf_config, "rope_theta", 1_000_000.0), + num_hidden_layers=resolved_layers, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + rope_table_len=rope_table_len, + num_devices=num_devices, + mesh_device=mesh_device, + ) + model = build_qwen25_coder_32b_model( + mesh_device=mesh_device, + config=config, + weights=weights, + precision=precision, + cache_path=cache_path, + paged_attention_config=paged_attention_config, + ) + runtime = Qwen25Coder32BRuntimeConfig( + model_name=hf_model, + model_cache_path=cache_path, + max_prefill_chunk_size=4096, + max_context_len=int(getattr(hf_config, "max_position_embeddings", max_seq_len)), + max_seq_len=max_seq_len, + trace_prefill_supported_seq_lens=_trace_seq_lens(num_devices, 4096, max_seq_len), + n_layers=resolved_layers, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + max_batch_size=max_batch_size, + cluster_shape=list(mesh_device.shape), + kv_cache_dtype=precision.kv_cache_dtype, + batched_prefill_batched_extract=not os.environ.get("DISABLE_BATCHED_EXTRACT"), + ) + tokenizer = load_tokenizer(hf_model, hf_revision) + return Qwen25Coder32BForCausalLM( + model=model, + tokenizer=tokenizer, + runtime_config=runtime, + instruct=instruct, + ) diff --git a/code/models/common/models/qwen25_coder_32b/model.py b/code/models/common/models/qwen25_coder_32b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..67d7a78cb598a78b25659814d94227e422410b1b --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/model.py @@ -0,0 +1,1244 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Qwen2.5-Coder-32B-Instruct — native TTTv2 stack (``Embedding1D``, ``RMSNorm1D``, +``Attention1D``, ``MLP1D``, ``RotarySetup1D``, ``LMHead1D``). Targets T3K (mesh ``(1, 8)``). + +Tensor layout contracts: + - **Prefill** hidden states: ``[1, 1, S, dim]`` TILE, ``S % 128 == 0``. + - **Decode** hidden states: ``[1, 1, B, dim]`` TILE (``B`` padded to tile in modules). + +Executor contract (``EagerLLMExecutor`` / ``TracedLLMExecutor``): pre-embedded forwards, +``set_kv_cache``, ``rope_setup``, ``page_table`` through attention, ``model_args`` holds a +:class:`Qwen25Coder32BExecutorRuntimeConfig` (not v1 ``ModelArgs``). +""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, List + +import torch +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.models.qwen25_coder_32b import weight_utils +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D, prepare_rot_idxs +from models.common.modules.sampling.sampling_1d import Sampling1D +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + +# Pinned HF revision SHA for Qwen/Qwen2.5-Coder-32B-Instruct (resolved 2026-05-19). +DEFAULT_HF_REVISION = "381fc969f78efac66bc87ff7ddeadb7e73c218a7" + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + """Minimal LazyWeight; ``Attention1D`` / ``MLP1D`` / ``Embedding1D`` resolvers fill mesh + memory.""" + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +@dataclass +class Qwen25Coder32BPagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (``block_size`` / ``max_num_blocks`` only).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Qwen25Coder32BExecutorRuntimeConfig: + """Engine-facing runtime knobs. Exposed as ``model.model_args`` for shared ``EagerLLMExecutor``.""" + + n_layers: int + n_kv_heads: int + head_dim: int + max_batch_size: int + max_seq_len: int + cluster_shape: list[int] + max_prefill_chunk_size: int = 4096 + model_cache_path: Path | None = None + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + optimizations: Any = None + # Batched prefill (parity caveat #12): fuse equal-length users into batched passes to close the + # batch-32 TTFT gap. ``supports_batched_prefill`` is the per-model opt-in (the shared engine only + # batches models whose prefill_forward threads ``batch_size`` — Qwen2.5-Coder-32B does, below). + # Qwen2.5-Coder is a standard dense Qwen2.5 attention (NO QK-norm), so every prefill op is + # row-independent and the batched fold is bit-safe (same as the qwen25_7b port). + # ``max_prefill_batch_size`` is the largest supported padded wave; 32 folds batch-32 prefill in ONE + # 32-user pass (TTTv1 structural parity) so the eager norm+lm_head tail + full-vocab readback run + # once instead of 4×. At the natural 128 bucket the fold is 32*128=4096=2*2048, an exact multiple of + # MAX_QKV_MM_SEQ_LEN (reshape-safe), and 4096 % mlp_prefill_len_cutoff(1024) == 0 for the FF reshape; + # the DRAM guard (padded_batch*seq < 128K) passes with 4096. This model already runs the 1024 engine + # default for the FF cutoff (_resolve_coder_32b_wh_tuning), so the fold is the only prefill lever + # applied here. ``disable_batched_prefill`` is the escape hatch back to the sequential loop; + # ``max_prefill_chunk_size`` (above) drives the #45234 chunked-prompt decline. + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + # When True (default), batched prefill runs norm+lm_head ONCE per group over the gathered last-token + # rows (TTTv1 parity); False falls back to the bit-identical per-slot path (one lm_head per user). + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int) -> bool: + # Mirror TTTv1's ModelArgs.get_trace_prefill_supported_seq_lens: on the only supported SKU + # (T3K, 8 devices) the device default is [128, 1024]. Prefill compiles, captures and replays + # for these seq lens on hardware. + # + # Gate tracing to COLD prefill (num_cached_tokens == 0). The shared traced-prefill path does + # not thread the cached start position / chunk_start into the trace body, so a prefix-cached + # request (num_cached_tokens > 0) that pads to a trace-eligible length would replay with + # start_pos=0 -> wrong RoPE indices and KV written at cache offset 0 (silently wrong output). + # The eager prefix-cache path offsets correctly, so fall through to it; full traced + # prefix-cache prefill is the upstream work item (issue #32056). + if num_cached_tokens != 0: + return False + num_devices = int(self.cluster_shape[0]) * int(self.cluster_shape[1]) + allowed = {1: (128,), 2: (128, 1024), 8: (128, 1024)}.get(num_devices, (128,)) + return ( + prefill_seq_len in allowed + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +@dataclass +class Qwen25Coder32BConfig: + """Resolved hyper-parameters for a loaded HF Qwen2.5-Coder-32B checkpoint.""" + + hf_model_id: str + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + rope_theta: float + num_hidden_layers: int + max_batch_size: int + max_seq_len: int + rope_table_len: int + num_devices: int = 8 + mesh_device: ttnn.MeshDevice | None = None + n_layers: int | None = None + block_configs: list[Any] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.n_layers is None: + self.n_layers = self.num_hidden_layers + + +@dataclass(frozen=True) +class Qwen25Coder32BLayerWeights: + wqkv: torch.Tensor + wo: torch.Tensor + wqkv_bias: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Qwen25Coder32BWeights: + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Qwen25Coder32BLayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +_QWEN_ATTN_HIFI4_FP32_KERNEL = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=True, +) +"""HiFi4 + fp32 dest acc for Qwen2.5 attention matmuls (LI_QKV, LI_O, SDPA). + +TTTv1 ``DecodersPrecision.accuracy("Qwen2.5-Coder-32B-Instruct")`` resolves to the +non-Llama / non-Mistral branch in ``model_config.py:160-177`` which forces all attention +ops to ``HIFI4``. The TTTv2 ``Attention1D`` default is ``HIFI2`` with fp16 accumulation; +without this override, attention QKV / WO / SDPA produce a broad per-layer drift vs HF +(same regression debugged on the Qwen2.5-7B port). Used in ``QWEN25_CODER_32B_ACCURACY``. +""" + + +_LOFI_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, +) +"""LoFi + packer L1 acc for the MLP FF1/FF3 matmuls in performance mode. + +Mirrors TTTv1 ``DecodersPrecision.performance("Qwen2.5-Coder-32B-Instruct")``: the +non-Qwen2.5-7B branch at ``model_config.py:208-218`` sets ``FF1_FF3 → BFP4`` and +``LI_FF1_FF3 → LOFI``. This single delta is the bulk of the perf-mode throughput uplift. +""" + + +_TTTV1_HIFI2_COMPUTE_KERNEL_CFG = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=True, + fp32_dest_acc_en=True, + packer_l1_acc=True, +) +"""TTTv1's ``compute_kernel_config_hifi2`` for HIFI2 attention ops in performance mode. + +TTTv2's ``Attention1D`` default for HIFI2 ops is ``compute_kernel_hifi2_fp16`` +(``math_approx=False``, ``fp32_dest_acc=False``) — that matches TTTv1's HIFI2_FP16 +setting, **not** TTTv1's plain HIFI2 (``math_approx=True``, ``fp32_dest_acc=True``). +For Qwen2.5-Coder-32B perf mode, the TTTv1 OpFidelity defaults +(``model_config.py:292-297``) resolve attention decode kernels to plain HIFI2, so the +perf recipe explicitly pins this kernel instead of falling back to TTTv2's fp16 variant. +""" + + +@dataclass(frozen=True) +class Qwen25Coder32BPrecisionConfig: + """Per-layer precision + math-fidelity recipe for Qwen2.5-Coder-32B-Instruct on T3K. + + Mirrors the fields TTTv1's ``DecodersPrecision`` distinguishes for Qwen2.5-Coder-32B. + The base model name resolves to ``Qwen2.5-Coder-32B`` (via ``common.get_base_model_name``), + which falls into the standard Qwen2.5-family branch — **not** the Qwen2.5-7B / Qwen2.5-VL-7B + special case in ``model_config.py:187``. As a result: + + * **Accuracy** (``model_config.py:160-177``): BF16 ``WQKV`` / ``WO`` + ``HIFI4`` on every + ``LI_QKV`` / ``SDPA`` / ``LI_O``. FF and LM head stay at engine defaults (BFP8 FF + + ``HIFI2_FP16``). **KV cache is BFP8, not BF16** (diverges from TTTv1) — BF16 KV deadlocks + the traced on-device-topk decode replay at full depth; BFP8 clears it and is loss-free + here (see ``kv_cache_dtype`` below). + * **Performance** (``model_config.py:208-218``, non-7B branch): only ``FF1_FF3 → BFP4`` + and ``LI_FF1_FF3 → LOFI``. Everything else reverts to TTTv1 defaults + (BFP8 attention + ``HIFI2`` attention kernels + BFP8 KV cache). + + Two module-level recipes are exposed: :data:`QWEN25_CODER_32B_ACCURACY` (default) and + :data:`QWEN25_CODER_32B_PERFORMANCE`. Pass one to :meth:`Qwen25Coder32B.from_pretrained` + via ``precision=``; use ``dataclasses.replace(QWEN25_CODER_32B_ACCURACY, ...)`` to + customize a single field. Defaults below mirror the accuracy recipe so ``Qwen25Coder32BPrecisionConfig()`` + is equivalent to :data:`QWEN25_CODER_32B_ACCURACY`. + + The ``mlp_w2_dtype`` / ``mlp_ff2_compute_kernel_cfg`` fields are absent because TTTv1 leaves + them at engine defaults (``BFP8`` / ``HIFI2_FP16``) in *both* recipes for this model, and + those defaults coincide with the TTTv2 ``MLP1D`` defaults. + """ + + # Attention weight dtypes: accuracy uses BF16 (default), performance overrides BFP8. + wqkv_dtype: ttnn.DataType = ttnn.bfloat16 + wo_dtype: ttnn.DataType = ttnn.bfloat16 + # KV cache: BFP8 in BOTH recipes. TTTv1's accuracy config nominally uses BF16 KV; BFP8 is kept + # here because it is loss-free for this model (teacher-forcing top1 98.6% / top5 100.0%, >= the + # 95/99 targets -- the accuracy recipe's precision comes from BF16 attention weights + HIFI4 + + # fp32 dest acc, not the KV dtype) AND it avoids BF16's doubled per-layer KV read traffic in the + # memory-bound decode step. NOTE: an earlier bringup switched BF16 -> BFP8 KV to work around a + # batch-1·accuracy·on_device_topk decode-trace-replay hang, attributing it to BF16 KV's larger + # footprint tripping a "cumulative traffic ceiling". That was a layout band-aid: the true root + # cause was on-device sampling buffers being allocated while a trace was already active and then + # clobbered on replay, fixed generically in executor.py + # (TracedLLMExecutor._prealloc_sampling_buffers, called before the prefill trace is captured). + # BFP8 KV is now retained on its own merits (loss-free + less decode traffic), not for the hang. + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + # MLP FF1/FF3 weight dtype. Accuracy keeps BFP8 default; performance overrides BFP4. + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat8_b + + # MLP FF1/FF3 compute kernel. ``None`` → MLP1D default (HIFI2_FP16); performance overrides LOFI. + mlp_ff1_3_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + # Attention compute kernels. Accuracy sets HIFI4 + fp32 dest acc on every stage; performance + # leaves them at the Attention1D default (HIFI2 fp16 dest acc). + attn_li_qkv_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = _QWEN_ATTN_HIFI4_FP32_KERNEL + attn_sdpa_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = _QWEN_ATTN_HIFI4_FP32_KERNEL + attn_li_o_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = _QWEN_ATTN_HIFI4_FP32_KERNEL + + # Not in TTTv1 DecodersPrecision; TTTv2 accuracy tightens to bf16 to lock in top-1 (matches the + # Mistral-7B port). Performance keeps BFP8. + lm_head_dtype: ttnn.DataType = ttnn.bfloat16 + + +# TTTv1 ``DecodersPrecision.accuracy("Qwen2.5-Coder-32B-Instruct")`` (``model_config.py:160-177``): +# BF16 attention weights + HIFI4 + fp32_dest_acc on every attention stage. FF and LM head sit at +# TTTv2 defaults (BFP8 + HIFI2_FP16); LM head tightens to bf16 (TTTv2 addition). KV cache is BFP8 +# (NOT TTTv1's BF16) to avoid the full-depth decode-trace-replay deadlock — see kv_cache_dtype above. +QWEN25_CODER_32B_ACCURACY = Qwen25Coder32BPrecisionConfig() + +# TTTv1 ``DecodersPrecision.performance("Qwen2.5-Coder-32B-Instruct")`` (``model_config.py:208-218``, +# non-7B branch): FF1_FF3 → BFP4 and LI_FF1_FF3 → LOFI; everything else reverts to engine defaults +# (BFP8 attention, HIFI2 attention kernels, BFP8 KV cache, BFP8 LM head). +QWEN25_CODER_32B_PERFORMANCE = Qwen25Coder32BPrecisionConfig( + wqkv_dtype=ttnn.bfloat8_b, + wo_dtype=ttnn.bfloat8_b, + kv_cache_dtype=ttnn.bfloat8_b, + mlp_w1_w3_dtype=ttnn.bfloat4_b, + mlp_ff1_3_compute_kernel_cfg=_LOFI_COMPUTE_KERNEL_CFG, + # Pin TTTv1's plain HIFI2 (math_approx=True, fp32_dest_acc=True). Without this, the + # Attention1D default for the three attention ops is HIFI2_FP16, which mismatches + # TTTv1's compute_kernel_config_hifi2 (model_config.py:728-733). + attn_li_qkv_kernel_cfg=_TTTV1_HIFI2_COMPUTE_KERNEL_CFG, + attn_sdpa_kernel_cfg=_TTTV1_HIFI2_COMPUTE_KERNEL_CFG, + attn_li_o_kernel_cfg=_TTTV1_HIFI2_COMPUTE_KERNEL_CFG, + lm_head_dtype=ttnn.bfloat8_b, +) + + +def _slice_last_token_tile(x: ttnn.Tensor, last_token_idx: int) -> ttnn.Tensor: + """Slice the 32-row tile containing ``last_token_idx`` from ``[1, 1, S, W]``. + + Width-sharded LM matmul M tile rows must equal ``LMHead1D`` program-config tile rows. + """ + floor = (last_token_idx // 32) * 32 + return ttnn.slice(x, (0, 0, floor, 0), (1, 1, floor + 32, x.shape[-1])) + + +def _post_attn_norm_decode_configs( + mlp: MLP1D, + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding so its output matches MLP1D's W1/W3 input. + + MLP1D decode uses ``_dram_shard_core_grid_k_n(dim, padded_hidden / num_devices)`` for W1/W3 + inputs, but the default RMSNorm program config is derived from ``_compute_norm_core_grid(dim)`` + alone. Mismatched DRAM-width-shard between RMSNorm output and MLP1D W1/W3 input silently + corrupts decode activations (observed on the 7B port — same shape pattern on T3K Coder-32B). + """ + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + return program_config, mlp.config.decode_input_memcfg + + +# Cast coder's two per-layer prefill norm-reconstruction all-gathers (pre-attn + post-attn) from bf16 to +# bfloat8_b, mirroring the sibling qwen25_72b H003 (commit 783d2945e2b, model.py:262). Every coder prefill +# decoder layer runs two full-activation all-gathers at bf16 to reconstruct the dim-fractured RMSNorm output +# to full ``dim`` before Attention1D / MLP1D (the QKV / W1/W3 matmuls expect width ``dim``). Each moves a +# 5120-wide activation across 8 devices, x2/layer x64 layers = 128 bf16 all-gathers on the batched b32-ci +# prefill critical path. Every matmul these gathers feed already quantizes activations to bf8_b (QKV/WO) or +# bf4_b (W1/W3), so the bf16 gather precision is largely unconsumed; casting the gather input to bf8_b halves +# each collective's cross-device payload (AllGatherAsync is payload-bound at dim=5120). PREFILL-ONLY: the +# decode + tail-norm call sites of _all_gather_rmsnorm_tensor pass no dtype -> byte-identical decode/tail. +# Unlike 72b there is no third fused pre-WO gather to cast: coder's dim=5120 fails the fused all-gather auto- +# gate ((5120//32//8) % 8 = 4 != 0), so only these two norm gathers apply (the fused hook is a no-op here). +# NOT bit-exact vs bf16 (mantissa drop) -> gated on token-accuracy + eval-32, NOT the 32-user byte-compare +# (INVALID on coder's tie-heavy on_device_topk). coder's perf recipe is BFP4-FF/LoFi (less precision-tolerant +# than 72b's >70B bf4/bf8 recipe), so the accuracy profile is re-gated separately. A/B escape hatch: set +# DISABLE_PREFILL_AG_BF8=1 to keep the gathers at bf16 (byte-identical to the pre-change HEAD). None disables. +_PREFILL_AG_CCL_DTYPE: ttnn.DataType | None = ( + None if os.environ.get("DISABLE_PREFILL_AG_BF8") == "1" else ttnn.bfloat8_b +) + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, + x: ttnn.Tensor, + *, + memory_config: ttnn.MemoryConfig | None = None, + dtype: ttnn.DataType | None = None, +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + if memory_config is not None: + return ttnn.to_memory_config(x, memory_config) + return x + + if memory_config is None: + memory_config = x.memory_config() + + # Prefill-only opt-in (``dtype`` is set only at the two per-layer prefill call sites): cast the gather + # input (bf16) to a smaller CCL dtype (bfloat8_b) to halve this collective's cross-device payload. Decode + # and tail-norm call sites leave ``dtype=None`` -> byte-identical. An explicit typecast is required + # (to_memory_config does not cast an already-DRAM tensor). + if dtype is not None and x.dtype != dtype: + x_cast = ttnn.typecast(x, dtype) + ttnn.deallocate(x) + x = x_cast + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + + +@dataclass +class _Qwen25Coder32BWHTuning: + """Wormhole-specific L1 / cutoff tuning resolved at build time. + + Precision-vs-fidelity knobs live on :class:`Qwen25Coder32BPrecisionConfig`. This + dataclass only carries non-mode-dependent L1 footprint controls: MLP prefill cutoff + and the optional W1→DRAM spill on decode. + """ + + mlp_prefill_len_cutoff: int | None = None + mlp_decode_spill_w1_to_dram: bool = False + # Use ttnn.experimental.minimal_matmul for the QKV + W2 prefill matmuls above seq_len > 128 (TTTv1 + # parity, PLAN_01). A/B escape hatch: set DISABLE_MINIMAL_MATMUL=1 to force ttnn.linear. On a 32B the + # batch-32-ci prefill is matmul-compute-bound (~80% of FLOPs = the 3 MLP matmuls), so minimal_matmul + # is a real prefill-TTFT win — it narrows the batch-32-ci TTFT gap vs TTTv1 (which uses minimal_matmul + # for the same matmuls). The shared plumbing (attention_1d.use_minimal_qkv_matmul / mlp_1d + # use_minimal_w2_matmul, both gated seq_len>128) is already in the base; this flag just engages it. + # Coder is a Qwen2.5 arch: the QKV bias is added AFTER the matmul, so it is unchanged by the minimal + # path (same as the mistral_7b / deepseek Qwen2 ports). Independent of the FF-hidden DRAM-shard pad + # (a decode fix); minimal_matmul is prefill-only, so decode throughput is unchanged. + prefill_minimal_matmul: bool = True + # Cast coder's per-layer prefill attention output reduce-scatter from bf16 to bfloat8_b, matching + # TTTv1's ccl_dtype (model_config.py:1070). coder's dim=5120 fails the fused all-gather auto-gate + # ((5120//32//8) % 8 = 4 != 0), so attention runs a SEPARATE per-layer reduce_scatter on the bf16 WO + # output (attention_1d._all_reduce_output_prefill), x64 layers; a bf16 reduce moves 2x the cross-device + # bytes of a bf8_b one and reduce_scatter on this shape is dtype-bound (PLAN_03: 2210us -> 1534us/layer + # at dim=5120). TTTv1 already reduces this at bf8_b, so the result matches the shipped reference numerics + # (not a degradation below it). PREFILL-ONLY: the decode reduce (_all_reduce_output_decode) never reads + # prefill_reduce_ccl_dtype -> decode byte-identical. A/B escape hatch: set DISABLE_PREFILL_REDUCE_BF8=1 + # to keep the reduce at bf16 (byte-identical to the pre-change HEAD). Not bit-exact vs bf16 (mantissa + # drop) -> gated on token-accuracy + eval-32, not the 32-user byte-compare (INVALID on coder's tie-heavy + # on_device_topk). + prefill_reduce_ccl_bf8: bool = True + + +def _resolve_qwen_coder_wh_tuning(*, num_dev: int, max_batch_size: int) -> _Qwen25Coder32BWHTuning: + """Pick WH L1 tuning knobs for Qwen2.5-Coder-32B-Instruct on T3K. + + Use TTTv1's Wormhole default (``prefill_len_cutoff=1024`` at ``model_config.py:516``). + Coder-32B is not in the "reduce to 512" override list (``model_config.py:583-589``), so + TTTv1 keeps 1024 for this model on T3K. ``mlp_decode_spill_w1_to_dram`` is off on T3K + because per-device FF shards (5120×3456 per chip) are smaller than 7B-on-N300. + """ + t = _Qwen25Coder32BWHTuning( + mlp_prefill_len_cutoff=1024, + mlp_decode_spill_w1_to_dram=False, + ) + t.prefill_minimal_matmul = not os.environ.get("DISABLE_MINIMAL_MATMUL") + t.prefill_reduce_ccl_bf8 = not os.environ.get("DISABLE_PREFILL_REDUCE_BF8") + logger.info( + f"L1 tuning for Qwen2.5-Coder-32B on {num_dev} device(s): " + f"prefill_len_cutoff={t.mlp_prefill_len_cutoff}, " + f"decode_spill_w1_to_dram={t.mlp_decode_spill_w1_to_dram}, " + f"prefill_minimal_matmul={t.prefill_minimal_matmul}, " + f"prefill_reduce_ccl_bf8={t.prefill_reduce_ccl_bf8}" + ) + return t + + +def _build_decoder_layer( + *, + idx: int, + weights: Qwen25Coder32BLayerWeights, + qcfg: Qwen25Coder32BConfig, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Qwen25Coder32BPrecisionConfig, + executor_mode: bool, + paged_cfg: Qwen25Coder32BPagedAttentionConfig | None, + cache_path: Path | None, + wh: _Qwen25Coder32BWHTuning, +) -> Qwen25Coder32BDecoderLayer: + """Construct one decoder layer (attention + MLP + the two RMSNorms) from an HF layer.""" + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, dtype=precision.wqkv_dtype, cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None + ) + lazy_wo = _lazy( + weights.wo, dtype=precision.wo_dtype, cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None + ) + + def _qk_norm_cfg(weight: torch.Tensor | None, name: str) -> RMSNorm1DConfig | None: + if weight is None: + return None + lw = _lazy( + weight.unsqueeze(0).unsqueeze(0).unsqueeze(0).to(torch.bfloat16), + dtype=ttnn.bfloat16, + cache=(cache_path / "attn", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=qcfg.rms_norm_eps, + decode_in_sharded=False, + decode_out_sharded=False, + prefill_distributed=False, + tt_ccl=tt_ccl, + ) + + bias_lw = ( + LazyWeight( + source=weights.wqkv_bias.to(torch.bfloat16), + dtype=ttnn.bfloat16, + cache_dir_weight_name=(cache_path / "attn", f"{prefix}_bias") if cache_path else None, + ) + if weights.wqkv_bias is not None + else None + ) + + attn = Attention1D.from_config( + Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=qcfg.n_heads, + n_kv_heads=qcfg.n_kv_heads, + head_dim=qcfg.head_dim, + max_batch_size=qcfg.max_batch_size, + max_seq_len=qcfg.max_seq_len, + q_norm_config=_qk_norm_cfg(None, "qn"), + k_norm_config=_qk_norm_cfg(None, "kn"), + wqkv_bias=bias_lw, + use_vllm_paged_kv_cache=executor_mode, + paged_attention_config=paged_cfg, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + li_qkv_prefill_compute_kernel_cfg=precision.attn_li_qkv_kernel_cfg, + li_qkv_decode_compute_kernel_cfg=precision.attn_li_qkv_kernel_cfg, + sdpa_decode_compute_kernel_cfg=precision.attn_sdpa_kernel_cfg, + li_o_prefill_compute_kernel_cfg=precision.attn_li_o_kernel_cfg, + li_o_decode_compute_kernel_cfg=precision.attn_li_o_kernel_cfg, + prefill_qkv_minimal_matmul=wh.prefill_minimal_matmul, + # Route the folded-batch prefill WO projection through minimal_matmul (completes PLAN_01's + # QKV+FF2 minimal plumbing on coder). WO is the least-efficient prefill matmul on ttnn.linear; + # the minimal op (already used for QKV/FF2) recovers most of that gap and makes TTTv2 beat + # TTTv1's ttnn.linear WO (a shared op on both stacks). coder's WO is the non-fused shape + # (dim=5120 fails the fused all-gather auto-gate), so this runs every layer. Gated by the + # same DISABLE_MINIMAL_MATMUL escape hatch (OFF => byte-identical to the pre-change HEAD). + prefill_wo_minimal_matmul=wh.prefill_minimal_matmul, + # Cast the per-layer prefill attention output reduce-scatter (bf16 WO output) to bfloat8_b, + # matching TTTv1's ccl_dtype. coder is on the NON-fused path (dim=5120 fails the fused + # all-gather auto-gate), so this reduce runs every layer; a bf8_b reduce halves the + # cross-device payload. Prefill-only (the decode reduce ignores this field). Gated by the + # DISABLE_PREFILL_REDUCE_BF8 escape hatch (None => byte-identical to the pre-change HEAD). + prefill_reduce_ccl_dtype=(ttnn.bfloat8_b if wh.prefill_reduce_ccl_bf8 else None), + ) + ) + + w1, w2, w3 = weights.w1, weights.w2, weights.w3 + # Pad the FF hidden dim to a grid-friendly per-device size so the DRAM-sharded decode FF + # matmuls (W1/W3/W2) use a full multi-core grid instead of 4. The decode grid must divide both + # the K-tile and N-tile counts (in0 is K-width-sharded on dim, weights are N-sharded on hidden); + # with the raw per-device hidden 27648/8 -> 3456 = 108 tiles (2^2*27), gcd(dim_tiles=160, 108)=4 + # -> only 4 DRAM readers stream the FF weights, which dominate memory-bound decode. Padding per + # device to 128 tiles (4096, total 32768) gives gcd(160, 128)=32 cores. The extra columns are + # zeros (silu(0)=0, mul->0, and W2 contracts the padded rows to 0), so decode output is unchanged. + _ff_align = TILE_SIZE * TILE_SIZE * num_dev # 32*32*8 = 8192 on T3K + _ff_pad = math.ceil(w1.shape[-1] / _ff_align) * _ff_align + if _ff_pad != w1.shape[-1]: + w1 = torch.nn.functional.pad(w1, (0, _ff_pad - w1.shape[-1])) + w3 = torch.nn.functional.pad(w3, (0, _ff_pad - w3.shape[-1])) + w2 = torch.nn.functional.pad(w2, (0, 0, 0, _ff_pad - w2.shape[-2])) + mlp = MLP1D.from_config( + MLP1DConfig( + w1=_lazy( + w1, dtype=precision.mlp_w1_w3_dtype, cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None + ), + w2=_lazy(w2, dtype=ttnn.bfloat8_b, cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None), + w3=_lazy( + w3, dtype=precision.mlp_w1_w3_dtype, cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=qcfg.max_batch_size, + prefill_len_cutoff=wh.mlp_prefill_len_cutoff, + w1_w3_dtype=precision.mlp_w1_w3_dtype, + w2_dtype=ttnn.bfloat8_b, + ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=precision.mlp_ff1_3_compute_kernel_cfg, + decode_spill_w1_to_dram_before_w3=wh.mlp_decode_spill_w1_to_dram, + prefill_w2_minimal_matmul=wh.prefill_minimal_matmul, + ) + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + mlp, + dim=qcfg.dim, + # Use the padded FF hidden so the post-attn RMSNorm decode output is width-sharded on the + # SAME (32-core) grid as MLP1D's W1/W3 decode input; a mismatch silently corrupts decode. + hidden_dim=_ff_pad, + num_devices=num_dev, + max_batch_size=qcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1D: + lw = _lazy( + weight.to(torch.bfloat16), + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1D.from_config( + RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=qcfg.rms_norm_eps, + max_batch_size=qcfg.max_batch_size, + tt_ccl=tt_ccl, + **extra, + ) + ) + + return Qwen25Coder32BDecoderLayer( + input_layernorm=_build_norm(weights.attention_norm, "pre_attn"), + self_attn=attn, + post_attention_layernorm=_build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ), + mlp=mlp, + ) + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + hf_lm_head: torch.nn.Module, + qcfg: Qwen25Coder32BConfig, + lm_head_dtype: ttnn.DataType, + cache_path: Path | None, +) -> LMHead1D: + """Build the vocab-sharded LM head with DRAM-matmul program configs. + + LM head DRAM matmul is sized for decode batch tiles (``max_batch_size``). Prefill logits + use a single 32-row tile via ``post_process_prefill_output`` / :func:`_slice_last_token_tile`. + """ + lm_w = hf_lm_head.weight.detach().to(torch.bfloat16).clone() + lm_splits, lm_split_sizes, lm_weights_memcfgs = weight_utils.build_lm_head_lazy_weights( + mesh_device, + lm_w, + dim=qcfg.dim, + vocab_size=qcfg.vocab_size, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(qcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(qcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, qcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(qcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1D.from_config( + LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=qcfg.dim, + max_batch_size=qcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + ) + ) + + +def build_qwen25_coder_32b_model( + *, + mesh_device: ttnn.MeshDevice, + config: Qwen25Coder32BConfig, + weights: Qwen25Coder32BWeights, + precision: Qwen25Coder32BPrecisionConfig, + cache_path: Path | None, + paged_attention_config: Qwen25Coder32BPagedAttentionConfig | None, +) -> Qwen25Coder32B: + """Build the TT tensor graph from provider-neutral Qwen2.5-Coder-32B dimensions and tensors.""" + + ttnn.SetDefaultDevice(mesh_device) + num_devices = mesh_device.get_num_devices() + if num_devices != 8: + raise ValueError( + f"Qwen2.5-Coder-32B-Instruct port targets T3K (mesh (1, 8) = 8 devices) only. " + f"Got mesh_device with {num_devices} device(s). Open a T3K mesh with MESH_DEVICE=T3K." + ) + if config.n_heads % num_devices != 0 or config.n_kv_heads % num_devices != 0: + raise ValueError( + f"Checkpoint heads ({config.n_heads}/{config.n_kv_heads}) must be divisible by " + f"device count ({num_devices})" + ) + if len(weights.layers) != config.num_hidden_layers: + raise ValueError(f"Expected {config.num_hidden_layers} decoder layer weights, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) + topology = default_topology(mesh_device) + emb = Embedding1D.from_config( + Embedding1DConfig( + weights=_lazy( + weights.embedding, + dtype=ttnn.bfloat16, + cache=(cache_path / "embedding", "tok_embeddings") if cache_path else None, + ), + mesh_device=mesh_device, + embed_scale=1.0, + ) + ) + rope_setup = RotarySetup1D.from_config( + Rope1DConfig( + cos_matrix=_lazy( + weights.rope_cos, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "cos") if cache_path else None + ), + sin_matrix=_lazy( + weights.rope_sin, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "sin") if cache_path else None + ), + max_batch_size=config.max_batch_size, + head_dim=config.head_dim, + device=mesh_device, + use_qk_fused=False, + ) + ) + + wh = _resolve_qwen_coder_wh_tuning(num_dev=num_devices, max_batch_size=config.max_batch_size) + layers = [ + _build_decoder_layer( + idx=idx, + weights=weights.layers[idx], + qcfg=config, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + executor_mode=paged_attention_config is not None, + paged_cfg=paged_attention_config, + cache_path=cache_path, + wh=wh, + ) + for idx in range(config.num_hidden_layers) + ] + final_norm = RMSNorm1D.from_config( + RMSNorm1DConfig( + weight=_lazy( + weights.final_norm, dtype=ttnn.bfloat16, cache=(cache_path / "norm", "final") if cache_path else None + ), + mesh_device=mesh_device, + eps=config.rms_norm_eps, + max_batch_size=config.max_batch_size, + tt_ccl=tt_ccl, + ) + ) + lm_head = _build_lm_head( + mesh_device=mesh_device, + hf_lm_head=type("_LMHeadWeight", (), {"weight": weights.lm_head})(), + qcfg=config, + lm_head_dtype=precision.lm_head_dtype, + cache_path=cache_path, + ) + return Qwen25Coder32B(config, emb, rope_setup, layers, final_norm, lm_head, mesh_device) + + +class Qwen25Coder32BDecoderLayer(LightweightModule): + def __init__( + self, + *, + input_layernorm: RMSNorm1D, + self_attn: Attention1D, + post_attention_layernorm: RMSNorm1D, + mlp: MLP1D, + ): + super().__init__() + self.input_layernorm = input_layernorm + self.self_attn = self_attn + self.post_attention_layernorm = post_attention_layernorm + self.mlp = mlp + self.attention_norm = input_layernorm + self.attention = self_attn + self.ff_norm = post_attention_layernorm + self.feed_forward = mlp + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + *, + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + # Match Llama ``TransformerBlock1D``: fractured embed / norm activations must be + # all-gathered to full ``dim`` before Attention1D / MLP1D (QKV matmul expects width ``dim``). + r = self.input_layernorm.prefill_forward(x) + r = _all_gather_rmsnorm_tensor(self.input_layernorm, r, dtype=_PREFILL_AG_CCL_DTYPE) + r = self.self_attn.prefill_forward( + r, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + r2 = self.post_attention_layernorm.prefill_forward(h) + r2 = _all_gather_rmsnorm_tensor(self.post_attention_layernorm, r2, dtype=_PREFILL_AG_CCL_DTYPE) + r2 = self.mlp.prefill_forward(r2) + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def decode_forward( + self, + x: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + xa = _all_gather_rmsnorm_tensor( + self.input_layernorm, x, memory_config=self.input_layernorm.config.decode_memory_config + ) + r = self.input_layernorm.forward(xa, "decode") + r = self.self_attn.forward(r, current_pos, rot_mats, mode="decode", page_table=page_table) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + hf = _all_gather_rmsnorm_tensor( + self.post_attention_layernorm, h, memory_config=self.post_attention_layernorm.config.decode_memory_config + ) + r2 = self.post_attention_layernorm.forward(hf, "decode") + r2 = self.mlp.forward(r2, "decode") + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + +class Qwen25Coder32B(LightweightModule): + """ + Full decoder for Qwen2.5-Coder-32B-Instruct (TTTv2 modules only) on T3K. + + Prefill/decode on **embedded** activations match ``EagerLLMExecutor``. Token embedding + is ``embed_prefill`` / ``embed_decode``. Bind KV with ``set_kv_cache`` before first forward. + """ + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__( + self, + cfg: Qwen25Coder32BConfig, + embed: Embedding1D, + rope_setup: RotarySetup1D, + layers: List[Qwen25Coder32BDecoderLayer], + norm: RMSNorm1D, + lm_head: LMHead1D, + mesh_device: ttnn.MeshDevice, + ): + super().__init__() + self.cfg = cfg + self.config = cfg + self.config.mesh_device = mesh_device + self.embed = embed + self.embedding = self.embed + self.rope_setup = rope_setup + self.layers = layers + self.norm = norm + self.lm_head = lm_head + self.mesh_device = mesh_device + self.model_args: Qwen25Coder32BExecutorRuntimeConfig | None = None + + self.vocab_size = cfg.vocab_size + self.n_layers = cfg.num_hidden_layers + self.num_devices = mesh_device.get_num_devices() + self.tt_ccl = get_tt_ccl(mesh_device) if self.num_devices > 1 else None + self.config.block_configs = [ + type("_BlockConfig", (), {"attention_config": layer.self_attn.config})() for layer in self.layers + ] + + # On-device sampling. The model owns its sampler; callers only pick behavior per request + # via ``sampling_params`` (the executor routes greedy/argmax vs the top-k op path). Buffers + # are lazy -- nothing materializes until the first on-device sampled decode -- so this is + # harmless when ``sampling_params is None`` (the host-argmax path, which stays the demo + # default). Qwen2.5-Coder-32B is a T3K-only port (8 devices), where Sampling1D's all-gather + # uses a barrier-free Ring and is trace-capture-safe. + # + # Coder's LM head exposes 128 padded rows, beyond the top-k sampling + # kernel's 32-user limit. Route greedy requests through force-argmax. + self.supports_on_device_sampling = self.num_devices >= 1 + self.sampling = ( + Sampling1D( + vocab_size=self.vocab_size, + mesh_device=mesh_device, + tt_ccl=self.tt_ccl, + max_batch_size=_nearest_32(cfg.max_batch_size), + allow_force_argmax=True, + pad_to_power_of_2=True, + ) + if self.supports_on_device_sampling + else None + ) + + @property + def n_kv_heads(self) -> int: + return self.cfg.n_kv_heads + + def iter_executor_named_modules(self): + if not hasattr(self, "layers"): + return + for index, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{index}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.self_attn.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].self_attn, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in getattr(self.config, "block_configs", ())) + for config in tuple({id(item): item for item in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, + block_size=block_size, + max_num_blocks=max_num_blocks, + ) + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + @classmethod + def from_pretrained( + cls, + mesh_device: ttnn.MeshDevice, + hf_model_id: str = "Qwen/Qwen2.5-Coder-32B-Instruct", + *, + revision: str | None = DEFAULT_HF_REVISION, + max_batch_size: int = 32, + max_seq_len: int = 4096, + num_layers: int | None = None, + cache_dir: Path | str | None = None, + precision: Qwen25Coder32BPrecisionConfig = QWEN25_CODER_32B_ACCURACY, + block_size: int = 32, + executor_mode: bool = False, + ) -> Qwen25Coder32B: + """Compatibility constructor; provider loading lives in ``hf_adaptor``.""" + + from models.common.models.qwen25_coder_32b.hf_adaptor import from_pretrained + + optimizations: str | Qwen25Coder32BPrecisionConfig + if precision == QWEN25_CODER_32B_PERFORMANCE: + optimizations = "performance" + elif precision == QWEN25_CODER_32B_ACCURACY: + optimizations = "accuracy" + else: + optimizations = precision + product = from_pretrained( + mesh_device, + hf_model=hf_model_id, + hf_revision=revision, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + optimizations=optimizations, + n_layers=num_layers, + paged_attention_config=( + Qwen25Coder32BPagedAttentionConfig( + block_size=block_size, + max_num_blocks=((max_seq_len + block_size - 1) // block_size) * max_batch_size, + ) + if executor_mode + else None + ), + cache_dir=cache_dir, + ) + return product.model + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.self_attn.config.kv_cache = None + if hasattr(layer.self_attn, "kv_cache"): + layer.self_attn.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for index, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {index} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {index} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.self_attn.config.kv_cache = pair + if hasattr(layer.self_attn, "kv_cache"): + layer.self_attn.kv_cache = pair + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). The batched path always + # returns the full hidden state (get_last_token == -1); the executor does per-slot last-token + # extraction + norm/lm_head so those stages stay bit-identical to the single-user path. + x = x_embed + for layer in self.layers: + x = layer.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + + if get_last_token == -1: + return x + + # Slice + deallocate the full-sequence buffer before norm/LM head reduces peak L1. + old = x + if last_token_slice is None: + x_tile = _slice_last_token_tile(old, get_last_token) + else: + x_tile = ttnn.slice( + old, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(old.shape[2]) // 32, + ) + ttnn.deallocate(old) + if last_token_index is not None: + if x_tile.dtype != ttnn.bfloat16: + old_tile = x_tile + x_tile = ttnn.typecast(x_tile, ttnn.bfloat16) + ttnn.deallocate(old_tile) + old_tile = x_tile + x_tile = ttnn.embedding(last_token_index, x_tile, layout=ttnn.TILE_LAYOUT) + x_tile = ttnn.unsqueeze_to_4D(x_tile) + ttnn.deallocate(old_tile) + return self._last_tile_logits(x_tile) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + x = _slice_last_token_tile(hidden_states, last_token_idx) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + return self._last_tile_logits(x) + + def _last_tile_logits(self, x_tile: ttnn.Tensor) -> ttnn.Tensor: + """Final-norm + all-gather + LM-head on a 32-row tile. ``x_tile`` shape ``[1, 1, 32, dim]``.""" + x = self.norm.prefill_forward(x_tile) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + del last_token_slice, last_token_index + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(self.mesh_device), + ) + x = ttnn.matmul(selector, hidden_states, memory_config=ttnn.DRAM_MEMORY_CONFIG) + ttnn.deallocate(selector) + return self._last_tile_logits(x) + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for layer in self.layers: + x = layer.decode_forward(x, current_pos, rot_mats, page_table=page_table) + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + return self.lm_head.forward(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) + + def prefill_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, start_pos: int = 0, user_id: int = 0) -> ttnn.Tensor: + """Legacy path: embed + RoPE + blocks + final norm (no page table). For tests only.""" + x = self.embed_prefill(token_ids_tt) + seq_len = x.shape[2] + assert seq_len % 128 == 0, "prefill seq_len must be divisible by 128" + rot = self.rope_setup.prefill_forward(start_pos, seq_len) + h = x + for layer in self.layers: + h = layer.prefill_forward(h, rot, user_id=user_id, page_table=None) + h = self.norm.prefill_forward(h) + return _all_gather_rmsnorm_tensor(self.norm, h) + + def decode_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, current_pos: int) -> ttnn.Tensor: + """Legacy path: single-token decode without paged ``page_table``.""" + x = self.embed.forward(token_ids_tt) + x = ttnn.unsqueeze_to_4D(x) + pos = torch.tensor([current_pos], dtype=torch.long) + rot_idxs = prepare_rot_idxs(self.rope_setup.config, pos, on_host=False) + rot = self.rope_setup.decode_forward(rot_idxs) + cur = ttnn.from_torch( + pos, + device=self.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(self.mesh_device), + ) + h = x + for layer in self.layers: + h = layer.decode_forward(h, cur, rot, page_table=None) + h = _all_gather_rmsnorm_tensor(self.norm, h, memory_config=self.norm.config.decode_memory_config) + return self.norm.forward(h, "decode") + + def lm_logits(self, hidden: ttnn.Tensor) -> ttnn.Tensor: + """Project last hidden to logits (vocab-sharded on multi-device). + + Skip the explicit interleaved→shard if the caller already produced a sharded + input (``decode_from_token_ids`` returns the decode-mode RMSNorm's width-sharded + output, which already matches ``LMHead1D.config.input_memcfg``). + """ + x = hidden + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded() and not x.memory_config().is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + return self.lm_head.forward(x) diff --git a/code/models/common/models/qwen25_coder_32b/weight_utils.py b/code/models/common/models/qwen25_coder_32b/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..342f13eaa1e96b4de5e41722655eadd5b4788383 --- /dev/null +++ b/code/models/common/models/qwen25_coder_32b/weight_utils.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side HF weight layout helpers for Qwen2.5-Coder-32B-Instruct (TTTv2 port). + +Q/K linear weights use HuggingFace head layout; TTNN RoPE expects Meta layout. +This duplicates the small permute helpers from ``load_checkpoints.reverse_permute`` +without importing ``models/tt_transformers``. Logic is shape-agnostic and mirrors +the Qwen2.5-7B port; the Qwen2 family checkpoint exposes QKV biases but no +``q_norm`` / ``k_norm`` (those are Qwen3). +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.modules.lazy_weight import LazyWeight + + +def reverse_permute(tensor: torch.Tensor, n_heads: int, dim1: int, dim2: int) -> torch.Tensor: + """HF Q/K weight rows → Meta layout (RoPE-compatible).""" + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def reverse_permute_1d(tensor: torch.Tensor) -> torch.Tensor: + """Undo HF split of last dim into interleaved complex pairs for norm/bias vectors.""" + shape = tensor.shape + dim = shape[-1] + assert dim % 2 == 0, "Last dimension must be even" + reals = tensor[..., : dim // 2] + imags = tensor[..., dim // 2 :] + interleaved = torch.stack((reals, imags), dim=-1).flatten(start_dim=len(shape) - 1) + return interleaved + + +def permute_hf_rope_to_meta_tables(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert HF cos/sin tables to Meta-style tables for ``RotarySetup1D`` (see attention tests).""" + if len(cos.shape) == 3: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + return cos, sin + + +def build_rope_cos_sin_torch( + rotary_emb: Any, table_len: int, head_dim: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Build ``[1, 1, table_len, head_dim]`` cos/sin tensors (Meta layout) from HF rotary module.""" + x = torch.zeros(1, 1, table_len, head_dim, dtype=dtype) + position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0) + with torch.no_grad(): + cos_hf, sin_hf = rotary_emb(x, position_ids) + cos_m, sin_m = permute_hf_rope_to_meta_tables(cos_hf.float(), sin_hf.float()) + return cos_m.to(dtype), sin_m.to(dtype) + + +def attention_wqkv_wo_from_hf_layer( + self_attn: Any, + num_devices: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Mirror ``get_attention_weights_from_ref_model`` in ``test_attention_1d.py`` (HF module API).""" + wq_raw = self_attn.q_proj.weight + wk_raw = self_attn.k_proj.weight + wv_raw = self_attn.v_proj.weight + wo_raw = self_attn.o_proj.weight + + dim = wq_raw.shape[1] + cfg = self_attn.config + n_heads = cfg.num_attention_heads + n_kv_heads = cfg.num_key_value_heads + head_dim = cfg.hidden_size // n_heads + n_heads_times_head_dim = n_heads * head_dim + n_kv_heads_times_head_dim = n_kv_heads * head_dim + + wq_meta = reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim) + wk_meta = reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim) + wv_meta = wv_raw + wo_meta = wo_raw + + wq = wq_meta.T + wk = wk_meta.T + wv = wv_meta.T + wo = wo_meta.T + + qkv_list = [] + for i in range(num_devices): + wq_chunk = torch.chunk(wq, num_devices, dim=1)[i] + wk_chunk = torch.chunk(wk, num_devices, dim=1)[i] + wv_chunk = torch.chunk(wv, num_devices, dim=1)[i] + qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1) + qkv_list.append(qkv) + wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0).clone() + wo = wo.unsqueeze(0).unsqueeze(0).clone() + + q_norm = None + k_norm = None + if getattr(self_attn, "q_norm", None) is not None: + q_norm = reverse_permute_1d(self_attn.q_norm.weight).clone() + if getattr(self_attn, "k_norm", None) is not None: + k_norm = reverse_permute_1d(self_attn.k_norm.weight).clone() + + wqkv_bias = None + if getattr(self_attn.q_proj, "bias", None) is not None: + bq_raw = self_attn.q_proj.bias + bk_raw = self_attn.k_proj.bias + bv_raw = self_attn.v_proj.bias + bq_meta = reverse_permute_1d(bq_raw.view(n_heads, head_dim)).view(-1) + bk_meta = reverse_permute_1d(bk_raw.view(n_kv_heads, head_dim)).view(-1) + bv_meta = bv_raw + qkv_bias_list = [] + for i in range(num_devices): + bq_chunk = torch.chunk(bq_meta, num_devices, dim=0)[i] + bk_chunk = torch.chunk(bk_meta, num_devices, dim=0)[i] + bv_chunk = torch.chunk(bv_meta, num_devices, dim=0)[i] + qkv_bias_list.append(torch.cat([bq_chunk, bk_chunk, bv_chunk], dim=-1)) + wqkv_bias = torch.cat(qkv_bias_list, dim=-1).clone() + + return wqkv, wo, q_norm, k_norm, wqkv_bias + + +def mlp_weights_from_hf_layer(mlp: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (w1, w2, w3) TTNN layouts: gate^T, down^T, up^T for ``MLP1D``.""" + w1 = mlp.gate_proj.weight.T.contiguous().clone() + w3 = mlp.up_proj.weight.T.contiguous().clone() + w2 = mlp.down_proj.weight.T.contiguous().clone() + return w1, w2, w3 + + +def rms_weight_torch(layernorm: Any) -> torch.Tensor: + return layernorm.weight.detach().float().clone() + + +def embed_tokens_torch(embed: Any) -> torch.Tensor: + w = embed.weight.detach().float().clone() + return w.unsqueeze(0).unsqueeze(0) + + +def build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """ + Column-split LM head for ``LMHead1D`` (adapted from ``LMHead1D.from_model_args`` math, no v1 args). + ``lm_head_weight`` shape: ``[vocab_size, dim]`` (HF). + + Returns ``(lazy_weights, split_sizes, weights_memcfgs)`` for ``LMHead1DConfig`` DRAM-sharded matmuls. + """ + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / 32) * 32 + if vocab_size < padded_vocab_size: + pad = padded_vocab_size - vocab_size + torch_w = torch.cat([torch_w, torch.zeros(torch_w.shape[0], pad, dtype=torch_w.dtype)], dim=-1) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + tile = ttnn.TILE_SIZE + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (tile * dram_size.x)) * (tile * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights: list[LazyWeight] = [] + weights_memcfgs: list[ttnn.MemoryConfig] = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_splits.append(torch_w[:, start:end]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + name = f"lm_head_split_{i}_{combined.shape[-1]}" + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, name) if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs diff --git a/code/models/common/models/qwen2_7b/README.md b/code/models/common/models/qwen2_7b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b4f461486153f17c35544d5efadd82c458348882 --- /dev/null +++ b/code/models/common/models/qwen2_7b/README.md @@ -0,0 +1,44 @@ +# Qwen2-7B with TTTv2 + +This directory contains the Qwen2-7B TTTv2 product path. + +## Construction path + +```text +HF checkpoint + -> hf_adaptor.py: provider configuration, tokenizer, and weights + -> model.py: Qwen2 tensor graph composed from TTTv2 modules + -> executor.py: thin typed entry point into qwen2_executor.py + -> generator.py: vLLM boundary, DP composition, and dispatch +``` + +The model composes reusable embedding, rotary, RMSNorm, attention, MLP, +LM-head, optional sampling, and TT collective modules. Qwen2 architecture and +device tuning remain model-owned. + +## Executor composition + +The local `executor.py` retains the public Qwen2 class/config/builder names. +`models/common/models/qwen2_executor.py` supplies Qwen2-family policy and +composes `models/common/models/executor.py::ModelExecutor`. + +The shared owner constructs paged-KV management, output reading, +prefill/decode runtimes, eager/trace execution, warmup, and ordered cleanup. +The family layer preserves Qwen2-7B's narrow request signatures and Q128 +top-k tile-end priming before traced warmup and after eager warmup. This +structural refactor does not add Qwen3 native sampling state. + +## vLLM, DP, and ownership + +The generator builds one executor per lane and configures `VLLMAdapter`. DP1 +uses one lane; larger DP uses `LaneGroupExecutor`. Executors own TT resources; +the generator and adapter own dispatch policy only. + +## Tests + +- `models/common/tests/models/qwen2_7b/test_hf_adaptor.py` +- `models/common/tests/models/qwen2_7b/test_demo_contract.py` +- `models/common/tests/demos/qwen2_7b/demo.py` +- `models/common/tests/models/test_qwen2_executor_family.py` +- `models/common/tests/llm_runtime/test_executor_integration.py` +- `models/common/tests/llm_runtime/test_model_executor.py` diff --git a/code/models/common/models/qwen2_7b/__init__.py b/code/models/common/models/qwen2_7b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..78dc51889282b317a56ab5d181145ed4aa7527c0 --- /dev/null +++ b/code/models/common/models/qwen2_7b/__init__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +from models.common.models.qwen2_7b.executor import Qwen2Executor, Qwen2ExecutorConfig +from models.common.models.qwen2_7b.generator import Qwen2Generator, Qwen2GeneratorConfig +from models.common.models.qwen2_7b.hf_adaptor import Qwen2ForCausalLM, Qwen2RuntimeConfig +from models.common.models.qwen2_7b.model import ( + QWEN2_7B_ACCURACY, + QWEN2_7B_PERFORMANCE, + Qwen2PagedAttentionConfig, + Qwen2_7B, + Qwen2_7BPrecisionConfig, + Qwen2_7BTransformerConfig, +) + +__all__ = [ + "QWEN2_7B_ACCURACY", + "QWEN2_7B_PERFORMANCE", + "Qwen2_7B", + "Qwen2_7BPrecisionConfig", + "Qwen2_7BTransformerConfig", + "Qwen2Executor", + "Qwen2ExecutorConfig", + "Qwen2ForCausalLM", + "Qwen2PagedAttentionConfig", + "Qwen2Generator", + "Qwen2GeneratorConfig", + "Qwen2RuntimeConfig", +] diff --git a/code/models/common/models/qwen2_7b/executor.py b/code/models/common/models/qwen2_7b/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..470d04590a60cf54e1e61aeeeda4b61a16d547fb --- /dev/null +++ b/code/models/common/models/qwen2_7b/executor.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen2-7B construction entry point over the shared Qwen2 executor.""" + +from models.common.models.qwen2_7b.hf_adaptor import Qwen2ForCausalLM +from models.common.models.qwen2_executor import Qwen2Executor, Qwen2ExecutorConfig +from models.common.models.qwen2_executor import build_qwen2_7b_executor as _build_qwen2_7b_executor + + +def build_qwen2_7b_executor(llm: Qwen2ForCausalLM, config: Qwen2ExecutorConfig) -> Qwen2Executor: + """Build one executor around an already-loaded Qwen2 model adapter.""" + + return _build_qwen2_7b_executor(llm, config) diff --git a/code/models/common/models/qwen2_7b/generator.py b/code/models/common/models/qwen2_7b/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..07d3ea4cb18d5a76d35ca0eb70aae52b50ffc308 --- /dev/null +++ b/code/models/common/models/qwen2_7b/generator.py @@ -0,0 +1,481 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM construction and compatibility delegation for Qwen2-7B.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +import ttnn +from models.common.llm_runtime.config import PagedKVCacheConfig, TraceConfig, TraceMode, WarmupConfig +from models.common.llm_runtime.lane_group import LaneGroupExecutor +from models.common.llm_runtime.vllm_adapter import NormalizedPrefillKwargs, VLLMAdapter, VLLMAdapterConfig +from models.common.models.qwen2_7b.executor import Qwen2ExecutorConfig, build_qwen2_7b_executor +from models.common.models.qwen2_7b.hf_adaptor import from_pretrained +from models.common.models.qwen2_7b.model import Qwen2PagedAttentionConfig + +_PROVISIONAL_BLOCK_SIZE = 32 + + +@dataclass(frozen=True) +class Qwen2GeneratorConfig: + """Validated construction inputs for one vLLM-facing Qwen generator.""" + + hf_model: str + mesh_device: Any + max_batch_size: int + max_seq_len: int + hf_revision: str | None = None + n_layers: int | None = None + tt_data_parallel: int = 1 + optimizations: Any = "performance" + trace_mode: TraceMode = "all" + device_sampling_enabled: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.hf_model, str) or not self.hf_model: + raise ValueError("hf_model must be a non-empty string") + if self.mesh_device is None: + raise ValueError("mesh_device is required") + _validate_positive_int("max_batch_size", self.max_batch_size) + _validate_positive_int("max_seq_len", self.max_seq_len) + _validate_positive_int("tt_data_parallel", self.tt_data_parallel) + if self.n_layers is not None: + _validate_positive_int("n_layers", self.n_layers) + if self.max_batch_size % self.tt_data_parallel != 0: + raise ValueError( + f"max_batch_size={self.max_batch_size} must be divisible by " + f"tt_data_parallel={self.tt_data_parallel}" + ) + if not isinstance(self.device_sampling_enabled, bool): + raise TypeError("device_sampling_enabled must be bool") + TraceConfig(mode=self.trace_mode) + + +class Qwen2Generator: + """Adapt vLLM's model interface to the model-owned execution target. + + vLLM constructs this facade with `initialize_vllm_model`, resolves + KV capacity through `allocate_kv_cache`, warms the configured + programs, and then calls `prefill_forward` and + `decode_forward`. Each forward call is normalized by + `VLLMAdapter`, dispatched to the eager or traced executor, and + delegated to ``Qwen2Executor`` or ``LaneGroupExecutor``. + + This class owns dispatch policy but no TT resources. `cleanup` + delegates to the target that owns those resources. + """ + + model_capabilities = { + "supports_prefix_caching": True, + "supports_async_decode": True, + "supports_sample_on_device": True, + "accepts_trace_mode": True, + } + requires_prefill_trace_warmup = True + + def __init__(self, target: Any, adapter: VLLMAdapter): + self.target = target + self._adapter = adapter + + # Public vLLM API + + @property + def model(self): + return self.target.model + + @property + def model_args(self): + return self.target.model_args + + @property + def mesh_device(self): + return self.target.mesh_device + + @property + def cache_path(self): + return self.target.cache_path + + @property + def already_warmed_up_prefill(self): + return self.target.already_warmed_up_prefill + + @already_warmed_up_prefill.setter + def already_warmed_up_prefill(self, value): + self.target.already_warmed_up_prefill = value + + @classmethod + def get_max_tokens_all_users( + cls, + model_name: str = "", + num_devices: int = 1, + tt_data_parallel: int = 1, + max_model_len: int = 0, + max_num_seqs: int = 1, + ) -> int: + """Return the unpadded per-submesh KV token budget for vLLM sizing.""" + + return int(max_model_len) + + @classmethod + def initialize_vllm_model( + cls, + hf_config, + mesh_device, + max_batch_size, + max_seq_len, + n_layers=None, + tt_data_parallel=1, + optimizations="performance", + trace_mode: TraceMode = "all", + hf_revision: str | None = None, + device_sampling_enabled: bool = True, + ): + """Build the configured single-lane or data-parallel Qwen target.""" + + hf_model = getattr(hf_config, "_name_or_path", None) + if not hf_model: + raise ValueError("hf_config must provide a non-empty _name_or_path") + return build_qwen2_7b_generator( + Qwen2GeneratorConfig( + hf_model=str(hf_model), + hf_revision=hf_revision, + mesh_device=mesh_device, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + n_layers=n_layers, + tt_data_parallel=tt_data_parallel, + optimizations=optimizations, + trace_mode=trace_mode, + device_sampling_enabled=device_sampling_enabled, + ) + ) + + def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): + """Resolve the late vLLM capacity, then allocate a borrowed cache handle.""" + + supplied = (kv_cache_shape is not None, dtype is not None, num_layers is not None) + if not any(supplied): + return self.target.allocate_kv_cache() + if not all(supplied): + raise TypeError("kv_cache_shape, dtype, and num_layers must be supplied together") + + resolved = self._adapter.resolve_legacy_kv_cache_config(kv_cache_shape, dtype, num_layers) + self.target.configure_paged_kv_cache(resolved) + return self.target.allocate_kv_cache() + + def compile_prefill( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + ) -> None: + """Normalize a vLLM prefill call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.compile_prefill(execution=execution, **normalized) + + def compile_decode( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + ) -> None: + """Normalize a vLLM decode call and compile its selected target.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.compile_decode(execution=execution, **normalized) + + def prefill_forward( + self, + tokens: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + prompt_lens: Sequence[int] | torch.Tensor | None = None, # ↓ Sequence metadata + start_pos: torch.Tensor | None = None, + empty_slots: Sequence[int] | None = None, # ↓ Lane routing + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM prefill call.""" + + normalized, trace_requested = self._adapter.normalize_prefill( + tokens, + page_table, + enable_trace=enable_trace, + prompt_lens=prompt_lens, + start_pos=start_pos, + empty_slots=empty_slots, + kv_cache=kv_cache, + sampling_params=sampling_params, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_prefill_execution(normalized, trace_requested) + return self.target.prefill_forward(execution=execution, **normalized) + + def decode_forward( + self, + tokens: torch.Tensor, + start_pos: torch.Tensor, + page_table: torch.Tensor, + *, + enable_trace: bool, # ↓ Required policy + kv_cache: Any = None, # ↓ Borrowed resources + sampling_params: Any = None, # ↓ Sampling + reset_batch: bool = False, # ↓ State transition + read_from_device: bool = True, # ↓ Output policy + **compatibility_kwargs: Any, # ↓ Compatibility + ) -> Any: + """Normalize and dispatch one vLLM decode call.""" + + normalized, trace_requested = self._adapter.normalize_decode( + tokens, + start_pos, + page_table, + enable_trace=enable_trace, + kv_cache=kv_cache, + sampling_params=sampling_params, + reset_batch=reset_batch, + compatibility_kwargs=compatibility_kwargs, + ) + execution = self._select_execution("decode", trace_requested) + return self.target.decode_forward( + execution=execution, + read_from_device=read_from_device, + **normalized, + ) + + def read_decode_output( + self, + tt_out: Any, + *, + async_read: bool = False, + ) -> Any: + """Delegate vLLM's raw decode-output read.""" + + return self.target.read_decode_output(tt_out=tt_out, async_read=async_read) + + def process_decode_output_host( + self, + tt_out: Any, + *, + is_tokens: bool = False, + ) -> tuple[Any, Any]: + """Delegate vLLM's asynchronous host-output completion.""" + + return self.target.process_decode_output_host(tt_out=tt_out, is_tokens=is_tokens) + + def warmup_model_prefill( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_prefill( + kv_cache=kv_cache, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def warmup_model_decode( + self, + *, + kv_cache: Any, # ↓ Borrowed resources + max_batch_size: int, # ↓ Coverage dimensions + num_blocks: int, + can_sample_on_device: bool, # ↓ Execution policy + enable_trace: bool, + ) -> None: + return self.target.warmup_model_decode( + kv_cache=kv_cache, + max_batch_size=max_batch_size, + num_blocks=num_blocks, + can_sample_on_device=can_sample_on_device, + enable_trace=enable_trace, + ) + + def cleanup(self): + """Release every resource owned by the concrete target.""" + + return self.target.cleanup() + + # Private implementation + + def _select_prefill_execution( + self, + normalized: NormalizedPrefillKwargs, + trace_requested: bool, + ): + # Static trace intent is authoritative. Eligibility and configured + # coverage are preflighted by the selected execution target; this + # facade must never turn a required trace miss into eager KV writes. + return self._select_execution("prefill", trace_requested) + + def _select_execution(self, operation: str, enable_trace: bool): + if not enable_trace: + return self.target.eager_execution + execution = getattr(self.target, f"traced_{operation}_execution") + if execution is None: + raise RuntimeError(f"vLLM requested unavailable traced {operation} execution") + return execution + + +def build_qwen2_7b_generator(config: Qwen2GeneratorConfig) -> Qwen2Generator: + """Construct lane-local models/executors and compose their shared target surface.""" + + per_lane_max_batch_size = config.max_batch_size // config.tt_data_parallel + submeshes = ( + [config.mesh_device] + if config.tt_data_parallel == 1 + else list(_create_submeshes(config.mesh_device, config.tt_data_parallel)) + ) + if len(submeshes) != config.tt_data_parallel: + raise ValueError(f"Expected {config.tt_data_parallel} submeshes, got {len(submeshes)}") + + max_num_blocks = ( + config.max_seq_len + _PROVISIONAL_BLOCK_SIZE - 1 + ) // _PROVISIONAL_BLOCK_SIZE + per_lane_max_batch_size + lanes = [] + try: + for submesh in submeshes: + paged_attention_config = Qwen2PagedAttentionConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + ) + llm = from_pretrained( + mesh_device=submesh, + hf_model=config.hf_model, + hf_revision=config.hf_revision, + instruct="Instruct" in config.hf_model, + max_batch_size=per_lane_max_batch_size, + max_seq_len=config.max_seq_len, + optimizations=config.optimizations, + n_layers=config.n_layers, + dtype=ttnn.bfloat8_b, + paged_attention_config=paged_attention_config, + ) + model_kv_cache_dtypes, _, _, _ = _model_kv_metadata(llm.model) + executor_config = Qwen2ExecutorConfig( + trace=TraceConfig(mode=config.trace_mode), + warmup=WarmupConfig(), + paged_kv_cache=PagedKVCacheConfig( + block_size=_PROVISIONAL_BLOCK_SIZE, + max_num_blocks=max_num_blocks, + dtype=model_kv_cache_dtypes[0], + ), + device_sampling_enabled=config.device_sampling_enabled, + ) + lanes.append(build_qwen2_7b_executor(llm, executor_config)) + + adapter = _build_vllm_adapter(lanes[0]) + except BaseException as primary: + _cleanup_after_construction_failure(lanes, primary) + raise + + target = lanes[0] if config.tt_data_parallel == 1 else LaneGroupExecutor(lanes, mesh_device=config.mesh_device) + return Qwen2Generator(target, adapter) + + +def _build_vllm_adapter(lane) -> VLLMAdapter: + model_kv_cache_dtypes, num_layers, kv_heads_per_device, head_dim = _model_kv_metadata(lane.model) + return VLLMAdapter( + VLLMAdapterConfig.resolve( + trace=lane.config.trace, + paged_kv_cache=lane.config.paged_kv_cache, + expected_num_layers=num_layers, + expected_kv_heads_per_device=kv_heads_per_device, + expected_head_dim=head_dim, + model_kv_cache_dtype=model_kv_cache_dtypes, + ) + ) + + +def _model_kv_metadata(model) -> tuple[tuple[Any, ...], int, int, int]: + layers = tuple(getattr(model, "layers", ())) + if not layers: + raise ValueError("Qwen2-7B model must contain at least one attention layer") + + attention_configs = tuple(layer.attention.config for layer in layers) + model_config = model.config + num_layers = int(model_config.n_layers) + if len(attention_configs) != num_layers: + raise ValueError(f"Model config declares {num_layers} layers but exposes {len(attention_configs)}") + + num_devices = int(model_config.num_devices) + n_kv_heads = int(attention_configs[0].n_kv_heads) + if n_kv_heads % num_devices != 0: + raise ValueError(f"n_kv_heads={n_kv_heads} must be divisible by num_devices={num_devices}") + + head_dim = int(attention_configs[0].head_dim) + if any( + int(attention_config.n_kv_heads) != n_kv_heads or int(attention_config.head_dim) != head_dim + for attention_config in attention_configs + ): + raise ValueError("Every Qwen layer must expose the same KV head shape") + + return ( + tuple(attention_config.kv_cache_dtype for attention_config in attention_configs), + num_layers, + n_kv_heads // num_devices, + head_dim, + ) + + +def _create_submeshes(mesh_device, tt_data_parallel): + from models.tt_transformers.tt.generator import create_submeshes + + return create_submeshes(mesh_device, tt_data_parallel) + + +def _cleanup_after_construction_failure(lanes, primary): + failures = [] + for lane in lanes: + try: + lane.cleanup() + except BaseException as error: + failures.append(error) + if failures: + setattr(primary, "cleanup_failures", failures) + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") diff --git a/code/models/common/models/qwen2_7b/hf_adaptor.py b/code/models/common/models/qwen2_7b/hf_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..23646bac62a81e0b6f40ead008c84b8cda2a2dd3 --- /dev/null +++ b/code/models/common/models/qwen2_7b/hf_adaptor.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face product adaptor for the TTTv2 Qwen2-7B path.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import ttnn +from models.common.models.qwen2_7b import weight_utils +from models.common.models.qwen2_7b.model import ( + QWEN2_7B_ACCURACY, + QWEN2_7B_PERFORMANCE, + Qwen2_7B, + Qwen2_7BLayerWeights, + Qwen2_7BModelParameters, + Qwen2_7BPrecisionConfig, + Qwen2_7BWeights, + Qwen2PagedAttentionConfig, + build_qwen2_7b_transformer_config, +) + +DEFAULT_HF_MODEL = "Qwen/Qwen2-7B-Instruct" + + +@dataclass(frozen=True) +class Qwen2GenerationConfig: + max_decode_tokens: int = 128 + temperature: float = 0.0 + top_k: int = 32 + top_p: float = 0.08 + stop_token_ids: tuple[int, ...] = () + + +@dataclass(frozen=True) +class Qwen2RuntimeConfig: + model_name: str + model_cache_path: Path | None + max_prefill_chunk_size: int + max_context_len: int + max_seq_len: int + trace_prefill_supported_seq_lens: tuple[int, ...] + supports_batched_prefill: bool = True + max_prefill_batch_size: int = 32 + disable_batched_prefill: bool = False + batched_prefill_batched_extract: bool = True + + def can_enable_trace(self, prefill_seq_len: int, num_cached_tokens: int = 0) -> bool: + return ( + prefill_seq_len in self.trace_prefill_supported_seq_lens + and prefill_seq_len <= self.max_prefill_chunk_size + and prefill_seq_len <= self.max_seq_len + ) + + +def _chat_template_ids(encoded): + if hasattr(encoded, "keys") and "input_ids" in encoded: + encoded = encoded["input_ids"] + if hasattr(encoded, "ids"): + return list(encoded.ids) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if isinstance(encoded, (list, tuple)) and len(encoded) == 1 and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return list(encoded) + + +def encode_prompt(tokenizer, prompt_text, system_prompt_text=None, *, instruct=True): + if instruct: + chat = [] + if isinstance(prompt_text, str): + if system_prompt_text: + chat.append({"role": "system", "content": system_prompt_text}) + if prompt_text: + chat.append({"role": "user", "content": prompt_text}) + else: + chat = prompt_text + try: + return _chat_template_ids(tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True)) + except ValueError: + pass + return tokenizer.encode(prompt_text, add_special_tokens=False) + + +@dataclass +class Qwen2ForCausalLM: + model: Qwen2_7B + tokenizer: Any + runtime_config: Qwen2RuntimeConfig + instruct: bool = True + generation_config: Qwen2GenerationConfig = field(default_factory=Qwen2GenerationConfig) + + def __post_init__(self): + self.model.model_args = self.runtime_config + if not self.generation_config.stop_token_ids: + stops = tuple(getattr(self.tokenizer, "stop_tokens", []) or []) + self.generation_config = Qwen2GenerationConfig( + max_decode_tokens=self.generation_config.max_decode_tokens, + temperature=self.generation_config.temperature, + top_k=self.generation_config.top_k, + top_p=self.generation_config.top_p, + stop_token_ids=stops, + ) + + @property + def model_name(self): + return self.runtime_config.model_name + + @property + def model_cache_path(self): + return self.runtime_config.model_cache_path + + @property + def max_seq_len(self): + return self.model.config.max_seq_len + + @property + def max_context_len(self): + return self.runtime_config.max_context_len + + def encode_prompt(self, prompt_text, system_prompt_text=None, instruct=None): + return encode_prompt( + self.tokenizer, + prompt_text, + system_prompt_text, + instruct=self.instruct if instruct is None else instruct, + ) + + def encode_chat(self, messages): + return self.encode_prompt(messages, instruct=True) + + +def load_tokenizer(hf_model: str, hf_revision: str | None = None): + tokenizer = AutoTokenizer.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=os.getenv("CI") == "true", + ) + eos = getattr(tokenizer, "eos_token_id", None) + stops = [] if eos is None else ([eos] if isinstance(eos, int) else list(eos)) + im_start = tokenizer.convert_tokens_to_ids("<|im_start|>") + if isinstance(im_start, int) and im_start >= 0 and im_start not in stops: + stops.append(im_start) + tokenizer.stop_tokens = stops + return tokenizer + + +def _trace_seq_lens(num_devices: int, max_prefill_chunk_size: int, max_seq_len: int) -> tuple[int, ...]: + if num_devices != 2: + raise ValueError(f"Qwen2-7B supports TP2/N300 only, got {num_devices} devices") + allowed = (128, 1024) + return tuple(length for length in allowed if length <= min(max_prefill_chunk_size, max_seq_len)) + + +def _cache_path(hf_model: str, mesh_device, cache_dir: Path | str | None) -> Path: + if cache_dir is not None: + path = Path(cache_dir) + elif os.getenv("TT_CACHE_PATH"): + path = Path(os.environ["TT_CACHE_PATH"]) + else: + path = Path("model_cache") / hf_model / "N300" + path.mkdir(parents=True, exist_ok=True) + return path + + +def convert_hf_model_weights( + hf, + hf_config, + *, + n_layers: int, + num_devices: int, + rope_table_len: int, + head_dim: int, +) -> Qwen2_7BWeights: + """Extract and convert all Hugging Face tensors consumed by the TT builder.""" + base = hf.model + rope_cos, rope_sin = weight_utils.build_rope_cos_sin_torch( + base.rotary_emb, rope_table_len, head_dim, torch.bfloat16 + ) + layers = [] + for layer in base.layers[:n_layers]: + wqkv, wo, q_norm, k_norm, wqkv_bias = weight_utils.attention_wqkv_wo_from_hf_layer(layer.self_attn, num_devices) + if q_norm is not None or k_norm is not None: + raise ValueError("Qwen2-7B does not support QK norm weights") + if wqkv_bias is None: + raise ValueError("Qwen2-7B requires QKV projection bias") + w1, w2, w3 = weight_utils.mlp_weights_from_hf_layer(layer.mlp) + layers.append( + Qwen2_7BLayerWeights( + wqkv=wqkv, + wo=wo, + wqkv_bias=wqkv_bias, + w1=w1, + w2=w2, + w3=w3, + attention_norm=weight_utils.rms_weight_torch(layer.input_layernorm).to(torch.bfloat16), + ff_norm=weight_utils.rms_weight_torch(layer.post_attention_layernorm).to(torch.bfloat16), + ) + ) + + if hf_config.tie_word_embeddings: + raise ValueError("Qwen2-7B requires an untied LM head") + return Qwen2_7BWeights( + embedding=weight_utils.embed_tokens_torch(base.embed_tokens), + rope_cos=rope_cos, + rope_sin=rope_sin, + layers=tuple(layers), + final_norm=weight_utils.rms_weight_torch(base.norm).to(torch.bfloat16), + lm_head=hf.lm_head.weight.detach().to(torch.bfloat16).clone(), + ) + + +def from_pretrained( + mesh_device, + *, + hf_model: str = DEFAULT_HF_MODEL, + hf_revision: str | None = None, + instruct: bool = True, + max_batch_size: int = 32, + max_seq_len: int = 4096, + optimizations: str | Qwen2_7BPrecisionConfig = "accuracy", + n_layers: int | None = None, + dtype=ttnn.bfloat8_b, + paged_attention_config: Qwen2PagedAttentionConfig | None = None, + cache_dir: Path | str | None = None, +) -> Qwen2ForCausalLM: + del dtype + if mesh_device.get_num_devices() != 2: + raise ValueError(f"Qwen2-7B supports TP2/N300 only, got {mesh_device.get_num_devices()} devices") + ttnn.SetDefaultDevice(mesh_device) + hf_config = AutoConfig.from_pretrained( + hf_model, + revision=hf_revision, + local_files_only=os.getenv("CI") == "true", + ) + hf = AutoModelForCausalLM.from_pretrained( + hf_model, + revision=hf_revision, + torch_dtype=torch.bfloat16, + local_files_only=os.getenv("CI") == "true", + ) + hf.eval() + resolved_layers = hf_config.num_hidden_layers if n_layers is None else n_layers + precision = ( + optimizations + if isinstance(optimizations, Qwen2_7BPrecisionConfig) + else (QWEN2_7B_PERFORMANCE if optimizations == "performance" else QWEN2_7B_ACCURACY) + ) + if not isinstance(precision, Qwen2_7BPrecisionConfig): + raise TypeError("optimizations must be 'accuracy', 'performance', or Qwen2_7BPrecisionConfig") + cache_path = _cache_path(hf_model, mesh_device, cache_dir) + if paged_attention_config is None: + block_size = 32 + paged_attention_config = Qwen2PagedAttentionConfig( + block_size=block_size, + max_num_blocks=((max_seq_len + block_size - 1) // block_size) * max_batch_size, + ) + head_dim = hf_config.hidden_size // hf_config.num_attention_heads + params = Qwen2_7BModelParameters( + dim=hf_config.hidden_size, + n_heads=hf_config.num_attention_heads, + n_kv_heads=hf_config.num_key_value_heads, + head_dim=head_dim, + hidden_dim=hf_config.intermediate_size, + vocab_size=hf_config.vocab_size, + rms_norm_eps=hf_config.rms_norm_eps, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + ) + rope_table_len = ((max(max_seq_len * 2, 8192) + 127) // 128) * 128 + weights = convert_hf_model_weights( + hf, + hf_config, + n_layers=resolved_layers, + num_devices=mesh_device.get_num_devices(), + rope_table_len=rope_table_len, + head_dim=head_dim, + ) + model_config = build_qwen2_7b_transformer_config( + mesh_device=mesh_device, + params=params, + weights=weights, + n_layers=resolved_layers, + precision=precision, + cache_path=cache_path, + paged_attention_config=paged_attention_config, + model_name=Path(hf_model).name, + ) + tokenizer = load_tokenizer(hf_model, hf_revision) + model = Qwen2_7B(model_config) + max_prefill_chunk_size = 2048 + runtime_config = Qwen2RuntimeConfig( + model_name=Path(hf_model).name, + model_cache_path=cache_path, + max_prefill_chunk_size=max_prefill_chunk_size, + max_context_len=int(hf_config.max_position_embeddings), + max_seq_len=max_seq_len, + trace_prefill_supported_seq_lens=_trace_seq_lens( + mesh_device.get_num_devices(), max_prefill_chunk_size, max_seq_len + ), + disable_batched_prefill=bool(os.getenv("DISABLE_BATCHED_PREFILL")), + ) + del hf + return Qwen2ForCausalLM( + model=model, + tokenizer=tokenizer, + runtime_config=runtime_config, + instruct=instruct, + ) diff --git a/code/models/common/models/qwen2_7b/model.py b/code/models/common/models/qwen2_7b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..efd949a602147c1826def377bae60435cb15a604 --- /dev/null +++ b/code/models/common/models/qwen2_7b/model.py @@ -0,0 +1,1166 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Qwen2-7B-Instruct — native TTTv2 stack (``Embedding1D``, ``RMSNorm1D``, +``Attention1D``, ``MLP1D``, ``RotarySetup1D``, ``LMHead1D``). + +Tensor layout contracts: + - **Prefill** hidden states: ``[1, 1, S, dim]`` TILE, ``S % 128 == 0``. + - **Decode** hidden states: ``[1, 1, B, dim]`` TILE (``B`` padded to tile in modules). + +The tensor model is provider-neutral. Hugging Face loading and conversion live +in :mod:`models.common.models.qwen2_7b.hf_adaptor`. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.attention.attention_1d import ( + Attention1D, + Attention1DConfig, + _dram_matmul_config, + _dram_shard_core_grid, +) +from models.common.modules.embedding.embedding_1d import Embedding1D, Embedding1DConfig +from models.common.modules.lazy_weight import LazyWeight +from models.common.modules.lm_head.lm_head_1d import LMHead1D, LMHead1DConfig, _nearest_32 +from models.common.modules.mlp.mlp_1d import MLP1D, MLP1DConfig, _dram_shard_core_grid_k_n +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig, _create_sharded_norm_program_config +from models.common.modules.rope.rope_1d import Rope1DConfig, RotarySetup1D, prepare_rot_idxs +from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE, get_padded_hidden_dim + + +def _lazy( + tensor: torch.Tensor, + *, + dtype: ttnn.DataType, + cache: tuple[Path, str] | None, +) -> LazyWeight: + """Minimal LazyWeight; ``Attention1D`` / ``MLP1D`` / ``Embedding1D`` resolvers fill mesh + memory.""" + return LazyWeight(source=tensor, dtype=dtype, cache_dir_weight_name=cache) + + +@dataclass +class Qwen2PagedAttentionConfig: + """Paged KV layout for ``Attention1D`` (``block_size`` / ``max_num_blocks`` only).""" + + block_size: int + max_num_blocks: int + + +@dataclass +class Qwen2_7BTransformerConfig: + """Complete provider-neutral tensor-graph configuration.""" + + n_layers: int + vocab_size: int + max_batch_size: int + max_seq_len: int + dim: int + num_devices: int + mesh_device: ttnn.MeshDevice + embedding_config: Embedding1DConfig + rope_config: Rope1DConfig + block_configs: list[Qwen2_7BDecoderLayerConfig] + norm_config: RMSNorm1DConfig + lm_head_config: LMHead1DConfig + sampling_config: Sampling1DConfig | None = None + decode_residual_memcfg: ttnn.MemoryConfig | None = None + prefill_residual_memcfg: ttnn.MemoryConfig | None = None + activation_dtypes: list[ttnn.DataType | None] = field(default_factory=list) + tt_ccl: Any = None + cache_path: str | None = None + + +@dataclass(frozen=True) +class Qwen2_7BModelParameters: + dim: int + n_heads: int + n_kv_heads: int + head_dim: int + hidden_dim: int + vocab_size: int + rms_norm_eps: float + max_batch_size: int + max_seq_len: int + + +@dataclass(frozen=True) +class Qwen2_7BLayerWeights: + wqkv: torch.Tensor + wo: torch.Tensor + wqkv_bias: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w3: torch.Tensor + attention_norm: torch.Tensor + ff_norm: torch.Tensor + + +@dataclass(frozen=True) +class Qwen2_7BWeights: + embedding: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + layers: tuple[Qwen2_7BLayerWeights, ...] + final_norm: torch.Tensor + lm_head: torch.Tensor + + +@dataclass(frozen=True) +class Qwen2_7BPrecisionConfig: + """Per-layer precision + math-fidelity recipe for Qwen2-7B-Instruct. + + Mirrors the fields TTTv1's ``DecodersPrecision`` actually distinguishes for + Qwen2-7B (Llama-family group in ``model_config.py`` for both ``accuracy()`` + and ``performance()``). Two module-level recipes are exposed: + :data:`QWEN2_7B_ACCURACY` and :data:`QWEN2_7B_PERFORMANCE`. Pass one to + the product adaptor via ``optimizations=``; use + ``dataclasses.replace(QWEN2_7B_PERFORMANCE, perf_decode_tuning=False)`` to + customize a single field (e.g. disable aggressive decode math during + teacher-forcing parity runs). + + Attention compute-kernel configs (HIFI4 + fp32 dest acc for LI_QKV / LI_O / + SDPA prefill, plus the decode variants and the LoFi perf-decode kernel) are + resolved inside ``_resolve_qwen_wh_tuning`` from + :attr:`perf_decode_tuning`; they are not exposed here because TTTv1 ships + them as a coupled bundle for this model. + """ + + # Qwen2-7B keeps BF16 WQKV/WO in BOTH TTTv1 recipes — TTTv2's + # ``Attention1D`` default of BFP8 silently downgrades QKV/WO precision and + # broadens per-layer divergence on this model (see + # debugging/numerical_divergence_vs_hf_2026-05-14.md). + wqkv_dtype: ttnn.DataType = ttnn.bfloat16 + wo_dtype: ttnn.DataType = ttnn.bfloat16 + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + + # MLP weights: BFP8 in both modes for Qwen2-7B (TTTv1 does not drop these + # to BFP4 for this checkpoint). + mlp_w1_w3_dtype: ttnn.DataType = ttnn.bfloat8_b + mlp_w2_dtype: ttnn.DataType = ttnn.bfloat8_b + + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + + # Decode-side throughput tuning: LoFi attention decode (LI_QKV / SDPA / + # LI_O), HiFi2 MLP decode FF, SDPA exp-approx, no W1→DRAM spill for small + # batch. Performance recipe only — keep ``False`` in accuracy mode and in + # any teacher-forcing path (aggressive decode math drops top-1). + perf_decode_tuning: bool = False + + +# TTTv1 ``DecodersPrecision.accuracy("Qwen2-7B-Instruct")`` (Llama-family +# group in ``model_config.py``): BF16 attention, BFP8 MLP, BF16 LM head, BF16 +# KV cache, HiFi4 throughout. ``perf_decode_tuning`` stays off — decode keeps +# HiFi4 + fp32 dest acc. +QWEN2_7B_ACCURACY = Qwen2_7BPrecisionConfig( + kv_cache_dtype=ttnn.bfloat16, + lm_head_dtype=ttnn.bfloat16, +) + +# TTTv1 ``DecodersPrecision.performance("Qwen2-7B-Instruct")``: same dtypes +# as accuracy except ``kv_cache_dtype`` / ``lm_head_dtype`` drop to BFP8, plus +# the decode-side throughput bundle (LoFi attn decode + SDPA exp-approx + HiFi2 +# MLP decode FF) is engaged via ``perf_decode_tuning=True``. +QWEN2_7B_PERFORMANCE = Qwen2_7BPrecisionConfig( + perf_decode_tuning=True, +) + + +def _qwen_wh_mlp_matmul_compute_kernel() -> ttnn.WormholeComputeKernelConfig: + """HiFi4 lowers L1 circular-buffer footprint vs HiFi2 for wide FF matmuls on Wormhole.""" + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + + +def _qwen_wh_mlp_decode_matmul_compute_kernel() -> ttnn.WormholeComputeKernelConfig: + """HiFi2 decode FF matmuls: faster steady-state decode while prefill keeps HiFi4 (``_qwen_wh_mlp_matmul_compute_kernel``).""" + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + + +def _qwen_wh_decode_attn_lofi_kernel() -> ttnn.WormholeComputeKernelConfig: + """LoFi decode attention matmuls + SDPA when ``perf_decode_tuning`` is enabled (performance demo only).""" + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.LoFi, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=False, + ) + + +def _qwen_wh_attn_hifi4_kernel() -> ttnn.WormholeComputeKernelConfig: + """HiFi4 + fp32 dest acc for Qwen2-7B attention matmuls (LI_QKV, LI_O, SDPA). + + Matches the TTTv1 ``ModelArgs.compute_kernel_config_hifi4`` used by both + ``ModelOptimizations.performance`` and ``ModelOptimizations.accuracy`` for + Qwen2-7B (see ``models/tt_transformers/tt/model_config.py``). + The TTTv2 ``Attention1D`` defaults are HiFi2 with fp16 accumulation; that + silently downgrades attention prefill QKV/WO and decode QKV/SDPA/WO matmul + precision for this model, producing a broad per-layer divergence vs HF. + """ + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=True, + ) + + +def _slice_last_token_tile(x: ttnn.Tensor, last_token_idx: int) -> ttnn.Tensor: + """Slice the 32-row tile containing ``last_token_idx`` from ``[1, 1, S, W]``. + + Width-sharded LM matmul M tile rows must equal ``LMHead1D`` program-config tile rows. + """ + floor = (last_token_idx // 32) * 32 + return ttnn.slice(x, (0, 0, floor, 0), (1, 1, floor + 32, x.shape[-1])) + + +def _post_attn_norm_decode_configs( + *, + dim: int, + hidden_dim: int, + num_devices: int, + max_batch_size: int, +) -> tuple[Any, ttnn.MemoryConfig]: + """Resolve post-attention RMSNorm decode sharding so its output matches MLP1D's W1/W3 input. + + MLP1D decode uses ``_dram_shard_core_grid_k_n(dim, padded_hidden / num_devices)`` for W1/W3 + inputs, but the default RMSNorm program config is derived from ``_compute_norm_core_grid(dim)`` + alone. On N300 Qwen2-7B that DRAM-width-shard mismatch silently corrupts decode activations. + """ + padded_hidden = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + grid = _dram_shard_core_grid_k_n(dim, padded_hidden // num_devices) + tile_padded_batch_rows = TILE_SIZE * math.ceil(max_batch_size / TILE_SIZE) + program_config = _create_sharded_norm_program_config(dim, grid, tile_padded_batch_rows, TILE_SIZE) + memory_config = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // grid.num_cores), + grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return program_config, memory_config + + +def _all_gather_rmsnorm_tensor( + norm: RMSNorm1D, x: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None +) -> ttnn.Tensor: + cfg = norm.config + if cfg.mesh_device.get_num_devices() == 1 or x.shape[-1] == cfg.weight.source.numel(): + return x + + if memory_config is None: + memory_config = x.memory_config() + + tt_ccl = cfg.tt_ccl or get_tt_ccl(cfg.mesh_device) + return ttnn.experimental.all_gather_async( + x, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=default_topology(cfg.mesh_device), + memory_config=memory_config, + barrier_semaphore=tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + + +@dataclass +class _Qwen2WHTuning: + """Wormhole-specific MLP / LM-head / attention tuning resolved at build time. + + Populated for Qwen2-7B on N150 / N300 only; otherwise all fields stay ``None`` + (modules fall back to library defaults). + """ + + mlp_prefill_len_cutoff: int | None = None + mlp_ff_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + mlp_decode_spill_w1_to_dram: bool = False + mlp_decode_ff_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + lm_head_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + perf_decode_sdpa_cfg: ttnn.SDPAProgramConfig | None = None + # Attention prefill compute kernels: always HiFi4 + fp32 dest acc on Qwen + # (Attention1D defaults are HiFi2/fp16 acc — too lossy for this model). + attn_li_qkv_prefill_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_li_o_prefill_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + # Attention decode compute kernels: HiFi4 + fp32 dest acc by default; LoFi + # under ``perf_decode_tuning`` (performance demo only). + attn_li_qkv_decode_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_sdpa_decode_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + attn_li_o_decode_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + +def _resolve_qwen_wh_tuning( + *, hf_model_id: str, num_dev: int, max_batch_size: int, perf_decode_tuning: bool +) -> _Qwen2WHTuning: + """Pick WH tuning knobs for Qwen2-7B-Instruct on N150 / N300. + + ``get_padded_prefill_len`` maps 129..1024 tokens to a 1024-wide tile. ``MLP1D`` then + reshapes/chunks using ``prefill_len_cutoff``. + + ``mlp_prefill_len_cutoff=1024`` = the shared engine's own Wormhole default (``mlp_1d.py``: + ``512 if is_blackhole() else 1024``) and TTTv1's ``prefill_len_cutoff``. For the folded batch-32-ci + FF prefill (``[1,1,B*S,dim]``, B*S=32*128=4096 at the natural 128 bucket) this tiles the FF matmul as + 4 chunks of 1024 (``per_core_M=4``) instead of 16 chunks of 256 (``per_core_M=1``) — 4× fewer / 4× + larger sub-matmuls on the ~80%-FLOP FF block, matching TTTv1's blocking. + + This supersedes an earlier note here that "cutoff 512 still trips ``validate_circular_buffer_region`` + on WH (prefill multicast + LM DRAM matmul vs L1)". That is no longer true: 1024 was re-tested on + device on N300 (2026-07-27) and capture/replay is clean — no circular-buffer assert — despite this + checkpoint having the widest per-device FF shard of the family (intermediate 18944 → 9472/device). + The L1 headroom that makes it fit is the HiFi4 FF/LM configuration set below (HiFi4 shrinks the CB + footprint vs HiFi2), which was tuned after that observation was recorded. Fallbacks, in order, if a + future config change reintroduces the assert: 512, then 256 (the previous value) — the fold width in + ``max_prefill_batch_size`` is independent of this knob and keeps its own win either way. + + Output is unchanged by the cutoff: ``in0_block_w`` and the K-contraction order are independent of the + M-tiling, so only ``per_core_M`` changes. Decode never reads ``prefill_len_cutoff``. + + Decode W1→DRAM before W3 avoids L1 overlap between W1 activations and W3 matmul CBs + (N300 batch 32). + """ + t = _Qwen2WHTuning() + model_slug = hf_model_id.split("/")[-1] + if not (model_slug.startswith("Qwen2-7B") and num_dev in (1, 2)): + return t + + t.mlp_prefill_len_cutoff = 1024 + t.mlp_ff_compute_kernel_cfg = _qwen_wh_mlp_matmul_compute_kernel() + t.mlp_decode_spill_w1_to_dram = max_batch_size >= 16 + t.lm_head_compute_kernel_cfg = _qwen_wh_mlp_matmul_compute_kernel() + t.mlp_decode_ff_compute_kernel_cfg = ( + _qwen_wh_mlp_decode_matmul_compute_kernel() if perf_decode_tuning else t.mlp_ff_compute_kernel_cfg + ) + + # Attention prefill always uses HiFi4 + fp32 dest acc on Qwen (matches TTTv1). + attn_hifi4 = _qwen_wh_attn_hifi4_kernel() + t.attn_li_qkv_prefill_kernel_cfg = attn_hifi4 + t.attn_li_o_prefill_kernel_cfg = attn_hifi4 + # Attention decode defaults to HiFi4 + fp32 dest acc; LoFi only under perf_decode_tuning. + t.attn_li_qkv_decode_kernel_cfg = attn_hifi4 + t.attn_sdpa_decode_kernel_cfg = attn_hifi4 + t.attn_li_o_decode_kernel_cfg = attn_hifi4 + + if perf_decode_tuning: + lo = _qwen_wh_decode_attn_lofi_kernel() + t.perf_decode_sdpa_cfg = ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=(8, 8), + exp_approx_mode=True, + q_chunk_size=0, + k_chunk_size=0, + ) + t.attn_li_qkv_decode_kernel_cfg = lo + t.attn_sdpa_decode_kernel_cfg = lo + t.attn_li_o_decode_kernel_cfg = lo + logger.info( + f"MLP/LM/attention tuning for {hf_model_id} on {num_dev} device(s): " + f"prefill_len_cutoff={t.mlp_prefill_len_cutoff}, FF prefill HiFi4, attn prefill+decode HiFi4+fp32, " + f"decode spill W1→DRAM={t.mlp_decode_spill_w1_to_dram}, " + f"perf_decode_tuning={perf_decode_tuning}" + ) + return t + + +@dataclass +class Qwen2_7BDecoderLayerConfig: + attention_norm_config: RMSNorm1DConfig + attention_config: Attention1DConfig + ff_norm_config: RMSNorm1DConfig + mlp_config: MLP1DConfig + + +def _build_decoder_layer( + *, + idx: int, + weights: Qwen2_7BLayerWeights, + qcfg: Qwen2_7BModelParameters, + mesh_device: ttnn.MeshDevice, + tt_ccl: Any, + topology: Any, + num_dev: int, + precision: Qwen2_7BPrecisionConfig, + paged_cfg: Qwen2PagedAttentionConfig, + cache_path: Path | None, + wh: _Qwen2WHTuning, +) -> Qwen2_7BDecoderLayerConfig: + """Build one decoder-layer config from provider-converted tensors.""" + prefix = f"layer{idx}" + + lazy_wqkv = _lazy( + weights.wqkv, + dtype=precision.wqkv_dtype, + cache=(cache_path / "attn", f"{prefix}_wqkv") if cache_path else None, + ) + lazy_wo = _lazy( + weights.wo, + dtype=precision.wo_dtype, + cache=(cache_path / "attn", f"{prefix}_wo") if cache_path else None, + ) + + bias_lw = LazyWeight( + source=weights.wqkv_bias.to(torch.bfloat16), + dtype=ttnn.bfloat16, + cache_dir_weight_name=(cache_path / "attn", f"{prefix}_bias") if cache_path else None, + ) + + attention_config = Attention1DConfig( + wqkv=lazy_wqkv, + wo=lazy_wo, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + n_heads=qcfg.n_heads, + n_kv_heads=qcfg.n_kv_heads, + head_dim=qcfg.head_dim, + max_batch_size=qcfg.max_batch_size, + max_seq_len=qcfg.max_seq_len, + q_norm_config=None, + k_norm_config=None, + wqkv_bias=bias_lw, + use_vllm_paged_kv_cache=True, + paged_attention_config=paged_cfg, + kv_cache=None, + kv_cache_dtype=precision.kv_cache_dtype, + decode_sdpa_prg_config=wh.perf_decode_sdpa_cfg, + li_qkv_decode_compute_kernel_cfg=wh.attn_li_qkv_decode_kernel_cfg, + sdpa_decode_compute_kernel_cfg=wh.attn_sdpa_decode_kernel_cfg, + li_o_decode_compute_kernel_cfg=wh.attn_li_o_decode_kernel_cfg, + li_qkv_prefill_compute_kernel_cfg=wh.attn_li_qkv_prefill_kernel_cfg, + li_o_prefill_compute_kernel_cfg=wh.attn_li_o_prefill_kernel_cfg, + ) + + mlp_config = MLP1DConfig( + w1=_lazy( + weights.w1, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w1") if cache_path else None, + ), + w2=_lazy( + weights.w2, + dtype=precision.mlp_w2_dtype, + cache=(cache_path / "mlp", f"{prefix}_w2") if cache_path else None, + ), + w3=_lazy( + weights.w3, + dtype=precision.mlp_w1_w3_dtype, + cache=(cache_path / "mlp", f"{prefix}_w3") if cache_path else None, + ), + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + max_batch_size=qcfg.max_batch_size, + prefill_len_cutoff=wh.mlp_prefill_len_cutoff, + ff1_3_compute_kernel_cfg=wh.mlp_ff_compute_kernel_cfg, + ff2_compute_kernel_cfg=wh.mlp_ff_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=wh.mlp_decode_ff_compute_kernel_cfg, + decode_ff2_compute_kernel_cfg=wh.mlp_decode_ff_compute_kernel_cfg, + decode_spill_w1_to_dram_before_w3=wh.mlp_decode_spill_w1_to_dram, + ) + + post_attn_decode_program_config, post_attn_decode_memory_config = _post_attn_norm_decode_configs( + dim=qcfg.dim, + hidden_dim=qcfg.hidden_dim, + num_devices=num_dev, + max_batch_size=qcfg.max_batch_size, + ) + + def _build_norm(weight: torch.Tensor, name: str, **extra: Any) -> RMSNorm1DConfig: + lw = _lazy( + weight, + dtype=ttnn.bfloat16, + cache=(cache_path / "norm", f"{prefix}_{name}") if cache_path else None, + ) + return RMSNorm1DConfig( + weight=lw, + mesh_device=mesh_device, + eps=qcfg.rms_norm_eps, + max_batch_size=qcfg.max_batch_size, + tt_ccl=tt_ccl, + **extra, + ) + + return Qwen2_7BDecoderLayerConfig( + attention_norm_config=_build_norm(weights.attention_norm, "pre_attn"), + attention_config=attention_config, + ff_norm_config=_build_norm( + weights.ff_norm, + "post_attn", + decode_program_config=post_attn_decode_program_config, + decode_memory_config=post_attn_decode_memory_config, + ), + mlp_config=mlp_config, + ) + + +def _build_lm_head( + *, + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + qcfg: Qwen2_7BModelParameters, + lm_head_dtype: ttnn.DataType, + lm_head_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None, + cache_path: Path | None, +) -> LMHead1DConfig: + """Build the vocab-sharded LM head with DRAM-matmul program configs. + + LM head DRAM matmul is sized for decode batch tiles (``max_batch_size``). Prefill logits + use a single 32-row tile via ``post_process_prefill_output`` / :func:`_slice_last_token_tile`. + """ + lm_splits, lm_split_sizes, lm_weights_memcfgs = _build_lm_head_lazy_weights( + mesh_device, + lm_head_weight, + dim=qcfg.dim, + vocab_size=qcfg.vocab_size, + dtype=lm_head_dtype, + cache_dir=cache_path / "lm_head" if cache_path else None, + ) + lm_head_core_grid = _dram_shard_core_grid(qcfg.dim) + tile = ttnn.TILE_SIZE + tile_padded_batch_rows = tile * math.ceil(qcfg.max_batch_size / tile) + lm_prog_configs = [ + _dram_matmul_config(tile_padded_batch_rows, qcfg.dim, ss, lm_head_core_grid.num_cores) for ss in lm_split_sizes + ] + lm_input_memcfg = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, _nearest_32(qcfg.dim // lm_head_core_grid.num_cores)), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return LMHead1DConfig( + output_weights=lm_splits, + mesh_device=mesh_device, + dim=qcfg.dim, + max_batch_size=qcfg.max_batch_size, + lm_head_dtype=lm_head_dtype, + program_configs=lm_prog_configs, + compute_kernel_config=lm_head_compute_kernel_cfg, + input_memcfg=lm_input_memcfg, + weights_memcfgs=lm_weights_memcfgs, + ) + + +def _build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """Column-split a provider-neutral LM-head tensor for ``LMHead1D``.""" + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / TILE_SIZE) * TILE_SIZE + if vocab_size < padded_vocab_size: + torch_w = torch.cat( + [torch_w, torch.zeros(torch_w.shape[0], padded_vocab_size - vocab_size, dtype=torch_w.dtype)], dim=-1 + ) + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (TILE_SIZE * dram_size.x)) * (TILE_SIZE * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights = [] + weights_memcfgs = [] + for index, split_size in enumerate(split_sizes): + device_splits = [] + for device_index in range(num_devices): + start = device_index * size_per_device + sum(split_sizes[:index]) + device_splits.append(torch_w[:, start : start + split_size]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, f"lm_head_split_{index}_{combined.shape[-1]}") if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs + + +def build_qwen2_7b_transformer_config( + *, + mesh_device: ttnn.MeshDevice, + params: Qwen2_7BModelParameters, + weights: Qwen2_7BWeights, + n_layers: int, + precision: Qwen2_7BPrecisionConfig, + cache_path: Path, + paged_attention_config: Qwen2PagedAttentionConfig, + model_name: str = "Qwen2-7B-Instruct", +) -> Qwen2_7BTransformerConfig: + """Build the TT graph from provider-neutral dimensions and converted tensors.""" + num_devices = mesh_device.get_num_devices() + if num_devices != 2: + raise ValueError(f"Qwen2-7B supports TP2/N300 only, got {num_devices} devices") + if params.n_heads % num_devices or params.n_kv_heads % num_devices: + raise ValueError( + f"Checkpoint heads ({params.n_heads}/{params.n_kv_heads}) " + f"must be divisible by device count ({num_devices})" + ) + if len(weights.layers) != n_layers: + raise ValueError(f"Expected {n_layers} decoder layer weight sets, got {len(weights.layers)}") + + tt_ccl = get_tt_ccl(mesh_device) + topology = default_topology(mesh_device) + wh = _resolve_qwen_wh_tuning( + hf_model_id=model_name, + num_dev=num_devices, + max_batch_size=params.max_batch_size, + perf_decode_tuning=precision.perf_decode_tuning, + ) + embedding_config = Embedding1DConfig( + weights=_lazy(weights.embedding, dtype=ttnn.bfloat16, cache=(cache_path / "embedding", "tok_embeddings")), + mesh_device=mesh_device, + embed_scale=1.0, + ) + rope_config = Rope1DConfig( + cos_matrix=_lazy(weights.rope_cos, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "cos")), + sin_matrix=_lazy(weights.rope_sin, dtype=ttnn.bfloat16, cache=(cache_path / "rope", "sin")), + max_batch_size=params.max_batch_size, + head_dim=params.head_dim, + device=mesh_device, + use_qk_fused=False, + ) + block_configs = [ + _build_decoder_layer( + idx=index, + weights=weights.layers[index], + qcfg=params, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=topology, + num_dev=num_devices, + precision=precision, + paged_cfg=paged_attention_config, + cache_path=cache_path, + wh=wh, + ) + for index in range(n_layers) + ] + norm_config = RMSNorm1DConfig( + weight=_lazy(weights.final_norm, dtype=ttnn.bfloat16, cache=(cache_path / "norm", "final")), + mesh_device=mesh_device, + eps=params.rms_norm_eps, + max_batch_size=params.max_batch_size, + tt_ccl=tt_ccl, + ) + lm_head_config = _build_lm_head( + mesh_device=mesh_device, + lm_head_weight=weights.lm_head, + qcfg=params, + lm_head_dtype=precision.lm_head_dtype, + lm_head_compute_kernel_cfg=wh.lm_head_compute_kernel_cfg, + cache_path=cache_path, + ) + sampling_config = Sampling1DConfig( + vocab_size=params.vocab_size, + valid_vocab_size=params.vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=_nearest_32(params.max_batch_size), + allow_force_argmax=False, + pad_to_power_of_2=True, + ) + return Qwen2_7BTransformerConfig( + n_layers=n_layers, + vocab_size=params.vocab_size, + max_batch_size=params.max_batch_size, + max_seq_len=params.max_seq_len, + dim=params.dim, + num_devices=num_devices, + mesh_device=mesh_device, + embedding_config=embedding_config, + rope_config=rope_config, + block_configs=block_configs, + norm_config=norm_config, + lm_head_config=lm_head_config, + sampling_config=sampling_config, + decode_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + prefill_residual_memcfg=ttnn.DRAM_MEMORY_CONFIG, + activation_dtypes=[None] * n_layers, + tt_ccl=tt_ccl, + cache_path=str(cache_path), + ) + + +class Qwen2_7BDecoderLayer(LightweightModule): + def __init__( + self, + *, + input_layernorm: RMSNorm1D, + self_attn: Attention1D, + post_attention_layernorm: RMSNorm1D, + mlp: MLP1D, + ): + super().__init__() + self.input_layernorm = input_layernorm + self.self_attn = self_attn + self.post_attention_layernorm = post_attention_layernorm + self.mlp = mlp + self.attention_norm = input_layernorm + self.attention = self_attn + self.ff_norm = post_attention_layernorm + self.feed_forward = mlp + + @classmethod + def from_config(cls, config: Qwen2_7BDecoderLayerConfig) -> Qwen2_7BDecoderLayer: + return cls( + input_layernorm=RMSNorm1D.from_config(config.attention_norm_config), + self_attn=Attention1D.from_config(config.attention_config), + post_attention_layernorm=RMSNorm1D.from_config(config.ff_norm_config), + mlp=MLP1D.from_config(config.mlp_config), + ) + + def prefill_forward( + self, + x: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + *, + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # For batched prefill (batch_size > 1) x is the folded [1,1,B*S,dim] hidden state; norm, + # residual add and MLP are row-independent so they treat B*S as one long sequence unchanged. + # Only attention unfolds the batch axis internally (see Attention1D.prefill_forward). + # Match Llama ``TransformerBlock1D``: fractured embed / norm activations must be + # all-gathered to full ``dim`` before Attention1D / MLP1D (QKV matmul expects width ``dim``). + r = self.input_layernorm.prefill_forward(x) + r = _all_gather_rmsnorm_tensor(self.input_layernorm, r) + r = self.self_attn.prefill_forward( + r, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + r2 = self.post_attention_layernorm.prefill_forward(h) + r2 = _all_gather_rmsnorm_tensor(self.post_attention_layernorm, r2) + r2 = self.mlp.prefill_forward(r2) + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def decode_forward( + self, + x: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + xa = _all_gather_rmsnorm_tensor( + self.input_layernorm, x, memory_config=self.input_layernorm.config.decode_memory_config + ) + r = self.input_layernorm.forward(xa, "decode") + r = self.self_attn.forward(r, current_pos, rot_mats, mode="decode", page_table=page_table) + h = ttnn.add(x, r, memory_config=ttnn.DRAM_MEMORY_CONFIG) + hf = _all_gather_rmsnorm_tensor( + self.post_attention_layernorm, h, memory_config=self.post_attention_layernorm.config.decode_memory_config + ) + r2 = self.post_attention_layernorm.forward(hf, "decode") + r2 = self.mlp.forward(r2, "decode") + return ttnn.add(h, r2, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + +class Qwen2_7B(LightweightModule): + """ + Full decoder for Qwen2-7B-Instruct (TTTv2 modules only). + + Prefill/decode on **embedded** activations match ``EagerLLMExecutor``. Token embedding + is ``embed_prefill`` / ``embed_decode``. Bind KV with ``set_kv_cache`` before first forward. + """ + + decode_residual_memcfg = ttnn.DRAM_MEMORY_CONFIG + + def __init__(self, config: Qwen2_7BTransformerConfig): + from tqdm import tqdm + + super().__init__() + self.config = config + self.embed = Embedding1D.from_config(config.embedding_config) + self.embedding = self.embed + self.rope_setup = RotarySetup1D.from_config(config.rope_config) + self.layers = [ + Qwen2_7BDecoderLayer.from_config(config.block_configs[index]) + for index in tqdm(range(config.n_layers), desc="Building layers") + ] + self.norm = RMSNorm1D.from_config(config.norm_config) + self.lm_head = LMHead1D.from_config(config.lm_head_config) + self.sampling = Sampling1D.from_config(config.sampling_config) if config.sampling_config is not None else None + self.supports_on_device_sampling = self.sampling is not None + self.mesh_device = config.mesh_device + self.tt_ccl = config.tt_ccl or get_tt_ccl(config.mesh_device) + self.vocab_size = config.vocab_size + self.n_layers = config.n_layers + self.num_devices = config.num_devices + self.decode_residual_memcfg = config.decode_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.prefill_residual_memcfg = config.prefill_residual_memcfg or ttnn.DRAM_MEMORY_CONFIG + self.activation_dtypes = config.activation_dtypes or [None] * config.n_layers + self.model_args = None + + @property + def n_kv_heads(self) -> int: + return self.config.block_configs[0].attention_config.n_kv_heads + + def iter_executor_named_modules(self): + if not hasattr(self, "layers"): + return + for index, layer in enumerate(self.layers): + for suffix, submodule in ( + ("attn_norm", getattr(layer, "attention_norm", None)), + ("attention", getattr(layer, "attention", None)), + ("ff_norm", getattr(layer, "ff_norm", None)), + ("mlp", getattr(layer, "feed_forward", None)), + ): + if submodule is not None: + yield f"layer[{index}].{suffix}", submodule + if hasattr(self, "norm"): + yield "final_norm", self.norm + if hasattr(self, "lm_head"): + yield "lm_head", self.lm_head + + def configure_paged_attention(self, *, block_size: int, max_num_blocks: int) -> None: + for name, value in (("block_size", block_size), ("max_num_blocks", max_num_blocks)): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + live_configs = tuple(layer.attention.config for layer in self.layers) + for layer_index, config in enumerate(live_configs): + if not config.use_vllm_paged_kv_cache or config.paged_attention_config is None: + raise RuntimeError(f"Model layer {layer_index} is not configured for externally managed paged KV cache") + if config.kv_cache is not None or getattr(self.layers[layer_index].attention, "kv_cache", None) is not None: + raise RuntimeError(f"Model layer {layer_index} already has a bound KV cache") + construction_configs = tuple(block.attention_config for block in self.config.block_configs) + for config in tuple({id(item): item for item in (*construction_configs, *live_configs)}.values()): + config.paged_attention_config = replace( + config.paged_attention_config, + block_size=block_size, + max_num_blocks=max_num_blocks, + ) + + def set_kv_cache(self, kv_cache: list | None) -> None: + if kv_cache is None: + for layer in self.layers: + layer.attention.config.kv_cache = None + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = None + return + if len(kv_cache) != len(self.layers): + raise ValueError(f"kv_cache has {len(kv_cache)} entries but model has {len(self.layers)} layers") + cache_pairs = [] + for index, value in enumerate(kv_cache): + try: + pair = tuple(value) + except TypeError as error: + raise TypeError(f"kv_cache layer {index} must provide an iterable K/V tensor pair") from error + if len(pair) != 2: + raise ValueError(f"kv_cache layer {index} must contain exactly two K/V tensors") + cache_pairs.append(pair) + for layer, pair in zip(self.layers, cache_pairs): + layer.attention.config.kv_cache = pair + if hasattr(layer.attention, "kv_cache"): + layer.attention.kv_cache = pair + + def prepare_prefill_rot_mats(self, position_indices: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]: + self.rope_setup.load_device_weights() + cos = None + sin = None + try: + cos = ttnn.embedding(position_indices, self.rope_setup.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(position_indices, self.rope_setup.sin_matrix, layout=ttnn.TILE_LAYOUT) + return ttnn.unsqueeze_to_4D(cos), ttnn.unsqueeze_to_4D(sin) + except BaseException: + for tensor in (sin, cos): + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException: + pass + raise + + def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + x = ttnn.unsqueeze_to_4D(x) + return ttnn.to_memory_config(x, self.decode_residual_memcfg) + + def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor: + x = self.embed.forward(tokens) + return ttnn.unsqueeze_to_4D(x) + + def prefill_forward( + self, + x_embed: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + get_last_token: int = -1, + batch_size: int = 1, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + # batch_size > 1: x_embed is the folded [1,1,B*S,dim] tensor (B users). The batched path always + # returns the full hidden state when no runtime last-token slice is supplied; the executor does + # per-slot last-token extraction + norm/lm_head so those stages stay bit-identical to the single-user path. + x = x_embed + for layer in self.layers: + x = layer.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + chunk_start_idx_tensor=chunk_start_idx_tensor, + ) + + if get_last_token == -1 and last_token_slice is None: + return x + + # Slice + deallocate the full-sequence buffer before norm/LM head reduces peak L1. + old = x + if last_token_slice is None: + x_tile = _slice_last_token_tile(old, get_last_token) + else: + x_tile = ttnn.slice( + old, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(old.shape[2]) // 32, + ) + ttnn.deallocate(old) + if last_token_index is not None: + if x_tile.dtype != ttnn.bfloat16: + old_tile = x_tile + x_tile = ttnn.typecast(x_tile, ttnn.bfloat16) + ttnn.deallocate(old_tile) + old_tile = x_tile + x_tile = ttnn.embedding(last_token_index, x_tile, layout=ttnn.TILE_LAYOUT) + x_tile = ttnn.unsqueeze_to_4D(x_tile) + ttnn.deallocate(old_tile) + return self._last_tile_logits(x_tile) + + def post_process_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + if last_token_slice is None: + x = _slice_last_token_tile(hidden_states, last_token_idx) + else: + x = ttnn.slice( + hidden_states, + last_token_slice[0], + last_token_slice[1], + slice_dim=2, + num_devices=int(hidden_states.shape[2]) // 32, + ) + if last_token_index is not None: + if x.dtype != ttnn.bfloat16: + old = x + x = ttnn.typecast(x, ttnn.bfloat16) + ttnn.deallocate(old) + old = x + x = ttnn.embedding(last_token_index, x, layout=ttnn.TILE_LAYOUT) + x = ttnn.unsqueeze_to_4D(x) + ttnn.deallocate(old) + return self._last_tile_logits(x) + + def _last_tile_logits(self, x_tile: ttnn.Tensor) -> ttnn.Tensor: + """Final-norm + all-gather + LM-head on a 32-row tile. ``x_tile`` shape ``[1, 1, 32, dim]``.""" + x = self.norm.prefill_forward(x_tile) + x = _all_gather_rmsnorm_tensor(self.norm, x) + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + x = self.lm_head.forward(x) + return ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) + + def post_process_batched_prefill_output( + self, + hidden_states: ttnn.Tensor, + last_token_idx_list: list[int], + padded_batch: int, + prefill_seq_len: int, + last_token_slice: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, + last_token_index: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + del last_token_slice, last_token_index + fold_len = padded_batch * prefill_seq_len + selector = torch.zeros(1, 1, 32, fold_len, dtype=torch.bfloat16) + for local_row, last_token_idx in enumerate(last_token_idx_list): + selector[0, 0, local_row, local_row * prefill_seq_len + last_token_idx] = 1.0 + selector = ttnn.from_torch( + selector, + device=self.mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), + ) + x = ttnn.matmul( + selector, + hidden_states, + dtype=ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + compute_kernel_config=ttnn.init_device_compute_kernel_config( + self.mesh_device.arch(), + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ), + ) + ttnn.deallocate(selector) + return self._last_tile_logits(x) + + def decode_forward( + self, + x_embed: ttnn.Tensor, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + x = x_embed + for layer in self.layers: + x = layer.decode_forward(x, current_pos, rot_mats, page_table=page_table) + x = _all_gather_rmsnorm_tensor(self.norm, x, memory_config=self.norm.config.decode_memory_config) + x = self.norm.decode_forward(x) + return self.lm_head.forward(x) + + def gather_and_untilize_logits(self, logits: ttnn.Tensor) -> ttnn.Tensor: + if self.num_devices > 1 and self.tt_ccl is not None: + logits = ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + memory_config=logits.memory_config(), + topology=default_topology(self.mesh_device), + barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=24, + num_workers_per_link=4, + num_buffers_per_channel=2, + ) + return ttnn.untilize(logits, use_multicore=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + def increment_positions(self, current_pos: ttnn.Tensor, rot_mat_idxs: ttnn.Tensor) -> None: + ttnn.plus_one(current_pos, skip_negative_entries=True) + ttnn.plus_one(rot_mat_idxs) + + def prefill_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, start_pos: int = 0, user_id: int = 0) -> ttnn.Tensor: + """Legacy path: embed + RoPE + blocks + final norm (no page table). For tests only.""" + x = self.embed_prefill(token_ids_tt) + seq_len = x.shape[2] + assert seq_len % 128 == 0, "prefill seq_len must be divisible by 128" + rot = self.rope_setup.prefill_forward(start_pos, seq_len) + h = x + for layer in self.layers: + h = layer.prefill_forward(h, rot, user_id=user_id, page_table=None) + h = self.norm.prefill_forward(h) + return _all_gather_rmsnorm_tensor(self.norm, h) + + def decode_from_token_ids(self, token_ids_tt: ttnn.Tensor, *, current_pos: int) -> ttnn.Tensor: + """Legacy path: single-token decode without paged ``page_table``.""" + x = self.embed.forward(token_ids_tt) + x = ttnn.unsqueeze_to_4D(x) + pos = torch.tensor([current_pos], dtype=torch.long) + rot_idxs = prepare_rot_idxs(self.rope_setup.config, pos, on_host=False) + rot = self.rope_setup.decode_forward(rot_idxs) + cur = ttnn.from_torch( + pos, + device=self.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(self.mesh_device), + ) + h = x + for layer in self.layers: + h = layer.decode_forward(h, cur, rot, page_table=None) + h = _all_gather_rmsnorm_tensor(self.norm, h, memory_config=self.norm.config.decode_memory_config) + return self.norm.forward(h, "decode") + + def lm_logits(self, hidden: ttnn.Tensor) -> ttnn.Tensor: + """Project last hidden to logits (vocab-sharded on multi-device). + + Skip the explicit interleaved→shard if the caller already produced a sharded + input (``decode_from_token_ids`` returns the decode-mode RMSNorm's width-sharded + output, which already matches ``LMHead1D.config.input_memcfg``). + """ + x = hidden + lm_head_memcfg = self.lm_head.config.input_memcfg + if lm_head_memcfg is not None and lm_head_memcfg.is_sharded() and not x.memory_config().is_sharded(): + x = ttnn.interleaved_to_sharded(x, lm_head_memcfg) + return self.lm_head.forward(x) + + +__all__ = [ + "QWEN2_7B_ACCURACY", + "QWEN2_7B_PERFORMANCE", + "Qwen2PagedAttentionConfig", + "Qwen2_7B", + "Qwen2_7BDecoderLayer", + "Qwen2_7BDecoderLayerConfig", + "Qwen2_7BLayerWeights", + "Qwen2_7BModelParameters", + "Qwen2_7BPrecisionConfig", + "Qwen2_7BTransformerConfig", + "Qwen2_7BWeights", + "build_qwen2_7b_transformer_config", +] diff --git a/code/models/common/models/qwen2_7b/weight_utils.py b/code/models/common/models/qwen2_7b/weight_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..62ecf0d9cfd14583e96e3b4e946caa3615693395 --- /dev/null +++ b/code/models/common/models/qwen2_7b/weight_utils.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Host-side HF weight layout helpers for Qwen2.x (TTTv2 port). + +Q/K linear weights use HuggingFace head layout; TTNN RoPE expects Meta layout. +This duplicates the small permute helpers from ``load_checkpoints.reverse_permute`` +without importing ``models/tt_transformers``. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +import ttnn +from models.common.modules.lazy_weight import LazyWeight + + +def reverse_permute(tensor: torch.Tensor, n_heads: int, dim1: int, dim2: int) -> torch.Tensor: + """HF Q/K weight rows → Meta layout (RoPE-compatible).""" + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +def reverse_permute_1d(tensor: torch.Tensor) -> torch.Tensor: + """Undo HF split of last dim into interleaved complex pairs for norm/bias vectors.""" + shape = tensor.shape + dim = shape[-1] + assert dim % 2 == 0, "Last dimension must be even" + reals = tensor[..., : dim // 2] + imags = tensor[..., dim // 2 :] + interleaved = torch.stack((reals, imags), dim=-1).flatten(start_dim=len(shape) - 1) + return interleaved + + +def permute_hf_rope_to_meta_tables(cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert HF cos/sin tables to Meta-style tables for ``RotarySetup1D`` (see attention tests).""" + if len(cos.shape) == 3: + cos = cos.squeeze(0) + sin = sin.squeeze(0) + cos = cos[:, : cos.shape[1] // 2] + cos = torch.stack((cos, cos), dim=-1).flatten(-2) + sin = sin[:, : sin.shape[1] // 2] + sin = torch.stack((sin, sin), dim=-1).flatten(-2) + cos = cos.unsqueeze(0).unsqueeze(0) + sin = sin.unsqueeze(0).unsqueeze(0) + return cos, sin + + +def build_rope_cos_sin_torch( + rotary_emb: Any, table_len: int, head_dim: int, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Build ``[1, 1, table_len, head_dim]`` cos/sin tensors (Meta layout) from HF rotary module.""" + x = torch.zeros(1, 1, table_len, head_dim, dtype=dtype) + position_ids = torch.arange(table_len, dtype=torch.long).unsqueeze(0) + with torch.no_grad(): + cos_hf, sin_hf = rotary_emb(x, position_ids) + cos_m, sin_m = permute_hf_rope_to_meta_tables(cos_hf.float(), sin_hf.float()) + return cos_m.to(dtype), sin_m.to(dtype) + + +def attention_wqkv_wo_from_hf_layer( + self_attn: Any, + num_devices: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Mirror ``get_attention_weights_from_ref_model`` in ``test_attention_1d.py`` (HF module API).""" + wq_raw = self_attn.q_proj.weight + wk_raw = self_attn.k_proj.weight + wv_raw = self_attn.v_proj.weight + wo_raw = self_attn.o_proj.weight + + dim = wq_raw.shape[1] + cfg = self_attn.config + n_heads = cfg.num_attention_heads + n_kv_heads = cfg.num_key_value_heads + head_dim = cfg.hidden_size // n_heads + n_heads_times_head_dim = n_heads * head_dim + n_kv_heads_times_head_dim = n_kv_heads * head_dim + + wq_meta = reverse_permute(wq_raw, n_heads, n_heads_times_head_dim, dim) + wk_meta = reverse_permute(wk_raw, n_kv_heads, n_kv_heads_times_head_dim, dim) + wv_meta = wv_raw + wo_meta = wo_raw + + wq = wq_meta.T + wk = wk_meta.T + wv = wv_meta.T + wo = wo_meta.T + + qkv_list = [] + for i in range(num_devices): + wq_chunk = torch.chunk(wq, num_devices, dim=1)[i] + wk_chunk = torch.chunk(wk, num_devices, dim=1)[i] + wv_chunk = torch.chunk(wv, num_devices, dim=1)[i] + qkv = torch.cat([wq_chunk, wk_chunk, wv_chunk], dim=-1) + qkv_list.append(qkv) + wqkv = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0).clone() + wo = wo.unsqueeze(0).unsqueeze(0).clone() + + q_norm = None + k_norm = None + if getattr(self_attn, "q_norm", None) is not None: + q_norm = reverse_permute_1d(self_attn.q_norm.weight).clone() + if getattr(self_attn, "k_norm", None) is not None: + k_norm = reverse_permute_1d(self_attn.k_norm.weight).clone() + + wqkv_bias = None + if getattr(self_attn.q_proj, "bias", None) is not None: + bq_raw = self_attn.q_proj.bias + bk_raw = self_attn.k_proj.bias + bv_raw = self_attn.v_proj.bias + bq_meta = reverse_permute_1d(bq_raw.view(n_heads, head_dim)).view(-1) + bk_meta = reverse_permute_1d(bk_raw.view(n_kv_heads, head_dim)).view(-1) + bv_meta = bv_raw + qkv_bias_list = [] + for i in range(num_devices): + bq_chunk = torch.chunk(bq_meta, num_devices, dim=0)[i] + bk_chunk = torch.chunk(bk_meta, num_devices, dim=0)[i] + bv_chunk = torch.chunk(bv_meta, num_devices, dim=0)[i] + qkv_bias_list.append(torch.cat([bq_chunk, bk_chunk, bv_chunk], dim=-1)) + wqkv_bias = torch.cat(qkv_bias_list, dim=-1).clone() + + return wqkv, wo, q_norm, k_norm, wqkv_bias + + +def mlp_weights_from_hf_layer(mlp: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (w1, w2, w3) TTNN layouts: gate^T, down^T, up^T for ``MLP1D``.""" + w1 = mlp.gate_proj.weight.T.contiguous().clone() + w3 = mlp.up_proj.weight.T.contiguous().clone() + w2 = mlp.down_proj.weight.T.contiguous().clone() + return w1, w2, w3 + + +def rms_weight_torch(layernorm: Any) -> torch.Tensor: + return layernorm.weight.detach().float().clone() + + +def embed_tokens_torch(embed: Any) -> torch.Tensor: + w = embed.weight.detach().float().clone() + return w.unsqueeze(0).unsqueeze(0) + + +def build_lm_head_lazy_weights( + mesh_device: ttnn.MeshDevice, + lm_head_weight: torch.Tensor, + *, + dim: int, + vocab_size: int, + max_columns_per_device: int = 8192, + dtype: ttnn.DataType = ttnn.bfloat8_b, + cache_dir: Path | None = None, +) -> tuple[list[LazyWeight], list[int], list[ttnn.MemoryConfig]]: + """ + Column-split LM head for ``LMHead1D`` (adapted from ``LMHead1D.from_model_args`` math, no v1 args). + ``lm_head_weight`` shape: ``[vocab_size, dim]`` (HF). + + Returns ``(lazy_weights, split_sizes, weights_memcfgs)`` for ``LMHead1DConfig`` DRAM-sharded matmuls. + """ + num_devices = mesh_device.get_num_devices() + torch_w = lm_head_weight.T.contiguous().to(torch.bfloat16) + padded_vocab_size = math.ceil(vocab_size / 32) * 32 + if vocab_size < padded_vocab_size: + pad = padded_vocab_size - vocab_size + torch_w = torch.cat([torch_w, torch.zeros(torch_w.shape[0], pad, dtype=torch_w.dtype)], dim=-1) + + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + tile = ttnn.TILE_SIZE + + def dram_sharded_memcfg(k: int, n: int) -> ttnn.MemoryConfig: + padded_n = math.ceil(n / (tile * dram_size.x)) * (tile * dram_size.x) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_n // dram_size.x), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + output_weights: list[LazyWeight] = [] + weights_memcfgs: list[ttnn.MemoryConfig] = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_splits.append(torch_w[:, start:end]) + combined = torch.cat(device_splits, dim=-1) + mem_cfg = dram_sharded_memcfg(dim, math.ceil(combined.shape[-1] / num_devices)) + weights_memcfgs.append(mem_cfg) + name = f"lm_head_split_{i}_{combined.shape[-1]}" + output_weights.append( + LazyWeight( + source=combined, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=(cache_dir, name) if cache_dir else None, + ) + ) + return output_weights, split_sizes, weights_memcfgs diff --git a/code/models/common/modules/attention/attention_1d.py b/code/models/common/modules/attention/attention_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..0170f2cd382eb48267c26aa66f023b80648699a2 --- /dev/null +++ b/code/models/common/modules/attention/attention_1d.py @@ -0,0 +1,2541 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2-style Attention module for 1D-topology devices: N150 (1x1), N300 (1x2), T3K (1x8). + +Single unified Attention1D class with separate forward methods: + - decode_forward(): For decode mode (single token per user) + - prefill_forward(): For prefill mode (multiple tokens) + - forward(x, mode, **kwargs): Dispatcher that calls the appropriate method + +Execution paths: + Decode: QKV matmul → all_reduce → create_qkv_heads → rotary → KV cache update → SDPA → concat_heads → WO matmul → all_reduce + Prefill: [reshape] → QKV matmul → all_reduce → create_qkv_heads → rotary → KV cache fill → SDPA → concat_heads → [reshape] → all_gather → WO matmul → reduce_scatter + +Key design decisions: + - No static branching on topology in forward() - TG excluded at module level + - Paged vs non-paged KV cache selected at runtime based on page_table presence + - Fused all-gather matmul used for Ring topology (T3K 1x8) when dimensions allow + +Weight format requirements: + Q and K weights MUST be in Meta format, not HuggingFace format. HuggingFace stores + Q/K weights with a different head layout that is incompatible with TTNN's RoPE + implementation. Use `reverse_permute` from `models.tt_transformers.tt.load_checkpoints` + to convert HF weights to Meta format before passing to Attention1D: + + from models.tt_transformers.tt.load_checkpoints import reverse_permute + wq_meta = reverse_permute(wq_hf, n_heads, n_heads * head_dim, dim) + wk_meta = reverse_permute(wk_hf, n_kv_heads, n_kv_heads * head_dim, dim) + + The `from_model_args` factory handles this automatically via `convert_hf_to_meta`. + When using `from_config` with weights extracted directly from HuggingFace models, + you must apply this transformation manually. +""" + +import math +from dataclasses import dataclass, field, replace +from functools import lru_cache +from pathlib import Path +from typing import Callable, Optional + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight +from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig +from models.common.modules.tt_ccl import ( + CCL_CHUNKS_PER_SYNC, + CCL_NUM_BUFFERS_PER_CHANNEL, + CCL_NUM_WORKERS_PER_LINK, + TT_CCL, + default_topology, + get_tt_ccl, +) +from models.common.tensor_utils import ( + TILE_SIZE, + get_out_subblock_w, + get_rot_transformation_mat, + nearest_32, + zeros_like_kv_cache, + zeros_like_paged_cache, +) + +# ============================================================================= +# Constants +# ============================================================================= + +MAX_QKV_MM_SEQ_LEN = 2048 # Maximum sequence length for single QKV matmul + +# Maximum sequence length for WO matmul operation on Wormhole/Blackhole. +# Sequences longer than this are reshaped to [1, seq_len // MAX_MM_SEQ_LEN, MAX_MM_SEQ_LEN, dim] +# to fit on device and parallelize computation. After matmul, reshaped back. +# Source: TTTv1 model_config.py "MAX_MM_SEQ_LEN": 1024 +MAX_MM_SEQ_LEN = 1024 + +# ============================================================================= +# Attention1DConfig dataclass +# ============================================================================= + + +@dataclass +class Attention1DConfig: + """ + Central configuration for Attention1D - the single source of truth for all settings. + + Simple usage (all defaults): + config = Attention1DConfig(wqkv, wo) + + Override any field: + config = Attention1DConfig(wqkv, wo, n_heads=32, head_dim=128) + + Full customization: + config = Attention1DConfig( + wqkv, wo, + mesh_device=custom_device, + decode_xqkv_prg_config=my_program_config, + ... + ) + """ + + # Required: weights (LazyWeight) + # IMPORTANT: Q and K weights must be in Meta format (not HuggingFace format). + # HF weights require `reverse_permute` transformation for RoPE compatibility. + # See module docstring for details. + wqkv: LazyWeight # Combined QKV projection weight (Q/K in Meta format) + wo: LazyWeight # Output projection weight + + # Optional: Q/K normalization configs (e.g., for Qwen models) + # Composed sub-module pattern: RMSNorm1DConfig instead of raw weights + q_norm_config: RMSNorm1DConfig | None = None + k_norm_config: RMSNorm1DConfig | None = None + + # Optional: QKV bias (e.g., for Qwen models) + wqkv_bias: LazyWeight | None = None + + # Device and collectives + mesh_device: ttnn.MeshDevice | None = None + tt_ccl: TT_CCL | None = None + topology: Optional[ttnn.Topology] = None # None = auto-detect + num_reduce_scatter_links: int | None = None + num_all_gather_links: int | None = None + decode_agmm_num_links: int = 1 + decode_agmm_chunks_per_sync: int = CCL_CHUNKS_PER_SYNC + decode_agmm_num_workers_per_link: int = CCL_NUM_WORKERS_PER_LINK + + # Optional CCL dtype for the prefill output reduce-scatter. When the fused all-gather+WO path is + # not selected, attention runs a separate reduce-scatter on the bf16 WO output; reducing it in a + # smaller dtype (e.g. bfloat8_b) halves the cross-device payload — matching TTTv1's bfloat8_b + # ccl_dtype (model_config.py) and the MLP w2 reduce. None keeps the WO output dtype (no cast, + # default — byte-identical to leaving the reduce untouched). + prefill_reduce_ccl_dtype: ttnn.DataType | None = None + + # Optional CCL dtype for the prefill all-gather. When the fused all-gather+WO path IS selected, + # attention all-gathers the bf16 concat-heads output before the WO matmul; casting that input to a + # smaller dtype (e.g. bfloat8_b) halves the collective's cross-device payload. Model files that also + # own norm-reconstruction all-gathers may thread this same dtype through them. None keeps the input + # dtype (no cast, default — byte-identical). to_memory_config does not cast an already-DRAM tensor, + # so an explicit typecast is required (mirrors prefill_reduce_ccl_dtype). + prefill_ag_ccl_dtype: ttnn.DataType | None = None + + # Model dimensions (derived from weights if None) + dim: int | None = None + n_heads: int | None = None + n_kv_heads: int | None = None + head_dim: int | None = None + qkv_size: int | None = None # head_dim * (2 * n_kv_heads + n_heads) + + # Batch and sequence config + max_batch_size: int = 32 + max_seq_len: int = 128 * 1024 + + # Attention config + scale: float | None = None # Default: head_dim ** -0.5 + sliding_window: int | None = None # For sliding window attention + use_qk_fused: bool = False # Fused Q/K rotary embedding + + # KV cache config + # + # The KV cache is a static configuration, allocated once and reused for all forward calls. + # This reflects how inference engines like vLLM manage KV caches: + # - Cache shape is determined by static model config (num_kv_heads, head_dim, block_size) + # - The cache is allocated once during model initialization (allocate_vllm_kv_cache) + # - The same cache tensors are passed to every forward call + # + # kv_cache options: + # - None: Set externally before forward (e.g., via from_model_args or direct assignment) + # - tuple[LazyWeight, LazyWeight]: (keys, values) backed by cache files, resolved lazily + # - tuple[ttnn.Tensor, ttnn.Tensor]: Pre-allocated (keys, values) tensors (e.g., from vLLM) + # When enabled, KV cache allocation and capacity are owned by the caller. + use_vllm_paged_kv_cache: bool = False + kv_cache: "tuple[LazyWeight, LazyWeight] | tuple[ttnn.Tensor, ttnn.Tensor] | None" = None + paged_attention_config: "PagedAttentionConfig | None" = None # type: ignore + kv_cache_dtype: ttnn.DataType = ttnn.bfloat8_b + # Threshold for sharding KV cache during prefill to handle update_cache memory limitations. + min_kv_prefill_shard_seqlen: int | None = None + + # Weight dtypes (None = auto-resolved based on device arch / model config) + wqkv_dtype: ttnn.DataType | None = None # QKV projection weight dtype + wo_dtype: ttnn.DataType | None = None # Output projection weight dtype + activation_dtype: ttnn.DataType | None = None # Intermediate activation dtype (post-RoPE, etc.) + + # Weight memory configs (None = DRAM interleaved by default) + wqkv_memcfg: ttnn.MemoryConfig | None = None # QKV weight placement + wo_memcfg: ttnn.MemoryConfig | None = None # Output weight placement + + # Decode program configs (None = auto-derived in _resolve_attention1d_config) + decode_input_memcfg: ttnn.MemoryConfig | None = None # DRAM-sharded input for decode + decode_xqkv_prg_config: "ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig | None" = None # QKV matmul + decode_sdpa_prg_config: ttnn.SDPAProgramConfig | None = None # Scaled dot-product attention + decode_attn_output_prg_config: "ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig | None" = ( + None # WO matmul + ) + decode_residual_memcfg: ttnn.MemoryConfig | None = None # Residual add output placement + decode_create_qkv_head_memcfg: ttnn.MemoryConfig | None = None # QKV head split output + decode_scores_memcfg: Callable[[int], ttnn.MemoryConfig] | None = None # SDPA scores; f(batch_size) + + # Prefill program configs (Callable factories: seq_len → config) + prefill_input_memcfg: ttnn.MemoryConfig | None = None # DRAM interleaved input for prefill + prefill_xqkv_prg_config: Callable[ + [int], ttnn.MatmulMultiCoreReuseMultiCastProgramConfig + ] | None = None # f(seq_len) + prefill_sdpa_prg_config: Callable[[int, int | None], ttnn.SDPAProgramConfig] | None = None # f(seq_len, chunk_size) + prefill_wo_prg_config: Callable[[int], ttnn.MatmulMultiCoreReuseMultiCastProgramConfig] | None = None # f(seq_len) + prefill_kv_memcfg: Callable[[int], ttnn.MemoryConfig] | None = None # f(seq_len) for KV cache write + + # WO-matmul prefill M-chunk cutoff. The folded [1,1,seq_len,dim] prefill activation (seq_len = B*S + # when batched) is reshaped into [1, seq_len//cutoff, cutoff, dim] so the WO matmul tiles at + # per_core_M = ceil(cutoff / 256). None → MAX_MM_SEQ_LEN (1024 = TTTv1 model_config.py parity); + # resolved+written back in _resolve_attention1d_config so forward and the prg config agree. A + # per-model override (e.g. 2048) regroups the folded batch-32 prefill into fewer/larger M-chunks — + # bit-identical, since only per_core_M / the chunk count change (the K-contraction, in0_block_w, and + # n_dim are all independent of M-tiling). Deliberately DECOUPLED from the fused all_gather_matmul + # n_dim, which stays on MAX_MM_SEQ_LEN so the fused output width / sharding is untouched. + wo_prefill_len_cutoff: int | None = None + + # Optional: use ttnn.experimental.minimal_matmul (instead of ttnn.linear) for the QKV prefill + # matmul above seq_len > 128 — matches TTTv1's long-prefill path (attention.py L907-913), which is + # ~2x faster on the large folded-batch QKV matmul. Default OFF so untouched models / decode stay + # byte-identical; the caller opts in. The factory yields a ttnn.MinimalMatmulConfig keyed on the + # folded seq_len (sibling to prefill_xqkv_prg_config, NOT a replacement — both coexist). + prefill_qkv_minimal_matmul: bool = False + prefill_xqkv_minimal_matmul_config: Callable[[int], "ttnn.MinimalMatmulConfig"] | None = None # f(seq_len) + + # Optional: use ttnn.experimental.minimal_matmul (instead of ttnn.linear) for the WO prefill + # matmul above seq_len > 128 — the completion of PLAN_01's QKV+FF2-only minimal plumbing. On the + # folded batched prefill the WO projection is the least-efficient prefill matmul on ttnn.linear + # (a per-op profile measured it ~3x slower than the same-fidelity QKV minimal_matmul); switching + # it to minimal_matmul recovers most of that gap. Default OFF so untouched models / decode stay + # byte-identical; the caller opts in. The factory yields a ttnn.MinimalMatmulConfig keyed on the + # folded seq_len (sibling to prefill_wo_prg_config, NOT a replacement — both coexist). NOT + # byte-identical vs the ttnn.linear WO (minimal accumulation differs) — same numerical class as + # the already-shipped QKV/FF2 minimal path; gate on token-accuracy + eval-32, not a byte-compare. + prefill_wo_minimal_matmul: bool = False + prefill_wo_minimal_matmul_config: Callable[[int], "ttnn.MinimalMatmulConfig"] | None = None # f(seq_len) + + # Fused all-gather matmul (Ring topology only, decode path) + use_fused_all_gather_matmul: bool | None = None # None = auto-detect based on topology + dim + decode_all_gather_matmul_prg_config: "ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig | None" = ( + None # Fused AG+WO + ) + decode_all_gather_matmul_memcfg: ttnn.MemoryConfig | None = None # Output placement for fused AG+WO + + # Separate all-gather (non-Ring topology or non-fused decode path) + decode_gather_users_memcfg: ttnn.MemoryConfig | None = None # Output placement for standalone all_gather + + # Architecture-sensitive caller requests. ``None`` selects mesh/SKU + # defaults during construction; explicit compatible values win. + li_qkv_decode_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + sdpa_decode_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + li_o_decode_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + li_qkv_prefill_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + sdpa_prefill_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + li_o_prefill_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + prefill_qkv_grid: tuple[int, int] | None = None + dram_shard_grid_width: int | None = None + decode_create_qkv_head_grid: ttnn.CoreGrid | None = None + decode_transformation_core_grid: ttnn.CoreCoord | None = None + + # Transformation matrices for rotary embedding (static, computed once). + # These are NOT LazyWeights because they are: + # - Small fixed-size tensors (32x32), not large model weights + # - Computed programmatically (get_rot_transformation_mat), not loaded from checkpoints + transformation_mat_decode: ttnn.Tensor | None = None + transformation_mat_prefill: ttnn.Tensor | None = None + + # Internal: pre-computed bias LazyWeights (materialized in __init__) + _wqkv_bias_decode: list[LazyWeight] | None = field(default=None, repr=False) + _wqkv_bias_prefill: LazyWeight | None = field(default=None, repr=False) + + def use_minimal_qkv_matmul(self, seq_len: int) -> bool: + """Whether the QKV prefill matmul should use ttnn.experimental.minimal_matmul. + + Mirrors TTTv1 ``use_minimal_qkv_prefill_matmul`` (model_config.py:1659) — minimal_matmul + only above ``seq_len > 128`` (the folded-batch / long-prompt regime), and only when the + per-model opt-in is set. ``seq_len`` is the folded ``B*S`` length. + """ + return bool(self.prefill_qkv_minimal_matmul) and seq_len > 128 + + def use_minimal_wo_matmul(self, seq_len: int) -> bool: + """Whether the WO prefill matmul should use ttnn.experimental.minimal_matmul. + + Mirrors ``use_minimal_qkv_matmul`` — minimal_matmul only above ``seq_len > 128`` (the + folded-batch / long-prompt regime), and only when the per-model opt-in is set. ``seq_len`` + is the folded ``B*S`` length. + """ + return bool(self.prefill_wo_minimal_matmul) and seq_len > 128 + + def is_resolved(self) -> bool: + """Check if all required fields are resolved.""" + required = [ + "wqkv", + "wo", + "mesh_device", + "dim", + "n_heads", + "n_kv_heads", + "head_dim", + "qkv_size", + "scale", + "decode_xqkv_prg_config", # Required for all 1D (DRAM-sharded matmul) + "decode_attn_output_prg_config", # Required for all 1D (DRAM-sharded matmul) + "decode_sdpa_prg_config", + "decode_residual_memcfg", + "prefill_xqkv_prg_config", + "prefill_sdpa_prg_config", + "prefill_wo_prg_config", + ] + + # Multi-device needs CCL and topology + if self.mesh_device and self.mesh_device.get_num_devices() > 1: + required.extend(["tt_ccl", "topology"]) + + return all(getattr(self, f) is not None for f in required) + + +# ============================================================================= +# Attention1D Class +# ============================================================================= + + +class Attention1D(LightweightModule): + """ + Attention for 1D mesh topologies (N150, N300, T3K). + + Simple API (90% of users) - requires weights, model dimensions, and memory bounds: + attn = Attention1D(wqkv, wo, n_heads=32, n_kv_heads=8, head_dim=128, + max_batch_size=1, max_seq_len=2048) + + Power API (10% of users) - full customization via config: + config = Attention1DConfig(wqkv, wo, n_heads=32, n_kv_heads=8, head_dim=128, + sliding_window=4096, q_norm_config=..., ...) + attn = Attention1D.from_config(config) + + Execution paths: + Decode: QKV matmul → all_reduce → create_qkv_heads → rotary → KV cache → SDPA → concat_heads → WO matmul → all_reduce + Prefill: [reshape] → QKV matmul → all_reduce → create_qkv_heads → rotary → KV cache → SDPA → concat_heads → WO matmul + + KV Cache Management: + The KV cache is a static configuration stored in Attention1DConfig.kv_cache, + allocated once and reused for all forward calls. This design reflects how + inference engines like vLLM manage KV caches: the cache shape is determined + by static model config and the same tensors are reused for all forward calls. + + Options: + - from_model_args(): Automatically allocates KV cache (TTTv1 compatibility) + - from_config() with kv_cache: Pre-allocated (keys, values) tensors (vLLM integration) + - from_config() with kv_cache as LazyWeight tuple: Cache backed by files + """ + + def __init__( + self, + wqkv: LazyWeight, + wo: LazyWeight, + n_heads: int, + n_kv_heads: int, + head_dim: int, + max_batch_size: int, + max_seq_len: int, + ): + """ + Simple API for 90% of users - requires weights, model dimensions, and memory bounds. + + Args: + wqkv: Combined QKV projection weight, sharded on dim=-1 + wo: Output projection weight, sharded on dim=-2 + n_heads: Number of attention heads (from model config) + n_kv_heads: Number of key/value heads for GQA (from model config) + head_dim: Dimension per head (from model config) + max_batch_size: Maximum batch size for KV cache allocation + max_seq_len: Maximum sequence length for KV cache allocation + + Other settings (mesh_device, topology, program configs, etc.) are derived + automatically. Use from_config() for full customization. + + Example: + attn = Attention1D(wqkv, wo, n_heads=32, n_kv_heads=8, head_dim=128, + max_batch_size=1, max_seq_len=2048) + """ + super().__init__() + self.config = resolve_attention1d_arch_config( + Attention1DConfig( + wqkv=wqkv, + wo=wo, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + head_dim=head_dim, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + ) + ) + self._device_weights_loaded = False + self._bind_forward_methods() + + @classmethod + def from_config(cls, config: Attention1DConfig): + """ + Power API for 10% of users - any level of customization via config. + + Override any subset of fields in Attention1DConfig: + config = Attention1DConfig(wqkv, wo, n_heads=32, head_dim=128) + attn = Attention1D.from_config(config) + """ + if not isinstance(config, Attention1DConfig): + raise TypeError("Attention1D.from_config expects Attention1DConfig") + instance = object.__new__(cls) + super(Attention1D, instance).__init__() + instance.config = resolve_attention1d_arch_config(config) + instance._device_weights_loaded = False + instance._bind_forward_methods() + return instance + + @classmethod + def _from_config_with_arch(cls, config: Attention1DConfig, arch): + """Internal construction path for callers that already queried the mesh architecture.""" + instance = object.__new__(cls) + super(Attention1D, instance).__init__() + instance.config = resolve_attention1d_arch_config(config, _arch=arch) + instance._device_weights_loaded = False + instance._bind_forward_methods() + return instance + + def prefill_forward( + self, + x: ttnn.Tensor | LazyWeight, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int | list[int] = 0, + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + chunk_start_idx_tensor: ttnn.Tensor | None = None, + batch_size: Optional[int] = None, # todo)) work on removing this argument + ) -> ttnn.Tensor: + """ + Prefill forward - multiple tokens. + + Args: + x: Input tensor, shape (1, 1, seq_len, dim), TILE_LAYOUT. For batched prefill + (``batch_size > 1``) the leading dims hold ``batch_size`` users' tokens folded + into the sequence axis: shape ``(1, 1, batch_size * S, dim)``. + Memory config: DRAM interleaved (default ``prefill_input_memcfg``). + If ``LazyWeight``, it is automatically placed with the correct memory config. + rot_mats: Tuple of (cos, sin) rotation matrices for rotary embedding, + each shape (1, 1, head_dim, head_dim), TILE_LAYOUT, DRAM interleaved. + user_id: User ID for KV cache fill (selects which user's cache to write). For batched + prefill this is the *list* of valid slot indices, one per user in the batch. + page_table: Page table for paged attention (optional). For batched prefill this is the + full ``[batch_size, num_blocks]`` table; each user's row is selected by ``batch_idx``. + chunk_page_table: Page table for chunked prefill (optional). + chunk_start_idx: Start index for chunked prefill (optional). + chunk_start_idx_tensor: Runtime device start-index tensor for + position-general chunked prefill (optional). + batch_size: Optional number of users batched into one forward pass. ``1`` = the single-user path + (unchanged). ``> 1`` unfolds ``[1,1,B*S,*] -> [B,1,S,*]`` only for the + create_qkv_heads → rotary → KV-fill → SDPA → concat_heads window, mirroring TTTv1. + + Returns: + Output tensor (1, 1, seq_len, dim), DRAM interleaved. For batched prefill the output is + the folded ``(1, 1, batch_size * S, dim)`` hidden state. + + Note: + seq_len must be divisible by 128 and > 0. + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="prefill") + cfg = self.config + + if batch_size is None: # todo)) work on removing this argument + batch_size = x.shape[0] + if batch_size > 1: + x = ttnn.reshape(x, [1, 1, x.shape[-2] * x.shape[-3] * x.shape[0], -1]) + + seq_len = x.shape[-2] + original_seq_len = seq_len + assert seq_len % 128 == 0 and seq_len > 0, "seq_len must be divisible by 128" + assert seq_len % batch_size == 0, f"folded seq_len {seq_len} must be divisible by batch_size {batch_size}" + # Per-user sequence length (== seq_len when batch_size == 1). + per_user_seq_len = seq_len // batch_size + if batch_size > 1: + assert chunk_start_idx is None, "batched prefill does not support chunked SDPA" + + num_devices = cfg.mesh_device.get_num_devices() + n_local_heads = cfg.n_heads // num_devices + n_local_kv_heads = cfg.n_kv_heads // num_devices + + # --- STAGE 1: Reshape for long sequences --- + # The QKV matmul runs on the folded [1,1,seq_len,dim] tensor (seq_len = B*S), exactly like a + # single long sequence — so the batch axis is folded into the sequence here (TTTv1 parity). + if seq_len > MAX_QKV_MM_SEQ_LEN: + if seq_len % MAX_QKV_MM_SEQ_LEN != 0: + raise ValueError(f"seq_len {seq_len} must be divisible by {MAX_QKV_MM_SEQ_LEN}") + x = ttnn.reshape(x, [1, seq_len // MAX_QKV_MM_SEQ_LEN, MAX_QKV_MM_SEQ_LEN, -1]) + + # --- STAGE 2: QKV Matmul --- + # Above seq_len > 128 the folded-batch QKV matmul is large; minimal_matmul is ~2x faster than + # ttnn.linear there (TTTv1 parity). Opt-in via prefill_qkv_minimal_matmul; output shape is + # identical to the ttnn.linear path so everything downstream is unchanged. + if cfg.use_minimal_qkv_matmul(seq_len): + xqkv_fused = ttnn.experimental.minimal_matmul( + x, + self.wqkv, + compute_kernel_config=self.config.li_qkv_prefill_compute_kernel_cfg, + config=cfg.prefill_xqkv_minimal_matmul_config(seq_len), + ) + else: + xqkv_fused = ttnn.linear( + x, + self.wqkv, + dtype=cfg.activation_dtype or ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + compute_kernel_config=self.config.li_qkv_prefill_compute_kernel_cfg, + program_config=cfg.prefill_xqkv_prg_config(seq_len), + ) + + # Add bias if present + if self.wqkv_bias_prefill is not None: + xqkv_fused = xqkv_fused + self.wqkv_bias_prefill + + # --- STAGE 3: No all-reduce for 1D topologies --- + # For 1D meshes (1xN), QKV weights are sharded only on axis 1 (columns). + # After matmul, each device has complete rows - no partial sums to reduce. + # TTTv1's tt_all_reduce with cluster_axis=1 is a no-op for 1D meshes. + + # Reshape back + if seq_len > MAX_QKV_MM_SEQ_LEN: + xqkv_fused = ttnn.reshape(xqkv_fused, [1, 1, seq_len, -1]) + + if original_seq_len != seq_len: + xqkv_fused = xqkv_fused[:, :, :original_seq_len, :] + seq_len = original_seq_len + + if batch_size > 1: + xqkv_fused = ttnn.reshape(xqkv_fused, [batch_size, 1, seq_len // batch_size, -1]) + + ttnn.deallocate(x) + + # --- STAGE 4: Create QKV Heads --- + q_heads_pre_rot, k_heads_pre_rot, v_heads = ttnn.experimental.nlp_create_qkv_heads( + xqkv_fused, + num_heads=n_local_heads, + num_kv_heads=n_local_kv_heads, + transpose_k_heads=False, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + ttnn.deallocate(xqkv_fused) + + # --- STAGE 5: Q/K Normalization (optional) --- + if self.q_norm is not None: + q_heads_pre_rot = self.q_norm.prefill_forward(q_heads_pre_rot) + if self.k_norm is not None: + k_heads_pre_rot = self.k_norm.prefill_forward(k_heads_pre_rot) + + # --- STAGE 6: Rotary Embedding --- + # Explicit dtype check is required - ttnn.typecast is NOT a no-op when source == target types. + # Without this check, typecast would allocate a new tensor and run an identity copy kernel. + if q_heads_pre_rot.dtype != ttnn.bfloat16: + q_heads_pre_rot = ttnn.typecast(q_heads_pre_rot, dtype=ttnn.bfloat16) + + q_heads = ttnn.experimental.rotary_embedding_llama( + q_heads_pre_rot, rot_mats[0], rot_mats[1], cfg.transformation_mat_prefill, is_decode_mode=False + ) + ttnn.deallocate(q_heads_pre_rot) + + if k_heads_pre_rot.dtype != ttnn.bfloat16: + k_heads_pre_rot = ttnn.typecast(k_heads_pre_rot, dtype=ttnn.bfloat16) + + k_heads = ttnn.experimental.rotary_embedding_llama( + k_heads_pre_rot, rot_mats[0], rot_mats[1], cfg.transformation_mat_prefill, is_decode_mode=False + ) + ttnn.deallocate(k_heads_pre_rot) + + # --- STAGE 7: Typecast to cache dtype --- + keys, values = self.kv_cache + + k_heads_cache_dtype = ttnn.typecast(k_heads, dtype=keys.dtype) + ttnn.deallocate(k_heads) + + v_heads_cache_dtype = ttnn.typecast(v_heads, dtype=values.dtype) + ttnn.deallocate(v_heads) + + # --- STAGE 8: Shard KV for long sequences --- + # The shard memcfg is derived for the folded seq_len and assumes a [1,*,S,*] layout, so it + # only applies to single-user prefill. Batched prefill is always paged (page_table set), so + # this branch is naturally skipped; the extra batch_size==1 guard makes that explicit. + if seq_len >= cfg.min_kv_prefill_shard_seqlen and page_table is None and batch_size == 1: + k_fill = ttnn.interleaved_to_sharded(k_heads_cache_dtype, cfg.prefill_kv_memcfg(seq_len)) + v_fill = ttnn.interleaved_to_sharded(v_heads_cache_dtype, cfg.prefill_kv_memcfg(seq_len)) + else: + k_fill = k_heads_cache_dtype + v_fill = v_heads_cache_dtype + + # --- STAGE 9: KV Cache Fill --- + # Method bound at construction based on paged_attention_config (see _bind_forward_methods). + # Batched prefill fills one user per call in a loop: paged_fill_cache reads batch_idx_ptr[0] + # for all rows, so a per-row batch_idx in a single call is unsupported (TTTv1 attention.py + # L1013-1040). Device-verified correct (kernel probe E: per-slot readback pcc=1.0). + if batch_size > 1: + self._kv_fill_prefill_batched(keys, values, k_fill, v_fill, user_id, page_table, chunk_page_table) + else: + self._kv_fill_prefill(keys, values, k_fill, v_fill, user_id, page_table, chunk_page_table) + + # Deallocate sharded k_fill/v_fill only in non-paged mode (paged mode uses sliced views) + is_paged = cfg.paged_attention_config is not None + if seq_len >= cfg.min_kv_prefill_shard_seqlen and not is_paged and batch_size == 1: + ttnn.deallocate(k_fill) + ttnn.deallocate(v_fill) + + # --- STAGE 10: SDPA --- + q_heads_sdpa = ttnn.typecast(q_heads, dtype=cfg.activation_dtype or ttnn.bfloat8_b) + ttnn.deallocate(q_heads) + + # Chunked vs non-chunked is a RUNTIME decision (based on seq_len vs max_prefill_chunk_size), + # so this branching is valid under TTTv2 principles. Valid combinations: + # - Paged + Chunked: production vLLM serving (long prompts) + # - Paged + Non-chunked: short prompts in vLLM + # - Non-paged + Non-chunked: simple testing (uses else branch with contiguous KV) + # Invalid combinations (rejected at config time in _resolve_attention1d_config): + # - Non-paged + Chunked: chunked_sdpa requires page_table + # - sliding_window + Chunked: chunked_sdpa does not implement window masking + # Llama3_8B stages one device-resident chunk-start tensor before entering the layer loop and + # reuses it here. + # todo)) Other models still receive a Python scalar from the legacy shared executor. Retain + # that compatibility path until their executor/model plumbing is updated to stage and forward + # the tensor once per chunk. + if chunk_start_idx is not None or chunk_start_idx_tensor is not None: + if chunk_start_idx_tensor is None: + attn_output = ttnn.transformer.chunked_scaled_dot_product_attention( + input_tensor_q=q_heads_sdpa, + input_tensor_k=keys, + input_tensor_v=values, + page_table_tensor=page_table, + chunk_start_idx=chunk_start_idx, + compute_kernel_config=self.config.sdpa_prefill_compute_kernel_cfg, + program_config=cfg.prefill_sdpa_prg_config(seq_len, chunk_start_idx), + ) + else: + block_size = cfg.paged_attention_config.block_size + attn_output = ttnn.transformer.chunked_scaled_dot_product_attention( + input_tensor_q=q_heads_sdpa, + input_tensor_k=keys, + input_tensor_v=values, + page_table_tensor=page_table, + chunk_start_idx_tensor=chunk_start_idx_tensor, + compute_kernel_config=self.config.sdpa_prefill_compute_kernel_cfg, + program_config=cfg.prefill_sdpa_prg_config(seq_len, block_size), + ) + else: + sdpa_seq_len = seq_len // batch_size if batch_size > 1 else seq_len + attn_output = ttnn.transformer.scaled_dot_product_attention( + q_heads_sdpa, + k_heads_cache_dtype, + v_heads_cache_dtype, + is_causal=True, + sliding_window_size=cfg.sliding_window, + scale=cfg.scale, + compute_kernel_config=self.config.sdpa_prefill_compute_kernel_cfg, + program_config=cfg.prefill_sdpa_prg_config(sdpa_seq_len, None), + ) + + ttnn.deallocate(q_heads_sdpa) + ttnn.deallocate(k_heads_cache_dtype) + ttnn.deallocate(v_heads_cache_dtype) + + # --- STAGE 11: Reshape and Concat Heads --- + if batch_size == 1: + attn_output = ttnn.reshape(attn_output, [1, n_local_heads, -1, cfg.head_dim]) + + attn_output_concat = ttnn.experimental.nlp_concat_heads(attn_output, memory_config=ttnn.DRAM_MEMORY_CONFIG) + ttnn.deallocate(attn_output) + + if batch_size > 1: + attn_output_concat = ttnn.reshape(attn_output_concat, [1, 1, seq_len, -1]) + + # --- STAGE 12: Reshape for long sequences (to fit WO matmul on device) --- + # wo_prefill_len_cutoff (resolved: None → MAX_MM_SEQ_LEN) sets the WO M-chunk size; a per-model + # override regroups the folded prefill into fewer/larger M-chunks (bit-identical M-reblocking). + wo_cutoff = cfg.wo_prefill_len_cutoff + if seq_len > wo_cutoff: + attn_output_concat = ttnn.reshape(attn_output_concat, [1, seq_len // wo_cutoff, wo_cutoff, -1]) + + # --- STAGE 13: All-Gather for Ring topology --- + # Method bound at construction based on use_fused_all_gather_matmul (see _bind_forward_methods) + attn_output_concat = self._all_gather_before_wo_prefill(attn_output_concat) + + # --- STAGE 14: WO Matmul --- + # Above seq_len > 128 the folded-batch WO matmul is large and is the least-efficient prefill + # matmul on ttnn.linear; minimal_matmul (already used for QKV/FF2) recovers most of the gap. + # Opt-in via prefill_wo_minimal_matmul; output shape matches the ttnn.linear path so STAGE-15 + # and the all-reduce/no-op are unchanged. + if cfg.use_minimal_wo_matmul(seq_len): + output = ttnn.experimental.minimal_matmul( + attn_output_concat, + self.wo, + compute_kernel_config=self.config.li_o_prefill_compute_kernel_cfg, + config=cfg.prefill_wo_minimal_matmul_config(seq_len), + ) + else: + output = ttnn.linear( + attn_output_concat, + self.wo, + compute_kernel_config=self.config.li_o_prefill_compute_kernel_cfg, + dtype=cfg.activation_dtype or ttnn.bfloat8_b, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + program_config=cfg.prefill_wo_prg_config(seq_len), + ) + + # --- STAGE 15: Reshape back (undo long sequence reshape) --- + if seq_len > wo_cutoff: + output = ttnn.reshape(output, [1, 1, seq_len, -1]) + + ttnn.deallocate(attn_output_concat) + + # --- STAGE 16: All-Reduce output --- + # Method bound at construction based on use_fused_all_gather_matmul (see _bind_forward_methods) + output = self._reduce_after_wo_prefill(output) + + return output + + def decode_forward( + self, + x: ttnn.Tensor | LazyWeight, + current_pos: ttnn.Tensor, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + page_table: ttnn.Tensor | None = None, + ) -> ttnn.Tensor: + """ + Decode forward - single token per user. + + Args: + x: Input tensor, shape (seq_len, 1, batch, dim), TILE_LAYOUT. + Memory config: DRAM WIDTH_SHARDED (default ``decode_input_memcfg``). + If ``LazyWeight``, it is automatically placed with the correct memory config. + current_pos: Current position tensor, shape (batch_size,) on host. + rot_mats: Tuple of (cos, sin) rotation matrices for rotary embedding, + each shape (1, batch, head_dim, head_dim), TILE_LAYOUT, L1 interleaved. + page_table: Page table for paged attention (optional). + + Returns: + Output tensor (seq_len, 1, batch, dim), DRAM WIDTH_SHARDED + (``decode_residual_memcfg``). + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="decode") + cfg = self.config + + num_devices = cfg.mesh_device.get_num_devices() + n_local_heads = cfg.n_heads // num_devices + n_local_kv_heads = cfg.n_kv_heads // num_devices + + # --- STAGE 1: QKV Matmul --- + # All 1D topologies use DRAM-sharded matmul with L1_WIDTH_SHARDED output (matches TTTv1) + xqkv_fused_sharded = ttnn.linear( + x, + self.wqkv, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + program_config=cfg.decode_xqkv_prg_config, + compute_kernel_config=self.config.li_qkv_decode_compute_kernel_cfg, + dtype=cfg.activation_dtype or ttnn.bfloat16, + ) + + # Add bias if present + if self.wqkv_bias_decode: + num_tiles = int(math.ceil(xqkv_fused_sharded.shape[-2] / TILE_SIZE)) + xqkv_fused_sharded = xqkv_fused_sharded + self.wqkv_bias_decode[num_tiles - 1] + + ttnn.deallocate(x) + + # --- STAGE 2: Convert QKV from sharded to interleaved (matches TTTv1) --- + xqkv_fused = self._all_reduce_qkv_decode(xqkv_fused_sharded) + ttnn.deallocate(xqkv_fused_sharded) + + # Reshape for create_qkv_heads + fqkv_shape = xqkv_fused.shape + xqkv_fused = ttnn.reshape(xqkv_fused, (1, 1, cfg.max_batch_size, fqkv_shape[3]), (1, 1, 32, fqkv_shape[3])) + + # --- STAGE 3: Create QKV Heads --- + q_heads_pre_rot, k_heads_pre_rot, v_heads = ttnn.experimental.nlp_create_qkv_heads_decode( + xqkv_fused, + num_heads=n_local_heads, + num_kv_heads=n_local_kv_heads, + overlap_qk_coregrid=self._decode_overlap_qk_coregrid, + memory_config=cfg.decode_create_qkv_head_memcfg, + ) + ttnn.deallocate(xqkv_fused) + + # --- STAGE 4: Q/K Normalization (optional) --- + # Workaround: RMSNorm doesn't support HEIGHT_SHARDED inputs (TTTv1 norm_reshard pattern) + if self.q_norm is not None: + q_mem_cfg = q_heads_pre_rot.memory_config() + q_heads_pre_rot = ttnn.to_memory_config(q_heads_pre_rot, ttnn.L1_MEMORY_CONFIG, dtype=q_heads_pre_rot.dtype) + q_heads_pre_rot = self.q_norm.decode_forward(q_heads_pre_rot) + q_heads_pre_rot = ttnn.to_memory_config(q_heads_pre_rot, q_mem_cfg, dtype=q_heads_pre_rot.dtype) + if self.k_norm is not None: + k_mem_cfg = k_heads_pre_rot.memory_config() + k_heads_pre_rot = ttnn.to_memory_config(k_heads_pre_rot, ttnn.L1_MEMORY_CONFIG, dtype=k_heads_pre_rot.dtype) + k_heads_pre_rot = self.k_norm.decode_forward(k_heads_pre_rot) + k_heads_pre_rot = ttnn.to_memory_config(k_heads_pre_rot, k_mem_cfg, dtype=k_heads_pre_rot.dtype) + + # --- STAGE 5: Rotary Embedding --- + # Method bound at construction based on use_qk_fused (see _bind_forward_methods) + q_heads, k_heads = self._rotary_embed_decode(q_heads_pre_rot, k_heads_pre_rot, rot_mats) + ttnn.deallocate(q_heads_pre_rot) + ttnn.deallocate(k_heads_pre_rot) + + # --- STAGE 6: KV Cache Update --- + # Method bound at construction based on use_qk_fused (see _bind_forward_methods) + keys, values = self.kv_cache + self._kv_update_decode(keys, values, k_heads, v_heads, current_pos, page_table) + ttnn.deallocate(k_heads) + ttnn.deallocate(v_heads) + + # --- STAGE 7: SDPA --- + # Method bound at construction based on paged_attention_config (see _bind_forward_methods) + attn_output = self._sdpa_decode(q_heads, keys, values, current_pos, page_table) + + ttnn.deallocate(q_heads) + + # --- STAGE 8: Reshape for concat heads --- + attn_output_sharded = ttnn.to_memory_config( + attn_output, memory_config=cfg.decode_scores_memcfg(cfg.max_batch_size) + ) + + # --- STAGE 9: Concat Heads --- + attn_output_cat = ttnn.experimental.nlp_concat_heads_decode(attn_output_sharded, num_heads=n_local_heads) + ttnn.deallocate(attn_output_sharded) + ttnn.deallocate(attn_output) + + # --- STAGE 10: All-Gather + WO Matmul --- + # Method bound at construction based on use_fused_all_gather_matmul (see _bind_forward_methods) + dense_out = self._all_gather_wo_decode(attn_output_cat) + + ttnn.deallocate(attn_output_cat) + + # --- STAGE 11: Finalize output --- + # Method bound at construction based on use_fused_all_gather_matmul (see _bind_forward_methods) + # Fused path: returns dense_out as-is + # Non-fused path: reduce-scatter + final memory config + return self._finalize_decode_output(dense_out) + + def forward( + self, + x: ttnn.Tensor | LazyWeight, + current_pos: ttnn.Tensor | None, + rot_mats: tuple[ttnn.Tensor, ttnn.Tensor], + user_id: int | list[int] = 0, + mode: str = "decode", + page_table: ttnn.Tensor | None = None, + chunk_page_table: ttnn.Tensor | None = None, + chunk_start_idx: int | None = None, + batch_size: int = 1, + ) -> ttnn.Tensor: + """Dispatch to the appropriate forward method based on mode.""" + if mode == "prefill": + return self.prefill_forward( + x, + rot_mats, + user_id=user_id, + page_table=page_table, + chunk_page_table=chunk_page_table, + chunk_start_idx=chunk_start_idx, + batch_size=batch_size, + ) + else: + return self.decode_forward( + x, + current_pos, + rot_mats, + page_table=page_table, + ) + + def _bind_forward_methods(self): + """ + Bind forward method variants based on static config. + + This eliminates runtime if-statements in the hot path by pre-binding: + - _sdpa_decode: paged or non-paged SDPA for decode + - _kv_fill_prefill: paged or non-paged KV cache fill for prefill + - _all_gather_before_wo_prefill: fused all-gather or no-op + - _reduce_after_wo_prefill: reduce-scatter or no-op + - _all_gather_wo_decode: fused or separate all-gather + WO matmul + - _finalize_decode_output: fused (direct return) or non-fused (reduce + memcfg) + - _rotary_embed_decode: fused QK or separate rotary embedding + - _kv_update_decode: fused or separate KV cache update + + Note: page_table is still passed as a forward argument for vLLM compatibility + (vLLM dynamically manages page tables), but the method binding is static. + """ + cfg = self.config + is_paged = cfg.paged_attention_config is not None + + # Bind paged vs non-paged methods + if is_paged: + self._sdpa_decode = self._sdpa_decode_paged + self._kv_fill_prefill = self._kv_fill_prefill_paged + else: + self._sdpa_decode = self._sdpa_decode_non_paged + self._kv_fill_prefill = self._kv_fill_prefill_non_paged + + # Bind fused all-gather methods (based on use_fused_all_gather_matmul + topology) + use_fused = cfg.use_fused_all_gather_matmul and cfg.topology == ttnn.Topology.Ring + if use_fused: + self._all_gather_before_wo_prefill = self._all_gather_before_wo_prefill_fused + self._reduce_after_wo_prefill = self._reduce_after_wo_prefill_fused + self._all_gather_wo_decode = self._fused_all_gather_wo_decode + self._finalize_decode_output = self._finalize_decode_output_fused + else: + self._all_gather_before_wo_prefill = self._all_gather_before_wo_prefill_noop + self._reduce_after_wo_prefill = self._reduce_after_wo_prefill_non_fused + self._all_gather_wo_decode = self._separate_all_gather_wo_decode + self._finalize_decode_output = self._finalize_decode_output_non_fused + + # Bind fused QK methods (rotary embedding and KV cache update) + self._decode_overlap_qk_coregrid = not cfg.use_qk_fused + if cfg.use_qk_fused: + self._rotary_embed_decode = self._rotary_embed_decode_fused + self._kv_update_decode = self._kv_update_decode_fused + else: + self._rotary_embed_decode = self._rotary_embed_decode_nonfused + self._kv_update_decode = self._kv_update_decode_nonfused + + # ========================================================================= + # Bound SDPA Decode Methods (paged vs non-paged) + # ========================================================================= + + def _sdpa_decode_paged(self, q_heads, keys, values, current_pos, page_table) -> ttnn.Tensor: + """Paged SDPA decode - uses page_table for KV cache lookup.""" + cfg = self.config + return ttnn.transformer.paged_scaled_dot_product_attention_decode( + q_heads, + keys, + values, + page_table_tensor=page_table, + cur_pos_tensor=current_pos, + scale=cfg.scale, + sliding_window_size=cfg.sliding_window, + program_config=cfg.decode_sdpa_prg_config, + compute_kernel_config=self.config.sdpa_decode_compute_kernel_cfg, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + + def _sdpa_decode_non_paged(self, q_heads, keys, values, current_pos, page_table) -> ttnn.Tensor: + """Non-paged SDPA decode - contiguous KV cache (page_table ignored).""" + cfg = self.config + return ttnn.transformer.scaled_dot_product_attention_decode( + q_heads, + keys, + values, + cur_pos_tensor=current_pos, + scale=cfg.scale, + sliding_window_size=cfg.sliding_window, + program_config=cfg.decode_sdpa_prg_config, + compute_kernel_config=self.config.sdpa_decode_compute_kernel_cfg, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + + # ========================================================================= + # Bound KV Fill Prefill Methods (paged vs non-paged) + # ========================================================================= + + def _kv_fill_prefill_paged(self, keys, values, k_fill, v_fill, user_id, page_table, chunk_page_table) -> None: + """Paged KV cache fill for prefill - uses page_table for block allocation.""" + block_size = keys.shape[2] + fill_page_table = chunk_page_table if chunk_page_table is not None else page_table + page_len = fill_page_table.shape[1] * block_size + + if k_fill.shape[0] > 1: + # paged_fill_cache applies batch_idx to the passed page table. For batched + # prefill, pass one user's page-table row at a time and always write row 0. + valid_slots = user_id if isinstance(user_id, (list, tuple)) else list(range(k_fill.shape[0])) + seq_len_per_user = k_fill.shape[2] + for slot in valid_slots: + k_user = k_fill[slot : slot + 1, :, :, :] + v_user = v_fill[slot : slot + 1, :, :, :] + page_table_user = fill_page_table[slot : slot + 1, :] + page_len_user = page_table_user.shape[1] * block_size + k_user_sliced = k_user[:, :, :page_len_user, :] if page_len_user < seq_len_per_user else k_user + v_user_sliced = v_user[:, :, :page_len_user, :] if page_len_user < seq_len_per_user else v_user + ttnn.experimental.paged_fill_cache(keys, k_user_sliced, page_table_user, batch_idx=0) + ttnn.experimental.paged_fill_cache(values, v_user_sliced, page_table_user, batch_idx=0) + return + + k_fill_sliced = k_fill[:, :, :page_len, :] if page_len < k_fill.shape[2] else k_fill + v_fill_sliced = v_fill[:, :, :page_len, :] if page_len < v_fill.shape[2] else v_fill + + if k_fill_sliced.dtype != keys.dtype or v_fill_sliced.dtype != values.dtype: + raise ValueError( + "paged KV cache prefill requires K/V fill tensors to already match the cache dtype: " + f"k_fill={k_fill_sliced.dtype}, keys={keys.dtype}, " + f"v_fill={v_fill_sliced.dtype}, values={values.dtype}" + ) + + ttnn.experimental.paged_fill_cache(keys, k_fill_sliced, fill_page_table, batch_idx=user_id) + ttnn.experimental.paged_fill_cache(values, v_fill_sliced, fill_page_table, batch_idx=user_id) + + def _kv_fill_prefill_non_paged(self, keys, values, k_fill, v_fill, user_id, page_table, chunk_page_table) -> None: + """Non-paged KV cache fill for prefill - contiguous cache (page_table ignored).""" + cfg = self.config + ttnn.fill_cache(keys, k_fill, user_id % cfg.max_batch_size) + ttnn.fill_cache(values, v_fill, user_id % cfg.max_batch_size) + + def _kv_fill_prefill_batched(self, keys, values, k_fill, v_fill, user_id, page_table, chunk_page_table) -> None: + """Batched prefill KV fill — one fill call per user (loop), mirroring TTTv1 attention.py + L1013-1040. + + ``k_fill`` / ``v_fill`` are ``[batch_size, n_kv_heads, per_user_seq_len, head_dim]``. + Eager execution supplies active local ``user_id`` rows, while trace capture deliberately + supplies every padded row to keep one stable graph. Planner-owned padded page-table rows are + all ``-1``, the paged-fill skip sentinel, so those captured calls issue no KV writes. + ``batch_idx`` is scalar because ``paged_fill_cache`` reads ``batch_idx_ptr[0]`` for all rows. + """ + cfg = self.config + is_paged = cfg.paged_attention_config is not None + valid_slots = user_id if isinstance(user_id, (list, tuple)) else list(range(k_fill.shape[0])) + + if is_paged: + block_size = keys.shape[2] + fill_page_table = chunk_page_table if chunk_page_table is not None else page_table + page_len = fill_page_table.shape[1] * block_size + seq_len_per_user = k_fill.shape[2] + for slot_idx in valid_slots: + k_user = k_fill[slot_idx : slot_idx + 1, :, :, :] + v_user = v_fill[slot_idx : slot_idx + 1, :, :, :] + k_user = k_user[:, :, :page_len, :] if page_len < seq_len_per_user else k_user + v_user = v_user[:, :, :page_len, :] if page_len < seq_len_per_user else v_user + ttnn.experimental.paged_fill_cache(keys, k_user, fill_page_table, batch_idx=slot_idx) + ttnn.experimental.paged_fill_cache(values, v_user, fill_page_table, batch_idx=slot_idx) + else: + for slot_idx in valid_slots: + ttnn.fill_cache(keys, k_fill[slot_idx : slot_idx + 1, :, :, :], slot_idx % cfg.max_batch_size) + ttnn.fill_cache(values, v_fill[slot_idx : slot_idx + 1, :, :, :], slot_idx % cfg.max_batch_size) + + # ========================================================================= + # Bound All-Gather/Reduce Methods for Prefill (fused vs non-fused) + # ========================================================================= + + def _all_gather_before_wo_prefill_fused(self, attn_output_concat: ttnn.Tensor) -> ttnn.Tensor: + """Fused path: all-gather before WO matmul (Ring topology).""" + cfg = self.config + # Optionally cast the gather input (bf16) to a smaller CCL dtype (e.g. bfloat8_b) to halve this + # per-layer collective's cross-device payload. None (default) leaves the dtype unchanged. an + # explicit typecast is required (to_memory_config does not cast an already-DRAM tensor). + if cfg.prefill_ag_ccl_dtype is not None and attn_output_concat.dtype != cfg.prefill_ag_ccl_dtype: + attn_output_cast = ttnn.typecast(attn_output_concat, cfg.prefill_ag_ccl_dtype) + ttnn.deallocate(attn_output_concat) + attn_output_concat = attn_output_cast + return ttnn.experimental.all_gather_async( + attn_output_concat, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=cfg.topology, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + + def _all_gather_before_wo_prefill_noop(self, attn_output_concat: ttnn.Tensor) -> ttnn.Tensor: + """Non-fused path: no all-gather before WO matmul.""" + return attn_output_concat + + def _reduce_after_wo_prefill_fused(self, output: ttnn.Tensor) -> ttnn.Tensor: + """Fused path: no reduce after WO matmul (already complete from all-gather).""" + return output + + def _reduce_after_wo_prefill_non_fused(self, output: ttnn.Tensor) -> ttnn.Tensor: + """Non-fused path: reduce-scatter after WO matmul.""" + return self._all_reduce_output_prefill(output) + + # ========================================================================= + # Bound Finalize Methods for Decode (fused vs non-fused) + # ========================================================================= + + def _finalize_decode_output_fused(self, dense_out: ttnn.Tensor) -> ttnn.Tensor: + """Fused path: output is already complete, return as-is.""" + return dense_out + + def _finalize_decode_output_non_fused(self, dense_out: ttnn.Tensor) -> ttnn.Tensor: + """Non-fused path: reduce-scatter and apply final memory config.""" + cfg = self.config + + dense_out_reduced = self._all_reduce_output_decode(dense_out) + + # Only deallocate if a new tensor was created (multi-device case). + # For single device, _all_reduce_output_decode returns input unchanged. + if cfg.mesh_device.get_num_devices() > 1: + ttnn.deallocate(dense_out) + + return ttnn.to_memory_config(dense_out_reduced, cfg.decode_residual_memcfg) + + # ========================================================================= + # Bound Rotary Embedding Methods for Decode (fused QK vs separate) + # ========================================================================= + + def _rotary_embed_decode_fused( + self, q_heads_pre_rot: ttnn.Tensor, k_heads_pre_rot: ttnn.Tensor, rot_mats: tuple[ttnn.Tensor, ttnn.Tensor] + ) -> tuple[ttnn.Tensor, ttnn.Tensor]: + """Fused QK rotary embedding - single kernel for both Q and K. + + The fused kernel requires Q and K on non-overlapping core grids. + After create_qkv_heads_decode with interleaved input, Q is already on the + correct cores (first batch cores), so only K needs resharding to a + non-overlapping grid — saving 1 dispatch vs the original 2-reshard approach. + """ + cfg = self.config + k_heads_pre_rot = self._reshard_k_for_fused(k_heads_pre_rot, q_heads_pre_rot.shape[1]) + return ttnn.experimental.rotary_embedding_llama_fused_qk( + q_heads_pre_rot, k_heads_pre_rot, rot_mats[0], rot_mats[1], cfg.transformation_mat_decode + ) + + def _rotary_embed_decode_nonfused( + self, q_heads_pre_rot: ttnn.Tensor, k_heads_pre_rot: ttnn.Tensor, rot_mats: tuple[ttnn.Tensor, ttnn.Tensor] + ) -> tuple[ttnn.Tensor, ttnn.Tensor]: + """Separate rotary embedding - independent kernels for Q and K.""" + cfg = self.config + q_heads = ttnn.experimental.rotary_embedding_llama( + q_heads_pre_rot, rot_mats[0], rot_mats[1], cfg.transformation_mat_decode, is_decode_mode=True + ) + k_heads = ttnn.experimental.rotary_embedding_llama( + k_heads_pre_rot, rot_mats[0], rot_mats[1], cfg.transformation_mat_decode, is_decode_mode=True + ) + return q_heads, k_heads + + # ========================================================================= + # Bound KV Cache Update Methods for Decode (fused vs separate) + # ========================================================================= + + def _kv_update_decode_fused( + self, + keys: ttnn.Tensor, + values: ttnn.Tensor, + k_heads: ttnn.Tensor, + v_heads: ttnn.Tensor, + current_pos: ttnn.Tensor, + page_table: ttnn.Tensor | None, + ) -> None: + """Fused KV cache update - single kernel for both K and V.""" + ttnn.experimental.paged_fused_update_cache( + keys, k_heads, values, v_heads, update_idxs_tensor=current_pos, page_table=page_table + ) + + def _kv_update_decode_nonfused( + self, + keys: ttnn.Tensor, + values: ttnn.Tensor, + k_heads: ttnn.Tensor, + v_heads: ttnn.Tensor, + current_pos: ttnn.Tensor, + page_table: ttnn.Tensor | None, + ) -> None: + """Separate KV cache update - independent kernels for K and V.""" + ttnn.experimental.paged_update_cache(keys, k_heads, update_idxs_tensor=current_pos, page_table=page_table) + ttnn.experimental.paged_update_cache(values, v_heads, update_idxs_tensor=current_pos, page_table=page_table) + + def load_device_weights(self): + """Materialize LazyWeights onto device. Called automatically on first forward; idempotent.""" + if self._device_weights_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device weights!" + + cfg = self.config + self.wqkv = cfg.wqkv.get_device_weight() + self.wo = cfg.wo.get_device_weight() + + # Initialize Q/K norm RMSNorm1D instances if configs present + if cfg.q_norm_config is not None: + self.q_norm = RMSNorm1D.from_config(cfg.q_norm_config) + self.q_norm.load_device_weights() + else: + self.q_norm = None + + if cfg.k_norm_config is not None: + self.k_norm = RMSNorm1D.from_config(cfg.k_norm_config) + self.k_norm.load_device_weights() + else: + self.k_norm = None + + # Materialize bias LazyWeights + if cfg._wqkv_bias_decode is not None: + self.wqkv_bias_decode = [bias.get_device_weight() for bias in cfg._wqkv_bias_decode] + else: + self.wqkv_bias_decode = None + if cfg._wqkv_bias_prefill is not None: + self.wqkv_bias_prefill = cfg._wqkv_bias_prefill.get_device_weight() + else: + self.wqkv_bias_prefill = None + + if cfg.kv_cache is not None: + keys, values = cfg.kv_cache + if isinstance(keys, LazyWeight): + keys = keys.get_device_weight() + if isinstance(values, LazyWeight): + values = values.get_device_weight() + self.kv_cache = (keys, values) + else: + assert not hasattr(self, "kv_cache") or self.kv_cache is None, ( + "kv_cache was set directly on Attention1D instead of via config.kv_cache. " + "Use model.set_kv_cache() to bind kv_cache through the config." + ) + self.kv_cache = None + + self._device_weights_loaded = True + + # ========================================================================= + # Internal CCL methods + # ========================================================================= + + def _all_reduce_qkv_decode(self, xqkv: ttnn.Tensor) -> ttnn.Tensor: + """ + All-reduce QKV for decode mode. + + For 1D topologies, tt_all_reduce with cluster_axis=1 returns input as-is (no reduction). + We just need to convert from sharded to interleaved (matches TTTv1 non-TG path). + """ + # All 1D topologies: convert L1_WIDTH_SHARDED to L1 interleaved (matches TTTv1) + # bfloat16 is required by nlp_create_qkv_heads_decode + return ttnn.sharded_to_interleaved(xqkv, ttnn.L1_MEMORY_CONFIG, ttnn.bfloat16) + + def _all_reduce_output_decode(self, output: ttnn.Tensor) -> ttnn.Tensor: + """ + Final all-reduce for decode output. + + For 1D topologies (1xN), this is reduce_scatter only. + """ + cfg = self.config + + # Single device: no all-reduce needed + if cfg.mesh_device.get_num_devices() == 1: + return output + + # For 1D topologies: reduce_scatter across devices on axis 1 + # Convert sharded to interleaved first if needed + if output.is_sharded(): + output_interleaved = ttnn.sharded_to_interleaved(output, ttnn.L1_MEMORY_CONFIG) + output.deallocate(True) + else: + output_interleaved = output + + reduced = ttnn.experimental.reduce_scatter_minimal_async( + output_interleaved, + persistent_output_buffers=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_rs_semaphore_handles(), + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + num_links=cfg.num_reduce_scatter_links, + memory_config=cfg.decode_residual_memcfg, + intermediate_memory_config=ttnn.DRAM_MEMORY_CONFIG, + topology=cfg.topology, + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + output_interleaved.deallocate(True) + + return reduced + + def _all_reduce_output_prefill(self, output: ttnn.Tensor) -> ttnn.Tensor: + """ + Final all-reduce for prefill output. + + For 1D topologies (1xN), this is reduce_scatter only. + """ + cfg = self.config + + # Single device: no all-reduce needed + if cfg.mesh_device.get_num_devices() == 1: + return output + + # For 1D topologies: reduce_scatter across devices + # Prefill uses interleaved memory, ensure we're in DRAM. + output = ttnn.to_memory_config(output, ttnn.DRAM_MEMORY_CONFIG) + + # When the fused all-gather+WO path is off, the WO output (bf16) is reduce-scattered here. + # Optionally cast it to a smaller CCL dtype first to halve the cross-device payload (matches + # TTTv1's bfloat8_b ccl_dtype). None (default) leaves the dtype unchanged. to_memory_config + # does not cast an already-DRAM tensor, so an explicit typecast is required. + if cfg.prefill_reduce_ccl_dtype is not None and output.dtype != cfg.prefill_reduce_ccl_dtype: + output_cast = ttnn.typecast(output, cfg.prefill_reduce_ccl_dtype) + ttnn.deallocate(output) + output = output_cast + + reduced = ttnn.experimental.reduce_scatter_minimal_async( + output, + persistent_output_buffers=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_rs_semaphore_handles(), + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + num_links=cfg.num_reduce_scatter_links, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + intermediate_memory_config=ttnn.DRAM_MEMORY_CONFIG, + topology=cfg.topology, + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + output.deallocate(True) + + return reduced + + def _fused_all_gather_wo_decode(self, attn_output_cat: ttnn.Tensor) -> ttnn.Tensor: + """Fused all-gather matmul for Ring topology.""" + cfg = self.config + + attn_output_cat = ttnn.to_memory_config(attn_output_cat, cfg.decode_all_gather_matmul_memcfg) + + _, dense_out = ttnn.experimental.all_gather_matmul_async( + attn_output_cat, + self.wo, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(), + all_gather_core_grid_offset=(0, 4), + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + num_links=cfg.decode_agmm_num_links, + memory_config_ag=cfg.decode_all_gather_matmul_memcfg, + memory_config_mm=cfg.decode_residual_memcfg, + program_config=cfg.decode_all_gather_matmul_prg_config, + compute_kernel_config=self.config.li_o_decode_compute_kernel_cfg, + chunks_per_sync=cfg.decode_agmm_chunks_per_sync, + num_workers_per_link=cfg.decode_agmm_num_workers_per_link, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + + return ttnn.to_memory_config(dense_out, cfg.decode_residual_memcfg) + + def _separate_all_gather_wo_decode(self, attn_output_cat: ttnn.Tensor) -> ttnn.Tensor: + """ + Separate all-gather then WO matmul (non-Ring or when fused not available). + + For 1D topologies with cluster_axis=1, all-gather is a no-op since there's only + 1 device on axis 0. The attention heads are already local to each device. + """ + cfg = self.config + + # For 1D topologies (1xN mesh), all-gather with cluster_axis=1 is a no-op + # because axis 0 has only 1 device. Skip the CCL call entirely. + # This matches TTTv1 behavior where tt_all_gather returns input unchanged. + + # WO matmul + dense_out = ttnn.linear( + attn_output_cat, + self.wo, + program_config=cfg.decode_attn_output_prg_config, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + compute_kernel_config=self.config.li_o_decode_compute_kernel_cfg, + ) + + return dense_out + + def _reshard_k_for_fused(self, k_tensor: ttnn.Tensor, q_batch: int) -> ttnn.Tensor: + """Move K tensor to non-overlapping core grid for fused QK rotary embedding. + + After create_qkv_heads_decode with interleaved input, Q and K share the + same core grid. The fused rotary kernel requires non-overlapping grids. + Q's grid (first ``q_batch`` cores) is already correct, so we only move K + to the next ``q_batch`` cores — one dispatch instead of two. + """ + n_kv_heads = k_tensor.shape[2] + row_size = 8 + k_start_core = ttnn.CoreCoord(q_batch % row_size, q_batch // row_size) + k_core_grid = ttnn.CoreRangeSet({_num_to_corerange(q_batch, start_core=k_start_core)}) + k_mem_config = ttnn.create_sharded_memory_config( + shape=(nearest_32(n_kv_heads), self.config.head_dim), + core_grid=k_core_grid, + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + return ttnn.to_memory_config(k_tensor, k_mem_config) + + # ========================================================================= + # Factory method for backward compatibility + # ========================================================================= + + @classmethod + def from_model_args( + cls, + mesh_device, + tt_ccl, + args, + state_dict, + weight_cache_path, + layer_num: int, + transformation_mats: dict[str, ttnn.Tensor], + paged_attention_config=None, + use_paged_kv_cache: bool = False, + ): + """Factory method for backward compatibility with ModelArgs.""" + import torch + + from models.tt_transformers.tt.model_config import OpGroup, TensorGroup + + if args.is_galaxy: + raise ValueError("Attention1D cannot be used for Galaxy devices. Use Attention2D instead.") + + configuration = args + model_config = configuration.get_model_config() + decoders_opt = model_config.get("DECODERS_OPTIMIZATIONS") + + layer_name = configuration.get_state_dict_prefix("Attention", layer_num) + + wq_str = f"{layer_name}.wq" + wk_str = f"{layer_name}.wk" + wv_str = f"{layer_name}.wv" + wo_str = f"{layer_name}.wo" + q_norm_str = f"{layer_name}.q_norm" + k_norm_str = f"{layer_name}.k_norm" + + # Get dtypes + wqkv_dtype = decoders_opt.get_tensor_dtype(decoder_id=layer_num, tensor=TensorGroup.WQKV) + wo_dtype = decoders_opt.get_tensor_dtype(decoder_id=layer_num, tensor=TensorGroup.WO) + kv_cache_dtype = decoders_opt.get_tensor_dtype(decoder_id=layer_num, tensor=TensorGroup.KV_CACHE) + activation_dtype = decoders_opt.get_tensor_dtype(decoder_id=layer_num, tensor=TensorGroup.ACTIVATION) + + # Get compute kernel configs + li_qkv_decode_cfg = decoders_opt.get_math_fidelity( + decoder_id=layer_num, op=OpGroup.LI_QKV_DECODE, configuration=configuration + ) + sdpa_decode_cfg = decoders_opt.get_math_fidelity( + decoder_id=layer_num, op=OpGroup.SDPA_DECODE, configuration=configuration + ) + li_o_decode_cfg = decoders_opt.get_math_fidelity( + decoder_id=layer_num, op=OpGroup.LI_O_DECODE, configuration=configuration + ) + li_qkv_prefill_cfg = decoders_opt.get_math_fidelity( + decoder_id=layer_num, op=OpGroup.LI_QKV_PREFILL, configuration=configuration + ) + sdpa_prefill_cfg = decoders_opt.get_math_fidelity( + decoder_id=layer_num, op=OpGroup.SDPA_PREFILL, configuration=configuration + ) + li_o_prefill_cfg = decoders_opt.get_math_fidelity( + decoder_id=layer_num, op=OpGroup.LI_O_PREFILL, configuration=configuration + ) + + # Build combined QKV weight + num_devices = mesh_device.get_num_devices() + qkv_list = [] + for i in range(num_devices): + wq_selected = torch.chunk(state_dict[f"{wq_str}.weight"], num_devices, dim=0)[i] + wk_selected = torch.chunk(state_dict[f"{wk_str}.weight"], num_devices, dim=0)[i] + wv_selected = torch.chunk(state_dict[f"{wv_str}.weight"], num_devices, dim=0)[i] + + wq = torch.transpose(wq_selected, -2, -1) + wk = torch.transpose(wk_selected, -2, -1) + wv = torch.transpose(wv_selected, -2, -1) + + qkv = torch.cat([wq, wk, wv], dim=-1) + qkv_list.append(qkv) + + qkv_cat = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0) + + # Create LazyWeights + wqkv_mem_config = configuration.create_dram_sharded_mem_config( + configuration.dim, configuration.qkv_size // num_devices + ) + + wqkv = LazyWeight( + source=qkv_cat, + dtype=wqkv_dtype, + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + memory_config=wqkv_mem_config, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + cache_dir_weight_name=(Path(weight_cache_path) / layer_name, "wqkv_sharded") if weight_cache_path else None, + ) + + pt_wo = state_dict[f"{wo_str}.weight"].transpose(-1, -2).unsqueeze(0).unsqueeze(0) + wo_mem_config = configuration.create_dram_sharded_mem_config( + (configuration.n_heads * configuration.head_dim) // num_devices, configuration.dim + ) + + use_fused_all_gather_matmul = getattr(configuration, "use_fused_all_gather_matmul", False) + + wo = LazyWeight( + source=pt_wo, + dtype=wo_dtype, + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG if use_fused_all_gather_matmul else wo_mem_config, + mesh_mapper_config=ttnn.MeshMapperConfig( + # For fused: shard on dim -1 (width) - matches TTTv1 dims=(2, 3) + # For non-fused: shard on dim -2 (height) - matches TTTv1 dims=(3, 2) + placements=[ttnn.PlacementShard(-1 if use_fused_all_gather_matmul else -2)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + cache_dir_weight_name=( + ( + Path(weight_cache_path) / layer_name, + "wo_width_sharded" if use_fused_all_gather_matmul else "wo", + ) + if weight_cache_path + else None + ), + ) + + # Q/K norm configs (optional) - using RMSNorm1DConfig composition pattern + q_norm_config = None + k_norm_config = None + + # Get compute kernel config for Q/K normalization + # Q/K norm uses interleaved path (not sharded like regular decode RMSNorm) + arch = mesh_device.arch() + architecture = _attention_architecture(mesh_device, arch) + qk_norm_compute_kernel = ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=False, + ) + + if f"{q_norm_str}.weight" in state_dict: + q_norm_torch = state_dict[f"{q_norm_str}.weight"] + # Note: add_unit_offset is handled by RMSNorm1DConfig resolution + q_norm_torch = q_norm_torch.reshape(1, 1, -1, TILE_SIZE) + + q_norm_lazy = LazyWeight( + source=q_norm_torch, + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(Path(weight_cache_path) / layer_name, "q_norm") if weight_cache_path else None, + ) + + q_norm_common = RMSNorm1DConfig( + weight=q_norm_lazy, + mesh_device=mesh_device, + eps=configuration.norm_eps, + add_unit_offset=configuration.rms_norm_add_unit_offset, + decode_in_sharded=False, # Q/K heads are interleaved after create_qkv_heads + decode_out_sharded=False, + prefill_distributed=False, # Q/K norm doesn't need distributed prefill + ) + q_norm_config = replace(q_norm_common, compute_kernel_config=qk_norm_compute_kernel) + + if f"{k_norm_str}.weight" in state_dict: + k_norm_torch = state_dict[f"{k_norm_str}.weight"] + # Note: add_unit_offset is handled by RMSNorm1DConfig resolution + k_norm_torch = k_norm_torch.reshape(1, 1, -1, TILE_SIZE) + + k_norm_lazy = LazyWeight( + source=k_norm_torch, + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(Path(weight_cache_path) / layer_name, "k_norm") if weight_cache_path else None, + ) + + k_norm_common = RMSNorm1DConfig( + weight=k_norm_lazy, + mesh_device=mesh_device, + eps=configuration.norm_eps, + add_unit_offset=configuration.rms_norm_add_unit_offset, + decode_in_sharded=False, # Q/K heads are interleaved after create_qkv_heads + decode_out_sharded=False, + prefill_distributed=False, # Q/K norm doesn't need distributed prefill + ) + k_norm_config = replace(k_norm_common, compute_kernel_config=qk_norm_compute_kernel) + + # Handle QKV bias + wqkv_bias = None + if f"{wq_str}.bias" in state_dict: + qkv_bias = torch.concat( + [ + torch.concat( + [ + torch.chunk(state_dict[f"{wq_str}.bias"], num_devices)[i], + torch.chunk(state_dict[f"{wk_str}.bias"], num_devices)[i], + torch.chunk(state_dict[f"{wv_str}.bias"], num_devices)[i], + ], + dim=-1, + ) + for i in range(num_devices) + ], + dim=-1, + ) + wqkv_bias = LazyWeight(source=qkv_bias) + + # Determine scale + if configuration.query_pre_attn_scalar is not None: + scale = configuration.query_pre_attn_scalar**-0.5 + else: + scale = configuration.head_dim**-0.5 + + # Build config + # Note: kv_cache is created by default in _resolve_attention1d_config if not provided + # and use_paged_kv_cache=False. For paged cache (vLLM), set use_paged_kv_cache=True. + config = Attention1DConfig( + wqkv=wqkv, + wo=wo, + q_norm_config=q_norm_config, + k_norm_config=k_norm_config, + wqkv_bias=wqkv_bias, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + topology=configuration.ccl_topology(), + dim=configuration.dim, + n_heads=configuration.n_heads, + n_kv_heads=configuration.n_kv_heads, + head_dim=configuration.head_dim, + qkv_size=configuration.qkv_size, + max_batch_size=configuration.max_batch_size, + max_seq_len=configuration.max_seq_len, + scale=scale, + sliding_window=configuration.sliding_window if hasattr(configuration, "sliding_window") else None, + use_qk_fused=getattr(configuration, "use_qk_fused", False), + use_vllm_paged_kv_cache=use_paged_kv_cache, + paged_attention_config=paged_attention_config, + kv_cache_dtype=kv_cache_dtype, + min_kv_prefill_shard_seqlen=configuration.min_kv_prefill_shard_seqlen, + wqkv_dtype=wqkv_dtype, + wo_dtype=wo_dtype, + activation_dtype=activation_dtype, + decode_xqkv_prg_config=model_config.get("XQKV_DECODE_PROGCFG"), + decode_sdpa_prg_config=model_config.get("SDPA_DECODE_PROGCFG"), + decode_attn_output_prg_config=model_config.get("ATTN_OUTPUT_PROGCFG"), + decode_residual_memcfg=model_config.get("DECODE_RESIDUAL_MEMCFG"), + decode_create_qkv_head_memcfg=model_config.get("CREATE_QKV_DECODE_SHARD"), + decode_scores_memcfg=model_config.get("SCORES_BATCHED_MM_OUTPUT_MEMCFG"), + prefill_xqkv_prg_config=model_config.get("XQKV_PREFILL_PROGCFG"), + prefill_sdpa_prg_config=model_config.get("SDPA_PROGCFG"), + prefill_wo_prg_config=model_config.get("WO_PREFILL_PROGCFG"), + prefill_kv_memcfg=model_config.get("KV_PREFILL_MEM_CFG"), + use_fused_all_gather_matmul=use_fused_all_gather_matmul, + decode_all_gather_matmul_prg_config=model_config.get("ATTN_ALL_GATHER_MATMUL_PROGCFG"), + decode_all_gather_matmul_memcfg=model_config.get("ATTN_ALL_GATHER_MATMUL_OUTPUT_MEMCFG"), + # Use provided transformation matrices if available; otherwise auto-created in resolve + transformation_mat_decode=transformation_mats.get("decode") if transformation_mats else None, + transformation_mat_prefill=transformation_mats.get("prefill") if transformation_mats else None, + li_qkv_decode_compute_kernel_cfg=li_qkv_decode_cfg, + sdpa_decode_compute_kernel_cfg=sdpa_decode_cfg, + li_o_decode_compute_kernel_cfg=li_o_decode_cfg, + li_qkv_prefill_compute_kernel_cfg=li_qkv_prefill_cfg, + sdpa_prefill_compute_kernel_cfg=sdpa_prefill_cfg, + li_o_prefill_compute_kernel_cfg=li_o_prefill_cfg, + prefill_qkv_grid=(8, 8) if architecture == "wormhole" else (8, 10), + dram_shard_grid_width=(8 if architecture == "wormhole" else configuration.dram_shard_grid_width), + decode_create_qkv_head_grid=None if architecture == "wormhole" else ttnn.CoreGrid(y=4, x=8), + decode_transformation_core_grid=( + mesh_device.compute_with_storage_grid_size() if architecture == "wormhole" else ttnn.CoreCoord(8, 8) + ), + ) + return cls._from_config_with_arch(config, arch) + + +# ============================================================================= +# Config resolution +# ============================================================================= + + +_ATTENTION_COMPUTE_SLOT_NAMES = ( + "li_qkv_decode_compute_kernel_cfg", + "sdpa_decode_compute_kernel_cfg", + "li_o_decode_compute_kernel_cfg", + "li_qkv_prefill_compute_kernel_cfg", + "sdpa_prefill_compute_kernel_cfg", + "li_o_prefill_compute_kernel_cfg", +) + + +def _attention_mesh_device(config: Attention1DConfig) -> ttnn.MeshDevice: + mesh_device = config.mesh_device or config.wqkv.device or ttnn.GetDefaultDevice() + if mesh_device is None: + raise ValueError("Attention1D requires a mesh_device on the config, weight, or default device") + return mesh_device + + +def _attention_architecture(mesh_device: ttnn.MeshDevice, arch) -> str: + """Return the supported architecture name using only the concrete mesh.""" + if arch == ttnn.device.Arch.WORMHOLE_B0: + return "wormhole" + if arch == ttnn.device.Arch.BLACKHOLE: + return "blackhole" + raise ValueError(f"Unsupported Attention1D architecture: {arch}") + + +def _shared_attention_compute_defaults(arch) -> dict[str, object]: + """Build the architecture-neutral TTTv2 six-slot baseline for ``arch``. + + This preserves simple/common-config construction. Model numerical recipes + are supplied explicitly by model-owned architecture wrappers and take + precedence over these shared defaults. + """ + + def make(fidelity, *, fp32=False): + return ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=fidelity, + math_approx_mode=False, + fp32_dest_acc_en=fp32, + packer_l1_acc=True, + ) + + return { + "li_qkv_decode_compute_kernel_cfg": make(ttnn.MathFidelity.HiFi2), + "sdpa_decode_compute_kernel_cfg": make(ttnn.MathFidelity.HiFi2), + "li_o_decode_compute_kernel_cfg": make(ttnn.MathFidelity.HiFi2), + "li_qkv_prefill_compute_kernel_cfg": make(ttnn.MathFidelity.HiFi2), + "sdpa_prefill_compute_kernel_cfg": make(ttnn.MathFidelity.HiFi4, fp32=True), + "li_o_prefill_compute_kernel_cfg": make(ttnn.MathFidelity.HiFi2), + } + + +def _wormhole_attention_architecture_fields(mesh_device: ttnn.MeshDevice) -> dict[str, object]: + """Return the frozen Wormhole compatibility geometry. + + Blackhole geometry is deliberately absent: model/SKU overrides are + expressed as optional fields on ``Attention1DConfig``. + """ + return { + "prefill_qkv_grid": (8, 8), + "dram_shard_grid_width": 8, + "decode_create_qkv_head_grid": None, + "decode_transformation_core_grid": mesh_device.compute_with_storage_grid_size(), + } + + +def _blackhole_attention_mesh_overlay(mesh_device: ttnn.MeshDevice) -> dict[str, object]: + """Resolve effective Blackhole mesh/SKU geometry for the simple API.""" + return { + "prefill_qkv_grid": (8, 10), + "dram_shard_grid_width": mesh_device.dram_grid_size().x, + "decode_create_qkv_head_grid": ttnn.CoreGrid(y=4, x=8), + "decode_transformation_core_grid": ttnn.CoreCoord(8, 8), + } + + +def resolve_attention1d_arch_config(config: Attention1DConfig, *, _arch=None) -> Attention1DConfig: + """Resolve and validate Attention1D's architecture composition exactly once.""" + if not isinstance(config, Attention1DConfig): + raise TypeError("resolve_attention1d_arch_config expects Attention1DConfig") + common = config + mesh_device = _attention_mesh_device(common) + arch = mesh_device.arch() if _arch is None else _arch + architecture = _attention_architecture(mesh_device, arch) + + defaults = _shared_attention_compute_defaults(arch) + slots = { + name: defaults[name] if getattr(common, name) is None else getattr(common, name) + for name in _ATTENTION_COMPUTE_SLOT_NAMES + } + default_fields = ( + _wormhole_attention_architecture_fields(mesh_device) + if architecture == "wormhole" + else _blackhole_attention_mesh_overlay(mesh_device) + ) + architecture_fields = { + name: default if getattr(common, name) is None else getattr(common, name) + for name, default in default_fields.items() + } + + missing = [name for name, value in slots.items() if value is None] + if missing: + raise ValueError(f"Incomplete {architecture} Attention1D recipe; missing: {', '.join(missing)}") + kernel_arch = ttnn.device.Arch.WORMHOLE_B0 if architecture == "wormhole" else ttnn.device.Arch.BLACKHOLE + copied_slots = {} + for name, value in slots.items(): + try: + copied_slots[name] = ttnn.init_device_compute_kernel_config(kernel_arch, value) + except (TypeError, ValueError) as error: + raise ValueError(f"Invalid {architecture} Attention1D {name}: {error}") from error + slots = copied_slots + dram_width = architecture_fields["dram_shard_grid_width"] + expected_dram_width = 8 if architecture == "wormhole" else mesh_device.dram_grid_size().x + if dram_width <= 0: + raise ValueError("Attention1D SKU overlay requires a positive dram_shard_grid_width") + if dram_width != expected_dram_width: + raise ValueError(f"Attention1D SKU DRAM width {dram_width} does not match resolved width {expected_dram_width}") + compute_grid = mesh_device.compute_with_storage_grid_size() + qkv_x, qkv_y = architecture_fields["prefill_qkv_grid"] + if qkv_x <= 0 or qkv_y <= 0 or qkv_x > compute_grid.x or qkv_y > compute_grid.y: + raise ValueError( + f"Attention1D prefill QKV grid {(qkv_x, qkv_y)} exceeds compute grid " f"{(compute_grid.x, compute_grid.y)}" + ) + for field_name in ("decode_create_qkv_head_grid", "decode_transformation_core_grid"): + grid = architecture_fields[field_name] + if grid is None: + continue + if grid.x <= 0 or grid.y <= 0 or grid.x > compute_grid.x or grid.y > compute_grid.y: + raise ValueError( + f"Attention1D {field_name} {(grid.x, grid.y)} exceeds compute grid " + f"{(compute_grid.x, compute_grid.y)}" + ) + + return _resolve_attention1d_config( + common, + _architecture=architecture, + _architecture_fields=architecture_fields, + _resolved_fields={**slots, **architecture_fields}, + ) + + +def _resolve_attention1d_config( + config: Attention1DConfig, + *, + _architecture: str | None = None, + _architecture_fields: dict[str, object] | None = None, + _resolved_fields: dict[str, object] | None = None, +) -> Attention1DConfig: + """Materialize the config with sensible defaults.""" + to_set = dict(_resolved_fields or {}) + + # --- Phase 1: Model dimensions (fail-fast validation, no device needed) --- + # n_heads, n_kv_heads, head_dim MUST be provided - they cannot be reliably inferred + # from weight shapes alone (different models have different configurations). + + n_heads = config.n_heads + n_kv_heads = config.n_kv_heads + head_dim = config.head_dim + + if n_heads is None: + raise ValueError( + "n_heads must be provided. It cannot be reliably inferred from weights. " + "Get this value from your model's config.json (num_attention_heads)." + ) + if n_kv_heads is None: + raise ValueError( + "n_kv_heads must be provided. It cannot be reliably inferred from weights. " + "Get this value from your model's config.json (num_key_value_heads)." + ) + if head_dim is None: + raise ValueError( + "head_dim must be provided. It cannot be reliably inferred from weights. " + "Typically: head_dim = hidden_size // num_attention_heads." + ) + + # Reject sliding_window + paged attention (chunked prefill doesn't support window masking) + if config.sliding_window is not None and config.paged_attention_config is not None: + raise ValueError( + "Chunked prefill with sliding_window attention is not supported. " + "chunked_scaled_dot_product_attention does not implement sliding window masking. " + "Set sliding_window=None or disable paged attention / chunked prefill." + ) + + # --- Phase 2: Device and foundational fields --- + + # Derive mesh_device + mesh_device = _attention_mesh_device(config) + if config.mesh_device is None: + to_set["mesh_device"] = mesh_device + + assert mesh_device is not None, "mesh_device must be available!" + if _architecture is None: + raise ValueError("Attention1D internal config resolution requires a construction-time architecture") + architecture = _architecture + is_bh = architecture == "blackhole" + if _architecture_fields is None: + if is_bh: + raise ValueError("Blackhole Attention1D requires explicit SKU overlay fields") + architecture_fields = _wormhole_attention_architecture_fields(mesh_device) + else: + architecture_fields = _architecture_fields + + num_devices = mesh_device.get_num_devices() + + if num_devices > 1: + if n_heads % num_devices != 0: + raise ValueError( + f"n_heads ({n_heads}) must be divisible by mesh device count ({num_devices}) " + "for Attention1D tensor-parallel sharding." + ) + if n_kv_heads % num_devices != 0: + raise ValueError( + f"n_kv_heads ({n_kv_heads}) must be divisible by mesh device count ({num_devices}) " + "for Attention1D tensor-parallel sharding." + ) + + # Derive tt_ccl + if config.tt_ccl is None and num_devices > 1: + to_set["tt_ccl"] = get_tt_ccl(mesh_device) + + # Auto-detect topology + if config.topology is None: + to_set["topology"] = default_topology(mesh_device) + + topology = to_set.get("topology", config.topology) + + # Auto-detect num_links (same approach as TTTv1 ccl.py tt_all_reduce) + tt_ccl = to_set.get("tt_ccl", config.tt_ccl) + if config.num_reduce_scatter_links is None and num_devices > 1: + to_set["num_reduce_scatter_links"] = tt_ccl.get_num_links() + if config.num_all_gather_links is None and num_devices > 1: + to_set["num_all_gather_links"] = tt_ccl.get_num_links() + + # --- Phase 3: Derived dimensions --- + + # dim CAN be reliably inferred from weight shapes + dim = config.dim + if dim is None: + # wqkv shape is (1, 1, dim, qkv_size_per_device * num_devices) or (dim, qkv_size) + wqkv_shape = config.wqkv.source.shape + dim = wqkv_shape[-2] if len(wqkv_shape) == 4 else wqkv_shape[0] + to_set["dim"] = dim + + # qkv_size is derived from the required dimensions + qkv_size = config.qkv_size + if qkv_size is None: + qkv_size = head_dim * (2 * n_kv_heads + n_heads) + to_set["qkv_size"] = qkv_size + + if config.scale is None: + to_set["scale"] = head_dim**-0.5 + + if config.min_kv_prefill_shard_seqlen is None: + to_set["min_kv_prefill_shard_seqlen"] = (TILE_SIZE * 8 * 8) // (n_kv_heads // num_devices) + + # --- Phase 4: Dtypes --- + + if config.wqkv_dtype is None: + to_set["wqkv_dtype"] = ttnn.bfloat8_b + if config.wo_dtype is None: + to_set["wo_dtype"] = ttnn.bfloat8_b + if config.activation_dtype is None: + to_set["activation_dtype"] = ttnn.bfloat16 + + # --- Phase 5: Program configs --- + + tile_size = TILE_SIZE + tile_padded_batch_rows = tile_size * math.ceil(config.max_batch_size / tile_size) + n_local_heads = n_heads // num_devices + + # Decode configs - use DRAM-sharded matmul for all 1D topologies (matches TTTv1) + if config.decode_xqkv_prg_config is None: + attn_input_grid = _dram_shard_core_grid(dim) + to_set["decode_xqkv_prg_config"] = _dram_matmul_config( + m=tile_padded_batch_rows, + k=dim, + n=qkv_size // num_devices, + num_cores=attn_input_grid.num_cores, + ) + + if config.decode_sdpa_prg_config is None: + to_set["decode_sdpa_prg_config"] = ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=(8, 8), + exp_approx_mode=False, + q_chunk_size=0, + k_chunk_size=0, + ) + + if config.decode_attn_output_prg_config is None: + # Use DRAM-sharded matmul for all 1D topologies (matches TTTv1) + to_set["decode_attn_output_prg_config"] = _dram_matmul_config( + m=tile_padded_batch_rows, + k=(n_heads * head_dim) // num_devices, + n=dim, + num_cores=n_local_heads, + ) + + if config.decode_residual_memcfg is None: + residual_grid = _dram_shard_core_grid(dim // num_devices) + to_set["decode_residual_memcfg"] = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // residual_grid.num_cores // num_devices), + residual_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + if config.decode_create_qkv_head_memcfg is None: + to_set["decode_create_qkv_head_memcfg"] = ( + ttnn.create_sharded_memory_config( + shape=(TILE_SIZE, head_dim), + core_grid=architecture_fields["decode_create_qkv_head_grid"], + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + if is_bh + else ttnn.L1_HEIGHT_SHARDED_MEMORY_CONFIG + ) + + if config.decode_scores_memcfg is None: + + def scores_memcfg(batch_size): + return ttnn.create_sharded_memory_config( + shape=(math.ceil(n_local_heads / 32) * 32, head_dim), + core_grid=ttnn.CoreRangeSet({_num_to_corerange(batch_size)}), + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + to_set["decode_scores_memcfg"] = scores_memcfg + + # Prefill configs + # DRAM shard grid width: on Wormhole always 8 (despite 12 physical DRAM cores); + # on Blackhole use actual DRAM grid width (7 for P100, 8 for P150). + # Matching per_core_N to this width avoids silent PCC issues on P100. + dram_shard_grid_width = architecture_fields["dram_shard_grid_width"] + qkv_prefill_grid = architecture_fields["prefill_qkv_grid"] + + if config.prefill_xqkv_prg_config is None: + + @lru_cache + def xqkv_prefill_prg_config(seq_len: int): + return ttnn.MatmulMultiCoreReuseMultiCastProgramConfig( + compute_with_storage_grid_size=qkv_prefill_grid, + in0_block_w=1, + out_subblock_h=1, + out_subblock_w=1, + per_core_M=max(1, 8 if seq_len >= MAX_QKV_MM_SEQ_LEN else math.ceil(seq_len / tile_size / 8)), + per_core_N=math.ceil(qkv_size / num_devices / 32 / dram_shard_grid_width), + transpose_mcast=False, + fused_activation=None, + fuse_batch=seq_len <= MAX_QKV_MM_SEQ_LEN, + ) + + to_set["prefill_xqkv_prg_config"] = xqkv_prefill_prg_config + + # minimal_matmul config for QKV (only materialized when the opt-in is set). Block sizes and grid + # mirror TTTv1 (model_config.py:1626-1632): a full 8x8 (8x10 on BH) compute grid, 8/8/8 blocks. + if config.prefill_qkv_minimal_matmul and config.prefill_xqkv_minimal_matmul_config is None: + minimal_qkv_grid = ttnn.CoreCoord(*qkv_prefill_grid) + + @lru_cache + def xqkv_minimal_matmul_config(seq_len: int): + return ttnn.MinimalMatmulConfig( + M_block_size=8, + K_block_size=8, + N_block_size=8, + compute_with_storage_grid_size=minimal_qkv_grid, + ) + + to_set["prefill_xqkv_minimal_matmul_config"] = xqkv_minimal_matmul_config + + # minimal_matmul config for WO (only materialized when the opt-in is set). Same block sizes / grid + # as the QKV minimal config — the folded WO matmul is the same shape class (full-K activation @ a + # PlacementShard(-1) column-sharded weight), so it tiles identically. + if config.prefill_wo_minimal_matmul and config.prefill_wo_minimal_matmul_config is None: + minimal_wo_grid = ttnn.CoreCoord(*qkv_prefill_grid) + + @lru_cache + def wo_minimal_matmul_config(seq_len: int): + return ttnn.MinimalMatmulConfig( + M_block_size=8, + K_block_size=8, + N_block_size=8, + compute_with_storage_grid_size=minimal_wo_grid, + ) + + to_set["prefill_wo_minimal_matmul_config"] = wo_minimal_matmul_config + + if config.prefill_sdpa_prg_config is None: + + @lru_cache + def sdpa_prg_config(seq_len: int, chunk_start_idx: int | None): + q_chunk = 256 if seq_len >= 2048 else 64 + k_chunk = 256 if seq_len >= 2048 else 64 + + if chunk_start_idx is not None and chunk_start_idx != 0: + q_chunk = min(q_chunk, chunk_start_idx & -chunk_start_idx) + k_chunk = min(k_chunk, chunk_start_idx & -chunk_start_idx) + + return ttnn.SDPAProgramConfig( + compute_with_storage_grid_size=(8, 8), + exp_approx_mode=False, + q_chunk_size=q_chunk, + k_chunk_size=k_chunk, + ) + + to_set["prefill_sdpa_prg_config"] = sdpa_prg_config + + # Resolve the WO-matmul prefill M-chunk cutoff (None → MAX_MM_SEQ_LEN = TTTv1 parity) and write it + # back, so prefill_forward's STAGE-12/15 reshape and the prg config below agree. Only the M-tiling + # reads it; the fused all_gather_matmul n_dim stays on MAX_MM_SEQ_LEN (decoupled — see the field doc). + wo_prefill_len_cutoff = config.wo_prefill_len_cutoff + if wo_prefill_len_cutoff is None: + wo_prefill_len_cutoff = MAX_MM_SEQ_LEN + to_set["wo_prefill_len_cutoff"] = wo_prefill_len_cutoff + + if config.prefill_wo_prg_config is None: + use_fused = to_set.get("use_fused_all_gather_matmul", config.use_fused_all_gather_matmul) + if use_fused is None: + use_fused = ( + num_devices == 8 + and (dim // tile_size // num_devices) % num_devices == 0 + and num_devices > 1 + and topology == ttnn.Topology.Ring + ) + to_set["use_fused_all_gather_matmul"] = use_fused + + k_dim = (n_heads * head_dim) // num_devices + # n_dim intentionally stays on the MAX_MM_SEQ_LEN constant (NOT wo_prefill_len_cutoff): it drives + # the fused all_gather_matmul output width / sharding, which must be untouched by the M-cutoff. + n_dim = MAX_MM_SEQ_LEN if use_fused and MAX_MM_SEQ_LEN % (dim // num_devices) == 0 else dim + prefill_rows = 8 + + @lru_cache + def wo_prefill_prg_config(seq_len: int): + num_rows = min(seq_len, wo_prefill_len_cutoff) + grid_size = _find_prefill_grid(prefill_rows, k_dim // tile_size) + return _matmul_config( + m=num_rows, + k=k_dim, + n=n_dim, + grid_size=grid_size, + in0_block_w=1, + fuse_batch=seq_len <= wo_prefill_len_cutoff, + per_core_n=math.ceil(n_dim / (tile_size * dram_shard_grid_width)) if not use_fused else None, + ) + + to_set["prefill_wo_prg_config"] = wo_prefill_prg_config + + if config.prefill_kv_memcfg is None: + n_local_kv_heads = n_kv_heads // num_devices + + @lru_cache + def kv_prefill_memcfg(seq_len: int): + # Shard the combined (n_kv_heads * seq_len) dimension across 64 cores + # This ensures tile-aligned shard shapes (multiples of 32) + num_cores = 8 * 8 # 64 cores + shard_height = (n_local_kv_heads * seq_len) // num_cores + return ttnn.create_sharded_memory_config( + shape=(shard_height, head_dim), + core_grid=ttnn.CoreGrid(y=8, x=8), + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + to_set["prefill_kv_memcfg"] = kv_prefill_memcfg + + # Fused all-gather matmul configs + if config.use_fused_all_gather_matmul is None: + use_fused = ( + num_devices == 8 + and (dim // tile_size // num_devices) % num_devices == 0 + and num_devices > 1 + and topology == ttnn.Topology.Ring + ) + to_set["use_fused_all_gather_matmul"] = use_fused + + use_fused = to_set.get("use_fused_all_gather_matmul", config.use_fused_all_gather_matmul) + + if use_fused and config.decode_all_gather_matmul_prg_config is None: + do_core_grid_size = (8, 1) + do_per_core_N = dim // num_devices // tile_size // (do_core_grid_size[0] * do_core_grid_size[1]) + + to_set["decode_all_gather_matmul_prg_config"] = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig( + compute_with_storage_grid_size=do_core_grid_size, + in0_block_w=dim // tile_size // (do_core_grid_size[0] * do_core_grid_size[1]), + out_subblock_h=1, + out_subblock_w=get_out_subblock_w(do_per_core_N, out_subblock_h=1), + per_core_M=tile_padded_batch_rows // tile_size, + per_core_N=do_per_core_N, + fuse_batch=True, + fused_activation=None, + mcast_in0=True, + ) + + if use_fused and config.decode_all_gather_matmul_memcfg is None: + to_set["decode_all_gather_matmul_memcfg"] = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // 8), + ttnn.CoreGrid(x=8, y=1), + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + # Separate all-gather memory config (non-Ring or non-fused path) + # Matches TTTv1 GATHER_USERS_MEMCFG: shape=(tile_size * mesh_cols, dim // 8 // users_core_grid.num_cores) + if not use_fused and config.decode_gather_users_memcfg is None and num_devices > 1: + mesh_cols = list(mesh_device.shape)[1] + users_core_grid = _dram_shard_core_grid(dim // 8) + to_set["decode_gather_users_memcfg"] = ttnn.create_sharded_memory_config( + (tile_size * mesh_cols, dim // 8 // users_core_grid.num_cores), + users_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + # --- Phase 7: Input/output memory configs --- + + if config.decode_input_memcfg is None: + # All 1D topologies use WIDTH_SHARDED for DRAM-sharded matmul (matches TTTv1) + attn_input_grid = _dram_shard_core_grid(dim) + to_set["decode_input_memcfg"] = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // attn_input_grid.num_cores), + attn_input_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + if config.prefill_input_memcfg is None: + to_set["prefill_input_memcfg"] = ttnn.DRAM_MEMORY_CONFIG + + # --- Phase 8: Resolve LazyWeights --- + + wqkv_dtype = to_set.get("wqkv_dtype", config.wqkv_dtype) + wo_dtype = to_set.get("wo_dtype", config.wo_dtype) + + wqkv_mem_config = config.wqkv_memcfg + if wqkv_mem_config is None: + # Use DRAM-sharded config for all 1D topologies (matches TTTv1 behavior) + # Note: For multi-device, the per-device portion uses qkv_size/num_devices + dram_grid_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_grid_size.x - 1, dram_grid_size.y - 1))} + ) + # For multi-device, each device has qkv_size/num_devices columns + per_device_qkv_size = qkv_size // num_devices + wqkv_mem_config = _create_dram_sharded_mem_config( + k=dim, n=per_device_qkv_size, dram_grid=dram_grid, dram_cores=dram_grid_size.x + ) + + to_set["wqkv"] = resolve_lazy_weight( + config.wqkv, + device=mesh_device, + memory_config=wqkv_mem_config, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + dtype=wqkv_dtype, + ) + + wo_mem_config = config.wo_memcfg + if wo_mem_config is None: + # Fused all-gather matmul needs INTERLEAVED (TTTv1: dims=(2,3) for fused) + # Non-fused uses DRAM-sharded for all 1D topologies (TTTv1 behavior) + if use_fused: + wo_mem_config = ttnn.DRAM_MEMORY_CONFIG + else: + dram_grid_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_grid_size.x - 1, dram_grid_size.y - 1))} + ) + # For multi-device, each device has (n_heads * head_dim) / num_devices input dim + per_device_hidden = (n_heads * head_dim) // num_devices + wo_mem_config = _create_dram_sharded_mem_config( + k=per_device_hidden, n=dim, dram_grid=dram_grid, dram_cores=dram_grid_size.x + ) + + to_set["wo"] = resolve_lazy_weight( + config.wo, + device=mesh_device, + memory_config=wo_mem_config, + mesh_mapper_config=ttnn.MeshMapperConfig( + # For fused: shard on dim -1 (width) - matches TTTv1 dims=(2, 3) + # For non-fused: shard on dim -2 (height) - matches TTTv1 dims=(3, 2) + placements=[ttnn.PlacementShard(-1 if use_fused else -2)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + dtype=wo_dtype, + ) + + # Resolve Q/K norm configs if present + # RMSNorm1DConfig needs mesh_device propagated from attention config + if config.q_norm_config is not None: + q_norm_cfg = config.q_norm_config + if q_norm_cfg.mesh_device is None: + q_norm_cfg = replace(q_norm_cfg, mesh_device=mesh_device) + to_set["q_norm_config"] = q_norm_cfg + + if config.k_norm_config is not None: + k_norm_cfg = config.k_norm_config + if k_norm_cfg.mesh_device is None: + k_norm_cfg = replace(k_norm_cfg, mesh_device=mesh_device) + to_set["k_norm_config"] = k_norm_cfg + + # --- Phase 9: Handle QKV bias --- + + if config.wqkv_bias is not None: + # Extract source tensor from LazyWeight for custom transformation + qkv_bias = config.wqkv_bias.source + + # Mesh mapper config for sharding bias on last dimension + # For 1D mesh: [Replicate on dim 0, Shard on dim -1] + bias_mesh_mapper_config = ( + ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape(1, num_devices), + ) + if num_devices > 1 + else None # Single device: replicate (no sharding) + ) + + # Prefill bias: transform to 4D shape (1, 1, 1, qkv_size) + # Use replace() to create new LazyWeight with transformed source, then resolve_lazy_weight + prefill_bias_lazy = resolve_lazy_weight( + replace(config.wqkv_bias, source=qkv_bias.unsqueeze(0).unsqueeze(0).unsqueeze(0)), + device=mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=bias_mesh_mapper_config, + ) + to_set["_wqkv_bias_prefill"] = prefill_bias_lazy + + # Decode bias - one per batch size multiple + # Match TTTv1 pattern: 2D tensor (batch_size, qkv_size) + wqkv_bias_decode = [] + for batch_size in range(tile_size, tile_padded_batch_rows + tile_size, tile_size): + decode_bias_lazy = resolve_lazy_weight( + replace(config.wqkv_bias, source=qkv_bias.unsqueeze(0).expand(batch_size, -1)), + device=mesh_device, + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=bias_mesh_mapper_config, + ) + wqkv_bias_decode.append(decode_bias_lazy) + to_set["_wqkv_bias_decode"] = wqkv_bias_decode + + # --- Phase 10: Create transformation matrices for rotary embedding --- + + if config.transformation_mat_decode is None: + use_qk_fused = config.use_qk_fused + max_batch_size = config.max_batch_size + doubled_batch_size = max_batch_size * 2 if use_qk_fused else max_batch_size + + # Get core grid for batch distribution + core_grid = architecture_fields["decode_transformation_core_grid"] + batch_grid = ttnn.num_cores_to_corerangeset(doubled_batch_size, core_grid, row_wise=True) + + # Create transformation matrix repeated across batch cores + trans_mat = get_rot_transformation_mat().repeat(1, 1, doubled_batch_size, 1) + trans_mat_mem_config = ttnn.create_sharded_memory_config( + shape=(TILE_SIZE, TILE_SIZE), + core_grid=batch_grid, + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + transformation_mat_decode = ttnn.from_torch( + trans_mat, + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + dtype=ttnn.bfloat16, + memory_config=trans_mat_mem_config, + mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device), + ) + to_set["transformation_mat_decode"] = transformation_mat_decode + + if config.transformation_mat_prefill is None: + # Prefill uses simpler DRAM config (replicated across devices) + prefill_trans_mat = get_rot_transformation_mat() + transformation_mat_prefill = ttnn.from_torch( + prefill_trans_mat, + device=mesh_device, + layout=ttnn.TILE_LAYOUT, + dtype=ttnn.bfloat16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device), + ) + to_set["transformation_mat_prefill"] = transformation_mat_prefill + + # --- Phase 11: Resolve KV cache --- + # KV cache is a static configuration, allocated once and reused for all forward calls. + # If use_paged_kv_cache=True, the cache is managed externally (e.g., by vLLM) and not created here. + # Attention1D does not validate KV-cache DRAM/token capacity; callers own sizing. + # Allocation will fail at device memory allocation time if the cache does not fit. + if not config.use_vllm_paged_kv_cache: + n_local_kv_heads = n_kv_heads // num_devices + kv_cache_dtype = config.kv_cache_dtype + kv_cache_defaults = dict( + device=mesh_device, + dtype=kv_cache_dtype, + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate()], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + ) + + if config.kv_cache is None: + # Create default kv_cache LazyWeights + if config.paged_attention_config: + cache_k = zeros_like_paged_cache(config.paged_attention_config, n_local_kv_heads, head_dim) + cache_v = zeros_like_paged_cache(config.paged_attention_config, n_local_kv_heads, head_dim) + else: + cache_k = zeros_like_kv_cache(config.max_batch_size, n_local_kv_heads, config.max_seq_len, head_dim) + cache_v = zeros_like_kv_cache(config.max_batch_size, n_local_kv_heads, config.max_seq_len, head_dim) + kv_cache = (LazyWeight(source=cache_k), LazyWeight(source=cache_v)) + else: + kv_cache = config.kv_cache + + # Resolve defaults for LazyWeights; pass through pre-allocated tensors as-is + resolved = [] + for i, entry in enumerate(kv_cache): + if isinstance(entry, LazyWeight): + resolved.append(resolve_lazy_weight(entry, **kv_cache_defaults)) + elif isinstance(entry, ttnn.Tensor): + resolved.append(entry) + else: + raise TypeError(f"kv_cache[{i}] must be LazyWeight or ttnn.Tensor, got {type(entry).__name__}") + to_set["kv_cache"] = tuple(resolved) + + return replace(config, **to_set) + + +# ============================================================================= +# Helper functions +# ============================================================================= + + +def _find_largest_divisor(n: int, max_divisor: int = 8) -> int: + """Find largest divisor of n up to max_divisor.""" + for i in range(max_divisor, 0, -1): + if n % i == 0: + return i + return 1 + + +def _find_grid(n_tiles: int, max_rows: int = 8, max_cols: int = 8) -> tuple[int, int]: + """Find grid dimensions (rows, cols) that evenly divide n_tiles.""" + max_cores = max_rows * max_cols + target = 32 + possible_cores = [k for k in range(1, max_cores + 1) if n_tiles % k == 0] + possible_cores.sort(key=lambda x: abs(x - target)) + + for cores in possible_cores: + for rows in range(1, max_rows + 1): + if cores % rows == 0: + cols = cores // rows + if cols <= max_cols: + return rows, cols + + raise AssertionError(f"Cannot find grid for {n_tiles} tiles within {max_rows}x{max_cols}") + + +def _find_prefill_grid(row_tiles: int, col_tiles: int, max_rows: int = 8, max_cols: int = 8) -> tuple[int, int]: + """Find grid where row_tiles divides rows and col_tiles divides cols.""" + cols = next((i for i in range(max_cols, 0, -1) if col_tiles % i == 0), None) + rows = next((i for i in range(max_rows, 0, -1) if row_tiles % i == 0), None) + assert cols is not None and rows is not None + return rows, cols + + +def _dram_shard_core_grid(k: int, tile_size: int = TILE_SIZE) -> ttnn.CoreGrid: + """Get core grid for DRAM sharding based on K dimension.""" + rows, cols = _find_grid(k // tile_size) + return ttnn.CoreGrid(x=cols, y=rows) + + +def _dram_matmul_config( + m: int, k: int, n: int, num_cores: int, tile_size: int = TILE_SIZE +) -> ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig: + """Create DRAM-sharded matmul program config.""" + return ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig( + in0_block_w=_find_largest_divisor(k // (tile_size * num_cores)), + per_core_M=math.ceil(m / tile_size), + per_core_N=math.ceil(n / (tile_size * num_cores)), + fused_activation=None, + ) + + +def _matmul_config( + m: int, + k: int, + n: int, + grid_size: tuple[int, int], + tile_size: int = TILE_SIZE, + in0_block_w: int = None, + fuse_batch: bool = False, + fused_activation=None, + per_core_m: int = None, + per_core_n: int = None, +) -> ttnn.MatmulMultiCoreReuseMultiCastProgramConfig: + """Create multicast matmul program config.""" + if per_core_m is None: + per_core_m = math.ceil(m / (tile_size * grid_size[1])) + if per_core_n is None: + per_core_n = math.ceil(n / (tile_size * grid_size[0])) + + out_subblock_h = 1 + out_subblock_w = get_out_subblock_w(per_core_n, out_subblock_h) + + if in0_block_w is None: + in0_block_w = _find_largest_divisor(k // (tile_size * grid_size[1])) + + return ttnn.MatmulMultiCoreReuseMultiCastProgramConfig( + compute_with_storage_grid_size=grid_size, + in0_block_w=in0_block_w, + out_subblock_h=out_subblock_h, + out_subblock_w=out_subblock_w, + per_core_M=per_core_m, + per_core_N=per_core_n, + transpose_mcast=False, + fused_activation=fused_activation, + fuse_batch=fuse_batch, + ) + + +def _create_dram_sharded_mem_config( + k: int, n: int, dram_grid: ttnn.CoreRangeSet, tile_size: int = TILE_SIZE, dram_cores: int = 12 +) -> ttnn.MemoryConfig: + """Create DRAM-sharded memory config for weight tensors.""" + padded_size = math.ceil(n / (tile_size * dram_cores)) * (tile_size * dram_cores) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_size // dram_cores), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + +def _num_to_corerange(num_cores: int, start_core: ttnn.CoreCoord = None) -> ttnn.CoreRange: + """Convert number of cores to CoreRange.""" + if start_core is None: + start_core = ttnn.CoreCoord(0, 0) + + if num_cores == 1: + return ttnn.CoreRange(start_core, start_core) + + # Arrange cores in rows of 8 + row_size = 8 + start_x, start_y = start_core.x, start_core.y + + # Calculate end coordinates + total_cores_with_start = start_x + num_cores + end_y = start_y + (total_cores_with_start - 1) // row_size + end_x = (total_cores_with_start - 1) % row_size + + return ttnn.CoreRange(start_core, ttnn.CoreCoord(end_x, end_y)) + + +def _load_input_device_tensor(x: ttnn.Tensor | LazyWeight, config: Attention1DConfig, mode: str) -> ttnn.Tensor: + """Resolve input tensor to ttnn.Tensor if LazyWeight, otherwise return as-is.""" + assert mode in ("decode", "prefill"), f"mode must be 'decode' or 'prefill', got {mode}" + + if isinstance(x, LazyWeight): + mem_cfg = config.decode_input_memcfg if mode == "decode" else config.prefill_input_memcfg + resolved_x = resolve_lazy_weight( + x, + device=config.mesh_device, + memory_config=mem_cfg, + mesh_mapper_config=None, # replicated + layout=ttnn.TILE_LAYOUT, + ) + return resolved_x.get_device_weight() + + assert isinstance(x, ttnn.Tensor), f"x must be ttnn.Tensor or LazyWeight, got {type(x)}" + return x diff --git a/code/models/common/modules/embedding/embedding_1d.py b/code/models/common/modules/embedding/embedding_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..37940e756c1edc7e896a66289641bcdab44632db --- /dev/null +++ b/code/models/common/modules/embedding/embedding_1d.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2-style Embedding module for 1D-topology devices: N150 (1x1), N300 (1x2), T3K (1x8). + +Single unified Embedding1D class: + - forward(x): Token embedding lookup, optionally scaled by embed_scale. + +Execution path: + embedding(x, weights) → [multiply(embed_scale) if embed_scale != 1.0] + +This module replaces both TTTv1 Embedding and ScaledEmbedding classes. +""" + +from dataclasses import dataclass, replace +from pathlib import Path + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight + +# ============================================================================= +# Top-level config dataclass +# ============================================================================= + + +@dataclass +class Embedding1DConfig: + """ + Central configuration for Embedding1D - the single source of truth for all settings. + + Simple usage (all defaults): + config = Embedding1DConfig(weights=lazy_weights) + + With scale: + config = Embedding1DConfig(weights=lazy_weights, embed_scale=3072**0.5) + + Full customization: + config = Embedding1DConfig( + weights=lazy_weights, + mesh_device=custom_device, + weights_memcfg=custom_memcfg, + ) + """ + + # Required: embedding weight table (vocab_size, dim), sharded on dim=-1 + weights: LazyWeight + + # Optional: device (derived from weights if None) + mesh_device: ttnn.MeshDevice | None = None + + # Optional: scaling factor applied after lookup (1.0 = no scaling) + embed_scale: float = 1.0 + + # Optional: power-user overrides (None = compute defaults) + weights_dtype: ttnn.DataType | None = None + weights_memcfg: ttnn.MemoryConfig | None = None + output_memcfg: ttnn.MemoryConfig | None = None + + def is_resolved(self) -> bool: + """Check if all fields are resolved.""" + return all(getattr(self, f) is not None for f in self.__dataclass_fields__) + + +# ============================================================================= +# Embedding1D - Unified Embedding for 1D-topology devices +# ============================================================================= + + +class Embedding1D(LightweightModule): + """ + Embedding for non-TG devices supporting token lookup with optional scaling. + + Replaces both TTTv1 Embedding and ScaledEmbedding. + + Simple API (90% of users): + emb = Embedding1D(weights) + + With scaling: + emb = Embedding1D(weights, embed_scale=3072**0.5) + + Power API (10% of users): + config = Embedding1DConfig(weights=lazy_w, weights_memcfg=custom_cfg) + emb = Embedding1D.from_config(config) + + Execution path: + embedding(x, weights) → [multiply(embed_scale) if embed_scale != 1.0] + """ + + def __init__(self, weights: LazyWeight, embed_scale: float = 1.0): + """ + Simple API - derives all config from weights. + + Args: + weights: Embedding weight table (vocab_size, dim), sharded on last dim. + embed_scale: Scale factor applied after lookup (default 1.0 = no scaling). + """ + super().__init__() + self.config = _resolve_embedding1d_config(Embedding1DConfig(weights=weights, embed_scale=embed_scale)) + self._device_weights_loaded = False + + @classmethod + def from_config(cls, config: Embedding1DConfig): + """ + Power API - any level of customization via config. + + Override any subset of fields in Embedding1DConfig: + config = Embedding1DConfig(weights=w, embed_scale=2.0) + emb = Embedding1D.from_config(config) + """ + instance = object.__new__(cls) + super(Embedding1D, instance).__init__() + instance.config = _resolve_embedding1d_config(config) + instance._device_weights_loaded = False + return instance + + def load_device_weights(self): + if self._device_weights_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device weights!" + self.weights = self.config.weights.get_device_weight() + self._device_weights_loaded = True + + def forward(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Token embedding lookup. + + Args: + x: Token IDs tensor (uint32), shape [1, 1, 1, seq_len]. + + Returns: + Embedded tokens, shape [1, 1, seq_len, dim]. + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config) + + out = ttnn.embedding( + x, + self.weights, + layout=ttnn.TILE_LAYOUT, + memory_config=self.config.output_memcfg, + ) + + if self.config.embed_scale != 1.0: + out = ttnn.multiply(out, self.config.embed_scale, memory_config=self.config.output_memcfg) + + return out + + # [INFO] this is the entry point for TTTv1 model_config.py and will retire with TTTv1 + @classmethod + def from_model_args( + cls, + mesh_device, + args, + weight_cache_path, + state_dict, + dtype, + embed_scale: float = 1.0, + ): + """Factory method for backward compatibility with ModelArgs. + + Args: + mesh_device: The mesh device to use. + args: Model arguments (ModelArgs instance). + weight_cache_path: Path for weight caching. + state_dict: The state dictionary containing weights. + dtype: Data type for weights. + embed_scale: Scale factor applied after lookup (default 1.0). + """ + if args.is_galaxy: + raise ValueError("Embedding1D cannot be used for Galaxy devices.") + + base_name = args.get_state_dict_prefix("", None) + "tok_embeddings.weight" + torch_weight = state_dict[base_name].unsqueeze(0).unsqueeze(0) + + cache_dir = None if args.dummy_weights else Path(weight_cache_path) / "embedding" + + num_devices = mesh_device.get_num_devices() + + weights = LazyWeight( + source=torch_weight, + dtype=dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(cache_dir, "tok_embeddings") if cache_dir else None, + ) + + config = Embedding1DConfig( + weights=weights, + mesh_device=mesh_device, + embed_scale=embed_scale, + weights_dtype=dtype, + weights_memcfg=ttnn.DRAM_MEMORY_CONFIG, + output_memcfg=ttnn.DRAM_MEMORY_CONFIG, + ) + return cls.from_config(config) + + +# ============================================================================= +# Config resolution +# ============================================================================= + + +def _resolve_embedding1d_config(config: Embedding1DConfig) -> Embedding1DConfig: + """Materialize the config to known good defaults using replace pattern.""" + + to_set = {} + + # --- Phase 1: Foundational fields --- + + # Derive mesh_device from weights + mesh_device = config.mesh_device + if mesh_device is None: + mesh_device = config.weights.device + if mesh_device is None: + mesh_device = ttnn.GetDefaultDevice() + if config.mesh_device is None: + to_set["mesh_device"] = mesh_device + + assert mesh_device is not None, "mesh_device must be available at this point!" + + # --- Phase 2: Default configs --- + + if config.weights_dtype is None: + to_set["weights_dtype"] = ttnn.bfloat16 + + if config.weights_memcfg is None: + to_set["weights_memcfg"] = ttnn.DRAM_MEMORY_CONFIG + + if config.output_memcfg is None: + to_set["output_memcfg"] = ttnn.DRAM_MEMORY_CONFIG + + # --- Phase 3: Resolve LazyWeight --- + + num_devices = mesh_device.get_num_devices() + weights_memcfg = ( + config.weights_memcfg + if config.weights_memcfg is not None + else to_set.get("weights_memcfg", ttnn.DRAM_MEMORY_CONFIG) + ) + weights_dtype = ( + config.weights_dtype if config.weights_dtype is not None else to_set.get("weights_dtype", ttnn.bfloat16) + ) + + to_set["weights"] = resolve_lazy_weight( + config.weights, + device=mesh_device, + memory_config=weights_memcfg, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=weights_dtype, + ) + + resolved_config = replace(config, **to_set) + assert resolved_config.is_resolved(), "Config must be resolved!" + assert resolved_config.weights.is_resolved(), "Weights must be resolved!" + + return resolved_config + + +def _load_input_device_tensor(x: ttnn.Tensor | LazyWeight, config: Embedding1DConfig) -> ttnn.Tensor: + """Resolve the input tensor to ttnn tensor if x is a LazyWeight, otherwise return as-is.""" + if isinstance(x, LazyWeight): + resolved_x = resolve_lazy_weight( + x, + device=config.mesh_device, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=None, # replicated + layout=ttnn.ROW_MAJOR_LAYOUT, + ) + return resolved_x.get_device_weight() + + assert isinstance(x, ttnn.Tensor), "x must be a ttnn tensor at this point!" + return x diff --git a/code/models/common/modules/lm_head/lm_head_1d.py b/code/models/common/modules/lm_head/lm_head_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..235a584bd257eb2a67a9173b0d4f19555bf429c1 --- /dev/null +++ b/code/models/common/modules/lm_head/lm_head_1d.py @@ -0,0 +1,567 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2-style LM Head module for 1D-topology devices: N150 (1x1), N300 (1x2), T3K (1x8). + +Computes logits over the vocabulary by splitting the output projection into +weight chunks that fit in L1, running linear ops per chunk, and concatenating. + +Each device holds the correct logits for its vocab shard (column-parallel linear, +no partial sums). The caller handles all_gather for multi-device argmax. + +Execution path: + for each (weight, pc): linear(x, weight) → sharded_to_interleaved → append + → concat +""" + +import math +from dataclasses import dataclass, replace +from pathlib import Path +from typing import List + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight +from models.common.tensor_utils import TILE_SIZE +from models.common.tensor_utils import nearest_32 as _nearest_32 + +# ============================================================================= +# Config dataclass +# ============================================================================= + + +@dataclass +class LMHead1DConfig: + """ + Configuration for LMHead1D. + + Simple usage (pre-split weights): + config = LMHead1DConfig(output_weights=[w1, w2, w3]) + + The program_configs and other fields are auto-computed from weights if None. + """ + + # Required: output projection weights (already split for L1 fit) + output_weights: List[LazyWeight] + + # Optional: device + mesh_device: ttnn.MeshDevice | None = None + + # Optional: derived from weights if None + dim: int | None = None + + # Optional: batch/tile config + max_batch_size: int = 32 + + # Optional: power-user overrides + program_configs: List | None = None + compute_kernel_config: ttnn.DeviceComputeKernelConfig | None = None + # Logical output width for each split on one device. A physical weight shard + # may be tile-padded beyond this width; forward trims it before concat. + output_split_sizes: List[int] | None = None + lm_head_dtype: ttnn.DataType = ttnn.bfloat8_b + output_memcfg: ttnn.MemoryConfig | None = None + + # Input memory config (None = DRAM interleaved for simple API, width-sharded for from_model_args) + input_memcfg: ttnn.MemoryConfig | None = None + + # Weight memory configs (None = auto-compute) + weights_memcfgs: List[ttnn.MemoryConfig] | None = None + + def is_resolved(self) -> bool: + return all(getattr(self, f) is not None for f in self.__dataclass_fields__) + + +# ============================================================================= +# LMHead1D +# ============================================================================= + + +class LMHead1D(LightweightModule): + """ + LM Head for non-TG (1D) devices. + + Splits vocabulary projection into L1-sized chunks, runs linear per chunk, + and concatenates. Each device holds correct logits for its vocab shard; + caller handles all_gather for multi-device argmax. + + Simple API: + lm_head = LMHead1D(output_weights=[w1, w2]) + + Power API: + config = LMHead1DConfig(output_weights=[w1, w2], lm_head_dtype=ttnn.bfloat16) + lm_head = LMHead1D.from_config(config) + + Execution path: + for each (w, pc): linear(x, w) → sharded_to_interleaved → concat + """ + + def __init__(self, output_weights: List[LazyWeight]): + super().__init__() + self.config = resolve_lm_head_1d_arch_config(LMHead1DConfig(output_weights=output_weights)) + self._device_weights_loaded = False + + @classmethod + def from_config(cls, config: LMHead1DConfig): + instance = object.__new__(cls) + super(LMHead1D, instance).__init__() + instance.config = resolve_lm_head_1d_arch_config(config) + instance._device_weights_loaded = False + return instance + + def load_device_weights(self): + if self._device_weights_loaded: + return + self.output_weights = [w.get_device_weight() for w in self.config.output_weights] + self._device_weights_loaded = True + + def forward(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Compute logits over vocabulary. + + Args: + x: Input hidden states, shape [1, 1, batch_rows, dim]. + + Returns: + Logits tensor, shape [1, 1, batch_rows, padded_vocab / num_devices] per device. + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config) + cfg = self.config + + outputs = [] + for weight, pc, logical_width in zip(self.output_weights, cfg.program_configs, cfg.output_split_sizes): + if pc is not None: + # DRAM-sharded path (from_model_args): width-sharded output → interleaved + output = ttnn.linear( + x, + weight, + compute_kernel_config=cfg.compute_kernel_config, + program_config=pc, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + dtype=cfg.lm_head_dtype, + ) + output = ttnn.sharded_to_interleaved(output, memory_config=cfg.output_memcfg) + else: + # Auto path (simple API): interleaved output directly + output = ttnn.linear( + x, + weight, + compute_kernel_config=cfg.compute_kernel_config, + memory_config=cfg.output_memcfg, + dtype=cfg.lm_head_dtype, + ) + if output.shape[-1] != logical_width: + output = ttnn.slice( + output, + (0, 0, 0, 0), + (output.shape[0], output.shape[1], output.shape[2], logical_width), + ) + outputs.append(output) + + # Concatenate splits + output = ttnn.concat(outputs, dim=-1, memory_config=cfg.output_memcfg) + + return output + + # [INFO] this is the entry point for TTTv1 model_config.py and will retire with TTTv1 + @classmethod + def from_model_args( + cls, + mesh_device, + args, + state_dict, + state_dict_prefix, + weight_cache_path, + max_columns_per_device, + dtype=None, + model_config=None, + tt_ccl=None, + ): + """Factory method for backward compatibility with ModelArgs. + + Note: tt_ccl is accepted for signature compatibility with TTTv1 LMHead + but is not used -- 1D LMHead does not need CCL (column-parallel, no partial sums). + """ + if args.is_galaxy: + raise ValueError("LMHead1D cannot be used for Galaxy devices.") + + import torch + + vocab_size = args.vocab_size + num_devices = mesh_device.get_num_devices() + dim = args.dim + padded_vocab_size = math.ceil(vocab_size / (TILE_SIZE * num_devices)) * (TILE_SIZE * num_devices) + size_per_device = padded_vocab_size // num_devices + num_splits = math.ceil(size_per_device / max_columns_per_device) + split_sizes = [min(size_per_device, max_columns_per_device)] * (num_splits - 1) + split_sizes.append(size_per_device - sum(split_sizes)) + + # Build output weights + source_weight = state_dict[f"{state_dict_prefix}output.weight"] + if tuple(source_weight.shape) != (vocab_size, dim): + raise ValueError( + f"LMHead1D output weight must have shape {(vocab_size, dim)}, got {tuple(source_weight.shape)}" + ) + torch_output_weights = source_weight.permute(1, 0) + if vocab_size < padded_vocab_size: + padding_size = padded_vocab_size - vocab_size + torch_output_weights = torch.cat( + [ + torch_output_weights, + torch.zeros(torch_output_weights.shape[0], padding_size, dtype=torch_output_weights.dtype), + ], + dim=-1, + ) + + # DRAM grid for weight memory configs + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1))} + ) + + cache_dir = None if args.dummy_weights else Path(weight_cache_path) / "lm_head" + + output_weights = [] + weights_memcfgs = [] + for i, split_size in enumerate(split_sizes): + device_splits = [] + physical_split_size = math.ceil(split_size / TILE_SIZE) * TILE_SIZE + for device_idx in range(num_devices): + start = device_idx * size_per_device + sum(split_sizes[:i]) + end = start + split_size + device_split = torch_output_weights[:, start:end] + if split_size < physical_split_size: + device_split = torch.cat( + [ + device_split, + torch.zeros(dim, physical_split_size - split_size, dtype=device_split.dtype), + ], + dim=-1, + ) + device_splits.append(device_split) + combined_split = torch.cat(device_splits, dim=-1) + + mem_cfg = _create_dram_sharded_mem_config( + k=dim, + n=math.ceil(combined_split.shape[-1] / num_devices), + dram_grid=dram_grid, + tile_size=TILE_SIZE, + dram_cores=dram_size.x, + ) + weights_memcfgs.append(mem_cfg) + + w_dtype = dtype if dtype is not None else ttnn.bfloat8_b + output_weights.append( + LazyWeight( + source=combined_split, + dtype=w_dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=mem_cfg, + cache_dir_weight_name=( + (cache_dir, f"output_split_{i}_logical_{split_size}_physical_{combined_split.shape[-1]}") + if cache_dir + else None + ), + ) + ) + + # Program configs - use the args.lm_head_core_grid which TTTv1 carefully computes + tile_padded_batch_rows = TILE_SIZE * math.ceil(args.max_batch_size / TILE_SIZE) + lm_head_core_grid = args.lm_head_core_grid + program_configs = [ + args.dram_matmul_config( + tile_padded_batch_rows, + dim, + ss, + lm_head_core_grid.num_cores, + ) + for ss in split_sizes + ] + + # Get input memory config: width-sharded across lm_head_core_grid + # (DRAM-sharded matmul requires width-sharded input in L1) + input_memcfg = ttnn.create_sharded_memory_config( + ( + tile_padded_batch_rows, + _nearest_32(dim // lm_head_core_grid.num_cores), + ), + lm_head_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + config = LMHead1DConfig( + output_weights=output_weights, + mesh_device=mesh_device, + dim=dim, + max_batch_size=args.max_batch_size, + program_configs=program_configs, + output_split_sizes=split_sizes, + lm_head_dtype=getattr(args, "lm_head_dtype", ttnn.bfloat8_b), + output_memcfg=ttnn.L1_MEMORY_CONFIG, + input_memcfg=input_memcfg, + weights_memcfgs=weights_memcfgs, + ) + return cls.from_config(config) + + +# ============================================================================= +# Config resolution +# ============================================================================= + + +def _compute_kernel_config_hifi2(arch) -> ttnn.DeviceComputeKernelConfig: + """Construct the TTTv1 LM-head recipe for the selected architecture.""" + return ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + + +def _validate_lm_head_compute_kernel_config(compute_kernel_config) -> None: + if compute_kernel_config is None: + raise ValueError("LMHead1D architecture config requires compute_kernel_config") + if not isinstance(compute_kernel_config, ttnn.WormholeComputeKernelConfig): + raise TypeError( + "LMHead1D compute_kernel_config must be a concrete TTNN compute kernel config, " + f"got {type(compute_kernel_config).__name__}" + ) + + +def _derive_lm_head_mesh_device(config: LMHead1DConfig) -> ttnn.MeshDevice: + if not config.output_weights: + raise ValueError("LMHead1D requires at least one output weight") + mesh_device = config.mesh_device + if mesh_device is None: + mesh_device = config.output_weights[0].device + if mesh_device is None: + mesh_device = ttnn.GetDefaultDevice() + if mesh_device is None: + raise ValueError("LMHead1D mesh_device must be available") + for split_index, weight in enumerate(config.output_weights): + if weight.device is not None and weight.device is not mesh_device: + raise ValueError(f"LMHead1D mesh_device must match output weight {split_index} device") + return mesh_device + + +def _validate_lm_head_program_configs(config: LMHead1DConfig) -> None: + num_splits = len(config.output_weights) + if num_splits == 0: + raise ValueError("LMHead1D requires at least one output weight") + if config.program_configs is None or len(config.program_configs) != num_splits: + raise ValueError("LMHead1D requires exactly one program config per output-weight split") + if config.weights_memcfgs is None or len(config.weights_memcfgs) != num_splits: + raise ValueError("LMHead1D requires exactly one weight memory config per output-weight split") + if config.output_split_sizes is None or len(config.output_split_sizes) != num_splits: + raise ValueError("LMHead1D requires exactly one logical output width per output-weight split") + + num_devices = config.mesh_device.get_num_devices() + physical_widths = [] + for split_index, (weight, logical_width) in enumerate(zip(config.output_weights, config.output_split_sizes)): + source_shape = weight.source.shape + if len(source_shape) < 2 or source_shape[-2] != config.dim: + raise ValueError( + f"LMHead1D split {split_index} weight K dimension must equal dim={config.dim}, got {source_shape}" + ) + combined_width = source_shape[-1] + if combined_width % num_devices: + raise ValueError( + f"LMHead1D split {split_index} combined physical width {combined_width} " + f"must be divisible by {num_devices} devices" + ) + physical_width = combined_width // num_devices + if physical_width % TILE_SIZE: + raise ValueError( + f"LMHead1D split {split_index} physical per-device width {physical_width} must be tile aligned" + ) + if not isinstance(logical_width, int) or logical_width <= 0 or logical_width > physical_width: + raise ValueError( + f"LMHead1D split {split_index} logical width must satisfy 0 < logical <= {physical_width}, " + f"got {logical_width}" + ) + physical_widths.append(physical_width) + + has_explicit_program = any(program_config is not None for program_config in config.program_configs) + if has_explicit_program: + if config.input_memcfg is None or not config.input_memcfg.is_sharded(): + raise ValueError("LMHead1D DRAM-sharded programs require a sharded input_memcfg") + if config.input_memcfg.buffer_type != ttnn.BufferType.L1: + raise ValueError("LMHead1D DRAM-sharded programs require input_memcfg in L1") + if config.input_memcfg.memory_layout != ttnn.TensorMemoryLayout.WIDTH_SHARDED: + raise ValueError("LMHead1D DRAM-sharded programs require width-sharded input_memcfg") + num_compute_cores = config.input_memcfg.shard_spec.grid.num_cores() + + for split_index, program_config in enumerate(config.program_configs): + if program_config is None: + continue + if not isinstance(program_config, ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig): + raise TypeError( + "LMHead1D explicit program configs must use the DRAM-sharded matmul path; " + f"split {split_index} has {type(program_config).__name__}" + ) + if program_config.in0_block_w <= 0 or program_config.per_core_M <= 0 or program_config.per_core_N <= 0: + raise ValueError(f"LMHead1D split {split_index} program dimensions must be positive") + dim_tiles = config.dim // TILE_SIZE + if dim_tiles % program_config.in0_block_w != 0: + raise ValueError(f"LMHead1D split {split_index} dim tiles must be divisible by in0_block_w") + expected_per_core_m = math.ceil(config.max_batch_size / TILE_SIZE) + if program_config.per_core_M != expected_per_core_m: + raise ValueError( + f"LMHead1D split {split_index} per_core_M={program_config.per_core_M} " + f"does not match padded batch tiles {expected_per_core_m}" + ) + + weight_memcfg = config.weights_memcfgs[split_index] + if not weight_memcfg.is_sharded(): + raise ValueError(f"LMHead1D split {split_index} DRAM-sharded program requires sharded weight memory") + if weight_memcfg.buffer_type != ttnn.BufferType.DRAM: + raise ValueError(f"LMHead1D split {split_index} weight memory must reside in DRAM") + if weight_memcfg.memory_layout != ttnn.TensorMemoryLayout.WIDTH_SHARDED: + raise ValueError(f"LMHead1D split {split_index} weight memory must be width sharded") + covered_width = program_config.per_core_N * TILE_SIZE * num_compute_cores + if covered_width < physical_widths[split_index]: + raise ValueError( + f"LMHead1D split {split_index} program covers {covered_width} columns, " + f"below physical width {physical_widths[split_index]}" + ) + + +def resolve_lm_head_1d_arch_config(config: LMHead1DConfig) -> LMHead1DConfig: + """Return an independent, fully resolved config after one architecture query.""" + mesh_device = _derive_lm_head_mesh_device(config) + arch = mesh_device.arch() + + if arch not in (ttnn.device.Arch.WORMHOLE_B0, ttnn.device.Arch.BLACKHOLE): + raise ValueError(f"Unsupported LMHead1D architecture: {arch}") + requested_compute = config.compute_kernel_config or _compute_kernel_config_hifi2(arch) + try: + compute_kernel_config = ttnn.init_device_compute_kernel_config(arch, requested_compute) + except (TypeError, ValueError) as error: + raise ValueError(f"Invalid LMHead1D compute recipe for {arch}: {error}") from error + + _validate_lm_head_compute_kernel_config(compute_kernel_config) + resolved_common = _resolve_lm_head_1d_config(config) + resolved_common = replace(resolved_common, compute_kernel_config=compute_kernel_config) + _validate_lm_head_program_configs(resolved_common) + return resolved_common + + +def _resolve_lm_head_1d_config(config: LMHead1DConfig) -> LMHead1DConfig: + """Resolve defaults for LMHead1DConfig.""" + if not config.output_weights: + raise ValueError("LMHead1D requires at least one output weight") + to_set = {} + + # Mesh device + mesh_device = _derive_lm_head_mesh_device(config) + if config.mesh_device is None: + to_set["mesh_device"] = mesh_device + + # Dim + dim = config.dim + if dim is None: + dim = config.output_weights[0].source.shape[-2] + to_set["dim"] = dim + + # Output memcfg + if config.output_memcfg is None: + to_set["output_memcfg"] = ttnn.L1_MEMORY_CONFIG + + # Input memcfg + if config.input_memcfg is None: + to_set["input_memcfg"] = ttnn.DRAM_MEMORY_CONFIG + + # Program configs + num_devices = mesh_device.get_num_devices() + + if config.program_configs is None: + # Use None program configs for auto-resolve (let ttnn.linear auto-select). + # DRAM-sharded program configs require matching DRAM-sharded weight memory configs, + # which are only set up correctly via from_model_args. + pcs = [None for _ in config.output_weights] + to_set["program_configs"] = pcs + + if config.output_split_sizes is None: + split_sizes = [] + for split_index, weight in enumerate(config.output_weights): + combined_width = weight.source.shape[-1] + if combined_width % num_devices: + raise ValueError( + f"LMHead1D split {split_index} combined physical width {combined_width} " + f"must be divisible by {num_devices} devices" + ) + split_sizes.append(combined_width // num_devices) + to_set["output_split_sizes"] = split_sizes + elif len(config.output_split_sizes) != len(config.output_weights): + raise ValueError("LMHead1D requires exactly one logical output split size per output weight") + + # Weight memory configs + resolve LazyWeights + if config.weights_memcfgs is None: + # Use regular DRAM for auto-resolve (DRAM sharded requires DRAM-core-aligned padding + # which is handled by from_model_args when it explicitly provides weights_memcfgs) + memcfgs = [ttnn.DRAM_MEMORY_CONFIG for _ in config.output_weights] + to_set["weights_memcfgs"] = memcfgs + + weights_memcfgs = ( + config.weights_memcfgs if config.weights_memcfgs is not None else to_set.get("weights_memcfgs", []) + ) + + resolved_weights = [] + for i, w in enumerate(config.output_weights): + mem_cfg = weights_memcfgs[i] if i < len(weights_memcfgs) else ttnn.DRAM_MEMORY_CONFIG + resolved_weights.append( + resolve_lazy_weight( + w, + device=mesh_device, + memory_config=mem_cfg, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + dtype=config.lm_head_dtype, + ) + ) + to_set["output_weights"] = resolved_weights + + resolved = replace(config, **to_set) + return resolved + + +def _load_input_device_tensor(x: ttnn.Tensor | LazyWeight, config: LMHead1DConfig) -> ttnn.Tensor: + """Resolve input tensor using config's input_memcfg.""" + if isinstance(x, LazyWeight): + resolved_x = resolve_lazy_weight( + x, + device=config.mesh_device, + memory_config=config.input_memcfg, + mesh_mapper_config=None, + layout=ttnn.TILE_LAYOUT, + ) + return resolved_x.get_device_weight() + assert isinstance(x, ttnn.Tensor) + return x + + +# ============================================================================= +# Config helper functions (adapted from TTTv1 model_config.py) +# ============================================================================= + + +def _create_dram_sharded_mem_config( + k: int, n: int, dram_grid: ttnn.CoreRangeSet, tile_size: int = TILE_SIZE, dram_cores: int = 12 +) -> ttnn.MemoryConfig: + padded_size = math.ceil(n / (tile_size * dram_cores)) * (tile_size * dram_cores) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_size // dram_cores), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) diff --git a/code/models/common/modules/mlp/mlp_1d.py b/code/models/common/modules/mlp/mlp_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..87570219df81d10f956e957dea1323fae867ba24 --- /dev/null +++ b/code/models/common/modules/mlp/mlp_1d.py @@ -0,0 +1,1194 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2-style MLP module for 1D-topology devices: N150 (1x1), N300 (1x2), T3K (1x8). + +Single unified MLP1D class with separate forward methods: + - decode_forward(): For decode mode + - prefill_forward(): For prefill mode + - forward(x, mode): Dispatcher that calls the appropriate method + +Execution paths: + Decode: linear(w1) → linear(w3) → mul+silu → reshard → linear(w2) → all_reduce(sharded) → reshard + Prefill: [reshape] → linear(w1) → linear(w3) → mul+silu → linear(w2) → all_reduce → reshape + +""" + +import math +from dataclasses import dataclass, replace +from functools import lru_cache +from pathlib import Path +from typing import Callable, Optional + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight +from models.common.modules.tt_ccl import ( + CCL_CHUNKS_PER_SYNC, + CCL_NUM_BUFFERS_PER_CHANNEL, + CCL_NUM_WORKERS_PER_LINK, + TT_CCL, + default_topology, + get_tt_ccl, +) +from models.common.tensor_utils import TILE_SIZE, get_out_subblock_w, get_padded_hidden_dim, pad_dim_to_size +from models.tt_transformers.tt.common import Mode + +# ============================================================================= +# Top-level config dataclass +# ============================================================================= + + +@dataclass +class MLP1DConfig: + """ + Central configuration for MLP1D - the single source of truth for all settings. + + Simple usage (all defaults): + config = MLP1DConfig(w1, w2, w3) + + Override any field: + config = MLP1DConfig(w1, w2, w3, max_batch_size=64, topology=ttnn.Topology.Ring) + + Full customization: + config = MLP1DConfig( + w1, w2, w3, + mesh_device=custom_device, + decode_w1_w3_prg_config=my_program_config, + ... + ) + """ + + # Required: weights (LazyWeight) + w1: LazyWeight + w2: LazyWeight + w3: LazyWeight + + # Optional: device and collectives + mesh_device: ttnn.MeshDevice | None = None + tt_ccl: TT_CCL | None = None + topology: Optional[ttnn.Topology] = None # None = auto-detect + num_reduce_scatter_links: int = 1 + decode_rs_memory_config: ttnn.MemoryConfig = ttnn.L1_MEMORY_CONFIG + decode_rs_chunks_per_sync: int = 1 + decode_rs_num_workers_per_link: int = 1 + + # Optional: derived from weights if None + dim: int | None = None + hidden_dim: int | None = None + + # Optional: sensible defaults + max_batch_size: int = 32 + mlp_activation_type: ttnn.UnaryOpType = ttnn.UnaryOpType.SILU + + # Optional: power-user overrides (None = compute defaults) + w1_w3_memcfg: ttnn.MemoryConfig | None = None + w2_memcfg: ttnn.MemoryConfig | None = None + + decode_input_memcfg: ttnn.MemoryConfig | None = None + decode_w1_w3_prg_config: ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig | None = None + decode_w2_prg_config: ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig | None = None + decode_mlp2_input_memcfg: ttnn.MemoryConfig | None = None + decode_residual_memcfg: ttnn.MemoryConfig | None = None + + prefill_input_memcfg: ttnn.MemoryConfig | None = None + prefill_w1_w3_prg_config: Callable[[int], ttnn.MatmulMultiCoreReuseMultiCastProgramConfig] | None = None + prefill_w2_prg_config: Callable[[int], ttnn.MatmulMultiCoreReuseMultiCastProgramConfig] | None = None + + # Optional: use ttnn.experimental.minimal_matmul (instead of ttnn.linear) for the W2 down-proj + # prefill matmul above seq_len > 128 — matches TTTv1 (mlp.py L275-281), ~2.5x faster on the large + # folded-batch down-proj. Default OFF so untouched models / decode stay byte-identical; the caller + # opts in. FF1/FF3 stay on ttnn.linear (TTTv1 does too). The factory yields a ttnn.MinimalMatmulConfig + # keyed on the folded seq_len (sibling to prefill_w2_prg_config). + prefill_w2_minimal_matmul: bool = False + prefill_w2_minimal_matmul_config: Callable[[int], "ttnn.MinimalMatmulConfig"] | None = None + + w1_w3_dtype: ttnn.DataType | None = None + w2_dtype: ttnn.DataType | None = None + activation_dtype: ttnn.DataType | None = None + # If True, move W1 decode output to DRAM before W3 linear so W3 CB validation does not overlap W1 L1. + decode_spill_w1_to_dram_before_w3: bool = False + + linear_dtype: ttnn.DataType | None = None + mul_dtype: ttnn.DataType | None = None + + # Architecture-sensitive requests. ``None`` selects the concrete mesh/SKU + # default during construction; explicit compatible values are copied and + # preserved in the internal resolved state. + ff1_3_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + ff2_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + decode_ff1_3_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + decode_ff2_compute_kernel_cfg: ttnn.DeviceComputeKernelConfig | None = None + prefill_len_cutoff: int | None = None + prefill_dram_shard_grid_width: int | None = None + prefill_ff1_ff3_grid: tuple[int, int] | None = None + prefill_ff2_grid: tuple[int, int] | None = None + + def use_minimal_w2_matmul(self, seq_len: int) -> bool: + """Whether the W2 down-proj prefill matmul should use minimal_matmul. + + Mirrors TTTv1 (mlp.py:275) — minimal_matmul only above ``seq_len > 128`` and only when the + per-model opt-in is set. ``seq_len`` is the folded ``B*S`` length. + """ + return bool(self.prefill_w2_minimal_matmul) and seq_len > 128 + + def is_resolved(self) -> bool: + """Check if all fields except optional ones are resolved.""" + # activation_dtype is optional override for linear_dtype and mul_dtype. + # prefill_w2_minimal_matmul_config is only materialized when the minimal-matmul opt-in is set. + optional = { + "activation_dtype", + "prefill_w2_minimal_matmul_config", + } + # topology: None for single_device (CCL not needed) + if self.mesh_device and self.mesh_device.get_num_devices() == 1: + optional.add("topology") + + return all(getattr(self, f) is not None for f in self.__dataclass_fields__ if f not in optional) + + +# ============================================================================= +# MLP1D - Unified MLP for 1D-topology devices (Linear or Ring) with decode and prefill modes +# ============================================================================= + + +class MLP1D(LightweightModule): + """ + MLP for non-TG devices supporting both decode and prefill modes. + + Simple API (90% of users): + mlp = MLP1D(w1, w2, w3) + + Power API (10% of users) - any level of customization via config: + config = MLP1DConfig(w1, w2, w3, max_batch_size=64, topology=ttnn.Topology.Ring) + mlp = MLP1D.from_config(config) + + Execution paths: + Decode: linear(w1) → linear(w3) → mul+silu → reshard → linear(w2) → all_reduce(sharded) → reshard + Prefill: [reshape] → linear(w1) → linear(w3) → mul+silu → linear(w2) → all_reduce → reshape + """ + + def __init__(self, w1: LazyWeight, w2: LazyWeight, w3: LazyWeight): + """ + Simple API for 90% of users - derives all config from weights. + + Args: + w1: Gate projection weight (dim, hidden_dim), sharded on dim=-1 + w2: Down projection weight (hidden_dim, dim), sharded on dim=-2 + w3: Up projection weight (dim, hidden_dim), sharded on dim=-1 + + The mesh_device is derived from w1.device(). tt_ccl is created/cached automatically. + All other settings use sensible defaults. + """ + super().__init__() + self.config = resolve_mlp1d_arch_config(MLP1DConfig(w1=w1, w2=w2, w3=w3)) + self._device_weights_loaded = False + + @classmethod + def from_config(cls, config: MLP1DConfig): + """ + Power API for 10% of users - any level of customization via config. + + Override any subset of fields in MLP1DConfig: + config = MLP1DConfig(w1, w2, w3, max_batch_size=64) + mlp = MLP1D.from_config(config) + + Full customization: + config = MLP1DConfig( + w1, w2, w3, + mesh_device=custom_device, + decode_w1_w3_prg_config=my_program_config, + ... + ) + mlp = MLP1D.from_config(config) + """ + if not isinstance(config, MLP1DConfig): + raise TypeError("MLP1D.from_config expects MLP1DConfig") + instance = object.__new__(cls) + super(MLP1D, instance).__init__() + instance.config = resolve_mlp1d_arch_config(config) + instance._device_weights_loaded = False + return instance + + def load_device_weights(self): + """Materialize LazyWeights onto device. Called automatically on first forward; idempotent.""" + if self._device_weights_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device weights!" + + self.w1 = self.config.w1.get_device_weight() + self.w2 = self.config.w2.get_device_weight() + self.w3 = self.config.w3.get_device_weight() + + self._device_weights_loaded = True + + def decode_forward(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Decode forward - NO if-else, fully flattened. + + Execution path: + linear(w1) → linear(w3) → mul+silu → reshard → linear(w2) → all_reduce(sharded) → reshard + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="decode") + cfg = self.config + + # --- STAGE 1: W1/W3 Linear (L1 sharded) --- + w1_out = ttnn.linear( + x, + self.w1, + dtype=cfg.linear_dtype, + core_grid=None, + compute_kernel_config=cfg.decode_ff1_3_compute_kernel_cfg, + program_config=cfg.decode_w1_w3_prg_config, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + ) + if cfg.decode_spill_w1_to_dram_before_w3: + w1_dram = ttnn.to_memory_config(w1_out, ttnn.DRAM_MEMORY_CONFIG) + ttnn.deallocate(w1_out) + w1_out = w1_dram + + w3_out = ttnn.linear( + x, + self.w3, + dtype=cfg.linear_dtype, + core_grid=None, + compute_kernel_config=cfg.decode_ff1_3_compute_kernel_cfg, + program_config=cfg.decode_w1_w3_prg_config, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + ) + ttnn.deallocate(x) + + # --- STAGE 2: No CCL for non-TG --- + + # --- STAGE 3: Activation + Multiply --- + mul_out_memcfg = ttnn.DRAM_MEMORY_CONFIG if cfg.decode_spill_w1_to_dram_before_w3 else w1_out.memory_config() + w2_in = ttnn.mul( + w1_out, + w3_out, + input_tensor_a_activations=[cfg.mlp_activation_type], + dtype=cfg.mul_dtype, + memory_config=mul_out_memcfg, + ) + + # --- STAGE 3.5: Reshard for w2 --- + w2_in = ttnn.to_memory_config(w2_in, cfg.decode_mlp2_input_memcfg) + + ttnn.deallocate(w3_out) + ttnn.deallocate(w1_out) + + # --- STAGE 4: No all_gather for non-TG --- + + # --- STAGE 5: W2 Linear --- + w2_out = ttnn.linear( + w2_in, + self.w2, + compute_kernel_config=cfg.decode_ff2_compute_kernel_cfg, + dtype=cfg.linear_dtype, + program_config=cfg.decode_w2_prg_config, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + core_grid=None, + ) + ttnn.deallocate(w2_in) + + # --- STAGE 6: Final All-Reduce (decode: sharded=True, no runtime branching) --- + w2_out_reduced = self._all_reduce_decode(w2_out) + + # --- STAGE 7: Reshape + Final memory config --- + original_shape = w2_out_reduced.shape + w2_out_reduced = ttnn.reshape( + w2_out_reduced, (1, 1, original_shape[-4] * original_shape[-3] * original_shape[-2], original_shape[-1]) + ) + w2_out_reduced = ttnn.to_memory_config(w2_out_reduced, cfg.decode_residual_memcfg) + + return w2_out_reduced + + def prefill_forward(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Prefill forward - minimal runtime logic for seq_len-dependent configs. + + Execution path: + [reshape if seq_len >= cutoff] → linear(w1) → linear(w3) → mul+silu → linear(w2) → all_reduce → reshape + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="prefill") + + cfg = self.config + seq_len = x.shape[-2] + + # Seq_len-dependent: reshape for long sequences + if seq_len >= cfg.prefill_len_cutoff: + assert ( + seq_len % cfg.prefill_len_cutoff == 0 + ), f"seq_len ({seq_len}) must be divisible by prefill_len_cutoff ({cfg.prefill_len_cutoff})" + x = ttnn.reshape(x, [1, seq_len // cfg.prefill_len_cutoff, cfg.prefill_len_cutoff, -1]) + + # Seq_len-dependent: get program configs by calling methods on config + pc_w1_w3 = cfg.prefill_w1_w3_prg_config(seq_len) + pc_w2 = cfg.prefill_w2_prg_config(seq_len) + + # --- STAGE 1: W1/W3 Linear (DRAM) --- + w1_out = ttnn.linear( + x, + self.w1, + dtype=cfg.linear_dtype, + core_grid=None, + compute_kernel_config=cfg.ff1_3_compute_kernel_cfg, + program_config=pc_w1_w3, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + w3_out = ttnn.linear( + x, + self.w3, + dtype=cfg.linear_dtype, + core_grid=None, + compute_kernel_config=cfg.ff1_3_compute_kernel_cfg, + program_config=pc_w1_w3, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + ttnn.deallocate(x) + + # --- STAGE 2: No CCL for non-TG --- + + # --- STAGE 3: Activation + Multiply --- + w2_in = ttnn.mul( + w1_out, + w3_out, + input_tensor_a_activations=[cfg.mlp_activation_type], + dtype=cfg.mul_dtype, + memory_config=w1_out.memory_config(), + ) + + # --- STAGE 3.5: No reshard for prefill --- + + ttnn.deallocate(w3_out) + ttnn.deallocate(w1_out) + + # --- STAGE 4: No all_gather for non-TG --- + + # --- STAGE 5: W2 Linear --- + # Above seq_len > 128 the folded-batch down-proj is large; minimal_matmul is ~2.5x faster than + # ttnn.linear there (TTTv1 parity). Opt-in via prefill_w2_minimal_matmul; output shape is + # identical to the ttnn.linear path. FF1/FF3 stay on ttnn.linear (matches TTTv1). + if cfg.use_minimal_w2_matmul(seq_len): + w2_out = ttnn.experimental.minimal_matmul( + w2_in, + self.w2, + compute_kernel_config=cfg.ff2_compute_kernel_cfg, + config=cfg.prefill_w2_minimal_matmul_config(seq_len), + ) + else: + w2_out = ttnn.linear( + w2_in, + self.w2, + compute_kernel_config=cfg.ff2_compute_kernel_cfg, + dtype=cfg.linear_dtype, + program_config=pc_w2, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + core_grid=None, + ) + ttnn.deallocate(w2_in) + + # --- STAGE 6: Final All-Reduce (prefill: sharded=False) --- + w2_out_reduced = self._all_reduce_prefill(w2_out) + + # --- STAGE 7: Reshape (no final memory config change for prefill) --- + original_shape = w2_out_reduced.shape + w2_out_reduced = ttnn.reshape( + w2_out_reduced, (1, 1, original_shape[-4] * original_shape[-3] * original_shape[-2], original_shape[-1]) + ) + + return w2_out_reduced + + def forward(self, x: ttnn.Tensor | LazyWeight, mode: str | Mode) -> ttnn.Tensor: + """Dispatch to the appropriate forward method based on mode.""" + if isinstance(mode, Mode): + mode = mode.value + if mode == "decode": + return self.decode_forward(x) + else: + return self.prefill_forward(x) + + def _all_reduce_decode(self, w2_out: ttnn.Tensor) -> ttnn.Tensor: + """All-reduce for decode mode (sharded input).""" + cfg = self.config + if cfg.mesh_device.get_num_devices() == 1: + return w2_out + + original_shape = w2_out.shape + if original_shape[0] != 1 or original_shape[1] != 1: + w2_out = ttnn.reshape( + w2_out, (1, 1, original_shape[-4] * original_shape[-3] * original_shape[-2], original_shape[-1]) + ) + + reduced = ttnn.experimental.reduce_scatter_minimal_async( + w2_out, + persistent_output_buffers=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_rs_semaphore_handles(), + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + num_links=cfg.num_reduce_scatter_links, + memory_config=w2_out.memory_config(), + intermediate_memory_config=cfg.decode_rs_memory_config, + topology=cfg.topology, + chunks_per_sync=cfg.decode_rs_chunks_per_sync, + num_workers_per_link=cfg.decode_rs_num_workers_per_link, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + w2_out.deallocate(True) + return reduced + + def _all_reduce_prefill(self, w2_out: ttnn.Tensor) -> ttnn.Tensor: + """All-reduce for prefill mode (interleaved input).""" + cfg = self.config + if cfg.mesh_device.get_num_devices() == 1: + return w2_out + + original_shape = w2_out.shape + if original_shape[0] != 1 or original_shape[1] != 1: + w2_out = ttnn.reshape( + w2_out, (1, 1, original_shape[-4] * original_shape[-3] * original_shape[-2], original_shape[-1]) + ) + + if w2_out.is_sharded(): + w2_out_sharded = w2_out + w2_out = ttnn.sharded_to_interleaved(w2_out_sharded, ttnn.L1_MEMORY_CONFIG) + w2_out_sharded.deallocate(True) + + reduced = ttnn.experimental.reduce_scatter_minimal_async( + w2_out, + persistent_output_buffers=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_rs_semaphore_handles(), + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + num_links=cfg.num_reduce_scatter_links, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + intermediate_memory_config=ttnn.DRAM_MEMORY_CONFIG, + topology=cfg.topology, + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + w2_out.deallocate(True) + return reduced + + # [INFO] this is the entry point for TTTv1 model_config.py and will retire with TTTv1 + @classmethod + def from_model_args( + cls, + mesh_device, + tt_ccl, + args, + state_dict, + weight_cache_path, + layer_num: int, + dtype=None, + model_config=None, + state_dict_prefix: Optional[str] = None, + prefetcher=None, + ): + """Factory method for backward compatibility with ModelArgs. + + Args: + mesh_device: The mesh device to use. + tt_ccl: The TT CCL instance. + args: Model arguments (ModelArgs instance). + state_dict: The state dictionary containing weights. + weight_cache_path: Path for weight caching. + layer_num: The layer number. + dtype: Optional data type for weights (for signature compatibility with TTTv1 MLP). + model_config: Optional model config dict. If None, calls args.get_model_config(). + state_dict_prefix: Optional prefix for state dict keys. + + Note: + The `dtype` parameter is accepted for signature compatibility with TTTv1 MLP + but is not used directly - dtype is determined by the DecodersPrecision config. + """ + if args.is_galaxy: + raise ValueError("MLP1D cannot be used for Galaxy devices.") + + import torch + + from models.tt_transformers.tt.model_config import OpGroup, TensorGroup + + # Get model_config for overrides - use passed model_config if provided + if model_config is None: + model_config = args.get_model_config() + decoders_opt = model_config.get("DECODERS_OPTIMIZATIONS") + effective_layer_num = max(layer_num, 0) + + # Extract settings from args/model_config + ccl_topology = args.ccl_topology() + + if state_dict_prefix is None: + state_dict_prefix = args.get_state_dict_prefix("MLP", layer_num) + + # Get dtypes from optimizer config + ff1_3_dtype = decoders_opt.get_tensor_dtype(decoder_id=effective_layer_num, tensor=TensorGroup.FF1_FF3) + ff2_dtype = decoders_opt.get_tensor_dtype(decoder_id=effective_layer_num, tensor=TensorGroup.FF2) + activation_dtype = decoders_opt.get_tensor_dtype(decoder_id=effective_layer_num, tensor=TensorGroup.ACTIVATION) + + # Get compute kernel configs + ff1_3_compute_kernel_cfg = decoders_opt.get_math_fidelity( + decoder_id=effective_layer_num, op=OpGroup.LI_FF1_FF3, configuration=args + ) + ff2_compute_kernel_cfg = decoders_opt.get_math_fidelity( + decoder_id=effective_layer_num, op=OpGroup.LI_FF2, configuration=args + ) + + # Get decode program configs from model_config + decode_w1_w3_prg_config = args.get_mlp_ff1_3_prg_config(Mode.DECODE, None, None) + decode_w2_prg_config = args.get_mlp_ff2_prg_config(Mode.DECODE, None, None) + decode_mlp2_input_memcfg = args.get_mlp_binary_mult_mem_config(Mode.DECODE) + decode_residual_memcfg = args.get_mlp_output_mem_config(Mode.DECODE, None) + mlp_rs_cfg = model_config.get("MLP_RS_CONFIG", {}) + + # Compute memory configs for weights + num_devices = mesh_device.get_num_devices() + tile_size = TILE_SIZE + dram_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + { + ttnn.CoreRange( + ttnn.CoreCoord(0, 0), + ttnn.CoreCoord(dram_size.x - 1, dram_size.y - 1), + ) + } + ) + + w1_w3_mem_config = _create_dram_sharded_mem_config( + k=args.dim, + n=args.hidden_dim // num_devices, + dram_grid=dram_grid, + tile_size=tile_size, + dram_cores=dram_size.x, + ) + w2_mem_config = _create_dram_sharded_mem_config( + k=args.hidden_dim // num_devices, + n=args.dim, + dram_grid=dram_grid, + tile_size=tile_size, + dram_cores=dram_size.x, + ) + + cache_dir = None if args.dummy_weights else Path(weight_cache_path) / state_dict_prefix + + # Create LazyWeights + def make_weight_source(name: str, shard_dim: int): + tensor = torch.transpose(state_dict[f"{state_dict_prefix}.{name}.weight"], -2, -1) + return pad_dim_to_size(tensor, dim=shard_dim, size=args.hidden_dim) + + w1 = LazyWeight( + source=make_weight_source("w1", -1), + dtype=ff1_3_dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=w1_w3_mem_config, + cache_dir_weight_name=(cache_dir, "w1_sharded") if cache_dir else None, + ) + w2 = LazyWeight( + source=make_weight_source("w2", -2), + dtype=ff2_dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-2)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=w2_mem_config, + cache_dir_weight_name=(cache_dir, "w2_sharded") if cache_dir else None, + ) + w3 = LazyWeight( + source=make_weight_source("w3", -1), + dtype=ff1_3_dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=w1_w3_mem_config, + cache_dir_weight_name=(cache_dir, "w3_sharded") if cache_dir else None, + ) + + # Create config with all the overrides and use from_config + config = MLP1DConfig( + w1=w1, + w2=w2, + w3=w3, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + dim=args.dim, + hidden_dim=args.hidden_dim, + max_batch_size=args.max_batch_size, + mlp_activation_type=getattr(args, "mlp_activation_type", ttnn.UnaryOpType.SILU), + topology=ccl_topology, + decode_rs_memory_config=mlp_rs_cfg.get("rs_memory_config", ttnn.L1_MEMORY_CONFIG), + decode_rs_chunks_per_sync=mlp_rs_cfg.get("chunks_per_sync", 1), + decode_rs_num_workers_per_link=mlp_rs_cfg.get("num_workers_per_link", 1), + decode_w1_w3_prg_config=decode_w1_w3_prg_config, + decode_w2_prg_config=decode_w2_prg_config, + decode_mlp2_input_memcfg=decode_mlp2_input_memcfg, + decode_residual_memcfg=decode_residual_memcfg, + w1_w3_dtype=ff1_3_dtype, + w2_dtype=ff2_dtype, + activation_dtype=activation_dtype, + decode_spill_w1_to_dram_before_w3=False, + ff1_3_compute_kernel_cfg=ff1_3_compute_kernel_cfg, + ff2_compute_kernel_cfg=ff2_compute_kernel_cfg, + decode_ff1_3_compute_kernel_cfg=ff1_3_compute_kernel_cfg, + decode_ff2_compute_kernel_cfg=ff2_compute_kernel_cfg, + prefill_len_cutoff=args.prefill_len_cutoff, + ) + return cls.from_config(config) + + +# ============================================================================= +# Config helper functions (adapted from TTTv1 model_config.py) +# ============================================================================= + + +def _find_largest_divisor(n: int, max_divisor: int = 8) -> int: + """Find largest divisor of n up to max_divisor.""" + for i in range(max_divisor, 0, -1): + if n % i == 0: + return i + return 1 + + +def _find_grid(n_tiles: int, max_rows: int = 8, max_cols: int = 8) -> tuple[int, int]: + """Find grid dimensions (rows, cols) that evenly divide n_tiles.""" + max_cores = max_rows * max_cols + target = max_cores // 2 # prefer half the grid for balanced utilization + possible_cores = [k for k in range(1, max_cores + 1) if n_tiles % k == 0] + possible_cores.sort(key=lambda x: abs(x - target)) + + for cores in possible_cores: + for rows in range(1, max_rows + 1): + if cores % rows == 0: + cols = cores // rows + if cols <= max_cols: + return rows, cols + + raise AssertionError(f"Cannot find grid for {n_tiles} tiles within {max_rows}x{max_cols}") + + +def _find_grid_k_n(k_tiles: int, n_tiles: int, max_rows: int = 8, max_cols: int = 8) -> tuple[int, int]: + """Find grid that evenly divides both K and N tile counts.""" + max_cores = max_rows * max_cols + possible_cores = [c for c in range(1, max_cores + 1) if k_tiles % c == 0 and n_tiles % c == 0] + possible_cores.sort(reverse=True) + + for cores in possible_cores: + for rows in range(1, max_rows + 1): + if cores % rows == 0: + cols = cores // rows + if cols <= max_cols: + return rows, cols + + raise AssertionError(f"Cannot find grid for K={k_tiles}, N={n_tiles} tiles") + + +def _find_prefill_grid(row_tiles: int, col_tiles: int, max_rows: int = 8, max_cols: int = 8) -> tuple[int, int]: + """Find grid where row_tiles divides rows and col_tiles divides cols.""" + cols = next((i for i in range(max_cols, 0, -1) if col_tiles % i == 0), None) + rows = next((i for i in range(max_rows, 0, -1) if row_tiles % i == 0), None) + assert cols is not None and rows is not None + return rows, cols + + +def _dram_shard_core_grid(k: int, tile_size: int = TILE_SIZE) -> ttnn.CoreGrid: + """Get core grid for DRAM sharding based on K dimension.""" + rows, cols = _find_grid(k // tile_size) + return ttnn.CoreGrid(x=cols, y=rows) + + +def _dram_shard_core_grid_k_n(k: int, n: int, tile_size: int = TILE_SIZE) -> ttnn.CoreGrid: + """Get core grid for DRAM sharding based on K and N dimensions.""" + rows, cols = _find_grid_k_n(k // tile_size, n // tile_size) + return ttnn.CoreGrid(x=cols, y=rows) + + +def _dram_matmul_config( + m: int, k: int, n: int, num_cores: int, tile_size: int = TILE_SIZE, fused_activation=None +) -> ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig: + """Create DRAM-sharded matmul program config.""" + return ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig( + in0_block_w=_find_largest_divisor(k // (tile_size * num_cores)), + per_core_M=math.ceil(m / tile_size), + per_core_N=math.ceil(n / (tile_size * num_cores)), + fused_activation=fused_activation, + ) + + +def _matmul_config( + m: int, + k: int, + n: int, + grid_size: tuple[int, int], + tile_size: int = TILE_SIZE, + in0_block_w: int = None, + fuse_batch: bool = False, + fused_activation=None, + per_core_m: int = None, + per_core_n: int = None, +) -> ttnn.MatmulMultiCoreReuseMultiCastProgramConfig: + """Create multicast matmul program config.""" + if per_core_m is None: + per_core_m = math.ceil(m / (tile_size * grid_size[1])) + if per_core_n is None: + per_core_n = math.ceil(n / (tile_size * grid_size[0])) + + out_subblock_h = 1 + out_subblock_w = get_out_subblock_w(per_core_n, out_subblock_h) + + if in0_block_w is None: + in0_block_w = _find_largest_divisor(k // (tile_size * grid_size[1])) + + return ttnn.MatmulMultiCoreReuseMultiCastProgramConfig( + compute_with_storage_grid_size=grid_size, + in0_block_w=in0_block_w, + out_subblock_h=out_subblock_h, + out_subblock_w=out_subblock_w, + per_core_M=per_core_m, + per_core_N=per_core_n, + transpose_mcast=False, + fused_activation=fused_activation, + fuse_batch=fuse_batch, + ) + + +def _copy_compute_kernel_config(arch, config=None) -> ttnn.DeviceComputeKernelConfig: + """Create an independent concrete compute config for ``arch``.""" + if config is None: + values = { + "math_fidelity": ttnn.MathFidelity.HiFi2, + "math_approx_mode": False, + "fp32_dest_acc_en": False, + "packer_l1_acc": True, + "dst_full_sync_en": False, + } + else: + required = ( + "math_fidelity", + "math_approx_mode", + "fp32_dest_acc_en", + "packer_l1_acc", + "dst_full_sync_en", + ) + missing = [field for field in required if not hasattr(config, field)] + if missing: + raise ValueError(f"MLP1D compute config is missing fields: {', '.join(missing)}") + values = {field: getattr(config, field) for field in required} + if hasattr(config, "throttle_level"): + values["throttle_level"] = config.throttle_level + + candidate = ttnn.WormholeComputeKernelConfig(**values) + return ttnn.init_device_compute_kernel_config(arch, candidate) + + +def _compute_kernel_config_hifi2_fp16(arch) -> ttnn.DeviceComputeKernelConfig: + """Default compute kernel config for MLP (HiFi2 with FP16 accumulation).""" + return _copy_compute_kernel_config(arch) + + +def _resolve_mlp1d_mesh(config: MLP1DConfig): + mesh_device = config.mesh_device or getattr(config.w1, "device", None) or ttnn.GetDefaultDevice() + if mesh_device is None: + raise ValueError("MLP1D requires a mesh_device or a weight associated with a mesh device") + for name in ("w1", "w2", "w3"): + weight_device = getattr(getattr(config, name), "device", None) + if weight_device is not None and weight_device != mesh_device: + raise ValueError(f"MLP1D {name} device must match mesh_device") + return mesh_device + + +def _validate_mlp1d_arch_config(config: MLP1DConfig, arch) -> None: + for field in ( + "ff1_3_compute_kernel_cfg", + "ff2_compute_kernel_cfg", + "decode_ff1_3_compute_kernel_cfg", + "decode_ff2_compute_kernel_cfg", + ): + if getattr(config, field) is None: + raise ValueError(f"MLP1D resolved config requires {field}") + + if config.prefill_len_cutoff <= 0 or config.prefill_len_cutoff % TILE_SIZE: + raise ValueError("MLP1D prefill_len_cutoff must be a positive multiple of the tile size") + if config.prefill_dram_shard_grid_width <= 0: + raise ValueError("MLP1D prefill DRAM shard grid width must be positive") + + mesh_device = _resolve_mlp1d_mesh(config) + dram_width = mesh_device.dram_grid_size().x + expected_dram_width = 8 if arch == ttnn.device.Arch.WORMHOLE_B0 else dram_width + if config.prefill_dram_shard_grid_width != expected_dram_width: + raise ValueError( + "MLP1D prefill DRAM shard grid width does not match the resolved architecture/SKU " + f"({config.prefill_dram_shard_grid_width} != {expected_dram_width})" + ) + + compute_grid = mesh_device.compute_with_storage_grid_size() + dim = config.dim if config.dim is not None else config.w1.source.shape[-2] + hidden_dim = config.hidden_dim if config.hidden_dim is not None else config.w1.source.shape[-1] + grid_tile_shapes = { + "prefill_ff1_ff3_grid": (8, dim // TILE_SIZE), + "prefill_ff2_grid": ( + 8, + get_padded_hidden_dim(hidden_dim, mesh_device.get_num_devices(), TILE_SIZE) // TILE_SIZE, + ), + } + for field, tile_shape in grid_tile_shapes.items(): + grid = getattr(config, field) + if len(grid) != 2 or min(grid) <= 0: + raise ValueError(f"MLP1D {field} must contain two positive dimensions") + if grid[0] > compute_grid.x or grid[1] > compute_grid.y: + raise ValueError(f"MLP1D {field} {grid} exceeds mesh compute grid ({compute_grid.x}, {compute_grid.y})") + if tile_shape[0] % grid[0] or tile_shape[1] % grid[1]: + raise ValueError(f"MLP1D {field} {grid} must evenly divide tile shape {tile_shape}") + + +def resolve_mlp1d_arch_config(config: MLP1DConfig) -> MLP1DConfig: + """Return a fully resolved config after one concrete-mesh architecture query.""" + if not isinstance(config, MLP1DConfig): + raise TypeError("resolve_mlp1d_arch_config expects MLP1DConfig") + mesh_device = _resolve_mlp1d_mesh(config) + arch = mesh_device.arch() + if arch not in (ttnn.device.Arch.WORMHOLE_B0, ttnn.device.Arch.BLACKHOLE): + raise ValueError(f"Unsupported MLP1D architecture: {arch}") + + dim = config.dim if config.dim is not None else config.w1.source.shape[-2] + hidden_dim = config.hidden_dim if config.hidden_dim is not None else config.w1.source.shape[-1] + if dim <= 0 or dim % TILE_SIZE: + raise ValueError("MLP1D dim must be a positive multiple of the tile size") + if hidden_dim <= 0: + raise ValueError("MLP1D hidden_dim must be positive") + num_devices = mesh_device.get_num_devices() + padded_hidden_dim = get_padded_hidden_dim(hidden_dim, num_devices, TILE_SIZE) + # Simple construction retains the legacy TTTv1 cutoff (1024 on Wormhole, + # 512 on Blackhole). Blackhole uses the concrete SKU's DRAM width (7 on + # P100, 8 on P150), while Wormhole's accepted recipe intentionally uses 8. + default_prefill_len_cutoff = 1024 if arch == ttnn.device.Arch.WORMHOLE_B0 else 512 + default_prefill_dram_shard_grid_width = ( + 8 if arch == ttnn.device.Arch.WORMHOLE_B0 else mesh_device.dram_grid_size().x + ) + default_prefill_ff1_ff3_grid = _find_prefill_grid(8, dim // TILE_SIZE) + default_prefill_ff2_grid = _find_prefill_grid(8, padded_hidden_dim // TILE_SIZE) + resolved = replace( + config, + ff1_3_compute_kernel_cfg=_copy_compute_kernel_config(arch, config.ff1_3_compute_kernel_cfg), + ff2_compute_kernel_cfg=_copy_compute_kernel_config(arch, config.ff2_compute_kernel_cfg), + decode_ff1_3_compute_kernel_cfg=_copy_compute_kernel_config(arch, config.decode_ff1_3_compute_kernel_cfg), + decode_ff2_compute_kernel_cfg=_copy_compute_kernel_config(arch, config.decode_ff2_compute_kernel_cfg), + prefill_len_cutoff=( + default_prefill_len_cutoff if config.prefill_len_cutoff is None else config.prefill_len_cutoff + ), + prefill_dram_shard_grid_width=( + default_prefill_dram_shard_grid_width + if config.prefill_dram_shard_grid_width is None + else config.prefill_dram_shard_grid_width + ), + prefill_ff1_ff3_grid=( + default_prefill_ff1_ff3_grid if config.prefill_ff1_ff3_grid is None else config.prefill_ff1_ff3_grid + ), + prefill_ff2_grid=(default_prefill_ff2_grid if config.prefill_ff2_grid is None else config.prefill_ff2_grid), + ) + _validate_mlp1d_arch_config(resolved, arch) + return _resolve_mlp1d_config(resolved) + + +def _create_dram_sharded_mem_config( + k: int, n: int, dram_grid: ttnn.CoreRangeSet, tile_size: int = TILE_SIZE, dram_cores: int = 12 +) -> ttnn.MemoryConfig: + """Create DRAM-sharded memory config for weight tensors.""" + padded_size = math.ceil(n / (tile_size * dram_cores)) * (tile_size * dram_cores) + shard_spec = ttnn.ShardSpec(dram_grid, (k, padded_size // dram_cores), ttnn.ShardOrientation.ROW_MAJOR) + return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, shard_spec) + + +def _resolve_mlp1d_config(config: MLP1DConfig) -> MLP1DConfig: + """Materialize non-architecture fields after architecture requests resolve.""" + + to_set = {} + + # --- Phase 1: Foundational fields (order matters due to dependencies) --- + + # Derive dimensions from weights + # w1 is expected in TTNN layout: (dim, hidden_dim) - caller must transpose if needed + dim = config.dim + if config.dim is None: + dim = config.w1.source.shape[-2] + to_set["dim"] = dim + + hidden_dim = config.hidden_dim + if config.hidden_dim is None: + hidden_dim = config.w1.source.shape[-1] + to_set["hidden_dim"] = hidden_dim + + # Derive and validate the concrete mesh without another architecture query. + mesh_device = _resolve_mlp1d_mesh(config) + if config.mesh_device is None: + to_set["mesh_device"] = mesh_device + + # Derive tt_ccl + tt_ccl = config.tt_ccl + if config.tt_ccl is None: + tt_ccl = get_tt_ccl(mesh_device) + to_set["tt_ccl"] = tt_ccl + + assert tt_ccl is not None, "tt_ccl must be available at this point!" + assert tt_ccl.mesh_device == mesh_device, "tt_ccl must match the device of mesh_device!" + + # Auto-detect topology + topology = config.topology + if config.topology is None: + topology = default_topology(mesh_device) + to_set["topology"] = topology + + # --- Phase 2: Derived fields --- + + num_devices = mesh_device.get_num_devices() + tile_size = TILE_SIZE + tile_padded_batch_rows = tile_size * math.ceil(config.max_batch_size / tile_size) + + # Compute padded hidden_dim for memory configs (must match auto-padding in LazyWeight) + padded_hidden_dim = get_padded_hidden_dim(hidden_dim, num_devices, tile_size) + + # Always computed (not user-overridable None fields) + w1_w3_dtype = config.w1_w3_dtype + if config.w1_w3_dtype is None: + w1_w3_dtype = ttnn.bfloat8_b + to_set["w1_w3_dtype"] = w1_w3_dtype + w2_dtype = config.w2_dtype + if config.w2_dtype is None: + w2_dtype = ttnn.bfloat8_b + to_set["w2_dtype"] = w2_dtype + + if config.linear_dtype is None: + to_set["linear_dtype"] = config.activation_dtype or ttnn.bfloat16 + if config.mul_dtype is None: + to_set["mul_dtype"] = config.activation_dtype or ttnn.bfloat8_b + prefill_len_cutoff = config.prefill_len_cutoff + + # --- Phase 3: Decode program configs --- + # Note: Use padded_hidden_dim to match auto-padding in LazyWeight + + mlp_core_grid = _dram_shard_core_grid_k_n(dim, padded_hidden_dim // num_devices) + + if config.decode_input_memcfg is None: + to_set["decode_input_memcfg"] = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // mlp_core_grid.num_cores), # Shard shape: 1 shard per core + mlp_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + if config.decode_w1_w3_prg_config is None: + to_set["decode_w1_w3_prg_config"] = _dram_matmul_config( + m=tile_padded_batch_rows, + k=dim, + n=padded_hidden_dim // num_devices, + num_cores=mlp_core_grid.num_cores, + ) + + mlp2_core_grid = _dram_shard_core_grid_k_n(padded_hidden_dim // num_devices, dim) + + if config.decode_w2_prg_config is None: + to_set["decode_w2_prg_config"] = _dram_matmul_config( + m=tile_padded_batch_rows, + k=padded_hidden_dim // num_devices, + n=dim, + num_cores=mlp2_core_grid.num_cores, + ) + + if config.decode_mlp2_input_memcfg is None: + to_set["decode_mlp2_input_memcfg"] = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, padded_hidden_dim // num_devices // mlp2_core_grid.num_cores), + mlp2_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + if config.decode_residual_memcfg is None: + residual_grid = _dram_shard_core_grid(dim // num_devices) + to_set["decode_residual_memcfg"] = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // residual_grid.num_cores // num_devices), + residual_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + # --- Phase 4: Prefill program configs --- + + # Matching per_core_N to the resolved SKU width avoids silent PCC issues + # on P100 while preserving the accepted Wormhole width of 8. + dram_shard_grid_width = config.prefill_dram_shard_grid_width + + if config.prefill_input_memcfg is None: + to_set["prefill_input_memcfg"] = ttnn.DRAM_MEMORY_CONFIG + + if config.prefill_w1_w3_prg_config is None: + prefill_mlp_grid_size = config.prefill_ff1_ff3_grid + n_w1_w3 = padded_hidden_dim // num_devices + + @lru_cache + def w1_w3_prg_config(seq_len: int): + return _matmul_config( + m=min(seq_len, prefill_len_cutoff), + k=dim, + n=n_w1_w3, + grid_size=prefill_mlp_grid_size, + per_core_n=math.ceil(n_w1_w3 / (tile_size * dram_shard_grid_width)), + ) + + to_set["prefill_w1_w3_prg_config"] = lambda seq_len: w1_w3_prg_config(seq_len) + + if config.prefill_w2_prg_config is None: + n_w2 = dim + grid_size = config.prefill_ff2_grid + + @lru_cache + def w2_prg_config(seq_len: int): + return _matmul_config( + m=min(seq_len, prefill_len_cutoff), + k=padded_hidden_dim, + n=n_w2, + grid_size=grid_size, + per_core_n=math.ceil(n_w2 / (tile_size * dram_shard_grid_width)), + ) + + to_set["prefill_w2_prg_config"] = lambda seq_len: w2_prg_config(seq_len) + + # minimal_matmul config for W2 (only materialized when the opt-in is set). Block sizes mirror TTTv1 + # (mlp.py via model_config.py:1329-1334: 8/8/8 blocks); the compute grid is the W2 prefill grid + # (find_prefill_grid over padded_hidden_dim), equivalent to TTTv1's mlp2_grid(seq_len). + if config.prefill_w2_minimal_matmul and config.prefill_w2_minimal_matmul_config is None: + minimal_w2_grid = config.prefill_ff2_grid + + @lru_cache + def w2_minimal_matmul_config(seq_len: int): + return ttnn.MinimalMatmulConfig( + M_block_size=8, + K_block_size=8, + N_block_size=8, + compute_with_storage_grid_size=ttnn.CoreCoord(minimal_w2_grid[0], minimal_w2_grid[1]), + ) + + to_set["prefill_w2_minimal_matmul_config"] = lambda seq_len: w2_minimal_matmul_config(seq_len) + + # --- Phase 5: Weight memory configs --- + + dram_grid_size = mesh_device.dram_grid_size() + dram_grid = ttnn.CoreRangeSet( + { + ttnn.CoreRange( + ttnn.CoreCoord(0, 0), + ttnn.CoreCoord(dram_grid_size.x - 1, dram_grid_size.y - 1), + ) + } + ) + + w1_w3_memcfg = config.w1_w3_memcfg + if w1_w3_memcfg is None: + w1_w3_memcfg = _create_dram_sharded_mem_config( + k=dim, + n=padded_hidden_dim // num_devices, + dram_grid=dram_grid, + tile_size=tile_size, + dram_cores=dram_grid_size.x, + ) + to_set["w1_w3_memcfg"] = w1_w3_memcfg + + w2_memcfg = config.w2_memcfg + if w2_memcfg is None: + w2_memcfg = _create_dram_sharded_mem_config( + k=padded_hidden_dim // num_devices, + n=dim, + dram_grid=dram_grid, + tile_size=tile_size, + dram_cores=dram_grid_size.x, + ) + to_set["w2_memcfg"] = w2_memcfg + + # --- Phase 6: Resolve LazyWeights --- + + w1_w3_mesh_mapper_config = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ) + to_set["w1"] = resolve_lazy_weight( + config.w1, + device=mesh_device, + memory_config=w1_w3_memcfg, + mesh_mapper_config=w1_w3_mesh_mapper_config, + layout=ttnn.TILE_LAYOUT, + dtype=w1_w3_dtype, + ) + + to_set["w2"] = resolve_lazy_weight( + config.w2, + device=mesh_device, + memory_config=w2_memcfg, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(-2)], + mesh_shape_override=ttnn.MeshShape([num_devices]), + ), + layout=ttnn.TILE_LAYOUT, + dtype=w2_dtype, + ) + + to_set["w3"] = resolve_lazy_weight( + config.w3, + device=mesh_device, + memory_config=w1_w3_memcfg, + mesh_mapper_config=w1_w3_mesh_mapper_config, + layout=ttnn.TILE_LAYOUT, + dtype=w1_w3_dtype, + ) + + # --- Final: Create new config with all resolved fields --- + # todo)) using the current if is None else to_set[] does not seem to be worth the saving in space. Maybe cleaner to build all the fields in to_set and then replace the config with them with: + # to_override_set = {k: v for k, v in kwargs.items() if getattr(config, k, None) is None} + # return replace(config, **to_override_set) + resolved_config = replace(config, **to_set) + + assert all( + [resolved_config.w1.is_resolved(), resolved_config.w2.is_resolved(), resolved_config.w3.is_resolved()] + ), "All weights must be resolved!" + assert resolved_config.is_resolved(), "Config must be resolved!" + # check that the padded shapes match the config + assert resolved_config.w1.padded_shape == (dim, padded_hidden_dim), "w1 padded_shape does not match the config!" + assert resolved_config.w2.padded_shape == (padded_hidden_dim, dim), "w2 padded_shape does not match the config!" + assert resolved_config.w3.padded_shape == (dim, padded_hidden_dim), "w3 padded_shape does not match the config!" + + return resolved_config + + +def _load_input_device_tensor(x: ttnn.Tensor | LazyWeight, config: MLP1DConfig, mode: str) -> ttnn.Tensor: + """Resolve the input tensor to ttnn tensor if x is a LazyWeight, otherwise sanity check x and then return as is""" + assert mode in ["decode", "prefill"], "mode must be one of decode or prefill!" + mem_cfg = config.decode_input_memcfg if mode == "decode" else config.prefill_input_memcfg + if isinstance(x, LazyWeight): + # resolve in place + resolved_x = resolve_lazy_weight( + x, + device=config.mesh_device, + memory_config=mem_cfg, + mesh_mapper_config=None, # replicated + layout=ttnn.TILE_LAYOUT, + ) + return resolved_x.get_device_weight() + + assert isinstance(x, ttnn.Tensor), "x must be a ttnn tensor at this point!" + if x.memory_config() != mem_cfg: + raise ValueError("Input tensor memory config does not match the config!") + + return x diff --git a/code/models/common/modules/mlp/mlp_2d.py b/code/models/common/modules/mlp/mlp_2d.py new file mode 100644 index 0000000000000000000000000000000000000000..edc156730ff7bf918b2c0d56f4590f66218adba1 --- /dev/null +++ b/code/models/common/modules/mlp/mlp_2d.py @@ -0,0 +1,823 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2-style MLP module for TG (Galaxy) devices with 2D mesh topology. + +Single unified MLP2D class with separate forward methods: + - decode_forward(): For decode mode + - prefill_forward(): For prefill mode + - forward(x, mode): Dispatcher that calls the appropriate method + +Execution paths: + - Unified: linear → linear → reduce_scatter(×2) → mul+silu → all_gather → linear → all_reduce + +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any, Callable, Optional + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight +from models.common.modules.tt_ccl import ( + CCL_CHUNKS_PER_SYNC, + CCL_NUM_BUFFERS_PER_CHANNEL, + CCL_NUM_WORKERS_PER_LINK, + TT_CCL, + default_topology, + get_tt_ccl, +) +from models.common.tensor_utils import pad_dim_to_size +from models.common.utility_functions import is_blackhole + +# ============================================================================= +# Top-level config dataclass +# ============================================================================= + + +@dataclass +class MLP2DConfig: + """ + Central configuration for MLP2D - the single source of truth for all settings. + + None fields are populated with derived defaults during config resolution + (inside ``MLP2D.__init__`` or ``MLP2D.from_config``). + + Simple usage (all defaults): + config = MLP2DConfig(w1, w2, w3) + + Override any field: + config = MLP2DConfig(w1, w2, w3, max_batch_size=64) + + Full customization: + config = MLP2DConfig( + w1, w2, w3, + mesh_device=custom_device, + decode_w1_w3_prg_config=my_program_config, + ... + ) + """ + + # Required: weights (LazyWeight) + w1: LazyWeight + w2: LazyWeight + w3: LazyWeight + + # Optional: device and collectives + mesh_device: ttnn.MeshDevice | None = None + tt_ccl: TT_CCL | None = None + topology: Optional[ttnn.Topology] = None # None = auto-detect + num_reduce_scatter_links: int = 1 + num_all_gather_links: int = 2 + + # Optional: derived from weights if None + dim: int | None = None + hidden_dim: int | None = None + + # Optional: sensible defaults + max_batch_size: int = 32 + mlp_activation_type: ttnn.UnaryOpType = ttnn.UnaryOpType.SILU + + # Optional: power-user overrides (None = compute defaults) + w1_w3_memcfg: ttnn.MemoryConfig | None = None + w2_memcfg: ttnn.MemoryConfig | None = None + + # Decode settings + decode_input_memcfg: ttnn.MemoryConfig | None = None + decode_w1_w3_prg_config: ttnn.MatmulMultiCoreReuseMultiCastProgramConfig | None = None + decode_w2_prg_config: ttnn.MatmulMultiCoreReuseMultiCastProgramConfig | None = None + ff1_out_reduce_scatter_memcfg: ttnn.MemoryConfig | None = None + ff2_out_reduce_scatter_memcfg: ttnn.MemoryConfig | None = None + sharded_attn_input_memcfg: ttnn.MemoryConfig | None = None + + # Prefill settings + prefill_input_memcfg: ttnn.MemoryConfig | None = None + prefill_w1_w3_prg_config: Callable[[int], ttnn.MatmulMultiCoreReuseMultiCastProgramConfig] | None = None + prefill_w2_prg_config: Callable[[int], ttnn.MatmulMultiCoreReuseMultiCastProgramConfig] | None = None + + # Dtypes & Kernels + w1_w3_dtype: ttnn.DataType | None = None + w2_dtype: ttnn.DataType | None = None + activation_dtype: ttnn.DataType | None = None + ccl_dtype: ttnn.DataType | None = None + mul_dtype: ttnn.DataType | None = None + + ff1_3_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + ff2_compute_kernel_cfg: ttnn.WormholeComputeKernelConfig | None = None + + prefill_len_cutoff: int | None = None + + def is_resolved(self) -> bool: + """Check if all fields except optional ones are resolved.""" + # These fields are optional overrides; they can stay None to let TTNN use defaults. + optional = { + "activation_dtype", + "decode_w1_w3_prg_config", + "decode_w2_prg_config", + "ff1_out_reduce_scatter_memcfg", + "ff2_out_reduce_scatter_memcfg", + "sharded_attn_input_memcfg", + "prefill_w1_w3_prg_config", + "prefill_w2_prg_config", + } + # topology: None for single_device (CCL not needed) + if self.mesh_device and self.mesh_device.get_num_devices() == 1: + optional.add("topology") + + return all(getattr(self, f) is not None for f in self.__dataclass_fields__ if f not in optional) + + +# ============================================================================= +# MLP2D - Unified MLP for 2D-topology devices (Galaxy) +# ============================================================================= + + +class MLP2D(LightweightModule): + """ + MLP for TG (Galaxy) devices supporting both decode and prefill modes. + + Execution paths: + Unified: linear → linear → reduce_scatter(×2) → mul+silu → all_gather → linear → all_reduce + """ + + def __init__(self, w1: LazyWeight, w2: LazyWeight, w3: LazyWeight): + """ + Simple API for 90% of users - derives all config from weights. + + Args: + w1: Gate projection weight (dim, hidden_dim) + w2: Down projection weight (hidden_dim, dim) + w3: Up projection weight (dim, hidden_dim) + """ + super().__init__() + self.config = _resolve_mlp2d_config(MLP2DConfig(w1=w1, w2=w2, w3=w3)) + self._device_weights_loaded = False + + @classmethod + def from_config(cls, config: MLP2DConfig): + """ + Power API for 10% of users - any level of customization via config. + """ + # bypass the __init__ method of the base class for power users who want to customize the config + instance = object.__new__(cls) + super(MLP2D, instance).__init__() + instance.config = _resolve_mlp2d_config(config) + instance._device_weights_loaded = False + return instance + + def load_device_weights(self): + """Materialize LazyWeights onto device. Called automatically on first forward; idempotent.""" + if self._device_weights_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device weights!" + + self.w1 = self.config.w1.get_device_weight() + self.w2 = self.config.w2.get_device_weight() + self.w3 = self.config.w3.get_device_weight() + + self._device_weights_loaded = True + + def _all_reduce_tg( + self, + input_tensor: ttnn.Tensor, + cluster_axis: int, + dim: int, + sharded: bool, + memory_config: Any, + reduce_scatter_memory_config: Any = None, + ) -> ttnn.Tensor: + """ + All-reduce for TG (Galaxy) devices along specified cluster axis. + """ + cfg = self.config + # Ensure dim 0 and 1 are 1 + original_shape = input_tensor.shape + if original_shape[0] != 1 or original_shape[1] != 1: + input_tensor = ttnn.reshape( + input_tensor, (1, 1, original_shape[-4] * original_shape[-3] * original_shape[-2], original_shape[-1]) + ) + + # Cast to CCL dtype + if input_tensor.dtype != cfg.ccl_dtype: + input_tensor = ttnn.to_memory_config(input_tensor, ttnn.L1_MEMORY_CONFIG, cfg.ccl_dtype) + if sharded and memory_config is not None: + input_tensor = ttnn.to_memory_config(input_tensor, memory_config, cfg.ccl_dtype) + + if not sharded: + input_tensor = ttnn.to_memory_config(input_tensor, ttnn.DRAM_MEMORY_CONFIG) + + input_mem_cfg = input_tensor.memory_config() + # In composite all-reduce (RS + AG), the RS output memcfg can be different from the final desired memcfg. + # If not provided, fall back to the input tensor's memory config (this guarantees shard height matches). + rs_mem_cfg = ttnn.DRAM_MEMORY_CONFIG if not sharded else (reduce_scatter_memory_config or input_mem_cfg) + + reduced_tensor = ttnn.experimental.reduce_scatter_minimal_async( + input_tensor, + persistent_output_buffers=None, + dim=dim, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_rs_semaphore_handles(cluster_axis), + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), + num_links=cfg.num_reduce_scatter_links, + cluster_axis=cluster_axis, + memory_config=rs_mem_cfg, + intermediate_memory_config=ttnn.DRAM_MEMORY_CONFIG, + topology=cfg.topology, + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + + reduced_tensor = ttnn.experimental.all_gather_async( + reduced_tensor, + persistent_output_buffer=None, + dim=dim, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis), + num_links=cfg.num_all_gather_links, + cluster_axis=cluster_axis, + topology=cfg.topology, + memory_config=input_mem_cfg, + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + + reduced_tensor = ttnn.reshape(reduced_tensor, original_shape) + # Preserve requested sharding on the final output (when provided). + if sharded and memory_config is not None: + reduced_tensor = ttnn.to_memory_config(reduced_tensor, memory_config) + return reduced_tensor + + def _reduce_scatter_axis1(self, tensor: ttnn.Tensor, memory_config: Any) -> ttnn.Tensor: + """Reduce scatter along cluster axis 1.""" + cfg = self.config + cluster_axis = 1 + return ttnn.experimental.reduce_scatter_minimal_async( + tensor, + persistent_output_buffers=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_rs_semaphore_handles(cluster_axis), + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), + num_links=cfg.num_reduce_scatter_links, + cluster_axis=cluster_axis, + memory_config=memory_config, + intermediate_memory_config=ttnn.DRAM_MEMORY_CONFIG, + topology=ttnn.Topology.Linear, + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + + def _all_gather_axis1(self, tensor: ttnn.Tensor, memory_config: Any) -> ttnn.Tensor: + """All gather along cluster axis 1.""" + cfg = self.config + cluster_axis = 1 + return ttnn.experimental.all_gather_async( + tensor, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis), + num_links=2, + cluster_axis=cluster_axis, + topology=ttnn.Topology.Linear, + memory_config=memory_config, + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), + chunks_per_sync=CCL_CHUNKS_PER_SYNC, + num_workers_per_link=CCL_NUM_WORKERS_PER_LINK, + num_buffers_per_channel=CCL_NUM_BUFFERS_PER_CHANNEL, + ) + + def decode_forward(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Decode forward for TG. + + Unified Path: linear → linear → reduce_scatter(×2) → mul+silu → all_gather → linear → all_reduce + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="decode") + cfg = self.config + + # --- STAGE 1: W1/W3 Linear (L1 sharded) --- + w1_out = ttnn.linear( + x, + self.w1, + dtype=ttnn.bfloat8_b, + core_grid=None, + compute_kernel_config=cfg.ff1_3_compute_kernel_cfg, + program_config=cfg.decode_w1_w3_prg_config, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + ) + w3_out = ttnn.linear( + x, + self.w3, + dtype=ttnn.bfloat8_b, + core_grid=None, + compute_kernel_config=cfg.ff1_3_compute_kernel_cfg, + program_config=cfg.decode_w1_w3_prg_config, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + ) + ttnn.deallocate(x) + + # --- STAGE 2: CCL after W1/W3 (reduce_scatter) --- + input_mem_cfg = w1_out.memory_config() + + w1_out = self._reduce_scatter_axis1(w1_out, cfg.ff1_out_reduce_scatter_memcfg) + w3_out = self._reduce_scatter_axis1(w3_out, cfg.ff1_out_reduce_scatter_memcfg) + + # --- STAGE 3: Activation + Multiply --- + w2_in = ttnn.mul( + w1_out, + w3_out, + input_tensor_a_activations=[cfg.mlp_activation_type], + dtype=cfg.mul_dtype, + memory_config=w1_out.memory_config(), + ) + + ttnn.deallocate(w3_out) + ttnn.deallocate(w1_out) + + # --- STAGE 4: All-gather before W2 --- + w2_in = self._all_gather_axis1(w2_in, input_mem_cfg) + w2_in = ttnn.to_memory_config(w2_in, ttnn.L1_MEMORY_CONFIG) + + # --- STAGE 5: W2 Linear --- + w2_out = ttnn.linear( + w2_in, + self.w2, + compute_kernel_config=cfg.ff2_compute_kernel_cfg, + dtype=cfg.ccl_dtype, + program_config=cfg.decode_w2_prg_config, + memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, + core_grid=None, + ) + ttnn.deallocate(w2_in) + + # --- STAGE 6: Final All-Reduce --- + w2_out_reduced = self._all_reduce_tg( + w2_out, + cluster_axis=0, + dim=3, + sharded=True, + memory_config=cfg.ff2_out_reduce_scatter_memcfg, + reduce_scatter_memory_config=cfg.ff2_out_reduce_scatter_memcfg, + ) + + # --- STAGE 7: Reshape + Final memory config --- + original_shape = w2_out_reduced.shape + w2_out_reduced = ttnn.reshape( + w2_out_reduced, (1, 1, original_shape[-4] * original_shape[-3] * original_shape[-2], original_shape[-1]) + ) + # NOTE: For direct-API usage (e.g. unit tests) decode configs may leave this unset. + if cfg.sharded_attn_input_memcfg is not None: + w2_out_reduced = ttnn.to_memory_config(w2_out_reduced, cfg.sharded_attn_input_memcfg) + + return w2_out_reduced + + def prefill_forward(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Prefill forward for TG. + + Unified Path: [reshape] → linear → linear → reduce_scatter(×2) → mul+silu → all_gather → linear → all_reduce → reshape + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="prefill") + cfg = self.config + seq_len = x.shape[-2] + + # Seq_len-dependent: reshape for long sequences + if seq_len >= cfg.prefill_len_cutoff: + assert ( + seq_len % cfg.prefill_len_cutoff == 0 + ), f"seq_len ({seq_len}) must be divisible by prefill_len_cutoff ({cfg.prefill_len_cutoff})" + x = ttnn.reshape(x, [1, seq_len // cfg.prefill_len_cutoff, cfg.prefill_len_cutoff, -1]) + + # Seq_len-dependent: get program configs (None = let TTNN pick defaults) + pc_w1_w3 = cfg.prefill_w1_w3_prg_config(seq_len) if cfg.prefill_w1_w3_prg_config else None + pc_w2 = cfg.prefill_w2_prg_config(seq_len) if cfg.prefill_w2_prg_config else None + + # --- STAGE 1: W1/W3 Linear (DRAM) --- + w1_out = ttnn.linear( + x, + self.w1, + dtype=ttnn.bfloat8_b, + core_grid=None, + compute_kernel_config=cfg.ff1_3_compute_kernel_cfg, + program_config=pc_w1_w3, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + w3_out = ttnn.linear( + x, + self.w3, + dtype=ttnn.bfloat8_b, + core_grid=None, + compute_kernel_config=cfg.ff1_3_compute_kernel_cfg, + program_config=pc_w1_w3, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + ttnn.deallocate(x) + + # --- STAGE 2: CCL after W1/W3 (reduce_scatter for prefill) --- + input_mem_cfg = w1_out.memory_config() + + w1_out = self._reduce_scatter_axis1(w1_out, None) # None mem_config for prefill + w3_out = self._reduce_scatter_axis1(w3_out, None) + + # --- STAGE 3: Activation + Multiply --- + w2_in = ttnn.mul( + w1_out, + w3_out, + input_tensor_a_activations=[cfg.mlp_activation_type], + dtype=cfg.mul_dtype, + memory_config=w1_out.memory_config(), + ) + + ttnn.deallocate(w3_out) + ttnn.deallocate(w1_out) + + # --- STAGE 4: All-gather before W2 --- + w2_in = self._all_gather_axis1(w2_in, input_mem_cfg) + # No L1 conversion for prefill + + # --- STAGE 5: W2 Linear --- + w2_out = ttnn.linear( + w2_in, + self.w2, + compute_kernel_config=cfg.ff2_compute_kernel_cfg, + dtype=cfg.ccl_dtype, + program_config=pc_w2, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + core_grid=None, + ) + ttnn.deallocate(w2_in) + + # --- STAGE 6: Final All-Reduce --- + w2_out_reduced = self._all_reduce_tg( + w2_out, + cluster_axis=0, + dim=3, + sharded=False, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + + # --- STAGE 7: Reshape (no final memory config change for prefill) --- + original_shape = w2_out_reduced.shape + w2_out_reduced = ttnn.reshape( + w2_out_reduced, (1, 1, original_shape[-4] * original_shape[-3] * original_shape[-2], original_shape[-1]) + ) + + return w2_out_reduced + + def forward(self, x: ttnn.Tensor | LazyWeight, mode: str) -> ttnn.Tensor: + """Dispatch to the appropriate forward method based on mode.""" + if mode == "decode": + return self.decode_forward(x) + else: + return self.prefill_forward(x) + + @classmethod + def from_model_args( + cls, + mesh_device, + tt_ccl, + args, + state_dict, + weight_cache_path, + layer_num: int, + state_dict_prefix: Optional[str] = None, + ): + """Factory method for backward compatibility with ModelArgs.""" + # MLP2D requires Galaxy topology (4x8 or 8x4) due to Galaxy-specific CCL operations + valid_shapes = [(4, 8), (8, 4)] + shape_tuple = tuple(args.cluster_shape) + if shape_tuple not in valid_shapes: + # IMPORTANT: do this validation before touching mesh_device/tt_ccl/model_config + # so negative tests don't need to open a mesh device or initialize fabric. + raise ValueError( + f"MLP2D requires Galaxy topology (8x4). Got cluster_shape={args.cluster_shape}. " + "For non-Galaxy devices, use MLP1D instead." + ) + + import torch + + from models.tt_transformers.tt.model_config import OpGroup, TensorGroup + + # Get model_config for overrides + model_config = args.get_model_config() + decoders_opt = model_config.get("DECODERS_OPTIMIZATIONS") + effective_layer_num = max(layer_num, 0) + + # Extract settings + ccl_topology = args.ccl_topology() + + if state_dict_prefix is None: + state_dict_prefix = args.get_state_dict_prefix("MLP", layer_num) + + # Get dtypes + ff1_3_dtype = decoders_opt.get_tensor_dtype(decoder_id=effective_layer_num, tensor=TensorGroup.FF1_FF3) + ff2_dtype = decoders_opt.get_tensor_dtype(decoder_id=effective_layer_num, tensor=TensorGroup.FF2) + activation_dtype = decoders_opt.get_tensor_dtype(decoder_id=effective_layer_num, tensor=TensorGroup.ACTIVATION) + + # Get compute kernel configs + ff1_3_compute_kernel_cfg = decoders_opt.get_math_fidelity( + decoder_id=effective_layer_num, op=OpGroup.LI_FF1_FF3, configuration=args + ) + ff2_compute_kernel_cfg = decoders_opt.get_math_fidelity( + decoder_id=effective_layer_num, op=OpGroup.LI_FF2, configuration=args + ) + + # Get decode program configs from model_config + # Note: Handling legacy small dim behavior by setting configs to None if implicit check failed. + # TTTv1 TG configs (FF1_3_TG_PROGCFG, FF2_TG_PROGCFG) assume shard layouts that only + # work for models with dim >= 8192 (e.g. Llama-70B). Smaller models fall back to TTNN defaults. + _MIN_DIM_FOR_TG_DECODE_CONFIGS = 8192 + + decode_w1_w3_prg_config = model_config.get("FF1_3_TG_PROGCFG") + if args.dim < _MIN_DIM_FOR_TG_DECODE_CONFIGS: + # TT-Transformers TG FF2 config assumes a specific intermediate sharding that + # doesn't match this MLP2D implementation for small dim. Let TTNN pick defaults. + decode_w1_w3_prg_config = None + + decode_w2_prg_config = model_config.get("FF2_TG_PROGCFG") + if args.dim < _MIN_DIM_FOR_TG_DECODE_CONFIGS: + decode_w2_prg_config = None + + # Memory configs + ff1_out_reduce_scatter_memcfg = model_config.get("FF1_OUT_REDUCE_SCATTER_MEMCFG") + + # TT-Transformers config uses shard height 32*4 here; MLP2D tensors are height 32. + # Passing this into to_memory_config can TT_FATAL on shard-height mismatch. + ff2_out_reduce_scatter_memcfg = model_config.get("FF2_OUT_REDUCE_SCATTER_MEMCFG") + if args.dim < _MIN_DIM_FOR_TG_DECODE_CONFIGS: + # Some TT-Transformers configs size this as shard_height=32*cluster_rows (e.g. 256 on 8x4), + # but MLP2D decode tensors here are height=32. Use the attention-input sharding instead. + ff2_out_reduce_scatter_memcfg = model_config.get("SHARDED_ATTN_INPUT_MEMCFG") + + sharded_attn_input_memcfg = model_config.get("SHARDED_ATTN_INPUT_MEMCFG") + + # Prefill configs + prefill_w1_w3_prg_config_factory = model_config.get("PREFILL_MLP_W1_W3_PRG_CONFIG") + prefill_w2_prg_config_factory = model_config.get("PREFILL_MLP_W2_PRG_CONFIG") + + cache_dir = None if args.dummy_weights else Path(weight_cache_path) / state_dict_prefix + hidden_dim_string = f".hidden_dim_{args.hidden_dim}" if args.hidden_dim != args.unpadded_hidden_dim else "" + + def make_weight_source(name: str, pad_dim: int): + tensor = torch.transpose(state_dict[f"{state_dict_prefix}.{name}.weight"], -2, -1) + return pad_dim_to_size(tensor, dim=pad_dim, size=args.hidden_dim) + + # 2D sharding dims for weights + w1_shard_dims = (-1, -2) + w2_shard_dims = (-2, -1) + + w1 = LazyWeight( + source=make_weight_source("w1", -1), + dtype=ff1_3_dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(w1_shard_dims[0]), ttnn.PlacementShard(w1_shard_dims[1])], + mesh_shape_override=ttnn.MeshShape(args.cluster_shape), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, # TG uses DRAM for weights + cache_dir_weight_name=(cache_dir, f"w1_sharded{hidden_dim_string}") if cache_dir else None, + ) + w2 = LazyWeight( + source=make_weight_source("w2", -2), + dtype=ff2_dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(w2_shard_dims[0]), ttnn.PlacementShard(w2_shard_dims[1])], + mesh_shape_override=ttnn.MeshShape(args.cluster_shape), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(cache_dir, f"w2_sharded{hidden_dim_string}") if cache_dir else None, + ) + w3 = LazyWeight( + source=make_weight_source("w3", -1), + dtype=ff1_3_dtype, + device=mesh_device, + mesh_mapper_config=ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(w1_shard_dims[0]), ttnn.PlacementShard(w1_shard_dims[1])], + mesh_shape_override=ttnn.MeshShape(args.cluster_shape), + ), + layout=ttnn.TILE_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(cache_dir, f"w3_sharded{hidden_dim_string}") if cache_dir else None, + ) + + config = MLP2DConfig( + w1=w1, + w2=w2, + w3=w3, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + dim=args.dim, + hidden_dim=args.hidden_dim, + max_batch_size=args.max_batch_size, + mlp_activation_type=getattr(args, "mlp_activation_type", ttnn.UnaryOpType.SILU), + topology=ccl_topology, + decode_w1_w3_prg_config=decode_w1_w3_prg_config, + decode_w2_prg_config=decode_w2_prg_config, + ff1_out_reduce_scatter_memcfg=ff1_out_reduce_scatter_memcfg, + ff2_out_reduce_scatter_memcfg=ff2_out_reduce_scatter_memcfg, + sharded_attn_input_memcfg=sharded_attn_input_memcfg, + prefill_w1_w3_prg_config=prefill_w1_w3_prg_config_factory, + prefill_w2_prg_config=prefill_w2_prg_config_factory, + w1_w3_dtype=ff1_3_dtype, + w2_dtype=ff2_dtype, + activation_dtype=activation_dtype, + ccl_dtype=args.ccl_dtype, + ff1_3_compute_kernel_cfg=ff1_3_compute_kernel_cfg, + ff2_compute_kernel_cfg=ff2_compute_kernel_cfg, + ) + + return cls.from_config(config) + + +# ============================================================================= +# Helper functions +# ============================================================================= + + +# todo)) work with the CCL team to find opportunity to simplify this --> e.g., build into TTNN APIs? + + +def _compute_kernel_config_hifi2_fp16() -> ttnn.WormholeComputeKernelConfig: + """Default compute kernel config for MLP (HiFi2 with FP16 accumulation).""" + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=True, + ) + + +def _resolve_mlp2d_config(config: MLP2DConfig) -> MLP2DConfig: + """Materialize the config to known good defaults using replace pattern.""" + to_set = {} + + # --- Phase 1: Foundational fields --- + + # Derive dimensions + dim = config.dim + if config.dim is None: + dim = config.w1.source.shape[-2] + to_set["dim"] = dim + + hidden_dim = config.hidden_dim + if config.hidden_dim is None: + hidden_dim = config.w1.source.shape[-1] + to_set["hidden_dim"] = hidden_dim + + # Derive mesh_device + mesh_device = config.mesh_device + if mesh_device is None: + mesh_device = config.w1.device + if mesh_device is None: + mesh_device = ttnn.GetDefaultDevice() + if config.mesh_device is None: + to_set["mesh_device"] = mesh_device + + assert mesh_device is not None + cluster_shape = list(mesh_device.shape) + # MLP2D is designed for 2D mesh topologies (cluster_shape[0] > 1 and cluster_shape[1] > 1) + # Note: from_model_args() enforces Galaxy (4x8 or 8x4) because it uses model_config.py + # which has Galaxy-specific hardcoded values. Direct MLP2DConfig usage is more flexible. + assert cluster_shape[0] > 1 and cluster_shape[1] > 1, ( + f"MLP2D requires 2D mesh (both cluster_shape dimensions > 1). " + f"Got cluster_shape={cluster_shape}. For 1D meshes, use MLP1D instead." + ) + + # Derive tt_ccl + tt_ccl = config.tt_ccl + if config.tt_ccl is None: + tt_ccl = get_tt_ccl(mesh_device) + to_set["tt_ccl"] = tt_ccl + + # Auto-detect topology + topology = config.topology + if config.topology is None: + topology = default_topology(mesh_device) + to_set["topology"] = topology + + # --- Phase 2: Dtypes and Tunings --- + + w1_w3_dtype = config.w1_w3_dtype or ttnn.bfloat8_b + to_set["w1_w3_dtype"] = w1_w3_dtype + w2_dtype = config.w2_dtype or ttnn.bfloat8_b + to_set["w2_dtype"] = w2_dtype + + if config.ccl_dtype is None: + to_set["ccl_dtype"] = ttnn.bfloat8_b + if config.mul_dtype is None: + to_set["mul_dtype"] = config.activation_dtype or ttnn.bfloat8_b + + if config.prefill_len_cutoff is None: + to_set["prefill_len_cutoff"] = 512 if is_blackhole() else 1024 + + # Compute kernel configs + if config.ff1_3_compute_kernel_cfg is None: + to_set["ff1_3_compute_kernel_cfg"] = _compute_kernel_config_hifi2_fp16() + if config.ff2_compute_kernel_cfg is None: + to_set["ff2_compute_kernel_cfg"] = _compute_kernel_config_hifi2_fp16() + + # --- Phase 2.5: Input Memory Configs --- + + if config.decode_input_memcfg is None: + to_set["decode_input_memcfg"] = ttnn.L1_MEMORY_CONFIG + + if config.prefill_input_memcfg is None: + to_set["prefill_input_memcfg"] = ttnn.DRAM_MEMORY_CONFIG + + # --- Phase 3: Prefill Program Configs --- + # NOTE: prefill_w1_w3_prg_config and prefill_w2_prg_config are optional. + # When None, TTNN picks defaults. Only from_model_args (Power API) provides these. + # This keeps the Simple API working without complex 2D-aware program config generation. + + # --- Phase 4: Resolve Weights (always 2D sharded for MLP2D) --- + + # TG weights use DRAM interleaved (no specific shard memory config on weights themselves) + w1_w3_memcfg = config.w1_w3_memcfg or ttnn.DRAM_MEMORY_CONFIG + to_set["w1_w3_memcfg"] = w1_w3_memcfg + w2_memcfg = config.w2_memcfg or ttnn.DRAM_MEMORY_CONFIG + to_set["w2_memcfg"] = w2_memcfg + + # MLP2D ALWAYS uses 2D sharding - this is fundamental to how 2D mesh MLP works. + # w1/w3: shard dims (-1, -2) = N sharded on mesh axis 0, K sharded on mesh axis 1 + # w2: shard dims (-2, -1) = K sharded on mesh axis 0, N sharded on mesh axis 1 + w1_w3_shard_dims = (-1, -2) + w2_shard_dims = (-2, -1) + + def get_weight_mesh_mapper(lazy_weight: LazyWeight, shard_dims: tuple[int, int]): + """Return existing mesh_mapper_config if set, else create 2D shard mapper.""" + existing = getattr(lazy_weight, "mesh_mapper_config", None) + if existing is not None: + return existing + # Default: apply 2D sharding + return ttnn.MeshMapperConfig( + placements=[ttnn.PlacementShard(shard_dims[0]), ttnn.PlacementShard(shard_dims[1])], + mesh_shape_override=ttnn.MeshShape(cluster_shape), + ) + + to_set["w1"] = resolve_lazy_weight( + config.w1, + device=mesh_device, + memory_config=w1_w3_memcfg, + mesh_mapper_config=get_weight_mesh_mapper(config.w1, w1_w3_shard_dims), + layout=ttnn.TILE_LAYOUT, + dtype=w1_w3_dtype, + ) + + to_set["w2"] = resolve_lazy_weight( + config.w2, + device=mesh_device, + memory_config=w2_memcfg, + mesh_mapper_config=get_weight_mesh_mapper(config.w2, w2_shard_dims), + layout=ttnn.TILE_LAYOUT, + dtype=w2_dtype, + ) + + to_set["w3"] = resolve_lazy_weight( + config.w3, + device=mesh_device, + memory_config=w1_w3_memcfg, + mesh_mapper_config=get_weight_mesh_mapper(config.w3, w1_w3_shard_dims), + layout=ttnn.TILE_LAYOUT, + dtype=w1_w3_dtype, + ) + + resolved_config = replace(config, **to_set) + assert resolved_config.is_resolved(), "Config must be resolved!" + return resolved_config + + +def _load_input_device_tensor(x: ttnn.Tensor | LazyWeight, config: MLP2DConfig, mode: str) -> ttnn.Tensor: + """Resolve the input tensor to ttnn tensor if x is a LazyWeight, otherwise return as is.""" + assert mode in ["decode", "prefill"], "mode must be one of decode or prefill!" + mem_cfg = config.decode_input_memcfg if mode == "decode" else config.prefill_input_memcfg + if isinstance(x, LazyWeight): + # For MLP2D, input must be sharded to match weight sharding: + # - w1/w3 shard dims = (-1, -2): K sharded on mesh axis 1 + # - Input [batch, 1, seq, K]: shard K (dim -1) on mesh axis 1, replicate on axis 0 + cluster_shape = list(config.mesh_device.shape) + input_mesh_mapper = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape(cluster_shape), + ) + resolved_x = resolve_lazy_weight( + x, + device=config.mesh_device, + memory_config=mem_cfg, + mesh_mapper_config=input_mesh_mapper, + layout=ttnn.TILE_LAYOUT, + ) + return resolved_x.get_device_weight() + + assert isinstance(x, ttnn.Tensor), "x must be a ttnn tensor at this point!" + return x diff --git a/code/models/common/modules/moe/configs/mistral_large_3.yaml b/code/models/common/modules/moe/configs/mistral_large_3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fcd84b4396b6002fa2184016849dac3b619ae449 --- /dev/null +++ b/code/models/common/modules/moe/configs/mistral_large_3.yaml @@ -0,0 +1,28 @@ +# mistral_large_3 +mesh_shape: [16, 4] +cluster_axis: 0 +batch_per_device: 32 +hidden_size: 7168 +select_experts_k: 4 +num_routed_experts: 128 +num_shared_experts: 1 +has_bias: false + +# gate (routing front-end): n_group 1 = generalized op (ungrouped) / 8 = deepseek op (grouped top-8). +# Mistral Large 3 router = linear → top-4 (by raw logit) → softmax over the selected, NO bias, no scale. +# TTMoEGate's score_func=softmax is exactly that (rank by logit, then exp-over-selected; ≡ softmax-front + +# renorm, Z cancels), scale 1.0. 128 experts → padded to the op's 256-face. (Like gpt_oss but bias-free.) +n_group: 1 +score_func: softmax +routed_scaling_factor: 1.0 + +compute: + intermediate_size: 4096 + activation_type: SILU + +reduce: + shared_expert_scale: 1.0 + +experts: + expert_mapping: sequential + shared_expert_ids_to_devices: fully_replicated diff --git a/code/models/common/modules/moe/tt_moe_decode.py b/code/models/common/modules/moe/tt_moe_decode.py new file mode 100644 index 0000000000000000000000000000000000000000..184d7ba9977f887b7d41090dc9921d2634ee6a1d --- /dev/null +++ b/code/models/common/modules/moe/tt_moe_decode.py @@ -0,0 +1,1020 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. + +# SPDX-License-Identifier: Apache-2.0 + +"""Configurable Mixture-of-Experts decode block (`TTMoEDecode`). + +A single forward step of a decode-time MoE layer on a 2D device mesh, wrapping the +ttnn `all_to_all_dispatch_metadata` → `moe_compute` → `deepseek_moe_fast_reduce_nc_fused` +→ reduce-scatter pipeline. All op kwargs (memory configs, cluster topology, splits, +shared-expert plumbing) are driven by `TTMoEDecodeConfig`, which derives sane defaults +from a minimal YAML per model. This module's job is just to wire the configured pieces +together and expose a `forward(x, scores, indices)` interface. + +Pipeline overview (one decode step): + 1. `all_to_all_dispatch_metadata`: route each token to its `select_experts_k` chosen + routed experts (cross-cluster send) plus all shared experts (local broadcast). + 2. `moe_compute`: per-device per-token-slot matmul `x @ w0`, `x @ w1`, activation + (SiLU/SWIGLU/GELU), `intermediate @ w2`, optionally with bias. + 3. `deepseek_moe_fast_reduce_nc_fused`: score-weighted combine of the per-expert + outputs back to per-token results, with a fixed scalar `shared_expert_scale` + applied to shared-expert contributions. + 4. Reduce-scatter across the replicated mesh axis to produce per-device output + chunks of width `hidden_size / num_replicated`. + +Weight ownership: routed experts are sharded across the dispatch axis; shared experts +are replicated. Weights upload as `bfloat4_b` (with bf16 intermediate tiles for bias). + +The two private classes `_TTMoEDecodeExpertState` and `_TTMoEDecodeBuffers` exist to +keep weight/mapping init and per-iteration scratch separated from the forward logic. +""" + +from __future__ import annotations + +import torch +from loguru import logger +from ttnn.experimental.moe_compute_utils import ( + auto_output_width_shard_dim, + effective_matmul_ring_size, + map_shared_experts, +) + +import ttnn +from models.common.modules.moe.tt_moe_decode_config import TTMoEDecodeConfig + + +def _tt_to_torch_dtype(tt_dtype): + """Map a ttnn dtype to the closest host torch dtype for buffer allocation. + + Only the dtypes this module actually uses are handled — `bfloat8_b` falls back to + `torch.bfloat16` since torch has no native 8-bit float; the host tensor is just a + placeholder that gets reinterpreted at upload time. + """ + if tt_dtype == ttnn.bfloat16 or tt_dtype == ttnn.bfloat8_b: + return torch.bfloat16 + if tt_dtype == ttnn.float32: + return torch.float32 + if tt_dtype == ttnn.uint16: + return torch.uint16 + raise ValueError(f"Unsupported tt dtype: {tt_dtype}") + + +class _TTMoEDecodeExpertState: + """Owns per-layer routed + shared expert weights and the global expert-mapping table. + + Holds three uploaded ttnn tensors after init: + - `tt_expert_mapping`: `[num_devices, num_experts]` lookup of which linearized mesh + coord owns each expert, replicated to every device. Used by both `dispatch` and + `fast_reduce` ops. + - `tt_w0_w1`: interleaved-and-tile-reordered w0/w1 weights for `moe_compute`'s + consumption, sharded across mesh devices along the expert dim. `bfloat4_b`. + - `tt_w2`: same idea, but with the ring-rotated N-tile layout w2 needs. `bfloat4_b`. + + Bias support: when `has_bias=True`, biases are folded into the same prepared weight + tensors via `prepare_w0_w1_tensor_with_bias` / `prepare_w2_tensor_with_bias` (one + extra K/N tile each). Bias + shared experts together is not supported (raises). + """ + + def _load_weights(self): + # TODO (AM) eventually support loading weights from path + pass + + @staticmethod + def _validate( + torch_w0: "torch.Tensor", + torch_w1: "torch.Tensor", + torch_w2: "torch.Tensor", + mesh_device: ttnn.MeshDevice, + mesh_shape: tuple[int, int], + cluster_axis: int, + has_bias: bool, + num_routed_experts: int, + expert_mapping: list[int], + num_shared_experts: int, + shared_expert_ids_to_devices: dict[int, list[int]] | None, + shared_id_to_torch_w0: dict[int, "torch.Tensor"] | None, + shared_id_to_torch_w1: dict[int, "torch.Tensor"] | None, + shared_id_to_torch_w2: dict[int, "torch.Tensor"] | None, + torch_b0: "torch.Tensor" | None, + torch_b1: "torch.Tensor" | None, + torch_b2: "torch.Tensor" | None, + ) -> None: + """Fail fast on (weights, biases, shared experts) ↔ config mismatch. + + Cross-checks tensor shapes against each other (w0 is the source of truth for + `L`, `H`, `N`) and against the config-derived counts. Catches the common + misconfigurations that otherwise surface as opaque shape errors deep inside + the bf4 preparers or kernels. + """ + # --- mesh / topology sanity --- + if cluster_axis not in (0, 1): + raise ValueError(f"cluster_axis must be 0 or 1, got {cluster_axis}") + num_devices = mesh_device.get_num_devices() + if mesh_shape[0] * mesh_shape[1] != num_devices: + raise ValueError( + f"mesh_shape {mesh_shape} (= {mesh_shape[0] * mesh_shape[1]} devices) does " + f"not match mesh_device.get_num_devices() = {num_devices}" + ) + if num_routed_experts % num_devices != 0: + raise ValueError( + f"num_routed_experts ({num_routed_experts}) must be divisible by num_devices ({num_devices})" + ) + + # --- routed weight shape cross-checks (w0 is the source of truth for L, H, N) --- + if torch_w0.ndim != 4: + raise ValueError(f"torch_w0 must be 4D [L, E, H, N], got shape {tuple(torch_w0.shape)}") + L, E, H, N = torch_w0.shape + if E != num_routed_experts: + raise ValueError(f"torch_w0.shape[1] = {E} does not match num_routed_experts ({num_routed_experts})") + if tuple(torch_w1.shape) != (L, E, H, N): + raise ValueError(f"torch_w1 shape {tuple(torch_w1.shape)} must match torch_w0 shape ({L}, {E}, {H}, {N})") + if tuple(torch_w2.shape) != (L, E, N, H): + raise ValueError(f"torch_w2 shape {tuple(torch_w2.shape)} must be (L, E, N, H) = ({L}, {E}, {N}, {H})") + + # --- expert_mapping --- + if len(expert_mapping) != num_routed_experts: + raise ValueError( + f"expert_mapping length {len(expert_mapping)} != num_routed_experts ({num_routed_experts})" + ) + bad = [(e, d) for e, d in enumerate(expert_mapping) if not (0 <= d < num_devices)] + if bad: + raise ValueError(f"expert_mapping has out-of-range device ids (0..{num_devices - 1}): {bad[:5]}") + + # --- bias presence and shapes --- + bias_tensors = (torch_b0, torch_b1, torch_b2) + any_bias = any(b is not None for b in bias_tensors) + all_bias = all(b is not None for b in bias_tensors) + if has_bias and not all_bias: + missing = [name for name, b in zip(("b0", "b1", "b2"), bias_tensors) if b is None] + raise ValueError(f"has_bias=True but {missing} not provided") + if not has_bias and any_bias: + raise ValueError("has_bias=False but one or more of torch_b0/b1/b2 was provided") + if has_bias: + if tuple(torch_b0.shape) != (L, E, N) or tuple(torch_b1.shape) != (L, E, N): + raise ValueError( + f"b0/b1 must be (L, E, N) = ({L}, {E}, {N}); got " + f"{tuple(torch_b0.shape)} and {tuple(torch_b1.shape)}" + ) + if tuple(torch_b2.shape) != (L, E, H): + raise ValueError(f"b2 must be (L, E, H) = ({L}, {E}, {H}); got {tuple(torch_b2.shape)}") + + # --- shared experts --- + if num_shared_experts > 0: + if shared_expert_ids_to_devices is None: + raise ValueError(f"num_shared_experts={num_shared_experts} but shared_expert_ids_to_devices is None") + if len(shared_expert_ids_to_devices) != num_shared_experts: + raise ValueError( + f"shared_expert_ids_to_devices has {len(shared_expert_ids_to_devices)} entries " + f"but num_shared_experts={num_shared_experts}" + ) + + shared_experts_per_device = [0] * num_devices + for edl in shared_expert_ids_to_devices.values(): + for d in edl: + shared_experts_per_device[d] += 1 + if len(set(shared_experts_per_device)) > 1 or 0 in shared_experts_per_device: + raise ValueError( + "Every device, should have the same number of, and at least 1 shared expert:" + f" {shared_experts_per_device=}" + ) + + expected_ids = set(range(num_routed_experts, num_routed_experts + num_shared_experts)) + if set(shared_expert_ids_to_devices.keys()) != expected_ids: + raise ValueError( + f"shared expert ids must be contiguous after routed: expected {sorted(expected_ids)}, " + f"got {sorted(shared_expert_ids_to_devices.keys())}" + ) + shared_dicts = (shared_id_to_torch_w0, shared_id_to_torch_w1, shared_id_to_torch_w2) + if any(d is None for d in shared_dicts): + raise ValueError("num_shared_experts > 0 but shared_id_to_torch_w0/w1/w2 not all provided") + for name, d in zip(("w0", "w1", "w2"), shared_dicts): + if set(d.keys()) != expected_ids: + raise ValueError( + f"shared_id_to_torch_{name} keys {sorted(d.keys())} != " + f"shared expert ids {sorted(expected_ids)}" + ) + expected_shape = (L, 1, N, H) if name == "w2" else (L, 1, H, N) + for sid, t in d.items(): + if tuple(t.shape) != expected_shape: + raise ValueError( + f"shared_id_to_torch_{name}[{sid}] shape {tuple(t.shape)} != expected {expected_shape}" + ) + if has_bias: + # Mirrors the runtime check in __init__; surface it here too so it fails + # before any device upload. + raise NotImplementedError("bias + shared experts is not yet supported") + else: + if shared_expert_ids_to_devices: + raise ValueError( + f"num_shared_experts=0 but shared_expert_ids_to_devices is non-empty: " + f"{shared_expert_ids_to_devices}" + ) + if any(d is not None for d in (shared_id_to_torch_w0, shared_id_to_torch_w1, shared_id_to_torch_w2)): + raise ValueError("num_shared_experts=0 but one or more shared_id_to_torch_w* dicts were provided") + + @staticmethod + def _init_expert_mapping( + torch_expert_mapping: "torch.Tensor", + shared_expert_ids_to_devices: dict[int, list[int]] | None, + mesh_device: ttnn.MeshDevice, + mesh_shape: tuple[int, int], + cluster_axis: int, + ): + """Build and upload the `[num_devices, num_experts]` expert-mapping table. + + `mapping[d, e]` = linearized mesh coord of the device that owns expert `e`. + Routed experts have the same value across all source rows `d` (ownership is + global), so we just `repeat` the 1D input. Shared experts let different source + devices pick different replicas based on cluster distance — `map_shared_experts` + rewrites those columns to the nearest replica per source row. + + Matches the "new format" used by `test_all_to_all_dispatch_metadata_6U.py` and + `gen_expert_mapping` in `test_moe_compute_6U.py`. Replicated to every device + because both dispatch and fast-reduce need the full lookup locally. + """ + if torch_expert_mapping.ndim != 1: + raise ValueError( + f"expected 1D expert_mapping (length=num_experts), got shape {tuple(torch_expert_mapping.shape)}" + ) + num_devices = mesh_device.get_num_devices() + mapping_2d = torch_expert_mapping.to(torch.int32).unsqueeze(0).repeat(num_devices, 1) + + if shared_expert_ids_to_devices is not None: + mapping_2d = map_shared_experts(mapping_2d, shared_expert_ids_to_devices, mesh_shape, cluster_axis) + + return ttnn.from_torch( + mapping_2d, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=ttnn.uint16, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device), + ) + + @staticmethod + def _device_reorder_weights( + torch_expert_mapping: "torch.Tensor", + torch_w0: "torch.Tensor", + torch_w1: "torch.Tensor", + torch_w2: "torch.Tensor", + torch_b0: "torch.Tensor" | None, + torch_b1: "torch.Tensor" | None, + torch_b2: "torch.Tensor" | None, + ): + """Permute the expert dim so that `ShardTensorToMesh(dim=experts)` lands each + expert on its assigned device. + + The host weights come in routed-expert-id order (`[L, num_experts, ...]`), but + sharding by the expert dim splits contiguously — so without reordering, device 0 + gets experts [0..E/D), device 1 gets [E/D..2E/D), etc. `expert_mapping[e]` tells + us the *target* device for expert `e`; `argsort` (stable) groups experts that + share a target device into contiguous chunks in the right order. Same permutation + applies to biases when present. + """ + perm = torch.argsort(torch_expert_mapping, stable=True) + mapped_tensors = [t[:, perm, :, :] for t in (torch_w0, torch_w1, torch_w2)] + + if torch_b0 is not None: + mapped_tensors += [t[:, perm, :] for t in (torch_b0, torch_b1, torch_b2)] + else: + mapped_tensors += [None] * 3 + + return tuple(mapped_tensors) + + def __init__( + self, + mesh_device: ttnn.MeshDevice, + torch_w0: "torch.Tensor", + torch_w1: "torch.Tensor", + torch_w2: "torch.Tensor", + *, + mesh_shape: tuple[int, int], + cluster_axis: int, + has_bias: bool, + num_routed_experts: int, + expert_mapping: list[int], + num_shared_experts: int, + shared_expert_ids_to_devices: dict[int, list[int]] | None = None, + shared_id_to_torch_w0: dict[int, "torch.Tensor"] | None = None, + shared_id_to_torch_w1: dict[int, "torch.Tensor"] | None = None, + shared_id_to_torch_w2: dict[int, "torch.Tensor"] | None = None, + torch_b0: "torch.Tensor" | None = None, + torch_b1: "torch.Tensor" | None = None, + torch_b2: "torch.Tensor" | None = None, + ): + """Prepare and upload all expert state to the mesh. + + Pipeline: argsort-permute routed weights on host to match the device assignment + (`_device_reorder_weights`), upload them to the mesh sharded on the experts dim + (dim 1) as bf16, optionally splice in shared experts on-device so each device holds + `routed_per_device + shared_per_device` slots (`ttnn.experimental.add_shared_expert_weights`), + then run the C++ on-device packers (`ttnn.experimental.prepare_*`) and quantize to + `bfloat4_b` (`ttnn.experimental.quantize_weights_via_host`) — see + `_init_total_expert_weights_impl`. + + Routed weight shapes (host, post-permute): `w0/w1 = [L, num_routed, H, N]`, + `w2 = [L, num_routed, N, H]`. Shared weights are dicts keyed by global expert id; + each value has shape `[L, 1, ...]` matching the routed layout. + + Biases match the routed shape minus the matmul-output dim (`[L, num_routed, N]` + for `b0/b1`, `[L, num_routed, H]` for `b2`). Shared experts + bias is not yet + supported because the shared splice doesn't carry bias rows — would need a parallel API. + """ + self._validate( + torch_w0, + torch_w1, + torch_w2, + mesh_device, + mesh_shape, + cluster_axis, + has_bias, + num_routed_experts, + expert_mapping, + num_shared_experts, + shared_expert_ids_to_devices, + shared_id_to_torch_w0, + shared_id_to_torch_w1, + shared_id_to_torch_w2, + torch_b0, + torch_b1, + torch_b2, + ) + + # An empty mapping ({} when num_shared_experts==0) means "no shared experts" — + # normalize to None so the shared-expert paths below are skipped entirely. + if not shared_expert_ids_to_devices: + shared_expert_ids_to_devices = None + + num_routed = torch_w0.shape[1] + logger.info( + f"Initializing expert state: routed_experts={num_routed} num_shared={num_shared_experts} " + f"mesh_shape={mesh_shape} cluster_axis={cluster_axis} has_bias={has_bias}" + ) + + torch_expert_mapping = torch.tensor(expert_mapping, dtype=torch.int) + + ( + mapped_torch_w0, + mapped_torch_w1, + mapped_torch_w2, + mapped_torch_b0, + mapped_torch_b1, + mapped_torch_b2, + ) = self._device_reorder_weights( + torch_expert_mapping, torch_w0, torch_w1, torch_w2, torch_b0, torch_b1, torch_b2 + ) + + num_devices = mesh_device.get_num_devices() + num_layers = torch_w0.shape[0] + hidden_size = torch_w0.shape[-2] + intermediate_size = torch_w0.shape[-1] + routed_per_device = num_routed // num_devices + + # Upload the (reordered) routed weights to the mesh sharded on the experts dim (dim 1), + # bf16 / ROW_MAJOR, so the C++ `ttnn.experimental.prepare_*` helpers can do the layout + # transform on-device. Each device receives its assigned contiguous expert chunk. + def _shard_experts(t): + return ttnn.from_torch( + t, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=ttnn.bfloat16, + mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=1), + ) + + tt_w0 = _shard_experts(mapped_torch_w0) + tt_w1 = _shard_experts(mapped_torch_w1) + tt_w2 = _shard_experts(mapped_torch_w2) + + if shared_expert_ids_to_devices is not None: + if has_bias: + # add_shared_expert_weights only handles weights; extending it to bias + # requires a parallel API and per-shared-expert bias tensors. + raise NotImplementedError("bias + shared experts is not yet supported") + + logger.info(f"Adding shared expert weights for {len(shared_expert_ids_to_devices)} shared experts") + # The C++ add_shared_expert_weights consumes pre-arranged shared *device* tensors + # (per device: its assigned shared experts in sorted-id order; then concatenated + # across devices). Arrange on host, upload sharded on dim 1, splice on-device. + tt_shared_w0 = _shard_experts( + self._arrange_shared(shared_id_to_torch_w0, shared_expert_ids_to_devices, num_devices) + ) + tt_shared_w1 = _shard_experts( + self._arrange_shared(shared_id_to_torch_w1, shared_expert_ids_to_devices, num_devices) + ) + tt_shared_w2 = _shard_experts( + self._arrange_shared(shared_id_to_torch_w2, shared_expert_ids_to_devices, num_devices) + ) + routed_w0, routed_w1, routed_w2 = tt_w0, tt_w1, tt_w2 + tt_w0, tt_w1, tt_w2 = ttnn.experimental.add_shared_expert_weights( + routed_w0, + routed_w1, + routed_w2, + tt_shared_w0, + tt_shared_w1, + tt_shared_w2, + cluster_axis=cluster_axis, + ) + for t in (routed_w0, routed_w1, routed_w2, tt_shared_w0, tt_shared_w1, tt_shared_w2): + ttnn.deallocate(t) + + total_shared_slots = sum(len(devs) for devs in shared_expert_ids_to_devices.values()) + experts_per_device = (num_routed + total_shared_slots) // num_devices + else: + experts_per_device = routed_per_device + + tt_b0 = tt_b1 = tt_b2 = None + if has_bias: + tt_b0 = _shard_experts(mapped_torch_b0) + tt_b1 = _shard_experts(mapped_torch_b1) + tt_b2 = _shard_experts(mapped_torch_b2) + + self.tt_expert_mapping = self._init_expert_mapping( + torch_expert_mapping, shared_expert_ids_to_devices, mesh_device, mesh_shape, cluster_axis + ) + logger.info("Uploaded expert mapping to mesh") + + self.tt_w0_w1, self.tt_w2 = self._init_total_expert_weights_impl( + tt_w0, + tt_w1, + tt_w2, + mesh_device, + num_layers, + experts_per_device, + hidden_size, + intermediate_size, + has_bias, + tt_b0, + tt_b1, + tt_b2, + ) + + @staticmethod + def _arrange_shared( + shared_dict: dict[int, "torch.Tensor"], + shared_expert_ids_to_devices: dict[int, list[int]], + num_devices: int, + ) -> "torch.Tensor": + """Arrange a `{shared_expert_id: tensor}` dict into the per-device-stacked layout the + C++ `add_shared_expert_weights` expects. + + For each device (in order), concatenate its assigned shared experts (in sorted global-id + order) along the experts dim (dim 1), then concatenate across devices. Sharding the + result on dim 1 then lands each device's shared experts on it, in slot order — matching + how the device-side splice appends shared experts after routed ones per device. + """ + device_to_shared: list[list[int]] = [[] for _ in range(num_devices)] + for sid in sorted(shared_expert_ids_to_devices): + for d in shared_expert_ids_to_devices[sid]: + device_to_shared[d].append(sid) + per_device = [torch.cat([shared_dict[sid] for sid in device_to_shared[d]], dim=1) for d in range(num_devices)] + return torch.cat(per_device, dim=1) + + @staticmethod + def _init_total_expert_weights_impl( + tt_w0: ttnn.Tensor, + tt_w1: ttnn.Tensor, + tt_w2: ttnn.Tensor, + mesh_device: ttnn.MeshDevice, + num_layers: int, + experts_per_device: int, + hidden_size: int, + intermediate_size: int, + has_bias: bool, + tt_b0: ttnn.Tensor | None = None, + tt_b1: ttnn.Tensor | None = None, + tt_b2: ttnn.Tensor | None = None, + ) -> tuple[ttnn.Tensor, ttnn.Tensor]: + """Pack and quantize the combined routed+shared weight device tensors to `bfloat4_b`. + + Inputs are multi-device tensors sharded on the experts dim (dim 1), bf16 / ROW_MAJOR. + Runs the C++ on-device packers (`ttnn.experimental.prepare_*`, with the `_with_bias` + variants when bias rows are folded in), fetches the DRAM-sharded mem configs from + `ttnn.experimental.get_weight_mem_configs`, then quantizes to `bfloat4_b` onto those + configs via `ttnn.experimental.quantize_weights_via_host` (host round-trip). + + `experts_per_device` is the *combined* (routed + shared) per-device expert count — the + experts dim of the inputs, divided across the mesh. Returns the two device tensors + `(tt_w0_w1, tt_w2)` ready to feed `moe_compute`. + """ + logger.info( + f"Preparing expert weights on device: per_device={experts_per_device} " + f"hidden={hidden_size} intermediate={intermediate_size} has_bias={has_bias}" + ) + + if has_bias: + tt_w0_w1_prepped = ttnn.experimental.prepare_w0_w1_tensor_with_bias( + tt_w0, tt_w1, tt_b0, tt_b1, L=num_layers, E=experts_per_device, K=hidden_size, N=intermediate_size + ) + tt_w2_prepped = ttnn.experimental.prepare_w2_tensor_with_bias( + tt_w2, tt_b2, L=num_layers, E=experts_per_device, N=intermediate_size, K=hidden_size + ) + else: + tt_w0_w1_prepped = ttnn.experimental.prepare_w0_w1_tensor_for_moe_compute( + tt_w0, tt_w1, L=num_layers, E=experts_per_device, K=hidden_size, N=intermediate_size + ) + tt_w2_prepped = ttnn.experimental.prepare_w2_tensor_for_moe_compute( + tt_w2, L=num_layers, E=experts_per_device, N=intermediate_size, K=hidden_size + ) + + # Raw uploaded inputs are consumed by the packers; free them before the host round-trip. + for t in (tt_w0, tt_w1, tt_w2): + ttnn.deallocate(t) + if has_bias: + for t in (tt_b0, tt_b1, tt_b2): + ttnn.deallocate(t) + + # has_bias grows the padded K / N by a tile to accommodate the bias row. + weight_mem_configs = ttnn.experimental.get_weight_mem_configs( + mesh_device, + num_layers=num_layers, + experts_per_device=experts_per_device, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + has_bias=has_bias, + ) + + tt_w0_w1 = ttnn.experimental.quantize_weights_via_host( + tt_w0_w1_prepped, dtype=ttnn.bfloat4_b, memory_config=weight_mem_configs.w0_w1 + ) + ttnn.deallocate(tt_w0_w1_prepped) + tt_w2 = ttnn.experimental.quantize_weights_via_host( + tt_w2_prepped, dtype=ttnn.bfloat4_b, memory_config=weight_mem_configs.w2 + ) + ttnn.deallocate(tt_w2_prepped) + logger.info("Prepared and quantized w0/w1 and w2 to bfloat4_b on mesh") + + return tt_w0_w1, tt_w2 + + +class _TTMoEDecodeBuffers: + """Persistent buffers and semaphores for the MoE decode pipeline. + + Allocates the dispatch output triple (sparse buffer, expert indices, expert + scores), the dispatch/combine cross-device semaphores, and the combine + output buffer once in __init__, for reuse across forward() calls. + + Shapes and memory configs mirror those used in test_optimized_moe_decode_block.py. + """ + + SPARSE_BUFFER_DTYPE = ttnn.bfloat16 + INDICES_DTYPE = ttnn.uint16 + SCORES_DTYPE = ttnn.bfloat16 + COMBINE_OUTPUT_DTYPE = ttnn.bfloat16 + + def __init__( + self, + mesh_device: ttnn.MeshDevice, + *, + mesh_shape: tuple[int, int], + cluster_axis: int, + batch_per_device: int, + hidden_size: int, + effective_experts_k: int, + shard_dim: int, + compute_tilize_drain_core: ttnn.CoreCoord, + ): + """Allocate the persistent buffers and semaphores reused across `forward()` calls. + + Allocates four ttnn tensors and two global semaphores: + - `dispatch_global_semaphore` / `combine_global_semaphore`: cross-device sync + points. Single-use per forward (no double buffering needed — combine syncs + after fully reading the dispatch output, dispatch syncs at end of pipeline). + - `tt_dispatch_output_tensors` triple: sparse buffer (DRAM, hidden-wide token slots), + indices and scores (both L1 height-sharded on the drain core, narrow). + - `tt_combine_output`: DRAM `[effective_experts_k, batch_per_device, hidden_size]` + intermediate after `moe_compute`, before the post-combine tilize. + + `effective_experts_k = select_experts_k + num_shared_experts` — the K dimension of + the per-token expert-output stack. Sized from config; passed in to keep this class + agnostic of where the value came from. + """ + # --- derived sizes (seq=1 for decode) --- + num_dispatch_devices = mesh_shape[cluster_axis] + total_tokens = batch_per_device * num_dispatch_devices + tokens_per_device = batch_per_device + + shard_dims = (shard_dim, None) if cluster_axis == 0 else (None, shard_dim) + + # --- global semaphores (one each, no double buffering required — + # combine syncs after reading dispatch output, dispatch syncs at end) --- + compute_grid_size = mesh_device.compute_with_storage_grid_size() + worker_cores = ttnn.CoreRangeSet( + {ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(compute_grid_size.x - 1, compute_grid_size.y - 1))} + ) + self.dispatch_global_semaphore = ttnn.create_global_semaphore(mesh_device, worker_cores, 0) + self.combine_global_semaphore = ttnn.create_global_semaphore(mesh_device, worker_cores, 0) + + # --- dispatch output buffers --- + # Sparse buffer: DRAM, row-major, sharded along cluster axis + sparse_buffer = ttnn.from_torch( + torch.zeros( + [num_dispatch_devices, total_tokens, hidden_size], + dtype=_tt_to_torch_dtype(self.SPARSE_BUFFER_DTYPE), + ), + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=self.SPARSE_BUFFER_DTYPE, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=shard_dims, mesh_shape=mesh_shape), + ) + + # Indices / scores share an L1 height-sharded mem config on the drain core + shard_spec = ttnn.ShardSpec( + ttnn.CoreRangeSet({ttnn.CoreRange(compute_tilize_drain_core, compute_tilize_drain_core)}), + [total_tokens, effective_experts_k], + ttnn.ShardOrientation.ROW_MAJOR, + ) + l1_height_sharded = ttnn.MemoryConfig( + ttnn.TensorMemoryLayout.HEIGHT_SHARDED, + ttnn.BufferType.L1, + shard_spec, + ) + + indices = ttnn.from_torch( + torch.zeros( + [num_dispatch_devices, total_tokens, effective_experts_k], + dtype=_tt_to_torch_dtype(self.INDICES_DTYPE), + ), + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=self.INDICES_DTYPE, + memory_config=l1_height_sharded, + mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=shard_dims, mesh_shape=mesh_shape), + ) + + scores = ttnn.from_torch( + torch.zeros( + [num_dispatch_devices, total_tokens, effective_experts_k], + dtype=_tt_to_torch_dtype(self.SCORES_DTYPE), + ), + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=self.SCORES_DTYPE, + memory_config=l1_height_sharded, + mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=shard_dims, mesh_shape=mesh_shape), + ) + + self.tt_dispatch_output_tensors = (sparse_buffer, indices, scores) + self.tt_combine_output = ttnn.from_torch( + torch.zeros( + [effective_experts_k, tokens_per_device, hidden_size], + dtype=_tt_to_torch_dtype(self.COMBINE_OUTPUT_DTYPE), + ), + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + dtype=self.COMBINE_OUTPUT_DTYPE, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + + +class TTMoEDecode: + """MoE decode block: token dispatch → expert compute → score-weighted combine → reduce-scatter. + + Constructed once per layer (or per shared layer slot); `forward()` drives one + decode step. All shape / topology / memory-config decisions live in `config` + (`TTMoEDecodeConfig`) — this class only orchestrates the ttnn op sequence. + + Two RS branches are auto-selected from `mesh_shape[1 - cluster_axis]`: + - `== DEEPSEEK_RS_DP_DIM (8)`: fused `deepseek_moe_reduce_scatter` consuming the + pre-split list of N outputs from `fast_reduce_nc_fused`. + - `== SKIP_RS_DP_DIM (1)`: no replication, RS is a no-op. + - else: generic `ttnn.reduce_scatter` over the single fast-reduce output. + """ + + DEEPSEEK_RS_DP_DIM: int = 8 + SKIP_RS_DP_DIM: int = 1 + + def __init__( + self, + mesh_device: ttnn.MeshDevice, + config: TTMoEDecodeConfig, + torch_w0: torch.Tensor, + torch_w1: torch.Tensor, + torch_w2: torch.Tensor, + shared_id_to_torch_w0: dict[int, torch.Tensor] | None = None, + shared_id_to_torch_w1: dict[int, torch.Tensor] | None = None, + shared_id_to_torch_w2: dict[int, torch.Tensor] | None = None, + torch_b0: torch.Tensor | None = None, + torch_b1: torch.Tensor | None = None, + torch_b2: torch.Tensor | None = None, + ): + """Upload weights / biases / shared experts to the mesh and allocate scratch buffers. + + Routed weight shapes: `w0/w1 = [L, num_routed_experts, hidden_size, intermediate_size]`, + `w2 = [L, num_routed_experts, intermediate_size, hidden_size]`. Shared weights are + dicts keyed by global expert id (in `[num_routed, num_routed + num_shared)`), each + value `[L, 1, ...]` matching the routed layout. + + Bias shapes (only when `config.has_bias`): `b0/b1 = [L, num_routed_experts, intermediate_size]`, + `b2 = [L, num_routed_experts, hidden_size]`. Bias + shared experts together raises + `NotImplementedError`. + """ + self.config = config + self.expert_state = _TTMoEDecodeExpertState( + mesh_device, + torch_w0, + torch_w1, + torch_w2, + **config.experts.model_dump(), + shared_id_to_torch_w0=shared_id_to_torch_w0, + shared_id_to_torch_w1=shared_id_to_torch_w1, + shared_id_to_torch_w2=shared_id_to_torch_w2, + torch_b0=torch_b0, + torch_b1=torch_b1, + torch_b2=torch_b2, + ) + buffers_dict = config.buffers.model_dump() + if buffers_dict.get("compute_tilize_drain_core") is None: + matmul_ring_size = effective_matmul_ring_size(mesh_device) + buffers_dict["compute_tilize_drain_core"] = ttnn.experimental.get_moe_tilize_drain_core( + mesh_device, + config.compute.output_height_shard_dim, + auto_output_width_shard_dim(config.hidden_size, matmul_ring_size=matmul_ring_size), + config.hidden_size, + mux_core_range_set=config.compute.mux_core_range_set, + ) + else: + raise ValueError( + "compute_tilize_drain_core is not user-configurable; omit it to resolve dynamically at runtime" + ) + self.buffers = _TTMoEDecodeBuffers(mesh_device, **buffers_dict) + + @property + def _num_fast_reduce_outputs(self) -> int: + """Number of outputs fast_reduce_nc_fused will produce — N for the deepseek + RS-list path, 1 otherwise (downstream ttnn.reduce_scatter takes a single tensor).""" + return self.config.num_fast_reduce_outputs + + @property + def _pre_split_chunk(self) -> int: + """Logical per-fast-reduce-output width (hidden_size / num_fast_reduce_outputs). + For the single-output path this is just hidden_size.""" + return self.config.hidden_size // self._num_fast_reduce_outputs + + @property + def _padded_pre_split_chunk(self) -> int: + """Aligned per-fast-reduce-output width = config.reduce.split_size.""" + return self.config.reduce.split_size + + @property + def _post_rs_logical_chunk(self) -> int: + """Logical per-device width after RS — what the model expects downstream.""" + return self.config.hidden_size // self.config.mesh_shape[1 - self.config.cluster_axis] + + @property + def _needs_fast_reduce_padding(self) -> bool: + return self._pre_split_chunk != self._padded_pre_split_chunk + + def _pad_for_fast_reduce(self, tt_x: ttnn.Tensor) -> ttnn.Tensor: + """Interleave-pad the H dim so each post-RS per-device chunk is tile-aligned. + + Downstream RS splits the H dim evenly into `num_replicated` chunks, so padding + must be inserted at each device-chunk boundary — not just appended at the end — + or device d>0 ends up with data shifted by `d * (padded - logical)` positions. + + Layout produced: `[chunk_0, pad_0, ..., chunk_{R-1}, pad_{R-1}]`, R=num_replicated. + Each `chunk_d` is `hidden / R` wide; each `pad_d` brings it up to TILE_SIZE alignment. + """ + if not self._needs_fast_reduce_padding: + return tt_x + num_replicated = self.config.mesh_shape[1 - self.config.cluster_axis] + chunk = self._post_rs_logical_chunk + padded_chunk = self._padded_pre_split_chunk * self._num_fast_reduce_outputs // num_replicated + shape = list(tt_x.shape) + reshaped = ttnn.reshape(tt_x, shape[:-1] + [num_replicated, chunk]) + padded = ttnn.pad( + reshaped, + padding=[(0, 0)] * len(shape) + [(0, padded_chunk - chunk)], + value=0.0, + ) + return ttnn.reshape(padded, shape[:-1] + [num_replicated * padded_chunk]) + + def _unpad_after_reduce_scatter(self, tt_final: ttnn.Tensor) -> ttnn.Tensor: + """Slice the trailing padding off each device's post-RS tensor. + + After RS each device holds at least `_post_rs_logical_chunk` valid columns followed + by zero padding (inserted before tilize). No-op when the original hidden splits + evenly without padding. + """ + if not self._needs_fast_reduce_padding: + return tt_final + chunk = self._post_rs_logical_chunk + shape = list(tt_final.shape) + start = [0] * len(shape) + end = shape[:-1] + [chunk] + return ttnn.slice(tt_final, start, end) + + def _format_dispatch_inputs( + self, + tt_x: ttnn.Tensor, + tt_indices: ttnn.Tensor, + tt_scores: ttnn.Tensor, + ): + """Coerce each dispatch input into the memory config the dispatch op needs. + + For each of (`tt_x`, `tt_indices`, `tt_scores`) returns a `(tensor, dealloc_flag)` + pair: if a `to_memory_config` was necessary, the returned tensor is a fresh + allocation that `forward()` should deallocate after the dispatch op; if the input + already matched, the original is returned with `dealloc=False` to leave caller + ownership intact. + """ + if tt_x.memory_config() != self.config.dispatch_input_memory_config: + tt_dispatch_input_tensor_bundle = ( + ttnn.to_memory_config(tt_x, memory_config=self.config.dispatch_input_memory_config), + True, + ) + else: + tt_dispatch_input_tensor_bundle = tt_x, False + + if tt_indices.memory_config() != self.config.dispatch_input_memory_config: + tt_dispatch_input_expert_indices_tensor_bundle = ( + ttnn.to_memory_config( + tt_indices, + memory_config=self.config.dispatch_input_memory_config, + ), + True, + ) + else: + tt_dispatch_input_expert_indices_tensor_bundle = tt_indices, False + + if tt_scores.memory_config() != self.config.dispatch_input_expert_scores_memory_config: + tt_dispatch_input_expert_scores_tensor_bundle = ( + ttnn.to_memory_config( + tt_scores, + memory_config=self.config.dispatch_input_expert_scores_memory_config, + ), + True, + ) + else: + tt_dispatch_input_expert_scores_tensor_bundle = tt_scores, False + + return ( + tt_dispatch_input_tensor_bundle, + tt_dispatch_input_expert_indices_tensor_bundle, + tt_dispatch_input_expert_scores_tensor_bundle, + ) + + def forward( + self, tt_x: ttnn.Tensor, tt_scores: ttnn.Tensor, tt_indices: ttnn.Tensor, layer_id: int = 0 + ) -> ttnn.Tensor: + """Run one decode-step MoE forward. + + Inputs (sharded along the dispatch axis): + - `tt_x`: `[1, batch_per_device, 1, hidden_size]` activations per device. + - `tt_indices`: `[batch_per_device, 1, 1, select_experts_k]` chosen routed experts + per token (uint16). + - `tt_scores`: `[batch_per_device, 1, 1, select_experts_k]` per-(token, k) score + for the routed combine; shared experts use `config.reduce.shared_expert_scale` + uniformly, not this tensor. + - `layer_id`: which slice of the layered weight tensors to use (currently + assumed `0` since the rest of the test/module pipeline is `num_layers=1`). + + Output: `[1, 1, batch_per_device, hidden_size / num_replicated]` per device, + i.e. each device holds its post-reduce-scatter chunk of the combined hidden dim. + + Pipeline matches the reference test_optimized_moe_decode_block: + 1. `all_to_all_dispatch_metadata` → per-device sparse buffer of inbound tokens. + 2. `moe_compute` → per-(k, token) expert output stack, optionally with bias. + 3. Tilize (`deepseek_moe_post_combine_tilize` when batch_per_device == TILE_SIZE + and an NdShard config is available; else `tilize_with_val_padding` fallback). + 4. `deepseek_moe_fast_reduce_nc_fused` → score-weighted sum over k, with shared + experts scaled by `shared_expert_scale`. + 5. Reduce-scatter across the replicated axis (3 variants — see class docstring). + 6. Strip the per-device-chunk padding inserted before tilize, if any. + """ + ( + (tt_dispatch_input_tensor, dealloc_input), + (tt_dispatch_input_expert_indices_tensor, dealloc_indices), + (tt_dispatch_input_expert_scores_tensor, dealloc_scores), + ) = self._format_dispatch_inputs(tt_x, tt_indices, tt_scores) + + ( + tt_dispatch_output_sparse_buffer, + tt_dispatch_output_expert_indices, + tt_dispatch_output_expert_scores, + ) = ttnn.experimental.all_to_all_dispatch_metadata( + tt_dispatch_input_tensor, + tt_dispatch_input_expert_indices_tensor, + tt_dispatch_input_expert_scores_tensor, + self.expert_state.tt_expert_mapping, + **self.config.dispatch.model_dump(), + # shared_expert_ids + # cluster_axi + # num_links + # drain_sync_tilizer_core + # worker_mode + # dispatch_algorithm + output_tensors=self.buffers.tt_dispatch_output_tensors, + cross_device_semaphore=self.buffers.dispatch_global_semaphore, + ) + + if dealloc_input: + ttnn.deallocate(tt_dispatch_input_tensor) + + if dealloc_scores: + ttnn.deallocate(tt_dispatch_input_expert_scores_tensor) + + ( + _, + _, + _, + tt_l1_compute_output, + _, + tt_combine_output, + ) = ttnn.experimental.moe_compute( + tt_dispatch_output_sparse_buffer, + tt_dispatch_output_expert_indices, + tt_dispatch_output_expert_scores, + self.expert_state.tt_expert_mapping, + self.expert_state.tt_w0_w1, + self.expert_state.tt_w2, + layer_id=layer_id, + # output_height_shard_dim + # cluster_axis + # mux_core_range_set + # has_bias + # activation_type + **self.config.compute.model_dump(), + optional_output_tensor=self.buffers.tt_combine_output, + optional_cross_device_semaphore=self.buffers.combine_global_semaphore, + ) + ttnn.deallocate(tt_l1_compute_output) + + # unsqueeze + # [select_experts_k, tokens_per_device, hidden_size] -> [select_experts_k, 1, tokens_per_device, hidden_size] + # Note: this does not reallocate, aliases tt_combine_output so don't dealloc + tt_unsqueezed_output = ttnn.unsqueeze(tt_combine_output, dim=1) + + # When hidden / num_replicated isn't tile-aligned, interleave-pad each per-device + # chunk up to split_size so fast_reduce_nc_fused's 128-divisibility check passes. + # No-op when already aligned. + tt_unsqueezed_output = self._pad_for_fast_reduce(tt_unsqueezed_output) + + if self.config.use_post_combine_tilize: + tt_tilized_compute_output = ttnn.experimental.deepseek_moe_post_combine_tilize( + tt_unsqueezed_output, + # output_memory_config, + **self.config.post_combine_tilize.model_dump(), + ) + + else: + output_tensor_shape = list(tt_unsqueezed_output.shape) + output_tensor_shape[2] = ((output_tensor_shape[2] + ttnn.TILE_SIZE - 1) // ttnn.TILE_SIZE) * ttnn.TILE_SIZE + tt_tilized_compute_output = ttnn.tilize_with_val_padding( + tt_unsqueezed_output, + output_tensor_shape=output_tensor_shape, + pad_value=0.0, + # memory_config + **self.config.tilize_with_val_padding.model_dump(), + ) + + # scale with scores and accumulate + tt_fast_reduce_output_tensors = ttnn.experimental.deepseek_moe_fast_reduce_nc_fused( + tt_tilized_compute_output, + tt_dispatch_input_expert_indices_tensor, + self.expert_state.tt_expert_mapping, + # reduce_dim + # cluster_axis + # split_size + # output_memory_config + # num_shared_experts + # shared_expert_scale + **self.config.reduce.model_dump(), + scores_tensor=tt_scores, + ) + ttnn.deallocate(tt_tilized_compute_output) + + if dealloc_indices: + ttnn.deallocate(tt_dispatch_input_expert_indices_tensor) + + # [select_experts_k, tokens_per_device, hidden_size // num_replicated_devices] final per device shape + if ( + self.config.mesh_shape[1 - self.config.cluster_axis] == self.DEEPSEEK_RS_DP_DIM + and self.config.topology == ttnn.Topology.Ring + ): + tt_final_output = ttnn.experimental.deepseek_moe_reduce_scatter( + tt_fast_reduce_output_tensors, + # output_memory_config + # dim + # num_links + # topology + # cluster_axis + **self.config.deepseek_moe_reduce_scatter.model_dump(), + ) + for t in tt_fast_reduce_output_tensors: + ttnn.deallocate(t) + + # note: in this path the output is L1 sharded as if set up for deepseek_moe_reduce_scatter. Likely better to + # switch this to something more generic. + elif self.config.mesh_shape[1 - self.config.cluster_axis] == self.SKIP_RS_DP_DIM: + tt_final_output = tt_fast_reduce_output_tensors[0] + else: + tt_final_output = ttnn.reduce_scatter( + tt_fast_reduce_output_tensors[0], **self.config.reduce_scatter.model_dump() + ) + for t in tt_fast_reduce_output_tensors: + ttnn.deallocate(t) + + # Strip the per-chunk padding we inserted before tilize. No-op when no padding was applied. + return self._unpad_after_reduce_scatter(tt_final_output) diff --git a/code/models/common/modules/moe/tt_moe_decode_config.py b/code/models/common/modules/moe/tt_moe_decode_config.py new file mode 100644 index 0000000000000000000000000000000000000000..47649d830f84f57a99e40bb2afcda072bfb07b71 --- /dev/null +++ b/code/models/common/modules/moe/tt_moe_decode_config.py @@ -0,0 +1,774 @@ +# 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) diff --git a/code/models/common/modules/moe/tt_moe_decode_config_schemas.py b/code/models/common/modules/moe/tt_moe_decode_config_schemas.py new file mode 100644 index 0000000000000000000000000000000000000000..7823d11abc66de08e87f0a7de8947b6a2cf9d09b --- /dev/null +++ b/code/models/common/modules/moe/tt_moe_decode_config_schemas.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +"""YAML/JSON-friendly field types for ttnn objects used in `TTMoEDecodeConfig`. + +Each alias is `Annotated[ttnn.X, BeforeValidator, PlainSerializer(when_used="json")]` +so: + - `model_dump()` (default `mode="python"`) returns the raw ttnn objects — keeps + `**cfg.x.model_dump()` splatting into ttnn ops working unchanged. + - `model_dump(mode="json")` (and JSON/YAML serialization) returns dicts/strings + produced by each type's `to_json()`. Combined with `exclude_defaults=True` + + `exclude_none=True`, only user-set non-default values appear. + - Loading the same dict back through `model_validate` rehydrates the ttnn + objects via `from_json`. + +Most ttnn types here already expose `to_json` / `from_json` on the C++ side; we +just hook those into pydantic. Enums fall back to name-based round-tripping. +""" + +from __future__ import annotations + +import enum +import json +from typing import Annotated, Any + +from pydantic import BeforeValidator, PlainSerializer +from ttnn.operations.ccl import MoEActivationFunction + +import ttnn + + +def _ttnn_json_annotated(cls): + """Annotated alias for any ttnn type that exposes `to_json` / `from_json`.""" + + def _validate(v: Any) -> Any: + if isinstance(v, cls): + return v + if isinstance(v, str): + return cls.from_json(v) + if isinstance(v, (dict, list)): + return cls.from_json(json.dumps(v)) + raise ValueError(f"cannot convert {v!r} to {cls.__name__}") + + def _serialize(v: Any) -> Any: + # Parse to dict/list so yaml/json downstream handle it natively + return json.loads(v.to_json()) + + return Annotated[cls, BeforeValidator(_validate), PlainSerializer(_serialize, when_used="json")] + + +def _enum_annotated(enum_cls: type[enum.Enum]): + """Annotated alias for an enum: serialize as name, accept name/int/instance.""" + + def _validate(v: Any) -> Any: + if isinstance(v, enum_cls): + return v + if isinstance(v, str): + return enum_cls[v] + if isinstance(v, int): + return enum_cls(v) + raise ValueError(f"cannot convert {v!r} to {enum_cls.__name__}") + + def _serialize(v: Any) -> str: + return v.name + + return Annotated[enum_cls, BeforeValidator(_validate), PlainSerializer(_serialize, when_used="json")] + + +# Most ttnn types round-trip through to_json / from_json cleanly: +CoreCoord = _ttnn_json_annotated(ttnn.CoreCoord) +CoreRange = _ttnn_json_annotated(ttnn.CoreRange) +CoreRangeSet = _ttnn_json_annotated(ttnn.CoreRangeSet) + + +# MemoryConfig needs a small workaround: when it carries an `NdShardSpec`, ttnn's +# `to_json` can't serialize the `Shape` (it emits an error string for the +# `span` field). Pull the shape out via the live object and +# overwrite that field with a plain list before yielding the dict. `from_json` +# already accepts a list there. +def _validate_memory_config(v: Any) -> Any: + if isinstance(v, ttnn.MemoryConfig): + return v + if isinstance(v, str): + return ttnn.MemoryConfig.from_json(v) + if isinstance(v, (dict, list)): + return ttnn.MemoryConfig.from_json(json.dumps(v)) + raise ValueError(f"cannot convert {v!r} to ttnn.MemoryConfig") + + +def _serialize_memory_config(v: ttnn.MemoryConfig) -> dict: + d = json.loads(v.to_json()) + nd = d.get("nd_shard_spec") + if isinstance(nd, dict) and isinstance(nd.get("shard_shape"), str): + # the bad string is `ttsl::json::to_json_t: Unsupported type span` + nd["shard_shape"] = list(v.nd_shard_spec.shard_shape) + return d + + +MemoryConfig = Annotated[ + ttnn.MemoryConfig, + BeforeValidator(_validate_memory_config), + PlainSerializer(_serialize_memory_config, when_used="json"), +] + +# enums (name-based) +Topology = _enum_annotated(ttnn.Topology) +WorkerMode = _enum_annotated(ttnn.WorkerMode) +DispatchAlgorithm = _enum_annotated(ttnn.DispatchAlgorithm) +ActivationFunction = _enum_annotated(MoEActivationFunction) diff --git a/code/models/common/modules/moe/tt_moe_gate.py b/code/models/common/modules/moe/tt_moe_gate.py new file mode 100644 index 0000000000000000000000000000000000000000..0a9651048ec4d28f0981a1ced86faabccf75a7a4 --- /dev/null +++ b/code/models/common/modules/moe/tt_moe_gate.py @@ -0,0 +1,756 @@ +# 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 diff --git a/code/models/common/modules/moe/tt_moe_gate_config.py b/code/models/common/modules/moe/tt_moe_gate_config.py new file mode 100644 index 0000000000000000000000000000000000000000..775647a2c119273e863d364508310142fe4497f0 --- /dev/null +++ b/code/models/common/modules/moe/tt_moe_gate_config.py @@ -0,0 +1,95 @@ +# 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}) diff --git a/code/models/common/modules/rmsnorm/rmsnorm_1d.py b/code/models/common/modules/rmsnorm/rmsnorm_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..32e482f531dc50ad4e3210c70ac27ab0be7027a2 --- /dev/null +++ b/code/models/common/modules/rmsnorm/rmsnorm_1d.py @@ -0,0 +1,790 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +RMSNorm1D: RMSNorm for 1D mesh topologies (single-chip, n150, n300, T3K row/column). + +Execution paths: + - Path 1: _decode_local - Local decode (sharded or interleaved) - always for 1D decode + - Path 2: _prefill_local - Local prefill (interleaved) - single-device or prefill_distributed=False + - Path 3: _prefill_1d_distributed - 1D distributed prefill - multi-device CCL + +Key design: + - decode_forward always uses local sharded path (no 1D distributed decode) + - prefill_forward switches between local and distributed based on prefill_distributed config + - Config provides program_config/memory_config for sharded, None for interleaved + - from_model_args() sets prefill_distributed=True when dim >= 4096 for backward compatibility +""" + +import math +from dataclasses import dataclass, replace +from typing import Optional + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight +from models.common.modules.tt_ccl import default_topology, get_tt_ccl +from models.common.tensor_utils import TILE_SIZE + +# ============================================================================= +# Constants +# ============================================================================= + +SHARD_HEIGHT = TILE_SIZE # Current ttnn.rms_norm requires shard height = single tile + + +# ============================================================================= +# RMSNorm1DConfig +# ============================================================================= + + +@dataclass +class RMSNorm1DConfig: + """ + Configuration for RMSNorm1D - only fields relevant to 1D mesh topologies. + + Paths: + - Path 1: Local decode (sharded) - set decode_program_config & decode_memory_config + - Path 2: Local prefill (interleaved) - no extra config needed + - Path 3: 1D distributed prefill - requires tt_ccl for Ring all-gather + + Simple usage: + config = RMSNorm1DConfig(weight) + + Override any field: + config = RMSNorm1DConfig(weight, eps=1e-6, prefill_distributed=True) + """ + + # Required: weight + weight: LazyWeight + + # Device and collectives + mesh_device: ttnn.MeshDevice | None = None + tt_ccl: "TT_CCL | None" = None # type: ignore + + # Normalization settings + eps: float = 1e-5 + add_unit_offset: bool = False + + # Distributed control (None = auto-detect: True for multi-device, False for single-device) + prefill_distributed: bool | None = None + + # Batch size for decode + max_batch_size: int = 32 + + # Local decode configs (Path 1) + decode_in_sharded: bool = True # If True, use sharded decode; if False, use interleaved + decode_out_sharded: bool = True # If True, output stays sharded (only valid if in_sharded=True) + decode_program_config: ttnn.LayerNormShardedMultiCoreProgramConfig | None = None + decode_memory_config: ttnn.MemoryConfig | None = None + + # Input memory configs (optional, for prepare_input_tensor) + decode_input_memcfg: ttnn.MemoryConfig | None = None + prefill_input_memcfg: ttnn.MemoryConfig | None = None + + # Compute kernel config (None = architecture default resolved from mesh_device) + compute_kernel_config: ttnn.DeviceComputeKernelConfig | None = None + + # Internal: distributed weight LazyWeight (sharded across devices for Path 3) + # Note: Uses LazyWeight with explicit mesh_mapper_config to shard weights correctly. + # The get_device_weight() call respects the mesh_mapper_config for proper distribution. + _weight_distributed: LazyWeight | None = None + + def is_resolved(self) -> bool: + """Check if all required fields are resolved.""" + required = ["weight", "mesh_device", "compute_kernel_config"] + + # Multi-device distributed prefill needs CCL and distributed weight + num_devices = self.mesh_device.get_num_devices() if self.mesh_device else 1 + if num_devices > 1 and self.prefill_distributed: + required.extend(["tt_ccl", "_weight_distributed"]) + + # Decode always needs sharded config + required.extend(["decode_program_config", "decode_memory_config"]) + + return all(getattr(self, f) is not None for f in required) + + +# ============================================================================= +# RMSNorm1D +# ============================================================================= + + +class RMSNorm1D(LightweightModule): + """ + RMSNorm for 1D mesh topologies. + + Execution paths (bound at construction): + - decode_forward -> _decode_local (Path 1) - always local sharded + - prefill_forward -> _prefill_local (Path 2) or _prefill_1d_distributed (Path 3) + + Simple API: + norm = RMSNorm1D(weight) + + Power API: + config = RMSNorm1DConfig(weight, max_batch_size=64) + norm = RMSNorm1D.from_config(config) + """ + + def __init__(self, weight: LazyWeight): + """ + Simple API - derives all config from weight. + + Args: + weight: RMSNorm weight tensor of shape (dim,) or (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) + + Note: Use from_config() to customize eps, add_unit_offset, or other settings. + """ + super().__init__() + self._initialize_from_config(resolve_rmsnorm_1d_arch_config(RMSNorm1DConfig(weight=weight))) + + @classmethod + def from_config(cls, config: RMSNorm1DConfig): + """Power API - any level of customization via config.""" + instance = object.__new__(cls) + super(RMSNorm1D, instance).__init__() + instance._initialize_from_config(resolve_rmsnorm_1d_arch_config(config)) + return instance + + def _initialize_from_config(self, config: RMSNorm1DConfig) -> None: + self.config = config + self._device_weights_loaded = False + self._bind_forward_methods() + + def _bind_forward_methods(self): + """Bind decode_forward and prefill_forward based on config.""" + cfg = self.config + + # Decode: sharded or interleaved based on config + if cfg.decode_in_sharded: + self.decode_forward = self._decode_local_sharded + else: + self.decode_forward = self._decode_local_interleaved + + # Prefill: local or 1D distributed based on dim + if cfg.prefill_distributed: + self.prefill_forward = self._prefill_1d_distributed + else: + self.prefill_forward = self._prefill_local + + def load_device_weights(self): + """Load weights to device lazily.""" + if self._device_weights_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device weights!" + + cfg = self.config + + # Always load replicated weight (decode always needs it) + self.weight = cfg.weight.get_device_weight() + + # Load sharded weight if distributed prefill is enabled + if cfg.prefill_distributed and cfg._weight_distributed is not None: + self.weight_distributed = cfg._weight_distributed.get_device_weight() + + self._device_weights_loaded = True + + # ========================================================================= + # Path 1a: Local decode - sharded + # ========================================================================= + + def _decode_local_sharded(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Sharded decode - standard LLM decode path. + + Execution: to_sharded -> rms_norm(sharded) with program_config and memory_config. + Output: sharded if decode_out_sharded=True, interleaved otherwise + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="decode") + cfg = self.config + + assert cfg.decode_in_sharded, "decode_in_sharded must be True for sharded decode" + + x = ttnn.to_memory_config(x, memory_config=cfg.decode_memory_config) + + x = ttnn.rms_norm( + x, + epsilon=cfg.eps, + weight=self.weight, + program_config=cfg.decode_program_config, + memory_config=cfg.decode_memory_config, + compute_kernel_config=cfg.compute_kernel_config, + ) + + if not cfg.decode_out_sharded: + x = ttnn.sharded_to_interleaved(x) + + return x + + # ========================================================================= + # Path 1b: Local decode - interleaved + # ========================================================================= + + def _decode_local_interleaved(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Interleaved decode - for vision encoder decode. + + Execution: rms_norm with program_config=None, memory_config=None + Output: always interleaved (decode_out_sharded must be False) + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="decode") + cfg = self.config + + assert not cfg.decode_in_sharded, "decode_in_sharded must be False for interleaved decode" + assert not cfg.decode_out_sharded, "decode_out_sharded must be False when decode_in_sharded=False" + + return ttnn.rms_norm( + x, + epsilon=cfg.eps, + weight=self.weight, + program_config=None, + memory_config=None, + compute_kernel_config=cfg.compute_kernel_config, + ) + + # ========================================================================= + # Path 2: Local prefill (interleaved) + # ========================================================================= + + def _prefill_local(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + Local prefill - interleaved RMSNorm when prefill_distributed=False. + + Execution: rms_norm(interleaved) without program_config + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="prefill") + cfg = self.config + + return ttnn.rms_norm( + x, + epsilon=cfg.eps, + weight=self.weight, + program_config=None, + memory_config=None, + compute_kernel_config=cfg.compute_kernel_config, + ) + + # ========================================================================= + # Path 3: 1D distributed prefill + # ========================================================================= + + def _prefill_1d_distributed(self, x: ttnn.Tensor | LazyWeight) -> ttnn.Tensor: + """ + 1D distributed prefill - when prefill_distributed=True. + + Execution: + rms_norm_pre_all_gather -> all_gather(stats) -> rms_norm_post_all_gather + """ + self.load_device_weights() + x = _load_input_device_tensor(x, self.config, mode="prefill") + cfg = self.config + + # Run distributed rmsnorm part 1 + tt_stats = ttnn.rms_norm_pre_all_gather(x, compute_kernel_config=cfg.compute_kernel_config, dtype=ttnn.bfloat16) + + # AllGather stats + tt_stats = ttnn.experimental.all_gather_async( + tt_stats, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(), + num_links=1, + topology=default_topology(cfg.mesh_device), + memory_config=ttnn.DRAM_MEMORY_CONFIG, + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + + # Run distributed rmsnorm part 2 (uses sharded weight) + tt_out = ttnn.rms_norm_post_all_gather( + x, + tt_stats, + epsilon=cfg.eps, + weight=self.weight_distributed, + compute_kernel_config=cfg.compute_kernel_config, + ) + tt_stats.deallocate(True) + + return tt_out + + # ========================================================================= + # Forward dispatcher + # ========================================================================= + + def forward(self, x: "ttnn.Tensor | LazyWeight", mode: str) -> ttnn.Tensor: + """ + Dispatch to the appropriate forward method based on mode. + + Args: + x: Input tensor (ttnn.Tensor or LazyWeight wrapping torch tensor) + mode: "decode" or "prefill" + + Note: Sharding behavior for decode is controlled via config: + - decode_in_sharded: True for sharded (LLM), False for interleaved (vision) + - decode_out_sharded: True to keep output sharded, False to convert to interleaved + """ + if mode == "decode": + return self.decode_forward(x) + else: + return self.prefill_forward(x) + + # [INFO] this is the entry point for TTTv1 model_config.py and will retire with TTTv1 + @classmethod + def from_model_args( + cls, + mesh_device, + tt_ccl, + args, + state_dict, + weight_cache_path, + layer_num: int, + weight_key: str, + state_dict_prefix: Optional[str] = None, + sharded_program_config=None, + sharded_output_config=None, + ): + """Factory method for backward compatibility with ModelArgs.""" + + if args.is_galaxy: + raise ValueError("RMSNorm1D cannot be used for Galaxy devices. Use RMSNorm2D instead.") + + # Build weight name from state_dict_prefix and weight_key + if state_dict_prefix: + weight_name = f"{state_dict_prefix}{weight_key}.weight" + else: + if layer_num is None: + weight_name = f"{weight_key}.weight" + else: + weight_name = f"layers.{layer_num}.{weight_key}.weight" + + # Transform weight to expected shape: (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) + dim = args.dim + torch_weight = ( + state_dict[weight_name].unsqueeze(0).view(1, 1, dim).reshape([1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT]) + ) + + # Add offset before creating LazyWeight + if args.rms_norm_add_unit_offset: + torch_weight = torch_weight + 1.0 + + # Create LazyWeight + lazy_weight = LazyWeight( + source=torch_weight, + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(weight_cache_path, weight_name) if weight_cache_path else None, + mesh_mapper_config=( + ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate()], + mesh_shape_override=ttnn.MeshShape([mesh_device.get_num_devices()]), + ) + if mesh_device.get_num_devices() > 1 + else None + ), + ) + + # Preserve the legacy bridge's established numerical policy. Native TTTv2 + # model profiles supply their own explicit architecture wrapper instead. + arch = mesh_device.arch() + compute_kernel_config = _legacy_rmsnorm_compute_kernel_config(arch, getattr(args, "base_model_name", None)) + + # Determine prefill_distributed based on dim and num_devices. + # Keep dim=4096 on the distributed path so multichip prefill norms match + # fractured activations produced by TTTv2 integration. + num_devices = mesh_device.get_num_devices() + prefill_distributed = num_devices > 1 and dim >= 4096 + + # Build config + config = RMSNorm1DConfig( + weight=lazy_weight, + eps=args.norm_eps, + add_unit_offset=False, # Already applied above + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=args.max_batch_size, + prefill_distributed=prefill_distributed, + decode_program_config=sharded_program_config, + decode_memory_config=sharded_output_config, + compute_kernel_config=compute_kernel_config, + ) + + instance = object.__new__(cls) + super(RMSNorm1D, instance).__init__() + instance._initialize_from_config(_resolve_rmsnorm_1d_arch_config(config, mesh_device=mesh_device, arch=arch)) + return instance + + +# ============================================================================= +# Config resolution +# ============================================================================= + + +def _rmsnorm_compute_kernel_config(arch) -> ttnn.DeviceComputeKernelConfig: + """Construct the TTTv1-matched RMSNorm recipe for the concrete device.""" + return ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=True, + ) + + +def _legacy_rmsnorm_compute_kernel_config(arch, base_model_name: str | None) -> ttnn.DeviceComputeKernelConfig: + """Preserve the retiring ModelArgs bridge's established per-model policy.""" + return ttnn.init_device_compute_kernel_config( + arch, + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=base_model_name not in ("Qwen2.5-7B", "Qwen2.5-VL-7B"), + packer_l1_acc=False, + ) + + +def _validate_rmsnorm_compute_kernel_config(compute_kernel_config) -> None: + if compute_kernel_config is None: + raise ValueError("RMSNorm1D architecture config requires compute_kernel_config") + if not isinstance(compute_kernel_config, ttnn.WormholeComputeKernelConfig): + raise TypeError( + "RMSNorm1D compute_kernel_config must be a concrete TTNN compute kernel config, " + f"got {type(compute_kernel_config).__name__}" + ) + + +def _derive_rmsnorm_mesh_device(config: RMSNorm1DConfig) -> ttnn.MeshDevice: + mesh_device = config.mesh_device + if mesh_device is None: + mesh_device = config.weight.device + if mesh_device is None: + mesh_device = ttnn.GetDefaultDevice() + if mesh_device is None: + raise ValueError("RMSNorm1D mesh_device must be available") + if config.weight.device is not None and config.weight.device is not mesh_device: + raise ValueError("RMSNorm1D mesh_device must match weight.device") + return mesh_device + + +def _validate_rmsnorm_program_config( + config: RMSNorm1DConfig, compute_kernel_config: ttnn.DeviceComputeKernelConfig +) -> None: + if not config.decode_in_sharded: + if config.decode_out_sharded: + raise ValueError("decode_out_sharded=True requires decode_in_sharded=True") + return + + program_config = config.decode_program_config + memory_config = config.decode_memory_config + if program_config is None or memory_config is None: + raise ValueError("Sharded RMSNorm1D decode requires decode_program_config and decode_memory_config") + if not isinstance(program_config, ttnn.LayerNormShardedMultiCoreProgramConfig): + raise TypeError("decode_program_config must be LayerNormShardedMultiCoreProgramConfig") + + if program_config.block_h <= 0 or program_config.block_w <= 0 or program_config.subblock_w <= 0: + raise ValueError("RMSNorm1D decode program block and subblock dimensions must be positive") + if program_config.block_w % program_config.subblock_w != 0: + raise ValueError("RMSNorm1D decode block_w must be divisible by subblock_w") + + grid = program_config.compute_with_storage_grid_size + if grid.x <= 0 or grid.y <= 0: + raise ValueError("RMSNorm1D decode program grid dimensions must be positive") + device_grid = config.mesh_device.compute_with_storage_grid_size() + if grid.x > device_grid.x or grid.y > device_grid.y: + raise ValueError("RMSNorm1D decode program grid exceeds the device compute grid") + + dim = config.weight.source.numel() + num_cores = grid.x * grid.y + tile_padded_batch_rows = TILE_SIZE * math.ceil(config.max_batch_size / TILE_SIZE) + expected_block_h = tile_padded_batch_rows // TILE_SIZE + expected_block_w = dim // num_cores // TILE_SIZE + if program_config.block_h != expected_block_h or program_config.block_w != expected_block_w: + raise ValueError( + "RMSNorm1D decode program blocks do not match batch/hidden geometry: " + f"got ({program_config.block_h}, {program_config.block_w}), " + f"expected ({expected_block_h}, {expected_block_w})" + ) + + if compute_kernel_config.dst_full_sync_en: + max_subblock_w = 8 if compute_kernel_config.fp32_dest_acc_en else 16 + else: + max_subblock_w = 4 if compute_kernel_config.fp32_dest_acc_en else 8 + if program_config.subblock_w > max_subblock_w: + raise ValueError( + "RMSNorm1D decode subblock_w exceeds destination-register capacity: " + f"{program_config.subblock_w} > {max_subblock_w}" + ) + + if not memory_config.is_sharded(): + raise ValueError("Sharded RMSNorm1D decode requires a sharded decode_memory_config") + if memory_config.memory_layout != ttnn.TensorMemoryLayout.WIDTH_SHARDED: + raise ValueError("RMSNorm1D decode_memory_config must be width sharded") + expected_shard_shape = [tile_padded_batch_rows, dim // num_cores] + if list(memory_config.shard_spec.shape) != expected_shard_shape: + raise ValueError( + "RMSNorm1D decode memory shard does not match program geometry: " + f"got {list(memory_config.shard_spec.shape)}, expected {expected_shard_shape}" + ) + + +def resolve_rmsnorm_1d_arch_config(config: RMSNorm1DConfig) -> RMSNorm1DConfig: + """Return a new fully resolved common config after one architecture query.""" + mesh_device = _derive_rmsnorm_mesh_device(config) + return _resolve_rmsnorm_1d_arch_config(config, mesh_device=mesh_device, arch=mesh_device.arch()) + + +def _resolve_rmsnorm_1d_arch_config(config: RMSNorm1DConfig, *, mesh_device: ttnn.MeshDevice, arch) -> RMSNorm1DConfig: + if arch not in (ttnn.device.Arch.WORMHOLE_B0, ttnn.device.Arch.BLACKHOLE): + raise ValueError(f"Unsupported RMSNorm1D architecture: {arch}") + + requested_compute = config.compute_kernel_config or _rmsnorm_compute_kernel_config(arch) + try: + compute_kernel_config = ttnn.init_device_compute_kernel_config(arch, requested_compute) + except (TypeError, ValueError) as error: + raise ValueError(f"Invalid RMSNorm1D compute recipe for {arch}: {error}") from error + + _validate_rmsnorm_compute_kernel_config(compute_kernel_config) + resolved_common = _resolve_1d_config(config) + resolved_common = replace(resolved_common, compute_kernel_config=compute_kernel_config) + _validate_rmsnorm_program_config(resolved_common, compute_kernel_config) + return resolved_common + + +def _resolve_1d_config(config: RMSNorm1DConfig) -> RMSNorm1DConfig: + """Resolve config defaults for RMSNorm1D.""" + + to_set = {} + + # --- Phase 1: Foundational fields --- + + # Derive dim from weight (works for any shape: [dim], [1, dim], [1, 1, dim//32, 32], etc.) + dim = config.weight.source.numel() + + # Derive mesh_device from weight + mesh_device = _derive_rmsnorm_mesh_device(config) + if config.mesh_device is None: + to_set["mesh_device"] = mesh_device + + num_devices = mesh_device.get_num_devices() + + # --- Phase 2: Auto-detect distributed prefill --- + # Weight shape is (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT), shard on dim 2 + # Need at least num_devices tiles to shard across devices + num_weight_tiles = dim // SHARD_HEIGHT + can_shard_weights = num_weight_tiles >= num_devices + + if config.prefill_distributed is None: + if num_devices == 1: + to_set["prefill_distributed"] = False + elif not can_shard_weights: + # Can't shard weight tiles across num_devices, use local prefill + to_set["prefill_distributed"] = False + else: + to_set["prefill_distributed"] = True + elif config.prefill_distributed and num_devices == 1: + raise ValueError("RMSNorm1D distributed prefill requires more than one device") + elif config.prefill_distributed and not can_shard_weights: + raise ValueError( + f"RMSNorm1D distributed prefill requires at least {num_devices} weight tiles, got {num_weight_tiles}" + ) + + prefill_distributed = to_set.get("prefill_distributed", config.prefill_distributed) + + # Derive tt_ccl (only needed for distributed prefill) + if config.tt_ccl is None and num_devices > 1 and prefill_distributed: + to_set["tt_ccl"] = get_tt_ccl(mesh_device) + + # --- Phase 3: Local decode configs (Path 1) --- + + tile_size = TILE_SIZE + tile_padded_batch_rows = tile_size * math.ceil(config.max_batch_size / tile_size) + + if config.decode_program_config is None: + norm_core_grid = _compute_norm_core_grid(dim) + to_set["decode_program_config"] = _create_sharded_norm_program_config( + dim, norm_core_grid, tile_padded_batch_rows, tile_size + ) + + if config.decode_memory_config is None: + norm_core_grid = _compute_norm_core_grid(dim) + to_set["decode_memory_config"] = ttnn.create_sharded_memory_config( + (tile_padded_batch_rows, dim // norm_core_grid.num_cores), + norm_core_grid, + ttnn.ShardStrategy.WIDTH, + ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + # --- Phase 5: Resolve LazyWeight (replicated for decode and local prefill) --- + + if num_devices == 1: + mesh_mapper_config_replicated = None + mesh_mapper_config_sharded = None + else: + # For 1D mesh topology, use 2D placement: [Replicate, Replicate] + # This replicates the weight to all devices + mesh_mapper_config_replicated = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate(), ttnn.PlacementReplicate()], + mesh_shape_override=ttnn.MeshShape(1, num_devices), + ) + # For distributed prefill, shard on tensor dim 2 (hidden tiles) + # Weight shape is (1, 1, dim // 32, 32), so dim 2 is the tiles dimension + mesh_mapper_config_sharded = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(2)], + mesh_shape_override=ttnn.MeshShape(1, num_devices), + ) + + # Reshape weight to expected shape (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) + assert dim % SHARD_HEIGHT == 0, "dim must be divisible by SHARD_HEIGHT" + expected_shape = (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) + if config.weight.source.shape != expected_shape: + transformed_source = config.weight.source.reshape(*expected_shape) + else: + transformed_source = config.weight.source + + # Apply unit offset if requested (e.g., for Gemma models) + if config.add_unit_offset: + transformed_source = transformed_source + 1.0 + + # Create a new LazyWeight with the transformed source + transformed_weight = replace(config.weight, source=transformed_source) + + resolved_weight = resolve_lazy_weight( + transformed_weight, + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=mesh_mapper_config_replicated, + ) + to_set["weight"] = resolved_weight + + # --- Phase 7: Create distributed weight LazyWeight (for Path 3) --- + # + # Note: We create a LazyWeight with explicit mesh_mapper_config (ShardTensor2dMesh) + # to shard weights across devices. The get_device_weight() call at line 182 respects + # this config for proper distributed RMSNorm alignment. + + if prefill_distributed and num_devices > 1: + # Create LazyWeight with sharded mesh_mapper_config + # The mesh_mapper_config with ShardTensor2dMesh ensures proper weight distribution + weight_distributed = replace( + transformed_weight, + dtype=ttnn.bfloat16, # RMSNorm weights are always bfloat16 + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=mesh_mapper_config_sharded, + ) + to_set["_weight_distributed"] = weight_distributed + + return replace(config, **to_set) + + +# ============================================================================= +# Helper functions +# ============================================================================= + + +def _compute_norm_core_grid(dim: int, tile_size: int = TILE_SIZE) -> ttnn.CoreGrid: + """Compute core grid for RMSNorm that evenly divides dim.""" + n_tiles = dim // tile_size + rows, cols = _find_grid(n_tiles) + return ttnn.CoreGrid(x=cols, y=rows) + + +def _find_grid(n_tiles: int, max_rows: int = 8, max_cols: int = 8) -> tuple[int, int]: + """Find grid dimensions (rows, cols) that evenly divide n_tiles.""" + max_cores = max_rows * max_cols + target = 32 # based on TTTv1 tuned configs + possible_cores = [k for k in range(1, max_cores + 1) if n_tiles % k == 0] + possible_cores.sort(key=lambda x: abs(x - target)) + + for cores in possible_cores: + for rows in range(1, max_rows + 1): + if cores % rows == 0: + cols = cores // rows + if cols <= max_cols: + return rows, cols + + raise AssertionError(f"Cannot find grid for {n_tiles} tiles within {max_rows}x{max_cols}") + + +def _create_sharded_norm_program_config( + dim: int, grid: ttnn.CoreGrid, tile_padded_batch_rows: int, tile_size: int = TILE_SIZE +) -> ttnn.LayerNormShardedMultiCoreProgramConfig: + """Create LayerNormShardedMultiCoreProgramConfig for RMSNorm.""" + block_w = dim // grid.num_cores // tile_size + subblock_w = 4 + while subblock_w > 0: + if block_w % subblock_w == 0: + break + subblock_w -= 1 + + return ttnn.LayerNormShardedMultiCoreProgramConfig( + compute_with_storage_grid_size=[grid.x, grid.y], + subblock_w=subblock_w, + block_h=tile_padded_batch_rows // tile_size, + block_w=block_w, + inplace=False, + ) + + +# ============================================================================= +# Input Tensor Utilities +# ============================================================================= + + +def _load_input_device_tensor(x: ttnn.Tensor | LazyWeight, config: RMSNorm1DConfig, mode: str) -> ttnn.Tensor: + """ + Resolve the input tensor to ttnn tensor if x is a LazyWeight, otherwise return as is. + + For distributed prefill, uses sharded mesh_mapper_config to shard input across devices. + For decode/non-distributed prefill, replicates input (mesh_mapper_config=None). + + Args: + x: Input tensor (ttnn.Tensor or LazyWeight wrapping torch tensor) + config: Resolved RMSNorm1DConfig + mode: "decode" or "prefill" + + Returns: + ttnn.Tensor ready for the actual forward computation + """ + assert mode in ("decode", "prefill"), f"mode must be 'decode' or 'prefill', got {mode}" + + if isinstance(x, LazyWeight): + # Determine memory config based on mode (default to DRAM if not specified) + if mode == "decode": + mem_cfg = config.decode_input_memcfg or ttnn.DRAM_MEMORY_CONFIG + else: + mem_cfg = config.prefill_input_memcfg or ttnn.DRAM_MEMORY_CONFIG + + # For distributed prefill, shard input on last dim; otherwise replicate + is_distributed = mode == "prefill" and config.prefill_distributed + if is_distributed: + num_devices = config.mesh_device.get_num_devices() + mesh_mapper_config = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape(1, num_devices), + ) + else: + mesh_mapper_config = None + + resolved_x = resolve_lazy_weight( + x, + device=config.mesh_device, + memory_config=mem_cfg, + mesh_mapper_config=mesh_mapper_config, + layout=ttnn.TILE_LAYOUT, + ) + return resolved_x.get_device_weight() + + # Already a ttnn.Tensor - return as is + assert isinstance(x, ttnn.Tensor), f"x must be ttnn.Tensor or LazyWeight, got {type(x)}" + return x diff --git a/code/models/common/modules/rmsnorm/rmsnorm_2d.py b/code/models/common/modules/rmsnorm/rmsnorm_2d.py new file mode 100644 index 0000000000000000000000000000000000000000..2079d7266b8027d1858b34195787302f11d65b4d --- /dev/null +++ b/code/models/common/modules/rmsnorm/rmsnorm_2d.py @@ -0,0 +1,511 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +RMSNorm2D: RMSNorm for 2D mesh topologies (TG, Galaxy). + +Execution paths: + - decode_forward - Sharded, Linear topology, cluster_axis=1 + - prefill_forward - Interleaved, Linear topology, cluster_axis=1 + +Key design: + - Both decode and prefill are ALWAYS distributed for 2D mesh + - Uses Linear topology with cluster_axis=1 (gather across columns) + - Decode uses sharded configs, Prefill uses interleaved + - Separate weight for distributed path (sharded across columns, replicated across rows) +""" + +from dataclasses import dataclass, replace +from typing import Optional + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight +from models.common.modules.tt_ccl import get_tt_ccl +from models.common.tensor_utils import TILE_SIZE + +# ============================================================================= +# Constants +# ============================================================================= + +SHARD_HEIGHT = TILE_SIZE + + +# ============================================================================= +# RMSNorm2DConfig +# ============================================================================= + + +@dataclass +class RMSNorm2DConfig: + """ + Configuration for RMSNorm2D - only fields relevant to 2D mesh topologies. + + Both paths are always distributed with Linear topology and cluster_axis=1. + + Paths: + - Path 4: 2D distributed decode (sharded) - requires decode_* configs + - Path 5: 2D distributed prefill (interleaved) - no program_config needed + + Simple usage: + config = RMSNorm2DConfig(weight, cluster_shape=(4, 8)) + + Override any field: + config = RMSNorm2DConfig(weight, cluster_shape=(4, 8), eps=1e-6) + """ + + # Required: weight and cluster_shape + weight: LazyWeight + cluster_shape: tuple[int, int] | None = None # (rows, cols), e.g. (4, 8) for TG + + # Normalization settings + eps: float = 1e-5 + add_unit_offset: bool = False + + # Device and collectives + mesh_device: ttnn.MeshDevice | None = None + tt_ccl: "TT_CCL | None" = None # type: ignore + + # Batch size for decode + max_batch_size: int = 32 + + # Input memory configs (for _load_input_device_tensor_2d) + decode_input_memcfg: ttnn.MemoryConfig | None = None + prefill_input_memcfg: ttnn.MemoryConfig | None = None + + # 2D distributed decode configs - sharded + decode_progcfg: ttnn.LayerNormShardedMultiCoreProgramConfig | None = None + decode_stats_memcfg: ttnn.MemoryConfig | None = None + + # Compute kernel config (only for prefill - decode uses program_config) + compute_kernel_config_prefill: ttnn.WormholeComputeKernelConfig | None = None + + def is_resolved(self) -> bool: + """Check if all required fields are resolved.""" + required = [ + "mesh_device", + "cluster_shape", + "tt_ccl", + # Decode sharded configs + "decode_input_memcfg", + "decode_progcfg", + "decode_stats_memcfg", + # Compute config + "compute_kernel_config_prefill", + # Weight (2D always uses distributed/sharded) + "weight", + ] + return all(getattr(self, f) is not None for f in required) + + +# ============================================================================= +# RMSNorm2D +# ============================================================================= + + +class RMSNorm2D(LightweightModule): + """ + RMSNorm for 2D mesh topologies (TG, Galaxy). + + Both decode and prefill are always distributed with Linear topology. + + Execution paths (bound at construction): + - decode_forward -> _decode_2d_distributed (Path 4) - sharded + - prefill_forward -> _prefill_2d_distributed (Path 5) - interleaved + + Simple API: + norm = RMSNorm2D(weight) + + Power API: + config = RMSNorm2DConfig(weight, cluster_shape=(4, 8), max_batch_size=64) + norm = RMSNorm2D.from_config(config) + """ + + def __init__(self, weight: LazyWeight): + """ + Simple API - derives all config from weight. + + Args: + weight: RMSNorm weight tensor of shape (dim,) or (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) + + Note: Use from_config() to customize eps, add_unit_offset, cluster_shape, or other settings. + """ + super().__init__() + self.config = _resolve_2d_config(RMSNorm2DConfig(weight=weight)) + self._device_weights_loaded = False + + @classmethod + def from_config(cls, config: RMSNorm2DConfig): + """Power API - any level of customization via config.""" + instance = object.__new__(cls) + super(RMSNorm2D, instance).__init__() + instance.config = _resolve_2d_config(config) + instance._device_weights_loaded = False + return instance + + def load_device_weights(self): + """Load weights to device lazily.""" + if self._device_weights_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device weights!" + + # Load weight (sharded across columns, replicated across rows) + self.weight = self.config.weight.get_device_weight() + + self._device_weights_loaded = True + + def decode_forward(self, x: "ttnn.Tensor | LazyWeight") -> ttnn.Tensor: + """ + 2D distributed decode - sharded for TG/Galaxy. + + Uses Linear topology, cluster_axis=1, with program_config. + + Execution: + to_sharded -> rms_norm_pre_all_gather(sharded) -> all_gather(cluster_axis=1) + -> rms_norm_post_all_gather(sharded) + """ + self.load_device_weights() + x = _load_input_device_tensor_2d(x, self.config, mode="decode") + cfg = self.config + + # Convert to sharded memory config + x = ttnn.to_memory_config(x, memory_config=cfg.decode_input_memcfg) + + # Run distributed rmsnorm part 1 (sharded) + tt_stats = ttnn.rms_norm_pre_all_gather(x, program_config=cfg.decode_progcfg) + + # All gather stats along cluster axis 1 (columns) + cluster_axis = 1 + tt_stats = ttnn.experimental.all_gather_async( + tt_stats, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis), + num_links=1, + cluster_axis=cluster_axis, + topology=ttnn.Topology.Linear, + memory_config=cfg.decode_stats_memcfg, + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + + # Run distributed rmsnorm part 2 + tt_out = ttnn.rms_norm_post_all_gather( + x, + epsilon=cfg.eps, + weight=self.weight, + program_config=cfg.decode_progcfg, + stats=tt_stats, + ) + tt_stats.deallocate(True) + + return tt_out + + def prefill_forward(self, x: "ttnn.Tensor | LazyWeight") -> ttnn.Tensor: + """ + 2D distributed prefill - interleaved for TG/Galaxy. + + Uses Linear topology, cluster_axis=1. + + Execution: + rms_norm_pre_all_gather -> reshape -> all_gather(cluster_axis=1) + -> rms_norm_post_all_gather + """ + self.load_device_weights() + x = _load_input_device_tensor_2d(x, self.config, mode="prefill") + cfg = self.config + + # Run distributed rmsnorm part 1 + tt_stats = ttnn.rms_norm_pre_all_gather( + x, compute_kernel_config=cfg.compute_kernel_config_prefill, dtype=ttnn.bfloat16 + ) + + # Reshape stats for all_gather (preserve batch dimension for multi-batch prefill) + batch = x.shape[0] + padded_shape = (batch, 1, x.shape[-2], 32) + tt_stats = ttnn.reshape(tt_stats, ttnn.Shape(padded_shape)) + + # All gather stats along cluster axis 1 (columns) + cluster_axis = 1 + tt_stats_gathered = ttnn.experimental.all_gather_async( + tt_stats, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis), + num_links=1, + cluster_axis=cluster_axis, + topology=ttnn.Topology.Linear, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), + chunks_per_sync=10, + num_workers_per_link=2, + num_buffers_per_channel=2, + ) + tt_stats.deallocate(True) + + # Run distributed rmsnorm part 2 + tt_out = ttnn.rms_norm_post_all_gather( + x, + tt_stats_gathered, + epsilon=cfg.eps, + weight=self.weight, + compute_kernel_config=cfg.compute_kernel_config_prefill, + ) + tt_stats_gathered.deallocate(True) + + return tt_out + + # ========================================================================= + # Forward dispatcher + # ========================================================================= + + def forward(self, x: "ttnn.Tensor | LazyWeight", mode: str) -> ttnn.Tensor: + """Dispatch to decode_forward or prefill_forward based on mode.""" + if mode == "decode": + return self.decode_forward(x) + else: + return self.prefill_forward(x) + + # ========================================================================= + # Factory method for TTTv1 compatibility + # ========================================================================= + + @classmethod + def from_model_args( + cls, + mesh_device, + tt_ccl, + args, + state_dict, + weight_cache_path, + layer_num: int, + weight_key: str, + state_dict_prefix: Optional[str] = None, + ): + """Factory method for backward compatibility with ModelArgs.""" + # Validate TG topology + cluster_shape = tuple(mesh_device.shape) + valid_shapes = [(4, 8), (8, 4)] + if cluster_shape not in valid_shapes: + raise ValueError( + f"RMSNorm2D requires Galaxy topology (4x8 or 8x4), got {cluster_shape}. " + "Use RMSNorm1D for non-Galaxy devices." + ) + + # Build weight name from state_dict_prefix and weight_key + if state_dict_prefix: + weight_name = f"{state_dict_prefix}{weight_key}.weight" + else: + if layer_num is None: + weight_name = f"{weight_key}.weight" + else: + weight_name = f"layers.{layer_num}.{weight_key}.weight" + + # Transform weight to expected shape: (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) + dim = args.dim + torch_weight = ( + state_dict[weight_name].unsqueeze(0).view(1, 1, dim).reshape([1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT]) + ) + + # Add offset before creating LazyWeight + if args.rms_norm_add_unit_offset: + torch_weight = torch_weight + 1.0 + + # Create LazyWeight (mesh_mapper_config is handled by _resolve_2d_config) + lazy_weight = LazyWeight( + source=torch_weight, + dtype=ttnn.bfloat16, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + cache_dir_weight_name=(weight_cache_path, weight_name + "_2d") if weight_cache_path else None, + ) + + # Build config (dim derived from weight in _resolve_2d_config) + config = RMSNorm2DConfig( + weight=lazy_weight, + cluster_shape=cluster_shape, + eps=args.norm_eps, + add_unit_offset=False, # Already applied above + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=args.max_batch_size, + ) + + return cls.from_config(config) + + +# ============================================================================= +# Config resolution +# ============================================================================= + + +def _resolve_2d_config(config: RMSNorm2DConfig) -> RMSNorm2DConfig: + """Resolve config defaults for RMSNorm2D.""" + + to_set = {} + + # --- Phase 1: Foundational fields --- + + # Derive dim from weight (works for any shape: [dim], [1, dim], [1, 1, dim//32, 32], etc.) + dim = config.weight.source.numel() + + # Derive mesh_device from weight + mesh_device = config.mesh_device + if mesh_device is None: + mesh_device = config.weight.device + if mesh_device is None: + mesh_device = ttnn.GetDefaultDevice() + if config.mesh_device is None: + to_set["mesh_device"] = mesh_device + + assert mesh_device is not None, "mesh_device must be available!" + + # Derive cluster_shape from mesh_device if not provided + cluster_shape = config.cluster_shape + if cluster_shape is None: + cluster_shape = tuple(mesh_device.shape) + to_set["cluster_shape"] = cluster_shape + + assert len(cluster_shape) == 2 and all( + d > 1 for d in cluster_shape + ), f"cluster_shape must be 2D with both dims > 1, got {cluster_shape}" + + num_rows, num_cols = cluster_shape + + # Derive tt_ccl (always needed for 2D) + if config.tt_ccl is None: + to_set["tt_ccl"] = get_tt_ccl(mesh_device) + + # --- Phase 2: Compute kernel config (prefill only - decode uses program_config) --- + + if config.compute_kernel_config_prefill is None: + # 2D uses fp32=False, packer_l1=False (from DistributedNorm.ln_cfg) + to_set["compute_kernel_config_prefill"] = ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi2, + math_approx_mode=False, + fp32_dest_acc_en=False, + packer_l1_acc=False, + ) + + # --- Phase 3: 2D distributed decode configs (Path 4) --- + + hidden_size_per_device = dim // num_cols + core_grid_ln = ( + min(4, hidden_size_per_device // TILE_SIZE // 8), + 8, + ) + num_cores_ln = core_grid_ln[0] * core_grid_ln[1] + + if config.decode_input_memcfg is None: + to_set["decode_input_memcfg"] = ttnn.create_sharded_memory_config( + shape=(1, 1, 32, hidden_size_per_device), + core_grid=ttnn.CoreGrid(y=core_grid_ln[0], x=core_grid_ln[1]), + strategy=ttnn.ShardStrategy.WIDTH, + ) + + if config.decode_progcfg is None: + to_set["decode_progcfg"] = ttnn.LayerNormShardedMultiCoreProgramConfig( + compute_with_storage_grid_size=(core_grid_ln[1], core_grid_ln[0]), + subblock_w=(hidden_size_per_device // num_cores_ln) // TILE_SIZE, + block_h=1, + block_w=(hidden_size_per_device // num_cores_ln) // TILE_SIZE, + inplace=False, + ) + + if config.decode_stats_memcfg is None: + # Stats memory: 32 x (32 * num_cols) where num_cols is the cluster width + to_set["decode_stats_memcfg"] = ttnn.create_sharded_memory_config( + shape=[1, 1, 32, 32 * num_cols], + core_grid=ttnn.CoreGrid(y=1, x=1), + strategy=ttnn.ShardStrategy.WIDTH, + ) + + # --- Phase 4: Resolve weight (sharded across columns, replicated across rows) --- + # 2D always uses distributed paths + # Weight shape must be (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT), shard on dim 2 + + # Reshape weight to expected shape if needed (handles [dim], [1, dim], etc.) + assert dim % SHARD_HEIGHT == 0, f"dim must be divisible by SHARD_HEIGHT={SHARD_HEIGHT}, got {dim}" + expected_shape = (1, 1, dim // SHARD_HEIGHT, SHARD_HEIGHT) + if config.weight.source.shape != expected_shape: + transformed_source = config.weight.source.reshape(*expected_shape) + else: + transformed_source = config.weight.source + + # Apply unit offset if requested (e.g., for Gemma models) + if config.add_unit_offset: + transformed_source = transformed_source + 1.0 + + # Create a new LazyWeight with the transformed source + transformed_weight = replace(config.weight, source=transformed_source) + + # Shard on tensor dim 2 across columns (cluster axis 1), replicate across rows (cluster axis 0) + mesh_mapper_config = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(2)], + mesh_shape_override=ttnn.MeshShape(list(cluster_shape)), + ) + + resolved_weight = resolve_lazy_weight( + transformed_weight, + device=mesh_device, + layout=ttnn.ROW_MAJOR_LAYOUT, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=mesh_mapper_config, + ) + to_set["weight"] = resolved_weight + + return replace(config, **to_set) + + +# ============================================================================= +# Input Tensor Utilities +# ============================================================================= + + +def _load_input_device_tensor_2d(x: "ttnn.Tensor | LazyWeight", config: RMSNorm2DConfig, mode: str) -> ttnn.Tensor: + """ + Resolve the input tensor to ttnn tensor if x is a LazyWeight, otherwise return as is. + + For 2D distributed, shards input on last dim across columns (cluster_axis=1). + + Args: + x: Input tensor (ttnn.Tensor or LazyWeight wrapping torch tensor) + config: Resolved RMSNorm2DConfig + mode: "decode" or "prefill" + + Returns: + ttnn.Tensor ready for the actual forward computation + """ + assert mode in ("decode", "prefill"), f"mode must be 'decode' or 'prefill', got {mode}" + + if isinstance(x, LazyWeight): + num_rows, num_cols = config.cluster_shape + + # Determine memory config based on mode + if mode == "decode": + mem_cfg = config.decode_input_memcfg or ttnn.DRAM_MEMORY_CONFIG + else: + mem_cfg = config.prefill_input_memcfg or ttnn.DRAM_MEMORY_CONFIG + + # Shard input on last dim across columns, replicate across rows + mesh_mapper_config = ttnn.MeshMapperConfig( + placements=[ttnn.PlacementReplicate(), ttnn.PlacementShard(-1)], + mesh_shape_override=ttnn.MeshShape(num_rows, num_cols), + ) + + resolved_x = resolve_lazy_weight( + x, + device=config.mesh_device, + memory_config=mem_cfg, + mesh_mapper_config=mesh_mapper_config, + layout=ttnn.TILE_LAYOUT, + ) + return resolved_x.get_device_weight() + + # Already a ttnn.Tensor - return as is + assert isinstance(x, ttnn.Tensor), f"x must be ttnn.Tensor or LazyWeight, got {type(x)}" + return x diff --git a/code/models/common/modules/rope/rope_1d.py b/code/models/common/modules/rope/rope_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..9384325d1876968081375229e6f13e66af385eeb --- /dev/null +++ b/code/models/common/modules/rope/rope_1d.py @@ -0,0 +1,510 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTTv2-style Rotary Position Embedding (RoPE) setup for 1D-topology devices. + +RotarySetup1D takes pre-computed cos/sin rotation matrices (as LazyWeight) and +provides efficient position-based lookup at runtime. + +Core API: + - decode_forward(rot_idxs): Look up cos/sin rotation matrices from ttnn.Tensor indices + - prefill_forward(start_pos, seq_len): Slice cos/sin matrices for prefill + - get_both_trans_mats(): Decode + prefill transformation matrices (read-only) + +Backward-compatible API (will retire with TTTv1): + - get_rot_idxs(position_idxs): Convert torch position indices to ttnn tensor + - get_rot_mats(position_idxs): Combined get_rot_idxs + decode_forward + +Helper: + - prepare_rot_idxs(config, position_idxs, on_host): Standalone torch→ttnn index conversion + +Note: torch is used at construction time for building transformation matrices. + The decode_forward() and prefill_forward() methods use only ttnn ops. +""" + +from dataclasses import dataclass, replace +from typing import Any, Dict, List, Optional, Tuple, Union + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_weight import LazyWeight, resolve_lazy_weight +from models.common.tensor_utils import TILE_SIZE, get_rot_transformation_mat, nearest_32 + +# ============================================================================= +# Config dataclass +# ============================================================================= + + +@dataclass +class Rope1DConfig: + """ + Configuration for RotarySetup1D. + + Simple usage: + rope = RotarySetup1D(cos_lw, sin_lw, max_batch_size=32) + + Power API: + config = Rope1DConfig(cos_lw, sin_lw, max_batch_size=32, + device=mesh_device, use_qk_fused=True) + rope = RotarySetup1D.from_config(config) + """ + + # Required: cos/sin embedding lookup tables as LazyWeight + # source shape: [1, 1, max_seq_len, head_dim] + cos_matrix: LazyWeight + sin_matrix: LazyWeight + + # Required scalars + max_batch_size: int + + # Optional (derived from weights if None) + head_dim: int | None = None # derived from cos_matrix.source.shape[-1] + device: Any | None = None # derived from cos_matrix.device + use_qk_fused: bool = False + datatype: ttnn.DataType = ttnn.bfloat16 + + # Derived (resolved by _resolve_rope_config, None = auto) + batch_size_per_device_group: int | None = None + core_grid: Any | None = None + batch_grid: Any | None = None + decode_trans_mat_mem_config: Optional[ttnn.MemoryConfig] = None + cos_sin_shard_mem_config: Optional[ttnn.MemoryConfig] = None + + # Internal: transformation matrix LazyWeights (built by _resolve_rope_config) + _decode_trans_mat: Optional[LazyWeight] = None + _prefill_trans_mat: Optional[LazyWeight] = None + + +# ============================================================================= +# RotarySetup1D +# ============================================================================= + + +class RotarySetup1D(LightweightModule): + """ + Rotary position embedding setup for non-TG (1D) devices. + + Takes pre-computed cos/sin lookup tables as LazyWeight and provides + efficient position-based rotation matrix lookup at runtime. + + Simple API: + rope = RotarySetup1D(cos_lw, sin_lw, max_batch_size=32) + rot_idxs = prepare_rot_idxs(rope.config, position_idxs) + [cos, sin] = rope.decode_forward(rot_idxs) + [cos, sin] = rope.prefill_forward(start_pos=0, seq_len=128) + + Power API: + config = Rope1DConfig(cos_lw, sin_lw, max_batch_size=32, use_qk_fused=True) + rope = RotarySetup1D.from_config(config) + + Backward-compatible API: + rope.get_rot_mats(position_idxs) # accepts torch.Tensor or ttnn.Tensor + rope.get_rot_idxs(position_idxs, on_host=True) + """ + + def __init__( + self, + cos_matrix: LazyWeight, + sin_matrix: LazyWeight, + max_batch_size: int, + ) -> None: + """ + Simple API — takes cos/sin as LazyWeight, derives all other config. + + Args: + cos_matrix: LazyWeight with source shape [1, 1, max_seq_len, head_dim]. + sin_matrix: LazyWeight with source shape [1, 1, max_seq_len, head_dim]. + max_batch_size: Maximum batch size for decode mode. + """ + super().__init__() + config = Rope1DConfig( + cos_matrix=cos_matrix, + sin_matrix=sin_matrix, + max_batch_size=max_batch_size, + ) + self.config = _resolve_rope_config(config) + self._device_weights_loaded = False + + @classmethod + def from_config(cls, config: Rope1DConfig): + """Power API — construct from config dataclass with full customization.""" + instance = object.__new__(cls) + super(RotarySetup1D, instance).__init__() + instance.config = _resolve_rope_config(config) + instance._device_weights_loaded = False + return instance + + def load_device_weights(self): + """Load cos/sin and transformation matrices to device. Called lazily on first use.""" + if self._device_weights_loaded: + return + + cfg = self.config + + self.cos_matrix = cfg.cos_matrix.get_device_weight() + self.sin_matrix = cfg.sin_matrix.get_device_weight() + self.transformation_mat = cfg._decode_trans_mat.get_device_weight() + self.transformation_mat_prefill = cfg._prefill_trans_mat.get_device_weight() + + self._device_weights_loaded = True + + # ---- Core API ---- + + def decode_forward(self, rot_idxs: ttnn.Tensor) -> List[ttnn.Tensor]: + """Look up cos/sin rotation matrices for given position indices (decode mode). + + Args: + rot_idxs: ttnn.Tensor of shape [1, batch], dtype uint32. + + Returns: + [cos, sin] — height-sharded rotation matrices. + """ + self.load_device_weights() + cfg = self.config + + # Send to device if needed + if rot_idxs.device != cfg.device: + rot_idxs = ttnn.to_device(rot_idxs, cfg.device, memory_config=ttnn.DRAM_MEMORY_CONFIG) + + # Embedding lookup: [1, batch] → [1, batch, head_dim] + cos = ttnn.embedding(rot_idxs, self.cos_matrix, layout=ttnn.TILE_LAYOUT) + sin = ttnn.embedding(rot_idxs, self.sin_matrix, layout=ttnn.TILE_LAYOUT) + + # Reshape: [1, batch, head_dim] → [1, 1, batch, head_dim] → [1, batch, 1[32], head_dim] + cos = ttnn.unsqueeze_to_4D(cos) + sin = ttnn.unsqueeze_to_4D(sin) + cos = ttnn.transpose(cos, 1, 2) + sin = ttnn.transpose(sin, 1, 2) + + # Trim to actual batch size. + # When batch_size_per_device_group is not a multiple of TILE_SIZE (e.g., batch=1), + # the embedding produces nearest_32(batch) rows that need trimming. + # When tile-aligned (e.g., batch=32), this is a no-op slice. + cos = cos[:, : cfg.batch_size_per_device_group, :, :] + sin = sin[:, : cfg.batch_size_per_device_group, :, :] + + # Shard to batch cores + cos = ttnn.interleaved_to_sharded(cos, cfg.cos_sin_shard_mem_config) + sin = ttnn.interleaved_to_sharded(sin, cfg.cos_sin_shard_mem_config) + + return [cos, sin] + + def prefill_forward(self, start_pos: int, seq_len: int, pad_to: int | None = None) -> List[ttnn.Tensor]: + """Slice cos/sin matrices for prefill mode. + + This replaces the duplicated cos/sin slicing logic found in every model's + prepare_inputs_prefill(). Models no longer need to reach into + rope_setup.cos_matrix / sin_matrix directly. + + Args: + start_pos: Starting position in the sequence. + seq_len: Number of positions to slice (typically the padded input length S). + pad_to: If set, zero-pad the sequence dim to this length (for SDPA alignment). + Ignored when pad_to <= seq_len. + + Returns: + [cos_slice, sin_slice] — interleaved DRAM tensors, + shape [1, 1, seq_len (or pad_to), head_dim]. + """ + self.load_device_weights() + + mat_len = self.cos_matrix.shape[2] + end_pos = start_pos + seq_len + assert mat_len >= end_pos, f"Requested range [{start_pos}:{end_pos}) exceeds cos/sin table length {mat_len}" + + cos_slice = self.cos_matrix[:, :, start_pos:end_pos, :] + sin_slice = self.sin_matrix[:, :, start_pos:end_pos, :] + + if pad_to is not None and pad_to > seq_len: + pad_len = pad_to - seq_len + padding = [(0, 0)] * 4 + padding[2] = (0, pad_len) + cos_slice = ttnn.pad(cos_slice, padding=padding, value=0.0) + sin_slice = ttnn.pad(sin_slice, padding=padding, value=0.0) + + return [cos_slice, sin_slice] + + def forward(self, mode: str, **kwargs) -> List[ttnn.Tensor]: + """Dispatch to decode_forward or prefill_forward based on mode. + + Args: + mode: "decode" or "prefill". + **kwargs: Forwarded to the underlying method. + decode: rot_idxs (ttnn.Tensor) + prefill: start_pos (int), seq_len (int), pad_to (int | None) + """ + if mode == "decode": + return self.decode_forward(**kwargs) + elif mode == "prefill": + return self.prefill_forward(**kwargs) + else: + raise ValueError(f"Unknown mode: {mode!r}. Expected 'decode' or 'prefill'.") + + def get_both_trans_mats(self) -> Dict[str, ttnn.Tensor]: + """Return both decode and prefill transformation matrices.""" + self.load_device_weights() + return {"decode": self.transformation_mat, "prefill": self.transformation_mat_prefill} + + # ---- Backward-compatible API (will retire with TTTv1) ---- + + def get_rot_idxs(self, position_idxs: "torch.Tensor", on_host: bool = False) -> ttnn.Tensor: + """Convert torch position indices to ttnn tensor. Delegates to prepare_rot_idxs().""" + return prepare_rot_idxs(self.config, position_idxs, on_host=on_host) + + def get_rot_mats( + self, + position_idxs: "Union[torch.Tensor, ttnn.Tensor]", + return_rot_idxs: bool = False, + ) -> "Union[List[ttnn.Tensor], Tuple[List[ttnn.Tensor], ttnn.Tensor]]": + """Look up cos/sin from torch or ttnn position indices. + + For new code, prefer: rot_idxs = prepare_rot_idxs(...); rope.decode_forward(rot_idxs) + """ + import torch + + if isinstance(position_idxs, torch.Tensor): + rot_idxs = self.get_rot_idxs(position_idxs) + else: + rot_idxs = position_idxs + + cos_sin = self.decode_forward(rot_idxs) + + if return_rot_idxs: + return cos_sin, rot_idxs + return cos_sin + + # ---- Factory for TTTv1 backward compatibility ---- + + @classmethod + def from_model_args(cls, device: Any, args, model_name: str = "unknown"): + """Factory method for backward compatibility with ModelArgs. + + Args: + device: MeshDevice or single device. + args: Model arguments (ModelArgs instance). + model_name: Model name for logging (default "unknown"). + """ + if args.is_galaxy: + raise ValueError("RotarySetup1D cannot be used for Galaxy devices.") + + from models.tt_transformers.tt.common import rope_scaling_model_factory + + rope_scaling = None + if hasattr(args, "rope_scaling_params") and args.rope_scaling_params is not None: + rope_scaling = rope_scaling_model_factory( + args.rope_scaling_params, getattr(args, "original_max_context_len", None) + ) + + # Compute cos/sin torch tensors using TTTv1 reference + from models.tt_transformers.tt.rope import compute_gather_cos_sin + + cos_torch, sin_torch = compute_gather_cos_sin( + dhead=args.head_dim, end=2 * args.max_seq_len, theta=args.rope_theta, rope_scaling=rope_scaling + ) + + # Wrap in LazyWeight + cos_lw = LazyWeight(source=cos_torch, device=device) + sin_lw = LazyWeight(source=sin_torch, device=device) + + config = Rope1DConfig( + cos_matrix=cos_lw, + sin_matrix=sin_lw, + max_batch_size=args.max_batch_size, + head_dim=args.head_dim, + device=device, + use_qk_fused=getattr(args, "use_qk_fused", False), + datatype=ttnn.bfloat16, + ) + return cls.from_config(config) + + +# ============================================================================= +# Config resolver +# ============================================================================= + + +def _resolve_rope_config(config: Rope1DConfig) -> Rope1DConfig: + """Resolve None fields in Rope1DConfig to sensible defaults. + + Fills in: + - head_dim: derived from cos_matrix.source.shape[-1] + - device: derived from cos_matrix.device if None + - batch_size_per_device_group, core_grid, batch_grid + - decode_trans_mat_mem_config, cos_sin_shard_mem_config + - cos_matrix/sin_matrix LazyWeight fields: device, memory_config, layout, dtype + - _decode_trans_mat, _prefill_trans_mat: internal transformation matrix LazyWeights + """ + to_set = {} + + # --- Phase 1: Foundational fields --- + + head_dim = config.head_dim + if head_dim is None: + head_dim = config.cos_matrix.source.shape[-1] + to_set["head_dim"] = head_dim + + device = config.device + if device is None: + device = config.cos_matrix.device + if device is None: + device = ttnn.GetDefaultDevice() + if config.device is None: + to_set["device"] = device + + assert device is not None, "device must be available at this point!" + + # --- Phase 2: Derived scalars --- + + batch_size_per_device_group = config.batch_size_per_device_group + if batch_size_per_device_group is None: + batch_size_per_device_group = config.max_batch_size * 2 if config.use_qk_fused else config.max_batch_size + to_set["batch_size_per_device_group"] = batch_size_per_device_group + + core_grid = config.core_grid + if core_grid is None: + core_grid = device.compute_with_storage_grid_size() + to_set["core_grid"] = core_grid + + batch_grid = config.batch_grid + if batch_grid is None: + batch_grid = ttnn.num_cores_to_corerangeset(batch_size_per_device_group, core_grid, row_wise=True) + to_set["batch_grid"] = batch_grid + + # --- Phase 3: Memory configs --- + + if config.decode_trans_mat_mem_config is None: + to_set["decode_trans_mat_mem_config"] = ttnn.create_sharded_memory_config( + shape=(TILE_SIZE, TILE_SIZE), + core_grid=batch_grid, + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + if config.cos_sin_shard_mem_config is None: + to_set["cos_sin_shard_mem_config"] = ttnn.create_sharded_memory_config( + shape=(TILE_SIZE, head_dim), + core_grid=batch_grid, + strategy=ttnn.ShardStrategy.HEIGHT, + orientation=ttnn.ShardOrientation.ROW_MAJOR, + use_height_and_width_as_shard_shape=True, + ) + + # --- Phase 4: Resolve LazyWeights --- + + to_set["cos_matrix"] = resolve_lazy_weight( + config.cos_matrix, + device=device, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=None, # replicated + layout=ttnn.TILE_LAYOUT, + dtype=config.datatype, + ) + to_set["sin_matrix"] = resolve_lazy_weight( + config.sin_matrix, + device=device, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=None, # replicated + layout=ttnn.TILE_LAYOUT, + dtype=config.datatype, + ) + + # --- Phase 5: Transformation matrix LazyWeights --- + + decode_trans_mat_mem_config = to_set.get("decode_trans_mat_mem_config", config.decode_trans_mat_mem_config) + + assert config._decode_trans_mat is None, "_decode_trans_mat is internal and should not be set by the user" + decode_trans_mat_source = get_rot_transformation_mat(dhead=TILE_SIZE).repeat( + 1, + 1, + batch_size_per_device_group, + 1, + ) + to_set["_decode_trans_mat"] = LazyWeight( + source=decode_trans_mat_source, + device=device, + memory_config=decode_trans_mat_mem_config, + mesh_mapper_config=None, # replicated + layout=ttnn.TILE_LAYOUT, + dtype=config.datatype, + ) + + assert config._prefill_trans_mat is None, "_prefill_trans_mat is internal and should not be set by the user" + prefill_trans_mat_source = get_rot_transformation_mat(dhead=TILE_SIZE) + to_set["_prefill_trans_mat"] = LazyWeight( + source=prefill_trans_mat_source, + device=device, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper_config=None, # replicated + layout=ttnn.TILE_LAYOUT, + dtype=config.datatype, + ) + + return replace(config, **to_set) + + +# ============================================================================= +# Standalone helper function +# ============================================================================= + + +def prepare_rot_idxs( + config: Rope1DConfig, + position_idxs: "torch.Tensor", + on_host: bool = False, +) -> ttnn.Tensor: + """Convert torch position indices to a TTNN rotation index tensor. + + This is the standalone equivalent of RotarySetup1D.get_rot_idxs(). + Useful for pre-allocating host tensors for trace mode. + + Args: + config: Resolved Rope1DConfig (device must be set). + position_idxs: 1D torch tensor of position indices, shape [batch]. + on_host: If True, return tensor on host (for later copy_host_to_device_tensor). + + Returns: + ttnn tensor of shape [1, nearest_32(batch)], dtype uint32. + """ + import torch + + assert isinstance(position_idxs, torch.Tensor), "Position ids must be a torch tensor" + assert len(position_idxs.shape) == 1, "position idxs must be a [batch] tensor" + assert config.device is not None, "config.device must be set (use a resolved config)" + assert config.batch_size_per_device_group is not None, "config must be resolved (_resolve_rope_config)" + + if config.use_qk_fused: + position_idxs = position_idxs.repeat(2) + assert ( + position_idxs.shape[0] == config.batch_size_per_device_group + ), "Position idxs must match batch_size_per_device_group" + + batch = position_idxs.shape[0] + position_idxs = position_idxs.reshape(1, batch) + assert torch.min(position_idxs) >= 0, "Position idxs must be non-negative" + + # Pad to nearest tile boundary + pad_size = nearest_32(batch) - batch + position_idxs = torch.nn.functional.pad(position_idxs, (0, pad_size), "constant", 0) + + device = config.device + + if on_host: + rot_idxs = ttnn.as_tensor( + position_idxs, + dtype=ttnn.uint32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(device), + ) + else: + rot_idxs = ttnn.as_tensor( + position_idxs, + dtype=ttnn.uint32, + layout=ttnn.ROW_MAJOR_LAYOUT, + device=device, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(device), + ) + + return rot_idxs diff --git a/code/models/common/modules/sampling/params.py b/code/models/common/modules/sampling/params.py new file mode 100644 index 0000000000000000000000000000000000000000..f120ff1973fba8a8a643a56a2e4a2da3dd6ce35b --- /dev/null +++ b/code/models/common/modules/sampling/params.py @@ -0,0 +1,574 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Topology-neutral preparation of TTTv2 sampling request parameters. + +This module deliberately contains no mesh or TTNN policy. A caller resolves +the sampler capabilities from its ``Sampling1DConfig`` and passes them here. +The resulting immutable value contains every request-owned sampling field in +slot order and is safe to slice or retain across eager and trace lifecycles. +""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Sequence +from dataclasses import dataclass +from numbers import Integral +from typing import Any, Literal + +import torch + +from models.common.sampling.sampling_params import SamplingParams + +SamplingPath = Literal["argmax", "topk"] +SamplingRowPath = Literal["inactive", "argmax", "topk"] +LogProbMode = Literal["none", "sampled_token", "top_n"] + + +@dataclass(frozen=True) +class PreparedSamplingParams: + """One normalized device-sampling request. + + All row-owned tuples have exactly ``batch_size`` entries. ``temperature`` + holds the inverse temperature consumed by ``Sampling1D``. Greedy rows use + the device representation ``top_k=1``, ``top_p=0`` and ``temperature=1``. + + ``sampling_path`` is batch-wide because the common runtime selects one + program for the entire lane. A mixed greedy/stochastic request therefore + uses ``topk`` even when force-argmax is available. + """ + + top_k: tuple[int, ...] + top_p: tuple[float, ...] + temperature: tuple[float, ...] + presence_penalty: tuple[float, ...] + frequency_penalty: tuple[float, ...] + repetition_penalty: tuple[float, ...] + seeds: tuple[int | None, ...] + enable_log_probs: tuple[bool, ...] + num_logprobs: tuple[int, ...] + logprob_modes: tuple[LogProbMode, ...] + greedy_mask: tuple[bool, ...] + row_paths: tuple[SamplingRowPath, ...] + active_mask: tuple[bool, ...] + sampling_path: SamplingPath + active_rows: int + batch_size: int + max_device_top_k: int + prompt_tokens: Any | None = None + output_tokens: Any | None = None + slot_remap: Any | None = None + + def __post_init__(self) -> None: + if self.active_rows <= 0 or self.active_rows > self.batch_size: + raise ValueError("active_rows must be in [1, batch_size]") + row_fields = ( + "top_k", + "top_p", + "temperature", + "presence_penalty", + "frequency_penalty", + "repetition_penalty", + "seeds", + "enable_log_probs", + "num_logprobs", + "logprob_modes", + "greedy_mask", + "row_paths", + "active_mask", + ) + for name in row_fields: + if len(getattr(self, name)) != self.batch_size: + raise ValueError(f"{name} must contain exactly batch_size entries") + if sum(self.active_mask) != self.active_rows: + raise ValueError("active_rows must equal the number of active_mask entries") + + @property + def penalties_enabled(self) -> bool: + return ( + any(active and value != 0.0 for active, value in zip(self.active_mask, self.presence_penalty)) + or any(active and value != 0.0 for active, value in zip(self.active_mask, self.frequency_penalty)) + or any(active and value != 1.0 for active, value in zip(self.active_mask, self.repetition_penalty)) + ) + + @property + def log_probs_enabled(self) -> bool: + return any(active and mode != "none" for active, mode in zip(self.active_mask, self.logprob_modes)) + + @property + def all_active_rows_greedy(self) -> bool: + return all(greedy for active, greedy in zip(self.active_mask, self.greedy_mask) if active) + + @property + def all_active_rows_argmax(self) -> bool: + return all(path == "argmax" for active, path in zip(self.active_mask, self.row_paths) if active) + + +_DEFAULTS: dict[str, Any] = { + "temperature": 0.0, + "top_p": 1.0, + "top_k": 1, + "presence_penalty": 0.0, + "frequency_penalty": 0.0, + "repetition_penalty": 1.0, + "seed": None, + "enable_log_probs": False, + "num_logprobs": 0, +} + + +def prepare_sampling_params( + sampling_params: SamplingParams, + batch_size: int, + *, + max_device_top_k: int, + allow_force_argmax: bool, + prompt_tokens: Any | None = None, + output_tokens: Any | None = None, + slot_remap: Any | None = None, +) -> PreparedSamplingParams: + """Normalize, validate, and classify a TTTv2 sampling request. + + Unsupported stochastic ``top_k`` values are rejected rather than clamped. + The check intentionally happens after greedy rows are normalized, because + vLLM commonly represents unrestricted ``top_k`` as the vocabulary size even + for a request whose temperature is zero. + """ + + _validate_policy(batch_size, max_device_top_k, allow_force_argmax) + _validate_sampling_value(sampling_params) + + temperature_input = _as_sequence(getattr(sampling_params, "temperature"), "temperature") + active_rows = len(temperature_input) + if active_rows > batch_size: + raise ValueError(f"temperature describes {active_rows} active rows, exceeding batch_size={batch_size}") + + temperature = _normalize_per_row(temperature_input, "temperature", active_rows, batch_size) + top_p = _normalize_per_row(getattr(sampling_params, "top_p"), "top_p", active_rows, batch_size) + top_k = _normalize_per_row(getattr(sampling_params, "top_k"), "top_k", active_rows, batch_size) + presence_penalty = _normalize_per_row( + getattr(sampling_params, "presence_penalty", _DEFAULTS["presence_penalty"]), + "presence_penalty", + active_rows, + batch_size, + ) + frequency_penalty = _normalize_per_row( + getattr(sampling_params, "frequency_penalty", _DEFAULTS["frequency_penalty"]), + "frequency_penalty", + active_rows, + batch_size, + ) + repetition_penalty = _normalize_per_row( + getattr(sampling_params, "repetition_penalty", _DEFAULTS["repetition_penalty"]), + "repetition_penalty", + active_rows, + batch_size, + ) + seeds = _normalize_seeds(getattr(sampling_params, "seed", None), batch_size) + enable_log_probs = _normalize_output_field( + getattr(sampling_params, "enable_log_probs", False), + "enable_log_probs", + batch_size, + ) + num_logprobs_value = getattr(sampling_params, "num_logprobs", 0) + num_logprobs = _normalize_output_field( + 0 if num_logprobs_value is None else num_logprobs_value, + "num_logprobs", + batch_size, + ) + + row_paths: list[SamplingRowPath] = ["inactive"] * batch_size + active_mask = [row < active_rows for row in range(batch_size)] + greedy_mask = [False] * batch_size + logprob_modes: list[LogProbMode] = ["none"] * batch_size + for row in range(batch_size): + top_p[row] = min(max(float(top_p[row]), 0.0), 1.0) + repetition_penalty[row] = float(repetition_penalty[row]) or 1.0 + + if row >= active_rows: + temperature[row] = 1.0 + top_k[row] = 1 + top_p[row] = 0.0 + enable_log_probs[row] = False + num_logprobs[row] = 0 + continue + + is_greedy = float(temperature[row]) == 0.0 + greedy_mask[row] = is_greedy + if is_greedy: + temperature[row] = 1.0 + top_k[row] = 1 + top_p[row] = 0.0 + else: + temperature[row] = 1.0 / float(temperature[row]) + top_k[row] = _exact_top_k(top_k[row], row=row, max_device_top_k=max_device_top_k) + + enabled = bool(enable_log_probs[row]) + count = int(num_logprobs[row]) + if not enabled: + count = 0 + num_logprobs[row] = 0 + elif count < 0: + raise ValueError(f"sampling_params.num_logprobs[{row}] must be non-negative, got {count}") + logprob_modes[row] = "none" if not enabled else ("sampled_token" if count == 0 else "top_n") + row_paths[row] = "argmax" if is_greedy and allow_force_argmax and not enabled else "topk" + + sampling_path: SamplingPath = "argmax" if all(path == "argmax" for path in row_paths[:active_rows]) else "topk" + return PreparedSamplingParams( + top_k=tuple(int(value) for value in top_k), + top_p=tuple(float(value) for value in top_p), + temperature=tuple(float(value) for value in temperature), + presence_penalty=tuple(float(value) for value in presence_penalty), + frequency_penalty=tuple(float(value) for value in frequency_penalty), + repetition_penalty=tuple(float(value) for value in repetition_penalty), + seeds=tuple(None if value is None else int(value) for value in seeds), + enable_log_probs=tuple(bool(value) for value in enable_log_probs), + num_logprobs=tuple(int(value) for value in num_logprobs), + logprob_modes=tuple(logprob_modes), + greedy_mask=tuple(greedy_mask), + row_paths=tuple(row_paths), + active_mask=tuple(active_mask), + sampling_path=sampling_path, + active_rows=active_rows, + batch_size=batch_size, + max_device_top_k=max_device_top_k, + prompt_tokens=prompt_tokens, + output_tokens=output_tokens, + slot_remap=slot_remap, + ) + + +def format_sampling_params( + sampling_params: SamplingParams, + batch_size: int, + *, + max_device_top_k: int, + allow_force_argmax: bool, + prompt_tokens: Any | None = None, + output_tokens: Any | None = None, + slot_remap: Any | None = None, +) -> PreparedSamplingParams: + """Compatibility spelling for callers that describe this step as formatting.""" + + return prepare_sampling_params( + sampling_params, + batch_size, + max_device_top_k=max_device_top_k, + allow_force_argmax=allow_force_argmax, + prompt_tokens=prompt_tokens, + output_tokens=output_tokens, + slot_remap=slot_remap, + ) + + +def slice_prepared_sampling_params( + prepared: PreparedSamplingParams, + rows: Sequence[int], +) -> PreparedSamplingParams: + """Slice a complete prepared request, including sampling-owned history. + + State tensors and sequences are indexed on their leading request dimension. + Values with one leading row broadcast to the selected rows. Slot-remap + values themselves are preserved; this function only selects which request + rows are assigned to the destination lane. + """ + + if not isinstance(prepared, PreparedSamplingParams): + raise TypeError("prepared must be PreparedSamplingParams") + selected = tuple(int(row) for row in rows) + if not selected: + raise ValueError("prepared sampling rows cannot be empty") + if any(row < 0 or row >= prepared.batch_size for row in selected): + raise ValueError(f"prepared sampling rows must be in [0, {prepared.batch_size})") + + active_mask = tuple(prepared.active_mask[row] for row in selected) + active_rows = sum(active_mask) + if active_rows == 0: + raise ValueError("prepared sampling slice must include at least one active row") + row_paths = tuple(prepared.row_paths[row] for row in selected) + active_paths = tuple(path for active, path in zip(active_mask, row_paths) if active) + sampling_path: SamplingPath = "argmax" if all(path == "argmax" for path in active_paths) else "topk" + + def select_tuple(value: tuple[Any, ...]) -> tuple[Any, ...]: + return tuple(value[row] for row in selected) + + return PreparedSamplingParams( + top_k=select_tuple(prepared.top_k), + top_p=select_tuple(prepared.top_p), + temperature=select_tuple(prepared.temperature), + presence_penalty=select_tuple(prepared.presence_penalty), + frequency_penalty=select_tuple(prepared.frequency_penalty), + repetition_penalty=select_tuple(prepared.repetition_penalty), + seeds=select_tuple(prepared.seeds), + enable_log_probs=select_tuple(prepared.enable_log_probs), + num_logprobs=select_tuple(prepared.num_logprobs), + logprob_modes=select_tuple(prepared.logprob_modes), + greedy_mask=select_tuple(prepared.greedy_mask), + row_paths=row_paths, + active_mask=active_mask, + sampling_path=sampling_path, + active_rows=active_rows, + batch_size=len(selected), + max_device_top_k=prepared.max_device_top_k, + prompt_tokens=_slice_request_state(prepared.prompt_tokens, selected, "prompt_tokens"), + output_tokens=_slice_request_state(prepared.output_tokens, selected, "output_tokens"), + slot_remap=_slice_request_state(prepared.slot_remap, selected, "slot_remap"), + ) + + +def place_prepared_sampling_params( + prepared: PreparedSamplingParams, + slots: Sequence[int], +) -> PreparedSamplingParams: + """Place request-ordered active rows into lane-local destination slots. + + Prefill parameters arrive in request order while device K/P/T, seed, and + penalty state are slot indexed. This conversion preserves inactive safe + defaults and expands prompt/output history to the fixed lane capacity. + """ + + if not isinstance(prepared, PreparedSamplingParams): + raise TypeError("prepared must be PreparedSamplingParams") + sources = tuple(row for row, active in enumerate(prepared.active_mask) if active) + destinations = tuple(int(slot) for slot in slots) + if len(destinations) != len(sources): + raise ValueError(f"expected {len(sources)} destination slots, got {len(destinations)}") + if len(set(destinations)) != len(destinations): + raise ValueError("destination slots must be unique") + if any(slot < 0 or slot >= prepared.batch_size for slot in destinations): + raise ValueError(f"destination slots must be in [0, {prepared.batch_size})") + + def place(values: tuple[Any, ...], default: Any) -> tuple[Any, ...]: + result = [default] * prepared.batch_size + for source, destination in zip(sources, destinations): + result[destination] = values[source] + return tuple(result) + + row_paths = place(prepared.row_paths, "inactive") + active_mask = tuple(path != "inactive" for path in row_paths) + active_paths = tuple(path for path in row_paths if path != "inactive") + sampling_path: SamplingPath = "argmax" if all(path == "argmax" for path in active_paths) else "topk" + return PreparedSamplingParams( + top_k=place(prepared.top_k, 1), + top_p=place(prepared.top_p, 0.0), + temperature=place(prepared.temperature, 1.0), + presence_penalty=place(prepared.presence_penalty, 0.0), + frequency_penalty=place(prepared.frequency_penalty, 0.0), + repetition_penalty=place(prepared.repetition_penalty, 1.0), + seeds=place(prepared.seeds, None), + enable_log_probs=place(prepared.enable_log_probs, False), + num_logprobs=place(prepared.num_logprobs, 0), + logprob_modes=place(prepared.logprob_modes, "none"), + greedy_mask=place(prepared.greedy_mask, False), + row_paths=row_paths, + active_mask=active_mask, + sampling_path=sampling_path, + active_rows=prepared.active_rows, + batch_size=prepared.batch_size, + max_device_top_k=prepared.max_device_top_k, + prompt_tokens=_place_request_state( + prepared.prompt_tokens, + sources=sources, + destinations=destinations, + capacity=prepared.batch_size, + name="prompt_tokens", + ), + output_tokens=_place_request_state( + prepared.output_tokens, + sources=sources, + destinations=destinations, + capacity=prepared.batch_size, + name="output_tokens", + ), + slot_remap=prepared.slot_remap, + ) + + +def slice_sampling_params(sampling_params: SamplingParams, rows: Sequence[int]) -> SamplingParams: + """Return request parameters for ``rows`` without mutating the caller value.""" + + _validate_sampling_value(sampling_params) + selected = tuple(int(row) for row in rows) + if not selected: + raise ValueError("sampling parameter rows cannot be empty") + if any(row < 0 for row in selected): + raise ValueError("sampling parameter rows must be non-negative") + + def slice_value(value: Any, name: str) -> Any: + normalized = _host_value(value) + if not _is_sequence(normalized): + return normalized + values = list(normalized) + if not values: + raise ValueError(f"sampling_params.{name} cannot be empty") + if len(values) == 1: + return [values[0] for _ in selected] + try: + return [values[row] for row in selected] + except IndexError as error: + raise ValueError(f"sampling_params.{name} does not cover rows {selected}") from error + + updates = { + field.name: slice_value(getattr(sampling_params, field.name), field.name) + for field in dataclasses.fields(sampling_params) + } + return dataclasses.replace(sampling_params, **updates) + + +def _validate_policy(batch_size: int, max_device_top_k: int, allow_force_argmax: bool) -> None: + if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + if not isinstance(max_device_top_k, int) or isinstance(max_device_top_k, bool) or max_device_top_k <= 0: + raise ValueError("max_device_top_k must be a positive integer") + if not isinstance(allow_force_argmax, bool): + raise TypeError("allow_force_argmax must be bool") + + +def _validate_sampling_value(sampling_params: Any) -> None: + if not dataclasses.is_dataclass(sampling_params) or isinstance(sampling_params, type): + raise TypeError("sampling_params must be a dataclass instance") + for name in ("temperature", "top_k", "top_p"): + if not hasattr(sampling_params, name): + raise TypeError(f"sampling_params must define {name}") + + +def _host_value(value: Any) -> Any: + if isinstance(value, torch.Tensor): + if value.ndim == 0: + return value.item() + return value.reshape(-1).tolist() + return value + + +def _is_sequence(value: Any) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)) + + +def _as_sequence(value: Any, name: str) -> list[Any]: + value = _host_value(value) + values = list(value) if _is_sequence(value) else [value] + if not values: + raise ValueError(f"sampling_params.{name} cannot be empty") + return values + + +def _normalize_per_row(value: Any, name: str, active_rows: int, batch_size: int) -> list[Any]: + value = _host_value(value) + if not _is_sequence(value): + values = [value] * active_rows + else: + values = list(value) + if not values: + raise ValueError(f"sampling_params.{name} cannot be empty") + if len(values) != 1 and len(values) < active_rows: + raise ValueError( + f"sampling_params.{name} has {len(values)} entries but temperature describes " + f"{active_rows} active rows" + ) + if len(values) > batch_size: + raise ValueError(f"sampling_params.{name} has {len(values)} entries, exceeding batch_size={batch_size}") + return values + [_DEFAULTS[name]] * (batch_size - len(values)) + + +def _normalize_seeds(value: Any, batch_size: int) -> list[int | None]: + value = _host_value(value) + if value is None: + values: list[int | None] = [] + elif _is_sequence(value): + values = list(value) + else: + # Seed is request-owned and is never implicitly broadcast to sibling rows. + values = [value] + if len(values) > batch_size: + raise ValueError(f"sampling_params.seed has {len(values)} entries, exceeding batch_size={batch_size}") + normalized = [None if item is None or int(item) == -1 else int(item) for item in values] + return normalized + [None] * (batch_size - len(normalized)) + + +def _normalize_output_field(value: Any, name: str, batch_size: int) -> list[Any]: + value = _host_value(value) + if not _is_sequence(value): + return [value] * batch_size + values = list(value) + if not values: + raise ValueError(f"sampling_params.{name} cannot be empty") + if len(values) == 1: + return values * batch_size + if len(values) > batch_size: + raise ValueError(f"sampling_params.{name} has {len(values)} entries, exceeding batch_size={batch_size}") + return values + [_DEFAULTS[name]] * (batch_size - len(values)) + + +def _slice_request_state(value: Any, rows: tuple[int, ...], name: str) -> Any: + if value is None: + return None + if isinstance(value, torch.Tensor): + if value.ndim == 0: + return value + selected = (0,) * len(rows) if int(value.shape[0]) == 1 else rows + if max(selected) >= int(value.shape[0]): + raise ValueError(f"{name} does not cover rows {rows}") + indices = torch.tensor(selected, dtype=torch.long, device=value.device) + return value.index_select(0, indices) + if _is_sequence(value): + values = list(value) + if not values: + raise ValueError(f"{name} cannot be empty") + selected = (0,) * len(rows) if len(values) == 1 else rows + try: + sliced = [values[row] for row in selected] + except IndexError as error: + raise ValueError(f"{name} does not cover rows {rows}") from error + return tuple(sliced) if isinstance(value, tuple) else sliced + raise TypeError(f"{name} must be a row-indexed tensor or sequence") + + +def _place_request_state( + value: Any, + *, + sources: tuple[int, ...], + destinations: tuple[int, ...], + capacity: int, + name: str, +) -> Any: + if value is None: + return None + selected = _slice_request_state(value, sources, name) + if isinstance(selected, torch.Tensor): + if selected.ndim == 0: + selected = selected.reshape(1) + fill_value = False if selected.dtype == torch.bool else -1 + placed = torch.full( + (capacity, *selected.shape[1:]), + fill_value, + dtype=selected.dtype, + device=selected.device, + ) + indices = torch.tensor(destinations, dtype=torch.long, device=selected.device) + placed.index_copy_(0, indices, selected) + return placed + values = list(selected) + exemplar = values[0] if values else -1 + if _is_sequence(exemplar): + inactive = tuple(-1 for _ in exemplar) if isinstance(exemplar, tuple) else [-1 for _ in exemplar] + else: + inactive = -1 + placed = [inactive for _ in range(capacity)] + for value_row, destination in zip(values, destinations): + placed[destination] = value_row + if isinstance(selected, tuple): + return tuple(placed) + return placed + + +def _exact_top_k(value: Any, *, row: int, max_device_top_k: int) -> int: + if isinstance(value, bool) or not isinstance(value, Integral): + raise TypeError(f"sampling_params.top_k[{row}] must be an integer, got {value!r}") + top_k = int(value) + if not 1 <= top_k <= max_device_top_k: + raise ValueError( + f"stochastic sampling_params.top_k[{row}]={top_k} is outside the device-supported " + f"range [1, {max_device_top_k}]; route this request to host sampling" + ) + return top_k diff --git a/code/models/common/modules/sampling/penalties_1d.py b/code/models/common/modules/sampling/penalties_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..8d2e25560373d9f32305d2e0d6da68744b7bb285 --- /dev/null +++ b/code/models/common/modules/sampling/penalties_1d.py @@ -0,0 +1,717 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Penalties1D: Presence/frequency/repetition penalty transforms for 1D mesh topologies. + +TTTv2 module — declarative config, lazy buffer allocation, no mutable module state. +Penalty state (PenaltyParams, PenaltyAccumulator) is caller-constructed and passed +as arguments to forward methods. + +See also: models/common/sampling/tt_penalties.py (TTTv1 source) +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any, Optional + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_buffer import LazyBuffer, resolve_lazy_buffer + +# --------------------------------------------------------------------------- +# Caller-constructed penalty state dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class PenaltyParams: + """Per-request penalty constants. Set before the decode loop, read-only during it. + + Caller allocates tensors directly (via LazyBuffer, ttnn.from_torch, or any other means) + and constructs this dataclass. The module does NOT provide a convenience creation method. + """ + + prompt_mask: ttnn.Tensor # [max_batch_size, vocab_per_device], int32, sharded + presence_penalties: ttnn.Tensor # [max_batch_size, 1], bfloat16 + frequency_penalties: ttnn.Tensor # [max_batch_size, 1], bfloat16 + repetition_penalties: ttnn.Tensor # [max_batch_size, 1], bfloat16 + inverse_repetition_penalties: ttnn.Tensor # [max_batch_size, 1], bfloat16 (precomputed 1/rep) + + +@dataclass +class PenaltyAccumulator: + """Per-step accumulator state. Mutated by update_output_tokens() after each sampled token. + + Caller allocates tensors directly and constructs this dataclass. + """ + + output_mask: ttnn.Tensor # [max_batch_size, vocab_per_device], int32, sharded + output_counts: ttnn.Tensor # [max_batch_size, vocab_per_device], int32, sharded + output_counts_gathered: ttnn.Tensor # [max_batch_size, vocab_size], int32, replicated + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +@dataclass +class Penalties1DConfig: + """Declarative config for Penalties1D. + + Buffer fields use the triple-type pattern: ``LazyBuffer | ttnn.Tensor | None``. + + - ``None`` → auto-filled by ``_resolve_penalties1d_config()`` with topology-aware defaults. + - ``LazyBuffer`` → declarative spec, materialized lazily in ``load_device_buffers()``. + - ``ttnn.Tensor`` → pre-allocated device tensor, used directly (power user bypass). + """ + + vocab_size: int # Required. Caller pre-pads to be divisible by num_devices. + mesh_device: Optional[ttnn.MeshDevice] = None # None → GetDefaultDevice() + max_batch_size: int = 32 + sub_core_grids: Any = None # From args.sub_core_grids; passed to ttnn ops via op_kwargs + + # --- Persistent buffer specs (LazyBuffer | ttnn.Tensor | None) --- + # Sharded vocab buffers: [max_batch_size, vocab_size], int32, TILE, sharded across devices + prompt_mask: LazyBuffer | ttnn.Tensor | None = None + output_mask: LazyBuffer | ttnn.Tensor | None = None + output_counts: LazyBuffer | ttnn.Tensor | None = None + # Replicated vocab buffers: [max_batch_size, vocab_size], int32, replicated + output_counts_gathered: LazyBuffer | ttnn.Tensor | None = None + zeros: LazyBuffer | ttnn.Tensor | None = None + # Utility buffers + decode_src: LazyBuffer | ttnn.Tensor | None = None # [max_batch_size, 1], int32, ROW_MAJOR, ones + # BF16 penalty param buffers: [max_batch_size, 1], bfloat16, TILE, replicated + presence_penalties: LazyBuffer | ttnn.Tensor | None = None + frequency_penalties: LazyBuffer | ttnn.Tensor | None = None + repetition_penalties: LazyBuffer | ttnn.Tensor | None = None + inverse_repetition_penalties: LazyBuffer | ttnn.Tensor | None = None + + @staticmethod + def _buf_resolved(buf) -> bool: + if buf is None: + return False + if isinstance(buf, ttnn.Tensor): + return True + return buf.is_resolved() + + def is_resolved(self) -> bool: + return self.mesh_device is not None and all( + self._buf_resolved(getattr(self, f)) + for f in ( + "prompt_mask", + "output_mask", + "output_counts", + "output_counts_gathered", + "zeros", + "decode_src", + "presence_penalties", + "frequency_penalties", + "repetition_penalties", + "inverse_repetition_penalties", + ) + ) + + +# --------------------------------------------------------------------------- +# Module +# --------------------------------------------------------------------------- + + +class Penalties1D(LightweightModule): + """Presence/frequency/repetition penalty transforms for 1D mesh topologies. + + Pure compute pipeline — all penalty state (PenaltyParams, PenaltyAccumulator) is + caller-constructed and passed as arguments to forward methods. + """ + + def __init__(self, vocab_size: int, mesh_device: ttnn.MeshDevice | None = None, **kwargs): + """Happy path — minimal required args, config auto-resolved.""" + super().__init__() + self.config = _resolve_penalties1d_config( + Penalties1DConfig(vocab_size=vocab_size, mesh_device=mesh_device, **kwargs) + ) + self._device_buffers_loaded = False + + @classmethod + def from_config(cls, config: Penalties1DConfig) -> Penalties1D: + """Power path — fully custom config.""" + instance = object.__new__(cls) + super(Penalties1D, instance).__init__() + instance.config = _resolve_penalties1d_config(config) + instance._device_buffers_loaded = False + return instance + + @classmethod + def from_model_args(cls, mesh_device, args) -> Penalties1D: + """Backward compat factory for TTTv1 model args.""" + cluster_shape = mesh_device.shape + if min(cluster_shape) > 1: + raise ValueError( + f"Penalties1D only supports 1D mesh topologies, got shape {cluster_shape}. " + "Use TTPenalties for Galaxy (2D) topologies." + ) + padded_vocab_size = getattr(args, "padded_vocab_size", None) + vocab_size = padded_vocab_size if padded_vocab_size is not None else args.vocab_size + sub_core_grids = getattr(args, "sub_core_grids", None) + return cls(vocab_size=vocab_size, mesh_device=mesh_device, sub_core_grids=sub_core_grids) + + # -- Device buffers (idempotent) ------------------------------------------ + + def load_device_buffers(self): + """Materialize module-owned buffers. Called on first forward; idempotent.""" + if self._device_buffers_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device buffers!" + cfg = self.config + + try: + # Module-owned buffers + self._decode_src = _materialize(cfg.decode_src) + self._zeros = _materialize(cfg.zeros) + + # Derived topology fields + self._cluster_shape = cfg.mesh_device.shape + self._num_devices = max(self._cluster_shape[-1], self._cluster_shape[-2]) + self._op_kwargs = {"sub_core_grids": cfg.sub_core_grids} if cfg.sub_core_grids else {} + self._replicate_mapper = ttnn.ShardTensor2dMesh( + cfg.mesh_device, dims=(None, None), mesh_shape=self._cluster_shape + ) + self._use_low_perf_tilize = cfg.sub_core_grids is not None + + # Slice tensors for scatter → slice (port from tt_penalties.py:117-139) + self._slice_start, self._slice_end = self._build_slice_tensors() + except BaseException as primary: + try: + self.release() + except BaseException as cleanup_error: + _attach_cleanup_failures(primary, (cleanup_error,)) + raise + + self._device_buffers_loaded = True + + def release(self) -> None: + """Release every module-owned buffer idempotently. + + Preallocated ``ttnn.Tensor`` config fields are caller-owned and are not + deallocated here. LazyBuffer fields and the slice tensors constructed + by :meth:`load_device_buffers` are owned by this module. + """ + + failures = [] + for name in ( + "prompt_mask", + "output_mask", + "output_counts", + "output_counts_gathered", + "zeros", + "decode_src", + "presence_penalties", + "frequency_penalties", + "repetition_penalties", + "inverse_repetition_penalties", + ): + specification = getattr(self.config, name, None) + if isinstance(specification, LazyBuffer): + try: + specification.release() + except BaseException as error: + failures.append(error) + + for name in ("_slice_start", "_slice_end"): + tensor = getattr(self, name, None) + if tensor is not None: + try: + ttnn.deallocate(tensor) + except BaseException as error: + failures.append(error) + else: + setattr(self, name, None) + + for name in ("_decode_src", "_zeros"): + specification = getattr(self.config, name.removeprefix("_"), None) + if isinstance(specification, LazyBuffer) and specification._value is None: + setattr(self, name, None) + + self._device_buffers_loaded = False + if failures: + _raise_cleanup_failures(failures) + + # -- Lifecycle methods (per-request setup) --------------------------------- + + def init_prompt_penalties( + self, + params: PenaltyParams, + accum: PenaltyAccumulator, + prompt_tokens: "torch.Tensor", + ) -> None: + """Record prompt token positions into params.prompt_mask via scatter_add. + + Called once per request before the decode loop. + Port of TTPenalties.reset_prompt_tokens (tt_penalties.py:194-212). + """ + # todo)) can we get rid of the torch import here? --> will rethink the boundaries of the module when active development is done on the TTTv1 side + import torch + + self.load_device_buffers() + + prompt_tokens_2d = prompt_tokens.reshape(-1, prompt_tokens.shape[-1]) + prompt_tokens_2d = self._pad_batch_to_max(prompt_tokens_2d, pad_value=-1) + + src_host = (prompt_tokens_2d != -1).to(torch.int32) + idx_host = torch.where(prompt_tokens_2d == -1, torch.zeros_like(prompt_tokens_2d), prompt_tokens_2d) + + prompt_tokens_tt = ttnn.from_torch( + idx_host, + device=self.config.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=self._replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + src_tt = ttnn.from_torch( + src_host, + device=self.config.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=self._replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + self._token_bin_counts_and_mask(new_tokens=prompt_tokens_tt, src=src_tt, mask=params.prompt_mask) + + # -- Forward methods ------------------------------------------------------ + + def decode_forward( + self, + logits: ttnn.Tensor, + params: PenaltyParams, + accum: PenaltyAccumulator, + ) -> ttnn.Tensor: + """Apply presence/frequency/repetition penalties to logits. + + Port of apply_penalties() at tt_penalties.py:32-75. + """ + self.load_device_buffers() + op = self._op_kwargs + + original_shape = logits.shape + logits = ttnn.reshape(logits, (-1, original_shape[-1])) + + # Presence: logits -= typecast(output_mask, bf16) * presence + presence_term = ttnn.multiply( + ttnn.typecast(accum.output_mask, ttnn.bfloat16, **op), params.presence_penalties, **op + ) + presence_term_bf16 = ttnn.typecast(presence_term, ttnn.bfloat16, **op) + presence_term.deallocate() + logits = ttnn.subtract(logits, presence_term_bf16, output_tensor=logits, **op) + presence_term_bf16.deallocate() + + # Frequency: logits -= typecast(output_counts, bf16) * frequency + output_counts_bf16 = ttnn.typecast(accum.output_counts, ttnn.bfloat16, **op) + freq_term = ttnn.multiply(output_counts_bf16, params.frequency_penalties, **op) + output_counts_bf16.deallocate() + freq_term_bf16 = ttnn.typecast(freq_term, ttnn.bfloat16, **op) + freq_term.deallocate() + logits = ttnn.subtract(logits, freq_term_bf16, output_tensor=logits, **op) + freq_term_bf16.deallocate() + + # Repetition: sign-dependent scaling + combined_mask_int32 = ttnn.add(params.prompt_mask, accum.output_mask, **op) + combined_mask = ttnn.typecast(combined_mask_int32, ttnn.bfloat16, **op) + combined_mask_int32.deallocate() + penalties = ttnn.where(combined_mask, params.repetition_penalties, 1.0, **op) + inverse_penalties = ttnn.where(combined_mask, params.inverse_repetition_penalties, 1.0, **op) + combined_mask.deallocate() + + logits_bf16 = ttnn.typecast(logits, ttnn.bfloat16, **op) + logits_gt1 = ttnn.gt(logits_bf16, 0, **op) + logits_bf16.deallocate() + scaling = ttnn.where(logits_gt1, inverse_penalties, penalties, **op) + logits_gt1.deallocate() + penalties.deallocate() + inverse_penalties.deallocate() + logits = ttnn.multiply(logits, scaling, output_tensor=logits, **op) + scaling.deallocate() + + return ttnn.reshape(logits, original_shape) + + def forward(self, logits, params=None, accum=None, **kwargs) -> ttnn.Tensor: + """Dispatcher. If params/accum are None, returns logits unchanged.""" + if params is None or accum is None: + return logits + if "prompt_tokens" in kwargs: + self.init_prompt_penalties(params, accum, kwargs["prompt_tokens"]) + return logits + return self.decode_forward(logits, params, accum) + + # -- Accumulator operations (called AFTER sampling) ----------------------- + + def update_output_tokens(self, accum: PenaltyAccumulator, new_tokens: ttnn.Tensor) -> None: + """Update accum with newly sampled tokens. + + Port of TTPenalties.update_output_tokens (tt_penalties.py:247-264). + Called after each decode step, between decode_forward and the next step. + """ + self.load_device_buffers() + + if (new_tokens.shape[-1] == self.config.max_batch_size and new_tokens.shape[-2] == 1) or ( + new_tokens.shape[-2] == self.config.max_batch_size and new_tokens.shape[-1] == 1 + ): + # Standard decode: [..., 1, B] or row-sharded: [..., B, 1] → [B, 1] + new_tokens = ttnn.reshape(new_tokens, [self.config.max_batch_size, 1], **self._op_kwargs) + src = self._decode_src + else: + # todo)) can we get rid of the torch import here? --> will rethink the boundaries of the module when active development is done on the TTTv1 side + import torch + + src = ttnn.from_torch( + torch.ones(self.config.max_batch_size, new_tokens.shape[-1], dtype=torch.int32), + device=self.config.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=self._replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + self._token_bin_counts_and_mask( + new_tokens=new_tokens, + counts=accum.output_counts_gathered, + src=src, + counts_sliced=accum.output_counts, + mask=accum.output_mask, + ) + + def reset_output_tokens(self, accum: PenaltyAccumulator, tokens: "torch.Tensor | None" = None) -> None: + """Zero out accumulator buffers. Optionally re-initialize from provided tokens. + + Port of TTPenalties.reset_output_tokens (tt_penalties.py:214-245). + """ + self.load_device_buffers() + op = self._op_kwargs + + accum.output_mask = ttnn.mul(accum.output_mask, 0, output_tensor=accum.output_mask, **op) + accum.output_counts = ttnn.mul(accum.output_counts, 0, output_tensor=accum.output_counts, **op) + accum.output_counts_gathered = ttnn.mul( + accum.output_counts_gathered, 0, output_tensor=accum.output_counts_gathered, **op + ) + + if tokens is not None: + # todo)) can we get rid of the torch import here? --> will rethink the boundaries of the module when active development is done on the TTTv1 side + import torch + + tokens_2d = tokens.reshape(-1, tokens.shape[-1]) + tokens_2d = self._pad_batch_to_max(tokens_2d, pad_value=-1) + src_host = (tokens_2d != -1).to(torch.int32) + idx_host = torch.where(tokens_2d == -1, torch.zeros_like(tokens_2d), tokens_2d) + + tokens_tt = ttnn.from_torch( + idx_host, + device=self.config.mesh_device, + dtype=ttnn.uint32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=self._replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + src_tt = ttnn.from_torch( + src_host, + device=self.config.mesh_device, + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + mesh_mapper=self._replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + self._token_bin_counts_and_mask( + new_tokens=tokens_tt, + counts=accum.output_counts_gathered, + src=src_tt, + counts_sliced=accum.output_counts, + mask=accum.output_mask, + ) + + # -- Private helpers ------------------------------------------------------ + + def _build_slice_tensors(self): + """Derive slice_start/slice_end from vocab_size and num_devices. + + Port of tt_penalties.py:117-139. + """ + # todo)) can we get rid of the torch import here? --> will rethink the boundaries of the module when active development is done on the TTTv1 side + import torch + + cfg = self.config + vocab_per_dev = cfg.vocab_size // self._num_devices + d = torch.arange(self._num_devices, dtype=torch.int32) + + if self._cluster_shape[-1] == self._num_devices: + shard_dims_slice = (None, 0) + else: + shard_dims_slice = (0, None) + + start_1d = torch.empty(2 * self._num_devices, dtype=torch.int32) + start_1d[0::2] = 0 + start_1d[1::2] = d * vocab_per_dev + + padded_batch_size = ((cfg.max_batch_size + ttnn.TILE_SIZE - 1) // ttnn.TILE_SIZE) * ttnn.TILE_SIZE + end_1d = torch.empty(2 * self._num_devices, dtype=torch.int32) + end_1d[0::2] = padded_batch_size + end_1d[1::2] = (d + 1) * vocab_per_dev + + slice_start = None + try: + slice_start = ttnn.from_torch( + start_1d, + device=cfg.mesh_device, + mesh_mapper=ttnn.ShardTensor2dMesh( + cfg.mesh_device, + dims=shard_dims_slice, + mesh_shape=self._cluster_shape, + ), + ) + slice_end = ttnn.from_torch( + end_1d, + device=cfg.mesh_device, + mesh_mapper=ttnn.ShardTensor2dMesh( + cfg.mesh_device, + dims=shard_dims_slice, + mesh_shape=self._cluster_shape, + ), + ) + except BaseException as primary: + if slice_start is not None: + try: + ttnn.deallocate(slice_start) + except BaseException as cleanup_error: + _attach_cleanup_failures(primary, (cleanup_error,)) + raise + return slice_start, slice_end + + def _pad_batch_to_max(self, tokens_2d: "torch.Tensor", pad_value: int) -> "torch.Tensor": + """Pad/truncate first dim to max_batch_size.""" + # todo)) can we get rid of the torch import here? --> will rethink the boundaries of the module when active development is done on the TTTv1 side + import torch + + if tokens_2d.dim() != 2: + raise ValueError(f"Expected 2D tensor [B, S], got {tokens_2d.shape}") + B, S = tokens_2d.shape + if B < self.config.max_batch_size: + pad = torch.full((self.config.max_batch_size - B, S), pad_value, dtype=tokens_2d.dtype) + return torch.cat([tokens_2d, pad], dim=0) + if B > self.config.max_batch_size: + return tokens_2d[: self.config.max_batch_size] + return tokens_2d + + def _token_bin_counts_and_mask(self, new_tokens, src, counts=None, mask=None, counts_sliced=None): + """Scatter tokens into histogram, slice per-device, compute mask. + + Port of TTPenalties.token_bin_counts_and_mask (tt_penalties.py:266-289). + """ + op = self._op_kwargs + counts_new = ttnn.scatter_add(self._zeros, 1, new_tokens, src, **op) + + new_tokens.deallocate() + counts_new = self._tilize_counts(counts_new) + if counts is not None: + counts = ttnn.add(counts, counts_new, output_tensor=counts, **op) + else: + counts = counts_new + + if counts.shape[-2] % ttnn.TILE_SIZE: + # Dynamic tiled slice requires tile-aligned logical input/output + # heights. Pure padded views share the original buffers, so the + # caller-owned logical-B state remains authoritative. + counts_padded = ttnn.reshape( + counts, + counts.padded_shape, + counts.padded_shape, + skip_padding_fill=True, + ) + counts_sliced_padded = None + if counts_sliced is not None: + counts_sliced_padded = ttnn.reshape( + counts_sliced, + counts_sliced.padded_shape, + counts_sliced.padded_shape, + skip_padding_fill=True, + ) + sliced = ttnn.slice( + counts_padded, + self._slice_start, + self._slice_end, + output_tensor=counts_sliced_padded, + slice_dim=1, + num_devices=self._num_devices, + **op, + ) + if counts_sliced is None: + logical_shape = list(sliced.padded_shape) + logical_shape[-2] = self.config.max_batch_size + counts_sliced = ttnn.reshape( + sliced, + logical_shape, + sliced.padded_shape, + skip_padding_fill=True, + ) + else: + # Preserve the established batch-32 slice exactly. + counts_sliced = ttnn.slice( + counts, + self._slice_start, + self._slice_end, + output_tensor=counts_sliced, + slice_dim=1, + num_devices=self._num_devices, + **op, + ) + + mask = ttnn.gt(counts_sliced, 0, output_tensor=mask, **op) + return counts, mask + + def _tilize_counts(self, counts): + """Tilize one histogram while preserving non-tile batch heights.""" + + if counts.padded_shape[-2] % ttnn.TILE_SIZE == 0: + return ttnn.tilize(counts, **self._op_kwargs, use_low_perf=self._use_low_perf_tilize) + return ttnn.to_layout(counts, ttnn.TILE_LAYOUT, **self._op_kwargs) + + +# --------------------------------------------------------------------------- +# Config resolution +# --------------------------------------------------------------------------- + + +def _resolve_penalties1d_config(config: Penalties1DConfig) -> Penalties1DConfig: + """Fill None fields in config with topology-aware defaults. + + Power users who set fields explicitly will NOT have them overwritten. + Mirrors the ``_resolve_mlp1d_config()`` pattern. + """ + import torch # lazy import — only needed for source tensor construction + + to_set: dict = {} + + # Phase 1: Device + mesh_device = config.mesh_device or ttnn.GetDefaultDevice() + to_set["mesh_device"] = mesh_device + + # Phase 2: Topology → shard_dims (port from tt_penalties.py:97-103) + cluster_shape = mesh_device.shape + num_devices = max(cluster_shape[-1], cluster_shape[-2]) + if cluster_shape[-1] == num_devices: + shard_dims = (None, 1) + else: + shard_dims = (1, None) + + B = config.max_batch_size + V = config.vocab_size + + # Build mesh mappers + shard_mapper = ttnn.ShardTensor2dMesh(mesh_device, dims=shard_dims, mesh_shape=cluster_shape) + replicate_mapper = None # None → replicate_tensor_to_mesh_mapper in LazyBuffer + + def _resolve_buf(field_val, defaults, source_factory): + """Resolve a single buffer field: None → LazyBuffer, LazyBuffer → fill, ttnn.Tensor → passthrough.""" + if field_val is None: + return LazyBuffer(source=source_factory(), **defaults) + if isinstance(field_val, ttnn.Tensor): + return field_val + return resolve_lazy_buffer(field_val, **defaults) + + # Phase 3: Sharded vocab buffers — [B, V], int32, TILE, sharded + sharded_vocab_defaults = dict( + dtype=ttnn.int32, + layout=ttnn.TILE_LAYOUT, + device=mesh_device, + mesh_mapper=shard_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + zeros_BV = lambda: torch.zeros(B, V, dtype=torch.int32) + to_set["prompt_mask"] = _resolve_buf(config.prompt_mask, sharded_vocab_defaults, zeros_BV) + to_set["output_mask"] = _resolve_buf(config.output_mask, sharded_vocab_defaults, zeros_BV) + to_set["output_counts"] = _resolve_buf(config.output_counts, sharded_vocab_defaults, zeros_BV) + + # Phase 4: Replicated vocab buffers — [B, V], int32, replicated + replicated_vocab_defaults = dict( + dtype=ttnn.int32, + layout=ttnn.TILE_LAYOUT, + device=mesh_device, + mesh_mapper=replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + to_set["output_counts_gathered"] = _resolve_buf(config.output_counts_gathered, replicated_vocab_defaults, zeros_BV) + + replicated_vocab_rm_defaults = dict( + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + device=mesh_device, + mesh_mapper=replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + to_set["zeros"] = _resolve_buf(config.zeros, replicated_vocab_rm_defaults, zeros_BV) + + # Phase 5: Utility buffers + decode_src_defaults = dict( + dtype=ttnn.int32, + layout=ttnn.ROW_MAJOR_LAYOUT, + device=mesh_device, + mesh_mapper=replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + to_set["decode_src"] = _resolve_buf( + config.decode_src, decode_src_defaults, lambda: torch.ones(B, 1, dtype=torch.int32) + ) + + # Phase 6: BF16 penalty param buffers — [B, 1], bfloat16, TILE, replicated + bf16_param_defaults = dict( + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + device=mesh_device, + mesh_mapper=replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + zeros_B1 = lambda: torch.zeros(B, 1, dtype=torch.float32) + for field_name in ( + "presence_penalties", + "frequency_penalties", + "repetition_penalties", + "inverse_repetition_penalties", + ): + to_set[field_name] = _resolve_buf(getattr(config, field_name), bf16_param_defaults, zeros_B1) + + resolved = replace(config, **to_set) + assert resolved.is_resolved(), "Config not fully resolved after _resolve_penalties1d_config" + return resolved + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _materialize(buf): + """Materialize a buffer field: ttnn.Tensor passthrough, LazyBuffer → get_device_buffer().""" + if isinstance(buf, ttnn.Tensor): + return buf + return buf.get_device_buffer() + + +def _attach_cleanup_failures(primary, failures): + if not failures: + return + previous = tuple(getattr(primary, "cleanup_failures", ())) + primary.cleanup_failures = previous + tuple(failures) + add_note = getattr(primary, "add_note", None) + if callable(add_note): + add_note(f"cleanup also encountered {len(failures)} failure(s)") + + +def _raise_cleanup_failures(failures): + primary = failures[0] + if len(failures) > 1: + previous = tuple(getattr(primary, "cleanup_failures", ())) + primary.cleanup_failures = previous + tuple(failures[1:]) + add_note = getattr(primary, "add_note", None) + if callable(add_note): + add_note(f"cleanup also encountered {len(failures) - 1} additional failure(s)") + raise primary diff --git a/code/models/common/modules/sampling/sampling_1d.py b/code/models/common/modules/sampling/sampling_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..6c4178d7a270d7fac6eeb26a81938afecc7ee903 --- /dev/null +++ b/code/models/common/modules/sampling/sampling_1d.py @@ -0,0 +1,886 @@ +# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +""" +Sampling1D: Top-k/top-p/temperature sampling for 1D mesh topologies. + +TTTv2 module — declarative config, lazy buffer allocation, k/p/temp as per-call args. +No mutable sampling state stored on the module. + +See also: models/common/sampling/tt_sampling.py (TTTv1 source) +""" + +from __future__ import annotations + +import inspect +import sys +from dataclasses import dataclass, replace +from typing import Any, Optional + +from loguru import logger + +import ttnn +from models.common.lightweightmodule import LightweightModule +from models.common.modules.lazy_buffer import LazyBuffer, resolve_lazy_buffer +from models.common.modules.tt_ccl import get_tt_ccl +from models.common.sampling.vocab_padding import ( + build_invalid_vocab_mask, + build_tail_invalid_vocab_mask, + get_vocab_shard_dims, +) + +# --------------------------------------------------------------------------- +# Power-of-2 helpers (local copies; keep this TTTv2 module self-contained) +# --------------------------------------------------------------------------- + + +def _is_power_of_2(n: int) -> bool: + return n > 0 and (n & (n - 1)) == 0 + + +def _upper_power_of_2(n: int) -> int: + if n <= 1: + return 1 + return 1 << (n - 1).bit_length() + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +@dataclass +class Sampling1DConfig: + """Declarative config for Sampling1D. + + Buffer fields use the triple-type pattern: ``LazyBuffer | ttnn.Tensor | None``. + k/p/temp are NOT stored here — they are per-call forward args. + """ + + vocab_size: int # Required. Padded logits width; caller pre-pads to be divisible by num_devices. + valid_vocab_size: Optional[int] = None # Real token vocabulary size. Defaults to vocab_size. + mesh_device: Optional[ttnn.MeshDevice] = None # None → GetDefaultDevice() + tt_ccl: Any = None # None → get_tt_ccl(mesh_device) if multi-device + max_batch_size: int = 32 + max_top_k: int = 32 + sub_core_grids: Any = None + sub_core_grid_topk: Any = None + start_core: Optional[ttnn.CoreCoord] = None # None → CoreCoord(0,0) + num_gather_links: int = 1 + sampling_memory_config: Optional[ttnn.MemoryConfig] = None # None → DRAM_MEMORY_CONFIG + allow_force_argmax: bool = False + num_argmax_gather_links: Optional[int] = None # None → same as num_gather_links + ag_topology: Optional[ttnn.Topology] = None # None → Topology.Linear + argmax_chunks_per_sync: int = 10 + argmax_num_workers_per_link: int = 1 + # Pad each per-device logit shard up to the next power of 2 before ttnn.topk. Big device-perf + # win for non-power-of-2 vocab on the multi-device path (TTTv1's pad_logits_to_power_of_2). + # Strict TTTv1 parity: only the multi-device path is padded — the 1×1 multi_step split path + # is left unpadded. + pad_to_power_of_2: bool = False + + # --- Persistent buffer specs (LazyBuffer | ttnn.Tensor | None) --- + # Static index buffers (computed from vocab_size + num_devices, never mutated) + index_offsets: LazyBuffer | ttnn.Tensor | None = None # [1,1,32,max_top_k*num_devices], int32, TILE + invalid_vocab_mask: LazyBuffer | ttnn.Tensor | None = None # Full fallback mask, sharded like logits + invalid_vocab_tail_mask: LazyBuffer | ttnn.Tensor | None = None # Compact [1,1,32,tail] local mask + invalid_vocab_tail_width: int = 0 + # Seed/ID buffers (seeds mutable via LazyBuffer.update(), user_ids static) + seeds: LazyBuffer | ttnn.Tensor | None = None # [32], uint32, ROW_MAJOR + user_ids: LazyBuffer | ttnn.Tensor | None = None # [32], uint32, ROW_MAJOR + + @staticmethod + def _buf_resolved(buf) -> bool: + if buf is None: + return False + if isinstance(buf, ttnn.Tensor): + return True + return buf.is_resolved() + + def is_resolved(self) -> bool: + if self.mesh_device is None: + return False + if self.mesh_device.get_num_devices() > 1 and self.tt_ccl is None: + return False + required_buffers = ["index_offsets", "seeds", "user_ids"] + if not all(self._buf_resolved(getattr(self, f)) for f in required_buffers): + return False + if self.valid_vocab_size is not None and self.valid_vocab_size < self.vocab_size: + return self._buf_resolved(self.invalid_vocab_mask) or self._buf_resolved(self.invalid_vocab_tail_mask) + return True + + +# --------------------------------------------------------------------------- +# Module +# --------------------------------------------------------------------------- + + +class Sampling1D(LightweightModule): + """Top-k/top-p/temperature sampling for 1D mesh topologies. + + k/p/temp are per-call forward args — NOT stored on the module. This eliminates + mutable sampling state. + """ + + def __init__(self, vocab_size: int, mesh_device: ttnn.MeshDevice | None = None, **kwargs): + """Happy path — minimal required args.""" + super().__init__() + self.config = _resolve_sampling1d_config( + Sampling1DConfig(vocab_size=vocab_size, mesh_device=mesh_device, **kwargs) + ) + self._device_buffers_loaded = False + self._bind_strategy() + + @classmethod + def from_config(cls, config: Sampling1DConfig) -> Sampling1D: + """Power path — fully custom config.""" + instance = object.__new__(cls) + super(Sampling1D, instance).__init__() + instance.config = _resolve_sampling1d_config(config) + instance._device_buffers_loaded = False + instance._bind_strategy() + return instance + + # -- Strategy binding (TTTv2: no if-else in forward) ---------------------- + + def _bind_strategy(self): + """Bind self._topk to the correct strategy based on mesh topology.""" + cluster_shape = self.config.mesh_device.shape + self._multi_step_reduction = list(cluster_shape) == [1, 1] + if self._multi_step_reduction: + self._topk = self._topk_single_device + else: + self._topk = self._topk_multi_device + + # Argmax strategy: single vs multi-device + num_devices = self.config.mesh_device.get_num_devices() + if num_devices > 1: + self._pre_argmax_gather = self._argmax_all_gather + else: + self._pre_argmax_gather = self._argmax_noop + + # Memory config strategy for top-k post-processing + cfg = self.config + if cfg.sampling_memory_config is not None and cfg.sampling_memory_config != ttnn.DRAM_MEMORY_CONFIG: + self._prepare_topk_memory = self._topk_memory_sharded_roundtrip + else: + self._prepare_topk_memory = self._topk_memory_noop + + # CCL introspection (port from TTSampling.__init__ lines 77-91) + self._line_all_gather = getattr(self.config.tt_ccl, "line_all_gather", None) if self.config.tt_ccl else None + self._line_all_gather_supports_buffer_key = False + if callable(self._line_all_gather): + try: + sig = inspect.signature(self._line_all_gather) + params = sig.parameters + self._line_all_gather_supports_buffer_key = "buffer_key" in params or any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values() + ) + except (TypeError, ValueError): + logger.warning("Unable to inspect line_all_gather signature") + + # -- Device buffers (idempotent) ------------------------------------------ + + def load_device_buffers(self): + """Materialize all LazyBuffer fields from resolved config.""" + if self._device_buffers_loaded: + return + + assert self.config.is_resolved(), "config must be resolved before loading device buffers!" + cfg = self.config + + try: + self._index_offsets = _materialize(cfg.index_offsets) + self._invalid_vocab_mask = ( + _materialize(cfg.invalid_vocab_mask) if cfg.invalid_vocab_mask is not None else None + ) + self._invalid_vocab_tail_mask = ( + _materialize(cfg.invalid_vocab_tail_mask) if cfg.invalid_vocab_tail_mask is not None else None + ) + self._invalid_vocab_tail_width = cfg.invalid_vocab_tail_width + self._seeds = _materialize(cfg.seeds) + self._user_ids = _materialize(cfg.user_ids) + from models.common.utils import LogProbsCalculator # lazy: transitively imports torch + + self._log_probs_calculator = LogProbsCalculator(cfg.mesh_device, cfg.sub_core_grids, cfg.tt_ccl) + + # Pre-compute static sub_core_grids for ttnn.sampling() + self._sampling_sub_core_grids = ( + ttnn.num_cores_to_corerangeset_in_subcoregrids( + cfg.start_core, cfg.max_batch_size, cfg.sub_core_grids, row_wise=True + ) + if cfg.sub_core_grids is not None + else None + ) + except BaseException as primary: + try: + self.release() + except BaseException as cleanup_error: + failures = (cleanup_error,) + tuple(getattr(cleanup_error, "cleanup_failures", ())) + _attach_cleanup_failures(primary, failures) + raise + + self._device_buffers_loaded = True + + def release(self) -> None: + """Release model-owned buffers idempotently and permit rematerialization.""" + failures = [] + calculator = getattr(self, "_log_probs_calculator", None) + if calculator is not None: + try: + calculator.release() + except BaseException as error: + failures.append(error) + else: + self._log_probs_calculator = None + + for name in ( + "index_offsets", + "invalid_vocab_mask", + "invalid_vocab_tail_mask", + "seeds", + "user_ids", + ): + specification = getattr(self.config, name, None) + if isinstance(specification, LazyBuffer): + try: + specification.release() + except BaseException as error: + failures.append(error) + if not isinstance(specification, LazyBuffer) or specification._value is None: + setattr(self, f"_{name}", None) + + self._sampling_sub_core_grids = None + self._device_buffers_loaded = False + if failures: + _raise_cleanup_failures(failures) + + # -- Forward methods ------------------------------------------------------ + + def decode_forward( + self, + logits: ttnn.Tensor, + *, + k: ttnn.Tensor | None = None, + p: ttnn.Tensor | None = None, + temp: ttnn.Tensor | None = None, + seeds: ttnn.Tensor | None = None, + tt_out_tok: ttnn.Tensor | None = None, + enable_log_probs: bool | list[bool] = False, + ): + """Sample tokens from logits. + + Args: + logits: Input logits tensor (sharded across devices) + k, p, temp: Per-call sampling parameters. All None + allow_force_argmax → argmax path. + All provided → top-k sampling path. + seeds: Optional per-call seed override. If None, uses config seeds buffer. + tt_out_tok: Optional output tensor to write results to. + enable_log_probs: Per-call logprobs toggle (bool, or per-user list of bool — if any + user is enabled the whole batch computes logprobs). Refreshed every call, so no + mutable logprobs state is stored on the module. Only the top-k path emits logprobs + (sampled-token logprob); the argmax path never does (see ``_sample_argmax``). The + calculator additionally requires a multi-device shard (T3K 1×8) — it returns ``None`` + on 1×1/1×2 even when enabled. + + Returns: + (token_ids, log_probs_or_none) + """ + self.load_device_buffers() + cfg = self.config + + # Per-call logprobs mode — refreshed each forward from the arg, never persisted on the + # module (TTTv2 principle: no mutable sampling state). Mirrors main's reset_params path. + self._log_probs_calculator.set_log_probs_mode(enable_log_probs) + + # Route: argmax or top-k + if k is None and p is None and temp is None: + if cfg.allow_force_argmax: + return self._sample_argmax(logits, tt_out_tok) + else: + raise ValueError("k/p/temp are all None but allow_force_argmax is False") + if k is None or p is None or temp is None: + raise ValueError("k, p, temp must all be provided, or all be None (for argmax)") + + return self._sample_topk(logits, k, p, temp, seeds, tt_out_tok) + + def forward(self, logits, **kwargs): + """Dispatcher.""" + return self.decode_forward(logits, **kwargs) + + # -- Argmax path (port of the argmax branch of TTSampling.forward) -------- + + def _sample_argmax(self, logits, tt_out_tok): + slice_valid_vocab = self._can_slice_valid_vocab_for_argmax() + if not slice_valid_vocab: + logits = self._mask_invalid_vocab_logits(logits) + logits = self._pre_argmax_gather(logits) + if slice_valid_vocab: + logits = self._slice_valid_vocab_for_argmax(logits) + x_untilized = ttnn.untilize(logits, use_multicore=True) + tt_out_tok = ttnn.argmax( + x_untilized, + dim=-1, + output_tensor=tt_out_tok, + keepdim=False, + ) + # Argmax path never emits logprobs (main's contract: force-argmax is disabled whenever + # logprobs are requested). Return None unconditionally — do not call the calculator. + return tt_out_tok, None + + def _get_argmax_all_gather_config(self, cluster_axis): + """Clamp the tuned all-gather config to what the actual submesh supports. + + Port of main's ``_get_force_argmax_all_gather_config`` (#44246): bound num_links to the + links available on the submesh, and force Linear topology below 8 devices (Ring routing + like D0→D12 only wraps cleanly on T3K-class 8-device groups). + """ + cfg = self.config + num_links = cfg.num_argmax_gather_links + if hasattr(cfg.tt_ccl, "get_num_links"): + num_links = min(num_links, cfg.tt_ccl.get_num_links(cluster_axis)) + topology = cfg.ag_topology + if cfg.mesh_device.get_num_devices() < 8: + topology = ttnn.Topology.Linear + return max(1, num_links), topology + + def _argmax_all_gather(self, logits): + """Multi-device: all-gather logits before argmax. + + On ring-capable meshes (e.g. T3K 1×8) use Ring topology with no barrier semaphore to match the + model's logits gather and avoid trace-capture issues seen with some barrier-based configurations. + For other meshes, fall back to the clamped Linear+barrier path. + """ + cfg = self.config + cluster_axis = None if 1 in cfg.mesh_device.shape else 1 + num_links, topology = self._get_argmax_all_gather_config(cluster_axis) + kwargs = {} + if cluster_axis is not None: + kwargs["cluster_axis"] = cluster_axis + return ttnn.experimental.all_gather_async( + logits, + persistent_output_buffer=None, + dim=3, + multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(cluster_axis), + num_links=num_links, + memory_config=logits.memory_config(), + topology=topology, + barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(cluster_axis), + chunks_per_sync=cfg.argmax_chunks_per_sync, + num_workers_per_link=cfg.argmax_num_workers_per_link, + num_buffers_per_channel=2, + **kwargs, + ) + + @staticmethod + def _argmax_noop(logits): + """Single-device: no gather needed.""" + return logits + + # -- Top-k memory strategies (bound at init, no if-else in forward) ------- + + def _topk_memory_sharded_roundtrip(self, topk_values, topk_indices_int32): + """Non-DRAM sampling_memory_config: round-trip through sharded memory.""" + cfg = self.config + topk_values_sharded = ttnn.to_memory_config( + topk_values, memory_config=cfg.sampling_memory_config, dtype=ttnn.bfloat16 + ) + topk_values = ttnn.to_memory_config(topk_values_sharded, memory_config=ttnn.DRAM_MEMORY_CONFIG) + ttnn.deallocate(topk_values_sharded) + topk_indices_int32 = ttnn.to_memory_config(topk_indices_int32, cfg.sampling_memory_config) + return topk_values, topk_indices_int32 + + @staticmethod + def _topk_memory_noop(topk_values, topk_indices_int32): + """DRAM memory config: no extra round-trip needed.""" + return topk_values, topk_indices_int32 + + # -- Top-k sampling (port of the top-k branch of TTSampling.forward) ------ + + def _sample_topk(self, logits, k, p, temp, seeds, tt_out_tok): + cfg = self.config + + x_bf16 = ttnn.typecast(logits, dtype=ttnn.bfloat16, sub_core_grids=cfg.sub_core_grids) + x_bf16 = self._mask_invalid_vocab_logits(x_bf16) + + # Strategy-dispatched top-k + topk_values, topk_indices = self._topk(x_bf16) + + # Convert indices to int32 + topk_indices_int32 = ttnn.typecast(topk_indices, dtype=ttnn.int32, sub_core_grids=cfg.sub_core_grids) + + topk_values, topk_indices_int32 = self._prepare_topk_memory(topk_values, topk_indices_int32) + + # Add device offsets for global vocabulary indices + index_offsets, sliced_offsets = self._slice_user_rows( + self._index_offsets, int(topk_indices_int32.shape[2]), cfg.sampling_memory_config + ) + topk_global_indices = ttnn.add( + index_offsets, + topk_indices_int32, + dtype=ttnn.int32, + memory_config=cfg.sampling_memory_config, + sub_core_grids=cfg.sub_core_grids, + ) + ttnn.deallocate(topk_indices_int32) + if sliced_offsets: + ttnn.deallocate(index_offsets) + + # Use distinct names so we can free the interleaved intermediate after untilize + topk_global_indices_interleaved = ttnn.to_memory_config(topk_global_indices, ttnn.DRAM_MEMORY_CONFIG) + topk_global_indices = ttnn.untilize( + topk_global_indices_interleaved, use_multicore=True, sub_core_grids=cfg.sub_core_grids + ) + ttnn.deallocate(topk_global_indices_interleaved) + + # Seed the RNG + seeds_tensor = seeds if seeds is not None else self._seeds + ttnn.manual_seed( + seeds=seeds_tensor, + user_ids=self._user_ids, + sub_core_grids=self._sampling_sub_core_grids, + ) + + # Sample + tt_out_tok = ttnn.sampling( + topk_values, + topk_global_indices, + k=k, + p=p, + temp=temp, + sub_core_grids=self._sampling_sub_core_grids, + output_tensor=tt_out_tok, + ) + + ttnn.deallocate(topk_values) + ttnn.deallocate(topk_global_indices) + + # Logprobs (old path: single sampled-token logprob). Gated on enable_log_probs so the + # disabled case incurs zero extra ops. The calculator itself returns None unless the mesh + # is multi-device with num_devices ∈ {8, 32} (T3K 1×8). + if self._log_probs_calculator.enable_log_probs: + log_probs = self._log_probs_calculator.calculate_log_probs(logits, tt_out_tok) + else: + log_probs = None + return tt_out_tok, log_probs + + def _mask_invalid_vocab_logits(self, logits): + if self._invalid_vocab_tail_mask is not None: + return self._mask_invalid_vocab_tail_logits(logits) + if self._invalid_vocab_mask is None: + return logits + return ttnn.add( + logits, + self._invalid_vocab_mask, + memory_config=logits.memory_config(), + sub_core_grids=self.config.sub_core_grids, + ) + + def _mask_invalid_vocab_tail_logits(self, logits): + cfg = self.config + tail_width = self._invalid_vocab_tail_width + local_width = logits.shape[-1] + valid_width = local_width - tail_width + if tail_width <= 0 or valid_width < 0: + return self._mask_invalid_vocab_logits_fallback(logits) + if valid_width == 0: + return ttnn.add( + logits, + self._invalid_vocab_tail_mask, + memory_config=logits.memory_config(), + sub_core_grids=cfg.sub_core_grids, + ) + + valid_logits = ttnn.slice( + logits, + [0, 0, 0, 0], + [logits.shape[0], logits.shape[1], logits.shape[2], valid_width], + memory_config=logits.memory_config(), + sub_core_grids=cfg.sub_core_grids, + ) + tail_logits = ttnn.slice( + logits, + [0, 0, 0, valid_width], + [logits.shape[0], logits.shape[1], logits.shape[2], local_width], + memory_config=logits.memory_config(), + sub_core_grids=cfg.sub_core_grids, + ) + masked_tail_logits = ttnn.add( + tail_logits, + self._invalid_vocab_tail_mask, + memory_config=logits.memory_config(), + sub_core_grids=cfg.sub_core_grids, + ) + masked_logits = ttnn.concat( + [valid_logits, masked_tail_logits], + dim=3, + memory_config=logits.memory_config(), + sub_core_grids=cfg.sub_core_grids, + ) + ttnn.deallocate(valid_logits) + ttnn.deallocate(tail_logits) + ttnn.deallocate(masked_tail_logits) + return masked_logits + + def _mask_invalid_vocab_logits_fallback(self, logits): + if self._invalid_vocab_mask is None: + return logits + return ttnn.add( + logits, + self._invalid_vocab_mask, + memory_config=logits.memory_config(), + sub_core_grids=self.config.sub_core_grids, + ) + + def _can_slice_valid_vocab_for_argmax(self): + cfg = self.config + return cfg.valid_vocab_size < cfg.vocab_size and cfg.valid_vocab_size % ttnn.TILE_SIZE == 0 + + def _slice_valid_vocab_for_argmax(self, logits): + cfg = self.config + if not self._can_slice_valid_vocab_for_argmax() or logits.shape[-1] != cfg.vocab_size: + return logits + return ttnn.slice( + logits, + [0, 0, 0, 0], + [logits.shape[0], logits.shape[1], logits.shape[2], cfg.valid_vocab_size], + memory_config=logits.memory_config(), + sub_core_grids=cfg.sub_core_grids, + ) + + def _slice_user_rows(self, tensor, active_batch, memory_config): + if int(tensor.shape[2]) == active_batch: + return tensor, False + return ( + ttnn.slice( + tensor, + [0, 0, 0, 0], + [tensor.shape[0], tensor.shape[1], active_batch, tensor.shape[3]], + memory_config=memory_config, + sub_core_grids=self.config.sub_core_grids, + ), + True, + ) + + # -- Top-k strategies (bound at init, no if-else in forward) -------------- + + def _topk_single_device(self, x_bf16): + """Split vocab in half → two topk → concat. Port of tt_sampling.py:346-371.""" + cfg = self.config + x_list = ttnn.split(x_bf16, x_bf16.shape[-1] // 2, dim=3) + + values_parts = [] + indices_parts = [] + for i in range(len(x_list)): + vals, idxs = ttnn.topk( + x_list[i], + k=cfg.max_top_k, + dim=-1, + sub_core_grids=cfg.sub_core_grid_topk, + ) + values_parts.append(vals) + indices_parts.append(idxs) + x_list[i].deallocate() + + gathered_values = ttnn.concat(values_parts, dim=3) + gathered_indices = ttnn.concat(indices_parts, dim=3) + + for v, i in zip(values_parts, indices_parts): + ttnn.deallocate(v) + ttnn.deallocate(i) + + return gathered_values, gathered_indices + + def _topk_multi_device(self, x_bf16): + """Local topk → all_gather across devices. Port of tt_sampling.py:372-421.""" + cfg = self.config + cluster_shape = cfg.mesh_device.shape + + # Pad the per-device shard up to the next power of 2 so ttnn.topk hits its fast path. + # Padded entries get -inf so they are never selected. Mirrors the padding in TTSampling.forward. + if cfg.pad_to_power_of_2 and not _is_power_of_2(x_bf16.shape[-1]): + padded_width = _upper_power_of_2(x_bf16.shape[-1]) + x_bf16 = ttnn.pad( + x_bf16, + [(0, 0), (0, 0), (0, 0), (0, padded_width - x_bf16.shape[-1])], + value=-sys.float_info.max, + sub_core_grids=cfg.sub_core_grids, + ) + + topk_values, topk_indices = ttnn.topk( + x_bf16, + k=cfg.max_top_k, + dim=-1, + sub_core_grids=cfg.sub_core_grid_topk, + ) + + # For 1D meshes use cluster_axis=None + sampling_cluster_axis = None if 1 in cluster_shape else 0 + + # Gather values + gathered_values = self._perform_all_gather( + topk_values, + dim=3, + cluster_axis=sampling_cluster_axis, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + num_links=cfg.num_gather_links, + buffer_key="SAMPLING_VALUES", + ) + ttnn.deallocate(topk_values) + + # Gather indices + gathered_indices = self._perform_all_gather( + topk_indices, + dim=3, + cluster_axis=sampling_cluster_axis, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + num_links=cfg.num_gather_links, + buffer_key="SAMPLING_INDICES", + ) + ttnn.deallocate(topk_indices) + + return gathered_values, gathered_indices + + # -- CCL helper ----------------------------------------------------------- + + def _perform_all_gather(self, tensor, dim, cluster_axis, memory_config, num_links, buffer_key=None): + """Flexible all-gather: prefer line_all_gather if available, else ttnn.all_gather. + + Port of TTSampling._perform_all_gather. + """ + if callable(self._line_all_gather): + kwargs = { + "dim": dim, + "cluster_axis": cluster_axis, + "memory_config": memory_config, + "num_links": num_links, + } + if self._line_all_gather_supports_buffer_key and buffer_key is not None: + kwargs["buffer_key"] = buffer_key + return self._line_all_gather(tensor, **kwargs) + + return ttnn.all_gather( + tensor, + dim=dim, + num_links=num_links, + memory_config=memory_config, + cluster_axis=cluster_axis, + topology=ttnn.Topology.Linear, + ) + + # -- (Backward compat) Model args factory ---------------------------------- + + @classmethod + def from_model_args(cls, mesh_device, tt_ccl, args, model_config=None) -> Sampling1D: + """Backward compat factory for TTTv1 model args.""" + cluster_shape = mesh_device.shape + if min(cluster_shape) > 1: + raise ValueError( + f"Sampling1D only supports 1D mesh topologies, got shape {cluster_shape}. " + "Use TTSampling for Galaxy (2D) topologies." + ) + padded_vocab_size = getattr(args, "padded_vocab_size", None) + vocab_size = padded_vocab_size if padded_vocab_size is not None else args.vocab_size + valid_vocab_size = getattr(args, "vocab_size", vocab_size) + + # Extract config from model_config dict + mc = model_config or getattr(args, "model_config", {}) + num_gather_links = 1 + if "GALAXY_NUM_LINKS" in mc: + max_links = mc["GALAXY_NUM_LINKS"] + max_top_k = getattr(args, "max_top_k", 32) + num_gather_links = min(max_top_k // 32, max_links) if max_top_k // 32 <= max_links else max_links + + sampling_memory_config = mc.get("DECODE_SAMPLING_INPUT_MEMCFG", ttnn.DRAM_MEMORY_CONFIG) + + allow_force_argmax = False + num_argmax_gather_links = num_gather_links + ag_topology = ttnn.Topology.Linear + argmax_chunks_per_sync = 10 + argmax_num_workers_per_link = 1 + if "SAMPLING_AG_CONFIG" in mc: + ag_cfg = mc["SAMPLING_AG_CONFIG"] + allow_force_argmax = ag_cfg.get("allow_force_argmax", False) + num_argmax_gather_links = ag_cfg.get("num_links", num_gather_links) + ag_topology = ag_cfg.get("topology", ttnn.Topology.Linear) + argmax_chunks_per_sync = ag_cfg.get("chunks_per_sync", 10) + + config = Sampling1DConfig( + vocab_size=vocab_size, + valid_vocab_size=valid_vocab_size, + mesh_device=mesh_device, + tt_ccl=tt_ccl, + max_batch_size=getattr(args, "max_batch_size", 32), + max_top_k=getattr(args, "max_top_k", 32), + sub_core_grids=getattr(args, "sub_core_grids", None), + sub_core_grid_topk=getattr(args, "sub_core_grid_topk", None), + start_core=getattr(args, "start_core", ttnn.CoreCoord(0, 0)), + num_gather_links=num_gather_links, + sampling_memory_config=sampling_memory_config, + allow_force_argmax=allow_force_argmax, + num_argmax_gather_links=num_argmax_gather_links, + ag_topology=ag_topology, + argmax_chunks_per_sync=argmax_chunks_per_sync, + argmax_num_workers_per_link=argmax_num_workers_per_link, + pad_to_power_of_2=getattr(args, "pad_logits_to_power_of_2", False), + ) + return cls.from_config(config) + + +# --------------------------------------------------------------------------- +# Config resolution +# --------------------------------------------------------------------------- + + +def _resolve_sampling1d_config(config: Sampling1DConfig) -> Sampling1DConfig: + """Fill None fields with topology-aware defaults.""" + import torch + + to_set: dict = {} + + # Phase 1: Device and CCL + mesh_device = config.mesh_device or ttnn.GetDefaultDevice() + to_set["mesh_device"] = mesh_device + + cluster_shape = mesh_device.shape + num_devices = mesh_device.get_num_devices() + multi_step_reduction = list(cluster_shape) == [1, 1] + + if num_devices > 1 and config.tt_ccl is None: + to_set["tt_ccl"] = get_tt_ccl(mesh_device) + + # Phase 2: Scalar config defaults + valid_vocab_size = config.valid_vocab_size if config.valid_vocab_size is not None else config.vocab_size + if valid_vocab_size > config.vocab_size: + raise ValueError(f"valid_vocab_size ({valid_vocab_size}) must be <= vocab_size ({config.vocab_size})") + to_set["valid_vocab_size"] = valid_vocab_size + if config.start_core is None: + to_set["start_core"] = ttnn.CoreCoord(0, 0) + if config.sampling_memory_config is None: + to_set["sampling_memory_config"] = ttnn.DRAM_MEMORY_CONFIG + if config.num_argmax_gather_links is None: + to_set["num_argmax_gather_links"] = config.num_gather_links + if config.ag_topology is None: + to_set["ag_topology"] = ttnn.Topology.Linear + + # Phase 3: Buffer specs + B = config.max_batch_size + K = config.max_top_k + V = config.vocab_size + replicate_mapper = ttnn.ShardTensor2dMesh(mesh_device, dims=(None, None), mesh_shape=cluster_shape) + + # num_devices_in_mesh for index computation + if multi_step_reduction: + num_devices_in_mesh = 2 + else: + num_devices_in_mesh = max(cluster_shape[0], cluster_shape[1]) + per_device_vocab = V // num_devices_in_mesh + + def _resolve_buf(field_val, defaults, source_factory): + if field_val is None: + return LazyBuffer(source=source_factory(), **defaults) + if isinstance(field_val, ttnn.Tensor): + return field_val + return resolve_lazy_buffer(field_val, **defaults) + + # index_offsets: [1, 1, B, K * num_devices_in_mesh] + def _make_index_offsets(): + offsets = torch.ones(1, 1, B, K * num_devices_in_mesh, dtype=torch.int64) + for device_id in range(num_devices_in_mesh): + offsets[:, :, :, device_id * K : (device_id + 1) * K] = device_id * per_device_vocab + return offsets + + idx_defaults = dict( + dtype=ttnn.int32, + layout=ttnn.TILE_LAYOUT, + device=mesh_device, + mesh_mapper=replicate_mapper, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + to_set["index_offsets"] = _resolve_buf(config.index_offsets, idx_defaults, _make_index_offsets) + + vocab_shard_dims = get_vocab_shard_dims(cluster_shape) + invalid_vocab_defaults = dict( + dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, + device=mesh_device, + mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=vocab_shard_dims, mesh_shape=cluster_shape), + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + invalid_vocab_tail_mask = build_tail_invalid_vocab_mask( + valid_vocab_size, + V, + B, + cluster_shape, + tile_size=ttnn.TILE_SIZE, + ) + if config.invalid_vocab_mask is None and ( + invalid_vocab_tail_mask is not None or config.invalid_vocab_tail_mask is not None + ): + to_set["invalid_vocab_tail_width"] = ( + invalid_vocab_tail_mask.tail_width + if invalid_vocab_tail_mask is not None + else config.invalid_vocab_tail_width + ) + to_set["invalid_vocab_tail_mask"] = _resolve_buf( + config.invalid_vocab_tail_mask, + invalid_vocab_defaults, + lambda: invalid_vocab_tail_mask.mask, + ) + else: + invalid_vocab_mask = build_invalid_vocab_mask(valid_vocab_size, V, B) + if invalid_vocab_mask is not None or config.invalid_vocab_mask is not None: + to_set["invalid_vocab_mask"] = _resolve_buf( + config.invalid_vocab_mask, + invalid_vocab_defaults, + lambda: invalid_vocab_mask, + ) + + # seeds and user_ids: [B], uint32, ROW_MAJOR + seed_defaults = dict( + dtype=ttnn.uint32, + layout=ttnn.ROW_MAJOR_LAYOUT, + device=mesh_device, + memory_config=ttnn.DRAM_MEMORY_CONFIG, + ) + to_set["seeds"] = _resolve_buf( + config.seeds, seed_defaults, lambda: torch.arange(B, dtype=torch.int64).to(torch.int32) + ) + to_set["user_ids"] = _resolve_buf( + config.user_ids, seed_defaults, lambda: torch.arange(B, dtype=torch.int64).to(torch.int32) + ) + + resolved = replace(config, **to_set) + assert resolved.is_resolved(), "Config not fully resolved after _resolve_sampling1d_config" + return resolved + + +# -- Helper functions ------------------------------------------------------- + + +def _materialize(buf): + """Materialize a buffer field: ttnn.Tensor passthrough, LazyBuffer → get_device_buffer().""" + if isinstance(buf, ttnn.Tensor): + return buf + return buf.get_device_buffer() + + +def _attach_cleanup_failures(primary, failures): + if not failures: + return + previous = tuple(getattr(primary, "cleanup_failures", ())) + primary.cleanup_failures = previous + tuple(failures) + add_note = getattr(primary, "add_note", None) + if callable(add_note): + add_note(f"cleanup also encountered {len(failures)} failure(s)") + + +def _raise_cleanup_failures(failures): + primary = failures[0] + if len(failures) > 1: + previous = tuple(getattr(primary, "cleanup_failures", ())) + primary.cleanup_failures = previous + tuple(failures[1:]) + add_note = getattr(primary, "add_note", None) + if callable(add_note): + add_note(f"cleanup also encountered {len(failures) - 1} additional failure(s)") + raise primary diff --git a/code/models/common/modules/sampling/sampling_state_1d.py b/code/models/common/modules/sampling/sampling_state_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..9fd350ce2a24e26b820ed72d025d8a28ca8babd7 --- /dev/null +++ b/code/models/common/modules/sampling/sampling_state_1d.py @@ -0,0 +1,937 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC +# SPDX-License-Identifier: Apache-2.0 + +"""Per-lane TTTv2 sampling state orchestration. + +``Sampling1D`` owns sampling programs and a stable seed buffer. This module +borrows that sampler, constructs the matching ``Penalties1D`` module, and +returns all mutable request state to the caller in ``SamplingState1DState``. +It does not import or fall back to the legacy sampling generator, penalties, +or seed manager. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass, replace +from typing import Any + +import torch + +from models.common.modules.lazy_buffer import LazyBuffer +from models.common.modules.sampling.params import PreparedSamplingParams, place_prepared_sampling_params +from models.common.modules.sampling.penalties_1d import ( + Penalties1D, + Penalties1DConfig, + PenaltyAccumulator, + PenaltyParams, + _materialize, +) +from models.common.modules.sampling.sampling_1d import Sampling1D +from models.common.modules.sampling.seed_manager_1d import SeedManager1D, SeedState + + +@dataclass(frozen=True) +class SamplingStaticIdentity1D: + """Sampling choices that must participate in trace/program identity.""" + + sampling_path: str + penalties_enabled: bool + log_probs_enabled: bool + logprob_modes: tuple[str, ...] + + +@dataclass +class SamplingState1DState: + """Mutable per-lane state owned by the runtime caller.""" + + seed_state: SeedState + penalty_params: PenaltyParams + penalty_accumulator: PenaltyAccumulator + active_mask: tuple[bool, ...] + static_identity: SamplingStaticIdentity1D | None = None + pending_sample_id: int | None = None + next_sample_id: int = 0 + penalty_history_valid: bool = True + released: bool = False + + @property + def active_slots(self) -> tuple[int, ...]: + return tuple(slot for slot, active in enumerate(self.active_mask) if active) + + +class SamplingState1D: + """Compose one borrowed sampler with native seed and penalty state. + + One instance belongs to one executor lane. The ``Sampling1D`` argument is + borrowed and is never released here. The topology-matched ``Penalties1D`` + instance is owned by this controller. Mutable request state is returned by + :meth:`create_state` rather than retained on the controller. + + ``salt_duplicate_seeds`` is forwarded to the native ``SeedManager1D``. + vLLM-facing executors disable it so independent requests that share a seed + stay bit-identical. + """ + + def __init__( + self, + sampling: Sampling1D, + *, + penalties_factory: Callable[[Penalties1DConfig], Penalties1D] | None = None, + seed_manager_factory: Callable[[Any], SeedManager1D] | None = None, + salt_duplicate_seeds: bool = True, + ) -> None: + sampling_config = getattr(sampling, "config", None) + if sampling_config is None: + raise TypeError("sampling must expose a resolved Sampling1D config") + self._validate_1d_topology(sampling_config.mesh_device) + + penalties_config = Penalties1DConfig( + vocab_size=int(sampling_config.vocab_size), + mesh_device=sampling_config.mesh_device, + max_batch_size=int(sampling_config.max_batch_size), + sub_core_grids=sampling_config.sub_core_grids, + ) + if not isinstance(salt_duplicate_seeds, bool): + raise TypeError("salt_duplicate_seeds must be bool") + make_penalties = penalties_factory or Penalties1D.from_config + make_seed_manager = seed_manager_factory or ( + lambda config: SeedManager1D(config, salt_duplicate_seeds=salt_duplicate_seeds) + ) + + self.sampling = sampling + self.salt_duplicate_seeds = salt_duplicate_seeds + self.penalties = make_penalties(penalties_config) + try: + self.seed_manager = make_seed_manager(sampling_config) + self._validate_component_contracts() + except BaseException as primary: + try: + self.penalties.release() + except BaseException as cleanup_error: + self._attach_cleanup_failures(primary, (cleanup_error,)) + raise + self._state_leased = False + + @property + def sampling_config(self): + return self.sampling.config + + @property + def penalties_config(self): + return self.penalties.config + + def create_state(self) -> SamplingState1DState: + """Materialize buffers and return fresh caller-owned lane state.""" + + if self._state_leased: + raise RuntimeError("one SamplingState1DState is already live for this lane") + state = None + try: + self.penalties.load_device_buffers() + self._write_noop_penalty_params() + state = SamplingState1DState( + seed_state=self.seed_manager.create_state(), + penalty_params=PenaltyParams( + prompt_mask=_materialize(self.penalties.config.prompt_mask), + presence_penalties=_materialize(self.penalties.config.presence_penalties), + frequency_penalties=_materialize(self.penalties.config.frequency_penalties), + repetition_penalties=_materialize(self.penalties.config.repetition_penalties), + inverse_repetition_penalties=_materialize(self.penalties.config.inverse_repetition_penalties), + ), + penalty_accumulator=PenaltyAccumulator( + output_mask=_materialize(self.penalties.config.output_mask), + output_counts=_materialize(self.penalties.config.output_counts), + output_counts_gathered=_materialize(self.penalties.config.output_counts_gathered), + ), + active_mask=(False,) * int(self.sampling.config.max_batch_size), + ) + self.seed_manager.reset(state.seed_state) + self._rebuild_penalty_history( + state, + prompt_tokens=None, + output_tokens=None, + active_mask=state.active_mask, + ) + except BaseException as primary: + cleanup_failures = [] + if state is not None: + try: + self.seed_manager.reset(state.seed_state) + except BaseException as error: + cleanup_failures.append(error) + try: + self.penalties.release() + except BaseException as error: + cleanup_failures.append(error) + self._attach_cleanup_failures(primary, cleanup_failures) + raise + self._state_leased = True + return state + + def static_identity(self, prepared: PreparedSamplingParams) -> SamplingStaticIdentity1D: + """Return the trace/program identity for one prepared request.""" + + self._validate_prepared(prepared) + active_modes = tuple(mode for active, mode in zip(prepared.active_mask, prepared.logprob_modes) if active) + return SamplingStaticIdentity1D( + sampling_path=prepared.sampling_path, + penalties_enabled=prepared.penalties_enabled, + log_probs_enabled=prepared.log_probs_enabled, + logprob_modes=active_modes, + ) + + def admit( + self, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + slots: Iterable[int] | None = None, + ) -> None: + """Register prefill admissions and rebuild complete penalty history.""" + + self._require_idle(state) + self._validate_prepared(prepared) + if prepared.slot_remap is not None: + raise ValueError("consume slot_remap with apply_slot_remap before admission") + self._validate_prepared_history(prepared) + active_slots = self._active_slots(prepared) + admission_slots = active_slots if slots is None else tuple(int(slot) for slot in slots) + if len(set(admission_slots)) != len(admission_slots): + raise ValueError("admission slots must be unique") + if any(slot not in active_slots for slot in admission_slots): + raise ValueError("admission slots must be active prepared-sampling rows") + + request_seeds = tuple(prepared.seeds[slot] for slot in admission_slots) + self.seed_manager.admit(state.seed_state, request_seeds, admission_slots) + self.seed_manager.synchronize( + state.seed_state, + prepared.seeds, + active_slots, + reset_batch=True, + ) + self._synchronize_penalties_and_identity(state, prepared, rebuild_history=True) + + def admit_prefill( + self, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + slots: Iterable[int], + positions: Sequence[int] | None = None, + ) -> None: + """Admit one request-ordered prefill into its persistent decode slot.""" + + self._require_idle(state) + self._validate_prepared(prepared) + self._validate_prepared_history(prepared) + active_sources = self._active_slots(prepared) + destination_slots = tuple(int(slot) for slot in slots) + if len(active_sources) != 1 or len(destination_slots) != 1: + raise ValueError("native device-sampled prefill currently requires exactly one active request") + destination = destination_slots[0] + if prepared.slot_remap is not None: + self.seed_manager.apply_slot_remap(state.seed_state, prepared.slot_remap) + placed = place_prepared_sampling_params(prepared, destination_slots) + execution_prepared = _broadcast_prefill_prepared(prepared, active_sources[0]) + self.seed_manager.admit( + state.seed_state, + (prepared.seeds[active_sources[0]],), + destination_slots, + ) + self._write_penalty_params(execution_prepared) + self._rebuild_penalty_history( + state, + prompt_tokens=execution_prepared.prompt_tokens, + output_tokens=execution_prepared.output_tokens, + active_mask=execution_prepared.active_mask, + ) + position = None if positions is None else int(tuple(positions)[0]) + if prepared.sampling_path == "topk": + self.seed_manager.refresh_prefill_replicated( + state.seed_state, + destination, + position=position, + ) + else: + self.seed_manager.restore_defaults(state.seed_state) + state.active_mask = tuple(state.seed_state.active) + state.static_identity = self.static_identity(placed) + state.penalty_history_valid = not prepared.penalties_enabled + + def prefill_forward( + self, + logits, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + k=None, + p=None, + temp=None, + tt_out_tok=None, + count_tokens: bool = True, + ): + """Apply request-ordered prefill penalties, sample, and count once.""" + + self._require_idle(state) + self._validate_prepared(prepared) + if prepared.penalties_enabled: + logits = self.penalties.decode_forward(logits, state.penalty_params, state.penalty_accumulator) + output = self.sampling.decode_forward( + logits, + k=k, + p=p, + temp=temp, + seeds=self.seed_manager.get_seed_device_buffer() if prepared.sampling_path == "topk" else None, + tt_out_tok=tt_out_tok, + enable_log_probs=list(prepared.enable_log_probs), + ) + if prepared.penalties_enabled and count_tokens: + self.penalties.update_output_tokens(state.penalty_accumulator, output[0]) + state.penalty_history_valid = False + return output + + def apply_slot_remap( + self, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + ) -> PreparedSamplingParams: + """Move state, rebuild histories, and return a consumed-remap value.""" + + self._require_idle(state) + self._validate_prepared(prepared) + if prepared.slot_remap is None: + raise ValueError("prepared sampling state does not contain a slot_remap") + self._validate_prepared_history(prepared) + active_slots = self._active_slots(prepared) + projected_active, projected_seeds = self._project_seed_membership(state.seed_state, prepared.slot_remap) + new_slots = self._validate_remapped_decode_membership( + prepared, + active_slots, + projected_active=projected_active, + projected_seeds=projected_seeds, + ) + self.seed_manager.apply_slot_remap(state.seed_state, prepared.slot_remap) + if new_slots: + self.seed_manager.admit( + state.seed_state, + tuple(prepared.seeds[slot] for slot in new_slots), + new_slots, + ) + self.seed_manager.synchronize( + state.seed_state, + prepared.seeds, + active_slots, + reset_batch=False, + ) + self._synchronize_penalties_and_identity(state, prepared, rebuild_history=True) + return replace(prepared, slot_remap=None) + + def synchronize_decode( + self, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + reset_batch: bool, + ) -> PreparedSamplingParams: + """Synchronize one decode boundary without resetting survivor streams. + + New requests are admitted by prefill. ``reset_batch`` permits changed + active membership and rebuilds host-provided prompt/output histories; + unchanged slots retain their SeedState counters, salts, and RNG state. + """ + + self._require_idle(state) + self._validate_prepared(prepared) + if reset_batch: + self._validate_prepared_history(prepared) + if not state.penalty_history_valid and not reset_batch: + raise RuntimeError("decode after sampled prefill requires reset_batch=True with complete output history") + if prepared.slot_remap is not None: + prepared = self.apply_slot_remap(state, prepared) + active_slots = self._active_slots(prepared) + self.seed_manager.synchronize( + state.seed_state, + prepared.seeds, + active_slots, + reset_batch=bool(reset_batch), + ) + identity = self.static_identity(prepared) + rebuild_history = bool(reset_batch) + if state.static_identity is None and active_slots: + if not reset_batch: + raise RuntimeError("initial decode sampling state requires prefill admission or reset_batch=True") + rebuild_history = True + if rebuild_history: + self._synchronize_penalties_and_identity(state, prepared, rebuild_history=True) + else: + state.active_mask = prepared.active_mask + state.static_identity = identity + self._write_penalty_params(prepared) + return prepared + + @staticmethod + def _project_seed_membership(seed_state: SeedState, remap: Sequence[int] | torch.Tensor): + """Project active/seed membership through one remap without mutation.""" + + flat = remap.reshape(-1).tolist() if isinstance(remap, torch.Tensor) else list(remap) + capacity = len(seed_state.active) + if len(flat) != capacity: + raise ValueError(f"slot remap must contain {capacity} entries") + normalized = tuple(int(slot) for slot in flat) + if any(slot < 0 or slot >= capacity for slot in normalized): + raise ValueError("slot remap source is outside the seed-state capacity") + moves = tuple((old_slot, new_slot) for new_slot, old_slot in enumerate(normalized) if old_slot != new_slot) + moved_sources = tuple(old_slot for old_slot, _ in moves) + if len(set(moved_sources)) != len(moved_sources): + raise ValueError("slot remap cannot copy one seed stream into multiple destinations") + + active_before = tuple(seed_state.active) + seeds_before = tuple(seed_state.request_seeds) + projected_active = list(active_before) + projected_seeds = list(seeds_before) + for old_slot, new_slot in moves: + projected_active[new_slot] = active_before[old_slot] + projected_seeds[new_slot] = seeds_before[old_slot] if active_before[old_slot] else None + moved_destinations = {new_slot for _, new_slot in moves} + for old_slot in set(moved_sources).difference(moved_destinations): + projected_active[old_slot] = False + projected_seeds[old_slot] = None + return tuple(projected_active), tuple(projected_seeds) + + @staticmethod + def _validate_remapped_decode_membership( + prepared: PreparedSamplingParams, + active_slots: Sequence[int], + *, + projected_active: Sequence[bool], + projected_seeds: Sequence[int | None], + ) -> tuple[int, ...]: + """Return genuinely new destinations or reject a changed survivor.""" + + changed_slots = tuple( + slot for slot in active_slots if projected_active[slot] and projected_seeds[slot] != prepared.seeds[slot] + ) + if changed_slots: + raise RuntimeError( + "changed active seed slots require reset_batch=True or an explicit prefill admission: " + f"{list(changed_slots)}" + ) + new_slots = tuple(slot for slot in active_slots if not projected_active[slot]) + if new_slots and not any(projected_active[slot] for slot in active_slots): + raise RuntimeError( + "new or changed active seed slots require reset_batch=True or an explicit prefill admission: " + f"{list(new_slots)}" + ) + return new_slots + + def cleanup( + self, + state: SamplingState1DState, + live_slots: Iterable[int], + *, + prepared: PreparedSamplingParams | None = None, + ) -> None: + """Remove completed requests and clear or rebuild their penalty state.""" + + self._require_idle(state) + live = tuple(int(slot) for slot in live_slots) + if len(set(live)) != len(live): + raise ValueError("live_slots must be unique") + capacity = int(self.sampling.config.max_batch_size) + if any(slot < 0 or slot >= capacity for slot in live): + raise ValueError("live_slots contain a row outside the lane capacity") + if live and prepared is None: + raise ValueError("prepared sampling state is required to preserve live penalty history") + if prepared is not None: + self._validate_prepared(prepared) + if prepared.slot_remap is not None: + raise ValueError("consume slot_remap before cleanup") + self._validate_prepared_history(prepared) + if set(live) != set(self._active_slots(prepared)): + raise ValueError("live_slots must match the active rows in prepared sampling state") + self.seed_manager.cleanup(state.seed_state, live) + if not live: + self._write_noop_penalty_params() + inactive_mask = (False,) * capacity + self._rebuild_penalty_history( + state, + prompt_tokens=None, + output_tokens=None, + active_mask=inactive_mask, + ) + state.active_mask = inactive_mask + state.static_identity = None + state.penalty_history_valid = True + return + self._synchronize_penalties_and_identity(state, prepared, rebuild_history=True) + + def reset( + self, + state: SamplingState1DState, + prepared: PreparedSamplingParams | None = None, + ) -> None: + """Clear the lane, optionally admitting one complete replacement batch.""" + + self._require_idle(state) + self.seed_manager.reset(state.seed_state) + self._write_noop_penalty_params() + inactive_mask = (False,) * int(self.sampling.config.max_batch_size) + self._rebuild_penalty_history( + state, + prompt_tokens=None, + output_tokens=None, + active_mask=inactive_mask, + ) + state.active_mask = inactive_mask + state.static_identity = None + state.penalty_history_valid = True + if prepared is not None: + self.admit(state, prepared) + + def refresh_dynamic_inputs( + self, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + positions=None, + advance_seeds: bool = True, + ) -> None: + """Refresh penalty constants and stable seeds before eager/trace replay.""" + + self._require_idle(state) + self._validate_prepared(prepared) + if prepared.slot_remap is not None: + raise RuntimeError("apply slot_remap before refreshing dynamic sampling inputs") + identity = self.static_identity(prepared) + if state.static_identity is None: + raise RuntimeError("sampling state must be admitted before refresh") + if identity != state.static_identity: + raise RuntimeError( + "sampling static identity changed; select the matching trace/program and " + "admit, remap, or reset the lane before replay" + ) + + active_slots = self._active_slots(prepared) + self.seed_manager.synchronize( + state.seed_state, + prepared.seeds, + active_slots, + reset_batch=False, + ) + self._write_penalty_params(prepared) + if prepared.sampling_path == "topk" and advance_seeds: + self.seed_manager.refresh(state.seed_state, active_slots, positions=positions) + else: + self.seed_manager.restore_defaults(state.seed_state) + state.active_mask = prepared.active_mask + + def before_sampling( + self, + logits, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + positions=None, + advance_seeds: bool = True, + ): + """Apply penalties before sampling and open one exactly-once update.""" + + self._require_idle(state) + if advance_seeds: + self.refresh_dynamic_inputs( + state, + prepared, + positions=positions, + advance_seeds=True, + ) + else: + # Compile-only execution and trace capture prepare state before + # entering the device body. Do not issue host-to-device writes + # here: a captured write would replay after the runtime refresh and + # restore capture-time penalty/seed values over the live request. + self._validate_prepared(prepared) + if prepared.slot_remap is not None: + raise RuntimeError("apply slot_remap before entering the sampling body") + identity = self.static_identity(prepared) + if state.static_identity != identity: + raise RuntimeError("sampling body state does not match its static trace/program identity") + if prepared.penalties_enabled: + logits = self.penalties.decode_forward( + logits, + state.penalty_params, + state.penalty_accumulator, + ) + sample_id = state.next_sample_id + state.next_sample_id += 1 + state.pending_sample_id = sample_id + return logits, sample_id + + def after_sampling( + self, + state: SamplingState1DState, + sampled_tokens, + *, + sample_id: int, + count_tokens: bool = True, + ) -> None: + """Record sampled tokens exactly once for the matching sampling step.""" + + self._validate_state(state) + if state.pending_sample_id != int(sample_id): + raise RuntimeError( + f"sample_id {sample_id} is not the pending sampling step " f"({state.pending_sample_id})" + ) + if count_tokens and state.static_identity is not None and state.static_identity.penalties_enabled: + # Leave the step pending if the device update fails: retrying blindly + # could count a partially applied token twice. + self.penalties.update_output_tokens(state.penalty_accumulator, sampled_tokens) + state.pending_sample_id = None + + def cancel_sampling(self, state: SamplingState1DState, *, sample_id: int) -> None: + """Cancel a step only when sampling produced no token.""" + + self._validate_state(state) + if state.pending_sample_id != int(sample_id): + raise RuntimeError(f"cannot cancel non-pending sample_id {sample_id}") + state.pending_sample_id = None + + def decode_forward( + self, + logits, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + k=None, + p=None, + temp=None, + positions=None, + tt_out_tok=None, + count_tokens: bool = True, + advance_seeds: bool = True, + ): + """Apply penalties, sample, and update output history in order. + + Eager compile-only warmups set ``count_tokens=False`` so they do not + admit a phantom generated token. Trace capture keeps it true because + capture records the update for replay without executing it immediately. + """ + + if prepared.sampling_path == "argmax": + if any(value is not None for value in (k, p, temp)): + raise ValueError("argmax sampling requires k, p, and temp to be None") + elif any(value is None for value in (k, p, temp)): + raise ValueError("topk sampling requires k, p, and temp device tensors") + + penalized_logits, sample_id = self.before_sampling( + logits, + state, + prepared, + positions=positions, + advance_seeds=advance_seeds, + ) + try: + output = self.sampling.decode_forward( + penalized_logits, + k=k, + p=p, + temp=temp, + seeds=(self.seed_manager.get_seed_device_buffer() if prepared.sampling_path == "topk" else None), + tt_out_tok=tt_out_tok, + enable_log_probs=list(prepared.enable_log_probs), + ) + except BaseException: + self.cancel_sampling(state, sample_id=sample_id) + raise + sampled_tokens = output[0] + self.after_sampling( + state, + sampled_tokens, + sample_id=sample_id, + count_tokens=count_tokens, + ) + return output + + def release(self, state: SamplingState1DState | None = None) -> None: + """Release controller-owned penalties, never the borrowed sampler.""" + + if state is None: + if self._state_leased: + raise ValueError("the live caller-owned state is required for release") + self.penalties.release() + return + if state.released: + self.penalties.release() + self._state_leased = False + return + self._require_idle(state) + failures = [] + try: + self.seed_manager.reset(state.seed_state) + except BaseException as error: + failures.append(error) + finally: + state.active_mask = (False,) * int(self.sampling.config.max_batch_size) + state.static_identity = None + state.penalty_history_valid = True + state.released = True + try: + self.penalties.release() + except BaseException as error: + failures.append(error) + finally: + self._state_leased = False + if failures: + primary = failures[0] + self._attach_cleanup_failures(primary, failures[1:]) + raise primary + + # Internal helpers ----------------------------------------------------- + + @staticmethod + def _validate_1d_topology(mesh_device) -> None: + shape = tuple(int(value) for value in mesh_device.shape) + if len(shape) != 2 or min(shape) != 1: + raise ValueError(f"SamplingState1D only supports 1D mesh topologies, got shape {shape}") + + def _validate_component_contracts(self) -> None: + sampling = self.sampling.config + penalties = self.penalties.config + if penalties.mesh_device is not sampling.mesh_device: + raise ValueError("Penalties1D must borrow the Sampling1D mesh_device") + if int(penalties.vocab_size) != int(sampling.vocab_size): + raise ValueError("Penalties1D and Sampling1D vocab_size must match") + if int(penalties.max_batch_size) != int(sampling.max_batch_size): + raise ValueError("Penalties1D and Sampling1D max_batch_size must match") + if penalties.sub_core_grids is not sampling.sub_core_grids: + raise ValueError("Penalties1D must borrow the Sampling1D sub_core_grids") + if int(self.seed_manager.max_batch_size) != int(sampling.max_batch_size): + raise ValueError("SeedManager1D and Sampling1D max_batch_size must match") + if self.seed_manager.seed_buffer is not sampling.seeds: + raise ValueError("SeedManager1D must borrow the Sampling1D seeds LazyBuffer") + + def _validate_state(self, state: SamplingState1DState) -> None: + if not isinstance(state, SamplingState1DState): + raise TypeError("state must be a caller-owned SamplingState1DState") + if state.released: + raise RuntimeError("sampling state was released") + capacity = int(self.sampling.config.max_batch_size) + if len(state.active_mask) != capacity or state.seed_state.capacity != capacity: + raise ValueError("sampling state capacity does not match the lane") + + def _require_idle(self, state: SamplingState1DState) -> None: + self._validate_state(state) + if state.pending_sample_id is not None: + raise RuntimeError(f"sampling step {state.pending_sample_id} must be completed or cancelled first") + + def _validate_prepared(self, prepared: PreparedSamplingParams) -> None: + if not isinstance(prepared, PreparedSamplingParams): + raise TypeError("prepared must be PreparedSamplingParams") + sampling = self.sampling.config + if prepared.batch_size != int(sampling.max_batch_size): + raise ValueError("prepared batch_size does not match the Sampling1D lane capacity") + if prepared.max_device_top_k != int(sampling.max_top_k): + raise ValueError("prepared max_device_top_k does not match Sampling1D.max_top_k") + for slot, active in enumerate(prepared.active_mask): + if not active: + continue + repetition = float(prepared.repetition_penalty[slot]) + if repetition <= 0.0: + raise ValueError("active repetition penalties must be positive") + + @staticmethod + def _active_slots(prepared: PreparedSamplingParams) -> tuple[int, ...]: + return tuple(slot for slot, active in enumerate(prepared.active_mask) if active) + + def _validate_prepared_history(self, prepared: PreparedSamplingParams) -> None: + capacity = int(self.sampling.config.max_batch_size) + prompt = self._history_tensor( + prepared.prompt_tokens, + name="prompt_tokens", + capacity=capacity, + active_mask=prepared.active_mask, + ) + if prompt is None and any( + active and float(value) != 1.0 for active, value in zip(prepared.active_mask, prepared.repetition_penalty) + ): + raise ValueError("prompt_tokens are required when repetition penalty is enabled") + self._history_tensor( + prepared.output_tokens, + name="output_tokens", + capacity=capacity, + active_mask=prepared.active_mask, + ) + + def _synchronize_penalties_and_identity( + self, + state: SamplingState1DState, + prepared: PreparedSamplingParams, + *, + rebuild_history: bool, + ) -> None: + self._write_penalty_params(prepared) + if rebuild_history: + self._rebuild_penalty_history( + state, + prompt_tokens=prepared.prompt_tokens, + output_tokens=prepared.output_tokens, + active_mask=prepared.active_mask, + ) + state.active_mask = prepared.active_mask + state.static_identity = self.static_identity(prepared) + state.penalty_history_valid = True + + def _write_noop_penalty_params(self) -> None: + capacity = int(self.sampling.config.max_batch_size) + self._update_penalty_buffer("presence_penalties", torch.zeros(capacity, 1)) + self._update_penalty_buffer("frequency_penalties", torch.zeros(capacity, 1)) + self._update_penalty_buffer("repetition_penalties", torch.ones(capacity, 1)) + self._update_penalty_buffer("inverse_repetition_penalties", torch.ones(capacity, 1)) + + def _write_penalty_params(self, prepared: PreparedSamplingParams) -> None: + presence = [] + frequency = [] + repetition = [] + for slot, active in enumerate(prepared.active_mask): + presence.append(float(prepared.presence_penalty[slot]) if active else 0.0) + frequency.append(float(prepared.frequency_penalty[slot]) if active else 0.0) + repetition.append(float(prepared.repetition_penalty[slot]) if active else 1.0) + inverse = [1.0 / value for value in repetition] + self._update_penalty_buffer("presence_penalties", torch.tensor(presence, dtype=torch.float32).reshape(-1, 1)) + self._update_penalty_buffer("frequency_penalties", torch.tensor(frequency, dtype=torch.float32).reshape(-1, 1)) + self._update_penalty_buffer( + "repetition_penalties", torch.tensor(repetition, dtype=torch.float32).reshape(-1, 1) + ) + self._update_penalty_buffer( + "inverse_repetition_penalties", torch.tensor(inverse, dtype=torch.float32).reshape(-1, 1) + ) + + def _update_penalty_buffer(self, name: str, source: torch.Tensor) -> None: + specification = getattr(self.penalties.config, name) + if not isinstance(specification, LazyBuffer) and not callable(getattr(specification, "update", None)): + raise TypeError(f"Penalties1DConfig.{name} must be a mutable LazyBuffer") + specification.update(source) + + def _rebuild_penalty_history( + self, + state: SamplingState1DState, + *, + prompt_tokens, + output_tokens, + active_mask: tuple[bool, ...], + ) -> None: + capacity = int(self.sampling.config.max_batch_size) + prompt = self._history_tensor( + prompt_tokens, + name="prompt_tokens", + capacity=capacity, + active_mask=active_mask, + ) + repetition_needs_prompt = any( + active and float(value) != 1.0 for active, value in zip(active_mask, self._current_repetition_values()) + ) + if prompt is None and repetition_needs_prompt: + raise ValueError("prompt_tokens are required when repetition penalty is enabled") + if prompt is None: + prompt = torch.full((capacity, 1), -1, dtype=torch.int64) + output = self._history_tensor( + output_tokens, + name="output_tokens", + capacity=capacity, + active_mask=active_mask, + ) + self.penalties.init_prompt_penalties( + state.penalty_params, + state.penalty_accumulator, + prompt, + ) + self.penalties.reset_output_tokens(state.penalty_accumulator, output) + + @staticmethod + def _history_tensor(value, *, name: str, capacity: int, active_mask: tuple[bool, ...]): + if value is None: + return None + if isinstance(value, torch.Tensor): + tensor = value + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + if len(value) == 0: + return None + tensor = torch.as_tensor(value) + else: + raise TypeError("sampling token history must be a torch.Tensor or sequence") + if tensor.numel() == 0: + return None + if tensor.ndim == 0: + tensor = tensor.reshape(1, 1) + elif tensor.ndim == 1: + tensor = tensor.reshape(1, -1) + rows = int(tensor.shape[0]) + if rows > capacity: + raise ValueError(f"{name} has {rows} rows, exceeding lane capacity {capacity}") + active_slots = tuple(slot for slot, active in enumerate(active_mask) if active) + if active_slots and rows <= max(active_slots): + raise ValueError(f"{name} does not cover active slot {max(active_slots)}") + return tensor + + def _current_repetition_values(self) -> tuple[float, ...]: + source = getattr(self.penalties.config.repetition_penalties, "source", None) + if not isinstance(source, torch.Tensor): + raise TypeError("repetition penalty buffer must retain a torch source") + return tuple(float(value) for value in source.reshape(-1).tolist()) + + @staticmethod + def _attach_cleanup_failures(primary, failures) -> None: + if not failures: + return + previous = tuple(getattr(primary, "cleanup_failures", ())) + primary.cleanup_failures = previous + tuple(failures) + add_note = getattr(primary, "add_note", None) + if callable(add_note): + add_note(f"cleanup also encountered {len(failures)} failure(s)") + + +def _broadcast_prefill_prepared( + prepared: PreparedSamplingParams, + source_row: int, +) -> PreparedSamplingParams: + """Broadcast one request over every physical prefill sampling row.""" + + capacity = prepared.batch_size + + def broadcast(values): + return (values[int(source_row)],) * capacity + + return replace( + prepared, + top_k=broadcast(prepared.top_k), + top_p=broadcast(prepared.top_p), + temperature=broadcast(prepared.temperature), + presence_penalty=broadcast(prepared.presence_penalty), + frequency_penalty=broadcast(prepared.frequency_penalty), + repetition_penalty=broadcast(prepared.repetition_penalty), + seeds=broadcast(prepared.seeds), + enable_log_probs=broadcast(prepared.enable_log_probs), + num_logprobs=broadcast(prepared.num_logprobs), + logprob_modes=broadcast(prepared.logprob_modes), + greedy_mask=broadcast(prepared.greedy_mask), + row_paths=broadcast(prepared.row_paths), + active_mask=(True,) * capacity, + active_rows=capacity, + prompt_tokens=_broadcast_prefill_history(prepared.prompt_tokens, source_row, capacity), + output_tokens=_broadcast_prefill_history(prepared.output_tokens, source_row, capacity), + slot_remap=None, + ) + + +def _broadcast_prefill_history(value, source_row: int, capacity: int): + if value is None: + return None + if isinstance(value, torch.Tensor): + if value.ndim == 0: + value = value.reshape(1, 1) + row = 0 if int(value.shape[0]) == 1 else int(source_row) + return value[row : row + 1].expand((capacity, *value.shape[1:])).clone() + values = list(value) + row = values[0 if len(values) == 1 else int(source_row)] + repeated = [row for _ in range(capacity)] + return tuple(repeated) if isinstance(value, tuple) else repeated diff --git a/code/models/common/modules/sampling/seed_manager_1d.py b/code/models/common/modules/sampling/seed_manager_1d.py new file mode 100644 index 0000000000000000000000000000000000000000..dc9a7ad58f853ffdc93cc60e0d6609112be1e387 --- /dev/null +++ b/code/models/common/modules/sampling/seed_manager_1d.py @@ -0,0 +1,651 @@ +# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +# SPDX-License-Identifier: Apache-2.0 + +"""Caller-owned request seed state for :mod:`sampling_1d`. + +``Sampling1D`` owns a stable seed ``LazyBuffer`` because traces capture its +device handle. Request seeds, however, are invocation state. ``SeedManager1D`` +bridges those two lifetimes without storing request state on the module or on +the manager: every mutable lifecycle field lives in a caller-owned +``SeedState`` passed to each operation. + +This module is intentionally independent of the legacy sampling generator and +state classes. It only borrows the ``Sampling1DConfig.seeds`` ``LazyBuffer``. +""" + +from __future__ import annotations + +import copy +import random +import secrets +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Iterable, Sequence + +import torch + +from models.common.modules.lazy_buffer import LazyBuffer + +if TYPE_CHECKING: + from models.common.modules.sampling.sampling_1d import Sampling1DConfig + + +# The device reserves UINT32_MAX as its "advance the existing RNG" sentinel. +# Seeds generated here are always ordinary positive seed values. +MAX_UINT32 = 2**32 - 1 +DEVICE_SEED_MAX = 1_000_000 +_UINT64_MASK = (1 << 64) - 1 + + +def _hash_request_seed_to_device_seed(seed: int, counter: int, salt: int = 0) -> int: + """Return a stable, bounded device seed for one request token. + + ``salt`` distinguishes simultaneous requests carrying the same explicit + request seed. A unique request keeps salt zero, preserving the familiar + slot-independent ``(seed, counter)`` stream. + """ + + value = (int(seed) & _UINT64_MASK) ^ ((int(counter) + 0x9E3779B97F4A7C15) & _UINT64_MASK) + value ^= (int(salt) * 0xD1B54A32D192ED03) & _UINT64_MASK + value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & _UINT64_MASK + value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & _UINT64_MASK + value = (value ^ (value >> 31)) & _UINT64_MASK + return (value % DEVICE_SEED_MAX) + 1 + + +@dataclass(frozen=True) +class SeedSlotState: + """Immutable checkpoint for a request temporarily removed from a lane. + + Callers use :meth:`SeedManager1D.suspend` before preemption and pass the + checkpoint back to :meth:`SeedManager1D.resume`. The Python RNG state is + included so an unseeded request continues its varied stream instead of + being reinitialized after resume. + """ + + request_seed: int | None + token_counter: int + salt: int + unseeded_rng_state: object + last_absolute_position: int | None + current_device_seed: int | None + + +@dataclass(frozen=True) +class SeedStateSnapshot: + """Immutable, assertion-friendly view of a mutable :class:`SeedState`.""" + + active: tuple[bool, ...] + request_seeds: tuple[int | None, ...] + token_counters: tuple[int, ...] + salts: tuple[int, ...] + unseeded_rng_states: tuple[object, ...] + last_absolute_positions: tuple[int | None, ...] + current_device_seeds: tuple[int | None, ...] + buffer_is_default: bool + + @property + def active_slots(self) -> tuple[int, ...]: + return tuple(slot for slot, active in enumerate(self.active) if active) + + +@dataclass +class SeedState: + """Mutable per-lane seed state owned by the runtime caller. + + The parallel lists are slot indexed. ``request_seeds`` deliberately uses + ``None`` for an unseeded request, so ``active`` separately distinguishes an + admitted unseeded request from an empty slot. + """ + + active: list[bool] + request_seeds: list[int | None] + token_counters: list[int] + salts: list[int] + unseeded_rngs: list[random.Random] + last_absolute_positions: list[int | None] + current_device_seeds: list[int | None] + buffer_is_default: bool = True + + def __post_init__(self) -> None: + lengths = { + len(self.active), + len(self.request_seeds), + len(self.token_counters), + len(self.salts), + len(self.unseeded_rngs), + len(self.last_absolute_positions), + len(self.current_device_seeds), + } + if len(lengths) != 1: + raise ValueError("all SeedState slot fields must have the same capacity") + if any(counter < 0 for counter in self.token_counters): + raise ValueError("SeedState token counters must be nonnegative") + if any(salt < 0 for salt in self.salts): + raise ValueError("SeedState salts must be nonnegative") + + @classmethod + def create( + cls, + capacity: int, + *, + entropy_factory: Callable[[int], int] = secrets.randbits, + ) -> "SeedState": + """Create empty caller-owned state with independent unseeded RNGs.""" + + capacity = int(capacity) + if capacity <= 0: + raise ValueError("SeedState capacity must be positive") + return cls( + active=[False] * capacity, + request_seeds=[None] * capacity, + token_counters=[0] * capacity, + salts=[0] * capacity, + unseeded_rngs=[random.Random(entropy_factory(64)) for _ in range(capacity)], + last_absolute_positions=[None] * capacity, + current_device_seeds=[None] * capacity, + ) + + @property + def capacity(self) -> int: + return len(self.active) + + def snapshot(self) -> SeedStateSnapshot: + """Return an immutable view without exposing mutable RNG objects.""" + + return SeedStateSnapshot( + active=tuple(self.active), + request_seeds=tuple(self.request_seeds), + token_counters=tuple(self.token_counters), + salts=tuple(self.salts), + unseeded_rng_states=tuple(rng.getstate() for rng in self.unseeded_rngs), + last_absolute_positions=tuple(self.last_absolute_positions), + current_device_seeds=tuple(self.current_device_seeds), + buffer_is_default=bool(self.buffer_is_default), + ) + + +class SeedManager1D: + """Update a ``Sampling1DConfig`` seed buffer from caller-owned state. + + The manager retains only the borrowed buffer, its construction defaults, + and an entropy provider. It never retains a ``SeedState``. + """ + + def __init__( + self, + sampling_config: "Sampling1DConfig", + *, + entropy_factory: Callable[[int], int] = secrets.randbits, + salt_duplicate_seeds: bool = True, + ) -> None: + seed_buffer = getattr(sampling_config, "seeds", None) + if ( + seed_buffer is None + or not callable(getattr(seed_buffer, "update", None)) + or not callable(getattr(seed_buffer, "get_device_buffer", None)) + ): + raise TypeError("Sampling1DConfig.seeds must be a mutable LazyBuffer-compatible object") + source = getattr(seed_buffer, "source", None) + if not isinstance(source, torch.Tensor): + raise TypeError("Sampling1DConfig.seeds.source must be a torch.Tensor") + + capacity = int(getattr(sampling_config, "max_batch_size", 0)) + if capacity <= 0: + raise ValueError("Sampling1DConfig.max_batch_size must be positive") + if source.numel() != capacity: + raise ValueError( + "Sampling1DConfig seed-buffer source size does not match max_batch_size " + f"({source.numel()} != {capacity})" + ) + if not isinstance(salt_duplicate_seeds, bool): + raise TypeError("salt_duplicate_seeds must be bool") + + # When False, concurrent slots sharing a request seed keep salt 0, so two + # independent requests carrying the same seed stay bit-identical (the + # OpenAI/vLLM reproducibility contract). vLLM v1 already assigns child i + # ``seed + i``, so n>1 completions never arrive here sharing a seed. + # Demo callers that intentionally replicate one seed across slots keep + # the default and remain salted. + self.salt_duplicate_seeds = salt_duplicate_seeds + self._seed_buffer = seed_buffer + self._default_source = source.detach().clone() + self._default_values = tuple(int(value) for value in self._default_source.reshape(-1).tolist()) + self._capacity = capacity + self._entropy_factory = entropy_factory + + @property + def max_batch_size(self) -> int: + return self._capacity + + @property + def default_values(self) -> tuple[int, ...]: + return self._default_values + + @property + def seed_buffer(self) -> LazyBuffer: + """Return the borrowed persistent buffer specification.""" + + return self._seed_buffer + + def create_state(self) -> SeedState: + """Convenience constructor; ownership of the returned state is the caller's.""" + + return SeedState.create(self._capacity, entropy_factory=self._entropy_factory) + + def get_seed_device_buffer(self): + """Return the stable device handle captured by eager/trace execution.""" + + return self._seed_buffer.get_device_buffer() + + def admit( + self, + state: SeedState, + request_seeds: int | None | Sequence[int | None] | torch.Tensor, + slots: Iterable[int], + ) -> None: + """Register a simultaneous prefill admission in request order. + + Admission always starts a new request stream, even if a target slot + previously held the same integer seed. When salting is enabled, + equal-seed salts are allocated after all target slots have been + cleared, so the simultaneous admission set and surviving live requests + determine collision-free salts. When ``salt_duplicate_seeds`` is + False, every explicit seed keeps salt 0. + """ + + self._validate_state(state) + normalized_slots = self._normalize_slots(slots, label="admission slot") + seeds = self._request_ordered_seeds(request_seeds, len(normalized_slots)) + + for slot in normalized_slots: + self._clear_slot(state, slot) + for slot, seed in zip(normalized_slots, seeds): + self._register_slot(state, slot, seed) + + def synchronize( + self, + state: SeedState, + slot_seeds: int | None | Sequence[int | None] | torch.Tensor, + active_slots: Iterable[int], + *, + reset_batch: bool, + ) -> None: + """Synchronize slot-indexed decode state at a batch boundary. + + With ``reset_batch=False`` every active slot must already be registered + with the same request seed; the call may only remove completed slots. + With ``reset_batch=True`` new or changed slots are admitted, while an + unchanged running request preserves its counter, salt, RNG, and current + seed. This prevents an unrelated admission from splicing a survivor + onto a finished equal-seed sibling's stream. + """ + + self._validate_state(state) + active = self._normalize_slots(active_slots, label="active slot") + desired = {slot: self._slot_indexed_seed(slot_seeds, slot) for slot in active} + changed = [slot for slot in active if not state.active[slot] or state.request_seeds[slot] != desired[slot]] + if changed and not reset_batch: + raise RuntimeError( + "new or changed active seed slots require reset_batch=True or an explicit admit() call: " f"{changed}" + ) + + live = set(active) + removed = [slot for slot in range(self._capacity) if state.active[slot] and slot not in live] + for slot in removed: + self._clear_slot(state, slot) + + # Clear all replacements before assigning any salt, making allocation + # depend on the complete surviving/admission set rather than mutation order. + for slot in changed: + self._clear_slot(state, slot) + for slot in active: + if slot in changed: + self._register_slot(state, slot, desired[slot]) + + if removed or changed: + self._write_current_values(state) + + def refresh( + self, + state: SeedState, + active_slots: Iterable[int], + *, + positions=None, + ) -> tuple[int, ...]: + """Advance active streams and refresh the persistent seed buffer. + + ``positions`` is slot-indexed and denotes the zero-based absolute decode + position. An explicit request seed is hashed with ``position + 1``. + Repeating a refresh for the same absolute position is idempotent and + does not double-advance counters. Unseeded rows reuse their cached draw + at the same position; at a new position their caller-owned RNG advances + exactly once. + + If ``positions`` is omitted, each call consumes the slot's next + sequential token counter. + """ + + self._validate_state(state) + active = self._normalize_slots(active_slots, label="active slot") + for slot in active: + if not state.active[slot]: + raise RuntimeError(f"seed slot {slot} must be admitted or synchronized before refresh") + + values = list(self._default_values) + for slot in active: + position = self._position_for_slot(positions, slot) + if positions is not None and position is None: + raise ValueError(f"absolute positions do not cover active seed slot {slot}") + request_seed = state.request_seeds[slot] + if request_seed is not None: + if position is None: + counter = state.token_counters[slot] + state.token_counters[slot] = counter + 1 + state.last_absolute_positions[slot] = None + else: + if position < 0: + raise ValueError("active seed positions must be nonnegative") + counter = position + 1 + # Absolute-position refresh is deliberately idempotent. + state.token_counters[slot] = counter + 1 + state.last_absolute_positions[slot] = position + device_seed = _hash_request_seed_to_device_seed(request_seed, counter, state.salts[slot]) + else: + if position is not None and position < 0: + raise ValueError("active seed positions must be nonnegative") + repeated_position = ( + position is not None + and state.last_absolute_positions[slot] == position + and state.current_device_seeds[slot] is not None + ) + if repeated_position: + device_seed = state.current_device_seeds[slot] + else: + device_seed = state.unseeded_rngs[slot].randint(1, DEVICE_SEED_MAX) + state.token_counters[slot] += 1 + state.last_absolute_positions[slot] = position + + state.current_device_seeds[slot] = device_seed + values[slot] = device_seed + + self._write_values(state, values) + return tuple(values) + + def refresh_prefill_replicated( + self, + state: SeedState, + slot: int, + *, + position: int | None = None, + ) -> int: + """Advance one admitted prefill stream and replicate its device seed. + + Single-request prefill logits may occupy a sequence-tile row rather + than the request's persistent decode slot. Replicating the one draw + makes every physical sampling row observe that request stream while + keeping the caller-owned counter attached to its real slot. + """ + + slot = self._validate_slot(slot, label="prefill slot") + positions = None if position is None else {slot: int(position)} + values = self.refresh(state, (slot,), positions=positions) + device_seed = int(values[slot]) + self._write_values(state, (device_seed,) * self._capacity) + return device_seed + + def apply_slot_remap(self, state: SeedState, remap: Sequence[int] | torch.Tensor) -> None: + """Move complete seeded and unseeded request state during compaction. + + ``remap[new_slot] = old_slot``. Identity entries are no-ops. A moved + source that is not also a move destination is vacated. + """ + + self._validate_state(state) + flat = remap.reshape(-1).tolist() if isinstance(remap, torch.Tensor) else list(remap) + if len(flat) != self._capacity: + raise ValueError(f"slot remap must contain {self._capacity} entries") + normalized = [int(slot) for slot in flat] + if any(slot < 0 or slot >= self._capacity for slot in normalized): + raise ValueError("slot remap source is outside the seed-state capacity") + moves = [(old_slot, new_slot) for new_slot, old_slot in enumerate(normalized) if old_slot != new_slot] + if not moves: + return + moved_sources_list = [old_slot for old_slot, _ in moves] + if len(set(moved_sources_list)) != len(moved_sources_list): + raise ValueError("slot remap cannot copy one seed stream into multiple destinations") + + snapshots = [ + self._checkpoint_slot(state, slot) if state.active[slot] else None for slot in range(self._capacity) + ] + moved_sources = {old_slot for old_slot, _ in moves} + moved_destinations = {new_slot for _, new_slot in moves} + for old_slot, new_slot in moves: + checkpoint = snapshots[old_slot] + if checkpoint is None: + self._clear_slot(state, new_slot) + else: + self._restore_checkpoint(state, new_slot, checkpoint) + for old_slot in moved_sources - moved_destinations: + self._clear_slot(state, old_slot) + self._write_current_values(state) + + def suspend(self, state: SeedState, slot: int) -> SeedSlotState: + """Detach and return an immutable checkpoint for a preempted request.""" + + self._validate_state(state) + slot = self._validate_slot(slot, label="suspend slot") + if not state.active[slot]: + raise RuntimeError(f"cannot suspend inactive seed slot {slot}") + checkpoint = self._checkpoint_slot(state, slot) + self._clear_slot(state, slot) + self._write_current_values(state) + return checkpoint + + def resume(self, state: SeedState, slot: int, checkpoint: SeedSlotState) -> None: + """Restore a checkpoint into an empty slot without changing its stream.""" + + self._validate_state(state) + slot = self._validate_slot(slot, label="resume slot") + if state.active[slot]: + raise RuntimeError(f"cannot resume into active seed slot {slot}") + if self.salt_duplicate_seeds and checkpoint.request_seed is not None: + collision = any( + other != slot + and state.active[other] + and state.request_seeds[other] == checkpoint.request_seed + and state.salts[other] == checkpoint.salt + for other in range(self._capacity) + ) + if collision: + raise RuntimeError("cannot resume an equal-seed checkpoint after its salt was reused") + self._restore_checkpoint(state, slot, checkpoint) + self._write_current_values(state) + + def cleanup(self, state: SeedState, live_slots: Iterable[int]) -> None: + """Remove completed/inactive slot state and clear stale buffer rows.""" + + self._validate_state(state) + live = set(self._normalize_slots(live_slots, label="live slot")) + removed = [slot for slot in range(self._capacity) if state.active[slot] and slot not in live] + if not removed: + return + for slot in removed: + self._clear_slot(state, slot) + self._write_current_values(state) + + def restore_defaults(self, state: SeedState) -> None: + """Restore construction defaults without discarding request streams.""" + + self._validate_state(state) + self._write_values(state, self._default_values) + + def reset(self, state: SeedState) -> None: + """Clear every request stream and restore construction defaults.""" + + self._validate_state(state) + for slot in range(self._capacity): + self._clear_slot(state, slot, reseed_unseeded=True) + self._write_values(state, self._default_values) + + # Internal state operations ------------------------------------------------- + + def _validate_state(self, state: SeedState) -> None: + if not isinstance(state, SeedState): + raise TypeError("state must be a caller-owned SeedState") + if state.capacity != self._capacity: + raise ValueError( + f"SeedState capacity {state.capacity} does not match seed-buffer capacity {self._capacity}" + ) + + def _validate_slot(self, slot: int, *, label: str) -> int: + slot = int(slot) + if slot < 0 or slot >= self._capacity: + raise ValueError(f"{label} {slot} is outside the seed-state capacity") + return slot + + def _normalize_slots(self, slots: Iterable[int], *, label: str) -> tuple[int, ...]: + normalized = tuple(self._validate_slot(slot, label=label) for slot in slots) + if len(set(normalized)) != len(normalized): + raise ValueError(f"{label}s must be unique") + return normalized + + def _request_ordered_seeds(self, seeds, count: int) -> tuple[int | None, ...]: + if isinstance(seeds, torch.Tensor): + if seeds.ndim == 0: + seeds = seeds.item() + else: + seeds = seeds.reshape(-1).tolist() + if isinstance(seeds, (list, tuple)): + if len(seeds) != count: + raise ValueError(f"expected {count} request seeds, got {len(seeds)}") + return tuple(self._normalize_seed(seed) for seed in seeds) + return tuple(self._normalize_seed(seeds) for _ in range(count)) + + def _slot_indexed_seed(self, seeds, slot: int) -> int | None: + if seeds is None: + return None + if isinstance(seeds, torch.Tensor): + if seeds.ndim == 0: + return self._normalize_seed(seeds.item()) + flat = seeds.reshape(-1) + return None if slot >= flat.numel() else self._normalize_seed(flat[slot].item()) + if isinstance(seeds, (list, tuple)): + return None if slot >= len(seeds) else self._normalize_seed(seeds[slot]) + return self._normalize_seed(seeds) + + @staticmethod + def _normalize_seed(seed) -> int | None: + if seed is None: + return None + if isinstance(seed, torch.Tensor): + if seed.numel() != 1: + raise ValueError("each request seed tensor must contain exactly one value") + seed = seed.item() + return int(seed) + + @staticmethod + def _position_for_slot(positions, slot: int) -> int | None: + if positions is None: + return None + if isinstance(positions, dict): + value = positions.get(slot) + return None if value is None else int(value) + if isinstance(positions, torch.Tensor): + if positions.ndim == 0: + return int(positions.item()) + flat = positions.reshape(-1) + return None if slot >= flat.numel() else int(flat[slot].item()) + if isinstance(positions, (list, tuple)): + if slot >= len(positions) or positions[slot] is None: + return None + return int(positions[slot]) + return int(positions) + + def _next_free_salt(self, state: SeedState, slot: int, seed: int) -> int: + if not self.salt_duplicate_seeds: + return 0 + taken = { + state.salts[other] + for other in range(self._capacity) + if other != slot and state.active[other] and state.request_seeds[other] == seed + } + salt = 0 + while salt in taken: + salt += 1 + return salt + + def _register_slot(self, state: SeedState, slot: int, seed: int | None) -> None: + state.active[slot] = True + state.request_seeds[slot] = seed + state.token_counters[slot] = 0 + state.salts[slot] = 0 if seed is None else self._next_free_salt(state, slot, seed) + state.unseeded_rngs[slot].seed(self._entropy_factory(64) if seed is None else int(seed)) + state.last_absolute_positions[slot] = None + state.current_device_seeds[slot] = None + + def _clear_slot(self, state: SeedState, slot: int, *, reseed_unseeded: bool = False) -> None: + state.active[slot] = False + state.request_seeds[slot] = None + state.token_counters[slot] = 0 + state.salts[slot] = 0 + if reseed_unseeded: + state.unseeded_rngs[slot].seed(self._entropy_factory(64)) + state.last_absolute_positions[slot] = None + state.current_device_seeds[slot] = None + + @staticmethod + def _checkpoint_slot(state: SeedState, slot: int) -> SeedSlotState: + return SeedSlotState( + request_seed=state.request_seeds[slot], + token_counter=state.token_counters[slot], + salt=state.salts[slot], + unseeded_rng_state=copy.deepcopy(state.unseeded_rngs[slot].getstate()), + last_absolute_position=state.last_absolute_positions[slot], + current_device_seed=state.current_device_seeds[slot], + ) + + @staticmethod + def _restore_checkpoint(state: SeedState, slot: int, checkpoint: SeedSlotState) -> None: + state.active[slot] = True + state.request_seeds[slot] = checkpoint.request_seed + state.token_counters[slot] = checkpoint.token_counter + state.salts[slot] = checkpoint.salt + rng = random.Random() + rng.setstate(checkpoint.unseeded_rng_state) + state.unseeded_rngs[slot] = rng + state.last_absolute_positions[slot] = checkpoint.last_absolute_position + state.current_device_seeds[slot] = checkpoint.current_device_seed + + def _write_current_values(self, state: SeedState) -> None: + values = list(self._default_values) + for slot in range(self._capacity): + if state.active[slot] and state.current_device_seeds[slot] is not None: + values[slot] = state.current_device_seeds[slot] + self._write_values(state, values) + + def _write_values(self, state: SeedState, values: Sequence[int]) -> None: + if len(values) != self._capacity: + raise ValueError(f"expected {self._capacity} seed values, got {len(values)}") + try: + normalized = [int(value) for value in values] + except (TypeError, ValueError) as error: + raise ValueError("device seed values must be integer-like") from error + if any(value < 0 or value > MAX_UINT32 for value in normalized): + raise ValueError("device seed values must be in [0, UINT32_MAX]") + + # Materialize before update. LazyBuffer.update() only replaces the + # future source when unmaterialized, but request values must reach the + # stable handle captured by Sampling1D eager/trace execution now. + self.get_seed_device_buffer() + source = torch.tensor( + normalized, + dtype=self._default_source.dtype, + device=self._default_source.device, + ).reshape(self._default_source.shape) + try: + self._seed_buffer.update(source) + finally: + # Request state must never become the construction source used after + # Sampling1D.release() and later rematerialization. + self._seed_buffer.source = self._default_source.detach().clone() + state.buffer_is_default = tuple(normalized) == self._default_values