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 4, 10:00 UTC
[ROCm][Deepseek] dsv3.2 further optimization (#41217)
Signed-off-by: ganyi <ygan@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com>
Pleaplusone committed
May 1, 2026 at 22:06 UTC
bc635fad2389e228a31d6bc6e698caf53d395e13
6 files changed
+293
-73
docs/design/attention_backends.md
+1
-1
@@ -216,7 +216,7 @@ configuration.
216
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
217
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
218
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
219
-| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
219
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
220
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
221
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
222
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
vllm/model_executor/models/deepseek_v2.py
+38
-23
@@ -674,30 +674,45 @@ class Indexer(nn.Module):
674
) -> torch.Tensor:
675
q, _ = self.wq_b(qr)
676
q = q.view(-1, self.n_head, self.head_dim)
677
- q_pe, q_nope = torch.split(
678
- q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
679
- )
680
- # Fused wk + weights_proj: one GEMM, then split
681
- kw, _ = self.wk_weights_proj(hidden_states)
682
- k = kw[:, : self.head_dim]
683
- weights = kw[:, self.head_dim :]
684
-
685
- k = self.k_norm(k)
686
- k_pe, k_nope = torch.split(
687
- k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
688
- )
677
690
- q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
691
- # Note: RoPE (NeoX) can introduce extra leading dimensions during compilation
692
- # so we need to reshape back to token-flattened shapes
693
- q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
694
- k_pe = k_pe.reshape(-1, 1, self.rope_dim)
695
-
696
- # `rotary_emb` is shape-preserving; `q_pe` is already
697
- # [num_tokens, n_head, rope_dim].
698
- q = torch.cat([q_pe, q_nope], dim=-1)
699
- # `k_pe` is [num_tokens, 1, rope_dim] (MQA).
700
- k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
678
+ if current_platform.is_rocm():
679
+ # This path should works on all platform, will remove extra
680
+ # branches in the future
681
+ # Fused wk + weights_proj: one GEMM, then split
682
+ kw, _ = self.wk_weights_proj(hidden_states)
683
+ k = kw[:, : self.head_dim]
684
+ weights = kw[:, self.head_dim :]
685
+
686
+ k = self.k_norm(k)
687
+
688
+ rotary_emb(
689
+ positions, q[..., : self.rope_dim], k[..., : self.rope_dim].unsqueeze(1)
690
+ )
691
+ else:
692
+ q_pe, q_nope = torch.split(
693
+ q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
694
+ )
695
+ # Fused wk + weights_proj: one GEMM, then split
696
+ kw, _ = self.wk_weights_proj(hidden_states)
697
+ k = kw[:, : self.head_dim]
698
+ weights = kw[:, self.head_dim :]
699
+
700
+ k = self.k_norm(k)
701
+ k_pe, k_nope = torch.split(
702
+ k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
703
+ )
704
+
705
+ q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
706
+ # Note: RoPE (NeoX) can introduce extra leading dimensions during
707
+ # compilation so we need to reshape back to token-flattened shapes
708
+ q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
709
+ k_pe = k_pe.reshape(-1, 1, self.rope_dim)
710
+
711
+ # `rotary_emb` is shape-preserving; `q_pe` is already
712
+ # [num_tokens, n_head, rope_dim].
713
+ q = torch.cat([q_pe, q_nope], dim=-1)
714
+ # `k_pe` is [num_tokens, 1, rope_dim] (MQA).
715
+ k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
716
717
# we only quant q here since k quant is fused with cache insertion
718
q = q.view(-1, self.head_dim)
vllm/v1/attention/backends/mla/indexer.py
+1
-1
@@ -122,7 +122,7 @@ class DeepseekV32IndexerBackend(AttentionBackend):
122
123
@staticmethod
124
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
125
- return [1 if current_platform.is_rocm() else 64]
125
+ return [1, 64] if current_platform.is_rocm() else [64]
126
127
@classmethod
128
def get_supported_head_sizes(cls) -> list[int]:
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+4
@@ -396,6 +396,7 @@ class AiterMLAHelper:
396
"""
397
398
_AITER_MIN_MLA_HEADS: Final = 16
399
+ _AITER_UNSUPPORTED_HEADS = [32]
400
401
@staticmethod
402
def check_num_heads_validity(num_heads: int):
@@ -419,6 +420,9 @@ class AiterMLAHelper:
420
421
@staticmethod
422
def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
423
+ assert num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS, (
424
+ f"unsupported head_num: {num_heads}"
425
+ )
426
return (
427
q
428
if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+227
-29
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, ClassVar
7
import numpy as np
8
import torch
9
10
+from vllm import _custom_ops as ops
11
from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.config import VllmConfig
13
from vllm.config.cache import CacheDType
@@ -14,6 +15,7 @@ from vllm.logger import init_logger
15
from vllm.model_executor.layers.attention.mla_attention import (
16
get_mla_dims,
17
)
18
+from vllm.platforms import current_platform
19
from vllm.triton_utils import tl, triton
20
from vllm.v1.attention.backend import (
21
AttentionBackend,
@@ -25,9 +27,6 @@ from vllm.v1.attention.backend import (
27
MultipleOf,
28
SparseMLAAttentionImpl,
29
)
28
-from vllm.v1.attention.backends.mla.flashmla_sparse import (
29
- triton_convert_req_index_to_global_index,
30
-)
30
from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
31
AiterMLAHelper,
32
)
@@ -38,6 +37,188 @@ if TYPE_CHECKING:
37
logger = init_logger(__name__)
38
39
40
+@triton.jit
41
+def _convert_req_index_to_global_index_kernel(
42
+ req_id_ptr, # int32 [num_tokens]
43
+ block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]
44
+ token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
45
+ cu_seqlens_ptr, # int32 [num_tokens + 1]
46
+ out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
47
+ # shapes (compile-time where possible)
48
+ max_num_blocks_per_req: tl.constexpr,
49
+ BLOCK_SIZE: tl.constexpr,
50
+ BLOCK_N: tl.constexpr, # tile width along columns
51
+ # strides (in elements)
52
+ bt_stride0,
53
+ bt_stride1,
54
+ ti_stride0,
55
+ ti_stride1,
56
+):
57
+ # program_id(0) -> token_id (row)
58
+ # program_id(1) -> tile index along columns
59
+ token_id = tl.program_id(0)
60
+ tile_id = tl.program_id(1)
61
+
62
+ # Each program covers BLOCK_N consecutive columns
63
+ indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)
64
+
65
+ # Load request id for this token (no mask: grid is exact)
66
+ req = tl.load(req_id_ptr + token_id)
67
+
68
+ # Load cumulative sequence lengths to get starting index of this request
69
+ seq_start = tl.load(cu_seqlens_ptr + token_id)
70
+ seq_end = tl.load(cu_seqlens_ptr + token_id + 1)
71
+
72
+ if tile_id * BLOCK_N + seq_start >= seq_end:
73
+ return
74
+
75
+ # Load token indices for this tile
76
+ ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
77
+ tok = tl.load(ti_ptr) # int32
78
+
79
+ # Only token == -1 should propagate as -1
80
+ is_invalid_tok = tok < 0
81
+
82
+ # Compute block id and in-block offset
83
+ block_id = tok // BLOCK_SIZE
84
+ inblock_off = tok % BLOCK_SIZE
85
+
86
+ # Guard block_table access
87
+ valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
88
+ bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
89
+ base = tl.load(bt_ptr, mask=valid_block, other=0)
90
+
91
+ # # If token == -1 OR block_id OOB, output 0; else base * BLOCK_SIZE + offset
92
+ out_val = tl.where(
93
+ is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off
94
+ )
95
+ out_ptr_ij = out_ptr + seq_start + indice_id
96
+ out_ptr_ij_mask = (seq_start + indice_id) < seq_end
97
+
98
+ # store the results with mask
99
+ tl.store(out_ptr_ij, out_val, mask=out_ptr_ij_mask)
100
+
101
+
102
+def triton_convert_req_index_to_global_index(
103
+ req_id: torch.Tensor, # int32 [num_tokens]
104
+ block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req]
105
+ token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS]
106
+ cu_seqlens: torch.Tensor, # int32 [num_tokens + 1]
107
+ paged_kv_indices: torch.Tensor, # int32 [num_tokens * topk] out_buffer
108
+ BLOCK_SIZE: int = 64,
109
+ NUM_TOPK_TOKENS: int = 2048,
110
+ BLOCK_N: int = 128, # tile width along columns
111
+):
112
+ """
113
+ out[token_id, indice_id] =
114
+ block_table[req_id[token_id],
115
+ token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE
116
+ + token_indices[token_id, indice_id] % BLOCK_SIZE
117
+
118
+ Only when token_indices[token_id, indice_id] == -1 do we output -1.
119
+ For safety, we also output -1 if the derived block_id would be
120
+ out-of-bounds.
121
+ """
122
+ assert req_id.dtype == torch.int32
123
+ assert block_table.dtype == torch.int32
124
+ assert token_indices.dtype == torch.int32
125
+ assert token_indices.shape[1] == NUM_TOPK_TOKENS
126
+ assert NUM_TOPK_TOKENS % BLOCK_N == 0, (
127
+ f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible byBLOCK_N ({BLOCK_N})"
128
+ )
129
+ # print("req_id: ", req_id, flush=True)
130
+ num_tokens = req_id.shape[0]
131
+ _, max_num_blocks_per_req = block_table.shape
132
+ tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N
133
+
134
+ # Ensure contiguous tensors on the same device
135
+ req_id_c = req_id.contiguous()
136
+ block_table_c = block_table.contiguous()
137
+ token_indices_c = token_indices.contiguous()
138
+
139
+ # Strides in elements
140
+ bt_stride0, bt_stride1 = block_table_c.stride()
141
+ ti_stride0, ti_stride1 = token_indices_c.stride()
142
+
143
+ # Exact 2D grid: tokens × column tiles
144
+ grid = (num_tokens, tiles_per_row)
145
+
146
+ _convert_req_index_to_global_index_kernel[grid](
147
+ req_id_c,
148
+ block_table_c,
149
+ token_indices_c,
150
+ cu_seqlens,
151
+ paged_kv_indices,
152
+ # shapes / constexprs
153
+ max_num_blocks_per_req,
154
+ BLOCK_SIZE,
155
+ BLOCK_N,
156
+ # strides
157
+ bt_stride0,
158
+ bt_stride1,
159
+ ti_stride0,
160
+ ti_stride1,
161
+ )
162
+ return
163
+
164
+
165
+@triton.jit
166
+def generate_sparse_seqlen_kernel(
167
+ seq_len_ptr, # [num_seq]
168
+ cu_query_lens_ptr, # [num_seq]
169
+ out_ptr, # [num_query_tokens]
170
+ topk_token: tl.constexpr,
171
+ BLOCK_SIZE: tl.constexpr,
172
+):
173
+ seq_id = tl.program_id(0)
174
+ query_offset = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
175
+ query_start = tl.load(cu_query_lens_ptr + seq_id)
176
+ query_end = tl.load(cu_query_lens_ptr + seq_id + 1)
177
+ if query_start + tl.program_id(1) * BLOCK_SIZE > query_end:
178
+ return
179
+ query_len = query_end - query_start
180
+ query_mask = query_offset + query_start < query_end
181
+ seq_len = tl.load(seq_len_ptr + seq_id)
182
+ # Just return since the out_ptr is zero initialized.
183
+ if seq_len == 0:
184
+ return
185
+ context_start_point = seq_len - query_len
186
+ sparse_seqlen = context_start_point + query_offset
187
+ sparse_seqlen_masked = tl.where(
188
+ sparse_seqlen + 1 < topk_token, sparse_seqlen + 1, topk_token
189
+ )
190
+ tl.store(
191
+ out_ptr + query_start + query_offset, sparse_seqlen_masked, mask=query_mask
192
+ )
193
+
194
+
195
+def generate_sparse_seqlen_triton(
196
+ query_lens: torch.Tensor,
197
+ seq_lens: torch.Tensor,
198
+ cu_query_lens: torch.Tensor,
199
+ topk_token: int,
200
+ num_tokens: int,
201
+ max_query_len: int,
202
+):
203
+ num_seqs = query_lens.size(0)
204
+ # zero initialize the tensor to make sure invalid positions will be zero
205
+ out = torch.zeros([num_tokens], dtype=torch.int32, device=query_lens.device)
206
+ block_size = 64
207
+ num_block_per_row = triton.cdiv(max_query_len, block_size)
208
+ grid = (
209
+ num_seqs,
210
+ num_block_per_row,
211
+ )
212
+ generate_sparse_seqlen_kernel[grid](
213
+ seq_lens,
214
+ cu_query_lens,
215
+ out,
216
+ topk_token,
217
+ block_size,
218
+ )
219
+ return out
220
+
221
+
222
@triton.jit
223
def fetch_id_to_ragged_kernel(
224
in_tensor_ptr, # [num_seq, topk]
@@ -86,11 +267,13 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
267
"auto",
268
"float16",
269
"bfloat16",
270
+ "fp8",
271
+ "fp8_e4m3",
272
]
273
274
@staticmethod
275
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
93
- return [1]
276
+ return [1, 64]
277
278
@staticmethod
279
def get_name() -> str:
@@ -144,7 +327,7 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
327
paged_kv_last_page_len: torch.Tensor
328
paged_kv_indices: torch.Tensor
329
paged_kv_indptr: torch.Tensor
147
- paged_kv_indptr_rest: torch.Tensor
330
+ attn_out_dtype: torch.dtype
331
332
block_size: int = 1
333
topk_tokens: int = 2048
@@ -167,6 +350,7 @@ class ROCMAiterMLASparseMetadataBuilder(
350
):
351
self.kv_cache_spec = kv_cache_spec
352
self.model_config = vllm_config.model_config
353
+ self.model_dtype = vllm_config.model_config.dtype
354
parallel_config = vllm_config.parallel_config
355
self.device = device
356
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -174,9 +358,6 @@ class ROCMAiterMLASparseMetadataBuilder(
358
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
359
self.mla_dims = get_mla_dims(self.model_config)
360
self.topk_tokens = vllm_config.model_config.hf_config.index_topk
177
- self.topk_tokens_tensor = torch.tensor(
178
- [self.topk_tokens], device=device, dtype=torch.int32
179
- )
361
self.max_model_len_tensor = torch.tensor(
362
[self.model_config.max_model_len], device=device, dtype=torch.int32
363
)
@@ -222,18 +403,33 @@ class ROCMAiterMLASparseMetadataBuilder(
403
)
404
# Zero-fill for cudagraphs
405
self.req_id_per_token_buffer.fill_(0)
406
+ self.paged_kv_indices.fill_(0)
407
+ self.paged_kv_indptr.fill_(0)
408
self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
409
torch.from_numpy(req_id_per_token), non_blocking=True
410
)
228
- self.paged_kv_indices.fill_(0)
229
- self.paged_kv_indptr.fill_(0)
411
+ query_lens = (
412
+ common_attn_metadata.query_start_loc[1:]
413
+ - common_attn_metadata.query_start_loc[:-1]
414
+ )
415
+ seq_lens = common_attn_metadata.seq_lens
416
+ sparse_seqlen = generate_sparse_seqlen_triton(
417
+ query_lens,
418
+ seq_lens,
419
+ common_attn_metadata.query_start_loc,
420
+ self.topk_tokens,
421
+ num_tokens,
422
+ common_attn_metadata.max_query_len,
423
+ )
424
+
425
+ torch.cumsum(sparse_seqlen, dim=0, out=self.paged_kv_indptr[1 : num_tokens + 1])
426
+ self.paged_kv_indptr[num_tokens + 1 :].fill_(self.paged_kv_indptr[num_tokens])
427
428
req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
429
qo_indptr = self.qo_indptr[: num_tokens + 1]
430
paged_kv_last_page_len = self.paged_kv_last_page_len[:num_tokens]
234
- paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens]
431
paged_kv_indptr = self.paged_kv_indptr[: num_tokens + 1]
236
- paged_kv_indptr_rest = self.paged_kv_indptr[num_tokens + 1 :]
432
+ paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens]
433
434
metadata = ROCMAiterMLASparseMetadata(
435
num_reqs=common_attn_metadata.num_reqs,
@@ -245,12 +441,12 @@ class ROCMAiterMLASparseMetadataBuilder(
441
block_table=common_attn_metadata.block_table_tensor,
442
req_id_per_token=req_id_per_token,
443
block_size=self.kv_cache_spec.block_size,
444
+ attn_out_dtype=self.model_dtype,
445
topk_tokens=self.topk_tokens,
446
qo_indptr=qo_indptr,
447
paged_kv_last_page_len=paged_kv_last_page_len,
448
paged_kv_indices=paged_kv_indices,
449
paged_kv_indptr=paged_kv_indptr,
253
- paged_kv_indptr_rest=paged_kv_indptr_rest,
450
)
451
return metadata
452
@@ -314,29 +510,20 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
510
assert indexer is not None
511
self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
512
317
- def _forward_bf16_kv(
513
+ def _forward_mla(
514
self,
515
+ layer: AttentionLayer,
516
q: torch.Tensor, # [sq, heads, d_qk]
517
kv_c_and_k_pe_cache: torch.Tensor, # [blocks, heads, d_qk]
321
- topk_indices: torch.Tensor, # [sq, topk]
518
attn_metadata: ROCMAiterMLASparseMetadata,
519
) -> torch.Tensor:
520
num_tokens = q.shape[0]
521
mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
522
output = torch.empty(
523
[num_tokens, mla_num_heads, self.kv_lora_rank],
328
- dtype=q.dtype,
524
+ dtype=attn_metadata.attn_out_dtype,
525
device=q.device,
526
)
331
- seq_len = (topk_indices != -1).sum(dim=-1)
332
- torch.cumsum(seq_len, dim=0, out=attn_metadata.paged_kv_indptr[1:])
333
- attn_metadata.paged_kv_indptr_rest.fill_(attn_metadata.paged_kv_indptr[-1])
334
- fetch_id_to_ragged_triton(
335
- topk_indices,
336
- attn_metadata.paged_kv_indptr,
337
- attn_metadata.paged_kv_indices,
338
- attn_metadata.topk_tokens,
339
- )
527
528
rocm_aiter_ops.mla_decode_fwd(
529
q,
@@ -348,6 +535,8 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
535
attn_metadata.paged_kv_indptr,
536
attn_metadata.paged_kv_indices,
537
attn_metadata.paged_kv_last_page_len,
538
+ q_scale=layer._q_scale,
539
+ kv_scale=layer._k_scale,
540
)
541
542
return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)
@@ -366,23 +555,32 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
555
if isinstance(q, tuple):
556
q = torch.cat(q, dim=-1)
557
369
- num_actual_toks = q.shape[0]
558
+ num_actual_toks = attn_metadata.num_actual_tokens
559
560
# Get topk indices
561
assert self.topk_indices_buffer is not None
562
topk_indices = self.topk_indices_buffer[:num_actual_toks]
563
375
- topk_indices_global = triton_convert_req_index_to_global_index(
564
+ triton_convert_req_index_to_global_index(
565
attn_metadata.req_id_per_token,
566
attn_metadata.block_table,
567
topk_indices,
568
+ attn_metadata.paged_kv_indptr,
569
+ attn_metadata.paged_kv_indices,
570
BLOCK_SIZE=attn_metadata.block_size,
571
NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
572
)
573
574
+ # write the latent and rope to kv cache
575
+ fp8_attention = self.kv_cache_dtype.startswith("fp8")
576
+ if fp8_attention:
577
+ original_q_shape = q.shape
578
+ kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(current_platform.fp8_dtype())
579
+ q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale)
580
+ q = q.view(original_q_shape)
581
mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
384
- attn_out = self._forward_bf16_kv(
385
- mla_padded_q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata
582
+ attn_out = self._forward_mla(
583
+ layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata
584
)
585
586
return attn_out, None
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+22
-19
@@ -13,9 +13,6 @@ from vllm.utils.torch_utils import LayerNameType
13
from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
14
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
15
16
-if current_platform.is_cuda_alike():
17
- from vllm import _custom_ops as ops
18
-
16
17
@triton.jit
18
def _indexer_k_quant_and_cache_kernel(
@@ -97,7 +94,8 @@ def indexer_k_quant_and_cache_triton(
94
# In real layout, we store the first portion as kv cache value
95
# and second portion as kv cache scale
96
kv_cache = kv_cache.view(num_blocks, -1)
100
- kv_cache_value = kv_cache[:, : block_size * head_dim]
97
+ fp8_dtype = current_platform.fp8_dtype()
98
+ kv_cache_value = kv_cache[:, : block_size * head_dim].view(fp8_dtype)
99
kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32)
100
head_tile_size = head_tile_size // kv_cache.element_size()
101
grid = (num_tokens,)
@@ -111,7 +109,7 @@ def indexer_k_quant_and_cache_triton(
109
block_size,
110
num_tokens,
111
head_dim,
114
- "NHD",
112
+ "SHUFFLE",
113
block_tile_size,
114
head_tile_size,
115
IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
@@ -212,7 +210,7 @@ def cp_gather_indexer_k_quant_cache_triton(
210
block_table_stride,
211
k_cache_value.stride(0),
212
k_cache_scale.stride(0),
215
- "NHD",
213
+ "SHUFFLE",
214
head_dim,
215
block_tile_size,
216
head_tile_size,
@@ -325,33 +323,38 @@ def rocm_fp8_paged_mqa_logits(
323
from vllm._aiter_ops import rocm_aiter_ops
324
325
aiter_paged_mqa_logits_module = None
326
+ # if rocm_aiter_ops.is_enabled():
327
+ batch_size, next_n, heads, head_dim = q_fp8.shape
328
+ num_blocks, block_size, _, _ = kv_cache_fp8.shape
329
+
330
if rocm_aiter_ops.is_enabled():
331
aiter_paged_mqa_logits_module = paged_mqa_logits_module()
332
333
if aiter_paged_mqa_logits_module is not None:
332
- deepgemm_fp8_paged_mqa_logits_stage1 = (
333
- aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1
334
+ deepgemm_fp8_paged_mqa_logits = (
335
+ aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
336
)
337
batch_size, next_n, heads, _ = q_fp8.shape
336
- out_qk = torch.full(
337
- (heads, batch_size * next_n, max_model_len),
338
+ out_logits = torch.full(
339
+ [batch_size * next_n, max_model_len],
340
float("-inf"),
341
device="cuda",
342
dtype=torch.float32,
343
)
342
- # TODO: 1. Replace _stage1 and out_qk.sum with another fused variant;
343
- # 2. Remove ChunkQ when AITER PR #2891 merged
344
- deepgemm_fp8_paged_mqa_logits_stage1(
344
+ deepgemm_fp8_paged_mqa_logits(
345
q_fp8,
346
kv_cache_fp8,
347
weights,
348
- out_qk,
348
+ out_logits,
349
context_lens,
350
block_tables,
351
max_model_len,
352
- ChunkQ=heads,
352
+ ChunkK=256,
353
+ Preshuffle=block_size == 64,
354
+ KVBlockSize=block_size,
355
+ WavePerEU=2,
356
)
354
- return out_qk.sum(dim=0)
357
+ return out_logits
358
else:
359
return fp8_paged_mqa_logits_torch(
360
q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
@@ -540,7 +543,7 @@ def rocm_aiter_sparse_attn_indexer(
543
num_tokens = slot_mapping.shape[0]
544
k = k[:num_tokens]
545
543
- ops.indexer_k_quant_and_cache(
546
+ indexer_k_quant_and_cache_triton(
547
k,
548
kv_cache,
549
slot_mapping,
@@ -563,13 +566,13 @@ def rocm_aiter_sparse_attn_indexer(
566
device=k.device,
567
dtype=torch.uint8,
568
)
566
-
567
- ops.cp_gather_indexer_k_quant_cache(
569
+ cp_gather_indexer_k_quant_cache_triton(
570
kv_cache,
571
k_fp8,
572
k_scale,
573
chunk.block_table,
574
chunk.cu_seq_lens,
575
+ chunk.token_to_seq,
576
)
577
578
logits = rocm_fp8_mqa_logits(