diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 76ba3949762f..22230f7393ca 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -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 + + # 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)