From ee7536060fac89059a0f075e1596fba34b7c314c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=8E=E7=94=9F=E3=81=AE=E7=94=B7?= Date: Mon, 13 Jul 2026 14:38:55 +0900 Subject: [PATCH] Pass attention context to optimized backends --- comfy/ldm/cosmos/predict2.py | 31 +++++++++++++++++++++++++++---- comfy/ldm/lightricks/model.py | 11 +++++++++-- comfy/ldm/modules/attention.py | 5 +++-- comfy/ldm/wan/ar_model.py | 12 ++++++++++-- comfy/ldm/wan/model.py | 1 + 5 files changed, 50 insertions(+), 10 deletions(-) diff --git a/comfy/ldm/cosmos/predict2.py b/comfy/ldm/cosmos/predict2.py index aec874815..f560bded0 100644 --- a/comfy/ldm/cosmos/predict2.py +++ b/comfy/ldm/cosmos/predict2.py @@ -37,7 +37,13 @@ class GPT2FeedForward(nn.Module): return x -def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor: +def torch_attention_op( + q_B_S_H_D: torch.Tensor, + k_B_S_H_D: torch.Tensor, + v_B_S_H_D: torch.Tensor, + transformer_options: Optional[dict] = {}, + is_self_attention: bool = False, +) -> torch.Tensor: """Computes multi-head attention using PyTorch's native implementation. This function provides a PyTorch backend alternative to Transformer Engine's attention operation. @@ -64,7 +70,16 @@ def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H q_B_H_S_D = rearrange(q_B_S_H_D, "b ... h k -> b h ... k").view(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1]) k_B_H_S_D = rearrange(k_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1]) v_B_H_S_D = rearrange(v_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1]) - return optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True, transformer_options=transformer_options) + return optimized_attention( + q_B_H_S_D, + k_B_H_S_D, + v_B_H_S_D, + in_q_shape[-2], + skip_reshape=True, + transformer_options=transformer_options, + is_self_attention=is_self_attention, + attention_token_shape=transformer_options.get("attention_token_shape", tuple(in_q_shape[1:-2])), + ) class Attention(nn.Module): @@ -173,7 +188,13 @@ class Attention(nn.Module): return q, k, v def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor: - result = self.attn_op(q, k, v, transformer_options=transformer_options) # [B, S, H, D] + result = self.attn_op( + q, + k, + v, + transformer_options=transformer_options, + is_self_attention=self.is_selfattn, + ) # [B, S, H * D] return self.output_dropout(self.output_proj(result)) def forward( @@ -492,6 +513,8 @@ class Block(nn.Module): gate_mlp_B_T_1_1_D = rearrange(gate_mlp_B_T_D, "b t d -> b t 1 1 d") B, T, H, W, D = x_B_T_H_W_D.shape + self_attn_options = transformer_options.copy() + self_attn_options["attention_token_shape"] = (T, H, W) def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D): return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D @@ -508,7 +531,7 @@ class Block(nn.Module): rearrange(normalized_x_B_T_H_W_D.to(compute_dtype), "b t h w d -> b (t h w) d"), None, rope_emb=rope_emb_L_1_1_D, - transformer_options=transformer_options, + transformer_options=self_attn_options, ), "b (t h w) d -> b t h w d", t=T, diff --git a/comfy/ldm/lightricks/model.py b/comfy/ldm/lightricks/model.py index 9953b6679..9266c3c3c 100644 --- a/comfy/ldm/lightricks/model.py +++ b/comfy/ldm/lightricks/model.py @@ -454,6 +454,7 @@ class CrossAttention(nn.Module): def forward(self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}): q = self.to_q(x) + is_self_attention = context is None context = x if context is None else context k = self.to_k(context) v = self.to_v(context) @@ -466,11 +467,11 @@ class CrossAttention(nn.Module): k = apply_rotary_emb(k, pe if k_pe is None else k_pe) if mask is None: - out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options) + out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention) elif isinstance(mask, GuideAttentionMask): out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options) else: - out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options) + out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention) # Apply per-head gating if enabled if self.to_gate_logits is not None: @@ -970,6 +971,12 @@ class LTXBaseModel(torch.nn.Module, ABC): merged_args = {**transformer_options, **kwargs} x, pixel_coords, additional_args = self._process_input(x, keyframe_idxs, denoise_mask, **merged_args) merged_args.update(additional_args) + transformer_options.pop("attention_token_grid", None) + orig_shape = additional_args.get("orig_shape") + if isinstance(x, torch.Tensor) and isinstance(orig_shape, (tuple, list)) and len(orig_shape) == 5: + token_grid = tuple(orig_shape[-3:]) + if math.prod(token_grid) == x.shape[1]: + transformer_options["attention_token_grid"] = token_grid # Prepare timestep and context timestep, embedded_timestep, prompt_timestep = self._prepare_timestep(timestep, batch_size, input_dtype, **merged_args) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index e6500cff4..bd6a08525 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -844,6 +844,7 @@ class CrossAttention(nn.Module): def forward(self, x, context=None, value=None, mask=None, transformer_options={}): q = self.to_q(x) + is_self_attention = context is None and value is None context = default(context, x) k = self.to_k(context) if value is not None: @@ -853,9 +854,9 @@ class CrossAttention(nn.Module): v = self.to_v(context) if mask is None: - out = optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options) + out = optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention) else: - out = optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options) + out = optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options, is_self_attention=is_self_attention) return self.to_out(out) diff --git a/comfy/ldm/wan/ar_model.py b/comfy/ldm/wan/ar_model.py index d72f53602..0d6fb313e 100644 --- a/comfy/ldm/wan/ar_model.py +++ b/comfy/ldm/wan/ar_model.py @@ -62,6 +62,7 @@ class CausalWanSelfAttention(nn.Module): v.view(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options, + is_self_attention=True, ) else: end = kv_cache["end"] @@ -78,6 +79,8 @@ class CausalWanSelfAttention(nn.Module): kv_cache["v"][:, :new_end].view(b, new_end, n * d), heads=self.num_heads, transformer_options=transformer_options, + is_self_attention=True, + is_kv_cached_attention=True, ) x = self.o(x) @@ -170,7 +173,7 @@ class CausalWanModel(WanModel): device=device, dtype=dtype, operations=operations) def forward_block(self, x, timestep, context, start_frame, - kv_caches, crossattn_caches, clip_fea=None): + kv_caches, crossattn_caches, clip_fea=None, transformer_options={}): """ Forward one temporal block for autoregressive inference. @@ -191,6 +194,7 @@ class CausalWanModel(WanModel): x = self.patch_embedding(x.float()).to(x.dtype) grid_sizes = x.shape[2:] + transformer_options["grid_sizes"] = grid_sizes x = x.flatten(2).transpose(1, 2) # Per-frame time embedding @@ -212,11 +216,15 @@ class CausalWanModel(WanModel): freqs = self.rope_encode(t, h, w, t_start=start_frame, device=x.device, dtype=x.dtype) # Transformer blocks + transformer_options["total_blocks"] = len(self.blocks) + transformer_options["block_type"] = "double" for i, block in enumerate(self.blocks): + transformer_options["block_index"] = i x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, kv_cache=kv_caches[i], - crossattn_cache=crossattn_caches[i]) + crossattn_cache=crossattn_caches[i], + transformer_options=transformer_options) # Head x = self.head(x, e) diff --git a/comfy/ldm/wan/model.py b/comfy/ldm/wan/model.py index 1c9782a38..8d7fa0bfa 100644 --- a/comfy/ldm/wan/model.py +++ b/comfy/ldm/wan/model.py @@ -86,6 +86,7 @@ class WanSelfAttention(nn.Module): self.v(x).view(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options, + is_self_attention=True, ) if "attn1_patch" in patches: