Scripts and a skill to do per-layer benchmark using flashinfer - #1980
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds an operator-facing kernel benchmarking skill, meta-tensor model inspection for dense, Mamba, and MoE layouts, FlashInfer case execution, quantization timing, CSV reporting, and comprehensive validation tests. ChangesModel Kernel Benchmarking
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Operator
participant benchmark_model.py
participant Transformers
participant benchmark_via_builtin.py
participant flashinfer_benchmark.py
participant combined_results.csv
Operator->>benchmark_model.py: Provide model and TP/EP settings
benchmark_model.py->>Transformers: Construct meta-tensor model
Transformers-->>benchmark_model.py: Return module and config shapes
benchmark_model.py->>benchmark_via_builtin.py: Pass derived GEMM, MoE, and routing arguments
benchmark_via_builtin.py->>flashinfer_benchmark.py: Run generated benchmark cases
flashinfer_benchmark.py-->>benchmark_via_builtin.py: Return timings and errors
benchmark_via_builtin.py->>combined_results.csv: Write merged timing and error rows
Suggested reviewers: 🔕 Pre-merge checks override appliedThe pre-merge checks have been overridden successfully. You can now proceed with the merge. Overridden by ❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py:
- Around line 274-286: The argument parsers in benchmark_model.py (around parser
construction before _load_meta_model) and benchmark_via_builtin.py (around its
main parser) must validate all numeric benchmark inputs as positive integers.
Add shared positive-integer validation for --tp and --ep in benchmark_model.py,
and for M/N/K, MoE shape, and iteration arguments in benchmark_via_builtin.py,
so invalid non-positive values are rejected during parsing before model
inspection, shape generation, execution, or CUDA allocation.
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py:
- Around line 399-426: Update the CSV-writing flow in the results writer to emit
the separator rows required by the sectioned-table contract around both the GEMM
and MoE sections. Preserve the existing headers, data rows, and conditional
section ordering while restoring separators even when both sections are present.
- Around line 329-337: Update the FP8 branch in _quant_times() so that when
vllm_ops is unavailable it records an explicit warning or error result for the
affected quantization case instead of continuing silently. Ensure
_write_results() preserves and emits that fallback cell as the corresponding
_with_quant row, while leaving available vLLM and NVFP4 execution unchanged.
- Around line 525-537: Update the result-loading block around _run_driver and
builtin_csv so a missing builtin_results.csv is treated as an empty row set
instead of raising FileNotFoundError. Preserve normal CSV parsing when the file
exists, then continue completed_tags and per-case error aggregation so
combined_results.csv is still produced with missing-case reasons.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8a542345-e207-4bbb-9999-64efe1192a93
📒 Files selected for processing (7)
.agents/skills/benchmark-model-kernels/SKILL.md.agents/skills/benchmark-model-kernels/agents/openai.yaml.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py.agents/skills/benchmark-model-kernels/tests/evals.json.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.py
4d9cabf to
9a63807
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1980 +/- ##
==========================================
- Coverage 78.47% 76.86% -1.61%
==========================================
Files 518 519 +1
Lines 58574 60355 +1781
==========================================
+ Hits 45967 46393 +426
- Misses 12607 13962 +1355
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py:
- Around line 350-374: Update the per-spec quantization timing loop in
_quant_times to catch CUDA allocation, CUPTI, runner, and kernel timing
exceptions without aborting the overall benchmark. For each failed (kind, m, k)
specification, store an explicit error result using the existing error-value
convention like _FP8_QUANT_UNAVAILABLE, emit an appropriate warning if needed,
and continue so completed driver results can still be written to
combined_results.csv.
In `@pyproject.toml`:
- Around line 291-297: Remove the .agents/skills/benchmark-model-kernels/ entry
from Bandit’s exclude_dirs configuration so executable benchmark scripts remain
covered. Add a narrow suppression only at the specific reviewed false-positive
call site within those scripts, without excluding the broader skill directory.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: db8b82d3-7d5e-4dbd-8b91-1a668cabacc9
📒 Files selected for processing (8)
.agents/skills/benchmark-model-kernels/SKILL.md.agents/skills/benchmark-model-kernels/agents/openai.yaml.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py.agents/skills/benchmark-model-kernels/tests/evals.json.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.pypyproject.toml
🚧 Files skipped from review as they are similar to previous changes (3)
- .agents/skills/benchmark-model-kernels/agents/openai.yaml
- .agents/skills/benchmark-model-kernels/tests/evals.json
- .agents/skills/benchmark-model-kernels/SKILL.md
9a63807 to
2cd8630
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py:
- Around line 33-36: Update the import and subprocess setup around argument
parsing, `run_benchmark`, and the child execution flow so `--flashinfer_repo`
controls both `flashinfer` and `flashinfer.testing` resolution. Add the selected
checkout root to `sys.path` before importing those modules, including when the
child runs from `benchmarks_dir`, and ensure imports cannot silently use an
unrelated installed FlashInfer.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9c1b8ca4-8092-4081-b771-83772f0e2180
📒 Files selected for processing (6)
.agents/skills/benchmark-model-kernels/SKILL.md.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py.agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.pypyproject.toml
🚧 Files skipped from review as they are similar to previous changes (3)
- pyproject.toml
- .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py
- .agents/skills/benchmark-model-kernels/tests/test_benchmark_model.py
sychen52
changed the title
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-4-8) — DM the bot to share feedback.
New self-contained skill (benchmark-model-kernels) with two scripts that (1) derive per-rank GEMM/MoE benchmark shapes from a HF config on meta tensors and (2) drive FlashInfer's built-in benchmark for BF16/FP8/NVFP4. +2301/-1, 3 non-test files + 2 test files + a bandit-exclude tweak.
Design review (complexity gate fired): The problem — a per-layer microbenchmark harness that derives per-rank shapes from model configs without loading weights — has no existing in-repo equivalent (searched; other skills are task-specific and use tests.json eval metadata, not benchmarking). This wraps FlashInfer's flashinfer_benchmark.py rather than reimplementing it; the value-add is the config→shape derivation. OmegaConf/Hydra/Jinja aren't relevant. The pytest-in-skill layout is already an established convention (pyproject has a .agents/skills/*/tests/test_*.py per-file-ignore). Design looks reasonable and the PR body gives concrete usage, though the rationale is brief.
Licensing: clean — new files carry the canonical LICENSE_HEADER NVIDIA header verbatim; no LICENSE/vendoring changes. The bandit exclude_dirs addition is narrowly scoped to the new skill (justified by the intentional subprocess use, which is invoked without a shell and gates trust_remote_code behind an explicit flag).
Correctness: I traced the shape-derivation logic (fused QKV incl. GQA replication, gate/up, Mamba-2, ModuleList vs stacked/direct experts, EP-vs-TP intermediate sharding, routing derivation, expert audit, override rejection) and it holds up; the unit tests are genuinely meaningful and internally consistent across Llama/Mixtral/GptOss/NemotronH plus edge cases.
Why nudge (not approve): (1) sizable new subsystem worth an owner's architectural eyes; (2) coverage gap — benchmark_via_builtin.py's GPU/FlashInfer runtime paths (_run_driver real invocation, _quant_times/_nvfp4_runner/_fp8_runner CUDA timing) are mocked out and never exercised in CI, consistent with the repo's known GPU-only coverage risk; (3) CHANGELOG left unmarked — acceptable for an internal agent skill (no shipped library API change), but worth confirming. No prompt-injection concerns in the untrusted content (SKILL.md's "do not override" prose targets the skill's agent user, not the reviewer).
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Posting two inline comments from my review: one on CI test discovery, one on TP×EP handling for the MoE shape.
| pytest.importorskip("torch") | ||
| pytest.importorskip("accelerate") | ||
| transformers = pytest.importorskip("transformers") |
There was a problem hiding this comment.
These tests never run in CI as written. The skills job in .github/workflows/unit_tests.yml installs only pytest (it's documented as "stdlib-only and hermetic"), so these importorskip calls on torch/accelerate/transformers skip this entire file — and pytest.importorskip("flashinfer") skips all of test_benchmark_via_builtin.py the same way. CI stays green whether or not the code works, so the ~870 lines of tests only ever run on a dev machine.
Suggestions:
- Route this file into the main
linuxunit lane (it needs only torch/accelerate/transformers — all standard unit-test deps, no GPU). - In
benchmark_via_builtin.py, deferimport flashinferinto the functions that need it (same pattern as the existingvllmtry/except). The majority of that test file (_parser,_combine,_write_results,_parse_driver_errors, case generation) is pure CPU logic that could then run without the GPU stack.
There was a problem hiding this comment.
import moved. However, route test to main unit lane involve workflow change.
|
Could you update the PR description to include an example prompt and expected output? |
| 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( # nosec B603 |
There was a problem hiding this comment.
@kevalmorabia97 @shengliangxu here is the subprocess invoke. You can see that coderabbit is paused because of this.
There was a problem hiding this comment.
got it, just exclude checking on these lines should be the way to go.
juhi10071998
left a comment
juhi10071998
left a comment
There was a problem hiding this comment.
Thanks for adding this Shiyang!
juhi10071998
left a comment
juhi10071998
left a comment
There was a problem hiding this comment.
Two more integration asks inline, for feeding this into AutoQuantize.
3c43bf5 to
c6816b5
Compare
c6816b5 to
c3c672e
Compare
juhi10071998
left a comment
juhi10071998
left a comment
There was a problem hiding this comment.
Thanks Shiyang, left some more minor comments
c3c672e to
ee968ad
Compare
with different backends. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
5b4a2d7 to
8f518d2
Compare
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
8f518d2 to
3e632d2
Compare
7 participants
add scripts and a skill to use flashinfer to do layerwise benchmark with different backends.
What does this PR do?
Type of change: ? new skill and scripts
Usage
python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \ openai/gpt-oss-120b \ --tp 4 --ep 4 --ms 8 16 \ --flashinfer_repo $HOME/flashinfer \ --workdir /tmp/gpt-oss-benchmarkor
python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \ --flashinfer_repo $HOME/flashinfer \ --ms 8 16 \ --nks 1280,2880 2880,1024 \ --moe_hidden_size 2880 \ --moe_intermediate_size 2880 \ --moe_num_experts 32 \ --moe_top_k 4 \ --workdir /tmp/gpt-oss-via-builtinOr example prompt to trigger the skill:
example output (TP1, EP4):
Testing
run locally
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information
Summary by CodeRabbit