[ruibeikaa]

Problem and result

Everything below was measured on four GP100 dies (two PH402 SKU 200 boards, 32 GB HBM2 per die, 128 GB in all) with every expert resident in VRAM: no CPU share, no PCIe share, no adaptive tier. A rig that keeps experts in RAM or streams them over PCIe was not measured. HIP, SYCL and Metal were not built.

A decode window's QSA select runs one CTA per query for the top-k (block_topk_wide_kernel past ~135K cells) and the warp-shuffle block scorer (block_scores_multi_kernel), in each of the 12 QSA layers. Both grow linearly with the context: on one GP100 die at 300K cells a top-k call takes 0.39 ms with 1-5 SMs busy and a scores call 0.22 ms (two queries), about 7.4 ms of a window over the four dies. At short contexts the same work is small, and an extra kernel is pure overhead.

This adds a multi-CTA decode top-k and a shuffle-free decode scorer, which return the same ids and the same score bits as the kernels they stand for. With STRATA_SELECT_LONG=1 the verifier keeps a second set of window graphs that uses them once a window reaches STRATA_SELECT_LONG_AT cells (default 65536); below that the graphs are today's. On four GP100 dies with --pipeline-windows 3, the 400 tokens decoded at 115K context went 51.0 -> 52.5 tok/s (+2.9%; 30.75 -> 29.86 ms a window), the prompts below 64K stayed where they were, and every reply and engine state hash stayed the same on 15 requests. Unset (the default), nothing is allocated or captured and every graph is the one before.

Related discussion

None (opt-in kernels in the STRATA_EXPERIMENTAL_SM60 build, no default change). This branch is main plus three commits and needs none of my other PRs. The main e88ee39 run below was on a branch that also carried #3455, #3760, #3819 and #3751, all on in the arm before this one.

Approach

• qsa_select.cu, inside !__HIPCC__ && STRATA_EXPERIMENTAL_SM60:
  • block_topk_multi_*: six kernels over several CTAs per query (one per 2,048 blocks of capacity, at most 32), exchanging through caller scratch - no clusters, no spin waits, capturable, nothing carried between calls. The digit rule is block_topk_kernel's as a scan and the emit is block_topk_wide_kernel's row emit offset by the earlier CTAs' counts, so the ids are the same.
  • block_scores_dx_kernel<1..8>: the multi scorer's per-head sums built in registers as the same 32-leaf xor-butterfly tree, without __shfl_xor_sync, then relu over the heads in the warp kernel's order; the tail block stays block_scores_tail_kernel's. The score bits are the same.
  • New public functions qsa_topk_multi_scratch_ints, qsa_block_topk_long, qsa_block_scores_long. Outside the build gate, or outside their conditions (scratch given, no active-block count, at most 16 / 8 queries), they forward to qsa_block_topk / qsa_block_scores. The existing declarations and bodies are unchanged. STRATA_TOPK_MULTI=0 / STRATA_SCORES_DX=0 turn either half off inside the long graphs (both default on at compute capability 6.x).
• verify.cpp:
  • exec_[2][9], the Fresh prompts over --short-read decode !!!! (token 0) when the GPU holds 100% of the experts (48 GB card) #871 doorbell set and the launched flags get the variant as first index; sel_long_ picks it before a capture.
  • On an overlapped layer split, run() hands the choice to the next stage before capturing that stage's window, and prelaunch() derives it from the window's position, so a stage never launches the other variant.
  • capture_all captures both variants with the switch on, and pl_launch picks by the window's context, so --pipeline-windows takes the long select too.
  • Batch rows keep the one-CTA select. The scratch is carved from the arena only with the switch on, in a build that has the kernels.
• Metal: three stubs (its windows run through the Metal engine and never reach them). SYCL has its own verifier and only sees the new declarations. HIP: the forwards.

Validation

Rig: four GP100 dies (compute capability 6.0, application clocks 1050 MHz graphics / 715 MHz memory, 140 W cap), Windows 11, driver 581.80 (TCC), CUDA 12.9, VS 2022. Build -DSTRATA_EXPERIMENTAL_SM60=ON -DCMAKE_CUDA_ARCHITECTURES=60, ggml at the pinned 3cf03257f, unpatched. Model: Swift 1.5 IQ3_S (abliterated GSQ-RCO GGUF, dense matrices in Q8_0).

Passed:

• main e88ee39 (engine 0.1.45.3) plus my GP100 stack, one exe for every arm, one engine start per arm, greedy; a ladder in which each arm adds one switch to the arm before. Arm S4 = [Engine] GP100 (sm_60): 4-bit slots for resident IQ2_S / IQ3_XXS / IQ3_S gate/up experts, decode +9.7% (opt-in, STRATA_EXPERT_NIBBLE=1) #3455, [Engine] A third decode window in flight on a layer split (--pipeline-windows 3, opt-in): a host-side scheduler, a slot pool, per-stage undo, an 11-output draft chain #3760, [Engine] Layer split: a stage's commit without the host wait, as on one GPU (STRATA_SPLIT_COMMIT_ASYNC=1, opt-in) #3819 and [Engine] sm_60: the MTP drafter's draft head and projections on the Pascal Q8_0 GEMV too (STRATA_MTP_Q8_SM60=1, opt-in) #3751 on; arm S5 = S4 + STRATA_SELECT_LONG=1 (threshold at its default).
  • Engine arguments, the same in every arm: --serve --expert-cache auto --prefill auto --spec 4 --spec-min-p 0.85 --mtp <drafter> --max-context 1048576 --kv int8 --layer-split 12,24,36 --mtp-window 8192 --pipeline-windows 3 --conversation-cache-mib 32768 --conversation-cache-slots 8 --rope-scaling yarn --rope-scale 4. Environment, the same in every arm: STRATA_Q8_SM60=1 STRATA_STAGE_TRIM=1 STRATA_PREFILL_PIPE=70 STRATA_MMQ_RESIDENT_SORT_NE=1 STRATA_DF_BRANCH=1 STRATA_FUSE_DECODE=1 STRATA_FUSE_WAIT=2 STRATA_SIDE_MIXER=2 STRATA_GR_UP75=1 STRATA_B6_MMVQ_ROWS=1 STRATA_ADAPT_LAG=2 STRATA_ADAPT_NOWAIT=1 STRATA_DECODE_TIMING=1. Each start's switches were checked in the engine process's environment.
  • Requests: five short prompts (77-82 tokens in, 400-512 out), a 9.8K-token document + 300 out, a fresh 26K-token read + 400 out and an 88K-token append to it (115K context) + 400 out - the eight benchmark prompts, decode speed = their tokens out / decode time - then prompt reads of 240 / 1.15K / 7.7K / 33K tokens and a 26K-token conversation with a 1.1K-token append (16 out each). A parity pass with STRATA_STATE_HASH=1 and a 250K-token read added (15 requests), and a timing pass without debug switches (14 requests).
  • Parity, S5 against S4: every reply hash and every STRATA_STATE_HASH item equal on all 15 requests, including the two that decode past 64K cells, where the long graphs run (400 tokens at 115K, 16 tokens after the 250K read).
  • Decode, S4 -> S5, timing pass / parity pass: the 115K request 51.0 -> 52.5 / 50.1 -> 52.4 tok/s; the five short prompts 58.65 -> 58.63 / 58.40 -> 58.42; the 26K read's 400 tokens 50.3 -> 50.3 / 50.2 -> 50.2; all eight token-weighted 56.20 -> 56.39 (+0.34%) / 55.90 -> 56.22. The 16 tokens after the 250K read (parity pass only, too few to weigh): 46.3 -> 50.0. For scale, the unchanged engine ran twice in the same timing pass: 47.21 and 47.33 tok/s.
  • Prompt reads within 0.6%; start-up 278 -> 277 s; peak VRAM +108 to +112 MiB per die (25,427 / 25,113 / 25,789 / 29,301 -> 25,537 / 25,221 / 25,897 / 29,413 MiB: the second graph set and the scratch).
• On another day, with STRATA_SM60_TABLE=0 in both arms, adding STRATA_SELECT_LONG=1 together with STRATA_IQ_STAGE_TABLES=1 took the 115K request 52.0 -> 53.6 tok/s and the five short prompts +0.15%, the same replies; the two switches were not separated there.
• Older tree (0.1.41-based, my GP100 branch of that time, serial decode, not this revision): per call on one die at 300K cells and two queries, top-k 390 -> 71 us and scores 217 -> 154 us; in the engine at a 300K context 51.9 -> 47.1 ms a window (38.6 -> 41.8 tok/s), short contexts unchanged, 8 replies the same. The two kernels were checked there byte for byte against qsa_block_topk_ref and block_scores_kernel: contexts 1K-1M, 1-5 queries, 1/3/7/32 CTAs per query, capacities 131K-1M, tie-heavy scores (7 values, mostly zeros, -0.0 / 0.0 / negatives, NaN).
• This branch, rebased onto main 5c8dd32: sm_60 with STRATA_EXPERIMENTAL_SM60=ON (strata) and sm_86 without it (strata), both built without errors. The sm_60 exe contains block_scores_dx_kernel and block_topk_multi_*; the sm_86 exe contains neither.

Failed: none observed.

Not run:

• HIP, SYCL and Metal builds (no such machine here; the kernels are behind the CUDA-only gate, Metal has the stubs, HIP the forwards).
• A GPU run of this branch on main 5c8dd32: the numbers above are from the revisions named there; the rebased branch was only built.
• Serial decode on main, including the hand-over of the variant to the next stage on an overlapped split (Fresh prompts over --short-read decode !!!! (token 0) when the GPU holds 100% of the experts (48 GB card) #871): every arm on e88ee39 had --pipeline-windows 3, where decode windows go through pl_launch. Serial decode with the long select was measured only on the older tree, which predates Fresh prompts over --short-read decode !!!! (token 0) when the GPU holds 100% of the experts (48 GB card) #871.
• A kernel-level parity run of the two kernels on this revision (done on the older tree; here the evidence is the end-to-end hashes).
• Thresholds other than 65536, decode past 250K on this revision, cards other than GP100, packs other than Swift 1.5.

Limits

• Past the threshold only. Below it the extra kernels cost more than they save: on the older tree about 0.5 ms a window more at 1K and break-even at 30K, which is why the threshold exists.
• GP100 builds with STRATA_EXPERIMENTAL_SM60 only. In another build the switch still keeps the second graph set, with the plain kernels in it: no effect but the graphs' memory.
• +108 to +112 MiB of VRAM per die with the switch on.

CI note: the red metal-link check is not caused by this PR. The job fails while compiling upstream's metal/src/mtl_rt.mm (lines 69 and 92: MTLGPUFamilyMetal4, MTLCompileOptions.mathMode, MTLMathModeFast / MTLMathModeSafe), which the macos-14 runner's SDK does not provide, so it stops before linking. The only file this PR changes under metal/ is metal/src/engine/stubs_0145.cpp (a link stub for the new entry point).

[Niko1221]

Thank you. Fixed in 0.2.7: Multi-CTA decode top-k and a shuffle-free decode scorer, opt-in STRATA_SELECT_LONG=1. Your name is on the commit.

Release: https://github.com/Niko1221/Strata/releases/tag/v0.2.7