Function sp_attn_forward_causal

wan/distributed/sequence_parallel.py:414–550

r""" Sequence-parallel causal self-attention using Ulysses all-to-all. Input x is sequence-chunked (each GPU holds padded_seq_lens/world_size tokens). all-to-all gathers the full sequence and scatters heads, so each GPU operates on the full sequence with local_heads. After RoPE, pa

(
    self,
    x,
    seq_lens,
    grid_sizes,
    freqs,
    kv_cache=None,
    current_start=0,
    max_attention_size=1_000_000,
    frame_seqlen=None,
    seq_lens_int=None)

def sp_attn_forward_causal(
    self,
    x,
    seq_lens,
    grid_sizes,
    freqs,
    kv_cache=None,
    current_start=0,
    max_attention_size=1_000_000,
    frame_seqlen=None,
    seq_lens_int=None):
    r"""
    Sequence-parallel causal self-attention using Ulysses all-to-all.

    Input x is sequence-chunked (each GPU holds padded_seq_lens/world_size tokens).
    all-to-all gathers the full sequence and scatters heads, so each GPU
    operates on the full sequence with local_heads.  After RoPE, padding
    tokens are sliced off so that only the actual ``seq_lens`` tokens enter
    the KV cache and flash attention.  The output is padded back to
    ``padded_seq_lens`` before the reverse all-to-all.

    Args:
        x(Tensor):              Shape [B, padded_seq_lens // world_size, C]
        seq_lens(Tensor):       Total number of valid tokens (full sequence, not per-rank).
        grid_sizes(Tensor):     Shape [B, 3], the second dimension contains (F, H, W).
        freqs(Tensor):          Rope freqs, shape [1024, C / num_heads / 2].
        kv_cache(dict):         Self-attention KV cache. Contains keys ``k``, ``v``
                                (Tensor of shape [B, kv_size, local_heads, head_dim]),
                                ``global_end_index``, and ``local_end_index``
                                (scalar Tensors tracking cache position).
        current_start(int):     Token offset of the current chunk in the full sequence.
                                Used to index into the KV cache and compute positional
                                embeddings correctly.
        max_attention_size(int): Maximum number of KV tokens each query can attend to.
                                Limits the effective context window of self-attention
                                to control memory usage.
    """
    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
    sp_size = get_world_size()

    # query, key, value function
    def qkv_fn(x):
        q = self.norm_q(self.q(x)).view(b, s, n, d)
        k = self.norm_k(self.k(x)).view(b, s, n, d)
        v = self.v(x).view(b, s, n, d)
        return q, k, v

    q, k, v = qkv_fn(x)

    # all-to-all: gather sequence, scatter heads
    # [B, s/p, N, d] -> [B, s, N/p, d]
    q = all_to_all(q, scatter_dim=2, gather_dim=1)
    k = all_to_all(k, scatter_dim=2, gather_dim=1)
    v = all_to_all(v, scatter_dim=2, gather_dim=1)

    # padded_seq_lens = s * sp_size may exceed seq_lens due to SP padding
    padded_seq_lens = s * sp_size
    if seq_lens_int is None: