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 abe11b207..4f72dcc1f 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 @@ -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") 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 index 3e1ce351e..8b6ffbf44 100644 --- a/python/sglang/multimodal_gen/test/unit/test_video_sparse_attention.py +++ b/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