From 382d62a21f4fa9c8c633eaf2bb96bc5c049a9f28 Mon Sep 17 00:00:00 2001 From: wangyuzhan Date: Mon, 14 Sep 2026 19:12:11 +0800 Subject: [PATCH] Add optimized SM90 support --- README.md | 7 +++ csrc/api.cpp | 21 ++++++--- ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ ...pancy_1_b_4096_b2_4096_tma_16_cluster_8.cu | 10 ++++ csrc/cuda_kernels/v3_cluster/topk_select.cuh | 2 +- scripts/generate_instantiations.py | 3 +- setup.py | 46 +++++++++++++------ tests/kernelkit/platform.py | 12 +++-- 14 files changed, 146 insertions(+), 25 deletions(-) create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu create mode 100644 csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu diff --git a/README.md b/README.md index fd016c3..c467ecd 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,13 @@ git submodule update --init --recursive pip install -v . ``` +The extension is compiled for `sm_100a` and `sm_103a` by default. To build for +Hopper GPUs such as H20/H100, select `sm_90a` explicitly: + +```bash +DEEP_SELECT_CUDA_ARCHS=90a pip install -v . +``` + ## Usage ```python diff --git a/csrc/api.cpp b/csrc/api.cpp index b1d9285..85cb0a7 100644 --- a/csrc/api.cpp +++ b/csrc/api.cpp @@ -76,10 +76,9 @@ void topk( int64_t cur_stride = tensor.stride(0); uint64_t itemsize = tensor.dtype().itemsize(); TORCH_CHECK(cur_stride * itemsize % alignment_requirement_bytes == 0, - std::format("{}.stride(0) (currently {} numbers) must be a multiple of {} Bytes ({} numbers)", - tensor_name, cur_stride, - alignment_requirement_bytes, alignment_requirement_bytes / itemsize - ) + tensor_name, ".stride(0) (currently ", cur_stride, + " numbers) must be a multiple of ", alignment_requirement_bytes, + " Bytes (", alignment_requirement_bytes / itemsize, " numbers)" ); }; check_dim0_stride("input", input, INPUT_STRIDE_ALIGNMENT_REQUIREMENT); @@ -125,12 +124,20 @@ void topk( TORCH_CHECK((uint32_t)vocab_size < MAX_VOCAB_SIZE, "vocab_size must be < 2^23 for bfloat16 input"); TORCH_CHECK(topk <= 4096, "topk must be <= 4096"); - if (batch_size <= 6 && (uint32_t)vocab_size >= 512u * 1024u && topk <= 1024) { // TODO Tune + const uint32_t cluster_min_vocab_size = + device_prop->major == 9 ? 128u * 1024u : 512u * 1024u; + if (batch_size <= 6 && (uint32_t)vocab_size >= cluster_min_vocab_size && + topk <= 1024) { INTEGER_TYPE_SWITCH(output_index_t, OutIdxT, [&]() { BOOL_SWITCH(sorted_index, SORTED_INDEX, [&]() { BOOL_SWITCH(return_value, RETURN_VALUE, [&]() { - topk_select_bf16_cluster::run_topk_select_kernel< - TopkSelectConfig>(args); + if (device_prop->major == 9) { + topk_select_bf16_cluster::run_topk_select_kernel< + TopkSelectConfig>(args); + } else { + topk_select_bf16_cluster::run_topk_select_kernel< + TopkSelectConfig>(args); + } }); }); }); diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..f180078 --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..70945d4 --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..035821a --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..f53d410 --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i32_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..0d81c00 --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..659e9f3 --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si0_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..0ad3844 --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv0_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu new file mode 100644 index 0000000..b3099a7 --- /dev/null +++ b/csrc/cuda_kernels/v3_cluster/instantiations/value_bf16_outidx_i64_sv0_si1_rv1_maxtopk_1024_numthreads_256_occupancy_1_b_4096_b2_4096_tma_16_cluster_8.cu @@ -0,0 +1,10 @@ +#include "../topk_select.cuh" + +namespace topk_select_bf16_cluster { + +template +void run_topk_select_kernel< + TopkSelectConfig +>(const TopkSelectArgs &args); + +} // topk_select_bf16_cluster diff --git a/csrc/cuda_kernels/v3_cluster/topk_select.cuh b/csrc/cuda_kernels/v3_cluster/topk_select.cuh index bdbfef5..456c579 100644 --- a/csrc/cuda_kernels/v3_cluster/topk_select.cuh +++ b/csrc/cuda_kernels/v3_cluster/topk_select.cuh @@ -54,7 +54,7 @@ public: static_assert(Config::target_occupancy == 1); static_assert(MAX_TOPK == 512 || MAX_TOPK == 1024); - static_assert(Config::cluster_size == 16); + static_assert(Config::cluster_size == 8 || Config::cluster_size == 16); static_assert(NUM_ELEMS_PER_ROUND * sizeof(ValueT) % 1024 == 0); static_assert((uint64_t)(MAX_VOCAB_SIZE / NUM_ELEMS_PER_SEG) * (MAX_VOCAB_SIZE / NUM_ELEMS_PER_SEG) + BF16Base::PERM_ADD_BASE <= 0xFFFFFFFFull); diff --git a/scripts/generate_instantiations.py b/scripts/generate_instantiations.py index ba00a9f..f6920c6 100644 --- a/scripts/generate_instantiations.py +++ b/scripts/generate_instantiations.py @@ -145,12 +145,13 @@ def remove_and_remake_dir(): generate_instantiations(instantiation_dir, "topk_select_fp32", configs) elif instantiation_dir == "csrc/cuda_kernels/v3_cluster/instantiations": remove_and_remake_dir() - # The cluster tier is bf16 + mk1024 only and always uses a 16-CTA cluster. + # SM90 supports the portable 8-CTA cluster; keep the SM100 16-CTA files. configs = [] for out_idx_t in ["int32_t", "int64_t"]: for si in [False, True]: for rv in [False, True]: configs.append(TopkSelectConfigs("nv_bfloat16", out_idx_t, False, si, rv, 1024, 256, 1, 4096, 4096, 16, 16)) + configs.append(TopkSelectConfigs("nv_bfloat16", out_idx_t, False, si, rv, 1024, 256, 1, 4096, 4096, 16, 8)) generate_instantiations(instantiation_dir, "topk_select_bf16_cluster", configs) else: raise ValueError(f"Invalid `instantiation_dir: {instantiation_dir}") diff --git a/setup.py b/setup.py index 9e5b4a6..b3a8545 100644 --- a/setup.py +++ b/setup.py @@ -9,6 +9,12 @@ exec(open("deep_select/__version__.py").read()) +CUDA_ARCHS = [ + arch.strip() + for arch in os.getenv("DEEP_SELECT_CUDA_ARCHS", "100a,103a").split(",") + if arch.strip() +] + CUDA_SOURCES = [ "csrc/api.cpp", @@ -97,6 +103,11 @@ ] +if "90a" in CUDA_ARCHS: + CUDA_SOURCES += sorted(str(path) for path in Path( + "csrc/cuda_kernels/v3_cluster/instantiations" + ).glob("*_tma_16_cluster_8.cu")) + def build_on_cuda_platform(): from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME @@ -113,18 +124,26 @@ def append_nvcc_threads(nvcc_extra_args): nvcc_version_number = nvcc_version.split("release ")[1].split(",")[0].strip() major, minor = map(int, nvcc_version_number.split(".")) print(f"Compiling using NVCC {major}.{minor}") - if major < 12 or (major == 12 and minor <= 8): - raise RuntimeError("sm100 compilation requires NVCC 12.9 or higher.") - - cc_flag = [ - # Currently skip build for sm80 and sm90 to speed up compilation - # "-gencode", "arch=compute_80,code=sm_80", - # "-gencode", "arch=compute_90a,code=sm_90a", - - # Compile sm100 and sm103 separately to give the compiler more room for optimization - "-gencode", "arch=compute_100a,code=sm_100a", - "-gencode", "arch=compute_103a,code=sm_103a", - ] + cuda_archs = CUDA_ARCHS + supported_cuda_archs = {"90a", "100a", "103a"} + invalid_cuda_archs = set(cuda_archs) - supported_cuda_archs + if invalid_cuda_archs: + raise ValueError( + f"Invalid DEEP_SELECT_CUDA_ARCHS entries: {sorted(invalid_cuda_archs)}. " + f"Available values are {sorted(supported_cuda_archs)}" + ) + if not cuda_archs: + raise ValueError("DEEP_SELECT_CUDA_ARCHS must contain at least one architecture") + if any(arch in {"100a", "103a"} for arch in cuda_archs): + if major < 12 or (major == 12 and minor <= 8): + raise RuntimeError("sm100/sm103 compilation requires NVCC 12.9 or higher.") + elif major < 12: + raise RuntimeError("sm90 compilation requires NVCC 12.0 or higher.") + + cc_flag = [] + for arch in cuda_archs: + cc_flag += ["-gencode", f"arch=compute_{arch},code=sm_{arch}"] + print(f"CUDA architectures: {', '.join(cuda_archs)}") this_dir = os.path.dirname(os.path.abspath(__file__)) @@ -182,7 +201,6 @@ def run(self): return (ext_modules, SpillCheckBuildExtension) -build_target_platform = kk.get_current_platform() overrided_platform = os.environ.get('DEEP_SELECT_BUILD_TARGET_PLATFORM', None) if overrided_platform is not None: overrided_platform_dict = { @@ -191,6 +209,8 @@ def run(self): if overrided_platform not in overrided_platform_dict: raise ValueError(f"Invalid `DEEP_SELECT_BUILD_TARGET_PLATFORM`: {overrided_platform}. Available values are {list(overrided_platform_dict.keys())}") build_target_platform = overrided_platform_dict[overrided_platform] +else: + build_target_platform = kk.get_current_platform() print(f"Build target: {build_target_platform}") try: diff --git a/tests/kernelkit/platform.py b/tests/kernelkit/platform.py index 60fc997..dbd223b 100644 --- a/tests/kernelkit/platform.py +++ b/tests/kernelkit/platform.py @@ -1,5 +1,6 @@ import enum import functools +import shutil import subprocess @@ -13,9 +14,14 @@ def get_current_platform() -> Platform: """ Get the current platform via `lspci` """ - output = subprocess.check_output(["lspci"], text=True) - if "3D controller: NVIDIA Corporation Device" in output: - return Platform.CUDA + if shutil.which("lspci"): + output = subprocess.check_output(["lspci"], text=True) + if "NVIDIA" in output: + return Platform.CUDA + elif shutil.which("nvidia-smi"): + output = subprocess.check_output(["nvidia-smi", "-L"], text=True) + if "GPU " in output: + return Platform.CUDA return Platform.CPU_ONLY