recipes/docker-images/glm53-kindling/experimental

Experimental speedups

Compose overrides that stack on compose/glm53.yaml. Each one mounts files over the image this repo builds (vLLM nightly ddd6fbca); several replace whole vLLM or FlashInfer files, so they only match that image.

How to use it

Requirements: the image built from this tree (image/build.sh) and the normal four-node setup from the main README (the same compose/glm53.yaml and a per-node compose/.env on every box). The overrides replace files inside that image, so they do not apply to other images.

1. Put both ConnectX PCIe roots on the fabric: tag both interfaces rdma in mentatd's MENTAT_ANNOUNCE_IFACES (main README, step 4), or set FABRIC_SUBNETS to both subnets in each node's compose/.env. The entrypoint derives NCCL_IB_HCA and NCCL_IB_GID_INDEX from it, and arx and arxbig use the same two devices. Each rank's log shows them (arx all-reduce: rank r/4 on [dev0, dev1]). With only one device they log a warning and fall back to NCCL.

2. Start the stack with every override, on every node: ./glm53 up -d from the repo root. The script holds the list of compose files, in order (later files win).

3. The first boot loads the checkpoint normally (~8 min), compiles the CUDA extensions, and writes a processed-weight snapshot per rank under weight-snapshots in the cache mount (CACHE_HOME, or the glm53_cache volume), ~48 GB per node. Later boots restore it (~3.5 min).

4. Check: bash smoketest/run.sh http://<head>:8002 should pass 8/8, and each rank's log should show arx all-reduce: rank r/4 and arxbig all-gather: rank r/4.

Turning pieces off

Leave a compose file out to drop that piece, or set its switch in .env (every node must use the same values):

Switch	Default	What it does
VLLM_ARXBIG	1	arxbig RDMA collectives for prefill (arx.yaml)
VLLM_ARXBIG_AG	0	route prefill all-gathers through arxbig; off because its output sits in pinned memory, where GEMMs reading it run 3.7x slower
VLLM_ARXBIG_RS	1	RDMA reduce-scatter buffers (~0.5 GB pinned per rank); needed by VLLM_GLM_SP_MOE_FUSED
VLLM_ARX_RING	0 (1 with TP=RING4)	for boxes cabled in a ring with no switch: arx and arxbig open QPs only to rank r-1 and r+1, and at four ranks the rank between two others relays their data, with the same results as over the switch. arx takes rings of 2, 3 or 4; arxbig takes 2 or 4, so a ring of 3 needs VLLM_ARXBIG=0
ARX_RING_PREV_HCAS	from mentat with TP=RING4	with VLLM_ARX_RING=1: the two RDMA devices of the port cabled to rank r-1. Each device uses the GID of its own IPv4 address. Device r must share its subnet with device r of the neighbour's list for that cable: root 0 then root 1 when every box is cabled the same way, root 1 first on a box whose cabling is mirrored
ARX_RING_NEXT_HCAS	unset	the same for the port cabled to rank r+1
VLLM_GLM_SP_TP	1	sequence parallelism for forwards of VLLM_GLM_SP_MIN_TOKENS (1024) or more (sp.yaml)
VLLM_GLM_SP_FP8_GATHER	1	gather KDA attention inputs as FP8
VLLM_GLM_SP_MOE_FUSED	1	MoE combine feeding the RDMA reduce-scatter; needs arx.yaml with VLLM_ARXBIG_RS=1 and megamoe.yaml with VLLM_MOE_PREFILL=1
VLLM_GLM_SP_MOE_QUANT_GATHER	1	with VLLM_GLM_SP_MOE_FUSED: route each rank's own rows and gather the MoE input as NVFP4 + FP8 instead of bf16 (same results)
VLLM_GLM_ARX_PREFETCH	1	L2 prefetch during decode all-reduces; needs arx.yaml
VLLM_MEGAMOE	1	decode MoE kernel for batches of up to VLLM_MEGAMOE_MAX_TOKENS (8)
VLLM_MOE_PREFILL	1	prefill MoE kernel for batches of VLLM_MOE_PREFILL_MIN_TOKENS (1024) or more
VLLM_MOE_PREFILL_Y8	1	FP8 per-expert rows in the fused prefill MoE
VLLM_TRITON_SPARSE_MLA	1	Triton sparse MLA instead of FlashInfer's (fixes.yaml)
MAX_NUM_SEQS	50	requests decoding at once (fixes.yaml); about 50 KDA states fit in the pinned KV pool, against the stock 32
VLLM_GLM5NEXT_RECOVERSSM	1	with recoverssm.yaml: KDA drafts verified from one saved state per request, bit-exact; raises MAX_NUM_SEQS to 64 at TP=4 and 16 at TP=2
VLLM_GLM5NEXT_DRAFT_POOL	1	the drafter's KV in its own small pool instead of a page in every KV block: ~40% more KV tokens (fixes.yaml)
VLLM_DENSE_W4	in_proj, o_proj, shared experts, drafter	regex of dense layers stored as NVFP4 (the rest are FP8); add \|lm_head$ for the opt-in NVFP4 lm_head
VLLM_DENSE_FP8_LM_HEAD	1	FP8 lm_head
VLLM_ADAPTIVE_K_ONLINE	1	adaptive-k learns the step cost while serving; 0 uses the fixed table below
VLLM_ADAPTIVE_K_MODEL	fitted on this stack	starting step cost: ms, ms per expert touched, ms per verified token
VLLM_ADAPTIVE_K_COST_MS	measured on this stack	with VLLM_ADAPTIVE_K_ONLINE=0: step cost per verify length

Debug switches: VLLM_MOE_PREFILL_CHECK=N and VLLM_GLM_SP_MOE_CHECK=N also run the stock path on the first N prefill batches and log the difference.

Snapshot version and tag

A snapshot holds processed weights. Its key covers the checkpoint, the TP layout, the MoE backend and max_num_batched_tokens. Two kinds of change need more:

• A code change to how weights are processed (weight_snapshot.py, dense_fp8.py, the loader patches): bump SNAPSHOT_VERSION in weight_snapshot.py. Snapshot names start with v<version>-, and each boot deletes the snapshots of every other version from its snapshot directory, for all TP sizes.
• A knob that changes processing (VLLM_DENSE_W4, VLLM_DENSE_FP8*): a new VLLM_WEIGHT_SNAPSHOT_TAG (set in fp8.yaml). Snapshots under the old tag stay on disk until the next version bump.

A shape change fails the restore with an error. A same-shape change (such as which layers are NVFP4) would restore stale weights without one.

Quality check

experimental/quality/quality.py (GSM8K and HumanEval against a running server; see its docstring for the data files and the no-network HumanEval run). experimental/quality/prefill_block.py measures what a long prefill does to other requests: gaps in a stream that is already generating, and time to first token for short requests arriving meanwhile. experimental/quality/agent_tools.py runs five small agent tasks (list, read, search and write files in an in-memory tree) streamed and not streamed, and checks the result; python3 agent_tools.py http://<head>:8002.

Results

Single stream, thinking off, 512 tokens (dev/repro/decode.py prompts):

	structured	code	prose
stock, DFlash2 k=7	121.9	91.3	38.7
all overrides	170.3	120.8	65.8

Concurrent streams, aggregate tok/s at 1/2/4/8 streams: stock 85.5/65.4/99.5/147.8, all overrides 126/102/146/198 (16 streams: 271). With a different prompt per stream, code runs 317 tok/s at 16 streams and 451 at 32, and 50 mixed streams run 404. RigMark (reasoning=low) code / prose / structured 107.9 / 61.7 / 157.1 tok/s. Needle retrieval 12/12 up to 507k tokens. KV cache 3.63M tokens per rank at the 26 GiB pin and 50 sequences (2.63M before the drafter got its own pool).

Cold prefill, tok/s (random-word prompts, nothing cached):

	32k	128k
stock	2,730	2,679
all overrides	4,946	4,750

Boot goes from about 8 minutes to about 3.5 once snapshots exist.

The pieces

• arx (arx/): all-reduce over RDMA for tensors up to 256 KB, without NCCL. It takes 13-27 us where NCCL takes 80-90. The GPU publishes a sequence number, and a CPU thread posts the RDMA writes over both ConnectX roots.
• snapshot (snapshot/weight_snapshot.py, base_loader.py): the first boot writes each rank's processed weights to disk. Later boots build the model on dummy weights through the normal path, then overwrite every tensor in place (about 15 s for 46 GiB). modelopt.py and flashinfer_cutlass_moe.py are left from an earlier design and may no longer be needed.
• fp8 (snapshot/dense_fp8.py): the NVFP4 checkpoint leaves about 4 GB per rank of dense linears in bf16, and decode reads them every step. This converts them to FP8 after load (per-channel weight scales, per-token activation scales, CUTLASS scaled_mm). Prefill-sized batches go through in 2048-row pieces when the weight is over 16 MiB: CUTLASS rereads the whole weight for every row of tiles, and past L2 that comes from DRAM, so one call on a 16k x 6416 x 4096 GEMM runs at 68 TFLOPS and the pieces at 169. +14-20% decode, and about +0.5% NLL on prose. Layers matching VLLM_DENSE_W4 (by default the KDA in_proj and every drafter layer, attention o_proj and the shared experts) go to NVFP4 instead, through megamoe/megadense4.cu (W4A16 at the 4-bit roofline, about twice as fast as the FP8 GEMM). They keep an FP8 copy for batches of more than 32 tokens. That is +6-10% decode. The in_proj costs about as much NLL again as FP8 did. The drafter layers only change acceptance. NVFP4 on every dense layer cost 2-3x more NLL, so it is not the default.
• megamoe (megamoe/): an NVFP4 MoE kernel for batches of up to 8 tokens. It reads the CUTLASS backend's own tensors, so CUTLASS still handles prefill. Activations stay 16-bit, which is exact in the weights, where CUTLASS rounds them to FP4. With VLLM_MOE_PREFILL=1 (the default here), batches of 1024+ tokens take megamoe/moe_prefill.cu: the same W4A4 math as CUTLASS on the same tensors, with the token gather, SwiGLU and FP4 requantization folded into the fc1 GEMM (hand-written block-scaled mma.sync NVFP4). About 18 ms per layer at 16k tokens against CUTLASS's 22-26, equally close to an fp32 reference. +7-8% prefill.
• adaptive-k (adaptive-k/adaptive_k.py): a scheduler that picks how many of the 7 DFlash2 drafts each step verifies (2, 3, 4, 5 or 7). It uses recent per-position acceptance, per request and averaged over all requests, and a step cost it learns from measured step times: a fixed part, the MoE experts the step's tokens touch, and a per-token part. Weight reads follow the experts touched, so a 128-token step costs far less than a straight line through small ones, and at 16+ streams the scheduler verifies more drafts than a fixed table would allow. The drafter uses Triton attention, which avoids a mid-step host sync. The scheduler picks the step's draft count before the drafts exist. Once they do, draft-trunc/ cuts each request's drafts where their predicted survival falls below VLLM_DRAFT_TRUNC_TAU (0.3). The prediction is vLLM's acceptance estimator, which scores each draft from the logit of its most likely token (AUC 0.91-0.93 per position on logged steps). Dead drafts keep the verify shape, route to expert -1 (the MoE kernels skip them) and are rejected, so greedy output is unchanged, token for token and logprob for logprob. A cut step grades only the drafts the estimator expected to survive, and fitting on those steps drifted it low over a long mixed workload. So every 4th step (VLLM_DRAFT_TRUNC_EXPLORE) skips the cut, and only those steps train the estimator. On TP=4, paired with gate/bench/bench.py after a long mixed warm-up and net of drift: code +3.3%, prose +3.2%, mixed streams +3.5 to 6.9%, structured -1.1%. VLLM_DRAFT_TRUNC_TAU=0 turns it off. draft-trunc/ holds whole copies of five vLLM files, so it checks the image's vLLM commit and stops the engine on any other.
• recoverssm (fixes/recoverssm.py, compose/recoverssm.yaml; VLLM_GLM5NEXT_RECOVERSSM=0 turns it off): the KDA layers keep one recurrent state per request instead of one per draft position, and after sampling replay the accepted tokens through the stock kernel, so outputs are bit-identical. Ported from vLLM's Kimi-K3 RecoverSSM. At TP=2, 16 requests run at once instead of 3, and 8 streams decode 44% faster. At TP=4, 61 run at 64 streams instead of 49 (+28% aggregate).
• fixes (fixes/): things found by profiling.
  • The fused mHC post + pre-norm decode kernel read its projection weights in fp32, 1.5 MB a call, though the checkpoint stores them in bf16. It now reads a cached bf16 copy (mhc_tilelang.py, tilelang_kernels.py): 8.2 us a call instead of 11.9, ~1% faster decode at one stream, and greedy logprobs identical to the fp32 path.
  • FlashInfer's MLA planner cloned a 136 MB metadata buffer on every step, only to roll it back on error.
  • The GLM indexer's head gate ran a fp32 GEMM that took 69 us per layer.
  • FlashInfer's sparse MLA runs at about 10 TFLOPS on GB10 whatever the indices. gb10_sparse_mla.py is a Triton kernel for GLM's NoPE MLA (one program per query, all 16 heads, 32 gathered keys per block) that runs 3.3x faster. VLLM_TRITON_SPARSE_MLA=0 goes back to FlashInfer. At decode sizes (a request's 1 + 7 draft tokens) one program per query left most of the 48 SMs idle, and FlashInfer decoded 3-6% faster; below 96 query tokens the kernel now splits each token's rows across programs and combines them (4.4x faster at 8 tokens), which recovers most of that without FlashInfer's per-step host sync and keeps Triton's prefill.
  • FlashKDA needs dense q, k and v, and the KDA short conv wrote them as one [tokens, 3 x 2048] tensor, so each KDA layer copied three 56 MB tensors per 16k-token step. causal_conv1d.py gains an out_group argument that writes each group of features as its own dense tensor, and kda.py uses it. The outputs are bit-identical; prefill +2-3%.
  • The DFlash2 drafter's KV rode the target's pool: every 2304-token block carried the drafter's pages for its five layers, 30% of the KV cache, though a sliding-window layer only ever holds its 2048-token window. kv_cache_utils.py, kv_cache_coordinator.py, kv_cache_interface.py, worker_utils.py and warmup.py give the drafter's group its own pool (sized for max_num_seqs windows plus one prefill chunk, 105 blocks, 0.58 GiB) and keep its block ids out of the target's zeroing and warmup. At the same 26 GiB pin: 1,416 -> 1,976 blocks, 2.63M -> 3.67M KV tokens at 32 sequences (3.63M at the 50 fixes.yaml sets). Draft acceptance and speed are unchanged. VLLM_GLM5NEXT_DRAFT_POOL=0 goes back; VLLM_GLM5NEXT_DRAFT_POOL_BLOCKS sets the pool's size.
• sp (fixes/model.py): sequence parallelism for prefill. A forward of at least 1024 tokens keeps the residual stream split across the TP ranks, so mHC and the norms run on a quarter of the tokens, with an all-gather before and a reduce-scatter after attention and the MLP. +17% prefill. Decode stays plain TP: applied to every batch, SP cost decode 10-15%. The KDA layers' attention inputs are gathered as per-token FP8, half the bytes, since in_proj's FP8 GEMM would quantize the same rows the same way (+2%).
• arxbig (arx/arxbig.cu, VLLM_ARXBIG, VLLM_ARXBIG_RS): all-gather and reduce-scatter over RoCE for prefill-sized SP collectives, from pinned buffers the ConnectX writes directly. The all-gather is ~10% faster than NCCL's (187 vs 165 Gb/s), but its result lands in pinned memory, and on GB10 a GEMM that reads its operand from pinned memory runs 3.7x slower, so prefill gathers stay on NCCL (VLLM_ARXBIG_AG=0). The point is the reduce-scatter. With VLLM_GLM_SP_MOE_FUSED, MoE layers under SP run the router, the shared expert and moe_prefill's fc1/fc2 in model.py, and one kernel writes shared + scaled routed sum (fp32, rounded once) straight into the send buffer, publishing rows as it goes, so the network runs under the finalize. It replaces the runner's scale and add passes and NCCL's reduce-scatter: ~12.5 -> ~8.4 ms per MoE layer at 16k tokens, prefill +2.5-3%. With VLLM_MOE_PREFILL_Y8 (default) fc2 writes its per-expert rows as e4m3 with a scale per 128 columns, halving the 1.07 GB each layer writes and reads back: prefill +5%, GSM8K/HumanEval/NLL/tool-call results unchanged. With VLLM_GLM_SP_MOE_QUANT_GATHER (default) each rank routes its own rows and quantizes them the way the experts would (NVFP4 for the routed experts, per-token FP8 for the shared expert's GEMM), and the gather moves those instead of bf16: 0.56 + 1 bytes per value instead of 2, and a quarter of the fp32 router GEMM per rank. The layer output matches the bf16 gather to within 1e-4 relative; prefill +3-4%. All ranks write to all peers at once and this fabric has no PFC, so incast drops packets and go-back-N retransmits make individual calls vary (4-9 ms); per-destination serialization was worse (two QPs cannot fill a link).
• ring (TP=RING4; image/fabric_ring.py, ring mode in arx/): TP=4 on four boxes cabled in a loop. A ConnectX-7 can't forward RoCE for another box, because RoCE addressed to its MAC goes to its own RDMA engine, so every collective uses only neighbour links. mentat 0.17's ring claim places rank i on member i in cable order and gives each rank its interface toward each neighbour. fabric_ring.py runs at each rank's Python start and turns that into NCCL's ring algorithm with a channel graph that receives on the port toward the previous rank and sends on the port toward the next, and into arx's devices. The graph carries the speeds NCCL computes for its own ring: with others it picked other protocols, the reductions summed in another order, and a greedy text diverged. arx and arxbig relay data for the opposite box through a neighbour, raw, so every rank sums the same values in rank order and results stay bit-identical. Each cable needs its own subnet for each PCIe root. On a switch, with both neighbours on the one port, arx/test_ring.py passed bit for bit, decode was ~1% slower from the relay and prefill 3-4% slower (2026-09-28). Not yet measured on cables. Three boxes, one cable per pair, also form a ring, and there both other ranks are neighbours, so arx sends every partial direct and relays nothing. arx's ring mode takes 3 ranks; arxbig's does not, because its ring reduce-scatter splits the rank two away across both neighbours, so a ring of 3 runs with VLLM_ARXBIG=0. The entrypoint has no three-box ring mode (only TP=RING4), so the NCCL ring and ARX_RING_*_HCAS are set by hand. fabric_ring.py lists each port's devices in PCI order, root 0 first, which fits boxes that are all cabled the same way; on a box cabled mirrored, root 0 of a port faces the neighbour's root 1, and its lists must put root 1 first. On cables, arx/test_ring.py --mode ring at WORLD_SIZE=3, cut down to the arx all-reduce checks, the all-reduce burst and the latency bench, passed 85 checks per rank with 12.2 / 19.7 / 42.8 / 69.1 us at 8 / 64 / 256 / 512 KB (#52).
• arx prefetch (VLLM_GLM_ARX_PREFETCH): while a decode all-reduce waits for its peers, its threads ask L2 for the weights the next kernels read (the router gate and shared expert after attention, the next in_proj after the MoE). About -0.6 ms per 32 ms step.

dense_fp8.simulate_nvfp4 is a diagnostic. It rounds the dense weights through NVFP4 to measure the quality cost: about +1% NLL on prose.

Sources

Most of fixes/, and a few files elsewhere, are vLLM or FlashInfer files with our changes, bind-mounted over the originals in the image. They keep their Apache-2.0 headers. The rest is ours.

Origin	Files
vLLM, changed	arx/cuda_communicator.py; fixes/attention.py, causal_conv1d.py, flashinfer_mla_sparse_sm90.py, gdn_attn.py, kda.py, kv_cache_coordinator.py, kv_cache_interface.py, kv_cache_utils.py, model.py, tilelang_kernels.py, warmup.py, worker_utils.py; snapshot/base_loader.py, flashinfer_cutlass_moe.py, modelopt.py
FlashInfer, changed	fixes/_fa_common.py
ours	arx/arx.py, arx_vllm.cu, arxbig.cu; adaptive-k/adaptive_k.py; fixes/gb10_sparse_mla.py, recoverssm.py (after vLLM's Kimi-K3 RecoverSSM); megamoe/; snapshot/weight_snapshot.py, dense_fp8.py; quality/

Other people's work in here:

• arx and arxbig take the RDMA path MTU from the ports rather than assuming 4096, from Chuck's #6.
• tilelang_kernels.py carries vLLM's copy of SGLang's mHC kernel, and causal_conv1d.py vLLM's adaptation of Tri Dao's causal-conv1d, both credited in the file.
• The decode split in gb10_sparse_mla.py is the Flash-Decoding scheme (Dao et al., 2023): split each query's keys across programs, then merge the partial softmax results.

Quality

quality/quality.py (GSM8K first 250 and HumanEval, greedy, thinking off; HumanEval programs run by quality/run_he.py in a no-network container), against stock on the same boxes:

	stock	all overrides
GSM8K (250)	97.2%	97.2%
HumanEval pass@1	156/164	156/164
count to 200, thinking off, 5 runs	0 corrupt	0 corrupt
tool call at 42k context, 40 greedy runs	10 diverge	0 diverge
NLL vs bf16 dense layers	within ±0.006	prose +0.008 to +0.025, code within ±0.007