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

Restore kernel files to upstream to keep their explanatory comments

Browse source
This commit is contained in:
Daniel Han 2026-06-18 08:23:21 +00:00
commit c8a76c78ea
10 changed files with 92 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/kernels/__init__.py
View file

@ -44,7 +44,7 @@ from .fast_lora import (
apply_lora_o,
fast_lora_forward,
)
from .fp8 import * # Patch FP8Linear forwards before model creation to cover compiled non-fast-inference models
from .fp8 import * # Patch FbgmemFP8Linear/FP8Linear forwards before model creation, so compiled non-fast-inference models are covered too
from .utils import (
fast_dequantize,
fast_gemv,

13
unsloth/kernels/cross_entropy_loss.py
View file

@ -173,7 +173,8 @@ def _chunked_cross_entropy_forward(
logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))
if chunk_idx == 0:
# logsumexp(chunked_logsumexp) - x; do the -x separately
# logsumexp(chunked_logsumexp) - x
# Do the -x separately
if label_idx != -100:
x = tl.load(logits_ptr + label_idx).to(tl.float32)
# Go logit scaling for Cohere: t * x
@ -312,7 +313,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
BLOCK_SIZE: int
num_warps: int
if n_chunks == 1:
# Small vocabs <= 65336 (Llama, Mistral)
# For small vocabs <= 65336 like Llama, Mistral
BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
if is_cdna():
num_warps = num_warps // 2
@ -334,7 +335,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
num_warps = num_warps,
)
else:
# Large vocabs > 65336 (Gemma 256K)
# For large vocabs > 65336 like Gemma 256K
logsumexp = torch.empty(
(
n_rows,
@ -365,10 +366,11 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
LOGIT_SCALE = logit_scaling,
num_warps = 32 if not is_cdna() else 16,
)
# logsumexp(chunked_logsumexp) - x; do the -x separately
# logsumexp(chunked_logsumexp) - x
# Do the -x separately
logsumexp = torch.logsumexp(logsumexp, dim = 1) # Row sum
losses += logsumexp
losses.masked_fill_(labels == -100, 0) # Mask padding out
losses.masked_fill_(labels == -100, 0) # Don't forget to mask padding out!
ctx.save_for_backward(logits, logsumexp, labels)
ctx.DO_SOFTCAPPING = DO_SOFTCAPPING
@ -456,6 +458,7 @@ if (Version(torch.__version__) < Version("2.4.0")) and not hasattr(
fast_cross_entropy_loss = torch._disable_dynamo(fast_cross_entropy_loss)
# Patch CE Losses in transformers
def patch_loss_functions(torch_compile = True):
_patch_loss_functions(fast_cross_entropy_loss, torch_compile = torch_compile)

15
unsloth/kernels/fast_lora.py
View file

@ -40,6 +40,8 @@ class LoRA_MLP(torch.autograd.Function):
i = h @ W
### Backpropagation chain rule
See our blog post for more details
df = sigmoid(e) * (1 - f) + f
dC/dW = h.T @ dY
dC/dU = X.T @ (D @ W.T * f)
@ -58,6 +60,8 @@ class LoRA_MLP(torch.autograd.Function):
### Gate projection LoRA weights
dC/dAg = X.T @ (D @ W.T * df * g) @ B.T
dC/dBg = A.T @ X.T @ (D @ W.T * df * g)
Don't forget to see our blog post for more details!
"""
@staticmethod
@ -339,6 +343,8 @@ class LoRA_QKV(torch.autograd.Function):
V = X @ Wv = X @ Wv + X @ Av @ Bv
### Backpropagation chain rule
See our blogpost for more details.
dC/dWq = X.T @ D(Wq)
dC/dWk = X.T @ D(Wk)
dC/dWv = X.T @ D(Wv)
@ -381,8 +387,9 @@ class LoRA_QKV(torch.autograd.Function):
):
dtype = X.dtype
# bitsandbytes 8-bit matmul expects 2D; TorchInductor/AOTAutograd fails on
# 3D during backward, so flatten the sequence dim.
# bitsandbytes 8-bit matmul expects 2D inputs.
# TorchInductor/AOTAutograd fails on 3D tensors during backward,
# so we explicitly flatten the sequence dimension.
orig_shape = X.shape
X_for_matmul = X
if X.dim() == 3:
@ -391,6 +398,7 @@ class LoRA_QKV(torch.autograd.Function):
K = matmul_lora(X_for_matmul, KW, KW_quant, KA, KB, KS)
V = matmul_lora(X_for_matmul, VW, VW_quant, VA, VB, VS)
# Restore original shape after matmul
if len(orig_shape) == 3:
Q = Q.view(orig_shape[0], orig_shape[1], -1)
K = K.view(orig_shape[0], orig_shape[1], -1)
@ -452,6 +460,7 @@ class LoRA_QKV(torch.autograd.Function):
QA, QB, KA, KB, VA, VB = QA.t(), QB.t(), KA.t(), KB.t(), VA.t(), VB.t()
### Weight projection LoRA weights
# See our blogpost for more details.
d_QA = torch.empty_like(QA)
d_QB = torch.empty_like(QB)
d_KA = torch.empty_like(KA)
@ -622,6 +631,7 @@ class LoRA_W(torch.autograd.Function):
d_B = torch.empty_like(B)
### Weight projection LoRA weights
# Weight projection
# d_A = X.t() @ (dY @ B.t())
# d_B = (A.t() @ X.t()) @ dY
# d_A *= S
@ -629,6 +639,7 @@ class LoRA_W(torch.autograd.Function):
d_A.addmm_(X.t(), dY @ B.t(), alpha = S, beta = 0)
d_B.addmm_(A.t() @ X.t(), dY, alpha = S, beta = 0)
# Get derivative for dX
W = fast_dequantize(W.t(), W_quant)
dX = dY @ W.t()
del W

12
unsloth/kernels/fp8.py
View file

@ -184,7 +184,10 @@ def _w8a8_block_fp8_matmul(
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
):
"""Triton block-wise quantized matmul: C = A @ B."""
"""Triton-accelerated function used to perform linear operations (dot
product) on input tensors `A` and `B` with block-wise quantization, and
store the result in output tensor `C`.
"""
pid = tl.program_id(axis = 0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
@ -329,7 +332,7 @@ class FP8BlockQuantLinear(torch.autograd.Function):
def forward(ctx, X, weight, weight_scale):
m, n = weight.shape
# Saved for backward before any transformation
# Original scale, saved for backward before any transformation
original_weight_scale = weight_scale
# Per-tensor quant: expand scalar to (ceil(m/128), ceil(n/128)) block shape
@ -339,6 +342,7 @@ class FP8BlockQuantLinear(torch.autograd.Function):
num_blocks_n = triton.cdiv(n, block_size[1])
weight_scale = weight_scale.expand(num_blocks_m, num_blocks_n).contiguous()
else:
# Block quantization path
p, q = weight_scale.shape
block_size = getattr(weight, "block_size", None) or getattr(
weight_scale, "block_size", [128, 128]
@ -347,7 +351,7 @@ class FP8BlockQuantLinear(torch.autograd.Function):
if triton.cdiv(m, block_size[0]) != p or triton.cdiv(n, block_size[1]) != q:
if triton.cdiv(m, block_size[0]) == q and triton.cdiv(n, block_size[1]) == p:
weight_scale = weight_scale.T
original_weight_scale = weight_scale
original_weight_scale = weight_scale # Update for transposed case
else:
raise ValueError(
f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {block_size}"
@ -366,7 +370,7 @@ class FP8BlockQuantLinear(torch.autograd.Function):
output_dtype = X.dtype,
)
ctx.weight = weight
ctx.weight_scale = original_weight_scale
ctx.weight_scale = original_weight_scale # Save original for backward
return output.to(X.dtype)
@staticmethod

14
unsloth/kernels/geglu.py
View file

@ -49,6 +49,7 @@ def _exact_forward_kernel(
f_row = f_row.to(g_row.dtype) # Exact copy from HF
h_row = f_row * g_row
# Store h
tl.store(h + offsets, h_row, mask = mask)
@ -96,13 +97,17 @@ def _exact_backward_kernel(
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 = 1/2 * e * (1 + erf(1/sqrt(2) * e)); reuse f_partial_row below
# Break e_row away for re-use
# f = 1/2 * e * (1 + erf(1/sqrt(2) * e))
f_partial_row = 0.5 * (tl.math.erf(tl.math.rsqrt(2.0) * e_row) + 1.0)
f_row = f_partial_row * e_row
f_row = f_row.to(DW_row.dtype)
# h = f * g
h_row = f_row * g_row
# df = DW * f
df_row = DW_row * f_row
# dg = DW * g
dg_row = DW_row * g_row
# df/de = 1/2 * (1 + erf(1/sqrt(2) * e)) + 1/sqrt(2*pi) * e * exp(-1/2 * e^2)
@ -112,6 +117,7 @@ def _exact_backward_kernel(
de_row = dg_row.to(tl.float32) * df_de
de_row = de_row.to(DW_row.dtype)
# Store derivatives in buffers
tl.store(DW + offsets, h_row, mask = mask) # h = f * g
tl.store(e + offsets, df_row, mask = mask) # df = DW * f
tl.store(g + offsets, de_row, mask = mask) # de
@ -157,6 +163,7 @@ def _approx_forward_kernel(
f_row = f_row.to(g_row.dtype) # Exact copy from HF
h_row = f_row * g_row
# Store h
tl.store(h + offsets, h_row, mask = mask)
@ -208,6 +215,7 @@ def _approx_backward_kernel(
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)
# See https://www.desmos.com/calculator/nqprfoni6x
s = 0.7978845608028654 # math.sqrt(2 / math.pi)
a = s * e_row # a = sqrt(2 / pi) * x
b = a * 0.044715 * e_row * e_row # b = a * 0.044715 * x^2
@ -220,13 +228,17 @@ def _approx_backward_kernel(
# f = 1/2 * e * (1 + tanh( sqrt(2/pi) * (x + 0.044715 * x^3 ) ))
f_row = T2 * e_row
f_row = f_row.to(DW_row.dtype)
# h = f * g
h_row = f_row * g_row
# df = DW * f
df_row = DW_row * f_row
# dg = DW * g
dg_row = DW_row * g_row
de_row = dg_row.to(tl.float32) * df_de
de_row = de_row.to(DW_row.dtype)
# Store derivatives in buffers
tl.store(DW + offsets, h_row, mask = mask) # h = f * g
tl.store(e + offsets, df_row, mask = mask) # df = DW * f
tl.store(g + offsets, de_row, mask = mask) # de

12
unsloth/kernels/layernorm.py
View file

@ -45,17 +45,17 @@ def layernorm_forward(
r += row_idx
mu += row_idx
# torchtune Fp32LayerNorm: all modules are float32
# https://pytorch.org/torchtune/stable/_modules/torchtune/modules/layer_norm.html#Fp32LayerNorm
# According to https://pytorch.org/torchtune/stable/_modules/torchtune/modules/layer_norm.html#Fp32LayerNorm, all modules
# are in float32!
X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32)
W_row = tl.load(W + col_offsets, mask = mask, other = 0).to(tl.float32)
b_row = tl.load(b + col_offsets, mask = mask, other = 0).to(tl.float32)
mean_X = tl.sum(X_row, axis = 0) / n_cols
# Mask out masked positions, where (X - mean) collapses to -mean
# (X[0] - mean) == -mean so we need to mask it out
XX = tl.where(mask, X_row - mean_X, 0)
row_var = tl.sum(XX * XX, axis = 0) / n_cols
# Explicit float32 scalar for correct type promotion on HIP/ROCm
# Explicit float32 scalar to ensure correct type promotion on HIP/ROCm
eps_f32 = tl.full((), eps, tl.float32)
inv_var = tl.math.rsqrt(row_var + eps_f32)
tl.store(r, inv_var)
@ -88,8 +88,8 @@ def layernorm_backward(
r += row_idx
mu += row_idx
# torchtune Fp32LayerNorm: all modules are float32
# https://pytorch.org/torchtune/stable/_modules/torchtune/modules/layer_norm.html#Fp32LayerNorm
# According to https://pytorch.org/torchtune/stable/_modules/torchtune/modules/layer_norm.html#Fp32LayerNorm, all modules
# are in float32!
dY_row = tl.load(dY + col_offsets, mask = mask, other = 0).to(tl.float32)
X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32)
W_row = tl.load(W + col_offsets, mask = mask, other = 0).to(tl.float32)

1
unsloth/kernels/rms_layernorm.py
View file

@ -98,6 +98,7 @@ def _rms_layernorm_backward(
X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32)
W_row = tl.load(W + col_offsets, mask = mask, other = 0).to(tl.float32)
# Get saved row variance
inv_var = tl.load(r).to(tl.float32)
normed = X_row * inv_var

11
unsloth/kernels/rope_embedding.py
View file

@ -114,7 +114,11 @@ def _rope_embedding(
BACKWARD_PASS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""RoPE embedding: Q * cos + rotate_half(Q) * sin."""
"""
Calculates the RoPE Embedding quickly
RoPE is Q * cos + rotate_half(Q) * sin
See our blog post for more info
"""
ROPE_GROUP_SIZE = 4
row_position = tl.program_id(0)
group_head_position = tl.program_id(1)
@ -134,7 +138,7 @@ def _rope_embedding(
)
if BACKWARD_PASS:
# Backward negates sin (rotation transpose).
# See our blog post for more info.
sin1 = -sin1
# [TODO] Autotune ROPE_GROUP_SIZE to be 1, 2, 4, 8
@ -177,7 +181,8 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
n_rows, n_cols = Q.shape
assert seq_len <= cos.shape[0]
# [TODO] head_dim//2 blocksize causes concurrency / nondeterminism issues.
# [TODO] Changing blocksize to head_dim//2 seems to have
# some concurrency / un-deterministic issues.
BLOCK_SIZE, num_warps = calculate_settings(head_dim // 2) # (head_dim//2)
# group_size = 4 # 4 or 8, too large group_size can hurt performance.

13
unsloth/kernels/swiglu.py
View file

@ -17,7 +17,7 @@ import triton.language as tl
import torch
from .utils import calculate_settings, torch_gpu_device
# signed int32 max is 2**31-1, so num_elements cannot exceed it
# signed int32 max is 2**31-1 so num_elements cannot exceed 2**31
NUM_INT32_ELEMENTS = 2**31
SAFE_INT32_BUFFER_MULTIPLIER = 4
BLOCK_SIZE = 1024
@ -37,10 +37,13 @@ def _fg_kernel(e, g, h, n_elements, BLOCK_SIZE: tl.constexpr, LONG_INDEXING: tl.
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 * tl.sigmoid(e_row) # e_row / (1 + tl.exp(-e_row))
f_row = f_row.to(g_row.dtype) # Exact copy from HF
# h = f * g
h_row = f_row * g_row
# Store h
tl.store(h + offsets, h_row, mask = mask)
@ -84,17 +87,23 @@ def _DWf_DW_dfg_kernel(DW, e, g, n_elements, BLOCK_SIZE: tl.constexpr, LONG_INDE
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)
# e = e.float()
# se = 1.0 / (1.0 + torch.exp(-e))
se_row = tl.sigmoid(e_row) # 1.0 / (1.0 + tl.exp(-e_row))
# f = (se * e).to(dtype)
f_row = se_row * e_row
f_row = f_row.to(DW_row.dtype)
# h = f * g
h_row = f_row * g_row
# df = DW * f
df_row = DW_row * f_row
# dg = DW * g
dg_row = DW_row * g_row
# de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
de_row = dg_row.to(tl.float32) * se_row * (1.0 + e_row * (1.0 - se_row))
de_row = de_row.to(DW_row.dtype)
# Reuse input buffers to store outputs h, df, de
# Store derivatives in buffers
tl.store(DW + offsets, h_row, mask = mask) # h = f * g
tl.store(e + offsets, df_row, mask = mask) # df = DW * f
tl.store(g + offsets, de_row, mask = mask) # de

34
unsloth/kernels/utils.py
View file

@ -164,7 +164,8 @@ elif DEVICE_TYPE == "mlx":
elif hasattr(torch._C, "_cuda_getCurrentRawStream"):
_gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
else:
# CPU-only torch wheel: _get_tensor_stream only runs during real GPU work, so no-op is safe.
# CPU-only torch wheel (no compiled CUDA backend). _get_tensor_stream
# is only invoked during real GPU work, so a no-op binding is safe.
def _gpu_getCurrentRawStream(_index = 0):
return 0
@ -176,13 +177,16 @@ def _get_tensor_stream(tensor: torch_Tensor) -> c_void_p:
return c_void_p(_gpu_getCurrentRawStream(tensor.device.index))
# Get array of CUDA streams and other buffers
global CUDA_STREAMS
global XPU_STREAMS
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
# DEVICE_COUNT == 0 (CPU-only): empty containers, only indexed during real GPU
# work, so they just need to exist for the module to import cleanly.
# DEVICE_COUNT == 0 = no visible accelerator (e.g. CPU-only CI runner).
# The consumer functions below only index these arrays during real GPU
# work, so empty containers are safe -- they just need to be defined so
# the module imports cleanly.
if DEVICE_TYPE == "xpu":
if DEVICE_COUNT > 0:
_XPU_STREAMS = {
@ -236,8 +240,8 @@ cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_n
cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4
if DEVICE_TYPE == "xpu":
# xpu inference gemv, see:
# https://github.com/bitsandbytes-foundation/bitsandbytes/blob/c3b8de268fdb55a88f92feada23fc811a1e6877a/bitsandbytes/backends/xpu/ops.py#L115
# for xpu, inference gemv using above link
cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemv_4bit_inference_fp16
cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemv_4bit_inference_bf16
else:
@ -259,7 +263,7 @@ torch_float16 = torch.float16
torch_bfloat16 = torch.bfloat16
# Float8Tensor from torchao if available
# Check whether torchao can be imported to get Float8Tensor
if importlib.util.find_spec("torchao") is not None:
try:
from torchao.quantization import Float8Tensor
@ -288,7 +292,7 @@ def get_lora_parameters(proj):
) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
W = base_layer.weight
# QAT: fake-quantize base layer weights
# Optionally apply fake quantization to base layer weights for QAT
if hasattr(base_layer, "weight_fake_quantizer"):
weight_fake_quantizer = getattr(base_layer, "weight_fake_quantizer", None)
if weight_fake_quantizer is not None:
@ -302,7 +306,7 @@ def get_lora_parameters(proj):
W_quant = getattr(base_layer, "weight_scale", None)
if getattr(base_layer, "quant_method", None) == "fp8":
# Stash fp8 block_size on W/W_quant to pass downstream
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
@ -315,7 +319,7 @@ def get_lora_parameters(proj):
adapter = getattr(proj, "active_adapter", ("default"))
adapter = adapter[0]
# QAT: fake-quantize lora weights
# Optionally apply fake quantization to lora weights for QAT
lora_A_linear = proj.lora_A[adapter]
lora_B_linear = proj.lora_B[adapter]
A = lora_A_linear.weight
@ -357,7 +361,7 @@ def get_lora_parameters_bias(proj):
return W, W_quant, None, None, None, base_layer.bias
if getattr(base_layer, "quant_method", None) == "fp8":
# Stash fp8 block_size on W/W_quant to pass downstream
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
@ -427,8 +431,9 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
XPU_STREAM = XPU_STREAMS[device_index]
n_elements_absmax = absmax.numel()
# Create weight matrix
if use_global_buffer:
# Reuse buffers for faster inference
# Use same buffers for faster inference
size = shape[0] * shape[1]
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
@ -479,6 +484,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
)
out_absmax += offset
# Dequantize W
fx = (
cdequantize_blockwise_fp16_nf4
if dtype == torch_float16
@ -538,8 +544,9 @@ elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
n_elements_absmax = absmax.numel()
# Create weight matrix
if use_global_buffer:
# Reuse buffers for faster inference
# Use same buffers for faster inference
size = shape[0] * shape[1]
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
@ -591,6 +598,7 @@ elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
)
out_absmax += offset
# Dequantize W
fx = (
cdequantize_blockwise_fp16_nf4
if dtype == torch_float16
@ -648,6 +656,7 @@ else:
n_elements_absmax = absmax.numel()
device = W.device
# Create weight matrix
if out is None:
out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False)
else:
@ -657,6 +666,7 @@ else:
n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False
)
# Do dequantization
ptr_out_absmax = get_ptr(out_absmax)
cdequantize_blockwise_fp32(
get_ptr(code2),
@ -1030,6 +1040,7 @@ def fast_linear_forward(
W = fast_dequantize(W.t(), W_quant, use_global_buffer = True)
out = torch_matmul(X, W, out = out)
# Add in LoRA weights
if lora_A is not None:
out_dim = out.shape[2]
dtype = X.dtype
@ -1090,6 +1101,7 @@ def matmul_lora(
del W
if A is not None:
# LoRA is enabled
A, B = A.t(), B.t()
XA = torch_matmul(X, A.to(dtype))
out.addmm_(XA, B.to(dtype), alpha = s)