Method forward

wan/modules/model_fast.py:172–213

r""" Args: x(Tensor): Shape [B, L1, C] context(Tensor): Shape [B, L2, C] context_lens(Tensor): Shape [B] cross_attn_first_call(bool, optional): If provided, used as the "first call this generation" gate instead of reading

(self, x, context, context_lens, crossattn_cache=None,
                cross_attn_first_call=None)

class WanCrossAttention(WanSelfAttention):

    def forward(self, x, context, context_lens, crossattn_cache=None,
                cross_attn_first_call=None):
        r"""
        Args:
            x(Tensor): Shape [B, L1, C]
            context(Tensor): Shape [B, L2, C]
            context_lens(Tensor): Shape [B]
            cross_attn_first_call(bool, optional): If provided, used as the
                "first call this generation" gate instead of reading
                crossattn_cache["is_init"].item() (which forces a CPU↔GPU
                sync). Caller (pipeline) tracks this as a Python bool.
        """
        b, n, d = x.size(0), self.num_heads, self.head_dim

        # compute query, key, value
        q = self.norm_q(self.q(x)).view(b, -1, n, d)

        if crossattn_cache is not None:
            if cross_attn_first_call is None:
                is_first = crossattn_cache["is_init"].item() == 0
            else:
                is_first = cross_attn_first_call
            if is_first:
                crossattn_cache["is_init"].fill_(1)
                k = self.norm_k(self.k(context)).view(b, -1, n, d)
                v = self.v(context).view(b, -1, n, d)
                crossattn_cache["k"].copy_(k)
                crossattn_cache["v"].copy_(v)
            else:
                k = crossattn_cache["k"]
                v = crossattn_cache["v"]
        else:
            k = self.norm_k(self.k(context)).view(b, -1, n, d)
            v = self.v(context).view(b, -1, n, d)

        # compute attention
        x = flash_attention(q, k, v, k_lens=context_lens)

        # output
        x = x.flatten(2)
        x = self.o(x)
        return x


class CausalWanAttentionBlock(nn.Module):