[Support for Cross-Attention] How to use cross-attention with magiattention

[zelingao98]

Dear all,

thanks for the great work!
I am using this great magiattention.
However, I find it seems that magiattention only support self-attention.
e.g.,
q: [32*1024(seq_len), 12(num_head), 128(channel)]

k: [32*512(seq_len), 12(num_head), 128(channel)]

v: [32*512(seq_len), 12(num_head), 128(channel)]
the code will say something like:

[default0]:[rank0]: File "/opt/conda/lib/python3.10/site-packages/magi_attention/meta/_calc_dispatch_meta.py", line 241, in _calc_self_attn_dispatch_meta_from_qk_ranges
[default0]:[rank0]: assert total_seqlen_q == total_seqlen_k and num_chunks_q == num_chunks_k, (
[default0]:[rank0]: AssertionError: For self-attn, total_seqlen_q=32768 should be the same as total_seqlen_k=16384, as well as num_chunks_q=64 and num_chunks_k=32

Pls correct me if I am wrong.

Best!

[littsk]

Hello!

Thank you for your interest in magiattention and for reaching out with your question. Your observation is correct; the error message you're encountering indicates that the current version of magiattention is optimized for self-attention, where the sequence lengths of the query (Q), key (K), and value (V) tensors are expected to be the same.

However, you can still perform cross-attention by using a workaround. Here is a polished version of your proposed answer, which correctly outlines the current approach for handling cross-attention.

Using Cross-Attention with magiattention

At present, the distributed dispatching and sharding of key (K) and value (V) tensors for cross-attention is not yet supported in magiattention. The assertion error you are seeing occurs because the system anticipates that the sequence lengths for Q, K, and V will be identical, as is the case with self-attention.

A temporary and effective workaround for the current release is to replicate the K and V tensors across all ranks. This approach allows you to perform cross-attention by providing the complete K and V tensors to each process, while still distributing the query (Q) tensor.

We are actively working to provide a more comprehensive and configurable solution for cross-attention in future versions. We appreciate your support and feedback, and we encourage contributions and issue reports from our community.

Python Example for Cross-Attention

Here is a Python example demonstrating the basic usage of cross-attention with torchrun by replicating the K and V tensors:

# basic_xattn_torchrun.py

import os
import torch
import torch.distributed as dist

from magi_attention import init_dist_attn_runtime_mgr
from magi_attention.common.enum import AttnMaskType
from magi_attention.common.ranges import AttnRanges
from magi_attention.config import DistAttnConfig
from magi_attention.functional.flex_flash_attn import flex_flash_attn_func
from magi_attention.testing import assert_close


def main():
    assert torch.cuda.is_available(), "CUDA is required."

    # 1) Initialize the process group from environment variables provided by torchrun
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group(backend="nccl", init_method="env://")

    # 2) Define sequence ranges: Q=4096, cross-attention K/V=256
    # Note: For cross-attention, K and V are replicated on each rank, not sharded.
    q_ranges = AttnRanges.from_ranges([[0, 4096]])
    k_ranges = AttnRanges.from_ranges([[0, 4096]])

    self_attn_q_ranges = AttnRanges.from_ranges([[0, 4096]])
    self_attn_k_ranges = AttnRanges.from_ranges([[0, 4096]])
    xattn_q_ranges = AttnRanges.from_ranges([[0, 4096]])
    xattn_k_ranges = AttnRanges.from_ranges([[0, 256]])

    total_seqlen_q = 4096
    total_seqlen_k = 4096
    total_seqlen_xattn_k = 256
    chunk_size = 256

    # 3) Build the DistAttnRuntimeMgr with the default process group as the communication group
    dist_attn_runtime_mgr = init_dist_attn_runtime_mgr(
        q_ranges=q_ranges,
        k_ranges=k_ranges,
        attn_mask_type=[AttnMaskType.FULL] * len(q_ranges),
        total_seqlen_q=total_seqlen_q,
        total_seqlen_k=total_seqlen_k,
        chunk_size=chunk_size,
        cp_group=dist.distributed_c10d._get_default_group(),
        cp_mesh=None,
        is_same_source=True,
        is_q_permutable=True,
        is_k_permutable=True,
        dist_attn_config=DistAttnConfig(),
    )

    # 4) Prepare input tensors and synchronize them across all ranks
    # Shapes are in the format [Sequence Length, Heads, Dimension] with H=1, D=128
    total_q = torch.randn(total_seqlen_q, 1, 128, device=device, dtype=torch.float16)
    total_k = torch.randn(total_seqlen_k, 1, 128, device=device, dtype=torch.float16)
    total_v = torch.randn(total_seqlen_k, 1, 128, device=device, dtype=torch.float16)
    xattn_k = torch.randn(total_seqlen_xattn_k, 1, 128, device=device, dtype=torch.float16)
    xattn_v = torch.randn(total_seqlen_xattn_k, 1, 128, device=device, dtype=torch.float16)

    dist.all_reduce(total_q)
    dist.all_reduce(total_k)
    dist.all_reduce(total_v)
    dist.all_reduce(xattn_k)
    dist.all_reduce(xattn_v)

    # 5) Establish a ground truth by calling the attention function with explicit ranges
    # do self-attn
    total_o_ref_self_attn, _ = flex_flash_attn_func(
        q=total_q,
        k=total_k,
        v=total_v,
        q_ranges=self_attn_q_ranges.to_tensor(device=device),
        k_ranges=self_attn_k_ranges.to_tensor(device=device),
        attn_type_map=torch.zeros(len(self_attn_q_ranges), dtype=torch.int32, device=device),
        max_seqlen_q=self_attn_q_ranges.max_seqlen,
        max_seqlen_k=self_attn_k_ranges.max_seqlen,
    )
    # do cross-attn
    total_o_ref_xattn, _ = flex_flash_attn_func(
        q=total_q,
        k=xattn_k,
        v=xattn_v,
        q_ranges=xattn_q_ranges.to_tensor(device=device),
        k_ranges=xattn_k_ranges.to_tensor(device=device),
        attn_type_map=torch.zeros(len(xattn_q_ranges), dtype=torch.int32, device=device),
        max_seqlen_q=xattn_q_ranges.max_seqlen,
        max_seqlen_k=xattn_k_ranges.max_seqlen,
    )

    # 6) Use host-side arguments: Q is dispatched, while K and V are replicated
    # dispatch q, k, v
    local_q = dist_attn_runtime_mgr.dispatch_qo(total_q)
    local_k = dist_attn_runtime_mgr.dispatch_kv(total_k)
    local_v = dist_attn_runtime_mgr.dispatch_kv(total_v)
    # do self-attn
    local_o_self_attn, _ = dist_attn_runtime_mgr.calc_attn(local_q, local_k, local_v)
    # do cross-attn
    host_xattn_attn_arg = dist_attn_runtime_mgr.get_xattn_args(
        xattn_q_ranges,
        xattn_k_ranges,
        attn_mask_type=[AttnMaskType.FULL] * len(xattn_k_ranges),
        return_host_only=True,
    )
    
    local_o_xattn, _ = flex_flash_attn_func(
        q=local_q,
        k=xattn_k,  # Replicated on all ranks
        v=xattn_v,  # Replicated on all ranks
        **host_xattn_attn_arg.to_ffa_args(),
    )
    total_o_self_attn = dist_attn_runtime_mgr.undispatch_qo(local_o_self_attn)
    total_o_xattn = dist_attn_runtime_mgr.undispatch_qo(local_o_xattn)

    assert_close(
        total_o_self_attn, total_o_ref_self_attn, test_case="self-attn forward (host args) vs ground truth", atol=1e-3, rtol=1e-2,
    )
    assert_close(
        total_o_xattn, total_o_ref_xattn, test_case="cross-attn forward (host args) vs ground truth", atol=1e-3, rtol=1e-2,
    )

    if dist.get_rank() == 0:
        print("Cross-attention outputs successfully aligned with ground truth. K/V are replicated per rank ✅")

    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Launch Commands

You can run the example using torchrun with the following commands:

Single-Node (4 GPUs):

torchrun --standalone --nproc_per_node=4 basic_xattn_torchrun.py

Single-Node (Explicit Rendezvous):

torchrun --nproc_per_node=4 --rdzv_backend=c10d --rdzv_endpoint=127.0.0.1:29500 basic_xattn_torchrun.py

Multi-Node (2 nodes, 4 GPUs each):

On node 0:

torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 \
  --rdzv_backend=c10d --rdzv_endpoint=<node0_host>:29500 \
  basic_xattn_torchrun.py

On node 1:

torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 \
  --rdzv_backend=c10d --rdzv_endpoint=<node0_host>:29500 \
  basic_xattn_torchrun.py

Key Points

• For cross-attention, the key (K) and value (V) tensors should be fully replicated on every rank.
• Only the query (Q) tensor is partitioned and managed by the dispatch_qo(...) and undispatch_qo(...) functions.

[zelingao98]

# basic_xattn_torchrun.py

import os
import torch
import torch.distributed as dist

from magi_attention import init_dist_attn_runtime_mgr
from magi_attention.common.enum import AttnMaskType
from magi_attention.common.ranges import AttnRanges
from magi_attention.config import DistAttnConfig
from magi_attention.functional.flex_flash_attn import flex_flash_attn_func
from magi_attention.api.functools import compute_pad_size, pad_at_dim, unpad_at_dim, apply_padding
# from magi_attention.testing import assert_close
from magi_attention.api import magi_attn_interface

import torch.nn.functional as torch_F
from einops import rearrange


def compute_ranges(seq_len_list):
    ranges = []
    start = 0
    for seq_len in seq_len_list:
        ranges.append([start, start + seq_len])
        start += seq_len
    return ranges


def main_selfattn_crossattn():
    assert torch.cuda.is_available(), "CUDA is required."

    # 1) Initialize the process group from environment variables provided by torchrun
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group(backend="nccl", init_method="env://")

    # 2) Define sequence ranges: Q=4096, cross-attention K/V=256
    batch_size = 1
    q_seq_len = 4096
    kv_seq_len = 512
    num_head = 12
    head_dim = 128
    chunk_size = 256

    # create data on each rank
    local_q = torch.randn(batch_size, q_seq_len, num_head, head_dim, dtype=torch.bfloat16, device=device)
    local_xattn_k = torch.randn(batch_size, kv_seq_len, num_head, head_dim, dtype=torch.bfloat16, device=device)
    local_xattn_v = torch.randn(batch_size, kv_seq_len, num_head, head_dim, dtype=torch.bfloat16, device=device)

    # accumulate all q k v on all ranks
    # dist.all_reduce(total_q)
    # dist.all_reduce(xattn_k)
    # dist.all_reduce(xattn_v)
    # todo now we gather all datas from each rank
    world_size = dist.get_world_size()
    q_gather = [torch.zeros([batch_size, q_seq_len, num_head, head_dim], dtype=torch.bfloat16, device=device) for _ in range(world_size)]
    k_gather = [torch.zeros([batch_size, kv_seq_len, num_head, head_dim], dtype=torch.bfloat16, device=device) for _ in range(world_size)]
    v_gather = [torch.zeros([batch_size, kv_seq_len, num_head, head_dim], dtype=torch.bfloat16, device=device) for _ in range(world_size)]
    torch.distributed.all_gather(q_gather, local_q)
    torch.distributed.all_gather(k_gather, local_xattn_k)
    torch.distributed.all_gather(v_gather, local_xattn_v)
    batch_size = batch_size * world_size
    total_q = torch.cat(q_gather, dim=0)
    xattn_k = torch.cat(k_gather, dim=0)
    xattn_v = torch.cat(v_gather, dim=0)

    # note for sdpa self-attention
    total_q_sdpa = rearrange(total_q, "b s n d -> b n s d")
    xattn_k_sdpa = rearrange(xattn_k, "b s n d -> b n s d")
    xattn_v_sdpa = rearrange(xattn_v, "b s n d -> b n s d")
    total_o_ref_sdpa_selfattn = torch_F.scaled_dot_product_attention(total_q_sdpa, total_q_sdpa, total_q_sdpa)
    total_o_ref_sdpa_selfattn = rearrange(total_o_ref_sdpa_selfattn, "b n s d -> (b s) n d")
    total_o_ref_sdpa_crossattn = torch_F.scaled_dot_product_attention(total_q_sdpa, xattn_k_sdpa, xattn_v_sdpa)
    total_o_ref_sdpa_crossattn = rearrange(total_o_ref_sdpa_crossattn, "b n s d -> (b s) n d")

    # note for single self-attention
    total_q = rearrange(total_q, "b s n d -> (b s) n d")
    xattn_k = rearrange(xattn_k, "b s n d -> (b s) n d")
    xattn_v = rearrange(xattn_v, "b s n d -> (b s) n d")
    q_ranges_selfattn = AttnRanges.from_ranges(compute_ranges([q_seq_len]*batch_size))
    kv_ranges_selfattn = AttnRanges.from_ranges(compute_ranges([q_seq_len]*batch_size))
    total_o_ref_selfattn, _ = flex_flash_attn_func(
        q=total_q,
        k=total_q,
        v=total_q,
        q_ranges=q_ranges_selfattn.to_tensor(device=device),
        k_ranges=kv_ranges_selfattn.to_tensor(device=device),
        attn_type_map=torch.zeros(len(q_ranges_selfattn), dtype=torch.int32, device=device),
        max_seqlen_q=q_ranges_selfattn.max_seqlen,
        max_seqlen_k=kv_ranges_selfattn.max_seqlen,
    )

    xattn_q_ranges = AttnRanges.from_ranges(compute_ranges([q_seq_len]*batch_size))
    xattn_k_ranges = AttnRanges.from_ranges(compute_ranges([kv_seq_len]*batch_size))
    total_o_ref_crossattn, _ = flex_flash_attn_func(
        q=total_q,
        k=xattn_k,
        v=xattn_v,
        q_ranges=xattn_q_ranges.to_tensor(device=device),
        k_ranges=xattn_k_ranges.to_tensor(device=device),
        attn_type_map=torch.zeros(len(xattn_q_ranges), dtype=torch.int32, device=device),
        max_seqlen_q=xattn_q_ranges.max_seqlen,
        max_seqlen_k=xattn_k_ranges.max_seqlen,
    )

    if dist.get_rank() == 0:
        print("single self attention")
        print("self attention")
        print('mean diff', (total_o_ref_selfattn - total_o_ref_sdpa_selfattn).abs().mean())
        print('max diff', (total_o_ref_selfattn - total_o_ref_sdpa_selfattn).abs().max())
        print("cross attention")
        print('mean diff', (total_o_ref_crossattn - total_o_ref_sdpa_crossattn).abs().mean())
        print('max diff', (total_o_ref_crossattn - total_o_ref_sdpa_crossattn).abs().max())

    # note for cp cross-attention
    xattn_q_ranges = AttnRanges.from_ranges(compute_ranges([q_seq_len]*batch_size))
    xattn_k_ranges = AttnRanges.from_ranges(compute_ranges([kv_seq_len]*batch_size))
    xattn_mask_type = [AttnMaskType.FULL] * len(xattn_k_ranges)

    q_ranges = AttnRanges.from_ranges(compute_ranges([q_seq_len]*batch_size))
    k_ranges = AttnRanges.from_ranges(compute_ranges([q_seq_len]*batch_size))
    total_seqlen_q = q_seq_len * batch_size
    total_seqlen_k = q_seq_len * batch_size
    attn_mask_type = [AttnMaskType.FULL] * len(q_ranges)
    
    pad_size = compute_pad_size(total_seqlen_q, cp_size=dist.get_world_size(), chunk_size=chunk_size)
    q_ranges, k_ranges, attn_mask_type = apply_padding(q_ranges, k_ranges, attn_mask_type, total_seqlen_q, pad_size)
    xattn_q_ranges, xattn_k_ranges, xattn_mask_type = apply_padding(xattn_q_ranges, xattn_k_ranges, xattn_mask_type, total_seqlen_q, pad_size)
    total_seqlen_q += pad_size
    total_seqlen_k += pad_size

    dist_attn_runtime_mgr = init_dist_attn_runtime_mgr(
        q_ranges=q_ranges,
        k_ranges=k_ranges,
        attn_mask_type=attn_mask_type,
        total_seqlen_q=total_seqlen_q,
        total_seqlen_k=total_seqlen_k,
        chunk_size=chunk_size,
        cp_group=dist.distributed_c10d._get_default_group(),
        cp_mesh=None,
        is_same_source=True,
        is_q_permutable=True,
        is_k_permutable=True,
        dist_attn_config=DistAttnConfig(),
    )

    # dispatch q, k, v
    total_q = pad_at_dim(total_q, dim=0, pad_size=pad_size)
    local_q = dist_attn_runtime_mgr.dispatch_qo(total_q)
    local_o_selfattn, _ = dist_attn_runtime_mgr.calc_attn(local_q, local_q, local_q)
    total_o_cp_selfattn = dist_attn_runtime_mgr.undispatch_qo(local_o_selfattn)
    total_o_cp_selfattn = unpad_at_dim(total_o_cp_selfattn, dim=0, pad_size=pad_size)

    host_xattn_attn_arg = dist_attn_runtime_mgr.get_xattn_args(
        xattn_q_ranges,
        xattn_k_ranges,
        attn_mask_type=xattn_mask_type,
        return_host_only=True,
    )

    local_q = dist_attn_runtime_mgr.dispatch_qo(total_q)
    local_o_cp_crossattn, _ = flex_flash_attn_func(
        q=local_q,
        k=xattn_k,  # Replicated on all ranks
        v=xattn_v,  # Replicated on all ranks
        **host_xattn_attn_arg.to_ffa_args(),
    )
    total_o_cp_crossattn = dist_attn_runtime_mgr.undispatch_qo(local_o_cp_crossattn)
    total_o_cp_crossattn = unpad_at_dim(total_o_cp_crossattn, dim=0, pad_size=pad_size)

    if dist.get_rank() == 0:
        print("cp self attention")
        print("self attention")
        print('mean diff', (total_o_cp_selfattn - total_o_ref_sdpa_selfattn).abs().mean())
        print('max diff', (total_o_cp_selfattn - total_o_ref_sdpa_selfattn).abs().max())
        print("cross attention")
        print('mean diff', (total_o_cp_crossattn - total_o_ref_sdpa_crossattn).abs().mean())
        print('max diff', (total_o_cp_crossattn - total_o_ref_sdpa_crossattn).abs().max())

if __name__ == "__main__":
    main_selfattn_crossattn()

How to run

torchrun --standalone --nproc_per_node=4 basic_xattn_torchrun.py

Confusion

Result Precision

My precision is as follows:

Everything is okay on my side with magiattention. I just curious whether this result is right or not.

Run with Context Parallel

When we train the model with magiattention of cp mode, we firstly gather all data and then input then into attention.
My question is whether such an operation will occur in each rank as each rank has its own input.