5 ms·
Good point. But the overall point about Mojo availing a different level of abstraction as compared to Python still stands: I imagine that no amount of magic/op
by davidatbu 1y ago
Good point. But the overall point about Mojo availing a different level of abstraction as compared to Python still stands: I imagine that no amount of magic/operator-fusion/etc in `torch.compile()` would let one get reasonable performance for an implementation of, say, flash-attn. One would have to use CUDA/Triton/Mojo/etc.
- boroboro4 1y agoBut python is already operating fully on different level of abstraction - you mention triton yourself, and there is new python cuda api too (the one similar to triton). More to this - flash attention 4 is actually written in python. Somehow python managed to be both high level and low level language for GPUs…
- davidatbu 1y agoIIUC, triton uses Python syntax, but it has a separate compiler (which is kinda what Mojo is doing, except Mojo's syntax is a superset of Python's, instead of a subset, like Triton). I think it's fair to describe it as a different language (otherwise, we'd also have to describe Mojo also as "Python"). Triton's website and repo describes itself as "the Triton language and compiler" (as opposed to, I dunno, "Write GPU kernels in Python"). Also, flash attention is at v3-beta right now? [0] And it requires one of CUDA/Triton/ROCm? [0] https://github.com/Dao-AILab/flash-attention https://github.com/Dao-AILab/flash-attention But maybe I'm out of the loop? Where do you see that flash attention 4 is written in Python?
- boroboro4 1y agoFrom this perspective PyTorch is separate language, at least as soon as you start using torch.compile (only subset of PyTorch python will be compilable). That’s strength of python - it’s great for describing things and later for analyzing them (and compiling, for example). Just to be clear here - you use triton from plain python, it runs compilation inside. Just like I’m pretty sure not all mojo can be used to write kernels? I might be wrong here, but it would be very hard to fit general purpose code into kernels (and to be frank pointless, constrains bring speed). As for flash attention there was a leak: https://www.reddit.com/r/LocalLLaMA/comments/1mt9htu/flashattention_4_leak/ https://www.reddit.com/r/LocalLLaMA/comments/1mt9htu/flashat...