| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | import vllm._custom_ops as ops |
| 9 | from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant |
| 10 | from vllm.distributed import cleanup_dist_env_and_memory |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.platforms.rocm import on_gfx950 |
| 13 | from vllm.utils.platform_utils import num_compute_units |
| 14 | |
| 15 | # Global per-test cleanup costs more than the tests themselves. |
| 16 | # These tests can just cleanup once at the end |
| 17 | pytestmark = pytest.mark.skip_global_cleanup |
| 18 | |
| 19 | SEEDS = [0] |
| 20 | |
| 21 | # These options are independent in the kernels, so the sets below cover every |
| 22 | # pair of them instead of the full product. bias_mode: 0 = none, 1 = (m,), |
| 23 | # 2 = (n, m). |
| 24 | OPTIONS_WVSPLITKRC = [ |
| 25 | # dtype, padded_a, bias_mode, xnorm |
| 26 | (torch.float16, False, 0, False), |
| 27 | (torch.float16, True, 1, True), |
| 28 | (torch.float16, True, 2, False), |
| 29 | (torch.bfloat16, True, 0, True), |
| 30 | (torch.bfloat16, False, 1, False), |
| 31 | (torch.bfloat16, False, 2, True), |
| 32 | ] |
| 33 | |
| 34 | OPTIONS_WVSPLITK = [ |
| 35 | # dtype, padded_a, padded_b, bias_mode, xnorm |
| 36 | (torch.float16, False, False, 0, False), |
| 37 | (torch.float16, False, True, 1, True), |
| 38 | (torch.float16, True, False, 2, True), |
| 39 | (torch.bfloat16, True, True, 0, True), |
| 40 | (torch.bfloat16, True, False, 1, False), |
| 41 | (torch.bfloat16, False, True, 2, False), |
| 42 | ] |
| 43 | |
| 44 | OPTIONS_WVSPLITK_FP8 = [ |
| 45 | # dtype, padded_a, padded_b, biased, xnorm |
| 46 | (torch.float16, False, False, True, True), |
| 47 | (torch.bfloat16, True, False, False, True), |
| 48 | (torch.bfloat16, False, True, True, True), |
| 49 | (torch.float16, True, False, True, False), |
| 50 | (torch.float16, False, True, False, False), |
| 51 | (torch.bfloat16, True, True, False, False), |
| 52 | ] |
| 53 | |
| 54 | DTYPES = [torch.bfloat16, torch.float16] |
| 55 | |
| 56 | # kWvSlots in csrc/rocm/skinny_gemms.cu; beyond it streams take overflow |
| 57 | # workspaces. |
| 58 | WVSPLITKRC_SLOTS = 8 |
| 59 | |
| 60 | # The CU count is the grid size. Production passes num_compute_units(); a small |
| 61 | # value manufactures better concurrency for these tests. |
| 62 | WVSPLITKRC_TEST_CU = 16 |
| 63 | |
| 64 | # Gate spin, ~125 ms on gfx950 -- must outlast the host enqueue loop it gates. |
| 65 | _GATE_CYCLES = 300_000_000 |
| 66 | |
| 67 | # Specific (N, K, M) combinations for targeted testing |
| 68 | NKM_FACTORS_LLMM1 = [ |
| 69 | # Small, medium, large cases |
| 70 | (1, 8, 16), |
| 71 | (1, 32, 64), |
| 72 | (1, 128, 256), |
| 73 | (1, 512, 1024), |
| 74 | (1, 2048, 4096), |
| 75 | # Edge cases with specific K sizes |
| 76 | (1, 6144, 1024), |
| 77 | (1, 8192, 2048), |
| 78 | # Very large case |
| 79 | (1, 4096, 8192), |
| 80 | ] |
| 81 | |
| 82 | NKM_FACTORS_WVSPLITK = [ |
| 83 | # Different batch sizes with key dimensions |
| 84 | (1, 32, 16), |
| 85 | (1, 64, 64), |
| 86 | (2, 256, 256), |
| 87 | (3, 1024, 1024), |
| 88 | (4, 4096, 4096), |
| 89 | (4, 4096, 4096 + 1), |
| 90 | (4, 4096 + 16, 4096), |
| 91 | (4, 4096 + 16, 4096 + 1), |
| 92 | # Extended K values |
| 93 | (1, 9216, 512), |
| 94 | (2, 10240, 1024), |
| 95 | (4, 16384, 8192), |
| 96 | (4, 16384 * 2, 8192), |
| 97 | (4, 16384 * 2, 8192 + 1), |
| 98 | (4, 16384 * 2 + 16, 8192), |
| 99 | (4, 16384 * 2 + 16, 8192 + 1), |
| 100 | # Minimum M constraint validation (m >= 8) |
| 101 | (1, 64, 8), |
| 102 | (2, 128, 8), |
| 103 | (4, 256, 8), |
| 104 | ] |
| 105 | |
| 106 | # N is bucketed up to its next ^2 (16/32/64/128) and the remainder is masked |
| 107 | N_FACTORS_WVSPLITKRC = [ |
| 108 | 13, |
| 109 | 16, |
| 110 | 17, |
| 111 | 25, |
| 112 | 29, |
| 113 | 31, |
| 114 | 32, |
| 115 | 41, |
| 116 | 51, |
| 117 | 64, |
| 118 | 71, |
| 119 | 81, |
| 120 | 91, |
| 121 | 103, |
| 122 | 117, |
| 123 | 128, |
| 124 | ] |
| 125 | # K shards are 512 wide, evenly divided or not, +8 for a partial 8-element load |
| 126 | K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8] |
| 127 | # M tiles are 64 rows, +16 for a partial tile |
| 128 | M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16] |
| 129 | |
| 130 | # (N, K, M) with more K-shards than the readback can stage in one LDS pass. |
| 131 | # Spans both CHUNKK values and both N-tile counts. |
| 132 | NKM_FACTORS_WVSPLITKRC_LARGE_K = [ |
| 133 | (128, 6144, 128), |
| 134 | (96, 8192, 128), |
| 135 | # Keep 24 K-shards (batched readback) within the DPX budget of 128 CUs. |
| 136 | (128, 12288, 64), |
| 137 | (32, 12288, 128), |
| 138 | ] |
| 139 | |
| 140 | NKM_FACTORS_WVSPLITK_FP8 = [ |
| 141 | # FP8-specific cases with K % 16 == 0 |
| 142 | (1, 16, 16), |
| 143 | (1, 32, 16 + 16), |
| 144 | (1, 64, 64), |
| 145 | (1, 64, 64 + 16), |
| 146 | (1, 64 + 16, 64), |
| 147 | (1, 64 + 16, 64 + 16), |
| 148 | (4, 64, 64), |
| 149 | (4, 64, 64 + 16), |
| 150 | (4, 64 + 16, 64), |
| 151 | (4, 64 + 16, 64 + 16), |
| 152 | (2, 512, 512), |
| 153 | (3, 512, 512), |
| 154 | (3, 512, 512 + 16), |
| 155 | (4, 512, 512), |
| 156 | (3, 2048, 2048), |
| 157 | (3, 2048, 2048 + 16), |
| 158 | (4, 2048 + 16, 2048), |
| 159 | (4, 2048 + 16, 2048 + 16), |
| 160 | (4, 4096, 4096), |
| 161 | (4, 16400, 2048), |
| 162 | (4, 16400, 2048 + 16), |
| 163 | # Extended FP8 dimensions not covered by WVSPLITK |
| 164 | (1, 14336, 1024), |
| 165 | (2, 24576, 2048), |
| 166 | (4, 32768, 28672), |
| 167 | (4, 32768 * 2, 28672), |
| 168 | (4, 32768 * 2, 28672 + 16), |
| 169 | (4, 32768 * 2 + 16, 28672), |
| 170 | (4, 32768 * 2 + 16, 28672 + 16), |
| 171 | ] |
| 172 | |
| 173 | |
| 174 | @pytest.fixture(scope="module", autouse=True) |
| 175 | def cleanup_after_all_tests(): |
| 176 | yield |
| 177 | cleanup_dist_env_and_memory() |
| 178 | |
| 179 | |
| 180 | def pad_fp8(weight): |
| 181 | num_pad = 256 // weight.element_size() |
| 182 | import torch.nn.functional as F |
| 183 | |
| 184 | return F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad] |
| 185 | |
| 186 | |
| 187 | def make_bias(bias_mode, n, m, dtype): |
| 188 | if bias_mode == 0: |
| 189 | return None |
| 190 | shape = (m,) if bias_mode == 1 else (n, m) |
| 191 | return torch.rand(shape, dtype=dtype, device="cuda") * 2 - 1 |
| 192 | |
| 193 | |
| 194 | @pytest.mark.parametrize("n", N_FACTORS_WVSPLITKRC) |
| 195 | @pytest.mark.parametrize("k", K_FACTORS_WVSPLITKRC) |
| 196 | @pytest.mark.parametrize("m", M_FACTORS_WVSPLITKRC) |
| 197 | @pytest.mark.parametrize("dtype,padded_a,bias_mode,xnorm", OPTIONS_WVSPLITKRC) |
| 198 | @pytest.mark.parametrize("seed", SEEDS) |
| 199 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm") |
| 200 | @pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950") |
| 201 | def test_rocm_wvsplitkrc_kernel(n, k, m, dtype, padded_a, bias_mode, xnorm, seed): |
| 202 | torch.manual_seed(seed) |
| 203 | cu_count = num_compute_units() |
| 204 | |
| 205 | # Next ^2 of n |
| 206 | N_p2 = 1 << (n - 1).bit_length() |
| 207 | # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), |
| 208 | # and each working on a 512-shard of K, how many CUs would we need? |
| 209 | rndup_cus = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512) |
| 210 | # How many of 4 waves in a group can work on same 16 Ms at same time? |
| 211 | # This reduces the Ms each group works on, i.e. increasing the number of CUs needed. |
| 212 | GrpsShrB = min(N_p2 // 16, 4) |
| 213 | # Given the above, how many CUs would we need? |
| 214 | CuNeeded = rndup_cus * GrpsShrB |
| 215 | # Deterministic reduction stores one float workspace value per K shard. |
| 216 | fits_wvsplitkrc = (N_p2 * m * ((k + 512 - 1) // 512)) <= 128 * 1024 * 12 |
| 217 | fits_wvsplitkrc &= CuNeeded <= cu_count |
| 218 | |
| 219 | if not fits_wvsplitkrc: |
| 220 | pytest.skip("Too large for wvSplitKrc") |
| 221 | |
| 222 | xavier = ( |
| 223 | math.sqrt(2 / k) if xnorm else 1 |
| 224 | ) # normalize to avoid large output-bias deltas |
| 225 | A = torch.randn(n, k, dtype=dtype, device="cuda") * xavier |
| 226 | B = torch.randn(m, k, dtype=dtype, device="cuda") * xavier |
| 227 | if padded_a: |
| 228 | A = pad_fp8(A) |
| 229 | |
| 230 | BIAS = make_bias(bias_mode, n, m, dtype) |
| 231 | |
| 232 | ref_out = torch.nn.functional.linear(A, B, BIAS) |
| 233 | out = ops.wvSplitKrc(A, B, cu_count, BIAS) |
| 234 | |
| 235 | if xnorm: |
| 236 | # The O(1) bias lifts outputs to ~O(1), where one bf16 ULP (~3.9e-3, the |
| 237 | # worst measured divergence) exceeds 1e-3. Bump atol to 5e-3 only for |
| 238 | # biased bf16 (above that ULP, still under finfo(bf16).eps); keep 1e-3 |
| 239 | # otherwise. |
| 240 | atol = 5e-3 if (dtype == torch.bfloat16 and BIAS is not None) else 1e-3 |
| 241 | torch.testing.assert_close(out, ref_out, atol=atol, rtol=1e-8) |
| 242 | else: |
| 243 | torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2) |
| 244 | |
| 245 | |
| 246 | @pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC_LARGE_K) |
| 247 | @pytest.mark.parametrize("dtype", DTYPES) |
| 248 | @pytest.mark.parametrize("seed", SEEDS) |
| 249 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm") |
| 250 | @pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950") |
| 251 | def test_rocm_wvsplitkrc_large_k(n, k, m, dtype, seed): |
| 252 | """K large enough that the split-K readback must stage LDS in batches.""" |
| 253 | torch.manual_seed(seed) |
| 254 | |
| 255 | xavier = math.sqrt(2 / k) |
| 256 | A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier |
| 257 | B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier |
| 258 | |
| 259 | ref_out = torch.nn.functional.linear(A, B, None) |
| 260 | out = ops.wvSplitKrc(A, B, num_compute_units(), None) |
| 261 | |
| 262 | torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2) |
| 263 | |
| 264 | |
| 265 | @pytest.mark.parametrize("n_streams", [WVSPLITKRC_SLOTS, WVSPLITKRC_SLOTS + 4]) |
| 266 | @pytest.mark.parametrize("dtype", DTYPES) |
| 267 | @pytest.mark.parametrize("seed", SEEDS) |
| 268 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm") |
| 269 | @pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950") |
| 270 | def test_rocm_wvsplitkrc_multistream(n_streams, dtype, seed): |
| 271 | """Concurrent streams must not share a split-K workspace. |
| 272 | |
| 273 | The workspace is cleared by protocol rather than per invocation, so sharing |
| 274 | one corrupts the reduction and returns wrong numbers rather than failing. |
| 275 | n_streams > WVSPLITKRC_SLOTS additionally exercises the overflow path. |
| 276 | """ |
| 277 | torch.manual_seed(seed) |
| 278 | |
| 279 | n, k, m = 16, 2048, 256 |
| 280 | iters = 4 |
| 281 | xavier = math.sqrt(2 / k) |
| 282 | |
| 283 | B = torch.randn(m, k, dtype=dtype, device="cuda") * xavier |
| 284 | As = [ |
| 285 | torch.randn(n, k, dtype=dtype, device="cuda") * xavier for _ in range(n_streams) |
| 286 | ] |
| 287 | refs = [torch.nn.functional.linear(A, B, None) for A in As] |
| 288 | torch.accelerator.synchronize() |
| 289 | |
| 290 | streams = [torch.Stream() for _ in range(n_streams)] |
| 291 | outs: list[list[torch.Tensor]] = [[] for _ in range(n_streams)] |
| 292 | _release_together(streams) |
| 293 | # Round-robin the enqueues so kernels from different streams are in flight |
| 294 | # at the same time. |
| 295 | for _ in range(iters): |
| 296 | for i, s in enumerate(streams): |
| 297 | with torch.cuda.stream(s): |
| 298 | outs[i].append(ops.wvSplitKrc(As[i], B, WVSPLITKRC_TEST_CU, None)) |
| 299 | for s in streams: |
| 300 | torch.accelerator.current_stream().wait_stream(s) |
| 301 | torch.accelerator.synchronize() |
| 302 | |
| 303 | for i in range(n_streams): |
| 304 | for out in outs[i]: |
| 305 | torch.testing.assert_close(out, refs[i], atol=1e-3, rtol=1e-2) |
| 306 | |
| 307 | |
| 308 | def _release_together(streams): |
| 309 | """Release every stream at once, once everything below is enqueued. |
| 310 | |
| 311 | An event recorded on an idle stream is already complete, so waiting on it |
| 312 | gates nothing -- the spin, not the event, is what holds the streams. |
| 313 | """ |
| 314 | blocker = torch.Stream() |
| 315 | with torch.cuda.stream(blocker): |
| 316 | torch.cuda._sleep(_GATE_CYCLES) |
| 317 | gate = torch.Event() |
| 318 | gate.record() |
| 319 | for s in streams: |
| 320 | s.wait_event(gate) |
| 321 | |
| 322 | |
| 323 | @pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1) |
| 324 | @pytest.mark.parametrize("dtype", DTYPES) |
| 325 | @pytest.mark.parametrize("rows_per_block", [2, 4, 8, 16]) |
| 326 | @pytest.mark.parametrize("seed", SEEDS) |
| 327 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm") |
| 328 | @torch.inference_mode() |
| 329 | def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed): |
| 330 | torch.manual_seed(seed) |
| 331 | # TODO: Zero-centering the inputs causes errors for LLMM1! |
| 332 | # Without that the numbers quickly saturate, and may |
| 333 | # be giving false matches. |
| 334 | A = torch.rand(n, k, dtype=dtype, device="cuda") |
| 335 | B = torch.rand(m, k, dtype=dtype, device="cuda") |
| 336 | |
| 337 | ref_out = torch.matmul(A, B.t()) |
| 338 | out = ops.LLMM1(B, A, rows_per_block) |
| 339 | |
| 340 | torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2) |
| 341 | |
| 342 | |
| 343 | @pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK) |
| 344 | @pytest.mark.parametrize( |
| 345 | "dtype,padded_a,padded_b,bias_mode,xnorm", |
| 346 | OPTIONS_WVSPLITK, |
| 347 | ) |
| 348 | @pytest.mark.parametrize("seed", SEEDS) |
| 349 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm") |
| 350 | def test_rocm_wvsplitk_kernel( |
| 351 | n, k, m, dtype, padded_a, padded_b, bias_mode, xnorm, seed |
| 352 | ): |
| 353 | torch.manual_seed(seed) |
| 354 | cu_count = num_compute_units() |
| 355 | |
| 356 | xavier = ( |
| 357 | math.sqrt(2 / k) if xnorm else 1 |
| 358 | ) # normalize to avoid large output-bias deltas |
| 359 | A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier |
| 360 | B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier |
| 361 | |
| 362 | BIAS = make_bias(bias_mode, n, m, dtype) |
| 363 | |
| 364 | if padded_a: |
| 365 | A = pad_fp8(A) |
| 366 | if padded_b: |
| 367 | B = pad_fp8(B) |
| 368 | |
| 369 | ref_out = torch.nn.functional.linear(A, B, BIAS) |
| 370 | out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS) |
| 371 | |
| 372 | # Accumulation error in fp16 GEMM scales with sqrt(K) |
| 373 | atol = torch.finfo(dtype).eps * math.sqrt(k) |
| 374 | torch.testing.assert_close(out, ref_out, atol=atol, rtol=1e-2) |
| 375 | |
| 376 | |
| 377 | @pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8) |
| 378 | @pytest.mark.parametrize( |
| 379 | "dtype,padded_a,padded_b,biased,xnorm", |
| 380 | OPTIONS_WVSPLITK_FP8, |
| 381 | ) |
| 382 | @pytest.mark.parametrize("seed", SEEDS) |
| 383 | @pytest.mark.skipif( |
| 384 | not (current_platform.is_rocm() and current_platform.supports_fp8()), |
| 385 | reason="only test for rocm fp8", |
| 386 | ) |
| 387 | def test_rocm_wvsplitk_fp8_kernel( |
| 388 | n, k, m, dtype, padded_a, padded_b, biased, xnorm, seed |
| 389 | ): |
| 390 | torch.manual_seed(seed) |
| 391 | |
| 392 | xavier = math.sqrt(2 / k) if xnorm else 1 # normalize to avoid large deltas |
| 393 | A = (torch.rand(n, k, device="cuda") * 2 - 1) * xavier |
| 394 | B = (torch.rand(m, k, device="cuda") * 2 - 1) * xavier |
| 395 | |
| 396 | A, scale_a = ref_dynamic_per_tensor_fp8_quant(A) |
| 397 | B, scale_b = ref_dynamic_per_tensor_fp8_quant(B) |
| 398 | if padded_b: |
| 399 | B = pad_fp8(B) |
| 400 | if padded_a: |
| 401 | A = pad_fp8(A) |
| 402 | |
| 403 | BIAS = None if (not biased) else (torch.rand(m, dtype=dtype, device="cuda") * 2 - 1) |
| 404 | |
| 405 | ref_out = torch._scaled_mm( |
| 406 | A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS |
| 407 | ) |
| 408 | out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, num_compute_units(), BIAS) |
| 409 | |
| 410 | if xnorm: |
| 411 | torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8) |
| 412 | elif k >= 32 * 1024: |
| 413 | # wider pytrch thresh for large-K & no xnorm |
| 414 | torch.testing.assert_close(out, ref_out, atol=0.07, rtol=5e-2) |
| 415 | else: |
| 416 | torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2) |