diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 10b435dd57..4b22f8c3e5 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/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 \ diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index 246055dc7b..7fe15d0e3b 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -25,9 +25,9 @@ def _rope_embedding( cos, cos_row_stride, sin, sin_row_stride, seqlen, - head_dim, - n_heads, - BACKWARD_PASS, + head_dim : tl.constexpr, + n_heads : tl.constexpr, + BACKWARD_PASS : tl.constexpr, BLOCK_SIZE : tl.constexpr, ): """ @@ -144,7 +144,8 @@ class Fast_RoPE_Embedding(torch.autograd.Function): pass pass - +# [TODO] Unsure why RoPE Embedding is not torch.compiling properly +@torch.compiler.disable def fast_rope_embedding(Q, K, cos, sin): Q = Fast_RoPE_Embedding.apply(Q.transpose(1, 2), cos, sin).transpose(1, 2) K = Fast_RoPE_Embedding.apply(K.transpose(1, 2), cos, sin).transpose(1, 2)