| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import TYPE_CHECKING, ClassVar |
| 6 | |
| 7 | import numpy as np |
| 8 | import torch |
| 9 | |
| 10 | from vllm import _custom_ops as ops |
| 11 | from vllm._aiter_ops import rocm_aiter_ops |
| 12 | from vllm.config import VllmConfig, get_current_vllm_config |
| 13 | from vllm.config.cache import CacheDType |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 16 | get_mla_dims, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.attention.sparse_mla_attention import ( |
| 19 | SharedTopkIndicesBuffer, |
| 20 | ) |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.triton_utils import tl, triton |
| 23 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 24 | from vllm.utils.torch_utils import np_to_pinned_tensor |
| 25 | from vllm.v1.attention.backend import ( |
| 26 | AttentionBackend, |
| 27 | AttentionCGSupport, |
| 28 | AttentionLayer, |
| 29 | AttentionMetadata, |
| 30 | AttentionMetadataBuilder, |
| 31 | CommonAttentionMetadata, |
| 32 | MLAAttentionImpl, |
| 33 | MultipleOf, |
| 34 | ) |
| 35 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( |
| 36 | AiterMLAHelper, |
| 37 | ) |
| 38 | from vllm.v1.attention.backends.utils import split_decodes_and_prefills |
| 39 | from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( |
| 40 | rocm_sparse_attn_prefill, |
| 41 | ) |
| 42 | from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout |
| 43 | from vllm.v1.worker.workspace import current_workspace_manager |
| 44 | |
| 45 | if TYPE_CHECKING: |
| 46 | from vllm.model_executor.models.deepseek_v2 import Indexer |
| 47 | logger = init_logger(__name__) |
| 48 | |
| 49 | |
| 50 | def _use_rocm_sparse_triton( |
| 51 | *, |
| 52 | kv_cache_dtype: str, |
| 53 | head_size: int, |
| 54 | kv_lora_rank: int, |
| 55 | num_prefills: int, |
| 56 | num_decodes: int, |
| 57 | num_decode_tokens: int, |
| 58 | max_query_len: int, |
| 59 | ) -> bool: |
| 60 | """Select the rope-free BF16 path not supported by AITER sparse MLA. |
| 61 | |
| 62 | The ragged Triton kernel indexes metadata per query token, so multi-token |
| 63 | speculative verification rows have the same capability requirements as |
| 64 | plain decode rows. |
| 65 | """ |
| 66 | return ( |
| 67 | not kv_cache_dtype.startswith("fp8") |
| 68 | and head_size == kv_lora_rank |
| 69 | and (num_prefills > 0 or num_decodes > 0) |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | @triton.jit |
| 74 | def _fit_kpool_indices_kernel( |
| 75 | token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS + TAIL_WIDTH] |
| 76 | out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] |
| 77 | ti_stride0, |
| 78 | ti_stride1, |
| 79 | out_stride0, |
| 80 | out_stride1, |
| 81 | NUM_TOPK_TOKENS: tl.constexpr, |
| 82 | TAIL_WIDTH: tl.constexpr, |
| 83 | BLOCK_T: tl.constexpr, |
| 84 | BLOCK_TAIL: tl.constexpr, |
| 85 | ): |
| 86 | row_ptr = token_indices_ptr + tl.program_id(0) * ti_stride0 |
| 87 | columns = tl.arange(0, BLOCK_T) |
| 88 | column_mask = columns < NUM_TOPK_TOKENS |
| 89 | history = tl.load(row_ptr + columns * ti_stride1, mask=column_mask, other=-1) |
| 90 | |
| 91 | tail_columns = tl.arange(0, BLOCK_TAIL) |
| 92 | tail = tl.load( |
| 93 | row_ptr + (NUM_TOPK_TOKENS + tail_columns) * ti_stride1, |
| 94 | mask=tail_columns < TAIL_WIDTH, |
| 95 | other=-1, |
| 96 | ) |
| 97 | |
| 98 | valid_history = tl.sum((history >= 0).to(tl.int32)) |
| 99 | valid_tail = tl.sum((tail >= 0).to(tl.int32)) |
| 100 | keep_history = tl.minimum(valid_history, NUM_TOPK_TOKENS - valid_tail) |
| 101 | |
| 102 | tail_offsets = columns - keep_history |
| 103 | tail_values = tl.load( |
| 104 | row_ptr + (NUM_TOPK_TOKENS + tail_offsets) * ti_stride1, |
| 105 | mask=column_mask & (tail_offsets >= 0) & (tail_offsets < TAIL_WIDTH), |
| 106 | other=-1, |
| 107 | ) |
| 108 | fitted = tl.where(columns < keep_history, history, tail_values) |
| 109 | fitted = tl.where(columns < keep_history + valid_tail, fitted, -1) |
| 110 | tl.store( |
| 111 | out_ptr + tl.program_id(0) * out_stride0 + columns * out_stride1, |
| 112 | fitted, |
| 113 | mask=column_mask, |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | def fit_kpool_indices_to_aiter( |
| 118 | token_indices: torch.Tensor, topk_tokens: int |
| 119 | ) -> torch.Tensor: |
| 120 | """Keep the live kpool tail while fitting AITER's fixed top-k width.""" |
| 121 | if token_indices.shape[1] < topk_tokens: |
| 122 | raise ValueError("token_indices width must be at least topk_tokens") |
| 123 | if token_indices.shape[1] == topk_tokens: |
| 124 | return token_indices |
| 125 | |
| 126 | num_tokens, width = token_indices.shape |
| 127 | tail_width = width - topk_tokens |
| 128 | output = torch.empty( |
| 129 | (num_tokens, topk_tokens), |
| 130 | dtype=token_indices.dtype, |
| 131 | device=token_indices.device, |
| 132 | ) |
| 133 | if num_tokens == 0: |
| 134 | return output |
| 135 | |
| 136 | _fit_kpool_indices_kernel[(num_tokens,)]( |
| 137 | token_indices, |
| 138 | output, |
| 139 | token_indices.stride(0), |
| 140 | token_indices.stride(1), |
| 141 | output.stride(0), |
| 142 | output.stride(1), |
| 143 | NUM_TOPK_TOKENS=topk_tokens, |
| 144 | TAIL_WIDTH=tail_width, |
| 145 | BLOCK_T=triton.next_power_of_2(topk_tokens), |
| 146 | BLOCK_TAIL=triton.next_power_of_2(tail_width), |
| 147 | ) |
| 148 | return output |
| 149 | |
| 150 | |
| 151 | @triton.jit |
| 152 | def _convert_req_index_to_global_index_kernel( |
| 153 | req_id_ptr, # int32 [num_tokens] |
| 154 | block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] |
| 155 | token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] |
| 156 | cu_seqlens_ptr, # int32 [num_tokens + 1] |
| 157 | out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] |
| 158 | # shapes (compile-time where possible) |
| 159 | max_num_blocks_per_req: tl.constexpr, |
| 160 | BLOCK_SIZE: tl.constexpr, |
| 161 | BLOCK_N: tl.constexpr, # tile width along columns |
| 162 | # strides (in elements) |
| 163 | bt_stride0, |
| 164 | bt_stride1, |
| 165 | ti_stride0, |
| 166 | ti_stride1, |
| 167 | ): |
| 168 | # program_id(0) -> token_id (row) |
| 169 | # program_id(1) -> tile index along columns |
| 170 | token_id = tl.program_id(0) |
| 171 | tile_id = tl.program_id(1) |
| 172 | |
| 173 | # Each program covers BLOCK_N consecutive columns |
| 174 | indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) |
| 175 | |
| 176 | # Load request id for this token (no mask: grid is exact) |
| 177 | req = tl.load(req_id_ptr + token_id) |
| 178 | |
| 179 | # Load cumulative sequence lengths to get starting index of this request |
| 180 | seq_start = tl.load(cu_seqlens_ptr + token_id) |
| 181 | seq_end = tl.load(cu_seqlens_ptr + token_id + 1) |
| 182 | |
| 183 | if tile_id * BLOCK_N + seq_start >= seq_end: |
| 184 | return |
| 185 | |
| 186 | # Load token indices for this tile |
| 187 | ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 |
| 188 | tok = tl.load(ti_ptr) # int32 |
| 189 | |
| 190 | # Only token == -1 should propagate as -1 |
| 191 | is_invalid_tok = tok < 0 |
| 192 | |
| 193 | # Compute block id and in-block offset |
| 194 | block_id = tok // BLOCK_SIZE |
| 195 | inblock_off = tok % BLOCK_SIZE |
| 196 | |
| 197 | # Guard block_table access |
| 198 | valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) |
| 199 | bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 |
| 200 | base = tl.load(bt_ptr, mask=valid_block, other=0) |
| 201 | |
| 202 | # # If token == -1 OR block_id OOB, output 0; else base * BLOCK_SIZE + offset |
| 203 | out_val = tl.where( |
| 204 | is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off |
| 205 | ) |
| 206 | out_ptr_ij = out_ptr + seq_start + indice_id |
| 207 | out_ptr_ij_mask = (seq_start + indice_id) < seq_end |
| 208 | |
| 209 | # store the results with mask |
| 210 | tl.store(out_ptr_ij, out_val, mask=out_ptr_ij_mask) |
| 211 | |
| 212 | |
| 213 | def triton_convert_req_index_to_global_index( |
| 214 | req_id: torch.Tensor, # int32 [num_tokens] |
| 215 | block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] |
| 216 | token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] |
| 217 | cu_seqlens: torch.Tensor, # int32 [num_tokens + 1] |
| 218 | paged_kv_indices: torch.Tensor, # int32 [num_tokens * topk] out_buffer |
| 219 | BLOCK_SIZE: int = 64, |
| 220 | NUM_TOPK_TOKENS: int = 2048, |
| 221 | BLOCK_N: int = 128, # tile width along columns |
| 222 | ): |
| 223 | """out[token_id, indice_id] = |
| 224 | block_table[req_id[token_id], |
| 225 | token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE |
| 226 | + token_indices[token_id, indice_id] % BLOCK_SIZE |
| 227 | |
| 228 | Only when token_indices[token_id, indice_id] == -1 do we output -1. |
| 229 | For safety, we also output -1 if the derived block_id would be |
| 230 | out-of-bounds. |
| 231 | """ |
| 232 | assert req_id.dtype == torch.int32 |
| 233 | assert block_table.dtype == torch.int32 |
| 234 | assert token_indices.dtype == torch.int32 |
| 235 | assert token_indices.shape[1] == NUM_TOPK_TOKENS |
| 236 | assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( |
| 237 | f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible byBLOCK_N ({BLOCK_N})" |
| 238 | ) |
| 239 | # print("req_id: ", req_id, flush=True) |
| 240 | num_tokens = req_id.shape[0] |
| 241 | _, max_num_blocks_per_req = block_table.shape |
| 242 | tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N |
| 243 | |
| 244 | # Ensure contiguous tensors on the same device |
| 245 | req_id_c = req_id.contiguous() |
| 246 | block_table_c = block_table.contiguous() |
| 247 | token_indices_c = token_indices.contiguous() |
| 248 | |
| 249 | # Strides in elements |
| 250 | bt_stride0, bt_stride1 = block_table_c.stride() |
| 251 | ti_stride0, ti_stride1 = token_indices_c.stride() |
| 252 | |
| 253 | # Exact 2D grid: tokens × column tiles |
| 254 | grid = (num_tokens, tiles_per_row) |
| 255 | |
| 256 | _convert_req_index_to_global_index_kernel[grid]( |
| 257 | req_id_c, |
| 258 | block_table_c, |
| 259 | token_indices_c, |
| 260 | cu_seqlens, |
| 261 | paged_kv_indices, |
| 262 | # shapes / constexprs |
| 263 | max_num_blocks_per_req, |
| 264 | BLOCK_SIZE, |
| 265 | BLOCK_N, |
| 266 | # strides |
| 267 | bt_stride0, |
| 268 | bt_stride1, |
| 269 | ti_stride0, |
| 270 | ti_stride1, |
| 271 | ) |
| 272 | return |
| 273 | |
| 274 | |
| 275 | @triton.jit |
| 276 | def generate_sparse_seqlen_kernel( |
| 277 | seq_len_ptr, # [num_seq] |
| 278 | cu_query_lens_ptr, # [num_seq] |
| 279 | out_ptr, # [num_query_tokens] |
| 280 | topk_token: tl.constexpr, |
| 281 | BLOCK_SIZE: tl.constexpr, |
| 282 | ): |
| 283 | seq_id = tl.program_id(0) |
| 284 | query_offset = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 285 | query_start = tl.load(cu_query_lens_ptr + seq_id) |
| 286 | query_end = tl.load(cu_query_lens_ptr + seq_id + 1) |
| 287 | if query_start + tl.program_id(1) * BLOCK_SIZE > query_end: |
| 288 | return |
| 289 | query_len = query_end - query_start |
| 290 | query_mask = query_offset + query_start < query_end |
| 291 | seq_len = tl.load(seq_len_ptr + seq_id) |
| 292 | # Just return since the out_ptr is zero initialized. |
| 293 | if seq_len == 0: |
| 294 | return |
| 295 | context_start_point = seq_len - query_len |
| 296 | sparse_seqlen = context_start_point + query_offset |
| 297 | sparse_seqlen_masked = tl.where( |
| 298 | sparse_seqlen + 1 < topk_token, sparse_seqlen + 1, topk_token |
| 299 | ) |
| 300 | tl.store( |
| 301 | out_ptr + query_start + query_offset, sparse_seqlen_masked, mask=query_mask |
| 302 | ) |
| 303 | |
| 304 | |
| 305 | def generate_sparse_seqlen_triton( |
| 306 | query_lens: torch.Tensor, |
| 307 | seq_lens: torch.Tensor, |
| 308 | cu_query_lens: torch.Tensor, |
| 309 | topk_token: int, |
| 310 | num_tokens: int, |
| 311 | max_query_len: int, |
| 312 | ): |
| 313 | num_seqs = query_lens.size(0) |
| 314 | # zero initialize the tensor to make sure invalid positions will be zero |
| 315 | out = torch.zeros([num_tokens], dtype=torch.int32, device=query_lens.device) |
| 316 | block_size = 64 |
| 317 | num_block_per_row = triton.cdiv(max_query_len, block_size) |
| 318 | grid = ( |
| 319 | num_seqs, |
| 320 | num_block_per_row, |
| 321 | ) |
| 322 | generate_sparse_seqlen_kernel[grid]( |
| 323 | seq_lens, |
| 324 | cu_query_lens, |
| 325 | out, |
| 326 | topk_token, |
| 327 | block_size, |
| 328 | ) |
| 329 | return out |
| 330 | |
| 331 | |
| 332 | class ROCMAiterMLASparseBackend(AttentionBackend): |
| 333 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] |
| 334 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 335 | "auto", |
| 336 | "float16", |
| 337 | "bfloat16", |
| 338 | "fp8", |
| 339 | "fp8_e4m3", |
| 340 | ] |
| 341 | |
| 342 | @staticmethod |
| 343 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 344 | return [1, MultipleOf(16)] |
| 345 | |
| 346 | @staticmethod |
| 347 | def get_name() -> str: |
| 348 | return "ROCM_AITER_MLA_SPARSE" |
| 349 | |
| 350 | @staticmethod |
| 351 | def get_metadata_cls() -> type["ROCMAiterMLASparseMetadata"]: |
| 352 | return ROCMAiterMLASparseMetadata |
| 353 | |
| 354 | @staticmethod |
| 355 | def get_builder_cls() -> type["ROCMAiterMLASparseMetadataBuilder"]: |
| 356 | return ROCMAiterMLASparseMetadataBuilder |
| 357 | |
| 358 | @staticmethod |
| 359 | def get_impl_cls() -> type["ROCMAiterMLASparseImpl"]: |
| 360 | return ROCMAiterMLASparseImpl |
| 361 | |
| 362 | @classmethod |
| 363 | def is_mla(cls) -> bool: |
| 364 | return True |
| 365 | |
| 366 | @classmethod |
| 367 | def is_sparse(cls) -> bool: |
| 368 | return True |
| 369 | |
| 370 | @classmethod |
| 371 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 372 | # Global index conversion assumes contiguous pages within each layer. |
| 373 | return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC) |
| 374 | |
| 375 | @classmethod |
| 376 | def supports_sink(cls) -> bool: |
| 377 | from vllm.platforms.rocm import on_mi3xx |
| 378 | |
| 379 | return on_mi3xx() |
| 380 | |
| 381 | |
| 382 | @dataclass |
| 383 | class ROCMAiterMLASparseMetadata(AttentionMetadata): |
| 384 | num_reqs: int |
| 385 | max_query_len: int |
| 386 | max_seq_len: int |
| 387 | |
| 388 | num_actual_tokens: int # Number of tokens excluding padding. |
| 389 | query_start_loc: torch.Tensor |
| 390 | slot_mapping: torch.Tensor |
| 391 | |
| 392 | block_table: torch.Tensor |
| 393 | req_id_per_token: torch.Tensor |
| 394 | |
| 395 | qo_indptr: torch.Tensor |
| 396 | paged_kv_last_page_len: torch.Tensor |
| 397 | paged_kv_indices: torch.Tensor |
| 398 | paged_kv_indptr: torch.Tensor |
| 399 | attn_out_dtype: torch.dtype |
| 400 | |
| 401 | block_size: int = 1 |
| 402 | topk_tokens: int = 2048 |
| 403 | |
| 404 | # Fields read by the shared MLA forward. This impl has no dense-MHA prefill |
| 405 | # path (supports_dense_mha_prefill=False), so it always runs the MQA path; |
| 406 | num_decodes: int = 0 |
| 407 | num_prefills: int = 0 |
| 408 | num_decode_tokens: int = 0 |
| 409 | prefill_max_seq_len: int = 0 |
| 410 | prefill: object = None |
| 411 | |
| 412 | # Persistent MLA metadata (only populated when persistent mode is enabled, |
| 413 | # i.e. when the aiter sparse decode kernel supports work-stealing splits). |
| 414 | work_meta_data: torch.Tensor | None = None |
| 415 | work_indptr: torch.Tensor | None = None |
| 416 | work_info_set: torch.Tensor | None = None |
| 417 | reduce_indptr: torch.Tensor | None = None |
| 418 | reduce_final_map: torch.Tensor | None = None |
| 419 | reduce_partial_map: torch.Tensor | None = None |
| 420 | |
| 421 | |
| 422 | @dataclass |
| 423 | class ROCMAiterMLASparseMetadataBuilder( |
| 424 | AttentionMetadataBuilder[ROCMAiterMLASparseMetadata] |
| 425 | ): |
| 426 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 427 | |
| 428 | def __init__( |
| 429 | self, |
| 430 | kv_cache_spec: AttentionSpec, |
| 431 | layer_names: list[str], |
| 432 | vllm_config: VllmConfig, |
| 433 | device: torch.device, |
| 434 | ): |
| 435 | self.kv_cache_spec = kv_cache_spec |
| 436 | self.model_config = vllm_config.model_config |
| 437 | self.model_dtype = vllm_config.model_config.dtype |
| 438 | self.kv_cache_dtype = vllm_config.cache_config.cache_dtype |
| 439 | parallel_config = vllm_config.parallel_config |
| 440 | self.device = device |
| 441 | max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 442 | |
| 443 | self.vllm_config = vllm_config |
| 444 | self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) |
| 445 | |
| 446 | self.num_heads = self.model_config.get_num_attention_heads(parallel_config) |
| 447 | self.mla_dims = get_mla_dims(self.model_config) |
| 448 | self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk |
| 449 | attention_context = vllm_config.compilation_config.static_forward_context |
| 450 | # Sink decode must use AITER's nonpersistent path. In particular, |
| 451 | # gfx942 has no persistent+LSE kernel, and its metadata heuristic |
| 452 | # terminates for HY-V4's TP1 H64 shape. The Triton sparse MLA kernel |
| 453 | # never reads the persistent metadata either. |
| 454 | self._use_persistent_metadata = all( |
| 455 | getattr(attention_context[name].impl, "sinks", None) is None |
| 456 | and not getattr(attention_context[name].impl, "use_aiter_sparse_mla", False) |
| 457 | for name in layer_names |
| 458 | ) |
| 459 | # Bounds the KV-split heuristic (see `_sparse_decode_max_split`). |
| 460 | self._num_compute_units = current_platform.num_compute_units() |
| 461 | self.max_model_len_tensor = torch.tensor( |
| 462 | [self.model_config.max_model_len], device=device, dtype=torch.int32 |
| 463 | ) |
| 464 | # this is ignored by `flash_mla_with_kvcache` if indices not None |
| 465 | self.dummy_block_table = torch.empty( |
| 466 | (1, 1), dtype=torch.int32, device=self.device |
| 467 | ) |
| 468 | |
| 469 | self.req_id_per_token_buffer = torch.zeros( |
| 470 | (vllm_config.scheduler_config.max_num_batched_tokens,), |
| 471 | dtype=torch.int32, |
| 472 | device=device, |
| 473 | ) |
| 474 | self.qo_indptr = torch.arange( |
| 475 | 0, max_num_batched_tokens + 1, dtype=torch.int32, device=device |
| 476 | ) |
| 477 | self.paged_kv_last_page_len = torch.ones( |
| 478 | max_num_batched_tokens, dtype=torch.int32, device=device |
| 479 | ) |
| 480 | |
| 481 | # These two needs to be calculated in runtime, |
| 482 | # but we still needs to prepare the buffer |
| 483 | self.paged_kv_indices = torch.zeros( |
| 484 | [max_num_batched_tokens * self.topk_tokens], |
| 485 | dtype=torch.int32, |
| 486 | device=device, |
| 487 | ) |
| 488 | self.paged_kv_indptr = torch.zeros( |
| 489 | [max_num_batched_tokens + 1], dtype=torch.int32, device=device |
| 490 | ) |
| 491 | |
| 492 | # ----- Persistent MLA metadata buffers ----- |
| 493 | # The aiter sparse decode kernel supports a "persistent" path that |
| 494 | # uses precomputed work-splitting metadata for better load balancing |
| 495 | # across CUs. Mirrors the approach used in rocm_aiter_mla.py. |
| 496 | # |
| 497 | # In the sparse case each query token is its own "batch" entry in the |
| 498 | # qo_indptr (qo_indptr = [0, 1, 2, ..., num_tokens]) and max_qo_len=1. |
| 499 | # We pad get_mla_metadata_info_v1's batch_size to max_num_batched_tokens |
| 500 | # so the buffers are large enough for any decode shape we might see. |
| 501 | from aiter import dtypes, get_mla_metadata_info_v1 |
| 502 | |
| 503 | # Keep metadata sizing consistent with the padded tensor shape passed |
| 504 | # to the sparse decode kernel. |
| 505 | self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads( |
| 506 | self.num_heads |
| 507 | ) |
| 508 | |
| 509 | q_dtype = self.model_dtype |
| 510 | kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") |
| 511 | if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): |
| 512 | kv_cache_dtype_str = "fp8" |
| 513 | else: |
| 514 | kv_cache_dtype_str = "bf16" |
| 515 | kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16) |
| 516 | |
| 517 | ( |
| 518 | (work_meta_data_size, work_meta_data_type), |
| 519 | (work_indptr_size, work_indptr_type), |
| 520 | (work_info_set_size, work_info_set_type), |
| 521 | (reduce_indptr_size, reduce_indptr_type), |
| 522 | (reduce_final_map_size, reduce_final_map_type), |
| 523 | (reduce_partial_map_size, reduce_partial_map_type), |
| 524 | ) = get_mla_metadata_info_v1( |
| 525 | max_num_batched_tokens, |
| 526 | 1, |
| 527 | self._num_attention_heads, |
| 528 | q_dtype, |
| 529 | kv_dtype, |
| 530 | is_sparse=True, |
| 531 | fast_mode=True, |
| 532 | ) |
| 533 | self._mla_work_meta_data = torch.empty( |
| 534 | work_meta_data_size, dtype=work_meta_data_type, device=device |
| 535 | ) |
| 536 | self._mla_work_indptr = torch.empty( |
| 537 | work_indptr_size, dtype=work_indptr_type, device=device |
| 538 | ) |
| 539 | self._mla_work_info_set = torch.empty( |
| 540 | work_info_set_size, dtype=work_info_set_type, device=device |
| 541 | ) |
| 542 | self._mla_reduce_indptr = torch.empty( |
| 543 | reduce_indptr_size, dtype=reduce_indptr_type, device=device |
| 544 | ) |
| 545 | self._mla_reduce_final_map = torch.empty( |
| 546 | reduce_final_map_size, dtype=reduce_final_map_type, device=device |
| 547 | ) |
| 548 | self._mla_reduce_partial_map = torch.empty( |
| 549 | reduce_partial_map_size, |
| 550 | dtype=reduce_partial_map_type, |
| 551 | device=device, |
| 552 | ) |
| 553 | |
| 554 | self._prev_req_extent: int = 0 |
| 555 | self._prev_indices_extent: int = 0 |
| 556 | self._prev_metadata_key: tuple | None = None |
| 557 | |
| 558 | def _sparse_decode_max_split(self, max_seq_len: int) -> int: |
| 559 | """Cap ``max_split_per_batch`` for the aiter sparse-MLA decode reduce. |
| 560 | |
| 561 | The reduce only covers the selected tokens per row (``<= topk_tokens``), |
| 562 | so aiter's default (``-1`` => split across every CU) over-fragments it. |
| 563 | Mirror ``triton_mla.py``: aim for a minimum work per split, round to a |
| 564 | power of two, and cap by the CU count. Numerics are unchanged. |
| 565 | """ |
| 566 | effective_len = min(max_seq_len, self.topk_tokens) |
| 567 | min_work_per_split = 128 |
| 568 | ideal_splits = triton.next_power_of_2( |
| 569 | max(1, effective_len // min_work_per_split) |
| 570 | ) |
| 571 | return min(ideal_splits, self._num_compute_units) |
| 572 | |
| 573 | def build( |
| 574 | self, |
| 575 | common_prefix_len: int, |
| 576 | common_attn_metadata: CommonAttentionMetadata, |
| 577 | fast_build: bool = False, |
| 578 | ) -> ROCMAiterMLASparseMetadata: |
| 579 | num_tokens = common_attn_metadata.num_actual_tokens |
| 580 | (num_decodes, num_prefills, num_decode_tokens, _) = split_decodes_and_prefills( |
| 581 | common_attn_metadata, |
| 582 | decode_threshold=self.reorder_batch_threshold or 1, |
| 583 | ) |
| 584 | starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32) |
| 585 | seg_lengths = np.diff(starts) |
| 586 | req_id_per_token = np.repeat( |
| 587 | np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths |
| 588 | ) |
| 589 | # Only re-zero the shrink-tail. paged_kv_indptr is fully rewritten |
| 590 | # by the cumsum below. paged_kv_indices entries past new_indices_extent |
| 591 | # are never read (the attention kernel only touches the ranges |
| 592 | # defined by paged_kv_indptr). |
| 593 | new_req_extent = int(req_id_per_token.shape[0]) |
| 594 | new_indices_extent = num_tokens * self.topk_tokens |
| 595 | if self._prev_req_extent > new_req_extent: |
| 596 | self.req_id_per_token_buffer[new_req_extent : self._prev_req_extent].fill_( |
| 597 | 0 |
| 598 | ) |
| 599 | if self._prev_indices_extent > new_indices_extent: |
| 600 | self.paged_kv_indices[new_indices_extent : self._prev_indices_extent].fill_( |
| 601 | 0 |
| 602 | ) |
| 603 | self._prev_req_extent = new_req_extent |
| 604 | self._prev_indices_extent = new_indices_extent |
| 605 | self.req_id_per_token_buffer[:new_req_extent].copy_( |
| 606 | np_to_pinned_tensor(req_id_per_token), non_blocking=True |
| 607 | ) |
| 608 | query_lens = ( |
| 609 | common_attn_metadata.query_start_loc[1:] |
| 610 | - common_attn_metadata.query_start_loc[:-1] |
| 611 | ) |
| 612 | seq_lens = common_attn_metadata.seq_lens |
| 613 | sparse_seqlen = generate_sparse_seqlen_triton( |
| 614 | query_lens, |
| 615 | seq_lens, |
| 616 | common_attn_metadata.query_start_loc, |
| 617 | self.topk_tokens, |
| 618 | num_tokens, |
| 619 | common_attn_metadata.max_query_len, |
| 620 | ) |
| 621 | |
| 622 | torch.cumsum(sparse_seqlen, dim=0, out=self.paged_kv_indptr[1 : num_tokens + 1]) |
| 623 | self.paged_kv_indptr[num_tokens + 1 :].fill_(self.paged_kv_indptr[num_tokens]) |
| 624 | |
| 625 | req_id_per_token = self.req_id_per_token_buffer[:num_tokens] |
| 626 | qo_indptr = self.qo_indptr[: num_tokens + 1] |
| 627 | paged_kv_last_page_len = self.paged_kv_last_page_len[:num_tokens] |
| 628 | paged_kv_indptr = self.paged_kv_indptr[: num_tokens + 1] |
| 629 | paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens] |
| 630 | |
| 631 | # ----- Compute persistent MLA metadata ----- |
| 632 | # The AITER sparse decode kernel uses qseqlen=1 (each query token is |
| 633 | # treated as its own batch entry). Build its persistent work metadata |
| 634 | # only when AITER is selected and no layer needs the nonpersistent LSE |
| 635 | # path for attention sinks. The output is a deterministic function of |
| 636 | # the per-request query and context lengths (both clamped to |
| 637 | # topk_tokens, past which per-token KV length saturates) and num_heads; |
| 638 | # fingerprint those CPU-side and skip the launch when nothing changed. |
| 639 | head_size = self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim |
| 640 | use_triton_sparse = _use_rocm_sparse_triton( |
| 641 | kv_cache_dtype=self.kv_cache_dtype, |
| 642 | head_size=head_size, |
| 643 | kv_lora_rank=self.mla_dims.kv_lora_rank, |
| 644 | num_prefills=num_prefills, |
| 645 | num_decodes=num_decodes, |
| 646 | num_decode_tokens=num_decode_tokens, |
| 647 | max_query_len=common_attn_metadata.max_query_len, |
| 648 | ) |
| 649 | work_meta_data = None |
| 650 | work_indptr = None |
| 651 | work_info_set = None |
| 652 | reduce_indptr = None |
| 653 | reduce_final_map = None |
| 654 | reduce_partial_map = None |
| 655 | if self._use_persistent_metadata and not use_triton_sparse: |
| 656 | num_reqs = common_attn_metadata.num_reqs |
| 657 | with gpu_sync_allowed(): |
| 658 | seq_lens_cpu = common_attn_metadata.seq_lens[:num_reqs].cpu().numpy() |
| 659 | clamped_seq_lens = np.minimum( |
| 660 | seq_lens_cpu, |
| 661 | self.topk_tokens, |
| 662 | ) |
| 663 | clamped_context_lens = np.minimum( |
| 664 | seq_lens_cpu - seg_lengths, |
| 665 | self.topk_tokens, |
| 666 | ) |
| 667 | metadata_key = ( |
| 668 | num_tokens, |
| 669 | int(common_attn_metadata.max_query_len), |
| 670 | self._num_attention_heads, |
| 671 | clamped_seq_lens.tobytes(), |
| 672 | clamped_context_lens.tobytes(), |
| 673 | seg_lengths.tobytes(), |
| 674 | ) |
| 675 | if metadata_key != self._prev_metadata_key: |
| 676 | from aiter import get_mla_metadata_v1 |
| 677 | |
| 678 | max_split_per_batch = self._sparse_decode_max_split( |
| 679 | int(common_attn_metadata.max_seq_len) |
| 680 | ) |
| 681 | get_mla_metadata_v1( |
| 682 | qo_indptr, |
| 683 | paged_kv_indptr, |
| 684 | paged_kv_last_page_len, |
| 685 | self._num_attention_heads, |
| 686 | 1, |
| 687 | True, |
| 688 | self._mla_work_meta_data, |
| 689 | self._mla_work_info_set, |
| 690 | self._mla_work_indptr, |
| 691 | self._mla_reduce_indptr, |
| 692 | self._mla_reduce_final_map, |
| 693 | self._mla_reduce_partial_map, |
| 694 | page_size=1, |
| 695 | kv_granularity=16, |
| 696 | max_seqlen_qo=1, |
| 697 | uni_seqlen_qo=1, |
| 698 | fast_mode=True, |
| 699 | max_split_per_batch=max_split_per_batch, |
| 700 | ) |
| 701 | torch.cuda.current_stream(self.device).synchronize() |
| 702 | self._prev_metadata_key = metadata_key |
| 703 | work_meta_data = self._mla_work_meta_data |
| 704 | work_indptr = self._mla_work_indptr |
| 705 | work_info_set = self._mla_work_info_set |
| 706 | reduce_indptr = self._mla_reduce_indptr |
| 707 | reduce_final_map = self._mla_reduce_final_map |
| 708 | reduce_partial_map = self._mla_reduce_partial_map |
| 709 | |
| 710 | metadata = ROCMAiterMLASparseMetadata( |
| 711 | num_reqs=common_attn_metadata.num_reqs, |
| 712 | max_query_len=common_attn_metadata.max_query_len, |
| 713 | max_seq_len=common_attn_metadata.max_seq_len, |
| 714 | num_actual_tokens=common_attn_metadata.num_actual_tokens, |
| 715 | query_start_loc=common_attn_metadata.query_start_loc, |
| 716 | slot_mapping=common_attn_metadata.slot_mapping, |
| 717 | block_table=common_attn_metadata.block_table_tensor, |
| 718 | req_id_per_token=req_id_per_token, |
| 719 | block_size=self.kv_cache_spec.block_size, |
| 720 | attn_out_dtype=self.model_dtype, |
| 721 | topk_tokens=self.topk_tokens, |
| 722 | num_decodes=num_decodes, |
| 723 | num_prefills=num_prefills, |
| 724 | num_decode_tokens=num_decode_tokens, |
| 725 | qo_indptr=qo_indptr, |
| 726 | paged_kv_last_page_len=paged_kv_last_page_len, |
| 727 | paged_kv_indices=paged_kv_indices, |
| 728 | paged_kv_indptr=paged_kv_indptr, |
| 729 | work_meta_data=work_meta_data, |
| 730 | work_indptr=work_indptr, |
| 731 | work_info_set=work_info_set, |
| 732 | reduce_indptr=reduce_indptr, |
| 733 | reduce_final_map=reduce_final_map, |
| 734 | reduce_partial_map=reduce_partial_map, |
| 735 | ) |
| 736 | return metadata |
| 737 | |
| 738 | |
| 739 | # Take from |
| 740 | # https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla_prefill.py#L72 |
| 741 | def reference_mla_sparse_prefill( |
| 742 | q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, sm_scale: float, d_v: int |
| 743 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 744 | import math |
| 745 | |
| 746 | def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor: |
| 747 | return torch.logsumexp(a * math.log(2), dim=dim) * math.log2(math.e) |
| 748 | |
| 749 | skv = kv.shape[0] |
| 750 | sq = q.shape[0] |
| 751 | topk = indices.shape[-1] |
| 752 | dqk = q.shape[-1] |
| 753 | indices = indices[:, 0, :] # [s_q, topk] |
| 754 | invalid_indices_mask = (indices < 0) | (indices >= skv) |
| 755 | indices[invalid_indices_mask] = 0 |
| 756 | qs = q # [s_q, h_q, d_qk] |
| 757 | kvs = kv[:, 0, :][indices].view(sq, topk, dqk) # [s_q, topk, d_qk] |
| 758 | |
| 759 | attn_score = (qs @ kvs.transpose(1, 2)).float() # [s_q, h_q, topk] |
| 760 | attn_score.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf")) |
| 761 | attn_score *= sm_scale * math.log2(math.e) |
| 762 | lse = log2sumexp2(attn_score, dim=-1) # [s_q, h_q] |
| 763 | attn_score = torch.exp2(attn_score - lse.unsqueeze(-1)) # [s_q, h_q, topk] |
| 764 | result = attn_score.to(q.dtype) @ kvs[:, :, :d_v] |
| 765 | return (result, lse) |
| 766 | |
| 767 | |
| 768 | class ROCMAiterMLASparseImpl( |
| 769 | MLAAttentionImpl[ROCMAiterMLASparseMetadata], SharedTopkIndicesBuffer |
| 770 | ): |
| 771 | is_sparse = True |
| 772 | supports_dense_mha_prefill = False |
| 773 | supports_dcp = False |
| 774 | use_aiter_sparse_mla = False |
| 775 | |
| 776 | def __init__( |
| 777 | self, |
| 778 | num_heads: int, |
| 779 | head_size: int, |
| 780 | scale: float, |
| 781 | num_kv_heads: int, |
| 782 | alibi_slopes: list[float] | None, |
| 783 | sliding_window: int | None, |
| 784 | kv_cache_dtype: str, |
| 785 | logits_soft_cap: float | None, |
| 786 | attn_type: str, |
| 787 | kv_sharing_target_layer_name: str | None, |
| 788 | # MLA Specific Arguments |
| 789 | topk_indices_buffer: torch.Tensor | None = None, |
| 790 | indexer: "Indexer | None" = None, |
| 791 | **mla_args, |
| 792 | ) -> None: |
| 793 | AiterMLAHelper.check_num_heads_validity(num_heads) |
| 794 | |
| 795 | self.num_heads = num_heads |
| 796 | self.head_size = head_size |
| 797 | self.scale = float(scale) |
| 798 | self.num_kv_heads = num_kv_heads |
| 799 | sinks = mla_args.pop("sinks", None) |
| 800 | if sinks is not None: |
| 801 | if sinks.dtype != torch.float32: |
| 802 | raise ValueError( |
| 803 | f"ROCm AITER MLA sinks must be float32, got {sinks.dtype}" |
| 804 | ) |
| 805 | if sinks.ndim != 1 or sinks.numel() != num_heads: |
| 806 | raise ValueError( |
| 807 | "ROCm AITER MLA sinks must have shape " |
| 808 | f"({num_heads},), got {tuple(sinks.shape)}" |
| 809 | ) |
| 810 | if not sinks.is_contiguous(): |
| 811 | raise ValueError("ROCm AITER MLA sinks must be contiguous") |
| 812 | self.sinks: torch.Tensor | None = sinks |
| 813 | self.kv_cache_dtype = kv_cache_dtype |
| 814 | self.kv_lora_rank: int = mla_args["kv_lora_rank"] |
| 815 | self.softmax_scale = scale |
| 816 | self.init_topk_indices_buffer(indexer, topk_indices_buffer) |
| 817 | |
| 818 | vllm_config = get_current_vllm_config() |
| 819 | max_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 820 | q_concat_shape = (max_tokens, num_heads, head_size) |
| 821 | (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( |
| 822 | (q_concat_shape, vllm_config.model_config.dtype), |
| 823 | ) |
| 824 | self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] |
| 825 | |
| 826 | if rocm_aiter_ops.is_triton_sparse_mla_enabled(): |
| 827 | reason = self._aiter_sparse_mla_unsupported_reason(vllm_config) |
| 828 | if reason is None: |
| 829 | self.use_aiter_sparse_mla = True |
| 830 | else: |
| 831 | logger.warning_once( |
| 832 | "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; " |
| 833 | "using the default sparse MLA path instead.", |
| 834 | reason, |
| 835 | ) |
| 836 | |
| 837 | def _aiter_sparse_mla_unsupported_reason( |
| 838 | self, vllm_config: VllmConfig |
| 839 | ) -> str | None: |
| 840 | model_dtype = vllm_config.model_config.dtype |
| 841 | if model_dtype != torch.bfloat16: |
| 842 | return f"the model dtype is {model_dtype}, not bfloat16" |
| 843 | cache_dtype = self.kv_cache_dtype |
| 844 | if cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"): |
| 845 | return f"kv-cache dtype {cache_dtype!r} is neither bfloat16 nor e4m3 fp8" |
| 846 | # Context parallelism merges each rank's partials with a per-token LSE. |
| 847 | if self.dcp_world_size > 1 or self.pcp_world_size > 1: |
| 848 | return "context parallelism is not supported" |
| 849 | return None |
| 850 | |
| 851 | def record_logical_topk_ready(self) -> None: |
| 852 | # This impl shares the top-k indices buffer via SharedTopkIndicesBuffer |
| 853 | # but does not participate in sparse-MLA index groups. |
| 854 | pass |
| 855 | |
| 856 | def _forward_mla( |
| 857 | self, |
| 858 | layer: AttentionLayer, |
| 859 | q: torch.Tensor, # [sq, heads, d_qk] |
| 860 | kv_c_and_k_pe_cache: torch.Tensor, # [blocks, heads, d_qk] |
| 861 | attn_metadata: ROCMAiterMLASparseMetadata, |
| 862 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 863 | num_tokens = q.shape[0] |
| 864 | base_mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads) |
| 865 | mla_num_heads = base_mla_num_heads |
| 866 | need_lse = self.sinks is not None |
| 867 | from vllm.platforms.rocm import on_gfx942 |
| 868 | |
| 869 | # Keep sink attention available for dtypes/head shapes without an |
| 870 | # AITER return-LSE kernel. Sink layers never need persistent metadata. |
| 871 | triton_sink_fallback = ( |
| 872 | need_lse |
| 873 | and q.dtype == kv_c_and_k_pe_cache.dtype |
| 874 | and ( |
| 875 | q.dtype == torch.float16 |
| 876 | or ( |
| 877 | q.dtype == torch.bfloat16 |
| 878 | and ( |
| 879 | mla_num_heads > 128 |
| 880 | or (on_gfx942() and 32 < mla_num_heads <= 64) |
| 881 | ) |
| 882 | ) |
| 883 | ) |
| 884 | ) |
| 885 | |
| 886 | if triton_sink_fallback or _use_rocm_sparse_triton( |
| 887 | kv_cache_dtype=self.kv_cache_dtype, |
| 888 | head_size=q.shape[-1], |
| 889 | kv_lora_rank=self.kv_lora_rank, |
| 890 | num_prefills=attn_metadata.num_prefills, |
| 891 | num_decodes=attn_metadata.num_decodes, |
| 892 | num_decode_tokens=attn_metadata.num_decode_tokens, |
| 893 | max_query_len=attn_metadata.max_query_len, |
| 894 | ): |
| 895 | output = torch.empty( |
| 896 | [num_tokens, q.shape[1], self.kv_lora_rank], |
| 897 | dtype=attn_metadata.attn_out_dtype, |
| 898 | device=q.device, |
| 899 | ) |
| 900 | triton_sinks = None |
| 901 | if self.sinks is not None: |
| 902 | triton_sinks = AiterMLAHelper.get_mla_padded_q( |
| 903 | self.num_heads, |
| 904 | self.sinks.reshape(1, self.num_heads, 1), |
| 905 | q.shape[1], |
| 906 | ).reshape(-1) |
| 907 | rocm_sparse_attn_prefill( |
| 908 | q=q, |
| 909 | kv=kv_c_and_k_pe_cache.view(-1, 1, q.shape[-1]), |
| 910 | indices=None, |
| 911 | topk_length=None, |
| 912 | scale=self.scale, |
| 913 | head_dim=q.shape[-1], |
| 914 | nope_head_dim=self.kv_lora_rank, |
| 915 | rope_head_dim=q.shape[-1] - self.kv_lora_rank, |
| 916 | attn_sink=triton_sinks, |
| 917 | output=output, |
| 918 | ragged_indices=attn_metadata.paged_kv_indices, |
| 919 | ragged_indptr=attn_metadata.paged_kv_indptr, |
| 920 | ) |
| 921 | output = AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output) |
| 922 | return output, None |
| 923 | |
| 924 | # AITER's nonpersistent return-LSE dispatch has discrete head kernels. |
| 925 | supported_head_buckets: tuple[int, ...] | None = None |
| 926 | head_dtype_name = "" |
| 927 | if need_lse: |
| 928 | if ( |
| 929 | q.dtype == torch.bfloat16 |
| 930 | and kv_c_and_k_pe_cache.dtype == torch.bfloat16 |
| 931 | ): |
| 932 | supported_head_buckets = (16, 32, 64, 128) |
| 933 | head_dtype_name = "BF16" |
| 934 | elif ( |
| 935 | q.dtype == current_platform.fp8_dtype() |
| 936 | and kv_c_and_k_pe_cache.dtype == current_platform.fp8_dtype() |
| 937 | ): |
| 938 | supported_head_buckets = (16, 128) |
| 939 | head_dtype_name = "FP8" |
| 940 | else: |
| 941 | raise ValueError( |
| 942 | "ROCm AITER MLA attention sinks require query and KV to " |
| 943 | "both use BF16 or both use FP8, got " |
| 944 | f"query={q.dtype}, KV={kv_c_and_k_pe_cache.dtype}" |
| 945 | ) |
| 946 | |
| 947 | if supported_head_buckets is not None: |
| 948 | from vllm.platforms.rocm import on_mi3xx |
| 949 | |
| 950 | if on_mi3xx(): |
| 951 | supported_heads = next( |
| 952 | ( |
| 953 | heads |
| 954 | for heads in supported_head_buckets |
| 955 | if heads >= mla_num_heads |
| 956 | ), |
| 957 | None, |
| 958 | ) |
| 959 | if supported_heads is None: |
| 960 | raise ValueError( |
| 961 | "ROCm AITER MLA attention sinks support at most 128 " |
| 962 | f"padded local {head_dtype_name} heads; increase " |
| 963 | "tensor_parallel_size" |
| 964 | ) |
| 965 | if supported_heads != mla_num_heads: |
| 966 | q = AiterMLAHelper.get_mla_padded_q( |
| 967 | mla_num_heads, q, supported_heads |
| 968 | ) |
| 969 | mla_num_heads = supported_heads |
| 970 | output = torch.empty( |
| 971 | [num_tokens, mla_num_heads, self.kv_lora_rank], |
| 972 | dtype=attn_metadata.attn_out_dtype, |
| 973 | device=q.device, |
| 974 | ) |
| 975 | |
| 976 | if need_lse: |
| 977 | # gfx942 has no persistent MLA code object that writes final LSE. |
| 978 | # The split-KV path consumes the same ragged indices and is exact. |
| 979 | lse = rocm_aiter_ops.mla_decode_fwd_lse( |
| 980 | q, |
| 981 | kv_c_and_k_pe_cache, |
| 982 | output, |
| 983 | self.scale, |
| 984 | attn_metadata.qo_indptr, |
| 985 | 1, |
| 986 | attn_metadata.paged_kv_indptr, |
| 987 | attn_metadata.paged_kv_indices, |
| 988 | attn_metadata.paged_kv_last_page_len, |
| 989 | q_scale=layer._q_scale, |
| 990 | kv_scale=layer._k_scale, |
| 991 | ) |
| 992 | else: |
| 993 | # Preserve the persistent work-stealing fast path for models that |
| 994 | # do not use sinks. |
| 995 | mla_kwargs: dict = dict( |
| 996 | q_scale=layer._q_scale, |
| 997 | kv_scale=layer._k_scale, |
| 998 | ) |
| 999 | if attn_metadata.work_meta_data is not None: |
| 1000 | mla_kwargs.update( |
| 1001 | work_meta_data=attn_metadata.work_meta_data, |
| 1002 | work_indptr=attn_metadata.work_indptr, |
| 1003 | work_info_set=attn_metadata.work_info_set, |
| 1004 | reduce_indptr=attn_metadata.reduce_indptr, |
| 1005 | reduce_final_map=attn_metadata.reduce_final_map, |
| 1006 | reduce_partial_map=attn_metadata.reduce_partial_map, |
| 1007 | ) |
| 1008 | |
| 1009 | rocm_aiter_ops.mla_decode_fwd( |
| 1010 | q, |
| 1011 | kv_c_and_k_pe_cache, |
| 1012 | output, |
| 1013 | self.scale, |
| 1014 | attn_metadata.qo_indptr, |
| 1015 | 1, |
| 1016 | attn_metadata.paged_kv_indptr, |
| 1017 | attn_metadata.paged_kv_indices, |
| 1018 | attn_metadata.paged_kv_last_page_len, |
| 1019 | **mla_kwargs, |
| 1020 | ) |
| 1021 | lse = None |
| 1022 | |
| 1023 | if mla_num_heads != base_mla_num_heads: |
| 1024 | if mla_num_heads % base_mla_num_heads == 0: |
| 1025 | head_stride = mla_num_heads // base_mla_num_heads |
| 1026 | output = output[:, ::head_stride] |
| 1027 | if lse is not None: |
| 1028 | lse = lse[:, ::head_stride] |
| 1029 | else: |
| 1030 | output = output[:, :base_mla_num_heads] |
| 1031 | if lse is not None: |
| 1032 | lse = lse[:, :base_mla_num_heads] |
| 1033 | |
| 1034 | output = AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output) |
| 1035 | if lse is not None: |
| 1036 | lse = AiterMLAHelper.get_mla_unpadded_o( |
| 1037 | self.num_heads, lse.unsqueeze(-1) |
| 1038 | ).squeeze(-1) |
| 1039 | |
| 1040 | if self.sinks is not None: |
| 1041 | assert lse is not None |
| 1042 | # Empty ragged rows have only sink mass and no value contribution. |
| 1043 | # AITER can return NaN output/LSE for those rows; do not multiply it |
| 1044 | # by a zero normalization factor and propagate the NaN. |
| 1045 | has_keys = ( |
| 1046 | attn_metadata.paged_kv_indptr[1:] > attn_metadata.paged_kv_indptr[:-1] |
| 1047 | ).unsqueeze(-1) |
| 1048 | lse = torch.where(has_keys, lse, float("-inf")) |
| 1049 | sink_lse = torch.logaddexp(lse, self.sinks) |
| 1050 | sink_scale = torch.exp(lse - sink_lse) |
| 1051 | output = torch.where( |
| 1052 | has_keys.unsqueeze(-1), |
| 1053 | output.float() * sink_scale.unsqueeze(-1), |
| 1054 | 0.0, |
| 1055 | ).to(output.dtype) |
| 1056 | lse = sink_lse |
| 1057 | |
| 1058 | return output, lse |
| 1059 | |
| 1060 | def _forward_mla_aiter( |
| 1061 | self, |
| 1062 | layer: AttentionLayer, |
| 1063 | q: torch.Tensor, # [sq, heads, d_qk], not head-padded |
| 1064 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1065 | attn_metadata: ROCMAiterMLASparseMetadata, |
| 1066 | ) -> torch.Tensor: |
| 1067 | """_forward_mla on aiter's Triton sparse MLA kernel. |
| 1068 | |
| 1069 | Reads the same index stream, but needs no q head padding and no |
| 1070 | persistent metadata, and its launch depends on shapes only, so it is |
| 1071 | CUDA-graph capturable. |
| 1072 | """ |
| 1073 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 1074 | output = torch.empty( |
| 1075 | [q.shape[0], self.num_heads, self.kv_lora_rank], |
| 1076 | dtype=attn_metadata.attn_out_dtype, |
| 1077 | device=q.device, |
| 1078 | ) |
| 1079 | rocm_aiter_ops.triton_sparse_mla_fwd( |
| 1080 | q[:num_actual_tokens], |
| 1081 | kv_c_and_k_pe_cache.view(-1, 1, 1, q.shape[-1]), |
| 1082 | output[:num_actual_tokens], |
| 1083 | self.scale, |
| 1084 | attn_metadata.paged_kv_indptr, |
| 1085 | attn_metadata.paged_kv_indices, |
| 1086 | kv_lora_rank=self.kv_lora_rank, |
| 1087 | qk_rope_head_dim=self.qk_rope_head_dim, |
| 1088 | q_scale=layer._q_scale, |
| 1089 | kv_scale=layer._k_scale, |
| 1090 | attn_sink=self.sinks, |
| 1091 | # triton_convert_req_index_to_global_index writes 0, never -1, |
| 1092 | # for an invalid top-k entry, so no slot in the stream is negative. |
| 1093 | has_invalid=False, |
| 1094 | ) |
| 1095 | return output |
| 1096 | |
| 1097 | def forward_mqa( |
| 1098 | self, |
| 1099 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 1100 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1101 | attn_metadata: ROCMAiterMLASparseMetadata, |
| 1102 | layer: AttentionLayer, |
| 1103 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 1104 | # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use |
| 1105 | # MQA 576/512 approach for both prefill and decode |
| 1106 | |
| 1107 | fp8_attention = self.kv_cache_dtype.startswith("fp8") |
| 1108 | if isinstance(q, tuple): |
| 1109 | ql_nope, q_pe = q |
| 1110 | if fp8_attention: |
| 1111 | q = layer._decode_concat_quant_fp8_op( # type: ignore[attr-defined] |
| 1112 | ql_nope, q_pe, layer._q_scale |
| 1113 | ) |
| 1114 | else: |
| 1115 | q = self.q_concat_buffer[: ql_nope.shape[0]] |
| 1116 | if q_pe.shape[-1] == 0: |
| 1117 | q.copy_(ql_nope) |
| 1118 | elif q.dtype == torch.float16: |
| 1119 | torch.cat((ql_nope, q_pe), dim=-1, out=q) |
| 1120 | else: |
| 1121 | ops.concat_mla_q(ql_nope, q_pe, q) |
| 1122 | |
| 1123 | num_actual_toks = attn_metadata.num_actual_tokens |
| 1124 | |
| 1125 | # Get topk indices |
| 1126 | assert self.topk_indices_buffer is not None |
| 1127 | topk_indices = fit_kpool_indices_to_aiter( |
| 1128 | self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens |
| 1129 | ) |
| 1130 | |
| 1131 | triton_convert_req_index_to_global_index( |
| 1132 | attn_metadata.req_id_per_token, |
| 1133 | attn_metadata.block_table, |
| 1134 | topk_indices, |
| 1135 | attn_metadata.paged_kv_indptr, |
| 1136 | attn_metadata.paged_kv_indices, |
| 1137 | BLOCK_SIZE=attn_metadata.block_size, |
| 1138 | NUM_TOPK_TOKENS=attn_metadata.topk_tokens, |
| 1139 | ) |
| 1140 | |
| 1141 | # write the latent and rope to kv cache |
| 1142 | if fp8_attention: |
| 1143 | kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(current_platform.fp8_dtype()) |
| 1144 | if q.dtype != current_platform.fp8_dtype(): |
| 1145 | original_q_shape = q.shape |
| 1146 | q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale) |
| 1147 | q = q.view(original_q_shape) |
| 1148 | if self.use_aiter_sparse_mla: |
| 1149 | output = self._forward_mla_aiter( |
| 1150 | layer, q, kv_c_and_k_pe_cache, attn_metadata |
| 1151 | ) |
| 1152 | return output, None |
| 1153 | mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q) |
| 1154 | return self._forward_mla( |
| 1155 | layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata |
| 1156 | ) |