One multiply, one RTX 5090, eight kernels

A hand-written CUDA kernel for a 4096-by-4096 linear layer with a batch of 16 vectors ran in 29 microseconds on an NVIDIA GeForce RTX 5090, down from 748 microseconds for the first naive kernel, according to a lab benchmark built for the article and timed with cold weights as the median of 500 samples of 100 consecutive launches each. The progression across eight kernels cut time to 324 microseconds after moving data through shared memory and registers, 244 with FP16 inputs without tensor cores, 155 with the tensor-core kernel, and 47 with a background copy, while a direct cuBLAS call took 14.40 microseconds with warm weights and PyTorch 2.14.1+cu130 took 14.47 microseconds on the host clock. The gap shows how far hand-written kernels remain from NVIDIA's cuBLAS library on this multiply, which reads a large weight matrix and does little arithmetic per value.

A linear layer in PyTorch, given a small batch of vectors, runs on the GPU as one matrix multiply. This article follows one such multiply through eight CUDA kernels on one RTX 5090, from 748 microseconds for the first kernel to 29 for the last, and looks at what changes inside the GPU with each kernel. Every time on the page was measured on that card, and every machine instruction shown is copied from the compiled kernels. The multiply the-multiply The layer followed here has 4096 inputs and 4096 outputs, and the batch is 16 vectors. The input X holds the batch, one row per vector: 16 rows of 4096 values. The weights W are stored the way PyTorch stores them, as out in : 4096 rows, one per output feature, each 4096 values long. Each value of the output Y is the dot product of one row of X with one row of W , so Y has 16 rows of 4096 values. Each value of W takes part in only 16 of those products, one for each row of X , so the multiply reads a large matrix and does little arithmetic with each value it reads. The measurements come from a lab written for this article: CUDA kernels for this multiply and the programs that time them, run on one RTX 5090. The card reports 170 streaming multiprocessors SMs , the units that run a kernel's threads. In PyTorch the layer is one call to F.linear . The lab timed PyTorch with warm weights only, ones the GPU has just read, and not with cold weights, ones it has not. In FP16, PyTorch 2.14.1+cu130 took 14.47 microseconds on the host clock, a time that includes the overhead of Python itself. A direct call to cuBLAS, NVIDIA's library of matrix routines, took 14.40 with warm weights, timed on the GPU with CUDA events as every kernel was. The two times differ by less than these measurements can resolve. The direct call was also timed with cold weights, so it is the reference the kernels are compared with. The question for the rest of the article is how close hand-written kernels come to cuBLAS on this multiply, what each kernel on the way changes inside the GPU, and what the time of the fastest kernels is spent on. The kernels are numbered Kernel 1 to Kernel 8. Basic tiling is covered in CUDA matrix multiplication optimization https://www.abhik.ai/articles/cuda-matrix-multiplication-optimization and is not explained again here; this article goes on from it to tensor cores and to how the work is spread across the GPU. From 748 to 29 microseconds from--to--microseconds Each kernel was timed with cold and with warm weights. Cold times are the headline numbers here and warm times are shown beside them; Where the time goes where-the-time-goes explains how each was measured. Times on this GPU fall on steps of about 2 microseconds, so a difference smaller than one step is not a result. The steps belong to each launch rather than to the timer, and their cause is not known. The figure sets the cold and warm time of each kernel on a logarithmic scale, beside a few lines of that kernel's machine code. Pressing Play, or Next, adds one kernel at a time. The horizontal axis is logarithmic, so equal distances are equal ratios. Times fall on steps of about 2.0 µs, so differences smaller than one step are not resolved. Each cuBLAS FP16 marker is drawn only beside the bars it compares with. Recorded output. NVIDIA GeForce RTX 5090, driver 580.173.02, Build cuda 13.0.r13.0/compiler.36424714 0, compiled with -O3 -std=c++17 -arch=sm 120. Each time is the median of 500 samples, and each sample is 100 consecutive launches. Kernels appear one at a time. Cold means the weights were not read just before the launch. In the complete figure, moving the data through shared memory and then through registers Kernels 2 and 3 took the time from 748 to 324 microseconds, and FP16 inputs without tensor cores Kernel 4 took it to 244. The tensor-core kernel Kernel 5 reached 155, and the background copy of Kernel 6 reached 47. Splitting the inner dimension across more blocks Kernel 7 reached 29. Kernel 8 is the result of a sweep over tile, split and K-slice sizes, terms the next two sections explain. Its cold time is 28.68 microseconds against 28.69 for Kernel 7, and its warm time is 12.30 against 12.31. Those pairs are the same within one timing step, so tuning did not change the result, and the figure shows Kernel 8 as a kernel that gained nothing. The dashed and dotted lines in the figure are cuBLAS FP16 with cold and with warm weights; the comparison with it comes later, under Against cuBLAS FP16 against-cublas-fp16 . The bars show how much each kernel gained, not what changed inside the GPU to bring it about. The next section looks inside one SM at Kernel 6, and the section after it looks at the whole chip, where Kernel 7 changes how many blocks there are. Inside one SM inside-one-sm This section follows Kernel 6, whose code Kernels 7 and 8 reuse with other tile and split sizes. Kernel 6 multiplies on the tensor cores, as Kernel 5 does, copies each piece of its inputs into shared memory without passing it through registers, and requests the next piece before it uses the current one. The section first shows how the multiply is cut into blocks and where each block runs, then what one block does on its SM. Tiles, blocks and SMs tiles-blocks-and-sms A kernel does not compute the output in one piece. It cuts the work into tiles. W has 4096 rows, one per output feature; in Kernel 6 a tile is a group of 128 of those rows, so the 4096 rows make 32 tiles. Each tile produces its own 128 columns of the output Y , and no tile needs anything from another. Kernel 6 gives each tile one thread block, a group of threads that cooperate, so a launch of this multiply is 32 blocks. The GPU places each block on an SM, which holds the block's threads and its shared memory while they run. Filling the chip filling-the-chip sets that count against the SMs of the chip. The figure draws 8 of the 32 tiles and 8 of the 170 SMs. Stepping through it places each tile's block on an SM. 8 of 32 tiles and 8 of 170 SMs are shown. Each tile is 128 rows of W, the weights stored as out in , and each block is placed on an SM. The lab did not record which SM ran which block, so the pairing drawn here is only to show the idea. Authored illustration. The tile and block counts are those of Kernel 6 and follow from the matrix size. The SM count of 170 is a recorded value reported by the GPU. Motion does not model execution time. Inside one block inside-one-block A block does not hold its whole tile of W at once. It walks along the inner dimension of the multiply one K-slice at a time: a K-slice is a strip of W and the matching strip of X , small enough to sit in the SM's shared memory, a small store that every thread in the block can read. In Kernel 6 a K-slice is 32 columns wide, so the 4096 columns of the inner dimension make 128 K-slices, and the loop of the kernel makes one pass for each. Kernel 6 keeps two buffers in shared memory. Every pass but the last first requests the next K-slice into one buffer, with a copy that does not block the warp. Each pass then waits until the current K-slice has arrived in its buffer, and every warp of the block reaches a barrier. Only then does it multiply that K-slice on the tensor cores, the units that run the matrix-multiply instructions. A second barrier ends the pass. The figure shortens the loop to three K-slices and follows one block of Kernel 6 through them, one step for each thing the code does, in the order the code issues the work rather than in measured time. Beside the stage is the loop's machine code, copied from the compiled kernel. Selecting a line of code opens it below the stage, starting with HMMA, the tensor-core multiply. Order of work in one Kernel 6 block Kernel 6, its loop shortened to three K-slices: the order in which one block requests each K-slice, waits for it and multiplies it, and the machine instructions behind each step. Select a line to open it. Inner loop of Kernel 6, as compiled Before the loop, request K-slice 0 into buffer 0. The copy does not block the warp. HMMA.16816.F32 R12, R20, R26, R12 One tensor-core multiply: A, a 16 x 16 piece of the K-slice of X, times B, a 16 x 8 piece of the K-slice of W, added into D, a 16 x 8 running sum. The 32 threads of a warp issue it together. Paired source line PTX mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3}; 32 threads of the warp Thread t0 holds 8 cells of A, 4 of B and 4 of D. No thread holds a whole matrix. Choose a thread, or fire the instruction. Source: NVIDIA PTX documentation: matrix fragments for mma.m16n8k16 https://docs.nvidia.com/cuda/parallel-thread-execution/index.html warp-level-matrix-fragment-mma-16816-float Authored illustration. The stage and timeline show the order in which the Kernel 6 code issues work, not when the hardware performs it, and their widths do not model GPU execution time. The loop is shortened to three K-slices. The instruction lines are copied from the loop body in the compiled output of Kernel 6 Build cuda 13.0.r13.0/compiler.36424714 0, compiled with -O3 -std=c++17 -arch=sm 120 . A pass runs only one of the two DEPBAR lines: 0x1 after it requests the next K-slice, and 0x0 on the last pass. Every pass but the last requests the next K-slice before it waits for the current one, so when the wait ends, at most one request is still outstanding. The order alone does not say whether a copy and a multiply overlap in time. That is a question for measurement, and the figure cannot answer it. The lab timed two cut-down forms of Kernel 6, called probes: a loads-only probe that loads every K-slice but multiplies nothing, and a math-only probe that multiplies but loads nothing after the first K-slice. With cold weights the loads-only probe took 37.03 microseconds, the math-only probe took 22.55, and the full kernel took 47.21. So the full kernel is 10.18 microseconds above the loads-only probe, less than the 22.55 of the math-only probe by more than one timing step, which is consistent with overlap of loading and arithmetic. That reading applies to Kernel 6 with cold weights only. Where the time goes where-the-time-goes returns to the probes for the other cases, after the next section turns from one block to the whole launch. Filling the chip filling-the-chip Kernel 6 launches the 32 blocks counted in Tiles, blocks and SMs tiles-blocks-and-sms , for the 170 SMs the GPU reports. Each block is placed on one SM, so at most 32 SMs hold a block of Kernel 6. Kernel 7 keeps the same tiles and also splits the inner dimension: the 4096 columns of each tile are cut into 8 ranges of 512, and each range gets its own block, which multiplies its range and adds its partial sums into the output. That makes 256 blocks, more than there are SMs. The figure sets the blocks of each launch against the SMs of the chip. Choosing a launch shows its block count against the SM count, with its cold and warm times. Authored illustration. The number of lit cells is the block count or the SM count, whichever is smaller. Which cells are lit is an illustration; the lab did not observe which SM ran which block. Recorded output. NVIDIA GeForce RTX 5090, driver 580.173.02, Build cuda 13.0.r13.0/compiler.36424714 0, compiled with -O3 -std=c++17 -arch=sm 120. Each time is the median of 500 samples, and each sample is 100 consecutive launches. The cold and warm times are those of the chosen kernel, and the SM count is the one the GPU reports. Kernel 7 gets its extra blocks by splitting the inner dimension, but smaller tiles would also give more blocks. The sweep of tile, split and K-slice sizes was run with warm weights only, and each of its times is the median of 200 samples of 100 launches, fewer samples than the 500 behind the figures. It includes a pair of settings that launch the same 128 blocks but differ in tile size as well as in split. One cuts W into 128 tiles of 32 rows and does not split the inner dimension; it took 38.96 microseconds. The other has 32 tiles of 128 rows and splits the inner dimension of each into 4 ranges; it took 14.37. The setting that splits is the faster of the two, but because the pair also differs in tile size, it does not isolate the split. Against Kernel 6, Kernel 7 took 28.69 microseconds instead of 47.21 with cold weights and 12.31 instead of 43.06 with warm weights: 1.65 times as fast cold and 3.50 times as fast warm. The split gained less when the weights were cold. The next section looks at what a cold launch spends its time on. Where the time goes where-the-time-goes Every time so far comes in two forms, and here the difference matters. A warm launch rereads the same weights as the launch before it. A cold launch reads weights the GPU has not just read: the lab rotates through several identical copies of the weights, so each launch reads a different copy from the one before. Cold is the closer of the two to a real model, where each layer has its own weights and reads them once per forward pass. The probes from Inside one block inside-one-block were run for Kernel 7 as well as Kernel 6, each with warm and with cold weights. The math-only probe multiplies its one K-slice as many times as the full kernel multiplies K-slices, so it does the same arithmetic without reading the rest of W . If the full kernel takes as long as the loads-only probe, the multiply adds no measurable time to reading the weights. The figure sets the full kernel beside its two probes on one scale, opening on Kernel 7 with cold weights. Choosing another case shows its three times. In this measurement the full kernel takes about as long as loading alone, so the tensor-core math adds no measurable time on top of the memory reads. Both probes are the full kernel with one part removed, and both still include the launch, the barriers and the waits, so their sum overstates the total. Recorded output. NVIDIA GeForce RTX 5090, driver 580.173.02, Build cuda 13.0.r13.0/compiler.36424714 0, compiled with -O3 -std=c++17 -arch=sm 120. Each time is the median of 500 samples, and each sample is 100 consecutive launches. Each bar is its own measurement, and all three use one linear scale for every case. For Kernel 7 with cold weights the full kernel took 28.69 microseconds, the loads-only probe 28.36 and the math-only probe 10.25. The full kernel and the loads-only probe are the same within one timing step, so in this measurement the tensor-core math adds no measurable time on top of loading. With warm weights the full Kernel 7 kernel is about one timing step above either probe, but each of those launches also includes the zeroing of the output and the launch itself, so nothing is concluded from that case. Kernel 6 with warm weights does give a reading: its full kernel costs about the sum of its loads-only and math-only probes, so no overlap is visible in that case. Against cuBLAS FP16 against-cublas-fp16 The multiply the-multiply asked how close hand-written kernels come to cuBLAS on this multiply, whose FP16 times are the dashed and dotted lines in the figure Eight kernels, one multiply . The answer has three parts. With cold weights, cuBLAS FP16 is two timing steps faster: 24.65 microseconds against 28.68 for Kernel 8. With warm weights, Kernel 8 is one timing step ahead, which is the resolution of the method: 12.30 against 14.40. And the two are not the same computation: cuBLAS FP16 accumulates and returns FP16, while these kernels accumulate and return FP32. This answer, like every result above, comes from one card and one shape of multiply, and the lab left some questions open. The last section lists them. What was not measured what-was-not-measured Each item below is something the lab did not measure, followed by what the page therefore does not say. - The cause of the timing step. The lab did not find why times fall on steps of about 2 microseconds, or whether the 1.96 microseconds that an empty kernel takes is part of every time, so every time on the page is accurate to one step and none has launch overhead subtracted. - A sweep on cold weights. The sweep of tile, split and K-slice sizes ran with warm weights only, and Kernel 8 was chosen on that warm data, so the page does not say which setting is fastest with cold weights. - Other GPUs, batch sizes and matrix sizes. Every time is for one RTX 5090, a batch of 16 and a weight matrix of 4096 by 4096, so the page does not say how any result carries over to another card or shape. - Kernel 6's two changes, one at a time. Kernel 6 copies without passing through registers and also requests the next K-slice before it uses the current one, and the lab timed the two only together, so the page does not say how much each one gained. - PyTorch with cold weights. The lab timed PyTorch with warm weights only and on its own random inputs, so the page sets it only beside the warm time of the direct cuBLAS call. - Block placement. The lab did not record which SM ran which block, so the figures that place blocks on SMs are illustrations, and the page does not say which SMs held a block. - The split on its own. The pair of sweep settings with 128 blocks differs in tile size as well as in split, so that pair does not isolate the split, and the page does not say how much of its difference the split alone made.