| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import functools |
| 5 | from collections.abc import Callable |
| 6 | from dataclasses import dataclass |
| 7 | from typing import Any, cast |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.distributed import ( |
| 12 | get_tensor_model_parallel_world_size, |
| 13 | tensor_model_parallel_all_reduce, |
| 14 | ) |
| 15 | from vllm.forward_context import get_forward_context |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation |
| 18 | from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic |
| 19 | from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import ( |
| 20 | RocmSparseAttnIndexer, |
| 21 | RocmSparseMQAIndexer, |
| 22 | ) |
| 23 | from vllm.models.deepseek_v41.attention import ( |
| 24 | DeepseekV4Attention, |
| 25 | DeepseekV4Indexer, |
| 26 | DeepseekV4IndexerCache, |
| 27 | _replace_layer_index, |
| 28 | ) |
| 29 | from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache |
| 30 | from vllm.models.deepseek_v41.sparse_mla import ( |
| 31 | DeepseekV4FlashMLAMetadata, |
| 32 | DeepseekV4SparseMLABackend, |
| 33 | DeepseekV4SparseMLAMetadataBuilder, |
| 34 | DeepseekV41SparseSWAMetadataBuilder, |
| 35 | ) |
| 36 | from vllm.platforms import current_platform |
| 37 | from vllm.platforms.rocm import _ON_GFX950 |
| 38 | from vllm.triton_utils import tl, triton |
| 39 | from vllm.utils.multi_stream_utils import execute_in_parallel |
| 40 | from vllm.v1.attention.backend import ( |
| 41 | CommonAttentionMetadata, |
| 42 | MultipleOf, |
| 43 | ) |
| 44 | from vllm.v1.attention.backends.mla.sparse_swa import ( |
| 45 | DeepseekSparseSWABackend, |
| 46 | DeepseekSparseSWAMetadata, |
| 47 | ) |
| 48 | from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( |
| 49 | build_prefill_topk_ragged_indices, |
| 50 | build_ragged_indices_from_dense, |
| 51 | rocm_inv_rope_einsum, |
| 52 | rocm_inverse_rope_mxfp8_rows, |
| 53 | rocm_inverse_rope_rows_, |
| 54 | rocm_mxfp8_wo_a_bmm, |
| 55 | rocm_sparse_attn_decode, |
| 56 | rocm_sparse_attn_prefill, |
| 57 | ) |
| 58 | from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import ( |
| 59 | rocm_paged_mxfp4_cache_layout, |
| 60 | ) |
| 61 | from vllm.v1.worker.workspace import current_workspace_manager |
| 62 | |
| 63 | logger = init_logger(__name__) |
| 64 | |
| 65 | |
| 66 | def _trust_dsv4_extra_cache_nan_free( |
| 67 | kv_cache_dtype: str, |
| 68 | has_kv_transfer: bool, |
| 69 | has_extra_cache: bool, |
| 70 | ) -> bool: |
| 71 | return ( |
| 72 | _ON_GFX950 |
| 73 | and kv_cache_dtype == "fp8_ds_mla" |
| 74 | and not has_kv_transfer |
| 75 | and has_extra_cache |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool: |
| 80 | from vllm._aiter_ops import rocm_aiter_ops |
| 81 | |
| 82 | if not rocm_aiter_ops.is_triton_sparse_mla_enabled(): |
| 83 | return False |
| 84 | # The aiter kernel does not scrub NaNs, so like the trust check above it |
| 85 | # needs a cache only the canonical writer has touched. |
| 86 | if has_kv_transfer: |
| 87 | logger.warning_once( |
| 88 | "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV " |
| 89 | "connector; DeepSeek V4 keeps the Triton sparse attention." |
| 90 | ) |
| 91 | return False |
| 92 | return True |
| 93 | |
| 94 | |
| 95 | def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: |
| 96 | lengths = lengths.to(dtype=torch.int32).contiguous() |
| 97 | indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) |
| 98 | torch.cumsum(lengths, dim=0, out=indptr[1:]) |
| 99 | return indptr |
| 100 | |
| 101 | |
| 102 | def apply_pre_quantized_block_scaled_mm( |
| 103 | linear: torch.nn.Module, |
| 104 | x_fp8: torch.Tensor, |
| 105 | x_scale: torch.Tensor, |
| 106 | ) -> torch.Tensor: |
| 107 | """Block-scaled fp8 GEMM on pre-quantized activations. |
| 108 | |
| 109 | The fused q/kv norm kernel writes fp8 qr + per-1x128 scales; this |
| 110 | drives the linear's block-scaled GEMM directly with them, bypassing |
| 111 | apply_weights which would re-quantize the fp8 input. Only valid for |
| 112 | the wq_b-style column/replicated linears: their output is the local |
| 113 | TP shard, so no all-reduce is needed. |
| 114 | """ |
| 115 | from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import ( |
| 116 | FP8BlockParams, |
| 117 | ) |
| 118 | |
| 119 | params = FP8BlockParams.from_layer(linear) |
| 120 | weight_scale = ( |
| 121 | params.weight_scale |
| 122 | if params.weight_scale_inv is None |
| 123 | else params.weight_scale_inv |
| 124 | ) |
| 125 | kernel = linear.quant_method.fp8_linear |
| 126 | out = kernel.apply_block_scaled_mm( |
| 127 | A=x_fp8, B=params.weight, As=x_scale, Bs=weight_scale |
| 128 | ) |
| 129 | return out.to(dtype=kernel.config.out_dtype) |
| 130 | |
| 131 | |
| 132 | # ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes |
| 133 | # do not touch the shared DeepSeek V4 cache utilities. |
| 134 | _SPARSE_PREFILL_TOPK_ALIGNMENT = 128 |
| 135 | |
| 136 | |
| 137 | @triton.jit |
| 138 | def _combine_topk_swa_indices_kernel( |
| 139 | combined_indices_ptr, |
| 140 | combined_indices_stride, |
| 141 | combined_lens_ptr, |
| 142 | topk_indices_ptr, |
| 143 | topk_indices_stride, |
| 144 | query_start_loc_ptr, |
| 145 | seq_lens_ptr, |
| 146 | gather_lens_ptr, |
| 147 | M, |
| 148 | N, |
| 149 | TOP_K: tl.constexpr, |
| 150 | COMPRESS_RATIO: tl.constexpr, |
| 151 | WINDOW_SIZE: tl.constexpr, |
| 152 | TOPK_WIDTH: tl.constexpr, |
| 153 | PADDED_TOP_K: tl.constexpr, |
| 154 | ): |
| 155 | batch_idx = tl.program_id(0) |
| 156 | worker_id = tl.program_id(1) |
| 157 | num_workers = tl.num_programs(1) |
| 158 | |
| 159 | base = tl.load(query_start_loc_ptr) |
| 160 | query_start = tl.load(query_start_loc_ptr + batch_idx) - base |
| 161 | query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base |
| 162 | query_len = query_end - query_start |
| 163 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 164 | gather_len = tl.load(gather_lens_ptr + batch_idx) |
| 165 | start_pos = seq_len - query_len |
| 166 | gather_start = seq_len - gather_len |
| 167 | |
| 168 | for token_idx in range(query_start + worker_id, query_end, num_workers): |
| 169 | token_idx_in_query = token_idx - query_start |
| 170 | pos = start_pos + token_idx_in_query |
| 171 | topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) |
| 172 | swa_len = tl.minimum(pos + 1, WINDOW_SIZE) |
| 173 | |
| 174 | topk_offset = tl.arange(0, PADDED_TOP_K) |
| 175 | topk_mask = topk_offset < topk_len |
| 176 | safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0) |
| 177 | topk_indices = tl.load( |
| 178 | topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset, |
| 179 | mask=topk_mask, |
| 180 | other=-1, |
| 181 | ) |
| 182 | valid_topk = (topk_indices >= 0) & (topk_indices < N) |
| 183 | topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1) |
| 184 | tl.store( |
| 185 | combined_indices_ptr + token_idx * combined_indices_stride + topk_offset, |
| 186 | topk_indices, |
| 187 | mask=topk_mask, |
| 188 | ) |
| 189 | |
| 190 | swa_offset = tl.arange(0, WINDOW_SIZE) |
| 191 | tl.store( |
| 192 | combined_indices_ptr |
| 193 | + token_idx * combined_indices_stride |
| 194 | + topk_len |
| 195 | + swa_offset, |
| 196 | M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, |
| 197 | mask=swa_offset < swa_len, |
| 198 | ) |
| 199 | |
| 200 | tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) |
| 201 | |
| 202 | |
| 203 | def combine_topk_swa_indices( |
| 204 | topk_indices: torch.Tensor, |
| 205 | query_start_loc: torch.Tensor, |
| 206 | seq_lens: torch.Tensor, |
| 207 | gather_lens: torch.Tensor, |
| 208 | window_size: int, |
| 209 | compress_ratio: int, |
| 210 | topk: int, |
| 211 | M: int, |
| 212 | N: int, |
| 213 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 214 | """Combine compressed-attention and sliding-window indices with Torch. |
| 215 | |
| 216 | The Triton implementation inherited from DeepSeek V4 launches a |
| 217 | two-dimensional grid with 128 workers per request. On gfx950 it can issue |
| 218 | an out-of-bounds access for V4.1's mixed prefill metadata (including the |
| 219 | synthetic mixed-token warmup). This path is prefill-only and the tensors |
| 220 | are small, so use ordinary Torch indexing until a gfx950-safe fused kernel |
| 221 | is available. |
| 222 | """ |
| 223 | topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() |
| 224 | num_tokens = topk_indices.shape[0] |
| 225 | combined_topk = ( |
| 226 | (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) |
| 227 | // _SPARSE_PREFILL_TOPK_ALIGNMENT |
| 228 | * _SPARSE_PREFILL_TOPK_ALIGNMENT |
| 229 | ) |
| 230 | combined_indices = torch.full( |
| 231 | (num_tokens, combined_topk), |
| 232 | fill_value=-1, |
| 233 | dtype=torch.int32, |
| 234 | device=topk_indices.device, |
| 235 | ) |
| 236 | combined_lens = torch.empty( |
| 237 | num_tokens, dtype=torch.int32, device=topk_indices.device |
| 238 | ) |
| 239 | |
| 240 | # query_start_loc may have a non-zero base for a narrowed mixed batch. |
| 241 | query_lens = query_start_loc[1:] - query_start_loc[:-1] |
| 242 | req_ids = torch.repeat_interleave( |
| 243 | torch.arange(seq_lens.shape[0], device=seq_lens.device), query_lens |
| 244 | ) |
| 245 | query_starts = query_start_loc[:-1] - query_start_loc[0] |
| 246 | token_offsets = torch.arange(num_tokens, device=seq_lens.device) - ( |
| 247 | torch.repeat_interleave(query_starts, query_lens) |
| 248 | ) |
| 249 | positions = seq_lens[req_ids] - query_lens[req_ids] + token_offsets |
| 250 | |
| 251 | logical_topk_width = min(topk, topk_indices.shape[1]) |
| 252 | topk_lens = torch.minimum( |
| 253 | (positions + 1) // compress_ratio, |
| 254 | torch.full_like(positions, logical_topk_width), |
| 255 | ).clamp_min(0) |
| 256 | topk_offsets = torch.arange(logical_topk_width, device=seq_lens.device) |
| 257 | topk_mask = topk_offsets[None, :] < topk_lens[:, None] |
| 258 | topk_values = topk_indices[:, :logical_topk_width].to(torch.int32) |
| 259 | topk_valid = topk_mask & (topk_values >= 0) & (topk_values < N) |
| 260 | combined_indices[:, :logical_topk_width] = torch.where( |
| 261 | topk_valid, |
| 262 | topk_values + (M * req_ids).to(torch.int32)[:, None], |
| 263 | -1, |
| 264 | ) |
| 265 | |
| 266 | swa_lens = torch.minimum( |
| 267 | positions + 1, torch.full_like(positions, window_size) |
| 268 | ).clamp_min(0) |
| 269 | swa_offsets = torch.arange(window_size, device=seq_lens.device) |
| 270 | swa_mask = swa_offsets[None, :] < swa_lens[:, None] |
| 271 | swa_columns = topk_lens[:, None] + swa_offsets[None, :] |
| 272 | gather_starts = seq_lens - gather_lens |
| 273 | swa_values = ( |
| 274 | M * req_ids[:, None] |
| 275 | + N |
| 276 | + swa_offsets[None, :] |
| 277 | + positions[:, None] |
| 278 | - swa_lens[:, None] |
| 279 | + 1 |
| 280 | - gather_starts[req_ids, None] |
| 281 | ).to(torch.int32) |
| 282 | rows = torch.arange(num_tokens, device=seq_lens.device)[:, None].expand_as( |
| 283 | swa_columns |
| 284 | ) |
| 285 | flat_dst = rows[swa_mask] * combined_topk + swa_columns[swa_mask] |
| 286 | combined_indices.view(-1).index_copy_(0, flat_dst, swa_values[swa_mask]) |
| 287 | combined_lens.copy_((topk_lens + swa_lens).to(torch.int32)) |
| 288 | return combined_indices, combined_lens |
| 289 | |
| 290 | |
| 291 | @triton.jit |
| 292 | def _compute_topk_lens_kernel( |
| 293 | topk_lens_ptr, |
| 294 | topk_indices_ptr, |
| 295 | topk_indices_stride, |
| 296 | topk, |
| 297 | is_valid_token_ptr, |
| 298 | TRITON_BLOCK_SIZE: tl.constexpr, |
| 299 | ): |
| 300 | token_idx = tl.program_id(0) |
| 301 | is_valid_token = tl.load(is_valid_token_ptr + token_idx) |
| 302 | |
| 303 | count = tl.zeros((), dtype=tl.int32) |
| 304 | for i in range(0, topk, TRITON_BLOCK_SIZE): |
| 305 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 306 | mask = offset < topk |
| 307 | local_idx = tl.load( |
| 308 | topk_indices_ptr + token_idx * topk_indices_stride + offset, |
| 309 | mask=mask, |
| 310 | other=-1, |
| 311 | ) |
| 312 | count += tl.sum((local_idx >= 0).to(tl.int32), axis=0) |
| 313 | |
| 314 | tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) |
| 315 | |
| 316 | |
| 317 | @triton.jit |
| 318 | def _pack_global_topk_ragged_kernel( |
| 319 | global_topk_ragged_ptr, |
| 320 | topk_indptr_ptr, |
| 321 | topk_indices_ptr, |
| 322 | topk_indices_stride, |
| 323 | token_to_req_indices_ptr, |
| 324 | block_table_ptr, |
| 325 | block_table_stride, |
| 326 | block_size, |
| 327 | topk, |
| 328 | BLOCK_SIZE: tl.constexpr, |
| 329 | ): |
| 330 | token_idx = tl.program_id(0) |
| 331 | block_idx = tl.program_id(1) |
| 332 | offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 333 | |
| 334 | out_start = tl.load(topk_indptr_ptr + token_idx) |
| 335 | out_end = tl.load(topk_indptr_ptr + token_idx + 1) |
| 336 | out_len = out_end - out_start |
| 337 | if block_idx * BLOCK_SIZE >= out_len: |
| 338 | return |
| 339 | |
| 340 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 341 | mask = (offset < out_len) & (offset < topk) |
| 342 | local_idx = tl.load( |
| 343 | topk_indices_ptr + token_idx * topk_indices_stride + offset, |
| 344 | mask=mask, |
| 345 | other=-1, |
| 346 | ) |
| 347 | valid = mask & (local_idx >= 0) |
| 348 | block_indices = local_idx // block_size |
| 349 | block_numbers = tl.load( |
| 350 | block_table_ptr + req_idx * block_table_stride + block_indices, |
| 351 | mask=valid, |
| 352 | other=0, |
| 353 | ) |
| 354 | block_offsets = local_idx % block_size |
| 355 | slot_ids = tl.where(valid, block_numbers * block_size + block_offsets, -1) |
| 356 | tl.store(global_topk_ragged_ptr + out_start + offset, slot_ids, mask=mask) |
| 357 | |
| 358 | |
| 359 | def compute_global_topk_ragged_indices_and_indptr( |
| 360 | topk_indices: torch.Tensor, |
| 361 | token_to_req_indices: torch.Tensor, |
| 362 | block_table: torch.Tensor, |
| 363 | block_size: int, |
| 364 | is_valid_token: torch.Tensor, |
| 365 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 366 | topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() |
| 367 | num_tokens = topk_indices.shape[0] |
| 368 | topk = topk_indices.shape[1] |
| 369 | |
| 370 | topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) |
| 371 | _compute_topk_lens_kernel[(num_tokens,)]( |
| 372 | topk_lens, |
| 373 | topk_indices, |
| 374 | topk_indices.stride(0), |
| 375 | topk, |
| 376 | is_valid_token, |
| 377 | TRITON_BLOCK_SIZE=1024, |
| 378 | ) |
| 379 | |
| 380 | topk_indptr = _build_indptr_from_lengths(topk_lens) |
| 381 | global_topk_ragged = torch.empty( |
| 382 | num_tokens * topk, |
| 383 | dtype=torch.int32, |
| 384 | device=topk_indices.device, |
| 385 | ) |
| 386 | if global_topk_ragged.numel() > 0: |
| 387 | block = 128 |
| 388 | _pack_global_topk_ragged_kernel[(num_tokens, triton.cdiv(topk, block))]( |
| 389 | global_topk_ragged, |
| 390 | topk_indptr, |
| 391 | topk_indices, |
| 392 | topk_indices.stride(0), |
| 393 | token_to_req_indices, |
| 394 | block_table, |
| 395 | block_table.stride(0), |
| 396 | block_size, |
| 397 | topk, |
| 398 | BLOCK_SIZE=block, |
| 399 | ) |
| 400 | return global_topk_ragged, topk_indptr, topk_lens |
| 401 | |
| 402 | |
| 403 | def _copy_ragged_to_graph_buffers( |
| 404 | ragged_indices: torch.Tensor, |
| 405 | ragged_indptr: torch.Tensor, |
| 406 | ragged_indices_buffer: torch.Tensor, |
| 407 | ragged_indptr_buffer: torch.Tensor, |
| 408 | num_rows: int, |
| 409 | max_entries_per_row: int, |
| 410 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 411 | """Copy dynamic ragged metadata into persistent CUDA graph buffers. |
| 412 | |
| 413 | FULL decode graphs capture kernel argument addresses. Keep the returned |
| 414 | tensors backed by stable storage, while indptr continues to bound reads. |
| 415 | """ |
| 416 | indptr_out = ragged_indptr_buffer[: num_rows + 1] |
| 417 | indptr_out.copy_(ragged_indptr, non_blocking=True) |
| 418 | |
| 419 | max_entries = max(num_rows * max_entries_per_row, 1) |
| 420 | ragged_out = ragged_indices_buffer[:max_entries] |
| 421 | source_entries = ragged_indices.numel() |
| 422 | if source_entries > 0: |
| 423 | ragged_out[:source_entries].copy_(ragged_indices, non_blocking=True) |
| 424 | if _ON_GFX950: |
| 425 | # Preserve the graph-stable base pointer while exposing source capacity |
| 426 | # to the sync-free split selector; indptr still carries the true NNZ. |
| 427 | ragged_out = ragged_out[: max(source_entries, 1)] |
| 428 | return ragged_out, indptr_out |
| 429 | |
| 430 | |
| 431 | # (ragged_indices, ragged_indptr, lens) as consumed by rocm_sparse_attn_decode. |
| 432 | _TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 433 | |
| 434 | |
| 435 | @dataclass |
| 436 | class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): |
| 437 | decode_swa_ragged_indices: torch.Tensor | None = None |
| 438 | decode_swa_ragged_indptr: torch.Tensor | None = None |
| 439 | # Built only for the aiter prefill, which reads the SWA cache in place. |
| 440 | prefill_swa_ragged_indices: torch.Tensor | None = None |
| 441 | prefill_swa_ragged_indptr: torch.Tensor | None = None |
| 442 | |
| 443 | |
| 444 | class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): |
| 445 | # Keep fused multi-step decode disabled until update_draft_decode_metadata() |
| 446 | # also refreshes the ROCm-specific ragged SWA indices and indptrs. |
| 447 | supports_draft_decode_metadata_update = False |
| 448 | |
| 449 | def __init__(self, *args, **kwargs): |
| 450 | super().__init__(*args, **kwargs) |
| 451 | max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens |
| 452 | # The non-causal (DSpark draft) path widens each token's SWA index list |
| 453 | # to ``noncausal_index_width`` (>= window_size), so size the persistent |
| 454 | # ragged buffer to the wider bound to cover both causal and non-causal. |
| 455 | swa_index_width = max(self.window_size, self.noncausal_index_width) |
| 456 | self.decode_swa_ragged_indices_buffer = torch.empty( |
| 457 | max_tokens * swa_index_width, |
| 458 | dtype=torch.int32, |
| 459 | device=self.device, |
| 460 | ) |
| 461 | self.decode_swa_ragged_indptr_buffer = torch.empty( |
| 462 | max_tokens + 1, |
| 463 | dtype=torch.int32, |
| 464 | device=self.device, |
| 465 | ) |
| 466 | self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled( |
| 467 | self.vllm_config.kv_transfer_config is not None |
| 468 | ) |
| 469 | |
| 470 | def build( |
| 471 | self, |
| 472 | common_prefix_len: int, |
| 473 | common_attn_metadata: CommonAttentionMetadata, |
| 474 | fast_build: bool = False, |
| 475 | replay_start: torch.Tensor | None = None, |
| 476 | ) -> DeepseekV4ROCMAiterSparseSWAMetadata: |
| 477 | base = super().build( |
| 478 | common_prefix_len=common_prefix_len, |
| 479 | common_attn_metadata=common_attn_metadata, |
| 480 | fast_build=fast_build, |
| 481 | replay_start=replay_start, |
| 482 | ) |
| 483 | |
| 484 | prefill_ragged_indices = None |
| 485 | prefill_ragged_indptr = None |
| 486 | if ( |
| 487 | self._use_aiter_sparse_mla |
| 488 | and base.num_prefill_tokens > 0 |
| 489 | and base.prefill_swa_indices is not None |
| 490 | and base.prefill_swa_lens is not None |
| 491 | ): |
| 492 | # Paged-coordinate rows, so every layer's prefill shares one build. |
| 493 | prefill_ragged_indices, prefill_ragged_indptr = ( |
| 494 | build_ragged_indices_from_dense( |
| 495 | base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1), |
| 496 | base.prefill_swa_lens, |
| 497 | ) |
| 498 | ) |
| 499 | |
| 500 | ragged_indices = None |
| 501 | ragged_indptr = None |
| 502 | if ( |
| 503 | base.num_decode_tokens > 0 |
| 504 | and base.decode_swa_indices is not None |
| 505 | and base.decode_swa_lens is not None |
| 506 | ): |
| 507 | ragged_indices, ragged_indptr = build_ragged_indices_from_dense( |
| 508 | base.decode_swa_indices.reshape( |
| 509 | base.num_decode_tokens, base.decode_swa_width |
| 510 | ), |
| 511 | base.decode_swa_lens, |
| 512 | ) |
| 513 | ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( |
| 514 | ragged_indices, |
| 515 | ragged_indptr, |
| 516 | self.decode_swa_ragged_indices_buffer, |
| 517 | self.decode_swa_ragged_indptr_buffer, |
| 518 | base.num_decode_tokens, |
| 519 | base.decode_swa_width, |
| 520 | ) |
| 521 | |
| 522 | return DeepseekV4ROCMAiterSparseSWAMetadata( |
| 523 | **vars(base), |
| 524 | decode_swa_ragged_indices=ragged_indices, |
| 525 | decode_swa_ragged_indptr=ragged_indptr, |
| 526 | prefill_swa_ragged_indices=prefill_ragged_indices, |
| 527 | prefill_swa_ragged_indptr=prefill_ragged_indptr, |
| 528 | ) |
| 529 | |
| 530 | |
| 531 | @functools.cache |
| 532 | def _aiter_indexer_cache_ops() -> tuple[Callable[..., None], Callable[..., tuple]]: |
| 533 | """The aiter indexer key writer and query quantizer.""" |
| 534 | from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import ( |
| 535 | k_norm_rope_mxfp4_cache, |
| 536 | ) |
| 537 | from aiter.ops.triton.rope.q_rope_mxfp4_quant import q_rope_mxfp4_quant |
| 538 | |
| 539 | return k_norm_rope_mxfp4_cache, q_rope_mxfp4_quant |
| 540 | |
| 541 | |
| 542 | def rocm_mxfp4_indexer_k_store( |
| 543 | k_pre: torch.Tensor, |
| 544 | positions: torch.Tensor, |
| 545 | cos_sin_cache: torch.Tensor, |
| 546 | rms_norm_weight: torch.Tensor, |
| 547 | rms_norm_eps: float, |
| 548 | k_cache: torch.Tensor, |
| 549 | kv_slot_mapping: torch.Tensor, |
| 550 | compress_ratio: int, |
| 551 | use_fp4_cache: bool, |
| 552 | *, |
| 553 | num_heads: int, |
| 554 | ) -> None: |
| 555 | """`indexer_k_norm_rope_store` for the ROCm MXFP4 cache: aiter's cache op |
| 556 | writes the key in the order its MQA-logits kernel reads with ``num_heads`` |
| 557 | query heads.""" |
| 558 | assert use_fp4_cache, "the ROCm indexer cache op writes MXFP4 only" |
| 559 | layout = rocm_paged_mxfp4_cache_layout(num_heads, k_pre.shape[1], k_cache.shape[1]) |
| 560 | k_norm_rope_mxfp4_cache, _ = _aiter_indexer_cache_ops() |
| 561 | k_norm_rope_mxfp4_cache( |
| 562 | k_pre, |
| 563 | positions, |
| 564 | cos_sin_cache, |
| 565 | rms_norm_weight, |
| 566 | rms_norm_eps, |
| 567 | k_cache, |
| 568 | kv_slot_mapping, |
| 569 | compress_ratio, |
| 570 | shuffle=layout, |
| 571 | ) |
| 572 | |
| 573 | |
| 574 | def rocm_mxfp4_indexer_q_quant( |
| 575 | positions: torch.Tensor, |
| 576 | index_q: torch.Tensor, |
| 577 | index_q_cos_sin_cache: torch.Tensor, |
| 578 | index_weights: torch.Tensor, |
| 579 | index_weights_softmax_scale: float, |
| 580 | index_weights_head_scale: float, |
| 581 | use_fp4: bool = True, |
| 582 | weights_out_dtype: torch.dtype = torch.float32, |
| 583 | ) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]: |
| 584 | """`fused_indexer_q_rope_quant` for the ROCm MXFP4 indexer, on aiter's op: |
| 585 | ((packed [T, H, D // 2], e8m0 as one int32 per head [T, H]), fp32 |
| 586 | weights).""" |
| 587 | assert use_fp4 and weights_out_dtype == torch.float32, ( |
| 588 | "the ROCm indexer quantizes Q to MXFP4 and scores with fp32 weights" |
| 589 | ) |
| 590 | _, q_rope_mxfp4_quant = _aiter_indexer_cache_ops() |
| 591 | q_packed, q_scale, weights_out = q_rope_mxfp4_quant( |
| 592 | index_q, |
| 593 | positions, |
| 594 | index_q_cos_sin_cache, |
| 595 | index_weights, |
| 596 | index_weights_softmax_scale * index_weights_head_scale, |
| 597 | ) |
| 598 | return (q_packed, q_scale.view(torch.int32).squeeze(-1)), weights_out |
| 599 | |
| 600 | |
| 601 | class DeepseekV41RocmMxfp4Indexer(DeepseekV4Indexer): |
| 602 | """The indexer on aiter's paged MXFP4 cache: its K store and Q quant write |
| 603 | in the order aiter's MQA-logits kernel reads, and its layers score with it. |
| 604 | ``_produce_k`` and ``forward_q`` mirror ``DeepseekV4Indexer``'s except for |
| 605 | those two calls.""" |
| 606 | |
| 607 | mqa_cls = RocmSparseMQAIndexer |
| 608 | attn_cls = RocmSparseAttnIndexer |
| 609 | |
| 610 | def _produce_k( |
| 611 | self, |
| 612 | latent: torch.Tensor | None, |
| 613 | positions: torch.Tensor, |
| 614 | rotary_emb: torch.nn.Module, |
| 615 | ) -> None: |
| 616 | attn_metadata = get_forward_context().attn_metadata |
| 617 | if not isinstance(attn_metadata, dict) or latent is None: |
| 618 | return |
| 619 | assert self.owns_k |
| 620 | indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) |
| 621 | k_pre, _ = self.wk(latent) |
| 622 | rocm_mxfp4_indexer_k_store( |
| 623 | k_pre, |
| 624 | positions, |
| 625 | rotary_emb.cos_sin_cache, |
| 626 | self.k_norm.weight, |
| 627 | self.k_norm.variance_epsilon, |
| 628 | self.k_cache.kv_cache, |
| 629 | indexer_metadata.slot_mapping, |
| 630 | self.compress_ratio, |
| 631 | self.use_fp4_kv, |
| 632 | num_heads=self.n_head, |
| 633 | ) |
| 634 | |
| 635 | def forward_q( |
| 636 | self, |
| 637 | qr: torch.Tensor | QuantizedActivation, |
| 638 | qr_scale: torch.Tensor | None, |
| 639 | indexer_weights: torch.Tensor, |
| 640 | positions: torch.Tensor, |
| 641 | rotary_emb: torch.nn.Module, |
| 642 | ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: |
| 643 | q = self._wq_b_proj(qr, qr_scale).view(-1, self.n_head, self.head_dim) |
| 644 | (q, q_scale), weights = rocm_mxfp4_indexer_q_quant( |
| 645 | positions, |
| 646 | q, |
| 647 | rotary_emb.cos_sin_cache, |
| 648 | indexer_weights, |
| 649 | self.softmax_scale, |
| 650 | self.n_head**-0.5, |
| 651 | use_fp4=self.use_fp4_kv, |
| 652 | weights_out_dtype=self.indexer_weights_dtype, |
| 653 | ) |
| 654 | return q, q_scale, weights |
| 655 | |
| 656 | |
| 657 | class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): |
| 658 | @staticmethod |
| 659 | def get_name() -> str: |
| 660 | return "ROCM_FLASHMLA_SPARSE_DSV4" |
| 661 | |
| 662 | @staticmethod |
| 663 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 664 | return [64, 128] |
| 665 | |
| 666 | @staticmethod |
| 667 | def get_builder_cls() -> type[DeepseekV4SparseMLAMetadataBuilder]: |
| 668 | return DeepseekV4SparseMLAMetadataBuilder |
| 669 | |
| 670 | |
| 671 | class DeepseekV41ROCMAiterSparseSWABackend(DeepseekSparseSWABackend): |
| 672 | @staticmethod |
| 673 | def get_builder_cls() -> type["DeepseekV4ROCMAiterSparseSWAMetadataBuilder"]: |
| 674 | return DeepseekV4ROCMAiterSparseSWAMetadataBuilder |
| 675 | |
| 676 | |
| 677 | class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): |
| 678 | """ROCm sparse MLA attention layer for DeepSeek V4.1.""" |
| 679 | |
| 680 | backend_cls = DeepseekV4ROCMAiterMLASparseBackend |
| 681 | swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend |
| 682 | _use_aiter_sparse_mla = False |
| 683 | |
| 684 | def _indexer_cls( |
| 685 | self, k_cache: DeepseekV4IndexerCache | None |
| 686 | ) -> type[DeepseekV4Indexer]: |
| 687 | if k_cache is not None and k_cache.rocm_mxfp4: |
| 688 | return DeepseekV41RocmMxfp4Indexer |
| 689 | return super()._indexer_cls(k_cache) |
| 690 | |
| 691 | def __init__(self, *args, **kwargs): |
| 692 | vllm_config = args[0] if args else kwargs["vllm_config"] |
| 693 | super().__init__(*args, **kwargs) |
| 694 | # CUDA executes WO_A with a quantized grouped-BMM kernel. ROCm's |
| 695 | # correctness path below dequantizes WO_A once and uses torch.einsum, |
| 696 | # so retain the ordinary MXFP8 linear kernel for post-load processing |
| 697 | # instead of asking for the CUDA-only BMM kernel. |
| 698 | self.wo_a.is_bmm = False |
| 699 | self._has_kv_transfer = vllm_config.kv_transfer_config is not None |
| 700 | # Block scale for the preshuffled weight; None = not preshuffled. |
| 701 | self._wqa_wkv_scale: torch.Tensor | None = None |
| 702 | self._wo_b_scale: torch.Tensor | None = None |
| 703 | self._fused_compressor_weight: torch.Tensor | None |
| 704 | self.register_buffer("_fused_compressor_weight", None, persistent=False) |
| 705 | self._fused_compressor_split_sizes: tuple[int, int] | None = None |
| 706 | # Decode ragged topk metadata is a pure function of the indices its |
| 707 | # index source published, so every consumer of a source rebuilds the |
| 708 | # same thing. The source owns one cache per compress ratio, refreshed |
| 709 | # when it runs; consumers below it read through. |
| 710 | self._topk_ragged_cache: dict[int, _TopkRagged] = {} |
| 711 | self._prefill_topk_ragged_cache: dict[ |
| 712 | str, tuple[torch.Tensor, torch.Tensor] |
| 713 | ] = {} |
| 714 | self._index_source_prefix: str | None = None |
| 715 | if self.compress_ratio > 0: |
| 716 | assert self.index_source_layer_id is not None |
| 717 | self._index_source_prefix = _replace_layer_index( |
| 718 | self.prefix, self.index_source_layer_id |
| 719 | ) |
| 720 | if self._index_source_prefix not in self._static_forward_context: |
| 721 | raise NotImplementedError( |
| 722 | f"Index source {self._index_source_prefix} not found on " |
| 723 | "this rank; PP splits inside a v4.1 index-sharing group " |
| 724 | "are not supported." |
| 725 | ) |
| 726 | self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer) |
| 727 | if self.compressor is None and self.indexer is None: |
| 728 | # Dense layers have nothing to overlap; keep the base serial path. |
| 729 | self.aux_stream_list = None |
| 730 | |
| 731 | @classmethod |
| 732 | def get_padded_num_q_heads(cls, num_heads: int) -> int: |
| 733 | return num_heads |
| 734 | |
| 735 | def prepare_attn_preshuffle(self) -> None: |
| 736 | from vllm._aiter_ops import rocm_aiter_ops |
| 737 | |
| 738 | if not rocm_aiter_ops.is_enabled(): |
| 739 | return |
| 740 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 741 | _upcast_e8m0_to_fp32, |
| 742 | ) |
| 743 | from vllm.model_executor.utils import replace_parameter |
| 744 | |
| 745 | def _prep(linear) -> torch.Tensor | None: |
| 746 | w = getattr(linear, "weight", None) |
| 747 | if w is None or w.dim() != 2: |
| 748 | return None |
| 749 | # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold. |
| 750 | if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0: |
| 751 | return None |
| 752 | ws = getattr(linear, "weight_scale_inv", None) # per-block scale |
| 753 | if ws is None: |
| 754 | return None |
| 755 | if ws.dtype == torch.float8_e8m0fnu: |
| 756 | ws = _upcast_e8m0_to_fp32(ws).contiguous() |
| 757 | # Shuffle the weight in place (single weight, no unshuffled copy). |
| 758 | replace_parameter( |
| 759 | linear, |
| 760 | "weight", |
| 761 | rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), |
| 762 | ) |
| 763 | return ws |
| 764 | |
| 765 | self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) |
| 766 | self._wo_b_scale = _prep(self.wo_b) |
| 767 | |
| 768 | def prepare_compressor_gemm_fusion(self) -> bool: |
| 769 | # V4.1 derives index keys from the source compressor's emitted latent |
| 770 | # and has no nested ``indexer.compressor``. Keep the projections |
| 771 | # separate and use the shared linear/PyTorch correctness path. |
| 772 | return False |
| 773 | |
| 774 | def _bpre_attn_gemm( |
| 775 | self, |
| 776 | weight: torch.Tensor, |
| 777 | scale: torch.Tensor, |
| 778 | x: torch.Tensor, |
| 779 | reduce_tp: bool, |
| 780 | ) -> torch.Tensor: |
| 781 | from vllm._aiter_ops import rocm_aiter_ops |
| 782 | |
| 783 | x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True) |
| 784 | out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( |
| 785 | x_fp8, weight, x_scale, scale, output_dtype=x.dtype |
| 786 | ) |
| 787 | if reduce_tp and get_tensor_model_parallel_world_size() > 1: |
| 788 | out = tensor_model_parallel_all_reduce(out) |
| 789 | return out |
| 790 | |
| 791 | def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 792 | if self._wqa_wkv_scale is not None and hidden_states.dim() == 2: |
| 793 | return self._bpre_attn_gemm( |
| 794 | self.fused_wqa_wkv.weight, self._wqa_wkv_scale, hidden_states, False |
| 795 | ) |
| 796 | return super()._fused_wqa_wkv_gemm(hidden_states) |
| 797 | |
| 798 | def _run_parallel_input_projections( |
| 799 | self, hidden_states: torch.Tensor |
| 800 | ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: |
| 801 | if self.indexer is not None and self.compressor is None: |
| 802 | # Serial: qr must sit on the default stream before the q-side |
| 803 | # fork below can consume it (ROCm). |
| 804 | aux_streams = self.aux_stream_list |
| 805 | self.aux_stream_list = None |
| 806 | try: |
| 807 | return super()._run_parallel_input_projections(hidden_states) |
| 808 | finally: |
| 809 | self.aux_stream_list = aux_streams |
| 810 | return super()._run_parallel_input_projections(hidden_states) |
| 811 | |
| 812 | def forward( |
| 813 | self, |
| 814 | positions: torch.Tensor, |
| 815 | hidden_states: torch.Tensor, |
| 816 | llama_4_scaling: torch.Tensor | None = None, |
| 817 | ) -> torch.Tensor: |
| 818 | if ( |
| 819 | self.indexer is None |
| 820 | or self._prepare_and_attn_fn == self._prepare_and_attn_eager |
| 821 | or not current_platform.enable_multi_stream_overlap( |
| 822 | self.aux_stream_list, get_forward_context().attn_metadata |
| 823 | ) |
| 824 | ): |
| 825 | # Sequential fallback: no forks outside the capture region, where |
| 826 | # HIP event sync is unreliable; MRV1 keeps the input prep in its |
| 827 | # wide eager region (#51430). |
| 828 | aux_streams = self.aux_stream_list |
| 829 | self.aux_stream_list = None |
| 830 | try: |
| 831 | return super().forward(positions, hidden_states, llama_4_scaling) |
| 832 | finally: |
| 833 | self.aux_stream_list = aux_streams |
| 834 | |
| 835 | attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states) |
| 836 | if self.compressor is not None: |
| 837 | q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_full( |
| 838 | hidden_states, positions |
| 839 | ) |
| 840 | else: |
| 841 | q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_reindex( |
| 842 | hidden_states, positions |
| 843 | ) |
| 844 | self._sparse_indexer_and_attn( |
| 845 | hidden_states, |
| 846 | index_q, |
| 847 | index_q_scale, |
| 848 | index_weights, |
| 849 | q, |
| 850 | kv, |
| 851 | positions, |
| 852 | attn_out, |
| 853 | ) |
| 854 | return self._o_proj(attn_out, positions) |
| 855 | |
| 856 | def _forward_csa2_full( |
| 857 | self, hidden_states: torch.Tensor, positions: torch.Tensor |
| 858 | ) -> tuple[ |
| 859 | torch.Tensor, |
| 860 | torch.Tensor, |
| 861 | torch.Tensor | None, |
| 862 | torch.Tensor | None, |
| 863 | torch.Tensor | None, |
| 864 | ]: |
| 865 | """CSA layers: one fork replaces the base pipeline's three. |
| 866 | |
| 867 | Each fork/join is a HIP event pair per side stream on ROCm, so the |
| 868 | base three-fork pipeline is expensive there; the compressor chain and |
| 869 | the indexer weights projection run fully in parallel with the default |
| 870 | chain. The indexer K write and q-side stay serial after the join: |
| 871 | ``wk`` consumes the aux-0 latent and the q-side's wq_b consumes the |
| 872 | default-stream qr. |
| 873 | """ |
| 874 | attn_metadata = get_forward_context().attn_metadata |
| 875 | compressor = self.compressor |
| 876 | indexer = self.indexer |
| 877 | assert compressor is not None and indexer is not None |
| 878 | aux_streams = self.aux_stream_list |
| 879 | assert aux_streams is not None and len(aux_streams) >= 2 |
| 880 | |
| 881 | def default_chain(): |
| 882 | qr_kv = self._fused_wqa_wkv_gemm(hidden_states) |
| 883 | qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) |
| 884 | q = self._wq_b_proj(qr, qr_scale).view( |
| 885 | -1, self.n_local_heads, self.head_dim |
| 886 | ) |
| 887 | return ( |
| 888 | self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata), |
| 889 | qr, |
| 890 | qr_scale, |
| 891 | kv, |
| 892 | ) |
| 893 | |
| 894 | def compressor_chain(): |
| 895 | kv_score = torch.mm( |
| 896 | hidden_states, |
| 897 | compressor.fused_wkv_wgate.weight.T, |
| 898 | out_dtype=torch.float32, |
| 899 | ) |
| 900 | latent = compressor(kv_score, positions) |
| 901 | compressor.insert_cache(latent, positions, self.rotary_emb) |
| 902 | return latent |
| 903 | |
| 904 | def indexer_weights_chain(): |
| 905 | # ReplicatedLinear returns (output, bias); bias is None. |
| 906 | weights, _ = indexer.weights_proj(hidden_states) |
| 907 | return weights |
| 908 | |
| 909 | (q, qr, qr_scale, kv), (latent, indexer_weights) = execute_in_parallel( |
| 910 | default_chain, |
| 911 | [compressor_chain, indexer_weights_chain], |
| 912 | self.ln_events[0], |
| 913 | self.ln_events[1:3], |
| 914 | aux_streams[:2], |
| 915 | enable=True, |
| 916 | default_first=True, |
| 917 | ) |
| 918 | |
| 919 | indexer._produce_k(latent, positions, self.indexer_rotary_emb) |
| 920 | index_q, index_q_scale, index_weights_out = indexer.forward_q( |
| 921 | qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb |
| 922 | ) |
| 923 | return q, kv, index_q, index_q_scale, index_weights_out |
| 924 | |
| 925 | def _forward_csa2_reindex( |
| 926 | self, hidden_states: torch.Tensor, positions: torch.Tensor |
| 927 | ) -> tuple[ |
| 928 | torch.Tensor, |
| 929 | torch.Tensor, |
| 930 | torch.Tensor | None, |
| 931 | torch.Tensor | None, |
| 932 | torch.Tensor | None, |
| 933 | ]: |
| 934 | """Indexer-only layers: the q-side overlaps the SWA q path. |
| 935 | |
| 936 | The input projections run serial first: the q-side consumes qr, so |
| 937 | it can only overlap the SWA q projection and KV insert. |
| 938 | """ |
| 939 | indexer = self.indexer |
| 940 | assert indexer is not None |
| 941 | attn_metadata = get_forward_context().attn_metadata |
| 942 | # Serial on these layers via the _run_parallel_input_projections |
| 943 | # override: qr must sit on the default stream before the fork. |
| 944 | qr_kv, _, indexer_weights = self._run_parallel_input_projections(hidden_states) |
| 945 | assert indexer_weights is not None |
| 946 | qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) |
| 947 | |
| 948 | def swa_q_chain(): |
| 949 | q = self._wq_b_proj(qr, qr_scale).view( |
| 950 | -1, self.n_local_heads, self.head_dim |
| 951 | ) |
| 952 | return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) |
| 953 | |
| 954 | def indexer_q_chain(): |
| 955 | return indexer.forward_q( |
| 956 | qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb |
| 957 | ) |
| 958 | |
| 959 | aux_streams = self.aux_stream_list |
| 960 | assert aux_streams is not None |
| 961 | q, aux_results = execute_in_parallel( |
| 962 | swa_q_chain, |
| 963 | [indexer_q_chain], |
| 964 | self.ln_events[0], |
| 965 | self.ln_events[1:2], |
| 966 | aux_streams[:1], |
| 967 | enable=True, |
| 968 | default_first=True, |
| 969 | ) |
| 970 | index_q, index_q_scale, index_weights_out = aux_results[0] |
| 971 | return q, kv, index_q, index_q_scale, index_weights_out |
| 972 | |
| 973 | @functools.cached_property |
| 974 | def _wq_b_uses_aiter_block_scaled(self) -> bool: |
| 975 | """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel. |
| 976 | |
| 977 | Cached: the linear kernels and the aiter env gates are fixed once |
| 978 | the model is built, so this is evaluated at the first forward |
| 979 | only. |
| 980 | |
| 981 | The fused norm+quant path is only valid if the quant and GEMM it |
| 982 | replaces are exactly the aiter ones; otherwise fall back to the |
| 983 | shared path. |
| 984 | """ |
| 985 | from vllm._aiter_ops import rocm_aiter_ops |
| 986 | from vllm.model_executor.kernels.linear.scaled_mm import ( |
| 987 | Fp8BlockScaledMMLinearKernel, |
| 988 | ) |
| 989 | |
| 990 | if not rocm_aiter_ops.is_linear_fp8_enabled(): |
| 991 | return False |
| 992 | |
| 993 | linears = [self.wq_b] |
| 994 | if self.indexer is not None: |
| 995 | linears.append(self.indexer.wq_b) |
| 996 | for linear in linears: |
| 997 | kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) |
| 998 | if not isinstance(kernel, Fp8BlockScaledMMLinearKernel): |
| 999 | return False |
| 1000 | return True |
| 1001 | |
| 1002 | def _split_qkv_and_norm( |
| 1003 | self, qr_kv: torch.Tensor |
| 1004 | ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: |
| 1005 | """Fuse q/kv RMSNorm + per-1x128 fp8 q quant into one aiter kernel. |
| 1006 | |
| 1007 | The shared path norms q and kv in one triton kernel and the wq_b |
| 1008 | linears then re-read the bf16 qr to quantize it. The aiter kernel |
| 1009 | computes both RMSNorms (fp32 accumulate) and the fp8 group quant |
| 1010 | in a single pass, writing fp8 qr + group scales directly; both |
| 1011 | wq_b GEMMs (attention and indexer) then consume that pair and |
| 1012 | skip their own input quant. kv stays bf16: the fused insert |
| 1013 | kernel RoPE/quantizes it itself. Falls back to the shared path |
| 1014 | when the aiter linear path is not active. |
| 1015 | """ |
| 1016 | qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) |
| 1017 | if not ( |
| 1018 | qr.dim() == 2 |
| 1019 | and qr.shape[0] > 0 |
| 1020 | and self.q_lora_rank % 128 == 0 |
| 1021 | and self._wq_b_uses_aiter_block_scaled |
| 1022 | ): |
| 1023 | return super()._split_qkv_and_norm(qr_kv) |
| 1024 | |
| 1025 | from vllm._aiter_ops import rocm_aiter_ops |
| 1026 | |
| 1027 | return rocm_aiter_ops.fused_qk_rmsnorm_group_quant( |
| 1028 | q=qr, |
| 1029 | q_weight=self.q_norm.weight.data, |
| 1030 | q_epsilon=self.eps, |
| 1031 | kv=kv, |
| 1032 | kv_weight=self.kv_norm.weight.data, |
| 1033 | kv_epsilon=self.eps, |
| 1034 | group_size=128, |
| 1035 | transpose_scale=False, |
| 1036 | ) |
| 1037 | |
| 1038 | def _alloc_attn_out( |
| 1039 | self, num_tokens: int, hidden_states: torch.Tensor |
| 1040 | ) -> torch.Tensor | QuantizedActivation: |
| 1041 | if not _ON_GFX950: |
| 1042 | return super()._alloc_attn_out(num_tokens, hidden_states) |
| 1043 | # wo_a's MXFP8 input: the decode reduce writes it directly, prefill |
| 1044 | # rows are rotated and quantized after their bf16 attention. |
| 1045 | width = self.n_local_heads * self.head_dim |
| 1046 | data = torch.empty( |
| 1047 | (num_tokens, width), dtype=torch.float8_e4m3fn, device=hidden_states.device |
| 1048 | ) |
| 1049 | scale = torch.empty( |
| 1050 | (num_tokens, width // 32), dtype=torch.uint8, device=hidden_states.device |
| 1051 | ) |
| 1052 | return QuantizedActivation( |
| 1053 | data=data, |
| 1054 | scale=scale, |
| 1055 | orig_dtype=hidden_states.dtype, |
| 1056 | orig_shape=data.shape, |
| 1057 | quant_key=kMxfp8Dynamic, |
| 1058 | ) |
| 1059 | |
| 1060 | def _o_proj( |
| 1061 | self, attn_out: torch.Tensor | QuantizedActivation, positions: torch.Tensor |
| 1062 | ) -> torch.Tensor: |
| 1063 | if isinstance(attn_out, QuantizedActivation): |
| 1064 | z = rocm_mxfp8_wo_a_bmm( |
| 1065 | attn_out.data, |
| 1066 | attn_out.scale, |
| 1067 | self.wo_a, |
| 1068 | self.n_local_groups, |
| 1069 | self.o_lora_rank, |
| 1070 | ) |
| 1071 | return self._wo_b_after_wo_a(z) |
| 1072 | o = attn_out[:, : self.n_local_heads, :] |
| 1073 | # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. |
| 1074 | z = rocm_inv_rope_einsum( |
| 1075 | self.rotary_emb, |
| 1076 | o, |
| 1077 | positions, |
| 1078 | self.rope_head_dim, |
| 1079 | self.n_local_groups, |
| 1080 | self.o_lora_rank, |
| 1081 | self.wo_a, |
| 1082 | inverse_rope=False, |
| 1083 | ) |
| 1084 | return self._wo_b_after_wo_a(z.flatten(1)) |
| 1085 | |
| 1086 | def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor: |
| 1087 | if self._wo_b_scale is not None and zf.dim() == 2: |
| 1088 | return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) |
| 1089 | return self.wo_b(zf) |
| 1090 | |
| 1091 | def forward_mqa( |
| 1092 | self, |
| 1093 | q: torch.Tensor, |
| 1094 | kv: torch.Tensor, |
| 1095 | positions: torch.Tensor, |
| 1096 | output: torch.Tensor | QuantizedActivation, |
| 1097 | ) -> None: |
| 1098 | mxfp8_out = output if isinstance(output, QuantizedActivation) else None |
| 1099 | if mxfp8_out is None: |
| 1100 | assert isinstance(output, torch.Tensor) |
| 1101 | assert output.shape == q.shape, ( |
| 1102 | f"output buffer shape {output.shape} must match q shape {q.shape}" |
| 1103 | ) |
| 1104 | assert output.dtype == q.dtype, ( |
| 1105 | f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" |
| 1106 | ) |
| 1107 | |
| 1108 | forward_context = get_forward_context() |
| 1109 | attn_metadata = forward_context.attn_metadata |
| 1110 | |
| 1111 | if attn_metadata is None: |
| 1112 | # Warmup dummy run: no real metadata. Reserve the workspace |
| 1113 | # _forward_prefill would; the dequantize / topk / sparse_fwd kernels |
| 1114 | # are skipped this step. The aiter prefill reads the caches in place, |
| 1115 | # so it asks only for the bf16 rows of the MXFP8 output. |
| 1116 | swa_only = self.compress_ratio == 0 |
| 1117 | N = ( |
| 1118 | 0 |
| 1119 | if swa_only |
| 1120 | else (self.max_model_len + self.compress_ratio - 1) |
| 1121 | // self.compress_ratio |
| 1122 | ) |
| 1123 | M = N + self.window_size + self.max_num_batched_tokens |
| 1124 | shapes = self._prefill_workspace_shapes( |
| 1125 | M, |
| 1126 | self.max_num_batched_tokens, |
| 1127 | q, |
| 1128 | gather=not self._use_aiter_sparse_mla, |
| 1129 | ) |
| 1130 | if shapes: |
| 1131 | current_workspace_manager().get_simultaneous(*shapes) |
| 1132 | if mxfp8_out is None: |
| 1133 | output.zero_() |
| 1134 | else: |
| 1135 | mxfp8_out.data.zero_() |
| 1136 | mxfp8_out.scale.zero_() |
| 1137 | return |
| 1138 | |
| 1139 | assert isinstance(attn_metadata, dict) |
| 1140 | rocm_metadata = cast( |
| 1141 | DeepseekV4FlashMLAMetadata | None, |
| 1142 | attn_metadata.get(self.compressed_cache_prefix) |
| 1143 | if self.compressed_cache_prefix is not None |
| 1144 | else None, |
| 1145 | ) |
| 1146 | swa_metadata = cast( |
| 1147 | DeepseekV4ROCMAiterSparseSWAMetadata | None, |
| 1148 | attn_metadata.get(self.swa_cache_layer.prefix), |
| 1149 | ) |
| 1150 | assert swa_metadata is not None |
| 1151 | |
| 1152 | swa_only = self.compress_ratio == 0 |
| 1153 | self_kv_cache = None if swa_only else self._compressed_kv_cache() |
| 1154 | swa_kv_cache = self.swa_cache_layer.kv_cache |
| 1155 | |
| 1156 | num_decodes = swa_metadata.num_decodes |
| 1157 | num_prefills = swa_metadata.num_prefills |
| 1158 | num_decode_tokens = swa_metadata.num_decode_tokens |
| 1159 | |
| 1160 | if num_prefills > 0: |
| 1161 | self._forward_prefill( |
| 1162 | q=q[num_decode_tokens:], |
| 1163 | positions=positions[num_decode_tokens:], |
| 1164 | compressed_k_cache=self_kv_cache, |
| 1165 | swa_k_cache=swa_kv_cache, |
| 1166 | output=( |
| 1167 | output[num_decode_tokens:] |
| 1168 | if isinstance(output, torch.Tensor) |
| 1169 | else None |
| 1170 | ), |
| 1171 | attn_metadata=rocm_metadata, |
| 1172 | swa_metadata=swa_metadata, |
| 1173 | mxfp8_out=( |
| 1174 | None |
| 1175 | if mxfp8_out is None |
| 1176 | else ( |
| 1177 | mxfp8_out.data[num_decode_tokens:], |
| 1178 | mxfp8_out.scale[num_decode_tokens:], |
| 1179 | ) |
| 1180 | ), |
| 1181 | ) |
| 1182 | if mxfp8_out is not None: |
| 1183 | # The reduce epilogue rotates and quantizes every decode row and |
| 1184 | # the prefill path already did its own, so nothing is owed here. |
| 1185 | if num_decodes > 0: |
| 1186 | self._forward_decode( |
| 1187 | q=q[:num_decode_tokens], |
| 1188 | positions=positions[:num_decode_tokens], |
| 1189 | kv_cache=self_kv_cache, |
| 1190 | swa_metadata=swa_metadata, |
| 1191 | attn_metadata=rocm_metadata, |
| 1192 | swa_only=swa_only, |
| 1193 | output=None, |
| 1194 | output_mxfp8=( |
| 1195 | mxfp8_out.data[:num_decode_tokens], |
| 1196 | mxfp8_out.scale[:num_decode_tokens], |
| 1197 | ), |
| 1198 | ) |
| 1199 | return |
| 1200 | assert isinstance(output, torch.Tensor) |
| 1201 | rotated = 0 |
| 1202 | if num_decodes > 0: |
| 1203 | rotated = self._forward_decode( |
| 1204 | q=q[:num_decode_tokens], |
| 1205 | positions=positions[:num_decode_tokens], |
| 1206 | kv_cache=self_kv_cache, |
| 1207 | swa_metadata=swa_metadata, |
| 1208 | attn_metadata=rocm_metadata, |
| 1209 | swa_only=swa_only, |
| 1210 | output=output[:num_decode_tokens], |
| 1211 | ) |
| 1212 | # Only the decode reduce rotates its own rows, and only the leading |
| 1213 | # `rotated` of them; prefill rows and any decode path that did not |
| 1214 | # fuse still owe the standalone pass. Settle that here rather than in |
| 1215 | # _o_proj: the split is batch-dependent and _o_proj runs compiled, |
| 1216 | # where such a value freezes at its trace-time value. |
| 1217 | rocm_inverse_rope_rows_( |
| 1218 | output[rotated:, : self.n_local_heads, :], |
| 1219 | positions[rotated:], |
| 1220 | self.rotary_emb.cos_sin_cache, |
| 1221 | self.rope_head_dim, |
| 1222 | ) |
| 1223 | |
| 1224 | def _decode_topk_ragged( |
| 1225 | self, |
| 1226 | swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, |
| 1227 | attn_metadata: DeepseekV4FlashMLAMetadata | None, |
| 1228 | num_decodes: int, |
| 1229 | num_decode_tokens: int, |
| 1230 | ) -> _TopkRagged: |
| 1231 | """Ragged form of the topk indices this layer's index source published. |
| 1232 | |
| 1233 | The packing depends only on the shared ``topk_indices_buffer`` and |
| 1234 | per-step metadata, apart from the layer's own compress ratio, so the |
| 1235 | source memoizes one result per ratio for all the layers below it. |
| 1236 | """ |
| 1237 | assert attn_metadata is not None |
| 1238 | assert swa_metadata.is_valid_token is not None |
| 1239 | assert self.topk_indices_buffer is not None |
| 1240 | assert self._index_source_prefix is not None |
| 1241 | |
| 1242 | source = self._static_forward_context[self._index_source_prefix] |
| 1243 | if source is self: |
| 1244 | # Fresh indices as of this layer; drop what the last step cached. |
| 1245 | source._topk_ragged_cache = {} |
| 1246 | cached = source._topk_ragged_cache.get(self.compress_ratio) |
| 1247 | if cached is not None: |
| 1248 | return cached |
| 1249 | |
| 1250 | built = compute_global_topk_ragged_indices_and_indptr( |
| 1251 | self.topk_indices_buffer[:num_decode_tokens], |
| 1252 | swa_metadata.token_to_req_indices, |
| 1253 | attn_metadata.block_table[:num_decodes], |
| 1254 | attn_metadata.block_size // self.compress_ratio, |
| 1255 | swa_metadata.is_valid_token[:num_decode_tokens], |
| 1256 | ) |
| 1257 | source._topk_ragged_cache[self.compress_ratio] = built |
| 1258 | return built |
| 1259 | |
| 1260 | def _forward_decode( |
| 1261 | self, |
| 1262 | q: torch.Tensor, |
| 1263 | positions: torch.Tensor, |
| 1264 | kv_cache: torch.Tensor | None, |
| 1265 | swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, |
| 1266 | attn_metadata: DeepseekV4FlashMLAMetadata | None, |
| 1267 | swa_only: bool, |
| 1268 | output: torch.Tensor | None, |
| 1269 | output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 1270 | ) -> int: |
| 1271 | """Returns how many leading rows the decode epilogue inverse-RoPE'd.""" |
| 1272 | num_decodes = swa_metadata.num_decodes |
| 1273 | num_decode_tokens = swa_metadata.num_decode_tokens |
| 1274 | |
| 1275 | topk_lens = None |
| 1276 | topk_ragged_indices = None |
| 1277 | topk_ragged_indptr = None |
| 1278 | if not swa_only: |
| 1279 | ( |
| 1280 | topk_ragged_indices, |
| 1281 | topk_ragged_indptr, |
| 1282 | topk_lens, |
| 1283 | ) = self._decode_topk_ragged( |
| 1284 | swa_metadata=swa_metadata, |
| 1285 | attn_metadata=attn_metadata, |
| 1286 | num_decodes=num_decodes, |
| 1287 | num_decode_tokens=num_decode_tokens, |
| 1288 | ) |
| 1289 | |
| 1290 | if self._use_aiter_sparse_mla: |
| 1291 | assert swa_metadata.decode_swa_ragged_indices is not None |
| 1292 | assert swa_metadata.decode_swa_ragged_indptr is not None |
| 1293 | out = output |
| 1294 | if output_mxfp8 is not None: |
| 1295 | # The aiter kernel has no MXFP8 epilogue: write bf16 rows, then |
| 1296 | # rotate and quantize them as the decode reduce would have. |
| 1297 | out = self._mxfp8_bf16_rows(q.shape[0], q) |
| 1298 | assert out is not None |
| 1299 | self._aiter_sparse_mla( |
| 1300 | q, |
| 1301 | out, |
| 1302 | self.swa_cache_layer.kv_cache, |
| 1303 | swa_metadata.decode_swa_ragged_indices, |
| 1304 | swa_metadata.decode_swa_ragged_indptr, |
| 1305 | kv_cache, |
| 1306 | topk_ragged_indices, |
| 1307 | topk_ragged_indptr, |
| 1308 | ) |
| 1309 | if output_mxfp8 is not None: |
| 1310 | rocm_inverse_rope_mxfp8_rows( |
| 1311 | out, |
| 1312 | positions, |
| 1313 | self.rotary_emb.cos_sin_cache, |
| 1314 | self.rope_head_dim, |
| 1315 | output_mxfp8[0], |
| 1316 | output_mxfp8[1], |
| 1317 | ) |
| 1318 | return q.shape[0] |
| 1319 | # The aiter kernel leaves the inverse RoPE to the caller. |
| 1320 | return 0 |
| 1321 | |
| 1322 | return rocm_sparse_attn_decode( |
| 1323 | q=q, |
| 1324 | kv_cache=kv_cache, |
| 1325 | swa_k_cache=self.swa_cache_layer.kv_cache, |
| 1326 | swa_only=swa_only, |
| 1327 | topk_indices=None, |
| 1328 | topk_lens=topk_lens, |
| 1329 | swa_indices=swa_metadata.decode_swa_indices, |
| 1330 | swa_lens=swa_metadata.decode_swa_lens, |
| 1331 | swa_ragged_indices=swa_metadata.decode_swa_ragged_indices, |
| 1332 | swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr, |
| 1333 | topk_ragged_indices=topk_ragged_indices, |
| 1334 | topk_ragged_indptr=topk_ragged_indptr, |
| 1335 | attn_sink=self.attn_sink, |
| 1336 | scale=self.scale, |
| 1337 | head_dim=self.head_dim, |
| 1338 | nope_head_dim=self.nope_head_dim, |
| 1339 | rope_head_dim=self.rope_head_dim, |
| 1340 | output=output, |
| 1341 | output_mxfp8=output_mxfp8, |
| 1342 | inv_rope_positions=positions, |
| 1343 | inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache, |
| 1344 | extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free( |
| 1345 | self.kv_cache_dtype, |
| 1346 | self._has_kv_transfer, |
| 1347 | not swa_only and kv_cache is not None, |
| 1348 | ), |
| 1349 | ) |
| 1350 | |
| 1351 | def _forward_prefill( |
| 1352 | self, |
| 1353 | q: torch.Tensor, |
| 1354 | positions: torch.Tensor, |
| 1355 | compressed_k_cache: torch.Tensor | None, |
| 1356 | swa_k_cache: torch.Tensor, |
| 1357 | output: torch.Tensor | None, |
| 1358 | attn_metadata: DeepseekV4FlashMLAMetadata | None, |
| 1359 | swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, |
| 1360 | mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 1361 | ) -> None: |
| 1362 | """Sparse prefill into bf16 ``output``, or into ``mxfp8_out``. |
| 1363 | |
| 1364 | With ``mxfp8_out`` the rows go to a bf16 workspace first and are then |
| 1365 | inverse-RoPE'd and quantized into it in one pass. |
| 1366 | """ |
| 1367 | if self._use_aiter_sparse_mla: |
| 1368 | self._forward_prefill_aiter( |
| 1369 | q=q, |
| 1370 | positions=positions, |
| 1371 | compressed_k_cache=compressed_k_cache, |
| 1372 | swa_k_cache=swa_k_cache, |
| 1373 | output=output, |
| 1374 | attn_metadata=attn_metadata, |
| 1375 | swa_metadata=swa_metadata, |
| 1376 | mxfp8_out=mxfp8_out, |
| 1377 | ) |
| 1378 | return |
| 1379 | |
| 1380 | swa_only = attn_metadata is None |
| 1381 | |
| 1382 | num_prefill_tokens = swa_metadata.num_prefill_tokens |
| 1383 | num_decodes = swa_metadata.num_decodes |
| 1384 | num_decode_tokens = swa_metadata.num_decode_tokens |
| 1385 | |
| 1386 | seq_lens = swa_metadata.prefill_seq_lens |
| 1387 | gather_lens = swa_metadata.prefill_gather_lens |
| 1388 | assert seq_lens is not None |
| 1389 | assert gather_lens is not None |
| 1390 | |
| 1391 | query_start_loc_cpu = swa_metadata.query_start_loc_cpu |
| 1392 | query_start_loc = swa_metadata.query_start_loc |
| 1393 | assert query_start_loc_cpu is not None |
| 1394 | assert query_start_loc is not None |
| 1395 | prefill_token_base = query_start_loc_cpu[num_decodes] |
| 1396 | |
| 1397 | # Local indices filled by the index source; SWA-only layers pass |
| 1398 | # top_k=0 and never read them. |
| 1399 | assert self.topk_indices_buffer is not None |
| 1400 | topk_indices = self.topk_indices_buffer[num_decode_tokens:] |
| 1401 | topk_indices = topk_indices[:num_prefill_tokens] |
| 1402 | if not swa_only: |
| 1403 | top_k = topk_indices.shape[-1] |
| 1404 | N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio |
| 1405 | else: |
| 1406 | top_k = 0 |
| 1407 | N = 0 |
| 1408 | |
| 1409 | M = N + self.window_size + self.max_num_batched_tokens |
| 1410 | chunk_plan = swa_metadata.get_prefill_chunk_plan( |
| 1411 | compress_ratio=self.compress_ratio, |
| 1412 | prefill_chunk_size=self.PREFILL_CHUNK_SIZE, |
| 1413 | has_compressed=not swa_only, |
| 1414 | ) |
| 1415 | assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" |
| 1416 | |
| 1417 | workspace = current_workspace_manager().get_simultaneous( |
| 1418 | *self._prefill_workspace_shapes(M, num_prefill_tokens, q) |
| 1419 | ) |
| 1420 | kv_rows = workspace[0].view(-1, q.shape[-1]) |
| 1421 | if mxfp8_out is not None: |
| 1422 | output = workspace[1] |
| 1423 | assert output is not None |
| 1424 | for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan: |
| 1425 | chunk_size = chunk_end - chunk_start |
| 1426 | kv = kv_rows[: chunk_size * chunk_M].view(chunk_size, chunk_M, -1) |
| 1427 | if not swa_only: |
| 1428 | assert attn_metadata is not None |
| 1429 | assert compressed_k_cache is not None |
| 1430 | block_table = attn_metadata.block_table[num_decodes:] |
| 1431 | # compressed_k_cache is OCP on every platform (Triton encoder). |
| 1432 | dequantize_and_gather_k_cache( |
| 1433 | kv[:chunk_size], |
| 1434 | compressed_k_cache, |
| 1435 | seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, |
| 1436 | gather_lens=None, |
| 1437 | block_table=block_table[chunk_start:chunk_end], |
| 1438 | block_size=attn_metadata.block_size // self.compress_ratio, |
| 1439 | offset=0, |
| 1440 | use_fnuz=False, |
| 1441 | ) |
| 1442 | |
| 1443 | swa_block_table = swa_metadata.block_table[num_decodes:] |
| 1444 | dequantize_and_gather_k_cache( |
| 1445 | kv[:chunk_size], |
| 1446 | swa_k_cache, |
| 1447 | seq_lens=seq_lens[chunk_start:chunk_end], |
| 1448 | gather_lens=gather_lens[chunk_start:chunk_end], |
| 1449 | block_table=swa_block_table[chunk_start:chunk_end], |
| 1450 | block_size=swa_metadata.block_size, |
| 1451 | offset=chunk_N, |
| 1452 | use_fnuz=current_platform.is_fp8_fnuz(), |
| 1453 | ) |
| 1454 | |
| 1455 | query_start = ( |
| 1456 | query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base |
| 1457 | ) |
| 1458 | query_end = ( |
| 1459 | query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base |
| 1460 | ) |
| 1461 | |
| 1462 | combined_indices, combined_lens = combine_topk_swa_indices( |
| 1463 | topk_indices[query_start:query_end], |
| 1464 | query_start_loc[ |
| 1465 | num_decodes + chunk_start : num_decodes + chunk_end + 1 |
| 1466 | ], |
| 1467 | seq_lens[chunk_start:chunk_end], |
| 1468 | gather_lens[chunk_start:chunk_end], |
| 1469 | self.window_size, |
| 1470 | self.compress_ratio, |
| 1471 | top_k, |
| 1472 | chunk_M, |
| 1473 | chunk_N, |
| 1474 | ) |
| 1475 | rocm_sparse_attn_prefill( |
| 1476 | q=q[query_start:query_end], |
| 1477 | kv=kv.view(-1, 1, q.shape[-1]), |
| 1478 | indices=combined_indices, |
| 1479 | topk_length=combined_lens, |
| 1480 | scale=self.scale, |
| 1481 | head_dim=self.head_dim, |
| 1482 | nope_head_dim=self.nope_head_dim, |
| 1483 | rope_head_dim=self.rope_head_dim, |
| 1484 | attn_sink=self.attn_sink, |
| 1485 | output=output[query_start:query_end], |
| 1486 | ) |
| 1487 | if mxfp8_out is not None: |
| 1488 | rocm_inverse_rope_mxfp8_rows( |
| 1489 | output, |
| 1490 | positions[:num_prefill_tokens], |
| 1491 | self.rotary_emb.cos_sin_cache, |
| 1492 | self.rope_head_dim, |
| 1493 | mxfp8_out[0][:num_prefill_tokens], |
| 1494 | mxfp8_out[1][:num_prefill_tokens], |
| 1495 | ) |
| 1496 | |
| 1497 | def _prefill_workspace_shapes( |
| 1498 | self, M: int, num_prefill_tokens: int, q: torch.Tensor, gather: bool = True |
| 1499 | ) -> list[tuple[tuple[int, ...], torch.dtype]]: |
| 1500 | """The prefill gather buffer, plus bf16 rows for the MXFP8 output. |
| 1501 | |
| 1502 | One ``get_simultaneous`` call: separate calls alias the same memory. |
| 1503 | ``gather=False`` leaves out the gather buffer, which the aiter kernel, |
| 1504 | reading the caches in place, never needs. |
| 1505 | """ |
| 1506 | shapes: list[tuple[tuple[int, ...], torch.dtype]] = [] |
| 1507 | if gather: |
| 1508 | shapes.append(((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)) |
| 1509 | if _ON_GFX950: |
| 1510 | shapes.append( |
| 1511 | ( |
| 1512 | (num_prefill_tokens, self.n_local_heads, self.head_dim), |
| 1513 | torch.bfloat16, |
| 1514 | ) |
| 1515 | ) |
| 1516 | return shapes |
| 1517 | |
| 1518 | def _prefill_topk_ragged( |
| 1519 | self, |
| 1520 | swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, |
| 1521 | attn_metadata: DeepseekV4FlashMLAMetadata, |
| 1522 | compressed_k_cache: torch.Tensor, |
| 1523 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1524 | """Prefill counterpart of ``_decode_topk_ragged``: layers that share an |
| 1525 | index source and a compressed cache reuse one build per step.""" |
| 1526 | assert swa_metadata.token_to_req_indices is not None |
| 1527 | assert swa_metadata.query_start_loc is not None |
| 1528 | assert swa_metadata.seq_lens is not None |
| 1529 | assert swa_metadata.is_valid_token is not None |
| 1530 | assert self.topk_indices_buffer is not None |
| 1531 | assert self._index_source_prefix is not None |
| 1532 | assert self.compressed_cache_prefix is not None |
| 1533 | |
| 1534 | source = self._static_forward_context[self._index_source_prefix] |
| 1535 | if source is self: |
| 1536 | source._prefill_topk_ragged_cache = {} |
| 1537 | cached = source._prefill_topk_ragged_cache.get(self.compressed_cache_prefix) |
| 1538 | if cached is not None: |
| 1539 | return cached |
| 1540 | |
| 1541 | num_decode_tokens = swa_metadata.num_decode_tokens |
| 1542 | num_prefill_tokens = swa_metadata.num_prefill_tokens |
| 1543 | built = build_prefill_topk_ragged_indices( |
| 1544 | self.topk_indices_buffer[ |
| 1545 | num_decode_tokens : num_decode_tokens + num_prefill_tokens |
| 1546 | ], |
| 1547 | swa_metadata.token_to_req_indices, |
| 1548 | swa_metadata.query_start_loc, |
| 1549 | swa_metadata.seq_lens, |
| 1550 | swa_metadata.is_valid_token, |
| 1551 | attn_metadata.block_table, |
| 1552 | block_size=attn_metadata.block_size // self.compress_ratio, |
| 1553 | compress_ratio=self.compress_ratio, |
| 1554 | num_compressed=(self.max_model_len + self.compress_ratio - 1) |
| 1555 | // self.compress_ratio, |
| 1556 | token_offset=num_decode_tokens, |
| 1557 | num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1], |
| 1558 | ) |
| 1559 | source._prefill_topk_ragged_cache[self.compressed_cache_prefix] = built |
| 1560 | return built |
| 1561 | |
| 1562 | def _forward_prefill_aiter( |
| 1563 | self, |
| 1564 | q: torch.Tensor, |
| 1565 | positions: torch.Tensor, |
| 1566 | compressed_k_cache: torch.Tensor | None, |
| 1567 | swa_k_cache: torch.Tensor, |
| 1568 | output: torch.Tensor | None, |
| 1569 | attn_metadata: DeepseekV4FlashMLAMetadata | None, |
| 1570 | swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, |
| 1571 | mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 1572 | ) -> None: |
| 1573 | """Prefill straight off the paged caches, with no bf16 gather. |
| 1574 | |
| 1575 | With ``mxfp8_out`` the kernel writes bf16 rows to the workspace, which |
| 1576 | are then inverse-RoPE'd and quantized into it as ``_forward_prefill`` |
| 1577 | does. |
| 1578 | """ |
| 1579 | num_prefill_tokens = swa_metadata.num_prefill_tokens |
| 1580 | assert swa_metadata.prefill_swa_ragged_indices is not None |
| 1581 | assert swa_metadata.prefill_swa_ragged_indptr is not None |
| 1582 | |
| 1583 | topk_indices = None |
| 1584 | topk_indptr = None |
| 1585 | if attn_metadata is None: |
| 1586 | compressed_k_cache = None |
| 1587 | else: |
| 1588 | assert compressed_k_cache is not None |
| 1589 | topk_indices, topk_indptr = self._prefill_topk_ragged( |
| 1590 | swa_metadata, attn_metadata, compressed_k_cache |
| 1591 | ) |
| 1592 | |
| 1593 | if mxfp8_out is not None: |
| 1594 | output = self._mxfp8_bf16_rows(num_prefill_tokens, q) |
| 1595 | assert output is not None |
| 1596 | self._aiter_sparse_mla( |
| 1597 | q[:num_prefill_tokens], |
| 1598 | output[:num_prefill_tokens], |
| 1599 | swa_k_cache, |
| 1600 | swa_metadata.prefill_swa_ragged_indices, |
| 1601 | swa_metadata.prefill_swa_ragged_indptr, |
| 1602 | compressed_k_cache, |
| 1603 | topk_indices, |
| 1604 | topk_indptr, |
| 1605 | ) |
| 1606 | if mxfp8_out is not None: |
| 1607 | rocm_inverse_rope_mxfp8_rows( |
| 1608 | output[:num_prefill_tokens], |
| 1609 | positions[:num_prefill_tokens], |
| 1610 | self.rotary_emb.cos_sin_cache, |
| 1611 | self.rope_head_dim, |
| 1612 | mxfp8_out[0][:num_prefill_tokens], |
| 1613 | mxfp8_out[1][:num_prefill_tokens], |
| 1614 | ) |
| 1615 | |
| 1616 | def _mxfp8_bf16_rows(self, num_tokens: int, q: torch.Tensor) -> torch.Tensor: |
| 1617 | """bf16 attention rows for the aiter kernel to write when the layer's |
| 1618 | output is MXFP8, taken from the workspace the warmup reserved.""" |
| 1619 | shapes = self._prefill_workspace_shapes(0, num_tokens, q, gather=False) |
| 1620 | assert shapes, "an MXFP8 attention output is gfx950-only" |
| 1621 | return current_workspace_manager().get_simultaneous(*shapes)[0] |
| 1622 | |
| 1623 | def _aiter_sparse_mla( |
| 1624 | self, |
| 1625 | q: torch.Tensor, |
| 1626 | output: torch.Tensor, |
| 1627 | swa_k_cache: torch.Tensor, |
| 1628 | swa_indices: torch.Tensor, |
| 1629 | swa_indptr: torch.Tensor, |
| 1630 | compressed_k_cache: torch.Tensor | None, |
| 1631 | topk_indices: torch.Tensor | None, |
| 1632 | topk_indptr: torch.Tensor | None, |
| 1633 | ) -> None: |
| 1634 | from vllm._aiter_ops import rocm_aiter_ops |
| 1635 | |
| 1636 | # Both fp8 caches are read in place: the SWA window as the main segment, |
| 1637 | # the top-k compressed tokens as the extra one. |
| 1638 | rocm_aiter_ops.triton_sparse_mla_fwd( |
| 1639 | q, |
| 1640 | swa_k_cache, |
| 1641 | output, |
| 1642 | self.scale, |
| 1643 | swa_indptr, |
| 1644 | swa_indices, |
| 1645 | attn_sink=self.attn_sink[: q.shape[1]], |
| 1646 | extra_kv_buffer=compressed_k_cache, |
| 1647 | extra_kv_indptr=topk_indptr, |
| 1648 | extra_kv_indices=topk_indices, |
| 1649 | ) |