Function sp_dit_forward_causal

wan/distributed/sequence_parallel.py:249–411

x: A list of videos each with shape [C, T, H, W]. t: [B]. context: A list of text embeddings each with shape [L, C]. seq_len: Maximum sequence length for positional encoding. y: Conditional video inputs for

(
    self,
    x,
    t,
    context,
    seq_len,
    y=None,
    dit_cond_dict=None,
    kv_cache=None,
    crossattn_cache=None,
    current_start=0,
    max_attention_size=1_000_000,
    frame_seqlen=None,
    cross_attn_first_call=None,
)

def sp_dit_forward_causal(
    self,
    x,
    t,
    context,
    seq_len,
    y=None,
    dit_cond_dict=None,
    kv_cache=None,
    crossattn_cache=None,
    current_start=0,
    max_attention_size=1_000_000,
    frame_seqlen=None,
    cross_attn_first_call=None,
):
    """
    x:                  A list of videos each with shape [C, T, H, W].
    t:                  [B].
    context:            A list of text embeddings each with shape [L, C].
    seq_len:            Maximum sequence length for positional encoding.
    y:                  Conditional video inputs for image-to-video mode, same shape as x.
    dit_cond_dict:      Dictionary of conditioning signals. May contain key ``c2ws_plucker_emb``
                        with camera Plucker embeddings of shape [B, C, F, H, W] for camera control.
    kv_cache:           Per-layer self-attention KV cache. Each dict 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).
    crossattn_cache:    Per-layer cross-attention KV cache. Each dict contains keys ``k``, ``v``
                        (Tensor of shape [B, text_len, num_heads, head_dim]) and ``is_init`` (bool).
    current_start:      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: Maximum number of KV tokens each query can attend to. Limits the
                        effective context window of self-attention to control memory usage.

    Note: This function uses sequence parallel (SP) with Ulysses attention.
    The sequence is chunked across GPUs. Inside attention, all-to-all redistributes
    from sequence-parallel to head-parallel, applies causal RoPE on the full sequence,
    updates the KV cache, runs attention, then all-to-all back.
    """

    assert len(x) == 1

    if self.model_type == 'i2v':
        assert y is not None
    # params
    device = self.patch_embedding.weight.device
    if self.freqs.device != device:
        self.freqs = self.freqs.to(device)

    if y is not None:
        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]

    # embeddings
    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
    grid_sizes = torch.stack(
        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
    x = [u.flatten(2).transpose(1, 2) for u in x]
    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)[0]
    assert seq_lens.max() <= seq_len