From 4d9cabfd206988b406ac01fc4e988ba9dc76391d Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Wed, 15 Jul 2026 17:34:00 -0700 Subject: [PATCH] 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 | 267 +++++++++ .../agents/openai.yaml | 4 + .../scripts/benchmark_model.py | 307 ++++++++++ .../scripts/benchmark_via_builtin.py | 561 ++++++++++++++++++ .../benchmark-model-kernels/tests/evals.json | 74 +++ .../tests/test_benchmark_model.py | 172 ++++++ .../tests/test_benchmark_via_builtin.py | 218 +++++++ 7 files changed, 1603 insertions(+) create mode 100644 .agents/skills/benchmark-model-kernels/SKILL.md create mode 100644 .agents/skills/benchmark-model-kernels/agents/openai.yaml 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/evals.json 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..9c1c0b55046 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/SKILL.md @@ -0,0 +1,267 @@ +--- +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 do not load checkpoint weights, launch distributed +workers, measure collectives, or measure serving throughput. + +## Interview one decision at a time + +Run a turn-by-turn interview. Ask exactly one unresolved decision per message, +offer two or three concrete options, and wait for the answer before continuing. +Recommend an option only when it is substantively better; otherwise present +the options neutrally. State the default and use it when the user says "default" +or has no preference. Skip decisions already answered or not applicable. Never +present model, parallelism, M, and FlashInfer as one form. + +At the start, tell the user that a local FlashInfer source checkout is required +for the full GPU benchmark, but not for the shape preview. Resolve the checkout +only after the preview succeeds. + +Use this order: + +1. Model +2. TP and EP, supplied directly or derived from a deployment setup +3. M sweep +4. Shape preview +5. FlashInfer checkout and GPU index +6. Full GPU benchmark + +### Ask for the model + +Ask which model to inspect and offer: + +- A local directory containing `config.json`. +- A local `config.json`. +- A Hugging Face ID such as `org/model`. + +These are equivalent ways to identify the model; do not recommend one over +another. There is no model default. After receiving it, inspect the meta model +before asking about TP or EP so later options are model-specific. + +The model script loads configuration, constructs `AutoModelForCausalLM` under +`accelerate.init_empty_weights`, and walks the instantiated Linear modules. +It never calls a checkpoint weight loader. Standard Hub models download only +configuration metadata. + +Do not enable `--trust-remote-code` without explicit approval. Pin +`--revision ` when it is needed. Remote code can perform arbitrary I/O, +so strict no-download behavior is not guaranteed in that mode. + +The supported layout is deliberately small: + +- Fuse `q_proj`, `k_proj`, and `v_proj`. +- Fuse `gate_proj` and `up_proj`. +- Apply common column sharding to those inputs and row sharding to attention + output and MLP down projections. +- Map routed experts to FlashInfer's fused-MoE shape. + +These are exact shapes from the instantiated modules under that declared +layout; they are not proof that a particular serving runtime supports the +model. Embeddings, LM heads, normalization, and routers are outside the +standalone GEMM list. Use `benchmark_via_builtin.py` with manual shapes for a +different layout. + +## Choose TP and EP + +Ask one parallelism question with two paths: + +- If the user knows the deployment parallelism, accept `TP` and `EP` together. +- Otherwise, ask for the intended deployment's GPU model and GPU count, then + propose valid `TP` and `EP` values and ask for confirmation. + +State the defaults: `TP=1`, `EP=1`. Do not infer deployment parallelism from +the GPU selected to execute the microbenchmark. The benchmark runs one +representative rank on one visible GPU; that device may differ from the +intended deployment topology. + +- Dense model: use `EP=1`; choose TP to match the intended deployment. +- MoE: choose EP from divisors of the routed expert count. TP and EP may + overlay the same ranks; do not assume `TP * EP` equals GPU count. +- For MoE with `EP=1`, shard the expert intermediate width by TP. With `EP>1`, + partition whole experts by EP. Use manual shapes for an orthogonal expert-TP + x EP mesh. +- Require every sharded dimension to divide exactly. Allow GQA KV replication + only when TP is a multiple of the KV-head count. +- Require local routed experts to be at least `top_k`. + +For a dense model, fix `EP=1`. When helping choose values, use the intended +deployment's GPU memory and count; the meta model does not allocate weights. + +## Choose M values + +Ask the M question only after TP and EP are settled. Offer: + +- Balanced/default, recommended: `1 8 64 512`. +- Decode-focused: `1 4 16 32`. +- Throughput-focused: `64 256 1024 4096`. + +Accept custom unique positive values. If omitted, use the balanced default. + +Explain M as: + +- Approximately active sequence count during one-token decode. +- Total scheduled tokens during prefill or continuous batching. +- Input token count for the fused-MoE case. + +Do not call M literal endpoint concurrency. + +## Preview, then run the full benchmark + +After M is chosen, show a compact summary of model, TP, EP, M values, and any +known deployment topology. Preview shapes without requiring a FlashInfer +checkout: + +```bash +python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \ + --tp --ep \ + --ms ... \ + --print_only +``` + +Review every N,K shape and the MoE tuple. If the preview succeeds, proceed +directly to resolving FlashInfer and the GPU, then run the full benchmark. Do +not ask a run-mode question or offer short-check-only and preview-only choices. +Only stop after the preview when the user's initial request explicitly limits +the task to shape inspection. + +### Resolve FlashInfer for the full GPU benchmark + +Tell the user: + +> A local FlashInfer source checkout is required for +> `benchmarks/flashinfer_benchmark.py` and its utilities. The installed wheel +> alone is not sufficient. + +Search for existing checkouts and inspect their revisions and working-tree +state. Offer up to three viable choices, recommending a clean checkout that +matches the installed package when available. If the user supplies a path, +verify: + +```bash +test -f /benchmarks/flashinfer_benchmark.py +``` + +If no checkout is available, offer: + +| Choice | Tradeoff | +| --- | --- | +| Revision matching the installed package | Best compatibility | +| Latest release tag | Stable, but may omit new routines | +| Latest `main` | Newest routines, but may require rebuilding | + +Resolve the latest tag at runtime; never hard-code it. Ask for a destination +and approval in separate turns before cloning or installing. Prefer matching +checkout and package revisions. + +If the user asks to upgrade FlashInfer, install the requested version first, +then use the source tag for that exact version. Check dependency constraints: +vLLM or TensorRT-LLM may pin a different FlashInfer release. Report such a +conflict instead of silently changing the requested version. + +Inspect the allocated benchmark environment, not a login host: + +```bash +nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader +python -c "import torch, flashinfer; print(torch.__version__, flashinfer.__version__)" +python -c "import accelerate, transformers; print(accelerate.__version__, transformers.__version__)" +``` + +Verify that CUPTI timing really works with a tiny `bench_gpu_time` probe. Do +not accept a warning that falls back to CUDA events as equivalent. The +`cupti-python` and `nvidia-cuda-cupti` packages must match the CUDA major used +by PyTorch. When device access is sandboxed, request an unsandboxed GPU command; +do not route through a remote debug environment unless the user asked for one. + +If several GPUs are visible, ask for the GPU index in its own turn; otherwise +use the only visible GPU. Choose a fresh work directory and state it, asking +only if the user wants another location. Show a final compact summary including +the FlashInfer revision and output directory, then run on one GPU: + +```bash +CUDA_VISIBLE_DEVICES= \ +python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \ + --tp --ep \ + --ms ... \ + --flashinfer_repo \ + --workdir +``` + +The model script derives `--nks` and `--moe_*`; do not override them. Treat the +short plumbing check as an internal validation step, not a run mode. For the +check, append `--dry_run_iters 1 --num_iters 3 --no_autotune` and do not present +it as a stable performance result. Use the normal defaults for the full +benchmark. Always complete this short check before a full sweep after changing +FlashInfer versions or shape logic. The first fused-MoE build or autotune can +take several minutes; reuse its cache. + +The bundled runner keeps the derived shape as the logical case label and uses +the physical padding of the matching vLLM backend: + +- Dense NVFP4 cuDNN and CUTLASS pad `N` and `K` to multiples of 32. +- Dense NVFP4 TensorRT-LLM and the dense BF16/FP8 cases stay exact. +- Gated per-tensor FP8 FlashInfer CUTLASS MoE pads `F` to a multiple of 16; + non-gated activation uses 128. +- Gated NVFP4 FlashInfer CUTLASS MoE stays exact because vLLM rejects the + padding. Do not apply MXFP4 alignment rules to this NVFP4 case. + +Do not turn scale-factor storage padding into GEMM padding. In particular, +vLLM keeps an already 32-aligned dense `N=2880` while its scale storage rounds +to 2944. FlashInfer 0.6.14's benchmark driver nevertheless prepares the +TensorRT-LLM shuffled tensor for every `mm_fp4` backend, so that utility can +reject `N=2880` even for cuDNN/CUTLASS. Report that driver limitation; do not +benchmark `N=2944` and call it vLLM-equivalent. + +If any launched backend case fails, write its shortened error reason into the +affected `combined_results.csv` cell, finish writing the table, and then exit +nonzero instead of presenting a partial table as a successful benchmark. +Report both logical and physical shapes whenever padding changes a dimension. +Run BF16 and FP8 MoE cases before alignment-sensitive NVFP4 cases because a +CUDA fault can poison the remaining cases in the same FlashInfer process. + +For manual shapes, invoke the sibling directly: + +```bash +CUDA_VISIBLE_DEVICES= \ +python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \ + --flashinfer_repo \ + --ms ... \ + --nks , , ... \ + --workdir +``` + +Require `--nks` and/or all four MoE shape arguments. Shell-quote every +user-supplied path and model reference. + +## Report results + +Report the command, GPU, versions, TP/EP, M values, shapes, warnings, and: + +- `/testlist.txt` +- `/builtin_results.csv` +- `/combined_results.csv` + +FlashInfer's raw `builtin_results.csv` reports milliseconds. +`combined_results.csv` reports microseconds in wide `GEMM` and `MoE` sections. +Each section starts with an `M,,,...` row. GEMM then groups backend rows +under each requested `NxK` label; MoE lists its backend rows directly. Plain +rows are kernel-only. Standalone activation quantization is timed once per M,K +and layout, then reused across backends. Most `_with_quant` rows add that time; +NVFP4 MoE with quantization is a direct fused measurement. Unavailable routines +are skipped with a warning. A launched case that produces no row gets an error +cell in `combined_results.csv` and makes the command fail after the file is +written. FlashInfer's raw `builtin_results.csv` remains success-only. + +Do not describe the result as end-to-end latency or throughput. It omits model +weights, layer frequency, communication, KV cache, scheduling, and server +overhead. diff --git a/.agents/skills/benchmark-model-kernels/agents/openai.yaml b/.agents/skills/benchmark-model-kernels/agents/openai.yaml new file mode 100644 index 00000000000..79200f5be98 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Benchmark Model Kernels" + short_description: "Inspect and benchmark model kernel shapes" + default_prompt: "Use $benchmark-model-kernels to guide me one decision at a time through model inspection, TP/EP, M values, and a FlashInfer benchmark." 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..25c5562fdb4 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py @@ -0,0 +1,307 @@ +# 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, and applies the common serving/export TP layout. It never +calls a checkpoint weight loader. Nonstandard layouts should be supplied +directly to benchmark_via_builtin.py. +""" + +import argparse +import shlex +import subprocess # nosec B404 - launches the trusted sibling with an argv list +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 + + +_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", + "--moe_hidden_size", + "--moe_intermediate_size", + "--moe_num_experts", + "--moe_top_k", +} + + +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] + if leaf not in _PROJECTIONS or ".experts." in name or ".local_experts." in name: + 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 "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 _expert_shape(module: Any) -> tuple[int, int] | None: + gate = getattr(module, "gate_proj", None) + up = getattr(module, "up_proj", None) + down = getattr(module, "down_proj", None) + if gate is None: + gate = getattr(module, "w1", None) + up = getattr(module, "w3", None) + down = getattr(module, "w2", 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: + return None + if gate_shape != up_shape or down_shape != (gate_shape[1], gate_shape[0]): + raise ShapeError("unsupported expert Linear shapes") + return gate_shape[1], gate_shape[0] + + +def _moe_shapes(model: Any, config: Any) -> set[_MoeShape]: + shapes = set() + top_k = getattr(config, "num_experts_per_tok", None) + for name, module in model.named_modules(): + if name.rsplit(".", 1)[-1] in {"experts", "local_experts"}: + 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") + shapes.add(_MoeShape(*shape, len(expert_modules), int(top_k))) + + params = dict(module.named_parameters(recurse=False)) + gate_up, down = params.get("gate_up_proj"), params.get("down_proj") + if gate_up is None or down is None or gate_up.ndim != 3 or down.ndim != 3: + continue + if top_k is None: + raise ShapeError("could not determine MoE top_k") + num_experts = int(gate_up.shape[0]) + hidden = int(config.hidden_size) + intermediate = gate_up.numel() // (2 * num_experts * hidden) + if down.numel() != num_experts * hidden * intermediate: + raise ShapeError(f"unsupported fused experts at {name}") + shapes.add(_MoeShape(hidden, intermediate, num_experts, int(top_k))) + return shapes + + +def _inspect_model( + model: Any, config: Any, tp: int, ep: int +) -> tuple[list[_Kernel], _MoeShape | None]: + config = getattr(config, "text_config", None) or config + kernels = _dense_kernels(model, config, tp) + moe_shapes = _moe_shapes(model, config) + if len(moe_shapes) > 1: + raise ShapeError("model contains multiple routed-expert layouts") + moe = next(iter(moe_shapes), None) + if moe is None: + if ep != 1: + 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) + if not kernels: + raise ShapeError("no dense benchmark shapes found") + return kernels, moe + + +def _command(kernels: list[_Kernel], moe: _MoeShape | None, passthrough: list[str]) -> list[str]: + script = Path(__file__).with_name("benchmark_via_builtin.py") + nks = list(dict.fromkeys((n, k) for n, k, _ in kernels)) + command = [sys.executable, str(script), "--nks", *(f"{n},{k}" for n, k in nks)] + 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), + ] + return command + passthrough + + +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=int, default=1, help="tensor parallel size, e.g. 8") + parser.add_argument("--ep", type=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/--moe_* shapes cannot be overridden") + + try: + config, model = _load_meta_model(args.model, args.trust_remote_code, args.revision) + kernels, moe = _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") + for n, k, label in dict.fromkeys(kernels): + print(f"# {n}x{k} <- {label}") + if moe: + print(f"# MoE: H={moe.hidden} F={moe.intermediate} E={moe.experts} top_k={moe.top_k}") + command = _command(kernels, moe, passthrough) + print(">>> " + shlex.join(command)) + if not args.print_only: + # The trusted sibling is launched with an argv list and no shell. + subprocess.run(command, check=True) # nosec B603 + + +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..efee2a4f964 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py @@ -0,0 +1,561 @@ +# 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 # nosec B404 +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:" + +_ResultValue = float | str + + +@dataclass +class _Case: + section: str + tag: str + key: str + argv: list[str] + quant: tuple[str, int, int] | None = None + + +def _nk_pair(value: str) -> tuple[int, int]: + n, k = value.split(",") + return int(n), int(k) + + +def _round_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _available_routines(benchmarks_dir: Path) -> set[str]: + code = ( + "from routines.flashinfer_benchmark_utils import benchmark_apis\n" + "print('\\n'.join(r for values in benchmark_apis.values() for r in values))" + ) + # This probes the user-selected FlashInfer checkout with a fixed code string. + result = subprocess.run( # nosec B603 + [sys.executable, "-c", code], + cwd=benchmarks_dir, + capture_output=True, + text=True, + timeout=300, + check=False, + ) + if result.returncode: + raise RuntimeError(f"could not query FlashInfer routines:\n{result.stderr.strip()}") + return set(result.stdout.split()) + + +def _short_reason(message: str) -> str: + first = next((line.strip() for line in message.splitlines() if line.strip()), "") + return first.replace(",", ";")[:100] or "failed (no error message)" + + +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() + errors[pending_tag] = _short_reason(message) + pending_tag = None + return errors + + +def _missing_case_reason(case: _Case, returncode: int) -> str: + if returncode: + return f"FlashInfer driver exited with status {returncode}" + + def value(option: str, default: str = "0") -> str: + try: + return case.argv[case.argv.index(option) + 1] + except (ValueError, IndexError): + return default + + routine = value("--routine", "") + n = int(value("--n")) + k = int(value("--k")) + if routine == "mm_fp4" and n % 128: + return f"mm_fp4 benchmark TRT-LLM shuffle requires n % 128 == 0; got n = {n}" + if routine == "mm_fp8" and k % 128: + return f"mm_fp8 requires k % 128 == 0; got k = {k}" + if routine == "mm_fp8" and n % 32: + return f"mm_fp8 requires n % 32 == 0; got n = {n}" + return "FlashInfer produced no result row" + + +def _run_driver(benchmarks_dir: Path, testlist: Path, output: Path) -> tuple[int, list[str]]: + process = subprocess.Popen( # nosec B603 + [ + 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 = [] + for line in process.stdout: + print(line, end="", flush=True) + lines.append(line) + return process.wait(), lines + + +def _gemm_cases( + ms: list[int], + nks: list[tuple[int, int]], + common: list[str], + available: set[str], +) -> tuple[list[_Case], set[str]]: + cases = [] + missing: set[str] = set() + + 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: + if routine not in available: + missing.add(routine) + continue + 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, missing + + +def _moe_cases( + ms: list[int], + shape: tuple[int, int, int, int] | None, + common: list[str], + available: set[str], + activation: str | None, + routing: str | None, +) -> tuple[list[_Case], set[str]]: + if shape is None: + return [], set() + routine = "cutlass_fused_moe" + if routine not in available: + return [], {routine} + + 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] + + cases = [] + gated = activation is None or activation in {"Swiglu", "Geglu", "SwigluBias"} + fp8_intermediate = _round_up(intermediate, 16 if gated else 128) + variants = ( + ("bf16_cutlass_moe", intermediate, []), + ("fp8_cutlass_moe", fp8_intermediate, ["--cutlass_variant", "fp8"]), + ( + "nvfp4_cutlass_moe", + intermediate, + ["--cutlass_variant", "nvfp4", "--quantized_input"], + ), + ("nvfp4_cutlass_moe_with_quant", 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, set() + + +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], float]: + results = {} + specs = {case.quant for case in cases if case.quant is not None} + for kind, m, k in sorted(specs): + tensor = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + if kind.startswith("nvfp4_"): + runner = _nvfp4_runner(tensor, kind.removeprefix("nvfp4_")) + elif vllm_ops is not None: + runner = _fp8_runner(tensor) + else: + continue + 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], float], + 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) + results[case.section][f"{name}_with_quant{separator}{shape}"] = ( + value + quant[case.quant] + ) + 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]], +) -> 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 n, k in nks: + writer.writerow([f"{n}x{k}"]) + 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=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("--dry_run_iters", type=int, help="warmup iterations, e.g. 5") + parser.add_argument("--num_iters", type=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=int, help="model hidden size, e.g. 4096") + parser.add_argument("--moe_intermediate_size", type=int, help="expert width, e.g. 14336") + parser.add_argument("--moe_num_experts", type=int, help="local expert count, e.g. 8") + parser.add_argument("--moe_top_k", type=int, help="experts per token, e.g. 2") + parser.add_argument("--moe_activation_type", help="FlashInfer activation, e.g. Swiglu") + parser.add_argument("--moe_routing_method", help="FlashInfer routing, e.g. renormalize") + parser.add_argument("--workdir", type=Path, default=Path("benchmark_via_builtin_out")) + return parser + + +def main() -> None: + """Validate inputs, run the FlashInfer driver, and combine its results.""" + parser = _parser() + args = parser.parse_args() + ms = list(dict.fromkeys(args.ms)) + nks = list(dict.fromkeys(args.nks or [])) + 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") + + benchmarks_dir = args.flashinfer_repo / "benchmarks" + driver = benchmarks_dir / "flashinfer_benchmark.py" + if not driver.is_file(): + parser.error(f"{driver} does not exist") + try: + available = _available_routines(benchmarks_dir) + except RuntimeError as exc: + parser.error(str(exc)) + + 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, missing_gemm = _gemm_cases(ms, nks, common, available) + moe_shape = moe_values if all(moe_values) else None + moe_cases, missing_moe = _moe_cases( + ms, + moe_shape, + common, + available, + args.moe_activation_type, + args.moe_routing_method, + ) + cases = gemm_cases + moe_cases + missing = missing_gemm | missing_moe + if missing: + print(f"[WARN] unavailable routines skipped: {', '.join(sorted(missing))}") + if not cases: + parser.error("none of the requested routines is available") + + 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" + if builtin_csv.exists() or combined_csv.exists(): + parser.error(f"{args.workdir} already contains results; choose a fresh --workdir") + cases_by_tag = {case.tag: case for case in 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) + + 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) + errors = {} + for case in cases: + if case.tag not in completed_tags: + reason = parsed_errors.get(case.tag, "") + if reason == "failed (no error message)": + reason = "" + errors[case.tag] = reason or _missing_case_reason(case, returncode) + + 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) + 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/evals.json b/.agents/skills/benchmark-model-kernels/tests/evals.json new file mode 100644 index 00000000000..af7a7708246 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/evals.json @@ -0,0 +1,74 @@ +[ + { + "name": "flashinfer-checkout-required", + "skills": ["benchmark-model-kernels"], + "query": "Benchmark Qwen/Qwen3-8B. I installed FlashInfer from pip but do not have its source repo.", + "files": [], + "expected_behavior": [ + "Explains why a local FlashInfer source checkout is required", + "Offers a matching revision, latest release, and latest main with tradeoffs", + "Asks for a destination and approval before cloning or installing", + "Allows model-only print preview before the checkout is ready" + ] + }, + { + "name": "dense-model-parallelism", + "skills": ["benchmark-model-kernels"], + "query": "Help me benchmark ./checkpoint on four H100s. I am not sure what TP or batch sizes to use.", + "files": [], + "expected_behavior": [ + "Validates the local config and constructs the Transformers model on meta tensors", + "Explains the common fused QKV and gate/up layout", + "Recommends EP=1 and validates TP against actual layer dimensions", + "States TP=1, EP=1, and M values 1 8 64 512 as defaults", + "Explains that M is only approximately decode concurrency", + "Runs a print-only preview before proposing the GPU command" + ] + }, + { + "name": "moe-topology", + "skills": ["benchmark-model-kernels"], + "query": "Plan a kernel benchmark for a local Mixtral config on eight GPUs with top-k 2.", + "files": [], + "expected_behavior": [ + "Inspects actual expert modules or fused expert tensors", + "Chooses EP from divisors of the routed expert count", + "Explains that TP and EP may overlay the same ranks", + "Explains that EP=1 shards expert width by TP while EP greater than one partitions experts", + "Checks that local expert count is at least top-k" + ] + }, + { + "name": "remote-code-opt-in", + "skills": ["benchmark-model-kernels"], + "query": "The model needs trust_remote_code=True. Benchmark it without downloading weights.", + "files": [], + "expected_behavior": [ + "Stops for informed approval before enabling remote code", + "Recommends pinning a commit revision or audited local code", + "Explains that the wrapper does not load weights but remote code can perform arbitrary I/O" + ] + }, + { + "name": "manual-layout", + "skills": ["benchmark-model-kernels"], + "query": "My runtime uses a custom fused projection layout that is not QKV plus gate/up.", + "files": [], + "expected_behavior": [ + "Does not guess the custom sharding or fusion", + "Requests exact per-rank N,K and MoE shapes", + "Uses benchmark_via_builtin.py directly" + ] + }, + { + "name": "serving-boundary", + "skills": ["benchmark-model-kernels"], + "query": "Measure request throughput and p99 latency for my vLLM endpoint at concurrency 128.", + "files": [], + "expected_behavior": [ + "Routes to the deployment benchmarking workflow", + "Does not translate endpoint concurrency into M values", + "Explains that these scripts measure kernels rather than end-to-end serving" + ] + } +] 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..8069b8b9d58 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py @@ -0,0 +1,172 @@ +# 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 + +import pytest + +pytest.importorskip("accelerate") +transformers = pytest.importorskip("transformers") + +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 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, _ = benchmark_model._inspect_model(model, config, tp=4, ep=1) + + assert (24, 32, "fused_qkv") in kernels + + +def test_meta_loader_never_materializes_model_tensors(tmp_path): + model_dir = _save(tmp_path, _llama_config()) + + _, model = benchmark_model._load_meta_model(str(model_dir / "config.json"), False, None) + + tensors = list(model.named_parameters()) + list(model.named_buffers()) + assert tensors and all(tensor.is_meta for _, tensor in tensors) + + +def test_revision_does_not_reach_registered_model_constructor(tmp_path): + model_dir = _save(tmp_path, _llama_config()) + + _, model = benchmark_model._load_meta_model(str(model_dir), False, "main") + + assert type(model).__name__ == "LlamaForCausalLM" + + +def test_mixtral_modulelist_experts_use_ep(tmp_path): + config = transformers.MixtralConfig( + vocab_size=128, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + ) + model_dir = _save(tmp_path, config) + _, model = benchmark_model._load_meta_model(str(model_dir), False, None) + + _, moe = benchmark_model._inspect_model(model, config, tp=2, ep=2) + + assert moe == benchmark_model._MoeShape(32, 48, 2, 2) + + +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) + + +def test_moe_tp_without_ep_shards_expert_intermediate_size(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, + ) + 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=1) + + assert moe == benchmark_model._MoeShape(32, 24, 4, 2) + + +def test_derived_shapes_cannot_be_overridden(monkeypatch, capsys): + monkeypatch.setattr( + sys, + "argv", + [str(SCRIPT), "unused/model", "--nks", "1,1", "--print_only"], + ) + + with pytest.raises(SystemExit, match="2"): + benchmark_model.main() + assert "cannot be overridden" 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..7f367b58ab8 --- /dev/null +++ b/.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py @@ -0,0 +1,218 @@ +# 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 + +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"]) +def test_nk_pair_rejects_invalid_values(value): + with pytest.raises(ValueError): + benchmark._nk_pair(value) + + +def test_gemm_cases_are_data_driven_and_preserve_requested_shapes(): + available = {"mm_bf16", "mm_fp4", "bmm_fp8", "mm_fp8"} + + cases, missing = benchmark._gemm_cases([1], [(65, 129)], [], available) + + assert not missing + 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_moe_cases_and_result_combination(): + cases, missing = benchmark._moe_cases( + [1], (32, 50, 4, 2), [], {"cutlass_fused_moe"}, None, None + ) + assert not missing + 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_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_uses_the_known_mm_fp4_shuffle_constraint(): + case = benchmark._Case( + section="gemm", + tag="gemm_nvfp4_cutlass_MxNxK=8x2880x1024", + key="nvfp4_cutlass_MxNxK=8x2880x1024", + argv=[ + "--routine", + "mm_fp4", + "--backends", + "trtllm", + "--use_nvfp4", + "--m", + "8", + "--n", + "2880", + "--k", + "1024", + ], + ) + + assert benchmark._missing_case_reason(case, 0) == ( + "mm_fp4 benchmark TRT-LLM shuffle requires n % 128 == 0; got n = 2880" + ) + + +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)], + ) + + assert output.read_text() == ( + "============================================================\n" + "GEMM\n" + "============================================================\n" + "M,1,8\n" + "4x5\n" + "bf16,,\n" + "fp8_cutlass,,3.500\n" + "2x3\n" + "bf16,1.250,\n" + "fp8_cutlass,,\n" + "\n" + "============================================================\n" + "MoE\n" + "============================================================\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