From dae59bbd90b5568d7d630e442a88fe085564c848 Mon Sep 17 00:00:00 2001 From: Bijay Pathak Date: Sat, 23 May 2026 19:56:48 +0000 Subject: [PATCH 1/3] [EP] Add SM80/A100 build support with SM90 feature guards --- build.sh | 2 ++ ep/include/ep_configs.cuh | 3 +++ ep/src/internode.cu | 36 ++++++++++++++++++------------------ ep/src/internode_ll.cu | 6 +++--- 4 files changed, 26 insertions(+), 21 deletions(-) diff --git a/build.sh b/build.sh index 52aba50ee..d6ae51609 100755 --- a/build.sh +++ b/build.sh @@ -430,6 +430,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 e1618329a..a25c3c5e6 100644 --- a/ep/include/ep_configs.cuh +++ b/ep/include/ep_configs.cuh @@ -88,5 +88,8 @@ typedef uint16_t __hip_fp8x2_storage_t; 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 diff --git a/ep/src/internode.cu b/ep/src/internode.cu index 88a82cad7..f86ca9d4b 100644 --- a/ep/src/internode.cu +++ b/ep/src/internode.cu @@ -623,7 +623,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 = @@ -915,7 +915,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 @@ -1328,7 +1328,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); @@ -1347,7 +1347,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 @@ -1533,7 +1533,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); @@ -1587,7 +1587,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 @@ -1770,7 +1770,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); @@ -1810,7 +1810,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 @@ -1833,7 +1833,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; @@ -1871,7 +1871,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, @@ -1899,7 +1899,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) { @@ -1914,7 +1914,7 @@ __global__ void cached_notify( } } -#if defined(__NVCC__) +#if defined(__NVCC__) && !defined(DISABLE_SM90_FEATURES) tma_store_fence(); __syncwarp(); @@ -2290,7 +2290,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 = @@ -2419,7 +2419,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); @@ -2454,7 +2454,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(); @@ -2558,7 +2558,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; @@ -2726,7 +2726,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, @@ -3015,7 +3015,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 1cb908b89..dd929ebd7 100644 --- a/ep/src/internode_ll.cu +++ b/ep/src/internode_ll.cu @@ -811,7 +811,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; @@ -899,7 +899,7 @@ __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); -#else +#elif !defined(DISABLE_SM90_FEATURES) // Prefetch if (elect_one_sync(lane_id)) tma_load_and_arrive(0, cpy_src_int4_ptr, get_num_tma_bytes(0)); @@ -1010,7 +1010,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(); From 321488583a3194b87e8414364045c35bd0ce705f Mon Sep 17 00:00:00 2001 From: Manoj Gopalakrishnan Date: Thu, 28 May 2026 05:33:28 +0000 Subject: [PATCH 2/3] Add A100 (SM 8.0) support for low-latency internode kernels - Add software grid barrier (cuda_grid_barrier) for non-SM90 path using cooperative launch with atomicAdd/atomicExch reset pattern - Use cudaLaunchKernelEx with cooperative attribute for non-SM90 GPUs - Guard TMA intrinsics with DISABLE_SM90_FEATURES to avoid ISA errors - Add warp-copy fallback for TMA store paths on pre-SM90 - Handle CUDA 13+ cuda_fp8.h include compatibility in ep_configs.cuh Signed-off-by: Manoj Gopalakrishnan --- ep/include/ep_configs.cuh | 7 ++++++- ep/include/ep_launch.cuh | 25 +++++++------------------ ep/include/ep_utils.cuh | 36 ++++++++++++++++++++++++++++++++++++ ep/src/internode_ll.cu | 13 ++++++++++++- 4 files changed, 61 insertions(+), 20 deletions(-) diff --git a/ep/include/ep_configs.cuh b/ep/include/ep_configs.cuh index a25c3c5e6..fb4c13a09 100644 --- a/ep/include/ep_configs.cuh +++ b/ep/include/ep_configs.cuh @@ -82,7 +82,11 @@ 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; @@ -93,3 +97,4 @@ typedef uint16_t __nv_fp8x2_storage_t; #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/include/ep_utils.cuh b/ep/include/ep_utils.cuh index bb7fd854b..8854d8d9a 100644 --- a/ep/include/ep_utils.cuh +++ b/ep/include/ep_utils.cuh @@ -153,6 +153,42 @@ __device__ __forceinline__ void st_release_sys_global(dtype_t const* ptr, } // namespace amd #endif +// Software grid barrier for CUDA (non-SM90) path. +// Mirrors amd::grid_sync_then_zero: last arriver clears the counter so the +// barrier is reusable across iterations without an extra clean kernel. +#if !defined(__HIP_PLATFORM_AMD__) && !defined(__HIPCC__) +__device__ __forceinline__ void cuda_grid_barrier(int* bar_ptr, + int num_blocks) { + // Ensure all threads in this block have completed their prior memory + // operations before thread 0 signals arrival on behalf of the block. + __syncthreads(); + if (threadIdx.x == 0) { + // Release fence: ensure all prior writes (e.g. packed_recv_count clean) + // are visible to other blocks before we signal arrival. + __threadfence(); + + unsigned int val = atomicAdd(reinterpret_cast(bar_ptr), 1u); + if (val == static_cast(num_blocks - 1)) { + // Last arriver: all blocks have fenced their writes. Reset counter. + // Use release semantics so the store of 0 is visible after all prior + // writes from this (last) block. + atomicExch(reinterpret_cast(bar_ptr), 0u); + } else { + // Spin until last arriver resets counter to 0. + while (*(unsigned int volatile*)bar_ptr != 0u) + ; + } + + // Acquire fence: ensure we see all writes from all blocks that arrived + // before us (transitively through the atomic total order). + __threadfence(); + } + // Broadcast barrier completion from thread 0 to all threads in the block + // before any thread proceeds past the barrier. + __syncthreads(); +} +#endif + __forceinline__ __device__ int get_lane_id() { int lane_id; #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) diff --git a/ep/src/internode_ll.cu b/ep/src/internode_ll.cu index dd929ebd7..9bb1a12d6 100644 --- a/ep/src/internode_ll.cu +++ b/ep/src/internode_ll.cu @@ -340,6 +340,8 @@ __global__ __launch_bounds__(1024, 1) void dispatch( // Reset counter after sync so send-only launches (return_recv_hook) do not // leave a stale value that deadlocks the next dispatch. amd::grid_sync_then_zero(grid_sync_barrier_ptr, num_sms); +#elif defined(DISABLE_SM90_FEATURES) + cuda_grid_barrier(grid_sync_barrier_ptr, num_sms); #else cg::this_grid().sync(); #endif @@ -464,6 +466,8 @@ LOW_LATENCY_DISPATCH_RECV: if (phases & LOW_LATENCY_SEND_PHASE) #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) amd::grid_sync_then_zero(grid_sync_barrier_ptr, num_sms); +#elif defined(DISABLE_SM90_FEATURES) + cuda_grid_barrier(grid_sync_barrier_ptr, num_sms); #else cg::this_grid().sync(); #endif @@ -899,7 +903,12 @@ __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) +#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)) tma_load_and_arrive(0, cpy_src_int4_ptr, get_num_tma_bytes(0)); @@ -1144,6 +1153,8 @@ LOW_LATENCY_COMBINE_RECV: } #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) amd::grid_sync_then_zero(grid_sync_barrier_ptr, num_sms); +#elif defined(DISABLE_SM90_FEATURES) + cuda_grid_barrier(grid_sync_barrier_ptr, num_sms); #else cg::this_grid().sync(); #endif From d2519ffc1ce3c8b39eb6b5668a7429d5a9641a00 Mon Sep 17 00:00:00 2001 From: Bijay K Pathak Date: Tue, 7 Jul 2026 20:20:01 -0700 Subject: [PATCH 3/3] [EP] Use cg::this_grid().sync() for A100 low-latency kernels Removes the custom cuda_grid_barrier software grid barrier for the DISABLE_SM90_FEATURES (A100) path in favor of cg::this_grid().sync(), which is already safe to use since the launch config for that path sets cudaLaunchAttributeCooperative. Co-authored-by: Manoj Gopalakrishnan --- ep/include/ep_utils.cuh | 36 ------------------------------------ ep/src/internode_ll.cu | 6 ------ 2 files changed, 42 deletions(-) diff --git a/ep/include/ep_utils.cuh b/ep/include/ep_utils.cuh index 8854d8d9a..bb7fd854b 100644 --- a/ep/include/ep_utils.cuh +++ b/ep/include/ep_utils.cuh @@ -153,42 +153,6 @@ __device__ __forceinline__ void st_release_sys_global(dtype_t const* ptr, } // namespace amd #endif -// Software grid barrier for CUDA (non-SM90) path. -// Mirrors amd::grid_sync_then_zero: last arriver clears the counter so the -// barrier is reusable across iterations without an extra clean kernel. -#if !defined(__HIP_PLATFORM_AMD__) && !defined(__HIPCC__) -__device__ __forceinline__ void cuda_grid_barrier(int* bar_ptr, - int num_blocks) { - // Ensure all threads in this block have completed their prior memory - // operations before thread 0 signals arrival on behalf of the block. - __syncthreads(); - if (threadIdx.x == 0) { - // Release fence: ensure all prior writes (e.g. packed_recv_count clean) - // are visible to other blocks before we signal arrival. - __threadfence(); - - unsigned int val = atomicAdd(reinterpret_cast(bar_ptr), 1u); - if (val == static_cast(num_blocks - 1)) { - // Last arriver: all blocks have fenced their writes. Reset counter. - // Use release semantics so the store of 0 is visible after all prior - // writes from this (last) block. - atomicExch(reinterpret_cast(bar_ptr), 0u); - } else { - // Spin until last arriver resets counter to 0. - while (*(unsigned int volatile*)bar_ptr != 0u) - ; - } - - // Acquire fence: ensure we see all writes from all blocks that arrived - // before us (transitively through the atomic total order). - __threadfence(); - } - // Broadcast barrier completion from thread 0 to all threads in the block - // before any thread proceeds past the barrier. - __syncthreads(); -} -#endif - __forceinline__ __device__ int get_lane_id() { int lane_id; #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) diff --git a/ep/src/internode_ll.cu b/ep/src/internode_ll.cu index f600f23a9..5d778126c 100644 --- a/ep/src/internode_ll.cu +++ b/ep/src/internode_ll.cu @@ -352,8 +352,6 @@ __global__ __launch_bounds__(1024, 1) void dispatch( // Reset counter after sync so send-only launches (return_recv_hook) do not // leave a stale value that deadlocks the next dispatch. amd::grid_sync_then_zero(grid_sync_barrier_ptr, num_sms); -#elif defined(DISABLE_SM90_FEATURES) - cuda_grid_barrier(grid_sync_barrier_ptr, num_sms); #else cg::this_grid().sync(); #endif @@ -478,8 +476,6 @@ LOW_LATENCY_DISPATCH_RECV: if (phases & LOW_LATENCY_SEND_PHASE) #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) amd::grid_sync_then_zero(grid_sync_barrier_ptr, num_sms); -#elif defined(DISABLE_SM90_FEATURES) - cuda_grid_barrier(grid_sync_barrier_ptr, num_sms); #else cg::this_grid().sync(); #endif @@ -1165,8 +1161,6 @@ LOW_LATENCY_COMBINE_RECV: } #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) amd::grid_sync_then_zero(grid_sync_barrier_ptr, num_sms); -#elif defined(DISABLE_SM90_FEATURES) - cuda_grid_barrier(grid_sync_barrier_ptr, num_sms); #else cg::this_grid().sync(); #endif