Update rms_layernorm.py

2024-11-04 23:20:32 -08:00 · 2024-11-04 23:20:32 -08:00 · 40858e4f76
commit 40858e4f76
parent a3429d208d
1 changed files with 2 additions and 2 deletions
--- a/unsloth/kernels/rms_layernorm.py
+++ b/unsloth/kernels/rms_layernorm.py
@ -171,8 +171,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
                num_warps  = num_warps,
            )
        ctx.eps = eps
-        ctx.BLOCK_SIZE = BLOCK_SIZE
-        ctx.num_warps  = num_warps
+        ctx.BLOCK_SIZE = 4096
+        ctx.num_warps  = 16
        ctx.GEMMA = gemma
        ctx.save_for_backward(X, W, r)
        return Y.view(*shape)