fix(diffusion): size VSA top-k from padded blocks (#32695)

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

@@ -124,7 +124,6 @@ def get_non_pad_index(
 
 
 class VideoSparseAttentionBackend(AttentionBackend):
-
     accept_output_buffer: bool = True
 
     @staticmethod
@@ -170,8 +169,13 @@ class VideoSparseAttentionMetadata(AttentionMetadata):
     max_seqlen_k: int = 0
 
 
-class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
+def _compute_cur_topk(attn_metadata: VideoSparseAttentionMetadata) -> int:
+    num_kv_blocks = attn_metadata.variable_block_sizes.numel()
+    cur_topk = math.ceil((1 - attn_metadata.VSA_sparsity) * num_kv_blocks)
+    return max(1, min(cur_topk, num_kv_blocks))
 
+
+class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
     def __init__(self):
         pass
 
@@ -230,7 +234,6 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
 
 
 class VideoSparseAttentionImpl(AttentionImpl):
-
     def __init__(
         self,
         num_heads: int,
@@ -308,12 +311,7 @@ class VideoSparseAttentionImpl(AttentionImpl):
         value = value.transpose(1, 2).contiguous()
         gate_compress = gate_compress.transpose(1, 2).contiguous()
 
-        VSA_sparsity = attn_metadata.VSA_sparsity
-
-        cur_topk = math.ceil(
-            (1 - VSA_sparsity)
-            * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE))
-        )
+        cur_topk = _compute_cur_topk(attn_metadata)
 
         if video_sparse_attn is None:
             raise NotImplementedError("video_sparse_attn is not installed")

python/sglang/multimodal_gen/test/unit/test_video_sparse_attention.py

@@ -1,8 +1,15 @@
+import math
+
 import torch
 
+from sglang.multimodal_gen.runtime.layers.attention.backends import (
+    video_sparse_attn as vsa_module,
+)
 from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import (
+    VSA_TILE_SIZE,
     VideoSparseAttentionImpl,
     VideoSparseAttentionMetadataBuilder,
+    _compute_cur_topk,
 )
 
 
@@ -37,3 +44,63 @@ def test_video_sparse_attention_tile_buffer_reuse_and_untile():
     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)
+
+
+def test_vsa_forward_cur_topk_uses_padded_kv_block_count(monkeypatch):
+    metadata = VideoSparseAttentionMetadataBuilder().build(
+        current_timestep=0,
+        raw_latent_shape=(5, 32, 32),
+        patch_size=(1, 1, 1),
+        VSA_sparsity=0.75,
+        device=torch.device("cpu"),
+    )
+    num_kv_blocks = metadata.variable_block_sizes.numel()
+    block_elements = math.prod(VSA_TILE_SIZE)
+    padded_seq_len = num_kv_blocks * block_elements
+    expected_topk = math.ceil((1 - metadata.VSA_sparsity) * num_kv_blocks)
+    unpadded_topk = math.ceil(
+        (1 - metadata.VSA_sparsity) * (metadata.total_seq_length / block_elements)
+    )
+    captured = {}
+
+    def fake_video_sparse_attn(
+        query,
+        key,
+        value,
+        variable_block_sizes,
+        topk,
+        block_size,
+        compress_attn_weight,
+    ):
+        captured["topk"] = topk
+        captured["block_size"] = block_size
+        assert torch.equal(variable_block_sizes, metadata.variable_block_sizes)
+        return query
+
+    monkeypatch.setattr(vsa_module, "video_sparse_attn", fake_video_sparse_attn)
+
+    query = torch.ones(1, padded_seq_len, 1, 1)
+    output = object.__new__(VideoSparseAttentionImpl).forward(
+        query, query, query, query, metadata
+    )
+
+    assert unpadded_topk < expected_topk
+    assert captured["topk"] == expected_topk
+    assert captured["block_size"] == VSA_TILE_SIZE
+    assert output.shape == query.shape
+
+
+def test_vsa_cur_topk_clamps_to_valid_block_range():
+    metadata = VideoSparseAttentionMetadataBuilder().build(
+        current_timestep=0,
+        raw_latent_shape=(5, 32, 32),
+        patch_size=(1, 1, 1),
+        VSA_sparsity=1.0,
+        device=torch.device("cpu"),
+    )
+    num_kv_blocks = metadata.variable_block_sizes.numel()
+
+    assert _compute_cur_topk(metadata) == 1
+
+    metadata.VSA_sparsity = -0.01
+    assert _compute_cur_topk(metadata) == num_kv_blocks