chore: Refactor benchmark imports to be lazy-loaded

[bkryu]

📌 Description

No changes to library code. Will not launch internal CI

#2354 added moe_comm microbenchmarks to the benchmark harness. However, the PR causes failures in non-MPI enabled environments because flashinfer_benchmark.py blindly imports routines for all modules.

Current PR refactors flashinfer_benchmark.py to use lazy imports for routine modules instead of importing all modules at startup. This avoids loading unnecessary dependencies when running specific benchmarks. Users without MPI installed can now run all non-MPI benchmarks without import errors.

🔍 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

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

• Chores
  • Optimized benchmark execution for improved efficiency with deferred loading of benchmark routines, reducing unnecessary resource consumption during startup.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[gemini-code-assist]

Summary of Changes

Hello @bkryu, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical issue in the flashinfer_benchmark.py script where it would fail in environments lacking MPI due to its eager loading of all benchmark routine modules. By implementing lazy loading for these modules, the script now only imports necessary dependencies when a specific benchmark routine is called. This enhancement significantly improves the script's robustness and usability across diverse environments by preventing unnecessary dependency-related import errors, particularly for users without MPI installations.

Highlights

• Lazy Imports Implementation: Refactored flashinfer_benchmark.py to utilize lazy loading for all routine modules, preventing unnecessary dependency imports at startup.
• MPI Environment Compatibility: Resolved import errors in non-MPI enabled environments by ensuring moe_comm and its mpi4py dependency are only loaded when the moe_comm benchmark is explicitly selected.
• Improved Benchmark Usability: Enabled users without MPI installed to successfully run all non-MPI benchmarks without encountering import-related failures.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[coderabbitai]

📝 Walkthrough

Walkthrough

The change refactors import statements in the benchmark module, converting top-level imports of routine modules into lazy, in-function imports. This defers loading of heavy dependencies (gemm, moe, moe_comm, norm, quantization) until their specific routines are executed, maintaining the same routing logic and outcomes.

Changes

Cohort / File(s)	Summary
Lazy import refactoring benchmarks/flashinfer_benchmark.py	Removed top-level imports of parse/run functions for routine modules (gemm, moe, moe_comm, norm, quantization) and moved them into in-function imports within both run_test and parse_args branches to defer dependency loading until specific routines are requested

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• chore: Add TRTLLM MoE A2A benchmark #2354: Main PR's change directly extends the same routing/parse logic in benchmarks/flashinfer_benchmark.py with moe_comm imports and dispatch modifications.

Suggested reviewers

• Anerudhan
• yzh119
• kahyunnam
• jiahanc
• cyx-6

Poem

🐰 Imports deferred, like carrots in a warren,
Load them later, not at boot—lazy laden!
Heavy modules stay tucked away,
Until their day comes to play,
Faster starts with wisdom in tow! 🚀

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: refactoring benchmark imports to be lazy-loaded instead of eager imports at startup.
Description check	✅ Passed	The description is comprehensive and covers the required template sections with clear context about the motivation, changes, and impact of the PR.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request effectively refactors the benchmark imports to be lazy-loaded, addressing the issue of unnecessary dependencies (like mpi4py) causing failures in non-MPI environments. By moving the routine-specific imports into the if/elif blocks within run_test and parse_args functions, the code now only loads modules when they are actually required. This significantly improves the robustness and flexibility of the benchmark harness, allowing users to run non-MPI benchmarks without encountering import errors. The changes are well-implemented and align perfectly with the stated objective.

[yzh119]

Recent public CI failures might be fixed when #2386 was checked in.

[yongwww]

#2386 was merged, please rebase the pr on top of latest main.

[yzh119]

The changes in this file are not tested on CI so I don't think rebase will bring any benefits. Let's bypass and merge directly.