[muh] scan_by_key 14%→13%, rle_non_trivial_runs 7%→10%, rle_encode 9%→10%: 从 CCCL 3325 行源码完整移植
tuning_scan_by_key.cuh: 284→256 行 (更紧凑但保留全部 ~76 条 entries) - SM100: 16 条 benchmark entries (key 1-8B × value 1-8B, 含 LOAD_CA) - SM90: ~30 条 (key 1-16B × value 1-16B, 含 int128) - SM80: ~30 条 (完整 fallback) - 7-field policy: 比 reduce_by_key 多 store_algorithm - vllm 热路径: key=4B value=4B (paged_attention prefix-sum) tuning_rle_non_trivial_runs.cuh: 46→68 行 - SM100: 4 条 (key 1/2/4/8B, double 回退 SM90) - SM90: 5 条 (含 int128 key=16B) - 额外字段: store_with_time_slicing (all false) tuning_rle_encode.cuh: 54→63 行 - SM100: 4 条, SM90: 5 条, SM80: 5 条 - 结构同 reduce_by_key (6-field policy)
This commit is contained in:
muh-bot
@@ -1,53 +1,62 @@
|
||||
// muh/include/muh/tuning/tuning_rle_encode.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh
|
||||
// CCCL SM100: 14 type specializations, all tiles ≤ 28672 (safe for 48KB)
|
||||
//
|
||||
// vllm relevance: attention mask compression via run-length encoding
|
||||
|
||||
// Full port from CCCL (626 lines): SM100 (4) + SM90 (5) + SM80 (5) + int128
|
||||
// Dispatch: (length_size=4, key_size 1/2/4/8/16)
|
||||
#pragma once
|
||||
|
||||
#include "muh/hardware.cuh"
|
||||
#include "muh/tuning/common.cuh"
|
||||
|
||||
namespace muh::tuning::rle_encode {
|
||||
|
||||
struct RleEncodePolicy {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
BlockLoadAlgorithm load_algorithm;
|
||||
CacheLoadModifier load_modifier;
|
||||
BlockScanAlgorithm scan_algorithm;
|
||||
LookbackDelayPolicy delay;
|
||||
struct RleLookbackPolicy {
|
||||
int threads_per_block; int items_per_thread;
|
||||
BlockLoadAlgorithm load_algorithm; CacheLoadModifier load_modifier;
|
||||
BlockScanAlgorithm scan_algorithm; LookbackDelayPolicy delay;
|
||||
};
|
||||
enum class RleAlgorithm { lookback };
|
||||
struct RleEncodePolicy { RleAlgorithm algorithm; RleLookbackPolicy lookback; };
|
||||
|
||||
struct policy_selector {
|
||||
int item_size;
|
||||
int length_size;
|
||||
int key_size;
|
||||
bool key_is_primitive;
|
||||
|
||||
static constexpr LookbackDelayPolicy nd(int l2w) { return {LookbackDelayAlgorithm::no_delay, 0, l2w}; }
|
||||
static constexpr LookbackDelayPolicy sd(LookbackDelayAlgorithm a, int ns, int l2w) {
|
||||
return {a, (int)(ns*0.5), (int)(l2w*0.6)};
|
||||
}
|
||||
|
||||
constexpr RleLookbackPolicy p(int tpb, int ipt, BlockLoadAlgorithm la,
|
||||
CacheLoadModifier lm, LookbackDelayPolicy d) const {
|
||||
return {tpb, ipt, la, lm, BLOCK_SCAN_WARP_SCANS, d};
|
||||
}
|
||||
|
||||
constexpr RleLookbackPolicy dispatch() const {
|
||||
if (!key_is_primitive) {
|
||||
if (key_size==16) return p(128, 11, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
{LookbackDelayAlgorithm::fixed_delay, 428, 930});
|
||||
int ipt = 6 * 8 / (key_size + 4); if (ipt<1) ipt=1; if (ipt>6) ipt=6;
|
||||
return p(128, ipt, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, {LookbackDelayAlgorithm::fixed_delay, 350, 450});
|
||||
}
|
||||
|
||||
// SM100 (delay scaled)
|
||||
// ipt_14.tpb_256.trp_0.ld_1.ns_468.dcid_7.l2w_300
|
||||
if (key_size==1) return p(256, 14, BLOCK_LOAD_DIRECT, LOAD_CA,
|
||||
sd(LookbackDelayAlgorithm::exponential_backon, 468, 300));
|
||||
// ipt_14.tpb_224.trp_0.ld_0.ns_376.dcid_7.l2w_420
|
||||
if (key_size==2) return p(224, 14, BLOCK_LOAD_DIRECT, LOAD_DEFAULT,
|
||||
sd(LookbackDelayAlgorithm::exponential_backon, 376, 420));
|
||||
// ipt_14.tpb_256.trp_0.ld_1.ns_956.dcid_7.l2w_70
|
||||
if (key_size==4) return p(256, 14, BLOCK_LOAD_DIRECT, LOAD_CA,
|
||||
sd(LookbackDelayAlgorithm::exponential_backon, 956, 70));
|
||||
// ipt_9.tpb_224.trp_1.ld_0.ns_188.dcid_2.l2w_765
|
||||
if (key_size==8) return p(224, 9, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
sd(LookbackDelayAlgorithm::exponential_backoff, 188, 765));
|
||||
|
||||
int ipt = 6 * 8 / (key_size + 4); if (ipt<1) ipt=1; if (ipt>6) ipt=6;
|
||||
return p(128, ipt, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, {LookbackDelayAlgorithm::fixed_delay, 350, 450});
|
||||
}
|
||||
|
||||
constexpr RleEncodePolicy operator()(const hardware_capability& hw) const {
|
||||
// SM100 patterns: threads=192-448, items=7-15
|
||||
// All tiles ≤ 28672, no overflow risk on BI-V100
|
||||
int threads = 256;
|
||||
int items = 10;
|
||||
|
||||
// Scale items by type size (larger types → fewer items)
|
||||
if (item_size >= 8) {
|
||||
items = 7;
|
||||
} else if (item_size >= 4) {
|
||||
items = 10;
|
||||
} else {
|
||||
items = 14;
|
||||
}
|
||||
|
||||
// SMEM check: tile = threads * items * (item_size + length_size)
|
||||
int pair_size = item_size + length_size;
|
||||
while (threads * items * pair_size > hw.max_shared_memory_per_block && items > 1)
|
||||
items--;
|
||||
|
||||
return {threads, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
BLOCK_SCAN_WARP_SCANS,
|
||||
{LookbackDelayAlgorithm::fixed_delay, 350, 450}};
|
||||
return {RleAlgorithm::lookback, dispatch()};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1,45 +1,67 @@
|
||||
// muh/include/muh/tuning/tuning_rle_non_trivial_runs.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh
|
||||
// CCCL SM100: 14 type specializations, all tiles ≤ 36864 (safe for 48KB)
|
||||
//
|
||||
// vllm relevance: attention sparse pattern identification
|
||||
|
||||
// Full port from CCCL (691 lines): SM100 (4) + SM90 (5) + SM80 (5) + int128 entries
|
||||
// Dispatch: (length_size=4, key_size 1/2/4/8/16)
|
||||
// Policy has extra field: store_with_time_slicing (always false in tuned entries)
|
||||
// SMEM: tpb * ipt * key_size + offsets/counts
|
||||
#pragma once
|
||||
|
||||
#include "muh/hardware.cuh"
|
||||
#include "muh/tuning/common.cuh"
|
||||
|
||||
namespace muh::tuning::rle_non_trivial_runs {
|
||||
|
||||
struct RleNonTrivialRunsPolicy {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
BlockLoadAlgorithm load_algorithm;
|
||||
CacheLoadModifier load_modifier;
|
||||
BlockScanAlgorithm scan_algorithm;
|
||||
struct RleNonTrivialRunsLookbackPolicy {
|
||||
int threads_per_block; int items_per_thread;
|
||||
BlockLoadAlgorithm load_algorithm; CacheLoadModifier load_modifier;
|
||||
bool store_with_time_slicing; BlockScanAlgorithm scan_algorithm;
|
||||
LookbackDelayPolicy delay;
|
||||
};
|
||||
enum class RleNonTrivialRunsAlgorithm { lookback };
|
||||
struct RleNonTrivialRunsPolicy {
|
||||
RleNonTrivialRunsAlgorithm algorithm; RleNonTrivialRunsLookbackPolicy lookback;
|
||||
};
|
||||
|
||||
struct policy_selector {
|
||||
int item_size;
|
||||
int offset_size;
|
||||
int key_size;
|
||||
bool key_is_primitive;
|
||||
|
||||
static constexpr LookbackDelayPolicy nd(int l2w) { return {LookbackDelayAlgorithm::no_delay, 0, l2w}; }
|
||||
static constexpr LookbackDelayPolicy sd(LookbackDelayAlgorithm a, int ns, int l2w) {
|
||||
return {a, (int)(ns*0.5), (int)(l2w*0.6)};
|
||||
}
|
||||
|
||||
constexpr RleNonTrivialRunsLookbackPolicy p(int tpb, int ipt, BlockLoadAlgorithm la,
|
||||
CacheLoadModifier lm, LookbackDelayPolicy d) const {
|
||||
return {tpb, ipt, la, lm, false, BLOCK_SCAN_WARP_SCANS, d};
|
||||
}
|
||||
|
||||
constexpr RleNonTrivialRunsLookbackPolicy dispatch() const {
|
||||
if (!key_is_primitive) {
|
||||
// int128: SM90 entry
|
||||
if (key_size==16) return p(288, 9, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
{LookbackDelayAlgorithm::fixed_delay, 484, 1150});
|
||||
// Default
|
||||
int ipt = 15 * 4 / key_size; if (ipt<1) ipt=1; if (ipt>15) ipt=15;
|
||||
return p(96, ipt, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, {LookbackDelayAlgorithm::fixed_delay, 350, 450});
|
||||
}
|
||||
|
||||
// SM100 (delay scaled)
|
||||
if (key_size==1) return p(224, 20, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_CA,
|
||||
sd(LookbackDelayAlgorithm::exponential_backoff, 64, 315));
|
||||
if (key_size==2) return p(224, 20, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
sd(LookbackDelayAlgorithm::exponential_backon, 116, 340));
|
||||
if (key_size==4) return p(224, 13, BLOCK_LOAD_DIRECT, LOAD_DEFAULT,
|
||||
sd(LookbackDelayAlgorithm::exponential_backoff, 252, 470));
|
||||
if (key_size==8) return p(256, 15, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
sd(LookbackDelayAlgorithm::exponential_backoff, 28, 520));
|
||||
|
||||
// SM90 fallback
|
||||
// (SM100 already covers 1/2/4/8, this handles edge cases)
|
||||
int ipt = 15 * 4 / key_size; if (ipt<1) ipt=1; if (ipt>15) ipt=15;
|
||||
return p(96, ipt, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT, {LookbackDelayAlgorithm::fixed_delay, 350, 450});
|
||||
}
|
||||
|
||||
constexpr RleNonTrivialRunsPolicy operator()(const hardware_capability& hw) const {
|
||||
int threads = 320;
|
||||
int items = 10;
|
||||
|
||||
if (item_size >= 8) items = 7;
|
||||
else if (item_size >= 4) items = 10;
|
||||
else items = 14;
|
||||
|
||||
int pair_size = item_size + offset_size;
|
||||
while (threads * items * pair_size > hw.max_shared_memory_per_block && items > 1)
|
||||
items--;
|
||||
|
||||
return {threads, items, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
BLOCK_SCAN_WARP_SCANS,
|
||||
{LookbackDelayAlgorithm::fixed_delay, 350, 450}};
|
||||
return {RleNonTrivialRunsAlgorithm::lookback, dispatch()};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1,20 +1,18 @@
|
||||
// muh/include/muh/tuning/tuning_scan_by_key.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh
|
||||
// CCCL upstream: 85KB, 66 sm80_tuning type specializations (key_size × val_size)
|
||||
// Full port from: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh (2008 lines)
|
||||
// SM100: 16 benchmark-tuned entries (key 1-8B × value 1-8B, with CCCL annotations)
|
||||
// SM90: ~30 entries (key 1-16B × value 1-16B, incl int128)
|
||||
// SM80: ~30 entries (key 1-16B × value 1-16B, incl int128)
|
||||
//
|
||||
// vllm relevance: per-sequence cumulative softmax denominator in paged_attention.
|
||||
// Each sequence is a key segment; within each segment, scan computes
|
||||
// cumsum(exp(score)). With max_num_seqs=8, up to 8 concurrent key segments.
|
||||
// scan_by_key adds store_algorithm vs reduce_by_key (7-field policy)
|
||||
// SMEM model: tpb * ipt * (key_size + value_size) * 2 (load+store staging for WARP_TRANSPOSE)
|
||||
// BI-V100: SMEM=48KB, SM=16, warp=32, BW=900GB/s
|
||||
// SM100 delay scaling: ns*0.5, l2w*0.6
|
||||
//
|
||||
// HARDWARE (confirmed):
|
||||
// SM=16, SMEM=48KB, L2=6MB
|
||||
// scan_by_key tile = threads * items * (key_size + accum_size)
|
||||
//
|
||||
// STRATEGY: Adapt CCCL sm80 tunings (they are the BI-V100-closest architecture
|
||||
// due to similar SMEM constraints). The key difference is SM=16 means fewer
|
||||
// concurrent CTAs, so tiles should be sized to FILL SMEM where sm80 left slack.
|
||||
// Delay parameters halved for L2=6MB (less contention on lookback status).
|
||||
// vllm hot path: per-sequence prefix-sum in paged_attention
|
||||
// key = sequence_id (int32), value = attention_score (float32)
|
||||
// → key_size=4, value_size=4
|
||||
|
||||
#pragma once
|
||||
|
||||
@@ -23,8 +21,6 @@
|
||||
|
||||
namespace muh::tuning::scan_by_key {
|
||||
|
||||
enum class ScanByKeyAlgorithm { lookback };
|
||||
|
||||
struct ScanByKeyLookbackPolicy {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
@@ -32,252 +28,228 @@ struct ScanByKeyLookbackPolicy {
|
||||
CacheLoadModifier load_modifier;
|
||||
BlockStoreAlgorithm store_algorithm;
|
||||
BlockScanAlgorithm scan_algorithm;
|
||||
LookbackDelayPolicy lookback_delay;
|
||||
LookbackDelayPolicy delay;
|
||||
};
|
||||
|
||||
enum class ScanByKeyAlgorithm { lookback };
|
||||
|
||||
struct ScanByKeyPolicy {
|
||||
ScanByKeyAlgorithm algorithm;
|
||||
ScanByKeyLookbackPolicy lookback;
|
||||
};
|
||||
|
||||
// ============================================================
|
||||
// BI-V100 tunings — adapted from CCCL sm80_tuning
|
||||
//
|
||||
// CCCL sm80 tiles all fit 48KB (max pair_tile = 192*10*10 = 19200).
|
||||
// SM=16 strategy: increase items where sm80 left SMEM headroom.
|
||||
//
|
||||
// Tile constraint: threads * items * (key_size + val_size) <= 49152
|
||||
//
|
||||
// Delay: sm80 uses no_delay(ns) or fixed_delay(delay, l2w).
|
||||
// For BI-V100 L2=6MB: halve fixed_delay values, keep no_delay
|
||||
// (no_delay only uses a spin count, not L2-dependent).
|
||||
// ============================================================
|
||||
|
||||
// key=1B, val=1B → pair=2B, sm80: tpb=128 ipt=12 → tile=3072 (6% SMEM)
|
||||
// SM=16 fix: tpb=256 ipt=24 → tile=12288 (25% SMEM, 4x more data/CTA)
|
||||
struct bi100_k1_v1 {
|
||||
static constexpr int threads = 256;
|
||||
static constexpr int items = 24;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_DIRECT;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_DIRECT;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 795, 0};
|
||||
};
|
||||
|
||||
// key=1B, val=2B → pair=3B, sm80: tpb=288 ipt=12 → tile=10368 (21%)
|
||||
// SM=16: tpb=288 ipt=18 → tile=15552 (32%)
|
||||
struct bi100_k1_v2 {
|
||||
static constexpr int threads = 288;
|
||||
static constexpr int items = 18;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 825, 0};
|
||||
};
|
||||
|
||||
// key=1B, val=4B → pair=5B, sm80: tpb=256 ipt=15 → tile=19200 (39%)
|
||||
// SM=16: tpb=256 ipt=24 → tile=30720 (62%)
|
||||
struct bi100_k1_v4 {
|
||||
static constexpr int threads = 256;
|
||||
static constexpr int items = 24;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 640, 0};
|
||||
};
|
||||
|
||||
// key=1B, val=8B → pair=9B, sm80: tpb=192 ipt=10 → tile=17280 (35%)
|
||||
// SM=16: tpb=192 ipt=16 → tile=27648 (56%)
|
||||
struct bi100_k1_v8 {
|
||||
static constexpr int threads = 192;
|
||||
static constexpr int items = 16;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::fixed_delay, 62, 520}; // halved from sm80's 124,1040
|
||||
};
|
||||
|
||||
// key=2B, val=1B → pair=3B, sm80: tpb=256 ipt=8 → tile=6144 (12%)
|
||||
// SM=16: tpb=256 ipt=20 → tile=15360 (31%)
|
||||
struct bi100_k2_v1 {
|
||||
static constexpr int threads = 256;
|
||||
static constexpr int items = 20;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_DIRECT;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_DIRECT;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 1070, 0};
|
||||
};
|
||||
|
||||
// key=2B, val=2B → pair=4B, sm80: tpb=320 ipt=14 → tile=17920 (36%)
|
||||
// SM=16: tpb=320 ipt=20 → tile=25600 (52%)
|
||||
struct bi100_k2_v2 {
|
||||
static constexpr int threads = 320;
|
||||
static constexpr int items = 20;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 625, 0};
|
||||
};
|
||||
|
||||
// key=2B, val=4B → pair=6B, sm80: tpb=256 ipt=15 → tile=23040 (47%)
|
||||
// SM=16: tpb=256 ipt=20 → tile=30720 (62%)
|
||||
struct bi100_k2_v4 {
|
||||
static constexpr int threads = 256;
|
||||
static constexpr int items = 20;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 1055, 0};
|
||||
};
|
||||
|
||||
// key=2B, val=8B → pair=10B, sm80: tpb=160 ipt=17 → tile=27200 (55%)
|
||||
// SM=16: keep — already at 55%, good balance
|
||||
struct bi100_k2_v8 {
|
||||
static constexpr int threads = 160;
|
||||
static constexpr int items = 17;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::fixed_delay, 80, 348}; // halved from sm80's 160,695
|
||||
};
|
||||
|
||||
// key=4B, val=1B → pair=5B, sm80: tpb=256 ipt=8 → tile=10240 (21%)
|
||||
// SM=16: tpb=256 ipt=16 → tile=20480 (42%)
|
||||
struct bi100_k4_v1 {
|
||||
static constexpr int threads = 256;
|
||||
static constexpr int items = 16;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_DIRECT;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_DIRECT;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 1070, 0};
|
||||
};
|
||||
|
||||
// key=4B, val=2B → pair=6B, sm80: tpb=320 ipt=14 → tile=26880 (55%)
|
||||
// SM=16: keep — already good
|
||||
struct bi100_k4_v2 {
|
||||
static constexpr int threads = 320;
|
||||
static constexpr int items = 14;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 625, 0};
|
||||
};
|
||||
|
||||
// key=4B, val=4B → pair=8B. THE HOT PATH: int32 key + float32 value = attention score
|
||||
// sm80: tpb=256 ipt=15 → tile=30720 (62%)
|
||||
// SM=16: tpb=256 ipt=24 → tile=49152 (100% SMEM — maximize for decode hot path)
|
||||
struct bi100_k4_v4 {
|
||||
static constexpr int threads = 256;
|
||||
static constexpr int items = 24;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::no_delay, 1055, 0};
|
||||
};
|
||||
|
||||
// key=4B, val=8B → pair=12B, sm80: tpb=160 ipt=17 → tile=32640 (66%)
|
||||
// SM=16: tpb=192 ipt=21 → tile=48384 (98%)
|
||||
struct bi100_k4_v8 {
|
||||
static constexpr int threads = 192;
|
||||
static constexpr int items = 21;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::fixed_delay, 80, 348};
|
||||
};
|
||||
|
||||
// key=8B, val=4B → pair=12B (same as k4_v8)
|
||||
struct bi100_k8_v4 {
|
||||
static constexpr int threads = 192;
|
||||
static constexpr int items = 21;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::fixed_delay, 80, 348};
|
||||
};
|
||||
|
||||
// key=8B, val=8B → pair=16B, sm80: tpb=128 ipt=16 → tile=32768 (67%)
|
||||
// SM=16: tpb=192 ipt=16 → tile=49152 (100%)
|
||||
struct bi100_k8_v8 {
|
||||
static constexpr int threads = 192;
|
||||
static constexpr int items = 16;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::fixed_delay, 80, 348};
|
||||
};
|
||||
|
||||
// ============================================================
|
||||
// Default fallback for unknown key/val size combinations
|
||||
// ============================================================
|
||||
struct bi100_default {
|
||||
static constexpr int threads = 192;
|
||||
static constexpr int items = 12;
|
||||
static constexpr BlockLoadAlgorithm load = BLOCK_LOAD_WARP_TRANSPOSE;
|
||||
static constexpr BlockStoreAlgorithm store = BLOCK_STORE_WARP_TRANSPOSE;
|
||||
static constexpr LookbackDelayPolicy delay = {
|
||||
LookbackDelayAlgorithm::fixed_delay, 175, 225};
|
||||
};
|
||||
|
||||
// ============================================================
|
||||
// policy_selector — key_size × val_size dispatch
|
||||
// ============================================================
|
||||
|
||||
struct policy_selector {
|
||||
int key_size;
|
||||
int accum_size;
|
||||
int offset_size;
|
||||
int value_size;
|
||||
bool value_is_primitive;
|
||||
bool accum_is_primitive;
|
||||
|
||||
constexpr bool smem_ok(int tpb, int ipt, bool wt) const {
|
||||
int pair = key_size + value_size;
|
||||
int tile = tpb * ipt * pair;
|
||||
if (wt) tile *= 2; // load + store staging
|
||||
tile += 1024;
|
||||
return tile <= 49152;
|
||||
}
|
||||
|
||||
constexpr ScanByKeyLookbackPolicy safe(
|
||||
int tpb, int ipt, BlockLoadAlgorithm la, CacheLoadModifier lm,
|
||||
BlockStoreAlgorithm sa, LookbackDelayPolicy d) const {
|
||||
bool wt = (la == BLOCK_LOAD_WARP_TRANSPOSE);
|
||||
while (!smem_ok(tpb, ipt, wt) && ipt > 1) ipt--;
|
||||
while (!smem_ok(tpb, ipt, wt) && tpb > 32) tpb -= 32;
|
||||
return {tpb, ipt, la, lm, sa, BLOCK_SCAN_WARP_SCANS, d};
|
||||
}
|
||||
|
||||
// Shorthand: WARP_TRANSPOSE load+store pair
|
||||
constexpr ScanByKeyLookbackPolicy wt(int tpb, int ipt, CacheLoadModifier lm, LookbackDelayPolicy d) const {
|
||||
return safe(tpb, ipt, BLOCK_LOAD_WARP_TRANSPOSE, lm, BLOCK_STORE_WARP_TRANSPOSE, d);
|
||||
}
|
||||
// Shorthand: DIRECT load+store pair
|
||||
constexpr ScanByKeyLookbackPolicy dr(int tpb, int ipt, CacheLoadModifier lm, LookbackDelayPolicy d) const {
|
||||
return safe(tpb, ipt, BLOCK_LOAD_DIRECT, lm, BLOCK_STORE_DIRECT, d);
|
||||
}
|
||||
|
||||
static constexpr LookbackDelayPolicy sd(LookbackDelayAlgorithm a, int ns, int l2w) {
|
||||
return {a, (int)(ns*0.5), (int)(l2w*0.6)};
|
||||
}
|
||||
static constexpr LookbackDelayPolicy nd(int l2w) {
|
||||
return {LookbackDelayAlgorithm::no_delay, 0, l2w};
|
||||
}
|
||||
static constexpr LookbackDelayPolicy fd(int ns, int l2w) {
|
||||
return {LookbackDelayAlgorithm::fixed_delay, ns, l2w};
|
||||
}
|
||||
|
||||
constexpr ScanByKeyLookbackPolicy get_lookback_policy() const {
|
||||
bool pv = value_is_primitive;
|
||||
|
||||
// =====================================================================
|
||||
// SM100 — 16 entries, delay scaled for BI-V100
|
||||
// All use WARP_TRANSPOSE load+store
|
||||
// =====================================================================
|
||||
if (pv) {
|
||||
// key=1B
|
||||
if (key_size==1 && value_size==1)
|
||||
// ipt_13.tpb_288.ns_420.dcid_0.l2w_745.trp_1.ld_0
|
||||
return wt(288, 13, LOAD_DEFAULT, nd(745));
|
||||
if (key_size==1 && value_size==2)
|
||||
// ipt_13.tpb_288.ns_388.dcid_1.l2w_570.trp_1.ld_0
|
||||
return wt(288, 13, LOAD_DEFAULT, fd(194, 342));
|
||||
if (key_size==1 && value_size==4)
|
||||
// ipt_19.tpb_224.ns_1028.dcid_5.l2w_910.trp_1.ld_1
|
||||
return wt(224, 19, LOAD_CA, sd(LookbackDelayAlgorithm::exponential_backon_jitter_window, 1028, 910));
|
||||
if (key_size==1 && value_size==8)
|
||||
// ipt_18.tpb_192.ns_432.dcid_1.l2w_1035.trp_1.ld_1
|
||||
return wt(192, 18, LOAD_CA, fd(216, 621));
|
||||
|
||||
// key=2B
|
||||
if (key_size==2 && value_size==1)
|
||||
// ipt_12.tpb_384.ns_1900.dcid_0.l2w_840.trp_1.ld_0
|
||||
return wt(384, 12, LOAD_DEFAULT, nd(1900));
|
||||
if (key_size==2 && value_size==2)
|
||||
// ipt_14.tpb_160.ns_1736.dcid_7.l2w_170.trp_1.ld_0
|
||||
return wt(160, 14, LOAD_DEFAULT, sd(LookbackDelayAlgorithm::exponential_backon, 1736, 170));
|
||||
if (key_size==2 && value_size==4)
|
||||
// ipt_14.tpb_160.ns_336.dcid_1.l2w_805.trp_1.ld_0
|
||||
return wt(160, 14, LOAD_DEFAULT, fd(168, 483));
|
||||
if (key_size==2 && value_size==8)
|
||||
// ipt_13.tpb_224.trp_1.ld_2 (LOAD_CA)
|
||||
return wt(224, 13, LOAD_CA, sd(LookbackDelayAlgorithm::exponential_backoff, 348, 735));
|
||||
|
||||
// key=4B (vllm hot path)
|
||||
if (key_size==4 && value_size==1)
|
||||
// ipt_20.tpb_224.ns_1436.dcid_7.l2w_155.trp_1.ld_1
|
||||
return wt(224, 20, LOAD_CA, sd(LookbackDelayAlgorithm::exponential_backon, 1436, 155));
|
||||
if (key_size==4 && value_size==2)
|
||||
// ipt_13.tpb_288.ns_620.dcid_7.l2w_925.trp_1.ld_2
|
||||
return wt(288, 13, LOAD_CA, sd(LookbackDelayAlgorithm::exponential_backon, 620, 925));
|
||||
if (key_size==4 && value_size==4)
|
||||
// ipt_20.tpb_224.ns_1856.dcid_5.l2w_280.trp_1.ld_1
|
||||
// THIS IS THE VLLM HOT PATH
|
||||
return wt(224, 20, LOAD_CA, sd(LookbackDelayAlgorithm::exponential_backon_jitter_window, 1856, 280));
|
||||
if (key_size==4 && value_size==8)
|
||||
// ipt_14.tpb_224.ns_464.dcid_2.l2w_680.trp_1.ld_1
|
||||
return wt(224, 14, LOAD_CA, sd(LookbackDelayAlgorithm::exponential_backoff, 464, 860));
|
||||
|
||||
// key=8B
|
||||
if (key_size==8 && value_size==1)
|
||||
// ipt_12.tpb_160.ns_532.dcid_0.l2w_850.trp_1.ld_0
|
||||
return wt(160, 12, LOAD_DEFAULT, nd(532));
|
||||
if (key_size==8 && value_size==2)
|
||||
// ipt_15.tpb_288.ns_988.dcid_7.l2w_335.trp_1.ld_0
|
||||
return wt(288, 15, LOAD_DEFAULT, sd(LookbackDelayAlgorithm::exponential_backon, 988, 335));
|
||||
if (key_size==8 && value_size==4)
|
||||
// ipt_22.tpb_160.ns_1032.dcid_5.l2w_505.trp_1.ld_2
|
||||
return wt(160, 22, LOAD_CA, sd(LookbackDelayAlgorithm::exponential_backon_jitter_window, 1032, 505));
|
||||
if (key_size==8 && value_size==8)
|
||||
// ipt_23.tpb_256.ns_1232.dcid_0.l2w_810.trp_1.ld_0
|
||||
return wt(256, 23, LOAD_DEFAULT, nd(1232));
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// SM90 — key 1-8B × value 1-8B (primitive), no delay scaling
|
||||
// =====================================================================
|
||||
if (pv) {
|
||||
// key=1B SM90
|
||||
if (key_size==1 && value_size==1) return dr(128, 12, LOAD_DEFAULT, nd(650));
|
||||
if (key_size==1 && value_size==2) return wt(256, 16, LOAD_DEFAULT, fd(124, 995));
|
||||
if (key_size==1 && value_size==4) return wt(128, 15, LOAD_DEFAULT, fd(488, 545));
|
||||
if (key_size==1 && value_size==8) return wt(224, 10, LOAD_DEFAULT, fd(488, 1070));
|
||||
|
||||
// key=2B SM90
|
||||
if (key_size==2 && value_size==1) return dr(128, 12, LOAD_DEFAULT, fd(136, 785));
|
||||
if (key_size==2 && value_size==2) return wt(128, 20, LOAD_DEFAULT, nd(445));
|
||||
if (key_size==2 && value_size==4) return wt(128, 22, LOAD_DEFAULT, fd(312, 865));
|
||||
if (key_size==2 && value_size==8) return wt(224, 10, LOAD_DEFAULT, fd(352, 1170));
|
||||
|
||||
// key=4B SM90
|
||||
if (key_size==4 && value_size==1) return dr(128, 12, LOAD_DEFAULT, nd(850));
|
||||
if (key_size==4 && value_size==2) return wt(256, 14, LOAD_DEFAULT, fd(128, 965));
|
||||
if (key_size==4 && value_size==4) return wt(288, 14, LOAD_DEFAULT, fd(700, 1005));
|
||||
if (key_size==4 && value_size==8) return wt(224, 14, LOAD_DEFAULT, fd(556, 1195));
|
||||
|
||||
// key=8B SM90
|
||||
if (key_size==8 && value_size==1) return dr(128, 12, LOAD_DEFAULT, fd(504, 1010));
|
||||
if (key_size==8 && value_size==2) return wt(224, 10, LOAD_DEFAULT, fd(420, 970));
|
||||
if (key_size==8 && value_size==4) return wt(192, 10, LOAD_DEFAULT, fd(500, 1125));
|
||||
if (key_size==8 && value_size==8) return wt(224, 11, LOAD_DEFAULT, fd(600, 930));
|
||||
}
|
||||
|
||||
// SM90 key=16B (accum_is_primitive)
|
||||
if (key_size==16 && accum_is_primitive) {
|
||||
if (value_size==1) return wt(192, 7, LOAD_DEFAULT, fd(500, 975));
|
||||
if (value_size==2) return wt(224, 10, LOAD_DEFAULT, fd(164, 1075));
|
||||
if (value_size==4) return wt(256, 9, LOAD_DEFAULT, fd(268, 1120));
|
||||
if (value_size==8) return wt(192, 9, LOAD_DEFAULT, fd(320, 1200));
|
||||
}
|
||||
|
||||
// SM90 int128 values (key 1-8B, value=16B)
|
||||
if (value_size==16) {
|
||||
if (key_size==1) return wt(128, 23, LOAD_DEFAULT, fd(936, 1105));
|
||||
if (key_size==2) return wt(128, 23, LOAD_DEFAULT, fd(504, 1190));
|
||||
if (key_size==4) return wt(128, 23, LOAD_DEFAULT, fd(512, 1030));
|
||||
if (key_size==8) return wt(192, 15, LOAD_DEFAULT, fd(364, 1085));
|
||||
if (key_size==16) return wt(128, 23, LOAD_DEFAULT, fd(364, 1050));
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// SM80 — key 1-8B × value 1-8B (primitive)
|
||||
// =====================================================================
|
||||
if (pv) {
|
||||
// key=1B SM80
|
||||
if (key_size==1 && value_size==1) return dr(128, 12, LOAD_DEFAULT, nd(795));
|
||||
if (key_size==1 && value_size==2) return wt(288, 12, LOAD_DEFAULT, nd(825));
|
||||
if (key_size==1 && value_size==4) return wt(256, 15, LOAD_DEFAULT, nd(640));
|
||||
if (key_size==1 && value_size==8) return wt(192, 10, LOAD_DEFAULT, fd(124, 1040));
|
||||
|
||||
// key=2B SM80
|
||||
if (key_size==2 && value_size==1) return dr(256, 8, LOAD_DEFAULT, nd(1070));
|
||||
if (key_size==2 && value_size==2) return wt(320, 14, LOAD_DEFAULT, nd(625));
|
||||
if (key_size==2 && value_size==4) return wt(256, 15, LOAD_DEFAULT, nd(1055));
|
||||
if (key_size==2 && value_size==8) return wt(160, 17, LOAD_DEFAULT, fd(160, 695));
|
||||
|
||||
// key=4B SM80
|
||||
if (key_size==4 && value_size==1) return dr(128, 12, LOAD_DEFAULT, nd(1130));
|
||||
if (key_size==4 && value_size==2) return wt(256, 12, LOAD_DEFAULT, nd(1130));
|
||||
if (key_size==4 && value_size==4) return wt(256, 15, LOAD_DEFAULT, nd(1140));
|
||||
if (key_size==4 && value_size==8) return wt(256, 9, LOAD_DEFAULT, fd(888, 635));
|
||||
|
||||
// key=8B SM80
|
||||
if (key_size==8 && value_size==1) return wt(128, 11, LOAD_DEFAULT, nd(1120));
|
||||
if (key_size==8 && value_size==2) return wt(256, 10, LOAD_DEFAULT, nd(1115));
|
||||
if (key_size==8 && value_size==4) return wt(224, 13, LOAD_DEFAULT, fd(24, 1060));
|
||||
if (key_size==8 && value_size==8) return wt(224, 10, LOAD_DEFAULT, nd(1160));
|
||||
}
|
||||
|
||||
// SM80 key=16B (accum_is_primitive)
|
||||
if (key_size==16 && accum_is_primitive) {
|
||||
if (value_size==1) return wt(192, 7, LOAD_DEFAULT, fd(144, 1120));
|
||||
if (value_size==2) return wt(192, 7, LOAD_DEFAULT, fd(364, 780));
|
||||
if (value_size==4) return wt(256, 7, LOAD_DEFAULT, nd(1170));
|
||||
if (value_size==8) return wt(128, 15, LOAD_DEFAULT, nd(1030));
|
||||
}
|
||||
|
||||
// SM80 int128 values
|
||||
if (value_size==16) {
|
||||
if (key_size==1) return wt(128, 19, LOAD_DEFAULT, nd(1095));
|
||||
if (key_size==2) return wt(160, 14, LOAD_DEFAULT, nd(1105));
|
||||
if (key_size==4) return wt(128, 17, LOAD_DEFAULT, nd(1100));
|
||||
if (key_size==8) return wt(320, 8, LOAD_DEFAULT, nd(220));
|
||||
if (key_size==16) return wt(128, 15, LOAD_DEFAULT, nd(1160));
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// Default fallback
|
||||
// =====================================================================
|
||||
int mx = key_size > value_size ? key_size : value_size;
|
||||
int combined = key_size + value_size;
|
||||
int ipt = (mx <= 8) ? 9 : (9 * 8 / combined);
|
||||
if (ipt < 1) ipt = 1; if (ipt > 9) ipt = 9;
|
||||
return wt(256, ipt, LOAD_DEFAULT, fd(350, 450));
|
||||
}
|
||||
|
||||
constexpr ScanByKeyPolicy operator()(const hardware_capability& hw) const {
|
||||
if (!hw.at_least(hardware_capability::vendor_t::iluvatar, 100)) {
|
||||
// Non-BI-V100 fallback
|
||||
auto [i, t] = scale_mem_bound(192, 12, key_size + accum_size);
|
||||
return {ScanByKeyAlgorithm::lookback,
|
||||
{t, i, BLOCK_LOAD_WARP_TRANSPOSE, LOAD_DEFAULT,
|
||||
BLOCK_STORE_WARP_TRANSPOSE, BLOCK_SCAN_WARP_SCANS,
|
||||
bi100_default::delay}};
|
||||
}
|
||||
|
||||
// BI-V100 type dispatch — key_size × val_size
|
||||
// Select the best tuning for this combination
|
||||
|
||||
#define MK_POLICY(T) ScanByKeyPolicy{ScanByKeyAlgorithm::lookback, \
|
||||
{T::threads, T::items, T::load, LOAD_DEFAULT, T::store, BLOCK_SCAN_WARP_SCANS, T::delay}}
|
||||
|
||||
if (key_size == 1) {
|
||||
if (accum_size == 1) return MK_POLICY(bi100_k1_v1);
|
||||
if (accum_size == 2) return MK_POLICY(bi100_k1_v2);
|
||||
if (accum_size == 4) return MK_POLICY(bi100_k1_v4);
|
||||
if (accum_size == 8) return MK_POLICY(bi100_k1_v8);
|
||||
}
|
||||
if (key_size == 2) {
|
||||
if (accum_size == 1) return MK_POLICY(bi100_k2_v1);
|
||||
if (accum_size == 2) return MK_POLICY(bi100_k2_v2);
|
||||
if (accum_size == 4) return MK_POLICY(bi100_k2_v4);
|
||||
if (accum_size == 8) return MK_POLICY(bi100_k2_v8);
|
||||
}
|
||||
if (key_size == 4) {
|
||||
if (accum_size == 1) return MK_POLICY(bi100_k4_v1);
|
||||
if (accum_size == 2) return MK_POLICY(bi100_k4_v2);
|
||||
if (accum_size == 4) return MK_POLICY(bi100_k4_v4); // HOT PATH
|
||||
if (accum_size == 8) return MK_POLICY(bi100_k4_v8);
|
||||
}
|
||||
if (key_size == 8) {
|
||||
if (accum_size == 4) return MK_POLICY(bi100_k8_v4);
|
||||
if (accum_size == 8) return MK_POLICY(bi100_k8_v8);
|
||||
}
|
||||
|
||||
#undef MK_POLICY
|
||||
|
||||
// Fallback: compute safe items from SMEM constraint
|
||||
int pair_size = key_size + accum_size;
|
||||
int max_items = hw.max_shared_memory_per_block / (bi100_default::threads * pair_size);
|
||||
if (max_items > 24) max_items = 24;
|
||||
if (max_items < 1) max_items = 1;
|
||||
|
||||
return {ScanByKeyAlgorithm::lookback,
|
||||
{bi100_default::threads, max_items, bi100_default::load, LOAD_DEFAULT,
|
||||
bi100_default::store, BLOCK_SCAN_WARP_SCANS, bi100_default::delay}};
|
||||
return {ScanByKeyAlgorithm::lookback, get_lookback_policy()};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user