diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d10c23038a..13c56a526f 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -215,7 +215,7 @@ configuration. | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | -| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 351fdba085..d1a38322ae 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -44,7 +44,11 @@ class AiterMLABackend(MLACommonBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [1] + # The aiter MLA decode kernel always operates with page_size=1 + # internally (the wrapper flattens kv_buffer via .view(-1, 1, 1, H)). + # We support any kernel_block_size by expanding block-level indices + # into per-token flat indices in the metadata builder. + return [MultipleOf(1)] @staticmethod def get_name() -> str: @@ -74,6 +78,8 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): attn_out_dtype: torch.dtype = torch.bfloat16 # The max query output length: int max_qo_len: int | None = None + # Whether persistent MLA metadata was computed (only for qseqlen=1) + has_persistent_metadata: bool = False @dataclass @@ -105,7 +111,16 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): self.compilation_config = vllm_config.compilation_config self.decode_attn_out_dtype = vllm_config.model_config.dtype - # kernel block size is always 1. + + # Store the kernel block size from the spec. When kernel_block_size=1 + # (no spec-dec), behavior is identical to the original. When > 1 + # (e.g. 16 with Eagle3), we expand block-level indices into per-token + # flat indices since the aiter kernel always uses page_size=1 internally. + self.kernel_block_size = kv_cache_spec.block_size + + # In the flat view (.view(-1,1,1,H)), each token is its own page, + # so max_num_pages_per_req = max_model_len regardless of + # kernel_block_size. max_num_pages_per_req = vllm_config.model_config.max_model_len max_num_reqs = vllm_config.scheduler_config.max_num_seqs max_num_pages = max_num_reqs * max_num_pages_per_req @@ -115,8 +130,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # so we can only use the persistent buffer if a cudagraph is actually # being used. - # paged_kv_last_page_len is always 1s (kernel block size is always 1), - # so we create it once and reuse slices in both eager and cudagraph modes. + # paged_kv_last_page_len is always 1s (the aiter kernel always sees + # page_size=1 after .view(-1,1,1,H) flattening), so we create it + # once and reuse slices in both eager and cudagraph modes. self.paged_kv_last_page_len = torch.ones( max_num_reqs, dtype=torch.int32, device=device ) @@ -196,14 +212,14 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): num_decode_tokens: int, dcp_tot_seq_lens_device: torch.Tensor | None, ) -> AiterMLADecodeMetadata: - # kernel block size is always 1, although the kv block size is not 1. device = self.device num_reqs = seq_lens_device.size(0) - # kernel block size is always 1, so each page has exactly 1 token. - # last_page_len is always 1 - just slice the pre-initialized buffer. + # The aiter kernel always operates with page_size=1 (the wrapper + # flattens kv_buffer). last_page_len is always 1. paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs] + # indptr: cumsum of seq_lens (one page per token in the flat view) paged_kv_indptr = torch.cat( [ torch.zeros(1, dtype=seq_lens_device.dtype, device=device), @@ -215,11 +231,19 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): self.paged_kv_indices.fill_(-1) - _copy_page_indices_kernel[(num_reqs,)]( + + # Expand block_table entries into per-token flat indices. + # When kernel_block_size=1, this degrades to a direct copy (identical + # to the original _copy_page_indices_kernel). + # When kernel_block_size=K>1, block_table entry b covering K tokens + # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). + _expand_page_indices_kernel[(num_reqs,)]( self.paged_kv_indices, block_table_tensor, block_table_tensor.stride(0), paged_kv_indptr, + seq_lens_device, + KERNEL_BLOCK_SIZE=self.kernel_block_size, BLOCK_SIZE=1024, ) paged_kv_indices = self.paged_kv_indices @@ -245,27 +269,37 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): 0, num_reqs + 1, step=1, dtype=torch.int32, device=device ) - from aiter import get_mla_metadata_v1 + # The aiter MLA ASM kernel only supports qseqlen=1 (single-token + # decode). With speculative decoding, the verification step has + # qseqlen > 1 (e.g. 8 for spec7). get_mla_metadata_v1 calls + # get_heuristic_kernel_mla which fails for qseqlen > 1. + # We track whether persistent metadata was successfully computed + # so forward_mqa can skip passing it (falling back to the kernel + # computing its own metadata internally, like v0.18.0). + has_persistent_metadata = False + if max_qo_len == 1: + from aiter import get_mla_metadata_v1 - get_mla_metadata_v1( - qo_indptr, - paged_kv_indptr, - paged_kv_last_page_len, - self._num_attention_heads, - 1, - True, - self._mla_work_meta_data, - self._mla_work_info_set, - self._mla_work_indptr, - self._mla_reduce_indptr, - self._mla_reduce_final_map, - self._mla_reduce_partial_map, - page_size=1, - kv_granularity=16, - max_seqlen_qo=max_qo_len, - uni_seqlen_qo=max_qo_len, - fast_mode=True, - ) + get_mla_metadata_v1( + qo_indptr, + paged_kv_indptr, + paged_kv_last_page_len, + self._num_attention_heads, + 1, + True, + self._mla_work_meta_data, + self._mla_work_info_set, + self._mla_work_indptr, + self._mla_reduce_indptr, + self._mla_reduce_final_map, + self._mla_reduce_partial_map, + page_size=1, + kv_granularity=16, + max_seqlen_qo=max_qo_len, + uni_seqlen_qo=max_qo_len, + fast_mode=True, + ) + has_persistent_metadata = True attn_metadata = AiterMLADecodeMetadata( block_table=block_table_tensor, @@ -277,6 +311,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): dcp_tot_seq_lens=dcp_tot_seq_lens_device, max_qo_len=max_qo_len, attn_out_dtype=self.decode_attn_out_dtype, + has_persistent_metadata=has_persistent_metadata, ) return attn_metadata @@ -290,41 +325,67 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) - attn_metadata.work_meta_data = self._mla_work_meta_data - attn_metadata.work_indptr = self._mla_work_indptr - attn_metadata.work_info_set = self._mla_work_info_set - attn_metadata.reduce_indptr = self._mla_reduce_indptr - attn_metadata.reduce_final_map = self._mla_reduce_final_map - attn_metadata.reduce_partial_map = self._mla_reduce_partial_map + if ( + attn_metadata.decode is not None + and attn_metadata.decode.has_persistent_metadata + ): + attn_metadata.work_meta_data = self._mla_work_meta_data + attn_metadata.work_indptr = self._mla_work_indptr + attn_metadata.work_info_set = self._mla_work_info_set + attn_metadata.reduce_indptr = self._mla_reduce_indptr + attn_metadata.reduce_final_map = self._mla_reduce_final_map + attn_metadata.reduce_partial_map = self._mla_reduce_partial_map return attn_metadata @triton.jit -def _copy_page_indices_kernel( +def _expand_page_indices_kernel( page_indices, block_table, block_table_stride, - cu_num_blocks, + cu_num_tokens, + seq_lens, + KERNEL_BLOCK_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): - """Copy block table rows into a flat page_indices buffer using indptr. - Avoids blocking boolean mask indexing (tensor[mask]) which has - data-dependent output size and forces sync. - This is the same kernel as introduced in backends/flashinfer.py. + """Expand block table entries into per-token flat page indices. + + The aiter MLA kernel always operates with page_size=1 internally + (kv_buffer is flattened via .view(-1, 1, 1, H)). This kernel converts + block-level indices from the block table into individual token positions + in the flattened KV buffer. + + When KERNEL_BLOCK_SIZE=1: block_idx=t, offset=0, flat=block_id + (equivalent to a direct copy -- no regression from the original kernel). + + When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens) + is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). """ req_idx = tl.program_id(0) row_ptr = block_table + req_idx * block_table_stride - start_idx = tl.load(cu_num_blocks + req_idx) - end_idx = tl.load(cu_num_blocks + req_idx + 1) - num_blocks = end_idx - start_idx + start_idx = tl.load(cu_num_tokens + req_idx) + num_tokens = tl.load(seq_lens + req_idx) offset = tl.arange(0, BLOCK_SIZE) - for i in tl.range(0, num_blocks, BLOCK_SIZE): - block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks) + for i in tl.range(0, num_tokens, BLOCK_SIZE): + token_offsets = i + offset + mask = token_offsets < num_tokens + + # Which block in the block table does this token belong to? + block_idx = token_offsets // KERNEL_BLOCK_SIZE + # Offset within that block + offset_in_block = token_offsets % KERNEL_BLOCK_SIZE + + # Load the block ID from the block table + block_ids = tl.load(row_ptr + block_idx, mask=mask) + + # Compute flat index in the flattened kv_buffer + flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block + tl.store( - page_indices + start_idx + i + offset, - block_ids, - mask=i + offset < num_blocks, + page_indices + start_idx + token_offsets, + flat_indices, + mask=mask, ) @@ -426,6 +487,24 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) + # Build kwargs for mla_decode_fwd. Pass persistent metadata only + # when it was successfully computed (qseqlen=1 decode steps). + # For multi-token verification steps (spec-dec), the kernel falls + # back to computing metadata internally. + mla_kwargs = dict( + q_scale=layer._q_scale, + kv_scale=layer._k_scale, + ) + if attn_metadata.work_meta_data is not None: + mla_kwargs.update( + work_meta_data=attn_metadata.work_meta_data, + work_indptr=attn_metadata.work_indptr, + work_info_set=attn_metadata.work_info_set, + reduce_indptr=attn_metadata.reduce_indptr, + reduce_final_map=attn_metadata.reduce_final_map, + reduce_partial_map=attn_metadata.reduce_partial_map, + ) + rocm_aiter_ops.mla_decode_fwd( q, kv_buffer, @@ -436,14 +515,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): attn_metadata.decode.paged_kv_indptr, attn_metadata.decode.paged_kv_indices, attn_metadata.decode.paged_kv_last_page_len, - q_scale=layer._q_scale, - kv_scale=layer._k_scale, - work_meta_data=attn_metadata.work_meta_data, - work_indptr=attn_metadata.work_indptr, - work_info_set=attn_metadata.work_info_set, - reduce_indptr=attn_metadata.reduce_indptr, - reduce_final_map=attn_metadata.reduce_final_map, - reduce_partial_map=attn_metadata.reduce_partial_map, + **mla_kwargs, ) if self._needs_head_repeat: