fix xqa mha_sm90.cu

[qsang-nv]

📌 Description

This PR fixes 2 things:
1 CUDA_ARCH is not defined in host code so xqa_wrapper.cu should use another macro
2 spec_q_seq_len is not included in jit module name, which could lead to crashes in unittests because test could use wrong cached ops

🔍 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

• New Features

  • Enhanced support for NVIDIA Hopper (SM90) GPUs with optimized compilation configuration.

• Bug Fixes

  • Refined speculative decoding behavior to correctly handle edge cases with sink parameters.

• Performance Improvements

  • Improved kernel selection logic for SM90-specific multi-head attention operations.

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

[coderabbitai]

Walkthrough

This PR refactors SM90 MHA kernel selection from runtime architecture checks to compile-time macro-based configuration. It updates the JIT compilation pipeline to conditionally include SM90-specific sources and flags, extends XQA operation registration with speculative query sequence length tracking, and refines runtime conditions for SM90 inference path selection to account for sink conflicts.

Changes

Cohort / File(s)	Summary
SM90 MHA Selection Mechanism csrc/xqa/xqa_wrapper.cu	Replaces __CUDA_ARCH__ == 900 architecture check with USE_SM90_MHA macro-based condition to determine SM90-specific MHA path selection at compile time rather than runtime.
JIT Compilation Configuration flashinfer/jit/xqa.py	Adds conditional SM90 detection and compilation flag handling: includes mha_sm90.cu and tensorMap.cpp only when SM90 is available, propagates USE_SM90_MHA flag into CUDA compile flags, and extends XQA kernel names with spec_q_seq_len_{q_seq_len} component.
Public API Registration & Runtime Conditions flashinfer/xqa.py	Updates XQA and fake-op registration strings to append spec_q_seq_len_{q_seq_len} suffix; refines run_sm90_fp8_mha runtime selection logic to only force it False when sinks are provided (speculative decoding + sinks conflict) rather than always for q_seq_len > 1.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

• Compilation flag propagation: Verify that USE_SM90_MHA flag is correctly assembled and passed through the build pipeline across all three files
• API signature consistency: Confirm that the new spec_q_seq_len suffix is applied uniformly to both xqa and fake-op registrations
• Runtime condition logic: Review the conditional logic change for run_sm90_fp8_mha to ensure the sinks-based decision correctly handles all cases (sinks=None vs. sinks provided with speculative decoding)

Possibly related PRs

• fix flaky xqa test #2126: Modifies kernel code within mha_sm90.cu (pointer-to-index access), complementing this PR's SM90 source inclusion and build configuration changes
• feat: add xqa fp8 mha and fp8 kv cache #1769: Adds run_sm90_fp8_mha runtime selection and launcher changes to xqa_wrapper.cu, related to the SM90 MHA selection refinements in this PR

Suggested reviewers

• yzh119
• cyx-6
• wenscarl
• nvmbreughe

Poem

🐰 A macro now guides where architecture once spoke,
SM90's path rebuilt in compile-time cloak,
Specs and sinks dance in runtime's choice,
Register names echo with sequenced voice,
Hopper hops swiftly through optimized code! ✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Title check	❓ Inconclusive	The title 'fix xqa mha_sm90.cu' is vague and does not clearly convey what was actually fixed. It mentions a file but not the specific issues addressed.	Revise the title to be more descriptive, e.g., 'Fix SM90 macro handling and include spec_q_seq_len in JIT module names' to better communicate the two distinct fixes.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description check	✅ Passed	The PR description adequately explains both fixes with clear bullet points about the CUDA_ARCH issue and the spec_q_seq_len omission, despite the checklist items being unchecked.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 23ff744 and acdb839.

📒 Files selected for processing (3)

• csrc/xqa/xqa_wrapper.cu (1 hunks)
• flashinfer/jit/xqa.py (2 hunks)
• flashinfer/xqa.py (3 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

flashinfer/jit/xqa.py (1)

flashinfer/jit/core.py (1)

• gen_jit_spec (315-381)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Deploy Docs

🔇 Additional comments (6)

csrc/xqa/xqa_wrapper.cu (1)

71-75: Critical fix: Correct use of compile-time macro instead of device-only macro.

The replacement of __CUDA_ARCH__ with USE_SM90_MHA is essential and correct. __CUDA_ARCH__ is only defined during device code compilation and would be undefined in host code, causing incorrect behavior. The new USE_SM90_MHA macro is properly set via compiler flags from the JIT compilation pipeline.

flashinfer/jit/xqa.py (3)

119-119: JIT spec name correctly extended with spec_q_seq_len.

The addition of spec_q_seq_len_{q_seq_len} to the JIT spec name addresses the second issue mentioned in the PR objectives. This prevents cache collisions when different q_seq_len values are used, which was causing crashes in unittests.

---

131-132: Compile flag correctly propagated.

The flag_sm90_mha is properly appended to extra_cuda_cflags, ensuring the USE_SM90_MHA macro is defined during compilation and available in the C++ code.

---

110-116: Verify that tensorMap.cpp is SM90-specific before conditional inclusion.

The code conditionally includes tensorMap.cpp only when has_sm90 is True. This assumption should be validated: if tensorMap.cpp contains functionality needed for non-SM90 architectures (SM10/SM11/SM12), the conditional logic requires adjustment. Confirm that tensorMap functionality is exclusive to SM90 before merging.

flashinfer/xqa.py (2)

61-61: Op registration names correctly extended.

The operation registration names now include spec_q_seq_len_{q_seq_len}, matching the JIT spec name extension in flashinfer/jit/xqa.py. This consistency is critical for correct operation dispatch and addresses the cache collision issue mentioned in the PR objectives.

Also applies to: 114-114

---

303-306: Verify test coverage for speculative decoding with sinks condition.

The refined SM90 MHA disable condition is logically sound—disabling it only when both sinks is not None and q_seq_len > 1 is more targeted than always disabling for speculative decoding. However, confirm that this behavioral change is intentional in this PR and includes appropriate test coverage for the sinks + speculative decoding scenario, as current FlashInfer test coverage for speculative decoding is incomplete across parameter combinations.

---

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]

/bot run

[flashinfer-bot]

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

[flashinfer-bot]

[FAILED] Pipeline #39489427: 4/20 passed