| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Kernel test utils.""" |
| 4 | |
| 5 | from collections.abc import Sequence |
| 6 | from typing import Any |
| 7 | from unittest.mock import patch |
| 8 | |
| 9 | import torch |
| 10 | from torch._prims_common import TensorLikeType |
| 11 | |
| 12 | from tests.kernels.quant_utils import native_w8a8_block_matmul |
| 13 | from vllm.model_executor.custom_op import op_registry |
| 14 | from vllm.model_executor.layers.activation import SiluAndMul |
| 15 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 16 | from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input |
| 17 | |
| 18 | # For now, disable "test_aot_dispatch_dynamic" since there are some |
| 19 | # bugs related to this test in PyTorch 2.4. |
| 20 | DEFAULT_OPCHECK_TEST_UTILS: tuple[str, ...] = ( |
| 21 | "test_schema", |
| 22 | "test_autograd_registration", |
| 23 | "test_faketensor", |
| 24 | ) |
| 25 | |
| 26 | ALL_OPCHECK_TEST_UTILS: tuple[str, ...] = ( |
| 27 | "test_schema", |
| 28 | "test_autograd_registration", |
| 29 | "test_faketensor", |
| 30 | "test_aot_dispatch_dynamic", |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | def _assert_accurate( |
| 35 | actual: torch.Tensor, |
| 36 | expected: torch.Tensor, |
| 37 | atol: float, |
| 38 | rtol: float = 0.0, |
| 39 | pass_rate: float = 0.99999, |
| 40 | max_violation_factor: float = 3.0, |
| 41 | ) -> None: |
| 42 | """Check numeric accuracy with pass-rate, max-error, and mean-error bounds.""" |
| 43 | a = actual.detach().float().flatten() |
| 44 | e = expected.detach().float().flatten() |
| 45 | |
| 46 | abs_err = (a - e).abs() |
| 47 | tol = atol + rtol * e.abs() |
| 48 | |
| 49 | rate = (abs_err <= tol).float().mean().item() |
| 50 | assert rate >= pass_rate, ( |
| 51 | f"Accuracy pass rate {rate:.6f} < {pass_rate} (atol={atol}, rtol={rtol})" |
| 52 | ) |
| 53 | |
| 54 | max_err = abs_err.max().item() |
| 55 | assert max_err <= max_violation_factor * atol, ( |
| 56 | f"Max absolute error {max_err:.6f} exceeds {max_violation_factor} * atol={atol}" |
| 57 | ) |
| 58 | |
| 59 | mean_err = abs_err.mean().item() |
| 60 | assert mean_err <= atol * 0.25, ( |
| 61 | f"Mean absolute error {mean_err:.6f} >= atol * 0.25 = {atol * 0.25:.6f}" |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | def _assert_deterministic( |
| 66 | fn, |
| 67 | *args, |
| 68 | n_runs: int = 4, |
| 69 | **kwargs, |
| 70 | ) -> None: |
| 71 | """Verify that repeated calls produce bitwise-identical tensor outputs.""" |
| 72 | |
| 73 | def _collect(result: Any) -> list[torch.Tensor]: |
| 74 | if isinstance(result, torch.Tensor): |
| 75 | return [result.detach().clone()] |
| 76 | if isinstance(result, (tuple, list)): |
| 77 | return [t.detach().clone() for t in result if isinstance(t, torch.Tensor)] |
| 78 | raise TypeError(f"Unexpected return type {type(result)}") |
| 79 | |
| 80 | reference = _collect(fn(*args, **kwargs)) |
| 81 | |
| 82 | for run in range(1, n_runs): |
| 83 | outputs = _collect(fn(*args, **kwargs)) |
| 84 | for idx, (ref, out) in enumerate(zip(reference, outputs)): |
| 85 | assert torch.equal(ref, out), ( |
| 86 | f"Run {run}: output[{idx}] differs from run 0 " |
| 87 | f"(max diff = {(out.float() - ref.float()).abs().max().item():.2e})" |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | # Copied/modified from torch._refs.__init__.py |
| 92 | def fp8_allclose( |
| 93 | a: TensorLikeType, |
| 94 | b: TensorLikeType, |
| 95 | rtol: float = 1e-05, |
| 96 | atol: float = 1e-08, |
| 97 | equal_nan: bool = False, |
| 98 | ) -> bool: |
| 99 | """Reference implementation of torch.allclose.""" |
| 100 | torch._refs._check_close_args(name="torch.allclose", a=a, b=b, rtol=rtol, atol=atol) |
| 101 | |
| 102 | return bool( |
| 103 | torch.all( |
| 104 | torch.isclose( |
| 105 | a.double(), b.double(), rtol=rtol, atol=atol, equal_nan=equal_nan |
| 106 | ) |
| 107 | ).item() |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | def bf16_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: |
| 112 | """Representable-step distance between two bf16 tensors. |
| 113 | |
| 114 | Reinterprets the bf16 bit patterns under the IEEE-754 total ordering so |
| 115 | that adjacent representable values differ by exactly 1. |
| 116 | """ |
| 117 | |
| 118 | def key(t: torch.Tensor) -> torch.Tensor: |
| 119 | u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF |
| 120 | return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000) |
| 121 | |
| 122 | return (key(a) - key(b)).abs() |
| 123 | |
| 124 | |
| 125 | def fp8_ulp_distance(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: |
| 126 | """Representable-step distance between two 8-bit fp8 tensors. |
| 127 | |
| 128 | Reinterprets the fp8 bytes under a sign-magnitude total ordering so that |
| 129 | adjacent representable values differ by exactly 1. Inputs must already share |
| 130 | the same fp8 encoding (e.g. both FP8_STORE_DTYPE). |
| 131 | """ |
| 132 | |
| 133 | def key(t: torch.Tensor) -> torch.Tensor: |
| 134 | u = t.contiguous().view(torch.uint8).to(torch.int64) |
| 135 | return torch.where(u >= 0x80, 0xFF - u, u + 0x80) |
| 136 | |
| 137 | return (key(a) - key(b)).abs() |
| 138 | |
| 139 | |
| 140 | # Marlin MoE test utils |
| 141 | |
| 142 | |
| 143 | def stack_and_dev(tensors: list[torch.Tensor]): |
| 144 | dev = tensors[0].device |
| 145 | return torch.stack(tensors, dim=0).to(dev) |
| 146 | |
| 147 | |
| 148 | def compute_max_diff(output, output_ref): |
| 149 | return torch.mean(torch.abs(output - output_ref)) / torch.mean( |
| 150 | torch.abs(output_ref) |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | def torch_experts( |
| 155 | a: torch.Tensor, |
| 156 | w1: torch.Tensor, |
| 157 | w2: torch.Tensor, |
| 158 | topk_weight: torch.Tensor, |
| 159 | topk_ids: torch.Tensor, |
| 160 | global_num_experts: int = -1, |
| 161 | b_bias1: torch.Tensor | None = None, |
| 162 | b_bias2: torch.Tensor | None = None, |
| 163 | expert_map: torch.Tensor | None = None, |
| 164 | w1_scale: torch.Tensor | None = None, |
| 165 | w2_scale: torch.Tensor | None = None, |
| 166 | a1_scale: torch.Tensor | None = None, |
| 167 | a2_scale: torch.Tensor | None = None, |
| 168 | quant_dtype: torch.dtype | None = None, |
| 169 | per_act_token_quant=False, |
| 170 | block_shape: list[int] | None = None, |
| 171 | apply_router_weights_on_input: bool = False, |
| 172 | activation: MoEActivation = MoEActivation.SILU, |
| 173 | ) -> torch.Tensor: |
| 174 | assert ( |
| 175 | global_num_experts == -1 |
| 176 | or (global_num_experts == w1.shape[0] and expert_map is None) |
| 177 | or (expert_map is not None and global_num_experts == expert_map.shape[0]) |
| 178 | ) |
| 179 | |
| 180 | if quant_dtype in [torch.float16, torch.bfloat16]: |
| 181 | quant_dtype = None |
| 182 | quant_input_only = quant_dtype is not None and w1_scale is None and w2_scale is None |
| 183 | if quant_input_only: |
| 184 | assert a1_scale is None and a2_scale is None |
| 185 | assert per_act_token_quant |
| 186 | |
| 187 | M, K = a.shape |
| 188 | topk = topk_ids.shape[1] |
| 189 | |
| 190 | if apply_router_weights_on_input: |
| 191 | assert topk == 1 |
| 192 | a = a * topk_weight.to(a.dtype) |
| 193 | |
| 194 | a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K) |
| 195 | |
| 196 | out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device) |
| 197 | |
| 198 | if a1_scale: |
| 199 | assert not per_act_token_quant and block_shape is None |
| 200 | a, a_scale = moe_kernel_quantize_input( |
| 201 | a, a1_scale, quant_dtype, per_act_token_quant, block_shape |
| 202 | ) |
| 203 | |
| 204 | if quant_input_only: |
| 205 | a = (a.float() * a_scale.view(-1, 1)).to(w1.dtype) |
| 206 | |
| 207 | num_experts = w1.shape[0] |
| 208 | |
| 209 | topk_ids = topk_ids.view(-1) |
| 210 | if expert_map is not None: |
| 211 | topk_ids = expert_map[topk_ids] |
| 212 | |
| 213 | f32 = torch.float32 |
| 214 | |
| 215 | act = op_registry[activation.custom_op_name] |
| 216 | |
| 217 | for i in range(num_experts): |
| 218 | mask = topk_ids == i |
| 219 | if mask.sum(): |
| 220 | if quant_dtype is None: |
| 221 | tmp1 = a[mask] @ w1[i].transpose(0, 1) |
| 222 | if b_bias1 is not None: |
| 223 | tmp1 = tmp1 + b_bias1[i].view(1, -1).to(tmp1.dtype) |
| 224 | tmp2 = act()(tmp1) |
| 225 | out[mask] = tmp2 @ w2[i].transpose(0, 1) |
| 226 | if b_bias2 is not None: |
| 227 | out[mask] = out[mask] + b_bias2[i].view(1, -1).to(tmp1.dtype) |
| 228 | elif quant_input_only: |
| 229 | tmp1 = a[mask] @ w1[i].transpose(0, 1) |
| 230 | tmp2 = SiluAndMul()(tmp1) |
| 231 | tmp2, tmp2_scale = moe_kernel_quantize_input( |
| 232 | tmp2, None, quant_dtype, per_act_token_quant |
| 233 | ) |
| 234 | tmp2 = (tmp2.float() * tmp2_scale.view(-1, 1)).to(w2.dtype) |
| 235 | out[mask] = tmp2 @ w2[i].transpose(0, 1) |
| 236 | elif block_shape is not None: |
| 237 | # block quantized |
| 238 | assert ( |
| 239 | a_scale is not None |
| 240 | and w1_scale is not None |
| 241 | and w2_scale is not None |
| 242 | ) |
| 243 | tmp1 = native_w8a8_block_matmul( |
| 244 | a[mask], w1[i], a_scale[mask], w1_scale[i], block_shape, out.dtype |
| 245 | ) |
| 246 | if b_bias1 is not None: |
| 247 | tmp1 = tmp1 + b_bias1[i].view(1, -1).to(tmp1.dtype) |
| 248 | tmp2 = SiluAndMul()(tmp1) |
| 249 | tmp2, b_scale = moe_kernel_quantize_input( |
| 250 | tmp2, a2_scale, quant_dtype, per_act_token_quant, block_shape |
| 251 | ) |
| 252 | |
| 253 | out[mask] = native_w8a8_block_matmul( |
| 254 | tmp2, w2[i], b_scale, w2_scale[i], block_shape, out.dtype |
| 255 | ) |
| 256 | if b_bias2 is not None: |
| 257 | out[mask] = out[mask] + b_bias2[i].view(1, -1).to(tmp1.dtype) |
| 258 | else: |
| 259 | assert ( |
| 260 | a_scale is not None |
| 261 | and w1_scale is not None |
| 262 | and w2_scale is not None |
| 263 | ) |
| 264 | scales = a_scale if a_scale.numel() == 1 else a_scale[mask] |
| 265 | |
| 266 | tmp1 = a[mask].to(f32) * scales |
| 267 | w1_dq = (w1[i].to(f32) * w1_scale[i]).transpose(0, 1) |
| 268 | tmp1 = (tmp1 @ w1_dq).to(out.dtype) |
| 269 | if b_bias1 is not None: |
| 270 | tmp1 = tmp1 + b_bias1[i].view(1, -1).to(out.dtype) |
| 271 | |
| 272 | tmp2 = act()(tmp1).to(out.dtype) |
| 273 | |
| 274 | tmp2, b_scale = moe_kernel_quantize_input( |
| 275 | tmp2, a2_scale, quant_dtype, per_act_token_quant, block_shape |
| 276 | ) |
| 277 | assert b_scale is not None |
| 278 | |
| 279 | tmp2 = tmp2.to(f32) * b_scale |
| 280 | w2_dq = (w2[i].to(f32) * w2_scale[i]).transpose(0, 1) |
| 281 | out[mask] = (tmp2 @ w2_dq).to(out.dtype) |
| 282 | if b_bias2 is not None: |
| 283 | out[mask] = out[mask] + b_bias2[i].view(1, -1).to(out.dtype) |
| 284 | |
| 285 | if apply_router_weights_on_input: |
| 286 | return out |
| 287 | else: |
| 288 | return ( |
| 289 | (out.view(M, -1, w2.shape[1]).to(f32) * topk_weight.view(M, -1, 1)) |
| 290 | .sum(dim=1) |
| 291 | .to(out.dtype) |
| 292 | ) |
| 293 | |
| 294 | |
| 295 | def torch_moe( |
| 296 | a: torch.Tensor, |
| 297 | w1: torch.Tensor, |
| 298 | w2: torch.Tensor, |
| 299 | score: torch.Tensor, |
| 300 | topk: int, |
| 301 | b_bias1: torch.Tensor | None = None, |
| 302 | b_bias2: torch.Tensor | None = None, |
| 303 | global_num_experts: int = -1, |
| 304 | expert_map: torch.Tensor | None = None, |
| 305 | activation: MoEActivation = MoEActivation.SILU, |
| 306 | ) -> torch.Tensor: |
| 307 | score = torch.softmax(score, dim=-1, dtype=torch.float32) |
| 308 | topk_weight, topk_ids = torch.topk(score, topk) |
| 309 | return torch_experts( |
| 310 | a, |
| 311 | w1, |
| 312 | w2, |
| 313 | topk_weight, |
| 314 | topk_ids, |
| 315 | global_num_experts, |
| 316 | b_bias1, |
| 317 | b_bias2, |
| 318 | expert_map, |
| 319 | activation=activation, |
| 320 | ) |
| 321 | |
| 322 | |
| 323 | def torch_moe_single(a, w, score, topk): |
| 324 | B, D = a.shape |
| 325 | a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) |
| 326 | out = torch.zeros(B * topk, w.shape[1], dtype=a.dtype, device=a.device) |
| 327 | score = torch.softmax(score, dim=-1, dtype=torch.float32) |
| 328 | _, topk_ids = torch.topk(score, topk) |
| 329 | topk_ids = topk_ids.view(-1) |
| 330 | for i in range(w.shape[0]): |
| 331 | mask = topk_ids == i |
| 332 | if mask.sum(): |
| 333 | out[mask] = a[mask] @ w[i].transpose(0, 1) |
| 334 | return (out.view(B, -1, w.shape[1])).sum(dim=1) |
| 335 | |
| 336 | |
| 337 | # A special version of op check that has a restricted default set of test_utils |
| 338 | # and a patched version of allclose that supports fp8 types. |
| 339 | def opcheck( |
| 340 | op: torch._ops.OpOverload |
| 341 | | torch._ops.OpOverloadPacket |
| 342 | | torch._library.custom_ops.CustomOpDef, |
| 343 | args: tuple[Any, ...], |
| 344 | kwargs: dict[str, Any] | None = None, |
| 345 | *, |
| 346 | test_utils: str | Sequence[str] = ALL_OPCHECK_TEST_UTILS, |
| 347 | raise_exception: bool = True, |
| 348 | cond: bool = True, |
| 349 | ) -> dict[str, str]: |
| 350 | with patch("torch.allclose", new=fp8_allclose): |
| 351 | return ( |
| 352 | torch.library.opcheck( |
| 353 | op, args, kwargs, test_utils=test_utils, raise_exception=raise_exception |
| 354 | ) |
| 355 | if cond |
| 356 | else {} |
| 357 | ) |
| 358 | |
| 359 | |
| 360 | # For testing quantized linear kernels |
| 361 | def to_fp8(tensor: torch.Tensor): |
| 362 | finfo = torch.finfo(torch.float8_e4m3fn) |
| 363 | return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to( |
| 364 | dtype=torch.float8_e4m3fn |
| 365 | ) |
| 366 | |
| 367 | |
| 368 | def to_int8(tensor: torch.Tensor): |
| 369 | return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8) |
| 370 | |
| 371 | |
| 372 | def baseline_scaled_mm( |
| 373 | a: torch.Tensor, |
| 374 | b: torch.Tensor, |
| 375 | scale_a: torch.Tensor, |
| 376 | scale_b: torch.Tensor, |
| 377 | out_dtype: type[torch.dtype], |
| 378 | bias: torch.Tensor | None = None, |
| 379 | ) -> torch.Tensor: |
| 380 | # We treat N-dimensional group scaling as extended numpy-style broadcasting |
| 381 | # in numpy simply stretches dimensions with an extent of 1 to match |
| 382 | # the target shape by repeating the data along that dimension (broadcasting) |
| 383 | # , we extend these semantics to say if the extent of a dimension in the |
| 384 | # source shape is not 1 and does not match the target shape we repeat each |
| 385 | # element along that dimension src_shape[dim] // target_shape[dim] times |
| 386 | # example if we have: |
| 387 | # a = [[1, 2], and target_shape = (2, 4) |
| 388 | # [3, 4]] |
| 389 | # then we would expand a to: |
| 390 | # a = [[1, 1, 2, 2], |
| 391 | # [3, 3, 4, 4]] |
| 392 | # NOTE this function does not explicitly broadcast dimensions |
| 393 | # with an extent of 1, since this can be done implicitly by pytorch |
| 394 | def group_broadcast(t, shape): |
| 395 | for i, s in enumerate(shape): |
| 396 | if t.shape[i] != s and t.shape[i] != 1: |
| 397 | assert s % t.shape[i] == 0 |
| 398 | t = ( |
| 399 | t.unsqueeze(i + 1) |
| 400 | .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :]) |
| 401 | .flatten(i, i + 1) |
| 402 | ) |
| 403 | return t |
| 404 | |
| 405 | scale_a = group_broadcast(scale_a, a.shape) |
| 406 | scale_b = group_broadcast(scale_b, b.shape) |
| 407 | |
| 408 | output = torch.mm( |
| 409 | (scale_a * a.to(dtype=torch.float32)), (scale_b * b.to(dtype=torch.float32)) |
| 410 | ).to(out_dtype) |
| 411 | |
| 412 | if bias is not None: |
| 413 | output = output + bias |
| 414 | |
| 415 | return output |