From ae7cb78e9122b26536e8c25b1c79a2bd642e86da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 4 Nov 2024 23:24:33 -0800 Subject: [PATCH] Update rms_layernorm.py --- unsloth/kernels/rms_layernorm.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 821f6b540d..53f1efa9f5 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -16,6 +16,7 @@ import triton import triton.language as tl import torch from .utils import calculate_settings +next_power_of_2 = triton.next_power_of_2 @triton.jit @@ -142,9 +143,15 @@ class Fast_RMS_Layernorm(torch.autograd.Function): n_rows : int n_cols : int n_rows, n_cols = X.shape - BLOCK_SIZE : int - num_warps : int - # BLOCK_SIZE, num_warps = calculate_settings(n_cols) + BLOCK_SIZE : int = next_power_of_2(n_cols) + MAX_FUSED_SIZE : int = 65536 + if BLOCK_SIZE > MAX_FUSED_SIZE: + raise RuntimeError(f"Cannot launch Triton kernel since n = {n_cols} exceeds "\ + f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.") + num_warps : int = 4 + if BLOCK_SIZE >= 32768: num_warps = 32 + elif BLOCK_SIZE >= 8192: num_warps = 16 + elif BLOCK_SIZE >= 2048: num_warps = 8 Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0") r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") @@ -157,8 +164,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function): r, r.stride(0), n_cols = int(n_cols), eps = float(eps), - BLOCK_SIZE = 4096, - num_warps = 16, + BLOCK_SIZE = BLOCK_SIZE, + num_warps = num_warps, ) else: _gemma_rms_layernorm_forward[(n_rows,)]( @@ -171,8 +178,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function): num_warps = num_warps, ) ctx.eps = eps - ctx.BLOCK_SIZE = 4096 - ctx.num_warps = 16 + ctx.BLOCK_SIZE = BLOCK_SIZE + ctx.num_warps = num_warps ctx.GEMMA = gemma ctx.save_for_backward(X, W, r) return Y.view(*shape)