Plain loops at the speed of CUDA

Aditya Kumar · 13 November 2026

TL;DR — In Vx, a GPU kernel is two nested for loops: one across blocks and one across lanes. We timed such loops on an NVIDIA A100 against CUDA. A loop that scales a 4096 by 16384 grid ran in 0.343 ms, and the same loop in CUDA took 0.388 ms. Three kernels from llm.c, Andrej Karpathy's GPT-2 in C and CUDA, written as across loops, came within 5% of llm.c's own CUDA, and the softmax was faster. Getting there took two things a program says, lanes::sum and fast_tanh, and one compiler change. Every number here is from one A100, and we timed some in different ways. The last section says which.

The promise of across loops is that you write the loop you mean, and the compiler makes the GPU kernel. A promise like that is only worth something if the kernel is fast. So we rented an A100 and compared Vx with the CUDA a person would write.

One loop

This is the first across example in the book, on a larger grid. Each block takes a row, and 256 lanes go through its 16,384 columns:

spawn on(Topology::Device) {
  for y in 0..4096 across blocks {
    for x in 0..16384 across lanes(256) {
      grid[y][x] = grid[y][x] * 2.0 + (y as f32);
    }
  }
}

The same loop in CUDA, written the usual way, is one block per row, with each thread going through the columns in steps of the block size:

__global__ void scale_rows(float *grid) {
  int y = blockIdx.x;
  for (int x = threadIdx.x; x < COLS; x += blockDim.x) {
    size_t i = (size_t)y * COLS + x;
    grid[i] = grid[i] * 2.0f + (float)y;
  }
}
// scale_rows<<<4096, 256>>>(grid);

The loop reads and writes 512 MiB, so its speed is the speed of the memory. Vx was 12% faster because it knows something the CUDA kernel does not: the block size. The CUDA kernel reads blockDim.x when it runs. Vx launches every kernel with the block size the program gave, so the compiler puts the number 256 into the GPU code. A CUDA programmer can get the same by writing 256 in place of blockDim.x. We did not time that version.

Three kernels from llm.c

A loop that only scales numbers is a low bar. llm.c is a better one. It has several CUDA versions of each kernel of GPT-2, from simple to fast, and a harness that times them. We wrote three of them in Vx: GELU, layernorm and softmax, on llm.c's sizes. This is GELU in llm.c's first version:

__global__ void gelu_forward_kernel1(floatX* out, const floatX* inp, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        float xi = inp[i];
        float cube = 0.044715f * xi * xi * xi;
        out[i] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)));
    }
}

And in Vx:

pub fn gelu_forward(out : &mut Tensor<f32, [?], Memory::HBM>, inp : &Tensor<f32, [?], Memory::HBM>, n : i32) -> void {
  spawn on(Topology::Device) {
    for b in 0..(n + 255) / 256 across blocks {
      for l in 0..256 across lanes {
        let i = b * 256 + l;
        if i < n {
          let x = inp[i];
          let cube = 0.044715 * x * x * x;
          out[i] = 0.5 * x * (1.0 + (0.7978846 * (x + cube)).fast_tanh());
        }
      }
    }
  }
  return;
}

The Vx function is also an ordinary function. The compiler checks that two passes of a loop cannot write the same element, and the same source runs on the CPU, one lane at a time, when there is no GPU.

The times, on an A100, for the fastest version of each kernel in llm.c:

Each Vx result matched a CPU computation of the same function. The llm.c harness checks its own.

What it took

The first versions were not this fast. Three changes closed the gap, and two of them are in the program, not the compiler.

lanes::sum and lanes::max. In layernorm and softmax, each lane computes a part, and then the parts become one number for the row. Our first versions combined the 256 parts in a plain loop, which one lane runs alone. llm.c uses a tree in shared memory, and so does lanes::sum. Changing only those lines took layernorm from 0.088 to 0.043 ms, and softmax, with the next change too, from 0.130 to 0.050 ms.

fast_tanh and fast_exp. Vx's tanh and exp are exact: an f32 result is the float nearest to the true value, and the CPU and the GPU give the same bits. To round correctly, they compute in f64, and a GPU does f64 arithmetic more slowly. fast_tanh and fast_exp use a few f32 instructions on the GPU, as CUDA's do. GELU went from 0.060 to 0.047 ms. The new book chapter Exact and fast math says when to use which.

The block size in the GPU code. This is the compiler change (#3354). A lane loop steps by the block size, which the GPU code read when it ran. So in GELU, where the lane loop covers exactly one block, LLVM did not know that the loop runs once. It kept the loop, and two more instructions for each load and store. Vx always launches a kernel with the block size the program gave, so the compiler now writes that number into the GPU code. Every kernel we timed got 4% to 10% faster, and gave the same results:

The change is on main, and it will be in the next release. Vx 0.0.4 gives the "before" numbers.

How we measured

A comparison is only as good as its method, so here is ours, with its limits.

• The machine. One A100-SXM4-80GB, CUDA 12.8, driver 595.91. Each kernel ran on one size, and the GPU had no other work.
• Single launches. With VX_TIME_KERNEL=1, Vx times each launch on its own with CUDA events. The tables give the median of 50 launches, or 10 for the loop of the first section. We timed the CUDA version of that loop in the same way.
• llm.c's own timing. llm.c runs a kernel many times in a row and gives the average. On this A100, that gave about 4 µs less per launch than timing a single launch. So the layernorm and softmax comparisons favour llm.c. For GELU, we timed llm.c's kernel as Vx times its own: the 0.046 ms in the table.
• 1000 launches in a row. To check that GELU really is level, we loaded the GPU code from Vx and the code from nvcc into one program, and ran each 1000 times in a row, 7 times over. The spread was under 1%. Vx was 2% slower.
• Fast math. The llm.c numbers are without --use_fast_math, except where the table says so. With it, llm.c was up to 9% faster. Vx uses fast functions only where the program calls them.
• Caches. layernorm's input is 25 MB, which fits in the A100's 40 MB L2 cache. That may help a version that reads its input three times, as ours does.

What is left

• Half-precision math. An across loop cannot call tanh or fast_tanh on f16 or bf16 yet (#3353). llm.c's fast versions use bf16.
• Other GPUs. These are numbers from one A100. A GPU with slower f64, such as an L4, will show a larger gap between the exact and the fast functions.
• More kernels. Matrix products in llm.c use cuBLAS, and attention uses more of the GPU than a lane loop shows. Those are the next comparisons.

The measurements that found the compiler change are in #3352. The book's Writing parallel code chapter shows how to write across loops and lanes::.

Vx is Apache 2.0 with the LLVM exception — install it and try the examples.