[SM90] Add optimized SM90 support.

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

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<nv_bfloat16, OutIdxT, false, SORTED_INDEX, RETURN_VALUE, 1024, 256, 1, 4096, 4096, 16, 512, 16>>(args);
+                        if (device_prop->major == 9) {
+                            topk_select_bf16_cluster::run_topk_select_kernel<
+                                TopkSelectConfig<nv_bfloat16, OutIdxT, false, SORTED_INDEX, RETURN_VALUE, 1024, 256, 1, 4096, 4096, 16, 512, 8>>(args);
+                        } else {
+                            topk_select_bf16_cluster::run_topk_select_kernel<
+                                TopkSelectConfig<nv_bfloat16, OutIdxT, false, SORTED_INDEX, RETURN_VALUE, 1024, 256, 1, 4096, 4096, 16, 512, 16>>(args);
+                        }
                     });
                 });
             });

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<nv_bfloat16, int32_t, false, false, false, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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<nv_bfloat16, int32_t, false, false, true, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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<nv_bfloat16, int32_t, false, true, false, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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<nv_bfloat16, int32_t, false, true, true, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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<nv_bfloat16, int64_t, false, false, false, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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<nv_bfloat16, int64_t, false, false, true, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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<nv_bfloat16, int64_t, false, true, false, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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<nv_bfloat16, int64_t, false, true, true, 1024, 256, 1, 4096, 4096, 16, 512, 8>
+>(const TopkSelectArgs &args);
+
+}   // topk_select_bf16_cluster

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);

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}")

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:

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