6 ms·
Highly efficient matrix transpose in Mojo
- sestep 1y agoI'm not an expert in this space, but is this meaningful? I'd assume that it's more common to fuse together transposition with an operation that precedes or follows it (e.g. matmul), which should be far more efficient than materializing the entire transposition in memory if it's just an intermediate value.
- musebox35 1y agoMatrix transpose is a canonical example of a memory bound operation and often used to showcase optimization in a particular programming language or library. See for example the cutlass matrix transpose tutorial from Jay Shah of flash attention 3 paper: https://research.colfax-intl.com/tutorial-matrix-transpose-in-cutlass/ https://research.colfax-intl.com/tutorial-matrix-transpose-i...
- saagarjha 1y agoUnfortunately the issue (alluded to in the blog post you linked) is that transposes do absolutely no work but memory loads. Sure, they test that you can swizzle your accesses, but modern accelerators are all about pipelining and feeding matrix multiply units, which is considerably harder than loading from memory as fast as possible. Actually, even the Mojo post barely beats CUDA for most of its kernels, because you can hit memory bandwidth for transpose on the latest hardware using techniques from 5-10 years ago. This is definitely not true for more interesting operations.
- musebox35 1y agoI totally agree that the resulting kernel will be rarely useful. I just wanted to highlight that it is a commonly used educational exercise to showcase how to optimize for memory throughput. If the post showed how to fuse a transpose + rmsnorm epilogue to a gemm then the kernel would be more functional but the blog post would be much harder to follow for newcomers. Jay Shah’s later articles contain examples that involve epilogue fusion. IMHO, understanding how to write an efficient transpose helps with following the more involved ones.
- saagarjha 1y agoIt's less that the result is kind of useless and more that hitting memory throughput on a simple algorithm like this is not very difficult. It takes a complex example to actually have trouble doing this.
- simon_vtr 1y agoThat was exactly my reason to write this blogpost and optimise transpose. It is a simple educational yet not trivial example to learn the basics.
- colesantiago 1y agoDoes anyone use Mojo in production at all or are even hiring for Mojo?
- melodyogonna 1y agoModular (the company behind Mojo) uses it in production. I imagine that if they have any clients then those also use Mojo in production - albeit indirectly - since all the GPU kernels used by Modular are written in Mojo.
- htrp 1y agoLeft unsaid, the 14% improvement in performance came at the cost of increasing dev time by 35%
- bravetraveler 1y agoReminds me of this, lol: > "From the moment I understood the weakness of my flesh, it disgusted me. I craved the strength and certainty of steel." 14% all the time vs 35% some of the time edit: Closing numbers are far less impressive than those buried in the middle of the post. Confusing; bye everyone
- vlan121 1y agoMojos compiler is closed source. Thats a big no-no
- dgurchenkov 1y agoI work on Mojo. The whole compiler, runtime etc. will get open sourced, most likely within a year. It is just a matter of time and us getting all the required work done. https://docs.modular.com/mojo/faq/#open-source https://docs.modular.com/mojo/faq/#open-source
- almostgotcaught 1y ago> runtime Are you talking about your libc equivalent or MAX?
- dgurchenkov 1y agoBoth. Mojo standard library is already open source. Mojo at the moment does not need a runtime (but if it ever needs one it'd get open sourced). My point was, Mojo as a whole, as a programming language & a reference implementation, will definitely get open sourced. MAX itself is a bigger beast to work with, and I am out of my depth to talk about it. I think it'll get open sourced as well, just the timeline might be different (shorter or longer, IDK).
- xiphias2 1y ago,,will get open sourced'' means closed source, parent wrote the same
- GeekyBear 1y agoChris Lattner (the CEO of Modular) was previously the technical lead behind the creation of LLVM, Clang and Swift, all of which were open sourced. He has a bit of a track record already.
- 1y ago
- voronar 1y agoMr. Mojo Risin'
- arjvik 1y agoWhere's the 14%? Looks like their final kernels show a 0.14% improvement of Mojo over the equivalent CUDA kernel?
- jsnell 1y agoThe "Switching to Mojo gave a 14% improvement over CUDA" title is editorialized, the original is "Highly efficient matrix transpose in Mojo". Also, the improvement is 0.14%, not 14% making the editorialized linkbait particularly egregious.
- atomicapple 1y agoI think the OP based the title off of "This kernel archives 1437.55 GB/s compared to the 1251.76 GB/s we get in CUDA" (14.8%) and not the final kernels for whatever reason
- deleted 1y ago[deleted]
- baal80spam 1y ago0.14% is within the limits of statistical error. So this is a nothing-"article".
- jsnell 1y agoI don't think that's fair. The article promised a highly efficient kernel and seems to have delivered exactly that, which isn't "nothing". My beef is entirely with the submitted title.
- deleted 1y ago[deleted]
- jebarker 1y agoYeah, it seems like the blog post is just meant to be an example of how to do something in Mojo and not a dunk on CUDA.
- timmyd 1y agoFWIW I didnt take the blog as a dunk on CUDA, just as an impressive outcome from the blog writer in Mojo. It's awesome to see this on Hopper - if it makes it go faster thats awesome.
- noracists 1y agoslop
- deleted 1y ago[deleted]
- londons_explore 1y agoWhy do we ever need to transpose a matrix? Isn't it better to simply combine the transposition with whatever next operation one wishes to do with the matrix?
- throwawayabcdef 1y agoThe next operation might need the data in column major order to read it fast. So you might have to transpose first. And these maybe be concurrent stages of a processing pipeline.
- viraptor 1y agoNow I'm curious, how many times do you have to fully read the matrix in GPU for the total impact of reading columns to be higher than one-off actual transpose and then sequential row reads? I know it depends on lots of things, I'm after a rough estimate.
- saagarjha 1y agoIt's quite rare. Usually problems are tiled anyway and you can amortize the cost of having data in the "wrong" layout by loading coalesced in whatever is the best layout for your data and then transposing inside your tile, which gives you access to much faster memory.
- stephencanon 1y agoThe one pure transpose case that does come up occasionally is an in-place non-square transpose, where there is a rich literature of very fussy algorithms. If someone managed to make any headway with compiler optimization there, I'd be interested.
- hogepodge 1y agoYou're right that a good graph compiler will do this for you. There still may be times, like if you're interfacing with another library, where you'll need to switch a matrix between row major or column major layouts.
- almostgotcaught 1y agoAs someone said below - you'd never write just a transpose kernel - it'll be fused into something else.
- saagarjha 1y agoLook the frontier AI companies need something other than reversing binary trees to give interview candidates
- almostgotcaught 1y agoNo one is going to ask this on an interview. Used to be matmul. These days it's FA.
- saagarjha 1y agoI think these would be rather difficult to fit in a standard 1-hour interview slot
- ByteDrifter 1y ago[dead]
- melodyogonna 1y agoI wonder if there is a reason for not using the high level abstractions provided by Modular
- saagarjha 1y agoMost interesting algorithms (e.g. with dynamic shapes, mixed computation) are typically better scheduled by hand.
- Q6T46nT668w6i3m 1y agoSure, but Modular’s mission was to provide abstractions to minimize these types of optimizations.
- totalperspectiv 1y agoI’d also add that Mojo is new, and people are still feeling it out by trying to 1:1 things with Cuda.
- saagarjha 1y ago> This kernel archives a bandwidth of 1056.08 GB/s which is faster than the 875.46 GB/s we archived using CUDA. I believe that to be the reason because we use the PTX api for TMA transfers in Mojo. I can't say for sure because I couldn't find the CUDA kernel but I kind of doubt this is true. You can hit memory bandwidth on Hopper without using TMA at all, which is mostly designed for accelerating asynchronous copies and reducing memory pressure. If all you are doing is a transpose you don't need any of this to go fast (though it might simplify your indexing code…?)
- simon_vtr 1y agoThe kernels I mention in CUDA use all the equivalent logic like the Mojo kernels. You can find them on my GitHub: https://github.com/simveit/effective_transpose https://github.com/simveit/effective_transpose You may want to provide a faster kernel on H100 via PR and I will merge after checking it’s faster.
- iandanforth 1y agoI'm probably just ignorant but shouldn't the graphic of the tiled transpose have the green vector column-oriented in the final matrix?
- somethingsome 1y agoThe colors are reading writing operations ;) You have global memory and shared memory, the global is slower. You read in rows in the global memory (faster than reading columns) You write in columns in the shared memory (slower than in rows, but the shared memory is fast, this is the transpose operation) You read in rows in the shared memory (very fast) You write in rows in the global memory (faster than writing in columns) The idea behind that tiling is to hide the slow part in a memory that is faster.
- daft_pink 1y agoI think Mojo’s lack of being a true open product and existing to drive profits at Modular has really held it back. It’s just really impractical to use a licensed programming language in 2025.
- deleted 1y ago[deleted]
- totalperspectiv 1y agoMy impression is that this is on purpose on their part. They’ve repeatedly stated that by 2026 they will open source the compiler, and I think they’ve wanted a slow adoption ramp in order to spend some more time getting it right first. Possibly rose-tinted glasses on my part, but I’m optimistic for 2026. Chris Lattner has a pretty strong track record of getting these things right.
- veidr 1y agoYeah, and he's clearly trying to avoid what happened to Swift[1]. Although the danger of "corporate owner priorities dictate releasing half-baked/awful changes" risk is still there, Lattner himself has more influence within Modular (obviously, as co-founder and CEO) than he did at Apple, so it may work out better this time. [1]: https://news.ycombinator.com/item?id=30416070 https://news.ycombinator.com/item?id=30416070
- melodyogonna 1y agoYeah, Mojo's development has been pretty transparent. Chris publishes technical documents for most features and takes community feedback into account. A recent example is here: https://forum.modular.com/t/variable-bindings-proposal-discussion/1579 https://forum.modular.com/t/variable-bindings-proposal-discu... Btw, Mojo's development is a masterclass in language development and community building, it's been fun watching Chris go back to fix technical debts in existing features rather than proceeding with adding new features.
- GeekyBear 1y ago
- thunkingdeep 1y agoIs the word archive used in place of achieve? I’m not sure if there is a terminology issue that I don’t understand in this post…
- totalperspectiv 1y agoIn the coarse graining code, you use an @parameter-for. Doesn’t that lead to some pretty large code size unrolling that? Or is that less of an issue on GPU? Great write up! I learned a lot!
- simon_vtr 1y agoIt doesn’t. The batch size is just 8. This is a very good trick and often needed to archive peak performance in memory bound kernels. You can checkout the equivalent code in cuda aswell :)
- graycat 1y agoFast matrix transpose? Agree for a transposed matrix, just change the indexing arithmetic that converts row i and column j to an offset in the storage for the matrix and then remember that this is a transposed matrix. Some software object semantics could make this easy for other software to use.
- jjtheblunt 1y agoi think the problem with changing the indexing arithmetic is that you could end up with arithmetic incompatible with vector instructions in hardware that you're hoping to use for parallelism.
- graycat 1y ago> vector instructions Gee, for the polar decomposition, Gauss-Seidel, etc., looked really hard for those in my IBM PC/XT and couldn't find any!!!
- olaf 1y agoAFAIK ‚let‘ was removed from the language, for me it‘s a big turn-off that the Python compatibility aspect has such a high priority. Or did I overlook something?