vllm_omni.diffusion.models.magi2.attention ¶

Native MAGI-2 packed attention with sinks and Ulysses exchange.

The sink and context-parallel math is adapted from SandAI's Apache-2.0 MAGI-2 preview implementation. This version uses vLLM's bundled FlashAttention extension and vLLM-Omni's existing Ulysses process group; the PyTorch path is an exact, portable oracle for small tests.

logger `module-attribute` ¶

logger = logging.getLogger(__name__)

Magi2PackedAttentionKernel ¶

Bases: Module

Model kernel plugged into the shared diffusion Attention contract.

softcap `instance-attribute` ¶

softcap = softcap

forward ¶

forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> Tensor

VarlenHandler `dataclass` ¶

Packed-sequence metadata consumed by MAGI-2 attention.

cu_seqlens_k `instance-attribute` ¶

cu_seqlens_k: Tensor | None

cu_seqlens_q `instance-attribute` ¶

cu_seqlens_q: Tensor | None

max_seqlen_k `class-attribute` `instance-attribute` ¶

max_seqlen_k: int | None = None

max_seqlen_q `class-attribute` `instance-attribute` ¶

max_seqlen_q: int | None = None

resolved ¶

resolved(
    q_tokens: int, k_tokens: int
) -> tuple[Tensor, Tensor, int, int]

apply_rotary_emb ¶

apply_rotary_emb(
    x: Tensor,
    cos: Tensor,
    sin: Tensor,
    *,
    interleaved: bool = False,
) -> Tensor

Apply the released element-wise RoPE layout to [...,H,D].

correct_out_lse_with_sink ¶

correct_out_lse_with_sink(
    out: Tensor, lse: Tensor, sink: Tensor | None
) -> tuple[Tensor, Tensor]

Add zero-valued attention sinks to an already-computed softmax.

FlashAttention returns out[T,H,D] and conventionally lse[H,T]. A MAGI sink contains additional logits [num_sink,H] whose values are zero vectors, so only the denominator and LSE change.

packed_attention_with_sink ¶

packed_attention_with_sink(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    varlen: VarlenHandler,
    *,
    softcap: float = -1.0,
    sink: Tensor | None = None,
) -> Tensor

Run packed attention on one rank after Ulysses head exchange.

rotate_half ¶

rotate_half(x: Tensor, interleaved: bool = False) -> Tensor

torch_varlen_attention_with_sink ¶

torch_varlen_attention_with_sink(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    *,
    cu_seqlens_q: Tensor,
    cu_seqlens_k: Tensor,
    softcap: float = -1.0,
    sink: Tensor | None = None,
) -> Tensor

Reference packed attention, including GQA and sink logits.

ulysses_packed_attention_with_sink ¶

ulysses_packed_attention_with_sink(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    varlen: VarlenHandler,
    split_sizes: list[int] | Tensor,
    *,
    softcap: float = -1.0,
    sink: Tensor | None = None,
    group: Magi2ParallelGroup | None = None,
) -> Tensor

MAGI-2 attention with overlapping Ulysses CP/head exchange.