[muh] add tuning_batched_topk.cuh: BI-V100 tuning header for batched_topk
Translated from CCCL cub/device/dispatch/tuning/tuning_batched_topk.cuh. Uses hardware_capability dispatch instead of compute_capability.
This commit is contained in:
41
muh/include/muh/tuning/tuning_batched_topk.cuh
Normal file
41
muh/include/muh/tuning/tuning_batched_topk.cuh
Normal file
@@ -0,0 +1,41 @@
|
||||
// muh/include/muh/tuning/tuning_batched_topk.cuh — BI-V100
|
||||
//
|
||||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh
|
||||
// CCCL: extends topk with batch-level parallelism
|
||||
//
|
||||
// vllm relevance: multi-sequence parallel decode top-k sampling
|
||||
// SMEM risk: same as topk (handled by bits_per_pass)
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "muh/hardware.cuh"
|
||||
#include "muh/tuning/common.cuh"
|
||||
#include "muh/tuning/tuning_topk.cuh"
|
||||
|
||||
namespace muh::tuning::batched_topk {
|
||||
|
||||
struct BatchedTopkPolicy {
|
||||
topk::TopkPolicy per_batch;
|
||||
int max_batches_per_block;
|
||||
};
|
||||
|
||||
struct policy_selector {
|
||||
int key_size;
|
||||
int max_k;
|
||||
|
||||
constexpr BatchedTopkPolicy operator()(const hardware_capability& hw) const {
|
||||
auto base = topk::policy_selector{key_size}(hw);
|
||||
|
||||
// Batches per block: limited by SMEM
|
||||
// Each batch needs: bits_per_pass buckets * sizeof(int) for histogram
|
||||
int buckets = 1 << base.bits_per_pass;
|
||||
int hist_smem = buckets * 4; // sizeof(int)
|
||||
int max_batches = hw.max_shared_memory_per_block / hist_smem;
|
||||
if (max_batches < 1) max_batches = 1;
|
||||
if (max_batches > 32) max_batches = 32; // cap for occupancy
|
||||
|
||||
return {base, max_batches};
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace muh::tuning::batched_topk
|
||||
Reference in New Issue
Block a user