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: