From 0c40600cb5e59b9fba2c8f562225ef267ae52a0b Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Wed, 15 Jul 2026 17:34:00 -0700 Subject: [PATCH 1/4] add scripts and a skill to use flashinfer to do layerwise benchmark with different backends. Signed-off-by: Shiyang Chen --- .../skills/benchmark-model-kernels/SKILL.md | 134 ++++ .../scripts/benchmark_model.py | 643 +++++++++++++++++ .../scripts/benchmark_via_builtin.py | 650 ++++++++++++++++++ .../tests/test_benchmark_model.py | 495 +++++++++++++ .../tests/test_benchmark_via_builtin.py | 372 ++++++++++ pyproject.toml | 8 +- 6 files changed, 2301 insertions(+), 1 deletion(-) create mode 100644 .agents/skills/benchmark-model-kernels/SKILL.md create mode 100644 .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py create mode 100644 .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py create mode 100644 .agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py create mode 100644 .agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py diff --git a/.agents/skills/benchmark-model-kernels/SKILL.md b/.agents/skills/benchmark-model-kernels/SKILL.md new file mode 100644 index 00000000000..1750147aedf --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/SKILL.md @@ -0,0 +1,134 @@ +--- +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. +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..e2ca4ab3282 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py @@ -0,0 +1,643 @@ +# 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: + 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 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..e14bd9aed8d --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py @@ -0,0 +1,650 @@ +# 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. +""" + +import argparse +import csv +import shlex +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +import flashinfer +import numpy as np +import torch +from flashinfer.testing import bench_gpu_time + +try: + from vllm import _custom_ops as vllm_ops +except ImportError: + vllm_ops = None + +_FP8_MAX = torch.finfo(torch.float8_e4m3fn).max +_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): + 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): + scale = tensor.abs().max().float() / _FP8_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 + 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..4011978ab52 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py @@ -0,0 +1,495 @@ +# 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_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..547a1a7c704 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py @@ -0,0 +1,372 @@ +# 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 + +pytest.importorskip("flashinfer") + +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 = [ From 326fcc38502073750c674cf52c9f946387f0da81 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Fri, 17 Jul 2026 18:38:51 -0700 Subject: [PATCH 2/4] address review comments Signed-off-by: Shiyang Chen --- .../skills/benchmark-model-kernels/SKILL.md | 4 +- .../scripts/benchmark_model.py | 14 +++++++ .../scripts/benchmark_via_builtin.py | 22 ++++++++--- .../tests/test_benchmark_model.py | 38 +++++++++++++++++++ .../tests/test_benchmark_via_builtin.py | 2 - 5 files changed, 71 insertions(+), 9 deletions(-) diff --git a/.agents/skills/benchmark-model-kernels/SKILL.md b/.agents/skills/benchmark-model-kernels/SKILL.md index 1750147aedf..02a32797cb2 100644 --- a/.agents/skills/benchmark-model-kernels/SKILL.md +++ b/.agents/skills/benchmark-model-kernels/SKILL.md @@ -34,7 +34,9 @@ order: 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. + 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: diff --git a/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py index e2ca4ab3282..98cf92d5e51 100644 --- a/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py @@ -516,6 +516,13 @@ def _inspect_model( 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: @@ -617,6 +624,13 @@ def main() -> None: 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}" diff --git a/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py index e14bd9aed8d..8840a48cf86 100644 --- a/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py @@ -22,6 +22,8 @@ benchmark driver and utilities. """ +from __future__ import annotations + import argparse import csv import shlex @@ -29,18 +31,16 @@ import sys from dataclasses import dataclass from pathlib import Path +from typing import TYPE_CHECKING -import flashinfer -import numpy as np -import torch -from flashinfer.testing import bench_gpu_time +if TYPE_CHECKING: + import torch try: from vllm import _custom_ops as vllm_ops except ImportError: vllm_ops = None -_FP8_MAX = torch.finfo(torch.float8_e4m3fn).max _ERROR_CASE_PREFIX = "[ERROR] Error running test:" _ERROR_MESSAGE_PREFIX = "[ERROR] Error:" _FP8_QUANT_UNAVAILABLE = "ERROR: vLLM is unavailable for FP8 activation quantization" @@ -319,6 +319,8 @@ def _moe_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 @@ -331,7 +333,9 @@ def kernel(value, scale): def _fp8_runner(tensor: torch.Tensor): - scale = tensor.abs().max().float() / _FP8_MAX + 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) @@ -353,6 +357,12 @@ def _quant_times( 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_")) diff --git a/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py index 4011978ab52..b26b3f99382 100644 --- a/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py @@ -407,6 +407,44 @@ def test_command_names_gemm_shapes_and_merges_duplicates(): ] +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 = [] 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 index 547a1a7c704..6652282afeb 100644 --- a/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py @@ -20,8 +20,6 @@ import pytest -pytest.importorskip("flashinfer") - 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 From 4a387899585210d14f7a8f89a5feff68055de20d Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 20 Jul 2026 10:15:51 -0700 Subject: [PATCH 3/4] revert to nosec Signed-off-by: Shiyang Chen --- .../scripts/benchmark_via_builtin.py | 4 ++-- pyproject.toml | 8 +------- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py index 8840a48cf86..832c89d2ef4 100644 --- a/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py @@ -27,7 +27,7 @@ import argparse import csv import shlex -import subprocess +import subprocess # nosec B404 import sys from dataclasses import dataclass from pathlib import Path @@ -145,7 +145,7 @@ 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( + process = subprocess.Popen( # nosec B603 [ sys.executable, "flashinfer_benchmark.py", diff --git a/pyproject.toml b/pyproject.toml index d94479ea7f4..ccec4e1b82b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -288,13 +288,7 @@ module = ["examples.diffusers.fastgen.preprocess.*"] ignore_errors = true [tool.bandit] -exclude_dirs = [ - ".agents/skills/benchmark-model-kernels/", - ".github/", - "examples/", - "noxfile.py", - "tests/", -] +exclude_dirs = [".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 = [ From 8f518d2782d88cda558ac418ef98ab42e573a038 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Wed, 22 Jul 2026 16:23:15 -0700 Subject: [PATCH 4/4] more backends and address comments Signed-off-by: Shiyang Chen --- .../skills/benchmark-model-kernels/SKILL.md | 55 +- .../scripts/benchmark_model.py | 269 +++-- .../scripts/benchmark_via_builtin.py | 927 +++++++++++------- .../tests/test_benchmark_model.py | 221 +++-- .../tests/test_benchmark_via_builtin.py | 337 ++++--- 5 files changed, 1147 insertions(+), 662 deletions(-) diff --git a/.agents/skills/benchmark-model-kernels/SKILL.md b/.agents/skills/benchmark-model-kernels/SKILL.md index 02a32797cb2..1ff4f5af571 100644 --- a/.agents/skills/benchmark-model-kernels/SKILL.md +++ b/.agents/skills/benchmark-model-kernels/SKILL.md @@ -45,15 +45,16 @@ order: --tp --ep --ms ... --print_only ``` - Review the printed shapes, MoE tuple, and derived routing with the user. + 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, and - verify CUPTI timing with a tiny `bench_gpu_time(..., enable_cupti=True)` + 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 @@ -78,13 +79,9 @@ order: 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. + 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 @@ -92,11 +89,14 @@ Do not restate, re-derive, or override these — the code enforces them: ## 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_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. @@ -117,20 +117,29 @@ python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \ --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. +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 +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 index 98cf92d5e51..2cd7aa0019d 100644 --- a/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py @@ -16,9 +16,9 @@ """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. +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. @@ -26,6 +26,7 @@ import argparse import importlib.util +import re import shlex import sys from dataclasses import dataclass @@ -53,15 +54,6 @@ class _ExpertShape: 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 = { @@ -82,14 +74,19 @@ class _MoeRouting: "--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"} +# 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: @@ -121,7 +118,15 @@ def _load_meta_model(model_ref: str, trust_remote_code: bool, revision: str | No 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) + 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 @@ -144,6 +149,12 @@ def _divide(value: int, size: int, label: str) -> int: 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], @@ -164,7 +175,8 @@ def _fused_qkv( 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" + 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]: @@ -172,11 +184,15 @@ def _dense_kernels(model: Any, config: Any, tp: int) -> list[_Kernel]: 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", "routers"} for part in parts[:-1]) + or any(part in _ROUTER_PATH_PARTS for part in parts[:-1]) ): continue shape = _linear_shape(module) @@ -205,26 +221,33 @@ def _dense_kernels(model: Any, config: Any, tp: int) -> list[_Kernel]: ) if "o_proj" in layers: n, k = layers["o_proj"] - kernels.append((n, _divide(k, tp, f"{parent}.o_proj"), "attention_out")) + 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}") - n = _divide(gate_n + up_n, tp, f"{parent}.gate_up") - kernels.append((n, gate_k, "fused_gate_up")) + # 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, "up")) + 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, "fused_gate_up")) + 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"), "down")) + kernels.append( + (n, _divide(k, tp, f"{parent}.down_proj"), _path_label(parent, "down_proj")) + ) return list(dict.fromkeys(kernels)) @@ -268,8 +291,80 @@ def _mamba_kernels(model: Any, tp: int) -> list[_Kernel]: local_in = 2 * local_intermediate + 2 * local_groups * state + local_heads kernels.extend( [ - (local_in, hidden, "mamba_in"), - (hidden, local_intermediate, "mamba_out"), + (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)) @@ -330,9 +425,16 @@ def _moe_activation(config: Any, gated: bool) -> str | None: if gated: if normalized in {"silu", "swiglu", "swish"}: return "Swiglu" - if normalized == "geglu" or normalized.startswith("gelu") or normalized == "quick_gelu": + # 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}") + 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", @@ -346,8 +448,28 @@ def _moe_activation(config: Any, gated: bool) -> str | None: 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: - for attr in ("num_experts_per_tok", "moe_top_k"): + # 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) @@ -420,21 +542,6 @@ def _moe_shapes(model: Any, config: Any) -> set[_MoeShape]: 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 @@ -452,6 +559,19 @@ def _unsupported_decoder_linears( 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(): @@ -468,11 +588,11 @@ def _unsupported_decoder_linears( leaf = parts[-1] if not in_decoder: continue - if any(part in {"router", "routers"} for part in parts): + 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 _ROUTER_NAMES or name in mamba_projections: + 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()] @@ -480,9 +600,11 @@ def _unsupported_decoder_linears( def _inspect_model( model: Any, config: Any, tp: int, ep: int -) -> tuple[list[_Kernel], _MoeShape | None, _MoeRouting | None, list[str]]: +) -> 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) + kernels = ( + _dense_kernels(model, config, tp) + _mamba_kernels(model, tp) + _gdn_kernels(model, tp) + ) problems = [] try: moe_shapes = _moe_shapes(model, config) @@ -511,7 +633,6 @@ def _inspect_model( # 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") @@ -530,31 +651,26 @@ def _inspect_model( 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 + return kernels, moe, 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 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", @@ -568,16 +684,6 @@ def _command( ] 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 @@ -607,7 +713,7 @@ def main() -> None: 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) + kernels, moe, problems = _inspect_model(model, config, args.tp, args.ep) except ShapeError as exc: parser.error(str(exc)) @@ -615,7 +721,10 @@ def main() -> None: 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") + 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: @@ -631,14 +740,6 @@ def main() -> None: ) 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: @@ -646,7 +747,7 @@ def main() -> None: "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) + 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: diff --git a/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py index 832c89d2ef4..b2eed9daef0 100644 --- a/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py @@ -16,16 +16,18 @@ """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. +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 @@ -34,6 +36,8 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: + from typing import TextIO + import torch try: @@ -55,27 +59,88 @@ "SwigluStep", "Identity", ) -_MOE_ROUTINGS = ("renormalize", "deepseek_v3", "llama4", "renormalize_naive", "topk") - _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 - key: str + backend: str + m: int + n: int | None + k: int | None 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 + 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: @@ -101,59 +166,72 @@ def _nk_pair(value: str) -> tuple[int, int]: def _named_nks( nks: list[tuple[int, int]], names: list[str] | None -) -> tuple[list[tuple[int, int]], 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") - - unique_nks = list(dict.fromkeys(nks)) if names is None: - return unique_nks, None + return {(n, k): [f"{n}x{k}"] for n, k in dict.fromkeys(nks)} - names_by_nk: dict[tuple[int, int], list[str]] = {} + labels_by_nk: dict[tuple[int, int], list[str]] = {} for nk, name in zip(nks, names, strict=True): - labels = names_by_nk.setdefault(nk, []) + labels = labels_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] + return labels_by_nk 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 +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): - 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. + 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", - "--testlist", - str(testlist.resolve()), - "--output_path", - str(output.resolve()), - ], + [sys.executable, "flashinfer_benchmark.py", *argv], cwd=benchmarks_dir, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, @@ -162,159 +240,295 @@ def _run_driver( ) 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) + 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 = [] + 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: - 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" + # 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 backend != "trtllm": + if driver_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, + add( + m, n, k, - ["--batch_size", "1"], - "fp8_static", + f"nvfp4_{row_suffix}", + "mm_fp4", + driver_backend, + run_n, + run_k, + extra, + _QuantSpec("nvfp4", layout, m, run_k), ) - 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, - ) + 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: 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 { +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", } - 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"]), + # 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"), ) - 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, - ) - ) + 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 @@ -322,6 +536,13 @@ 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 ) @@ -344,124 +565,133 @@ def kernel(value, value_scale): return kernel, (tensor, scale) -def _quant_times( +def _attach_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 +) -> 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 - 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 + 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 _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 "" +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, - results: dict[str, dict[str, _ResultValue]], - ms: list[int], - nks: list[tuple[int, int]], - nk_names: list[str] | None = None, + 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") - gemm = results["gemm"] + if header: + writer.writerow([header]) 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"] + 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"]) - 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)]) + 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: @@ -502,27 +732,56 @@ def _parser() -> argparse.ArgumentParser: 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 _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: - nks, nk_names = _named_nks(args.nks or [], args.nk_names) + labels_by_nk = _named_nks(args.nks or [], args.nk_names) except ValueError as exc: parser.error(str(exc)) moe_values = ( @@ -533,30 +792,19 @@ def main(argv: list[str] | None = None) -> None: ) 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): + 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 all(moe_values) and args.moe_top_k > args.moe_num_experts: + if moe_shape is not None and moe_shape.top_k > moe_shape.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" @@ -573,29 +821,9 @@ def main(argv: list[str] | None = None) -> None: 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 + 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" @@ -604,56 +832,31 @@ def main(argv: list[str] | None = None) -> None: 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, + 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, ) - results = _combine(cases_by_tag, rows, quant, errors) - _write_results(combined_csv, results, ms, nks, nk_names) + 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}") - if errors: - failed = [cases_by_tag[tag].key for tag in errors if tag in cases_by_tag] + 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 returncode: - raise RuntimeError(f"FlashInfer driver exited with status {returncode}") if __name__ == "__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 index b26b3f99382..02c7f80b159 100644 --- a/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py @@ -94,10 +94,14 @@ def test_llama_meta_walk_fuses_common_projections(tmp_path, monkeypatch, capsys) 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 ( + "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. @@ -106,9 +110,13 @@ def test_gqa_kv_heads_are_replicated_when_tp_exceeds_kv_heads(tmp_path): 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) + kernels, _, problems = benchmark_model._inspect_model(model, config, tp=4, ep=1) - assert (24, 32, "fused_qkv") in kernels + 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 == [] @@ -144,10 +152,9 @@ def test_mixtral_modulelist_experts_use_ep(tmp_path): 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) + _, moe, _ = 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): @@ -166,7 +173,7 @@ def test_gpt_oss_direct_expert_tensors_are_inspected(tmp_path): 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) + _, moe, _ = benchmark_model._inspect_model(model, config, tp=2, ep=2) assert moe == benchmark_model._MoeShape(32, 48, 2, 2, "Swiglu") @@ -177,34 +184,88 @@ def test_nemotron_h_mamba_and_stacked_experts_are_inspected(tmp_path): _, 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) + 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, "mamba_in") in kernels - assert (32, 16, "mamba_out") in kernels - assert (20, 32, "up") in kernels - assert (32, 20, "down") in kernels + 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 == [] - # 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, []) + command = benchmark_model._command(kernels, moe, []) 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) + _, _, 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) @@ -213,11 +274,10 @@ def test_expert_audit_problem_is_not_masked_by_per_rank_validation(tmp_path): # 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) + kernels, moe, 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) @@ -242,12 +302,12 @@ def __init__(self): mlp_hidden_act="relu2", ) - kernels, moe, routing, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) + 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, routing, []) + command = benchmark_model._command(kernels, moe, []) assert "--nks" not in command assert command[:2] == ["--moe_hidden_size", "32"] @@ -268,13 +328,54 @@ def __init__(self): } +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" ) - with pytest.raises(benchmark_model.ShapeError, match="unsupported gated MoE activation"): - benchmark_model._moe_activation(SimpleNamespace(hidden_act="relu"), True) + 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(): @@ -293,8 +394,8 @@ def __init__(self): model.mixer = Mixer() assert benchmark_model._mamba_kernels(model, tp=2) == [ - (42, 32, "mamba_in"), - (32, 16, "mamba_out"), + (42, 32, "mixer.in_proj"), + (32, 16, "mixer.out_proj"), ] @@ -313,9 +414,9 @@ def __init__(self): ("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) + kernels, _, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) - assert (48, 32, "up") in kernels + assert (48, 32, "layers.*.up_proj") in kernels assert problems == ["unsupported decoder Linear GEMM layout(s): layers.0.unknown_proj (48x32)"] @@ -337,7 +438,7 @@ def __init__(self): benchmark_model.main() captured = capsys.readouterr() - assert "# 48x32 <- up" in captured.out + 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 @@ -364,7 +465,7 @@ def __init__(self): num_experts_per_tok=2, ) - _, moe, _, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) + _, moe, problems = benchmark_model._inspect_model(model, config, tp=1, ep=1) assert moe is None assert problems == [ @@ -372,38 +473,24 @@ def __init__(self): ] -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(): +def test_command_keeps_same_shape_kernels_as_separate_named_pairs(): kernels = [ - (64, 32, "fused_qkv"), - (64, 32, "fused_gate_up"), - (32, 64, "down"), + (64, 32, "a_proj|b_proj"), + (64, 32, "c_proj"), + (32, 64, "down_proj"), ] - command = benchmark_model._command(kernels, None, None, []) + command = benchmark_model._command(kernels, None, []) assert command == [ "--nks", "64,32", + "64,32", "32,64", "--nk_names", - "fused_qkv/fused_gate_up", - "down", + "a_proj|b_proj", + "c_proj", + "down_proj", ] @@ -464,10 +551,10 @@ def test_runner_is_invoked_in_process(tmp_path, monkeypatch): "128,32", "32,64", "--nk_names", - "fused_qkv", - "attention_out", - "fused_gate_up", - "down", + "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", @@ -497,12 +584,14 @@ def __init__(self): 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) + kernels, moe, 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")] + assert kernels == [ + (48, 32, "layers.*.mlp.up_proj"), + (32, 48, "layers.*.mlp.down_proj"), + ] @pytest.mark.parametrize( 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 index 6652282afeb..e558c478233 100644 --- a/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py @@ -61,19 +61,20 @@ def test_parser_rejects_zero_for_numeric_options(option, capsys): 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} == { + 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.key.split("_MxNxK=", 1)[1] for case in cases} == {"1x65x129"} + assert {(case.m, case.n, case.k) for case in cases} == {(1, 65, 129)} physical_shapes = { - case.key.split("_MxNxK=", 1)[0]: ( + case.backend: ( case.argv[case.argv.index("--n") + 1], case.argv[case.argv.index("--k") + 1], ) @@ -81,6 +82,7 @@ def test_gemm_cases_are_data_driven_and_preserve_requested_shapes(): } 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") @@ -88,203 +90,231 @@ def test_gemm_cases_are_data_driven_and_preserve_requested_shapes(): 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), + 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_merge_duplicate_shapes(): - nks, names = benchmark._named_nks( +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 nks == [(32, 64), (128, 256)] - assert names == ["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_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", +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.key.split("_M=", 1)[0]: case.argv[case.argv.index("--intermediate_size") + 1] + (case.backend, case.with_quant): 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", + ("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", } - - fp8 = next(case for case in cases if case.key == "fp8_cutlass_moe_M=1") + 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 - rows = [{"case_tag": fp8.tag, "median_time": "0.001"}] - quant = {fp8.quant: 2.0} - - results = benchmark._combine({fp8.tag: fp8}, rows, quant) + fp8.result = 1.0 + fp8.quant_result = 2.0 - assert results["moe"]["fp8_cutlass_moe_M=1"] == 1.0 - assert results["moe"]["fp8_cutlass_moe_with_quant_M=1"] == 3.0 + assert benchmark._output_rows(fp8) == [(False, 1.0), (True, 3.0)] def test_swiglustep_uses_gated_fp8_alignment(): - cases = benchmark._moe_cases([1], (32, 50, 4, 2), [], "SwigluStep", "topk") + cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2, "SwigluStep"), []) - fp8 = next(case for case in cases if case.key == "fp8_cutlass_moe_M=1") + 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], (32, 50, 4, 2), [], "Relu2", "topk") + cases = benchmark._moe_cases([1], benchmark._MoeShape(32, 50, 4, 2, "Relu2"), []) intermediate_sizes = { - case.key.split("_M=", 1)[0]: case.argv[case.argv.index("--intermediate_size") + 1] + (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_moe": "50", - "fp8_cutlass_moe": "128", - "nvfp4_cutlass_moe": "128", - "nvfp4_cutlass_moe_with_quant": "128", + ("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_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", + backend="fp8_cutlass", + m=1, + n=32, + k=64, argv=[], - quant=("fp8_static", 1, 64), + quant=benchmark._QuantSpec("fp8", "static", 1, 64), + result=1.0, ) 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 - ) + benchmark._attach_quant_times([case], 1, 1, False) output = tmp_path / "combined_results.csv" - benchmark._write_results(output, results, [1], [(32, 64)]) + benchmark._write_results(output, [case], {(32, 64): ["32x64"]}) - assert quant == {case.quant: benchmark._FP8_QUANT_UNAVAILABLE} + assert case.quant_result == 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 "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() ) - 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", + backend="fp8_trtllm", + m=8, + n=1280, + k=2880, argv=[], - quant=("fp8_static", 8, 2880), + 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", ] - errors = benchmark._parse_driver_errors(output) - results = benchmark._combine({case.tag: case}, [], {}, errors) + 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, results, [8], [(1280, 2880)]) + benchmark._write_results(csv_path, [case], {(1280, 2880): ["1280x2880"]}) 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() + 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(): - tag = "gemm_nvfp4_cutlass_MxNxK=8x2880x1024" output = [ - f"[ERROR] Error running test: --routine mm_fp4 --case_tag {tag}\n", + "[ERROR] Error running test: --routine mm_fp4 --case_tag gemm_nvfp4\n", "[ERROR] Error:\n", ] - assert benchmark._parse_driver_errors(output) == {tag: ""} - - -def test_write_results_preserves_the_original_sectioned_tables(tmp_path): + 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, - { - "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"], + 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" - "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" + "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" - "M,1,8\n" - "fp8_cutlass_moe_with_quant,,2.500\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" ) @@ -326,7 +356,7 @@ def test_missing_builtin_results_still_writes_combined_errors( 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(benchmark, "_run_case", lambda *_: (returncode, [])) monkeypatch.setattr( sys, "argv", @@ -348,11 +378,49 @@ def test_missing_builtin_results_still_writes_combined_errors( assert not (workdir / "builtin_results.csv").exists() combined = (workdir / "combined_results.csv").read_text() - assert f"bf16,ERROR: {expected_reason}" in combined + 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_driver_streams_and_persists_the_driver_output(tmp_path, capsys): +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( @@ -360,11 +428,26 @@ def test_run_driver_streams_and_persists_the_driver_output(tmp_path, capsys): ) driver_log = tmp_path / "driver.log" - returncode, lines = benchmark._run_driver( - benchmarks_dir, tmp_path / "testlist.txt", tmp_path / "out.csv", 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" + )