Repository navigation
Conversation
…FX906_TOPK_SPLIT=1)
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.
Problem and result
At long context, a decode window's indexer top-k on gfx906 runs one CTA per query: HIP takes
block_topk_wide_kernelpast the register kernel's reach (82922f5), and the multi-CTA split from #2072 isCUDA-only. A window with 1–6 queries therefore keeps 1–6 of the MI50's 60 CUs busy.
This PR adds a two-stage split for the gfx906 build, enabled with
STRATA_GFX906_TOPK_SPLIT=1. With the variableunset, the dispatch is unchanged.
qsa_select_benchon one MI50 at this PR's base, 3 queries, capacity 262,144, the decode dispatch(
QSA_BENCH_TOPK_DECODE=1), 200 repetitions, two runs each with identical results:Approach
n_bid, with one CTA per chunk. In LDS,block_topk_kernel's rule picks each chunk's candidates. Every block thewhole query selects is among its chunk's candidates, because a query's selection restricted to a chunk is a prefix of
that chunk's order.
reference's.
The split is taken only for decode calls (no active-block count) past the register kernel's reach, with
nq <= 16and a geometry that fits the candidate buffer; anything else falls through to the existing dispatch.Limits
top-k at the same time.
STRATA_HIP_GFX906; CUDA and wave32 HIP builds areunaffected.
Validation
qsa_select_benchfor gfx906, CUDA sm_86 and wave32 HIP gfx1201; no warnings fromqsa_select.cu.src/core/dflash2.cpp, fixedseparately in [Engine] gfx906: fix the build of src/core/dflash2.cpp #3986), so there is no end-to-end run on main. On a 0.1.43.1 base with the same kernels, decode at
258K went from 37.8 tok/s with the wide kernel to 39.4 tok/s with the split (Flash-Next IQ3_S, 2× MI50), and the
greedy answer was byte-identical to the reference build.
src/kernels/qsa_select_bench.cppgainsQSA_BENCH_TOPK_DECODE=1, which calls the top-k without an active-blockcount, as a decode window does, so the decode dispatch can be measured.