Reduce redundant Top-K host setup

csrc/3rdparty/kerutils/include/kerutils/host/host.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstdio>
+#include <array>
 #include <exception>
 #include <memory>
 #include <string>
@@ -97,81 +98,70 @@ inline __host__ __device__ constexpr T find_next_power_of_2(const T& x) {
     return find_next_power_of_2<T, LOWER_BOUND*2>(x);
 }
 
-// A wrapper for make_tensor_map
-static inline CUtensorMap make_tensor_map(
-    const std::vector<uint64_t> &size,
-    const std::vector<uint64_t> &strides,   // PAY ATTENTION: In BYTES
-    const std::vector<uint32_t> &box_size,
-    void* global_ptr,
-    CUtensorMapDataType data_type,
-    CUtensorMapSwizzle swizzle_mode,
-    CUtensorMapL2promotion l2_promotion,
-    CUtensorMapInterleave interleave_mode = CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-    CUtensorMapFloatOOBfill oob_fill = CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
-    const std::vector<uint32_t> &element_strides_ = {}
-) {
-    int dim = size.size();
-    KU_ASSERT(dim >= 1);
-
-    std::vector<uint32_t> element_strides;
-    if (element_strides_.empty()) {
-        for (int i = 0; i < dim; ++i)
-            element_strides.push_back(1);
-    } else {
-        element_strides = element_strides_;
-    }
-    KU_ASSERT(strides.size() == (uint32_t)dim-1 && box_size.size() == (uint32_t)dim && element_strides.size() == (uint32_t)dim);
-
-    auto call_cuTensorMapEncodeTiled = [&]<typename... Args>(Args... args) {
+// Driver symbols are independent of tensor storage and the current device.
+// Inline linkage shares this cache across translation units. A failed lookup
+// throws before initialization completes, so a later call retries it.
+inline PFN_cuTensorMapEncodeTiled_v12000 get_tensor_map_encoder() {
+    static const auto encoder = [] {
         cudaDriverEntryPointQueryResult cuda_status;
         void* pfn = nullptr;
-#if (__CUDACC_VER_MAJOR__ > 12)
+#if CUDA_VERSION >= 13000
         KU_CUDA_CHECK(cudaGetDriverEntryPointByVersion(
-            "cuTensorMapEncodeTiled",
-            &pfn, 12000,
-            cudaEnableDefault,
-            &cuda_status));
+            "cuTensorMapEncodeTiled", &pfn, 12000, cudaEnableDefault, &cuda_status));
 #else
         KU_CUDA_CHECK(cudaGetDriverEntryPoint(
-            "cuTensorMapEncodeTiled",
-            &pfn,
-            cudaEnableDefault,
-            &cuda_status));
+            "cuTensorMapEncodeTiled", &pfn, cudaEnableDefault, &cuda_status));
 #endif
-        if (cuda_status != cudaDriverEntryPointSuccess) {
-            KU_ASSERT(false, "Failed to load `cuTensorMapEncodeTiled`. cuda_status = %d", cuda_status);
-        }
-        return reinterpret_cast<decltype(&cuTensorMapEncodeTiled)>(pfn)(args...); \
-    };
+        KU_ASSERT(cuda_status == cudaDriverEntryPointSuccess && pfn != nullptr,
+                  "Failed to load `cuTensorMapEncodeTiled`. cuda_status = %d", cuda_status);
+        return reinterpret_cast<PFN_cuTensorMapEncodeTiled_v12000>(pfn);
+    }();
+    return encoder;
+}
+
+namespace detail {
 
+inline CUtensorMap encode_tensor_map(
+    int dim,
+    const uint64_t* size,
+    const uint64_t* strides,
+    const uint32_t* box_size,
+    void* global_ptr,
+    CUtensorMapDataType data_type,
+    CUtensorMapSwizzle swizzle_mode,
+    CUtensorMapL2promotion l2_promotion,
+    CUtensorMapInterleave interleave_mode,
+    CUtensorMapFloatOOBfill oob_fill,
+    const uint32_t* element_strides
+) {
     CUtensorMap result;
-    CUresult ret_code = call_cuTensorMapEncodeTiled(
+    CUresult ret_code = get_tensor_map_encoder()(
         &result,
         data_type,
         dim,
         global_ptr,
-        size.data(),
-        strides.data(),
-        box_size.data(),
-        element_strides.data(),
+        size,
+        strides,
+        box_size,
+        element_strides,
         interleave_mode,
         swizzle_mode,
         l2_promotion,
         oob_fill
     );
     if (ret_code != CUresult::CUDA_SUCCESS) {
-        auto print_vector = [&](auto t, const char* fmt, const char end='\n') {
-            for (auto elem : t) {
-                printf(fmt, elem);
+        auto print_vector = [&](auto t, int count, const char* fmt, const char end='\n') {
+            for (int i = 0; i < count; ++i) {
+                printf(fmt, t[i]);
             }
             printf("%c", end);
         };
         fprintf(stderr, "Failed to create tensormap\n");
         fprintf(stderr, "Dim: %d\n", dim);
-        printf("size: "); print_vector(size, "%lu ");
-        printf("strides: "); print_vector(strides, "%lu ");
-        printf("box_size: "); print_vector(box_size, "%u ");
-        printf("element_strides: "); print_vector(element_strides, "%u ");
+        printf("size: "); print_vector(size, dim, "%lu ");
+        printf("strides: "); print_vector(strides, dim-1, "%lu ");
+        printf("box_size: "); print_vector(box_size, dim, "%u ");
+        printf("element_strides: "); print_vector(element_strides, dim, "%u ");
         printf("global ptr: 0x%lx\n", (int64_t)global_ptr);
         printf("data_type: %d\n", (int)data_type);
         printf("swizzle_mode: %d\n", (int)swizzle_mode);
@@ -183,6 +173,54 @@ static inline CUtensorMap make_tensor_map(
     return result;
 }
 
+} // namespace detail
+
+// Preserve the dynamically ranked interface for existing callers.
+static inline CUtensorMap make_tensor_map(
+    const std::vector<uint64_t>& size,
+    const std::vector<uint64_t>& strides, // In bytes
+    const std::vector<uint32_t>& box_size,
+    void* global_ptr,
+    CUtensorMapDataType data_type,
+    CUtensorMapSwizzle swizzle_mode,
+    CUtensorMapL2promotion l2_promotion,
+    CUtensorMapInterleave interleave_mode = CU_TENSOR_MAP_INTERLEAVE_NONE,
+    CUtensorMapFloatOOBfill oob_fill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
+    const std::vector<uint32_t>& element_strides_ = {}
+) {
+    int dim = size.size();
+    KU_ASSERT(dim >= 1);
+    std::vector<uint32_t> element_strides = element_strides_.empty()
+        ? std::vector<uint32_t>(dim, 1) : element_strides_;
+    KU_ASSERT(strides.size() == (uint32_t)dim-1 && box_size.size() == (uint32_t)dim && element_strides.size() == (uint32_t)dim);
+    return detail::encode_tensor_map(dim, size.data(), strides.data(), box_size.data(),
+        global_ptr, data_type, swizzle_mode, l2_promotion, interleave_mode, oob_fill, element_strides.data());
+}
+
+// Fixed rank keeps metadata on the stack; the descriptor is still rebuilt for
+// every call's pointer, dimensions, and strides.
+template<size_t Rank>
+inline CUtensorMap make_tensor_map(
+    const std::array<uint64_t, Rank>& size,
+    const std::array<uint64_t, Rank-1>& strides, // In bytes
+    const std::array<uint32_t, Rank>& box_size,
+    void* global_ptr,
+    CUtensorMapDataType data_type,
+    CUtensorMapSwizzle swizzle_mode,
+    CUtensorMapL2promotion l2_promotion,
+    CUtensorMapInterleave interleave_mode = CU_TENSOR_MAP_INTERLEAVE_NONE,
+    CUtensorMapFloatOOBfill oob_fill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
+    const std::array<uint32_t, Rank>& element_strides = [] {
+        std::array<uint32_t, Rank> ones;
+        ones.fill(1);
+        return ones;
+    }()
+) {
+    static_assert(Rank >= 1 && Rank <= 5);
+    return detail::encode_tensor_map(Rank, size.data(), strides.data(), box_size.data(),
+        global_ptr, data_type, swizzle_mode, l2_promotion, interleave_mode, oob_fill, element_strides.data());
+}
+
 // Given strides (in number of elements), this function converts their datatype in uint64_t and then multiplies by elem_size
 template<typename T>
 static inline std::vector<uint64_t> make_stride_helper(const std::vector<T> &strides_in_elems, size_t elem_size) {
@@ -279,4 +317,4 @@ void launch_kernel(const KernelLaunchConfig &cfg, KernelFunc kernel, Args&&... a
 
 #endif  // KERUTILS_IS_BUILD_ON_CUDA
 
-}   // namespace kerutils
+}   // namespace kerutils

csrc/cuda_kernels/common_parts.cuh

@@ -411,7 +411,7 @@ public:
         constexpr CUtensorMapDataType dtype = std::is_same_v<ValueT, nv_bfloat16>
             ? CU_TENSOR_MAP_DATA_TYPE_BFLOAT16
             : CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
-        return ku::make_tensor_map(
+        return ku::make_tensor_map<3>(
             // Split args.vocab_size into two dims because
             //   - TMA requires innermost box dim <= swizzle size
             //   - To avoid OOB as we have `INPUT_STRIDE_ALIGNMENT_REQUIREMENT`
@@ -420,7 +420,7 @@ public:
                 ku::ceil_div((uint64_t)args.vocab_size, (uint64_t)NUM_ELEMS_PER_TMA_ROW),
                 args.batch_size
             },
-            ku::make_stride_helper<uint64_t>({NUM_ELEMS_PER_TMA_ROW, args.stride_input_batch}, sizeof(ValueT)),
+            {uint64_t(NUM_ELEMS_PER_TMA_ROW) * sizeof(ValueT), uint64_t(args.stride_input_batch) * sizeof(ValueT)},
             {
                 NUM_ELEMS_PER_TMA_ROW,
                 NUM_TMA_ROWS_PER_SEG,

csrc/cuda_kernels/v3/topk_select.cuh

@@ -130,7 +130,6 @@ void run_topk_select_kernel(const TopkSelectArgs &args) {
     auto kernel = topk_kernel<Kernel>;
     constexpr size_t smem_size = sizeof(typename Kernel::SharedMemoryPlan);
     KU_ASSERT(smem_size * Kernel::TARGET_OCCUPANCY <= args.shared_memory_size_per_sm);
-    KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
 
     KU_ASSERT(args.stride_input_batch % 8 == 0, "stride_input_batch must be 16B-aligned");
     typename Kernel::TmaParams tma_params = {Kernel::make_topk_tensor_map(args)};

csrc/cuda_kernels/v3_cluster/topk_select.cuh

@@ -286,10 +286,6 @@ void run_topk_select_kernel(const TopkSelectArgs &args) {
     auto kernel = topk_kernel<Kernel>;
     constexpr size_t smem_size = sizeof(typename Kernel::SharedMemoryPlanBF16Cluster);
     KU_ASSERT(smem_size * Kernel::TARGET_OCCUPANCY <= args.shared_memory_size_per_sm);
-    KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
-    if constexpr (Kernel::CLUSTER_SIZE > 8) {
-        KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1));
-    }
 
     KU_ASSERT(args.stride_input_batch % 8 == 0, "stride_input_batch must be 16B-aligned");
     typename Kernel::TmaParams tma_params = {Kernel::make_topk_tensor_map(args)};

csrc/cuda_kernels/v3_fp32/topk_select.cuh

@@ -589,7 +589,6 @@ void run_topk_select_kernel(const TopkSelectArgs &args) {
     auto kernel = topk_kernel<Kernel>;
     constexpr size_t smem_size = sizeof(typename Kernel::SharedMemoryPlanFP32);
     KU_ASSERT(smem_size * Kernel::TARGET_OCCUPANCY <= args.shared_memory_size_per_sm);
-    KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
 
     KU_ASSERT(args.stride_input_batch % 4 == 0, "stride_input_batch must be 16B-aligned");
     typename Kernel::TmaParams tma_params = {Kernel::make_topk_tensor_map(args)};

tests/host/README.md

@@ -0,0 +1,17 @@
+# Host setup regression tests
+
+These tests use real CUDA headers and a mocked driver resolver and encoder. They require CUDA 13's compiler, but do not initialize CUDA or require a GPU. The mock records encoder arguments rather than validating device-specific tensor-map constraints.
+
+From the repository root, build into a directory outside the checkout:
+
+```sh
+nvcc -std=c++20 -O2 --cudart shared -I csrc/3rdparty/kerutils/include -Xcompiler -pthread tests/host/tensor_map.cu tests/host/tensor_map_peer.cu tests/host/allocations.cpp -o /path/to/artifacts/tensor-map-test
+/path/to/artifacts/tensor-map-test
+/path/to/artifacts/tensor-map-test retry
+```
+
+The normal process tests concurrent cold initialization, one lookup shared across two translation units, current per-call tensor arguments, retry after an encoder failure, and vector/array equivalence for ranks one through five with default and explicit element strides. The separate `retry` process additionally tests runtime lookup failure, missing symbol, and a null function pointer before successful initialization. Assertion diagnostics from these intentional failures are expected.
+
+For the baseline regression, compile the same sources with `-DTEST_BASELINE` and point the include path at baseline kerutils headers. Run without `retry`: the single-lookup assertion must fail after sixteen concurrent calls. This mode excludes the new array overload and does not attempt the baseline's unsafe null-pointer call.
+
+Run either binary with `benchmark` to measure rank-three metadata preparation through its actual headers. It warms up 1,000 calls, then reports seven rounds of 200,000 calls, allocations, and resolver calls. The baseline uses the original vector/stride-helper path; the candidate uses fixed arrays. Both use the same recording mock encoder. This isolates host preparation and does not measure driver encoding, kernel launch, or end-to-end Top-K latency. Alternate baseline/candidate runs to check timing noise. The candidate additionally asserts zero steady-state allocations and lookups.

tests/host/allocations.cpp

@@ -0,0 +1,13 @@
+#include <cstdlib>
+#include <new>
+
+thread_local bool count_allocations = false;
+thread_local size_t allocations = 0;
+
+void* operator new(size_t size) {
+    if (count_allocations) ++allocations;
+    if (auto* ptr = std::malloc(size ? size : 1)) return ptr;
+    throw std::bad_alloc();
+}
+void operator delete(void* ptr) noexcept { std::free(ptr); }
+void operator delete(void* ptr, size_t) noexcept { std::free(ptr); }