The core of muh is now C++ headers that mirror CCCL's tuning architecture:
muh/include/muh/
├── hardware.cuh — hardware_capability descriptor (replaces cuda::compute_capability)
├── muh.cuh — top-level include + scoring formula
└── tuning/
├── common.cuh — shared types, compatible with CCCL's common.cuh
├── tuning_reduce.cuh — P0: attention reduction (5 type specializations)
├── tuning_topk.cuh — P0: sampling top-k/top-p (2B/4B key specializations)
├── tuning_scan.cuh — P0: prefix scan (6 lookback + 6 lookahead specializations)
├── tuning_transform.cuh — P1: activation elementwise (SiLU/GELU/RMSNorm)
├── tuning_batch_memcpy.cuh — P1: KV cache block copy
└── tuning_for.cuh — P2: RoPE position encoding
Architecture:
- Each tuning header has a policy_selector struct with operator()(hardware_capability)
- Dispatches on muh::hardware_capability instead of cuda::compute_capability
- bi100_* structs hold per-type tuning values (initialized from CCCL SM100 reference)
- When CCCL headers are available, re-exports their enum types
- When standalone, provides compatible enum definitions
Python files (extract.py, parse.py, gen_yaml.py, gen_patch.py) remain as tooling.
The C++ headers are what actually gets compiled into the vllm binary.
52 lines
1.5 KiB
Plaintext
52 lines
1.5 KiB
Plaintext
// muh/include/muh/tuning/tuning_for.cuh — BI-V100 for-each tuning
|
||
//
|
||
// Mirrors: cccl_upstream/cub/cub/device/dispatch/tuning/tuning_for.cuh
|
||
//
|
||
// vllm impact: RoPE position encoding, simple elementwise kernels
|
||
// Competition weight: contributes to Output TPS
|
||
|
||
#pragma once
|
||
|
||
#include "muh/hardware.cuh"
|
||
#include "muh/tuning/common.cuh"
|
||
|
||
namespace muh::tuning::for_each {
|
||
|
||
/// For-each policy
|
||
struct ForPolicy {
|
||
int threads_per_block;
|
||
int items_per_thread;
|
||
};
|
||
|
||
// ============================================================
|
||
// BI-V100 tuning
|
||
//
|
||
// CCCL reference from tuning_for.cuh:
|
||
// threads_per_block: 256 default, can be runtime-determined if set <1
|
||
// items_per_thread: typically 1-4 for simple elementwise
|
||
//
|
||
// The for_each kernel is extremely simple — it's a parallel_for
|
||
// with no shared memory, no reduction, no scan. Tuning is purely
|
||
// about occupancy (threads × items = tile_size).
|
||
// ============================================================
|
||
|
||
struct bi100_default {
|
||
static constexpr int threads = 256;
|
||
static constexpr int items = 4;
|
||
};
|
||
|
||
// ============================================================
|
||
// policy_selector
|
||
// ============================================================
|
||
|
||
struct policy_selector {
|
||
constexpr ForPolicy operator()(const hardware_capability& hw) const {
|
||
if (hw.at_least(hardware_capability::vendor_t::iluvatar, 100)) {
|
||
return {bi100_default::threads, bi100_default::items};
|
||
}
|
||
return {256, 4};
|
||
}
|
||
};
|
||
|
||
} // namespace muh::tuning::for_each
|