benchmark-model-kernels

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 <commit> when it is needed.

2. TP and EP — accept both directly, or derive them from the intended deployment's GPU model and count and confirm. Defaults: TP=1, EP=1. Derive parallelism from the intended deployment, never from the GPU that runs the microbenchmark. The script validates every sharding rule (divisibility, GQA replication, expert partitioning) and errors loudly.

3. M sweep — balanced default (recommended): 1 8 64 512; decode-focused: 1 4 16 32; throughput-focused: 64 256 1024 4096. M is roughly the tokens scheduled per step (decode: active sequences), not endpoint concurrency. With EP, an MoE row models one rank's share of the global batch: a global batch of B tokens corresponds to the column M = B/EP, so do not compare different EP values at the same M.

4. Shape preview — always run this before anything else; no FlashInfer or GPU needed:

   python "$SKILL_DIR/scripts/benchmark_model.py" <model> \
     --tp <tp> --ep <ep> --ms <m1> <m2> ... --print_only

   Review the printed shapes and the MoE tuple with the user. # unsupported: lines mean the list is partial and the script exits nonzero — handle those via Manual supplements below.

5. FlashInfer checkout and GPU — the full benchmark needs a FlashInfer source checkout containing benchmarks/flashinfer_benchmark.py; the installed wheel alone is not enough. Prefer a clean checkout matching the installed flashinfer version; ask before cloning or installing anything. On the benchmark machine, check nvidia-smi and package versions, verify the target GPU is idle (concurrent work on the same GPU skews timings), and verify CUPTI timing with a tiny bench_gpu_time(..., enable_cupti=True) probe — a warning that falls back to CUDA events is a failure (the cupti-python/nvidia-cuda-cupti packages must match PyTorch's CUDA major). Ask for a GPU index only when several are visible. Pick a fresh workdir and state it.

6. Full benchmark:

   CUDA_VISIBLE_DEVICES=<gpu-index> \
   python "$SKILL_DIR/scripts/benchmark_model.py" <model> \
     --tp <tp> --ep <ep> --ms <m1> <m2> ... \
     --flashinfer_repo <flashinfer-repo> --workdir <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 (as N,K,module-name triples) and all --moe_* arguments; overrides are rejected.
• Row labels are logical shapes; the runner applies vLLM's physical padding. Report both shapes whenever they differ.
• A failed case writes FlashInfer's error message (or a pointer to driver.log) into its combined_results.csv cell, and the command exits nonzero after the table is written. Never present a partial table as a successful benchmark; read driver.log before rerunning anything.

Known backend and driver limits

• mm_fp4 TensorRT-LLM needs N % 128 == 0 (shuffled weight layout): its cell reports no result row, or on stock 0.6.x drivers an empty assertion that fails every mm_fp4 backend for that shape.
• mm_fp8 (trtllm_low_latency) needs K % 128 == 0.
• MoE rows cover FlashInfer's CUTLASS fused MoE, the trtllm-gen NVFP4 and per-tensor FP8 MoE, and the CuteDSL NVFP4 MoE. The CuteDSL MoE row appears only for Swiglu models (the kernel supports nothing else); the cutedsl and trtllm backends require recent GPUs and report per-case errors elsewhere.
• Gated NVFP4 CUTLASS MoE with 2F % 128 != 0 per rank fails — vLLM raises instead of padding — often as a CUDA error: misaligned address. Prefer EP over TP for the experts to keep the per-rank width legal.
• Do not benchmark a padded shape to dodge these limits and call it vLLM-equivalent; report the limit instead.

Manual supplements

For each # unsupported: layout from the preview: inspect that module's forward path and the intended runtime's TP/EP sharding, derive the per-rank shape (never guess sharding from a weight shape alone), and benchmark only the missing shape:

CUDA_VISIBLE_DEVICES=<gpu-index> \
python "$SKILL_DIR/scripts/benchmark_via_builtin.py" \
  --flashinfer_repo <flashinfer-repo> --ms <m1> <m2> ... \
  --nks <n>,<k>,<name> --workdir <workdir>

Label these rows as manual supplements; this is not a substitute for running benchmark_model.py first. Shell-quote user-supplied paths.

Report

Report the command, GPU, versions, TP/EP, M values, shapes (logical and physical where they differ), warnings, and the artifacts: testlist.txt, driver.log (full driver output), builtin_results.csv (milliseconds, success-only), and combined_results.csv (long form, microseconds: columns module_name, M, N, K, backend, with_quant, runtime in GEMM and MoE sections; modules fused into one GEMM are joined with | in module_name, while distinct same-shape modules appear as duplicated rows sharing one measurement; MoE rows keep the H= F= E= top_k= parameter line and leave N/K empty). In quantization-recipe terms: bf16 rows are the unquantized W16A16 baseline, fp8 rows are per-tensor W8A8, and nvfp4 rows are W4A4. Plain quantized rows time the kernel with pre-quantized activations; *_with_quant rows add a separately measured activation-quantization time in the scale-factor layout that backend consumes, except the NVFP4 CUTLASS MoE row, which is a single fused measurement. MoE routing is synthetic: uniform expert distribution everywhere (real skewed routing is slower), and the trtllm-gen rows, which route in-kernel, use a fixed renormalize method regardless of the model's routing scheme. CUTLASS and CuteDSL MoE rows receive precomputed expert indices and exclude routing-selection cost entirely. No row includes the router GEMM itself. These are kernel times only — never describe them as end-to-end latency or throughput; they omit weights, layer frequency, communication, KV cache, and scheduling.