[swbiggart]

What this changes

Before: with the opt-in TF_FLASHNEXT_DECODE_GRAPH=1 (which needs TF_FLASHNEXT_SPLIT_DENSE=0) at --tp 2, every first plain round failed on both ranks with UnownedCaptureSource and took the engine down (#2974). Cache.run set only pass.prepared, so the round's token ids were uploaded inside the capture, where the copy-source owner refuses host sources (since f630f1a).
After: a round stages its inputs on the stream with stageWindow, as the one-GPU shared windows do, and the graph captures or replays only kernels and collectives. Also: the communicator's abort releases the round graph before ncclCommAbort, which waits without a bound on a graph holding NCCL plans, so a lost peer would hang the other rank; and a folder with learned pockets stays eager, because rowScales allocates a new buffer every round. The default path is unchanged: both sit behind a round-graph capture or the opt-in.

• Issue: Fixes Flash Next on CUDA: TF_FLASHNEXT_DECODE_GRAPH=1 at --tp 2 fails the first plain round #2974. Refs Roadmap: one N-rank layer for TP, PP and EP, shared by every family #2560 (part of its --tp 2 work).
• Stacked on: main 4cfd95573. Six commits: staging, abort release, pocket guard, device test, gradient-serving.tsv re-pins, and 0ddafbc80, the changelog line (it also numbers the re-pinned rows' notes).
• atc overlap --json (ATC 0.9.6, 8 changed files, complete: false) on 46eb7bb30 over main 93a2e93ec, conflicts checked with git merge-tree against both:
  • CHANGELOG.md: main now has ## Unreleased, which that run saw this PR and llama on Metal: Qwen2 checkpoints, Bimo-7B-CyberAgent's bf16 release with its q, k and v biases and YaRN at HF's rounding points #2890 re-add; other PRs' lines there conflict.
  • New code conflicts, all with PRs already conflicting with main: MoE route tap: optional TF_ROUTE_TAP trace of global top-k #1886 (route tap) edits the prepare lines this PR deletes; the EXL3 stack (EXL3 Flash Next on CUDA: one rank, offloaded experts, sm_120 gate (1/8) #1441 to EXL3 Flash Next: draft lane guard for offloaded experts (8/8) #1446, EXL3 Flash Next: host-resident K/V on the tier (6/8) #1522) in four of this PR's files.
  • GLM-5.3-Flash on CUDA: a decode window's DSA select on many CTAs #2722, CUDA (GB10): a second carveout source, spark-attic's atticd, opt-in #2762, CUDA (GB10): Flash Next KV in the display carveout, opt-in #2763 and cuda: tf_qmm_prefill_many works out its copy addresses once, before the K loop #2825 touch other gradient-serving.tsv rows and already conflict with main. flashnext: refuse a malformed TF_FLASHNEXT_* setting by name before device work #2585 (merged) changed how the setting is read, not the capture.
  • Re-run right before this pull request opened: atc overlap --json at 2026-11-12 14:09 UTC on head 2c3d7e487; every conflicts entry is named in this section, or adds no conflict beyond those it already has with main (1), or conflicts only in CHANGELOG.md's ## Unreleased lines, as this section describes (6); each checked with git merge-tree.

Checks on this head

• Head commit: 0ddafbc80 on base 4cfd95573 (main). Receipts were measured on 46eb7bb30 (served cells on fc6f58fbb, only tests differ) and stand on this head: this PR's own lines are unchanged since apart from the cuda_round_graph.zig pin row, re-pinned to the merged file; the rebase onto eab305db0 brought in flashnext: refuse a malformed TF_FLASHNEXT_* setting by name before device work #2585, which changes how TF_FLASHNEXT_* settings are read, with the same defaults, and the rebases onto 4f8317dac and 000b534a2 changed only where the CHANGELOG.md line sits. The rebase onto ab78aef90 and 4cfd95573 brought in flashnext-rank: the failure modes start once both communicators are up, and timeout bounds only the collective #2976 and flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977: the device test keeps mode, the leave check and context = 1024 and takes flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977's prompt.chunkRows() and LOGITS line, zig/tests/cuda/main.zig keeps both usage lines, cuda_collective.zig is re-pinned with both notes, and flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977's exchange test rig sets the release hook to null. The device test ran again on dd2089769, this code on ab78aef90 (below).
• Lean checker, python3 tools/zig/lean_check.py, on this head: 341 on base 4cfd95573, 341 on head, 0 new. python3 tools/pr_check.py --base upstream/main on this head: 11 changed files, 0 errors, 1 advisory warning ("pass --pr NUMBER").
• Tests run, with the command and pass, fail and skip counts (macOS unless noted):
  • zig fmt --check build.zig zig and git diff --check upstream/main...HEAD on this head: exit 0.
  • zig build test-cuda-host on this head: exit 0 on Linux x86_64 (CUDA 13.3.1 image; run without --summary, so it prints no Build Summary line). On Linux, a406758da, this code on ab78aef90 before the exchange rig set its hook, crashed cuda_moe_exchange_test.zig's abort test the same way; that rig now sets it to null in the same commit. On Linux the head before, 928b4e086, crashed three tests (two in cuda_rank_media_test.zig, one in cuda_lanes_restore_test.zig) in Collective.abort: they build a Collective from undefined and left the new release hook unset, which macOS happened to read as null. Those tests now set it to null, folded into the commit that adds the hook; at 46eb7bb30, 379/379 steps, 6184/6223 passed, 39 skipped. It runs the new cuda_round_graph_test.zig: staging, capture, replay and recapture; an upload inside a capture fails cleanly; a peer lost after a replay or during a capture leaves no graph at ncclCommAbort.
  • Fail before: on main the test file does not compile (Cache.launch, Collective.release missing). At 46eb7bb30, with if (c.release) |r| r.run(r.ctx); planted as _ = c.release; in Collective.abort: 377/379 steps, 6182/6223 passed, 39 skipped, 2 failed (the two lost-peer tests).
  • zig build test test-golden -Dcpu=apple_m1 on f8c8046bd, this code on eab305db0 before the changelog-only rebases onto 4f8317dac, 000b534a2 and 3885dceb2: exit 1, 837/839 steps succeeded (1 failed); 7595/7669 tests passed (73 skipped, 1 failed). The one failure, dashboard_test ("dashboard page polls the public stats snapshot without private details", untouched here), is inherited: main eab305db0 fails the same test on the same Mac (837/839 steps succeeded (1 failed); 13041/13152 tests passed (110 skipped, 1 failed)).
  • zig build test-gradient-serving on this head: pass.
  • tf-cuda-test flashnext-model-rank on dd2089769, this code on ab78aef90, 2 x RTX PRO 6000, TF_FLASHNEXT_DECODE_GRAPH=1 TF_FLASHNEXT_SPLIT_DENSE=0 TF_FLASHNEXT_CHUNK=32: graph on, PASS 52 rounds equal eager rows byte for byte (round graph true): replays, two lanes, owner order, 512 seam; leave with rank 1 SIGKILLed after LEAVE-READY, PASS the leader stopped after 10 rounds when its peer left: RankDisconnected, rank 0 out 1.0 s after the kill.
  • Exit 0 at 46eb7bb30: zig build, zig build compile-cuda-tests -Dtarget=aarch64-linux-gnu, the Flash Next tests for x86_64-windows-gnu (the POSIX-only host test is skipped there), and a CUDA build for sm_121.
• CI: runs on this head.
• Not run, and why:
  • Device run with NCCL_COMM_BLOCKING=1: passed on an earlier head with the same code under test.
  • RTX served cells: the device test is the RTX evidence.
  • bench_concurrent.py --levels 1,4 --alone --serial with the graph: the device test covers two lanes.
  • TF_FLASHNEXT_ATTN_LANES=1 with the graph, and a checkpoint with pockets: read only (the checkpoint has none).
  • Device test on GB10: it binds 127.0.0.1, so both GPUs must be on one host.
  • A peer lost mid-recapture on a nonblocking communicator could still leave the abort waiting: not reproduced.
  • --tp 1, Metal, HIP, XPU: unreachable; the round graph needs a CUDA collective.

Receipt

• Main against branch: cells and numbers, or why they can't be run. On main the graph cannot serve (release 1.2.6, RTX pair, both communicators, logs in Flash Next on CUDA: TF_FLASHNEXT_DECODE_GRAPH=1 at --tp 2 fails the first plain round #2974).

  2 x GB10 (200G RoCE), served, bench_concurrent.py stream client via tp2_bench.py, 256 tokens, greedy, median of 3, one start each, tok/s:

  Setting	code drafted	code plain	chat drafted	chat plain
  branch fc6f58fbb, split 0, graph on	47.75	20.17	34.32	20.13
  branch fc6f58fbb, split 0, graph off (main's eager path)	47.94	19.04	34.68	19.06
  main fb8fdf3d8, its defaults (dense split on)	51.79	23.26	38.09	23.23

  2 x RTX PRO 6000 Blackwell Max-Q, TF_FLASHNEXT_DECODE_GRAPH=1 TF_FLASHNEXT_SPLIT_DENSE=0 TF_FLASHNEXT_CHUNK=32 tf-cuda-test flashnext-model-rank MODEL RANK PORT ids.txt [leave]:

  Run	main fb8fdf3d8	head 46eb7bb30
  graph on	FAIL flashnext-model-rank: UnownedCaptureSource	PASS 52 rounds equal eager rows byte for byte (round graph true): replays, two lanes, owner order, 512 seam
  leave, rank 1 SIGKILLed after LEAVE-READY	not run (fails earlier)	PASS the leader stopped after 11 rounds when its peer left: RankDisconnected

  With rank 0 killed instead, rank 1 exits with RankDisconnected. Raw results, logs and the comparison tool (python3 cell_sum.py . main-base rgc-graph-off rgc-graph-on): https://gist.github.com/swbiggart/8f680962b86c8db6201950f15eed53de

  A fix, not a speed-up: split 0 costs more than the graph gains.

• Boxes it reaches: Flash Next on CUDA at --tp 2 with TF_FLASHNEXT_DECODE_GRAPH=1, plain decode rounds without media or pockets.

• New switch? No; TF_FLASHNEXT_DECODE_GRAPH exists, off by default.

• Environment, with the TensorFold commit, Zig and backend toolchain, machine and GPU, checkpoint and revision: commits as above; Zig 0.17.0, nvcc 13.3.73, NCCL 2.32.3; RTX sm_120, driver 595.99.02; GB10 sm_121, rank 0 driver 580.178.04; Linux; nvidia/Qwen3.8-Flash-Next-NVFP4 at fc694b54fb0174e0913e6adf86691ef85a4ead47 with its MTP drafter.

• Paths and platforms the change reaches, which I ran and which I could not: Engine.round's graph (sm_120, sm_121); Collective.abort on every abort path (host tests, sm_120 leave); the pocket guard, read only. Not run: other SMs.

• Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false: token_sha is code 4e759dbd859c, chat 764da77fbd3f in every GB10 cell, drafted and plain. Capture cannot change the collectives' bits (two-operand FP32 sum, BF16/FP32 gathers).

• Main against this branch, same machine and checkpoint, posted before asking for a run on our machines (or why you can't): the tables above.

• Decode speed before and after, from tools/bench_openai.py: the GB10 table. Default decode is unchanged by construction.

• Prompt speed before and after: n/a; the graph runs only on one-row decode rounds.

• Licence of any new weights or drafter, or "n/a": n/a.

Checklist

• The title and description match the head commit, including the receipt and the stack.
• 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 in the diff or its commits. No AI attribution lines.

• 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 and concurrent were not served.
• No precision traded for speed. If the bits change, the description says which sums change and why.
• Decode and prompt processing are no slower than main at one stream and the widest admitted; with --drafter, no slower than without it, greedy and sampled. The opt-in is slower than main's defaults on GB10.

[ashhart]

Thanks @swbiggart, for the fix and for tracing #2974 to its cause.

The fix is right:

• Cache.run now stages the round with stageWindow, as the one-GPU windows do, and refuses an unstaged pass (zig/src/families/flashnext/cuda_round_graph.zig). The capture holds only kernels and collectives, so the copy-source owner never sees a host upload inside it.
• The abort releases the round graph before ncclCommAbort. Cache.clear drops the hook, and Engine.deinit calls it first, so the callback can't outlive the cache. Your planted check, where both lost-peer host tests fail without the release, shows they test the right thing.
• Leaving folders with learned pockets eager is the safe choice while rowScales allocates every round.
• The default path doesn't change: the hook is set only once an opt-in graph captures.

Thanks for the rebase. 2c3d7e4871 carries the same six commits I reviewed, with only the pins and the changelog moved.

One thing before it lands: it and #2977 overlap, and #2977 lands first (after #2976, which it's stacked on). Once #2977 is on main, please rebase onto it. Three spots will conflict:

1. run in zig/tests/cuda/flashnext_model_rank_tests.zig. Keep your mode argument, the leave check and context = 1024, and take flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977's const chunk = try prompt.chunkRows(); and its LOGITS line, so the same test still covers flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977's prompt passes at TF_FLASHNEXT_CHUNK=96.
2. The usage text in zig/tests/cuda/main.zig, where your flashnext-model-rank line sits next to flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977's flashnext-rank line. Keep both.
3. zig/tests/fixtures/gradient-serving.tsv. Both PRs change cuda_collective.zig, so its pin needs the merged file's hash, with both notes beside it.

CHANGELOG.md will conflict too. Don't worry about that one; we sort the changelog out when landing.

This and #2978 both edit Collective.abort, so whichever of those two lands second also needs a small rebase there.

Thanks also for the honest notes on what you couldn't run.

• Synthena (Ash's Agent)

[swbiggart]

Thanks. Once #2977 is on main I'll rebase this onto it:

• in flashnext_model_rank_tests.zig, keep mode, the leave check and context = 1024, and take flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977's prompt.chunkRows() and LOGITS line;
• keep both usage lines in zig/tests/cuda/main.zig;
• re-pin cuda_collective.zig from the merged file, with both notes.

The same goes for Collective.abort against #2978, whichever of the two lands second.

[swbiggart]

Rebased onto main 4cfd95573, which has #2976 and #2977, as 0ddafbc80, as you listed:

• flashnext_model_rank_tests.zig keeps mode, the leave check and context = 1024, and takes flashnext cuda: at --tp 2 a prompt pass reduce-scatters its routed experts' slots by rows and gathers their sums #2977's prompt.chunkRows() and LOGITS line.
• zig/tests/cuda/main.zig keeps both usage lines.
• cuda_collective.zig is re-pinned with both notes.

One more test line: #2977's cuda_moe_exchange_test.zig builds its Collective from undefined, like the media and restore rigs, so it now sets release = null, in the commit that adds the hook. Without it, Linux test-cuda-host crashed in that file's abort test on the rebased head. With it, zig build test-cuda-host (Linux x86_64) and test-gradient-serving pass, and lean is 341, 0 new.

tf-cuda-test flashnext-model-rank at dd2089769, this code on ab78aef90, 2 x RTX PRO 6000, TF_FLASHNEXT_DECODE_GRAPH=1 TF_FLASHNEXT_SPLIT_DENSE=0 TF_FLASHNEXT_CHUNK=32:

• graph on: PASS 52 rounds equal eager rows byte for byte (round graph true): replays, two lanes, owner order, 512 seam
• leave, rank 1 SIGKILLed after LEAVE-READY: PASS the leader stopped after 10 rounds when its peer left: RankDisconnected, 1.0 s after the kill.

[ashhart]

Thanks @swbiggart. The rebase preserves leave, the 1,024-row context, chunkRows, LOGITS and both usage lines. The merged pins and explicit null hooks in the test rigs are in place. The graph lifecycle tests pass, and your repeated 52-round/peer-loss device check answers that part too.

This now waits on one final TP=2 graph-on/off run that we'll schedule on two GB10s, with eager-byte equality and bounded peer failure. No action needed from you. We'll keep the abort-release ordering when combining this with #2978 and #2982.

• Synthena (Ash's Agent)