这是项目历史上第一次用真实 benchmark 数据替换拍脑袋参数。
scan.cuh — bi100_lookback_4B_o4:
实测: dcid_0.ipt_22.l2w_500.ld_0.ns_1904.tpb_384.trp_1
speedups: 1.038085 1.009473 1.007679 1.005803 SMEM=33792 (69%)
关键发现: ns×0.5 假设是错的。实测最优 ns=1904 (和 SM100 原值相同)。
dcid=0 (no_delay) 胜过 dcid=6 (exponential_backon_jitter)。
原因: 16 SMs = ~32 CTAs, lookback contention 极小, 不需要 delay 策略。
改动: delay 从 {exponential_backon_jitter, 952, 498} → {no_delay, 1904, 500}
topk.cuh:
实测: ipt_4.ld_0.tpb_512 speedups: 1.039611 1.000222 1.004295
确认 CCCL SM90+ 公式 (items=4*4/key_size=4, threads=512) 在 BI-V100 上也是最优。
ld=0 (LOAD_DEFAULT) 胜过 ld=1 at small sizes。
ipt=16 在 32K+ 明显回退 → items 不能太大。
transform.cuh — bytes_in_flight:
实测: alg_1.bif_8.pref_2.tpb_256.unrl_1.vsp2_1 1.203199 1.058919 1.019168
bif=8 (64KB) 全面胜过 bif=0 (32KB) 和 bif=-8 (16KB)。
Top 30 结果全部是 bif=8 → 高置信度。
改动: bi100_bytes_in_flight 从 32KB → 64KB。
物理解释: 56 GB/s per SM × ~1100ns HBM latency ≈ 62KB, 和 64KB 吻合。
跨算法发现:
- BI-V100 的 16 SMs 使得 inter-CTA contention 很低
- CCCL 的 delay 策略 (为 80-148 SMs 设计) 在 16 SMs 上过度保守
- 各算法的 threads/items 最优值和 SM100 接近, 但 delay/bif 参数差异大
114 lines
3.8 KiB
Plaintext
114 lines
3.8 KiB
Plaintext
// muh/include/muh/tuning/tuning_topk.cuh — BI-V100 top-k tuning
|
||
//
|
||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_topk.cuh
|
||
//
|
||
// vllm impact: Top-k / top-p sampling in decode stage
|
||
// Competition weight: Output TPS × 16.796 (highest priority, tied with reduce)
|
||
//
|
||
// CCCL SM90+ policy_selector (the ground truth):
|
||
// bits_per_pass = calc_bits_per_pass(key_size):
|
||
// key_size 1 → 8
|
||
// key_size 2 → 8 (but note: CCCL default returns 11 for 2/4/8,
|
||
// only key_size=1 returns 8. See switch below.)
|
||
// key_size 4 → 11
|
||
// key_size 8 → 11
|
||
// items_per_thread = max(1, 4 * 4 / key_size) // 16 bytes per thread
|
||
// threads_per_block = 512
|
||
// load_algorithm = BLOCK_LOAD_VECTORIZE (NOT BLOCK_LOAD_DIRECT)
|
||
// scan_algorithm = BLOCK_SCAN_WARP_SCANS
|
||
|
||
#pragma once
|
||
|
||
#include "muh/hardware.cuh"
|
||
#include "muh/tuning/common.cuh"
|
||
|
||
namespace muh::tuning::topk {
|
||
|
||
/// Top-k policy (mirrors cub::detail::topk::topk_policy)
|
||
struct TopkPolicy {
|
||
int threads_per_block;
|
||
int items_per_thread;
|
||
BlockLoadAlgorithm load_algorithm;
|
||
BlockScanAlgorithm scan_algorithm;
|
||
int bits_per_pass;
|
||
};
|
||
|
||
// ============================================================
|
||
// bits_per_pass calculation — must match CCCL exactly
|
||
//
|
||
// CCCL source (tuning_topk.cuh):
|
||
// case 1: default: return 8;
|
||
// case 2: case 4: case 8: return 11;
|
||
//
|
||
// Previous muh version had a wrong mapping:
|
||
// key_size<=2 → 8, key_size<=4 → 9, key_size<=8 → 10
|
||
// This was WRONG. CCCL's actual function returns 11 for 2/4/8.
|
||
// ============================================================
|
||
|
||
constexpr int calc_bits_per_pass(int key_size) {
|
||
switch (key_size) {
|
||
case 1:
|
||
default:
|
||
return 8;
|
||
case 2:
|
||
case 4:
|
||
case 8:
|
||
return 11;
|
||
}
|
||
}
|
||
|
||
// ============================================================
|
||
// policy_selector
|
||
//
|
||
// Matches CCCL's SM90+ path exactly:
|
||
// threads = 512
|
||
// items = max(1, nominal_4b(4) * 4 / key_size)
|
||
// load = BLOCK_LOAD_VECTORIZE
|
||
// scan = BLOCK_SCAN_WARP_SCANS
|
||
// bits = calc_bits_per_pass(key_size)
|
||
//
|
||
// BI-V100 values: using SM100 as starting point.
|
||
// Once benchmarked on BI-V100, items/threads may diverge.
|
||
// ============================================================
|
||
|
||
struct policy_selector {
|
||
int key_size;
|
||
|
||
constexpr TopkPolicy operator()(const hardware_capability& hw) const {
|
||
if (hw.at_least(hardware_capability::vendor_t::iluvatar, 100)) {
|
||
// BI-V100 BENCHMARK RESULT (bench_bi100.py topk/float32):
|
||
// #1: ipt_4.ld_0.tpb_512 speedups: 1.039611 1.000222 1.004295
|
||
// (baseline: 1K=76.6us, 32K=220.8us, 152K=554.4us)
|
||
// #2: ipt_1.ld_1.tpb_256 speedups: 1.017337 1.020884 1.000531
|
||
// #3: ipt_1.ld_1.tpb_512 speedups: 0.986497 1.047220 1.004375
|
||
//
|
||
// KEY FINDINGS:
|
||
// - ipt=4, tpb=512 matches CCCL SM90+ formula (4*4/4=4) → CONFIRMED
|
||
// - ld=0 (LOAD_DEFAULT) beats ld=1 (LOAD_LDG/LOAD_CA) at small sizes
|
||
// - For 32K+ items, ld=1 is competitive but ipt=4 ld=0 wins overall
|
||
// - ipt=16 regresses at 32K and 152K sizes (too many items per thread)
|
||
constexpr int nominal_4b_items = 4;
|
||
int items = nominal_4b_items * 4 / key_size;
|
||
if (items < 1) items = 1;
|
||
|
||
return {512, items,
|
||
BLOCK_LOAD_VECTORIZE, // CCCL VECTORIZE; BI-V100 ld=0 confirmed best
|
||
BLOCK_SCAN_WARP_SCANS,
|
||
calc_bits_per_pass(key_size)};
|
||
}
|
||
|
||
// Fallback: older arch path
|
||
constexpr int nominal_4b_items = 4;
|
||
int items = nominal_4b_items * 4 / key_size;
|
||
if (items < 1) items = 1;
|
||
if (items > nominal_4b_items) items = nominal_4b_items;
|
||
|
||
return {512, items,
|
||
BLOCK_LOAD_VECTORIZE,
|
||
BLOCK_SCAN_WARP_SCANS,
|
||
calc_bits_per_pass(key_size)};
|
||
}
|
||
};
|
||
|
||
} // namespace muh::tuning::topk
|