Fix triton 3.6.0 + torch 2.9.x torch.compile crash (missing cluster_dims) (#4001)

Co-authored-by: Daniel Hanchen <danielhanchen@users.noreply.github.com>

unsloth/__init__.py

@@ -128,6 +128,7 @@ from .import_fixes import (
     check_vllm_torch_sm100_compatibility,
     fix_vllm_guided_decoding_params,
     fix_vllm_pdl_blackwell,
+    fix_triton_compiled_kernel_missing_attrs,
     fix_rocm_triton_key_error,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
@@ -148,6 +149,7 @@ fix_vllm_aimv2_issue()
 check_vllm_torch_sm100_compatibility()
 fix_vllm_guided_decoding_params()
 fix_vllm_pdl_blackwell()
+fix_triton_compiled_kernel_missing_attrs()
 fix_rocm_triton_key_error()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
@@ -166,6 +168,7 @@ del fix_vllm_aimv2_issue
 del check_vllm_torch_sm100_compatibility
 del fix_vllm_guided_decoding_params
 del fix_vllm_pdl_blackwell
+del fix_triton_compiled_kernel_missing_attrs
 del fix_rocm_triton_key_error
 del ignore_logger_messages
 del patch_ipykernel_hf_xet

unsloth/import_fixes.py

@@ -799,6 +799,54 @@ def fix_huggingface_hub():
         )
 
 
+def fix_triton_compiled_kernel_missing_attrs():
+    """
+    Triton 3.6.0+ removed direct `num_ctas` and `cluster_dims` attributes from
+    CompiledKernel, but torch 2.9.x Inductor still expects them in
+    torch/_inductor/runtime/triton_heuristics.py make_launcher() (line ~1757).
+
+    The scope dict eagerly evaluates:
+        binary.metadata.num_ctas, *binary.metadata.cluster_dims
+    when hasattr(binary, "metadata") is True, but metadata lacks cluster_dims.
+    This crashes before reaching the new launch path that doesn't need cta_args.
+
+    Upstream fix: pytorch/pytorch@97bd4db added hasattr guards.
+    We monkey-patch CompiledKernel.__init__ to inject the missing attributes
+    so the older hasattr(binary, "num_ctas") branch succeeds instead.
+    """
+    try:
+        import torch
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    try:
+        import triton
+        import triton.compiler.compiler as triton_compiler
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    # Only needed when the CompiledKernel class lacks num_ctas as a direct attr
+    # but has metadata (triton >= 3.6.0 with torch < 2.10)
+    _ck_cls = triton_compiler.CompiledKernel
+    if hasattr(_ck_cls, "num_ctas"):
+        return  # Old triton with direct attrs -- no patch needed
+
+    _orig_init = _ck_cls.__init__
+
+    def _patched_init(self, *args, **kwargs):
+        _orig_init(self, *args, **kwargs)
+        if not hasattr(self, "num_ctas"):
+            self.num_ctas = getattr(self.metadata, "num_ctas", 1)
+        if not hasattr(self, "cluster_dims") and not hasattr(self, "clusterDims"):
+            self.cluster_dims = (1, 1, 1)
+
+    _ck_cls.__init__ = _patched_init
+    logger.info(
+        "Unsloth: Patched triton CompiledKernel with num_ctas/cluster_dims "
+        "for torch.compile compatibility."
+    )
+
+
 def fix_rocm_triton_key_error():
     """
     ROCm + torch.compile can fail if Triton lacks `triton_key`.