Hongsheng Liu
d4096e01fe
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
294 lines
11 KiB
Python
294 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
# Copyright (c) 2026 SandAI. All Rights Reserved.
|
|
|
|
"""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.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from functools import cache
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
|
from vllm_omni.diffusion.attention.backends.utils.fa import (
|
|
resolve_vllm_flash_attn_version,
|
|
vllm_flash_attn_varlen_with_lse,
|
|
)
|
|
|
|
from .parallel import (
|
|
Magi2ParallelGroup,
|
|
get_magi2_ulysses_group,
|
|
scatter_heads_gather_seqlen,
|
|
scatter_seqlen_gather_heads,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@cache
|
|
def _resolve_flash_attn_version() -> int:
|
|
"""Apply MAGI-2's operator override to the shared FA2/FA3/FA4 resolver."""
|
|
|
|
requested = os.environ.get("MAGI2_FLASH_ATTN_VERSION")
|
|
version = resolve_vllm_flash_attn_version(requested)
|
|
logger.info("MAGI-2 selected FlashAttention %d", version)
|
|
return version
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VarlenHandler:
|
|
"""Packed-sequence metadata consumed by MAGI-2 attention."""
|
|
|
|
cu_seqlens_q: torch.Tensor | None
|
|
cu_seqlens_k: torch.Tensor | None
|
|
max_seqlen_q: int | None = None
|
|
max_seqlen_k: int | None = None
|
|
|
|
def resolved(self, q_tokens: int, k_tokens: int) -> tuple[torch.Tensor, torch.Tensor, int, int]:
|
|
if self.cu_seqlens_q is None:
|
|
cu_q = torch.tensor([0, q_tokens], device="cpu", dtype=torch.int32)
|
|
else:
|
|
cu_q = self.cu_seqlens_q
|
|
if self.cu_seqlens_k is None:
|
|
cu_k = torch.tensor([0, k_tokens], device="cpu", dtype=torch.int32)
|
|
else:
|
|
cu_k = self.cu_seqlens_k
|
|
max_q = self.max_seqlen_q
|
|
max_k = self.max_seqlen_k
|
|
if max_q is None:
|
|
max_q = int(torch.diff(cu_q).max().item()) if cu_q.numel() > 1 else 0
|
|
if max_k is None:
|
|
max_k = int(torch.diff(cu_k).max().item()) if cu_k.numel() > 1 else 0
|
|
return cu_q, cu_k, max_q, max_k
|
|
|
|
|
|
def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
|
|
if not interleaved:
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
x1, x2 = x[..., ::2], x[..., 1::2]
|
|
return torch.stack((-x2, x1), dim=-1).flatten(-2)
|
|
|
|
|
|
def apply_rotary_emb(
|
|
x: torch.Tensor,
|
|
cos: torch.Tensor,
|
|
sin: torch.Tensor,
|
|
*,
|
|
interleaved: bool = False,
|
|
) -> torch.Tensor:
|
|
"""Apply the released element-wise RoPE layout to ``[...,H,D]``."""
|
|
|
|
rotary_dim = cos.shape[-1] * 2
|
|
if rotary_dim > x.shape[-1]:
|
|
raise ValueError(f"RoPE dimension {rotary_dim} exceeds head dimension {x.shape[-1]}")
|
|
if interleaved:
|
|
cos = cos.unsqueeze(-2).repeat_interleave(2, dim=-1)
|
|
sin = sin.unsqueeze(-2).repeat_interleave(2, dim=-1)
|
|
else:
|
|
cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2)
|
|
sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)
|
|
rotated = x[..., :rotary_dim] * cos + rotate_half(x[..., :rotary_dim], interleaved) * sin
|
|
return torch.cat((rotated, x[..., rotary_dim:]), dim=-1)
|
|
|
|
|
|
def correct_out_lse_with_sink(
|
|
out: torch.Tensor,
|
|
lse: torch.Tensor,
|
|
sink: torch.Tensor | None,
|
|
) -> tuple[torch.Tensor, torch.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.
|
|
"""
|
|
|
|
if sink is None or sink.numel() == 0:
|
|
return out, lse
|
|
if out.ndim != 3 or lse.ndim != 2 or sink.ndim != 2:
|
|
raise ValueError(
|
|
"expected out[T,H,D], lse[H,T], sink[num_sink,H], got "
|
|
f"{tuple(out.shape)}, {tuple(lse.shape)}, {tuple(sink.shape)}"
|
|
)
|
|
old_lse = lse.float().transpose(0, 1)
|
|
sink_lse = torch.logsumexp(sink.float(), dim=0).unsqueeze(0)
|
|
if old_lse.shape[-1] != sink_lse.shape[-1]:
|
|
raise ValueError("attention sink and FlashAttention head counts differ")
|
|
new_lse = torch.logaddexp(old_lse, sink_lse)
|
|
delta = old_lse - new_lse
|
|
delta = torch.where(torch.isfinite(delta), delta, torch.full_like(delta, -torch.inf))
|
|
corrected = out * torch.exp(delta).unsqueeze(-1).to(out.dtype)
|
|
return corrected, new_lse.transpose(0, 1).contiguous()
|
|
|
|
|
|
def _repeat_kv_heads(tensor: torch.Tensor, query_heads: int) -> torch.Tensor:
|
|
kv_heads = tensor.shape[1]
|
|
if kv_heads == query_heads:
|
|
return tensor
|
|
if query_heads % kv_heads:
|
|
raise ValueError(f"query heads {query_heads} must be divisible by KV heads {kv_heads}")
|
|
return tensor.repeat_interleave(query_heads // kv_heads, dim=1)
|
|
|
|
|
|
def torch_varlen_attention_with_sink(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
*,
|
|
cu_seqlens_q: torch.Tensor,
|
|
cu_seqlens_k: torch.Tensor,
|
|
softcap: float = -1.0,
|
|
sink: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Reference packed attention, including GQA and sink logits."""
|
|
|
|
if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:
|
|
raise ValueError("packed attention expects q/k/v shaped [tokens,heads,dim]")
|
|
if cu_seqlens_q.numel() != cu_seqlens_k.numel():
|
|
raise ValueError("query and key cumulative-length arrays must contain the same batch count")
|
|
output = torch.empty_like(q)
|
|
scale = q.shape[-1] ** -0.5
|
|
for batch_idx in range(cu_seqlens_q.numel() - 1):
|
|
q_start, q_end = (int(v) for v in cu_seqlens_q[batch_idx : batch_idx + 2].tolist())
|
|
k_start, k_end = (int(v) for v in cu_seqlens_k[batch_idx : batch_idx + 2].tolist())
|
|
q_part = q[q_start:q_end].float()
|
|
k_part = _repeat_kv_heads(k[k_start:k_end], q.shape[1]).float()
|
|
v_part = _repeat_kv_heads(v[k_start:k_end], q.shape[1]).float()
|
|
scores = torch.einsum("qhd,khd->hqk", q_part, k_part) * scale
|
|
if softcap > 0:
|
|
scores = softcap * torch.tanh(scores / softcap)
|
|
if sink is not None and sink.numel() > 0:
|
|
sink_scores = sink.float().transpose(0, 1).unsqueeze(1).expand(-1, q_part.shape[0], -1)
|
|
probabilities = torch.softmax(torch.cat((scores, sink_scores), dim=-1), dim=-1)[..., : k_part.shape[0]]
|
|
else:
|
|
probabilities = torch.softmax(scores, dim=-1)
|
|
output[q_start:q_end] = torch.einsum("hqk,khd->qhd", probabilities, v_part).to(output.dtype)
|
|
return output
|
|
|
|
|
|
def packed_attention_with_sink(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
varlen: VarlenHandler,
|
|
*,
|
|
softcap: float = -1.0,
|
|
sink: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Run packed attention on one rank after Ulysses head exchange."""
|
|
|
|
cu_q, cu_k, max_q, max_k = varlen.resolved(q.shape[0], k.shape[0])
|
|
cu_q = cu_q.to(device=q.device, dtype=torch.int32).contiguous()
|
|
cu_k = cu_k.to(device=q.device, dtype=torch.int32).contiguous()
|
|
if q.is_cuda:
|
|
out, lse = vllm_flash_attn_varlen_with_lse(
|
|
q,
|
|
k,
|
|
v,
|
|
cu_seqlens_q=cu_q,
|
|
cu_seqlens_k=cu_k,
|
|
max_seqlen_q=max_q,
|
|
max_seqlen_k=max_k,
|
|
softcap=softcap,
|
|
deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1",
|
|
fa_version=_resolve_flash_attn_version(),
|
|
)
|
|
return correct_out_lse_with_sink(out, lse, sink)[0]
|
|
return torch_varlen_attention_with_sink(
|
|
q,
|
|
k,
|
|
v,
|
|
cu_seqlens_q=cu_q,
|
|
cu_seqlens_k=cu_k,
|
|
softcap=softcap,
|
|
sink=sink,
|
|
)
|
|
|
|
|
|
def ulysses_packed_attention_with_sink(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
varlen: VarlenHandler,
|
|
split_sizes: list[int] | torch.Tensor,
|
|
*,
|
|
softcap: float = -1.0,
|
|
sink: torch.Tensor | None = None,
|
|
group: Magi2ParallelGroup | None = None,
|
|
) -> torch.Tensor:
|
|
"""MAGI-2 attention with overlapping Ulysses CP/head exchange."""
|
|
|
|
group = group or get_magi2_ulysses_group()
|
|
if isinstance(split_sizes, torch.Tensor):
|
|
split_sizes = [int(v) for v in split_sizes.detach().cpu().tolist()]
|
|
if group.world_size > 1:
|
|
q, k, v = scatter_heads_gather_seqlen((q, k, v), split_sizes, group)
|
|
if sink is not None:
|
|
if sink.shape[-1] % group.world_size:
|
|
raise ValueError("attention sink heads must divide across Ulysses ranks")
|
|
sink = torch.chunk(sink, group.world_size, dim=-1)[group.rank].contiguous()
|
|
output = packed_attention_with_sink(q, k, v, varlen, softcap=softcap, sink=sink)
|
|
if group.world_size > 1:
|
|
output = scatter_seqlen_gather_heads(output.contiguous(), split_sizes, group)
|
|
assert isinstance(output, torch.Tensor)
|
|
return output
|
|
|
|
|
|
class Magi2PackedAttentionKernel(nn.Module):
|
|
"""Model kernel plugged into the shared diffusion Attention contract."""
|
|
|
|
def __init__(self, softcap: float) -> None:
|
|
super().__init__()
|
|
self.softcap = softcap
|
|
|
|
def forward(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AttentionMetadata | None = None,
|
|
) -> torch.Tensor:
|
|
if attn_metadata is None:
|
|
raise ValueError("MAGI-2 packed attention requires attention metadata")
|
|
varlen = attn_metadata.extra.get("magi2_varlen")
|
|
split_sizes = attn_metadata.extra.get("magi2_split_sizes")
|
|
sink = attn_metadata.extra.get("magi2_sink")
|
|
if not isinstance(varlen, VarlenHandler):
|
|
raise TypeError("magi2_varlen must be a VarlenHandler")
|
|
if not isinstance(split_sizes, (list, torch.Tensor)):
|
|
raise TypeError("magi2_split_sizes must be a list or tensor")
|
|
if sink is not None and not isinstance(sink, torch.Tensor):
|
|
raise TypeError("magi2_sink must be a tensor or None")
|
|
return ulysses_packed_attention_with_sink(
|
|
query,
|
|
key,
|
|
value,
|
|
varlen,
|
|
split_sizes,
|
|
softcap=self.softcap,
|
|
sink=sink,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"VarlenHandler",
|
|
"Magi2PackedAttentionKernel",
|
|
"apply_rotary_emb",
|
|
"correct_out_lse_with_sink",
|
|
"packed_attention_with_sink",
|
|
"torch_varlen_attention_with_sink",
|
|
"ulysses_packed_attention_with_sink",
|
|
]
|