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

fix(fp8): use int64 offsets in weight_dequant_kernel (#6884)

Browse source
This commit is contained in:
Anas Khan 2026-07-06 19:41:41 +05:30 • committed by GitHub
commit f4d1dc541f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/fp8.py
View file

@ -68,7 +68,9 @@ def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
n = tl.cdiv(N, BLOCK_SIZE)
offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs = offs_m[:, None] * N + offs_n[None, :]
# tl.arange is int32, so offs_m * N overflows for tensors with more than
# 2**31 elements (e.g. flattened MoE expert stacks); index in int64.
offs = offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64)
mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
x = tl.load(x_ptr + offs, mask = mask).to(tl.float32)
s = tl.load(s_ptr + pid_m * n + pid_n)