wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-18 20:38:15 +08:00
Code Issues Actions 27 Packages Projects Releases Wiki Activity

Pass attention context to optimized backends

Browse Source
This commit is contained in:
野生の男 2026-07-13 14:38:55 +09:00
parent 917faef771
commit ee7536060f
5 changed files with 50 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
comfy/ldm/cosmos/predict2.py
View File

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

11
comfy/ldm/lightricks/model.py
View File

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

5
comfy/ldm/modules/attention.py
View File

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

12
comfy/ldm/wan/ar_model.py
View File

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

1
comfy/ldm/wan/model.py
View File

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