Repository navigation
Conversation
…ds with a histogram per warp (CUDA) - 9-12x at 262K-524K cells, a 243K-token prompt +26% on an RTX 3060 On CUDA the register kernel holds 33,792 blocks (~135K cells). Past it qsa_block_topk took the original kernel: 256 threads, each reading its own stretch of the scores on all six passes, one shared histogram updated with atomics. The dispatch follows the capacity (--max-context), so with a context past ~135K that was every prompt batch and every decode window on a card without the sm_75 active-bound dispatch (Niko1221#512), whatever the prompt's length, and on sm_75 every prompt batch past 33,792 blocks and every decode window. block_topk_wide_kernel is the register kernel's layout (1,024 threads, a histogram per warp, warp scans) with the keys read from memory: - the four radix passes: thread t reads blocks t, t + 1024, ... - a histogram has no order, so a warp reads 32 neighbours at a time; - the count and the output: warp w holds `per` rows of 32 consecutive blocks, and the cells are written in the order (warp, row, lane); a row without a selected block is skipped. The selection rule is the original kernel's (radix threshold, ties to the lowest index, cells ascending): identical ids. One launch per call, no scratch memory, no capacity limit. Dispatch (CUDA; HIP is unchanged and does not compile the kernel): where the blocks do not fit the registers the new kernel, except a call that names fewer than 4,608 blocks (a prompt batch below ~18K cells), which keeps the original kernel: its fixed cost is lower there. A decode window (no block count) takes the new kernel at any context. The register kernel's range, the sm_75 active-bound dispatch and the sm_90+ cluster path are unchanged. STRATA_TOPK_OLD=1 keeps the original kernel, as before. qsa_topk_parity (new, with the kernel parity tests): the dispatched top-k against qsa_block_topk_ref on both sides of the register kernel's reach, at short contexts under a long capacity and up to 524,288 cells, called as decode and as the prompt path call it, with continuous scores, 64 score levels and equal scores. RTX 3060 12 GB (sm_86), 256 queries with their block count under a 262,144-cell capacity (524,288 for the last two), original -> dispatched: 16K cells 0.199 -> 0.199 ms (the original kernel) 32K 0.845 -> 0.328 131K 4.70 -> 0.97 262K 16.09 -> 1.71 393K 26.02 -> 2.49 524K 37.53 -> 3.24 4 queries without a count (a decode window): 0.031 -> 0.028 ms at 8K, 0.176 -> 0.083 at 131K, 0.309 -> 0.127 at 262K. IQ3_XXS on that card (--resident-experts, --max-context 262144, --kv int8 --kv-resident 32768), tools/needle_bench.py, STRATA_TOPK_OLD=1 against the default: tok/s of the prompt, and the prompt path's selection phase (STRATA_PREFILL_TIMING): 16,471 tokens 1004.5, 1016.2 -> 989.0, 1013.6 195 -> 195 ms 31,970 1043.3, 1055.5 -> 1060.1, 1060.5 900 -> 731 61,333 932.9 -> 1013.8 4,061 -> 2,428 122,755 904.7, 904.2 -> 992.7, 989.6 20,019 -> 9,034 185,729 819.7 -> 954.7 50,055 -> 20,761 243,176 743.3 -> 933.7 96,326 -> 35,544 The needles are found in every run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
RTX 5090 (sm_120) numbers for #603 and #575, the same text on both PRs. Both of you measured on an RTX 3060. Setup: RTX 5090 32 GB on PCIe 5, Ryzen 9 9950X3D, Windows 11, CUDA 13. The three engines were built from v0.1.38 ( Prompt speed (tokens/s, round 1 / round 2):
First-token logits: byte-identical to v0.1.38 for both PRs at both lengths ( Decode (a 1,000-token chat, On this card: #603 is neutral at 32K and slightly ahead at 125K. #575's split costs 2.5-4.3% of prompt speed even at 125K; the fixed cost #603 describes would explain it. Neither changes decode here. I can't run kernel-level timings ( |
|
Tested on 2× RTX 3090 (sm_86, NVLink), Qwen3.8-Flash-Next IQ3_S, Exactness: our gate (3 greedy prompts with Speed (flashbench, FB_TEMP 1.0, 2 repeats; same night, same box):
So on Ampere at a 262K window: short/medium prompts unchanged, +4.3 % prompt reading at 128K (sampled run +3.9 %), decode +1-3 % (sampled 32K +5.6 %). Consistent with the select fallback being hit at every prompt batch once the window is past the register kernel's reach. Thanks — we're keeping it in our build. |
|
Thanks! Taken into 0.1.39. qsa_topk_parity --selftest and qsa_topk_active_parity pass on the RTX 5070, and output is byte-identical. Our test configs (32K and 131K windows) stay inside the register kernel's reach, so there is no speed change there; the 262K gains are yours and sergqwer's. We took it over #575: they cover the same range, and yours is faster on the cards both of you measured. |
217 upstream commits (0.1.38 -> 0.1.39: Niko1221#646 decode kernels and zero-doorbell verify graph, Niko1221#603 long-context top-k, Niko1221#583 byte-budget ring, Niko1221#598 pipelined layer split, Niko1221#578 helper GPUs, Niko1221#663, Niko1221#540, Niko1221#600 in the experimental build, batch slots Niko1221#465, the SYCL port, older-CPU builds, ...). 23 files conflicted (~70 hunks); every hunk was resolved by hand and checked against both sides. Resolution rules: - The prompt path's stage overlap: upstream's Niko1221#598 (equivalent to our e85f7ee; run() drains the whole chain on every exit), with our release_inherited_pin in the stage threads and our early ids copy for the host grouping (a5a3254) kept. - STRATA_FP16=load's form-aware weights kept wherever upstream's new code assumed BF16: the shared-expert gate (sigmoid+scale fused as upstream), GDN alpha/beta (our kernels; upstream's new gdn_step_commit_kernel added), the hc v3 kernels (upstream's EXACT_T with our dot8w). - verify.cpp: upstream's all-resident fast path, plan copy and Niko1221#578 rows, with our 2-GPU fetch overlap and partner fetch back in the non-resident branch; mapped buffers stay portable (the partner reads them). - Volta prompt attention: upstream's dispatch (the m8n8k4 kernel on sm_70 in this build) with our split hi/lo accumulator chains inside it; the wmma kernel stays opt-in (STRATA_ATTN_WMMA=1 without STRATA_ATTN_M884). - qsa_decode_attn.cu: upstream's (our adf56cc imported the same PR Niko1221#540; upstream's is it plus review fixes, bit-exact). - POWER / aarch64: upstream's new x86 intrinsics in verify.cpp through the portable strata_store_fence / strata_cpu_pause; its AVX2 start-up checks and log line behind x86 guards; stubs for the new CPU probes; the router lookahead keeps portable.cpp's VSX dot on non-x86 (upstream's dispatch would have sent POWER to its scalar fallback); the pool's ThreadAffinity (any CPU count) with our note_host_pin; SMT2 selection with Niko1221#642's sort. - setup.py / tests: both sides (CUDA 12 engine, ISA floor, ARM / non-x86). Checked: - every line either side added since 0.1.38 is in the merge or replaced on purpose (script over all 45 files both sides changed); - all parity tests pass (iq, native_expert on the real shards, mmvq_multi, gr, gdn, gdn_rec, qsa, qsa_topk, elementwise, router_dot, shared_expert, kv_*, sampler, rope, s2/s_gemv, gemm_bf16, ...); - teacher-forced 2K and 32K logs byte-identical to the pre-merge build (4 GPUs, launch config); - 4 GPUs: decode text identical, same windows and acceptance, 28.25 -> 26.9 ms/window (~86 -> ~90 tok/s); prompts 7.8K 3,992 -> 4,038, 65K 7,102 -> 7,172, 123K 7,323 -> 7,334 tok/s; --prefill-ring 320 applies; - 2 GPUs: decode texts identical to the pre-merge binary on three prompts, decode 61.5-71.0 -> 63.5-72.7 tok/s, 65K prompt 3,975 -> 4,083. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016ZHWKWn7UYZV64KwpGrWuN
4 participants
What
On CUDA the register kernel (
block_topk_reg_kernel) holds 33,792 blocks (~135K cells). Past that,qsa_block_topktakes the original kernel (block_topk_kernel): 256 threads, each reading its own stretch of the scores on all six passes, and one shared histogram updated with atomics. The dispatch follows the capacity (--max-context), so with--max-context 262144:This adds one kernel for that range,
block_topk_wide_kernel: the register kernel's layout (1,024 threads, a histogram per warp, warp scans) with the keys read from memory instead of registers.treads blockst,t + 1024, ...: a histogram has no order, so a warp reads 32 neighbours at a time.wholdsperrows of 32 consecutive blocks, and the cells are written in the order (warp, row, lane). The per-warp counts give each warp its first output position and the tied cells before it; a row without a selected block is skipped.STRATA_TOPK_OLD=1keeps the original kernel, as before.qsa_topk_parity(new, next to the other kernel parity tests) compares the dispatched top-k withqsa_block_topk_ref.Relation to #575
#575 covers the same range with a split: the register kernel on two halves of the blocks, then a merge kernel. I built #575 (
cb29a2a) and this branch on the same base (99f3dbd, 0.1.38) and measured both on the same RTX 3060; the numbers are in the tables below.--max-context: the 8K, 16K and 32K prompts spend more time in the selection than stock does.I have no RTX 3090, so I cannot say how the two compare on #575's test card.
Measured
Setup: RTX 3060 12 GB (sm_86) on PCIe x16, Ryzen 5 5600X, 48 GB DDR4-2933, Debian 13, CUDA 13.4, base
99f3dbd(0.1.38). Qwen3.8-Flash-Next IQ3_XXS on that one card,--resident-experts --max-context 262144 --kv int8 --kv-resident 32768 --prefill auto --spec 4.Kernel (
qsa_topk_parity CONTEXT 256 50 CAPACITY: 256 queries with their block count, capacity 262,144 cells, the last two rows 524,288; ms per call). "Register" is #512's dispatch forced on this card withSTRATA_TOPK_ACTIVE_ANY=1; it stops at 33,792 blocks, and its column is from a separate run of the same tool.Between ~20K and ~110K cells the register kernel is faster than this one by up to 0.15 ms per call; I left #512's restriction to sm_75 as it is.
A decode window's call (4 queries, no block count, capacity 262,144): 0.031 -> 0.029 ms at 2,100 cells, 0.031 -> 0.028 at 8,192, 0.054 -> 0.034 at 32,768, 0.176 -> 0.083 at 131,072, 0.309 -> 0.127 at 262,144. #575 leaves these calls on the original kernel.
Prompts (
tools/needle_bench.py --depths 50, a fresh engine per run,STRATA_PREFILL_TIMING=1in every run; "original" is this branch withSTRATA_TOPK_OLD=1). Several values in a cell are separate runs. Where a prompt reused a 16,384- or 32,768-token checkpoint of the previous prompt in the same engine, the tokens read are in brackets; the reuse is the same in every arm.Prompt speed, tok/s:
The 243K prompt takes 260 s instead of 327 s.
The selection phase of the prompt path (
qsa selectin the timing line: the block scores and the top-k), ms:The 7,884-token prompt is the first of each engine and its tok/s varies from run to run in every arm (600 to 925 overall), so it is in the selection table only. With #512's dispatch forced on this card (
STRATA_TOPK_ACTIVE_ANY=1, one run) the selection takes 96, 286, 810 and 8,842 ms at 7,884, 16,471, 31,970 and 122,755 tokens, and the 122,755-token prompt reads at 998.2 tok/s.Other setups (the same model and flags,
STRATA_TOPK_OLD=1against the default, tok/s of the prompt):--mmap-expertsDecode (the x16 card, one long prompt, then 512 tokens generated three times, greedy, tok/s):
Repeated requests of one prompt do not always produce the same text, also within one arm, so only requests with identical text compare like for like. At 243K the first two requests produced identical text in both arms: 15.3 -> 14.9 s and 14.3 -> 14.0 s for the 512 tokens. At 122K no request pair has identical text, and the arms differ by less than the requests of one arm do.
Test
qsa_topk_parity --selfteston both cards: 9 context / capacity pairs (9K; both sides of the register kernel's reach at 135,160 and 135,176 cells; 2,100, 9K and 32K under a 262K capacity; 262K; 300K under 524K; 524K), each called as a decode window calls it (8 queries, no block count) and as the prompt path does (64 queries with their block count), each with continuous scores, 64 score levels and equal scores. All 18 cases: identical ids. It also passes withSTRATA_TOPK_ACTIVE_ANY=1, where a counted call past 33,792 blocks takes the new kernel.qsa_topk_active_parity(prefill: use active QSA top-k bounds on Turing with large KV capacity #512's test) passes on both cards; with this change its call without a bound takes the new kernel at the 262,144 capacity.decode_cluster_parityskips on these cards (below sm_90).Not run: a HIP build (the kernel and its dispatch branch are inside
#if !defined(__HIPCC__)), and a card other than the RTX 3060.🤖 Generated with Claude Code