@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0ba46d4b11

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 17:01 UTC

[ROCm][DSV4] Enable Tilelang MHC replacing torch/triton mhc (#43679)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

TJian committed May 28, 2026 at 15:05 UTC 0ba46d4b11d2b224e39f71af443335dbecf3a878

13 files changed +716 -99

requirements/build/rocm.txt

@@ -16,3 +16,4 @@ wheel

16 jinja2>=3.1.6

17 amdsmi==7.0.2

18 timm>=1.0.17

19 +tilelang==0.1.10

requirements/rocm.txt

@@ -22,3 +22,4 @@ timm>=1.0.17

22 # amd-quark: required for Quark quantization on ROCm

23 # To be consistent with test_quark.py

24 amd-quark>=0.8.99

25 +tilelang==0.1.10

requirements/test/rocm.in

@@ -43,6 +43,7 @@ schemathesis>=3.39.15 # Required for openai schema test

43 # quantization

44 bitsandbytes==0.49.2

45 buildkite-test-collector==0.1.9

46 +tilelang==0.1.10

48 genai_perf>=0.0.8

49 tritonclient>=2.51.0

requirements/test/rocm.txt

+21 -2

@@ -43,7 +43,9 @@ anyio==4.13.0

43 # starlette

44 # watchfiles

45 apache-tvm-ffi==0.1.10

46 - # via xgrammar

46 + # via

47 + # tilelang

48 + # xgrammar

49 arctic-inference==0.1.1

50 # via -r requirements/test/rocm.in

51 argcomplete==3.6.3

@@ -129,7 +131,9 @@ click==8.3.1

131 # typer

132 # uvicorn

133 cloudpickle==3.1.2

132 - # via -r requirements/test/../common.txt

134 + # via

135 + # -r requirements/test/../common.txt

136 + # tilelang

137 colorama==0.4.6

138 # via

139 # perceptron

@@ -511,6 +515,8 @@ mistral-common==1.11.2

515 # -c requirements/common.txt

516 # -r requirements/test/../common.txt

517 # -r requirements/test/rocm.in

518 +ml-dtypes==0.5.4

519 + # via tilelang

520 model-hosting-container-standards==0.1.14

521 # via

522 # -c requirements/common.txt

@@ -587,6 +593,7 @@ numpy==2.2.6

593 # lm-eval

594 # matplotlib

595 # mistral-common

596 + # ml-dtypes

597 # mteb

598 # numba

599 # opencv-python-headless

@@ -610,6 +617,7 @@ numpy==2.2.6

617 # statsmodels

618 # tensorizer

619 # tifffile

620 + # tilelang

621 # torchvision

622 # transformers

623 # tritonclient

@@ -811,6 +819,7 @@ psutil==7.2.2

819 # accelerate

820 # peft

821 # tensorizer

822 + # tilelang

823 py==1.11.0

824 # via pytest-forked

825 py-cpuinfo==9.0.0

@@ -1192,6 +1201,10 @@ tiktoken==0.12.0

1201 # gpt-oss

1202 # lm-eval

1203 # mistral-common

1204 +tilelang==0.1.10

1205 + # via

1206 + # -c requirements/rocm.txt

1207 + # -r requirements/test/rocm.in

1208 timm==1.0.17

1209 # via

1210 # -c requirements/rocm.txt

@@ -1208,6 +1221,8 @@ tomli==2.4.0

1221 # via schemathesis

1222 tomli-w==1.2.0

1223 # via schemathesis

1224 +torch-c-dlpack-ext==0.1.5

1225 + # via tilelang

1226 tqdm==4.67.3

1227 # via

1228 # -r requirements/test/../common.txt

@@ -1225,6 +1240,7 @@ tqdm==4.67.3

1240 # pqdm

1241 # segmentation-models-pytorch

1242 # sentence-transformers

1243 + # tilelang

1244 # transformers

1245 transformers==5.5.3

1246 # via

@@ -1293,6 +1309,7 @@ typing-extensions==4.15.0

1309 # sentence-transformers

1310 # sqlalchemy

1311 # starlette

1312 + # tilelang

1313 # torch

1314 # typeguard

1315 # typing-inspection

@@ -1359,6 +1376,8 @@ yarl==1.23.0

1376 # via

1377 # aiohttp

1378 # schemathesis

1379 +z3-solver==4.15.4.0

1380 + # via tilelang

1381 zipp==3.23.0

1382 # via importlib-metadata

1383

tests/kernels/test_mhc_kernels.py

+166 -2

@@ -4,7 +4,12 @@ import pytest

4 import torch

6 import vllm.model_executor.kernels.mhc # noqa: F401

7 +from vllm.model_executor.kernels.mhc.tilelang import (

8 + _tilelang_hc_prenorm_gemm,

9 + _torch_hc_prenorm_gemm,

10 +)

11 from vllm.platforms import current_platform

12 +from vllm.utils.import_utils import has_tilelang

13 from vllm.utils.torch_utils import set_random_seed

15 DEVICE = current_platform.device_type

@@ -92,8 +97,128 @@ def hc_head_ref(

99 @pytest.mark.skipif(

95 - not current_platform.is_cuda(),

96 - reason="CUDA required",

100 + not (current_platform.is_cuda_alike() and has_tilelang()),

101 + reason="CUDA or ROCm and tilelang required",

102 +)

103 +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])

104 +@pytest.mark.parametrize("hidden_size", [4096, 7168])

105 +@pytest.mark.parametrize("hc_mult", [4])

106 +def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):

107 + torch.set_default_device(DEVICE)

108 + set_random_seed(0)

109 +

110 + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)

111 + hc_mult2 = hc_mult * hc_mult

112 + hc_mult3 = 2 * hc_mult + hc_mult2

113 + fn = (

114 + torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float)

115 + * 1e-4

116 + * (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1))

117 + ).flatten(1, 2)

118 + hc_scale = torch.randn((3,), dtype=torch.float) * 0.1

119 + hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1

120 +

121 + hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6

122 + sinkhorn_repeat = 20

123 + hc_post_alpha = 1.0

124 +

125 + ref = mhc_pre_ref(

126 + residual,

127 + fn,

128 + hc_scale,

129 + hc_base,

130 + rms_eps,

131 + hc_pre_eps,

132 + hc_sinkhorn_eps,

133 + hc_post_alpha,

134 + sinkhorn_repeat,

135 + )

136 + out = torch.ops.vllm.mhc_pre_tilelang(

137 + residual,

138 + fn,

139 + hc_scale,

140 + hc_base,

141 + rms_eps,

142 + hc_pre_eps,

143 + hc_sinkhorn_eps,

144 + hc_post_alpha,

145 + sinkhorn_repeat,

146 + )

147 +

148 + for actual, expected in zip(out, ref, strict=True):

149 + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)

150 +

151 +

152 +@pytest.mark.skipif(

153 + not (current_platform.is_cuda_alike() and has_tilelang()),

154 + reason="CUDA or ROCm and tilelang required",

155 +)

156 +@pytest.mark.parametrize(

157 + ("num_tokens", "hidden_size"),

158 + [

159 + (1, 1280),

160 + (512, 1280),

161 + (2048, 1280),

162 + (1, 4096),

163 + (64, 4096),

164 + (512, 4096),

165 + (2048, 4096),

166 + (1, 7168),

167 + (64, 7168),

168 + (512, 7168),

169 + (2048, 7168),

170 + ],

171 +)

172 +def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):

173 + torch.set_default_device(DEVICE)

174 + set_random_seed(0)

175 +

176 + hc_mult = 4

177 + hc_mult3 = 2 * hc_mult + hc_mult * hc_mult

178 + x = torch.randn((num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16)

179 + fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4

180 + out_ref = torch.empty((1, num_tokens, hc_mult3), dtype=torch.float32)

181 + sqrsum_ref = torch.empty((1, num_tokens), dtype=torch.float32)

182 + out = torch.empty_like(out_ref)

183 + sqrsum = torch.empty_like(sqrsum_ref)

184 +

185 + _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref)

186 + _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)

187 +

188 + torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)

189 + torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4)

190 +

191 +

192 +@pytest.mark.skipif(

193 + not (current_platform.is_cuda_alike() and has_tilelang()),

194 + reason="CUDA or ROCm and tilelang required",

195 +)

196 +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])

197 +@pytest.mark.parametrize("hidden_size", [4096, 7168])

198 +@pytest.mark.parametrize("hc_mult", [4])

199 +def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult):

200 + torch.set_default_device(DEVICE)

201 + set_random_seed(0)

202 +

203 + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)

204 + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)

205 + post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)

206 + comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)

207 +

208 + ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix)

209 + out = torch.ops.vllm.mhc_post_tilelang(

210 + x,

211 + residual,

212 + post_layer_mix,

213 + comb_res_mix,

214 + )

215 +

216 + torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)

217 +

218 +

219 +@pytest.mark.skipif(

220 + not (current_platform.is_cuda_alike() and has_tilelang()),

221 + reason="CUDA or ROCm and tilelang required",

222 )

223 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])

224 @pytest.mark.parametrize("hidden_size", [4096, 7168])

@@ -196,3 +321,42 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):

321

322 out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)

323 torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)

324 +

325 +

326 +@pytest.mark.skipif(

327 + not (current_platform.is_cuda_alike() and has_tilelang()),

328 + reason="CUDA or ROCm and tilelang required",

329 +)

330 +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])

331 +@pytest.mark.parametrize("hidden_size", [4096, 7168])

332 +@pytest.mark.parametrize("hc_mult", [4])

333 +def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):

334 + torch.set_default_device(DEVICE)

335 + set_random_seed(0)

336 +

337 + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)

338 + fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4

339 + hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1

340 + hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1

341 + rms_eps = hc_eps = 1e-6

342 +

343 + out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)

344 + out.fill_(float("nan"))

345 +

346 + result = torch.ops.vllm.hc_head_fused_kernel_tilelang(

347 + residual,

348 + fn,

349 + hc_scale,

350 + hc_base,

351 + out,

352 + hidden_size,

353 + rms_eps,

354 + hc_eps,

355 + hc_mult,

356 + )

357 +

358 + assert result is None

359 + assert not torch.isnan(out).any()

360 +

361 + out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)

362 + torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)

vllm/_tilelang_ops.py

+224 -40

@@ -2,7 +2,7 @@

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

3 import math

4 from functools import cache

5 -from typing import TYPE_CHECKING

5 +from typing import TYPE_CHECKING, Any

7 import torch

@@ -10,8 +10,9 @@ from vllm.platforms import current_platform

10 from vllm.utils.import_utils import has_tilelang

11 from vllm.utils.math_utils import cdiv

13 -# tilelang is only available on CUDA platforms

14 -if TYPE_CHECKING or current_platform.is_cuda():

13 +# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so

14 +# registering the Python wrapper modules does not require TileLang everywhere.

15 +if TYPE_CHECKING or current_platform.is_cuda_alike():

16 if not has_tilelang():

17 raise ImportError(

18 "tilelang is required for mhc but is not installed. Install it with "

@@ -23,6 +24,8 @@ else:

24 tilelang = None # type: ignore[assignment]

25 T = None # type: ignore[assignment]

27 +ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()

28 +

30 @cache

31 def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:

@@ -37,12 +40,17 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:

40 return split_k

43 +pass_configs: dict[tilelang.PassConfigKey, Any] = {

44 + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,

45 + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,

46 +}

47 +

48 +if current_platform.is_cuda():

49 + pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10

50 +

51 +

52 @tilelang.jit(

41 - pass_configs={

42 - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,

43 - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,

44 - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,

45 - },

53 + pass_configs=pass_configs,

54 )

55 def mhc_pre_big_fuse_tilelang(

56 gemm_out_mul,

@@ -78,7 +86,8 @@ def mhc_pre_big_fuse_tilelang(

86 layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]

88 with T.Kernel(num_tokens, threads=96) as i:

81 - T.pdl_sync()

89 + if ENABLE_PDL:

90 + T.pdl_sync()

91 ##################################################################

92 # _pre_norm_fn_fwd_norm

93 rms = T.alloc_fragment(1, T.float32)

@@ -174,18 +183,16 @@ def mhc_pre_big_fuse_tilelang(

183 ol[i1_h] += pre * xl[i_hc, i1_h]

184

185 T.copy(ol, layer_input[i, i0_h * hidden_block])

177 - T.pdl_trigger()

186 +

187 + if ENABLE_PDL:

188 + T.pdl_trigger()

189

190

191 # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478

192

193

194 @tilelang.jit(

184 - pass_configs={

185 - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,

186 - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,

187 - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,

188 - },

195 + pass_configs=pass_configs,

196 )

197 def mhc_pre_big_fuse_with_norm_tilelang(

198 gemm_out_mul,

@@ -230,7 +237,8 @@ def mhc_pre_big_fuse_with_norm_tilelang(

237 T.clear(mixes)

238 rms[0] = 0

239

233 - T.pdl_sync()

240 + if ENABLE_PDL:

241 + T.pdl_sync()

242

243 for i_split in T.serial(n_splits):

244 rms[0] += gemm_out_sqrsum[i_split, i]

@@ -341,15 +349,12 @@ def mhc_pre_big_fuse_with_norm_tilelang(

349

350 T.copy(ol, layer_input[i, i0_h * hidden_block])

351

344 - T.pdl_trigger()

352 + if ENABLE_PDL:

353 + T.pdl_trigger()

354

355

356 @tilelang.jit(

348 - pass_configs={

349 - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,

350 - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,

351 - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,

352 - },

357 + pass_configs=pass_configs,

358 )

359 def mhc_fused_tilelang(

360 comb_mix,

@@ -390,8 +395,8 @@ def mhc_fused_tilelang(

395

396 with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):

397 tid = T.get_thread_binding()

393 - warp_id = T.get_warp_idx()

394 - lane = T.get_lane_idx()

398 + warp_id = tid // 32

399 + lane = tid % 32

400

401 s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)

402 s_post = T.alloc_shared((hc,), T.float32)

@@ -407,7 +412,8 @@ def mhc_fused_tilelang(

412 T.clear(sqr)

413 h_split_start = i_ks * h_per_split

414

410 - T.pdl_sync()

415 + if ENABLE_PDL:

416 + T.pdl_sync()

417

418 T.copy(post_mix[i_n, 0], s_post)

419 T.copy(comb_mix[i_n, 0, 0], s_comb)

@@ -466,15 +472,12 @@ def mhc_fused_tilelang(

472 v2 += s_warp[w, tile_n]

473 rp_out[i_ks, i_n] = v2

474

469 - T.pdl_trigger()

475 + if ENABLE_PDL:

476 + T.pdl_trigger()

477

478

479 @tilelang.jit(

473 - pass_configs={

474 - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,

475 - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,

476 - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,

477 - },

480 + pass_configs=pass_configs,

481 )

482 def mhc_post_tilelang(

483 a,

@@ -507,7 +510,8 @@ def mhc_post_tilelang(

510

511 a_local = T.alloc_fragment((hc, hc), T.float32)

512 c_local = T.alloc_fragment(hc, T.float32)

510 - T.pdl_sync()

513 + if ENABLE_PDL:

514 + T.pdl_sync()

515 T.copy(a[i_n, 0, 0], a_local)

516 T.copy(c[i_n, 0], c_local)

517

@@ -523,15 +527,193 @@ def mhc_post_tilelang(

527 x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]

528

529 T.copy(x_local, x[i_n, 0, i0_h * h_blk])

526 - T.pdl_trigger()

530 + if ENABLE_PDL:

531 + T.pdl_trigger()

532 +

533 +

534 +@tilelang.jit(

535 + pass_configs=pass_configs,

536 +)

537 +def hc_prenorm_gemm_tilelang(

538 + x,

539 + fn,

540 + out,

541 + sqrsum,

542 + hidden_size: int,

543 + hc_mult: int = 4,

544 + n_out: int = 24,

545 + n_thr: int = 512,

546 + tile_n: int = 12,

547 + n_splits: int = 1,

548 +) -> tilelang.JITKernel:

549 + num_tokens = T.dynamic("num_tokens")

550 + hc_hidden_size = hc_mult * hidden_size

551 + k_per_split = hc_hidden_size // n_splits

552 + k_iters = k_per_split // n_thr

553 + n_tiles = T.ceildiv(n_out, tile_n)

554 +

555 + x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type]

556 + fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type]

557 + out: T.Tensor((n_splits, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type]

558 + sqrsum: T.Tensor((n_splits, num_tokens), T.float32) # type: ignore[no-redef, valid-type]

559 +

560 + with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as (

561 + i_n,

562 + i_t,

563 + i_s,

564 + ):

565 + tid = T.get_thread_binding()

566 + acc = T.alloc_local((tile_n,), T.float32)

567 + sqr = T.alloc_local((1,), T.float32)

568 + T.clear(acc)

569 + T.clear(sqr)

570 +

571 + if ENABLE_PDL:

572 + T.pdl_sync()

573 +

574 + for it in T.serial(k_iters):

575 + i_k = i_s * k_per_split + it * n_thr + tid

576 + x_val = x[i_n, i_k]

577 + for i_o in T.unroll(tile_n):

578 + out_idx = i_t * tile_n + i_o

579 + if out_idx < n_out:

580 + acc[i_o] += x_val * fn[out_idx, i_k]

581 + if i_t == 0:

582 + sqr[0] += x_val * x_val

583 +

584 + for i_o in T.unroll(tile_n):

585 + acc[i_o] = T.warp_reduce_sum(acc[i_o])

586 + if i_t == 0:

587 + sqr[0] = T.warp_reduce_sum(sqr[0])

588 +

589 + lane = tid % 32

590 + warp_id = tid // 32

591 + num_warps = n_thr // 32

592 + warp_acc = T.alloc_shared((num_warps, tile_n), T.float32)

593 + warp_sqr = T.alloc_shared(num_warps, T.float32)

594 +

595 + if lane == 0:

596 + for i_o in T.unroll(tile_n):

597 + warp_acc[warp_id, i_o] = acc[i_o]

598 + if i_t == 0:

599 + warp_sqr[warp_id] = sqr[0]

600 + T.sync_threads()

601 +

602 + if warp_id == 0:

603 + if lane < tile_n:

604 + reduced_acc = T.alloc_var(T.float32, init=0.0)

605 + for i_w in T.unroll(num_warps):

606 + reduced_acc += warp_acc[i_w, lane]

607 + out_idx = i_t * tile_n + lane

608 + if out_idx < n_out:

609 + out[i_s, i_n, out_idx] = reduced_acc

610 + if lane == 0 and i_t == 0:

611 + reduced_sqr = T.alloc_var(T.float32, init=0.0)

612 + for i_w in T.unroll(num_warps):

613 + reduced_sqr += warp_sqr[i_w]

614 + sqrsum[i_s, i_n] = reduced_sqr

615 +

616 + if ENABLE_PDL:

617 + T.pdl_trigger()

618 +

619 +

620 +@tilelang.jit(

621 + pass_configs=pass_configs,

622 +)

623 +def hc_prenorm_gemm_block_m_tilelang(

624 + x,

625 + fn,

626 + out,

627 + sqrsum,

628 + hidden_size: int,

629 + hc_mult: int = 4,

630 + n_out: int = 24,

631 + n_thr: int = 512,

632 + tile_n: int = 12,

633 + block_m: int = 2,

634 +) -> tilelang.JITKernel:

635 + num_tokens = T.dynamic("num_tokens")

636 + hc_hidden_size = hc_mult * hidden_size

637 + k_iters = hc_hidden_size // n_thr

638 + n_tiles = T.ceildiv(n_out, tile_n)

639 + m_tiles = T.ceildiv(num_tokens, block_m)

640 +

641 + x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type]

642 + fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type]

643 + out: T.Tensor((1, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type]

644 + sqrsum: T.Tensor((1, num_tokens), T.float32) # type: ignore[no-redef, valid-type]

645 +

646 + with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t):

647 + tid = T.get_thread_binding()

648 + acc = T.alloc_local((block_m, tile_n), T.float32)

649 + sqr = T.alloc_local((block_m,), T.float32)

650 + T.clear(acc)

651 + T.clear(sqr)

652 +

653 + if ENABLE_PDL:

654 + T.pdl_sync()

655 +

656 + for it in T.serial(k_iters):

657 + i_k = it * n_thr + tid

658 + fn_val = T.alloc_local((tile_n,), T.float32)

659 + for i_o in T.unroll(tile_n):

660 + out_idx = i_t * tile_n + i_o

661 + if out_idx < n_out:

662 + fn_val[i_o] = fn[out_idx, i_k]

663 + else:

664 + fn_val[i_o] = 0.0

665 + for i_m in T.unroll(block_m):

666 + token_idx = i_mt * block_m + i_m

667 + if token_idx < num_tokens:

668 + x_val = x[token_idx, i_k]

669 + for i_o in T.unroll(tile_n):

670 + acc[i_m, i_o] += x_val * fn_val[i_o]

671 + if i_t == 0:

672 + sqr[i_m] += x_val * x_val

673 +

674 + for i_m in T.unroll(block_m):

675 + for i_o in T.unroll(tile_n):

676 + acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o])

677 + if i_t == 0:

678 + sqr[i_m] = T.warp_reduce_sum(sqr[i_m])

679 +

680 + lane = tid % 32

681 + warp_id = tid // 32

682 + num_warps = n_thr // 32

683 + warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32)

684 + warp_sqr = T.alloc_shared((num_warps, block_m), T.float32)

685 +

686 + if lane == 0:

687 + for i_m in T.unroll(block_m):

688 + for i_o in T.unroll(tile_n):

689 + warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o]

690 + if i_t == 0:

691 + warp_sqr[warp_id, i_m] = sqr[i_m]

692 + T.sync_threads()

693 +

694 + if warp_id == 0:

695 + for i_m in T.unroll(block_m):

696 + token_idx = i_mt * block_m + i_m

697 + if token_idx < num_tokens:

698 + if lane < tile_n:

699 + reduced_acc = T.alloc_var(T.float32, init=0.0)

700 + for i_w in T.unroll(num_warps):

701 + reduced_acc += warp_acc[i_w, i_m, lane]

702 + out_idx = i_t * tile_n + lane

703 + if out_idx < n_out:

704 + out[0, token_idx, out_idx] = reduced_acc

705 + if lane == 0 and i_t == 0:

706 + reduced_sqr = T.alloc_var(T.float32, init=0.0)

707 + for i_w in T.unroll(num_warps):

708 + reduced_sqr += warp_sqr[i_w, i_m]

709 + sqrsum[0, token_idx] = reduced_sqr

710 +

711 + if ENABLE_PDL:

712 + T.pdl_trigger()

713

714

715 @tilelang.jit(

530 - pass_configs={

531 - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,

532 - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,

533 - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,

534 - },

716 + pass_configs=pass_configs,

717 )

718 def hc_head_fuse_tilelang(

719 residual,

@@ -566,7 +748,8 @@ def hc_head_fuse_tilelang(

748 out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]

749

750 with T.Kernel(num_tokens, threads=n_thr) as i:

569 - T.pdl_sync()

751 + if ENABLE_PDL:

752 + T.pdl_sync()

753

754 # ------------------------------------------------------------------

755 # Pass 1 – for each residual channel m_c and h_block:

@@ -624,4 +807,5 @@ def hc_head_fuse_tilelang(

807

808 T.copy(ol, out[i, i0_h * h_block], disable_tma=True)

809

627 - T.pdl_trigger()

810 + if ENABLE_PDL:

811 + T.pdl_trigger()

vllm/model_executor/kernels/mhc/tilelang.py

+145 -27

@@ -5,6 +5,88 @@ import torch

5 from vllm.utils.torch_utils import direct_register_custom_op

8 +def _torch_hc_prenorm_gemm(

9 + x: torch.Tensor,

10 + fn: torch.Tensor,

11 + out: torch.Tensor,

12 + sqrsum: torch.Tensor,

13 +) -> None:

14 + assert out.shape[0] == 1

15 + assert sqrsum.shape[0] == 1

16 + x_float = x.float()

17 + out[0].copy_(x_float @ fn.t())

18 + sqrsum[0].copy_(x_float.square().sum(dim=-1))

19 +

20 +

21 +def _tilelang_hc_prenorm_gemm(

22 + x: torch.Tensor,

23 + fn: torch.Tensor,

24 + out: torch.Tensor,

25 + sqrsum: torch.Tensor,

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:

32 + from vllm._tilelang_ops 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,

55 + )

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(

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,

74 + )

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 + )

88 +

89 +

90 def mhc_pre_tilelang(

91 residual: torch.Tensor,

92 fn: torch.Tensor,

@@ -80,10 +162,16 @@ def mhc_pre_tilelang(

162 residual_flat = residual.view(-1, hc_mult, hidden_size)

163 num_tokens = residual_flat.shape[0]

164

83 - # these numbers are from deepgemm kernel impl

84 - block_k = 64

85 - block_m = 64

86 - n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))

165 + from vllm.utils.deep_gemm import is_deep_gemm_supported

166 +

167 + use_deep_gemm = is_deep_gemm_supported()

168 + if use_deep_gemm:

169 + # these numbers are from deepgemm kernel impl

170 + block_k = 64

171 + block_m = 64

172 + n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))

173 + else:

174 + n_splits = 1

175

176 post_mix = torch.empty(

177 num_tokens, hc_mult, dtype=torch.float32, device=residual.device

@@ -102,13 +190,24 @@ def mhc_pre_tilelang(

190 n_splits, num_tokens, dtype=torch.float32, device=residual.device

191 )

192

105 - tf32_hc_prenorm_gemm(

106 - residual_flat.view(num_tokens, hc_mult * hidden_size),

107 - fn,

108 - gemm_out_mul,

109 - gemm_out_sqrsum,

110 - n_splits,

111 - )

193 + residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size)

194 + if use_deep_gemm:

195 + tf32_hc_prenorm_gemm(

196 + residual_2d,

197 + fn,

198 + gemm_out_mul,

199 + gemm_out_sqrsum,

200 + n_splits,

201 + )

202 + else:

203 + _tilelang_hc_prenorm_gemm(

204 + residual_2d,

205 + fn,

206 + gemm_out_mul,

207 + gemm_out_sqrsum,

208 + hidden_size,

209 + hc_mult,

210 + )

211

212 if norm_weight is None:

213 mhc_pre_big_fuse_tilelang(

@@ -304,16 +403,24 @@ def mhc_fused_post_pre_tilelang(

403 post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)

404 comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)

405

307 - fma_token_threshold = 16

308 - if num_tokens <= fma_token_threshold:

406 + from vllm.utils.deep_gemm import is_deep_gemm_supported

407 +

408 + use_deep_gemm = is_deep_gemm_supported()

409 + use_small_fma = num_tokens <= 16

410 + if use_small_fma:

411 # TODO(gnovack): investigate autotuning these heuristics

412 tile_n = 2 if num_tokens < 8 else 3

413 n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4

414 else:

313 - # these number are from deepgemm kernel impl

314 - block_k = 64

315 - block_m = 64

316 - n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))

415 + if use_deep_gemm:

416 + # these number are from deepgemm kernel impl

417 + block_k = 64

418 + block_m = 64

419 + n_splits = compute_num_split(

420 + block_k, hc_hidden_size, cdiv(num_tokens, block_m)

421 + )

422 + else:

423 + n_splits = 1

424

425 gemm_out_mul = torch.empty(

426 n_splits,

@@ -348,7 +455,7 @@ def mhc_fused_post_pre_tilelang(

455 device=residual.device,

456 )

457

351 - if num_tokens <= fma_token_threshold:

458 + if use_small_fma:

459 mhc_fused_tilelang(

460 comb_res_mix_flat,

461 residual_flat,

@@ -375,15 +482,26 @@ def mhc_fused_post_pre_tilelang(

482 residual.shape[-1],

483 )

484

378 - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm

379 -

380 - tf32_hc_prenorm_gemm(

381 - residual_cur.view(num_tokens, hc_mult * hidden_size),

382 - fn,

383 - gemm_out_mul,

384 - gemm_out_sqrsum,

385 - n_splits,

386 - )

485 + residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size)

486 + if use_deep_gemm:

487 + from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm

488 +

489 + tf32_hc_prenorm_gemm(

490 + residual_cur_2d,

491 + fn,

492 + gemm_out_mul,

493 + gemm_out_sqrsum,

494 + n_splits,

495 + )

496 + else:

497 + _tilelang_hc_prenorm_gemm(

498 + residual_cur_2d,

499 + fn,

500 + gemm_out_mul,

501 + gemm_out_sqrsum,

502 + hidden_size,

503 + hc_mult,

504 + )

505

506 if norm_weight is None:

507 mhc_pre_big_fuse_tilelang(

vllm/model_executor/layers/mhc.py

+126 -20

@@ -3,8 +3,12 @@

3 import torch

5 # this import will also register the custom ops

6 +# import vllm.model_executor.kernels.mhc # noqa: F401

7 import vllm.model_executor.kernels.mhc as mhc_kernels

8 from vllm.model_executor.custom_op import CustomOp

9 +from vllm.utils.import_utils import has_tilelang

10 +

11 +HAS_TILELANG = has_tilelang()

14 # --8<-- [start:mhc_pre]

@@ -85,6 +89,52 @@ class MHCPreOp(CustomOp):

89 # sinkhorn_repeat,

90 # )

91 # else:

92 + if HAS_TILELANG:

93 + return torch.ops.vllm.mhc_pre_tilelang(

94 + residual,

95 + fn,

96 + hc_scale,

97 + hc_base,

98 + rms_eps,

99 + hc_pre_eps,

100 + hc_sinkhorn_eps,

101 + hc_post_mult_value,

102 + sinkhorn_repeat,

103 + n_splits,

104 + norm_weight,

105 + norm_eps,

106 + )

107 + else:

108 + return self.forward_native(

109 + residual,

110 + fn,

111 + hc_scale,

112 + hc_base,

113 + rms_eps,

114 + hc_pre_eps,

115 + hc_sinkhorn_eps,

116 + hc_post_mult_value,

117 + sinkhorn_repeat,

118 + n_splits,

119 + norm_weight,

120 + norm_eps,

121 + )

122 +

123 + def forward_native(

124 + self,

125 + residual: torch.Tensor,

126 + fn: torch.Tensor,

127 + hc_scale: torch.Tensor,

128 + hc_base: torch.Tensor,

129 + rms_eps: float,

130 + hc_pre_eps: float,

131 + hc_sinkhorn_eps: float,

132 + hc_post_mult_value: float,

133 + sinkhorn_repeat: int,

134 + n_splits: int = 1,

135 + norm_weight: torch.Tensor | None = None,

136 + norm_eps: float = 0.0,

137 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

138 return mhc_kernels.mhc_pre_torch(

139 residual,

140 fn,

@@ -97,9 +147,6 @@ class MHCPreOp(CustomOp):

147 sinkhorn_repeat,

148 )

149

100 - def forward_native(self, *args, **kwargs):

101 - raise NotImplementedError("Native implementation of mhc_pre is not available")

102 -

150

151 # --8<-- [start:mhc_post]

152 @CustomOp.register("mhc_post")

@@ -147,6 +194,20 @@ class MHCPostOp(CustomOp):

194 # comb_res_mix,

195 # )

196 # else:

197 + if HAS_TILELANG:

198 + return torch.ops.vllm.mhc_post_tilelang(

199 + x, residual, post_layer_mix, comb_res_mix

200 + )

201 + else:

202 + return self.forward_native(x, residual, post_layer_mix, comb_res_mix)

203 +

204 + def forward_native(

205 + self,

206 + x: torch.Tensor,

207 + residual: torch.Tensor,

208 + post_layer_mix: torch.Tensor,

209 + comb_res_mix: torch.Tensor,

210 + ) -> torch.Tensor:

211 return mhc_kernels.mhc_post_torch(

212 x,

213 residual,

@@ -154,9 +215,6 @@ class MHCPostOp(CustomOp):

215 comb_res_mix,

216 )

217

157 - def forward_native(self, *args, **kwargs):

158 - raise NotImplementedError("Native implementation of mhc_post is not available")

159 -

218

219 # --8<-- [start:hc_head]

220 @CustomOp.register("hc_head")

@@ -220,17 +278,32 @@ class HCHeadOp(CustomOp):

278 out = torch.empty(

279 num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device

280 )

223 - torch.ops.vllm.hc_head_triton(

224 - hs_flat,

225 - hc_fn,

226 - hc_scale,

227 - hc_base,

228 - out,

229 - hidden_size,

230 - rms_norm_eps,

231 - hc_eps,

232 - hc_mult,

233 - )

281 +

282 + if HAS_TILELANG:

283 + torch.ops.vllm.hc_head_fused_kernel_tilelang(

284 + hs_flat,

285 + hc_fn,

286 + hc_scale,

287 + hc_base,

288 + out,

289 + hidden_size,

290 + rms_norm_eps,

291 + hc_eps,

292 + hc_mult,

293 + )

294 + else:

295 + torch.ops.vllm.hc_head_triton(

296 + hs_flat,

297 + hc_fn,

298 + hc_scale,

299 + hc_base,

300 + out,

301 + hidden_size,

302 + rms_norm_eps,

303 + hc_eps,

304 + hc_mult,

305 + )

306 +

307 return out.view(*outer_shape, hidden_size)

308

309 def forward_native(self, *args, **kwargs):

@@ -290,9 +363,42 @@ class MHCFusedPostPreOp(CustomOp):

363 norm_eps,

364 )

365

293 - def forward_hip(self, *args, **kwargs):

294 - raise NotImplementedError(

295 - "Hip implementation of mhc_fused_post_pre is not available"

366 + def forward_hip(

367 + self,

368 + x: torch.Tensor,

369 + residual: torch.Tensor,

370 + post_layer_mix: torch.Tensor,

371 + comb_res_mix: torch.Tensor,

372 + fn: torch.Tensor,

373 + hc_scale: torch.Tensor,

374 + hc_base: torch.Tensor,

375 + rms_eps: float,

376 + hc_pre_eps: float,

377 + hc_sinkhorn_eps: float,

378 + hc_post_mult_value: float,

379 + sinkhorn_repeat: int,

380 + n_splits: int = 1,

381 + tile_n: int = 1,

382 + norm_weight: torch.Tensor | None = None,

383 + norm_eps: float = 0.0,

384 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

385 + return torch.ops.vllm.mhc_fused_post_pre_tilelang(

386 + x,

387 + residual,

388 + post_layer_mix,

389 + comb_res_mix,

390 + fn,

391 + hc_scale,

392 + hc_base,

393 + rms_eps,

394 + hc_pre_eps,

395 + hc_sinkhorn_eps,

396 + hc_post_mult_value,

397 + sinkhorn_repeat,

398 + n_splits,

399 + tile_n,

400 + norm_weight,

401 + norm_eps,

402 )

403

404 def forward_native(self, *args, **kwargs):

vllm/models/deepseek_v4/amd/model.py

+11 -7

@@ -54,6 +54,7 @@ from vllm.models.deepseek_v4.attention import (

54 )

55 from vllm.platforms import current_platform

56 from vllm.sequence import IntermediateTensors

57 +from vllm.utils.import_utils import has_tilelang

60 class DeepseekV4MLP(nn.Module):

@@ -473,6 +474,7 @@ class DeepseekV4DecoderLayer(nn.Module):

474 self.mhc_pre = MHCPreOp()

475 self.mhc_post = MHCPostOp()

476 self.mhc_fused_post_pre = MHCFusedPostPreOp()

477 + self.has_tilelang = has_tilelang()

478

479 def hc_pre(

480 self,

@@ -503,7 +505,7 @@ class DeepseekV4DecoderLayer(nn.Module):

505 ):

506 return self.mhc_post(x, residual, post, comb)

507

506 - def _forward_cuda(

508 + def _forward_fused_post_pre(

509 self,

510 x: torch.Tensor,

511 positions: torch.Tensor,

@@ -555,7 +557,7 @@ class DeepseekV4DecoderLayer(nn.Module):

557 x = self.ffn(x, input_ids)

558 return x, residual, post_mix, res_mix

559

558 - def _forward_rocm(

560 + def _forward_unfused_post_pre(

561 self,

562 x: torch.Tensor,

563 positions: torch.Tensor,

@@ -594,12 +596,13 @@ class DeepseekV4DecoderLayer(nn.Module):

596 ) -> tuple[

597 torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None

598 ]:

597 - if current_platform.is_rocm():

598 - return self._forward_rocm(

599 + if not self.has_tilelang:

600 + return self._forward_unfused_post_pre(

601 x, positions, input_ids, post_mix, res_mix, residual

602 )

601 -

602 - return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)

603 + return self._forward_fused_post_pre(

604 + x, positions, input_ids, post_mix, res_mix, residual

605 + )

606

607

608 @support_torch_compile

@@ -682,6 +685,7 @@ class DeepseekV4Model(nn.Module):

685 requires_grad=False,

686 )

687 self.hc_head_op = HCHeadOp()

688 + self.has_tilelang = has_tilelang()

689 # Pre-hc_head residual stream buffer for the MTP draft. Stable

690 # address (outside the cudagraph pool) so the copy_ in forward()

691 # refreshes it correctly across captured shapes.

@@ -748,7 +752,7 @@ class DeepseekV4Model(nn.Module):

752 res_mix,

753 residual,

754 )

751 - if layer is not None and current_platform.is_cuda():

755 + if layer is not None and self.has_tilelang:

756 hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)

757

758 if not get_pp_group().is_last_rank:

vllm/models/deepseek_v4/amd/mtp.py

+3 -1

@@ -39,6 +39,7 @@ from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weigh

39 from vllm.model_executor.models.utils import maybe_prefix

40 from vllm.platforms import current_platform

41 from vllm.sequence import IntermediateTensors

42 +from vllm.utils.import_utils import has_tilelang

44 from .model import DeepseekV4DecoderLayer

@@ -118,6 +119,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):

119 )

120

121 self.hc_head_op = HCHeadOp()

122 + self.has_tilelang = has_tilelang()

123

124 def forward(

125 self,

@@ -144,7 +146,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):

146 hidden_states, residual, post_mix, res_mix = self.mtp_block(

147 positions=positions, x=hidden_states, input_ids=None

148 )

147 - if current_platform.is_cuda():

149 + if self.has_tilelang:

150 hidden_states = self.mtp_block.hc_post(

151 hidden_states, residual, post_mix, res_mix

152 )

vllm/platforms/cuda.py

@@ -592,6 +592,15 @@ class CudaPlatformBase(Platform):

592 default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm

593 )

594

595 + @classmethod

596 + def is_arch_support_pdl(cls) -> bool:

597 + try:

598 + device = torch.cuda.current_device()

599 + major, _ = torch.cuda.get_device_capability(device)

600 + except Exception:

601 + return False

602 + return major >= 9

603 +

604

605 # NVML utils

606 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,

vllm/platforms/interface.py

@@ -1016,6 +1016,13 @@ class Platform:

1016 # Native always used by default. Platforms can override this behavior.

1017 return IrOpPriorityConfig.with_default(["native"])

1018

1019 + @classmethod

1020 + def is_arch_support_pdl(cls) -> bool:

1021 + """

1022 + Does the current platform support PDL (Programmatic Dependent Launch)?

1023 + """

1024 + return False

1025 +

1026

1027 class UnspecifiedPlatform(Platform):

1028 _enum = PlatformEnum.UNSPECIFIED

vllm/utils/import_utils.py

@@ -430,6 +430,7 @@ def has_triton_kernels() -> bool:

430 return is_available

431

432

433 +@cache

434 def has_tilelang() -> bool:

435 """Whether the optional `tilelang` package is available."""

436 return _has_module("tilelang")