diff --git a/unsloth/kernels/swiglu.py b/unsloth/kernels/swiglu.py index 4e9b7ba2ae..3106656f32 100644 --- a/unsloth/kernels/swiglu.py +++ b/unsloth/kernels/swiglu.py @@ -24,11 +24,11 @@ def _fg_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,): offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < n_elements - e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32) + e_row = tl.load(e + offsets, mask = mask, other = 0)#.to(tl.float32) g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32) # f = e * sigmoid(e) - f_row = e_row / (1 + tl.exp(-e_row)) + f_row = e_row / (1 + tl.exp(-e_row.to(tl.float32)).to(g_row.dtype)) f_row = f_row.to(g_row.dtype) # Exact copy from HF # h = f * g h_row = f_row * g_row