types

unsloth/kernels/rms_layernorm.py

@@ -191,7 +191,7 @@ pass
 
 
 # [TODO] Unsure why RMS Layernorm is not torch.compiling properly
-@torch.compiler.disable
+# @torch.compiler.disable
 def fast_rms_layernorm(layernorm, X : torch.Tensor, gemma : bool = False):
     W : torch.Tensor = layernorm.weight
     eps : float = layernorm.variance_epsilon if \

unsloth/kernels/rope_embedding.py

@@ -16,7 +16,7 @@ import triton
 import triton.language as tl
 import torch
 from .utils import calculate_settings
-ROPE_GROUP_SIZE = 4
+ROPE_GROUP_SIZE : int = 4
 
 @triton.heuristics({"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),})
 @triton.jit