[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>

tests/model_executor/test_jit_warmup.py

@@ -105,6 +105,39 @@ def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> No
     ]
 
 
+def test_warmup_range_uses_custom_advancement() -> None:
+    keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
+        tokens=WarmupIntRange(
+            1,
+            10,
+            advance=lambda value: _next_power_of_2(value) + 1,
+        ),
+        cfg=_config(),
+    )
+
+    assert [key.block_size for key in keys] == [1, 2, 4, 8, 16]
+
+
+def test_warmup_range_validates_custom_advancement() -> None:
+    trace_dispatch = ToyKernel()._trace_dispatch(ToyKernel().dispatch)
+
+    with pytest.raises(ValueError, match="both step and advance"):
+        trace_dispatch(
+            tokens=WarmupIntRange(
+                1,
+                10,
+                step=2,
+                advance=lambda value: value + 1,
+            ),
+            cfg=_config(),
+        )
+    with pytest.raises(ValueError, match="must return a greater value"):
+        trace_dispatch(
+            tokens=WarmupIntRange(1, 10, advance=lambda value: value),
+            cfg=_config(),
+        )
+
+
 def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
     cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
 
@@ -143,6 +176,38 @@ def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
     ]
 
 
+def test_trace_dispatch_filters_with_traced_predicate() -> None:
+    class PredicateKernel(ToyKernel):
+        def _is_valid_warmup_dispatch(
+            self,
+            *,
+            tokens: int,
+            lanes: int,
+            max_work: int,
+        ) -> bool:
+            block_size = _next_power_of_2(tokens)
+            return block_size * lanes <= max_work
+
+        def get_warmup_keys(
+            self, max_tokens: int, cfg: Any
+        ) -> list[ToyKernel.CompileKey]:
+            return self._trace_dispatch(self.dispatch)(
+                tokens=WarmupIntRange(1, max_tokens + 1),
+                lanes=(1, 2),
+                cfg=cfg,
+                max_work=4,
+                _when=self._is_valid_warmup_dispatch,
+            )
+
+    assert PredicateKernel().get_warmup_keys(5, _config()) == [
+        ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
+        ToyKernel.CompileKey(1, 2, 1, ("base", "default", -1, 1, 1), True),
+        ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
+        ToyKernel.CompileKey(2, 4, 1, ("base", "default", -2, 2, 4), True),
+        ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
+    ]
+
+
 def test_zip_inputs_validates_input_rows() -> None:
     with pytest.raises(ValueError, match="requires at least one"):
         zip_inputs()

tests/models/inkling/test_fa4_rel_attention.py

@@ -2,8 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Correctness test for the Inkling FA4 relative-attention score-mod kernel.
 
-Checks ``inkling_fa4_rel_attention`` against a pure-PyTorch reference that
-implements the relative bias exactly as documented in the Inkling architecture
+Checks ``INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference
+that implements the relative bias exactly as documented in the Inkling architecture
 guide::
 
     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 (
     compute_log_scaling_tau,
 )
 from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
+    INKLING_FA4_REL_ATTENTION_KERNEL,
     _use_sheared_bias,
     inkling_fa4_num_splits,
-    inkling_fa4_rel_attention,
 )
 from vllm.platforms import current_platform
 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
         num_kv_heads=num_kv_heads,
         max_kv_len=max(kv_lens),
     )
-    out = inkling_fa4_rel_attention(
+    out = INKLING_FA4_REL_ATTENTION_KERNEL(
         q,
         key_cache,
         value_cache,

tests/models/inkling/test_fa4_warmup.py

@@ -1,17 +1,47 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
 import torch
 
+import vllm.models.inkling.nvidia.ops.fa4_rel_attention as fa4_rel_attention
+from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig
 from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
+    InklingFA4RelAttentionKernel,
+    _num_warps_bucket,
     bucket_max_seqlen_q,
     inkling_fa4_num_splits,
 )
-from vllm.models.inkling.nvidia.ops.fa4_warmup import (
-    InklingFA4WarmupConfig,
-    _iter_compile_units,
-    _num_warps_bucket,
-)
+
+
+def _vllm_config_from_reference_config(config: SimpleNamespace) -> SimpleNamespace:
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=InklingMMConfig(
+                text_config=InklingModelConfig(
+                    num_attention_heads=config.num_heads,
+                    num_key_value_heads=config.num_kv_heads,
+                    head_dim=config.head_dim,
+                    rel_extent=config.rel_extent,
+                    swa_num_attention_heads=config.num_heads,
+                    swa_num_key_value_heads=config.num_kv_heads,
+                    swa_head_dim=config.head_dim,
+                    sliding_window_size=config.rel_extent,
+                )
+            ),
+            max_model_len=config.max_kv_len,
+            dtype=config.dtype,
+        ),
+        cache_config=SimpleNamespace(
+            cache_dtype="auto",
+            block_size=config.block_size,
+        ),
+        scheduler_config=SimpleNamespace(
+            max_num_seqs=config.max_num_reqs,
+            max_num_batched_tokens=config.max_num_batched_tokens,
+        ),
+    )
 
 
 def test_bucket_max_seqlen_q():
@@ -28,8 +58,13 @@ def test_bucket_max_seqlen_q():
     ]
 
 
-def test_warmup_enumerates_every_runtime_compile_class():
-    config = InklingFA4WarmupConfig(
+def test_warmup_enumerates_every_runtime_compile_class(monkeypatch):
+    monkeypatch.setattr(
+        fa4_rel_attention,
+        "get_tensor_model_parallel_world_size",
+        lambda: 1,
+    )
+    config = SimpleNamespace(
         num_heads=16,
         num_kv_heads=2,
         head_dim=128,
@@ -43,27 +78,47 @@ def test_warmup_enumerates_every_runtime_compile_class():
         max_num_reqs=64,
         max_num_batched_tokens=192,
     )
-    warmed_keys = {unit.key[-1] for unit in _iter_compile_units(config)}
+    vllm_config = _vllm_config_from_reference_config(config)
+    kernel = InklingFA4RelAttentionKernel()
+    warmed_keys = set(kernel.get_warmup_keys(vllm_config))
+    expected_keys: set[InklingFA4RelAttentionKernel.CompileKey] = set()
 
-    for query_len in range(1, config.max_num_batched_tokens + 1):
-        max_seqlen_q = bucket_max_seqlen_q(query_len)
-        max_num_reqs = min(
-            config.max_num_reqs,
-            config.max_num_batched_tokens - query_len + 1,
-        )
-        for num_reqs in range(1, max_num_reqs + 1):
-            num_splits = inkling_fa4_num_splits(
-                is_local=config.is_local,
-                batch_size=num_reqs,
-                max_query_len=max_seqlen_q,
-                num_heads=config.num_heads,
-                num_kv_heads=config.num_kv_heads,
-                max_kv_len=config.max_kv_len,
+    for is_local in (False, True):
+        max_kv_len = config.rel_extent if is_local else config.max_kv_len
+        window_size = (config.rel_extent - 1, 0) if is_local else (-1, -1)
+        for query_len in range(1, config.max_num_batched_tokens + 1):
+            max_seqlen_q = bucket_max_seqlen_q(query_len)
+            max_num_reqs = min(
+                config.max_num_reqs,
+                config.max_num_batched_tokens - query_len + 1,
             )
-            runtime_key = (
-                max_seqlen_q,
-                num_splits,
-                _num_warps_bucket(num_reqs) if num_splits > 1 else None,
-                num_reqs > 1024,
-            )
-            assert runtime_key in warmed_keys
+            for num_reqs in range(1, max_num_reqs + 1):
+                num_splits = inkling_fa4_num_splits(
+                    is_local=is_local,
+                    batch_size=num_reqs,
+                    max_query_len=max_seqlen_q,
+                    num_heads=config.num_heads,
+                    num_kv_heads=config.num_kv_heads,
+                    max_kv_len=max_kv_len,
+                )
+                expected_keys.add(
+                    InklingFA4RelAttentionKernel.CompileKey(
+                        is_local=is_local,
+                        num_heads=config.num_heads,
+                        num_kv_heads=config.num_kv_heads,
+                        head_dim=config.head_dim,
+                        rel_extent=config.rel_extent,
+                        dtype=config.dtype,
+                        kv_dtype=config.kv_dtype,
+                        block_size=config.block_size,
+                        window_size=window_size,
+                        max_seqlen_q=max_seqlen_q,
+                        num_splits=num_splits,
+                        num_warps_bucket=(
+                            _num_warps_bucket(num_reqs) if num_splits > 1 else None
+                        ),
+                        large_num_reqs=num_reqs > 1024,
+                    )
+                )
+
+    assert warmed_keys == expected_keys

vllm/config/kernel.py

@@ -180,7 +180,7 @@ class KernelConfig:
     # to the shared JIT warmup infrastructure.
     # https://github.com/vllm-project/vllm/pull/47451
     enable_cutedsl_warmup: bool = True
-    """If True, run CuTeDSL compile warmup during kernel warmup."""
+    """Deprecated: run legacy CuTeDSL warmup providers."""
 
     enable_jit_warmup: bool = True
     """If True, run JIT compile warmup during kernel warmup."""

vllm/model_executor/warmup/cutedsl_warmup.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Run registered CuTeDSL warmup compile units."""
+"""Deprecated compatibility registry for legacy CuTeDSL warmups."""
 
 # TODO(roberto): Remove this compatibility registry after registered CuTeDSL
 # warmups are migrated to the shared JIT warmup infrastructure.

vllm/model_executor/warmup/fa4_cutedsl_warmup.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warm up FA4 CuTeDSL MLA prefill compile keys."""
+"""Warm up FA4 CuTeDSL kernels."""
 
 from __future__ import annotations
 
@@ -12,7 +12,7 @@ if TYPE_CHECKING:
     from vllm.v1.worker.gpu_worker import Worker
 
 
-def fa4_cutedsl_warmup(worker: Worker) -> None:
+def _warm_fa4_mla_prefill(worker: Worker) -> None:
     runner = worker.model_runner
     if runner.is_pooling_model:
         return
@@ -32,3 +32,22 @@ def fa4_cutedsl_warmup(worker: Worker) -> None:
     from vllm.v1.attention.backends.mla.prefill import flash_attn
 
     flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config)
+
+
+def _warm_inkling_fa4_rel_attention(worker: Worker) -> None:
+    from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig
+    from vllm.models.inkling.nvidia.ops.fa4_rel_attention import (
+        INKLING_FA4_REL_ATTENTION_KERNEL,
+    )
+
+    vllm_config = worker.vllm_config
+    hf_config = vllm_config.model_config.hf_config
+    if not isinstance(hf_config, (InklingMMConfig, InklingModelConfig)):
+        return
+
+    INKLING_FA4_REL_ATTENTION_KERNEL.warmup(vllm_config)
+
+
+def fa4_cutedsl_warmup(worker: Worker) -> None:
+    _warm_fa4_mla_prefill(worker)
+    _warm_inkling_fa4_rel_attention(worker)

vllm/model_executor/warmup/jit_warmup.py

@@ -28,13 +28,29 @@ CompileKeyT = TypeVar("CompileKeyT")
 
 @dataclass(frozen=True)
 class WarmupIntRange:
+    """Expand integers with range semantics or a custom monotonic progression."""
+
     start: int
     stop: int
     step: int = 1
+    advance: Callable[[int], int] | None = None
 
 
 WarmupValues = Any
 CompileKeyDispatchFn = Callable[..., CompileKeyT]
+WarmupPredicateFn = Callable[..., bool]
+_LocalExprs = tuple[tuple[str, ast.AST], ...]
+
+
+def _eval_local_exprs(
+    local_exprs: _LocalExprs,
+    values: Mapping[str, Any],
+    globals_: Mapping[str, Any],
+) -> dict[str, Any]:
+    evaluated_values = dict(values)
+    for name, expr in local_exprs:
+        evaluated_values[name] = _eval_dispatch_expr(expr, evaluated_values, globals_)
+    return evaluated_values
 
 
 @dataclass(frozen=True)
@@ -46,7 +62,20 @@ class _WarmupInputRows:
 
 def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:
     if isinstance(values, WarmupIntRange):
-        return tuple(range(values.start, values.stop, values.step))
+        if values.advance is None:
+            return tuple(range(values.start, values.stop, values.step))
+        if values.step != 1:
+            raise ValueError("WarmupIntRange cannot set both step and advance")
+
+        expanded: list[int] = []
+        value = values.start
+        while value < values.stop:
+            expanded.append(value)
+            next_value = values.advance(value)
+            if next_value <= value:
+                raise ValueError("WarmupIntRange.advance must return a greater value")
+            value = next_value
+        return tuple(expanded)
     if isinstance(values, (list, tuple)):
         return tuple(values)
     return (values,)
@@ -75,20 +104,6 @@ def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows:
     return _WarmupInputRows(rows=tuple(input_rows))
 
 
-def _expand_warmup_value_grid(
-    values: Mapping[str, WarmupValues],
-    input_names: frozenset[str],
-) -> tuple[dict[str, Any], ...]:
-    names = tuple(name for name in values if name in input_names)
-    if not names:
-        return ({},)
-
-    expanded_values = tuple(_expand_warmup_values(values[name]) for name in names)
-    return tuple(
-        dict(zip(names, value_set)) for value_set in itertools.product(*expanded_values)
-    )
-
-
 def _expand_warmup_input_rows(
     rows: tuple[Mapping[str, WarmupValues], ...],
     input_names: frozenset[str],
@@ -117,7 +132,7 @@ def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]:
 
 @dataclass(frozen=True)
 class _CompileKeyDispatchTrace:
-    local_exprs: tuple[tuple[str, ast.AST], ...]
+    local_exprs: _LocalExprs
     field_exprs: tuple[tuple[str, ast.AST], ...]
     globals: Mapping[str, Any]
     input_names: frozenset[str]
@@ -128,11 +143,9 @@ class _CompileKeyDispatchTrace:
         compile_key_type: type[CompileKeyT],
         kwargs: Mapping[str, Any],
     ) -> CompileKeyT:
-        dispatch_values = {**self.defaults, **kwargs}
-        for name, expr in self.local_exprs:
-            dispatch_values[name] = _eval_dispatch_expr(
-                expr, dispatch_values, self.globals
-            )
+        dispatch_values = _eval_local_exprs(
+            self.local_exprs, {**self.defaults, **kwargs}, self.globals
+        )
         return compile_key_type(
             **{
                 field: _eval_dispatch_expr(expr, dispatch_values, self.globals)
@@ -141,6 +154,23 @@ class _CompileKeyDispatchTrace:
         )
 
 
+@dataclass(frozen=True)
+class _WarmupPredicateTrace:
+    local_exprs: _LocalExprs
+    return_expr: ast.AST
+    globals: Mapping[str, Any]
+    input_names: frozenset[str]
+    defaults: Mapping[str, Any]
+
+    def matches(self, kwargs: Mapping[str, Any]) -> bool:
+        dispatch_values = _eval_local_exprs(
+            self.local_exprs, {**self.defaults, **kwargs}, self.globals
+        )
+        return bool(
+            _eval_dispatch_expr(self.return_expr, dispatch_values, self.globals)
+        )
+
+
 _BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
     ast.Add: operator.add,
     ast.Sub: operator.sub,
@@ -317,10 +347,10 @@ def _collect_input_names(
     }
 
 
-def _collect_dispatch_body(
-    fn: CompileKeyDispatchFn[Any],
+def _collect_expression_body(
+    fn: Callable[..., Any],
     function_def: ast.FunctionDef,
-) -> tuple[list[tuple[str, ast.AST]], ast.Call]:
+) -> tuple[list[tuple[str, ast.AST]], ast.AST]:
     local_exprs: list[tuple[str, ast.AST]] = []
     for statement in function_def.body:
         if (
@@ -335,7 +365,8 @@ def _collect_dispatch_body(
                 statement.targets[0], ast.Name
             ):
                 raise _dispatch_expr_error(
-                    statement, "Dispatch assignments must target one local name"
+                    statement,
+                    "Traced warmup helper assignments must target one local name",
                 )
             local_exprs.append((statement.targets[0].id, statement.value))
             continue
@@ -343,29 +374,45 @@ def _collect_dispatch_body(
         if isinstance(statement, ast.AnnAssign):
             if statement.value is None:
                 raise _dispatch_expr_error(
-                    statement, "Dispatch annotations must assign a value"
+                    statement, "Traced warmup helper annotations must assign a value"
                 )
             if not isinstance(statement.target, ast.Name):
                 raise _dispatch_expr_error(
-                    statement, "Dispatch assignments must target one local name"
+                    statement,
+                    "Traced warmup helper assignments must target one local name",
                 )
             local_exprs.append((statement.target.id, statement.value))
             continue
 
-        if isinstance(statement, ast.Return) and isinstance(statement.value, ast.Call):
+        if isinstance(statement, ast.Return) and statement.value is not None:
             return local_exprs, statement.value
 
         if isinstance(statement, ast.Return):
             raise _dispatch_expr_error(
-                statement, "Dispatch must return one CompileKey(...) call"
+                statement, "Traced warmup helper must return an expression"
             )
 
         raise _dispatch_expr_error(
             statement,
-            "Dispatch may only contain local assignments before CompileKey return",
+            "Traced warmup helper may only contain local assignments before return",
         )
 
-    raise ValueError(f"Expected {fn.__name__} to return one CompileKey(...) call")
+    raise ValueError(f"Expected {fn.__name__} to return an expression")
+
+
+def _function_trace_inputs(
+    fn: Callable[..., Any],
+) -> tuple[dict[str, Any], set[str]]:
+    signature = inspect.signature(fn)
+    defaults = {
+        name: parameter.default
+        for name, parameter in signature.parameters.items()
+        if parameter.default is not inspect.Parameter.empty
+    }
+    bound_self = getattr(fn, "__self__", None)
+    if bound_self is not None:
+        defaults["self"] = bound_self
+    return defaults, set(signature.parameters)
 
 
 def _trace_compile_key_dispatch(
@@ -375,21 +422,20 @@ def _trace_compile_key_dispatch(
     globals_ = source_fn.__globals__
     function_def = get_function_source_node(fn)
 
-    local_exprs, return_call = _collect_dispatch_body(fn, function_def)
+    local_exprs, return_expr = _collect_expression_body(fn, function_def)
+    if not isinstance(return_expr, ast.Call):
+        raise _dispatch_expr_error(
+            return_expr,
+            "Dispatch must return one CompileKey(...) call",
+        )
 
     field_exprs: list[tuple[str, ast.AST]] = []
-    signature = inspect.signature(fn)
-    defaults = {
-        name: parameter.default
-        for name, parameter in signature.parameters.items()
-        if parameter.default is not inspect.Parameter.empty
-    }
-    candidate_names = set(signature.parameters)
+    defaults, candidate_names = _function_trace_inputs(fn)
     input_names: set[str] = set()
     local_names = {name for name, _ in local_exprs}
     for _, expr in local_exprs:
         input_names.update(_collect_input_names(expr, candidate_names))
-    for keyword in return_call.keywords:
+    for keyword in return_expr.keywords:
         if keyword.arg is None:
             raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey")
         field_exprs.append((keyword.arg, keyword.value))
@@ -406,6 +452,30 @@ def _trace_compile_key_dispatch(
     )
 
 
+def _trace_warmup_predicate(
+    fn: WarmupPredicateFn,
+) -> _WarmupPredicateTrace:
+    source_fn = getattr(fn, "__func__", fn)
+    globals_ = source_fn.__globals__
+    function_def = get_function_source_node(fn)
+
+    local_exprs, return_expr = _collect_expression_body(fn, function_def)
+    defaults, candidate_names = _function_trace_inputs(fn)
+    input_names: set[str] = set()
+    local_names = {name for name, _ in local_exprs}
+    for _, expr in local_exprs:
+        input_names.update(_collect_input_names(expr, candidate_names))
+    input_names.update(_collect_input_names(return_expr, candidate_names, local_names))
+
+    return _WarmupPredicateTrace(
+        tuple(local_exprs),
+        return_expr,
+        globals_,
+        frozenset(input_names),
+        defaults,
+    )
+
+
 class VllmJitKernel(Generic[CompileKeyT], ABC):
     """Kernel wrapper that owns dispatch, warmup keys, and compilation."""
 
@@ -424,6 +494,7 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
 
         def traced(
             *input_groups: _WarmupInputRows,
+            _when: WarmupPredicateFn | None = None,
             **kwargs: WarmupValues,
         ) -> list[CompileKeyT]:
             for group in input_groups:
@@ -432,24 +503,50 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
                         "_trace_dispatch positional arguments must be "
                         "zip_inputs(...) groups"
                     )
+            predicate_trace = (
+                _trace_warmup_predicate(_when) if _when is not None else None
+            )
+            input_names = compile_key_dispatch_trace.input_names
+            if predicate_trace is not None:
+                input_names = input_names | predicate_trace.input_names
             expanded_input_groups = tuple(
-                _expand_warmup_input_rows(
-                    group.rows, compile_key_dispatch_trace.input_names
-                )
+                _expand_warmup_input_rows(group.rows, input_names)
                 for group in input_groups
             )
-            expanded_kwargs = _expand_warmup_value_grid(
-                kwargs, compile_key_dispatch_trace.input_names
+            # Expand independent keyword inputs into cartesian-product axes.
+            expanded_kwarg_axes = tuple(
+                (name, _expand_warmup_values(value))
+                for name, value in kwargs.items()
+                if name in input_names
             )
-            dispatch_value_groups = (*expanded_input_groups, expanded_kwargs)
-            return list(
-                dict.fromkeys(
-                    compile_key_dispatch_trace.compile_key(
-                        self.CompileKey, _merge_warmup_kwargs(dispatch_values)
+            dispatch_value_axes = (
+                *expanded_input_groups,
+                *(values for _, values in expanded_kwarg_axes),
+            )
+            input_group_count = len(expanded_input_groups)
+            kwarg_names = tuple(name for name, _ in expanded_kwarg_axes)
+            compile_keys: dict[CompileKeyT, None] = {}
+            for dispatch_value_set in itertools.product(*dispatch_value_axes):
+                dispatch_values = _merge_warmup_kwargs(
+                    (
+                        *dispatch_value_set[:input_group_count],
+                        dict(
+                            zip(
+                                kwarg_names,
+                                dispatch_value_set[input_group_count:],
+                            )
+                        ),
                     )
-                    for dispatch_values in itertools.product(*dispatch_value_groups)
                 )
-            )
+                if predicate_trace is not None and not predicate_trace.matches(
+                    dispatch_values
+                ):
+                    continue
+                compile_key = compile_key_dispatch_trace.compile_key(
+                    self.CompileKey, dispatch_values
+                )
+                compile_keys[compile_key] = None
+            return list(compile_keys)
 
         return traced
 

vllm/models/inkling/nvidia/attention.py

@@ -35,11 +35,10 @@ from vllm.v1.kv_cache_interface import (
 from ..configs import InklingModelConfig
 from .layernorm import InklingRMSNorm
 from .ops.fa4_rel_attention import (
+    INKLING_FA4_REL_ATTENTION_KERNEL,
     bucket_max_seqlen_q,
     inkling_fa4_num_splits,
-    inkling_fa4_rel_attention,
 )
-from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup
 from .ops.qkvr_prep import fused_qkvr_prep
 from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata
 from .short_conv import InklingShortConv
@@ -175,25 +174,6 @@ class InklingAttention(nn.Module, AttentionLayerBase):
         compilation_config.static_forward_context[prefix] = self
         self.kv_cache = torch.tensor([])  # replaced by bind_kv_cache
 
-        register_fa4_warmup(
-            InklingFA4WarmupConfig(
-                num_heads=self.num_heads,
-                num_kv_heads=self.num_kv_heads,
-                head_dim=self.head_dim,
-                rel_extent=self.rel_extent,
-                window_size=self.window_size,
-                is_local=self.is_local,
-                max_kv_len=self._max_kv_len,
-                dtype=vllm_config.model_config.dtype,
-                kv_dtype=self.kv_cache_torch_dtype,
-                block_size=vllm_config.cache_config.block_size,
-                max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
-                max_num_batched_tokens=(
-                    vllm_config.scheduler_config.max_num_batched_tokens
-                ),
-            )
-        )
-
     def get_attn_backend(self) -> type[AttentionBackend]:
         return FlashAttentionBackend
 
@@ -309,7 +289,7 @@ class InklingAttention(nn.Module, AttentionLayerBase):
             num_kv_heads=self.num_kv_heads,
             max_kv_len=self._max_kv_len,
         )
-        inkling_fa4_rel_attention(
+        INKLING_FA4_REL_ATTENTION_KERNEL(
             q[:nt],
             key_cache,
             value_cache,

vllm/models/inkling/nvidia/ops/__init__.py

@@ -15,11 +15,13 @@ from .sconv import fused_sconv
 _LAZY_EXPORTS = {
     "silu_and_mul_triton": "silu_and_mul",
     "sink_silu_mul_epilogue": "silu_and_mul",
-    "inkling_fa4_rel_attention": "fa4_rel_attention",
+    "INKLING_FA4_REL_ATTENTION_KERNEL": "fa4_rel_attention",
 }
 
 if TYPE_CHECKING:
-    from .fa4_rel_attention import inkling_fa4_rel_attention  # noqa: F401
+    from .fa4_rel_attention import (  # noqa: F401
+        INKLING_FA4_REL_ATTENTION_KERNEL,
+    )
     from .silu_and_mul import (  # noqa: F401
         silu_and_mul_triton,
         sink_silu_mul_epilogue,

vllm/models/inkling/nvidia/ops/fa4_rel_attention.py

@@ -3,12 +3,23 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from dataclasses import dataclass
 from functools import cache
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.warmup.jit_warmup import (
+    VllmJitKernel,
+    WarmupIntRange,
+    zip_inputs,
+)
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
 
 
 def bucket_max_seqlen_q(max_seqlen_q: int) -> int:
@@ -98,73 +109,328 @@ def inkling_fa4_num_splits(
     )
 
 
-def inkling_fa4_rel_attention(
-    q: torch.Tensor,
-    key_cache: torch.Tensor,
-    value_cache: torch.Tensor,
-    *,
-    block_table: torch.Tensor,
-    cache_seqlens: torch.Tensor,
-    cu_seqlens_q: torch.Tensor,
-    max_seqlen_q: int,
-    softmax_scale: float,
-    causal: bool,
-    window_size: tuple[int, int],
-    rel_extent: int,
-    rel_logits: torch.Tensor,
-    num_splits: int = 32,
-    out: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias.
+def _num_warps_bucket(num_reqs: int) -> int:
+    num_warps = min((num_reqs + 30) // 31, 32)
+    return 1 << (num_warps - 1).bit_length()
 
-    ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache``
-    are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``;
-    ``block_table`` is the per-request page table and ``cache_seqlens`` the
-    per-request KV lengths (``seqused_k``). ``rel_logits`` is
-    ``(num_tokens, num_heads, rel_extent)``.
 
-    Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's
-    sheared relative-bias layout.
-    """
-    # cute uses (None, None) to mean "no window".
-    cute_window = (None, None) if window_size == (-1, -1) else window_size
+class InklingFA4RelAttentionKernel(
+    VllmJitKernel["InklingFA4RelAttentionKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        is_local: bool
+        num_heads: int
+        num_kv_heads: int
+        head_dim: int
+        rel_extent: int
+        dtype: torch.dtype
+        kv_dtype: torch.dtype
+        block_size: int
+        window_size: tuple[int, int]
+        max_seqlen_q: int
+        num_splits: int
+        num_warps_bucket: int | None
+        large_num_reqs: bool
 
-    rel_logits = rel_logits.contiguous()
-    flash_attn_varlen_func: Callable[..., Any]
-    if _use_sheared_bias():
-        from vllm.third_party.tml_fa4 import (
-            flash_attn_varlen_func as tml_flash_attn_varlen_func,
+    @staticmethod
+    def kernel(
+        q: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        *,
+        block_table: torch.Tensor,
+        cache_seqlens: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        max_seqlen_q: int,
+        softmax_scale: float,
+        causal: bool,
+        window_size: tuple[int, int],
+        rel_extent: int,
+        rel_logits: torch.Tensor,
+        num_splits: int = 32,
+        out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias.
+
+        ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache``
+        are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``;
+        ``block_table`` is the per-request page table and ``cache_seqlens`` the
+        per-request KV lengths (``seqused_k``). ``rel_logits`` is
+        ``(num_tokens, num_heads, rel_extent)``.
+
+        Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's
+        sheared relative-bias layout.
+        """
+        # cute uses (None, None) to mean "no window".
+        cute_window = (None, None) if window_size == (-1, -1) else window_size
+
+        rel_logits = rel_logits.contiguous()
+        flash_attn_varlen_func: Callable[..., Any]
+        if _use_sheared_bias():
+            from vllm.third_party.tml_fa4 import (
+                flash_attn_varlen_func as tml_flash_attn_varlen_func,
+            )
+
+            flash_attn_varlen_func = tml_flash_attn_varlen_func
+            bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
+        else:
+            from vllm.vllm_flash_attn.cute import (
+                flash_attn_varlen_func as cute_flash_attn_varlen_func,
+            )
+
+            flash_attn_varlen_func = cute_flash_attn_varlen_func
+            bias_kwargs = {
+                "score_mod": _get_score_mod(rel_extent),
+                "aux_tensors": [rel_logits],
+            }
+
+        ret = flash_attn_varlen_func(
+            q=q,
+            k=key_cache,
+            v=value_cache,
+            cu_seqlens_q=cu_seqlens_q,
+            seqused_k=cache_seqlens,
+            max_seqlen_q=max_seqlen_q,
+            page_table=block_table,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=cute_window,
+            num_splits=num_splits,
+            return_lse=False,
+            out=out,
+            **bias_kwargs,
+        )
+        if isinstance(ret, tuple):
+            return ret[0]
+        return ret
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        is_local: bool,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        rel_extent: int,
+        dtype: torch.dtype,
+        kv_dtype: torch.dtype,
+        block_size: int,
+        window_size: tuple[int, int],
+        max_kv_len: int,
+        query_len: int,
+        num_reqs: int,
+    ) -> CompileKey:
+        max_seqlen_q = bucket_max_seqlen_q(query_len)
+        num_splits = inkling_fa4_num_splits(
+            is_local=is_local,
+            batch_size=num_reqs,
+            max_query_len=max_seqlen_q,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            max_kv_len=max_kv_len,
+        )
+        return self.CompileKey(
+            is_local=is_local,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            rel_extent=rel_extent,
+            dtype=dtype,
+            kv_dtype=kv_dtype,
+            block_size=block_size,
+            window_size=window_size,
+            max_seqlen_q=max_seqlen_q,
+            num_splits=num_splits,
+            num_warps_bucket=(_num_warps_bucket(num_reqs) if num_splits > 1 else None),
+            large_num_reqs=num_reqs > 1024,
         )
 
-        flash_attn_varlen_func = tml_flash_attn_varlen_func
-        bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
-    else:
-        from vllm.vllm_flash_attn.cute import (
-            flash_attn_varlen_func as cute_flash_attn_varlen_func,
+    def _is_valid_warmup_dispatch(
+        self,
+        *,
+        query_len: int,
+        num_reqs: int,
+        max_num_batched_tokens: int,
+    ) -> bool:
+        return query_len + num_reqs <= max_num_batched_tokens + 1
+
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
+        max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        if max_num_reqs <= 0 or max_num_batched_tokens <= 0:
+            return []
+
+        hf_config = vllm_config.model_config.hf_config
+        get_text_config = getattr(hf_config, "get_text_config", None)
+        config = get_text_config() if callable(get_text_config) else hf_config
+
+        tp_size = get_tensor_model_parallel_world_size()
+        dtype = vllm_config.model_config.dtype
+        kv_dtype = kv_cache_dtype_str_to_dtype(
+            vllm_config.cache_config.cache_dtype,
+            vllm_config.model_config,
+        )
+        block_size = vllm_config.cache_config.block_size
+        local_extent = config.sliding_window_size
+
+        global_num_kv_heads = config.num_key_value_heads
+        local_num_kv_heads = config.swa_num_key_value_heads
+        assert config.num_attention_heads % tp_size == 0
+        assert config.swa_num_attention_heads % tp_size == 0
+        if global_num_kv_heads >= tp_size:
+            assert global_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % global_num_kv_heads == 0
+        if local_num_kv_heads >= tp_size:
+            assert local_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % local_num_kv_heads == 0
+
+        return self._trace_dispatch(self.dispatch)(
+            zip_inputs(
+                dict(
+                    is_local=False,
+                    num_heads=config.num_attention_heads // tp_size,
+                    num_kv_heads=max(1, global_num_kv_heads // tp_size),
+                    head_dim=config.head_dim,
+                    rel_extent=config.rel_extent,
+                    dtype=dtype,
+                    kv_dtype=kv_dtype,
+                    block_size=block_size,
+                    window_size=(-1, -1),
+                    max_kv_len=vllm_config.model_config.max_model_len,
+                ),
+                dict(
+                    is_local=True,
+                    num_heads=config.swa_num_attention_heads // tp_size,
+                    num_kv_heads=max(1, local_num_kv_heads // tp_size),
+                    head_dim=config.swa_head_dim,
+                    rel_extent=local_extent,
+                    dtype=dtype,
+                    kv_dtype=kv_dtype,
+                    block_size=block_size,
+                    window_size=(local_extent - 1, 0),
+                    max_kv_len=local_extent,
+                ),
+            ),
+            query_len=WarmupIntRange(
+                1,
+                max_num_batched_tokens + 1,
+                advance=lambda value: bucket_max_seqlen_q(value) + 1,
+            ),
+            num_reqs=WarmupIntRange(1, max_num_reqs + 1),
+            max_num_batched_tokens=max_num_batched_tokens,
+            _when=self._is_valid_warmup_dispatch,
         )
 
-        flash_attn_varlen_func = cute_flash_attn_varlen_func
-        bias_kwargs = {
-            "score_mod": _get_score_mod(rel_extent),
-            "aux_tensors": [rel_logits],
-        }
+    def compile(self, compile_key: CompileKey) -> None:
+        from torch._subclasses.fake_tensor import FakeTensorMode
 
-    ret = flash_attn_varlen_func(
-        q=q,
-        k=key_cache,
-        v=value_cache,
-        cu_seqlens_q=cu_seqlens_q,
-        seqused_k=cache_seqlens,
-        max_seqlen_q=max_seqlen_q,
-        page_table=block_table,
-        softmax_scale=softmax_scale,
-        causal=causal,
-        window_size=cute_window,
-        num_splits=num_splits,
-        return_lse=False,
-        out=out,
-        **bias_kwargs,
-    )
-    if isinstance(ret, tuple):
-        return ret[0]
-    return ret
+        if compile_key.num_splits == 1 or compile_key.num_warps_bucket is None:
+            num_reqs = 1025 if compile_key.large_num_reqs else 1
+        else:
+            min_num_warps = (
+                1
+                if compile_key.num_warps_bucket == 1
+                else compile_key.num_warps_bucket // 2 + 1
+            )
+            min_num_reqs = (min_num_warps - 1) * 31 + 1
+            num_reqs = max(
+                min_num_reqs,
+                1025 if compile_key.large_num_reqs else 1,
+            )
+
+        with FakeTensorMode():
+            device = torch.accelerator.current_accelerator()
+            total_q = compile_key.max_seqlen_q + num_reqs - 1
+            q = torch.empty(
+                total_q,
+                compile_key.num_heads,
+                compile_key.head_dim,
+                dtype=compile_key.dtype,
+                device=device,
+            )
+            kv = torch.empty(
+                1,
+                2,
+                compile_key.block_size,
+                compile_key.num_kv_heads,
+                compile_key.head_dim,
+                dtype=compile_key.kv_dtype,
+                device=device,
+            )
+            key_cache, value_cache = kv.unbind(1)
+            self.kernel(
+                q,
+                key_cache,
+                value_cache,
+                block_table=torch.empty(
+                    num_reqs,
+                    1,
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                cache_seqlens=torch.empty(
+                    num_reqs,
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                cu_seqlens_q=torch.empty(
+                    num_reqs + 1,
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                max_seqlen_q=compile_key.max_seqlen_q,
+                softmax_scale=compile_key.head_dim**-1,
+                causal=True,
+                window_size=compile_key.window_size,
+                rel_extent=compile_key.rel_extent,
+                rel_logits=torch.empty(
+                    total_q,
+                    compile_key.num_heads,
+                    compile_key.rel_extent,
+                    dtype=compile_key.dtype,
+                    device=device,
+                ),
+                num_splits=compile_key.num_splits,
+                out=torch.empty_like(q),
+            )
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        *,
+        block_table: torch.Tensor,
+        cache_seqlens: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        max_seqlen_q: int,
+        softmax_scale: float,
+        causal: bool,
+        window_size: tuple[int, int],
+        rel_extent: int,
+        rel_logits: torch.Tensor,
+        num_splits: int = 32,
+        out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.kernel(
+            q,
+            key_cache,
+            value_cache,
+            block_table=block_table,
+            cache_seqlens=cache_seqlens,
+            cu_seqlens_q=cu_seqlens_q,
+            max_seqlen_q=max_seqlen_q,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            rel_extent=rel_extent,
+            rel_logits=rel_logits,
+            num_splits=num_splits,
+            out=out,
+        )
+
+
+INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel()

vllm/models/inkling/nvidia/ops/fa4_warmup.py

@@ -1,154 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Startup compilation of the FA4 kernels used by Inkling."""
-
-from __future__ import annotations
-
-from collections.abc import Iterator
-from dataclasses import dataclass
-from functools import partial
-
-import torch
-
-from vllm.model_executor.warmup.cutedsl_warmup import (
-    CuTeDSLCompileUnit,
-    register_cutedsl_warmup_provider,
-)
-
-from .fa4_rel_attention import (
-    bucket_max_seqlen_q,
-    inkling_fa4_num_splits,
-    inkling_fa4_rel_attention,
-)
-
-
-@dataclass(frozen=True)
-class InklingFA4WarmupConfig:
-    num_heads: int
-    num_kv_heads: int
-    head_dim: int
-    rel_extent: int
-    window_size: tuple[int, int]
-    is_local: bool
-    max_kv_len: int
-    dtype: torch.dtype
-    kv_dtype: torch.dtype
-    block_size: int
-    max_num_reqs: int
-    max_num_batched_tokens: int
-
-
-def _num_warps_bucket(num_reqs: int) -> int:
-    num_warps = min((num_reqs + 30) // 31, 32)
-    return 1 << (num_warps - 1).bit_length()
-
-
-def _compile(config: InklingFA4WarmupConfig, max_seqlen_q: int, num_reqs: int) -> None:
-    from torch._subclasses.fake_tensor import FakeTensorMode
-
-    num_splits = inkling_fa4_num_splits(
-        is_local=config.is_local,
-        batch_size=num_reqs,
-        max_query_len=max_seqlen_q,
-        num_heads=config.num_heads,
-        num_kv_heads=config.num_kv_heads,
-        max_kv_len=config.max_kv_len,
-    )
-    with FakeTensorMode():
-        device = torch.accelerator.current_accelerator()
-        total_q = max_seqlen_q + num_reqs - 1
-        q = torch.empty(
-            total_q,
-            config.num_heads,
-            config.head_dim,
-            dtype=config.dtype,
-            device=device,
-        )
-        kv = torch.empty(
-            1,
-            2,
-            config.block_size,
-            config.num_kv_heads,
-            config.head_dim,
-            dtype=config.kv_dtype,
-            device=device,
-        )
-        key_cache, value_cache = kv.unbind(1)
-        inkling_fa4_rel_attention(
-            q,
-            key_cache,
-            value_cache,
-            block_table=torch.empty(num_reqs, 1, dtype=torch.int32, device=device),
-            cache_seqlens=torch.empty(num_reqs, dtype=torch.int32, device=device),
-            cu_seqlens_q=torch.empty(num_reqs + 1, dtype=torch.int32, device=device),
-            max_seqlen_q=max_seqlen_q,
-            softmax_scale=1.0 / config.head_dim,
-            causal=True,
-            window_size=config.window_size,
-            rel_extent=config.rel_extent,
-            rel_logits=torch.empty(
-                total_q,
-                config.num_heads,
-                config.rel_extent,
-                dtype=config.dtype,
-                device=device,
-            ),
-            num_splits=num_splits,
-            out=torch.empty_like(q),
-        )
-
-
-def _iter_compile_units(
-    config: InklingFA4WarmupConfig,
-) -> Iterator[CuTeDSLCompileUnit]:
-    max_bucket = bucket_max_seqlen_q(config.max_num_batched_tokens)
-    max_seqlen_q = 1
-    while max_seqlen_q <= max_bucket:
-        min_query_len = max_seqlen_q // 2 + 1
-        max_num_reqs = min(
-            config.max_num_reqs,
-            config.max_num_batched_tokens - min_query_len + 1,
-        )
-        seen: set[tuple[int, int, int | None, bool]] = set()
-        for num_reqs in range(1, max_num_reqs + 1):
-            num_splits = inkling_fa4_num_splits(
-                is_local=config.is_local,
-                batch_size=num_reqs,
-                max_query_len=max_seqlen_q,
-                num_heads=config.num_heads,
-                num_kv_heads=config.num_kv_heads,
-                max_kv_len=config.max_kv_len,
-            )
-            key = (
-                max_seqlen_q,
-                num_splits,
-                _num_warps_bucket(num_reqs) if num_splits > 1 else None,
-                num_reqs > 1024,
-            )
-            if key in seen:
-                continue
-            seen.add(key)
-            yield CuTeDSLCompileUnit(
-                name="inkling_fa4",
-                key=("inkling_fa4", config, key),
-                compile=partial(_compile, config, max_seqlen_q, num_reqs),
-            )
-        max_seqlen_q *= 2
-
-
-class _WarmupProvider:
-    def __init__(self) -> None:
-        self.configs: set[InklingFA4WarmupConfig] = set()
-
-    def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
-        return tuple(
-            unit for config in self.configs for unit in _iter_compile_units(config)
-        )
-
-
-_PROVIDER = _WarmupProvider()
-
-
-def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None:
-    _PROVIDER.configs.add(config)
-    register_cutedsl_warmup_provider(_PROVIDER)

vllm/v1/attention/backends/mla/prefill/flash_attn.py

@@ -13,7 +13,7 @@ import vllm.envs as envs
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticTensorSym,
 )
-from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
+from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.fa_utils import (
     compile_flash_attn_varlen_func_from_specs,
@@ -57,12 +57,6 @@ FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True)
 FA4_MLA_PREFILL_LSE_OPTIONS = (False, True)
 
 
-@dataclass(frozen=True)
-class _FA4MLAPrefillShapeProbe:
-    max_seqlen_q: int
-    max_seqlen_k: int
-
-
 class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
     """FA4 MLA prefill compile-key wrapper used by generic JIT warmup."""
 
@@ -92,6 +86,16 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
         return_softmax_lse: bool = False
         num_splits: int = 0
 
+    @staticmethod
+    def kernel(
+        *args: Any,
+        runtime_kernel: Callable[..., Any] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        kernel = flash_attn_varlen_func if runtime_kernel is None else runtime_kernel
+        assert kernel is not None
+        return kernel(*args, **kwargs)
+
     def dispatch(  # type: ignore[override]
         self,
         *,
@@ -99,7 +103,8 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
         dtype: torch.dtype,
         num_heads: int,
         mla_dims: "MLADims",
-        shape_probe: _FA4MLAPrefillShapeProbe,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
         requires_v_padding: bool,
         causal: bool,
         fa_version: int,
@@ -109,17 +114,17 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
     ) -> CompileKey:
         return self.CompileKey(
             q_shape=(
-                batch_size * shape_probe.max_seqlen_q,
+                batch_size * max_seqlen_q,
                 num_heads,
                 mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
             ),
             k_shape=(
-                batch_size * shape_probe.max_seqlen_k,
+                batch_size * max_seqlen_k,
                 num_heads,
                 mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
             ),
             v_shape=(
-                batch_size * shape_probe.max_seqlen_k,
+                batch_size * max_seqlen_k,
                 num_heads,
                 (
                     mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
@@ -128,8 +133,8 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
                 ),
             ),
             q_dtype=dtype,
-            max_seqlen_q=shape_probe.max_seqlen_q,
-            max_seqlen_k=shape_probe.max_seqlen_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
             softmax_scale=(mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim)
             ** -0.5,
             causal=causal,
@@ -150,6 +155,28 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
             num_splits=num_splits,
         )
 
+    def _is_valid_warmup_shape_probe(
+        self,
+        *,
+        max_seqlen_k: int,
+        num_splits: int,
+        is_sm90: bool,
+        qk_head_dim: int,
+        effective_v_head_dim: int,
+    ) -> bool:
+        long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
+        very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
+        return (
+            max_seqlen_k < long_k
+            or (num_splits != 1 and max_seqlen_k == long_k)
+            or (
+                num_splits != 1
+                and max_seqlen_k == very_long_k
+                and not is_sm90
+                and qk_head_dim != effective_v_head_dim
+            )
+        )
+
     def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]:
         from vllm.model_executor.layers.attention.mla_attention import (
             get_mla_dims,
@@ -185,58 +212,49 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
         effective_v_head_dim = (
             qk_head_dim if requires_v_padding else mla_dims.v_head_dim
         )
-        shape_probes = [
-            _FA4MLAPrefillShapeProbe(1, FA4_MLA_PREFILL_K_TILE),
-            _FA4MLAPrefillShapeProbe(
-                FA4_MLA_PREFILL_Q_TILE + 1,
-                4 * FA4_MLA_PREFILL_K_TILE,
+        shape_probes = (
+            dict(max_seqlen_q=1, max_seqlen_k=FA4_MLA_PREFILL_K_TILE),
+            dict(
+                max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1,
+                max_seqlen_k=4 * FA4_MLA_PREFILL_K_TILE,
             ),
-        ]
-        if num_splits != 1:
-            shape_probes.extend(
-                (
-                    _FA4MLAPrefillShapeProbe(
-                        1,
-                        FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
-                    ),
-                    _FA4MLAPrefillShapeProbe(
-                        FA4_MLA_PREFILL_Q_TILE + 1,
-                        FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
-                    ),
-                )
-            )
-            if not is_sm90 and qk_head_dim != effective_v_head_dim:
-                shape_probes.extend(
-                    (
-                        _FA4MLAPrefillShapeProbe(
-                            1,
-                            FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
-                        ),
-                        _FA4MLAPrefillShapeProbe(
-                            FA4_MLA_PREFILL_Q_TILE + 1,
-                            FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
-                        ),
-                    )
-                )
+            dict(
+                max_seqlen_q=1,
+                max_seqlen_k=FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
+            ),
+            dict(
+                max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1,
+                max_seqlen_k=FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
+            ),
+            dict(
+                max_seqlen_q=1,
+                max_seqlen_k=FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS
+                * FA4_MLA_PREFILL_K_TILE,
+            ),
+            dict(
+                max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1,
+                max_seqlen_k=FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS
+                * FA4_MLA_PREFILL_K_TILE,
+            ),
+        )
 
         return self._trace_dispatch(self.dispatch)(
+            zip_inputs(*shape_probes),
             batch_size=FA4_MLA_PREFILL_COMPILE_BATCH_SIZE,
             dtype=dtype,
             num_heads=num_heads,
             mla_dims=mla_dims,
-            shape_probe=shape_probes,
             requires_v_padding=requires_v_padding,
+            is_sm90=is_sm90,
+            qk_head_dim=qk_head_dim,
+            effective_v_head_dim=effective_v_head_dim,
             causal=FA4_MLA_PREFILL_CAUSAL_OPTIONS,
             return_lse=FA4_MLA_PREFILL_LSE_OPTIONS,
             num_splits=num_splits,
             fa_version=fa_version,
+            _when=self._is_valid_warmup_shape_probe,
         )
 
-    @staticmethod
-    def kernel(*args: Any, **kwargs: Any) -> Any:
-        assert flash_attn_varlen_func is not None
-        return flash_attn_varlen_func(*args, **kwargs)
-
     def compile(self, compile_key: CompileKey) -> None:
         assert compile_flash_attn_varlen_func_from_specs is not None
         window_size = (
@@ -271,11 +289,13 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
         runtime_kernel: Callable[..., Any] | None = None,
         **kwargs: Any,
     ) -> Any:
-        kernel = self.kernel if runtime_kernel is None else runtime_kernel
-        return kernel(q=q, k=k, v=v, **kwargs)
-
-
-FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()
+        return self.kernel(
+            q=q,
+            k=k,
+            v=v,
+            runtime_kernel=runtime_kernel,
+            **kwargs,
+        )
 
 
 class FlashAttnPrefillBackend(MLAPrefillBackend):
@@ -470,3 +490,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
             causal=False,  # Context is unmasked
             return_softmax_lse=True,
         )
+
+
+FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()