[6/N][warmup][DSv4] Migrate sampling, and DFlash JIT kernels (#56323)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

.buildkite/test_areas/models_basic.yaml

@@ -41,7 +41,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Basic Models (Other)"
   device: h200_35gb
   key: basic-models-tests-other
-  timeout_in_minutes: 35
+  timeout_in_minutes: 38
   source_file_dependencies:
   - vllm/
   - "!vllm/distributed/kv_transfer/"

.buildkite/test_areas/samplers.yaml

@@ -5,7 +5,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Samplers"
   device: h200_35gb
   key: samplers-test
-  timeout_in_minutes: 40
+  timeout_in_minutes: 60
   source_file_dependencies:
   - vllm/model_executor/layers
   - vllm/sampling_metadata.py

tests/model_executor/test_jit_warmup.py

@@ -171,6 +171,19 @@ def test_warmup_range_validates_custom_advancement() -> None:
         )
 
 
+def test_warmup_cases_support_lambda_advancement() -> None:
+    def warmup_inputs() -> dict[str, Any]:
+        value: Any = WarmupIntRange(1, 9, advance=lambda value: value * 2)
+        return dict(value=value)
+
+    assert list(ToyKernel()._expand_warmup_cases(warmup_inputs)) == [
+        {"value": 1},
+        {"value": 2},
+        {"value": 4},
+        {"value": 8},
+    ]
+
+
 def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
     cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
 
@@ -577,6 +590,25 @@ def test_registry_records_only_inside_model_setup_context() -> None:
     assert kernel.compiled == []
 
 
+def test_registry_capture() -> None:
+    class Owner:
+        @JitWarmupRegistry.capture
+        def __init__(
+            self,
+            vllm_config: Any,
+            kernel: RecordingToyKernel,
+        ) -> None:
+            kernel.register_warmup(3, vllm_config)
+
+    kernel = RecordingToyKernel()
+    config = _config()
+    owner = Owner(config, kernel)
+
+    assert len(owner.jit_warmup_registry) == 1  # type: ignore[attr-defined]
+    kernel.register_warmup(5, config)
+    assert len(owner.jit_warmup_registry) == 1  # type: ignore[attr-defined]
+
+
 def test_registry_expands_requests_and_deduplicates_owner_keys() -> None:
     registry = JitWarmupRegistry(_config())
     kernel = RecordingToyKernel()

tests/model_executor/test_jit_warmup_triton_launcher.py

@@ -4,22 +4,68 @@
 from dataclasses import dataclass
 from typing import Any
 
-from vllm.model_executor.warmup.jit_warmup import kernel_launcher
+import pytest
+import torch
+
+from vllm.model_executor.warmup import jit_warmup_triton_helper
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupChoices,
+    WarmupIntRange,
+    kernel_launcher,
+)
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
     LaunchSpec,
+    TritonJitKey,
+    TritonWarmupTensor,
     VllmTritonJitKernel,
+    triton_kernel_dispatcher_with_warmup,
+    triton_warmup_inputs,
 )
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _binder_test_kernel(x, value, BLOCK: tl.constexpr):
+    pass
+
+
+@triton.jit
+def _second_binder_test_kernel(x, value, BLOCK: tl.constexpr):
+    pass
+
+
+@triton.jit
+def _pointer_group_test_kernel(x_ptr, value, BLOCK: tl.constexpr):
+    pass
 
 
 class _FakeTritonKernel:
-    arg_names = ("first", "second", "CONST")
+    arg_names: tuple[str, ...] = ("first", "second", "CONST")
 
     def __init__(self) -> None:
         self.warmup_calls: list[dict[str, Any]] = []
+        self.runtime_calls: list[tuple[Any, tuple[Any, ...], dict[str, Any]]] = []
 
     def warmup(self, **kwargs: Any) -> None:
         self.warmup_calls.append(kwargs)
 
+    def __getitem__(self, grid: Any) -> Any:
+        def launch(*args: Any, **kwargs: Any) -> None:
+            self.runtime_calls.append((grid, args, kwargs))
+
+        return launch
+
+
+def _patch_key_deriver(
+    monkeypatch: pytest.MonkeyPatch,
+    derive: Any,
+) -> None:
+    monkeypatch.setattr(
+        jit_warmup_triton_helper,
+        "_triton_key_deriver",
+        lambda kernel: lambda kwargs: derive(kernel, kwargs),
+    )
+
 
 class _TestTritonKernel(VllmTritonJitKernel["_TestTritonKernel.CompileKey"]):
     kernel = _FakeTritonKernel()
@@ -91,6 +137,291 @@ def test_triton_launcher_supports_cpu_function_wrappers() -> None:
     assert calls == [("runtime", 2, 7)]
 
 
+def test_triton_kernel_decorator_returns_launcher(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+    kernel.arg_names = (
+        "FIRST_PTR",
+        "FIRST_STRIDE0",
+        "aliased_ptr",
+        "aliased_stride_0",
+        "SECOND",
+        "CONST",
+    )
+
+    def warmup_inputs() -> dict[str, Any]:
+        return dict(
+            first=TritonWarmupTensor(torch.float32, shape=(2, 3), strides=(5, 1)),
+            second=WarmupChoices(1, 2),
+        )
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def launch(first: str, second: int = 2, config: int = 7) -> LaunchSpec:
+        return (2,), dict(aliased_ptr=first, CONST=config)
+
+    def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]:
+        return {TritonJitKey(id(kernel), "fake", 0, kwargs["SECOND"])}
+
+    _patch_key_deriver(monkeypatch, fake_keys)
+
+    keys = launch.get_warmup_keys()
+    keys_with_config = launch.get_warmup_keys(vllm_config=object())
+    assert keys_with_config == keys
+    assert [dict(key.inputs)["second"] for key in keys] == [1, 2]
+    launch.compile(keys[0])
+    assert kernel.warmup_calls == [
+        {
+            "grid": (1,),
+            "FIRST_PTR": TritonWarmupTensor(
+                torch.float32, shape=(2, 3), strides=(5, 1)
+            ),
+            "FIRST_STRIDE0": 5,
+            "aliased_ptr": TritonWarmupTensor(
+                torch.float32, shape=(2, 3), strides=(5, 1)
+            ),
+            "aliased_stride_0": 5,
+            "SECOND": 1,
+            "CONST": 7,
+        }
+    ]
+    first = torch.empty_strided((2, 3), (5, 1))
+    launch(first)
+    assert kernel.runtime_calls == [
+        (
+            (2,),
+            (),
+            {
+                "FIRST_PTR": first,
+                "FIRST_STRIDE0": 5,
+                "aliased_ptr": first,
+                "aliased_stride_0": 5,
+                "SECOND": 2,
+                "CONST": 7,
+            },
+        )
+    ]
+    assert launch.__name__ == "launch"
+    with pytest.raises(TypeError, match="unexpected keyword"):
+        launch(first, 1, 7, stale_constexpr=True)
+
+
+def test_triton_kernel_decorator_exhausts_large_ranges_before_deduplication(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+    dispatched: list[int] = []
+
+    def warmup_inputs() -> dict[str, Any]:
+        tokens: Any = WarmupIntRange(1, 8193)
+        return dict(first="warmup", second=tokens)
+
+    def specialization(second: int) -> int:
+        return 7 if second <= 97 else 8
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def dispatch(first: str, second: int) -> LaunchSpec:
+        dispatched.append(second)
+        return (second,), dict(CONST=specialization(second))
+
+    def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]:
+        return {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])}
+
+    _patch_key_deriver(monkeypatch, fake_keys)
+
+    keys = dispatch.get_warmup_keys()
+    assert dispatched == list(range(1, 8193))
+    assert len(keys) == 2
+    assert {specialization(dict(key.inputs)["second"]) for key in keys} == {7, 8}
+
+
+def test_triton_kernel_decorator_propagates_dispatch_assertions(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+
+    def warmup_inputs() -> dict[str, Any]:
+        return dict(first="warmup", second=WarmupChoices(1, 2))
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def dispatch(first: str, second: int) -> LaunchSpec:
+        assert second != 2, "broken dispatch"
+        return (1,), dict(CONST=second)
+
+    _patch_key_deriver(
+        monkeypatch,
+        lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])},
+    )
+
+    with pytest.raises(AssertionError, match="broken dispatch"):
+        dispatch.get_warmup_keys()
+
+
+def test_triton_kernel_dispatch_uses_cuda_fake_tensors(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+
+    def warmup_inputs() -> dict[str, Any]:
+        return dict(
+            first=TritonWarmupTensor(torch.float32, shape=(2, 3)),
+            second=1,
+        )
+
+    @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
+    def dispatch(first: torch.Tensor, second: int) -> LaunchSpec:
+        assert isinstance(first, torch.Tensor)
+        assert first.is_cuda
+        assert first[0].is_contiguous()
+        return (first.shape[0],), dict(CONST=first[0].numel())
+
+    _patch_key_deriver(
+        monkeypatch,
+        lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])},
+    )
+
+    keys = dispatch.get_warmup_keys()
+    assert len(keys) == 1
+    assert dict(keys[0].inputs)["first"].device.type == "cuda"
+
+
+def test_triton_kernel_decorates_native_launchers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _FakeTritonKernel()
+
+    def warmup_inputs(kernel: Any = kernel) -> dict[str, Any]:
+        return triton_warmup_inputs(
+            kernel,
+            "warmup",
+            WarmupChoices(1, 2),
+            grid=(2,),
+            CONST=7,
+        )
+
+    launcher = triton_kernel_dispatcher_with_warmup(warmup_inputs=warmup_inputs)(kernel)
+    launcher[(3,)]("runtime", 4, CONST=8)
+
+    assert kernel.runtime_calls == [((3,), ("runtime", 4), {"CONST": 8})]
+
+    _patch_key_deriver(
+        monkeypatch,
+        lambda kernel, kwargs: {
+            TritonJitKey(id(kernel), "fake", 0, (kwargs["second"], kwargs["CONST"]))
+        },
+    )
+
+    keys = launcher.get_warmup_keys()
+    assert len(keys) == 2
+    launcher.compile(keys[0])
+    assert kernel.warmup_calls == [
+        {"grid": (1,), "first": "warmup", "second": 1, "CONST": 7}
+    ]
+
+
+def test_triton_warmup_inputs_expands_explicit_pointer_dtypes() -> None:
+    inputs = triton_warmup_inputs(
+        _pointer_group_test_kernel,
+        grid=(1,),
+        pointer_dtypes={torch.float32: ("x_ptr",)},
+        value=2,
+        BLOCK=16,
+    )
+
+    assert inputs["x_ptr"] == TritonWarmupTensor(torch.float32)
+    assert inputs["value"] == 2
+
+    with pytest.raises(ValueError, match="not a pointer: value"):
+        triton_warmup_inputs(
+            _pointer_group_test_kernel,
+            grid=(1,),
+            pointer_dtypes={torch.float32: ("value",)},
+            x_ptr=TritonWarmupTensor(torch.float32),
+            BLOCK=16,
+        )
+    with pytest.raises(ValueError, match="Missing Triton pointer inputs: x_ptr"):
+        triton_warmup_inputs(
+            _pointer_group_test_kernel,
+            grid=(1,),
+            pointer_dtypes={},
+            value=2,
+            BLOCK=16,
+        )
+
+
+def test_triton_key_derivation_applies_wrappers_and_runtime_options(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    knobs = triton.knobs
+    driver = triton.runtime.driver
+
+    device = "test-device"
+    monkeypatch.setattr(driver.active, "get_current_device", lambda: device)
+    calls: list[dict[str, Any]] = []
+
+    def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any:
+        calls.append(dict(BLOCK=BLOCK, **options))
+        specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)]
+        return {}, specialization, options
+
+    cache: dict[Any, Any] = {}
+    _binder_test_kernel.device_caches[device] = ({}, cache, None, None, binder)
+    heuristic_kernel = triton.heuristics(
+        {"BLOCK": lambda args: 16 if args["value"] <= 16 else 32}
+    )(_binder_test_kernel)
+
+    keys = jit_warmup_triton_helper._triton_key_deriver(heuristic_kernel)(
+        {"x": TritonWarmupTensor(tl.float32), "value": 2}
+    )
+
+    assert len(keys) == 1
+    expected_options = {
+        "debug": knobs.runtime.debug,
+        "instrumentation_mode": knobs.compilation.instrumentation_mode,
+    }
+    if hasattr(knobs.compilation, "fpsan_homomorphic_casts"):
+        expected_options["fpsan_homomorphic_casts"] = (
+            knobs.compilation.fpsan_homomorphic_casts
+        )
+    assert calls == [{"BLOCK": 16, **expected_options}]
+
+
+def test_triton_key_derivation_covers_autotune_configs_and_jit_identity(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    driver = triton.runtime.driver
+
+    device = "test-device"
+    monkeypatch.setattr(driver.active, "get_current_device", lambda: device)
+
+    def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any:
+        specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)]
+        return {}, specialization, options
+
+    for kernel in (_binder_test_kernel, _second_binder_test_kernel):
+        kernel.device_caches[device] = ({}, {}, None, None, binder)
+    autotuned = triton.autotune(
+        configs=[
+            triton.Config({"BLOCK": 16}, num_warps=2),
+            triton.Config({"BLOCK": 32}, num_warps=4),
+        ],
+        key=["value"],
+    )(_binder_test_kernel)
+    inputs = {"x": TritonWarmupTensor(tl.float32), "value": 2}
+
+    autotune_keys = jit_warmup_triton_helper._triton_key_deriver(autotuned)(inputs)
+    other_kernel_keys = jit_warmup_triton_helper._triton_key_deriver(
+        _second_binder_test_kernel
+    )(inputs | {"BLOCK": 16})
+
+    assert len(autotune_keys) == 2
+    assert {key.jit_function_key for key in autotune_keys} == {
+        _binder_test_kernel.cache_key
+    }
+    assert {key.jit_function_id for key in autotune_keys} == {id(_binder_test_kernel)}
+    assert autotune_keys.isdisjoint(other_kernel_keys)
+
+
 def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:
     from vllm.v1.attention.backends.utils import PAD_SLOT_ID
     from vllm.v1.worker.block_table import ComputeSlotMappingKernel

tests/model_executor/test_spec_decode_rejection_warmup.py

@@ -1,121 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm.config.watermarking import WatermarkConfig, derive_watermark_key
-from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
-    spec_decode_rejection_warmup,
-)
-
-MASTER_KEY = 20250910
-
-
-def _worker(
-    watermark_config: WatermarkConfig | None,
-    draft_sample_method: str = "greedy",
-) -> SimpleNamespace:
-    return SimpleNamespace(
-        device=torch.device("cpu"),
-        vllm_config=SimpleNamespace(
-            speculative_config=SimpleNamespace(
-                num_speculative_tokens=2,
-                rejection_sample_method="standard",
-                draft_sample_method=draft_sample_method,
-            ),
-            model_config=SimpleNamespace(
-                get_vocab_size=lambda: 128, dtype=torch.bfloat16
-            ),
-            watermark_config=watermark_config,
-        ),
-    )
-
-
-def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
-    calls: list[dict] = []
-
-    def fake_rejection_sample(**kwargs):
-        calls.append(kwargs)
-        return None, None
-
-    monkeypatch.setattr(
-        "vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils.rejection_sample",
-        fake_rejection_sample,
-    )
-    return calls
-
-
-def test_warmup_token_id_dtypes_match_the_runtime_buffers(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    calls = _capture_calls(monkeypatch)
-    spec_decode_rejection_warmup(_worker(None))
-
-    assert calls
-    for kwargs in calls:
-        assert kwargs["draft_sampled"].dtype == torch.int32
-        assert kwargs["pos"].dtype == torch.int64
-
-
-def test_warmup_index_mapping_dtypes_match_the_runtime_buffers(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    calls = _capture_calls(monkeypatch)
-    spec_decode_rejection_warmup(_worker(None))
-
-    assert calls
-    for kwargs in calls:
-        assert kwargs["idx_mapping"].dtype == torch.int64
-        assert kwargs["expanded_idx_mapping"].dtype == torch.int64
-        assert kwargs["expanded_local_pos"].dtype == torch.int32
-
-
-def test_warmup_watermark_args_match_the_runtime_kernel_arguments(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    watermark_config = WatermarkConfig(
-        key=MASTER_KEY, algorithm="dual_key_gumbel", context_width=3
-    )
-    calls = _capture_calls(monkeypatch)
-    spec_decode_rejection_warmup(
-        _worker(watermark_config, draft_sample_method="probabilistic")
-    )
-
-    assert calls
-    key_b = derive_watermark_key(MASTER_KEY, b"key_b")
-    for kwargs in calls:
-        contexts = kwargs["contexts"]
-        assert contexts.dtype == torch.int32
-        assert contexts.shape == (3, 3) and contexts.stride() == (3, 1)
-        assert kwargs["watermarking"].dtype == torch.bool
-        assert kwargs["watermark_key"] == key_b
-        assert key_b & 0xFFFFFFFF >= 2**31 and key_b >> 32 >= 2**31
-
-
-def test_warmup_is_skipped_when_the_recovery_key_cannot_be_derived(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    calls = _capture_calls(monkeypatch)
-    monkeypatch.setattr(
-        "vllm.v1.watermarking.spec_decode.speculative_target_watermark_key",
-        lambda _: (_ for _ in ()).throw(NotImplementedError("no philox")),
-    )
-    spec_decode_rejection_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))
-
-    assert calls == []
-
-
-def test_warmup_draft_logits_match_the_runtime_draft_sample_method(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    calls = _capture_calls(monkeypatch)
-    spec_decode_rejection_warmup(_worker(None))
-    assert calls
-    assert all(kwargs["draft_logits"] is None for kwargs in calls)
-
-    calls = _capture_calls(monkeypatch)
-    spec_decode_rejection_warmup(_worker(None, draft_sample_method="probabilistic"))
-    assert calls
-    assert all(kwargs["draft_logits"] is not None for kwargs in calls)

tests/v1/sample/test_topk_topp_sampler.py

@@ -79,6 +79,10 @@ def test_rocm_aiter_sampler_defers_import_when_generators_force_native(
         def is_cuda():
             return False
 
+        @staticmethod
+        def is_cuda_alike():
+            return True
+
         @staticmethod
         def is_cpu():
             return False

tests/v1/shutdown/test_delete.py

@@ -11,6 +11,7 @@ from tests.v1.shutdown.utils import (
     SHUTDOWN_TEST_TIMEOUT_SEC,
 )
 from vllm import LLM, SamplingParams
+from vllm.config import KernelConfig
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.platforms import current_platform
 from vllm.sampling_params import RequestOutputKind
@@ -39,7 +40,10 @@ async def test_async_llm_delete(
         pytest.skip(reason="Not enough CUDA devices")
 
     engine_args = AsyncEngineArgs(
-        model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size
+        model=model,
+        enforce_eager=True,
+        tensor_parallel_size=tensor_parallel_size,
+        kernel_config=KernelConfig(enable_jit_warmup=False),
     )
 
     # Instantiate AsyncLLM; make request to complete any deferred
@@ -94,7 +98,10 @@ def test_llm_delete(
         # Instantiate LLM; make request to complete any deferred
         # initialization; then delete instance
         llm = LLM(
-            model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size
+            model=model,
+            enforce_eager=True,
+            tensor_parallel_size=tensor_parallel_size,
+            kernel_config=KernelConfig(enable_jit_warmup=False),
         )
         if send_one_request:
             llm.generate(
@@ -122,7 +129,10 @@ def test_llm_delete_inprocess(
     with monkeypatch.context() as m:
         m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
 
-        with VllmRunner(model) as vllm_model:
+        with VllmRunner(
+            model,
+            kernel_config=KernelConfig(enable_jit_warmup=False),
+        ) as vllm_model:
             if send_one_request:
                 vllm_model.generate(
                     ["Hello my name is"],

tests/v1/spec_decode/test_dflash_prepare_inputs.py

@@ -43,7 +43,7 @@ def _run_prepare(
         num_scheduled_tokens=np.array([4], dtype=np.int32),
         positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
         query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
-        idx_mapping=torch.tensor([2], dtype=torch.int32, device=device),
+        idx_mapping=torch.tensor([2], dtype=torch.int64, device=device),
     )
     query_slot_mapping = torch.full(
         (max_num_tokens,), -2, dtype=torch.int64, device=device

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -157,6 +157,7 @@ def test_mm_support_configured_after_model_load(monkeypatch):
         speculator.device = device
         speculator.max_num_tokens = 4
         speculator.max_num_reqs = 2
+        speculator.num_speculative_steps = 1
         speculator.hidden_size = 3
         speculator.dtype = torch.float32
         speculator.draft_model_config = draft_model_config

vllm/model_executor/layers/mamba/ops/scatter_states.py

@@ -1,8 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from typing import Any
+
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import WarmupIntRange
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
@@ -33,11 +41,39 @@ def _scatter_states_kernel(
     tl.store(state_ptr + state_idx * stride_state_batch + offsets, values, mask=mask)
 
 
+def _scatter_states_warmup_inputs(
+    *,
+    state_shape: tuple[int, ...],
+    state_dtype: torch.dtype,
+    indices_dtype: torch.dtype,
+    max_num_tokens: int,
+) -> dict[str, object]:
+    num_tokens: Any = WarmupIntRange(1, max_num_tokens + 1)
+    return dict(
+        state=TritonWarmupTensor(
+            state_dtype,
+            shape=(1,) + state_shape,
+        ),
+        src=TritonWarmupTensor(
+            state_dtype,
+            shape=(num_tokens,) + state_shape,
+        ),
+        indices=TritonWarmupTensor(
+            indices_dtype,
+            shape=(num_tokens,),
+        ),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_scatter_states_kernel,
+    warmup_inputs=_scatter_states_warmup_inputs,
+)
 def scatter_states(
     state: torch.Tensor,
     src: torch.Tensor,
     indices: torch.Tensor,
-) -> None:
+) -> DispatchSpec:
     """Scatter ``src`` rows into ``state`` at ``indices`` (in place).
 
     Equivalent to ``state[indices] = src`` but non-atomic and bandwidth-bound,
@@ -58,13 +94,10 @@ def scatter_states(
     assert src[0].is_contiguous()
     block_size = min(triton.next_power_of_2(row_size), 1024)
     grid = (triton.cdiv(row_size, block_size), indices.numel())
-    _scatter_states_kernel[grid](
-        state,
-        src,
-        indices,
-        state.stride(0),
-        src.stride(0),
-        indices.stride(0),
+    return grid, dict(
+        stride_state_batch=state.stride(0),
+        stride_src_batch=src.stride(0),
+        stride_indices=indices.stride(0),
         row_size=row_size,
         BLOCK_SIZE=block_size,
         num_warps=8,

vllm/model_executor/warmup/jit_warmup.py

@@ -21,6 +21,7 @@ from typing import Any, Generic, TypeVar, cast
 __all__ = [
     "JitWarmupRegistry",
     "VllmJitKernel",
+    "WarmupChoices",
     "WarmupIntRange",
     "get_ast_full_name",
     "get_function_source_node",
@@ -62,6 +63,21 @@ class WarmupIntRange:
     advance: Callable[[int], int] | None = None
 
 
+@dataclass(frozen=True)
+class WarmupChoices:
+    """Expand an explicit finite set of values in a traced warmup method."""
+
+    values: tuple[Any, ...]
+
+    def __init__(self, *values: Any) -> None:
+        object.__setattr__(self, "values", values)
+
+
+def _when(condition: bool) -> None:
+    """Filter symbolic cases in an AST-traced warmup-cases method."""
+    raise RuntimeError("_when() is only valid in an AST-traced warmup method")
+
+
 WarmupValues = Any
 CompileKeyDispatchFn = Callable[..., CompileKeyT]
 WarmupPredicateFn = Callable[..., bool]
@@ -87,6 +103,8 @@ class _WarmupInputRows:
 
 
 def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:
+    if isinstance(values, WarmupChoices):
+        return values.values
     if isinstance(values, WarmupIntRange):
         if values.advance is None:
             return tuple(range(values.start, values.stop, values.step))
@@ -278,6 +296,27 @@ class _DispatchExprEvaluator(ast.NodeVisitor):
     def visit_Constant(self, node: ast.Constant) -> Any:
         return node.value
 
+    def visit_Lambda(self, node: ast.Lambda) -> Callable[..., Any]:
+        arguments = node.args
+        if (
+            arguments.posonlyargs
+            or arguments.vararg is not None
+            or arguments.kwonlyargs
+            or arguments.kwarg is not None
+            or arguments.defaults
+        ):
+            raise _dispatch_expr_error(node, "Traced lambdas require positional args")
+        names = tuple(argument.arg for argument in arguments.args)
+
+        def evaluate(*values: Any) -> Any:
+            if len(values) != len(names):
+                raise TypeError(f"Expected {len(names)} lambda arguments")
+            return type(self)(
+                dict(self.values) | dict(zip(names, values)), self.globals
+            ).eval(node.body)
+
+        return evaluate
+
     def visit_IfExp(self, node: ast.IfExp) -> Any:
         return self.visit(node.body if self.visit(node.test) else node.orelse)
 
@@ -333,7 +372,12 @@ class _DispatchExprEvaluator(ast.NodeVisitor):
         return op(self.visit(node.left), self.visit(node.right))
 
     def visit_Call(self, node: ast.Call) -> Any:
-        args = [self.visit(arg) for arg in node.args]
+        args: list[Any] = []
+        for arg in node.args:
+            if isinstance(arg, ast.Starred):
+                args.extend(self.visit(arg.value))
+            else:
+                args.append(self.visit(arg))
         fn = self.visit(node.func)
         call_kwargs: dict[str, Any] = {}
         for keyword in node.keywords:
@@ -388,6 +432,45 @@ def _eval_dispatch_expr(
     return _DispatchExprEvaluator(kwargs, globals_).eval(node)
 
 
+def _validate_dispatch_expr(node: ast.AST) -> None:
+    """Enforce the dispatch AST subset before compiling traced code."""
+    allowed_nodes = (
+        ast.Name,
+        ast.Constant,
+        ast.Lambda,
+        ast.arguments,
+        ast.arg,
+        ast.IfExp,
+        ast.Tuple,
+        ast.List,
+        ast.BoolOp,
+        ast.And,
+        ast.Or,
+        ast.Compare,
+        *tuple(_CMP_OPS),
+        ast.UnaryOp,
+        ast.Not,
+        ast.USub,
+        ast.BinOp,
+        *tuple(_BIN_OPS),
+        ast.Call,
+        ast.keyword,
+        ast.Starred,
+        ast.Attribute,
+        ast.Subscript,
+        ast.Load,
+    )
+    for child in ast.walk(node):
+        if not isinstance(child, allowed_nodes):
+            raise _dispatch_expr_error(child, "Unsupported dispatch expression")
+        if isinstance(child, ast.Call) and any(
+            keyword.arg is None for keyword in child.keywords
+        ):
+            raise _dispatch_expr_error(
+                child, "Dispatch helper calls cannot use **kwargs"
+            )
+
+
 def _collect_input_names(
     node: ast.AST,
     candidate_names: set[str],
@@ -406,6 +489,23 @@ def _collect_input_names(
     }
 
 
+def _named_assignment(
+    statement: ast.stmt,
+    subject: str,
+) -> tuple[str, ast.AST] | None:
+    match statement:
+        case ast.Assign(targets=[ast.Name(id=name)], value=value):
+            return name, value
+        case ast.AnnAssign(target=ast.Name(id=name), value=value) if value is not None:
+            return name, value
+        case ast.Assign() | ast.AnnAssign():
+            raise _dispatch_expr_error(
+                statement, f"{subject} assignments require one name"
+            )
+        case _:
+            return None
+
+
 def _collect_expression_body(
     fn: Callable[..., Any],
     function_def: ast.FunctionDef | ast.Lambda,
@@ -422,37 +522,18 @@ def _collect_expression_body(
         ):
             continue
 
-        if isinstance(statement, ast.Assign):
-            if len(statement.targets) != 1 or not isinstance(
-                statement.targets[0], ast.Name
-            ):
-                raise _dispatch_expr_error(
-                    statement,
-                    "Traced warmup helper assignments must target one local name",
-                )
-            local_exprs.append((statement.targets[0].id, statement.value))
+        if (
+            assignment := _named_assignment(statement, "Traced warmup helper")
+        ) is not None:
+            local_exprs.append(assignment)
             continue
 
-        if isinstance(statement, ast.AnnAssign):
-            if statement.value is None:
-                raise _dispatch_expr_error(
-                    statement, "Traced warmup helper annotations must assign a value"
-                )
-            if not isinstance(statement.target, ast.Name):
-                raise _dispatch_expr_error(
-                    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 statement.value is not None:
-            return local_exprs, statement.value
-
         if isinstance(statement, ast.Return):
-            raise _dispatch_expr_error(
-                statement, "Traced warmup helper must return an expression"
-            )
+            if statement.value is None:
+                raise _dispatch_expr_error(
+                    statement, "Traced warmup helper must return an expression"
+                )
+            return local_exprs, statement.value
 
         raise _dispatch_expr_error(
             statement,
@@ -574,10 +655,17 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
     bind_launch_inputs = True
 
     def __init__(self) -> None:
-        self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
+        dispatch = type(self).dispatch
+        self._dispatch_trace = (
+            None
+            if dispatch is VllmJitKernel.dispatch
+            else _trace_compile_key_dispatch(self.dispatch)
+        )
         self._compiled_cache: dict[Any, Any] = {}
 
     def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
+        if self._dispatch_trace is None:
+            raise TypeError(f"{type(self).__name__} does not define dispatch()")
         return self._dispatch_trace.compile_key(self.CompileKey, kwargs)
 
     def _get_or_compile(
@@ -683,7 +771,184 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
 
         return traced
 
-    @abstractmethod
+    def _expand_warmup_cases(
+        self,
+        cases_fn: Callable[..., Any],
+        *args: Any,
+        _value_expander: Callable[[WarmupValues], tuple[Any, ...]] = (
+            _expand_warmup_values
+        ),
+        **kwargs: Any,
+    ) -> Iterator[Mapping[str, Any]]:
+        """Expand symbolic domains declared inside a warmup-cases method."""
+        function_def = get_function_source_node(cases_fn)
+        if isinstance(function_def, ast.Lambda):
+            raise _dispatch_expr_error(
+                function_def, "Warmup cases must be a function definition"
+            )
+        signature = inspect.signature(cases_fn)
+        bound = signature.bind(*args, **kwargs)
+        bound.apply_defaults()
+        static_values = dict(bound.arguments)
+        bound_self = getattr(cases_fn, "__self__", None)
+        if bound_self is not None:
+            static_values["self"] = bound_self
+        globals_ = getattr(cases_fn, "__func__", cases_fn).__globals__
+
+        domains: list[tuple[str, tuple[Any, ...]]] = []
+        local_exprs: list[tuple[str, ast.AST]] = []
+        predicates: list[ast.AST] = []
+        return_expr: ast.AST | None = None
+        for statement in function_def.body:
+            if (
+                isinstance(statement, ast.Expr)
+                and isinstance(statement.value, ast.Constant)
+                and isinstance(statement.value.value, str)
+            ):
+                continue
+            assignment = _named_assignment(statement, "Warmup case")
+            if assignment is not None:
+                name, value_expr = assignment
+                call_name = (
+                    (get_ast_full_name(value_expr.func) or "").split(".")[-1]
+                    if isinstance(value_expr, ast.Call)
+                    else ""
+                )
+                if call_name in {"WarmupIntRange", "WarmupChoices"}:
+                    value = _eval_dispatch_expr(value_expr, static_values, globals_)
+                    domains.append((name, _value_expander(value)))
+                else:
+                    local_exprs.append((name, value_expr))
+                continue
+            if (
+                isinstance(statement, ast.Expr)
+                and isinstance(statement.value, ast.Call)
+                and (get_ast_full_name(statement.value.func) or "").endswith("_when")
+            ):
+                if len(statement.value.args) != 1 or statement.value.keywords:
+                    raise _dispatch_expr_error(
+                        statement, "_when requires exactly one positional expression"
+                    )
+                predicates.append(statement.value.args[0])
+                continue
+            if isinstance(statement, ast.Return):
+                if return_expr is not None or statement.value is None:
+                    raise _dispatch_expr_error(
+                        statement, "Warmup cases require one final return expression"
+                    )
+                return_expr = statement.value
+                continue
+            raise _dispatch_expr_error(
+                statement,
+                "AST-traced warmup cases support assignments, _when, and return",
+            )
+
+        if return_expr is None or not isinstance(return_expr, ast.Call):
+            raise ValueError("AST-traced warmup cases must return one case(...) call")
+        if any(keyword.arg is None for keyword in return_expr.keywords):
+            raise _dispatch_expr_error(
+                return_expr, "Warmup case expressions do not support **kwargs"
+            )
+
+        class _InlineDomainRewriter(ast.NodeTransformer):
+            def visit_Call(self, node: ast.Call) -> ast.AST:
+                call_name = (get_ast_full_name(node.func) or "").split(".")[-1]
+                if call_name not in {"WarmupIntRange", "WarmupChoices"}:
+                    return self.generic_visit(node)
+                name = f"__warmup_domain_{len(domains)}"
+                domain = _eval_dispatch_expr(node, static_values, globals_)
+                domains.append((name, _value_expander(domain)))
+                return ast.copy_location(ast.Name(id=name, ctx=ast.Load()), node)
+
+        return_expr = cast(ast.Call, _InlineDomainRewriter().visit(return_expr))
+
+        domain_names = tuple(name for name, _ in domains)
+        dynamic_names = set(domain_names)
+        dynamic_local_exprs: list[tuple[str, ast.AST]] = []
+        for name, expr in local_exprs:
+            referenced_names = {
+                node.id for node in ast.walk(expr) if isinstance(node, ast.Name)
+            }
+            if referenced_names & dynamic_names:
+                dynamic_local_exprs.append((name, expr))
+                dynamic_names.add(name)
+            else:
+                static_values[name] = _eval_dispatch_expr(expr, static_values, globals_)
+
+        for _, expr in dynamic_local_exprs:
+            _validate_dispatch_expr(expr)
+        for predicate in predicates:
+            _validate_dispatch_expr(predicate)
+        _validate_dispatch_expr(return_expr)
+
+        case_body: list[ast.stmt] = [
+            ast.Assign(
+                targets=[ast.Name(id=name, ctx=ast.Store())],
+                value=cast(ast.expr, expr),
+            )
+            for name, expr in dynamic_local_exprs
+        ]
+        if predicates:
+            predicate = (
+                predicates[0]
+                if len(predicates) == 1
+                else ast.BoolOp(
+                    op=ast.And(),
+                    values=[cast(ast.expr, value) for value in predicates],
+                )
+            )
+            case_body.append(
+                ast.If(
+                    test=ast.UnaryOp(op=ast.Not(), operand=cast(ast.expr, predicate)),
+                    body=[ast.Return(value=ast.Constant(value=None))],
+                    orelse=[],
+                )
+            )
+        case_body.append(ast.Return(value=cast(ast.expr, return_expr)))
+        # FunctionDef fields vary across the supported Python versions, so no
+        # single constructor overload matches every mypy target.
+        case_function = ast.FunctionDef(  # type: ignore[call-overload]
+            name="__vllm_warmup_case",
+            args=ast.arguments(
+                posonlyargs=[],
+                args=[ast.arg(arg=name) for name in domain_names],
+                kwonlyargs=[],
+                kw_defaults=[],
+                defaults=[],
+            ),
+            body=case_body,
+            decorator_list=[],
+            returns=None,
+            type_comment=None,
+        )
+        case_globals = dict(globals_)
+        case_globals.update(static_values)
+        exec(
+            compile(
+                ast.fix_missing_locations(
+                    ast.Module(body=[case_function], type_ignores=[])
+                ),
+                "<jit-warmup>",
+                "exec",
+            ),
+            case_globals,
+        )
+        evaluate_case = cast(Callable[..., Any], case_globals[case_function.name])
+
+        domain_values = tuple(values for _, values in domains)
+        for values in itertools.product(*domain_values):
+            case = evaluate_case(*values)
+            if case is None:
+                continue
+            if not isinstance(case, Mapping):
+                raise TypeError("AST-traced warmup cases must return a mapping")
+            if any(
+                isinstance(value, WarmupChoices | WarmupIntRange)
+                for value in case.values()
+            ):
+                raise TypeError("Warmup domains must be direct expressions")
+            yield case
+
     def dispatch(self, **kwargs: Any) -> CompileKeyT:
         """Build one compile key from one concrete dispatch point."""
         raise NotImplementedError
@@ -751,6 +1016,19 @@ class JitWarmupRegistry:
             list[tuple[tuple[Any, ...], dict[str, Any]]],
         ] = {}
 
+    @classmethod
+    def capture(cls, init_fn: Callable[..., None]) -> Callable[..., None]:
+        """Collect warmup registrations made while the decorated callable runs."""
+
+        @wraps(init_fn)
+        def wrapped(instance: Any, vllm_config: Any, *args: Any, **kwargs: Any) -> None:
+            registry = cls(vllm_config)
+            instance.jit_warmup_registry = registry
+            with registry.activate():
+                init_fn(instance, vllm_config, *args, **kwargs)
+
+        return wrapped
+
     @contextmanager
     def activate(self) -> Iterator[None]:
         """Collect registrations made in this context."""
@@ -808,9 +1086,10 @@ class JitWarmupRegistry:
                 if (
                     not args
                     and not kwargs
-                    and inspect.signature(kernel.get_warmup_keys).parameters
+                    and "vllm_config"
+                    in inspect.signature(kernel.get_warmup_keys).parameters
                 ):
-                    args = (self.vllm_config,)
+                    kwargs = {"vllm_config": self.vllm_config}
                 for compile_key in kernel.get_warmup_keys(*args, **kwargs):
                     compile_keys[compile_key] = None
             if compile_keys:

vllm/model_executor/warmup/jit_warmup_triton_helper.py

@@ -3,25 +3,130 @@
 import ast
 import inspect
 from abc import abstractmethod
-from collections.abc import Callable, Mapping
-from dataclasses import dataclass
-from functools import cached_property
-from typing import Any, ClassVar, Generic, TypeVar, cast
+from collections.abc import Callable, Hashable, Iterable, Mapping
+from contextlib import nullcontext, suppress
+from dataclasses import dataclass, field
+from functools import cache, cached_property, update_wrapper
+from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload
 
 from vllm.model_executor.warmup.jit_warmup import (
     VllmJitKernel,
     get_ast_full_name,
     get_function_source_node,
 )
+from vllm.platforms import current_platform
+from vllm.triton_utils import triton
 
 CompileKeyT = TypeVar("CompileKeyT")
-# ``(grid, launch_kwargs)`` or ``(grid, launch_kwargs, outputs)`` for
-# self-allocating kernels. ``grid=None`` skips the launch (e.g. empty batch)
-# but still returns the declared outputs.
+P = ParamSpec("P")
+# ``kernel_launcher`` also supports a ``None`` grid when no launch is needed.
 LaunchSpec = (
     tuple[tuple[int, ...] | None, dict[str, Any]]
     | tuple[tuple[int, ...] | None, dict[str, Any], Any]
 )
+# Dispatchers select a concrete launch configuration for a Triton kernel.
+DispatchSpec = (
+    tuple[tuple[int, ...], dict[str, Any]] | tuple[tuple[int, ...], dict[str, Any], Any]
+)
+
+
+@dataclass(frozen=True)
+class _LaunchBindingPlan:
+    input_targets: tuple[tuple[int, str], ...]
+    stride_targets: tuple[tuple[str, str, int], ...]
+
+
+@cache
+def _launch_binding_plan(
+    arg_names: tuple[str, ...],
+    input_names: tuple[str, ...],
+    launch_names: tuple[str, ...],
+) -> _LaunchBindingPlan:
+    arg_name_set = set(arg_names)
+    launch_name_set = set(launch_names)
+
+    def kernel_arg(input_name: str) -> str | None:
+        pointer_name = f"{input_name}_ptr"
+        candidates = (
+            input_name,
+            pointer_name,
+            input_name.upper(),
+            pointer_name.upper(),
+        )
+        return next(
+            (candidate for candidate in candidates if candidate in arg_name_set),
+            None,
+        )
+
+    input_targets = tuple(
+        (index, target)
+        for index, name in enumerate(input_names)
+        if (target := kernel_arg(name)) is not None and target not in launch_name_set
+    )
+    bound_names = launch_name_set | {target for _, target in input_targets}
+    stride_targets = []
+    for target in arg_names:
+        if target in bound_names:
+            continue
+        name, separator, dim = target.rpartition("_stride")
+        if not separator:
+            name, separator, dim = target.rpartition("_STRIDE")
+        dim = dim.removeprefix("_")
+        if separator and (not dim or dim.isdigit()):
+            pointer = f"{name}_PTR" if separator == "_STRIDE" else f"{name}_ptr"
+            source = name if name in bound_names else pointer
+            if source in bound_names:
+                stride_targets.append((target, source, int(dim or 0)))
+    return _LaunchBindingPlan(input_targets, tuple(stride_targets))
+
+
+def triton_warmup_inputs(
+    kernel: Any,
+    *args: Any,
+    grid: tuple[int, ...],
+    pointer_dtypes: (
+        Mapping[Any, Iterable[str]] | Iterable[tuple[Any, Iterable[str]]] | None
+    ) = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Build launcher inputs from Triton's native positional argument order."""
+    arg_names = tuple(kernel.arg_names)
+    if len(args) > len(arg_names):
+        raise ValueError(
+            f"Received {len(args)} positional inputs for {len(arg_names)} "
+            "Triton kernel arguments"
+        )
+    inputs = dict(zip(arg_names, args))
+    duplicate_names = inputs.keys() & kwargs.keys()
+    if duplicate_names:
+        raise ValueError(
+            f"Triton inputs passed twice: {', '.join(sorted(duplicate_names))}"
+        )
+    inputs.update(kwargs)
+    if pointer_dtypes is not None:
+        function_def = get_function_source_node(kernel)
+        if not isinstance(function_def, ast.FunctionDef):
+            raise ValueError("Expected Triton kernel to be defined as a function")
+        pointer_names = _pointer_arg_names(function_def, arg_names)
+        groups = (
+            pointer_dtypes.items()
+            if isinstance(pointer_dtypes, Mapping)
+            else pointer_dtypes
+        )
+        for dtype, names in groups:
+            for name in names:
+                if name not in arg_names:
+                    raise ValueError(f"Unknown Triton pointer argument: {name}")
+                if name not in pointer_names:
+                    raise ValueError(f"Triton argument is not a pointer: {name}")
+                if name in inputs:
+                    raise ValueError(f"Triton input passed twice: {name}")
+                inputs[name] = TritonWarmupTensor(dtype)
+        missing_pointers = pointer_names - inputs.keys()
+        if missing_pointers:
+            names = ", ".join(sorted(missing_pointers))
+            raise ValueError(f"Missing Triton pointer inputs: {names}")
+    return {"grid": grid, **inputs}
 
 
 def triton_scalar_specialization_rep(value: int) -> int:
@@ -72,6 +177,7 @@ class TritonWarmupTensor:
     aligned: bool = True
     shape: tuple[int, ...] = (1,)
     strides: tuple[int, ...] | None = None
+    init: Any = 0
 
     def data_ptr(self) -> int:
         return 0 if self.aligned else 1
@@ -79,6 +185,12 @@ class TritonWarmupTensor:
     def ptr_range(self) -> int:
         return 0
 
+    @property
+    def device(self) -> Any:
+        import torch
+
+        return torch.device(current_platform.device_type)
+
     def stride(self, dim: int | None = None) -> int | tuple[int, ...]:
         if self.strides is None:
             strides: list[int] = []
@@ -95,9 +207,10 @@ class TritonWarmupTensor:
 class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
     """Triton owner whose runtime launch specification is reused for warmup."""
 
-    kernel: ClassVar[Any]
+    kernel: Any
     _warming = False
     _warming_compile_key: CompileKeyT | None = None
+    _run_autotune = False
 
     @abstractmethod
     def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
@@ -126,6 +239,21 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
             f"Cannot inspect kernel parameters for {type(self.kernel).__name__}"
         )
 
+    def _prepare_launch_kwargs(
+        self,
+        input_names: tuple[str, ...],
+        input_values: tuple[Any, ...],
+        launch_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        plan = _launch_binding_plan(
+            self._kernel_arg_names, input_names, tuple(launch_kwargs)
+        )
+        for index, target in plan.input_targets:
+            launch_kwargs[target] = input_values[index]
+        for target, source, dim in plan.stride_targets:
+            launch_kwargs[target] = launch_kwargs[source].stride(dim)
+        return launch_kwargs
+
     def launch(
         self,
         launch_spec: LaunchSpec,
@@ -136,15 +264,20 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
         else:
             grid, launch_kwargs = launch_spec
             outputs = None
-        kwargs = dict(launch_kwargs)
+        kwargs = self._prepare_launch_kwargs(
+            tuple(inputs), tuple(inputs.values()), dict(launch_kwargs)
+        )
         kernel = kwargs.pop("kernel", None) or self.kernel
         runtime_launcher = kwargs.pop("_runtime_launcher", None)
         runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0)
-        for name, value in inputs.items():
-            target = name if name in self._kernel_arg_names else f"{name}_ptr"
-            if target in self._kernel_arg_names and target not in kwargs:
-                kwargs[target] = value
         if self._warming:
+            if self._run_autotune:
+                assert grid is not None
+                kernel[grid](**kwargs)
+                return outputs
+            kwargs = {
+                name: _triton_metadata_arg(value) for name, value in kwargs.items()
+            }
             if (
                 self._warming_compile_key is not None
                 and "launch_pdl" in kwargs
@@ -168,6 +301,367 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
         return outputs
 
 
+@dataclass(frozen=True)
+class TritonJitKey:
+    """Process-local identity of one Triton JIT specialization."""
+
+    jit_function_id: int
+    jit_function_key: Hashable
+    device: Hashable
+    cache_key: Hashable
+
+
+@dataclass(frozen=True)
+class TritonCompileKey:
+    """Triton-derived compile key with one non-comparing replay input."""
+
+    jit_keys: frozenset[TritonJitKey]
+    inputs: tuple[tuple[str, Any], ...] = field(compare=False, hash=False, repr=False)
+
+
+def _triton_key_deriver(
+    kernel: Any,
+) -> Callable[[Mapping[str, Any]], set[TritonJitKey]]:
+    """Prepare Triton's key derivation once for one kernel and device."""
+    knobs = triton.knobs
+    Autotuner = triton.runtime.autotuner.Autotuner
+    Heuristics = triton.runtime.autotuner.Heuristics
+    driver = triton.runtime.driver
+    JITFunction = triton.runtime.jit.JITFunction
+    compute_cache_key = triton.runtime.jit.compute_cache_key
+
+    if isinstance(kernel, Heuristics):
+        derive_inner = _triton_key_deriver(kernel.fn)
+
+        def derive_heuristic(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:
+            heuristic_kwargs = dict(kwargs)
+            for name, heuristic in kernel.values.items():
+                heuristic_kwargs[name] = heuristic(heuristic_kwargs)
+            return derive_inner(heuristic_kwargs)
+
+        return derive_heuristic
+
+    if isinstance(kernel, Autotuner):
+        derive_inner = _triton_key_deriver(kernel.fn)
+
+        def derive_autotuned(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:
+            previous_nargs = getattr(kernel, "nargs", None)
+            kernel.nargs = {}
+            try:
+                configs = kernel.prune_configs(dict(kwargs))
+            finally:
+                kernel.nargs = previous_nargs
+            keys: set[TritonJitKey] = set()
+            for config in configs:
+                conflicts = kwargs.keys() & config.kwargs.keys()
+                if conflicts:
+                    names = ", ".join(sorted(conflicts))
+                    raise ValueError(f"Conflicting autotune parameters: {names}")
+                keys.update(derive_inner(dict(kwargs) | config.all_kwargs()))
+            return keys
+
+        return derive_autotuned
+
+    if not isinstance(kernel, JITFunction):
+        raise TypeError(f"Unsupported Triton kernel wrapper: {type(kernel).__name__}")
+
+    device = cast(Hashable, driver.active.get_current_device())
+    _, kernel_key_cache, _, _, binder = kernel.device_caches[device]
+    jit_function_id = id(kernel)
+    jit_function_key = cast(Hashable, kernel.cache_key)
+    debug = kernel.debug or knobs.runtime.debug
+    instrumentation_mode = knobs.compilation.instrumentation_mode
+    fpsan_homomorphic_casts = getattr(
+        knobs.compilation, "fpsan_homomorphic_casts", None
+    )
+    inspection_hook = knobs.runtime.add_stages_inspection_hook
+
+    def derive_jit(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:
+        binder_kwargs = dict(kwargs)
+        binder_kwargs["debug"] = binder_kwargs.get("debug", debug) or debug
+        binder_kwargs["instrumentation_mode"] = instrumentation_mode
+        if fpsan_homomorphic_casts is not None:
+            binder_kwargs["fpsan_homomorphic_casts"] = fpsan_homomorphic_casts
+        _, specialization, options = binder(**binder_kwargs)
+        if inspection_hook is not None:
+            _, inspection_hash = inspection_hook()
+            specialization.append(f'("custom_pipeline", {inspection_hash})')
+        cache_key = cast(
+            Hashable,
+            compute_cache_key(kernel_key_cache, specialization, options),
+        )
+        return {
+            TritonJitKey(
+                jit_function_id,
+                jit_function_key,
+                device,
+                cache_key,
+            )
+        }
+
+    return derive_jit
+
+
+WarmupCases = Mapping[str, Any] | Iterable[Mapping[str, Any]]
+
+
+class _AutomaticTritonJitKernel(VllmTritonJitKernel[TritonCompileKey]):
+    def warmup_inputs(self, compile_key: TritonCompileKey) -> dict[str, Any]:
+        return dict(compile_key.inputs)
+
+    def _provider_cases(
+        self,
+        provider: Callable[..., WarmupCases],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Iterable[Mapping[str, Any]]:
+        cases_node = get_function_source_node(provider)
+        if any(
+            isinstance(node, ast.Call)
+            and (get_ast_full_name(node.func) or "").split(".")[-1]
+            in {"WarmupIntRange", "WarmupChoices", "_when"}
+            for node in ast.walk(cases_node)
+        ):
+            return self._expand_warmup_cases(provider, *args, **kwargs)
+        cases = provider(*args, **kwargs)
+        return (cases,) if isinstance(cases, Mapping) else cases
+
+
+def _is_autotuned(kernel: Any) -> bool:
+    Autotuner = triton.runtime.autotuner.Autotuner
+
+    while kernel is not None:
+        if isinstance(kernel, Autotuner):
+            return True
+        kernel = getattr(kernel, "fn", None)
+    return False
+
+
+def _triton_metadata_arg(value: Any) -> Any:
+    from torch._subclasses.fake_tensor import FakeTensor
+
+    if not isinstance(value, FakeTensor):
+        return value
+    return TritonWarmupTensor(
+        value.dtype,
+        aligned=getattr(value, "_vllm_warmup_aligned", True),
+        shape=tuple(value.shape),
+        strides=tuple(value.stride()),
+    )
+
+
+def _materialize_warmup_case(
+    case: Mapping[str, Any],
+    *,
+    real: bool,
+    mode: Any,
+    cache: dict[TritonWarmupTensor, Any],
+) -> dict[str, Any]:
+    import torch
+
+    def materialize(value: Any) -> Any:
+        if not isinstance(value, TritonWarmupTensor):
+            return value
+        if not real:
+            try:
+                return cache[value]
+            except KeyError:
+                pass
+            except TypeError:
+                # Unhashable initializers are uncommon and safe to materialize directly.
+                pass
+
+        strides = cast(tuple[int, ...], value.stride())
+        with mode:
+            tensor = torch.empty_strided(
+                value.shape,
+                strides,
+                dtype=value.dtype,
+                device=value.device,
+            )
+        if real:
+            if callable(value.init):
+                value.init(tensor)
+            else:
+                tensor.fill_(value.init)
+        else:
+            tensor._vllm_warmup_aligned = value.aligned
+            with suppress(TypeError):
+                cache[value] = tensor
+        return tensor
+
+    return {name: materialize(value) for name, value in case.items()}
+
+
+class _DecoratedTritonJitKernel(_AutomaticTritonJitKernel):
+    def __init__(
+        self,
+        kernel: Any,
+        warmup_inputs: Callable[..., WarmupCases],
+        dispatch: Callable[..., DispatchSpec] | None,
+    ) -> None:
+        self.kernel = kernel
+        self._run_autotune = _is_autotuned(kernel)
+        self._warmup_inputs_fn = warmup_inputs
+        self._dispatch_fn = dispatch
+        self._dispatch_arg_names: tuple[str, ...] = ()
+        self._dispatch_defaults: dict[str, Any] = {}
+        if dispatch is not None:
+            parameters = tuple(inspect.signature(dispatch).parameters.values())
+            if any(
+                parameter.kind
+                in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
+                for parameter in parameters
+            ):
+                raise TypeError("Triton dispatchers require explicit parameters")
+            self._dispatch_arg_names = tuple(parameter.name for parameter in parameters)
+            self._dispatch_defaults = {
+                parameter.name: parameter.default
+                for parameter in parameters
+                if parameter.default is not inspect.Parameter.empty
+            }
+        super().__init__()
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if self._dispatch_fn is None:
+            grid = kwargs.pop("grid") if not args else args[0]
+            kernel_args = () if not args else args[1:]
+            inputs = triton_warmup_inputs(
+                self.kernel, *kernel_args, grid=grid, **kwargs
+            )
+            inputs.pop("grid")
+            return self.launch((grid, inputs), {})
+        spec = self._dispatch_fn(*args, **kwargs)
+        input_values = args
+        if len(args) != len(self._dispatch_arg_names):
+            input_values += tuple(
+                kwargs[name] if name in kwargs else self._dispatch_defaults[name]
+                for name in self._dispatch_arg_names[len(args) :]
+            )
+        return self.launch(
+            spec,
+            dict(zip(self._dispatch_arg_names, input_values)),
+        )
+
+    def __getitem__(self, grid: tuple[int, ...]) -> Callable[..., Any]:
+        return cast(Callable[..., Any], self.kernel[grid])
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.kernel, name)
+
+    def get_warmup_keys(
+        self, *, vllm_config: Any | None = None, **kwargs: Any
+    ) -> list[TritonCompileKey]:
+        keys: list[TritonCompileKey] = []
+        seen_jit_keys: set[frozenset[TritonJitKey]] = set()
+        if "vllm_config" in inspect.signature(self._warmup_inputs_fn).parameters:
+            kwargs.setdefault("vllm_config", vllm_config)
+        cases = self._provider_cases(self._warmup_inputs_fn, **kwargs)
+        derive_keys = _triton_key_deriver(self.kernel)
+        materialize = self._run_autotune
+        mode: Any = nullcontext() if self._run_autotune else None
+        materialized: dict[TritonWarmupTensor, Any] = {}
+        for case in cases:
+            inputs = (
+                _materialize_warmup_case(
+                    case,
+                    real=self._run_autotune,
+                    mode=mode,
+                    cache=materialized,
+                )
+                if materialize
+                else case
+            )
+            if self._dispatch_fn is None:
+                prepared = dict(inputs)
+                prepared.pop("grid")
+            else:
+                dispatch_inputs = {
+                    name: inputs[name]
+                    for name in self._dispatch_arg_names
+                    if name in inputs
+                }
+                try:
+                    _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2]
+                except (AssertionError, AttributeError, TypeError):
+                    if materialize:
+                        raise
+                    from torch._subclasses.fake_tensor import FakeTensorMode
+
+                    materialize = True
+                    mode = FakeTensorMode()
+                    inputs = _materialize_warmup_case(
+                        case,
+                        real=False,
+                        mode=mode,
+                        cache=materialized,
+                    )
+                    dispatch_inputs = {
+                        name: inputs[name]
+                        for name in self._dispatch_arg_names
+                        if name in inputs
+                    }
+                    _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2]
+                prepared = self._prepare_launch_kwargs(
+                    tuple(inputs),
+                    tuple(inputs.values()),
+                    dict(launch_kwargs),
+                )
+            if materialize:
+                prepared = {
+                    name: _triton_metadata_arg(value)
+                    for name, value in prepared.items()
+                }
+            jit_keys = frozenset(derive_keys(prepared))
+            if jit_keys and jit_keys not in seen_jit_keys:
+                seen_jit_keys.add(jit_keys)
+                keys.append(TritonCompileKey(jit_keys, tuple(inputs.items())))
+        return keys
+
+
+class TritonKernelDispatcher(Protocol[P]):
+    """Callable Triton launcher with compile-time warmup support."""
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Any: ...
+
+    def register_warmup(self, *args: Any, **kwargs: Any) -> None: ...
+
+
+@overload
+def triton_kernel_dispatcher_with_warmup(
+    *,
+    warmup_inputs: Callable[..., WarmupCases],
+    kernel: None = None,
+) -> Callable[[Any], Any]: ...
+
+
+@overload
+def triton_kernel_dispatcher_with_warmup(
+    *,
+    warmup_inputs: Callable[..., WarmupCases],
+    kernel: Any,
+) -> Callable[[Callable[P, DispatchSpec]], TritonKernelDispatcher[P]]: ...
+
+
+def triton_kernel_dispatcher_with_warmup(
+    *,
+    warmup_inputs: Callable[..., WarmupCases],
+    kernel: Any | None = None,
+) -> Callable[[Any], Any]:
+    """Decorate a dispatch function or a native Triton kernel for warmup."""
+
+    def decorate(dispatch: Any) -> Any:
+        owner = _DecoratedTritonJitKernel(
+            dispatch if kernel is None else kernel,
+            warmup_inputs,
+            None if kernel is None else dispatch,
+        )
+        update_wrapper(owner, dispatch, updated=())
+        return owner
+
+    return decorate
+
+
 @dataclass(frozen=True)
 class TritonPointerInputVariant:
     # Named pointer-alignment variant for compile-only Triton warmup.

vllm/model_executor/warmup/kernel_warmup.py

@@ -41,9 +41,6 @@ from vllm.model_executor.warmup.qwen_vl_triton_warmup import qwen_vl_triton_warm
 from vllm.model_executor.warmup.replayssm_warmup import (
     replayssm_autotune_warmup,
 )
-from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
-    spec_decode_rejection_warmup,
-)
 from vllm.model_executor.warmup.watermark_sample_warmup import (
     watermark_sample_warmup,
 )
@@ -178,7 +175,10 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
         logger.info("JIT kernel warmup starting.")
         jit_warmup_start = time.perf_counter()
         try:
-            worker.model_runner.jit_warmup_registry.warmup()
+            registry = (
+                worker.model_runner.jit_warmup_registry  # type: ignore[attr-defined]
+            )
+            registry.warmup()
         except Exception:
             logger.exception(
                 "JIT kernel warmup failed after %.2fs.",
@@ -209,7 +209,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     # Run next so input-prep kernels JIT against pristine runner state.
     if enable_jit_warmup:
         kimi_k3_triton_warmup(worker)
-        spec_decode_rejection_warmup(worker)
         watermark_sample_warmup(worker)
         qwen4_exp_qsa_triton_warmup(worker)
 

vllm/model_executor/warmup/spec_decode_rejection_warmup.py

@@ -1,139 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warm up speculative-decoding rejection kernels."""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Any
-
-import torch
-
-from vllm.logger import init_logger
-
-if TYPE_CHECKING:
-    from vllm.v1.worker.gpu_worker import Worker
-
-logger = init_logger(__name__)
-
-
-@torch.inference_mode()
-def spec_decode_rejection_warmup(worker: Worker) -> None:
-    spec_config = worker.vllm_config.speculative_config
-    if spec_config is None:
-        return
-
-    from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
-        rejection_sample,
-    )
-
-    model_config = worker.vllm_config.model_config
-    vocab_size = model_config.get_vocab_size()
-    num_spec = spec_config.num_speculative_tokens
-    if num_spec <= 0 or vocab_size <= 0:
-        return
-
-    rejection_method = getattr(spec_config, "rejection_sample_method", None)
-    use_block_verification = rejection_method == "block"
-    use_synthetic = rejection_method == "synthetic"
-    # HAS_DRAFT_LOGITS is a constexpr, so the warmup must match the runtime.
-    use_draft_logits = (
-        getattr(spec_config, "draft_sample_method", "greedy") == "probabilistic"
-    )
-
-    device = worker.device
-    num_reqs = 1
-    tokens_per_req = num_spec + 1
-    num_logits = num_reqs * tokens_per_req
-
-    watermark_config = getattr(worker.vllm_config, "watermark_config", None)
-    watermark_kwargs: dict[str, Any] = {}
-    if watermark_config is not None:
-        from vllm.v1.watermarking.spec_decode import speculative_target_watermark_key
-
-        try:
-            watermark_key = speculative_target_watermark_key(watermark_config)
-        except Exception:
-            logger.warning(
-                "Skipping spec-decode rejection sampler warmup: could not derive "
-                "the watermark recovery key.",
-                exc_info=True,
-            )
-            return
-        # Tensor dtypes and integer widths affect Triton specialization.
-        watermark_kwargs = {
-            "contexts": torch.zeros(
-                (num_logits, watermark_config.context_width),
-                dtype=torch.int32,
-                device=device,
-            ),
-            "watermarking": torch.zeros(num_reqs, dtype=torch.bool, device=device),
-            "watermark_key": watermark_key,
-        }
-
-    # Sampling-parameter processing may promote either distribution to fp32.
-    model_dtype = model_config.dtype
-    warmup_dtype_pairs: set[tuple[torch.dtype, torch.dtype | None]]
-    if use_draft_logits:
-        warmup_dtype_pairs = {
-            (model_dtype, model_dtype),
-            (torch.float32, torch.float32),
-            (torch.float32, model_dtype),
-            (model_dtype, torch.float32),
-        }
-    else:
-        warmup_dtype_pairs = {(model_dtype, None), (torch.float32, None)}
-
-    logger.info(
-        "Warming up spec-decode rejection sampler kernels "
-        "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s, watermark=%s).",
-        vocab_size,
-        num_spec,
-        [(str(t), str(d)) for t, d in warmup_dtype_pairs],
-        use_block_verification,
-        watermark_config is not None,
-    )
-    for tgt_dtype, draft_dtype in warmup_dtype_pairs:
-        target_logits = torch.zeros(
-            (num_logits, vocab_size), dtype=tgt_dtype, device=device
-        )
-        draft_logits = (
-            torch.zeros(
-                (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device
-            )
-            if draft_dtype is not None
-            else None
-        )
-        synthetic_rates = (
-            torch.full((num_spec,), 0.5, dtype=torch.float32, device=device)
-            if use_synthetic
-            else None
-        )
-        try:
-            rejection_sample(
-                target_logits=target_logits,
-                draft_logits=draft_logits,
-                draft_sampled=torch.zeros(num_logits, dtype=torch.int32, device=device),
-                cu_num_logits=torch.tensor(
-                    [0, num_logits], dtype=torch.int32, device=device
-                ),
-                pos=torch.zeros(num_logits, dtype=torch.int64, device=device),
-                idx_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device),
-                expanded_idx_mapping=torch.zeros(
-                    num_logits, dtype=torch.int64, device=device
-                ),
-                expanded_local_pos=torch.arange(
-                    num_logits, dtype=torch.int32, device=device
-                ),
-                temperature=torch.zeros(num_reqs, dtype=torch.float32, device=device),
-                seed=torch.full((num_reqs,), 42, dtype=torch.int64, device=device),
-                num_speculative_steps=num_spec,
-                synthetic_conditional_rates=synthetic_rates,
-                use_fp64=False,
-                use_block_verification=use_block_verification,
-                **watermark_kwargs,
-            )
-        except Exception:
-            logger.warning(
-                "Skipping spec-decode rejection sampler warmup.", exc_info=True
-            )
-            return

vllm/models/glm5next/nvidia/kda.py

@@ -175,6 +175,15 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
         projection_size = self.head_dim * self.num_heads
         self.local_projection_size = divide(projection_size, self.tp_size)
 
+        _, recurrent_state_shape = self.get_state_shape()
+        _, recurrent_state_dtype = self.get_state_dtype()
+        scatter_states.register_warmup(
+            state_shape=recurrent_state_shape,
+            state_dtype=recurrent_state_dtype,
+            indices_dtype=torch.int32,
+            max_num_tokens=vllm_config.scheduler_config.max_num_seqs,
+        )
+
         # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches).
         # Order matches checkpoint's fused_qkvbfg_a_proj convention.
         # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks.

vllm/v1/sample/ops/topk_topp_sampler.py

@@ -13,11 +13,27 @@ from vllm.platforms import CpuArchEnum, current_platform
 from vllm.triton_utils import HAS_TRITON
 
 if HAS_TRITON:
-    from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
+    from vllm.v1.sample.ops.topk_topp_triton import (
+        _topk_topp,
+        _topp_split_mask,
+        _topp_split_stats,
+        _topp_split_step,
+        apply_top_k_top_p_triton,
+    )
 
 logger = init_logger(__name__)
 
 
+def register_top_k_top_p_warmups() -> None:
+    """Register every native accelerator sampling kernel used at runtime."""
+    if HAS_TRITON and not current_platform.is_cpu():
+        _topk_topp.register_warmup()
+        if current_platform.is_cuda_alike():
+            _topp_split_stats.register_warmup()
+            _topp_split_step.register_warmup()
+            _topp_split_mask.register_warmup()
+
+
 def _skip_aiter_sampler_on_gfx1250() -> bool:
     # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU.
     from vllm.platforms.rocm import on_gfx1250
@@ -128,6 +144,9 @@ class TopKTopPSampler(nn.Module):
         else:
             self.forward = self.forward_native
 
+        # Every accelerator backend can fall back to native sampling at runtime.
+        register_top_k_top_p_warmups()
+
     def forward_native(
         self,
         logits: torch.Tensor,

vllm/v1/sample/ops/topk_topp_triton.py

@@ -9,8 +9,21 @@ using Pivot-based Truncation and Selection" By Park et al.
 
 """
 
+from typing import Any
+
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupChoices,
+    WarmupIntRange,
+    _when,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.math_utils import next_power_of_2
@@ -863,6 +876,90 @@ def _topk_topp_kernel(
                 tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n)
 
 
+def _max_sampler_batch_size(vllm_config: Any) -> int:
+    return vllm_config.scheduler_config.max_num_seqs * (
+        vllm_config.num_speculative_tokens + 1
+    )
+
+
+def _topk_topp_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    split_enabled = current_platform.is_cuda_alike()
+    mode: Any = WarmupChoices((True, False), (True, True), (False, True))
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(1, _max_sampler_batch_size(vllm_config) + 1)
+    topk_enabled = mode[0]
+    topp_enabled = mode[1]
+    _when(topk_enabled or not (split_enabled and batch_size <= _SPLIT_MAX_BATCH))
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        buffer=TritonWarmupTensor(torch.float32),
+        percentile_to_std_table=TritonWarmupTensor(torch.float32),
+        normal_cdf_to_sigma_table=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if topk_enabled else None,
+        p=TritonWarmupTensor(torch.float32) if topp_enabled else None,
+        mask_value=float("-inf"),
+        num_sm=num_compute_units(),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topk_topp_kernel,
+    warmup_inputs=_topk_topp_warmup_inputs,
+)
+def _topk_topp(
+    logits: torch.Tensor,
+    buffer: torch.Tensor,
+    percentile_to_std_table: torch.Tensor,
+    normal_cdf_to_sigma_table: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor | None,
+    mask_value: float,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    topk_enabled = k is not None
+    topp_enabled = p is not None
+    k_ptr = k if k is not None else logits
+    p_ptr = p if p is not None else logits
+    num_programs = min(num_sm, batch_size)
+    # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.
+    # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot
+    # approximation; use smaller tiles for accurate top-p results.
+    if logits.device.type == "cpu":
+        block_size, block_size_trunc, num_warps = 256, 128, 4
+    elif logits.device.type == "xpu":
+        block_size, block_size_trunc, num_warps = 4096, 2048, 4
+    else:
+        block_size, block_size_trunc, num_warps = 8192, 4096, 8
+        # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so
+        # per-tile latency bounds kernel latency, and Triton's default of 4
+        # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is
+        # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.
+    split_covers_ponly = (
+        current_platform.is_cuda_alike()
+        and topp_enabled
+        and batch_size <= _SPLIT_MAX_BATCH
+    )
+    return (num_programs,), dict(
+        K=k_ptr,
+        P=p_ptr,
+        BATCH_SIZE=batch_size,
+        VOCAB_SIZE=vocab_size,
+        BLOCK_SIZE=block_size,
+        BLOCK_SIZE_TRUNC=block_size_trunc,
+        TOPK_ENABLED=topk_enabled,
+        TOPP_ENABLED=topp_enabled,
+        SPLIT_COVERS_PONLY=split_covers_ponly,
+        num_warps=num_warps,
+    )
+
+
 # ---------------------------------------------------------------------------
 # Split-row pipeline for small batches.
 #
@@ -892,6 +989,13 @@ _SPLIT_ROUNDS = 5
 _TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {}
 
 
+def _topp_split_count(batch_size: int, num_sm: int) -> int:
+    splits = 1
+    while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:
+        splits *= 2
+    return splits
+
+
 @triton.jit
 def _topp_sb_stats_kernel(
     LOGITS,
@@ -939,6 +1043,51 @@ def _topp_sb_stats_kernel(
     tl.store(STATS + base + 3, cnt)
 
 
+def _topp_split_stats_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    num_sm = num_compute_units()
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(
+        1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1
+    )
+    has_k: Any = WarmupChoices(False, True)
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        stats=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if has_k else None,
+        p=TritonWarmupTensor(torch.float32),
+        num_sm=num_sm,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topp_sb_stats_kernel,
+    warmup_inputs=_topp_split_stats_warmup_inputs,
+)
+def _topp_split_stats(
+    logits: torch.Tensor,
+    stats: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    splits = _topp_split_count(batch_size, num_sm)
+    return (logits.shape[0] * splits,), dict(
+        K=k if k is not None else logits,
+        HAS_K=k is not None,
+        VOCAB_SIZE=vocab_size,
+        S=splits,
+        BLOCK=8192,
+        num_warps=8,
+    )
+
+
 @triton.jit
 def _topp_sb_row_stats(STATS, row, S: tl.constexpr):
     """Combine per-slice partials into row (max, sum_exp, min, n_finite)."""
@@ -1183,6 +1332,58 @@ def _topp_sb_step_kernel(
     tl.store(pbase + fidx * 3 + 2, nmin)
 
 
+def _topp_split_step_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    num_sm = num_compute_units()
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(
+        1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1
+    )
+    round: Any = WarmupIntRange(0, _SPLIT_ROUNDS)
+    has_k: Any = WarmupChoices(False, True)
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        stats=TritonWarmupTensor(torch.float32),
+        parts=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if has_k else None,
+        p=TritonWarmupTensor(torch.float32),
+        round=round,
+        num_sm=num_sm,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topp_sb_step_kernel,
+    warmup_inputs=_topp_split_step_warmup_inputs,
+)
+def _topp_split_step(
+    logits: torch.Tensor,
+    stats: torch.Tensor,
+    parts: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor,
+    round: int,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    splits = _topp_split_count(batch_size, num_sm)
+    return (logits.shape[0] * splits,), dict(
+        K=k if k is not None else logits,
+        HAS_K=k is not None,
+        S=splits,
+        F=_SPLIT_FANOUT,
+        NUM_ROUNDS=_SPLIT_ROUNDS,
+        VOCAB_SIZE=vocab_size,
+        BLOCK=2048,
+        num_warps=8,
+    )
+
+
 @triton.jit
 def _topp_sb_mask_kernel(
     LOGITS,
@@ -1313,6 +1514,57 @@ def _topp_sb_mask_kernel(
         tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n)
 
 
+def _topp_split_mask_warmup_inputs(vllm_config: Any) -> dict[str, Any]:
+    vocab_size = vllm_config.model_config.get_vocab_size()
+    num_sm = num_compute_units()
+    logits_stride: Any = WarmupChoices(16, 2)
+    batch_size: Any = WarmupIntRange(
+        1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1
+    )
+    has_k: Any = WarmupChoices(False, True)
+    logits = TritonWarmupTensor(
+        torch.float32,
+        shape=(batch_size, vocab_size),
+        strides=(logits_stride, 1),
+    )
+    return dict(
+        logits=logits,
+        stats=TritonWarmupTensor(torch.float32),
+        parts=TritonWarmupTensor(torch.float32),
+        k=TritonWarmupTensor(torch.int32) if has_k else None,
+        p=TritonWarmupTensor(torch.float32),
+        mask_value=float("-inf"),
+        num_sm=num_sm,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_topp_sb_mask_kernel,
+    warmup_inputs=_topp_split_mask_warmup_inputs,
+)
+def _topp_split_mask(
+    logits: torch.Tensor,
+    stats: torch.Tensor,
+    parts: torch.Tensor,
+    k: torch.Tensor | None,
+    p: torch.Tensor,
+    mask_value: float,
+    num_sm: int,
+) -> DispatchSpec:
+    batch_size, vocab_size = logits.shape
+    splits = _topp_split_count(batch_size, num_sm)
+    return (logits.shape[0] * splits,), dict(
+        K=k if k is not None else logits,
+        HAS_K=k is not None,
+        S=splits,
+        F=_SPLIT_FANOUT,
+        NUM_ROUNDS=_SPLIT_ROUNDS,
+        VOCAB_SIZE=vocab_size,
+        BLOCK=8192,
+        num_warps=8,
+    )
+
+
 def _apply_topp_split(
     logits: torch.Tensor,
     k: torch.Tensor | None,
@@ -1321,11 +1573,6 @@ def _apply_topp_split(
     num_sm: int,
 ) -> None:
     """Split-row top-p pipeline for small batches; masks logits in place."""
-    batch_size, vocab_size = logits.shape
-    splits = 1
-    while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:
-        splits *= 2
-
     ws = _TRITON_SPLIT_CACHE.get(logits.device)
     if ws is None:
         ws = {
@@ -1336,53 +1583,31 @@ def _apply_topp_split(
         }
         _TRITON_SPLIT_CACHE[logits.device] = ws
 
-    k_ptr = k if k is not None else logits  # dummy pointer when HAS_K=False
-    has_k = k is not None
-
-    _topp_sb_stats_kernel[(batch_size * splits,)](
+    _topp_split_stats(
         logits,
-        logits.stride(0),
         ws["stats"],
-        k_ptr,
+        k,
         p,
-        HAS_K=has_k,
-        VOCAB_SIZE=vocab_size,
-        S=splits,
-        BLOCK=8192,
-        num_warps=8,
+        num_sm,
     )
     for round_i in range(_SPLIT_ROUNDS):
-        _topp_sb_step_kernel[(batch_size * splits,)](
+        _topp_split_step(
             logits,
-            logits.stride(0),
             ws["stats"],
             ws["parts"],
-            k_ptr,
+            k,
             p,
             round_i,
-            HAS_K=has_k,
-            S=splits,
-            F=_SPLIT_FANOUT,
-            NUM_ROUNDS=_SPLIT_ROUNDS,
-            VOCAB_SIZE=vocab_size,
-            BLOCK=2048,
-            num_warps=8,
+            num_sm,
         )
-    _topp_sb_mask_kernel[(batch_size * splits,)](
+    _topp_split_mask(
         logits,
-        logits.stride(0),
         ws["stats"],
         ws["parts"],
-        k_ptr,
+        k,
         p,
-        HAS_K=has_k,
-        MASK_VALUE=mask_value,
-        S=splits,
-        F=_SPLIT_FANOUT,
-        NUM_ROUNDS=_SPLIT_ROUNDS,
-        VOCAB_SIZE=vocab_size,
-        BLOCK=8192,
-        num_warps=8,
+        mask_value,
+        num_sm,
     )
 
 
@@ -1477,39 +1702,15 @@ def apply_top_k_top_p_triton(
     else:
         normal_cdf_to_sigma_table, percentile_to_std_table = tables
 
-    # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.
-    # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot
-    # approximation; use smaller tiles for accurate top-p results.
-    launch_kwargs = {}
-    if logits.device.type == "cpu":
-        block_size, block_size_trunc = 256, 128
-    elif logits.device.type == "xpu":
-        block_size, block_size_trunc = 4096, 2048
-    else:
-        block_size, block_size_trunc = 8192, 4096
-        # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so
-        # per-tile latency bounds kernel latency, and Triton's default of 4
-        # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is
-        # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.
-        launch_kwargs["num_warps"] = 8
-
-    _topk_topp_kernel[(NUM_PROGRAMS,)](
+    _topk_topp(
         logits,
-        logits.stride(0),
         buffer,
         percentile_to_std_table,
         normal_cdf_to_sigma_table,
-        k_ptr,
-        p_ptr,
-        BATCH_SIZE=batch_size,
-        MASK_VALUE=mask_value,
-        VOCAB_SIZE=vocab_size,
-        BLOCK_SIZE=block_size,
-        BLOCK_SIZE_TRUNC=block_size_trunc,
-        TOPK_ENABLED=topk_enabled,
-        TOPP_ENABLED=topp_enabled,
-        SPLIT_COVERS_PONLY=use_split,
-        **launch_kwargs,
+        k_ptr if topk_enabled else None,
+        p_ptr if topp_enabled else None,
+        mask_value,
+        num_sm,
     )
     if use_split:
         _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm)

vllm/v1/sample/rejection_sampler.py

@@ -5,13 +5,19 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from dataclasses import replace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn as nn
 
 from vllm.config.model import PROCESSED_LOGPROBS_MODES
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup import WarmupChoices
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.triton_utils import tl, triton
 from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput
 from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor
@@ -88,6 +94,14 @@ class RejectionSampler(nn.Module):
                 device=device,
             )
         self.synthetic_mode = self.synthetic_conditional_rates is not None
+        _rejection_greedy_sample.register_warmup(synthetic_mode=self.synthetic_mode)
+        _rejection_random_sample.register_warmup(
+            synthetic_mode=self.synthetic_mode,
+        )
+        _expand.register_warmup()
+        _sample_recovered_tokens.register_warmup(
+            use_fp64_gumbel=self.use_fp64_gumbel,
+        )
 
     def forward(
         self,
@@ -453,7 +467,7 @@ def rejection_sample(
     if not sampling_metadata.all_random:
         # Rejection sampling for greedy sampling requests.
         target_argmax = target_logits.argmax(dim=-1)
-        rejection_greedy_sample_kernel[(batch_size,)](
+        _rejection_greedy_sample(
             output_token_ids,
             cu_num_draft_tokens,
             draft_token_ids,
@@ -463,7 +477,6 @@ def rejection_sample(
             max_spec_len,
             uniform_probs,
             synthetic_conditional_rates,
-            SYNTHETIC_MODE=synthetic_mode,
         )
         if sampling_metadata.all_greedy:
             return output_token_ids
@@ -488,7 +501,7 @@ def rejection_sample(
 
     # Rejection sampling for random sampling requests.
     assert uniform_probs is not None
-    rejection_random_sample_kernel[(batch_size,)](
+    _rejection_random_sample(
         output_token_ids,
         cu_num_draft_tokens,
         draft_token_ids,
@@ -499,10 +512,7 @@ def rejection_sample(
         uniform_probs,
         is_greedy,
         max_spec_len,
-        vocab_size,
         synthetic_conditional_rates,
-        NO_DRAFT_PROBS=draft_probs is None,
-        SYNTHETIC_MODE=synthetic_mode,
     )
     return output_token_ids
 
@@ -594,13 +604,12 @@ def expand_batch_to_tokens(
     batch_size = x.shape[0]
     assert cu_num_tokens.shape[0] == batch_size
     expanded_x = x.new_empty(num_tokens)
-    expand_kernel[(batch_size,)](
+    _expand(
         expanded_x,
         x,
         cu_num_tokens,
         replace_from,
         replace_to,
-        MAX_NUM_TOKENS=MAX_SPEC_LEN,  # To avoid recompilation.
     )
     return expanded_x
 
@@ -694,18 +703,14 @@ def sample_recovered_tokens(
     inv_q = q.reciprocal()
 
     recovered_token_ids = torch.empty_like(draft_token_ids)
-    BLOCK_SIZE = 8192
-    sample_recovered_tokens_kernel[(batch_size, max_spec_len)](
+    _sample_recovered_tokens(
         recovered_token_ids,
         cu_num_draft_tokens,
         draft_token_ids,
         draft_probs,
         target_probs,
         inv_q,
-        vocab_size,
-        BLOCK_SIZE,
-        NO_DRAFT_PROBS=draft_probs is None,
-        USE_FP64_GUMBEL=use_fp64_gumbel,
+        max_spec_len=max_spec_len,
     )
     return recovered_token_ids
 
@@ -769,6 +774,50 @@ def rejection_greedy_sample_kernel(
         )
 
 
+def _rejection_greedy_sample_warmup_inputs(
+    *, synthetic_mode: bool
+) -> dict[str, object]:
+    all_greedy = WarmupChoices(False, True)
+    int32 = TritonWarmupTensor(torch.int32)
+    return dict(
+        output_token_ids=int32,
+        cu_num_draft_tokens=int32,
+        draft_token_ids=int32,
+        target_argmax=TritonWarmupTensor(torch.int64),
+        bonus_token_ids=TritonWarmupTensor(torch.int64),
+        is_greedy=None if all_greedy else TritonWarmupTensor(torch.bool),
+        max_spec_len=5,
+        uniform_probs=(
+            None
+            if all_greedy and not synthetic_mode
+            else TritonWarmupTensor(torch.float64)
+        ),
+        synthetic_conditional_rates=(
+            TritonWarmupTensor(torch.float32) if synthetic_mode else None
+        ),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=rejection_greedy_sample_kernel,
+    warmup_inputs=_rejection_greedy_sample_warmup_inputs,
+)
+def _rejection_greedy_sample(
+    output_token_ids: torch.Tensor,
+    cu_num_draft_tokens: torch.Tensor,
+    draft_token_ids: torch.Tensor,
+    target_argmax: torch.Tensor,
+    bonus_token_ids: torch.Tensor,
+    is_greedy: torch.Tensor | None,
+    max_spec_len: int,
+    uniform_probs: torch.Tensor | None,
+    synthetic_conditional_rates: torch.Tensor | None,
+) -> DispatchSpec:
+    return (output_token_ids.shape[0],), dict(
+        SYNTHETIC_MODE=synthetic_conditional_rates is not None,
+    )
+
+
 # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
 @triton.jit(do_not_specialize=["max_spec_len"])
 def rejection_random_sample_kernel(
@@ -845,6 +894,57 @@ def rejection_random_sample_kernel(
         )
 
 
+def _rejection_random_sample_warmup_inputs(
+    *, synthetic_mode: bool
+) -> dict[str, object]:
+    vocab_size: Any = WarmupChoices(2, 16)
+    no_draft_probs = WarmupChoices(False, True)
+    int32 = TritonWarmupTensor(torch.int32)
+    probabilities = TritonWarmupTensor(
+        torch.float32,
+        shape=(1, vocab_size),
+    )
+    return dict(
+        output_token_ids=int32,
+        cu_num_draft_tokens=int32,
+        draft_token_ids=int32,
+        draft_probs=None if no_draft_probs else probabilities,
+        target_probs=probabilities,
+        bonus_token_ids=TritonWarmupTensor(torch.int64),
+        recovered_token_ids=int32,
+        uniform_probs=TritonWarmupTensor(torch.float64),
+        is_greedy=TritonWarmupTensor(torch.bool),
+        max_spec_len=5,
+        synthetic_conditional_rates=(
+            TritonWarmupTensor(torch.float32) if synthetic_mode else None
+        ),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=rejection_random_sample_kernel,
+    warmup_inputs=_rejection_random_sample_warmup_inputs,
+)
+def _rejection_random_sample(
+    output_token_ids: torch.Tensor,
+    cu_num_draft_tokens: torch.Tensor,
+    draft_token_ids: torch.Tensor,
+    draft_probs: torch.Tensor | None,
+    target_probs: torch.Tensor,
+    bonus_token_ids: torch.Tensor,
+    recovered_token_ids: torch.Tensor,
+    uniform_probs: torch.Tensor,
+    is_greedy: torch.Tensor,
+    max_spec_len: int,
+    synthetic_conditional_rates: torch.Tensor | None,
+) -> DispatchSpec:
+    return (output_token_ids.shape[0],), dict(
+        vocab_size=target_probs.shape[1],
+        NO_DRAFT_PROBS=draft_probs is None,
+        SYNTHETIC_MODE=synthetic_conditional_rates is not None,
+    )
+
+
 # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
 @triton.jit(do_not_specialize=["replace_from", "replace_to"])
 def expand_kernel(
@@ -869,6 +969,31 @@ def expand_kernel(
     tl.store(output_ptr + start_idx + offset, src_val, mask=offset < num_tokens)
 
 
+def _expand_warmup_inputs() -> dict[str, object]:
+    dtype = WarmupChoices(torch.float32, torch.int32, torch.int64)
+    data = TritonWarmupTensor(dtype)
+    return dict(
+        output=data,
+        input=data,
+        cu_num_tokens=TritonWarmupTensor(torch.int32),
+        replace_from=0,
+        replace_to=0,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=expand_kernel, warmup_inputs=_expand_warmup_inputs
+)
+def _expand(
+    output: torch.Tensor,
+    input: torch.Tensor,
+    cu_num_tokens: torch.Tensor,
+    replace_from: int,
+    replace_to: int,
+) -> DispatchSpec:
+    return (input.shape[0],), dict(MAX_NUM_TOKENS=MAX_SPEC_LEN)
+
+
 @triton.jit
 def sample_recovered_tokens_kernel(
     output_token_ids_ptr,  # [num_tokens]
@@ -951,3 +1076,49 @@ def sample_recovered_tokens_kernel(
 
     recovered_id = tl.minimum(recovered_id, vocab_size - 1)
     tl.store(output_token_ids_ptr + token_idx, recovered_id)
+
+
+def _sample_recovered_tokens_warmup_inputs(
+    *, use_fp64_gumbel: bool
+) -> dict[str, object]:
+    vocab_size: Any = WarmupChoices(2, 16)
+    no_draft_probs = WarmupChoices(False, True)
+    int32 = TritonWarmupTensor(torch.int32)
+    probabilities = TritonWarmupTensor(
+        torch.float32,
+        shape=(1, vocab_size),
+    )
+    return dict(
+        output_token_ids=int32,
+        cu_num_draft_tokens=int32,
+        draft_token_ids=int32,
+        draft_probs=None if no_draft_probs else probabilities,
+        target_probs=probabilities,
+        inv_q=TritonWarmupTensor(
+            torch.float64 if use_fp64_gumbel else torch.float32,
+            shape=(1, vocab_size),
+        ),
+        max_spec_len=5,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=sample_recovered_tokens_kernel,
+    warmup_inputs=_sample_recovered_tokens_warmup_inputs,
+)
+def _sample_recovered_tokens(
+    output_token_ids: torch.Tensor,
+    cu_num_draft_tokens: torch.Tensor,
+    draft_token_ids: torch.Tensor,
+    draft_probs: torch.Tensor | None,
+    target_probs: torch.Tensor,
+    inv_q: torch.Tensor,
+    *,
+    max_spec_len: int,
+) -> DispatchSpec:
+    return (cu_num_draft_tokens.shape[0], max_spec_len), dict(
+        vocab_size=target_probs.shape[1],
+        BLOCK_SIZE=8192,
+        NO_DRAFT_PROBS=draft_probs is None,
+        USE_FP64_GUMBEL=inv_q.dtype == torch.float64,
+    )

vllm/v1/spec_decode/dflash.py

@@ -13,7 +13,7 @@ from vllm.logger import init_logger
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
 from vllm.v1.spec_decode.utils import (
-    copy_and_expand_dflash_inputs_kernel,
+    _copy_and_expand_dflash_inputs,
     next_power_of_2,
 )
 
@@ -79,6 +79,13 @@ class DFlashProposer(SpecDecodeBaseProposer):
         self.dflash_causal = not dflash_has_any_non_causal(
             self.draft_model_config.hf_config
         )
+        _copy_and_expand_dflash_inputs.register_warmup(
+            block_table_stride=(self.max_model_len + self.block_size - 1)
+            // self.block_size,
+            parallel_drafting_token_id=self.parallel_drafting_token_id,
+            block_size=self.block_size,
+            num_speculative_tokens=self.num_speculative_tokens,
+        )
 
     @override
     def _create_draft_vllm_config(self) -> VllmConfig:
@@ -137,11 +144,9 @@ class DFlashProposer(SpecDecodeBaseProposer):
         max_ctx_per_req = cad.max_query_len
         max_tokens_per_req = max_ctx_per_req + num_query_per_req
         BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req))
-        num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE
-        grid = (batch_size, num_blocks)
 
         has_num_rejected = num_rejected_tokens_gpu is not None
-        copy_and_expand_dflash_inputs_kernel[grid](
+        _copy_and_expand_dflash_inputs(
             # Inputs
             next_token_ids_ptr=next_token_ids,
             target_positions_ptr=target_positions,
@@ -166,8 +171,9 @@ class DFlashProposer(SpecDecodeBaseProposer):
             num_query_per_req=num_query_per_req,
             num_speculative_tokens=self.num_speculative_tokens,
             total_input_tokens=num_context,
-            BLOCK_SIZE=BLOCK_SIZE,
-            HAS_NUM_REJECTED=has_num_rejected,
+            max_tokens_per_req=max_tokens_per_req,
+            triton_block_size=BLOCK_SIZE,
+            has_num_rejected=has_num_rejected,
         )
 
         query_slot_mapping = self._slot_mapping_buffer[:num_query_total]

vllm/v1/spec_decode/llm_base_proposer.py

@@ -52,10 +52,11 @@ from vllm.v1.sample.sampler import _SAMPLING_EPS
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.spec_decode.utils import (
     PADDING_SLOT_ID,
+    _copy_and_expand_eagle_inputs,
+    _eagle_prepare_inputs_padded,
+    _eagle_prepare_next_token_padded,
+    _eagle_step_slot_mapping_metadata,
     compute_new_slot_mapping,
-    copy_and_expand_eagle_inputs_kernel,
-    eagle_prepare_inputs_padded_kernel,
-    eagle_prepare_next_token_padded_kernel,
     eagle_step_update_slot_mapping_and_metadata,
     extend_all_queries_by_N,
     next_power_of_2,
@@ -140,6 +141,23 @@ class SpecDecodeBaseProposer:
 
         self.max_batch_size = vllm_config.scheduler_config.max_num_seqs
         self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        _eagle_step_slot_mapping_metadata.register_warmup(
+            max_model_len=self.max_model_len,
+            max_batch_size=self.max_batch_size,
+        )
+        _eagle_prepare_inputs_padded.register_warmup(max_batch_size=self.max_batch_size)
+        _eagle_prepare_next_token_padded.register_warmup(
+            vocab_size=vllm_config.model_config.get_vocab_size(),
+            num_sampled_tokens_per_req=self.num_speculative_tokens + 1,
+            max_batch_size=self.max_batch_size,
+        )
+        if self.needs_extra_input_slots:
+            _copy_and_expand_eagle_inputs.register_warmup(
+                parallel_drafting_token_id=self.parallel_drafting_token_id,
+                num_padding_slots_per_request=self.extra_slots_per_request,
+                shift_input_ids=self.pass_hidden_states_to_model,
+                max_num_tokens=self.max_num_tokens,
+            )
         self.token_arange_np = np.arange(self.max_num_tokens, dtype=np.int32)
 
         # Can be specialized by methods like DFlash to reduce the limit
@@ -872,13 +890,12 @@ class SpecDecodeBaseProposer:
             )
 
             # Kernel grid: one program per request (row)
-            grid = (batch_size, num_blocks)
             query_start_loc = cad.query_start_loc
             query_end_loc = cad.query_start_loc[1:] - 1
             if num_rejected_tokens_gpu is not None:
                 query_end_loc = query_end_loc - num_rejected_tokens_gpu
 
-            copy_and_expand_eagle_inputs_kernel[grid](
+            _copy_and_expand_eagle_inputs(
                 # (Padded) Inputs from the target model
                 target_token_ids_ptr=target_token_ids,
                 target_positions_ptr=target_positions,
@@ -901,6 +918,8 @@ class SpecDecodeBaseProposer:
                 num_padding_slots_per_request=self.extra_slots_per_request,
                 shift_input_ids=self.pass_hidden_states_to_model,
                 BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,
+                batch_size=batch_size,
+                num_blocks=num_blocks,
             )
             if self.pass_hidden_states_to_model:
                 assert self.parallel_drafting_hidden_state_tensor is not None
@@ -1065,11 +1084,10 @@ class SpecDecodeBaseProposer:
         valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)
 
         # Kernel grid: one program per request (row)
-        grid = (batch_size,)
 
         # Find the next power of 2 for block sizes
         BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens)
-        eagle_prepare_next_token_padded_kernel[grid](
+        _eagle_prepare_next_token_padded(
             sampled_token_ids,
             discard_request_mask,
             backup_tokens_gpu,
@@ -1108,8 +1126,7 @@ class SpecDecodeBaseProposer:
             (num_reqs,), dtype=torch.int32, device=device
         )
 
-        grid = (num_reqs,)
-        eagle_prepare_inputs_padded_kernel[grid](
+        _eagle_prepare_inputs_padded(
             spec_decode_metadata.cu_num_draft_tokens,
             valid_sampled_tokens_count,
             common_attn_metadata.query_start_loc,

vllm/v1/spec_decode/utils.py

@@ -1,7 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupChoices,
+    WarmupIntRange,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import (
@@ -85,6 +96,49 @@ def eagle_step_slot_mapping_metadata_kernel(
     tl.store(seq_lens_ptr + req_idx, new_seq_len)
 
 
+def _eagle_step_warmup_inputs(*, max_model_len: int, max_batch_size: int):
+    block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2)
+    n_blocks_per_req = triton.cdiv(max_model_len, block_size)
+    return dict(
+        positions=TritonWarmupTensor(torch.int64),
+        block_table=TritonWarmupTensor(
+            torch.int32, shape=(1, n_blocks_per_req), strides=(n_blocks_per_req, 1)
+        ),
+        block_table_stride=n_blocks_per_req,
+        seq_lens=TritonWarmupTensor(torch.int32),
+        out_clamped_positions=TritonWarmupTensor(torch.int64),
+        out_slot_mapping=TritonWarmupTensor(torch.int64),
+        block_size=block_size,
+        max_model_len=max_model_len,
+        n_blocks_per_req=n_blocks_per_req,
+        PAD_ID=PADDING_SLOT_ID,
+        batch_size=WarmupIntRange(1, max_batch_size + 1),
+        input_batch_size=1,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=eagle_step_slot_mapping_metadata_kernel,
+    warmup_inputs=_eagle_step_warmup_inputs,
+)
+def _eagle_step_slot_mapping_metadata(
+    positions: torch.Tensor,
+    block_table: torch.Tensor,
+    block_table_stride: int,
+    seq_lens: torch.Tensor,
+    out_clamped_positions: torch.Tensor,
+    out_slot_mapping: torch.Tensor,
+    *,
+    block_size: int,
+    max_model_len: int,
+    n_blocks_per_req: int,
+    PAD_ID: int,
+    batch_size: int,
+    input_batch_size: int,
+) -> DispatchSpec:
+    return (input_batch_size,), {}
+
+
 def eagle_step_update_slot_mapping_and_metadata(
     positions_1d: torch.Tensor,
     block_table_tensor: torch.Tensor,
@@ -118,7 +172,7 @@ def eagle_step_update_slot_mapping_and_metadata(
         input_batch_size = batch_size
 
     n_blocks_per_req = block_table_tensor.shape[1]
-    eagle_step_slot_mapping_metadata_kernel[(input_batch_size,)](
+    _eagle_step_slot_mapping_metadata(
         positions_1d,
         block_table_tensor,
         block_table_tensor.stride(0),
@@ -130,6 +184,7 @@ def eagle_step_update_slot_mapping_and_metadata(
         n_blocks_per_req=n_blocks_per_req,
         PAD_ID=PADDING_SLOT_ID,
         batch_size=batch_size,
+        input_batch_size=input_batch_size,
     )
 
 
@@ -175,6 +230,33 @@ def eagle_prepare_inputs_padded_kernel(
     tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens)
 
 
+def _eagle_prepare_inputs_warmup_inputs(*, max_batch_size: int):
+    int32 = TritonWarmupTensor(torch.int32)
+    return dict(
+        cu_num_draft_tokens=int32,
+        valid_sampled_tokens_count=int32,
+        query_start_loc_gpu=int32,
+        token_indices_to_sample=int32,
+        num_rejected_tokens_gpu=int32,
+        num_reqs=WarmupIntRange(1, max_batch_size + 1),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=eagle_prepare_inputs_padded_kernel,
+    warmup_inputs=_eagle_prepare_inputs_warmup_inputs,
+)
+def _eagle_prepare_inputs_padded(
+    cu_num_draft_tokens: torch.Tensor,
+    valid_sampled_tokens_count: torch.Tensor,
+    query_start_loc_gpu: torch.Tensor,
+    token_indices_to_sample: torch.Tensor,
+    num_rejected_tokens_gpu: torch.Tensor,
+    num_reqs: int,
+) -> DispatchSpec:
+    return (num_reqs,), {}
+
+
 @triton.jit
 def eagle_prepare_next_token_padded_kernel(
     sampled_token_ids_ptr,  # [num_reqs, num_sampled_tokens_per_req]
@@ -238,6 +320,49 @@ def eagle_prepare_next_token_padded_kernel(
         tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)
 
 
+def _eagle_prepare_next_token_warmup_inputs(
+    *, vocab_size: int, num_sampled_tokens_per_req: int, max_batch_size: int
+):
+    int32 = TritonWarmupTensor(torch.int32)
+    sampled_width: Any = WarmupIntRange(1, num_sampled_tokens_per_req + 1)
+    return dict(
+        sampled_token_ids=TritonWarmupTensor(
+            torch.int32,
+            shape=(1, sampled_width),
+            strides=(sampled_width, 1),
+        ),
+        discard_request_mask=TritonWarmupTensor(torch.bool),
+        backup_next_token_ids=int32,
+        next_token_ids=int32,
+        valid_sampled_tokens_count=int32,
+        vocab_size=vocab_size,
+        num_sampled_tokens_per_req=sampled_width,
+        num_reqs=WarmupIntRange(1, max_batch_size + 1),
+        stride_sampled_token_ids=sampled_width,
+        BLOCK_SIZE_TOKENS=next_power_of_2(sampled_width),
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=eagle_prepare_next_token_padded_kernel,
+    warmup_inputs=_eagle_prepare_next_token_warmup_inputs,
+)
+def _eagle_prepare_next_token_padded(
+    sampled_token_ids: torch.Tensor,
+    discard_request_mask: torch.Tensor,
+    backup_next_token_ids: torch.Tensor,
+    next_token_ids: torch.Tensor,
+    valid_sampled_tokens_count: torch.Tensor,
+    vocab_size: int,
+    num_sampled_tokens_per_req: int,
+    num_reqs: int,
+    stride_sampled_token_ids: int,
+    *,
+    BLOCK_SIZE_TOKENS: int,
+) -> DispatchSpec:
+    return (num_reqs,), {}
+
+
 def compute_new_slot_mapping(
     cad: CommonAttentionMetadata,
     new_positions: torch.Tensor,
@@ -454,6 +579,68 @@ def copy_and_expand_eagle_inputs_kernel(
     )
 
 
+def _copy_and_expand_eagle_warmup_inputs(
+    *,
+    parallel_drafting_token_id: int,
+    num_padding_slots_per_request: int,
+    shift_input_ids: bool,
+    max_num_tokens: int,
+):
+    int32 = TritonWarmupTensor(torch.int32)
+    int64 = TritonWarmupTensor(torch.int64)
+    boolean = TritonWarmupTensor(torch.bool)
+    return dict(
+        target_token_ids_ptr=int32,
+        target_positions_ptr=int64,
+        next_token_ids_ptr=int32,
+        out_input_ids_ptr=int32,
+        out_positions_ptr=int64,
+        out_is_rejected_token_mask_ptr=boolean,
+        out_is_masked_token_mask_ptr=boolean,
+        out_new_token_indices_ptr=int32,
+        out_hidden_state_mapping_ptr=int32,
+        query_start_loc_ptr=int32,
+        query_end_loc_ptr=int32,
+        padding_token_id=0,
+        parallel_drafting_token_id=parallel_drafting_token_id,
+        total_input_tokens=WarmupIntRange(1, max_num_tokens + 1),
+        num_padding_slots_per_request=num_padding_slots_per_request,
+        shift_input_ids=shift_input_ids,
+        BLOCK_SIZE_TOKENS=WarmupIntRange(1, 257, advance=lambda value: value * 2),
+        batch_size=1,
+        num_blocks=1,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=copy_and_expand_eagle_inputs_kernel,
+    warmup_inputs=_copy_and_expand_eagle_warmup_inputs,
+)
+def _copy_and_expand_eagle_inputs(
+    target_token_ids_ptr: torch.Tensor,
+    target_positions_ptr: torch.Tensor,
+    next_token_ids_ptr: torch.Tensor,
+    out_input_ids_ptr: torch.Tensor,
+    out_positions_ptr: torch.Tensor,
+    out_is_rejected_token_mask_ptr: torch.Tensor,
+    out_is_masked_token_mask_ptr: torch.Tensor,
+    out_new_token_indices_ptr: torch.Tensor,
+    out_hidden_state_mapping_ptr: torch.Tensor,
+    query_start_loc_ptr: torch.Tensor,
+    query_end_loc_ptr: torch.Tensor,
+    padding_token_id: int,
+    parallel_drafting_token_id: int,
+    total_input_tokens: int,
+    num_padding_slots_per_request: int,
+    shift_input_ids: bool,
+    *,
+    BLOCK_SIZE_TOKENS: int,
+    batch_size: int,
+    num_blocks: int,
+) -> DispatchSpec:
+    return (batch_size, num_blocks), {}
+
+
 @triton.jit
 def copy_and_expand_dflash_inputs_kernel(
     # Inputs
@@ -562,6 +749,75 @@ def copy_and_expand_dflash_inputs_kernel(
     )
 
 
+def _copy_and_expand_dflash_warmup_inputs(
+    *,
+    block_table_stride: int,
+    parallel_drafting_token_id: int,
+    block_size: int,
+    num_speculative_tokens: int,
+):
+    int32 = TritonWarmupTensor(torch.int32)
+    int64 = TritonWarmupTensor(torch.int64)
+    triton_block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2)
+    has_num_rejected = WarmupChoices(False, True)
+    return dict(
+        next_token_ids_ptr=int32,
+        target_positions_ptr=int64,
+        out_input_ids_ptr=int32,
+        out_context_positions_ptr=int64,
+        out_query_positions_ptr=int64,
+        out_context_slot_mapping_ptr=int64,
+        out_query_slot_mapping_ptr=int64,
+        out_token_indices_ptr=int32,
+        block_table_ptr=int32,
+        query_start_loc_ptr=int32,
+        num_rejected_tokens_ptr=(int32 if has_num_rejected else 0),
+        block_table_stride=block_table_stride,
+        parallel_drafting_token_id=parallel_drafting_token_id,
+        block_size=block_size,
+        num_query_per_req=1 + num_speculative_tokens,
+        num_speculative_tokens=num_speculative_tokens,
+        total_input_tokens=WarmupChoices(1, 2, 16),
+        max_tokens_per_req=triton_block_size,
+        triton_block_size=triton_block_size,
+        has_num_rejected=has_num_rejected,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=copy_and_expand_dflash_inputs_kernel,
+    warmup_inputs=_copy_and_expand_dflash_warmup_inputs,
+)
+def _copy_and_expand_dflash_inputs(
+    *,
+    next_token_ids_ptr: torch.Tensor,
+    target_positions_ptr: torch.Tensor,
+    out_input_ids_ptr: torch.Tensor,
+    out_context_positions_ptr: torch.Tensor,
+    out_query_positions_ptr: torch.Tensor,
+    out_context_slot_mapping_ptr: torch.Tensor,
+    out_query_slot_mapping_ptr: torch.Tensor,
+    out_token_indices_ptr: torch.Tensor,
+    block_table_ptr: torch.Tensor,
+    query_start_loc_ptr: torch.Tensor,
+    num_rejected_tokens_ptr: torch.Tensor | int,
+    block_table_stride: int,
+    parallel_drafting_token_id: int,
+    block_size: int,
+    num_query_per_req: int,
+    num_speculative_tokens: int,
+    total_input_tokens: int,
+    max_tokens_per_req: int,
+    triton_block_size: int,
+    has_num_rejected: bool,
+) -> DispatchSpec:
+    num_reqs = query_start_loc_ptr.shape[0] - 1
+    num_blocks = triton.cdiv(max_tokens_per_req, triton_block_size)
+    return (num_reqs, num_blocks), dict(
+        BLOCK_SIZE=triton_block_size,
+    )
+
+
 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
 def update_num_computed_tokens_for_batch_change(
     num_computed_tokens: torch.Tensor,

vllm/v1/worker/cpu_model_runner.py

@@ -80,16 +80,15 @@ class CPUModelRunner(GPUModelRunner):
 
         # Speculative decoding fallbacks
         import vllm.v1.sample.rejection_sampler
-        import vllm.v1.spec_decode.llm_base_proposer
         import vllm.v1.spec_decode.utils as spec_decode_utils
 
-        vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = (
+        spec_decode_utils._eagle_prepare_inputs_padded.kernel = (
             cpu_tl.eagle_prepare_inputs_padded_kernel
         )
-        vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_next_token_padded_kernel = (
+        spec_decode_utils._eagle_prepare_next_token_padded.kernel = (
             cpu_tl.eagle_prepare_next_token_padded_kernel
         )
-        vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = (
+        spec_decode_utils._copy_and_expand_eagle_inputs.kernel = (
             cpu_tl.copy_and_expand_eagle_inputs_kernel
         )
         spec_decode_utils.copy_and_expand_dflash_inputs_kernel = (

vllm/v1/worker/gpu/model_runner.py

@@ -188,6 +188,7 @@ logger = init_logger(__name__)
 
 
 class GPUModelRunner(LoRAModelRunnerMixin):
+    @JitWarmupRegistry.capture
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.vllm_config = vllm_config
         self.model_config = vllm_config.model_config
@@ -202,7 +203,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.speculative_config is not None and self.speculative_config.use_dspark()
         )
         self.observability_config = vllm_config.observability_config
-        self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
 
         self.device = device
         self.dtype = self.model_config.dtype
@@ -399,7 +399,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             logger.info_once("Loading model from scratch...")
 
             # Capture warmup providers selected while constructing the model.
-            with self.jit_warmup_registry.activate():
+            with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
                 self.model = model_loader.load_model(
                     vllm_config=self.vllm_config,
                     model_config=self.vllm_config.model_config,
@@ -461,47 +461,50 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         # Initialize samplers. Model states may override via custom_sampler().
         if self.is_last_pp_rank and not self.is_pooling_model:
-            sampler_kwargs: dict[str, Any] = {
-                "max_num_reqs": self.max_num_reqs,
-                "vocab_size": self.vocab_size,
-                "device": self.device,
-                "req_states": self.req_states,
-                "logprobs_mode": self.model_config.logprobs_mode,
-                "num_speculative_tokens": self.decode_query_len,
-                "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
-                "enable_trace_replay": self.model_config.enable_trace_replay,
-                "reasoning_config": self.vllm_config.reasoning_config,
-                "return_sampling_mask": self.model_config.return_sampling_mask,
-            }
-            if self.vllm_config.watermark_config is None:
-                self.sampler = Sampler(**sampler_kwargs)
-            else:
-                wm_config = self.vllm_config.watermark_config
-                watermarker = create_watermarker(wm_config)
-                if self.speculative_config is not None:
-                    watermarker = create_speculative_target_watermarker(watermarker)
-                self.sampler = GPUWatermarkSampler(
-                    watermarker,
-                    deduplicate_contexts=wm_config.deduplicate_contexts,
-                    deduplicate_contexts_max_history=(
-                        wm_config.deduplicate_contexts_max_history
-                    ),
-                    **sampler_kwargs,
-                )
-            custom = self.model_state.custom_sampler(self.sampler)
+            with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
+                sampler_kwargs: dict[str, Any] = {
+                    "max_num_reqs": self.max_num_reqs,
+                    "vocab_size": self.vocab_size,
+                    "device": self.device,
+                    "req_states": self.req_states,
+                    "logprobs_mode": self.model_config.logprobs_mode,
+                    "num_speculative_tokens": self.decode_query_len,
+                    "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
+                    "enable_trace_replay": self.model_config.enable_trace_replay,
+                    "reasoning_config": self.vllm_config.reasoning_config,
+                    "return_sampling_mask": self.model_config.return_sampling_mask,
+                }
+                if self.vllm_config.watermark_config is None:
+                    self.sampler = Sampler(**sampler_kwargs)
+                else:
+                    wm_config = self.vllm_config.watermark_config
+                    watermarker = create_watermarker(wm_config)
+                    if self.speculative_config is not None:
+                        watermarker = create_speculative_target_watermarker(watermarker)
+                    self.sampler = GPUWatermarkSampler(
+                        watermarker,
+                        deduplicate_contexts=wm_config.deduplicate_contexts,
+                        deduplicate_contexts_max_history=(
+                            wm_config.deduplicate_contexts_max_history
+                        ),
+                        **sampler_kwargs,
+                    )
+                custom = self.model_state.custom_sampler(self.sampler)
 
-            if custom:
-                self.vllm_config._check_watermarking_unsupported(custom_sampler=True)
-                self.sampler, self.rejection_sampler = custom
-            elif self.speculative_config is not None:
-                self.rejection_sampler = RejectionSampler(
-                    self.sampler,
-                    self.speculative_config,
-                    self.device,
-                    watermark_key=speculative_target_watermark_key(
-                        self.vllm_config.watermark_config
-                    ),
-                )
+                if custom:
+                    self.vllm_config._check_watermarking_unsupported(
+                        custom_sampler=True
+                    )
+                    self.sampler, self.rejection_sampler = custom
+                elif self.speculative_config is not None:
+                    self.rejection_sampler = RejectionSampler(
+                        self.sampler,
+                        self.speculative_config,
+                        self.device,
+                        watermark_key=speculative_target_watermark_key(
+                            self.vllm_config.watermark_config
+                        ),
+                    )
             self.prompt_logprobs_worker = PromptLogprobsWorker(
                 self.max_num_reqs,
                 logprobs_mode=self.model_config.logprobs_mode,
@@ -738,7 +741,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.speculator.init_cudagraph_manager(cudagraph_mode)
 
         # Capture warmup providers that depend on allocated KV-cache strides.
-        with self.jit_warmup_registry.activate():
+        with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
             kv_caches_dict = init_kv_cache(
                 self.compilation_config.static_forward_context,
                 self.kv_cache_config,
@@ -924,6 +927,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # during actual execution.
         assert self.sampler is not None
         self.sampler(logits, dummy_input_batch)
+        if self.rejection_sampler is not None:
+            # Warm head-dtype logits; scheduler warmup covers processed FP32 logits.
+            assert self.speculator is not None
+            self.rejection_sampler(
+                logits, dummy_input_batch, self.speculator.draft_logits
+            )
 
     @torch.inference_mode()
     def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:

vllm/v1/worker/gpu/sample/states.py

@@ -4,7 +4,10 @@ import numpy as np
 import torch
 
 from vllm.sampling_params import SamplingParams
-from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.sample.ops.topk_topp_sampler import (
+    apply_top_k_top_p,
+    register_top_k_top_p_warmups,
+)
 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
 from vllm.v1.worker.gpu.sample.gumbel import apply_temperature
 from vllm.v1.worker.gpu.sample.min_p import apply_min_p
@@ -18,6 +21,7 @@ class SamplingStates:
     def __init__(self, max_num_reqs: int, vocab_size: int):
         self.max_num_reqs = max_num_reqs
         self.vocab_size = vocab_size
+        register_top_k_top_p_warmups()
 
         self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
         self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32)

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import copy
+from types import SimpleNamespace
 from typing import Any
 
+import numpy as np
 import torch
 import torch.nn as nn
 
@@ -10,6 +12,15 @@ from vllm.config import VllmConfig, replace
 from vllm.config.compilation import CUDAGraphMode
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup import (
+    WarmupChoices,
+    WarmupIntRange,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -108,6 +119,7 @@ class DFlashSpeculator(DraftModelSpeculator):
 
         self.query_cudagraph_manager: DFlashCudaGraphManager | None = None
         self.draft_kv_cache_group_id: int = -1
+        prepare_dflash_inputs.register_warmup(speculator=self)
 
     @property
     def attn_vllm_config(self) -> VllmConfig:
@@ -682,6 +694,64 @@ def _prepare_dflash_inputs_kernel(
                 tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)
 
 
+def _prepare_dflash_warmup_inputs(*, speculator: DFlashSpeculator) -> dict[str, Any]:
+    int32 = TritonWarmupTensor(torch.int32)
+    int64 = TritonWarmupTensor(torch.int64)
+    float32 = TritonWarmupTensor(torch.float32)
+    input_buffers = SimpleNamespace(
+        input_ids=int32, positions=int64, query_start_loc=int32, seq_lens=int32
+    )
+    gid: Any = WarmupChoices(*speculator.draft_kv_cache_group_ids)
+    max_target_query_len: Any = WarmupIntRange(1, 257, advance=lambda value: value * 2)
+    block_table_stride = int(speculator.block_tables.input_block_tables[gid].stride(0))
+    block_table = TritonWarmupTensor(
+        torch.int32,
+        shape=(1, block_table_stride),
+        strides=(block_table_stride, 1),
+    )
+    input_batch = SimpleNamespace(
+        num_reqs=1,
+        num_scheduled_tokens=np.array([max_target_query_len]),
+        positions=int64,
+        query_start_loc=int32,
+        idx_mapping=int64,
+    )
+    return dict(
+        input_buffers=input_buffers,
+        query_slot_mapping=int64,
+        context_positions=int64,
+        context_slot_mapping=int64,
+        sample_indices=int64,
+        sample_pos=int64,
+        sample_idx_mapping=int32,
+        temperature=float32,
+        seeds=int64,
+        input_batch=input_batch,
+        num_sampled=int32,
+        num_rejected=int32,
+        last_sampled=int64,
+        next_prefill_tokens=int32,
+        input_temperature=float32,
+        input_seeds=int64,
+        block_table=block_table,
+        block_size=speculator.block_tables.kernel_block_sizes[gid],
+        cp_rank=speculator.block_tables.cp_rank,
+        cp_size=speculator.block_tables.cp_size,
+        cp_interleave=speculator.block_tables.cp_interleave,
+        parallel_drafting_token_id=speculator.parallel_drafting_token_id,
+        num_query_per_req=speculator.num_query_per_req,
+        num_speculative_steps=speculator.num_speculative_steps,
+        max_num_reqs=speculator.max_num_reqs,
+        max_num_tokens=speculator.max_num_tokens,
+        max_model_len=speculator.max_model_len,
+        sample_from_anchor=speculator.sample_from_anchor,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_prepare_dflash_inputs_kernel,
+    warmup_inputs=_prepare_dflash_warmup_inputs,
+)
 def prepare_dflash_inputs(
     input_buffers: InputBuffers,
     query_slot_mapping: torch.Tensor,
@@ -718,50 +788,33 @@ def prepare_dflash_inputs(
     max_num_tokens: int,
     max_model_len: int,
     sample_from_anchor: bool = False,
-) -> None:
+) -> DispatchSpec:
     num_reqs = input_batch.num_reqs
     assert num_reqs > 0
     # Cover the longest possible per-request span (ctx + query). Use the max
     # per-request query length, not the total token count across the batch.
     max_target_query_len = int(input_batch.num_scheduled_tokens.max())
     max_tokens_per_req = max_target_query_len + num_query_per_req
-    BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
-    num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
-    _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
-        input_buffers.input_ids,
-        input_buffers.positions,
-        input_buffers.query_start_loc,
-        input_buffers.seq_lens,
-        query_slot_mapping,
-        context_positions,
-        context_slot_mapping,
-        sample_indices,
-        sample_pos,
-        sample_idx_mapping,
-        temperature,
-        seeds,
-        input_batch.positions,
-        input_batch.query_start_loc,
-        input_batch.idx_mapping,
-        last_sampled,
-        next_prefill_tokens,
-        num_sampled,
-        num_rejected,
-        input_temperature,
-        input_seeds,
-        block_table,
-        block_table.stride(0),
-        parallel_drafting_token_id,
-        block_size,
-        num_query_per_req,
-        num_speculative_steps,
-        max_num_reqs,
-        max_num_tokens,
-        max_model_len,
-        cp_rank,
-        SAMPLE_FROM_ANCHOR=sample_from_anchor,
+    block = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
+    num_blocks = triton.cdiv(max_tokens_per_req, block)
+    return (num_reqs, num_blocks), dict(
+        out_input_ids_ptr=input_buffers.input_ids,
+        out_query_positions_ptr=input_buffers.positions,
+        out_query_start_loc_ptr=input_buffers.query_start_loc,
+        out_seq_lens_ptr=input_buffers.seq_lens,
+        out_query_slot_mapping_ptr=query_slot_mapping,
+        out_context_positions_ptr=context_positions,
+        out_context_slot_mapping_ptr=context_slot_mapping,
+        out_sample_indices_ptr=sample_indices,
+        out_sample_pos_ptr=sample_pos,
+        out_sample_idx_mapping_ptr=sample_idx_mapping,
+        out_temperature_ptr=temperature,
+        out_seeds_ptr=seeds,
+        target_positions_ptr=input_batch.positions,
+        target_query_start_loc_ptr=input_batch.query_start_loc,
+        idx_mapping_ptr=input_batch.idx_mapping,
+        temperature_ptr=input_temperature,
+        seeds_ptr=input_seeds,
         PAD_SLOT_ID=PAD_SLOT_ID,
-        CP_SIZE=cp_size,
-        CP_INTERLEAVE=cp_interleave,
-        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_SIZE=block,
     )

vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py

@@ -9,6 +9,11 @@ from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    DispatchSpec,
+    TritonWarmupTensor,
+    triton_kernel_dispatcher_with_warmup,
+)
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
@@ -69,6 +74,9 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
 
         self.cudagraph_manager: SpeculatorCudaGraphManager | None = None
 
+        _shift_input_ids.register_warmup()
+        _shift_input_embeds.register_warmup(speculator=self)
+
     def load_draft_model(
         self,
         target_model: nn.Module,
@@ -1008,6 +1016,32 @@ def _shift_input_ids_kernel(
     tl.store(input_ids_ptr + last_token_index, draft_token)
 
 
+def _shift_input_ids_warmup_inputs() -> dict[str, Any]:
+    return dict(
+        input_ids=TritonWarmupTensor(torch.int32),
+        idx_mapping=TritonWarmupTensor(torch.int32),
+        query_start_loc=TritonWarmupTensor(torch.int32),
+        last_token_indices=TritonWarmupTensor(torch.int64),
+        draft_tokens=TritonWarmupTensor(torch.int64),
+        num_reqs=1,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_shift_input_ids_kernel,
+    warmup_inputs=_shift_input_ids_warmup_inputs,
+)
+def _shift_input_ids(
+    input_ids: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    last_token_indices: torch.Tensor,
+    draft_tokens: torch.Tensor,
+    num_reqs: int,
+) -> DispatchSpec:
+    return (num_reqs,), dict(BLOCK_SIZE=1024)
+
+
 @triton.jit
 def _shift_input_embeds_kernel(
     input_embeds_ptr,
@@ -1065,6 +1099,44 @@ def _shift_input_embeds_kernel(
     )
 
 
+def _shift_input_embeds_warmup_inputs(
+    *, speculator: MultiModuleMTPSpeculator
+) -> dict[str, Any]:
+    hidden = TritonWarmupTensor(
+        speculator.dtype,
+        shape=(1, speculator.hidden_size),
+        strides=(speculator.hidden_size, 1),
+    )
+    return dict(
+        input_embeds=hidden,
+        draft_embeds=hidden,
+        idx_mapping=TritonWarmupTensor(torch.int32),
+        query_start_loc=TritonWarmupTensor(torch.int32),
+        last_token_indices=TritonWarmupTensor(torch.int64),
+        num_reqs=1,
+    )
+
+
+@triton_kernel_dispatcher_with_warmup(
+    kernel=_shift_input_embeds_kernel,
+    warmup_inputs=_shift_input_embeds_warmup_inputs,
+)
+def _shift_input_embeds(
+    input_embeds: torch.Tensor,
+    draft_embeds: torch.Tensor,
+    idx_mapping: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    last_token_indices: torch.Tensor,
+    num_reqs: int,
+) -> DispatchSpec:
+    hidden_size = input_embeds.shape[-1]
+    return (num_reqs, triton.cdiv(hidden_size, 256)), dict(
+        hidden_size=hidden_size,
+        BLOCK_SIZE_Q=16,
+        BLOCK_SIZE_H=256,
+    )
+
+
 def update_draft_inputs(
     draft_tokens: torch.Tensor,
     draft_embeds: torch.Tensor | None,
@@ -1074,29 +1146,21 @@ def update_draft_inputs(
     idx_mapping: torch.Tensor,
     num_reqs: int,
 ) -> None:
-    _shift_input_ids_kernel[(num_reqs,)](
+    _shift_input_ids(
         input_buffers.input_ids,
         idx_mapping,
         input_buffers.query_start_loc,
         last_token_indices,
         draft_tokens,
-        BLOCK_SIZE=1024,
+        num_reqs=num_reqs,
     )
     if input_embeds is not None:
         assert draft_embeds is not None
-        hidden_size = input_embeds.shape[-1]
-        hidden_block_size = 256
-        _shift_input_embeds_kernel[
-            (num_reqs, triton.cdiv(hidden_size, hidden_block_size))
-        ](
+        _shift_input_embeds(
             input_embeds,
-            input_embeds.stride(0),
             draft_embeds,
-            draft_embeds.stride(0),
             idx_mapping,
             input_buffers.query_start_loc,
             last_token_indices,
-            hidden_size,
-            BLOCK_SIZE_Q=16,
-            BLOCK_SIZE_H=hidden_block_size,
+            num_reqs=num_reqs,
         )

vllm/v1/worker/gpu_model_runner.py

@@ -498,6 +498,7 @@ class ExecuteModelState(NamedTuple):
 class GPUModelRunner(
     LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin
 ):
+    @JitWarmupRegistry.capture
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -514,7 +515,6 @@ class GPUModelRunner(
         self.scheduler_config = vllm_config.scheduler_config
         self.speculative_config = vllm_config.speculative_config
         self.observability_config = vllm_config.observability_config
-        self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
 
         model_config = self.model_config
         cache_config = self.cache_config
@@ -750,36 +750,35 @@ class GPUModelRunner(
             self.parallel_config.cp_kv_cache_interleave_size
         )
         # Capture warmup providers registered by the initial placeholder InputBatch
-        with self.jit_warmup_registry.activate():
-            self.input_batch = InputBatch(
-                max_num_reqs=self.max_num_reqs,
-                # We need to use the encoder length for encoder-decoder
-                # because of KV cache for cross-attention.
-                max_model_len=max(self.max_model_len, self.max_encoder_len),
-                max_num_batched_tokens=self.max_num_tokens,
-                device=self.device,
-                vocab_size=self.model_config.get_vocab_size(),
-                block_sizes=[placeholder_block_size],
-                kernel_block_sizes=[placeholder_block_size],
-                max_num_blocks_per_req=[placeholder_max_num_blocks],
-                num_spec_tokens=self.num_spec_tokens,
-                logitsprocs=build_logitsprocs(
-                    self.vllm_config,
-                    self.device,
-                    PIN_MEMORY,
-                    self.is_pooling_model,
-                    custom_logitsprocs,
-                ),
-                # We currently don't know whether a particular custom logits processor
-                # uses output token ids so we set this conservatively. Thinking-budget
-                # tracking is requested dynamically when a budgeted request is in the
-                # batch.
-                logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
-                is_pooling_model=self.is_pooling_model,
-                cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
-                reasoning_config=self.vllm_config.reasoning_config,
-                use_replayssm=self.cache_config.use_replayssm,
-            )
+        self.input_batch = InputBatch(
+            max_num_reqs=self.max_num_reqs,
+            # We need to use the encoder length for encoder-decoder
+            # because of KV cache for cross-attention.
+            max_model_len=max(self.max_model_len, self.max_encoder_len),
+            max_num_batched_tokens=self.max_num_tokens,
+            device=self.device,
+            vocab_size=self.model_config.get_vocab_size(),
+            block_sizes=[placeholder_block_size],
+            kernel_block_sizes=[placeholder_block_size],
+            max_num_blocks_per_req=[placeholder_max_num_blocks],
+            num_spec_tokens=self.num_spec_tokens,
+            logitsprocs=build_logitsprocs(
+                self.vllm_config,
+                self.device,
+                PIN_MEMORY,
+                self.is_pooling_model,
+                custom_logitsprocs,
+            ),
+            # We currently don't know whether a particular custom logits processor
+            # uses output token ids so we set this conservatively. Thinking-budget
+            # tracking is requested dynamically when a budgeted request is in the
+            # batch.
+            logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
+            is_pooling_model=self.is_pooling_model,
+            cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
+            reasoning_config=self.vllm_config.reasoning_config,
+            use_replayssm=self.cache_config.use_replayssm,
+        )
 
         # Separate cuda stream for overlapping transfer of sampled token ids from
         # GPU to CPU when async scheduling is enabled.
@@ -5332,7 +5331,7 @@ class GPUModelRunner(
                     self.load_config.load_format = "dummy"
                 model_loader = get_model_loader(self.load_config)
                 # Capture warmup providers selected while constructing the model.
-                with self.jit_warmup_registry.activate():
+                with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
                     self.model = model_loader.load_model(
                         vllm_config=self.vllm_config, model_config=self.model_config
                     )
@@ -7304,7 +7303,7 @@ class GPUModelRunner(
                 self.parallel_config.cp_kv_cache_interleave_size
             )
             # Capture warmup providers registered after final KV-cache geometry is known
-            with self.jit_warmup_registry.activate():
+            with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
                 self.input_batch = InputBatch(
                     max_num_reqs=self.max_num_reqs,
                     max_model_len=max_model_len,
@@ -7454,7 +7453,7 @@ class GPUModelRunner(
         # Reinitialize need to after initialize_attn_backend
         self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
         # Capture warmup providers that depend on allocated KV-cache strides.
-        with self.jit_warmup_registry.activate():
+        with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
             kv_caches = self.initialize_kv_cache_tensors(
                 kv_cache_config,
                 kernel_block_sizes,