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:
- Sparse attention. Attention compresses keys and values into a small latent vector per token, as DeepSeek's MLA does. A separate indexer scores every earlier token, and each query attends only to the 2,048 that score highest. Below 2,048 tokens this selects everything; above it, which tokens survive depends on an exact top-k.
- Experts. After the first three layers, the feed-forward part is 256 experts, of which a router picks 8 per token, plus one shared expert.
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:
| Compared | Model | Result |
|---|---|---|
| Vx GPU vs official, logits and generation | 6 layers, 24-token prompts, 3 seeds | max difference 8.8e-6; top-1 equal at every position; 16 greedy tokens identical |
| Every stage of every layer vs official | 78 layers, hidden 256 | logits within 1.0e-5; 0 sparse-attention selections differ |
| 78 layers, hidden 1024, 64 heads | within 4.8e-6 | |
| full width: hidden 6144, 256 experts, 4 layers | within 2.2e-6 | |
| Long sequence vs official | 2,300 tokens, top-2,048 selection | within 4.3e-6; 0 selections differ |
| Prompt in chunks vs official | 100-token prompt, 4 chunkings | relative error 2.5e-6; greedy match in each |
| GPU-resident decode vs the GPU model above | 6 layers, and full width at 4 layers, 48 tokens | bit-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:
- INT4, from the compressed-tensors checkpoints: 4-bit weights with a BF16 scale per group of 128. Its tests include tensors read from a published GLM INT4 checkpoint.
- FP8, the official format: E4M3 bytes with a float32 scale per 128 by 128 block.
- BF16, for the router and output layer, which the quantized checkpoints keep at 16 bits.
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 trained model. Even in INT4, GLM-5 is larger than one GPU. The plan is a hybrid: attention and the dense layers on the GPU, the experts in host memory, which on the A100 machine we rented is 944 GB. That is the step that runs the real checkpoint end to end, from a prompt to text.
- Full-width runs of the INT4, FP8 and BF16 paths, and the multi-token prediction layer, which is not implemented.
- Speed. The kernels are plain matrix-vector products, launched one at a time. Fusing them, and using tensor cores for the quantized formats, comes after the trained model runs.
- Moving from the PTX build script to Vx's own
spawn on(..)GPU path, once that has warp shuffles and tensor-core instructions.
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.