Skip to content

Commit 6ed481e

Browse files
CUDA: Check PTX version on host side to guard PDL dispatch (ggml-org#23530)
* CUDA: Check PTX version on host side to guard PDL dispatch Checking on `__CUDA_ARCH_LIST__` alone is insufficient for JIT, as this variable doesn't differentiate between compiling for say sm_90, sm_90a or sm_90f (so forward-jittable PTX vs. arch/family-specific PTX). Thus, one can have a bug when compiling with `DCMAKE_CUDA_ARCHITECTURES="89;90a"`, where current code would wrongly dispatch to PDL on sm_90/sm_120 in forward-JIT mode. This PR fixes this issue by checking `cudaFuncAttributes::ptxVersion` of the incoming kernel at runtime. A check on ptxVersion alone is sufficient, as device-codes will always be >= ptxVersion (and any violation of this would be a severe bug in CUDA/nvcc), see: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/#gpu-code-code-code * Implement MurmurHash3 mixer for better hash distribution Magic constants were taken from boost: https://github.com/boostorg/container_hash/blob/2698b43803c012601e6bb1a6116e83767b97986c/include/boost/container_hash/detail/hash_mix.hpp#L19-L65 * Update ggml/src/ggml-cuda/common.cuh Co-authored-by: Johannes Gäßler <johannesg@5d6.de> * Address review comments, make seed non-zero * Apply code-formatting * Replace std::size_t -> size_t for consistency --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
1 parent cb47092 commit 6ed481e

1 file changed

Lines changed: 58 additions & 2 deletions

File tree

ggml/src/ggml-cuda/common.cuh

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#include <cstdint>
88
#include <cstdlib>
99
#include <memory>
10+
#include <mutex>
1011

1112
#if defined(GGML_USE_HIP)
1213
#define GGML_COMMON_DECL_HIP
@@ -1552,6 +1553,62 @@ struct ggml_cuda_pdl_config {
15521553
ggml_cuda_pdl_config& operator=(ggml_cuda_pdl_config&&) = delete;
15531554

15541555
};
1556+
1557+
static bool ggml_cuda_kernel_can_use_pdl(const void * kernel) {
1558+
const int device = ggml_cuda_get_device();
1559+
1560+
struct cache_key {
1561+
int device;
1562+
const void * kernel;
1563+
1564+
bool operator==(const cache_key & other) const { return device == other.device && kernel == other.kernel; }
1565+
};
1566+
1567+
struct cache_key_hash {
1568+
// MurmurHash3 mixing function for better hash distribution (vs. just std::hash which in some implementations simply returns the identity)
1569+
static size_t hash_mix(size_t x) {
1570+
std::uint64_t y = x;
1571+
const std::uint64_t m = 0xe9846af9b1a615d;
1572+
1573+
y ^= y >> 32;
1574+
y *= m;
1575+
y ^= y >> 32;
1576+
y *= m;
1577+
y ^= y >> 28;
1578+
1579+
return static_cast<size_t>(y);
1580+
}
1581+
1582+
size_t operator()(const cache_key & key) const {
1583+
// Use a nonzero seed to avoid mapping all-zero keys to zero
1584+
size_t h = 42;
1585+
h = hash_mix(h + key.device);
1586+
h = hash_mix(h + reinterpret_cast<size_t>(key.kernel));
1587+
return h;
1588+
}
1589+
};
1590+
1591+
static std::mutex cache_mutex;
1592+
static std::unordered_map<cache_key, bool, cache_key_hash> cache;
1593+
1594+
const cache_key key = { device, kernel };
1595+
std::lock_guard<std::mutex> lock(cache_mutex);
1596+
const auto it = cache.find(key);
1597+
if (it != cache.end()) {
1598+
return it->second;
1599+
}
1600+
1601+
cudaFuncAttributes attr = {};
1602+
CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel));
1603+
1604+
// PDL device-side primitives are emitted only for PTX versions >= 90.
1605+
// We have to guard on a loaded kernel's PTX version so a kernel forward-JIT'ed
1606+
// from pre-Hopper PTX to a Hopper-or-newer GPU does not opt into PDL.
1607+
const bool can_use_pdl = attr.ptxVersion >= 90;
1608+
cache.emplace(key, can_use_pdl);
1609+
return can_use_pdl;
1610+
}
1611+
15551612
#endif //defined(GGML_CUDA_USE_PDL)
15561613

15571614

@@ -1564,8 +1621,7 @@ static __inline__ void ggml_cuda_kernel_launch(Kernel kernel, const ggml_cuda_ke
15641621
return env == nullptr || std::atoi(env) != 0;
15651622
}();
15661623

1567-
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
1568-
if (env_pdl_enabled && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_HOPPER) {
1624+
if (env_pdl_enabled && ggml_cuda_kernel_can_use_pdl(reinterpret_cast<const void *>(kernel))) {
15691625
auto pdl_cfg = ggml_cuda_pdl_config(launch_params);
15701626

15711627
CUDA_CHECK(cudaLaunchKernelEx(&pdl_cfg.cfg, kernel, std::forward<Args>(args)... ));

0 commit comments

Comments
 (0)