[Attention] Tune CUTLASS MLA num_splits

[MatthewBonanni]

Purpose

Tune the num_splits heuristic for CUTLASS_MLA to achieve some speedup now that #26026 has fixed the hang. Based on experiments performed using the tools introduced in #26835, this is the optimal num_splits policy:

Following the optimal policy would yield this speedup:

As a simpler alternative, we implement a heuristic yielding the following policy:

This results in the following speedup:

---

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 new heuristic for determining num_splits in CUTLASS MLA to improve performance. The new logic is based on the ratio of sequence length to batch size. While this is a reasonable approach for performance tuning, my review has identified a critical concern. The change removes a safeguard that was in place to prevent kernel hangs when the batch size is greater than one. Reintroducing this hang would be a critical issue, and it's not clear from the pull request description if the underlying problem has been resolved. I have left a comment detailing this concern.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ 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
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

[LucasWilkinson]

LGTM! thanks for doing this!