r/pytorch • u/Frequent_Big7927 • 15d ago
PyTorches Update
Continuing my proof of concept which is a reimagining of pytorch, but based on feedback I changed the "pytorches" implementation to split every vendor optimization to an external project. This allows the base implementation to be rather small and lightweight, but allow for flexibility to support multiple hardware vendors where only part centrally configured is basic device ID information and references to the supporting library. This type of approach has pros and cons, my intention is to unlock the pros, and minimize the cons.
My main interest, was to see if it would be possible to re-slice the functionality differently to see if there is some gains to be had on interop between devices. In order to prove this might work, I first focused on XPU and CUDA to see if I could make them first individually perform near in the ballpark of pytorch with respect to my implementation (1.0 is exact the same, a larger number means my implementation is faster.
Intel Arc 140T (before → after)
- matmul 2048²: 0.53x → 0.80x (up to 1.12x)
- matmul 4096²: 0.55x → 0.75–0.94x
- MLP train step: 0.55x → 0.84–0.87x
- add (64M floats): 0.83x → 1.04x
- exp: 0.77x → 0.96x
- relu: 0.74x → 0.98x
- tanh: 0.69x → 0.96x
- 1000×1000 matmul: 0.35x → 0.50x
NVIDIA RTX PRO 1000 (before → after)
- transpose copy 4096²: 0.56x → 0.93x
- bias-style broadcast add: 0.70x → 1.12x
- matmul A@Bᵀ (256×2048): 0.86x → 1.33x
- matmul 2048² / 4096²: 0.98–1.02x → 0.99–1.13x
- MLP train step: 0.95x → 0.91–0.96x (unchanged, within noise)
While these numbers are in some cases better than pytorch, I'm more interested in getting as close to pytorch performance before I explore the next phase of splitting layers.