| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Tests for the ROCm Hybrid W4A16 kernel (HIP skinny + Triton prefill). |
| 5 | |
| 6 | Run `pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py`. |
| 7 | """ |
| 8 | |
| 9 | import importlib |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.torch_utils import set_random_seed |
| 16 | |
| 17 | if not current_platform.is_rocm(): |
| 18 | pytest.skip("ROCm only", allow_module_level=True) |
| 19 | |
| 20 | pytest.importorskip("triton") |
| 21 | |
| 22 | from vllm.platforms.rocm import on_gfx1x # noqa: E402 |
| 23 | |
| 24 | device = "cuda" |
| 25 | |
| 26 | hybrid_module = importlib.import_module( |
| 27 | "vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16" |
| 28 | ) |
| 29 | RDNAHybridW4A16LinearKernel = hybrid_module.RDNAHybridW4A16LinearKernel |
| 30 | pack_int4_exllama_shuffle = hybrid_module.pack_int4_exllama_shuffle |
| 31 | SUPPORTED_GROUP_SIZES = hybrid_module.SUPPORTED_GROUP_SIZES |
| 32 | MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE |
| 33 | |
| 34 | |
| 35 | # --------------------------------------------------------------------------- |
| 36 | # Reference implementation |
| 37 | # --------------------------------------------------------------------------- |
| 38 | |
| 39 | |
| 40 | def _pack_zp_rows_for_kernel(zp_nkg: torch.Tensor) -> torch.Tensor: |
| 41 | """Pack raw uint4 zero points along N: [N, G] int32 -> [N//8, G] int32. |
| 42 | |
| 43 | Row n's nibble lands in word[n//8] at bits 4*(n%8). |
| 44 | """ |
| 45 | assert zp_nkg.dtype == torch.int32 |
| 46 | N, G = zp_nkg.shape |
| 47 | assert N % 8 == 0 |
| 48 | shifts = (torch.arange(8, device=zp_nkg.device, dtype=torch.int32) * 4)[:, None] |
| 49 | return torch.sum( |
| 50 | (zp_nkg.view(N // 8, 8, G) & 0xF) << shifts, dim=1, dtype=torch.int32 |
| 51 | ).contiguous() |
| 52 | |
| 53 | |
| 54 | def _rdna_hybrid_w4a16_reference( |
| 55 | x_mk: torch.Tensor, |
| 56 | w_int4_nk: torch.Tensor, |
| 57 | scales_nkg: torch.Tensor, |
| 58 | zp_nkg: torch.Tensor | None, |
| 59 | group_size: int, |
| 60 | bias: torch.Tensor | None, |
| 61 | ) -> torch.Tensor: |
| 62 | """Reference for the Hybrid W4A16 op. |
| 63 | |
| 64 | x_mk: [M, K] fp16/bf16 |
| 65 | w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15] |
| 66 | scales_nkg: [N, K//G] fp16/bf16 |
| 67 | zp_nkg: [N, K//G] int32 raw zero points in [0, 15], or None for |
| 68 | symmetric (uint4b8, dequant subtracts 8) |
| 69 | """ |
| 70 | G = group_size |
| 71 | N, K = w_int4_nk.shape |
| 72 | assert K % G == 0 |
| 73 | s_full = scales_nkg.repeat_interleave(G, dim=1).to(torch.float32) # [N, K] |
| 74 | if zp_nkg is None: |
| 75 | z_full = torch.full((N, K), 8.0, device=x_mk.device, dtype=torch.float32) |
| 76 | else: |
| 77 | z_full = zp_nkg.repeat_interleave(G, dim=1).to(torch.float32) |
| 78 | w_fp = (w_int4_nk.to(torch.float32) - z_full) * s_full # [N, K] |
| 79 | out = x_mk.to(torch.float32) @ w_fp.t() # [M, N] |
| 80 | if bias is not None: |
| 81 | out = out + bias.to(torch.float32) |
| 82 | return out.to(x_mk.dtype) |
| 83 | |
| 84 | |
| 85 | # --------------------------------------------------------------------------- |
| 86 | # Forward correctness: decode (HIP skinny) + prefill (Triton) |
| 87 | # --------------------------------------------------------------------------- |
| 88 | |
| 89 | |
| 90 | @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") |
| 91 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 92 | @pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES) |
| 93 | @pytest.mark.parametrize("has_zp", [False, True]) |
| 94 | @pytest.mark.parametrize( |
| 95 | "M", |
| 96 | [1, MAX_SKINNY_BATCH_SIZE, MAX_SKINNY_BATCH_SIZE + 1, 64], |
| 97 | ids=["M=1_decode", "M=5_decode", "M=6_prefill", "M=64_prefill"], |
| 98 | ) |
| 99 | def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M): |
| 100 | """Smoke test the registered custom op for both decode and prefill batches. |
| 101 | |
| 102 | Verifies the dispatch logic in `_rdna_hybrid_w4a16_apply_impl`: |
| 103 | - M <= MAX_SKINNY_BATCH_SIZE: HIP wvSplitK_int4_g |
| 104 | - M > MAX_SKINNY_BATCH_SIZE: Triton prefill kernel |
| 105 | """ |
| 106 | if not torch.cuda.is_available(): |
| 107 | pytest.skip("CUDA/HIP device not available") |
| 108 | |
| 109 | set_random_seed(0) |
| 110 | |
| 111 | K, N = 1024, 256 |
| 112 | assert K % group_size == 0 and K % 8 == 0 and N % 8 == 0 |
| 113 | |
| 114 | # Activations. |
| 115 | x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) |
| 116 | |
| 117 | # Weights as raw uint4 in [N, K], packed to ExLlama shuffle [N, K//8]. |
| 118 | w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) |
| 119 | w_q_i32 = pack_int4_exllama_shuffle(w_int4_nk).contiguous() # [N, K//8] int32 |
| 120 | w_q = w_q_i32.view(torch.int8) # same bytes viewed as int8 [N, K//2] |
| 121 | |
| 122 | # Scales [N, K//G] in act dtype. |
| 123 | scales_nkg = ( |
| 124 | 0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32) |
| 125 | ).to(dtype) |
| 126 | |
| 127 | # Optional raw zero points, [N, K//G] int32 for the reference and packed |
| 128 | # along N for the op. |
| 129 | if has_zp: |
| 130 | zp_nkg = torch.randint( |
| 131 | 0, 16, (N, K // group_size), device=device, dtype=torch.int32 |
| 132 | ) |
| 133 | w_zp = _pack_zp_rows_for_kernel(zp_nkg) |
| 134 | else: |
| 135 | zp_nkg = None |
| 136 | w_zp = None |
| 137 | |
| 138 | from vllm.utils.platform_utils import num_compute_units |
| 139 | |
| 140 | out = torch.ops.vllm.rdna_hybrid_w4a16_apply( |
| 141 | x_mk, |
| 142 | w_q, |
| 143 | scales_nkg, |
| 144 | w_zp, |
| 145 | None, # bias |
| 146 | num_compute_units(), |
| 147 | group_size, |
| 148 | ) |
| 149 | |
| 150 | ref = _rdna_hybrid_w4a16_reference( |
| 151 | x_mk, w_int4_nk, scales_nkg, zp_nkg, group_size, bias=None |
| 152 | ) |
| 153 | |
| 154 | torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2) |
| 155 | |
| 156 | |
| 157 | @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") |
| 158 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 159 | @pytest.mark.parametrize("M", [1, MAX_SKINNY_BATCH_SIZE + 1]) |
| 160 | def test_rdna_hybrid_w4a16_apply_with_bias(dtype, M): |
| 161 | """Bias is added correctly on both decode and prefill paths.""" |
| 162 | if not torch.cuda.is_available(): |
| 163 | pytest.skip("CUDA/HIP device not available") |
| 164 | |
| 165 | set_random_seed(0) |
| 166 | K, N, G = 1024, 128, 128 |
| 167 | |
| 168 | x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) |
| 169 | w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) |
| 170 | w_q_i32 = pack_int4_exllama_shuffle(w_int4_nk).contiguous() |
| 171 | w_q = w_q_i32.view(torch.int8) |
| 172 | scales_nkg = ( |
| 173 | 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32) |
| 174 | ).to(dtype) |
| 175 | bias = torch.randn(N, device=device, dtype=dtype) * 0.1 |
| 176 | |
| 177 | from vllm.utils.platform_utils import num_compute_units |
| 178 | |
| 179 | out = torch.ops.vllm.rdna_hybrid_w4a16_apply( |
| 180 | x_mk, |
| 181 | w_q, |
| 182 | scales_nkg, |
| 183 | None, |
| 184 | bias, |
| 185 | num_compute_units(), |
| 186 | G, |
| 187 | ) |
| 188 | ref = _rdna_hybrid_w4a16_reference(x_mk, w_int4_nk, scales_nkg, None, G, bias=bias) |
| 189 | |
| 190 | torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2) |
| 191 | |
| 192 | |
| 193 | # --------------------------------------------------------------------------- |
| 194 | # pack_int4_exllama_shuffle round-trips correctly |
| 195 | # --------------------------------------------------------------------------- |
| 196 | |
| 197 | |
| 198 | def test_pack_int4_exllama_shuffle_layout(): |
| 199 | """Pack 8 K-values per int32 in interleave [0,2,4,6,1,3,5,7] order.""" |
| 200 | if not torch.cuda.is_available(): |
| 201 | pytest.skip("CUDA/HIP device not available") |
| 202 | set_random_seed(0) |
| 203 | N, K = 4, 16 |
| 204 | w = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) |
| 205 | packed = pack_int4_exllama_shuffle(w) |
| 206 | assert packed.shape == (N, K // 8) and packed.dtype == torch.int32 |
| 207 | |
| 208 | # Manual unshuffle using ExLlama shifts [0,16,4,20,8,24,12,28]. |
| 209 | shifts = torch.tensor( |
| 210 | [0, 16, 4, 20, 8, 24, 12, 28], device=device, dtype=torch.int32 |
| 211 | ) |
| 212 | unshuffled = (packed.unsqueeze(-1) >> shifts) & 0xF |
| 213 | unshuffled = unshuffled.reshape(N, K) |
| 214 | torch.testing.assert_close(unshuffled, w) |
| 215 | |
| 216 | |
| 217 | # --------------------------------------------------------------------------- |
| 218 | # process_weights_after_loading: layout repack and zp normalization |
| 219 | # --------------------------------------------------------------------------- |
| 220 | |
| 221 | |
| 222 | def _pack_int4_along_k_to_ckpt(w_int4_kn: torch.Tensor) -> torch.Tensor: |
| 223 | """Pack int4 values along K into CT checkpoint layout: [K,N] -> [N, K//8].""" |
| 224 | assert w_int4_kn.dtype == torch.int32 |
| 225 | K, N = w_int4_kn.shape |
| 226 | assert K % 8 == 0 |
| 227 | out = torch.zeros((N, K // 8), dtype=torch.int32, device=w_int4_kn.device) |
| 228 | for i in range(8): |
| 229 | out |= (w_int4_kn[i::8, :].t() & 0xF) << (i * 4) |
| 230 | return out.contiguous() |
| 231 | |
| 232 | |
| 233 | def _pack_int4_along_n_for_zp(zp_int4_gn: torch.Tensor) -> torch.Tensor: |
| 234 | """Pack int4 zero points along N: [G, N] -> [G, N//8] int32 (CT layout).""" |
| 235 | assert zp_int4_gn.dtype == torch.int32 |
| 236 | G, N = zp_int4_gn.shape |
| 237 | assert N % 8 == 0 |
| 238 | shifts = torch.arange(8, device=zp_int4_gn.device, dtype=torch.int32) * 4 |
| 239 | return torch.sum( |
| 240 | (zp_int4_gn.view(G, N // 8, 8) & 0xF) << shifts, dim=2, dtype=torch.int32 |
| 241 | ).contiguous() |
| 242 | |
| 243 | |
| 244 | def _build_dummy_layer( |
| 245 | w_ckpt_nk8: torch.Tensor, |
| 246 | scales_ckpt_nkg: torch.Tensor, |
| 247 | zeros_ckpt: torch.Tensor | None, |
| 248 | ): |
| 249 | from vllm.model_executor.parameter import ( |
| 250 | GroupQuantScaleParameter, |
| 251 | PackedColumnParameter, |
| 252 | PackedvLLMParameter, |
| 253 | ) |
| 254 | |
| 255 | weight_loader = lambda *args, **kwargs: None |
| 256 | |
| 257 | class DummyLayer(torch.nn.Module): |
| 258 | pass |
| 259 | |
| 260 | layer = DummyLayer() |
| 261 | layer.register_parameter( |
| 262 | "weight_packed", |
| 263 | PackedvLLMParameter( |
| 264 | data=w_ckpt_nk8, |
| 265 | weight_loader=weight_loader, |
| 266 | input_dim=1, |
| 267 | output_dim=0, |
| 268 | packed_factor=8, |
| 269 | packed_dim=1, |
| 270 | ), |
| 271 | ) |
| 272 | layer.register_parameter( |
| 273 | "weight_scale", |
| 274 | GroupQuantScaleParameter( |
| 275 | data=scales_ckpt_nkg, |
| 276 | weight_loader=weight_loader, |
| 277 | input_dim=1, |
| 278 | output_dim=0, |
| 279 | ), |
| 280 | ) |
| 281 | if zeros_ckpt is not None: |
| 282 | layer.register_parameter( |
| 283 | "weight_zero_point", |
| 284 | PackedColumnParameter( |
| 285 | data=zeros_ckpt, |
| 286 | weight_loader=weight_loader, |
| 287 | output_dim=0, |
| 288 | packed_factor=8, |
| 289 | packed_dim=0, |
| 290 | ), |
| 291 | ) |
| 292 | return layer |
| 293 | |
| 294 | |
| 295 | @pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES) |
| 296 | def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_init): |
| 297 | """uint4b8 (symmetric): w_q -> [N, K//8] int8 ExLlama shuffle, no zp param.""" |
| 298 | if not torch.cuda.is_available(): |
| 299 | pytest.skip("CUDA/HIP device not available") |
| 300 | |
| 301 | from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( |
| 302 | MPLinearLayerConfig, |
| 303 | ) |
| 304 | from vllm.scalar_type import scalar_types |
| 305 | |
| 306 | set_random_seed(0) |
| 307 | |
| 308 | K, N = 256, 128 |
| 309 | G = group_size |
| 310 | assert K % G == 0 |
| 311 | |
| 312 | # Reference unpacked weights, then pack into CT checkpoint layout [N, K//8]. |
| 313 | w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) |
| 314 | w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn) |
| 315 | scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16) |
| 316 | |
| 317 | layer = _build_dummy_layer(w_ckpt_nk8, scales_ckpt_nkg, zeros_ckpt=None) |
| 318 | |
| 319 | config = MPLinearLayerConfig( |
| 320 | full_weight_shape=(K, N), |
| 321 | partition_weight_shape=(K, N), |
| 322 | weight_type=scalar_types.uint4b8, |
| 323 | act_type=torch.float16, |
| 324 | group_size=G, |
| 325 | zero_points=False, |
| 326 | ) |
| 327 | kernel = RDNAHybridW4A16LinearKernel( |
| 328 | config, |
| 329 | w_q_param_name="weight_packed", |
| 330 | w_s_param_name="weight_scale", |
| 331 | w_zp_param_name=None, |
| 332 | ) |
| 333 | kernel.process_weights_after_loading(layer) |
| 334 | |
| 335 | # Skinny weight is stored once as int8 [N, K//2]; the Triton path |
| 336 | # reinterprets it as int32 [N, K//8] via a view (no separate parameter). |
| 337 | assert layer.weight_packed.dtype == torch.int8 |
| 338 | assert tuple(layer.weight_packed.shape) == (N, K // 2) |
| 339 | w_q_i32 = layer.weight_packed.view(torch.int32) |
| 340 | assert tuple(w_q_i32.shape) == (N, K // 8) |
| 341 | |
| 342 | expected_packed = pack_int4_exllama_shuffle(w_int4_kn.t().contiguous()) |
| 343 | torch.testing.assert_close(w_q_i32, expected_packed) |
| 344 | |
| 345 | # Scales: [N, K//G] (skinny layout, no transpose since CT already had it). |
| 346 | assert tuple(layer.weight_scale.shape) == (N, K // G) |
| 347 | torch.testing.assert_close(layer.weight_scale, scales_ckpt_nkg) |
| 348 | |
| 349 | |
| 350 | @pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES) |
| 351 | def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init): |
| 352 | """uint4 (asymmetric): zero points are packed [N//8, K//G] int32.""" |
| 353 | if not torch.cuda.is_available(): |
| 354 | pytest.skip("CUDA/HIP device not available") |
| 355 | |
| 356 | from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( |
| 357 | MPLinearLayerConfig, |
| 358 | ) |
| 359 | from vllm.scalar_type import scalar_types |
| 360 | |
| 361 | set_random_seed(0) |
| 362 | |
| 363 | K, N = 256, 128 |
| 364 | G = group_size |
| 365 | assert K % G == 0 and N % 8 == 0 |
| 366 | |
| 367 | w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) |
| 368 | w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn) |
| 369 | scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16) |
| 370 | |
| 371 | # CT zero-point layout is N-packed: [N//8, K//G] int32. |
| 372 | zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32) |
| 373 | zeros_packed_gn8 = _pack_int4_along_n_for_zp(zeros_int4_gn) # [K//G, N//8] |
| 374 | zeros_ckpt_n8kg = zeros_packed_gn8.t().contiguous() # [N//8, K//G] |
| 375 | |
| 376 | layer = _build_dummy_layer(w_ckpt_nk8, scales_ckpt_nkg, zeros_ckpt=zeros_ckpt_n8kg) |
| 377 | |
| 378 | config = MPLinearLayerConfig( |
| 379 | full_weight_shape=(K, N), |
| 380 | partition_weight_shape=(K, N), |
| 381 | weight_type=scalar_types.uint4, |
| 382 | act_type=torch.float16, |
| 383 | group_size=G, |
| 384 | zero_points=True, |
| 385 | ) |
| 386 | kernel = RDNAHybridW4A16LinearKernel( |
| 387 | config, |
| 388 | w_q_param_name="weight_packed", |
| 389 | w_s_param_name="weight_scale", |
| 390 | w_zp_param_name="weight_zero_point", |
| 391 | ) |
| 392 | kernel.process_weights_after_loading(layer) |
| 393 | |
| 394 | # Zero-points: [N//8, K//G] int32, row n's nibble at word[n//8] |
| 395 | # bits 4*(n%8). |
| 396 | assert layer.weight_zero_point.dtype == torch.int32 |
| 397 | assert tuple(layer.weight_zero_point.shape) == (N // 8, K // G) |
| 398 | expected_zp = _pack_zp_rows_for_kernel(zeros_int4_gn.t().contiguous()) |
| 399 | torch.testing.assert_close(layer.weight_zero_point, expected_zp) |
| 400 | |
| 401 | # Quantized weights match symmetric path's layout regardless of zp. |
| 402 | w_q_i32 = layer.weight_packed.view(torch.int32) |
| 403 | assert tuple(w_q_i32.shape) == (N, K // 8) |
| 404 | expected_packed = pack_int4_exllama_shuffle(w_int4_kn.t().contiguous()) |
| 405 | torch.testing.assert_close(w_q_i32, expected_packed) |
| 406 | |
| 407 | |
| 408 | # --------------------------------------------------------------------------- |
| 409 | # can_implement enforces the supported-group-size policy |
| 410 | # --------------------------------------------------------------------------- |
| 411 | |
| 412 | |
| 413 | @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") |
| 414 | @pytest.mark.parametrize( |
| 415 | "group_size,expected_ok", [(32, True), (64, True), (128, True), (256, False)] |
| 416 | ) |
| 417 | def test_hybrid_can_implement_group_size(group_size, expected_ok): |
| 418 | from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( |
| 419 | MPLinearLayerConfig, |
| 420 | ) |
| 421 | from vllm.scalar_type import scalar_types |
| 422 | |
| 423 | K, N = 1024, 256 |
| 424 | config = MPLinearLayerConfig( |
| 425 | full_weight_shape=(K, N), |
| 426 | partition_weight_shape=(K, N), |
| 427 | weight_type=scalar_types.uint4b8, |
| 428 | act_type=torch.float16, |
| 429 | group_size=group_size, |
| 430 | zero_points=False, |
| 431 | ) |
| 432 | ok, _ = RDNAHybridW4A16LinearKernel.can_implement(config) |
| 433 | assert ok is expected_ok |
| 434 | |
| 435 | |
| 436 | # --------------------------------------------------------------------------- |
| 437 | # Tests for the HIP wvSplitK_int4_g decode kernel |
| 438 | # --------------------------------------------------------------------------- |
| 439 | |
| 440 | |
| 441 | def _hip_skinny_reference( |
| 442 | a_mk: torch.Tensor, |
| 443 | w_int4_nk: torch.Tensor, |
| 444 | scales_nkg: torch.Tensor, |
| 445 | *, |
| 446 | group_size: int, |
| 447 | zp_bias: int, |
| 448 | ) -> torch.Tensor: |
| 449 | """Reference for symmetric HIP skinny: C = A @ (W - zp_bias) * S.""" |
| 450 | K = a_mk.shape[1] |
| 451 | N = w_int4_nk.shape[0] |
| 452 | num_groups = K // group_size |
| 453 | |
| 454 | w_fp = (w_int4_nk.to(torch.float32) - zp_bias).view(N, num_groups, group_size) |
| 455 | s = scales_nkg.to(torch.float32).unsqueeze(-1) |
| 456 | w_dequant = (w_fp * s).view(N, K) |
| 457 | |
| 458 | return (a_mk.to(torch.float32) @ w_dequant.t()).to(a_mk.dtype) |
| 459 | |
| 460 | |
| 461 | @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") |
| 462 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 463 | @pytest.mark.parametrize( |
| 464 | "M,K,N,G", |
| 465 | [ |
| 466 | (1, 256, 256, 32), |
| 467 | (1, 256, 256, 64), |
| 468 | (1, 512, 256, 128), |
| 469 | (2, 512, 256, 64), |
| 470 | (3, 256, 512, 64), |
| 471 | ], |
| 472 | ) |
| 473 | def test_hip_skinny_wvSplitK_int4_g(dtype, M, K, N, G): |
| 474 | """Test HIP wvSplitK_int4_g kernel directly via _custom_ops.""" |
| 475 | import vllm._custom_ops as ops |
| 476 | from vllm.utils.platform_utils import num_compute_units |
| 477 | |
| 478 | set_random_seed(0) |
| 479 | |
| 480 | a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) |
| 481 | w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) |
| 482 | |
| 483 | b_packed_i32 = pack_int4_exllama_shuffle(w_int4_nk) |
| 484 | b_packed_i8 = b_packed_i32.view(torch.int8) |
| 485 | |
| 486 | scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to( |
| 487 | dtype |
| 488 | ) |
| 489 | |
| 490 | cu_count = num_compute_units() |
| 491 | out = ops.wvSplitK_int4_g(b_packed_i8, a, scales, cu_count, G) |
| 492 | |
| 493 | ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8) |
| 494 | |
| 495 | torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) |
| 496 | |
| 497 | |
| 498 | # --------------------------------------------------------------------------- |
| 499 | # Tests for the full hybrid dispatch (HIP decode + Triton prefill) |
| 500 | # --------------------------------------------------------------------------- |
| 501 | |
| 502 | |
| 503 | @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") |
| 504 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) |
| 505 | @pytest.mark.parametrize( |
| 506 | "M,K,N,G", |
| 507 | [ |
| 508 | (1, 256, 256, 64), |
| 509 | (1, 512, 256, 32), |
| 510 | (1, 512, 256, 128), |
| 511 | (32, 512, 256, 64), |
| 512 | (64, 1024, 256, 128), |
| 513 | ], |
| 514 | ) |
| 515 | def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G): |
| 516 | """Test the full hybrid dispatch via the custom op.""" |
| 517 | from vllm.utils.platform_utils import num_compute_units |
| 518 | |
| 519 | set_random_seed(0) |
| 520 | |
| 521 | a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) |
| 522 | w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) |
| 523 | |
| 524 | b_packed_i32 = pack_int4_exllama_shuffle(w_int4_nk) |
| 525 | b_packed_i8 = b_packed_i32.view(torch.int8) |
| 526 | |
| 527 | scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to( |
| 528 | dtype |
| 529 | ) |
| 530 | |
| 531 | cu_count = num_compute_units() |
| 532 | out = torch.ops.vllm.rdna_hybrid_w4a16_apply( |
| 533 | a, b_packed_i8, scales, None, None, cu_count, G |
| 534 | ) |
| 535 | |
| 536 | ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8) |
| 537 | |
| 538 | torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) |
| 539 | |
| 540 | |
| 541 | @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") |
| 542 | @pytest.mark.skipif( |
| 543 | not hasattr(torch.ops, "_rocm_C") |
| 544 | or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"), |
| 545 | reason="wvSplitK_int4_g not built", |
| 546 | ) |
| 547 | def test_wvsplitk_int4_g_rejects_unpacked_zero_points(): |
| 548 | """Act-dtype zero points raise instead of being read as packed words. |
| 549 | |
| 550 | Both layouts are 2D with compatible extents, so only the dtype separates |
| 551 | them; misreading one as the other returns wrong numbers silently. |
| 552 | """ |
| 553 | import vllm._custom_ops as ops |
| 554 | from vllm.utils.platform_utils import num_compute_units |
| 555 | |
| 556 | K, N, G, M = 256, 64, 128, 1 |
| 557 | a = torch.randn((M, K), device=device, dtype=torch.float16) |
| 558 | w = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view( |
| 559 | torch.int8 |
| 560 | ) |
| 561 | scales = torch.rand((N, K // G), device=device, dtype=torch.float16) |
| 562 | zp_unpacked = torch.zeros((N, K // G), device=device, dtype=torch.float16) |
| 563 | |
| 564 | with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"): |
| 565 | ops.wvSplitK_int4_g(w, a, scales, num_compute_units(), G, zp_unpacked, None) |