diff --git a/build.sh b/build.sh index 11620007e..acfcd0e24 100755 --- a/build.sh +++ b/build.sh @@ -439,6 +439,8 @@ else -e HOST_GLIBC_VER="${HOST_GLIBC_VER}" \ -e UCCL_RETAG_TO_HOST_GLIBC="${UCCL_RETAG_TO_HOST_GLIBC:-0}" \ -e UCCL_LOCAL_VERSION="${UCCL_LOCAL_VERSION:-}" \ + -e DISABLE_SM90_FEATURES="${DISABLE_SM90_FEATURES:-0}" \ + -e DISABLE_AGGRESSIVE_PTX_INSTRS="${DISABLE_AGGRESSIVE_PTX_INSTRS:-0}" \ -w /io \ "$IMAGE_NAME" \ /bin/bash /io/build_inner.sh diff --git a/ep/include/ep_configs.cuh b/ep/include/ep_configs.cuh index 7c1d9f92f..718882293 100644 --- a/ep/include/ep_configs.cuh +++ b/ep/include/ep_configs.cuh @@ -84,11 +84,19 @@ typedef uint16_t __hip_fp8x2_storage_t; #ifndef DISABLE_SM90_FEATURES #include #else -// Ampere does not support FP8 features +// Ampere does not support FP8 features, but CUDA 13+ always provides the +// header. +#if __CUDACC_VER_MAJOR__ >= 13 +#include +#else #define __NV_E4M3 0 #define __NV_E5M2 1 typedef int __nv_fp8_interpretation_t; typedef int __nv_fp8x4_e4m3; typedef uint8_t __nv_fp8_storage_t; +typedef uint16_t __nv_fp8x2_storage_t; +#define __NV_SATFINITE 0 +#define __nv_cvt_float2_to_fp8x2(a, b, c) ((uint16_t)0) +#endif #endif #endif diff --git a/ep/include/ep_launch.cuh b/ep/include/ep_launch.cuh index 564b0e8e6..b80131e56 100644 --- a/ep/include/ep_launch.cuh +++ b/ep/include/ep_launch.cuh @@ -15,30 +15,19 @@ cfg.attrs = attr; \ cfg.numAttrs = 2 #else -#define SETUP_LAUNCH_CONFIG(sms, threads, stream) \ - int __num_sms = (sms); \ - int __num_threads = (threads); \ - auto __stream = (stream) +#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \ + cudaLaunchConfig_t cfg = {(num_sms), (num_threads), 0, stream, nullptr, 0}; \ + cudaLaunchAttribute attr[1]; \ + attr[0].id = cudaLaunchAttributeCooperative; \ + attr[0].val.cooperative = 1; \ + cfg.attrs = attr; \ + cfg.numAttrs = 1 #endif #endif #ifndef LAUNCH_KERNEL -#ifndef DISABLE_SM90_FEATURES #define LAUNCH_KERNEL(config, kernel, ...) \ CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__)) -#else -#define LAUNCH_KERNEL(config, kernel, ...) \ - do { \ - kernel<<<__num_sms, __num_threads, 0, __stream>>>(__VA_ARGS__); \ - cudaError_t e = cudaGetLastError(); \ - if (e != cudaSuccess) { \ - EPException cuda_exception("CUDA", __FILE__, __LINE__, \ - cudaGetErrorString(e)); \ - fprintf(stderr, "%s\n", cuda_exception.what()); \ - throw cuda_exception; \ - } \ - } while (0) -#endif #endif #ifndef SET_SHARED_MEMORY_FOR_TMA diff --git a/ep/src/internode.cu b/ep/src/internode.cu index 38f6b160f..e6bff2a28 100644 --- a/ep/src/internode.cu +++ b/ep/src/internode.cu @@ -649,7 +649,7 @@ __global__ void __launch_bounds__( }; // TMA stuffs -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) extern __shared__ __align__(1024) uint8_t smem_tma_buffer[]; auto tma_buffer = smem_tma_buffer + target_rank * kNumTMABytesPerWarp; auto tma_mbarrier = @@ -941,7 +941,7 @@ __global__ void __launch_bounds__( } __syncwarp(); -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) // Release the transaction in the window if (is_token_in_rank_uint64 != 0) { // Acquire lock first @@ -1354,7 +1354,7 @@ __global__ void __launch_bounds__( reinterpret_cast(dst_shifted), reinterpret_cast(shifted), ld_nc_global, st_na_global); -#else +#elif !defined(DISABLE_SM90_FEATURES) if (lane_id == 0) { tma_load_1d(tma_buffer, shifted, tma_mbarrier, num_bytes_per_token, false); @@ -1373,7 +1373,7 @@ __global__ void __launch_bounds__( if ((++num_tokens_sent) == num_max_rdma_chunked_send_tokens) src_rdma_tail = i + 1; -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) tma_store_wait(); __syncwarp(); #endif @@ -1559,7 +1559,7 @@ __global__ void __launch_bounds__( reinterpret_cast(shifted + hidden_bytes), ld_nc_global, st_na_global); -#else +#elif !defined(DISABLE_SM90_FEATURES) if (lane_id == 0) { tma_load_1d(tma_buffer, shifted, tma_mbarrier, tma_load_bytes); mbarrier_arrive_and_expect_tx(tma_mbarrier, tma_load_bytes); @@ -1613,7 +1613,7 @@ __global__ void __launch_bounds__( st_na_global(recv_topk_weights + recv_idx, weight_value); } -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) // Wait TMA to be finished tma_store_wait(); #endif @@ -1796,7 +1796,7 @@ __global__ void cached_notify( } else if (sm_id == 1) { if (is_cached_dispatch) return; -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) EP_DEVICE_ASSERT(num_warps >= num_channels); #endif EP_DEVICE_ASSERT(num_rdma_ranks <= WARP_SIZE); @@ -1836,7 +1836,7 @@ __global__ void cached_notify( } else { if (is_cached_dispatch) return; -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) EP_DEVICE_ASSERT(num_warps >= num_channels); #endif EP_DEVICE_ASSERT(rdma_channel_prefix_matrix != nullptr and @@ -1859,7 +1859,7 @@ __global__ void cached_notify( EP_STATIC_ASSERT(num_bytes_per_token % 16 == 0, "num_bytes_per_token should be divisible by 16"); -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) // TMA stuffs extern __shared__ __align__(1024) uint8_t smem_tma_buffer[]; auto tma_buffer = smem_tma_buffer + warp_id * kNumTMABytesPerWarp; @@ -1897,7 +1897,7 @@ __global__ void cached_notify( auto batch_start_idx = max(token_start_idx, batch_end_idx - num_tokens_per_batch); -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) if (lane_id == 0) { tma_load_1d( tma_buffer, @@ -1925,7 +1925,7 @@ __global__ void cached_notify( } else { last_head = current_head; } -#else +#elif !defined(DISABLE_SM90_FEATURES) auto current_head = reinterpret_cast(tma_buffer) [(token_idx - batch_start_idx) * NUM_MAX_NVL_PEERS + lane_id]; if (current_head < 0) { @@ -1940,7 +1940,7 @@ __global__ void cached_notify( } } -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) tma_store_fence(); __syncwarp(); @@ -2316,7 +2316,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * WARP_SIZE, 1) channel_id, num_channels, nvl_rank) .advance_also(local_buffer_ptr); -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) // TMA stuffs extern __shared__ __align__(1024) uint8_t smem_tma_buffer[]; auto tma_buffer = @@ -2445,7 +2445,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * WARP_SIZE, 1) sizeof(SourceMeta) + lane_id * sizeof(float)), ld_nc_global(topk_weights + token_idx * num_topk + lane_id)); -#else +#elif !defined(DISABLE_SM90_FEATURES) if (lane_id == 0) { tma_store_wait(); tma_load_1d(tma_buffer, shifted_x, tma_mbarrier, hidden_bytes); @@ -2480,7 +2480,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * WARP_SIZE, 1) } // Move queue tail -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) tma_store_wait(); #endif __syncwarp(); @@ -2584,7 +2584,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * WARP_SIZE, 1) EP_STATIC_ASSERT(kNumWarpsPerForwarder == 1 or kNumRDMARanks + 2 <= 16, "Barriers are not enough"); -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) // TMA stuffs constexpr int kNumStages = 2; constexpr int kNumTMALoadBytes = sizeof(int4) * 32; @@ -2752,7 +2752,7 @@ __global__ void __launch_bounds__((kNumForwarders + 1) * WARP_SIZE, 1) nullptr, nullptr, num_max_nvl_chunked_recv_tokens_per_rdma, get_addr_fn, recv_tw_fn, nullptr, dummy_tma_phases); -#else +#elif !defined(DISABLE_SM90_FEATURES) combine_token( expected_head >= 0, expected_head, lane_id, hidden_int4, num_topk, @@ -3041,7 +3041,7 @@ void combine(cudaDataType_t type, void* combined_x, constexpr int kNumTMABytesPerSenderWarp = 16384; constexpr int kNumTMABytesPerForwarderWarp = 9248; -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) constexpr int smem_size = std::max(kNumTMABytesPerSenderWarp * NUM_MAX_NVL_PEERS, kNumTMABytesPerForwarderWarp * kNumCombineForwarderWarps); diff --git a/ep/src/internode_ll.cu b/ep/src/internode_ll.cu index deb37acb6..5d778126c 100644 --- a/ep/src/internode_ll.cu +++ b/ep/src/internode_ll.cu @@ -823,7 +823,7 @@ __global__ __launch_bounds__(1024, 1) void combine( int offset, num_tokens_to_send; unpack2(layout, num_tokens_to_send, offset); -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) // TMA stuffs constexpr int kNumTMABufferBytes = sizeof(int4) * WARP_SIZE * kNumUnrolls; constexpr int kNumStages = 3; @@ -911,6 +911,11 @@ __global__ __launch_bounds__(1024, 1) void combine( UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, cpy_dst_int4_ptr, cpy_src_int4_ptr, ld_nc_global, st_na_global); +#elif defined(DISABLE_SM90_FEATURES) + // Non-SM90 NVIDIA path: simple warp copy (no TMA available) + UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, cpy_dst_int4_ptr, + cpy_src_int4_ptr, ld_nc_global, st_na_global); + #else // Prefetch if (elect_one_sync(lane_id)) @@ -1022,7 +1027,7 @@ __global__ __launch_bounds__(1024, 1) void combine( #endif } -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) // Flush all stores tma_store_wait(); __syncwarp();