[BlivionIaG]

Flash Next EXL3: TP=2 and MTP drafts

Draft PR; validation is in progress. Pushed head: 56fa728c378a5fdc485bde7de6ee35210be4ece5. Cumulative family slice, based on main 9e8c368b0.

Adds two-rank EXL3 execution and MTP draft verification. The proposed canonical expert reduction uses 64 fixed leaves of eight routed IDs and a fixed balanced FP32 tree, so the order is independent of the rank partition. Shared expert 512 joins the last leaf before the final tree. The draft vocabulary head is expanded to BF16 at load.

The full-vocabulary rank now reserves the shared MTP head-logit scratch needed by batched draft verification. The former zero-row cap caused CUDA_ERROR_INVALID_VALUE under concurrent long prompts. Before the rebase, a regression failed with the old sizing and passed with the correction. The row bound and live launch guard now belong to each engine/allocation; a simulated process-wide 32-row guard failed the two-owner admission regression before the rebase. The reduction changes the previous TP=1 summation order; that tradeoff still needs explicit maintainer acceptance. No FP64 hardware reference is claimed.

Projection forks belong to their Ops owner and retire with it. This removes the retained 160-byte global-map allocation; partially initialized branches unwind before publication. The cleanup regressions and mutation result were recorded before the rebase; they have not been requalified on this head.

Rebased onto current main

Rebased onto main 9e8c368b0.

Validation and remaining gates

On the rebased chain tip #1446 (1ddba433a9894c45233de702ea2b4b0b81732b19) plus the expert branch, host gates passed: zig fmt --check, zig build test (9,566 passed, 55 skipped), test-flashnext-exl3-host (258), test-kernel-contract (150), test-golden, check-native-cuda-host, and the K/V steps (101). These are tip results, not per-head CI results.

The device matrix was run twice, on the rebase onto main 2761c24d and again on the tip rebased onto main 9e8c368b0 (source bd5d57abd, chain plus the experimental opt-in expert tooling). TP1 and TP2 covered c=1 docs, c=1 code and c=4, with 16,384 input and 1,024 output tokens, drafts 6, 16 expert lines, graphs and async enabled, actual graph replay witnessed on every rank, no prefix cache and no hidden warmups. Full outputs matched TP1 vs TP2, drafted vs plain, and concurrent vs solo, and no CUDA or graph-instantiation errors were logged. These runs are instrumented and are not serving-speed measurements. The current heads differ from the tested tip only by moving blocks between files to satisfy the repository's file-length check; host gates on the current expert-branch tip pass again.

The matched performance comparison to #1113 and uninstrumented ABAB remain outstanding. Earlier validation and resource results were recorded on pre-rebase sources and do not establish results for the rebased head.

Performance work is deferred on purpose

This chain lands the architecture first: a working, bit-exact base on upstream main (TP=2 equals TP=1, drafted equals plain, prompt rows equal decode rows, NVFP4 bytes unchanged). The performance optimizations of the origin draft (#1113) are not all ported here. We ported only the ones that are exact and cheap; the rest follow in separate PRs once this base is merged, each with its own byte-equality receipt, so nothing in this chain has to be rewritten.

Known gaps against #1113 on the same hardware (numbers in the tables below): deep prefill is slower than #1113's non-exact version, drafted decode and the expert-line count at 2 streams are below it, and the two-stream prose case with drafts on is slower than plain. Candidates for the follow-ups: tile tuning, overlapping the partial all-reduce with attention, trimming the deep-prefill chunk buffers, a cheaper draft probe.

This PR remains a draft while qualification and maintainer acceptance are outstanding. Fixes #1419.

[ashhart]

Thanks for this. Three things block it, two of which came up on #1113:

1. TP=2 has to give the same tokens as TP=1. An FP64 receipt won't settle it. Each rank now sums its own experts
   before the fp32 all-reduce (exl3/tp/cuda_exl3_tp.zig:263-291), which regroups the sum, and you report different
   text on two of three prompts. One way through is to exchange the per-slot fp32 expert outputs, sum them in pick
   order on each rank, and show equal hashes at TP=1 and TP=2. Could you also say what regroups in the head? As far as
   I can see, the logits are gathered whole before the argmax.
2. Please remove TF_FLASHNEXT_ATTN_SPLIT (exl3/model/cuda_exl3_heads.zig:26-32,181). It's an opt-in mode that
   changes tokens.
3. Offloaded engines shouldn't serve plain by default. plainByDefault() (exl3/model/cuda_exl3_pass.zig:114-118)
   turns drafts off unless TF_FLASHNEXT_DRAFT_ENABLE=1.

This PR also changes NVFP4 Flash Next. Please take those changes out, or give them their own PR with NVFP4 receipts:

• tp_draft_snapshots = true (cuda_lanes.zig:32) and --tp in the CUDA serves list (native/cuda.zig:35) admit
  NVFP4 TP=2 with drafts and the prompt cache. docs/multi-rank-receipts.md still says the CUDA build serves neither
  flag, and the refusal test (native/cuda_flashnext_refusal_test.zig:23) no longer checks Flash Next.
• TF_FLASHNEXT_DRAFTS and _CONFIDENCE are renamed (cuda_lanes_load.zig:75-76), so existing settings are ignored
  without a warning.
• Every NCCL wait now spins for 3 ms (cuda_collective.zig:14,85).

Two smaller points:

• The draft vocabulary is decoded to bf16 at load (exl3/tp/cuda_exl3_mtp.zig:164-180). That only affects proposals,
  but please say so in the description.
• fn_exl3_tp_mask and fn_exl3_tp_place need runner entries.