fix(fp8): use int64 offsets in weight_dequant_kernel

[anxkhn]

What

weight_dequant_kernel in unsloth/kernels/fp8.py builds its flat load/store
offset as:

offs = offs_m[:, None] * N + offs_n[None, :]

offs_m and offs_n come from tl.arange(0, BLOCK_SIZE), which is int32 in
Triton, and N does not promote them. So offs_m * N is computed in 32-bit and
overflows for any 2D tensor with more than 2**31 (2,147,483,648) elements: the
offset wraps to a negative value and tl.load/tl.store read and write out of
bounds, crashing with:

RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

This is hit on the FP8 block-quantized MoE dequant path, where a layer's stacked
expert weight is flattened to a single 2D tensor before dequant. For example a
[256, 4096, 6144] gate_up_proj stack becomes [1048576, 6144], about 6.4e9
elements, roughly 3x over the int32 limit. Smaller MoEs stay under 2**31, which
is why this has not shown up until larger expert stacks.

Fixes #6830.

Fix

Cast the offset arithmetic to int64 so the multiply is a widening 64-bit op:

offs = offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64)

The overflow is in the element product offs_m * N, so both operands are cast
before the multiply. The mask (offs_m < M) & (offs_n < N) stays int32 on
purpose: those are per-block indices bounded by a single matrix dimension, never
the flattened product, so they cannot overflow. The scale-pointer offset
pid_m * n + pid_n is tiny and is left unchanged.

This matches how the sibling kernels already index in 64-bit:
unsloth/kernels/swiglu.py (the LONG_INDEXING int64 path, with the
"signed int32 max is 2**31-1" note) and unsloth/kernels/cross_entropy_loss.py
(row_idx * ...to(tl.int64)). The mlp kernels were fixed for this same class in
PR #3614, which did not touch fp8.py.

Testing

• ruff check unsloth/kernels/fp8.py passes.
• The failure is a CUDA-only illegal memory access that only triggers above the
  2**31-element threshold, so it needs a large-tensor GPU run to reproduce. I
  do not have that hardware, so I have not run a GPU regression myself. The fix is
  a one-line indexing change, and the reporter verified it end to end in IMA crash, with the probe repro #6830
  (GLM-5.2 FP8 training, output bitwise-deterministic, max abs diff 4.2e-04 vs a
  pure-PyTorch dequant reference). A CI/maintainer run on GPU would fully confirm.

[gemini-code-assist]

Code Review

This pull request updates the weight_dequant_kernel in unsloth/kernels/fp8.py to cast offsets to 64-bit integers (tl.int64) before computing the flat index. This prevents integer overflow when dealing with large tensors containing more than 2**31 elements, such as flattened MoE expert stacks. There are no review comments, and I have no additional feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[danielhanchen]

@codex review

[danielhanchen]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Can't wait for the next one!

Reviewed commit: 0d7fa30b8a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[danielhanchen]

Reproduced the int32 offset overflow (illegal memory access above 2^31 elements) on a large tensor and confirmed the int64 cast resolves it with numerically correct dequant output. Thanks.