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