add scripts and a skill to use flashinfer to do layerwise benchmark#1980
Open
sychen52 wants to merge 1 commit into
Open
add scripts and a skill to use flashinfer to do layerwise benchmark#1980sychen52 wants to merge 1 commit into
sychen52 wants to merge 1 commit into
Conversation
with different backends. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Contributor
📝 WalkthroughWalkthroughAdds the ChangesModel Kernel Benchmarking
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Operator
participant benchmark_model.py
participant benchmark_via_builtin.py
participant flashinfer_benchmark.py
participant combined_results.csv
Operator->>benchmark_model.py: Select model, TP/EP, and M values
benchmark_model.py->>benchmark_model.py: Inspect meta-tensor module shapes
benchmark_model.py->>benchmark_via_builtin.py: Provide derived GEMM and MoE shapes
benchmark_via_builtin.py->>flashinfer_benchmark.py: Run benchmark case matrix
flashinfer_benchmark.py-->>benchmark_via_builtin.py: Return kernel timings and errors
benchmark_via_builtin.py->>combined_results.csv: Write merged benchmark results
Caution Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional.
❌ Failed checks (1 error)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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
Comment on lines
+274
to
+286
| 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) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Use positive-integer argument validation across both benchmark entry points.
Invalid numeric inputs currently reach division, shape generation, driver execution, or CUDA allocation.
.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py#L274-L286: reject non-positive--tpand--ep..agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py#L463-L476: reject non-positive M/N/K, MoE shape, and iteration values.
📍 Affects 2 files
.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py#L274-L286(this comment).agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py#L463-L476
🤖 Prompt for 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.
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_model.py around
lines 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.
Comment on lines
+329
to
+337
| 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 |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
sed -n '250,380p' .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py && printf '\n---\n' && rg -n "_with_quant|vllm_ops|nvfp4_|quantization timing|unavailable|skip" .agents/skills/benchmark-model-kernels -SRepository: NVIDIA/Model-Optimizer
Length of output: 10454
🏁 Script executed:
sed -n '1,80p' .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py && printf '\n---\n' && sed -n '480,540p' .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py && printf '\n---\n' && sed -n '1,240p' .agents/skills/benchmark-model-kernels/tests/test_benchmark_via_builtin.pyRepository: NVIDIA/Model-Optimizer
Length of output: 12030
🏁 Script executed:
sed -n '380,470p' .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.pyRepository: NVIDIA/Model-Optimizer
Length of output: 3667
Emit an explicit FP8 quantization fallback when vLLM is unavailable
_quant_times() currently skips FP8 quant timing when vllm_ops is missing, and _write_results() then omits the corresponding _with_quant row entirely. Add an explicit warning or error cell so the missing quant result is visible instead of silently disappearing.
🤖 Prompt for 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.
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py
around lines 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.
Comment on lines
+399
to
+426
| 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)]) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Restore the section separators required by the CSV contract test.
test_write_results_preserves_the_original_sectioned_tables expects separator rows around both sections, while this writer emits only GEMM and MoE. The test therefore fails deterministically.
Proposed fix
gemm = results["gemm"]
if gemm:
+ writer.writerow(["=" * 60])
writer.writerow(["GEMM"])
+ writer.writerow(["=" * 60])
writer.writerow(["M", *ms])
...
if gemm:
writer.writerow([])
+ writer.writerow(["=" * 60])
writer.writerow(["MoE"])
+ writer.writerow(["=" * 60])
writer.writerow(["M", *ms])📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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)]) | |
| with path.open("w", newline="") as stream: | |
| writer = csv.writer(stream, lineterminator="\n") | |
| gemm = results["gemm"] | |
| if gemm: | |
| writer.writerow(["=" * 60]) | |
| writer.writerow(["GEMM"]) | |
| writer.writerow(["=" * 60]) | |
| 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(["=" * 60]) | |
| writer.writerow(["MoE"]) | |
| writer.writerow(["=" * 60]) | |
| 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)]) |
🤖 Prompt for 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.
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py
around lines 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.
Comment on lines
+525
to
+537
| 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) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Handle a driver failure that produces no builtin_results.csv.
If the child exits before creating the output, Line 527 raises FileNotFoundError; consequently, combined_results.csv is never written with per-case errors as promised. Treat a missing CSV as zero completed rows and continue through failure aggregation.
Proposed fix
- with builtin_csv.open(newline="") as stream:
- rows = list(csv.DictReader(stream))
+ if builtin_csv.is_file():
+ with builtin_csv.open(newline="") as stream:
+ rows = list(csv.DictReader(stream))
+ else:
+ rows = []📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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) | |
| returncode, driver_output = _run_driver(benchmarks_dir, testlist, builtin_csv) | |
| if builtin_csv.is_file(): | |
| with builtin_csv.open(newline="") as stream: | |
| rows = list(csv.DictReader(stream)) | |
| else: | |
| rows = [] | |
| 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) |
🤖 Prompt for 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.
In @.agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py
around lines 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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-builtinTesting
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
New Features
Tests