[atakann]

What this changes

Part of #1490, after the profile posted there. TF_GLM_GRAPHS=1 runs GLM's verify windows on CUDA as CUDA graphs on each rank: one stream's windows (Engine.verify) and several streams' windows in one pass (Engine.verifyShared). It is off by default, and with it unset every pass runs eagerly as on main.

• Forward.run is now stage (each rank's state check, the rings' rebase when due, the ids upload) and then kernels (the embedding to the head). The eager path calls both, so it makes the same launches in the same order as before.
• cuda/keyed_graph.zig is the cache from flashnext/cuda_window_graph.zig as a shared helper. A key's shape picks the held graph, and moving holds every other host value. The same key replays. New values record the pass again and update the held graph in place (cuGraphExecUpdate), and other launches get a new executable. A new slot runs eagerly once before its capture. Flash Next keeps its own copy in this PR; it can move onto the shared one in a follow-up.
• glm/cuda_window_graph.zig builds the keys. A key holds every host value the window's launches take: the rows, each stream's slot and buffers, its KDA layers' state halves, and its length's host geometry (the 512-key chunks, each latent layer's dense and sparse branches, the selection's pool bucket and score grid). So a held graph moves in place at most once every 256 tokens, and the rounds between replay it.
• The ranks need no agreement. A graph and its eager pass launch the same collectives (NCCL and the RoCE gather) in the same order, and the RoCE gather keeps its sequence on the device, so a replay posts the right one.
• Not graphed: windows with TF_GLM_PROF, image rows or DFlash2's taps, prompt chunks, the MTP head's draft steps and the sampling, which reads the logits on the host.

Receipt

• Environment: two GB10s (driver 580.173.02), TP=2 over RoCE, nvidia/GLM-5.3-Flash-FP4-Local @ 76348dd, --context 32768 --parallel 4, default drafts (MTP + copies), zig build native -Dsm=121 (ReleaseFast), Zig 0.17.0, nvcc 13.3. Measured on main 6ec5fa4c plus this change. Rebased on 09e2f7fa: one commit touches these files (36a87df7, the test list in cuda/root.zig), and the tests below ran on the rebased tree.

• Exactness, one binary, TF_GLM_GRAPHS unset against =1: 14 requests (7 prompts of 93 to 7,808 tokens, each with drafts and with "draft": false), 700 tokens each with ignore_eos. The replies cross the 256- and 512-token seams, the dense limit and the 8,192-token pool bucket.

  • Greedy, one request at a time: 14 of 14 equal.
  • Greedy, four requests at a time (shared windows): 14 of 14 equal, and equal to the replies one at a time.
  • Sampled (temperature 0.7, a fixed seed a request): 14 of 14 equal. 12 of the 14 differ from the greedy replies.
  • In both arms, each prompt's drafted reply equals its "draft": false reply (7 of 7).

• Decode speed, the same binary, N concurrent greedy requests of 256 tokens, the mean of two runs:

  streams	unset (tok/s)	TF_GLM_GRAPHS=1 (tok/s)	change
  1	70.1	72.4	+3.3%
  2	86.8	89.9	+3.6%
  4	133.5	138.2	+3.5%

  This matches the idle time in the GLM-5.3-Flash on CUDA: capture decode and verify rounds as CUDA graphs (one rank and TP=2) #1490 profile (3.6% to 5.3%), most of it the GPU's own gap between queued launches.

• Prompt speed: prompt chunks run eagerly, and run makes the same launches as before. First-token times in the load runs are the same in both arms.

• Memory: the lowest MemAvailable during the graph run was 19.1 GiB on each rank, about 1.4 GiB below the run without graphs.

• Tests run, with the local .zig-cache cleared, so no result came from the cache:

  • zig build test --test-timeout 10m: 547/547 steps, 7143/7202 tests passed, 59 skipped, 0 failed.
  • test-kernel-contract 182/182, test-server-security test-pull-security 29/29 (server-security 24 passed), fuzz-cortheon 20/20. test-golden: 322/335 and 55/55 equal to server goldens, 0 unexpected differences.
  • zig fmt --check build.zig zig, check-release-paths, a plain zig build, git diff --check and the two Python unittests from the No GPU workflow are clean. Lean: 425 problems, 425 on the base, 0 new.
  • New host tests: cuda/keyed_graph_test.zig (Flash Next's cache test on the shared cache: replays, updates in place, a new executable on new launches; in test-runtime-host, 285/285) and glm/cuda_window_graph_test.zig (test-glm-graphs-host: the next tokens replay; a 512-key chunk, a score grid or the next request in a slot moves the graph in place; a KDA half, the sparse branch, other rows or another slot take another graph; a shared key follows its segments' order). Each fails by name when an expected count or key relation is changed.

• Not run, and why: TP=1 on one GB10, because the FP4-Local pack does not fit one GB10 (asked on GLM-5.3-Flash on CUDA: capture decode and verify rounds as CUDA graphs (one rank and TP=2) #1490 which checkpoint to use). Other GPUs and Windows hosts. Graphs under memory pressure: this cache has no floor yet (at most 128 graphs, cleared when full), so it stays opt-in.

Overlap: this merges cleanly with #2257 (32-row windows) and #2142. #2162 (TF_GLM_L2PF) touches the same lines of cuda_engine.zig (an import and verifyShared's forward), and its prefetch runs on a side stream during a window. If it lands first, I'll rebase and keep windows eager under TF_GLM_L2PF until the two are measured together.

zig build -Dgpu=spark fails on main (the -Dgpu path returns before profile.wire), so these builds used zig build native -Dsm=121. That fix is a separate PR.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (resumed prompts not tested apart: this change does not touch prompt passes or restores)
• No precision traded for speed. If the bits change, the description says which sums change and why. (no bits change)
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (no kernel change)
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks @atakann, this is a careful piece of work. The profile came first, the keys hold every host value a window's launches take (stricter than the gate asked), and at TP=2 greedy, sampled, four-at-a-time and drafted replies all equal eager on one binary, across the 256-token, 512-token and 8,192-pool seams. Off by default with the eager path making the same launches is what we need.

On TP=1: none of the GLM-5.3-Flash checkpoints we serve fits one GB10, so that half of the gate can't be reached there. We'll take TP=2 as the gate for this PR, and note that on #1490.

A few things before it can queue:

1. A CHANGELOG line under Unreleased: TF_GLM_GRAPHS=1 (opt-in) replays verify windows as CUDA graphs on each rank, with your TP=2 numbers and the 1.4 GiB lower MemAvailable. Plus a row in the GLM CUDA recipe.
2. The switch like GLM's others: register TF_GLM_GRAPHS with the profile registry, so the startup line and capabilities report it, and refuse any value other than unset, 0 or 1 by name before the weights load.
3. A device test of graph against eager: one window replayed and run eagerly from the same state, with the logits' bytes compared, in the shape of cuda/graph_cache_regression_test.zig. That's the parity test the gate asked for, and it guards later changes to the keys.

[atakann]

Thanks @ashhart. Items 1 and 2 are in b1c436a, rebased on 34ea8cb:

1. A CHANGELOG line under Unreleased with the TP=2 numbers and the 1.4 GiB, and a TF_GLM_GRAPHS entry in the GLM CUDA recipe.
2. TF_GLM_GRAPHS is a row in the profile registry (glm5_next, identical, .one), so the startup line and capabilities report it. cuda_ranks.Settings reads it before the weights load and refuses any value other than unset, 0 or 1 by name (TF_GLM_GRAPHS: 0 or 1). The engine takes it from the settings. cuda_ranks_test covers 1, unset and a typo.

With the local .zig-cache cleared: zig build test 547/547 steps, 7143/7202 tests passed, 59 skipped, 0 failed; kernel contract 182/182, goldens 322/335 and 55/55 with 0 unexpected differences, server and pull security 29/29, fuzz-cortheon 20/20; fmt, release paths, a plain build, the Python unittests and lean (0 new) are clean.

Item 3, the device test, is next. I'll post its shape here before the code.

[atakann]

For item 3, the shape I'll take: tf-cuda-test glm-graph-check, beside glm-shared-check and in its form. Twin streams on a glm-synth checkpoint (TF_GLM_LAYERS=2), the same prompts on both: each round one twin's window runs through the graph cache and the other's eagerly from the same state, and the logits' bytes and the drawn tokens are compared, then both keep the same rows. Lone windows and shared windows of 2 and 3 streams, for enough rounds to cross a 512-key chunk, so replays and in-place updates both run, with the cache's counts checked at the end. TP=1 on one GB10 first, then the same command at TP=2.

[atakann]

Item 3 is in 5a039df, rebased on 882ca0a: tf-cuda-test glm-graph-check CKPT RANK WORLD IDFILE (zig/tests/cuda/glm_graph_tests.zig), in glm-paged-check's form.

Three graphed streams and their eager twins get the same prompts (500, 2,330 and 3,990 tokens). Each round runs one twin's window through the graph cache and the other's with the caches off, from the same state, and compares the logits' bytes, the drawn tokens and the drafts; both then keep the same rows. 48 lone rounds a stream with one key each (they cross the 512- and 4,096-key seams), 48 shared rounds of all three (one crosses a score grid's), then 90 mixed rounds of 1 to 4 rows, lone and 2 or 3 streams shared. At the end the twins' whole states are compared, and each cache must have captured, updated and replayed.

On a GB10 at TP=1 with the FP4-Local checkpoint and TF_GLM_LAYERS=5:

PASS graphs at TP=1: 1619 window rows over 282 rounds (lone and 2 or 3 streams shared, 1 to 4 rows each, past 512- and 4,096-key seams and a score grid's) equal their eager twins' logits, tokens, drafts and states byte for byte; lone 17 captured, 4 updated, 135 replayed; shared 26, 12, 44

A planted control: with the 512-key chunks left out of the key, it fails at round 4, the first lone window past 512 keys (FAIL graphs: round 4 (1 streams, 4 rows): logits or tokens differ from the eager twins'). Lean: 0 new. TP=2 runs the same command on both ranks; I'll add it when both GB10s are free.

[atakann]

glm-graph-check at TP=2 on two GB10s, the same binary on both ranks (TF_GLM_LAYERS=5, NCCL over the direct link):

rank 0: PASS graphs at TP=2: 1619 window rows over 282 rounds (lone and 2 or 3 streams shared, 1 to 4 rows each, past 512- and 4,096-key seams and a score grid's) equal their eager twins' logits, tokens, drafts and states byte for byte; lone 17 captured, 4 updated, 135 replayed; shared 26, 12, 44
rank 1: PASS rank 1 followed rank 0's graphed and eager windows to the end

One note on what TP=2 covers: rank 0 turns its caches off for the twins' windows, and rank 1 keeps its own caches on for both, so on rank 1 the twins' windows run through graphs too. Rank 0's comparison therefore holds rank 0's graphs against its eager launches, with rank 1 graphed in both arms.

[ashhart]

Thanks @atakann, all three are in:

• The switch reads through the registry and refuses anything but unset, 0 or 1 before load.
• The recipe and CHANGELOG lines carry the numbers and the memory.
• glm-graph-check is the parity test we wanted. It compares 1,619 window rows over 282 rounds, lone and shared, across the 512- and 4,096-key seams and a score grid. They equal their eager twins byte for byte at TP=1 and TP=2, and a planted control fails where it should.

Running it at TP=1 on the cut-down checkpoint also covers the TP=1 half of #1490's gate for the window code.

Two small things for the landing. We'll do both ourselves when we prepare it, unless you'd rather push them:

1. core/profile.zig: main now ends both lists with TF_GLM_PROMPT_TAIL (glm cuda: TF_GLM_PROMPT_TAIL=1 keeps raw prompts' tail in the unpaged prompt cache, as the KV pages' store does #2226), so TF_GLM_GRAPHS goes after it. The rest applies as is.
2. The CHANGELOG line: "Thanks to @atakann." at the end.

Before it lands, we'll run main against this branch at TP=2 on two GB10s with the switch off, to confirm the default path is unchanged.

Optional, not blocking: the TP=2 check compares rank 0's graphs against eager while rank 1 replays in both arms. One run with the roles swapped would compare rank 1's graphs directly too.

[ashhart]

Our GLM-5.3-Flash NVFP4 run at TP=2 on two GB10s, main against this branch with TF_GLM_GRAPHS unset (the default), base new new base:

• Tokens are equal in all 52 checked cells.
• Prefill at 8K and 32K, and every decode cell, are level: 1,474-1,478 against 1,473-1,478 tok/s, and code at one and two streams 55.5 and 80.1-80.2 against 55.5 and 80.1-80.5.
• 2K prefill is unsettled: 879 and 862 tok/s against base's 1,613 and 963. That length runs in two modes on this pair (about 1,620 and about 960 tok/s, seen in another PR's run too), and both of this branch's runs happened to land in the slow one.

We've queued a prefill-only repeat at 2K and 8K, six runs in a balanced order, to tell the mode from the change. This lands if that's level.