From bc0db93289b844135099cd1c16b8372feeac71e6 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 29 Sep 2026 03:44:07 +0300 Subject: [PATCH] perf(qsa): reduce block top-k candidates in levels past 16 splits --- python/freetoken/kernel/triton/qsa/topk.py | 226 ++++++++++++++++++--- tests/models/qwen4_exp/test_qsa_kernels.py | 39 +++- 2 files changed, 225 insertions(+), 40 deletions(-) diff --git a/python/freetoken/kernel/triton/qsa/topk.py b/python/freetoken/kernel/triton/qsa/topk.py index ae57566cab..0f37a0f9d9 100644 --- a/python/freetoken/kernel/triton/qsa/topk.py +++ b/python/freetoken/kernel/triton/qsa/topk.py @@ -17,10 +17,18 @@ top-k is a subset of that union -- a global winner beats at most k-1 columns, so it beats all but at most k-1 of its own slice -- so the answer stays exact. ``_split_plan`` derives the geometry from the buffer width alone, which is fixed when a graph is captured. + +A row wide enough to leave more candidates than one program holds resident (262,144 columns, +a 1M-token window, leaves 32 x k) gets reduction levels between the two phases: each program +of a level takes ``_MAX_RESIDENT`` consecutive candidates of the level before and writes their +own top-k, until the survivors fit the merge. The same subset argument holds per level, and +candidates stay in ascending column order throughout, so ties still go to the lowest column. """ from __future__ import annotations +from dataclasses import dataclass + import torch import triton import triton.language as tl @@ -125,6 +133,44 @@ def _compact( return above, equal +@triton.jit +def _live_candidates( + visible, + CHUNK: tl.constexpr, + N_SPLITS: tl.constexpr, + TOP_K: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, +): + """Candidate slots level ``LEVEL`` filled for a row: its splits past the row's visible tail + are one skipped suffix, so every later stage keeps costing what the live part costs.""" + live = tl.minimum(tl.cdiv(visible, CHUNK), N_SPLITS) * TOP_K + for _ in tl.static_range(LEVEL): + live = tl.cdiv(live, RESIDENT) * TOP_K + return live + + +@triton.jit +def _emit_candidates( + key, + column, + k_eff, + key_out, + col_out, + TOP_K: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +): + """Write a resident tile's own top-k as ranked (key, column) candidates; returns how many.""" + prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES) + rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) + write = take & (rank < TOP_K) + tl.store(key_out + rank, key.to(tl.int32, bitcast=True), mask=write) + tl.store(col_out + rank, column, mask=write) + return above + tl.minimum(equal, ties) + + @triton.jit def _qsa_block_topk_kernel( logits_ptr, @@ -225,17 +271,95 @@ def _qsa_topk_split_kernel( if k_eff > 0: offsets = tl.arange(0, CHUNK) key = _load_keys(logits_ptr + row.to(tl.int64) * stride_logits_row + base, offsets, limit) - prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES) - rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) - write = take & (rank < TOP_K) - tl.store(key_ptr + slot + rank, key.to(tl.int32, bitcast=True), mask=write) - tl.store(col_ptr + slot + rank, (base + offsets).to(tl.int32), mask=write) - emitted = above + tl.minimum(equal, ties) + emitted = _emit_candidates( + key, (base + offsets).to(tl.int32), k_eff, key_ptr + slot, col_ptr + slot, + TOP_K, BINS, RADIX, PASSES, + ) # A key of 0 is the dead-column sentinel, so the merge needs no separate count per slot. pad = tl.arange(0, PAD_K) tl.store(key_ptr + slot + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) +@triton.jit +def _reduce_tile( + in_key, + in_col, + out_key, + out_col, + count, + k_eff, + TOP_K: tl.constexpr, + BLOCK: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +): + """Top-k of ``count`` candidates held as one resident tile; returns how many it wrote.""" + tl.static_assert(BLOCK <= 0xFFFF, "packed cumsum keeps 16 bits per half") + offsets = tl.arange(0, BLOCK) + inside = offsets < count + key = tl.load(in_key + offsets, mask=inside, other=0).to(tl.uint32, bitcast=True) + column = tl.load(in_col + offsets, mask=inside, other=-1) + return _emit_candidates(key, column, k_eff, out_key, out_col, TOP_K, BINS, RADIX, PASSES) + + +@triton.jit +def _qsa_topk_reduce_kernel( + visible_ptr, + in_key_ptr, + in_col_ptr, + out_key_ptr, + out_col_ptr, + stride_scratch_row, + num_columns, + TOP_K: tl.constexpr, + PAD_K: tl.constexpr, + CHUNK: tl.constexpr, + N_SPLITS: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, + BLOCK_SMALL: tl.constexpr, + BLOCK_MID: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +) -> None: + """A reduction level: one program per (row, ``RESIDENT`` candidates of level ``LEVEL - 1``), + writing their own top-k as this level's candidates.""" + row = tl.program_id(0) + split = tl.program_id(1) + visible = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0) + k_eff = tl.minimum(visible, TOP_K) + live = _live_candidates(visible, CHUNK, N_SPLITS, TOP_K, LEVEL - 1, RESIDENT) + base = split * RESIDENT + count = tl.minimum(tl.maximum(live - base, 0), RESIDENT) + source = row.to(tl.int64) * stride_scratch_row + base + slot = row.to(tl.int64) * stride_scratch_row + split * TOP_K + in_key, in_col = in_key_ptr + source, in_col_ptr + source + out_key, out_col = out_key_ptr + slot, out_col_ptr + slot + + emitted = 0 + # As in the merge: residency costs the whole tile, so a short range takes a narrower one. + if count > 0: + if count <= BLOCK_SMALL: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, BLOCK_SMALL, BINS, RADIX, PASSES, + ) + elif count <= BLOCK_MID: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, BLOCK_MID, BINS, RADIX, PASSES, + ) + else: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, RESIDENT, BINS, RADIX, PASSES, + ) + pad = tl.arange(0, PAD_K) + tl.store(out_key + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) + + @triton.jit def _merge_tile( key_row, @@ -274,6 +398,8 @@ def _qsa_topk_merge_kernel( PAD_K: tl.constexpr, CHUNK: tl.constexpr, N_SPLITS: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, BLOCK_SMALL: tl.constexpr, BLOCK_MID: tl.constexpr, BLOCK_FULL: tl.constexpr, @@ -281,13 +407,11 @@ def _qsa_topk_merge_kernel( RADIX: tl.constexpr, PASSES: tl.constexpr, ) -> None: - """Phase 2: one program per row over the candidates phase 1 left behind.""" + """Phase 2: one program per row over the candidates the last level ``LEVEL`` left behind.""" row = tl.program_id(0) limit = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0) k_eff = tl.minimum(limit, TOP_K) - # Candidates sit chunk-major, so the splits past the visible tail are one skipped suffix - # and the merge keeps costing what the live part of the row costs. - candidates = tl.minimum(tl.cdiv(limit, CHUNK), N_SPLITS) * TOP_K + candidates = _live_candidates(limit, CHUNK, N_SPLITS, TOP_K, LEVEL, RESIDENT) key_row = key_ptr + row.to(tl.int64) * stride_scratch_row col_row = col_ptr + row.to(tl.int64) * stride_scratch_row out_row = out_ptr + row.to(tl.int64) * stride_out_row @@ -325,27 +449,37 @@ def _qsa_topk_merge_kernel( tl.store(out_row + pad, -1, mask=(pad >= emitted) & (pad < TOP_K)) -def _split_plan(columns: int, top_k: int) -> tuple[int, int] | None: - """``(chunk, n_splits)`` for the split+merge path, or None to keep the one-program path.""" +@dataclass(frozen=True) +class _SplitPlan: + """``chunk`` columns per first-level program, and the program count of every level: the + first splits the row, each later one reduces the candidates of the level before it.""" + + chunk: int + levels: tuple[int, ...] + + +def _split_plan(columns: int, top_k: int) -> _SplitPlan | None: + """The split geometry of a ``columns``-wide buffer, or None to keep the one-program path.""" if top_k <= 0 or columns <= _MIN_CHUNK: return None max_splits = _MAX_RESIDENT // triton.next_power_of_2(top_k) if max_splits < 2: return None - chunk = max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits))) - if chunk > _MAX_RESIDENT: - return None + chunk = min(_MAX_RESIDENT, max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits)))) n_splits = -(-columns // chunk) # Merging n_splits*top_k candidates has to be cheaper than scanning the row once. if n_splits < 2 or n_splits * top_k >= columns: return None - return chunk, n_splits + levels = [n_splits] + while levels[-1] * top_k > _MAX_RESIDENT: + levels.append(-(-levels[-1] * top_k // _MAX_RESIDENT)) + return _SplitPlan(chunk, tuple(levels)) def qsa_block_topk_scratch_width(columns: int, top_k: int) -> int: """int32 columns of scratch ``qsa_block_topk`` wants per row; 0 when it needs none.""" plan = _split_plan(columns, top_k) - return 0 if plan is None else 2 * plan[1] * top_k + return 0 if plan is None else 2 * top_k * sum(plan.levels) def qsa_block_topk( @@ -396,24 +530,31 @@ def qsa_block_topk( ) return out - chunk, n_splits = plan - half = n_splits * top_k + width = 2 * top_k * sum(plan.levels) if scratch is None: - scratch = torch.empty((rows, 2 * half), dtype=torch.int32, device=logits.device) + scratch = torch.empty((rows, width), dtype=torch.int32, device=logits.device) elif ( scratch.ndim != 2 or scratch.shape[0] < rows - or scratch.shape[1] < 2 * half + or scratch.shape[1] < width or scratch.dtype != torch.int32 or scratch.stride(1) != 1 ): raise ValueError( f"QSA block top-k needs a row-contiguous int32 scratch of at least " - f"[{rows}, {2 * half}], got {tuple(scratch.shape)} {scratch.dtype}" + f"[{rows}, {width}], got {tuple(scratch.shape)} {scratch.dtype}" ) - keys = scratch[:rows, :half] - cols = scratch[:rows, half : 2 * half] - _qsa_topk_split_kernel[(rows, n_splits)]( + # Each level's candidates: its keys, then their columns, level after level along the row. + candidates: list[tuple[torch.Tensor, torch.Tensor]] = [] + offset = 0 + for programs in plan.levels: + half = programs * top_k + candidates.append( + (scratch[:rows, offset : offset + half], scratch[:rows, offset + half : offset + 2 * half]) + ) + offset += 2 * half + keys, cols = candidates[0] + _qsa_topk_split_kernel[(rows, plan.levels[0])]( logits, visible, keys, @@ -423,14 +564,39 @@ def qsa_block_topk( columns, TOP_K=top_k, PAD_K=pad_k, - CHUNK=chunk, + CHUNK=plan.chunk, BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1, ) - merge_block = triton.next_power_of_2(half) + for level in range(1, len(plan.levels)): + (in_keys, in_cols), (out_keys, out_cols) = candidates[level - 1], candidates[level] + _qsa_topk_reduce_kernel[(rows, plan.levels[level])]( + visible, + in_keys, + in_cols, + out_keys, + out_cols, + keys.stride(0), + columns, + TOP_K=top_k, + PAD_K=pad_k, + CHUNK=plan.chunk, + N_SPLITS=plan.levels[0], + LEVEL=level, + RESIDENT=_MAX_RESIDENT, + BLOCK_SMALL=min(1024, _MAX_RESIDENT), + BLOCK_MID=min(4096, _MAX_RESIDENT), + BINS=_BINS, + RADIX=_RADIX, + PASSES=_PASSES, + num_warps=8, + num_stages=1, + ) + keys, cols = candidates[-1] + merge_block = triton.next_power_of_2(keys.shape[1]) _qsa_topk_merge_kernel[(rows,)]( visible, keys, @@ -441,8 +607,10 @@ def qsa_block_topk( columns, TOP_K=top_k, PAD_K=pad_k, - CHUNK=chunk, - N_SPLITS=n_splits, + CHUNK=plan.chunk, + N_SPLITS=plan.levels[0], + LEVEL=len(plan.levels) - 1, + RESIDENT=_MAX_RESIDENT, BLOCK_SMALL=min(1024, merge_block), BLOCK_MID=min(4096, merge_block), BLOCK_FULL=merge_block, diff --git a/tests/models/qwen4_exp/test_qsa_kernels.py b/tests/models/qwen4_exp/test_qsa_kernels.py index 4b678b3bdc..6715519f05 100644 --- a/tests/models/qwen4_exp/test_qsa_kernels.py +++ b/tests/models/qwen4_exp/test_qsa_kernels.py @@ -331,7 +331,17 @@ def test_torch_topk_env_picks_the_fallback(monkeypatch): # Block top-k, split + merge path (wide buffers) # -------------------------------------------------------------------------------------- -SPLIT_CHUNK = 4096 # _split_plan's chunk for every buffer these tests use +# (columns, _split_plan's chunk). A 1M-token window scores 262,144 compressed rows, and its 32 +# splits leave more candidates than one program holds, so a reduction level runs before the merge. +SPLIT_SHAPES = [(65536, 4096), (262144, 8192)] + + +def _split_chunk(columns: int, width: int) -> int: + from freetoken.kernel.triton.qsa.topk import _split_plan + + plan = _split_plan(columns, width) + assert plan is not None, f"{columns} columns must take the split path" + return plan.chunk def _policy_topk_blocks(logits, visible, width): @@ -348,29 +358,34 @@ def _policy_topk_blocks(logits, visible, width): return out -def _split_topk_case(n_blocks: int, width: int, bs: int, mode: str, seed: int): +def _split_topk_case(n_blocks: int, chunk: int, width: int, bs: int, mode: str, seed: int): if mode != "boundary": return _topk_case(n_blocks, bs, mode, seed) # width - 212 columns beat the tie, so the 212 remaining winners start 100 columns below - # a chunk boundary and run past it into a chunk that is all tie. + # a chunk boundary and run past it into a chunk that is all tie. The last row cuts at the + # middle split boundary, which is where a reduction level divides the candidates. logits = torch.zeros(bs, n_blocks, device="cuda") + splits = n_blocks // chunk for row in range(bs): - cut = SPLIT_CHUNK * (1 + row % (n_blocks // SPLIT_CHUNK - 1)) + cut = chunk * (splits // 2 if row == bs - 1 else 1 + row % (splits - 1)) logits[row, : width - 212] = 2.0 logits[row, cut - 100 :] = 1.0 return logits, torch.full((bs,), n_blocks, dtype=torch.int32, device="cuda") @requires_cuda -@pytest.mark.parametrize("n_blocks", [65536]) +@pytest.mark.parametrize("n_blocks,chunk", SPLIT_SHAPES) @pytest.mark.parametrize("bs", [4]) @pytest.mark.parametrize("mode", ["random", "boundary", "ragged", "dead"]) -def test_block_topk_split_path_matches_torch_topk(n_blocks: int, bs: int, mode: str): +def test_block_topk_split_path_matches_torch_topk(n_blocks: int, chunk: int, bs: int, mode: str): from freetoken.kernel.triton.qsa import qsa_block_topk, qsa_block_topk_scratch_width width = 512 assert qsa_block_topk_scratch_width(n_blocks, width) > 0, "case must take the split path" - logits, visible = _split_topk_case(n_blocks, width, bs, mode, seed=n_blocks + bs + len(mode)) + assert _split_chunk(n_blocks, width) == chunk + logits, visible = _split_topk_case( + n_blocks, chunk, width, bs, mode, seed=n_blocks + bs + len(mode) + ) blocks = torch.empty(bs, width, dtype=torch.int32, device=logits.device) qsa_block_topk(logits, visible, blocks) @@ -431,8 +446,10 @@ def test_block_topk_split_path_cost_tracks_live_blocks(): """A wide buffer with a short row must not pay for the splits past its visible tail.""" from freetoken.kernel.triton.qsa import qsa_block_topk, qsa_block_topk_scratch_width - # wide enough that the split work dwarfs the 20-launch floor even at boosted clocks - rows, columns, width = 1, 262144, 512 + # enough rows that the live splits fill the card, so their work outweighs the launch floor + rows, columns, width = 8, 262144, 512 + assert qsa_block_topk_scratch_width(columns, width) > 0, "case must take the split path" + chunk = _split_chunk(columns, width) device = torch.device("cuda") logits = torch.randn(rows, columns, device=device) visible = torch.full((rows,), columns, dtype=torch.int32, device=device) @@ -463,8 +480,8 @@ def replay_us(live: int) -> float: best = min(best, start.elapsed_time(stop) * 1000.0 / 20) return best - full, short = replay_us(columns), replay_us(SPLIT_CHUNK) - assert full > 2.0 * short, f"{columns} live {full:.1f}us vs {SPLIT_CHUNK} live {short:.1f}us" + full, short = replay_us(columns), replay_us(chunk) + assert full > 2.0 * short, f"{columns} live {full:.1f}us vs {chunk} live {short:.1f}us" @requires_cuda