Intelligence
Batch Size and MoE Routing
Changing only the batch size changes which experts a Mixture-of-Experts model selects, and differently on AMD RDNA3 and RDNA4.
Mixture-of-Experts models use a router to pick which experts process each token. I found that changing only the batch size changes the floating-point results enough to change which experts get selected, and that this happens differently on AMD RDNA3 and RDNA4 GPUs.
Different batch sizes lead PyTorch to use different GEMM kernels, which add numbers in a different order. Since floating-point addition isn’t associative, the results drift slightly, and that’s enough to flip close routing decisions. When I forced batch size 1 to use the kernel shapes PyTorch chooses for batch size 10, the differences became minimal.
This matters for reproducibility, and possibly security, since how your input gets routed can depend on what else is in the batch. I’m still working on it, updates soon!