bench: Unify PDL behavior, add missing norm routines, and misc improvements - #3435
Conversation
📝 WalkthroughWalkthroughThis PR extends the benchmark infrastructure with Programmatic Dependent Launch (PDL) support and introduces three new RMSNorm benchmark variants. It adds a ChangesPDL Support and RMSNorm Benchmark Enhancements
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds benchmarks for several new normalization routines, including fused_add_rmsnorm, gemma_rmsnorm, and gemma_fused_add_rmsnorm, defaulting their backend to cute-dsl. It also introduces a global --enable_pdl flag with a warning mechanism for unsupported routines, programmatically filters backends for bmm_fp8 via trial calls, and pre-allocates scale tensors on-device to ensure compatibility with CUDA graph capture. The reviewer feedback suggests using a list comprehension instead of list.remove() for filtering backends in gemm.py, and recommends supporting 3D inputs (via num_heads) in both testFusedAddRmsnorm and testGemmaFusedAddRmsnorm to maintain consistency with other RMSNorm benchmarks.
| for backend in backends_to_remove: | ||
| backends.remove(backend) |
There was a problem hiding this comment.
Instead of iterating and calling .remove() on the list, which can be inefficient and prone to issues if there are duplicates, you can use a list comprehension to filter backends in a more Pythonic and efficient way.
backends = [b for b in backends if b not in backends_to_remove]There was a problem hiding this comment.
Current changes match the the pattern used by testMmFp4 and other sibling test functions in this file — backends comes from a fixed argparse choices= set so duplicates aren't possible, and the list is ~5 items. Leaving as-is for consistency.
| f"Unsupported input dtype: {args.input_dtype}. Supported dtypes are bfloat16, float16." | ||
| ) | ||
|
|
||
| input_shape = (batch_size, hidden_size) |
There was a problem hiding this comment.
To make fused_add_rmsnorm consistent with rmsnorm and gemma_rmsnorm, we should support 3D inputs by extracting num_heads from args and setting the input_shape accordingly.
| input_shape = (batch_size, hidden_size) | |
| num_heads = args.num_heads | |
| if num_heads is not None: | |
| input_shape = (batch_size, num_heads, hidden_size) | |
| else: | |
| input_shape = (batch_size, hidden_size) |
There was a problem hiding this comment.
Skipping this — flashinfer.fused_add_rmsnorm and flashinfer.gemma_fused_add_rmsnorm are 2D-only per their docstrings and impl (no if input.dim(== 3: branch like rmsnorm has, no QK fused-add variant in flashinfer/norm/). Exposing --num_heads here would let users construct 3D inputs that the underlying API doesn't actually support. If we want a 3D fused-add benchmark, the API needs to grow that support first.
| f"Unsupported input dtype: {args.input_dtype}. Supported dtypes are bfloat16, float16." | ||
| ) | ||
|
|
||
| input_shape = (batch_size, hidden_size) |
There was a problem hiding this comment.
To make gemma_fused_add_rmsnorm consistent with rmsnorm and gemma_rmsnorm, we should support 3D inputs by extracting num_heads from args and setting the input_shape accordingly.
| input_shape = (batch_size, hidden_size) | |
| num_heads = args.num_heads | |
| if num_heads is not None: | |
| input_shape = (batch_size, num_heads, hidden_size) | |
| else: | |
| input_shape = (batch_size, hidden_size) |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
benchmarks/routines/allreduce_comm.py (1)
291-299:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHonor
--enable_pdlin allreduce benchmark execution.Line 461 adds a PDL-support warning, but both benchmark paths still hardcode
launch_with_pdl=True, so the CLI flag is ignored.Suggested fix
def run_allreduce(inp): allreduce_fusion( input=inp, workspace=workspace, pattern=pattern_code, - launch_with_pdl=True, + launch_with_pdl=args.enable_pdl, output=output, use_oneshot=use_oneshot, ) return output @@ def run_allreduce(inp): allreduce_fusion( input=inp, workspace=workspace, pattern=pattern_code, - launch_with_pdl=True, + launch_with_pdl=args.enable_pdl, residual_out=residual_out, norm_out=norm_out, residual_in=residual, rms_gamma=norm_weight,Also applies to: 309-320, 461-461
🤖 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 `@benchmarks/routines/allreduce_comm.py` around lines 291 - 299, The calls to allreduce_fusion currently hardcode launch_with_pdl=True, ignoring the CLI flag; update each call site (e.g., inside the run_allreduce wrapper where allreduce_fusion is invoked and the other similar invocation blocks) to pass the actual enable_pdl boolean flag (or the variable that holds the CLI option) as launch_with_pdl instead of True so the benchmark honors --enable_pdl; locate and change all occurrences (the run_allreduce function and the other allreduce_fusion call sites noted in the diff) to use launch_with_pdl=enable_pdl.
🤖 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.
Outside diff comments:
In `@benchmarks/routines/allreduce_comm.py`:
- Around line 291-299: The calls to allreduce_fusion currently hardcode
launch_with_pdl=True, ignoring the CLI flag; update each call site (e.g., inside
the run_allreduce wrapper where allreduce_fusion is invoked and the other
similar invocation blocks) to pass the actual enable_pdl boolean flag (or the
variable that holds the CLI option) as launch_with_pdl instead of True so the
benchmark honors --enable_pdl; locate and change all occurrences (the
run_allreduce function and the other allreduce_fusion call sites noted in the
diff) to use launch_with_pdl=enable_pdl.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5ecc8c26-bd2d-4cd3-9e3b-e69cbd9bea37
📒 Files selected for processing (15)
benchmarks/README.mdbenchmarks/flashinfer_benchmark.pybenchmarks/routines/allreduce_comm.pybenchmarks/routines/attention.pybenchmarks/routines/flashinfer_benchmark_utils.pybenchmarks/routines/gemm.pybenchmarks/routines/mamba.pybenchmarks/routines/mixed_comm.pybenchmarks/routines/moe.pybenchmarks/routines/moe_comm.pybenchmarks/routines/norm.pybenchmarks/routines/quantization.pybenchmarks/routines/rope.pybenchmarks/routines/sampling.pybenchmarks/samples/sample_testlist.txt
3 participants
📌 Description
Reshapes
flashinfer_benchmark.pyto (a) make--enable_pdlbehave uniformly across every routine, (b) add three previously-missing norm benchmarks, and (c) fix a few backend-label / API-routing bugs surfaced along the way.--enable_pdluniform semanticsA single rule across all 50+ routines:
--enable_pdlomitted → PDL off;--enable_pdlset → PDL on; routines whose API can't accept PDL emit a one-shot[WARNING]and ignore the flag.--enable_pdlto the shared parser inflashinfer_benchmark.py; removed three duplicate declarations inparse_norm_args/parse_gemm_args/parse_quantization_args.warn_if_pdl_unsupported(args, routine)helper inflashinfer_benchmark_utils.py.enable_pdl=None→ auto-enable on H100/B200, but benchmark wasn't passing it): all 4 attention wrappers (decode / prefill paged / prefill ragged / MLA) + 4 directtrtllm/MLA call sites,
mm_fp4(API default wasTrue!),softmax, plus 4 MoE routines (trtllm_fp4_block_scale_moe,trtllm_fp8_per_tensor_scale_moe,cutlass_fused_moe,cute_dsl_fp4_block_scale_moe).enable_pdl=True:moe.py(trtllm_fp8_block_scale_moe) andmoe_comm.py(trtllm_fp8_block_scale_routed_moein a2a dispatch).enable_pdl=Falseblocking opt-in: twotrtllm_ragged_attention_deepseekcall sites inattention.py(cute-dsl and trtllm-native branches).mla_rope_quantize_fp8,rope_quantize_fp8,rope_quantize_fp8_append_paged_kv_cache).warn_if_pdl_unsupportedcalls to routines whose APIs don't acceptenable_pdl: 6 gemm, 5 rope, 14 sampling, 1 quantization (nvfp4_batched_quantize), 1 mamba, 1 moe (b12x_fused_moe),allreduce_fusion, mixed_comm.
--enable_pdl).New norm benchmarks
Three rmsnorm-family routines added to
benchmarks/routines/norm.py, ported from a side branch and stripped of unrelated tooling:fused_add_rmsnormgemma_rmsnormgemma_fused_add_rmsnormRegistered in
flashinfer_benchmark_utils.py(benchmark_apis,routine_cc_to_supported_backends), documented inbenchmarks/README.md(bullets + support matrix), and exercised by newsamples/sample_testlist.txtentries.Backend label corrections (cuda → cute-dsl)
Six rmsnorm-family routines (
rmsnorm,rmsnorm_quant,fused_add_rmsnorm,fused_add_rmsnorm_quant,gemma_rmsnorm,gemma_fused_add_rmsnorm) actually route to CuTe-DSL kernels by default in flashinfer(env-var fallback to CUDA JIT via
FLASHINFER_USE_CUDA_NORM=1), but the benchmark was labeling them ascuda. Updated the cc-support map andrun_backenddispatch to usecute-dsl, with acuda → cute-dslauto-remap matching the existing FP4 routine pattern so sample lines don't need explicit
--backends. README support matrix updated.rmsnorm_quantCUDA-graph capture fixtestRmsnormQuant/testFusedAddRmsnormQuantwere constructingscaleas a Pythonfloat, which forced flashinfer's_normalize_scale_tensorto allocate a host tensor on every call — illegal underCUDA-graph capture (and explicitly deprecated by a
FutureWarning). Now pre-allocatescaleas a shape-(1,)GPUfloat32tensor once in the test function; CSVscalecolumn unchanged thanks to.item()on emit.
testBmmFp8refactor (programmatic backend probing)bmm_fp8test has been checking hard-coded backend support, leading to gaps (e.g. cutlass is supported on SM120 but was not runnable) Aligned withtestMmFp4: removed the hard-codedroutine_cc_to_supported_backends["bmm_fp8"]table and theif backend in [...]whitelist insiderun_backend. The test now probes each backend with a trialflashinfer.gemm.bmm_fp8(...)call and relies on the@backend_requirementdecorator to raiseBackendSupportedErrorfor unsupported (routine, backend, cc) combos. Future SM/backend additions no longer needto touch the benchmark utils.
Misc
samples/sample_testlist.txt: added 9 new norm-variant sample lines (3 routines × 3 configs each); removed thermsnorm_quant_fp8_e5m2line (cute-dsl path is missing the dtype-string entry upstream —separate bug, not introduced here).
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
--enable_pdlflag to enable Programmatic Dependent Launch optimization in benchmarks.Documentation