Skip to content

[Engine] Add opt-in streamed-expert layer-major prefill - #1

Draft
xdCloudy wants to merge 39 commits into
mainfrom
proposal/qwen-layer-major-prefill-ninfer-review
Draft

xdCloudy wants to merge 39 commits into
mainfrom
proposal/qwen-layer-major-prefill-ninfer-review

Conversation

@xdCloudy

@xdCloudy xdCloudy commented Oct 27, 2026 •

Copy link
Copy Markdown
Owner

Summary

This PR turns the layer-major proposal into an opt-in implementation for Qwen streamed-expert prefill. With STRATA_PREFILL_LAYER_MAJOR=1, eligible prompts process several ordinary chunks through one layer before advancing, group routed expert work across those chunks, and stage each streamed expert at most once per layer/super-chunk. The feature remains disabled by default.

The implementation reuses the Strata layer-major work from Tomas Vizdal in issue #1670, through the current-Strata adaptation at durikim1-cpu/Strata commit ab57cca4. Attribution is retained. The later experimental cache-arena commit was not needed for this bounded scratch allocator.

NInfer-all's Qwen notes and executor were reviewed for bounded spans, grouped expert work, and overlapped staging. No NInfer source code was copied. Strata's MIT licensing remains unchanged; no NInfer Apache-2.0 notice is required for code reuse.

Design and limits

  • Reuses current Strata attention, GDN, QSA, PLE, expert cache, MMQ, and stream abstractions. A stream synchronization in scratch release keeps asynchronous copies from outliving their source storage.
  • Supports the tested Qwen streamed-expert geometry: 2560-wide embeddings, 512 experts, top-10 routing, 640-wide routed/shared FFNs, and the supported hyper-connection layout. It is single-GPU.
  • Falls back for unsupported geometry, layer splits or helper GPUs, image rows, DFlash2 taps, batch yields, forced chunks, route observers/debug dumps, non-MMQ or fused expert layouts, and insufficient scratch memory.
  • A 16,384-token minimum prompt guard was added after an active low-cache 8K run produced a different greedy continuation after token one. The divergence was reproduced with different cache sizes and super-chunk targets but is not explained. 4K and 8K prompts currently use the ordinary path, even when scratch would fit.
  • The prompt-length gate is not sufficient to guarantee parity at every cache capacity.
  • Layer-major now also falls back below 822 allocated expert-cache slots. The known-divergent --expert-cache 500 setup allocates 656 slots; two reversed-order 16K fallback runs matched automatic for all eight generated tokens. The lowest active configuration with repeated evidence was 822 allocated slots (--expert-cache 625), on one prompt; this is a conservative guard, not a general parity guarantee above the floor. The current committed binary was also checked at exactly 822 slots in two reversed-order rounds: layer-major remained active in both and matched all eight generated IDs each time. Automatic prefill measured 372.88/381.88 tok/s and layer-major 375.82/376.30 tok/s (means 377.38 vs 376.06), so no throughput gain is claimed. Reported expert H2D was 341,247.4 MB automatic vs 77,521.3 MB layer-major per run; peak VRAM was 6,380–6,393 MiB vs 6,548–6,552 MiB. This is repeatability evidence for one prompt and cache configuration, not a general parity guarantee.
  • A one-chunk schedule control at 500 slots (--layer-major-tokens 1280, one 1,280-token chunk per super-chunk) still diverged after the shared first generated token with layer-major active. This rules out multi-chunk aggregation within a super-chunk as the sole cause; only two output tokens were generated, so this remains diagnostic evidence. Updated measurements are in the benchmark results.
  • Four 16K checks at 500 expert-cache slots diverged after the shared first token; the result persisted with host/VRAM residual placement, frozen adaptive tiering, and 4,096-token super-chunks. At 625 slots, a two-round retest matched all eight generated IDs in all four arm runs with layer-major active twice. The 625/750 results do not establish a safe minimum; the cause of the 500-slot divergence remains unknown.
  • Expert accumulation is deterministic by expert/group order but can differ from router-order accumulation by FP32 rounding. Broader state and decode checks remain open.

Validation

Completed checks:

  • CUDA 13.3 builds for strata, strata_prefill, and layer_major_kernel_test at sm_75 and native RTX 3060 sm_86.
  • Nine CPU CTest checks covering model family, prefill planning/share rules, memory guard, file expert source, expert profile save/version, expert layout, and batch rows.
  • RTX 3060 layer_major_kernel_test, including shared-expert initialization, routed scatter accumulation, and copy_i32.
  • Post-guard 4K/8K matched A/B on the recorded earlier revision: both layer-major arms explicitly fell back with reason “a short prompt”; all eight greedy output IDs matched automatic prefill at each size. An earlier fresh 8K sample on this corpus was inconclusive because its outputs differed. A new two-round retest on the current executable at 8,192 and 16,383 tokens matched all eight IDs across automatic and feature-flag arms in both rounds; both feature-flag arms explicitly fell back with a short prompt, so the optimized path was active zero times. This confirms fallback parity for docs/DETAILS.md at these sizes, not active layer-major numerical correctness. Automatic/fallback rates were 303.67/307.41 vs 264.27/284.92 tok/s at 8K and 393.15/336.84 vs 378.99/375.42 tok/s at 16,383. The executable was built at f3feced2; the branch head adds benchmark records only. Structured rows are under cachefloor_fallback_current_head_validation in the results JSON.
  • 16K matched A/B: after the guard fix, the exact 16,384-token minimum was active in two reversed-order rounds and produced exact 8-token greedy parity in both. A fresh rebuild and two-round retest at revision 9b17dc454fd53e4df12259272ab4872df95b69a5 again matched all eight greedy IDs in both rounds. Automatic prefill measured 395.46 and 462.69 tok/s; layer-major measured 553.62 and 552.37 tok/s. Reported H2D was 200,688.6 MB automatic vs 102,274.4 MB layer-major per round. A second 16K active-path retest used docs/DETAILS.md as a different prompt body and also matched all eight IDs in both rounds; its body hash and run rows are in the results JSON. That sample measured 384.69/438.88 tok/s automatic and 413.90/423.66 tok/s layer-major, so it supports parity on another corpus but not a general speed claim.
  • A 16K cache-capacity sweep found a correctness limitation on docs/DETAILS.md: at 500 expert-cache slots, automatic returned [248068, 198, ...] while layer-major returned [248068, 271, ...]. Four active-path checks diverged across host and VRAM residual placement, frozen adaptation, and 4,096-token super-chunks. At 625 and 750 slots, the initial one-round pairs matched exactly and agreed with the 1,000-slot continuation. A repeat at 625 slots in two reversed-order rounds matched all eight generated IDs in all four arm runs, with layer-major active twice. Automatic prefill measured 381.61 and 365.42 tok/s; layer-major measured 379.08 and 392.08 tok/s. Reported expert H2D was 341,247.4 MB automatic vs 77,521.3 MB layer-major in both rounds; this sample does not establish a throughput gain or a generally safe capacity. Automatic itself changes continuation at 500 slots, so the minimum safe capacity remains unknown. The 500-slot timings and H2D data are diagnostic only, not performance claims. A forced-follow diagnostic found different residual hashes in the first verify window at position 16,383 while both paths still chose token 248068; at position 16,384 the first-row argmax flipped from 198 to 271. This locates the difference before decoding but does not identify its source operation. Compact hash files and exact forced IDs are in the linked results. Raw FP32 logits show KL 1.36e-12 at the first output position, then KL 0.0495 at the second; automatic's 198-over-271 margin is 0.2896, while layer-major's 271-over-198 margin is 0.3419. A router-order scratch experiment left the output mismatch and increased KL to 0.128, so changing accumulation order alone in that experiment did not restore parity; the experiment was reverted.
  • 32K matched A/B: exact 8-token parity in both rounds. On an RTX 3060 12 GB (driver 610.74, CUDA 13.3, sm_86), IQ3_XXS Qwen3.8-Flash-Next, 1,000 expert slots, two interleaved rounds with reversed arm order, mean prefill was 544.3 tok/s automatic vs 625.5 tok/s layer-major (+14.9%). Reported expert H2D bytes averaged 382,393 MB vs 177,567 MB (−53.6%). These are measured results for this setup; the historical 4–6× report is not claimed.
  • 64K matched A/B: exact 8-token parity in both reversed-order rounds with layer-major active. On the RTX 3060 12 GB, IQ3_XXS, 1,000 requested expert slots and a 65,544 context, mean prefill was 564.03 tok/s automatic vs 524.73 tok/s layer-major (−7.0%); expert H2D was 749,483.4 MB vs 351,578.4 MB per run (−53.1%). Peak VRAM was 8,459 MiB vs 8,681–8,683 MiB. This is one prompt/configuration and shows less transfer without faster prefill; file cache was not cleared and no cold-cache or NVMe-read claim is made. Structured rows are under active_64k_cache1000_two_round_validation in the results JSON.
  • bench/layer_major_ab.py records exact prompts, commands, logs, output IDs, prefill/decode rates, H2D bytes, and sampled peak RAM/VRAM. It defaults to two reversed-order rounds and requires an active layer-major path unless --allow-fallback is requested.

The 32K comparison used 1,000 cache slots to leave scratch VRAM available; automatic cache sizing did not leave enough scratch and safely fell back. OS file cache was not cleared, so the results are not cold-cache tests or physical NVMe-read measurements. Decode rates varied substantially; no decode speedup or non-regression claim is made.

A temporary per-layer diagnostic at the 16K/500-slot divergence found a tiny layer-0 residual difference (RMS 9.38e-10) growing to RMS 1.19e-4 by layer 14. At QSA layer 15, the attention projection difference reached RMS 0.00413; the selected-cell-list hashes differed (2,051 cells each), followed by the first top-10 expert-order swap. This narrows where the difference is amplified but leaves the initial rounding source unresolved. Using STRATA_SELECT_OLD=1 in both arms matched eight tokens in two reversed-order rounds, but it also changed automatic prefill’s continuation; applying it only to layer-major still diverged from default automatic. This is diagnostic evidence, not a fix. Details are in the linked benchmark results. A test-only active-path run restricted each MMQ group to one expert and still diverged on token 2, so multi-expert grouping alone is not the cause; the temporary override was removed. A layer-0 stage capture found identical post-attention residuals, router IDs/weights, shared-expert output/gate, and all ten routed expert output vectors for the final prompt row. The combined MoE output differs by max 2.38e-7 (RMS 6.39e-9) because automatic applies weighted FP32 FMAs in router order before adding the shared term, while layer-major adds per-expert group sums to an accumulator seeded with the shared term. This identifies the arithmetic-order source; later QSA can amplify it into a greedy flip. A scratch-free direct-FMA accumulator experiment that removed per-group sum rounding still diverged, confirming that group arrival order remains the mismatch; that experiment was reverted.

Remaining validation

The PR stays draft because numerical divergence at 8K and at 16K with a 500-slot cache remains unresolved, and broader end-to-end validation is incomplete. Still needed: tail and recurrent-state coverage, larger multi-token and decode-continuation comparisons, cancellation/recovery, tool/multi-turn workloads, memory-tier sweeps, cold-cache disk measurements, and an NInfer side-by-side run. HIP/ROCm validation is unavailable in the current Windows environment. The all-target Windows build also remains blocked by the existing POSIX setenv call in src/kernels/gdn_chunk_parity.cpp; the engine, prefill, and targeted CUDA test targets build successfully.

The previous upstream submission attempt was denied by GitHub permissions. The exact fork comparison is main...proposal/qwen-layer-major-prefill-ninfer-review. Current upstream PR #4283 skips unused expert copies in the ordinary streamed walk inside run_impl. The active layer-major path returns before run_impl, and it already stages only experts routed across its super-chunk; the two optimizations are therefore behaviorally separate. If the layer-major gate falls back to the ordinary path, both opt-ins can apply there. Both branches touch prefill.cpp, so their source changes still need careful integration.

Structured per-run measurements, prompt hashes, generated IDs and pre-gate mismatch records are in the RTX 3060 results README and results.json. The benchmark host was Windows x64 with an Intel Core i5-14600K (14 cores/20 logical processors), CUDA 13.3.73, and MSBuild 17.14.51.

The boundary immediately below the gate was also checked: at 16,383 tokens the optimized arm fell back with reason a short prompt, and all eight greedy output IDs matched. Its structured result is included in the same results JSON.

Further 8K diagnostics locate the greedy flip: the first output remains ID 248068, then the second-token logits narrowly favor 198 by 0.151 for automatic prefill and 271 by 0.961 for layer-major. The residual hash differs at the first verify window after prefill. This identifies the immediate near-tie decision, but the cause of the prompt-state/logit difference remains unresolved; structured details are in the linked results JSON.

An STRATA_PREFILL_MMQ=0 control at 500 slots matched outputs only because the optimized path fell back (the experts are not on MMQ); it does not exercise layer-major and is not parity evidence.

Current upstream comparison (2026-10-27)

Upstream main has since merged and expanded layer-major prefill in PR #3976, so this branch now overlaps with code already in upstream and should be reviewed as a comparison/diagnostic branch rather than assumed to be a fresh port. I built upstream commit 5e8a4cca and ran the same 16,384-token docs/DETAILS.md prompt at 500 requested expert-cache slots for two reversed-order rounds. Current upstream's layer-major path was active both times, matched all eight generated IDs in both rounds, and measured 377.02 tok/s mean versus 304.10 tok/s automatic on this sample. Upstream logged 116,472.2 MB of prompt copies for layer-major versus 543,615.2 MB for automatic; that is the engine's prompt-copies metric, not direct H2D bytes. This result does not validate the older branch implementation that diverged at the same requested cache size, and it does not establish general performance.

The comparison and all raw run logs are recorded under upstream_main_5e8a_16k_cache500_two_round_validation in the branch results JSON. The separate harness parser now accepts the current upstream layer-major log format and preserves its prompt copies metric independently from direct H2D measurements.

Because the core scheduling feature is now present upstream and the low-cache mismatch was not reproduced on that newer implementation, the remaining value of this PR needs maintainer review against current main. The PR remains draft; its historical branch results continue to describe the tested branch revisions only.

A further current-upstream validation covered 4K, 8K, 16K, 32K, and 64K with 500 requested cache slots, host residuals, and --adapt-every 0. Across two reversed-order rounds, layer-major was active and matched all eight generated IDs at each size. Mean automatic/layer-major prefill rates were 94.93/232.38, 217.99/241.41, 276.68/382.71, 327.83/403.63, and 338.78/491.82 tok/s, respectively. Upstream's prompt copies counter was 74.7%–83.6% lower in the layer-major arm; it is not direct H2D. At 64K, sampled peak VRAM was 7,522–7,528 MiB and sampled system RAM delta was 43.0–43.9 GiB. These remain measurements of one model, prompt body, build, and machine, with no cold-cache clearing. The structured records are under upstream_main_5e8a_cache500_multisize_two_round_validation in the results JSON.

xdCloudy and others added 2 commits October 27, 2026 00:51
Port the layer-major expert scheduling work into the current prefill path, retain bounded scratch with fallback, and guard prompt-cache callbacks while the later super-chunk state is live. Keep the mode opt-in and document its supported configurations and outstanding validation.

Co-authored-by: Tomas Vizdal <tomasvizdal@outlook.cz>

Co-authored-by: durikim1-cpu <236957742+durikim1-cpu@users.noreply.github.com>

Co-authored-by: libratechw <libratechwriter@gmail.com>
@xdCloudy xdCloudy changed the title [Engine] Scope Qwen layer-major prefill with NInfer and prior Strata work [Engine] Add opt-in streamed-expert layer-major prefill Oct 27, 2026
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.

1 participant