code/models/tt_dit/blocks/attention.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /blocks /attention.py

Add files using upload-large-folder tool

629b123 verified 11 days ago

15 kB

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

import ttnn

from ..layers.linear import ColParallelLinear
from ..layers.module import Module, Parameter, UnregisteredModule
from ..layers.normalization import RMSNorm
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


# adapted from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/models/attention_processor.py
class Attention(Module):
    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,
        eps: float,
        mesh_device: ttnn.MeshDevice,
        ccl_manager: CCLManager | None,
        parallel_config: DiTParallelConfig,
        padding_config: PaddingConfig | None,
        k_chunk_size: int = 512,
        q_chunk_size: int = 128,
        is_fsdp: 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.padding_config = padding_config
        self.use_spatial_weights_for_prompt = use_spatial_weights_for_prompt

        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)

        self.sdpa_worker_grid = (
            self.mesh_device.compute_with_storage_grid_size().x,
            self.mesh_device.compute_with_storage_grid_size().y - 1,
        )

        self.sdpa_program_config = ttnn.SDPAProgramConfig(
            compute_with_storage_grid_size=self.sdpa_worker_grid,
            q_chunk_size=q_chunk_size,
            k_chunk_size=k_chunk_size,
            exp_approx_mode=False,  # NOTE: False is more correct
        )
        self.sdpa_compute_kernel_config = ttnn.WormholeComputeKernelConfig(
            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.to_qkv = ColParallelLinear(query_dim, 3 * padded_inner_dim, mesh_axis=tp_axis, **common_args)

        self.norm_q = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)
        self.norm_k = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)

        self.to_out = (
            ColParallelLinear(padded_inner_dim, out_dim, 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, mesh_axis=tp_axis, **common_args
            )

            self.norm_added_q = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)
            self.norm_added_k = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)

            self.to_add_out = (
                ColParallelLinear(padded_inner_dim, out_dim, 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
        )

    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])

    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 forward(
        self,
        *,
        spatial: ttnn.Tensor,
        prompt: ttnn.Tensor | None = None,
        spatial_sequence_length: int,
        spatial_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None,
        prompt_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None,
    ) -> tuple[ttnn.Tensor, ttnn.Tensor | None]:
        """Forward pass of the model.
        Args:
            spatial: Tensor with shape [batch_size, spatial_sequence_length / sp_factor, query_dim].
            prompt: Tensor with shape [batch_size, prompt_sequence_length, query_dim] (not sharded!).
            spatial_rope: Tuple of two tensors with shape [spatial_sequence_length / sp_factor, head_dim].
            prompt_rope: Tuple of two tensors with shape [prompt_sequence_length, head_dim] (not sharded!).
        """
        assert len(spatial.shape) == 3
        if prompt is not None:
            assert len(prompt.shape) == 3
        for t in spatial_rope or ():
            assert len(t.shape) == 2
        for t in prompt_rope or ():
            assert len(t.shape) == 2

        device_grid = self.mesh_device.compute_with_storage_grid_size()
        core_grid = ttnn.CoreGrid(x=device_grid.x, y=device_grid.y)
        tp_axis = self.parallel_config.tensor_parallel.mesh_axis

        qkv = self.to_qkv(
            spatial
        )  # [batch_size, spatial_sequence_length / sp_factor, 3 * n_local_heads * head_dim (in this order)]
        local_heads = self.n_local_heads
        # Explicit memory_config: forward the input qkv layout to the inner reshapes so the
        # Q/K/V outputs keep the spec the downstream SDPA op is tuned for (matches the
        # layout main was implicitly producing pre-#42543).
        q, k, v = ttnn.transformer.split_query_key_value_and_split_heads(
            qkv, num_heads=local_heads, transpose_key=False, memory_config=qkv.memory_config()
        )  # [batch_size, n_local_heads, spatial_sequence_length / sp_factor, head_dim]

        q = self.norm_q(q)
        k = self.norm_k(k)

        if spatial_rope is not None:
            q = _apply_rope(q, spatial_rope)
            k = _apply_rope(k, spatial_rope)

        if self.add_qkv_proj is not None:
            add_qkv = self.add_qkv_proj(prompt)
            add_q, add_k, add_v = ttnn.transformer.split_query_key_value_and_split_heads(
                add_qkv, num_heads=local_heads, transpose_key=False, memory_config=add_qkv.memory_config()
            )
            add_q = self.norm_added_q(add_q)
            add_k = self.norm_added_k(add_k)

            if prompt_rope is not None:
                add_q = _apply_rope(add_q, prompt_rope)
                add_k = _apply_rope(add_k, prompt_rope)

            if self.context_head_factors is not None:
                add_q = add_q * self.context_head_factors.data
        else:
            shape = [1, self.n_local_heads, 0, self.head_dim]
            add_q = add_k = add_v = ttnn.zeros(shape, device=self.mesh_device, layout=q.layout, dtype=q.dtype)

        if self.parallel_config.sequence_parallel.factor > 1:
            spatial, prompt, _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=spatial_sequence_length,
                program_config=self.sdpa_program_config,
                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.sdpa_worker_grid[1]),
            )
        else:
            assert (
                spatial_sequence_length == spatial.shape[1]
            ), "spatial sequence must not be padded without sequence parallelism"

            spatial, prompt = 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 prompt is not None:
            prompt = ttnn.transformer.concatenate_heads(prompt)

        if self.to_out is not None:
            spatial = self.ccl_manager.all_gather_persistent_buffer(
                spatial, dim=2, mesh_axis=tp_axis, use_hyperparams=True
            )
            spatial = self.to_out(spatial)

        if self.to_add_out is not None:
            prompt = self.ccl_manager.all_gather_persistent_buffer(
                prompt, dim=2, mesh_axis=tp_axis, use_hyperparams=True
            )
            prompt = self.to_add_out(prompt)

        return spatial, prompt

    @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))


def _apply_rope(x: ttnn.Tensor, freqs_cis: tuple[ttnn.Tensor, ttnn.Tensor]) -> ttnn.Tensor:
    cos, sin = freqs_cis
    cos = cos.reshape([1, 1, *cos.shape])
    sin = sin.reshape([1, 1, *sin.shape])

    return x * cos + ttnn.alt_complex_rotate90(x) * sin

	# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

	# SPDX-License-Identifier: Apache-2.0

	from __future__ import annotations

	from typing import TYPE_CHECKING

	import torch

	import ttnn

	from ..layers.linear import ColParallelLinear
	from ..layers.module import Module, Parameter, UnregisteredModule
	from ..layers.normalization import RMSNorm
	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


	# adapted from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/models/attention_processor.py
	class Attention(Module):
	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,
	eps: float,
	mesh_device: ttnn.MeshDevice,
	ccl_manager: CCLManager \| None,
	parallel_config: DiTParallelConfig,
	padding_config: PaddingConfig \| None,
	k_chunk_size: int = 512,
	q_chunk_size: int = 128,
	is_fsdp: 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.padding_config = padding_config
	self.use_spatial_weights_for_prompt = use_spatial_weights_for_prompt

	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)

	self.sdpa_worker_grid = (
	self.mesh_device.compute_with_storage_grid_size().x,
	self.mesh_device.compute_with_storage_grid_size().y - 1,
	)

	self.sdpa_program_config = ttnn.SDPAProgramConfig(
	compute_with_storage_grid_size=self.sdpa_worker_grid,
	q_chunk_size=q_chunk_size,
	k_chunk_size=k_chunk_size,
	exp_approx_mode=False, # NOTE: False is more correct
	)
	self.sdpa_compute_kernel_config = ttnn.WormholeComputeKernelConfig(
	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.to_qkv = ColParallelLinear(query_dim, 3 * padded_inner_dim, mesh_axis=tp_axis, **common_args)

	self.norm_q = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)
	self.norm_k = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)

	self.to_out = (
	ColParallelLinear(padded_inner_dim, out_dim, 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, mesh_axis=tp_axis, **common_args
	)

	self.norm_added_q = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)
	self.norm_added_k = RMSNorm(embedding_dim=head_dim, norm_eps=eps, bias=False, mesh_device=mesh_device)

	self.to_add_out = (
	ColParallelLinear(padded_inner_dim, out_dim, 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
	)

	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])

	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 forward(
	self,
	*,
	spatial: ttnn.Tensor,
	prompt: ttnn.Tensor \| None = None,
	spatial_sequence_length: int,
	spatial_rope: tuple[ttnn.Tensor, ttnn.Tensor] \| None = None,
	prompt_rope: tuple[ttnn.Tensor, ttnn.Tensor] \| None = None,
	) -> tuple[ttnn.Tensor, ttnn.Tensor \| None]:
	"""Forward pass of the model.

	Args:
	spatial: Tensor with shape [batch_size, spatial_sequence_length / sp_factor, query_dim].
	prompt: Tensor with shape [batch_size, prompt_sequence_length, query_dim] (not sharded!).
	spatial_rope: Tuple of two tensors with shape [spatial_sequence_length / sp_factor, head_dim].
	prompt_rope: Tuple of two tensors with shape [prompt_sequence_length, head_dim] (not sharded!).
	"""
	assert len(spatial.shape) == 3
	if prompt is not None:
	assert len(prompt.shape) == 3
	for t in spatial_rope or ():
	assert len(t.shape) == 2
	for t in prompt_rope or ():
	assert len(t.shape) == 2

	device_grid = self.mesh_device.compute_with_storage_grid_size()
	core_grid = ttnn.CoreGrid(x=device_grid.x, y=device_grid.y)
	tp_axis = self.parallel_config.tensor_parallel.mesh_axis

	qkv = self.to_qkv(
	spatial
	) # [batch_size, spatial_sequence_length / sp_factor, 3 * n_local_heads * head_dim (in this order)]
	local_heads = self.n_local_heads
	# Explicit memory_config: forward the input qkv layout to the inner reshapes so the
	# Q/K/V outputs keep the spec the downstream SDPA op is tuned for (matches the
	# layout main was implicitly producing pre-#42543).
	q, k, v = ttnn.transformer.split_query_key_value_and_split_heads(
	qkv, num_heads=local_heads, transpose_key=False, memory_config=qkv.memory_config()
	) # [batch_size, n_local_heads, spatial_sequence_length / sp_factor, head_dim]

	q = self.norm_q(q)
	k = self.norm_k(k)

	if spatial_rope is not None:
	q = _apply_rope(q, spatial_rope)
	k = _apply_rope(k, spatial_rope)

	if self.add_qkv_proj is not None:
	add_qkv = self.add_qkv_proj(prompt)
	add_q, add_k, add_v = ttnn.transformer.split_query_key_value_and_split_heads(
	add_qkv, num_heads=local_heads, transpose_key=False, memory_config=add_qkv.memory_config()
	)
	add_q = self.norm_added_q(add_q)
	add_k = self.norm_added_k(add_k)

	if prompt_rope is not None:
	add_q = _apply_rope(add_q, prompt_rope)
	add_k = _apply_rope(add_k, prompt_rope)

	if self.context_head_factors is not None:
	add_q = add_q * self.context_head_factors.data
	else:
	shape = [1, self.n_local_heads, 0, self.head_dim]
	add_q = add_k = add_v = ttnn.zeros(shape, device=self.mesh_device, layout=q.layout, dtype=q.dtype)

	if self.parallel_config.sequence_parallel.factor > 1:
	spatial, prompt, _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=spatial_sequence_length,
	program_config=self.sdpa_program_config,
	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.sdpa_worker_grid[1]),
	)
	else:
	assert (
	spatial_sequence_length == spatial.shape[1]
	), "spatial sequence must not be padded without sequence parallelism"

	spatial, prompt = 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 prompt is not None:
	prompt = ttnn.transformer.concatenate_heads(prompt)

	if self.to_out is not None:
	spatial = self.ccl_manager.all_gather_persistent_buffer(
	spatial, dim=2, mesh_axis=tp_axis, use_hyperparams=True
	)
	spatial = self.to_out(spatial)

	if self.to_add_out is not None:
	prompt = self.ccl_manager.all_gather_persistent_buffer(
	prompt, dim=2, mesh_axis=tp_axis, use_hyperparams=True
	)
	prompt = self.to_add_out(prompt)

	return spatial, prompt

	@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))


	def _apply_rope(x: ttnn.Tensor, freqs_cis: tuple[ttnn.Tensor, ttnn.Tensor]) -> ttnn.Tensor:
	cos, sin = freqs_cis
	cos = cos.reshape([1, 1, *cos.shape])
	sin = sin.reshape([1, 1, *sin.shape])

	return x * cos + ttnn.alt_complex_rotate90(x) * sin