[Bugfix] Disable TMA on Blackwell GPUs to fix Triton autotuner OOM in fla/solve_trilfix: disable TMA on Blackwell (sm_12x) to prevent Triton autotuner OO…

[Rks2302]

Summary

Fixes Triton autotuner OOM crash in fla/ops/solve_tril.py when running
Qwen3.5 models on Blackwell GPUs (RTX 5090, compute capability sm_12x).

Root Cause

is_tma_supported evaluates to True on any GPU with compute capability >= 9,
which includes Blackwell (sm_12x). During first inference, the Triton autotuner
benchmarks the merge_fn kernel in solve_tril with TMA enabled, causing
oversized descriptor buffer allocations that OOM even when model weights fit
comfortably in VRAM.

Error

RuntimeError: Triton Error [CUDA]: out of memory
File "fla/ops/solve_tril.py", line 545, in solve_tril
merge_fn[NT, B * H](..., USE_TMA=is_tma_supported)
File "triton/runtime/autotuner.py"
timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}

Fix

Add upper bound < 12 to restrict TMA only to Hopper (sm_90x). TMA works
correctly on Hopper but causes Triton autotuner OOM on Blackwell (sm_12x).

Testing

• GPU: NVIDIA RTX 5090 (Blackwell, sm_120)
• vLLM: 0.17.0
• Model: Qwen3.5-35B-A3B-AWQ, Qwen3.5-27B-AWQ
• CUDA: 12.8

After this fix, Qwen3.5 AWQ models run successfully on RTX 5090 without
--enforce-eager. Full inference pipeline verified working.

Related Issues

• [Bug]: Qwen 3.5 27B AWQ 4bit capturing CUDA graph fails #35743 (Qwen3.5 AWQ CUDA graph failures on Blackwell)
• [Bug]: AssertionError in causal_conv1d_update when capturing CUDA graphs for Qwen3.5/GDN layers #35945 (causal_conv1d AssertionError - related GDN bug)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request addresses an out-of-memory error on Blackwell GPUs by disabling Tensor Memory Access (TMA) for this architecture. The fix is correct and targeted. My review includes a suggestion to refactor the implementation of the check to improve its readability and efficiency by avoiding a redundant function call.

[gemini-code-assist]

While this fix is correct, the implementation can be improved for readability and efficiency. The expression for is_tma_supported now calls torch.cuda.get_device_capability(0)[0] twice and the formatting makes the line very long. It's better to combine the two compute capability checks into a single range check to avoid the redundant call and make the condition clearer.

Suggested change

	is_tma_supported = (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9) and (
	hasattr(triton.language, "_experimental_make_tensor_descriptor")
	or hasattr(triton.language, "make_tensor_descriptor")
	)
	) and torch.cuda.get_device_capability(0)[0] < 12 # Disable on Blackwell (sm_12x): Triton autotuner OOM
	# Disable on Blackwell (sm_12x): Triton autotuner OOM
	is_tma_supported = (is_nvidia and 9 <= torch.cuda.get_device_capability(0)[0] < 12) and (
	hasattr(triton.language, "_experimental_make_tensor_descriptor")
	or hasattr(triton.language, "make_tensor_descriptor")
	)

[mergify]

Hi @Rks2302, 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

[mergify]

Hi @Rks2302, 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

[ZJY0516]

I remember Hopper also had this OOM issue. We should find a better way to both avoid OOM and maintain performance.