Repository navigation
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
--expert-cache 500setup 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.--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.Validation
Completed checks:
strata,strata_prefill, andlayer_major_kernel_testatsm_75and native RTX 3060sm_86.layer_major_kernel_test, including shared-expert initialization, routed scatter accumulation, andcopy_i32.a short prompt, so the optimized path was active zero times. This confirms fallback parity fordocs/DETAILS.mdat 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 atf3feced2; the branch head adds benchmark records only. Structured rows are undercachefloor_fallback_current_head_validationin the results JSON.9b17dc454fd53e4df12259272ab4872df95b69a5again 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 useddocs/DETAILS.mdas 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.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.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.active_64k_cache1000_two_round_validationin the results JSON.bench/layer_major_ab.pyrecords 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-fallbackis 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=1in 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 max2.38e-7(RMS6.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
setenvcall insrc/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 beforerun_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 touchprefill.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=0control 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
5e8a4ccaand ran the same 16,384-tokendocs/DETAILS.mdprompt 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 ofprompt copiesfor 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_validationin the branch results JSON. The separate harness parser now accepts the current upstream layer-major log format and preserves itsprompt copiesmetric 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'sprompt copiescounter 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 underupstream_main_5e8a_cache500_multisize_two_round_validationin the results JSON.