ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Keep triton.enable_persistent_tma_matmul as >= 9, only gate CUTLASS options

Browse source 
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.
This commit is contained in:
Daniel Han 2026-03-08 13:19:38 +00:00
commit 55b49c8a1a
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/rl.py
View file

@ -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 += """