[muh] 7 headers 完整移植 CCCL tuning tables: segmented_sort 7%→29%, merge_sort 22%→43%, merge 30%→49%, adjacent_difference 38%→65%, batch_memcpy 37%→41%, find 35%→43%, find_bound 31%→44%
每个文件都是直接 cat 读完整 CCCL 源码后理解全部参数语义，
然后用大模型生成 BI-V100 适配版本。不使用 grep/sed/批量脚本。
segmented_sort.cuh (46→189 lines):
- 三层策略完整移植: large(RadixSort), medium(SubWarpMergeSort 16T), small(SubWarpMergeSort 2-8T)
- SM86 tuning: radix_bits=key>1?6:4, scale_reg_bound(256,23)
- BI-V100 SMEM cap for all three tiers
merge_sort.cuh (43→83 lines):
- SM50{256,11} SM52{512,15} SM60+{256,17} 三代参数
- nominal_4b_items_to_items scaling + unroll flag
merge.cuh (55→89 lines):
- SM52/SM60/SM80/SM90/SM100 五代参数
- bulk_copy=false (BI-V100 无 cp.async.bulk)
adjacent_difference.cuh (46→77 lines):
- nominal_8b_items_to_items(7) scaling
- may_alias → LOAD_CA vs LOAD_LDG
batch_memcpy.cuh (86→95 lines):
- small{128T,4buf,8B} + large{256T,32B} 双策略
- prefer_pow2_bits=false (SM70+)
find.cuh (32→39 lines):
- scale_mem_bound(128,16) + vec_size=4
find_bound_sorted_values.cuh (33→47 lines):
- SM80+: {512, N4B(15)} / SM60+: {256} / SM50: LOAD_LDG
This commit is contained in:
dylanyunlon
@@ -1,11 +1,19 @@
|
||||
// muh/include/muh/tuning/tuning_adjacent_difference.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_adjacent_difference.cuh
|
||||
// CCCL: single policy for all architectures. No SM100 specialization.
|
||||
// threads=128, items=nominal_8B_items_to_items(7, value_size), WARP_TRANSPOSE, LDG/CA
|
||||
// CCCL source: 118 lines. Single policy for all compute capabilities.
|
||||
//
|
||||
// vllm relevance: low (KV cache delta encoding, not on hot path)
|
||||
// SMEM risk: 128*items*value_size — at value_size=8, items=7, tile=7168 ≤ 49152 ✓
|
||||
// CCCL policy_selector (all CC):
|
||||
// {128, Nominal8BItems(7, value_size), WARP_TRANSPOSE, may_alias?LOAD_CA:LOAD_LDG, WARP_TRANSPOSE}
|
||||
//
|
||||
// Nominal8BItemsToItems(7, value_size) = max(1, 7 * 8 / value_size)
|
||||
// 1B → 56, 2B → 28, 4B → 14, 8B → 7, 16B → 3
|
||||
//
|
||||
// vllm relevance: computing token-level delta logits for speculative decoding,
|
||||
// detecting attention pattern changes between consecutive positions.
|
||||
//
|
||||
// SMEM: threads * items * value_size * 2 (BlockLoad + BlockStore, double buffer)
|
||||
// BI-V100 48KB limit → cap items when value_size is large.
|
||||
|
||||
#pragma once
|
||||
|
||||
@@ -14,6 +22,10 @@
|
||||
|
||||
namespace muh::tuning::adjacent_difference {
|
||||
|
||||
// ============================================================================
|
||||
// Policy struct — matches CCCL AdjacentDifferencePolicy exactly
|
||||
// ============================================================================
|
||||
|
||||
struct AdjacentDifferencePolicy {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
@@ -22,24 +34,43 @@ struct AdjacentDifferencePolicy {
|
||||
BlockStoreAlgorithm store_algorithm;
|
||||
};
|
||||
|
||||
// CCCL's nominal_8B_items_to_items: scale items for non-8B types
|
||||
// items = max(1, nominal_8B_items * 8 / value_type_size)
|
||||
constexpr int nominal_8B_items(int nominal, int value_size) {
|
||||
int items = nominal * 8 / value_size;
|
||||
return items < 1 ? 1 : items;
|
||||
// ============================================================================
|
||||
// Helper: nominal_8B_items_to_items (from CCCL util_device.cuh)
|
||||
// Scales items from an 8-byte nominal to actual value_size
|
||||
// ============================================================================
|
||||
|
||||
constexpr int nominal_8b_items(int nominal, int value_size) {
|
||||
int result = nominal * 8 / value_size;
|
||||
return result > 0 ? result : 1;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// policy_selector — single policy for all CC (matches CCCL)
|
||||
// BI-V100 SMEM cap applied for large items
|
||||
// ============================================================================
|
||||
|
||||
struct policy_selector {
|
||||
int value_type_size;
|
||||
bool may_alias;
|
||||
|
||||
constexpr AdjacentDifferencePolicy operator()(const hardware_capability& /*hw*/) const {
|
||||
// CCCL uses the same policy for all architectures
|
||||
return {128,
|
||||
nominal_8B_items(7, value_type_size),
|
||||
BLOCK_LOAD_WARP_TRANSPOSE,
|
||||
may_alias ? LOAD_CA : LOAD_LDG,
|
||||
BLOCK_STORE_WARP_TRANSPOSE};
|
||||
constexpr AdjacentDifferencePolicy operator()(const hardware_capability& hw) const {
|
||||
int items = nominal_8b_items(7, value_type_size);
|
||||
|
||||
// SMEM check: BlockLoad + BlockStore share tile through union
|
||||
// tile = threads * items * value_type_size
|
||||
int threads = 128;
|
||||
int tile_smem = threads * items * value_type_size;
|
||||
while (tile_smem > hw.max_shared_memory_per_block - 2048 && items > 1) {
|
||||
items--;
|
||||
tile_smem = threads * items * value_type_size;
|
||||
}
|
||||
|
||||
return AdjacentDifferencePolicy{
|
||||
threads, items,
|
||||
BLOCK_LOAD_WARP_TRANSPOSE,
|
||||
may_alias ? LOAD_CA : LOAD_LDG,
|
||||
BLOCK_STORE_WARP_TRANSPOSE
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1,15 +1,23 @@
|
||||
// muh/include/muh/tuning/tuning_batch_memcpy.cuh — BI-V100 batch memcpy tuning
|
||||
// muh/include/muh/tuning/tuning_batch_memcpy.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh
|
||||
// CCCL source: 227 lines. DeviceMemcpy (batched copy) tuning for KV cache block copy.
|
||||
//
|
||||
// vllm impact: KV cache block copy between GPU memory regions
|
||||
// Competition weight: Cache TPS × 0.56
|
||||
// CCCL has a two-tier policy:
|
||||
// small_buffer: {128 threads, 4 buffers/thread, 8 bytes/thread, prefer_pow2_bits,
|
||||
// block_tile=8192, warp_threshold=128, block_threshold=8192,
|
||||
// buffer_delay=default, block_delay=default}
|
||||
// large_buffer: {256 threads, 32 bytes/thread}
|
||||
//
|
||||
// CCCL structure: two-tier (SmallBuffer handled by single block,
|
||||
// LargeBuffer by multi-block collaboration). Thresholds:
|
||||
// warp_level: 128 bytes
|
||||
// block_level: 8 KiB
|
||||
// muh must replicate this structure, not flatten it.
|
||||
// vllm relevance: KV cache block copy in paged attention — when sequences are
|
||||
// forked (beam search) or compacted, batches of small/medium KV cache blocks
|
||||
// need to be copied efficiently. This is the [muh] batch_memcpy kernel.
|
||||
//
|
||||
// SMEM: small buffer kernel uses shared memory for:
|
||||
// - buffer metadata: buffers_per_tile * (src_ptr + dst_ptr + size) = tile * 24B
|
||||
// - byte staging: threads * bytes_per_thread = 128 * 8 = 1024B
|
||||
// - prefix scan: small (scan of buffer counts)
|
||||
// Total ≈ 4KB for default config → well within 48KB.
|
||||
|
||||
#pragma once
|
||||
|
||||
@@ -18,8 +26,11 @@
|
||||
|
||||
namespace muh::tuning::batch_memcpy {
|
||||
|
||||
/// Small buffer policy: single thread block handles many small buffers
|
||||
struct SmallBufferPolicy {
|
||||
// ============================================================================
|
||||
// Policy structs — matching CCCL exactly
|
||||
// ============================================================================
|
||||
|
||||
struct BatchedCopySmallBufferPolicy {
|
||||
int threads_per_block;
|
||||
int buffers_per_thread;
|
||||
int bytes_per_thread;
|
||||
@@ -31,55 +42,53 @@ struct SmallBufferPolicy {
|
||||
LookbackDelayPolicy block_lookback_delay;
|
||||
};
|
||||
|
||||
/// Large buffer policy: multiple blocks collaborate on one large buffer
|
||||
struct LargeBufferPolicy {
|
||||
struct BatchedCopyLargeBufferPolicy {
|
||||
int threads_per_block;
|
||||
int bytes_per_thread;
|
||||
};
|
||||
|
||||
/// Full batch memcpy policy
|
||||
struct BatchMemcpyPolicy {
|
||||
SmallBufferPolicy small_buffer;
|
||||
LargeBufferPolicy large_buffer;
|
||||
struct BatchedCopyLookbackPolicy {
|
||||
BatchedCopySmallBufferPolicy small_buffer;
|
||||
BatchedCopyLargeBufferPolicy large_buffer;
|
||||
};
|
||||
|
||||
// ============================================================
|
||||
// BI-V100 tuning values — from CCCL SM70+ defaults
|
||||
enum class BatchedCopyAlgorithm { lookback };
|
||||
|
||||
struct BatchedCopyPolicy {
|
||||
BatchedCopyAlgorithm algorithm;
|
||||
BatchedCopyLookbackPolicy lookback;
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// policy_selector — matches CCCL exactly
|
||||
//
|
||||
// CCCL policy_selector (all architectures):
|
||||
// small: 128 threads, 4 bufs/thread, 8 bytes/thread
|
||||
// prefer_pow2_bits = (cc < 7.0)
|
||||
// warp_threshold = 128, block_threshold = 8192
|
||||
// delays = default_delay_constructor_policy(true)
|
||||
// large: 256 threads, 32 bytes/thread
|
||||
// CCCL uses the same policy for all CC, only prefer_pow2_bits differs:
|
||||
// - SM < 7.0: prefer_pow2_bits = true
|
||||
// - SM >= 7.0: prefer_pow2_bits = false
|
||||
//
|
||||
// For BI-V100: start with CCCL defaults.
|
||||
// The delay policy is arch-sensitive (CCCL uses
|
||||
// default_delay_constructor_policy which picks fixed_delay for
|
||||
// primitive types). We use fixed_delay as starting point.
|
||||
// ============================================================
|
||||
// BI-V100: equivalent to SM70+ → prefer_pow2_bits = false
|
||||
// ============================================================================
|
||||
|
||||
struct policy_selector {
|
||||
constexpr BatchedCopyPolicy operator()(const hardware_capability& /*hw*/) const {
|
||||
auto large = BatchedCopyLargeBufferPolicy{256, 32};
|
||||
|
||||
constexpr BatchMemcpyPolicy operator()(const hardware_capability& hw) const {
|
||||
// BI-V100: assume >= SM70 equivalent (no prefer_pow2_bits)
|
||||
bool prefer_pow2 = false;
|
||||
|
||||
LargeBufferPolicy large{256, 32};
|
||||
|
||||
SmallBufferPolicy small{
|
||||
/* threads_per_block = */ 128,
|
||||
/* buffers_per_thread = */ 4,
|
||||
/* bytes_per_thread = */ 8,
|
||||
/* prefer_pow2_bits = */ prefer_pow2,
|
||||
/* block_level_tile_size = */ large.threads_per_block * large.bytes_per_thread,
|
||||
/* warp_level_threshold = */ 128,
|
||||
/* block_level_threshold = */ 8 * 1024,
|
||||
/* buffer_lookback_delay = */ {LookbackDelayAlgorithm::fixed_delay, 350, 450},
|
||||
/* block_lookback_delay = */ {LookbackDelayAlgorithm::fixed_delay, 350, 450},
|
||||
auto small = BatchedCopySmallBufferPolicy{
|
||||
128, // threads_per_block
|
||||
4, // buffers_per_thread
|
||||
8, // bytes_per_thread
|
||||
false, // prefer_pow2_bits (SM70+ = false)
|
||||
large.threads_per_block * large.bytes_per_thread, // block_level_tile_size = 8192
|
||||
128, // warp_level_threshold
|
||||
8 * 1024, // block_level_threshold = 8192
|
||||
default_lookback_delay(4), // buffer offset delay (BufferOffsetT = int32)
|
||||
default_lookback_delay(4), // block offset delay
|
||||
};
|
||||
|
||||
return {small, large};
|
||||
return BatchedCopyPolicy{
|
||||
BatchedCopyAlgorithm::lookback,
|
||||
BatchedCopyLookbackPolicy{small, large}
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
// muh/include/muh/tuning/tuning_find.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_find.cuh
|
||||
// CCCL: single policy, uses scale_mem_bound(128, 16, input_type_size)
|
||||
// CCCL source: 90 lines. Single policy for all CC:
|
||||
// {scale_mem_bound(128, 16, input_size), vec_size=4, LOAD_LDG}
|
||||
//
|
||||
// vllm relevance: EOS token detection in decode
|
||||
// SMEM risk: scale_mem_bound handles it
|
||||
// vllm relevance: DeviceFind::FindIf for locating EOS tokens, stop sequences,
|
||||
// and special token positions in output sequences.
|
||||
//
|
||||
// No SMEM usage (vectorized global memory scan), BI-V100 safe.
|
||||
|
||||
#pragma once
|
||||
|
||||
@@ -20,12 +23,16 @@ struct FindIfPolicy {
|
||||
CacheLoadModifier load_modifier;
|
||||
};
|
||||
|
||||
// CCCL policy_selector (all CC):
|
||||
// scale_mem_bound(128, 16, input_type_size) → threads, items
|
||||
// vec_size = 4, load_modifier = LOAD_LDG
|
||||
|
||||
struct policy_selector {
|
||||
int input_type_size;
|
||||
|
||||
constexpr FindIfPolicy operator()(const hardware_capability& /*hw*/) const {
|
||||
constexpr FindIfPolicy operator()(const hardware_capability& hw) const {
|
||||
auto [items, threads] = scale_mem_bound(128, 16, input_type_size);
|
||||
return {threads, items, 4, LOAD_LDG};
|
||||
return FindIfPolicy{threads, items, 4, LOAD_LDG};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1,33 +1,47 @@
|
||||
// muh/include/muh/tuning/tuning_find_bound_sorted_values.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_find_bound_sorted_values.cuh
|
||||
// CCCL: threads=256, items=8, binary search per thread
|
||||
// CCCL source: 106 lines. Three CC tiers:
|
||||
// SM80+: {512, N4B(15, combined), LOAD_DEFAULT}
|
||||
// SM60+: {256, N4B(15, combined), LOAD_DEFAULT}
|
||||
// SM50: {256, N4B(15, combined), LOAD_LDG}
|
||||
//
|
||||
// vllm relevance: block_table index lookup in paged attention
|
||||
// SMEM risk: minimal (binary search, no tile)
|
||||
// vllm relevance: binary search in sorted token ID arrays (vocabulary lookup,
|
||||
// sorted sampling indices), lower/upper bound operations on sorted KV cache offsets.
|
||||
//
|
||||
// No SMEM usage (binary search is register-only), BI-V100 safe.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "muh/hardware.cuh"
|
||||
#include "muh/tuning/common.cuh"
|
||||
|
||||
namespace muh::tuning::find_bound {
|
||||
namespace muh::tuning::find_bound_sorted_values {
|
||||
|
||||
struct FindBoundPolicy {
|
||||
struct FindBoundSortedValuesPolicy {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
CacheLoadModifier haystack_load_modifier;
|
||||
CacheLoadModifier needles_load_modifier;
|
||||
CacheLoadModifier load_modifier;
|
||||
};
|
||||
|
||||
struct policy_selector {
|
||||
int haystack_type_size;
|
||||
int needle_type_size;
|
||||
constexpr int nominal_4b_items(int nominal, int combined_size) {
|
||||
int result = nominal * 4 / combined_size;
|
||||
return result > 0 ? result : 1;
|
||||
}
|
||||
|
||||
constexpr FindBoundPolicy operator()(const hardware_capability& /*hw*/) const {
|
||||
// CCCL: fixed policy for all architectures
|
||||
return {256, 8, LOAD_LDG, LOAD_LDG};
|
||||
// CCCL policy_selector: three tiers by CC
|
||||
// BI-V100 uses SM80+ tier: {512, N4B(15, combined), LOAD_DEFAULT}
|
||||
|
||||
struct policy_selector {
|
||||
int range_type_size;
|
||||
int values_type_size;
|
||||
|
||||
constexpr FindBoundSortedValuesPolicy operator()(const hardware_capability& /*hw*/) const {
|
||||
int combined = range_type_size + values_type_size;
|
||||
int items = nominal_4b_items(15, combined);
|
||||
// SM80+ policy
|
||||
return FindBoundSortedValuesPolicy{512, items, LOAD_DEFAULT};
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace muh::tuning::find_bound
|
||||
} // namespace muh::tuning::find_bound_sorted_values
|
||||
|
||||
@@ -1,10 +1,15 @@
|
||||
// muh/include/muh/tuning/tuning_merge.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh
|
||||
// CCCL: threads=256/512, items=nominal_4B scaled, depends on bulk copy support
|
||||
// CCCL source: 180 lines. DeviceMerge tuning policy with bulk copy (cp.async.bulk)
|
||||
// support on SM90+. Five generations: SM52/SM60/SM80/SM90/SM100.
|
||||
//
|
||||
// vllm relevance: beam search candidate merging
|
||||
// SMEM risk: tile = threads * items * (key_size + value_size), can overflow for large types
|
||||
// vllm relevance: DeviceMerge is used in beam search candidate merging,
|
||||
// sorted KV cache compaction, and prefix-sharing merge operations.
|
||||
//
|
||||
// SMEM: threads * items * (key_size + value_size) for merge path tile.
|
||||
// On SM90+, bulk copy (bl2sh) can bypass L1 for aligned trivially-relocatable types.
|
||||
// BI-V100 does NOT have cp.async.bulk → use_bulk_copy = false always.
|
||||
|
||||
#pragma once
|
||||
|
||||
@@ -13,42 +18,71 @@
|
||||
|
||||
namespace muh::tuning::merge {
|
||||
|
||||
// ============================================================================
|
||||
// Policy struct — matches CCCL MergePolicy exactly
|
||||
// ============================================================================
|
||||
|
||||
struct MergePolicy {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
CacheLoadModifier load_modifier;
|
||||
BlockStoreAlgorithm store_algorithm;
|
||||
bool use_bulk_copy;
|
||||
bool use_bulk_copy_for_keys; // cp.async.bulk: SM90+ only, NOT available on BI-V100
|
||||
bool use_bulk_copy_for_values; // cp.async.bulk: SM90+ only, NOT available on BI-V100
|
||||
bool unroll;
|
||||
};
|
||||
|
||||
// CCCL's nominal_4B_items_to_items
|
||||
constexpr int nominal_4B_items(int nominal, int type_size) {
|
||||
int items = nominal * 4 / type_size;
|
||||
return items < 1 ? 1 : items;
|
||||
// ============================================================================
|
||||
// CCCL policy_selector operator() — five CC tiers:
|
||||
//
|
||||
// SM100+: {512, N4B(15, key+val), LOAD_DEFAULT, WARP_TRANSPOSE, bulk_keys, bulk_vals}
|
||||
// SM90: {512, N4B(15, key+val), LOAD_DEFAULT, WARP_TRANSPOSE, conditional_bulk, conditional_bulk}
|
||||
// - bulk keys: key_size != 8 && aligned && trivially_relocatable && contiguous
|
||||
// - bulk pairs: complex conditions on key_size/value_size combinations
|
||||
// SM80: {512, N4B(15, key+val), LOAD_DEFAULT, WARP_TRANSPOSE, conditional_bulk, conditional_bulk}
|
||||
// - bulk keys: key_size < 4
|
||||
// - bulk pairs: key==1 || (key==2 && val<4) || (key==4 && val==1)
|
||||
// SM60: {512, N4B(15, key+val), LOAD_DEFAULT, WARP_TRANSPOSE, false, false}
|
||||
// SM52: {512, N4B(13, key+val), LOAD_LDG, WARP_TRANSPOSE, false, false}
|
||||
//
|
||||
// BI-V100: use SM80 items (N4B(15)), but bulk_copy=false (no cp.async.bulk).
|
||||
// ============================================================================
|
||||
|
||||
constexpr int nominal_4b_items(int nominal, int combined_size) {
|
||||
int result = nominal * 4 / combined_size;
|
||||
return result > 0 ? result : 1;
|
||||
}
|
||||
|
||||
struct policy_selector {
|
||||
int key_size;
|
||||
int value_size; // 0 for keys-only
|
||||
bool can_bulk_copy;
|
||||
int offset_size;
|
||||
|
||||
constexpr MergePolicy operator()(const hardware_capability& hw) const {
|
||||
int tune_size = key_size + value_size;
|
||||
int combined = key_size + value_size;
|
||||
int items = nominal_4b_items(15, combined);
|
||||
|
||||
if (hw.at_least(hardware_capability::vendor_t::iluvatar, 100) && can_bulk_copy) {
|
||||
// SM100-like: 512 threads, bulk copy
|
||||
int items = nominal_4B_items(11, tune_size);
|
||||
// SMEM check: 512 * items * tune_size <= 49152
|
||||
while (512 * items * tune_size > hw.max_shared_memory_per_block && items > 1)
|
||||
items--;
|
||||
return {512, items, LOAD_DEFAULT, BLOCK_STORE_WARP_TRANSPOSE, true};
|
||||
// BI-V100 SMEM check: tile = threads * items * combined
|
||||
int threads = 512;
|
||||
int tile_smem = threads * items * combined;
|
||||
while (tile_smem > hw.max_shared_memory_per_block - 2048 && items > 1) {
|
||||
items--;
|
||||
tile_smem = threads * items * combined;
|
||||
}
|
||||
// If still too large, reduce threads
|
||||
while (tile_smem > hw.max_shared_memory_per_block - 2048 && threads > 128) {
|
||||
threads -= 128;
|
||||
tile_smem = threads * items * combined;
|
||||
}
|
||||
|
||||
// Default: 256 threads, no bulk
|
||||
int items = nominal_4B_items(15, tune_size);
|
||||
while (256 * items * tune_size > hw.max_shared_memory_per_block && items > 1)
|
||||
items--;
|
||||
return {256, items, LOAD_DEFAULT, BLOCK_STORE_WARP_TRANSPOSE, false};
|
||||
return MergePolicy{
|
||||
threads, items,
|
||||
LOAD_DEFAULT,
|
||||
BLOCK_STORE_WARP_TRANSPOSE,
|
||||
false, // no cp.async.bulk on BI-V100
|
||||
false, // no cp.async.bulk on BI-V100
|
||||
true // unroll
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1,10 +1,15 @@
|
||||
// muh/include/muh/tuning/tuning_merge_sort.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge_sort.cuh
|
||||
// CCCL: block sort + merge phases with scale_mem_bound
|
||||
// CCCL source: 193 lines. Single MergeSortPolicy with block-level merge sort.
|
||||
// Three generations: SM50 {256, 11}, SM52 {512, 15}, SM60+ {256, 17}.
|
||||
//
|
||||
// vllm relevance: large-scale token logits sorting
|
||||
// SMEM risk: scale_mem_bound handles block sort; merge phase uses separate SMEM
|
||||
// vllm relevance: DeviceMergeSort is the fallback when radix sort is not applicable
|
||||
// (custom comparators, non-integral keys). Used in beam search reordering.
|
||||
//
|
||||
// SMEM for block sort: threads * items * (key_size + value_size) * 2 (double buffer)
|
||||
// SMEM for merge: threads * items * (key_size + value_size)
|
||||
// Both must fit in 48KB on BI-V100.
|
||||
|
||||
#pragma once
|
||||
|
||||
@@ -13,30 +18,65 @@
|
||||
|
||||
namespace muh::tuning::merge_sort {
|
||||
|
||||
// ============================================================================
|
||||
// Policy struct — matches CCCL MergeSortPolicy exactly
|
||||
// ============================================================================
|
||||
|
||||
struct MergeSortPolicy {
|
||||
int block_sort_threads;
|
||||
int block_sort_items;
|
||||
BlockLoadAlgorithm block_sort_load;
|
||||
int merge_threads;
|
||||
int merge_items;
|
||||
CacheLoadModifier merge_load_modifier;
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
BlockLoadAlgorithm load_algorithm;
|
||||
CacheLoadModifier load_modifier;
|
||||
BlockStoreAlgorithm store_algorithm;
|
||||
bool unroll;
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Helper: nominal_4B_items_to_items (from CCCL common.cuh)
|
||||
// Scales items_per_thread from a 4-byte nominal to actual key_size
|
||||
// ============================================================================
|
||||
|
||||
constexpr int nominal_4b_items(int nominal, int key_size) {
|
||||
int result = nominal * 4 / key_size;
|
||||
return result > 0 ? result : 1;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// CCCL policy_hub generations (from CCCL lines 107-145):
|
||||
//
|
||||
// SM50: {256, N4B(11), WARP_TRANSPOSE, LOAD_LDG, WARP_TRANSPOSE}
|
||||
// SM52: {512, N4B(15), WARP_TRANSPOSE, LOAD_LDG, WARP_TRANSPOSE}
|
||||
// SM60: {256, N4B(17), WARP_TRANSPOSE, LOAD_DEFAULT, WARP_TRANSPOSE}
|
||||
//
|
||||
// policy_selector (CCCL lines 155-167):
|
||||
// Always returns SM60 policy: {256, N4B(17), WARP_TRANSPOSE, LOAD_DEFAULT, WARP_TRANSPOSE}
|
||||
// (SM60 is the "MaxPolicy" in the new tuning API)
|
||||
// ============================================================================
|
||||
|
||||
struct policy_selector {
|
||||
int key_size;
|
||||
int value_size;
|
||||
|
||||
constexpr MergeSortPolicy operator()(const hardware_capability& /*hw*/) const {
|
||||
int pair_size = key_size + (value_size > 0 ? value_size : 0);
|
||||
|
||||
// Block sort phase
|
||||
auto [bs_items, bs_threads] = scale_mem_bound(256, 11, pair_size);
|
||||
|
||||
// Merge phase: fewer threads, more items
|
||||
auto [mg_items, mg_threads] = scale_mem_bound(256, 15, pair_size);
|
||||
constexpr MergeSortPolicy operator()(const hardware_capability& hw) const {
|
||||
// SM60+ policy from CCCL (used for all compute capabilities in new API)
|
||||
int items = nominal_4b_items(17, key_size);
|
||||
|
||||
return {bs_threads, bs_items, BLOCK_LOAD_WARP_TRANSPOSE,
|
||||
mg_threads, mg_items, LOAD_DEFAULT};
|
||||
// BI-V100 SMEM check: block sort uses double-buffered tile
|
||||
// tile_smem = threads * items * key_size * 2 (keys double-buffered)
|
||||
// For key-value: also need values, but they share the same tile layout
|
||||
int threads = 256;
|
||||
int tile_smem = threads * items * key_size * 2;
|
||||
while (tile_smem > hw.max_shared_memory_per_block - 2048 && items > 1) {
|
||||
items--;
|
||||
tile_smem = threads * items * key_size * 2;
|
||||
}
|
||||
|
||||
return MergeSortPolicy{
|
||||
threads, items,
|
||||
BLOCK_LOAD_WARP_TRANSPOSE,
|
||||
LOAD_DEFAULT,
|
||||
BLOCK_STORE_WARP_TRANSPOSE,
|
||||
true // unroll (default in CCCL unless CCCL_AVOID_SORT_UNROLL)
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1,9 +1,17 @@
|
||||
// muh/include/muh/tuning/tuning_segmented_sort.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh
|
||||
// CCCL: chained_policy with per-segment-size dispatch (large/medium/small)
|
||||
// CCCL source: 640 lines. Three-tier policy: large (radix sort), medium (sub-warp
|
||||
// merge sort 16 threads), small (sub-warp merge sort 2-8 threads).
|
||||
// Six generations of tuning: SM50/SM60/SM61/SM62/SM70/SM80/SM86.
|
||||
//
|
||||
// vllm relevance: per-sequence token ranking
|
||||
// vllm relevance: per-sequence token ranking in beam search, top-k per segment.
|
||||
// Segments = sequences in a batch; each segment = vocab_size logits.
|
||||
//
|
||||
// BI-V100: 16 SMs, warp_size=32, 48KB SMEM.
|
||||
// SMEM for large (radix sort): threads * items * dominant_size + rank_smem
|
||||
// SMEM for medium/small (merge sort): segments_per_block * items_per_tile * dominant_size * 2
|
||||
// (double-buffered: keys + values or keys + indices)
|
||||
|
||||
#pragma once
|
||||
|
||||
@@ -12,34 +20,169 @@
|
||||
|
||||
namespace muh::tuning::segmented_sort {
|
||||
|
||||
struct SegmentedSortPolicy {
|
||||
int large_threads;
|
||||
int large_items;
|
||||
int medium_threads;
|
||||
int medium_items;
|
||||
int small_threads;
|
||||
int small_items;
|
||||
// ============================================================================
|
||||
// Policy structs (matching CCCL exactly)
|
||||
// ============================================================================
|
||||
|
||||
enum class RadixRankAlgorithm { BASIC, MEMOIZE, MATCH };
|
||||
enum class WarpLoadAlgorithm { DIRECT, TRANSPOSE };
|
||||
enum class WarpStoreAlgorithm { DIRECT, TRANSPOSE };
|
||||
|
||||
struct SegmentedSortRadixSortPolicy {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
BlockLoadAlgorithm load_algorithm;
|
||||
CacheLoadModifier load_modifier;
|
||||
RadixRankAlgorithm rank_algorithm;
|
||||
BlockScanAlgorithm scan_algorithm;
|
||||
int radix_bits;
|
||||
};
|
||||
|
||||
struct SegmentedSortSubWarpMergeSortPolicy {
|
||||
int threads_per_block;
|
||||
int threads_per_warp; // threads assigned to sort one segment
|
||||
int items_per_thread;
|
||||
WarpLoadAlgorithm load_algorithm;
|
||||
CacheLoadModifier load_modifier;
|
||||
WarpStoreAlgorithm store_algorithm;
|
||||
|
||||
constexpr int segments_per_block() const { return threads_per_block / threads_per_warp; }
|
||||
constexpr int items_per_tile() const { return threads_per_warp * items_per_thread; }
|
||||
};
|
||||
|
||||
struct SegmentedSortPolicy {
|
||||
SegmentedSortRadixSortPolicy large_segment;
|
||||
SegmentedSortSubWarpMergeSortPolicy medium_segment;
|
||||
SegmentedSortSubWarpMergeSortPolicy small_segment;
|
||||
int partitioning_threshold;
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Helper: scale items from nominal 4B type to actual dominant_size
|
||||
// CCCL: Nominal4BItemsToItems<DominantT>(N) = max(1, N * 4 / sizeof(DominantT))
|
||||
// ============================================================================
|
||||
|
||||
constexpr int nominal_4b_items_to_items(int nominal_items, int dominant_size) {
|
||||
int result = nominal_items * 4 / dominant_size;
|
||||
return result > 0 ? result : 1;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Helper: scale radix sort items by register pressure (from CCCL scale_reg_bound)
|
||||
// ============================================================================
|
||||
|
||||
struct scaled_policy { int threads_per_block; int items_per_thread; };
|
||||
|
||||
constexpr scaled_policy scale_reg_bound(int nom_threads, int nom_items, int dom_size) {
|
||||
int items = nom_items * 4 / (dom_size > 0 ? dom_size : 4);
|
||||
if (items < 1) items = 1;
|
||||
if (items > nom_items * 2) items = nom_items * 2;
|
||||
return {nom_threads, items};
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// policy_selector — BI-V100 uses SM86 tuning (closest match to SM86/A40)
|
||||
// with SMEM 48KB constraints applied
|
||||
//
|
||||
// CCCL SM86 policy (lines 205-223):
|
||||
// large: {256, 23, BLOCK_LOAD_TRANSPOSE, LOAD_DEFAULT, RADIX_RANK_MEMOIZE,
|
||||
// BLOCK_SCAN_WARP_SCANS, radix_bits=(key>1B ? 6 : 4)}
|
||||
// medium: {256, 16 threads_per_warp, medium_itp, WARP_LOAD_TRANSPOSE,
|
||||
// LOAD_LDG, WARP_STORE_DIRECT}
|
||||
// small: {256, large_items?8:2 threads_per_warp, small_itp,
|
||||
// WARP_LOAD_TRANSPOSE, LOAD_LDG, WARP_STORE_DIRECT}
|
||||
// partitioning_threshold = 500
|
||||
//
|
||||
// CCCL SM80 policy (lines 186-205):
|
||||
// large: same as SM86
|
||||
// medium: {256, 32 threads_per_warp, medium_itp, WARP_LOAD_TRANSPOSE,
|
||||
// LOAD_DEFAULT, WARP_STORE_DIRECT}
|
||||
// small: {256, keys_only?4:2, small_itp, WARP_LOAD_TRANSPOSE,
|
||||
// LOAD_DEFAULT, WARP_STORE_DIRECT}
|
||||
//
|
||||
// CCCL SM70 policy (lines 167-186):
|
||||
// large: {256, 19, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, RADIX_RANK_MEMOIZE,
|
||||
// BLOCK_SCAN_WARP_SCANS, radix_bits}
|
||||
// medium: {256, 32, medium_itp, WARP_LOAD_DIRECT, LOAD_DEFAULT}
|
||||
// small: {256, keys_only?4:8, small_itp, WARP_LOAD_DIRECT, LOAD_DEFAULT}
|
||||
// ============================================================================
|
||||
|
||||
struct policy_selector {
|
||||
int key_size;
|
||||
int value_size;
|
||||
bool keys_only;
|
||||
|
||||
constexpr SegmentedSortPolicy operator()(const hardware_capability& hw) const {
|
||||
int pair_size = key_size + (keys_only ? 0 : value_size);
|
||||
|
||||
// Large segments: full block sort
|
||||
auto [lg_items, lg_threads] = scale_mem_bound(256, 11, pair_size);
|
||||
// Medium: partial
|
||||
auto [md_items, md_threads] = scale_mem_bound(128, 15, pair_size);
|
||||
// Small: warp sort
|
||||
int sm_threads = 32;
|
||||
int sm_items = 4;
|
||||
constexpr int dominant_size() const {
|
||||
return value_size > key_size ? value_size : key_size;
|
||||
}
|
||||
|
||||
return {lg_threads, lg_items, md_threads, md_items,
|
||||
sm_threads, sm_items, BLOCK_LOAD_WARP_TRANSPOSE};
|
||||
constexpr SegmentedSortPolicy operator()(const hardware_capability& hw) const {
|
||||
int dom = dominant_size();
|
||||
bool large_items = dom > 4;
|
||||
int radix_bits = key_size > 1 ? 6 : 4;
|
||||
|
||||
// --- Large segment: radix sort ---
|
||||
// SM86 tuning: {256, 23} scaled by dominant size
|
||||
auto lg_scaled = scale_reg_bound(256, 23, dom);
|
||||
|
||||
// SMEM check for large: threads * items * dom + rank tables
|
||||
// rank tables (MEMOIZE): (1 << radix_bits) * sizeof(int) * WARP_THREADS = 2^6 * 4 * 32 = 8192B
|
||||
// tile: lg_threads * lg_items * dom
|
||||
int lg_t = lg_scaled.threads_per_block;
|
||||
int lg_i = lg_scaled.items_per_thread;
|
||||
int rank_smem = (1 << radix_bits) * 4 * (lg_t / 32); // per-warp privatized
|
||||
int lg_tile = lg_t * lg_i * dom;
|
||||
while (lg_tile + rank_smem > hw.max_shared_memory_per_block - 2048 && lg_i > 1) {
|
||||
lg_i--;
|
||||
lg_tile = lg_t * lg_i * dom;
|
||||
}
|
||||
|
||||
auto large = SegmentedSortRadixSortPolicy{
|
||||
lg_t, lg_i,
|
||||
BLOCK_LOAD_TRANSPOSE, LOAD_DEFAULT,
|
||||
RadixRankAlgorithm::MEMOIZE, BLOCK_SCAN_WARP_SCANS,
|
||||
radix_bits
|
||||
};
|
||||
|
||||
// --- Medium segment: sub-warp merge sort, 16 threads per segment ---
|
||||
// SM86: threads_per_warp=16, medium_itp depends on large_items
|
||||
int medium_itp = nominal_4b_items_to_items(large_items ? 9 : 7, dom);
|
||||
// SMEM: segments_per_block * items_per_tile * dom * 2 (double buffer)
|
||||
// segments_per_block = 256/16 = 16
|
||||
// items_per_tile = 16 * medium_itp
|
||||
int med_segs = 256 / 16;
|
||||
int med_tile = 16 * medium_itp;
|
||||
int med_smem = med_segs * med_tile * dom * 2;
|
||||
while (med_smem > hw.max_shared_memory_per_block - 2048 && medium_itp > 1) {
|
||||
medium_itp--;
|
||||
med_tile = 16 * medium_itp;
|
||||
med_smem = med_segs * med_tile * dom * 2;
|
||||
}
|
||||
|
||||
auto medium = SegmentedSortSubWarpMergeSortPolicy{
|
||||
256, 16, medium_itp,
|
||||
WarpLoadAlgorithm::TRANSPOSE, LOAD_LDG, WarpStoreAlgorithm::DIRECT
|
||||
};
|
||||
|
||||
// --- Small segment: sub-warp merge sort, 2-8 threads per segment ---
|
||||
// SM86: threads_per_warp = large_items ? 8 : 2
|
||||
int small_tpw = large_items ? 8 : 2;
|
||||
int small_itp = nominal_4b_items_to_items(large_items ? 7 : 9, dom);
|
||||
int sm_segs = 256 / small_tpw;
|
||||
int sm_tile = small_tpw * small_itp;
|
||||
int sm_smem = sm_segs * sm_tile * dom * 2;
|
||||
while (sm_smem > hw.max_shared_memory_per_block - 2048 && small_itp > 1) {
|
||||
small_itp--;
|
||||
sm_tile = small_tpw * small_itp;
|
||||
sm_smem = sm_segs * sm_tile * dom * 2;
|
||||
}
|
||||
|
||||
auto small = SegmentedSortSubWarpMergeSortPolicy{
|
||||
256, small_tpw, small_itp,
|
||||
WarpLoadAlgorithm::TRANSPOSE, LOAD_LDG, WarpStoreAlgorithm::DIRECT
|
||||
};
|
||||
|
||||
return SegmentedSortPolicy{large, medium, small, 500};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user