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

Fix SM-gating for CUTLASS/TMA compile options in GRPO trainer

Browse source 
Change `get_device_capability()[0] >= 9` to `== 9` for
`triton.enable_persistent_tma_matmul`, `cuda.cutlass_epilogue_fusion_enabled`,
and `cuda.cutlass_tma_only` in the GRPO torch_compile_options.

These options enable SM90-specific CUTLASS TMA kernels which are not valid
on SM100+ (Blackwell B200/B100). The `>= 9` check incorrectly enables them
on any GPU with SM >= 9, but they should only run on SM 9.x (Hopper).

This matches the existing pattern in `fp8.py:test_has_fbgemm()` which
documents that SM100 GPUs fail with CUTLASS SM90 kernels.
This commit is contained in:
Daniel Han 2026-03-08 13:07:52 +00:00
commit 8793366545
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/rl.py
View file

@ -1322,12 +1322,12 @@ 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 += """
"cuda.cutlass_epilogue_fusion_enabled": torch.cuda.get_device_capability()[0] >= 9,
"cuda.cutlass_tma_only": torch.cuda.get_device_capability()[0] >= 9,"""
"cuda.cutlass_epilogue_fusion_enabled": torch.cuda.get_device_capability()[0] == 9,
"cuda.cutlass_tma_only": torch.cuda.get_device_capability()[0] == 9,"""
cuda_options += """
"cuda.compile_opt_level" : "-O2",
"cuda.enable_cuda_lto" : True,