bench(moe_deepseek): scope autotune(True) to pre-warm only

[leejnau]

📌 Description

Scopes autotune(True) to wrap only the per-backend pre-warm pass instead of the entire benchmark loop, so the measurement iterations themselves run with is_tuning_mode=False.

Verified on B200 in NGC rc14 across all four measurement modes (graphs × CUPTI, on × off) at --num-tokens 128 --ep 8. Default mode (graphs + CUPTI) is byte-identical: CuteDSL 0.147 ms / TRTLLM 0.143–0.144 ms on both branches, matching the #3292 baseline. The CUDA-events fallback path (--no-cuda-graph --no-cupti) drops from CuteDSL 1.389 ms → 0.197 ms (~7×); at N=1 EP=16, CUTLASS drops from 2.162 ms → 0.051 ms (42×). Graphs-on / CUPTI-off and graphs-off / CUPTI-on modes track accordingly — no regression in any graphs-enabled path, substantial improvement in every graphs-disabled path.

Also updates the now-stale run_benchmark docstring.

🔍 Related Issues

#3292
#3126

🚀 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

Release Notes

• Refactor
  • Modified benchmark autotuning to apply during the pre-warm phase only, rather than across the entire token sweep.
  • Added configurable parameter to benchmark functions for controlling autotuning behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 7b91dbe0-e034-4dd4-959e-503a8f9ca58d

📥 Commits

Reviewing files that changed from the base of the PR and between 6a09cc0 and 4e471fd.

📒 Files selected for processing (1)

• benchmarks/bench_moe_deepseek.py

---

📝 Walkthrough

Walkthrough

The benchmark refactors autotuning scope from a global wrapper around the token-count measurement loop to localized pre-warm phases within each backend. Three backend functions and their orchestration layer receive a do_autotune parameter to control this conditional behavior.

Changes

Autotune Scope Relocation

Layer / File(s)	Summary
Backend autotune parameter and pre-warm scoping benchmarks/bench_moe_deepseek.py	bench_cute_dsl, bench_cutlass, and bench_trtllm each add a do_autotune parameter, introduce local imports for contextlib and autotune, and wrap their pre-warm invocation with a conditional context manager that applies autotune(True) only when enabled.
Orchestration refactoring to remove global autotune wrapping benchmarks/bench_moe_deepseek.py	run_benchmark documentation and token-count loop are updated to remove the outer autotune(True) wrapper, while _benchmark_single gains a do_autotune parameter that is forwarded to all three backend bench_* calls, threading autotune control through the call hierarchy.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• flashinfer-ai/flashinfer#2072: Both PRs modify autotune usage and propagate an autotune flag into MoE execution paths by switching from always-on autotune(True) to conditional autotune during warmup and threading the flag through calls.
• flashinfer-ai/flashinfer#3292: Both PRs modify autotune handling and the warmup/benchmark flow in benchmarks/bench_moe_deepseek.py, touching the same functions and changing where autotune is applied.

Suggested reviewers

• yzh119
• sricketts
• bkryu
• nv-yunzheq
• yyihuang
• aleozlx

Poem

🐰 Autotuning now blooms in each backend's morning glow,
No longer wrapped around the token stream's flow,
A parameter threads through, so simple, so neat,
From orchestration's hands to pre-warm's greeting sweet! ✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title precisely and concisely describes the main change: scoping autotune(True) to pre-warm only instead of the entire benchmark loop.
Description check	✅ Passed	The description includes detailed technical context, verification results, related issues, and completed pre-commit and test checklist items matching the required template.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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 refactors the autotuning logic in the DeepSeek-V3 MoE benchmarks by moving the autotune context from the main benchmark loop to the pre-warm phase of individual backend functions. This change ensures that autotuning occurs on the default stream to avoid illegal memory access errors during CUDA graph capture and prevents host-side cache lookups from interfering with GPU timing measurements. I have no feedback to provide as there were no review comments.