GLM on an A100, in Vx

TL;DR — The decoder of GLM-5, a mixture-of-experts model with about 750 billion parameters, now runs on an NVIDIA A100 through 28 GPU kernels written in Vx. No small trained version of the model exists, so every check uses random weights in the model's real shapes. Against the official Hugging Face implementation, the logits agree to within 1e-5 at every layer, up to 78 layers deep and 2,300 tokens long, and 16 tokens of greedy generation are identical. A decode path that keeps the weights on the GPU gives bit-identical logits at the model's full width, and runs INT4, FP8 and BF16 weights as they are stored. The GPU runs also found a bug the CPU tests never could: kernels read stale data when several processes shared the GPU. The full trained model has not run yet.

glm-vx is an experimental serving engine for GLM-5 whose compute kernels are written in Vx. It started on 4 October with CPU kernels. This post is about the GPU work tracked in glm-vx#1: what runs on the GPU, how we know it computes the right thing, and what the two A100 runs, on 15 and 20 October, found.

The model

GLM-5 (Hugging Face class GlmMoeDsaForCausalLM) has 78 decoder layers. Each layer has two parts that make it more than a plain transformer:

The hidden size is 6144, and the context length in the official config is 1,048,576 tokens. The official weights are FP8, about 756 GB. Community INT4 versions exist, in the compressed-tensors format.

What runs on the GPU

The GPU kernels are 1,406 lines of Vx in seven files: matrix-vector and matrix products in float32, TF32 on tensor cores, INT4, FP8 and BF16; RMSNorm, RoPE, SwiGLU and softmax; the sparse-attention indexer scores and its top-k; and MLA decode attention. Each is an ordinary Vx function, written per thread, with a few GPU intrinsics declared as extern. Here is the BF16 matrix-vector product, one warp per row:

extern "C" {
  fn vx_gpu_thread_idx_x() -> i32;
  fn vx_gpu_block_idx_x() -> i32;
  fn vx_gpu_block_dim_x() -> i32;
  fn vx_gpu_shuffle_down_f32(value: f32, delta: i32) -> f32;
}

// out[N] = W[N, K] x[K]. Weights are raw bfloat16 bits; the sum is float32.
pub fn glm_vx_gpu_bf16_matvec(out: *mut f32, weight: *const u16, x: *const f32, rows: i32, cols: i32) -> void {
  if rows <= 0 { return; }
  unsafe {
    let tid: i32 = vx_gpu_thread_idx_x();
    let lane: i32 = tid & 31;
    let warps: i32 = vx_gpu_block_dim_x() / 32;
    let row: i64 = (vx_gpu_block_idx_x() as i64) * (warps as i64) + ((tid / 32) as i64);
    if row < (rows as i64) {
      let k: i64 = cols as i64;
      let mut sum: f32 = 0.0;
      let mut col: i64 = lane as i64;
      while col < k {
        sum = sum + glm_vx_bf16_decode(weight[row * k + col] as i32) * x[col];
        col = col + 32;
      }
      sum = sum + vx_gpu_shuffle_down_f32(sum, 16);
      sum = sum + vx_gpu_shuffle_down_f32(sum, 8);
      sum = sum + vx_gpu_shuffle_down_f32(sum, 4);
      sum = sum + vx_gpu_shuffle_down_f32(sum, 2);
      sum = sum + vx_gpu_shuffle_down_f32(sum, 1);
      if lane == 0 { out[row] = sum; }
    }
    return;
  }
}

A build script turns these into PTX. vxc --emit-llvm gives MLIR's LLVM dialect, mlir-translate turns that into LLVM IR, the script retargets it to nvptx64 and defines the intrinsics (thread and block indices, warp shuffles, the TF32 mma.sync), and LLVM's opt and llc produce PTX for sm_80. Python loads the PTX through the CUDA driver API and launches the kernels. This is not Vx's own GPU path, spawn on(..), which does not have warp shuffles or tensor-core instructions yet. The same Vx source also compiles for the CPU, and a CPU emulator of the GPU's thread model runs the kernels in the test suite, so most bugs are found before a GPU is rented.

The same build on the rented GPU machine and on a development machine without a GPU gives byte-identical PTX. Rebuilding with a compiler 44 commits newer, on 20 October, changed no byte of it either.

How do we know it is right?

The usual check for an inference engine is to load the trained model, give it a prompt, and compare the output with a reference implementation. That needs the trained model to fit, and no trained GLM-5 checkpoint smaller than about 744 billion parameters has been published. So we compare on random weights instead, in the model's real shapes, against the official implementation: the GlmMoeDsa classes in Hugging Face Transformers, unmodified, in float32 on the CPU. A random model is a harder test than it sounds. Nothing in it is smooth or redundant, so a wrong index, a transposed matrix or a misplaced scale changes the output by a lot, not by a little.

The checks form a chain. Each one compares against the step above it:

ComparedModelResult
Vx GPU vs official, logits and generation6 layers, 24-token prompts, 3 seedsmax difference 8.8e-6; top-1 equal at every position; 16 greedy tokens identical
Every stage of every layer vs official78 layers, hidden 256logits within 1.0e-5; 0 sparse-attention selections differ
78 layers, hidden 1024, 64 headswithin 4.8e-6
full width: hidden 6144, 256 experts, 4 layerswithin 2.2e-6
Long sequence vs official2,300 tokens, top-2,048 selectionwithin 4.3e-6; 0 selections differ
Prompt in chunks vs official100-token prompt, 4 chunkingsrelative error 2.5e-6; greedy match in each
GPU-resident decode vs the GPU model above6 layers, and full width at 4 layers, 48 tokensbit-identical

"Every stage" means the norms, the attention output and the feed-forward output of each layer, at every position, and the list of tokens each query's sparse attention selected. The tolerance is 2e-4, and the worst error seen is 20 times smaller.

One early failure was not an error at all. With only two indexer heads, the ReLU in the indexer can zero every head for some key, and two keys then tie exactly. PyTorch's topk and our kernel broke the tie differently, a different token was attended to, and a logit jumped by 4.0. With 16 heads the ties go away. A trained checkpoint can tie too, so the tie rule now matches the model's exact argsort(-scores, kind='stable'), on the GPU as well.

Keeping the weights on the GPU

The first GPU path kept the model in host memory and sent each layer's weights to the GPU when it was needed. That is simple, and at full width it is slow: 29 seconds per token, nearly all of it copying. The resident path uploads each weight once and leaves it there, with activations and the attention caches on the GPU too. Per token, the host sends the token's embedding row and the RoPE angles, and reads back the ids of the experts the router picked and the logits. At full width it takes 10.8 ms per token for 4 layers. These are single runs of a test script on random weights, not benchmarks.

Both paths launch the same kernels in the same order, so their logits must be equal bit for bit, and they are. A bitwise test is stricter than a tolerance: a difference of one bit points at a mistake in how the launches are ordered, which a tolerance would hide.

INT4, FP8 and BF16

The trained model will not fit on a GPU in float32, so the resident path keeps each weight matrix in the format it is stored in and decodes it inside the matrix-vector kernel:

Each is compared with the float32 path on the exact decoded weights. On 20 October, on the A100, FP8 and BF16 were bit-identical to float32. INT4 differs only in the order of its sums: the top token matched at every position, with a largest difference of 5.2e-6, and the weights took 3.8 times less memory.

The BF16 result being bit-identical needed a decode that is exact for every one of the 65,536 encodings. The PTX build has no bit cast from an integer to a float, so the kernel builds the value from the exponent and mantissa with products of powers of two, keeping subnormals, infinities and NaN. A test runs all 65,536 encodings.

What the GPU found

A 2,300-token comparison failed once, at position 1,428, and passed when rerun. Three identical runs started at the same time gave three different answers. Running 16 processes on one A100, each checking every GPU operation against an independent answer, showed wrong results in up to a fifth of some operations.

The cause was in the upload. Copying from ordinary host memory to the GPU can return before the copy has landed, and the copy is ordered against the default stream, not the streams our kernels ran on. A kernel launched right after an upload could read the old contents of its buffer. On an idle GPU the copy almost always won the race, and on a shared one it often lost. The fix is one synchronization after the copy. With it, 16 processes show no wrong result, and identical concurrent runs give identical logits. Every GPU run now includes 8 processes sharing the GPU, because this bug only shows under contention.

The 20 October run found two smaller things: BF16 tests that only run on a GPU were missing an import, and the saved toolchain archive for new machines held links into system libraries rather than the libraries themselves. Both are fixed.

What is left

The evidence for each number here, with the commands that produced it, is in the evidence folders of glm-vx, and progress is logged in glm-vx#1.

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