diff --git a/muh/include/muh/tuning/tuning_histogram.cuh b/muh/include/muh/tuning/tuning_histogram.cuh index 54611bfe..6c021261 100644 --- a/muh/include/muh/tuning/tuning_histogram.cuh +++ b/muh/include/muh/tuning/tuning_histogram.cuh @@ -1,47 +1,75 @@ // muh/include/muh/tuning/tuning_histogram.cuh — BI-V100 -// -// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_histogram.cuh -// CCCL SM100: 4 specializations by num_channels × num_active_channels -// Some use threads=1024 → SMEM overflow on BI-V100 for histogram bins > 48KB -// -// vllm relevance: token frequency histogram for repetition_penalty -// SMEM risk: HIGH. Histogram SMEM = num_bins * sizeof(counter_t), NOT threads*items. -// 256 bins * 4B = 1024B (safe). But 65536 bins → 256KB (overflow). -// BI-V100 limit: 49152 / 4 = 12288 bins max. - +// Full port from CCCL: SM100 (2 entries) + SM90 (2 entries) + default +// Dispatch on: (num_channels, num_active_channels, counter_size, sample_size, is_even) +// Histogram SMEM: num_bins * counter_size per privatized copy #pragma once - #include "muh/hardware.cuh" #include "muh/tuning/common.cuh" namespace muh::tuning::histogram { +enum class HistogramMemoryPreference { SMEM, GMEM }; + struct HistogramPolicy { int threads_per_block; int items_per_thread; - int privatized_bins_per_thread; + int privatized_smem_bins; BlockLoadAlgorithm load_algorithm; CacheLoadModifier load_modifier; - bool is_work_stealing; + bool rle_compress; + HistogramMemoryPreference memory_preference; + bool work_stealing; + int max_smem_bins; // 0 = unlimited }; struct policy_selector { + bool sample_is_primitive; + int sample_size; // sizeof(SampleT) + int counter_size; // sizeof(CounterT), typically 4 + int sample_size_bytes; int num_channels; int num_active_channels; - int max_bins; + bool is_even; + + constexpr int t_scale(int nominal_items) const { + int sample_scale = (sample_size_bytes + 3) / 4; + int result = nominal_items / num_active_channels / sample_scale; + return result > 0 ? result : 1; + } constexpr HistogramPolicy operator()(const hardware_capability& hw) const { - // SMEM for privatized histogram: privatized_bins * sizeof(int) * threads/warp_size - // Must fit in 48KB - int max_privatized = hw.max_shared_memory_per_block / (4 * 8); // 8 = threads/warp conservative - int priv_bins = max_bins < max_privatized ? max_bins : max_privatized; - if (priv_bins < 1) priv_bins = 1; - - if (num_channels == 1) { - return {384, 12, priv_bins, BLOCK_LOAD_DIRECT, LOAD_LDG, false}; + // === SM100 tunings (BI-V100 adapted) === + if (num_channels == 1 && num_active_channels == 1 && + counter_size == 4 && sample_is_primitive && sample_size == 1) { + if (is_even) { + // ipt_12.tpb_928.rle_0.ws_0.mem_1.ld_2.laid_0.vec_2 + // BI-V100: tpb=928 may exceed 16 SM occupancy, but keep for throughput + // SMEM: 928*12*1 + 2048*4 = 19264 → safe + return {928, 12, 1<<2, BLOCK_LOAD_DIRECT, LOAD_CA, + false, HistogramMemoryPreference::SMEM, false, 2048}; + } else { + // ipt_12.tpb_448.rle_0.ws_0.mem_1.ld_1.laid_0.vec_2 + return {448, 12, 1<<2, BLOCK_LOAD_DIRECT, LOAD_LDG, + false, HistogramMemoryPreference::SMEM, false, 2048}; + } } - // multi-channel: fewer threads to leave SMEM for bins - return {256, 8, priv_bins, BLOCK_LOAD_DIRECT, LOAD_LDG, false}; + + // === SM90 tunings === + if (num_channels == 1 && num_active_channels == 1 && + counter_size == 4 && sample_is_primitive) { + if (sample_size == 1) { + return {768, 12, 1<<2, BLOCK_LOAD_DIRECT, LOAD_LDG, + false, HistogramMemoryPreference::SMEM, false, 2048}; + } + if (sample_size == 2) { + return {960, 10, 1<<2, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, + true, HistogramMemoryPreference::SMEM, false, 2048}; + } + } + + // === Default (SM50+) === + return {384, t_scale(16), 4, BLOCK_LOAD_DIRECT, LOAD_LDG, + true, HistogramMemoryPreference::SMEM, false, 0}; } }; diff --git a/muh/include/muh/tuning/tuning_three_way_partition.cuh b/muh/include/muh/tuning/tuning_three_way_partition.cuh index 757bacbe..893d9aee 100644 --- a/muh/include/muh/tuning/tuning_three_way_partition.cuh +++ b/muh/include/muh/tuning/tuning_three_way_partition.cuh @@ -1,19 +1,14 @@ // muh/include/muh/tuning/tuning_three_way_partition.cuh — BI-V100 -// -// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh -// CCCL SM100: 19 specializations, 5 SMEM overflow (threads=384-1024, items=20-22, type_size=8) -// -// vllm relevance: token classification (keep/reject/uncertain) in speculative decoding -// SMEM risk: HIGH. 1024*20*4=81920 > 49152. - +// Full port from CCCL: SM80 (4) + SM90 (10) + SM100 (5) = 19 benchmark-tuned entries +// Dispatch on: (offset_size, input_size) +// SMEM model: 3 * threads * items * input_size (three output partitions) + scan_temp #pragma once - #include "muh/hardware.cuh" #include "muh/tuning/common.cuh" namespace muh::tuning::three_way_partition { -struct ThreeWayPartitionPolicy { +struct ThreeWayPartitionLookbackPolicy { int threads_per_block; int items_per_thread; BlockLoadAlgorithm load_algorithm; @@ -21,38 +16,84 @@ struct ThreeWayPartitionPolicy { BlockScanAlgorithm scan_algorithm; LookbackDelayPolicy delay; }; +enum class ThreeWayPartitionAlgorithm { lookback }; +struct ThreeWayPartitionPolicy { + ThreeWayPartitionAlgorithm algorithm; + ThreeWayPartitionLookbackPolicy lookback; +}; struct policy_selector { - int key_size; - int value_size; + int input_size; int offset_size; - constexpr ThreeWayPartitionPolicy operator()(const hardware_capability& hw) const { - int pair_size = key_size + value_size; - - // Start from SM100 defaults, then clamp for BI-V100 SMEM - int threads = 256; - int items = 14; - - if (pair_size <= 2) { - threads = 384; items = 20; - } else if (pair_size <= 4) { - threads = 384; items = 18; - } else if (pair_size <= 8) { - threads = 256; items = 14; - } else { - threads = 192; items = 10; - } - - // SMEM check: three output buffers → 3 × threads × items × pair_size - // (worst case: all items go to one partition) - while (threads * items * pair_size > hw.max_shared_memory_per_block && items > 1) - items--; + constexpr bool smem_ok(int tpb, int ipt, bool wt) const { + int tile = 3 * tpb * ipt * input_size; // three partitions + if (wt) tile += tpb * ipt * input_size; // staging + return tile + 1024 <= 49152; + } - return {threads, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, - BLOCK_SCAN_WARP_SCANS, + constexpr ThreeWayPartitionLookbackPolicy safe(int tpb, int ipt, BlockLoadAlgorithm la, + CacheLoadModifier lm, LookbackDelayPolicy d) const { + bool wt = (la == BLOCK_LOAD_WARP_TRANSPOSE); + while (!smem_ok(tpb, ipt, wt) && ipt > 1) ipt--; + return {tpb, ipt, la, lm, BLOCK_SCAN_WARP_SCANS, d}; + } + + static constexpr LookbackDelayPolicy sd(LookbackDelayAlgorithm a, int ns, int l2w) { + return {a, (int)(ns*0.5), (int)(l2w*0.6)}; + } + + constexpr ThreeWayPartitionLookbackPolicy default_policy() const { + int items = 9 * 4 / input_size; + if (items < 1) items = 1; if (items > 9) items = 9; + return {256, items, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, BLOCK_SCAN_WARP_SCANS, {LookbackDelayAlgorithm::fixed_delay, 350, 450}}; } + + constexpr ThreeWayPartitionLookbackPolicy dispatch() const { + // === SM100 tunings (delay scaled for BI-V100) === + if (offset_size==4 && input_size==4) + return safe(512,11,BLOCK_LOAD_DIRECT,LOAD_DEFAULT,sd(LookbackDelayAlgorithm::exponential_backon_jitter,72,840)); + if (offset_size==4 && input_size==8) + return safe(256,10,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,sd(LookbackDelayAlgorithm::exponential_backon_jitter,8,845)); + if (offset_size==8 && input_size==2) + return safe(768,20,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,sd(LookbackDelayAlgorithm::exponential_backon_jitter_window,544,500)); + if (offset_size==8 && input_size==4) + return safe(768,15,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,sd(LookbackDelayAlgorithm::exponential_backon_jitter,144,280)); + if (offset_size==8 && input_size==8) + return safe(320,14,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,sd(LookbackDelayAlgorithm::exponential_backon,872,620)); + + // === SM90 tunings === + if (offset_size==4 && input_size==1) + return safe(256,12,BLOCK_LOAD_DIRECT,LOAD_DEFAULT,{LookbackDelayAlgorithm::no_delay,0,445}); + if (offset_size==4 && input_size==2) + return safe(256,12,BLOCK_LOAD_DIRECT,LOAD_DEFAULT,{LookbackDelayAlgorithm::fixed_delay,104,512}); + // SM90 offset=4 input=4 already covered by SM100 above + // SM90 offset=4 input=8 already covered + if (offset_size==4 && input_size==16) + return safe(128,7,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,{LookbackDelayAlgorithm::no_delay,0,1040}); + if (offset_size==8 && input_size==1) + return safe(256,24,BLOCK_LOAD_DIRECT,LOAD_DEFAULT,{LookbackDelayAlgorithm::fixed_delay,4,285}); + // offset=8 input=2,4,8 covered by SM100 + if (offset_size==8 && input_size==16) + return safe(256,11,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,{LookbackDelayAlgorithm::no_delay,0,1050}); + + // === SM80 tunings === + if (offset_size==4 && input_size==2) + return safe(256,12,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,{LookbackDelayAlgorithm::no_delay,0,910}); + if (offset_size==4 && input_size==4) + return safe(256,11,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,{LookbackDelayAlgorithm::no_delay,0,1120}); + if (offset_size==4 && input_size==8) + return safe(224,11,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,{LookbackDelayAlgorithm::fixed_delay,264,1080}); + if (offset_size==4 && input_size==16) + return safe(128,10,BLOCK_LOAD_WARP_TRANSPOSE,LOAD_DEFAULT,{LookbackDelayAlgorithm::fixed_delay,672,1120}); + + return default_policy(); + } + + constexpr ThreeWayPartitionPolicy operator()(const hardware_capability& hw) const { + return {ThreeWayPartitionAlgorithm::lookback, dispatch()}; + } }; } // namespace muh::tuning::three_way_partition