@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 243 lines 8.69 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """SM120 implementation variant for ``FLASHINFER_MLA_SPARSE_SM120``."""
4
5 from typing import TYPE_CHECKING, cast
6
7 import torch
8
9 from vllm.model_executor.layers.attention.sparse_mla_attention import (
10 SparseMLACommonImpl,
11 )
12 from vllm.v1.attention.backend import AttentionLayer, AttentionType
13 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
14 FlashInferMLASparseMetadata,
15 _get_workspace_buffer,
16 )
17 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
18 from vllm.v1.attention.backends.mla.sparse_utils import (
19 triton_convert_req_index_to_global_index,
20 )
21
22 if TYPE_CHECKING:
23 from vllm.model_executor.models.deepseek_v2 import Indexer
24
25
26 def _kv_scale_format_for_model(model_type: str | None) -> str:
27 if model_type is not None and model_type.startswith("glm"):
28 return "arbitrary_fp32"
29 return "pow2_fp32"
30
31
32 class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
33 """SM120 FlashInfer sparse-MLA implementation."""
34
35 is_sparse = True
36 supports_dense_mha_prefill = False
37
38 def __init__(
39 self,
40 num_heads: int,
41 head_size: int,
42 scale: float,
43 num_kv_heads: int,
44 alibi_slopes: list[float] | None,
45 sliding_window: int | None,
46 kv_cache_dtype: str,
47 logits_soft_cap: float | None,
48 attn_type: str,
49 kv_sharing_target_layer_name: str | None,
50 indexer: "Indexer | None" = None,
51 **mla_args,
52 ) -> None:
53 if any([alibi_slopes, sliding_window, logits_soft_cap]):
54 raise NotImplementedError(
55 "FLASHINFER_MLA_SPARSE_SM120 does not support alibi_slopes / "
56 "sliding_window / logits_soft_cap"
57 )
58 if attn_type != AttentionType.DECODER:
59 raise NotImplementedError(
60 "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention"
61 )
62
63 if kv_cache_dtype != "fp8_ds_mla":
64 raise NotImplementedError(
65 "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla "
66 f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}."
67 )
68
69 topk_indices_buffer = mla_args.pop("topk_indices_buffer", None)
70 super().__init__(
71 num_heads,
72 head_size,
73 scale,
74 num_kv_heads,
75 alibi_slopes,
76 sliding_window,
77 kv_cache_dtype,
78 logits_soft_cap,
79 attn_type,
80 kv_sharing_target_layer_name,
81 indexer=indexer,
82 topk_indices_buffer=topk_indices_buffer,
83 **mla_args,
84 )
85 from vllm.config import get_current_vllm_config
86
87 vllm_config = get_current_vllm_config()
88 model_type = None
89 if vllm_config.model_config is not None:
90 model_type = getattr(
91 vllm_config.model_config.hf_text_config, "model_type", None
92 )
93 self.kv_scale_format = _kv_scale_format_for_model(model_type)
94
95 from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
96
97 if not has_flashinfer_sparse_mla_sm120():
98 raise RuntimeError(
99 "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's "
100 "sparse MLA decode API."
101 )
102 assert self.topk_indices_buffer is not None
103
104 self.supports_quant_query_input = False
105 self._workspace_buffer: torch.Tensor | None = None
106
107 def forward_mqa(
108 self,
109 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
110 kv_c_and_k_pe_cache: torch.Tensor,
111 attn_metadata: FlashInferMLASparseMetadata,
112 layer: AttentionLayer,
113 ) -> tuple[torch.Tensor, torch.Tensor | None]:
114 if isinstance(q, tuple):
115 q = torch.cat(q, dim=-1)
116
117 num_actual_toks = q.shape[0]
118
119 assert self.topk_indices_buffer is not None
120 topk_indices = self.topk_indices_buffer[:num_actual_toks]
121 index_group = self.index_group
122 if isinstance(index_group, HiSparseMLAIndexGroup):
123 num_decode_tokens = attn_metadata.num_decode_tokens
124 outputs = []
125 if num_decode_tokens:
126 topk_indices_physical = cast(
127 torch.Tensor,
128 index_group.convert_logical_to_physical_topk(
129 self.index_group_index,
130 topk_indices[:num_decode_tokens],
131 attn_metadata,
132 block_stride_rows=None,
133 return_valid_counts=False,
134 ),
135 )
136 outputs.append(
137 self._run_mqa_kernel(
138 q[:num_decode_tokens],
139 index_group.physical_kv_cache(self.index_group_index),
140 topk_indices_physical,
141 )
142 )
143 if num_decode_tokens < num_actual_toks:
144 cache = index_group.cache(self.index_group_index)
145 if num_decode_tokens == 0 and cache.all_context_pages_resident:
146 topk_indices_physical = cast(
147 torch.Tensor,
148 index_group.convert_logical_to_physical_topk(
149 self.index_group_index,
150 topk_indices,
151 attn_metadata,
152 block_stride_rows=None,
153 return_valid_counts=False,
154 ),
155 )
156 prefill_cache = index_group.physical_kv_cache(
157 self.index_group_index
158 )
159 else:
160 prefill_cache, block_table, req_ids = (
161 index_group.stage_prefill_rows(
162 self.index_group_index,
163 kv_c_and_k_pe_cache,
164 attn_metadata,
165 )
166 )
167 topk_indices_physical = cast(
168 torch.Tensor,
169 triton_convert_req_index_to_global_index(
170 req_ids,
171 block_table,
172 topk_indices[num_decode_tokens:],
173 BLOCK_SIZE=attn_metadata.block_size,
174 NUM_TOPK_TOKENS=topk_indices.shape[1],
175 ),
176 )
177 outputs.append(
178 self._run_mqa_kernel(
179 q[num_decode_tokens:],
180 prefill_cache,
181 topk_indices_physical,
182 )
183 )
184 output = torch.cat(outputs) if len(outputs) > 1 else outputs[0]
185 return output, None
186
187 topk_indices_physical = cast(
188 torch.Tensor,
189 triton_convert_req_index_to_global_index(
190 attn_metadata.req_id_per_token[:num_actual_toks],
191 attn_metadata.block_table,
192 topk_indices,
193 BLOCK_SIZE=attn_metadata.block_size,
194 NUM_TOPK_TOKENS=topk_indices.shape[1],
195 ),
196 )
197 return (
198 self._run_mqa_kernel(
199 q,
200 kv_c_and_k_pe_cache,
201 topk_indices_physical,
202 ),
203 None,
204 )
205
206 def _run_mqa_kernel(
207 self,
208 q: torch.Tensor,
209 kv_cache: torch.Tensor,
210 topk_indices_physical: torch.Tensor,
211 ) -> torch.Tensor:
212 num_actual_toks = q.shape[0]
213
214 output = q.new_empty(
215 (num_actual_toks, self.num_heads, self.kv_lora_rank),
216 dtype=q.dtype,
217 )
218
219 if self._workspace_buffer is None:
220 self._workspace_buffer = _get_workspace_buffer(q.device)
221
222 from vllm.utils.flashinfer import (
223 flashinfer_trtllm_batch_decode_with_kv_cache_mla,
224 )
225
226 sparse_capacity = topk_indices_physical.shape[1]
227 out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
228 query=q.unsqueeze(1),
229 kv_cache=kv_cache.view(torch.uint8).unsqueeze(1),
230 workspace_buffer=self._workspace_buffer,
231 qk_nope_head_dim=self.qk_nope_head_dim,
232 kv_lora_rank=self.kv_lora_rank,
233 qk_rope_head_dim=self.qk_rope_head_dim,
234 block_tables=topk_indices_physical.unsqueeze(1),
235 seq_lens=None,
236 max_seq_len=sparse_capacity,
237 out=output.unsqueeze(1),
238 bmm1_scale=self.scale,
239 bmm2_scale=1.0,
240 sparse_mla_top_k=sparse_capacity,
241 kv_scale_format=self.kv_scale_format,
242 )
243 return out.squeeze(1)