# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from typing import TYPE_CHECKING import torch from loguru import logger import ttnn from models.common.utility_functions import is_blackhole from ..layers.linear import ColParallelLinear from ..layers.module import Module, Parameter, UnregisteredModule from ..layers.normalization import DistributedRMSNorm from ..utils.matmul import get_matmul_config, get_matmul_core_grid from ..utils.mochi import get_rot_transformation_mat from ..utils.padding import PaddingConfig, pad_weight_tensor from ..utils.substate import pop_substate if TYPE_CHECKING: pass from ..parallel.config import DiTParallelConfig from ..parallel.manager import CCLManager # Best compute core grid per shape from flux2 dual-grid sweep (2026-05-27). # See flux2_dual_grid_mm_sweep.md "Best grid per shape". _FLUX2_MATMUL_CORE_GRIDS: dict[tuple[int, int, int], tuple[int, int]] = { (1024, 128, 768): ttnn.CoreCoord(11, 10), (512, 15360, 768): ttnn.CoreCoord(12, 9), (1024, 6144, 768): ttnn.CoreCoord(11, 10), (512, 6144, 768): ttnn.CoreCoord(12, 9), (1024, 6144, 4608): ttnn.CoreCoord(12, 9), (512, 6144, 4608): ttnn.CoreCoord(12, 9), (1024, 6144, 2304): ttnn.CoreCoord(12, 9), (512, 6144, 2304): ttnn.CoreCoord(12, 9), } # adapted from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/models/attention_processor.py class Attention(Module): # SDPA chunk sizes keyed by (is_blackhole, sp_factor, tp_factor). Resolution priority for # non-ring: caller overrides > class map > constructor args > class default. sdpa_chunk_size_map: dict[tuple, tuple[int, int]] = {} default_sdpa_chunk_size: tuple[int, int] = (128, 512) ring_sdpa_chunk_size_map: dict[tuple, tuple[int, int]] = { # -1 is the default resolution. (False, 2, 4): {-1: (256, 256)}, (False, 8, 4): {-1: (256, 256)}, (True, 2, 2): {-1: (128, 512)}, (True, 4, 8): { -1: (256, 512), # default 4096: (128, 256), # 1024×1024 — 18.0% util (ring_sdpa_sweep.md) 4608: (128, 256), # 1024×1024 — 18.0% util (ring_sdpa_sweep.md) 4096 * 4: (256, 512), # 2048×2048 — 55.4% util (ring_sdpa_sweep_2048.md) (4096 + 128) * 4: (256, 512), # 2048×2048 — 55.4% util (ring_sdpa_sweep_2048.md) 4096 * 4: (256, 512), # 2048×2048 — 55.4% util (ring_sdpa_sweep_2048.md) 4096 * 16: (192, 512), # 4096×4096 — 67.4% util (ring_sdpa_sweep_4096.md) }, (True, 8, 4): { -1: (256, 512), # default 16384: (320, 384), # 2048×2048 — 43.3% util (ring_sdpa_sweep_8x4_2048.md) 65536: (256, 512), # 4096×4096 — 65.7% util (ring_sdpa_sweep_8x4_4096.md) }, } default_ring_sdpa_chunk_size: tuple[int, int] = {-1: (256, 256)} def __init__( self, *, query_dim: int, head_dim: int, heads: int, out_dim: int, added_kv_proj_dim: int, context_pre_only: bool = False, pre_only: bool = False, use_spatial_weights_for_prompt: bool = False, context_head_scaling: bool = False, proj_bias: bool = True, eps: float, mesh_device: ttnn.MeshDevice, ccl_manager: CCLManager | None, parallel_config: DiTParallelConfig, padding_config: PaddingConfig | None, k_chunk_size: int | None = None, q_chunk_size: int | None = None, sdpa_chunk_size_overrides: dict | None = None, per_head_norm: bool = False, is_fsdp: bool = False, shard_prompt: bool = False, ) -> None: super().__init__() self.head_dim = head_dim self.pre_only = pre_only self.mesh_device = mesh_device self.ccl_manager = ccl_manager self.parallel_config = parallel_config self.shard_prompt = shard_prompt self.padding_config = padding_config self.use_spatial_weights_for_prompt = use_spatial_weights_for_prompt # True = per-head RMS (head_dim divisor, stats stay local); False = global RMS over # padded_inner_dim (AG of stats). self.per_head_norm = per_head_norm self.padded_heads = padding_config.target_heads if padding_config is not None else heads self.n_local_heads = self.padded_heads // self.parallel_config.tensor_parallel.factor tp_axis = parallel_config.tensor_parallel.mesh_axis padded_inner_dim = head_dim * self.padded_heads # FSDP: shard weights on sequence parallel axis to reduce memory fsdp_mesh_axis = parallel_config.sequence_parallel.mesh_axis if is_fsdp else None common_args = dict(mesh_device=mesh_device, ccl_manager=ccl_manager, fsdp_mesh_axis=fsdp_mesh_axis) full_grid = self.mesh_device.compute_with_storage_grid_size() # Reserve last row for CCL. self.sdpa_worker_grid = (full_grid.x, full_grid.y - 1) chunk_lookup = {**self.sdpa_chunk_size_map, **(sdpa_chunk_size_overrides or {})} resolved_q_chunk, resolved_k_chunk = chunk_lookup.get( ( is_blackhole(), parallel_config.sequence_parallel.factor, parallel_config.tensor_parallel.factor, ), ( q_chunk_size if q_chunk_size is not None else self.default_sdpa_chunk_size[0], k_chunk_size if k_chunk_size is not None else self.default_sdpa_chunk_size[1], ), ) self.sdpa_program_config = ttnn.SDPAProgramConfig( compute_with_storage_grid_size=self.sdpa_worker_grid, q_chunk_size=resolved_q_chunk, k_chunk_size=resolved_k_chunk, exp_approx_mode=False, # NOTE: False is more correct ) self.ring_sdpa_worker_grid = None # (full_grid.x, full_grid.y - 5) self.ring_sdpa_program_config = {} self.sdpa_compute_kernel_config = ttnn.init_device_compute_kernel_config( self.mesh_device.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, math_approx_mode=False, fp32_dest_acc_en=False, # NOTE: Set to True if there's a correctness issue ) self.mm_compute_kernel_config = ttnn.init_device_compute_kernel_config( self.mesh_device.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, math_approx_mode=True, fp32_dest_acc_en=True, packer_l1_acc=True, ) # chunks=3 splits QKV inside minimal_matmul_split (faster than matmul + ttnn.chunk). self.to_qkv = ColParallelLinear( query_dim, 3 * padded_inner_dim, bias=proj_bias, mesh_axis=tp_axis, chunks=3, **common_args ) # Fuses head-split + RMSNorm + RoPE. Weight is tiled per-head → [padded_inner_dim] in # _prepare_torch_state. self.norm_q = DistributedRMSNorm( embedding_dim=padded_inner_dim, norm_eps=eps, mesh_axis=tp_axis, mesh_device=mesh_device, ccl_manager=ccl_manager, ) self.norm_k = DistributedRMSNorm( embedding_dim=padded_inner_dim, norm_eps=eps, mesh_axis=tp_axis, mesh_device=mesh_device, ccl_manager=ccl_manager, ) self.to_out = ( ColParallelLinear(padded_inner_dim, out_dim, bias=proj_bias, mesh_axis=tp_axis, **common_args) if not self.pre_only else None ) if use_spatial_weights_for_prompt: self.add_qkv_proj = UnregisteredModule(self.to_qkv) self.norm_added_q = UnregisteredModule(self.norm_q) self.norm_added_k = UnregisteredModule(self.norm_k) self.to_add_out = UnregisteredModule(self.to_out) if self.to_out is not None else None elif added_kv_proj_dim > 0: self.add_qkv_proj = ColParallelLinear( added_kv_proj_dim, 3 * padded_inner_dim, bias=proj_bias, mesh_axis=tp_axis, chunks=3, **common_args ) self.norm_added_q = DistributedRMSNorm( embedding_dim=padded_inner_dim, norm_eps=eps, mesh_axis=tp_axis, mesh_device=mesh_device, ccl_manager=ccl_manager, ) self.norm_added_k = DistributedRMSNorm( embedding_dim=padded_inner_dim, norm_eps=eps, mesh_axis=tp_axis, mesh_device=mesh_device, ccl_manager=ccl_manager, ) self.to_add_out = ( ColParallelLinear(padded_inner_dim, out_dim, bias=proj_bias, mesh_axis=tp_axis, **common_args) if not context_pre_only else None ) else: self.add_qkv_proj = None self.norm_added_q = None self.norm_added_k = None self.to_add_out = None self.context_head_factors = ( Parameter(total_shape=[self.padded_heads, 1, 1], device=mesh_device, mesh_axes=[tp_axis, None, None]) if context_head_scaling and self.add_qkv_proj is not None else None ) assert not (use_spatial_weights_for_prompt and context_head_scaling), ( "context_head_scaling is not supported with use_spatial_weights_for_prompt=True " "(concatenated path keeps q/k/v combined; per-head scaling on add_q is never applied)" ) # Empty joint input for ring SDPA when there is no prompt stream. self.dummy_joint_input = ttnn.zeros( [1, self.n_local_heads, 0, self.head_dim], device=mesh_device, layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16, ) self.trans_mat = ttnn.from_torch( get_rot_transformation_mat(), device=mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, ) @staticmethod def get_core_grid(M, full_grid) -> ttnn.CoreCoord: if M <= 1152: return ttnn.CoreCoord(full_grid.x, full_grid.y - 2) else: return ttnn.CoreCoord(full_grid.x, full_grid.y - 1) def get_ring_sdpa_program_config(self, per_device_seq_len: int) -> ttnn.SDPAProgramConfig: if per_device_seq_len not in self.ring_sdpa_program_config: device_parallel_key = ( is_blackhole(), self.parallel_config.sequence_parallel.factor, self.parallel_config.tensor_parallel.factor, ) ring_chunk_size = self.ring_sdpa_chunk_size_map.get(device_parallel_key, self.default_ring_sdpa_chunk_size) if per_device_seq_len in ring_chunk_size: ring_chunk_size = ring_chunk_size[per_device_seq_len] else: logger.warning( f"No ring SDPA chunk size found for resolution {per_device_seq_len}, using default {ring_chunk_size[-1]}" ) ring_chunk_size = ring_chunk_size[-1] self.ring_sdpa_program_config[per_device_seq_len] = ttnn.SDPAProgramConfig( compute_with_storage_grid_size=self.ring_sdpa_worker_grid, q_chunk_size=ring_chunk_size[0], k_chunk_size=ring_chunk_size[1], exp_approx_mode=False, # NOTE: False is more correct ) return self.ring_sdpa_program_config[per_device_seq_len] def get_ring_sdpa_core_grid(self, resolved_seq_len: int) -> ttnn.CoreCoord: full_grid = self.mesh_device.compute_with_storage_grid_size() if resolved_seq_len == 4096: return (full_grid.x, full_grid.y - 5) elif resolved_seq_len == 4608: return (full_grid.x, full_grid.y - 5) elif resolved_seq_len == 4096 * 4: return (full_grid.x, full_grid.y - 1) elif resolved_seq_len == (4096 + 128) * 4: return (full_grid.x, full_grid.y - 1) elif resolved_seq_len == 4096 * 16: return (full_grid.x, full_grid.y - 1) else: return (full_grid.x, full_grid.y - 1) def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None: weight, bias = self._reshape_and_merge_qkv( pop_substate(state, "to_q"), pop_substate(state, "to_k"), pop_substate(state, "to_v"), ) if weight is not None: state["to_qkv.weight"] = weight if bias is not None: state["to_qkv.bias"] = bias weight, bias = self._reshape_and_merge_qkv( pop_substate(state, "add_q_proj"), pop_substate(state, "add_k_proj"), pop_substate(state, "add_v_proj"), ) if weight is not None: state["add_qkv_proj.weight"] = weight if bias is not None: state["add_qkv_proj.bias"] = bias if "to_out.0.weight" in state: state["to_out.weight"] = state.pop("to_out.0.weight") if "to_out.0.bias" in state: state["to_out.bias"] = state.pop("to_out.0.bias") if self.padding_config is not None: if "to_out.weight" in state: weight = state["to_out.weight"].T weight = pad_weight_tensor(weight, self.padding_config, pad_input_dim=True) state["to_out.weight"] = weight.T if "to_add_out.weight" in state: weight = state["to_add_out.weight"].T weight = pad_weight_tensor(weight, self.padding_config, pad_input_dim=True) state["to_add_out.weight"] = weight.T if "context_head_factors" in state: factors = state["context_head_factors"] if self.padding_config is not None: pad = (0, self.padding_config.head_padding) factors = torch.nn.functional.pad(factors, pad) state["context_head_factors"] = factors.reshape([-1, 1, 1]) # HF stores per-head weight [head_dim]; DistributedRMSNorm expects [padded_inner_dim]. for key in ["norm_q.weight", "norm_k.weight", "norm_added_q.weight", "norm_added_k.weight"]: if key in state and state[key].shape[0] == self.head_dim: state[key] = state[key].repeat(self.padded_heads) def _reshape_and_merge_qkv( self, q_state: dict[str, torch.Tensor], k_state: dict[str, torch.Tensor], v_state: dict[str, torch.Tensor], ) -> tuple[torch.Tensor | None, torch.Tensor | None]: # Rearrange QKV projections such column-fracturing shards the heads def _merge_tensors(q, k, v): n_dev = self.parallel_config.tensor_parallel.factor q, k, v = q.T, k.T, v.T # Pad QKV weights and biases to match the padded heads if self.padding_config is not None: q = pad_weight_tensor(q, self.padding_config, pad_output_dim=True) k = pad_weight_tensor(k, self.padding_config, pad_output_dim=True) v = pad_weight_tensor(v, self.padding_config, pad_output_dim=True) q = q.reshape(q.shape[0], n_dev, self.n_local_heads, self.head_dim) k = k.reshape(k.shape[0], n_dev, self.n_local_heads, self.head_dim) v = v.reshape(v.shape[0], n_dev, self.n_local_heads, self.head_dim) qkv = torch.cat([q, k, v], dim=2) qkv = qkv.reshape(qkv.shape[0], 3 * self.padded_heads * self.head_dim) qkv = qkv.T return qkv if "weight" in q_state and "weight" in k_state and "weight" in v_state: weight = _merge_tensors(q_state["weight"], k_state["weight"], v_state["weight"]) else: weight = None if "bias" in q_state and "bias" in k_state and "bias" in v_state: bias = _merge_tensors( q_state["bias"].unsqueeze(-1), k_state["bias"].unsqueeze(-1), v_state["bias"].unsqueeze(-1) ) bias = bias.squeeze(-1) else: bias = None return weight, bias def _to_out_fused_addcmul( self, x: ttnn.Tensor, addcmul_residual: ttnn.Tensor | None, addcmul_gate: ttnn.Tensor | None, to_out_module: ColParallelLinear, parallel_config, tp_axis: int | None = None, ) -> ttnn.Tensor: """Fused projection + addcmul: output = residual + proj(x) * gate. Ring (parallel_config, tp>1): all_gather_minimal_matmul_async fuses AG+matmul+addcmul. Linear (parallel_config=None): gathers x on tp_axis if K-dim mismatches, then dit_minimal_matmul_addcmul_fused. """ if ( to_out_module.fsdp_mesh_axis is not None and to_out_module.mesh_device.shape[to_out_module.fsdp_mesh_axis] > 1 ): unsqueezed_weight = ttnn.unsqueeze_to_4D(to_out_module.weight.data) weight = self.ccl_manager.all_gather_persistent_buffer( unsqueezed_weight, dim=2, mesh_axis=to_out_module.fsdp_mesh_axis ) weight = ttnn.reshape(weight, (weight.shape[-2], weight.shape[-1])) else: weight = to_out_module.weight.data bias = to_out_module.bias.data if to_out_module.bias is not None else None if parallel_config is not None and parallel_config.tensor_parallel.factor > 1: M = x.padded_shape[-2] K = weight.padded_shape[-2] N = weight.padded_shape[-1] full_grid = self.mesh_device.compute_with_storage_grid_size() core_grid = self.get_core_grid(M, full_grid) # core_grid = _FLUX2_MATMUL_CORE_GRIDS.get((M, K, N)) or ttnn.CoreCoord(full_grid.x, full_grid.y - 1) # print(f"Using core grid {core_grid} for matmul config (M, K, N) = ({M}, {K}, {N})", flush=True) matmul_config = get_matmul_config( M, K, N, core_grid, use_heuristic=True, num_k_shards=self.parallel_config.tensor_parallel.factor ) ag_persistent_buffer = self.ccl_manager.get_ag_ping_pong_buffer( x.shape, -1, parallel_config.tensor_parallel.mesh_axis, dtype=x.get_dtype() ) ag_global_semaphores = self.ccl_manager.get_ag_ping_pong_semaphore( parallel_config.tensor_parallel.mesh_axis ) return ttnn.experimental.all_gather_minimal_matmul_async( input_tensor=x, weight_tensor=weight, bias_tensor=bias, config=matmul_config, compute_kernel_config=self.mm_compute_kernel_config, persistent_output_buffer=ag_persistent_buffer, multi_device_global_semaphore=ag_global_semaphores, num_links=self.ccl_manager.num_links, topology=self.ccl_manager.topology, cluster_axis=parallel_config.tensor_parallel.mesh_axis, barrier_semaphore=None, force_transpose=True, num_workers_per_link=full_grid.x // self.ccl_manager.num_links, num_buffers_per_channel=48 if not is_blackhole() else 24, scalar=1.0, addcmul_input_tensor1=addcmul_residual, addcmul_input_tensor2=addcmul_gate, )[0] else: if tp_axis is not None and x.padded_shape[-1] != weight.padded_shape[-2]: x = self.ccl_manager.all_gather_persistent_buffer(x, dim=-1, mesh_axis=tp_axis, use_hyperparams=True) M = x.padded_shape[-2] K = x.padded_shape[-1] N = weight.padded_shape[-1] # core_grid = _FLUX2_MATMUL_CORE_GRIDS.get((M, K, N)) or self.mesh_device.compute_with_storage_grid_size() # core_grid = self.mesh_device.compute_with_storage_grid_size() core_grid = get_matmul_core_grid(self.mesh_device) matmul_config = get_matmul_config( M, K, N, core_grid, use_heuristic=True, num_k_shards=self.parallel_config.tensor_parallel.factor ) return ttnn.experimental.dit_minimal_matmul_addcmul_fused( x, weight, 1.0, addcmul_residual, addcmul_gate, bias_tensor=bias, config=matmul_config, compute_kernel_config=self.mm_compute_kernel_config, ) def _project_out( self, input: ttnn.Tensor, proj_op: ColParallelLinear | None, addcmul_residual: ttnn.Tensor | None, addcmul_gate: ttnn.Tensor | None, is_ring: bool, tp_axis: int, ) -> ttnn.Tensor: """Project with optional fused addcmul. Returns input unchanged if proj_op is None.""" if input is None or proj_op is None: return input return self._to_out_fused_addcmul( input, addcmul_residual, addcmul_gate, proj_op, self.parallel_config if is_ring else None, None if is_ring else tp_axis, ) def forward( self, *, sequence_1: ttnn.Tensor, sequence_1_length: int, sequence_2: ttnn.Tensor | None = None, sequence_2_length: int = 0, sequence_1_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, sequence_2_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, addcmul_spatial_residual: ttnn.Tensor | None = None, addcmul_spatial_gate: ttnn.Tensor | None = None, addcmul_prompt_residual: ttnn.Tensor | None = None, addcmul_prompt_gate: ttnn.Tensor | None = None, ) -> tuple[ttnn.Tensor, ttnn.Tensor | None]: """Forward with optional fused to_out + addcmul. Ring callers pass spatial/prompt fractured on TP (fused AG+matmul inside); Linear callers must all-gather first. addcmul_*_residual+gate fuse the out projection's residual+gate when both are provided. """ assert len(sequence_1.shape) == 3 if sequence_2 is not None: assert len(sequence_2.shape) == 3 for t in sequence_1_rope or (): assert len(t.shape) == 4, f"rope tensors must be pre-shaped to (1, 1, seq, head_dim), got {t.shape}" for t in sequence_2_rope or (): assert len(t.shape) == 4, f"rope tensors must be pre-shaped to (1, 1, seq, head_dim), got {t.shape}" sp_axis = self.parallel_config.sequence_parallel.mesh_axis shard_prompt = self.shard_prompt and self.parallel_config.sequence_parallel.factor > 1 def _split_heads(x: ttnn.Tensor) -> ttnn.Tensor: # V only — Q/K head-split is fused inside DistributedRMSNorm. out, _, _ = ttnn.experimental.nlp_create_qkv_heads( ttnn.unsqueeze(x, 0), num_heads=self.n_local_heads, num_kv_heads=0, transpose_k_heads=False ) return out # Sequence 1 processing Typically spatial or spatial+prompt. q_flat_c, k_flat_c, v_flat_c = self.to_qkv( sequence_1, compute_kernel_config=self.mm_compute_kernel_config, parallel_config=self.parallel_config, core_grid=self.get_core_grid(sequence_1.shape[-2], self.mesh_device.compute_with_storage_grid_size()), use_heuristic_mmcfg=True, ) # Combined RoPE: concat spatial and prompt tensors along the seq dim (dim=2). sequence_1_cos, sequence_1_sin = sequence_1_rope trans_mat = self.trans_mat if sequence_1_rope is not None else None # Single norm_q/norm_k/v on the full combined [spatial+prompt] shard. q = self.norm_q( ttnn.unsqueeze(q_flat_c, 0), num_heads_per_device=self.n_local_heads, rope_cos=sequence_1_cos, rope_sin=sequence_1_sin, trans_mat=trans_mat, per_head_norm=self.per_head_norm, ) k = self.norm_k( ttnn.unsqueeze(k_flat_c, 0), num_heads_per_device=self.n_local_heads, rope_cos=sequence_1_cos, rope_sin=sequence_1_sin, trans_mat=trans_mat, per_head_norm=self.per_head_norm, ) v = _split_heads(v_flat_c) # [1, H, M_s+M_p, D] # No joint stream — dummies handled in the post-processing block below. add_q = add_k = add_v = self.dummy_joint_input persistent_joint_k = persistent_joint_v = None if self.add_qkv_proj is not None and sequence_2 is not None: # Rope tensors arrive pre-shaped as (1, 1, seq, head_dim); no per-call reshape needed. prompt_cos = sequence_2_rope[0] if sequence_2_rope is not None else None prompt_sin = sequence_2_rope[1] if sequence_2_rope is not None else None prompt_trans_mat = self.trans_mat if sequence_2_rope is not None else None add_q_flat, add_k_flat, add_v_flat = self.add_qkv_proj( sequence_2, compute_kernel_config=self.mm_compute_kernel_config, parallel_config=self.parallel_config, core_grid=self.get_core_grid(sequence_2.shape[-2], self.mesh_device.compute_with_storage_grid_size()), use_heuristic_mmcfg=True, ) add_q = self.norm_added_q( ttnn.unsqueeze(add_q_flat, 0), num_heads_per_device=self.n_local_heads, rope_cos=prompt_cos, rope_sin=prompt_sin, trans_mat=prompt_trans_mat, per_head_norm=self.per_head_norm, ) add_k = self.norm_added_k( ttnn.unsqueeze(add_k_flat, 0), num_heads_per_device=self.n_local_heads, rope_cos=prompt_cos, rope_sin=prompt_sin, trans_mat=prompt_trans_mat, per_head_norm=self.per_head_norm, ) add_v = _split_heads(add_v_flat) if self.context_head_factors is not None: add_q = add_q * self.context_head_factors.data if shard_prompt: persistent_joint_k = self.ccl_manager.get_ag_ping_pong_buffer(add_k.shape, 2, sp_axis) persistent_joint_v = self.ccl_manager.get_ag_ping_pong_buffer(add_v.shape, 2, sp_axis) if self.ring_sdpa_worker_grid is None: self.ring_sdpa_worker_grid = self.get_ring_sdpa_core_grid(sequence_1_length) if self.parallel_config.sequence_parallel.factor > 1: spatial, sequence_2, _lse = ttnn.transformer.ring_joint_scaled_dot_product_attention( q, k, v, add_q, add_k, add_v, persistent_output_buffer_k=self.ccl_manager.get_ag_ping_pong_buffer( k.shape, 2, self.parallel_config.sequence_parallel.mesh_axis ), persistent_output_buffer_v=self.ccl_manager.get_ag_ping_pong_buffer( v.shape, 2, self.parallel_config.sequence_parallel.mesh_axis ), joint_strategy="rear", logical_n=sequence_1_length, logical_l=sequence_2_length, program_config=self.get_ring_sdpa_program_config(sequence_1_length), compute_kernel_config=self.sdpa_compute_kernel_config, dim=2, multi_device_global_semaphore=self.ccl_manager.get_ag_ping_pong_semaphore( self.parallel_config.sequence_parallel.mesh_axis ), num_links=self.ccl_manager.num_links, cluster_axis=self.parallel_config.sequence_parallel.mesh_axis, mesh_device=self.mesh_device, topology=self.ccl_manager.topology, subdevice_id=self.ccl_manager.ccl_sub_device_id, ccl_core_grid_offset=(0, self.ring_sdpa_worker_grid[1]), persistent_output_buffer_joint_k=persistent_joint_k, persistent_output_buffer_joint_v=persistent_joint_v, ) else: assert ( sequence_1_length == spatial.shape[1] ), "spatial sequence must not be padded without sequence parallelism" spatial, sequence_2 = ttnn.transformer.joint_scaled_dot_product_attention( q, k, v, add_q, add_k, add_v, joint_strategy="rear", program_config=self.sdpa_program_config, compute_kernel_config=self.sdpa_compute_kernel_config, ) spatial = ttnn.transformer.concatenate_heads(spatial) if self.to_out is not None: spatial = self.to_out( x=spatial, addcmul_a=addcmul_spatial_residual, addcmul_b=addcmul_spatial_gate, parallel_config=self.parallel_config, core_grid=self.get_core_grid(spatial.shape[-2], self.mesh_device.compute_with_storage_grid_size()), ) if sequence_2 is not None and sequence_2.shape[2] > 0: sequence_2 = ttnn.transformer.concatenate_heads(sequence_2) if self.to_add_out is not None: sequence_2 = self.to_add_out( x=sequence_2, addcmul_a=addcmul_prompt_residual, addcmul_b=addcmul_prompt_gate, parallel_config=self.parallel_config, core_grid=self.get_core_grid( sequence_2.shape[-2], self.mesh_device.compute_with_storage_grid_size() ), ) return spatial, sequence_2 @classmethod def spatial_sequence_padding_length(cls, *, length: int, sp_factor: int, k_chunk_size: int = 512) -> int: if sp_factor == 1: return 0 divisor = k_chunk_size * sp_factor return -length % divisor @classmethod def pad_spatial_sequence(cls, x: torch.Tensor, /, *, sp_factor: int, k_chunk_size: int = 512) -> torch.Tensor: padding_len = cls.spatial_sequence_padding_length( length=x.shape[-2], sp_factor=sp_factor, k_chunk_size=k_chunk_size ) return torch.nn.functional.pad(x, (0, 0, 0, padding_len))