| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shape configurations for IR op benchmarks.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | NUM_TOKENS = [1, 2, 4, 16, 64, 256, 1024, 4096, 16384] |
| 8 | COMMON_HIDDEN_SIZES = [ |
| 9 | 2048, # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n |
| 10 | 3072, # Gemma 7B/9B |
| 11 | 4096, # Llama 3 8B, Qwen 3 8B, Mistral 7B |
| 12 | 5120, # Llama 4 Scout 17B-16E |
| 13 | 7168, # DeepSeek V3 |
| 14 | 8192, # Llama 3 70B |
| 15 | 16384, # Llama 3 405B |
| 16 | ] |
| 17 | |
| 18 | # Each entry maps an op name to a list of kwarg dicts that will be passed |
| 19 | # to that op's registered input generator via op.generate_inputs(**kwargs). |
| 20 | SHAPE_CONFIGS: dict[str, list[dict]] = { |
| 21 | op_name: [ |
| 22 | {"num_tokens": n, "hidden_size": d, "dtype": dtype} |
| 23 | for dtype in [torch.float16, torch.bfloat16, torch.float32] |
| 24 | for d in COMMON_HIDDEN_SIZES |
| 25 | for n in NUM_TOKENS |
| 26 | ] |
| 27 | for op_name in ("rms_norm", "fused_add_rms_norm") |
| 28 | } |