A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 21:00 UTC
[ROCm] Add AMD GPU support on Deepseek v3.2 and SparseMLA (#26670)
Signed-off-by: ganyi <ygan@amd.com>
Pleaplusone committed
Nov 20, 2025 at 18:54 UTC
06c20c9904644d8f65523bb747756b2eae706b8e
9 files changed
+583
-15
csrc/cache_kernels.cu
+4
@@ -552,7 +552,11 @@ __global__ void indexer_k_quant_and_cache_kernel(
552
#ifndef USE_ROCM
553
__syncwarp();
554
#endif
555
+#if defined(__gfx942__)
556
+ float scale = fmaxf(amax, 1e-4) / 224.0f;
557
+#else
558
float scale = fmaxf(amax, 1e-4) / 448.0f;
559
+#endif
560
if (use_ue8m0) {
561
scale = exp2f(ceilf(log2f(scale)));
562
}
vllm/attention/ops/rocm_aiter_mla_sparse.py
new
+210
@@ -0,0 +1,210 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import importlib
4
+from functools import lru_cache
5
+
6
+import torch
7
+
8
+from vllm._aiter_ops import rocm_aiter_ops
9
+from vllm.logger import init_logger
10
+from vllm.platforms import current_platform
11
+
12
+logger = init_logger(__name__)
13
+
14
+
15
+# Take from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L84
16
+def fp8_mqa_logits_torch(
17
+ q: torch.Tensor,
18
+ kv: tuple[torch.Tensor, torch.Tensor],
19
+ weights: torch.Tensor,
20
+ cu_seqlen_ks: torch.Tensor,
21
+ cu_seqlen_ke: torch.Tensor,
22
+) -> torch.Tensor:
23
+ """Compute FP8 MQA logits for a single sequence without KV paging.
24
+
25
+ Args:
26
+ q: Query tensor of shape [M, H, D]. Casted to
27
+ `torch.float8_e4m3fn` by caller.
28
+ kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
29
+ dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
30
+ [N, 1]) with dtype `torch.float32`.
31
+ weights: weights of shape [M, H], dtype `torch.float32`.
32
+ cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
33
+ shape [M], dtype int32.
34
+ cu_seqlen_ke: End indices (exclusive) for valid K per query position,
35
+ shape [M], dtype int32.
36
+
37
+ Returns:
38
+ Logits tensor of shape [M, N], dtype `torch.float32`.
39
+ """
40
+ kv, scale = kv
41
+ seq_len_kv = kv.shape[0]
42
+ k = kv.to(torch.bfloat16)
43
+ q = q.to(torch.bfloat16)
44
+
45
+ mask_lo = (
46
+ torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None]
47
+ )
48
+ mask_hi = (
49
+ torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
50
+ )
51
+ mask = mask_lo & mask_hi
52
+
53
+ score = torch.einsum("mhd,nd->hmn", q, k).float() * scale
54
+ logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
55
+ logits = logits.masked_fill(~mask, float("-inf"))
56
+
57
+ return logits
58
+
59
+
60
+def rocm_fp8_mqa_logits(
61
+ q: torch.Tensor,
62
+ kv: tuple[torch.Tensor, torch.Tensor],
63
+ weights: torch.Tensor,
64
+ cu_seqlen_ks: torch.Tensor,
65
+ cu_seqlen_ke: torch.Tensor,
66
+) -> torch.Tensor:
67
+ """Compute FP8 MQA logits for a single sequence without KV paging.
68
+
69
+ Args:
70
+ q: Query tensor of shape [M, H, D]. Casted to
71
+ `torch.float8_e4m3fn` by caller.
72
+ kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
73
+ dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
74
+ [N, 1]) with dtype `torch.float32`.
75
+ weights: weights of shape [M, H], dtype `torch.float32`.
76
+ cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
77
+ shape [M], dtype int32.
78
+ cu_seqlen_ke: End indices (exclusive) for valid K per query position,
79
+ shape [M], dtype int32.
80
+
81
+ Returns:
82
+ Logits tensor of shape [M, N], dtype `torch.float32`.
83
+ """
84
+
85
+ # TODO(ganyi): Temporarily workaround, will remove the module check and reference
86
+ # path after aiter merge this kernel into main
87
+ @lru_cache
88
+ def has_mqa_logits_module():
89
+ return importlib.util.find_spec("aiter.ops.triton.fp8_mqa_logits") is not None
90
+
91
+ if rocm_aiter_ops.is_enabled() and has_mqa_logits_module():
92
+ from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
93
+
94
+ kv, scale = kv
95
+ return fp8_mqa_logits(q, kv, scale, weights, cu_seqlen_ks, cu_seqlen_ke)
96
+ else:
97
+ return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
98
+
99
+
100
+# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156
101
+def fp8_paged_mqa_logits_torch(
102
+ q: torch.Tensor,
103
+ kv_cache: torch.Tensor,
104
+ weights: torch.Tensor,
105
+ context_lens: torch.Tensor,
106
+ block_tables: torch.Tensor,
107
+ max_model_len: int,
108
+):
109
+ from vllm.utils.math_utils import cdiv
110
+
111
+ fp8_dtype = current_platform.fp8_dtype()
112
+ batch_size, next_n, _, dim = q.size()
113
+ kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
114
+ scale = scale.contiguous().view(torch.float)
115
+ q = q.float()
116
+ kv_cache = kv_cache.view(fp8_dtype).float() * scale
117
+ num_block, block_size, _, dim = kv_cache.size()
118
+ logits = torch.full(
119
+ [batch_size * next_n, max_model_len],
120
+ float("-inf"),
121
+ device=q.device,
122
+ dtype=torch.float32,
123
+ )
124
+ context_lens = context_lens.tolist()
125
+ for i in range(batch_size):
126
+ context_len = context_lens[i]
127
+ q_offsets = torch.arange(context_len - next_n, context_len, device="cuda")
128
+ weight_slice = (
129
+ weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
130
+ )
131
+ for block_rk in range(cdiv(context_len, block_size)):
132
+ block_idx = block_tables[i][block_rk]
133
+ qx, kx = q[i], kv_cache[block_idx]
134
+ k_offsets = torch.arange(
135
+ block_rk * block_size, (block_rk + 1) * block_size, device="cuda"
136
+ )
137
+ mask = (k_offsets[None, :] < context_len) & (
138
+ k_offsets[None, :] <= q_offsets[:, None]
139
+ )
140
+ s = torch.where(
141
+ mask[None, :, :],
142
+ (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(
143
+ logits.dtype
144
+ ),
145
+ float("-inf"),
146
+ )
147
+ s = torch.relu(s) * weight_slice[..., None]
148
+ s = s.sum(dim=0)
149
+ logits[
150
+ i * next_n : (i + 1) * next_n,
151
+ block_rk * block_size : (block_rk + 1) * block_size,
152
+ ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf"))
153
+ return logits
154
+
155
+
156
+def rocm_fp8_paged_mqa_logits(
157
+ q_fp8: torch.Tensor,
158
+ kv_cache_fp8: torch.Tensor,
159
+ weights: torch.Tensor,
160
+ context_lens: torch.Tensor,
161
+ block_tables: torch.Tensor,
162
+ schedule_metadata: torch.Tensor,
163
+ max_model_len: int,
164
+) -> torch.Tensor:
165
+ """Compute FP8 MQA logits using paged KV-cache.
166
+
167
+ Args:
168
+ q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to
169
+ `torch.float8_e4m3fn` by caller.
170
+ kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape
171
+ [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
172
+ 4 bytes per (block,pos) store the `float` dequant scale.
173
+ weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
174
+ context_lens: Tensor of shape [B], dtype int32; effective context length
175
+ for each batch element.
176
+ block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
177
+ block indices to physical blocks in the paged cache.
178
+ schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
179
+ used to distribute work across SMs.
180
+ max_model_len: Maximum sequence length used to size the logits output.
181
+
182
+ Returns:
183
+ Logits tensor of shape [B * next_n, max_model_len], dtype
184
+ `torch.float32`.
185
+ """
186
+
187
+ if rocm_aiter_ops.is_enabled():
188
+ from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits_stage1
189
+
190
+ batch_size, next_n, heads, _ = q_fp8.shape
191
+ out_qk = torch.full(
192
+ (heads, batch_size * next_n, max_model_len),
193
+ float("-inf"),
194
+ device="cuda",
195
+ dtype=torch.float32,
196
+ )
197
+ deepgemm_fp8_paged_mqa_logits_stage1(
198
+ q_fp8,
199
+ kv_cache_fp8,
200
+ weights,
201
+ out_qk,
202
+ context_lens,
203
+ block_tables,
204
+ max_model_len,
205
+ )
206
+ return out_qk.sum(dim=0)
207
+ else:
208
+ return fp8_paged_mqa_logits_torch(
209
+ q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
210
+ )
vllm/model_executor/models/deepseek_v2.py
+18
-4
@@ -594,6 +594,7 @@ def sparse_attn_indexer(
594
) -> torch.Tensor:
595
# careful! this will be None in dummy run
596
attn_metadata = get_forward_context().attn_metadata
597
+ fp8_dtype = current_platform.fp8_dtype()
598
# assert isinstance(attn_metadata, dict)
599
if not isinstance(attn_metadata, dict):
600
return sparse_attn_indexer_fake(
@@ -633,7 +634,7 @@ def sparse_attn_indexer(
634
k_fp8 = torch.empty(
635
[chunk.total_seq_lens, head_dim],
636
device=k.device,
636
- dtype=torch.float8_e4m3fn,
637
+ dtype=fp8_dtype,
638
)
639
k_scale = torch.empty(
640
[chunk.total_seq_lens, 4],
@@ -647,7 +648,12 @@ def sparse_attn_indexer(
648
chunk.block_table,
649
chunk.cu_seq_lens,
650
)
650
- logits = fp8_mqa_logits(
651
+ fp8_mqa_logits_func = fp8_mqa_logits
652
+ if current_platform.is_rocm():
653
+ from vllm.attention.ops.rocm_aiter_mla_sparse import rocm_fp8_mqa_logits
654
+
655
+ fp8_mqa_logits_func = rocm_fp8_mqa_logits
656
+ logits = fp8_mqa_logits_func(
657
q_fp8[chunk.token_start : chunk.token_end],
658
(k_fp8, k_scale.view(torch.float32)),
659
weights[chunk.token_start : chunk.token_end],
@@ -692,7 +698,14 @@ def sparse_attn_indexer(
698
next_n = padded_q_fp8_decode_tokens.shape[1]
699
assert batch_size == decode_metadata.seq_lens.shape[0]
700
num_padded_tokens = batch_size * next_n
695
- logits = fp8_paged_mqa_logits(
701
+ fp8_paged_mqa_logits_func = fp8_paged_mqa_logits
702
+ if current_platform.is_rocm():
703
+ from vllm.attention.ops.rocm_aiter_mla_sparse import (
704
+ rocm_fp8_paged_mqa_logits,
705
+ )
706
+
707
+ fp8_paged_mqa_logits_func = rocm_fp8_paged_mqa_logits
708
+ logits = fp8_paged_mqa_logits_func(
709
padded_q_fp8_decode_tokens,
710
kv_cache,
711
weights[:num_padded_tokens],
@@ -749,7 +762,8 @@ def sparse_attn_indexer_fake(
762
_flattened_kv = torch.empty(
763
[total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
764
)
752
- _k_fp8 = _flattened_kv[..., :head_dim].view(torch.float8_e4m3fn).contiguous()
765
+ fp8_dtype = current_platform.fp8_dtype()
766
+ _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()
767
_k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous()
768
return topk_indices_buffer
769
vllm/platforms/rocm.py
+12
-1
@@ -225,7 +225,18 @@ class RocmPlatform(Platform):
225
from vllm.attention.backends.registry import AttentionBackendEnum
226
227
if use_sparse:
228
- raise NotImplementedError("Sparse Attention is not supported on ROCm.")
228
+ if kv_cache_dtype.startswith("fp8"):
229
+ raise ValueError(
230
+ "ROCMAiterMLASparseBackend doesn't support fp8 kv_cache_dtype."
231
+ )
232
+ assert block_size == 1, (
233
+ "Sparse MLA backend on ROCm only supports block size 1 for now."
234
+ )
235
+ logger.info_once("Using Sparse MLA backend on V1 engine.")
236
+ return (
237
+ "vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse."
238
+ "ROCMAiterMLASparseBackend"
239
+ )
240
241
if use_mla:
242
if selected_backend is None:
vllm/utils/deep_gemm.py
+3
-2
@@ -325,6 +325,7 @@ DEFAULT_BLOCK_SIZE = [128, 128]
325
def per_block_cast_to_fp8(
326
x: torch.Tensor, block_size: list[int] = DEFAULT_BLOCK_SIZE, use_ue8m0: bool = False
327
) -> tuple[torch.Tensor, torch.Tensor]:
328
+ fp8_dtype = current_platform.fp8_dtype()
329
assert x.dim() == 2
330
m, n = x.shape
331
block_m, block_n = block_size
@@ -334,9 +335,9 @@ def per_block_cast_to_fp8(
335
x_padded[:m, :n] = x
336
x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
337
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
337
- sf = x_amax / 448.0
338
+ sf = x_amax / 224.0 if current_platform.is_fp8_fnuz() else x_amax / 448.0
339
sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf
339
- x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
340
+ x_scaled = (x_view * (1.0 / sf)).to(fp8_dtype)
341
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
342
x_view.size(0), x_view.size(2)
343
)
vllm/v1/attention/backends/mla/flashmla_sparse.py
+1
-1
@@ -168,7 +168,7 @@ def _convert_req_index_to_global_index_kernel(
168
inblock_off = tok % BLOCK_SIZE
169
170
# Guard block_table access
171
- valid_block = block_id < max_num_blocks_per_req
171
+ valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
172
bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
173
base = tl.load(bt_ptr, mask=valid_block, other=0)
174
vllm/v1/attention/backends/mla/indexer.py
+9
-6
@@ -11,7 +11,8 @@ from vllm.attention.backends.abstract import (
11
)
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
14
-from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata
14
+from vllm.platforms import current_platform
15
+from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata, is_deep_gemm_supported
16
from vllm.v1.attention.backends.utils import (
17
AttentionCGSupport,
18
AttentionMetadataBuilder,
@@ -23,7 +24,9 @@ logger = init_logger(__name__)
24
25
26
class DeepseekV32IndexerBackend(AttentionBackend):
26
- supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]
27
+ supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [
28
+ 1 if current_platform.is_rocm() else 64
29
+ ]
30
31
@classmethod
32
def get_supported_head_sizes(cls) -> list[int]:
@@ -328,10 +331,10 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
331
requires_padding = (decode_lens_cpu.max() > decode_lens_cpu.min()).item()
332
333
seq_lens = common_attn_metadata.seq_lens[:num_decodes]
331
-
332
- self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
333
- seq_lens, self.kv_cache_spec.block_size, self.num_sms
334
- )
334
+ if is_deep_gemm_supported():
335
+ self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
336
+ seq_lens, self.kv_cache_spec.block_size, self.num_sms
337
+ )
338
decode_metadata = DeepSeekV32IndexerDecodeMetadata(
339
block_table=common_attn_metadata.block_table_tensor[:num_decodes, ...],
340
seq_lens=common_attn_metadata.seq_lens[:num_decodes],
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
new
+325
@@ -0,0 +1,325 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+from typing import TYPE_CHECKING, ClassVar, Optional
6
+
7
+import numpy as np
8
+import torch
9
+
10
+from vllm import _custom_ops as ops
11
+from vllm._aiter_ops import rocm_aiter_ops
12
+from vllm.attention.backends.abstract import (
13
+ AttentionBackend,
14
+ AttentionLayer,
15
+ AttentionMetadata,
16
+)
17
+from vllm.attention.backends.utils import get_mla_dims
18
+from vllm.config import VllmConfig
19
+from vllm.logger import init_logger
20
+from vllm.v1.attention.backends.mla.common import (
21
+ MLACommonBaseImpl,
22
+)
23
+from vllm.v1.attention.backends.mla.flashmla_sparse import (
24
+ triton_convert_req_index_to_global_index,
25
+)
26
+from vllm.v1.attention.backends.utils import (
27
+ AttentionCGSupport,
28
+ AttentionMetadataBuilder,
29
+ CommonAttentionMetadata,
30
+)
31
+from vllm.v1.kv_cache_interface import AttentionSpec
32
+
33
+if TYPE_CHECKING:
34
+ from vllm.model_executor.models.deepseek_v2 import Indexer
35
+logger = init_logger(__name__)
36
+
37
+
38
+class ROCMAiterMLASparseBackend(AttentionBackend):
39
+ accept_output_buffer: bool = True
40
+
41
+ @staticmethod
42
+ def get_name() -> str:
43
+ return "ROCM_AITER_MLA_SPARSE"
44
+
45
+ @staticmethod
46
+ def get_metadata_cls() -> type[AttentionMetadata]:
47
+ return ROCMAiterMLASparseMetadata
48
+
49
+ @staticmethod
50
+ def get_builder_cls() -> type["ROCMAiterMLASparseMetadataBuilder"]:
51
+ return ROCMAiterMLASparseMetadataBuilder
52
+
53
+ @staticmethod
54
+ def get_impl_cls() -> type["ROCMAiterMLASparseImpl"]:
55
+ return ROCMAiterMLASparseImpl
56
+
57
+ @staticmethod
58
+ def get_kv_cache_shape(
59
+ num_blocks: int,
60
+ block_size: int,
61
+ num_kv_heads: int, # assumed to be 1 for MLA
62
+ head_size: int,
63
+ cache_dtype_str: str = "auto",
64
+ ) -> tuple[int, ...]:
65
+ return (num_blocks, block_size, head_size)
66
+
67
+ @classmethod
68
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
69
+ return [torch.bfloat16]
70
+
71
+ @classmethod
72
+ def get_supported_head_sizes(cls) -> list[int]:
73
+ return [576]
74
+
75
+
76
+@dataclass
77
+class ROCMAiterMLASparseMetadata:
78
+ num_reqs: int
79
+ max_query_len: int
80
+ max_seq_len: int
81
+
82
+ num_actual_tokens: int # Number of tokens excluding padding.
83
+ query_start_loc: torch.Tensor
84
+ slot_mapping: torch.Tensor
85
+
86
+ block_table: torch.Tensor
87
+ req_id_per_token: torch.Tensor
88
+ block_size: int = 1
89
+ topk_tokens: int = 2048
90
+
91
+
92
+@dataclass
93
+class ROCMAiterMLASparseMetadataBuilder(
94
+ AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]
95
+):
96
+ cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
97
+
98
+ def __init__(
99
+ self,
100
+ kv_cache_spec: AttentionSpec,
101
+ layer_names: list[str],
102
+ vllm_config: VllmConfig,
103
+ device: torch.device,
104
+ ):
105
+ self.kv_cache_spec = kv_cache_spec
106
+ self.model_config = vllm_config.model_config
107
+ parallel_config = vllm_config.parallel_config
108
+ self.device = device
109
+
110
+ self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
111
+ self.mla_dims = get_mla_dims(self.model_config)
112
+ self.topk_tokens = vllm_config.model_config.hf_config.index_topk
113
+ self.topk_tokens_tensor = torch.tensor(
114
+ [self.topk_tokens], device=device, dtype=torch.int32
115
+ )
116
+ self.max_model_len_tensor = torch.tensor(
117
+ [self.model_config.max_model_len], device=device, dtype=torch.int32
118
+ )
119
+ # this is ignored by `flash_mla_with_kvcache` if indices not None
120
+ self.dummy_block_table = torch.empty(
121
+ (1, 1), dtype=torch.int32, device=self.device
122
+ )
123
+
124
+ self.req_id_per_token_buffer = torch.empty(
125
+ (vllm_config.scheduler_config.max_num_batched_tokens,),
126
+ dtype=torch.int32,
127
+ device=device,
128
+ )
129
+
130
+ def build(
131
+ self,
132
+ common_prefix_len: int,
133
+ common_attn_metadata: CommonAttentionMetadata,
134
+ fast_build: bool = False,
135
+ ) -> ROCMAiterMLASparseMetadata:
136
+ num_tokens = common_attn_metadata.num_actual_tokens
137
+ starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
138
+ seg_lengths = np.diff(starts)
139
+ req_id_per_token = np.repeat(
140
+ np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
141
+ )
142
+ # Zero-fill for cudagraphs
143
+ self.req_id_per_token_buffer.fill_(0)
144
+ self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
145
+ torch.from_numpy(req_id_per_token), non_blocking=True
146
+ )
147
+ req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
148
+
149
+ metadata = ROCMAiterMLASparseMetadata(
150
+ num_reqs=common_attn_metadata.num_reqs,
151
+ max_query_len=common_attn_metadata.max_query_len,
152
+ max_seq_len=common_attn_metadata.max_seq_len,
153
+ num_actual_tokens=common_attn_metadata.num_actual_tokens,
154
+ query_start_loc=common_attn_metadata.query_start_loc,
155
+ slot_mapping=common_attn_metadata.slot_mapping,
156
+ block_table=common_attn_metadata.block_table_tensor,
157
+ req_id_per_token=req_id_per_token,
158
+ block_size=self.kv_cache_spec.block_size,
159
+ topk_tokens=self.topk_tokens,
160
+ )
161
+ return metadata
162
+
163
+
164
+# Take from
165
+# https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla_prefill.py#L72
166
+def reference_mla_sparse_prefill(
167
+ q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, sm_scale: float, d_v: int
168
+) -> tuple[torch.Tensor, torch.Tensor]:
169
+ import math
170
+
171
+ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor:
172
+ return torch.logsumexp(a * math.log(2), dim=dim) * math.log2(math.e)
173
+
174
+ skv = kv.shape[0]
175
+ sq = q.shape[0]
176
+ topk = indices.shape[-1]
177
+ dqk = q.shape[-1]
178
+ indices = indices[:, 0, :] # [s_q, topk]
179
+ invalid_indices_mask = (indices < 0) | (indices >= skv)
180
+ indices[invalid_indices_mask] = 0
181
+ qs = q # [s_q, h_q, d_qk]
182
+ kvs = kv[:, 0, :][indices].view(sq, topk, dqk) # [s_q, topk, d_qk]
183
+
184
+ attn_score = (qs @ kvs.transpose(1, 2)).float() # [s_q, h_q, topk]
185
+ attn_score.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf"))
186
+ attn_score *= sm_scale * math.log2(math.e)
187
+ lse = log2sumexp2(attn_score, dim=-1) # [s_q, h_q]
188
+ attn_score = torch.exp2(attn_score - lse.unsqueeze(-1)) # [s_q, h_q, topk]
189
+ result = attn_score.to(q.dtype) @ kvs[:, :, :d_v]
190
+ return (result, lse)
191
+
192
+
193
+class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):
194
+ def __init__(
195
+ self,
196
+ num_heads: int,
197
+ head_size: int,
198
+ scale: float,
199
+ num_kv_heads: int,
200
+ alibi_slopes: list[float] | None,
201
+ sliding_window: int | None,
202
+ kv_cache_dtype: str,
203
+ logits_soft_cap: float | None,
204
+ attn_type: str,
205
+ kv_sharing_target_layer_name: str | None,
206
+ # MLA Specific Arguments
207
+ topk_indice_buffer: torch.Tensor | None = None,
208
+ indexer: Optional["Indexer"] = None,
209
+ **mla_args,
210
+ ) -> None:
211
+ super().__init__(
212
+ num_heads,
213
+ head_size,
214
+ scale,
215
+ num_kv_heads,
216
+ alibi_slopes,
217
+ sliding_window,
218
+ kv_cache_dtype,
219
+ logits_soft_cap,
220
+ attn_type,
221
+ kv_sharing_target_layer_name,
222
+ **mla_args,
223
+ )
224
+ self.softmax_scale = scale
225
+ assert indexer is not None
226
+ self.topk_indices_buffer = indexer.topk_indices_buffer
227
+ self.is_fp8bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
228
+
229
+ def _forward_bf16_kv(
230
+ self,
231
+ q: torch.Tensor,
232
+ kv_c_and_k_pe_cache: torch.Tensor,
233
+ topk_indices: torch.Tensor,
234
+ attn_metadata: ROCMAiterMLASparseMetadata,
235
+ ) -> torch.Tensor:
236
+ num_tokens = q.shape[0]
237
+ kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
238
+ -1, 1, kv_c_and_k_pe_cache.shape[-1]
239
+ )
240
+
241
+ topk_indices = topk_indices.view(num_tokens, 1, -1)
242
+ output = reference_mla_sparse_prefill(
243
+ q, kv_c_and_k_pe_cache, topk_indices, self.softmax_scale, 512
244
+ )[0]
245
+ return output[:, : self.num_heads, :]
246
+
247
+ def forward(
248
+ self,
249
+ layer: AttentionLayer,
250
+ q: torch.Tensor,
251
+ k_c_normed: torch.Tensor, # key in unified attn
252
+ k_pe: torch.Tensor, # value in unified attn
253
+ kv_cache: torch.Tensor,
254
+ attn_metadata: ROCMAiterMLASparseMetadata,
255
+ output: torch.Tensor | None = None,
256
+ output_scale: torch.Tensor | None = None,
257
+ output_block_scale: torch.Tensor | None = None,
258
+ ) -> torch.Tensor:
259
+ # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use
260
+ # MQA 576/512 approach for both prefill and decode
261
+
262
+ assert output is not None, "Output tensor must be provided."
263
+
264
+ if output_scale is not None or output_block_scale is not None:
265
+ raise NotImplementedError(
266
+ "fused output quantization is not yet supported for ROCMAiterMLASparse"
267
+ )
268
+
269
+ if attn_metadata is None:
270
+ # The zero fill is required when used with DP + EP
271
+ # to ensure all ranks within a DP group compute the
272
+ # same expert outputs.
273
+ return output.fill_(0)
274
+
275
+ num_actual_toks = attn_metadata.num_actual_tokens
276
+
277
+ # Inputs and outputs may be padded for CUDA graphs
278
+
279
+ q = q[:num_actual_toks, ...]
280
+ k_c_normed = k_c_normed[:num_actual_toks, ...]
281
+ k_pe = k_pe[:num_actual_toks, ...]
282
+
283
+ q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
284
+ # Convert from (B, N, P) to (N, B, P)
285
+ q_nope = q_nope.transpose(0, 1)
286
+ if self.is_fp8bmm_enabled:
287
+ # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
288
+ ql_nope = rocm_aiter_ops.triton_fp8_bmm(
289
+ q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
290
+ )
291
+ else:
292
+ # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
293
+ ql_nope = torch.bmm(q_nope, self.W_UK_T)
294
+ # Convert from (N, B, L) to (B, N, L)
295
+ ql_nope = ql_nope.transpose(0, 1)
296
+
297
+ topk_indices = self.topk_indices_buffer[:num_actual_toks]
298
+
299
+ topk_indices_global = triton_convert_req_index_to_global_index(
300
+ attn_metadata.req_id_per_token,
301
+ attn_metadata.block_table,
302
+ topk_indices,
303
+ BLOCK_SIZE=attn_metadata.block_size,
304
+ NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
305
+ )
306
+
307
+ q = torch.cat([ql_nope, q_pe], dim=-1)
308
+
309
+ # write the latent and rope to kv cache
310
+ if kv_cache.numel() > 0:
311
+ ops.concat_and_cache_mla(
312
+ k_c_normed,
313
+ k_pe.squeeze(1),
314
+ kv_cache,
315
+ attn_metadata.slot_mapping.flatten(),
316
+ kv_cache_dtype=self.kv_cache_dtype,
317
+ scale=layer._k_scale,
318
+ )
319
+
320
+ attn_out = self._forward_bf16_kv(
321
+ q, kv_cache, topk_indices_global, attn_metadata
322
+ )
323
+
324
+ self._v_up_proj(attn_out, out=output[:num_actual_toks])
325
+ return output
vllm/v1/worker/utils.py
+1
-1
@@ -316,7 +316,7 @@ def bind_kv_cache(
316
# TODO - analyze where runner_kv_caches is used and the right
317
# way to ensure it properly reflects multiple attention layers
318
# in the same decoder block.
319
- if current_platform.is_cuda() or current_platform.is_xpu():
319
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
320
# We know that the GPU runner is not impacted by this
321
# case. Some test code depends on runner_kv_caches, but
322
# not in a way that's impacted by ignoring this.