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

[9/N][warmup][DSv4] Migrate MHC TileLang kernels (#50178)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Codex <codex@openai.com>

Roberto L. Castro committed Sep 13, 2026 at 02:25 UTC e52be1a62d3879b1202f4f355d3c3472b560c6f2
9 files changed +1454 -489
tests/kernels/test_mhc_jit_warmup.py new
+296
@@ -0,0 +1,296 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Validate JIT dispatch against pre-contract behavior."""
4 +
5 +from typing import Any
6 +
7 +import pytest
8 +
9 +from vllm.platforms import current_platform
10 +
11 +if not current_platform.is_cuda_alike():
12 + pytest.skip("NVIDIA dispatch tests require CUDA", allow_module_level=True)
13 +
14 +from vllm.model_executor.kernels.mhc.tilelang_kernels import (
15 + HcHeadFusedTileLangKernel,
16 + HcPrenormGemmTileLangKernel,
17 + MhcFusedTileLangKernel,
18 + MhcPostTileLangKernel,
19 + MhcPreBigFuseTileLangKernel,
20 +)
21 +from vllm.model_executor.warmup import jit_warmup_tilelang_helper
22 +
23 +
24 +@pytest.mark.parametrize(
25 + ("kwargs", "expected"),
26 + [
27 + (
28 + dict(
29 + num_tokens=64,
30 + hc_hidden_size=4096,
31 + hidden_size=2048,
32 + hc_mult=2,
33 + n_out=128,
34 + ),
35 + (2048, 2, 128, 1024, 4, 1, False, 1),
36 + ),
37 + (
38 + dict(
39 + num_tokens=1024,
40 + hc_hidden_size=4096,
41 + hidden_size=2048,
42 + hc_mult=2,
43 + n_out=128,
44 + ),
45 + (2048, 2, 128, 512, 12, 1, True, 2),
46 + ),
47 + (
48 + dict(
49 + num_tokens=64,
50 + hc_hidden_size=4096,
51 + hidden_size=2048,
52 + hc_mult=2,
53 + n_out=128,
54 + n_thr=256,
55 + tile_n=8,
56 + n_splits=4,
57 + ),
58 + (2048, 2, 128, 256, 8, 4, False, 1),
59 + ),
60 + ],
61 +)
62 +def test_hc_prenorm_gemm_dispatch_matches_legacy_runtime_config(
63 + kwargs: dict[str, Any],
64 + expected: tuple[int, int, int, int, int, int, bool, int],
65 +) -> None:
66 + kernel = HcPrenormGemmTileLangKernel()
67 +
68 + assert kernel.dispatch(**kwargs) == kernel.CompileKey(*expected)
69 +
70 +
71 +@pytest.mark.parametrize(
72 + ("is_broadcast", "use_norm_weight", "expected_use_norm", "expected_eps"),
73 + [
74 + (False, False, False, 0.0),
75 + (False, True, True, 1.0e-5),
76 + (True, False, True, 2.0e-5),
77 + ],
78 +)
79 +def test_mhc_pre_big_fuse_dispatch_matches_legacy_runtime_config(
80 + is_broadcast: bool,
81 + use_norm_weight: bool,
82 + expected_use_norm: bool,
83 + expected_eps: float,
84 +) -> None:
85 + kernel = MhcPreBigFuseTileLangKernel()
86 +
87 + assert kernel.dispatch(
88 + hidden_size=4096,
89 + hc_mult=4,
90 + n_splits=2,
91 + is_broadcast=is_broadcast,
92 + use_norm_weight=use_norm_weight,
93 + rms_eps=1.0e-6,
94 + hc_pre_eps=2.0e-6,
95 + hc_sinkhorn_eps=3.0e-6,
96 + hc_post_mult_value=0.5,
97 + sinkhorn_repeat=3,
98 + norm_eps=1.0e-5,
99 + broadcast_norm_eps=2.0e-5,
100 + ) == kernel.CompileKey(
101 + hidden_size=4096,
102 + hc_mult=4,
103 + n_splits=2,
104 + use_norm_weight=expected_use_norm,
105 + is_broadcast=is_broadcast,
106 + rms_eps=1.0e-6,
107 + hc_pre_eps=2.0e-6,
108 + hc_sinkhorn_eps=3.0e-6,
109 + hc_post_mult_value=0.5,
110 + sinkhorn_repeat=3,
111 + norm_eps=expected_eps,
112 + )
113 +
114 +
115 +@pytest.mark.parametrize(
116 + ("num_tokens", "hidden_size", "expected_n_splits", "expected_tile_n"),
117 + [(4, 4096, 8, 2), (4, 8192, 4, 2), (8, 4096, 4, 3)],
118 +)
119 +def test_mhc_fused_dispatch_matches_legacy_runtime_config(
120 + num_tokens: int,
121 + hidden_size: int,
122 + expected_n_splits: int,
123 + expected_tile_n: int,
124 +) -> None:
125 + kernel = MhcFusedTileLangKernel()
126 +
127 + assert kernel.dispatch(
128 + num_tokens=num_tokens,
129 + hidden_size=hidden_size,
130 + hc_mult=4,
131 + ) == kernel.CompileKey(
132 + hidden_size=hidden_size,
133 + hc_mult=4,
134 + n_splits=expected_n_splits,
135 + tile_n=expected_tile_n,
136 + )
137 +
138 +
139 +@pytest.mark.parametrize(
140 + ("kernel", "compile_key"),
141 + [
142 + (
143 + HcPrenormGemmTileLangKernel(),
144 + HcPrenormGemmTileLangKernel.CompileKey(
145 + hidden_size=2048,
146 + hc_mult=2,
147 + n_out=128,
148 + n_thr=1024,
149 + tile_n=4,
150 + n_splits=1,
151 + use_block_m=False,
152 + block_m=1,
153 + ),
154 + ),
155 + (
156 + HcPrenormGemmTileLangKernel(),
157 + HcPrenormGemmTileLangKernel.CompileKey(
158 + hidden_size=2048,
159 + hc_mult=2,
160 + n_out=128,
161 + n_thr=512,
162 + tile_n=12,
163 + n_splits=1,
164 + use_block_m=False,
165 + block_m=1,
166 + ),
167 + ),
168 + (
169 + HcPrenormGemmTileLangKernel(),
170 + HcPrenormGemmTileLangKernel.CompileKey(
171 + hidden_size=2048,
172 + hc_mult=2,
173 + n_out=128,
174 + n_thr=512,
175 + tile_n=12,
176 + n_splits=1,
177 + use_block_m=True,
178 + block_m=2,
179 + ),
180 + ),
181 + (
182 + MhcPreBigFuseTileLangKernel(),
183 + MhcPreBigFuseTileLangKernel.CompileKey(
184 + hidden_size=4096,
185 + hc_mult=4,
186 + n_splits=2,
187 + use_norm_weight=True,
188 + is_broadcast=False,
189 + rms_eps=1.0e-6,
190 + hc_pre_eps=2.0e-6,
191 + hc_sinkhorn_eps=3.0e-6,
192 + hc_post_mult_value=0.5,
193 + sinkhorn_repeat=3,
194 + norm_eps=1.0e-5,
195 + ),
196 + ),
197 + (
198 + MhcPreBigFuseTileLangKernel(),
199 + MhcPreBigFuseTileLangKernel.CompileKey(
200 + hidden_size=4096,
201 + hc_mult=4,
202 + n_splits=1,
203 + use_norm_weight=False,
204 + is_broadcast=False,
205 + rms_eps=1.0e-6,
206 + hc_pre_eps=2.0e-6,
207 + hc_sinkhorn_eps=3.0e-6,
208 + hc_post_mult_value=0.5,
209 + sinkhorn_repeat=3,
210 + norm_eps=0.0,
211 + ),
212 + ),
213 + (
214 + MhcPreBigFuseTileLangKernel(),
215 + MhcPreBigFuseTileLangKernel.CompileKey(
216 + hidden_size=4096,
217 + hc_mult=4,
218 + n_splits=2,
219 + use_norm_weight=True,
220 + is_broadcast=True,
221 + rms_eps=1.0e-6,
222 + hc_pre_eps=2.0e-6,
223 + hc_sinkhorn_eps=3.0e-6,
224 + hc_post_mult_value=0.5,
225 + sinkhorn_repeat=3,
226 + norm_eps=1.0e-5,
227 + ),
228 + ),
229 + (
230 + MhcPostTileLangKernel(),
231 + MhcPostTileLangKernel.CompileKey(hidden_size=4096, hc_mult=4),
232 + ),
233 + (
234 + MhcFusedTileLangKernel(),
235 + MhcFusedTileLangKernel.CompileKey(
236 + hidden_size=4096,
237 + hc_mult=4,
238 + n_splits=4,
239 + tile_n=3,
240 + ),
241 + ),
242 + (
243 + MhcFusedTileLangKernel(),
244 + MhcFusedTileLangKernel.CompileKey(
245 + hidden_size=4096,
246 + hc_mult=4,
247 + n_splits=8,
248 + tile_n=2,
249 + ),
250 + ),
251 + (
252 + MhcFusedTileLangKernel(),
253 + MhcFusedTileLangKernel.CompileKey(
254 + hidden_size=8192,
255 + hc_mult=4,
256 + n_splits=4,
257 + tile_n=2,
258 + ),
259 + ),
260 + (
261 + HcHeadFusedTileLangKernel(),
262 + HcHeadFusedTileLangKernel.CompileKey(
263 + hidden_size=4096,
264 + hc_mult=4,
265 + rms_eps=1.0e-6,
266 + hc_eps=2.0e-6,
267 + ),
268 + ),
269 + ],
270 +)
271 +def test_tilelang_warmup_inputs_reproduce_compile_key(
272 + monkeypatch: pytest.MonkeyPatch,
273 + kernel: Any,
274 + compile_key: Any,
275 +) -> None:
276 + compiled: list[Any] = []
277 + single_kernel = isinstance(
278 + kernel,
279 + (MhcPostTileLangKernel, MhcFusedTileLangKernel, HcHeadFusedTileLangKernel),
280 + )
281 + expected_kernel = kernel.kernel() if single_kernel else kernel.kernel(compile_key)
282 + if single_kernel:
283 + monkeypatch.setattr(
284 + kernel,
285 + "dispatch",
286 + lambda **kwargs: pytest.fail("single-kernel launch must not dispatch"),
287 + )
288 + monkeypatch.setattr(
289 + jit_warmup_tilelang_helper,
290 + "compile_tilelang",
291 + lambda jit_impl, *args, **kwargs: compiled.append(jit_impl),
292 + )
293 +
294 + kernel.compile(compile_key)
295 +
296 + assert compiled == [expected_kernel]
tests/kernels/test_mhc_kernels.py
+11 -2
@@ -10,10 +10,12 @@ import torch.nn.functional as F
10 import vllm.model_executor.kernels.mhc # noqa: F401
11 import vllm.model_executor.layers.mhc as mhc_layers
12 from vllm.model_executor.kernels.mhc.tilelang import (
13 - _tilelang_hc_prenorm_gemm,
13 _torch_hc_prenorm_gemm,
14 mhc_pre_delayed_tilelang,
15 )
16 +from vllm.model_executor.kernels.mhc.tilelang_kernels import (
17 + _HC_PRENORM_GEMM_TILELANG_KERNEL,
18 +)
19 from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch
20 from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
21 from vllm.model_executor.layers.mhc import (
@@ -485,7 +487,14 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
487 sqrsum = torch.empty_like(sqrsum_ref)
488
489 _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref)
488 - _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)
490 + _HC_PRENORM_GEMM_TILELANG_KERNEL(
491 + x,
492 + fn,
493 + out,
494 + sqrsum,
495 + hidden_size,
496 + hc_mult,
497 + )
498
499 torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)
500 torch.testing.assert_close(sqrsum, sqrsum_ref, atol=1e-2, rtol=1e-6)
vllm/model_executor/kernels/mhc/tilelang.py
+117 -264
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 +from vllm.utils.math_utils import cdiv
6 from vllm.utils.torch_utils import direct_register_custom_op
7
8
@@ -18,73 +19,56 @@ def _torch_hc_prenorm_gemm(
19 sqrsum[0].copy_(x_float.square().sum(dim=-1))
20
21
21 -def _tilelang_hc_prenorm_gemm(
22 +def _hc_prenorm_gemm_outputs(
23 x: torch.Tensor,
24 fn: torch.Tensor,
24 - out: torch.Tensor,
25 - sqrsum: torch.Tensor,
25 + *,
26 hidden_size: int,
27 hc_mult: int,
28 - tile_n: int = 12,
29 - n_thr: int = 512,
30 - n_splits: int = 1,
31 -) -> None:
28 + use_tilelang_fallback: bool = True,
29 +) -> tuple[torch.Tensor, torch.Tensor]:
30 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
33 - hc_prenorm_gemm_block_m_tilelang,
34 - hc_prenorm_gemm_tilelang,
35 - )
36 -
37 - assert out.shape[0] == n_splits
38 - assert sqrsum.shape[0] == n_splits
39 - assert x.shape[1] == hc_mult * hidden_size
40 - assert x.shape[1] % n_splits == 0
41 - assert (x.shape[1] // n_splits) % n_thr == 0
42 - use_default_config = tile_n == 12 and n_thr == 512
43 - if n_splits == 1 and use_default_config and x.shape[0] >= 1024:
44 - hc_prenorm_gemm_block_m_tilelang(
45 - x,
46 - fn,
47 - out,
48 - sqrsum,
49 - hidden_size,
50 - hc_mult,
51 - fn.shape[0],
52 - n_thr,
53 - tile_n,
54 - 2,
31 + compute_num_split,
32 + )
33 + from vllm.utils.deep_gemm import (
34 + is_deep_gemm_supported,
35 + tf32_hc_prenorm_gemm,
36 + )
37 +
38 + use_deep_gemm = is_deep_gemm_supported() or not use_tilelang_fallback
39 + num_tokens = x.shape[0]
40 + n_splits = (
41 + compute_num_split(64, x.shape[1], cdiv(num_tokens, 64)) if use_deep_gemm else 1
42 + )
43 + out = torch.empty(
44 + n_splits,
45 + num_tokens,
46 + fn.shape[0],
47 + dtype=torch.float32,
48 + device=x.device,
49 + )
50 + sqrsum = torch.empty(
51 + n_splits,
52 + num_tokens,
53 + dtype=torch.float32,
54 + device=x.device,
55 + )
56 + if use_deep_gemm:
57 + tf32_hc_prenorm_gemm(x, fn, out, sqrsum, n_splits)
58 + else:
59 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
60 + _HC_PRENORM_GEMM_TILELANG_KERNEL,
61 )
56 - return
57 - if (
58 - n_splits == 1
59 - and use_default_config
60 - and x.shape[0] < 128
61 - and x.shape[1] % 1024 == 0
62 - ):
63 - hc_prenorm_gemm_tilelang(
62 +
63 + _HC_PRENORM_GEMM_TILELANG_KERNEL(
64 x,
65 fn,
66 out,
67 sqrsum,
68 hidden_size,
69 hc_mult,
70 - fn.shape[0],
71 - 1024,
72 - 4,
73 - n_splits,
70 )
75 - return
76 - hc_prenorm_gemm_tilelang(
77 - x,
78 - fn,
79 - out,
80 - sqrsum,
81 - hidden_size,
82 - hc_mult,
83 - fn.shape[0],
84 - n_thr,
85 - tile_n,
86 - n_splits,
87 - )
71 + return out, sqrsum
72
73
74 def mhc_pre_delayed_tilelang(
@@ -127,6 +111,7 @@ def mhc_pre_delayed_tilelang(
111 (tokens, hc_mult, hc_mult), (tokens, hidden_size), and (tokens, hc_mult).
112 """
113 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
114 + _HC_PRENORM_GEMM_TILELANG_KERNEL,
115 mhc_pre_big_fuse_tilelang,
116 )
117 from vllm.model_executor.kernels.mhc.warmup import (
@@ -186,7 +171,14 @@ def mhc_pre_delayed_tilelang(
171 if use_deep_gemm:
172 tf32_hc_prenorm_gemm(x, fn, mixes, sqrsum, n_splits)
173 else:
189 - _tilelang_hc_prenorm_gemm(x, fn, mixes, sqrsum, input_size, 1)
174 + _HC_PRENORM_GEMM_TILELANG_KERNEL(
175 + x,
176 + fn,
177 + mixes,
178 + sqrsum,
179 + input_size,
180 + 1,
181 + )
182 if norm_weight is not None:
183 assert norm_weight.shape == (hidden_size,)
184 assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous()
@@ -298,7 +290,8 @@ def mhc_pre_tilelang(
290 hc_sinkhorn_eps: sinkhorn epsilon
291 hc_post_mult_value: post-mix multiplier value
292 sinkhorn_repeat: number of sinkhorn iterations
301 - n_splits: split-k factor;
293 + n_splits: retained for the shared MHC operator API; the active GEMM
294 + backend selects its split factor internally.
295 norm_weight: optional RMSNorm weight, shape (hidden_size,), dtype
296 torch.bfloat16. When provided, RMSNorm is fused into the
297 layer_input write path of the big_fuse kernel.
@@ -311,12 +304,8 @@ def mhc_pre_tilelang(
304 layer_input: shape (..., hidden_size), dtype torch.bfloat16
305 """
306 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
314 - compute_num_split,
315 - mhc_pre_big_fuse_tilelang,
316 - mhc_pre_big_fuse_with_norm_tilelang,
307 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
308 )
318 - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
319 - from vllm.utils.math_utils import cdiv
309
310 assert residual.dtype == torch.bfloat16
311 assert fn.dtype == torch.float32
@@ -346,17 +335,6 @@ def mhc_pre_tilelang(
335 residual_flat = residual.view(-1, hc_mult, hidden_size)
336 num_tokens = residual_flat.shape[0]
337
349 - from vllm.utils.deep_gemm import is_deep_gemm_supported
350 -
351 - use_deep_gemm = is_deep_gemm_supported()
352 - if use_deep_gemm:
353 - # these numbers are from deepgemm kernel impl
354 - block_k = 64
355 - block_m = 64
356 - n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
357 - else:
358 - n_splits = 1
359 -
338 post_mix = torch.empty(
339 num_tokens, hc_mult, dtype=torch.float32, device=residual.device
340 )
@@ -367,78 +345,30 @@ def mhc_pre_tilelang(
345 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
346 )
347
370 - gemm_out_mul = torch.empty(
371 - n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
348 + gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
349 + residual_flat.view(num_tokens, hc_mult * hidden_size),
350 + fn,
351 + hidden_size=hidden_size,
352 + hc_mult=hc_mult,
353 )
373 - gemm_out_sqrsum = torch.empty(
374 - n_splits, num_tokens, dtype=torch.float32, device=residual.device
354 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
355 + gemm_out_mul,
356 + gemm_out_sqrsum,
357 + hc_scale,
358 + hc_base,
359 + residual_flat,
360 + post_mix,
361 + comb_mix,
362 + layer_input,
363 + rms_eps,
364 + hc_pre_eps,
365 + hc_sinkhorn_eps,
366 + hc_post_mult_value,
367 + sinkhorn_repeat,
368 + norm_weight=norm_weight,
369 + norm_eps=norm_eps,
370 )
371
377 - residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size)
378 - if use_deep_gemm:
379 - tf32_hc_prenorm_gemm(
380 - residual_2d,
381 - fn,
382 - gemm_out_mul,
383 - gemm_out_sqrsum,
384 - n_splits,
385 - )
386 - else:
387 - _tilelang_hc_prenorm_gemm(
388 - residual_2d,
389 - fn,
390 - gemm_out_mul,
391 - gemm_out_sqrsum,
392 - hidden_size,
393 - hc_mult,
394 - )
395 -
396 - if norm_weight is None:
397 - mhc_pre_big_fuse_tilelang(
398 - gemm_out_mul,
399 - gemm_out_sqrsum,
400 - hc_scale,
401 - hc_base,
402 - residual_flat,
403 - post_mix,
404 - comb_mix,
405 - layer_input,
406 - # Pre-mix buffers are unused in V4 mode.
407 - post_mix,
408 - post_mix,
409 - hidden_size,
410 - rms_eps,
411 - hc_pre_eps,
412 - hc_sinkhorn_eps,
413 - hc_post_mult_value,
414 - sinkhorn_repeat,
415 - n_splits,
416 - hc_mult,
417 - )
418 - else:
419 - mhc_pre_big_fuse_with_norm_tilelang(
420 - gemm_out_mul,
421 - gemm_out_sqrsum,
422 - hc_scale,
423 - hc_base,
424 - residual_flat,
425 - post_mix,
426 - comb_mix,
427 - layer_input,
428 - norm_weight,
429 - post_mix,
430 - post_mix,
431 - hidden_size,
432 - rms_eps,
433 - hc_pre_eps,
434 - hc_sinkhorn_eps,
435 - hc_post_mult_value,
436 - sinkhorn_repeat,
437 - norm_eps,
438 - n_splits,
439 - hc_mult,
440 - )
441 -
372 return (
373 post_mix.view(*outer_shape, hc_mult, 1),
374 comb_mix.view(*outer_shape, hc_mult, hc_mult),
@@ -505,11 +435,10 @@ def mhc_pre_broadcast_tilelang(
435 fn_broadcast: torch.Tensor | None = None,
436 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
437 """First-layer mHC pre for a residual broadcast from ``(T, H)``."""
438 + # n_splits is retained for the shared API; split selection is internal.
439 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
509 - compute_num_split,
510 - mhc_pre_big_fuse_broadcast_with_norm_tilelang,
440 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
441 )
512 - from vllm.utils.math_utils import cdiv
442
443 assert norm_weight is not None, "broadcast mHC pre currently requires fused RMSNorm"
444 assert residual.dtype == torch.bfloat16
@@ -537,8 +466,6 @@ def mhc_pre_broadcast_tilelang(
466 residual_flat = residual
467 num_tokens = residual.shape[0]
468
540 - n_splits = compute_num_split(64, hidden_size, cdiv(num_tokens, 64))
541 -
469 residual_out = torch.empty(
470 num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=residual.device
471 )
@@ -551,42 +478,31 @@ def mhc_pre_broadcast_tilelang(
478 layer_input = torch.empty(
479 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
480 )
554 - gemm_out_mul = torch.empty(
555 - n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
556 - )
557 - gemm_out_sqrsum = torch.empty(
558 - n_splits, num_tokens, dtype=torch.float32, device=residual.device
559 - )
481
561 - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
562 -
563 - tf32_hc_prenorm_gemm(
482 + gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
483 residual_flat,
484 fn_broadcast,
566 - gemm_out_mul,
567 - gemm_out_sqrsum,
568 - n_splits,
485 + hidden_size=hidden_size,
486 + hc_mult=hc_mult,
487 + use_tilelang_fallback=False,
488 )
570 - mhc_pre_big_fuse_broadcast_with_norm_tilelang(
489 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
490 gemm_out_mul,
491 gemm_out_sqrsum,
492 hc_scale,
493 hc_base,
494 residual_flat,
576 - residual_out,
495 post_mix,
496 comb_mix,
497 layer_input,
580 - norm_weight,
581 - hidden_size,
498 rms_eps,
499 hc_pre_eps,
500 hc_sinkhorn_eps,
501 hc_post_mult_value,
502 sinkhorn_repeat,
587 - norm_eps,
588 - n_splits,
589 - hc_mult,
503 + residual_out=residual_out,
504 + norm_weight=norm_weight,
505 + norm_eps=norm_eps,
506 )
507 return (
508 residual_out,
@@ -603,11 +519,11 @@ def mhc_post_tilelang(
519 comb_res_mix: torch.Tensor,
520 ) -> torch.Tensor:
521 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
606 - mhc_post_tilelang as _mhc_post_kernel,
522 + _MHC_POST_TILELANG_KERNEL,
523 )
524
525 out = torch.empty_like(residual)
610 - _mhc_post_kernel(
526 + _MHC_POST_TILELANG_KERNEL(
527 comb_res_mix,
528 residual,
529 post_layer_mix.squeeze(-1),
@@ -644,6 +560,9 @@ def mhc_fused_post_pre_tilelang(
560 RMSNorm'd activation (fused into the kernel); otherwise it is the
561 raw pre-norm activation as before.
562
563 + ``n_splits`` and ``tile_n`` are retained for the shared MHC operator API.
564 + The TileLang path selects both values internally from the runtime shape.
565 +
566 Returns:
567 residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
568 post_mix_cur: shape (..., hc_mult, 1)
@@ -652,13 +571,10 @@ def mhc_fused_post_pre_tilelang(
571 """
572
573 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
655 - compute_num_split,
656 - mhc_fused_tilelang,
657 - mhc_post_tilelang,
658 - mhc_pre_big_fuse_tilelang,
659 - mhc_pre_big_fuse_with_norm_tilelang,
574 + _MHC_FUSED_TILELANG_KERNEL,
575 + _MHC_POST_TILELANG_KERNEL,
576 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
577 )
661 - from vllm.utils.math_utils import cdiv
578
579 assert residual.dtype == torch.bfloat16
580 assert x.dtype == torch.bfloat16
@@ -692,48 +608,14 @@ def mhc_fused_post_pre_tilelang(
608 if not norm_weight.is_contiguous():
609 norm_weight = norm_weight.contiguous()
610
695 - assert n_splits in (1, 2, 4, 8)
696 - assert hidden_size % n_splits == 0
697 -
611 residual_flat = residual.view(-1, hc_mult, hidden_size)
612 num_tokens = residual_flat.shape[0]
613 x_flat = x.view(num_tokens, hidden_size)
614 post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
615 comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
616
704 - from vllm.utils.deep_gemm import is_deep_gemm_supported
705 -
706 - use_deep_gemm = is_deep_gemm_supported()
617 use_small_fma = num_tokens <= 16
708 - if use_small_fma:
709 - # TODO(gnovack): investigate autotuning these heuristics
710 - tile_n = 2 if num_tokens < 8 else 3
711 - n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
712 - else:
713 - if use_deep_gemm:
714 - # these number are from deepgemm kernel impl
715 - block_k = 64
716 - block_m = 64
717 - n_splits = compute_num_split(
718 - block_k, hc_hidden_size, cdiv(num_tokens, block_m)
719 - )
720 - else:
721 - n_splits = 1
722 -
723 - gemm_out_mul = torch.empty(
724 - n_splits,
725 - num_tokens,
726 - hc_mult3,
727 - dtype=torch.float32,
728 - device=residual.device,
729 - )
730 - gemm_out_sqrsum = torch.empty(
731 - n_splits,
732 - num_tokens,
733 - dtype=torch.float32,
734 - device=residual.device,
735 - )
736 - residual_cur = torch.empty_like(residual_flat)
618 +
619 post_mix_cur = torch.empty(
620 num_tokens,
621 hc_mult,
@@ -754,55 +636,17 @@ def mhc_fused_post_pre_tilelang(
636 )
637
638 if use_small_fma:
757 - mhc_fused_tilelang(
639 + gemm_out_mul, gemm_out_sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL(
640 comb_res_mix_flat,
641 residual_flat,
642 post_layer_mix_flat,
643 x_flat,
644 fn.view(hc_mult3, hc_mult, hidden_size),
763 - gemm_out_mul,
764 - gemm_out_sqrsum,
765 - residual_cur,
645 hc_mult,
646 hidden_size,
647 hc_mult3,
769 - tile_n=tile_n,
770 - n_splits=n_splits,
648 )
772 - else:
773 - mhc_post_tilelang(
774 - comb_res_mix_flat,
775 - residual_flat,
776 - post_layer_mix_flat,
777 - x_flat,
778 - residual_cur,
779 - residual.shape[-2],
780 - residual.shape[-1],
781 - )
782 -
783 - residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size)
784 - if use_deep_gemm:
785 - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
786 -
787 - tf32_hc_prenorm_gemm(
788 - residual_cur_2d,
789 - fn,
790 - gemm_out_mul,
791 - gemm_out_sqrsum,
792 - n_splits,
793 - )
794 - else:
795 - _tilelang_hc_prenorm_gemm(
796 - residual_cur_2d,
797 - fn,
798 - gemm_out_mul,
799 - gemm_out_sqrsum,
800 - hidden_size,
801 - hc_mult,
802 - )
803 -
804 - if norm_weight is None:
805 - mhc_pre_big_fuse_tilelang(
649 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
650 gemm_out_mul,
651 gemm_out_sqrsum,
652 hc_scale,
@@ -811,20 +655,32 @@ def mhc_fused_post_pre_tilelang(
655 post_mix_cur,
656 comb_mix_cur,
657 layer_input_cur,
814 - # Pre-mix buffers are unused in V4 mode.
815 - post_mix_cur,
816 - post_mix_cur,
817 - hidden_size,
658 rms_eps,
659 hc_pre_eps,
660 hc_sinkhorn_eps,
661 hc_post_mult_value,
662 sinkhorn_repeat,
823 - n_splits,
824 - hc_mult,
663 + norm_weight=norm_weight,
664 + norm_eps=norm_eps,
665 )
666 else:
827 - mhc_pre_big_fuse_with_norm_tilelang(
667 + residual_cur = torch.empty_like(residual_flat)
668 + _MHC_POST_TILELANG_KERNEL(
669 + comb_res_mix_flat,
670 + residual_flat,
671 + post_layer_mix_flat,
672 + x_flat,
673 + residual_cur,
674 + residual.shape[-2],
675 + residual.shape[-1],
676 + )
677 + gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
678 + residual_cur.view(num_tokens, hc_mult * hidden_size),
679 + fn,
680 + hidden_size=hidden_size,
681 + hc_mult=hc_mult,
682 + )
683 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
684 gemm_out_mul,
685 gemm_out_sqrsum,
686 hc_scale,
@@ -833,18 +689,13 @@ def mhc_fused_post_pre_tilelang(
689 post_mix_cur,
690 comb_mix_cur,
691 layer_input_cur,
836 - norm_weight,
837 - post_mix_cur,
838 - post_mix_cur,
839 - hidden_size,
692 rms_eps,
693 hc_pre_eps,
694 hc_sinkhorn_eps,
695 hc_post_mult_value,
696 sinkhorn_repeat,
845 - norm_eps,
846 - n_splits,
847 - hc_mult,
697 + norm_weight=norm_weight,
698 + norm_eps=norm_eps,
699 )
700
701 return (
@@ -926,9 +777,11 @@ def hc_head_fused_kernel_tilelang(
777 )
778 if num_tokens == 0:
779 return out
929 - from vllm.model_executor.kernels.mhc.tilelang_kernels import hc_head_fuse_tilelang
780 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
781 + _HC_HEAD_FUSED_TILELANG_KERNEL,
782 + )
783
931 - hc_head_fuse_tilelang(
784 + _HC_HEAD_FUSED_TILELANG_KERNEL(
785 hs_flat,
786 fn,
787 hc_scale,
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+755
@@ -4,10 +4,19 @@
4 from __future__ import annotations
5
6 import math
7 +from dataclasses import dataclass
8 from functools import cache
9 +from typing import Any
10
11 import torch
12
13 +from vllm.model_executor.warmup.jit_warmup import WarmupIntRange, zip_inputs
14 +from vllm.model_executor.warmup.jit_warmup_tilelang_helper import (
15 + TileLangLaunchSpec,
16 + VllmTileLangJitKernel,
17 + kernel_launcher,
18 + make_tilelang_warmup_tensor,
19 +)
20 from vllm.platforms import current_platform
21 from vllm.tilelang_utils import T, tilelang, tilelang_jit
22 from vllm.utils.math_utils import cdiv
@@ -28,6 +37,10 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
37 return split_k
38
39
40 +def _mhc_fused_n_splits(num_tokens: int, hidden_size: int) -> int:
41 + return 8 if num_tokens < 8 and hidden_size <= 4096 else 4
42 +
43 +
44 @tilelang_jit
45 def mhc_pre_big_fuse_tilelang(
46 gemm_out_mul,
@@ -987,3 +1000,745 @@ def hc_head_fuse_tilelang(
1000
1001 if ENABLE_PDL:
1002 T.pdl_trigger()
1003 +
1004 +
1005 +class HcPrenormGemmTileLangKernel(
1006 + VllmTileLangJitKernel["HcPrenormGemmTileLangKernel.CompileKey"]
1007 +):
1008 + @dataclass(frozen=True)
1009 + class CompileKey:
1010 + hidden_size: int
1011 + hc_mult: int
1012 + n_out: int
1013 + n_thr: int
1014 + tile_n: int
1015 + n_splits: int
1016 + use_block_m: bool
1017 + block_m: int
1018 +
1019 + @staticmethod
1020 + def kernel(compile_key: CompileKey) -> Any:
1021 + if compile_key.use_block_m:
1022 + return hc_prenorm_gemm_block_m_tilelang
1023 + return hc_prenorm_gemm_tilelang
1024 +
1025 + def dispatch( # type: ignore[override]
1026 + self,
1027 + *,
1028 + num_tokens: int,
1029 + hc_hidden_size: int,
1030 + n_thr: int = 512,
1031 + tile_n: int = 12,
1032 + n_splits: int = 1,
1033 + **compile_key_fields: int,
1034 + ) -> CompileKey:
1035 + use_default_config = n_splits == 1 and tile_n == 12 and n_thr == 512
1036 + use_block_m = use_default_config and num_tokens >= 1024
1037 + use_small_tile = (
1038 + use_default_config
1039 + and not use_block_m
1040 + and num_tokens < 128
1041 + and hc_hidden_size % 1024 == 0
1042 + )
1043 + effective_n_thr = 1024 if use_small_tile else n_thr
1044 + effective_tile_n = 4 if use_small_tile else tile_n
1045 + return self.CompileKey(
1046 + **compile_key_fields,
1047 + n_thr=effective_n_thr,
1048 + tile_n=effective_tile_n,
1049 + n_splits=n_splits,
1050 + use_block_m=use_block_m,
1051 + block_m=2 if use_block_m else 1,
1052 + )
1053 +
1054 + def get_warmup_keys(
1055 + self,
1056 + vllm_config: Any,
1057 + *,
1058 + hidden_size: int,
1059 + hc_mult: int,
1060 + n_out: int,
1061 + ) -> list[CompileKey]:
1062 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
1063 + if max_tokens <= 0:
1064 + return []
1065 + return self._trace_dispatch(self.dispatch)(
1066 + num_tokens=WarmupIntRange(
1067 + 1,
1068 + max_tokens + 1,
1069 + advance=lambda value: (
1070 + 128 if value < 128 else 1024 if value < 1024 else max_tokens + 1
1071 + ),
1072 + ),
1073 + hc_hidden_size=hidden_size * hc_mult,
1074 + hidden_size=hidden_size,
1075 + hc_mult=hc_mult,
1076 + n_out=n_out,
1077 + )
1078 +
1079 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1080 + num_tokens = 128
1081 + if compile_key.use_block_m:
1082 + num_tokens = 1024
1083 + elif compile_key.n_thr == 1024:
1084 + num_tokens = 1
1085 + hc_hidden_size = compile_key.hidden_size * compile_key.hc_mult
1086 + return dict(
1087 + x=make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hc_hidden_size),
1088 + fn=make_tilelang_warmup_tensor(
1089 + torch.float32, compile_key.n_out, hc_hidden_size
1090 + ),
1091 + out=make_tilelang_warmup_tensor(
1092 + torch.float32,
1093 + compile_key.n_splits,
1094 + num_tokens,
1095 + compile_key.n_out,
1096 + ),
1097 + sqrsum=make_tilelang_warmup_tensor(
1098 + torch.float32, compile_key.n_splits, num_tokens
1099 + ),
1100 + hidden_size=compile_key.hidden_size,
1101 + hc_mult=compile_key.hc_mult,
1102 + tile_n=compile_key.tile_n,
1103 + n_thr=compile_key.n_thr,
1104 + n_splits=compile_key.n_splits,
1105 + )
1106 +
1107 + @kernel_launcher
1108 + def __call__(
1109 + self,
1110 + x: Any,
1111 + fn: Any,
1112 + out: Any,
1113 + sqrsum: Any,
1114 + hidden_size: int,
1115 + hc_mult: int,
1116 + tile_n: int = 12,
1117 + n_thr: int = 512,
1118 + n_splits: int = 1,
1119 + ) -> TileLangLaunchSpec:
1120 + assert out.shape[0] == n_splits
1121 + assert sqrsum.shape[0] == n_splits
1122 + assert x.shape[1] == hc_mult * hidden_size
1123 + assert x.shape[1] % n_splits == 0
1124 + assert (x.shape[1] // n_splits) % n_thr == 0
1125 + compile_key = self.dispatch(
1126 + num_tokens=x.shape[0],
1127 + hc_hidden_size=x.shape[1],
1128 + hidden_size=hidden_size,
1129 + hc_mult=hc_mult,
1130 + n_out=fn.shape[0],
1131 + n_thr=n_thr,
1132 + tile_n=tile_n,
1133 + n_splits=n_splits,
1134 + )
1135 + kernel_arg = (
1136 + compile_key.block_m if compile_key.use_block_m else compile_key.n_splits
1137 + )
1138 + return (compile_key,), (
1139 + x,
1140 + fn,
1141 + out,
1142 + sqrsum,
1143 + hidden_size,
1144 + hc_mult,
1145 + fn.shape[0],
1146 + compile_key.n_thr,
1147 + compile_key.tile_n,
1148 + kernel_arg,
1149 + )
1150 +
1151 +
1152 +class MhcPreBigFuseTileLangKernel(
1153 + VllmTileLangJitKernel["MhcPreBigFuseTileLangKernel.CompileKey"]
1154 +):
1155 + @dataclass(frozen=True)
1156 + class CompileKey:
1157 + hidden_size: int
1158 + hc_mult: int
1159 + n_splits: int
1160 + use_norm_weight: bool
1161 + is_broadcast: bool
1162 + rms_eps: float
1163 + hc_pre_eps: float
1164 + hc_sinkhorn_eps: float
1165 + hc_post_mult_value: float
1166 + sinkhorn_repeat: int
1167 + norm_eps: float
1168 +
1169 + @staticmethod
1170 + def kernel(compile_key: CompileKey) -> Any:
1171 + if compile_key.is_broadcast:
1172 + return mhc_pre_big_fuse_broadcast_with_norm_tilelang
1173 + if compile_key.use_norm_weight:
1174 + return mhc_pre_big_fuse_with_norm_tilelang
1175 + return mhc_pre_big_fuse_tilelang
1176 +
1177 + @staticmethod
1178 + def _kernel_args(
1179 + compile_key: CompileKey,
1180 + *,
1181 + gemm_out_mul: Any,
1182 + gemm_out_sqrsum: Any,
1183 + hc_scale: Any,
1184 + hc_base: Any,
1185 + residual: Any,
1186 + residual_out: Any | None,
1187 + post_mix: Any,
1188 + comb_mix: Any,
1189 + layer_input: Any,
1190 + norm_weight: Any | None,
1191 + ) -> tuple[Any, ...]:
1192 + common_args = (
1193 + gemm_out_mul,
1194 + gemm_out_sqrsum,
1195 + hc_scale,
1196 + hc_base,
1197 + residual,
1198 + )
1199 + output_args = (post_mix, comb_mix, layer_input)
1200 + eps_args = (
1201 + compile_key.rms_eps,
1202 + compile_key.hc_pre_eps,
1203 + compile_key.hc_sinkhorn_eps,
1204 + compile_key.hc_post_mult_value,
1205 + compile_key.sinkhorn_repeat,
1206 + )
1207 + middle_args: tuple[Any, ...]
1208 + if compile_key.is_broadcast:
1209 + assert residual_out is not None
1210 + assert norm_weight is not None
1211 + middle_args = (residual_out, *output_args, norm_weight)
1212 + elif compile_key.use_norm_weight:
1213 + assert norm_weight is not None
1214 + middle_args = (*output_args, norm_weight, post_mix, post_mix)
1215 + else:
1216 + middle_args = (*output_args, post_mix, post_mix)
1217 + norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else ()
1218 + return (
1219 + *common_args,
1220 + *middle_args,
1221 + compile_key.hidden_size,
1222 + *eps_args,
1223 + *norm_eps_args,
1224 + compile_key.n_splits,
1225 + compile_key.hc_mult,
1226 + )
1227 +
1228 + def dispatch( # type: ignore[override]
1229 + self,
1230 + *,
1231 + hidden_size: int,
1232 + hc_mult: int,
1233 + n_splits: int,
1234 + is_broadcast: bool,
1235 + use_norm_weight: bool,
1236 + use_pre_gemm_splits: bool = False,
1237 + pre_gemm_k: int = 0,
1238 + sinkhorn_repeat: int,
1239 + norm_eps: float,
1240 + broadcast_norm_eps: float = 0.0,
1241 + num_tokens: int = 1,
1242 + use_fused_tilelang: bool = False,
1243 + **compile_key_fields: float,
1244 + ) -> CompileKey:
1245 + pre_gemm_n_splits = (
1246 + compute_num_split(64, pre_gemm_k, (num_tokens + 63) // 64)
1247 + if use_pre_gemm_splits
1248 + else n_splits
1249 + )
1250 + fused_n_splits = _mhc_fused_n_splits(num_tokens, hidden_size)
1251 + actual_n_splits = fused_n_splits if use_fused_tilelang else pre_gemm_n_splits
1252 + actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps
1253 + actual_use_norm_weight = use_norm_weight or is_broadcast
1254 + return self.CompileKey(
1255 + **compile_key_fields,
1256 + hidden_size=hidden_size,
1257 + hc_mult=hc_mult,
1258 + n_splits=actual_n_splits,
1259 + use_norm_weight=actual_use_norm_weight,
1260 + is_broadcast=is_broadcast,
1261 + sinkhorn_repeat=sinkhorn_repeat,
1262 + norm_eps=actual_norm_eps if actual_use_norm_weight else 0.0,
1263 + )
1264 +
1265 + def get_warmup_keys(
1266 + self,
1267 + vllm_config: Any,
1268 + *,
1269 + hidden_size: int,
1270 + hc_mult: int,
1271 + use_norm_weight: bool,
1272 + rms_eps: float,
1273 + hc_pre_eps: float,
1274 + hc_sinkhorn_eps: float,
1275 + hc_post_mult_value: float,
1276 + sinkhorn_repeat: int,
1277 + norm_eps: float | tuple[float, ...],
1278 + broadcast_norm_eps: float | tuple[float, ...] | None = None,
1279 + include_pre_gemm_splits: bool = False,
1280 + include_broadcast_splits: bool = False,
1281 + ) -> list[CompileKey]:
1282 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
1283 + if max_tokens <= 0:
1284 + return []
1285 +
1286 + warmup_cases = zip_inputs(
1287 + dict(
1288 + is_broadcast=False,
1289 + use_pre_gemm_splits=False,
1290 + pre_gemm_k=0,
1291 + use_fused_tilelang=False,
1292 + ),
1293 + dict(
1294 + is_broadcast=False,
1295 + use_pre_gemm_splits=include_pre_gemm_splits,
1296 + pre_gemm_k=hc_mult * hidden_size,
1297 + use_fused_tilelang=False,
1298 + ),
1299 + dict(
1300 + is_broadcast=include_broadcast_splits,
1301 + use_pre_gemm_splits=include_broadcast_splits,
1302 + pre_gemm_k=hidden_size,
1303 + use_fused_tilelang=False,
1304 + ),
1305 + dict(
1306 + is_broadcast=False,
1307 + use_pre_gemm_splits=False,
1308 + pre_gemm_k=0,
1309 + use_fused_tilelang=True,
1310 + ),
1311 + )
1312 + return self._trace_dispatch(self.dispatch)(
1313 + warmup_cases,
1314 + num_tokens=WarmupIntRange(
1315 + 1,
1316 + max_tokens + 1,
1317 + advance=lambda value: 8 if value < 8 else cdiv(value, 64) * 64 + 1,
1318 + ),
1319 + n_splits=1,
1320 + hidden_size=hidden_size,
1321 + hc_mult=hc_mult,
1322 + use_norm_weight=use_norm_weight,
1323 + rms_eps=rms_eps,
1324 + hc_pre_eps=hc_pre_eps,
1325 + hc_sinkhorn_eps=hc_sinkhorn_eps,
1326 + hc_post_mult_value=hc_post_mult_value,
1327 + sinkhorn_repeat=sinkhorn_repeat,
1328 + norm_eps=norm_eps,
1329 + broadcast_norm_eps=broadcast_norm_eps
1330 + if broadcast_norm_eps is not None
1331 + else norm_eps,
1332 + _when=lambda *, use_fused_tilelang, num_tokens: (
1333 + not use_fused_tilelang or num_tokens <= 16
1334 + ),
1335 + )
1336 +
1337 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1338 + hidden_size = compile_key.hidden_size
1339 + hc_mult = compile_key.hc_mult
1340 + hc_mult3 = hc_mult * (2 + hc_mult)
1341 + num_tokens = 1
1342 +
1343 + gemm_out_mul = make_tilelang_warmup_tensor(
1344 + torch.float32, compile_key.n_splits, num_tokens, hc_mult3
1345 + )
1346 + gemm_out_sqrsum = make_tilelang_warmup_tensor(
1347 + torch.float32, compile_key.n_splits, num_tokens
1348 + )
1349 + hc_scale = make_tilelang_warmup_tensor(torch.float32, 3)
1350 + hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult3)
1351 + post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
1352 + comb_mix = make_tilelang_warmup_tensor(
1353 + torch.float32, num_tokens, hc_mult * hc_mult
1354 + )
1355 + layer_input = make_tilelang_warmup_tensor(
1356 + torch.bfloat16, num_tokens, hidden_size
1357 + )
1358 + residual_shape = (
1359 + (num_tokens, hidden_size)
1360 + if compile_key.is_broadcast
1361 + else (num_tokens, hc_mult, hidden_size)
1362 + )
1363 + residual = make_tilelang_warmup_tensor(torch.bfloat16, *residual_shape)
1364 + residual_out = (
1365 + make_tilelang_warmup_tensor(
1366 + torch.bfloat16, num_tokens, hc_mult, hidden_size
1367 + )
1368 + if compile_key.is_broadcast
1369 + else None
1370 + )
1371 + norm_weight = (
1372 + make_tilelang_warmup_tensor(torch.bfloat16, hidden_size)
1373 + if compile_key.is_broadcast or compile_key.use_norm_weight
1374 + else None
1375 + )
1376 + return dict(
1377 + gemm_out_mul=gemm_out_mul,
1378 + gemm_out_sqrsum=gemm_out_sqrsum,
1379 + hc_scale=hc_scale,
1380 + hc_base=hc_base,
1381 + residual=residual,
1382 + post_mix=post_mix,
1383 + comb_mix=comb_mix,
1384 + layer_input=layer_input,
1385 + rms_eps=compile_key.rms_eps,
1386 + hc_pre_eps=compile_key.hc_pre_eps,
1387 + hc_sinkhorn_eps=compile_key.hc_sinkhorn_eps,
1388 + hc_post_mult_value=compile_key.hc_post_mult_value,
1389 + sinkhorn_repeat=compile_key.sinkhorn_repeat,
1390 + residual_out=residual_out,
1391 + norm_weight=norm_weight,
1392 + norm_eps=compile_key.norm_eps,
1393 + )
1394 +
1395 + @kernel_launcher
1396 + def __call__(
1397 + self,
1398 + gemm_out_mul: Any,
1399 + gemm_out_sqrsum: Any,
1400 + hc_scale: Any,
1401 + hc_base: Any,
1402 + residual: Any,
1403 + post_mix: Any,
1404 + comb_mix: Any,
1405 + layer_input: Any,
1406 + rms_eps: float,
1407 + hc_pre_eps: float,
1408 + hc_sinkhorn_eps: float,
1409 + hc_post_mult_value: float,
1410 + sinkhorn_repeat: int,
1411 + *,
1412 + residual_out: Any | None = None,
1413 + norm_weight: Any | None = None,
1414 + norm_eps: float = 0.0,
1415 + ) -> TileLangLaunchSpec:
1416 + is_broadcast = residual_out is not None
1417 + use_norm_weight = norm_weight is not None
1418 + n_splits = gemm_out_mul.shape[0]
1419 + if is_broadcast:
1420 + assert residual_out is not None
1421 + hidden_size = residual.shape[-1]
1422 + hc_mult = residual_out.shape[-2]
1423 + else:
1424 + hc_mult = residual.shape[-2]
1425 + hidden_size = residual.shape[-1]
1426 + compile_key = self.dispatch(
1427 + hidden_size=hidden_size,
1428 + hc_mult=hc_mult,
1429 + n_splits=n_splits,
1430 + is_broadcast=is_broadcast,
1431 + use_norm_weight=use_norm_weight,
1432 + rms_eps=rms_eps,
1433 + hc_pre_eps=hc_pre_eps,
1434 + hc_sinkhorn_eps=hc_sinkhorn_eps,
1435 + hc_post_mult_value=hc_post_mult_value,
1436 + sinkhorn_repeat=sinkhorn_repeat,
1437 + norm_eps=norm_eps,
1438 + broadcast_norm_eps=norm_eps,
1439 + )
1440 + return (
1441 + (compile_key,),
1442 + self._kernel_args(
1443 + compile_key,
1444 + gemm_out_mul=gemm_out_mul,
1445 + gemm_out_sqrsum=gemm_out_sqrsum,
1446 + hc_scale=hc_scale,
1447 + hc_base=hc_base,
1448 + residual=residual,
1449 + residual_out=residual_out,
1450 + post_mix=post_mix,
1451 + comb_mix=comb_mix,
1452 + layer_input=layer_input,
1453 + norm_weight=norm_weight,
1454 + ),
1455 + )
1456 +
1457 +
1458 +class MhcPostTileLangKernel(VllmTileLangJitKernel["MhcPostTileLangKernel.CompileKey"]):
1459 + @dataclass(frozen=True)
1460 + class CompileKey:
1461 + hidden_size: int
1462 + hc_mult: int
1463 +
1464 + @staticmethod
1465 + def kernel() -> Any:
1466 + return mhc_post_tilelang
1467 +
1468 + def dispatch( # type: ignore[override]
1469 + self,
1470 + **compile_key_fields: int,
1471 + ) -> CompileKey:
1472 + return self.CompileKey(**compile_key_fields)
1473 +
1474 + def get_warmup_keys(
1475 + self,
1476 + *,
1477 + hidden_size: int,
1478 + hc_mult: int,
1479 + ) -> list[CompileKey]:
1480 + return self._trace_dispatch(self.dispatch)(
1481 + hidden_size=hidden_size,
1482 + hc_mult=hc_mult,
1483 + )
1484 +
1485 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1486 + num_tokens = 1
1487 + hidden_size = compile_key.hidden_size
1488 + hc_mult = compile_key.hc_mult
1489 + comb_mix = make_tilelang_warmup_tensor(
1490 + torch.float32, num_tokens, hc_mult, hc_mult
1491 + )
1492 + residual = make_tilelang_warmup_tensor(
1493 + torch.bfloat16, num_tokens, hc_mult, hidden_size
1494 + )
1495 + post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
1496 + layer_input = make_tilelang_warmup_tensor(
1497 + torch.bfloat16, num_tokens, hidden_size
1498 + )
1499 + out = make_tilelang_warmup_tensor(
1500 + torch.bfloat16, num_tokens, hc_mult, hidden_size
1501 + )
1502 + return dict(
1503 + comb_mix=comb_mix,
1504 + residual=residual,
1505 + post_mix=post_mix,
1506 + layer_input=layer_input,
1507 + out=out,
1508 + hc_mult=hc_mult,
1509 + hidden_size=hidden_size,
1510 + )
1511 +
1512 + @kernel_launcher
1513 + def __call__(
1514 + self,
1515 + comb_mix: Any,
1516 + residual: Any,
1517 + post_mix: Any,
1518 + layer_input: Any,
1519 + out: Any,
1520 + hc_mult: int,
1521 + hidden_size: int,
1522 + ) -> TileLangLaunchSpec:
1523 + return (), (
1524 + comb_mix,
1525 + residual,
1526 + post_mix,
1527 + layer_input,
1528 + out,
1529 + hc_mult,
1530 + hidden_size,
1531 + )
1532 +
1533 +
1534 +class MhcFusedTileLangKernel(
1535 + VllmTileLangJitKernel["MhcFusedTileLangKernel.CompileKey"]
1536 +):
1537 + @dataclass(frozen=True)
1538 + class CompileKey:
1539 + hidden_size: int
1540 + hc_mult: int
1541 + n_splits: int
1542 + tile_n: int
1543 +
1544 + @staticmethod
1545 + def kernel() -> Any:
1546 + return mhc_fused_tilelang
1547 +
1548 + def dispatch( # type: ignore[override]
1549 + self,
1550 + *,
1551 + num_tokens: int,
1552 + hidden_size: int,
1553 + hc_mult: int,
1554 + ) -> CompileKey:
1555 + # TODO(gnovack): investigate autotuning these heuristics
1556 + tile_n = 2 if num_tokens < 8 else 3
1557 + n_splits = _mhc_fused_n_splits(num_tokens, hidden_size)
1558 + return self.CompileKey(
1559 + hidden_size=hidden_size,
1560 + hc_mult=hc_mult,
1561 + n_splits=n_splits,
1562 + tile_n=tile_n,
1563 + )
1564 +
1565 + def get_warmup_keys(
1566 + self,
1567 + vllm_config: Any,
1568 + *,
1569 + hidden_size: int,
1570 + hc_mult: int,
1571 + ) -> list[CompileKey]:
1572 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
1573 + return self._trace_dispatch(self.dispatch)(
1574 + num_tokens=WarmupIntRange(1, max_tokens + 1),
1575 + hidden_size=hidden_size,
1576 + hc_mult=hc_mult,
1577 + _when=lambda *, num_tokens: num_tokens <= 16,
1578 + )
1579 +
1580 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1581 + num_tokens = 1 if compile_key.tile_n == 2 else 8
1582 + hidden_size = compile_key.hidden_size
1583 + hc_mult = compile_key.hc_mult
1584 + hc_mult3 = hc_mult * (2 + hc_mult)
1585 + comb_mix = make_tilelang_warmup_tensor(
1586 + torch.float32, num_tokens, hc_mult, hc_mult
1587 + )
1588 + residual_in = make_tilelang_warmup_tensor(
1589 + torch.bfloat16, num_tokens, hc_mult, hidden_size
1590 + )
1591 + post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
1592 + x_in = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size)
1593 + weight_t = make_tilelang_warmup_tensor(
1594 + torch.float32, hc_mult3, hc_mult, hidden_size
1595 + )
1596 + return dict(
1597 + comb_mix=comb_mix,
1598 + residual_in=residual_in,
1599 + post_mix=post_mix,
1600 + x_in=x_in,
1601 + weight_t=weight_t,
1602 + hc_mult=hc_mult,
1603 + hidden_size=hidden_size,
1604 + hc_mult3=hc_mult3,
1605 + )
1606 +
1607 + @kernel_launcher
1608 + def __call__(
1609 + self,
1610 + comb_mix: Any,
1611 + residual_in: Any,
1612 + post_mix: Any,
1613 + x_in: Any,
1614 + weight_t: Any,
1615 + hc_mult: int,
1616 + hidden_size: int,
1617 + hc_mult3: int,
1618 + ) -> TileLangLaunchSpec:
1619 + num_tokens = residual_in.shape[0]
1620 + tile_n = 2 if num_tokens < 8 else 3
1621 + n_splits = _mhc_fused_n_splits(num_tokens, hidden_size)
1622 + yp_out = residual_in.new_empty(
1623 + (n_splits, num_tokens, hc_mult3), dtype=torch.float32
1624 + )
1625 + rp_out = residual_in.new_empty((n_splits, num_tokens), dtype=torch.float32)
1626 + residual_out = residual_in.new_empty(residual_in.shape)
1627 + return (
1628 + (),
1629 + (
1630 + comb_mix,
1631 + residual_in,
1632 + post_mix,
1633 + x_in,
1634 + weight_t,
1635 + yp_out,
1636 + rp_out,
1637 + residual_out,
1638 + hc_mult,
1639 + hidden_size,
1640 + hc_mult3,
1641 + ),
1642 + dict(tile_n=tile_n, split_k=n_splits),
1643 + (yp_out, rp_out, residual_out),
1644 + )
1645 +
1646 +
1647 +class HcHeadFusedTileLangKernel(
1648 + VllmTileLangJitKernel["HcHeadFusedTileLangKernel.CompileKey"]
1649 +):
1650 + @dataclass(frozen=True)
1651 + class CompileKey:
1652 + hidden_size: int
1653 + hc_mult: int
1654 + rms_eps: float
1655 + hc_eps: float
1656 +
1657 + @staticmethod
1658 + def kernel() -> Any:
1659 + return hc_head_fuse_tilelang
1660 +
1661 + def dispatch( # type: ignore[override]
1662 + self,
1663 + *,
1664 + hidden_size: int,
1665 + hc_mult: int,
1666 + rms_eps: float,
1667 + hc_eps: float,
1668 + ) -> CompileKey:
1669 + return self.CompileKey(
1670 + hidden_size=hidden_size,
1671 + hc_mult=hc_mult,
1672 + rms_eps=rms_eps,
1673 + hc_eps=hc_eps,
1674 + )
1675 +
1676 + def get_warmup_keys(
1677 + self,
1678 + *,
1679 + hidden_size: int,
1680 + hc_mult: int,
1681 + rms_eps: float,
1682 + hc_eps: float,
1683 + ) -> list[CompileKey]:
1684 + return self._trace_dispatch(self.dispatch)(
1685 + hidden_size=hidden_size,
1686 + hc_mult=hc_mult,
1687 + rms_eps=rms_eps,
1688 + hc_eps=hc_eps,
1689 + )
1690 +
1691 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1692 + num_tokens = 1
1693 + hidden_size = compile_key.hidden_size
1694 + hc_mult = compile_key.hc_mult
1695 + residual = make_tilelang_warmup_tensor(
1696 + torch.bfloat16, num_tokens, hc_mult, hidden_size
1697 + )
1698 + fn = make_tilelang_warmup_tensor(torch.float32, hc_mult, hc_mult * hidden_size)
1699 + hc_scale = make_tilelang_warmup_tensor(torch.float32, 1)
1700 + hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult)
1701 + out = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size)
1702 + return dict(
1703 + residual=residual,
1704 + fn=fn,
1705 + hc_scale=hc_scale,
1706 + hc_base=hc_base,
1707 + out=out,
1708 + hidden_size=hidden_size,
1709 + rms_eps=compile_key.rms_eps,
1710 + hc_eps=compile_key.hc_eps,
1711 + hc_mult=hc_mult,
1712 + )
1713 +
1714 + @kernel_launcher
1715 + def __call__(
1716 + self,
1717 + residual: Any,
1718 + fn: Any,
1719 + hc_scale: Any,
1720 + hc_base: Any,
1721 + out: Any,
1722 + hidden_size: int,
1723 + rms_eps: float,
1724 + hc_eps: float,
1725 + hc_mult: int,
1726 + ) -> TileLangLaunchSpec:
1727 + return (), (
1728 + residual,
1729 + fn,
1730 + hc_scale,
1731 + hc_base,
1732 + out,
1733 + hidden_size,
1734 + rms_eps,
1735 + hc_eps,
1736 + hc_mult,
1737 + )
1738 +
1739 +
1740 +_HC_PRENORM_GEMM_TILELANG_KERNEL = HcPrenormGemmTileLangKernel()
1741 +_MHC_PRE_BIG_FUSE_TILELANG_KERNEL = MhcPreBigFuseTileLangKernel()
1742 +_MHC_POST_TILELANG_KERNEL = MhcPostTileLangKernel()
1743 +_MHC_FUSED_TILELANG_KERNEL = MhcFusedTileLangKernel()
1744 +_HC_HEAD_FUSED_TILELANG_KERNEL = HcHeadFusedTileLangKernel()
vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py deleted
-211
@@ -1,211 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Warm up DeepSeek V4 mHC TileLang kernels before serving requests.
4 -
5 -Ported from lucifer1004/vllm-jasl with the two env-var knobs removed
6 -(`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`).
7 -Gating is intrinsic: non-DSv4 models and layers without hc_* attributes
8 -return early, so the warmup is a no-op except where it's needed.
9 -"""
10 -
11 -import time
12 -from collections.abc import Iterable
13 -
14 -import torch
15 -
16 -from vllm.logger import init_logger
17 -from vllm.tracing import instrument
18 -
19 -logger = init_logger(__name__)
20 -
21 -_AUTO_WARMUP_MAX_TOKENS = 16_384
22 -_DEFAULT_TOKEN_SIZE_CANDIDATES = (
23 - 1,
24 - 2,
25 - 4,
26 - 8,
27 - 16,
28 - 32,
29 - 64,
30 - 128,
31 - 256,
32 - 512,
33 - 1024,
34 - 2048,
35 - 4096,
36 - 8192,
37 - 16_384,
38 -)
39 -
40 -
41 -def _normalize_token_sizes(
42 - token_sizes: Iterable[int],
43 - *,
44 - max_tokens: int,
45 -) -> list[int]:
46 - return sorted({size for size in token_sizes if 1 <= size <= max_tokens})
47 -
48 -
49 -def _select_mhc_warmup_token_sizes(
50 - *,
51 - max_tokens: int,
52 - cudagraph_capture_sizes: list[int],
53 -) -> list[int]:
54 - if max_tokens <= 0:
55 - return []
56 -
57 - max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS)
58 - candidates = list(_DEFAULT_TOKEN_SIZE_CANDIDATES)
59 - candidates.extend(cudagraph_capture_sizes)
60 - candidates.append(max_auto_tokens)
61 - return _normalize_token_sizes(candidates, max_tokens=max_auto_tokens)
62 -
63 -
64 -def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None:
65 - for module in model.modules():
66 - if module.__class__.__name__ != "DeepseekV4DecoderLayer":
67 - continue
68 - if all(
69 - hasattr(module, attr)
70 - for attr in (
71 - "hc_pre",
72 - "hc_post",
73 - "hc_attn_fn",
74 - "hc_attn_scale",
75 - "hc_attn_base",
76 - "hc_ffn_fn",
77 - "hc_ffn_scale",
78 - "hc_ffn_base",
79 - )
80 - ):
81 - return module
82 - return None
83 -
84 -
85 -def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None:
86 - for module in model.modules():
87 - if module.__class__.__name__ != "DeepseekV4Model":
88 - continue
89 - if all(
90 - hasattr(module, attr)
91 - for attr in ("hc_head_fn", "hc_head_scale", "hc_head_base")
92 - ):
93 - return module
94 - return None
95 -
96 -
97 -def _warmup_layer_mhc(
98 - layer: torch.nn.Module,
99 - token_sizes: list[int],
100 -) -> None:
101 - max_tokens = max(token_sizes)
102 - hidden_size = int(layer.hidden_size)
103 - hc_mult = int(layer.hc_mult)
104 - device = layer.hc_attn_fn.device
105 - residual = torch.zeros(
106 - max_tokens,
107 - hc_mult,
108 - hidden_size,
109 - dtype=torch.bfloat16,
110 - device=device,
111 - )
112 -
113 - for size in token_sizes:
114 - residual_slice = residual[:size]
115 - for fn, scale, base in (
116 - (layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base),
117 - (layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base),
118 - ):
119 - layer_input, post_mix, comb_mix = layer.hc_pre(
120 - residual_slice,
121 - fn,
122 - scale,
123 - base,
124 - )
125 - layer.hc_post(layer_input, residual_slice, post_mix, comb_mix)
126 -
127 -
128 -def _warmup_hc_head(
129 - model: torch.nn.Module,
130 - token_sizes: list[int],
131 -) -> None:
132 - # Upstream a8887c208 ("[DSV4] aiter mhc support (ROCm)") refactored
133 - # ``hc_head`` from a free function into the ``HCHeadOp`` CustomOp
134 - # instance attached to the model as ``hc_head_op``. We call through
135 - # that instance so the warmup exercises the same dispatched
136 - # implementation as the inference path.
137 - hc_head_op = getattr(model, "hc_head_op", None)
138 - if hc_head_op is None:
139 - return
140 -
141 - max_tokens = max(token_sizes)
142 - hidden_size = int(model.config.hidden_size)
143 - hc_mult = int(model.hc_mult)
144 - device = model.hc_head_fn.device
145 - hidden_states = torch.zeros(
146 - max_tokens,
147 - hc_mult,
148 - hidden_size,
149 - dtype=torch.bfloat16,
150 - device=device,
151 - )
152 -
153 - for size in token_sizes:
154 - hc_head_op(
155 - hidden_states[:size],
156 - model.hc_head_fn,
157 - model.hc_head_scale,
158 - model.hc_head_base,
159 - model.rms_norm_eps,
160 - model.hc_eps,
161 - )
162 -
163 -
164 -@instrument(span_name="DeepSeek V4 mHC warmup")
165 -def deepseek_v4_mhc_warmup(
166 - model: torch.nn.Module,
167 - *,
168 - max_tokens: int,
169 - cudagraph_capture_sizes: list[int] | None = None,
170 -) -> None:
171 - # Cheap model-type gate before walking ``model.modules()``. The class
172 - # walk below is O(num_layers) and shows up in startup time on very
173 - # large checkpoints; bail out for any model that is not DeepSeek V4.
174 - config = getattr(model, "config", None)
175 - model_type = getattr(config, "model_type", None) if config is not None else None
176 - if model_type is not None and model_type not in (
177 - "deepseek_v4",
178 - "deepseek_v41",
179 - ):
180 - return
181 -
182 - layer = _find_first_mhc_layer(model)
183 - if layer is None:
184 - return
185 -
186 - device = layer.hc_attn_fn.device
187 - if device.type != "cuda":
188 - return
189 -
190 - deepseek_model = _find_deepseek_v4_model(model)
191 - token_sizes = _select_mhc_warmup_token_sizes(
192 - max_tokens=max_tokens,
193 - cudagraph_capture_sizes=cudagraph_capture_sizes or [],
194 - )
195 - if not token_sizes:
196 - return
197 -
198 - started = time.perf_counter()
199 - logger.info(
200 - "Warming up DeepSeek V4 mHC TileLang kernels for token sizes: %s",
201 - token_sizes,
202 - )
203 - with torch.inference_mode():
204 - _warmup_layer_mhc(layer, token_sizes)
205 - if deepseek_model is not None:
206 - _warmup_hc_head(deepseek_model, token_sizes)
207 - torch.accelerator.synchronize()
208 - logger.info(
209 - "DeepSeek V4 mHC TileLang warmup finished in %.2f seconds.",
210 - time.perf_counter() - started,
211 - )
vllm/model_executor/warmup/jit_warmup_tilelang_helper.py new
+172
@@ -0,0 +1,172 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Compile-only helpers for TileLang JIT warmup."""
4 +
5 +from __future__ import annotations
6 +
7 +import logging
8 +from abc import abstractmethod
9 +from collections.abc import Callable, Iterator, Mapping, MutableMapping
10 +from contextlib import contextmanager
11 +from dataclasses import dataclass, field
12 +from functools import wraps
13 +from typing import Any, ClassVar, Generic, TypeAlias, TypeVar, cast
14 +
15 +import torch
16 +
17 +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
18 +
19 +CompileKeyT = TypeVar("CompileKeyT")
20 +# (kernel_args, launch_args), optionally followed by keywords and a result.
21 +TileLangLaunchSpec: TypeAlias = (
22 + tuple[tuple[Any, ...], tuple[Any, ...]]
23 + | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None]
24 + | tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None, Any]
25 +)
26 +
27 +
28 +@contextmanager
29 +def _quiet_tilelang_warmup_logs() -> Iterator[None]:
30 + loggers = [
31 + logging.getLogger("tilelang.cache.kernel_cache"),
32 + logging.getLogger("tilelang.jit.kernel"),
33 + ]
34 + previous_levels = [logger.level for logger in loggers]
35 + try:
36 + for logger in loggers:
37 + logger.setLevel(logging.ERROR)
38 + yield
39 + finally:
40 + for logger, level in zip(loggers, previous_levels):
41 + logger.setLevel(level)
42 +
43 +
44 +@dataclass(frozen=True)
45 +class TileLangWarmupTensor:
46 + """Minimal tensor-like object accepted by TileLang compile().
47 +
48 + TileLang builds its cache key and TIR from tensor dtype, shape and stride.
49 + This object deliberately has no storage, so compile-only warmup does not
50 + allocate GPU memory and does not launch the kernel.
51 + """
52 +
53 + dtype: torch.dtype
54 + shape: tuple[int, ...] = (1,)
55 + strides: tuple[int, ...] | None = field(default=None)
56 +
57 + def stride(self) -> tuple[int, ...]:
58 + if self.strides is not None:
59 + return self.strides
60 +
61 + strides: list[int] = []
62 + stride = 1
63 + for size in reversed(self.shape):
64 + strides.append(stride)
65 + stride *= size
66 + return tuple(reversed(strides))
67 +
68 + def new_empty(
69 + self,
70 + shape: tuple[int, ...],
71 + *,
72 + dtype: torch.dtype | None = None,
73 + device: Any = None,
74 + ) -> TileLangWarmupTensor:
75 + return TileLangWarmupTensor(dtype or self.dtype, shape)
76 +
77 +
78 +def make_tilelang_warmup_tensor(
79 + dtype: torch.dtype,
80 + *shape: int,
81 + strides: tuple[int, ...] | None = None,
82 +) -> TileLangWarmupTensor:
83 + return TileLangWarmupTensor(dtype=dtype, shape=tuple(shape), strides=strides)
84 +
85 +
86 +def _tilelang_call_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
87 + call_kwargs = dict(kwargs)
88 + tune_params = call_kwargs.pop("__tune_params", {})
89 + if isinstance(tune_params, dict):
90 + call_kwargs.update(tune_params)
91 + return call_kwargs
92 +
93 +
94 +def compile_tilelang(jit_impl: Any, *args: Any, **kwargs: Any) -> None:
95 + """Compile one TileLang specialization and populate its call cache.
96 +
97 + TileLang's ``compile()`` materializes the kernel without launching it.
98 + We also store the compiled kernel in ``_kernel_cache`` using the same
99 + parsed key as ``__call__`` so runtime does not report a cache miss for an
100 + already materialized specialization.
101 + """
102 +
103 + with _quiet_tilelang_warmup_logs():
104 + compiled = jit_impl.compile(*args, **kwargs)
105 + func = getattr(jit_impl, "func", None)
106 + parse_args = getattr(func, "parse_args", None)
107 + cache = getattr(jit_impl, "_kernel_cache", None)
108 + if not callable(parse_args) or not isinstance(cache, MutableMapping):
109 + return
110 +
111 + key, _ = parse_args(*args, **_tilelang_call_kwargs(kwargs))
112 + cache[key] = compiled
113 +
114 +
115 +class VllmTileLangJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
116 + """TileLang owner whose runtime launch specification is reused for warmup."""
117 +
118 + kernel: ClassVar[Any]
119 + _warming_key: CompileKeyT | None = None
120 +
121 + @abstractmethod
122 + def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
123 + """Return runtime-shaped inputs that reproduce one compile key."""
124 + raise NotImplementedError
125 +
126 + def compile(self, compile_key: CompileKeyT) -> None:
127 + self._warming_key = compile_key
128 + try:
129 + cast(Callable[..., Any], self)(**self.warmup_inputs(compile_key))
130 + finally:
131 + self._warming_key = None
132 +
133 + def launch(
134 + self,
135 + jit_impl: Any,
136 + args: tuple[Any, ...],
137 + kwargs: Mapping[str, Any] | None = None,
138 + ) -> Any:
139 + call_kwargs = dict(kwargs or {})
140 + if self._warming_key is not None:
141 + return compile_tilelang(jit_impl, *args, **call_kwargs)
142 + return jit_impl(*args, **call_kwargs)
143 +
144 +
145 +def kernel_launcher(
146 + call_fn: Callable[..., TileLangLaunchSpec],
147 +) -> Callable[..., Any]:
148 + """Launch TileLang from declarative kernel and runtime argument tuples."""
149 +
150 + @wraps(call_fn)
151 + def wrapper(
152 + self: VllmTileLangJitKernel[Any],
153 + *args: Any,
154 + **kwargs: Any,
155 + ) -> Any:
156 + launch_spec = call_fn(self, *args, **kwargs)
157 + kernel_args, launch_args = launch_spec[:2]
158 + if self._warming_key is not None and kernel_args not in (
159 + (),
160 + (self._warming_key,),
161 + ):
162 + raise RuntimeError(
163 + "TileLang warmup inputs produced a different compile key: "
164 + f"expected {self._warming_key!r}, got kernel arguments "
165 + f"{kernel_args!r}"
166 + )
167 + launch_kwargs = launch_spec[2] if len(launch_spec) > 2 else None
168 + output = launch_spec[3] if len(launch_spec) > 3 else None
169 + result = self.launch(self.kernel(*kernel_args), launch_args, launch_kwargs)
170 + return output if len(launch_spec) > 3 else result
171 +
172 + return wrapper
vllm/model_executor/warmup/kernel_warmup.py
-12
@@ -17,9 +17,6 @@ from vllm.logger import init_logger
17 from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
18 from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
19 from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
20 -from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
21 - deepseek_v4_mhc_warmup,
22 -)
20 from vllm.model_executor.warmup.flashinfer_autotune_cache import (
21 resolve_flashinfer_autotune_file,
22 write_flashinfer_autotune_cache,
@@ -197,15 +194,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
194 compilation_config = worker.vllm_config.compilation_config
195 cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or [])
196
200 - # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
201 - # layer per token; warm them across token sizes first so the first real
202 - # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside).
203 - deepseek_v4_mhc_warmup(
204 - worker.get_model(),
205 - max_tokens=worker.scheduler_config.max_num_batched_tokens,
206 - cudagraph_capture_sizes=cudagraph_capture_sizes,
207 - )
208 -
197 # Run next so input-prep kernels JIT against pristine runner state.
198 if enable_jit_warmup:
199 kimi_k3_triton_warmup(worker)
vllm/models/deepseek_v4/nvidia/model.py
+59
@@ -1210,6 +1210,53 @@ class DeepseekV4DecoderLayer(nn.Module):
1210 requires_grad=False,
1211 )
1212
1213 + if vllm_config.kernel_config.enable_jit_warmup:
1214 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
1215 + _HC_PRENORM_GEMM_TILELANG_KERNEL,
1216 + _MHC_FUSED_TILELANG_KERNEL,
1217 + _MHC_POST_TILELANG_KERNEL,
1218 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
1219 + )
1220 + from vllm.utils.deep_gemm import is_deep_gemm_supported
1221 +
1222 + include_pre_gemm_splits = is_deep_gemm_supported()
1223 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL.register_warmup(
1224 + vllm_config,
1225 + hidden_size=self.hidden_size,
1226 + hc_mult=self.hc_mult,
1227 + use_norm_weight=True,
1228 + include_pre_gemm_splits=include_pre_gemm_splits,
1229 + include_broadcast_splits=(
1230 + get_pp_group().is_first_rank and extract_layer_index(prefix) == 0
1231 + ),
1232 + rms_eps=self.rms_norm_eps,
1233 + hc_pre_eps=self.hc_eps,
1234 + hc_sinkhorn_eps=self.hc_eps,
1235 + hc_post_mult_value=self.hc_post_alpha,
1236 + sinkhorn_repeat=self.hc_sinkhorn_iters,
1237 + norm_eps=(
1238 + self.attn_norm.variance_epsilon,
1239 + self.ffn_norm.variance_epsilon,
1240 + ),
1241 + broadcast_norm_eps=self.attn_norm.variance_epsilon,
1242 + )
1243 + if not include_pre_gemm_splits:
1244 + _HC_PRENORM_GEMM_TILELANG_KERNEL.register_warmup(
1245 + vllm_config,
1246 + hidden_size=self.hidden_size,
1247 + hc_mult=self.hc_mult,
1248 + n_out=self.hc_mult * (2 + self.hc_mult),
1249 + )
1250 + _MHC_POST_TILELANG_KERNEL.register_warmup(
1251 + hidden_size=self.hidden_size,
1252 + hc_mult=self.hc_mult,
1253 + )
1254 + _MHC_FUSED_TILELANG_KERNEL.register_warmup(
1255 + vllm_config,
1256 + hidden_size=self.hidden_size,
1257 + hc_mult=self.hc_mult,
1258 + )
1259 +
1260 def forward(
1261 self,
1262 x: torch.Tensor,
@@ -1402,6 +1449,18 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1449 else:
1450 self._mtp_hidden_buffer = None
1451
1452 + if vllm_config.kernel_config.enable_jit_warmup and get_pp_group().is_last_rank:
1453 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
1454 + _HC_HEAD_FUSED_TILELANG_KERNEL,
1455 + )
1456 +
1457 + _HC_HEAD_FUSED_TILELANG_KERNEL.register_warmup(
1458 + hidden_size=config.hidden_size,
1459 + hc_mult=self.hc_mult,
1460 + rms_eps=self.rms_norm_eps,
1461 + hc_eps=self.hc_eps,
1462 + )
1463 +
1464 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1465 return self.embed_tokens(input_ids)
1466
vllm/models/glm5next/nvidia/model.py
+44
@@ -403,6 +403,50 @@ class Glm5NextDecoderLayer(nn.Module):
403 self.mhc_post_op = MHCPostOp()
404 self.mhc_fused_post_pre_op = MHCFusedPostPreOp()
405
406 + if vllm_config.kernel_config.enable_jit_warmup:
407 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
408 + _HC_PRENORM_GEMM_TILELANG_KERNEL,
409 + _MHC_FUSED_TILELANG_KERNEL,
410 + _MHC_POST_TILELANG_KERNEL,
411 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
412 + )
413 + from vllm.utils.deep_gemm import is_deep_gemm_supported
414 +
415 + include_pre_gemm_splits = is_deep_gemm_supported()
416 + _MHC_PRE_BIG_FUSE_TILELANG_KERNEL.register_warmup(
417 + vllm_config,
418 + hidden_size=self.hidden_size,
419 + hc_mult=self.n,
420 + use_norm_weight=True,
421 + include_pre_gemm_splits=include_pre_gemm_splits,
422 + include_broadcast_splits=False,
423 + rms_eps=self.rms_norm_eps,
424 + hc_pre_eps=self.hc_eps,
425 + hc_sinkhorn_eps=self.hc_eps,
426 + hc_post_mult_value=self.mhc_post_mult_value,
427 + sinkhorn_repeat=self.mhc_sinkhorn_iterations,
428 + norm_eps=(
429 + self.input_layernorm.variance_epsilon,
430 + self.post_attention_layernorm.variance_epsilon,
431 + ),
432 + )
433 + if not include_pre_gemm_splits:
434 + _HC_PRENORM_GEMM_TILELANG_KERNEL.register_warmup(
435 + vllm_config,
436 + hidden_size=self.hidden_size,
437 + hc_mult=self.n,
438 + n_out=self.n * (2 + self.n),
439 + )
440 + _MHC_POST_TILELANG_KERNEL.register_warmup(
441 + hidden_size=self.hidden_size,
442 + hc_mult=self.n,
443 + )
444 + _MHC_FUSED_TILELANG_KERNEL.register_warmup(
445 + vllm_config,
446 + hidden_size=self.hidden_size,
447 + hc_mult=self.n,
448 + )
449 +
450 def forward(
451 self,
452 positions: torch.Tensor,