From aed894c190aa3db74c502d521b3bc3935ad85de1 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:24:24 +0200 Subject: [PATCH] [6/N][warmup][DSv4] Migrate sampling, and DFlash JIT kernels (#56323) Signed-off-by: Lucas Wilkinson Signed-off-by: LopezCastroRoberto Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .buildkite/test_areas/models_basic.yaml | 2 +- .buildkite/test_areas/samplers.yaml | 2 +- tests/model_executor/test_jit_warmup.py | 32 ++ .../test_jit_warmup_triton_launcher.py | 335 ++++++++++- .../test_spec_decode_rejection_warmup.py | 121 ---- tests/v1/sample/test_topk_topp_sampler.py | 4 + tests/v1/shutdown/test_delete.py | 16 +- .../spec_decode/test_dflash_prepare_inputs.py | 2 +- .../test_gpu_autoregressive_speculator.py | 1 + .../layers/mamba/ops/scatter_states.py | 49 +- vllm/model_executor/warmup/jit_warmup.py | 345 ++++++++++-- .../warmup/jit_warmup_triton_helper.py | 520 +++++++++++++++++- vllm/model_executor/warmup/kernel_warmup.py | 9 +- .../warmup/spec_decode_rejection_warmup.py | 139 ----- vllm/models/glm5next/nvidia/kda.py | 9 + vllm/v1/sample/ops/topk_topp_sampler.py | 21 +- vllm/v1/sample/ops/topk_topp_triton.py | 333 ++++++++--- vllm/v1/sample/rejection_sampler.py | 201 ++++++- vllm/v1/spec_decode/dflash.py | 18 +- vllm/v1/spec_decode/llm_base_proposer.py | 35 +- vllm/v1/spec_decode/utils.py | 258 ++++++++- vllm/v1/worker/cpu_model_runner.py | 7 +- vllm/v1/worker/gpu/model_runner.py | 95 ++-- vllm/v1/worker/gpu/sample/states.py | 6 +- .../gpu/spec_decode/dflash/speculator.py | 131 +++-- .../multi_module_mtp/speculator.py | 88 ++- vllm/v1/worker/gpu_model_runner.py | 67 ++- 27 files changed, 2288 insertions(+), 558 deletions(-) delete mode 100644 tests/model_executor/test_spec_decode_rejection_warmup.py delete mode 100644 vllm/model_executor/warmup/spec_decode_rejection_warmup.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index bf453a12d8..870cc9b393 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.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/" diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index cfde3717c8..dae7e96468 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.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 diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py index 62ecb05f6b..7d38f4b40d 100644 --- a/tests/model_executor/test_jit_warmup.py +++ b/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() diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index e76041532f..b36361e637 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/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 diff --git a/tests/model_executor/test_spec_decode_rejection_warmup.py b/tests/model_executor/test_spec_decode_rejection_warmup.py deleted file mode 100644 index 6abfb1a4d1..0000000000 --- a/tests/model_executor/test_spec_decode_rejection_warmup.py +++ /dev/null @@ -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) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index fdf7041692..5a0ce5cc08 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/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 diff --git a/tests/v1/shutdown/test_delete.py b/tests/v1/shutdown/test_delete.py index 991ad68918..3a513945ea 100644 --- a/tests/v1/shutdown/test_delete.py +++ b/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"], diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 16d6d8e516..ce47cb8ae6 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/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 diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 920c9a25f0..a38265a6f7 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/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 diff --git a/vllm/model_executor/layers/mamba/ops/scatter_states.py b/vllm/model_executor/layers/mamba/ops/scatter_states.py index 0102b6cf12..04b5cb61c2 100644 --- a/vllm/model_executor/layers/mamba/ops/scatter_states.py +++ b/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, diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index fc9d6fcdf8..72626964ae 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/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=[]) + ), + "", + "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: diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index b019404cc9..dab055902a 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/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. diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index bbb0a66d53..e56ae50ebb 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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) diff --git a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py b/vllm/model_executor/warmup/spec_decode_rejection_warmup.py deleted file mode 100644 index 6f6f08a2d5..0000000000 --- a/vllm/model_executor/warmup/spec_decode_rejection_warmup.py +++ /dev/null @@ -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 diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index e6ba6ac5c6..192327cf4a 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/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. diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index d714976d85..aeb82dec42 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/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, diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 7c1f080fcd..695ba624f7 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/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) diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index a1baf0a808..5aa8345473 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/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, + ) diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 42d8ef035c..284db7dee7 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/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] diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 2edbd27d6e..c0e456521c 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/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, diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py index 65b9408a89..fbfe443842 100644 --- a/vllm/v1/spec_decode/utils.py +++ b/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, diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 67ddf83bae..04ca2b9194 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/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 = ( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0ce0288f31..14e0ac4629 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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: diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index fe4dee6a6b..08cae82128 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/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) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index a5710829e0..10f4346b54 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/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, ) diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 1c6ae315cf..1c37ccb6c2 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/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, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d6a27b1f8b..97a59bd273 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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,