@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7863fff6e5
Fork 0

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 3, 22:45 UTC

[ROCm][DSv4] implement flash sparse mla with triton kernels (#41812)

Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>

Hexiang Wang committed May 12, 2026 at 00:27 UTC 7863fff6e591161a03a7df68b0879b9fa61cfc0b
6 files changed +1849 -212
tests/kernels/attention/test_rocm_triton_attn_dsv4.py new
+377
@@ -0,0 +1,377 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +
9 +pytestmark = pytest.mark.skipif(
10 + not current_platform.is_rocm(), reason="Only used by ROCm"
11 +)
12 +
13 +NOPE_HEAD_DIM = 448
14 +ROPE_HEAD_DIM = 64
15 +HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
16 +
17 +
18 +def _ref_global_topk_ragged(
19 + topk_indices: torch.Tensor,
20 + token_to_req_indices: torch.Tensor,
21 + block_table: torch.Tensor,
22 + block_size: int,
23 + is_valid_token: torch.Tensor,
24 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
25 + topk = topk_indices.reshape(topk_indices.shape[0], -1)
26 + valid = (topk >= 0) & is_valid_token[:, None]
27 + lens = valid.sum(dim=1, dtype=torch.int32)
28 + indptr = torch.zeros(lens.shape[0] + 1, dtype=torch.int32, device=topk.device)
29 + torch.cumsum(lens, dim=0, out=indptr[1:])
30 +
31 + safe_topk = torch.clamp(topk, min=0)
32 + block_indices = safe_topk // block_size
33 + block_offsets = safe_topk % block_size
34 + req_indices = token_to_req_indices[:, None].expand_as(topk)
35 + slot_ids = block_table[req_indices, block_indices] * block_size + block_offsets
36 +
37 + offsets = torch.arange(topk.shape[1], dtype=torch.int32, device=topk.device)
38 + positions = indptr[:-1, None] + offsets[None, :]
39 + return slot_ids[valid], positions[valid].to(torch.long), indptr, lens
40 +
41 +
42 +def _ref_sparse_prefill_ragged(
43 + q: torch.Tensor,
44 + kv: torch.Tensor,
45 + rows: list[list[int]],
46 + scale: float,
47 + attn_sink: torch.Tensor | None,
48 +) -> torch.Tensor:
49 + q_f32 = q.float()
50 + kv_f32 = kv.float()
51 + out = torch.empty_like(q_f32)
52 +
53 + for query_idx in range(q.shape[0]):
54 + row_indices = rows[query_idx]
55 + for head_idx in range(q.shape[1]):
56 + if row_indices:
57 + selected_kv = kv_f32[row_indices]
58 + scores = torch.mv(selected_kv, q_f32[query_idx, head_idx]) * scale
59 + if attn_sink is not None:
60 + scores_with_sink = torch.cat(
61 + [scores, attn_sink[head_idx].float().reshape(1)]
62 + )
63 + probs = torch.softmax(scores_with_sink, dim=0)[:-1]
64 + else:
65 + probs = torch.softmax(scores, dim=0)
66 + out[query_idx, head_idx] = torch.sum(
67 + probs[:, None] * selected_kv, dim=0
68 + )
69 + else:
70 + out[query_idx, head_idx] = 0
71 + return out.to(torch.bfloat16)
72 +
73 +
74 +def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
75 + assert kv.shape[-1] == HEAD_DIM
76 + num_tokens = kv.shape[0]
77 + num_blocks = (num_tokens + block_size - 1) // block_size
78 + cache = torch.zeros(
79 + (num_blocks, block_size, 584),
80 + dtype=torch.uint8,
81 + device=kv.device,
82 + )
83 + cache_flat = cache.view(torch.uint8).flatten()
84 + kv_nope_fp8 = (
85 + kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
86 + )
87 + kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
88 +
89 + for slot in range(num_tokens):
90 + block_idx = slot // block_size
91 + pos = slot % block_size
92 + block_base = block_idx * cache.stride(0)
93 + token_base = block_base + pos * 576
94 + scale_base = block_base + block_size * 576 + pos * 8
95 + cache_flat[token_base : token_base + NOPE_HEAD_DIM].copy_(kv_nope_fp8[slot])
96 + cache_flat[
97 + token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
98 + ].copy_(kv_rope_u8[slot])
99 + cache_flat[scale_base : scale_base + 7].fill_(127)
100 + return cache
101 +
102 +
103 +def _read_fp8_ds_mla_cache(
104 + cache: torch.Tensor, slot: int, block_size: int
105 +) -> torch.Tensor:
106 + cache_flat = cache.view(torch.uint8).flatten()
107 + block_idx = slot // block_size
108 + pos = slot % block_size
109 + block_base = block_idx * cache.stride(0)
110 + token_base = block_base + pos * 576
111 +
112 + nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
113 + nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
114 + rope_u8 = cache_flat[
115 + token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
116 + ]
117 + rope = rope_u8.view(torch.bfloat16).to(torch.float32)
118 + return torch.cat([nope, rope])
119 +
120 +
121 +def _ref_sparse_decode_ragged(
122 + q: torch.Tensor,
123 + main_cache: torch.Tensor,
124 + main_rows: list[list[int]],
125 + scale: float,
126 + attn_sink: torch.Tensor | None,
127 + block_size: int,
128 + extra_cache: torch.Tensor | None = None,
129 + extra_rows: list[list[int]] | None = None,
130 +) -> torch.Tensor:
131 + q_f32 = q.float()
132 + out = torch.empty_like(q_f32)
133 +
134 + for query_idx in range(q.shape[0]):
135 + row_kv = [
136 + _read_fp8_ds_mla_cache(main_cache, int(slot), block_size)
137 + for slot in main_rows[query_idx]
138 + ]
139 + if extra_cache is not None and extra_rows is not None:
140 + row_kv.extend(
141 + _read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
142 + for slot in extra_rows[query_idx]
143 + )
144 +
145 + kv = torch.stack(row_kv).to(q.device)
146 + for head_idx in range(q.shape[1]):
147 + scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale
148 + if attn_sink is not None:
149 + scores_with_sink = torch.cat(
150 + [scores, attn_sink[head_idx].float().reshape(1)]
151 + )
152 + probs = torch.softmax(scores_with_sink, dim=0)[:-1]
153 + else:
154 + probs = torch.softmax(scores, dim=0)
155 + out[query_idx, head_idx] = torch.sum(probs[:, None] * kv, dim=0)
156 + return out.to(torch.bfloat16)
157 +
158 +
159 +def _ref_combine_topk_swa_ragged(
160 + device: torch.device,
161 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
162 + expected_ragged = torch.tensor(
163 + [
164 + 100,
165 + 101,
166 + 7,
167 + 8,
168 + 9,
169 + 110,
170 + 111,
171 + 8,
172 + 9,
173 + 10,
174 + 120,
175 + 121,
176 + 122,
177 + 9,
178 + 10,
179 + 11,
180 + 150,
181 + 27,
182 + 28,
183 + 29,
184 + 160,
185 + 161,
186 + 28,
187 + 29,
188 + 30,
189 + ],
190 + dtype=torch.int32,
191 + device=device,
192 + )
193 + expected_lens = torch.tensor([5, 5, 6, 4, 5], dtype=torch.int32, device=device)
194 + expected_indptr = torch.zeros(6, dtype=torch.int32, device=device)
195 + torch.cumsum(expected_lens, dim=0, out=expected_indptr[1:])
196 + return expected_ragged, expected_indptr, expected_lens
197 +
198 +
199 +@torch.inference_mode()
200 +def test_compute_global_topk_ragged_indices_and_indptr() -> None:
201 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
202 + compute_global_topk_ragged_indices_and_indptr,
203 + )
204 +
205 + device = torch.device("cuda")
206 + block_size = 4
207 + topk_indices = torch.tensor(
208 + [
209 + [0, 3, 4, -1],
210 + [5, 8, -1, -1],
211 + [2, 7, 9, -1],
212 + ],
213 + dtype=torch.int32,
214 + device=device,
215 + )
216 + token_to_req_indices = torch.tensor([0, 1, 1], dtype=torch.int32, device=device)
217 + block_table = torch.tensor(
218 + [
219 + [10, 11, 12],
220 + [20, 21, 22],
221 + ],
222 + dtype=torch.int32,
223 + device=device,
224 + )
225 + is_valid_token = torch.tensor([True, False, True], dtype=torch.bool, device=device)
226 +
227 + actual_ragged, actual_indptr, actual_lens = (
228 + compute_global_topk_ragged_indices_and_indptr(
229 + topk_indices,
230 + token_to_req_indices,
231 + block_table,
232 + block_size,
233 + is_valid_token,
234 + )
235 + )
236 + expected_values, expected_positions, expected_indptr, expected_lens = (
237 + _ref_global_topk_ragged(
238 + topk_indices,
239 + token_to_req_indices,
240 + block_table,
241 + block_size,
242 + is_valid_token,
243 + )
244 + )
245 +
246 + torch.testing.assert_close(actual_ragged[expected_positions], expected_values)
247 + torch.testing.assert_close(actual_indptr, expected_indptr)
248 + torch.testing.assert_close(actual_lens, expected_lens)
249 +
250 +
251 +@torch.inference_mode()
252 +def test_sparse_attn_prefill_ragged_kernel() -> None:
253 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
254 + _rocm_sparse_attn_prefill_ragged_triton,
255 + )
256 +
257 + device = torch.device("cuda")
258 + torch.manual_seed(0)
259 + q = torch.randn(3, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
260 + kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
261 + indices = torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device)
262 + indptr = torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device)
263 + attn_sink = torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device)
264 + scale = HEAD_DIM**-0.5
265 +
266 + actual = _rocm_sparse_attn_prefill_ragged_triton(
267 + q=q,
268 + kv=kv,
269 + indices=indices,
270 + indptr=indptr,
271 + scale=scale,
272 + attn_sink=attn_sink,
273 + nope_head_dim=NOPE_HEAD_DIM,
274 + rope_head_dim=ROPE_HEAD_DIM,
275 + )
276 + expected = _ref_sparse_prefill_ragged(
277 + q, kv, [[0, 2], [1, 3, 4], []], scale, attn_sink
278 + )
279 +
280 + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
281 +
282 +
283 +@torch.inference_mode()
284 +def test_sparse_attn_decode_ragged_kernel() -> None:
285 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
286 + _rocm_sparse_attn_decode_ragged_triton,
287 + )
288 +
289 + device = torch.device("cuda")
290 + torch.manual_seed(1)
291 + block_size = 4
292 + q = torch.randn(2, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
293 + main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
294 + extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
295 + main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
296 + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
297 + main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
298 + main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
299 + extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
300 + extra_indptr = torch.tensor([0, 1, 3], dtype=torch.int32, device=device)
301 + attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
302 + scale = HEAD_DIM**-0.5
303 +
304 + actual = _rocm_sparse_attn_decode_ragged_triton(
305 + q=q,
306 + main_cache=main_cache,
307 + main_indices=main_indices,
308 + main_indptr=main_indptr,
309 + scale=scale,
310 + attn_sink=attn_sink,
311 + nope_head_dim=NOPE_HEAD_DIM,
312 + rope_head_dim=ROPE_HEAD_DIM,
313 + extra_cache=extra_cache,
314 + extra_indices=extra_indices,
315 + extra_indptr=extra_indptr,
316 + )
317 + expected = _ref_sparse_decode_ragged(
318 + q=q,
319 + main_cache=main_cache,
320 + main_rows=[[0, 2], [4, 1]],
321 + scale=scale,
322 + attn_sink=attn_sink,
323 + block_size=block_size,
324 + extra_cache=extra_cache,
325 + extra_rows=[[1], [3, 0]],
326 + )
327 +
328 + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
329 +
330 +
331 +@torch.inference_mode()
332 +def test_combine_topk_swa_indices_ragged() -> None:
333 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
334 + combine_topk_swa_indices_ragged,
335 + )
336 +
337 + device = torch.device("cuda")
338 + topk_indices = torch.tensor(
339 + [
340 + [100, 101, 102, 103],
341 + [110, 111, 112, 113],
342 + [120, 121, 122, 123],
343 + [130, 131, 132, 133],
344 + [140, 141, 142, 143],
345 + ],
346 + dtype=torch.int32,
347 + device=device,
348 + )
349 + query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
350 + seq_lens = torch.tensor([6, 4], dtype=torch.int32, device=device)
351 + gather_lens = torch.tensor([4, 3], dtype=torch.int32, device=device)
352 + window_size = 3
353 + compress_ratio = 2
354 + topk = 4
355 + M = 20
356 + N = 8
357 +
358 + actual_ragged, actual_indptr, actual_lens = combine_topk_swa_indices_ragged(
359 + topk_indices,
360 + query_start_loc,
361 + seq_lens,
362 + gather_lens,
363 + window_size,
364 + compress_ratio,
365 + topk,
366 + M,
367 + N,
368 + )
369 + expected_ragged, expected_indptr, expected_lens = _ref_combine_topk_swa_ragged(
370 + device
371 + )
372 +
373 + torch.testing.assert_close(
374 + actual_ragged[: expected_ragged.numel()], expected_ragged
375 + )
376 + torch.testing.assert_close(actual_indptr, expected_indptr)
377 + torch.testing.assert_close(actual_lens, expected_lens)
vllm/model_executor/layers/deepseek_v4_attention.py
+24 -46
@@ -28,11 +28,7 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
28 fused_inv_rope_fp8_quant,
29 fused_q_kv_rmsnorm,
30 )
31 -from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
32 - rocm_forward_decode_fallback,
33 - rocm_inv_rope_einsum,
34 - rocm_sparse_attn_prefill,
35 -)
31 +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum
32
33 if TYPE_CHECKING:
34 from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -725,6 +721,12 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
721 self.kv_cache = torch.tensor([])
722
723 def get_attn_backend(self) -> type[AttentionBackend]:
724 + if current_platform.is_rocm():
725 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
726 + DeepseekV4ROCMAiterMLASparseBackend,
727 + )
728 +
729 + return DeepseekV4ROCMAiterMLASparseBackend
730 return DeepseekV4FlashMLASparseBackend
731
732 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
@@ -757,6 +759,14 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
759 f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
760 )
761
762 + if current_platform.is_rocm():
763 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
764 + DeepseekV4ROCMAiterMLASparseImpl,
765 + )
766 +
767 + DeepseekV4ROCMAiterMLASparseImpl.forward(self, q, kv, positions, output)
768 + return
769 +
770 # Get SWA and indexer metadata from forward context
771 forward_context = get_forward_context()
772 attn_metadata = forward_context.attn_metadata
@@ -839,25 +849,6 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
849 swa_indices = swa_metadata.decode_swa_indices
850 swa_lens = swa_metadata.decode_swa_lens
851
842 - if current_platform.is_rocm():
843 - rocm_forward_decode_fallback(
844 - q=q,
845 - kv_cache=kv_cache,
846 - swa_k_cache=self.swa_cache_layer.kv_cache,
847 - swa_only=swa_only,
848 - topk_indices=topk_indices,
849 - topk_lens=topk_lens,
850 - swa_indices=swa_indices,
851 - swa_lens=swa_lens,
852 - attn_sink=self.attn_sink,
853 - scale=self.scale,
854 - head_dim=self.head_dim,
855 - nope_head_dim=self.nope_head_dim,
856 - rope_head_dim=self.rope_head_dim,
857 - output=output,
858 - )
859 - return
860 -
852 # We treat queries in the same seq as different queries
853 # and later we only attend by generated indices.
854 # q arrives pre-padded to self.padded_heads by the outer wrapper.
@@ -1021,28 +1012,15 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
1012 M,
1013 N,
1014 )
1024 -
1025 - if current_platform.is_rocm():
1026 - rocm_sparse_attn_prefill(
1027 - q=q[query_start:query_end],
1028 - kv=kv.view(-1, 1, q.shape[-1]),
1029 - indices=combined_indices.unsqueeze(1),
1030 - topk_length=combined_lens,
1031 - scale=self.scale,
1032 - head_dim=self.head_dim,
1033 - attn_sink=self.attn_sink,
1034 - output=output[query_start:query_end],
1035 - )
1036 - else:
1037 - output_chunk, _, _ = flash_mla_sparse_fwd(
1038 - q=q[query_start:query_end],
1039 - kv=kv.view(-1, 1, q.shape[-1]),
1040 - indices=combined_indices.unsqueeze(1),
1041 - sm_scale=self.scale,
1042 - attn_sink=self.attn_sink,
1043 - topk_length=combined_lens,
1044 - out=output[query_start:query_end],
1045 - )
1015 + flash_mla_sparse_fwd(
1016 + q=q[query_start:query_end],
1017 + kv=kv.view(-1, 1, q.shape[-1]),
1018 + indices=combined_indices.unsqueeze(1),
1019 + sm_scale=self.scale,
1020 + attn_sink=self.attn_sink,
1021 + topk_length=combined_lens,
1022 + out=output[query_start:query_end],
1023 + )
1024
1025
1026 class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
vllm/v1/attention/backends/mla/flashmla_sparse.py
+2 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 -from typing import TYPE_CHECKING, ClassVar
4 +from typing import TYPE_CHECKING, Any, ClassVar
5
6 import numpy as np
7 import torch
@@ -112,7 +112,7 @@ class FlashMLASparseBackend(AttentionBackend):
112 return FlashMLASparseMetadataBuilder
113
114 @staticmethod
115 - def get_impl_cls() -> type["FlashMLASparseImpl"]:
115 + def get_impl_cls() -> type[SparseMLAAttentionImpl[Any]]:
116 return FlashMLASparseImpl
117
118 @classmethod
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse_dsv4.py new
+682
@@ -0,0 +1,682 @@
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, cast
6 +
7 +import torch
8 +
9 +from vllm.forward_context import get_forward_context
10 +from vllm.triton_utils import tl, triton
11 +from vllm.v1.attention.backend import (
12 + AttentionLayer,
13 + CommonAttentionMetadata,
14 + SparseMLAAttentionImpl,
15 +)
16 +from vllm.v1.attention.backends.mla.flashmla_sparse import (
17 + DeepseekV4FlashMLASparseBackend,
18 + FlashMLASparseMetadata,
19 + FlashMLASparseMetadataBuilder,
20 +)
21 +from vllm.v1.attention.backends.mla.sparse_swa import (
22 + DeepseekSparseSWAMetadata,
23 + DeepseekSparseSWAMetadataBuilder,
24 +)
25 +from vllm.v1.attention.ops.deepseek_v4_ops import dequantize_and_gather_k_cache
26 +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
27 + build_ragged_indices_from_dense,
28 + rocm_sparse_attn_decode,
29 + rocm_sparse_attn_prefill,
30 +)
31 +from vllm.v1.worker.workspace import current_workspace_manager
32 +
33 +if TYPE_CHECKING:
34 + from vllm.model_executor.layers.deepseek_v4_attention import (
35 + DeepseekV4MLAAttention,
36 + )
37 +
38 +
39 +def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
40 + lengths = lengths.to(dtype=torch.int32).contiguous()
41 + indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
42 + torch.cumsum(lengths, dim=0, out=indptr[1:])
43 + return indptr
44 +
45 +
46 +@triton.jit
47 +def _compute_topk_lens_kernel(
48 + topk_lens_ptr,
49 + topk_indices_ptr,
50 + topk_indices_stride,
51 + topk,
52 + is_valid_token_ptr,
53 + TRITON_BLOCK_SIZE: tl.constexpr,
54 +):
55 + token_idx = tl.program_id(0)
56 + is_valid_token = tl.load(is_valid_token_ptr + token_idx)
57 +
58 + count = tl.zeros((), dtype=tl.int32)
59 + for i in range(0, topk, TRITON_BLOCK_SIZE):
60 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
61 + mask = offset < topk
62 + local_idx = tl.load(
63 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
64 + mask=mask,
65 + other=-1,
66 + )
67 + count += tl.sum((local_idx >= 0).to(tl.int32), axis=0)
68 +
69 + tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))
70 +
71 +
72 +@triton.jit
73 +def _pack_global_topk_ragged_kernel(
74 + global_topk_ragged_ptr,
75 + topk_indptr_ptr,
76 + topk_indices_ptr,
77 + topk_indices_stride,
78 + token_to_req_indices_ptr,
79 + block_table_ptr,
80 + block_table_stride,
81 + block_size,
82 + topk,
83 + BLOCK_SIZE: tl.constexpr,
84 +):
85 + token_idx = tl.program_id(0)
86 + block_idx = tl.program_id(1)
87 + offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
88 +
89 + out_start = tl.load(topk_indptr_ptr + token_idx)
90 + out_end = tl.load(topk_indptr_ptr + token_idx + 1)
91 + out_len = out_end - out_start
92 + if block_idx * BLOCK_SIZE >= out_len:
93 + return
94 +
95 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
96 + mask = (offset < out_len) & (offset < topk)
97 + local_idx = tl.load(
98 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
99 + mask=mask,
100 + other=-1,
101 + )
102 + valid = mask & (local_idx >= 0)
103 + block_indices = local_idx // block_size
104 + block_numbers = tl.load(
105 + block_table_ptr + req_idx * block_table_stride + block_indices,
106 + mask=valid,
107 + other=0,
108 + )
109 + block_offsets = local_idx % block_size
110 + slot_ids = tl.where(valid, block_numbers * block_size + block_offsets, -1)
111 + tl.store(global_topk_ragged_ptr + out_start + offset, slot_ids, mask=mask)
112 +
113 +
114 +def compute_global_topk_ragged_indices_and_indptr(
115 + topk_indices: torch.Tensor,
116 + token_to_req_indices: torch.Tensor,
117 + block_table: torch.Tensor,
118 + block_size: int,
119 + is_valid_token: torch.Tensor,
120 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
121 + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
122 + num_tokens = topk_indices.shape[0]
123 + topk = topk_indices.shape[1]
124 +
125 + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
126 + _compute_topk_lens_kernel[(num_tokens,)](
127 + topk_lens,
128 + topk_indices,
129 + topk_indices.stride(0),
130 + topk,
131 + is_valid_token,
132 + TRITON_BLOCK_SIZE=1024,
133 + )
134 +
135 + topk_indptr = _build_indptr_from_lengths(topk_lens)
136 + global_topk_ragged = torch.empty(
137 + num_tokens * topk,
138 + dtype=torch.int32,
139 + device=topk_indices.device,
140 + )
141 + if global_topk_ragged.numel() > 0:
142 + block = 128
143 + _pack_global_topk_ragged_kernel[(num_tokens, triton.cdiv(topk, block))](
144 + global_topk_ragged,
145 + topk_indptr,
146 + topk_indices,
147 + topk_indices.stride(0),
148 + token_to_req_indices,
149 + block_table,
150 + block_table.stride(0),
151 + block_size,
152 + topk,
153 + BLOCK_SIZE=block,
154 + )
155 + return global_topk_ragged, topk_indptr, topk_lens
156 +
157 +
158 +@triton.jit
159 +def _compute_combined_lens_kernel(
160 + combined_lens_ptr,
161 + query_start_loc_ptr,
162 + seq_lens_ptr,
163 + TOP_K: tl.constexpr,
164 + COMPRESS_RATIO: tl.constexpr,
165 + WINDOW_SIZE: tl.constexpr,
166 +):
167 + batch_idx = tl.program_id(0)
168 + worker_id = tl.program_id(1)
169 + num_workers = tl.num_programs(1)
170 +
171 + base = tl.load(query_start_loc_ptr)
172 + query_start = tl.load(query_start_loc_ptr + batch_idx) - base
173 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
174 + query_len = query_end - query_start
175 + seq_len = tl.load(seq_lens_ptr + batch_idx)
176 + start_pos = seq_len - query_len
177 +
178 + for token_idx in range(query_start + worker_id, query_end, num_workers):
179 + token_idx_in_query = token_idx - query_start
180 + pos = start_pos + token_idx_in_query
181 + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
182 + swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
183 + tl.store(combined_lens_ptr + token_idx, topk_len + swa_len)
184 +
185 +
186 +@triton.jit
187 +def _combine_topk_swa_indices_ragged_kernel(
188 + combined_ragged_ptr,
189 + combined_indptr_ptr,
190 + topk_indices_ptr,
191 + topk_indices_stride,
192 + query_start_loc_ptr,
193 + seq_lens_ptr,
194 + gather_lens_ptr,
195 + M,
196 + N,
197 + topk_width,
198 + TOP_K: tl.constexpr,
199 + COMPRESS_RATIO: tl.constexpr,
200 + WINDOW_SIZE: tl.constexpr,
201 + BLOCK_SIZE: tl.constexpr,
202 +):
203 + batch_idx = tl.program_id(0)
204 + worker_id = tl.program_id(1)
205 + block_idx = tl.program_id(2)
206 + num_workers = tl.num_programs(1)
207 +
208 + base = tl.load(query_start_loc_ptr)
209 + query_start = tl.load(query_start_loc_ptr + batch_idx) - base
210 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
211 + query_len = query_end - query_start
212 + seq_len = tl.load(seq_lens_ptr + batch_idx)
213 + gather_len = tl.load(gather_lens_ptr + batch_idx)
214 + start_pos = seq_len - query_len
215 + gather_start = seq_len - gather_len
216 +
217 + for token_idx in range(query_start + worker_id, query_end, num_workers):
218 + token_idx_in_query = token_idx - query_start
219 + pos = start_pos + token_idx_in_query
220 + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
221 + swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
222 + combined_len = topk_len + swa_len
223 +
224 + offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
225 + if block_idx * BLOCK_SIZE < combined_len:
226 + out_start = tl.load(combined_indptr_ptr + token_idx)
227 + topk_mask = (offset < topk_len) & (offset < topk_width)
228 + topk_vals = tl.load(
229 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
230 + mask=topk_mask,
231 + other=-1,
232 + )
233 + tl.store(
234 + combined_ragged_ptr + out_start + offset,
235 + topk_vals + M * batch_idx,
236 + mask=topk_mask,
237 + )
238 +
239 + swa_offset = offset - topk_len
240 + swa_mask = (offset >= topk_len) & (swa_offset < swa_len)
241 + tl.store(
242 + combined_ragged_ptr + out_start + offset,
243 + M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start,
244 + mask=swa_mask,
245 + )
246 +
247 +
248 +def combine_topk_swa_indices_ragged(
249 + topk_indices: torch.Tensor,
250 + query_start_loc: torch.Tensor,
251 + seq_lens: torch.Tensor,
252 + gather_lens: torch.Tensor,
253 + window_size: int,
254 + compress_ratio: int,
255 + topk: int,
256 + M: int,
257 + N: int,
258 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
259 + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
260 + num_tokens = topk_indices.shape[0]
261 + num_reqs = seq_lens.shape[0]
262 + combined_lens = torch.empty(
263 + num_tokens, dtype=torch.int32, device=topk_indices.device
264 + )
265 +
266 + num_workers = 128
267 + _compute_combined_lens_kernel[(num_reqs, num_workers)](
268 + combined_lens,
269 + query_start_loc,
270 + seq_lens,
271 + TOP_K=topk,
272 + COMPRESS_RATIO=compress_ratio,
273 + WINDOW_SIZE=window_size,
274 + )
275 +
276 + combined_indptr = _build_indptr_from_lengths(combined_lens)
277 + combined_ragged = torch.empty(
278 + num_tokens * (topk + window_size),
279 + dtype=torch.int32,
280 + device=topk_indices.device,
281 + )
282 + if combined_ragged.numel() > 0:
283 + block = 128
284 + _combine_topk_swa_indices_ragged_kernel[
285 + (num_reqs, num_workers, triton.cdiv(topk + window_size, block))
286 + ](
287 + combined_ragged,
288 + combined_indptr,
289 + topk_indices,
290 + topk_indices.stride(0),
291 + query_start_loc,
292 + seq_lens,
293 + gather_lens,
294 + M,
295 + N,
296 + topk_indices.shape[-1],
297 + TOP_K=topk,
298 + COMPRESS_RATIO=compress_ratio,
299 + WINDOW_SIZE=window_size,
300 + BLOCK_SIZE=block,
301 + )
302 + return combined_ragged, combined_indptr, combined_lens
303 +
304 +
305 +@dataclass
306 +class DeepseekV4ROCMAiterMLASparseMetadata(FlashMLASparseMetadata):
307 + """ROCm-specific DeepSeek V4 metadata carrying ragged decode topk."""
308 +
309 + c128a_decode_topk_ragged_indices: torch.Tensor | None = None
310 + c128a_decode_topk_ragged_indptr: torch.Tensor | None = None
311 +
312 +
313 +@dataclass
314 +class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
315 + decode_swa_ragged_indices: torch.Tensor | None = None
316 + decode_swa_ragged_indptr: torch.Tensor | None = None
317 +
318 +
319 +class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder):
320 + def build(
321 + self,
322 + common_prefix_len: int,
323 + common_attn_metadata: CommonAttentionMetadata,
324 + fast_build: bool = False,
325 + ) -> DeepseekV4ROCMAiterMLASparseMetadata:
326 + base = super().build(
327 + common_prefix_len=common_prefix_len,
328 + common_attn_metadata=common_attn_metadata,
329 + fast_build=fast_build,
330 + )
331 +
332 + ragged_indices = None
333 + ragged_indptr = None
334 + dense_decode = base.c128a_global_decode_topk_indices
335 + decode_lens = base.c128a_decode_topk_lens
336 + if dense_decode is not None and decode_lens is not None:
337 + ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
338 + dense_decode.reshape(dense_decode.shape[0], -1),
339 + decode_lens,
340 + )
341 +
342 + return DeepseekV4ROCMAiterMLASparseMetadata(
343 + **vars(base),
344 + c128a_decode_topk_ragged_indices=ragged_indices,
345 + c128a_decode_topk_ragged_indptr=ragged_indptr,
346 + )
347 +
348 +
349 +class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
350 + def build(
351 + self,
352 + common_prefix_len: int,
353 + common_attn_metadata: CommonAttentionMetadata,
354 + fast_build: bool = False,
355 + ) -> DeepseekV4ROCMAiterSparseSWAMetadata:
356 + base = super().build(
357 + common_prefix_len=common_prefix_len,
358 + common_attn_metadata=common_attn_metadata,
359 + fast_build=fast_build,
360 + )
361 +
362 + ragged_indices = None
363 + ragged_indptr = None
364 + if (
365 + base.num_decode_tokens > 0
366 + and base.decode_swa_indices is not None
367 + and base.decode_swa_lens is not None
368 + ):
369 + ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
370 + base.decode_swa_indices.reshape(base.num_decode_tokens, -1),
371 + base.decode_swa_lens,
372 + )
373 +
374 + return DeepseekV4ROCMAiterSparseSWAMetadata(
375 + **vars(base),
376 + decode_swa_ragged_indices=ragged_indices,
377 + decode_swa_ragged_indptr=ragged_indptr,
378 + )
379 +
380 +
381 +class DeepseekV4ROCMAiterMLASparseImpl(
382 + SparseMLAAttentionImpl[DeepseekV4ROCMAiterMLASparseMetadata]
383 +):
384 + """ROCm sparse MLA implementation used by DeepSeek V4's custom MLA layer."""
385 +
386 + _PREFILL_CHUNK_SIZE = 4
387 +
388 + def __init__(
389 + self,
390 + num_heads: int,
391 + head_size: int,
392 + scale: float,
393 + num_kv_heads: int,
394 + alibi_slopes: list[float] | None,
395 + sliding_window: int | None,
396 + kv_cache_dtype: str,
397 + logits_soft_cap: float | None,
398 + attn_type: str,
399 + kv_sharing_target_layer_name: str | None,
400 + **_: object,
401 + ) -> None:
402 + self.num_heads = num_heads
403 + self.head_size = head_size
404 + self.scale = float(scale)
405 + self.num_kv_heads = num_kv_heads
406 + self.kv_cache_dtype = kv_cache_dtype
407 +
408 + def forward_mqa(
409 + self,
410 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
411 + kv_c_and_k_pe_cache: torch.Tensor,
412 + attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata,
413 + layer: AttentionLayer,
414 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
415 + raise NotImplementedError(
416 + "DeepseekV4ROCMAiterMLASparseImpl is driven by "
417 + "DeepseekV4MLAAttention.forward."
418 + )
419 +
420 + @classmethod
421 + def forward(
422 + cls,
423 + layer: "DeepseekV4MLAAttention",
424 + q: torch.Tensor,
425 + kv: torch.Tensor,
426 + positions: torch.Tensor,
427 + output: torch.Tensor,
428 + ) -> None:
429 + assert output.shape == q.shape, (
430 + f"output buffer shape {output.shape} must match q shape {q.shape}"
431 + )
432 + assert output.dtype == q.dtype, (
433 + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
434 + )
435 +
436 + forward_context = get_forward_context()
437 + attn_metadata = forward_context.attn_metadata
438 + assert isinstance(attn_metadata, dict)
439 + rocm_metadata = cast(
440 + DeepseekV4ROCMAiterMLASparseMetadata | None,
441 + attn_metadata.get(layer.prefix),
442 + )
443 + swa_metadata = cast(
444 + DeepseekV4ROCMAiterSparseSWAMetadata | None,
445 + attn_metadata.get(layer.swa_cache_layer.prefix),
446 + )
447 + assert swa_metadata is not None
448 +
449 + swa_only = layer.compress_ratio <= 1
450 + self_kv_cache = layer.kv_cache if not swa_only else None
451 + swa_kv_cache = layer.swa_cache_layer.kv_cache
452 +
453 + num_decodes = swa_metadata.num_decodes
454 + num_prefills = swa_metadata.num_prefills
455 + num_decode_tokens = swa_metadata.num_decode_tokens
456 +
457 + if num_prefills > 0:
458 + cls._forward_prefill(
459 + layer=layer,
460 + q=q[num_decode_tokens:],
461 + positions=positions[num_decode_tokens:],
462 + compressed_k_cache=self_kv_cache,
463 + swa_k_cache=swa_kv_cache,
464 + output=output[num_decode_tokens:],
465 + attn_metadata=rocm_metadata,
466 + swa_metadata=swa_metadata,
467 + )
468 + if num_decodes > 0:
469 + cls._forward_decode(
470 + layer=layer,
471 + q=q[:num_decode_tokens],
472 + kv_cache=self_kv_cache,
473 + swa_metadata=swa_metadata,
474 + attn_metadata=rocm_metadata,
475 + swa_only=swa_only,
476 + output=output[:num_decode_tokens],
477 + )
478 +
479 + @classmethod
480 + def _forward_decode(
481 + cls,
482 + layer: "DeepseekV4MLAAttention",
483 + q: torch.Tensor,
484 + kv_cache: torch.Tensor | None,
485 + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
486 + attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
487 + swa_only: bool,
488 + output: torch.Tensor,
489 + ) -> None:
490 + num_decodes = swa_metadata.num_decodes
491 + num_decode_tokens = swa_metadata.num_decode_tokens
492 +
493 + topk_indices = None
494 + topk_lens = None
495 + topk_ragged_indices = None
496 + topk_ragged_indptr = None
497 + if not swa_only:
498 + assert attn_metadata is not None
499 + assert swa_metadata.is_valid_token is not None
500 + block_size = attn_metadata.block_size // layer.compress_ratio
501 + is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
502 + if layer.compress_ratio == 4:
503 + assert layer.topk_indices_buffer is not None
504 + (
505 + topk_ragged_indices,
506 + topk_ragged_indptr,
507 + topk_lens,
508 + ) = compute_global_topk_ragged_indices_and_indptr(
509 + layer.topk_indices_buffer[:num_decode_tokens],
510 + swa_metadata.token_to_req_indices,
511 + attn_metadata.block_table[:num_decodes],
512 + block_size,
513 + is_valid,
514 + )
515 + else:
516 + topk_indices = attn_metadata.c128a_global_decode_topk_indices
517 + topk_lens = attn_metadata.c128a_decode_topk_lens
518 + topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices
519 + topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr
520 +
521 + rocm_sparse_attn_decode(
522 + q=q,
523 + kv_cache=kv_cache,
524 + swa_k_cache=layer.swa_cache_layer.kv_cache,
525 + swa_only=swa_only,
526 + topk_indices=topk_indices,
527 + topk_lens=topk_lens,
528 + swa_indices=swa_metadata.decode_swa_indices,
529 + swa_lens=swa_metadata.decode_swa_lens,
530 + swa_ragged_indices=swa_metadata.decode_swa_ragged_indices,
531 + swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr,
532 + topk_ragged_indices=topk_ragged_indices,
533 + topk_ragged_indptr=topk_ragged_indptr,
534 + attn_sink=layer.attn_sink,
535 + scale=layer.scale,
536 + head_dim=layer.head_dim,
537 + nope_head_dim=layer.nope_head_dim,
538 + rope_head_dim=layer.rope_head_dim,
539 + output=output,
540 + )
541 +
542 + @classmethod
543 + def _forward_prefill(
544 + cls,
545 + layer: "DeepseekV4MLAAttention",
546 + q: torch.Tensor,
547 + positions: torch.Tensor,
548 + compressed_k_cache: torch.Tensor | None,
549 + swa_k_cache: torch.Tensor,
550 + output: torch.Tensor,
551 + attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
552 + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
553 + ) -> None:
554 + swa_only = attn_metadata is None
555 +
556 + num_prefills = swa_metadata.num_prefills
557 + num_prefill_tokens = swa_metadata.num_prefill_tokens
558 + num_decodes = swa_metadata.num_decodes
559 + num_decode_tokens = swa_metadata.num_decode_tokens
560 +
561 + seq_lens = swa_metadata.prefill_seq_lens
562 + gather_lens = swa_metadata.prefill_gather_lens
563 + assert seq_lens is not None
564 + assert gather_lens is not None
565 +
566 + query_start_loc_cpu = swa_metadata.query_start_loc_cpu
567 + query_start_loc = swa_metadata.query_start_loc
568 + assert query_start_loc_cpu is not None
569 + assert query_start_loc is not None
570 + prefill_token_base = query_start_loc_cpu[num_decodes]
571 +
572 + if not swa_only:
573 + if layer.compress_ratio == 4:
574 + assert layer.topk_indices_buffer is not None
575 + topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
576 + topk_indices = topk_indices[:num_prefill_tokens]
577 + else:
578 + assert attn_metadata is not None
579 + topk_indices = attn_metadata.c128a_prefill_topk_indices
580 + assert topk_indices is not None
581 + top_k = topk_indices.shape[-1]
582 + N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio
583 + else:
584 + assert layer.topk_indices_buffer is not None
585 + topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
586 + top_k = 0
587 + N = 0
588 +
589 + M = N + layer.window_size + layer.max_num_batched_tokens
590 + num_chunks = (num_prefills + cls._PREFILL_CHUNK_SIZE - 1) // (
591 + cls._PREFILL_CHUNK_SIZE
592 + )
593 +
594 + workspace_manager = current_workspace_manager()
595 + kv = workspace_manager.get_simultaneous(
596 + ((cls._PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
597 + )[0]
598 + for chunk_idx in range(num_chunks):
599 + chunk_start = chunk_idx * cls._PREFILL_CHUNK_SIZE
600 + chunk_end = min(chunk_start + cls._PREFILL_CHUNK_SIZE, num_prefills)
601 + chunk_size = chunk_end - chunk_start
602 + if not swa_only:
603 + assert attn_metadata is not None
604 + assert compressed_k_cache is not None
605 + block_table = attn_metadata.block_table[num_decodes:]
606 + dequantize_and_gather_k_cache(
607 + kv[:chunk_size],
608 + compressed_k_cache,
609 + seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio,
610 + gather_lens=None,
611 + block_table=block_table[chunk_start:chunk_end],
612 + block_size=attn_metadata.block_size // layer.compress_ratio,
613 + offset=0,
614 + )
615 +
616 + swa_block_table = swa_metadata.block_table[num_decodes:]
617 + dequantize_and_gather_k_cache(
618 + kv[:chunk_size],
619 + swa_k_cache,
620 + seq_lens=seq_lens[chunk_start:chunk_end],
621 + gather_lens=gather_lens[chunk_start:chunk_end],
622 + block_table=swa_block_table[chunk_start:chunk_end],
623 + block_size=swa_metadata.block_size,
624 + offset=N,
625 + )
626 +
627 + query_start = (
628 + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
629 + )
630 + query_end = (
631 + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
632 + )
633 +
634 + combined_ragged_indices, combined_ragged_indptr, combined_lens = (
635 + combine_topk_swa_indices_ragged(
636 + topk_indices[query_start:query_end],
637 + query_start_loc[
638 + num_decodes + chunk_start : num_decodes + chunk_end + 1
639 + ],
640 + seq_lens[chunk_start:chunk_end],
641 + gather_lens[chunk_start:chunk_end],
642 + layer.window_size,
643 + layer.compress_ratio,
644 + top_k,
645 + M,
646 + N,
647 + )
648 + )
649 + rocm_sparse_attn_prefill(
650 + q=q[query_start:query_end],
651 + kv=kv.view(-1, 1, q.shape[-1]),
652 + indices=torch.empty(
653 + q[query_start:query_end].shape[0],
654 + 1,
655 + 0,
656 + dtype=torch.int32,
657 + device=q.device,
658 + ),
659 + topk_length=combined_lens,
660 + scale=layer.scale,
661 + head_dim=layer.head_dim,
662 + nope_head_dim=layer.nope_head_dim,
663 + rope_head_dim=layer.rope_head_dim,
664 + attn_sink=layer.attn_sink,
665 + output=output[query_start:query_end],
666 + ragged_indices=combined_ragged_indices,
667 + ragged_indptr=combined_ragged_indptr,
668 + )
669 +
670 +
671 +class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend):
672 + @staticmethod
673 + def get_name() -> str:
674 + return "ROCM_V4_FLASHMLA_SPARSE"
675 +
676 + @staticmethod
677 + def get_builder_cls() -> type["DeepseekV4ROCMAiterMLASparseMetadataBuilder"]:
678 + return DeepseekV4ROCMAiterMLASparseMetadataBuilder
679 +
680 + @staticmethod
681 + def get_impl_cls() -> type["DeepseekV4ROCMAiterMLASparseImpl"]:
682 + return DeepseekV4ROCMAiterMLASparseImpl
vllm/v1/attention/backends/mla/sparse_swa.py
+6
@@ -112,6 +112,12 @@ class DeepseekSparseSWABackend(AttentionBackend):
112
113 @staticmethod
114 def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]:
115 + if current_platform.is_rocm():
116 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
117 + DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
118 + )
119 +
120 + return DeepseekV4ROCMAiterSparseSWAMetadataBuilder
121 return DeepseekSparseSWAMetadataBuilder
122
123 @staticmethod
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+758 -164
@@ -905,185 +905,757 @@ def rocm_inv_rope_einsum(
905 return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
906
907
908 -def rocm_ref_sparse_attn_prefill(
908 +_DSV4_SPARSE_NOPE_DIM = 448
909 +_DSV4_SPARSE_ROPE_DIM = 64
910 +
911 +
912 +def _validate_dsv4_sparse_dims(
913 + head_dim: int,
914 + nope_head_dim: int,
915 + rope_head_dim: int,
916 + op_name: str,
917 +) -> None:
918 + assert head_dim == nope_head_dim + rope_head_dim, (
919 + f"{op_name} expected head_dim={nope_head_dim + rope_head_dim}, got {head_dim}"
920 + )
921 + assert (
922 + nope_head_dim == _DSV4_SPARSE_NOPE_DIM
923 + and rope_head_dim == _DSV4_SPARSE_ROPE_DIM
924 + ), (
925 + f"{op_name} expects {_DSV4_SPARSE_NOPE_DIM} NoPE dims and "
926 + f"{_DSV4_SPARSE_ROPE_DIM} RoPE dims"
927 + )
928 +
929 +
930 +@triton.jit
931 +def _pack_dense_prefix_to_ragged_kernel(
932 + indices_ptr,
933 + lengths_ptr,
934 + indptr_ptr,
935 + out_ptr,
936 + indices_stride0,
937 + num_rows_limit,
938 + row_width,
939 + BLOCK_SIZE: tl.constexpr,
940 +):
941 + row_idx = tl.program_id(0)
942 + block_idx = tl.program_id(1)
943 + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
944 +
945 + row_len = tl.load(lengths_ptr + row_idx)
946 + if block_idx * BLOCK_SIZE >= row_len:
947 + return
948 +
949 + mask = offsets < row_len
950 + vals = tl.load(
951 + indices_ptr + row_idx * indices_stride0 + offsets,
952 + mask=mask & (offsets < row_width),
953 + other=-1,
954 + ).to(tl.int32)
955 + if num_rows_limit >= 0:
956 + vals = tl.where((vals >= 0) & (vals < num_rows_limit), vals, -1)
957 +
958 + out_start = tl.load(indptr_ptr + row_idx)
959 + tl.store(out_ptr + out_start + offsets, vals, mask=mask)
960 +
961 +
962 +def build_ragged_indices_from_dense(
963 + indices: torch.Tensor,
964 + lengths: torch.Tensor,
965 + num_rows: int = -1,
966 +) -> tuple[torch.Tensor, torch.Tensor]:
967 + indices = indices.reshape(indices.shape[0], -1)
968 + lengths = lengths.to(device=indices.device, dtype=torch.int32).reshape(-1)
969 + assert lengths.numel() == indices.shape[0], (
970 + f"Expected one length per row, got {lengths.shape} for indices {indices.shape}"
971 + )
972 +
973 + max_width = indices.shape[1] if indices.ndim == 2 else 0
974 + lengths = lengths.clamp(min=0, max=max_width).contiguous()
975 +
976 + indptr = torch.empty(indices.shape[0] + 1, dtype=torch.int32, device=indices.device)
977 + indptr[0] = 0
978 + torch.cumsum(lengths, dim=0, out=indptr[1:])
979 +
980 + if indices.numel() == 0:
981 + flat = torch.empty(0, dtype=torch.int32, device=indices.device)
982 + else:
983 + flat = torch.empty(
984 + int(indptr[-1].item()), dtype=torch.int32, device=indices.device
985 + )
986 + if flat.numel() > 0:
987 + block_size = 128
988 + _pack_dense_prefix_to_ragged_kernel[
989 + (indices.shape[0], triton.cdiv(max_width, block_size))
990 + ](
991 + indices,
992 + lengths,
993 + indptr,
994 + flat,
995 + indices.stride(0),
996 + int(num_rows),
997 + max_width,
998 + BLOCK_SIZE=block_size,
999 + )
1000 +
1001 + return flat, indptr
1002 +
1003 +
1004 +def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor:
1005 + if x.dtype == torch.int32 and x.ndim == 1 and x.is_contiguous():
1006 + return x
1007 + return x.to(torch.int32).contiguous()
1008 +
1009 +
1010 +@triton.jit
1011 +def _sparse_attn_prefill_ragged_kernel(
1012 + q_ptr,
1013 + kv_ptr,
1014 + kv_indices_ptr,
1015 + kv_indptr_ptr,
1016 + attn_sink_ptr,
1017 + out_ptr,
1018 + q_stride_t,
1019 + q_stride_h,
1020 + q_stride_d,
1021 + kv_stride_n,
1022 + kv_stride_d,
1023 + out_stride_t,
1024 + out_stride_h,
1025 + out_stride_d,
1026 + num_heads,
1027 + head_dim,
1028 + num_kv,
1029 + scale,
1030 + HAS_ATTN_SINK: tl.constexpr,
1031 + BLOCK_H: tl.constexpr,
1032 + BLOCK_D: tl.constexpr,
1033 + BLOCK_K: tl.constexpr,
1034 +):
1035 + query_idx = tl.program_id(0)
1036 + pid_h = tl.program_id(1)
1037 +
1038 + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
1039 + dim_offsets = tl.arange(0, BLOCK_D)
1040 + head_mask = head_offsets < num_heads
1041 + dim_mask = dim_offsets < head_dim
1042 +
1043 + q = tl.load(
1044 + q_ptr
1045 + + query_idx * q_stride_t
1046 + + head_offsets[:, None] * q_stride_h
1047 + + dim_offsets[None, :] * q_stride_d,
1048 + mask=head_mask[:, None] & dim_mask[None, :],
1049 + other=0.0,
1050 + )
1051 +
1052 + neg_large = -3.4028234663852886e38
1053 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
1054 + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
1055 + acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
1056 +
1057 + kv_start = tl.load(kv_indptr_ptr + query_idx)
1058 + kv_end = tl.load(kv_indptr_ptr + query_idx + 1)
1059 + kv_len = kv_end - kv_start
1060 +
1061 + k_offsets = tl.arange(0, BLOCK_K)
1062 + for k_start in tl.range(0, kv_len, BLOCK_K):
1063 + k_pos = k_start + k_offsets
1064 + in_range = k_pos < kv_len
1065 + slot = tl.load(kv_indices_ptr + kv_start + k_pos, mask=in_range, other=-1)
1066 + valid = in_range & (slot >= 0) & (slot < num_kv)
1067 +
1068 + kv = tl.load(
1069 + kv_ptr + slot[:, None] * kv_stride_n + dim_offsets[None, :] * kv_stride_d,
1070 + mask=valid[:, None] & dim_mask[None, :],
1071 + other=0.0,
1072 + )
1073 + kv = tl.where(valid[:, None] & dim_mask[None, :], kv, 0.0)
1074 +
1075 + scores = tl.dot(q, tl.trans(kv)) * scale
1076 + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
1077 +
1078 + m_block = tl.max(scores, axis=1)
1079 + m_new = tl.maximum(m_i, m_block)
1080 + alpha = tl.exp(m_i - m_new)
1081 + p = tl.exp(scores - m_new[:, None])
1082 + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
1083 + l_new = l_i * alpha + tl.sum(p, axis=1)
1084 +
1085 + acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
1086 + m_i = m_new
1087 + l_i = l_new
1088 +
1089 + if HAS_ATTN_SINK:
1090 + sink = tl.load(
1091 + attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large
1092 + ).to(tl.float32)
1093 + m_final = tl.maximum(m_i, sink)
1094 + alpha = tl.exp(m_i - m_final)
1095 + l_final = l_i * alpha + tl.exp(sink - m_final)
1096 + denom = tl.maximum(l_final, 1.0e-30)
1097 + out = tl.where(
1098 + l_final[:, None] > 0.0,
1099 + (acc * alpha[:, None]) / denom[:, None],
1100 + 0.0,
1101 + )
1102 + else:
1103 + denom = tl.maximum(l_i, 1.0e-30)
1104 + out = tl.where(l_i[:, None] > 0.0, acc / denom[:, None], 0.0)
1105 +
1106 + tl.store(
1107 + out_ptr
1108 + + query_idx * out_stride_t
1109 + + head_offsets[:, None] * out_stride_h
1110 + + dim_offsets[None, :] * out_stride_d,
1111 + out,
1112 + mask=head_mask[:, None] & dim_mask[None, :],
1113 + )
1114 +
1115 +
1116 +@triton.jit
1117 +def _sparse_attn_decode_ragged_kernel(
1118 + q_ptr,
1119 + main_cache_ptr,
1120 + main_indices_ptr,
1121 + main_indptr_ptr,
1122 + extra_cache_ptr,
1123 + extra_indices_ptr,
1124 + extra_indptr_ptr,
1125 + attn_sink_ptr,
1126 + out_ptr,
1127 + q_stride0,
1128 + q_stride1,
1129 + out_stride0,
1130 + out_stride1,
1131 + main_cache_stride0,
1132 + extra_cache_stride0,
1133 + main_num_rows,
1134 + extra_num_rows,
1135 + main_block_size,
1136 + extra_block_size,
1137 + scale,
1138 + num_heads,
1139 + HAS_ATTN_SINK: tl.constexpr,
1140 + HAS_EXTRA: tl.constexpr,
1141 + NOPE_DIM: tl.constexpr,
1142 + NOPE_BLOCK: tl.constexpr,
1143 + ROPE_DIM: tl.constexpr,
1144 + IS_FNUZ: tl.constexpr,
1145 + BLOCK_H: tl.constexpr,
1146 + BLOCK_K: tl.constexpr,
1147 +):
1148 + query_idx = tl.program_id(0)
1149 + pid_h = tl.program_id(1)
1150 +
1151 + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
1152 + head_mask = head_offsets < num_heads
1153 + nope_offsets = tl.arange(0, NOPE_BLOCK)
1154 + nope_mask = nope_offsets < NOPE_DIM
1155 + rope_offsets = tl.arange(0, ROPE_DIM)
1156 +
1157 + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1
1158 + q_nope = tl.load(
1159 + q_row_ptr + nope_offsets[None, :],
1160 + mask=head_mask[:, None] & nope_mask[None, :],
1161 + other=0.0,
1162 + )
1163 + q_rope = tl.load(
1164 + q_row_ptr + NOPE_DIM + rope_offsets[None, :],
1165 + mask=head_mask[:, None],
1166 + other=0.0,
1167 + )
1168 +
1169 + neg_large = -3.4028234663852886e38
1170 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
1171 + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
1172 + acc_nope = tl.zeros((BLOCK_H, NOPE_BLOCK), dtype=tl.float32)
1173 + acc_rope = tl.zeros((BLOCK_H, ROPE_DIM), dtype=tl.float32)
1174 + k_offsets = tl.arange(0, BLOCK_K)
1175 +
1176 + main_start = tl.load(main_indptr_ptr + query_idx)
1177 + main_end = tl.load(main_indptr_ptr + query_idx + 1)
1178 + main_len = main_end - main_start
1179 +
1180 + zero_nope = tl.zeros((BLOCK_K, NOPE_BLOCK), dtype=tl.bfloat16)
1181 + zero_rope = tl.zeros((BLOCK_K, ROPE_DIM), dtype=tl.bfloat16)
1182 +
1183 + for k_start in tl.range(0, main_len, BLOCK_K):
1184 + k_pos = k_start + k_offsets
1185 + in_range = k_pos < main_len
1186 + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1)
1187 + valid = in_range & (slot >= 0) & (slot < main_num_rows)
1188 + safe_slot = tl.where(valid, slot, 0)
1189 +
1190 + block_idx = safe_slot // main_block_size
1191 + pos_in_block = safe_slot % main_block_size
1192 + cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0
1193 + token_data_ptr = cache_block_ptr + pos_in_block * 576
1194 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8
1195 +
1196 + x_uint8 = tl.load(
1197 + token_data_ptr[:, None] + nope_offsets[None, :],
1198 + mask=valid[:, None] & nope_mask[None, :],
1199 + other=0,
1200 + )
1201 + if IS_FNUZ:
1202 + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)
1203 + else:
1204 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
1205 + encoded_scales = tl.load(
1206 + token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
1207 + mask=valid[:, None] & nope_mask[None, :],
1208 + other=127,
1209 + )
1210 + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
1211 + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
1212 + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
1213 + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
1214 +
1215 + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
1216 + k_rope = tl.load(
1217 + rope_ptr[:, None] + rope_offsets[None, :],
1218 + mask=valid[:, None],
1219 + other=0.0,
1220 + )
1221 + k_rope = tl.where(valid[:, None], k_rope, zero_rope)
1222 + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
1223 +
1224 + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(q_rope, tl.trans(k_rope))
1225 + scores *= scale
1226 + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
1227 +
1228 + m_block = tl.max(scores, axis=1)
1229 + m_new = tl.maximum(m_i, m_block)
1230 + alpha = tl.exp(m_i - m_new)
1231 + p = tl.exp(scores - m_new[:, None])
1232 + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
1233 + l_new = l_i * alpha + tl.sum(p, axis=1)
1234 +
1235 + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
1236 + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
1237 + m_i = m_new
1238 + l_i = l_new
1239 +
1240 + if HAS_EXTRA:
1241 + extra_start = tl.load(extra_indptr_ptr + query_idx)
1242 + extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
1243 + extra_len = extra_end - extra_start
1244 +
1245 + for k_start in tl.range(0, extra_len, BLOCK_K):
1246 + k_pos = k_start + k_offsets
1247 + in_range = k_pos < extra_len
1248 + slot = tl.load(
1249 + extra_indices_ptr + extra_start + k_pos, mask=in_range, other=-1
1250 + )
1251 + valid = in_range & (slot >= 0) & (slot < extra_num_rows)
1252 + safe_slot = tl.where(valid, slot, 0)
1253 +
1254 + block_idx = safe_slot // extra_block_size
1255 + pos_in_block = safe_slot % extra_block_size
1256 + cache_block_ptr = (
1257 + extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0
1258 + )
1259 + token_data_ptr = cache_block_ptr + pos_in_block * 576
1260 + token_scale_ptr = (
1261 + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8
1262 + )
1263 +
1264 + x_uint8 = tl.load(
1265 + token_data_ptr[:, None] + nope_offsets[None, :],
1266 + mask=valid[:, None] & nope_mask[None, :],
1267 + other=0,
1268 + )
1269 + if IS_FNUZ:
1270 + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)
1271 + else:
1272 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
1273 + encoded_scales = tl.load(
1274 + token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
1275 + mask=valid[:, None] & nope_mask[None, :],
1276 + other=127,
1277 + )
1278 + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
1279 + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
1280 + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
1281 + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
1282 +
1283 + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
1284 + k_rope = tl.load(
1285 + rope_ptr[:, None] + rope_offsets[None, :],
1286 + mask=valid[:, None],
1287 + other=0.0,
1288 + )
1289 + k_rope = tl.where(valid[:, None], k_rope, zero_rope)
1290 + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
1291 +
1292 + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(
1293 + q_rope,
1294 + tl.trans(k_rope),
1295 + )
1296 + scores *= scale
1297 + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
1298 +
1299 + m_block = tl.max(scores, axis=1)
1300 + m_new = tl.maximum(m_i, m_block)
1301 + alpha = tl.exp(m_i - m_new)
1302 + p = tl.exp(scores - m_new[:, None])
1303 + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
1304 + l_new = l_i * alpha + tl.sum(p, axis=1)
1305 +
1306 + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
1307 + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
1308 + m_i = m_new
1309 + l_i = l_new
1310 +
1311 + if HAS_ATTN_SINK:
1312 + sink = tl.load(
1313 + attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large
1314 + ).to(tl.float32)
1315 + m_final = tl.maximum(m_i, sink)
1316 + alpha = tl.exp(m_i - m_final)
1317 + l_final = l_i * alpha + tl.exp(sink - m_final)
1318 + denom = tl.maximum(l_final, 1.0e-30)
1319 + out_nope = tl.where(
1320 + l_final[:, None] > 0.0,
1321 + (acc_nope * alpha[:, None]) / denom[:, None],
1322 + 0.0,
1323 + )
1324 + out_rope = tl.where(
1325 + l_final[:, None] > 0.0,
1326 + (acc_rope * alpha[:, None]) / denom[:, None],
1327 + 0.0,
1328 + )
1329 + else:
1330 + denom = tl.maximum(l_i, 1.0e-30)
1331 + out_nope = tl.where(l_i[:, None] > 0.0, acc_nope / denom[:, None], 0.0)
1332 + out_rope = tl.where(l_i[:, None] > 0.0, acc_rope / denom[:, None], 0.0)
1333 +
1334 + out_row_ptr = (
1335 + out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1
1336 + )
1337 + tl.store(
1338 + out_row_ptr + nope_offsets[None, :],
1339 + out_nope,
1340 + mask=head_mask[:, None] & nope_mask[None, :],
1341 + )
1342 + tl.store(
1343 + out_row_ptr + NOPE_DIM + rope_offsets[None, :],
1344 + out_rope,
1345 + mask=head_mask[:, None],
1346 + )
1347 +
1348 +
1349 +def _rocm_sparse_attn_prefill_ragged_triton(
1350 q: torch.Tensor,
1351 kv: torch.Tensor,
1352 indices: torch.Tensor,
912 - topk_length: torch.Tensor | None,
1353 + indptr: torch.Tensor,
1354 scale: float,
914 - head_dim: int,
1355 attn_sink: torch.Tensor | None,
1356 + nope_head_dim: int,
1357 + rope_head_dim: int,
1358 ) -> torch.Tensor:
917 - indices = indices.clone().squeeze(1)
918 - s_q, h_q, d_qk = q.shape
919 - topk = indices.shape[-1]
920 - s_kv = kv.shape[0]
921 - if topk_length is not None:
922 - mask = torch.arange(topk, device=indices.device).unsqueeze(
923 - 0
924 - ) >= topk_length.unsqueeze(1)
925 - indices[mask] = -1
926 - invalid_mask = (indices < 0) | (indices >= s_kv)
927 - indices[invalid_mask] = 0
928 -
929 - qf = q.float()
930 - gathered_kv = kv.index_select(0, indices.flatten()).reshape(s_q, topk, d_qk).float()
931 - scores = qf @ gathered_kv.transpose(1, 2)
932 - scores *= scale
933 - scores[invalid_mask.unsqueeze(1).expand_as(scores)] = float("-inf")
934 -
935 - orig_lse = torch.logsumexp(scores, dim=-1)
936 - lse_for_o = orig_lse
937 - if attn_sink is not None:
938 - lse_for_o = torch.logsumexp(
939 - torch.stack(
940 - [orig_lse, attn_sink[:h_q].view(1, h_q).expand_as(orig_lse)],
941 - dim=0,
942 - ),
943 - dim=0,
944 - )
945 - lse_for_o = lse_for_o.clone()
946 - lse_for_o[lse_for_o == float("-inf")] = float("+inf")
947 - probs = torch.exp(scores - lse_for_o.unsqueeze(-1))
948 - out = probs @ gathered_kv[..., :head_dim]
949 - lonely_q_mask = orig_lse == float("-inf")
950 - out[lonely_q_mask.unsqueeze(-1).expand_as(out)] = 0.0
951 - return out.to(torch.bfloat16)
1359 + assert q.ndim == 3, f"expected q=[sq,h,d], got {q.shape}"
1360 + assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}"
1361 + assert indices.ndim == 1, f"expected indices=[nnz], got {indices.shape}"
1362 + assert indptr.ndim == 1, f"expected indptr=[sq+1], got {indptr.shape}"
1363 + assert q.is_cuda and kv.is_cuda and indices.is_cuda and indptr.is_cuda
1364 +
1365 + indices = _as_int32_contiguous_1d(indices)
1366 + indptr = _as_int32_contiguous_1d(indptr)
1367 + has_attn_sink = attn_sink is not None
1368 + if attn_sink is None:
1369 + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32)
1370 + else:
1371 + attn_sink = attn_sink.contiguous()
1372
1373 + num_queries, num_heads, head_dim = q.shape
1374 + assert indptr.numel() == num_queries + 1, (
1375 + f"expected indptr shape [{num_queries + 1}], got {indptr.shape}"
1376 + )
1377 + _validate_dsv4_sparse_dims(
1378 + head_dim,
1379 + nope_head_dim,
1380 + rope_head_dim,
1381 + "_rocm_sparse_attn_prefill_ragged_triton",
1382 + )
1383
954 -def rocm_sparse_attn_prefill(
1384 + block_h = 16
1385 + block_d = triton.next_power_of_2(head_dim)
1386 + block_k = 16 if head_dim >= 256 else 32
1387 + out = torch.empty_like(q, dtype=torch.bfloat16)
1388 + _sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
1389 + q,
1390 + kv,
1391 + indices,
1392 + indptr,
1393 + attn_sink,
1394 + out,
1395 + q.stride(0),
1396 + q.stride(1),
1397 + q.stride(2),
1398 + kv.stride(0),
1399 + kv.stride(1),
1400 + out.stride(0),
1401 + out.stride(1),
1402 + out.stride(2),
1403 + num_heads,
1404 + head_dim,
1405 + kv.shape[0],
1406 + float(scale),
1407 + HAS_ATTN_SINK=has_attn_sink,
1408 + BLOCK_H=block_h,
1409 + BLOCK_D=block_d,
1410 + BLOCK_K=block_k,
1411 + num_warps=8,
1412 + )
1413 + return out
1414 +
1415 +
1416 +def _rocm_sparse_attn_prefill_triton(
1417 q: torch.Tensor,
1418 kv: torch.Tensor,
1419 indices: torch.Tensor,
958 - topk_length: torch.Tensor | None,
1420 scale: float,
960 - head_dim: int,
1421 attn_sink: torch.Tensor | None,
962 - output: torch.Tensor,
963 -) -> None:
964 - output_chunk = rocm_ref_sparse_attn_prefill(
1422 + nope_head_dim: int,
1423 + rope_head_dim: int,
1424 + topk_length: torch.Tensor | None = None,
1425 +) -> torch.Tensor:
1426 + ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
1427 + indices,
1428 + topk_length
1429 + if topk_length is not None
1430 + else (indices >= 0).sum(dim=-1, dtype=torch.int32),
1431 + num_rows=kv.shape[0],
1432 + )
1433 + return _rocm_sparse_attn_prefill_ragged_triton(
1434 q=q,
1435 kv=kv,
967 - indices=indices,
968 - topk_length=topk_length,
1436 + indices=ragged_indices,
1437 + indptr=ragged_indptr,
1438 scale=scale,
970 - head_dim=head_dim,
1439 attn_sink=attn_sink,
1440 + nope_head_dim=nope_head_dim,
1441 + rope_head_dim=rope_head_dim,
1442 )
973 - output.copy_(output_chunk.to(output.dtype))
1443
1444
976 -def rocm_dequantize_blocked_k_cache(
977 - quant_k_cache: torch.Tensor,
978 - head_dim: int,
1445 +def _rocm_sparse_attn_decode_ragged_triton(
1446 + q: torch.Tensor,
1447 + main_cache: torch.Tensor,
1448 + main_indices: torch.Tensor,
1449 + main_indptr: torch.Tensor,
1450 + scale: float,
1451 + attn_sink: torch.Tensor | None,
1452 nope_head_dim: int,
1453 rope_head_dim: int,
1454 + extra_cache: torch.Tensor | None = None,
1455 + extra_indices: torch.Tensor | None = None,
1456 + extra_indptr: torch.Tensor | None = None,
1457 ) -> torch.Tensor:
982 - fp8_dtype = current_platform.fp8_dtype()
983 - tile_size = 64
984 - num_tiles = nope_head_dim // tile_size
985 -
986 - num_blocks, block_size, _ = quant_k_cache.shape
987 - quant_k_cache = quant_k_cache.view(num_blocks, -1)
988 - input_nope_rope = quant_k_cache[
989 - :, : block_size * (nope_head_dim + 2 * rope_head_dim)
990 - ].view(num_blocks, block_size, nope_head_dim + 2 * rope_head_dim)
991 - input_nope = input_nope_rope[:, :, :nope_head_dim].view(fp8_dtype)
992 - input_rope = input_nope_rope[:, :, nope_head_dim:].view(torch.bfloat16)
993 - input_scale = (
994 - quant_k_cache[:, block_size * (nope_head_dim + 2 * rope_head_dim) :]
995 - .view(num_blocks, block_size, 8)[:, :, :num_tiles]
996 - .view(torch.float8_e8m0fnu)
1458 + assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}"
1459 + assert main_cache.ndim == 3, (
1460 + f"expected main_cache=[blocks,block,bytes], got {main_cache.shape}"
1461 + )
1462 + assert main_indices.ndim == 1, (
1463 + f"expected main_indices=[nnz], got {main_indices.shape}"
1464 + )
1465 + assert main_indptr.ndim == 1, f"expected main_indptr=[b+1], got {main_indptr.shape}"
1466 + assert (
1467 + q.is_cuda
1468 + and main_cache.is_cuda
1469 + and main_indices.is_cuda
1470 + and main_indptr.is_cuda
1471 )
1472
999 - result = torch.empty(
1000 - (num_blocks, block_size, 1, head_dim),
1001 - dtype=torch.bfloat16,
1002 - device=quant_k_cache.device,
1473 + main_indices = _as_int32_contiguous_1d(main_indices)
1474 + main_indptr = _as_int32_contiguous_1d(main_indptr)
1475 + has_attn_sink = attn_sink is not None
1476 + if attn_sink is None:
1477 + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32)
1478 + else:
1479 + attn_sink = attn_sink.contiguous()
1480 +
1481 + num_queries, num_heads, head_dim = q.shape
1482 + assert main_indptr.numel() == num_queries + 1, (
1483 + f"expected main_indptr shape [{num_queries + 1}], got {main_indptr.shape}"
1484 + )
1485 + _validate_dsv4_sparse_dims(
1486 + head_dim,
1487 + nope_head_dim,
1488 + rope_head_dim,
1489 + "_rocm_sparse_attn_decode_ragged_triton",
1490 + )
1491 +
1492 + has_extra = (
1493 + extra_cache is not None
1494 + and extra_indices is not None
1495 + and extra_indptr is not None
1496 + )
1497 + if has_extra:
1498 + assert extra_cache is not None
1499 + assert extra_indices is not None
1500 + assert extra_indptr is not None
1501 + assert extra_indices.ndim == 1, (
1502 + f"expected extra_indices=[nnz], got {extra_indices.shape}"
1503 + )
1504 + assert extra_indptr.ndim == 1, (
1505 + f"expected extra_indptr=[b+1], got {extra_indptr.shape}"
1506 + )
1507 + extra_indices = _as_int32_contiguous_1d(extra_indices)
1508 + extra_indptr = _as_int32_contiguous_1d(extra_indptr)
1509 + assert extra_indptr.numel() == num_queries + 1, (
1510 + f"expected extra_indptr shape [{num_queries + 1}], got {extra_indptr.shape}"
1511 + )
1512 + else:
1513 + extra_cache = main_cache
1514 + extra_indices = torch.empty(0, device=q.device, dtype=torch.int32)
1515 + extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32)
1516 +
1517 + block_h = 16
1518 + block_k = 16 if head_dim >= 256 else 32
1519 + out = torch.empty_like(q, dtype=torch.bfloat16)
1520 + _sparse_attn_decode_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
1521 + q,
1522 + main_cache,
1523 + main_indices,
1524 + main_indptr,
1525 + extra_cache,
1526 + extra_indices,
1527 + extra_indptr,
1528 + attn_sink,
1529 + out,
1530 + q.stride(0),
1531 + q.stride(1),
1532 + out.stride(0),
1533 + out.stride(1),
1534 + main_cache.stride(0),
1535 + extra_cache.stride(0),
1536 + main_cache.shape[0] * main_cache.shape[1],
1537 + extra_cache.shape[0] * extra_cache.shape[1],
1538 + main_cache.shape[1],
1539 + extra_cache.shape[1],
1540 + scale,
1541 + num_heads,
1542 + HAS_ATTN_SINK=has_attn_sink,
1543 + HAS_EXTRA=has_extra,
1544 + NOPE_DIM=nope_head_dim,
1545 + NOPE_BLOCK=triton.next_power_of_2(nope_head_dim),
1546 + ROPE_DIM=rope_head_dim,
1547 + IS_FNUZ=current_platform.is_fp8_fnuz(),
1548 + BLOCK_H=block_h,
1549 + BLOCK_K=block_k,
1550 + num_warps=8,
1551 )
1004 - result[..., nope_head_dim:] = input_rope.unsqueeze(2)
1005 - for tile_idx in range(num_tiles):
1006 - cur_nope = input_nope[
1007 - ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
1008 - ].to(torch.bfloat16)
1009 - cur_scales = input_scale[:, :, tile_idx].to(torch.bfloat16).unsqueeze(-1)
1010 - result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
1011 - cur_nope * cur_scales
1012 - ).unsqueeze(2)
1013 - return result
1014 -
1015 -
1016 -def rocm_ref_sparse_attn_decode(
1552 + return out
1553 +
1554 +
1555 +def _rocm_sparse_attn_decode_triton(
1556 q: torch.Tensor,
1018 - blocked_k: torch.Tensor,
1019 - indices_in_kvcache: torch.Tensor,
1020 - topk_length: torch.Tensor | None,
1557 + main_cache: torch.Tensor,
1558 + main_indices: torch.Tensor,
1559 scale: float,
1022 - head_dim: int,
1560 attn_sink: torch.Tensor | None,
1024 - extra_blocked_k: torch.Tensor | None = None,
1025 - extra_indices_in_kvcache: torch.Tensor | None = None,
1026 - extra_topk_length: torch.Tensor | None = None,
1561 + nope_head_dim: int,
1562 + rope_head_dim: int,
1563 + extra_cache: torch.Tensor | None = None,
1564 + extra_indices: torch.Tensor | None = None,
1565 + main_lengths: torch.Tensor | None = None,
1566 + extra_lengths: torch.Tensor | None = None,
1567 + main_ragged_indices: torch.Tensor | None = None,
1568 + main_ragged_indptr: torch.Tensor | None = None,
1569 + extra_ragged_indices: torch.Tensor | None = None,
1570 + extra_ragged_indptr: torch.Tensor | None = None,
1571 ) -> torch.Tensor:
1028 - b, s_q, h_q, d_qk = q.shape
1029 -
1030 - def process_scope(
1031 - cur_blocked_k: torch.Tensor,
1032 - cur_indices: torch.Tensor,
1033 - cur_topk_length: torch.Tensor | None,
1034 - ) -> tuple[torch.Tensor, torch.Tensor]:
1035 - cur_indices = cur_indices.reshape(b, s_q, -1)
1036 - topk = cur_indices.size(-1)
1037 - fixed_indices = torch.clamp_min(cur_indices, 0)
1038 - gathered_kv = (
1039 - cur_blocked_k.view(-1, d_qk)
1040 - .index_select(0, fixed_indices.view(-1))
1041 - .view(b, s_q, topk, d_qk)
1572 + if main_ragged_indices is None or main_ragged_indptr is None:
1573 + main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense(
1574 + main_indices,
1575 + main_lengths
1576 + if main_lengths is not None
1577 + else (main_indices >= 0).sum(dim=-1, dtype=torch.int32),
1578 + num_rows=main_cache.shape[0] * main_cache.shape[1],
1579 + )
1580 +
1581 + if (
1582 + (extra_ragged_indices is None or extra_ragged_indptr is None)
1583 + and extra_cache is not None
1584 + and extra_indices is not None
1585 + ):
1586 + extra_ragged_indices, extra_ragged_indptr = build_ragged_indices_from_dense(
1587 + extra_indices,
1588 + extra_lengths
1589 + if extra_lengths is not None
1590 + else (extra_indices >= 0).sum(dim=-1, dtype=torch.int32),
1591 + num_rows=extra_cache.shape[0] * extra_cache.shape[1],
1592 )
1043 - invalid_mask = cur_indices == -1
1044 - if cur_topk_length is not None:
1045 - cur_topk_length = cur_topk_length.reshape(b)
1046 - invalid_mask |= torch.arange(0, topk, device=invalid_mask.device).view(
1047 - 1, 1, topk
1048 - ) >= cur_topk_length.view(b, 1, 1)
1049 - return gathered_kv, invalid_mask
1050 -
1051 - gathered_kv, invalid_mask = process_scope(
1052 - blocked_k, indices_in_kvcache, topk_length
1593 +
1594 + return _rocm_sparse_attn_decode_ragged_triton(
1595 + q=q,
1596 + main_cache=main_cache,
1597 + main_indices=main_ragged_indices,
1598 + main_indptr=main_ragged_indptr,
1599 + scale=scale,
1600 + attn_sink=attn_sink,
1601 + nope_head_dim=nope_head_dim,
1602 + rope_head_dim=rope_head_dim,
1603 + extra_cache=extra_cache,
1604 + extra_indices=extra_ragged_indices,
1605 + extra_indptr=extra_ragged_indptr,
1606 + )
1607 +
1608 +
1609 +def rocm_sparse_attn_prefill(
1610 + q: torch.Tensor,
1611 + kv: torch.Tensor,
1612 + indices: torch.Tensor,
1613 + topk_length: torch.Tensor | None,
1614 + scale: float,
1615 + head_dim: int,
1616 + nope_head_dim: int,
1617 + rope_head_dim: int,
1618 + attn_sink: torch.Tensor | None,
1619 + output: torch.Tensor,
1620 + ragged_indices: torch.Tensor | None = None,
1621 + ragged_indptr: torch.Tensor | None = None,
1622 +) -> None:
1623 + assert kv.ndim == 3 and kv.shape[1] == 1, (
1624 + f"ROCm Triton sparse prefill expects kv=[skv,1,d], got {kv.shape}"
1625 + )
1626 + _validate_dsv4_sparse_dims(
1627 + head_dim,
1628 + nope_head_dim,
1629 + rope_head_dim,
1630 + "rocm_sparse_attn_prefill",
1631 )
1054 - if extra_blocked_k is not None:
1055 - assert extra_indices_in_kvcache is not None
1056 - gathered_kv1, invalid_mask1 = process_scope(
1057 - extra_blocked_k, extra_indices_in_kvcache, extra_topk_length
1632 + if ragged_indices is not None and ragged_indptr is not None:
1633 + output_chunk = _rocm_sparse_attn_prefill_ragged_triton(
1634 + q=q,
1635 + kv=kv.squeeze(1),
1636 + indices=ragged_indices,
1637 + indptr=ragged_indptr,
1638 + scale=scale,
1639 + attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
1640 + nope_head_dim=nope_head_dim,
1641 + rope_head_dim=rope_head_dim,
1642 )
1059 - gathered_kv = torch.cat([gathered_kv, gathered_kv1], dim=2)
1060 - invalid_mask = torch.cat([invalid_mask, invalid_mask1], dim=2)
1061 -
1062 - gathered_kv = gathered_kv.view(b * s_q, -1, d_qk).float()
1063 - gathered_kv[gathered_kv != gathered_kv] = 0.0
1064 - qf = q.float().view(b * s_q, h_q, d_qk)
1065 - attn_weight = qf @ gathered_kv.transpose(-1, -2)
1066 - attn_weight *= scale
1067 - attn_weight[
1068 - invalid_mask.view(b * s_q, 1, -1).expand(b * s_q, h_q, invalid_mask.size(-1))
1069 - ] = float("-inf")
1070 - lse = attn_weight.logsumexp(dim=-1)
1071 - attn_weight = torch.exp(attn_weight - lse.unsqueeze(-1))
1072 - output = attn_weight @ gathered_kv[..., :head_dim]
1073 - output = output.view(b, s_q, h_q, head_dim)
1074 - lse = lse.view(b, s_q, h_q)
1075 -
1076 - if attn_sink is not None:
1077 - output *= (1.0 / (1.0 + torch.exp(attn_sink.view(1, 1, h_q) - lse))).unsqueeze(
1078 - -1
1643 + else:
1644 + indices_2d = indices.reshape(indices.shape[0], -1)
1645 + output_chunk = _rocm_sparse_attn_prefill_triton(
1646 + q=q,
1647 + kv=kv.squeeze(1),
1648 + indices=indices_2d,
1649 + scale=scale,
1650 + attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
1651 + nope_head_dim=nope_head_dim,
1652 + rope_head_dim=rope_head_dim,
1653 + topk_length=topk_length,
1654 )
1080 -
1081 - lonely_q_mask = lse == float("-inf")
1082 - output[lonely_q_mask.unsqueeze(-1).expand_as(output)] = 0.0
1083 - return output.squeeze(1).to(torch.bfloat16)
1655 + output.copy_(output_chunk.to(output.dtype))
1656
1657
1086 -def rocm_forward_decode_fallback(
1658 +def rocm_sparse_attn_decode(
1659 q: torch.Tensor,
1660 kv_cache: torch.Tensor | None,
1661 swa_k_cache: torch.Tensor,
@@ -1092,6 +1664,10 @@ def rocm_forward_decode_fallback(
1664 topk_lens: torch.Tensor | None,
1665 swa_indices: torch.Tensor,
1666 swa_lens: torch.Tensor,
1667 + swa_ragged_indices: torch.Tensor | None,
1668 + swa_ragged_indptr: torch.Tensor | None,
1669 + topk_ragged_indices: torch.Tensor | None,
1670 + topk_ragged_indptr: torch.Tensor | None,
1671 attn_sink: torch.Tensor | None,
1672 scale: float,
1673 head_dim: int,
@@ -1099,31 +1675,49 @@ def rocm_forward_decode_fallback(
1675 rope_head_dim: int,
1676 output: torch.Tensor,
1677 ) -> None:
1102 - blocked_swa = rocm_dequantize_blocked_k_cache(
1103 - swa_k_cache,
1104 - head_dim=head_dim,
1105 - nope_head_dim=nope_head_dim,
1106 - rope_head_dim=rope_head_dim,
1678 + assert swa_k_cache.dtype == torch.uint8, (
1679 + "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, "
1680 + f"got {swa_k_cache.dtype}"
1681 + )
1682 + _validate_dsv4_sparse_dims(
1683 + head_dim,
1684 + nope_head_dim,
1685 + rope_head_dim,
1686 + "rocm_sparse_attn_decode",
1687 )
1108 - blocked_extra = None
1688 +
1689 + main_indices = swa_indices.reshape(swa_indices.shape[0], -1)
1690 +
1691 + extra_cache = None
1692 + extra_indices = None
1693 if not swa_only:
1694 assert kv_cache is not None
1111 - blocked_extra = rocm_dequantize_blocked_k_cache(
1112 - kv_cache,
1113 - head_dim=head_dim,
1114 - nope_head_dim=nope_head_dim,
1115 - rope_head_dim=rope_head_dim,
1695 + assert topk_indices is not None or (
1696 + topk_ragged_indices is not None and topk_ragged_indptr is not None
1697 )
1117 - attn_out = rocm_ref_sparse_attn_decode(
1118 - q=q.unsqueeze(1),
1119 - blocked_k=blocked_swa,
1120 - indices_in_kvcache=swa_indices.unsqueeze(1),
1121 - topk_length=swa_lens,
1698 + assert kv_cache.dtype == torch.uint8, (
1699 + "ROCm Triton sparse decode expects uint8 fp8_ds_mla extra cache, "
1700 + f"got {kv_cache.dtype}"
1701 + )
1702 + extra_cache = kv_cache
1703 + if topk_indices is not None:
1704 + extra_indices = topk_indices.reshape(topk_indices.shape[0], -1)
1705 +
1706 + attn_out = _rocm_sparse_attn_decode_triton(
1707 + q=q,
1708 + main_cache=swa_k_cache,
1709 + main_indices=main_indices,
1710 scale=scale,
1123 - head_dim=head_dim,
1124 - attn_sink=attn_sink[: q.shape[1]] if attn_sink is not None else None,
1125 - extra_blocked_k=blocked_extra,
1126 - extra_indices_in_kvcache=topk_indices,
1127 - extra_topk_length=topk_lens,
1711 + attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
1712 + nope_head_dim=nope_head_dim,
1713 + rope_head_dim=rope_head_dim,
1714 + extra_cache=extra_cache,
1715 + extra_indices=extra_indices,
1716 + main_lengths=swa_lens,
1717 + extra_lengths=topk_lens,
1718 + main_ragged_indices=swa_ragged_indices,
1719 + main_ragged_indptr=swa_ragged_indptr,
1720 + extra_ragged_indices=topk_ragged_indices,
1721 + extra_ragged_indptr=topk_ragged_indptr,
1722 )
1723 output.copy_(attn_out.to(output.dtype))