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

Code before:
            # Use CPU to avoid GPU sync; breaking async scheduling
            requires_padding = (decode_lens_cpu.max() > decode_lens_cpu.min()).item()

            seq_lens = common_attn_metadata.seq_lens[:num_decodes]
            block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]

            # Padded CUDA graph requests have block_table entries of -1.
            # Clamp to 0 to prevent OOB access in the DeepGEMM kernel.
            # This is safe because padded requests have seq_lens=0, so the
            # kernel produces no meaningful output for those rows.
            block_table.clamp_(min=0)

            max_decode_len = int(decode_lens_cpu.max().item())
            if max_decode_len > 2:
                # fp8_paged_mqa_logits does not support next_n > 2.
                # Flatten multi-token decode requests into single-token
                # batch entries, expanding seq_lens and block tables so
                # the kernel always sees next_n=1.
                #
                # For request i with seq_len=S, decode_len=L, create L
                # entries with seqlens [S-L+1, S-L+2, ..., S].
                actual_expanded = int(decode_lens_cpu.sum().item())
                expanded_base = torch.repeat_interleave(
                    seq_lens - decode_lens, decode_lens
                )
                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

Review comment (suggestion):
reuse `common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]`

Code after:
                common_attn_metadata.query_start_loc[: num_decodes + 1],
                out=self.decode_lens_buffer[:num_decodes],
            )
            decode_lens = self.decode_lens_buffer[:num_decodes]
            decode_lens_cpu = torch.diff(
                common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
            )

            seq_lens = common_attn_metadata.seq_lens[:num_decodes]
            block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]

            # Padded CUDA graph requests have block_table entries of -1.
            # Clamp to 0 to prevent OOB access in the DeepGEMM kernel.
            # This is safe because padded requests have seq_lens=0, so the
            # kernel produces no meaningful output for those rows.
            block_table.clamp_(min=0)

            max_decode_len = int(decode_lens_cpu.max().item())
            if max_decode_len > 1:
                # Flatten multi-token decode requests into single-token
                # batch entries, expanding seq_lens and block tables so
                # the kernel always sees next_n=1.

                # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
                # 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