@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7ca017778f

[Feat][Perf] Add new warmup infrastructure for JITs (#47451)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Roberto L. Castro committed Jul 20, 2026 at 22:21 UTC 7ca017778fc04cfb1a96080e7c484b1b44fc3a64
14 files changed +1855 -693
.buildkite/test_areas/model_executor.yaml
+4
@@ -10,7 +10,9 @@ steps:
10 - vllm/engine/arg_utils.py
11 - vllm/config/model.py
12 - vllm/model_executor
13 + - vllm/model_executor/warmup
14 - tests/model_executor
15 + - tests/model_executor/test_jit_warmup.py
16 - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
17 commands:
18 - apt-get update && apt-get install -y curl libsodium23
@@ -34,7 +36,9 @@ steps:
36 - vllm/engine/arg_utils.py
37 - vllm/config/model.py
38 - vllm/model_executor
39 + - vllm/model_executor/warmup
40 - tests/model_executor
41 + - tests/model_executor/test_jit_warmup.py
42 - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
43 - vllm/_aiter_ops.py
44 - vllm/platforms/rocm.py
tests/model_executor/test_jit_warmup.py new
+306
@@ -0,0 +1,306 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import ast
5 +from dataclasses import dataclass
6 +from types import SimpleNamespace
7 +from typing import Any, cast
8 +
9 +import pytest
10 +
11 +from vllm.model_executor.warmup.jit_warmup import (
12 + VllmJitKernel,
13 + WarmupIntRange,
14 + get_ast_full_name,
15 + zip_inputs,
16 +)
17 +
18 +
19 +def _next_power_of_2(value: int) -> int:
20 + return 1 << max(0, value - 1).bit_length()
21 +
22 +
23 +def _round_up(value: int, *, multiple: int) -> int:
24 + return ((value + multiple - 1) // multiple) * multiple
25 +
26 +
27 +def _config(
28 + *,
29 + bias: int = 0,
30 + disabled: bool = False,
31 + name: str = "base",
32 + vectorized: bool = False,
33 +) -> SimpleNamespace:
34 + return SimpleNamespace(
35 + bias=bias,
36 + disabled=disabled,
37 + name=name,
38 + vectorized=vectorized,
39 + )
40 +
41 +
42 +class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
43 + @dataclass(frozen=True)
44 + class CompileKey:
45 + block_size: int
46 + work: int
47 + vector_width: int
48 + descriptor: tuple[object, ...]
49 + enabled: bool
50 +
51 + def dispatch( # type: ignore[override]
52 + self,
53 + *,
54 + tokens: int,
55 + cfg: Any,
56 + lanes: int = 1,
57 + mode: str = "default",
58 + debug: int = 0,
59 + ) -> CompileKey:
60 + block_size = _next_power_of_2(tokens)
61 + work: int = block_size * lanes + cfg.bias
62 + return self.CompileKey(
63 + block_size=block_size,
64 + work=work,
65 + vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
66 + descriptor=(
67 + cfg.name,
68 + mode,
69 + -block_size,
70 + block_size % 3,
71 + block_size**2,
72 + ),
73 + enabled=not cfg.disabled,
74 + )
75 +
76 + def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
77 + return self._trace_dispatch(self.dispatch)(
78 + tokens=WarmupIntRange(1, max_tokens + 1),
79 + cfg=cfg,
80 + # This argument is intentionally unused by dispatch expressions.
81 + debug=WarmupIntRange(0, 100),
82 + )
83 +
84 + def compile(self, compile_key: CompileKey) -> None:
85 + pass
86 +
87 +
88 +class RecordingToyKernel(ToyKernel):
89 + def __init__(self) -> None:
90 + self.compiled: list[ToyKernel.CompileKey] = []
91 + super().__init__()
92 +
93 + def compile(self, compile_key: ToyKernel.CompileKey) -> None:
94 + self.compiled.append(compile_key)
95 +
96 +
97 +def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
98 + cfg = _config()
99 +
100 + assert ToyKernel().get_warmup_keys(5, cfg) == [
101 + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
102 + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
103 + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
104 + ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
105 + ]
106 +
107 +
108 +def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
109 + cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
110 +
111 + assert ToyKernel().compile_key(
112 + {
113 + "tokens": 4,
114 + "cfg": cfg,
115 + "lanes": 2,
116 + }
117 + ) == ToyKernel.CompileKey(
118 + block_size=4,
119 + work=11,
120 + vector_width=4,
121 + descriptor=("cfg", "default", -4, 1, 16),
122 + enabled=False,
123 + )
124 +
125 +
126 +def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
127 + cfg = _config(vectorized=True)
128 +
129 + keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
130 + zip_inputs(
131 + dict(tokens=1, mode="small"),
132 + dict(tokens=4, mode="wide"),
133 + ),
134 + cfg=cfg,
135 + lanes=(1, 2),
136 + )
137 +
138 + assert keys == [
139 + ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
140 + ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
141 + ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
142 + ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
143 + ]
144 +
145 +
146 +def test_zip_inputs_validates_input_rows() -> None:
147 + with pytest.raises(ValueError, match="requires at least one"):
148 + zip_inputs()
149 + with pytest.raises(ValueError, match="rows must be mappings"):
150 + zip_inputs(cast(Any, ("tokens", 1)))
151 + with pytest.raises(ValueError, match="at least one dispatch input name"):
152 + zip_inputs({})
153 + with pytest.raises(ValueError, match="dispatch input names must be strings"):
154 + zip_inputs(cast(Any, {1: 2}))
155 + with pytest.raises(ValueError, match="same dispatch input names"):
156 + zip_inputs({"tokens": 1}, {"mode": "small"})
157 +
158 +
159 +def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
160 + kernel = ToyKernel()
161 +
162 + with pytest.raises(TypeError, match="zip_inputs"):
163 + kernel._trace_dispatch(kernel.dispatch)(
164 + cast(Any, {"tokens": 1}),
165 + cfg=_config(),
166 + )
167 +
168 + with pytest.raises(ValueError, match="specified more than once"):
169 + kernel._trace_dispatch(kernel.dispatch)(
170 + zip_inputs(dict(tokens=1, mode="small")),
171 + tokens=2,
172 + cfg=_config(),
173 + )
174 +
175 +
176 +def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
177 + class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
178 + @dataclass(frozen=True)
179 + class CompileKey:
180 + value: int
181 +
182 + def dispatch( # type: ignore[override]
183 + self,
184 + *,
185 + tokens: int,
186 + block_size: int,
187 + ) -> CompileKey:
188 + return self.CompileKey(value=_round_up(tokens, multiple=block_size))
189 +
190 + def get_warmup_keys(self) -> list[CompileKey]:
191 + return []
192 +
193 + def compile(self, compile_key: CompileKey) -> None:
194 + pass
195 +
196 + class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
197 + @dataclass(frozen=True)
198 + class CompileKey:
199 + value: int
200 +
201 + def dispatch( # type: ignore[override]
202 + self,
203 + *,
204 + tokens: int,
205 + block_size: int,
206 + ) -> CompileKey:
207 + return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
208 +
209 + def get_warmup_keys(self) -> list[CompileKey]:
210 + return []
211 +
212 + def compile(self, compile_key: CompileKey) -> None:
213 + pass
214 +
215 + assert HelperKernel().compile_key(
216 + {
217 + "tokens": 5,
218 + "block_size": 4,
219 + }
220 + ) == HelperKernel.CompileKey(value=8)
221 + with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
222 + StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
223 +
224 +
225 +def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
226 + class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
227 + @dataclass(frozen=True)
228 + class CompileKey:
229 + value: int
230 +
231 + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
232 + if value > 0:
233 + value = 1
234 + return self.CompileKey(value=value)
235 +
236 + def get_warmup_keys(self) -> list[CompileKey]:
237 + return []
238 +
239 + def compile(self, compile_key: CompileKey) -> None:
240 + pass
241 +
242 + class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
243 + @dataclass(frozen=True)
244 + class CompileKey:
245 + value: int
246 +
247 + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
248 + return self.CompileKey(**{"value": value})
249 +
250 + def get_warmup_keys(self) -> list[CompileKey]:
251 + return []
252 +
253 + def compile(self, compile_key: CompileKey) -> None:
254 + pass
255 +
256 + with pytest.raises(ValueError, match="local assignments"):
257 + BranchKernel()
258 + with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
259 + KwargsReturnKernel()
260 +
261 +
262 +def test_dispatch_reports_unsupported_expression_with_context() -> None:
263 + class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
264 + @dataclass(frozen=True)
265 + class CompileKey:
266 + value: object
267 +
268 + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
269 + return self.CompileKey(value={value})
270 +
271 + def get_warmup_keys(self) -> list[CompileKey]:
272 + return []
273 +
274 + def compile(self, compile_key: CompileKey) -> None:
275 + pass
276 +
277 + with pytest.raises(ValueError) as exc_info:
278 + UnsupportedKernel().compile_key({"value": 1})
279 +
280 + message = str(exc_info.value)
281 + assert "Unsupported dispatch expression" in message
282 + assert "{value}" in message
283 + assert "Supported dispatch expressions" in message
284 +
285 +
286 +def test_warmup_compiles_all_returned_keys_in_order() -> None:
287 + kernel = RecordingToyKernel()
288 + cfg = _config()
289 +
290 + kernel.warmup(3, cfg)
291 +
292 + assert kernel.compiled == [
293 + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
294 + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
295 + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
296 + ]
297 +
298 +
299 +def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
300 + dotted_expr = ast.parse("foo.bar.baz").body[0]
301 + call_expr = ast.parse("foo()").body[0]
302 + assert isinstance(dotted_expr, ast.Expr)
303 + assert isinstance(call_expr, ast.Expr)
304 +
305 + assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
306 + assert get_ast_full_name(call_expr.value) is None
vllm/config/kernel.py
+8
@@ -175,9 +175,15 @@ class KernelConfig:
175 enable_flashinfer_autotune: bool = None # type: ignore[assignment]
176 """If True, run FlashInfer autotuning during kernel warmup."""
177
178 + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
179 + # to the shared JIT warmup infrastructure.
180 + # https://github.com/vllm-project/vllm/pull/47451
181 enable_cutedsl_warmup: bool = True
182 """If True, run CuTeDSL compile warmup during kernel warmup."""
183
184 + enable_jit_warmup: bool = True
185 + """If True, run JIT compile warmup during kernel warmup."""
186 +
187 enable_bf16x3_router_gemm: bool = False
188 """If True, use the experimental SM100 BF16x3 CuteDSL router GEMM."""
189
@@ -250,6 +256,7 @@ class KernelConfig:
256 """
257 ignored_factors = {
258 "enable_cutedsl_warmup",
259 + "enable_jit_warmup",
260 "enable_flashinfer_autotune",
261 "ir_op_priority", # handled separately below
262 }
@@ -260,6 +267,7 @@ class KernelConfig:
267 @field_validator(
268 "enable_flashinfer_autotune",
269 "enable_cutedsl_warmup",
270 + "enable_jit_warmup",
271 mode="wrap",
272 )
273 @classmethod
vllm/model_executor/warmup/cutedsl_warmup.py
+4
@@ -2,6 +2,10 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Run registered CuTeDSL warmup compile units."""
4
5 +# TODO(roberto): Remove this compatibility registry after registered CuTeDSL
6 +# warmups are migrated to the shared JIT warmup infrastructure.
7 +# https://github.com/vllm-project/vllm/pull/47451
8 +
9 from __future__ import annotations
10
11 import time
vllm/model_executor/warmup/fa4_cutedsl_config.py deleted
-204
@@ -1,204 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""FA4 MLA prefill CuTeDSL compile warmup config."""
4 -
5 -from __future__ import annotations
6 -
7 -from collections.abc import Hashable, Iterator
8 -from dataclasses import dataclass
9 -from typing import TYPE_CHECKING, Literal
10 -
11 -import torch
12 -
13 -if TYPE_CHECKING:
14 - from vllm.v1.attention.backends.fa_utils import (
15 - FlashAttentionCuTeDSLCompileSpec,
16 - )
17 -
18 -FA4ArchitectureFamily = Literal["sm90", "sm100f", "sm120"]
19 -
20 -FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16)
21 -
22 -# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan
23 -# covers qhead_per_kvhead=1.
24 -# Batch is not a current FA4 MLA-prefill key field. Use b1 for compile-only
25 -# specs because it is the conservative case for Split-KV shape heuristics.
26 -# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but vLLM
27 -# does not use them in this backend yet; they need a separate
28 -# num_kv_heads/qv/top-k-aware warmup plan if wired in later.
29 -FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1
30 -FA4_MLA_PREFILL_Q_TILE = 128
31 -FA4_MLA_PREFILL_K_TILE = 128
32 -FA4_MLA_PREFILL_LONG_K_BLOCKS = 32
33 -FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64
34 -FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True)
35 -FA4_MLA_PREFILL_LSE_OPTIONS = (False, True)
36 -
37 -
38 -@dataclass(frozen=True)
39 -class FA4MLAPrefillCompileContext:
40 - dtype: torch.dtype
41 - num_heads: int
42 - qk_head_dim: int
43 - v_head_dim: int
44 - kv_nope_head_dim: int
45 - requires_v_padding: bool
46 - scale: float
47 - num_splits: int
48 - fa_version: int
49 -
50 - # Return the V head dim FA4 sees.
51 - @property
52 - def effective_v_head_dim(self) -> int:
53 - if self.requires_v_padding:
54 - return self.qk_head_dim
55 - return self.v_head_dim
56 -
57 -
58 -@dataclass(frozen=True)
59 -class FA4MLAPrefillCompileRequest:
60 - """One compile-only FA4 MLA prefill request."""
61 -
62 - key: Hashable
63 - compile_spec: FlashAttentionCuTeDSLCompileSpec
64 -
65 - # Compile this request.
66 - def compile(self) -> None:
67 - self.compile_spec.compile()
68 -
69 -
70 -# Yield deduped compile requests.
71 -def iter_fa4_mla_prefill_compile_requests(
72 - ctx: FA4MLAPrefillCompileContext,
73 -) -> Iterator[FA4MLAPrefillCompileRequest]:
74 - """Yield compile requests for this fixed MLA backend.
75 -
76 - FA4 dedupes duplicate atomic kernel selections in its own JIT cache.
77 - """
78 - seen: set[Hashable] = set()
79 - for compile_spec in iter_fa4_mla_prefill_compile_specs(ctx):
80 - key = compile_spec.request_key()
81 - if key in seen:
82 - continue
83 - seen.add(key)
84 - yield FA4MLAPrefillCompileRequest(
85 - key=key,
86 - compile_spec=compile_spec,
87 - )
88 -
89 -
90 -# Build compile specs for this setup.
91 -def iter_fa4_mla_prefill_compile_specs(
92 - ctx: FA4MLAPrefillCompileContext,
93 -) -> Iterator[FlashAttentionCuTeDSLCompileSpec]:
94 - """Yield compile-only FA4 MLA prefill requests for this fixed setup."""
95 -
96 - arch_family = _fa4_architecture_family_from_compute_capability(
97 - *torch.cuda.get_device_capability()
98 - )
99 - if not _supports_fa4_mla_prefill(ctx, arch_family):
100 - return
101 -
102 - from vllm.v1.attention.backends.fa_utils import (
103 - FlashAttentionCuTeDSLCompileSpec,
104 - )
105 -
106 - batch_size = FA4_MLA_PREFILL_COMPILE_BATCH_SIZE
107 - v_stride = None
108 - if not ctx.requires_v_padding:
109 - v_stride = (
110 - ctx.num_heads * ctx.kv_nope_head_dim,
111 - ctx.kv_nope_head_dim,
112 - 1,
113 - )
114 -
115 - for _, max_seqlen_q, max_seqlen_k in _shape_probes_for_context(ctx, arch_family):
116 - total_q_tokens = batch_size * max_seqlen_q
117 - total_kv_tokens = batch_size * max_seqlen_k
118 - for causal in FA4_MLA_PREFILL_CAUSAL_OPTIONS:
119 - for return_lse in FA4_MLA_PREFILL_LSE_OPTIONS:
120 - yield FlashAttentionCuTeDSLCompileSpec(
121 - q_shape=(total_q_tokens, ctx.num_heads, ctx.qk_head_dim),
122 - k_shape=(total_kv_tokens, ctx.num_heads, ctx.qk_head_dim),
123 - v_shape=(
124 - total_kv_tokens,
125 - ctx.num_heads,
126 - ctx.effective_v_head_dim,
127 - ),
128 - v_stride=v_stride,
129 - q_dtype=ctx.dtype,
130 - cu_seqlens_q_shape=(batch_size + 1,),
131 - cu_seqlens_k_shape=(batch_size + 1,),
132 - max_seqlen_q=max_seqlen_q,
133 - max_seqlen_k=max_seqlen_k,
134 - softmax_scale=ctx.scale,
135 - causal=causal,
136 - return_softmax_lse=return_lse,
137 - num_splits=ctx.num_splits,
138 - fa_version=ctx.fa_version,
139 - )
140 -
141 -
142 -# Pick one q/k point per current FA4 MLA-prefill shape regime.
143 -def _shape_probes_for_context(
144 - ctx: FA4MLAPrefillCompileContext,
145 - arch_family: FA4ArchitectureFamily,
146 -) -> tuple[tuple[str, int, int], ...]:
147 - q_stage1_q = 1
148 - q_stage2_q = FA4_MLA_PREFILL_Q_TILE + 1
149 - # FA4 never auto-splits when ceil(max_seqlen_k / tile_n) <= 4.
150 - no_split_k = 4 * FA4_MLA_PREFILL_K_TILE
151 - long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
152 - # Diff-head-dim Blackwell Split-KV switches tile_n at 64 K blocks.
153 - very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
154 -
155 - base_probes = (
156 - ("q_stage1", q_stage1_q, FA4_MLA_PREFILL_K_TILE),
157 - ("q_stage2", q_stage2_q, no_split_k),
158 - )
159 - # SM120 currently rejects Split-KV in FA4; num_splits=1 also has no split
160 - # shape regimes on any architecture.
161 - if ctx.num_splits == 1 or arch_family == "sm120":
162 - return base_probes
163 -
164 - long_k_probes = (
165 - ("q_stage1_long_k", q_stage1_q, long_k),
166 - ("q_stage2_long_k", q_stage2_q, long_k),
167 - )
168 -
169 - # SM90 does not have the SM100 q_stage or diff-head-dim tile_n=64 branch.
170 - # Same-dim SM100-family MLA also does not need the very-long-K probe.
171 - if arch_family == "sm90" or ctx.qk_head_dim == ctx.effective_v_head_dim:
172 - return (*base_probes, *long_k_probes)
173 -
174 - very_long_k_probes = (
175 - ("q_stage1_very_long_k", q_stage1_q, very_long_k),
176 - ("q_stage2_very_long_k", q_stage2_q, very_long_k),
177 - )
178 - return (*base_probes, *long_k_probes, *very_long_k_probes)
179 -
180 -
181 -# Check whether this setup can use FA4 MLA prefill.
182 -def _supports_fa4_mla_prefill(
183 - ctx: FA4MLAPrefillCompileContext,
184 - arch_family: FA4ArchitectureFamily,
185 -) -> bool:
186 - return (
187 - ctx.dtype in FA4_STANDARD_DTYPES
188 - and ctx.num_heads > 0
189 - and (arch_family != "sm120" or ctx.num_splits == 1)
190 - )
191 -
192 -
193 -# Map CUDA capability to the FA4 arch family used by warmup checks.
194 -def _fa4_architecture_family_from_compute_capability(
195 - major: int,
196 - minor: int,
197 -) -> FA4ArchitectureFamily:
198 - if (major, minor) == (9, 0):
199 - return "sm90"
200 - if major == 10:
201 - return "sm100f"
202 - if (major, minor) == (12, 0):
203 - return "sm120"
204 - raise ValueError(f"FA4 warmup does not know CUDA capability {major}.{minor}")
vllm/model_executor/warmup/fa4_cutedsl_warmup.py new
+30
@@ -0,0 +1,30 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Warm up FA4 CuTeDSL MLA prefill compile keys."""
4 +
5 +from __future__ import annotations
6 +
7 +from typing import TYPE_CHECKING
8 +
9 +from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
10 +
11 +if TYPE_CHECKING:
12 + from vllm.v1.worker.gpu_worker import Worker
13 +
14 +
15 +def fa4_cutedsl_warmup(worker: Worker) -> None:
16 + runner = worker.model_runner
17 + if runner.is_pooling_model:
18 + return
19 +
20 + vllm_config = runner.vllm_config
21 + if not vllm_config.model_config.use_mla:
22 + return
23 +
24 + backend_cls = get_mla_prefill_backend(vllm_config)
25 + if backend_cls.get_name() != "FLASH_ATTN":
26 + return
27 +
28 + from vllm.v1.attention.backends.mla.prefill import flash_attn
29 +
30 + flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config)
vllm/model_executor/warmup/jit_warmup.py new
+474
@@ -0,0 +1,474 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared interfaces and tracing helpers for explicit JIT warmup keys."""
4 +
5 +from __future__ import annotations
6 +
7 +import ast
8 +import inspect
9 +import itertools
10 +import operator
11 +import textwrap
12 +from abc import ABC, abstractmethod
13 +from collections.abc import Callable, Iterable, Mapping
14 +from dataclasses import dataclass
15 +from typing import Any, Generic, TypeVar
16 +
17 +__all__ = [
18 + "VllmJitKernel",
19 + "WarmupIntRange",
20 + "get_ast_full_name",
21 + "get_function_source_node",
22 + "zip_inputs",
23 +]
24 +
25 +
26 +CompileKeyT = TypeVar("CompileKeyT")
27 +
28 +
29 +@dataclass(frozen=True)
30 +class WarmupIntRange:
31 + start: int
32 + stop: int
33 + step: int = 1
34 +
35 +
36 +WarmupValues = Any
37 +CompileKeyDispatchFn = Callable[..., CompileKeyT]
38 +
39 +
40 +@dataclass(frozen=True)
41 +class _WarmupInputRows:
42 + """Warmup dispatch inputs expanded in lockstep."""
43 +
44 + rows: tuple[Mapping[str, WarmupValues], ...]
45 +
46 +
47 +def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:
48 + if isinstance(values, WarmupIntRange):
49 + return tuple(range(values.start, values.stop, values.step))
50 + if isinstance(values, (list, tuple)):
51 + return tuple(values)
52 + return (values,)
53 +
54 +
55 +def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows:
56 + """Group row-wise dispatch inputs that should be expanded in lockstep."""
57 + if not rows:
58 + raise ValueError("zip_inputs requires at least one dispatch input row")
59 + if not all(isinstance(row, Mapping) for row in rows):
60 + raise ValueError("zip_inputs rows must be mappings")
61 +
62 + first_names = frozenset(rows[0])
63 + if not first_names:
64 + raise ValueError("zip_inputs rows require at least one dispatch input name")
65 + if not all(isinstance(name, str) for name in first_names):
66 + raise ValueError("zip_inputs dispatch input names must be strings")
67 +
68 + input_rows: list[Mapping[str, WarmupValues]] = []
69 + for row in rows:
70 + names = frozenset(row)
71 + if names != first_names:
72 + raise ValueError("zip_inputs rows must use the same dispatch input names")
73 + input_rows.append(dict(row))
74 +
75 + return _WarmupInputRows(rows=tuple(input_rows))
76 +
77 +
78 +def _expand_warmup_value_grid(
79 + values: Mapping[str, WarmupValues],
80 + input_names: frozenset[str],
81 +) -> tuple[dict[str, Any], ...]:
82 + names = tuple(name for name in values if name in input_names)
83 + if not names:
84 + return ({},)
85 +
86 + expanded_values = tuple(_expand_warmup_values(values[name]) for name in names)
87 + return tuple(
88 + dict(zip(names, value_set)) for value_set in itertools.product(*expanded_values)
89 + )
90 +
91 +
92 +def _expand_warmup_input_rows(
93 + rows: tuple[Mapping[str, WarmupValues], ...],
94 + input_names: frozenset[str],
95 +) -> tuple[dict[str, Any], ...]:
96 + active_names = frozenset(name for name in rows[0] if name in input_names)
97 + if not active_names:
98 + return ({},)
99 +
100 + return tuple(
101 + {name: value for name, value in row.items() if name in active_names}
102 + for row in rows
103 + )
104 +
105 +
106 +def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]:
107 + merged: dict[str, Any] = {}
108 + for part in parts:
109 + for name, value in part.items():
110 + if name in merged:
111 + raise ValueError(
112 + f"Warmup dispatch input '{name}' is specified more than once"
113 + )
114 + merged[name] = value
115 + return merged
116 +
117 +
118 +@dataclass(frozen=True)
119 +class _CompileKeyDispatchTrace:
120 + local_exprs: tuple[tuple[str, ast.AST], ...]
121 + field_exprs: tuple[tuple[str, ast.AST], ...]
122 + globals: Mapping[str, Any]
123 + input_names: frozenset[str]
124 + defaults: Mapping[str, Any]
125 +
126 + def compile_key(
127 + self,
128 + compile_key_type: type[CompileKeyT],
129 + kwargs: Mapping[str, Any],
130 + ) -> CompileKeyT:
131 + dispatch_values = {**self.defaults, **kwargs}
132 + for name, expr in self.local_exprs:
133 + dispatch_values[name] = _eval_dispatch_expr(
134 + expr, dispatch_values, self.globals
135 + )
136 + return compile_key_type(
137 + **{
138 + field: _eval_dispatch_expr(expr, dispatch_values, self.globals)
139 + for field, expr in self.field_exprs
140 + }
141 + )
142 +
143 +
144 +_BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
145 + ast.Add: operator.add,
146 + ast.Sub: operator.sub,
147 + ast.Mult: operator.mul,
148 + ast.FloorDiv: operator.floordiv,
149 + ast.Mod: operator.mod,
150 + ast.Pow: operator.pow,
151 +}
152 +_CMP_OPS: dict[type[ast.cmpop], Callable[[Any, Any], bool]] = {
153 + ast.Eq: operator.eq,
154 + ast.NotEq: operator.ne,
155 + ast.Lt: operator.lt,
156 + ast.LtE: operator.le,
157 + ast.Gt: operator.gt,
158 + ast.GtE: operator.ge,
159 +}
160 +
161 +
162 +def _dispatch_expr_source(node: ast.AST) -> str:
163 + try:
164 + return ast.unparse(node)
165 + except Exception:
166 + return ast.dump(node)
167 +
168 +
169 +def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError:
170 + return ValueError(
171 + f"{reason}: {_dispatch_expr_source(node)}. "
172 + "Supported dispatch expressions are names, constants, attributes, "
173 + "tuple/list literals, conditional expressions, comparisons, boolean "
174 + "operators, unary not/minus, arithmetic, and calls without **kwargs."
175 + )
176 +
177 +
178 +class _DispatchExprEvaluator(ast.NodeVisitor):
179 + def __init__(
180 + self,
181 + values: Mapping[str, Any],
182 + globals_: Mapping[str, Any],
183 + ) -> None:
184 + self.values = values
185 + self.globals = globals_
186 +
187 + def eval(self, node: ast.AST) -> Any:
188 + return self.visit(node)
189 +
190 + def generic_visit(self, node: ast.AST) -> Any:
191 + raise _dispatch_expr_error(node, "Unsupported dispatch expression")
192 +
193 + def visit_Name(self, node: ast.Name) -> Any:
194 + if node.id in self.values:
195 + return self.values[node.id]
196 + if node.id in self.globals:
197 + return self.globals[node.id]
198 + raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'")
199 +
200 + def visit_Constant(self, node: ast.Constant) -> Any:
201 + return node.value
202 +
203 + def visit_IfExp(self, node: ast.IfExp) -> Any:
204 + return self.visit(node.body if self.visit(node.test) else node.orelse)
205 +
206 + def visit_Tuple(self, node: ast.Tuple) -> tuple[Any, ...]:
207 + return tuple(self.visit(elt) for elt in node.elts)
208 +
209 + def visit_List(self, node: ast.List) -> list[Any]:
210 + return [self.visit(elt) for elt in node.elts]
211 +
212 + def visit_BoolOp(self, node: ast.BoolOp) -> Any:
213 + if isinstance(node.op, ast.And):
214 + result = None
215 + for value in node.values:
216 + result = self.visit(value)
217 + if not result:
218 + return result
219 + return result
220 + if isinstance(node.op, ast.Or):
221 + result = None
222 + for value in node.values:
223 + result = self.visit(value)
224 + if result:
225 + return result
226 + return result
227 + raise _dispatch_expr_error(node, "Unsupported dispatch boolean operator")
228 +
229 + def visit_Compare(self, node: ast.Compare) -> bool:
230 + left = self.visit(node.left)
231 + for op_node, comparator in zip(node.ops, node.comparators):
232 + right = self.visit(comparator)
233 + op = _CMP_OPS.get(type(op_node))
234 + if op is None:
235 + raise _dispatch_expr_error(
236 + node, "Unsupported dispatch comparison operator"
237 + )
238 + if not op(left, right):
239 + return False
240 + left = right
241 + return True
242 +
243 + def visit_UnaryOp(self, node: ast.UnaryOp) -> Any:
244 + operand = self.visit(node.operand)
245 + if isinstance(node.op, ast.Not):
246 + return not operand
247 + if isinstance(node.op, ast.USub):
248 + return -operand
249 + raise _dispatch_expr_error(node, "Unsupported dispatch unary operator")
250 +
251 + def visit_BinOp(self, node: ast.BinOp) -> Any:
252 + op = _BIN_OPS.get(type(node.op))
253 + if op is None:
254 + raise _dispatch_expr_error(node, "Unsupported dispatch binary operator")
255 + return op(self.visit(node.left), self.visit(node.right))
256 +
257 + def visit_Call(self, node: ast.Call) -> Any:
258 + args = [self.visit(arg) for arg in node.args]
259 + fn = self.visit(node.func)
260 + call_kwargs: dict[str, Any] = {}
261 + for keyword in node.keywords:
262 + if keyword.arg is None:
263 + raise _dispatch_expr_error(
264 + node, "Dispatch helper calls cannot use **kwargs"
265 + )
266 + call_kwargs[keyword.arg] = self.visit(keyword.value)
267 + return fn(*args, **call_kwargs)
268 +
269 + def visit_Attribute(self, node: ast.Attribute) -> Any:
270 + return getattr(self.visit(node.value), node.attr)
271 +
272 +
273 +def get_ast_full_name(node: ast.AST) -> str | None:
274 + if isinstance(node, ast.Name):
275 + return node.id
276 + if isinstance(node, ast.Attribute):
277 + parent = get_ast_full_name(node.value)
278 + if parent is not None:
279 + return f"{parent}.{node.attr}"
280 + return None
281 +
282 +
283 +def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef:
284 + source_fn = getattr(fn, "fn", fn)
285 + source = textwrap.dedent(inspect.getsource(source_fn))
286 + tree = ast.parse(source)
287 + function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
288 + if len(function_defs) != 1:
289 + name = getattr(source_fn, "__name__", type(source_fn).__name__)
290 + raise ValueError(f"Expected one function in {name}, found {len(function_defs)}")
291 + return function_defs[0]
292 +
293 +
294 +def _eval_dispatch_expr(
295 + node: ast.AST,
296 + kwargs: Mapping[str, Any],
297 + globals_: Mapping[str, Any],
298 +) -> Any:
299 + return _DispatchExprEvaluator(kwargs, globals_).eval(node)
300 +
301 +
302 +def _collect_input_names(
303 + node: ast.AST,
304 + candidate_names: set[str],
305 + local_names: set[str] | None = None,
306 +) -> set[str]:
307 + if local_names is None:
308 + local_names = set()
309 + return {
310 + child.id
311 + for child in ast.walk(node)
312 + if (
313 + isinstance(child, ast.Name)
314 + and child.id in candidate_names
315 + and child.id not in local_names
316 + )
317 + }
318 +
319 +
320 +def _collect_dispatch_body(
321 + fn: CompileKeyDispatchFn[Any],
322 + function_def: ast.FunctionDef,
323 +) -> tuple[list[tuple[str, ast.AST]], ast.Call]:
324 + local_exprs: list[tuple[str, ast.AST]] = []
325 + for statement in function_def.body:
326 + if (
327 + isinstance(statement, ast.Expr)
328 + and isinstance(statement.value, ast.Constant)
329 + and isinstance(statement.value.value, str)
330 + ):
331 + continue
332 +
333 + if isinstance(statement, ast.Assign):
334 + if len(statement.targets) != 1 or not isinstance(
335 + statement.targets[0], ast.Name
336 + ):
337 + raise _dispatch_expr_error(
338 + statement, "Dispatch assignments must target one local name"
339 + )
340 + local_exprs.append((statement.targets[0].id, statement.value))
341 + continue
342 +
343 + if isinstance(statement, ast.AnnAssign):
344 + if statement.value is None:
345 + raise _dispatch_expr_error(
346 + statement, "Dispatch annotations must assign a value"
347 + )
348 + if not isinstance(statement.target, ast.Name):
349 + raise _dispatch_expr_error(
350 + statement, "Dispatch assignments must target one local name"
351 + )
352 + local_exprs.append((statement.target.id, statement.value))
353 + continue
354 +
355 + if isinstance(statement, ast.Return) and isinstance(statement.value, ast.Call):
356 + return local_exprs, statement.value
357 +
358 + if isinstance(statement, ast.Return):
359 + raise _dispatch_expr_error(
360 + statement, "Dispatch must return one CompileKey(...) call"
361 + )
362 +
363 + raise _dispatch_expr_error(
364 + statement,
365 + "Dispatch may only contain local assignments before CompileKey return",
366 + )
367 +
368 + raise ValueError(f"Expected {fn.__name__} to return one CompileKey(...) call")
369 +
370 +
371 +def _trace_compile_key_dispatch(
372 + fn: CompileKeyDispatchFn[Any],
373 +) -> _CompileKeyDispatchTrace:
374 + source_fn = getattr(fn, "__func__", fn)
375 + globals_ = source_fn.__globals__
376 + function_def = get_function_source_node(fn)
377 +
378 + local_exprs, return_call = _collect_dispatch_body(fn, function_def)
379 +
380 + field_exprs: list[tuple[str, ast.AST]] = []
381 + signature = inspect.signature(fn)
382 + defaults = {
383 + name: parameter.default
384 + for name, parameter in signature.parameters.items()
385 + if parameter.default is not inspect.Parameter.empty
386 + }
387 + candidate_names = set(signature.parameters)
388 + input_names: set[str] = set()
389 + local_names = {name for name, _ in local_exprs}
390 + for _, expr in local_exprs:
391 + input_names.update(_collect_input_names(expr, candidate_names))
392 + for keyword in return_call.keywords:
393 + if keyword.arg is None:
394 + raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey")
395 + field_exprs.append((keyword.arg, keyword.value))
396 + input_names.update(
397 + _collect_input_names(keyword.value, candidate_names, local_names)
398 + )
399 +
400 + return _CompileKeyDispatchTrace(
401 + tuple(local_exprs),
402 + tuple(field_exprs),
403 + globals_,
404 + frozenset(input_names),
405 + defaults,
406 + )
407 +
408 +
409 +class VllmJitKernel(Generic[CompileKeyT], ABC):
410 + """Kernel wrapper that owns dispatch, warmup keys, and compilation."""
411 +
412 + CompileKey: type[CompileKeyT]
413 +
414 + def __init__(self) -> None:
415 + self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
416 +
417 + def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
418 + return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs)
419 +
420 + def _trace_dispatch(
421 + self, dispatch: CompileKeyDispatchFn[CompileKeyT]
422 + ) -> Callable[..., list[CompileKeyT]]:
423 + compile_key_dispatch_trace = _trace_compile_key_dispatch(dispatch)
424 +
425 + def traced(
426 + *input_groups: _WarmupInputRows,
427 + **kwargs: WarmupValues,
428 + ) -> list[CompileKeyT]:
429 + for group in input_groups:
430 + if not isinstance(group, _WarmupInputRows):
431 + raise TypeError(
432 + "_trace_dispatch positional arguments must be "
433 + "zip_inputs(...) groups"
434 + )
435 + expanded_input_groups = tuple(
436 + _expand_warmup_input_rows(
437 + group.rows, compile_key_dispatch_trace.input_names
438 + )
439 + for group in input_groups
440 + )
441 + expanded_kwargs = _expand_warmup_value_grid(
442 + kwargs, compile_key_dispatch_trace.input_names
443 + )
444 + dispatch_value_groups = (*expanded_input_groups, expanded_kwargs)
445 + return list(
446 + dict.fromkeys(
447 + compile_key_dispatch_trace.compile_key(
448 + self.CompileKey, _merge_warmup_kwargs(dispatch_values)
449 + )
450 + for dispatch_values in itertools.product(*dispatch_value_groups)
451 + )
452 + )
453 +
454 + return traced
455 +
456 + @abstractmethod
457 + def dispatch(self, **kwargs: Any) -> CompileKeyT:
458 + """Build one compile key from one concrete dispatch point."""
459 + raise NotImplementedError
460 +
461 + @abstractmethod
462 + def get_warmup_keys(self, *args: Any, **kwargs: Any) -> list[CompileKeyT]:
463 + """Return compile keys that should be warmed for this kernel."""
464 + raise NotImplementedError
465 +
466 + @abstractmethod
467 + def compile(self, compile_key: CompileKeyT) -> None:
468 + """Compile one warmup key."""
469 + raise NotImplementedError
470 +
471 + def warmup(self, *args: Any, **kwargs: Any) -> None:
472 + """Compile this kernel's warmup keys."""
473 + for compile_key in self.get_warmup_keys(*args, **kwargs):
474 + self.compile(compile_key)
vllm/model_executor/warmup/jit_warmup_triton_helper.py new
+183
@@ -0,0 +1,183 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import ast
4 +import inspect
5 +from collections.abc import Callable
6 +from dataclasses import dataclass
7 +from typing import Any
8 +
9 +from vllm.model_executor.warmup.jit_warmup import (
10 + get_ast_full_name,
11 + get_function_source_node,
12 +)
13 +
14 +
15 +@dataclass(frozen=True)
16 +class TritonWarmupTensor:
17 + # Compile-only tensor descriptor for Triton pointer specialization.
18 + dtype: Any
19 + aligned: bool = True
20 + shape: tuple[int, ...] = (1,)
21 +
22 + def data_ptr(self) -> int:
23 + return 0 if self.aligned else 1
24 +
25 + def ptr_range(self) -> int:
26 + return 0
27 +
28 + def stride(self) -> tuple[int, ...]:
29 + strides: list[int] = []
30 + stride = 1
31 + for size in reversed(self.shape):
32 + strides.append(stride)
33 + stride *= size
34 + return tuple(reversed(strides))
35 +
36 +
37 +@dataclass(frozen=True)
38 +class TritonPointerInputVariant:
39 + # Named pointer-alignment variant for compile-only Triton warmup.
40 + alignments: tuple[tuple[str, bool], ...]
41 +
42 + @classmethod
43 + def from_alignment(cls, **aligned: bool) -> "TritonPointerInputVariant":
44 + return cls(tuple(aligned.items()))
45 +
46 + def is_aligned(self, name: str) -> bool:
47 + for alignment_name, aligned in self.alignments:
48 + if alignment_name == name:
49 + return aligned
50 + raise KeyError(f"Unknown Triton pointer input variant: {name}")
51 +
52 + def pointer(
53 + self,
54 + name: str,
55 + dtype: Any,
56 + shape: tuple[int, ...] = (1,),
57 + ) -> TritonWarmupTensor:
58 + return TritonWarmupTensor(dtype, aligned=self.is_aligned(name), shape=shape)
59 +
60 +
61 +def _literal_str_refs(node: ast.AST) -> tuple[str | int, ...]:
62 + if isinstance(node, ast.Constant) and isinstance(node.value, str | int):
63 + return (node.value,)
64 + if isinstance(node, ast.List | ast.Tuple):
65 + refs: list[str | int] = []
66 + for elt in node.elts:
67 + if isinstance(elt, ast.Constant) and isinstance(elt.value, str | int):
68 + refs.append(elt.value)
69 + else:
70 + raise ValueError(
71 + f"Unsupported Triton specialization ref: {ast.dump(elt)}"
72 + )
73 + return tuple(refs)
74 + raise ValueError(f"Unsupported Triton specialization refs: {ast.dump(node)}")
75 +
76 +
77 +def _normalize_arg_refs(
78 + refs: tuple[str | int, ...],
79 + arg_names: tuple[str, ...],
80 +) -> frozenset[str]:
81 + names: set[str] = set()
82 + for ref in refs:
83 + if isinstance(ref, int):
84 + names.add(arg_names[ref])
85 + else:
86 + names.add(ref)
87 + return frozenset(names)
88 +
89 +
90 +def _decorator_keyword_refs(
91 + function_def: ast.FunctionDef,
92 + keyword_name: str,
93 +) -> tuple[str | int, ...]:
94 + for decorator in function_def.decorator_list:
95 + if not isinstance(decorator, ast.Call):
96 + continue
97 + decorator_name = get_ast_full_name(decorator.func)
98 + if decorator_name not in ("triton.jit", "jit"):
99 + continue
100 + for keyword in decorator.keywords:
101 + if keyword.arg == keyword_name:
102 + return _literal_str_refs(keyword.value)
103 + return ()
104 +
105 +
106 +def _triton_do_not_specialize_args(
107 + kernel: Callable[..., Any],
108 + function_def: ast.FunctionDef,
109 + arg_names: tuple[str, ...],
110 +) -> frozenset[str]:
111 + refs = getattr(kernel, "do_not_specialize", None)
112 + if refs is not None:
113 + return _normalize_arg_refs(tuple(refs), arg_names)
114 + return _normalize_arg_refs(
115 + _decorator_keyword_refs(function_def, "do_not_specialize"),
116 + arg_names,
117 + )
118 +
119 +
120 +def _triton_constexpr_arg_names(
121 + kernel: Callable[..., Any],
122 + function_def: ast.FunctionDef,
123 + arg_names: tuple[str, ...],
124 +) -> frozenset[str]:
125 + constexprs = getattr(kernel, "constexprs", None)
126 + if constexprs is not None:
127 + return frozenset(arg_names[index] for index in constexprs)
128 +
129 + names: set[str] = set()
130 + for arg in function_def.args.args + function_def.args.kwonlyargs:
131 + if arg.annotation is None:
132 + continue
133 + annotation = get_ast_full_name(arg.annotation)
134 + if annotation in ("tl.constexpr", "triton.language.constexpr", "constexpr"):
135 + names.add(arg.arg)
136 + return frozenset(names)
137 +
138 +
139 +def _leftmost_name(node: ast.AST) -> str | None:
140 + if isinstance(node, ast.Name):
141 + return node.id
142 + if isinstance(node, ast.BinOp):
143 + return _leftmost_name(node.left)
144 + return None
145 +
146 +
147 +def _pointer_arg_names(
148 + function_def: ast.FunctionDef,
149 + arg_names: tuple[str, ...],
150 +) -> frozenset[str]:
151 + candidate_names = set(arg_names)
152 + pointer_names = {name for name in arg_names if name.endswith("_ptr")}
153 + for node in ast.walk(function_def):
154 + if not isinstance(node, ast.Call):
155 + continue
156 + if get_ast_full_name(node.func) not in ("tl.load", "tl.store"):
157 + continue
158 + if not node.args:
159 + continue
160 + name = _leftmost_name(node.args[0])
161 + if name in candidate_names:
162 + pointer_names.add(name)
163 + return frozenset(pointer_names)
164 +
165 +
166 +def trace_triton_kernel_specialization_args(
167 + kernel: Callable[..., Any],
168 +) -> tuple[str, ...]:
169 + function_def = get_function_source_node(kernel)
170 + source_fn = getattr(kernel, "fn", kernel)
171 + arg_names = tuple(inspect.signature(source_fn).parameters)
172 + constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names)
173 + do_not_specialize_args = _triton_do_not_specialize_args(
174 + kernel, function_def, arg_names
175 + )
176 + pointer_args = _pointer_arg_names(function_def, arg_names)
177 +
178 + return tuple(
179 + name
180 + for name in arg_names
181 + if name in constexpr_args
182 + or (name not in pointer_args and name not in do_not_specialize_args)
183 + )
vllm/model_executor/warmup/kernel_warmup.py
+11 -2
@@ -17,6 +17,9 @@ from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
17 from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
18 deepseek_v4_mhc_warmup,
19 )
20 +from vllm.model_executor.warmup.fa4_cutedsl_warmup import (
21 + fa4_cutedsl_warmup,
22 +)
23 from vllm.model_executor.warmup.flashinfer_autotune_cache import (
24 resolve_flashinfer_autotune_file,
25 write_flashinfer_autotune_cache,
@@ -27,7 +30,7 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
30 )
31 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
32 from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
30 - sparse_mla_triton_warmup_if_needed,
33 + sparse_mla_triton_warmup,
34 )
35 from vllm.model_executor.warmup.v1_block_table_warmup import (
36 warm_v1_block_table_kernels,
@@ -94,7 +97,6 @@ def kernel_warmup(worker: "Worker"):
97 )
98
99 # Run next so input-prep kernels JIT against pristine runner state.
97 - sparse_mla_triton_warmup_if_needed(worker)
100 flashinfer_sparse_mla_decode_autotune_warmup(worker)
101 deepseek_v4_sparse_mla_attention_warmup(worker)
102
@@ -156,8 +158,15 @@ def kernel_warmup(worker: "Worker"):
158 )
159
160 if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
161 + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
162 + # to the shared JIT warmup infrastructure.
163 + # https://github.com/vllm-project/vllm/pull/47451
164 cutedsl_warmup()
165
166 + if worker.vllm_config.kernel_config.enable_jit_warmup:
167 + fa4_cutedsl_warmup(worker)
168 + sparse_mla_triton_warmup(worker)
169 +
170
171 def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
172 if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None:
vllm/model_executor/warmup/sparse_mla_triton_warmup.py
+27 -253
@@ -4,11 +4,10 @@
4
5 from typing import TYPE_CHECKING
6
7 -import torch
8 -
7 from vllm.logger import init_logger
8
9 if TYPE_CHECKING:
10 + from vllm.config import VllmConfig
11 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
12 from vllm.v1.worker.gpu_worker import Worker
13
@@ -31,47 +30,7 @@ _GENERIC_SPARSE_MLA_BACKENDS = frozenset(
30 )
31 _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})
32
34 -_SPARSE_PREFILL_METADATA_NUM_PREFILLS = (1, 2, 4, 8)
35 -_SPARSE_PREFILL_METADATA_NUM_DECODES = (0, 1, 2)
36 -_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS = (4, 128)
37 -_PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS = (2, 3)
38 -_PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS = (
39 - # query_slice_start offset, query_slice_stop offset
40 - (0, 0),
41 - (0, -1),
42 - (1, 0),
43 - (1, -1),
44 -)
45 -_COMBINE_TOPK_SWA_INPUT_VARIANTS = (
46 - # offset_topk, offset_query_and_seq, offset_gather
47 - (False, False, False),
48 - (False, True, False),
49 - (True, True, True),
50 -)
51 -_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES = (
52 - # compress_ratio, topk, topk_width, N
53 - (1, 0, 512, 512),
54 - (4, 512, 512, 512 * 4),
55 - # DSv4-Pro C4A traffic uses top-k 1024 with N=1024.
56 - (4, 1024, 1024, 1024),
57 - (128, 8192, 8192, 8192 * 128),
58 - # Real C128A traffic also specializes N=1 in one call path.
59 - (128, 8192, 8192, 1),
60 -)
61 -
62 -
63 -def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int:
64 - return max(0, min(num_tokens, max_tokens))
65 -
66 -
67 -def _next_power_of_2(x: int) -> int:
68 - return 1 << (x - 1).bit_length()
69 -
70 -
71 -def _hf_config_int(runner: "GPUModelRunner", name: str, default: int) -> int:
72 - model_config = getattr(runner.vllm_config, "model_config", None)
73 - hf_config = getattr(model_config, "hf_config", None)
74 - return int(getattr(hf_config, name, default) or default)
33 +_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})
34
35
36 def _attention_backend_name(backend: object) -> str | None:
@@ -96,242 +55,57 @@ def _has_attention_backend(
55 return False
56
57
99 -def _warm_sparse_swa_prefill_metadata_kernel(
100 - device: torch.device,
101 - window_size: int,
102 - prefill_tokens: int,
58 +def _compile_sparse_swa_prefill_metadata_kernel(
59 + vllm_config: "VllmConfig",
60 ) -> None:
61 from vllm.v1.attention.backends.mla.sparse_swa import (
105 - _compute_prefill_metadata_kernel,
62 + _COMPUTE_PREFILL_METADATA_KERNEL,
63 )
64
108 - for num_prefills in _SPARSE_PREFILL_METADATA_NUM_PREFILLS:
109 - for num_decodes in _SPARSE_PREFILL_METADATA_NUM_DECODES:
110 - query_lens = [1] * num_decodes
111 - query_lens += [prefill_tokens] * num_prefills
112 - query_start_locs = [0]
113 - for query_len in query_lens:
114 - query_start_locs.append(query_start_locs[-1] + query_len)
115 - query_start_loc = torch.tensor(
116 - query_start_locs,
117 - dtype=torch.int32,
118 - device=device,
119 - )
120 - seq_lens = torch.tensor(
121 - [1] * num_decodes + [window_size + q for q in query_lens[num_decodes:]],
122 - dtype=torch.int32,
123 - device=device,
124 - )
125 - prefill_gather_lens = torch.empty(
126 - num_prefills, dtype=torch.int32, device=device
127 - )
128 - _compute_prefill_metadata_kernel[(1,)](
129 - prefill_gather_lens,
130 - seq_lens,
131 - query_start_loc,
132 - num_prefills,
133 - num_decodes,
134 - window_size,
135 - BLOCK_SIZE=_next_power_of_2(num_prefills),
136 - )
65 + _COMPUTE_PREFILL_METADATA_KERNEL.warmup(vllm_config)
66
67
139 -def _warm_prefill_chunk_metadata_kernel(
140 - device: torch.device,
141 - compress_ratio: int,
142 - query_len: int,
68 +def _compile_prefill_chunk_metadata_kernel(
69 + vllm_config: "VllmConfig",
70 ) -> None:
144 - from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata
145 -
146 - num_reqs = 2
147 - query_start_loc_cpu = torch.arange(
148 - 0, (num_reqs + 1) * query_len, query_len, dtype=torch.int32
71 + from vllm.v1.attention.backends.mla.indexer import (
72 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL,
73 )
150 - query_start_loc = query_start_loc_cpu.to(device=device)
74
152 - uncompressed_seq_lens_cpu = torch.tensor(
153 - [
154 - compress_ratio * multiplier + query_len
155 - for multiplier in _PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS
156 - ],
157 - dtype=torch.int32,
158 - )
159 - compressed_seq_lens_cpu = uncompressed_seq_lens_cpu // compress_ratio
160 - uncompressed_seq_lens = uncompressed_seq_lens_cpu.to(device=device)
161 - compressed_seq_lens = compressed_seq_lens_cpu.to(device=device)
162 - block_table = torch.zeros(
163 - (num_reqs, int(compressed_seq_lens_cpu.max().item())),
164 - dtype=torch.int32,
165 - device=device,
166 - )
167 -
168 - offset_uncompressed_seq_lens = torch.empty(
169 - num_reqs + 1, dtype=torch.int32, device=device
170 - )[1:]
171 - offset_uncompressed_seq_lens.copy_(uncompressed_seq_lens)
172 - query_slices = tuple(
173 - slice(start, num_reqs * query_len + stop)
174 - for start, stop in _PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS
175 - )
176 - for warmup_uncompressed_seq_lens in (
177 - uncompressed_seq_lens,
178 - offset_uncompressed_seq_lens,
179 - ):
180 - for query_slice in query_slices:
181 - build_prefill_chunk_metadata(
182 - 0,
183 - num_reqs,
184 - query_start_loc,
185 - query_start_loc_cpu,
186 - warmup_uncompressed_seq_lens,
187 - compressed_seq_lens,
188 - compressed_seq_lens_cpu,
189 - block_table,
190 - compress_ratio,
191 - query_slice=query_slice,
192 - )
75 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL.warmup(vllm_config)
76
77
195 -def _warm_combine_topk_swa_indices_kernel(
196 - device: torch.device,
197 - num_tokens: int,
198 - window_size: int,
199 - compress_ratio: int,
200 - topk: int,
201 - topk_width: int,
202 - n: int,
78 +def _compile_combine_topk_swa_indices_kernel(
79 + vllm_config: "VllmConfig",
80 ) -> None:
204 - from vllm.models.deepseek_v4.common.ops.cache_utils import combine_topk_swa_indices
205 -
206 - if num_tokens <= 0:
207 - return
208 -
209 - def _make_topk_indices(*, offset: bool) -> torch.Tensor:
210 - if offset:
211 - topk_storage = torch.full(
212 - (num_tokens * topk_width + 1,),
213 - -1,
214 - dtype=torch.int32,
215 - device=device,
216 - )
217 - topk_indices = topk_storage[1:].reshape(num_tokens, topk_width)
218 - else:
219 - topk_indices = torch.full(
220 - (num_tokens, topk_width), -1, dtype=torch.int32, device=device
221 - )
222 - if topk > 0:
223 - topk_indices.copy_(
224 - torch.arange(num_tokens * topk_width, dtype=torch.int32, device=device)
225 - .reshape(num_tokens, topk_width)
226 - .remainder(topk_width)
227 - )
228 - return topk_indices
229 -
230 - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
231 - seq_lens = torch.tensor(
232 - [window_size + num_tokens], dtype=torch.int32, device=device
81 + from vllm.models.deepseek_v4.common.ops.cache_utils import (
82 + _COMBINE_TOPK_SWA_INDICES_KERNEL,
83 )
234 - gather_lens = torch.tensor(
235 - [min(window_size + num_tokens, window_size + num_tokens - 1)],
236 - dtype=torch.int32,
237 - device=device,
238 - )
239 - offset_query_start_loc = torch.empty(3, dtype=torch.int32, device=device)[1:]
240 - offset_query_start_loc.copy_(query_start_loc)
241 - offset_seq_lens = torch.empty(2, dtype=torch.int32, device=device)[1:]
242 - offset_seq_lens.copy_(seq_lens)
243 - offset_gather_lens = torch.empty(2, dtype=torch.int32, device=device)[1:]
244 - offset_gather_lens.copy_(gather_lens)
245 -
246 - for (
247 - offset_topk,
248 - offset_query_and_seq,
249 - offset_gather,
250 - ) in _COMBINE_TOPK_SWA_INPUT_VARIANTS:
251 - warmup_topk_indices = _make_topk_indices(offset=offset_topk)
252 - warmup_query_start_loc = (
253 - offset_query_start_loc if offset_query_and_seq else query_start_loc
254 - )
255 - warmup_seq_lens = offset_seq_lens if offset_query_and_seq else seq_lens
256 - warmup_gather_lens = offset_gather_lens if offset_gather else gather_lens
257 - n_values = (n,) if n == 1 else (n, n + 1)
258 - for m in (window_size + num_tokens, topk_width):
259 - for n_value in n_values:
260 - combine_topk_swa_indices(
261 - warmup_topk_indices,
262 - warmup_query_start_loc,
263 - warmup_seq_lens,
264 - warmup_gather_lens,
265 - window_size,
266 - compress_ratio,
267 - topk,
268 - M=m,
269 - N=n_value,
270 - )
271 -
272 -
273 -@torch.inference_mode()
274 -def sparse_mla_triton_warmup(
275 - runner: "GPUModelRunner",
276 - num_tokens: int,
277 - *,
278 - compress_ratios: tuple[int, ...],
279 - combine_topk_swa_cases: tuple[tuple[int, int, int, int], ...] = (),
280 -) -> None:
281 - device = getattr(runner, "device", torch.device("cuda"))
282 - window_size = _hf_config_int(runner, "sliding_window", 128)
283 -
284 - _warm_sparse_swa_prefill_metadata_kernel(device, window_size, num_tokens)
285 - for compress_ratio in compress_ratios:
286 - _warm_prefill_chunk_metadata_kernel(device, compress_ratio, num_tokens)
287 - for compress_ratio, topk, topk_width, n in combine_topk_swa_cases:
288 - _warm_combine_topk_swa_indices_kernel(
289 - device,
290 - num_tokens,
291 - window_size,
292 - compress_ratio,
293 - topk,
294 - topk_width,
295 - n,
296 - )
297 -
84
299 -def deepseek_v4_sparse_triton_warmup(
300 - runner: "GPUModelRunner",
301 - num_tokens: int,
302 -) -> None:
303 - sparse_mla_triton_warmup(
304 - runner,
305 - num_tokens,
306 - compress_ratios=_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS,
307 - combine_topk_swa_cases=_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES,
308 - )
85 + _COMBINE_TOPK_SWA_INDICES_KERNEL.warmup(vllm_config)
86
87
311 -def sparse_mla_triton_warmup_if_needed(worker: "Worker") -> None:
88 +def sparse_mla_triton_warmup(worker: "Worker") -> None:
89 runner = worker.model_runner
90 if runner.is_pooling_model:
91 return
92
93 max_tokens = worker.scheduler_config.max_num_batched_tokens
317 - num_tokens = _clamp_warmup_tokens(8, max_tokens)
318 - if num_tokens <= 0:
94 + max_num_prefills = min(worker.scheduler_config.max_num_seqs, max_tokens)
95 + if max_tokens <= 0 or max_num_prefills <= 0:
96 return
97
98 + vllm_config = runner.vllm_config
99 try:
100 if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS):
323 - deepseek_v4_sparse_triton_warmup(runner, num_tokens)
101 + _compile_sparse_swa_prefill_metadata_kernel(vllm_config)
102 + _compile_prefill_chunk_metadata_kernel(vllm_config)
103 + _compile_combine_topk_swa_indices_kernel(vllm_config)
104 elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS):
325 - sparse_mla_triton_warmup(
326 - runner,
327 - num_tokens,
328 - compress_ratios=(1,),
329 - )
105 + _compile_sparse_swa_prefill_metadata_kernel(vllm_config)
106 + _compile_prefill_chunk_metadata_kernel(vllm_config)
107 elif _has_attention_backend(runner, _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS):
331 - _warm_prefill_chunk_metadata_kernel(
332 - getattr(runner, "device", torch.device("cuda")),
333 - compress_ratio=1,
334 - query_len=num_tokens,
335 - )
108 + _compile_prefill_chunk_metadata_kernel(vllm_config)
109 +
110 except Exception:
111 logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True)
vllm/models/deepseek_v4/common/ops/cache_utils.py
+239 -71
@@ -14,14 +14,23 @@ preparation.
14 window indices for sparse prefill.
15 """
16
17 +from dataclasses import dataclass
18 +from typing import Any
19 +
20 import torch
21
22 from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 get_fp8_min_max,
24 )
25 +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs
26 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
27 + TritonPointerInputVariant,
28 + TritonWarmupTensor,
29 +)
30 from vllm.platforms import current_platform
31 from vllm.triton_utils import tl, triton
32 from vllm.utils.import_utils import has_cutedsl
33 +from vllm.utils.math_utils import next_power_of_2
34
35
36 @triton.jit
@@ -528,7 +537,6 @@ def combine_topk_swa_indices(
537 N: int,
538 ) -> tuple[torch.Tensor, torch.Tensor]:
539 num_tokens = topk_indices.shape[0]
531 - num_reqs = seq_lens.shape[0]
540 combined_topk = (
541 (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
542 // _SPARSE_PREFILL_TOPK_ALIGNMENT
@@ -544,13 +552,10 @@ def combine_topk_swa_indices(
552 num_tokens, dtype=torch.int32, device=topk_indices.device
553 )
554
547 - NUM_WORKERS = 128
548 - _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)](
555 + _COMBINE_TOPK_SWA_INDICES_KERNEL(
556 combined_indices,
550 - combined_indices.stride(0),
557 combined_lens,
558 topk_indices,
553 - topk_indices.stride(0),
559 query_start_loc,
560 seq_lens,
561 gather_lens,
@@ -559,82 +564,245 @@ def combine_topk_swa_indices(
564 TOP_K=topk,
565 COMPRESS_RATIO=compress_ratio,
566 WINDOW_SIZE=window_size,
562 - PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
567 )
568 return combined_indices, combined_lens
569
570
567 -@triton.jit
568 -def _combine_topk_swa_indices_kernel(
569 - combined_indices_ptr,
570 - combined_indices_stride,
571 - combined_lens_ptr,
572 - topk_indices_ptr,
573 - topk_indices_stride,
574 - query_start_loc_ptr,
575 - seq_lens_ptr,
576 - gather_lens_ptr,
577 - M,
578 - N,
579 - TOP_K: tl.constexpr,
580 - COMPRESS_RATIO: tl.constexpr,
581 - WINDOW_SIZE: tl.constexpr,
582 - PADDED_TOP_K: tl.constexpr,
571 +_COMBINE_TOPK_SWA_NUM_WORKERS = 128
572 +
573 +
574 +# Representative pointer alignment variants for Triton pointer specialization.
575 +_COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs(
576 + dict(
577 + topk_indices=True,
578 + query_start_loc=True,
579 + seq_lens=True,
580 + gather_lens=True,
581 + ),
582 + dict(
583 + topk_indices=True,
584 + query_start_loc=False,
585 + seq_lens=False,
586 + gather_lens=True,
587 + ),
588 + dict(
589 + topk_indices=False,
590 + query_start_loc=False,
591 + seq_lens=False,
592 + gather_lens=False,
593 + ),
594 +)
595 +
596 +
597 +_DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS = zip_inputs(
598 + # DSv4-Flash / SWA-only and C4A.
599 + dict(compress_ratio=1, topk=0, topk_width=512),
600 + dict(compress_ratio=4, topk=512, topk_width=512),
601 + # DSv4-Pro C4A.
602 + dict(compress_ratio=4, topk=1024, topk_width=1024),
603 + # DSv4-Pro C128A.
604 + dict(compress_ratio=128, topk=8192, topk_width=8192),
605 +)
606 +
607 +
608 +def _hf_config_int(vllm_config: Any, name: str, default: int) -> int:
609 + model_config = getattr(vllm_config, "model_config", None)
610 + hf_config = getattr(model_config, "hf_config", None)
611 + return int(getattr(hf_config, name, default) or default)
612 +
613 +
614 +def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int:
615 + scheduler_config = getattr(vllm_config, "scheduler_config", None)
616 + return int(getattr(scheduler_config, name, default) or default)
617 +
618 +
619 +class CombineTopkSwaIndicesKernel(
620 + VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"]
621 ):
584 - batch_idx = tl.program_id(0)
585 - worker_id = tl.program_id(1)
586 - num_workers = tl.num_programs(1)
622 + @dataclass(frozen=True)
623 + class CompileKey:
624 + TOP_K: int
625 + COMPRESS_RATIO: int
626 + WINDOW_SIZE: int
627 + PADDED_TOP_K: int
628 + input_variant: TritonPointerInputVariant
629 +
630 + @staticmethod
631 + @triton.jit(
632 + do_not_specialize=[
633 + "combined_indices_stride",
634 + "topk_indices_stride",
635 + "M",
636 + "N",
637 + ]
638 + )
639 + def kernel(
640 + combined_indices_ptr,
641 + combined_indices_stride,
642 + combined_lens_ptr,
643 + topk_indices_ptr,
644 + topk_indices_stride,
645 + query_start_loc_ptr,
646 + seq_lens_ptr,
647 + gather_lens_ptr,
648 + M,
649 + N,
650 + TOP_K: tl.constexpr,
651 + COMPRESS_RATIO: tl.constexpr,
652 + WINDOW_SIZE: tl.constexpr,
653 + PADDED_TOP_K: tl.constexpr,
654 + ):
655 + batch_idx = tl.program_id(0)
656 + worker_id = tl.program_id(1)
657 + num_workers = tl.num_programs(1)
658 +
659 + # query_start_loc is a global tensor; rebase to chunk-local offsets
660 + # by subtracting the chunk's starting value.
661 + base = tl.load(query_start_loc_ptr)
662 + query_start = tl.load(query_start_loc_ptr + batch_idx) - base
663 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
664 + query_len = query_end - query_start
665 + seq_len = tl.load(seq_lens_ptr + batch_idx)
666 + gather_len = tl.load(gather_lens_ptr + batch_idx)
667 + start_pos = seq_len - query_len
668 + # The SWA portion of the gathered buffer starts from position
669 + # (seq_len - gather_len), not position 0. We need this offset
670 + # to correctly index into the gathered buffer.
671 + gather_start = seq_len - gather_len
672 +
673 + for token_idx in range(query_start + worker_id, query_end, num_workers):
674 + # topk_len is fully determined by the query token's absolute position:
675 + # both the C4A indexer and the C128A metadata builder emit
676 + # min((pos + 1) // compress_ratio, topk_tokens) valid entries.
677 + # Caller passes TOP_K=0 for SWA-only layers to zero this out.
678 + token_idx_in_query = token_idx - query_start
679 + pos = start_pos + token_idx_in_query
680 + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
681 + swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
682 +
683 + offset = tl.arange(0, PADDED_TOP_K)
684 + mask = offset < topk_len
685 + topk_indices = tl.load(
686 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
687 + mask=mask,
688 + )
689 + tl.store(
690 + combined_indices_ptr + token_idx * combined_indices_stride + offset,
691 + topk_indices + M * batch_idx,
692 + mask=mask,
693 + )
694 + offset = tl.arange(0, WINDOW_SIZE)
695 + # Index into gathered buffer: N + (position - gather_start)
696 + # For positions [pos - swa_len + 1, pos], the buffer indices are:
697 + # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start]
698 + tl.store(
699 + combined_indices_ptr
700 + + token_idx * combined_indices_stride
701 + + topk_len
702 + + offset,
703 + M * batch_idx + N + offset + pos - swa_len + 1 - gather_start,
704 + mask=offset < swa_len,
705 + )
706
588 - # query_start_loc is a global tensor; rebase to chunk-local offsets
589 - # by subtracting the chunk's starting value.
590 - base = tl.load(query_start_loc_ptr)
591 - query_start = tl.load(query_start_loc_ptr + batch_idx) - base
592 - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
593 - query_len = query_end - query_start
594 - seq_len = tl.load(seq_lens_ptr + batch_idx)
595 - gather_len = tl.load(gather_lens_ptr + batch_idx)
596 - start_pos = seq_len - query_len
597 - # The SWA portion of the gathered buffer starts from position
598 - # (seq_len - gather_len), not position 0. We need this offset
599 - # to correctly index into the gathered buffer.
600 - gather_start = seq_len - gather_len
601 -
602 - for token_idx in range(query_start + worker_id, query_end, num_workers):
603 - # topk_len is fully determined by the query token's absolute position:
604 - # both the C4A indexer and the C128A metadata builder emit
605 - # min((pos + 1) // compress_ratio, topk_tokens) valid entries.
606 - # Caller passes TOP_K=0 for SWA-only layers to zero this out.
607 - token_idx_in_query = token_idx - query_start
608 - pos = start_pos + token_idx_in_query
609 - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
610 - swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
611 -
612 - offset = tl.arange(0, PADDED_TOP_K)
613 - mask = offset < topk_len
614 - topk_indices = tl.load(
615 - topk_indices_ptr + token_idx * topk_indices_stride + offset,
616 - mask=mask,
707 + combined_len = topk_len + swa_len
708 + tl.store(combined_lens_ptr + token_idx, combined_len)
709 +
710 + def dispatch( # type: ignore[override]
711 + self,
712 + *,
713 + topk_width: int,
714 + topk_indices: bool,
715 + query_start_loc: bool,
716 + seq_lens: bool,
717 + gather_lens: bool,
718 + topk: int,
719 + compress_ratio: int,
720 + WINDOW_SIZE: int,
721 + ) -> CompileKey:
722 + padded_topk = next_power_of_2(topk_width)
723 + input_variant = TritonPointerInputVariant.from_alignment(
724 + topk_indices=topk_indices,
725 + query_start_loc=query_start_loc,
726 + seq_lens=seq_lens,
727 + gather_lens=gather_lens,
728 )
618 - tl.store(
619 - combined_indices_ptr + token_idx * combined_indices_stride + offset,
620 - topk_indices + M * batch_idx,
621 - mask=mask,
729 + return self.CompileKey(
730 + TOP_K=topk,
731 + COMPRESS_RATIO=compress_ratio,
732 + WINDOW_SIZE=WINDOW_SIZE,
733 + PADDED_TOP_K=padded_topk,
734 + input_variant=input_variant,
735 )
623 - offset = tl.arange(0, WINDOW_SIZE)
624 - # Index into gathered buffer: N + (position - gather_start)
625 - # For positions [pos - swa_len + 1, pos], the buffer indices are:
626 - # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start]
627 - tl.store(
628 - combined_indices_ptr
629 - + token_idx * combined_indices_stride
630 - + topk_len
631 - + offset,
632 - M * batch_idx + N + offset + pos - swa_len + 1 - gather_start,
633 - mask=offset < swa_len,
736 +
737 + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
738 + if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0:
739 + return []
740 +
741 + window_size = _hf_config_int(vllm_config, "sliding_window", 128)
742 + return self._trace_dispatch(self.dispatch)(
743 + _DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS,
744 + _COMBINE_TOPK_SWA_POINTER_INPUTS,
745 + WINDOW_SIZE=window_size,
746 + )
747 +
748 + def compile(self, compile_key: CompileKey) -> None:
749 + warmup = getattr(self.kernel, "warmup", None)
750 + assert warmup is not None
751 + int32_ptr = TritonWarmupTensor(torch.int32)
752 + input_variant = compile_key.input_variant
753 + warmup(
754 + int32_ptr,
755 + 1, # do not specialize combined_indices_stride
756 + int32_ptr,
757 + input_variant.pointer("topk_indices", torch.int32),
758 + 1, # do not specialize topk_indices_stride
759 + input_variant.pointer("query_start_loc", torch.int32),
760 + input_variant.pointer("seq_lens", torch.int32),
761 + input_variant.pointer("gather_lens", torch.int32),
762 + 1, # do not specialize M
763 + 1, # do not specialize N
764 + TOP_K=compile_key.TOP_K,
765 + COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
766 + WINDOW_SIZE=compile_key.WINDOW_SIZE,
767 + PADDED_TOP_K=compile_key.PADDED_TOP_K,
768 + grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),
769 )
770
636 - combined_len = topk_len + swa_len
637 - tl.store(combined_lens_ptr + token_idx, combined_len)
771 + def __call__(
772 + self,
773 + combined_indices: torch.Tensor,
774 + combined_lens: torch.Tensor,
775 + topk_indices: torch.Tensor,
776 + query_start_loc: torch.Tensor,
777 + seq_lens: torch.Tensor,
778 + gather_lens: torch.Tensor,
779 + M: int,
780 + N: int,
781 + *,
782 + TOP_K: int,
783 + COMPRESS_RATIO: int,
784 + WINDOW_SIZE: int,
785 + ) -> None:
786 + num_reqs = seq_lens.shape[0]
787 + self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](
788 + combined_indices,
789 + combined_indices.stride(0),
790 + combined_lens,
791 + topk_indices,
792 + topk_indices.stride(0),
793 + query_start_loc,
794 + seq_lens,
795 + gather_lens,
796 + M,
797 + N,
798 + TOP_K=TOP_K,
799 + COMPRESS_RATIO=COMPRESS_RATIO,
800 + WINDOW_SIZE=WINDOW_SIZE,
801 + PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]),
802 + )
803 +
804 +
805 +_COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel()
806
807
808 def build_flashinfer_mixed_sparse_indices(
vllm/v1/attention/backends/mla/indexer.py
+215 -79
@@ -8,6 +8,14 @@ import vllm.envs as envs
8 from vllm.config import VllmConfig
9 from vllm.distributed import get_dcp_group, get_pcp_group
10 from vllm.logger import init_logger
11 +from vllm.model_executor.warmup.jit_warmup import (
12 + VllmJitKernel,
13 + WarmupIntRange,
14 +)
15 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
16 + TritonPointerInputVariant,
17 + TritonWarmupTensor,
18 +)
19 from vllm.platforms import current_platform
20 from vllm.triton_utils import tl, triton
21 from vllm.utils.deep_gemm import (
@@ -190,6 +198,211 @@ class DeepseekV32IndexerPrefillChunkMetadata:
198 max_local_total_seq_lens: int = 0
199
200
201 +_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS = (
202 + TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
203 + TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
204 +)
205 +
206 +
207 +class BuildPrefillChunkMetadataKernel(
208 + VllmJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
209 +):
210 + BLOCK_SIZE = 1024
211 +
212 + @dataclass(frozen=True)
213 + class CompileKey:
214 + query_slice_start: int
215 + query_slice_stop: int
216 + DCP_RANK: int
217 + DCP_WORLD: int
218 + DCP_INTERLEAVE: int
219 + BLOCK_SIZE: int
220 + COMPRESS_RATIO: int
221 + input_variant: TritonPointerInputVariant
222 +
223 + @staticmethod
224 + @triton.jit
225 + def kernel(
226 + # Inputs
227 + query_start_loc_ptr,
228 + uncompressed_seq_lens_ptr,
229 + cu_compressed_seq_lens_ptr,
230 + # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP,
231 + # aliases cu_compressed_seq_lens_ptr otherwise.
232 + row_start_cu_compressed_seq_lens_ptr,
233 + # Outputs
234 + token_to_seq_ptr,
235 + cu_compressed_seq_len_ks_ptr,
236 + cu_compressed_seq_len_ke_ptr,
237 + query_slice_start,
238 + query_slice_stop,
239 + DCP_RANK,
240 + DCP_WORLD,
241 + DCP_INTERLEAVE,
242 + BLOCK_SIZE: tl.constexpr,
243 + COMPRESS_RATIO: tl.constexpr,
244 + ):
245 + batch_idx = tl.program_id(0)
246 +
247 + query_start = tl.load(query_start_loc_ptr + batch_idx)
248 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
249 + query_len = query_end - query_start
250 +
251 + seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
252 + seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
253 + compressed_seq_len = seq_end - seq_start
254 +
255 + # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start
256 + # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr).
257 + row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx)
258 +
259 + uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
260 + start_pos = uncompressed_seq_len - query_len
261 +
262 + for i in range(0, query_len, BLOCK_SIZE):
263 + offset = i + tl.arange(0, BLOCK_SIZE)
264 + abs_pos = query_start + offset
265 + mask = (
266 + (offset < query_len)
267 + & (abs_pos >= query_slice_start)
268 + & (abs_pos < query_slice_stop)
269 + )
270 + out_pos = abs_pos - query_slice_start
271 +
272 + # cu_seq_len_ks: row start in the gathered K buffer.
273 + tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask)
274 +
275 + # cu_seq_len_ke: row start + per-token context length. Under DCP the
276 + # global per-token length is sharded across ranks.
277 + global_ctx = start_pos + 1 + offset
278 + len_per_token = global_ctx // COMPRESS_RATIO
279 + if DCP_WORLD > 1:
280 + # Per-rank local context length under interleave-aware DCP, matching
281 + # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
282 + base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE
283 + remainder = len_per_token - base * DCP_WORLD
284 + remainder = tl.minimum(
285 + tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE
286 + )
287 + len_per_token = base + remainder
288 + tl.store(
289 + cu_compressed_seq_len_ke_ptr + out_pos,
290 + row_start + len_per_token,
291 + mask=mask,
292 + )
293 +
294 + # Compute token_to_seq
295 + for i in range(0, compressed_seq_len, BLOCK_SIZE):
296 + offset = i + tl.arange(0, BLOCK_SIZE)
297 + mask = offset < compressed_seq_len
298 + tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
299 +
300 + def dispatch( # type: ignore[override]
301 + self,
302 + *,
303 + query_slice_start: int,
304 + query_slice_stop: int,
305 + DCP_RANK: int,
306 + DCP_WORLD: int,
307 + DCP_INTERLEAVE: int,
308 + BLOCK_SIZE: int,
309 + COMPRESS_RATIO: int,
310 + input_variant: TritonPointerInputVariant,
311 + ) -> CompileKey:
312 + return self.CompileKey(
313 + query_slice_start=query_slice_start,
314 + query_slice_stop=query_slice_stop,
315 + DCP_RANK=DCP_RANK,
316 + DCP_WORLD=DCP_WORLD,
317 + DCP_INTERLEAVE=DCP_INTERLEAVE,
318 + BLOCK_SIZE=BLOCK_SIZE,
319 + COMPRESS_RATIO=COMPRESS_RATIO,
320 + input_variant=input_variant,
321 + )
322 +
323 + def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
324 + max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8))
325 + hf_config = vllm_config.model_config.hf_config
326 + parallel_config = vllm_config.parallel_config
327 + dcp_world = parallel_config.decode_context_parallel_size
328 + dcp_interleave = parallel_config.cp_kv_cache_interleave_size
329 + dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
330 + compress_ratios = tuple(
331 + int(ratio)
332 + for ratio in (getattr(hf_config, "compress_ratios", None) or (1,))
333 + )
334 + return self._trace_dispatch(self.dispatch)(
335 + query_slice_start=WarmupIntRange(0, 2),
336 + query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
337 + DCP_RANK=dcp_rank,
338 + DCP_WORLD=dcp_world,
339 + DCP_INTERLEAVE=dcp_interleave,
340 + BLOCK_SIZE=self.BLOCK_SIZE,
341 + COMPRESS_RATIO=list(compress_ratios),
342 + input_variant=_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS,
343 + )
344 +
345 + def compile(self, compile_key: CompileKey) -> None:
346 + warmup = getattr(self.kernel, "warmup", None)
347 + assert warmup is not None
348 + int32_ptr = TritonWarmupTensor(torch.int32)
349 + warmup(
350 + int32_ptr,
351 + compile_key.input_variant.pointer("uncompressed_seq_lens", torch.int32),
352 + int32_ptr,
353 + int32_ptr,
354 + int32_ptr,
355 + int32_ptr,
356 + int32_ptr,
357 + compile_key.query_slice_start,
358 + compile_key.query_slice_stop,
359 + compile_key.DCP_RANK,
360 + compile_key.DCP_WORLD,
361 + compile_key.DCP_INTERLEAVE,
362 + BLOCK_SIZE=compile_key.BLOCK_SIZE,
363 + COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
364 + grid=(1,),
365 + )
366 +
367 + def __call__(
368 + self,
369 + query_start_loc: torch.Tensor,
370 + uncompressed_seq_lens: torch.Tensor,
371 + cu_compressed_seq_lens: torch.Tensor,
372 + row_start_cu_compressed_seq_lens: torch.Tensor,
373 + token_to_seq: torch.Tensor,
374 + cu_compressed_seq_len_ks: torch.Tensor,
375 + cu_compressed_seq_len_ke: torch.Tensor,
376 + query_slice_start: int,
377 + query_slice_stop: int,
378 + DCP_RANK: int,
379 + DCP_WORLD: int,
380 + DCP_INTERLEAVE: int,
381 + *,
382 + num_reqs: int,
383 + COMPRESS_RATIO: int,
384 + ) -> None:
385 + self.kernel[(num_reqs,)](
386 + query_start_loc,
387 + uncompressed_seq_lens,
388 + cu_compressed_seq_lens,
389 + row_start_cu_compressed_seq_lens,
390 + token_to_seq,
391 + cu_compressed_seq_len_ks,
392 + cu_compressed_seq_len_ke,
393 + query_slice_start,
394 + query_slice_stop,
395 + DCP_RANK,
396 + DCP_WORLD,
397 + DCP_INTERLEAVE,
398 + BLOCK_SIZE=self.BLOCK_SIZE,
399 + COMPRESS_RATIO=COMPRESS_RATIO,
400 + )
401 +
402 +
403 +_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
404 +
405 +
406 @dataclass
407 class DeepseekV32IndexerPrefillMetadata:
408 chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
@@ -841,7 +1054,7 @@ def build_prefill_chunk_metadata(
1054
1055 # Under DCP the kernel writes this rank's local row bounds into
1056 # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
844 - _build_prefill_chunk_metadata_kernel[(num_reqs,)](
1057 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL(
1058 query_start_loc,
1059 uncompressed_seq_lens[start_idx:end_idx],
1060 cu_seq_lens,
@@ -854,7 +1067,7 @@ def build_prefill_chunk_metadata(
1067 dcp_rank,
1068 dcp_world_size,
1069 cp_kv_cache_interleave_size,
857 - BLOCK_SIZE=1024,
1070 + num_reqs=num_reqs,
1071 COMPRESS_RATIO=compress_ratio,
1072 )
1073
@@ -881,80 +1094,3 @@ def build_prefill_chunk_metadata(
1094 local_total_seq_lens=local_total_seq_lens,
1095 max_local_total_seq_lens=max_local_total_seq_lens,
1096 )
884 -
885 -
886 -@triton.jit
887 -def _build_prefill_chunk_metadata_kernel(
888 - # Inputs
889 - query_start_loc_ptr,
890 - uncompressed_seq_lens_ptr,
891 - cu_compressed_seq_lens_ptr,
892 - # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP,
893 - # aliases cu_compressed_seq_lens_ptr otherwise.
894 - row_start_cu_compressed_seq_lens_ptr,
895 - # Outputs
896 - token_to_seq_ptr,
897 - cu_compressed_seq_len_ks_ptr,
898 - cu_compressed_seq_len_ke_ptr,
899 - query_slice_start,
900 - query_slice_stop,
901 - DCP_RANK,
902 - DCP_WORLD,
903 - DCP_INTERLEAVE,
904 - BLOCK_SIZE: tl.constexpr,
905 - COMPRESS_RATIO: tl.constexpr,
906 -):
907 - batch_idx = tl.program_id(0)
908 -
909 - query_start = tl.load(query_start_loc_ptr + batch_idx)
910 - query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
911 - query_len = query_end - query_start
912 -
913 - seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
914 - seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
915 - compressed_seq_len = seq_end - seq_start
916 -
917 - # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start
918 - # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr).
919 - row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx)
920 -
921 - uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
922 - start_pos = uncompressed_seq_len - query_len
923 -
924 - for i in range(0, query_len, BLOCK_SIZE):
925 - offset = i + tl.arange(0, BLOCK_SIZE)
926 - abs_pos = query_start + offset
927 - mask = (
928 - (offset < query_len)
929 - & (abs_pos >= query_slice_start)
930 - & (abs_pos < query_slice_stop)
931 - )
932 - out_pos = abs_pos - query_slice_start
933 -
934 - # cu_seq_len_ks: row start in the gathered K buffer.
935 - tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask)
936 -
937 - # cu_seq_len_ke: row start + per-token context length. Under DCP the
938 - # global per-token length is sharded across ranks.
939 - global_ctx = start_pos + 1 + offset
940 - len_per_token = global_ctx // COMPRESS_RATIO
941 - if DCP_WORLD > 1:
942 - # Per-rank local context length under interleave-aware DCP, matching
943 - # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
944 - base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE
945 - remainder = len_per_token - base * DCP_WORLD
946 - remainder = tl.minimum(
947 - tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE
948 - )
949 - len_per_token = base + remainder
950 - tl.store(
951 - cu_compressed_seq_len_ke_ptr + out_pos,
952 - row_start + len_per_token,
953 - mask=mask,
954 - )
955 -
956 - # Compute token_to_seq
957 - for i in range(0, compressed_seq_len, BLOCK_SIZE):
958 - offset = i + tl.arange(0, BLOCK_SIZE)
959 - mask = offset < compressed_seq_len
960 - tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+246 -50
@@ -3,7 +3,9 @@
3 """FlashAttention backend for MLA prefill."""
4
5 import functools
6 -from typing import TYPE_CHECKING
6 +from collections.abc import Callable
7 +from dataclasses import dataclass
8 +from typing import TYPE_CHECKING, Any
9
10 import torch
11
@@ -11,14 +13,7 @@ import vllm.envs as envs
13 from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 kFp8StaticTensorSym,
15 )
14 -from vllm.model_executor.warmup.cutedsl_warmup import (
15 - CuTeDSLCompileUnit,
16 - register_cutedsl_warmup_provider,
17 -)
18 -from vllm.model_executor.warmup.fa4_cutedsl_config import (
19 - FA4MLAPrefillCompileContext,
20 - iter_fa4_mla_prefill_compile_requests,
21 -)
16 +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
17 from vllm.platforms import current_platform
18 from vllm.v1.attention.backends.fa_utils import (
19 compile_flash_attn_varlen_func_from_specs,
@@ -32,6 +27,7 @@ from vllm.v1.attention.backends.mla.prefill.base import (
27
28 if TYPE_CHECKING:
29 from vllm.config import VllmConfig
30 + from vllm.model_executor.layers.attention.mla_attention import MLADims
31 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
32
33 if is_flash_attn_varlen_func_available():
@@ -40,6 +36,245 @@ else:
36 flash_attn_varlen_func = None # type: ignore[assignment]
37
38
39 +FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16)
40 +
41 +# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan
42 +# covers qhead_per_kvhead=1. Batch is not a current FA4 MLA-prefill key field.
43 +# Use b1 for compile-only specs because it is the conservative case for
44 +# Split-KV shape heuristics.
45 +# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but
46 +# vLLM does not use them in this backend yet; they need a separate
47 +# num_kv_heads/qv/top-k-aware warmup plan if wired in later.
48 +FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1
49 +FA4_MLA_PREFILL_Q_TILE = 128
50 +FA4_MLA_PREFILL_K_TILE = 128
51 +FA4_MLA_PREFILL_LONG_K_BLOCKS = 32
52 +FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64
53 +FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True)
54 +FA4_MLA_PREFILL_LSE_OPTIONS = (False, True)
55 +
56 +
57 +@dataclass(frozen=True)
58 +class _FA4MLAPrefillShapeProbe:
59 + max_seqlen_q: int
60 + max_seqlen_k: int
61 +
62 +
63 +class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
64 + """FA4 MLA prefill compile-key wrapper used by generic JIT warmup."""
65 +
66 + @dataclass(frozen=True)
67 + class CompileKey:
68 + """High-level FA4 compile-only key used by vLLM warmup.
69 +
70 + This is not the CuTeDSL cache key. FA4 owns the selector that maps these
71 + serving inputs to the actual compile-static fields: tile sizes, q_stage,
72 + Split-KV, scheduler choice, layout-presence booleans, dtype/head dims,
73 + arch, and related fields.
74 + """
75 +
76 + q_shape: tuple[int, ...]
77 + k_shape: tuple[int, ...]
78 + v_shape: tuple[int, ...]
79 + q_dtype: torch.dtype
80 + max_seqlen_q: int
81 + max_seqlen_k: int
82 + softmax_scale: float
83 + causal: bool
84 + fa_version: int
85 + v_stride: tuple[int, ...] | None = None
86 + cu_seqlens_q_shape: tuple[int, ...] | None = None
87 + cu_seqlens_k_shape: tuple[int, ...] | None = None
88 + window_size: tuple[int, int] | None = None
89 + return_softmax_lse: bool = False
90 + num_splits: int = 0
91 +
92 + def dispatch( # type: ignore[override]
93 + self,
94 + *,
95 + batch_size: int,
96 + dtype: torch.dtype,
97 + num_heads: int,
98 + mla_dims: "MLADims",
99 + shape_probe: _FA4MLAPrefillShapeProbe,
100 + requires_v_padding: bool,
101 + causal: bool,
102 + fa_version: int,
103 + window_size: tuple[int, int] | None = None,
104 + return_lse: bool = False,
105 + num_splits: int = 0,
106 + ) -> CompileKey:
107 + return self.CompileKey(
108 + q_shape=(
109 + batch_size * shape_probe.max_seqlen_q,
110 + num_heads,
111 + mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
112 + ),
113 + k_shape=(
114 + batch_size * shape_probe.max_seqlen_k,
115 + num_heads,
116 + mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
117 + ),
118 + v_shape=(
119 + batch_size * shape_probe.max_seqlen_k,
120 + num_heads,
121 + (
122 + mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
123 + if requires_v_padding
124 + else mla_dims.v_head_dim
125 + ),
126 + ),
127 + q_dtype=dtype,
128 + max_seqlen_q=shape_probe.max_seqlen_q,
129 + max_seqlen_k=shape_probe.max_seqlen_k,
130 + softmax_scale=(mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim)
131 + ** -0.5,
132 + causal=causal,
133 + fa_version=fa_version,
134 + v_stride=(
135 + None
136 + if requires_v_padding
137 + else (
138 + num_heads * (mla_dims.qk_nope_head_dim + mla_dims.v_head_dim),
139 + mla_dims.qk_nope_head_dim + mla_dims.v_head_dim,
140 + 1,
141 + )
142 + ),
143 + cu_seqlens_q_shape=(batch_size + 1,),
144 + cu_seqlens_k_shape=(batch_size + 1,),
145 + window_size=window_size,
146 + return_softmax_lse=return_lse,
147 + num_splits=num_splits,
148 + )
149 +
150 + def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]:
151 + from vllm.model_executor.layers.attention.mla_attention import (
152 + get_mla_dims,
153 + )
154 +
155 + mla_dims = get_mla_dims(vllm_config.model_config)
156 + qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
157 + fa_version = get_flash_attn_version(head_size=qk_head_dim)
158 + if fa_version != 4:
159 + return []
160 +
161 + dtype = vllm_config.model_config.dtype
162 + if dtype not in FA4_STANDARD_DTYPES:
163 + dtype = torch.bfloat16
164 +
165 + is_sm90 = current_platform.is_device_capability(90)
166 + is_sm100_family = current_platform.is_device_capability_family(100)
167 + is_sm120 = current_platform.is_device_capability_family(120)
168 + if not (is_sm90 or is_sm100_family or is_sm120):
169 + return []
170 +
171 + num_splits = 1 if envs.VLLM_BATCH_INVARIANT else 0
172 + if is_sm120 and num_splits != 1:
173 + return []
174 +
175 + num_heads = vllm_config.model_config.get_num_attention_heads(
176 + vllm_config.parallel_config
177 + )
178 + if num_heads <= 0:
179 + return []
180 +
181 + requires_v_padding = False
182 + effective_v_head_dim = (
183 + qk_head_dim if requires_v_padding else mla_dims.v_head_dim
184 + )
185 + shape_probes = [
186 + _FA4MLAPrefillShapeProbe(1, FA4_MLA_PREFILL_K_TILE),
187 + _FA4MLAPrefillShapeProbe(
188 + FA4_MLA_PREFILL_Q_TILE + 1,
189 + 4 * FA4_MLA_PREFILL_K_TILE,
190 + ),
191 + ]
192 + if num_splits != 1:
193 + shape_probes.extend(
194 + (
195 + _FA4MLAPrefillShapeProbe(
196 + 1,
197 + FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
198 + ),
199 + _FA4MLAPrefillShapeProbe(
200 + FA4_MLA_PREFILL_Q_TILE + 1,
201 + FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
202 + ),
203 + )
204 + )
205 + if not is_sm90 and qk_head_dim != effective_v_head_dim:
206 + shape_probes.extend(
207 + (
208 + _FA4MLAPrefillShapeProbe(
209 + 1,
210 + FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
211 + ),
212 + _FA4MLAPrefillShapeProbe(
213 + FA4_MLA_PREFILL_Q_TILE + 1,
214 + FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
215 + ),
216 + )
217 + )
218 +
219 + return self._trace_dispatch(self.dispatch)(
220 + batch_size=FA4_MLA_PREFILL_COMPILE_BATCH_SIZE,
221 + dtype=dtype,
222 + num_heads=num_heads,
223 + mla_dims=mla_dims,
224 + shape_probe=shape_probes,
225 + requires_v_padding=requires_v_padding,
226 + causal=FA4_MLA_PREFILL_CAUSAL_OPTIONS,
227 + return_lse=FA4_MLA_PREFILL_LSE_OPTIONS,
228 + num_splits=num_splits,
229 + fa_version=fa_version,
230 + )
231 +
232 + @staticmethod
233 + def kernel(*args: Any, **kwargs: Any) -> Any:
234 + assert flash_attn_varlen_func is not None
235 + return flash_attn_varlen_func(*args, **kwargs)
236 +
237 + def compile(self, compile_key: CompileKey) -> None:
238 + assert compile_flash_attn_varlen_func_from_specs is not None
239 + window_size = (
240 + list(compile_key.window_size)
241 + if compile_key.window_size is not None
242 + else None
243 + )
244 + compile_flash_attn_varlen_func_from_specs(
245 + q_shape=compile_key.q_shape,
246 + k_shape=compile_key.k_shape,
247 + v_shape=compile_key.v_shape,
248 + q_dtype=compile_key.q_dtype,
249 + v_stride=compile_key.v_stride,
250 + cu_seqlens_q_shape=compile_key.cu_seqlens_q_shape,
251 + cu_seqlens_k_shape=compile_key.cu_seqlens_k_shape,
252 + max_seqlen_q=compile_key.max_seqlen_q,
253 + max_seqlen_k=compile_key.max_seqlen_k,
254 + softmax_scale=compile_key.softmax_scale,
255 + causal=compile_key.causal,
256 + window_size=window_size,
257 + return_softmax_lse=compile_key.return_softmax_lse,
258 + fa_version=compile_key.fa_version,
259 + num_splits=compile_key.num_splits,
260 + )
261 +
262 + def __call__(
263 + self,
264 + *,
265 + q: torch.Tensor,
266 + k: torch.Tensor,
267 + v: torch.Tensor,
268 + runtime_kernel: Callable[..., Any] | None = None,
269 + **kwargs: Any,
270 + ) -> Any:
271 + kernel = self.kernel if runtime_kernel is None else runtime_kernel
272 + return kernel(q=q, k=k, v=v, **kwargs)
273 +
274 +
275 +FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()
276 +
277 +
278 class FlashAttnPrefillBackend(MLAPrefillBackend):
279 """FlashAttention backend for MLA prefill."""
280
@@ -127,46 +362,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
362
363 # Track whether we're using vllm's FA or upstream (for ROCm)
364 self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu()
130 - if self.vllm_flash_attn_version == 4:
131 - register_cutedsl_warmup_provider(self)
132 -
133 - def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
134 - if self.vllm_flash_attn_version != 4:
135 - return ()
136 - if compile_flash_attn_varlen_func_from_specs is None:
137 - raise RuntimeError(
138 - "FA4 compile-only API is unavailable; CuTeDSL warmup does not "
139 - "fall back to synthetic forward passes."
140 - )
141 -
142 - dtype = self.vllm_config.model_config.dtype
143 - if dtype not in self.supported_dtypes:
144 - dtype = torch.bfloat16
145 -
146 - qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
147 - ctx = FA4MLAPrefillCompileContext(
148 - dtype=dtype,
149 - num_heads=self.num_heads,
150 - qk_head_dim=qk_head_dim,
151 - v_head_dim=self.v_head_dim,
152 - kv_nope_head_dim=self.qk_nope_head_dim + self.v_head_dim,
153 - requires_v_padding=self.requires_v_padding,
154 - scale=self.scale,
155 - num_splits=1 if envs.VLLM_BATCH_INVARIANT else 0,
156 - fa_version=self.vllm_flash_attn_version,
157 - )
158 - compile_requests = tuple(iter_fa4_mla_prefill_compile_requests(ctx))
159 - if not compile_requests:
160 - return ()
161 -
162 - return tuple(
163 - CuTeDSLCompileUnit(
164 - name="fa4_mla_prefill",
165 - key=request.key,
166 - compile=request.compile,
167 - )
168 - for request in compile_requests
169 - )
365
366 def supports_quant_output(self, quant_key: "QuantKey") -> bool:
367 device_capability = current_platform.get_device_capability()
@@ -207,10 +402,11 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
402 if envs.VLLM_BATCH_INVARIANT:
403 kwargs["num_splits"] = 1
404
210 - attn_out = self.flash_attn_varlen_func(
405 + attn_out = FA4_MLA_PREFILL_KERNEL(
406 q=q,
407 k=k,
408 v=maybe_padded_v,
409 + runtime_kernel=self.flash_attn_varlen_func,
410 softmax_scale=softmax_scale,
411 **kwargs,
412 )
vllm/v1/attention/backends/mla/sparse_swa.py
+108 -34
@@ -7,9 +7,14 @@ import torch
7
8 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
9 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10 +from vllm.model_executor.warmup.jit_warmup import (
11 + VllmJitKernel,
12 + WarmupIntRange,
13 +)
14 +from vllm.model_executor.warmup.jit_warmup_triton_helper import TritonWarmupTensor
15 from vllm.platforms import current_platform
16 from vllm.triton_utils import tl, triton
12 -from vllm.utils.math_utils import cdiv
17 +from vllm.utils.math_utils import cdiv, next_power_of_2
18 from vllm.v1.attention.backend import (
19 AttentionBackend,
20 AttentionCGSupport,
@@ -279,6 +284,107 @@ class DeepseekSparseSWAMetadata:
284 return chunk_plan
285
286
287 +class ComputePrefillMetadataKernel(
288 + VllmJitKernel["ComputePrefillMetadataKernel.CompileKey"]
289 +):
290 + @dataclass(frozen=True)
291 + class CompileKey:
292 + BLOCK_SIZE: int
293 +
294 + @staticmethod
295 + @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])
296 + def kernel(
297 + # Outputs
298 + prefill_gather_lens_ptr,
299 + # Inputs
300 + seq_lens_ptr,
301 + query_start_loc_ptr,
302 + num_prefills,
303 + num_decodes,
304 + window_size,
305 + BLOCK_SIZE: tl.constexpr,
306 + ):
307 + """Compute prefill gather_lens in a single pass."""
308 + offset = tl.arange(0, BLOCK_SIZE)
309 + mask = offset < num_prefills
310 + # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
311 + # masked-off lanes even though the load mask gates the actual read, so
312 + # clamp the offset. Caller guarantees num_prefills > 0.
313 + safe_offset = tl.minimum(offset, num_prefills - 1)
314 +
315 + seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
316 + qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
317 + qsl_end = tl.load(
318 + query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask
319 + )
320 +
321 + query_len = qsl_end - qsl_start
322 + prefix_len = seq_len - query_len
323 + gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
324 +
325 + tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
326 +
327 + def dispatch( # type: ignore[override]
328 + self,
329 + *,
330 + num_prefills: int,
331 + ) -> CompileKey:
332 + return self.CompileKey(
333 + BLOCK_SIZE=next_power_of_2(num_prefills),
334 + )
335 +
336 + def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
337 + scheduler_config = vllm_config.scheduler_config
338 + max_prefills = max(
339 + 1,
340 + min(
341 + scheduler_config.max_num_seqs,
342 + scheduler_config.max_num_batched_tokens,
343 + ),
344 + )
345 + return self._trace_dispatch(self.dispatch)(
346 + num_prefills=WarmupIntRange(1, max_prefills + 1),
347 + )
348 +
349 + def compile(self, compile_key: CompileKey) -> None:
350 + warmup = getattr(self.kernel, "warmup", None)
351 + assert warmup is not None
352 + int32_ptr = TritonWarmupTensor(torch.int32)
353 + warmup(
354 + int32_ptr,
355 + int32_ptr,
356 + int32_ptr,
357 + compile_key.BLOCK_SIZE,
358 + 0,
359 + 1,
360 + BLOCK_SIZE=compile_key.BLOCK_SIZE,
361 + grid=(1,),
362 + )
363 +
364 + def __call__(
365 + self,
366 + prefill_gather_lens: torch.Tensor,
367 + seq_lens: torch.Tensor,
368 + query_start_loc: torch.Tensor,
369 + num_prefills: int,
370 + num_decodes: int,
371 + window_size: int,
372 + ) -> None:
373 + compile_key = self.dispatch(num_prefills=num_prefills)
374 + self.kernel[(1,)](
375 + prefill_gather_lens,
376 + seq_lens,
377 + query_start_loc,
378 + num_prefills,
379 + num_decodes,
380 + window_size,
381 + BLOCK_SIZE=compile_key.BLOCK_SIZE,
382 + )
383 +
384 +
385 +_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
386 +
387 +
388 class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
389 """Builds metadata for DeepseekV4 SWA cache.
390
@@ -602,14 +708,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
708 pfx_gather_lens = torch.empty(
709 num_prefills, dtype=torch.int32, device=seq_lens.device
710 )
605 - _compute_prefill_metadata_kernel[(1,)](
711 + _COMPUTE_PREFILL_METADATA_KERNEL(
712 pfx_gather_lens,
713 seq_lens,
714 query_start_loc,
715 num_prefills,
716 num_decodes,
717 self.window_size,
612 - BLOCK_SIZE=triton.next_power_of_2(num_prefills),
718 )
719
720 result["prefill_seq_lens"] = seq_lens[num_decodes:]
@@ -626,37 +731,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
731 return result
732
733
629 -@triton.jit
630 -def _compute_prefill_metadata_kernel(
631 - # Outputs
632 - prefill_gather_lens_ptr,
633 - # Inputs
634 - seq_lens_ptr,
635 - query_start_loc_ptr,
636 - num_prefills,
637 - num_decodes,
638 - window_size,
639 - BLOCK_SIZE: tl.constexpr,
640 -):
641 - """Compute prefill gather_lens in a single pass."""
642 - offset = tl.arange(0, BLOCK_SIZE)
643 - mask = offset < num_prefills
644 - # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
645 - # masked-off lanes even though the load mask gates the actual read, so
646 - # clamp the offset. Caller guarantees num_prefills > 0.
647 - safe_offset = tl.minimum(offset, num_prefills - 1)
648 -
649 - seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
650 - qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
651 - qsl_end = tl.load(query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask)
652 -
653 - query_len = qsl_end - qsl_start
654 - prefix_len = seq_len - query_len
655 - gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
656 -
657 - tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
658 -
659 -
734 @triton.jit(do_not_specialize=["token_offset"])
735 def _compute_swa_indices_and_lens_kernel(
736 swa_indices_ptr,