ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update rms_layernorm.py

Browse source
This commit is contained in:
Daniel Han 2024-11-04 23:27:50 -08:00
commit 22da266219
1 changed files with 7 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

11
unsloth/kernels/rms_layernorm.py
View file

@ -15,8 +15,11 @@
import triton
import triton.language as tl
import torch
from .utils import calculate_settings
next_power_of_2 = triton.next_power_of_2
from .utils import (
calculate_settings,
MAX_FUSED_SIZE,
next_power_of_2,
)
@triton.jit
@ -143,8 +146,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
n_rows : int
n_cols : int
n_rows, n_cols = X.shape
BLOCK_SIZE : int = next_power_of_2(n_cols)
MAX_FUSED_SIZE : int = 65536
BLOCK_SIZE : int = n_cols
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}.")