[Diffusion][NPU]Add attention backends for diffusion models for Ascend NPU (#23482)
Co-authored-by: Napkin-AI <arseniy.mironov.dev@gmail.com>
This commit is contained in:
Arseniy Mironov
co-authored by
Napkin-AI
Napkin-AI
parent
58b5fe3e29
commit
45a85efc3a
@@ -27,6 +27,9 @@ class AdapterArchConfig(ArchConfig):
|
||||
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
|
||||
AttentionBackendEnum.VMOBA_ATTN,
|
||||
AttentionBackendEnum.SAGE_ATTN_3,
|
||||
AttentionBackendEnum.LASER_ATTN,
|
||||
AttentionBackendEnum.BLOCK_SPARSE_ATTN,
|
||||
AttentionBackendEnum.RAIN_FUSION_ATTN,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -35,6 +35,9 @@ class DiTArchConfig(ArchConfig):
|
||||
AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
|
||||
AttentionBackendEnum.VMOBA_ATTN,
|
||||
AttentionBackendEnum.SAGE_ATTN_3,
|
||||
AttentionBackendEnum.LASER_ATTN,
|
||||
AttentionBackendEnum.BLOCK_SPARSE_ATTN,
|
||||
AttentionBackendEnum.RAIN_FUSION_ATTN,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import attentions # noqa: F401
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
|
||||
LaserAttentionBackend,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
BSA_BLOCK_SIZE = 128
|
||||
|
||||
|
||||
class BlockSparseAttentionBackend(AttentionBackend):
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [32, 64, 96, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_enum() -> AttentionBackendEnum:
|
||||
return AttentionBackendEnum.BLOCK_SPARSE_ATTN
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["BlockSparseAttentionImpl"]:
|
||||
return BlockSparseAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["BlockSparseAttentionMetadata"]:
|
||||
return BlockSparseAttentionMetadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["BlockSparseAttentionMetadataBuilder"]:
|
||||
return BlockSparseAttentionMetadataBuilder
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockSparseAttentionMetadata(AttentionMetadata):
|
||||
current_timestep: int
|
||||
skip_first_steps: int
|
||||
sparsity: float
|
||||
block_frame_stride: int
|
||||
|
||||
|
||||
class BlockSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def build(
|
||||
self,
|
||||
current_timestep: int,
|
||||
skip_first_steps: int,
|
||||
sparsity: float,
|
||||
raw_latent_shape: list[int],
|
||||
patch_size: tuple[int, int, int],
|
||||
**kwargs: dict[str, Any],
|
||||
) -> BlockSparseAttentionMetadata:
|
||||
"""
|
||||
Builds BlockSparseAttention metadata.
|
||||
|
||||
Args:
|
||||
current_timestep: The current diffusion timestep.
|
||||
skip_first_steps: Number of initial timesteps to skip before applying
|
||||
sparsity. Must be non‑negative.
|
||||
sparsity: Fraction of tokens to drop (block‑wise) in the block sparse
|
||||
attention mechanism. Must be in the range [0.0, 1.0).
|
||||
raw_latent_shape: Shape of the latent tensor before patching.
|
||||
patch_size: Patch size as (T, height, width). Only the height
|
||||
and width components are used to divide the latent dimensions.
|
||||
**kwargs: Additional keyword arguments (ignored, but accepted for
|
||||
compatibility with base class or calling conventions).
|
||||
|
||||
Returns:
|
||||
BlockSparseAttentionMetadata
|
||||
Note:
|
||||
The `block_frame_stride` is needed to set the first blocks to be non‑sparse.
|
||||
"""
|
||||
if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
|
||||
raise ValueError(
|
||||
(
|
||||
"Invalid attention metadata values."
|
||||
f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
|
||||
f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
|
||||
)
|
||||
)
|
||||
|
||||
if sparsity == 0.0:
|
||||
logger.warning(
|
||||
(
|
||||
"Sparsity is set to 0.0, which means no tokens will be dropped."
|
||||
"For better performance use Laser Attention or increase sparsity."
|
||||
)
|
||||
)
|
||||
|
||||
if len(raw_latent_shape) >= 5:
|
||||
latent_height, latent_width = raw_latent_shape[3:5]
|
||||
else:
|
||||
latent_height, latent_width = raw_latent_shape[-2:]
|
||||
|
||||
latent_height //= patch_size[1]
|
||||
latent_width //= patch_size[2]
|
||||
|
||||
frame_stride = latent_height * latent_width
|
||||
block_frame_stride = (frame_stride + BSA_BLOCK_SIZE - 1) // BSA_BLOCK_SIZE
|
||||
|
||||
return BlockSparseAttentionMetadata(
|
||||
current_timestep=current_timestep,
|
||||
skip_first_steps=skip_first_steps,
|
||||
sparsity=sparsity,
|
||||
block_frame_stride=block_frame_stride,
|
||||
)
|
||||
|
||||
|
||||
class BlockSparseAttentionImpl(AttentionImpl):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
self.num_heads = num_heads
|
||||
self.num_kv_heads = num_kv_heads or num_heads
|
||||
self.block_size = BSA_BLOCK_SIZE
|
||||
self.stride = 8
|
||||
self.default_tokens = 214748647
|
||||
|
||||
self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
|
||||
num_heads,
|
||||
head_size,
|
||||
causal,
|
||||
softmax_scale,
|
||||
num_kv_heads,
|
||||
prefix,
|
||||
**extra_impl_args,
|
||||
)
|
||||
|
||||
def _get_estimate_mask(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
sparsity: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return torch.ops.attentions.sparse_block_estimate(
|
||||
query=query,
|
||||
key=key,
|
||||
actual_seq_lengths=None,
|
||||
actual_seq_lengths_kv=None,
|
||||
input_layout="BNSD",
|
||||
stride=self.stride,
|
||||
sparse_size=self.block_size,
|
||||
num_heads=query.shape[1],
|
||||
num_key_value_heads=key.shape[1],
|
||||
scale_value=self.softmax_scale / self.stride,
|
||||
threshold=1.0,
|
||||
causal=self.causal,
|
||||
keep_sink=True,
|
||||
keep_recent=True,
|
||||
row_sparse=1.0 - sparsity,
|
||||
)
|
||||
|
||||
def _block_sparse_attention(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
smask: torch.Tensor,
|
||||
sct: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return torch.ops.attentions.block_sparse_attention(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
sparse_mask=smask,
|
||||
sparse_count_table=sct,
|
||||
input_layout="BNSD",
|
||||
sparse_size=self.block_size,
|
||||
num_heads=query.shape[1],
|
||||
num_key_value_heads=key.shape[1],
|
||||
scale_value=self.softmax_scale,
|
||||
causal=self.causal,
|
||||
inner_precise=1,
|
||||
pre_tokens=self.default_tokens,
|
||||
next_tokens=self.default_tokens,
|
||||
actual_seq_lengths=None,
|
||||
actual_seq_lengths_kv=None,
|
||||
)
|
||||
|
||||
def _get_smask(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
block_frame_stride: int,
|
||||
sparsity: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
smask, sct = self._get_estimate_mask(
|
||||
query,
|
||||
key,
|
||||
sparsity,
|
||||
)
|
||||
|
||||
seq_len = smask.shape[2]
|
||||
|
||||
# Set the first blocks to be non-sparse to ensure the quality of the first few steps
|
||||
smask[:, :, :block_frame_stride, :seq_len] = 1
|
||||
smask[:, :, :seq_len, :block_frame_stride] = 1
|
||||
smask = smask.to(torch.int8)
|
||||
sct = smask.sum(dim=-1, dtype=torch.int32)
|
||||
return smask, sct
|
||||
|
||||
def _adaptive_block_sparse_attention(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
block_frame_stride: int,
|
||||
sparsity: float,
|
||||
) -> torch.Tensor:
|
||||
# TODO Currently implementation for BSND input layout has quality issues
|
||||
# When the implementation is improved, transposes can be removed
|
||||
q = query.permute(0, 2, 1, 3).contiguous()
|
||||
k = key.permute(0, 2, 1, 3).contiguous()
|
||||
v = value.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
smask, sct = self._get_smask(
|
||||
q,
|
||||
k,
|
||||
block_frame_stride,
|
||||
sparsity,
|
||||
)
|
||||
output = self._block_sparse_attention(q, k, v, smask, sct)
|
||||
output = output.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
return output
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
|
||||
output = self.laser_attn_impl.forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_metadata,
|
||||
)
|
||||
else:
|
||||
output = self._adaptive_block_sparse_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_metadata.block_frame_stride,
|
||||
attn_metadata.sparsity,
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,191 @@
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPABackend
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# Import to use torch.ops.attentions, install package with sgl_kernel_npu
|
||||
try:
|
||||
import attentions # noqa: F401
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
(
|
||||
"The required 'attentions' package is not installed."
|
||||
"The package can be installed with sgl_kernel_npu"
|
||||
)
|
||||
) from e
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class LaserAttentionBackend(AttentionBackend):
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [32, 64, 96, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_enum() -> AttentionBackendEnum:
|
||||
return AttentionBackendEnum.LASER_ATTN
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["LaserAttentionImpl"]:
|
||||
return LaserAttentionImpl
|
||||
|
||||
|
||||
class LaserAttentionImpl(AttentionImpl):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.softmax_scale = softmax_scale
|
||||
|
||||
# After preprocess input layout should be BNSD.
|
||||
self.seqlen_base = 256
|
||||
self.seqlen_index = 2
|
||||
self.dim_index = 3
|
||||
self.dim_base = 128
|
||||
self.max_token = 2**31 - 1
|
||||
self.seq_len_pad_base = 256
|
||||
|
||||
# the laser attention operator has issues with small seq_len
|
||||
self.min_seqlen = 2048
|
||||
self.sdpa_impl = SDPABackend.get_impl_cls()(
|
||||
num_heads,
|
||||
head_size,
|
||||
causal,
|
||||
softmax_scale,
|
||||
num_kv_heads,
|
||||
prefix,
|
||||
**extra_impl_args,
|
||||
)
|
||||
|
||||
def _pad(self, input_tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Pad the input tensor along the sequence length and head dimension.
|
||||
to multiples of base values. self.seqlen_index and self.dim_index should be positive integers.
|
||||
"""
|
||||
|
||||
seq_len = input_tensor.size(self.seqlen_index)
|
||||
head_dim = input_tensor.size(self.dim_index)
|
||||
|
||||
pad_seq = 0
|
||||
if seq_len % self.seqlen_base != 0:
|
||||
pad_seq = ((seq_len // self.seqlen_base) + 1) * self.seqlen_base - seq_len
|
||||
|
||||
pad_dim = 0
|
||||
if head_dim % self.dim_base != 0:
|
||||
pad_dim = ((head_dim // self.dim_base) + 1) * self.dim_base - head_dim
|
||||
|
||||
if pad_seq == 0 and pad_dim == 0:
|
||||
return input_tensor
|
||||
|
||||
pad_list = [0] * (2 * input_tensor.ndim)
|
||||
|
||||
pad_list[len(pad_list) - 2 * self.seqlen_index - 1] = pad_seq
|
||||
pad_list[len(pad_list) - 2 * self.dim_index - 1] = pad_dim
|
||||
|
||||
return torch.nn.functional.pad(input_tensor, pad_list)
|
||||
|
||||
def _la_preprocess_input(
|
||||
self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
# Currently BSND input layout is not supported
|
||||
q = query.transpose(1, 2)
|
||||
k = key.transpose(1, 2)
|
||||
v = value.transpose(1, 2)
|
||||
|
||||
if q.dtype != torch.float16:
|
||||
q = q.to(torch.float16)
|
||||
k = k.to(torch.float16)
|
||||
v = v.to(torch.float16)
|
||||
|
||||
q = self._pad(q)
|
||||
k = self._pad(k)
|
||||
v = self._pad(v)
|
||||
|
||||
return q, k, v
|
||||
|
||||
def _la_postprocess_output(
|
||||
self,
|
||||
attention_out: torch.Tensor,
|
||||
dtype: torch.dtype,
|
||||
qseqlen: int,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
if dtype != attention_out.dtype:
|
||||
attention_out = attention_out.to(dtype)
|
||||
|
||||
attention_out = attention_out[:, :, :qseqlen, :head_dim]
|
||||
attention_out = attention_out.transpose(1, 2).contiguous()
|
||||
return attention_out
|
||||
|
||||
def _laser_attention(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
head_num: int,
|
||||
pre_tokens: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return torch.ops.attentions.la(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
atten_mask=None,
|
||||
alibi_mask=None,
|
||||
drop_mask=None,
|
||||
scale_value=self.softmax_scale,
|
||||
head_num=head_num,
|
||||
input_layout="BNSD",
|
||||
keep_prob=1.0,
|
||||
pre_tokens=pre_tokens,
|
||||
next_tokens=1,
|
||||
is_highPrecision=True,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
q_seqlen, head_dim = query.shape[1], query.shape[3]
|
||||
kv_seqlen = key.shape[1]
|
||||
|
||||
if q_seqlen < self.min_seqlen or kv_seqlen != q_seqlen:
|
||||
output = self.sdpa_impl.forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_metadata,
|
||||
)
|
||||
else:
|
||||
pre_tokens = self.max_token
|
||||
if kv_seqlen % self.seq_len_pad_base != 0:
|
||||
pre_tokens = (
|
||||
kv_seqlen // self.seq_len_pad_base + 1
|
||||
) * self.seq_len_pad_base - kv_seqlen
|
||||
|
||||
q, k, v = self._la_preprocess_input(query, key, value)
|
||||
_, la_output = self._laser_attention(q, k, v, q.shape[1], pre_tokens)
|
||||
output = self._la_postprocess_output(
|
||||
la_output, query.dtype, q_seqlen, head_dim
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,414 @@
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional
|
||||
|
||||
import attentions # noqa: F401
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
|
||||
LaserAttentionBackend,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class RainFusionAttentionBackend(AttentionBackend):
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [32, 64, 96, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_enum() -> AttentionBackendEnum:
|
||||
return AttentionBackendEnum.RAIN_FUSION_ATTN
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["RainFusionAttentionImpl"]:
|
||||
return RainFusionAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["RainFusionAttentionMetadata"]:
|
||||
return RainFusionAttentionMetadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["RainFusionAttentionMetadataBuilder"]:
|
||||
return RainFusionAttentionMetadataBuilder
|
||||
|
||||
|
||||
@dataclass
|
||||
class RainFusionAttentionMetadata(AttentionMetadata):
|
||||
current_timestep: int
|
||||
skip_first_steps: int
|
||||
sparsity: float
|
||||
latent_shape: list[int]
|
||||
|
||||
|
||||
class RainFusionAttentionMetadataBuilder(AttentionMetadataBuilder):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def build(
|
||||
self,
|
||||
current_timestep: int,
|
||||
skip_first_steps: int,
|
||||
sparsity: float,
|
||||
raw_latent_shape: list[int],
|
||||
patch_size: tuple[int, int, int],
|
||||
**kwargs: dict[str, Any],
|
||||
) -> RainFusionAttentionMetadata:
|
||||
if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
|
||||
raise ValueError(
|
||||
(
|
||||
"Invalid attention metadata values."
|
||||
f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
|
||||
f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
|
||||
)
|
||||
)
|
||||
|
||||
if sparsity == 0.0:
|
||||
logger.warning(
|
||||
(
|
||||
"Sparsity is set to 0.0, which means no tokens will be dropped."
|
||||
"For better performance use Laser Attention or increase sparsity."
|
||||
)
|
||||
)
|
||||
|
||||
latent_shape = raw_latent_shape[-3:]
|
||||
latent_shape = [latent_shape[i] // patch_size[i] for i in range(3)]
|
||||
|
||||
return RainFusionAttentionMetadata(
|
||||
current_timestep=current_timestep,
|
||||
skip_first_steps=skip_first_steps,
|
||||
sparsity=sparsity,
|
||||
latent_shape=latent_shape,
|
||||
)
|
||||
|
||||
|
||||
class RainFusionAttentionImpl(AttentionImpl):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
self.block_size = 128
|
||||
self.inner_precise = 0
|
||||
|
||||
self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
|
||||
num_heads,
|
||||
head_size,
|
||||
causal,
|
||||
softmax_scale,
|
||||
num_kv_heads,
|
||||
prefix,
|
||||
**extra_impl_args,
|
||||
)
|
||||
|
||||
def _avgpool(
|
||||
self, input_tensor: torch.Tensor, pool_size: int = 128
|
||||
) -> torch.Tensor:
|
||||
batch, seqlen, heads, dim = input_tensor.shape
|
||||
x = input_tensor.permute(0, 2, 3, 1).reshape(batch * heads, dim, seqlen)
|
||||
|
||||
pooled = torch.nn.functional.avg_pool1d(
|
||||
x, kernel_size=pool_size, stride=pool_size, ceil_mode=True
|
||||
)
|
||||
out = pooled.reshape(batch, heads, dim, -1).permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
return out
|
||||
|
||||
def _get_mask_index(self, mask: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, num_heads, seq_len, _ = mask.shape
|
||||
|
||||
mask_reshaped = mask.reshape(-1, seq_len)
|
||||
row_indices = torch.arange(
|
||||
seq_len, device=mask.device, dtype=torch.float32
|
||||
).unsqueeze(0)
|
||||
|
||||
sorted_vals = torch.where(mask_reshaped, row_indices, seq_len)
|
||||
sorted_vals, _ = torch.sort(sorted_vals, dim=-1)
|
||||
valid_count = mask_reshaped.sum(dim=-1, keepdim=True)
|
||||
keep_mask = row_indices < valid_count
|
||||
result = torch.where(keep_mask, sorted_vals, -1)
|
||||
|
||||
pos_matrix = result.reshape(batch_size, num_heads, seq_len, seq_len).to(
|
||||
torch.int64
|
||||
)
|
||||
return pos_matrix
|
||||
|
||||
def _get_blockwise_mask(
|
||||
self,
|
||||
qkv_pool: torch.Tensor,
|
||||
sparsity: float,
|
||||
scale: float,
|
||||
pool_size: int,
|
||||
latent_shape: tuple,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
first_frame_len = latent_shape[1] * latent_shape[2]
|
||||
|
||||
query_pool, key_pool, value_pool = torch.chunk(qkv_pool, 3, dim=0)
|
||||
attn_scores = (
|
||||
query_pool.permute(0, 2, 1, 3) @ key_pool.permute(0, 2, 3, 1) * scale
|
||||
)
|
||||
|
||||
keep_len = math.ceil(attn_scores.shape[-1] * (1 - sparsity))
|
||||
|
||||
topk_values, _ = torch.topk(attn_scores, k=keep_len, dim=-1)
|
||||
mask = attn_scores >= topk_values[..., -1:]
|
||||
|
||||
firstframe_block_num = (first_frame_len + pool_size - 1) // pool_size
|
||||
if firstframe_block_num > 0:
|
||||
mask[:, :, :firstframe_block_num, :] = True
|
||||
mask[:, :, :, :firstframe_block_num] = True
|
||||
|
||||
select_idx = self._get_mask_index(mask)
|
||||
select_idx = select_idx[0].transpose(0, 1)
|
||||
select_num_idx = mask[0].transpose(0, 1).sum(dim=-1)
|
||||
return select_idx, select_num_idx
|
||||
|
||||
def _rearrange_with_remaining(
|
||||
self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d
|
||||
or
|
||||
b n (f hn hb wn wb) d -> b n (f hn wn hb wb) d
|
||||
"""
|
||||
tq, hq, wq = latent_shape
|
||||
first_frame_len, frame_num = hq * wq, tq
|
||||
|
||||
b, s, n, d = tensor.shape
|
||||
|
||||
if (hq % 8 != 0) or (wq % 8 != 0):
|
||||
tensor_first = tensor[:, :first_frame_len, :, :]
|
||||
tensor = tensor[:, first_frame_len:, :, :]
|
||||
tensor_hwt = rearrange(
|
||||
tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
|
||||
)
|
||||
if hq % 8 != 0:
|
||||
tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
|
||||
tensor_h_r = tensor_h_r.reshape(b, frame_num - 1, -1, n, d)
|
||||
if wq % 8 != 0:
|
||||
tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
|
||||
tensor_w_r = tensor_w_r.reshape(b, frame_num - 1, -1, n, d)
|
||||
tensor_hwt = rearrange(
|
||||
tensor_hwt,
|
||||
"b f (hn hb) (wn wb) n d -> b f (hn wn hb wb) n d",
|
||||
f=frame_num - 1,
|
||||
hb=8,
|
||||
wb=8,
|
||||
hn=hq // 8,
|
||||
wn=wq // 8,
|
||||
)
|
||||
if hq % 8 != 0:
|
||||
tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
|
||||
if wq % 8 != 0:
|
||||
tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=2)
|
||||
tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
|
||||
tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
|
||||
else:
|
||||
tensor_hwt = rearrange(
|
||||
tensor,
|
||||
"b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d",
|
||||
f=frame_num,
|
||||
hb=8,
|
||||
wb=8,
|
||||
hn=hq // 8,
|
||||
wn=wq // 8,
|
||||
)
|
||||
|
||||
return tensor_hwt
|
||||
|
||||
def _inv_rearrange_with_remaining(
|
||||
self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
|
||||
) -> torch.Tensor:
|
||||
tq, hq, wq = latent_shape
|
||||
first_frame_len, frame_num = hq * wq, tq
|
||||
|
||||
b, s, n, d = tensor.shape
|
||||
|
||||
if (hq % 8 != 0) or (wq % 8 != 0):
|
||||
tensor_first = tensor[:, :first_frame_len, :, :]
|
||||
tensor = tensor[:, first_frame_len:, :, :]
|
||||
tensor_hwt = rearrange(
|
||||
tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
|
||||
)
|
||||
if hq % 8 != 0:
|
||||
tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
|
||||
if wq % 8 != 0:
|
||||
tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
|
||||
tensor_hwt = tensor_hwt.reshape(b, frame_num - 1, -1, n, d)
|
||||
tensor_hwt = rearrange(
|
||||
tensor_hwt,
|
||||
"b f (hn wn hb wb) n d -> b f (hn hb) (wn wb) n d",
|
||||
f=frame_num - 1,
|
||||
hb=8,
|
||||
wb=8,
|
||||
hn=hq // 8,
|
||||
wn=wq // 8,
|
||||
)
|
||||
if wq % 8 != 0:
|
||||
tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=3)
|
||||
if hq % 8 != 0:
|
||||
tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
|
||||
tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
|
||||
tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
|
||||
else:
|
||||
tensor_hwt = rearrange(
|
||||
tensor,
|
||||
"b (f hn wn hb wb) n h -> b (f hn hb wn wb) n h",
|
||||
f=frame_num,
|
||||
hb=8,
|
||||
wb=8,
|
||||
hn=hq // 8,
|
||||
wn=wq // 8,
|
||||
)
|
||||
|
||||
return tensor_hwt
|
||||
|
||||
def _do_tensor_rearrange_pooling(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
pool_size: int,
|
||||
latent_shape: tuple[int, int, int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Tensor block rearrangement + pooling operation
|
||||
"""
|
||||
tensor = torch.cat((query, key, value), dim=0)
|
||||
|
||||
tensor = self._rearrange_with_remaining(tensor, latent_shape)
|
||||
tensor_pool = self._avgpool(tensor, pool_size)
|
||||
|
||||
query_, key_, value_ = torch.chunk(tensor, 3, dim=0)
|
||||
return query_, key_, value_, tensor_pool
|
||||
|
||||
def _rain_fusion_attention(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
select_idx: torch.Tensor,
|
||||
select_num_idx: torch.Tensor,
|
||||
blockshape: List[int],
|
||||
scale: float = 1.0,
|
||||
head_num: int = 1,
|
||||
input_layout: str = "TND",
|
||||
actual_seq_lengths=Optional[torch.Tensor],
|
||||
actual_seq_lengths_kv=Optional[torch.Tensor],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return torch.ops.attentions.rainfusionattention(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
select_idx=select_idx,
|
||||
select_num_idx=select_num_idx,
|
||||
blockshape=blockshape,
|
||||
attn_mask=None,
|
||||
actual_seq_qlen=actual_seq_lengths,
|
||||
actual_seq_kvlen=actual_seq_lengths_kv,
|
||||
block_table=None,
|
||||
q_input_layout=input_layout,
|
||||
kv_input_layout=input_layout,
|
||||
head_num=head_num,
|
||||
mask_type=0,
|
||||
scale=scale,
|
||||
inner_precise=self.inner_precise,
|
||||
block_size=0,
|
||||
)
|
||||
|
||||
def _rain_fusion_sparse_attention(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
latent_shape: tuple[int, int, int],
|
||||
sparsity: float,
|
||||
):
|
||||
q, k, v, qkv_pool = self._do_tensor_rearrange_pooling(
|
||||
query, key, value, self.block_size, latent_shape
|
||||
)
|
||||
|
||||
select_idx, select_num_idx = self._get_blockwise_mask(
|
||||
qkv_pool,
|
||||
sparsity,
|
||||
self.softmax_scale,
|
||||
self.block_size,
|
||||
latent_shape,
|
||||
)
|
||||
|
||||
batch_size, seqlen_q, head_num, head_dim = q.shape
|
||||
seqlen_kv = k.shape[1]
|
||||
|
||||
layout = "TND"
|
||||
q = q.reshape(-1, head_num, head_dim)
|
||||
k = k.reshape(-1, head_num, head_dim)
|
||||
v = v.reshape(-1, head_num, head_dim)
|
||||
|
||||
actual_seq_lengths = [seqlen_q] * batch_size
|
||||
actual_seq_lengths_kv = [seqlen_kv] * batch_size
|
||||
|
||||
out, _ = self._rain_fusion_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
scale=self.softmax_scale,
|
||||
head_num=head_num,
|
||||
input_layout=layout,
|
||||
select_idx=select_idx,
|
||||
select_num_idx=select_num_idx,
|
||||
blockshape=[self.block_size, self.block_size],
|
||||
actual_seq_lengths=actual_seq_lengths,
|
||||
actual_seq_lengths_kv=actual_seq_lengths_kv,
|
||||
)
|
||||
|
||||
out = out.reshape(batch_size, seqlen_q, head_num, head_dim)
|
||||
out = self._inv_rearrange_with_remaining(out, latent_shape)
|
||||
return out
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
|
||||
output = self.laser_attn_impl.forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_metadata,
|
||||
)
|
||||
else:
|
||||
output = self._rain_fusion_sparse_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_metadata.latent_shape,
|
||||
attn_metadata.sparsity,
|
||||
)
|
||||
return output
|
||||
@@ -1582,6 +1582,35 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
attn_metadata = self.attn_metadata_builder.build(
|
||||
raw_latent_shape=batch.raw_latent_shape
|
||||
)
|
||||
elif self.attn_backend.get_enum() in [
|
||||
AttentionBackendEnum.BLOCK_SPARSE_ATTN,
|
||||
AttentionBackendEnum.RAIN_FUSION_ATTN,
|
||||
]:
|
||||
sparse_config = server_args.attention_backend_config
|
||||
|
||||
current_timestep = i
|
||||
skip_first_steps = sparse_config.get("skip_first_steps", 10)
|
||||
sparsity = sparse_config.get("sparsity", 0.2)
|
||||
|
||||
raw_latent_shape = batch.raw_latent_shape
|
||||
patch_size = server_args.pipeline_config.dit_config.patch_size
|
||||
|
||||
if isinstance(patch_size, int):
|
||||
patch_size_t = getattr(
|
||||
server_args.pipeline_config.dit_config, "patch_size_t", None
|
||||
)
|
||||
if patch_size_t is not None:
|
||||
patch_size = (patch_size_t, patch_size, patch_size)
|
||||
else:
|
||||
patch_size = (patch_size, patch_size, patch_size)
|
||||
|
||||
attn_metadata = self.attn_metadata_builder.build(
|
||||
current_timestep=current_timestep,
|
||||
skip_first_steps=skip_first_steps,
|
||||
sparsity=sparsity,
|
||||
raw_latent_shape=raw_latent_shape,
|
||||
patch_size=patch_size,
|
||||
)
|
||||
else:
|
||||
# attn_metadata can be None for SDPA attention backend
|
||||
return None
|
||||
|
||||
@@ -38,6 +38,9 @@ class AttentionBackendEnum(enum.Enum):
|
||||
AITER_SAGE = enum.auto()
|
||||
SLA_ATTN = enum.auto()
|
||||
SAGE_SLA_ATTN = enum.auto()
|
||||
LASER_ATTN = enum.auto()
|
||||
BLOCK_SPARSE_ATTN = enum.auto()
|
||||
RAIN_FUSION_ATTN = enum.auto()
|
||||
NO_ATTENTION = enum.auto()
|
||||
|
||||
def __str__(self):
|
||||
@@ -52,6 +55,9 @@ class AttentionBackendEnum(enum.Enum):
|
||||
AttentionBackendEnum.VMOBA_ATTN,
|
||||
AttentionBackendEnum.SLA_ATTN,
|
||||
AttentionBackendEnum.SAGE_SLA_ATTN,
|
||||
AttentionBackendEnum.LASER_ATTN,
|
||||
AttentionBackendEnum.BLOCK_SPARSE_ATTN,
|
||||
AttentionBackendEnum.RAIN_FUSION_ATTN,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -120,6 +120,57 @@ class NPUPlatformBase(Platform):
|
||||
logger.info("Using Ascend Flash Attention backend.")
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.ascend_fa.AscendFABackend"
|
||||
|
||||
elif selected_backend == AttentionBackendEnum.LASER_ATTN:
|
||||
try:
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import ( # noqa: F401
|
||||
LaserAttentionBackend,
|
||||
)
|
||||
|
||||
logger.info("Using Laser Attention backend")
|
||||
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn.LaserAttentionBackend"
|
||||
except ImportError as e:
|
||||
logger.error(f"Failed to import Laser Attention backend: {e}")
|
||||
raise ImportError(
|
||||
"Laser Attention backend is not installed. "
|
||||
"It requires the `attentions` module which can be installed along with sgl_kernel_npu. "
|
||||
"Manual installation from source is required. See https://github.com/sgl-project/sgl-kernel-npu."
|
||||
) from e
|
||||
|
||||
elif selected_backend == AttentionBackendEnum.BLOCK_SPARSE_ATTN:
|
||||
try:
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.block_sparse_attn import ( # noqa: F401
|
||||
BlockSparseAttentionBackend,
|
||||
)
|
||||
|
||||
logger.info("Using Block Sparse Attention backend")
|
||||
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.block_sparse_attn.BlockSparseAttentionBackend"
|
||||
except ImportError as e:
|
||||
logger.error(f"Failed to import Block Sparse Attention backend: {e}")
|
||||
raise ImportError(
|
||||
"Block Sparse Attention backend is not installed. "
|
||||
"It requires the `attentions` module which can be installed along with sgl_kernel_npu. "
|
||||
"Manual installation from source is required. See https://github.com/sgl-project/sgl-kernel-npu."
|
||||
) from e
|
||||
|
||||
elif selected_backend == AttentionBackendEnum.RAIN_FUSION_ATTN:
|
||||
try:
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.rain_fusion_attn import ( # noqa: F401
|
||||
RainFusionAttentionBackend,
|
||||
)
|
||||
|
||||
logger.info("Using Rain Fusion Attention backend")
|
||||
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.rain_fusion_attn.RainFusionAttentionBackend"
|
||||
except ImportError as e:
|
||||
logger.error(f"Failed to import Rain Fusion Attention backend: {e}")
|
||||
raise ImportError(
|
||||
"Rain Fusion Attention backend is not installed. "
|
||||
"It requires the `attentions` module which can be installed along with sgl_kernel_npu. "
|
||||
"Manual installation from source is required. See https://github.com/sgl-project/sgl-kernel-npu."
|
||||
) from e
|
||||
|
||||
logger.info("Using Torch SDPA backend.")
|
||||
return (
|
||||
"sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
|
||||
|
||||
Reference in New Issue
Block a user