@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6c20e89c02
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 9, 21:48 UTC

[ROCm][Deepseekv3.2] Refactor Sparse Indexer as CustomOp (#29287)

Signed-off-by: ganyi <ygan@amd.com>

Pleaplusone committed Jan 21, 2026 at 23:16 UTC 6c20e89c0209284c1f6c99ecd59c786a78f81b7d
8 files changed +982 -323
vllm/_aiter_ops.py
+12
@@ -9,6 +9,10 @@ from torch._ops import OpOverload
9 import vllm.envs as envs
10 from vllm.platforms import current_platform
11 from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
12 +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
13 + rocm_aiter_sparse_attn_indexer,
14 + rocm_aiter_sparse_attn_indexer_fake,
15 +)
16
17 _FP8_DTYPE = current_platform.fp8_dtype()
18
@@ -1091,6 +1095,14 @@ class rocm_aiter_ops:
1095 dispatch_key=current_platform.dispatch_key,
1096 )
1097
1098 + direct_register_custom_op(
1099 + op_name="rocm_aiter_sparse_attn_indexer",
1100 + op_func=rocm_aiter_sparse_attn_indexer,
1101 + mutates_args=["topk_indices_buffer"],
1102 + fake_impl=rocm_aiter_sparse_attn_indexer_fake,
1103 + dispatch_key=current_platform.dispatch_key,
1104 + )
1105 +
1106 _OPS_REGISTERED = True
1107
1108 @staticmethod
vllm/config/compilation.py
+1
@@ -611,6 +611,7 @@ class CompilationConfig:
611 "vllm::gdn_attention_core",
612 "vllm::kda_attention",
613 "vllm::sparse_attn_indexer",
614 + "vllm::rocm_aiter_sparse_attn_indexer",
615 ]
616
617 def compute_hash(self) -> str:
vllm/model_executor/layers/sparse_attn_indexer.py new
+318
@@ -0,0 +1,318 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Custom Sparse Attention Indexer layers."""
4 +
5 +import torch
6 +
7 +from vllm._aiter_ops import rocm_aiter_ops
8 +from vllm.forward_context import get_forward_context
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.custom_op import CustomOp
11 +from vllm.platforms import current_platform
12 +from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
13 +from vllm.utils.torch_utils import direct_register_custom_op
14 +from vllm.v1.attention.backends.mla.indexer import (
15 + DeepseekV32IndexerMetadata,
16 +)
17 +from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
18 +from vllm.v1.worker.workspace import current_workspace_manager
19 +
20 +if current_platform.is_cuda_alike():
21 + from vllm import _custom_ops as ops
22 +elif current_platform.is_xpu():
23 + from vllm._ipex_ops import ipex_ops as ops
24 +
25 +logger = init_logger(__name__)
26 +
27 +
28 +def sparse_attn_indexer(
29 + hidden_states: torch.Tensor,
30 + k_cache_prefix: str,
31 + kv_cache: torch.Tensor,
32 + q_fp8: torch.Tensor,
33 + k: torch.Tensor,
34 + weights: torch.Tensor,
35 + quant_block_size: int,
36 + scale_fmt: str | None,
37 + topk_tokens: int,
38 + head_dim: int,
39 + max_model_len: int,
40 + total_seq_lens: int,
41 + topk_indices_buffer: torch.Tensor,
42 +) -> torch.Tensor:
43 + # careful! this will be None in dummy run
44 + attn_metadata = get_forward_context().attn_metadata
45 + fp8_dtype = current_platform.fp8_dtype()
46 +
47 + # assert isinstance(attn_metadata, dict)
48 + if not isinstance(attn_metadata, dict):
49 + # Reserve workspace for indexer during profiling run
50 + current_workspace_manager().get_simultaneous(
51 + ((total_seq_lens, head_dim), torch.float8_e4m3fn),
52 + ((total_seq_lens, 4), torch.uint8),
53 + )
54 + return sparse_attn_indexer_fake(
55 + hidden_states,
56 + k_cache_prefix,
57 + kv_cache,
58 + q_fp8,
59 + k,
60 + weights,
61 + quant_block_size,
62 + scale_fmt,
63 + topk_tokens,
64 + head_dim,
65 + max_model_len,
66 + total_seq_lens,
67 + topk_indices_buffer,
68 + )
69 + attn_metadata = attn_metadata[k_cache_prefix]
70 + assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
71 + slot_mapping = attn_metadata.slot_mapping
72 + has_decode = attn_metadata.num_decodes > 0
73 + has_prefill = attn_metadata.num_prefills > 0
74 + num_decode_tokens = attn_metadata.num_decode_tokens
75 +
76 + ops.indexer_k_quant_and_cache(
77 + k,
78 + kv_cache,
79 + slot_mapping,
80 + quant_block_size,
81 + scale_fmt,
82 + )
83 +
84 + topk_indices_buffer[: hidden_states.shape[0]] = -1
85 + if has_prefill:
86 + prefill_metadata = attn_metadata.prefill
87 +
88 + # Get the full shared workspace buffers once (will allocate on first use)
89 + workspace_manager = current_workspace_manager()
90 + k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
91 + ((total_seq_lens, head_dim), fp8_dtype),
92 + ((total_seq_lens, 4), torch.uint8),
93 + )
94 + for chunk in prefill_metadata.chunks:
95 + k_fp8 = k_fp8_full[: chunk.total_seq_lens]
96 + k_scale = k_scale_full[: chunk.total_seq_lens]
97 + ops.cp_gather_indexer_k_quant_cache(
98 + kv_cache,
99 + k_fp8,
100 + k_scale,
101 + chunk.block_table,
102 + chunk.cu_seq_lens,
103 + )
104 +
105 + logits = fp8_mqa_logits(
106 + q_fp8[chunk.token_start : chunk.token_end],
107 + (k_fp8, k_scale.view(torch.float32).flatten()),
108 + weights[chunk.token_start : chunk.token_end],
109 + chunk.cu_seqlen_ks,
110 + chunk.cu_seqlen_ke,
111 + )
112 + num_rows = logits.shape[0]
113 +
114 + topk_indices = topk_indices_buffer[
115 + chunk.token_start : chunk.token_end, :topk_tokens
116 + ]
117 + torch.ops._C.top_k_per_row_prefill(
118 + logits,
119 + chunk.cu_seqlen_ks,
120 + chunk.cu_seqlen_ke,
121 + topk_indices,
122 + num_rows,
123 + logits.stride(0),
124 + logits.stride(1),
125 + topk_tokens,
126 + )
127 +
128 + if has_decode:
129 + decode_metadata = attn_metadata.decode
130 + # kv_cache size requirement [num_block, block_size, n_head, head_dim],
131 + # we only have [num_block, block_size, head_dim],
132 + kv_cache = kv_cache.unsqueeze(-2)
133 + decode_lens = decode_metadata.decode_lens
134 + if decode_metadata.requires_padding:
135 + # pad in edge case where we have short chunked prefill length <
136 + # decode_threshold since we unstrictly split
137 + # prefill and decode by decode_threshold
138 + # (currently set to 1 + speculative tokens)
139 + padded_q_fp8_decode_tokens = pack_seq_triton(
140 + q_fp8[:num_decode_tokens], decode_lens
141 + )
142 + else:
143 + padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
144 + decode_lens.shape[0], -1, *q_fp8.shape[1:]
145 + )
146 + # TODO: move and optimize below logic with triton kernels
147 + batch_size = padded_q_fp8_decode_tokens.shape[0]
148 + next_n = padded_q_fp8_decode_tokens.shape[1]
149 + assert batch_size == decode_metadata.seq_lens.shape[0]
150 + num_padded_tokens = batch_size * next_n
151 +
152 + logits = fp8_paged_mqa_logits(
153 + padded_q_fp8_decode_tokens,
154 + kv_cache,
155 + weights[:num_padded_tokens],
156 + decode_metadata.seq_lens,
157 + decode_metadata.block_table,
158 + decode_metadata.schedule_metadata,
159 + max_model_len=max_model_len,
160 + )
161 +
162 + num_rows = logits.shape[0]
163 +
164 + topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
165 + torch.ops._C.top_k_per_row_decode(
166 + logits,
167 + next_n,
168 + decode_metadata.seq_lens,
169 + topk_indices,
170 + num_rows,
171 + logits.stride(0),
172 + logits.stride(1),
173 + topk_tokens,
174 + )
175 +
176 + if decode_metadata.requires_padding:
177 + # if padded, we need to unpack
178 + # the topk indices removing padded tokens
179 + topk_indices = unpack_seq_triton(
180 + topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
181 + decode_lens,
182 + )
183 + topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (
184 + topk_indices
185 + )
186 +
187 + return topk_indices_buffer
188 +
189 +
190 +def sparse_attn_indexer_fake(
191 + hidden_states: torch.Tensor,
192 + k_cache_prefix: str,
193 + kv_cache: torch.Tensor,
194 + q_fp8: torch.Tensor,
195 + k: torch.Tensor,
196 + weights: torch.Tensor,
197 + quant_block_size: int,
198 + scale_fmt: str | None,
199 + topk_tokens: int,
200 + head_dim: int,
201 + max_model_len: int,
202 + total_seq_lens: int,
203 + topk_indices_buffer: torch.Tensor | None,
204 +) -> torch.Tensor:
205 + return topk_indices_buffer
206 +
207 +
208 +direct_register_custom_op(
209 + op_name="sparse_attn_indexer",
210 + op_func=sparse_attn_indexer,
211 + mutates_args=["topk_indices_buffer"],
212 + fake_impl=sparse_attn_indexer_fake,
213 + dispatch_key=current_platform.dispatch_key,
214 +)
215 +
216 +
217 +@CustomOp.register("sparse_attn_indexer")
218 +class SparseAttnIndexer(CustomOp):
219 + """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a
220 + separate custom op since it involves heavy custom kernels like `mqa_logits`,
221 + `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires
222 + specific memory layout or implementation for different hardware backends to
223 + achieve optimal performance.
224 +
225 + For now, the default native path will use CUDA backend path. Other platform
226 + may requires add the corresponding Custom Op name `sparse_attn_indexer` to
227 + `custom_ops` in `CompilationConfig` to enable the platform specific path.
228 + """
229 +
230 + def __init__(
231 + self,
232 + k_cache,
233 + quant_block_size: int,
234 + scale_fmt: str,
235 + topk_tokens: int,
236 + head_dim: int,
237 + max_model_len: int,
238 + max_total_seq_len: int,
239 + topk_indices_buffer: torch.Tensor,
240 + ):
241 + super().__init__()
242 + self.k_cache = k_cache
243 + self.quant_block_size = quant_block_size
244 + self.scale_fmt = scale_fmt
245 + self.topk_tokens = topk_tokens
246 + self.head_dim = head_dim
247 + self.max_model_len = max_model_len
248 + self.max_total_seq_len = max_total_seq_len
249 + self.topk_indices_buffer = topk_indices_buffer
250 +
251 + def forward_native(
252 + self,
253 + hidden_states: torch.Tensor,
254 + q_fp8: torch.Tensor,
255 + k: torch.Tensor,
256 + weights: torch.Tensor,
257 + ):
258 + if current_platform.is_cuda():
259 + return self.forward_cuda(hidden_states, q_fp8, k, weights)
260 + elif current_platform.is_rocm():
261 + return self.forward_hip(hidden_states, q_fp8, k, weights)
262 + else:
263 + raise NotImplementedError(
264 + "SparseAttnIndexer native forward is only implemented for "
265 + "CUDA and ROCm platform."
266 + )
267 +
268 + def forward_cuda(
269 + self,
270 + hidden_states: torch.Tensor,
271 + q_fp8: torch.Tensor,
272 + k: torch.Tensor,
273 + weights: torch.Tensor,
274 + ):
275 + return torch.ops.vllm.sparse_attn_indexer(
276 + hidden_states,
277 + self.k_cache.prefix,
278 + self.k_cache.kv_cache[0],
279 + q_fp8,
280 + k,
281 + weights,
282 + self.quant_block_size,
283 + self.scale_fmt,
284 + self.topk_tokens,
285 + self.head_dim,
286 + self.max_model_len,
287 + self.max_total_seq_len,
288 + self.topk_indices_buffer,
289 + )
290 +
291 + def forward_hip(
292 + self,
293 + hidden_states: torch.Tensor,
294 + q_fp8: torch.Tensor,
295 + k: torch.Tensor,
296 + weights: torch.Tensor,
297 + ):
298 + if rocm_aiter_ops.is_enabled():
299 + return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
300 + hidden_states,
301 + self.k_cache.prefix,
302 + self.k_cache.kv_cache[0],
303 + q_fp8,
304 + k,
305 + weights,
306 + self.quant_block_size,
307 + self.scale_fmt,
308 + self.topk_tokens,
309 + self.head_dim,
310 + self.max_model_len,
311 + self.max_total_seq_len,
312 + self.topk_indices_buffer,
313 + )
314 + else:
315 + raise RuntimeError(
316 + "Sparse attention indexer ROCm custom op requires ROCm "
317 + "Aiter ops to be enabled."
318 + )
vllm/model_executor/models/deepseek_v2.py
+14 -233
@@ -43,7 +43,6 @@ from vllm.distributed import (
43 get_tensor_model_parallel_world_size,
44 tensor_model_parallel_all_gather,
45 )
46 -from vllm.forward_context import get_forward_context
46 from vllm.logger import init_logger
47 from vllm.model_executor.layers.activation import SiluAndMul
48 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -63,6 +62,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
62 per_token_group_quant_fp8,
63 )
64 from vllm.model_executor.layers.rotary_embedding import get_rope
65 +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
66 from vllm.model_executor.layers.vocab_parallel_embedding import (
67 ParallelLMHead,
68 VocabParallelEmbedding,
@@ -74,16 +74,11 @@ from vllm.model_executor.model_loader.weight_utils import (
74 from vllm.model_executor.models.utils import sequence_parallel_chunk
75 from vllm.platforms import current_platform
76 from vllm.sequence import IntermediateTensors
77 -from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
78 -from vllm.utils.torch_utils import direct_register_custom_op
77 from vllm.v1.attention.backend import AttentionBackend
78 from vllm.v1.attention.backends.mla.indexer import (
79 DeepseekV32IndexerBackend,
82 - DeepseekV32IndexerMetadata,
80 )
84 -from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
81 from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
86 -from vllm.v1.worker.workspace import current_workspace_manager
82
83 from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP
84 from .utils import (
@@ -94,11 +89,6 @@ from .utils import (
89 maybe_prefix,
90 )
91
97 -if current_platform.is_cuda_alike():
98 - from vllm import _custom_ops as ops
99 -elif current_platform.is_xpu():
100 - from vllm._ipex_ops import ipex_ops as ops
101 -
92 logger = init_logger(__name__)
93
94
@@ -599,213 +589,6 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
589 return DeepseekV32IndexerBackend
590
591
602 -def sparse_attn_indexer(
603 - hidden_states: torch.Tensor,
604 - k_cache_prefix: str,
605 - kv_cache: torch.Tensor,
606 - q_fp8: torch.Tensor,
607 - k: torch.Tensor,
608 - weights: torch.Tensor,
609 - quant_block_size: int,
610 - scale_fmt: str | None,
611 - topk_tokens: int,
612 - head_dim: int,
613 - max_model_len: int,
614 - total_seq_lens: int,
615 - topk_indices_buffer: torch.Tensor | None,
616 -) -> torch.Tensor:
617 - # careful! this will be None in dummy run
618 - attn_metadata = get_forward_context().attn_metadata
619 - fp8_dtype = current_platform.fp8_dtype()
620 -
621 - # assert isinstance(attn_metadata, dict)
622 - if not isinstance(attn_metadata, dict):
623 - # Reserve workspace for indexer during profiling run
624 - current_workspace_manager().get_simultaneous(
625 - ((total_seq_lens, head_dim), torch.float8_e4m3fn),
626 - ((total_seq_lens, 4), torch.uint8),
627 - )
628 -
629 - return sparse_attn_indexer_fake(
630 - hidden_states,
631 - k_cache_prefix,
632 - kv_cache,
633 - q_fp8,
634 - k,
635 - weights,
636 - quant_block_size,
637 - scale_fmt,
638 - topk_tokens,
639 - head_dim,
640 - max_model_len,
641 - total_seq_lens,
642 - topk_indices_buffer,
643 - )
644 - attn_metadata = attn_metadata[k_cache_prefix]
645 - assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
646 - slot_mapping = attn_metadata.slot_mapping
647 - has_decode = attn_metadata.num_decodes > 0
648 - has_prefill = attn_metadata.num_prefills > 0
649 - num_decode_tokens = attn_metadata.num_decode_tokens
650 -
651 - ops.indexer_k_quant_and_cache(
652 - k,
653 - kv_cache,
654 - slot_mapping,
655 - quant_block_size,
656 - scale_fmt,
657 - )
658 -
659 - topk_indices_buffer[: hidden_states.shape[0]] = -1
660 - if has_prefill:
661 - prefill_metadata = attn_metadata.prefill
662 -
663 - # Get the full shared workspace buffers once (will allocate on first use)
664 - workspace_manager = current_workspace_manager()
665 - k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
666 - ((total_seq_lens, head_dim), fp8_dtype),
667 - ((total_seq_lens, 4), torch.uint8),
668 - )
669 -
670 - for chunk in prefill_metadata.chunks:
671 - k_fp8 = k_fp8_full[: chunk.total_seq_lens]
672 - k_scale = k_scale_full[: chunk.total_seq_lens]
673 - ops.cp_gather_indexer_k_quant_cache(
674 - kv_cache,
675 - k_fp8,
676 - k_scale,
677 - chunk.block_table,
678 - chunk.cu_seq_lens,
679 - )
680 - fp8_mqa_logits_func = fp8_mqa_logits
681 - if current_platform.is_rocm():
682 - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
683 - rocm_fp8_mqa_logits,
684 - )
685 -
686 - fp8_mqa_logits_func = rocm_fp8_mqa_logits
687 - logits = fp8_mqa_logits_func(
688 - q_fp8[chunk.token_start : chunk.token_end],
689 - (k_fp8, k_scale.view(torch.float32).flatten()),
690 - weights[chunk.token_start : chunk.token_end],
691 - chunk.cu_seqlen_ks,
692 - chunk.cu_seqlen_ke,
693 - )
694 - num_rows = logits.shape[0]
695 - topk_indices = topk_indices_buffer[
696 - chunk.token_start : chunk.token_end, :topk_tokens
697 - ]
698 - torch.ops._C.top_k_per_row_prefill(
699 - logits,
700 - chunk.cu_seqlen_ks,
701 - chunk.cu_seqlen_ke,
702 - topk_indices,
703 - num_rows,
704 - logits.stride(0),
705 - logits.stride(1),
706 - topk_tokens,
707 - )
708 -
709 - if has_decode:
710 - decode_metadata = attn_metadata.decode
711 - # kv_cache size requirement [num_block, block_size, n_head, head_dim],
712 - # we only have [num_block, block_size, head_dim],
713 - kv_cache = kv_cache.unsqueeze(-2)
714 - decode_lens = decode_metadata.decode_lens
715 - if decode_metadata.requires_padding:
716 - # pad in edge case where we have short chunked prefill length <
717 - # decode_threshold since we unstrictly split
718 - # prefill and decode by decode_threshold
719 - # (currently set to 1 + speculative tokens)
720 -
721 - # [num_decode_tokens, n_head, head_dim] -> [bs, 1+next_n, n_head, head_dim]
722 - padded_q_fp8_decode_tokens = pack_seq_triton(
723 - q_fp8[:num_decode_tokens], decode_lens
724 - )
725 - # [num_decode_tokens, n_head] -> [bs, 1+next_n, n_head]
726 - padded_weights = pack_seq_triton(weights[:num_decode_tokens], decode_lens)
727 - # [bs, 1+next_n, n_head] -> [bs * next_n, n_head]
728 - padded_weights = padded_weights.flatten(0, 1)
729 - else:
730 - padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
731 - decode_lens.shape[0], -1, *q_fp8.shape[1:]
732 - )
733 - padded_weights = weights
734 - # TODO: move and optimize below logic with triton kernels
735 - batch_size = padded_q_fp8_decode_tokens.shape[0]
736 - next_n = padded_q_fp8_decode_tokens.shape[1]
737 - assert batch_size == decode_metadata.seq_lens.shape[0]
738 - num_padded_tokens = batch_size * next_n
739 - fp8_paged_mqa_logits_func = fp8_paged_mqa_logits
740 - if current_platform.is_rocm():
741 - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
742 - rocm_fp8_paged_mqa_logits,
743 - )
744 -
745 - fp8_paged_mqa_logits_func = rocm_fp8_paged_mqa_logits
746 - logits = fp8_paged_mqa_logits_func(
747 - padded_q_fp8_decode_tokens,
748 - kv_cache,
749 - padded_weights[:num_padded_tokens],
750 - decode_metadata.seq_lens,
751 - decode_metadata.block_table,
752 - decode_metadata.schedule_metadata,
753 - max_model_len=max_model_len,
754 - )
755 - num_rows = logits.shape[0]
756 - topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
757 -
758 - torch.ops._C.top_k_per_row_decode(
759 - logits,
760 - next_n,
761 - decode_metadata.seq_lens,
762 - topk_indices,
763 - num_rows,
764 - logits.stride(0),
765 - logits.stride(1),
766 - topk_tokens,
767 - )
768 - if decode_metadata.requires_padding:
769 - # if padded, we need to unpack
770 - # the topk indices removing padded tokens
771 - topk_indices = unpack_seq_triton(
772 - topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
773 - decode_lens,
774 - )
775 - topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (
776 - topk_indices
777 - )
778 -
779 - return topk_indices_buffer
780 -
781 -
782 -def sparse_attn_indexer_fake(
783 - hidden_states: torch.Tensor,
784 - k_cache_prefix: str,
785 - kv_cache: torch.Tensor,
786 - q_fp8: torch.Tensor,
787 - k: torch.Tensor,
788 - weights: torch.Tensor,
789 - quant_block_size: int,
790 - scale_fmt: str | None,
791 - topk_tokens: int,
792 - head_dim: int,
793 - max_model_len: int,
794 - total_seq_lens: int,
795 - topk_indices_buffer: torch.Tensor | None,
796 -) -> torch.Tensor:
797 - return topk_indices_buffer
798 -
799 -
800 -direct_register_custom_op(
801 - op_name="sparse_attn_indexer",
802 - op_func=sparse_attn_indexer,
803 - mutates_args=["topk_indices_buffer"],
804 - fake_impl=sparse_attn_indexer_fake,
805 - dispatch_key=current_platform.dispatch_key,
806 -)
807 -
808 -
592 class Indexer(nn.Module):
593 def __init__(
594 self,
@@ -870,6 +653,16 @@ class Indexer(nn.Module):
653 from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
654
655 self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)
656 + self.indexer_op = SparseAttnIndexer(
657 + self.k_cache,
658 + self.quant_block_size,
659 + self.scale_fmt,
660 + self.topk_tokens,
661 + self.head_dim,
662 + self.max_model_len,
663 + self.max_total_seq_len,
664 + self.topk_indices_buffer,
665 + )
666
667 def forward(
668 self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
@@ -892,6 +685,8 @@ class Indexer(nn.Module):
685 q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
686 k_pe = k_pe.reshape(-1, 1, self.rope_dim)
687
688 + # `rotary_emb` is shape-preserving; `q_pe` is already
689 + # [num_tokens, n_head, rope_dim].
690 q = torch.cat([q_pe, q_nope], dim=-1)
691 # `k_pe` is [num_tokens, 1, rope_dim] (MQA).
692 k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
@@ -913,21 +708,7 @@ class Indexer(nn.Module):
708 )
709 weights = weights.squeeze(-1)
710
916 - return torch.ops.vllm.sparse_attn_indexer(
917 - hidden_states,
918 - self.k_cache.prefix,
919 - self.k_cache.kv_cache[0],
920 - q_fp8,
921 - k,
922 - weights,
923 - self.quant_block_size,
924 - self.scale_fmt,
925 - self.topk_tokens,
926 - self.head_dim,
927 - self.max_model_len,
928 - self.max_total_seq_len,
929 - self.topk_indices_buffer,
930 - )
711 + return self.indexer_op(hidden_states, q_fp8, k, weights)
712
713
714 class DeepseekV2MLAAttention(nn.Module):
vllm/platforms/rocm.py
+3
@@ -480,6 +480,9 @@ class RocmPlatform(Platform):
480 ):
481 compilation_config.custom_ops.append("+grouped_topk")
482
483 + # Default dispatch to rocm's sparse_attn_indexer implementation
484 + compilation_config.custom_ops.append("+sparse_attn_indexer")
485 +
486 @classmethod
487 def verify_model_arch(cls, model_arch: str) -> None:
488 if model_arch in _ROCM_UNSUPPORTED_MODELS:
vllm/v1/attention/backends/mla/indexer.py
+6
@@ -63,6 +63,7 @@ class DeepseekV32IndexerPrefillChunkMetadata:
63 cu_seqlen_ks: torch.Tensor
64 cu_seqlen_ke: torch.Tensor
65 cu_seq_lens: torch.Tensor
66 + token_to_seq: torch.Tensor
67 total_seq_lens: int
68 token_start: int
69 token_end: int
@@ -234,6 +235,10 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
235 token_start = query_start_loc_cpu[reqs_start].item()
236 token_end = query_start_loc_cpu[reqs_end].item()
237 total_seq_lens = seq_lens_cpu[reqs_start:reqs_end].sum()
238 + seq_idx = torch.arange(0, reqs_end - reqs_start, dtype=torch.int32)
239 + token_to_seq = torch.repeat_interleave(
240 + seq_idx, seq_lens_cpu[reqs_start:reqs_end]
241 + ).to(self.device)
242 assert total_seq_lens <= self.max_prefill_buffer_size
243 cu_seq_lens = (
244 torch.cat(
@@ -249,6 +254,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
254 cu_seqlen_ks=cu_seqlen_ks,
255 cu_seqlen_ke=cu_seqlen_ke,
256 cu_seq_lens=cu_seq_lens,
257 + token_to_seq=token_to_seq,
258 total_seq_lens=total_seq_lens,
259 block_table=block_table[reqs_start:reqs_end],
260 token_start=token_start,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+110 -10
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
15 MLACommonBaseImpl,
16 get_mla_dims,
17 )
18 +from vllm.triton_utils import tl, triton
19 from vllm.v1.attention.backend import (
20 AttentionBackend,
21 AttentionCGSupport,
@@ -33,6 +34,48 @@ if TYPE_CHECKING:
34 logger = init_logger(__name__)
35
36
37 +@triton.jit
38 +def fetch_id_to_ragged_kernel(
39 + in_tensor_ptr, # [num_seq, topk]
40 + cumsum_ptr, # [num_seq + 1]
41 + out_tensor_ptr, # [max_num_seq * topk]
42 + in_tensor_ptr_stride,
43 + TOPK: tl.constexpr,
44 + TOKEN_NUM: tl.constexpr,
45 + BLOCK_SIZE: tl.constexpr,
46 +):
47 + seq_id = tl.program_id(0)
48 + block_id = tl.program_id(1)
49 + offset = tl.arange(0, BLOCK_SIZE)
50 + token_start = tl.load(cumsum_ptr + seq_id)
51 + token_end = tl.load(cumsum_ptr + seq_id + 1)
52 + token_num = token_end - token_start
53 + row_offset = block_id * BLOCK_SIZE
54 + if row_offset >= token_num:
55 + return
56 + in_tensor_offset = seq_id * in_tensor_ptr_stride + row_offset + offset
57 + in_tensor_mask = (row_offset + offset) < TOPK
58 + in_tensor_val = tl.load(in_tensor_ptr + in_tensor_offset, mask=in_tensor_mask)
59 + out_tensor_offset = token_start + row_offset + offset
60 + out_tensor_mask = (out_tensor_offset < token_end) & in_tensor_mask
61 + tl.store(out_tensor_ptr + out_tensor_offset, in_tensor_val, mask=out_tensor_mask)
62 +
63 +
64 +def fetch_id_to_ragged_triton(
65 + in_tensor: torch.Tensor, cumsum: torch.Tensor, out_tensor: torch.Tensor, topk
66 +):
67 + num_tokens = in_tensor.size(0)
68 + block_size = 64
69 + num_block_per_row = triton.cdiv(topk, block_size)
70 + grid = (
71 + num_tokens,
72 + num_block_per_row,
73 + )
74 + fetch_id_to_ragged_kernel[grid](
75 + in_tensor, cumsum, out_tensor, in_tensor.stride(0), topk, num_tokens, block_size
76 + )
77 +
78 +
79 class ROCMAiterMLASparseBackend(AttentionBackend):
80 accept_output_buffer: bool = True
81
@@ -83,6 +126,13 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
126
127 block_table: torch.Tensor
128 req_id_per_token: torch.Tensor
129 +
130 + qo_indptr: torch.Tensor
131 + paged_kv_last_page_len: torch.Tensor
132 + paged_kv_indices: torch.Tensor
133 + paged_kv_indptr: torch.Tensor
134 + paged_kv_indptr_rest: torch.Tensor
135 +
136 block_size: int = 1
137 topk_tokens: int = 2048
138
@@ -91,7 +141,7 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
141 class ROCMAiterMLASparseMetadataBuilder(
142 AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]
143 ):
94 - cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
144 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
145
146 def __init__(
147 self,
@@ -104,6 +154,7 @@ class ROCMAiterMLASparseMetadataBuilder(
154 self.model_config = vllm_config.model_config
155 parallel_config = vllm_config.parallel_config
156 self.device = device
157 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
158
159 self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
160 self.mla_dims = get_mla_dims(self.model_config)
@@ -124,6 +175,23 @@ class ROCMAiterMLASparseMetadataBuilder(
175 dtype=torch.int32,
176 device=device,
177 )
178 + self.qo_indptr = torch.arange(
179 + 0, max_num_batched_tokens + 1, dtype=torch.int32, device=device
180 + )
181 + self.paged_kv_last_page_len = torch.ones(
182 + max_num_batched_tokens, dtype=torch.int32, device=device
183 + )
184 +
185 + # These two needs to be calculated in runtime,
186 + # but we still needs to prepare the buffer
187 + self.paged_kv_indices = torch.zeros(
188 + [max_num_batched_tokens * self.topk_tokens],
189 + dtype=torch.int32,
190 + device=device,
191 + )
192 + self.paged_kv_indptr = torch.zeros(
193 + [max_num_batched_tokens + 1], dtype=torch.int32, device=device
194 + )
195
196 def build(
197 self,
@@ -142,7 +210,15 @@ class ROCMAiterMLASparseMetadataBuilder(
210 self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
211 torch.from_numpy(req_id_per_token), non_blocking=True
212 )
213 + self.paged_kv_indices.fill_(0)
214 + self.paged_kv_indptr.fill_(0)
215 +
216 req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
217 + qo_indptr = self.qo_indptr[: num_tokens + 1]
218 + paged_kv_last_page_len = self.paged_kv_last_page_len[:num_tokens]
219 + paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens]
220 + paged_kv_indptr = self.paged_kv_indptr[: num_tokens + 1]
221 + paged_kv_indptr_rest = self.paged_kv_indptr[num_tokens + 1 :]
222
223 metadata = ROCMAiterMLASparseMetadata(
224 num_reqs=common_attn_metadata.num_reqs,
@@ -155,6 +231,11 @@ class ROCMAiterMLASparseMetadataBuilder(
231 req_id_per_token=req_id_per_token,
232 block_size=self.kv_cache_spec.block_size,
233 topk_tokens=self.topk_tokens,
234 + qo_indptr=qo_indptr,
235 + paged_kv_last_page_len=paged_kv_last_page_len,
236 + paged_kv_indices=paged_kv_indices,
237 + paged_kv_indptr=paged_kv_indptr,
238 + paged_kv_indptr_rest=paged_kv_indptr_rest,
239 )
240 return metadata
241
@@ -226,20 +307,39 @@ class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):
307
308 def _forward_bf16_kv(
309 self,
229 - q: torch.Tensor,
230 - kv_c_and_k_pe_cache: torch.Tensor,
231 - topk_indices: torch.Tensor,
310 + q: torch.Tensor, # [sq, heads, d_qk]
311 + kv_c_and_k_pe_cache: torch.Tensor, # [blocks, heads, d_qk]
312 + topk_indices: torch.Tensor, # [sq, topk]
313 attn_metadata: ROCMAiterMLASparseMetadata,
314 ) -> torch.Tensor:
315 num_tokens = q.shape[0]
235 - kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
236 - -1, 1, kv_c_and_k_pe_cache.shape[-1]
316 + output = torch.empty(
317 + [num_tokens, self.num_heads, self.kv_lora_rank],
318 + dtype=q.dtype,
319 + device=q.device,
320 + )
321 + seq_len = (topk_indices != -1).sum(dim=-1)
322 + torch.cumsum(seq_len, dim=0, out=attn_metadata.paged_kv_indptr[1:])
323 + attn_metadata.paged_kv_indptr_rest.fill_(attn_metadata.paged_kv_indptr[-1])
324 + fetch_id_to_ragged_triton(
325 + topk_indices,
326 + attn_metadata.paged_kv_indptr,
327 + attn_metadata.paged_kv_indices,
328 + attn_metadata.topk_tokens,
329 + )
330 +
331 + rocm_aiter_ops.mla_decode_fwd(
332 + q,
333 + kv_c_and_k_pe_cache,
334 + output,
335 + self.scale,
336 + attn_metadata.qo_indptr,
337 + 1,
338 + attn_metadata.paged_kv_indptr,
339 + attn_metadata.paged_kv_indices,
340 + attn_metadata.paged_kv_last_page_len,
341 )
342
239 - topk_indices = topk_indices.view(num_tokens, 1, -1)
240 - output = reference_mla_sparse_prefill(
241 - q, kv_c_and_k_pe_cache, topk_indices, self.softmax_scale, 512
242 - )[0]
343 return output[:, : self.num_heads, :]
344
345 def forward(
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+518 -80
@@ -1,100 +1,220 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 import importlib
4 -from functools import lru_cache
5
6 import torch
7
8 -from vllm._aiter_ops import rocm_aiter_ops
9 -from vllm.logger import init_logger
8 +from vllm.forward_context import get_forward_context
9 from vllm.platforms import current_platform
10 +from vllm.triton_utils import tl, triton
11 +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
12 +from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
13
12 -logger = init_logger(__name__)
14 +if current_platform.is_cuda_alike():
15 + from vllm import _custom_ops as ops
16
17
15 -# Take from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L84
16 -def fp8_mqa_logits_torch(
17 - q: torch.Tensor,
18 - kv: tuple[torch.Tensor, torch.Tensor],
19 - weights: torch.Tensor,
20 - cu_seqlen_ks: torch.Tensor,
21 - cu_seqlen_ke: torch.Tensor,
22 -) -> torch.Tensor:
23 - """Compute FP8 MQA logits for a single sequence without KV paging.
24 -
25 - Args:
26 - q: Query tensor of shape [M, H, D]. Casted to
27 - `torch.float8_e4m3fn` by caller.
28 - kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
29 - dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
30 - [N, 1]) with dtype `torch.float32`.
31 - weights: weights of shape [M, H], dtype `torch.float32`.
32 - cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
33 - shape [M], dtype int32.
34 - cu_seqlen_ke: End indices (exclusive) for valid K per query position,
35 - shape [M], dtype int32.
36 -
37 - Returns:
38 - Logits tensor of shape [M, N], dtype `torch.float32`.
39 - """
40 - k_fp8, scale = kv
41 - seq_len_kv = k_fp8.shape[0]
42 - k = k_fp8.to(torch.bfloat16)
43 - q = q.to(torch.bfloat16)
44 -
45 - mask_lo = (
46 - torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None]
47 - )
48 - mask_hi = (
49 - torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
50 - )
51 - mask = mask_lo & mask_hi
18 +@triton.jit
19 +def _indexer_k_quant_and_cache_kernel(
20 + k_ptr, # [num_tokens, head_dim]
21 + kv_cache_ptr, # [n_blks, blk_size//tile_block, head_dim // 16B, tile_block, 16B]
22 + # [n_blocks, blk_size, head_dim]
23 + kv_cache_scale_ptr, # [n_blks, blk_size]
24 + slot_mapping_ptr, # [num_tokens]
25 + kv_cache_scale_stride,
26 + kv_cache_value_stride,
27 + block_size,
28 + num_tokens,
29 + head_dim: tl.constexpr,
30 + LAYOUT: tl.constexpr,
31 + BLOCK_TILE_SIZE: tl.constexpr,
32 + HEAD_TILE_SIZE: tl.constexpr,
33 + IS_FNUZ: tl.constexpr,
34 + USE_UE8M0: tl.constexpr,
35 +):
36 + tid = tl.program_id(0)
37 + offset = tl.arange(0, head_dim)
38 + if LAYOUT == "SHUFFLE":
39 + tile_offset = (
40 + offset // HEAD_TILE_SIZE * BLOCK_TILE_SIZE * HEAD_TILE_SIZE
41 + + offset % HEAD_TILE_SIZE
42 + )
43 + else:
44 + tile_offset = offset
45 + tile_store_offset = tile_offset
46 + # for idx in tl.range(tid, num_tokens, n_program):
47 + src_ptr = k_ptr + tid * head_dim
48 + slot_id = tl.load(slot_mapping_ptr + tid)
49 + if slot_id < 0:
50 + return
51 + block_id = slot_id // block_size
52 + block_offset = slot_id % block_size
53 + tile_block_id = block_offset // BLOCK_TILE_SIZE
54 + tile_block_offset = block_offset % BLOCK_TILE_SIZE
55 + val = tl.load(src_ptr + offset)
56 + amax = tl.max(val.abs(), axis=-1).to(tl.float32)
57 + if IS_FNUZ:
58 + scale = tl.maximum(1e-4, amax) / 224.0
59 + else:
60 + scale = tl.maximum(1e-4, amax) / 448.0
61
53 - score = torch.einsum("mhd,nd->hmn", q, k).float() * scale
54 - logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
55 - logits = logits.masked_fill(~mask, float("-inf"))
62 + if USE_UE8M0:
63 + scale = tl.exp2(tl.ceil(tl.log2(scale)))
64
57 - return logits
65 + fp8_val = (val.to(tl.float32) / scale).to(kv_cache_ptr.type.element_ty)
66 + if LAYOUT == "SHUFFLE":
67 + dst_ptr = (
68 + kv_cache_ptr
69 + + block_id * kv_cache_value_stride
70 + + tile_block_id * BLOCK_TILE_SIZE * head_dim
71 + + tile_block_offset * HEAD_TILE_SIZE
72 + )
73 + else:
74 + dst_ptr = (
75 + kv_cache_ptr + block_id * kv_cache_value_stride + block_offset * head_dim
76 + )
77 + tl.store(dst_ptr + tile_store_offset, fp8_val)
78 + dst_scale_ptr = kv_cache_scale_ptr + block_id * kv_cache_scale_stride + block_offset
79 + tl.store(dst_scale_ptr, scale)
80
81
60 -def rocm_fp8_mqa_logits(
61 - q: torch.Tensor,
62 - kv: tuple[torch.Tensor, torch.Tensor],
63 - weights: torch.Tensor,
64 - cu_seqlen_ks: torch.Tensor,
65 - cu_seqlen_ke: torch.Tensor,
66 -) -> torch.Tensor:
67 - """Compute FP8 MQA logits for a single sequence without KV paging.
68 -
69 - Args:
70 - q: Query tensor of shape [M, H, D]. Casted to
71 - `torch.float8_e4m3fn` by caller.
72 - kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
73 - dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
74 - [N, 1]) with dtype `torch.float32`.
75 - weights: weights of shape [M, H], dtype `torch.float32`.
76 - cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
77 - shape [M], dtype int32.
78 - cu_seqlen_ke: End indices (exclusive) for valid K per query position,
79 - shape [M], dtype int32.
82 +def indexer_k_quant_and_cache_triton(
83 + k: torch.Tensor,
84 + kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4]
85 + slot_mapping: torch.Tensor,
86 + quant_block_size,
87 + scale_fmt,
88 + block_tile_size=16,
89 + head_tile_size=16,
90 +):
91 + num_blocks = kv_cache.shape[0]
92 + head_dim = k.shape[-1]
93 + num_tokens = slot_mapping.shape[0]
94 + block_size = kv_cache.shape[1]
95 + # In real layout, we store the first portion as kv cache value
96 + # and second portion as kv cache scale
97 + kv_cache = kv_cache.view(num_blocks, -1)
98 + kv_cache_value = kv_cache[:, : block_size * head_dim]
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,)
102 + _indexer_k_quant_and_cache_kernel[grid](
103 + k,
104 + kv_cache_value,
105 + kv_cache_scale,
106 + slot_mapping,
107 + kv_cache_scale.stride(0),
108 + kv_cache_value.stride(0),
109 + block_size,
110 + num_tokens,
111 + head_dim,
112 + "NHD",
113 + block_tile_size,
114 + head_tile_size,
115 + IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
116 + USE_UE8M0=scale_fmt == "ue8m0",
117 + )
118
81 - Returns:
82 - Logits tensor of shape [M, N], dtype `torch.float32`.
83 - """
119
85 - # TODO(ganyi): Temporarily workaround, will remove the module check and reference
86 - # path after aiter merge this kernel into main
87 - @lru_cache
88 - def has_mqa_logits_module():
89 - return importlib.util.find_spec("aiter.ops.triton.fp8_mqa_logits") is not None
120 +@triton.jit
121 +def _cp_gather_indexer_quant_cache_kernel(
122 + kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B]
123 + # [n_blks, blk_size, head_dim]
124 + kv_cache_scale_ptr, # [n_blks, blk_size]
125 + k_fp8_ptr, # [num_tokens, head_dim]
126 + k_scale_ptr, # [num_tokens]
127 + block_table_ptr, # [batch_size, block_table_stride]
128 + cu_seqlen_ptr, # [batch_size + 1]
129 + token_to_seq_ptr, # [num_tokens]
130 + block_size,
131 + block_table_stride,
132 + kv_cache_stride,
133 + kv_cache_scale_stride,
134 + LAYOUT: tl.constexpr,
135 + HEAD_DIM: tl.constexpr,
136 + BLOCK_TILE_SIZE: tl.constexpr,
137 + HEAD_TILE_SIZE: tl.constexpr,
138 +):
139 + tid = tl.program_id(0)
140 + offset = tl.arange(0, HEAD_DIM)
141 + batch_id = tl.load(token_to_seq_ptr + tid)
142 + batch_start = tl.load(cu_seqlen_ptr + batch_id)
143 + batch_end = tl.load(cu_seqlen_ptr + batch_id + 1)
144 + batch_offset = tid - batch_start
145 + if tid >= batch_end:
146 + return
147 + block_table_id = batch_offset // block_size
148 + block_offset = batch_offset % block_size
149 + block_table_offset = batch_id * block_table_stride + block_table_id
150 + block_id = tl.load(block_table_ptr + block_table_offset)
151 + tiled_block_id = block_offset // BLOCK_TILE_SIZE
152 + tiled_block_offset = block_offset % BLOCK_TILE_SIZE
153 + if LAYOUT == "SHUFFLE":
154 + src_cache_offset = (
155 + block_id * kv_cache_stride
156 + + tiled_block_id * HEAD_DIM * BLOCK_TILE_SIZE
157 + + tiled_block_offset * HEAD_TILE_SIZE
158 + )
159 + else:
160 + src_cache_offset = block_id * kv_cache_stride + block_offset * HEAD_DIM
161 + src_scale_offset = block_id * kv_cache_scale_stride + block_offset
162 + dst_offset = tid * HEAD_DIM
163 + src_scale_ptr = kv_cache_scale_ptr + src_scale_offset
164 + src_cache_ptr = kv_cache_ptr + src_cache_offset
165 + dst_k_ptr = k_fp8_ptr + dst_offset
166 + scale_val = tl.load(src_scale_ptr)
167 + tl.store(k_scale_ptr + tid, scale_val)
168 + if LAYOUT == "SHUFFLE":
169 + tiled_src_offset = (
170 + offset // HEAD_TILE_SIZE * HEAD_TILE_SIZE * BLOCK_TILE_SIZE
171 + + offset % HEAD_TILE_SIZE
172 + )
173 + else:
174 + tiled_src_offset = offset
175 + val = tl.load(src_cache_ptr + tiled_src_offset)
176 + tl.store(dst_k_ptr + offset, val)
177
91 - if rocm_aiter_ops.is_enabled() and has_mqa_logits_module():
92 - from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
178
94 - kv, scale = kv
95 - return fp8_mqa_logits(q, kv, scale, weights, cu_seqlen_ks, cu_seqlen_ke)
96 - else:
97 - return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
179 +def cp_gather_indexer_k_quant_cache_triton(
180 + k_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4]
181 + k_fp8: torch.Tensor,
182 + k_fp8_scale: torch.Tensor,
183 + block_table: torch.Tensor,
184 + cu_seqlen: torch.Tensor,
185 + token_to_seq: torch.Tensor,
186 + block_tile_size: int = 16,
187 + head_tile_size: int = 16,
188 +):
189 + num_tokens = k_fp8.size(0)
190 + block_size = k_cache.size(1)
191 + block_table_stride = block_table.stride(0)
192 + head_dim = k_fp8.shape[-1]
193 + num_blocks = k_cache.shape[0]
194 + # we assume the kv cache already been split to 2 portion
195 + k_cache = k_cache.view(num_blocks, -1)
196 + fp8_dtype = current_platform.fp8_dtype()
197 + k_cache_value = k_cache[:, : block_size * head_dim].view(fp8_dtype)
198 + k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)
199 + grid = (num_tokens,)
200 + k_fp8_scale = k_fp8_scale.view(torch.float32)
201 + _cp_gather_indexer_quant_cache_kernel[grid](
202 + k_cache_value,
203 + k_cache_scale,
204 + k_fp8,
205 + k_fp8_scale,
206 + block_table,
207 + cu_seqlen,
208 + token_to_seq,
209 + block_size,
210 + block_table_stride,
211 + k_cache_value.stride(0),
212 + k_cache_scale.stride(0),
213 + "NHD",
214 + head_dim,
215 + block_tile_size,
216 + head_tile_size,
217 + )
218
219
220 # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156
@@ -183,10 +303,38 @@ def rocm_fp8_paged_mqa_logits(
303 Logits tensor of shape [B * next_n, max_model_len], dtype
304 `torch.float32`.
305 """
306 + from vllm._aiter_ops import rocm_aiter_ops
307 +
308 + @functools.lru_cache
309 + def paged_mqa_logits_module():
310 + paged_mqa_logits_module_path = None
311 + if importlib.util.find_spec("aiter.ops.triton.pa_mqa_logits") is not None:
312 + paged_mqa_logits_module_path = "aiter.ops.triton.pa_mqa_logits"
313 + elif (
314 + importlib.util.find_spec("aiter.ops.triton.attention.pa_mqa_logits")
315 + is not None
316 + ):
317 + paged_mqa_logits_module_path = "aiter.ops.triton.attention.pa_mqa_logits"
318 +
319 + if paged_mqa_logits_module_path is not None:
320 + try:
321 + module = importlib.import_module(paged_mqa_logits_module_path)
322 + return module
323 + except ImportError:
324 + return None
325 + return None
326
327 + aiter_paged_mqa_logits_module = None
328 if rocm_aiter_ops.is_enabled():
188 - from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits_stage1
329 + aiter_paged_mqa_logits_module = paged_mqa_logits_module()
330 + # FIXME(ganyi): Temporarily disable the aiter path until nightly docker
331 + # update aiter to the fix PR.
332 + aiter_paged_mqa_logits_module = None
333
334 + if aiter_paged_mqa_logits_module is not None:
335 + deepgemm_fp8_paged_mqa_logits_stage1 = (
336 + aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1
337 + )
338 batch_size, next_n, heads, _ = q_fp8.shape
339 out_qk = torch.full(
340 (heads, batch_size * next_n, max_model_len),
@@ -208,3 +356,293 @@ def rocm_fp8_paged_mqa_logits(
356 return fp8_paged_mqa_logits_torch(
357 q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
358 )
359 +
360 +
361 +# Take from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L84
362 +def fp8_mqa_logits_torch(
363 + q: torch.Tensor,
364 + kv: tuple[torch.Tensor, torch.Tensor],
365 + weights: torch.Tensor,
366 + cu_seqlen_ks: torch.Tensor,
367 + cu_seqlen_ke: torch.Tensor,
368 +) -> torch.Tensor:
369 + """Compute FP8 MQA logits for a single sequence without KV paging.
370 +
371 + Args:
372 + q: Query tensor of shape [M, H, D]. Casted to
373 + `torch.float8_e4m3fn` by caller.
374 + kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
375 + dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
376 + [N, 1]) with dtype `torch.float32`.
377 + weights: weights of shape [M, H], dtype `torch.float32`.
378 + cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
379 + shape [M], dtype int32.
380 + cu_seqlen_ke: End indices (exclusive) for valid K per query position,
381 + shape [M], dtype int32.
382 +
383 + Returns:
384 + Logits tensor of shape [M, N], dtype `torch.float32`.
385 + """
386 + kv, scale = kv
387 + seq_len_kv = kv.shape[0]
388 + k = kv.to(torch.bfloat16)
389 + q = q.to(torch.bfloat16)
390 +
391 + mask_lo = (
392 + torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None]
393 + )
394 + mask_hi = (
395 + torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
396 + )
397 + mask = mask_lo & mask_hi
398 +
399 + score = torch.einsum("mhd,nd->hmn", q, k).float() * scale
400 + logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
401 + logits = logits.masked_fill(~mask, float("-inf"))
402 +
403 + return logits
404 +
405 +
406 +def rocm_fp8_mqa_logits(
407 + q: torch.Tensor,
408 + kv: tuple[torch.Tensor, torch.Tensor],
409 + weights: torch.Tensor,
410 + cu_seqlen_ks: torch.Tensor,
411 + cu_seqlen_ke: torch.Tensor,
412 +) -> torch.Tensor:
413 + """Compute FP8 MQA logits for a single sequence without KV paging.
414 +
415 + Args:
416 + q: Query tensor of shape [M, H, D]. Casted to
417 + `torch.float8_e4m3fn` by caller.
418 + kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
419 + dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
420 + [N, 1]) with dtype `torch.float32`.
421 + weights: weights of shape [M, H], dtype `torch.float32`.
422 + cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
423 + shape [M], dtype int32.
424 + cu_seqlen_ke: End indices (exclusive) for valid K per query position,
425 + shape [M], dtype int32.
426 +
427 + Returns:
428 + Logits tensor of shape [M, N], dtype `torch.float32`.
429 + """
430 +
431 + # TODO(ganyi): Temporarily workaround, will remove the module check and reference
432 + # path after aiter merge this kernel into main
433 + from vllm._aiter_ops import rocm_aiter_ops
434 +
435 + @functools.lru_cache
436 + def mqa_logits_module():
437 + mqa_logits_module_path = None
438 + if importlib.util.find_spec("aiter.ops.triton.fp8_mqa_logits") is not None:
439 + mqa_logits_module_path = "aiter.ops.triton.fp8_mqa_logits"
440 + elif (
441 + importlib.util.find_spec("aiter.ops.triton.attention.fp8_mqa_logits")
442 + is not None
443 + ):
444 + mqa_logits_module_path = "aiter.ops.triton.attention.fp8_mqa_logits"
445 +
446 + if mqa_logits_module_path is not None:
447 + try:
448 + module = importlib.import_module(mqa_logits_module_path)
449 + return module
450 + except ImportError:
451 + return None
452 + return None
453 +
454 + aiter_mqa_logits_module = None
455 + if rocm_aiter_ops.is_enabled():
456 + aiter_mqa_logits_module = mqa_logits_module()
457 +
458 + if aiter_mqa_logits_module is not None:
459 + fp8_mqa_logits = aiter_mqa_logits_module.fp8_mqa_logits
460 + kv, scale = kv
461 + return fp8_mqa_logits(q, kv, scale, weights, cu_seqlen_ks, cu_seqlen_ke)
462 + else:
463 + return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
464 +
465 +
466 +def rocm_aiter_sparse_attn_indexer_fake(
467 + hidden_states: torch.Tensor,
468 + k_cache_prefix: str,
469 + kv_cache: torch.Tensor,
470 + q_fp8: torch.Tensor,
471 + k: torch.Tensor,
472 + weights: torch.Tensor,
473 + quant_block_size: int,
474 + scale_fmt: str | None,
475 + topk_tokens: int,
476 + head_dim: int,
477 + max_model_len: int,
478 + total_seq_lens: int,
479 + topk_indices_buffer: torch.Tensor | None,
480 +) -> torch.Tensor:
481 + # profile run
482 + # NOTE(Chen): create the max possible flattened_kv. So that
483 + # profile_run can get correct memory usage.
484 + _flattened_kv = torch.empty(
485 + [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
486 + )
487 + fp8_dtype = current_platform.fp8_dtype()
488 + _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()
489 + _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous()
490 + return topk_indices_buffer
491 +
492 +
493 +def rocm_aiter_sparse_attn_indexer(
494 + hidden_states: torch.Tensor,
495 + k_cache_prefix: str,
496 + kv_cache: torch.Tensor,
497 + q_fp8: torch.Tensor,
498 + k: torch.Tensor,
499 + weights: torch.Tensor,
500 + quant_block_size: int,
501 + scale_fmt: str | None,
502 + topk_tokens: int,
503 + head_dim: int,
504 + max_model_len: int,
505 + total_seq_lens: int,
506 + topk_indices_buffer: torch.Tensor | None,
507 +) -> torch.Tensor:
508 + # careful! this will be None in dummy run
509 + attn_metadata = get_forward_context().attn_metadata
510 + fp8_dtype = current_platform.fp8_dtype()
511 + # assert isinstance(attn_metadata, dict)
512 + if not isinstance(attn_metadata, dict):
513 + return rocm_aiter_sparse_attn_indexer_fake(
514 + hidden_states,
515 + k_cache_prefix,
516 + kv_cache,
517 + q_fp8,
518 + k,
519 + weights,
520 + quant_block_size,
521 + scale_fmt,
522 + topk_tokens,
523 + head_dim,
524 + max_model_len,
525 + total_seq_lens,
526 + topk_indices_buffer,
527 + )
528 + attn_metadata = attn_metadata[k_cache_prefix]
529 + assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
530 + slot_mapping = attn_metadata.slot_mapping
531 + has_decode = attn_metadata.num_decodes > 0
532 + has_prefill = attn_metadata.num_prefills > 0
533 + num_decode_tokens = attn_metadata.num_decode_tokens
534 +
535 + ops.indexer_k_quant_and_cache(
536 + k,
537 + kv_cache,
538 + slot_mapping,
539 + quant_block_size,
540 + scale_fmt,
541 + )
542 +
543 + topk_indices_buffer[: hidden_states.shape[0]] = -1
544 + if has_prefill:
545 + prefill_metadata = attn_metadata.prefill
546 + for chunk in prefill_metadata.chunks:
547 + k_fp8 = torch.empty(
548 + [chunk.total_seq_lens, head_dim],
549 + device=k.device,
550 + dtype=fp8_dtype,
551 + )
552 + k_scale = torch.empty(
553 + [chunk.total_seq_lens, 4],
554 + device=k.device,
555 + dtype=torch.uint8,
556 + )
557 +
558 + ops.cp_gather_indexer_k_quant_cache(
559 + kv_cache,
560 + k_fp8,
561 + k_scale,
562 + chunk.block_table,
563 + chunk.cu_seq_lens,
564 + )
565 +
566 + logits = rocm_fp8_mqa_logits(
567 + q_fp8[chunk.token_start : chunk.token_end],
568 + (k_fp8, k_scale.view(torch.float32)),
569 + weights[chunk.token_start : chunk.token_end],
570 + chunk.cu_seqlen_ks,
571 + chunk.cu_seqlen_ke,
572 + )
573 + num_rows = logits.shape[0]
574 + assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
575 + topk_indices = topk_indices_buffer[
576 + chunk.token_start : chunk.token_end, :topk_tokens
577 + ]
578 + torch.ops._C.top_k_per_row_prefill(
579 + logits,
580 + chunk.cu_seqlen_ks,
581 + chunk.cu_seqlen_ke,
582 + topk_indices,
583 + num_rows,
584 + logits.stride(0),
585 + logits.stride(1),
586 + topk_tokens,
587 + )
588 +
589 + if has_decode:
590 + decode_metadata = attn_metadata.decode
591 + # kv_cache size requirement [num_block, block_size, n_head, head_dim],
592 + # we only have [num_block, block_size, head_dim],
593 + kv_cache = kv_cache.unsqueeze(-2)
594 + decode_lens = decode_metadata.decode_lens
595 + if decode_metadata.requires_padding:
596 + # pad in edge case where we have short chunked prefill length <
597 + # decode_threshold since we unstrictly split
598 + # prefill and decode by decode_threshold
599 + # (currently set to 1 + speculative tokens)
600 + padded_q_fp8_decode_tokens = pack_seq_triton(
601 + q_fp8[:num_decode_tokens], decode_lens
602 + )
603 + else:
604 + padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
605 + decode_lens.shape[0], -1, *q_fp8.shape[1:]
606 + )
607 + # TODO: move and optimize below logic with triton kernels
608 + batch_size = padded_q_fp8_decode_tokens.shape[0]
609 + next_n = padded_q_fp8_decode_tokens.shape[1]
610 + assert batch_size == decode_metadata.seq_lens.shape[0]
611 + num_padded_tokens = batch_size * next_n
612 +
613 + logits = rocm_fp8_paged_mqa_logits(
614 + padded_q_fp8_decode_tokens,
615 + kv_cache,
616 + weights[:num_padded_tokens],
617 + decode_metadata.seq_lens,
618 + decode_metadata.block_table,
619 + decode_metadata.schedule_metadata,
620 + max_model_len=max_model_len,
621 + )
622 +
623 + num_rows = logits.shape[0]
624 + assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
625 + topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]
626 + torch.ops._C.top_k_per_row_decode(
627 + logits,
628 + next_n,
629 + decode_metadata.seq_lens,
630 + topk_indices,
631 + num_rows,
632 + logits.stride(0),
633 + logits.stride(1),
634 + topk_tokens,
635 + )
636 +
637 + if decode_metadata.requires_padding:
638 + # if padded, we need to unpack
639 + # the topk indices removing padded tokens
640 + topk_indices = unpack_seq_triton(
641 + topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
642 + decode_lens,
643 + )
644 + topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (
645 + topk_indices
646 + )
647 +
648 + return topk_indices_buffer