| 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 | mhc_fused_post_pre_split_config, |
| 21 | require_fused_post_pre_config, |
| 22 | ) |
| 23 | from vllm.model_executor.warmup import jit_warmup_tilelang_helper |
| 24 | |
| 25 | |
| 26 | @pytest.mark.parametrize( |
| 27 | ("kwargs", "expected"), |
| 28 | [ |
| 29 | ( |
| 30 | dict( |
| 31 | num_tokens=64, |
| 32 | hc_hidden_size=4096, |
| 33 | hidden_size=2048, |
| 34 | hc_mult=2, |
| 35 | n_out=128, |
| 36 | ), |
| 37 | (2048, 2, 128, 1024, 4, 1, False, 1), |
| 38 | ), |
| 39 | ( |
| 40 | dict( |
| 41 | num_tokens=1024, |
| 42 | hc_hidden_size=4096, |
| 43 | hidden_size=2048, |
| 44 | hc_mult=2, |
| 45 | n_out=128, |
| 46 | ), |
| 47 | (2048, 2, 128, 512, 12, 1, True, 2), |
| 48 | ), |
| 49 | ( |
| 50 | dict( |
| 51 | num_tokens=64, |
| 52 | hc_hidden_size=4096, |
| 53 | hidden_size=2048, |
| 54 | hc_mult=2, |
| 55 | n_out=128, |
| 56 | n_thr=256, |
| 57 | tile_n=8, |
| 58 | n_splits=4, |
| 59 | ), |
| 60 | (2048, 2, 128, 256, 8, 4, False, 1), |
| 61 | ), |
| 62 | ], |
| 63 | ) |
| 64 | def test_hc_prenorm_gemm_dispatch_matches_legacy_runtime_config( |
| 65 | kwargs: dict[str, Any], |
| 66 | expected: tuple[int, int, int, int, int, int, bool, int], |
| 67 | ) -> None: |
| 68 | kernel = HcPrenormGemmTileLangKernel() |
| 69 | |
| 70 | assert kernel.dispatch(**kwargs) == kernel.CompileKey(*expected) |
| 71 | |
| 72 | |
| 73 | @pytest.mark.parametrize( |
| 74 | ("is_broadcast", "use_norm_weight", "expected_use_norm", "expected_eps"), |
| 75 | [ |
| 76 | (False, False, False, 0.0), |
| 77 | (False, True, True, 1.0e-5), |
| 78 | (True, False, True, 2.0e-5), |
| 79 | ], |
| 80 | ) |
| 81 | def test_mhc_pre_big_fuse_dispatch_matches_legacy_runtime_config( |
| 82 | is_broadcast: bool, |
| 83 | use_norm_weight: bool, |
| 84 | expected_use_norm: bool, |
| 85 | expected_eps: float, |
| 86 | ) -> None: |
| 87 | kernel = MhcPreBigFuseTileLangKernel() |
| 88 | |
| 89 | assert kernel.dispatch( |
| 90 | hidden_size=4096, |
| 91 | hc_mult=4, |
| 92 | n_splits=2, |
| 93 | is_broadcast=is_broadcast, |
| 94 | use_norm_weight=use_norm_weight, |
| 95 | rms_eps=1.0e-6, |
| 96 | hc_pre_eps=2.0e-6, |
| 97 | hc_sinkhorn_eps=3.0e-6, |
| 98 | hc_post_mult_value=0.5, |
| 99 | sinkhorn_repeat=3, |
| 100 | norm_eps=1.0e-5, |
| 101 | broadcast_norm_eps=2.0e-5, |
| 102 | ) == kernel.CompileKey( |
| 103 | hidden_size=4096, |
| 104 | hc_mult=4, |
| 105 | n_splits=2, |
| 106 | use_norm_weight=expected_use_norm, |
| 107 | is_broadcast=is_broadcast, |
| 108 | rms_eps=1.0e-6, |
| 109 | hc_pre_eps=2.0e-6, |
| 110 | hc_sinkhorn_eps=3.0e-6, |
| 111 | hc_post_mult_value=0.5, |
| 112 | sinkhorn_repeat=3, |
| 113 | norm_eps=expected_eps, |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | @pytest.mark.parametrize( |
| 118 | ("num_tokens", "hidden_size", "expected_tile_n", "expected_n_splits"), |
| 119 | [ |
| 120 | (1, 4096, 2, 8), |
| 121 | (4, 8192, 2, 8), |
| 122 | (8, 5120, 2, 8), |
| 123 | (16, 5120, 6, 8), |
| 124 | (32, 7168, 6, 8), |
| 125 | ], |
| 126 | ) |
| 127 | def test_mhc_fused_dispatch_matches_the_launch_config( |
| 128 | num_tokens: int, |
| 129 | hidden_size: int, |
| 130 | expected_tile_n: int, |
| 131 | expected_n_splits: int, |
| 132 | ) -> None: |
| 133 | """The compile key must be the config the launch will actually use. |
| 134 | |
| 135 | Both read mhc_fused_post_pre_split_config, so this pins the tuned bands |
| 136 | and guards against the two drifting apart again. |
| 137 | """ |
| 138 | kernel = MhcFusedTileLangKernel() |
| 139 | tile_n, n_splits, n_thr = require_fused_post_pre_config( |
| 140 | num_tokens, hidden_size, hc_mult=4 |
| 141 | ) |
| 142 | |
| 143 | assert (tile_n, n_splits) == (expected_tile_n, expected_n_splits) |
| 144 | assert kernel.dispatch( |
| 145 | num_tokens=num_tokens, |
| 146 | hidden_size=hidden_size, |
| 147 | hc_mult=4, |
| 148 | ) == kernel.CompileKey( |
| 149 | hidden_size=hidden_size, |
| 150 | hc_mult=4, |
| 151 | n_splits=n_splits, |
| 152 | tile_n=tile_n, |
| 153 | n_thr=n_thr, |
| 154 | ) |
| 155 | |
| 156 | |
| 157 | def test_mhc_fused_declines_shapes_it_cannot_tile() -> None: |
| 158 | """Above the token cutoff, and for a hidden size the block cannot split.""" |
| 159 | assert mhc_fused_post_pre_split_config(33, 5120, 4) is None |
| 160 | assert mhc_fused_post_pre_split_config(1, 5137, 4) is None |
| 161 | with pytest.raises(ValueError, match="does not cover num_tokens"): |
| 162 | require_fused_post_pre_config(33, 5120, 4) |
| 163 | |
| 164 | |
| 165 | @pytest.mark.parametrize( |
| 166 | ("kernel", "compile_key"), |
| 167 | [ |
| 168 | ( |
| 169 | HcPrenormGemmTileLangKernel(), |
| 170 | HcPrenormGemmTileLangKernel.CompileKey( |
| 171 | hidden_size=2048, |
| 172 | hc_mult=2, |
| 173 | n_out=128, |
| 174 | n_thr=1024, |
| 175 | tile_n=4, |
| 176 | n_splits=1, |
| 177 | use_block_m=False, |
| 178 | block_m=1, |
| 179 | ), |
| 180 | ), |
| 181 | ( |
| 182 | HcPrenormGemmTileLangKernel(), |
| 183 | HcPrenormGemmTileLangKernel.CompileKey( |
| 184 | hidden_size=2048, |
| 185 | hc_mult=2, |
| 186 | n_out=128, |
| 187 | n_thr=512, |
| 188 | tile_n=12, |
| 189 | n_splits=1, |
| 190 | use_block_m=False, |
| 191 | block_m=1, |
| 192 | ), |
| 193 | ), |
| 194 | ( |
| 195 | HcPrenormGemmTileLangKernel(), |
| 196 | HcPrenormGemmTileLangKernel.CompileKey( |
| 197 | hidden_size=2048, |
| 198 | hc_mult=2, |
| 199 | n_out=128, |
| 200 | n_thr=512, |
| 201 | tile_n=12, |
| 202 | n_splits=1, |
| 203 | use_block_m=True, |
| 204 | block_m=2, |
| 205 | ), |
| 206 | ), |
| 207 | ( |
| 208 | MhcPreBigFuseTileLangKernel(), |
| 209 | MhcPreBigFuseTileLangKernel.CompileKey( |
| 210 | hidden_size=4096, |
| 211 | hc_mult=4, |
| 212 | n_splits=2, |
| 213 | use_norm_weight=True, |
| 214 | is_broadcast=False, |
| 215 | rms_eps=1.0e-6, |
| 216 | hc_pre_eps=2.0e-6, |
| 217 | hc_sinkhorn_eps=3.0e-6, |
| 218 | hc_post_mult_value=0.5, |
| 219 | sinkhorn_repeat=3, |
| 220 | norm_eps=1.0e-5, |
| 221 | ), |
| 222 | ), |
| 223 | ( |
| 224 | MhcPreBigFuseTileLangKernel(), |
| 225 | MhcPreBigFuseTileLangKernel.CompileKey( |
| 226 | hidden_size=4096, |
| 227 | hc_mult=4, |
| 228 | n_splits=1, |
| 229 | use_norm_weight=False, |
| 230 | is_broadcast=False, |
| 231 | rms_eps=1.0e-6, |
| 232 | hc_pre_eps=2.0e-6, |
| 233 | hc_sinkhorn_eps=3.0e-6, |
| 234 | hc_post_mult_value=0.5, |
| 235 | sinkhorn_repeat=3, |
| 236 | norm_eps=0.0, |
| 237 | ), |
| 238 | ), |
| 239 | ( |
| 240 | MhcPreBigFuseTileLangKernel(), |
| 241 | MhcPreBigFuseTileLangKernel.CompileKey( |
| 242 | hidden_size=4096, |
| 243 | hc_mult=4, |
| 244 | n_splits=2, |
| 245 | use_norm_weight=True, |
| 246 | is_broadcast=True, |
| 247 | rms_eps=1.0e-6, |
| 248 | hc_pre_eps=2.0e-6, |
| 249 | hc_sinkhorn_eps=3.0e-6, |
| 250 | hc_post_mult_value=0.5, |
| 251 | sinkhorn_repeat=3, |
| 252 | norm_eps=1.0e-5, |
| 253 | ), |
| 254 | ), |
| 255 | ( |
| 256 | MhcPostTileLangKernel(), |
| 257 | MhcPostTileLangKernel.CompileKey(hidden_size=4096, hc_mult=4), |
| 258 | ), |
| 259 | ( |
| 260 | MhcFusedTileLangKernel(), |
| 261 | MhcFusedTileLangKernel.CompileKey( |
| 262 | hidden_size=4096, |
| 263 | hc_mult=4, |
| 264 | n_splits=8, |
| 265 | tile_n=6, |
| 266 | n_thr=128, |
| 267 | ), |
| 268 | ), |
| 269 | ( |
| 270 | MhcFusedTileLangKernel(), |
| 271 | MhcFusedTileLangKernel.CompileKey( |
| 272 | hidden_size=4096, |
| 273 | hc_mult=4, |
| 274 | n_splits=8, |
| 275 | tile_n=2, |
| 276 | n_thr=128, |
| 277 | ), |
| 278 | ), |
| 279 | ( |
| 280 | MhcFusedTileLangKernel(), |
| 281 | MhcFusedTileLangKernel.CompileKey( |
| 282 | hidden_size=8192, |
| 283 | hc_mult=4, |
| 284 | n_splits=8, |
| 285 | tile_n=2, |
| 286 | n_thr=128, |
| 287 | ), |
| 288 | ), |
| 289 | ( |
| 290 | HcHeadFusedTileLangKernel(), |
| 291 | HcHeadFusedTileLangKernel.CompileKey( |
| 292 | hidden_size=4096, |
| 293 | hc_mult=4, |
| 294 | rms_eps=1.0e-6, |
| 295 | hc_eps=2.0e-6, |
| 296 | ), |
| 297 | ), |
| 298 | ], |
| 299 | ) |
| 300 | def test_tilelang_warmup_inputs_reproduce_compile_key( |
| 301 | monkeypatch: pytest.MonkeyPatch, |
| 302 | kernel: Any, |
| 303 | compile_key: Any, |
| 304 | ) -> None: |
| 305 | compiled: list[Any] = [] |
| 306 | single_kernel = isinstance( |
| 307 | kernel, |
| 308 | (MhcPostTileLangKernel, MhcFusedTileLangKernel, HcHeadFusedTileLangKernel), |
| 309 | ) |
| 310 | expected_kernel = kernel.kernel() if single_kernel else kernel.kernel(compile_key) |
| 311 | if single_kernel: |
| 312 | monkeypatch.setattr( |
| 313 | kernel, |
| 314 | "dispatch", |
| 315 | lambda **kwargs: pytest.fail("single-kernel launch must not dispatch"), |
| 316 | ) |
| 317 | monkeypatch.setattr( |
| 318 | jit_warmup_tilelang_helper, |
| 319 | "compile_tilelang", |
| 320 | lambda jit_impl, *args, **kwargs: compiled.append(jit_impl), |
| 321 | ) |
| 322 | |
| 323 | kernel.compile(compile_key) |
| 324 | |
| 325 | assert compiled == [expected_kernel] |