Skip to content

[Engine] gfx906: opt-in two-stage split for the decode QSA top-k - #3987

Open
MarcoTool wants to merge 1 commit into
Niko1221:mainfrom
MarcoTool:gfx906-topk-split
Open

MarcoTool wants to merge 1 commit into
Niko1221:mainfrom
MarcoTool:gfx906-topk-split

Conversation

@MarcoTool

Copy link
Copy Markdown

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_kernel past the register kernel's reach (82922f5), and the multi-CTA split from #2072 is
CUDA-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 variable
unset, the dispatch is unchanged.

qsa_select_bench on 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:

context ref wide, current default split ids identical to ref
40K 0.175 ms 0.104 ms 0.076 ms 3/3
64K 0.266 ms 0.114 ms 0.081 ms 3/3
131K 0.516 ms 0.210 ms 0.086 ms 3/3
258K 0.744 ms 0.324 ms 0.118 ms 3/3

Approach

  • Stage A: each query's blocks are cut into the fewest chunks of at most 4,104 blocks, decided on the device from
    n_bid, with one CTA per chunk. In LDS, block_topk_kernel's rule picks each chunk's candidates. Every block the
    whole 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.
  • Stage B: one CTA per query runs the same rule over the candidates and writes the cells, so the ids equal the
    reference's.
  • Both stages pick each radix digit with a parallel scan of the 256 bins.

The split is taken only for decode calls (no active-block count) past the register kernel's reach, with
nq <= 16 and a geometry that fits the candidate buffer; anything else falls through to the existing dispatch.

Limits

  • The candidate buffers are device globals, like the CUDA split's, so two streams on one device must not run this
    top-k at the same time.
  • The two-stage kernels and the dispatch are compiled only under STRATA_HIP_GFX906; CUDA and wave32 HIP builds are
    unaffected.

Validation

  • Built qsa_select_bench for gfx906, CUDA sm_86 and wave32 HIP gfx1201; no warnings from qsa_select.cu.
  • Ran the benchmark above on an MI50 (ROCm 7.14, Ubuntu 22.04, kernel 6.8).
  • The full gfx906 engine does not build on this base for an unrelated reason (src/core/dflash2.cpp, fixed
    separately 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.cpp gains QSA_BENCH_TOPK_DECODE=1, which calls the top-k without an active-block
count, as a decode window does, so the decode dispatch can be measured.

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