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

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

Add files using upload-large-folder tool

629b123 verified 21 days ago

15.6 kB

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

# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import TYPE_CHECKING

import ttnn

from ..layers.feedforward import ParallelFeedForward
from ..layers.linear import ColParallelLinear, prepare_chunked_linear_output
from ..layers.module import Module
from ..layers.normalization import DistributedLayerNorm
from ..utils.substate import rename_substate
from ..utils.tensor import as_bf16
from .attention_opt import Attention

if TYPE_CHECKING:
    import torch

    from ..parallel.config import DiTParallelConfig
    from ..parallel.manager import CCLManager
    from ..utils.padding import PaddingConfig


# adapted from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/models/attention_processor.py
class TransformerBlock(Module):
    def __init__(
        self,
        *,
        dim: int,
        modulation_dim: int | None = None,
        num_heads: int,
        head_dim: int,
        context_pre_only: bool,
        add_attention_to_output: bool = True,
        context_head_scaling: bool = False,
        ff_activation_fn: str = "gelu",
        ff_mult: int = 4,
        ff_bias: bool = True,
        attention_proj_bias: bool = True,
        time_norm_affine: bool = True,
        mesh_device: ttnn.MeshDevice,
        ccl_manager: CCLManager | None,
        parallel_config: DiTParallelConfig,
        padding_config: PaddingConfig | None,
        attention_k_chunk_size: int = 512,
        attention_q_chunk_size: int = 128,
        is_fsdp: bool = False,
        shard_prompt: bool = False,
    ) -> None:
        super().__init__()

        if modulation_dim is None:
            modulation_dim = dim

        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.add_attention_to_output = add_attention_to_output

        self.context_pre_only = context_pre_only
        self.mesh_device = mesh_device
        self.ccl_manager = ccl_manager
        self.parallel_config = parallel_config

        # FSDP: shard weights on sequence parallel axis to reduce memory
        fsdp_mesh_axis = parallel_config.sequence_parallel.mesh_axis if is_fsdp else None

        self.norm1_linear = (
            ColParallelLinear(
                modulation_dim,
                6 * dim,
                bias=True,
                mesh_device=mesh_device,
                mesh_axis=parallel_config.tensor_parallel.mesh_axis,
                fsdp_mesh_axis=fsdp_mesh_axis,
                ccl_manager=ccl_manager,
            )
            if time_norm_affine
            else None
        )

        self.norm1_norm = DistributedLayerNorm(
            dim,
            norm_eps=1e-6,
            norm_elementwise_affine=False,
            bias=False,
            mesh_axis=parallel_config.tensor_parallel.mesh_axis,
            mesh_device=mesh_device,
            ccl_manager=ccl_manager,
        )

        context_norm_dim = 6 * dim if not context_pre_only else 2 * dim
        self.norm1_context_linear = (
            ColParallelLinear(
                modulation_dim,
                context_norm_dim,
                bias=True,
                mesh_device=mesh_device,
                mesh_axis=parallel_config.tensor_parallel.mesh_axis,
                fsdp_mesh_axis=fsdp_mesh_axis,
                ccl_manager=ccl_manager,
            )
            if time_norm_affine
            else None
        )
        self.norm1_context_norm = DistributedLayerNorm(
            dim,
            norm_eps=1e-6,
            norm_elementwise_affine=False,
            bias=False,
            mesh_axis=parallel_config.tensor_parallel.mesh_axis,
            mesh_device=mesh_device,
            ccl_manager=ccl_manager,
        )

        self.attn = Attention(
            query_dim=dim,
            head_dim=head_dim,
            heads=num_heads,
            out_dim=dim,
            added_kv_proj_dim=dim,
            context_pre_only=context_pre_only,
            context_head_scaling=context_head_scaling,
            proj_bias=attention_proj_bias,
            eps=1e-6,
            mesh_device=mesh_device,
            ccl_manager=ccl_manager,
            parallel_config=parallel_config,
            padding_config=padding_config,
            k_chunk_size=attention_k_chunk_size,
            q_chunk_size=attention_q_chunk_size,
            is_fsdp=is_fsdp,
            shard_prompt=shard_prompt,
            per_head_norm=True,
        )

        self.norm2 = DistributedLayerNorm(
            dim,
            norm_eps=1e-6,
            norm_elementwise_affine=False,
            bias=False,
            mesh_axis=parallel_config.tensor_parallel.mesh_axis,
            mesh_device=mesh_device,
            ccl_manager=ccl_manager,
        )

        self.ff = ParallelFeedForward(
            dim=dim,
            dim_out=dim,
            mult=ff_mult,
            bias=ff_bias,
            activation_fn=ff_activation_fn,
            mesh_device=mesh_device,
            mesh_axis=parallel_config.tensor_parallel.mesh_axis,
            fsdp_mesh_axis=fsdp_mesh_axis,
            ccl_manager=ccl_manager,
        )

        self.norm2_context = None
        self.ff_context = None

        if not context_pre_only:
            self.norm2_context = DistributedLayerNorm(
                dim,
                norm_eps=1e-6,
                norm_elementwise_affine=False,
                bias=False,
                mesh_axis=parallel_config.tensor_parallel.mesh_axis,
                mesh_device=mesh_device,
                ccl_manager=ccl_manager,
            )
            self.ff_context = ParallelFeedForward(
                dim=dim,
                dim_out=dim,
                mult=ff_mult,
                bias=ff_bias,
                activation_fn=ff_activation_fn,
                mesh_device=mesh_device,
                mesh_axis=parallel_config.tensor_parallel.mesh_axis,
                fsdp_mesh_axis=fsdp_mesh_axis,
                ccl_manager=ccl_manager,
            )

        self.ff_compute_kernel_config = ttnn.init_device_compute_kernel_config(
            mesh_device.arch(),
            math_fidelity=ttnn.MathFidelity.HiFi2,
            math_approx_mode=False,
            fp32_dest_acc_en=True,
            packer_l1_acc=True,
        )

        device_grid = self.mesh_device.compute_with_storage_grid_size()
        self.core_grid = ttnn.CoreGrid(x=device_grid.x, y=device_grid.y)

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        rename_substate(state, "norm1.linear", "norm1_linear")
        rename_substate(state, "norm1.norm", "norm1_norm")
        rename_substate(state, "norm1_context.linear", "norm1_context_linear")
        rename_substate(state, "norm1_context.norm", "norm1_context_norm")
        rename_substate(state, "ff.net.0.proj", "ff.ff1")
        rename_substate(state, "ff.net.2", "ff.ff2")
        rename_substate(state, "ff_context.net.0.proj", "ff_context.ff1")
        rename_substate(state, "ff_context.net.2", "ff_context.ff2")

        prepare_chunked_linear_output(
            state,
            prefix="norm1_linear",
            device_count=self.parallel_config.tensor_parallel.factor,
            chunks=6,
        )
        prepare_chunked_linear_output(
            state,
            prefix="norm1_context_linear",
            device_count=self.parallel_config.tensor_parallel.factor,
            chunks=2 if self.context_pre_only else 6,
        )

    def forward(
        self,
        spatial: ttnn.Tensor,
        prompt: ttnn.Tensor,
        time_embed: ttnn.Tensor,
        spatial_sequence_length: int,
        prompt_sequence_length: int,
        *,
        spatial_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None,
        prompt_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None,
        temb_mod_params_img: tuple[ttnn.Tensor, ...] | None = None,
        temb_mod_params_txt: tuple[ttnn.Tensor, ...] | None = None,
        skip_time_embed_activation_fn: bool = False,
    ) -> tuple[ttnn.Tensor, ttnn.Tensor | None]:
        """Run the model forward.
        Args:
            spatial: Tensor with shape [batch_size, spatial_sequence_length / sp_factor, query_dim / tp_factor].
            prompt: Tensor with shape [batch_size, prompt_sequence_length, query_dim / tp_factor] (sequence is not sharded!).
            time_embed: Tensor with shape [batch_size, 1, query_dim].
            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] (sequence is not sharded!).
        """
        assert len(spatial.shape) == 3
        assert len(prompt.shape) == 3

        tp_axis = self.parallel_config.tensor_parallel.mesh_axis

        if not skip_time_embed_activation_fn:
            time_embed = ttnn.silu(time_embed, memory_config=ttnn.DRAM_MEMORY_CONFIG)

        # Contract: scale params are consumed with +1 already applied (the norm computes
        # (1 + scale) * x + shift). External callers pre-add 1 once and share the tuple across
        # blocks; the internal-modulation path below pre-adds it here for parity.
        if self.norm1_linear is not None:
            assert temb_mod_params_img is None

            spatial_time = self.norm1_linear(
                time_embed,
                core_grid=Attention.get_core_grid(
                    time_embed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
                ),
            )
            s_shift_a, s_scale_a, s_gate_a, s_shift_f, s_scale_f, s_gate_f = _chunk_time3d(spatial_time, 6)
            temb_mod_params_img = (s_shift_a, s_scale_a + 1, s_gate_a, s_shift_f, s_scale_f + 1, s_gate_f)

        if self.norm1_context_linear is not None:
            assert temb_mod_params_txt is None

            prompt_time = self.norm1_context_linear(
                time_embed,
                core_grid=Attention.get_core_grid(
                    time_embed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
                ),
            )
            if self.context_pre_only:
                # layout: (scale_attn, shift_attn)
                p_scale_a, p_shift_a = _chunk_time3d(prompt_time, 2)
                temb_mod_params_txt = (p_scale_a + 1, p_shift_a)
            else:
                p_shift_a, p_scale_a, p_gate_a, p_shift_f, p_scale_f, p_gate_f = _chunk_time3d(prompt_time, 6)
                temb_mod_params_txt = (p_shift_a, p_scale_a + 1, p_gate_a, p_shift_f, p_scale_f + 1, p_gate_f)

        assert temb_mod_params_img is not None
        assert temb_mod_params_txt is not None

        (
            spatial_shift_attn,
            spatial_scale_attn,
            spatial_gate_attn,
            spatial_shift_ff,
            spatial_scale_ff,
            spatial_gate_ff,
        ) = temb_mod_params_img

        spatial_normed = ttnn.squeeze(
            self.norm1_norm(
                ttnn.unsqueeze(spatial, 0),
                dynamic_weight=spatial_scale_attn,
                dynamic_bias=spatial_shift_attn,
            ),
            0,
        )

        if self.context_pre_only:
            prompt_scale_attn, prompt_shift_attn = temb_mod_params_txt
            prompt_gate_attn = None
            prompt_shift_ff = None
            prompt_scale_ff = None
            prompt_gate_ff = None
        else:
            (
                prompt_shift_attn,
                prompt_scale_attn,
                prompt_gate_attn,
                prompt_shift_ff,
                prompt_scale_ff,
                prompt_gate_ff,
            ) = temb_mod_params_txt

        prompt_normed = ttnn.squeeze(
            self.norm1_context_norm(
                ttnn.unsqueeze(prompt, 0),
                dynamic_weight=prompt_scale_attn,
                dynamic_bias=prompt_shift_attn,
            ),
            0,
        )

        is_ring = self.ccl_manager.topology == ttnn.Topology.Ring

        # NOTE: workaround - addcmul is less accurate with fp32 gate input
        spatial_gate_attn = as_bf16(spatial_gate_attn)
        if prompt_gate_attn is not None:
            prompt_gate_attn = as_bf16(prompt_gate_attn)

        spatial, prompt_attn = self.attn.forward(
            sequence_1=spatial_normed,
            sequence_2=prompt_normed,
            sequence_1_rope=spatial_rope,
            sequence_2_rope=prompt_rope,
            sequence_1_length=spatial_sequence_length,
            sequence_2_length=prompt_sequence_length,
            addcmul_spatial_residual=spatial,
            addcmul_spatial_gate=spatial_gate_attn,
            addcmul_prompt_residual=prompt if not self.context_pre_only else None,
            addcmul_prompt_gate=prompt_gate_attn,
        )

        if prompt_attn is not None:
            prompt = prompt_attn

        spatial_normed = ttnn.squeeze(
            self.norm2(
                ttnn.unsqueeze(spatial, 0),
                dynamic_weight=spatial_scale_ff,
                dynamic_bias=spatial_shift_ff,
            ),
            0,
        )

        spatial_gate_ff = as_bf16(spatial_gate_ff)
        if is_ring:
            spatial = self.ff.forward_fused_addcmul(
                spatial_normed,
                spatial,
                spatial_gate_ff,
                scalar=1.0,
                compute_kernel_config=self.ff_compute_kernel_config,
                parallel_config=self.parallel_config,
                core_grid=Attention.get_core_grid(
                    spatial_normed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
                ),
            )
        else:
            spatial_normed = self.ccl_manager.all_gather_persistent_buffer(
                spatial_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True
            )
            spatial_ff = ttnn.squeeze(
                self.ff(ttnn.unsqueeze(spatial_normed, 0), compute_kernel_config=self.ff_compute_kernel_config), 0
            )
            spatial = ttnn.addcmul(spatial, spatial_ff, spatial_gate_ff)

        if self.context_pre_only:
            return spatial, None

        prompt_normed = ttnn.squeeze(
            self.norm2_context(
                ttnn.unsqueeze(prompt, 0),
                dynamic_weight=prompt_scale_ff,
                dynamic_bias=prompt_shift_ff,
            ),
            0,
        )

        prompt_gate_ff = as_bf16(prompt_gate_ff)
        if is_ring:
            prompt = self.ff_context.forward_fused_addcmul(
                prompt_normed,
                prompt,
                prompt_gate_ff,
                scalar=1.0,
                compute_kernel_config=self.ff_compute_kernel_config,
                parallel_config=self.parallel_config,
                core_grid=Attention.get_core_grid(
                    prompt_normed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
                ),
            )
        else:
            prompt_normed = self.ccl_manager.all_gather_persistent_buffer(
                prompt_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True
            )
            prompt_ff = ttnn.squeeze(
                self.ff_context(ttnn.unsqueeze(prompt_normed, 0), compute_kernel_config=self.ff_compute_kernel_config),
                0,
            )
            prompt = ttnn.addcmul(prompt, prompt_ff, prompt_gate_ff)

        return spatial, prompt


def _chunk_time3d(t: ttnn.Tensor, count: int) -> tuple[ttnn.Tensor]:
    size = t.shape[-1] // count
    return tuple(t[:, :, i * size : (i + 1) * size] for i in range(count))

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

	# SPDX-License-Identifier: Apache-2.0

	from __future__ import annotations

	from typing import TYPE_CHECKING

	import ttnn

	from ..layers.feedforward import ParallelFeedForward
	from ..layers.linear import ColParallelLinear, prepare_chunked_linear_output
	from ..layers.module import Module
	from ..layers.normalization import DistributedLayerNorm
	from ..utils.substate import rename_substate
	from ..utils.tensor import as_bf16
	from .attention_opt import Attention

	if TYPE_CHECKING:
	import torch

	from ..parallel.config import DiTParallelConfig
	from ..parallel.manager import CCLManager
	from ..utils.padding import PaddingConfig


	# adapted from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/models/attention_processor.py
	class TransformerBlock(Module):
	def __init__(
	self,
	*,
	dim: int,
	modulation_dim: int \| None = None,
	num_heads: int,
	head_dim: int,
	context_pre_only: bool,
	add_attention_to_output: bool = True,
	context_head_scaling: bool = False,
	ff_activation_fn: str = "gelu",
	ff_mult: int = 4,
	ff_bias: bool = True,
	attention_proj_bias: bool = True,
	time_norm_affine: bool = True,
	mesh_device: ttnn.MeshDevice,
	ccl_manager: CCLManager \| None,
	parallel_config: DiTParallelConfig,
	padding_config: PaddingConfig \| None,
	attention_k_chunk_size: int = 512,
	attention_q_chunk_size: int = 128,
	is_fsdp: bool = False,
	shard_prompt: bool = False,
	) -> None:
	super().__init__()

	if modulation_dim is None:
	modulation_dim = dim

	self.dim = dim
	self.num_heads = num_heads
	self.head_dim = head_dim
	self.add_attention_to_output = add_attention_to_output

	self.context_pre_only = context_pre_only
	self.mesh_device = mesh_device
	self.ccl_manager = ccl_manager
	self.parallel_config = parallel_config

	# FSDP: shard weights on sequence parallel axis to reduce memory
	fsdp_mesh_axis = parallel_config.sequence_parallel.mesh_axis if is_fsdp else None

	self.norm1_linear = (
	ColParallelLinear(
	modulation_dim,
	6 * dim,
	bias=True,
	mesh_device=mesh_device,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	fsdp_mesh_axis=fsdp_mesh_axis,
	ccl_manager=ccl_manager,
	)
	if time_norm_affine
	else None
	)

	self.norm1_norm = DistributedLayerNorm(
	dim,
	norm_eps=1e-6,
	norm_elementwise_affine=False,
	bias=False,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	mesh_device=mesh_device,
	ccl_manager=ccl_manager,
	)

	context_norm_dim = 6 * dim if not context_pre_only else 2 * dim
	self.norm1_context_linear = (
	ColParallelLinear(
	modulation_dim,
	context_norm_dim,
	bias=True,
	mesh_device=mesh_device,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	fsdp_mesh_axis=fsdp_mesh_axis,
	ccl_manager=ccl_manager,
	)
	if time_norm_affine
	else None
	)
	self.norm1_context_norm = DistributedLayerNorm(
	dim,
	norm_eps=1e-6,
	norm_elementwise_affine=False,
	bias=False,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	mesh_device=mesh_device,
	ccl_manager=ccl_manager,
	)

	self.attn = Attention(
	query_dim=dim,
	head_dim=head_dim,
	heads=num_heads,
	out_dim=dim,
	added_kv_proj_dim=dim,
	context_pre_only=context_pre_only,
	context_head_scaling=context_head_scaling,
	proj_bias=attention_proj_bias,
	eps=1e-6,
	mesh_device=mesh_device,
	ccl_manager=ccl_manager,
	parallel_config=parallel_config,
	padding_config=padding_config,
	k_chunk_size=attention_k_chunk_size,
	q_chunk_size=attention_q_chunk_size,
	is_fsdp=is_fsdp,
	shard_prompt=shard_prompt,
	per_head_norm=True,
	)

	self.norm2 = DistributedLayerNorm(
	dim,
	norm_eps=1e-6,
	norm_elementwise_affine=False,
	bias=False,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	mesh_device=mesh_device,
	ccl_manager=ccl_manager,
	)

	self.ff = ParallelFeedForward(
	dim=dim,
	dim_out=dim,
	mult=ff_mult,
	bias=ff_bias,
	activation_fn=ff_activation_fn,
	mesh_device=mesh_device,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	fsdp_mesh_axis=fsdp_mesh_axis,
	ccl_manager=ccl_manager,
	)

	self.norm2_context = None
	self.ff_context = None

	if not context_pre_only:
	self.norm2_context = DistributedLayerNorm(
	dim,
	norm_eps=1e-6,
	norm_elementwise_affine=False,
	bias=False,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	mesh_device=mesh_device,
	ccl_manager=ccl_manager,
	)
	self.ff_context = ParallelFeedForward(
	dim=dim,
	dim_out=dim,
	mult=ff_mult,
	bias=ff_bias,
	activation_fn=ff_activation_fn,
	mesh_device=mesh_device,
	mesh_axis=parallel_config.tensor_parallel.mesh_axis,
	fsdp_mesh_axis=fsdp_mesh_axis,
	ccl_manager=ccl_manager,
	)

	self.ff_compute_kernel_config = ttnn.init_device_compute_kernel_config(
	mesh_device.arch(),
	math_fidelity=ttnn.MathFidelity.HiFi2,
	math_approx_mode=False,
	fp32_dest_acc_en=True,
	packer_l1_acc=True,
	)

	device_grid = self.mesh_device.compute_with_storage_grid_size()
	self.core_grid = ttnn.CoreGrid(x=device_grid.x, y=device_grid.y)

	def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
	rename_substate(state, "norm1.linear", "norm1_linear")
	rename_substate(state, "norm1.norm", "norm1_norm")
	rename_substate(state, "norm1_context.linear", "norm1_context_linear")
	rename_substate(state, "norm1_context.norm", "norm1_context_norm")
	rename_substate(state, "ff.net.0.proj", "ff.ff1")
	rename_substate(state, "ff.net.2", "ff.ff2")
	rename_substate(state, "ff_context.net.0.proj", "ff_context.ff1")
	rename_substate(state, "ff_context.net.2", "ff_context.ff2")

	prepare_chunked_linear_output(
	state,
	prefix="norm1_linear",
	device_count=self.parallel_config.tensor_parallel.factor,
	chunks=6,
	)
	prepare_chunked_linear_output(
	state,
	prefix="norm1_context_linear",
	device_count=self.parallel_config.tensor_parallel.factor,
	chunks=2 if self.context_pre_only else 6,
	)

	def forward(
	self,
	spatial: ttnn.Tensor,
	prompt: ttnn.Tensor,
	time_embed: ttnn.Tensor,
	spatial_sequence_length: int,
	prompt_sequence_length: int,
	*,
	spatial_rope: tuple[ttnn.Tensor, ttnn.Tensor] \| None = None,
	prompt_rope: tuple[ttnn.Tensor, ttnn.Tensor] \| None = None,
	temb_mod_params_img: tuple[ttnn.Tensor, ...] \| None = None,
	temb_mod_params_txt: tuple[ttnn.Tensor, ...] \| None = None,
	skip_time_embed_activation_fn: bool = False,
	) -> tuple[ttnn.Tensor, ttnn.Tensor \| None]:
	"""Run the model forward.

	Args:
	spatial: Tensor with shape [batch_size, spatial_sequence_length / sp_factor, query_dim / tp_factor].
	prompt: Tensor with shape [batch_size, prompt_sequence_length, query_dim / tp_factor] (sequence is not sharded!).
	time_embed: Tensor with shape [batch_size, 1, query_dim].
	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] (sequence is not sharded!).
	"""
	assert len(spatial.shape) == 3
	assert len(prompt.shape) == 3

	tp_axis = self.parallel_config.tensor_parallel.mesh_axis

	if not skip_time_embed_activation_fn:
	time_embed = ttnn.silu(time_embed, memory_config=ttnn.DRAM_MEMORY_CONFIG)

	# Contract: scale params are consumed with +1 already applied (the norm computes
	# (1 + scale) * x + shift). External callers pre-add 1 once and share the tuple across
	# blocks; the internal-modulation path below pre-adds it here for parity.
	if self.norm1_linear is not None:
	assert temb_mod_params_img is None

	spatial_time = self.norm1_linear(
	time_embed,
	core_grid=Attention.get_core_grid(
	time_embed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
	),
	)
	s_shift_a, s_scale_a, s_gate_a, s_shift_f, s_scale_f, s_gate_f = _chunk_time3d(spatial_time, 6)
	temb_mod_params_img = (s_shift_a, s_scale_a + 1, s_gate_a, s_shift_f, s_scale_f + 1, s_gate_f)

	if self.norm1_context_linear is not None:
	assert temb_mod_params_txt is None

	prompt_time = self.norm1_context_linear(
	time_embed,
	core_grid=Attention.get_core_grid(
	time_embed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
	),
	)
	if self.context_pre_only:
	# layout: (scale_attn, shift_attn)
	p_scale_a, p_shift_a = _chunk_time3d(prompt_time, 2)
	temb_mod_params_txt = (p_scale_a + 1, p_shift_a)
	else:
	p_shift_a, p_scale_a, p_gate_a, p_shift_f, p_scale_f, p_gate_f = _chunk_time3d(prompt_time, 6)
	temb_mod_params_txt = (p_shift_a, p_scale_a + 1, p_gate_a, p_shift_f, p_scale_f + 1, p_gate_f)

	assert temb_mod_params_img is not None
	assert temb_mod_params_txt is not None

	(
	spatial_shift_attn,
	spatial_scale_attn,
	spatial_gate_attn,
	spatial_shift_ff,
	spatial_scale_ff,
	spatial_gate_ff,
	) = temb_mod_params_img

	spatial_normed = ttnn.squeeze(
	self.norm1_norm(
	ttnn.unsqueeze(spatial, 0),
	dynamic_weight=spatial_scale_attn,
	dynamic_bias=spatial_shift_attn,
	),
	0,
	)

	if self.context_pre_only:
	prompt_scale_attn, prompt_shift_attn = temb_mod_params_txt
	prompt_gate_attn = None
	prompt_shift_ff = None
	prompt_scale_ff = None
	prompt_gate_ff = None
	else:
	(
	prompt_shift_attn,
	prompt_scale_attn,
	prompt_gate_attn,
	prompt_shift_ff,
	prompt_scale_ff,
	prompt_gate_ff,
	) = temb_mod_params_txt

	prompt_normed = ttnn.squeeze(
	self.norm1_context_norm(
	ttnn.unsqueeze(prompt, 0),
	dynamic_weight=prompt_scale_attn,
	dynamic_bias=prompt_shift_attn,
	),
	0,
	)

	is_ring = self.ccl_manager.topology == ttnn.Topology.Ring

	# NOTE: workaround - addcmul is less accurate with fp32 gate input
	spatial_gate_attn = as_bf16(spatial_gate_attn)
	if prompt_gate_attn is not None:
	prompt_gate_attn = as_bf16(prompt_gate_attn)

	spatial, prompt_attn = self.attn.forward(
	sequence_1=spatial_normed,
	sequence_2=prompt_normed,
	sequence_1_rope=spatial_rope,
	sequence_2_rope=prompt_rope,
	sequence_1_length=spatial_sequence_length,
	sequence_2_length=prompt_sequence_length,
	addcmul_spatial_residual=spatial,
	addcmul_spatial_gate=spatial_gate_attn,
	addcmul_prompt_residual=prompt if not self.context_pre_only else None,
	addcmul_prompt_gate=prompt_gate_attn,
	)

	if prompt_attn is not None:
	prompt = prompt_attn

	spatial_normed = ttnn.squeeze(
	self.norm2(
	ttnn.unsqueeze(spatial, 0),
	dynamic_weight=spatial_scale_ff,
	dynamic_bias=spatial_shift_ff,
	),
	0,
	)

	spatial_gate_ff = as_bf16(spatial_gate_ff)
	if is_ring:
	spatial = self.ff.forward_fused_addcmul(
	spatial_normed,
	spatial,
	spatial_gate_ff,
	scalar=1.0,
	compute_kernel_config=self.ff_compute_kernel_config,
	parallel_config=self.parallel_config,
	core_grid=Attention.get_core_grid(
	spatial_normed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
	),
	)
	else:
	spatial_normed = self.ccl_manager.all_gather_persistent_buffer(
	spatial_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True
	)
	spatial_ff = ttnn.squeeze(
	self.ff(ttnn.unsqueeze(spatial_normed, 0), compute_kernel_config=self.ff_compute_kernel_config), 0
	)
	spatial = ttnn.addcmul(spatial, spatial_ff, spatial_gate_ff)

	if self.context_pre_only:
	return spatial, None

	prompt_normed = ttnn.squeeze(
	self.norm2_context(
	ttnn.unsqueeze(prompt, 0),
	dynamic_weight=prompt_scale_ff,
	dynamic_bias=prompt_shift_ff,
	),
	0,
	)

	prompt_gate_ff = as_bf16(prompt_gate_ff)
	if is_ring:
	prompt = self.ff_context.forward_fused_addcmul(
	prompt_normed,
	prompt,
	prompt_gate_ff,
	scalar=1.0,
	compute_kernel_config=self.ff_compute_kernel_config,
	parallel_config=self.parallel_config,
	core_grid=Attention.get_core_grid(
	prompt_normed.shape[-2], self.mesh_device.compute_with_storage_grid_size()
	),
	)
	else:
	prompt_normed = self.ccl_manager.all_gather_persistent_buffer(
	prompt_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True
	)
	prompt_ff = ttnn.squeeze(
	self.ff_context(ttnn.unsqueeze(prompt_normed, 0), compute_kernel_config=self.ff_compute_kernel_config),
	0,
	)
	prompt = ttnn.addcmul(prompt, prompt_ff, prompt_gate_ff)

	return spatial, prompt


	def _chunk_time3d(t: ttnn.Tensor, count: int) -> tuple[ttnn.Tensor]:
	size = t.shape[-1] // count
	return tuple(t[:, :, i * size : (i + 1) * size] for i in range(count))