[UX] Only perform FlashInfer autotuning if needed by kernels

[mgoin]

Purpose

Before this PR, we would always run a forward pass with max_num_batched_tokens simply if FlashInfer was installed and we were on Hopper or above. This is obviously too loose and affects startup time.

Introduces a register_flashinfer_kernel_autotune function that we can call when we know that a FlashInfer kernel that benefits from autotuning will be used at runtime. Then we can query flashinfer_autotune_needed to better gate the autotuning warmup pass only when it is needed.

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to conditionally run FlashInfer autotuning only when needed, which improves startup time. The approach is to register the usage of autotune-able kernels during model initialization and then check a flag during the warmup phase. The implementation is straightforward and correct across the modified files. I have one suggestion regarding thread safety in the new registration logic to prevent potential race conditions in concurrent environments.

[chatgpt-codex-connector]

Codex Review: Here are some suggestions.

Reply with @codex fix comments to fix any unresolved comments.

About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you open a pull request for review, mark a draft as ready, or comment "@codex review". If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex fix this CI failure" or "@codex address that feedback".

[nvpohanh]

cc @nvjullin

[nvjullin]

I'm somewhat concerned about the brittleness of the design. There are multiple modes of failure

1. A new flashinfer op X is integrated but didn't register autotune
2. A flashinfer op Y that didn't used to support autotuning now supports autotuning

Initial benchmarks might look correct because some other op Z happens to be ran together and registers autotune, but X and Y will have bad perf by themselves. The point being, since autotuning is global, complex logic to enable/disable it is very difficult to reason about. To exacerbate the problem, the bad perf is very hard to notice and root cause because it is not a hard error.

On the other hand, autotuning is an extra _dummy_run. If no autotuning ops are involved, it is just a forward pass. Flashinfer will not do any extra work for ops that doesn't require autotuning.
If there are numbers to suggest this is a large overhead, we can ask the FlashInfer team to provide an API to query whether autotuning is needed.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mgoin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!