From 71a3333c0f0ca109e8bcc6f0ee83721e89db5e7e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 4 Nov 2024 23:20:32 -0800
Subject: [PATCH] Update rms_layernorm.py

---
 unsloth/kernels/rms_layernorm.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py
index 23b7342c87..821f6b540d 100644
--- 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)