fix: ensure each CTA processes full numHeadsQPerKv for trtllm decode kernel

[dongjiyingdjy]

📌 Description

Skip candidates where kernelMeta.mStepQ < params.mNumHeadsQPerKv in GQA tile selection to avoid numTokensPerCtaQ=0, resulting divide-by-zero crash.

🔍 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

• Performance
  • Refined GPU tile-size selection to enforce a lower bound on candidates, producing more consistent and often improved modeling time and resource utilization during GPU-accelerated operations.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

A filtering condition in the GQA tile-size heuristic is tightened: candidate tileSizeQ values must satisfy both tileSizeQ ≤ defaultTileSizeQ and tileSizeQ ≥ mNumHeadsQPerKv, changing which candidates are evaluated during tile-size selection.

Changes

Cohort / File(s)	Summary
GQA Tile-Size Heuristic Constraint include/flashinfer/trtllm/fmha/fmhaKernels.cuh	Added lower-bound check tileSizeQ >= mNumHeadsQPerKv to the candidate filtering condition alongside the existing upper-bound check, narrowing candidate tileSizeQ set and affecting modeling time evaluation and final tile selection.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

🐇 I nibble code beneath the moonlight,

tiles trimmed to fit just right,
bounds aligned in tidy rows—
a quieter kernel softly grows.

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title accurately describes the main change: ensuring each CTA processes the full numHeadsQPerKv, which aligns with the filtering condition change that now requires tileSizeQ >= mNumHeadsQPerKv.
Description check	✅ Passed	The pull request description provides a clear, focused explanation of the bug fix (avoiding divide-by-zero crash by skipping invalid candidates), with the required Description section completed and Related Issues/Reviewer Notes sections appropriately marked as empty.

_{✏️ 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.}

[yzh119]

Hi @dongjiyingdjy would you mind explaning the context of this pull requests in the PR description?

[coderabbitai]

Caution

Docstrings generation - FAILED

No docstrings were generated.

[yzh119]

/bot run

[flashinfer-bot]

GitLab MR !250 has been created, and the CI pipeline #42126090 is currently running. I'll report back once the pipeline job completes.

[dongjiyingdjy]

Hi @dongjiyingdjy would you mind explaning the context of this pull requests in the PR description?

This PR ensures that all Q heads within the same group are in the same CTA. The previous tile select strategy did not account for this, which could cause Q heads from a single group to be split across multiple CTAs, leading to incorrect results.