add scripts and a skill to use flashinfer to do layerwise benchmark

.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 <commit>` 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 <model> \
+  --tp <tp> --ep <ep> \
+  --ms <m1> <m2> ... \
+  --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 <flashinfer-repo>/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=<gpu-index> \
+python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py <model> \
+  --tp <tp> --ep <ep> \
+  --ms <m1> <m2> ... \
+  --flashinfer_repo <flashinfer-repo> \
+  --workdir <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=<gpu-index> \
+python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \
+  --flashinfer_repo <flashinfer-repo> \
+  --ms <m1> <m2> ... \
+  --nks <n1>,<k1> <n2>,<k2> ... \
+  --workdir <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:
+
+- `<workdir>/testlist.txt`
+- `<workdir>/builtin_results.csv`
+- `<workdir>/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,<m1>,<m2>,...` 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.

.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."