diff --git a/.agents/skills/benchmark-model-kernels/SKILL.md b/.agents/skills/benchmark-model-kernels/SKILL.md new file mode 100644 index 00000000000..1ff4f5af571 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/SKILL.md @@ -0,0 +1,145 @@ +--- +name: benchmark-model-kernels +description: >- + Inspect Hugging Face decoder layers on meta tensors and plan or run per-rank + BF16, FP8, and NVFP4 GEMM or fused-MoE microbenchmarks with the bundled + scripts and a local FlashInfer checkout. Use when choosing a model, GPU, TP, + EP, or M/token-concurrency sweep; deriving common fused QKV and gate/up + shapes without loading checkpoint weights; or using FlashInfer benchmark + utilities. Do not use for end-to-end server throughput or request latency. +--- + +# Benchmark Model Kernels + +Plan a single-GPU microbenchmark for the per-rank shapes of an intended +deployment. The scripts never load checkpoint weights, launch distributed +workers, or measure collectives, serving throughput, or request latency. + +## Interview, one decision per message + +Ask exactly one unresolved decision per message with two or three concrete +options, state the default, and wait for the answer. Recommend an option only +when it is substantively better. Skip decisions already answered. Follow this +order: + +1. **Model** — a local directory, a `config.json`, or a Hub ID (equivalent; no + default). The script builds the model on meta tensors from configuration + only. Do not enable `--trust-remote-code` without explicit approval; pin + `--revision ` when it is needed. +2. **TP and EP** — accept both directly, or derive them from the intended + deployment's GPU model and count and confirm. Defaults: `TP=1`, `EP=1`. + Derive parallelism from the intended deployment, never from the GPU that + runs the microbenchmark. The script validates every sharding rule + (divisibility, GQA replication, expert partitioning) and errors loudly. +3. **M sweep** — balanced default (recommended): `1 8 64 512`; decode-focused: + `1 4 16 32`; throughput-focused: `64 256 1024 4096`. M is roughly the + tokens scheduled per step (decode: active sequences), not endpoint + concurrency. With EP, an MoE row models one rank's share of the global + batch: a global batch of B tokens corresponds to the column M = B/EP, so + do not compare different EP values at the same M. +4. **Shape preview** — always run this before anything else; no FlashInfer or + GPU needed: + + ```bash + python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \ + --tp --ep --ms ... --print_only + ``` + + Review the printed shapes and the MoE tuple with the user. + `# unsupported:` lines mean the list is partial and the script exits + nonzero — handle those via **Manual supplements** below. +5. **FlashInfer checkout and GPU** — the full benchmark needs a FlashInfer + *source checkout* containing `benchmarks/flashinfer_benchmark.py`; the + installed wheel alone is not enough. Prefer a clean checkout matching the + installed `flashinfer` version; ask before cloning or installing anything. + On the benchmark machine, check `nvidia-smi` and package versions, verify + the target GPU is idle (concurrent work on the same GPU skews timings), + and verify CUPTI timing with a tiny `bench_gpu_time(..., enable_cupti=True)` + probe — a warning that falls back to CUDA events is a failure (the + `cupti-python`/`nvidia-cuda-cupti` packages must match PyTorch's CUDA + major). Ask for a GPU index only when several are visible. Pick a fresh + workdir and state it. +6. **Full benchmark**: + + ```bash + CUDA_VISIBLE_DEVICES= \ + python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \ + --tp --ep --ms ... \ + --flashinfer_repo --workdir + ``` + + Run a short plumbing check first (append + `--dry_run_iters 1 --num_iters 3 --no_autotune`, throwaway workdir), + especially after changing FlashInfer versions or shape logic; never present + it as a performance result. Then run with defaults. The first fused-MoE + build or autotune can take several minutes; its cache is reused. + +## Rules the scripts own + +Do not restate, re-derive, or override these — the code enforces them: + +- `benchmark_model.py` derives `--nks`, `--nk_names`, and all `--moe_*` + arguments; overrides are rejected. +- Row labels are logical shapes; the runner applies vLLM's physical padding. + Report both shapes whenever they differ. +- A failed case writes FlashInfer's error message (or a pointer to + `driver.log`) into its `combined_results.csv` cell, and the command exits + nonzero after the table is written. Never present a partial table as a + successful benchmark; read `driver.log` before rerunning anything. + +## Known backend and driver limits + +- `mm_fp4` TensorRT-LLM needs `N % 128 == 0` (shuffled weight layout): its + cell reports no result row, or on stock 0.6.x drivers an empty assertion + that fails every `mm_fp4` backend for that shape. +- `mm_fp8` (trtllm_low_latency) needs `K % 128 == 0`. +- MoE rows cover FlashInfer's CUTLASS fused MoE, the trtllm-gen NVFP4 and + per-tensor FP8 MoE, and the CuteDSL NVFP4 MoE. The CuteDSL MoE row appears + only for Swiglu models (the kernel supports nothing else); the `cutedsl` and + `trtllm` backends require recent GPUs and report per-case errors elsewhere. +- Gated NVFP4 CUTLASS MoE with `2F % 128 != 0` per rank fails — vLLM raises + instead of padding — often as a `CUDA error: misaligned address`. Prefer EP + over TP for the experts to keep the per-rank width legal. +- Do not benchmark a padded shape to dodge these limits and call it + vLLM-equivalent; report the limit instead. + +## Manual supplements + +For each `# unsupported:` layout from the preview: inspect that module's +forward path and the intended runtime's TP/EP sharding, derive the per-rank +shape (never guess sharding from a weight shape alone), and benchmark only the +missing shape: + +```bash +CUDA_VISIBLE_DEVICES= \ +python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \ + --flashinfer_repo --ms ... \ + --nks , --nk_names --workdir +``` + +Label these rows as manual supplements; this is not a substitute for running +`benchmark_model.py` first. Shell-quote user-supplied paths. + +## Report + +Report the command, GPU, versions, TP/EP, M values, shapes (logical and +physical where they differ), warnings, and the artifacts: `testlist.txt`, +`driver.log` (full driver output), `builtin_results.csv` (milliseconds, +success-only), and `combined_results.csv` (long form, microseconds: columns `module_name, +M, N, K, backend, with_quant, runtime` in `GEMM` and `MoE` sections; modules +fused into one GEMM are joined with `|` in `module_name`, while distinct +same-shape modules appear as duplicated rows sharing one measurement; MoE +rows keep the `H= F= E= top_k=` parameter line and leave N/K empty). In quantization-recipe terms: +`bf16` rows are the unquantized W16A16 baseline, `fp8` rows are per-tensor +W8A8, and `nvfp4` rows are W4A4. Plain quantized rows time the kernel with +pre-quantized activations; `*_with_quant` rows add a separately measured +activation-quantization time in the scale-factor layout that backend +consumes, except the NVFP4 CUTLASS MoE row, which is a single fused +measurement. MoE routing is synthetic: uniform expert +distribution everywhere (real skewed routing is slower), and the trtllm-gen +rows, which route in-kernel, use a fixed `renormalize` method regardless of +the model's routing scheme. CUTLASS and CuteDSL MoE rows receive precomputed +expert indices and exclude routing-selection cost entirely. No row includes +the router GEMM itself. These are kernel times +only — never describe them as end-to-end latency or throughput; they omit +weights, layer frequency, communication, KV cache, and scheduling. diff --git a/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py new file mode 100644 index 00000000000..2cd7aa0019d --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py @@ -0,0 +1,758 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Derive per-rank benchmark shapes from a Transformers model on meta tensors. + +The script walks the instantiated model's Linear modules, fuses Q/K/V and +gate/up projections, recognizes Mamba 2, GatedDeltaNet, and common +routed-expert layouts, and applies the common serving/export TP layout. It +never calls a checkpoint weight loader. +When a decoder layout is unsupported, the derived shapes are still printed and +the script exits nonzero; benchmark the missing shapes directly with +benchmark_via_builtin.py. +""" + +import argparse +import importlib.util +import re +import shlex +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +class ShapeError(ValueError): + """A model layer cannot be represented by the supported benchmark layout.""" + + +@dataclass(frozen=True) +class _MoeShape: + hidden: int + intermediate: int + experts: int + top_k: int + activation: str | None = None + + +@dataclass(frozen=True) +class _ExpertShape: + hidden: int + intermediate: int + gated: bool + + +_Kernel = tuple[int, int, str] + +_PROJECTIONS = { + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "gate_up_proj", + "down_proj", +} +_RESERVED = { + "--nks", + "--nk_names", + "--moe_hidden_size", + "--moe_intermediate_size", + "--moe_num_experts", + "--moe_top_k", + "--moe_activation_type", +} + +# Modules deliberately outside the benchmarked GEMM list fall into two +# exclusion mechanisms: +# 1. Positional: embeddings, the LM head, and anything else outside the +# decoder blocks never enter the audit (the `layers.` position filter +# in _unsupported_decoder_linears). +# 2. Name-based: routing and gating projections that live inside decoder +# blocks are excluded by these names. They are never quantized in +# deployment recipes and vLLM dispatches them through specialized (often +# FP32-output) paths, so a standard GEMM row would not model them anyway. +_ROUTER_PATH_PARTS = {"router", "routers"} +_GATING_LEAF_NAMES = {"gate", "router", "router_proj", "shared_expert_gate"} + + +def _positive_int(value: str) -> int: + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError(f"expected a positive integer, got {value!r}") from exc + if parsed <= 0: + raise argparse.ArgumentTypeError(f"expected a positive integer, got {value!r}") + return parsed + + +def _load_meta_model(model_ref: str, trust_remote_code: bool, revision: str | None): + # Transformers and Accelerate are optional, heavy ModelOpt dependencies. + try: + from accelerate import init_empty_weights + from transformers import AutoConfig, AutoModelForCausalLM + except ImportError as exc: + raise ShapeError("install ModelOpt with the 'hf' extra") from exc + + path = Path(model_ref).expanduser() + ref = str(path) if path.exists() else model_ref + try: + config = AutoConfig.from_pretrained( + ref, trust_remote_code=trust_remote_code, revision=revision + ) + model_kwargs = {"trust_remote_code": trust_remote_code} + auto_map = getattr(config, "auto_map", {}) or {} + if revision and trust_remote_code and "AutoModelForCausalLM" in auto_map: + model_kwargs["code_revision"] = revision + with init_empty_weights(include_buffers=True): + try: + model = AutoModelForCausalLM.from_config(config, **model_kwargs) + except Exception: + text_config = getattr(config, "text_config", None) + if text_config is None: + raise + # Multimodal wrapper configs build their text decoder + # directly; vision towers are outside the benchmark scope. + model = AutoModelForCausalLM.from_config(text_config, **model_kwargs) + except Exception as exc: + raise ShapeError(f"could not construct {model_ref!r} on meta tensors: {exc}") from exc + + tensors = list(model.named_parameters()) + list(model.named_buffers()) + materialized = [name for name, tensor in tensors if not tensor.is_meta] + if materialized: + raise ShapeError(f"model construction allocated tensors: {', '.join(materialized[:3])}") + return config, model + + +def _linear_shape(module: Any) -> tuple[int, int] | None: + if hasattr(module, "out_features") and hasattr(module, "in_features"): + return int(module.out_features), int(module.in_features) + return None + + +def _divide(value: int, size: int, label: str) -> int: + if value % size: + raise ShapeError(f"{label}={value} is not divisible by {size}") + return value // size + + +def _path_label(parent: str, leaf: str) -> str: + """Label a GEMM by its normalized module path and vLLM fused-module name.""" + path = f"{parent}.{leaf}" if parent else leaf + return re.sub(r"(?<=\.)\d+(?=\.|$)", "*", path) + + +def _fused_qkv( + q: tuple[int, int], + k: tuple[int, int], + v: tuple[int, int], + head_dim: int, + tp: int, + parent: str, +) -> _Kernel: + if q[1] != k[1] or q[1] != v[1] or k[0] != v[0]: + raise ShapeError(f"unsupported Q/K/V shapes under {parent}") + q_heads = _divide(q[0], head_dim, f"{parent}.q_proj output") + kv_heads = _divide(k[0], head_dim, f"{parent}.k_proj output") + if kv_heads >= tp: + _divide(q_heads, tp, f"{parent}.q_proj heads") + _divide(kv_heads, tp, f"{parent}.k_proj heads") + local_n = _divide(q[0] + k[0] + v[0], tp, f"{parent}.qkv") + else: + local_q = _divide(q_heads, tp, f"{parent}.q_proj heads") * head_dim + _divide(tp, kv_heads, "TP/KV replication ratio") + local_n = local_q + 2 * head_dim + label = "|".join(_path_label(parent, leaf) for leaf in ("q_proj", "k_proj", "v_proj")) + return local_n, q[1], label + + +def _dense_kernels(model: Any, config: Any, tp: int) -> list[_Kernel]: + groups: dict[str, dict[str, tuple[int, int]]] = {} + for name, module in model.named_modules(): + leaf = name.rsplit(".", 1)[-1] + parts = name.split(".") + # Routed experts are excluded here and handled by _moe_shapes. Shared + # experts (for example `.shared_experts.up_proj`) intentionally do not + # match the filter: they run densely for every token, so their + # projections belong in the dense GEMM list. + if ( + leaf not in _PROJECTIONS + or ".experts." in name + or ".local_experts." in name + or any(part in _ROUTER_PATH_PARTS for part in parts[:-1]) + ): + continue + shape = _linear_shape(module) + if shape: + groups.setdefault(name.rpartition(".")[0], {})[leaf] = shape + + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + head_dim = config.hidden_size // config.num_attention_heads + kernels = [] + for parent, layers in groups.items(): + qkv = {"q_proj", "k_proj", "v_proj"} + present_qkv = qkv.intersection(layers) + if present_qkv: + if present_qkv != qkv: + raise ShapeError(f"incomplete Q/K/V projections under {parent}") + kernels.append( + _fused_qkv( + layers["q_proj"], + layers["k_proj"], + layers["v_proj"], + int(head_dim), + tp, + parent, + ) + ) + if "o_proj" in layers: + n, k = layers["o_proj"] + kernels.append((n, _divide(k, tp, f"{parent}.o_proj"), _path_label(parent, "o_proj"))) + + if "gate_proj" in layers and "up_proj" in layers: + gate_n, gate_k = layers["gate_proj"] + up_n, up_k = layers["up_proj"] + if gate_k != up_k: + raise ShapeError(f"gate/up inputs differ under {parent}") + # vLLM shards gate and up individually before merging them, so + # each output must divide by TP, not just their sum. + n = _divide(gate_n, tp, f"{parent}.gate_proj") + _divide(up_n, tp, f"{parent}.up_proj") + label = "|".join(_path_label(parent, leaf) for leaf in ("gate_proj", "up_proj")) + kernels.append((n, gate_k, label)) + elif "gate_proj" in layers: + raise ShapeError(f"gate projection has no matching up projection under {parent}") + elif "up_proj" in layers: + n, k = layers["up_proj"] + kernels.append((_divide(n, tp, f"{parent}.up_proj"), k, _path_label(parent, "up_proj"))) + if "gate_up_proj" in layers: + n, k = layers["gate_up_proj"] + kernels.append( + (_divide(n, tp, f"{parent}.gate_up_proj"), k, _path_label(parent, "gate_up_proj")) + ) + if "down_proj" in layers: + n, k = layers["down_proj"] + kernels.append( + (n, _divide(k, tp, f"{parent}.down_proj"), _path_label(parent, "down_proj")) + ) + return list(dict.fromkeys(kernels)) + + +def _mamba_layout( + module: Any, +) -> tuple[tuple[int, int], tuple[int, int], int, int, int, int] | None: + in_shape = _linear_shape(getattr(module, "in_proj", None)) + out_shape = _linear_shape(getattr(module, "out_proj", None)) + attrs = ( + getattr(module, "intermediate_size", None), + getattr(module, "num_heads", None), + getattr(module, "n_groups", None), + getattr(module, "ssm_state_size", None), + ) + if in_shape is None or out_shape is None or any(value is None for value in attrs): + return None + intermediate, heads, groups, state = (int(value) for value in attrs if value is not None) + return in_shape, out_shape, intermediate, heads, groups, state + + +def _mamba_kernels(model: Any, tp: int) -> list[_Kernel]: + kernels = [] + for name, module in model.named_modules(): + layout = _mamba_layout(module) + if layout is None: + continue + in_shape, out_shape, intermediate, heads, groups, state = layout + hidden = in_shape[1] + expected_in = 2 * intermediate + 2 * groups * state + heads + if in_shape[0] != expected_in or out_shape != (hidden, intermediate): + raise ShapeError(f"unsupported Mamba projection shapes under {name}") + + local_intermediate = _divide(intermediate, tp, f"{name}.intermediate_size") + local_heads = _divide(heads, tp, f"{name}.num_heads") + if groups % tp == 0: + local_groups = groups // tp + elif groups == 1: + local_groups = 1 + else: + raise ShapeError(f"{name}.n_groups={groups} is not divisible by TP={tp}") + local_in = 2 * local_intermediate + 2 * local_groups * state + local_heads + kernels.extend( + [ + (local_in, hidden, _path_label(name, "in_proj")), + (hidden, local_intermediate, _path_label(name, "out_proj")), + ] + ) + return list(dict.fromkeys(kernels)) + + +def _gdn_layout(module: Any) -> tuple[Any, ...] | None: + out_shape = _linear_shape(getattr(module, "out_proj", None)) + attrs = ( + getattr(module, "num_k_heads", None), + getattr(module, "num_v_heads", None), + getattr(module, "key_dim", None), + getattr(module, "value_dim", None), + ) + has_input = ( + getattr(module, "in_proj_qkvz", None) is not None + or getattr(module, "in_proj_qkv", None) is not None + ) + if out_shape is None or not has_input or any(value is None for value in attrs): + return None + return (out_shape, *(int(value) for value in attrs if value is not None)) + + +def _gdn_kernels(model: Any, tp: int) -> list[_Kernel]: + kernels = [] + for name, module in model.named_modules(): + layout = _gdn_layout(module) + if layout is None: + continue + out_shape, num_k_heads, num_v_heads, key_dim, value_dim = layout + hidden = out_shape[0] + fused_qkvz = _linear_shape(getattr(module, "in_proj_qkvz", None)) + qkvz_label = _path_label(name, "in_proj_qkvz") + ba_label = _path_label(name, "in_proj_ba") + if fused_qkvz is not None: + # Qwen3-Next stores qkvz and ba pre-fused. + ba = _linear_shape(getattr(module, "in_proj_ba", None)) + valid = fused_qkvz == (2 * key_dim + 2 * value_dim, hidden) and ba == ( + 2 * num_v_heads, + hidden, + ) + else: + # Qwen3.5 stores them split, but vLLM's shared GDN mixer merges + # qkv+z and b+a into the same two per-rank GEMMs either way. + qkvz_label = "|".join(_path_label(name, leaf) for leaf in ("in_proj_qkv", "in_proj_z")) + ba_label = "|".join(_path_label(name, leaf) for leaf in ("in_proj_b", "in_proj_a")) + shapes = { + leaf: _linear_shape(getattr(module, leaf, None)) + for leaf in ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a") + } + valid = ( + shapes["in_proj_qkv"] == (2 * key_dim + value_dim, hidden) + and shapes["in_proj_z"] == (value_dim, hidden) + and shapes["in_proj_b"] == (num_v_heads, hidden) + and shapes["in_proj_a"] == (num_v_heads, hidden) + ) + if not valid or out_shape != (hidden, value_dim): + raise ShapeError(f"unsupported GatedDeltaNet projection shapes under {name}") + _divide(num_k_heads, tp, f"{name}.linear_num_key_heads") + local_v_heads = _divide(num_v_heads, tp, f"{name}.linear_num_value_heads") + kernels.extend( + [ + ( + _divide(2 * key_dim + 2 * value_dim, tp, f"{name}.qkvz"), + hidden, + qkvz_label, + ), + (2 * local_v_heads, hidden, ba_label), + ( + hidden, + _divide(value_dim, tp, f"{name}.value_dim"), + _path_label(name, "out_proj"), + ), + ] + ) + return list(dict.fromkeys(kernels)) + + +def _expert_shape(module: Any) -> _ExpertShape | None: + gate = getattr(module, "gate_proj", None) + up = getattr(module, "up_proj", None) + down = getattr(module, "down_proj", None) + if gate is None and getattr(module, "w1", None) is not None: + gate = getattr(module, "w1", None) + up = getattr(module, "w3", None) + down = getattr(module, "w2", None) + if gate is not None: + gate_shape, up_shape, down_shape = map(_linear_shape, (gate, up, down)) + if gate_shape is None or up_shape is None or down_shape is None: + raise ShapeError("incomplete gated expert Linear layout") + if gate_shape != up_shape or down_shape != (gate_shape[1], gate_shape[0]): + raise ShapeError("unsupported gated expert Linear shapes") + return _ExpertShape(gate_shape[1], gate_shape[0], True) + + up_shape, down_shape = map(_linear_shape, (up, down)) + if up_shape is None and down_shape is None: + return None + if up_shape is None or down_shape is None or down_shape != (up_shape[1], up_shape[0]): + raise ShapeError("unsupported non-gated expert Linear shapes") + return _ExpertShape(up_shape[1], up_shape[0], False) + + +def _stacked_expert_shape( + first: Any, down: Any, factor: int, name: str, expected_hidden: int | None +) -> _ExpertShape: + if first.ndim != 3 or down.ndim != 3 or first.shape[0] != down.shape[0]: + raise ShapeError(f"unsupported stacked experts at {name}") + + first_shape = tuple(int(value) for value in first.shape) + down_shape = tuple(int(value) for value in down.shape) + candidates = [] + if first_shape[1] % factor == 0: + hidden, intermediate = first_shape[2], first_shape[1] // factor + if down_shape[1:] == (hidden, intermediate): + candidates.append((hidden, intermediate)) + if first_shape[2] % factor == 0: + hidden, intermediate = first_shape[1], first_shape[2] // factor + if down_shape[1:] == (intermediate, hidden): + candidates.append((hidden, intermediate)) + if expected_hidden is not None: + candidates = [candidate for candidate in candidates if candidate[0] == expected_hidden] + if len(set(candidates)) != 1: + raise ShapeError(f"unsupported stacked expert projection shapes at {name}") + hidden, intermediate = candidates[0] + return _ExpertShape(hidden, intermediate, factor == 2) + + +def _moe_activation(config: Any, gated: bool) -> str | None: + configured = getattr(config, "mlp_hidden_act", None) or getattr(config, "hidden_act", None) + normalized = str(configured).lower().replace("-", "_") + if gated: + if normalized in {"silu", "swiglu", "swish"}: + return "Swiglu" + # vLLM's FlashInfer MoE path serves only the tanh-approximation GELU + # ("gelu_tanh", "gelu_pytorch_tanh", and "gelu_new" share that + # formula). Exact gelu and quick_gelu take non-FlashInfer backends in + # vLLM, so they are rejected here rather than timed as a proxy. + if normalized in {"gelu_tanh", "gelu_pytorch_tanh", "gelu_new"}: + return "Geglu" + raise ShapeError( + f"unsupported gated MoE activation {configured!r}; vLLM's FlashInfer MoE " + "serves only SiLU/SwiGLU and tanh-GELU" + ) + activations = { + "gelu": "Gelu", + "identity": "Identity", + "relu": "Relu", + "relu2": "Relu2", + "relu_squared": "Relu2", + "silu": "Silu", + } + if normalized not in activations: + raise ShapeError(f"unsupported non-gated MoE activation {configured!r}") + return activations[normalized] + + +# Fallback copy of ModelOpt's _ACTIVE_MOE_TOP_K_ATTRS for environments without +# ModelOpt installed; a test asserts it stays in sync with the canonical list. +_MOE_TOP_K_ATTRS_FALLBACK = ( + "num_experts_per_tok", + "num_experts_per_token", + "moe_top_k", + "top_k", + "num_selected_experts", +) + + +def _top_k(config: Any) -> int | None: + # ModelOpt's AutoQuantize cost model owns the canonical attribute list, so + # benchmark rows and AutoQuantize agree on how a config declares top_k. + try: + from modelopt.torch.quantization._auto_quantize_cost import _ACTIVE_MOE_TOP_K_ATTRS + + attrs = _ACTIVE_MOE_TOP_K_ATTRS + except ImportError: + attrs = _MOE_TOP_K_ATTRS_FALLBACK + + for attr in attrs: + value = getattr(config, attr, None) + if value is not None: + return int(value) + return None + + +def _moe_shapes(model: Any, config: Any) -> set[_MoeShape]: + shapes = set() + top_k = _top_k(config) + for name, module in model.named_modules(): + expert_container = name.rsplit(".", 1)[-1] in {"experts", "local_experts"} + if expert_container: + expert_modules = list(module.children()) + shape = _expert_shape(expert_modules[0]) if expert_modules else None + if shape: + if top_k is None: + raise ShapeError("could not determine MoE top_k") + if any(_expert_shape(expert) != shape for expert in expert_modules[1:]): + raise ShapeError(f"experts under {name} do not share one Linear layout") + shapes.add( + _MoeShape( + shape.hidden, + shape.intermediate, + len(expert_modules), + top_k, + _moe_activation(config, shape.gated), + ) + ) + + params = dict(module.named_parameters(recurse=False)) + down = params.get("down_proj") + if params.get("gate_up_proj") is not None: + first, factor = params["gate_up_proj"], 2 + elif params.get("up_proj") is not None: + first, factor = params["up_proj"], 1 + else: + expert_params = [ + f"{param_name}{tuple(param.shape)}" + for param_name, param in params.items() + if param.ndim >= 2 + ] + if expert_container and expert_params: + raise ShapeError( + f"unsupported stacked expert parameters at {name}: " + ", ".join(expert_params) + ) + continue + if down is None: + raise ShapeError(f"stacked experts at {name} have no down projection") + if top_k is None: + raise ShapeError("could not determine MoE top_k") + expected_hidden = getattr(config, "moe_latent_size", None) or getattr( + config, "hidden_size", None + ) + shape = _stacked_expert_shape( + first, + down, + factor, + name, + int(expected_hidden) if expected_hidden is not None else None, + ) + shapes.add( + _MoeShape( + shape.hidden, + shape.intermediate, + int(first.shape[0]), + top_k, + _moe_activation(config, shape.gated), + ) + ) + return shapes + + +def _declared_expert_count(config: Any) -> int | None: + if _top_k(config) is None: + return None + for attr in ("n_routed_experts", "num_local_experts", "num_experts"): + value = getattr(config, attr, None) + if value is not None and int(value) > 0: + return int(value) + return None + + +def _unsupported_decoder_linears( + model: Any, routed_experts_handled: bool = False +) -> list[tuple[str, int, int]]: + mamba_projections = set() + for parent, module in model.named_modules(): + if _mamba_layout(module) is not None: + mamba_projections.update({f"{parent}.in_proj", f"{parent}.out_proj"}) + if _gdn_layout(module) is not None: + mamba_projections.update( + f"{parent}.{leaf}" + for leaf in ( + "in_proj_qkvz", + "in_proj_ba", + "in_proj_qkv", + "in_proj_z", + "in_proj_b", + "in_proj_a", + "out_proj", + ) + ) + + layouts: dict[tuple[str, int, int], str] = {} + for name, module in model.named_modules(): + shape = _linear_shape(module) + if shape is None: + continue + parts = name.split(".") + in_decoder = any( + part in {"block", "blocks", "h", "layer", "layers"} + and index + 1 < len(parts) + and parts[index + 1].isdigit() + for index, part in enumerate(parts) + ) + leaf = parts[-1] + if not in_decoder: + continue + if any(part in _ROUTER_PATH_PARTS for part in parts): + continue + if routed_experts_handled and any(part in {"experts", "local_experts"} for part in parts): + continue + if leaf in _PROJECTIONS or leaf in _GATING_LEAF_NAMES or name in mamba_projections: + continue + layouts.setdefault((leaf, *shape), name) + return [(name, n, k) for (leaf, n, k), name in layouts.items()] + + +def _inspect_model( + model: Any, config: Any, tp: int, ep: int +) -> tuple[list[_Kernel], _MoeShape | None, list[str]]: + config = getattr(config, "text_config", None) or config + kernels = ( + _dense_kernels(model, config, tp) + _mamba_kernels(model, tp) + _gdn_kernels(model, tp) + ) + problems = [] + try: + moe_shapes = _moe_shapes(model, config) + except ShapeError as exc: + moe_shapes = set() + problems.append(str(exc)) + declared_experts = _declared_expert_count(config) + if not problems and declared_experts is not None: + if not moe_shapes: + problems.append( + f"model declares {declared_experts} routed experts but no supported expert " + "GEMM layout was found" + ) + elif any(shape.experts != declared_experts for shape in moe_shapes): + found = sorted({shape.experts for shape in moe_shapes}) + problems.append( + f"model declares {declared_experts} routed experts but instantiated layouts " + f"have expert counts {found}" + ) + experts_recognized = bool(moe_shapes) + if len(moe_shapes) > 1: + problems.append("model contains multiple routed-expert layouts") + if problems: + # The expert audit is unresolved, so a derived MoE tuple is suspect; + # skip its per-rank validation so the audit findings are reported + # instead of a masking ShapeError. + moe_shapes = set() + moe = next(iter(moe_shapes), None) + if moe is None: + if ep != 1 and not problems: + raise ShapeError("EP requires routed experts") + else: + if ep != 1 and ep % tp: + raise ShapeError( + f"EP={ep} is not a multiple of TP={tp}; vLLM expert parallelism spans TP x DP, " + "so no modeled serving layout matches this combination — if it is intentional " + "(e.g. Megatron-style EP), benchmark the per-rank expert shape directly with " + "benchmark_via_builtin.py" + ) + local_experts = _divide(moe.experts, ep, "expert count") + intermediate = moe.intermediate + if ep == 1: + intermediate = _divide(intermediate, tp, "expert intermediate size") + if moe.top_k > local_experts: + raise ShapeError("top_k exceeds the per-rank expert count") + moe = _MoeShape(moe.hidden, intermediate, local_experts, moe.top_k, moe.activation) + unsupported = _unsupported_decoder_linears(model, routed_experts_handled=experts_recognized) + if unsupported: + details = ", ".join(f"{name} ({n}x{k})" for name, n, k in unsupported) + problems.append(f"unsupported decoder Linear GEMM layout(s): {details}") + if not kernels and moe is None and not problems: + raise ShapeError("no dense benchmark shapes found") + return kernels, moe, problems + + +def _command( + kernels: list[_Kernel], + moe: _MoeShape | None, + passthrough: list[str], +) -> list[str]: + command: list[str] = [] + if kernels: + # One pair per derived kernel: same-shape kernels from different + # modules keep separate names and become duplicated result rows. + command += ["--nks", *(f"{n},{k}" for n, k, _ in kernels)] + command += ["--nk_names", *(label for _, _, label in kernels)] + if moe: + command += [ + "--moe_hidden_size", + str(moe.hidden), + "--moe_intermediate_size", + str(moe.intermediate), + "--moe_num_experts", + str(moe.experts), + "--moe_top_k", + str(moe.top_k), + ] + if moe.activation: + command += ["--moe_activation_type", moe.activation] + return command + passthrough + + +def _load_runner() -> Any: + path = Path(__file__).with_name("benchmark_via_builtin.py") + spec = importlib.util.spec_from_file_location("benchmark_via_builtin", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def main() -> None: + """Parse arguments, derive shapes, and optionally run the benchmark.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model", help="Hub ID, local model directory, or config.json") + parser.add_argument("--tp", type=_positive_int, default=1, help="tensor parallel size, e.g. 8") + parser.add_argument("--ep", type=_positive_int, default=1, help="expert parallel size, e.g. 8") + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--revision", help="Hugging Face branch, tag, or commit") + parser.add_argument("--print_only", action="store_true") + args, passthrough = parser.parse_known_args() + for token in passthrough: + if token.split("=", 1)[0] in _RESERVED: + parser.error("derived --nks/--nk_names/--moe_* shapes cannot be overridden") + + try: + config, model = _load_meta_model(args.model, args.trust_remote_code, args.revision) + kernels, moe, problems = _inspect_model(model, config, args.tp, args.ep) + except ShapeError as exc: + parser.error(str(exc)) + + print( + f"# {type(model).__name__} ({getattr(config, 'model_type', '?')}), " + f"TP={args.tp}, EP={args.ep}" + ) + print( + "# layout: Transformers meta model; fused QKV and gate/up; " + "Mamba 2, GatedDeltaNet, and routed experts" + ) + for n, k, label in dict.fromkeys(kernels): + print(f"# {n}x{k} <- {label}") + if moe: + activation = f" activation={moe.activation}" if moe.activation else "" + print( + f"# MoE: H={moe.hidden} F={moe.intermediate} E={moe.experts} " + f"top_k={moe.top_k}{activation}" + ) + if args.ep > 1: + print( + f"# MoE sharding: EP={args.ep} partitions whole experts; " + "expert width stays intact (expert-TP=1)" + ) + elif args.tp > 1: + print(f"# MoE sharding: TP={args.tp} shards the expert intermediate width (EP=1)") + for problem in problems: + print(f"# unsupported: {problem}") + if problems: + parser.error( + "the derived shapes above are incomplete; validate each unsupported layout's " + "TP/EP sharding and benchmark it directly with benchmark_via_builtin.py" + ) + command = _command(kernels, moe, passthrough) + runner = Path(__file__).with_name("benchmark_via_builtin.py") + print(">>> " + shlex.join([sys.executable, str(runner), *command])) + if not args.print_only: + _load_runner().main(command) + + +if __name__ == "__main__": + main() diff --git a/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py new file mode 100644 index 00000000000..b2eed9daef0 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py @@ -0,0 +1,863 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run FlashInfer's built-in GEMM and fused-MoE microbenchmarks. + +Plain rows contain kernel time. Most ``*_with_quant`` rows add a separately +measured activation-quantization time in the scale-factor layout the backend +consumes; the NVFP4 CUTLASS MoE row is instead a single fused measurement. +Logical shapes label each case while backend-specific physical padding follows +vLLM. A local FlashInfer source checkout is required for its benchmark driver +and utilities. +""" + +from __future__ import annotations + +import argparse +import csv +import os +import shlex +import subprocess # nosec B404 +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import TextIO + + import torch + +try: + from vllm import _custom_ops as vllm_ops +except ImportError: + vllm_ops = None + +_ERROR_CASE_PREFIX = "[ERROR] Error running test:" +_ERROR_MESSAGE_PREFIX = "[ERROR] Error:" +_FP8_QUANT_UNAVAILABLE = "ERROR: vLLM is unavailable for FP8 activation quantization" +_MOE_ACTIVATIONS = ( + "Gelu", + "Relu", + "Silu", + "Swiglu", + "Geglu", + "SwigluBias", + "Relu2", + "SwigluStep", + "Identity", +) +_ResultValue = float | str + + +@dataclass(frozen=True, order=True) +class _QuantSpec: + """One shared activation-quantization measurement of an m-by-k BF16 tile. + + Attributes: + dtype: Quantized element type, ``"nvfp4"`` or ``"fp8"``. + layout: Scale-factor layout the consuming kernel expects: ``"128x4"``, + ``"8x4"``, or ``"linear"`` for NVFP4; ``"static"`` for FP8. + m: Token count of the activation tile. + k: Inner dimension of the activation tile. + """ + + dtype: str + layout: str + m: int + k: int + + +@dataclass +class _Case: + """One driver invocation and, once it has run, its outcome. + + Attributes: + section: ``"gemm"`` or ``"moe"``. + tag: Case label passed to the driver as ``--case_tag``; validates + returned rows and labels the artifacts. Never an internal key. + backend: Value of the output ``backend`` column. + m: Token count. + n: Logical output size; the driver may run a padded physical shape + from ``argv``. ``None`` for MoE cases. + k: Logical reduction size, as ``n``. + argv: Driver arguments, before ``--case_tag``/``--output_path``. + with_quant: ``with_quant`` column of the measured row itself; ``True`` + only for the fused NVFP4 CUTLASS MoE measurement. + quant: Spec of the activation-quantization time a derived + ``with_quant`` row adds to ``result``. + result: Median kernel time in microseconds, an ``ERROR: ...`` message, + or ``None`` until the case has run. + quant_result: Measured time (or error) of ``quant``, recorded by + ``_attach_quant_times``. + """ + + section: str + tag: str + backend: str + m: int + n: int | None + k: int | None + argv: list[str] + with_quant: bool = False + quant: _QuantSpec | None = None + result: _ResultValue | None = None + quant_result: _ResultValue | None = None + + +@dataclass(frozen=True) +class _MoeShape: + """The per-rank fused-MoE problem all MoE cases share. + + Attributes: + hidden: Model hidden size. + intermediate: Per-rank expert width. + experts: Per-rank expert count. + top_k: Experts activated per token. + activation: FlashInfer activation name; ``None`` means the driver + default (gated SwiGLU). + """ + + hidden: int + intermediate: int + experts: int + top_k: int + activation: str | None = None + + def label(self) -> str: + label = f"H={self.hidden} F={self.intermediate} E={self.experts} top_k={self.top_k}" + if self.activation: + label += f" activation={self.activation}" + return label + + +def _positive_int(value: str) -> int: + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError(f"{value!r} is not an integer") from exc + if parsed <= 0: + raise argparse.ArgumentTypeError(f"{value!r} is not a positive integer") + return parsed + + +def _nk_pair(value: str) -> tuple[int, int]: + try: + n, k = value.split(",") + except ValueError as exc: + raise argparse.ArgumentTypeError(f"expected a positive N,K pair, got {value!r}") from exc + try: + return _positive_int(n), _positive_int(k) + except argparse.ArgumentTypeError as exc: + raise argparse.ArgumentTypeError(f"expected a positive N,K pair, got {value!r}") from exc + + +def _named_nks( + nks: list[tuple[int, int]], names: list[str] | None +) -> dict[tuple[int, int], list[str]]: + """Map each unique N,K pair, in first-seen order, to its module labels. + + Unnamed shapes label themselves ``"NxK"``. + """ + if names is not None and len(names) != len(nks): + raise ValueError("--nk_names must contain exactly one name for each --nks pair") + if names is None: + return {(n, k): [f"{n}x{k}"] for n, k in dict.fromkeys(nks)} + + labels_by_nk: dict[tuple[int, int], list[str]] = {} + for nk, name in zip(nks, names, strict=True): + labels = labels_by_nk.setdefault(nk, []) + if name not in labels: + labels.append(name) + return labels_by_nk + + +def _round_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _parse_driver_error(lines: list[str]) -> str | None: + """Extract the driver's error message from one case's output. + + Returns the message, ``""`` when the driver reported an error without a + message, or ``None`` when no error was reported. Each case runs in its own + driver process, so any reported error belongs to that case. + """ + error = None + pending = False + for line in lines: + stripped = line.strip() + if stripped.startswith(_ERROR_CASE_PREFIX): + pending = True + elif stripped.startswith(_ERROR_MESSAGE_PREFIX) and pending: + error = stripped.removeprefix(_ERROR_MESSAGE_PREFIX).strip().replace(",", ";") + pending = False + return error + + +def _failure_message(case_output: list[str], returncode: int, driver_log: Path) -> str: + """Classify a case that produced no result row into an ``ERROR: ...`` cell.""" + message = _parse_driver_error(case_output) + if message: + return f"ERROR: {message}" + if message is not None: + return ( + "ERROR: FlashInfer reported an error without a message (empty exception); " + f"see {driver_log}" + ) + if returncode: + return ( + f"ERROR: FlashInfer driver exited with status {returncode} for this case; " + f"see {driver_log}" + ) + return f"ERROR: FlashInfer produced no result row and no error message; see {driver_log}" + + +def _run_case(benchmarks_dir: Path, argv: list[str], log: TextIO) -> tuple[int, list[str]]: + # Each case gets its own driver process: a fatal CUDA fault (for example a + # misaligned address) permanently poisons the CUDA context, so sharing one + # process would fail every later case (verified empirically). This invokes + # the explicitly selected FlashInfer checkout without a shell. + process = subprocess.Popen( # nosec B603 + [sys.executable, "flashinfer_benchmark.py", *argv], + cwd=benchmarks_dir, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert process.stdout is not None + lines = [] + for line in process.stdout: + print(line, end="", flush=True) + log.write(line) + lines.append(line) + return process.wait(), lines + + +def _gpu_description() -> str: + try: + import torch + + # The driver name can be a placeholder on pre-release GPUs, so record + # compute capability, SM count, and memory to pin down the exact part. + properties = torch.cuda.get_device_properties(0) + name = ( + f"{properties.name} (sm_{properties.major}{properties.minor} / " + f"{properties.multi_processor_count} SMs / " + f"{properties.total_memory / (1 << 30):.0f} GiB)" + ) + except Exception: + return "unknown GPU" + watts = "unknown power limit" + try: + import pynvml + + pynvml.nvmlInit() + try: + # NVML does not honor CUDA_VISIBLE_DEVICES, so map the first + # visible device back to its physical NVML handle. + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")[0].strip() + if visible.startswith(("GPU-", "MIG-")): + handle = pynvml.nvmlDeviceGetHandleByUUID(visible) + else: + handle = pynvml.nvmlDeviceGetHandleByIndex(int(visible) if visible else 0) + limit = pynvml.nvmlDeviceGetPowerManagementLimit(handle) + watts = f"{limit / 1000:.0f} W power limit" + finally: + pynvml.nvmlShutdown() + except Exception: + pass + return f"{name}; {watts}" + + +def _environment_header(flashinfer_repo: Path) -> str: + try: + import flashinfer + + version = flashinfer.__version__ + except Exception: + version = "unknown" + try: + # Reads the revision of the explicitly selected checkout, no shell. + result = subprocess.run( # nosec B603 B607 + ["git", "-C", str(flashinfer_repo), "rev-parse", "HEAD"], + capture_output=True, + text=True, + timeout=10, + check=False, + ) + revision = result.stdout.strip() or "unknown" + except OSError: + revision = "unknown" + return ( + f"flashinfer {version}; checkout {flashinfer_repo.resolve()} @ {revision}; " + f"{_gpu_description()}" + ) + + +def _write_builtin(path: Path, rows: list[dict[str, str]]) -> None: + fieldnames: dict[str, None] = {} + for row in rows: + for key in row: + fieldnames.setdefault(key, None) + with path.open("w", newline="") as stream: + writer = csv.DictWriter( + stream, fieldnames=list(fieldnames), restval="", lineterminator="\n" + ) + writer.writeheader() + writer.writerows(rows) + + +def _gemm_cases( + ms: list[int], + nks: list[tuple[int, int]], + common: list[str], +) -> list[_Case]: + cases: list[_Case] = [] + + def add( + m: int, + n: int, + k: int, + backend: str, + routine: str, + driver_backend: str, + run_n: int | None = None, + run_k: int | None = None, + extra: list[str] | None = None, + quant: _QuantSpec | None = None, + ) -> None: + cases.append( + _Case( + section="gemm", + tag=f"gemm_{backend}_MxNxK={m}x{n}x{k}", + backend=backend, + m=m, + n=n, + k=k, + argv=[ + "--routine", + routine, + "--backends", + driver_backend, + *(extra or []), + "--m", + str(m), + "--n", + str(run_n if run_n is not None else n), + "--k", + str(run_k if run_k is not None else k), + *common, + ], + quant=quant, + ) + ) + + for m in ms: + for n, k in nks: + # Physical padding follows vLLM: dense NVFP4 on cuDNN, CUTLASS, + # and CuteDSL pads N and K to multiples of 32; trtllm keeps the + # exact shape with the shuffled layout; BF16 and FP8 stay exact. + add(m, n, k, "bf16", "mm_bf16", "cudnn") + for row_suffix, driver_backend in ( + ("cudnn", "cudnn"), + ("cutlass", "cutlass"), + ("cutedsl", "cute-dsl"), + ("trtllm", "trtllm"), + ): + layout = "128x4" if driver_backend != "trtllm" or m > 32 else "8x4" + extra = ["--use_nvfp4"] + if layout == "128x4": + extra.append("--use_128x4_sf_layout") + run_n, run_k = n, k + if driver_backend != "trtllm": + run_n, run_k = _round_up(n, 32), _round_up(k, 32) + add( + m, + n, + k, + f"nvfp4_{row_suffix}", + "mm_fp4", + driver_backend, + run_n, + run_k, + extra, + _QuantSpec("nvfp4", layout, m, run_k), + ) + for driver_backend in ("cudnn", "cutlass"): + add( + m, + n, + k, + f"fp8_{driver_backend}", + "bmm_fp8", + driver_backend, + extra=["--batch_size", "1"], + quant=_QuantSpec("fp8", "static", m, k), + ) + add( + m, + n, + k, + "fp8_trtllm", + "mm_fp8", + "trtllm_low_latency", + quant=_QuantSpec("fp8", "static", m, k), + ) + return cases + + +def _moe_cases(ms: list[int], shape: _MoeShape, common: list[str]) -> list[_Case]: + cases: list[_Case] = [] + + def add( + backend: str, + routine: str, + intermediate: int, + hidden: int = shape.hidden, + extra: list[str] | None = None, + quant: tuple[str, str] | None = None, + with_quant: bool = False, + ) -> None: + suffix = "_with_quant" if with_quant else "" + cases.extend( + _Case( + section="moe", + tag=f"moe_{backend}_moe{suffix}_M={m}", + backend=backend, + m=m, + n=None, + k=None, + with_quant=with_quant, + quant=_QuantSpec(*quant, m, hidden) if quant else None, + argv=[ + "--routine", + routine, + "--num_tokens", + str(m), + "--hidden_size", + str(hidden), + "--num_experts", + str(shape.experts), + "--top_k", + str(shape.top_k), + *(["--activation-type", shape.activation] if shape.activation else []), + "--intermediate_size", + str(intermediate), + *(extra or []), + *common, + ], + ) + for m in ms + ) + + gated = shape.activation is None or shape.activation in { + "Swiglu", + "Geglu", + "SwigluBias", + "SwigluStep", + } + # Pad a dimension only when vLLM pads it. FP8 per-tensor (CUTLASS and + # trtllm-gen) pads the intermediate to 16 gated / 128 non-gated. NVFP4 + # CUTLASS pads non-gated intermediate up to the 128-aligned swizzled scale + # rows but raises instead of padding gated, so gated stays exact and may + # fail like vLLM. NVFP4 trtllm-gen additionally pads hidden to 256. + fp8_intermediate = _round_up(shape.intermediate, 16 if gated else 128) + nvfp4_intermediate = shape.intermediate if gated else _round_up(shape.intermediate, 128) + + add("bf16_cutlass", "cutlass_fused_moe", shape.intermediate) + add( + "fp8_cutlass", + "cutlass_fused_moe", + fp8_intermediate, + extra=["--cutlass_variant", "fp8"], + quant=("fp8", "static"), + ) + add( + "nvfp4_cutlass", + "cutlass_fused_moe", + nvfp4_intermediate, + extra=["--cutlass_variant", "nvfp4", "--quantized_input"], + ) + # The NVFP4 CUTLASS with_quant row is its own fused driver measurement + # (unquantized input), not a derived base-plus-quant-time row. + add( + "nvfp4_cutlass", + "cutlass_fused_moe", + nvfp4_intermediate, + extra=["--cutlass_variant", "nvfp4"], + with_quant=True, + ) + # Routing is synthetic in this benchmark (uniform random logits), so the + # trtllm-gen rows, which route in-kernel, use a fixed renormalize method + # to stay comparable across models; the model's real routing scheme is + # not derivable from its config alone. CUTLASS and CuteDSL rows receive + # precomputed indices and have no routing stage to time. + add( + "fp8_trtllm", + "trtllm_fp8_per_tensor_scale_moe", + fp8_intermediate, + extra=["--routing_method", "renormalize"], + quant=("fp8", "static"), + ) + add( + "nvfp4_trtllm", + "trtllm_fp4_block_scale_moe", + fp8_intermediate, + hidden=_round_up(shape.hidden, 256), + extra=["--routing_method", "renormalize"], + quant=("nvfp4", "linear"), + ) + if shape.activation in (None, "Swiglu"): + # FlashInfer's CuteDSL fused MoE supports only gated Swiglu. + add( + "nvfp4_cutedsl", + "cute_dsl_fp4_block_scale_moe", + shape.intermediate, + quant=("nvfp4", "linear"), + ) + return cases + + +def _nvfp4_runner(tensor: torch.Tensor, layout: str): + import flashinfer + + global_scale = (448 * 6) / tensor.float().abs().nan_to_num().max() + if layout == "linear": + # The trtllm-gen and CuteDSL fused-MoE kernels consume activation + # scale factors in linear (unswizzled) layout. + def linear_kernel(value, scale): + return flashinfer.fp4_quantize(value, scale, is_sf_swizzled_layout=False) + + return linear_kernel, (tensor, global_scale) + sf_layout = ( + flashinfer.SfLayout.layout_128x4 if layout == "128x4" else flashinfer.SfLayout.layout_8x4 + ) + + def kernel(value, scale): + return flashinfer.nvfp4_quantize(value, scale, sfLayout=sf_layout, do_shuffle=False) + + return kernel, (tensor, global_scale) + + +def _fp8_runner(tensor: torch.Tensor): + import torch + + scale = tensor.abs().max().float() / torch.finfo(torch.float8_e4m3fn).max + + def kernel(value, value_scale): + quantized, _ = vllm_ops.scaled_fp8_quant(value.contiguous(), value_scale) + return quantized + + return kernel, (tensor, scale) + + +def _attach_quant_times( + cases: list[_Case], dry_runs: int, iterations: int, cuda_graph: bool +) -> None: + """Measure each distinct quantization spec once and attach shared times.""" + specs = sorted( + {case.quant for case in cases if case.quant is not None and isinstance(case.result, float)} + ) + results: dict[_QuantSpec, _ResultValue] = {} + if vllm_ops is None and any(spec.dtype == "fp8" for spec in specs): + print(f"[WARN] {_FP8_QUANT_UNAVAILABLE.removeprefix('ERROR: ')}") + results = {spec: _FP8_QUANT_UNAVAILABLE for spec in specs if spec.dtype == "fp8"} + specs = [spec for spec in specs if spec.dtype != "fp8"] + if specs: + # The GPU stack is imported lazily so shape planning, result parsing, + # and their tests work without FlashInfer or torch installed. + import numpy as np + import torch + from flashinfer.testing import bench_gpu_time + + for spec in specs: + tensor = torch.randn(spec.m, spec.k, device="cuda", dtype=torch.bfloat16) + kernel, inputs = ( + _nvfp4_runner(tensor, spec.layout) if spec.dtype == "nvfp4" else _fp8_runner(tensor) + ) + times = bench_gpu_time( + fn=kernel, + input_args=inputs, + dry_run_iters=dry_runs, + repeat_iters=iterations, + enable_cupti=True, + use_cuda_graph=cuda_graph, + cold_l2_cache=True, + sleep_after_run=True, + ) + results[spec] = float(np.median(times)) * 1000 + for case in cases: + if case.quant is not None and isinstance(case.result, float): + case.quant_result = results[case.quant] + + +def _format_result(value: _ResultValue) -> str: + if isinstance(value, float): + return f"{value:.3f}" + return value + + +def _output_rows(case: _Case) -> list[tuple[bool, _ResultValue]]: + """Expand a case into its (with_quant, runtime) output rows. + + A case with a quant spec gets a derived ``with_quant`` row adding the + shared activation-quantization time; an error on either measurement + propagates into the derived row. + """ + assert case.result is not None + rows = [(case.with_quant, case.result)] + if case.quant is None: + return rows + if isinstance(case.result, str): + rows.append((True, case.result)) + return rows + quant_result = case.quant_result + assert quant_result is not None + rows.append( + (True, case.result + quant_result if isinstance(quant_result, float) else quant_result) + ) + return rows + + +def _write_results( + path: Path, + cases: list[_Case], + labels_by_nk: dict[tuple[int, int], list[str]], + header: str | None = None, + moe_label: str | None = None, +) -> None: + columns = ["module_name", "M", "N", "K", "backend", "with_quant", "runtime"] + gemm = [case for case in cases if case.section == "gemm" and case.result is not None] + moe = [case for case in cases if case.section == "moe" and case.result is not None] + with path.open("w", newline="") as stream: + writer = csv.writer(stream, lineterminator="\n") + if header: + writer.writerow([header]) + if gemm: + writer.writerow(["GEMM"]) + writer.writerow(columns) + for (n, k), labels in labels_by_nk.items(): + group = sorted( + (case for case in gemm if (case.n, case.k) == (n, k)), + key=lambda case: (case.backend, case.m), + ) + # Modules fused into one GEMM are joined with "|" inside one + # name; distinct same-shape modules each get their own rows, + # duplicating the shared measurement. + for label in labels: + for case in group: + for with_quant, value in _output_rows(case): + writer.writerow( + [ + label, + case.m, + n, + k, + case.backend, + with_quant, + _format_result(value), + ] + ) + if moe: + if gemm: + writer.writerow([]) + writer.writerow(["MoE"]) + if moe_label: + writer.writerow([moe_label]) + writer.writerow(columns) + for case in sorted(moe, key=lambda case: (case.backend, case.m, case.with_quant)): + for with_quant, value in _output_rows(case): + writer.writerow( + [ + "experts", + case.m, + "", + "", + case.backend, + with_quant, + _format_result(value), + ] + ) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--flashinfer_repo", + type=Path, + required=True, + help="checkout containing benchmarks/flashinfer_benchmark.py", + ) + parser.add_argument( + "--ms", + type=_positive_int, + nargs="+", + default=[1, 8, 64, 512], + help="token counts, for example: 1 8 64 512", + ) + parser.add_argument("--nks", type=_nk_pair, nargs="+", help="GEMM N,K pairs, e.g. 4096,4096") + parser.add_argument( + "--nk_names", + nargs="+", + help="optional names parallel to --nks, e.g. qkv_proj o_proj", + ) + parser.add_argument("--dry_run_iters", type=_positive_int, help="warmup iterations, e.g. 5") + parser.add_argument("--num_iters", type=_positive_int, help="timed iterations, e.g. 30") + parser.add_argument("--no_cuda_graph", action="store_true") + parser.add_argument("--no_autotune", action="store_true") + parser.add_argument( + "--moe_hidden_size", type=_positive_int, help="model hidden size, e.g. 4096" + ) + parser.add_argument( + "--moe_intermediate_size", type=_positive_int, help="expert width, e.g. 14336" + ) + parser.add_argument("--moe_num_experts", type=_positive_int, help="local expert count, e.g. 8") + parser.add_argument("--moe_top_k", type=_positive_int, help="experts per token, e.g. 2") + parser.add_argument( + "--moe_activation_type", + choices=_MOE_ACTIVATIONS, + help="FlashInfer activation, e.g. Swiglu", + ) + parser.add_argument("--workdir", type=Path, default=Path("benchmark_via_builtin_out")) + return parser + + +def _execute_cases( + cases: list[_Case], benchmarks_dir: Path, workdir: Path, driver_log: Path, header: str +) -> list[dict[str, str]]: + """Run each case in its own driver process and record outcomes on it. + + Returns the raw driver CSV rows for ``builtin_results.csv``. + """ + case_csv = workdir / "case_result.csv" + rows: list[dict[str, str]] = [] + with driver_log.open("w") as log: + print(header, flush=True) + log.write(header + "\n") + for case in cases: + marker = f"[CASE] {case.tag}\n" + print(marker, end="", flush=True) + log.write(marker) + case_csv.unlink(missing_ok=True) + returncode, case_output = _run_case( + benchmarks_dir, + [*case.argv, "--case_tag", case.tag, "--output_path", str(case_csv.resolve())], + log, + ) + case_rows = [] + if case_csv.is_file(): + with case_csv.open(newline="") as stream: + # A row that does not carry this case's tag cannot be + # trusted as this case's measurement; fail the case. + case_rows = [ + row for row in csv.DictReader(stream) if row.get("case_tag") == case.tag + ] + if case_rows: + rows.extend(case_rows) + case.result = float(case_rows[-1]["median_time"]) * 1000 + else: + case.result = _failure_message(case_output, returncode, driver_log) + case_csv.unlink(missing_ok=True) + return rows + + +def main(argv: list[str] | None = None) -> None: + """Validate inputs, run the FlashInfer driver, and combine its results.""" + parser = _parser() + args = parser.parse_args(argv) + ms = list(dict.fromkeys(args.ms)) + try: + labels_by_nk = _named_nks(args.nks or [], args.nk_names) + except ValueError as exc: + parser.error(str(exc)) + moe_values = ( + args.moe_hidden_size, + args.moe_intermediate_size, + args.moe_num_experts, + args.moe_top_k, + ) + if any(moe_values) and not all(moe_values): + parser.error("all four --moe_* shape arguments are required together") + moe_shape = None + if all(moe_values): + moe_shape = _MoeShape( + args.moe_hidden_size, + args.moe_intermediate_size, + args.moe_num_experts, + args.moe_top_k, + args.moe_activation_type, + ) + if not labels_by_nk and moe_shape is None: + parser.error("pass --nks and/or all four --moe_* shape arguments") + if moe_shape is not None and moe_shape.top_k > moe_shape.experts: + parser.error("--moe_top_k cannot exceed --moe_num_experts") + + benchmarks_dir = args.flashinfer_repo / "benchmarks" + driver = benchmarks_dir / "flashinfer_benchmark.py" + if not driver.is_file(): + parser.error(f"{driver} does not exist") + + common = [] + if args.dry_run_iters is not None: + common += ["--dry_run_iters", str(args.dry_run_iters)] + if args.num_iters is not None: + common += ["--num_iters", str(args.num_iters)] + if not args.no_autotune: + common.append("--autotune") + if args.no_cuda_graph: + common.append("--no_cuda_graph") + + cases = _gemm_cases(ms, list(labels_by_nk), common) + if moe_shape is not None: + cases += _moe_cases(ms, moe_shape, common) + + args.workdir.mkdir(parents=True, exist_ok=True) + testlist = args.workdir / "testlist.txt" + builtin_csv = args.workdir / "builtin_results.csv" + combined_csv = args.workdir / "combined_results.csv" + driver_log = args.workdir / "driver.log" + if builtin_csv.exists() or combined_csv.exists(): + parser.error(f"{args.workdir} already contains results; choose a fresh --workdir") + testlist.write_text( + "\n".join(shlex.join([*case.argv, "--case_tag", case.tag]) for case in cases) + "\n" + ) + + header = _environment_header(args.flashinfer_repo) + rows = _execute_cases(cases, benchmarks_dir, args.workdir, driver_log, header) + if rows: + _write_builtin(builtin_csv, rows) + + _attach_quant_times( + cases, + args.dry_run_iters if args.dry_run_iters is not None else 5, + args.num_iters if args.num_iters is not None else 30, + not args.no_cuda_graph, + ) + moe_label = moe_shape.label() if moe_shape is not None else None + _write_results(combined_csv, cases, labels_by_nk, header, moe_label) + print(f"Wrote {combined_csv}") + failed = [case.tag for case in cases if isinstance(case.result, str)] + if failed: + raise RuntimeError( + "FlashInfer failed benchmark cases: " + + ", ".join(failed) + + f"; wrote failure details to {combined_csv}" + ) + + +if __name__ == "__main__": + main() diff --git a/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py new file mode 100644 index 00000000000..02c7f80b159 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py @@ -0,0 +1,622 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("accelerate") +transformers = pytest.importorskip("transformers") + +from torch import nn + +SCRIPT = Path(__file__).parents[1] / "scripts" / "benchmark_model.py" +SPEC = importlib.util.spec_from_file_location("benchmark_model", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark_model = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = benchmark_model +SPEC.loader.exec_module(benchmark_model) + + +def _save(tmp_path, config): + config.save_pretrained(tmp_path) + return tmp_path + + +def _preview(model_ref, monkeypatch, capsys, *, tp=1, ep=1): + monkeypatch.setattr( + sys, + "argv", + [str(SCRIPT), str(model_ref), "--tp", str(tp), "--ep", str(ep), "--print_only"], + ) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + benchmark_model.main() + return capsys.readouterr().out + + +def _llama_config(): + return transformers.LlamaConfig( + vocab_size=128, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=64, + ) + + +def _nemotron_h_config(*, n_groups=2): + return transformers.NemotronHConfig( + vocab_size=128, + hidden_size=32, + layers_block_type=["mamba", "moe"], + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + intermediate_size=40, + use_mamba_kernels=False, + ssm_state_size=4, + mamba_num_heads=4, + mamba_head_dim=8, + n_groups=n_groups, + conv_kernel=4, + expand=1, + n_routed_experts=4, + n_shared_experts=1, + moe_intermediate_size=48, + moe_shared_expert_intermediate_size=40, + num_experts_per_tok=2, + ) + + +def test_llama_meta_walk_fuses_common_projections(tmp_path, monkeypatch, capsys): + model_dir = _save(tmp_path, _llama_config()) + + output = _preview(model_dir, monkeypatch, capsys, tp=2) + + assert "layout: Transformers meta model; fused QKV and gate/up" in output + assert ( + "32x32 <- model.layers.*.self_attn.q_proj|model.layers.*.self_attn.k_proj|model.layers.*.self_attn.v_proj" + in output + ) + assert "32x16 <- model.layers.*.self_attn.o_proj" in output + assert "64x32 <- model.layers.*.mlp.gate_proj|model.layers.*.mlp.up_proj" in output + # Same-shape kernels keep separate --nks/--nk_names pairs. + assert "--nks 32,32 32,16 64,32 32,32" in output + assert "128x32" not in output # The output head is outside this benchmark. + + +def test_gqa_kv_heads_are_replicated_when_tp_exceeds_kv_heads(tmp_path): + config = _llama_config() + model_dir = _save(tmp_path, config) + _, model = benchmark_model._load_meta_model(str(model_dir), False, None) + + kernels, _, problems = benchmark_model._inspect_model(model, config, tp=4, ep=1) + + assert ( + 24, + 32, + "model.layers.*.self_attn.q_proj|model.layers.*.self_attn.k_proj|model.layers.*.self_attn.v_proj", + ) in kernels + assert problems == [] + + +def test_meta_loader_never_materializes_model_tensors(tmp_path): + model_dir = _save(tmp_path, _llama_config()) + + _, model = benchmark_model._load_meta_model(str(model_dir / "config.json"), False, None) + + tensors = list(model.named_parameters()) + list(model.named_buffers()) + assert tensors and all(tensor.is_meta for _, tensor in tensors) + + +def test_revision_does_not_reach_registered_model_constructor(tmp_path): + model_dir = _save(tmp_path, _llama_config()) + + _, model = benchmark_model._load_meta_model(str(model_dir), False, "main") + + assert type(model).__name__ == "LlamaForCausalLM" + + +def test_mixtral_modulelist_experts_use_ep(tmp_path): + config = transformers.MixtralConfig( + vocab_size=128, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + ) + model_dir = _save(tmp_path, config) + _, model = benchmark_model._load_meta_model(str(model_dir), False, None) + + _, moe, _ = benchmark_model._inspect_model(model, config, tp=2, ep=2) + + assert moe == benchmark_model._MoeShape(32, 48, 2, 2, "Swiglu") + + +def test_gpt_oss_direct_expert_tensors_are_inspected(tmp_path): + config = transformers.GptOssConfig( + vocab_size=128, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + ) + model_dir = _save(tmp_path, config) + _, model = benchmark_model._load_meta_model(str(model_dir), False, None) + + _, moe, _ = benchmark_model._inspect_model(model, config, tp=2, ep=2) + + assert moe == benchmark_model._MoeShape(32, 48, 2, 2, "Swiglu") + + +def test_nemotron_h_mamba_and_stacked_experts_are_inspected(tmp_path): + config = _nemotron_h_config() + model_dir = _save(tmp_path, config) + _, model = benchmark_model._load_meta_model(str(model_dir), False, None) + + experts = next(module for name, module in model.named_modules() if name.endswith(".experts")) + kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1) + + assert experts.up_proj.ndim == experts.down_proj.ndim == 3 + assert (42, 32, "model.layers.*.mixer.in_proj") in kernels + assert (32, 16, "model.layers.*.mixer.out_proj") in kernels + assert (20, 32, "model.layers.*.mixer.shared_experts.up_proj") in kernels + assert (32, 20, "model.layers.*.mixer.shared_experts.down_proj") in kernels + assert moe == benchmark_model._MoeShape(32, 24, 4, 2, "Relu2") + assert problems == [] + command = benchmark_model._command(kernels, moe, []) + assert command[command.index("--moe_activation_type") + 1] == "Relu2" + + with pytest.raises(benchmark_model.ShapeError, match=r"n_groups=2.*TP=4"): + benchmark_model._inspect_model(model, config, tp=4, ep=1) + + config.n_routed_experts = 8 + _, _, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1) + assert any("declares 8" in problem and "[4]" in problem for problem in problems) + + +@pytest.mark.parametrize( + ("config_cls_name", "model_cls_name"), + [ + ("Qwen3NextConfig", "Qwen3NextForCausalLM"), + ("Qwen3_5MoeTextConfig", "Qwen3_5MoeForCausalLM"), + ], +) +def test_gated_delta_net_kernels_are_derived(config_cls_name, model_cls_name): + config_cls = getattr(transformers, config_cls_name, None) + model_cls = getattr(transformers, model_cls_name, None) + if config_cls is None or model_cls is None: + pytest.skip(f"transformers does not provide {config_cls_name}") + assert config_cls is not None and model_cls is not None + from accelerate import init_empty_weights + + config = config_cls( + vocab_size=128, + hidden_size=32, + num_hidden_layers=1, + layer_types=["linear_attention"], + linear_num_key_heads=2, + linear_num_value_heads=4, + linear_key_head_dim=8, + linear_value_head_dim=8, + linear_conv_kernel_dim=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + intermediate_size=32, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=16, + shared_expert_intermediate_size=16, + decoder_sparse_step=1, + max_position_embeddings=64, + ) + with init_empty_weights(include_buffers=True): + model = model_cls(config) + + kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1) + + # key_dim = 2 heads x 8 = 16, value_dim = 4 heads x 8 = 32; vLLM's shared + # GDN mixer runs qkv+z and b+a as two fused per-rank GEMMs (both the + # pre-fused Qwen3-Next and split Qwen3.5 checkpoint layouts). + prefix = "model.layers.*.linear_attn" + if config_cls_name == "Qwen3NextConfig": + qkvz_label, ba_label = f"{prefix}.in_proj_qkvz", f"{prefix}.in_proj_ba" + else: + qkvz_label = f"{prefix}.in_proj_qkv|{prefix}.in_proj_z" + ba_label = f"{prefix}.in_proj_b|{prefix}.in_proj_a" + assert (48, 32, qkvz_label) in kernels + assert (4, 32, ba_label) in kernels + assert (32, 16, f"{prefix}.out_proj") in kernels + assert problems == [] + assert moe == benchmark_model._MoeShape(32, 8, 4, 2, "Swiglu") + + with pytest.raises( + benchmark_model.ShapeError, match=r"linear_num_key_heads=2 is not divisible by 4" + ): + benchmark_model._inspect_model(model, config, tp=4, ep=1) + + +def test_expert_audit_problem_is_not_masked_by_per_rank_validation(tmp_path): + config = _nemotron_h_config() + model_dir = _save(tmp_path, config) + _, model = benchmark_model._load_meta_model(str(model_dir), False, None) + config.n_routed_experts = 8 + + # EP=3 does not divide the instantiated expert count; the audit mismatch + # must still be reported instead of a masking divisibility error. + kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=3) + + assert kernels + assert moe is None + assert any("declares 8" in problem for problem in problems) + + +def test_moe_only_model_benchmarks_without_dense_kernels(): + class Expert(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(32, 48, bias=False) + self.down_proj = nn.Linear(48, 32, bias=False) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.experts = nn.ModuleList([Expert() for _ in range(4)]) + + model = nn.Module() + model.layers = nn.ModuleList([Block()]) + config = SimpleNamespace( + hidden_size=32, + num_attention_heads=4, + num_experts_per_tok=2, + mlp_hidden_act="relu2", + ) + + kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) + + assert kernels == [] + assert problems == [] + assert moe == benchmark_model._MoeShape(32, 48, 4, 2, "Relu2") + command = benchmark_model._command(kernels, moe, []) + assert "--nks" not in command + assert command[:2] == ["--moe_hidden_size", "32"] + + +def test_legacy_nongated_modulelist_experts_are_inspected(): + class Expert(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(32, 48, bias=False) + self.down_proj = nn.Linear(48, 32, bias=False) + + model = nn.Module() + model.experts = nn.ModuleList([Expert() for _ in range(4)]) + config = SimpleNamespace(num_experts_per_tok=2, mlp_hidden_act="relu2") + + assert benchmark_model._moe_shapes(model, config) == { + benchmark_model._MoeShape(32, 48, 4, 2, "Relu2") + } + + +def test_gate_and_up_projections_must_shard_individually(): + class Mlp(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(32, 6, bias=False) + self.up_proj = nn.Linear(32, 6, bias=False) + self.down_proj = nn.Linear(6, 32, bias=False) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.mlp = Mlp() + + model = nn.Module() + model.layers = nn.ModuleList([Block()]) + config = SimpleNamespace(hidden_size=32, num_attention_heads=4) + + # The summed width (12) divides by TP=4, but vLLM shards gate and up + # individually, so the per-projection width (6) must divide too. + with pytest.raises(benchmark_model.ShapeError, match=r"gate_proj=6 is not divisible by 4"): + benchmark_model._inspect_model(model, config, tp=4, ep=1) + + +def test_top_k_fallback_matches_the_modelopt_list(): + auto_quantize_cost = pytest.importorskip("modelopt.torch.quantization._auto_quantize_cost") + + assert benchmark_model._MOE_TOP_K_ATTRS_FALLBACK == auto_quantize_cost._ACTIVE_MOE_TOP_K_ATTRS + + +def test_top_k_covers_the_modelopt_attribute_aliases(): + assert benchmark_model._top_k(SimpleNamespace(num_selected_experts=2)) == 2 + assert benchmark_model._top_k(SimpleNamespace(num_experts_per_token=4)) == 4 + assert benchmark_model._top_k(SimpleNamespace(top_k=6)) == 6 + assert benchmark_model._top_k(SimpleNamespace(num_experts_per_tok=8, top_k=50)) == 8 + assert benchmark_model._top_k(SimpleNamespace()) is None + + +def test_gated_moe_activation_is_derived_or_rejected(): + assert ( + benchmark_model._moe_activation(SimpleNamespace(hidden_act="gelu_pytorch_tanh"), True) + == "Geglu" + ) + assert benchmark_model._moe_activation(SimpleNamespace(hidden_act="gelu_new"), True) == "Geglu" + # Exact gelu and quick_gelu are not served by vLLM's FlashInfer MoE path, + # so they must be rejected instead of timed via the tanh-GELU kernel. + for activation in ("gelu", "quick_gelu", "relu"): + with pytest.raises(benchmark_model.ShapeError, match="unsupported gated MoE activation"): + benchmark_model._moe_activation(SimpleNamespace(hidden_act=activation), True) + + +def test_mamba_single_group_is_replicated_across_tp(): + class Mixer(nn.Module): + intermediate_size = 32 + num_heads = 4 + n_groups = 1 + ssm_state_size = 4 + + def __init__(self): + super().__init__() + self.in_proj = nn.Linear(32, 76, bias=False) + self.out_proj = nn.Linear(32, 32, bias=False) + + model = nn.Module() + model.mixer = Mixer() + + assert benchmark_model._mamba_kernels(model, tp=2) == [ + (42, 32, "mixer.in_proj"), + (32, 16, "mixer.out_proj"), + ] + + +def test_unrecognized_decoder_linear_is_reported(): + class Block(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(32, 48, bias=False) + self.down_proj = nn.Linear(48, 32, bias=False) + self.unknown_proj = nn.Linear(32, 48, bias=False) + + model = nn.Module() + model.layers = nn.ModuleList([Block()]) + + assert benchmark_model._unsupported_decoder_linears(model) == [ + ("layers.0.unknown_proj", 48, 32) + ] + config = SimpleNamespace(hidden_size=32, num_attention_heads=4) + kernels, _, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) + + assert (48, 32, "layers.*.up_proj") in kernels + assert problems == ["unsupported decoder Linear GEMM layout(s): layers.0.unknown_proj (48x32)"] + + +def test_partial_inventory_is_printed_when_the_audit_fails(monkeypatch, capsys): + class Block(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(32, 48, bias=False) + self.down_proj = nn.Linear(48, 32, bias=False) + self.unknown_proj = nn.Linear(32, 48, bias=False) + + model = nn.Module() + model.layers = nn.ModuleList([Block()]) + config = SimpleNamespace(hidden_size=32, num_attention_heads=4, model_type="test") + monkeypatch.setattr(benchmark_model, "_load_meta_model", lambda *_: (config, model)) + monkeypatch.setattr(sys, "argv", [str(SCRIPT), "unused/model", "--print_only"]) + + with pytest.raises(SystemExit, match="2"): + benchmark_model.main() + + captured = capsys.readouterr() + assert "# 48x32 <- layers.*.up_proj" in captured.out + assert "# unsupported: unsupported decoder Linear GEMM layout(s)" in captured.out + assert "unknown_proj (48x32)" in captured.out + assert "benchmark_via_builtin.py" in captured.err + + +def test_declared_moe_without_supported_experts_is_reported(): + class Mlp(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(32, 48, bias=False) + self.down_proj = nn.Linear(48, 32, bias=False) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.mlp = Mlp() + + model = nn.Module() + model.layers = nn.ModuleList([Block()]) + config = SimpleNamespace( + hidden_size=32, + num_attention_heads=4, + num_local_experts=4, + num_experts_per_tok=2, + ) + + _, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) + + assert moe is None + assert problems == [ + "model declares 4 routed experts but no supported expert GEMM layout was found" + ] + + +def test_command_keeps_same_shape_kernels_as_separate_named_pairs(): + kernels = [ + (64, 32, "a_proj|b_proj"), + (64, 32, "c_proj"), + (32, 64, "down_proj"), + ] + + command = benchmark_model._command(kernels, None, []) + + assert command == [ + "--nks", + "64,32", + "64,32", + "32,64", + "--nk_names", + "a_proj|b_proj", + "c_proj", + "down_proj", + ] + + +def test_moe_sharding_interpretation_is_printed(monkeypatch, capsys): + class Expert(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(32, 48, bias=False) + self.down_proj = nn.Linear(48, 32, bias=False) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.experts = nn.ModuleList([Expert() for _ in range(6)]) + + model = nn.Module() + model.layers = nn.ModuleList([Block()]) + config = SimpleNamespace( + hidden_size=32, + num_attention_heads=4, + num_experts_per_tok=2, + mlp_hidden_act="relu2", + model_type="test", + ) + monkeypatch.setattr(benchmark_model, "_load_meta_model", lambda *_: (config, model)) + monkeypatch.setattr( + sys, "argv", [str(SCRIPT), "unused/model", "--tp", "2", "--ep", "2", "--print_only"] + ) + + benchmark_model.main() + + output = capsys.readouterr().out + assert "# MoE sharding: EP=2 partitions whole experts" in output + + # EP that is not a multiple of TP matches no modeled serving layout. + with pytest.raises( + benchmark_model.ShapeError, match=r"EP=2 is not a multiple of TP=4.*benchmark_via_builtin" + ): + benchmark_model._inspect_model(model, config, tp=4, ep=2) + + +def test_runner_is_invoked_in_process(tmp_path, monkeypatch): + model_dir = _save(tmp_path, _llama_config()) + launched = [] + runner = SimpleNamespace(main=launched.append) + monkeypatch.setattr(benchmark_model, "_load_runner", lambda: runner) + monkeypatch.setattr(sys, "argv", [str(SCRIPT), str(model_dir), "--ms", "1", "16"]) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + + benchmark_model.main() + + assert launched == [ + [ + "--nks", + "64,32", + "32,32", + "128,32", + "32,64", + "--nk_names", + "model.layers.*.self_attn.q_proj|model.layers.*.self_attn.k_proj|model.layers.*.self_attn.v_proj", + "model.layers.*.self_attn.o_proj", + "model.layers.*.mlp.gate_proj|model.layers.*.mlp.up_proj", + "model.layers.*.mlp.down_proj", + "--ms", + "1", + "16", + ] + ] + + +def test_router_gate_projection_is_not_treated_as_an_mlp(): + class Mlp(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(32, 48, bias=False) + self.down_proj = nn.Linear(48, 32, bias=False) + + class Router(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(32, 4, bias=False) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.mlp = Mlp() + self.router = Router() + + model = nn.Module() + model.layers = nn.ModuleList([Block()]) + config = SimpleNamespace(hidden_size=32, num_attention_heads=4) + + kernels, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) + + assert moe is None + assert problems == [] + assert kernels == [ + (48, 32, "layers.*.mlp.up_proj"), + (32, 48, "layers.*.mlp.down_proj"), + ] + + +@pytest.mark.parametrize( + ("option", "value"), [("--nks", "1,1"), ("--moe_activation_type", "Relu2")] +) +def test_derived_shapes_cannot_be_overridden(monkeypatch, capsys, option, value): + monkeypatch.setattr( + sys, + "argv", + [str(SCRIPT), "unused/model", option, value, "--print_only"], + ) + + with pytest.raises(SystemExit, match="2"): + benchmark_model.main() + assert "cannot be overridden" in capsys.readouterr().err + + +@pytest.mark.parametrize(("option", "value"), [("--tp", "0"), ("--ep", "-1")]) +def test_parallel_sizes_must_be_positive(monkeypatch, capsys, option, value): + monkeypatch.setattr( + sys, + "argv", + [str(SCRIPT), "unused/model", option, value, "--print_only"], + ) + + with pytest.raises(SystemExit, match="2"): + benchmark_model.main() + assert "expected a positive integer" in capsys.readouterr().err diff --git a/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py new file mode 100644 index 00000000000..e558c478233 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py @@ -0,0 +1,453 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import importlib.util +import sys +from pathlib import Path + +import pytest + +SCRIPT = Path(__file__).parents[1] / "scripts" / "benchmark_via_builtin.py" +SPEC = importlib.util.spec_from_file_location("benchmark_via_builtin", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = benchmark +SPEC.loader.exec_module(benchmark) + + +@pytest.mark.parametrize("value", ["1", "1,2,3", "a,2", "0,2", "2,-1"]) +def test_nk_pair_rejects_invalid_values(value): + with pytest.raises(argparse.ArgumentTypeError): + benchmark._nk_pair(value) + + +@pytest.mark.parametrize("value", ["0", "-1"]) +def test_positive_int_rejects_non_positive_values(value): + with pytest.raises(argparse.ArgumentTypeError, match="positive integer"): + benchmark._positive_int(value) + + +@pytest.mark.parametrize( + "option", + [ + "--ms", + "--dry_run_iters", + "--num_iters", + "--moe_hidden_size", + "--moe_intermediate_size", + "--moe_num_experts", + "--moe_top_k", + ], +) +def test_parser_rejects_zero_for_numeric_options(option, capsys): + with pytest.raises(SystemExit, match="2"): + benchmark._parser().parse_args(["--flashinfer_repo", "/unused", option, "0"]) + assert "not a positive integer" in capsys.readouterr().err + + +def test_gemm_cases_are_data_driven_and_preserve_requested_shapes(): + cases = benchmark._gemm_cases([1], [(65, 129)], []) + + assert {case.backend for case in cases} == { + "bf16", + "nvfp4_cudnn", + "nvfp4_cutlass", + "nvfp4_cutedsl", + "nvfp4_trtllm", + "fp8_cudnn", + "fp8_cutlass", + "fp8_trtllm", + } + assert len({case.tag for case in cases}) == len(cases) + assert {(case.m, case.n, case.k) for case in cases} == {(1, 65, 129)} + physical_shapes = { + case.backend: ( + case.argv[case.argv.index("--n") + 1], + case.argv[case.argv.index("--k") + 1], + ) + for case in cases + } + assert physical_shapes["nvfp4_cudnn"] == ("96", "160") + assert physical_shapes["nvfp4_cutlass"] == ("96", "160") + assert physical_shapes["nvfp4_cutedsl"] == ("96", "160") + assert physical_shapes["nvfp4_trtllm"] == ("65", "129") + assert all( + shape == ("65", "129") + for name, shape in physical_shapes.items() + if not name.startswith("nvfp4_") + ) + assert {case.quant for case in cases if case.quant is not None} == { + benchmark._QuantSpec("nvfp4", "128x4", 1, 160), + benchmark._QuantSpec("nvfp4", "8x4", 1, 129), + benchmark._QuantSpec("fp8", "static", 1, 129), + } + + +def test_nk_names_are_parallel_and_map_duplicate_shapes(): + labels_by_nk = benchmark._named_nks( + [(32, 64), (32, 64), (128, 256)], + ["o_proj", "out_proj", "qkv_proj"], + ) + + assert labels_by_nk == {(32, 64): ["o_proj", "out_proj"], (128, 256): ["qkv_proj"]} + # Unnamed shapes are deduplicated in first-seen order and label themselves. + assert benchmark._named_nks([(32, 64), (32, 64), (2, 3)], None) == { + (32, 64): ["32x64"], + (2, 3): ["2x3"], + } + with pytest.raises(ValueError, match="exactly one name"): + benchmark._named_nks([(32, 64)], ["o_proj", "out_proj"]) + + +def test_moe_cases_and_derived_with_quant_rows(): + cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2), []) + assert [(case.backend, case.with_quant) for case in cases] == [ + ("bf16_cutlass", False), + ("fp8_cutlass", False), + ("nvfp4_cutlass", False), + ("nvfp4_cutlass", True), + ("fp8_trtllm", False), + ("nvfp4_trtllm", False), + ("nvfp4_cutedsl", False), + ] + intermediate_sizes = { + (case.backend, case.with_quant): case.argv[case.argv.index("--intermediate_size") + 1] + for case in cases + } + assert intermediate_sizes == { + ("bf16_cutlass", False): "50", + ("nvfp4_cutlass", False): "50", + ("nvfp4_cutlass", True): "50", + ("fp8_cutlass", False): "64", + ("fp8_trtllm", False): "64", + ("nvfp4_trtllm", False): "64", + ("nvfp4_cutedsl", False): "50", + } + hidden_sizes = {case.backend: case.argv[case.argv.index("--hidden_size") + 1] for case in cases} + # Only the trtllm-gen NVFP4 MoE pads hidden (vLLM pads it to 256). + assert hidden_sizes["nvfp4_trtllm"] == "256" + assert all(value == "32" for row, value in hidden_sizes.items() if row != "nvfp4_trtllm") + routines = {case.backend: case.argv[case.argv.index("--routine") + 1] for case in cases} + assert routines["fp8_trtllm"] == "trtllm_fp8_per_tensor_scale_moe" + assert routines["nvfp4_trtllm"] == "trtllm_fp4_block_scale_moe" + assert routines["nvfp4_cutedsl"] == "cute_dsl_fp4_block_scale_moe" + # Only the trtllm-gen rows route in-kernel; they use a fixed renormalize + # method so rows stay comparable across models. + for case in cases: + if case.backend.endswith("_trtllm"): + assert case.argv[case.argv.index("--routing_method") + 1] == "renormalize" + else: + assert "--routing_method" not in case.argv + quants = {(case.backend, case.with_quant): case.quant for case in cases} + # FP8 rows share the static activation-quant timing; the trtllm-gen and + # CuteDSL NVFP4 rows use the linear-scale-layout quantize their kernels + # consume (trtllm at the 256-padded hidden). The NVFP4 CUTLASS pair is a + # direct fused measurement instead. + assert quants[("fp8_cutlass", False)] == benchmark._QuantSpec("fp8", "static", 1, 32) + assert quants[("fp8_trtllm", False)] == benchmark._QuantSpec("fp8", "static", 1, 32) + assert quants[("nvfp4_trtllm", False)] == benchmark._QuantSpec("nvfp4", "linear", 1, 256) + assert quants[("nvfp4_cutedsl", False)] == benchmark._QuantSpec("nvfp4", "linear", 1, 32) + assert quants[("bf16_cutlass", False)] is None + assert quants[("nvfp4_cutlass", False)] is None + assert quants[("nvfp4_cutlass", True)] is None + + fp8 = next(case for case in cases if case.backend == "fp8_cutlass") + assert fp8.quant is not None + fp8.result = 1.0 + fp8.quant_result = 2.0 + + assert benchmark._output_rows(fp8) == [(False, 1.0), (True, 3.0)] + + +def test_swiglustep_uses_gated_fp8_alignment(): + cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2, "SwigluStep"), []) + + fp8 = next(case for case in cases if case.backend == "fp8_cutlass") + assert fp8.argv[fp8.argv.index("--intermediate_size") + 1] == "64" + + +def test_non_gated_moe_pads_fp8_and_nvfp4_intermediate_to_128(): + cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2, "Relu2"), []) + + intermediate_sizes = { + (case.backend, case.with_quant): case.argv[case.argv.index("--intermediate_size") + 1] + for case in cases + } + # The Swiglu-only CuteDSL MoE row is not emitted for non-gated activations. + assert intermediate_sizes == { + ("bf16_cutlass", False): "50", + ("fp8_cutlass", False): "128", + ("nvfp4_cutlass", False): "128", + ("nvfp4_cutlass", True): "128", + ("fp8_trtllm", False): "128", + ("nvfp4_trtllm", False): "128", + } + + +def test_unavailable_fp8_quantization_is_written_as_an_error(monkeypatch, capsys, tmp_path): + case = benchmark._Case( + section="gemm", + tag="gemm_fp8_cutlass_MxNxK=1x32x64", + backend="fp8_cutlass", + m=1, + n=32, + k=64, + argv=[], + quant=benchmark._QuantSpec("fp8", "static", 1, 64), + result=1.0, + ) + monkeypatch.setattr(benchmark, "vllm_ops", None) + + benchmark._attach_quant_times([case], 1, 1, False) + output = tmp_path / "combined_results.csv" + benchmark._write_results(output, [case], {(32, 64): ["32x64"]}) + + assert case.quant_result == benchmark._FP8_QUANT_UNAVAILABLE + assert "[WARN] vLLM is unavailable for FP8 activation quantization" in capsys.readouterr().out + assert "32x64,1,32,64,fp8_cutlass,False,1.000\n" in output.read_text() + assert f"32x64,1,32,64,fp8_cutlass,True,{benchmark._FP8_QUANT_UNAVAILABLE}\n" in ( + output.read_text() + ) + + +def test_driver_errors_are_added_to_kernel_and_with_quant_rows(tmp_path): + case = benchmark._Case( + section="gemm", + tag="gemm_fp8_trtllm_MxNxK=8x1280x2880", + backend="fp8_trtllm", + m=8, + n=1280, + k=2880, + argv=[], + quant=benchmark._QuantSpec("fp8", "static", 8, 2880), + ) + output = [ + f"[ERROR] Error running test: --routine mm_fp8 --case_tag {case.tag}\n", + "[ERROR] Error: K must be divisible by 128, got 2880\n", + ] + + message = benchmark._parse_driver_error(output) + assert message == "K must be divisible by 128; got 2880" + case.result = f"ERROR: {message}" + csv_path = tmp_path / "combined_results.csv" + benchmark._write_results(csv_path, [case], {(1280, 2880): ["1280x2880"]}) + + expected = "ERROR: K must be divisible by 128; got 2880" + assert f"1280x2880,8,1280,2880,fp8_trtllm,False,{expected}\n" in csv_path.read_text() + assert f"1280x2880,8,1280,2880,fp8_trtllm,True,{expected}\n" in csv_path.read_text() + + +def test_empty_driver_error_has_no_synthetic_reason(): + output = [ + "[ERROR] Error running test: --routine mm_fp4 --case_tag gemm_nvfp4\n", + "[ERROR] Error:\n", + ] + + assert benchmark._parse_driver_error(output) == "" + # An error line without any message line is not a driver-reported error. + assert benchmark._parse_driver_error(output[:1]) is None + assert benchmark._parse_driver_error([]) is None + + +def test_write_results_emits_long_form_rows(tmp_path): + cases = [ + benchmark._Case("gemm", "a", "bf16", 1, 2, 3, [], result=1.25), + benchmark._Case( + "gemm", + "b", + "fp8_cutlass", + 8, + 4, + 5, + [], + quant=benchmark._QuantSpec("fp8", "static", 8, 5), + result=3.5, + quant_result=1.0, + ), + benchmark._Case( + "moe", + "c", + "fp8_cutlass", + 8, + None, + None, + [], + quant=benchmark._QuantSpec("fp8", "static", 8, 32), + result=2.0, + quant_result=0.5, + ), + # A case that never produced a result or an error emits no row. + benchmark._Case("gemm", "d", "fp8_trtllm", 1, 2, 3, []), + ] + output = tmp_path / "combined_results.csv" + benchmark._write_results( + output, + cases, + {(4, 5): ["q_proj|k_proj|v_proj"], (2, 3): ["in_proj", "out_proj"]}, + header="flashinfer test-header", + moe_label="H=32 F=50 E=4 top_k=2 activation=Relu2", + ) + + assert output.read_text() == ( + "flashinfer test-header\n" + "GEMM\n" + "module_name,M,N,K,backend,with_quant,runtime\n" + "q_proj|k_proj|v_proj,8,4,5,fp8_cutlass,False,3.500\n" + "q_proj|k_proj|v_proj,8,4,5,fp8_cutlass,True,4.500\n" + "in_proj,1,2,3,bf16,False,1.250\n" + "out_proj,1,2,3,bf16,False,1.250\n" + "\n" + "MoE\n" + "H=32 F=50 E=4 top_k=2 activation=Relu2\n" + "module_name,M,N,K,backend,with_quant,runtime\n" + "experts,8,,,fp8_cutlass,False,2.000\n" + "experts,8,,,fp8_cutlass,True,2.500\n" + ) + + +def test_top_k_cannot_exceed_expert_count(monkeypatch, capsys): + monkeypatch.setattr( + sys, + "argv", + [ + str(SCRIPT), + "--flashinfer_repo", + "/unused", + "--moe_hidden_size", + "4", + "--moe_intermediate_size", + "8", + "--moe_num_experts", + "1", + "--moe_top_k", + "2", + ], + ) + + with pytest.raises(SystemExit, match="2"): + benchmark.main() + assert "--moe_top_k cannot exceed --moe_num_experts" in capsys.readouterr().err + + +@pytest.mark.parametrize( + ("returncode", "expected_reason"), + [ + (0, "FlashInfer produced no result row"), + (1, "FlashInfer driver exited with status 1"), + ], +) +def test_missing_builtin_results_still_writes_combined_errors( + monkeypatch, tmp_path, returncode, expected_reason +): + benchmarks_dir = tmp_path / "flashinfer" / "benchmarks" + benchmarks_dir.mkdir(parents=True) + (benchmarks_dir / "flashinfer_benchmark.py").write_text("") + workdir = tmp_path / "results" + monkeypatch.setattr(benchmark, "_run_case", lambda *_: (returncode, [])) + monkeypatch.setattr( + sys, + "argv", + [ + str(SCRIPT), + "--flashinfer_repo", + str(benchmarks_dir.parent), + "--ms", + "1", + "--nks", + "2,3", + "--workdir", + str(workdir), + ], + ) + + with pytest.raises(RuntimeError, match="FlashInfer failed benchmark cases"): + benchmark.main() + + assert not (workdir / "builtin_results.csv").exists() + combined = (workdir / "combined_results.csv").read_text() + assert f"2x3,1,2,3,bf16,False,ERROR: {expected_reason}" in combined + assert "driver.log" in combined + # The reproducibility header leads both the combined CSV and driver.log. + assert combined.splitlines()[0].startswith("flashinfer ") + assert (workdir / "driver.log").read_text().startswith("flashinfer ") + + +def test_case_rows_with_foreign_tags_are_treated_as_failures(monkeypatch, tmp_path): + benchmarks_dir = tmp_path / "flashinfer" / "benchmarks" + benchmarks_dir.mkdir(parents=True) + (benchmarks_dir / "flashinfer_benchmark.py").write_text("") + workdir = tmp_path / "results" + + def fake_run_case(benchmarks_dir, argv, log): + output = Path(argv[argv.index("--output_path") + 1]) + output.write_text("case_tag,median_time\nsomeone_else,0.001\n") + return 0, [] + + monkeypatch.setattr(benchmark, "_run_case", fake_run_case) + monkeypatch.setattr( + sys, + "argv", + [ + str(SCRIPT), + "--flashinfer_repo", + str(benchmarks_dir.parent), + "--ms", + "1", + "--nks", + "2,3", + "--workdir", + str(workdir), + ], + ) + + with pytest.raises(RuntimeError, match="FlashInfer failed benchmark cases"): + benchmark.main() + + combined = (workdir / "combined_results.csv").read_text() + assert "no result row" in combined + + +def test_run_case_streams_and_appends_to_the_driver_log(tmp_path, capsys): + benchmarks_dir = tmp_path / "benchmarks" + benchmarks_dir.mkdir() + (benchmarks_dir / "flashinfer_benchmark.py").write_text( + "print('line one')\nprint('line two')\n" + ) + driver_log = tmp_path / "driver.log" + + with driver_log.open("w") as log: + returncode, lines = benchmark._run_case(benchmarks_dir, ["--unused"], log) + + assert returncode == 0 + assert lines == ["line one\n", "line two\n"] + assert driver_log.read_text() == "line one\nline two\n" + assert "line one" in capsys.readouterr().out + + +def test_write_builtin_merges_heterogeneous_row_columns(tmp_path): + path = tmp_path / "builtin_results.csv" + + benchmark._write_builtin( + path, + [ + {"routine": "mm_bf16", "median_time": "0.004", "case_tag": "a"}, + {"routine": "cutlass_fused_moe", "case_tag": "b", "num_experts": "8"}, + ], + ) + + assert path.read_text() == ( + "routine,median_time,case_tag,num_experts\nmm_bf16,0.004,a,\ncutlass_fused_moe,,b,8\n" + )