# 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