Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions vllm/v1/attention/backends/mla/indexer.py
Original file line number Diff line number Diff line change
Expand Up @@ -1261,7 +1261,6 @@ def build(
assert num_decode_tokens + num_prefill_tokens == num_tokens

compressed_slot_mapping = slot_mapping
compressed_seq_lens = seq_lens
indexer_block_table = block_table
if self.compress_ratio > 1:
kernel_block_size = self.kernel_block_size
Expand Down Expand Up @@ -1297,10 +1296,12 @@ def build(
self.compressed_slot_mapping_buffer[:padded_num_tokens],
dim=0,
)
compressed_seq_lens = seq_lens // self.compress_ratio

prefill_metadata = None
if num_prefills > 0:
compressed_seq_lens = (
seq_lens // self.compress_ratio if self.compress_ratio > 1 else seq_lens
)
# This CPU value is an upper bound for async-spec extend rows. It
# is safe for chunking/allocation because CUDA metadata below is
# built from exact device seq_lens and gather ignores the tail.
Expand Down
2 changes: 1 addition & 1 deletion vllm/v1/attention/backends/mla/sparse_swa.py
Original file line number Diff line number Diff line change
Expand Up @@ -613,7 +613,7 @@ def build(
)

is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
is_valid_token.copy_(slot_mapping >= 0)
torch.ge(slot_mapping, 0, out=is_valid_token)
if replay_start is None:
# Graph captures build without it; their dummy batches replay nothing.
replay_start = self.no_replay_start
Expand Down
Loading