diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 4357ad63aa..b068d6a5fc 100644 --- a/unsloth/__init__.py +++ b/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 diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 97e74dfb57..cd8875b5bf 100644 --- a/unsloth/import_fixes.py +++ b/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`.