[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