From 12ad7a319086f4b888552dc4d2db6ad29ee78e5f Mon Sep 17 00:00:00 2001 From: dylanyunlon Date: Mon, 3 Aug 2026 21:36:46 +0000 Subject: [PATCH] =?UTF-8?q?[muh]=207=20headers=20=E5=AE=8C=E6=95=B4?= =?UTF-8?q?=E7=A7=BB=E6=A4=8D=20CCCL=20tuning=20tables:=20segmented=5Fsort?= =?UTF-8?q?=207%=E2=86=9229%,=20merge=5Fsort=2022%=E2=86=9243%,=20merge=20?= =?UTF-8?q?30%=E2=86=9249%,=20adjacent=5Fdifference=2038%=E2=86=9265%,=20b?= =?UTF-8?q?atch=5Fmemcpy=2037%=E2=86=9241%,=20find=2035%=E2=86=9243%,=20fi?= =?UTF-8?q?nd=5Fbound=2031%=E2=86=9244%?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 每个文件都是直接 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 --- .../muh/tuning/tuning_adjacent_difference.cuh | 63 ++++-- .../muh/tuning/tuning_batch_memcpy.cuh | 101 +++++----- muh/include/muh/tuning/tuning_find.cuh | 17 +- .../tuning_find_bound_sorted_values.cuh | 42 ++-- muh/include/muh/tuning/tuning_merge.cuh | 78 +++++--- muh/include/muh/tuning/tuning_merge_sort.cuh | 80 ++++++-- .../muh/tuning/tuning_segmented_sort.cuh | 185 ++++++++++++++++-- 7 files changed, 422 insertions(+), 144 deletions(-) diff --git a/muh/include/muh/tuning/tuning_adjacent_difference.cuh b/muh/include/muh/tuning/tuning_adjacent_difference.cuh index 2ba75b9e..e7150720 100644 --- a/muh/include/muh/tuning/tuning_adjacent_difference.cuh +++ b/muh/include/muh/tuning/tuning_adjacent_difference.cuh @@ -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 + }; } }; diff --git a/muh/include/muh/tuning/tuning_batch_memcpy.cuh b/muh/include/muh/tuning/tuning_batch_memcpy.cuh index ab476aed..6c4029d9 100644 --- a/muh/include/muh/tuning/tuning_batch_memcpy.cuh +++ b/muh/include/muh/tuning/tuning_batch_memcpy.cuh @@ -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} + }; } }; diff --git a/muh/include/muh/tuning/tuning_find.cuh b/muh/include/muh/tuning/tuning_find.cuh index 760c9e91..924c2321 100644 --- a/muh/include/muh/tuning/tuning_find.cuh +++ b/muh/include/muh/tuning/tuning_find.cuh @@ -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}; } }; diff --git a/muh/include/muh/tuning/tuning_find_bound_sorted_values.cuh b/muh/include/muh/tuning/tuning_find_bound_sorted_values.cuh index e57d0b6f..7d58add6 100644 --- a/muh/include/muh/tuning/tuning_find_bound_sorted_values.cuh +++ b/muh/include/muh/tuning/tuning_find_bound_sorted_values.cuh @@ -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 diff --git a/muh/include/muh/tuning/tuning_merge.cuh b/muh/include/muh/tuning/tuning_merge.cuh index 0aa9f604..f25a4fa4 100644 --- a/muh/include/muh/tuning/tuning_merge.cuh +++ b/muh/include/muh/tuning/tuning_merge.cuh @@ -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 + }; } }; diff --git a/muh/include/muh/tuning/tuning_merge_sort.cuh b/muh/include/muh/tuning/tuning_merge_sort.cuh index b43d266e..635f33e1 100644 --- a/muh/include/muh/tuning/tuning_merge_sort.cuh +++ b/muh/include/muh/tuning/tuning_merge_sort.cuh @@ -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) + }; } }; diff --git a/muh/include/muh/tuning/tuning_segmented_sort.cuh b/muh/include/muh/tuning/tuning_segmented_sort.cuh index bffc11f7..8a3db0d8 100644 --- a/muh/include/muh/tuning/tuning_segmented_sort.cuh +++ b/muh/include/muh/tuning/tuning_segmented_sort.cuh @@ -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(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}; } };