| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared test utilities for vLLM IR op correctness tests.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.ir.op import IrOp |
| 8 | |
| 9 | NUM_TOKENS = [1, 8, 17, 32, 512, 2048] |
| 10 | COMMON_HIDDEN_SIZES = [ |
| 11 | 2048, # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n |
| 12 | 4096, # Llama 3 8B, Qwen 3 8B |
| 13 | 5120, # Llama 4 Scout 17B-16E |
| 14 | 7168, # DeepSeek V3 |
| 15 | 8192, # Llama 3 70B |
| 16 | ] |
| 17 | |
| 18 | |
| 19 | def clone_args(args: tuple) -> tuple: |
| 20 | return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args) |
| 21 | |
| 22 | |
| 23 | def supported_providers(op: IrOp) -> list[str]: |
| 24 | return [ |
| 25 | name for name, impl in op.impls.items() if name != "native" and impl.supported |
| 26 | ] |
| 27 | |
| 28 | |
| 29 | def assert_close(op: IrOp, actual, expected): |
| 30 | if isinstance(actual, torch.Tensor): |
| 31 | tol = op.get_tolerance(actual.dtype) |
| 32 | try: |
| 33 | torch.testing.assert_close(actual, expected, **tol) |
| 34 | except AssertionError as e: |
| 35 | raise AssertionError( |
| 36 | f"{e}\n\nTo adjust tolerance, use:\n" |
| 37 | f" ir.ops.{op.name}.override_tolerance(" |
| 38 | f"{actual.dtype}, atol=..., rtol=...)" |
| 39 | ) from None |
| 40 | elif isinstance(actual, (tuple, list)): |
| 41 | for a, ex in zip(actual, expected): |
| 42 | assert_close(op, a, ex) |
| 43 | else: |
| 44 | assert actual == expected |