[Bugfix] fix DeepSeek R1 with CUTLASS MLA Broken on B200

[chaunceyjiang]

Purpose

FIX #33627

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 aims to fix an issue with DeepSeek R1 using CUTLASS MLA on B200 GPUs by providing a default value for q_pad_num_heads. However, the current implementation is not correct as the default value is set in a place where it will not be used by the padding logic. The fix needs to be applied in a different location to be effective. I've left a critical comment explaining the issue with the current approach.

[mergify]

Hi @chaunceyjiang, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[chaunceyjiang]

/cc @MatthewBonanni @LucasWilkinson PTAL.

[MatthewBonanni]

Thanks for the fix! LGTM once comment is addressed

[MatthewBonanni]

Ah actually, I think the current state of the PR will force q_pad_num_heads to be None rather than MAX_HEADS as it should. Can you verify correctness?

[MatthewBonanni]

I think the proper fix would be:

• Remove q_pad_num_heads entirely from the argument list of MLAAttention.__init__().
• When calling impl_cls() in MLAAttention.__init__(), don't pass q_pad_num_heads.
• Assign self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) after the impl has been constructed
• Leave cutlass_mla.py as is