Implementing Real-Time Tensor-Parallelism Communication Over RDMA (RoCEv2) for Distributed LLM Inference — code snippets

snippet-1.sh

#!/usr/bin/env bash
# Production initialization script for configuring RoCEv2 lossless network interfaces (Mellanox ConnectX)
set -euo pipefail

INTERFACE="eth0"
# RoCEv2 uses DSCP value 26 (CS3) or 46 (EF) commonly. We will map traffic class to DSCP 26.
PRIORITY=3
DSCP=26

echo "=== Configuring sysctl parameters for high-performance RDMA ==="
sysctl -w net.core.rmem_max=134217728
sysctl -w net.core.wmem_max=134217728
sysctl -w net.core.rmem_default=67108864
sysctl -w net.core.wmem_default=67108864
sysctl -w net.ipv4.tcp_rmem="4096 87380 134217728"
sysctl -w net.ipv4.tcp_wmem="4096 65536 134217728"
# Increase maximum pending packet queue
sysctl -w net.core.netdev_max_backlog=250000

echo "=== Enabling Priority Flow Control (PFC) on priority ${PRIORITY} ==="
# Enable PFC on priority 3, disable on others
mlnx_qos -i "$INTERFACE" --pfc "${PRIORITY}"

echo "=== Mapping DSCP ${DSCP} to Priority ${PRIORITY} ==="
mlnx_qos -i "$INTERFACE" --dscp2prio set,${DSCP},${PRIORITY}

echo "=== Configuring ECN (Explicit Congestion Notification) for DCQCN ==="
# Enable ECN marking for priority 3 on Mellanox card
# ECN parameters reside in /sys/kernel/debug/mlx5 or are managed via mstflint/mcxadm
# Here we enable ECN for priority 3 (RoCEv2)
DEV_DIR="/sys/class/net/${INTERFACE}/device/infiniband"
IB_DEV=$(ls "$DEV_DIR" | head -n 1)

# Enable DCQCN (congestion control engine) for priority 3
echo 1 > "/sys/kernel/debug/mlx5/$(lspci -d 15b3: | head -n 1 | awk '{print $1}')/cc_params/ecn_enable" || true
echo "RoCEv2 PFC and ECN configured on interface ${INTERFACE} (IB device: ${IB_DEV}) with DSCP ${DSCP}"

snippet-2.txt

#include <iostream>
#include <vector>
#include <cuda_runtime.h>
#include <infiniband/verbs.h>

#define CUDA_CHECK(cmd) do { \
    cudaError_t err = cmd; \
    if (err != cudaSuccess) { \
        std::cerr << "CUDA Error: " << cudaGetErrorString(err) << " at line " << __LINE__ << std::endl; \
        exit(EXIT_FAILURE); \
    } \
} while(0)

struct GPUMemoryRegion {
    void* gpu_ptr;
    size_t size;
    struct ibv_mr* mr;
};

// Register a CUDA GPU memory region with InfiniBand HCA for GPUDirect RDMA
GPUMemoryRegion register_gpu_memory_with_hca(struct ibv_pd* pd, size_t size) {
    GPUMemoryRegion region;
    region.size = size;

    // 1. Allocate memory directly on the GPU (HBM3/GDDR6)
    CUDA_CHECK(cudaMalloc(&region.gpu_ptr, size));
    
    // Ensure all kernels are completed and memory is fully mapped
    CUDA_CHECK(cudaDeviceSynchronize());

    // 2. Access flags required for Tensor Parallelism All-Reduce:
    // Local Write (for receiving data locally), Remote Write (for peers writing to us), 
    // Remote Read (for peers reading from us).
    int access_flags = IBV_ACCESS_LOCAL_WRITE | 
                       IBV_ACCESS_REMOTE_WRITE | 
                       IBV_ACCESS_REMOTE_READ;

    // 3. Register the CUDA GPU pointer directly with the Protection Domain.
    // The nvidia-peermem kernel module handles translation of these virtual addresses
    // to physical GPU BAR1 addresses.
    region.mr = ibv_reg_mr(pd, region.gpu_ptr, size, access_flags);
    if (!region.mr) {
        std::cerr << "Fatal: Failed to register GPU memory region. Check if nvidia-peermem is loaded!" << std::endl;
        cudaFree(region.gpu_ptr);
        exit(EXIT_FAILURE);
    }

    std::cout << "Successfully registered GPUDirect RDMA Memory Region:" << std::endl;
    std::cout << "  GPU Virtual Address: " << region.gpu_ptr << std::endl;
    std::cout << "  Size: " << size << " bytes" << std::endl;
    std::cout << "  Local Key (lkey): 0x" << std::hex << region.mr->lkey << std::endl;
    std::cout << "  Remote Key (rkey): 0x" << std::hex << region.mr->rkey << std::dec << std::endl;

    return region;
}

void cleanup_gpu_memory(GPUMemoryRegion& region) {
    if (region.mr) {
        ibv_dereg_mr(region.mr);
    }
    if (region.gpu_ptr) {
        cudaFree(region.gpu_ptr);
    }
}

snippet-3.txt

#include <iostream>
#include <cstring>
#include <infiniband/verbs.h>

// Modifies the Queue Pair state to RTR (Ready to Receive) and RTS (Ready to Send)
// configured specifically for RoCEv2 (routed UDP encapsulation)
bool transition_qp_to_rts(struct ibv_qp* qp, 
                          uint32_t remote_qp_num, 
                          union ibv_gid remote_gid, 
                          uint8_t local_sgid_index, 
                          uint8_t traffic_class) {
    
    // 1. Transition to INIT
    struct ibv_qp_attr init_attr;
    std::memset(&init_attr, 0, sizeof(init_attr));
    init_attr.qp_state = IBV_QPS_INIT;
    init_attr.pkey_index = 0;
    init_attr.port_num = 1; // Physical port 1
    init_attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | 
                                IBV_ACCESS_REMOTE_WRITE | 
                                IBV_ACCESS_REMOTE_READ;

    int mask = IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS;
    if (ibv_modify_qp(qp, &init_attr, mask)) {
        std::cerr << "Failed to modify QP to INIT. errno: " << errno << std::endl;
        return false;
    }

    // 2. Transition to RTR (Ready to Receive)
    struct ibv_qp_attr rtr_attr;
    std::memset(&rtr_attr, 0, sizeof(rtr_attr));
    rtr_attr.qp_state = IBV_QPS_RTR;
    rtr_attr.path_mtu = IBV_MTU_4096; // Standard max payload for LLM tensors
    rtr_attr.dest_qp_num = remote_qp_num;
    rtr_attr.rq_psn = 0; // Starting packet sequence number
    rtr_attr.max_dest_rd_atomic = 1;
    rtr_attr.min_rnr_timer = 12; // 12 corresponds to 4.096ms delay before retry

    // Configure the Address Vector for RoCEv2 (L3 Routing)
    rtr_attr.ah_attr.is_global = 1; // Must be global (GRH enabled) for RoCEv2
    rtr_attr.ah_attr.port_num = 1;
    rtr_attr.ah_attr.sl = 0; // Service level (mapped to priority via configuration)
    rtr_attr.ah_attr.src_path_bits = 0;
    
    // Set Destination GID (IP address representation of remote node)
    std::memcpy(&rtr_attr.ah_attr.grh.dgid, &remote_gid, sizeof(remote_gid));
    
    // RoCEv2 properties in Global Routing Header (GRH)
    rtr_attr.ah_attr.grh.sgid_index = local_sgid_index; // Index of local RoCEv2 IP
    rtr_attr.ah_attr.grh.hop_limit = 64; // IP Time-To-Live
    
    // Crucial: Traffic class encodes the DSCP bits for switch-level Priority Flow Control
    // The format is: (DSCP << 2) | ECN. E.g., DSCP 26 (0x1A) -> 26 << 2 = 104 (0x68)
    rtr_attr.ah_attr.grh.traffic_class = traffic_class;

    mask = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | 
           IBV_QP_DEST_QPN | IBV_QP_RQ_PSN | 
           IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER;

    if (ibv_modify_qp(qp, &rtr_attr, mask)) {
        std::cerr << "Failed to modify QP to RTR. errno: " << errno << std::endl;
        return false;
    }

    // 3. Transition to RTS (Ready to Send)
    struct ibv_qp_attr rts_attr;
    std::memset(&rts_attr, 0, sizeof(rts_attr));
    rts_attr.qp_state = IBV_QPS_RTS;
    rts_attr.sq_psn = 0;
    rts_attr.timeout = 18; // Packet lifetime timeout: ~4s retry timer
    rts_attr.retry_cnt = 7; // Hardware Go-Back-N retries before failing
    rts_attr.rnr_retry = 7; // Receiver-Not-Ready retry count
    rts_attr.max_rd_atomic = 1;

    mask = IBV_QP_STATE | IBV_QP_SQ_PSN | IBV_QP_TIMEOUT | 
           IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY | IBV_QP_MAX_QP_RD_ATOMIC;

    if (ibv_modify_qp(qp, &rts_attr, mask)) {
        std::cerr << "Failed to modify QP to RTS. errno: " << errno << std::endl;
        return false;
    }

    std::cout << "QP transitioned successfully to RTS mode for RoCEv2." << std::endl;
    return true;
}

snippet-4.txt

#include <iostream>
#include <infiniband/verbs.h>
#include <cuda_runtime.h>

// Struct representing remote peer's memory region details exchanged out-of-band (e.g., over TCP control socket)
struct RemoteMR {
    uint64_t remote_addr;
    uint32_t rkey;
};

// Submits a zero-copy GPUDirect RDMA Write with Immediate to the remote GPU
bool post_gpudirect_rdma_write(struct ibv_qp* qp,
                               void* local_gpu_ptr, // Pointer to CUDA HBM3 memory
                               size_t size,
                               uint32_t local_lkey,
                               const RemoteMR& remote_peer,
                               uint32_t immediate_val) {
    
    // 1. Define Scatter/Gather Element (SGE) pointing directly to the CUDA buffer
    struct ibv_sge sge;
    std::memset(&sge, 0, sizeof(sge));
    sge.addr = reinterpret_cast<uint64_t>(local_gpu_ptr);
    sge.length = size;
    sge.lkey = local_lkey;

    // 2. Define Work Request (WR) for RDMA Write with Immediate
    struct ibv_send_wr wr;
    struct ibv_send_wr* bad_wr = nullptr;
    std::memset(&wr, 0, sizeof(wr));

    wr.wr_id = 0xDEADC0DE; // Identifier for completion handling
    wr.next = nullptr;
    wr.sg_list = &sge;
    wr.num_sge = 1;
    wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM; // Bypass receiver CPU, drop directly into memory, post immediate value to CQ
    wr.send_flags = IBV_SEND_SIGNALED;      // Ensure we get a local completion event
    
    // Set immediate value to notify the receiver
    wr.imm_data = htonl(immediate_val);

    // Target specifications (Remote virtual address and Rkey)
    wr.wr.rdma.remote_addr = remote_peer.remote_addr;
    wr.wr.rdma.rkey = remote_peer.rkey;

    // 3. Post the send request directly to the HCA's TX queue.
    // The HCA hardware will initiate a PCIe DMA read from the CUDA physical address
    // and stream it over the network to the peer node.
    int ret = ibv_post_send(qp, &wr, &bad_wr);
    if (ret != 0) {
        std::cerr << "Failed to post send work request. errno: " << ret << std::endl;
        return false;
    }

    return true;
}

// Helper to poll local completion queue to ensure GPU-to-GPU data transmission completed
bool wait_for_local_completion(struct ibv_cq* cq) {
    struct ibv_wc wc;
    int ne;
    do {
        // Poll completion queue
        ne = ibv_poll_cq(cq, 1, &wc);
    } while (ne == 0);

    if (ne < 0) {
        std::cerr << "Error polling completion queue." << std::endl;
        return false;
    }

    if (wc.status != IBV_WC_SUCCESS) {
        std::cerr << "Work request failed with status: " << ibv_wc_status_str(wc.status) << std::endl;
        return false;
    }

    return true;
}

snippet-6.sh

#!/usr/bin/env bash
# Production diagnostics script for troubleshooting RoCEv2 network and GPUDirect RDMA issues
set -euo pipefail

INTERFACE="eth0"
IB_DEVICE="mlx5_0"

echo "=== System Status: Verifying InfiniBand Device ==="
ibv_devinfo -d "$IB_DEVICE" | grep -E "hca_id|transport|state|phys_port_cnt"

echo "=== GID Table Check (Locating RoCEv2 GID index) ==="
# GID tables show supported protocols. RoCEv2 shows up as an IP address string
# in the GID table, and its type must be RoCE v2 (typically index 3 or 5)
show_gids | grep "$INTERFACE" || {
    echo "Warning: show_gids utility not found. Directly reading GID types from sysfs:"
    for i in {0..16}; do
        if [ -f "/sys/class/infiniband/${IB_DEVICE}/ports/1/gid_attrs/types/$i" ]; then
            TYPE=$(cat "/sys/class/infiniband/${IB_DEVICE}/ports/1/gid_attrs/types/$i" 2>/dev/null || true)
            IP=$(cat "/sys/class/infiniband/${IB_DEVICE}/ports/1/gids/$i" 2>/dev/null || true)
            if [ -n "$TYPE" ]; then
                echo "  Index $i: GID Type = $TYPE | IP/GID = $IP"
            fi
        fi
    done
}

echo "=== Checking PFC (Priority Flow Control) Pause Counters ==="
# If tx_pause or rx_pause on the designated priority (e.g. priority 3) is incrementing rapidly, 
# you have a network bottleneck causing buffer congestion.
ethtool -S "$INTERFACE" | grep -iE "pause|prio3" || echo "No explicit PFC counters visible via ethtool."

echo "=== Checking ECN and RoCE Congestion Counters ==="
# Look for CNP (Congestion Notification Packets) sent/received
HW_COUNTERS="/sys/class/infiniband/${IB_DEVICE}/ports/1/hw_counters"
if [ -d "$HW_COUNTERS" ]; then
    echo "RoCEv2 hardware-level congestion counters:"
    for counter in np_ecn_marked_roce np_cnp_sent rx_icrc_encap np_cnp_rx; do
        if [ -f "$HW_COUNTERS/$counter" ]; then
            printf "  %-25s : %s\n" "$counter" "$(cat "$HW_COUNTERS/$counter")"
        fi
    done
else
    echo "Hardware counters directory not found."
fi

echo "=== PCIe Topology: Checking GPUDirect Peer-to-Peer Capability ==="
# Verifies if PCIe Access Control Services (ACS) is enabled or if GPUs are grouped under a PCIe Switch.
# NVLink or PIX (direct PCIe switch bridge) is required for peer-to-peer mapping.
nvidia-smi topo -m