diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py index c6c25ebf2..793e2c52b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py +++ b/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 - 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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py index 2ee9a17df..abe11b207 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py +++ b/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], - 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( - ( - x.shape[0], - t_padded_size * h_padded_size * w_padded_size, - x.shape[-2], - x.shape[-1], - ), - device=x.device, - dtype=x.dtype, + target_shape = ( + x.shape[0], + t_padded_size * h_padded_size * w_padded_size, + x.shape[-2], + x.shape[-1], ) - 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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 89f5c14e9..782e99028 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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__() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 1f598d0d9..55fe2118b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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 ( diff --git a/python/sglang/multimodal_gen/test/unit/test_video_sparse_attention.py b/python/sglang/multimodal_gen/test/unit/test_video_sparse_attention.py new file mode 100644 index 000000000..3e1ce351e --- /dev/null +++ b/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)