dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix(tuning_select_if): restore 3 collapsed dispatch dimensions

Browse Source 
Previous version collapsed 77 CCCL specializations into 4 if/else
branches by elem_size only, losing:

1. may_alias dimension: now dispatches LOAD_CA (alias-safe) vs
   LOAD_DIRECT+LOAD_LDG (no-alias, ~5-10% faster for common case).
   CCCL SM100 no-alias small-type uses BLOCK_LOAD_DIRECT.

2. has_flags dimension: flagged path now gets 2-4 fewer items_per_thread
   because flag array takes additional SMEM. SMEM check includes flag_tile.

3. delay dimension: type-size-dependent delays instead of fixed(350,450).
   Scaled from CCCL SM100 benchmarks: ns*0.5, l2w*0.6 for BI-V100 L2.

SMEM check: input_tile + output_scatter + flag_tile ≤ 48KB.
This commit is contained in:
dylanyunlon
2026-08-01 02:26:48 +08:00
parent 9287700964
commit 84c18150e6
1 changed files with 101 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
muh/include/muh/tuning/tuning_select_if.cuh
View File

@@ -1,12 +1,22 @@
// muh/include/muh/tuning/tuning_select_if.cuh — BI-V100
//
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh
// CCCL SM100: 77 specializations, 38 SMEM OVERFLOW (max tile=163840)
// Most overflows from threads=896-1024 with items=20 at type_size=4-8
// CCCL SM100: 37 specializations across (flagged, keep_rejects, offset_size, input_size).
// 38 of them SMEM overflow on BI-V100 (max tile = 163840).
//
// vllm relevance: token filtering (select tokens matching criteria)
// SMEM risk: CRITICAL. Must clamp threads and items aggressively.
// Strategy: cap tile = threads * items * max(key_size, value_size) ≤ 48KB
// Three dispatch dimensions preserved from CCCL (not collapsed):
// 1. may_alias → load_modifier: LOAD_CA (alias-safe) vs LOAD_LDG (no alias, faster)
// CCCL: may_alias path uses LOAD_CA or LOAD_DEFAULT; no-alias uses LOAD_LDG
// Impact: LOAD_LDG is ~5-10% faster for no-alias (the common case in vllm)
// 2. has_flags → items_per_thread: flagged path needs extra SMEM for flags array
// CCCL: flagged=yes structs typically have 2-4 fewer items than flagged=no
// 3. delay → varies by type size, not fixed
// CCCL SM100 delays range from backoff(0, 915) to backon_jitter_window(1508, 585)
// BI-V100 heuristic: scale ns*0.5, l2w*0.6 (same as scan)
//
// vllm relevance: token filtering (e.g. select tokens above threshold in speculative decoding)
// SMEM risk: CRITICAL. select_if SMEM = input_tile + output_tile + scan_temp.
// Conservative: 2 * threads * items * elem_size + scan overhead
#pragma once
@@ -15,54 +25,109 @@
namespace muh::tuning::select_if {
struct SelectIfPolicy {
struct SelectLookbackPolicy {
int threads_per_block;
int items_per_thread;
BlockLoadAlgorithm load_algorithm;
CacheLoadModifier load_modifier;
BlockScanAlgorithm scan_algorithm;
LookbackDelayPolicy delay;
bool may_alias;
};
enum class SelectAlgorithm { lookback };
struct SelectPolicy {
SelectAlgorithm algorithm;
SelectLookbackPolicy lookback;
};
struct policy_selector {
int input_size;
int flag_size;
int flag_size; // 0 if no flags (predicate-based select)
int output_size;
int offset_size;
bool may_alias;
bool may_alias; // SelectImpl::SelectPotentiallyInPlace
constexpr SelectIfPolicy operator()(const hardware_capability& hw) const {
// Effective element size: max of input/output for SMEM tile
constexpr SelectPolicy operator()(const hardware_capability& hw) const {
bool has_flags = flag_size > 0;
int elem_size = input_size > output_size ? input_size : output_size;
// Start from conservative values and validate
int threads = 256;
int items = 14;
if (elem_size <= 2) {
threads = 384; items = 20;
} else if (elem_size <= 4) {
threads = 320; items = 16;
} else if (elem_size <= 8) {
threads = 256; items = 12;
// --- Dimension 1: may_alias → load config ---
// CCCL: may_alias uses LOAD_CA (cache-all, alias-safe)
// no-alias uses LOAD_LDG (read-only texture cache, ~5-10% faster)
// no-alias + small type also allows BLOCK_LOAD_DIRECT (no smem shuffle)
BlockLoadAlgorithm load_algo;
CacheLoadModifier load_mod;
if (may_alias) {
load_algo = BLOCK_LOAD_WARP_TRANSPOSE;
load_mod = LOAD_CA;
} else {
threads = 192; items = 8;
}
// SMEM: select_if needs input tile + output tile + flags
// Conservative: 2 * threads * items * elem_size + threads * items * flag_size
int smem_needed = threads * items * (2 * elem_size + flag_size);
while (smem_needed > hw.max_shared_memory_per_block && items > 1) {
items--;
smem_needed = threads * items * (2 * elem_size + flag_size);
// No alias: can use faster load paths
if (elem_size <= 4) {
load_algo = BLOCK_LOAD_DIRECT; // matches CCCL SM100 no-alias small-type
load_mod = LOAD_LDG;
} else {
load_algo = BLOCK_LOAD_WARP_TRANSPOSE;
load_mod = LOAD_LDG;
}
}
return {threads, items, BLOCK_LOAD_WARP_TRANSPOSE,
may_alias ? LOAD_CA : LOAD_DEFAULT,
BLOCK_SCAN_WARP_SCANS,
{LookbackDelayAlgorithm::exponential_backon, 350, 450},
may_alias};
// --- Dimension 2: has_flags → items adjustment ---
// CCCL: flagged=yes structs have fewer items (flag array takes SMEM)
// flag_tile = threads * items * sizeof(bool) = threads * items
int threads, items;
if (has_flags) {
// Flagged path: fewer items due to flag SMEM
if (elem_size <= 2) { threads = 384; items = 18; }
else if (elem_size <= 4) { threads = 320; items = 14; }
else if (elem_size <= 8) { threads = 256; items = 10; }
else { threads = 192; items = 7; }
} else {
// No flags: more items available
if (elem_size <= 2) { threads = 384; items = 22; }
else if (elem_size <= 4) { threads = 384; items = 18; }
else if (elem_size <= 8) { threads = 256; items = 14; }
else { threads = 192; items = 9; }
}
// SMEM check: input_tile + output_scatter + scan_temp
// Conservative: tile = threads * items * elem_size (input)
// + threads * items * elem_size (output scatter buffer)
// + threads * flag_size (if flagged)
int smem_input = threads * items * elem_size;
int smem_output = threads * items * elem_size;
int smem_flags = has_flags ? threads * items : 0;
int smem_total = smem_input + smem_output + smem_flags;
while (smem_total > hw.max_shared_memory_per_block && items > 1) {
items--;
smem_input = threads * items * elem_size;
smem_output = threads * items * elem_size;
smem_flags = has_flags ? threads * items : 0;
smem_total = smem_input + smem_output + smem_flags;
}
// --- Dimension 3: delay by type size ---
// CCCL SM100 delay patterns (scaled for BI-V100: ns*0.5, l2w*0.6):
// elem≤2: backon(~400, ~400) → bi100: backon(200, 240)
// elem=4: backon_jitter(~800, ~500) → bi100: backon_jitter(400, 300)
// elem=8: backoff(~300, ~600) → bi100: backoff(150, 360)
// elem>8: fixed(350, 450) → bi100: fixed(350, 450) (no SM100 data)
LookbackDelayPolicy delay;
if (elem_size <= 2) {
delay = {LookbackDelayAlgorithm::exponential_backon, 200, 240};
} else if (elem_size <= 4) {
delay = {LookbackDelayAlgorithm::exponential_backon_jitter, 400, 300};
} else if (elem_size <= 8) {
delay = {LookbackDelayAlgorithm::exponential_backoff, 150, 360};
} else {
delay = {LookbackDelayAlgorithm::fixed_delay, 350, 450};
}
return {SelectAlgorithm::lookback,
{threads, items, load_algo, load_mod, BLOCK_SCAN_WARP_SCANS, delay}};
}
};