[diffusion] refactor: move SLA to attention_backend folder (#17020)
This commit is contained in:
HuangJi
@@ -14,10 +14,7 @@ from sglang.multimodal_gen.runtime.layers.attention.layer import (
|
||||
USPAttention,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
|
||||
from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import (
|
||||
MinimalA2AAttnOp,
|
||||
SparseLinearAttention,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import MinimalA2AAttnOp
|
||||
|
||||
__all__ = [
|
||||
"USPAttention",
|
||||
@@ -25,7 +22,6 @@ __all__ = [
|
||||
"UlyssesAttention",
|
||||
"UlyssesAttention_VSA",
|
||||
"MinimalA2AAttnOp",
|
||||
"SparseLinearAttention",
|
||||
"AttentionBackend",
|
||||
"AttentionMetadata",
|
||||
"AttentionMetadataBuilder",
|
||||
|
||||
@@ -0,0 +1,363 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from turbo_layer.py for Attention Backend integration
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class SparseLinearAttentionBackend(AttentionBackend):
|
||||
"""Sparse Linear Attention Backend for efficient attention computation."""
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [64, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_enum() -> AttentionBackendEnum:
|
||||
return AttentionBackendEnum.SLA_ATTN
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["SparseLinearAttentionImpl"]:
|
||||
return SparseLinearAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["SparseLinearAttentionMetadata"]:
|
||||
return SparseLinearAttentionMetadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["SparseLinearAttentionMetadataBuilder"]:
|
||||
return SparseLinearAttentionMetadataBuilder
|
||||
|
||||
|
||||
@dataclass
|
||||
class SparseLinearAttentionMetadata(AttentionMetadata):
|
||||
"""Metadata for Sparse Linear Attention computation."""
|
||||
|
||||
# Basic attention parameters
|
||||
current_timestep: int
|
||||
|
||||
# Sparse attention configuration
|
||||
topk_ratio: float = 0.1
|
||||
|
||||
|
||||
class SparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder):
|
||||
"""Builder for SparseLinearAttentionMetadata."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def build(
|
||||
self,
|
||||
current_timestep: int,
|
||||
topk_ratio: float = 0.1,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> SparseLinearAttentionMetadata:
|
||||
return SparseLinearAttentionMetadata(
|
||||
current_timestep=current_timestep,
|
||||
topk_ratio=topk_ratio,
|
||||
)
|
||||
|
||||
|
||||
class SparseLinearAttentionImpl(AttentionImpl, nn.Module):
|
||||
"""Implementation of sparse linear attention for the backend."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool = False,
|
||||
softmax_scale: float | None = None,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
# SLA-specific parameters - matched to TurboDiffusion defaults
|
||||
topk_ratio: float = 0.1, # TurboDiffusion uses topk=0.1
|
||||
feature_map: str = "softmax",
|
||||
BLKQ: int = 128, # TurboDiffusion uses BLKQ=128
|
||||
BLKK: int = 64, # TurboDiffusion uses BLKK=64
|
||||
use_bf16: bool = True,
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
nn.Module.__init__(self)
|
||||
|
||||
# SLA-specific config
|
||||
self.topk_ratio = topk_ratio
|
||||
self.BLKQ = BLKQ
|
||||
self.BLKK = BLKK
|
||||
self.dtype = torch.bfloat16 if use_bf16 else torch.float16
|
||||
|
||||
# Learnable linear projection for combining sparse + linear attention
|
||||
self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32)
|
||||
|
||||
# Feature map for linear attention
|
||||
# Type annotation for callables
|
||||
self.feature_map_q: Callable[[torch.Tensor], torch.Tensor]
|
||||
self.feature_map_k: Callable[[torch.Tensor], torch.Tensor]
|
||||
if feature_map == "elu":
|
||||
self.feature_map_q = lambda x: F.elu(x) + 1
|
||||
self.feature_map_k = lambda x: F.elu(x) + 1
|
||||
elif feature_map == "relu":
|
||||
self.feature_map_q = F.relu
|
||||
self.feature_map_k = F.relu
|
||||
elif feature_map == "softmax":
|
||||
self.feature_map_q = lambda x: F.softmax(x, dim=-1)
|
||||
self.feature_map_k = lambda x: F.softmax(x, dim=-1)
|
||||
else:
|
||||
raise ValueError(f"Unknown feature map: {feature_map}")
|
||||
|
||||
self._init_weights()
|
||||
|
||||
def _init_weights(self) -> None:
|
||||
"""Initialize projection weights to zero for residual-like behavior."""
|
||||
with torch.no_grad():
|
||||
nn.init.zeros_(self.proj_l.weight)
|
||||
nn.init.zeros_(self.proj_l.bias) # type: ignore[arg-type]
|
||||
|
||||
def _calc_linear_attention_with_torch(self, q, k, v):
|
||||
kv = torch.matmul(k.transpose(-1, -2), v)
|
||||
k_sum = torch.sum(k, dim=-2, keepdim=True)
|
||||
return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: SparseLinearAttentionMetadata = None,
|
||||
) -> torch.Tensor:
|
||||
"""Forward pass for sparse linear attention.
|
||||
|
||||
Args:
|
||||
query: query tensor of shape (B, H, L, D)
|
||||
key: key tensor of shape (B, H, L, D)
|
||||
value: value tensor of shape (B, H, L, D)
|
||||
attn_metadata: attention metadata containing configuration
|
||||
Returns:
|
||||
output tensor of shape (B, H, L, D)
|
||||
"""
|
||||
dtype = query.dtype
|
||||
|
||||
# Transpose for computation
|
||||
query = query.transpose(1, 2).contiguous()
|
||||
key = key.transpose(1, 2).contiguous()
|
||||
value = value.transpose(1, 2).contiguous()
|
||||
|
||||
# Get sparse attention map
|
||||
sparse_map, lut, real_topk = get_block_map(
|
||||
query, key, topk_ratio=self.topk_ratio, BLKQ=self.BLKQ, BLKK=self.BLKK
|
||||
)
|
||||
|
||||
# Convert to computation dtype
|
||||
query = query.to(self.dtype)
|
||||
key = key.to(self.dtype)
|
||||
value = value.to(self.dtype)
|
||||
|
||||
# Sparse attention computation
|
||||
o_s = _attention.apply(
|
||||
query, key, value, sparse_map, lut, real_topk, self.BLKQ, self.BLKK
|
||||
)
|
||||
|
||||
# Apply feature maps
|
||||
query = self.feature_map_q(query).contiguous().to(self.dtype) # c_q
|
||||
key = self.feature_map_k(key).contiguous().to(self.dtype) # c_k
|
||||
# Linear attention computation
|
||||
o_l = self._calc_linear_attention_with_torch(query, key, value)
|
||||
|
||||
# Apply projection and combine results
|
||||
with torch.amp.autocast("cuda", dtype=self.dtype):
|
||||
o_l = self.proj_l(o_l)
|
||||
|
||||
# Combine sparse and linear attention
|
||||
output = (o_s + o_l).to(dtype).transpose(1, 2)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class _attention(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None):
|
||||
assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
|
||||
assert k_block_id.is_contiguous() and lut.is_contiguous()
|
||||
|
||||
# We recommend the following two settings
|
||||
assert BLOCK_M == 64 or BLOCK_M == 128
|
||||
assert BLOCK_N == 64
|
||||
|
||||
B, H, L, D = q.shape
|
||||
if qk_scale is None:
|
||||
qk_scale = D**-0.5
|
||||
|
||||
M_BLOCKS = triton.cdiv(L, BLOCK_M)
|
||||
|
||||
o_s = torch.empty_like(v)
|
||||
lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32)
|
||||
|
||||
grid = (M_BLOCKS, B * H)
|
||||
_attn_fwd[grid](
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
qk_scale,
|
||||
topk,
|
||||
lut,
|
||||
lse,
|
||||
o_s,
|
||||
L,
|
||||
M_BLOCKS,
|
||||
D,
|
||||
BLOCK_M,
|
||||
BLOCK_N,
|
||||
num_warps=4 if q.shape[-1] == 64 else 8,
|
||||
num_stages=3,
|
||||
)
|
||||
|
||||
ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s)
|
||||
ctx.qk_scale = qk_scale
|
||||
ctx.topk = topk
|
||||
ctx.BLOCK_M = BLOCK_M
|
||||
ctx.BLOCK_N = BLOCK_N
|
||||
return o_s
|
||||
|
||||
|
||||
def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
|
||||
arg_k = k - torch.mean(
|
||||
k, dim=-2, keepdim=True
|
||||
) # smooth-k technique in SageAttention
|
||||
pooled_qblocks = mean_pool(q, BLKQ)
|
||||
pooled_kblocks = mean_pool(arg_k, BLKK)
|
||||
pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
|
||||
|
||||
K = pooled_score.shape[-1]
|
||||
topk = min(K, int(topk_ratio * K))
|
||||
lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
|
||||
|
||||
sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
|
||||
sparse_map.scatter_(-1, lut, 1)
|
||||
return sparse_map, lut, topk
|
||||
|
||||
|
||||
def mean_pool(x, BLK):
|
||||
assert x.is_contiguous()
|
||||
|
||||
B, H, L, D = x.shape
|
||||
L_BLOCKS = (L + BLK - 1) // BLK
|
||||
x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
|
||||
|
||||
grid = (L_BLOCKS, B * H)
|
||||
compress_kernel[grid](x, x_mean, L, D, BLK)
|
||||
return x_mean
|
||||
|
||||
|
||||
@triton.jit
|
||||
def compress_kernel(
|
||||
X,
|
||||
XM,
|
||||
L: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
BLOCK_L: tl.constexpr,
|
||||
):
|
||||
idx_l = tl.program_id(0)
|
||||
idx_bh = tl.program_id(1)
|
||||
|
||||
offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
|
||||
offs_d = tl.arange(0, D)
|
||||
|
||||
x_offset = idx_bh * L * D
|
||||
xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
|
||||
x = tl.load(
|
||||
X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L
|
||||
)
|
||||
|
||||
nx = min(BLOCK_L, L - idx_l * BLOCK_L)
|
||||
x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
|
||||
tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_fwd(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
qk_scale: tl.constexpr,
|
||||
topk: tl.constexpr,
|
||||
LUT,
|
||||
LSE,
|
||||
OS,
|
||||
L: tl.constexpr,
|
||||
M_BLOCKS: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
idx_m = tl.program_id(0).to(tl.int64)
|
||||
idx_bh = tl.program_id(1).to(tl.int64)
|
||||
|
||||
qkv_offset = idx_bh * L * D
|
||||
lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk
|
||||
lse_offset = idx_bh * L
|
||||
offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
offs_n = tl.arange(0, BLOCK_N)
|
||||
offs_d = tl.arange(0, D)
|
||||
|
||||
Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
|
||||
K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None]
|
||||
V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :]
|
||||
OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
|
||||
LUT_ptr = LUT + lut_offset
|
||||
LSE_ptrs = LSE + lse_offset + offs_m
|
||||
|
||||
m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
|
||||
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
|
||||
o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
|
||||
|
||||
q = tl.load(Q_ptrs, mask=offs_m[:, None] < L)
|
||||
for block_idx in tl.range(topk):
|
||||
idx_n = tl.load(LUT_ptr + block_idx)
|
||||
n_mask = offs_n < L - idx_n * BLOCK_N
|
||||
|
||||
k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
|
||||
qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # = 1 / ln(2)
|
||||
if L - idx_n * BLOCK_N < BLOCK_N:
|
||||
qk = tl.where(n_mask[None, :], qk, float("-inf"))
|
||||
|
||||
v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
|
||||
local_m = tl.max(qk, 1)
|
||||
new_m = tl.maximum(m_i, local_m)
|
||||
qk = qk - new_m[:, None]
|
||||
|
||||
p = tl.math.exp2(qk)
|
||||
l_ij = tl.sum(p, 1)
|
||||
alpha = tl.math.exp2(m_i - new_m)
|
||||
o_s = o_s * alpha[:, None]
|
||||
o_s += tl.dot(p.to(v.dtype), v)
|
||||
|
||||
l_i = l_i * alpha + l_ij
|
||||
m_i = new_m
|
||||
|
||||
o_s = o_s / l_i[:, None]
|
||||
tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
|
||||
|
||||
m_i += tl.math.log2(l_i)
|
||||
tl.store(LSE_ptrs, m_i, mask=offs_m < L)
|
||||
@@ -1,18 +1,31 @@
|
||||
# copy and modify from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/rcm/utils/a2a_cp.py and https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py
|
||||
|
||||
from typing import Any, Callable, List, Tuple, Union
|
||||
from typing import Any, Callable, List, Tuple, Type, Union
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
from torch.distributed import ProcessGroup
|
||||
from torch.nn import Module
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
AttentionImpl,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn import (
|
||||
SparseLinearAttentionBackend,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import (
|
||||
ForwardContext,
|
||||
get_forward_context,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.utils import get_compute_dtype
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def post_all2all(local_seq_2_local_head, seq_world_size):
|
||||
def post_func(input):
|
||||
@@ -110,128 +123,6 @@ def async_a2a_communicate(
|
||||
return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_fwd(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
qk_scale: tl.constexpr,
|
||||
topk: tl.constexpr,
|
||||
LUT,
|
||||
LSE,
|
||||
OS,
|
||||
L: tl.constexpr,
|
||||
M_BLOCKS: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
idx_m = tl.program_id(0).to(tl.int64)
|
||||
idx_bh = tl.program_id(1).to(tl.int64)
|
||||
|
||||
qkv_offset = idx_bh * L * D
|
||||
lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk
|
||||
lse_offset = idx_bh * L
|
||||
offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
offs_n = tl.arange(0, BLOCK_N)
|
||||
offs_d = tl.arange(0, D)
|
||||
|
||||
Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
|
||||
K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None]
|
||||
V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :]
|
||||
OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
|
||||
LUT_ptr = LUT + lut_offset
|
||||
LSE_ptrs = LSE + lse_offset + offs_m
|
||||
|
||||
m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
|
||||
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
|
||||
o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
|
||||
|
||||
q = tl.load(Q_ptrs, mask=offs_m[:, None] < L)
|
||||
for block_idx in tl.range(topk):
|
||||
idx_n = tl.load(LUT_ptr + block_idx)
|
||||
n_mask = offs_n < L - idx_n * BLOCK_N
|
||||
|
||||
k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
|
||||
qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # = 1 / ln(2)
|
||||
if L - idx_n * BLOCK_N < BLOCK_N:
|
||||
qk = tl.where(n_mask[None, :], qk, float("-inf"))
|
||||
|
||||
v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
|
||||
local_m = tl.max(qk, 1)
|
||||
new_m = tl.maximum(m_i, local_m)
|
||||
qk = qk - new_m[:, None]
|
||||
|
||||
p = tl.math.exp2(qk)
|
||||
l_ij = tl.sum(p, 1)
|
||||
alpha = tl.math.exp2(m_i - new_m)
|
||||
o_s = o_s * alpha[:, None]
|
||||
o_s += tl.dot(p.to(v.dtype), v)
|
||||
|
||||
l_i = l_i * alpha + l_ij
|
||||
m_i = new_m
|
||||
|
||||
o_s = o_s / l_i[:, None]
|
||||
tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
|
||||
|
||||
m_i += tl.math.log2(l_i)
|
||||
tl.store(LSE_ptrs, m_i, mask=offs_m < L)
|
||||
|
||||
|
||||
def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
|
||||
arg_k = k - torch.mean(
|
||||
k, dim=-2, keepdim=True
|
||||
) # smooth-k technique in SageAttention
|
||||
pooled_qblocks = mean_pool(q, BLKQ)
|
||||
pooled_kblocks = mean_pool(arg_k, BLKK)
|
||||
pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
|
||||
|
||||
K = pooled_score.shape[-1]
|
||||
topk = min(K, int(topk_ratio * K))
|
||||
lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
|
||||
|
||||
sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
|
||||
sparse_map.scatter_(-1, lut, 1)
|
||||
return sparse_map, lut, topk
|
||||
|
||||
|
||||
def mean_pool(x, BLK):
|
||||
assert x.is_contiguous()
|
||||
|
||||
B, H, L, D = x.shape
|
||||
L_BLOCKS = (L + BLK - 1) // BLK
|
||||
x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
|
||||
|
||||
grid = (L_BLOCKS, B * H)
|
||||
compress_kernel[grid](x, x_mean, L, D, BLK)
|
||||
return x_mean
|
||||
|
||||
|
||||
@triton.jit
|
||||
def compress_kernel(
|
||||
X,
|
||||
XM,
|
||||
L: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
BLOCK_L: tl.constexpr,
|
||||
):
|
||||
idx_l = tl.program_id(0)
|
||||
idx_bh = tl.program_id(1)
|
||||
|
||||
offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
|
||||
offs_d = tl.arange(0, D)
|
||||
|
||||
x_offset = idx_bh * L * D
|
||||
xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
|
||||
x = tl.load(
|
||||
X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L
|
||||
)
|
||||
|
||||
nx = min(BLOCK_L, L - idx_l * BLOCK_L)
|
||||
x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
|
||||
tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))
|
||||
|
||||
|
||||
class _SeqAllToAll(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
@@ -301,7 +192,7 @@ class DistributedAttention(torch.nn.Module):
|
||||
self.stream = None
|
||||
|
||||
def forward(
|
||||
self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs
|
||||
self, query: Tensor, key: Tensor, value: Tensor, ctx_attn_metadata
|
||||
) -> Tensor:
|
||||
"""forward
|
||||
|
||||
@@ -309,22 +200,21 @@ class DistributedAttention(torch.nn.Module):
|
||||
query (Tensor): query input to the layer
|
||||
key (Tensor): key input to the layer
|
||||
value (Tensor): value input to the layer
|
||||
args: other args
|
||||
|
||||
Returns:
|
||||
* output (Tensor): context output
|
||||
"""
|
||||
if self.pg is None:
|
||||
return self.local_attn(query, key, value, *args, **kwargs)
|
||||
return self.local_attn(query, key, value, ctx_attn_metadata)
|
||||
pg_size = dist.get_world_size(self.pg)
|
||||
if pg_size < 2:
|
||||
return self.local_attn(query, key, value, *args, **kwargs)
|
||||
return self.local_attn(query, key, value, ctx_attn_metadata)
|
||||
|
||||
query_layer, key_layer, value_layer = _SeqAllToAllQKV.apply(
|
||||
self.pg, query, key, value, pg_size, self.stream, True
|
||||
)
|
||||
context_layer = self.local_attn(
|
||||
query_layer, key_layer, value_layer, *args, **kwargs
|
||||
query_layer, key_layer, value_layer, ctx_attn_metadata
|
||||
)
|
||||
|
||||
output = _SeqAllToAll.apply(self.pg, context_layer, False)
|
||||
@@ -336,8 +226,30 @@ class DistributedAttention(torch.nn.Module):
|
||||
|
||||
|
||||
class MinimalA2AAttnOp(DistributedAttention):
|
||||
def __init__(self, local_attn=None, *args, **kwargs):
|
||||
del args, kwargs
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
attention_type: str,
|
||||
topk: float,
|
||||
supported_attention_backends: set[AttentionBackendEnum] | None = None,
|
||||
):
|
||||
dtype = get_compute_dtype()
|
||||
attn_backend = get_attn_backend(
|
||||
head_size, dtype, supported_attention_backends=supported_attention_backends
|
||||
)
|
||||
# Maintained for compatibility purposes; can be removed when CI allows setting Attention_backend or when TurboWan supports FA.
|
||||
if attn_backend is not SparseLinearAttentionBackend:
|
||||
logger.warning(
|
||||
"TurboWan now only supports `sla_attn` and has been automatically set to `sla_attn`. Please set --attention-backend to `sla_attn`."
|
||||
)
|
||||
attn_backend = SparseLinearAttentionBackend
|
||||
impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls()
|
||||
local_attn = impl_cls(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
topk_ratio=topk,
|
||||
)
|
||||
super(MinimalA2AAttnOp, self).__init__(local_attn)
|
||||
|
||||
def set_context_parallel_group(self, process_group, ranks, stream):
|
||||
@@ -347,157 +259,7 @@ class MinimalA2AAttnOp(DistributedAttention):
|
||||
def forward(
|
||||
self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs
|
||||
) -> Tensor:
|
||||
results = super().forward(query, key, value, *args, **kwargs)
|
||||
forward_context: ForwardContext = get_forward_context()
|
||||
ctx_attn_metadata = forward_context.attn_metadata
|
||||
results = super().forward(query, key, value, ctx_attn_metadata)
|
||||
return rearrange(results, "b ... h l -> b ... (h l)")
|
||||
|
||||
|
||||
class SparseLinearAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
head_dim,
|
||||
topk,
|
||||
feature_map="softmax",
|
||||
BLKQ=64,
|
||||
BLKK=64,
|
||||
use_bf16=True,
|
||||
tie_feature_map_qk=True,
|
||||
):
|
||||
R"""
|
||||
Args:
|
||||
head_dim: dimension of each head.
|
||||
topk: ratio of keys selected for sparse attention, shared across all queries.
|
||||
feature_map: feature map for linear attention, one of ['hedgehog', 'elu', 'relu', 'softmax'].
|
||||
BLKQ: block size for query.
|
||||
BLKK: block size for key.
|
||||
use_bf16: whether to use bfloat16 (default) or float16 for computation. The conversion to bf16/fp16 is done inside the module.
|
||||
tie_feature_map_qk: whether to use the same feature map for query and key.
|
||||
"""
|
||||
super().__init__()
|
||||
self.dtype = torch.bfloat16 if use_bf16 else torch.float16
|
||||
self.topk = topk
|
||||
self.BLKQ = BLKQ
|
||||
self.BLKK = BLKK
|
||||
self.proj_l = nn.Linear(head_dim, head_dim, dtype=torch.float32)
|
||||
|
||||
if feature_map == "elu":
|
||||
self.feature_map_q = self._elu_feature_map
|
||||
self.feature_map_k = self.elu_feature_map
|
||||
elif feature_map == "relu":
|
||||
self.feature_map_q = nn.ReLU()
|
||||
self.feature_map_k = nn.ReLU()
|
||||
elif feature_map == "softmax":
|
||||
self.feature_map_q = self._softmax_feature_map
|
||||
self.feature_map_k = self._softmax_feature_map
|
||||
else:
|
||||
raise NotImplementedError(f"Not supported feature map {feature_map}.")
|
||||
|
||||
if tie_feature_map_qk:
|
||||
self.feature_map_k = self.feature_map_q
|
||||
|
||||
self.init_weights_()
|
||||
|
||||
def init_weights_(self):
|
||||
with torch.no_grad():
|
||||
nn.init.zeros_(self.proj_l.weight)
|
||||
nn.init.zeros_(self.proj_l.bias)
|
||||
|
||||
def forward(self, q, k, v, return_sparsity=False):
|
||||
R"""
|
||||
Args:
|
||||
q: queries of shape (B, H, L, D).
|
||||
k: keys of shape (B, H, L, D).
|
||||
v: values of shape (B, H, L, D).
|
||||
return_sparsity: whether to return the actual sparsity.
|
||||
"""
|
||||
dtype = q.dtype
|
||||
|
||||
q = q.transpose(1, 2).contiguous()
|
||||
k = k.transpose(1, 2).contiguous()
|
||||
v = v.transpose(1, 2).contiguous()
|
||||
|
||||
sparse_map, lut, real_topk = get_block_map(
|
||||
q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK
|
||||
)
|
||||
|
||||
q = q.to(self.dtype)
|
||||
k = k.to(self.dtype)
|
||||
v = v.to(self.dtype)
|
||||
o_s = _attention.apply(
|
||||
q, k, v, sparse_map, lut, real_topk, self.BLKQ, self.BLKK
|
||||
)
|
||||
|
||||
q = self.feature_map_q(q).contiguous().to(self.dtype) # c_q
|
||||
k = self.feature_map_k(k).contiguous().to(self.dtype) # c_k
|
||||
|
||||
o_l = self._torch_calc_linear(q, k, v)
|
||||
|
||||
with torch.amp.autocast("cuda", dtype=self.dtype):
|
||||
o_l = self.proj_l(o_l)
|
||||
o = (o_s + o_l).to(dtype).transpose(1, 2)
|
||||
|
||||
if return_sparsity:
|
||||
return o, real_topk / sparse_map.shape[-1]
|
||||
else:
|
||||
return o
|
||||
|
||||
def _torch_calc_linear(self, q, k, v):
|
||||
kv = torch.matmul(k.transpose(-1, -2), v)
|
||||
k_sum = torch.sum(k, dim=-2, keepdim=True)
|
||||
return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
|
||||
|
||||
def _calc_linear(self, q, k, v):
|
||||
kvsum = k.transpose(-1, -2) @ v
|
||||
ksum = torch.sum(k, dim=-2, keepdim=True)
|
||||
return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True))
|
||||
|
||||
def _softmax_feature_map(self, x):
|
||||
return F.softmax(x, dim=-1)
|
||||
|
||||
def _elu_feature_map(self, x):
|
||||
return F.elu(x) + 1
|
||||
|
||||
|
||||
class _attention(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None):
|
||||
assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
|
||||
assert k_block_id.is_contiguous() and lut.is_contiguous()
|
||||
|
||||
# We recommend the following two settings
|
||||
assert BLOCK_M == 64 or BLOCK_M == 128
|
||||
assert BLOCK_N == 64
|
||||
|
||||
B, H, L, D = q.shape
|
||||
if qk_scale is None:
|
||||
qk_scale = D**-0.5
|
||||
|
||||
M_BLOCKS = triton.cdiv(L, BLOCK_M)
|
||||
|
||||
o_s = torch.empty_like(v)
|
||||
lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32)
|
||||
|
||||
grid = (M_BLOCKS, B * H)
|
||||
_attn_fwd[grid](
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
qk_scale,
|
||||
topk,
|
||||
lut,
|
||||
lse,
|
||||
o_s,
|
||||
L,
|
||||
M_BLOCKS,
|
||||
D,
|
||||
BLOCK_M,
|
||||
BLOCK_N,
|
||||
num_warps=4 if q.shape[-1] == 64 else 8,
|
||||
num_stages=3,
|
||||
)
|
||||
|
||||
ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s)
|
||||
ctx.qk_scale = qk_scale
|
||||
ctx.topk = topk
|
||||
ctx.BLOCK_M = BLOCK_M
|
||||
ctx.BLOCK_N = BLOCK_N
|
||||
return o_s
|
||||
|
||||
@@ -18,7 +18,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention import (
|
||||
MinimalA2AAttnOp,
|
||||
SparseLinearAttention,
|
||||
UlyssesAttention_VSA,
|
||||
USPAttention,
|
||||
)
|
||||
@@ -311,9 +310,11 @@ class WanTransformerBlock(nn.Module):
|
||||
self.to_out = RowParallelLinear(dim, dim, bias=True, reduce_results=True)
|
||||
if attention_type == "sla":
|
||||
self.attn1 = MinimalA2AAttnOp(
|
||||
SparseLinearAttention(
|
||||
dim // num_heads, topk=sla_topk, BLKQ=128, BLKK=64
|
||||
)
|
||||
num_heads=divide(num_heads, get_tensor_model_parallel_world_size()),
|
||||
head_size=dim // num_heads,
|
||||
attention_type=attention_type,
|
||||
topk=sla_topk,
|
||||
supported_attention_backends={AttentionBackendEnum.SLA_ATTN},
|
||||
)
|
||||
else:
|
||||
self.attn1 = USPAttention(
|
||||
|
||||
@@ -245,6 +245,9 @@ class CudaPlatformBase(Platform):
|
||||
elif selected_backend == AttentionBackendEnum.TORCH_SDPA:
|
||||
logger.info("Using Torch SDPA backend")
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
|
||||
elif selected_backend == AttentionBackendEnum.SLA_ATTN:
|
||||
logger.info("Using Sparse Linear Attention backend")
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend"
|
||||
elif selected_backend in [
|
||||
AttentionBackendEnum.FA,
|
||||
]:
|
||||
|
||||
@@ -33,6 +33,7 @@ class AttentionBackendEnum(enum.Enum):
|
||||
VIDEO_SPARSE_ATTN = enum.auto()
|
||||
VMOBA_ATTN = enum.auto()
|
||||
AITER = enum.auto()
|
||||
SLA_ATTN = enum.auto()
|
||||
NO_ATTENTION = enum.auto()
|
||||
|
||||
def __str__(self):
|
||||
|
||||
Reference in New Issue
Block a user