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

python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py

@@ -323,8 +323,8 @@ class SparseLinearAttentionImpl(AttentionImpl, nn.Module):
         )
 
         # Apply feature maps
-        query = self.feature_map_q(query).contiguous().to(self.dtype)  # c_q
+        query = self.feature_map_q(query).to(self.dtype)  # c_q
-        key = self.feature_map_k(key).contiguous().to(self.dtype)  # c_k
+        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)
+        q_linear = self.feature_map_q(q).to(self.dtype)
-        k_linear = self.feature_map_k(k).contiguous().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

@@ -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],
+        attn_metadata: VideoSparseAttentionMetadata,
-        tile_partition_indices: torch.LongTensor,
-        non_pad_index: torch.LongTensor,
     ) -> 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]
-
+        target_shape = (
-        x_padded = torch.zeros(
+            x.shape[0],
-            (
+            t_padded_size * h_padded_size * w_padded_size,
-                x.shape[0],
+            x.shape[-2],
-                t_padded_size * h_padded_size * w_padded_size,
+            x.shape[-1],
-                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,
+        untile_combined_index: torch.LongTensor,
-        non_pad_index: torch.LongTensor,
     ) -> torch.Tensor:
-        x = x[:, non_pad_index][:, reverse_tile_partition_indices]
+        return x[:, untile_combined_index]
-        return x
 
     def preprocess_qkv(
         self,
         qkv: torch.Tensor,
         attn_metadata: VideoSparseAttentionMetadata,
     ) -> torch.Tensor:
-        return self.tile(
+        return self.tile(qkv, attn_metadata)
-            qkv,
-            attn_metadata.num_tiles,
-            attn_metadata.tile_partition_indices,
-            attn_metadata.non_pad_index,
-        )
 
     def postprocess_output(
         self,
         output: torch.Tensor,
         attn_metadata: VideoSparseAttentionMetadata,
     ) -> torch.Tensor:
-        return self.untile(
+        return self.untile(output, attn_metadata.untile_combined_index)
-            output,
-            attn_metadata.reverse_tile_partition_indices,
-            attn_metadata.non_pad_index,
-        )
 
     def forward(  # type: ignore[override]
         self,

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -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

@@ -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

@@ -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)