diff --git a/ep/README.md b/ep/README.md index 1ddd5a08c..9bc07983e 100644 --- a/ep/README.md +++ b/ep/README.md @@ -138,6 +138,7 @@ Notes: | UCCL_IB_TC | Traffic class in RDMA network | 104/0 (IB/EFA) | | UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC | Use relaxed atomics with manual `s_waitcnt vmcnt(0)` fences instead of acquire/release semantics. Required on AMD CDNA so the combine receiver actually sees the producer's tail-pointer updates over XGMI; without it the kernel deadlocks at scale. | 1 on AMD, 0 on CUDA | | UCCL_RDMA_ADAPTIVE_SLEEP | Enable adaptive sleeping on proxy threads, by putting the proxy threads into a sleeping state if there have been no new work requests / RDMA completion events after 120s. | null | +| UCCL_UDP_SPORT_BASE | Base UDP source port for RoCEv2 QPs on mlx5 devices. Valid range: 1-65535. | 0 (driver decides) | ## Results diff --git a/ep/include/proxy_ctx.hpp b/ep/include/proxy_ctx.hpp index 85f12d4ef..916c2f5a3 100644 --- a/ep/include/proxy_ctx.hpp +++ b/ep/include/proxy_ctx.hpp @@ -74,6 +74,7 @@ struct ProxyCtx { // std::vector extra_cqs; ibv_qp* ack_qp = nullptr; ibv_qp* recv_ack_qp = nullptr; + int udp_sport_base = 0; // EFA shared-QP model: when true, the QP and ack_recv_buf/mr fields above // alias another ProxyCtx (typically Proxy::ctx_) and must not be destroyed. bool qps_are_shared = false; diff --git a/ep/include/rdma_util.hpp b/ep/include/rdma_util.hpp index dd0383fa7..9913c85f4 100644 --- a/ep/include/rdma_util.hpp +++ b/ep/include/rdma_util.hpp @@ -13,12 +13,13 @@ #include #include #include +#include #include #include #include #include -void fill_local_gid(ProxyCtx& S, RDMAConnectionInfo* local_info) { +static inline void fill_local_gid(ProxyCtx& S, RDMAConnectionInfo* local_info) { if (!S.context) { fprintf(stderr, "Error: context not initialized when filling GID\n"); exit(1); @@ -103,6 +104,48 @@ static inline int get_roce_version_from_env() { return roce_version; } +static inline int get_udp_sport_base_from_env() { + static int udp_sport_base = 0; + char const* env = getenv("UCCL_UDP_SPORT_BASE"); + if (env) { + int val = atoi(env); + if (val < 1 || val > 65535) { + fprintf(stderr, + "UCCL_UDP_SPORT_BASE=%d out of the valid range [1, 65535], " + "ignoring\n", + val); + } else { + udp_sport_base = val; + } + } + return udp_sport_base; +} + +typedef int (*mlx5dv_modify_qp_udp_sport_fn)(struct ibv_qp*, uint16_t); + +static inline mlx5dv_modify_qp_udp_sport_fn load_mlx5dv_udp_sport_fn() { + static mlx5dv_modify_qp_udp_sport_fn fn = + []() -> mlx5dv_modify_qp_udp_sport_fn { + void* handle = dlopen("libmlx5.so", RTLD_NOW); + if (!handle) handle = dlopen("libmlx5.so.1", RTLD_NOW); + if (!handle) { + fprintf(stderr, "Failed to open libmlx5.so[.1]: %s\n", dlerror()); + return nullptr; + } + auto f = (mlx5dv_modify_qp_udp_sport_fn)dlsym(handle, + "mlx5dv_modify_qp_udp_sport"); + if (!f) { + fprintf(stderr, + "dlsym failed on mlx5dv_modify_qp_udp_sport: %s, " + "skipping UDP sport configuration\n", + dlerror()); + dlclose(handle); + } + return f; + }(); + return fn; +} + static sa_family_t envIbAddrFamily(void) { sa_family_t family = AF_INET; char const* env = getenv("UCCL_IB_ADDR_FAMILY"); diff --git a/ep/src/proxy.cpp b/ep/src/proxy.cpp index ff630403f..cbc3f6777 100644 --- a/ep/src/proxy.cpp +++ b/ep/src/proxy.cpp @@ -3,6 +3,7 @@ #include "d2h_queue_host.hpp" #include "ep_util.hpp" #include "rdma.hpp" +#include "rdma_util.hpp" #include "util/util.h" #include // for htonl, ntohl #include @@ -162,6 +163,7 @@ void Proxy::set_peers_meta(std::vector const& peers) { ctxs_for_all_ranks_.resize(peers.size()); for (size_t i = 0; i < peers.size(); ++i) { ctxs_for_all_ranks_[i] = std::make_unique(); + ctxs_for_all_ranks_[i]->udp_sport_base = get_udp_sport_base_from_env(); } } diff --git a/ep/src/rdma.cpp b/ep/src/rdma.cpp index 9d84df5db..e5f304620 100644 --- a/ep/src/rdma.cpp +++ b/ep/src/rdma.cpp @@ -7,7 +7,6 @@ #ifdef USE_DMABUF #include #include -#include #endif #include #include @@ -1338,6 +1337,26 @@ void modify_qp_to_rts(ProxyCtx& S, RDMAConnectionInfo* local_info) { } printf("ACK-QP modified to RTS state\n"); + + if (S.udp_sport_base != 0) { + auto fn = load_mlx5dv_udp_sport_fn(); + if (fn) { + auto set_sport = [&](ibv_qp* qp, int sport, char const* name) { + if (fn(qp, (uint16_t)sport)) + fprintf(stderr, "Warning: failed to set UDP sport %d on %s QP\n", + sport, name); + }; + set_sport(S.qp, S.udp_sport_base, "main"); + for (size_t r = 0; r < S.data_qps_by_channel.size(); ++r) + set_sport(S.data_qps_by_channel[r], S.udp_sport_base + r + 1, "data"); + set_sport(S.ack_qp, S.udp_sport_base + S.data_qps_by_channel.size() + 1, + "ACK"); + set_sport(S.recv_ack_qp, + S.udp_sport_base + S.data_qps_by_channel.size() + 2, + "recv ACK"); + printf("UDP source ports set starting from %d\n", S.udp_sport_base); + } + } } void post_receive_buffer_for_imm_on_qp(ProxyCtx& S, ibv_qp* qp) {