minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] Clean up VSA attention hot path (#25514)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-05-24 16:46:03 +08:00
committed by GitHub
parent 4c2b32bfbf
commit 0b65588c18
5 changed files with 88 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py
+4 -4
View File
@@ -323,8 +323,8 @@ class SparseLinearAttentionImpl(AttentionImpl, nn.Module):
)
# 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
query = self.feature_map_q(query).to(self.dtype) # c_q
key = self.feature_map_k(key).to(self.dtype) # c_k
# Linear attention computation
o_l = self._calc_linear_attention_with_torch(query, key, value)
@@ -681,8 +681,8 @@ class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module):
########## SPARGE END ##########
# Linear attention with feature maps
q_linear = self.feature_map_q(q).contiguous().to(self.dtype)
k_linear = self.feature_map_k(k).contiguous().to(self.dtype)
q_linear = self.feature_map_q(q).to(self.dtype)
k_linear = self.feature_map_k(k).to(self.dtype)
o_l = self._calc_linear_attention_with_torch(q_linear, k_linear, v)
# Project linear attention output and combine
python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py
+24 -26
View File
@@ -159,6 +159,8 @@ class VideoSparseAttentionMetadata(AttentionMetadata):
reverse_tile_partition_indices: torch.LongTensor
variable_block_sizes: torch.LongTensor
non_pad_index: torch.LongTensor
untile_combined_index: torch.LongTensor
tile_buf: torch.Tensor | None = None
# adaption for FastWan2.1-T2V-1.3B-Diffusers
# Sequence lengths for the forward batch
@@ -211,6 +213,7 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
non_pad_index = get_non_pad_index(
variable_block_sizes, math.prod(VSA_TILE_SIZE)
)
untile_combined_index = non_pad_index[reverse_tile_partition_indices]
return VideoSparseAttentionMetadata(
current_timestep=current_timestep,
@@ -222,6 +225,7 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
reverse_tile_partition_indices=reverse_tile_partition_indices,
variable_block_sizes=variable_block_sizes,
non_pad_index=non_pad_index,
untile_combined_index=untile_combined_index,
)
@@ -244,58 +248,52 @@ class VideoSparseAttentionImpl(AttentionImpl):
def tile(
self,
x: torch.Tensor,
num_tiles: list[int],
tile_partition_indices: torch.LongTensor,
non_pad_index: torch.LongTensor,
attn_metadata: VideoSparseAttentionMetadata,
) -> torch.Tensor:
num_tiles = attn_metadata.num_tiles
t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1]
w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
x_padded = torch.zeros(
(
target_shape = (
x.shape[0],
t_padded_size * h_padded_size * w_padded_size,
x.shape[-2],
x.shape[-1],
),
device=x.device,
dtype=x.dtype,
)
x_padded[:, non_pad_index] = x[:, tile_partition_indices]
return x_padded
buf = attn_metadata.tile_buf
if (
buf is None
or buf.shape != target_shape
or buf.dtype != x.dtype
or buf.device != x.device
):
buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
attn_metadata.tile_buf = buf
buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
return buf
def untile(
self,
x: torch.Tensor,
reverse_tile_partition_indices: torch.LongTensor,
non_pad_index: torch.LongTensor,
untile_combined_index: torch.LongTensor,
) -> torch.Tensor:
x = x[:, non_pad_index][:, reverse_tile_partition_indices]
return x
return x[:, untile_combined_index]
def preprocess_qkv(
self,
qkv: torch.Tensor,
attn_metadata: VideoSparseAttentionMetadata,
) -> torch.Tensor:
return self.tile(
qkv,
attn_metadata.num_tiles,
attn_metadata.tile_partition_indices,
attn_metadata.non_pad_index,
)
return self.tile(qkv, attn_metadata)
def postprocess_output(
self,
output: torch.Tensor,
attn_metadata: VideoSparseAttentionMetadata,
) -> torch.Tensor:
return self.untile(
output,
attn_metadata.reverse_tile_partition_indices,
attn_metadata.non_pad_index,
)
return self.untile(output, attn_metadata.untile_combined_index)
def forward( # type: ignore[override]
self,
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
+2
View File
@@ -624,6 +624,8 @@ class WanTransformerBlock_VSA(nn.Module):
added_kv_proj_dim: int | None = None,
supported_attention_backends: set[AttentionBackendEnum] | None = None,
prefix: str = "",
attention_type: str = "original",
sla_topk: float = 0.0,
quant_config: QuantizationConfig | None = None,
):
super().__init__()
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+15 -3
View File
@@ -221,12 +221,20 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
if not backends:
return None
if len(backends) > 1:
sparse_backends = {backend for backend in backends if backend.is_sparse}
selected_backend = (
sorted(sparse_backends, key=lambda backend: backend.name)[0]
if sparse_backends
else sorted(backends, key=lambda backend: backend.name)[0]
)
logger.warning(
"Multiple transformer attention backends detected: %s. "
"Using one backend for denoising metadata.",
"Using %s for denoising metadata.",
sorted(backend.name.lower() for backend in backends),
selected_backend.name.lower(),
)
return sorted(backends, key=lambda backend: backend.name)[0]
return selected_backend
return next(iter(backends))
def component_uses(
self, server_args: ServerArgs, stage_name: str | None = None
@@ -1474,12 +1482,16 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN
or self.attn_backend.get_enum() == AttentionBackendEnum.VIDEO_SPARSE_ATTN
):
attention_backend_config = server_args.attention_backend_config or {}
vsa_sparsity = attention_backend_config.get(
"VSA_sparsity", attention_backend_config.get("sparsity", 0.0)
)
attn_metadata = self.attn_metadata_builder.build(
current_timestep=i,
raw_latent_shape=batch.raw_latent_shape[2:5],
patch_size=server_args.pipeline_config.dit_config.patch_size,
STA_param=batch.STA_param,
VSA_sparsity=server_args.attention_backend_config.VSA_sparsity,
VSA_sparsity=vsa_sparsity,
device=get_local_torch_device(),
)
elif (
python/sglang/multimodal_gen/test/unit/test_video_sparse_attention.py
+39
View File
@@ -0,0 +1,39 @@
import torch
from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import (
VideoSparseAttentionImpl,
VideoSparseAttentionMetadataBuilder,
)
def test_video_sparse_attention_tile_buffer_reuse_and_untile():
metadata = VideoSparseAttentionMetadataBuilder().build(
current_timestep=0,
raw_latent_shape=(5, 7, 9),
patch_size=(1, 1, 1),
VSA_sparsity=0.5,
device=torch.device("cpu"),
)
impl = object.__new__(VideoSparseAttentionImpl)
total_seq_length = metadata.total_seq_length
x = torch.arange(2 * total_seq_length * 3 * 4, dtype=torch.float32).reshape(
2, total_seq_length, 3, 4
)
tiled = impl.preprocess_qkv(x, metadata)
assert metadata.tile_buf is tiled
assert torch.equal(
metadata.untile_combined_index,
metadata.non_pad_index[metadata.reverse_tile_partition_indices],
)
assert torch.equal(impl.postprocess_output(tiled, metadata), x)
next_x = x + 1
next_tiled = impl.preprocess_qkv(next_x, metadata)
assert next_tiled.data_ptr() == tiled.data_ptr()
assert torch.equal(impl.postprocess_output(next_tiled, metadata), next_x)
pad_mask = torch.ones(next_tiled.shape[1], dtype=torch.bool)
pad_mask[metadata.non_pad_index.cpu()] = False
assert torch.all(next_tiled[:, pad_mask] == 0)