[muh] three_way_partition 7%→13%, histogram 13%→21%: 继续从 CCCL 源码移植 SM80/SM90/SM100 tuning tables
tuning_three_way_partition.cuh: 58→99 行 - 移植 SM100 (5 entries) + SM90 (10) + SM80 (4) 共 19 条 - 按 (offset_size, input_size) 二维分派 - 三路划分 SMEM: 3 * tpb * ipt * input_size tuning_histogram.cuh: 48→76 行 - 移植 SM100 (2 entries) + SM90 (2) + default - privatized SMEM bins 保护 - 保留 CCCL benchmark 注释
This commit is contained in:
muh-bot
@@ -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};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user