[Perf][Attention] Avoid redundant sparse attention metadata operations (#57885)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com>
Woosuk Kwon committed
Sep 21, 2026 at 00:08 UTC
97dc6b19d2fe92b794b49e684677a2e1a8b2c540
2 files changed
+4
-3
vllm/v1/attention/backends/mla/indexer.py
+3
-2
@@ -1261,7 +1261,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1261
assert num_decode_tokens + num_prefill_tokens == num_tokens
1262
1263
compressed_slot_mapping = slot_mapping
1264
- compressed_seq_lens = seq_lens
1264
indexer_block_table = block_table
1265
if self.compress_ratio > 1:
1266
kernel_block_size = self.kernel_block_size
@@ -1297,10 +1296,12 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1296
self.compressed_slot_mapping_buffer[:padded_num_tokens],
1297
dim=0,
1298
)
1300
- compressed_seq_lens = seq_lens // self.compress_ratio
1299
1300
prefill_metadata = None
1301
if num_prefills > 0:
1302
+ compressed_seq_lens = (
1303
+ seq_lens // self.compress_ratio if self.compress_ratio > 1 else seq_lens
1304
+ )
1305
# This CPU value is an upper bound for async-spec extend rows. It
1306
# is safe for chunking/allocation because CUDA metadata below is
1307
# built from exact device seq_lens and gather ignores the tail.
vllm/v1/attention/backends/mla/sparse_swa.py
+1
-1
@@ -613,7 +613,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
613
)
614
615
is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
616
- is_valid_token.copy_(slot_mapping >= 0)
616
+ torch.ge(slot_mapping, 0, out=is_valid_token)
617
if replay_start is None:
618
# Graph captures build without it; their dummy batches replay nothing.
619
replay_start = self.no_replay_start