@tduncan / Nvidia-Optimization-Stack-VLLM / commits / aed894c190

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

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

Roberto L. Castro committed Sep 13, 2026 at 02:24 UTC aed894c190aa3db74c502d521b3bc3935ad85de1

27 files changed +2286 -556

.buildkite/test_areas/models_basic.yaml

+1 -1

@@ -41,7 +41,7 @@ steps:

41 - label: ":nvidia: (H200 MIG 35GB) Basic Models (Other)"

42 device: h200_35gb

43 key: basic-models-tests-other

44 - timeout_in_minutes: 35

44 + timeout_in_minutes: 38

45 source_file_dependencies:

46 - vllm/

47 - "!vllm/distributed/kv_transfer/"

.buildkite/test_areas/samplers.yaml

+1 -1

@@ -5,7 +5,7 @@ steps:

5 - label: ":nvidia: (H200 MIG 35GB) Samplers"

6 device: h200_35gb

7 key: samplers-test

8 - timeout_in_minutes: 40

8 + timeout_in_minutes: 60

9 source_file_dependencies:

10 - vllm/model_executor/layers

11 - vllm/sampling_metadata.py

tests/model_executor/test_jit_warmup.py

+32

@@ -171,6 +171,19 @@ def test_warmup_range_validates_custom_advancement() -> None:

171 )

172

173

174 +def test_warmup_cases_support_lambda_advancement() -> None:

175 + def warmup_inputs() -> dict[str, Any]:

176 + value: Any = WarmupIntRange(1, 9, advance=lambda value: value * 2)

177 + return dict(value=value)

178 +

179 + assert list(ToyKernel()._expand_warmup_cases(warmup_inputs)) == [

180 + {"value": 1},

181 + {"value": 2},

182 + {"value": 4},

183 + {"value": 8},

184 + ]

185 +

186 +

187 def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:

188 cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)

189

@@ -577,6 +590,25 @@ def test_registry_records_only_inside_model_setup_context() -> None:

590 assert kernel.compiled == []

591

592

593 +def test_registry_capture() -> None:

594 + class Owner:

595 + @JitWarmupRegistry.capture

596 + def __init__(

597 + self,

598 + vllm_config: Any,

599 + kernel: RecordingToyKernel,

600 + ) -> None:

601 + kernel.register_warmup(3, vllm_config)

602 +

603 + kernel = RecordingToyKernel()

604 + config = _config()

605 + owner = Owner(config, kernel)

606 +

607 + assert len(owner.jit_warmup_registry) == 1 # type: ignore[attr-defined]

608 + kernel.register_warmup(5, config)

609 + assert len(owner.jit_warmup_registry) == 1 # type: ignore[attr-defined]

610 +

611 +

612 def test_registry_expands_requests_and_deduplicates_owner_keys() -> None:

613 registry = JitWarmupRegistry(_config())

614 kernel = RecordingToyKernel()

tests/model_executor/test_jit_warmup_triton_launcher.py

+333 -2

@@ -4,22 +4,68 @@

4 from dataclasses import dataclass

5 from typing import Any

7 -from vllm.model_executor.warmup.jit_warmup import kernel_launcher

7 +import pytest

8 +import torch

9 +

10 +from vllm.model_executor.warmup import jit_warmup_triton_helper

11 +from vllm.model_executor.warmup.jit_warmup import (

12 + WarmupChoices,

13 + WarmupIntRange,

14 + kernel_launcher,

15 +)

16 from vllm.model_executor.warmup.jit_warmup_triton_helper import (

17 LaunchSpec,

18 + TritonJitKey,

19 + TritonWarmupTensor,

20 VllmTritonJitKernel,

21 + triton_kernel_dispatcher_with_warmup,

22 + triton_warmup_inputs,

23 )

24 +from vllm.triton_utils import tl, triton

25 +

26 +

27 +@triton.jit

28 +def _binder_test_kernel(x, value, BLOCK: tl.constexpr):

29 + pass

30 +

31 +

32 +@triton.jit

33 +def _second_binder_test_kernel(x, value, BLOCK: tl.constexpr):

34 + pass

35 +

36 +

37 +@triton.jit

38 +def _pointer_group_test_kernel(x_ptr, value, BLOCK: tl.constexpr):

39 + pass

42 class _FakeTritonKernel:

15 - arg_names = ("first", "second", "CONST")

43 + arg_names: tuple[str, ...] = ("first", "second", "CONST")

45 def __init__(self) -> None:

46 self.warmup_calls: list[dict[str, Any]] = []

47 + self.runtime_calls: list[tuple[Any, tuple[Any, ...], dict[str, Any]]] = []

49 def warmup(self, **kwargs: Any) -> None:

50 self.warmup_calls.append(kwargs)

52 + def __getitem__(self, grid: Any) -> Any:

53 + def launch(*args: Any, **kwargs: Any) -> None:

54 + self.runtime_calls.append((grid, args, kwargs))

55 +

56 + return launch

57 +

58 +

59 +def _patch_key_deriver(

60 + monkeypatch: pytest.MonkeyPatch,

61 + derive: Any,

62 +) -> None:

63 + monkeypatch.setattr(

64 + jit_warmup_triton_helper,

65 + "_triton_key_deriver",

66 + lambda kernel: lambda kwargs: derive(kernel, kwargs),

67 + )

68 +

70 class _TestTritonKernel(VllmTritonJitKernel["_TestTritonKernel.CompileKey"]):

71 kernel = _FakeTritonKernel()

@@ -91,6 +137,291 @@ def test_triton_launcher_supports_cpu_function_wrappers() -> None:

137 assert calls == [("runtime", 2, 7)]

138

139

140 +def test_triton_kernel_decorator_returns_launcher(

141 + monkeypatch: pytest.MonkeyPatch,

142 +) -> None:

143 + kernel = _FakeTritonKernel()

144 + kernel.arg_names = (

145 + "FIRST_PTR",

146 + "FIRST_STRIDE0",

147 + "aliased_ptr",

148 + "aliased_stride_0",

149 + "SECOND",

150 + "CONST",

151 + )

152 +

153 + def warmup_inputs() -> dict[str, Any]:

154 + return dict(

155 + first=TritonWarmupTensor(torch.float32, shape=(2, 3), strides=(5, 1)),

156 + second=WarmupChoices(1, 2),

157 + )

158 +

159 + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)

160 + def launch(first: str, second: int = 2, config: int = 7) -> LaunchSpec:

161 + return (2,), dict(aliased_ptr=first, CONST=config)

162 +

163 + def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]:

164 + return {TritonJitKey(id(kernel), "fake", 0, kwargs["SECOND"])}

165 +

166 + _patch_key_deriver(monkeypatch, fake_keys)

167 +

168 + keys = launch.get_warmup_keys()

169 + keys_with_config = launch.get_warmup_keys(vllm_config=object())

170 + assert keys_with_config == keys

171 + assert [dict(key.inputs)["second"] for key in keys] == [1, 2]

172 + launch.compile(keys[0])

173 + assert kernel.warmup_calls == [

174 + {

175 + "grid": (1,),

176 + "FIRST_PTR": TritonWarmupTensor(

177 + torch.float32, shape=(2, 3), strides=(5, 1)

178 + ),

179 + "FIRST_STRIDE0": 5,

180 + "aliased_ptr": TritonWarmupTensor(

181 + torch.float32, shape=(2, 3), strides=(5, 1)

182 + ),

183 + "aliased_stride_0": 5,

184 + "SECOND": 1,

185 + "CONST": 7,

186 + }

187 + ]

188 + first = torch.empty_strided((2, 3), (5, 1))

189 + launch(first)

190 + assert kernel.runtime_calls == [

191 + (

192 + (2,),

193 + (),

194 + {

195 + "FIRST_PTR": first,

196 + "FIRST_STRIDE0": 5,

197 + "aliased_ptr": first,

198 + "aliased_stride_0": 5,

199 + "SECOND": 2,

200 + "CONST": 7,

201 + },

202 + )

203 + ]

204 + assert launch.__name__ == "launch"

205 + with pytest.raises(TypeError, match="unexpected keyword"):

206 + launch(first, 1, 7, stale_constexpr=True)

207 +

208 +

209 +def test_triton_kernel_decorator_exhausts_large_ranges_before_deduplication(

210 + monkeypatch: pytest.MonkeyPatch,

211 +) -> None:

212 + kernel = _FakeTritonKernel()

213 + dispatched: list[int] = []

214 +

215 + def warmup_inputs() -> dict[str, Any]:

216 + tokens: Any = WarmupIntRange(1, 8193)

217 + return dict(first="warmup", second=tokens)

218 +

219 + def specialization(second: int) -> int:

220 + return 7 if second <= 97 else 8

221 +

222 + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)

223 + def dispatch(first: str, second: int) -> LaunchSpec:

224 + dispatched.append(second)

225 + return (second,), dict(CONST=specialization(second))

226 +

227 + def fake_keys(kernel: Any, kwargs: Any) -> set[TritonJitKey]:

228 + return {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])}

229 +

230 + _patch_key_deriver(monkeypatch, fake_keys)

231 +

232 + keys = dispatch.get_warmup_keys()

233 + assert dispatched == list(range(1, 8193))

234 + assert len(keys) == 2

235 + assert {specialization(dict(key.inputs)["second"]) for key in keys} == {7, 8}

236 +

237 +

238 +def test_triton_kernel_decorator_propagates_dispatch_assertions(

239 + monkeypatch: pytest.MonkeyPatch,

240 +) -> None:

241 + kernel = _FakeTritonKernel()

242 +

243 + def warmup_inputs() -> dict[str, Any]:

244 + return dict(first="warmup", second=WarmupChoices(1, 2))

245 +

246 + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)

247 + def dispatch(first: str, second: int) -> LaunchSpec:

248 + assert second != 2, "broken dispatch"

249 + return (1,), dict(CONST=second)

250 +

251 + _patch_key_deriver(

252 + monkeypatch,

253 + lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])},

254 + )

255 +

256 + with pytest.raises(AssertionError, match="broken dispatch"):

257 + dispatch.get_warmup_keys()

258 +

259 +

260 +def test_triton_kernel_dispatch_uses_cuda_fake_tensors(

261 + monkeypatch: pytest.MonkeyPatch,

262 +) -> None:

263 + kernel = _FakeTritonKernel()

264 +

265 + def warmup_inputs() -> dict[str, Any]:

266 + return dict(

267 + first=TritonWarmupTensor(torch.float32, shape=(2, 3)),

268 + second=1,

269 + )

270 +

271 + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)

272 + def dispatch(first: torch.Tensor, second: int) -> LaunchSpec:

273 + assert isinstance(first, torch.Tensor)

274 + assert first.is_cuda

275 + assert first[0].is_contiguous()

276 + return (first.shape[0],), dict(CONST=first[0].numel())

277 +

278 + _patch_key_deriver(

279 + monkeypatch,

280 + lambda kernel, kwargs: {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])},

281 + )

282 +

283 + keys = dispatch.get_warmup_keys()

284 + assert len(keys) == 1

285 + assert dict(keys[0].inputs)["first"].device.type == "cuda"

286 +

287 +

288 +def test_triton_kernel_decorates_native_launchers(

289 + monkeypatch: pytest.MonkeyPatch,

290 +) -> None:

291 + kernel = _FakeTritonKernel()

292 +

293 + def warmup_inputs(kernel: Any = kernel) -> dict[str, Any]:

294 + return triton_warmup_inputs(

295 + kernel,

296 + "warmup",

297 + WarmupChoices(1, 2),

298 + grid=(2,),

299 + CONST=7,

300 + )

301 +

302 + launcher = triton_kernel_dispatcher_with_warmup(warmup_inputs=warmup_inputs)(kernel)

303 + launcher[(3,)]("runtime", 4, CONST=8)

304 +

305 + assert kernel.runtime_calls == [((3,), ("runtime", 4), {"CONST": 8})]

306 +

307 + _patch_key_deriver(

308 + monkeypatch,

309 + lambda kernel, kwargs: {

310 + TritonJitKey(id(kernel), "fake", 0, (kwargs["second"], kwargs["CONST"]))

311 + },

312 + )

313 +

314 + keys = launcher.get_warmup_keys()

315 + assert len(keys) == 2

316 + launcher.compile(keys[0])

317 + assert kernel.warmup_calls == [

318 + {"grid": (1,), "first": "warmup", "second": 1, "CONST": 7}

319 + ]

320 +

321 +

322 +def test_triton_warmup_inputs_expands_explicit_pointer_dtypes() -> None:

323 + inputs = triton_warmup_inputs(

324 + _pointer_group_test_kernel,

325 + grid=(1,),

326 + pointer_dtypes={torch.float32: ("x_ptr",)},

327 + value=2,

328 + BLOCK=16,

329 + )

330 +

331 + assert inputs["x_ptr"] == TritonWarmupTensor(torch.float32)

332 + assert inputs["value"] == 2

333 +

334 + with pytest.raises(ValueError, match="not a pointer: value"):

335 + triton_warmup_inputs(

336 + _pointer_group_test_kernel,

337 + grid=(1,),

338 + pointer_dtypes={torch.float32: ("value",)},

339 + x_ptr=TritonWarmupTensor(torch.float32),

340 + BLOCK=16,

341 + )

342 + with pytest.raises(ValueError, match="Missing Triton pointer inputs: x_ptr"):

343 + triton_warmup_inputs(

344 + _pointer_group_test_kernel,

345 + grid=(1,),

346 + pointer_dtypes={},

347 + value=2,

348 + BLOCK=16,

349 + )

350 +

351 +

352 +def test_triton_key_derivation_applies_wrappers_and_runtime_options(

353 + monkeypatch: pytest.MonkeyPatch,

354 +) -> None:

355 + knobs = triton.knobs

356 + driver = triton.runtime.driver

357 +

358 + device = "test-device"

359 + monkeypatch.setattr(driver.active, "get_current_device", lambda: device)

360 + calls: list[dict[str, Any]] = []

361 +

362 + def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any:

363 + calls.append(dict(BLOCK=BLOCK, **options))

364 + specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)]

365 + return {}, specialization, options

366 +

367 + cache: dict[Any, Any] = {}

368 + _binder_test_kernel.device_caches[device] = ({}, cache, None, None, binder)

369 + heuristic_kernel = triton.heuristics(

370 + {"BLOCK": lambda args: 16 if args["value"] <= 16 else 32}

371 + )(_binder_test_kernel)

372 +

373 + keys = jit_warmup_triton_helper._triton_key_deriver(heuristic_kernel)(

374 + {"x": TritonWarmupTensor(tl.float32), "value": 2}

375 + )

376 +

377 + assert len(keys) == 1

378 + expected_options = {

379 + "debug": knobs.runtime.debug,

380 + "instrumentation_mode": knobs.compilation.instrumentation_mode,

381 + }

382 + if hasattr(knobs.compilation, "fpsan_homomorphic_casts"):

383 + expected_options["fpsan_homomorphic_casts"] = (

384 + knobs.compilation.fpsan_homomorphic_casts

385 + )

386 + assert calls == [{"BLOCK": 16, **expected_options}]

387 +

388 +

389 +def test_triton_key_derivation_covers_autotune_configs_and_jit_identity(

390 + monkeypatch: pytest.MonkeyPatch,

391 +) -> None:

392 + driver = triton.runtime.driver

393 +

394 + device = "test-device"

395 + monkeypatch.setattr(driver.active, "get_current_device", lambda: device)

396 +

397 + def binder(x: Any, value: int, BLOCK: int, **options: Any) -> Any:

398 + specialization = [("pointer", x.dtype), ("i32", value), ("constexpr", BLOCK)]

399 + return {}, specialization, options

400 +

401 + for kernel in (_binder_test_kernel, _second_binder_test_kernel):

402 + kernel.device_caches[device] = ({}, {}, None, None, binder)

403 + autotuned = triton.autotune(

404 + configs=[

405 + triton.Config({"BLOCK": 16}, num_warps=2),

406 + triton.Config({"BLOCK": 32}, num_warps=4),

407 + ],

408 + key=["value"],

409 + )(_binder_test_kernel)

410 + inputs = {"x": TritonWarmupTensor(tl.float32), "value": 2}

411 +

412 + autotune_keys = jit_warmup_triton_helper._triton_key_deriver(autotuned)(inputs)

413 + other_kernel_keys = jit_warmup_triton_helper._triton_key_deriver(

414 + _second_binder_test_kernel

415 + )(inputs | {"BLOCK": 16})

416 +

417 + assert len(autotune_keys) == 2

418 + assert {key.jit_function_key for key in autotune_keys} == {

419 + _binder_test_kernel.cache_key

420 + }

421 + assert {key.jit_function_id for key in autotune_keys} == {id(_binder_test_kernel)}

422 + assert autotune_keys.isdisjoint(other_kernel_keys)

423 +

424 +

425 def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:

426 from vllm.v1.attention.backends.utils import PAD_SLOT_ID

427 from vllm.v1.worker.block_table import ComputeSlotMappingKernel

tests/model_executor/test_spec_decode_rejection_warmup.py deleted

-121

@@ -1,121 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from types import SimpleNamespace

4 -

5 -import pytest

6 -import torch

7 -

8 -from vllm.config.watermarking import WatermarkConfig, derive_watermark_key

9 -from vllm.model_executor.warmup.spec_decode_rejection_warmup import (

10 - spec_decode_rejection_warmup,

11 -)

12 -

13 -MASTER_KEY = 20250910

14 -

15 -

16 -def _worker(

17 - watermark_config: WatermarkConfig | None,

18 - draft_sample_method: str = "greedy",

19 -) -> SimpleNamespace:

20 - return SimpleNamespace(

21 - device=torch.device("cpu"),

22 - vllm_config=SimpleNamespace(

23 - speculative_config=SimpleNamespace(

24 - num_speculative_tokens=2,

25 - rejection_sample_method="standard",

26 - draft_sample_method=draft_sample_method,

27 - ),

28 - model_config=SimpleNamespace(

29 - get_vocab_size=lambda: 128, dtype=torch.bfloat16

30 - ),

31 - watermark_config=watermark_config,

32 - ),

33 - )

34 -

35 -

36 -def _capture_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]:

37 - calls: list[dict] = []

38 -

39 - def fake_rejection_sample(**kwargs):

40 - calls.append(kwargs)

41 - return None, None

42 -

43 - monkeypatch.setattr(

44 - "vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils.rejection_sample",

45 - fake_rejection_sample,

46 - )

47 - return calls

48 -

49 -

50 -def test_warmup_token_id_dtypes_match_the_runtime_buffers(

51 - monkeypatch: pytest.MonkeyPatch,

52 -) -> None:

53 - calls = _capture_calls(monkeypatch)

54 - spec_decode_rejection_warmup(_worker(None))

55 -

56 - assert calls

57 - for kwargs in calls:

58 - assert kwargs["draft_sampled"].dtype == torch.int32

59 - assert kwargs["pos"].dtype == torch.int64

60 -

61 -

62 -def test_warmup_index_mapping_dtypes_match_the_runtime_buffers(

63 - monkeypatch: pytest.MonkeyPatch,

64 -) -> None:

65 - calls = _capture_calls(monkeypatch)

66 - spec_decode_rejection_warmup(_worker(None))

67 -

68 - assert calls

69 - for kwargs in calls:

70 - assert kwargs["idx_mapping"].dtype == torch.int64

71 - assert kwargs["expanded_idx_mapping"].dtype == torch.int64

72 - assert kwargs["expanded_local_pos"].dtype == torch.int32

73 -

74 -

75 -def test_warmup_watermark_args_match_the_runtime_kernel_arguments(

76 - monkeypatch: pytest.MonkeyPatch,

77 -) -> None:

78 - watermark_config = WatermarkConfig(

79 - key=MASTER_KEY, algorithm="dual_key_gumbel", context_width=3

80 - )

81 - calls = _capture_calls(monkeypatch)

82 - spec_decode_rejection_warmup(

83 - _worker(watermark_config, draft_sample_method="probabilistic")

84 - )

85 -

86 - assert calls

87 - key_b = derive_watermark_key(MASTER_KEY, b"key_b")

88 - for kwargs in calls:

89 - contexts = kwargs["contexts"]

90 - assert contexts.dtype == torch.int32

91 - assert contexts.shape == (3, 3) and contexts.stride() == (3, 1)

92 - assert kwargs["watermarking"].dtype == torch.bool

93 - assert kwargs["watermark_key"] == key_b

94 - assert key_b & 0xFFFFFFFF >= 2**31 and key_b >> 32 >= 2**31

95 -

96 -

97 -def test_warmup_is_skipped_when_the_recovery_key_cannot_be_derived(

98 - monkeypatch: pytest.MonkeyPatch,

99 -) -> None:

100 - calls = _capture_calls(monkeypatch)

101 - monkeypatch.setattr(

102 - "vllm.v1.watermarking.spec_decode.speculative_target_watermark_key",

103 - lambda _: (_ for _ in ()).throw(NotImplementedError("no philox")),

104 - )

105 - spec_decode_rejection_warmup(_worker(WatermarkConfig(key=MASTER_KEY)))

106 -

107 - assert calls == []

108 -

109 -

110 -def test_warmup_draft_logits_match_the_runtime_draft_sample_method(

111 - monkeypatch: pytest.MonkeyPatch,

112 -) -> None:

113 - calls = _capture_calls(monkeypatch)

114 - spec_decode_rejection_warmup(_worker(None))

115 - assert calls

116 - assert all(kwargs["draft_logits"] is None for kwargs in calls)

117 -

118 - calls = _capture_calls(monkeypatch)

119 - spec_decode_rejection_warmup(_worker(None, draft_sample_method="probabilistic"))

120 - assert calls

121 - assert all(kwargs["draft_logits"] is not None for kwargs in calls)

tests/v1/sample/test_topk_topp_sampler.py

@@ -79,6 +79,10 @@ def test_rocm_aiter_sampler_defers_import_when_generators_force_native(

79 def is_cuda():

80 return False

82 + @staticmethod

83 + def is_cuda_alike():

84 + return True

85 +

86 @staticmethod

87 def is_cpu():

88 return False

tests/v1/shutdown/test_delete.py

+13 -3

@@ -11,6 +11,7 @@ from tests.v1.shutdown.utils import (

11 SHUTDOWN_TEST_TIMEOUT_SEC,

12 )

13 from vllm import LLM, SamplingParams

14 +from vllm.config import KernelConfig

15 from vllm.engine.arg_utils import AsyncEngineArgs

16 from vllm.platforms import current_platform

17 from vllm.sampling_params import RequestOutputKind

@@ -39,7 +40,10 @@ async def test_async_llm_delete(

40 pytest.skip(reason="Not enough CUDA devices")

42 engine_args = AsyncEngineArgs(

42 - model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size

43 + model=model,

44 + enforce_eager=True,

45 + tensor_parallel_size=tensor_parallel_size,

46 + kernel_config=KernelConfig(enable_jit_warmup=False),

47 )

49 # Instantiate AsyncLLM; make request to complete any deferred

@@ -94,7 +98,10 @@ def test_llm_delete(

98 # Instantiate LLM; make request to complete any deferred

99 # initialization; then delete instance

100 llm = LLM(

97 - model=model, enforce_eager=True, tensor_parallel_size=tensor_parallel_size

101 + model=model,

102 + enforce_eager=True,

103 + tensor_parallel_size=tensor_parallel_size,

104 + kernel_config=KernelConfig(enable_jit_warmup=False),

105 )

106 if send_one_request:

107 llm.generate(

@@ -122,7 +129,10 @@ def test_llm_delete_inprocess(

129 with monkeypatch.context() as m:

130 m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

131

125 - with VllmRunner(model) as vllm_model:

132 + with VllmRunner(

133 + model,

134 + kernel_config=KernelConfig(enable_jit_warmup=False),

135 + ) as vllm_model:

136 if send_one_request:

137 vllm_model.generate(

138 ["Hello my name is"],

tests/v1/spec_decode/test_dflash_prepare_inputs.py

+1 -1

@@ -43,7 +43,7 @@ def _run_prepare(

43 num_scheduled_tokens=np.array([4], dtype=np.int32),

44 positions=torch.tensor(target_positions, dtype=torch.int64, device=device),

45 query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),

46 - idx_mapping=torch.tensor([2], dtype=torch.int32, device=device),

46 + idx_mapping=torch.tensor([2], dtype=torch.int64, device=device),

47 )

48 query_slot_mapping = torch.full(

49 (max_num_tokens,), -2, dtype=torch.int64, device=device

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -157,6 +157,7 @@ def test_mm_support_configured_after_model_load(monkeypatch):

157 speculator.device = device

158 speculator.max_num_tokens = 4

159 speculator.max_num_reqs = 2

160 + speculator.num_speculative_steps = 1

161 speculator.hidden_size = 3

162 speculator.dtype = torch.float32

163 speculator.draft_model_config = draft_model_config

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

+41 -8

@@ -1,8 +1,16 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +from typing import Any

5 +

6 import torch

8 +from vllm.model_executor.warmup.jit_warmup import WarmupIntRange

9 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

10 + DispatchSpec,

11 + TritonWarmupTensor,

12 + triton_kernel_dispatcher_with_warmup,

13 +)

14 from vllm.platforms import current_platform

15 from vllm.triton_utils import tl, triton

@@ -33,11 +41,39 @@ def _scatter_states_kernel(

41 tl.store(state_ptr + state_idx * stride_state_batch + offsets, values, mask=mask)

44 +def _scatter_states_warmup_inputs(

45 + *,

46 + state_shape: tuple[int, ...],

47 + state_dtype: torch.dtype,

48 + indices_dtype: torch.dtype,

49 + max_num_tokens: int,

50 +) -> dict[str, object]:

51 + num_tokens: Any = WarmupIntRange(1, max_num_tokens + 1)

52 + return dict(

53 + state=TritonWarmupTensor(

54 + state_dtype,

55 + shape=(1,) + state_shape,

56 + ),

57 + src=TritonWarmupTensor(

58 + state_dtype,

59 + shape=(num_tokens,) + state_shape,

60 + ),

61 + indices=TritonWarmupTensor(

62 + indices_dtype,

63 + shape=(num_tokens,),

64 + ),

65 + )

66 +

67 +

68 +@triton_kernel_dispatcher_with_warmup(

69 + kernel=_scatter_states_kernel,

70 + warmup_inputs=_scatter_states_warmup_inputs,

71 +)

72 def scatter_states(

73 state: torch.Tensor,

74 src: torch.Tensor,

75 indices: torch.Tensor,

40 -) -> None:

76 +) -> DispatchSpec:

77 """Scatter ``src`` rows into ``state`` at ``indices`` (in place).

79 Equivalent to ``state[indices] = src`` but non-atomic and bandwidth-bound,

@@ -58,13 +94,10 @@ def scatter_states(

94 assert src[0].is_contiguous()

95 block_size = min(triton.next_power_of_2(row_size), 1024)

96 grid = (triton.cdiv(row_size, block_size), indices.numel())

61 - _scatter_states_kernel[grid](

62 - state,

63 - src,

64 - indices,

65 - state.stride(0),

66 - src.stride(0),

67 - indices.stride(0),

97 + return grid, dict(

98 + stride_state_batch=state.stride(0),

99 + stride_src_batch=src.stride(0),

100 + stride_indices=indices.stride(0),

101 row_size=row_size,

102 BLOCK_SIZE=block_size,

103 num_warps=8,

vllm/model_executor/warmup/jit_warmup.py

+309 -30

@@ -21,6 +21,7 @@ from typing import Any, Generic, TypeVar, cast

21 __all__ = [

22 "JitWarmupRegistry",

23 "VllmJitKernel",

24 + "WarmupChoices",

25 "WarmupIntRange",

26 "get_ast_full_name",

27 "get_function_source_node",

@@ -62,6 +63,21 @@ class WarmupIntRange:

63 advance: Callable[[int], int] | None = None

66 +@dataclass(frozen=True)

67 +class WarmupChoices:

68 + """Expand an explicit finite set of values in a traced warmup method."""

69 +

70 + values: tuple[Any, ...]

71 +

72 + def __init__(self, *values: Any) -> None:

73 + object.__setattr__(self, "values", values)

74 +

75 +

76 +def _when(condition: bool) -> None:

77 + """Filter symbolic cases in an AST-traced warmup-cases method."""

78 + raise RuntimeError("_when() is only valid in an AST-traced warmup method")

79 +

80 +

81 WarmupValues = Any

82 CompileKeyDispatchFn = Callable[..., CompileKeyT]

83 WarmupPredicateFn = Callable[..., bool]

@@ -87,6 +103,8 @@ class _WarmupInputRows:

103

104

105 def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:

106 + if isinstance(values, WarmupChoices):

107 + return values.values

108 if isinstance(values, WarmupIntRange):

109 if values.advance is None:

110 return tuple(range(values.start, values.stop, values.step))

@@ -278,6 +296,27 @@ class _DispatchExprEvaluator(ast.NodeVisitor):

296 def visit_Constant(self, node: ast.Constant) -> Any:

297 return node.value

298

299 + def visit_Lambda(self, node: ast.Lambda) -> Callable[..., Any]:

300 + arguments = node.args

301 + if (

302 + arguments.posonlyargs

303 + or arguments.vararg is not None

304 + or arguments.kwonlyargs

305 + or arguments.kwarg is not None

306 + or arguments.defaults

307 + ):

308 + raise _dispatch_expr_error(node, "Traced lambdas require positional args")

309 + names = tuple(argument.arg for argument in arguments.args)

310 +

311 + def evaluate(*values: Any) -> Any:

312 + if len(values) != len(names):

313 + raise TypeError(f"Expected {len(names)} lambda arguments")

314 + return type(self)(

315 + dict(self.values) | dict(zip(names, values)), self.globals

316 + ).eval(node.body)

317 +

318 + return evaluate

319 +

320 def visit_IfExp(self, node: ast.IfExp) -> Any:

321 return self.visit(node.body if self.visit(node.test) else node.orelse)

322

@@ -333,7 +372,12 @@ class _DispatchExprEvaluator(ast.NodeVisitor):

372 return op(self.visit(node.left), self.visit(node.right))

373

374 def visit_Call(self, node: ast.Call) -> Any:

336 - args = [self.visit(arg) for arg in node.args]

375 + args: list[Any] = []

376 + for arg in node.args:

377 + if isinstance(arg, ast.Starred):

378 + args.extend(self.visit(arg.value))

379 + else:

380 + args.append(self.visit(arg))

381 fn = self.visit(node.func)

382 call_kwargs: dict[str, Any] = {}

383 for keyword in node.keywords:

@@ -388,6 +432,45 @@ def _eval_dispatch_expr(

432 return _DispatchExprEvaluator(kwargs, globals_).eval(node)

433

434

435 +def _validate_dispatch_expr(node: ast.AST) -> None:

436 + """Enforce the dispatch AST subset before compiling traced code."""

437 + allowed_nodes = (

438 + ast.Name,

439 + ast.Constant,

440 + ast.Lambda,

441 + ast.arguments,

442 + ast.arg,

443 + ast.IfExp,

444 + ast.Tuple,

445 + ast.List,

446 + ast.BoolOp,

447 + ast.And,

448 + ast.Or,

449 + ast.Compare,

450 + *tuple(_CMP_OPS),

451 + ast.UnaryOp,

452 + ast.Not,

453 + ast.USub,

454 + ast.BinOp,

455 + *tuple(_BIN_OPS),

456 + ast.Call,

457 + ast.keyword,

458 + ast.Starred,

459 + ast.Attribute,

460 + ast.Subscript,

461 + ast.Load,

462 + )

463 + for child in ast.walk(node):

464 + if not isinstance(child, allowed_nodes):

465 + raise _dispatch_expr_error(child, "Unsupported dispatch expression")

466 + if isinstance(child, ast.Call) and any(

467 + keyword.arg is None for keyword in child.keywords

468 + ):

469 + raise _dispatch_expr_error(

470 + child, "Dispatch helper calls cannot use **kwargs"

471 + )

472 +

473 +

474 def _collect_input_names(

475 node: ast.AST,

476 candidate_names: set[str],

@@ -406,6 +489,23 @@ def _collect_input_names(

489 }

490

491

492 +def _named_assignment(

493 + statement: ast.stmt,

494 + subject: str,

495 +) -> tuple[str, ast.AST] | None:

496 + match statement:

497 + case ast.Assign(targets=[ast.Name(id=name)], value=value):

498 + return name, value

499 + case ast.AnnAssign(target=ast.Name(id=name), value=value) if value is not None:

500 + return name, value

501 + case ast.Assign() | ast.AnnAssign():

502 + raise _dispatch_expr_error(

503 + statement, f"{subject} assignments require one name"

504 + )

505 + case _:

506 + return None

507 +

508 +

509 def _collect_expression_body(

510 fn: Callable[..., Any],

511 function_def: ast.FunctionDef | ast.Lambda,

@@ -422,38 +522,19 @@ def _collect_expression_body(

522 ):

523 continue

524

425 - if isinstance(statement, ast.Assign):

426 - if len(statement.targets) != 1 or not isinstance(

427 - statement.targets[0], ast.Name

428 - ):

429 - raise _dispatch_expr_error(

430 - statement,

431 - "Traced warmup helper assignments must target one local name",

432 - )

433 - local_exprs.append((statement.targets[0].id, statement.value))

525 + if (

526 + assignment := _named_assignment(statement, "Traced warmup helper")

527 + ) is not None:

528 + local_exprs.append(assignment)

529 continue

530

436 - if isinstance(statement, ast.AnnAssign):

531 + if isinstance(statement, ast.Return):

532 if statement.value is None:

533 raise _dispatch_expr_error(

439 - statement, "Traced warmup helper annotations must assign a value"

534 + statement, "Traced warmup helper must return an expression"

535 )

441 - if not isinstance(statement.target, ast.Name):

442 - raise _dispatch_expr_error(

443 - statement,

444 - "Traced warmup helper assignments must target one local name",

445 - )

446 - local_exprs.append((statement.target.id, statement.value))

447 - continue

448 -

449 - if isinstance(statement, ast.Return) and statement.value is not None:

536 return local_exprs, statement.value

537

452 - if isinstance(statement, ast.Return):

453 - raise _dispatch_expr_error(

454 - statement, "Traced warmup helper must return an expression"

455 - )

456 -

538 raise _dispatch_expr_error(

539 statement,

540 "Traced warmup helper may only contain local assignments before return",

@@ -574,10 +655,17 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):

655 bind_launch_inputs = True

656

657 def __init__(self) -> None:

577 - self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch)

658 + dispatch = type(self).dispatch

659 + self._dispatch_trace = (

660 + None

661 + if dispatch is VllmJitKernel.dispatch

662 + else _trace_compile_key_dispatch(self.dispatch)

663 + )

664 self._compiled_cache: dict[Any, Any] = {}

665

666 def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:

667 + if self._dispatch_trace is None:

668 + raise TypeError(f"{type(self).__name__} does not define dispatch()")

669 return self._dispatch_trace.compile_key(self.CompileKey, kwargs)

670

671 def _get_or_compile(

@@ -683,7 +771,184 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):

771

772 return traced

773

686 - @abstractmethod

774 + def _expand_warmup_cases(

775 + self,

776 + cases_fn: Callable[..., Any],

777 + *args: Any,

778 + _value_expander: Callable[[WarmupValues], tuple[Any, ...]] = (

779 + _expand_warmup_values

780 + ),

781 + **kwargs: Any,

782 + ) -> Iterator[Mapping[str, Any]]:

783 + """Expand symbolic domains declared inside a warmup-cases method."""

784 + function_def = get_function_source_node(cases_fn)

785 + if isinstance(function_def, ast.Lambda):

786 + raise _dispatch_expr_error(

787 + function_def, "Warmup cases must be a function definition"

788 + )

789 + signature = inspect.signature(cases_fn)

790 + bound = signature.bind(*args, **kwargs)

791 + bound.apply_defaults()

792 + static_values = dict(bound.arguments)

793 + bound_self = getattr(cases_fn, "__self__", None)

794 + if bound_self is not None:

795 + static_values["self"] = bound_self

796 + globals_ = getattr(cases_fn, "__func__", cases_fn).__globals__

797 +

798 + domains: list[tuple[str, tuple[Any, ...]]] = []

799 + local_exprs: list[tuple[str, ast.AST]] = []

800 + predicates: list[ast.AST] = []

801 + return_expr: ast.AST | None = None

802 + for statement in function_def.body:

803 + if (

804 + isinstance(statement, ast.Expr)

805 + and isinstance(statement.value, ast.Constant)

806 + and isinstance(statement.value.value, str)

807 + ):

808 + continue

809 + assignment = _named_assignment(statement, "Warmup case")

810 + if assignment is not None:

811 + name, value_expr = assignment

812 + call_name = (

813 + (get_ast_full_name(value_expr.func) or "").split(".")[-1]

814 + if isinstance(value_expr, ast.Call)

815 + else ""

816 + )

817 + if call_name in {"WarmupIntRange", "WarmupChoices"}:

818 + value = _eval_dispatch_expr(value_expr, static_values, globals_)

819 + domains.append((name, _value_expander(value)))

820 + else:

821 + local_exprs.append((name, value_expr))

822 + continue

823 + if (

824 + isinstance(statement, ast.Expr)

825 + and isinstance(statement.value, ast.Call)

826 + and (get_ast_full_name(statement.value.func) or "").endswith("_when")

827 + ):

828 + if len(statement.value.args) != 1 or statement.value.keywords:

829 + raise _dispatch_expr_error(

830 + statement, "_when requires exactly one positional expression"

831 + )

832 + predicates.append(statement.value.args[0])

833 + continue

834 + if isinstance(statement, ast.Return):

835 + if return_expr is not None or statement.value is None:

836 + raise _dispatch_expr_error(

837 + statement, "Warmup cases require one final return expression"

838 + )

839 + return_expr = statement.value

840 + continue

841 + raise _dispatch_expr_error(

842 + statement,

843 + "AST-traced warmup cases support assignments, _when, and return",

844 + )

845 +

846 + if return_expr is None or not isinstance(return_expr, ast.Call):

847 + raise ValueError("AST-traced warmup cases must return one case(...) call")

848 + if any(keyword.arg is None for keyword in return_expr.keywords):

849 + raise _dispatch_expr_error(

850 + return_expr, "Warmup case expressions do not support **kwargs"

851 + )

852 +

853 + class _InlineDomainRewriter(ast.NodeTransformer):

854 + def visit_Call(self, node: ast.Call) -> ast.AST:

855 + call_name = (get_ast_full_name(node.func) or "").split(".")[-1]

856 + if call_name not in {"WarmupIntRange", "WarmupChoices"}:

857 + return self.generic_visit(node)

858 + name = f"__warmup_domain_{len(domains)}"

859 + domain = _eval_dispatch_expr(node, static_values, globals_)

860 + domains.append((name, _value_expander(domain)))

861 + return ast.copy_location(ast.Name(id=name, ctx=ast.Load()), node)

862 +

863 + return_expr = cast(ast.Call, _InlineDomainRewriter().visit(return_expr))

864 +

865 + domain_names = tuple(name for name, _ in domains)

866 + dynamic_names = set(domain_names)

867 + dynamic_local_exprs: list[tuple[str, ast.AST]] = []

868 + for name, expr in local_exprs:

869 + referenced_names = {

870 + node.id for node in ast.walk(expr) if isinstance(node, ast.Name)

871 + }

872 + if referenced_names & dynamic_names:

873 + dynamic_local_exprs.append((name, expr))

874 + dynamic_names.add(name)

875 + else:

876 + static_values[name] = _eval_dispatch_expr(expr, static_values, globals_)

877 +

878 + for _, expr in dynamic_local_exprs:

879 + _validate_dispatch_expr(expr)

880 + for predicate in predicates:

881 + _validate_dispatch_expr(predicate)

882 + _validate_dispatch_expr(return_expr)

883 +

884 + case_body: list[ast.stmt] = [

885 + ast.Assign(

886 + targets=[ast.Name(id=name, ctx=ast.Store())],

887 + value=cast(ast.expr, expr),

888 + )

889 + for name, expr in dynamic_local_exprs

890 + ]

891 + if predicates:

892 + predicate = (

893 + predicates[0]

894 + if len(predicates) == 1

895 + else ast.BoolOp(

896 + op=ast.And(),

897 + values=[cast(ast.expr, value) for value in predicates],

898 + )

899 + )

900 + case_body.append(

901 + ast.If(

902 + test=ast.UnaryOp(op=ast.Not(), operand=cast(ast.expr, predicate)),

903 + body=[ast.Return(value=ast.Constant(value=None))],

904 + orelse=[],

905 + )

906 + )

907 + case_body.append(ast.Return(value=cast(ast.expr, return_expr)))

908 + # FunctionDef fields vary across the supported Python versions, so no

909 + # single constructor overload matches every mypy target.

910 + case_function = ast.FunctionDef( # type: ignore[call-overload]

911 + name="__vllm_warmup_case",

912 + args=ast.arguments(

913 + posonlyargs=[],

914 + args=[ast.arg(arg=name) for name in domain_names],

915 + kwonlyargs=[],

916 + kw_defaults=[],

917 + defaults=[],

918 + ),

919 + body=case_body,

920 + decorator_list=[],

921 + returns=None,

922 + type_comment=None,

923 + )

924 + case_globals = dict(globals_)

925 + case_globals.update(static_values)

926 + exec(

927 + compile(

928 + ast.fix_missing_locations(

929 + ast.Module(body=[case_function], type_ignores=[])

930 + ),

931 + "<jit-warmup>",

932 + "exec",

933 + ),

934 + case_globals,

935 + )

936 + evaluate_case = cast(Callable[..., Any], case_globals[case_function.name])

937 +

938 + domain_values = tuple(values for _, values in domains)

939 + for values in itertools.product(*domain_values):

940 + case = evaluate_case(*values)

941 + if case is None:

942 + continue

943 + if not isinstance(case, Mapping):

944 + raise TypeError("AST-traced warmup cases must return a mapping")

945 + if any(

946 + isinstance(value, WarmupChoices | WarmupIntRange)

947 + for value in case.values()

948 + ):

949 + raise TypeError("Warmup domains must be direct expressions")

950 + yield case

951 +

952 def dispatch(self, **kwargs: Any) -> CompileKeyT:

953 """Build one compile key from one concrete dispatch point."""

954 raise NotImplementedError

@@ -751,6 +1016,19 @@ class JitWarmupRegistry:

1016 list[tuple[tuple[Any, ...], dict[str, Any]]],

1017 ] = {}

1018

1019 + @classmethod

1020 + def capture(cls, init_fn: Callable[..., None]) -> Callable[..., None]:

1021 + """Collect warmup registrations made while the decorated callable runs."""

1022 +

1023 + @wraps(init_fn)

1024 + def wrapped(instance: Any, vllm_config: Any, *args: Any, **kwargs: Any) -> None:

1025 + registry = cls(vllm_config)

1026 + instance.jit_warmup_registry = registry

1027 + with registry.activate():

1028 + init_fn(instance, vllm_config, *args, **kwargs)

1029 +

1030 + return wrapped

1031 +

1032 @contextmanager

1033 def activate(self) -> Iterator[None]:

1034 """Collect registrations made in this context."""

@@ -808,9 +1086,10 @@ class JitWarmupRegistry:

1086 if (

1087 not args

1088 and not kwargs

811 - and inspect.signature(kernel.get_warmup_keys).parameters

1089 + and "vllm_config"

1090 + in inspect.signature(kernel.get_warmup_keys).parameters

1091 ):

813 - args = (self.vllm_config,)

1092 + kwargs = {"vllm_config": self.vllm_config}

1093 for compile_key in kernel.get_warmup_keys(*args, **kwargs):

1094 compile_keys[compile_key] = None

1095 if compile_keys:

vllm/model_executor/warmup/jit_warmup_triton_helper.py

+507 -13

@@ -3,25 +3,130 @@

3 import ast

4 import inspect

5 from abc import abstractmethod

6 -from collections.abc import Callable, Mapping

7 -from dataclasses import dataclass

8 -from functools import cached_property

9 -from typing import Any, ClassVar, Generic, TypeVar, cast

6 +from collections.abc import Callable, Hashable, Iterable, Mapping

7 +from contextlib import nullcontext, suppress

8 +from dataclasses import dataclass, field

9 +from functools import cache, cached_property, update_wrapper

10 +from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload

12 from vllm.model_executor.warmup.jit_warmup import (

13 VllmJitKernel,

14 get_ast_full_name,

15 get_function_source_node,

16 )

17 +from vllm.platforms import current_platform

18 +from vllm.triton_utils import triton

20 CompileKeyT = TypeVar("CompileKeyT")

18 -# ``(grid, launch_kwargs)`` or ``(grid, launch_kwargs, outputs)`` for

19 -# self-allocating kernels. ``grid=None`` skips the launch (e.g. empty batch)

20 -# but still returns the declared outputs.

21 +P = ParamSpec("P")

22 +# ``kernel_launcher`` also supports a ``None`` grid when no launch is needed.

23 LaunchSpec = (

24 tuple[tuple[int, ...] | None, dict[str, Any]]

25 | tuple[tuple[int, ...] | None, dict[str, Any], Any]

26 )

27 +# Dispatchers select a concrete launch configuration for a Triton kernel.

28 +DispatchSpec = (

29 + tuple[tuple[int, ...], dict[str, Any]] | tuple[tuple[int, ...], dict[str, Any], Any]

30 +)

31 +

32 +

33 +@dataclass(frozen=True)

34 +class _LaunchBindingPlan:

35 + input_targets: tuple[tuple[int, str], ...]

36 + stride_targets: tuple[tuple[str, str, int], ...]

37 +

38 +

39 +@cache

40 +def _launch_binding_plan(

41 + arg_names: tuple[str, ...],

42 + input_names: tuple[str, ...],

43 + launch_names: tuple[str, ...],

44 +) -> _LaunchBindingPlan:

45 + arg_name_set = set(arg_names)

46 + launch_name_set = set(launch_names)

47 +

48 + def kernel_arg(input_name: str) -> str | None:

49 + pointer_name = f"{input_name}_ptr"

50 + candidates = (

51 + input_name,

52 + pointer_name,

53 + input_name.upper(),

54 + pointer_name.upper(),

55 + )

56 + return next(

57 + (candidate for candidate in candidates if candidate in arg_name_set),

58 + None,

59 + )

60 +

61 + input_targets = tuple(

62 + (index, target)

63 + for index, name in enumerate(input_names)

64 + if (target := kernel_arg(name)) is not None and target not in launch_name_set

65 + )

66 + bound_names = launch_name_set | {target for _, target in input_targets}

67 + stride_targets = []

68 + for target in arg_names:

69 + if target in bound_names:

70 + continue

71 + name, separator, dim = target.rpartition("_stride")

72 + if not separator:

73 + name, separator, dim = target.rpartition("_STRIDE")

74 + dim = dim.removeprefix("_")

75 + if separator and (not dim or dim.isdigit()):

76 + pointer = f"{name}_PTR" if separator == "_STRIDE" else f"{name}_ptr"

77 + source = name if name in bound_names else pointer

78 + if source in bound_names:

79 + stride_targets.append((target, source, int(dim or 0)))

80 + return _LaunchBindingPlan(input_targets, tuple(stride_targets))

81 +

82 +

83 +def triton_warmup_inputs(

84 + kernel: Any,

85 + *args: Any,

86 + grid: tuple[int, ...],

87 + pointer_dtypes: (

88 + Mapping[Any, Iterable[str]] | Iterable[tuple[Any, Iterable[str]]] | None

89 + ) = None,

90 + **kwargs: Any,

91 +) -> dict[str, Any]:

92 + """Build launcher inputs from Triton's native positional argument order."""

93 + arg_names = tuple(kernel.arg_names)

94 + if len(args) > len(arg_names):

95 + raise ValueError(

96 + f"Received {len(args)} positional inputs for {len(arg_names)} "

97 + "Triton kernel arguments"

98 + )

99 + inputs = dict(zip(arg_names, args))

100 + duplicate_names = inputs.keys() & kwargs.keys()

101 + if duplicate_names:

102 + raise ValueError(

103 + f"Triton inputs passed twice: {', '.join(sorted(duplicate_names))}"

104 + )

105 + inputs.update(kwargs)

106 + if pointer_dtypes is not None:

107 + function_def = get_function_source_node(kernel)

108 + if not isinstance(function_def, ast.FunctionDef):

109 + raise ValueError("Expected Triton kernel to be defined as a function")

110 + pointer_names = _pointer_arg_names(function_def, arg_names)

111 + groups = (

112 + pointer_dtypes.items()

113 + if isinstance(pointer_dtypes, Mapping)

114 + else pointer_dtypes

115 + )

116 + for dtype, names in groups:

117 + for name in names:

118 + if name not in arg_names:

119 + raise ValueError(f"Unknown Triton pointer argument: {name}")

120 + if name not in pointer_names:

121 + raise ValueError(f"Triton argument is not a pointer: {name}")

122 + if name in inputs:

123 + raise ValueError(f"Triton input passed twice: {name}")

124 + inputs[name] = TritonWarmupTensor(dtype)

125 + missing_pointers = pointer_names - inputs.keys()

126 + if missing_pointers:

127 + names = ", ".join(sorted(missing_pointers))

128 + raise ValueError(f"Missing Triton pointer inputs: {names}")

129 + return {"grid": grid, **inputs}

130

131

132 def triton_scalar_specialization_rep(value: int) -> int:

@@ -72,6 +177,7 @@ class TritonWarmupTensor:

177 aligned: bool = True

178 shape: tuple[int, ...] = (1,)

179 strides: tuple[int, ...] | None = None

180 + init: Any = 0

181

182 def data_ptr(self) -> int:

183 return 0 if self.aligned else 1

@@ -79,6 +185,12 @@ class TritonWarmupTensor:

185 def ptr_range(self) -> int:

186 return 0

187

188 + @property

189 + def device(self) -> Any:

190 + import torch

191 +

192 + return torch.device(current_platform.device_type)

193 +

194 def stride(self, dim: int | None = None) -> int | tuple[int, ...]:

195 if self.strides is None:

196 strides: list[int] = []

@@ -95,9 +207,10 @@ class TritonWarmupTensor:

207 class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):

208 """Triton owner whose runtime launch specification is reused for warmup."""

209

98 - kernel: ClassVar[Any]

210 + kernel: Any

211 _warming = False

212 _warming_compile_key: CompileKeyT | None = None

213 + _run_autotune = False

214

215 @abstractmethod

216 def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:

@@ -126,6 +239,21 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):

239 f"Cannot inspect kernel parameters for {type(self.kernel).__name__}"

240 )

241

242 + def _prepare_launch_kwargs(

243 + self,

244 + input_names: tuple[str, ...],

245 + input_values: tuple[Any, ...],

246 + launch_kwargs: dict[str, Any],

247 + ) -> dict[str, Any]:

248 + plan = _launch_binding_plan(

249 + self._kernel_arg_names, input_names, tuple(launch_kwargs)

250 + )

251 + for index, target in plan.input_targets:

252 + launch_kwargs[target] = input_values[index]

253 + for target, source, dim in plan.stride_targets:

254 + launch_kwargs[target] = launch_kwargs[source].stride(dim)

255 + return launch_kwargs

256 +

257 def launch(

258 self,

259 launch_spec: LaunchSpec,

@@ -136,15 +264,20 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):

264 else:

265 grid, launch_kwargs = launch_spec

266 outputs = None

139 - kwargs = dict(launch_kwargs)

267 + kwargs = self._prepare_launch_kwargs(

268 + tuple(inputs), tuple(inputs.values()), dict(launch_kwargs)

269 + )

270 kernel = kwargs.pop("kernel", None) or self.kernel

271 runtime_launcher = kwargs.pop("_runtime_launcher", None)

272 runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0)

143 - for name, value in inputs.items():

144 - target = name if name in self._kernel_arg_names else f"{name}_ptr"

145 - if target in self._kernel_arg_names and target not in kwargs:

146 - kwargs[target] = value

273 if self._warming:

274 + if self._run_autotune:

275 + assert grid is not None

276 + kernel[grid](**kwargs)

277 + return outputs

278 + kwargs = {

279 + name: _triton_metadata_arg(value) for name, value in kwargs.items()

280 + }

281 if (

282 self._warming_compile_key is not None

283 and "launch_pdl" in kwargs

@@ -168,6 +301,367 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):

301 return outputs

302

303

304 +@dataclass(frozen=True)

305 +class TritonJitKey:

306 + """Process-local identity of one Triton JIT specialization."""

307 +

308 + jit_function_id: int

309 + jit_function_key: Hashable

310 + device: Hashable

311 + cache_key: Hashable

312 +

313 +

314 +@dataclass(frozen=True)

315 +class TritonCompileKey:

316 + """Triton-derived compile key with one non-comparing replay input."""

317 +

318 + jit_keys: frozenset[TritonJitKey]

319 + inputs: tuple[tuple[str, Any], ...] = field(compare=False, hash=False, repr=False)

320 +

321 +

322 +def _triton_key_deriver(

323 + kernel: Any,

324 +) -> Callable[[Mapping[str, Any]], set[TritonJitKey]]:

325 + """Prepare Triton's key derivation once for one kernel and device."""

326 + knobs = triton.knobs

327 + Autotuner = triton.runtime.autotuner.Autotuner

328 + Heuristics = triton.runtime.autotuner.Heuristics

329 + driver = triton.runtime.driver

330 + JITFunction = triton.runtime.jit.JITFunction

331 + compute_cache_key = triton.runtime.jit.compute_cache_key

332 +

333 + if isinstance(kernel, Heuristics):

334 + derive_inner = _triton_key_deriver(kernel.fn)

335 +

336 + def derive_heuristic(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:

337 + heuristic_kwargs = dict(kwargs)

338 + for name, heuristic in kernel.values.items():

339 + heuristic_kwargs[name] = heuristic(heuristic_kwargs)

340 + return derive_inner(heuristic_kwargs)

341 +

342 + return derive_heuristic

343 +

344 + if isinstance(kernel, Autotuner):

345 + derive_inner = _triton_key_deriver(kernel.fn)

346 +

347 + def derive_autotuned(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:

348 + previous_nargs = getattr(kernel, "nargs", None)

349 + kernel.nargs = {}

350 + try:

351 + configs = kernel.prune_configs(dict(kwargs))

352 + finally:

353 + kernel.nargs = previous_nargs

354 + keys: set[TritonJitKey] = set()

355 + for config in configs:

356 + conflicts = kwargs.keys() & config.kwargs.keys()

357 + if conflicts:

358 + names = ", ".join(sorted(conflicts))

359 + raise ValueError(f"Conflicting autotune parameters: {names}")

360 + keys.update(derive_inner(dict(kwargs) | config.all_kwargs()))

361 + return keys

362 +

363 + return derive_autotuned

364 +

365 + if not isinstance(kernel, JITFunction):

366 + raise TypeError(f"Unsupported Triton kernel wrapper: {type(kernel).__name__}")

367 +

368 + device = cast(Hashable, driver.active.get_current_device())

369 + _, kernel_key_cache, _, _, binder = kernel.device_caches[device]

370 + jit_function_id = id(kernel)

371 + jit_function_key = cast(Hashable, kernel.cache_key)

372 + debug = kernel.debug or knobs.runtime.debug

373 + instrumentation_mode = knobs.compilation.instrumentation_mode

374 + fpsan_homomorphic_casts = getattr(

375 + knobs.compilation, "fpsan_homomorphic_casts", None

376 + )

377 + inspection_hook = knobs.runtime.add_stages_inspection_hook

378 +

379 + def derive_jit(kwargs: Mapping[str, Any]) -> set[TritonJitKey]:

380 + binder_kwargs = dict(kwargs)

381 + binder_kwargs["debug"] = binder_kwargs.get("debug", debug) or debug

382 + binder_kwargs["instrumentation_mode"] = instrumentation_mode

383 + if fpsan_homomorphic_casts is not None:

384 + binder_kwargs["fpsan_homomorphic_casts"] = fpsan_homomorphic_casts

385 + _, specialization, options = binder(**binder_kwargs)

386 + if inspection_hook is not None:

387 + _, inspection_hash = inspection_hook()

388 + specialization.append(f'("custom_pipeline", {inspection_hash})')

389 + cache_key = cast(

390 + Hashable,

391 + compute_cache_key(kernel_key_cache, specialization, options),

392 + )

393 + return {

394 + TritonJitKey(

395 + jit_function_id,

396 + jit_function_key,

397 + device,

398 + cache_key,

399 + )

400 + }

401 +

402 + return derive_jit

403 +

404 +

405 +WarmupCases = Mapping[str, Any] | Iterable[Mapping[str, Any]]

406 +

407 +

408 +class _AutomaticTritonJitKernel(VllmTritonJitKernel[TritonCompileKey]):

409 + def warmup_inputs(self, compile_key: TritonCompileKey) -> dict[str, Any]:

410 + return dict(compile_key.inputs)

411 +

412 + def _provider_cases(

413 + self,

414 + provider: Callable[..., WarmupCases],

415 + *args: Any,

416 + **kwargs: Any,

417 + ) -> Iterable[Mapping[str, Any]]:

418 + cases_node = get_function_source_node(provider)

419 + if any(

420 + isinstance(node, ast.Call)

421 + and (get_ast_full_name(node.func) or "").split(".")[-1]

422 + in {"WarmupIntRange", "WarmupChoices", "_when"}

423 + for node in ast.walk(cases_node)

424 + ):

425 + return self._expand_warmup_cases(provider, *args, **kwargs)

426 + cases = provider(*args, **kwargs)

427 + return (cases,) if isinstance(cases, Mapping) else cases

428 +

429 +

430 +def _is_autotuned(kernel: Any) -> bool:

431 + Autotuner = triton.runtime.autotuner.Autotuner

432 +

433 + while kernel is not None:

434 + if isinstance(kernel, Autotuner):

435 + return True

436 + kernel = getattr(kernel, "fn", None)

437 + return False

438 +

439 +

440 +def _triton_metadata_arg(value: Any) -> Any:

441 + from torch._subclasses.fake_tensor import FakeTensor

442 +

443 + if not isinstance(value, FakeTensor):

444 + return value

445 + return TritonWarmupTensor(

446 + value.dtype,

447 + aligned=getattr(value, "_vllm_warmup_aligned", True),

448 + shape=tuple(value.shape),

449 + strides=tuple(value.stride()),

450 + )

451 +

452 +

453 +def _materialize_warmup_case(

454 + case: Mapping[str, Any],

455 + *,

456 + real: bool,

457 + mode: Any,

458 + cache: dict[TritonWarmupTensor, Any],

459 +) -> dict[str, Any]:

460 + import torch

461 +

462 + def materialize(value: Any) -> Any:

463 + if not isinstance(value, TritonWarmupTensor):

464 + return value

465 + if not real:

466 + try:

467 + return cache[value]

468 + except KeyError:

469 + pass

470 + except TypeError:

471 + # Unhashable initializers are uncommon and safe to materialize directly.

472 + pass

473 +

474 + strides = cast(tuple[int, ...], value.stride())

475 + with mode:

476 + tensor = torch.empty_strided(

477 + value.shape,

478 + strides,

479 + dtype=value.dtype,

480 + device=value.device,

481 + )

482 + if real:

483 + if callable(value.init):

484 + value.init(tensor)

485 + else:

486 + tensor.fill_(value.init)

487 + else:

488 + tensor._vllm_warmup_aligned = value.aligned

489 + with suppress(TypeError):

490 + cache[value] = tensor

491 + return tensor

492 +

493 + return {name: materialize(value) for name, value in case.items()}

494 +

495 +

496 +class _DecoratedTritonJitKernel(_AutomaticTritonJitKernel):

497 + def __init__(

498 + self,

499 + kernel: Any,

500 + warmup_inputs: Callable[..., WarmupCases],

501 + dispatch: Callable[..., DispatchSpec] | None,

502 + ) -> None:

503 + self.kernel = kernel

504 + self._run_autotune = _is_autotuned(kernel)

505 + self._warmup_inputs_fn = warmup_inputs

506 + self._dispatch_fn = dispatch

507 + self._dispatch_arg_names: tuple[str, ...] = ()

508 + self._dispatch_defaults: dict[str, Any] = {}

509 + if dispatch is not None:

510 + parameters = tuple(inspect.signature(dispatch).parameters.values())

511 + if any(

512 + parameter.kind

513 + in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)

514 + for parameter in parameters

515 + ):

516 + raise TypeError("Triton dispatchers require explicit parameters")

517 + self._dispatch_arg_names = tuple(parameter.name for parameter in parameters)

518 + self._dispatch_defaults = {

519 + parameter.name: parameter.default

520 + for parameter in parameters

521 + if parameter.default is not inspect.Parameter.empty

522 + }

523 + super().__init__()

524 +

525 + def __call__(self, *args: Any, **kwargs: Any) -> Any:

526 + if self._dispatch_fn is None:

527 + grid = kwargs.pop("grid") if not args else args[0]

528 + kernel_args = () if not args else args[1:]

529 + inputs = triton_warmup_inputs(

530 + self.kernel, *kernel_args, grid=grid, **kwargs

531 + )

532 + inputs.pop("grid")

533 + return self.launch((grid, inputs), {})

534 + spec = self._dispatch_fn(*args, **kwargs)

535 + input_values = args

536 + if len(args) != len(self._dispatch_arg_names):

537 + input_values += tuple(

538 + kwargs[name] if name in kwargs else self._dispatch_defaults[name]

539 + for name in self._dispatch_arg_names[len(args) :]

540 + )

541 + return self.launch(

542 + spec,

543 + dict(zip(self._dispatch_arg_names, input_values)),

544 + )

545 +

546 + def __getitem__(self, grid: tuple[int, ...]) -> Callable[..., Any]:

547 + return cast(Callable[..., Any], self.kernel[grid])

548 +

549 + def __getattr__(self, name: str) -> Any:

550 + return getattr(self.kernel, name)

551 +

552 + def get_warmup_keys(

553 + self, *, vllm_config: Any | None = None, **kwargs: Any

554 + ) -> list[TritonCompileKey]:

555 + keys: list[TritonCompileKey] = []

556 + seen_jit_keys: set[frozenset[TritonJitKey]] = set()

557 + if "vllm_config" in inspect.signature(self._warmup_inputs_fn).parameters:

558 + kwargs.setdefault("vllm_config", vllm_config)

559 + cases = self._provider_cases(self._warmup_inputs_fn, **kwargs)

560 + derive_keys = _triton_key_deriver(self.kernel)

561 + materialize = self._run_autotune

562 + mode: Any = nullcontext() if self._run_autotune else None

563 + materialized: dict[TritonWarmupTensor, Any] = {}

564 + for case in cases:

565 + inputs = (

566 + _materialize_warmup_case(

567 + case,

568 + real=self._run_autotune,

569 + mode=mode,

570 + cache=materialized,

571 + )

572 + if materialize

573 + else case

574 + )

575 + if self._dispatch_fn is None:

576 + prepared = dict(inputs)

577 + prepared.pop("grid")

578 + else:

579 + dispatch_inputs = {

580 + name: inputs[name]

581 + for name in self._dispatch_arg_names

582 + if name in inputs

583 + }

584 + try:

585 + _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2]

586 + except (AssertionError, AttributeError, TypeError):

587 + if materialize:

588 + raise

589 + from torch._subclasses.fake_tensor import FakeTensorMode

590 +

591 + materialize = True

592 + mode = FakeTensorMode()

593 + inputs = _materialize_warmup_case(

594 + case,

595 + real=False,

596 + mode=mode,

597 + cache=materialized,

598 + )

599 + dispatch_inputs = {

600 + name: inputs[name]

601 + for name in self._dispatch_arg_names

602 + if name in inputs

603 + }

604 + _, launch_kwargs = self._dispatch_fn(**dispatch_inputs)[:2]

605 + prepared = self._prepare_launch_kwargs(

606 + tuple(inputs),

607 + tuple(inputs.values()),

608 + dict(launch_kwargs),

609 + )

610 + if materialize:

611 + prepared = {

612 + name: _triton_metadata_arg(value)

613 + for name, value in prepared.items()

614 + }

615 + jit_keys = frozenset(derive_keys(prepared))

616 + if jit_keys and jit_keys not in seen_jit_keys:

617 + seen_jit_keys.add(jit_keys)

618 + keys.append(TritonCompileKey(jit_keys, tuple(inputs.items())))

619 + return keys

620 +

621 +

622 +class TritonKernelDispatcher(Protocol[P]):

623 + """Callable Triton launcher with compile-time warmup support."""

624 +

625 + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Any: ...

626 +

627 + def register_warmup(self, *args: Any, **kwargs: Any) -> None: ...

628 +

629 +

630 +@overload

631 +def triton_kernel_dispatcher_with_warmup(

632 + *,

633 + warmup_inputs: Callable[..., WarmupCases],

634 + kernel: None = None,

635 +) -> Callable[[Any], Any]: ...

636 +

637 +

638 +@overload

639 +def triton_kernel_dispatcher_with_warmup(

640 + *,

641 + warmup_inputs: Callable[..., WarmupCases],

642 + kernel: Any,

643 +) -> Callable[[Callable[P, DispatchSpec]], TritonKernelDispatcher[P]]: ...

644 +

645 +

646 +def triton_kernel_dispatcher_with_warmup(

647 + *,

648 + warmup_inputs: Callable[..., WarmupCases],

649 + kernel: Any | None = None,

650 +) -> Callable[[Any], Any]:

651 + """Decorate a dispatch function or a native Triton kernel for warmup."""

652 +

653 + def decorate(dispatch: Any) -> Any:

654 + owner = _DecoratedTritonJitKernel(

655 + dispatch if kernel is None else kernel,

656 + warmup_inputs,

657 + None if kernel is None else dispatch,

658 + )

659 + update_wrapper(owner, dispatch, updated=())

660 + return owner

661 +

662 + return decorate

663 +

664 +

665 @dataclass(frozen=True)

666 class TritonPointerInputVariant:

667 # Named pointer-alignment variant for compile-only Triton warmup.

vllm/model_executor/warmup/kernel_warmup.py

+4 -5

@@ -41,9 +41,6 @@ from vllm.model_executor.warmup.qwen_vl_triton_warmup import qwen_vl_triton_warm

41 from vllm.model_executor.warmup.replayssm_warmup import (

42 replayssm_autotune_warmup,

43 )

44 -from vllm.model_executor.warmup.spec_decode_rejection_warmup import (

45 - spec_decode_rejection_warmup,

46 -)

44 from vllm.model_executor.warmup.watermark_sample_warmup import (

45 watermark_sample_warmup,

46 )

@@ -178,7 +175,10 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):

175 logger.info("JIT kernel warmup starting.")

176 jit_warmup_start = time.perf_counter()

177 try:

181 - worker.model_runner.jit_warmup_registry.warmup()

178 + registry = (

179 + worker.model_runner.jit_warmup_registry # type: ignore[attr-defined]

180 + )

181 + registry.warmup()

182 except Exception:

183 logger.exception(

184 "JIT kernel warmup failed after %.2fs.",

@@ -209,7 +209,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):

209 # Run next so input-prep kernels JIT against pristine runner state.

210 if enable_jit_warmup:

211 kimi_k3_triton_warmup(worker)

212 - spec_decode_rejection_warmup(worker)

212 watermark_sample_warmup(worker)

213 qwen4_exp_qsa_triton_warmup(worker)

214

vllm/model_executor/warmup/spec_decode_rejection_warmup.py deleted

-139

@@ -1,139 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Warm up speculative-decoding rejection kernels."""

4 -

5 -from __future__ import annotations

6 -

7 -from typing import TYPE_CHECKING, Any

8 -

9 -import torch

10 -

11 -from vllm.logger import init_logger

12 -

13 -if TYPE_CHECKING:

14 - from vllm.v1.worker.gpu_worker import Worker

15 -

16 -logger = init_logger(__name__)

17 -

18 -

19 -@torch.inference_mode()

20 -def spec_decode_rejection_warmup(worker: Worker) -> None:

21 - spec_config = worker.vllm_config.speculative_config

22 - if spec_config is None:

23 - return

24 -

25 - from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (

26 - rejection_sample,

27 - )

28 -

29 - model_config = worker.vllm_config.model_config

30 - vocab_size = model_config.get_vocab_size()

31 - num_spec = spec_config.num_speculative_tokens

32 - if num_spec <= 0 or vocab_size <= 0:

33 - return

34 -

35 - rejection_method = getattr(spec_config, "rejection_sample_method", None)

36 - use_block_verification = rejection_method == "block"

37 - use_synthetic = rejection_method == "synthetic"

38 - # HAS_DRAFT_LOGITS is a constexpr, so the warmup must match the runtime.

39 - use_draft_logits = (

40 - getattr(spec_config, "draft_sample_method", "greedy") == "probabilistic"

41 - )

42 -

43 - device = worker.device

44 - num_reqs = 1

45 - tokens_per_req = num_spec + 1

46 - num_logits = num_reqs * tokens_per_req

47 -

48 - watermark_config = getattr(worker.vllm_config, "watermark_config", None)

49 - watermark_kwargs: dict[str, Any] = {}

50 - if watermark_config is not None:

51 - from vllm.v1.watermarking.spec_decode import speculative_target_watermark_key

52 -

53 - try:

54 - watermark_key = speculative_target_watermark_key(watermark_config)

55 - except Exception:

56 - logger.warning(

57 - "Skipping spec-decode rejection sampler warmup: could not derive "

58 - "the watermark recovery key.",

59 - exc_info=True,

60 - )

61 - return

62 - # Tensor dtypes and integer widths affect Triton specialization.

63 - watermark_kwargs = {

64 - "contexts": torch.zeros(

65 - (num_logits, watermark_config.context_width),

66 - dtype=torch.int32,

67 - device=device,

68 - ),

69 - "watermarking": torch.zeros(num_reqs, dtype=torch.bool, device=device),

70 - "watermark_key": watermark_key,

71 - }

72 -

73 - # Sampling-parameter processing may promote either distribution to fp32.

74 - model_dtype = model_config.dtype

75 - warmup_dtype_pairs: set[tuple[torch.dtype, torch.dtype | None]]

76 - if use_draft_logits:

77 - warmup_dtype_pairs = {

78 - (model_dtype, model_dtype),

79 - (torch.float32, torch.float32),

80 - (torch.float32, model_dtype),

81 - (model_dtype, torch.float32),

82 - }

83 - else:

84 - warmup_dtype_pairs = {(model_dtype, None), (torch.float32, None)}

85 -

86 - logger.info(

87 - "Warming up spec-decode rejection sampler kernels "

88 - "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s, watermark=%s).",

89 - vocab_size,

90 - num_spec,

91 - [(str(t), str(d)) for t, d in warmup_dtype_pairs],

92 - use_block_verification,

93 - watermark_config is not None,

94 - )

95 - for tgt_dtype, draft_dtype in warmup_dtype_pairs:

96 - target_logits = torch.zeros(

97 - (num_logits, vocab_size), dtype=tgt_dtype, device=device

98 - )

99 - draft_logits = (

100 - torch.zeros(

101 - (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device

102 - )

103 - if draft_dtype is not None

104 - else None

105 - )

106 - synthetic_rates = (

107 - torch.full((num_spec,), 0.5, dtype=torch.float32, device=device)

108 - if use_synthetic

109 - else None

110 - )

111 - try:

112 - rejection_sample(

113 - target_logits=target_logits,

114 - draft_logits=draft_logits,

115 - draft_sampled=torch.zeros(num_logits, dtype=torch.int32, device=device),

116 - cu_num_logits=torch.tensor(

117 - [0, num_logits], dtype=torch.int32, device=device

118 - ),

119 - pos=torch.zeros(num_logits, dtype=torch.int64, device=device),

120 - idx_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device),

121 - expanded_idx_mapping=torch.zeros(

122 - num_logits, dtype=torch.int64, device=device

123 - ),

124 - expanded_local_pos=torch.arange(

125 - num_logits, dtype=torch.int32, device=device

126 - ),

127 - temperature=torch.zeros(num_reqs, dtype=torch.float32, device=device),

128 - seed=torch.full((num_reqs,), 42, dtype=torch.int64, device=device),

129 - num_speculative_steps=num_spec,

130 - synthetic_conditional_rates=synthetic_rates,

131 - use_fp64=False,

132 - use_block_verification=use_block_verification,

133 - **watermark_kwargs,

134 - )

135 - except Exception:

136 - logger.warning(

137 - "Skipping spec-decode rejection sampler warmup.", exc_info=True

138 - )

139 - return

vllm/models/glm5next/nvidia/kda.py

@@ -175,6 +175,15 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

175 projection_size = self.head_dim * self.num_heads

176 self.local_projection_size = divide(projection_size, self.tp_size)

177

178 + _, recurrent_state_shape = self.get_state_shape()

179 + _, recurrent_state_dtype = self.get_state_dtype()

180 + scatter_states.register_warmup(

181 + state_shape=recurrent_state_shape,

182 + state_dtype=recurrent_state_dtype,

183 + indices_dtype=torch.int32,

184 + max_num_tokens=vllm_config.scheduler_config.max_num_seqs,

185 + )

186 +

187 # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches).

188 # Order matches checkpoint's fused_qkvbfg_a_proj convention.

189 # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks.

vllm/v1/sample/ops/topk_topp_sampler.py

+20 -1

@@ -13,11 +13,27 @@ from vllm.platforms import CpuArchEnum, current_platform

13 from vllm.triton_utils import HAS_TRITON

15 if HAS_TRITON:

16 - from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton

16 + from vllm.v1.sample.ops.topk_topp_triton import (

17 + _topk_topp,

18 + _topp_split_mask,

19 + _topp_split_stats,

20 + _topp_split_step,

21 + apply_top_k_top_p_triton,

22 + )

24 logger = init_logger(__name__)

27 +def register_top_k_top_p_warmups() -> None:

28 + """Register every native accelerator sampling kernel used at runtime."""

29 + if HAS_TRITON and not current_platform.is_cpu():

30 + _topk_topp.register_warmup()

31 + if current_platform.is_cuda_alike():

32 + _topp_split_stats.register_warmup()

33 + _topp_split_step.register_warmup()

34 + _topp_split_mask.register_warmup()

35 +

36 +

37 def _skip_aiter_sampler_on_gfx1250() -> bool:

38 # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU.

39 from vllm.platforms.rocm import on_gfx1250

@@ -128,6 +144,9 @@ class TopKTopPSampler(nn.Module):

144 else:

145 self.forward = self.forward_native

146

147 + # Every accelerator backend can fall back to native sampling at runtime.

148 + register_top_k_top_p_warmups()

149 +

150 def forward_native(

151 self,

152 logits: torch.Tensor,

vllm/v1/sample/ops/topk_topp_triton.py

+267 -66

@@ -9,8 +9,21 @@ using Pivot-based Truncation and Selection" By Park et al.

10 """

12 +from typing import Any

13 +

14 import torch

16 +from vllm.model_executor.warmup.jit_warmup import (

17 + WarmupChoices,

18 + WarmupIntRange,

19 + _when,

20 +)

21 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

22 + DispatchSpec,

23 + TritonWarmupTensor,

24 + triton_kernel_dispatcher_with_warmup,

25 +)

26 +from vllm.platforms import current_platform

27 from vllm.triton_utils import tl, triton

28 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

29 from vllm.utils.math_utils import next_power_of_2

@@ -863,6 +876,90 @@ def _topk_topp_kernel(

876 tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n)

877

878

879 +def _max_sampler_batch_size(vllm_config: Any) -> int:

880 + return vllm_config.scheduler_config.max_num_seqs * (

881 + vllm_config.num_speculative_tokens + 1

882 + )

883 +

884 +

885 +def _topk_topp_warmup_inputs(vllm_config: Any) -> dict[str, Any]:

886 + vocab_size = vllm_config.model_config.get_vocab_size()

887 + split_enabled = current_platform.is_cuda_alike()

888 + mode: Any = WarmupChoices((True, False), (True, True), (False, True))

889 + logits_stride: Any = WarmupChoices(16, 2)

890 + batch_size: Any = WarmupIntRange(1, _max_sampler_batch_size(vllm_config) + 1)

891 + topk_enabled = mode[0]

892 + topp_enabled = mode[1]

893 + _when(topk_enabled or not (split_enabled and batch_size <= _SPLIT_MAX_BATCH))

894 + logits = TritonWarmupTensor(

895 + torch.float32,

896 + shape=(batch_size, vocab_size),

897 + strides=(logits_stride, 1),

898 + )

899 + return dict(

900 + logits=logits,

901 + buffer=TritonWarmupTensor(torch.float32),

902 + percentile_to_std_table=TritonWarmupTensor(torch.float32),

903 + normal_cdf_to_sigma_table=TritonWarmupTensor(torch.float32),

904 + k=TritonWarmupTensor(torch.int32) if topk_enabled else None,

905 + p=TritonWarmupTensor(torch.float32) if topp_enabled else None,

906 + mask_value=float("-inf"),

907 + num_sm=num_compute_units(),

908 + )

909 +

910 +

911 +@triton_kernel_dispatcher_with_warmup(

912 + kernel=_topk_topp_kernel,

913 + warmup_inputs=_topk_topp_warmup_inputs,

914 +)

915 +def _topk_topp(

916 + logits: torch.Tensor,

917 + buffer: torch.Tensor,

918 + percentile_to_std_table: torch.Tensor,

919 + normal_cdf_to_sigma_table: torch.Tensor,

920 + k: torch.Tensor | None,

921 + p: torch.Tensor | None,

922 + mask_value: float,

923 + num_sm: int,

924 +) -> DispatchSpec:

925 + batch_size, vocab_size = logits.shape

926 + topk_enabled = k is not None

927 + topp_enabled = p is not None

928 + k_ptr = k if k is not None else logits

929 + p_ptr = p if p is not None else logits

930 + num_programs = min(num_sm, batch_size)

931 + # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.

932 + # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot

933 + # approximation; use smaller tiles for accurate top-p results.

934 + if logits.device.type == "cpu":

935 + block_size, block_size_trunc, num_warps = 256, 128, 4

936 + elif logits.device.type == "xpu":

937 + block_size, block_size_trunc, num_warps = 4096, 2048, 4

938 + else:

939 + block_size, block_size_trunc, num_warps = 8192, 4096, 8

940 + # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so

941 + # per-tile latency bounds kernel latency, and Triton's default of 4

942 + # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is

943 + # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.

944 + split_covers_ponly = (

945 + current_platform.is_cuda_alike()

946 + and topp_enabled

947 + and batch_size <= _SPLIT_MAX_BATCH

948 + )

949 + return (num_programs,), dict(

950 + K=k_ptr,

951 + P=p_ptr,

952 + BATCH_SIZE=batch_size,

953 + VOCAB_SIZE=vocab_size,

954 + BLOCK_SIZE=block_size,

955 + BLOCK_SIZE_TRUNC=block_size_trunc,

956 + TOPK_ENABLED=topk_enabled,

957 + TOPP_ENABLED=topp_enabled,

958 + SPLIT_COVERS_PONLY=split_covers_ponly,

959 + num_warps=num_warps,

960 + )

961 +

962 +

963 # ---------------------------------------------------------------------------

964 # Split-row pipeline for small batches.

965 #

@@ -892,6 +989,13 @@ _SPLIT_ROUNDS = 5

989 _TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {}

990

991

992 +def _topp_split_count(batch_size: int, num_sm: int) -> int:

993 + splits = 1

994 + while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:

995 + splits *= 2

996 + return splits

997 +

998 +

999 @triton.jit

1000 def _topp_sb_stats_kernel(

1001 LOGITS,

@@ -939,6 +1043,51 @@ def _topp_sb_stats_kernel(

1043 tl.store(STATS + base + 3, cnt)

1044

1045

1046 +def _topp_split_stats_warmup_inputs(vllm_config: Any) -> dict[str, Any]:

1047 + vocab_size = vllm_config.model_config.get_vocab_size()

1048 + num_sm = num_compute_units()

1049 + logits_stride: Any = WarmupChoices(16, 2)

1050 + batch_size: Any = WarmupIntRange(

1051 + 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1

1052 + )

1053 + has_k: Any = WarmupChoices(False, True)

1054 + logits = TritonWarmupTensor(

1055 + torch.float32,

1056 + shape=(batch_size, vocab_size),

1057 + strides=(logits_stride, 1),

1058 + )

1059 + return dict(

1060 + logits=logits,

1061 + stats=TritonWarmupTensor(torch.float32),

1062 + k=TritonWarmupTensor(torch.int32) if has_k else None,

1063 + p=TritonWarmupTensor(torch.float32),

1064 + num_sm=num_sm,

1065 + )

1066 +

1067 +

1068 +@triton_kernel_dispatcher_with_warmup(

1069 + kernel=_topp_sb_stats_kernel,

1070 + warmup_inputs=_topp_split_stats_warmup_inputs,

1071 +)

1072 +def _topp_split_stats(

1073 + logits: torch.Tensor,

1074 + stats: torch.Tensor,

1075 + k: torch.Tensor | None,

1076 + p: torch.Tensor,

1077 + num_sm: int,

1078 +) -> DispatchSpec:

1079 + batch_size, vocab_size = logits.shape

1080 + splits = _topp_split_count(batch_size, num_sm)

1081 + return (logits.shape[0] * splits,), dict(

1082 + K=k if k is not None else logits,

1083 + HAS_K=k is not None,

1084 + VOCAB_SIZE=vocab_size,

1085 + S=splits,

1086 + BLOCK=8192,

1087 + num_warps=8,

1088 + )

1089 +

1090 +

1091 @triton.jit

1092 def _topp_sb_row_stats(STATS, row, S: tl.constexpr):

1093 """Combine per-slice partials into row (max, sum_exp, min, n_finite)."""

@@ -1183,6 +1332,58 @@ def _topp_sb_step_kernel(

1332 tl.store(pbase + fidx * 3 + 2, nmin)

1333

1334

1335 +def _topp_split_step_warmup_inputs(vllm_config: Any) -> dict[str, Any]:

1336 + vocab_size = vllm_config.model_config.get_vocab_size()

1337 + num_sm = num_compute_units()

1338 + logits_stride: Any = WarmupChoices(16, 2)

1339 + batch_size: Any = WarmupIntRange(

1340 + 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1

1341 + )

1342 + round: Any = WarmupIntRange(0, _SPLIT_ROUNDS)

1343 + has_k: Any = WarmupChoices(False, True)

1344 + logits = TritonWarmupTensor(

1345 + torch.float32,

1346 + shape=(batch_size, vocab_size),

1347 + strides=(logits_stride, 1),

1348 + )

1349 + return dict(

1350 + logits=logits,

1351 + stats=TritonWarmupTensor(torch.float32),

1352 + parts=TritonWarmupTensor(torch.float32),

1353 + k=TritonWarmupTensor(torch.int32) if has_k else None,

1354 + p=TritonWarmupTensor(torch.float32),

1355 + round=round,

1356 + num_sm=num_sm,

1357 + )

1358 +

1359 +

1360 +@triton_kernel_dispatcher_with_warmup(

1361 + kernel=_topp_sb_step_kernel,

1362 + warmup_inputs=_topp_split_step_warmup_inputs,

1363 +)

1364 +def _topp_split_step(

1365 + logits: torch.Tensor,

1366 + stats: torch.Tensor,

1367 + parts: torch.Tensor,

1368 + k: torch.Tensor | None,

1369 + p: torch.Tensor,

1370 + round: int,

1371 + num_sm: int,

1372 +) -> DispatchSpec:

1373 + batch_size, vocab_size = logits.shape

1374 + splits = _topp_split_count(batch_size, num_sm)

1375 + return (logits.shape[0] * splits,), dict(

1376 + K=k if k is not None else logits,

1377 + HAS_K=k is not None,

1378 + S=splits,

1379 + F=_SPLIT_FANOUT,

1380 + NUM_ROUNDS=_SPLIT_ROUNDS,

1381 + VOCAB_SIZE=vocab_size,

1382 + BLOCK=2048,

1383 + num_warps=8,

1384 + )

1385 +

1386 +

1387 @triton.jit

1388 def _topp_sb_mask_kernel(

1389 LOGITS,

@@ -1313,6 +1514,57 @@ def _topp_sb_mask_kernel(

1514 tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n)

1515

1516

1517 +def _topp_split_mask_warmup_inputs(vllm_config: Any) -> dict[str, Any]:

1518 + vocab_size = vllm_config.model_config.get_vocab_size()

1519 + num_sm = num_compute_units()

1520 + logits_stride: Any = WarmupChoices(16, 2)

1521 + batch_size: Any = WarmupIntRange(

1522 + 1, min(_max_sampler_batch_size(vllm_config), _SPLIT_MAX_BATCH) + 1

1523 + )

1524 + has_k: Any = WarmupChoices(False, True)

1525 + logits = TritonWarmupTensor(

1526 + torch.float32,

1527 + shape=(batch_size, vocab_size),

1528 + strides=(logits_stride, 1),

1529 + )

1530 + return dict(

1531 + logits=logits,

1532 + stats=TritonWarmupTensor(torch.float32),

1533 + parts=TritonWarmupTensor(torch.float32),

1534 + k=TritonWarmupTensor(torch.int32) if has_k else None,

1535 + p=TritonWarmupTensor(torch.float32),

1536 + mask_value=float("-inf"),

1537 + num_sm=num_sm,

1538 + )

1539 +

1540 +

1541 +@triton_kernel_dispatcher_with_warmup(

1542 + kernel=_topp_sb_mask_kernel,

1543 + warmup_inputs=_topp_split_mask_warmup_inputs,

1544 +)

1545 +def _topp_split_mask(

1546 + logits: torch.Tensor,

1547 + stats: torch.Tensor,

1548 + parts: torch.Tensor,

1549 + k: torch.Tensor | None,

1550 + p: torch.Tensor,

1551 + mask_value: float,

1552 + num_sm: int,

1553 +) -> DispatchSpec:

1554 + batch_size, vocab_size = logits.shape

1555 + splits = _topp_split_count(batch_size, num_sm)

1556 + return (logits.shape[0] * splits,), dict(

1557 + K=k if k is not None else logits,

1558 + HAS_K=k is not None,

1559 + S=splits,

1560 + F=_SPLIT_FANOUT,

1561 + NUM_ROUNDS=_SPLIT_ROUNDS,

1562 + VOCAB_SIZE=vocab_size,

1563 + BLOCK=8192,

1564 + num_warps=8,

1565 + )

1566 +

1567 +

1568 def _apply_topp_split(

1569 logits: torch.Tensor,

1570 k: torch.Tensor | None,

@@ -1321,11 +1573,6 @@ def _apply_topp_split(

1573 num_sm: int,

1574 ) -> None:

1575 """Split-row top-p pipeline for small batches; masks logits in place."""

1324 - batch_size, vocab_size = logits.shape

1325 - splits = 1

1326 - while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm:

1327 - splits *= 2

1328 -

1576 ws = _TRITON_SPLIT_CACHE.get(logits.device)

1577 if ws is None:

1578 ws = {

@@ -1336,53 +1583,31 @@ def _apply_topp_split(

1583 }

1584 _TRITON_SPLIT_CACHE[logits.device] = ws

1585

1339 - k_ptr = k if k is not None else logits # dummy pointer when HAS_K=False

1340 - has_k = k is not None

1341 -

1342 - _topp_sb_stats_kernel[(batch_size * splits,)](

1586 + _topp_split_stats(

1587 logits,

1344 - logits.stride(0),

1588 ws["stats"],

1346 - k_ptr,

1589 + k,

1590 p,

1348 - HAS_K=has_k,

1349 - VOCAB_SIZE=vocab_size,

1350 - S=splits,

1351 - BLOCK=8192,

1352 - num_warps=8,

1591 + num_sm,

1592 )

1593 for round_i in range(_SPLIT_ROUNDS):

1355 - _topp_sb_step_kernel[(batch_size * splits,)](

1594 + _topp_split_step(

1595 logits,

1357 - logits.stride(0),

1596 ws["stats"],

1597 ws["parts"],

1360 - k_ptr,

1598 + k,

1599 p,

1600 round_i,

1363 - HAS_K=has_k,

1364 - S=splits,

1365 - F=_SPLIT_FANOUT,

1366 - NUM_ROUNDS=_SPLIT_ROUNDS,

1367 - VOCAB_SIZE=vocab_size,

1368 - BLOCK=2048,

1369 - num_warps=8,

1601 + num_sm,

1602 )

1371 - _topp_sb_mask_kernel[(batch_size * splits,)](

1603 + _topp_split_mask(

1604 logits,

1373 - logits.stride(0),

1605 ws["stats"],

1606 ws["parts"],

1376 - k_ptr,

1607 + k,

1608 p,

1378 - HAS_K=has_k,

1379 - MASK_VALUE=mask_value,

1380 - S=splits,

1381 - F=_SPLIT_FANOUT,

1382 - NUM_ROUNDS=_SPLIT_ROUNDS,

1383 - VOCAB_SIZE=vocab_size,

1384 - BLOCK=8192,

1385 - num_warps=8,

1609 + mask_value,

1610 + num_sm,

1611 )

1612

1613

@@ -1477,39 +1702,15 @@ def apply_top_k_top_p_triton(

1702 else:

1703 normal_cdf_to_sigma_table, percentile_to_std_table = tables

1704

1480 - # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.

1481 - # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot

1482 - # approximation; use smaller tiles for accurate top-p results.

1483 - launch_kwargs = {}

1484 - if logits.device.type == "cpu":

1485 - block_size, block_size_trunc = 256, 128

1486 - elif logits.device.type == "xpu":

1487 - block_size, block_size_trunc = 4096, 2048

1488 - else:

1489 - block_size, block_size_trunc = 8192, 4096

1490 - # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so

1491 - # per-tile latency bounds kernel latency, and Triton's default of 4

1492 - # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is

1493 - # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.

1494 - launch_kwargs["num_warps"] = 8

1495 -

1496 - _topk_topp_kernel[(NUM_PROGRAMS,)](

1705 + _topk_topp(

1706 logits,

1498 - logits.stride(0),

1707 buffer,

1708 percentile_to_std_table,

1709 normal_cdf_to_sigma_table,

1502 - k_ptr,

1503 - p_ptr,

1504 - BATCH_SIZE=batch_size,

1505 - MASK_VALUE=mask_value,

1506 - VOCAB_SIZE=vocab_size,

1507 - BLOCK_SIZE=block_size,

1508 - BLOCK_SIZE_TRUNC=block_size_trunc,

1509 - TOPK_ENABLED=topk_enabled,

1510 - TOPP_ENABLED=topp_enabled,

1511 - SPLIT_COVERS_PONLY=use_split,

1512 - **launch_kwargs,

1710 + k_ptr if topk_enabled else None,

1711 + p_ptr if topp_enabled else None,

1712 + mask_value,

1713 + num_sm,

1714 )

1715 if use_split:

1716 _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm)

vllm/v1/sample/rejection_sampler.py

+186 -15

@@ -5,13 +5,19 @@ from __future__ import annotations

6 from collections.abc import Sequence

7 from dataclasses import replace

8 -from typing import TYPE_CHECKING

8 +from typing import TYPE_CHECKING, Any

10 import torch

11 import torch.nn as nn

13 from vllm.config.model import PROCESSED_LOGPROBS_MODES

14 from vllm.logger import init_logger

15 +from vllm.model_executor.warmup.jit_warmup import WarmupChoices

16 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

17 + DispatchSpec,

18 + TritonWarmupTensor,

19 + triton_kernel_dispatcher_with_warmup,

20 +)

21 from vllm.triton_utils import tl, triton

22 from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput

23 from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor

@@ -88,6 +94,14 @@ class RejectionSampler(nn.Module):

94 device=device,

95 )

96 self.synthetic_mode = self.synthetic_conditional_rates is not None

97 + _rejection_greedy_sample.register_warmup(synthetic_mode=self.synthetic_mode)

98 + _rejection_random_sample.register_warmup(

99 + synthetic_mode=self.synthetic_mode,

100 + )

101 + _expand.register_warmup()

102 + _sample_recovered_tokens.register_warmup(

103 + use_fp64_gumbel=self.use_fp64_gumbel,

104 + )

105

106 def forward(

107 self,

@@ -453,7 +467,7 @@ def rejection_sample(

467 if not sampling_metadata.all_random:

468 # Rejection sampling for greedy sampling requests.

469 target_argmax = target_logits.argmax(dim=-1)

456 - rejection_greedy_sample_kernel[(batch_size,)](

470 + _rejection_greedy_sample(

471 output_token_ids,

472 cu_num_draft_tokens,

473 draft_token_ids,

@@ -463,7 +477,6 @@ def rejection_sample(

477 max_spec_len,

478 uniform_probs,

479 synthetic_conditional_rates,

466 - SYNTHETIC_MODE=synthetic_mode,

480 )

481 if sampling_metadata.all_greedy:

482 return output_token_ids

@@ -488,7 +501,7 @@ def rejection_sample(

501

502 # Rejection sampling for random sampling requests.

503 assert uniform_probs is not None

491 - rejection_random_sample_kernel[(batch_size,)](

504 + _rejection_random_sample(

505 output_token_ids,

506 cu_num_draft_tokens,

507 draft_token_ids,

@@ -499,10 +512,7 @@ def rejection_sample(

512 uniform_probs,

513 is_greedy,

514 max_spec_len,

502 - vocab_size,

515 synthetic_conditional_rates,

504 - NO_DRAFT_PROBS=draft_probs is None,

505 - SYNTHETIC_MODE=synthetic_mode,

516 )

517 return output_token_ids

518

@@ -594,13 +604,12 @@ def expand_batch_to_tokens(

604 batch_size = x.shape[0]

605 assert cu_num_tokens.shape[0] == batch_size

606 expanded_x = x.new_empty(num_tokens)

597 - expand_kernel[(batch_size,)](

607 + _expand(

608 expanded_x,

609 x,

610 cu_num_tokens,

611 replace_from,

612 replace_to,

603 - MAX_NUM_TOKENS=MAX_SPEC_LEN, # To avoid recompilation.

613 )

614 return expanded_x

615

@@ -694,18 +703,14 @@ def sample_recovered_tokens(

703 inv_q = q.reciprocal()

704

705 recovered_token_ids = torch.empty_like(draft_token_ids)

697 - BLOCK_SIZE = 8192

698 - sample_recovered_tokens_kernel[(batch_size, max_spec_len)](

706 + _sample_recovered_tokens(

707 recovered_token_ids,

708 cu_num_draft_tokens,

709 draft_token_ids,

710 draft_probs,

711 target_probs,

712 inv_q,

705 - vocab_size,

706 - BLOCK_SIZE,

707 - NO_DRAFT_PROBS=draft_probs is None,

708 - USE_FP64_GUMBEL=use_fp64_gumbel,

713 + max_spec_len=max_spec_len,

714 )

715 return recovered_token_ids

716

@@ -769,6 +774,50 @@ def rejection_greedy_sample_kernel(

774 )

775

776

777 +def _rejection_greedy_sample_warmup_inputs(

778 + *, synthetic_mode: bool

779 +) -> dict[str, object]:

780 + all_greedy = WarmupChoices(False, True)

781 + int32 = TritonWarmupTensor(torch.int32)

782 + return dict(

783 + output_token_ids=int32,

784 + cu_num_draft_tokens=int32,

785 + draft_token_ids=int32,

786 + target_argmax=TritonWarmupTensor(torch.int64),

787 + bonus_token_ids=TritonWarmupTensor(torch.int64),

788 + is_greedy=None if all_greedy else TritonWarmupTensor(torch.bool),

789 + max_spec_len=5,

790 + uniform_probs=(

791 + None

792 + if all_greedy and not synthetic_mode

793 + else TritonWarmupTensor(torch.float64)

794 + ),

795 + synthetic_conditional_rates=(

796 + TritonWarmupTensor(torch.float32) if synthetic_mode else None

797 + ),

798 + )

799 +

800 +

801 +@triton_kernel_dispatcher_with_warmup(

802 + kernel=rejection_greedy_sample_kernel,

803 + warmup_inputs=_rejection_greedy_sample_warmup_inputs,

804 +)

805 +def _rejection_greedy_sample(

806 + output_token_ids: torch.Tensor,

807 + cu_num_draft_tokens: torch.Tensor,

808 + draft_token_ids: torch.Tensor,

809 + target_argmax: torch.Tensor,

810 + bonus_token_ids: torch.Tensor,

811 + is_greedy: torch.Tensor | None,

812 + max_spec_len: int,

813 + uniform_probs: torch.Tensor | None,

814 + synthetic_conditional_rates: torch.Tensor | None,

815 +) -> DispatchSpec:

816 + return (output_token_ids.shape[0],), dict(

817 + SYNTHETIC_MODE=synthetic_conditional_rates is not None,

818 + )

819 +

820 +

821 # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.

822 @triton.jit(do_not_specialize=["max_spec_len"])

823 def rejection_random_sample_kernel(

@@ -845,6 +894,57 @@ def rejection_random_sample_kernel(

894 )

895

896

897 +def _rejection_random_sample_warmup_inputs(

898 + *, synthetic_mode: bool

899 +) -> dict[str, object]:

900 + vocab_size: Any = WarmupChoices(2, 16)

901 + no_draft_probs = WarmupChoices(False, True)

902 + int32 = TritonWarmupTensor(torch.int32)

903 + probabilities = TritonWarmupTensor(

904 + torch.float32,

905 + shape=(1, vocab_size),

906 + )

907 + return dict(

908 + output_token_ids=int32,

909 + cu_num_draft_tokens=int32,

910 + draft_token_ids=int32,

911 + draft_probs=None if no_draft_probs else probabilities,

912 + target_probs=probabilities,

913 + bonus_token_ids=TritonWarmupTensor(torch.int64),

914 + recovered_token_ids=int32,

915 + uniform_probs=TritonWarmupTensor(torch.float64),

916 + is_greedy=TritonWarmupTensor(torch.bool),

917 + max_spec_len=5,

918 + synthetic_conditional_rates=(

919 + TritonWarmupTensor(torch.float32) if synthetic_mode else None

920 + ),

921 + )

922 +

923 +

924 +@triton_kernel_dispatcher_with_warmup(

925 + kernel=rejection_random_sample_kernel,

926 + warmup_inputs=_rejection_random_sample_warmup_inputs,

927 +)

928 +def _rejection_random_sample(

929 + output_token_ids: torch.Tensor,

930 + cu_num_draft_tokens: torch.Tensor,

931 + draft_token_ids: torch.Tensor,

932 + draft_probs: torch.Tensor | None,

933 + target_probs: torch.Tensor,

934 + bonus_token_ids: torch.Tensor,

935 + recovered_token_ids: torch.Tensor,

936 + uniform_probs: torch.Tensor,

937 + is_greedy: torch.Tensor,

938 + max_spec_len: int,

939 + synthetic_conditional_rates: torch.Tensor | None,

940 +) -> DispatchSpec:

941 + return (output_token_ids.shape[0],), dict(

942 + vocab_size=target_probs.shape[1],

943 + NO_DRAFT_PROBS=draft_probs is None,

944 + SYNTHETIC_MODE=synthetic_conditional_rates is not None,

945 + )

946 +

947 +

948 # NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.

949 @triton.jit(do_not_specialize=["replace_from", "replace_to"])

950 def expand_kernel(

@@ -869,6 +969,31 @@ def expand_kernel(

969 tl.store(output_ptr + start_idx + offset, src_val, mask=offset < num_tokens)

970

971

972 +def _expand_warmup_inputs() -> dict[str, object]:

973 + dtype = WarmupChoices(torch.float32, torch.int32, torch.int64)

974 + data = TritonWarmupTensor(dtype)

975 + return dict(

976 + output=data,

977 + input=data,

978 + cu_num_tokens=TritonWarmupTensor(torch.int32),

979 + replace_from=0,

980 + replace_to=0,

981 + )

982 +

983 +

984 +@triton_kernel_dispatcher_with_warmup(

985 + kernel=expand_kernel, warmup_inputs=_expand_warmup_inputs

986 +)

987 +def _expand(

988 + output: torch.Tensor,

989 + input: torch.Tensor,

990 + cu_num_tokens: torch.Tensor,

991 + replace_from: int,

992 + replace_to: int,

993 +) -> DispatchSpec:

994 + return (input.shape[0],), dict(MAX_NUM_TOKENS=MAX_SPEC_LEN)

995 +

996 +

997 @triton.jit

998 def sample_recovered_tokens_kernel(

999 output_token_ids_ptr, # [num_tokens]

@@ -951,3 +1076,49 @@ def sample_recovered_tokens_kernel(

1076

1077 recovered_id = tl.minimum(recovered_id, vocab_size - 1)

1078 tl.store(output_token_ids_ptr + token_idx, recovered_id)

1079 +

1080 +

1081 +def _sample_recovered_tokens_warmup_inputs(

1082 + *, use_fp64_gumbel: bool

1083 +) -> dict[str, object]:

1084 + vocab_size: Any = WarmupChoices(2, 16)

1085 + no_draft_probs = WarmupChoices(False, True)

1086 + int32 = TritonWarmupTensor(torch.int32)

1087 + probabilities = TritonWarmupTensor(

1088 + torch.float32,

1089 + shape=(1, vocab_size),

1090 + )

1091 + return dict(

1092 + output_token_ids=int32,

1093 + cu_num_draft_tokens=int32,

1094 + draft_token_ids=int32,

1095 + draft_probs=None if no_draft_probs else probabilities,

1096 + target_probs=probabilities,

1097 + inv_q=TritonWarmupTensor(

1098 + torch.float64 if use_fp64_gumbel else torch.float32,

1099 + shape=(1, vocab_size),

1100 + ),

1101 + max_spec_len=5,

1102 + )

1103 +

1104 +

1105 +@triton_kernel_dispatcher_with_warmup(

1106 + kernel=sample_recovered_tokens_kernel,

1107 + warmup_inputs=_sample_recovered_tokens_warmup_inputs,

1108 +)

1109 +def _sample_recovered_tokens(

1110 + output_token_ids: torch.Tensor,

1111 + cu_num_draft_tokens: torch.Tensor,

1112 + draft_token_ids: torch.Tensor,

1113 + draft_probs: torch.Tensor | None,

1114 + target_probs: torch.Tensor,

1115 + inv_q: torch.Tensor,

1116 + *,

1117 + max_spec_len: int,

1118 +) -> DispatchSpec:

1119 + return (cu_num_draft_tokens.shape[0], max_spec_len), dict(

1120 + vocab_size=target_probs.shape[1],

1121 + BLOCK_SIZE=8192,

1122 + NO_DRAFT_PROBS=draft_probs is None,

1123 + USE_FP64_GUMBEL=inv_q.dtype == torch.float64,

1124 + )

vllm/v1/spec_decode/dflash.py

+12 -6

@@ -13,7 +13,7 @@ from vllm.logger import init_logger

13 from vllm.v1.attention.backend import CommonAttentionMetadata

14 from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer

15 from vllm.v1.spec_decode.utils import (

16 - copy_and_expand_dflash_inputs_kernel,

16 + _copy_and_expand_dflash_inputs,

17 next_power_of_2,

18 )

@@ -79,6 +79,13 @@ class DFlashProposer(SpecDecodeBaseProposer):

79 self.dflash_causal = not dflash_has_any_non_causal(

80 self.draft_model_config.hf_config

81 )

82 + _copy_and_expand_dflash_inputs.register_warmup(

83 + block_table_stride=(self.max_model_len + self.block_size - 1)

84 + // self.block_size,

85 + parallel_drafting_token_id=self.parallel_drafting_token_id,

86 + block_size=self.block_size,

87 + num_speculative_tokens=self.num_speculative_tokens,

88 + )

90 @override

91 def _create_draft_vllm_config(self) -> VllmConfig:

@@ -137,11 +144,9 @@ class DFlashProposer(SpecDecodeBaseProposer):

144 max_ctx_per_req = cad.max_query_len

145 max_tokens_per_req = max_ctx_per_req + num_query_per_req

146 BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req))

140 - num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE

141 - grid = (batch_size, num_blocks)

147

148 has_num_rejected = num_rejected_tokens_gpu is not None

144 - copy_and_expand_dflash_inputs_kernel[grid](

149 + _copy_and_expand_dflash_inputs(

150 # Inputs

151 next_token_ids_ptr=next_token_ids,

152 target_positions_ptr=target_positions,

@@ -166,8 +171,9 @@ class DFlashProposer(SpecDecodeBaseProposer):

171 num_query_per_req=num_query_per_req,

172 num_speculative_tokens=self.num_speculative_tokens,

173 total_input_tokens=num_context,

169 - BLOCK_SIZE=BLOCK_SIZE,

170 - HAS_NUM_REJECTED=has_num_rejected,

174 + max_tokens_per_req=max_tokens_per_req,

175 + triton_block_size=BLOCK_SIZE,

176 + has_num_rejected=has_num_rejected,

177 )

178

179 query_slot_mapping = self._slot_mapping_buffer[:num_query_total]

vllm/v1/spec_decode/llm_base_proposer.py

+26 -9

@@ -52,10 +52,11 @@ from vllm.v1.sample.sampler import _SAMPLING_EPS

52 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

53 from vllm.v1.spec_decode.utils import (

54 PADDING_SLOT_ID,

55 + _copy_and_expand_eagle_inputs,

56 + _eagle_prepare_inputs_padded,

57 + _eagle_prepare_next_token_padded,

58 + _eagle_step_slot_mapping_metadata,

59 compute_new_slot_mapping,

56 - copy_and_expand_eagle_inputs_kernel,

57 - eagle_prepare_inputs_padded_kernel,

58 - eagle_prepare_next_token_padded_kernel,

60 eagle_step_update_slot_mapping_and_metadata,

61 extend_all_queries_by_N,

62 next_power_of_2,

@@ -140,6 +141,23 @@ class SpecDecodeBaseProposer:

141

142 self.max_batch_size = vllm_config.scheduler_config.max_num_seqs

143 self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

144 + _eagle_step_slot_mapping_metadata.register_warmup(

145 + max_model_len=self.max_model_len,

146 + max_batch_size=self.max_batch_size,

147 + )

148 + _eagle_prepare_inputs_padded.register_warmup(max_batch_size=self.max_batch_size)

149 + _eagle_prepare_next_token_padded.register_warmup(

150 + vocab_size=vllm_config.model_config.get_vocab_size(),

151 + num_sampled_tokens_per_req=self.num_speculative_tokens + 1,

152 + max_batch_size=self.max_batch_size,

153 + )

154 + if self.needs_extra_input_slots:

155 + _copy_and_expand_eagle_inputs.register_warmup(

156 + parallel_drafting_token_id=self.parallel_drafting_token_id,

157 + num_padding_slots_per_request=self.extra_slots_per_request,

158 + shift_input_ids=self.pass_hidden_states_to_model,

159 + max_num_tokens=self.max_num_tokens,

160 + )

161 self.token_arange_np = np.arange(self.max_num_tokens, dtype=np.int32)

162

163 # Can be specialized by methods like DFlash to reduce the limit

@@ -872,13 +890,12 @@ class SpecDecodeBaseProposer:

890 )

891

892 # Kernel grid: one program per request (row)

875 - grid = (batch_size, num_blocks)

893 query_start_loc = cad.query_start_loc

894 query_end_loc = cad.query_start_loc[1:] - 1

895 if num_rejected_tokens_gpu is not None:

896 query_end_loc = query_end_loc - num_rejected_tokens_gpu

897

881 - copy_and_expand_eagle_inputs_kernel[grid](

898 + _copy_and_expand_eagle_inputs(

899 # (Padded) Inputs from the target model

900 target_token_ids_ptr=target_token_ids,

901 target_positions_ptr=target_positions,

@@ -901,6 +918,8 @@ class SpecDecodeBaseProposer:

918 num_padding_slots_per_request=self.extra_slots_per_request,

919 shift_input_ids=self.pass_hidden_states_to_model,

920 BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,

921 + batch_size=batch_size,

922 + num_blocks=num_blocks,

923 )

924 if self.pass_hidden_states_to_model:

925 assert self.parallel_drafting_hidden_state_tensor is not None

@@ -1065,11 +1084,10 @@ class SpecDecodeBaseProposer:

1084 valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)

1085

1086 # Kernel grid: one program per request (row)

1068 - grid = (batch_size,)

1087

1088 # Find the next power of 2 for block sizes

1089 BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens)

1072 - eagle_prepare_next_token_padded_kernel[grid](

1090 + _eagle_prepare_next_token_padded(

1091 sampled_token_ids,

1092 discard_request_mask,

1093 backup_tokens_gpu,

@@ -1108,8 +1126,7 @@ class SpecDecodeBaseProposer:

1126 (num_reqs,), dtype=torch.int32, device=device

1127 )

1128

1111 - grid = (num_reqs,)

1112 - eagle_prepare_inputs_padded_kernel[grid](

1129 + _eagle_prepare_inputs_padded(

1130 spec_decode_metadata.cu_num_draft_tokens,

1131 valid_sampled_tokens_count,

1132 common_attn_metadata.query_start_loc,

vllm/v1/spec_decode/utils.py

+257 -1

@@ -1,7 +1,18 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Any

4 +

5 import torch

7 +from vllm.model_executor.warmup.jit_warmup import (

8 + WarmupChoices,

9 + WarmupIntRange,

10 +)

11 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

12 + DispatchSpec,

13 + TritonWarmupTensor,

14 + triton_kernel_dispatcher_with_warmup,

15 +)

16 from vllm.platforms import current_platform

17 from vllm.triton_utils import tl, triton

18 from vllm.v1.attention.backends.utils import (

@@ -85,6 +96,49 @@ def eagle_step_slot_mapping_metadata_kernel(

96 tl.store(seq_lens_ptr + req_idx, new_seq_len)

99 +def _eagle_step_warmup_inputs(*, max_model_len: int, max_batch_size: int):

100 + block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2)

101 + n_blocks_per_req = triton.cdiv(max_model_len, block_size)

102 + return dict(

103 + positions=TritonWarmupTensor(torch.int64),

104 + block_table=TritonWarmupTensor(

105 + torch.int32, shape=(1, n_blocks_per_req), strides=(n_blocks_per_req, 1)

106 + ),

107 + block_table_stride=n_blocks_per_req,

108 + seq_lens=TritonWarmupTensor(torch.int32),

109 + out_clamped_positions=TritonWarmupTensor(torch.int64),

110 + out_slot_mapping=TritonWarmupTensor(torch.int64),

111 + block_size=block_size,

112 + max_model_len=max_model_len,

113 + n_blocks_per_req=n_blocks_per_req,

114 + PAD_ID=PADDING_SLOT_ID,

115 + batch_size=WarmupIntRange(1, max_batch_size + 1),

116 + input_batch_size=1,

117 + )

118 +

119 +

120 +@triton_kernel_dispatcher_with_warmup(

121 + kernel=eagle_step_slot_mapping_metadata_kernel,

122 + warmup_inputs=_eagle_step_warmup_inputs,

123 +)

124 +def _eagle_step_slot_mapping_metadata(

125 + positions: torch.Tensor,

126 + block_table: torch.Tensor,

127 + block_table_stride: int,

128 + seq_lens: torch.Tensor,

129 + out_clamped_positions: torch.Tensor,

130 + out_slot_mapping: torch.Tensor,

131 + *,

132 + block_size: int,

133 + max_model_len: int,

134 + n_blocks_per_req: int,

135 + PAD_ID: int,

136 + batch_size: int,

137 + input_batch_size: int,

138 +) -> DispatchSpec:

139 + return (input_batch_size,), {}

140 +

141 +

142 def eagle_step_update_slot_mapping_and_metadata(

143 positions_1d: torch.Tensor,

144 block_table_tensor: torch.Tensor,

@@ -118,7 +172,7 @@ def eagle_step_update_slot_mapping_and_metadata(

172 input_batch_size = batch_size

173

174 n_blocks_per_req = block_table_tensor.shape[1]

121 - eagle_step_slot_mapping_metadata_kernel[(input_batch_size,)](

175 + _eagle_step_slot_mapping_metadata(

176 positions_1d,

177 block_table_tensor,

178 block_table_tensor.stride(0),

@@ -130,6 +184,7 @@ def eagle_step_update_slot_mapping_and_metadata(

184 n_blocks_per_req=n_blocks_per_req,

185 PAD_ID=PADDING_SLOT_ID,

186 batch_size=batch_size,

187 + input_batch_size=input_batch_size,

188 )

189

190

@@ -175,6 +230,33 @@ def eagle_prepare_inputs_padded_kernel(

230 tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens)

231

232

233 +def _eagle_prepare_inputs_warmup_inputs(*, max_batch_size: int):

234 + int32 = TritonWarmupTensor(torch.int32)

235 + return dict(

236 + cu_num_draft_tokens=int32,

237 + valid_sampled_tokens_count=int32,

238 + query_start_loc_gpu=int32,

239 + token_indices_to_sample=int32,

240 + num_rejected_tokens_gpu=int32,

241 + num_reqs=WarmupIntRange(1, max_batch_size + 1),

242 + )

243 +

244 +

245 +@triton_kernel_dispatcher_with_warmup(

246 + kernel=eagle_prepare_inputs_padded_kernel,

247 + warmup_inputs=_eagle_prepare_inputs_warmup_inputs,

248 +)

249 +def _eagle_prepare_inputs_padded(

250 + cu_num_draft_tokens: torch.Tensor,

251 + valid_sampled_tokens_count: torch.Tensor,

252 + query_start_loc_gpu: torch.Tensor,

253 + token_indices_to_sample: torch.Tensor,

254 + num_rejected_tokens_gpu: torch.Tensor,

255 + num_reqs: int,

256 +) -> DispatchSpec:

257 + return (num_reqs,), {}

258 +

259 +

260 @triton.jit

261 def eagle_prepare_next_token_padded_kernel(

262 sampled_token_ids_ptr, # [num_reqs, num_sampled_tokens_per_req]

@@ -238,6 +320,49 @@ def eagle_prepare_next_token_padded_kernel(

320 tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)

321

322

323 +def _eagle_prepare_next_token_warmup_inputs(

324 + *, vocab_size: int, num_sampled_tokens_per_req: int, max_batch_size: int

325 +):

326 + int32 = TritonWarmupTensor(torch.int32)

327 + sampled_width: Any = WarmupIntRange(1, num_sampled_tokens_per_req + 1)

328 + return dict(

329 + sampled_token_ids=TritonWarmupTensor(

330 + torch.int32,

331 + shape=(1, sampled_width),

332 + strides=(sampled_width, 1),

333 + ),

334 + discard_request_mask=TritonWarmupTensor(torch.bool),

335 + backup_next_token_ids=int32,

336 + next_token_ids=int32,

337 + valid_sampled_tokens_count=int32,

338 + vocab_size=vocab_size,

339 + num_sampled_tokens_per_req=sampled_width,

340 + num_reqs=WarmupIntRange(1, max_batch_size + 1),

341 + stride_sampled_token_ids=sampled_width,

342 + BLOCK_SIZE_TOKENS=next_power_of_2(sampled_width),

343 + )

344 +

345 +

346 +@triton_kernel_dispatcher_with_warmup(

347 + kernel=eagle_prepare_next_token_padded_kernel,

348 + warmup_inputs=_eagle_prepare_next_token_warmup_inputs,

349 +)

350 +def _eagle_prepare_next_token_padded(

351 + sampled_token_ids: torch.Tensor,

352 + discard_request_mask: torch.Tensor,

353 + backup_next_token_ids: torch.Tensor,

354 + next_token_ids: torch.Tensor,

355 + valid_sampled_tokens_count: torch.Tensor,

356 + vocab_size: int,

357 + num_sampled_tokens_per_req: int,

358 + num_reqs: int,

359 + stride_sampled_token_ids: int,

360 + *,

361 + BLOCK_SIZE_TOKENS: int,

362 +) -> DispatchSpec:

363 + return (num_reqs,), {}

364 +

365 +

366 def compute_new_slot_mapping(

367 cad: CommonAttentionMetadata,

368 new_positions: torch.Tensor,

@@ -454,6 +579,68 @@ def copy_and_expand_eagle_inputs_kernel(

579 )

580

581

582 +def _copy_and_expand_eagle_warmup_inputs(

583 + *,

584 + parallel_drafting_token_id: int,

585 + num_padding_slots_per_request: int,

586 + shift_input_ids: bool,

587 + max_num_tokens: int,

588 +):

589 + int32 = TritonWarmupTensor(torch.int32)

590 + int64 = TritonWarmupTensor(torch.int64)

591 + boolean = TritonWarmupTensor(torch.bool)

592 + return dict(

593 + target_token_ids_ptr=int32,

594 + target_positions_ptr=int64,

595 + next_token_ids_ptr=int32,

596 + out_input_ids_ptr=int32,

597 + out_positions_ptr=int64,

598 + out_is_rejected_token_mask_ptr=boolean,

599 + out_is_masked_token_mask_ptr=boolean,

600 + out_new_token_indices_ptr=int32,

601 + out_hidden_state_mapping_ptr=int32,

602 + query_start_loc_ptr=int32,

603 + query_end_loc_ptr=int32,

604 + padding_token_id=0,

605 + parallel_drafting_token_id=parallel_drafting_token_id,

606 + total_input_tokens=WarmupIntRange(1, max_num_tokens + 1),

607 + num_padding_slots_per_request=num_padding_slots_per_request,

608 + shift_input_ids=shift_input_ids,

609 + BLOCK_SIZE_TOKENS=WarmupIntRange(1, 257, advance=lambda value: value * 2),

610 + batch_size=1,

611 + num_blocks=1,

612 + )

613 +

614 +

615 +@triton_kernel_dispatcher_with_warmup(

616 + kernel=copy_and_expand_eagle_inputs_kernel,

617 + warmup_inputs=_copy_and_expand_eagle_warmup_inputs,

618 +)

619 +def _copy_and_expand_eagle_inputs(

620 + target_token_ids_ptr: torch.Tensor,

621 + target_positions_ptr: torch.Tensor,

622 + next_token_ids_ptr: torch.Tensor,

623 + out_input_ids_ptr: torch.Tensor,

624 + out_positions_ptr: torch.Tensor,

625 + out_is_rejected_token_mask_ptr: torch.Tensor,

626 + out_is_masked_token_mask_ptr: torch.Tensor,

627 + out_new_token_indices_ptr: torch.Tensor,

628 + out_hidden_state_mapping_ptr: torch.Tensor,

629 + query_start_loc_ptr: torch.Tensor,

630 + query_end_loc_ptr: torch.Tensor,

631 + padding_token_id: int,

632 + parallel_drafting_token_id: int,

633 + total_input_tokens: int,

634 + num_padding_slots_per_request: int,

635 + shift_input_ids: bool,

636 + *,

637 + BLOCK_SIZE_TOKENS: int,

638 + batch_size: int,

639 + num_blocks: int,

640 +) -> DispatchSpec:

641 + return (batch_size, num_blocks), {}

642 +

643 +

644 @triton.jit

645 def copy_and_expand_dflash_inputs_kernel(

646 # Inputs

@@ -562,6 +749,75 @@ def copy_and_expand_dflash_inputs_kernel(

749 )

750

751

752 +def _copy_and_expand_dflash_warmup_inputs(

753 + *,

754 + block_table_stride: int,

755 + parallel_drafting_token_id: int,

756 + block_size: int,

757 + num_speculative_tokens: int,

758 +):

759 + int32 = TritonWarmupTensor(torch.int32)

760 + int64 = TritonWarmupTensor(torch.int64)

761 + triton_block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2)

762 + has_num_rejected = WarmupChoices(False, True)

763 + return dict(

764 + next_token_ids_ptr=int32,

765 + target_positions_ptr=int64,

766 + out_input_ids_ptr=int32,

767 + out_context_positions_ptr=int64,

768 + out_query_positions_ptr=int64,

769 + out_context_slot_mapping_ptr=int64,

770 + out_query_slot_mapping_ptr=int64,

771 + out_token_indices_ptr=int32,

772 + block_table_ptr=int32,

773 + query_start_loc_ptr=int32,

774 + num_rejected_tokens_ptr=(int32 if has_num_rejected else 0),

775 + block_table_stride=block_table_stride,

776 + parallel_drafting_token_id=parallel_drafting_token_id,

777 + block_size=block_size,

778 + num_query_per_req=1 + num_speculative_tokens,

779 + num_speculative_tokens=num_speculative_tokens,

780 + total_input_tokens=WarmupChoices(1, 2, 16),

781 + max_tokens_per_req=triton_block_size,

782 + triton_block_size=triton_block_size,

783 + has_num_rejected=has_num_rejected,

784 + )

785 +

786 +

787 +@triton_kernel_dispatcher_with_warmup(

788 + kernel=copy_and_expand_dflash_inputs_kernel,

789 + warmup_inputs=_copy_and_expand_dflash_warmup_inputs,

790 +)

791 +def _copy_and_expand_dflash_inputs(

792 + *,

793 + next_token_ids_ptr: torch.Tensor,

794 + target_positions_ptr: torch.Tensor,

795 + out_input_ids_ptr: torch.Tensor,

796 + out_context_positions_ptr: torch.Tensor,

797 + out_query_positions_ptr: torch.Tensor,

798 + out_context_slot_mapping_ptr: torch.Tensor,

799 + out_query_slot_mapping_ptr: torch.Tensor,

800 + out_token_indices_ptr: torch.Tensor,

801 + block_table_ptr: torch.Tensor,

802 + query_start_loc_ptr: torch.Tensor,

803 + num_rejected_tokens_ptr: torch.Tensor | int,

804 + block_table_stride: int,

805 + parallel_drafting_token_id: int,

806 + block_size: int,

807 + num_query_per_req: int,

808 + num_speculative_tokens: int,

809 + total_input_tokens: int,

810 + max_tokens_per_req: int,

811 + triton_block_size: int,

812 + has_num_rejected: bool,

813 +) -> DispatchSpec:

814 + num_reqs = query_start_loc_ptr.shape[0] - 1

815 + num_blocks = triton.cdiv(max_tokens_per_req, triton_block_size)

816 + return (num_reqs, num_blocks), dict(

817 + BLOCK_SIZE=triton_block_size,

818 + )

819 +

820 +

821 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

822 def update_num_computed_tokens_for_batch_change(

823 num_computed_tokens: torch.Tensor,

vllm/v1/worker/cpu_model_runner.py

+3 -4

@@ -80,16 +80,15 @@ class CPUModelRunner(GPUModelRunner):

81 # Speculative decoding fallbacks

82 import vllm.v1.sample.rejection_sampler

83 - import vllm.v1.spec_decode.llm_base_proposer

83 import vllm.v1.spec_decode.utils as spec_decode_utils

86 - vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = (

85 + spec_decode_utils._eagle_prepare_inputs_padded.kernel = (

86 cpu_tl.eagle_prepare_inputs_padded_kernel

87 )

89 - vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_next_token_padded_kernel = (

88 + spec_decode_utils._eagle_prepare_next_token_padded.kernel = (

89 cpu_tl.eagle_prepare_next_token_padded_kernel

90 )

92 - vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = (

91 + spec_decode_utils._copy_and_expand_eagle_inputs.kernel = (

92 cpu_tl.copy_and_expand_eagle_inputs_kernel

93 )

94 spec_decode_utils.copy_and_expand_dflash_inputs_kernel = (

vllm/v1/worker/gpu/model_runner.py

+53 -44

@@ -188,6 +188,7 @@ logger = init_logger(__name__)

188

189

190 class GPUModelRunner(LoRAModelRunnerMixin):

191 + @JitWarmupRegistry.capture

192 def __init__(self, vllm_config: VllmConfig, device: torch.device):

193 self.vllm_config = vllm_config

194 self.model_config = vllm_config.model_config

@@ -202,7 +203,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

203 self.speculative_config is not None and self.speculative_config.use_dspark()

204 )

205 self.observability_config = vllm_config.observability_config

205 - self.jit_warmup_registry = JitWarmupRegistry(vllm_config)

206

207 self.device = device

208 self.dtype = self.model_config.dtype

@@ -399,7 +399,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

399 logger.info_once("Loading model from scratch...")

400

401 # Capture warmup providers selected while constructing the model.

402 - with self.jit_warmup_registry.activate():

402 + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]

403 self.model = model_loader.load_model(

404 vllm_config=self.vllm_config,

405 model_config=self.vllm_config.model_config,

@@ -461,47 +461,50 @@ class GPUModelRunner(LoRAModelRunnerMixin):

461

462 # Initialize samplers. Model states may override via custom_sampler().

463 if self.is_last_pp_rank and not self.is_pooling_model:

464 - sampler_kwargs: dict[str, Any] = {

465 - "max_num_reqs": self.max_num_reqs,

466 - "vocab_size": self.vocab_size,

467 - "device": self.device,

468 - "req_states": self.req_states,

469 - "logprobs_mode": self.model_config.logprobs_mode,

470 - "num_speculative_tokens": self.decode_query_len,

471 - "use_fp64_gumbel": self.model_config.use_fp64_gumbel,

472 - "enable_trace_replay": self.model_config.enable_trace_replay,

473 - "reasoning_config": self.vllm_config.reasoning_config,

474 - "return_sampling_mask": self.model_config.return_sampling_mask,

475 - }

476 - if self.vllm_config.watermark_config is None:

477 - self.sampler = Sampler(**sampler_kwargs)

478 - else:

479 - wm_config = self.vllm_config.watermark_config

480 - watermarker = create_watermarker(wm_config)

481 - if self.speculative_config is not None:

482 - watermarker = create_speculative_target_watermarker(watermarker)

483 - self.sampler = GPUWatermarkSampler(

484 - watermarker,

485 - deduplicate_contexts=wm_config.deduplicate_contexts,

486 - deduplicate_contexts_max_history=(

487 - wm_config.deduplicate_contexts_max_history

488 - ),

489 - **sampler_kwargs,

490 - )

491 - custom = self.model_state.custom_sampler(self.sampler)

492 -

493 - if custom:

494 - self.vllm_config._check_watermarking_unsupported(custom_sampler=True)

495 - self.sampler, self.rejection_sampler = custom

496 - elif self.speculative_config is not None:

497 - self.rejection_sampler = RejectionSampler(

498 - self.sampler,

499 - self.speculative_config,

500 - self.device,

501 - watermark_key=speculative_target_watermark_key(

502 - self.vllm_config.watermark_config

503 - ),

504 - )

464 + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]

465 + sampler_kwargs: dict[str, Any] = {

466 + "max_num_reqs": self.max_num_reqs,

467 + "vocab_size": self.vocab_size,

468 + "device": self.device,

469 + "req_states": self.req_states,

470 + "logprobs_mode": self.model_config.logprobs_mode,

471 + "num_speculative_tokens": self.decode_query_len,

472 + "use_fp64_gumbel": self.model_config.use_fp64_gumbel,

473 + "enable_trace_replay": self.model_config.enable_trace_replay,

474 + "reasoning_config": self.vllm_config.reasoning_config,

475 + "return_sampling_mask": self.model_config.return_sampling_mask,

476 + }

477 + if self.vllm_config.watermark_config is None:

478 + self.sampler = Sampler(**sampler_kwargs)

479 + else:

480 + wm_config = self.vllm_config.watermark_config

481 + watermarker = create_watermarker(wm_config)

482 + if self.speculative_config is not None:

483 + watermarker = create_speculative_target_watermarker(watermarker)

484 + self.sampler = GPUWatermarkSampler(

485 + watermarker,

486 + deduplicate_contexts=wm_config.deduplicate_contexts,

487 + deduplicate_contexts_max_history=(

488 + wm_config.deduplicate_contexts_max_history

489 + ),

490 + **sampler_kwargs,

491 + )

492 + custom = self.model_state.custom_sampler(self.sampler)

493 +

494 + if custom:

495 + self.vllm_config._check_watermarking_unsupported(

496 + custom_sampler=True

497 + )

498 + self.sampler, self.rejection_sampler = custom

499 + elif self.speculative_config is not None:

500 + self.rejection_sampler = RejectionSampler(

501 + self.sampler,

502 + self.speculative_config,

503 + self.device,

504 + watermark_key=speculative_target_watermark_key(

505 + self.vllm_config.watermark_config

506 + ),

507 + )

508 self.prompt_logprobs_worker = PromptLogprobsWorker(

509 self.max_num_reqs,

510 logprobs_mode=self.model_config.logprobs_mode,

@@ -738,7 +741,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

741 self.speculator.init_cudagraph_manager(cudagraph_mode)

742

743 # Capture warmup providers that depend on allocated KV-cache strides.

741 - with self.jit_warmup_registry.activate():

744 + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]

745 kv_caches_dict = init_kv_cache(

746 self.compilation_config.static_forward_context,

747 self.kv_cache_config,

@@ -924,6 +927,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

927 # during actual execution.

928 assert self.sampler is not None

929 self.sampler(logits, dummy_input_batch)

930 + if self.rejection_sampler is not None:

931 + # Warm head-dtype logits; scheduler warmup covers processed FP32 logits.

932 + assert self.speculator is not None

933 + self.rejection_sampler(

934 + logits, dummy_input_batch, self.speculator.draft_logits

935 + )

936

937 @torch.inference_mode()

938 def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:

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

+5 -1

@@ -4,7 +4,10 @@ import numpy as np

4 import torch

6 from vllm.sampling_params import SamplingParams

7 -from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p

7 +from vllm.v1.sample.ops.topk_topp_sampler import (

8 + apply_top_k_top_p,

9 + register_top_k_top_p_warmups,

10 +)

11 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor

12 from vllm.v1.worker.gpu.sample.gumbel import apply_temperature

13 from vllm.v1.worker.gpu.sample.min_p import apply_min_p

@@ -18,6 +21,7 @@ class SamplingStates:

21 def __init__(self, max_num_reqs: int, vocab_size: int):

22 self.max_num_reqs = max_num_reqs

23 self.vocab_size = vocab_size

24 + register_top_k_top_p_warmups()

26 self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32)

27 self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32)

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

+92 -39

@@ -1,8 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import copy

4 +from types import SimpleNamespace

5 from typing import Any

7 +import numpy as np

8 import torch

9 import torch.nn as nn

@@ -10,6 +12,15 @@ from vllm.config import VllmConfig, replace

12 from vllm.config.compilation import CUDAGraphMode

13 from vllm.forward_context import BatchDescriptor, set_forward_context

14 from vllm.logger import init_logger

15 +from vllm.model_executor.warmup.jit_warmup import (

16 + WarmupChoices,

17 + WarmupIntRange,

18 +)

19 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

20 + DispatchSpec,

21 + TritonWarmupTensor,

22 + triton_kernel_dispatcher_with_warmup,

23 +)

24 from vllm.triton_utils import tl, triton

25 from vllm.v1.attention.backend import AttentionCGSupport

26 from vllm.v1.attention.backends.utils import PAD_SLOT_ID

@@ -108,6 +119,7 @@ class DFlashSpeculator(DraftModelSpeculator):

119

120 self.query_cudagraph_manager: DFlashCudaGraphManager | None = None

121 self.draft_kv_cache_group_id: int = -1

122 + prepare_dflash_inputs.register_warmup(speculator=self)

123

124 @property

125 def attn_vllm_config(self) -> VllmConfig:

@@ -682,6 +694,64 @@ def _prepare_dflash_inputs_kernel(

694 tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)

695

696

697 +def _prepare_dflash_warmup_inputs(*, speculator: DFlashSpeculator) -> dict[str, Any]:

698 + int32 = TritonWarmupTensor(torch.int32)

699 + int64 = TritonWarmupTensor(torch.int64)

700 + float32 = TritonWarmupTensor(torch.float32)

701 + input_buffers = SimpleNamespace(

702 + input_ids=int32, positions=int64, query_start_loc=int32, seq_lens=int32

703 + )

704 + gid: Any = WarmupChoices(*speculator.draft_kv_cache_group_ids)

705 + max_target_query_len: Any = WarmupIntRange(1, 257, advance=lambda value: value * 2)

706 + block_table_stride = int(speculator.block_tables.input_block_tables[gid].stride(0))

707 + block_table = TritonWarmupTensor(

708 + torch.int32,

709 + shape=(1, block_table_stride),

710 + strides=(block_table_stride, 1),

711 + )

712 + input_batch = SimpleNamespace(

713 + num_reqs=1,

714 + num_scheduled_tokens=np.array([max_target_query_len]),

715 + positions=int64,

716 + query_start_loc=int32,

717 + idx_mapping=int64,

718 + )

719 + return dict(

720 + input_buffers=input_buffers,

721 + query_slot_mapping=int64,

722 + context_positions=int64,

723 + context_slot_mapping=int64,

724 + sample_indices=int64,

725 + sample_pos=int64,

726 + sample_idx_mapping=int32,

727 + temperature=float32,

728 + seeds=int64,

729 + input_batch=input_batch,

730 + num_sampled=int32,

731 + num_rejected=int32,

732 + last_sampled=int64,

733 + next_prefill_tokens=int32,

734 + input_temperature=float32,

735 + input_seeds=int64,

736 + block_table=block_table,

737 + block_size=speculator.block_tables.kernel_block_sizes[gid],

738 + cp_rank=speculator.block_tables.cp_rank,

739 + cp_size=speculator.block_tables.cp_size,

740 + cp_interleave=speculator.block_tables.cp_interleave,

741 + parallel_drafting_token_id=speculator.parallel_drafting_token_id,

742 + num_query_per_req=speculator.num_query_per_req,

743 + num_speculative_steps=speculator.num_speculative_steps,

744 + max_num_reqs=speculator.max_num_reqs,

745 + max_num_tokens=speculator.max_num_tokens,

746 + max_model_len=speculator.max_model_len,

747 + sample_from_anchor=speculator.sample_from_anchor,

748 + )

749 +

750 +

751 +@triton_kernel_dispatcher_with_warmup(

752 + kernel=_prepare_dflash_inputs_kernel,

753 + warmup_inputs=_prepare_dflash_warmup_inputs,

754 +)

755 def prepare_dflash_inputs(

756 input_buffers: InputBuffers,

757 query_slot_mapping: torch.Tensor,

@@ -718,50 +788,33 @@ def prepare_dflash_inputs(

788 max_num_tokens: int,

789 max_model_len: int,

790 sample_from_anchor: bool = False,

721 -) -> None:

791 +) -> DispatchSpec:

792 num_reqs = input_batch.num_reqs

793 assert num_reqs > 0

794 # Cover the longest possible per-request span (ctx + query). Use the max

795 # per-request query length, not the total token count across the batch.

796 max_target_query_len = int(input_batch.num_scheduled_tokens.max())

797 max_tokens_per_req = max_target_query_len + num_query_per_req

728 - BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))

729 - num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)

730 - _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](

731 - input_buffers.input_ids,

732 - input_buffers.positions,

733 - input_buffers.query_start_loc,

734 - input_buffers.seq_lens,

735 - query_slot_mapping,

736 - context_positions,

737 - context_slot_mapping,

738 - sample_indices,

739 - sample_pos,

740 - sample_idx_mapping,

741 - temperature,

742 - seeds,

743 - input_batch.positions,

744 - input_batch.query_start_loc,

745 - input_batch.idx_mapping,

746 - last_sampled,

747 - next_prefill_tokens,

748 - num_sampled,

749 - num_rejected,

750 - input_temperature,

751 - input_seeds,

752 - block_table,

753 - block_table.stride(0),

754 - parallel_drafting_token_id,

755 - block_size,

756 - num_query_per_req,

757 - num_speculative_steps,

758 - max_num_reqs,

759 - max_num_tokens,

760 - max_model_len,

761 - cp_rank,

762 - SAMPLE_FROM_ANCHOR=sample_from_anchor,

798 + block = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))

799 + num_blocks = triton.cdiv(max_tokens_per_req, block)

800 + return (num_reqs, num_blocks), dict(

801 + out_input_ids_ptr=input_buffers.input_ids,

802 + out_query_positions_ptr=input_buffers.positions,

803 + out_query_start_loc_ptr=input_buffers.query_start_loc,

804 + out_seq_lens_ptr=input_buffers.seq_lens,

805 + out_query_slot_mapping_ptr=query_slot_mapping,

806 + out_context_positions_ptr=context_positions,

807 + out_context_slot_mapping_ptr=context_slot_mapping,

808 + out_sample_indices_ptr=sample_indices,

809 + out_sample_pos_ptr=sample_pos,

810 + out_sample_idx_mapping_ptr=sample_idx_mapping,

811 + out_temperature_ptr=temperature,

812 + out_seeds_ptr=seeds,

813 + target_positions_ptr=input_batch.positions,

814 + target_query_start_loc_ptr=input_batch.query_start_loc,

815 + idx_mapping_ptr=input_batch.idx_mapping,

816 + temperature_ptr=input_temperature,

817 + seeds_ptr=input_seeds,

818 PAD_SLOT_ID=PAD_SLOT_ID,

764 - CP_SIZE=cp_size,

765 - CP_INTERLEAVE=cp_interleave,

766 - BLOCK_SIZE=BLOCK_SIZE,

819 + BLOCK_SIZE=block,

820 )

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

+76 -12

@@ -9,6 +9,11 @@ from vllm.config import VllmConfig

9 from vllm.config.compilation import CUDAGraphMode

10 from vllm.forward_context import BatchDescriptor, set_forward_context

11 from vllm.logger import init_logger

12 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

13 + DispatchSpec,

14 + TritonWarmupTensor,

15 + triton_kernel_dispatcher_with_warmup,

16 +)

17 from vllm.triton_utils import tl, triton

18 from vllm.v1.attention.backends.utils import PAD_SLOT_ID

19 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

@@ -69,6 +74,9 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

75 self.cudagraph_manager: SpeculatorCudaGraphManager | None = None

77 + _shift_input_ids.register_warmup()

78 + _shift_input_embeds.register_warmup(speculator=self)

79 +

80 def load_draft_model(

81 self,

82 target_model: nn.Module,

@@ -1008,6 +1016,32 @@ def _shift_input_ids_kernel(

1016 tl.store(input_ids_ptr + last_token_index, draft_token)

1017

1018

1019 +def _shift_input_ids_warmup_inputs() -> dict[str, Any]:

1020 + return dict(

1021 + input_ids=TritonWarmupTensor(torch.int32),

1022 + idx_mapping=TritonWarmupTensor(torch.int32),

1023 + query_start_loc=TritonWarmupTensor(torch.int32),

1024 + last_token_indices=TritonWarmupTensor(torch.int64),

1025 + draft_tokens=TritonWarmupTensor(torch.int64),

1026 + num_reqs=1,

1027 + )

1028 +

1029 +

1030 +@triton_kernel_dispatcher_with_warmup(

1031 + kernel=_shift_input_ids_kernel,

1032 + warmup_inputs=_shift_input_ids_warmup_inputs,

1033 +)

1034 +def _shift_input_ids(

1035 + input_ids: torch.Tensor,

1036 + idx_mapping: torch.Tensor,

1037 + query_start_loc: torch.Tensor,

1038 + last_token_indices: torch.Tensor,

1039 + draft_tokens: torch.Tensor,

1040 + num_reqs: int,

1041 +) -> DispatchSpec:

1042 + return (num_reqs,), dict(BLOCK_SIZE=1024)

1043 +

1044 +

1045 @triton.jit

1046 def _shift_input_embeds_kernel(

1047 input_embeds_ptr,

@@ -1065,6 +1099,44 @@ def _shift_input_embeds_kernel(

1099 )

1100

1101

1102 +def _shift_input_embeds_warmup_inputs(

1103 + *, speculator: MultiModuleMTPSpeculator

1104 +) -> dict[str, Any]:

1105 + hidden = TritonWarmupTensor(

1106 + speculator.dtype,

1107 + shape=(1, speculator.hidden_size),

1108 + strides=(speculator.hidden_size, 1),

1109 + )

1110 + return dict(

1111 + input_embeds=hidden,

1112 + draft_embeds=hidden,

1113 + idx_mapping=TritonWarmupTensor(torch.int32),

1114 + query_start_loc=TritonWarmupTensor(torch.int32),

1115 + last_token_indices=TritonWarmupTensor(torch.int64),

1116 + num_reqs=1,

1117 + )

1118 +

1119 +

1120 +@triton_kernel_dispatcher_with_warmup(

1121 + kernel=_shift_input_embeds_kernel,

1122 + warmup_inputs=_shift_input_embeds_warmup_inputs,

1123 +)

1124 +def _shift_input_embeds(

1125 + input_embeds: torch.Tensor,

1126 + draft_embeds: torch.Tensor,

1127 + idx_mapping: torch.Tensor,

1128 + query_start_loc: torch.Tensor,

1129 + last_token_indices: torch.Tensor,

1130 + num_reqs: int,

1131 +) -> DispatchSpec:

1132 + hidden_size = input_embeds.shape[-1]

1133 + return (num_reqs, triton.cdiv(hidden_size, 256)), dict(

1134 + hidden_size=hidden_size,

1135 + BLOCK_SIZE_Q=16,

1136 + BLOCK_SIZE_H=256,

1137 + )

1138 +

1139 +

1140 def update_draft_inputs(

1141 draft_tokens: torch.Tensor,

1142 draft_embeds: torch.Tensor | None,

@@ -1074,29 +1146,21 @@ def update_draft_inputs(

1146 idx_mapping: torch.Tensor,

1147 num_reqs: int,

1148 ) -> None:

1077 - _shift_input_ids_kernel[(num_reqs,)](

1149 + _shift_input_ids(

1150 input_buffers.input_ids,

1151 idx_mapping,

1152 input_buffers.query_start_loc,

1153 last_token_indices,

1154 draft_tokens,

1083 - BLOCK_SIZE=1024,

1155 + num_reqs=num_reqs,

1156 )

1157 if input_embeds is not None:

1158 assert draft_embeds is not None

1087 - hidden_size = input_embeds.shape[-1]

1088 - hidden_block_size = 256

1089 - _shift_input_embeds_kernel[

1090 - (num_reqs, triton.cdiv(hidden_size, hidden_block_size))

1091 - ](

1159 + _shift_input_embeds(

1160 input_embeds,

1093 - input_embeds.stride(0),

1161 draft_embeds,

1095 - draft_embeds.stride(0),

1162 idx_mapping,

1163 input_buffers.query_start_loc,

1164 last_token_indices,

1099 - hidden_size,

1100 - BLOCK_SIZE_Q=16,

1101 - BLOCK_SIZE_H=hidden_block_size,

1165 + num_reqs=num_reqs,

1166 )

vllm/v1/worker/gpu_model_runner.py

+33 -34

@@ -498,6 +498,7 @@ class ExecuteModelState(NamedTuple):

498 class GPUModelRunner(

499 LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin

500 ):

501 + @JitWarmupRegistry.capture

502 def __init__(

503 self,

504 vllm_config: VllmConfig,

@@ -514,7 +515,6 @@ class GPUModelRunner(

515 self.scheduler_config = vllm_config.scheduler_config

516 self.speculative_config = vllm_config.speculative_config

517 self.observability_config = vllm_config.observability_config

517 - self.jit_warmup_registry = JitWarmupRegistry(vllm_config)

518

519 model_config = self.model_config

520 cache_config = self.cache_config

@@ -750,36 +750,35 @@ class GPUModelRunner(

750 self.parallel_config.cp_kv_cache_interleave_size

751 )

752 # Capture warmup providers registered by the initial placeholder InputBatch

753 - with self.jit_warmup_registry.activate():

754 - self.input_batch = InputBatch(

755 - max_num_reqs=self.max_num_reqs,

756 - # We need to use the encoder length for encoder-decoder

757 - # because of KV cache for cross-attention.

758 - max_model_len=max(self.max_model_len, self.max_encoder_len),

759 - max_num_batched_tokens=self.max_num_tokens,

760 - device=self.device,

761 - vocab_size=self.model_config.get_vocab_size(),

762 - block_sizes=[placeholder_block_size],

763 - kernel_block_sizes=[placeholder_block_size],

764 - max_num_blocks_per_req=[placeholder_max_num_blocks],

765 - num_spec_tokens=self.num_spec_tokens,

766 - logitsprocs=build_logitsprocs(

767 - self.vllm_config,

768 - self.device,

769 - PIN_MEMORY,

770 - self.is_pooling_model,

771 - custom_logitsprocs,

772 - ),

773 - # We currently don't know whether a particular custom logits processor

774 - # uses output token ids so we set this conservatively. Thinking-budget

775 - # tracking is requested dynamically when a budgeted request is in the

776 - # batch.

777 - logitsprocs_need_output_token_ids=bool(custom_logitsprocs),

778 - is_pooling_model=self.is_pooling_model,

779 - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,

780 - reasoning_config=self.vllm_config.reasoning_config,

781 - use_replayssm=self.cache_config.use_replayssm,

782 - )

753 + self.input_batch = InputBatch(

754 + max_num_reqs=self.max_num_reqs,

755 + # We need to use the encoder length for encoder-decoder

756 + # because of KV cache for cross-attention.

757 + max_model_len=max(self.max_model_len, self.max_encoder_len),

758 + max_num_batched_tokens=self.max_num_tokens,

759 + device=self.device,

760 + vocab_size=self.model_config.get_vocab_size(),

761 + block_sizes=[placeholder_block_size],

762 + kernel_block_sizes=[placeholder_block_size],

763 + max_num_blocks_per_req=[placeholder_max_num_blocks],

764 + num_spec_tokens=self.num_spec_tokens,

765 + logitsprocs=build_logitsprocs(

766 + self.vllm_config,

767 + self.device,

768 + PIN_MEMORY,

769 + self.is_pooling_model,

770 + custom_logitsprocs,

771 + ),

772 + # We currently don't know whether a particular custom logits processor

773 + # uses output token ids so we set this conservatively. Thinking-budget

774 + # tracking is requested dynamically when a budgeted request is in the

775 + # batch.

776 + logitsprocs_need_output_token_ids=bool(custom_logitsprocs),

777 + is_pooling_model=self.is_pooling_model,

778 + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,

779 + reasoning_config=self.vllm_config.reasoning_config,

780 + use_replayssm=self.cache_config.use_replayssm,

781 + )

782

783 # Separate cuda stream for overlapping transfer of sampled token ids from

784 # GPU to CPU when async scheduling is enabled.

@@ -5332,7 +5331,7 @@ class GPUModelRunner(

5331 self.load_config.load_format = "dummy"

5332 model_loader = get_model_loader(self.load_config)

5333 # Capture warmup providers selected while constructing the model.

5335 - with self.jit_warmup_registry.activate():

5334 + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]

5335 self.model = model_loader.load_model(

5336 vllm_config=self.vllm_config, model_config=self.model_config

5337 )

@@ -7304,7 +7303,7 @@ class GPUModelRunner(

7303 self.parallel_config.cp_kv_cache_interleave_size

7304 )

7305 # Capture warmup providers registered after final KV-cache geometry is known

7307 - with self.jit_warmup_registry.activate():

7306 + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]

7307 self.input_batch = InputBatch(

7308 max_num_reqs=self.max_num_reqs,

7309 max_model_len=max_model_len,

@@ -7454,7 +7453,7 @@ class GPUModelRunner(

7453 # Reinitialize need to after initialize_attn_backend

7454 self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)

7455 # Capture warmup providers that depend on allocated KV-cache strides.

7457 - with self.jit_warmup_registry.activate():

7456 + with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]

7457 kv_caches = self.initialize_kv_cache_tensors(

7458 kv_cache_config,

7459 kernel_block_sizes,