v4.3.4 update v2. (#2898)

This commit is contained in:
Junkai-Wu
2025-12-23 11:28:26 +08:00
committed by GitHub
parent 7f5fe3edf1
commit b7ecaa605d
4 changed files with 32 additions and 0 deletions

View File

@@ -471,6 +471,17 @@ make_im2col_tma_copy_desc(
tma_l2Promotion,
tma_oob_fill);
int driver_version = 0;
CUresult driver_version_result = cuDriverGetVersion(&driver_version);
assert(driver_version_result == CUDA_SUCCESS);
if (driver_version <= 13010) {
if (cute::bits_to_bytes(
cute::cosize(tensor_cwhdn.layout()) *
cute::sizeof_bits<typename EngineA::value_type>::value) < 131072) {
reinterpret_cast<uint64_t*>(&tma_desc)[1] &= ~(1llu << 21);
}
}
// The extra asserts help indicate the error's cause.
assert(encode_result != CUDA_ERROR_DEINITIALIZED);
assert(encode_result != CUDA_ERROR_NOT_INITIALIZED);

View File

@@ -1051,6 +1051,17 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The origin
smem_swizzle,
tma_l2Promotion,
tma_oobFill);
int driver_version = 0;
CUresult driver_version_result = cuDriverGetVersion(&driver_version);
assert(driver_version_result == CUDA_SUCCESS);
if (driver_version <= 13010) {
if (cute::bits_to_bytes(
cute::cosize(gtensor.layout()) *
cute::sizeof_bits<typename GEngine::value_type>::value) < 131072) {
reinterpret_cast<uint64_t*>(&tma_desc)[1] &= ~(1llu << 21);
}
}
if (result != CUDA_SUCCESS) {
std::cerr << "TMA Desc Addr: " << &tma_desc