7 ms·
Optimizing a WebGPU Matmul Kernel for 1 TFLOP
- billconan 2y agoWebGPU doesn't seem to talk about bank conflict, hiding some hardware details that might be necessary to write the best kernel. will it be able to match the perf of Cuda on the same hardware?
- zanussbaum 2y agogreat question, to me webGPU sits a hair high level than CUDA or Vulkan. so you don't have the exact same level of control but can get to 80% performance of it without having to write different kernels specific to the hardware
- brrrrrm 2y agoWebGPU cannot even come close unfortunately since they don't have support for hardware specific memory or warp-level primitives (like TMA or tensorcores). it's not like it gets 80% of perf, it gets < 30% of the peak perf for anything related to heavy compute matrix multiplications
- kayvr 2y agoI've heard the WebGPU workgroup wants to close the gap on tensor core support.
- zanussbaum 2y agoyou're definitely right, 80% was a bit of an overestimation, especially with respect to CUDA it would be cool to see if there's some way to get better access to those lower-level primitives but would be surprised it does seem like subgroup support are a step in the right direction though!
- Const-me 2y ago> don't have support for hardware specific memory I have no experience with WebGPU but if you mean group shared memory, I think the support is available. See the demo: https://compute.toys/view/25 https://compute.toys/view/25
- zanussbaum 2y agoi tried using workgroup shared memory and found it slower than just recomputing everything in each thread although i may have been doing something dumb i'm excited to try subgroups though: https://developer.chrome.com/blog/new-in-webgpu-128#experimenting_with_subgroups https://developer.chrome.com/blog/new-in-webgpu-128#experime...
- Const-me 2y agoCouple years ago, I wanted about the same thing in HLSL language, for a Direct3D 11.0 compute shader. Here’s the fastest version I managed to make back then: https://github.com/Const-me/Cgml/blob/master/Mistral/MistralShaders/mulMatTiled.hlsl https://github.com/Const-me/Cgml/blob/master/Mistral/Mistral... As you see, I have implemented 32×32 tiling, using thread groups of 32×8 threads, two groupshared buffers to load tiles of the input matrices, and I accumulate numbers into local variables, 32 / 8 = 4 accumulators per thread.
- deleted 2y ago[deleted]
- deleted 2y ago[deleted]
- deleted 2y ago[deleted]
- lostmsu 2y agoWhat's the perf like?
- Const-me 2y agoSorry, I have not benchmarked against cuBLAS or Eigen or similar, I did that thing for ML inference. I have implemented a profiler on top of D3D11_QUERY_TIMESTAMP and D3D11_QUERY_TIMESTAMP_DISJOINT queries, and tweaked the compute shader to minimize the time reported by these queries for my specific use case.
- shihab 2y agoGreat article! For context: this WebGPU version achieves ~17% of peak theoretical performance of M2. With CUDA (i.e. CuBLAS), you can reach ~75% of peak performance for same matrix config (without tensor core).
- zanussbaum 2y agothanks! and yes definitely not at CUDA levels :)
- weinzierl 2y ago75% can't be the best we can do. What would reach 100% or nearly 100%? Handcoded assembly?
- jsheard 2y agoWith GPUs it's not uncommon to run out of memory bandwidth before you max out the theoretical FLOPS. They may have a ton of bandwidth but it's never enough. That can lead you to some pretty counter-intuitive optimizations because it's often faster to do more compute work if it means you touch less memory in the process.
- wbl 2y agoShouldn't the roofline inform capacity assessments?
- KeplerBoy 2y agoSure, but rooflines don't account for stuff like memory granularity. You not only have to do a lot of bytes per flop to achieve the necessary arithmetic intensity, you also have to access those bytes in a coalesced way. I.e., you want to access consecutive bytes, which are ideally already in registers.
- stephencanon 2y agoFor sufficiently large GEMM you should never run out of bandwidth before you max out FLOPS if your blocking is organized correctly, because the arithmetic scales like O(n^3) while the memory access scales like O(n^2).
- maelito 2y agoWebGPU will make Web maps even more competitive than they are already. The smoothness of an iPhone map zoom, on any device.
- jsheard 2y ago> The smoothness of an iPhone map zoom, on any device. Any device except an iPhone, until Apple finally gets around to shipping WebGPU in Safari. Any year now...
- astlouis44 2y agoSafari is officially enabling support for WebGPU in iOS 18.2, which is rolling out within the first weeks of December.
- jsheard 2y agoWhere'd you hear that? It's not listed here: https://developer.apple.com/documentation/safari-release-notes/safari-18_2-release-notes https://developer.apple.com/documentation/safari-release-not...
- astlouis44 2y agoSource is here, from a Unity WebGPU thread. Look at the comment from October 27 from Brendan Duncan, a Unity employee: https://discussions.unity.com/t/early-access-to-the-new-webgpu-backend/933493/390 https://discussions.unity.com/t/early-access-to-the-new-webg... "I have found that WebGPU is enabled by default now with iOS 18.2. Apple has been working in the open on WebGPU. The WebKit source code has their latest WebGPU work in it. What hasn’t been known is their release schedule, but now with 18.2 it’s looking very promising that it will be on by default in that version."
- luketaylor 2y agoSource? Edit: I just pressed “Reset All to Defaults” under “WebKit Feature Flags” on my device running 18.2 beta, and the switch for WebGPU is on!! <3
- coffeeaddict1 2y agoYou can do slightly better fairly easily I think. See here for example https://github.com/AnswerDotAI/gpu.cpp/pull/35 https://github.com/AnswerDotAI/gpu.cpp/pull/35
- inglor 2y agoCan you explain why you did the naive algorithm here and not any of the fast matrix multiplication ones that trade multiplications for more additions? Just for educational purposes or is there a performance benefit in the technique?
- deleted 2y ago[deleted]
- zanussbaum 2y agoat least on my m2, the compiled kernel ends up using fast math anyways so using WGSL's fma didn't change anything about the actual kernel that gets run
- hedgehog 2y agoinglor is probably referring to Strassen or Coppersmith–Winograd.
- zanussbaum 2y agooh in that case it was because i didn't know about them :) something to try next!
- wbl 2y agoLast I checked the extra mems really hurt on a lot of cases especially for the more complex ones, but I'm no expert.
- saagarjha 2y agoBecause those algorithms are generally not worth implementing even though their algorithmic complexity is theoretically lower.
- pama 2y agoTo clarify the title: TFLOP/s is the unit the author goes after, not TFLOP. People in the threads compare CUDA performance on GPUs to WebAssembly performance: please recall that H100 has a theoretical performance of about 1000 TFLOP/s for bfloat16, and even moderately complicated algorithms in typical modern transformer architectures can reach about half of that performance.
- saagarjha 2y agoH100 can do well over 1500 TFLOPS in fp16.
- nulltype 2y agoWhich H100 and how much over 1500 TFLOP/s? The datasheet for the H100 SXM seems to indicate that it can only do ~1000 TFLOP/s peak.
- saagarjha 2y agoI just went to Nvidia’s site and downloaded the data sheet: https://resources.nvidia.com/en-us-tensor-core/nvidia-tensor-core-gpu-datasheet https://resources.nvidia.com/en-us-tensor-core/nvidia-tensor.... It says 1600/1900 in half precision?
- FL33TW00D 2y agoI wrote something similar a while back: https://github.com/FL33TW00D/wgpu-mm https://github.com/FL33TW00D/wgpu-mm Also does quantized matmuls.
- brrrrrm 2y agowould be fun to do a leaderboard of some specific size (e.g. 4096x4096x4096) just to get all the code and tricks in one spot for folks to learn about things
- mkeeter 2y agoFor a very deep dive into the subject, this is a great writeup: How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance (https://siboehm.com/articles/22/CUDA-MMM https://siboehm.com/articles/22/CUDA-MMM) (It's CUDA-specific, so there may be aspects that can't yet be ported to WGPU)
- zanussbaum 2y agothis was a huge inspiration for the post! i tried to highlight it in the blog but it might have gotten buried there are a few things that i wasn't able to figure out how to get access to/i wasn't sure if they were possible. for example, a lot of Simon's article takes advantage of the warp scheduler and warp tiling. i had a hard time finding information on if that's even possible with my M2/metal and the general memory access patterns. it seems like CUDA does have better documentation in this regard
- almostgotcaught 2y agoThat's a nice tutorial but just to be clear: that is not a deep dive in any sense. It's just the bog standard tricks. It doesn't cover MMA and WMMA, which today is table stakes for fast matmul. Also doesn't cover software pipelining. It's basically a good summary of the basics.
- saagarjha 2y agoIt’s a deep dive as of like 2015 probably. I don’t know if anyone has done something similar for modern GEMMs. Maybe the CUTLASS or Colfax people?
- jsbsjwbw 2y ago[flagged]