Skip to content

Fix multi-device mxfp4 dequantization race in _convert_moe_packed_tensors #47423

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
IlyasMoutawwakil merged 2 commits into from
Jul 22, 2026
+17 −12
Merged

Fix multi-device mxfp4 dequantization race in _convert_moe_packed_tensors #47423

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6de8651
Fix XPU multi-device mxfp4 dequantization race in
kaixuanliu Jul 20, 2026
a2308f8
update comment
kaixuanliu Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 17 additions & 12 deletions src/transformers/integrations/mxfp4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,18 +288,23 @@ def _convert_moe_packed_tensors(
exp = scales[r0:r1]
sub = out[r0:r1]

# This vector is only used to index into `lut`, but is hugeee in GPU memory so we delete it immediately
idx_lo = (blk & 0x0F).to(torch.int)
sub[:, 0::2] = lut[idx_lo]
del idx_lo

# This vector is only used to index into `lut`, but is hugeee in GPU memory so we delete it immediately
idx_hi = (blk >> 4).to(torch.int)
sub[:, 1::2] = lut[idx_hi]
del idx_hi

# Perform op
torch.ldexp(sub, exp, out=sub)
# With device_map="auto", tensors sitting on a non-current accelerator device are not
# ordered after their async H2D copy, so the compute below may read garbage and emit
# out-of-bounds `lut` indices (illegal memory access on CUDA, indexing abort on XPU).
# Aligning the active device with the tensor's device orders it correctly (no-op on CPU).
with on_device(blk.device):
# This vector is only used to index into `lut`, but is hugeee in GPU memory so we delete it immediately
idx_lo = (blk & 0x0F).to(torch.int)
sub[:, 0::2] = lut[idx_lo]
del idx_lo
Comment on lines +291 to +299

@IlyasMoutawwakil IlyasMoutawwakil Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

are you sure this is what's happening ? wdym by "tensors .. are not ordered" ? blk and lut are already on the same device at this point, the device context changes where new tensors with no device are created, like a torch.zeros or arange, why would it fix this ?

@kaixuanliu kaixuanliu Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, Good Questions. Let me update more background:
Here are facts from my side:

  1. With device_map="auto" across multiple accelerators, loading gpt-oss-20b aborts inside this loop: on XPU as an index out of bounds in the indexing kernel, on CUDA as an illegal memory access.
  2. During debugging, it shows that for a blk living on a non-current device (e.g. blk on xpu:1/cuda:1 while the process' current device is device 0), idx_lo = blk & 0x0F transiently produces values far outside [0, 15] (I saw 0–255 then garbage like [-2147483648, 2139160448]), which then indexes lut out of bounds. For tensors on the current device it is always correct.
  3. Wrapping only this compute in on_device(blk.device) makes the corruption disappear and load+generation succeed, on both XPU and CUDA.

By "not ordered" I mean a cross-device stream ordering gap, not tensor allocation. The loader copies blk H2D on a worker thread using its own device's stream, but the elementwise kernels here are launched while the process' current device is still device 0. Since kernel launch / default-stream selection follows the current device, these kernels aren't guaranteed to be ordered after the H2D copy that fills blk on device 1 — so they can read not-yet-valid memory. Setting the current device to blk.device makes them launch on that device's stream, which is properly ordered after the copy.


# This vector is only used to index into `lut`, but is hugeee in GPU memory so we delete it immediately
idx_hi = (blk >> 4).to(torch.int)
sub[:, 1::2] = lut[idx_hi]
del idx_hi

# Perform op
torch.ldexp(sub, exp, out=sub)
del blk, exp, sub

out = out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2)
Expand Down
Loading