Repository: vllm-project/vllm
File: vllm/v1/attention/backends/mla/indexer.py (Python)

Code before:
                )
                cumsum = torch.cumsum(decode_lens, dim=0, dtype=torch.int32)
                starts = torch.zeros_like(cumsum)
                starts[1:] = cumsum[:-1]
                expanded_starts = torch.repeat_interleave(starts, decode_lens)
                positions_within = (
                    torch.arange(
                        actual_expanded,
                        device=self.device,
                        dtype=torch.int32,
                    )
                    - expanded_starts
                )

                self.expanded_seq_lens_buffer[:actual_expanded] = (
                    expanded_base + positions_within + 1
                )
                self.expanded_seq_lens_buffer[:actual_expanded].clamp_(min=0)
                self.expanded_seq_lens_buffer[actual_expanded:] = 0
                seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]

                self._expanded_block_table_buffer[:actual_expanded] = (
                    torch.repeat_interleave(block_table, decode_lens, dim=0)
                )
                if actual_expanded < num_decode_tokens:
                    self._expanded_block_table_buffer[
                        actual_expanded:num_decode_tokens
                    ] = 0
                block_table = self._expanded_block_table_buffer[:num_decode_tokens]

                self.decode_lens_buffer[:num_decode_tokens] = 1
                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
                requires_padding = False
                offsets = None
                batch_size = num_decode_tokens
            else:
                next_n = 1 + self.num_speculative_tokens
                if next_n > 1:
                    offsets = torch.arange(
                        next_n, device=self.device, dtype=torch.int32
                    )
                else:
                    offsets = None
                batch_size = num_decodes

            # DeepGEMM is required for the paged MQA logits on CUDA devices
            if current_platform.is_cuda() and has_deep_gemm():
                self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
                    seq_lens,
                    self.kv_cache_spec.block_size,
                    self.num_sms,

Review comment (suggestion):
we may be able to get away with `self._expanded_block_table_buffer[actual_expanded:num_decode_tokens, 0] = 0`

Code after:
                # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
                # The context lengths are therefore
                # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].

                # 3 + 1 + 4 + 0 = 8
                actual_expanded = int(decode_lens_cpu.sum().item())

                # [7, 6, 8, 0] -> [7, 7, 7, 6, 8, 8, 8, 8]
                expanded_base = torch.repeat_interleave(
                    seq_lens - decode_lens, decode_lens
                )

                # [0, 3, 4, 8] -> [0, 0, 0, 3, 4, 4, 4, 4]
                expanded_starts = torch.repeat_interleave(
                    common_attn_metadata.query_start_loc[:num_decodes], decode_lens
                )

                # [0, 1, 2, 0, 0, 1, 2, 3]
                positions_within = (
                    self.arange_buffer[:actual_expanded] - expanded_starts
                )

                # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
                self.expanded_seq_lens_buffer[:actual_expanded] = (
                    expanded_base + positions_within + 1
                )
                self.expanded_seq_lens_buffer[actual_expanded:] = 0
                seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]

                # Give each of the flattened entries the same block table row as the
                # original request.
                self.expanded_block_table_buffer[:actual_expanded] = (
                    torch.repeat_interleave(block_table, decode_lens, dim=0)
                )
                if actual_expanded < num_decode_tokens:
                    self.expanded_block_table_buffer[
                        actual_expanded:num_decode_tokens, 0
                    ] = 0
                block_table = self.expanded_block_table_buffer[:num_decode_tokens]

                # All reqs now have decode_len=1
                self.decode_lens_buffer[:num_decode_tokens] = 1
                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
                offsets = None
                batch_size = num_decode_tokens
            else:
                next_n = 1 + self.num_speculative_tokens
                if next_n > 1:
                    offsets = torch.arange(
                        next_n, device=self.device, dtype=torch.int32
                    )