HelixML

9% faster GLM-5.3-Flash prefill with GPU-to-GPU IPC (and why TileLang didn't help)

Oct 26, 2026

We profiled GLM-5.3-Flash on two RTX PRO 6000 GPUs, with the non-uniform memory access (NUMA) check and the SGLang settings that reproduce the result. A TileLang kernel ran 2x faster in isolation and changed nothing in production. A chunked PCIe all-reduce made prompt reading 9% faster, and a memory-placement bug explained a phantom 27% regression.

Our production GLM-5.3-Flash replicas now read prompts about 9% faster, with no extra GPU memory and bit-identical output. Getting there took three findings:

  • A TileLang kernel that was twice as fast as the one it replaced, and changed nothing.
  • A 25-line SGLang patch that did help.
  • A "27% regression" in a newer engine release, which turned out to be Linux placing one buffer on the wrong CPU socket.

The setup is our GLM-5.3-Flash TP2 recipe. Each copy of the model is split across two 96 GB RTX PRO 6000 Blackwell GPUs and served by SGLang v0.4.3. Two copies run behind Ramjet.

Where two GPUs spend their time

Share of GPU time in one GLM-5.3-Flash TP2 replica. Generating one request: routed experts 32%, FP8 matmuls 29%, other 16-bit matmuls 13%, all-reduce 6%, sparse attention 3%, hyper-connections 7%, everything else 10%. Reading a prompt in 6,144-token chunks: routed experts 34%, FP8 matmuls 10%, other 16-bit matmuls 5%, all-reduce 24%, sparse attention 11%, hyper-connections 7%, everything else 9%.

Most large kernels already ran close to the card's measured limits. Two entries didn't:

  • Generation: the FP8 matmuls used far less memory bandwidth than the hardware allows.
  • Prompt reading: a quarter of the time went to the all-reduce, the step where the two GPUs sum their partial results after every layer.

TileLang: twice as fast, and no faster

During generation, a verify step processes about six tokens. At that size the FP8 matmul launches only 16 to 96 thread blocks on a GPU with 188 multiprocessors, so most of the chip sits idle.

We wrote a replacement in TileLang in about 200 lines. It keeps activations in 16-bit, which is what this checkpoint's weight-only FP8 format intends. Each weight tile is dequantized with its 128×128 block scale, and the long dimension is split across blocks so the whole chip streams weights at once:

for ko in T.Pipelined(k_blocks_per_split, num_stages=2):
    kb = bz * k_blocks_per_split + ko
    T.copy(A[0, kb * BK], A_s)                 # 16-bit activations
    T.copy(B[bx * BN, kb * BK], B_s)           # FP8 weight tile
    for i, k in T.Parallel(BN, BK):            # dequantize with the block scale
        B_h[i, k] = B_s[i, k] * S[nb, (kb * BK + k) // 128]
    T.gemm(A_s, B_h, acc, transpose_B=True)    # mma.sync on Blackwell SM120
Kernel time for six tokens, CUTLASS FP8 including activation quantization versus the TileLang weight-only kernel: shared expert gate and up 22.5 versus 11.6 microseconds, sparse-attention input projection 24.2 versus 12.1, shared expert down 14 versus 8, linear-attention output 24.8 versus 18, linear-attention query/key/value 47.4 versus 43.5. Inside the engine, one generation step took 18.9 ms with either kernel.

Inside the engine, the step time didn't move. Three things absorbed the gain:

  • The shared experts run on a second CUDA stream beside the routed experts, and our 512-block kernel took processors away from them.
  • Some weights were already warm in L2, which helped CUTLASS more in the engine than in our benchmark.
  • The all-reduce wait after each layer absorbed what was left.

The kernel was also eight times more accurate against an fp32 reference, but it bought no speed. It stays out of production.

Chunking the all-reduce: +9%

The two GPUs have no NVLink and reach each other only through the PCIe host bridge. NCCL therefore stages every all-reduce through a buffer in host memory.

SGLang ships a FlashInfer all-reduce that writes straight into the other GPU's memory over PCIe. It only accepts tensors up to 192 tokens, though, so every 6,144-token prompt chunk fell back to NCCL.

The same 50 MB all-reduce on two RTX PRO 6000 GPUs: NCCL through a host buffer on the far CPU socket 5.7 to 6.7 ms, on the near socket 2.30 ms, NCCL direct peer-to-peer 1.84 ms, the tuned FlashInfer PCIe all-reduce as one call 1.58 ms (needs 420 MiB extra per GPU), and the same kernel in 192-row chunks 1.46 ms, which production now runs.

Running the whole tensor as one call needs a 420 MiB buffer on each GPU, and these GPUs have about 1 GB of headroom. Sending it through the existing buffer in 192-row chunks needs no extra memory and turned out to be faster. The core of the patch is a loop:

# sglang/srt/distributed/parallel_state.py, behind SGLANG_PCIE_IPC_CHUNKED=1
if ws.supports(input_):
    return ws.all_reduce(input_)
rows = self._pcie_ipc_chunk_rows(input_)   # workspace elements // hidden size
out = torch.empty_like(input_)
for r0 in range(0, input_.shape[0], rows):
    ws.all_reduce(input_[r0 : r0 + rows], out=out[r0 : r0 + rows])
return out

The decision to chunk depends only on the tensor's shape, type and layout, so both GPUs always take the same path.

Cold prompt reading on one TP2 replica, tokens per second before and after the chunked PCIe all-reduce: 9,340 tokens 5,912 to 6,453 (+9.1%), 37,265 tokens 5,871 to 6,440 (+9.7%), 148,965 tokens 5,843 to 6,384 (+9.3%). Generation unchanged at about 19 ms per step.

We rolled out one replica at a time while the other kept serving. Each restart took about 16 minutes, mostly spent loading 164 GB of weights.

The 27% regression that wasn't

On the same replica, upstream's v0.6.1 read prompts 27% slower than v0.4.3. The profiler blamed NCCL: 6.87 ms per all-reduce inside v0.6.1, against 2.47 ms inside v0.4.3.

Both versions used the same NCCL version and kernel configuration, and the same benchmark run outside the engine took 2.3 ms in both images. We then reproduced SGLang's environment variables one at a time. NCCL_CUMEM_ENABLE=0, which SGLang sets by default, doubled the time on its own.

Diagram of the two-socket server. The GPU pair hangs off CPU socket 0, whose memory is full of 33 GB of page cache. With cuMem off, NCCL's host buffer often lands in socket 1's memory on first touch, so every all-reduce byte crosses the UPI link: 5.7 to 6.7 ms. With NCCL_CUMEM_ENABLE=1 the buffer sits next to the GPU: 2.3 to 2.4 ms. The chunked PCIe IPC all-reduce goes GPU to GPU with no host buffer: 1.46 ms.

With cuMem off, NCCL's host buffer is an ordinary file in /dev/shm, and Linux puts its pages on whichever socket has free memory when they are first written. Socket 0 was full of page cache from reading model weights. Because the kernel is set not to reclaim that cache for new allocations (vm.zone_reclaim_mode=0), the buffer often landed on socket 1.

Our v0.4.3 runs happened to start just after another engine had freed socket 0's memory:

Cold prompt reading in tokens per second, range of runs: v0.4.3 with the buffer on the near socket 5,843 to 5,912; v0.6.1 with the buffer on the far socket 4,264 to 4,600; v0.6.1 with NCCL_CUMEM_ENABLE=1 5,977 to 6,019; v0.4.3 plus the chunked IPC all-reduce, as in production, 6,384 to 6,453; v0.6.1 plus chunked IPC 6,509 to 6,635.

v0.4.3 was never faster. Any restart of a GLM replica on this server could lose a quarter of its prompt-reading speed, and nothing in the configuration would show it.

The chunked path doesn't use the host buffer at all. NCCL still runs a small all-gather of the output logits on every generation step, so production also sets NCCL_CUMEM_ENABLE=1. To check your own engine, look for nccl- lines in /proc/<pid>/numa_maps. The N0= and N1= counts show which socket holds the pages.

Getting the gain on your own engine

The NUMA problem (NUMA means non-uniform memory access: on a two-socket host, some memory is closer to a given CPU than the rest) is not specific to GLM or to our server. Any SGLang replica running on a two-socket host without NVLink can lose about a quarter of its prompt-reading speed to it. The check takes one command, against a running engine that is under load:

grep -A2 'nccl-' /proc/$(pgrep -f 'sglang.*scheduler')/numa_maps

Each matching line ends with N0= and N1= counts. If the large count sits on the socket your GPUs are not attached to (nvidia-smi topo -m shows which), the engine is paying the UPI round trip on every all-reduce. Restarting can move the buffer to the near socket or back to the far one, so a single slow run proves nothing until you check where it landed.

The fix is one environment variable. Current SGLang sets NCCL_CUMEM_ENABLE from --enable-symm-mem, which is off by default, and treats an existing value as authoritative. With cuMem off, NCCL places its shared-memory buffers as /dev/shm pages at first touch, after weight loading has filled the host with page cache:

NCCL_CUMEM_ENABLE=1

On our server that took the prompt-reading all-reduce from 5.7–6.7 ms to 2.3–2.4 ms. Upstream v0.6.1 with only this change measured 6,009 / 6,019 / 5,977 tokens per second at 9K / 37K / 149K tokens, against 4,264–4,600 for the slow runs. It costs nothing in memory and does not change numerics.

Our v0.4.3 image set the value with os.environ[...] rather than a default, so passing it in still lost. That line was corrected upstream by September 2026. If your engine is older, check the NCCL init log for NCCL_CUMEM_ENABLE : set by environment to 1 before trusting the result.

The chunked all-reduce is a separate step and buys less. It only applies to GPUs that reach each other through a PCIe host bridge, and it needs a derived image: the patch checks the SHA-256 of the SGLang files it edits and refuses to build against anything else. If you want to reproduce it, Ramjet's GLM-5.3-Flash recipe has the patch, the Dockerfile and the Compose file that turns it on.

Keep NCCL_CUMEM_ENABLE=1 set alongside it. The chunked path replaces NCCL for prompt reading, but NCCL still carries the logits all-gather on every generation step and any reduction the patch declines.

Two things that did not help on our hardware, in case you are deciding what to try first. NCCL_P2P_LEVEL=SYS got the all-reduce to 1.84 ms and needed about 1 GB of extra device memory, which ran out of headroom at our --mem-fraction-static=0.99. Running the FlashInfer kernel over the whole tensor as one call needs a 420 MiB workspace on each GPU and was slower than the chunked version. Neither is worth the memory on a card that has 1 GB spare.

What shipped

  • In production: v0.4.3 with the chunked all-reduce, SGLANG_PCIE_IPC_TUNE=1 and NCCL_CUMEM_ENABLE=1. Prompt reading is about 9% faster and generation is unchanged.
  • Not deployed: the TileLang kernel.
  • Next engine upgrade: the patch ported to v0.6.1 measured about 2% ahead of production. That isn't enough to justify an upgrade on its own.
  • Next experiment: the routed-expert kernel runs its 16-bit matmuls at 60% of the card's peak. An FP8-activation version could be faster, but it changes the model's numerics, so it needs a quality study first.

Could this go upstream?

We have not opened a pull request, and the case is weaker than the numbers suggest.

Upstream does have the building blocks. sgl-project/sglang#34528, merged in September 2026, added the PCIe-IPC all-reduce as an optional backend with its own tuning cache. Its notes explain why it stops at decode-sized tensors: on an 8-GPU host, handing prompt reading to those kernels cost 66% on time to first token and bought nothing on per-token time. It also carried a 45% per-token regression at 128K context with four concurrent requests. Whatever a patch for prompt reading proposes, that budget is the thing it has to beat.

Our measurements disagree with that conclusion on our hardware. The difference is the shape of the group. Their prefill numbers come from an 8-GPU host with GPUs split 4+4 across two sockets, where the 8-rank kernel keeps four ranks on each side and only the sum between the halves crosses the socket link. We run two GPUs on one socket and no NVLink, so the choice is this kernel or NCCL staging every all-reduce through a host buffer. On that shape, 32 chunked calls at 1.46 ms each beat one 2.3 ms NCCL call by enough to move end-to-end prompt reading 9%.

We have not measured the middle: four ranks on one socket, or eight ranks spread across both. The PR's numbers for other group sizes suggest the point where one path overtakes the other depends on how many pairs are already talking locally. That is the test a contribution needs, and we have not run it.

The patch also assumes the whole tensor is contiguous and splits it by rows, so it silently does nothing for the strided and non-contiguous reductions the rest of the engine issues. Turning that into a general contribution means handling those cases and deciding what happens when the workspace is too small, which is more design than the 25 lines suggest.

The case that is easy to make is the NUMA one. SGLang derives NCCL_CUMEM_ENABLE from --enable-symm-mem, which is off by default, and on a two-socket host that turns prompt-reading speed into a restart lottery with no configuration to show it. Changing that default, or warning when the buffer lands on the far socket, would help every multi-socket deployment. It has nothing to do with the chunked path, though, and it does not need the patch.

Methodology

  • Hardware: two RTX PRO 6000 Blackwell Server Edition GPUs (SM120, 96 GB), connected only by PCIe, on one socket of a two-socket Intel Xeon 6505P server.
  • Model: ormandj/GLM-5.3-Flash-W4A16-NVFP4-K32-Experts-FP8-WO at revision ee0989a9, with production flags including EAGLE speculative decoding.
  • Measurements:
    • All runs used one isolated replica on spare GPUs, never attached to the load balancer.
    • Prompts were salted and cold, with zero cached tokens confirmed from response usage.
    • Generation speed was compared per verify step, because tokens per second also moves with speculative acceptance.
    • Two baseline runs agreed to within 0.5%.
    • Every run was under our thermal guard.
  • Artifacts: the patch, Dockerfile, profiler scripts and experiment journal are in helixml/ramjet#323. The deployment is in deploy/glm53_flash_sm120.