Skip to content

EXL3 Flash Next: deep prefill (3/8) - #1443

Draft
BlivionIaG wants to merge 3 commits into
ashhart:mainfrom
BlivionIaG:pr/exl3-s3-deep-prefill
Draft

BlivionIaG wants to merge 3 commits into
ashhart:mainfrom
BlivionIaG:pr/exl3-s3-deep-prefill

Conversation

@BlivionIaG

@BlivionIaG BlivionIaG commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

Flash Next EXL3: exact deep prefill

Draft PR; validation is in progress. Pushed head: 505a649e7d1ca794c5d74691ca046cc742761bd7. Cumulative family slice, based on main 7c0688951584497dad8b0d5894b657d432fa917f.

Adds expert-major deep-prefill chunks using the canonical routed-expert reduction from #1442. The device equality gate compares last-row BF16 logits and 24 greedy continuation tokens against the batched path at multiple chunk sizes.

The GU/down contract records identify the replaced kernels. The final functional audit passes the 4,099-token deep/batched BF16 last-row and 24-token continuation-equality gate at chunk sizes 256/1,024/2,048; earlier-source numerical receipts retain their actual source identity.

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 pass, and removing stream cleanup produces exactly one failure (245 pass, one fail, no crashes).

This exact pushed prefix passes 177/177 executed host tests, 16/16 build steps (macOS arm64, Zig 0.17.0). Prefix host tests do not by themselves qualify its device execution.

Validation and remaining gates

Pushed family tip: 8743eaa3f17717312a25152d2e039c733b1ae590. Tested inference build source: 0184770e52e48e0c68479d4f0076dad989e9d42f, whose 9,031 tracked files were verified against its Git tree. The only subsequent changes are the snapshot validation client, its negative tests, and launch-count documentation; all native inference/build-input blobs are identical. Runtime receipts retain the actual build source instead of being relabeled. Native binary SHA256 246f5c3e3c171b60d2e6e6bf0e02cc99fd98acb6cbbb851028fc01a513e17b57; complete CUDA runner SHA256 87886278dc5f900dc1a7b8c9f3c7717e33d5eaf2c206e81db59b91b89e17b54f.

This update enables page-aware decode, verification, deep-prefill and conditional-head graph bindings. Host-tier frame ownership advances before capture/replay, and captured kernels consume the resulting device tables. Page geometry and tier windows participate in graph identity; TP rendezvous permits the legitimate rank-specific page strides. Pooled sequence graphs retire before window storage is released, and all engine graphs retire before engine buffers.

The immediate predecessor ce9eea00da9500121cb59b07e1bab29e54518790 passes all eight short-prompt comparisons: TP1/TP2, paged tier off/on, eager/graph. Two repetitions of 128-token code and prose replies match full token/output hashes, isolated replies and plain replies. Every graph case has actual decode and verification replay on every rank and conditional-head replay on rank 0. These short prompts do not qualify sparse host fetches. The tested inference source adds graph retirement ordering without changing inference arithmetic. All 36 current-build cold replies pass full token/output hashes across TP1/TP2, paged tier off/on, raw 2,048/8,192/32,700 tokens, and first/repeated/concurrent phases, with 64 outputs and zero cached prompt tokens. Both tier runs exercise sparse host fetches. The direct graph/snapshot, long-context replay and additional final functional gates below are complete; the full runner failed in its live resource stage without a final SUMMARY.

Completed device checks at the tested inference source

The following results use build source 0184770e52e48e0c68479d4f0076dad989e9d42f and the binary hashes above:

  • TP1 and TP2 direct paged graph/eager verification pass 32 decode rows and 32 variable-width verification rows with equal whole-BF16-row 64-bit Wyhash values and exact greedy token IDs. Graph mode reports 28 actual replays on rank 0 at each rank count. Injected capture/instantiation failures and two-stream comparisons also pass at both rank counts. This is hash equality for the logits, not a raw-byte comparison.
  • TP1 and TP2 direct snapshots compare actual BF16 logits, cache bytes and recurrent-state bytes with std.mem.eql. Two reversed-page lanes are saved at 3,075 and 2,821 tokens, continued by 8,192 prompt tokens and 16 decode steps, and restored under eager execution and three graph cycles. All byte comparisons pass on every participating rank. Each rank's graph cycles record 14/16/16 decode replays; the later cycles each record 144 prefill replays.
  • TP1 and TP2 direct paged prefill pass 8,192 seeded tokens through two eager and three graph runs, comparing the last prefill logits and 24 subsequent whole-BF16 decode-row 64-bit Wyhash values. Each rank asserts actual capture or replay. Rank 0 reports 72 layer-graph replays in each warmed graph pass. This prefill check uses row hashes; the snapshot gate above separately compares actual bytes.
  • All four TP1/TP2, tier-off/on long-context witnesses pass. Each matches full token/output hashes for six cold replies at 8,192/32,700 tokens and the following 128-token code/prose recovery requests, including isolated and plain comparisons. Every participating rank observes decode and verification replay; rank 0 also observes conditional-head replay. Latest prefill counters are 856 replays, 304 captures and 120 updates with tier off, and 93 replays, 96 captures and zero updates with tier on. Both tier-on cases exercise sparse host fetches. Counters span long and recovery requests: they are cumulative snapshots, not sums of repeated log lines or proof of per-request attribution.

Completed final functional audit

The independent GPU functional audit passes all ten listed checks, recording PR-build source 0184770e52e48e0c68479d4f0076dad989e9d42f and the separate current-main merge source 04383964f9b5d9fdae8491b87f2f1f704e1e8267. The source-specific direct graph and benchmark results above remain attributed to their actual builds. Additional completed gates are:

  • 48 served prompt-cache replies, with restored and cold full token hashes equal, including tier off/on.
  • 24 interrupted streams followed by 24 complete exact recovery replies.
  • The six three-distinct-long-prompt cases across TP1/TP2 and contiguous/paged tier-off/on execution: 18 complete concurrent replies, 18 equal isolated comparisons and 18 equal plain comparisons, without errors.
  • All 24 default-prompt 256-token replies and 12 short-workload sweep cells at widths 1/2/4 pass; the recorded rank/draft vectors match. These functional checks do not qualify the required 16,384-input/1,024-output c=1/c=8 workload.
  • At 4,099 raw prompt tokens, chunk sizes 256/1,024/2,048 preserve the last-row BF16 bytes and all 24 greedy continuation tokens.
  • Current-main merge 04383964f9b5d9fdae8491b87f2f1f704e1e8267 passes three distinct long replies in each of its three geometries, plus both ranks' actual-byte snapshots and observed graph replay.

This is a functional-GPU audit, not a waiver or pass for the broader host or global resource gates.

Independent numerical-phase audit and failed resource stage

The full CUDA runner at merge source 04383964f9b5d9fdae8491b87f2f1f704e1e8267, runner SHA256 658152c104d30d3aab30827eaa406db81899b6825b9748a1d6f886e5035b25a7, exited 1 during the live resource audit with UnknownContractImage: fp8_experts. It emitted no final runner SUMMARY. The resource audit did not complete, and the whole-suite gate is not passed and not waived.

A separate independent audit of the complete raw log (SHA256 4805bcd22d620f84c293d3d4f0043d8d4b39703d90a61712b18476d35a2ed033) covers the completed numerical phase: 441 registered runAll kernels, 315 planted controls red, six named KNOWN nonfinite annotations, and zero unexplained numerical failures. All supplemental numerical suites returned before execution entered live(). This is an independent numerical-phase audit, not a recovered runner SUMMARY or a whole-suite pass.

The six KNOWN annotations remain recorded:

  • Qwen35 decode attention with BF16 and FP8 caches, and Qwen35 contiguous attention with BF16 and FP8 caches: four annotations involving tf_q35_attn_partial, tf_q35_attn_partial8, tf_q35_attn_partial_legacy, tf_q35_attn_partial8_legacy and tf_q35_attn_merge. A NaN key in one chunk can leave the final row finite because the merge skips that chunk.
  • GLM tf_glm_scores and tf_glm_index_scores on pages out of order: two annotations where the ReLU fmaxf operation turns a NaN dot product into zero while the FP64 reference retains NaN.

The inherited broader host failure and stale global resource records remain unwaived. This outcome adds no performance result for the required 16,384-input/1,024-output c=1/c=8 workload, which remains pending.

The preceding pushed tip f94bcff5d1b0f0a506637e6ff70ecf1e17448c0d fixes retained borrowed recurrent spares after actual peer-window settlement. Its positive regression passes; removing its guard gives 246 pass, one fail, no crash. At that source, nine paged TP2 first/repeated/concurrent replies at 2,048/8,192/32,700 raw tokens plus 64 outputs match the TP1 reference with no CUDA error or host allocator report. That predecessor's paged path was eager despite its graph setting; those replies are not graph qualification.

  • Both exact updated graph/guard prefixes pass 322/322 host tests and 16/16 build steps. The snapshot gate now rejects missing, legacy-only, truncated or malformed full token SHA256 values: three client tests pass, along with all 55 benchmark-client regression tests. The stricter snapshot client SHA256 is 02a5440b2e3e7dee1f2e8a9f3cc00f3aab29f79e3b13fedab30206a714d9e5a3, recorded independently of the inference binary. At tested inference source 0184770e52e48e0c68479d4f0076dad989e9d42f, the full macOS host/contract/golden/fuzz build passes 596/596 steps, with 986 executed tests passed, four skipped and remaining work cached. The final Linux native/host build passes 276/276 steps, 1,102 executed tests passed and four skipped, using unchanged prebuilt kernels. The complete Flash Next host-test module cross-compiles for Windows. Lean remains at 425 inherited findings. All eight current pushed draft heads have nine successful CI checks each.
  • Hardware gates 1–4 have been exercised only on sm_120. Hardware execution on sm_80, sm_86, sm_89 and sm_121 has not been run. More than two physical ranks have not been run.
  • The offline gate-5 build covers sm_80/86/89/120/121. All 110 cells of this chain's 22 EXL3 launch records are current and within resource limits, including the corrected TP-mask register count of 10 on every SM. The full repository audit still fails with 1,235 stale-record problems outside those 22 launches; no actual launch-resource failure was reported. Verified unchanged upstream 7c0688951584497dad8b0d5894b657d432fa917f independently reports 1,237 stale-record problems. All shared launch resource measurements match this chain exactly; neither run is a whole-repository gate-5 pass. No five-SM whole-repository pass is claimed.
  • The broader CUDA-host suite at the preceding tip remains failed: 1,847 passed, 17 skipped, one failed, 465/467 steps. digest cache detects same-size rewrites even when mtime is restored fails at digest_cache_test.zig:87 on that tip and unchanged main 7c0688951584497dad8b0d5894b657d432fa917f; digest/build source blobs are unchanged by this update. Immediate same-size rewrites with restored mtime retain ctime on tmpfs and ordinary storage; delayed diagnostic rewrites change ctime. The failed gate is preserved without a source or test waiver.

The following receipts retain their actual preceding source identity a4dd9f9d02b9c47897f2078817772dee9165b0ff, rather than being attributed to the new head:

  • The registered GU/down prefill contract reports SUMMARY 2 kernels: 2 pass, 0 fail, 0 known; 1 controls: 1 red, 0 missed. Independent FP64 stress references, CPU trellis decoding, exact-integer and half-midpoint assertions, and FP16/BF16 simulations in sequential, pairwise and 32-lane orders cover six shapes and five-run repeatability per kernel. GU checks complete scratch bytes; down permits only two documented FP32 epilogue roundings. BF16 misses the down bound by at least 5061x; FP16 misses both kernels. These kernel and test sources are unchanged by this update.
  • All eight primary serving configurations complete: 36 cold replies match full hashes across rank counts, repeats and two cache geometries at 2,048/8,192/30,720 prompt tokens. The 1/2/4-stream sweeps and four distinct prompts contain 92 full-length replies, 92 isolated comparisons and 20 plain comparisons, with zero failed/unmeasured replies or mismatches. Four-stream observed lines are 16 at TP1 and 25/36 at TP2, from a request of 16.
  • TP1/TP2 eager/graph ABAB passes: all eight rank/phase vectors match full 128-token code/chat hashes, with per-rank replay coverage in both graph-enabled phases. TP1 prompt-cache cold/warm/branch/rewind matches its no-cache baseline on two pool lanes, 32 output tokens each; restored prefixes are 4,096 tokens. TP1 paged tier off/on matches at 2,048/8,192/32,700 raw prompt tokens plus 64 outputs.

The required exact 16,384-input/1,024-output performance workload at c=1 and c=8 remains pending, including eight active distinct inputs at c=8 and all four timing/rate metrics. The full CUDA runner on merge source 04383964f9b5d9fdae8491b87f2f1f704e1e8267 exited 1 during the live resource audit at UnknownContractImage: fp8_experts and emitted no final SUMMARY. The independent numerical-phase audit above is not a whole-suite pass. The inherited broader digest-cache host failure and 1,235 stale global resource records remain failed and unwaived. Canonical reduction-order acceptance and the explicitly deferred matched #1113 performance comparison remain outstanding. These PRs remain drafts.

The fresh official CUDA build of the separate current-main validation merge 04383964f9b5d9fdae8491b87f2f1f704e1e8267 preserves every instruction/control encoding word in all 31 pre-existing untiered attention/batch functions against baseline 7c0688951584497dad8b0d5894b657d432fa917f at sm_120/nvcc 13.4. This comparison excludes ELF metadata and the added plane/canonical EXL3 functions. Added plane functions have separate contract coverage. A matching NVFP4 checkpoint is unavailable, so no end-to-end NVFP4 serving receipt is claimed. The checkpoint manifest covers 24 files and 85,139,440,617 bytes, SHA256 4e7cbfc3b3d24c2e381b134e9203147dea4f18860026c6b9c17e4c29c54e8ec5.

Current observed main is d33dcfc767ee093393c6bfa00b3949bd7b4d8e88. Since the preceding observation db46ffd2a3ebc6c7600eb9e2cad5b1b7fa4992db, it adds an unrelated Qwen4 HC activation-contract correction and documentation, with no EXL3 runtime changes. The final pushed tip merges cleanly with this main locally. A separate local-only validation merge, 04383964f9b5d9fdae8491b87f2f1f704e1e8267, has all 9,049 tracked files verified and its official build passes 577/577 steps, with 2,549 tests passed and four skipped. Its native binary SHA256 is 6416b63defa3906bc2a4292201f3fefe47cfc4fd823aad15728af1e865b5e607; CUDA runner SHA256 is 658152c104d30d3aab30827eaa406db81899b6825b9748a1d6f886e5035b25a7. The final functional audit passes the integration three-distinct-long-prompt cases in all three tested geometries and both ranks' byte-level snapshots with actual graph replay. Its full runner exited 1 in the live resource stage without a final SUMMARY; this merge is not a pushed PR head. PR-build runtime receipts retain source 0184770e52e48e0c68479d4f0076dad989e9d42f; integration receipts retain source 04383964f9b5d9fdae8491b87f2f1f704e1e8267. Neither set is relabeled as the other.

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.

Required performance acceptance workload — pending

The required workload is exactly 16,384 input tokens and 1,024 output tokens per request, measured separately at c=1 and c=8. At c=8, all eight requests must be active together and use eight distinct 16,384-token inputs. Submitting eight requests that execute sequentially does not establish this concurrency case.

Both cases must record all four metrics: TTFT seconds, prefill tokens per second, decode tokens per second and mean ITL milliseconds. Qualification requires exact observed input/output lengths, zero cached input tokens, complete replies, full token/output SHA256 evidence and the promised drafted/plain, concurrent/solo and rank-count equality checks. The graph path must have observed graph capture and actual replay evidence; enabling a graph flag or executing eagerly does not by itself establish graph execution.

These required c=1/c=8 measurements are pending. The completed default-prompt 256-output measurements, 128-output sweeps at widths 1/2/4, earlier three-stream long-prompt checks and 32,700-input/64-output memory cases retain their original scopes; none is relabeled as the required 16,384-input/1,024-output qualification. A matched #1113 performance comparison remains unmeasured. The unchanged performance-deferral note above does not turn this target workload into a completed gate.

Prior/default short-workload measurements (256 output tokens)

These completed measurements cover the prior/default short-prompt workload, not the required 16,384-input/1,024-output performance gate. Build source 0184770e52e48e0c68479d4f0076dad989e9d42f, native SHA256 246f5c3e3c171b60d2e6e6bf0e02cc99fd98acb6cbbb851028fc01a513e17b57: four fresh single-stream TP1/TP2, plain/drafted configurations, 256 output tokens, three repetitions, greedy temperature 0, context 32,768, graphs and asynchronous instantiation enabled, prompt cache off. All 24 complete replies match full token/output hashes per prompt across rank counts and draft modes, with zero failed/unmeasured replies. Requested expert-cache lines are 72; observed lines differ as shown.

Medians, in TTFT seconds / prefill tokens per second / decode tokens per second / mean ITL milliseconds:

  • TP1 plain, prose, 72 observed lines: 0.536 / 57.782 / 26.464 / 37.788.
  • TP1 plain, code, 72 observed lines: 0.247 / 56.774 / 27.091 / 36.913.
  • TP1 drafted, prose, 72 observed lines: 0.573 / 54.147 / 28.745 / 34.788.
  • TP1 drafted, code, 72 observed lines: 0.267 / 52.343 / 33.382 / 29.956.
  • TP2 plain, prose, 93/91 observed lines: 0.188 / 164.471 / 41.258 / 24.238.
  • TP2 plain, code, 93/91 observed lines: 0.076 / 184.055 / 43.243 / 23.125.
  • TP2 drafted, prose, 79/88 observed lines: 0.232 / 133.801 / 39.557 / 25.280.
  • TP2 drafted, code, 79/88 observed lines: 0.110 / 127.288 / 61.317 / 16.309.

TP2 prose drafted decode is 39.557 versus plain 41.258 tokens/s, with mean ITL 25.280 versus 24.238 ms. Drafted cache geometry is 79/88 lines, versus 93/91 plain: these mode differences are measurements, not an isolated draft-overhead experiment. A matched comparison against #1113 remains unmeasured; no within-10% claim is made. The final functional audit separately passes 12 concurrency-sweep cells at widths 1/2/4 with 128-token replies.

Qualified two-rank host-tier memory comparison

Separate merge source 04383964f9b5d9fdae8491b87f2f1f704e1e8267, native SHA256 6416b63defa3906bc2a4292201f3fefe47cfc4fd823aad15728af1e865b5e607: TP2, 32,700 raw prompt tokens plus 64 outputs, uniform 16 expert lines per rank, prompt cache off, graphs/async enabled. The three phases submit one request at a time, including the one-request concurrency phase; two pool lanes are available. Tier-on uses 4,096 MiB host budget per rank, a 256-frame cap and slots setting 9. System-reported physical RAM is about 251.55 GiB. Both cases retain exact cold full hashes.

Owned requested bytes after all requests and before shutdown:

  • Tier off: device live sum 19,474,886,564; requested pinned live sum 49,057,599,720. Per-rank retained device/pinned peak counters, in bytes: rank 0 10,115,896,204/24,928,729,048; rank 1 9,550,145,562/24,128,870,672.
  • Tier on: device live sum 19,729,141,508; requested pinned live sum 57,315,256,712. Per-rank retained device/pinned peak counters, in bytes: rank 0 10,248,107,420/29,222,706,216; rank 1 9,672,189,290/28,092,550,496.

At this 256-frame setting, tier-on increases tracked device live bytes by 254,254,944 and requested pinned live bytes by 8,257,656,992. These runs do not establish a lower-frame-cap VRAM-saving or capacity configuration.

  • Tier off sampled GPU process maxima, rank 0/1: 10652/10170 MiB; sampled RSS maxima: 49,210,860/48,165,456 KiB.
  • Tier on sampled GPU process maxima, rank 0/1: 10654/10162 MiB; sampled RSS maxima: 52,998,848/51,596,000 KiB.

The counters cover requested DeviceBuffer/HostBuffer bytes only; driver/context/graph/NCCL internals and allocations outside those wrappers are excluded. Retained peaks are since the runtime's last reset, not physical high-water marks; observed sampling cannot exclude unseen resets. Endpoint readings are included in the derived maxima. Stable double reads do not hold the allocator mutex and are not atomic across ranks; HTTP gauges expose the same counters rather than an independent physical measurement. Per-rank peaks are not simultaneous cluster peaks; pinned sums are process accounting rather than unique physical pinned-page counts. RSS includes mapped/shared file pages, so its sum is not a physical RAM footprint. Cleanup zero counters are not established by this receipt. NVML/OS values are one-second samples. PCIe integrals remain device-wide estimates including model load and may miss bursts; OS VmPin/VmLck do not measure CUDA-pinned bytes.

Serving measurements recorded at the preceding tip

This preceding-source corpus uses 256-output default prompts and separate 64-output cold cases; it does not measure the required 16,384-input/1,024-output workload at c=1/c=8. These are independent GPU diagnostic measurements on preceding source a4dd9f9d02b9c47897f2078817772dee9165b0ff, using native SHA256 fba7887b03084d6cb60f4019e7850a86163e4b63d188c822d3dcfe6ac9103cf7; they are not relabeled as new-head runs; the inherited broad host gate remains failed. Two RTX 5060 Ti 16 GiB GPUs, driver 615.71.09, CUDA 13.4, Zig 0.17.0; context 32,768, greedy requests, graphs enabled, asynchronous graph instantiation disabled, prompt cache off. All eight primary configurations completed.

Decode: 256 output tokens, three repetitions, medians. Each row gives TTFT seconds / prefill tokens per second / decode tokens per second / mean ITL milliseconds. Cache lines are observed per rank, not just the requested setting:

  • TP1 plain, fibonacci-raw, 72 lines: 0.273 / 51.3 / 27.70 / 36.09.
  • TP1 plain, gpu-chat-no-think, 72 lines: 0.553 / 56.1 / 27.14 / 36.84.
  • TP1 drafted, fibonacci-raw, 72 lines: 0.287 / 48.9 / 31.21 / 32.04.
  • TP1 drafted, gpu-chat-no-think, 72 lines: 0.579 / 53.6 / 27.37 / 36.53.
  • TP2 plain, fibonacci-raw, 93/91 lines: 0.091 / 153.0 / 43.73 / 22.86.
  • TP2 plain, gpu-chat-no-think, 93/91 lines: 0.199 / 156.0 / 42.25 / 23.67.
  • TP2 drafted, fibonacci-raw, 79/88 lines: 0.129 / 108.8 / 51.27 / 19.51.
  • TP2 drafted, gpu-chat-no-think, 79/88 lines: 0.232 / 133.4 / 39.59 / 25.26.

Cold raw-token prompts: 64 output tokens; first request at each length, same four units. Requested 40 lines, expert reserve 512 MiB, server reserve 1.5 GiB; observed lines are 40 at TP1 and 67/76 at TP2. Prefill is prompt tokens / TTFT; mean ITL is stream-arrival span / 63, including zero gaps inside a chunk.

  • TP1, 2048 prompt tokens: 3.759 / 544.8 / 18.60 / 53.77.
  • TP1, 8192 prompt tokens: 8.084 / 1013.3 / 39.60 / 25.25.
  • TP1, 30720 prompt tokens: 30.008 / 1023.7 / 49.52 / 20.19.
  • TP2, 2048 prompt tokens: 3.121 / 656.2 / 35.68 / 28.02.
  • TP2, 8192 prompt tokens: 8.269 / 990.7 / 60.37 / 16.57.
  • TP2, 30720 prompt tokens: 30.098 / 1020.7 / 76.23 / 13.12.

A measured gap remains: TP2 prose drafted decode is 39.59 versus plain 42.25 tokens/s (-6.3%), with mean ITL 25.26 versus 23.67 ms (+6.7%). Their cache geometries differ: 79/88 drafted, 93/91 plain. Current performance deltas against #1113 are unmeasured; no within-10% claim is made.

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

@BlivionIaG
BlivionIaG force-pushed the pr/exl3-s3-deep-prefill branch 3 times, most recently from a23cd77 to 0b8304d Compare October 27, 2026 10:51
@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

The prefill speedup is impressive: the 16k-token prompt goes from 56 s to 10.5 s. The blocker is arithmetic.

The deep path is on by default for segments of 512 rows or more (exl3/prefill/cuda_exl3_prefill_state.zig:22,28;
cuda_exl3_prefill.zig:211-212). It runs dense projections of 64 rows or more on #806's candidate tiles
(cuda_exl3_prefill_tiles.zig:10-14) and experts on the new 64-member tiles, while decode and shorter segments use
Linear.matmul and runSlots. The docs prompt's hash moving from 1882ea29ebf6 to e6834625432c shows the
arithmetic changed.

Keeping cache states on the step grid (snap_grid) is a careful way to make resumes match fresh runs. Even so, a
prompt's tokens now depend on the step size and on whether the deep buffers could be reserved. When they can't, at
startup (cuda_exl3_prefill.zig:79) or at request time (:267), the batched pass takes over. Prompt rows also no
longer match serial decode, which is our oracle.

To land:

  • The prompt kernels need to produce the decode kernels' bytes for every row, whatever the chunk or segment length.
    Main's NVFP4 Flash Next prompt tiles do this; see the coverage section of docs/kernel-contract.md.
  • Register tf_exl3pf_gu and tf_exl3pf_down with replaces against the decode kernels, so the contract checks the
    bytes.
  • Add a test that shows the same tokens with deep prefill on and off, and at two step sizes.

Once prompt rows match decode, the grid is no longer needed for exactness.

The ring, overlap and prefetch work can stay as it is.

@BlivionIaG
BlivionIaG force-pushed the pr/exl3-s3-deep-prefill branch 4 times, most recently from 3011ec2 to 8f85bfb Compare October 27, 2026 17:56
@BlivionIaG
BlivionIaG force-pushed the pr/exl3-s3-deep-prefill branch from 8f85bfb to 5207d74 Compare October 27, 2026 20:27
@BlivionIaG BlivionIaG changed the title EXL3 Flash Next: deep prefill (3/7) EXL3 Flash Next: deep prefill (3/8) Oct 27, 2026
@BlivionIaG
BlivionIaG force-pushed the pr/exl3-s3-deep-prefill branch 7 times, most recently from 9fbf8e4 to c003aba Compare October 29, 2026 07:53
@BlivionIaG
BlivionIaG force-pushed the pr/exl3-s3-deep-prefill branch from c003aba to 4d106ad Compare October 29, 2026 10:36
Ranks split the experts by contiguous runs of leaves (the routed ids by eights, the shared expert in the last
leaf), the routers are replicated, and MTP drafts run on every rank.

The MoE branch has one canonical reduction order at any tensor-parallel degree: each leaf is an fp32 fma chain
over the row's slots in pick order starting at +0, and the 64 leaves are added in a fixed balanced binary tree.
One rank evaluates the whole tree. A rank of d computes the maximal complete subtrees of its run (one for a
power-of-two d), one fixed-size all-gather of those fp32 [rows, 2560] partials follows, and every rank finishes
the same tree, so TP=1 and TP=d write the same bits. Two ranks exchange 10 KB a row a layer. A host test emulates
d = 1 to 8 against the whole tree and against the definition without shortcuts, and mutations of the order, a
slot, a leaf and the tree shape must change bits.
Expert-major prefill for a Flash Next EXL3 checkpoint whose routed experts live in host memory: deep chunks of up to 4096 rows,
the experts streamed through a ring of pool lines while the rows compute, and the dense projections of a chunk on tensor tiles.

Every row is the decode kernels' bytes whatever the chunk or segment length:

* Dense projections of 64 rows or more run tf_exl3_promptd_*: the rows are rotated by Linear.matmul's own tf_exl3_rot_in, each
  64-row block sums the same chains matmul's plan cuts (sk splits of wk chains of k tiles, an fp32 mma chain from zero each,
  added in warp order and then split order) and finishes with its transform, scale and one rounding.
* The 64-member expert tiles (tf_exl3pf_gu, tf_exl3pf_down) take 32 members a pass and sum K as the grouped kernels do (four
  chains of a split, then the splits, then 0 + each), with the same rotation and epilogue code. Transposed mma operands give the
  same fp32 bits as the decode kernels' roles.
* The branch of a row is slice 2's canonical reduction at any tensor-parallel degree: 64 leaves of eight routed ids (the
  shared expert in the last), a leaf an fmaf chain in slot order from +0, a fixed balanced tree over the leaves. A rank sums
  the maximal complete subtrees of its contiguous leaf run (tf_exl3pf_tp_combine), the block partials travel in one explicit
  fp32 all-gather per layer (fixed size, [rows, 2560] a rank at a power of two) and every rank finishes the tree the same way
  (tf_exl3pf_tp_finish). One rank sums the whole tree and exchanges nothing; TP=1 and TP=2 write the same bits, and the host
  tests replay the leaf sums and the gather for 1 to 8 virtual ranks against the single-rank tree bit for bit.

Tests: exl3-prefill dense and prompt compare the tiles with Linear.matmul and runSlots byte for byte on the pack's layers;
exl3-prefill equal runs one prompt through the batched pass and through deep chunks of several sizes and compares the last
row's logits and the greedy tokens; exl3-tpprefill requires the same bytes for deep against batched and for two ranks against
one. tf_exl3pf_gu and tf_exl3pf_down are registered as replacements of the grouped decode form and pass gates 1 to 3 of the
kernel contract on random layers, with a control; gate 4 stays a named finding. tf-cuda-test exl3-prefill-contract runs those
gates alone.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

EXL3 Qwen3.8-Flash-Next on CUDA sm_120: expert offload, TP=2, graphs and MTP (pilot; the chain that replaces draft #1113)

2 participants