diff --git a/.agents/skills/benchmark-model-kernels/SKILL.md b/.agents/skills/benchmark-model-kernels/SKILL.md new file mode 100644 index 00000000000..02a32797cb2 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/SKILL.md @@ -0,0 +1,136 @@ +--- +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, MoE tuple, and derived routing 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, 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 including the routing method; overrides are rejected. +- Physical padding follows one rule: pad a dimension if and only if vLLM pads + it for that case; otherwise keep the exact shape and let it fail the same + way vLLM would. Row labels stay logical — report both shapes when padding + changes a dimension. +- BF16 and FP8 MoE cases run before NVFP4 because a CUDA fault poisons the + remaining cases in the same driver process. +- 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). Drivers + with the conditional-shuffle fix drop only that backend (its cell reports no + result row); stock 0.6.x drivers fail *all* backends for that shape with an + empty assertion. +- `mm_fp8` (trtllm_low_latency) needs `K % 128 == 0`. +- 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 +``` + +For a manual MoE supplement, also derive and pass the routing arguments +(`--moe_routing_method` plus the DeepSeek group/scaling/bias flags when +applicable) — the runner otherwise defaults to `topk`, a different routing +kernel. 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` (microseconds; GEMM rows grouped +under `: ` labels, MoE rows flat). `*_with_quant` rows add a +separately measured activation-quantization time, except NVFP4 MoE with +quantization, which is a single fused measurement. 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..98cf92d5e51 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py @@ -0,0 +1,657 @@ +# 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 and common routed-expert layouts, maps +config routing fields to FlashInfer's MoE routing method, 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 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 + + +@dataclass(frozen=True) +class _MoeRouting: + method: str + num_expert_group: int | None = None + topk_group: int | None = None + routed_scaling_factor: float | None = None + use_routing_bias: bool = False + + +_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", + "--moe_routing_method", + "--moe_num_expert_group", + "--moe_topk_group", + "--moe_routed_scaling_factor", + "--moe_use_routing_bias", +} + +_ROUTER_NAMES = {"gate", "router", "router_proj"} + + +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): + model = AutoModelForCausalLM.from_config(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 _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 + return local_n, q[1], "fused_qkv" + + +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(".") + if ( + leaf not in _PROJECTIONS + or ".experts." in name + or ".local_experts." in name + or any(part in {"router", "routers"} 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"), "attention_out")) + + 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}") + n = _divide(gate_n + up_n, tp, f"{parent}.gate_up") + kernels.append((n, gate_k, "fused_gate_up")) + 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, "up")) + if "gate_up_proj" in layers: + n, k = layers["gate_up_proj"] + kernels.append((_divide(n, tp, f"{parent}.gate_up_proj"), k, "fused_gate_up")) + if "down_proj" in layers: + n, k = layers["down_proj"] + kernels.append((n, _divide(k, tp, f"{parent}.down_proj"), "down")) + 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, "mamba_in"), + (hidden, local_intermediate, "mamba_out"), + ] + ) + 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" + if normalized == "geglu" or normalized.startswith("gelu") or normalized == "quick_gelu": + return "Geglu" + raise ShapeError(f"unsupported gated MoE activation {configured!r}") + 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] + + +def _top_k(config: Any) -> int | None: + for attr in ("num_experts_per_tok", "moe_top_k"): + 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 _moe_routing(model: Any, config: Any) -> _MoeRouting: + # DeepSeek-style group-limited routing is declared by these three config + # fields together; the score-correction bias is a tensor, not a config field. + groups = getattr(config, "n_group", None) + topk_group = getattr(config, "topk_group", None) + scaling = getattr(config, "routed_scaling_factor", None) + if groups is not None and topk_group is not None and scaling is not None: + tensors = list(model.named_parameters()) + list(model.named_buffers()) + use_bias = any(name.rsplit(".", 1)[-1] == "e_score_correction_bias" for name, _ in tensors) + return _MoeRouting("deepseek_v3", int(groups), int(topk_group), float(scaling), use_bias) + if getattr(config, "norm_topk_prob", False): + return _MoeRouting("renormalize") + return _MoeRouting("topk") + + +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"}) + + 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", "routers"} 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 _ROUTER_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, _MoeRouting | None, list[str]]: + config = getattr(config, "text_config", None) or config + kernels = _dense_kernels(model, config, tp) + _mamba_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) + routing = 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) + routing = _moe_routing(model, config) + 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, routing, problems + + +def _command( + kernels: list[_Kernel], + moe: _MoeShape | None, + routing: _MoeRouting | None, + passthrough: list[str], +) -> list[str]: + names: dict[tuple[int, int], list[str]] = {} + for n, k, label in kernels: + labels = names.setdefault((n, k), []) + if label not in labels: + labels.append(label) + command: list[str] = [] + if names: + command += ["--nks", *(f"{n},{k}" for n, k in names)] + command += ["--nk_names", *("/".join(labels) for labels in names.values())] + 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] + if routing: + command += ["--moe_routing_method", routing.method] + if routing.num_expert_group is not None: + command += ["--moe_num_expert_group", str(routing.num_expert_group)] + if routing.topk_group is not None: + command += ["--moe_topk_group", str(routing.topk_group)] + if routing.routed_scaling_factor is not None: + command += ["--moe_routed_scaling_factor", str(routing.routed_scaling_factor)] + if routing.use_routing_bias: + command.append("--moe_use_routing_bias") + 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, routing, 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 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)") + if routing: + groups = ( + f" n_group={routing.num_expert_group} topk_group={routing.topk_group}" + f" scaling={routing.routed_scaling_factor} bias={routing.use_routing_bias}" + if routing.method == "deepseek_v3" + else "" + ) + print(f"# MoE routing: {routing.method}{groups}") + 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, routing, 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..8840a48cf86 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py @@ -0,0 +1,660 @@ +# 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; NVFP4 MoE with quantization is measured +directly. 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 shlex +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + 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", +) +_MOE_ROUTINGS = ("renormalize", "deepseek_v3", "llama4", "renormalize_naive", "topk") + +_ResultValue = float | str + + +@dataclass +class _Case: + section: str + tag: str + key: str + argv: list[str] + quant: tuple[str, int, int] | None = None + + +def _index_cases(cases: list[_Case]) -> dict[str, _Case]: + indexed = {} + for case in cases: + if case.tag in indexed: + raise RuntimeError(f"duplicate benchmark case tag: {case.tag}") + indexed[case.tag] = case + return indexed + + +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 +) -> tuple[list[tuple[int, int]], list[str] | None]: + if names is not None and len(names) != len(nks): + raise ValueError("--nk_names must contain exactly one name for each --nks pair") + + unique_nks = list(dict.fromkeys(nks)) + if names is None: + return unique_nks, None + + names_by_nk: dict[tuple[int, int], list[str]] = {} + for nk, name in zip(nks, names, strict=True): + labels = names_by_nk.setdefault(nk, []) + if name not in labels: + labels.append(name) + return unique_nks, ["/".join(names_by_nk[nk]) for nk in unique_nks] + + +def _round_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _parse_driver_errors(lines: list[str]) -> dict[str, str]: + errors = {} + pending_tag = None + for line in lines: + stripped = line.strip() + if stripped.startswith(_ERROR_CASE_PREFIX): + command = stripped.removeprefix(_ERROR_CASE_PREFIX).strip() + try: + argv = shlex.split(command) + tag_index = argv.index("--case_tag") + pending_tag = argv[tag_index + 1] + except (ValueError, IndexError): + pending_tag = None + elif stripped.startswith(_ERROR_MESSAGE_PREFIX) and pending_tag is not None: + message = stripped.removeprefix(_ERROR_MESSAGE_PREFIX).strip().replace(",", ";") + errors[pending_tag] = message + pending_tag = None + return errors + + +def _run_driver( + benchmarks_dir: Path, testlist: Path, output: Path, driver_log: Path +) -> tuple[int, list[str]]: + # This invokes the explicitly selected FlashInfer checkout without a shell. + process = subprocess.Popen( + [ + sys.executable, + "flashinfer_benchmark.py", + "--testlist", + str(testlist.resolve()), + "--output_path", + str(output.resolve()), + ], + cwd=benchmarks_dir, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert process.stdout is not None + lines = [] + with driver_log.open("w") as log: + for line in process.stdout: + print(line, end="", flush=True) + log.write(line) + lines.append(line) + return process.wait(), lines + + +def _gemm_cases( + ms: list[int], + nks: list[tuple[int, int]], + common: list[str], +) -> list[_Case]: + cases = [] + + for m in ms: + for n, k in nks: + variants: list[tuple[str, str, str, int, int, list[str], str | None]] = [ + ("bf16", "mm_bf16", "cudnn", n, k, [], None) + ] + for backend in ("cudnn", "cutlass", "trtllm"): + layout = "128x4" if 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 backend != "trtllm": + run_n, run_k = _round_up(n, 32), _round_up(k, 32) + variants.append( + ( + f"nvfp4_{backend}", + "mm_fp4", + backend, + run_n, + run_k, + extra, + f"nvfp4_{layout}", + ) + ) + variants += [ + ( + f"fp8_{backend}", + "bmm_fp8", + backend, + n, + k, + ["--batch_size", "1"], + "fp8_static", + ) + for backend in ("cudnn", "cutlass") + ] + variants.append(("fp8_trtllm", "mm_fp8", "trtllm_low_latency", n, k, [], "fp8_static")) + + for name, routine, backend, run_n, run_k, extra, quant_kind in variants: + key = f"{name}_MxNxK={m}x{n}x{k}" + quant = (quant_kind, m, run_k) if quant_kind else None + cases.append( + _Case( + section="gemm", + tag=f"gemm_{key}", + key=key, + argv=[ + "--routine", + routine, + "--backends", + backend, + *extra, + "--m", + str(m), + "--n", + str(run_n), + "--k", + str(run_k), + *common, + ], + quant=quant, + ) + ) + return cases + + +def _moe_cases( + ms: list[int], + shape: tuple[int, int, int, int] | None, + common: list[str], + activation: str | None, + routing: str | None, + routing_args: list[str] | None = None, +) -> list[_Case]: + if shape is None: + return [] + routine = "cutlass_fused_moe" + + hidden, intermediate, experts, top_k = shape + shape_args = [ + "--hidden_size", + str(hidden), + "--num_experts", + str(experts), + "--top_k", + str(top_k), + ] + if activation: + shape_args += ["--activation-type", activation] + if routing: + shape_args += ["--routing_method", routing] + shape_args += routing_args or [] + + cases = [] + gated = activation is None or activation in { + "Swiglu", + "Geglu", + "SwigluBias", + "SwigluStep", + } + fp8_intermediate = _round_up(intermediate, 16 if gated else 128) + # Pad a dimension only when vLLM pads it: vLLM pads non-gated NVFP4 expert + # weights up to the 128-aligned swizzled scale rows, but raises instead of + # padding gated NVFP4, so the gated case stays exact and may fail like vLLM. + nvfp4_intermediate = intermediate if gated else _round_up(intermediate, 128) + variants = ( + ("bf16_cutlass_moe", intermediate, []), + ("fp8_cutlass_moe", fp8_intermediate, ["--cutlass_variant", "fp8"]), + ( + "nvfp4_cutlass_moe", + nvfp4_intermediate, + ["--cutlass_variant", "nvfp4", "--quantized_input"], + ), + ("nvfp4_cutlass_moe_with_quant", nvfp4_intermediate, ["--cutlass_variant", "nvfp4"]), + ) + for row, run_intermediate, extra in variants: + for m in ms: + key = f"{row}_M={m}" + quant = ("fp8_static", m, hidden) if row == "fp8_cutlass_moe" else None + cases.append( + _Case( + section="moe", + tag=f"moe_{key}", + key=key, + argv=[ + "--routine", + routine, + "--num_tokens", + str(m), + *shape_args, + "--intermediate_size", + str(run_intermediate), + *extra, + *common, + ], + quant=quant, + ) + ) + return cases + + +def _nvfp4_runner(tensor: torch.Tensor, layout: str): + import flashinfer + + global_scale = (448 * 6) / tensor.float().abs().nan_to_num().max() + 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 _quant_times( + cases: list[_Case], dry_runs: int, iterations: int, cuda_graph: bool +) -> dict[tuple[str, int, int], _ResultValue]: + results: dict[tuple[str, int, int], _ResultValue] = {} + specs = {case.quant for case in cases if case.quant is not None} + warned_fp8 = False + for kind, m, k in sorted(specs): + if not kind.startswith("nvfp4_") and vllm_ops is None: + results[(kind, m, k)] = _FP8_QUANT_UNAVAILABLE + if not warned_fp8: + print(f"[WARN] {_FP8_QUANT_UNAVAILABLE.removeprefix('ERROR: ')}") + warned_fp8 = True + continue + # 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 + + tensor = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + runner = ( + _nvfp4_runner(tensor, kind.removeprefix("nvfp4_")) + if kind.startswith("nvfp4_") + else _fp8_runner(tensor) + ) + kernel, inputs = runner + 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[(kind, m, k)] = float(np.median(times)) * 1000 + return results + + +def _combine( + cases_by_tag: dict[str, _Case], + rows: list[dict[str, str]], + quant: dict[tuple[str, int, int], _ResultValue], + errors: dict[str, str] | None = None, +) -> dict[str, dict[str, _ResultValue]]: + results: dict[str, dict[str, _ResultValue]] = {"gemm": {}, "moe": {}} + for row in rows: + case = cases_by_tag.get(row.get("case_tag", "")) + if case is None: + continue + value = float(row["median_time"]) * 1000 + results[case.section][case.key] = value + if case.quant is not None and case.quant in quant: + separator = "_MxNxK=" if case.section == "gemm" else "_M=" + name, shape = case.key.split(separator, 1) + quant_value = quant[case.quant] + results[case.section][f"{name}_with_quant{separator}{shape}"] = ( + value + quant_value if isinstance(quant_value, float) else quant_value + ) + for tag, reason in (errors or {}).items(): + case = cases_by_tag.get(tag) + if case is None: + continue + error = f"ERROR: {reason}" + results[case.section].setdefault(case.key, error) + if case.quant is not None: + separator = "_MxNxK=" if case.section == "gemm" else "_M=" + name, shape = case.key.split(separator, 1) + results[case.section].setdefault(f"{name}_with_quant{separator}{shape}", error) + return results + + +def _format_result(value: _ResultValue | None) -> str: + if value is None: + return "" + if isinstance(value, float): + return f"{value:.3f}" + return value + + +def _write_results( + path: Path, + results: dict[str, dict[str, _ResultValue]], + ms: list[int], + nks: list[tuple[int, int]], + nk_names: list[str] | None = None, +) -> None: + with path.open("w", newline="") as stream: + writer = csv.writer(stream, lineterminator="\n") + gemm = results["gemm"] + if gemm: + writer.writerow(["GEMM"]) + writer.writerow(["M", *ms]) + names = sorted({key.split("_MxNxK=", 1)[0] for key in gemm}) + for index, (n, k) in enumerate(nks): + shape = f"{n}x{k}" + label = f"{nk_names[index]}: {shape}" if nk_names is not None else shape + writer.writerow([label]) + for name in names: + cells = [gemm.get(f"{name}_MxNxK={m}x{n}x{k}") for m in ms] + writer.writerow( + [ + name, + *(_format_result(value) for value in cells), + ] + ) + + moe = results["moe"] + if moe: + if gemm: + writer.writerow([]) + writer.writerow(["MoE"]) + writer.writerow(["M", *ms]) + names = sorted({key.split("_M=", 1)[0] for key in moe}) + for name in names: + cells = [moe.get(f"{name}_M={m}") for m in ms] + writer.writerow([name, *(_format_result(value) for value in cells)]) + + +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( + "--moe_routing_method", + choices=_MOE_ROUTINGS, + default="topk", + help="FlashInfer routing, e.g. renormalize", + ) + parser.add_argument("--moe_num_expert_group", type=_positive_int) + parser.add_argument("--moe_topk_group", type=_positive_int) + parser.add_argument("--moe_routed_scaling_factor", type=float) + parser.add_argument("--moe_use_routing_bias", action="store_true") + parser.add_argument("--workdir", type=Path, default=Path("benchmark_via_builtin_out")) + return parser + + +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: + nks, nk_names = _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") + if not nks and not any(moe_values): + parser.error("pass --nks and/or all four --moe_* shape arguments") + if all(moe_values) and args.moe_top_k > args.moe_num_experts: + parser.error("--moe_top_k cannot exceed --moe_num_experts") + deepseek_values = ( + args.moe_num_expert_group, + args.moe_topk_group, + args.moe_routed_scaling_factor, + ) + if args.moe_routing_method == "deepseek_v3" and not all( + value is not None for value in deepseek_values + ): + parser.error( + "DeepSeekV3 routing requires --moe_num_expert_group, --moe_topk_group, " + "and --moe_routed_scaling_factor" + ) + if args.moe_routed_scaling_factor is not None and args.moe_routed_scaling_factor <= 0: + parser.error("--moe_routed_scaling_factor must be positive") + if ( + args.moe_num_expert_group is not None + and args.moe_topk_group is not None + and args.moe_topk_group > args.moe_num_expert_group + ): + parser.error("--moe_topk_group cannot exceed --moe_num_expert_group") + + 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") + + gemm_cases = _gemm_cases(ms, nks, common) + moe_shape = moe_values if all(moe_values) else None + moe_routing_args = [] + if args.moe_num_expert_group is not None: + moe_routing_args += ["--n_group", str(args.moe_num_expert_group)] + if args.moe_topk_group is not None: + moe_routing_args += ["--topk_group", str(args.moe_topk_group)] + if args.moe_routed_scaling_factor is not None: + moe_routing_args += [ + "--routed_scaling_factor", + str(args.moe_routed_scaling_factor), + ] + if args.moe_use_routing_bias: + moe_routing_args.append("--use_routing_bias") + moe_cases = _moe_cases( + ms, + moe_shape, + common, + args.moe_activation_type, + args.moe_routing_method, + moe_routing_args, + ) + cases = gemm_cases + moe_cases + + 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") + cases_by_tag = _index_cases(cases) + testlist.write_text( + "\n".join(shlex.join([*case.argv, "--case_tag", case.tag]) for case in cases) + "\n" + ) + + returncode, driver_output = _run_driver(benchmarks_dir, testlist, builtin_csv, driver_log) + + rows = [] + if builtin_csv.is_file(): + with builtin_csv.open(newline="") as stream: + rows = list(csv.DictReader(stream)) + completed_tags = {row.get("case_tag") for row in rows} + parsed_errors = _parse_driver_errors(driver_output) + missing_reason = ( + f"FlashInfer driver exited with status {returncode} before this case completed" + f" (an earlier CUDA fault can poison later cases); see {driver_log}" + if returncode + else f"FlashInfer produced no result row and no error message; see {driver_log}" + ) + empty_error_reason = ( + f"FlashInfer reported an error without a message (empty exception); see {driver_log}" + ) + errors = {} + for case in cases: + if case.tag in completed_tags: + continue + if case.tag in parsed_errors: + errors[case.tag] = parsed_errors[case.tag] or empty_error_reason + else: + errors[case.tag] = missing_reason + + completed_cases = [case for case in cases if case.tag in completed_tags] + quant = _quant_times( + completed_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, + ) + results = _combine(cases_by_tag, rows, quant, errors) + _write_results(combined_csv, results, ms, nks, nk_names) + print(f"Wrote {combined_csv}") + if errors: + failed = [cases_by_tag[tag].key for tag in errors if tag in cases_by_tag] + raise RuntimeError( + "FlashInfer failed benchmark cases: " + + ", ".join(failed) + + f"; wrote failure details to {combined_csv}" + ) + if returncode: + raise RuntimeError(f"FlashInfer driver exited with status {returncode}") + + +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..b26b3f99382 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py @@ -0,0 +1,533 @@ +# 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 <- fused_qkv" in output + assert "32x16 <- attention_out" in output + assert "64x32 <- fused_gate_up" in output + assert "--nks 32,32 32,16 64,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, "fused_qkv") 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, routing, _ = benchmark_model._inspect_model(model, config, tp=2, ep=2) + + assert moe == benchmark_model._MoeShape(32, 48, 2, 2, "Swiglu") + assert routing == benchmark_model._MoeRouting("topk") + + +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, routing, problems = benchmark_model._inspect_model(model, config, tp=2, ep=1) + + assert experts.up_proj.ndim == experts.down_proj.ndim == 3 + assert (42, 32, "mamba_in") in kernels + assert (32, 16, "mamba_out") in kernels + assert (20, 32, "up") in kernels + assert (32, 20, "down") in kernels + assert moe == benchmark_model._MoeShape(32, 24, 4, 2, "Relu2") + assert problems == [] + # NemotronH declares DeepSeek-style routing fields and a score-correction + # bias buffer on its router. + assert routing == benchmark_model._MoeRouting("deepseek_v3", 1, 1, 1.0, True) + command = benchmark_model._command(kernels, moe, routing, []) + assert command[command.index("--moe_activation_type") + 1] == "Relu2" + assert command[command.index("--moe_routing_method") + 1] == "deepseek_v3" + assert command[command.index("--moe_num_expert_group") + 1] == "1" + assert command[command.index("--moe_topk_group") + 1] == "1" + assert command[command.index("--moe_routed_scaling_factor") + 1] == "1.0" + assert "--moe_use_routing_bias" in command + + 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) + + +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, routing, problems = benchmark_model._inspect_model(model, config, tp=1, ep=3) + + assert kernels + assert moe is None + assert routing 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, routing, 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, routing, []) + 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_gated_moe_activation_is_derived_or_rejected(): + assert ( + benchmark_model._moe_activation(SimpleNamespace(hidden_act="gelu_pytorch_tanh"), True) + == "Geglu" + ) + with pytest.raises(benchmark_model.ShapeError, match="unsupported gated MoE activation"): + benchmark_model._moe_activation(SimpleNamespace(hidden_act="relu"), 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, "mamba_in"), + (32, 16, "mamba_out"), + ] + + +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, "up") 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 <- up" 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_moe_routing_is_derived_from_config_fields(): + model = nn.Module() + deepseek = SimpleNamespace(n_group=2, topk_group=1, routed_scaling_factor=2.5) + renormalize = SimpleNamespace(norm_topk_prob=True) + + assert benchmark_model._moe_routing(model, deepseek) == benchmark_model._MoeRouting( + "deepseek_v3", 2, 1, 2.5, False + ) + assert benchmark_model._moe_routing(model, renormalize) == benchmark_model._MoeRouting( + "renormalize" + ) + assert benchmark_model._moe_routing(model, SimpleNamespace()) == benchmark_model._MoeRouting( + "topk" + ) + + +def test_command_names_gemm_shapes_and_merges_duplicates(): + kernels = [ + (64, 32, "fused_qkv"), + (64, 32, "fused_gate_up"), + (32, 64, "down"), + ] + + command = benchmark_model._command(kernels, None, None, []) + + assert command == [ + "--nks", + "64,32", + "32,64", + "--nk_names", + "fused_qkv/fused_gate_up", + "down", + ] + + +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", + "fused_qkv", + "attention_out", + "fused_gate_up", + "down", + "--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, routing, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) + + assert moe is None + assert routing is None + assert problems == [] + assert kernels == [(48, 32, "up"), (32, 48, "down")] + + +@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..6652282afeb --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py @@ -0,0 +1,370 @@ +# 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.key.split("_MxNxK=", 1)[0] for case in cases} == { + "bf16", + "nvfp4_cudnn", + "nvfp4_cutlass", + "nvfp4_trtllm", + "fp8_cudnn", + "fp8_cutlass", + "fp8_trtllm", + } + assert len({case.tag for case in cases}) == len(cases) + assert {case.key.split("_MxNxK=", 1)[1] for case in cases} == {"1x65x129"} + physical_shapes = { + case.key.split("_MxNxK=", 1)[0]: ( + 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_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} == { + ("nvfp4_128x4", 1, 160), + ("nvfp4_8x4", 1, 129), + ("fp8_static", 1, 129), + } + + +def test_nk_names_are_parallel_and_merge_duplicate_shapes(): + nks, names = benchmark._named_nks( + [(32, 64), (32, 64), (128, 256)], + ["o_proj", "out_proj", "qkv_proj"], + ) + + assert nks == [(32, 64), (128, 256)] + assert names == ["o_proj/out_proj", "qkv_proj"] + with pytest.raises(ValueError, match="exactly one name"): + benchmark._named_nks([(32, 64)], ["o_proj", "out_proj"]) + + +def test_case_tags_are_unique_join_keys(): + case = benchmark._Case("gemm", "duplicate", "bf16_MxNxK=1x2x3", []) + + assert benchmark._index_cases([case]) == {case.tag: case} + with pytest.raises(RuntimeError, match="duplicate benchmark case tag: duplicate"): + benchmark._index_cases([case, case]) + + +def test_moe_cases_and_result_combination(): + cases = benchmark._moe_cases([1], (32, 50, 4, 2), [], None, None) + assert len(cases) == 4 + assert [case.key.split("_M=", 1)[0] for case in cases] == [ + "bf16_cutlass_moe", + "fp8_cutlass_moe", + "nvfp4_cutlass_moe", + "nvfp4_cutlass_moe_with_quant", + ] + intermediate_sizes = { + case.key.split("_M=", 1)[0]: case.argv[case.argv.index("--intermediate_size") + 1] + for case in cases + } + assert intermediate_sizes == { + "bf16_cutlass_moe": "50", + "nvfp4_cutlass_moe": "50", + "nvfp4_cutlass_moe_with_quant": "50", + "fp8_cutlass_moe": "64", + } + + fp8 = next(case for case in cases if case.key == "fp8_cutlass_moe_M=1") + assert fp8.quant is not None + rows = [{"case_tag": fp8.tag, "median_time": "0.001"}] + quant = {fp8.quant: 2.0} + + results = benchmark._combine({fp8.tag: fp8}, rows, quant) + + assert results["moe"]["fp8_cutlass_moe_M=1"] == 1.0 + assert results["moe"]["fp8_cutlass_moe_with_quant_M=1"] == 3.0 + + +def test_swiglustep_uses_gated_fp8_alignment(): + cases = benchmark._moe_cases([1], (32, 50, 4, 2), [], "SwigluStep", "topk") + + fp8 = next(case for case in cases if case.key == "fp8_cutlass_moe_M=1") + 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], (32, 50, 4, 2), [], "Relu2", "topk") + + intermediate_sizes = { + case.key.split("_M=", 1)[0]: case.argv[case.argv.index("--intermediate_size") + 1] + for case in cases + } + assert intermediate_sizes == { + "bf16_cutlass_moe": "50", + "fp8_cutlass_moe": "128", + "nvfp4_cutlass_moe": "128", + "nvfp4_cutlass_moe_with_quant": "128", + } + + +def test_moe_cases_forward_deepseek_routing_metadata(): + routing_args = [ + "--n_group", + "1", + "--topk_group", + "1", + "--routed_scaling_factor", + "2.5", + "--use_routing_bias", + ] + cases = benchmark._moe_cases( + [1], + (32, 48, 4, 2), + [], + "Relu2", + "deepseek_v3", + routing_args, + ) + + for case in cases: + assert case.argv[case.argv.index("--routing_method") + 1] == "deepseek_v3" + assert case.argv[case.argv.index("--n_group") + 1] == "1" + assert case.argv[case.argv.index("--topk_group") + 1] == "1" + assert case.argv[case.argv.index("--routed_scaling_factor") + 1] == "2.5" + assert "--use_routing_bias" in case.argv + + +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", + key="fp8_cutlass_MxNxK=1x32x64", + argv=[], + quant=("fp8_static", 1, 64), + ) + monkeypatch.setattr(benchmark, "vllm_ops", None) + + quant = benchmark._quant_times([case], 1, 1, False) + results = benchmark._combine( + {case.tag: case}, [{"case_tag": case.tag, "median_time": "0.001"}], quant + ) + output = tmp_path / "combined_results.csv" + benchmark._write_results(output, results, [1], [(32, 64)]) + + assert quant == {case.quant: benchmark._FP8_QUANT_UNAVAILABLE} + assert "[WARN] vLLM is unavailable for FP8 activation quantization" in capsys.readouterr().out + assert results["gemm"][case.key] == 1.0 + assert ( + results["gemm"]["fp8_cutlass_with_quant_MxNxK=1x32x64"] == benchmark._FP8_QUANT_UNAVAILABLE + ) + assert f"fp8_cutlass_with_quant,{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", + key="fp8_trtllm_MxNxK=8x1280x2880", + argv=[], + quant=("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", + ] + + errors = benchmark._parse_driver_errors(output) + results = benchmark._combine({case.tag: case}, [], {}, errors) + csv_path = tmp_path / "combined_results.csv" + benchmark._write_results(csv_path, results, [8], [(1280, 2880)]) + + expected = "ERROR: K must be divisible by 128; got 2880" + assert errors == {case.tag: "K must be divisible by 128; got 2880"} + assert results["gemm"][case.key] == expected + assert results["gemm"]["fp8_trtllm_with_quant_MxNxK=8x1280x2880"] == expected + assert f"fp8_trtllm,{expected}\n" in csv_path.read_text() + assert f"fp8_trtllm_with_quant,{expected}\n" in csv_path.read_text() + + +def test_empty_driver_error_has_no_synthetic_reason(): + tag = "gemm_nvfp4_cutlass_MxNxK=8x2880x1024" + output = [ + f"[ERROR] Error running test: --routine mm_fp4 --case_tag {tag}\n", + "[ERROR] Error:\n", + ] + + assert benchmark._parse_driver_errors(output) == {tag: ""} + + +def test_write_results_preserves_the_original_sectioned_tables(tmp_path): + output = tmp_path / "combined_results.csv" + benchmark._write_results( + output, + { + "gemm": { + "bf16_MxNxK=1x2x3": 1.25, + "fp8_cutlass_MxNxK=8x4x5": 3.5, + }, + "moe": {"fp8_cutlass_moe_with_quant_M=8": 2.5}, + }, + [1, 8], + [(4, 5), (2, 3)], + ["qkv_proj", "in_proj"], + ) + + assert output.read_text() == ( + "GEMM\n" + "M,1,8\n" + "qkv_proj: 4x5\n" + "bf16,,\n" + "fp8_cutlass,,3.500\n" + "in_proj: 2x3\n" + "bf16,1.250,\n" + "fp8_cutlass,,\n" + "\n" + "MoE\n" + "M,1,8\n" + "fp8_cutlass_moe_with_quant,,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_driver", 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"bf16,ERROR: {expected_reason}" in combined + assert "driver.log" in combined + + +def test_run_driver_streams_and_persists_the_driver_output(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" + + returncode, lines = benchmark._run_driver( + benchmarks_dir, tmp_path / "testlist.txt", tmp_path / "out.csv", driver_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 diff --git a/pyproject.toml b/pyproject.toml index ccec4e1b82b..d94479ea7f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -288,7 +288,13 @@ module = ["examples.diffusers.fastgen.preprocess.*"] ignore_errors = true [tool.bandit] -exclude_dirs = [".github/", "examples/", "noxfile.py", "tests/"] +exclude_dirs = [ + ".agents/skills/benchmark-model-kernels/", + ".github/", + "examples/", + "noxfile.py", + "tests/", +] # Do not change `skips`. It should be consistent with NVIDIA's Wheel-CI-CD bandit.yml config. # Use of `# nosec BXXX` requires special approval skips = [