[chore] Fix CI pre-commit mypy error

[jiahanc]

📌 Description

Fix CI pre-commit mypy error

mypy.....................................................................Failed
- hook id: mypy
- exit code: 1

flashinfer/jit/core.py: note: In member "__init__" of class "JitSpecRegistry":
flashinfer/jit/core.py:164:9: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
flashinfer/jit/core.py:165:9: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
flashinfer/autotuner.py: note: In member "__init__" of class "AutoTuner":
flashinfer/autotuner.py:579:9: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
flashinfer/autotuner.py:581:9: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
flashinfer/comm/allreduce.py: note: In function "allreduce_fusion":
flashinfer/comm/allreduce.py:713:13: error: Missing positional arguments "quant_out", "scale_out", "routed_scaling_factor" in call to "trtllm_moe_finalize_allreduce_fusion"  [call-arg]
flashinfer/aot.py: note: In function "main":
flashinfer/aot.py:979:5: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
Found 1 error in 1 file (checked 193 source files)

🔍 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

• Bug Fixes
  • Fixed an issue in MOE finalization to ensure correct parameter handling during fused finalization steps, improving stability and correctness of final outputs.

[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: 54dfb396-de07-4ace-a13f-fe3cc46f46ed

📥 Commits

Reviewing files that changed from the base of the PR and between 172cb0a and 4767d49.

📒 Files selected for processing (1)

• flashinfer/comm/allreduce.py

---

📝 Walkthrough

Walkthrough

The PR updates flashinfer/comm/allreduce.py TRTLLM MOE Finalize dispatch to explicitly pass quant_out=quant_out, scale_out=scale_out, and routed_scaling_factor=None to trtllm_moe_finalize_allreduce_fusion; return value and control flow remain unchanged.

Changes

Cohort / File(s)	Summary
MOE Finalize Dispatch flashinfer/comm/allreduce.py	Updated the kMoEFinalizeARResidualRMSNorm branch to forward quant_out=quant_out, scale_out=scale_out, and routed_scaling_factor=None into trtllm_moe_finalize_allreduce_fusion call.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• PR #2966: Updates TRT-LLM MOE finalize dispatch to forward newly-added optional parameters (quant_out, scale_out, routed_scaling_factor).
• PR #2982: Introduced the MOE finalize pattern and initial trtllm_moe_finalize_allreduce_fusion call; this PR aligns the dispatch with that API.

Suggested labels

run-ci

Suggested reviewers

• yzh119
• jimmyzho
• bkryu
• nv-yunzheq
• aleozlx

Poem

🐰 I hopped through code, with nimble feet,
Three tiny args made the call complete,
quant and scale, and routed too,
Now finalize hums a tidier tune,
Patch done — a little hop of glee!

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[chore] Fix CI pre-commit mypy error' directly corresponds to the main objective: fixing a mypy type-checking error in the CI pre-commit pipeline.
Description check	✅ Passed	The PR description follows the template with a clear 📌 Description section explaining the mypy error, includes the ✅ Pre-commit Checks and 🧪 Tests sections with items marked complete.
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 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 updates the allreduce_fusion function in flashinfer/comm/allreduce.py to include quant_out, scale_out, and routed_scaling_factor parameters in an internal call. A review comment suggests passing the actual quant_out and scale_out variables from the function arguments instead of hardcoding them to None to properly support quantization outputs.

[yzh119]

LGTM, do you have any idea which PR breaks the pre-commit CI on mainline? How does it escape check?

[jiahanc]

LGTM, do you have any idea which PR breaks the pre-commit CI on mainline? How does it escape check?

Relevant code was introduced in this PR #2982. pre-commit didnt fail in PR stage but failed right after it merged