@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 289 lines 9.49 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Warmup and autotune helpers for FlashInfer sparse MLA backends."""
4
5 from typing import TYPE_CHECKING, cast
6
7 import torch
8
9 from vllm.logger import init_logger
10 from vllm.model_executor.warmup.flashinfer_autotune_cache import (
11 resolve_flashinfer_autotune_file,
12 write_flashinfer_autotune_cache,
13 )
14 from vllm.platforms import current_platform
15 from vllm.utils.flashinfer import autotune as flashinfer_autotune
16 from vllm.utils.flashinfer import has_flashinfer
17 from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
18
19 if TYPE_CHECKING:
20 from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2GPUModelRunner
21 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
22 from vllm.v1.worker.gpu_worker import Worker
23
24 logger = init_logger(__name__)
25
26 _DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(
27 {
28 "FLASHMLA_SPARSE_DSV4",
29 "FLASHINFER_MLA_SPARSE_DSV4",
30 "ROCM_FLASHMLA_SPARSE_DSV4",
31 "FLASHMLA_SPARSE_DSV41",
32 "FLASHINFER_MLA_SPARSE_DSV41",
33 "DEEPSEEK_SPARSE_SWA",
34 }
35 )
36 _FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"})
37 _DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset(
38 {"FLASHINFER_MLA_SPARSE_DSV4", "FLASHINFER_MLA_SPARSE_DSV41"}
39 )
40
41 _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = {
42 "FLASHINFER_MLA_SPARSE_SM120": "DSv3.2",
43 "FLASHINFER_MLA_SPARSE_DSV4": "DSv4",
44 "FLASHINFER_MLA_SPARSE_DSV41": "DSv4.1",
45 }
46
47 _SPARSE_MLA_MIXED_WARMUP_TOKENS = 16
48
49
50 def autotune_hisparse_flashinfer_attention(runner: "GPUModelRunner") -> None:
51 """Autotune each HiSparse FlashInfer sparse-MLA configuration."""
52 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
53 FlashInferMLASparseImpl,
54 )
55
56 tuned: set[tuple[object, ...]] = set()
57 for layer in runner.vllm_config.compilation_config.static_forward_context.values():
58 impl = getattr(layer, "impl", None)
59 if not isinstance(impl, FlashInferMLASparseImpl):
60 continue
61 if getattr(layer, "hisparse_cache", None) is None:
62 continue
63 if impl.topk_indices_buffer is None:
64 continue
65 key = (
66 impl.kv_cache_dtype,
67 impl.num_heads,
68 impl.qk_nope_head_dim,
69 impl.qk_rope_head_dim,
70 impl.kv_lora_rank,
71 impl.topk_indices_buffer.shape[1],
72 )
73 if key in tuned:
74 continue
75 impl.autotune_hisparse_decode(layer)
76 tuned.add(key)
77
78
79 def _attention_backend_name(backend: object) -> str | None:
80 get_name = getattr(backend, "get_name", None)
81 if get_name is None:
82 return None
83 try:
84 return get_name()
85 except NotImplementedError:
86 return None
87
88
89 def _has_deepseek_v4_sparse_mla_backend(runner: "GPUModelRunner") -> bool:
90 for groups in getattr(runner, "attn_groups", []) or ():
91 for group in groups:
92 name = _attention_backend_name(getattr(group, "backend", None))
93 if name in _DEEPSEEK_V4_SPARSE_MLA_BACKENDS:
94 return True
95 return False
96
97
98 def _flashinfer_sparse_mla_decode_label(
99 runner: "GPUModelRunner",
100 allowed_backends: frozenset[str],
101 ) -> str | None:
102 for groups in getattr(runner, "attn_groups", []) or ():
103 for group in groups:
104 name = _attention_backend_name(getattr(group, "backend", None))
105 if name in allowed_backends:
106 return _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS.get(name)
107 return None
108
109
110 def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int:
111 return max(0, min(num_tokens, max_tokens))
112
113
114 def _uses_v2_model_runner(runner: "GPUModelRunner") -> bool:
115 vllm_config = getattr(runner, "vllm_config", None)
116 return bool(getattr(vllm_config, "use_v2_model_runner", False))
117
118
119 def _run_flashinfer_sparse_mla_decode_autotune(
120 worker: "Worker",
121 num_tokens: int,
122 allowed_backends: frozenset[str],
123 ) -> bool:
124 """Autotune FlashInfer's SM120 sparse-MLA decode path."""
125 runner = worker.model_runner
126 log_label = _flashinfer_sparse_mla_decode_label(runner, allowed_backends)
127 if log_label is None:
128 return False
129 if worker.vllm_config.kernel_config.enable_flashinfer_autotune is not True:
130 return False
131 if not has_flashinfer() or not current_platform.is_device_capability_family(120):
132 return False
133
134 try:
135 from flashinfer.autotuner import AutoTuner
136 except ImportError:
137 logger.warning(
138 "Skipping FlashInfer SM120 sparse MLA decode autotune because "
139 "FlashInfer autotuner is unavailable."
140 )
141 return False
142
143 from vllm.distributed.parallel_state import get_world_group
144
145 world = get_world_group()
146 is_leader = world.rank_in_group == 0
147 cache_path = resolve_flashinfer_autotune_file(runner)
148
149 dummy_run_kwargs = dict(
150 num_tokens=num_tokens,
151 skip_eplb=True,
152 is_profile=True,
153 force_attention=True,
154 create_mixed_batch=True,
155 )
156
157 if is_leader:
158 logger.info(
159 "Autotuning FlashInfer SM120 sparse MLA %s decode with cache: %s",
160 log_label,
161 cache_path,
162 )
163
164 with torch.inference_mode():
165 warmup_executed = True
166 if is_leader:
167 if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
168 v2_runner = cast("V2GPUModelRunner", runner)
169 warmup_executed = run_mixed_prefill_decode_warmup(
170 v2_runner,
171 worker.execute_model,
172 worker.sample_tokens,
173 num_tokens,
174 mixed_step_context=flashinfer_autotune(True, cache=str(cache_path)),
175 req_id_prefix="_sparse_mla_v2_warmup",
176 )
177 else:
178 with flashinfer_autotune(True, cache=str(cache_path)):
179 runner._dummy_run(**dummy_run_kwargs)
180 else:
181 if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
182 v2_runner = cast("V2GPUModelRunner", runner)
183 warmup_executed = run_mixed_prefill_decode_warmup(
184 v2_runner,
185 worker.execute_model,
186 worker.sample_tokens,
187 num_tokens,
188 req_id_prefix="_sparse_mla_v2_warmup",
189 )
190 else:
191 runner._dummy_run(**dummy_run_kwargs)
192
193 if not warmup_executed:
194 return False
195
196 tune_results: bytes | None = None
197 if is_leader and cache_path.exists():
198 with open(cache_path, "rb") as f:
199 tune_results = f.read()
200
201 tune_results = world.broadcast_object(tune_results, src=0)
202 if tune_results is None:
203 logger.warning(
204 "No FlashInfer SM120 sparse MLA %s decode autotune cache entries found. "
205 "Falling back to FlashInfer's default tactic heuristic.",
206 log_label,
207 )
208 world.barrier()
209 return True
210
211 write_flashinfer_autotune_cache(cache_path, tune_results)
212 world.barrier()
213
214 AutoTuner.get().load_configs(str(cache_path))
215 logger.info(
216 "FlashInfer SM120 sparse MLA %s decode autotune cache loaded on rank %d "
217 "from %s.",
218 log_label,
219 world.rank_in_group,
220 cache_path,
221 )
222 return True
223
224
225 def _flashinfer_sparse_mla_decode_autotune(
226 worker: "Worker",
227 num_tokens: int,
228 ) -> bool:
229 return _run_flashinfer_sparse_mla_decode_autotune(
230 worker, num_tokens, _FLASHINFER_MLA_SPARSE_BACKENDS
231 )
232
233
234 def _deepseek_v4_sparse_mla_decode_autotune(
235 worker: "Worker",
236 num_tokens: int,
237 ) -> bool:
238 return _run_flashinfer_sparse_mla_decode_autotune(
239 worker, num_tokens, _DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS
240 )
241
242
243 def flashinfer_sparse_mla_decode_autotune_warmup(worker: "Worker") -> None:
244 """Autotune generic FlashInfer sparse MLA decode when selected."""
245 runner = worker.model_runner
246 if runner.is_pooling_model:
247 return
248
249 max_tokens = worker.scheduler_config.max_num_batched_tokens
250 mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens)
251 if mixed_tokens <= 0:
252 return
253 _flashinfer_sparse_mla_decode_autotune(worker, mixed_tokens)
254
255
256 def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None:
257 """Warm DSv4 sparse-MLA mixed prefill+decode attention."""
258 runner = worker.model_runner
259 if runner.is_pooling_model or not _has_deepseek_v4_sparse_mla_backend(runner):
260 return
261
262 max_tokens = worker.scheduler_config.max_num_batched_tokens
263 mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens)
264 if mixed_tokens <= 0:
265 return
266
267 logger.info(
268 "Warming up DeepSeek V4 sparse MLA attention for mixed tokens=%s.",
269 mixed_tokens,
270 )
271 mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens)
272 if not mixed_warmup_done:
273 if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2:
274 v2_runner = cast("V2GPUModelRunner", runner)
275 run_mixed_prefill_decode_warmup(
276 v2_runner,
277 worker.execute_model,
278 worker.sample_tokens,
279 mixed_tokens,
280 req_id_prefix="_sparse_mla_v2_warmup",
281 )
282 else:
283 runner._dummy_run(
284 num_tokens=mixed_tokens,
285 skip_eplb=True,
286 is_profile=True,
287 force_attention=True,
288 create_mixed_batch=True,
289 )