@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6c95a641e9

[2/N][Feat][Perf] Add new warmup infrastructure for JITs. Add predicate filtering for JIT warmup, and migrate Inkling FA4 (#49315)

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

Roberto L. Castro committed Aug 11, 2026 at 13:46 UTC 6c95a641e95c0faa6f3aa802d1fd3cce3f3bc3ce
12 files changed +743 -390
tests/model_executor/test_jit_warmup.py
+65
@@ -105,6 +105,39 @@ def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> No
105 ]
106
107
108 +def test_warmup_range_uses_custom_advancement() -> None:
109 + keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
110 + tokens=WarmupIntRange(
111 + 1,
112 + 10,
113 + advance=lambda value: _next_power_of_2(value) + 1,
114 + ),
115 + cfg=_config(),
116 + )
117 +
118 + assert [key.block_size for key in keys] == [1, 2, 4, 8, 16]
119 +
120 +
121 +def test_warmup_range_validates_custom_advancement() -> None:
122 + trace_dispatch = ToyKernel()._trace_dispatch(ToyKernel().dispatch)
123 +
124 + with pytest.raises(ValueError, match="both step and advance"):
125 + trace_dispatch(
126 + tokens=WarmupIntRange(
127 + 1,
128 + 10,
129 + step=2,
130 + advance=lambda value: value + 1,
131 + ),
132 + cfg=_config(),
133 + )
134 + with pytest.raises(ValueError, match="must return a greater value"):
135 + trace_dispatch(
136 + tokens=WarmupIntRange(1, 10, advance=lambda value: value),
137 + cfg=_config(),
138 + )
139 +
140 +
141 def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
142 cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
143
@@ -143,6 +176,38 @@ def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
176 ]
177
178
179 +def test_trace_dispatch_filters_with_traced_predicate() -> None:
180 + class PredicateKernel(ToyKernel):
181 + def _is_valid_warmup_dispatch(
182 + self,
183 + *,
184 + tokens: int,
185 + lanes: int,
186 + max_work: int,
187 + ) -> bool:
188 + block_size = _next_power_of_2(tokens)
189 + return block_size * lanes <= max_work
190 +
191 + def get_warmup_keys(
192 + self, max_tokens: int, cfg: Any
193 + ) -> list[ToyKernel.CompileKey]:
194 + return self._trace_dispatch(self.dispatch)(
195 + tokens=WarmupIntRange(1, max_tokens + 1),
196 + lanes=(1, 2),
197 + cfg=cfg,
198 + max_work=4,
199 + _when=self._is_valid_warmup_dispatch,
200 + )
201 +
202 + assert PredicateKernel().get_warmup_keys(5, _config()) == [
203 + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
204 + ToyKernel.CompileKey(1, 2, 1, ("base", "default", -1, 1, 1), True),
205 + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
206 + ToyKernel.CompileKey(2, 4, 1, ("base", "default", -2, 2, 4), True),
207 + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
208 + ]
209 +
210 +
211 def test_zip_inputs_validates_input_rows() -> None:
212 with pytest.raises(ValueError, match="requires at least one"):
213 zip_inputs()
tests/models/inkling/test_fa4_rel_attention.py
+4 -4
@@ -2,8 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Correctness test for the Inkling FA4 relative-attention score-mod kernel.
4
5 -Checks ``inkling_fa4_rel_attention`` against a pure-PyTorch reference that
6 -implements the relative bias exactly as documented in the Inkling architecture
5 +Checks ``INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference
6 +that implements the relative bias exactly as documented in the Inkling architecture
7 guide::
8
9 logit(i, j, h) = (1 / head_dim) * dot(q[i, h], k[j, h]) + rel_bias(i, j, h)
@@ -23,9 +23,9 @@ from vllm.models.inkling.nvidia.attention import (
23 compute_log_scaling_tau,
24 )
25 from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
26 + INKLING_FA4_REL_ATTENTION_KERNEL,
27 _use_sheared_bias,
28 inkling_fa4_num_splits,
28 - inkling_fa4_rel_attention,
29 )
30 from vllm.platforms import current_platform
31 from vllm.platforms.interface import DeviceCapability
@@ -279,7 +279,7 @@ def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0
279 num_kv_heads=num_kv_heads,
280 max_kv_len=max(kv_lens),
281 )
282 - out = inkling_fa4_rel_attention(
282 + out = INKLING_FA4_REL_ATTENTION_KERNEL(
283 q,
284 key_cache,
285 value_cache,
tests/models/inkling/test_fa4_warmup.py
+84 -29
@@ -1,17 +1,47 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from types import SimpleNamespace
5 +
6 import torch
7
8 +import vllm.models.inkling.nvidia.ops.fa4_rel_attention as fa4_rel_attention
9 +from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig
10 from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
11 + InklingFA4RelAttentionKernel,
12 + _num_warps_bucket,
13 bucket_max_seqlen_q,
14 inkling_fa4_num_splits,
15 )
10 -from vllm.models.inkling.nvidia.ops.fa4_warmup import (
11 - InklingFA4WarmupConfig,
12 - _iter_compile_units,
13 - _num_warps_bucket,
14 -)
16 +
17 +
18 +def _vllm_config_from_reference_config(config: SimpleNamespace) -> SimpleNamespace:
19 + return SimpleNamespace(
20 + model_config=SimpleNamespace(
21 + hf_config=InklingMMConfig(
22 + text_config=InklingModelConfig(
23 + num_attention_heads=config.num_heads,
24 + num_key_value_heads=config.num_kv_heads,
25 + head_dim=config.head_dim,
26 + rel_extent=config.rel_extent,
27 + swa_num_attention_heads=config.num_heads,
28 + swa_num_key_value_heads=config.num_kv_heads,
29 + swa_head_dim=config.head_dim,
30 + sliding_window_size=config.rel_extent,
31 + )
32 + ),
33 + max_model_len=config.max_kv_len,
34 + dtype=config.dtype,
35 + ),
36 + cache_config=SimpleNamespace(
37 + cache_dtype="auto",
38 + block_size=config.block_size,
39 + ),
40 + scheduler_config=SimpleNamespace(
41 + max_num_seqs=config.max_num_reqs,
42 + max_num_batched_tokens=config.max_num_batched_tokens,
43 + ),
44 + )
45
46
47 def test_bucket_max_seqlen_q():
@@ -28,8 +58,13 @@ def test_bucket_max_seqlen_q():
58 ]
59
60
31 -def test_warmup_enumerates_every_runtime_compile_class():
32 - config = InklingFA4WarmupConfig(
61 +def test_warmup_enumerates_every_runtime_compile_class(monkeypatch):
62 + monkeypatch.setattr(
63 + fa4_rel_attention,
64 + "get_tensor_model_parallel_world_size",
65 + lambda: 1,
66 + )
67 + config = SimpleNamespace(
68 num_heads=16,
69 num_kv_heads=2,
70 head_dim=128,
@@ -43,27 +78,47 @@ def test_warmup_enumerates_every_runtime_compile_class():
78 max_num_reqs=64,
79 max_num_batched_tokens=192,
80 )
46 - warmed_keys = {unit.key[-1] for unit in _iter_compile_units(config)}
81 + vllm_config = _vllm_config_from_reference_config(config)
82 + kernel = InklingFA4RelAttentionKernel()
83 + warmed_keys = set(kernel.get_warmup_keys(vllm_config))
84 + expected_keys: set[InklingFA4RelAttentionKernel.CompileKey] = set()
85
48 - for query_len in range(1, config.max_num_batched_tokens + 1):
49 - max_seqlen_q = bucket_max_seqlen_q(query_len)
50 - max_num_reqs = min(
51 - config.max_num_reqs,
52 - config.max_num_batched_tokens - query_len + 1,
53 - )
54 - for num_reqs in range(1, max_num_reqs + 1):
55 - num_splits = inkling_fa4_num_splits(
56 - is_local=config.is_local,
57 - batch_size=num_reqs,
58 - max_query_len=max_seqlen_q,
59 - num_heads=config.num_heads,
60 - num_kv_heads=config.num_kv_heads,
61 - max_kv_len=config.max_kv_len,
86 + for is_local in (False, True):
87 + max_kv_len = config.rel_extent if is_local else config.max_kv_len
88 + window_size = (config.rel_extent - 1, 0) if is_local else (-1, -1)
89 + for query_len in range(1, config.max_num_batched_tokens + 1):
90 + max_seqlen_q = bucket_max_seqlen_q(query_len)
91 + max_num_reqs = min(
92 + config.max_num_reqs,
93 + config.max_num_batched_tokens - query_len + 1,
94 )
63 - runtime_key = (
64 - max_seqlen_q,
65 - num_splits,
66 - _num_warps_bucket(num_reqs) if num_splits > 1 else None,
67 - num_reqs > 1024,
68 - )
69 - assert runtime_key in warmed_keys
95 + for num_reqs in range(1, max_num_reqs + 1):
96 + num_splits = inkling_fa4_num_splits(
97 + is_local=is_local,
98 + batch_size=num_reqs,
99 + max_query_len=max_seqlen_q,
100 + num_heads=config.num_heads,
101 + num_kv_heads=config.num_kv_heads,
102 + max_kv_len=max_kv_len,
103 + )
104 + expected_keys.add(
105 + InklingFA4RelAttentionKernel.CompileKey(
106 + is_local=is_local,
107 + num_heads=config.num_heads,
108 + num_kv_heads=config.num_kv_heads,
109 + head_dim=config.head_dim,
110 + rel_extent=config.rel_extent,
111 + dtype=config.dtype,
112 + kv_dtype=config.kv_dtype,
113 + block_size=config.block_size,
114 + window_size=window_size,
115 + max_seqlen_q=max_seqlen_q,
116 + num_splits=num_splits,
117 + num_warps_bucket=(
118 + _num_warps_bucket(num_reqs) if num_splits > 1 else None
119 + ),
120 + large_num_reqs=num_reqs > 1024,
121 + )
122 + )
123 +
124 + assert warmed_keys == expected_keys
vllm/config/kernel.py
+1 -1
@@ -180,7 +180,7 @@ class KernelConfig:
180 # to the shared JIT warmup infrastructure.
181 # https://github.com/vllm-project/vllm/pull/47451
182 enable_cutedsl_warmup: bool = True
183 - """If True, run CuTeDSL compile warmup during kernel warmup."""
183 + """Deprecated: run legacy CuTeDSL warmup providers."""
184
185 enable_jit_warmup: bool = True
186 """If True, run JIT compile warmup during kernel warmup."""
vllm/model_executor/warmup/cutedsl_warmup.py
+1 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Run registered CuTeDSL warmup compile units."""
3 +"""Deprecated compatibility registry for legacy CuTeDSL warmups."""
4
5 # TODO(roberto): Remove this compatibility registry after registered CuTeDSL
6 # warmups are migrated to the shared JIT warmup infrastructure.
vllm/model_executor/warmup/fa4_cutedsl_warmup.py
+21 -2
@@ -1,6 +1,6 @@
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."""
3 +"""Warm up FA4 CuTeDSL kernels."""
4
5 from __future__ import annotations
6
@@ -12,7 +12,7 @@ if TYPE_CHECKING:
12 from vllm.v1.worker.gpu_worker import Worker
13
14
15 -def fa4_cutedsl_warmup(worker: Worker) -> None:
15 +def _warm_fa4_mla_prefill(worker: Worker) -> None:
16 runner = worker.model_runner
17 if runner.is_pooling_model:
18 return
@@ -32,3 +32,22 @@ def fa4_cutedsl_warmup(worker: Worker) -> None:
32 from vllm.v1.attention.backends.mla.prefill import flash_attn
33
34 flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config)
35 +
36 +
37 +def _warm_inkling_fa4_rel_attention(worker: Worker) -> None:
38 + from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig
39 + from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
40 + INKLING_FA4_REL_ATTENTION_KERNEL,
41 + )
42 +
43 + vllm_config = worker.vllm_config
44 + hf_config = vllm_config.model_config.hf_config
45 + if not isinstance(hf_config, (InklingMMConfig, InklingModelConfig)):
46 + return
47 +
48 + INKLING_FA4_REL_ATTENTION_KERNEL.warmup(vllm_config)
49 +
50 +
51 +def fa4_cutedsl_warmup(worker: Worker) -> None:
52 + _warm_fa4_mla_prefill(worker)
53 + _warm_inkling_fa4_rel_attention(worker)
vllm/model_executor/warmup/jit_warmup.py
+149 -52
@@ -28,13 +28,29 @@ CompileKeyT = TypeVar("CompileKeyT")
28
29 @dataclass(frozen=True)
30 class WarmupIntRange:
31 + """Expand integers with range semantics or a custom monotonic progression."""
32 +
33 start: int
34 stop: int
35 step: int = 1
36 + advance: Callable[[int], int] | None = None
37
38
39 WarmupValues = Any
40 CompileKeyDispatchFn = Callable[..., CompileKeyT]
41 +WarmupPredicateFn = Callable[..., bool]
42 +_LocalExprs = tuple[tuple[str, ast.AST], ...]
43 +
44 +
45 +def _eval_local_exprs(
46 + local_exprs: _LocalExprs,
47 + values: Mapping[str, Any],
48 + globals_: Mapping[str, Any],
49 +) -> dict[str, Any]:
50 + evaluated_values = dict(values)
51 + for name, expr in local_exprs:
52 + evaluated_values[name] = _eval_dispatch_expr(expr, evaluated_values, globals_)
53 + return evaluated_values
54
55
56 @dataclass(frozen=True)
@@ -46,7 +62,20 @@ class _WarmupInputRows:
62
63 def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:
64 if isinstance(values, WarmupIntRange):
49 - return tuple(range(values.start, values.stop, values.step))
65 + if values.advance is None:
66 + return tuple(range(values.start, values.stop, values.step))
67 + if values.step != 1:
68 + raise ValueError("WarmupIntRange cannot set both step and advance")
69 +
70 + expanded: list[int] = []
71 + value = values.start
72 + while value < values.stop:
73 + expanded.append(value)
74 + next_value = values.advance(value)
75 + if next_value <= value:
76 + raise ValueError("WarmupIntRange.advance must return a greater value")
77 + value = next_value
78 + return tuple(expanded)
79 if isinstance(values, (list, tuple)):
80 return tuple(values)
81 return (values,)
@@ -75,20 +104,6 @@ def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows:
104 return _WarmupInputRows(rows=tuple(input_rows))
105
106
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 -
107 def _expand_warmup_input_rows(
108 rows: tuple[Mapping[str, WarmupValues], ...],
109 input_names: frozenset[str],
@@ -117,7 +132,7 @@ def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]:
132
133 @dataclass(frozen=True)
134 class _CompileKeyDispatchTrace:
120 - local_exprs: tuple[tuple[str, ast.AST], ...]
135 + local_exprs: _LocalExprs
136 field_exprs: tuple[tuple[str, ast.AST], ...]
137 globals: Mapping[str, Any]
138 input_names: frozenset[str]
@@ -128,11 +143,9 @@ class _CompileKeyDispatchTrace:
143 compile_key_type: type[CompileKeyT],
144 kwargs: Mapping[str, Any],
145 ) -> 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 - )
146 + dispatch_values = _eval_local_exprs(
147 + self.local_exprs, {**self.defaults, **kwargs}, self.globals
148 + )
149 return compile_key_type(
150 **{
151 field: _eval_dispatch_expr(expr, dispatch_values, self.globals)
@@ -141,6 +154,23 @@ class _CompileKeyDispatchTrace:
154 )
155
156
157 +@dataclass(frozen=True)
158 +class _WarmupPredicateTrace:
159 + local_exprs: _LocalExprs
160 + return_expr: ast.AST
161 + globals: Mapping[str, Any]
162 + input_names: frozenset[str]
163 + defaults: Mapping[str, Any]
164 +
165 + def matches(self, kwargs: Mapping[str, Any]) -> bool:
166 + dispatch_values = _eval_local_exprs(
167 + self.local_exprs, {**self.defaults, **kwargs}, self.globals
168 + )
169 + return bool(
170 + _eval_dispatch_expr(self.return_expr, dispatch_values, self.globals)
171 + )
172 +
173 +
174 _BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
175 ast.Add: operator.add,
176 ast.Sub: operator.sub,
@@ -317,10 +347,10 @@ def _collect_input_names(
347 }
348
349
320 -def _collect_dispatch_body(
321 - fn: CompileKeyDispatchFn[Any],
350 +def _collect_expression_body(
351 + fn: Callable[..., Any],
352 function_def: ast.FunctionDef,
323 -) -> tuple[list[tuple[str, ast.AST]], ast.Call]:
353 +) -> tuple[list[tuple[str, ast.AST]], ast.AST]:
354 local_exprs: list[tuple[str, ast.AST]] = []
355 for statement in function_def.body:
356 if (
@@ -335,7 +365,8 @@ def _collect_dispatch_body(
365 statement.targets[0], ast.Name
366 ):
367 raise _dispatch_expr_error(
338 - statement, "Dispatch assignments must target one local name"
368 + statement,
369 + "Traced warmup helper assignments must target one local name",
370 )
371 local_exprs.append((statement.targets[0].id, statement.value))
372 continue
@@ -343,29 +374,45 @@ def _collect_dispatch_body(
374 if isinstance(statement, ast.AnnAssign):
375 if statement.value is None:
376 raise _dispatch_expr_error(
346 - statement, "Dispatch annotations must assign a value"
377 + statement, "Traced warmup helper annotations must assign a value"
378 )
379 if not isinstance(statement.target, ast.Name):
380 raise _dispatch_expr_error(
350 - statement, "Dispatch assignments must target one local name"
381 + statement,
382 + "Traced warmup helper assignments must target one local name",
383 )
384 local_exprs.append((statement.target.id, statement.value))
385 continue
386
355 - if isinstance(statement, ast.Return) and isinstance(statement.value, ast.Call):
387 + if isinstance(statement, ast.Return) and statement.value is not None:
388 return local_exprs, statement.value
389
390 if isinstance(statement, ast.Return):
391 raise _dispatch_expr_error(
360 - statement, "Dispatch must return one CompileKey(...) call"
392 + statement, "Traced warmup helper must return an expression"
393 )
394
395 raise _dispatch_expr_error(
396 statement,
365 - "Dispatch may only contain local assignments before CompileKey return",
397 + "Traced warmup helper may only contain local assignments before return",
398 )
399
368 - raise ValueError(f"Expected {fn.__name__} to return one CompileKey(...) call")
400 + raise ValueError(f"Expected {fn.__name__} to return an expression")
401 +
402 +
403 +def _function_trace_inputs(
404 + fn: Callable[..., Any],
405 +) -> tuple[dict[str, Any], set[str]]:
406 + signature = inspect.signature(fn)
407 + defaults = {
408 + name: parameter.default
409 + for name, parameter in signature.parameters.items()
410 + if parameter.default is not inspect.Parameter.empty
411 + }
412 + bound_self = getattr(fn, "__self__", None)
413 + if bound_self is not None:
414 + defaults["self"] = bound_self
415 + return defaults, set(signature.parameters)
416
417
418 def _trace_compile_key_dispatch(
@@ -375,21 +422,20 @@ def _trace_compile_key_dispatch(
422 globals_ = source_fn.__globals__
423 function_def = get_function_source_node(fn)
424
378 - local_exprs, return_call = _collect_dispatch_body(fn, function_def)
425 + local_exprs, return_expr = _collect_expression_body(fn, function_def)
426 + if not isinstance(return_expr, ast.Call):
427 + raise _dispatch_expr_error(
428 + return_expr,
429 + "Dispatch must return one CompileKey(...) call",
430 + )
431
432 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)
433 + defaults, candidate_names = _function_trace_inputs(fn)
434 input_names: set[str] = set()
435 local_names = {name for name, _ in local_exprs}
436 for _, expr in local_exprs:
437 input_names.update(_collect_input_names(expr, candidate_names))
392 - for keyword in return_call.keywords:
438 + for keyword in return_expr.keywords:
439 if keyword.arg is None:
440 raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey")
441 field_exprs.append((keyword.arg, keyword.value))
@@ -406,6 +452,30 @@ def _trace_compile_key_dispatch(
452 )
453
454
455 +def _trace_warmup_predicate(
456 + fn: WarmupPredicateFn,
457 +) -> _WarmupPredicateTrace:
458 + source_fn = getattr(fn, "__func__", fn)
459 + globals_ = source_fn.__globals__
460 + function_def = get_function_source_node(fn)
461 +
462 + local_exprs, return_expr = _collect_expression_body(fn, function_def)
463 + defaults, candidate_names = _function_trace_inputs(fn)
464 + input_names: set[str] = set()
465 + local_names = {name for name, _ in local_exprs}
466 + for _, expr in local_exprs:
467 + input_names.update(_collect_input_names(expr, candidate_names))
468 + input_names.update(_collect_input_names(return_expr, candidate_names, local_names))
469 +
470 + return _WarmupPredicateTrace(
471 + tuple(local_exprs),
472 + return_expr,
473 + globals_,
474 + frozenset(input_names),
475 + defaults,
476 + )
477 +
478 +
479 class VllmJitKernel(Generic[CompileKeyT], ABC):
480 """Kernel wrapper that owns dispatch, warmup keys, and compilation."""
481
@@ -424,6 +494,7 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
494
495 def traced(
496 *input_groups: _WarmupInputRows,
497 + _when: WarmupPredicateFn | None = None,
498 **kwargs: WarmupValues,
499 ) -> list[CompileKeyT]:
500 for group in input_groups:
@@ -432,24 +503,50 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
503 "_trace_dispatch positional arguments must be "
504 "zip_inputs(...) groups"
505 )
506 + predicate_trace = (
507 + _trace_warmup_predicate(_when) if _when is not None else None
508 + )
509 + input_names = compile_key_dispatch_trace.input_names
510 + if predicate_trace is not None:
511 + input_names = input_names | predicate_trace.input_names
512 expanded_input_groups = tuple(
436 - _expand_warmup_input_rows(
437 - group.rows, compile_key_dispatch_trace.input_names
438 - )
513 + _expand_warmup_input_rows(group.rows, input_names)
514 for group in input_groups
515 )
441 - expanded_kwargs = _expand_warmup_value_grid(
442 - kwargs, compile_key_dispatch_trace.input_names
516 + # Expand independent keyword inputs into cartesian-product axes.
517 + expanded_kwarg_axes = tuple(
518 + (name, _expand_warmup_values(value))
519 + for name, value in kwargs.items()
520 + if name in input_names
521 )
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)
522 + dispatch_value_axes = (
523 + *expanded_input_groups,
524 + *(values for _, values in expanded_kwarg_axes),
525 + )
526 + input_group_count = len(expanded_input_groups)
527 + kwarg_names = tuple(name for name, _ in expanded_kwarg_axes)
528 + compile_keys: dict[CompileKeyT, None] = {}
529 + for dispatch_value_set in itertools.product(*dispatch_value_axes):
530 + dispatch_values = _merge_warmup_kwargs(
531 + (
532 + *dispatch_value_set[:input_group_count],
533 + dict(
534 + zip(
535 + kwarg_names,
536 + dispatch_value_set[input_group_count:],
537 + )
538 + ),
539 )
450 - for dispatch_values in itertools.product(*dispatch_value_groups)
540 )
452 - )
541 + if predicate_trace is not None and not predicate_trace.matches(
542 + dispatch_values
543 + ):
544 + continue
545 + compile_key = compile_key_dispatch_trace.compile_key(
546 + self.CompileKey, dispatch_values
547 + )
548 + compile_keys[compile_key] = None
549 + return list(compile_keys)
550
551 return traced
552
vllm/models/inkling/nvidia/attention.py
+2 -22
@@ -35,11 +35,10 @@ from vllm.v1.kv_cache_interface import (
35 from ..configs import InklingModelConfig
36 from .layernorm import InklingRMSNorm
37 from .ops.fa4_rel_attention import (
38 + INKLING_FA4_REL_ATTENTION_KERNEL,
39 bucket_max_seqlen_q,
40 inkling_fa4_num_splits,
40 - inkling_fa4_rel_attention,
41 )
42 -from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup
42 from .ops.qkvr_prep import fused_qkvr_prep
43 from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata
44 from .short_conv import InklingShortConv
@@ -175,25 +174,6 @@ class InklingAttention(nn.Module, AttentionLayerBase):
174 compilation_config.static_forward_context[prefix] = self
175 self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache
176
178 - register_fa4_warmup(
179 - InklingFA4WarmupConfig(
180 - num_heads=self.num_heads,
181 - num_kv_heads=self.num_kv_heads,
182 - head_dim=self.head_dim,
183 - rel_extent=self.rel_extent,
184 - window_size=self.window_size,
185 - is_local=self.is_local,
186 - max_kv_len=self._max_kv_len,
187 - dtype=vllm_config.model_config.dtype,
188 - kv_dtype=self.kv_cache_torch_dtype,
189 - block_size=vllm_config.cache_config.block_size,
190 - max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
191 - max_num_batched_tokens=(
192 - vllm_config.scheduler_config.max_num_batched_tokens
193 - ),
194 - )
195 - )
196 -
177 def get_attn_backend(self) -> type[AttentionBackend]:
178 return FlashAttentionBackend
179
@@ -309,7 +289,7 @@ class InklingAttention(nn.Module, AttentionLayerBase):
289 num_kv_heads=self.num_kv_heads,
290 max_kv_len=self._max_kv_len,
291 )
312 - inkling_fa4_rel_attention(
292 + INKLING_FA4_REL_ATTENTION_KERNEL(
293 q[:nt],
294 key_cache,
295 value_cache,
vllm/models/inkling/nvidia/ops/__init__.py
+4 -2
@@ -15,11 +15,13 @@ from .sconv import fused_sconv
15 _LAZY_EXPORTS = {
16 "silu_and_mul_triton": "silu_and_mul",
17 "sink_silu_mul_epilogue": "silu_and_mul",
18 - "inkling_fa4_rel_attention": "fa4_rel_attention",
18 + "INKLING_FA4_REL_ATTENTION_KERNEL": "fa4_rel_attention",
19 }
20
21 if TYPE_CHECKING:
22 - from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401
22 + from .fa4_rel_attention import ( # noqa: F401
23 + INKLING_FA4_REL_ATTENTION_KERNEL,
24 + )
25 from .silu_and_mul import ( # noqa: F401
26 silu_and_mul_triton,
27 sink_silu_mul_epilogue,
vllm/models/inkling/nvidia/ops/fa4_rel_attention.py
+333 -67
@@ -3,12 +3,23 @@
3 from __future__ import annotations
4
5 from collections.abc import Callable
6 +from dataclasses import dataclass
7 from functools import cache
7 -from typing import Any
8 +from typing import TYPE_CHECKING, Any
9
10 import torch
11
12 +from vllm.distributed import get_tensor_model_parallel_world_size
13 +from vllm.model_executor.warmup.jit_warmup import (
14 + VllmJitKernel,
15 + WarmupIntRange,
16 + zip_inputs,
17 +)
18 from vllm.platforms import current_platform
19 +from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
20 +
21 +if TYPE_CHECKING:
22 + from vllm.config import VllmConfig
23
24
25 def bucket_max_seqlen_q(max_seqlen_q: int) -> int:
@@ -98,73 +109,328 @@ def inkling_fa4_num_splits(
109 )
110
111
101 -def inkling_fa4_rel_attention(
102 - q: torch.Tensor,
103 - key_cache: torch.Tensor,
104 - value_cache: torch.Tensor,
105 - *,
106 - block_table: torch.Tensor,
107 - cache_seqlens: torch.Tensor,
108 - cu_seqlens_q: torch.Tensor,
109 - max_seqlen_q: int,
110 - softmax_scale: float,
111 - causal: bool,
112 - window_size: tuple[int, int],
113 - rel_extent: int,
114 - rel_logits: torch.Tensor,
115 - num_splits: int = 32,
116 - out: torch.Tensor | None = None,
117 -) -> torch.Tensor:
118 - """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias.
119 -
120 - ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache``
121 - are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``;
122 - ``block_table`` is the per-request page table and ``cache_seqlens`` the
123 - per-request KV lengths (``seqused_k``). ``rel_logits`` is
124 - ``(num_tokens, num_heads, rel_extent)``.
125 -
126 - Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's
127 - sheared relative-bias layout.
128 - """
129 - # cute uses (None, None) to mean "no window".
130 - cute_window = (None, None) if window_size == (-1, -1) else window_size
131 -
132 - rel_logits = rel_logits.contiguous()
133 - flash_attn_varlen_func: Callable[..., Any]
134 - if _use_sheared_bias():
135 - from vllm.third_party.tml_fa4 import (
136 - flash_attn_varlen_func as tml_flash_attn_varlen_func,
112 +def _num_warps_bucket(num_reqs: int) -> int:
113 + num_warps = min((num_reqs + 30) // 31, 32)
114 + return 1 << (num_warps - 1).bit_length()
115 +
116 +
117 +class InklingFA4RelAttentionKernel(
118 + VllmJitKernel["InklingFA4RelAttentionKernel.CompileKey"]
119 +):
120 + @dataclass(frozen=True)
121 + class CompileKey:
122 + is_local: bool
123 + num_heads: int
124 + num_kv_heads: int
125 + head_dim: int
126 + rel_extent: int
127 + dtype: torch.dtype
128 + kv_dtype: torch.dtype
129 + block_size: int
130 + window_size: tuple[int, int]
131 + max_seqlen_q: int
132 + num_splits: int
133 + num_warps_bucket: int | None
134 + large_num_reqs: bool
135 +
136 + @staticmethod
137 + def kernel(
138 + q: torch.Tensor,
139 + key_cache: torch.Tensor,
140 + value_cache: torch.Tensor,
141 + *,
142 + block_table: torch.Tensor,
143 + cache_seqlens: torch.Tensor,
144 + cu_seqlens_q: torch.Tensor,
145 + max_seqlen_q: int,
146 + softmax_scale: float,
147 + causal: bool,
148 + window_size: tuple[int, int],
149 + rel_extent: int,
150 + rel_logits: torch.Tensor,
151 + num_splits: int = 32,
152 + out: torch.Tensor | None = None,
153 + ) -> torch.Tensor:
154 + """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias.
155 +
156 + ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache``
157 + are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``;
158 + ``block_table`` is the per-request page table and ``cache_seqlens`` the
159 + per-request KV lengths (``seqused_k``). ``rel_logits`` is
160 + ``(num_tokens, num_heads, rel_extent)``.
161 +
162 + Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's
163 + sheared relative-bias layout.
164 + """
165 + # cute uses (None, None) to mean "no window".
166 + cute_window = (None, None) if window_size == (-1, -1) else window_size
167 +
168 + rel_logits = rel_logits.contiguous()
169 + flash_attn_varlen_func: Callable[..., Any]
170 + if _use_sheared_bias():
171 + from vllm.third_party.tml_fa4 import (
172 + flash_attn_varlen_func as tml_flash_attn_varlen_func,
173 + )
174 +
175 + flash_attn_varlen_func = tml_flash_attn_varlen_func
176 + bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
177 + else:
178 + from vllm.vllm_flash_attn.cute import (
179 + flash_attn_varlen_func as cute_flash_attn_varlen_func,
180 + )
181 +
182 + flash_attn_varlen_func = cute_flash_attn_varlen_func
183 + bias_kwargs = {
184 + "score_mod": _get_score_mod(rel_extent),
185 + "aux_tensors": [rel_logits],
186 + }
187 +
188 + ret = flash_attn_varlen_func(
189 + q=q,
190 + k=key_cache,
191 + v=value_cache,
192 + cu_seqlens_q=cu_seqlens_q,
193 + seqused_k=cache_seqlens,
194 + max_seqlen_q=max_seqlen_q,
195 + page_table=block_table,
196 + softmax_scale=softmax_scale,
197 + causal=causal,
198 + window_size=cute_window,
199 + num_splits=num_splits,
200 + return_lse=False,
201 + out=out,
202 + **bias_kwargs,
203 )
204 + if isinstance(ret, tuple):
205 + return ret[0]
206 + return ret
207
139 - flash_attn_varlen_func = tml_flash_attn_varlen_func
140 - bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
141 - else:
142 - from vllm.vllm_flash_attn.cute import (
143 - flash_attn_varlen_func as cute_flash_attn_varlen_func,
208 + def dispatch( # type: ignore[override]
209 + self,
210 + *,
211 + is_local: bool,
212 + num_heads: int,
213 + num_kv_heads: int,
214 + head_dim: int,
215 + rel_extent: int,
216 + dtype: torch.dtype,
217 + kv_dtype: torch.dtype,
218 + block_size: int,
219 + window_size: tuple[int, int],
220 + max_kv_len: int,
221 + query_len: int,
222 + num_reqs: int,
223 + ) -> CompileKey:
224 + max_seqlen_q = bucket_max_seqlen_q(query_len)
225 + num_splits = inkling_fa4_num_splits(
226 + is_local=is_local,
227 + batch_size=num_reqs,
228 + max_query_len=max_seqlen_q,
229 + num_heads=num_heads,
230 + num_kv_heads=num_kv_heads,
231 + max_kv_len=max_kv_len,
232 + )
233 + return self.CompileKey(
234 + is_local=is_local,
235 + num_heads=num_heads,
236 + num_kv_heads=num_kv_heads,
237 + head_dim=head_dim,
238 + rel_extent=rel_extent,
239 + dtype=dtype,
240 + kv_dtype=kv_dtype,
241 + block_size=block_size,
242 + window_size=window_size,
243 + max_seqlen_q=max_seqlen_q,
244 + num_splits=num_splits,
245 + num_warps_bucket=(_num_warps_bucket(num_reqs) if num_splits > 1 else None),
246 + large_num_reqs=num_reqs > 1024,
247 )
248
146 - flash_attn_varlen_func = cute_flash_attn_varlen_func
147 - bias_kwargs = {
148 - "score_mod": _get_score_mod(rel_extent),
149 - "aux_tensors": [rel_logits],
150 - }
151 -
152 - ret = flash_attn_varlen_func(
153 - q=q,
154 - k=key_cache,
155 - v=value_cache,
156 - cu_seqlens_q=cu_seqlens_q,
157 - seqused_k=cache_seqlens,
158 - max_seqlen_q=max_seqlen_q,
159 - page_table=block_table,
160 - softmax_scale=softmax_scale,
161 - causal=causal,
162 - window_size=cute_window,
163 - num_splits=num_splits,
164 - return_lse=False,
165 - out=out,
166 - **bias_kwargs,
167 - )
168 - if isinstance(ret, tuple):
169 - return ret[0]
170 - return ret
249 + def _is_valid_warmup_dispatch(
250 + self,
251 + *,
252 + query_len: int,
253 + num_reqs: int,
254 + max_num_batched_tokens: int,
255 + ) -> bool:
256 + return query_len + num_reqs <= max_num_batched_tokens + 1
257 +
258 + def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
259 + max_num_reqs = vllm_config.scheduler_config.max_num_seqs
260 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
261 + if max_num_reqs <= 0 or max_num_batched_tokens <= 0:
262 + return []
263 +
264 + hf_config = vllm_config.model_config.hf_config
265 + get_text_config = getattr(hf_config, "get_text_config", None)
266 + config = get_text_config() if callable(get_text_config) else hf_config
267 +
268 + tp_size = get_tensor_model_parallel_world_size()
269 + dtype = vllm_config.model_config.dtype
270 + kv_dtype = kv_cache_dtype_str_to_dtype(
271 + vllm_config.cache_config.cache_dtype,
272 + vllm_config.model_config,
273 + )
274 + block_size = vllm_config.cache_config.block_size
275 + local_extent = config.sliding_window_size
276 +
277 + global_num_kv_heads = config.num_key_value_heads
278 + local_num_kv_heads = config.swa_num_key_value_heads
279 + assert config.num_attention_heads % tp_size == 0
280 + assert config.swa_num_attention_heads % tp_size == 0
281 + if global_num_kv_heads >= tp_size:
282 + assert global_num_kv_heads % tp_size == 0
283 + else:
284 + assert tp_size % global_num_kv_heads == 0
285 + if local_num_kv_heads >= tp_size:
286 + assert local_num_kv_heads % tp_size == 0
287 + else:
288 + assert tp_size % local_num_kv_heads == 0
289 +
290 + return self._trace_dispatch(self.dispatch)(
291 + zip_inputs(
292 + dict(
293 + is_local=False,
294 + num_heads=config.num_attention_heads // tp_size,
295 + num_kv_heads=max(1, global_num_kv_heads // tp_size),
296 + head_dim=config.head_dim,
297 + rel_extent=config.rel_extent,
298 + dtype=dtype,
299 + kv_dtype=kv_dtype,
300 + block_size=block_size,
301 + window_size=(-1, -1),
302 + max_kv_len=vllm_config.model_config.max_model_len,
303 + ),
304 + dict(
305 + is_local=True,
306 + num_heads=config.swa_num_attention_heads // tp_size,
307 + num_kv_heads=max(1, local_num_kv_heads // tp_size),
308 + head_dim=config.swa_head_dim,
309 + rel_extent=local_extent,
310 + dtype=dtype,
311 + kv_dtype=kv_dtype,
312 + block_size=block_size,
313 + window_size=(local_extent - 1, 0),
314 + max_kv_len=local_extent,
315 + ),
316 + ),
317 + query_len=WarmupIntRange(
318 + 1,
319 + max_num_batched_tokens + 1,
320 + advance=lambda value: bucket_max_seqlen_q(value) + 1,
321 + ),
322 + num_reqs=WarmupIntRange(1, max_num_reqs + 1),
323 + max_num_batched_tokens=max_num_batched_tokens,
324 + _when=self._is_valid_warmup_dispatch,
325 + )
326 +
327 + def compile(self, compile_key: CompileKey) -> None:
328 + from torch._subclasses.fake_tensor import FakeTensorMode
329 +
330 + if compile_key.num_splits == 1 or compile_key.num_warps_bucket is None:
331 + num_reqs = 1025 if compile_key.large_num_reqs else 1
332 + else:
333 + min_num_warps = (
334 + 1
335 + if compile_key.num_warps_bucket == 1
336 + else compile_key.num_warps_bucket // 2 + 1
337 + )
338 + min_num_reqs = (min_num_warps - 1) * 31 + 1
339 + num_reqs = max(
340 + min_num_reqs,
341 + 1025 if compile_key.large_num_reqs else 1,
342 + )
343 +
344 + with FakeTensorMode():
345 + device = torch.accelerator.current_accelerator()
346 + total_q = compile_key.max_seqlen_q + num_reqs - 1
347 + q = torch.empty(
348 + total_q,
349 + compile_key.num_heads,
350 + compile_key.head_dim,
351 + dtype=compile_key.dtype,
352 + device=device,
353 + )
354 + kv = torch.empty(
355 + 1,
356 + 2,
357 + compile_key.block_size,
358 + compile_key.num_kv_heads,
359 + compile_key.head_dim,
360 + dtype=compile_key.kv_dtype,
361 + device=device,
362 + )
363 + key_cache, value_cache = kv.unbind(1)
364 + self.kernel(
365 + q,
366 + key_cache,
367 + value_cache,
368 + block_table=torch.empty(
369 + num_reqs,
370 + 1,
371 + dtype=torch.int32,
372 + device=device,
373 + ),
374 + cache_seqlens=torch.empty(
375 + num_reqs,
376 + dtype=torch.int32,
377 + device=device,
378 + ),
379 + cu_seqlens_q=torch.empty(
380 + num_reqs + 1,
381 + dtype=torch.int32,
382 + device=device,
383 + ),
384 + max_seqlen_q=compile_key.max_seqlen_q,
385 + softmax_scale=compile_key.head_dim**-1,
386 + causal=True,
387 + window_size=compile_key.window_size,
388 + rel_extent=compile_key.rel_extent,
389 + rel_logits=torch.empty(
390 + total_q,
391 + compile_key.num_heads,
392 + compile_key.rel_extent,
393 + dtype=compile_key.dtype,
394 + device=device,
395 + ),
396 + num_splits=compile_key.num_splits,
397 + out=torch.empty_like(q),
398 + )
399 +
400 + def __call__(
401 + self,
402 + q: torch.Tensor,
403 + key_cache: torch.Tensor,
404 + value_cache: torch.Tensor,
405 + *,
406 + block_table: torch.Tensor,
407 + cache_seqlens: torch.Tensor,
408 + cu_seqlens_q: torch.Tensor,
409 + max_seqlen_q: int,
410 + softmax_scale: float,
411 + causal: bool,
412 + window_size: tuple[int, int],
413 + rel_extent: int,
414 + rel_logits: torch.Tensor,
415 + num_splits: int = 32,
416 + out: torch.Tensor | None = None,
417 + ) -> torch.Tensor:
418 + return self.kernel(
419 + q,
420 + key_cache,
421 + value_cache,
422 + block_table=block_table,
423 + cache_seqlens=cache_seqlens,
424 + cu_seqlens_q=cu_seqlens_q,
425 + max_seqlen_q=max_seqlen_q,
426 + softmax_scale=softmax_scale,
427 + causal=causal,
428 + window_size=window_size,
429 + rel_extent=rel_extent,
430 + rel_logits=rel_logits,
431 + num_splits=num_splits,
432 + out=out,
433 + )
434 +
435 +
436 +INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()
vllm/models/inkling/nvidia/ops/fa4_warmup.py deleted
-154
@@ -1,154 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Startup compilation of the FA4 kernels used by Inkling."""
4 -
5 -from __future__ import annotations
6 -
7 -from collections.abc import Iterator
8 -from dataclasses import dataclass
9 -from functools import partial
10 -
11 -import torch
12 -
13 -from vllm.model_executor.warmup.cutedsl_warmup import (
14 - CuTeDSLCompileUnit,
15 - register_cutedsl_warmup_provider,
16 -)
17 -
18 -from .fa4_rel_attention import (
19 - bucket_max_seqlen_q,
20 - inkling_fa4_num_splits,
21 - inkling_fa4_rel_attention,
22 -)
23 -
24 -
25 -@dataclass(frozen=True)
26 -class InklingFA4WarmupConfig:
27 - num_heads: int
28 - num_kv_heads: int
29 - head_dim: int
30 - rel_extent: int
31 - window_size: tuple[int, int]
32 - is_local: bool
33 - max_kv_len: int
34 - dtype: torch.dtype
35 - kv_dtype: torch.dtype
36 - block_size: int
37 - max_num_reqs: int
38 - max_num_batched_tokens: int
39 -
40 -
41 -def _num_warps_bucket(num_reqs: int) -> int:
42 - num_warps = min((num_reqs + 30) // 31, 32)
43 - return 1 << (num_warps - 1).bit_length()
44 -
45 -
46 -def _compile(config: InklingFA4WarmupConfig, max_seqlen_q: int, num_reqs: int) -> None:
47 - from torch._subclasses.fake_tensor import FakeTensorMode
48 -
49 - num_splits = inkling_fa4_num_splits(
50 - is_local=config.is_local,
51 - batch_size=num_reqs,
52 - max_query_len=max_seqlen_q,
53 - num_heads=config.num_heads,
54 - num_kv_heads=config.num_kv_heads,
55 - max_kv_len=config.max_kv_len,
56 - )
57 - with FakeTensorMode():
58 - device = torch.accelerator.current_accelerator()
59 - total_q = max_seqlen_q + num_reqs - 1
60 - q = torch.empty(
61 - total_q,
62 - config.num_heads,
63 - config.head_dim,
64 - dtype=config.dtype,
65 - device=device,
66 - )
67 - kv = torch.empty(
68 - 1,
69 - 2,
70 - config.block_size,
71 - config.num_kv_heads,
72 - config.head_dim,
73 - dtype=config.kv_dtype,
74 - device=device,
75 - )
76 - key_cache, value_cache = kv.unbind(1)
77 - inkling_fa4_rel_attention(
78 - q,
79 - key_cache,
80 - value_cache,
81 - block_table=torch.empty(num_reqs, 1, dtype=torch.int32, device=device),
82 - cache_seqlens=torch.empty(num_reqs, dtype=torch.int32, device=device),
83 - cu_seqlens_q=torch.empty(num_reqs + 1, dtype=torch.int32, device=device),
84 - max_seqlen_q=max_seqlen_q,
85 - softmax_scale=1.0 / config.head_dim,
86 - causal=True,
87 - window_size=config.window_size,
88 - rel_extent=config.rel_extent,
89 - rel_logits=torch.empty(
90 - total_q,
91 - config.num_heads,
92 - config.rel_extent,
93 - dtype=config.dtype,
94 - device=device,
95 - ),
96 - num_splits=num_splits,
97 - out=torch.empty_like(q),
98 - )
99 -
100 -
101 -def _iter_compile_units(
102 - config: InklingFA4WarmupConfig,
103 -) -> Iterator[CuTeDSLCompileUnit]:
104 - max_bucket = bucket_max_seqlen_q(config.max_num_batched_tokens)
105 - max_seqlen_q = 1
106 - while max_seqlen_q <= max_bucket:
107 - min_query_len = max_seqlen_q // 2 + 1
108 - max_num_reqs = min(
109 - config.max_num_reqs,
110 - config.max_num_batched_tokens - min_query_len + 1,
111 - )
112 - seen: set[tuple[int, int, int | None, bool]] = set()
113 - for num_reqs in range(1, max_num_reqs + 1):
114 - num_splits = inkling_fa4_num_splits(
115 - is_local=config.is_local,
116 - batch_size=num_reqs,
117 - max_query_len=max_seqlen_q,
118 - num_heads=config.num_heads,
119 - num_kv_heads=config.num_kv_heads,
120 - max_kv_len=config.max_kv_len,
121 - )
122 - key = (
123 - max_seqlen_q,
124 - num_splits,
125 - _num_warps_bucket(num_reqs) if num_splits > 1 else None,
126 - num_reqs > 1024,
127 - )
128 - if key in seen:
129 - continue
130 - seen.add(key)
131 - yield CuTeDSLCompileUnit(
132 - name="inkling_fa4",
133 - key=("inkling_fa4", config, key),
134 - compile=partial(_compile, config, max_seqlen_q, num_reqs),
135 - )
136 - max_seqlen_q *= 2
137 -
138 -
139 -class _WarmupProvider:
140 - def __init__(self) -> None:
141 - self.configs: set[InklingFA4WarmupConfig] = set()
142 -
143 - def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
144 - return tuple(
145 - unit for config in self.configs for unit in _iter_compile_units(config)
146 - )
147 -
148 -
149 -_PROVIDER = _WarmupProvider()
150 -
151 -
152 -def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None:
153 - _PROVIDER.configs.add(config)
154 - register_cutedsl_warmup_provider(_PROVIDER)
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+79 -56
@@ -13,7 +13,7 @@ import vllm.envs as envs
13 from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 kFp8StaticTensorSym,
15 )
16 -from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
16 +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs
17 from vllm.platforms import current_platform
18 from vllm.v1.attention.backends.fa_utils import (
19 compile_flash_attn_varlen_func_from_specs,
@@ -57,12 +57,6 @@ FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True)
57 FA4_MLA_PREFILL_LSE_OPTIONS = (False, True)
58
59
60 -@dataclass(frozen=True)
61 -class _FA4MLAPrefillShapeProbe:
62 - max_seqlen_q: int
63 - max_seqlen_k: int
64 -
65 -
60 class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
61 """FA4 MLA prefill compile-key wrapper used by generic JIT warmup."""
62
@@ -92,6 +86,16 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
86 return_softmax_lse: bool = False
87 num_splits: int = 0
88
89 + @staticmethod
90 + def kernel(
91 + *args: Any,
92 + runtime_kernel: Callable[..., Any] | None = None,
93 + **kwargs: Any,
94 + ) -> Any:
95 + kernel = flash_attn_varlen_func if runtime_kernel is None else runtime_kernel
96 + assert kernel is not None
97 + return kernel(*args, **kwargs)
98 +
99 def dispatch( # type: ignore[override]
100 self,
101 *,
@@ -99,7 +103,8 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
103 dtype: torch.dtype,
104 num_heads: int,
105 mla_dims: "MLADims",
102 - shape_probe: _FA4MLAPrefillShapeProbe,
106 + max_seqlen_q: int,
107 + max_seqlen_k: int,
108 requires_v_padding: bool,
109 causal: bool,
110 fa_version: int,
@@ -109,17 +114,17 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
114 ) -> CompileKey:
115 return self.CompileKey(
116 q_shape=(
112 - batch_size * shape_probe.max_seqlen_q,
117 + batch_size * max_seqlen_q,
118 num_heads,
119 mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
120 ),
121 k_shape=(
117 - batch_size * shape_probe.max_seqlen_k,
122 + batch_size * max_seqlen_k,
123 num_heads,
124 mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
125 ),
126 v_shape=(
122 - batch_size * shape_probe.max_seqlen_k,
127 + batch_size * max_seqlen_k,
128 num_heads,
129 (
130 mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
@@ -128,8 +133,8 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
133 ),
134 ),
135 q_dtype=dtype,
131 - max_seqlen_q=shape_probe.max_seqlen_q,
132 - max_seqlen_k=shape_probe.max_seqlen_k,
136 + max_seqlen_q=max_seqlen_q,
137 + max_seqlen_k=max_seqlen_k,
138 softmax_scale=(mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim)
139 ** -0.5,
140 causal=causal,
@@ -150,6 +155,28 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
155 num_splits=num_splits,
156 )
157
158 + def _is_valid_warmup_shape_probe(
159 + self,
160 + *,
161 + max_seqlen_k: int,
162 + num_splits: int,
163 + is_sm90: bool,
164 + qk_head_dim: int,
165 + effective_v_head_dim: int,
166 + ) -> bool:
167 + long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
168 + very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
169 + return (
170 + max_seqlen_k < long_k
171 + or (num_splits != 1 and max_seqlen_k == long_k)
172 + or (
173 + num_splits != 1
174 + and max_seqlen_k == very_long_k
175 + and not is_sm90
176 + and qk_head_dim != effective_v_head_dim
177 + )
178 + )
179 +
180 def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]:
181 from vllm.model_executor.layers.attention.mla_attention import (
182 get_mla_dims,
@@ -185,58 +212,49 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
212 effective_v_head_dim = (
213 qk_head_dim if requires_v_padding else mla_dims.v_head_dim
214 )
188 - shape_probes = [
189 - _FA4MLAPrefillShapeProbe(1, FA4_MLA_PREFILL_K_TILE),
190 - _FA4MLAPrefillShapeProbe(
191 - FA4_MLA_PREFILL_Q_TILE + 1,
192 - 4 * FA4_MLA_PREFILL_K_TILE,
215 + shape_probes = (
216 + dict(max_seqlen_q=1, max_seqlen_k=FA4_MLA_PREFILL_K_TILE),
217 + dict(
218 + max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1,
219 + max_seqlen_k=4 * FA4_MLA_PREFILL_K_TILE,
220 ),
194 - ]
195 - if num_splits != 1:
196 - shape_probes.extend(
197 - (
198 - _FA4MLAPrefillShapeProbe(
199 - 1,
200 - FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
201 - ),
202 - _FA4MLAPrefillShapeProbe(
203 - FA4_MLA_PREFILL_Q_TILE + 1,
204 - FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
205 - ),
206 - )
207 - )
208 - if not is_sm90 and qk_head_dim != effective_v_head_dim:
209 - shape_probes.extend(
210 - (
211 - _FA4MLAPrefillShapeProbe(
212 - 1,
213 - FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
214 - ),
215 - _FA4MLAPrefillShapeProbe(
216 - FA4_MLA_PREFILL_Q_TILE + 1,
217 - FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
218 - ),
219 - )
220 - )
221 + dict(
222 + max_seqlen_q=1,
223 + max_seqlen_k=FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
224 + ),
225 + dict(
226 + max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1,
227 + max_seqlen_k=FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
228 + ),
229 + dict(
230 + max_seqlen_q=1,
231 + max_seqlen_k=FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS
232 + * FA4_MLA_PREFILL_K_TILE,
233 + ),
234 + dict(
235 + max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1,
236 + max_seqlen_k=FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS
237 + * FA4_MLA_PREFILL_K_TILE,
238 + ),
239 + )
240
241 return self._trace_dispatch(self.dispatch)(
242 + zip_inputs(*shape_probes),
243 batch_size=FA4_MLA_PREFILL_COMPILE_BATCH_SIZE,
244 dtype=dtype,
245 num_heads=num_heads,
246 mla_dims=mla_dims,
227 - shape_probe=shape_probes,
247 requires_v_padding=requires_v_padding,
248 + is_sm90=is_sm90,
249 + qk_head_dim=qk_head_dim,
250 + effective_v_head_dim=effective_v_head_dim,
251 causal=FA4_MLA_PREFILL_CAUSAL_OPTIONS,
252 return_lse=FA4_MLA_PREFILL_LSE_OPTIONS,
253 num_splits=num_splits,
254 fa_version=fa_version,
255 + _when=self._is_valid_warmup_shape_probe,
256 )
257
235 - @staticmethod
236 - def kernel(*args: Any, **kwargs: Any) -> Any:
237 - assert flash_attn_varlen_func is not None
238 - return flash_attn_varlen_func(*args, **kwargs)
239 -
258 def compile(self, compile_key: CompileKey) -> None:
259 assert compile_flash_attn_varlen_func_from_specs is not None
260 window_size = (
@@ -271,11 +289,13 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
289 runtime_kernel: Callable[..., Any] | None = None,
290 **kwargs: Any,
291 ) -> Any:
274 - kernel = self.kernel if runtime_kernel is None else runtime_kernel
275 - return kernel(q=q, k=k, v=v, **kwargs)
276 -
277 -
278 -FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()
292 + return self.kernel(
293 + q=q,
294 + k=k,
295 + v=v,
296 + runtime_kernel=runtime_kernel,
297 + **kwargs,
298 + )
299
300
301 class FlashAttnPrefillBackend(MLAPrefillBackend):
@@ -470,3 +490,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
490 causal=False, # Context is unmasked
491 return_softmax_lse=True,
492 )
493 +
494 +
495 +FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()