From 55b49c8a1ac748a72126b6940a70ea959a69a37c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 8 Mar 2026 13:19:38 +0000
Subject: [PATCH] Keep triton.enable_persistent_tma_matmul as >= 9, only gate
 CUTLASS options

Triton TMA persistent matmul works on SM90+ (PyTorch's own
has_triton_tma_device() checks >= (9,0)). Only CUTLASS TMA-specific
kernels (cutlass_tma_only, cutlass_epilogue_fusion_enabled) are SM90-only
and need the == 9 gate.
---
 unsloth/models/rl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 7faa61fd59..af027ff9c7 100755
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -1322,7 +1322,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         if DEVICE_TYPE == "cuda":
             # CUDA-specific options (added to base options)
             cuda_options = """
-            "triton.enable_persistent_tma_matmul": torch.cuda.get_device_capability()[0] == 9,"""
+            "triton.enable_persistent_tma_matmul": torch.cuda.get_device_capability()[0] >= 9,"""
             # cutlass options were added in PyTorch 2.8.0
             if torch_version >= Version("2.8.0"):
                 cuda_options += """